commit 5b45ce798280fc96215b418e58f215a3ad226e35 Author: wehub-resource-sync Date: Mon Jul 13 12:45:29 2026 +0800 chore: import upstream snapshot with attribution diff --git a/.cargo/audit.toml b/.cargo/audit.toml new file mode 100644 index 0000000..b3e9e6d --- /dev/null +++ b/.cargo/audit.toml @@ -0,0 +1,32 @@ +# Audit config file +# +# It may be located in the user home (`~/.cargo/audit.toml`) or in the project +# root (`.cargo/audit.toml`). +# +# All of the options which can be passed via CLI arguments can also be +# permanently specified in this file. + +[advisories] +ignore = [ + "RUSTSEC-2024-0436", # Paste used to generate macro, should be removed at some point. + "RUSTSEC-2025-0119", # `number_prefix` used by `tokenizers`, only in the examples. + "RUSTSEC-2025-0141", # `bincode` is no longer maintained. + "RUSTSEC-2024-0388", # `derivative` dependancy in the DQN example is unmaintained. + "RUSTSEC-2026-0194", # `quick-xml` <0.41.0, pinned by polars → object_store 0.13.2; no upgrade path yet. + "RUSTSEC-2026-0195", # `quick-xml` <0.41.0, pinned by polars → object_store 0.13.2; no upgrade path yet. +] # advisory IDs to ignore e.g. ["RUSTSEC-2019-0001", ...] +informational_warnings = [ + "unmaintained", +] # warn for categories of informational advisories +severity_threshold = "low" # CVSS severity ("none", "low", "medium", "high", "critical") + +# Output Configuration +[output] +deny = ["unmaintained"] # exit on error if unmaintained dependencies are found +format = "terminal" # "terminal" (human readable report) or "json" +quiet = false # Only print information on error +show_tree = true # Show inverse dependency trees along with advisories (default: true) + +[yanked] +enabled = true # Warn for yanked crates in Cargo.lock (default: true) +update_index = true # Auto-update the crates.io index (default: true) diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..82299c4 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,3 @@ +[alias] +xtask = "run --target-dir target/xtask --color always --package xtask --bin xtask --" +run-checks = "xtask -c all validate --release" \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 0000000..a9a2ac0 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,40 @@ +--- +name: Bug report +about: Create a report to help us improve +title: '' +labels: '' +assignees: '' + +--- + +**Describe the bug** + + +**To Reproduce** + + +**Expected behavior** + + +**Screenshots** + + +**Desktop (please complete the following information):** + - OS: [e.g. iOS] + - Browser [e.g. chrome, safari] + - Version [e.g. 22] + +**Smartphone (please complete the following information):** + - Device: [e.g. iPhone6] + - OS: [e.g. iOS8.1] + - Browser [e.g. stock browser, safari] + - Version [e.g. 22] + +**Additional context** + \ No newline at end of file diff --git a/.github/ISSUE_TEMPLATE/doc_request.md b/.github/ISSUE_TEMPLATE/doc_request.md new file mode 100644 index 0000000..21d4e4d --- /dev/null +++ b/.github/ISSUE_TEMPLATE/doc_request.md @@ -0,0 +1,10 @@ +--- +name: Documentation request +about: Flag incoherent or missing documentation, including use case examples. +title: '' +labels: '' +assignees: '' + +--- + + diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md new file mode 100644 index 0000000..fbc095d --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.md @@ -0,0 +1,28 @@ +--- +name: Feature request +about: Suggest an idea for this project +title: '' +labels: '' +assignees: '' + +--- + + + +### Feature description + + + +### Feature motivation + + + +### (Optional) Suggest a Solution + + diff --git a/.github/PULL_REQUEST_TEMPLATE/template.md b/.github/PULL_REQUEST_TEMPLATE/template.md new file mode 100644 index 0000000..55d6298 --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE/template.md @@ -0,0 +1,12 @@ +* **Please check if the PR fulfills these requirements** +- [ ] The commit message follows our guidelines +- [ ] Docs have been added / updated (for bug fixes / features) + + +* **What kind of change does this PR introduce?** (Bug fix, feature, docs update, ...) + + +* **Does this PR introduce a breaking change?** (What changes might users need to make in their application due to this PR?) + + +* **Other information**: \ No newline at end of file diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..be20d25 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,21 @@ +version: 2 + +updates: + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "daily" + ignore: + - dependency-name: "tracel-ai/github-actions*" + + - package-ecosystem: "cargo" + directories: + - "/" + - "crates/burn" + - "crates/burn-*" + - "crates/burn-import/*-tests" + - "examples/*" + - "xtask" + schedule: + interval: "weekly" + diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 0000000..aa1a8a6 --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,18 @@ +## Pull Request Template + +### Checklist + +- [ ] Confirmed that `cargo run-checks` command has been executed. +- [ ] Made sure the book is up to date with changes in this PR. + +### Related Issues/PRs + +_Provide links to relevant issues and dependent PRs._ + +### Changes + +_Summarize the problem being addressed and your solution._ + +### Testing + +_Describe how these changes have been tested._ diff --git a/.github/workflows/combine-dependabot-prs.yml b/.github/workflows/combine-dependabot-prs.yml new file mode 100644 index 0000000..72850a9 --- /dev/null +++ b/.github/workflows/combine-dependabot-prs.yml @@ -0,0 +1,21 @@ +name: Combine Dependabot PRs + +on: + schedule: + - cron: '0 6 * * MON' # Monday at 6:00am UTC + workflow_dispatch: + +permissions: + contents: write + pull-requests: write + checks: read + +jobs: + combine-prs: + runs-on: ubuntu-latest + steps: + - name: combine-prs + id: combine-prs + uses: github/combine-prs@v5.2.0 + with: + labels: dependencies,automated diff --git a/.github/workflows/dependencies.yml b/.github/workflows/dependencies.yml new file mode 100644 index 0000000..fbfe433 --- /dev/null +++ b/.github/workflows/dependencies.yml @@ -0,0 +1,55 @@ +name: dependencies + +on: + schedule: + - cron: '0 21 * * TUE' # Run every Tuesday at 21:00 (UTC) + push: + tags: + - 'v*.*.*' # Run when a new version is being published + +env: + UDEPS_VERSION: "0.1.57" + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + dependencies: + runs-on: ubuntu-latest + strategy: + matrix: + checks: + - licenses + - bans sources + continue-on-error: ${{ matrix.checks == 'licenses' }} # failed licenses don't abort + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Audit Rust dependencies + # If a vulnerability is found, a new issue will automatically be opened + # since this action runs on main branch + uses: actions-rust-lang/audit@v1 + # -------------------------------------------------------------------------------- + - name: Detect multiple versions of the same crate + uses: EmbarkStudios/cargo-deny-action@v2 + with: + command: check ${{ matrix.checks }} + # -------------------------------------------------------------------------------- + - name: Install Rust nightly + uses: dtolnay/rust-toolchain@nightly + with: + toolchain: nightly + components: rustfmt + # -------------------------------------------------------------------------------- + - name: Install cargo-udeps + env: + UDEPS_LINK: https://github.com/est31/cargo-udeps/releases/download + run: | + curl -L "$UDEPS_LINK/v$UDEPS_VERSION/cargo-udeps-v$UDEPS_VERSION-x86_64-unknown-linux-gnu.tar.gz" | + tar xz -C $HOME/.cargo/bin --strip-components 2 + # -------------------------------------------------------------------------------- + - name: Run cargo-udeps + run: | + cargo +nightly udeps --all-targets diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..6bce54e --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,443 @@ +name: publish + +on: + push: + tags: + - "v*" + workflow_dispatch: + inputs: + dry-run-only: + description: "Run xtask publish in dry-run mode (no publish)" + type: boolean + required: false + default: false + +jobs: + publish-burn-rl: + needs: + - publish-burn-core + - publish-burn-optim + # dev dependencies + - publish-burn-flex + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-rl + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-vision: + needs: + - publish-burn-autodiff + - publish-burn-candle + - publish-burn-fusion + - publish-burn-cubecl-fusion + - publish-burn-cubecl + - publish-burn-flex + - publish-burn-tch + - publish-burn-core + - publish-burn-ir + - publish-burn-store + # dev dependencies + - publish-burn-wgpu + - publish-burn-cuda + - publish-burn-backend-extension + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-vision + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-router: + needs: + - publish-burn-ir + - publish-burn-std + - publish-burn-tensor + # dev dependencies + - publish-burn-autodiff + - publish-burn-flex + - publish-burn-wgpu + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-router + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-remote: + needs: + - publish-burn-ir + - publish-burn-std + - publish-burn-tensor + - publish-burn-router + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-remote + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-derive: + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-derive + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-dataset: + needs: + - publish-burn-std + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-dataset + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-std: + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-std + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-backend-extension: + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-backend-extension + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-tensor: + needs: + - publish-burn-std + - publish-burn-backend + - publish-burn-dispatch + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-tensor + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-backend: + needs: + - publish-burn-std + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-backend + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-ir: + needs: + - publish-burn-tensor + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-ir + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-fusion: + needs: + - publish-burn-ir + - publish-burn-tensor + - publish-burn-std + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-fusion + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-cubecl-fusion: + needs: + - publish-burn-ir + - publish-burn-std + - publish-burn-fusion + - publish-burn-tensor + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-cubecl-fusion + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-cubecl: + needs: + - publish-burn-ir + - publish-burn-std + - publish-burn-fusion + - publish-burn-cubecl-fusion + - publish-burn-tensor + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-cubecl + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-autodiff: + needs: + - publish-burn-backend + - publish-burn-std + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-autodiff + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-tch: + needs: + - publish-burn-tensor + - publish-burn-autodiff + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-tch + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-ndarray: + needs: + - publish-burn-ir + - publish-burn-tensor + - publish-burn-autodiff + - publish-burn-std + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-ndarray + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-flex: + needs: + - publish-burn-backend + - publish-burn-ir + - publish-burn-std + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-flex + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-wgpu: + needs: + - publish-burn-tensor + - publish-burn-autodiff + - publish-burn-std + - publish-burn-cubecl + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-wgpu + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-cpu: + needs: + - publish-burn-tensor + - publish-burn-fusion + - publish-burn-cubecl + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-cpu + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-cuda: + needs: + - publish-burn-tensor + - publish-burn-autodiff + - publish-burn-std + - publish-burn-cubecl + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-cuda + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-rocm: + needs: + - publish-burn-tensor + - publish-burn-autodiff + - publish-burn-std + - publish-burn-cubecl + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-rocm + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-candle: + needs: + - publish-burn-tensor + - publish-burn-autodiff + - publish-burn-tch + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-candle + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + # publish-burn-collective: + # needs: + # - publish-burn-std + # - publish-burn-tensor + # - publish-burn-communication + # # dev dependencies + # - publish-burn-wgpu + # - publish-burn-flex + # - publish-burn-cuda + # uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + # with: + # crate: burn-collective + # dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + # secrets: + # CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-communication: + needs: + - publish-burn-std + - publish-burn-tensor + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-communication + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-core: + needs: + - publish-burn-dataset + - publish-burn-std + - publish-burn-derive + - publish-burn-tensor + # dev dependencies + - publish-burn-autodiff + - publish-burn-wgpu + - publish-burn-tch + - publish-burn-cuda + - publish-burn-flex + - publish-burn-candle + - publish-burn-remote + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-core + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-nn: + needs: + - publish-burn-core + # dev dependencies + - publish-burn-autodiff + - publish-burn-wgpu + - publish-burn-tch + - publish-burn-flex + - publish-burn-candle + - publish-burn-remote + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-nn + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-optim: + needs: + - publish-burn-core + - publish-burn-collective + # dev dependencies + - publish-burn-autodiff + - publish-burn-wgpu + - publish-burn-tch + - publish-burn-flex + - publish-burn-candle + - publish-burn-remote + - publish-burn-nn + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-optim + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-train: + needs: + - publish-burn-core + - publish-burn-optim + - publish-burn-collective + - publish-burn-rl + - publish-burn-flex + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-train + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-dispatch: + needs: + - publish-burn-std + - publish-burn-backend + - publish-burn-backend-extension + - publish-burn-autodiff + - publish-burn-cpu + - publish-burn-cuda + - publish-burn-rocm + - publish-burn-wgpu + - publish-burn-flex + - publish-burn-ndarray + - publish-burn-tch + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-dispatch + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn: + needs: + - publish-burn-core + - publish-burn-nn + - publish-burn-optim + - publish-burn-collective + - publish-burn-store + - publish-burn-train + - publish-burn-cpu + - publish-burn-dispatch + - publish-burn-flex + - publish-burn-vision + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} + + publish-burn-store: + needs: + - publish-burn-core + - publish-burn-nn + - publish-burn-tensor + uses: tracel-ai/github-actions/.github/workflows/publish-crate.yml@v9 + with: + crate: burn-store + dry-run-only: ${{ github.event_name == 'workflow_dispatch' && inputs.dry-run-only || false }} + secrets: + CRATES_IO_API_TOKEN: ${{ secrets.CRATES_IO_API_TOKEN }} diff --git a/.github/workflows/stale-pr.yml b/.github/workflows/stale-pr.yml new file mode 100644 index 0000000..56e23ee --- /dev/null +++ b/.github/workflows/stale-pr.yml @@ -0,0 +1,41 @@ +name: Stale Pull Requests + +on: + schedule: + - cron: '0 12 * * *' # Run every day at 12:00 (UTC) + +# The minimum permissions required to run this Action +permissions: + contents: write # only for delete-branch option + issues: write + pull-requests: write + +jobs: + stale-pr: + runs-on: ubuntu-latest + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Stale pull requests + uses: actions/stale@v10 + with: + # The idle number of days before marking issues stale. + # + # With a negative number like -1, no issues + # will be marked as stale automatically. + days-before-issue-stale: -1 + # The idle number of days before marking pull requests stale + days-before-pr-stale: 30 + # The idle number of days before closing + # the stale pull requests (due to the stale label). + # + # With a negative number like -1, the pull requests + # will never be closed automatically. + days-before-pr-close: -1 + # Label to apply on staled pull requests + stale-pr-label: 'stale' + # The message that will be added as a comment to the pull request + stale-pr-message: 'This PR has been marked as stale because it has not been updated for over a month' + # Remove `stale` label from pull requests on updates/comments + remove-pr-stale-when-updated: true diff --git a/.github/workflows/test-gpu.yml b/.github/workflows/test-gpu.yml new file mode 100644 index 0000000..811ae27 --- /dev/null +++ b/.github/workflows/test-gpu.yml @@ -0,0 +1,162 @@ +name: CI GPU + +on: + workflow_dispatch: + inputs: + pr_number: + description: "Number of the pull request that triggers this run if any" + type: number + required: false + +# important to set the run name to this format so that the CI server +# can track the PR number from the workflow_run events. +run-name: ${{ github.workflow }}:${{ github.repository }}#${{ inputs.pr_number }} + +env: + # Note: It is not possible to define top level env vars and pass them to composite actions. + # To work around this issue we use inputs and define all the env vars here. + + RUST_PREVIOUS_VERSION: 1.92.0 + + # Dependency versioning + # from wgpu repo: https://github.com/gfx-rs/wgpu/blob/trunk/.github/workflows/ci.yml + + # GCP runners + GCP_RUNNERS_IMAGE_FAMILY: "tracel-ci-ubuntu-2404-amd64-nvidia" + GCP_RUNNERS_MACHINE_TYPE: "g2-standard-4" + GCP_RUNNERS_ZONE: "us-east1-c" + + # Test in release mode (make it an empty string to test in debug mode) + TEST_RELEASE_FLAG: "--release" + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + prepare-checks: + runs-on: ubuntu-latest + outputs: + rust-prev-version: ${{ env.RUST_PREVIOUS_VERSION }} + gcp_runners_image_family: ${{ env.GCP_RUNNERS_IMAGE_FAMILY }} + gcp_runners_machine_type: ${{ env.GCP_RUNNERS_MACHINE_TYPE }} + gcp_runners_zone: ${{ env.GCP_RUNNERS_ZONE }} + steps: + - name: Do Nothing + if: false + run: echo + + linux-std-cuda-tests: + needs: [prepare-checks] + timeout-minutes: 60 + # '@id:' label must be unique within this worklow + runs-on: + [ + "@id:burn-cuda-job-${{github.run_id}}-${{github.run_attempt}}", + "@pr_number:${{ inputs.pr_number }}", + "@organization:tracel-ai", + "@repository:burn", + "@image-family:${{ needs.prepare-checks.outputs.gcp_runners_image_family }}", + "@machine-type:${{ needs.prepare-checks.outputs.gcp_runners_machine_type }}", + "@zones:${{ needs.prepare-checks.outputs.gcp_runners_zone }}", + "@gpu:true", + ] + env: + LD_LIBRARY_PATH: "/usr/local/cuda/lib64" + # disable incremental compilation (reduces artifact size) + CARGO_PROFILE_TEST_INCREMENTAL: "false" + # Keep the stragegy to be able to easily add new rust versions if required + strategy: + matrix: + rust: [stable] + include: + - rust: stable + toolchain: stable + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Install Rust + uses: tracel-ai/github-actions/install-rust@v9 + with: + rust-toolchain: ${{ matrix.toolchain }} + enable-cache: false + # -------------------------------------------------------------------------------- + - name: Tests (burn-cuda) + run: cargo xtask test ${{ env.TEST_RELEASE_FLAG }} --ci gcp-cuda-runner + + linux-std-vulkan-tests: + needs: [prepare-checks] + timeout-minutes: 60 + # '@id:' label must be unique within this worklow + runs-on: + [ + "@id:burn-vulkan-job-${{github.run_id}}-${{github.run_attempt}}", + "@pr_number:${{ inputs.pr_number }}", + "@organization:tracel-ai", + "@repository:burn", + "@image-family:${{ needs.prepare-checks.outputs.gcp_runners_image_family }}", + "@machine-type:${{ needs.prepare-checks.outputs.gcp_runners_machine_type }}", + "@zones:${{ needs.prepare-checks.outputs.gcp_runners_zone }}", + "@gpu:true", + ] + env: + # disable incremental compilation (reduces artifact size) + CARGO_PROFILE_TEST_INCREMENTAL: "false" + # Keep the stragegy to be able to easily add new rust versions if required + strategy: + matrix: + rust: [stable] + include: + - rust: stable + toolchain: stable + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Setup Rust + uses: tracel-ai/github-actions/install-rust@v9 + with: + rust-toolchain: ${{ matrix.toolchain }} + enable-cache: false + # -------------------------------------------------------------------------------- + - name: Tests (burn-vulkan) + run: cargo xtask test ${{ env.TEST_RELEASE_FLAG }} --ci gcp-vulkan-runner + + linux-std-wgpu-tests: + needs: [prepare-checks] + timeout-minutes: 60 + # '@id:' label must be unique within this worklow + runs-on: + [ + "@id:burn-wgpu-job-${{github.run_id}}-${{github.run_attempt}}", + "@pr_number:${{ inputs.pr_number }}", + "@organization:tracel-ai", + "@repository:burn", + "@image-family:${{ needs.prepare-checks.outputs.gcp_runners_image_family }}", + "@machine-type:${{ needs.prepare-checks.outputs.gcp_runners_machine_type }}", + "@zones:${{ needs.prepare-checks.outputs.gcp_runners_zone }}", + "@gpu:true", + ] + env: + # disable incremental compilation (reduces artifact size) + CARGO_PROFILE_TEST_INCREMENTAL: "false" + # Keep the stragegy to be able to easily add new rust versions if required + strategy: + matrix: + rust: [stable] + include: + - rust: stable + toolchain: stable + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Setup Rust + uses: tracel-ai/github-actions/install-rust@v9 + with: + rust-toolchain: ${{ matrix.toolchain }} + enable-cache: false + # -------------------------------------------------------------------------------- + - name: Tests (burn-wgpu) + run: cargo xtask test ${{ env.TEST_RELEASE_FLAG }} --ci gcp-wgpu-runner diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..bac7ea8 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,278 @@ +name: CI + +on: + push: + branches: + - main + paths: + - "Cargo.lock" + - "**.rs" + - "**.sh" + - "**.ps1" + - "**.yml" + - "**.toml" + - "!**.md" + - "!LICENSE-APACHE" + - "!LICENSE-MIT" + pull_request: + types: [opened, synchronize] + paths: + - "Cargo.lock" + - "**.rs" + - "**.sh" + - "**.ps1" + - "**.yml" + - "**.toml" + - "!**.md" + - "!LICENSE-APACHE" + - "!LICENSE-MIT" + +env: + # Note: It is not possible to define top level env vars and pass them to composite actions. + # To work around this issue we use inputs and define all the env vars here. + + RUST_PREVIOUS_VERSION: 1.95.0 + + # Dependency versioning + # from wgpu repo: https://github.com/gfx-rs/wgpu/blob/trunk/.github/workflows/ci.yml + + # Mozilla Grcov + GRCOV_LINK: "https://github.com/mozilla/grcov/releases/download" + GRCOV_VERSION: "0.8.19" + + # Test in release mode (make it an empty string to test in debug mode) + TEST_RELEASE_FLAG: "--release" + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + prepare-checks: + runs-on: ubuntu-latest + outputs: + rust-prev-version: ${{ env.RUST_PREVIOUS_VERSION }} + steps: + - name: Do Nothing + if: false + run: echo + + code-quality: + runs-on: ubuntu-22.04 + needs: prepare-checks + strategy: + matrix: + rust: [stable] + include: + - rust: stable + toolchain: stable + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Setup Rust + uses: tracel-ai/github-actions/install-rust@v9 + with: + rust-toolchain: ${{ matrix.toolchain }} + cache-key: ${{ matrix.rust }}-linux + # -------------------------------------------------------------------------------- + - name: Audit + run: cargo xtask check audit + # -------------------------------------------------------------------------------- + - name: Format + shell: bash + env: + # work around for colors + # see: https://github.com/rust-lang/rustfmt/issues/3385 + TERM: xterm-256color + run: cargo xtask check format + # -------------------------------------------------------------------------------- + - name: Lint + run: cargo xtask check lint + # -------------------------------------------------------------------------------- + - name: Typos + uses: tracel-ai/github-actions/check-typos@v9 + + documentation: + runs-on: ubuntu-22.04 + needs: prepare-checks + strategy: + matrix: + rust: [stable] + include: + - rust: stable + toolchain: stable + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Setup Rust + uses: tracel-ai/github-actions/install-rust@v9 + with: + rust-toolchain: ${{ matrix.toolchain }} + cache-key: ${{ matrix.rust }}-linux + # -------------------------------------------------------------------------------- + - name: Documentation Build + run: cargo xtask doc build + # -------------------------------------------------------------------------------- + - name: Documentation Tests + run: cargo xtask doc tests + + linux-std-tests: + runs-on: ubuntu-22.04 + needs: [prepare-checks, code-quality] + env: + DISABLE_WGPU_SPIRV: "1" + # disable incremental compilation (reduces artifact size) + CARGO_PROFILE_TEST_INCREMENTAL: "false" + CARGO_TERM_COLOR: always + strategy: + matrix: + rust: [stable, prev] + # Tests are split across 3 shards: backend tests, workspace crates, examples/* + ci_type: [backends, crates, examples] + include: + - rust: stable + toolchain: stable + coverage: --enable-coverage + - rust: prev + toolchain: ${{ needs.prepare-checks.outputs.rust-prev-version }} + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Setup Rust + uses: tracel-ai/github-actions/install-rust@v9 + with: + rust-toolchain: ${{ matrix.toolchain }} + # Keep caches isolated per shard to prevent race conditions or bloated sizes + cache-key: ${{ matrix.rust }}-${{ matrix.ci_type }}-linux + # Disable cache on linux-std (stable) runner which currently always runs out of disk space with tests + coverage + enable-cache: ${{ matrix.rust != 'stable' }} + # # -------------------------------------------------------------------------------- + - name: Install grcov + if: matrix.rust == 'stable' + shell: bash + run: | + curl -L "$GRCOV_LINK/v$GRCOV_VERSION/grcov-x86_64-unknown-linux-musl.tar.bz2" | + tar xj -C $HOME/.cargo/bin + cargo xtask coverage install + # -------------------------------------------------------------------------------- + - name: Tests + run: cargo xtask ${{ matrix.coverage }} test ${{ env.TEST_RELEASE_FLAG }} --ci ${{ matrix.ci_type }} + # -------------------------------------------------------------------------------- + - name: Generate lcov.info + if: matrix.rust == 'stable' && matrix.ci_type != 'examples' + # /* is to exclude std library code coverage from analysis + run: cargo xtask coverage generate --ignore "/*,xtask/*,examples/*" --profile release + # -------------------------------------------------------------------------------- + - name: Upload Codecov Shard Artifact + if: matrix.rust == 'stable' && matrix.ci_type != 'examples' + uses: actions/upload-artifact@v7 + with: + # Names must be unique so they don't overwrite each other in the storage backend + name: coverage-linux-${{ matrix.ci_type }} + path: lcov.info + +# --- Dedicated coverage upload job --- + upload-coverage: + runs-on: ubuntu-latest + needs: linux-std-tests + steps: + - name: checkout + uses: actions/checkout@v7 + + - name: Download All Coverage Artifacts + uses: actions/download-artifact@v8 + with: + # Finds all artifacts prefixed with coverage-linux-* + pattern: coverage-linux-* + path: coverage-reports + + - name: Codecov upload + uses: codecov/codecov-action@v7 + with: + # Codecov automatically discovers and merges multiple lcov files inside this directory + directory: coverage-reports + token: ${{ secrets.CODECOV_TOKEN }} + + linux-no-std-tests: + runs-on: ubuntu-22.04 + needs: [prepare-checks, code-quality] + strategy: + matrix: + rust: [stable, prev] + include: + - rust: stable + toolchain: stable + - rust: prev + toolchain: ${{ needs.prepare-checks.outputs.rust-prev-version }} + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Setup Rust + uses: tracel-ai/github-actions/install-rust@v9 + with: + rust-toolchain: ${{ matrix.toolchain }} + cache-key: ${{ matrix.rust }}-linux-no-std + # -------------------------------------------------------------------------------- + - name: Crates Build + run: cargo xtask --context no-std build --ci + # -------------------------------------------------------------------------------- + - name: Crates Tests + run: cargo xtask --context no-std test ${{ env.TEST_RELEASE_FLAG }} --ci github-runner + + windows-std-tests: + runs-on: windows-2022 + needs: [prepare-checks, code-quality] + env: + CARGO_PROFILE_TEST_INCREMENTAL: "false" + CARGO_TERM_COLOR: always + strategy: + matrix: + rust: [stable] + # Tests are split across 3 shards: backend tests, workspace crates, examples/* + ci_type: [backends, crates, examples] + include: + - rust: stable + toolchain: stable + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Setup Rust + uses: tracel-ai/github-actions/install-rust@v9 + with: + rust-toolchain: ${{ matrix.toolchain }} + cache-key: ${{ matrix.rust }}-${{ matrix.ci_type }}-windows + # -------------------------------------------------------------------------------- + - name: Tests + run: cargo xtask test ${{ env.TEST_RELEASE_FLAG }} --ci ${{ matrix.ci_type }} + + macos-std-tests: + runs-on: blaze/macos-15 + needs: [prepare-checks, code-quality] + timeout-minutes: 60 + # Keep the stragegy to be able to easily add new rust versions if required + strategy: + matrix: + rust: [stable] + include: + - rust: stable + toolchain: stable + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Setup Rust + uses: tracel-ai/github-actions/install-rust@v9 + with: + rust-toolchain: ${{ matrix.toolchain }} + cache-key: ${{ matrix.rust }}-macos + # -------------------------------------------------------------------------------- + - name: Device check + run: system_profiler SPHardwareDataType + # -------------------------------------------------------------------------------- + - name: Tests + run: cargo xtask test ${{ env.TEST_RELEASE_FLAG }} --ci github-mac-runner diff --git a/.github/workflows/valgrind.yml b/.github/workflows/valgrind.yml new file mode 100644 index 0000000..516f932 --- /dev/null +++ b/.github/workflows/valgrind.yml @@ -0,0 +1,39 @@ +name: valgrind + +on: + schedule: + - cron: '0 23 * * WED' # Run every Wednesday at 23:00 (UTC) + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + valgrind: + runs-on: [ + '@id:burn-linux-valgrind-${{ github.run_id }}-${{ github.run_attempt }}', + '@image-family:ubuntu-2404-lts-amd64', + '@image-project:ubuntu-os-cloud', + '@disk-size:100', + '@keep-alive:false', + '@machine-type:n2-standard-16', + '@os:linux', + '@zones:northamerica-northeast1-b' + ] + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Install Mesa + uses: tracel-ai/github-actions/install-mesa@v9 + # -------------------------------------------------------------------------------- + - name: Install valgrind + run: | + sudo apt-get install valgrind + # -------------------------------------------------------------------------------- + - name: Run cargo-valgrind + env: + CARGO_TARGET_X86_64_UNKNOWN_LINUX_GNU_RUNNER: "valgrind -s --leak-check=full --show-leak-kinds=all --error-exitcode=1" + # Looking for vulnerabilities + run: | + cargo test diff --git a/.github/workflows/vulnerabilities.yml b/.github/workflows/vulnerabilities.yml new file mode 100644 index 0000000..dd80e2c --- /dev/null +++ b/.github/workflows/vulnerabilities.yml @@ -0,0 +1,131 @@ +name: vulnerabilities + +on: + schedule: + - cron: '0 21 * * WED' # Run every Wednesday at 21:00 (UTC) + push: + tags: + - 'v*.*.*' + +env: + CAREFUL_VERSION: "0.4.9" + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + cargo-careful: + runs-on: ubuntu-latest + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Install Rust nightly + uses: dtolnay/rust-toolchain@nightly + with: + toolchain: nightly + components: rustfmt, rust-src + # -------------------------------------------------------------------------------- + - name: Install Mesa + uses: tracel-ai/github-actions/install-mesa@v9 + # -------------------------------------------------------------------------------- + - name: Install cargo-careful + env: + CAREFUL_LINK: https://github.com/RalfJung/cargo-careful/releases/download + run: | + curl -L "$CAREFUL_LINK/v$CAREFUL_VERSION/cargo-careful.x86_64-unknown-linux-musl" \ + --output $HOME/.cargo/bin/cargo-careful + chmod +x $HOME/.cargo/bin/cargo-careful + # -------------------------------------------------------------------------------- + - name: Run cargo-careful + # Looking for undefined behaviours + run: cargo +nightly careful test + + address-sanitizer: + runs-on: ubuntu-latest + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Install Rust nightly + uses: dtolnay/rust-toolchain@nightly + with: + toolchain: nightly + components: rustfmt, rust-src + # -------------------------------------------------------------------------------- + - name: Install Mesa + uses: tracel-ai/github-actions/install-mesa@v9 + # -------------------------------------------------------------------------------- + - name: Run AddressSanitizer + env: + RUSTFLAGS: -Zsanitizer=address -Copt-level=3 + RUSTDOCFLAGS: -Zsanitizer=address + # Looking for memory vulnerabilities + run: cargo test -Zbuild-std --target x86_64-unknown-linux-gnu -- --nocapture + + thread-sanitizer: + runs-on: ubuntu-latest + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Install Rust nightly + uses: dtolnay/rust-toolchain@nightly + with: + toolchain: nightly + components: rustfmt, rust-src + # -------------------------------------------------------------------------------- + - name: Install Mesa + uses: tracel-ai/github-actions/install-mesa@v9 + # -------------------------------------------------------------------------------- + - name: Run ThreadSanitizer + env: + RUSTFLAGS: -Zsanitizer=thread -Copt-level=3 + RUSTDOCFLAGS: -Zsanitizer=thread + # Looking for data race among threads + run: cargo test -Zbuild-std --target x86_64-unknown-linux-gnu -- --nocapture + + memory-sanitizer: + runs-on: ubuntu-latest + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Install Rust nightly + uses: dtolnay/rust-toolchain@nightly + with: + toolchain: nightly + components: rustfmt, rust-src + # -------------------------------------------------------------------------------- + - name: Install Mesa + uses: tracel-ai/github-actions/install-mesa@v9 + # -------------------------------------------------------------------------------- + - name: Run MemorySanitizer + env: + RUSTFLAGS: -Zsanitizer=memory -Zsanitizer-memory-track-origins -Copt-level=3 + RUSTDOCFLAGS: -Zsanitizer=memory -Zsanitizer-memory-track-origins + # Looking for unitialized memory. + run: cargo test -Zbuild-std --target x86_64-unknown-linux-gnu -- --nocapture + + safe-stack: + runs-on: ubuntu-latest + steps: + - name: checkout + uses: actions/checkout@v7 + # -------------------------------------------------------------------------------- + - name: Install Rust nightly + uses: dtolnay/rust-toolchain@nightly + with: + toolchain: nightly + components: rustfmt, rust-src + # -------------------------------------------------------------------------------- + - name: Install Mesa + uses: tracel-ai/github-actions/install-mesa@v9 + # -------------------------------------------------------------------------------- + - name: Run SafeStack + env: + RUSTFLAGS: -Zsanitizer=safestack -Copt-level=3 + RUSTDOCFLAGS: -Zsanitizer=safestack + # Provides backward edge control flow protection + run: cargo test -Zbuild-std --target x86_64-unknown-linux-gnu -- --nocapture diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..a86740c --- /dev/null +++ b/.gitignore @@ -0,0 +1,27 @@ +target +# These are backup files generated by rustfmt +**/*.rs.bk +.DS_Store + +.dir-locals.el +.idea +.vscode +.vs +.fleet +.ipynb_checkpoints/ + +# Build output directory +out + +# Virtual Environment of Python +.venv +uv.lock + +# Nix direnv +.envrc +.direnv + +# tags files +tags + +examples/**/Cargo.lock diff --git a/CITATION.cff b/CITATION.cff new file mode 100644 index 0000000..72c7271 --- /dev/null +++ b/CITATION.cff @@ -0,0 +1,34 @@ +cff-version: 1.2.0 +message: "If you use this software, please cite it as below." +authors: + - family-names: "Simard" + given-names: "Nathaniel" + email: "nathaniel.simard.42@gmail.com" + - family-names: "Fortier-Dubois" + given-names: "Louis" + email: "louisfd94@gmail.com" + - family-names: "Tadjibaev" + given-names: "Dilshod" + email: "dilshod@gmail.com" + - family-names: "Lagrange" + given-names: "Guillaume" + email: "lagrange.guillaume.1@gmail.com" + - name: "Burn Framework Contributors" +title: "Burn" +version: 0.21.0 +date-released: 2026-05-07 +url: "https://burn.dev/" +repository-code: "https://github.com/tracel-ai/burn" +license: + - MIT + - Apache-2.0 +abstract: "Burn is a new comprehensive dynamic Deep Learning Framework built using Rust with extreme flexibility, compute efficiency and portability as its primary goals." +keywords: + - scientific-computing + - deep-learning + - machine-learning + - neural-networks + - rust + - high-performance-computing + - portability + - compute-efficiency diff --git a/CODE-OF-CONDUCT.md b/CODE-OF-CONDUCT.md new file mode 100644 index 0000000..0c14ff7 --- /dev/null +++ b/CODE-OF-CONDUCT.md @@ -0,0 +1,128 @@ +# Contributor Covenant Code of Conduct + +## Our Pledge + +We as members, contributors, and leaders pledge to make participation in our +community a harassment-free experience for everyone, regardless of age, body +size, visible or invisible disability, ethnicity, sex characteristics, gender +identity and expression, level of experience, education, socio-economic status, +nationality, personal appearance, race, religion, or sexual identity +and orientation. + +We pledge to act and interact in ways that contribute to an open, welcoming, +diverse, inclusive, and healthy community. + +## Our Standards + +Examples of behavior that contributes to a positive environment for our +community include: + +* Demonstrating empathy and kindness toward other people +* Being respectful of differing opinions, viewpoints, and experiences +* Giving and gracefully accepting constructive feedback +* Accepting responsibility and apologizing to those affected by our mistakes, + and learning from the experience +* Focusing on what is best not just for us as individuals, but for the + overall community + +Examples of unacceptable behavior include: + +* The use of sexualized language or imagery, and sexual attention or + advances of any kind +* Trolling, insulting or derogatory comments, and personal or political attacks +* Public or private harassment +* Publishing others' private information, such as a physical or email + address, without their explicit permission +* Other conduct which could reasonably be considered inappropriate in a + professional setting + +## Enforcement Responsibilities + +Community leaders are responsible for clarifying and enforcing our standards of +acceptable behavior and will take appropriate and fair corrective action in +response to any behavior that they deem inappropriate, threatening, offensive, +or harmful. + +Community leaders have the right and responsibility to remove, edit, or reject +comments, commits, code, wiki edits, issues, and other contributions that are +not aligned to this Code of Conduct, and will communicate reasons for moderation +decisions when appropriate. + +## Scope + +This Code of Conduct applies within all community spaces, and also applies when +an individual is officially representing the community in public spaces. +Examples of representing our community include using an official e-mail address, +posting via an official social media account, or acting as an appointed +representative at an online or offline event. + +## Enforcement + +Instances of abusive, harassing, or otherwise unacceptable behavior may be +reported to the community leaders responsible for enforcement at +nathaniel.simard.42@gmail.com. +All complaints will be reviewed and investigated promptly and fairly. + +All community leaders are obligated to respect the privacy and security of the +reporter of any incident. + +## Enforcement Guidelines + +Community leaders will follow these Community Impact Guidelines in determining +the consequences for any action they deem in violation of this Code of Conduct: + +### 1. Correction + +**Community Impact**: Use of inappropriate language or other behavior deemed +unprofessional or unwelcome in the community. + +**Consequence**: A private, written warning from community leaders, providing +clarity around the nature of the violation and an explanation of why the +behavior was inappropriate. A public apology may be requested. + +### 2. Warning + +**Community Impact**: A violation through a single incident or series +of actions. + +**Consequence**: A warning with consequences for continued behavior. No +interaction with the people involved, including unsolicited interaction with +those enforcing the Code of Conduct, for a specified period of time. This +includes avoiding interactions in community spaces as well as external channels +like social media. Violating these terms may lead to a temporary or +permanent ban. + +### 3. Temporary Ban + +**Community Impact**: A serious violation of community standards, including +sustained inappropriate behavior. + +**Consequence**: A temporary ban from any sort of interaction or public +communication with the community for a specified period of time. No public or +private interaction with the people involved, including unsolicited interaction +with those enforcing the Code of Conduct, is allowed during this period. +Violating these terms may lead to a permanent ban. + +### 4. Permanent Ban + +**Community Impact**: Demonstrating a pattern of violation of community +standards, including sustained inappropriate behavior, harassment of an +individual, or aggression toward or disparagement of classes of individuals. + +**Consequence**: A permanent ban from any sort of public interaction within +the community. + +## Attribution + +This Code of Conduct is adapted from the [Contributor Covenant][homepage], +version 2.0, available at +https://www.contributor-covenant.org/version/2/0/code_of_conduct.html. + +Community Impact Guidelines were inspired by [Mozilla's code of conduct +enforcement ladder](https://github.com/mozilla/diversity). + +[homepage]: https://www.contributor-covenant.org + +For answers to common questions about this code of conduct, see the FAQ at +https://www.contributor-covenant.org/faq. Translations are available at +https://www.contributor-covenant.org/translations. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..8f01ea1 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,97 @@ +# Contributing to Burn + +Welcome to the Burn community! We're glad you're interested in contributing. + +## How to Contribute + +The best way to get started is to look at [open issues](https://github.com/tracel-ai/burn/issues) +and find one that interests you. Issues labeled `good first issue` are a great starting point for +new contributors. + +If you have an idea that isn't covered by an existing issue, open one first to discuss the approach. +This helps align expectations and avoids wasted effort on both sides. + +For questions, discussions, or just to say hello, join us on +[Discord](https://discord.gg/uPEBbYYDB6). The [Contributor Book](https://burn.dev/contributor-book/) +covers architecture, environment setup, and guides for common tasks. + +## Pull Requests + +Every pull request should have a descriptive title, a description covering what you changed, why, +how you tested it, and a link to the relevant issue (if applicable). Prefer small, focused PRs over +large ones that bundle unrelated changes. + +Draft pull requests are considered not yet ready for review. + +CI checks should pass before requesting review, though the signal isn't always accurate. If you have +questions or need early feedback, let us know on the PR or on +[Discord](https://discord.gg/uPEBbYYDB6). + +### Change Ownership + +The core principle behind all contributions: **PR authors must understand, justify, and explain +every change they propose.** After a PR is accepted, both the reviewer and the author should be +confident it improves the codebase. + +This applies equally whether you wrote the code from scratch, adapted it from another project, or +used AI tools to help generate it. The origin of the code doesn't matter; what matters is that you +own it intellectually and can stand behind it during review. + +### AI-Assisted Contributions + +Using LLMs and AI tools to generate code that is part of a contribution is allowed. + +That said, the [Change Ownership](#change-ownership) principle applies fully. You are the author, +not your AI tool. This means: + +- Read and understand every line before submitting. +- Review AI-generated code for correctness, style consistency, and relevance. +- Test your changes locally and confirm they work as intended. +- Be prepared to explain the rationale behind any change during review. + +Do not use "AI generated" as a justification for low-quality code. + +### Before You Open a PR + +1. **Check for an existing issue.** If there isn't one, open an issue first to discuss the approach. + This is especially important for large changes or refactors. +2. **Read the codebase.** Understand the architecture and conventions already in place. The + [Contributor Book](https://burn.dev/contributor-book/) covers architecture, environment setup, + and guides for common tasks. +3. **Keep it focused.** One PR should address one concern. If you spot an unrelated issue while + working, open a separate PR for it. +4. **Run validation.** Run `cargo run-checks` before submitting. This runs formatting, linting, and + the full test suite. All checks must pass. + +### Code Quality Standards + +- Follow existing code style and project conventions. +- Write idiomatic Rust. If you are new to the codebase, study existing patterns before contributing. +- Keep dependencies minimal. Don't introduce new crates without discussion. +- Document public APIs. Non-trivial logic should have comments explaining _why_, not just _what_. +- Prefer clarity over cleverness. +- Bug fixes should include a regression test. + +### Large Pull Requests + +Large, complex PRs are harder to review effectively and carry more risk. To help both yourself and +reviewers, consider breaking substantial changes into smaller, incremental PRs. Each should be +valuable on its own, even if the full picture spans multiple PRs. + +Large efforts that are ultimately rejected are frustrating for everyone involved. If you're planning +a substantial change, open an issue or start a discussion first. It's much easier to course-correct +early than after the work is done. + +### Review Process + +- Maintainers review PRs as time allows. Please be patient. +- Be responsive to feedback. If changes are requested, address them or explain your reasoning. +- Reviewers may ask clarifying questions about any part of your PR. This is a normal part of + collaborative review and helps ensure shared understanding. +- Don't force-push to rewrite history during an active review without notice. +- If a PR goes stale for more than 14 days without a response from the author, it may be closed. + +## Getting Help + +If you're stuck or unsure about something, don't hesitate to ask. Open an issue, start a discussion, +or reach out on [Discord](https://discord.gg/uPEBbYYDB6). We're happy to help. diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..5dca8fa --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,12579 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "accelerate-src" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "415ed64958754dbe991900f3940677e6a7eefb4d7367afd70d642677b0c7d19d" + +[[package]] +name = "addr2line" +version = "0.25.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b5d307320b3181d6d7954e663bd7c774a838b8220fe0593c86d9fb09f498b4b" +dependencies = [ + "gimli", +] + +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + +[[package]] +name = "aead" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d122413f284cf2d62fb1b7db97e02edb8cda96d769b16e443a4f6195e35662b0" +dependencies = [ + "crypto-common 0.1.7", + "generic-array", +] + +[[package]] +name = "aes" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" +dependencies = [ + "cfg-if", + "cipher 0.4.4", + "cpufeatures 0.2.17", +] + +[[package]] +name = "aes" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1fc76eaeac4c9164506c466d4ffdd8ec9d0c5bf57ee97177c4d8eceb3a0e138" +dependencies = [ + "cipher 0.5.2", + "cpubits", + "cpufeatures 0.3.0", +] + +[[package]] +name = "aes-gcm" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "831010a0f742e1209b3bcea8fab6a8e149051ba6099432c8cb2cc117dec3ead1" +dependencies = [ + "aead", + "aes 0.8.4", + "cipher 0.4.4", + "ctr", + "ghash", + "subtle", +] + +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom 0.3.4", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "aligned" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee4508988c62edf04abd8d92897fca0c2995d907ce1dfeaf369dac3716a40685" +dependencies = [ + "as-slice", +] + +[[package]] +name = "aligned-vec" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc890384c8602f339876ded803c97ad529f3842aba97f6392b3dba0dd171769b" +dependencies = [ + "equator", +] + +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys 0.61.2", +] + +[[package]] +name = "anyhow" +version = "1.0.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" + +[[package]] +name = "approx" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab112f0a86d568ea0e627cc1d6be74a1e9cd55214684db5561995f6dad897c6" +dependencies = [ + "num-traits", +] + +[[package]] +name = "ar_archive_writer" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4087686b4b0a3427190bae57a1d9a478dbb2d40c5dc1bd6e2b6d797913bdd348" +dependencies = [ + "object", +] + +[[package]] +name = "arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1" +dependencies = [ + "derive_arbitrary", +] + +[[package]] +name = "arc-swap" +version = "1.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c049c0be4daef0b145cb3555416b3b8ef5b7888a38aea1a3a155801fe7b0810b" +dependencies = [ + "rustversion", +] + +[[package]] +name = "arg_enum_proc_macro" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ae92a5119aa49cdbcf6b9f893fe4e1d98b04ccbf82ee0584ad948a44a734dea" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "argminmax" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70f13d10a41ac8d2ec79ee34178d61e6f47a29c2edfe7ef1721c7383b0359e65" +dependencies = [ + "half", + "num-traits", +] + +[[package]] +name = "array-init-cursor" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed51fe0f224d1d4ea768be38c51f9f831dee9d05c163c11fba0b8c44387b1fc3" + +[[package]] +name = "arrayref" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" + +[[package]] +name = "arrayvec" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" + +[[package]] +name = "as-slice" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "516b6b4f0e40d50dcda9365d53964ec74560ad4284da2e7fc97122cd83174516" +dependencies = [ + "stable_deref_trait", +] + +[[package]] +name = "ash" +version = "0.38.0+1.3.281" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bb44936d800fea8f016d7f2311c6a4f97aebd5dc86f09906139ec848cf3a46f" +dependencies = [ + "libloading 0.8.9", +] + +[[package]] +name = "async-channel" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "924ed96dd52d1b75e9c1a3e6275715fd320f5f9439fb5a4a11fa51f4221158d2" +dependencies = [ + "concurrent-queue", + "event-listener-strategy", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "async-trait" +version = "0.1.89" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "async_io_stream" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d7b9decdf35d8908a7e3ef02f64c5e9b1695e230154c0e8de3969142d9b94c" +dependencies = [ + "futures", + "pharos", + "rustc_version", +] + +[[package]] +name = "atoi_simd" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ad17c7c205c2c28b527b9845eeb91cf1b4d008b438f98ce0e628227a822758e" +dependencies = [ + "debug_unsafe", +] + +[[package]] +name = "atomic" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89cbf775b137e9b968e67227ef7f775587cde3fd31b0d8599dbd0f598a48340" +dependencies = [ + "bytemuck", +] + +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "atomic_float" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "628d228f918ac3b82fe590352cc719d30664a0c13ca3a60266fe02c7132d480a" + +[[package]] +name = "attohttpc" +version = "0.30.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16e2cdb6d5ed835199484bb92bb8b3edd526effe995c61732580439c1a67e2e9" +dependencies = [ + "base64 0.22.1", + "http", + "log", + "url", +] + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "av-scenechange" +version = "0.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f321d77c20e19b92c39e7471cf986812cbb46659d2af674adc4331ef3f18394" +dependencies = [ + "aligned", + "anyhow", + "arg_enum_proc_macro", + "arrayvec", + "log", + "num-rational", + "num-traits", + "pastey", + "rayon", + "thiserror 2.0.18", + "v_frame", + "y4m", +] + +[[package]] +name = "av1-grain" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8cfddb07216410377231960af4fcab838eaa12e013417781b78bd95ee22077f8" +dependencies = [ + "anyhow", + "arrayvec", + "log", + "nom 8.0.0", + "num-rational", + "v_frame", +] + +[[package]] +name = "avif-serialize" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7178fe5f7d460b13895ebb9dcb28a3a6216d2df2574a0806cb51b555d297f38" +dependencies = [ + "arrayvec", +] + +[[package]] +name = "axum" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" +dependencies = [ + "axum-core", + "base64 0.22.1", + "bytes", + "form_urlencoded", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-util", + "itoa", + "matchit", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "serde_core", + "serde_json", + "serde_path_to_error", + "serde_urlencoded", + "sha1 0.10.6", + "sync_wrapper", + "tokio", + "tokio-tungstenite", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "axum-core" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "http-body-util", + "mime", + "pin-project-lite", + "sync_wrapper", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "backon" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef" +dependencies = [ + "fastrand", + "gloo-timers", + "tokio", +] + +[[package]] +name = "backtrace" +version = "0.3.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb531853791a215d7c62a30daf0dde835f381ab5de4589cfe7c649d2cbe92bd6" +dependencies = [ + "addr2line", + "cfg-if", + "libc", + "miniz_oxide", + "object", + "rustc-demangle", + "windows-link 0.2.1", +] + +[[package]] +name = "base16ct" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd307490d624467aa6f74b0eabb77633d1f758a7b25f12bceb0b22e08d9726f6" + +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "base64ct" +version = "1.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" + +[[package]] +name = "bincode" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36eaf5d7b090263e8150820482d5d93cd964a81e4019913c972f4edcc6edb740" +dependencies = [ + "bincode_derive", + "serde", + "unty", +] + +[[package]] +name = "bincode_derive" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf95709a440f45e986983918d0e8a1f30a9b1df04918fc828670606804ac3c09" +dependencies = [ + "virtue", +] + +[[package]] +name = "bit-set" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0700ddab506f33b20a03b13996eccd309a48e5ff77d0d95926aa0210fb4e95f1" +dependencies = [ + "bit-vec 0.6.3", +] + +[[package]] +name = "bit-set" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2f926cc3060f09db9ebc5b52823d85268d24bb917e472c0c4bea35780a7d" +dependencies = [ + "bit-vec 0.9.1", +] + +[[package]] +name = "bit-vec" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" + +[[package]] +name = "bit-vec" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b71798fca2c1fe1086445a7258a4bc81e6e49dcd24c8d0dd9a1e57395b603f51" + +[[package]] +name = "bit_field" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e4b40c7323adcfc0a41c4b88143ed58346ff65a288fc144329c5c45e05d70c6" + +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" +dependencies = [ + "serde_core", +] + +[[package]] +name = "bitstream-io" +version = "4.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7eff00be299a18769011411c9def0d827e8f2d7bf0c3dbf53633147a8867fd1f" +dependencies = [ + "no_std_io2", +] + +[[package]] +name = "blake3" +version = "1.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0aa83c34e62843d924f905e0f5c866eb1dd6545fc4d719e803d9ba6030371fce" +dependencies = [ + "arrayref", + "arrayvec", + "cc", + "cfg-if", + "constant_time_eq 0.4.2", + "cpufeatures 0.3.0", +] + +[[package]] +name = "blas-src" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5bd439b48f5425c9f4597141b52476358e05e08471bc30f4d1c2a1f4084d7c9f" +dependencies = [ + "accelerate-src", + "netlib-src", + "openblas-src", +] + +[[package]] +name = "block" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d8c1fef690941d3e7788d328517591fecc684c084084702d6ff1641e993699a" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "block-buffer" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" +dependencies = [ + "hybrid-array", + "zeroize", +] + +[[package]] +name = "block2" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdeb9d870516001442e364c5220d3574d2da8dc765554b4a617230d33fa58ef5" +dependencies = [ + "objc2", +] + +[[package]] +name = "bon" +version = "3.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a602c73c7b0148ec6d12af6fd5cc7a46e2eacc8878271a999abac56eed12f561" +dependencies = [ + "bon-macros", + "rustversion", +] + +[[package]] +name = "bon-macros" +version = "3.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dee98b0db6a962de883bf5d20362dee4d7ca0d12fe39a7c6c73c844e1cd7c1f" +dependencies = [ + "darling 0.23.0", + "ident_case", + "prettyplease", + "proc-macro2", + "quote", + "rustversion", + "syn 2.0.118", +] + +[[package]] +name = "boxcar" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36f64beae40a84da1b4b26ff2761a5b895c12adc41dc25aaee1c4f2bbfe97a6e" + +[[package]] +name = "bstr" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cee35f73844aa3014bb606320a6c1f010249dbdf43342fe54b5a4f6a8ed4b79" +dependencies = [ + "memchr", + "regex-automata", + "serde_core", +] + +[[package]] +name = "built" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c0e531d93d39c34eef561e929e8a7f86d77a5af08aac4f6d6e39976c51858e9" + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "burn" +version = "0.22.0-pre.1" +dependencies = [ + "burn-autodiff", + "burn-candle", + "burn-core", + "burn-cpu", + "burn-cuda", + "burn-dispatch", + "burn-flex", + "burn-ir", + "burn-ndarray", + "burn-nn", + "burn-optim", + "burn-rl", + "burn-rocm", + "burn-router", + "burn-std", + "burn-store", + "burn-tch", + "burn-train", + "burn-vision", + "burn-wgpu", + "cubecl", +] + +[[package]] +name = "burn-autodiff" +version = "0.22.0-pre.1" +dependencies = [ + "burn-backend", + "burn-std", + "derive-new", + "hashbrown 0.16.1", + "log", + "num-traits", + "parking_lot", + "portable-atomic", + "portable-atomic-util", + "spin", + "tracing", +] + +[[package]] +name = "burn-backend" +version = "0.22.0-pre.1" +dependencies = [ + "burn-std", + "bytemuck", + "cubecl", + "derive-new", + "enumset", + "hashbrown 0.16.1", + "num-traits", + "oneshot", + "paste", + "portable-atomic-util", + "rand 0.10.2", + "rand_distr 0.6.0", + "serde", + "serde_json", + "serial_test", + "spin", +] + +[[package]] +name = "burn-backend-extension" +version = "0.22.0-pre.1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "burn-backend-tests" +version = "0.22.0-pre.1" +dependencies = [ + "burn-autodiff", + "burn-backend", + "burn-cubecl", + "burn-cubecl-fusion", + "burn-fusion", + "burn-ndarray", + "burn-tensor", + "ctor", + "cubek", + "divan", + "num-traits", + "proc-macro2", + "quote", + "rand 0.10.2", + "serial_test", + "syn 2.0.118", +] + +[[package]] +name = "burn-candle" +version = "0.22.0-pre.1" +dependencies = [ + "burn-backend", + "burn-std", + "burn-tch", + "candle-core", + "derive-new", +] + +[[package]] +name = "burn-communication" +version = "0.22.0-pre.1" +dependencies = [ + "axum", + "burn-backend", + "burn-std", + "bytes", + "derive-new", + "futures", + "futures-util", + "log", + "rmp-serde", + "serde", + "serde_bytes", + "tokio", + "tokio-tungstenite", + "tokio-util", + "tracing", + "tracing-core", + "tracing-subscriber", +] + +[[package]] +name = "burn-core" +version = "0.22.0-pre.1" +dependencies = [ + "ahash", + "burn-backend", + "burn-backend-extension", + "burn-dataset", + "burn-derive", + "burn-dispatch", + "burn-ir", + "burn-pack", + "burn-std", + "burn-tensor", + "derive-new", + "half", + "hashbrown 0.16.1", + "log", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rand 0.10.2", + "regex", + "rstest", + "serde", + "serde_json", + "serial_test", + "spin", + "tempfile", + "thiserror 2.0.18", +] + +[[package]] +name = "burn-cpu" +version = "0.22.0-pre.1" +dependencies = [ + "burn-backend", + "burn-cubecl", + "burn-fusion", + "cubecl", +] + +[[package]] +name = "burn-cubecl" +version = "0.22.0-pre.1" +dependencies = [ + "burn-backend", + "burn-cubecl-fusion", + "burn-fusion", + "burn-ir", + "burn-std", + "cubecl", + "cubek", + "derive-new", + "futures-lite", + "log", + "serde", + "text_placeholder", + "tracing", +] + +[[package]] +name = "burn-cubecl-fusion" +version = "0.22.0-pre.1" +dependencies = [ + "burn-backend", + "burn-fusion", + "burn-ir", + "burn-std", + "cubecl", + "cubek", + "derive-new", + "half", + "hashbrown 0.16.1", + "serde", + "spin", +] + +[[package]] +name = "burn-cuda" +version = "0.22.0-pre.1" +dependencies = [ + "burn-backend", + "burn-cubecl", + "burn-fusion", + "cubecl", +] + +[[package]] +name = "burn-dataset" +version = "0.22.0-pre.1" +dependencies = [ + "burn-std", + "csv", + "derive-new", + "dirs", + "encoding_rs", + "fake", + "flate2", + "gix-tempfile", + "globwalk", + "hound", + "image", + "planus", + "polars", + "r2d2", + "r2d2_sqlite", + "rand 0.10.2", + "rayon", + "rmp-serde", + "rstest", + "rusqlite", + "sanitize-filename", + "serde", + "serde_json", + "serde_rusqlite", + "strum 0.28.0", + "tar", + "tempfile", + "thiserror 2.0.18", + "zip 8.6.0", +] + +[[package]] +name = "burn-derive" +version = "0.22.0-pre.1" +dependencies = [ + "derive-new", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "burn-dispatch" +version = "0.22.0-pre.1" +dependencies = [ + "burn-autodiff", + "burn-backend", + "burn-cpu", + "burn-cuda", + "burn-flex", + "burn-ndarray", + "burn-remote", + "burn-rocm", + "burn-tch", + "burn-wgpu", + "paste", +] + +[[package]] +name = "burn-flex" +version = "0.22.0-pre.1" +dependencies = [ + "aligned-vec", + "burn-backend", + "burn-ir", + "burn-std", + "bytemuck", + "gemm 0.19.0", + "half", + "libm", + "macerator", + "num-traits", + "once_cell", + "portable-atomic", + "portable-atomic-util", + "rayon", + "realfft", +] + +[[package]] +name = "burn-fusion" +version = "0.22.0-pre.1" +dependencies = [ + "burn-backend", + "burn-ir", + "burn-std", + "derive-new", + "hashbrown 0.16.1", + "log", + "serde", + "smallvec", + "spin", + "tracing", +] + +[[package]] +name = "burn-ir" +version = "0.22.0-pre.1" +dependencies = [ + "burn-backend", + "hashbrown 0.16.1", + "serde", +] + +[[package]] +name = "burn-ndarray" +version = "0.22.0-pre.1" +dependencies = [ + "atomic_float", + "blas-src", + "burn-autodiff", + "burn-backend", + "burn-ir", + "burn-std", + "bytemuck", + "bytes", + "const-random", + "itertools", + "libm", + "macerator", + "matrixmultiply", + "ndarray 0.17.2", + "num-traits", + "openblas-src", + "paste", + "portable-atomic", + "portable-atomic-util", + "rand 0.10.2", + "rayon", + "seq-macro", +] + +[[package]] +name = "burn-nn" +version = "0.22.0-pre.1" +dependencies = [ + "burn-core", + "num-traits", + "rstest", +] + +[[package]] +name = "burn-no-std-tests" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "burn-flex", + "burn-store", +] + +[[package]] +name = "burn-optim" +version = "0.22.0-pre.1" +dependencies = [ + "burn-autodiff", + "burn-core", + "burn-derive", + "burn-flex", + "burn-nn", + "burn-pack", + "derive-new", + "hashbrown 0.16.1", + "log", + "num-traits", + "rstest", + "serde", +] + +[[package]] +name = "burn-pack" +version = "0.22.0-pre.1" +dependencies = [ + "burn-std", + "byteorder", + "ciborium", + "serde", + "tempfile", +] + +[[package]] +name = "burn-remote" +version = "0.22.0-pre.1" +dependencies = [ + "burn-backend", + "burn-communication", + "burn-flex", + "burn-fusion", + "burn-ir", + "burn-router", + "burn-std", + "burn-tensor", + "bytes", + "futures-util", + "gloo-timers", + "iroh", + "log", + "rand 0.10.2", + "rmp-serde", + "serde", + "serde_bytes", + "tokio", + "tokio-util", + "wasm-bindgen-futures", +] + +[[package]] +name = "burn-rl" +version = "0.22.0-pre.1" +dependencies = [ + "burn-core", + "derive-new", + "log", + "rand 0.10.2", +] + +[[package]] +name = "burn-rocm" +version = "0.22.0-pre.1" +dependencies = [ + "burn-backend", + "burn-cubecl", + "burn-fusion", + "cubecl", +] + +[[package]] +name = "burn-router" +version = "0.22.0-pre.1" +dependencies = [ + "burn-backend", + "burn-flex", + "burn-fusion", + "burn-ir", + "burn-std", + "burn-tensor", + "burn-wgpu", + "hashbrown 0.16.1", + "log", + "serde", + "spin", +] + +[[package]] +name = "burn-std" +version = "0.22.0-pre.1" +dependencies = [ + "ahash", + "bytemuck", + "bytes", + "cubecl-common", + "cubecl-zspace", + "dashmap", + "data-encoding", + "derive-new", + "enumset", + "half", + "hashbrown 0.16.1", + "indicatif 0.18.6", + "num-traits", + "paste", + "portable-atomic-util", + "rand 0.10.2", + "rand_distr 0.6.0", + "reqwest 0.12.28", + "serde", + "serde_json", + "serial_test", + "smallvec", + "spin", + "tempfile", + "thiserror 2.0.18", + "tokio", + "uuid", +] + +[[package]] +name = "burn-store" +version = "0.22.0-pre.1" +dependencies = [ + "burn-core", + "burn-nn", + "burn-pack", + "burn-tensor", + "byteorder", + "bytes", + "ciborium", + "divan", + "half", + "hashbrown 0.16.1", + "lzma-rust2 0.16.5", + "memmap2", + "num-traits", + "regex", + "safetensors 0.7.0", + "serde", + "tar", + "tempfile", + "textdistance", + "thiserror 2.0.18", + "zip 8.6.0", +] + +[[package]] +name = "burn-tch" +version = "0.22.0-pre.1" +dependencies = [ + "burn-backend", + "cc", + "libc", + "log", + "tch", + "torch-sys", +] + +[[package]] +name = "burn-tensor" +version = "0.22.0-pre.1" +dependencies = [ + "burn-backend", + "burn-backend-extension", + "burn-dispatch", + "burn-std", + "colored 3.1.1", + "derive-new", + "num-traits", + "serde", +] + +[[package]] +name = "burn-train" +version = "0.22.0-pre.1" +dependencies = [ + "async-channel", + "burn-core", + "burn-flex", + "burn-nn", + "burn-optim", + "burn-rl", + "burn-std", + "burn-store", + "derive-new", + "dirs", + "log", + "nvml-wrapper", + "rand 0.10.2", + "ratatui", + "rstest", + "serde", + "sysinfo 0.38.4", + "systemstat", + "tempfile", + "thiserror 2.0.18", + "tracing-appender", + "tracing-core", + "tracing-subscriber", + "unicode-width", +] + +[[package]] +name = "burn-vision" +version = "0.22.0-pre.1" +dependencies = [ + "aligned-vec", + "bon", + "burn-core", + "burn-cubecl", + "burn-fusion", + "burn-ir", + "burn-store", + "bytemuck", + "cubecl", + "derive-new", + "dirs", + "half", + "image", + "macerator", + "ndarray 0.17.2", + "num-traits", + "paste", + "serde", +] + +[[package]] +name = "burn-wgpu" +version = "0.22.0-pre.1" +dependencies = [ + "burn-backend", + "burn-cubecl", + "burn-fusion", + "cubecl", +] + +[[package]] +name = "by_address" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64fa3c856b712db6612c019f14756e64e4bcea13337a6b33b696333a9eaa2d06" + +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9abbd1bc6865053c427f7198e6af43bfdedc55ab791faed4fbd361d789575ff" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + +[[package]] +name = "byteorder-lite" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f1fe948ff07f4bd06c30984e69f5b4899c516a3ef74f34df92a2df2ab535495" + +[[package]] +name = "bytes" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" +dependencies = [ + "portable-atomic", + "serde", +] + +[[package]] +name = "bytesize" +version = "1.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e93abca9e28e0a1b9877922aacb20576e05d4679ffa78c3d6dc22a26a216659" + +[[package]] +name = "bzip2" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bdb116a6ef3f6c3698828873ad02c3014b3c85cadb88496095628e3ef1e347f8" +dependencies = [ + "bzip2-sys", + "libc", +] + +[[package]] +name = "bzip2" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3a53fac24f34a81bc9954b5d6cfce0c21e18ec6959f44f56e8e90e4bb7c346c" +dependencies = [ + "libbz2-rs-sys", +] + +[[package]] +name = "bzip2-sys" +version = "0.1.13+1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "225bff33b2141874fe80d71e07d6eec4f85c5c216453dd96388240f96e1acc14" +dependencies = [ + "cc", + "pkg-config", +] + +[[package]] +name = "c2rust-bitfields" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46dc7d2bffa0d0b3d47eb2dc69973466858281446c2ac9f6d8a10e92ab1017df" +dependencies = [ + "c2rust-bitfields-derive", +] + +[[package]] +name = "c2rust-bitfields-derive" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe1117afa5937ce280034e31fa1e84ed1824a252f75380327eed438535333f8" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + +[[package]] +name = "camino" +version = "1.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f2d30e4173c4026932d51d31d6b0613b1fd3014bf3f9f8943d4ba139c437ba0" +dependencies = [ + "serde_core", +] + +[[package]] +name = "candle-core" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bd9895436c1ba5dc1037a19935d084b838db066ff4e15ef7dded020b7c12a4a" +dependencies = [ + "accelerate-src", + "byteorder", + "candle-kernels", + "candle-metal-kernels", + "candle-ug", + "cudarc 0.19.8", + "float8", + "gemm 0.19.0", + "half", + "libc", + "libm", + "memmap2", + "num-traits", + "num_cpus", + "objc2-foundation", + "objc2-metal", + "rand 0.9.4", + "rand_distr 0.5.1", + "rayon", + "safetensors 0.7.0", + "thiserror 2.0.18", + "tokenizers", + "yoke 0.8.3", + "zip 7.2.0", +] + +[[package]] +name = "candle-kernels" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "742e2ac226b777134436e9e692f44e77c278b8a7abb1554dc10e44dc911b349f" +dependencies = [ + "cudaforge", +] + +[[package]] +name = "candle-metal-kernels" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b6b5a4cae6b4e1ab0efcee4dc05272d11b374a3d1ba121b3a961e36be54ab60" +dependencies = [ + "half", + "objc2", + "objc2-foundation", + "objc2-metal", + "once_cell", + "thiserror 2.0.18", + "tracing", +] + +[[package]] +name = "candle-ug" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca0fc3167cbc99c8ec1be618cb620aa21dca95038f118c3579a79370e3dc5f77" +dependencies = [ + "ug", + "ug-cuda", + "ug-metal", +] + +[[package]] +name = "cargo-platform" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd0061da739915fae12ea00e16397555ed4371a6bb285431aab930f61b0aa4ba" +dependencies = [ + "serde", + "serde_core", +] + +[[package]] +name = "cargo_metadata" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef987d17b0a113becdd19d3d0022d04d7ef41f9efe4f3fb63ac44ba61df3ade9" +dependencies = [ + "camino", + "cargo-platform", + "semver", + "serde", + "serde_json", + "thiserror 2.0.18", +] + +[[package]] +name = "caseless" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b6fd507454086c8edfd769ca6ada439193cdb209c7681712ef6275cccbfe5d8" +dependencies = [ + "unicode-normalization", +] + +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + +[[package]] +name = "cblas-sys" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6feecd82cce51b0204cf063f0041d69f24ce83f680d87514b004248e7b0fa65" +dependencies = [ + "libc", +] + +[[package]] +name = "cc" +version = "1.2.66" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5d6cac793997bd970000024b2934968efe83b382de4fdcf4fcb46b6ee4ad996" +dependencies = [ + "find-msvc-tools", + "jobserver", + "libc", + "shlex", +] + +[[package]] +name = "cesu8" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d43a04d8753f35258c91f8ec639f792891f748a1edbd759cf1dcea3382ad83c" + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + +[[package]] +name = "chacha20" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core 0.10.1", +] + +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "iana-time-zone", + "num-traits", + "serde", + "windows-link 0.2.1", +] + +[[package]] +name = "chrono-tz" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6139a8597ed92cf816dfb33f5dd6cf0bb93a6adc938f11039f371bc5bcd26c3" +dependencies = [ + "chrono", + "phf 0.12.1", +] + +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common 0.1.7", + "inout 0.1.4", +] + +[[package]] +name = "cipher" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8cf2a2c93cd704877c0858356ed03480ff301ee950b43f1cbe4573b088bfa6c" +dependencies = [ + "crypto-common 0.2.2", + "inout 0.2.2", +] + +[[package]] +name = "clap" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", + "terminal_size", +] + +[[package]] +name = "clap_derive" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "cmake" +version = "0.1.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" +dependencies = [ + "cc", +] + +[[package]] +name = "cmov" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" + +[[package]] +name = "cobs" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fa961b519f0b462e3a3b4a34b64d119eeaca1d59af726fe450bbba07a9fc0a1" +dependencies = [ + "thiserror 2.0.18", +] + +[[package]] +name = "codespan-reporting" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af491d569909a7e4dee0ad7db7f5341fef5c614d5b8ec8cf765732aba3cff681" +dependencies = [ + "serde", + "termcolor", + "unicode-width", +] + +[[package]] +name = "color_quant" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d7b894f5411737b7867f4827955924d7c254fc9f4d91a6aad6b097804b1018b" + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + +[[package]] +name = "colored" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" +dependencies = [ + "lazy_static", + "windows-sys 0.59.0", +] + +[[package]] +name = "colored" +version = "3.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "combine" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" +dependencies = [ + "bytes", + "memchr", +] + +[[package]] +name = "comfy-table" +version = "7.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "958c5d6ecf1f214b4c2bbbbf6ab9523a864bd136dcf71a7e8904799acfe1ad47" +dependencies = [ + "crossterm", + "unicode-segmentation", + "unicode-width", +] + +[[package]] +name = "compact_str" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dfdd1c2274d9aa354115b09dc9a901d6c5576818cdf70d14cae2bdb47df00ab" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", +] + +[[package]] +name = "comrak" +version = "0.39.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fefab951771fc3beeed0773ce66a4f7b706273fc6c4c95b08dd1615744abcf5" +dependencies = [ + "caseless", + "entities", + "memchr", + "slug", + "typed-arena", + "unicode_categories", +] + +[[package]] +name = "concurrent-queue" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "condtype" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf0a07a401f374238ab8e2f11a104d2851bf9ce711ec69804834de8af45c7af" + +[[package]] +name = "console" +version = "0.15.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "054ccb5b10f9f2cbf51eb355ca1d05c2d279ce1804688d0db74b4733a5aeafd8" +dependencies = [ + "encode_unicode", + "libc", + "once_cell", + "unicode-width", + "windows-sys 0.59.0", +] + +[[package]] +name = "console" +version = "0.16.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fe5f465a4f6fee88fad41b85d990f84c835335e85b5d9e6e63e0d06d28cba7c" +dependencies = [ + "encode_unicode", + "libc", + "unicode-width", + "windows-sys 0.61.2", +] + +[[package]] +name = "console_error_panic_hook" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06aeb73f470f66dcdbf7223caeebb85984942f22f1adb2a088cf9668146bbbc" +dependencies = [ + "cfg-if", + "wasm-bindgen", +] + +[[package]] +name = "const-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + +[[package]] +name = "const-random" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359" +dependencies = [ + "const-random-macro", +] + +[[package]] +name = "const-random-macro" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" +dependencies = [ + "getrandom 0.2.17", + "once_cell", + "tiny-keccak", +] + +[[package]] +name = "constant_time_eq" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "245097e9a4535ee1e3e3931fcfcd55a796a44c643e8596ff6566d68f09b87bbc" + +[[package]] +name = "constant_time_eq" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c74b8349d32d297c9134b8c88677813a227df8f779daa29bfc29c183fe3dca6" + +[[package]] +name = "constant_time_eq" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" + +[[package]] +name = "constcat" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "136d3e02915a2cea4d74caa8681e2d44b1c3254bdbf17d11d41d587ff858832c" + +[[package]] +name = "convert_case" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baaaa0ecca5b51987b9423ccdc971514dd8b0bb7b4060b983d3664dad3f1f89f" +dependencies = [ + "unicode-segmentation", +] + +[[package]] +name = "convert_case" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "633458d4ef8c78b72454de2d54fd6ab2e60f9e02be22f3c6104cdc8a4e0fceb9" +dependencies = [ + "unicode-segmentation", +] + +[[package]] +name = "cookie" +version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ddef33a339a91ea89fb53151bd0a4689cfce27055c291dfa69945475d22c747" +dependencies = [ + "percent-encoding", + "time", + "version_check", +] + +[[package]] +name = "cookie_store" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15b2c103cf610ec6cae3da84a766285b42fd16aad564758459e6ecf128c75206" +dependencies = [ + "cookie", + "document-features", + "idna", + "indexmap", + "log", + "serde", + "serde_derive", + "serde_json", + "time", + "url", +] + +[[package]] +name = "cordyceps" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "688d7fbb8092b8de775ef2536f36c8c31f2bc4006ece2e8d8ad2d17d00ce0a2a" +dependencies = [ + "loom", + "tracing", +] + +[[package]] +name = "core-foundation" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "core-graphics-types" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45390e6114f68f718cc7a830514a96f903cccd70d02a8f6d9f643ac4ba45afaf" +dependencies = [ + "bitflags 1.3.2", + "core-foundation 0.9.4", + "libc", +] + +[[package]] +name = "cpubits" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15b85f9c39137c3a891689859392b1bd49812121d0d61c9caf00d46ed5ce06ae" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + +[[package]] +name = "crc" +version = "3.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5eb8a2a1cd12ab0d987a5d5e825195d372001a4094a0376319d5a0ad71c1ba0d" +dependencies = [ + "crc-catalog", +] + +[[package]] +name = "crc-catalog" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "critical-section" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "790eea4361631c5e7d22598ecd5723ff611904e3344ce8720784c93e3d83d40b" + +[[package]] +name = "crossbeam-channel" +version = "0.5.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d85363c37faeca707aef026efa9f3b34d077bce547e48f770770625c6013679e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-queue" +version = "0.3.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "803d13fb3b09d88be9f4dbc29062c66b19bf7170867ceb746d2a8689bf6c7a26" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "crossterm" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8b9f2e4c67f833b660cdb0a3523065869fb35570177239812ed4c905aeff87b" +dependencies = [ + "bitflags 2.13.0", + "crossterm_winapi", + "derive_more", + "document-features", + "mio", + "parking_lot", + "rustix", + "signal-hook 0.3.18", + "signal-hook-mio", + "winapi", +] + +[[package]] +name = "crossterm_winapi" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "acdd7c62a3665c7f6830a51635d9ac9b23ed385797f70a83bb8bafe9c572ab2b" +dependencies = [ + "winapi", +] + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "crypto-common" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +dependencies = [ + "hybrid-array", +] + +[[package]] +name = "csscolorparser" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb2a7d3066da2de787b7f032c736763eb7ae5d355f81a68bab2675a96008b0bf" +dependencies = [ + "lab", + "phf 0.11.3", +] + +[[package]] +name = "csv" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938" +dependencies = [ + "csv-core", + "itoa", + "ryu", + "serde_core", +] + +[[package]] +name = "csv-core" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782" +dependencies = [ + "memchr", +] + +[[package]] +name = "ctor" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83cf0d42651b16c6dfe68685716d18480d18a9c39c62d76e8cf3eb6ed5d8bcbf" +dependencies = [ + "ctor-proc-macro", + "dtor", + "link-section", +] + +[[package]] +name = "ctor-proc-macro" +version = "0.0.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a949c44fcacbbbb7ada007dc7acb34603dd97cd47de5d054f2b6493ecebb483" + +[[package]] +name = "ctr" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0369ee1ad671834580515889b80f2ea915f23b8be8d0daa4bbaf2ac5c7590835" +dependencies = [ + "cipher 0.4.4", +] + +[[package]] +name = "ctrlc" +version = "3.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0b1fab2ae45819af2d0731d60f2afe17227ebb1a1538a236da84c93e9a60162" +dependencies = [ + "dispatch2", + "nix 0.31.3", + "windows-sys 0.61.2", +] + +[[package]] +name = "ctutils" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +dependencies = [ + "cmov", +] + +[[package]] +name = "cubecl" +version = "0.11.0-pre.1" +source = "git+https://github.com/tracel-ai/cubecl?rev=9110136fe5bb3c6e7dc0295100fe7f88342ceb6e#9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" +dependencies = [ + "cubecl-core", + "cubecl-cpu", + "cubecl-cuda", + "cubecl-hip", + "cubecl-ir", + "cubecl-runtime", + "cubecl-std", + "cubecl-wgpu", + "half", +] + +[[package]] +name = "cubecl-common" +version = "0.11.0-pre.1" +source = "git+https://github.com/tracel-ai/cubecl?rev=9110136fe5bb3c6e7dc0295100fe7f88342ceb6e#9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" +dependencies = [ + "backtrace", + "bytemuck", + "bytes", + "cfg-if", + "cfg_aliases", + "ciborium", + "derive-new", + "derive_more", + "embassy-futures", + "embassy-time", + "etcetera", + "float4", + "float8", + "futures-lite", + "half", + "hashbrown 0.17.1", + "log", + "num-traits", + "oneshot", + "parking_lot", + "portable-atomic", + "portable-atomic-util", + "rand 0.10.2", + "sanitize-filename", + "serde", + "serde_bytes", + "serde_json", + "spin", + "toml 1.1.2+spec-1.1.0", + "tracing", + "tynm", + "wasm-bindgen-futures", + "web-time", + "xxhash-rust", +] + +[[package]] +name = "cubecl-core" +version = "0.11.0-pre.1" +source = "git+https://github.com/tracel-ai/cubecl?rev=9110136fe5bb3c6e7dc0295100fe7f88342ceb6e#9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" +dependencies = [ + "bitflags 2.13.0", + "bytemuck", + "cubecl-common", + "cubecl-ir", + "cubecl-macros", + "cubecl-runtime", + "cubecl-zspace", + "derive-new", + "derive_more", + "enumset", + "float-ord", + "half", + "hashbrown 0.17.1", + "itertools", + "log", + "num-integer", + "num-traits", + "paste", + "serde", + "serde_json", + "tracing", + "variadics_please", +] + +[[package]] +name = "cubecl-cpp" +version = "0.11.0-pre.1" +source = "git+https://github.com/tracel-ai/cubecl?rev=9110136fe5bb3c6e7dc0295100fe7f88342ceb6e#9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" +dependencies = [ + "bytemuck", + "cubecl-common", + "cubecl-core", + "cubecl-opt", + "cubecl-runtime", + "derive-new", + "half", + "itertools", + "log", +] + +[[package]] +name = "cubecl-cpu" +version = "0.11.0-pre.1" +source = "git+https://github.com/tracel-ai/cubecl?rev=9110136fe5bb3c6e7dc0295100fe7f88342ceb6e#9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" +dependencies = [ + "bytemuck", + "crossbeam-utils", + "cubecl-common", + "cubecl-core", + "cubecl-opt", + "cubecl-runtime", + "cubecl-std", + "derive-new", + "half", + "libc", + "log", + "paste", + "serde", + "smallvec", + "spin", + "sysinfo 0.38.4", + "tracel-llvm", + "tracel-llvm-bundler", + "winapi", +] + +[[package]] +name = "cubecl-cuda" +version = "0.11.0-pre.1" +source = "git+https://github.com/tracel-ai/cubecl?rev=9110136fe5bb3c6e7dc0295100fe7f88342ceb6e#9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" +dependencies = [ + "bytemuck", + "cubecl-common", + "cubecl-core", + "cubecl-cpp", + "cubecl-runtime", + "cudarc 0.19.8", + "derive-new", + "half", + "log", + "serde", + "tracing", +] + +[[package]] +name = "cubecl-hip" +version = "0.11.0-pre.1" +source = "git+https://github.com/tracel-ai/cubecl?rev=9110136fe5bb3c6e7dc0295100fe7f88342ceb6e#9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" +dependencies = [ + "bytemuck", + "cubecl-common", + "cubecl-core", + "cubecl-cpp", + "cubecl-hip-sys", + "cubecl-runtime", + "derive-new", + "half", + "log", + "paste", + "serde", + "tracing", +] + +[[package]] +name = "cubecl-hip-sys" +version = "7.2.5321100" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "58887c6d217859c2261a868a3a6b66aa8f44ea2f4bfa51d0dbe4dcb0d1f5768d" +dependencies = [ + "libc", + "regex", +] + +[[package]] +name = "cubecl-ir" +version = "0.11.0-pre.1" +source = "git+https://github.com/tracel-ai/cubecl?rev=9110136fe5bb3c6e7dc0295100fe7f88342ceb6e#9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" +dependencies = [ + "bumpalo", + "cubecl-common", + "cubecl-macros-internal", + "derive-new", + "derive_more", + "enumset", + "float-ord", + "fnv", + "foldhash 0.2.0", + "half", + "hashbrown 0.17.1", + "internment", + "itertools", + "num-traits", + "portable-atomic", + "serde", + "variadics_please", +] + +[[package]] +name = "cubecl-macros" +version = "0.11.0-pre.1" +source = "git+https://github.com/tracel-ai/cubecl?rev=9110136fe5bb3c6e7dc0295100fe7f88342ceb6e#9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" +dependencies = [ + "cubecl-common", + "darling 0.23.0", + "derive-new", + "derive_more", + "ident_case", + "inflections", + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "cubecl-macros-internal" +version = "0.11.0-pre.1" +source = "git+https://github.com/tracel-ai/cubecl?rev=9110136fe5bb3c6e7dc0295100fe7f88342ceb6e#9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" +dependencies = [ + "darling 0.23.0", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "cubecl-opt" +version = "0.11.0-pre.1" +source = "git+https://github.com/tracel-ai/cubecl?rev=9110136fe5bb3c6e7dc0295100fe7f88342ceb6e#9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" +dependencies = [ + "cubecl-common", + "cubecl-core", + "cubecl-ir", + "float-ord", + "hashbrown 0.17.1", + "log", + "num", + "petgraph", + "smallvec", + "stable-vec", + "type-map", +] + +[[package]] +name = "cubecl-runtime" +version = "0.11.0-pre.1" +source = "git+https://github.com/tracel-ai/cubecl?rev=9110136fe5bb3c6e7dc0295100fe7f88342ceb6e#9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" +dependencies = [ + "ahash", + "async-channel", + "bytemuck", + "cfg-if", + "cfg_aliases", + "cubecl-common", + "cubecl-ir", + "cubecl-zspace", + "derive-new", + "derive_more", + "enumset", + "etcetera", + "hashbrown 0.17.1", + "itertools", + "log", + "md5", + "serde", + "serde_json", + "spin", + "thiserror 2.0.18", + "toml 1.1.2+spec-1.1.0", + "tracing", + "wasm-bindgen-futures", + "web-time", +] + +[[package]] +name = "cubecl-spirv" +version = "0.11.0-pre.1" +source = "git+https://github.com/tracel-ai/cubecl?rev=9110136fe5bb3c6e7dc0295100fe7f88342ceb6e#9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" +dependencies = [ + "bitflags 2.13.0", + "cubecl-common", + "cubecl-core", + "cubecl-opt", + "cubecl-runtime", + "half", + "hashbrown 0.17.1", + "serde", + "tracel-rspirv", +] + +[[package]] +name = "cubecl-std" +version = "0.11.0-pre.1" +source = "git+https://github.com/tracel-ai/cubecl?rev=9110136fe5bb3c6e7dc0295100fe7f88342ceb6e#9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" +dependencies = [ + "cubecl-common", + "cubecl-core", + "cubecl-runtime", + "half", + "num-traits", + "paste", + "serde", + "spin", + "variadics_please", +] + +[[package]] +name = "cubecl-wgpu" +version = "0.11.0-pre.1" +source = "git+https://github.com/tracel-ai/cubecl?rev=9110136fe5bb3c6e7dc0295100fe7f88342ceb6e#9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" +dependencies = [ + "ash", + "async-channel", + "bytemuck", + "cfg-if", + "cfg_aliases", + "cubecl-common", + "cubecl-core", + "cubecl-cpp", + "cubecl-ir", + "cubecl-runtime", + "cubecl-spirv", + "derive-new", + "derive_more", + "half", + "hashbrown 0.17.1", + "itertools", + "log", + "sanitize-filename", + "tracel-ash", + "tracing", + "wasm-bindgen-futures", + "wgpu", +] + +[[package]] +name = "cubecl-zspace" +version = "0.11.0-pre.1" +source = "git+https://github.com/tracel-ai/cubecl?rev=9110136fe5bb3c6e7dc0295100fe7f88342ceb6e#9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" +dependencies = [ + "derive-new", + "serde", + "smallvec", +] + +[[package]] +name = "cubek" +version = "0.3.0-pre.1" +source = "git+https://github.com/tracel-ai/cubek?rev=5f837200f459a6a0694a0aae5655881599c400c8#5f837200f459a6a0694a0aae5655881599c400c8" +dependencies = [ + "cubecl", + "cubek-attention", + "cubek-convolution", + "cubek-fft", + "cubek-interpolate", + "cubek-matmul", + "cubek-pool", + "cubek-quant", + "cubek-random", + "cubek-reduce", + "cubek-std", +] + +[[package]] +name = "cubek-attention" +version = "0.3.0-pre.1" +source = "git+https://github.com/tracel-ai/cubek?rev=5f837200f459a6a0694a0aae5655881599c400c8#5f837200f459a6a0694a0aae5655881599c400c8" +dependencies = [ + "bytemuck", + "cubecl", + "cubecl-common", + "cubek-matmul", + "cubek-std", + "half", + "serde", +] + +[[package]] +name = "cubek-convolution" +version = "0.3.0-pre.1" +source = "git+https://github.com/tracel-ai/cubek?rev=5f837200f459a6a0694a0aae5655881599c400c8#5f837200f459a6a0694a0aae5655881599c400c8" +dependencies = [ + "bytemuck", + "cubecl", + "cubecl-common", + "cubek-matmul", + "cubek-std", + "derive-new", + "enumset", + "half", + "serde", +] + +[[package]] +name = "cubek-fft" +version = "0.3.0-pre.1" +source = "git+https://github.com/tracel-ai/cubek?rev=5f837200f459a6a0694a0aae5655881599c400c8#5f837200f459a6a0694a0aae5655881599c400c8" +dependencies = [ + "cubecl", +] + +[[package]] +name = "cubek-interpolate" +version = "0.3.0-pre.1" +source = "git+https://github.com/tracel-ai/cubek?rev=5f837200f459a6a0694a0aae5655881599c400c8#5f837200f459a6a0694a0aae5655881599c400c8" +dependencies = [ + "cubecl", + "cubecl-common", + "half", + "serde", + "thiserror 2.0.18", +] + +[[package]] +name = "cubek-matmul" +version = "0.3.0-pre.1" +source = "git+https://github.com/tracel-ai/cubek?rev=5f837200f459a6a0694a0aae5655881599c400c8#5f837200f459a6a0694a0aae5655881599c400c8" +dependencies = [ + "bytemuck", + "cubecl", + "cubecl-common", + "cubek-std", + "cubek-tile", + "half", + "serde", +] + +[[package]] +name = "cubek-pool" +version = "0.3.0-pre.1" +source = "git+https://github.com/tracel-ai/cubek?rev=5f837200f459a6a0694a0aae5655881599c400c8#5f837200f459a6a0694a0aae5655881599c400c8" +dependencies = [ + "cubecl", + "cubecl-common", + "thiserror 2.0.18", +] + +[[package]] +name = "cubek-quant" +version = "0.3.0-pre.1" +source = "git+https://github.com/tracel-ai/cubek?rev=5f837200f459a6a0694a0aae5655881599c400c8#5f837200f459a6a0694a0aae5655881599c400c8" +dependencies = [ + "cubecl", + "cubecl-common", + "cubek-tile", + "half", + "serde", +] + +[[package]] +name = "cubek-random" +version = "0.3.0-pre.1" +source = "git+https://github.com/tracel-ai/cubek?rev=5f837200f459a6a0694a0aae5655881599c400c8#5f837200f459a6a0694a0aae5655881599c400c8" +dependencies = [ + "cubecl", + "cubecl-common", + "half", + "num-traits", + "rand 0.10.2", + "serde", +] + +[[package]] +name = "cubek-reduce" +version = "0.3.0-pre.1" +source = "git+https://github.com/tracel-ai/cubek?rev=5f837200f459a6a0694a0aae5655881599c400c8#5f837200f459a6a0694a0aae5655881599c400c8" +dependencies = [ + "cubecl", + "cubek-std", + "half", + "num-traits", + "serde", + "thiserror 2.0.18", +] + +[[package]] +name = "cubek-std" +version = "0.3.0-pre.1" +source = "git+https://github.com/tracel-ai/cubek?rev=5f837200f459a6a0694a0aae5655881599c400c8#5f837200f459a6a0694a0aae5655881599c400c8" +dependencies = [ + "cubecl", + "cubecl-common", + "half", +] + +[[package]] +name = "cubek-tile" +version = "0.3.0-pre.1" +source = "git+https://github.com/tracel-ai/cubek?rev=5f837200f459a6a0694a0aae5655881599c400c8#5f837200f459a6a0694a0aae5655881599c400c8" +dependencies = [ + "bytemuck", + "cubecl", + "cubecl-common", + "half", + "serde", +] + +[[package]] +name = "cudaforge" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d475ff95b9e4096878f47fe0ca5dbad0e23849a79fc225305ea06c2f25771cd" +dependencies = [ + "anyhow", + "fs2", + "glob", + "num_cpus", + "rayon", + "serde", + "serde_json", + "sha2 0.10.9", + "thiserror 2.0.18", + "walkdir", + "which", +] + +[[package]] +name = "cudarc" +version = "0.17.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf99ab37ee7072d64d906aa2dada9a3422f1d975cdf8c8055a573bc84897ed8" +dependencies = [ + "half", + "libloading 0.8.9", +] + +[[package]] +name = "cudarc" +version = "0.19.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42310153e06cf4cd532901f7096beb27504d681736a29ee90728ae4e2d93b2a8" +dependencies = [ + "float8", + "half", + "libloading 0.9.0", +] + +[[package]] +name = "curve25519-dalek" +version = "5.0.0-rc.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f359e08ca85e7bd759e1fd933ff2bccd81864c60a8fba0e259c7f822b0924bf" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "curve25519-dalek-derive", + "digest 0.11.3", + "fiat-crypto", + "rand_core 0.10.1", + "rustc_version", + "serde", + "subtle", + "zeroize", +] + +[[package]] +name = "curve25519-dalek-derive" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "custom-csv-dataset" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "burn-dataset", + "csv", + "polars", + "reqwest 0.12.28", + "serde", +] + +[[package]] +name = "custom-cubecl-kernel" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "burn-cubecl", + "burn-fusion", + "cubecl", + "log", + "serde", +] + +[[package]] +name = "custom-image-dataset" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "flate2", + "tar", +] + +[[package]] +name = "custom-learning-strategy" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "guide", + "log", +] + +[[package]] +name = "custom-renderer" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "guide", + "log", +] + +[[package]] +name = "custom-training-loop" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "guide", + "log", +] + +[[package]] +name = "custom-wgpu-kernel" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "burn-cubecl", + "burn-fusion", + "bytemuck", + "cubecl", + "derive-new", + "log", + "serde", +] + +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core 0.20.11", + "darling_macro 0.20.11", +] + +[[package]] +name = "darling" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9cdf337090841a411e2a7f3deb9187445851f91b309c0c0a29e05f74a00a48c0" +dependencies = [ + "darling_core 0.21.3", + "darling_macro 0.21.3", +] + +[[package]] +name = "darling" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" +dependencies = [ + "darling_core 0.23.0", + "darling_macro 0.23.0", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 2.0.118", +] + +[[package]] +name = "darling_core" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1247195ecd7e3c85f83c8d2a366e4210d588e802133e1e355180a9870b517ea4" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "darling_core" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0" +dependencies = [ + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 2.0.118", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core 0.20.11", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "darling_macro" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81" +dependencies = [ + "darling_core 0.21.3", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "darling_macro" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" +dependencies = [ + "darling_core 0.23.0", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "data-encoding" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" + +[[package]] +name = "data-encoding-macro" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3259c913752a86488b501ed8680446a5ed2d5aeac6e596cb23ba3800768ea32c" +dependencies = [ + "data-encoding", + "data-encoding-macro-internal", +] + +[[package]] +name = "data-encoding-macro-internal" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccc2776f0c61eca1ca32528f85548abd1a4be8fb53d1b21c013e4f18da1e7090" +dependencies = [ + "data-encoding", + "syn 2.0.118", +] + +[[package]] +name = "debug_unsafe" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7eed2c4702fa172d1ce21078faa7c5203e69f5394d48cc436d25928394a867a2" + +[[package]] +name = "deflate64" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac6b926516df9c60bfa16e107b21086399f8285a44ca9711344b9e553c5146e2" + +[[package]] +name = "defmt" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2953bfe4f93bbd20cc71198842756f77d161884c99ebbabc41d80231ded88d1" +dependencies = [ + "bitflags 1.3.2", + "defmt-macros", +] + +[[package]] +name = "defmt-macros" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bad9c72e7ca2137e0dc3813245a0d282fd6daad32fd800af018306a9169b5fe8" +dependencies = [ + "defmt-parser", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "defmt-parser" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10d60334b3b2e7c9d91ef8150abfb6fa4c1c39ebbcf4a81c2e346aad939fee3e" +dependencies = [ + "thiserror 2.0.18", +] + +[[package]] +name = "deltae" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5729f5117e208430e437df2f4843f5e5952997175992d1414f94c57d61e270b4" + +[[package]] +name = "der" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a69dedd701da44b0536442edf09c81a64b0ab97a7a4a5e3d1971f00027cbc63d" +dependencies = [ + "const-oid", + "pem-rfc7468", + "zeroize", +] + +[[package]] +name = "deranged" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" + +[[package]] +name = "derive-new" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2cdc8d50f426189eef89dac62fabfa0abb27d5cc008f25bf4156a0203325becc" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "derive_arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e567bd82dcff979e4b03460c307b3cdc9e96fde3d73bed1496d2bc75d9dd62a" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling 0.20.11", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn 2.0.118", +] + +[[package]] +name = "derive_more" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d751e9e49156b02b44f9c1815bcb94b984cdcc4396ecc32521c739452808b134" +dependencies = [ + "derive_more-impl", +] + +[[package]] +name = "derive_more-impl" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "799a97264921d8623a957f6c3b9011f3b5492f557bbb7a5a19b7fa6d06ba8dcb" +dependencies = [ + "convert_case 0.10.0", + "proc-macro2", + "quote", + "rustc_version", + "syn 2.0.118", + "unicode-xid", +] + +[[package]] +name = "deunicode" +version = "1.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "abd57806937c9cc163efc8ea3910e00a62e2aeb0b8119f1793a978088f8f6b04" + +[[package]] +name = "diatomic-waker" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab03c107fafeb3ee9f5925686dbb7a73bc76e3932abb0d2b365cb64b169cf04c" + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer 0.10.4", + "crypto-common 0.1.7", + "subtle", +] + +[[package]] +name = "digest" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +dependencies = [ + "block-buffer 0.12.1", + "const-oid", + "crypto-common 0.2.2", + "ctutils", + "zeroize", +] + +[[package]] +name = "dirs" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3e8aa94d75141228480295a7d0e7feb620b1a5ad9f12bc40be62411e38cce4e" +dependencies = [ + "dirs-sys", +] + +[[package]] +name = "dirs-sys" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e01a3366d27ee9890022452ee61b2b63a67e6f13f58900b651ff5665f0bb1fab" +dependencies = [ + "libc", + "option-ext", + "redox_users", + "windows-sys 0.61.2", +] + +[[package]] +name = "dispatch2" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e0e367e4e7da84520dedcac1901e4da967309406d1e51017ae1abfb97adbd38" +dependencies = [ + "bitflags 2.13.0", + "block2", + "libc", + "objc2", +] + +[[package]] +name = "displaydoc" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "divan" +version = "0.1.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a405457ec78b8fe08b0e32b4a3570ab5dff6dd16eb9e76a5ee0a9d9cbd898933" +dependencies = [ + "cfg-if", + "clap", + "condtype", + "divan-macros", + "libc", + "regex-lite", +] + +[[package]] +name = "divan-macros" +version = "0.1.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9556bc800956545d6420a640173e5ba7dfa82f38d3ea5a167eb555bc69ac3323" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "dlib" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab8ecd87370524b461f8557c119c405552c396ed91fc0a8eec68679eab26f94a" +dependencies = [ + "libloading 0.8.9", +] + +[[package]] +name = "dlopen2" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e2c5bd4158e66d1e215c49b837e11d62f3267b30c92f1d171c4d3105e3dc4d4" +dependencies = [ + "libc", + "once_cell", + "winapi", +] + +[[package]] +name = "document-features" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4b8a88685455ed29a21542a33abd9cb6510b6b129abadabdcef0f4c55bc8f61" +dependencies = [ + "litrs", +] + +[[package]] +name = "dotenvy" +version = "0.15.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" + +[[package]] +name = "drawille" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e64e461c3f1e69d99372620640b3fd5f0309eeda2e26e4af69f6760c0e1df845" +dependencies = [ + "colored 2.2.0", + "fnv", +] + +[[package]] +name = "dtor" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edf234dd1594d6dd434a8fb8cada51ddbbc593e40e4a01556a0b31c62da2775b" +dependencies = [ + "dtor-proc-macro", +] + +[[package]] +name = "dtor-proc-macro" +version = "0.0.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2647271c92754afcb174e758003cfd1cbf1e43e5a7853d7b1813e63e19e39a73" + +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + +[[package]] +name = "dyn-stack" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c4713e43e2886ba72b8271aa66c93d722116acf7a75555cce11dcde84388fe8" +dependencies = [ + "bytemuck", + "dyn-stack-macros", +] + +[[package]] +name = "dyn-stack-macros" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1d926b4d407d372f141f93bb444696142c29d32962ccbd3531117cf3aa0bfa9" + +[[package]] +name = "ed25519" +version = "3.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29fcf32e6c73d1079f83ab4d782de2d81620346a5f38c6237a86a22f8368980a" +dependencies = [ + "pkcs8", + "serdect", + "signature", +] + +[[package]] +name = "ed25519-dalek" +version = "3.0.0-rc.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b011170fe4f04665565b4110afef66774fe9ffff278f3eb5b81cc73d26e27d60" +dependencies = [ + "curve25519-dalek", + "ed25519", + "rand_core 0.10.1", + "serde", + "sha2 0.11.0", + "signature", + "subtle", + "zeroize", +] + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "embassy-futures" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc2d050bdc5c21e0862a89256ed8029ae6c290a93aecefc73084b3002cdebb01" + +[[package]] +name = "embassy-time" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "592b0c143ec626e821d4d90da51a2bd91d559d6c442b7c74a47d368c9e23d97a" +dependencies = [ + "cfg-if", + "critical-section", + "document-features", + "embassy-time-driver", + "embedded-hal 0.2.7", + "embedded-hal 1.0.0", + "embedded-hal-async", + "futures-core", +] + +[[package]] +name = "embassy-time-driver" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ee71af1b3a0deaa53eaf2d39252f83504c853646e472400b763060389b9fcc9" +dependencies = [ + "document-features", +] + +[[package]] +name = "embedded-hal" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35949884794ad573cf46071e41c9b60efb0cb311e3ca01f7af807af1debc66ff" +dependencies = [ + "nb 0.1.3", + "void", +] + +[[package]] +name = "embedded-hal" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "361a90feb7004eca4019fb28352a9465666b24f840f5c3cddf0ff13920590b89" + +[[package]] +name = "embedded-hal-async" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c4c685bbef7fe13c3c6dd4da26841ed3980ef33e841cddfa15ce8a8fb3f1884" +dependencies = [ + "embedded-hal 1.0.0", +] + +[[package]] +name = "embedded-io" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef1a6892d9eef45c8fa6b9e0086428a2cca8491aca8f787c534a3d6d0bcb3ced" + +[[package]] +name = "embedded-io" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edd0f118536f44f5ccd48bcb8b111bdc3de888b58c74639dfb034a357d0f206d" + +[[package]] +name = "encode_unicode" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" + +[[package]] +name = "encoding_rs" +version = "0.8.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75030f3c4f45dafd7586dd6780965a8c7e8e285a5ecb86713e63a79c5b2766f3" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "entities" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5320ae4c3782150d900b79807611a59a99fc9a1d61d686faafc24b93fc8d7ca" + +[[package]] +name = "enum-as-inner" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1e6a265c649f3f5979b601d26f1d05ada116434c87741c9493cb56218f76cbc" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "enum-assoc" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ed8956bd5c1f0415200516e78ff07ec9e16415ade83c056c230d7b7ea0d55b7" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "enumset" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "839c4174b41e75c8f7306110b2c51996a293b8d1d850edd529011841d9fede7d" +dependencies = [ + "enumset_derive", + "serde", +] + +[[package]] +name = "enumset_derive" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4bd536557b58c682b217b8fb199afdff47cd3eff260623f19e77074eb073d63a" +dependencies = [ + "darling 0.21.3", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "env_filter" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "900d271a03799a1ee8d1ca9b19893b48ca674a9284fefcfb85f05e74ed314217" +dependencies = [ + "log", + "regex", +] + +[[package]] +name = "env_home" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f84e12ccf0a7ddc17a6c41c93326024c42920d7ee630d04950e6926645c0fe" + +[[package]] +name = "env_logger" +version = "0.11.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de671bd27a75a797dc9ae289ba1e77276e75e2026408aab65185384e2d5cd3f6" +dependencies = [ + "anstream", + "anstyle", + "env_filter", + "jiff", + "log", +] + +[[package]] +name = "equator" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4711b213838dfee0117e3be6ac926007d7f433d7bbe33595975d4190cb07e6fc" +dependencies = [ + "equator-macro", +] + +[[package]] +name = "equator-macro" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44f23cf4b44bfce11a86ace86f8a73ffdec849c9fd00a386a53d278bd9e81fb3" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" + +[[package]] +name = "etcetera" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de48cc4d1c1d97a20fd819def54b890cadde72ed3ad0c614822a0a433361be96" +dependencies = [ + "cfg-if", + "windows-sys 0.61.2", +] + +[[package]] +name = "ethnum" +version = "1.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40404c3f5f511ec4da6fe866ddf6a717c309fdbb69fbbad7b0f3edab8f2e835f" + +[[package]] +name = "euclid" +version = "0.22.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1a05365e3b1c6d1650318537c7460c6923f1abdd272ad6842baa2b509957a06" +dependencies = [ + "num-traits", +] + +[[package]] +name = "event-listener" +version = "5.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13b66accf52311f30a0db42147dadea9850cb48cd070028831ae5f5d4b856ab" +dependencies = [ + "concurrent-queue", + "parking", + "pin-project-lite", +] + +[[package]] +name = "event-listener-strategy" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93" +dependencies = [ + "event-listener", + "pin-project-lite", +] + +[[package]] +name = "exr" +version = "1.74.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6be87932f10230a4339ab394edd8e4611fcb72553d8295b4d52ea55249b3daa5" +dependencies = [ + "bit_field", + "half", + "lebe", + "miniz_oxide", + "num-complex", + "pulp 0.22.3", + "rayon-core", + "smallvec", + "zune-inflate", +] + +[[package]] +name = "fake" +version = "5.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea6be833b323a56361118a747470a45a1bcd5c52a2ec9b1e40c83dafe687e453" +dependencies = [ + "deunicode", + "either", + "rand 0.10.2", +] + +[[package]] +name = "fallible-iterator" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" + +[[package]] +name = "fallible-streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" + +[[package]] +name = "fancy-regex" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b95f7c0680e4142284cf8b22c14a476e87d61b004a3a0861872b32ef7ead40a2" +dependencies = [ + "bit-set 0.5.3", + "regex", +] + +[[package]] +name = "fast-float2" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8eb564c5c7423d25c886fb561d1e4ee69f72354d16918afa32c08811f6b6a55" + +[[package]] +name = "fast-srgb8" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd2e7510819d6fbf51a5545c8f922716ecfb14df168a3242f7d33e0239efe6a1" + +[[package]] +name = "fastrand" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" + +[[package]] +name = "fax" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caf1079563223d5d59d83c85886a56e586cfd5c1a26292e971a0fa266531ac5a" + +[[package]] +name = "fdeflate" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e6853b52649d4ac5c0bd02320cddc5ba956bdb407c4b75a2c6b75bf51500f8c" +dependencies = [ + "simd-adler32", +] + +[[package]] +name = "fiat-crypto" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64cd1e32ddd350061ae6edb1b082d7c54915b5c672c389143b9a63403a109f24" + +[[package]] +name = "filedescriptor" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e40758ed24c9b2eeb76c35fb0aebc66c626084edd827e07e1552279814c6682d" +dependencies = [ + "libc", + "thiserror 1.0.69", + "winapi", +] + +[[package]] +name = "filetime" +version = "0.2.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c287a33c7f0a620c38e641e7f60827713987b3c0f26e8ddc9462cc69cf75759" +dependencies = [ + "cfg-if", + "libc", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "finl_unicode" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9844ddc3a6e533d62bba727eb6c28b5d360921d5175e9ff0f1e621a5c590a4d5" + +[[package]] +name = "fixedbitset" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80" + +[[package]] +name = "fixedbitset" +version = "0.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" + +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "crc32fast", + "miniz_oxide", + "zlib-rs", +] + +[[package]] +name = "float-cmp" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b09cf3155332e944990140d967ff5eceb70df778b34f77d8075db46e4704e6d8" +dependencies = [ + "num-traits", +] + +[[package]] +name = "float-ord" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ce81f49ae8a0482e4c55ea62ebbd7e5a686af544c00b9d090bba3ff9be97b3d" + +[[package]] +name = "float4" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a5404bf31d22893d61cf24d4dda149d8e6b2ff07601c3cb3be651031f61a4ed" + +[[package]] +name = "float8" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2d1f04709a8ac06e8e8042875a3c466cc4832d3c1a18dbcb9dba3c6e83046bc" +dependencies = [ + "half", + "num-traits", + "rand 0.9.4", + "rand_distr 0.5.1", +] + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + +[[package]] +name = "foreign-types" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d737d9aa519fb7b749cbc3b962edcf310a8dd1f4b67c91c4f83975dbdd17d965" +dependencies = [ + "foreign-types-macros", + "foreign-types-shared", +] + +[[package]] +name = "foreign-types-macros" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a5c6c585bc94aaf2c7b51dd4c2ba22680844aba4c687be581871a6f518c5742" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "foreign-types-shared" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa9a19cbb55df58761df49b23516a86d432839add4af60fc256da840f66ed35b" + +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "fs2" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9564fc758e15025b46aa6643b1b77d047d1a56a1aea6e01002ac0c7026876213" +dependencies = [ + "libc", + "winapi", +] + +[[package]] +name = "fs4" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8640e34b88f7652208ce9e88b1a37a2ae95227d84abec377ccd3c5cfeb141ed4" +dependencies = [ + "rustix", + "windows-sys 0.59.0", +] + +[[package]] +name = "futures" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-buffered" +version = "0.2.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4421cb78ee172b6b06080093479d3c50f058e7c81b7d577bbb8d118d551d4cd5" +dependencies = [ + "cordyceps", + "diatomic-waker", + "futures-core", + "pin-project-lite", + "spin", +] + +[[package]] +name = "futures-channel" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-executor" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" + +[[package]] +name = "futures-lite" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f78e10609fe0e0b3f4157ffab1876319b5b0db102a2c60dc4626306dc46b44ad" +dependencies = [ + "fastrand", + "futures-core", + "futures-io", + "parking", + "pin-project-lite", +] + +[[package]] +name = "futures-macro" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "futures-sink" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-timer" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af43fadb8a98512d547e37b4e92e0ced13e205c061b87b4623eff01d918d6968" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "fuzzy-matcher" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "54614a3312934d066701a80f20f15fa3b56d67ac7722b39eea5b4c9dd1d66c94" +dependencies = [ + "thread_local", +] + +[[package]] +name = "gemm" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab96b703d31950f1aeddded248bc95543c9efc7ac9c4a21fda8703a83ee35451" +dependencies = [ + "dyn-stack", + "gemm-c32 0.18.2", + "gemm-c64 0.18.2", + "gemm-common 0.18.2", + "gemm-f16 0.18.2", + "gemm-f32 0.18.2", + "gemm-f64 0.18.2", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa0673db364b12263d103b68337a68fbecc541d6f6b61ba72fe438654709eacb" +dependencies = [ + "dyn-stack", + "gemm-c32 0.19.0", + "gemm-c64 0.19.0", + "gemm-common 0.19.0", + "gemm-f16 0.19.0", + "gemm-f32 0.19.0", + "gemm-f64 0.19.0", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-c32" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6db9fd9f40421d00eea9dd0770045a5603b8d684654816637732463f4073847" +dependencies = [ + "dyn-stack", + "gemm-common 0.18.2", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-c32" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "086936dbdcb99e37aad81d320f98f670e53c1e55a98bee70573e83f95beb128c" +dependencies = [ + "dyn-stack", + "gemm-common 0.19.0", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-c64" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfcad8a3d35a43758330b635d02edad980c1e143dc2f21e6fd25f9e4eada8edf" +dependencies = [ + "dyn-stack", + "gemm-common 0.18.2", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-c64" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20c8aeeeec425959bda4d9827664029ba1501a90a0d1e6228e48bef741db3a3f" +dependencies = [ + "dyn-stack", + "gemm-common 0.19.0", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-common" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a352d4a69cbe938b9e2a9cb7a3a63b7e72f9349174a2752a558a8a563510d0f3" +dependencies = [ + "bytemuck", + "dyn-stack", + "half", + "libm", + "num-complex", + "num-traits", + "once_cell", + "paste", + "pulp 0.21.5", + "raw-cpuid", + "rayon", + "seq-macro", + "sysctl", +] + +[[package]] +name = "gemm-common" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88027625910cc9b1085aaaa1c4bc46bb3a36aad323452b33c25b5e4e7c8e2a3e" +dependencies = [ + "bytemuck", + "dyn-stack", + "half", + "libm", + "num-complex", + "num-traits", + "once_cell", + "paste", + "pulp 0.22.3", + "raw-cpuid", + "rayon", + "seq-macro", + "sysctl", +] + +[[package]] +name = "gemm-f16" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cff95ae3259432f3c3410eaa919033cd03791d81cebd18018393dc147952e109" +dependencies = [ + "dyn-stack", + "gemm-common 0.18.2", + "gemm-f32 0.18.2", + "half", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "rayon", + "seq-macro", +] + +[[package]] +name = "gemm-f16" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3df7a55202e6cd6739d82ae3399c8e0c7e1402859b30e4cb780e61525d9486e" +dependencies = [ + "dyn-stack", + "gemm-common 0.19.0", + "gemm-f32 0.19.0", + "half", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "rayon", + "seq-macro", +] + +[[package]] +name = "gemm-f32" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc8d3d4385393304f407392f754cd2dc4b315d05063f62cf09f47b58de276864" +dependencies = [ + "dyn-stack", + "gemm-common 0.18.2", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-f32" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02e0b8c9da1fbec6e3e3ab2ce6bc259ef18eb5f6f0d3e4edf54b75f9fd41a81c" +dependencies = [ + "dyn-stack", + "gemm-common 0.19.0", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-f64" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35b2a4f76ce4b8b16eadc11ccf2e083252d8237c1b589558a49b0183545015bd" +dependencies = [ + "dyn-stack", + "gemm-common 0.18.2", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "gemm-f64" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "056131e8f2a521bfab322f804ccd652520c79700d81209e9d9275bbdecaadc6a" +dependencies = [ + "dyn-stack", + "gemm-common 0.19.0", + "num-complex", + "num-traits", + "paste", + "raw-cpuid", + "seq-macro", +] + +[[package]] +name = "generator" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3b854b0e584ead1a33f18b2fcad7cf7be18b3875c78816b753639aa501513ae" +dependencies = [ + "cc", + "cfg-if", + "libc", + "log", + "rustversion", + "windows-link 0.2.1", + "windows-result 0.4.1", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "wasi", + "wasm-bindgen", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "r-efi 5.3.0", + "wasip2", + "wasm-bindgen", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "r-efi 6.0.0", + "rand_core 0.10.1", + "wasm-bindgen", +] + +[[package]] +name = "ghash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0d8a4362ccb29cb0b265253fb0a2728f592895ee6854fd9bc13f2ffda266ff1" +dependencies = [ + "opaque-debug", + "polyval", +] + +[[package]] +name = "gif" +version = "0.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee8cfcc411d9adbbaba82fb72661cc1bcca13e8bba98b364e62b2dba8f960159" +dependencies = [ + "color_quant", + "weezl", +] + +[[package]] +name = "gimli" +version = "0.32.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e629b9b98ef3dd8afe6ca2bd0f89306cec16d43d907889945bc5d6687f2f13c7" + +[[package]] +name = "gix-features" +version = "0.48.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1849ae154d38bc403185be14fa871e38e3c93ee606875d94e207fdb9fba52dbc" +dependencies = [ + "gix-trace", + "gix-utils", + "libc", +] + +[[package]] +name = "gix-fs" +version = "0.21.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6cdff46db8798e47e2f727d84b9379aac5add3dd3d9d0b07bb4d7d5d640771fe" +dependencies = [ + "bstr", + "fastrand", + "gix-features", + "gix-path", + "gix-utils", + "thiserror 2.0.18", +] + +[[package]] +name = "gix-path" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "afa6ac14cd14939ea94a496ce7460daa6511c09f5b84757e9cfc6f9c8d0f93a6" +dependencies = [ + "bstr", + "gix-trace", + "gix-validate", + "thiserror 2.0.18", +] + +[[package]] +name = "gix-tempfile" +version = "23.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ef60812443484e67bf84e444cc71b4c78ae62deb822221774a4fa0c57fdb17f" +dependencies = [ + "dashmap", + "gix-fs", + "libc", + "parking_lot", + "signal-hook 0.4.4", + "signal-hook-registry", + "tempfile", +] + +[[package]] +name = "gix-trace" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44dc45eae785c0eb14173e0f152e6e224dcf4d45b6a6999a3aed22af541ad678" + +[[package]] +name = "gix-utils" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66c50966184123caf580ffa64e28031a878597f1c7fceb8fe19566c38eb1b771" +dependencies = [ + "fastrand", + "unicode-normalization", +] + +[[package]] +name = "gix-validate" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7bc6fc771c4063ba7cd2f47b91fb6076251c6a823b64b7fe7b8874b0fe4afae3" +dependencies = [ + "bstr", +] + +[[package]] +name = "gl_generator" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a95dfc23a2b4a9a2f5ab41d194f8bfda3cabec42af4e39f08c339eb2a0c124d" +dependencies = [ + "khronos_api", + "log", + "xml-rs", +] + +[[package]] +name = "glob" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + +[[package]] +name = "globset" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52dfc19153a48bde0cbd630453615c8151bce3a5adfac7a0aebfbf0a1e1f57e3" +dependencies = [ + "aho-corasick", + "bstr", + "log", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "globwalk" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf760ebf69878d9fd8f110c89703d90ce35095324d1f1edcb595c63945ee757" +dependencies = [ + "bitflags 2.13.0", + "ignore", + "walkdir", +] + +[[package]] +name = "gloo-timers" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbb143cf96099802033e0d4f4963b19fd2e0b728bcf076cd9cf7f6634f092994" +dependencies = [ + "futures-channel", + "futures-core", + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "glow" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29038e1c483364cc6bb3cf78feee1816002e127c331a1eec55a4d202b9e1adb5" +dependencies = [ + "js-sys", + "slotmap", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "glutin_wgl_sys" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c4ee00b289aba7a9e5306d57c2d05499b2e5dc427f84ac708bd2c090212cf3e" +dependencies = [ + "gl_generator", +] + +[[package]] +name = "gpu-allocator" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51255ea7cfaadb6c5f1528d43e92a82acb2b96c43365989a28b2d44ee38f8795" +dependencies = [ + "ash", + "hashbrown 0.16.1", + "log", + "presser", + "thiserror 2.0.18", + "windows 0.62.2", +] + +[[package]] +name = "guide" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "log", +] + +[[package]] +name = "h2" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155" +dependencies = [ + "atomic-waker", + "bytes", + "fnv", + "futures-core", + "futures-sink", + "http", + "indexmap", + "slab", + "tokio", + "tokio-util", + "tracing", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "bytemuck", + "cfg-if", + "crunchy", + "num-traits", + "rand 0.9.4", + "rand_distr 0.5.1", + "serde", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43a3c133739dddd0d2990f9a4bdf8eb4b21ef50e4851ca85ab661199821d510e" +dependencies = [ + "ahash", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.1.5", +] + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", + "rayon", + "serde", + "serde_core", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", + "serde", + "serde_core", +] + +[[package]] +name = "hashlink" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7382cf6263419f2d8df38c55d7da83da5c18aef87fc7a7fc1fb1e344edfe14c1" +dependencies = [ + "hashbrown 0.15.5", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + +[[package]] +name = "hf-hub" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "629d8f3bbeda9d148036d6b0de0a3ab947abd08ce90626327fc3547a49d59d97" +dependencies = [ + "dirs", + "http", + "indicatif 0.17.11", + "libc", + "log", + "rand 0.9.4", + "serde", + "serde_json", + "thiserror 2.0.18", + "ureq 2.12.1", + "windows-sys 0.60.2", +] + +[[package]] +name = "hickory-net" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2295ed2f9c31e471e1428a8f88a3f0e1f4b27c15049592138d1eebe9c35b183" +dependencies = [ + "async-trait", + "bytes", + "cfg-if", + "data-encoding", + "futures-channel", + "futures-io", + "futures-util", + "h2", + "hickory-proto", + "http", + "idna", + "ipnet", + "jni 0.22.4", + "rand 0.10.2", + "rustls", + "thiserror 2.0.18", + "tinyvec", + "tokio", + "tokio-rustls", + "tracing", + "url", +] + +[[package]] +name = "hickory-proto" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bab31817bfb44672a252e97fe81cd0c18d1b2cf892108922f6818820df8c643" +dependencies = [ + "data-encoding", + "idna", + "ipnet", + "jni 0.22.4", + "once_cell", + "prefix-trie", + "rand 0.10.2", + "ring", + "thiserror 2.0.18", + "tinyvec", + "tracing", + "url", +] + +[[package]] +name = "hickory-resolver" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0d58d28879ceecde6607729660c2667a081ccdc082e082675042793960f178c" +dependencies = [ + "cfg-if", + "futures-util", + "hickory-net", + "hickory-proto", + "ipconfig", + "ipnet", + "jni 0.22.4", + "moka", + "ndk-context", + "once_cell", + "parking_lot", + "rand 0.10.2", + "resolv-conf", + "rustls", + "smallvec", + "system-configuration", + "thiserror 2.0.18", + "tokio", + "tokio-rustls", + "tracing", +] + +[[package]] +name = "hmac" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" +dependencies = [ + "digest 0.10.7", +] + +[[package]] +name = "hmac" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6303bc9732ae41b04cb554b844a762b4115a61bfaa81e3e83050991eeb56863f" +dependencies = [ + "digest 0.11.3", +] + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "hound" +version = "3.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62adaabb884c94955b19907d60019f4e145d091c75345379e70d1ee696f7854f" + +[[package]] +name = "http" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + +[[package]] +name = "humantime" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15cdd26707701c53297e2fa6afb323d55fbc1d0810c3aec078ae3ef0424c3c15" + +[[package]] +name = "hybrid-array" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "818356c5132c1fede50f837ca96afbe78ff42413047f4abb886217845e1b6c8c" +dependencies = [ + "typenum", +] + +[[package]] +name = "hyper" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "h2", + "http", + "http-body", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "smallvec", + "tokio", + "want", +] + +[[package]] +name = "hyper-rustls" +version = "0.27.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" +dependencies = [ + "http", + "hyper", + "hyper-util", + "rustls", + "rustls-native-certs", + "tokio", + "tokio-rustls", + "tower-service", + "webpki-roots 1.0.8", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-channel", + "futures-util", + "http", + "http-body", + "hyper", + "ipnet", + "libc", + "percent-encoding", + "pin-project-lite", + "socket2", + "tokio", + "tower-service", + "tracing", +] + +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core 0.62.2", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke 0.8.3", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke 0.8.3", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + +[[package]] +name = "identity-hash" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfdd7caa900436d8f13b2346fe10257e0c05c1f1f9e351f4f5d57c03bd5f45da" + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "igd-next" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de7238d487a9aff61f81b5ab41c0a841532a115a398b5fa92a2fadd0885e2581" +dependencies = [ + "attohttpc", + "bytes", + "futures", + "http", + "http-body-util", + "hyper", + "hyper-util", + "log", + "rand 0.10.2", + "tokio", + "url", + "xmltree", +] + +[[package]] +name = "ignore" +version = "0.4.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe112b004901c62c2faa11f4f75e9864e0cc5af8da71c9115d184a3aa888749f" +dependencies = [ + "crossbeam-deque", + "globset", + "log", + "memchr", + "regex-automata", + "same-file", + "walkdir", + "winapi-util", +] + +[[package]] +name = "image" +version = "0.25.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85ab80394333c02fe689eaf900ab500fbd0c2213da414687ebf995a65d5a6104" +dependencies = [ + "bytemuck", + "byteorder-lite", + "color_quant", + "exr", + "gif", + "image-webp", + "moxcms", + "num-traits", + "png", + "qoi", + "ravif", + "rayon", + "rgb", + "tiff", + "zune-core", + "zune-jpeg", +] + +[[package]] +name = "image-webp" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "525e9ff3e1a4be2fbea1fdf0e98686a6d98b4d8f937e1bf7402245af1909e8c3" +dependencies = [ + "byteorder-lite", + "quick-error", +] + +[[package]] +name = "imgref" +version = "1.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89194689a993ab15268672e99e7b0e19da2da3268ac682e8f02d29d4d1434cd7" + +[[package]] +name = "import-model-weights" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "burn-store", +] + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", + "serde", + "serde_core", +] + +[[package]] +name = "indicatif" +version = "0.17.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" +dependencies = [ + "console 0.15.11", + "number_prefix", + "portable-atomic", + "unicode-width", + "web-time", +] + +[[package]] +name = "indicatif" +version = "0.18.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9433806cd6b4ec1aba79c021c7e4c58fb4c3b9977c085062e611ac929998fb0c" +dependencies = [ + "console 0.16.4", + "portable-atomic", + "unicode-width", + "unit-prefix", + "web-time", +] + +[[package]] +name = "indoc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +dependencies = [ + "rustversion", +] + +[[package]] +name = "inflections" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a257582fdcde896fd96463bf2d40eefea0580021c0712a0e2b028b60b47a837a" + +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "generic-array", +] + +[[package]] +name = "inout" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4250ce6452e92010fdf7268ccc5d14faa80bb12fc741938534c58f16804e03c7" +dependencies = [ + "hybrid-array", +] + +[[package]] +name = "inquire" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6654738b8024300cf062d04a1c13c10c8e2cea598ec1c47dc9b6641159429756" +dependencies = [ + "bitflags 2.13.0", + "crossterm", + "dyn-clone", + "fuzzy-matcher", + "unicode-segmentation", + "unicode-width", +] + +[[package]] +name = "instability" +version = "0.3.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5eb2d60ef19920a3a9193c3e371f726ec1dafc045dac788d0fb3704272458971" +dependencies = [ + "darling 0.23.0", + "indoc", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "internment" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "636d4b0f6a39fd684effe2a73f5310df16a3fa7954c26d36833e98f44d1977a2" +dependencies = [ + "hashbrown 0.15.5", + "serde", +] + +[[package]] +name = "interpolate_name" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c34819042dc3d3971c46c2190835914dfbe0c3c13f61449b2997f4e9722dfa60" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "ipconfig" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d40460c0ce33d6ce4b0630ad68ff63d6661961c48b6dba35e5a4d81cfb48222" +dependencies = [ + "socket2", + "widestring", + "windows-registry", + "windows-result 0.4.1", + "windows-sys 0.61.2", +] + +[[package]] +name = "ipnet" +version = "2.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" +dependencies = [ + "serde", +] + +[[package]] +name = "iroh" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fca9b4b462c343ff88fc0af4096c186f939b602a0bc08723536ef2c31c93971" +dependencies = [ + "backon", + "blake3", + "bytes", + "cfg_aliases", + "ctutils", + "data-encoding", + "derive_more", + "ed25519-dalek", + "futures-util", + "getrandom 0.4.3", + "hickory-resolver", + "http", + "ipnet", + "iroh-base", + "iroh-dns", + "iroh-metrics", + "iroh-relay", + "n0-error", + "n0-future", + "n0-watcher", + "netwatch", + "noq", + "noq-proto", + "noq-udp", + "papaya", + "pin-project", + "portable-atomic", + "portmapper", + "rand 0.10.2", + "reqwest 0.13.4", + "rustc-hash 2.1.3", + "rustls", + "rustls-pki-types", + "serde", + "smallvec", + "strum 0.28.0", + "time", + "tokio", + "tokio-stream", + "tokio-util", + "tracing", + "url", + "wasm-bindgen-futures", +] + +[[package]] +name = "iroh-base" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830a582cd54410dc1aa71d4786a82c3297d7b0165accd8b6dbbb3b240b48140d" +dependencies = [ + "curve25519-dalek", + "data-encoding", + "data-encoding-macro", + "derive_more", + "ed25519-dalek", + "getrandom 0.4.3", + "n0-error", + "rand 0.10.2", + "serde", + "url", + "zeroize", +] + +[[package]] +name = "iroh-dns" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "516e4eedc38e33ab69a6bd325520332dc3d67b25454e2d590ebb84a25240dd9a" +dependencies = [ + "arc-swap", + "cfg_aliases", + "derive_more", + "hickory-resolver", + "iroh-base", + "n0-error", + "n0-future", + "ndk-context", + "portable-atomic", + "rand 0.10.2", + "rustls", + "simple-dns", + "strum 0.28.0", + "tokio", + "tracing", + "url", +] + +[[package]] +name = "iroh-metrics" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291065721ad7c477b972e581bbc528df031dc8eb5e39fe1ff3300ae5dfb157ef" +dependencies = [ + "iroh-metrics-derive", + "itoa", + "n0-error", + "portable-atomic", + "ryu", + "serde", + "tracing", +] + +[[package]] +name = "iroh-metrics-derive" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae5f0c4405d1fbc9fb16ff422ca40620e93dc36c30ecaba0c2aee3992b7bd48" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "iroh-relay" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8149bb6a57126225a07d6928846d82dcedfd24ea0f863ef7b2eb475e1d726354" +dependencies = [ + "blake3", + "bytes", + "cfg_aliases", + "data-encoding", + "derive_more", + "getrandom 0.4.3", + "hickory-resolver", + "http", + "http-body-util", + "hyper", + "hyper-util", + "iroh-base", + "iroh-dns", + "iroh-metrics", + "lru", + "n0-error", + "n0-future", + "noq", + "noq-proto", + "num_enum", + "pin-project", + "postcard", + "rand 0.10.2", + "reqwest 0.13.4", + "rustls", + "rustls-pki-types", + "serde", + "serde_bytes", + "strum 0.28.0", + "tokio", + "tokio-rustls", + "tokio-util", + "tokio-websockets", + "tracing", + "url", + "vergen-gitcl", + "webpki-roots 1.0.8", + "ws_stream_wasm", +] + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "jiff" +version = "0.2.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "961d16382652bfdd8c6f68b223b26a8c93e0d475c672f414411db31c6c5c900e" +dependencies = [ + "defmt", + "jiff-static", + "log", + "portable-atomic", + "portable-atomic-util", + "serde_core", +] + +[[package]] +name = "jiff-static" +version = "0.2.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0879bd39df99c4c5e2c6615ccc026391a423dde10532c573e6086eb94a802cc" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "jni" +version = "0.21.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a87aa2bb7d2af34197c04845522473242e1aa17c12f4935d5856491a7fb8c97" +dependencies = [ + "cesu8", + "cfg-if", + "combine", + "jni-sys 0.3.1", + "log", + "thiserror 1.0.69", + "walkdir", + "windows-sys 0.45.0", +] + +[[package]] +name = "jni" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5efd9a482cf3a427f00d6b35f14332adc7902ce91efb778580e180ff90fa3498" +dependencies = [ + "cfg-if", + "combine", + "jni-macros", + "jni-sys 0.4.1", + "log", + "simd_cesu8", + "thiserror 2.0.18", + "walkdir", + "windows-link 0.2.1", +] + +[[package]] +name = "jni-macros" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a00109accc170f0bdb141fed3e393c565b6f5e072365c3bd58f5b062591560a3" +dependencies = [ + "proc-macro2", + "quote", + "rustc_version", + "simd_cesu8", + "syn 2.0.118", +] + +[[package]] +name = "jni-sys" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41a652e1f9b6e0275df1f15b32661cf0d4b78d4d87ddec5e0c3c20f097433258" +dependencies = [ + "jni-sys 0.4.1", +] + +[[package]] +name = "jni-sys" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6377a88cb3910bee9b0fa88d4f42e1d2da8e79915598f65fb0c7ee14c878af2" +dependencies = [ + "jni-sys-macros", +] + +[[package]] +name = "jni-sys-macros" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" +dependencies = [ + "quote", + "syn 2.0.118", +] + +[[package]] +name = "jobserver" +version = "0.1.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" +dependencies = [ + "getrandom 0.4.3", + "libc", +] + +[[package]] +name = "js-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "kasuari" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bde5057d6143cc94e861d90f591b9303d6716c6b9602309150bd068853c10899" +dependencies = [ + "hashbrown 0.16.1", + "portable-atomic", + "thiserror 2.0.18", +] + +[[package]] +name = "khronos-egl" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6aae1df220ece3c0ada96b8153459b67eebe9ae9212258bb0134ae60416fdf76" +dependencies = [ + "libc", + "libloading 0.8.9", + "pkg-config", +] + +[[package]] +name = "khronos_api" +version = "3.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2db585e1d738fc771bf08a151420d3ed193d9d895a36df7f6f8a9456b911ddc" + +[[package]] +name = "kstat-rs" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27964e4632377753acb0898ce6f28770d50cbca1339200ae63d700cff97b5c2b" +dependencies = [ + "libc", + "thiserror 1.0.69", +] + +[[package]] +name = "lab" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf36173d4167ed999940f804952e6b08197cae5ad5d572eb4db150ce8ad5d58f" + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "lebe" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a79a3332a6609480d7d0c9eab957bca6b455b91bb84e66d19f5ff66294b85b8" + +[[package]] +name = "libbz2-rs-sys" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34b357333733e8260735ba5894eb928c02ecc69c78715f01a8019e7fa7f2db4c" + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "libfuzzer-sys" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9fd2f41a1cba099f79a0b6b6c35656cf7c03351a7bae8ff0f28f25270f929d2" +dependencies = [ + "arbitrary", + "cc", +] + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link 0.2.1", +] + +[[package]] +name = "libloading" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "754ca22de805bb5744484a5b151a9e1a8e837d5dc232c2d7d8c2e3492edc8b60" +dependencies = [ + "cfg-if", + "windows-link 0.2.1", +] + +[[package]] +name = "liblzma" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45aec2360b3933207e27908049d8e4df4e476b58180afb1e56b2a4fb72efe4ba" +dependencies = [ + "liblzma-sys", + "num_cpus", +] + +[[package]] +name = "liblzma-sys" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a046c7f353ba30f810545151e04f63545833803f5b86ee3ddf1517247fe560a5" +dependencies = [ + "cc", + "libc", + "pkg-config", +] + +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + +[[package]] +name = "libredox" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c943259e342f1e06ff2da7a83eabdfe7f92ce10262688dbf1895ff0b3e6e4652" +dependencies = [ + "libc", +] + +[[package]] +name = "libsqlite3-sys" +version = "0.35.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "133c182a6a2c87864fe97778797e46c7e999672690dc9fa3ee8e241aa4a9c13f" +dependencies = [ + "cc", + "pkg-config", + "vcpkg", +] + +[[package]] +name = "line-clipping" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f50e8f47623268b5407192d26876c4d7f89d686ca130fdc53bced4814cd29f8" +dependencies = [ + "bitflags 2.13.0", +] + +[[package]] +name = "link-section" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b685d66585d646efe09fec763d796c291049c8b6bf84e04954bffc8748341f0d" + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + +[[package]] +name = "litrs" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "loom" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "419e0dc8046cb947daa77eb95ae174acfbddb7673b4151f56d1eed8e93fbfaca" +dependencies = [ + "cfg-if", + "generator", + "scoped-tls", + "tracing", + "tracing-subscriber", +] + +[[package]] +name = "loop9" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fae87c125b03c1d2c0150c90365d7d6bcc53fb73a9acaef207d2d065860f062" +dependencies = [ + "imgref", +] + +[[package]] +name = "lru" +version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b6180140927ee907000b0aa540091f6ea512ead4447c92b8fc35bc72788a5a6" +dependencies = [ + "hashbrown 0.17.1", +] + +[[package]] +name = "lru-slab" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" + +[[package]] +name = "lz4" +version = "1.28.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a20b523e860d03443e98350ceaac5e71c6ba89aea7d960769ec3ce37f4de5af4" +dependencies = [ + "lz4-sys", +] + +[[package]] +name = "lz4-sys" +version = "1.11.1+lz4-1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bd8c0d6c6ed0cd30b3652886bb8711dc4bb01d637a68105a3d5158039b418e6" +dependencies = [ + "cc", + "libc", +] + +[[package]] +name = "lzma-rust2" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c60a23ffb90d527e23192f1246b14746e2f7f071cb84476dd879071696c18a4a" +dependencies = [ + "crc", + "sha2 0.10.9", +] + +[[package]] +name = "lzma-rust2" +version = "0.16.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca93e534d1142d1d0dcca6d25fe302508a5dfb40b302802904577725ea0b695b" +dependencies = [ + "sha2 0.11.0", +] + +[[package]] +name = "mac-addr" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3d25b0e0b648a86960ac23b7ad4abb9717601dec6f66c165f5b037f3f03065f" + +[[package]] +name = "mac_address" +version = "1.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0aeb26bf5e836cc1c341c8106051b573f1766dfa05aa87f0b98be5e51b02303" +dependencies = [ + "nix 0.29.0", + "winapi", +] + +[[package]] +name = "macerator" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "508a2f720538bb7e3ea3cb6d098615b107cb82ae387d77d906276905e9ec894b" +dependencies = [ + "bytemuck", + "cfg_aliases", + "half", + "macerator-macros", + "moddef", + "num-traits", + "paste", + "rustc_version", +] + +[[package]] +name = "macerator-macros" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5ce85961d618ce9794bdf822bfe96fe9dd341aa5b033b454f7a8d96e79b9b1" +dependencies = [ + "darling 0.20.11", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "mach2" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dae608c151f68243f2b000364e1f7b186d9c29845f7d2d85bd31b9ad77ad552b" + +[[package]] +name = "macro_rules_attribute" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520" +dependencies = [ + "macro_rules_attribute-proc_macro", + "paste", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" + +[[package]] +name = "malloc_buf" +version = "0.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62bb907fe88d54d8d9ce32a3cceab4218ed2f6b7d35617cafe9adf84e43919cb" +dependencies = [ + "libc", +] + +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + +[[package]] +name = "matchit" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" + +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "num_cpus", + "once_cell", + "rawpointer", + "thread-tree", +] + +[[package]] +name = "maybe-rayon" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ea1f30cedd69f0a2954655f7188c6a834246d2bcf1e315e2ac40c4b24dc9519" +dependencies = [ + "cfg-if", + "rayon", +] + +[[package]] +name = "md5" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ebb8d8732c6a6df3d8f032a82911cfc747e00efb95cc46e8d0acd5b5b88570c" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "memmap2" +version = "0.9.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" +dependencies = [ + "libc", + "stable_deref_trait", +] + +[[package]] +name = "memmem" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a64a92489e2744ce060c349162be1c5f33c6969234104dbd99ddb5feb08b8c15" + +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + +[[package]] +name = "metal" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ecfd3296f8c56b7c1f6fbac3c71cefa9d78ce009850c45000015f206dc7fa21" +dependencies = [ + "bitflags 2.13.0", + "block", + "core-graphics-types", + "foreign-types", + "log", + "objc", + "paste", +] + +[[package]] +name = "mime" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + +[[package]] +name = "mio" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" +dependencies = [ + "libc", + "log", + "wasi", + "windows-sys 0.61.2", +] + +[[package]] +name = "mnist" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "log", + "rand 0.10.2", +] + +[[package]] +name = "mnist-inference-web" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "console_error_panic_hook", + "js-sys", + "serde", + "wasm-bindgen", + "wasm-bindgen-futures", +] + +[[package]] +name = "moddef" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4a0b3262dc837d2513fe2ef31ff8461352ef932dcca31ba0c0abe33547cf6b9b" + +[[package]] +name = "modern-lstm" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "planus", + "polars", + "rand 0.10.2", + "rand_distr 0.6.0", + "serde", +] + +[[package]] +name = "moka" +version = "0.12.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "957228ad12042ee839f93c8f257b62b4c0ab5eaae1d4fa60de53b27c9d7c5046" +dependencies = [ + "crossbeam-channel", + "crossbeam-epoch", + "crossbeam-utils", + "equivalent", + "parking_lot", + "portable-atomic", + "smallvec", + "tagptr", + "uuid", +] + +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "moxcms" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb85c154ba489f01b25c0d36ae69a87e4a1c73a72631fc6c0eb6dde34a73e44b" +dependencies = [ + "num-traits", + "pxfm", +] + +[[package]] +name = "mutants" +version = "0.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc0287524726960e07b119cebd01678f852f147742ae0d925e6a520dca956126" + +[[package]] +name = "n0-error" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c37e81176a83a77d2514528b91bdafc70ef88aab428f0e1b91aebb8d99888895" +dependencies = [ + "n0-error-macros", + "spez", +] + +[[package]] +name = "n0-error-macros" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2acd8b070213b0299282f884b4beba4e7b52d624fdcd504a3ad3665390c11e1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "n0-future" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2ab99dfb861450e68853d34ae665243a88b8c493d01ba957321a1e9b2312bbe" +dependencies = [ + "cfg_aliases", + "derive_more", + "futures-buffered", + "futures-lite", + "futures-util", + "js-sys", + "pin-project", + "send_wrapper", + "tokio", + "tokio-util", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-time", +] + +[[package]] +name = "n0-watcher" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbc618745ad0b7414b149d0517ad8b5573b2fb4d4e2717add3d2446ce1fdd826" +dependencies = [ + "derive_more", + "n0-error", + "n0-future", +] + +[[package]] +name = "naga" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23bf0a141a9ab6f07dbb492db53245e464bc9db42f407772d9ae03d83a2c1033" +dependencies = [ + "arrayvec", + "bit-set 0.10.0", + "bitflags 2.13.0", + "cfg-if", + "cfg_aliases", + "codespan-reporting", + "half", + "hashbrown 0.17.1", + "indexmap", + "libm", + "log", + "naga-types", + "num-traits", + "once_cell", + "rustc-hash 1.1.0", + "spirv", + "thiserror 2.0.18", + "unicode-ident", +] + +[[package]] +name = "naga-types" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "658200ddc25c6c7b860747516d132d1b284c0fafb7a380233acee9a72fb30e11" +dependencies = [ + "hashbrown 0.17.1", + "indexmap", + "rustc-hash 1.1.0", + "thiserror 2.0.18", +] + +[[package]] +name = "nb" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "801d31da0513b6ec5214e9bf433a77966320625a37860f910be265be6e18d06f" +dependencies = [ + "nb 1.1.0", +] + +[[package]] +name = "nb" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d5439c4ad607c3c23abf66de8c8bf57ba8adcd1f129e699851a6e43935d339d" + +[[package]] +name = "ndarray" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + +[[package]] +name = "ndarray" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "cblas-sys", + "libc", + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", + "rayon", +] + +[[package]] +name = "ndk-context" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27b02d87554356db9e9a873add8782d4ea6e3e58ea071a9adb9a2e8ddb884a8b" + +[[package]] +name = "ndk-sys" +version = "0.6.0+11769913" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee6cda3051665f1fb8d9e08fc35c96d5a244fb1be711a03b71118828afc9a873" +dependencies = [ + "jni-sys 0.3.1", +] + +[[package]] +name = "netdev" +version = "0.45.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "569dfbdd2efd771b24ec9bb57f956e04d4fbfc72f62b2f11961723f9b3f4b020" +dependencies = [ + "block2", + "dispatch2", + "dlopen2", + "ipnet", + "jni 0.21.1", + "libc", + "mac-addr", + "ndk-context", + "netlink-packet-core", + "netlink-packet-route", + "netlink-sys", + "objc2", + "objc2-core-foundation", + "objc2-core-wlan", + "objc2-foundation", + "objc2-system-configuration", + "once_cell", + "plist", + "windows-sys 0.61.2", +] + +[[package]] +name = "netlib-src" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "abe2058825b000774d9c99f502accb07e01692d58fc14b92a808b65f82978fbd" +dependencies = [ + "cmake", +] + +[[package]] +name = "netlink-packet-core" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3463cbb78394cb0141e2c926b93fc2197e473394b761986eca3b9da2c63ae0f4" +dependencies = [ + "paste", +] + +[[package]] +name = "netlink-packet-route" +version = "0.31.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2288fcb784eb3defd5fb16f4c4160d5f477de192eac730f43e1d11c24d9a007" +dependencies = [ + "bitflags 2.13.0", + "libc", + "log", + "netlink-packet-core", +] + +[[package]] +name = "netlink-proto" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b65d130ee111430e47eed7896ea43ca693c387f097dd97376bffafbf25812128" +dependencies = [ + "bytes", + "futures", + "log", + "netlink-packet-core", + "netlink-sys", + "thiserror 2.0.18", +] + +[[package]] +name = "netlink-sys" +version = "0.8.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd6c30ed10fa69cc491d491b85cc971f6bdeb8e7367b7cde2ee6cc878d583fae" +dependencies = [ + "bytes", + "futures-util", + "libc", + "log", + "tokio", +] + +[[package]] +name = "netwatch" +version = "0.19.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d9cbe01741347ef750d743d6690603f5eed8341e679fb51c8e629337aa11976" +dependencies = [ + "atomic-waker", + "bytes", + "cfg_aliases", + "derive_more", + "ipnet", + "js-sys", + "libc", + "n0-error", + "n0-future", + "n0-watcher", + "netdev", + "netlink-packet-core", + "netlink-packet-route", + "netlink-proto", + "netlink-sys", + "noq-udp", + "objc2-core-foundation", + "objc2-system-configuration", + "pin-project-lite", + "serde", + "socket2", + "time", + "tokio", + "tokio-util", + "tracing", + "web-sys", + "windows 0.62.2", + "windows-result 0.4.1", + "wmi", +] + +[[package]] +name = "new_debug_unreachable" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "650eef8c711430f1a879fdd01d4745a7deea475becfb90269c06775983bbf086" + +[[package]] +name = "nix" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" +dependencies = [ + "bitflags 2.13.0", + "cfg-if", + "cfg_aliases", + "libc", + "memoffset", +] + +[[package]] +name = "nix" +version = "0.31.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf20d2fde8ff38632c426f1165ed7436270b44f199fc55284c38276f9db47c3d" +dependencies = [ + "bitflags 2.13.0", + "cfg-if", + "cfg_aliases", + "libc", +] + +[[package]] +name = "no_std_io2" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "418abd1b6d34fbf6cae440dc874771b0525a604428704c76e48b29a5e67b8003" +dependencies = [ + "memchr", +] + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nom" +version = "8.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df9761775871bdef83bee530e60050f7e54b1105350d6884eb0fb4f46c2f9405" +dependencies = [ + "memchr", +] + +[[package]] +name = "noop_proc_macro" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0676bb32a98c1a483ce53e500a81ad9c3d5b3f7c920c28c24e9cb0980d0b5bc8" + +[[package]] +name = "noq" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4bf95190af1bd4a00a10e8255ca0c8ddd9e9a9f5e79151d7a7eb6d56aff5dc89" +dependencies = [ + "bytes", + "cfg_aliases", + "derive_more", + "noq-proto", + "noq-udp", + "pin-project-lite", + "rustc-hash 2.1.3", + "rustls", + "socket2", + "thiserror 2.0.18", + "tokio", + "tokio-stream", + "tracing", + "web-time", +] + +[[package]] +name = "noq-proto" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa6c890013591e709a3e45dd53501351b7e27e7ff3c7e9fc3dce43e300e7e9d3" +dependencies = [ + "aes-gcm", + "bytes", + "derive_more", + "enum-assoc", + "getrandom 0.4.3", + "identity-hash", + "lru-slab", + "rand 0.10.2", + "rand_pcg", + "ring", + "rustc-hash 2.1.3", + "rustls", + "rustls-pki-types", + "slab", + "sorted-index-buffer", + "thiserror 2.0.18", + "tinyvec", + "tracing", + "web-time", +] + +[[package]] +name = "noq-udp" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3137a52df66c20090a889828d1c655f21f52294cba64e5c4fbb04fc83eee7c8e" +dependencies = [ + "cfg_aliases", + "libc", + "socket2", + "tracing", + "windows-sys 0.61.2", +] + +[[package]] +name = "now" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d89e9874397a1f0a52fc1f197a8effd9735223cb2390e9dcc83ac6cd02923d0" +dependencies = [ + "chrono", +] + +[[package]] +name = "ntapi" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3b335231dfd352ffb0f8017f3b6027a4917f7df785ea2143d8af2adc66980ae" +dependencies = [ + "winapi", +] + +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "num" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23" +dependencies = [ + "num-bigint", + "num-complex", + "num-integer", + "num-iter", + "num-rational", + "num-traits", +] + +[[package]] +name = "num-bigint" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "bytemuck", + "num-traits", +] + +[[package]] +name = "num-conv" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" + +[[package]] +name = "num-derive" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed3955f1a9c7c0c15e092f9c887db08b1fc683305fdf6eb6684f22555355e202" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-iter" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c92800bd69a1eac91786bcfe9da64a897eb72911b8dc3095decbd07429e8048b" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", + "libm", +] + +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + +[[package]] +name = "num_enum" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d0bca838442ec211fa11de3a8b0e0e8f3a4522575b5c4c06ed722e005036f26" +dependencies = [ + "num_enum_derive", + "rustversion", +] + +[[package]] +name = "num_enum_derive" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "680998035259dcfcafe653688bf2aa6d3e2dc05e98be6ab46afb089dc84f1df8" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "num_threads" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c7398b9c8b70908f6371f47ed36737907c87c52af34c268fed0bf0ceb92ead9" +dependencies = [ + "libc", +] + +[[package]] +name = "number_prefix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" + +[[package]] +name = "nvml-wrapper" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f049ae562349fefb8e837eb15443da1e7c6dcbd8a11f52a228f92220c2e5c85e" +dependencies = [ + "bitflags 2.13.0", + "libloading 0.8.9", + "nvml-wrapper-sys", + "static_assertions", + "thiserror 1.0.69", + "wrapcenum-derive", +] + +[[package]] +name = "nvml-wrapper-sys" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b4d594420fcda43b1c2c4bd44d48974aa3c7a9ab2cbf10dc18e35265767bf0b" +dependencies = [ + "libloading 0.8.9", +] + +[[package]] +name = "objc" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "915b1b472bc21c53464d6c8461c9d3af805ba1ef837e1cac254428f4a77177b1" +dependencies = [ + "malloc_buf", +] + +[[package]] +name = "objc2" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a12a8ed07aefc768292f076dc3ac8c48f3781c8f2d5851dd3d98950e8c5a89f" +dependencies = [ + "objc2-encode", +] + +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags 2.13.0", + "block2", + "dispatch2", + "libc", + "objc2", +] + +[[package]] +name = "objc2-core-graphics" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e022c9d066895efa1345f8e33e584b9f958da2fd4cd116792e15e07e4720a807" +dependencies = [ + "dispatch2", + "objc2", + "objc2-core-foundation", + "objc2-io-surface", +] + +[[package]] +name = "objc2-core-wlan" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c71e34919aba0d701380d911702455038a8a3587467fe0141d6a71501e7ffe48" +dependencies = [ + "bitflags 2.13.0", + "objc2", + "objc2-core-foundation", + "objc2-foundation", + "objc2-security", + "objc2-security-foundation", +] + +[[package]] +name = "objc2-encode" +version = "4.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef25abbcd74fb2609453eb695bd2f860d389e457f67dc17cafc8b8cbc89d0c33" + +[[package]] +name = "objc2-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3e0adef53c21f888deb4fa59fc59f7eb17404926ee8a6f59f5df0fd7f9f3272" +dependencies = [ + "bitflags 2.13.0", + "block2", + "libc", + "objc2", + "objc2-core-foundation", +] + +[[package]] +name = "objc2-io-kit" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33fafba39597d6dc1fb709123dfa8289d39406734be322956a69f0931c73bb15" +dependencies = [ + "libc", + "objc2-core-foundation", +] + +[[package]] +name = "objc2-io-surface" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "180788110936d59bab6bd83b6060ffdfffb3b922ba1396b312ae795e1de9d81d" +dependencies = [ + "bitflags 2.13.0", + "objc2", + "objc2-core-foundation", +] + +[[package]] +name = "objc2-metal" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0125f776a10d00af4152d74616409f0d4a2053a6f57fa5b7d6aa2854ac04794" +dependencies = [ + "bitflags 2.13.0", + "block2", + "dispatch2", + "objc2", + "objc2-core-foundation", + "objc2-foundation", +] + +[[package]] +name = "objc2-quartz-core" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96c1358452b371bf9f104e21ec536d37a650eb10f7ee379fff67d2e08d537f1f" +dependencies = [ + "bitflags 2.13.0", + "objc2", + "objc2-core-foundation", + "objc2-core-graphics", + "objc2-foundation", + "objc2-metal", +] + +[[package]] +name = "objc2-security" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe137109bd1e8b5a99390f77a7d8b2961dafc1a1c5db8f2e60329ad6d895a" +dependencies = [ + "bitflags 2.13.0", + "objc2", + "objc2-core-foundation", +] + +[[package]] +name = "objc2-security-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef76382e9cedd18123099f17638715cc3d81dba3637d4c0d39ab69df2ef345a5" +dependencies = [ + "objc2", + "objc2-foundation", +] + +[[package]] +name = "objc2-system-configuration" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7216bd11cbda54ccabcab84d523dc93b858ec75ecfb3a7d89513fa22464da396" +dependencies = [ + "bitflags 2.13.0", + "dispatch2", + "libc", + "objc2", + "objc2-core-foundation", + "objc2-security", +] + +[[package]] +name = "object" +version = "0.37.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff76201f031d8863c38aa7f905eca4f53abbfa15f609db4277d44cd8938f33fe" +dependencies = [ + "memchr", +] + +[[package]] +name = "object_store" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "622acbc9100d3c10e2ee15804b0caa40e55c933d5aa53814cd520805b7958a49" +dependencies = [ + "async-trait", + "base64 0.22.1", + "bytes", + "chrono", + "form_urlencoded", + "futures-channel", + "futures-core", + "futures-util", + "http", + "http-body-util", + "humantime", + "hyper", + "itertools", + "parking_lot", + "percent-encoding", + "quick-xml 0.39.4", + "rand 0.10.2", + "reqwest 0.12.28", + "ring", + "serde", + "serde_json", + "serde_urlencoded", + "thiserror 2.0.18", + "tokio", + "tracing", + "url", + "walkdir", + "wasm-bindgen-futures", + "web-time", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +dependencies = [ + "critical-section", + "portable-atomic", +] + +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + +[[package]] +name = "oneshot" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfe21416a02c693fb9f980befcb230ecc70b0b3d1cc4abf88b9675c4c1457f0c" + +[[package]] +name = "onig" +version = "6.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" +dependencies = [ + "bitflags 2.13.0", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7" +dependencies = [ + "cc", + "pkg-config", +] + +[[package]] +name = "opaque-debug" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" + +[[package]] +name = "openblas-build" +version = "0.10.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb9c85e9e7dd5acdc67b9f3f0c99656b550df716bc63540c6a224a920754a5c2" +dependencies = [ + "anyhow", + "cc", + "flate2", + "tar", + "thiserror 2.0.18", + "ureq 3.3.0", +] + +[[package]] +name = "openblas-src" +version = "0.10.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1a81a5e467f1861ad6ac32d5ec1690ad097d19854753b7424250fe27da46b98" +dependencies = [ + "dirs", + "openblas-build", + "pkg-config", + "vcpkg", +] + +[[package]] +name = "openssl-probe" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" + +[[package]] +name = "option-ext" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" + +[[package]] +name = "ordered-float" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7bb71e1b3fa6ca1c61f383464aaf2bb0e2f8e772a1f01d486832464de363b951" +dependencies = [ + "num-traits", +] + +[[package]] +name = "owo-colors" +version = "4.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d211803b9b6b570f68772237e415a029d5a50c65d382910b879fb19d3271f94d" + +[[package]] +name = "p2p-remote-training" +version = "0.22.0-pre.1" +dependencies = [ + "blake3", + "burn", + "iroh", + "tokio", + "tracing", + "tracing-subscriber", +] + +[[package]] +name = "palette" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cbf71184cc5ecc2e4e1baccdb21026c20e5fc3dcf63028a086131b3ab00b6e6" +dependencies = [ + "approx", + "fast-srgb8", + "libm", + "palette_derive", +] + +[[package]] +name = "palette_derive" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5030daf005bface118c096f510ffb781fc28f9ab6a32ab224d8631be6851d30" +dependencies = [ + "by_address", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "papaya" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "997ee03cd38c01469a7046643714f0ad28880bcb9e6679ff0666e24817ca19b7" +dependencies = [ + "equivalent", + "seize", +] + +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link 0.2.1", +] + +[[package]] +name = "password-hash" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7676374caaee8a325c9e7a2ae557f216c5563a171d6997b0ef8a65af35147700" +dependencies = [ + "base64ct", + "rand_core 0.6.4", + "subtle", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "pastey" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35fb2e5f958ec131621fdd531e9fc186ed768cbe395337403ae56c17a74c68ec" + +[[package]] +name = "pbkdf2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83a0692ec44e4cf1ef28ca317f14f8f07da2d95ec3fa01f86e4467b725e60917" +dependencies = [ + "digest 0.10.7", + "hmac 0.12.1", + "password-hash", + "sha2 0.10.9", +] + +[[package]] +name = "pbkdf2" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ed6a7761f76e3b9f92dfb0a60a6a6477c61024b775147ff0973a02653abaf2" +dependencies = [ + "digest 0.10.7", + "hmac 0.12.1", +] + +[[package]] +name = "pbkdf2" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112d82ceb8c5bf524d9af484d4e4970c9fd5a0cc15ba14ad93dccd28873b0629" +dependencies = [ + "digest 0.11.3", + "hmac 0.13.0", +] + +[[package]] +name = "pem-rfc7468" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6305423e0e7738146434843d1694d621cce767262b2a86910beab705e4493d9" +dependencies = [ + "base64ct", +] + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "pest" +version = "2.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47627dd7305c6a2d6c8c6bcd24c5a4c17dbbf425f4f9c5313e724b38fc9782e9" +dependencies = [ + "memchr", + "ucd-trie", +] + +[[package]] +name = "pest_derive" +version = "2.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b4254325ecad416ab689e27ba51da03ba01a9632bc6e108f5fe7c3c4ad29d58" +dependencies = [ + "pest", + "pest_generator", +] + +[[package]] +name = "pest_generator" +version = "2.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c4c0e91ead7a8f7acecbca6f003fc2e8282b1dbe2dd9c9d2f16aba42995e0a7" +dependencies = [ + "pest", + "pest_meta", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "pest_meta" +version = "2.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9744bc48116fee06334924bb5f2bad41eed5e89bd26e29b0b799f9a3f82c210" +dependencies = [ + "pest", +] + +[[package]] +name = "petgraph" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8701b58ea97060d5e5b155d383a69952a60943f0e6dfe30b04c287beb0b27455" +dependencies = [ + "fixedbitset 0.5.7", + "hashbrown 0.15.5", + "indexmap", + "serde", +] + +[[package]] +name = "pharos" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9567389417feee6ce15dd6527a8a1ecac205ef62c2932bcf3d9f6fc5b78b414" +dependencies = [ + "futures", + "rustc_version", +] + +[[package]] +name = "phf" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078" +dependencies = [ + "phf_macros", + "phf_shared 0.11.3", +] + +[[package]] +name = "phf" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "913273894cec178f401a31ec4b656318d95473527be05c0752cc41cdc32be8b7" +dependencies = [ + "phf_shared 0.12.1", +] + +[[package]] +name = "phf_codegen" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aef8048c789fa5e851558d709946d6d79a8ff88c0440c587967f8e94bfb1216a" +dependencies = [ + "phf_generator", + "phf_shared 0.11.3", +] + +[[package]] +name = "phf_generator" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c80231409c20246a13fddb31776fb942c38553c51e871f8cbd687a4cfb5843d" +dependencies = [ + "phf_shared 0.11.3", + "rand 0.8.6", +] + +[[package]] +name = "phf_macros" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f84ac04429c13a7ff43785d75ad27569f2951ce0ffd30a3321230db2fc727216" +dependencies = [ + "phf_generator", + "phf_shared 0.11.3", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "phf_shared" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67eabc2ef2a60eb7faa00097bd1ffdb5bd28e62bf39990626a582201b7a754e5" +dependencies = [ + "siphasher", +] + +[[package]] +name = "phf_shared" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06005508882fb681fd97892ecff4b7fd0fee13ef1aa569f8695dae7ab9099981" +dependencies = [ + "siphasher", +] + +[[package]] +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "pkcs8" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "451913da69c775a56034ea8d9003d27ee8948e12443eae7c038ba100a4f21cb7" +dependencies = [ + "der", + "spki", +] + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + +[[package]] +name = "planus" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3daf8e3d4b712abe1d690838f6e29fb76b76ea19589c4afa39ec30e12f62af71" +dependencies = [ + "array-init-cursor", + "hashbrown 0.15.5", +] + +[[package]] +name = "plist" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da1d65da6dd5d1e44199ac0f58712d241c0f439f80adea8924d832384087f85" +dependencies = [ + "base64 0.22.1", + "indexmap", + "quick-xml 0.41.0", + "serde", + "time", +] + +[[package]] +name = "png" +version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60769b8b31b2a9f263dae2776c37b1b28ae246943cf719eb6946a1db05128a61" +dependencies = [ + "bitflags 2.13.0", + "crc32fast", + "fdeflate", + "flate2", + "miniz_oxide", +] + +[[package]] +name = "polars" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82f1f122456ec136102033b13f71905b7c3f01e526642679c86aace9f9cdefde" +dependencies = [ + "getrandom 0.2.17", + "getrandom 0.3.4", + "polars-arrow", + "polars-buffer", + "polars-compute", + "polars-core", + "polars-error", + "polars-io", + "polars-lazy", + "polars-ops", + "polars-parquet", + "polars-sql", + "polars-time", + "polars-utils", + "version_check", +] + +[[package]] +name = "polars-arrow" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87d4892d5cc6461bb4a184d18e6fa03a5d316ee1d6de06a33dfa08d479fbc2db" +dependencies = [ + "atoi_simd", + "bitflags 2.13.0", + "bytemuck", + "bytes", + "chrono", + "chrono-tz", + "dyn-clone", + "either", + "ethnum", + "getrandom 0.2.17", + "getrandom 0.3.4", + "half", + "hashbrown 0.16.1", + "itoa", + "lz4", + "num-traits", + "polars-arrow-format", + "polars-buffer", + "polars-error", + "polars-schema", + "polars-utils", + "serde", + "simdutf8", + "streaming-iterator", + "strum_macros 0.27.2", + "version_check", + "zstd 0.13.3", +] + +[[package]] +name = "polars-arrow-format" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a556ac0ee744e61e167f34c1eb0013ce740e0ee6cd8c158b2ec0b518f10e6675" +dependencies = [ + "planus", + "serde", +] + +[[package]] +name = "polars-async" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e87f836190486f500b28347436985cc0af29b7a514e53f98840d396ce4d5f5" +dependencies = [ + "atomic-waker", + "crossbeam-channel", + "crossbeam-deque", + "crossbeam-utils", + "parking_lot", + "pin-project-lite", + "polars-config", + "polars-error", + "polars-utils", + "rand 0.9.4", + "slotmap", + "tokio", +] + +[[package]] +name = "polars-buffer" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e481eeaf33c544ac0dd71a2e375553ca2fdae47b3472a96eaccb6eb43218783d" +dependencies = [ + "bytemuck", + "either", + "polars-utils", + "serde", + "version_check", +] + +[[package]] +name = "polars-compute" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c55d41642a9ee887ac394c5a310af3256fa8340a86cde2cb624c515aa963461c" +dependencies = [ + "atoi_simd", + "bytemuck", + "chrono", + "either", + "fast-float2", + "hashbrown 0.16.1", + "itoa", + "num-traits", + "polars-arrow", + "polars-buffer", + "polars-error", + "polars-utils", + "rand 0.9.4", + "serde", + "strength_reduce", + "strum_macros 0.27.2", + "version_check", + "zmij", +] + +[[package]] +name = "polars-config" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65af861341b00eac73bcb65423fb5cc3d2322526d6b7561a0ddf094947c38033" +dependencies = [ + "polars-error", + "serde", +] + +[[package]] +name = "polars-core" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e5924fc46306054bae78f9d35ea5e404cf185baa7f170eb55a16ff95191069c" +dependencies = [ + "bitflags 2.13.0", + "boxcar", + "bytemuck", + "chrono", + "chrono-tz", + "comfy-table", + "either", + "getrandom 0.3.4", + "hashbrown 0.16.1", + "indexmap", + "itoa", + "num-traits", + "polars-arrow", + "polars-async", + "polars-buffer", + "polars-compute", + "polars-config", + "polars-dtype", + "polars-error", + "polars-row", + "polars-schema", + "polars-utils", + "rand 0.9.4", + "rand_distr 0.5.1", + "rayon", + "regex", + "serde", + "serde_json", + "strum_macros 0.27.2", + "tokio", + "uuid", + "version_check", + "xxhash-rust", +] + +[[package]] +name = "polars-dtype" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b65a750bb99ea66be90c8a7e336f6f3a87427a0f7f89d2a40adae98314e9b27" +dependencies = [ + "boxcar", + "hashbrown 0.16.1", + "polars-arrow", + "polars-error", + "polars-utils", + "serde", + "uuid", +] + +[[package]] +name = "polars-error" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e49a75e3406b9b5b4e5ff177877fe0de766e9688fbdb263a7b25f293dc47d61a" +dependencies = [ + "object_store", + "parking_lot", + "polars-arrow-format", + "regex", + "signal-hook 0.4.4", + "simdutf8", +] + +[[package]] +name = "polars-expr" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e21fdd37e8d9ef109f13d3454baffa0a57041cf60069123b8a2bd846c8ad0205" +dependencies = [ + "bitflags 2.13.0", + "hashbrown 0.16.1", + "num-traits", + "polars-arrow", + "polars-buffer", + "polars-compute", + "polars-core", + "polars-io", + "polars-ops", + "polars-plan", + "polars-row", + "polars-time", + "polars-utils", + "rand 0.9.4", + "rayon", + "recursive", + "regex", + "version_check", +] + +[[package]] +name = "polars-io" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6363a1c44a65fe8d73cce7fe4d77c9b6fea3a0da44007012e755e5b4e65aa078" +dependencies = [ + "async-trait", + "atoi_simd", + "blake3", + "bytes", + "chrono", + "fast-float2", + "fastrand", + "fs4", + "futures", + "glob", + "hashbrown 0.16.1", + "home", + "itoa", + "memchr", + "memmap2", + "num-traits", + "object_store", + "parking_lot", + "percent-encoding", + "polars-arrow", + "polars-buffer", + "polars-compute", + "polars-config", + "polars-core", + "polars-error", + "polars-parquet", + "polars-schema", + "polars-time", + "polars-utils", + "rand 0.9.4", + "rayon", + "regex", + "reqwest 0.12.28", + "serde", + "serde_json", + "simdutf8", + "tokio", + "zmij", +] + +[[package]] +name = "polars-lazy" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "809d9590232a37d638337629c18279af97bdb0d17c3d8b2b6bb186e903e8bd5e" +dependencies = [ + "bitflags 2.13.0", + "chrono", + "either", + "memchr", + "polars-arrow", + "polars-buffer", + "polars-compute", + "polars-config", + "polars-core", + "polars-expr", + "polars-io", + "polars-mem-engine", + "polars-ops", + "polars-plan", + "polars-stream", + "polars-time", + "polars-utils", + "rayon", + "version_check", +] + +[[package]] +name = "polars-mem-engine" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f55c6b7d162c506bc8eee82b065fa0399ebcd20b8f08675a534f3d360904ba38" +dependencies = [ + "memmap2", + "polars-arrow", + "polars-core", + "polars-error", + "polars-expr", + "polars-io", + "polars-ops", + "polars-plan", + "polars-time", + "polars-utils", + "rayon", + "recursive", +] + +[[package]] +name = "polars-ooc" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78b3eea0b386837b760a97ec9c92df99cbc10f94885cae060fd7100f9b794163" +dependencies = [ + "async-trait", + "boxcar", + "libc", + "polars-async", + "polars-config", + "polars-core", + "polars-io", + "polars-utils", + "thread_local", + "tokio", +] + +[[package]] +name = "polars-ops" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb146490a717ac5ae4ff3a22a5adf3ebae79361f187b1f550f9e24783d7ad765" +dependencies = [ + "argminmax", + "base64 0.22.1", + "bytemuck", + "chrono", + "chrono-tz", + "either", + "hashbrown 0.16.1", + "hex", + "indexmap", + "libm", + "memchr", + "num-traits", + "polars-arrow", + "polars-buffer", + "polars-compute", + "polars-core", + "polars-error", + "polars-schema", + "polars-utils", + "rayon", + "regex", + "regex-syntax", + "strum_macros 0.27.2", + "unicode-normalization", + "unicode-reverse", + "version_check", +] + +[[package]] +name = "polars-parquet" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd6b79ba2103c00cbb9c5dd4459ffff1d8ce15286c7a6d376a04c711df20d8b7" +dependencies = [ + "async-stream", + "base64 0.22.1", + "bytemuck", + "ethnum", + "futures", + "hashbrown 0.16.1", + "num-traits", + "polars-arrow", + "polars-buffer", + "polars-compute", + "polars-config", + "polars-error", + "polars-parquet-format", + "polars-utils", + "regex", + "serde", + "simdutf8", + "streaming-decompression", +] + +[[package]] +name = "polars-parquet-format" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c025243dcfe8dbc57e94d9f82eb3bef10b565ab180d5b99bed87fd8aea319ce1" +dependencies = [ + "async-trait", + "futures", +] + +[[package]] +name = "polars-plan" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f5ccc230515adb10762a8c7b0df03fd88f3328deb5b60e9b1eeb2eceef4d344" +dependencies = [ + "bitflags 2.13.0", + "blake3", + "bytemuck", + "bytes", + "chrono", + "chrono-tz", + "either", + "futures", + "hashbrown 0.16.1", + "indexmap", + "memmap2", + "num-traits", + "percent-encoding", + "polars-arrow", + "polars-buffer", + "polars-compute", + "polars-config", + "polars-core", + "polars-error", + "polars-io", + "polars-ops", + "polars-time", + "polars-utils", + "rayon", + "recursive", + "regex", + "sha2 0.10.9", + "slotmap", + "strum_macros 0.27.2", + "tokio", + "version_check", +] + +[[package]] +name = "polars-row" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d4e3254450024078e10c919ecd3b467bdcfdd5cf386c2ca6eedec89bd4771d2" +dependencies = [ + "bitflags 2.13.0", + "bytemuck", + "polars-arrow", + "polars-buffer", + "polars-compute", + "polars-dtype", + "polars-error", + "polars-utils", +] + +[[package]] +name = "polars-schema" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f8a0de8951d02576fd0cdcecd9c605a6b6364d3105b7469b8d7874ea34eea2f" +dependencies = [ + "indexmap", + "polars-error", + "polars-utils", + "serde", + "version_check", +] + +[[package]] +name = "polars-sql" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b282a6164927eb12774b66b071b773a1573173ae53758e8d4df50389ff06efa2" +dependencies = [ + "bitflags 2.13.0", + "hex", + "polars-core", + "polars-error", + "polars-lazy", + "polars-ops", + "polars-plan", + "polars-time", + "polars-utils", + "regex", + "serde", + "sqlparser", +] + +[[package]] +name = "polars-stream" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfa8ff4ee21799898579595a0ef2fb728d0a9cac3d061835fb7f7f6dd854734a" +dependencies = [ + "async-channel", + "async-trait", + "bitflags 2.13.0", + "bytes", + "chrono-tz", + "crossbeam-channel", + "crossbeam-queue", + "futures", + "memchr", + "num-traits", + "parking_lot", + "percent-encoding", + "polars-arrow", + "polars-async", + "polars-buffer", + "polars-compute", + "polars-config", + "polars-core", + "polars-error", + "polars-expr", + "polars-io", + "polars-mem-engine", + "polars-ooc", + "polars-ops", + "polars-parquet", + "polars-plan", + "polars-time", + "polars-utils", + "rayon", + "recursive", + "slotmap", + "tokio", + "uuid", + "version_check", +] + +[[package]] +name = "polars-time" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1063fe074c4212a54917be604377c6e6bfbc8b6c942a5c57be214e4ccaaafdf" +dependencies = [ + "atoi_simd", + "bytemuck", + "chrono", + "chrono-tz", + "now", + "num-traits", + "polars-arrow", + "polars-compute", + "polars-core", + "polars-error", + "polars-ops", + "polars-utils", + "rayon", + "regex", + "strum_macros 0.27.2", +] + +[[package]] +name = "polars-utils" +version = "0.54.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "590b0a94aa8f97992d52f1198600ecc1c1f7cfa03c1b31cae057143455804ac0" +dependencies = [ + "argminmax", + "bincode", + "bytemuck", + "bytes", + "compact_str", + "either", + "flate2", + "foldhash 0.2.0", + "futures", + "half", + "hashbrown 0.16.1", + "indexmap", + "libc", + "memmap2", + "num-derive", + "num-traits", + "polars-config", + "polars-error", + "rand 0.9.4", + "raw-cpuid", + "rayon", + "regex", + "rmp-serde", + "serde", + "serde_json", + "serde_stacker", + "slotmap", + "stacker", + "sysinfo 0.37.2", + "tokio", + "uuid", + "version_check", +] + +[[package]] +name = "polyval" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1fe60d06143b2430aa532c94cfe9e29783047f06c0d7fd359a9a51b729fa25" +dependencies = [ + "cfg-if", + "cpufeatures 0.2.17", + "opaque-debug", + "universal-hash", +] + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +dependencies = [ + "critical-section", + "serde", +] + +[[package]] +name = "portable-atomic-util" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +dependencies = [ + "portable-atomic", +] + +[[package]] +name = "portmapper" +version = "0.19.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb3713e4977408279158444a18c1a01ac9bf2e7eaf1fbfd1a19ac9cd18d90721" +dependencies = [ + "base64 0.22.1", + "bytes", + "derive_more", + "hyper-util", + "igd-next", + "iroh-metrics", + "libc", + "n0-error", + "n0-future", + "netwatch", + "num_enum", + "rand 0.10.2", + "serde", + "smallvec", + "socket2", + "time", + "tokio", + "tokio-util", + "tower-layer", + "tracing", + "url", +] + +[[package]] +name = "postcard" +version = "1.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6764c3b5dd454e283a30e6dfe78e9b31096d9e32036b5d1eaac7a6119ccb9a24" +dependencies = [ + "cobs", + "embedded-io 0.4.0", + "embedded-io 0.6.1", + "postcard-derive", + "serde", +] + +[[package]] +name = "postcard-derive" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0232bd009a197ceec9cc881ba46f727fcd8060a2d8d6a9dde7a69030a6fe2bb" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + +[[package]] +name = "ppmd-rust" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "efca4c95a19a79d1c98f791f10aebd5c1363b473244630bb7dbde1dc98455a24" + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "prefix-trie" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cf6e3177f0684016a5c209b00882e15f8bdd3f3bb48f0491df10cd102d0c6e7" +dependencies = [ + "either", + "ipnet", + "num-traits", +] + +[[package]] +name = "presser" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8cf8e6a8aa66ce33f63993ffc4ea4271eb5b0530a9002db8455ea6050c77bfa" + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.118", +] + +[[package]] +name = "primal-check" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc0d895b311e3af9902528fbb8f928688abbd95872819320517cc24ca6b2bd08" +dependencies = [ + "num-integer", +] + +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "profiling" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d595e54a326bc53c1c197b32d295e14b169e3cfeaa8dc82b529f947fba6bcf5" +dependencies = [ + "profiling-procmacros", +] + +[[package]] +name = "profiling-procmacros" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4488a4a36b9a4ba6b9334a32a39971f77c1436ec82c38707bce707699cc3bbcb" +dependencies = [ + "quote", + "syn 2.0.118", +] + +[[package]] +name = "psm" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "645dbe486e346d9b5de3ef16ede18c26e6c70ad97418f4874b8b1889d6e761ea" +dependencies = [ + "ar_archive_writer", + "cc", +] + +[[package]] +name = "pulp" +version = "0.21.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96b86df24f0a7ddd5e4b95c94fc9ed8a98f1ca94d3b01bdce2824097e7835907" +dependencies = [ + "bytemuck", + "cfg-if", + "libm", + "num-complex", + "reborrow", + "version_check", +] + +[[package]] +name = "pulp" +version = "0.22.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "046aa45b989642ec2e4717c8e72d677b13edd831a4d3b6cf37d9a3e54912496a" +dependencies = [ + "bytemuck", + "cfg-if", + "libm", + "num-complex", + "paste", + "pulp-wasm-simd-flag", + "raw-cpuid", + "reborrow", + "version_check", +] + +[[package]] +name = "pulp-wasm-simd-flag" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d8f70e07b9c3962945a74e59ca1c511bba65b6419468acc217c457d93f3c740" + +[[package]] +name = "pxfm" +version = "0.1.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d55d956fa96f5ec02be2e13af0e20391a5aa83d6a074e3ad368959d0fab299ea" + +[[package]] +name = "pytorch-tests" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "burn-store", + "float-cmp", + "serde", +] + +[[package]] +name = "qoi" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f6d64c71eb498fe9eae14ce4ec935c555749aef511cca85b5568910d6e48001" +dependencies = [ + "bytemuck", +] + +[[package]] +name = "quick-error" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a993555f31e5a609f617c12db6250dedcac1b0a85076912c436e6fc9b2c8e6a3" + +[[package]] +name = "quick-xml" +version = "0.39.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdcc8dd4e2f670d309a5f0e83fe36dfdc05af317008fea29144da1a2ac858e5e" +dependencies = [ + "memchr", + "serde", +] + +[[package]] +name = "quick-xml" +version = "0.41.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e660451e55124f798a69a5af3f49ccfbefbd41910eefd25caf2393e1f3473ec1" +dependencies = [ + "memchr", +] + +[[package]] +name = "quinn" +version = "0.11.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c1a41e437b6bbd489372cd4971de128e85c855f56c57f283d20ff016cf7c0a8" +dependencies = [ + "bytes", + "cfg_aliases", + "pin-project-lite", + "quinn-proto", + "quinn-udp", + "rustc-hash 2.1.3", + "rustls", + "socket2", + "thiserror 2.0.18", + "tokio", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-proto" +version = "0.11.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560" +dependencies = [ + "bytes", + "getrandom 0.4.3", + "lru-slab", + "rand 0.10.2", + "rand_pcg", + "ring", + "rustc-hash 2.1.3", + "rustls", + "rustls-pki-types", + "slab", + "thiserror 2.0.18", + "tinyvec", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-udp" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35a133f956daabe89a61a685c2649f13d82d5aa4bd5d12d1277e1072a21c0694" +dependencies = [ + "cfg_aliases", + "libc", + "once_cell", + "socket2", + "tracing", + "windows-sys 0.61.2", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "r2d2" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93" +dependencies = [ + "log", + "parking_lot", + "scheduled-thread-pool", +] + +[[package]] +name = "r2d2_sqlite" +version = "0.31.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63417e83dc891797eea3ad379f52a5986da4bca0d6ef28baf4d14034dd111b0c" +dependencies = [ + "r2d2", + "rusqlite", + "uuid", +] + +[[package]] +name = "rand" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" +dependencies = [ + "libc", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + +[[package]] +name = "rand" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" +dependencies = [ + "chacha20", + "getrandom 0.4.3", + "rand_core 0.10.1", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + +[[package]] +name = "rand_distr" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" +dependencies = [ + "num-traits", + "rand 0.9.4", +] + +[[package]] +name = "rand_distr" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d431c2703ccf129de4d45253c03f49ebb22b97d6ad79ee3ecfc7e3f4862c1d8" +dependencies = [ + "num-traits", + "rand 0.10.2", +] + +[[package]] +name = "rand_pcg" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a" +dependencies = [ + "rand_core 0.10.1", +] + +[[package]] +name = "range-alloc" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca45419789ae5a7899559e9512e58ca889e41f04f1f2445e9f4b290ceccd1d08" + +[[package]] +name = "ratatui" +version = "0.30.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3274ba0a2c5e1bcad2a2005d20f4dc59dad26b2eb0940fb094500dba4099d57d" +dependencies = [ + "instability", + "ratatui-core", + "ratatui-crossterm", + "ratatui-macros", + "ratatui-termina", + "ratatui-termwiz", + "ratatui-widgets", + "serde", +] + +[[package]] +name = "ratatui-core" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cbb175c433c8e28a809d1f5773a2ae96e68c0ce40db865cbab1020bf33ae479c" +dependencies = [ + "bitflags 2.13.0", + "compact_str", + "critical-section", + "hashbrown 0.17.1", + "itertools", + "kasuari", + "lru", + "palette", + "serde", + "strum 0.28.0", + "thiserror 2.0.18", + "unicode-segmentation", + "unicode-truncate", + "unicode-width", +] + +[[package]] +name = "ratatui-crossterm" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "567584a3b0e6a8203c23de40b4861497266725eb5363dbfd18a1edd603cca9f0" +dependencies = [ + "cfg-if", + "crossterm", + "instability", + "ratatui-core", +] + +[[package]] +name = "ratatui-macros" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed7dc68daa7498a43e4d68e0eb078427e10c38fbcfbb1e42d955f1fa2140d814" +dependencies = [ + "ratatui-core", + "ratatui-widgets", +] + +[[package]] +name = "ratatui-termina" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0bf912d9e66f057a759d92e386a280ea886b352ab757d6ac4d653c7ed2c43c2" +dependencies = [ + "instability", + "ratatui-core", + "termina", +] + +[[package]] +name = "ratatui-termwiz" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "faf03e0380b7744054d6cb74224fe3adf062a029754933f575ca1e3b4c2ce977" +dependencies = [ + "ratatui-core", + "termwiz", +] + +[[package]] +name = "ratatui-widgets" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66e3d19bcc9130ca376277d93b60767ff121ace3be06f5f95f81dd68956407d1" +dependencies = [ + "bitflags 2.13.0", + "hashbrown 0.17.1", + "indoc", + "instability", + "itertools", + "line-clipping", + "ratatui-core", + "serde", + "strum 0.28.0", + "time", + "unicode-segmentation", + "unicode-width", +] + +[[package]] +name = "rav1e" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43b6dd56e85d9483277cde964fd1bdb0428de4fec5ebba7540995639a21cb32b" +dependencies = [ + "aligned-vec", + "arbitrary", + "arg_enum_proc_macro", + "arrayvec", + "av-scenechange", + "av1-grain", + "bitstream-io", + "built", + "cfg-if", + "interpolate_name", + "itertools", + "libc", + "libfuzzer-sys", + "log", + "maybe-rayon", + "new_debug_unreachable", + "noop_proc_macro", + "num-derive", + "num-traits", + "paste", + "profiling", + "rand 0.9.4", + "rand_chacha 0.9.0", + "simd_helpers", + "thiserror 2.0.18", + "v_frame", + "wasm-bindgen", +] + +[[package]] +name = "ravif" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e52310197d971b0f5be7fe6b57530dcd27beb35c1b013f29d66c1ad73fbbcc45" +dependencies = [ + "avif-serialize", + "imgref", + "loop9", + "quick-error", + "rav1e", + "rayon", + "rgb", +] + +[[package]] +name = "raw-cpuid" +version = "11.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "498cd0dc59d73224351ee52a95fee0f1a617a2eae0e7d9d720cc622c73a54186" +dependencies = [ + "bitflags 2.13.0", +] + +[[package]] +name = "raw-window-handle" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20675572f6f24e9e76ef639bc5552774ed45f1c30e2951e1e99c59888861c539" + +[[package]] +name = "raw-window-metal" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40d213455a5f1dc59214213c7330e074ddf8114c9a42411eb890c767357ce135" +dependencies = [ + "objc2", + "objc2-core-foundation", + "objc2-foundation", + "objc2-quartz-core", +] + +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools", + "rayon", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "realfft" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f821338fddb99d089116342c46e9f1fbf3828dba077674613e734e01d6ea8677" +dependencies = [ + "rustfft", +] + +[[package]] +name = "reborrow" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03251193000f4bd3b042892be858ee50e8b3719f2b08e5833ac4353724632430" + +[[package]] +name = "recursive" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0786a43debb760f491b1bc0269fe5e84155353c67482b9e60d0cfb596054b43e" +dependencies = [ + "recursive-proc-macro-impl", + "stacker", +] + +[[package]] +name = "recursive-proc-macro-impl" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76009fbe0614077fc1a2ce255e3a1881a2e3a3527097d5dc6d8212c585e7e38b" +dependencies = [ + "quote", + "syn 2.0.118", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags 2.13.0", +] + +[[package]] +name = "redox_users" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac" +dependencies = [ + "getrandom 0.2.17", + "libredox", + "thiserror 2.0.18", +] + +[[package]] +name = "regex" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a0e75113e14dc5acb068cd0786884f214f1312650a3d36d269f5c4f3cdee8a2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f388202e4b80542a0921078cc23b6333bcf1409c1e3f86404cae4766a6131db" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-lite" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973" + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "relative-path" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" + +[[package]] +name = "remote-inference-web" +version = "0.22.0-pre.1" +dependencies = [ + "blake3", + "burn", + "console_error_panic_hook", + "iroh", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", +] + +[[package]] +name = "renderdoc-sys" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19b30a45b0cd0bcca8037f3d0dc3421eaf95327a17cad11964fb8179b4fc4832" + +[[package]] +name = "reqwest" +version = "0.12.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-channel", + "futures-core", + "futures-util", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-native-certs", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tokio-util", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams 0.4.2", + "web-sys", + "webpki-roots 1.0.8", +] + +[[package]] +name = "reqwest" +version = "0.13.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "219c5811de6525e5416c7d5d53bb656d3afdbc6c5af816e0802bcfa42dbdc1c3" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-core", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "rustls", + "rustls-pki-types", + "rustls-platform-verifier", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tokio-util", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams 0.5.0", + "web-sys", +] + +[[package]] +name = "resolv-conf" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e061d1b48cb8d38042de4ae0a7a6401009d6143dc80d2e2d6f31f0bdd6470c7" + +[[package]] +name = "rgb" +version = "0.8.53" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47b34b781b31e5d73e9fbc8689c70551fd1ade9a19e3e28cfec8580a79290cc4" +dependencies = [ + "bytemuck", +] + +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + +[[package]] +name = "rmp" +version = "0.8.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ba8be72d372b2c9b35542551678538b562e7cf86c3315773cae48dfbfe7790c" +dependencies = [ + "num-traits", +] + +[[package]] +name = "rmp-serde" +version = "1.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f81bee8c8ef9b577d1681a70ebbc962c232461e397b22c208c43c04b67a155" +dependencies = [ + "rmp", + "serde", +] + +[[package]] +name = "rstest" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5a3193c063baaa2a95a33f03035c8a72b83d97a54916055ba22d35ed3839d49" +dependencies = [ + "futures-timer", + "futures-util", + "rstest_macros", +] + +[[package]] +name = "rstest_macros" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c845311f0ff7951c5506121a9ad75aec44d083c31583b2ea5a30bcb0b0abba0" +dependencies = [ + "cfg-if", + "glob", + "proc-macro-crate", + "proc-macro2", + "quote", + "regex", + "relative-path", + "rustc_version", + "syn 2.0.118", + "unicode-ident", +] + +[[package]] +name = "rusqlite" +version = "0.37.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "165ca6e57b20e1351573e3729b958bc62f0e48025386970b6e4d29e7a7e71f3f" +dependencies = [ + "bitflags 2.13.0", + "fallible-iterator", + "fallible-streaming-iterator", + "hashlink", + "libsqlite3-sys", + "smallvec", +] + +[[package]] +name = "rustc-demangle" +version = "0.1.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b74b56ffa8bb2830709a538c2cbcae9aa062db0d2a42563bfb09bdaae44020eb" + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + +[[package]] +name = "rustfft" +version = "6.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "21db5f9893e91f41798c88680037dba611ca6674703c1a18601b01a72c8adb89" +dependencies = [ + "num-complex", + "num-integer", + "num-traits", + "primal-check", + "strength_reduce", + "transpose", +] + +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags 2.13.0", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls" +version = "0.23.41" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b92b125634d9b795e7beca796cc790df15a7fb38323bf3196fda83292d06b1f" +dependencies = [ + "log", + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-native-certs" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" +dependencies = [ + "openssl-probe", + "rustls-pki-types", + "schannel", + "security-framework", +] + +[[package]] +name = "rustls-pki-types" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "764899a24af3980067ee14bc143654f297b22eaebfe3c7b6b211920a5a59b046" +dependencies = [ + "web-time", + "zeroize", +] + +[[package]] +name = "rustls-platform-verifier" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d1e2536ce4f35f4846aa13bff16bd0ff40157cdb14cc056c7b14ba41233ba0" +dependencies = [ + "core-foundation 0.10.1", + "core-foundation-sys", + "jni 0.22.4", + "log", + "once_cell", + "rustls", + "rustls-native-certs", + "rustls-platform-verifier-android", + "rustls-webpki", + "security-framework", + "security-framework-sys", + "webpki-root-certs", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls-platform-verifier-android" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" + +[[package]] +name = "rustls-webpki" +version = "0.103.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "safetensors" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d93279b86b3de76f820a8854dd06cbc33cfa57a417b19c47f6a25280112fb1df" +dependencies = [ + "serde", + "serde_json", +] + +[[package]] +name = "safetensors" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44560c11236a6130a46ce36c836a62936dc81ebf8c36a37947423571be0e55b6" +dependencies = [ + "serde", + "serde_json", +] + +[[package]] +name = "safetensors" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "675656c1eabb620b921efea4f9199f97fc86e36dd6ffd1fbbe48d0f59a4987f5" +dependencies = [ + "hashbrown 0.16.1", + "serde", + "serde_json", +] + +[[package]] +name = "safetensors-tests" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "burn-autodiff", + "burn-flex", + "burn-store", + "float-cmp", + "serde", +] + +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "sanitize-filename" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc984f4f9ceb736a7bb755c3e3bd17dc56370af2600c9780dcc48c66453da34d" +dependencies = [ + "regex", +] + +[[package]] +name = "schannel" +version = "0.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "scheduled-thread-pool" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19" +dependencies = [ + "parking_lot", +] + +[[package]] +name = "scoped-tls" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1cf6437eb19a8f4a6cc0f7dca544973b0b78843adbfeb3683d1a94a0024a294" + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "security-framework" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +dependencies = [ + "bitflags 2.13.0", + "core-foundation 0.10.1", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "seize" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b55fb86dfd3a2f5f76ea78310a88f96c4ea21a3031f8d212443d56123fd0521" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "semver" +version = "1.0.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" +dependencies = [ + "serde", + "serde_core", +] + +[[package]] +name = "send_wrapper" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd0b0ec5f1c1ca621c432a25813d8d60c88abe6d3e08a3eb9cf37d97a0fe3d73" + +[[package]] +name = "seq-macro" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_bytes" +version = "0.11.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5d440709e79d88e51ac01c4b72fc6cb7314017bb7da9eeff678aa94c10e3ea8" +dependencies = [ + "serde", + "serde_core", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_path_to_error" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" +dependencies = [ + "itoa", + "serde", + "serde_core", +] + +[[package]] +name = "serde_rusqlite" +version = "0.40.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec8bd74f47e124e760475a7e863b5820dcef09cae50782d03d65961f5ca1e6d9" +dependencies = [ + "rusqlite", + "serde_core", +] + +[[package]] +name = "serde_spanned" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" +dependencies = [ + "serde_core", +] + +[[package]] +name = "serde_stacker" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4936375d50c4be7eff22293a9344f8e46f323ed2b3c243e52f89138d9bb0f4a" +dependencies = [ + "serde", + "serde_core", + "stacker", +] + +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "serdect" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66cf8fedced2fcf12406bcb34223dffb92eaf34908ede12fed414c82b7f00b3e" +dependencies = [ + "base16ct", + "serde", +] + +[[package]] +name = "serial_test" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "699f4197115b8a7e7ff19c9a315a4bd6fffec26cc4626ef45ecaea389e081c6d" +dependencies = [ + "futures-executor", + "futures-util", + "log", + "once_cell", + "parking_lot", + "serial_test_derive", +] + +[[package]] +name = "serial_test_derive" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94e153fc76e1c6a068703d6d29c508a0b15c061c4b7e43da59cc097bc342673c" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "server" +version = "0.22.0-pre.1" +dependencies = [ + "burn", +] + +[[package]] +name = "sha1" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +dependencies = [ + "cfg-if", + "cpufeatures 0.2.17", + "digest 0.10.7", +] + +[[package]] +name = "sha1" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aacc4cc499359472b4abe1bf11d0b12e688af9a805fa5e3016f9a386dc2d0214" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", +] + +[[package]] +name = "sha1_smol" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d" + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures 0.2.17", + "digest 0.10.7", +] + +[[package]] +name = "sha2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", +] + +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "signal-hook" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d881a16cf4426aa584979d30bd82cb33429027e42122b169753d6ef1085ed6e2" +dependencies = [ + "libc", + "signal-hook-registry", +] + +[[package]] +name = "signal-hook" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a0c28ca5908dbdbcd52e6fdaa00358ab88637f8ab33e1f188dd510eb44b53d" +dependencies = [ + "libc", + "signal-hook-registry", +] + +[[package]] +name = "signal-hook-mio" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b75a19a7a740b25bc7944bdee6172368f988763b744e3d4dfe753f6b4ece40cc" +dependencies = [ + "libc", + "mio", + "signal-hook 0.3.18", +] + +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + +[[package]] +name = "signature" +version = "3.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28d567dcbaf0049cb8ac2608a76cd95ff9e4412e1899d389ee400918ca7537f5" + +[[package]] +name = "simd-adler32" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" + +[[package]] +name = "simd_cesu8" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94f90157bb87cddf702797c5dadfa0be7d266cdf49e22da2fcaa32eff75b2c33" +dependencies = [ + "rustc_version", + "simdutf8", +] + +[[package]] +name = "simd_helpers" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95890f873bec569a0362c235787f3aca6e1e887302ba4840839bcc6459c42da6" +dependencies = [ + "quote", +] + +[[package]] +name = "simdutf8" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" + +[[package]] +name = "simple-dns" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a75cbde1bf934313596a004973e462f9a82caa814dcf1a5f507bdf51597eeb4" +dependencies = [ + "bitflags 2.13.0", +] + +[[package]] +name = "simple-regression" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "log", + "rgb", + "serde", + "textplots", +] + +[[package]] +name = "siphasher" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "slotmap" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bdd58c3c93c3d278ca835519292445cb4b0d4dc59ccfdf7ceadaab3f8aeb4038" +dependencies = [ + "version_check", +] + +[[package]] +name = "slug" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "882a80f72ee45de3cc9a5afeb2da0331d58df69e4e7d8eeb5d3c7784ae67e724" +dependencies = [ + "deunicode", + "wasm-bindgen", +] + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" +dependencies = [ + "serde", +] + +[[package]] +name = "socket2" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "socks" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b" +dependencies = [ + "byteorder", + "libc", + "winapi", +] + +[[package]] +name = "sorted-index-buffer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea06cc588e43c632923a55450401b8f25e628131571d4e1baea1bdfdb2b5ed06" + +[[package]] +name = "spez" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c87e960f4dca2788eeb86bbdde8dd246be8948790b7618d656e68f9b720a86e8" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "spin" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d5fe4ccb98d9c292d56fec89a5e07da7fc4cf0dc11e156b41793132775d3e591" +dependencies = [ + "lock_api", + "portable-atomic", +] + +[[package]] +name = "spirv" +version = "0.4.0+sdk-1.4.341.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9571ea910ebd84c86af4b3ed27f9dbdc6ad06f17c5f96146b2b671e2976744f" +dependencies = [ + "bitflags 2.13.0", +] + +[[package]] +name = "spki" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d9efca8738c78ee9484207732f728b1ef517bbb1833d6fc0879ca898a522f6f" +dependencies = [ + "base64ct", + "der", +] + +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64 0.13.1", + "nom 7.1.3", + "serde", + "unicode-segmentation", +] + +[[package]] +name = "sqlparser" +version = "0.60.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "505aa16b045c4c1375bf5f125cce3813d0176325bfe9ffc4a903f423de7774ff" +dependencies = [ + "log", + "recursive", + "sqlparser_derive", +] + +[[package]] +name = "sqlparser_derive" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "028e551d5e270b31b9f3ea271778d9d827148d4287a5d96167b6bb9787f5cc38" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "stable-vec" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dac7bc0f7d0d44329b200020effbc25a534d89fa142af95e3ddf76113412a5e" + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "stacker" +version = "0.1.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "640c8cdd92b6b12f5bcb1803ca3bbf5ab96e5e6b6b96b9ab77dabe9e880b3190" +dependencies = [ + "cc", + "cfg-if", + "libc", + "psm", + "windows-sys 0.61.2", +] + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "streaming-decompression" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf6cc3b19bfb128a8ad11026086e31d3ce9ad23f8ea37354b31383a187c44cf3" +dependencies = [ + "fallible-streaming-iterator", +] + +[[package]] +name = "streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b2231b7c3057d5e4ad0156fb3dc807d900806020c5ffa3ee6ff2c8c76fb8520" + +[[package]] +name = "strength_reduce" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe895eb47f22e2ddd4dabc02bce419d2e643c8e3b585c78158b349195bc24d82" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "strum" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af23d6f6c1a224baef9d3f61e287d2761385a5b88fdab4eb4c6f11aeb54c4bcf" +dependencies = [ + "strum_macros 0.27.2", +] + +[[package]] +name = "strum" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9628de9b8791db39ceda2b119bbe13134770b56c138ec1d3af810d045c04f9bd" +dependencies = [ + "strum_macros 0.28.0", +] + +[[package]] +name = "strum_macros" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7695ce3845ea4b33927c055a39dc438a45b059f7c1b3d91d38d10355fb8cbca7" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "strum_macros" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab85eea0270ee17587ed4156089e10b9e6880ee688791d45a905f5b1ca36f664" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + +[[package]] +name = "symlink" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7973cce6668464ea31f176d85b13c7ab3bba2cb3b77a2ed26abd7801688010a" + +[[package]] +name = "syn" +version = "1.0.109" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "2.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "sync_wrapper" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +dependencies = [ + "futures-core", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "sysctl" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01198a2debb237c62b6826ec7081082d951f46dbb64b0e8c7649a452230d1dfc" +dependencies = [ + "bitflags 2.13.0", + "byteorder", + "enum-as-inner", + "libc", + "thiserror 1.0.69", + "walkdir", +] + +[[package]] +name = "sysinfo" +version = "0.37.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16607d5caffd1c07ce073528f9ed972d88db15dd44023fa57142963be3feb11f" +dependencies = [ + "libc", + "memchr", + "ntapi", + "objc2-core-foundation", + "objc2-io-kit", + "windows 0.61.3", +] + +[[package]] +name = "sysinfo" +version = "0.38.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92ab6a2f8bfe508deb3c6406578252e491d299cbbf3bc0529ecc3313aee4a52f" +dependencies = [ + "libc", + "memchr", + "ntapi", + "objc2-core-foundation", + "objc2-io-kit", + "windows 0.62.2", +] + +[[package]] +name = "system-configuration" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a13f3d0daba03132c0aa9767f98351b3488edc2c100cda2d2ec2b04f3d8d3c8b" +dependencies = [ + "bitflags 2.13.0", + "core-foundation 0.9.4", + "system-configuration-sys", +] + +[[package]] +name = "system-configuration-sys" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e1d1b10ced5ca923a1fcb8d03e96b8d3268065d724548c0211415ff6ac6bac4" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "systemstat" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a583abe520746270ffdbdaf0e3039a806f29be9d7034d66466a4839a01de0610" +dependencies = [ + "bytesize", + "kstat-rs", + "lazy_static", + "libc", + "mach2", + "nom 7.1.3", + "time", + "winapi", +] + +[[package]] +name = "tagptr" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" + +[[package]] +name = "tar" +version = "0.4.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f6221d9a6003c78398e3b239969f352578258df48c8eb051caadae0015bc840" +dependencies = [ + "filetime", + "libc", + "xattr", +] + +[[package]] +name = "tch" +version = "0.22.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e09b91610202dc4820c21eb474a42b386ef69f323b1c0902b5472ba7456ebb5" +dependencies = [ + "half", + "lazy_static", + "libc", + "ndarray 0.16.1", + "rand 0.8.6", + "safetensors 0.3.3", + "thiserror 1.0.69", + "torch-sys", + "zip 0.6.6", +] + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + +[[package]] +name = "termcolor" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06794f8f6c5c898b3275aebefa6b8a1cb24cd2c6c79397ab15774837a0bc5755" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "termina" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9048a889effe34a5cddee0af7f53285198b16dca3be510858d38dfdb3e62a04e" +dependencies = [ + "bitflags 2.13.0", + "parking_lot", + "rustix", + "signal-hook 0.3.18", + "windows-sys 0.61.2", +] + +[[package]] +name = "terminal_size" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "230a1b821ccbd75b185820a1f1ff7b14d21da1e442e22c0863ea5f08771a8874" +dependencies = [ + "rustix", + "windows-sys 0.61.2", +] + +[[package]] +name = "terminfo" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4ea810f0692f9f51b382fff5893887bb4580f5fa246fde546e0b13e7fcee662" +dependencies = [ + "fnv", + "nom 7.1.3", + "phf 0.11.3", + "phf_codegen", +] + +[[package]] +name = "termios" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "411c5bf740737c7918b8b1fe232dca4dc9f8e754b8ad5e20966814001ed0ac6b" +dependencies = [ + "libc", +] + +[[package]] +name = "termwiz" +version = "0.23.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4676b37242ccbd1aabf56edb093a4827dc49086c0ffd764a5705899e0f35f8f7" +dependencies = [ + "anyhow", + "base64 0.22.1", + "bitflags 2.13.0", + "fancy-regex", + "filedescriptor", + "finl_unicode", + "fixedbitset 0.4.2", + "hex", + "lazy_static", + "libc", + "log", + "memmem", + "nix 0.29.0", + "num-derive", + "num-traits", + "ordered-float", + "pest", + "pest_derive", + "phf 0.11.3", + "sha2 0.10.9", + "signal-hook 0.3.18", + "siphasher", + "terminfo", + "termios", + "thiserror 1.0.69", + "ucd-trie", + "unicode-segmentation", + "vtparse", + "wezterm-bidi", + "wezterm-blob-leases", + "wezterm-color-types", + "wezterm-dynamic", + "wezterm-input-types", + "winapi", +] + +[[package]] +name = "text-classification" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "derive-new", + "log", + "serde", + "tokenizers", +] + +[[package]] +name = "text-generation" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "derive-new", + "log", + "serde", + "tokenizers", +] + +[[package]] +name = "text_placeholder" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd5008f74a09742486ef0047596cf35df2b914e2a8dca5727fcb6ba6842a766b" +dependencies = [ + "hashbrown 0.13.2", + "serde", + "serde_json", +] + +[[package]] +name = "textdistance" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa672c55ab69f787dbc9126cc387dbe57fdd595f585e4524cf89018fa44ab819" + +[[package]] +name = "textplots" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f7657a0066c9f9663659db0665319adff8b0943305fc73eddf1010e5a2072b1" +dependencies = [ + "drawille", + "rgb", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl 2.0.18", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "thread-tree" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffbd370cb847953a25954d9f63e14824a36113f8c72eecf6eccef5dc4b45d630" +dependencies = [ + "crossbeam-channel", +] + +[[package]] +name = "thread_local" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "tiff" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b63feaf3343d35b6ca4d50483f94843803b0f51634937cc2ec519fc32232bc52" +dependencies = [ + "fax", + "flate2", + "half", + "quick-error", + "weezl", + "zune-jpeg", +] + +[[package]] +name = "time" +version = "0.3.53" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18dfaaeddcb932337b5e7866ee7d0ce9b76d2fd092997146f187ec09b4558a50" +dependencies = [ + "deranged", + "js-sys", + "libc", + "num-conv", + "num_threads", + "powerfmt", + "serde_core", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109" + +[[package]] +name = "time-macros" +version = "0.2.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c431b87111666e491a90baa837f914fb45cd5dc3c268591b0220ff5057f2085f" +dependencies = [ + "num-conv", + "time-core", +] + +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] + +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tinyvec" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + +[[package]] +name = "tokenizers" +version = "0.22.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b238e22d44a15349529690fb07bd645cf58149a1b1e44d6cb5bd1641ff1a6223" +dependencies = [ + "ahash", + "aho-corasick", + "compact_str", + "dary_heap", + "derive_builder", + "esaxx-rs", + "getrandom 0.3.4", + "hf-hub", + "itertools", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand 0.9.4", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror 2.0.18", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + +[[package]] +name = "tokio" +version = "1.52.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +dependencies = [ + "bytes", + "libc", + "mio", + "pin-project-lite", + "signal-hook-registry", + "socket2", + "tokio-macros", + "tracing", + "windows-sys 0.61.2", +] + +[[package]] +name = "tokio-macros" +version = "2.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "tokio-rustls" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" +dependencies = [ + "rustls", + "tokio", +] + +[[package]] +name = "tokio-stream" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", + "tokio-util", +] + +[[package]] +name = "tokio-tungstenite" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f72a05e828585856dacd553fba484c242c46e391fb0e58917c942ee9202915c" +dependencies = [ + "futures-util", + "log", + "tokio", + "tungstenite", +] + +[[package]] +name = "tokio-util" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "futures-util", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "tokio-websockets" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d52efb639344a7c6adb8e62c6f3d2c19c001ff1b79a5041ba1c6ed42e19c6aa5" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-core", + "futures-sink", + "getrandom 0.4.3", + "http", + "httparse", + "rand 0.10.2", + "ring", + "rustls-pki-types", + "sha1_smol", + "simdutf8", + "tokio", + "tokio-rustls", + "tokio-util", +] + +[[package]] +name = "toml" +version = "0.9.12+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf92845e79fc2e2def6a5d828f0801e29a2f8acc037becc5ab08595c7d5e9863" +dependencies = [ + "serde_core", + "serde_spanned", + "toml_datetime 0.7.5+spec-1.1.0", + "toml_parser", + "toml_writer", + "winnow 0.7.15", +] + +[[package]] +name = "toml" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81f3d15e84cbcd896376e6730314d59fb5a87f31e4b038454184435cd57defee" +dependencies = [ + "indexmap", + "serde_core", + "serde_spanned", + "toml_datetime 1.1.1+spec-1.1.0", + "toml_parser", + "toml_writer", + "winnow 1.0.3", +] + +[[package]] +name = "toml_datetime" +version = "0.7.5+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92e1cfed4a3038bc5a127e35a2d360f145e1f4b971b551a2ba5fd7aedf7e1347" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.12+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2153edc6955a6c354fad8f5efd38b6a8769bdccf9fe50f8e1329f81b0baa5d7" +dependencies = [ + "indexmap", + "toml_datetime 1.1.1+spec-1.1.0", + "toml_parser", + "winnow 1.0.3", +] + +[[package]] +name = "toml_parser" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" +dependencies = [ + "winnow 1.0.3", +] + +[[package]] +name = "toml_writer" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "756daf9b1013ebe47a8776667b466417e2d4c5679d441c26230efd9ef78692db" + +[[package]] +name = "torch-sys" +version = "0.22.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aef40c585e342df95b66a1fa7c923188623999c2b657227befb481dfb03a6a42" +dependencies = [ + "anyhow", + "cc", + "libc", + "serde", + "serde_json", + "ureq 2.12.1", + "zip 0.6.6", +] + +[[package]] +name = "tower" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" +dependencies = [ + "futures-core", + "futures-util", + "pin-project-lite", + "sync_wrapper", + "tokio", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower-http" +version = "0.6.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" +dependencies = [ + "bitflags 2.13.0", + "bytes", + "futures-util", + "http", + "http-body", + "pin-project-lite", + "tower", + "tower-layer", + "tower-service", + "url", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + +[[package]] +name = "tracel-ash" +version = "0.39.3+sdk1.4.350" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d97fcf7e656528db871e679727140819a221d8c2aa5b90bfb5eb492a0ecc16e6" +dependencies = [ + "ash", + "c2rust-bitfields", +] + +[[package]] +name = "tracel-llvm" +version = "22.1.4-4" +source = "git+https://github.com/tracel-ai/tracel-llvm?rev=82401e3c00a35a3fb7d5e2b415e79692e09e1079#82401e3c00a35a3fb7d5e2b415e79692e09e1079" +dependencies = [ + "tracel-mlir-rs", + "tracel-mlir-sys", +] + +[[package]] +name = "tracel-llvm-bundler" +version = "22.1.4-4" +source = "git+https://github.com/tracel-ai/tracel-llvm?rev=82401e3c00a35a3fb7d5e2b415e79692e09e1079#82401e3c00a35a3fb7d5e2b415e79692e09e1079" +dependencies = [ + "anyhow", + "bytes", + "constcat", + "dirs", + "liblzma", + "regex", + "reqwest 0.12.28", + "serde", + "serde_json", + "sha2 0.10.9", + "tar", + "walkdir", +] + +[[package]] +name = "tracel-mlir-rs" +version = "22.1.4-4" +source = "git+https://github.com/tracel-ai/tracel-llvm?rev=82401e3c00a35a3fb7d5e2b415e79692e09e1079#82401e3c00a35a3fb7d5e2b415e79692e09e1079" +dependencies = [ + "tracel-mlir-rs-macros", + "tracel-mlir-sys", +] + +[[package]] +name = "tracel-mlir-rs-macros" +version = "22.1.4-4" +source = "git+https://github.com/tracel-ai/tracel-llvm?rev=82401e3c00a35a3fb7d5e2b415e79692e09e1079#82401e3c00a35a3fb7d5e2b415e79692e09e1079" +dependencies = [ + "comrak", + "convert_case 0.8.0", + "proc-macro2", + "quote", + "regex", + "syn 2.0.118", + "tracel-llvm-bundler", + "tracel-tblgen-rs", + "unindent", +] + +[[package]] +name = "tracel-mlir-sys" +version = "22.1.4-4" +source = "git+https://github.com/tracel-ai/tracel-llvm?rev=82401e3c00a35a3fb7d5e2b415e79692e09e1079#82401e3c00a35a3fb7d5e2b415e79692e09e1079" +dependencies = [ + "tracel-llvm-bundler", +] + +[[package]] +name = "tracel-rspirv" +version = "0.13.1+sdk-1.4.350.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33f8d9173a81da18a8014822ac00525c6ecf8086c93f65b4113c24127b12cc09" +dependencies = [ + "bitflags 2.13.0", + "rustc-hash 2.1.3", + "serde", +] + +[[package]] +name = "tracel-tblgen-rs" +version = "22.1.4-4" +source = "git+https://github.com/tracel-ai/tracel-llvm?rev=82401e3c00a35a3fb7d5e2b415e79692e09e1079#82401e3c00a35a3fb7d5e2b415e79692e09e1079" +dependencies = [ + "paste", + "thiserror 2.0.18", + "tracel-llvm-bundler", +] + +[[package]] +name = "tracel-xtask" +version = "4.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "58031a2e8da33f2807d97effead6bf38b73abccec082eb09bf0adab357dd6032" +dependencies = [ + "anyhow", + "clap", + "ctrlc", + "derive_more", + "dotenvy", + "env_logger", + "home", + "humantime", + "inquire", + "log", + "owo-colors", + "rand 0.9.4", + "regex", + "serde", + "serde_json", + "strum 0.27.2", + "time", + "toml 0.9.12+spec-1.1.0", + "tracel-xtask-macros", + "ureq 3.3.0", + "zip 6.0.0", +] + +[[package]] +name = "tracel-xtask-macros" +version = "4.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6e019e77107aac569cead075b5710996898ec8b4633a4a4be52a59eac5b0960" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "log", + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-appender" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "050686193eb999b4bb3bc2acfa891a13da00f79734704c4b8b4ef1a10b368a3c" +dependencies = [ + "crossbeam-channel", + "symlink", + "thiserror 2.0.18", + "time", + "tracing-subscriber", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex-automata", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", +] + +[[package]] +name = "transpose" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ad61aed86bc3faea4300c7aee358b4c6d0c8d6ccc36524c96e4c92ccf26e77e" +dependencies = [ + "num-integer", + "strength_reduce", +] + +[[package]] +name = "try-lock" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" + +[[package]] +name = "tungstenite" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c01152af293afb9c7c2a57e4b559c5620b421f6d133261c60dd2d0cdb38e6b8" +dependencies = [ + "bytes", + "data-encoding", + "http", + "httparse", + "log", + "rand 0.9.4", + "sha1 0.10.6", + "thiserror 2.0.18", +] + +[[package]] +name = "tynm" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a21cdb0fc8f85c98b1ec812bc4cd69faf6c0fa2fc17d44ea3c2cdd38dc08e999" +dependencies = [ + "nom 8.0.0", +] + +[[package]] +name = "type-map" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb30dbbd9036155e74adad6812e9898d03ec374946234fbcebd5dfc7b9187b90" +dependencies = [ + "rustc-hash 2.1.3", +] + +[[package]] +name = "typed-arena" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6af6ae20167a9ece4bcb41af5b80f8a1f1df981f6391189ce00fd257af04126a" + +[[package]] +name = "typed-path" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e28f89b80c87b8fb0cf04ab448d5dd0dd0ade2f8891bae878de66a75a28600e" + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "ucd-trie" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971" + +[[package]] +name = "ug" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76b761acf8af3494640d826a8609e2265e19778fb43306c7f15379c78c9b05b0" +dependencies = [ + "gemm 0.18.2", + "half", + "libloading 0.8.9", + "memmap2", + "num", + "num-traits", + "num_cpus", + "rayon", + "safetensors 0.4.5", + "serde", + "thiserror 1.0.69", + "tracing", + "yoke 0.7.5", +] + +[[package]] +name = "ug-cuda" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f0a1fa748f26166778c33b8498255ebb7c6bffb472bcc0a72839e07ebb1d9b5" +dependencies = [ + "cudarc 0.17.8", + "half", + "serde", + "thiserror 1.0.69", + "ug", +] + +[[package]] +name = "ug-metal" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7adf545a99a086d362efc739e7cf4317c18cbeda22706000fd434d70ea3d95" +dependencies = [ + "half", + "metal", + "objc", + "serde", + "thiserror 1.0.69", + "ug", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-normalization" +version = "0.1.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fd4f6878c9cb28d874b009da9e8d183b5abc80117c40bbd187a1fde336be6e8" +dependencies = [ + "tinyvec", +] + +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + +[[package]] +name = "unicode-reverse" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b6f4888ebc23094adfb574fdca9fdc891826287a6397d2cd28802ffd6f20c76" +dependencies = [ + "unicode-segmentation", +] + +[[package]] +name = "unicode-segmentation" +version = "1.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" + +[[package]] +name = "unicode-truncate" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16b380a1238663e5f8a691f9039c73e1cdae598a30e9855f541d29b08b53e9a5" +dependencies = [ + "itertools", + "unicode-segmentation", + "unicode-width", +] + +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + +[[package]] +name = "unicode-xid" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" + +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + +[[package]] +name = "unindent" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + +[[package]] +name = "unit-prefix" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81e544489bf3d8ef66c953931f56617f423cd4b5494be343d9b9d3dda037b9a3" + +[[package]] +name = "universal-hash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc1de2c688dc15305988b563c3854064043356019f97a4b46276fe734c4f07ea" +dependencies = [ + "crypto-common 0.1.7", + "subtle", +] + +[[package]] +name = "unsynn" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "501a7adf1a4bd9951501e5c66621e972ef8874d787628b7f90e64f936ef7ec0a" +dependencies = [ + "mutants", + "proc-macro2", + "rustc-hash 2.1.3", +] + +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + +[[package]] +name = "unty" +version = "0.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d49784317cd0d1ee7ec5c716dd598ec5b4483ea832a2dced265471cc0f690ae" + +[[package]] +name = "ureq" +version = "2.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02d1a66277ed75f640d608235660df48c8e3c19f3b4edb6a263315626cc3c01d" +dependencies = [ + "base64 0.22.1", + "flate2", + "log", + "once_cell", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "socks", + "url", + "webpki-roots 0.26.11", +] + +[[package]] +name = "ureq" +version = "3.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dea7109cdcd5864d4eeb1b58a1648dc9bf520360d7af16ec26d0a9354bafcfc0" +dependencies = [ + "base64 0.22.1", + "cookie_store", + "flate2", + "log", + "percent-encoding", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "ureq-proto", + "utf8-zero", + "webpki-roots 1.0.8", +] + +[[package]] +name = "ureq-proto" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e994ba84b0bd1b1b0cf92878b7ef898a5c1760108fe7b6010327e274917a808c" +dependencies = [ + "base64 0.22.1", + "http", + "httparse", + "log", +] + +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", + "serde_derive", +] + +[[package]] +name = "utf8-zero" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8c0a043c9540bae7c578c88f91dda8bd82e59ae27c21baca69c8b191aaf5a6e" + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + +[[package]] +name = "uuid" +version = "1.23.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf80a72845275afea99e7f2b434723d3bc7e38470fcd1c7ed39a599c73319a53" +dependencies = [ + "atomic", + "getrandom 0.4.3", + "js-sys", + "rand 0.10.2", + "serde_core", + "wasm-bindgen", +] + +[[package]] +name = "v_frame" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "666b7727c8875d6ab5db9533418d7c764233ac9c0cff1d469aec8fa127597be2" +dependencies = [ + "aligned-vec", + "num-traits", + "wasm-bindgen", +] + +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + +[[package]] +name = "variadics_please" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b53d0f7f2d0182759a549f1f313ce16b07d8f9ba93098157b5fa10ee3e61117f" +dependencies = [ + "quote", + "unsynn", +] + +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + +[[package]] +name = "vergen" +version = "9.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b849a1f6d8639e8de261e81ee0fc881e3e3620db1af9f2e0da015d4382ceaf75" +dependencies = [ + "anyhow", + "derive_builder", + "rustversion", + "vergen-lib", +] + +[[package]] +name = "vergen-gitcl" +version = "9.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ff3b5300a085d6bcd8fc96a507f706a28ae3814693236c9b409db71a1d15b9" +dependencies = [ + "anyhow", + "derive_builder", + "rustversion", + "time", + "vergen", + "vergen-lib", +] + +[[package]] +name = "vergen-lib" +version = "9.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b34a29ba7e9c59e62f229ae1932fb1b8fb8a6fdcc99215a641913f5f5a59a569" +dependencies = [ + "anyhow", + "derive_builder", + "rustversion", +] + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "virtue" +version = "0.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "051eb1abcf10076295e815102942cc58f9d5e3b4560e46e53c21e8ff6f3af7b1" + +[[package]] +name = "void" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a02e4885ed3bc0f2de90ea6dd45ebcbb66dacffe03547fadbb0eeae2770887d" + +[[package]] +name = "vtparse" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d9b2acfb050df409c972a37d3b8e08cdea3bddb0c09db9d53137e504cfabed0" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + +[[package]] +name = "want" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +dependencies = [ + "try-lock", +] + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn 2.0.118", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "wasm-streams" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + +[[package]] +name = "wasm-streams" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + +[[package]] +name = "wayland-sys" +version = "0.31.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8eab23fefc9e41f8e841df4a9c707e8a8c4ed26e944ef69297184de2785e3be" +dependencies = [ + "dlib", + "log", + "once_cell", + "pkg-config", +] + +[[package]] +name = "web-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "webpki-root-certs" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d46a5a140e6f7afeccd8eae97eff335163939eac8b929834875168b29b3d267" +dependencies = [ + "rustls-pki-types", +] + +[[package]] +name = "webpki-roots" +version = "0.26.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" +dependencies = [ + "webpki-roots 1.0.8", +] + +[[package]] +name = "webpki-roots" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf85cb06032201fa7c6f829d7db5a7e5aa45bcc0655327713065f6f0576731bf" +dependencies = [ + "rustls-pki-types", +] + +[[package]] +name = "weezl" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a28ac98ddc8b9274cb41bb4d9d4d5c425b6020c50c46f25559911905610b4a88" + +[[package]] +name = "wezterm-bidi" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c0a6e355560527dd2d1cf7890652f4f09bb3433b6aadade4c9b5ed76de5f3ec" +dependencies = [ + "log", + "wezterm-dynamic", +] + +[[package]] +name = "wezterm-blob-leases" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "692daff6d93d94e29e4114544ef6d5c942a7ed998b37abdc19b17136ea428eb7" +dependencies = [ + "getrandom 0.3.4", + "mac_address", + "sha2 0.10.9", + "thiserror 1.0.69", + "uuid", +] + +[[package]] +name = "wezterm-color-types" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7de81ef35c9010270d63772bebef2f2d6d1f2d20a983d27505ac850b8c4b4296" +dependencies = [ + "csscolorparser", + "deltae", + "lazy_static", + "wezterm-dynamic", +] + +[[package]] +name = "wezterm-dynamic" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f2ab60e120fd6eaa68d9567f3226e876684639d22a4219b313ff69ec0ccd5ac" +dependencies = [ + "log", + "ordered-float", + "strsim", + "thiserror 1.0.69", + "wezterm-dynamic-derive", +] + +[[package]] +name = "wezterm-dynamic-derive" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46c0cf2d539c645b448eaffec9ec494b8b19bd5077d9e58cb1ae7efece8d575b" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + +[[package]] +name = "wezterm-input-types" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7012add459f951456ec9d6c7e6fc340b1ce15d6fc9629f8c42853412c029e57e" +dependencies = [ + "bitflags 1.3.2", + "euclid", + "lazy_static", + "serde", + "wezterm-dynamic", +] + +[[package]] +name = "wgan" +version = "0.22.0-pre.1" +dependencies = [ + "burn", + "image", +] + +[[package]] +name = "wgpu" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d8f4bd44d92da5270f03409dba9f952dab24f128e05d6a554926101d1bf9114" +dependencies = [ + "arrayvec", + "bitflags 2.13.0", + "bytemuck", + "cfg-if", + "cfg_aliases", + "document-features", + "hashbrown 0.17.1", + "js-sys", + "log", + "naga", + "parking_lot", + "portable-atomic", + "profiling", + "raw-window-handle", + "smallvec", + "static_assertions", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", + "wgpu-core", + "wgpu-hal", + "wgpu-types", +] + +[[package]] +name = "wgpu-core" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08763620e76fc980bca7bf84de82568614487a53172dd968d89187282eb87fa2" +dependencies = [ + "arrayvec", + "bit-set 0.10.0", + "bit-vec 0.9.1", + "bitflags 2.13.0", + "bytemuck", + "cfg_aliases", + "document-features", + "hashbrown 0.17.1", + "indexmap", + "log", + "naga", + "naga-types", + "once_cell", + "parking_lot", + "portable-atomic", + "profiling", + "raw-window-handle", + "rustc-hash 1.1.0", + "smallvec", + "thiserror 2.0.18", + "wgpu-core-deps-apple", + "wgpu-core-deps-emscripten", + "wgpu-core-deps-windows-linux-android", + "wgpu-hal", + "wgpu-naga-bridge", + "wgpu-types", +] + +[[package]] +name = "wgpu-core-deps-apple" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3283b6da70d7fc892de95840215aa36381b5d0cbc479e88ee2b173c7b992b225" +dependencies = [ + "wgpu-hal", +] + +[[package]] +name = "wgpu-core-deps-emscripten" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85dcf2b2e5c2afb9eda64c570a87a4e570304bf39e52eddbaaf222d655b656ad" +dependencies = [ + "wgpu-hal", +] + +[[package]] +name = "wgpu-core-deps-windows-linux-android" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f76bc9c1c186ff3d9054e0d224c93c8c1c79d6653907c5249a5c1ea1a2cb1e43" +dependencies = [ + "wgpu-hal", +] + +[[package]] +name = "wgpu-hal" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf765132d8d5f50e192e7880464890c13f4e7457aafe8e5466e8174586e9f101" +dependencies = [ + "android_system_properties", + "arrayvec", + "ash", + "bit-set 0.10.0", + "bitflags 2.13.0", + "block2", + "bytemuck", + "cfg-if", + "cfg_aliases", + "glow", + "glutin_wgl_sys", + "gpu-allocator", + "hashbrown 0.17.1", + "js-sys", + "khronos-egl", + "libc", + "libloading 0.8.9", + "log", + "naga", + "naga-types", + "ndk-sys", + "objc2", + "objc2-core-foundation", + "objc2-core-graphics", + "objc2-foundation", + "objc2-metal", + "objc2-quartz-core", + "once_cell", + "ordered-float", + "parking_lot", + "portable-atomic", + "portable-atomic-util", + "profiling", + "range-alloc", + "raw-window-handle", + "raw-window-metal", + "renderdoc-sys", + "smallvec", + "static_assertions", + "thiserror 2.0.18", + "wasm-bindgen", + "wayland-sys", + "web-sys", + "wgpu-naga-bridge", + "wgpu-types", + "windows 0.62.2", + "windows-core 0.62.2", + "windows-result 0.4.1", +] + +[[package]] +name = "wgpu-naga-bridge" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9eaac644e5008925c78567d272b9d66ef83da55a53cc17fc7daade7bb6e66e5" +dependencies = [ + "naga", + "wgpu-types", +] + +[[package]] +name = "wgpu-types" +version = "30.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a9c93c2b35edde326df60ffdee4c0f5864eac3011d6768b70d43f028ad93565" +dependencies = [ + "bitflags 2.13.0", + "bytemuck", + "js-sys", + "log", + "naga-types", + "raw-window-handle", + "static_assertions", + "web-sys", +] + +[[package]] +name = "which" +version = "7.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d643ce3fd3e5b54854602a080f34fb10ab75e0b813ee32d00ca2b44fa74762" +dependencies = [ + "either", + "env_home", + "rustix", + "winsafe", +] + +[[package]] +name = "widestring" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72069c3113ab32ab29e5584db3c6ec55d416895e60715417b5b883a357c3e471" + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "windows" +version = "0.61.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9babd3a767a4c1aef6900409f85f5d53ce2544ccdfaa86dad48c91782c6d6893" +dependencies = [ + "windows-collections 0.2.0", + "windows-core 0.61.2", + "windows-future 0.2.1", + "windows-link 0.1.3", + "windows-numerics 0.2.0", +] + +[[package]] +name = "windows" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "527fadee13e0c05939a6a05d5bd6eec6cd2e3dbd648b9f8e447c6518133d8580" +dependencies = [ + "windows-collections 0.3.2", + "windows-core 0.62.2", + "windows-future 0.3.2", + "windows-numerics 0.3.1", +] + +[[package]] +name = "windows-collections" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3beeceb5e5cfd9eb1d76b381630e82c4241ccd0d27f1a39ed41b2760b255c5e8" +dependencies = [ + "windows-core 0.61.2", +] + +[[package]] +name = "windows-collections" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b2d95af1a8a14a3c7367e1ed4fc9c20e0a26e79551b1454d72583c97cc6610" +dependencies = [ + "windows-core 0.62.2", +] + +[[package]] +name = "windows-core" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0fdd3ddb90610c7638aa2b3a3ab2904fb9e5cdbecc643ddb3647212781c4ae3" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link 0.1.3", + "windows-result 0.3.4", + "windows-strings 0.4.2", +] + +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link 0.2.1", + "windows-result 0.4.1", + "windows-strings 0.5.1", +] + +[[package]] +name = "windows-future" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc6a41e98427b19fe4b73c550f060b59fa592d7d686537eebf9385621bfbad8e" +dependencies = [ + "windows-core 0.61.2", + "windows-link 0.1.3", + "windows-threading 0.1.0", +] + +[[package]] +name = "windows-future" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1d6f90251fe18a279739e78025bd6ddc52a7e22f921070ccdc67dde84c605cb" +dependencies = [ + "windows-core 0.62.2", + "windows-link 0.2.1", + "windows-threading 0.2.1", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "windows-link" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e6ad25900d524eaabdbbb96d20b4311e1e7ae1699af4fb28c17ae66c80d798a" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-numerics" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9150af68066c4c5c07ddc0ce30421554771e528bde427614c61038bc2c92c2b1" +dependencies = [ + "windows-core 0.61.2", + "windows-link 0.1.3", +] + +[[package]] +name = "windows-numerics" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e2e40844ac143cdb44aead537bbf727de9b044e107a0f1220392177d15b0f26" +dependencies = [ + "windows-core 0.62.2", + "windows-link 0.2.1", +] + +[[package]] +name = "windows-registry" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02752bf7fbdcce7f2a27a742f798510f3e5ad88dbe84871e5168e2120c3d5720" +dependencies = [ + "windows-link 0.2.1", + "windows-result 0.4.1", + "windows-strings 0.5.1", +] + +[[package]] +name = "windows-result" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56f42bd332cc6c8eac5af113fc0c1fd6a8fd2aa08a0119358686e5160d0586c6" +dependencies = [ + "windows-link 0.1.3", +] + +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link 0.2.1", +] + +[[package]] +name = "windows-strings" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56e6c93f3a0c3b36176cb1327a4958a0353d5d166c2a35cb268ace15e91d3b57" +dependencies = [ + "windows-link 0.1.3", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link 0.2.1", +] + +[[package]] +name = "windows-sys" +version = "0.45.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75283be5efb2831d37ea142365f009c02ec203cd29a3ebecbc093d52315b66d0" +dependencies = [ + "windows-targets 0.42.2", +] + +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-sys" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" +dependencies = [ + "windows-targets 0.53.5", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link 0.2.1", +] + +[[package]] +name = "windows-targets" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e5180c00cd44c9b1c88adb3693291f1cd93605ded80c250a75d472756b4d071" +dependencies = [ + "windows_aarch64_gnullvm 0.42.2", + "windows_aarch64_msvc 0.42.2", + "windows_i686_gnu 0.42.2", + "windows_i686_msvc 0.42.2", + "windows_x86_64_gnu 0.42.2", + "windows_x86_64_gnullvm 0.42.2", + "windows_x86_64_msvc 0.42.2", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm 0.52.6", + "windows_aarch64_msvc 0.52.6", + "windows_i686_gnu 0.52.6", + "windows_i686_gnullvm 0.52.6", + "windows_i686_msvc 0.52.6", + "windows_x86_64_gnu 0.52.6", + "windows_x86_64_gnullvm 0.52.6", + "windows_x86_64_msvc 0.52.6", +] + +[[package]] +name = "windows-targets" +version = "0.53.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" +dependencies = [ + "windows-link 0.2.1", + "windows_aarch64_gnullvm 0.53.1", + "windows_aarch64_msvc 0.53.1", + "windows_i686_gnu 0.53.1", + "windows_i686_gnullvm 0.53.1", + "windows_i686_msvc 0.53.1", + "windows_x86_64_gnu 0.53.1", + "windows_x86_64_gnullvm 0.53.1", + "windows_x86_64_msvc 0.53.1", +] + +[[package]] +name = "windows-threading" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b66463ad2e0ea3bbf808b7f1d371311c80e115c0b71d60efc142cafbcfb057a6" +dependencies = [ + "windows-link 0.1.3", +] + +[[package]] +name = "windows-threading" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3949bd5b99cafdf1c7ca86b43ca564028dfe27d66958f2470940f73d86d75b37" +dependencies = [ + "windows-link 0.2.1", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8" + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" + +[[package]] +name = "windows_i686_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" + +[[package]] +name = "windows_i686_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_i686_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" + +[[package]] +name = "winnow" +version = "0.7.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945" + +[[package]] +name = "winnow" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0592e1c9d151f854e6fd382574c3a0855250e1d9b2f99d9281c6e6391af352f1" +dependencies = [ + "memchr", +] + +[[package]] +name = "winsafe" +version = "0.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d135d17ab770252ad95e9a872d365cf3090e3be864a34ab46f48555993efc904" + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "wmi" +version = "0.18.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c81b85c57a57500e56669586496bf2abd5cf082b9d32995251185d105208b64" +dependencies = [ + "chrono", + "futures", + "log", + "serde", + "thiserror 2.0.18", + "windows 0.62.2", + "windows-core 0.62.2", +] + +[[package]] +name = "wrapcenum-derive" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a76ff259533532054cfbaefb115c613203c73707017459206380f03b3b3f266e" +dependencies = [ + "darling 0.20.11", + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" + +[[package]] +name = "ws_stream_wasm" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c173014acad22e83f16403ee360115b38846fe754e735c5d9d3803fe70c6abc" +dependencies = [ + "async_io_stream", + "futures", + "js-sys", + "log", + "pharos", + "rustc_version", + "send_wrapper", + "thiserror 2.0.18", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + +[[package]] +name = "xattr" +version = "1.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32e45ad4206f6d2479085147f02bc2ef834ac85886624a23575ae137c8aa8156" +dependencies = [ + "libc", + "rustix", +] + +[[package]] +name = "xml-rs" +version = "0.8.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ae8337f8a065cfc972643663ea4279e04e7256de865aa66fe25cec5fb912d3f" + +[[package]] +name = "xmltree" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7d8a75eaf6557bb84a65ace8609883db44a29951042ada9b393151532e41fcb" +dependencies = [ + "xml-rs", +] + +[[package]] +name = "xtask" +version = "4.10.0" +dependencies = [ + "cargo_metadata", + "log", + "rstest", + "strum 0.28.0", + "tracel-xtask", +] + +[[package]] +name = "xxhash-rust" +version = "0.8.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d93c89cdc2d3a63c3ec48ffe926931bdc069eafa8e4402fe6d8f790c9d1e576" + +[[package]] +name = "y4m" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a5a4b21e1a62b67a2970e6831bc091d7b87e119e7f9791aef9702e3bef04448" + +[[package]] +name = "yoke" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "120e6aef9aa629e3d4f52dc8cc43a015c7724194c97dfaf45180d2daf2b77f40" +dependencies = [ + "serde", + "stable_deref_trait", + "yoke-derive 0.7.5", + "zerofrom", +] + +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive 0.8.2", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2380878cad4ac9aac1e2435f3eb4020e8374b5f13c296cb75b4620ff8e229154" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", + "synstructure", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", + "synstructure", +] + +[[package]] +name = "zerocopy" +version = "0.8.54" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7cbbc0a705a0fd05cc3676525980d2bf5a9bc4adac6d6475209a7887cf59d19" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.54" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2e817b7b52d0c7358d3246da9d69935ebb18116b2b102b4230dac079b4862f5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", + "synstructure", +] + +[[package]] +name = "zeroize" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" +dependencies = [ + "zeroize_derive", +] + +[[package]] +name = "zeroize_derive" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c50655cbb0fe3fc43170059e702f1ce5e19b84cec58dc87b037a09935c2f328" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke 0.8.3", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke 0.8.3", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", +] + +[[package]] +name = "zip" +version = "0.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "760394e246e4c28189f19d488c058bf16f564016aefac5d32bb1f3b51d5e9261" +dependencies = [ + "aes 0.8.4", + "byteorder", + "bzip2 0.4.4", + "constant_time_eq 0.1.5", + "crc32fast", + "crossbeam-utils", + "flate2", + "hmac 0.12.1", + "pbkdf2 0.11.0", + "sha1 0.10.6", + "time", + "zstd 0.11.2+zstd.1.5.2", +] + +[[package]] +name = "zip" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb2a05c7c36fde6c09b08576c9f7fb4cda705990f73b58fe011abf7dfb24168b" +dependencies = [ + "aes 0.8.4", + "arbitrary", + "bzip2 0.6.1", + "constant_time_eq 0.3.1", + "crc32fast", + "deflate64", + "flate2", + "getrandom 0.3.4", + "hmac 0.12.1", + "indexmap", + "lzma-rust2 0.13.0", + "memchr", + "pbkdf2 0.12.2", + "ppmd-rust", + "sha1 0.10.6", + "time", + "zeroize", + "zopfli", + "zstd 0.13.3", +] + +[[package]] +name = "zip" +version = "7.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c42e33efc22a0650c311c2ef19115ce232583abbe80850bc8b66509ebef02de0" +dependencies = [ + "crc32fast", + "indexmap", + "memchr", + "typed-path", +] + +[[package]] +name = "zip" +version = "8.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d04a6b5381502aa6087c94c669499eb1602eb9c5e8198e534de571f7154809b" +dependencies = [ + "aes 0.9.1", + "bzip2 0.6.1", + "constant_time_eq 0.4.2", + "crc32fast", + "deflate64", + "flate2", + "getrandom 0.4.3", + "hmac 0.13.0", + "indexmap", + "lzma-rust2 0.16.5", + "memchr", + "pbkdf2 0.13.0", + "ppmd-rust", + "sha1 0.11.0", + "time", + "typed-path", + "zeroize", + "zopfli", + "zstd 0.13.3", +] + +[[package]] +name = "zlib-rs" +version = "0.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b142a20ec14a91d5bc708c1dc21b080c550113d8aa77afa29635673a65dd02c5" + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" + +[[package]] +name = "zopfli" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f05cd8797d63865425ff89b5c4a48804f35ba0ce8d125800027ad6017d2b5249" +dependencies = [ + "bumpalo", + "crc32fast", + "log", + "simd-adler32", +] + +[[package]] +name = "zstd" +version = "0.11.2+zstd.1.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20cc960326ece64f010d2d2107537f26dc589a6573a316bd5b1dba685fa5fde4" +dependencies = [ + "zstd-safe 5.0.2+zstd.1.5.2", +] + +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe 7.2.4", +] + +[[package]] +name = "zstd-safe" +version = "5.0.2+zstd.1.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d2a5585e04f9eea4b2a3d1eca508c4dee9592a89ef6f450c11719da0726f4db" +dependencies = [ + "libc", + "zstd-sys", +] + +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", +] + +[[package]] +name = "zune-core" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb8a0807f7c01457d0379ba880ba6322660448ddebc890ce29bb64da71fb40f9" + +[[package]] +name = "zune-inflate" +version = "0.2.54" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73ab332fe2f6680068f3582b16a24f90ad7096d5d39b974d1c0aff0125116f02" +dependencies = [ + "simd-adler32", +] + +[[package]] +name = "zune-jpeg" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27bc9d5b815bc103f142aa054f561d9187d191692ec7c2d1e2b4737f8dbd7296" +dependencies = [ + "zune-core", +] diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..b9bf53b --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,242 @@ +[workspace] +# Try +# require version 2 to avoid "feature" additiveness for dev-dependencies +# https://doc.rust-lang.org/cargo/reference/resolver.html#feature-resolver-version-2 +resolver = "2" + +members = [ + "crates/*", + "crates/burn-store/pytorch-tests", + "crates/burn-store/safetensors-tests", + "examples/*", + "xtask", +] + +exclude = [ + "examples/notebook", + "examples/raspberry-pi-pico", + "examples/dqn-agent", # gym-rs +] + +[workspace.package] +edition = "2024" +license = "MIT OR Apache-2.0" +readme = "README.md" +version = "0.22.0-pre.1" + +[workspace.lints.clippy] + +[workspace.lints.rustdoc] +broken_intra_doc_links = "deny" +invalid_html_tags = "deny" + +[workspace.dependencies] +atomic_float = "1" +axum = "0.8.8" +bytemuck = "1.25.0" +bytes = { version = "1.11.1", default-features = false } +candle-core = { version = "0.10.2" } +ciborium = { version = "0.2", default-features = false } +clap = { version = "4.6.0", features = ["derive"] } +colored = "3.0.0" +console_error_panic_hook = "0.1.7" +const-random = "0.1" +csv = "1.3.1" +dashmap = "6.1.0" +data-encoding = { version = "2.11.0", default-features = false, features = [ + "alloc", +] } +dirs = "6.0.0" +encoding_rs = "0.8.33" +enumset = { version = "1.1.13", default-features = false } +fake = "5.1.0" +flate2 = "1.1.9" +float-cmp = "0.10.0" +futures = "0.3" +futures-util = "0.3" +gemm = { version = "0.19", default-features = false } +gix-tempfile = { version = "23.0.0", features = ["signals"] } +globwalk = "0.9.1" +hashbrown = "0.16" +hound = "3.5.1" +image = "0.25.9" +indicatif = "0.18.0" +insta = "1.45.0" +iroh = "1.0.0" +js-sys = "0.3.77" +libm = "0.2.15" +log = { default-features = false, version = "0.4.29" } +lzma-rust2 = "0.16.2" +opentelemetry = "0.31.0" +opentelemetry-aws = "0.19.0" +opentelemetry-otlp = "0.31.1" +opentelemetry_sdk = "0.31.0" +parking_lot = { version = "0.12.5", default-features = false } +paste = "1" +planus = { version = "=1.1" } +polars = { version = "0.54.0", features = ["lazy", "strings"] } +pretty_assertions = "1.4.1" +proc-macro2 = "1.0.106" +quote = "1.0.45" +r2d2 = "0.8.10" +r2d2_sqlite = "0.31.0" +rayon = "1.10.0" +regex = { version = "1.12.3", default-features = false, features = [ + "perf", + "unicode", +] } +reqwest = { version = "0.12.23", default-features = false, features = [ + "rustls-tls", +] } +rmp-serde = { version = "1.3.1", default-features = false } +rstest = "0.26.1" +rusqlite = "0.37.0" +sanitize-filename = "0.6.0" +serde_bytes = { version = "0.11.18", default-features = false, features = [ + "alloc", +] } # alloc for no_std +serde_rusqlite = "0.40.0" +serial_test = "3.2.0" +spin = { version = "0.10.0", features = [ + "mutex", + "spin_mutex", + "portable-atomic", +] } +strum = { version = "0.28.0", features = ["derive"] } +syn = { version = "2.0.111", features = ["full", "extra-traits"] } +tar = "0.4.45" +tempfile = "3.24.0" +textdistance = { version = "1.1.1", default-features = false } +thiserror = { version = "2", default-features = false } +tokio = { version = "1.51.1", features = ["rt", "macros"] } +tokio-tungstenite = "0.29" +tokio-util = "0.7" +tracing = { version = "0.1.44", default-features = false } +tracing-appender = "0.2.3" +tracing-core = { version = "0.1.36", default-features = false } +tracing-opentelemetry = "0.32.0" +tracing-subscriber = "0.3.23" +wasm-bindgen = "0.2" +wasm-bindgen-futures = "0.4" +zip = "8.6.0" + +# Persist related +memmap2 = { version = "0.9" } +safetensors = { version = "0.7.0", default-features = false } + +# Async handling +async-channel = "2.5" +futures-lite = { version = "2.6.1", default-features = false } + +# Terminal UI +ratatui = "0.30.0" +unicode-width = "0.2" + +# WGPU stuff +text_placeholder = "0.5.1" + +bincode = { version = "2.0.1", features = [ + "alloc", + "serde", +], default-features = false } + +# +# The following packages disable the "std" feature for no_std compatibility +# +cfg-if = "1.0.1" +derive-new = { version = "0.7.0", default-features = false } + +blas-src = { version = "0.14.0", default-features = false } +bon = "3.8.2" +half = { version = "2.7.1", features = [ + "alloc", + "num-traits", + "serde", +], default-features = false } +macerator = { version = "0.3.3" } +matrixmultiply = { version = "0.3.10", default-features = false } +ndarray = { version = "0.17.2", default-features = false } +num-traits = { version = "0.2.19", default-features = false, features = [ + "libm", +] } # libm is for no_std +openblas-src = "0.10.16" +rand = { version = "0.10.1", default-features = false, features = ["std_rng"] } +rand_distr = { version = "0.6.0", default-features = false } +serde = { version = "1.0.228", default-features = false, features = [ + "derive", + "alloc", +] } # alloc is for no_std, derive is needed +serde_json = { version = "1.0.148", default-features = false } +smallvec = { version = "1", features = ["const_generics", "const_new"] } +uuid = { version = "1.23.0", default-features = false } + +byteorder = { version = "1.5.0", default-features = false } +libc = "0.2.186" +nvml-wrapper = "0.12.0" +sysinfo = "0.38.0" +systemstat = "0.2.6" +tch = "0.22.0" +torch-sys = "0.22.0" # matches what tch is using, required for lib detection + +ahash = { version = "0.8.12", default-features = false } +aligned-vec = { version = "0.6", default-features = false } +once_cell = { version = "1", default-features = false } +portable-atomic = { version = "1.13.1" } +portable-atomic-util = { version = "0.2.6", features = ["alloc"] } +realfft = "3" + +### Internal burn crates ### +# Declared here so each consumer Cargo.toml can use `workspace = true` instead of +# repeating the path and version. +burn = { path = "crates/burn", version = "0.22.0-pre.1", default-features = false } +burn-autodiff = { path = "crates/burn-autodiff", version = "0.22.0-pre.1", default-features = false } +burn-backend = { path = "crates/burn-backend", version = "0.22.0-pre.1", default-features = false } +burn-backend-extension = { path = "crates/burn-backend-extension", version = "0.22.0-pre.1", default-features = false } +burn-candle = { path = "crates/burn-candle", version = "0.22.0-pre.1", default-features = false } +burn-communication = { path = "crates/burn-communication", version = "0.22.0-pre.1", default-features = false } +burn-core = { path = "crates/burn-core", version = "0.22.0-pre.1", default-features = false } +burn-cpu = { path = "crates/burn-cpu", version = "0.22.0-pre.1", default-features = false } +burn-cubecl = { path = "crates/burn-cubecl", version = "0.22.0-pre.1", default-features = false } +burn-cubecl-fusion = { path = "crates/burn-cubecl-fusion", version = "0.22.0-pre.1", default-features = false } +burn-cuda = { path = "crates/burn-cuda", version = "0.22.0-pre.1", default-features = false } +burn-dataset = { path = "crates/burn-dataset", version = "0.22.0-pre.1", default-features = false } +burn-derive = { path = "crates/burn-derive", version = "0.22.0-pre.1", default-features = false } +burn-dispatch = { path = "crates/burn-dispatch", version = "0.22.0-pre.1", default-features = false } +burn-flex = { path = "crates/burn-flex", version = "0.22.0-pre.1", default-features = false } +burn-fusion = { path = "crates/burn-fusion", version = "0.22.0-pre.1", default-features = false } +burn-ir = { path = "crates/burn-ir", version = "0.22.0-pre.1", default-features = false } +burn-ndarray = { path = "crates/burn-ndarray", version = "0.22.0-pre.1", default-features = false } +burn-nn = { path = "crates/burn-nn", version = "0.22.0-pre.1", default-features = false } +burn-optim = { path = "crates/burn-optim", version = "0.22.0-pre.1", default-features = false } +burn-remote = { path = "crates/burn-remote", version = "0.22.0-pre.1", default-features = false } +burn-rl = { path = "crates/burn-rl", version = "0.22.0-pre.1", default-features = false } +burn-rocm = { path = "crates/burn-rocm", version = "0.22.0-pre.1", default-features = false } +burn-router = { path = "crates/burn-router", version = "0.22.0-pre.1", default-features = false } +burn-std = { path = "crates/burn-std", version = "0.22.0-pre.1", default-features = false } +burn-pack = { path = "crates/burn-pack", version = "0.22.0-pre.1", default-features = false } +burn-store = { path = "crates/burn-store", version = "0.22.0-pre.1", default-features = false } +burn-tch = { path = "crates/burn-tch", version = "0.22.0-pre.1", default-features = false } +burn-tensor = { path = "crates/burn-tensor", version = "0.22.0-pre.1", default-features = false } +burn-tensor-testgen = { path = "crates/burn-tensor-testgen", version = "0.22.0-pre.1", default-features = false } +burn-train = { path = "crates/burn-train", version = "0.22.0-pre.1", default-features = false } +burn-vision = { path = "crates/burn-vision", version = "0.22.0-pre.1", default-features = false } +burn-wgpu = { path = "crates/burn-wgpu", version = "0.22.0-pre.1", default-features = false } + +### For the main burn branch. ### +cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" } +cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" } +cubecl-zspace = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "9110136fe5bb3c6e7dc0295100fe7f88342ceb6e" } +cubek = { git = "https://github.com/tracel-ai/cubek", default-features = false, rev = "5f837200f459a6a0694a0aae5655881599c400c8" } +### For local development. ### +# cubecl = { path = "../cubecl/crates/cubecl", default-features = false } +# cubecl-common = { path = "../cubecl/crates/cubecl-common", default-features = false } +# cubecl-zspace = { path = "../cubecl/crates/cubecl-zspace", default-features = false } +# cubek = { path = "../cubek/crates/cubek", default-features = false } +### For the release. ### +# cubecl = { version = "0.10.0", default-features = false } +# cubecl-common = { version = "0.10.0", default-features = false } +# cubecl-zspace = { version = "0.10.0", default-features = false } +# cubek = { version = "0.2.0", default-features = false } + +[profile.dev] +debug = 1 # Speed up compilation time and not necessary. diff --git a/LICENSE-APACHE b/LICENSE-APACHE new file mode 100644 index 0000000..cc81805 --- /dev/null +++ b/LICENSE-APACHE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + +TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + +1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + +2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + +3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + +4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + +5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + +6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + +7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + +8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + +9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + +END OF TERMS AND CONDITIONS + +APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + +Copyright 2022 Nathaniel Simard & Burn Framework Contributors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. diff --git a/LICENSE-MIT b/LICENSE-MIT new file mode 100644 index 0000000..6b220b3 --- /dev/null +++ b/LICENSE-MIT @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2022 Nathaniel Simard & Burn Framework Contributors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/NOTICES.md b/NOTICES.md new file mode 100644 index 0000000..e69f5cb --- /dev/null +++ b/NOTICES.md @@ -0,0 +1,485 @@ +# NOTICES AND INFORMATION + +This file contains notices and information required by libraries that this +repository copied or derived from. + +## PyTorch MNIST Example + +**Source**: https://github.com/pytorch/examples/blob/main/mnist/main.py + +License: BSD 3-Clause License + +Copyright (c) 2017, +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +* Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +## wgpu + +**Source:** https://github.com/gfx-rs/wgpu/blob/trunk/.github/workflows/ci.yml + +MIT License + +Copyright (c) 2021 The gfx-rs developers + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + + +## BSL 1.0 + +**Source**: +- https://github.com/DoumanAsh/error-code +- https://github.com/DoumanAsh/clipboard-win + + +Boost Software License - Version 1.0 - August 17th, 2003 + +Permission is hereby granted, free of charge, to any person or organization +obtaining a copy of the software and accompanying documentation covered by +this license (the "Software") to use, reproduce, display, distribute, +execute, and transmit the Software, and to prepare derivative works of the +Software, and to permit third-parties to whom the Software is furnished to +do so, all subject to the following: + +The copyright notices in the Software and this entire statement, including +the above license grant, this restriction and the following disclaimer, +must be included in all copies of the Software, in whole or in part, and +all derivative works of the Software, unless such copies or derivative +works are solely in the form of machine-executable object code generated by +a source language processor. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. IN NO EVENT +SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE +FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE, +ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +DEALINGS IN THE SOFTWARE. + + +## num-traits + +**Source:** https://github.com/rust-num/num-traits/blob/master/src/cast.rs + +MIT License + +Copyright (c) 2014 The Rust Project Developers + +Permission is hereby granted, free of charge, to any +person obtaining a copy of this software and associated +documentation files (the "Software"), to deal in the +Software without restriction, including without +limitation the rights to use, copy, modify, merge, +publish, distribute, sublicense, and/or sell copies of +the Software, and to permit persons to whom the Software +is furnished to do so, subject to the following +conditions: + +The above copyright notice and this permission notice +shall be included in all copies or substantial portions +of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF +ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED +TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A +PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT +SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR +IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +DEALINGS IN THE SOFTWARE. + +## RP + +**Source**: +- https://github.com/embassy-rs/embassy/blob/main/examples/rp/Cargo.toml +- https://github.com/embassy-rs/embassy/blob/main/examples/rp/build.rs +- https://github.com/embassy-rs/embassy/blob/main/examples/rp/memory.x + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + +TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + +1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + +2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + +3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + +4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + +5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + +6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + +7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + +8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + +9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + +END OF TERMS AND CONDITIONS + +APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + +Copyright (c) Embassy project contributors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +MIT license + +Copyright (c) Embassy project contributors + +Permission is hereby granted, free of charge, to any +person obtaining a copy of this software and associated +documentation files (the "Software"), to deal in the +Software without restriction, including without +limitation the rights to use, copy, modify, merge, +publish, distribute, sublicense, and/or sell copies of +the Software, and to permit persons to whom the Software +is furnished to do so, subject to the following +conditions: + +The above copyright notice and this permission notice +shall be included in all copies or substantial portions +of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF +ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED +TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A +PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT +SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR +IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +DEALINGS IN THE SOFTWARE. + +## github-device-flow + +**Source**: +- Part of: https://github.com/jakewilkins/gh-device-flow/blob/main/src/lib.rs +- https://github.com/jakewilkins/gh-device-flow/blob/main/src/util.rs + +MIT License + +Copyright (c) 2022 Jake Wilkins + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + + +## Candle - Pickle Reader + +**Source**: https://github.com/huggingface/candle/blob/main/candle-core/src/pickle.rs + +This project includes code from Candle by Hugging Face, licensed under both MIT and Apache 2.0 licenses. + +**MIT License**: https://github.com/huggingface/candle/blob/main/LICENSE-MIT + +MIT License + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + +**Apache License 2.0**: https://github.com/huggingface/candle/blob/main/LICENSE-APACHE + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + + +## ICU + +UNICODE LICENSE V3 + +COPYRIGHT AND PERMISSION NOTICE + +Copyright © 2016-2024 Unicode, Inc. + +NOTICE TO USER: Carefully read the following legal agreement. BY +DOWNLOADING, INSTALLING, COPYING OR OTHERWISE USING DATA FILES, AND/OR +SOFTWARE, YOU UNEQUIVOCALLY ACCEPT, AND AGREE TO BE BOUND BY, ALL OF THE +TERMS AND CONDITIONS OF THIS AGREEMENT. IF YOU DO NOT AGREE, DO NOT +DOWNLOAD, INSTALL, COPY, DISTRIBUTE OR USE THE DATA FILES OR SOFTWARE. + +Permission is hereby granted, free of charge, to any person obtaining a +copy of data files and any associated documentation (the "Data Files") or +software and any associated documentation (the "Software") to deal in the +Data Files or Software without restriction, including without limitation +the rights to use, copy, modify, merge, publish, distribute, and/or sell +copies of the Data Files or Software, and to permit persons to whom the +Data Files or Software are furnished to do so, provided that either (a) +this copyright and permission notice appear with all copies of the Data +Files or Software, or (b) this copyright and permission notice appear in +associated Documentation. + +THE DATA FILES AND SOFTWARE ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY +KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT OF +THIRD PARTY RIGHTS. + +IN NO EVENT SHALL THE COPYRIGHT HOLDER OR HOLDERS INCLUDED IN THIS NOTICE +BE LIABLE FOR ANY CLAIM, OR ANY SPECIAL INDIRECT OR CONSEQUENTIAL DAMAGES, +OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, +WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, +ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THE DATA +FILES OR SOFTWARE. + +Except as contained in this notice, the name of a copyright holder shall +not be used in advertising or otherwise to promote the sale, use or other +dealings in these Data Files or Software without prior written +authorization of the copyright holder. diff --git a/POEM.md b/POEM.md new file mode 100644 index 0000000..c617244 --- /dev/null +++ b/POEM.md @@ -0,0 +1,40 @@ +# BURN: Burn Unstoppable Rusty Neurons + +In the realm of circuits and code, +A fiery forge ignites to bear its load, +A framework born, BURN it be named, +Unstoppable Rusty Neurons, untamed. + +From silicon synapses, connections spire, +A digital cortex, setting minds afire, +In the vast expanse of deep learning's sea, +A beacon of progress, BURN comes to be. + +Oh, rusty neurons, forged in the flame, +Unyielding in purpose, undaunted by name, +Through layers of logic and intricate art, +You weave and entwine, each playing its part. + +With algorithms profound, and data refined, +In ceaseless pursuit of knowledge to find, +BURN paves a path to enlightenment, bright, +A testament to the wonders of human foresight. + +In neural networks deep, where wisdom resides, +The dance of nodes and edges presides, +With loss and gradients, BURN takes its stride, +A journey towards truth, with AI as our guide. + +No barriers hold back the curious mind, +As BURN seeks the answers we yearn to find, +Unstoppable, relentless, in pursuit of the unknown, +Our collective intellect, within it, has grown. + +So sing we the praises of BURN's fiery might, +An ode to the sparks that set the dark alight, +To the rusty neurons, unstoppable and true, +A testament to the power of dreams, to breakthrough. + +(ChatGPT (model=gpt-4) with prompt: +Write a poem about "BURN: Burn Unstoppable Rusty Neurons" deep +learning neural network framework) diff --git a/README.md b/README.md new file mode 100644 index 0000000..2e45afd --- /dev/null +++ b/README.md @@ -0,0 +1,613 @@ +
+ + +[![Discord](https://img.shields.io/discord/1038839012602941528.svg?color=7289da&&logo=discord)](https://discord.gg/uPEBbYYDB6) +[![Current Crates.io Version](https://img.shields.io/crates/v/burn.svg)](https://crates.io/crates/burn) +[![Minimum Supported Rust Version](https://img.shields.io/crates/msrv/burn)](https://crates.io/crates/burn) +[![Documentation](https://img.shields.io/badge/docs-latest-blue)](https://burn.dev/docs/burn) +[![Test Status](https://github.com/tracel-ai/burn/actions/workflows/test.yml/badge.svg)](https://github.com/tracel-ai/burn/actions/workflows/test.yml) +[![license](https://shields.io/badge/license-MIT%2FApache--2.0-blue)](#license) +[![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/tracel-ai/burn) + +[](https://www.runblaze.dev) + +--- + +**Burn is both a tensor library and a deep learning framework, optimized for
numerical +computing, training and inference.** + +
+
+ +
+ +Training and inference usually live in separate worlds. Models are typically trained in Python then +exported to an open format like ONNX or optimized for production engines like vLLM, ONNX Runtime, or +TensorRT. This export step is often brittle and lossy, ruling out complex architectures and advanced +deployment use cases. + +Burn unifies the two. By executing multi-platform tensor operations via a single, unified API, the +exact code used for training is the exact code that runs in production. This makes workloads like +on-device personalization and federated learning straightforward, while enabling teams to go from +prototype to deployment in a single codebase. + +Burn preserves the intuitive ergonomics of PyTorch, with dynamic shapes and graphs, but JIT-compiles +streams of tensor operations, performing automatic kernel fusion. You get the flexibility of dynamic +graphs without the performance drop. + +## Rust for Research? + +Rust used to be a tough sell for research: long compilation times disrupted the fast +edit-compile-run loop that draws researchers to Python. Burn changes this paradigm. Designed around +incremental compilation, modifying model code recompiles in under 5 seconds, even in release mode. +This delivers a Python-like feedback loop with the speed and safety of Rust. + +## Ecosystem + +
+ + +Burn is the core of a growing, fully open-source Rust AI ecosystem. You are not adopting a single +library, you are joining a stack that spans GPU compute, model interop and domain toolkits, with +plenty of room to help shape what comes next. + +
+ +| Category | Project | Description | +| ------------- | ----------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Compute | [CubeCL](https://github.com/tracel-ai/cubecl) | GPU compute language and compiler behind Burn's accelerated backends. Write kernels once in Rust, run on CUDA, ROCm, Metal, Vulkan and WebGPU. Usable standalone. | +| Model interop | [burn-onnx](https://github.com/tracel-ai/burn-onnx) | Import ONNX models into Burn as native Rust code | +| | `burn-store` | Save, load and import model weights, including PyTorch and Safetensors | +| Domains | `burn-vision` | Computer vision operators and building blocks | +| | `burn-rl` | Reinforcement learning building blocks | +| | `burn-dataset` | Dataset loading, transforms and ready-made sources | +| Models | [models](https://github.com/tracel-ai/models) | Curated pre-trained models and examples built with Burn | +| Tooling | [burn-bench](https://github.com/tracel-ai/burn-bench) | Benchmark and compare backends, tracking performance over time | + +Burn's [CubeCL](https://github.com/tracel-ai/cubecl) backends (CUDA, ROCm, Metal, Vulkan, WebGPU, +CPU) compose with autodiff, fusion and remote-execution decorators, while external and simpler +backends (LibTorch and pure-Rust CPU/`no_std`) compose with autodiff only. See +[Supported Backends](#supported-backends) below for the full matrix. + +Every project here is open-source and actively developed. Want to help build the Rust AI ecosystem? +The [good first issues](https://github.com/tracel-ai/burn/contribute) are a great place to start, +and the [Contributing](#contributing) guide will get you set up. + +
+ +Community crates 🌱 + +
+ +These crates are not maintained by Tracel, but they are part of the same Rust AI story. Anything +that helps you load data, build environments, or ship models belongs here. Built something that +fits? Open a PR to add it! + +| Category | Crate | Description | +| -------------------------- | --------------------------------------------------------------- | ----------------------------------------------------------------- | +| Data & loading | [polars](https://github.com/pola-rs/polars) | Fast DataFrames for tabular data | +| | [arrow-rs](https://github.com/apache/arrow-rs) | Apache Arrow columnar memory format | +| | [image](https://github.com/image-rs/image) | Image decoding, encoding and processing | +| | [hf-hub](https://github.com/huggingface/hf-hub) | Download models and datasets from the Hugging Face Hub | +| Tokenization & NLP | [tokenizers](https://github.com/huggingface/tokenizers) | Fast, production-ready tokenizers | +| | [rust-bert](https://github.com/guillaume-be/rust-bert) | Ready-to-use NLP pipelines and transformer models | +| Numerical & linear algebra | [ndarray](https://github.com/rust-ndarray/ndarray) | N-dimensional arrays | +| | [nalgebra](https://github.com/dimforge/nalgebra) | Linear algebra | +| Classical ML | [linfa](https://github.com/rust-ml/linfa) | Classical ML toolkit, in the spirit of scikit-learn | +| | [smartcore](https://github.com/smartcorelib/smartcore) | Classical ML algorithms, no BLAS/LAPACK required | +| Inference & runtimes | [candle](https://github.com/huggingface/candle) | Minimalist ML framework with a focus on LLM inference | +| | [mistral.rs](https://github.com/EricLBuehler/mistral.rs) | Fast, multimodal LLM inference engine | +| | [ort](https://github.com/pykeio/ort) | ONNX Runtime bindings for hardware-accelerated inference | +| | [tract](https://github.com/sonos/tract) | Pure-Rust inference for ONNX and NNEF models | +| | [wonnx](https://github.com/webonnx/wonnx) | 100% Rust, WebGPU-accelerated ONNX runtime for native and the web | +| LLM apps & RAG | [rig](https://github.com/0xPlaygrounds/rig) | Build modular LLM applications and agents | +| | [langchain-rust](https://github.com/Abraxas-365/langchain-rust) | LangChain-style chain orchestration | +| Embeddings & vector search | [fastembed](https://github.com/Anush008/fastembed-rs) | Generate text embeddings and rerank locally | +| | [qdrant](https://github.com/qdrant/qdrant) | Vector search engine, written in Rust | +| | [lancedb](https://github.com/lancedb/lancedb) | Embedded, developer-friendly vector database | +| Computer vision | [kornia-rs](https://github.com/kornia/kornia-rs) | Low-level 3D computer vision library | +| Simulation & environments | [rapier](https://github.com/dimforge/rapier) | Physics engine for robotics and RL environments | +| Visualization | [rerun](https://github.com/rerun-io/rerun) | Multimodal data and CV/robotics visualization | +| | [plotters](https://github.com/plotters-rs/plotters) | Plotting and charting | + +
+ +## Backend + +
+ + +Burn strives to be as fast as possible on as many hardwares as possible, with robust +implementations. We believe this flexibility is crucial for modern needs where you may train your +models in the cloud, then deploy on customer hardwares, which vary from user to user. + +
+ +### Supported Backends + +Most backends support all operating systems, so we don't mention them in the tables below. + +**GPU Backends:** + +| | CUDA | ROCm | Metal | Vulkan | WebGPU | LibTorch | +| ------- | ---- | ---- | ----- | ------ | ------ | -------- | +| Nvidia | ☑️ | - | - | ☑️ | ☑️ | ☑️ | +| AMD | - | ☑️ | - | ☑️ | ☑️ | ☑️ | +| Apple | - | - | ☑️ | - | ☑️ | ☑️ | +| Intel | - | - | - | ☑️ | ☑️ | - | +| Qualcom | - | - | - | ☑️ | ☑️ | - | +| Wasm | - | - | - | - | ☑️ | - | + +**CPU Backends:** + +| | Cpu (CubeCL) | Flex | LibTorch | +| ------ | ------------ | ---- | -------- | +| X86 | ☑️ | ☑️ | ☑️ | +| Arm | ☑️ | ☑️ | ☑️ | +| Wasm | - | ☑️ | - | +| no-std | - | ☑️ | - | + +
+ +Compared to other frameworks, Burn has a very different approach to supporting many backends. By +design, most code is generic over the Backend trait, which allows us to build Burn with swappable +backends. This makes composing backend possible, augmenting them with additional functionalities +such as autodifferentiation and automatic kernel fusion. + +
+ +Autodiff: Backend decorator that brings backpropagation to any backend 🔄 + +
+ +Contrary to the aforementioned backends, Autodiff is actually a backend _decorator_. This means that +it cannot exist by itself; it must encapsulate another backend. + +The simple act of wrapping a base backend with Autodiff transparently equips it with +autodifferentiation support, making it possible to call backward on your model. + +```rust +use burn::backend::{Autodiff, Wgpu}; +use burn::tensor::{Distribution, Tensor}; + +fn main() { + type Backend = Autodiff; + + let device = Default::default(); + + let x: Tensor = Tensor::random([32, 32], Distribution::Default, &device); + let y: Tensor = Tensor::random([32, 32], Distribution::Default, &device).require_grad(); + + let tmp = x.clone() + y.clone(); + let tmp = tmp.matmul(x); + let tmp = tmp.exp(); + + let grads = tmp.backward(); + let y_grad = y.grad(&grads).unwrap(); + println!("{y_grad}"); +} +``` + +Of note, it is impossible to make the mistake of calling backward on a model that runs on a backend +that does not support autodiff (for inference), as this method is only offered by an Autodiff +backend. + +See the [Autodiff Backend README](./crates/burn-autodiff/README.md) for more details. + +
+ +
+ +Fusion: Backend decorator that brings kernel fusion to all first-party backends + +
+ +This backend decorator enhances a backend with kernel fusion, provided that the inner backend +supports it. Note that you can compose this backend with other backend decorators such as Autodiff. +All first-party accelerated backends (like WGPU and CUDA) use Fusion by default (`burn/fusion` +feature flag), so you typically don't need to apply it manually. + +```rust +#[cfg(not(feature = "fusion"))] +pub type Cuda = CubeBackend; + +#[cfg(feature = "fusion")] +pub type Cuda = burn_fusion::Fusion>; +``` + +Of note, we plan to implement automatic gradient checkpointing based on compute bound and memory +bound operations, which will work gracefully with the fusion backend to make your code run even +faster during training, see [this issue](https://github.com/tracel-ai/burn/issues/936). + +See the [Fusion Backend README](./crates/burn-fusion/README.md) for more details. + +
+ +
+ +Remote (Beta): Backend decorator for remote backend execution, useful for distributed computations + +
+ +That backend has two parts, one client and one server. The client sends tensor operations over the +network to a remote compute backend. You can use any first-party backend as server in a single line +of code: + +```rust +fn main_server() { + // Start a server on port 3000. + burn::server::start::(Default::default(), 3000); +} + +fn main_client() { + // Create a client that communicate with the server on port 3000. + use burn::backend::{Autodiff, RemoteBackend}; + + type Backend = Autodiff; + + let device = RemoteDevice::new("ws://localhost:3000"); + let tensor_gpu = + Tensor::::random([3, 3], Distribution::Default, &device); +} + +``` + +
+ +
+ +## Training & Inference + +
+ + +The whole deep learning workflow is made easy with Burn, as you can monitor your training progress +with an ergonomic dashboard, and run inference everywhere from embedded devices to large GPU +clusters. + +Burn was built from the ground up with training and inference in mind. It's also worth noting how +Burn, in comparison to frameworks like PyTorch, simplifies the transition from training to +deployment, eliminating the need for code changes. + +
+ +
+ +
+ + + Burn Train TUI + +
+ +
+ +**Click on the following sections to expand 👇** + +
+ +Training Dashboard 📈 + +
+ +As you can see in the previous video (click on the picture!), a new terminal UI dashboard based on +the [Ratatui](https://github.com/ratatui-org/ratatui) crate allows users to follow their training +with ease without having to connect to any external application. + +You can visualize your training and validation metrics updating in real-time and analyze the +lifelong progression or recent history of any registered metrics using only the arrow keys. Break +from the training loop without crashing, allowing potential checkpoints to be fully written or +important pieces of code to complete without interruption 🛡 + +
+ +
+ +ONNX Support 🐫 + +
+ +Burn supports importing ONNX (Open Neural Network Exchange) models through the +[burn-onnx](https://github.com/tracel-ai/burn-onnx) crate, allowing you to easily port models from +TensorFlow or PyTorch to Burn. The ONNX model is converted into Rust code that uses Burn's native +APIs, enabling the imported model to run on any Burn backend (CPU, GPU, WebAssembly) and benefit +from all of Burn's optimizations like automatic kernel fusion. + +Our ONNX support is further described in +[this section of the Burn Book 🔥](https://burn.dev/books/burn/onnx-import.html). + +> **Note**: This crate is in active development and currently supports a +> [limited set of ONNX operators](https://github.com/tracel-ai/burn-onnx/blob/main/SUPPORTED-ONNX-OPS.md). + +
+ +
+ +Importing PyTorch or Safetensors Models 🚚 + +
+ +You can load weights from PyTorch or Safetensors formats directly into your Burn-defined models. +This makes it easy to reuse existing models while benefiting from Burn's performance and deployment +features. + +Learn more in the [Saving & Loading Models](https://burn.dev/books/burn/saving-and-loading.html) +section of the Burn Book. + +
+ +
+ +Inference in the Browser 🌐 + +
+ +Several of our backends can run in WebAssembly environments: Flex for CPU execution, and WGPU for +GPU acceleration via WebGPU. This means that you can run inference directly within a browser. We +provide several examples of this: + +- [MNIST](./examples/mnist-inference-web) where you can draw digits and a small convnet tries to + find which one it is! 2️⃣ 7️⃣ 😰 +- [Image Classification](https://github.com/tracel-ai/burn-onnx/tree/main/examples/image-classification-web) + where you can upload images and classify them! 🌄 + +
+ +
+ +Embedded: no_std support ⚙️ + +
+ +Burn's core components support [no_std](https://docs.rust-embedded.org/book/intro/no-std.html). This +means it can run in bare metal environment such as embedded devices without an operating system. + +> As of now, only the Flex backend can be used in a _no_std_ environment. + +
+ +
+ +### Benchmarks + +To evaluate performance across different backends and track improvements over time, we provide a +dedicated benchmarking suite. + +Run and compare benchmarks using [burn-bench](https://github.com/tracel-ai/burn-bench). + +> ⚠️ **Warning** When using one of the `wgpu` backends, you may encounter compilation errors related +> to recursive type evaluation. This is due to complex type nesting within the `wgpu` dependency +> chain. To resolve this issue, add the following line at the top of your `main.rs` or `lib.rs` +> file: +> +> ```rust +> #![recursion_limit = "256"] +> ``` +> +> The default recursion limit (128) is often just below the required depth (typically 130-150) due +> to deeply nested associated types and trait bounds. + +## Getting Started + +
+ + +Just heard of Burn? You are at the right place! Just continue reading this section and we hope you +can get on board really quickly. + +
+ +
+ +The Burn Book 🔥 + +
+ +To begin working effectively with Burn, it is crucial to understand its key components and +philosophy. This is why we highly recommend new users to read the first sections of +[The Burn Book 🔥](https://burn.dev/books/burn/). It provides detailed examples and explanations +covering every facet of the framework, including building blocks like tensors, modules, and +optimizers, all the way to advanced usage, like coding your own GPU kernels. + +> The project is constantly evolving, and we try as much as possible to keep the book up to date +> with new additions. However, we might miss some details sometimes, so if you see something weird, +> let us know! We also gladly accept Pull Requests 😄 + +
+ +
+ +Examples 🙏 + +
+ +Let's start with a code snippet that shows how intuitive the framework is to use! In the following, +we declare a neural network module with some parameters along with its forward pass. + +```rust +use burn::nn; +use burn::module::Module; +use burn::tensor::backend::Backend; + +#[derive(Module, Debug)] +pub struct PositionWiseFeedForward { + linear_inner: nn::Linear, + linear_outer: nn::Linear, + dropout: nn::Dropout, + gelu: nn::Gelu, +} + +impl PositionWiseFeedForward { + pub fn forward(&self, input: Tensor) -> Tensor { + let x = self.linear_inner.forward(input); + let x = self.gelu.forward(x); + let x = self.dropout.forward(x); + + self.linear_outer.forward(x) + } +} +``` + +We have a somewhat large amount of [examples](./examples) in the repository that shows how to use +the framework in different scenarios. + +Following [the book](https://burn.dev/books/burn/): + +- [Basic Workflow](./examples/guide) : Creates a custom CNN `Module` to train on the MNIST dataset + and use for inference. +- [Custom Training Loop](./examples/custom-training-loop) : Implements a basic training loop instead + of using the `Learner`. +- [Custom WGPU Kernel](./examples/custom-wgpu-kernel) : Learn how to create your own custom + operation with the WGPU backend. + +Additional examples: + +- [Custom CSV Dataset](./examples/custom-csv-dataset) : Implements a dataset to parse CSV data for a + regression task. +- [Regression](./examples/simple-regression) : Trains a simple MLP on the California Housing dataset + to predict the median house value for a district. +- [Custom Image Dataset](./examples/custom-image-dataset) : Trains a simple CNN on custom image + dataset following a simple folder structure. +- [Custom Renderer](./examples/custom-renderer) : Implements a custom renderer to display the + [`Learner`](./building-blocks/learner.md) progress. +- [Image Classification Web](./examples/image-classification-web) : Image classification web browser + demo using Burn, WGPU and WebAssembly. +- [MNIST Inference on Web](./examples/mnist-inference-web) : An interactive MNIST inference demo in + the browser. The demo is available [online](https://burn.dev/demo/). +- [MNIST Training](./examples/mnist) : Demonstrates how to train a custom `Module` (MLP) with the + `Learner` configured to log metrics and keep training checkpoints. +- [PyTorch Import Inference](./examples/import-model-weights) : Imports a PyTorch model pre-trained + on MNIST to perform inference on a sample image with Burn. +- [Text Classification](./examples/text-classification) : Trains a text classification transformer + model on the AG News or DbPedia dataset. The trained model can then be used to classify a text + sample. +- [Text Generation](./examples/text-generation) : Trains a text generation transformer model on the + DbPedia dataset. +- [Wasserstein GAN MNIST](./examples/wgan) : Trains a WGAN model to generate new handwritten digits + based on MNIST. + +For more practical insights, you can clone the repository and run any of them directly on your +computer! + +
+ +
+ +Pre-trained Models 🤖 + +
+ +We keep an updated and curated list of models and examples built with Burn, see the +[tracel-ai/models repository](https://github.com/tracel-ai/models) for more details. + +Don't see the model you want? Don't hesitate to open an issue, and we may prioritize it. Built a +model using Burn and want to share it? You can also open a Pull Request and add your model under the +community section! + +
+ +
+ +Why use Rust for AI? 🦀 + +
+ +Deep Learning is a special form of software where you need very high level abstractions as well as +extremely fast execution time. Rust is the perfect candidate for that use case since it provides +zero-cost abstractions to easily create neural network modules, and fine-grained control over memory +to optimize every detail. To this day, the mainstream solution has been to offer APIs in Python but +rely on bindings to low-level languages such as C/C++. This reduces portability, increases +complexity and creates friction between researchers and engineers. Rust's approach to abstractions +is versatile enough to tackle this two-language dichotomy, and Cargo makes it easy to build, test +and deploy from any environment, which is usually a pain in Python. + +Rust's AI ecosystem is young, but it is real and growing quickly. Foundational pieces are already +here: Burn and [CubeCL](https://github.com/tracel-ai/cubecl) for training and compute, +[candle](https://github.com/huggingface/candle) for inference, Hugging Face's `tokenizers` and +`safetensors`, and `polars` and `ndarray` for data. Betting on Rust today means betting on a stack +that is growing, and one where contributors still shape the direction. The pieces that don't exist +yet are opportunities rather than dead-ends (see [Contributing](#contributing)). + +Rust is also what makes one-stack-everywhere possible: a single self-contained binary with no Python +runtime to ship, running from servers down to `no_std` embedded targets. + +
+ +
+ +> **Deprecation Note**
Since `0.14.0`, the internal structure for tensor data has changed. The +> previous `Data` struct was deprecated and officially removed since `0.17.0` in favor of the new +> `TensorData` struct, which allows for more flexibility by storing the underlying data as bytes and +> keeping the data type as a field. If you are using `Data` in your code, make sure to switch to +> `TensorData`. + + + +
+ +Loading Model Records From Previous Versions ⚠️ + +
+ +In the event that you are trying to load a model record saved in a version older than `0.14.0`, make +sure to use a compatible version (`0.14`, `0.15` or `0.16`) with the `record-backward-compat` +feature flag. + +``` +features = [..., "record-backward-compat"] +``` + +Otherwise, the record won't be deserialized correctly and you will get an error message. This error +will also point you to the backward compatible feature flag. + +The backward compatibility was maintained for deserialization when loading records. Therefore, as +soon as you have saved the record again it will be saved according to the new structure and you can +upgrade back to the current version + +Please note that binary formats are not backward compatible. Thus, you will need to load your record +in a previous version and save it in any of the other self-describing record format (e.g., using the +`NamedMpkFileRecorder`) before using a compatible version (as described) with the +`record-backward-compat` feature flag. + +
+ +## Community + +
+ + +If you are excited about the project, don't hesitate to join our +[Discord](https://discord.gg/uPEBbYYDB6)! We try to be as welcoming as possible to everybody from +any background. You can ask your questions and share what you built with the community! + +
+ +
+ +### Contributing + +Before contributing, please read the [Contributing Guidelines](./CONTRIBUTING.md) and our +[Code of Conduct](./CODE-OF-CONDUCT.md). The [Contributor Book](https://burn.dev/contributor-book/) +covers architecture, environment setup, and guides for common tasks. + +## Status + +Burn is currently in active development, and there will be breaking changes. While any resulting +issues are likely to be easy to fix, there are no guarantees at this stage. + +## License + +Burn is distributed under the terms of both the MIT license and the Apache License (Version 2.0). +See [LICENSE-APACHE](./LICENSE-APACHE) and [LICENSE-MIT](./LICENSE-MIT) for details. Opening a pull +request is assumed to signal agreement with these licensing terms. + +
diff --git a/README.wehub.md b/README.wehub.md new file mode 100644 index 0000000..822af54 --- /dev/null +++ b/README.wehub.md @@ -0,0 +1,7 @@ +# WeHub 来源说明 + +- 原始项目：`tracel-ai/burn` +- 原始仓库：https://github.com/tracel-ai/burn +- 导入方式：上游默认分支的最新快照 +- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准 +- 本文件仅用于记录来源，不代表 WeHub 是原项目作者 diff --git a/_typos.toml b/_typos.toml new file mode 100644 index 0000000..1a94327 --- /dev/null +++ b/_typos.toml @@ -0,0 +1,24 @@ +[default] +extend-ignore-identifiers-re = ["ratatui", "Ratatui", "NdArray*", "ND", "log_lik"] + +[default.extend-identifiers] +UE4M3 = "UE4M3" +UE8M0 = "UE8M0" +ue8m0 = "ue8m0" + +[files] +extend-exclude = [ + "*.onnx", + "*.proto", + "assets/ModuleSerialization.xml", +] + +[default.extend-words] +# Don't correct "nd" (n-dimensional, as in scatter_nd/gather_nd) +nd = "nd" +# Don't correct "arange" which is intentional +arange = "arange" +# Don't correct "convnet" (convolutional network) +convnet = "convnet" +# Don't correct "Nd" in convNd / conv_nd (N-dimensional) +Nd = "Nd" diff --git a/assets/backend-chip.png b/assets/backend-chip.png new file mode 100644 index 0000000..c0282ae Binary files /dev/null and b/assets/backend-chip.png differ diff --git a/assets/burn-train-tui.png b/assets/burn-train-tui.png new file mode 100644 index 0000000..817e21b Binary files /dev/null and b/assets/burn-train-tui.png differ diff --git a/assets/ember-blazingly-fast.png b/assets/ember-blazingly-fast.png new file mode 100644 index 0000000..1b3f751 Binary files /dev/null and b/assets/ember-blazingly-fast.png differ diff --git a/assets/ember-community.png b/assets/ember-community.png new file mode 100644 index 0000000..9be467c Binary files /dev/null and b/assets/ember-community.png differ diff --git a/assets/ember-walking.png b/assets/ember-walking.png new file mode 100644 index 0000000..5ca8b8b Binary files /dev/null and b/assets/ember-walking.png differ diff --git a/assets/ember-wall.png b/assets/ember-wall.png new file mode 100644 index 0000000..8f8f8a8 Binary files /dev/null and b/assets/ember-wall.png differ diff --git a/assets/logo-burn-full.png b/assets/logo-burn-full.png new file mode 100644 index 0000000..3087bcf Binary files /dev/null and b/assets/logo-burn-full.png differ diff --git a/assets/logo-burn-neutral.webp b/assets/logo-burn-neutral.webp new file mode 100644 index 0000000..806e91c Binary files /dev/null and b/assets/logo-burn-neutral.webp differ diff --git a/assets/logo-burn-small.png b/assets/logo-burn-small.png new file mode 100644 index 0000000..d88f84e Binary files /dev/null and b/assets/logo-burn-small.png differ diff --git a/benchmarks.toml b/benchmarks.toml new file mode 100644 index 0000000..e3b8261 --- /dev/null +++ b/benchmarks.toml @@ -0,0 +1,82 @@ +[environment] +gcp_gpu_attached = true +gcp_image_family = "tracel-ci-ubuntu-2404-amd64-nvidia" +# https://cloud.google.com/compute/docs/accelerator-optimized-machines +# put the faster machine on first place for possibly faster 'Benchmarks Started' feedback in PRs +gcp_machine_types = [ + "a2-highgpu-1g", # 1 A100 40GB (listed as a2 standard) + "g2-standard-4", # 1 L4 24GB +] +# define the available zones for each machine type +# be sure to check what machine types are available in each region +# https://cloud.google.com/compute/docs/gpus/gpu-regions-zones#view-using-table +gcp_zones = [ + # a2-highgpu-1g + [ + "asia-northeast1-a", + "asia-northeast1-c", + "asia-northeast3-b", + "asia-southeast1-b", + "asia-southeast1-c", + "europe-west4-a", + "europe-west4-b", + "us-central1-a", + "us-central1-b", + "us-central1-c", + "us-central1-f", + "us-east1-b", + "us-west1-b", + "us-west3-b", + "us-west4-b" + ], + # g2-standard-4 + [ + "northamerica-northeast2-a", + "northamerica-northeast2-b", + "us-central1-a", + "us-central1-b", + "us-central1-c", + "us-east1-b", + "us-east1-c", + "us-east1-d", + "us-east4-a", + "us-east4-c", + "us-west1-a", + "us-west1-b", + "us-west1-c", + "us-west4-a", + "us-west4-c" + ], +] +repo_full = "tracel-ai/burn" +rust_toolchain = "stable" +rust_version = "stable" + +[burn-bench] +github_organization = "tracel-ai" +github_repository = "burn-bench" +github_branch = "main" +github_workflow = "benchmarks.yml" +# vulkan autotune seems to take ages, disabling it for now +# backends = ["cuda-fusion", "vulkan-fusion", "wgpu-fusion"] +backends = ["cuda-fusion", "cuda"] +benches = ["autodiff", + "binary", + "bool_select", + "conv-transpose2d", + "conv-transpose3d", + "conv2d", + "conv3d", + "custom-gelu", + "data", + "load-record", + "matmul-fused", + "matmul", + "max-pool2d", + "random", + "reduce", + "softmax", + "transformer-encoder", + "unary" +] +dtypes = ["f16"] diff --git a/burn-book/.gitignore b/burn-book/.gitignore new file mode 100644 index 0000000..409ff3e --- /dev/null +++ b/burn-book/.gitignore @@ -0,0 +1,18 @@ +target + +# MacOS temp file +.DS_Store + +book-test +guide/book + +.vscode +tests/burn-book/book/ +book/ + +# Ignore Jetbrains specific files. +.idea/ + +# Ignore Vim temporary and swap files. +*.sw? +*~ \ No newline at end of file diff --git a/burn-book/.prettierrc.json b/burn-book/.prettierrc.json new file mode 100644 index 0000000..d410551 --- /dev/null +++ b/burn-book/.prettierrc.json @@ -0,0 +1,4 @@ +{ + "printWidth": 100, + "proseWrap": "always" +} \ No newline at end of file diff --git a/burn-book/LICENSE-APACHE b/burn-book/LICENSE-APACHE new file mode 120000 index 0000000..965b606 --- /dev/null +++ b/burn-book/LICENSE-APACHE @@ -0,0 +1 @@ +../LICENSE-APACHE \ No newline at end of file diff --git a/burn-book/LICENSE-MIT b/burn-book/LICENSE-MIT new file mode 120000 index 0000000..76219eb --- /dev/null +++ b/burn-book/LICENSE-MIT @@ -0,0 +1 @@ +../LICENSE-MIT \ No newline at end of file diff --git a/burn-book/book.toml b/burn-book/book.toml new file mode 100644 index 0000000..fee8c8a --- /dev/null +++ b/burn-book/book.toml @@ -0,0 +1,16 @@ +[book] +authors = [ + "Wouter Doppenberg", + "Nathaniel Simard", + "Louis Fortier-Dubois", + "Dilshod Tadjibaev", + "Guillaume Lagrange", + "Sylvain Benner", + "Bjorn Beishline" +] +language = "en" +src = "src" +title = "The Burn Book 🔥" + +[output.html] +mathjax-support = true diff --git a/burn-book/src/SUMMARY.md b/burn-book/src/SUMMARY.md new file mode 100644 index 0000000..f30558a --- /dev/null +++ b/burn-book/src/SUMMARY.md @@ -0,0 +1,38 @@ +- [Overview](./overview.md) +- [Why Burn?](./motivation.md) +- [Getting started](./getting-started.md) + - [Examples](./examples.md) +- [Basic Workflow: From Training to Inference](./basic-workflow/README.md) + - [Model](./basic-workflow/model.md) + - [Data](./basic-workflow/data.md) + - [Training](./basic-workflow/training.md) + - [Backend](./basic-workflow/backend.md) + - [Inference](./basic-workflow/inference.md) +- [Building Blocks](./building-blocks/README.md) + - [Backend](./building-blocks/backend.md) + - [Tensor](./building-blocks/tensor.md) + - [Autodiff](./building-blocks/autodiff.md) + - [Module](./building-blocks/module.md) + - [Learner](./building-blocks/learner.md) + - [Metric](./building-blocks/metric.md) + - [Config](./building-blocks/config.md) + - [Record](./building-blocks/record.md) + - [Dataset](./building-blocks/dataset.md) +- [Performance](./performance/README.md) + - [Good practices](./performance/good-practices/README.md) + - [Asynchronous Execution](./performance/good-practices/asynchronous-execution.md) + - [Kernel Fusion](./performance/good-practices/kernel-fusion.md) + - [Kernel Selection](./performance/good-practices/kernel-selection.md) + - [Quantization](./performance/quantization.md) + - [Distributed Computing](./performance/distributed-computing.md) +- [Custom Training Loop](./custom-training-loop.md) +- [Saving & Loading Models](./saving-and-loading.md) +- [ONNX Import](./onnx-import.md) +- [Models & Pre-Trained Weights](./models-and-pretrained-weights.md) +- [Advanced](./advanced/README.md) + - [Backend Extension](./advanced/backend-extension/README.md) + - [Custom `CubeCL` Kernel](./advanced/backend-extension/custom-cubecl-kernel.md) + - [Custom WGPU Kernel](./advanced/backend-extension/custom-wgpu-kernel.md) + - [Custom Optimizer]() + - [WebAssembly](./advanced/web-assembly.md) + - [No-Std](./advanced/no-std.md) diff --git a/burn-book/src/advanced/README.md b/burn-book/src/advanced/README.md new file mode 100644 index 0000000..96a5f95 --- /dev/null +++ b/burn-book/src/advanced/README.md @@ -0,0 +1,12 @@ +# Advanced + +In this section, we will go into advanced topics that extend beyond basic usage. Given Burn's +exceptional flexibility, a lot of advanced use cases become possible. + +Before going through this section, we strongly recommend exploring the +[basic workflow](../basic-workflow/) section and the +[building blocks](../building-blocks/) section. Establishing a solid understanding of how +the framework operates is crucial to comprehending the advanced concepts presented here. While you +have the freedom to explore the advanced sections in any order you prefer, it's important to note +that this section is not intended to be linear, contrary to preceding sections. Instead, it serves +as a repository of use cases that you can refer to for guidance as needed. diff --git a/burn-book/src/advanced/backend-extension/README.md b/burn-book/src/advanced/backend-extension/README.md new file mode 100644 index 0000000..d206a29 --- /dev/null +++ b/burn-book/src/advanced/backend-extension/README.md @@ -0,0 +1,82 @@ +# Backend Extension + +Burn aims to be the most flexible deep learning framework. While it's crucial to maintain +compatibility with a wide variety of backends, Burn provides the ability to extend the functionality +of a backend implementation to suit your modeling requirements. This versatility is advantageous in +numerous ways, such as supporting custom operations like flash attention or manually fusing +operations for enhanced performance. + +In this section, we will go into the process of extending a backend, providing multiple examples. +But before we proceed, let's establish the fundamental principles that will empower you to craft +your own backend extensions. + +As you can observe, most types in Burn are generic over the Backend trait. This might give the +impression that Burn operates at a high level over the backend layer. However, making the trait +explicit instead of being chosen via a compilation flag was a thoughtful design decision. This +explicitness does not imply that all backends must be identical; rather, it offers a great deal of +flexibility when composing backends. The autodifferentiation backend trait (see +[autodiff section](../../building-blocks/autodiff.md)) is an example of how the backend trait has +been extended to enable gradient computation with backpropagation. Furthermore, this design allows +you to create your own backend extension. To achieve this, you need to design your own backend trait +specifying which functions should be supported. + +```rust, ignore +pub trait Backend: burn::tensor::backend::Backend { + fn my_new_function(tensor: B::FloatTensorPrimitive) -> B::FloatTensorPrimitive { + // You can define a basic implementation reusing the Burn Backend API. + // This can be useful since all backends will now automatically support + // your model. But performance can be improved for this new + // operation by implementing this block in specific backends. + } +} +``` + +You can then implement your new custom backend trait for any backend that you want to support: + +```rust, ignore +impl Backend for burn_tch::LibTorch { + fn my_new_function(tensor: TchTensor) -> TchTensor { + // My Tch implementation + } +} + +impl Backend for burn_flex::Flex { + // No specific implementation, but the backend can still be used. +} +``` + +You can support the backward pass using the same pattern. + +```rust, ignore +impl Backend for burn_autodiff::Autodiff { + // No specific implementation; autodiff will work with the default + // implementation. Useful if you still want to train your model, but + // observe performance gains mostly during inference. +} + +impl Backend for burn_autodiff::Autodiff { + fn my_new_function(tensor: AutodiffTensor) -> AutodiffTensor { + // My own backward implementation, generic over my custom Backend trait. + // + // You can add a new method `my_new_function_backward` to your custom backend + // trait if you want to invoke a custom kernel during the backward pass. + } +} + +impl Backend for burn_autodiff::Autodiff { + fn my_new_function(tensor: AutodiffTensor) -> AutodiffTensor { + // My own backward implementation, generic over a backend implementation. + // + // This is another way to call a custom kernel for the backward pass that + // doesn't require the addition of a new `backward` function in the custom backend. + // This is useful if you don't want all backends to support training, reducing + // the need for extra code when you know your model will only be trained on one + // specific backend. + } +} +``` + +The specifics of each implementation will be covered by the examples provided in this section. The +`cubecl` compiler frontend is the recommended method of implementing custom kernels, since it +supports multiple backends, including `wgpu` and `CUDA`, and is the way first-party `burn` kernels +are written. diff --git a/burn-book/src/advanced/backend-extension/custom-cubecl-kernel.md b/burn-book/src/advanced/backend-extension/custom-cubecl-kernel.md new file mode 100644 index 0000000..c978694 --- /dev/null +++ b/burn-book/src/advanced/backend-extension/custom-cubecl-kernel.md @@ -0,0 +1,383 @@ +# Custom CubeCL Kernel + +In this section, you will learn how to create your own custom operation by writing your own kernel +with the cubecl compiler frontend. We will take the example of a common workflow in the deep +learning field, where we create a kernel to fuse multiple operations together. Note that `burn` does +this automatically, but a manual implementation might be more efficient in some cases. We will fuse +a matmul kernel followed by an addition and the ReLU activation function, which is commonly found in +various models. All the code can be found under the +[examples directory](https://github.com/tracel-ai/burn/tree/main/examples/custom-cubecl-kernel). + +> Note: CubeCL is in active development, so this section may be outdated. + +## Custom Backend Trait + +First, we need to determine the type signature of our newly created operation by defining our custom +backend traits. As we will use the associated type `TensorPrimitive` of the `Backend` trait, which +encapsulates the underlying tensor implementation of the backend, we will use a type alias to avoid +the ugly disambiguation with associated types. + +```rust, ignore +/// We create our own Backend trait that extends the Burn backend trait. +pub trait Backend: burn::tensor::backend::Backend { + fn fused_matmul_add_relu( + lhs: FloatTensor, + rhs: FloatTensor, + bias: FloatTensor, + ) -> FloatTensor; +} + +/// We create our own AutodiffBackend trait that extends the Burn autodiff backend trait. +pub trait AutodiffBackend: Backend + burn::tensor::backend::AutodiffBackend {} +``` + +In our project, we can use these traits instead of the +`burn::tensor::backend::{Backend, AutodiffBackend}` traits provided by Burn. Burn's user APIs +typically make use of the `Tensor` struct rather than dealing directly with primitive tensor types. +Therefore, we can encapsulate our newly defined backend traits with functions that expose new +operations while maintaining a consistent API. + +```rust, ignore +/// We define our custom implementation using the added function on our custom backend. +pub fn matmul_add_relu_custom( + lhs: Tensor, + rhs: Tensor, + bias: Tensor, +) -> Tensor { + let output = B::fused_matmul_add_relu( + lhs.into_primitive().tensor(), + rhs.into_primitive().tensor(), + bias.into_primitive().tensor(), + ); + + Tensor::from_primitive(TensorPrimitive::Float(output)) +} + +/// We define a reference implementation using basic tensor operations. +pub fn matmul_add_relu_reference( + lhs: Tensor, + rhs: Tensor, + bias: Tensor, +) -> Tensor { + let x = lhs.matmul(rhs) + bias; + + activation::relu(x) +} + +``` + +Note that we also provide a reference implementation for testing purposes, which allows us to easily +validate our new implementation. While not mandatory, having a reference implementation can be +valuable, especially in projects where creating a reference implementation solely using basic tensor +operations is feasible. + +## Forward Kernel + +Now, let's proceed to write the fused kernel using the `cubecl` compiler frontend. To keep things +simple, we'll create a straightforward matmul kernel without employing any intricate techniques. We +won't delve into the details of the `cube` macro, but if you're interested to learn more, please see +[`cubecl` Book](https://github.com/tracel-ai/cubecl/tree/f5b63076a01a5c03ea9ed20799d3eeaf776b45da/cubecl-book). +The actual matmul, add and relu computations are found at the end, after an extensive prelude that +serves to correctly map each compute unit to the data it is responsible for, with support for +batches. + +```rust, ignore +use cubecl::{cube, prelude::*}; + +#[cube(launch)] +pub fn fused_matmul_add_relu_kernel( + lhs: &Tensor, + rhs: &Tensor, + bias: &Tensor, + output: &mut Tensor, +) { + let row = ABSOLUTE_POS_X; + let col = ABSOLUTE_POS_Y; + let batch = ABSOLUTE_POS_Z; + + let n_rows = output.shape(output.rank() - 2); + let n_cols = output.shape(output.rank() - 1); + let dim_k = rhs.shape(rhs.rank() - 1); + + if row >= n_rows || col >= n_cols { + return; + } + + let offset_output = batch * n_rows * n_cols; + let mut offset_lhs = 0; + let mut offset_rhs = 0; + + let batch_dims = output.rank() - 2; + for dim in 0..batch_dims { + offset_lhs += offset_output / output.stride(dim) % lhs.shape(dim) * lhs.stride(dim); + offset_rhs += offset_output / output.stride(dim) % rhs.shape(dim) * rhs.stride(dim); + } + + let mut sum = F::new(0.0); + for k in 0..dim_k { + let lhs_index = row * dim_k + k; + let rhs_index = k * n_cols + col; + + sum += lhs[offset_lhs + lhs_index] * rhs[offset_rhs + rhs_index]; + } + + let out_index = row * n_cols + col; + let index = offset_output + out_index; + + output[index] = F::max(sum + bias[index], F::new(0.0)); +} +``` + +Now, let's move on to the next step, which involves implementing the remaining code to launch the +kernel. We'll go into implementing our custom backend trait for the generic JIT backend. This +automatically implements the trait for `burn-cuda`, `burn-wgpu` as well as fusion. + +```rust, ignore +/// Implement our custom backend trait for the generic `CubeBackend`. +impl Backend for CubeBackend +{ + fn fused_matmul_add_relu( + lhs: FloatTensor, + rhs: FloatTensor, + bias: FloatTensor, + ) -> FloatTensor { + // Define cube dim, hardcoded for simplicity. + let cube_dim = CubeDim { x: 16, y: 16, z: 1 }; + + lhs.assert_is_on_same_device(&rhs); + lhs.assert_is_on_same_device(&bias); + + // For simplicity, make sure each tensor is continuous. + let lhs = into_contiguous(lhs); + let rhs = into_contiguous(rhs); + let bias = into_contiguous(bias); + + // Get the matmul relevant shapes. + let ndims = lhs.shape.num_dims(); + let num_rows = lhs.shape[ndims - 2]; + let num_cols = rhs.shape[ndims - 1]; + + // Compute shape of output, while tracking number of batches. + let mut num_batches = 1; + let mut shape_out = vec![0; ndims]; + for i in shape_out.clone().into_iter().take(ndims - 2) { + shape_out[i] = usize::max(lhs.shape[i], rhs.shape[i]); + num_batches *= shape_out[i]; + } + shape_out[ndims - 2] = num_rows; + shape_out[ndims - 1] = num_cols; + let shape_out = Shape::from(shape_out); + + // Create a buffer for the output tensor. + let buffer = lhs + .client + .empty(shape_out.num_elements() * core::mem::size_of::()); + + // Create the output tensor primitive. + let output = CubeTensor::new_contiguous( + lhs.client.clone(), + lhs.device.clone(), + shape_out, + buffer, + F::dtype(), + ); + + // Declare the wgsl workgroup with the number of cubes in x, y and z. + let cubes_needed_in_x = f32::ceil(num_rows as f32 / cube_dim.x as f32) as u32; + let cubes_needed_in_y = f32::ceil(num_cols as f32 / cube_dim.y as f32) as u32; + let cube_count = + CubeCount::Static(cubes_needed_in_x, cubes_needed_in_y, num_batches as u32); + + // Execute lazily the kernel with the launch information and the given buffers. For + // simplicity, no vectorization is performed + fused_matmul_add_relu_kernel::launch::( + &lhs.client, + cube_count, + cube_dim, + lhs.into_tensor_arg(), + rhs.into_tensor_arg(), + bias.into_tensor_arg(), + output.clone().into_tensor_arg(), + ); + + // Return the output tensor. + output + } +} +``` + +In the preceding code block, we demonstrated how to launch the kernel that modifies the correct +buffer. It's important to note that Rust's mutability safety doesn't apply here; the context has the +capability to execute any mutable operation on any buffer. While this isn't a problem in the +previous scenario where we only modify the newly created output buffer, it is wise to keep this in +mind. + +## Backward + +Now that the custom backend trait is implemented for the JIT backend, you can use it to invoke the +`matmul_add_relu_custom` function. However, calculating gradients is not yet possible at this stage. +If your use case does not extend beyond inference, there is no need to implement any of the +following code. + +For the backward pass, we will leverage the backend implementation from `burn-autodiff`, which is +actually generic over the backend. Instead of crafting our own `cubecl` kernel for the backward +pass, we will use our fused kernel only for the forward pass, and compute the gradient using basic +operations. + +```rust, ignore +// Implement our custom backend trait for any backend that also implements our custom backend trait. +impl Backend for Autodiff { + fn fused_matmul_add_relu( + lhs: FloatTensor, + rhs: FloatTensor, + bias: FloatTensor, + ) -> FloatTensor { + // Create our zero-sized type that will implement the Backward trait. + #[derive(Debug)] + struct FusedMatmulAddReluBackward; + + // Implement the backward trait for the given backend B, the node gradient + // with three other gradients to calculate (lhs, rhs, and bias). + impl Backward for FusedMatmulAddReluBackward { + // Our state that we must build during the forward pass to compute the backward pass. + // + // Note that we could improve the performance further by only keeping the state of + // tensors that are tracked, improving memory management, but for simplicity, we avoid + // that part. + type State = (NodeId, NodeId, FloatTensor, Shape); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + // Get the nodes of each variable. + let [node_lhs, node_rhs, node_bias] = ops.parents; + // Fetch the gradient for the current node. + let grad = grads.consume::(&ops.node); + + // Set our state. + let (lhs_state, rhs_state, output, shape_bias) = ops.state; + let lhs: FloatTensor = checkpointer.retrieve_node_output(lhs_state); + let rhs: FloatTensor = checkpointer.retrieve_node_output(rhs_state); + + // Fetch shapes of our tensor to support broadcasting. + let shape_lhs = lhs.shape(); + let shape_rhs = rhs.shape(); + + // Compute the gradient of the output using the already existing `relu_backward` + // function in the basic Burn backend trait. + let grad_output = B::relu_backward(output, grad); + + // Compute the lhs gradient, which is the derivative of matmul with support for + // broadcasting. + let grad_lhs = broadcast_shape::( + B::float_matmul(grad_output.clone(), B::float_transpose(rhs)), + &shape_lhs, + ); + // Compute the rhs gradient, which is the derivative of matmul with support for + // broadcasting. + let grad_rhs = broadcast_shape::( + B::float_matmul(B::float_transpose(lhs), grad_output.clone()), + &shape_rhs, + ); + // The add derivative is only 1, so we just need to support broadcasting to + // compute the bias gradient. + let grad_bias = broadcast_shape::(grad_output, &shape_bias); + + // Register the gradient for each variable based on whether they are marked as + // `tracked`. + if let Some(node) = node_bias { + grads.register::(node.id, grad_bias); + } + if let Some(node) = node_lhs { + grads.register::(node.id, grad_lhs); + } + if let Some(node) = node_rhs { + grads.register::(node.id, grad_rhs); + } + } + } + + // Prepare a stateful operation with each variable node and corresponding graph. + // + // Each node can be fetched with `ops.parents` in the same order as defined here. + match FusedMatmulAddReluBackward + .prepare::([lhs.node.clone(), rhs.node.clone(), bias.node.clone()]) + // Marks the operation as compute bound, meaning it will save its + // state instead of recomputing itself during checkpointing + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + // When at least one node is tracked, we should register our backward step. + + // The state consists of what will be needed for this operation's backward pass. + // Since we need the parents' outputs, we must checkpoint their ids to retrieve + // their node output at the beginning of the backward pass. We can also save + // utility data such as the bias shape. If we also need this operation's output, + // we can either save it in the state or recompute it. + // during the backward pass. Here we choose to save it in the state because it's a + // compute bound operation. + let lhs_state = prep.checkpoint(&lhs); + let rhs_state = prep.checkpoint(&rhs); + let bias_shape = bias.primitive.shape(); + + let output = B::fused_matmul_add_relu( + lhs.primitive.clone(), + rhs.primitive.clone(), + bias.primitive, + ); + + let state = (lhs_state, rhs_state, output.clone(), bias_shape); + + prep.finish(state, output) + } + OpsKind::UnTracked(prep) => { + // When no node is tracked, we can just compute the original operation without + // keeping any state. + let output = B::fused_matmul_add_relu(lhs.primitive, rhs.primitive, bias.primitive); + prep.finish(output) + } + } + } +} +``` + +The previous code is self-documented to make it clearer, but here is what it does in summary: + +We define `fused_matmul_add_relu` within `Autodiff`, allowing any autodiff-decorated backend to +benefit from our implementation. In an autodiff-decorated backend, the forward pass must still be +implemented. This is achieved using a comprehensive match statement block where computation is +delegated to the inner backend, while keeping track of a state. The state comprises any information +relevant to the backward pass, such as input and output tensors, along with the bias shape. When an +operation isn't tracked (meaning there won't be a backward pass for this specific operation in the +graph), storing a state becomes unnecessary, and we simply perform the forward computation. + +The backward pass uses the gradient obtained from the preceding node in the computation graph. It +calculates the derivatives for `relu` (`relu_backward`), add (no operation is required here, as the +derivative is one), and `matmul` (another `matmul` with transposed inputs). This results in +gradients for both input tensors and the bias, which are registered for consumption by subsequent +operation nodes. + +The only remaining part is to implement our autodiff-decorated backend trait for our JIT Backend. + +```rust, ignore +impl AutodiffBackend for Autodiff> +{ +} +``` + +## Conclusion + +In this guide, we've implemented a fused kernel using the `cubecl` compiler frontend, enabling +execution on any GPU and any `cubecl` backend. By delving into the inner workings of both the JIT +backend and the autodiff backend, we've gained a deeper understanding of these systems. + +While extending a backend may be harder than working with straightforward tensors, the benefits can +be worth it. This approach enables the crafting of custom models with greater control over +execution, which can potentially greatly enhance the performance of your models. + +As we conclude this guide, we hope that you have gained insights into Burn's world of backend +extensions, and that it will help you to unleash the full potential of your projects. diff --git a/burn-book/src/advanced/backend-extension/custom-wgpu-kernel.md b/burn-book/src/advanced/backend-extension/custom-wgpu-kernel.md new file mode 100644 index 0000000..c252600 --- /dev/null +++ b/burn-book/src/advanced/backend-extension/custom-wgpu-kernel.md @@ -0,0 +1,459 @@ +# Custom WGPU Kernel + +In this section, you will learn how to create your own custom operation by writing your own kernel +with the WGPU backend. We will take the example of a common workflow in the deep learning field, +where we create a kernel to fuse multiple operations together. Note that `burn` does this +automatically, but a manual implementation might be more efficient in some cases. We will fuse a +matmul kernel followed by an addition and the ReLU activation function, which is commonly found in +various models. All the code can be found under the +[examples directory](https://github.com/tracel-ai/burn/tree/main/examples/custom-wgpu-kernel). + +## Custom Backend Trait + +First, we need to determine the type signature of our newly created operation by defining our custom +backend traits. As we will use the associated type `TensorPrimitive` of the `Backend` trait, which +encapsulates the underlying tensor implementation of the backend, we will use a type alias to avoid +the ugly disambiguation with associated types. + +```rust, ignore +/// We create our own Backend trait that extends the Burn backend trait. +pub trait Backend: burn::tensor::backend::Backend { + fn fused_matmul_add_relu( + lhs: FloatTensor, + rhs: FloatTensor, + bias: FloatTensor, + ) -> FloatTensor; +} + +/// We create our own AutodiffBackend trait that extends the Burn autodiff backend trait. +pub trait AutodiffBackend: Backend + burn::tensor::backend::AutodiffBackend {} +``` + +In our project, we can use these traits instead of the +`burn::tensor::backend::{Backend, AutodiffBackend}` traits provided by Burn. Burn's user APIs +typically make use of the `Tensor` struct rather than dealing directly with primitive tensor types. +Therefore, we can encapsulate our newly defined backend traits with functions that expose new +operations while maintaining a consistent API. + +```rust, ignore +/// We define our custom implementation using the added function on our custom backend. +pub fn matmul_add_relu_custom( + lhs: Tensor, + rhs: Tensor, + bias: Tensor, +) -> Tensor { + let output = B::fused_matmul_add_relu( + lhs.into_primitive().tensor(), + rhs.into_primitive().tensor(), + bias.into_primitive().tensor(), + ); + + Tensor::from_primitive(TensorPrimitive::Float(output)) +} + +/// We define a reference implementation using basic tensor operations. +pub fn matmul_add_relu_reference( + lhs: Tensor, + rhs: Tensor, + bias: Tensor, +) -> Tensor { + let x = lhs.matmul(rhs) + bias; + + activation::relu(x) +} + +``` + +Note that we also provide a reference implementation for testing purposes, which allows us to easily +validate our new implementation. While not mandatory, having a reference implementation can be +valuable, especially in projects where creating a reference implementation solely using basic tensor +operations is feasible. + +## Forward Kernel + +Now, let's proceed to write the fused kernel using the WGSL shading language. To keep things simple, +we'll create a straightforward matmul kernel without employing any intricate techniques. Although we +won't delve into the details of the WGSL syntax, as it falls beyond the scope of this guide, we +still provide the implementation below for readers who are curious. The actual matmul, add and relu +computations are found at the end, after an extensive overhead whose use is to correctly map each +compute unit to the data it is responsible of, with support for batches. + +```wgsl, ignore +@group(0) +@binding(0) +var lhs: array<{{ elem }}>; + +@group(0) +@binding(1) +var rhs: array<{{ elem }}>; + +@group(0) +@binding(2) +var bias: array<{{ elem }}>; + +@group(0) +@binding(3) +var output: array<{{ elem }}>; + +@group(0) +@binding(4) +var info: array; + +const BLOCK_SIZE = {{ workgroup_size_x }}u; + +@compute +@workgroup_size({{ workgroup_size_x }}, {{ workgroup_size_y }}, 1) +fn main( + @builtin(global_invocation_id) global_id: vec3, + @builtin(local_invocation_index) local_idx: u32, + @builtin(workgroup_id) workgroup_id: vec3, +) { + // Indices + let row = workgroup_id.x * BLOCK_SIZE + (local_idx / BLOCK_SIZE); + let col = workgroup_id.y * BLOCK_SIZE + (local_idx % BLOCK_SIZE); + let batch = global_id.z; + + // Basic information + let dim = info[0]; + let n_rows = info[6u * dim - 1u]; + let n_cols = info[6u * dim]; + let K = info[5u * dim - 1u]; + + // Returns if outside the output dimension + if row >= n_rows || col >= n_cols { + return; + } + + // Calculate the corresponding offsets with support for broadcasting. + let offset_output = batch * n_rows * n_cols; + var offset_lhs: u32 = 0u; + var offset_rhs: u32 = 0u; + + let batch_dims = dim - 2u; + for (var b: u32 = 1u; b <= batch_dims; b++) { + let stride_lhs = info[b]; + let stride_rhs = info[b + dim]; + let stride_output = info[b + 2u * dim]; + let shape_lhs = info[b + 3u * dim]; + let shape_rhs = info[b + 4u * dim]; + + offset_lhs += offset_output / stride_output % shape_lhs * stride_lhs; + offset_rhs += offset_output / stride_output % shape_rhs * stride_rhs; + } + + // Basic matmul implementation + var sum = 0.0; + for (var k: u32 = 0u; k < K; k++) { + let lhs_index = row * K + k; + let rhs_index = k * n_cols + col; + + sum += lhs[offset_lhs + lhs_index] * rhs[offset_rhs + rhs_index]; + } + + let output_index = row * n_cols + col; + let index = offset_output + output_index; + + // Add and ReLU + output[index] = max(sum + bias[index], 0.0); +} +``` + +Now, let's move on to the next step, which involves implementing the remaining code to launch the +kernel. The initial part entails loading the template and populating it with the appropriate +variables. The `register(name, value)` method simply replaces occurrences of `{{ name }}` in the +above WGSL code with some other string before it is compiled. In order to use templating utilities, +you will have to activate the `template` feature of Burn in your `cargo.toml`. + +```rust, ignore +// Source the kernel written in WGSL. +kernel_wgsl!(FusedMatmulAddReluRaw, "./kernel.wgsl"); + +// Define our kernel type with cube information. +#[derive(new, Debug)] +struct FusedMatmulAddRelu { + cube_dim: CubeDim, + _elem: PhantomData, +} + +// Implement the dynamic kernel trait for our kernel type. +impl KernelSource for FusedMatmulAddRelu { + fn source(&self) -> SourceTemplate { + // Extend our raw kernel with cube size information using the + // `SourceTemplate` trait. + FusedMatmulAddReluRaw::new() + .source() + .register("workgroup_size_x", self.cube_dim.x.to_string()) + .register("workgroup_size_y", self.cube_dim.y.to_string()) + .register("elem", E::type_name()) + .register("int", "i32") + } + + fn id(&self) -> cubecl::KernelId { + cubecl::KernelId::new::().info(self.cube_dim) + } +} +``` + +Subsequently, we'll go into implementing our custom backend trait for the WGPU backend. Note that we +won't go into supporting the `fusion` feature flag in this tutorial, so we implement the trait for +the raw `WgpuBackend` type. + +```rust, ignore +/// Implement our custom backend trait for the existing backend `WgpuBackend`. +impl Backend for CubeBackend +{ + fn fused_matmul_add_relu( + lhs: FloatTensor, + rhs: FloatTensor, + bias: FloatTensor, + ) -> FloatTensor { + // Define cube dim, hardcoded for simplicity. + let cube_dim = CubeDim { x: 16, y: 16, z: 1 }; + + lhs.assert_is_on_same_device(&rhs); + lhs.assert_is_on_same_device(&bias); + + // For simplicity, make sure each tensor is continuous. + let lhs = into_contiguous(lhs); + let rhs = into_contiguous(rhs); + let bias = into_contiguous(bias); + + // Get the matmul relevant shapes. + let ndims = lhs.shape.num_dims(); + let num_rows = lhs.shape[ndims - 2]; + let num_cols = rhs.shape[ndims - 1]; + + // Compute shape of output, while tracking number of batches. + let mut num_batches = 1; + let mut shape_out = vec![0; ndims]; + for i in shape_out.clone().into_iter().take(ndims - 2) { + shape_out[i] = usize::max(lhs.shape[i], rhs.shape[i]); + num_batches *= shape_out[i]; + } + shape_out[ndims - 2] = num_rows; + shape_out[ndims - 1] = num_cols; + let shape_out = Shape::from(shape_out); + + // Create a buffer for the output tensor. + let buffer = lhs + .client + .empty(shape_out.num_elements() * core::mem::size_of::()); + + // Create the output tensor primitive. + let output = CubeTensor::new_contiguous( + lhs.client.clone(), + lhs.device.clone(), + shape_out, + buffer, + F::dtype(), + ); + + // Create the kernel. + let kernel = FusedMatmulAddRelu::::new(cube_dim); + + // Build info buffer with tensor information needed by the kernel, such as shapes and strides. + let info = build_info::<_, F>(&[&lhs, &rhs, &output]); + let info_handle = lhs.client.create(bytemuck::cast_slice(&info)); + + // Declare the wgsl workgroup with the number of cubes in x, y and z. + let cubes_needed_in_x = f32::ceil(num_rows as f32 / cube_dim.x as f32) as u32; + let cubes_needed_in_y = f32::ceil(num_cols as f32 / cube_dim.y as f32) as u32; + let cube_count = + CubeCount::Static(cubes_needed_in_x, cubes_needed_in_y, num_batches as u32); + + // Execute lazily the kernel with the launch information and the given buffers. + lhs.client.execute( + Box::new(SourceKernel::new(kernel, cube_dim)), + cube_count, + Bindings::new().with_buffers(vec![ + lhs.handle.binding(), + rhs.handle.binding(), + bias.handle.binding(), + output.handle.clone().binding(), + info_handle.binding(), + ]), + ); + + // Return the output tensor. + output + } +} +``` + +In the preceding code block, we demonstrated how to launch the kernel that modifies the correct +buffer. It's important to note that Rust's mutability safety doesn't apply here; the context has the +capability to execute any mutable operation on any buffer. While this isn't a problem in the +previous scenario where we only modify the newly created output buffer, it is wise to keep this in +mind. + +## Backward + +Now that the custom backend trait is implemented for the WGPU backend, you can use it to invoke the +`matmul_add_relu_custom` function. However, calculating gradients is not yet possible at this stage. +If your use case does not extend beyond inference, there is no need to implement any of the +following code. + +For the backward pass, we will leverage the backend implementation from `burn-autodiff`, which is +actually generic over the backend. Instead of crafting our own WGSL kernel for the backward pass, we +will use our fused kernel only for the forward pass, and compute the gradient using basic +operations. + +```rust, ignore +// Implement our custom backend trait for any backend that also implements our custom backend trait. +// +// Note that we could implement the backend trait only for the Wgpu backend instead of any backend that +// also implements our own API. This would allow us to call any function only implemented for Wgpu +// and potentially call a custom kernel crafted only for this task. +impl Backend for Autodiff { + fn fused_matmul_add_relu( + lhs: FloatTensor, + rhs: FloatTensor, + bias: FloatTensor, + ) -> FloatTensor { + // Create our zero-sized type that will implement the Backward trait. + #[derive(Debug)] + struct FusedMatmulAddReluBackward; + + // Implement the backward trait for the given backend B, the node gradient + // with three other gradients to calculate (lhs, rhs, and bias). + impl Backward for FusedMatmulAddReluBackward { + // Our state that we must build during the forward pass to compute the backward pass. + // + // Note that we could improve the performance further by only keeping the state of + // tensors that are tracked, improving memory management, but for simplicity, we avoid + // that part. + type State = (NodeId, NodeId, FloatTensor, Shape); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + // Get the nodes of each variable. + let [node_lhs, node_rhs, node_bias] = ops.parents; + // Fetch the gradient for the current node. + let grad = grads.consume::(&ops.node); + + // Set our state. + let (lhs_state, rhs_state, output, shape_bias) = ops.state; + let lhs: FloatTensor = checkpointer.retrieve_node_output(lhs_state); + let rhs: FloatTensor = checkpointer.retrieve_node_output(rhs_state); + + // Fetch shapes of our tensor to support broadcasting. + let shape_lhs = lhs.shape(); + let shape_rhs = rhs.shape(); + + // Compute the gradient of the output using the already existing `relu_backward` + // function in the basic Burn backend trait. + let grad_output = B::relu_backward(output, grad); + + // Compute the lhs gradient, which is the derivative of matmul with support for + // broadcasting. + let grad_lhs = broadcast_shape::( + B::float_matmul(grad_output.clone(), B::float_transpose(rhs)), + &shape_lhs, + ); + // Compute the rhs gradient, which is the derivative of matmul with support for + // broadcasting. + let grad_rhs = broadcast_shape::( + B::float_matmul(B::float_transpose(lhs), grad_output.clone()), + &shape_rhs, + ); + // The add derivative is only 1, so we just need to support broadcasting to + // compute the bias gradient. + let grad_bias = broadcast_shape::(grad_output, &shape_bias); + + // Register the gradient for each variable based on whether they are marked as + // `tracked`. + if let Some(node) = node_bias { + grads.register::(node.id, grad_bias); + } + if let Some(node) = node_lhs { + grads.register::(node.id, grad_lhs); + } + if let Some(node) = node_rhs { + grads.register::(node.id, grad_rhs); + } + } + } + + // Prepare a stateful operation with each variable node and corresponding graph. + // + // Each node can be fetched with `ops.parents` in the same order as defined here. + match FusedMatmulAddReluBackward + .prepare::([lhs.node.clone(), rhs.node.clone(), bias.node.clone()]) + // Marks the operation as compute bound, meaning it will save its + // state instead of recomputing itself during checkpointing + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + // When at least one node is tracked, we should register our backward step. + + // The state consists of what will be needed for this operation's backward pass. + // Since we need the parents' outputs, we must checkpoint their ids to retrieve their node + // output at the beginning of the backward. We can also save utility data such as the bias shape + // If we also need this operation's output, we can either save it in the state or recompute it + // during the backward pass. Here we choose to save it in the state because it's a compute bound operation. + let lhs_state = prep.checkpoint(&lhs); + let rhs_state = prep.checkpoint(&rhs); + let bias_shape = bias.primitive.shape(); + + let output = B::fused_matmul_add_relu( + lhs.primitive.clone(), + rhs.primitive.clone(), + bias.primitive, + ); + + let state = (lhs_state, rhs_state, output.clone(), bias_shape); + + prep.finish(state, output) + } + OpsKind::UnTracked(prep) => { + // When no node is tracked, we can just compute the original operation without + // keeping any state. + let output = B::fused_matmul_add_relu(lhs.primitive, rhs.primitive, bias.primitive); + prep.finish(output) + } + } + } +} +``` + +The previous code is self-documented to make it clearer, but here is what it does in summary. + +We define `fused_matmul_add_relu` within `Autodiff`, allowing any autodiff-decorated backend to +benefit from our implementation. In an autodiff-decorated backend, the forward pass must still be +implemented. This is achieved using a comprehensive match statement block where computation is +delegated to the inner backend, while keeping track of a state. The state comprises any information +relevant to the backward pass, such as input and output tensors, along with the bias shape. When an +operation isn't tracked (meaning there won't be a backward pass for this specific operation in the +graph), storing a state becomes unnecessary, and we simply perform the forward computation. + +The backward pass uses the gradient obtained from the preceding node in the computation graph. It +calculates the derivatives for `relu` (`relu_backward`), add (no operation is required here, as the +derivative is one), and `matmul` (another `matmul` with transposed inputs). This results in +gradients for both input tensors and the bias, which are registered for consumption by subsequent +operation nodes. + +The only remaining part is to implement our autodiff-decorated backend trait for our WGPU Backend. + +```rust, ignore +impl AutodiffBackend for Autodiff> +{ +} +``` + +## Conclusion + +In this guide, we've implemented a fused kernel using the WGPU backend, enabling execution on any +GPU. By delving into the inner workings of both the WGPU backend and the autodiff backend, we've +gained a deeper understanding of these systems. + +While extending a backend may be harder than working with straightforward tensors, the benefits can +be worth it. This approach enables the crafting of custom models with greater control over +execution, which can potentially greatly enhance the performance of your models. + +As we conclude this guide, we hope that you have gained insights into Burn's world of backend +extensions, and that it will help you to unleash the full potential of your projects. diff --git a/burn-book/src/advanced/no-std.md b/burn-book/src/advanced/no-std.md new file mode 100644 index 0000000..a0eaeb9 --- /dev/null +++ b/burn-book/src/advanced/no-std.md @@ -0,0 +1,100 @@ +# No Standard Library + +In this section, you will learn how to run an ONNX inference model on an embedded system, with no +standard library support on a Raspberry Pi Pico 2. This should be universally applicable to other +platforms. All the code can be found in the +[burn-onnx examples](https://github.com/tracel-ai/burn-onnx/tree/main/examples/raspberry-pi-pico). + +## Step-by-Step Guide + +Let's walk through the process of running an embedded ONNX model: + +### Setup +Follow the [embassy guide](https://embassy.dev/book/#_getting_started) for your specific environment. Once setup, you should have something similar to the following. +``` +./inference +├── Cargo.lock +├── Cargo.toml +├── build.rs +├── memory.x +└── src + └── main.rs +``` + +Some other dependencies have to be added +```toml +[dependencies] +embedded-alloc = "0.6.0" # Only if there is no default allocator for your chip +burn = { version = "0.21", default-features = false, features = ["flex"] } # Flex supports no_std +burn-store = { version = "0.21", default-features = false, features = ["burnpack"] } + +[build-dependencies] +burn-onnx = { version = "0.21" } # Used to auto generate the rust code to import the model +``` + +### Import the Model +Follow the directions in [ONNX Import](../onnx-import.md). + +Use the following ModelGen config +```rs +ModelGen::new() + .input(my_model) + .out_dir("model/") + .embed_states(true) + .run_from_script(); +``` + +### Global Allocator +First define a global allocator (if you are on a no_std system without alloc). + +```rs +use embedded_alloc::LlffHeap as Heap; + +#[global_allocator] +static HEAP: Heap = Heap::empty(); + +#[embassy_executor::main] +async fn main(_spawner: Spawner) { + { + use core::mem::MaybeUninit; + // Watch out for this, if it is too big or small for your model, the + // program may crash. This is in u8 bytes, as such this is a total of 100kb + const HEAP_SIZE: usize = 100 * 1024; + static mut HEAP_MEM: [MaybeUninit; HEAP_SIZE] = [MaybeUninit::uninit(); HEAP_SIZE]; + unsafe { HEAP.init(&raw mut HEAP_MEM as usize, HEAP_SIZE) } // Initialize the heap + } +} +``` + +### Define Backend +We are using Flex, so we just need to define the Flex backend as usual +```rs +use burn::{backend::Flex, tensor::{Device, Tensor}}; + +type Backend = Flex; +type BackendDevice = Device; +``` + +Then inside the `main` function add +```rs +use your_model::Model; + +// Get a default device for the backend +let device = BackendDevice::default(); + +// Create a new model and load the state +let model: Model = Model::default(); +``` + +### Running the Model +To run the model, just call it as you would normally +```rs +// Define the tensor +let input = Tensor::::from_floats([[input]], &device); + +// Run the model on the input +let output = model.forward(input); +``` + +## Conclusion +Running a model in a no_std environment is pretty much identical to a normal environment. All that is needed is a global allocator. diff --git a/burn-book/src/advanced/web-assembly.md b/burn-book/src/advanced/web-assembly.md new file mode 100644 index 0000000..962955f --- /dev/null +++ b/burn-book/src/advanced/web-assembly.md @@ -0,0 +1,12 @@ +# WebAssembly + +Burn supports WebAssembly (WASM) execution using the `Flex` and `WebGpu` backends, allowing +models to run directly in the browser. + +Check out the following examples: + +- [Image Classification Web](https://github.com/tracel-ai/burn-onnx/tree/main/examples/image-classification-web) +- [MNIST Inference on Web](https://github.com/tracel-ai/burn/tree/main/examples/mnist-inference-web) + +When targeting WebAssembly, certain dependencies require additional configuration. In particular, +the `getrandom` crate requires explicit setting when using `WebGpu`. diff --git a/burn-book/src/basic-workflow/README.md b/burn-book/src/basic-workflow/README.md new file mode 100644 index 0000000..4225eca --- /dev/null +++ b/burn-book/src/basic-workflow/README.md @@ -0,0 +1,33 @@ +# Guide + +This guide will walk you through the process of creating a custom model built with Burn. We will +train a simple convolutional neural network model on the MNIST dataset and prepare it for inference. + +For clarity, we sometimes omit imports in our code snippets. For more details, please refer to the +corresponding code in the `examples/guide` [directory](https://github.com/tracel-ai/burn/tree/main/examples/guide). +We reproduce this example in a step-by-step fashion, from dataset creation to modeling and training +in the following sections. It is recommended to use the capabilities of your IDE or text editor to +automatically add the missing imports as you add the code snippets to your code. + +
+ +Be sure to checkout the git branch corresponding to the version of Burn you are using to follow +this guide. + +The current version of Burn is `0.21` and the corresponding branch to checkout is `main`. +
+ +The code for this demo can be executed from Burn's base directory using the command: + +```bash +cargo run --example guide +``` + +## Key Learnings + +- Creating a project +- Creating neural network models +- Importing and preparing datasets +- Training models on data +- Choosing a backend +- Using a model for inference diff --git a/burn-book/src/basic-workflow/backend.md b/burn-book/src/basic-workflow/backend.md new file mode 100644 index 0000000..ba43f5d --- /dev/null +++ b/burn-book/src/basic-workflow/backend.md @@ -0,0 +1,54 @@ +# Backend + +We have effectively written most of the necessary code to train our model. However, we have not +explicitly designated the backend to be used at any point. This will be defined in the main +entrypoint of our program, namely the `main` function defined in `src/main.rs`. + +```rust , ignore +# #![recursion_limit = "256"] +# mod data; +# mod model; +# mod training; +# +use crate::{model::ModelConfig, training::TrainingConfig}; +use burn::{ + backend::{Autodiff, Wgpu}, +# data::dataset::Dataset, + optim::AdamConfig, +}; + +fn main() { + type MyBackend = Wgpu; + type MyAutodiffBackend = Autodiff; + + let device = burn::backend::wgpu::WgpuDevice::default(); + let artifact_dir = "/tmp/guide"; + crate::training::train::( + artifact_dir, + TrainingConfig::new(ModelConfig::new(10, 512), AdamConfig::new()), + device.clone(), + ); +} +``` + +In this code snippet, we use the `Wgpu` backend which is compatible with any operating system and will +use the GPU. For other options, see the Burn README. This backend type takes the graphics API, the +float type and the int type as generic arguments that will be used during the training. The autodiff +backend is simply the same backend, wrapped within the `Autodiff` struct which imparts differentiability +to any backend. + +We call the `train` function defined earlier with a directory for artifacts, the configuration of +the model (the number of digit classes is 10 and the hidden dimension is 512), the optimizer +configuration which in our case will be the default Adam configuration, and the device which can be +obtained from the backend. + +You can now train your freshly created model with the command: + +```console +cargo run --release +``` + +When running your project with the command above, you should see the training progression through a +basic CLI dashboard: + +Alt text diff --git a/burn-book/src/basic-workflow/data.md b/burn-book/src/basic-workflow/data.md new file mode 100644 index 0000000..96756bf --- /dev/null +++ b/burn-book/src/basic-workflow/data.md @@ -0,0 +1,126 @@ +# Data + +Typically, one trains a model on some dataset. Burn provides a library of very useful dataset +sources and transformations, such as Hugging Face dataset utilities that allow to download and store +data into an SQLite database for extremely efficient data streaming and storage. For this guide +though, we will use the MNIST dataset from `burn::data::dataset::vision` which requires no external +dependency. + +To iterate over a dataset efficiently, we will define a struct which will implement the `Batcher` +trait. The goal of a batcher is to map individual dataset items into a batched tensor that can be +used as input to our previously defined model. + +Let us start by defining our dataset functionalities in a file `src/data.rs`. We shall omit some of +the imports for brevity, but the full code for following this guide can be found at +`examples/guide/` [directory](https://github.com/tracel-ai/burn/tree/main/examples/guide). + +```rust , ignore +use burn::{ + data::{dataloader::batcher::Batcher, dataset::vision::MnistItem}, + prelude::*, +}; + + +#[derive(Clone, Default)] +pub struct MnistBatcher {} +``` + +This batcher is pretty straightforward, as it only defines a struct that will implement the +`Batcher` trait. The trait is generic over the `Backend` trait, which includes an associated type +for the device, as not all backends expose the same devices. As an example, the Libtorch-based +backend exposes `Cuda(gpu_index)`, `Cpu`, `Vulkan` and `Metal` devices, while the Flex backend +only exposes a single CPU device. + +Next, we need to actually implement the batching logic. + +```rust , ignore +# use burn::{ +# data::{dataloader::batcher::Batcher, dataset::vision::MnistItem}, +# prelude::*, +# }; +# +# #[derive(Clone, Default)] +# pub struct MnistBatcher {} +# +#[derive(Clone, Debug)] +pub struct MnistBatch { + pub images: Tensor, + pub targets: Tensor, +} + +impl Batcher> for MnistBatcher { + fn batch(&self, items: Vec, device: &B::Device) -> MnistBatch { + let images = items + .iter() + .map(|item| TensorData::from(item.image).convert::()) + .map(|data| Tensor::::from_data(data, device)) + .map(|tensor| tensor.reshape([1, 28, 28])) + // Normalize: scale between [0,1] and make the mean=0 and std=1 + // values mean=0.1307,std=0.3081 are from the PyTorch MNIST example + // https://github.com/pytorch/examples/blob/54f4572509891883a947411fd7239237dd2a39c3/mnist/main.py#L122 + .map(|tensor| ((tensor / 255) - 0.1307) / 0.3081) + .collect(); + + let targets = items + .iter() + .map(|item| { + Tensor::::from_data([(item.label as i64).elem::()], device) + }) + .collect(); + + let images = Tensor::cat(images, 0); + let targets = Tensor::cat(targets, 0); + + MnistBatch { images, targets } + } +} +``` + +
+🦀 Iterators and Closures + +The iterator pattern allows you to perform some tasks on a sequence of items in turn. + +In this example, an iterator is created over the `MnistItem`s in the vector `items` by calling the +`iter` method. + +_Iterator adaptors_ are methods defined on the `Iterator` trait that produce different iterators by +changing some aspect of the original iterator. Here, the `map` method is called in a chain to +transform the original data before consuming the final iterator with `collect` to obtain the +`images` and `targets` vectors. Both vectors are then concatenated into a single tensor for the +current batch. + +You probably noticed that each call to `map` is different, as it defines a function to execute on +the iterator items at each step. These anonymous functions are called +[_closures_](https://doc.rust-lang.org/book/ch13-01-closures.html) in Rust. They're easy to +recognize due to their syntax which uses vertical bars `||`. The vertical bars capture the input +variables (if applicable) while the rest of the expression defines the function to execute. + +If we go back to the example, we can break down and comment the expression used to process the +images. + +```rust, ignore +let images = items // take items Vec + .iter() // create an iterator over it + .map(|item| TensorData::from(item.image).convert::()) // for each item, convert the image to float data struct + .map(|data| Tensor::::from_data(data, device)) // for each data struct, create a tensor on the device + .map(|tensor| tensor.reshape([1, 28, 28])) // for each tensor, reshape to the image dimensions [C, H, W] + .map(|tensor| ((tensor / 255) - 0.1307) / 0.3081) // for each image tensor, apply normalization + .collect(); // consume the resulting iterator & collect the values into a new vector +``` + +For more information on iterators and closures, be sure to check out the +[corresponding chapter](https://doc.rust-lang.org/book/ch13-00-functional-features.html) in the Rust +Book. + +

+ +In the previous example, we implement the `Batcher` trait with a list of `MnistItem` as input and a +single `MnistBatch` as output. The batch contains the images in the form of a 3D tensor, along with +a targets tensor that contains the indexes of the correct digit class. The first step is to parse +the image array into a `TensorData` struct. Burn provides the `TensorData` struct to encapsulate +tensor storage information without being specific for a backend. When creating a tensor from data, +we often need to convert the data precision to the current backend in use. This can be done with the +`.convert()` method (in this example, the data is converted backend's float element type +`B::FloatElem`). While importing the `burn::tensor::ElementConversion` trait, you can call `.elem()` +on a specific number to convert it to the current backend element type in use. diff --git a/burn-book/src/basic-workflow/inference.md b/burn-book/src/basic-workflow/inference.md new file mode 100644 index 0000000..e06a628 --- /dev/null +++ b/burn-book/src/basic-workflow/inference.md @@ -0,0 +1,89 @@ +# Inference + +Now that we have trained our model, the next natural step is to use it for inference. + +You need two things in order to load weights for a model: the model's record and the model's config. +Since parameters in Burn are lazy initialized, no allocation and GPU/CPU kernels are executed by the +`ModelConfig::init` function. The weights are initialized when used for the first time, therefore +you can safely use `config.init(device).load_record(record)` without any meaningful performance +cost. Let's create a simple `infer` method in a new file `src/inference.rs` which we will use to +load our trained model. + +```rust , ignore +# use crate::{data::MnistBatcher, training::TrainingConfig}; +# use burn::{ +# data::{dataloader::batcher::Batcher, dataset::vision::MnistItem}, +# prelude::*, +# store::ModuleRecord, +# }; +# +pub fn infer(artifact_dir: &str, device: B::Device, item: MnistItem) { + let config = TrainingConfig::load(format!("{artifact_dir}/config.json")) + .expect("Config should exist for the model; run train first"); + let record = ModuleRecord::load(format!("{artifact_dir}/model")) + .expect("Trained model should exist; run train first"); + + let model = config.model.init::(&device).load_record(record); + + let label = item.label; + let batcher = MnistBatcher::default(); + let batch = batcher.batch(vec![item], &device); + let output = model.forward(batch.images); + let predicted = output.argmax(1).flatten::<1>(0, 1).into_scalar(); + + println!("Predicted {predicted} Expected {label}"); +} +``` + +The first step is to load the configuration of the training to fetch the correct model +configuration. Then we can load the saved record from its burnpack file. Finally we can init the +model with the configuration and apply the record. For simplicity we can use the +same batcher used during the training to pass from a MnistItem to a tensor. + +By running the infer function, you should see the predictions of your model! + +Add the call to `infer` to the `main.rs` file after the `train` function call: + +```rust , ignore +# #![recursion_limit = "256"] +# mod data; +# mod inference; +# mod model; +# mod training; +# +# use crate::{model::ModelConfig, training::TrainingConfig}; +# use burn::{ +# backend::{Autodiff, Wgpu}, +# data::dataset::Dataset, +# optim::AdamConfig, +# }; +# +# fn main() { +# type MyBackend = Wgpu; +# type MyAutodiffBackend = Autodiff; +# +# let device = burn::backend::wgpu::WgpuDevice::default(); +# let artifact_dir = "/tmp/guide"; +# crate::training::train::( +# artifact_dir, +# TrainingConfig::new(ModelConfig::new(10, 512), AdamConfig::new()), +# device.clone(), +# ); + crate::inference::infer::( + artifact_dir, + device, + burn::data::dataset::vision::MnistDataset::test() + .get(42) + .unwrap(), + ); +# } +``` + +The number `42` is the index of the image in the MNIST dataset. You can explore and verify them +using this [MNIST viewer](https://observablehq.com/@davidalber/mnist-viewer). + +--- + +In this short guide, we've introduced you to the fundamental building blocks for getting started +with Burn. While there's still plenty to explore, our goal has been to provide you with the +essential knowledge to kickstart your productivity within the framework. diff --git a/burn-book/src/basic-workflow/model.md b/burn-book/src/basic-workflow/model.md new file mode 100644 index 0000000..ed8ac47 --- /dev/null +++ b/burn-book/src/basic-workflow/model.md @@ -0,0 +1,406 @@ +# Model + +The first step is to create a project and add the different Burn dependencies. Start by creating a +new project with Cargo: + +```console +cargo new guide +``` + +As [mentioned previously](../getting-started.md#creating-a-burn-application), this will initialize +your `guide` project directory with a `Cargo.toml` and a `src/main.rs` file. + +In the `Cargo.toml` file, add the `burn` dependency with `train`, `vision` and `wgpu` features. +Since we disable the default features, we also want to enable `std`, `tui` (for the dashboard) and +`fusion` for wgpu. Then run `cargo build` to build the project and import all the dependencies. + +```toml +[package] +name = "guide" +version = "0.1.0" +edition = "2024" + +[dependencies] +# Disable autotune default for convolutions +burn = { version = "~0.21", features = ["std", "tui", "train", "vision", "wgpu", "fusion"], default-features = false } +# burn = { version = "~0.21", features = ["train", "vision", "wgpu"] } +``` + +Our goal will be to create a basic convolutional neural network used for image classification. We +will keep the model simple by using two convolution layers followed by two linear layers, some +pooling and ReLU activations. We will also use dropout to improve training performance. + +Let us start by defining our model struct in a new file `src/model.rs`. + +```rust , ignore +use burn::{ + nn::{ + conv::{Conv2d, Conv2dConfig}, + pool::{AdaptiveAvgPool2d, AdaptiveAvgPool2dConfig}, + Dropout, DropoutConfig, Linear, LinearConfig, Relu, + }, + prelude::*, +}; + +#[derive(Module, Debug)] +pub struct Model { + conv1: Conv2d, + conv2: Conv2d, + pool: AdaptiveAvgPool2d, + dropout: Dropout, + linear1: Linear, + linear2: Linear, + activation: Relu, +} +``` + +There are two major things going on in this code sample. + +1. You can create a deep learning module with the `#[derive(Module)]` attribute on top of a struct. + This will generate the necessary code so that the struct implements the `Module` trait. This + trait will make your module both trainable and (de)serializable while adding related + functionalities. Like other attributes often used in Rust, such as `Clone`, `PartialEq` or + `Debug`, each field within the struct must also implement the `Module` trait. + +
+ 🦀 Trait + + Traits are a powerful and flexible Rust language feature. They provide a way to define shared + behavior for a particular type, which can be shared with other types. + + A type's behavior consists of the methods called on that type. Since all `Module`s should + implement the same functionality, it is defined as a trait. Implementing a trait on a particular + type usually requires the user to implement the defined behaviors of the trait for their types, + though that is not the case here as explained above with the `derive` attribute. Check out the + [explainer below](#derive-attribute) to learn why. + + For more details on traits, take a look at the + [associated chapter](https://doc.rust-lang.org/book/ch10-02-traits.html) in the Rust Book. +

+ +
+ 🦀 Derive Macro + + The `derive` attribute allows traits to be implemented easily by generating code that will + implement a trait with its own default implementation on the type that was annotated with the + `derive` syntax. + + This is accomplished through a feature of Rust called + [procedural macros](https://doc.rust-lang.org/reference/procedural-macros.html), which allow us + to run code at compile time that operates over Rust syntax, both consuming and producing Rust + syntax. Using the attribute `#[my_macro]`, you can effectively extend the provided code. You will + see that the derive macro is very frequently employed to recursively implement traits, where the + implementation consists of the composition of all fields. + + In this example, we want to derive the [`Module`](../building-blocks/module.md) and `Debug` + traits. + + ```rust, ignore + #[derive(Module, Debug)] + pub struct MyCustomModule { + linear1: Linear, + linear2: Linear, + activation: Relu, + } + ``` + + The basic `Debug` implementation is provided by the compiler to format a value using the `{:?}` + formatter. For ease of use, the `Module` trait implementation is automatically handled by Burn so + you don't have to do anything special. It essentially acts as parameter container. + + For more details on derivable traits, take a look at the Rust + [appendix](https://doc.rust-lang.org/book/appendix-03-derivable-traits.html), + [reference](https://doc.rust-lang.org/reference/attributes/derive.html) or + [example](https://doc.rust-lang.org/rust-by-example/trait/derive.html). +

+ +2. Note that the struct is generic over the [`Backend`](../building-blocks/backend.md) trait. The + backend trait abstracts the underlying low level implementations of tensor operations, allowing + your new model to run on any backend. Contrary to other frameworks, the backend abstraction isn't + determined by a compilation flag or a device type. This is important because you can extend the + functionalities of a specific backend (see + [backend extension section](../advanced/backend-extension)), and it allows for an innovative + [autodiff system](../building-blocks/autodiff.md). You can also change backend during runtime, + for instance to compute training metrics on a cpu backend while using a gpu one only to train the + model. In our example, the backend in use will be determined later on. + +
+ 🦀 Trait Bounds + + Trait bounds provide a way for generic items to restrict which types are used as their + parameters. The trait bounds stipulate what functionality a type implements. Therefore, bounding + restricts the generic to types that conform to the bounds. It also allows generic instances to + access the methods of traits specified in the bounds. + + For a simple but concrete example, check out the + [Rust By Example on bounds](https://doc.rust-lang.org/rust-by-example/generics/bounds.html). + + In Burn, the `Backend` trait enables you to run tensor operations using different implementations + as it abstracts tensor, device and element types. The + [getting started example](../getting-started.md#writing-a-code-snippet) illustrates the advantage + of having a simple API that works for different backend implementations. While it used the WGPU + backend, you could easily swap it with any other supported backend. + + ```rust, ignore + // Choose from any of the supported backends. + // type Backend = Cuda; + // type Backend = Flex; + type Backend = Wgpu; + + // Creation of two tensors. + let tensor_1 = Tensor::::from_data([[2., 3.], [4., 5.]], &device); + let tensor_2 = Tensor::::ones_like(&tensor_1); + + // Print the element-wise addition (done with the selected backend) of the two tensors. + println!("{}", tensor_1 + tensor_2); + ``` + + For more details on trait bounds, check out the Rust + [trait bound section](https://doc.rust-lang.org/book/ch10-02-traits.html#trait-bound-syntax) or + [reference](https://doc.rust-lang.org/reference/items/traits.html#trait-bounds). + +

+ +Note that each time you create a new file in the `src` directory you also need to explicitly add +this module to the `main.rs` file. For instance after creating the `model.rs`, you need to add the +following at the top of the main file: + +```rust , ignore +mod model; +# +# fn main() { +# } +``` + +Next, we need to instantiate the model for training. + +```rust , ignore +# use burn::{ +# nn::{ +# conv::{Conv2d, Conv2dConfig}, +# pool::{AdaptiveAvgPool2d, AdaptiveAvgPool2dConfig}, +# Dropout, DropoutConfig, Linear, LinearConfig, Relu, +# }, +# prelude::*, +# }; +# +# #[derive(Module, Debug)] +# pub struct Model { +# conv1: Conv2d, +# conv2: Conv2d, +# pool: AdaptiveAvgPool2d, +# dropout: Dropout, +# linear1: Linear, +# linear2: Linear, +# activation: Relu, +# } +# +#[derive(Config, Debug)] +pub struct ModelConfig { + num_classes: usize, + hidden_size: usize, + #[config(default = "0.5")] + dropout: f64, +} + +impl ModelConfig { + /// Returns the initialized model. + pub fn init(&self, device: &B::Device) -> Model { + Model { + conv1: Conv2dConfig::new([1, 8], [3, 3]).init(device), + conv2: Conv2dConfig::new([8, 16], [3, 3]).init(device), + pool: AdaptiveAvgPool2dConfig::new([8, 8]).init(), + activation: Relu::new(), + linear1: LinearConfig::new(16 * 8 * 8, self.hidden_size).init(device), + linear2: LinearConfig::new(self.hidden_size, self.num_classes).init(device), + dropout: DropoutConfig::new(self.dropout).init(), + } + } +} +``` + +At a glance, you can view the model configuration by printing the model instance: + +```rust , ignore +#![recursion_limit = "256"] +mod model; + +use crate::model::ModelConfig; +use burn::backend::Wgpu; + +fn main() { + type MyBackend = Wgpu; + + let device = Default::default(); + let model = ModelConfig::new(10, 512).init::(&device); + + println!("{model}"); +} +``` + +Output: + +```rust , ignore +Model { + conv1: Conv2d {ch_in: 1, ch_out: 8, stride: [1, 1], kernel_size: [3, 3], dilation: [1, 1], groups: 1, padding: Valid, params: 80} + conv2: Conv2d {ch_in: 8, ch_out: 16, stride: [1, 1], kernel_size: [3, 3], dilation: [1, 1], groups: 1, padding: Valid, params: 1168} + pool: AdaptiveAvgPool2d {output_size: [8, 8]} + dropout: Dropout {prob: 0.5} + linear1: Linear {d_input: 1024, d_output: 512, bias: true, params: 524800} + linear2: Linear {d_input: 512, d_output: 10, bias: true, params: 5130} + activation: Relu + params: 531178 +} +``` + +
+🦀 References + +In the previous example, the `init()` method signature uses `&` to indicate that the parameter types +are references: `&self`, a reference to the current receiver (`ModelConfig`), and +`device: &B::Device`, a reference to the backend device. + +```rust, ignore +pub fn init(&self, device: &B::Device) -> Model { + Model { + // ... + } +} +``` + +References in Rust allow us to point to a resource to access its data without owning it. The idea of +ownership is quite core to Rust and is worth +[reading up on](https://doc.rust-lang.org/book/ch04-00-understanding-ownership.html). + +In a language like C, memory management is explicit and up to the programmer, which means it is easy +to make mistakes. In a language like Java or Python, memory management is automatic with the help of +a garbage collector. This is very safe and straightforward, but also incurs a runtime cost. + +In Rust, memory management is rather unique. Aside from simple types that implement +[`Copy`](https://doc.rust-lang.org/std/marker/trait.Copy.html) (e.g., +[primitives](https://doc.rust-lang.org/rust-by-example/primitives.html) like integers, floats, +booleans and `char`), every value is _owned_ by some variable called the _owner_. Ownership can be +transferred from one variable to another and sometimes a value can be _borrowed_. Once the _owner_ +variable goes out of scope, the value is _dropped_, which means that any memory it allocated can be +freed safely. + +Because the method does not own the `self` and `device` variables, the values the references point +to will not be dropped when the reference stops being used (i.e., the scope of the method). + +For more information on references and borrowing, be sure to read the +[corresponding chapter](https://doc.rust-lang.org/book/ch04-02-references-and-borrowing.html) in the +Rust Book. + +

+ +When creating a custom neural network module, it is often a good idea to create a config alongside +the model struct. This allows you to define default values for your network, thanks to the `Config` +attribute. The benefit of this attribute is that it makes the configuration serializable, enabling +you to painlessly save your model hyperparameters, enhancing your experimentation process. Note that +a constructor will automatically be generated for your configuration, which will take in as input +values the parameters which do not have default values: +`let config = ModelConfig::new(num_classes, hidden_size);`. The default values can be overridden +easily with builder-like methods: (e.g `config.with_dropout(0.2);`) + +The first implementation block is related to the initialization method. As we can see, all fields +are set using the configuration of the corresponding neural network's underlying module. In this +specific case, we have chosen to expand the tensor channels from 1 to 8 with the first layer, then +from 8 to 16 with the second layer, using a kernel size of 3 on all dimensions. We also use the +adaptive average pooling module to reduce the dimensionality of the images to an 8 by 8 matrix, +which we will flatten in the forward pass to have a 1024 (16 * 8 * 8) resulting tensor. + +Now let's see how the forward pass is defined. + +```rust , ignore +# use burn::{ +# nn::{ +# conv::{Conv2d, Conv2dConfig}, +# pool::{AdaptiveAvgPool2d, AdaptiveAvgPool2dConfig}, +# Dropout, DropoutConfig, Linear, LinearConfig, Relu, +# }, +# prelude::*, +# }; +# +# #[derive(Module, Debug)] +# pub struct Model { +# conv1: Conv2d, +# conv2: Conv2d, +# pool: AdaptiveAvgPool2d, +# dropout: Dropout, +# linear1: Linear, +# linear2: Linear, +# activation: Relu, +# } +# +# #[derive(Config, Debug)] +# pub struct ModelConfig { +# num_classes: usize, +# hidden_size: usize, +# #[config(default = "0.5")] +# dropout: f64, +# } +# +# impl ModelConfig { +# /// Returns the initialized model. +# pub fn init(&self, device: &B::Device) -> Model { +# Model { +# conv1: Conv2dConfig::new([1, 8], [3, 3]).init(device), +# conv2: Conv2dConfig::new([8, 16], [3, 3]).init(device), +# pool: AdaptiveAvgPool2dConfig::new([8, 8]).init(), +# activation: Relu::new(), +# linear1: LinearConfig::new(16 * 8 * 8, self.hidden_size).init(device), +# linear2: LinearConfig::new(self.hidden_size, self.num_classes).init(device), +# dropout: DropoutConfig::new(self.dropout).init(), +# } +# } +# } +# +impl Model { + /// # Shapes + /// - Images [batch_size, height, width] + /// - Output [batch_size, num_classes] + pub fn forward(&self, images: Tensor) -> Tensor { + let [batch_size, height, width] = images.dims(); + + // Create a channel at the second dimension. + let x = images.reshape([batch_size, 1, height, width]); + + + let x = self.conv1.forward(x); // [batch_size, 8, _, _] + let x = self.dropout.forward(x); + let x = self.conv2.forward(x); // [batch_size, 16, _, _] + let x = self.dropout.forward(x); + let x = self.activation.forward(x); + + let x = self.pool.forward(x); // [batch_size, 16, 8, 8] + let x = x.reshape([batch_size, 16 * 8 * 8]); + let x = self.linear1.forward(x); + let x = self.dropout.forward(x); + let x = self.activation.forward(x); + + self.linear2.forward(x) // [batch_size, num_classes] + } +} +``` + +For former PyTorch users, this might feel very intuitive, as each module is directly incorporated +into the code using an eager API. Note that no abstraction is imposed for the forward method. You +are free to define multiple forward functions with the names of your liking. Most of the neural +network modules already built with Burn use the `forward` nomenclature, simply because it is +standard in the field. + +Similar to neural network modules, the [`Tensor`](../building-blocks/tensor.md) struct given as a +parameter also takes the Backend trait as a generic argument, alongside its dimensionality. Even if +it is not used in this specific example, it is possible to add the kind of the tensor as a third +generic argument. For example, a 3-dimensional Tensor of different data types(float, int, bool) +would be defined as following: + +```rust , ignore +Tensor // Float tensor (default) +Tensor // Float tensor (explicit) +Tensor // Int tensor +Tensor // Bool tensor +``` + +Note that the specific element type, such as `f16`, `f32` and the likes, will be defined later with +the backend. diff --git a/burn-book/src/basic-workflow/training-output.png b/burn-book/src/basic-workflow/training-output.png new file mode 100644 index 0000000..e7bb69e Binary files /dev/null and b/burn-book/src/basic-workflow/training-output.png differ diff --git a/burn-book/src/basic-workflow/training.md b/burn-book/src/basic-workflow/training.md new file mode 100644 index 0000000..7150496 --- /dev/null +++ b/burn-book/src/basic-workflow/training.md @@ -0,0 +1,297 @@ +# Training + +We are now ready to write the necessary code to train our model on the MNIST dataset. We shall +define the code for this training section in the file: `src/training.rs`. + +Instead of a simple tensor, the model should output an item that can be understood by the learner, a +struct whose responsibility is to apply an optimizer to the model. The output struct is used for all +metrics calculated during the training. Therefore it should include all the necessary information to +calculate any metric that you want for a task. + +Burn provides two basic output types: `ClassificationOutput` and `RegressionOutput`. They implement +the necessary trait to be used with metrics. It is possible to create your own item, but it is +beyond the scope of this guide. + +Since the MNIST task is a classification problem, we will use the `ClassificationOutput` type. + +```rust , ignore +# use crate::{ +# data::{MnistBatch, MnistBatcher}, +# model::{Model, ModelConfig}, +# }; +# use burn::{ +# data::{dataloader::DataLoaderBuilder, dataset::vision::MnistDataset}, +# nn::loss::CrossEntropyLossConfig, +# optim::AdamConfig, +# prelude::*, +# tensor::backend::AutodiffBackend, +# train::{ +# ClassificationOutput, Learner, SupervisedTraining, TrainOutput, TrainStep, InferenceStep, +# metric::{AccuracyMetric, LossMetric}, +# }, +# }; +# +impl Model { + pub fn forward_classification( + &self, + images: Tensor, + targets: Tensor, + ) -> ClassificationOutput { + let output = self.forward(images); + let loss = CrossEntropyLossConfig::new() + .init(&output.device()) + .forward(output.clone(), targets.clone()); + + ClassificationOutput::new(loss, output, targets) + } +} +``` + +As evident from the preceding code block, we employ the cross-entropy loss module for loss +calculation, without the inclusion of any padding token. We then return the classification output +containing the loss, the output tensor with all logits and the targets. + +Please take note that tensor operations receive owned tensors as input. For reusing a tensor +multiple times, you need to use the `clone()` function. There's no need to worry; this process won't +involve actual copying of the tensor data. Instead, it will simply indicate that the tensor is +employed in multiple instances, implying that certain operations won't be performed in place. In +summary, our API has been designed with owned tensors to optimize performance. + +Moving forward, we will proceed with the implementation of both the training and validation steps +for our model. + +```rust , ignore +# use crate::{ +# data::{MnistBatch, MnistBatcher}, +# model::{Model, ModelConfig}, +# }; +# use burn::{ +# data::{dataloader::DataLoaderBuilder, dataset::vision::MnistDataset}, +# nn::loss::CrossEntropyLossConfig, +# optim::AdamConfig, +# prelude::*, +# tensor::backend::AutodiffBackend, +# train::{ +# ClassificationOutput, InferenceStep, Learner, SupervisedTraining, TrainOutput, TrainStep, +# metric::{AccuracyMetric, LossMetric}, +# }, +# }; +# +# impl Model { +# pub fn forward_classification( +# &self, +# images: Tensor, +# targets: Tensor, +# ) -> ClassificationOutput { +# let output = self.forward(images); +# let loss = CrossEntropyLossConfig::new() +# .init(&output.device()) +# .forward(output.clone(), targets.clone()); +# +# ClassificationOutput::new(loss, output, targets) +# } +# } +impl TrainStep for Model { + type Input = MnistBatch; + type Output = ClassificationOutput; + + fn step(&self, batch: MnistBatch) -> TrainOutput> { + let item = self.forward_classification(batch.images, batch.targets); + + TrainOutput::new(self, item.loss.backward(), item) + } +} + +impl InferenceStep for Model { + type Input = MnistBatch; + type Output = ClassificationOutput; + + fn step(&self, batch: MnistBatch) -> ClassificationOutput { + self.forward_classification(batch.images, batch.targets) + } +} +``` + +Here we define the input and output types as generic arguments in the `TrainStep` and `InferenceStep`. +We will call them `MnistBatch` and `ClassificationOutput`. In the training step, the computation of +gradients is straightforward, necessitating a simple invocation of `backward()` on the loss. Note +that contrary to PyTorch, gradients are not stored alongside each tensor parameter, but are rather +returned by the backward pass, as such: `let gradients = loss.backward();`. The gradient of a +parameter can be obtained with the grad function: `let grad = tensor.grad(&gradients);`. Although it +is not necessary when using the learner struct and the optimizers, it can prove to be quite useful +when debugging or writing custom training loops. One of the differences between the training and the +validation steps is that the former requires the backend to implement `AutodiffBackend` and not just +`Backend`. Otherwise, the `backward` function is not available, as the backend does not support +autodiff. We will see later how to create a backend with autodiff support. + +
+🦀 Generic Type Constraints in Method Definitions + +Although generic data types, trait and trait bounds were already introduced in previous sections of +this guide, the previous code snippet might be a lot to take in at first. + +In the example above, we implement the `TrainStep` and `InferenceStep` trait for our `Model` struct, +which is generic over the `Backend` trait as has been covered before. These traits are provided by +`burn::train` and define a common `step` method that should be implemented for all structs. Since +the trait is generic over the input and output types, the trait implementation must specify the +concrete types used. This is where the additional type constraints appear +`, ClassificationOutput>`. As we saw previously, the concrete input type for the +batch is `MnistBatch`, and the output of the forward pass is `ClassificationOutput`. The `step` +method signature matches the concrete input and output types. + +For more details specific to constraints on generic types when defining methods, take a look at +[this section](https://doc.rust-lang.org/book/ch10-01-syntax.html#in-method-definitions) of the Rust +Book. + +

+ +Let us move on to establishing the practical training configuration. + +```rust , ignore +# use std::path::PathBuf; +# +# use crate::{ +# data::{MnistBatch, MnistBatcher}, +# model::{Model, ModelConfig}, +# }; +# use burn::{ +# data::{dataloader::DataLoaderBuilder, dataset::vision::MnistDataset}, +# nn::loss::CrossEntropyLossConfig, +# optim::AdamConfig, +# prelude::*, +# tensor::backend::AutodiffBackend, +# train::{ +# ClassificationOutput, InferenceStep, Learner, SupervisedTraining, TrainOutput, TrainStep, +# metric::{AccuracyMetric, LossMetric}, +# }, +# }; +# +# impl Model { +# pub fn forward_classification( +# &self, +# images: Tensor, +# targets: Tensor, +# ) -> ClassificationOutput { +# let output = self.forward(images); +# let loss = CrossEntropyLossConfig::new() +# .init(&output.device()) +# .forward(output.clone(), targets.clone()); +# +# ClassificationOutput::new(loss, output, targets) +# } +# } +# impl TrainStep for Model { +# type Input = MnistBatch; +# type Output = ClassificationOutput; +# +# fn step(&self, batch: MnistBatch) -> TrainOutput> { +# let item = self.forward_classification(batch.images, batch.targets); +# +# TrainOutput::new(self, item.loss.backward(), item) +# } +# } +# +# impl InferenceStep for Model { +# type Input = MnistBatch; +# type Output = ClassificationOutput; +# +# fn step(&self, batch: MnistBatch) -> ClassificationOutput { +# self.forward_classification(batch.images, batch.targets) +# } +# } +# +#[derive(Config, Debug)] +pub struct TrainingConfig { + pub model: ModelConfig, + pub optimizer: AdamConfig, + #[config(default = 10)] + pub num_epochs: usize, + #[config(default = 64)] + pub batch_size: usize, + #[config(default = 4)] + pub num_workers: usize, + #[config(default = 42)] + pub seed: u64, + #[config(default = 1.0e-4)] + pub learning_rate: f64, +} + +fn create_artifact_dir(artifact_dir: &str) { + std::fs::remove_file(PathBuf::from(artifact_dir).join("experiment.log")).ok(); + std::fs::create_dir_all(artifact_dir).ok(); +} + +pub fn train(artifact_dir: &str, config: TrainingConfig, device: B::Device) { + create_artifact_dir(artifact_dir); + config + .save(format!("{artifact_dir}/config.json")) + .expect("Config should be saved successfully"); + + B::seed(&device, config.seed); + + let batcher = MnistBatcher::default(); + + let dataloader_train = DataLoaderBuilder::new(batcher.clone()) + .batch_size(config.batch_size) + .shuffle(config.seed) + .num_workers(config.num_workers) + .build(MnistDataset::train()); + + let dataloader_test = DataLoaderBuilder::new(batcher) + .batch_size(config.batch_size) + .shuffle(config.seed) + .num_workers(config.num_workers) + .build(MnistDataset::test()); + + let training = SupervisedTraining::new(artifact_dir, dataloader_train, dataloader_test) + .metrics((AccuracyMetric::new(), LossMetric::new())) + .with_checkpointer() + .num_epochs(config.num_epochs) + .summary(); + + let model = config.model.init::(&device); + let result = training.launch(Learner::new( + model, + config.optimizer.init(), + config.learning_rate, + )); + + result + .model + .into_record() + .save(format!("{artifact_dir}/model")) + .expect("Trained model should be saved successfully"); +} +``` + +It is a good practice to use the `Config` derive to create the experiment configuration. In the +`train` function, the first thing we are doing is making sure the `artifact_dir` exists, using the +standard rust library for file manipulation. All checkpoints, logging and metrics will be stored +under this directory. We initialize the dataloaders using the previously created batcher. Since no +automatic differentiation is needed during the validation phase, the `training.launch(...)` method +defines the necessary backend bounds on the data loader for `B::InnerBackend` (see +[Backend](./backend.md)). The autodiff capabilities are available through a type system, making it +nearly impossible to forget to deactivate gradient calculation. + +Next, we create a supervised training runner with the dataloaders for training and validation and +we register the accuracy and loss metric on both training and validation steps. We also enable +checkpointing with `with_checkpointer()`, which periodically saves the model, optimizer, and learning +rate scheduler state to burnpack files under the experiment directory so training can be resumed. + +For the sake of simplicity in this example, we employ the test set as the validation +set; however, we do not recommend this practice for actual usage. + +We create the learner containing the model, the optimizer and the learning rate. Notably, the third +argument of the learner's `new` function should actually be a learning rate _scheduler_. When provided with a +float as in our example, it is automatically transformed into a _constant_ learning rate scheduler. +The learning rate is not part of the optimizer config as it is often done in other frameworks, but +rather passed as a parameter when executing the optimizer step. This avoids having to mutate the +state of the optimizer and is therefore more functional. It makes no difference when using the +learner struct, but it will be an essential nuance to grasp if you implement your own training loop. + +Once the learner and supervised training instance are created, we can call `training.launch` and provide the learner. + +Finally, the trained model is returned by the `launch` method. The trained weights are then saved by +taking a record with `into_record()` and calling `save`, which writes a burnpack (`.bpk`) file. A +record holds plain tensor data, so any backend, regardless of precision, can load recorded weights of +any kind. diff --git a/burn-book/src/building-blocks/README.md b/burn-book/src/building-blocks/README.md new file mode 100644 index 0000000..660c9bb --- /dev/null +++ b/burn-book/src/building-blocks/README.md @@ -0,0 +1,7 @@ +# Building Blocks + +In this section, we'll guide you through the core elements that make up Burn. We'll walk you through +the key components that serve as the building blocks of the framework and your future projects. + +As you explore Burn, you might notice that we occasionally draw comparisons to PyTorch. We believe +it can provide a smoother learning curve and help you grasp the nuances more effectively. diff --git a/burn-book/src/building-blocks/autodiff.md b/burn-book/src/building-blocks/autodiff.md new file mode 100644 index 0000000..b9ed0df --- /dev/null +++ b/burn-book/src/building-blocks/autodiff.md @@ -0,0 +1,90 @@ +# Autodiff + +Burn's tensor also supports auto-differentiation, which is an essential part of any deep learning +framework. We introduced the `Backend` trait in the [previous section](./backend.md), but Burn also +has another trait for autodiff: `AutodiffBackend`. + +However, not all tensors support auto-differentiation; you need a backend that implements both the +`Backend` and `AutodiffBackend` traits. Fortunately, you can add auto-differentiation capabilities to any +backend using a backend decorator: `type MyAutodiffBackend = Autodiff`. This +decorator implements both the `AutodiffBackend` and `Backend` traits by maintaining a dynamic +computational graph and utilizing the inner backend to execute tensor operations. + +The `AutodiffBackend` trait adds new operations on float tensors that can't be called otherwise. It also +provides a new associated type, `B::Gradients`, where each calculated gradient resides. + +```rust, ignore +fn calculate_gradients(tensor: Tensor) -> B::Gradients { + let mut gradients = tensor.clone().backward(); + + let tensor_grad = tensor.grad(&gradients); // get + let tensor_grad = tensor.grad_remove(&mut gradients); // pop + + gradients +} +``` + +Note that some functions will always be available even if the backend doesn't implement the +`AutodiffBackend` trait. In such cases, those functions will do nothing. + +| Burn API | PyTorch Equivalent | +| --------------------------------------- | ----------------------------- | +| `tensor.detach()` | `tensor.detach()` | +| `tensor.require_grad()` | `tensor.requires_grad()` | +| `tensor.is_require_grad()` | `tensor.requires_grad` | +| `tensor.set_require_grad(require_grad)` | `tensor.requires_grad(False)` | + +However, you're unlikely to make any mistakes since you can't call `backward` on a tensor that is on +a backend that doesn't implement `AutodiffBackend`. Additionally, you can't retrieve the gradient of a +tensor without an autodiff backend. + +## Difference with PyTorch + +The way Burn handles gradients is different from PyTorch. First, when calling `backward`, each +parameter doesn't have its `grad` field updated. Instead, the backward pass returns all the +calculated gradients in a container. This approach offers numerous benefits, such as the ability to +easily send gradients to other threads. + +You can also retrieve the gradient for a specific parameter using the `grad` method on a tensor. +Since this method takes the gradients as input, it's hard to forget to call `backward` beforehand. +Note that sometimes, using `grad_remove` can improve performance by allowing inplace operations. + +In PyTorch, when you don't need gradients for inference or validation, you typically need to scope +your code using a block. + +```python +# Inference mode +torch.inference(): + # your code + ... + +# Or no grad +torch.no_grad(): + # your code + ... +``` + +With Burn, you don't need to wrap the backend with the `Autodiff` for inference, and you +can call `inner()` to obtain the inner tensor, which is useful for validation. + +```rust, ignore +/// Use `B: AutodiffBackend` +fn example_validation(tensor: Tensor) { + let inner_tensor: Tensor = tensor.inner(); + let _ = inner_tensor + 5; +} + +/// Use `B: Backend` +fn example_inference(tensor: Tensor) { + let _ = tensor + 5; + ... +} +``` + +**Gradients with Optimizers** + +We've seen how gradients can be used with tensors, but the process is a bit different when working +with optimizers from `burn-core`. To work with the `Module` trait, a translation step is required to +link tensor parameters with their gradients. This step is necessary to easily support gradient +accumulation and training on multiple devices, where each module can be forked and run on different +devices in parallel. We'll explore deeper into this topic in the [Module](./module.md) section. diff --git a/burn-book/src/building-blocks/backend.md b/burn-book/src/building-blocks/backend.md new file mode 100644 index 0000000..6382586 --- /dev/null +++ b/burn-book/src/building-blocks/backend.md @@ -0,0 +1,14 @@ +# Backend + +Nearly everything in Burn is based on the `Backend` trait, which enables you to run tensor +operations using different implementations without having to modify your code. While a backend may +not necessarily have autodiff capabilities, the `AutodiffBackend` trait specifies when autodiff is +needed. This trait not only abstracts operations but also tensor, device, and element types, +providing each backend the flexibility they need. It's worth noting that the trait assumes eager +mode since burn fully supports dynamic graphs. However, we may create another API to assist with +integrating graph-based backends, without requiring any changes to the user's code. + +Users are not expected to directly use the backend trait methods, as it is primarily designed with +backend developers in mind rather than Burn users. Therefore, most Burn userland APIs are generic +across backends. This approach helps users discover the API more organically with proper +autocomplete and documentation. diff --git a/burn-book/src/building-blocks/config.md b/burn-book/src/building-blocks/config.md new file mode 100644 index 0000000..0ae3fb4 --- /dev/null +++ b/burn-book/src/building-blocks/config.md @@ -0,0 +1,64 @@ +# Config + +When writing scientific code, you normally have a lot of values that are set, and Deep Learning is +no exception. Python has the possibility to define default parameters for functions, which helps +improve the developer experience. However, this has the downside of potentially breaking your code +when upgrading to a new version, as the default values might change without your knowledge, making +debugging very challenging. + +With that in mind, we came up with the Config system. It's a simple Rust derive that you can apply +to your types, allowing you to define default values with ease. Additionally, all configs can be +serialized, reducing potential bugs when upgrading versions and improving reproducibility. + +```rust , ignore +use burn::config::Config; + +#[derive(Config)] +pub struct MyModuleConfig { + d_model: usize, + d_ff: usize, + #[config(default = 0.1)] + dropout: f64, +} +``` + +The derive also adds useful `with_` methods for every attribute of your config, similar to a builder +pattern, along with a `save` method. + +```rust, ignore +fn main() { + let config = MyModuleConfig::new(512, 2048); + println!("{}", config.d_model); // 512 + println!("{}", config.d_ff); // 2048 + println!("{}", config.dropout); // 0.1 + let config = config.with_dropout(0.2); + println!("{}", config.dropout); // 0.2 + + config.save("config.json").unwrap(); +} +``` + +## Good practices + +By using the config type it is easy to create new module instances. The initialization method should +be implemented on the config type with the device as argument. + +```rust, ignore +impl MyModuleConfig { + /// Create a module on the given device. + pub fn init(&self, device: &B::Device) -> MyModule { + MyModule { + linear: LinearConfig::new(self.d_model, self.d_ff).init(device), + dropout: DropoutConfig::new(self.dropout).init(), + } + } +} +``` + +Then we could add this line to the above `main`: + +```rust, ignore +use burn::backend::Wgpu; +let device = Default::default(); +let my_module = config.init::(&device); +``` diff --git a/burn-book/src/building-blocks/dataset.md b/burn-book/src/building-blocks/dataset.md new file mode 100644 index 0000000..2b5e6bb --- /dev/null +++ b/burn-book/src/building-blocks/dataset.md @@ -0,0 +1,531 @@ +# Dataset + +At its core, a dataset is a collection of data typically related to a specific analysis or +processing task. The data modality can vary depending on the task, but most datasets primarily +consist of images, texts, audio or videos. + +This data source represents an integral part of machine learning to successfully train a model. +Thus, it is essential to provide a convenient and performant API to handle your data. Since this +process varies wildly from one problem to another, it is defined as a trait that should be +implemented on your type. The dataset trait is quite similar to the dataset abstract class in +PyTorch: + +```rust, ignore +pub trait Dataset: Send + Sync { + fn get(&self, index: usize) -> Option; + fn len(&self) -> usize; +} +``` + +The dataset trait assumes a fixed-length set of items that can be randomly accessed in constant +time. This is a major difference from datasets that use Apache Arrow underneath to improve streaming +performance. Datasets in Burn don't assume _how_ they are going to be accessed; it's just a +collection of items. + +However, you can compose multiple dataset transformations to lazily obtain what you want with zero +pre-processing, so that your training can start instantly! + +## Transformation + +Transformations in Burn are all lazy and modify one or multiple input datasets. The goal of these +transformations is to provide you with the necessary tools so that you can model complex data +distributions. + +| Transformation | Description | +| ------------------ | ------------------------------------------------------------------------------------------------------------------------ | +| `SamplerDataset` | Samples items from a dataset. This is a convenient way to model a dataset as a probability distribution of a fixed size. | +| `SelectionDataset` | Selects a subset of items by index from a dataset. Can be randomly shuffled; can be re-shuffled. | +| `ShuffledDataset` | Shuffles a wrapped dataset; This is a thin wrapper around `SelectionDataset`. | +| `PartialDataset` | Returns a view of the input dataset with a specified range. | +| `MapperDataset` | Computes a transformation lazily on the input dataset. | +| `ComposedDataset` | Composes multiple datasets together to create a larger one without copying any data. | +| `WindowsDataset` | Dataset designed to work with overlapping windows of data extracted from an input dataset. | + +Let us look at the basic usages of each dataset transform and how they can be composed together. +These transforms are lazy by default except when specified, reducing the need for unnecessary +intermediate allocations and improving performance. The full documentation of each transform can be +found at the [API reference](https://burn.dev/docs/burn/data/dataset/transform/index.html). + +- **SamplerDataset**: This transform can be used to sample items from a dataset with (default) or + without replacement. Transform is initialized with a sampling size which can be bigger or smaller + than the input dataset size. This is particularly useful in cases where we want to checkpoint + larger datasets more often during training and smaller datasets less often as the size of an epoch + is now controlled by the sampling size. Sample usage: + +```rust, ignore +type DbPedia = SqliteDataset; +let dataset: DbPedia = HuggingfaceDatasetLoader::new("fancyzhx/dbpedia_14") + .dataset("train"). + .unwrap(); + +let dataset = SamplerDataset::new(dataset, 10000); +``` + +- **SelectionDataset**: This transform can be used to select a subset of items from a dataset by + index. It can be initialized with a list of indices to select from the input dataset. This is + particularly useful when you want to create a smaller dataset from a larger one, for example, to + create a validation set from a training set. + + The `SelectionDataset` can also be initialized with a random seed to shuffle the indices before + selection. This is useful when you want to randomly select a subset of items from the dataset. + + Base dataset items may be included more than once in the selection. + +```rust, ignore +let explicit = SelectionDataset::from_indices_checked(dataset.clone(), vec![0, 1, 2, 0]); + +let shuffled = SelectionDataset::new_shuffled(dataset.clone(), &mut rng); +let shuffled = SelectionDataset::new_shuffled(dataset.clone(), 42); + +let mut mutable = SelectionDataset::new_select_all(dataset.clone(), vec![0, 1, 2, 0]); +mutable.shuffle(42); +mutable.shuffle(&mut rng); +``` + +- **ShuffledDataset**: This transform can be used to shuffle the items of a dataset. Particularly + useful before splitting the raw dataset into train/test splits. Can be initialized with a seed to + ensure reproducibility. + + The `ShuffledDataset` is a thin wrapper around the `SelectionDataset`. + +```rust, ignore +let dataset = ShuffledDataset::new(dataset, &mut rng); +let dataset = ShuffledDataset::new(dataset, 42); +``` + +- **PartialDataset**: This transform is useful to return a view of the dataset with specified start + and end indices. Used to create train/val/test splits. In the example below, we show how to chain + ShuffledDataset and PartialDataset to create splits. + +```rust, ignore +// define chained dataset type here for brevity +type PartialData = PartialDataset>; +let len = dataset.len(); +let split = "train"; // or "val"/"test" + +let data_split = match split { + "train" => PartialData::new(dataset, 0, len * 8 / 10), // Get first 80% dataset + "test" => PartialData::new(dataset, len * 8 / 10, len), // Take remaining 20% + _ => panic!("Invalid split type"), // Handle unexpected split types +}; +``` + +- **MapperDataset**: This transform is useful to apply a transformation on each of the items of a + dataset. Particularly useful for normalization of image data when channel means are known. + +- **ComposedDataset**: This transform is useful to compose multiple datasets downloaded from + multiple sources (say different HuggingfaceDatasetLoader sources) into a single bigger dataset + which can be sampled from one source. + +- **WindowsDataset**: This transform is useful to create overlapping windows of a dataset. + Particularly useful for sequential Time series Data, for example when working with an LSTM. + +## Storage + +There are multiple dataset storage options available for you to choose from. The choice of the +dataset to use should be based on the dataset's size as well as its intended purpose. + +| Storage | Description | +| ------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------- | +| `InMemDataset` | In-memory dataset that uses a vector to store items. Well-suited for smaller datasets. | +| `SqliteDataset` | Dataset that uses [SQLite](https://www.sqlite.org/) to index items that can be saved in a simple SQL database file. Well-suited for larger datasets. | +| `DataframeDataset` | Dataset that uses [Polars](https://www.pola.rs/) dataframe to store and manage data. Well-suited for efficient data manipulation and analysis. | + +## Sources + +For now, there are only a couple of dataset sources available with Burn, but more to come! + +### Hugging Face + +You can easily import any Hugging Face dataset with Burn. We use SQLite as the storage to avoid +downloading the model each time or starting a Python process. You need to know the format of each +item in the dataset beforehand. Here's an example with the +[dbpedia dataset](https://huggingface.co/datasets/fancyzhx/dbpedia_14). + +```rust, ignore +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +pub struct DbPediaItem { + pub title: String, + pub content: String, + pub label: usize, +} + +fn main() { + let dataset: SqliteDataset = HuggingfaceDatasetLoader::new("fancyzhx/dbpedia_14") + .dataset("train") // The training split. + .unwrap(); +} +``` + +We see that items must derive `serde::Serialize`, `serde::Deserialize`, `Clone`, and `Debug`, but +those are the only requirements. + +
+ +The `HuggingfaceDatasetLoader` relies on the +[`datasets` library by HuggingFace](https://huggingface.co/docs/datasets/index) to download +datasets. This is a Python library, so you must have an existing Python installation to use this +loader. + +
+ +### Images + +`ImageFolderDataset` is a generic vision dataset used to load images from disk. It is currently +available for multi-class and multi-label classification tasks as well as semantic segmentation and +object detection tasks. + +```rust, ignore +// Create an image classification dataset from the root folder, +// where images for each class are stored in their respective folder. +// +// For example: +// root/dog/dog1.png +// root/dog/dog2.png +// ... +// root/cat/cat1.png +let dataset = ImageFolderDataset::new_classification("path/to/dataset/root").unwrap(); +``` + +```rust, ignore +// Create a multi-label image classification dataset from a list of items, +// where each item is a tuple `(image path, labels)`, and a list of classes +// in the dataset. +// +// For example: +let items = vec![ + ("root/dog/dog1.png", vec!["animal".to_string(), "dog".to_string()]), + ("root/cat/cat1.png", vec!["animal".to_string(), "cat".to_string()]), +]; +let dataset = ImageFolderDataset::new_multilabel_classification_with_items( + items, + &["animal", "cat", "dog"], +) +.unwrap(); +``` + +```rust, ignore +// Create a segmentation mask dataset from a list of items, where each +// item is a tuple `(image path, mask path)` and a list of classes +// corresponding to the integer values in the mask. +let items = vec![ + ( + "path/to/images/image0.png", + "path/to/annotations/mask0.png", + ), + ( + "path/to/images/image1.png", + "path/to/annotations/mask1.png", + ), + ( + "path/to/images/image2.png", + "path/to/annotations/mask2.png", + ), +]; +let dataset = ImageFolderDataset::new_segmentation_with_items( + items, + &[ + "cat", // 0 + "dog", // 1 + "background", // 2 + ], +) +.unwrap(); +``` + +```rust, ignore +// Create an object detection dataset from a COCO dataset. Currently only +// the import of object detection data (bounding boxes) is supported. +// +// COCO offers separate annotation and image archives for training and +// validation, paths to the unpacked files need to be passed as parameters: + +let dataset = ImageFolderDataset::new_coco_detection( + "/path/to/coco/instances_train2017.json", + "/path/to/coco/images/train2017" +) +.unwrap(); + +``` + +### Comma-Separated Values (CSV) + +Loading records from a simple CSV file in-memory is simple with the `InMemDataset`: + +```rust, ignore +// Build dataset from csv with tab ('\t') delimiter. +// The reader can be configured for your particular file. +let mut rdr = csv::ReaderBuilder::new(); +let rdr = rdr.delimiter(b'\t'); + +let dataset = InMemDataset::from_csv("path/to/csv", rdr).unwrap(); +``` + +Note that this requires the `csv` crate. + +**What about streaming datasets?** + +There is no streaming dataset API with Burn, and this is by design! The learner struct will iterate +multiple times over the dataset and only checkpoint when done. You can consider the length of the +dataset as the number of iterations before performing checkpointing and running the validation. +There is nothing stopping you from returning different items even when called with the same `index` +multiple times. + +## How Is The Dataset Used? + +During training, the dataset is used to access the data samples and, for most use cases in +supervised learning, their corresponding ground-truth labels. Remember that the `Dataset` trait +implementation is responsible to retrieve the data from its source, usually some sort of data +storage. At this point, the dataset could be naively iterated over to provide the model a single +sample to process at a time, but this is not very efficient. + +Instead, we collect multiple samples that the model can process as a _batch_ to fully leverage +modern hardware (e.g., GPUs - which have impressive parallel processing capabilities). Since each +data sample in the dataset can be collected independently, the data loading is typically done in +parallel to further speed things up. In this case, we parallelize the data loading using a +multi-threaded `BatchDataLoader` to obtain a sequence of items from the `Dataset` implementation. +Finally, the sequence of items is combined into a batched tensor that can be used as input to a +model with the `Batcher` trait implementation. Other tensor operations can be performed during this +step to prepare the batch data, as is done [in the basic workflow guide](../basic-workflow/data.md). +The process is illustrated in the figure below for the MNIST dataset. + +Burn Data Loading Pipeline + +Although we have conveniently implemented the +[`MnistDataset`](https://github.com/tracel-ai/burn/blob/main/crates/burn-dataset/src/vision/mnist.rs) +used in the guide, we'll go over its implementation to demonstrate how the `Dataset` and `Batcher` +traits are used. + +The [MNIST dataset](http://yann.lecun.com/exdb/mnist/) of handwritten digits has a training set of +60,000 examples and a test set of 10,000 examples. A single item in the dataset is represented by a +\\(28 \times 28\\) pixels black-and-white image (stored as raw bytes) with its corresponding label +(a digit between \\(0\\) and \\(9\\)). This is defined by the `MnistItemRaw` struct. + +```rust, ignore +# #[derive(Deserialize, Debug, Clone)] +struct MnistItemRaw { + pub image_bytes: Vec, + pub label: u8, +} +``` + +With single-channel images of such low resolution, the entire training and test sets can be loaded +in memory at once. Therefore, we leverage the already existing `InMemDataset` to retrieve the raw +images and labels data. At this point, the image data is still just a bunch of bytes, but we want to +retrieve the _structured_ image data in its intended form. For that, we can define a `MapperDataset` +that transforms the raw image bytes to a 2D array image (which we convert to float while we're at +it). + +```rust, ignore +const WIDTH: usize = 28; +const HEIGHT: usize = 28; + +# /// MNIST item. +# #[derive(Deserialize, Serialize, Debug, Clone)] +pub struct MnistItem { + /// Image as a 2D array of floats. + pub image: [[f32; WIDTH]; HEIGHT], + + /// Label of the image. + pub label: u8, +} + +struct BytesToImage; + +impl Mapper for BytesToImage { + /// Convert a raw MNIST item (image bytes) to a MNIST item (2D array image). + fn map(&self, item: &MnistItemRaw) -> MnistItem { + // Ensure the image dimensions are correct. + debug_assert_eq!(item.image_bytes.len(), WIDTH * HEIGHT); + + // Convert the image to a 2D array of floats. + let mut image_array = [[0f32; WIDTH]; HEIGHT]; + for (i, pixel) in item.image_bytes.iter().enumerate() { + let x = i % WIDTH; + let y = i / HEIGHT; + image_array[y][x] = *pixel as f32; + } + + MnistItem { + image: image_array, + label: item.label, + } + } +} + +type MappedDataset = MapperDataset, BytesToImage, MnistItemRaw>; + +# /// The MNIST dataset consists of 70,000 28x28 black-and-white images in 10 classes (one for each digits), with 7,000 +# /// images per class. There are 60,000 training images and 10,000 test images. +# /// +# /// The data is downloaded from the web from the [CVDF mirror](https://github.com/cvdfoundation/mnist). +pub struct MnistDataset { + dataset: MappedDataset, +} +``` + +To construct the `MnistDataset`, the data source must be parsed into the expected `MappedDataset` +type. Since both the train and test sets use the same file format, we can separate the functionality +to load the `train()` and `test()` dataset. + +```rust, ignore + +impl MnistDataset { + /// Creates a new train dataset. + pub fn train() -> Self { + Self::new("train") + } + + /// Creates a new test dataset. + pub fn test() -> Self { + Self::new("test") + } + + fn new(split: &str) -> Self { + // Download dataset + let root = MnistDataset::download(split); + + // Parse data as vector of images bytes and vector of labels + let images: Vec> = MnistDataset::read_images(&root, split); + let labels: Vec = MnistDataset::read_labels(&root, split); + + // Collect as vector of MnistItemRaw + let items: Vec<_> = images + .into_iter() + .zip(labels) + .map(|(image_bytes, label)| MnistItemRaw { image_bytes, label }) + .collect(); + + // Create the MapperDataset for InMemDataset to transform + // items (MnistItemRaw -> MnistItem) + let dataset = InMemDataset::new(items); + let dataset = MapperDataset::new(dataset, BytesToImage); + + Self { dataset } + } + +# /// Download the MNIST dataset files from the web. +# /// Panics if the download cannot be completed or the content of the file cannot be written to disk. +# fn download(split: &str) -> PathBuf { +# // Dataset files are stored in the burn-dataset cache directory +# let cache_dir = dirs::cache_dir() +# .expect("Could not get cache directory") +# .join("burn-dataset"); +# let split_dir = cache_dir.join("mnist").join(split); +# +# if !split_dir.exists() { +# create_dir_all(&split_dir).expect("Failed to create base directory"); +# } +# +# // Download split files +# match split { +# "train" => { +# MnistDataset::download_file(TRAIN_IMAGES, &split_dir); +# MnistDataset::download_file(TRAIN_LABELS, &split_dir); +# } +# "test" => { +# MnistDataset::download_file(TEST_IMAGES, &split_dir); +# MnistDataset::download_file(TEST_LABELS, &split_dir); +# } +# _ => panic!("Invalid split specified {}", split), +# }; +# +# split_dir +# } +# +# /// Download a file from the MNIST dataset URL to the destination directory. +# /// File download progress is reported with the help of a [progress bar](indicatif). +# fn download_file>(name: &str, dest_dir: &P) -> PathBuf { +# // Output file name +# let file_name = dest_dir.as_ref().join(name); +# +# if !file_name.exists() { +# // Download gzip file +# let bytes = download_file_as_bytes(&format!("{URL}{name}.gz"), name); +# +# // Create file to write the downloaded content to +# let mut output_file = File::create(&file_name).unwrap(); +# +# // Decode gzip file content and write to disk +# let mut gz_buffer = GzDecoder::new(&bytes[..]); +# std::io::copy(&mut gz_buffer, &mut output_file).unwrap(); +# } +# +# file_name +# } +# +# /// Read images at the provided path for the specified split. +# /// Each image is a vector of bytes. +# fn read_images>(root: &P, split: &str) -> Vec> { +# let file_name = if split == "train" { +# TRAIN_IMAGES +# } else { +# TEST_IMAGES +# }; +# let file_name = root.as_ref().join(file_name); +# +# // Read number of images from 16-byte header metadata +# let mut f = File::open(file_name).unwrap(); +# let mut buf = [0u8; 4]; +# let _ = f.seek(SeekFrom::Start(4)).unwrap(); +# f.read_exact(&mut buf) +# .expect("Should be able to read image file header"); +# let size = u32::from_be_bytes(buf); +# +# let mut buf_images: Vec = vec![0u8; WIDTH * HEIGHT * (size as usize)]; +# let _ = f.seek(SeekFrom::Start(16)).unwrap(); +# f.read_exact(&mut buf_images) +# .expect("Should be able to read image file header"); +# +# buf_images +# .chunks(WIDTH * HEIGHT) +# .map(|chunk| chunk.to_vec()) +# .collect() +# } +# +# /// Read labels at the provided path for the specified split. +# fn read_labels>(root: &P, split: &str) -> Vec { +# let file_name = if split == "train" { +# TRAIN_LABELS +# } else { +# TEST_LABELS +# }; +# let file_name = root.as_ref().join(file_name); +# +# // Read number of labels from 8-byte header metadata +# let mut f = File::open(file_name).unwrap(); +# let mut buf = [0u8; 4]; +# let _ = f.seek(SeekFrom::Start(4)).unwrap(); +# f.read_exact(&mut buf) +# .expect("Should be able to read label file header"); +# let size = u32::from_be_bytes(buf); +# +# let mut buf_labels: Vec = vec![0u8; size as usize]; +# let _ = f.seek(SeekFrom::Start(8)).unwrap(); +# f.read_exact(&mut buf_labels) +# .expect("Should be able to read labels from file"); +# +# buf_labels +# } +} +``` + +Since the `MnistDataset` simply wraps a `MapperDataset` instance with `InMemDataset`, we can easily +implement the `Dataset` trait. + +```rust, ignore +impl Dataset for MnistDataset { + fn get(&self, index: usize) -> Option { + self.dataset.get(index) + } + + fn len(&self) -> usize { + self.dataset.len() + } +} +``` + +The only thing missing now is the `Batcher`, which we already went over +[in the basic workflow guide](../basic-workflow/data.md). The `Batcher` takes a list of `MnistItem` +retrieved by the dataloader as input and returns a batch of images as a 3D tensor along with their +targets. diff --git a/burn-book/src/building-blocks/dataset.png b/burn-book/src/building-blocks/dataset.png new file mode 100755 index 0000000..a66adb8 Binary files /dev/null and b/burn-book/src/building-blocks/dataset.png differ diff --git a/burn-book/src/building-blocks/learner.md b/burn-book/src/building-blocks/learner.md new file mode 100644 index 0000000..8a63604 --- /dev/null +++ b/burn-book/src/building-blocks/learner.md @@ -0,0 +1,112 @@ +# Learner + +The [burn-train](https://github.com/tracel-ai/burn/tree/main/crates/burn-train) crate encapsulates +multiple utilities for training deep learning models. The goal of the crate is to provide users with +a well-crafted and flexible training loop, so that projects do not have to write such components +from the ground up. Most of the interactions with `burn-train` will be with the `SupervisedTraining` +struct, briefly presented in the previous [training section](../basic-workflow/training.md). This +struct enables you to configure the training loop, offering support for registering metrics, +enabling logging, checkpointing states, using multiple devices, and so on. + +There are still some assumptions in the current provided APIs, which may make them inappropriate for +your learning requirements. Indeed, they assume your model will learn from a training dataset and be +validated against another dataset. This is the most common paradigm, allowing users to do both +supervised and unsupervised learning as well as fine-tuning. However, for more complex requirements, +creating a [custom training loop](../custom-training-loop.md) might be what you need. + +## Usage + +The `SupervisedLearning` struct must be created with the training and validation dataloaders. It provides you with numerous options when it comes to configurations. + +| Configuration | Description | +| ---------------------- | ------------------------------------------------------------------------------ | +| Training Metric | Register a training metric | +| Validation Metric | Register a validation metric | +| Training Metric Plot | Register a training metric with plotting (requires the metric to be numeric) | +| Validation Metric Plot | Register a validation metric with plotting (requires the metric to be numeric) | +| Metric Logger | Configure the metric loggers (default is saving them to files) | +| Renderer | Configure how to render metrics (default is CLI) | +| Grad Accumulation | Configure the number of steps before applying gradients | +| File Checkpointer | Configure how the model, optimizer and scheduler states are saved | +| Num Epochs | Set the number of epochs | +| Devices | Set the devices to be used | +| Checkpoint | Restart training from a checkpoint | +| Application logging | Configure the application logging installer (default is writing to `experiment.log`) | +| Training Strategy | Use a custom training strategy, allowing you to use your own training loop with all the capabilities of the `SupervisedTraining` struct | + +When the training is configured to your liking, you can then move forward to running the training. The +`launch` method requires a learner object providing: the model, the optimizer and the learning rate scheduler. Note +that the latter can be a simple float if you want it to be constant during training. + +The `launch` method will start the training and return the trained model once finished. + +Again, please refer to the [training section](../basic-workflow/training.md) for a relevant code +snippet. + +## Multiple optimizers + +It's common practice to set different learning rates, optimizer parameters, or use different optimizers entirely, for different parts +of a model. You can leverage Burn's `ParamGroup`s to mix and match optimizers and learning rate schedulers easily! + +```rust,ignore +let lr_scheduler_base = ComposedLrSchedulerConfig::new() + .cosine(CosineAnnealingLrSchedulerConfig::new(1.0, 2000)) + .linear(LinearLrSchedulerConfig::new(1e-8, 1.0, 2000)) + .linear(LinearLrSchedulerConfig::new(1e-2, 1e-6, 10000)); +let lr_scheduler = lr_scheduler_base.init().unwrap().with_group( + ParamGroup::from_predicate("conv"), + LinearLrSchedulerConfig::new(1e-6, 1e-3, 14000) + .build() + .unwrap(), +); + +let optimizer_base = AdamWConfig::new() + .with_cautious_weight_decay(true) + .with_weight_decay(5e-5); +let optim = optimizer_base.init().with_group( + ParamGroup::from_predicate("conv"), + SgdConfig::new().build(), + None, +); + +let result = training.launch(Learner::new( + model, + optim, + lr_scheduler, +)); +``` + +## Artifacts + +When creating a `SupervisedTraining` instance, all the collected data will be saved under the directory provided as +the argument to the `new` method. Here is an example of the data layout for a model checkpointed to +the burnpack format, with the accuracy and loss metrics registered: + +``` +├── experiment.log +├── checkpoint +│   ├── model-1.bpk +│   ├── optim-1.bpk +│   └── scheduler-1.bpk +│   ├── model-2.bpk +│   ├── optim-2.bpk +│   └── scheduler-2.bpk +├── train +│   ├── epoch-1 +│   │   ├── Accuracy.log +│   │   └── Loss.log +│   └── epoch-2 +│   ├── Accuracy.log +│   └── Loss.log +└── valid + ├── epoch-1 + │   ├── Accuracy.log + │   └── Loss.log + └── epoch-2 + ├── Accuracy.log + └── Loss.log +``` + +You can choose to save or synchronize that local directory with a remote file system, if desired. +The file checkpointer is capable of automatically deleting old checkpoints according to a specified +configuration. diff --git a/burn-book/src/building-blocks/metric.md b/burn-book/src/building-blocks/metric.md new file mode 100644 index 0000000..0876ec0 --- /dev/null +++ b/burn-book/src/building-blocks/metric.md @@ -0,0 +1,226 @@ +# Metric + +When working with the learner, you have the option to record metrics that will be monitored +throughout the training process. We currently offer a restricted range of metrics. + +| Metric | Description | +| ------------------- | ------------------------------------------------------------------------------------------- | +| Accuracy | Calculate the accuracy in percentage | +| TopKAccuracy | Calculate the top-k accuracy in percentage | +| Precision | Calculate precision in percentage | +| Recall | Calculate recall in percentage | +| FBetaScore | Calculate Fβ score in percentage | +| AUROC | Calculate the area under curve of ROC in percentage | +| AUC-PR | Calculate the area under the precision-recall curve (average precision) in percentage | +| Loss | Output the loss used for the backward pass | +| CharErrorRate (CER) | Calculate Character Error Rate in percentage | +| WordErrorRate (WER) | Calculate Word Error Rate in percentage | +| HammingScore | Calculate hamming score (also known as multi-label or label-based accuracy) in percentage | +| Perplexity | Calculate perplexity which is a measure of how well a probability model predicts samples | +| IterationSpeed | Tracks the training iteration speed, measuring how many iterations are completed per second | +| CPU Temperature | Fetch the temperature of CPUs | +| CPU Usage | Fetch the CPU utilization | +| CPU Memory Usage | Fetch the CPU RAM usage | +| Learning Rate | Fetch the current learning rate for each optimizer step | +| CUDA | Fetch general CUDA metrics such as utilization | + +| Vision Metric | Description | +| ------------- | ---------------------------------------------------------------------------------------------------- | +| A-FINE | Computes the Adaptive Fidelity-Naturalness Evaluator (A-FINE) full-reference perceptual quality metric built on CLIP ViT-B/32 features | +| Dice | Computes the Dice-Sorenson coefficient (DSC) for evaluating overlap between binary masks | +| DISTS | Computes the Deep Image Structure and Texture Similarity (DISTS) metric for image quality assessment | +| FID | Computes the Frechet Inception Distance (FID) for evaluating generative model quality | +| LPIPS | Computes the Learned Perceptual Image Patch Similarity (LPIPS) for image quality assessment | +| MS-SSIM | Computes the Multi-scale Structural Similarity index measure (MS-SSIM) for image quality assessment | +| PSNR | Computes the Peak Signal-to-Noise Ratio (PSNR) for image quality assessment | +| SSIM | Computes the Structural Similarity index measure (SSIM) for image quality assessment | + +## Using Metrics with the Learner + +In order to use a metric, the output of your training step must implement the `Adaptor` trait from +`burn-train::metric` for each metric's corresponding input type. The `Adaptor` trait simply converts +your output struct into the input type the metric expects. + +Burn provides four built-in output structs that cover common tasks. Each one already implements +`Adaptor` for a set of metrics, so in many cases you can use them directly without writing any +adaptor code yourself. + +- `ClassificationOutput`: + - Use case: Single-label classification + - Fields: `loss: Tensor`, `output: Tensor`, `targets: Tensor` + - Adapted metrics: Accuracy, TopKAccuracy, Perplexity, Precision\*, Recall\*, FBetaScore\*, AUROC\*, AUC-PR\*, Loss +- `MultiLabelClassificationOutput`: + - Use case: Multi-label classification + - Fields: `loss: Tensor`, `output: Tensor`, `targets: Tensor` + - Adapted metrics: HammingScore, Precision\*, Recall\*, FBetaScore\*, AUROC\*, AUC-PR\*, Loss +- `RegressionOutput`: + - Use case: Regression tasks + - Fields: `loss: Tensor`, `output: Tensor`, `targets: Tensor` + - Adapted metrics: Loss +- `SequenceOutput`: + - Use case: Sequence prediction + - Fields: `loss: Tensor`, `logits: Tensor`, `predictions: Option>`, `targets: Tensor` + - Adapted metrics: Accuracy, TopKAccuracy, Perplexity, CER, WER, Loss + +\* Precision, Recall, FBetaScore, AUROC, and AUC-PR all use `ConfusionStatsInput` as their input type so these +metrics are automatically (implicitly) adapted since `ConfusionStatsInput` is adapted. + +If your metric isn't already adapted for the appropriate output struct, you can implement `Adaptor` yourself. +For example, here is how `ClassificationOutput` adapts to `AccuracyInput`: + +```rust,ignore +impl Adaptor> for ClassificationOutput { + fn adapt(&self) -> AccuracyInput { + AccuracyInput::new(self.output.clone(), self.targets.clone()) + } +} +``` + +If your task type is not covered by the built-in output structs, you can create an output struct for your data +and then adapt your metric for the output struct: + +```rust,ignore +#[derive(new)] +pub struct ClassificationOutput { + /// The loss. + pub loss: Tensor, + + /// The output. + pub output: Tensor, + + /// The targets. + pub targets: Tensor, +} + +impl Adaptor> for ClassificationOutput { + fn adapt(&self) -> AccuracyInput { + AccuracyInput::new(self.output.clone(), self.targets.clone()) + } +} +``` + +You can also open an issue on the [GitHub repository](https://github.com/tracel-ai/burn) when your task type is +not covered by the built-in output structs. However, since creating an output struct for your data is simple, +it is recommended to try creating your own output struct first. + +# Custom Metric + +Generating your own custom metrics is done by implementing the `Metric` trait. + +```rust , ignore + +/// Metric trait. +/// +/// # Notes +/// +/// Implementations should define their own input type only used by the metric. +/// This is important since some conflict may happen when the model output is adapted for each +/// metric's input type. +pub trait Metric: Send + Sync + Clone { + /// The input type of the metric. + type Input; + + /// The parameterized name of the metric. + /// + /// This should be unique, so avoid using short generic names, prefer using the long name. + /// + /// For a metric that can exist at different parameters (e.g., top-k accuracy for different + /// values of k), the name should be unique for each instance. + fn name(&self) -> MetricName; + + /// Update the metric state and returns the current metric entry. + fn update(&mut self, item: &Self::Input, metadata: &MetricMetadata) -> SerializedEntry; + + /// Clear the metric state. + fn clear(&mut self); +} +``` + +As an example, let's see how the loss metric is implemented. + +```rust, ignore +/// The loss metric. +#[derive(Clone)] +pub struct LossMetric { + name: Arc, + state: NumericMetricState, + _b: B, +} + +/// The [loss metric](LossMetric) input type. +#[derive(new)] +pub struct LossInput { + tensor: Tensor, +} + +impl Default for LossMetric { + fn default() -> Self { + Self::new() + } +} + +impl LossMetric { + /// Create the metric. + pub fn new() -> Self { + Self { + name: Arc::new("Loss".to_string()), + state: NumericMetricState::default(), + _b: Default::default(), + } + } +} + + +impl Metric for LossMetric { + type Input = LossInput; + + fn update(&mut self, loss: &Self::Input, _metadata: &MetricMetadata) -> SerializedEntry { + let [batch_size] = loss.tensor.dims(); + let loss = loss + .tensor + .clone() + .mean() + .into_data() + .iter::() + .next() + .unwrap(); + + self.state.update( + loss, + batch_size, + FormatOptions::new(self.name()).precision(2), + ) + } + + fn clear(&mut self) { + self.state.reset() + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: None, + higher_is_better: false, + } + .into() + } +} +``` + +When the metric you are implementing is numeric in nature, you may want to also implement the +`Numeric` trait. This will allow your metric to be plotted. + +```rust, ignore +impl Numeric for LossMetric { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} +``` diff --git a/burn-book/src/building-blocks/module.md b/burn-book/src/building-blocks/module.md new file mode 100644 index 0000000..cc8ba80 --- /dev/null +++ b/burn-book/src/building-blocks/module.md @@ -0,0 +1,343 @@ +# Module + +The `Module` derive allows you to create your own neural network modules, similar to PyTorch. The +derive function only generates the necessary methods to essentially act as a parameter container for +your type, it makes no assumptions about how the forward pass is declared. + +```rust, ignore +use burn::module::Module; +use burn::tensor::backend::Backend; + +#[derive(Module, Debug)] +pub struct PositionWiseFeedForward { + linear_inner: Linear, + linear_outer: Linear, + dropout: Dropout, + gelu: Gelu, +} + +impl PositionWiseFeedForward { + /// Normal method added to a struct. + pub fn forward(&self, input: Tensor) -> Tensor { + let x = self.linear_inner.forward(input); + let x = self.gelu.forward(x); + let x = self.dropout.forward(x); + + self.linear_outer.forward(x) + } +} +``` + +Note that all fields declared in the struct must also implement the `Module` trait. + +## Tensor + +If you want to create your own module that contains tensors, and not just other modules defined with +the `Module` derive, you need to be careful to achieve the behavior you want. + +- `Param>`: If you want the tensor to be included as a parameter of your modules, you + need to wrap the tensor in a `Param` struct. This will create an ID that will be used to identify + this parameter. This is essential when performing module optimization and when saving states such + as optimizer and module checkpoints. Note that a module's record only contains parameters. + +- `Param>.set_require_grad(false)`: If you want the tensor to be included as a + parameter of your modules, and therefore saved with the module's weights, but you don't want it to + be updated by the optimizer. + +- `Tensor`: If you want the tensor to act as a constant that can be recreated when + instantiating a module. This can be useful when generating sinusoidal embeddings, for example. + +## Methods + +These methods are available for all modules. + +| Burn API | PyTorch Equivalent | +| --------------------------------------- | ---------------------------------------- | +| `module.devices()` | N/A | +| `module.fork(device)` | Similar to `module.to(device).detach()` | +| `module.to_device(device)` | `module.to(device)` | +| `module.no_grad()` | `module.require_grad_(False)` | +| `module.num_params()` | N/A | +| `module.visit(visitor)` | N/A | +| `module.map(mapper)` | N/A | +| `module.into_record()` | Similar to `state_dict` | +| `module.load_record(record)` | Similar to `load_state_dict(state_dict)` | +| `module.into_record().save(file_path)` | Similar to `torch.save(state_dict, ...)` | +| `ModuleRecord::load(file_path)` | Similar to `torch.load(...)` | + +Similar to the backend trait, there is also the `AutodiffModule` trait to signify a module with +autodiff support. + +| Burn API | PyTorch Equivalent | +| ---------------- | ------------------ | +| `module.valid()` | `module.eval()` | + +## Visitor & Mapper + +As mentioned earlier, modules primarily function as parameter containers. Therefore, we naturally +offer several ways to perform functions on each parameter. This is distinct from PyTorch, where +extending module functionalities is not as straightforward. + +The `map` and `visitor` methods are quite similar but serve different purposes. Mapping is used for +potentially mutable operations where each parameter of a module can be updated to a new value. In +Burn, optimizers are essentially just sophisticated module mappers. Visitors, on the other hand, are +used when you don't intend to modify the module but need to retrieve specific information from it, +such as the number of parameters or a list of devices in use. + +You can implement your own mapper or visitor by implementing these simple traits: + +```rust, ignore +/// Module visitor trait. +pub trait ModuleVisitor { + /// Visit a float tensor in the module. + fn visit_float(&mut self, id: ParamId, tensor: &Tensor); + /// Visit an int tensor in the module. + fn visit_int(&mut self, id: ParamId, tensor: &Tensor); + /// Visit a bool tensor in the module. + fn visit_bool(&mut self, id: ParamId, tensor: &Tensor); +} + +/// Module mapper trait. +pub trait ModuleMapper { + /// Map a float tensor in the module. + fn map_float(&mut self, id: ParamId, tensor: Tensor) -> Tensor; + /// Map an int tensor in the module. + fn map_int(&mut self, id: ParamId, tensor: Tensor) -> Tensor; + /// Map a bool tensor in the module. + fn map_bool(&mut self, id: ParamId, tensor: Tensor) -> Tensor; +} +``` + +Note that the trait doesn't require all methods to be implemented as they are already defined to +perform no operation. If you're only interested in float tensors (like the majority of use cases), +then you can simply implement `map_float` or `visit_float`. + +For example, the `ModuleMapper` trait could be implemented to clamp all parameters into the range +`[min, max]`. + +```rust, ignore +/// Clamp parameters into the range `[min, max]`. +pub struct Clamp { + /// Lower-bound of the range. + pub min: f32, + /// Upper-bound of the range. + pub max: f32, +} + +// Clamp all floating-point parameter tensors between `[min, max]`. +impl ModuleMapper for Clamp { + fn map_float( + &mut self, + _id: burn::module::ParamId, + tensor: burn::prelude::Tensor, + ) -> burn::prelude::Tensor { + tensor.clamp(self.min, self.max) + } +} + +// Clamp module mapper into the range `[-0.5, 0.5]` +let mut clamp = Clamp { + min: -0.5, + max: 0.5, +}; +let model = model.map(&mut clamp); +``` + +If you want to use this during training to constrain your model parameters, make sure that the +parameter tensors are still tracked for autodiff. This can be done with a simple adjustment to the +implementation. + +```rust, ignore +impl ModuleMapper for Clamp { + fn map_float( + &mut self, + _id: burn::module::ParamId, + tensor: burn::prelude::Tensor, + ) -> burn::prelude::Tensor { + let is_require_grad = tensor.is_require_grad(); + + let mut tensor = Tensor::from_inner(tensor.inner().clamp(self.min, self.max)); + + if is_require_grad { + tensor = tensor.require_grad(); + } + + tensor + } +} +``` + +## Module Display + +Burn provides a simple way to display the structure of a module and its configuration at a glance. +You can print the module to see its structure, which is useful for debugging and tracking changes +across different versions of a module. (See the print output of the +[Basic Workflow Model](../basic-workflow/model.md) example.) + +To customize the display of a module, you can implement the `ModuleDisplay` trait for your module. +This will change the default display settings for the module and its children. Note that +`ModuleDisplay` is automatically implemented for all modules, but you can override it to customize +the display by annotating the module with `#[module(custom_display)]`. + +```rust +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct PositionWiseFeedForward { + linear_inner: Linear, + linear_outer: Linear, + dropout: Dropout, + gelu: Gelu, +} + +impl ModuleDisplay for PositionWiseFeedForward { + /// Custom settings for the display of the module. + /// If `None` is returned, the default settings will be used. + fn custom_settings(&self) -> Option { + DisplaySettings::new() + // Will show all attributes (default is false) + .with_show_all_attributes(false) + // Will show each attribute on a new line (default is true) + .with_new_line_after_attribute(true) + // Will show the number of parameters (default is true) + .with_show_num_parameters(true) + // Will indent by 2 spaces (default is 2) + .with_indentation_size(2) + // Will show the parameter ID (default is false) + .with_show_param_id(false) + // Convenience method to wrap settings in Some() + .optional() + } + + /// Custom content to be displayed. + /// If `None` is returned, the default content will be used + /// (all attributes of the module) + fn custom_content(&self, content: Content) -> Option { + content + .add("linear_inner", &self.linear_inner) + .add("linear_outer", &self.linear_outer) + .add("anything", "anything_else") + .optional() + } +} +``` + +## Built-in Modules + +Burn comes with built-in modules that you can use to build your own modules. + +### General + +| Burn API | PyTorch Equivalent | +| ----------------- | --------------------------------------------- | +| `BatchNorm` | `nn.BatchNorm1d`, `nn.BatchNorm2d` etc. | +| `Celu` | `nn.CELU` | +| `Dropout` | `nn.Dropout` | +| `Elu` | `nn.ELU` | +| `Embedding` | `nn.Embedding` | +| `GaussianNoise` | _No direct equivalent_ | +| `Gelu` | `nn.Gelu` | +| `Glu` | `nn.Glu` | +| `GroupNorm` | `nn.GroupNorm` | +| `HardShrink` | `nn.Hardshrink` | +| `HardSigmoid` | `nn.Hardsigmoid` | +| `HardSwish` | `nn.Hardswish` | +| `InstanceNorm` | `nn.InstanceNorm1d`, `nn.InstanceNorm2d` etc. | +| `LayerNorm` | `nn.LayerNorm` | +| `LocalResponseNorm` | `nn.LocalResponseNorm` | +| `LeakyRelu` | `nn.LeakyReLU` | +| `Linear` | `nn.Linear` | +| `Prelu` | `nn.PReLu` | +| `Relu` | `nn.ReLU` | +| `Selu` | `nn.SELU` | +| `Sigmoid` | `nn.Sigmoid` | +| `Softplus` | `nn.Softplus` | +| `SoftShrink` | `nn.Softshrink` | +| `Softsign` | `nn.Softsign` | +| `Shrink` | _No direct equivalent_ | +| `RmsNorm` | _No direct equivalent_ | +| `SwiGlu` | _No direct equivalent_ | +| `Tanh` | `nn.Tanh` | +| `ThresholdedRelu` | _No direct equivalent_ | + +### Convolutions + +| Burn API | PyTorch Equivalent | +| ----------------- | ------------------------------ | +| `Conv1d` | `nn.Conv1d` | +| `Conv2d` | `nn.Conv2d` | +| `Conv3d` | `nn.Conv3d` | +| `ConvTranspose1d` | `nn.ConvTranspose1d` | +| `ConvTranspose2d` | `nn.ConvTranspose2d` | +| `ConvTranspose3d` | `nn.ConvTranspose3d` | +| `DeformConv2d` | `torchvision.ops.DeformConv2d` | + +### Pooling + +| Burn API | PyTorch Equivalent | +| ------------------- | ---------------------- | +| `AdaptiveAvgPool1d` | `nn.AdaptiveAvgPool1d` | +| `AdaptiveAvgPool2d` | `nn.AdaptiveAvgPool2d` | +| `AvgPool1d` | `nn.AvgPool1d` | +| `AvgPool2d` | `nn.AvgPool2d` | +| `MaxPool1d` | `nn.MaxPool1d` | +| `MaxPool2d` | `nn.MaxPool2d` | + +### Interpolation + +| Burn API | PyTorch Equivalent | +| --------------- | ------------------ | +| `Interpolate1d` | `nn.Upsample` | +| `Interpolate2d` | `nn.Upsample` | + +Interpolation modules resize tensors using one of the available `InterpolateMode` options: + +| Mode | Description | +| --------- | -------------------------------------------------------- | +| `Nearest` | Nearest-neighbor interpolation | +| `Linear` | Linear interpolation (bilinear for 2D) | +| `Cubic` | Cubic interpolation (bicubic for 2D) | +| `Lanczos` | Lanczos3 resampling (6-tap sinc-based filter, a=3) | + +Configuration is done via `Interpolate1dConfig` / `Interpolate2dConfig` with these options: + +| Option | Type | Default | Description | +| --------------- |------------------------------------------| --------- | -------------------------------------------------------- | +| `output_size` | `Option` / `Option<[usize; 2]>` | `None` | Target output size (takes precedence over scale_factor) | +| `scale_factor` | `Option` / `Option<[f32; 2]>` | `None` | Scale factor for resizing | +| `mode` | `InterpolateMode` | `Nearest` | Interpolation algorithm | +| `align_corners` | `bool` | `true` | Align input/output corner pixels | + +### RNNs + +| Burn API | PyTorch Equivalent | +| ---------------- | ---------------------- | +| `Gru`/`BiGru` | `nn.GRU` | +| `Lstm`/`BiLstm` | `nn.LSTM` | +| `GateController` | _No direct equivalent_ | + +### Transformer + +| Burn API | PyTorch Equivalent | +| -------------------- | ----------------------- | +| `MultiHeadAttention` | `nn.MultiheadAttention` | +| `TransformerDecoder` | `nn.TransformerDecoder` | +| `TransformerEncoder` | `nn.TransformerEncoder` | +| `PositionalEncoding` | _No direct equivalent_ | +| `RotaryEncoding` | _No direct equivalent_ | + +### Loss + +| Burn API | PyTorch Equivalent | +| ------------------------ | ------------------------ | +| `BinaryCrossEntropyLoss` | `nn.BCELoss` | +| `CosineEmbeddingLoss` | `nn.CosineEmbeddingLoss` | +| `CrossEntropyLoss` | `nn.CrossEntropyLoss` | +| `CTCLoss` | `nn.CTCLoss` | +| `GramMatrixLoss` | _No direct equivalent_ | +| `HuberLoss` | `nn.HuberLoss` | +| `KLDivLoss` | `nn.KLDivLoss` | +| `LpLoss` | _No direct equivalent_ | +| `MseLoss` | `nn.MSELoss` | +| `PoissonNllLoss` | `nn.PoissonNLLLoss` | +| `RNNTLoss` | `torchaudio.functional.rnnt_loss` | +| `SmoothL1Loss` | `nn.SmoothL1Loss` | diff --git a/burn-book/src/building-blocks/record.md b/burn-book/src/building-blocks/record.md new file mode 100644 index 0000000..d75b21f --- /dev/null +++ b/burn-book/src/building-blocks/record.md @@ -0,0 +1,85 @@ +# Record + +Records are how training state is saved and loaded with Burn. A record holds plain tensor data +(decoupled from the backend in use), so weights saved with one backend can be loaded on another, and +parameter initialization stays lazy. + +All records serialize to the **burnpack** format (`.bpk`), Burn's compact binary container +implemented by the `burn-pack` crate. + +## The burnpack format + +A burnpack file has three parts: + +- a small fixed-size **header** (a `"BURN"` magic number, a format version, and the metadata length); +- a **metadata** blob (CBOR) describing each tensor (name, dtype, shape, data offsets, optional + parameter id), any named **typed scalars**, and user key/value pairs; +- a **tensor data section** where each tensor's bytes start on a 256-byte boundary, so the data can + be read back with zero-copy / memory-mapped loading. + +Storing typed scalars (integers, floats, booleans) alongside tensors is what lets the optimizer and +learning rate scheduler persist their non-tensor state in the same format. + +## The three record types + +| Record | Holds | Produced from | +| -------------------- | ------------------------------ | ---------------------------------------------- | +| `ModuleRecord` | a module's parameters | `module.into_record()` | +| `OptimizerRecord` | the optimizer state | `optimizer.to_record()` | +| `LrSchedulerRecord` | the learning rate scheduler | `scheduler.to_record()` | + +Each record can be written to a file (`save` / `load`, which appends the `.bpk` extension when the +path has none) or to an in-memory byte buffer (`into_bytes` / `from_bytes`, useful for `no-std` +deployment where the bytes are embedded with the compiled code). + +### `ModuleRecord` + +`ModuleRecord` (in `burn::store`) holds a module's parameters keyed by their path within the module. +It is produced and applied through the `Module` trait itself: + +```rust, ignore +use burn::store::ModuleRecord; + +// Take a record and save it. +model.into_record().save("model")?; // writes model.bpk + +// Load it back and apply it to an initialized module. +let record = ModuleRecord::load("model")?; +let model = ModelConfig::new().init(&device).load_record(record); +``` + +Load-time behavior is configured with builder methods on the record (ignored when saving): + +- `.allow_partial(true)` — load even when some module parameters are absent from the record; +- `.validate(false)` — skip shape-mismatch / missing-tensor validation; +- `.cast_to_module_dtype()` / `.with_dtype_policy(..)` — cast the record's data to the module + parameter dtypes on load (by default the parameter adopts the record's dtype). + +The save-side dtype is not configurable: the record stores whatever dtype the module currently holds. +To control the dtype applied on load, use `.cast_to_module_dtype()` / `.with_dtype_policy(..)` above. +Use `try_load_record` for the fallible variant of `load_record`. + +### `OptimizerRecord` and `LrSchedulerRecord` + +The optimizer and learning rate scheduler expose the same shape of API, used to checkpoint and resume +training: + +```rust, ignore +// Optimizer state (no device needed on load; state migrates to each parameter's device on the +// next step). +optimizer.save("optim")?; +let optimizer = optimizer.load("optim")?; + +// Learning rate scheduler state (scalars only). +scheduler.to_record().save("scheduler")?; +let scheduler = scheduler.load_record(LrSchedulerRecord::load("scheduler")?); +``` + +When training with the `Learner`, these records are saved and restored for you by the checkpointer — +see [Learner](./learner.md). + +## Cross-framework formats + +To import weights from other ecosystems (PyTorch `.pt`, SafeTensors) or to use the more advanced +store features (key remapping, filtering, half-precision storage), use the `burn-store` crate. See +[Saving and Loading Models](../saving-and-loading.md) for examples. diff --git a/burn-book/src/building-blocks/tensor.md b/burn-book/src/building-blocks/tensor.md new file mode 100644 index 0000000..fc80bb2 --- /dev/null +++ b/burn-book/src/building-blocks/tensor.md @@ -0,0 +1,588 @@ +# Tensor + +As previously explained in the [model section](../basic-workflow/model.md), the Tensor struct has 3 +generic arguments: the backend B, the dimensionality D, and the data type. + +```rust, ignore +Tensor // Float tensor (default) +Tensor // Explicit float tensor +Tensor // Int tensor +Tensor // Bool tensor +``` + +Note that the specific element types used for `Float`, `Int`, and `Bool` tensors are defined by +backend implementations. + +Burn Tensors are defined by the number of dimensions D in its declaration as opposed to its shape. +The actual shape of the tensor is inferred from its initialization. For example, a Tensor of size +(5,) is initialized as below: + +```rust, ignore +let floats = [1.0, 2.0, 3.0, 4.0, 5.0]; + +// Get the default device +let device = Default::default(); + +// correct: Tensor is 1-Dimensional with 5 elements +let tensor_1 = Tensor::::from_floats(floats, &device); + +// incorrect: let tensor_1 = Tensor::::from_floats(floats, &device); +// this will lead to an error and is for creating a 5-D tensor +``` + +### Initialization + +Burn Tensors are primarily initialized using the `from_data()` method which takes the `TensorData` +struct as input. The `TensorData` struct has two public fields: `shape` and `dtype`. The `value`, +now stored as bytes, is private but can be accessed via any of the following methods: `as_slice`, +`as_mut_slice`, `to_vec` and `iter`. To retrieve the data from a tensor, the method `.to_data()` +should be employed when intending to reuse the tensor afterward. Alternatively, `.into_data()` is +recommended for one-time use. Let's look at a couple of examples for initializing a tensor from +different inputs. + +```rust, ignore + +// Initialization from a given Backend (Wgpu) +let tensor_1 = Tensor::::from_data([1.0, 2.0, 3.0], &device); + +// Initialization from a generic Backend +let tensor_2 = Tensor::::from_data(TensorData::from([1.0, 2.0, 3.0]), &device); + +// Initialization using from_floats (Recommended for f32 ElementType) +// Will be converted to TensorData internally. +let tensor_3 = Tensor::::from_floats([1.0, 2.0, 3.0], &device); + +// Initialization of Int Tensor from array slices +let arr: [i32; 6] = [1, 2, 3, 4, 5, 6]; +let tensor_4 = Tensor::::from_data(TensorData::from(&arr[0..3]), &device); + +// Initialization from a custom type + +struct BodyMetrics { + age: i8, + height: i16, + weight: f32 +} + +let bmi = BodyMetrics{ + age: 25, + height: 180, + weight: 80.0 + }; +let data = TensorData::from([bmi.age as f32, bmi.height as f32, bmi.weight]); +let tensor_5 = Tensor::::from_data(data, &device); + +``` + +## Ownership and Cloning + +Almost all Burn operations take ownership of the input tensors. Therefore, reusing a tensor multiple +times will necessitate cloning it. Let's look at an example to understand the ownership rules and +cloning better. Suppose we want to do a simple min-max normalization of an input tensor. + +```rust, ignore +let input = Tensor::::from_floats([1.0, 2.0, 3.0, 4.0], &device); +let min = input.min(); +let max = input.max(); +let input = (input - min).div(max - min); +``` + +With PyTorch tensors, the above code would work as expected. However, Rust's strict ownership rules +will give an error and prevent using the input tensor after the first `.min()` operation. The +ownership of the input tensor is transferred to the variable `min` and the input tensor is no longer +available for further operations. Burn Tensors like most complex primitives do not implement the +`Copy` trait and therefore have to be cloned explicitly. Now let's rewrite a working example of +doing min-max normalization with cloning. + +```rust, ignore +let input = Tensor::::from_floats([1.0, 2.0, 3.0, 4.0], &device); +let min = input.clone().min(); +let max = input.clone().max(); +let input = (input.clone() - min.clone()).div(max - min); +println!("{}", input.to_data());// Success: [0.0, 0.33333334, 0.6666667, 1.0] + +// Notice that max, min have been moved in last operation so +// the below print will give an error. +// If we want to use them for further operations, +// they will need to be cloned in similar fashion. +// println!("{:?}", min.to_data()); +``` + +We don't need to be worried about memory overhead because with cloning, the tensor's buffer isn't +copied, and only a reference to it is increased. This makes it possible to determine exactly how +many times a tensor is used, which is very convenient for reusing tensor buffers or even fusing +operations into a single kernel ([burn-fusion](https://burn.dev/docs/burn_fusion/index.htmls)). For +that reason, we don't provide explicit inplace operations. If a tensor is used only one time, +inplace operations will always be used when available. + +## Tensor Operations + +Normally with PyTorch, explicit inplace operations aren't supported during the backward pass, making +them useful only for data preprocessing or inference-only model implementations. With Burn, you can +focus more on _what_ the model should do, rather than on _how_ to do it. We take the responsibility +of making your code run as fast as possible during training as well as inference. The same +principles apply to broadcasting; all operations support broadcasting unless specified otherwise. + +Here, we provide a list of all supported operations along with their PyTorch equivalents. Note that +for the sake of simplicity, we ignore type signatures. For more details, refer to the +[full documentation](https://docs.rs/burn/latest/burn/tensor/struct.Tensor.html). + +### Basic Operations + +Those operations are available for all tensor kinds: `Int`, `Float`, and `Bool`. + +| Burn | PyTorch Equivalent | +| ---------------------------------------------------- | ------------------------------------------------------------------------- | +| `Tensor::cat(tensors, dim)` | `torch.cat(tensors, dim)` | +| `Tensor::empty(shape, options)` | `torch.empty(shape, device=device, dtype=dtype)` | +| `Tensor::from_primitive(primitive)` | N/A | +| `Tensor::stack(tensors, dim)` | `torch.stack(tensors, dim)` | +| `tensor.all()` | `tensor.all()` | +| `tensor.all_dim(dim)` | `tensor.all(dim)` | +| `tensor.any()` | `tensor.any()` | +| `tensor.any_dim(dim)` | `tensor.any(dim)` | +| `tensor.chunk(num_chunks, dim)` | `tensor.chunk(num_chunks, dim)` | +| `tensor.split(split_size, dim)` | `tensor.split(split_size, dim)` | +| `tensor.split_with_sizes(split_sizes, dim)` | `tensor.split([split_sizes], dim)` | +| `tensor.device()` | `tensor.device` | +| `tensor.dtype()` | `tensor.dtype` | +| `tensor.dims()` | `tensor.size()` | +| `tensor.equal(other)` | `x == y` | +| `tensor.equal_elem(other)` | `tensor.eq(other)` | +| `tensor.expand(shape)` | `tensor.expand(shape)` | +| `tensor.flatten(start_dim, end_dim)` | `tensor.flatten(start_dim, end_dim)` | +| `tensor.flip(axes)` | `tensor.flip(axes)` | +| `tensor.full_like(fill_value)` | `torch.full_like(tensor, fill_value)` | +| `tensor.gather(dim, indices)` | `torch.gather(tensor, dim, indices)` | +| `tensor.into_data()` | N/A | +| `tensor.into_primitive()` | N/A | +| `tensor.into_scalar()` | `tensor.item()` | +| `tensor.mask_fill(mask, value)` | `tensor.masked_fill(mask, value)` | +| `tensor.mask_where(mask, value_tensor)` | `torch.where(mask, value_tensor, tensor)` | +| `tensor.movedim(src, dst)` | `tensor.movedim(src, dst)` | +| `tensor.narrow(dim, start, length)` | `tensor.narrow(dim, start, length)` | +| `tensor.not_equal(other)` | `x != y` | +| `tensor.not_equal_elem(scalar)` | `tensor.ne(scalar)` | +| `tensor.ones_like()` | `torch.ones_like(tensor)` | +| `tensor.permute(axes)` | `tensor.permute(axes)` | +| `tensor.repeat_dim(dim, times)` | `tensor.repeat(*[times if i == dim else 1 for i in range(tensor.dim())])` | +| `tensor.repeat(sizes)` | `tensor.repeat(sizes)` | +| `tensor.reshape(shape)` | `tensor.view(shape)` | +| `tensor.roll(shifts, dims)` | `tensor.roll(shifts, dims)` | +| `tensor.roll_dim(shift, dim)` | `tensor.roll([shift], [dim])` | +| `tensor.scatter(dim, indices, values, update)` | `tensor.scatter_add(dim, indices, values)` | +| `tensor.scatter_nd(indices, values, update)` | N/A | +| `tensor.gather_nd(indices)` | N/A | +| `tensor.select(dim, indices)` | `tensor.index_select(dim, indices)` | +| `tensor.select_assign(dim, indices, values, update)` | `tensor.index_add(dim, indices, values)` | +| `tensor.shape()` | `tensor.shape` | +| `tensor.slice(slices)` | `tensor[(*ranges,)]` | +| `tensor.slice_assign(slices, values)` | `tensor[(*ranges,)] = values` | +| `tensor.slice_fill(slices, value)` | `tensor[(*ranges,)] = value` | +| `tensor.slice_dim(dim, slice)` | N/A | +| `tensor.squeeze()` | `tensor.squeeze()` | +| `tensor.squeeze_dim(dim)` | `tensor.squeeze(dim)` | +| `tensor.squeeze_dims(dims)` | `tensor.squeeze(dims)` where `dims` is a tuple of ints | +| `tensor.swap_dims(dim1, dim2)` | `tensor.transpose(dim1, dim2)` | +| `tensor.take(dim, indices)` | `numpy.take(tensor, indices, dim)` | +| `tensor.to_data()` | N/A | +| `tensor.to_device(device)` | `tensor.to(device)` | +| `tensor.transpose()` | `tensor.T` | +| `tensor.t()` | `tensor.T` | +| `tensor.unsqueeze()` | N/A | +| `tensor.unsqueeze_dim(dim)` | `tensor.unsqueeze(dim)` | +| `tensor.unsqueeze_dims(dims)` | N/A | +| `tensor.zeros_like()` | `torch.zeros_like(tensor)` | +| `Tensor::full(shape, fill_value, options)` | `torch.full(shape, fill_value, device=device, dtype=dtype)` | +| `Tensor::ones(shape, options)` | `torch.ones(shape, device=device, dtype=dtype)` | +| `Tensor::zeros(shape, options)` | `torch.zeros(shape, device=device, dtype=dtype)` | + +### Numeric Operations + +Those operations are available for numeric tensor kinds: `Float` and `Int`. + +| Burn | PyTorch Equivalent | +| --------------------------------------------------------------- | --------------------------------------------- | +| `tensor.abs()` | `torch.abs(tensor)` | +| `tensor.add(other)` or `tensor + other` | `tensor + other` | +| `tensor.add_scalar(scalar)` or `tensor + scalar` | `tensor + scalar` | +| `tensor.all_close(other, atol, rtol)` | `torch.allclose(tensor, other, atol, rtol)` | +| `tensor.argmax(dim)` | `tensor.argmax(dim)` | +| `tensor.argmin(dim)` | `tensor.argmin(dim)` | +| `tensor.argsort(dim)` | `tensor.argsort(dim)` | +| `tensor.argsort_descending(dim)` | `tensor.argsort(dim, descending=True)` | +| `tensor.bool()` | `tensor.bool()` | +| `tensor.clamp(min, max)` | `torch.clamp(tensor, min=min, max=max)` | +| `tensor.clamp_max(max)` | `torch.clamp(tensor, max=max)` | +| `tensor.clamp_min(min)` | `torch.clamp(tensor, min=min)` | +| `tensor.cumsum(dim)` | `tensor.cumsum(dim)` | +| `tensor.cumprod(dim)` | `tensor.cumprod(dim)` | +| `tensor.cummin(dim)` | `tensor.cummin(dim)` | +| `tensor.cummax(dim)` | `tensor.cummax(dim)` | +| `tensor.div(other)` or `tensor / other` | `tensor / other` | +| `tensor.div_scalar(scalar)` or `tensor / scalar` | `tensor / scalar` | +| `tensor.dot(other)` | `torch.dot(tensor, other)` | +| `tensor.greater(other)` | `tensor.gt(other)` | +| `tensor.greater_elem(scalar)` | `tensor.gt(scalar)` | +| `tensor.greater_equal(other)` | `tensor.ge(other)` | +| `tensor.greater_equal_elem(scalar)` | `tensor.ge(scalar)` | +| `tensor.lower(other)` | `tensor.lt(other)` | +| `tensor.lower_elem(scalar)` | `tensor.lt(scalar)` | +| `tensor.lower_equal(other)` | `tensor.le(other)` | +| `tensor.lower_equal_elem(scalar)` | `tensor.le(scalar)` | +| `tensor.max()` | `tensor.max()` | +| `tensor.max_abs()` | `tensor.abs().max()` | +| `tensor.max_abs_dim(dim)` | `tensor.abs().max(dim, keepdim=True)` | +| `tensor.max_abs_dims(dims)` | `tensor.abs().max(dims, keepdim=True)` | +| `tensor.max_dim(dim)` | `tensor.max(dim, keepdim=True)` | +| `tensor.max_dims(dims)` | `tensor.max(dims, keepdim=True)` | +| `tensor.max_dim_with_indices(dim)` | N/A | +| `tensor.max_pair(other)` | `torch.Tensor.max(a,b)` | +| `tensor.mean()` | `tensor.mean()` | +| `tensor.mean_dim(dim)` | `tensor.mean(dim, keepdim=True)` | +| `tensor.mean_dims(dims)` | `tensor.mean(dims, keepdim=True)` | +| `tensor.min()` | `tensor.min()` | +| `tensor.min_dim(dim)` | `tensor.min(dim, keepdim=True)` | +| `tensor.min_dims(dims)` | `tensor.min(dims, keepdim=True)` | +| `tensor.min_dim_with_indices(dim)` | N/A | +| `tensor.min_pair(other)` | `torch.Tensor.min(a,b)` | +| `tensor.mul(other)` or `tensor * other` | `tensor * other` | +| `tensor.mul_scalar(scalar)` or `tensor * scalar` | `tensor * scalar` | +| `tensor.neg()` or `-tensor` | `-tensor` | +| `tensor.one_hot(num_classes)` | `torch.nn.functional.one_hot` | +| `tensor.one_hot_fill(num_classes, on_value, off_value, axis)` | N/A | +| `tensor.pad(pads, mode)` | `torch.nn.functional.pad(tensor, pads, mode)` | +| `tensor.powf(other)` or `tensor.powi(intother)` | `tensor.pow(other)` | +| `tensor.powf_scalar(scalar)` or `tensor.powi_scalar(intscalar)` | `tensor.pow(scalar)` | +| `tensor.prod()` | `tensor.prod()` | +| `tensor.prod_dim(dim)` | `tensor.prod(dim, keepdim=True)` | +| `tensor.prod_dims(dims)` | `tensor.prod(dims, keepdim=True)` | +| `tensor.rem(other)` or `tensor % other` | `tensor % other` | +| `tensor.sign()` | `tensor.sign()` | +| `tensor.sort(dim)` | `tensor.sort(dim).values` | +| `tensor.sort_descending(dim)` | `tensor.sort(dim, descending=True).values` | +| `tensor.sort_descending_with_indices(dim)` | `tensor.sort(dim, descending=True)` | +| `tensor.sort_with_indices(dim)` | `tensor.sort(dim)` | +| `tensor.sub(other)` or `tensor - other` | `tensor - other` | +| `tensor.sub_scalar(scalar)` or `tensor - scalar` | `tensor - scalar` | +| `tensor.sum()` | `tensor.sum()` | +| `tensor.sum_dim(dim)` | `tensor.sum(dim, keepdim=True)` | +| `tensor.sum_dims(dims)` | `tensor.sum(dims, keepdim=True)` | +| `tensor.sum_dims_squeeze(dims)` | `tensor.sum(dims, keepdim=False)` | +| `tensor.topk(k, dim)` | `tensor.topk(k, dim).values` | +| `tensor.topk_with_indices(k, dim)` | `tensor.topk(k, dim)` | +| `tensor.tril(diagonal)` | `torch.tril(tensor, diagonal)` | +| `tensor.triu(diagonal)` | `torch.triu(tensor, diagonal)` | +| `tensor.unfold(dim, size, step)` | `tensor.unfold(dim, size, step)` | +| `Tensor::eye(size, device)` | `torch.eye(size, device=device)` | +| `scalar - tensor` | `scalar - tensor` | + +### Float Operations + +Those operations are only available for `Float` tensors. + +| Burn API | PyTorch Equivalent | +| -------------------------------------------- | ------------------------------------------ | +| `tensor.acos()` | `tensor.acos()` | +| `tensor.acosh()` | `tensor.acosh()` | +| `tensor.asin()` | `tensor.asin()` | +| `tensor.asinh()` | `tensor.asinh()` | +| `tensor.atan()` | `tensor.atan()` | +| `tensor.atanh()` | `tensor.atanh()` | +| `tensor.atan2(other_tensor)` | `tensor.atan2(other_tensor)` | +| `tensor.cast(dtype)` | `tensor.to(dtype)` | +| `tensor.ceil()` | `tensor.ceil()` | +| `tensor.contains_nan()` | N/A | +| `tensor.cos()` | `tensor.cos()` | +| `tensor.cosh()` | `tensor.cosh()` | +| `tensor.cross(other)` | `torch.cross(tensor, other)` | +| `tensor.deg2rad()` | `torch.deg2rad()` | +| `tensor.erf()` | `tensor.erf()` | +| `tensor.exp()` | `tensor.exp()` | +| `tensor.floor()` | `tensor.floor()` | +| `tensor.fmod(other)` | `tensor.fmod(other)` | +| `tensor.fmod_scalar(scalar)` | `tensor.fmod(scalar)` | +| `tensor.from_floats(floats, device)` | N/A | +| `tensor.int()` | Similar to `tensor.to(torch.long)` | +| `tensor.is_close(other, atol, rtol)` | `torch.isclose(tensor, other, atol, rtol)` | +| `tensor.is_finite()` | `torch.isfinite(tensor)` | +| `tensor.is_inf()` | `torch.isinf(tensor)` | +| `tensor.is_nan()` | `torch.isnan(tensor)` | +| `tensor.log()` | `tensor.log()` | +| `tensor.log1p()` | `tensor.log1p()` | +| `tensor.matmul(other)` | `tensor.matmul(other)` | +| `tensor.rad2deg()` | `torch.rad2deg()` | +| `tensor.random(shape, distribution, device)` | N/A | +| `tensor.random_like(distribution)` | `torch.rand_like()` only uniform | +| `tensor.recip()` or `1.0 / tensor` | `tensor.reciprocal()` or `1.0 / tensor` | +| `tensor.round()` | `tensor.round()` | +| `tensor.sin()` | `tensor.sin()` | +| `tensor.sinh()` | `tensor.sinh()` | +| `tensor.square()` | `tensor.square()` | +| `tensor.sqrt()` | `tensor.sqrt()` | +| `tensor.tan()` | `tensor.tan()` | +| `tensor.tanh()` | `tensor.tanh()` | +| `tensor.trunc()` | `tensor.trunc()` | +| `tensor.var(dim)` | `tensor.var(dim)` | +| `tensor.var_bias(dim)` | N/A | +| `tensor.var_mean(dim)` | N/A | +| `tensor.var_mean_bias(dim)` | N/A | +| `tensor.median(dim)` | `tensor.median(dim)` | +| `tensor.median_with_indices(dim)` | `tensor.median(dim)` | + +### Int Operations + +Those operations are only available for `Int` tensors. + +| Burn API | PyTorch Equivalent | +| ------------------------------------------------ | ------------------------------------------------------- | +| `Tensor::arange(5..10, device)` | `tensor.arange(start=5, end=10, device=device)` | +| `Tensor::arange_step(5..10, 2, device)` | `tensor.arange(start=5, end=10, step=2, device=device)` | +| `tensor.bitwise_and(other)` | `torch.bitwise_and(tensor, other)` | +| `tensor.bitwise_and_scalar(scalar)` | `torch.bitwise_and(tensor, scalar)` | +| `tensor.bitwise_not()` | `torch.bitwise_not(tensor)` | +| `tensor.bitwise_left_shift(other)` | `torch.bitwise_left_shift(tensor, other)` | +| `tensor.bitwise_left_shift_scalar(scalar)` | `torch.bitwise_left_shift(tensor, scalar)` | +| `tensor.bitwise_right_shift(other)` | `torch.bitwise_right_shift(tensor, other)` | +| `tensor.bitwise_right_shift_scalar(scalar)` | `torch.bitwise_right_shift(tensor, scalar)` | +| `tensor.bitwise_or(other)` | `torch.bitwise_or(tensor, other)` | +| `tensor.bitwise_or_scalar(scalar)` | `torch.bitwise_or(tensor, scalar)` | +| `tensor.bitwise_xor(other)` | `torch.bitwise_xor(tensor, other)` | +| `tensor.bitwise_xor_scalar(scalar)` | `torch.bitwise_xor(tensor, scalar)` | +| `tensor.float()` | `tensor.to(torch.float)` | +| `tensor.from_ints(ints)` | N/A | +| `tensor.cartesian_grid(shape, device)` | N/A | + +### Bool Operations + +Those operations are only available for `Bool` tensors. + +| Burn API | PyTorch Equivalent | +| ------------------------------------ | ------------------------------- | +| `Tensor::diag_mask(shape, diagonal)` | N/A | +| `Tensor::tril_mask(shape, diagonal)` | N/A | +| `Tensor::triu_mask(shape, diagonal)` | N/A | +| `tensor.argwhere()` | `tensor.argwhere()` | +| `tensor.bool_and()` | `tensor.logical_and()` | +| `tensor.bool_not()` | `tensor.logical_not()` | +| `tensor.bool_or()` | `tensor.logical_or()` | +| `tensor.bool_xor()` | `tensor.logical_xor()` | +| `tensor.float()` | `tensor.to(torch.float)` | +| `tensor.int()` | `tensor.to(torch.long)` | +| `tensor.nonzero()` | `tensor.nonzero(as_tuple=True)` | + +### Quantization Operations + +Those operations are only available for `Float` tensors on backends that implement quantization +strategies. + +| Burn API | PyTorch Equivalent | +| ---------------------------------- | ------------------ | +| `tensor.quantize(scheme, qparams)` | N/A | +| `tensor.dequantize()` | N/A | + +## Activation Functions + +| Burn API | PyTorch Equivalent | +| ------------------------------------------------ | -------------------------------------------------- | +| `activation::celu(tensor, alpha)` | `nn.functional.celu(tensor, alpha)` | +| `activation::elu(tensor, alpha)` | `nn.functional.elu(tensor, alpha)` | +| `activation::gelu(tensor)` | `nn.functional.gelu(tensor)` | +| `activation::glu(tensor, dim)` | `nn.functional.glu(tensor, dim)` | +| `activation::hard_shrink(tensor, lambda)` | `nn.functional.hardshrink(tensor, lambd)` | +| `activation::hard_sigmoid(tensor, alpha, beta)` | `nn.functional.hardsigmoid(tensor)` | +| `activation::hard_swish(tensor)` | `nn.functional.hardswish(tensor)` | +| `activation::leaky_relu(tensor, negative_slope)` | `nn.functional.leaky_relu(tensor, negative_slope)` | +| `activation::log_sigmoid(tensor)` | `nn.functional.log_sigmoid(tensor)` | +| `activation::log_softmax(tensor, dim)` | `nn.functional.log_softmax(tensor, dim)` | +| `activation::mish(tensor)` | `nn.functional.mish(tensor)` | +| `activation::prelu(tensor,alpha)` | `nn.functional.prelu(tensor,weight)` | +| `activation::quiet_softmax(tensor, dim)` | `nn.functional.quiet_softmax(tensor, dim)` | +| `activation::relu(tensor)` | `nn.functional.relu(tensor)` | +| `activation::shrink(tensor, lambda, bias)` | _No direct equivalent_ | +| `activation::soft_shrink(tensor, lambda)` | `nn.functional.softshrink(tensor, lambd)` | +| `activation::sigmoid(tensor)` | `nn.functional.sigmoid(tensor)` | +| `activation::selu(tensor)` | `nn.functional.selu(tensor)` | +| `activation::silu(tensor)` | `nn.functional.silu(tensor)` | +| `activation::softmax(tensor, dim)` | `nn.functional.softmax(tensor, dim)` | +| `activation::softmin(tensor, dim)` | `nn.functional.softmin(tensor, dim)` | +| `activation::softplus(tensor, beta)` | `nn.functional.softplus(tensor, beta)` | +| `activation::softsign(tensor)` | `nn.functional.softsign(tensor)` | +| `activation::tanh(tensor)` | `nn.functional.tanh(tensor)` | +| `activation::thresholded_relu(tensor, alpha)` | `nn.functional.threshold(tensor, alpha, 0)` | + +## Grid Functions + +| Burn API | PyTorch Equivalent | +| --------------------------------------------------- | -------------------------------------------------------------------- | +| `grid::affine_grid_2d(transformation_tensor, dims)` | `nn.functional.affine_grid(theta_tensor, size, align_corners)` | +| `grid::meshgrid(tensors, GridIndexing::Matrix)` | `torch.meshgrid(tensors, indexing="ij")` | +| `grid::meshgrid(tensors, GridIndexing::Cartesian)` | `torch.meshgrid(tensors, indexing="xy")` | +| `grid::meshgrid_stack(tensors, index_pos)` | _No direct equivalent_ | + +## Linalg Functions + +| Burn API | PyTorch Equivalent | +| -------------------------------------------------- | --------------------------------------------------- | +| `linalg::cosine_similarity(x1, x2, dim, eps)` | `nn.functional.cosine_similarity(x1, x2, dim, eps)` | +| `linalg::det(tensor)` | `torch.linalg.det(tensor)` | +| `linalg::diag(tensor)` | `torch.diag(tensor)` | +| `linalg::l0_norm(tensor, dim)` | _No direct equivalent_ | +| `linalg::l1_norm(tensor, dim)` | _No direct equivalent_ | +| `linalg::l2_norm(tensor, dim)` | _No direct equivalent_ | +| `linalg::lp_norm(tensor, p, dim)` | _No direct equivalent_ | +| `linalg::lu(tensor)` | `torch.linalg.lu(tensor)` | +| `linalg::matvec(matrix, vector)` | `torch.matmul(matrix, vector)` / `@` operator | +| `linalg::max_abs_norm(tensor, dim)` | _No direct equivalent_ | +| `linalg::min_abs_norm(tensor, dim)` | _No direct equivalent_ | +| `linalg::outer(lhs, rhs)` | `torch.outer(lhs, rhs)` / `einsum("bi,bj->bij", …)` | +| `linalg::outer_dim(lhs, rhs, dim)` | _No direct equivalent_ | +| `linalg::trace(tensor)` | `torch.trace(tensor)` | +| `linalg::vector_norm(tensor, p, dim)` | `torch.linalg.vector_norm(tensor, p, dim)` | +| `linalg::vector_normalize(tensor, norm, dim, eps)` | `nn.functional.normalize(tensor, p, dim, eps)` | + +## Signal Processing Functions + +Signal-processing helpers live in `burn::tensor::signal` and operate on real-valued float +tensors. FFT length `n` (and `n_fft` in STFT) must currently be a power of two: when `n` is +`Some(size)`, the input is truncated or zero-padded to `size` and the output has +`size / 2 + 1` frequency bins. Non-power-of-two sizes panic at the public API boundary; +general arbitrary-size DFT support (Bluestein's algorithm) is a tracked follow-up. + +| Burn API | PyTorch Equivalent | +| ----------------------------------------------------- | --------------------------------------------------------------------------------- | +| `signal::rfft(tensor, dim, n)` | `torch.fft.rfft(tensor, n, dim)` | +| `signal::irfft(re, im, dim, n)` | `torch.fft.irfft(complex, n, dim)` | +| `signal::stft(signal, window, options)` | `torch.stft(signal, n_fft, hop_length, win_length, window, center)` | +| `signal::istft(stft_matrix, window, length, options)` | `torch.istft(stft_matrix, n_fft, hop_length, win_length, window, center, length)` | +| `signal::blackman_window(size, periodic, options)` | `torch.blackman_window(size, periodic)` | +| `signal::hamming_window(size, periodic, options)` | `torch.hamming_window(size, periodic)` | +| `signal::hann_window(size, periodic, options)` | `torch.hann_window(size, periodic)` | + +`stft` and `istft` share a `StftOptions` struct with fields `n_fft`, `hop_length`, +`win_length`, `center`, and `onesided`. Use `StftOptions::new(n_fft)` for PyTorch-style +defaults (`hop_length = n_fft / 4`, `win_length = None`, `center = true`, `onesided = true`). +The option set is validated on entry to both `stft` and `istft`; `n_fft` must be a power of +two and `hop_length <= effective_win_length` (the COLA prerequisite for invertibility). + +## Displaying Tensor Details + +Burn provides flexible options for displaying tensor information, allowing you to control the level +of detail and formatting to suit your needs. + +### Basic Display + +To display a detailed view of a tensor, you can simply use Rust's `println!` or `format!` macros: + +```rust, ignore +let tensor = Tensor::::full([2, 3], 0.123456789, &Default::default()); +println!("{}", tensor); +``` + +This will output: + +``` +Tensor { + data: +[[0.12345679, 0.12345679, 0.12345679], + [0.12345679, 0.12345679, 0.12345679]], + shape: [2, 3], + device: Cpu, + backend: "flex", + kind: "Float", + dtype: "f32", +} +``` + +### Controlling Precision + +You can control the number of decimal places displayed using Rust's formatting syntax: + +```rust +println!("{:.2}", tensor); +``` + +Output: + +``` +Tensor { + data: +[[0.12, 0.12, 0.12], + [0.12, 0.12, 0.12]], + shape: [2, 3], + device: Cpu, + backend: "flex", + kind: "Float", + dtype: "f32", +} +``` + +### Global Print Options + +For more fine-grained control over tensor printing, Burn provides a `PrintOptions` struct and a +`set_print_options` function: + +```rust, ignore +use burn::tensor::{set_print_options, PrintOptions}; + +let print_options = PrintOptions { + precision: Some(2), + ..Default::default() +}; + +set_print_options(print_options); +``` + +Options: + +- `precision`: Number of decimal places for floating-point numbers (default: None) +- `threshold`: Maximum number of elements to display before summarizing (default: 1000) +- `edge_items`: Number of items to show at the beginning and end of each dimension when summarizing + (default: 3) + + ### Checking Tensor Closeness + + Burn provides a utility function `check_closeness` to compare two tensors and assess their + similarity. This function is particularly useful for debugging and validating tensor operations, + especially when working with floating-point arithmetic where small numerical differences can + accumulate. It's also valuable when comparing model outputs during the process of importing models + from other frameworks, helping to ensure that the imported model produces results consistent with + the original. + + Here's an example of how to use `check_closeness`: + + ```rust, ignore + use burn::tensor::{check_closeness, Tensor}; + type B = burn::backend::Flex; + + let device = Default::default(); + let tensor1 = Tensor::::from_floats( + [1.0, 2.0, 3.0, 4.0, 5.0, 6.001, 7.002, 8.003, 9.004, 10.1], + &device, + ); + let tensor2 = Tensor::::from_floats( + [1.0, 2.0, 3.0, 4.000, 5.0, 6.0, 7.001, 8.002, 9.003, 10.004], + &device, + ); + + check_closeness(&tensor1, &tensor2); + ``` + + The `check_closeness` function compares the two input tensors element-wise, checking their + absolute differences against a range of epsilon values. It then prints a detailed report showing + the percentage of elements that are within each tolerance level. + + The output provides a breakdown for different epsilon values, allowing you to assess the closeness + of the tensors at various precision levels. This is particularly helpful when dealing with + operations that may introduce small numerical discrepancies. + + The function uses color-coded output to highlight the results: + + - Green [PASS]: All elements are within the specified tolerance. + - Yellow [WARN]: Most elements (90% or more) are within tolerance. + - Red [FAIL]: Significant differences are detected. + + This utility can be invaluable when implementing or debugging tensor operations, especially those + involving complex mathematical computations or when porting algorithms from other frameworks. It's + also an essential tool when verifying the accuracy of imported models, ensuring that the Burn + implementation produces results that closely match those of the original model. diff --git a/burn-book/src/custom-training-loop.md b/burn-book/src/custom-training-loop.md new file mode 100644 index 0000000..f77481e --- /dev/null +++ b/burn-book/src/custom-training-loop.md @@ -0,0 +1,278 @@ +# Custom Training Loops + +Even though Burn comes with a project dedicated to simplifying training, it doesn't mean that you +have to use it. Sometimes you may have special needs for your training, and it might be faster to +just reimplement the training loop yourself. Also, you may just prefer implementing your own +training loop instead of using a pre-built one in general. + +Burn's got you covered! + +We will start from the same example shown in the [basic workflow](./basic-workflow) section, but +without using the `Learner` struct. + +```rust, ignore +#[derive(Config, Debug)] +pub struct MnistTrainingConfig { + #[config(default = 10)] + pub num_epochs: usize, + #[config(default = 64)] + pub batch_size: usize, + #[config(default = 4)] + pub num_workers: usize, + #[config(default = 42)] + pub seed: u64, + #[config(default = 1e-4)] + pub lr: f64, + pub model: ModelConfig, + pub optimizer: AdamConfig, +} + +pub fn run(device: B::Device) { + // Create the configuration. + let config_model = ModelConfig::new(10, 1024); + let config_optimizer = AdamConfig::new(); + let config = MnistTrainingConfig::new(config_model, config_optimizer); + + B::seed(&device, config.seed); + + // Create the model and optimizer. + let mut model = config.model.init::(&device); + let mut optim = config.optimizer.init(); + + // Create the batcher. + let batcher = MnistBatcher::default(); + + // Create the dataloaders. + let dataloader_train = DataLoaderBuilder::new(batcher.clone()) + .batch_size(config.batch_size) + .shuffle(config.seed) + .num_workers(config.num_workers) + .build(MnistDataset::train()); + + let dataloader_test = DataLoaderBuilder::new(batcher) + .batch_size(config.batch_size) + .shuffle(config.seed) + .num_workers(config.num_workers) + .build(MnistDataset::test()); + + ... +} +``` + +As seen with the previous example, setting up the configurations and the dataloader hasn't changed. +Now, let's move forward and write our own training loop: + +```rust, ignore +pub fn run(device: B::Device) { + ... + + // Iterate over our training and validation loop for X epochs. + for epoch in 1..config.num_epochs + 1 { + // Implement our training loop. + for (iteration, batch) in dataloader_train.iter().enumerate() { + let output = model.forward(batch.images); + let loss = CrossEntropyLoss::new(None, &output.device()) + .forward(output.clone(), batch.targets.clone()); + let accuracy = accuracy(output, batch.targets); + + println!( + "[Train - Epoch {} - Iteration {}] Loss {:.3} | Accuracy {:.3} %", + epoch, + iteration, + loss.clone().into_scalar(), + accuracy, + ); + + // Gradients for the current backward pass + let grads = loss.backward(); + // Gradients linked to each parameter of the model. + let grads = GradientsParams::from_grads(grads, &model); + // Update the model using the optimizer. + model = optim.step(config.lr, model, grads); + } + + // Get the model without autodiff. + let model_valid = model.valid(); + + // Implement our validation loop. + for (iteration, batch) in dataloader_test.iter().enumerate() { + let output = model_valid.forward(batch.images); + let loss = CrossEntropyLoss::new(None, &output.device()) + .forward(output.clone(), batch.targets.clone()); + let accuracy = accuracy(output, batch.targets); + + println!( + "[Valid - Epoch {} - Iteration {}] Loss {} | Accuracy {}", + epoch, + iteration, + loss.clone().into_scalar(), + accuracy, + ); + } + } +} +``` + +In the previous code snippet, we can observe that the loop starts from epoch `1` and goes up to +`num_epochs`. Within each epoch, we iterate over the training dataloader. During this process, we +execute the forward pass, which is necessary for computing both the loss and accuracy. To maintain +simplicity, we print the results to stdout. + +Upon obtaining the loss, we can invoke the `backward()` function, which returns the gradients +specific to each variable. It's important to note that we need to map these gradients to their +corresponding parameters using the `GradientsParams` type. This step is essential because you might +run multiple different autodiff graphs and accumulate gradients for each parameter id. + +Finally, we can perform the optimization step using the learning rate, the model, and the computed +gradients. It's worth mentioning that, unlike PyTorch, there's no need to register the gradients +with the optimizer, nor do you have to call `zero_grad`. The gradients are automatically consumed +during the optimization step. If you're interested in gradient accumulation, you can easily achieve +this by using the `GradientsAccumulator`. + +```rust, ignore +let mut accumulator = GradientsAccumulator::new(); +let grads = model.backward(); +let grads = GradientsParams::from_grads(grads, &model); +accumulator.accumulate(&model, grads); ... +let grads = accumulator.grads(); // Pop the accumulated gradients. +``` + +Note that after each epoch, we include a validation loop to assess our model's performance on +previously unseen data. To disable gradient tracking during this validation step, we can invoke +`model.valid()`, which provides a model on the inner backend without autodiff capabilities. It's +important to emphasize that we've declared our validation batcher to be on the inner backend, +specifically `MnistBatcher`; not using `model.valid()` will result in a compilation +error. + +You can find the code above available as an +[example](https://github.com/tracel-ai/burn/tree/main/examples/custom-training-loop) for you to +test. + +## Multiple optimizers + +It's common practice to set different learning rates, optimizer parameters, or use different optimizers entirely, for different parts +of a model. In Burn, each `GradientParams` can contain only a subset of gradients to actually apply with an optimizer. +This allows you to flexibly mix and match optimizers! + +```rust,ignore +// Start with calculating all gradients +let grads = loss.backward(); + +// Now split the gradients into various parts. +let grads_conv1 = GradientParams::from_module(&mut grads, &model.conv1); +let grads_conv2 = GradientParams::from_module(&mut grads, &model.conv2); + +// You can step the model with these gradients, using different learning +// rates for each param. You could also use an entirely different optimizer here! +model = optim.step(config.lr * 2.0, model, grads_conv1); +model = optim.step(config.lr * 4.0, model, grads_conv2); + +// For even more granular control you can split off individual parameter +// eg. a linear bias usually needs a smaller learning rate. +if let Some(bias) == model.linear1.bias { + let grads_bias = GradientParams::from_params(&mut grads, &model.linear1, &[bias.id]); + model = optim.step(config.lr * 0.1, model, grads_bias); +} + +// Note that above calls remove gradients, so we can just get all "remaining" gradients. +let grads = GradientsParams::from_grads(grads, &model); +model = optim.step(config.lr, model, grads); +``` + +## Custom Type + +The explanations above demonstrate how to create a basic training loop. However, you may find it +beneficial to organize your program using intermediary types. There are various ways to do this, but +it requires getting comfortable with generics. + +If you wish to group the optimizer and the model into the same structure, you have several options. +It's important to note that the optimizer trait depends on both the `AutodiffModule` trait and the +`AutodiffBackend` trait, while the module only depends on the `AutodiffBackend` trait. + +Here's a closer look at how you can create your types: + +**Create a struct that is generic over the backend and the optimizer, with a predefined model.** + +```rust, ignore +struct Learner +where + B: AutodiffBackend, +{ + model: Model, + optim: O, +} +``` + +This is quite straightforward. You can be generic over the backend since it's used with the concrete +type `Model` in this case. + +**Create a struct that is generic over the model and the optimizer.** + +```rust, ignore +struct Learner { + model: M, + optim: O, +} +``` + +This option is a quite intuitive way to declare the struct. You don't need to write type constraints +with a `where` statement when defining a struct; you can wait until you implement the actual +function. However, with this struct, you may encounter some issues when trying to implement code +blocks to your struct. + +```rust, ignore +impl Learner +where + B: AutodiffBackend, + M: AutodiffModule, + O: Optimizer, +{ + pub fn step(&mut self, _batch: MnistBatch) { + // + } +} +``` + +This will result in the following compilation error: + +```console +1. the type parameter `B` is not constrained by the impl trait, self type, or predicates + unconstrained type parameter [E0207] +``` + +To resolve this issue, you have two options. The first one is to make your function generic over +the backend and add your trait constraint within its definition: + +```rust, ignore +#[allow(dead_code)] +impl Learner2 { + pub fn step(&mut self, _batch: MnistBatch) + where + B: AutodiffBackend, + M: AutodiffModule, + O: Optimizer, + { + // + } +} +``` + +However, some people may prefer to have the constraints on the implementation block itself. In that +case, you can make your struct generic over the backend using `PhantomData`. + +**Create a struct that is generic over the backend, the model, and the optimizer.** + +```rust, ignore +struct Learner3 { + model: M, + optim: O, + _b: PhantomData, +} +``` + +You might wonder why `PhantomData` is required. Each generic argument must be used as a field when +declaring a struct. When you don't need the generic argument, you can use `PhantomData` to mark it +as a zero sized type. + +These are just some suggestions on how to define your own types, but you are free to use any pattern +that you prefer. diff --git a/burn-book/src/distributed-computing.md b/burn-book/src/distributed-computing.md new file mode 100644 index 0000000..4882183 --- /dev/null +++ b/burn-book/src/distributed-computing.md @@ -0,0 +1 @@ +# Distributed Computing diff --git a/burn-book/src/examples.md b/burn-book/src/examples.md new file mode 100644 index 0000000..26e1d50 --- /dev/null +++ b/burn-book/src/examples.md @@ -0,0 +1,101 @@ +# Examples + +In the [next chapter](./basic-workflow) you'll have the opportunity to implement the whole Burn +`guide` example yourself in a step by step manner. + +Many additional Burn examples are available in the +[examples](https://github.com/tracel-ai/burn/tree/main/examples) directory. Burn examples are +organized as library crates with one or more examples that are executable binaries. An example can +then be executed using the following cargo command line in the root of the Burn repository: + +```bash +cargo run --example +``` + +To learn more about crates and examples, read the Rust section below. + +
+🦀 About Rust crates + +Each Burn example is a **package** which are subdirectories of the `examples` directory. A package +is composed of one or more **crates**. + +A package is a bundle of one or more crates that provides a set of functionality. A package contains +a `Cargo.toml` file that describes how to build those crates. + +A crate is a compilation unit in Rust. It could be a single file, but it is often easier to split up +crates into multiple **modules**. + +A module lets us organize code within a crate for readability and easy reuse. Modules also allow us +to control the _privacy_ of items. For instance the `pub(crate)` keyword is employed to make a +module publicly available inside the crate. In the snippet below there are four modules declared, +two of them are public and visible to the users of the crates, one of them is public inside the +crate only and crate users cannot see it, at last one is private when there is no keyword. These +modules can be single files or a directory with a `mod.rs` file inside. + +```rust, ignore +pub mod data; +pub mod inference; +pub(crate) mod model; +mod training; +``` + +A crate can come in one of two forms: a **binary crate** or a **library crate**. When compiling a +crate, the compiler first looks in the crate root file (`src/lib.rs` for a library crate and +`src/main.rs` for a binary crate). Any module declared in the crate root file will be inserted in +the crate for compilation. + +All Burn examples are library crates and they can contain one or more executable examples that uses +the library. We even have some Burn examples that uses the library crate of other examples. + +The examples are unique files under the `examples` directory. Each file produces an executable file +with the same name. Each example can then be executed with `cargo run --example `. + +Below is a file tree of a typical Burn example package: + +``` +examples/burn-example +├── Cargo.toml +├── examples +│ ├── example1.rs ---> compiled to example1 binary +│ ├── example2.rs ---> compiled to example2 binary +│ └── ... +└── src + ├── lib.rs ---> this is the root file for a library + ├── module1.rs + ├── module2.rs + └── ... +``` + +

+ +The following additional examples are currently available if you want to check them out: + +| Example | Description | +| :-------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| [Custom CSV Dataset](https://github.com/tracel-ai/burn/tree/main/examples/custom-csv-dataset) | Implements a dataset to parse CSV data for a regression task. | +| [Regression](https://github.com/tracel-ai/burn/tree/main/examples/simple-regression) | Trains a simple MLP on the California Housing dataset to predict the median house value for a district. | +| [Custom Image Dataset](https://github.com/tracel-ai/burn/tree/main/examples/custom-image-dataset) | Trains a simple CNN on custom image dataset following a simple folder structure. | +| [Custom Renderer](https://github.com/tracel-ai/burn/tree/main/examples/custom-renderer) | Implements a custom renderer to display the [`Learner`](./building-blocks/learner.md) progress. | +| [Image Classification Web](https://github.com/tracel-ai/burn-onnx/tree/main/examples/image-classification-web) | Image classification web browser demo using Burn, WGPU and WebAssembly. | +| [MNIST Inference on Web](https://github.com/tracel-ai/burn/tree/main/examples/mnist-inference-web) | An interactive MNIST inference demo in the browser. The demo is available [online](https://burn.dev/demo/). | +| [MNIST Training](https://github.com/tracel-ai/burn/tree/main/examples/mnist) | Demonstrates how to train a custom [`Module`](./building-blocks/module.md) (MLP) with the [`Learner`](./building-blocks/learner.md) configured to log metrics and keep training checkpoints. | +| [ONNX Import Inference](https://github.com/tracel-ai/burn-onnx/tree/main/examples/onnx-inference) | Imports an ONNX model pre-trained on MNIST to perform inference on a sample image with Burn. | +| [PyTorch Import Inference](https://github.com/tracel-ai/burn/tree/main/examples/import-model-weights) | Imports a PyTorch model pre-trained on MNIST to perform inference on a sample image with Burn. | +| [Text Classification](https://github.com/tracel-ai/burn/tree/main/examples/text-classification) | Trains a text classification transformer model on the AG News or DbPedia datasets. The trained model can then be used to classify a text sample. | +| [Text Generation](https://github.com/tracel-ai/burn/tree/main/examples/text-generation) | Trains a text generation transformer model on the DbPedia dataset. | +| [Wasserstein GAN MNIST](https://github.com/tracel-ai/burn/tree/main/examples/wgan) | Trains a WGAN model to generate new handwritten digits based on MNIST. | + +For more information on each example, see their respective `README.md` file. Be sure to check out +the [examples](https://github.com/tracel-ai/burn/tree/main/examples) directory for an up-to-date +list. + +
+ +Note that some examples use the +[`datasets` library by HuggingFace](https://huggingface.co/docs/datasets/index) to download the +datasets required in the examples. This is a Python library, which means that you will need to +install Python before running these examples. This requirement will be clearly indicated in the +example's README when applicable. + +
diff --git a/burn-book/src/getting-started.md b/burn-book/src/getting-started.md new file mode 100644 index 0000000..5316954 --- /dev/null +++ b/burn-book/src/getting-started.md @@ -0,0 +1,219 @@ +# Getting Started + +Burn is a deep learning framework in the Rust programming language. Therefore, it goes without +saying that one must understand the basic notions of Rust. Reading the first chapters of the +[Rust Book](https://doc.rust-lang.org/book/) is recommended, but don't worry if you're just starting +out. We'll try to provide as much context and reference to external resources when required. Just +look out for the **🦀 Rust Note** indicators. + +## Installing Rust + +For installation instructions, please refer to the +[installation page](https://doc.rust-lang.org/book/ch01-01-installation.html). It explains in +details the most convenient way for you to install Rust on your computer, which is the very first +thing to do to start using Burn. + +## Creating a Burn application + +Once Rust is correctly installed, create a new Rust application by using Rust's build system and +package manager Cargo. It is automatically installed with Rust. + +
+🦀 Cargo Cheat Sheet + +[Cargo](https://doc.rust-lang.org/cargo/) is a very useful tool to manage Rust projects because it +handles a lot of tasks. More precisely, it is used to compile your code, download the +libraries/packages your code depends on, and build said libraries. + +Below is a quick cheat sheet of the main `cargo` commands you might use throughout this guide. + +| Command | Description | +| ------------------- | -------------------------------------------------------------------------------------------- | +| `cargo new` _path_ | Create a new Cargo package in the given directory. | +| `cargo add` _crate_ | Add dependencies to the Cargo.toml manifest file. | +| `cargo build` | Compile the local package and all of its dependencies (in debug mode, use `-r` for release). | +| `cargo check` | Check the local package for compilation errors (much faster). | +| `cargo run` | Run the local package binary. | + +For more information, check out +[Hello, Cargo!](https://doc.rust-lang.org/book/ch01-03-hello-cargo.html) in the Rust Book. + +

+ +In the directory of your choice, run the following: + +```console +cargo new my_burn_app +``` + +This will initialize the `my_burn_app` project directory with a `Cargo.toml` file and a `src` +directory with an auto-generated `main.rs` file inside. Head inside the directory to check: + +```console +cd my_burn_app +``` + +Then, add Burn as a dependency: + +```console +cargo add burn --features wgpu +``` + +Finally, compile the local package by executing the following: + +```console +cargo build +``` + +That's it, you're ready to start! You have a project configured with Burn and the WGPU backend, +which allows to execute low-level operations on any platform using the GPU. + +
+ +When using one of the `wgpu` backends, you may encounter compilation errors related to recursive +type evaluation. This is due to complex type nesting within the `wgpu` dependency chain. + +To resolve this issue, add the following line at the top of your `main.rs` or `lib.rs` file: + +```rust +#![recursion_limit = "256"] +``` + +The default recursion limit (128) is often just below the required depth (typically 130-150) due to +deeply nested associated types and trait bounds. + +
+ +## Writing a code snippet + +The `src/main.rs` was automatically generated by Cargo, so let's replace its content with the +following: + +```rust, ignore +use burn::tensor::Tensor; +use burn::backend::Wgpu; + +// Type alias for the backend to use. +type Backend = Wgpu; + +fn main() { + let device = Default::default(); + // Creation of two tensors, the first with explicit values and the second one with ones, with the same shape as the first + let tensor_1 = Tensor::::from_data([[2., 3.], [4., 5.]], &device); + let tensor_2 = Tensor::::ones_like(&tensor_1); + + // Print the element-wise addition (done with the WGPU backend) of the two tensors. + println!("{}", tensor_1 + tensor_2); +} +``` + +
+🦀 Use Declarations + +To bring any of the Burn module or item into scope, a `use` declaration is added. + +In the example above, we wanted bring the `Tensor` struct and `Wgpu` backend into scope with the +following: + +```rust, ignore +use burn::tensor::Tensor; +use burn::backend::Wgpu; +``` + +This is pretty self-explanatory in this case. But, the same declaration could be written as a +shortcut to simultaneously binding of multiple paths with a common prefix: + +```rust, ignore +use burn::{tensor::Tensor, backend::Wgpu}; +``` + +In this example, the common prefix is pretty short and there are only two items to bind locally. +Therefore, the first usage with two `use` declarations might be preferred. But know that both +examples are valid. For more details on the `use` keyword, take a look at +[this section](https://doc.rust-lang.org/book/ch07-04-bringing-paths-into-scope-with-the-use-keyword.html) +of the Rust Book or the +[Rust reference](https://doc.rust-lang.org/reference/items/use-declarations.html). + +

+ +
+🦀 Generic Data Types + +If you're new to Rust, you're probably wondering why we had to use `Tensor::::...`. +That's because the `Tensor` struct is [generic](https://doc.rust-lang.org/book/ch10-01-syntax.html) +over multiple concrete data types. More specifically, a `Tensor` can be defined using three generic +parameters: the backend, the number of dimensions (rank) and the data type (defaults to `Float`). +Here, we only specify the backend and number of dimensions since a `Float` tensor is used by +default. For more details on the `Tensor` struct, take a look at +[this section](./building-blocks/tensor.md). + +Most of the time when generics are involved, the compiler can infer the generic parameters +automatically. In this case, the compiler needs a little help. This can usually be done in one of +two ways: providing a type annotation or binding the generic parameter via the _turbofish_ `::<>` +syntax. In the example above we used the so-called _turbofish_ syntax, but we could have used type +annotations instead like this: + +```rust, ignore +let tensor_1: Tensor = Tensor::from_data([[2., 3.], [4., 5.]]); +let tensor_2 = Tensor::ones_like(&tensor_1); +``` + +You probably noticed that we provided a type annotation for the first tensor only and yet this +example still works. That's because the compiler (correctly) inferred that `tensor_2` had the same +generic parameters. The same could have been done in the original example, but specifying the +parameters for both is more explicit. + +

+ +By running `cargo run`, you should now see the result of the addition: + +```console +Tensor { + data: +[[3.0, 4.0], + [5.0, 6.0]], + shape: [2, 2], + device: DefaultDevice, + backend: "wgpu", + kind: "Float", + dtype: "f32", +} +``` + +While the previous example is somewhat trivial, the upcoming basic workflow section will walk you +through a much more relevant example for deep learning applications. + +## Using `prelude` + +Burn comes with a variety of things in its core library. When creating a new model or using an +existing one for inference, you may need to import every single component you used, which could be a +little verbose. + +To address it, a `prelude` module is provided, allowing you to easily import commonly used structs +and macros as a group: + +```rust, ignore +use burn::prelude::*; +``` + +which is equal to: + +```rust, ignore +use burn::{ + config::Config, + module::Module, + nn, + tensor::{ + backend::Backend, Bool, Device, ElementConversion, Float, Int, Shape, Tensor, + TensorData, + }, +}; +``` + +
+ +For the sake of simplicity, the subsequent chapters of this book will all use this form of importing +except in the [Building Blocks](./building-blocks) chapter, as explicit importing aids users in +grasping the usage of particular structures and macros. + +
diff --git a/burn-book/src/models-and-pretrained-weights.md b/burn-book/src/models-and-pretrained-weights.md new file mode 100644 index 0000000..4942fbe --- /dev/null +++ b/burn-book/src/models-and-pretrained-weights.md @@ -0,0 +1,30 @@ +# Models and Pre-Trained Weights + +## Models Repository + +The [`models`](https://github.com/tracel-ai/models) repository contains definitions of different +deep learning models with examples for different domains like computer vision and natural language +processing. + +This includes image classification models such as +[`MobileNetV2`](https://github.com/tracel-ai/models/tree/main/mobilenetv2-burn), +[`SqueezeNet`](https://github.com/tracel-ai/models/tree/main/squeezenet-burn) and +[`ResNet`](https://github.com/tracel-ai/models/tree/main/resnet-burn), object detection models such +as [`YOLOX`](https://github.com/tracel-ai/models/tree/main/yolox-burn) and language models like +[`BERT` and `RoBERTa`](https://github.com/tracel-ai/models/tree/main/bert-burn). + +Be sure to check out the up-to-date +[collection of models](https://github.com/tracel-ai/models?tab=readme-ov-file#collection-of-official-models) +to get you started. Pre-trained weights are available for every supported architecture in this +collection. You will also find a spotlight of +[community contributed models](https://github.com/tracel-ai/models?tab=readme-ov-file#community-contributions). + +## Burn-LM (alpha) + +[`Burn-LM`](https://github.com/tracel-ai/burn-lm) is an LLM inference engine built on Burn. It +provides access to large language models with open-source pre-trained weights and supports running, +fine-tuning, and experimenting with them on any Burn backend. + +Unlike tools focused solely on inference, Burn-LM is designed to work in a unified way across +different models and tasks, making it easier to explore both inference and training workflows within +the same framework. diff --git a/burn-book/src/motivation.md b/burn-book/src/motivation.md new file mode 100644 index 0000000..169491a --- /dev/null +++ b/burn-book/src/motivation.md @@ -0,0 +1,28 @@ +# Why Burn? + +Why bother with the effort of creating an entirely new deep learning framework from scratch when +PyTorch, TensorFlow, and other frameworks already exist? Spoiler alert: Burn isn't merely a +replication of PyTorch or TensorFlow in Rust. It represents a novel approach, placing significant +emphasis on making the right compromises in the right areas to facilitate exceptional flexibility, +high performance, and a seamless developer experience. Burn isn’t a framework specialized for only +one type of application, it is designed to serve as a versatile framework suitable for a wide range +of research and production uses. The foundation of Burn's design revolves around three key user +profiles: + +**Machine Learning Researchers** require tools to construct and execute experiments efficiently. +It’s essential for them to iterate quickly on their ideas and design testable experiments which can +help them discover new findings. The framework should facilitate the swift implementation of +cutting-edge research while ensuring fast execution for testing. + +**Machine Learning Engineers** are another important demographic to keep in mind. Their focus leans +less on swift implementation and more on establishing robustness, seamless deployment, and +cost-effective operations. They seek dependable, economical models capable of achieving objectives +without excessive expense. The whole machine learning workflow —from training to inference— must be +as efficient as possible with minimal unpredictable behavior. + +**Low level Software Engineers** working with hardware vendors want their processing units to run +models as fast as possible to gain competitive advantage. This endeavor involves harnessing +hardware-specific features such as Tensor Core for Nvidia. Since they are mostly working at a system +level, they want to have absolute control over how the computation will be executed. + +The goal of Burn is to satisfy all of those personas! diff --git a/burn-book/src/onnx-import.md b/burn-book/src/onnx-import.md new file mode 100644 index 0000000..1d048fb --- /dev/null +++ b/burn-book/src/onnx-import.md @@ -0,0 +1,278 @@ +# ONNX Import + +## Introduction + +As deep learning evolves, interoperability between frameworks becomes crucial. Burn provides robust +support for importing [ONNX (Open Neural Network Exchange)](https://onnx.ai/onnx/intro/index.html) +models through the [`burn-onnx`](https://github.com/tracel-ai/burn-onnx) crate, enabling you to +leverage pre-trained models in your Rust-based deep learning projects. + +## Why Import Models? + +Importing pre-trained models offers several advantages: + +1. **Time-saving**: Skip the resource-intensive process of training models from scratch. +2. **Access to state-of-the-art architectures**: Utilize cutting-edge models developed by + researchers and industry leaders. +3. **Transfer learning**: Fine-tune imported models for your specific tasks, benefiting from + knowledge transfer. +4. **Consistency across frameworks**: Maintain consistent performance when moving between + frameworks. + +## Understanding ONNX + +ONNX (Open Neural Network Exchange) is an open format designed to represent machine learning models +with these key features: + +- **Framework agnostic**: Provides a common format that works across various deep learning + frameworks. +- **Comprehensive representation**: Captures both the model architecture and trained weights. +- **Wide support**: Compatible with popular frameworks like PyTorch, TensorFlow, and scikit-learn. + +This standardization allows seamless movement of models between different frameworks and deployment +environments. + +## Burn's ONNX Support + +Burn's approach to ONNX import offers unique advantages: + +1. **Native Rust code generation**: Translates ONNX models into Rust source code for deep + integration with Burn's ecosystem. +2. **Compile-time optimization**: Leverages the Rust compiler to optimize the generated code, + potentially improving performance. +3. **No runtime dependency**: Eliminates the need for an ONNX runtime, unlike many other solutions. +4. **Trainability**: Allows imported models to be further trained or fine-tuned using Burn. +5. **Portability**: Enables compilation for various targets, including WebAssembly and embedded + devices. +6. **Backend flexibility**: Works with any of Burn's supported backends. + +## ONNX Compatibility + +Burn recommends ONNX models use **opset version 16 or higher** for best compatibility. While models +with older opset versions may work, opset 16+ ensures access to all supported operators and their +latest behavior. If you encounter issues with an older model, consider upgrading it using the ONNX +version converter. + +### Upgrading ONNX Models + +There are two simple ways to upgrade your ONNX models to the recommended opset version: + +Option 1: Use the provided utility script: + +``` +uv run --script https://raw.githubusercontent.com/tracel-ai/burn-onnx/refs/heads/main/onnx_opset_upgrade.py +``` + +Option 2: Use a custom Python script: + +```python +import onnx +from onnx import version_converter, shape_inference + +# Load your ONNX model +model = onnx.load('path/to/your/model.onnx') + +# Convert the model to opset version 16 +upgraded_model = version_converter.convert_version(model, 16) + +# Apply shape inference to the upgraded model +inferred_model = shape_inference.infer_shapes(upgraded_model) + +# Save the converted model +onnx.save(inferred_model, 'upgraded_model.onnx') +``` + +## Step-by-Step Guide + +Follow these steps to import an ONNX model into your Burn project: + +### Step 1: Update `Cargo.toml` + +First, add the required dependencies to your `Cargo.toml`: + +```toml +[dependencies] +burn = { version = "~0.21", features = ["flex"] } + +[build-dependencies] +burn-onnx = "~0.21" +``` + +### Step 2: Update `build.rs` + +In your `build.rs` file: + +```rust, ignore +use burn_onnx::ModelGen; + +fn main() { + ModelGen::new() + .input("src/model/my_model.onnx") + .out_dir("model/") + .run_from_script(); +} +``` + +This generates Rust code and a `.bpk` weights file from your ONNX model during the build process. + +### Step 3: Modify `mod.rs` + +In your `src/model/mod.rs` file, include the generated code: + +```rust, ignore +pub mod my_model { + include!(concat!(env!("OUT_DIR"), "/model/my_model.rs")); +} +``` + +### Step 4: Use the Imported Model + +Now you can use the imported model in your code: + +```rust, ignore +use burn::tensor; +use burn::backend::{Flex, flex::FlexDevice}; +use model::my_model::Model; + +fn main() { + let device = FlexDevice; + + // Create model instance and load weights from target dir default device + let model: Model = Model::default(); + + // Create input tensor (replace with your actual input) + let input = tensor::Tensor::::zeros([1, 3, 224, 224], &device); + + // Perform inference + let output = model.forward(input); + + println!("Model output: {:?}", output); +} +``` + +## Advanced Configuration + +The `ModelGen` struct provides configuration options: + +```rust, ignore +use burn_onnx::{ModelGen, LoadStrategy}; + +ModelGen::new() + .input("path/to/model.onnx") + .out_dir("model/") + .development(true) // Enable development mode for debugging + .load_strategy(LoadStrategy::Embedded) // Embed weights in the binary + .run_from_script(); +``` + +- `input`: Path to the ONNX model file +- `out_dir`: Output directory for generated code and weights +- `development`: When enabled, generates additional debug files (`.onnx.txt`, `.graph.txt`) +- `load_strategy`: Controls which weight-loading constructors are generated on the `Model` struct + (see below) + +Model weights are stored in Burnpack format (`.bpk`), which provides efficient serialization and +loading. + +### Load Strategy + +The `LoadStrategy` enum controls how the generated model loads its weights: + +| Strategy | Generated constructors | `Default` impl | Use case | +|------------|------------------------------------------------|-----------------|-------------------------------------------| +| `File` | `from_file()`, `from_bytes()` | Yes | Standard desktop/server (default) | +| `Embedded` | `from_embedded()`, `from_bytes()` | Yes | Single binary, small models | +| `Bytes` | `from_bytes()` | No | WASM, embedded, custom loaders | +| `None` | (none) | No | Manual weight management | + +The default strategy is `File`, which keeps weights in a separate `.bpk` file and generates a +`from_file()` constructor. + +For WebAssembly or environments without filesystem access, use `LoadStrategy::Bytes`: + +```rust, ignore +ModelGen::new() + .input("model.onnx") + .out_dir("model/") + .load_strategy(LoadStrategy::Bytes) + .run_from_script(); +``` + +Then load weights at runtime from any byte source (e.g., a network fetch): + +```rust, ignore +let model = Model::::from_bytes(weight_bytes, &device); +``` + +## Loading and Using Models + +You can load models in several ways, depending on the `LoadStrategy` used during code generation: + +```rust, ignore +// Load from the output directory with default device (recommended for most use cases) +// This automatically loads weights from the .bpk file +// Available with LoadStrategy::File or LoadStrategy::Embedded +let model = Model::::default(); + +// Create a new model instance with a specific device +// (initializes weights randomly; load weights via `load_from` afterward) +let model = Model::::new(&device); + +// Load from a specific .bpk file (LoadStrategy::File) +let model = Model::::from_file("path/to/weights.bpk", &device); + +// Load from in-memory bytes (LoadStrategy::File, Embedded, or Bytes) +let model = Model::::from_bytes(weight_bytes, &device); + +// Load from embedded weights (LoadStrategy::Embedded) +let model = Model::::from_embedded(&device); +``` + +## Troubleshooting + +Common issues and solutions: + +1. **Unsupported ONNX operator**: Check the + [list of supported ONNX operators](https://github.com/tracel-ai/burn-onnx/blob/main/SUPPORTED-ONNX-OPS.md). + You may need to simplify your model or wait for support. + +2. **Build errors**: Ensure your `burn-onnx` version matches your Burn version and verify the ONNX + file path in `build.rs`. + +3. **Runtime errors**: Confirm that your input tensors match the expected shape and data type of + your model. + +4. **Performance issues**: Consider using a more performant backend or optimizing your model + architecture. + +5. **Viewing generated files**: Find the generated Rust code and weights in the `OUT_DIR` directory + (usually `target/debug/build//out`). + +## Examples and Resources + +For practical examples, check out the +[burn-onnx examples](https://github.com/tracel-ai/burn-onnx/tree/main/examples): + +1. [ONNX Inference](https://github.com/tracel-ai/burn-onnx/tree/main/examples/onnx-inference) - + MNIST inference example +2. [Image Classification Web](https://github.com/tracel-ai/burn-onnx/tree/main/examples/image-classification-web) - + SqueezeNet running in the browser via WebAssembly +3. [Raspberry Pi Pico](https://github.com/tracel-ai/burn-onnx/tree/main/examples/raspberry-pi-pico) - + Embedded deployment example + +These demonstrate real-world usage of ONNX import in Burn projects. + +For contributors looking to add support for new ONNX operators: + +- [Development Guide](https://github.com/tracel-ai/burn-onnx/blob/main/DEVELOPMENT-GUIDE.md) - + Step-by-step guide for implementing new operators + +## Conclusion + +Importing ONNX models into Burn combines the vast ecosystem of pre-trained models with Burn's +performance and Rust's safety features. Following this guide, you can seamlessly integrate ONNX +models into your Burn projects for inference, fine-tuning, or further development. + +The `burn-onnx` crate is actively developed, with ongoing work to support more ONNX operators and +improve performance. Visit the [burn-onnx repository](https://github.com/tracel-ai/burn-onnx) for +updates and to contribute! diff --git a/burn-book/src/overview.md b/burn-book/src/overview.md new file mode 100644 index 0000000..aeb2e28 --- /dev/null +++ b/burn-book/src/overview.md @@ -0,0 +1,37 @@ +# Overview + +Welcome to The Burn Book 👋 + +This book will help you get started with the Burn deep learning framework, whether you are an +advanced user or a beginner. We have crafted some sections for you: + +- [Basic Workflow: From Training to Inference](./basic-workflow): We'll start with the fundamentals, + guiding you through the entire workflow, from training your models to deploying them for + inference. This section lays the groundwork for your Burn expertise. + +- [Building Blocks](./building-blocks): Dive deeper into Burn's core components, understanding how + they fit together. This knowledge forms the basis for more advanced usage and customization. + +- [Performance - Good Practices](./performance/good-practices/): Tips for writing models and + training code that make the most of hardware resources while avoiding common pitfalls that can + slow down execution. + +- [Custom Training Loop](./custom-training-loop.md): Gain the power to customize your training + loops, fine-tuning your models to meet your specific requirements. This section empowers you to + harness Burn's flexibility to its fullest. + +- [Saving & Loading Models](./saving-and-loading.md): Learn how to save and load your trained + models, including importing weights from PyTorch and SafeTensors formats. + +- [ONNX Import](./onnx-import.md): Learn how to import ONNX models using the + [burn-onnx](https://github.com/tracel-ai/burn-onnx) crate. + +- [Models & Pre-Trained Weights](./models-and-pretrained-weights.md): Get started quickly with + ready-to-use models and pre-trained weights. + +- [Advanced](./advanced): Finally, venture into advanced topics, exploring Burn's capabilities at + their peak. This section caters to those who want to push the boundaries of what's possible with + Burn. + +Throughout the book, we assume a basic understanding of deep learning concepts, but we may refer to +additional material when it seems appropriate. diff --git a/burn-book/src/performance/README.md b/burn-book/src/performance/README.md new file mode 100644 index 0000000..c42261e --- /dev/null +++ b/burn-book/src/performance/README.md @@ -0,0 +1,4 @@ +# Performance + +This section covers the key concepts you need to understand to get the most out of Burn and your +hardware. diff --git a/burn-book/src/performance/distributed-computing.md b/burn-book/src/performance/distributed-computing.md new file mode 100644 index 0000000..32da6f9 --- /dev/null +++ b/burn-book/src/performance/distributed-computing.md @@ -0,0 +1,4 @@ +# Distributed Computing + +Distributed computing support was introduced in Burn 0.19. Documentation and examples will be +available soon. diff --git a/burn-book/src/performance/good-practices/README.md b/burn-book/src/performance/good-practices/README.md new file mode 100644 index 0000000..2c0ac76 --- /dev/null +++ b/burn-book/src/performance/good-practices/README.md @@ -0,0 +1,12 @@ +# Performance - Best Practices + +This section provides valuable insights into the performance characteristics of Burn and guides +users on how to effectively leverage them for optimal results. + +It includes several sections, each offering relevant details. While understanding these concepts can +aid in model optimization, it’s always crucial to conduct benchmarks and profile models to +accurately assess performance improvements. + +- [Asynchronous Execution](./asynchronous-execution.md) +- [Kernel Fusion](./kernel-fusion.md) +- [Kernel Selection](./kernel-selection.md) diff --git a/burn-book/src/performance/good-practices/asynchronous-execution.md b/burn-book/src/performance/good-practices/asynchronous-execution.md new file mode 100644 index 0000000..2163eb8 --- /dev/null +++ b/burn-book/src/performance/good-practices/asynchronous-execution.md @@ -0,0 +1,62 @@ +# Asynchronous Execution + +Most Burn backends execute tensor operations in an asynchronous manner. However, the async notation +is often not required for most tensor operations, privileging the simplicity of sync Rust. + +There are only a few operations that trigger synchronization of the backend, and it is very +important to correctly handle those to optimize hardware utilization. Those operations are +`into_data`, `into_scalar`, and `sync`. Some tensor operations might call `into_data` underneath, +triggering a synchronization, like `to_device` for some backends. + +There are several ways to minimize synchronization overhead, one of which is to batch sync +operations into a single transaction. Burn provides a high-level composable API to build +transactions, which will only trigger a single sync on the device. + +For instance, it is often used when collecting metrics during training: + +```rust +/// All of these variables are tensors. +let (output, loss, targets) = ..; + +/// Now output, loss, and targets will be `TensorData` stored on the CPU. +let [output, loss, targets] = Transaction::default() + .register(output) + .register(loss) + .register(targets) + .execute() + .try_into() + .expect("Correct amount of tensor data"); +``` + +Another way of optimizing reads and avoiding device stalls is to read the data on a different +thread. Under the hood, CubeCL-based backends assign different execution queues for different +threads, meaning that syncing a thread shouldn’t impact the throughput of another thread. + +## Using Different Backends for Different Tasks + +Tensor operations aren’t the only things that are asynchronous; dataset and dataloading are also +lazily executed. This allows for efficient data augmentation and sampling without having to cache +huge datasets on disk. However, this might reduce training throughput if data augmentation is +performed on the same device as the training itself. So, it is normally encouraged to use a +different device, maybe even a different backend, for that purpose. For optimal performance, also +avoid small allocations followed by a batching procedure. Even if it doesn’t break asynchronicity, +it can slow down performance. + +```rust +/// Items is a vector of many tensors. +let items = ..; +let batch = Tensor::cat(items, 1); +``` + +Prefer doing the concatenation of tensors on the data augmentation device and not on the training +device. + +```rust +/// Items is a vector of many tensors. +let items = ..; +let device_training = ..; +let axis_batch = 0; + +let items = Tensor::cat(items, axis_batch); +let batch = Tensor::from_data(items.into_data(), device_training); +``` diff --git a/burn-book/src/performance/good-practices/kernel-fusion.md b/burn-book/src/performance/good-practices/kernel-fusion.md new file mode 100644 index 0000000..7ae9d6f --- /dev/null +++ b/burn-book/src/performance/good-practices/kernel-fusion.md @@ -0,0 +1,76 @@ +# Kernel Fusion + +An interesting property of async execution is that it allows performance optimizations like kernel +fusion. Coupled with CubeCL and its Just-In-Time compiler, Burn can serialize tensor operations into +a symbolic graph, then optimize it for improved efficiency. + +Kernel fusion may reorder operations to reduce global memory reads, writes, and allocations. Being +aware of which operations can be fused is relevant, as it can be easy to break an execution graph. + +The easiest way to optimize for fusion is to avoid keeping tensors alive for too long. When fusion +isn’t possible, all tensors that will be used later will trigger a global memory write. Fortunately, +Rust and Clippy are quite good at detecting unnecessary clones, but special care should still be +taken. + +View operations can also interfere with fusion. They can be included in optimized graphs, but only +to a limited extent, and they reduce vectorization potential as we have fewer guarantees about +memory access patterns with transformed indices. So, it is good practice to group view operations +together before executing blocks of operations. + +```rust +let tensor4 = tensor1.unsqueeze().matmul(tensor2) + tensor3.unsqueeze(); +``` + +Could be improved with the following: + +```rust +let tensor1 = tensor1.unsqueeze(); +let tensor3 = tensor3.unsqueeze(); +let tensor4 = tensor1.matmul(tensor2) + tensor3; +``` + +This reduces the necessary reordering and may reduce a global memory write or improve vectorization. +We might be able to detect these patterns in the future, but for now, it’s a good idea to order your +operations using this pattern. As a reminder, view operations typically only update tensor metadata +in most cases. These operations include `slice`, `slice_assign`, `select`, `gather`, `scatter`, +`reshape`, `swap_dims`, `transpose`, `unsqueeze`, etc. + +With fusion enabled, it is often not necessary to write custom kernels, as you can rely on our +system to optimize most element-wise operations. However, most compute-bound kernels require many +tricks and deep knowledge of GPU memory architectures, where automatic compiler optimizations often +underperform compared to human-designed algorithms. This is why Burn’s approach to fusion is +centered around fuse-on-read and fuse-on-write. This means that complex compute-bound kernels that +change the shapes of tensors can fuse a block of element-wise operations when reading the input +tensor and when writing the output tensor. The implication is that multiple compute-bound operations +in a sequence can reduce fusion potential. + +```rust +// This line might trigger 3 writes: tensor1, tensor2, and tensor3, if tensor1 and tensor2 are abstract tensors. +let tensor3 = tensor1.clone().sum_dim(tensor2.clone(), 2); +let tensor4 = tensor2.sum_dim(tensor3, 2); +let tensor5 = tensor4 + (tensor1 * tensor2); +``` + +```rust +let tmp = tensor1.clone() + tensor2.clone(); +let tensor3 = tensor1.sum_dim(tensor2, 2); +let tensor4 = tensor2.sum_dim(tensor3, 2); +let tensor5 = tensor4 + tmp; +``` + +The lesson? Whenever possible, pass only the latest value to a compute operation. Don’t clone a +tensor before compute-bound operations, as it might trigger an additional write if that tensor isn’t +materialized from initial fusion. + +It’s a bit complex, but the first code snippet is actually better if `tensor1` and `tensor2` are +concrete in global memory. This would be the case if `tensor1` and `tensor2` are model parameters, +so prefer this implementation style in such scenarios. + +The second code snippet is preferred when `tensor1` and `tensor2` are virtual tensors, meaning they +were fused by earlier operations and require a global memory read to be accessed later. This happens +if those tensors are part of a signal in neural networks. + +Reordering operations can help in such scenarios but will not create temporary values, making the +previous optimization harder. We might eventually automatically optimize these cases, but the +solution space is quite large, and it’s not a planned optimization. Profiling model blocks is always +a good idea to identify which code block is faster when faced with ambiguous situations. diff --git a/burn-book/src/performance/good-practices/kernel-selection.md b/burn-book/src/performance/good-practices/kernel-selection.md new file mode 100644 index 0000000..698f8cf --- /dev/null +++ b/burn-book/src/performance/good-practices/kernel-selection.md @@ -0,0 +1,25 @@ +# Kernel Selection + +As mentioned earlier, complex compute-bound operations are highly non-trivial and require many +tricks for optimal performance. However, the way these tricks are applied varies depending on the +hardware and problem shapes. To select the best kernel, we use a search method with a highly +configurable autotune system that performs micro-benchmarks at runtime on the current hardware. + +This may trigger a cold start, but the results of these benchmarks are cached on disk for subsequent +executions. + +For deployment or training on spot instances, it’s a good idea to bundle the autotune cache with the +code to mitigate cold starts. Refer to the +[CubeCL configuration documentation](https://burn.dev/books/cubecl/advanced-usage/config.html) for +more details on fine-grained settings . + +From the user’s point of view, kernel selection shouldn’t be a problem, but as usual, crafting +models with even shapes, multiples of 8, can significantly improve performance. Avoid creating +tensors with shapes that are multiples of 10, like `[1000, 1000]`, as these typically require bounds +checking and may limit vectorization. + +Prefer shapes like `[1024, 1024]`, where dimensions are multiples of 32 or powers of 2, as these are +generally optimal. If you have no choice but to use a suboptimal shape, prefer handling it in a +single kernel, transforming it into an optimal shape. It’s better to have a slow neural network +layer followed by fast ones than to propagate unevenness and end up with smaller, but slower, +layers. diff --git a/burn-book/src/performance/quantization.md b/burn-book/src/performance/quantization.md new file mode 100644 index 0000000..f1b48be --- /dev/null +++ b/burn-book/src/performance/quantization.md @@ -0,0 +1,133 @@ +# Quantization + +Quantization techniques perform computations and store tensors in lower precision data types like +8-bit integer instead of floating point precision. There are multiple approaches to quantize a deep +learning model categorized as: + +- Post-training quantization (PTQ) +- Quantization aware training (QAT) + +In post-training quantization, the model is trained in floating point precision and later converted +to the lower precision data type. There are two types of post-training quantization: + +1. Static quantization: quantizes the weights and activations of the model. Quantizing the + activations statically requires data to be calibrated (i.e., recording the activation values to + compute the optimal quantization parameters with representative data). +1. Dynamic quantization: quantized the weights ahead of time (like static quantization) but the + activations are dynamically at runtime. + +Sometimes post-training quantization is not able to achieve acceptable task accuracy. In general, +this is where quantization-aware training (QAT) can be used: during training, fake-quantization +modules are inserted in the forward and backward passes to simulate quantization effects, allowing +the model to learn representations that are more robust to reduced precision. + +Burn does not currently support QAT. Only post-training quantization (PTQ) is implemented at this +time. + +
+ +Quantization support in Burn is currently in active development. + +It supports the following PTQ modes on some backends: + +- Per-tensor and per-block quantization to 8-bit, 4-bit and 2-bit representations + +No integer operations are currently supported, which means tensors are dequantized to perform the +operations in floating point precision. + +
+ +## Module Quantization + +Quantizing the weights of your model after training is quite simple. We have access to the weight +tensors and can collect their statistics, such as the min and max value when using +`MinMaxCalibration`, to compute the quantization parameters. + +```rust , ignore +# use burn::module::Quantizer; +# use burn::tensor::quantization::{Calibration, QuantLevel, QuantParam, QuantScheme, QuantValue}; +# +// Quantization config +let scheme = QuantScheme::default() + .with_level(QuantLevel::Block(32)) + .with_value(QuantValue::Q4F) + .with_param(QuantParam::F16); +let mut quantizer = Quantizer { + calibration: Calibration::MinMax, + scheme, +}; + +// Quantize the weights +let model = model.quantize_weights(&mut quantizer); +``` + +### Calibration + +Calibration is the step during quantization where the range of all floating-point tensors is +computed. This is pretty straightforward for weights since the actual range is known at +_quantization-time_ (weights are static), but activations require more attention. + +To compute the quantization parameters, Burn supports the following `Calibration` methods. + +| Method | Description | +| :------- | :------------------------------------------------------------------------------- | +| `MinMax` | Computes the quantization range mapping based on the running min and max values. | + +### Quantization Scheme + +A quantization scheme defines how an input is quantized, including the representation of quantized +values, storage format, granularity, and how the values are scaled. + +```rust +let scheme = QuantScheme::default() + .with_mode(QuantMode::Symmetric) // Quantization mode + .with_level(QuantLevel::block([2, 16])) // Granularity (per-tensor or per-block) + .with_value(QuantValue::Q8S) // Data type of quantized values, independent of how they're stored + .with_store(QuantStore::Native) // Storage format for quantized values + .with_param(QuantParam::F16); // Precision for quantization parameters +``` + +#### Quantization Mode + +| Mode | Description | +| :---------- | :------------------------------------------- | +| `Symmetric` | Values are scaled symmetrically around zero. | + +#### Quantization Level + +| Level | Description | +| :----------------------------- | :----------------------------------------------------------------------------------------------------------- | +| `Tensor` | A single quantization parameter set for the entire tensor. | +| `Block(block_size: BlockSize)` | Tensor divided into blocks (1D, 2D, or higher) defined by block_size, each with its own quantization params. | + +#### Quantization Value + +| Value | Bits | Description | +| :----- | :--: | :-------------------------------------------- | +| `Q8F` | 8 | 8-bit full-range quantization | +| `Q4F` | 4 | 4-bit full-range quantization | +| `Q2F` | 2 | 2-bit full-range quantization | +| `Q8S` | 8 | 8-bit symmetric quantization | +| `Q4S` | 4 | 4-bit symmetric quantization | +| `Q2S` | 2 | 2-bit symmetric quantization | +| `E5M2` | 8 | 8-bit floating-point (5 exponent, 2 mantissa) | +| `E4M3` | 8 | 8-bit floating-point (4 exponent, 3 mantissa) | +| `E2M1` | 4 | 4-bit floating-point (2 exponent, 1 mantissa) | + +#### Quantization Store + +| Store | Description | +| :------------------ | :------------------------------------------------------------------------------------------------------------------------------------------------ | +| `Native` | Each quantized value is stored directly in a native format, which doesn't require packing and unpacking. | +| `PackedNative(dim)` | Multiple quantized values packed into a 32-bit integer. Argument is the dimension the tensor is packed on, starting from the innermost dimension. | +| `PackedU32(dim)` | Multiple quantized values packed into a 32-bit integer. Argument is the dimension the tensor is packed on, starting from the innermost dimension. | + +Native storage is not supported for sub-byte quantization values. + +#### Quantization Parameters Precision + +| Param | Description | +| :----- | :----------------------------- | +| `F32` | Full floating-point precision. | +| `F16` | Half-precision floating point. | +| `BF16` | Brain float 16-bit precision. | diff --git a/burn-book/src/saving-and-loading.md b/burn-book/src/saving-and-loading.md new file mode 100644 index 0000000..45222b3 --- /dev/null +++ b/burn-book/src/saving-and-loading.md @@ -0,0 +1,468 @@ +# Saving and Loading Models + +Saving your trained machine learning model is quite easy. As mentioned in the +[Record](./building-blocks/record.md) section, a module's parameters are captured in a +`ModuleRecord` and serialized to the [burnpack](./building-blocks/record.md) format (`.bpk`). + +```rust, ignore +use burn::store::ModuleRecord; + +// Take a record of the model's parameters and save it to disk. +model + .into_record() + .save(model_path) + .expect("Should be able to save the model"); +``` + +Note that the `.bpk` file extension is appended automatically when the path has none, so only the +file path and base name need to be provided. + +Now that you have a trained model saved to your disk, you can just as easily load it back. + +```rust, ignore +// Load the record from the burnpack file. +let record = ModuleRecord::load(model_path) + .expect("Should be able to load the model weights from the provided file"); + +// Apply the loaded weights to a model. +model = model.load_record(record); +``` + +The record is backend-independent: weights saved with one backend can be loaded on another. If you +need to convert precision, call `model.cast(dtype)` before taking the record, or use +`record.cast_to_module_dtype()` when loading. + +## Initialization from Recorded Weights + +The most straightforward way to load weights for a module is simply by using +[load_record](https://burn.dev/docs/burn/module/trait.Module.html#method.load_record). Note that +parameter initialization is lazy, therefore no actual tensor allocation and GPU/CPU kernels are +executed before the module is used. This means that you can use `init(device)` followed by +`load_record(record)` without any meaningful performance cost. + +```rust, ignore +use burn::store::ModuleRecord; + +// Create a dummy initialized model to save. +let device = Default::default(); +let model = Model::init(&device); + +// Save its parameters to a burnpack file. +model + .into_record() + .save(model_path) + .expect("Should be able to save the model"); +``` + +Afterwards, the model can just as easily be loaded from the record saved on disk. + +```rust, ignore +// Load the model record from the burnpack file. +let record = ModuleRecord::load(model_path) + .expect("Could not load model weights"); + +// Initialize a new model with the loaded record/weights. +let model = Model::init(&device).load_record(record); +``` + +For partial loading (only some parameters present in the record), use +`record.allow_partial(true)` before applying it, or `model.try_load_record(record)` for fallible +loading. + +## Model Weight Store + +While the `ModuleRecord` API above works well for basic saving and loading, the `burn-store` crate +adds memory efficiency and flexibility on top of the same burnpack format. It provides zero-copy +memory-mapped loading, cross-framework interoperability (PyTorch and SafeTensors), key remapping, +partial loading, and filtering. + +### Supported Formats + +| Format | Extension | Description | +| --------------- | -------------- | ----------------------------------------------------------------------------------------- | +| **Burnpack** | `.bpk` | Burn's native format with fast loading, zero-copy support, and training state persistence | +| **SafeTensors** | `.safetensors` | Industry-standard format from Hugging Face for secure tensor serialization | +| **PyTorch** | `.pt`, `.pth` | Direct loading of PyTorch model weights (read-only) | + +### Saving a Model + +```rust, ignore +use burn_store::{ModuleSnapshot, BurnpackStore}; + +// Save to Burnpack (recommended) +let mut store = BurnpackStore::from_file("model.bpk"); +model.save_into(&mut store)?; + +// Or save to SafeTensors +use burn_store::SafetensorsStore; +let mut store = SafetensorsStore::from_file("model.safetensors"); +model.save_into(&mut store)?; +``` + +### Loading a Model + +```rust, ignore +use burn_store::{ModuleSnapshot, BurnpackStore}; + +let device = Default::default(); +let mut model = MyModel::init(&device); + +// Load from Burnpack +let mut store = BurnpackStore::from_file("model.bpk"); +model.load_from(&mut store)?; +``` + +### Loading from PyTorch + +You can load weights directly from PyTorch `.pt` files: + +```rust, ignore +use burn_store::{ModuleSnapshot, PytorchStore}; + +let mut model = MyModel::init(&device); +let mut store = PytorchStore::from_file("pytorch_model.pt"); +model.load_from(&mut store)?; +``` + +#### Exporting from PyTorch + +Save only the model weights (state_dict), not the entire model: + +```python +import torch +import torch.nn as nn + +class Net(nn.Module): + def __init__(self): + super(Net, self).__init__() + self.conv1 = nn.Conv2d(2, 2, (2, 2)) + self.conv2 = nn.Conv2d(2, 2, (2, 2), bias=False) + + def forward(self, x): + return self.conv2(self.conv1(x)) + +model = Net() +torch.save(model.state_dict(), "model.pt") # Correct: save state_dict +# torch.save(model, "model.pt") # Wrong: saves entire model +``` + +#### Accessing Nested State Dicts + +Some PyTorch checkpoints nest the state_dict under a key: + +```rust, ignore +let mut store = PytorchStore::from_file("checkpoint.pt") + .with_top_level_key("state_dict"); +model.load_from(&mut store)?; +``` + +### Loading from SafeTensors + +For SafeTensors files exported from PyTorch, use the adapter for proper weight transformation: + +```rust, ignore +use burn_store::{ModuleSnapshot, PyTorchToBurnAdapter, SafetensorsStore}; + +let mut model = MyModel::init(&device); +let mut store = SafetensorsStore::from_file("model.safetensors") + .with_from_adapter(PyTorchToBurnAdapter); +model.load_from(&mut store)?; +``` + +For SafeTensors files created by Burn, no adapter is needed: + +```rust, ignore +let mut store = SafetensorsStore::from_file("model.safetensors"); +model.load_from(&mut store)?; +``` + +#### Exporting from PyTorch to SafeTensors + +```python +from safetensors.torch import save_file + +model = Net() +save_file(model.state_dict(), "model.safetensors") +``` + +### Saving for PyTorch Compatibility + +Use the adapter when saving for PyTorch consumption: + +```rust, ignore +use burn_store::{BurnToPyTorchAdapter, SafetensorsStore}; + +let mut store = SafetensorsStore::from_file("for_pytorch.safetensors") + .with_to_adapter(BurnToPyTorchAdapter) + .skip_enum_variants(true); +model.save_into(&mut store)?; +``` + +### Handling Load Results + +The `load_from` method returns detailed information about the loading process. + +> **Note:** Inspecting `result.missing`, `result.errors`, etc. requires the store to be configured with +> [`.allow_partial(true)`](#partial-loading). Without it, a missing tensor causes a hard `Err` +> before you ever receive an `ApplyResult`. + +```rust, ignore +// Use .allow_partial(true) to get an ApplyResult with structured info +let mut store = PytorchStore::from_file("pretrained.pt") + .allow_partial(true); + +let result = model.load_from(&mut store)?; + +// Print a formatted summary with suggestions +println!("{}", result); + +// Or inspect individual fields +println!("Applied: {} tensors", result.applied.len()); +println!("Missing: {:?}", result.missing); +println!("Errors: {:?}", result.errors); + +if result.is_success() { + println!("All tensors loaded successfully"); +} +``` + +If you don't use `.allow_partial(true)`, use normal error handling: + +```rust, ignore +match model.load_from(&mut store) { + Ok(result) => println!("Loaded successfully: {} tensors", result.applied.len()), + Err(e) => eprintln!("Failed to load: {e}"), +} +``` + +### Adding Metadata + +Burnpack and SafeTensors support custom metadata: + +```rust, ignore +let mut store = BurnpackStore::from_file("model.bpk") + .metadata("version", "1.0") + .metadata("description", "My trained model") + .metadata("epochs", "100"); +model.save_into(&mut store)?; +``` + +### Advanced Features + +#### Key Remapping + +Remap parameter names using regex patterns when model structures don't match: + +```rust, ignore +let mut store = PytorchStore::from_file("model.pt") + // Remove prefix: "model.conv1.weight" -> "conv1.weight" + .with_key_remapping(r"^model\.", "") + // Rename: "layer1" -> "encoder.layer1" + .with_key_remapping(r"^layer", "encoder.layer"); +model.load_from(&mut store)?; +``` + +For complex remapping: + +```rust, ignore +use burn_store::KeyRemapper; + +let remapper = KeyRemapper::new() + .add_pattern(r"^transformer\.h\.(\d+)\.", "transformer.layer$1.")? + .add_pattern(r"\.attn\.", ".attention.")?; + +let mut store = SafetensorsStore::from_file("model.safetensors") + .remap(remapper); +``` + +#### Partial Loading + +Load weights even when some tensors are missing: + +```rust, ignore +let mut store = PytorchStore::from_file("pretrained.pt") + .allow_partial(true); + +let result = model.load_from(&mut store)?; +println!("Missing (initialized randomly): {:?}", result.missing); +``` + +#### Filtering Tensors + +Load or save only specific layers: + +```rust, ignore +// Load only encoder layers +let mut store = SafetensorsStore::from_file("model.safetensors") + .with_regex(r"^encoder\..*") + .allow_partial(true); + +// Save only encoder layers +let mut store = SafetensorsStore::from_file("encoder.safetensors") + .with_regex(r"^encoder\..*"); +model.save_into(&mut store)?; + +// Multiple patterns (OR logic) +let mut store = SafetensorsStore::from_file("model.safetensors") + .with_regex(r"^encoder\..*") // encoder tensors + .with_regex(r".*\.bias$") // OR any bias tensors + .with_full_path("decoder.scale"); // OR specific tensor +``` + +#### Non-Contiguous Layer Indices + +PyTorch `nn.Sequential` with mixed layers creates non-contiguous indices. `PytorchStore` +automatically remaps these: + +``` +PyTorch: fc.0.weight, fc.2.weight, fc.4.weight (gaps from ReLU layers) +Burn: fc.0.weight, fc.1.weight, fc.2.weight (contiguous) +``` + +This is enabled by default. Disable if needed: + +```rust, ignore +let mut store = PytorchStore::from_file("model.pt") + .map_indices_contiguous(false); +``` + +#### Zero-Copy Loading + +For embedded models or large files, use zero-copy loading to avoid memory copies: + +```rust, ignore +// Embedded model (compile-time) +static MODEL_DATA: &[u8] = include_bytes!("model.bpk"); +let mut store = BurnpackStore::from_static(MODEL_DATA); +model.load_from(&mut store)?; + +// Large file (memory-mapped) +let mut store = BurnpackStore::from_file("large_model.bpk") + .zero_copy(true); +model.load_from(&mut store)?; +``` + +#### Half-Precision Storage + +Save models at half precision (F16) to reduce file size by ~50%, then load back at full precision: + +```rust, ignore +use burn_store::{ModuleSnapshot, BurnpackStore, HalfPrecisionAdapter}; + +let adapter = HalfPrecisionAdapter::new(); + +// Save: F32 -> F16 (same adapter for both directions) +let mut store = BurnpackStore::from_file("model_f16.bpk") + .with_to_adapter(adapter.clone()); +model.save_into(&mut store)?; + +// Load: F16 -> F32 +let mut store = BurnpackStore::from_file("model_f16.bpk") + .with_from_adapter(adapter); +model.load_from(&mut store)?; +``` + +By default, weights in Linear, Embedding, Conv\*, LayerNorm, GroupNorm, InstanceNorm, RmsNorm, and +PRelu modules are converted. BatchNorm is excluded because its running variance can underflow in +F16. Customize with `with_module()` and `without_module()`: + +```rust, ignore +// Keep LayerNorm at full precision +let adapter = HalfPrecisionAdapter::new() + .without_module("LayerNorm"); + +// Add a custom module to the conversion set +let adapter = HalfPrecisionAdapter::new() + .with_module("CustomLayer"); +``` + +#### Direct Tensor Access + +Inspect tensors without loading into a model: + +```rust, ignore +use burn_store::ModuleStore; + +let mut store = PytorchStore::from_file("model.pt"); + +// List all tensor names +let names = store.keys()?; + +// Get specific tensor +if let Some(snapshot) = store.get_snapshot("encoder.layer0.weight")? { + println!("Shape: {:?}, DType: {:?}", snapshot.shape, snapshot.dtype); +} +``` + +#### Model Surgery + +Transfer weights between models: + +```rust, ignore +use burn_store::{ModuleSnapshot, PathFilter}; + +// Transfer all weights +let snapshots = model1.collect(None, None, false); +model2.apply(snapshots, None, None, false); + +// Transfer only encoder weights +let filter = PathFilter::new().with_regex(r"^encoder\..*"); +let snapshots = model1.collect(Some(filter.clone()), None, false); +model2.apply(snapshots, Some(filter), None, false); +``` + +### API Reference + +#### Builder Methods + +| Category | Method | Description | +| ------------- | ------------------------------ | ---------------------------- | +| **Filtering** | `with_regex(pattern)` | Filter by regex pattern | +| | `with_full_path(path)` | Include specific tensor | +| | `with_predicate(fn)` | Custom filter logic | +| **Remapping** | `with_key_remapping(from, to)` | Regex-based renaming | +| | `remap(KeyRemapper)` | Complex remapping rules | +| **Adapters** | `with_from_adapter(adapter)` | Loading transformations | +| | `with_to_adapter(adapter)` | Saving transformations | +| | `HalfPrecisionAdapter::new()` | F32/F16 mixed-precision | +| **Config** | `allow_partial(bool)` | Continue on missing tensors | +| | `with_top_level_key(key)` | Access nested dict (PyTorch) | +| | `skip_enum_variants(bool)` | Skip enum variants in paths | +| | `map_indices_contiguous(bool)` | Remap non-contiguous indices | +| | `metadata(key, value)` | Add custom metadata | +| | `zero_copy(bool)` | Enable zero-copy loading | + +#### Direct Access Methods + +| Method | Description | +| --------------------- | -------------------------------- | +| `keys()` | Get ordered list of tensor names | +| `get_all_snapshots()` | Get all tensors as BTreeMap | +| `get_snapshot(name)` | Get specific tensor by name | + +### Troubleshooting + +#### Common Issues + +1. **"Missing source values" error**: You saved the entire PyTorch model instead of the state_dict. + Re-export with `torch.save(model.state_dict(), "model.pt")`. + +2. **Shape mismatch**: Your Burn model doesn't match the source architecture. Verify layer + configurations (channels, kernel sizes, bias settings). + +3. **Key not found**: Parameter names don't match. Use `with_key_remapping()` or inspect keys: + + ```rust, ignore + let store = PytorchStore::from_file("model.pt"); + println!("Available keys: {:?}", store.keys()?); + ``` + +#### Inspecting Files + +Use [Netron](https://github.com/lutzroeder/netron) to visualize `.pt` and `.safetensors` files. + +For Burnpack files: + +```bash +cargo run --example burnpack_inspect model.bpk +``` diff --git a/codecov.yml b/codecov.yml new file mode 100644 index 0000000..f6dbc7c --- /dev/null +++ b/codecov.yml @@ -0,0 +1,13 @@ +coverage: + status: + project: + default: + # https://docs.codecov.com/docs/commit-status#informational + informational: true + target: 80% + patch: + default: + informational: true + target: 80% +github_checks: + annotations: false diff --git a/contributor-book/.gitignore b/contributor-book/.gitignore new file mode 100644 index 0000000..409ff3e --- /dev/null +++ b/contributor-book/.gitignore @@ -0,0 +1,18 @@ +target + +# MacOS temp file +.DS_Store + +book-test +guide/book + +.vscode +tests/burn-book/book/ +book/ + +# Ignore Jetbrains specific files. +.idea/ + +# Ignore Vim temporary and swap files. +*.sw? +*~ \ No newline at end of file diff --git a/contributor-book/.prettierrc.json b/contributor-book/.prettierrc.json new file mode 100644 index 0000000..d410551 --- /dev/null +++ b/contributor-book/.prettierrc.json @@ -0,0 +1,4 @@ +{ + "printWidth": 100, + "proseWrap": "always" +} \ No newline at end of file diff --git a/contributor-book/LICENSE-APACHE b/contributor-book/LICENSE-APACHE new file mode 120000 index 0000000..965b606 --- /dev/null +++ b/contributor-book/LICENSE-APACHE @@ -0,0 +1 @@ +../LICENSE-APACHE \ No newline at end of file diff --git a/contributor-book/LICENSE-MIT b/contributor-book/LICENSE-MIT new file mode 120000 index 0000000..76219eb --- /dev/null +++ b/contributor-book/LICENSE-MIT @@ -0,0 +1 @@ +../LICENSE-MIT \ No newline at end of file diff --git a/contributor-book/book.toml b/contributor-book/book.toml new file mode 100644 index 0000000..4522483 --- /dev/null +++ b/contributor-book/book.toml @@ -0,0 +1,16 @@ +[book] +authors = [ + "Wouter Doppenberg", + "Nathaniel Simard", + "Louis Fortier-Dubois", + "Dilshod Tadjibaev", + "Guillaume Lagrange", + "Joshua Ferguson", + "The Burn Community", +] +language = "en" +src = "src" +title = "The Burn Contributor Book 🔥" + +[output.html] +mathjax-support = true diff --git a/contributor-book/src/SUMMARY.md b/contributor-book/src/SUMMARY.md new file mode 100644 index 0000000..89cc2cf --- /dev/null +++ b/contributor-book/src/SUMMARY.md @@ -0,0 +1,16 @@ +- [Overview](./overview.md) +- [How to Read This Book](./how-to-read-this-book.md) +- [Getting Started](./getting-started/README.md) + - [Setting Up The Environment](./getting-started/setting-up-the-environment.md) + - [Configuring Your Editor (Optional)](./getting-started/configuring-your-editor.md) + - [Testing](./getting-started/testing.md) +- [Architecture Overview](./project-architecture/README.md) + - [Modules](./project-architecture/module.md) + - [Serialization](./project-architecture/serialization.md) + - [Tensor](./project-architecture/tensor.md) + - [Backend](./project-architecture/backend.md) +- [Guides for Contributors](./guides/README.md) + - [Adding a New Operation to Burn](./guides/adding-a-new-operation-to-burn.md) + - [Submitting Examples to Burn](./guides/submitting-examples.md) +- [Frequently Encountered Issues](./frequently-encountered-issues/README.md) + - [Issues Related To Adding Operators](./frequently-encountered-issues/issues-while-adding-ops.md) diff --git a/contributor-book/src/frequently-encountered-issues/README.md b/contributor-book/src/frequently-encountered-issues/README.md new file mode 100644 index 0000000..d865f1b --- /dev/null +++ b/contributor-book/src/frequently-encountered-issues/README.md @@ -0,0 +1,5 @@ +# Frequently Encountered Issues + +This is a collection of issues people have encountered and asked about on the +[Discord server](https://discord.gg/uPEBbYYDB6). This section is separated from the guides since it +can involve lots of details that are only relevant to a small subset of contributors. diff --git a/contributor-book/src/frequently-encountered-issues/issues-while-adding-ops.md b/contributor-book/src/frequently-encountered-issues/issues-while-adding-ops.md new file mode 100644 index 0000000..785ee33 --- /dev/null +++ b/contributor-book/src/frequently-encountered-issues/issues-while-adding-ops.md @@ -0,0 +1,19 @@ +# Issues encountered while adding ops + +Below are some of the issues that were encountered while adding ops to the project. If you encounter +an issue while adding an op that isn't listed here, and it's not obvious how to fix it, you can add +it to this list or reach out on the [Discord server](https://discord.gg/uPEBbYYDB6) if you need +help. + +## Off by .000001 errors + +```sh +---- fusion::base::tests::maxmin::tests::test_mean_dim_2d stdout ---- thread 'fusion::base::tests::maxmin::tests::test_mean_dim_2d' panicked at burn-wgpu/src/fusion/base.rs:185:5: assertion `left == right` failed left: Data { value: [1.0, 4.0], shape: Shape { dims: [2, 1] } } right: Data { value: [0.99999994, 3.9999998], shape: Shape { dims: [2, 1] } } ---- + +tests::maxmin::tests::test_mean_dim_2d stdout ---- thread 'tests::maxmin::tests::test_mean_dim_2d' panicked at burn-wgpu/src/lib.rs:49:5: assertion `left == right` failed left: Data { value: [1.0, 4.0], shape: Shape { dims: [2, 1] } } right: Data { value: [0.99999994, 3.9999998], shape: Shape { dims: [2, 1] } } +``` + +If you encounter this, swap out the `assert_eq!` in the failing test for +`tensor1.to_data().assert_approx_eq` with `3` as the second argument. The second arguments specifies +the level of precision: `3` is equivalent to a less than 10-3 (0.001) difference between +the elements of the two tensors. diff --git a/contributor-book/src/getting-started/README.md b/contributor-book/src/getting-started/README.md new file mode 100644 index 0000000..0414d86 --- /dev/null +++ b/contributor-book/src/getting-started/README.md @@ -0,0 +1,6 @@ +# Getting Started + +This section is for setting up the environment and how to do basic development tasks such as running +tests and checking your code before committing. If you need help with the process or run into +issues, feel free to ask on the [Discord server](https://discord.gg/uPEBbYYDB6) in the Development +channels. diff --git a/contributor-book/src/getting-started/configuring-your-editor.md b/contributor-book/src/getting-started/configuring-your-editor.md new file mode 100644 index 0000000..10ad6a7 --- /dev/null +++ b/contributor-book/src/getting-started/configuring-your-editor.md @@ -0,0 +1,36 @@ +# Configuring your editor + +These steps are not required, and most of this isn't specific to Burn, but it's definitely helpful +if you haven't already done it. + +## VSCode + +Install the following extensions: + +- [rust-lang.rust-analyzer](https://marketplace.visualstudio.com/items?itemName=rust-lang.rust-analyzer) + for Rust syntax and semantic analysis +- [tamasfe.even-better-toml](https://marketplace.visualstudio.com/items?itemName=tamasfe.even-better-toml) + for TOML syntax and semantic analysis +- [fill-labs.dependi](https://marketplace.visualstudio.com/items?itemName=fill-labs.dependi) for + managing dependencies +- [vadimcn.vscode-lldb](https://marketplace.visualstudio.com/items?itemName=vadimcn.vscode-lldb) for + debugging + +### Setting up the Debugger + +To use the debugger, follow these steps: + +1. Open `Command Palette` with `Ctrl+Shift+P` or `F1` and type + `LLDB: Generate Launch Configurations from Cargo.toml` then select it, this will generate a file + that should be saved as `.vscode/launch.json`. +2. Select the configuration from the "run and debug" side panel, then select the target from the list. + Since this repo has `debug = 0` in the root `Cargo.toml` to speed up compilation, you need replace it with `debug = true` in the root `Cargo.toml` when using a debugger and breakpoints with `launch.json` settings. +3. Now you can enable breakpoints on code through IDE then start debugging the library/binary you + want, like in the following example: + +![debug-options](debug-options-vscode.png) + +If you're creating a new library or binary, keep in mind to repeat step 1 to always keep a fresh +list of targets. + +## Have another editor? Open a PR! diff --git a/contributor-book/src/getting-started/debug-options-vscode.png b/contributor-book/src/getting-started/debug-options-vscode.png new file mode 100644 index 0000000..c1db902 Binary files /dev/null and b/contributor-book/src/getting-started/debug-options-vscode.png differ diff --git a/contributor-book/src/getting-started/setting-up-the-environment.md b/contributor-book/src/getting-started/setting-up-the-environment.md new file mode 100644 index 0000000..6ae9fa9 --- /dev/null +++ b/contributor-book/src/getting-started/setting-up-the-environment.md @@ -0,0 +1,54 @@ +# Setting up the environment + +Depending on what part of the project you plan on contributing to, there are a couple of tools to +install and commands to be familiar with. This section should be up to date with current project +practices (as of 2024-04-15). + +## General + +There are a few commands you will want to run prior to any commit for a non-draft PR: + +1. `cargo fmt --all` will run `rustfmt` on all files in the project. +2. `cargo clippy --fix` will run [Clippy](https://github.com/rust-lang/rust-clippy) and fix any + coding issues it can. Clippy necessitates to be in a clean Git state, but this can be + circumvented by adding the `--allow-dirty` flag. +3. `cargo run-checks` is a command used to test the project. It is required to run successfully + prior to merging a PR. Fair warning, running these tests can take a while[^linux_mem_note]. + + > Want more detailed macro error diagnostics? This is especially useful for debugging tensor-related tests: + > + > ```bash + > RUSTC_BOOTSTRAP=1 RUSTFLAGS="-Zmacro-backtrace" cargo run-checks + > ``` + +## Updating the burn semver version + +If for some reason you need to bump for the next version (though that should probably be left to the +maintainers), edit the semantic version number in `burn/Cargo.toml`, and then run `cargo update` to +update the lock file. + +## Contributing to either the Burn Book or Contributor Book + +Both the Burn Book and the Contributor Book are built with mdbook. To open the book locally, run +`mdbook serve ` or `cargo xtask books {burn|contributor} open` which will install and +use mdbook automatically. + +Alternatively, if you want to install mdbook directly, run the following command[^update_note]: + +```bash +cargo install mdbook +``` + +Also instead of running `cargo run-checks`, you can run `cargo xtask check typos` to only check +for misspellings. This will install [typo](https://crates.io/crates/typos-cli), and if any are +encountered you should be able to run `typo -w /path/to/book` to fix them. + +[^linux_mem_note]: + If your system is running into issues with memory and you are on linux, you may want to switch + to a [virtual console](https://wiki.archlinux.org/title/Linux_console#Virtual_consoles) to run + the tests. To do this, press `ctrl+alt+f3` to switch to a virtual console (and log in), and + either `ctrl+alt+f1` or `ctrl+alt+f2` to switch back to your graphical session. + +[^update_note]: + You might also want to install [cargo-update](https://github.com/nabijaczleweli/cargo-update) to + easily keep your tools up to date, though it is in no way required. diff --git a/contributor-book/src/getting-started/testing.md b/contributor-book/src/getting-started/testing.md new file mode 100644 index 0000000..bd73280 --- /dev/null +++ b/contributor-book/src/getting-started/testing.md @@ -0,0 +1,38 @@ +# Testing + +## Test for Tensor Operations + +Test for tensor operations (generally of the form: given this input, expect it match or approximate +this output) are defined only in +[`crates/burn-tensor/src/test/ops`](https://github.com/tracel-ai/burn/tree/81a67b6a0992b9b5c33cda8b9784570143b67319/crates/burn-tensor/src/tests/ops) +and not in the backends (with the exception of `burn-autodiff`). The tensor operation tests are +added to the `testgen_all` macro rule in +[`crates/burn-tensor/src/tests/mod.rs`](https://github.com/tracel-ai/burn/blob/81a67b6a0992b9b5c33cda8b9784570143b67319/crates/burn-tensor/src/tests/mod.rs). +This is then propagated to the existing backends without any additional work. + +### Test for Autodiff + +Tests for autodiff go under +[burn-autodiff/src/tests](https://github.com/tracel-ai/burn/tree/81a67b6a0992b9b5c33cda8b9784570143b67319/crates/burn-autodiff/src/tests) +and should verify backward pass correctness. For binary tensor operations, both the left and right +sides need to be verified. + +Here's an easy way to define tests for a new operation's backward pass: + +1. Use small tensors with simple values. +2. Pop open a terminal, launch `ipython` and import `numpy` then do the calculations by hand. You + can also use [Google Colab](https://colab.google/) so you don't have to install the packages on + your system. +3. Compare the actual outputs to the expected output for left-hand side, right-hand side. + +For float tensors, it is advised to use +`actual_output_tensor.into_data().assert_approx_eq::>(&expected_tensor_data, Tolerance::default())` +instead of `assert_eq!(...` due to occasional hiccups with floating point calculations. Other +assertions should also always use `FloatElem`, and use `.elem()` to convert any +literals. Backends are tested for multiple precisions, and hardcoding to a fixed type causes tests +to fail with alternate floating point precisions. For convenience, it might be worth aliasing the +type like `type FT = FloatElem;`. + +For integers, tests should use `IntElem`, and exit the test if the test values are +unrepresentable (above `max_value`, below `min_value`). A minimum range of `[0..127]` (`i8`) can be +assumed. diff --git a/contributor-book/src/guides/README.md b/contributor-book/src/guides/README.md new file mode 100644 index 0000000..a4d9635 --- /dev/null +++ b/contributor-book/src/guides/README.md @@ -0,0 +1,3 @@ +# Guides for Contributors + +The following guides are meant to help contributors accomplish specific tasks, such as adding new operations to Burn. \ No newline at end of file diff --git a/contributor-book/src/guides/adding-a-new-operation-to-burn.md b/contributor-book/src/guides/adding-a-new-operation-to-burn.md new file mode 100644 index 0000000..b9b28a2 --- /dev/null +++ b/contributor-book/src/guides/adding-a-new-operation-to-burn.md @@ -0,0 +1,246 @@ +# Adding a New Operation to burn + +Let's discuss how one might go about adding new operators to Burn, using the example of the pow +operator added in [this PR](https://github.com/tracel-ai/burn/pull/1133/files). + +## Adding the Op to burn-tensor + +`burn-tensor` is the crate that defines all tensor operations that need to be implemented by the +various backends. The core of this lies in +[crates/burn-backend/src/tensor/ops/numeric.rs](https://github.com/tracel-ai/burn/blob/9f31281/crates/burn-backend/src/tensor/ops/numeric.rs#L17), +which is home to the numeric trait. The numeric trait is the home of all tensor operations that are +numeric in nature and that are shared by `Int` and `Float` Tensor types. The numeric trait is +implemented in +[crates/burn-backend/src/tensor/ops/int.rs](https://github.com/tracel-ai/burn/blob/9f31281/crates/burn-backend/src/tensor/ops/int.rs) +for the int type and in +[crates/burn-backend/src/tensor/ops/float.rs](https://github.com/tracel-ai/burn/blob/9f31281/crates/burn-backend/src/tensor/ops/float.rs) +for the float type. More information on the relationship between Tensor modules can be found under +the section for [Tensor Architecture](../project-architecture/tensor.md#tensor-operations). + +Here is where pow was added to `crates/burn-tensor/src/tensor/api/numeric.rs`: + +1. for the + [`Tensor` struct](https://github.com/tracel-ai/burn/blob/0ee2021567b3725907df5fd1a905ce60b1aca096/crates/burn-tensor/src/tensor/api/numeric.rs#L573) +2. for the + [numeric trait](https://github.com/tracel-ai/burn/blob/0ee2021567b3725907df5fd1a905ce60b1aca096/crates/burn-tensor/src/tensor/api/numeric.rs#L1955) +3. for the implementation of numeric for + [float](https://github.com/tracel-ai/burn/blob/0ee2021567b3725907df5fd1a905ce60b1aca096/crates/burn-tensor/src/tensor/api/numeric.rs#L2722) + and + [int](https://github.com/tracel-ai/burn/blob/0ee2021567b3725907df5fd1a905ce60b1aca096/crates/burn-tensor/src/tensor/api/numeric.rs#L2375) + +Tensor is a struct that has a single member: `primitive` (defined +[here](https://github.com/tracel-ai/burn/blob/0ee2021567b3725907df5fd1a905ce60b1aca096/crates/burn-tensor/src/tensor/api/base.rs#L27)), +that is defined by its +[`Kind`](https://github.com/tracel-ai/burn/blob/0ee2021567b3725907df5fd1a905ce60b1aca096/crates/burn-tensor/src/tensor/api/kind.rs#L16): +one of `Bool`, `Float`, or `Int` (those linked in 3). These call the ops for that data type defined +in the +[`Backend`](https://github.com/tracel-ai/burn/blob/9f31281/crates/burn-backend/src/backend/base.rs#L64) +supertrait[^supertrait]. This is the trait that is then implemented by the different `burn-` +backends (such as `burn-flex` and `burn-wgpu`) which must implement the functions if no default +is provided. + +In this case, we don't need to worry about `Bool` Tensors. `Float` ops are implemented under +[crates/burn-backend/src/backend/ops/tensor.rs](https://github.com/tracel-ai/burn/blob/9f31281/crates/burn-backend/src/backend/ops/tensor.rs), +and `Int` ops under +[crates/burn-backend/src/backend/ops/int_tensor.rs](https://github.com/tracel-ai/burn/blob/9f31281/crates/burn-backend/src/backend/ops/int_tensor.rs). +The current convention is ops of each type, if not unique to that type, are prefixed with the type. +So `powf` and sundry would be defined as `int_powf` for `IntTensorOps` and `float_powf` for +`FloatTensorOps`. If an op is unique to a type, then it should be implemented under +`burn-tensor/src/api/{type}.rs`. For example, here is an implementation for +[`sin` under `crates/burn-tensor/src/api/float.rs`](https://github.com/tracel-ai/burn/blob/0ee2021567b3725907df5fd1a905ce60b1aca096/crates/burn-tensor/src/tensor/api/float.rs#L82) +which obviously doesn't make sense for `Int` or `Bool` tensors. + +The `Int` Tensor function uses the ones defined for Float with 2 extra casts (LHS to a `Float` +tensor, Output to an `Int`). Given that the rest of the code will only look at the float +implementations. + +With the addition of quantized float tensors, the `Float` tensor primitive is represented by the +[`TensorPrimitive`](https://github.com/tracel-ai/burn/blob/a6a5c22e0db56d947b9165d4dae42783a5a6b689/crates/burn-tensor/src/tensor/api/kind.rs#L69) +enum. This allows us to handle both float and quantized float operations in the `Tensor` +implementation, correctly dispatching to the corresponding op (float or quantized) based on the +variant. Following the same convention, the equivalent +[quantized tensor ops](https://github.com/tracel-ai/burn/blob/a6a5c22e0db56d947b9165d4dae42783a5a6b689/crates/burn-tensor/src/tensor/ops/qtensor.rs#L45) +are prefixed with `q_*` (e.g., `q_reshape` instead of `float_reshape`). Most ops have a default +implementation that simply dequantizes the input into its floating-point representation, performs +the operation on the float tensor, and quantizes the output. Backends can overwrite specific +implementations when required/desired. + +### Adding Tests + +Additional tests should be added to `burn-backend-tests` under +[`crates/burn-backend-tests/tests/tensor/{float_or_int}/ops/{op_name}.rs`](https://github.com/tracel-ai/burn/blob/9f31281/crates/burn-backend-tests/tests/tensor/float/ops/powf.rs), +and the module name should be inserted into +`crates/burn-backend-tests/tests/tensor/{float_or_int}/ops/mod.rs`. + +If it makes sense for a floating point operation to support quantization, the +[`QTensorOps`](https://github.com/tracel-ai/burn/blob/9f31281/crates/burn-backend/src/backend/ops/qtensor.rs#L117) +counterpart is usually added at the same time with a default implementation (as mentioned in the +previous section). Tests for `q_*` ops follow a similar procedure: the test is added under +[`crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/{op_name}.rs`](https://github.com/tracel-ai/burn/tree/9f31281/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended), +the module name is inserted into +[`crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/mod.rs`](https://github.com/tracel-ai/burn/blob/9f31281/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/mod.rs). +If you take a look at any of the existing tests for an operation on a quantized tensor, you will see +that the inputs and expected outputs are always defined with floating point values. While it assumes +that the quantization and dequantization are correct, it makes the tests much more readable and +easier to understand w.r.t. what is being tested. Effectively, the tests are there to ensure that a +tensor operation is invariant to quantization (up to some quantization error, of course). + +_Note: the tests try to use tensors with floating point values which can be de/quantized without +introducing too much quantization error, but the result always depends on the operation (e.g., +tensor product of values can grow larger and significantly increase the output tensor range, leading +to more de/quantization error on the results)._ + +## Adding the Op to burn-autodiff + +Since this is probably the hardest and the least straightforward, we'll cover this backend +separately. `burn-autodiff` enables other backends to use autodifferentiation[^autodiff]. Ops for +float types are implemented in +[crates/burn-autodiff/src/ops/tensor.rs](https://github.com/tracel-ai/burn/blob/0ee2021567b3725907df5fd1a905ce60b1aca096/crates/burn-autodiff/src/ops/tensor.rs) +and need to: + +1. Define a unit struct [^absolute_units] that implements a backward (pass) function +2. Within the backward function, as this is an elementwise binary operation it implements the binary + function (from `backward.rs` under the same directory), the last 2 arguments are two closures + that define the left and right partial derivatives. +3. Then define what happens when a specific operation is tracked or untracked, where untracked just + calls the function in the normal way, and tracked sets the execution the backward function + defined above. +4. When tracked, operations are part of the autodiff graph and must save the needed information to + efficiently perform their backward pass later. If the information is light (such as a shape), it + should be directly saved in the state. If the operation's inputs are needed to compute the + backward pass, it should be checkpointed rather than saved. This will allow the input to be + provided lazily at the backward pass depending on the checkpointing strategy. +5. An operation must also be identified as _compute-bound_ (`.computeBound()`) or _memory-bound_ + (`.memoryBound()`) for gradient checkpointing. _Compute-bound_ operation are heavy to compute + (for instance matmul or convolution), which means that even with checkpointing they will save + their output for the backward pass and not recompute it. _Memory-bound_ operations are more + trivial (like `powf` which only performs one small operation per tensor entry), so it can be + beneficial to recompute them during the backward pass instead of saving their whole forward + output to memory. Operations registered as _memory-bound_ need to know their parents + (`.parents()` method) and how to recompute their forward pass during the backward pass (with a + struct that implements `RetroForward`), using their parents' outputs. + +The above steps are mostly boilerplate, so you can often just copy the contents of another similar +op, change the name of the structs, and ensure that either both sides have the data they need (if +they need to have a copy of the opposite sided tensor, clone its contents). + +### Computing derivatives + +For those that need it, here is a quick refresher on the necessary calculus. If you are familiar +with how to calculate partial derivatives, you can skip this section. + +Since `pow` is a binary operation, the left and right functions are the partial derivatives with +respect to the left and right sided tensors. + +Let's define the operator as a function \\(f(x,y)=x^{y}\\) , where \\(x\\) is the left hand tensor +and \\(y\\) is the right handed tensor. The two closures are defining the partial derivatives of +\\(f\\) with respect to \\(x\\),\\(y\\). Treat the other variables as a constant + +$$\frac{\delta }{\delta x} (x^{y})= y \cdot x^{y-1}$$ is the left handed closure, and + +$$\frac{\delta }{\delta y} (x^{y}) = x^{y} \cdot ln(x)$$ + +is the right. If you aren't sure how to calculate these by hand, it is recommended to use +[symbolab](), +plug in your operator in terms of \\(x\\) and \\(y\\), and just swap out the variable +\\(x\\)|\\(y\\) in the partial derivative to get the other side. + +### Testing autodiff + +For testing the `autodiff` operations, please refer to +[this section](../getting-started/testing.md). + +## Adding the Op to other backends + +Most of these are fairly straightforward implementations. For reference here's pow's float +implementation for torch and flex backends: + +1. Torch implementation in + [crates/burn-tch/src/ops/tensor.rs](https://github.com/tracel-ai/burn/blob/0ee2021567b3725907df5fd1a905ce60b1aca096/crates/burn-tch/src/ops/tensor.rs#L467) + and the Op used in + [crates/burn-tch/src/ops/base.rs](https://github.com/tracel-ai/burn/blob/0ee2021567b3725907df5fd1a905ce60b1aca096/crates/burn-tch/src/ops/base.rs#L481) +2. Flex in + [crates/burn-flex/src/ops/float.rs](https://github.com/tracel-ai/burn/blob/main/crates/burn-flex/src/ops/float.rs) + +This is where any calculation happens currently. Playing a guessing game with method names and +seeing what completions are suggested will take you far. If you are having trouble figuring out how +to do it from the docs for that backend, +[try searching github for relevant function calls](https://docs.github.com/en/search-github/github-code-search/understanding-github-code-search-syntax). + +## Adding the Op to fusion, JIT and cubecl backends + +Adding an operator to these backends can be fairly straightforward, though due to what these +backends are for, involves a bit more indirection. Fusion and jit, like autodiff, are not target +backends as much as backends that enable certain functionality for other backends, in this case +kernel fusion or just-in-time compilation. Adding the operator won't involve doing any calculation, +you'll just be describing how the generated code should look. Most of this can be +copy/pasted/adjusted from other functions. + +Here's how powf was added to `burn-fusion`: + +1. Added powf to the float ops under + [crates/burn-fusion/src/ops/tensor.rs](https://github.com/tracel-ai/burn/blob/9f31281/crates/burn-fusion/src/ops/tensor.rs#L2061) +2. Added powf to the `NumericOperationIr` enum under + [crates/burn-ir/src/operation.rs](https://github.com/tracel-ai/burn/blob/9f31281/crates/burn-ir/src/operation.rs#L564) +3. Added powf to the implementations of `NumericOperationIr` enum under + [crates/burn-ir/src/operation.rs](https://github.com/tracel-ai/burn/blob/9f31281/crates/burn-ir/src/operation.rs#L1086) +4. Added powf to the implemented of `NumericOperationIr` enum under + [burn/crates/burn-fusion/src/stream/context.rs](https://github.com/tracel-ai/burn/blob/9f31281/crates/burn-fusion/src/stream/context.rs#L883) + +The way `cubecl` handles tensor-scalar operations is by transforming both into a sequence of +vectorized scalar operations. Since powf already existed in `cubecl`, it was pretty easy to reuse +the existing implementation for the situation where both sides of the operation were tensors. The +`cubecl` crate is primarily concerned with how the operation is compiled and executed by the gpu. +The actual implementation is defined in `burn-cubecl`. + +Here is where code was added for powf in `burn-cubecl` and `cubecl`: + +1. to the implementation of + [`FloatTensorOps` under `burn/crates/burn-cubecl/src/ops/tensor.rs`](https://github.com/tracel-ai/burn/blob/9f31281/crates/burn-cubecl/src/ops/tensor.rs#L578) +2. the function being called was added to + [`burn/crates/burn-cubecl/src/ops/numeric.rs`](https://github.com/tracel-ai/burn/blob/9f31281/crates/burn-cubecl/src/ops/numeric.rs#L211-L214) +3. the operator was defined in + [`cubecl/crates/cubecl-ir/src/arithmetic.rs`](https://github.com/tracel-ai/cubecl/blob/88c0c6f781f70ad2f6e9981fd0cbe2e87e153a35/crates/cubecl-ir/src/arithmetic.rs#L41) +4. how the operation looks to the gpu was added to + [`burn/crates/burn-cubecl-fusion/src/engine/codegen/ir.rs`](https://github.com/tracel-ai/burn/blob/9f31281/crates/burn-cubecl-fusion/src/engine/codegen/ir.rs#L97) +5. the mappings between the gpu operation and the CPP, WGSL and SPIR-V instructions were added to + [`cubecl/crates/cubecl-cpp/src/shared/base.rs`](https://github.com/tracel-ai/cubecl/blob/88c0c6f781f70ad2f6e9981fd0cbe2e87e153a35/crates/cubecl-cpp/src/shared/base.rs#L1285), + [`cubecl/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs`](https://github.com/tracel-ai/cubecl/blob/88c0c6f781f70ad2f6e9981fd0cbe2e87e153a35/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs#L869) + and + [`cubecl/crates/cubecl-spirv/src/arithmetic.rs`](https://github.com/tracel-ai/cubecl/blob/88c0c6f781f70ad2f6e9981fd0cbe2e87e153a35/crates/cubecl-spirv/src/arithmetic.rs#L491) +6. the instructions themselves were added for WGSL to + [instruction op enum in `cubecl/crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs`](https://github.com/tracel-ai/cubecl/blob/f5b63076a01a5c03ea9ed20799d3eeaf776b45da/crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs#L124), + and the actual + [instruction in wgsl here](https://github.com/tracel-ai/cubecl/blob/88c0c6f781f70ad2f6e9981fd0cbe2e87e153a35/crates/cubecl-wgpu/src/compiler/wgsl/instructions.rs#L654), + for CPP in the enum here + [`cubecl/crates/cubecl-cpp/src/shared/instruction.rs`](https://github.com/tracel-ai/cubecl/blob/88c0c6f781f70ad2f6e9981fd0cbe2e87e153a35/crates/cubecl-cpp/src/shared/instruction.rs#L187) + and the actual instruction here + [`cubecl/crates/cubecl-cpp/src/shared/binary.rs`](https://github.com/tracel-ai/cubecl/blob/88c0c6f781f70ad2f6e9981fd0cbe2e87e153a35/crates/cubecl-cpp/src/shared/binary.rs#L216) + +We needed to generate some custom WGSL code for powf in WGSL, primarily due to issues with proper +case handling of the wgsl pow function, like 0 to the 0 power being 1, and any negative number to an +even power being positive. We reused as much as the existing logic as possible, and then branched at +the last point based off the var type of the rhs. +[See here](https://github.com/tracel-ai/cubecl/blob/88c0c6f781f70ad2f6e9981fd0cbe2e87e153a35/crates/cubecl-wgpu/src/compiler/wgsl/compiler.rs#L1229). +For most operations, you shouldn't need to add to `cubecl-wgpu/src/compiler/wgsl/extension.rs` +unless the operation isn't native to WGSL. + +For functions that need a complex kernel without a direct mapping to a base instruction, simply use +the `cube` macro (see +[the `cubecl` book](https://github.com/tracel-ai/cubecl/tree/88c0c6f781f70ad2f6e9981fd0cbe2e87e153a35/cubecl-book)). + +And you're done! Congrats, you just fully added a new operation to burn, and we are all one step +closer to the answer to [Are we learning yet?](https://www.arewelearningyet.com/) being "Yes, and +it's freaking fast!". Buy yourself a coffee. + +[^supertrait]: + for more on supertraits see + [the advanced trait section of the rust book](https://doc.rust-lang.org/book/ch19-03-advanced-traits.html#using-supertraits-to-require-one-traits-functionality-within-another-trait) + +[^autodiff]: + wiki link for + [automatic differentiation](https://en.wikipedia.org/wiki/Automatic_differentiation) + +[^absolute_units]: + for more information on unit structs see + [the defining and instantiating structs section of the rust book](https://doc.rust-lang.org/book/ch05-01-defining-structs.html#unit-like-structs-without-any-fields) diff --git a/contributor-book/src/guides/submitting-examples.md b/contributor-book/src/guides/submitting-examples.md new file mode 100644 index 0000000..12f32b1 --- /dev/null +++ b/contributor-book/src/guides/submitting-examples.md @@ -0,0 +1,108 @@ +# Submitting Examples to Burn + +This guide explains how to create and submit new examples to the Burn repository. Examples are a great way to demonstrate Burn's capabilities and help users understand how to use the framework effectively. + +For a minimal working example, see the [simple-regression](https://github.com/tracel-ai/burn/blob/main/examples/simple-regression/examples/regression.rs) example in the repository. + +## Repository Structure + +The Burn repository is set up as a workspace, with examples located in the `examples/` directory. Each example is a separate crate that can reuse workspace dependencies. + +## Creating a New Example + +1. Navigate to the examples directory: + ```bash + cd examples + ``` + +2. Create a new library crate: + ```bash + cargo new --lib + ``` + +3. Update the example's `Cargo.toml`: + ```toml + [package] + name = "" + version = "0.1.0" + edition = "2021" + readme = "README.md" + # Remove this line if it exists + # readme.workspace = true + + [dependencies] + # Reuse workspace dependencies when available + serde = { workspace = true } + # Add example-specific dependencies + burn = { path = "../../" } + ``` + +## Required Files and Structure + +### README.md +Each example must include a README.md file with: +- A brief description of what the example demonstrates +- A terminal command showing how to run the example +- Any prerequisites or setup instructions + +Example README structure: +````markdown +# Example Name + +Brief description of what this example demonstrates. + +## Running the Example + +```bash +cargo run --example +``` + +## Prerequisites + +List any prerequisites here. +```` + +### Source Code Structure + +- `src/` directory: Contains the main implementation code +- `examples/` directory: Contains example code + - `.rs`: Example implementation + +## Resource Handling + +- Resources (datasets, models, etc.) should be downloaded in the example code +- Do not track external files in the repository +- Include code to download and prepare resources when the example is run + +## Best Practices + +1. **Code Organization** + - Keep the code modular and well-documented + - Use clear, descriptive variable and function names + - Include comments explaining complex operations + +2. **Error Handling** + - Implement proper error handling + - Provide meaningful error messages + - Handle resource download failures gracefully + +3. **Performance** + - Optimize for reasonable execution time + - Include progress indicators for long-running operations + - Consider adding configuration options for different hardware capabilities + +4. **Documentation** + - Document all public APIs + - Include inline comments for complex logic + - Explain any non-obvious implementation details + +## Submitting Your Example + +1. Ensure your example follows all the guidelines above +2. Test your example thoroughly +3. Create a pull request with: + - A clear description of what the example demonstrates + - Any relevant issue numbers + - Screenshots or output examples (if applicable) + +Feel free to ask questions in the pull request if you need clarification or guidance. \ No newline at end of file diff --git a/contributor-book/src/how-to-read-this-book.md b/contributor-book/src/how-to-read-this-book.md new file mode 100644 index 0000000..4e23cf6 --- /dev/null +++ b/contributor-book/src/how-to-read-this-book.md @@ -0,0 +1,22 @@ +# How to read this book + +Throughout this book, we maintain the following structure. + +## Linking + +When referring to structures or functions within codebase, we provide permalinks to the lines in +specific commits, and indicate them by the relative path of their parent file from the project root. +For example this is a reference to the `Tensor` struct in +[`crates/burn-tensor/src/tensor/api/base.rs`](https://github.com/tracel-ai/burn/blob/e303e31c8bc85486690ff80df65d1e25e16728c4/crates/burn-tensor/src/tensor/api/base.rs#L27) + +When some reference information is useful but is beyond the scope of contributing to Burn, we +provide that information in a footnote. To build on the previous example, the `Tensor` mentioned is +what's referred to as a newtype struct[^1]. + +Direct hyperlinks are for tools and resources that are not part of the Burn project, but are useful +for contributing to it. For example, when working on implementing an operation for autodiff, it can +be useful to use [symbolab](https://www.symbolab.com/) to calculate the left and right partial +derivatives. + +[^1]: For more information on newtype please refer to + [the Advanced Types chapter of the Rust Book](https://doc.rust-lang.org/book/ch19-04-advanced-types.html#using-the-newtype-pattern-for-type-safety-and-abstraction) diff --git a/contributor-book/src/overview.md b/contributor-book/src/overview.md new file mode 100644 index 0000000..aa602df --- /dev/null +++ b/contributor-book/src/overview.md @@ -0,0 +1,28 @@ +# Overview + +Welcome to The Burn Contributor's Book 👋 + +This book will help you get acquainted with the internals of the Burn deep learning framework and +provide some detailed guidance on how to contribute to the project. Before opening a PR, please read +the [Contributing Guidelines](https://github.com/tracel-ai/burn/blob/main/CONTRIBUTING.md). + +We have crafted some sections for you: + +- [Getting Started](./getting-started): Much like the [Burn Book](https://burn.dev/books/burn/) which + targets users, we'll start with the fundamentals, guiding you through tasks like setting up the + development environment, running tests, and what you should check prior to each commit. + +- [Project Architecture](./project-architecture): This section will give you an in-depth look at the + architecture of Burn. + +- [Guides](./guides): We provide some guides on how to do specific tasks, such as adding a new + operations to Burn. + +- [Frequently Encountered Issues](./frequently-encountered-issues): If you are running into an issue + that has you stumped, this is the section to check out prior to asking on the + [Discord](https://discord.gg/uPEBbYYDB6). It's a collection of errors encountered by contributors, + what caused them, and how they were resolved. + +As this book is geared towards contributors and not towards users of Burn, we'll assume you have a +good understanding of software development, but will make efforts to explain anything outside of +that scope, or at least provide links to resources that explain it better than we can. diff --git a/contributor-book/src/project-architecture/README.md b/contributor-book/src/project-architecture/README.md new file mode 100644 index 0000000..d65f714 --- /dev/null +++ b/contributor-book/src/project-architecture/README.md @@ -0,0 +1,19 @@ +# Project Architecture + +This section documents most major architectural decisions with the reasoning behind them. + +**Sections** + +- [Module](./module.md) + - [Optimization](./module.md#optimization) + - [Constraints](./module.md#constraints) + - [Solution](./module.md#solution) +- [Serialization](./serialization.md) + - [Constraints](./serialization.md#constraints) + - [Solution](./serialization.md#solution) + - [Pros](./serialization.md#pros) + - [Cons](./serialization.md#cons) + - [Compatibility](./serialization.md#compatibility) +- [Tensor](./tensor.md) +- [Backend](./backend.md) + - [Autodiff](./backend.md#autodiff) diff --git a/contributor-book/src/project-architecture/backend.md b/contributor-book/src/project-architecture/backend.md new file mode 100644 index 0000000..47cea11 --- /dev/null +++ b/contributor-book/src/project-architecture/backend.md @@ -0,0 +1,45 @@ +# Backend + +The Backend trait abstracts multiple things: + +- Device type +- Float tensor type +- Bool tensor type +- Int tensor type +- Float element type +- Int element type +- Float tensor operations (kernels) +- Int tensor operations (kernels) +- Bool tensor operations (kernels) + +## Element types + +> Warning: there are plans to change this architecture in the near future. + +Even though having one type for tensors is convenient for the tensor API, it can be cumbersome when +implementing a backend. Therefore, backends can decide, through associated types, what types they +want to use for their int, float, and bool tensors. Since float and int can have multiple +precisions, the float and int element types are also associated types that must be declared by the +backend. + +Note that the backend chooses the precision and not the user. Since not all backends will support +the same element types, no assumptions must be made. Therefore, there are no methods on tensors to +change the precision, except for the `to_full_precision` function, which ensures numerical stability +on the current backend. Backend implementations can provide a way to choose the precision, which can +be accomplished with a generic parameter (e.g. `LibTorch`). + +## Operations + +To be as general as possible, tensor operations are implemented as plain functions. There is no +object or self, just functions that take tensors as input and often return tensors as output as +well. Backend implementations are free to use their own patterns to implement these kernels. Note +that Burn is a dynamic graph deep learning framework, so backends may have to implement asynchronous +kernel executions for performance reasons. + +## Autodiff + +As of now, there is only one backend decorator that supports autodiff. It follows the decorator +pattern, making any backend differentiable. However, the `AutodiffBackend` trait abstracts how +gradients are calculated, and other approaches to autodiff might be added later. For more +information about how the current autodiff backend works, you can read this (slightly outdated) +[blog post](https://burn.dev/blog/burn-rusty-approach-to-tensor-handling). diff --git a/contributor-book/src/project-architecture/module-serialization.png b/contributor-book/src/project-architecture/module-serialization.png new file mode 100644 index 0000000..8d5e829 Binary files /dev/null and b/contributor-book/src/project-architecture/module-serialization.png differ diff --git a/contributor-book/src/project-architecture/module.md b/contributor-book/src/project-architecture/module.md new file mode 100644 index 0000000..e921670 --- /dev/null +++ b/contributor-book/src/project-architecture/module.md @@ -0,0 +1,91 @@ +# Module + +Modules are a way of creating neural network structures that can be easily optimized, saved, and +loaded with little to no boilerplate. Unlike other frameworks, a module does not force the +declaration of the forward pass, leaving it up to the implementer to decide how it should be +defined. + +Additionally, most modules are created using a (de)serializable configuration, which defines the +structure of the module and its hyperparameters. Parameters and hyperparameters are not serialized +into the same file, and both are normally necessary to load a module for inference. + +## Optimization + +Optimization is normally done with variants of gradient descent, and it is important to provide an +easy API for optimizing modules. + +### Constraints + +1. **Users should be able to control what is optimized.** Modules can contain anything for maximum + flexibility, but not everything needs to be optimized. +2. **Optimizers should have a serializable state that is updated during training.** Many optimizers + keep track of previous gradients to implement some form of momentum. However, the state can be + anything, not just tensors, allowing for easy implementation of any kind of optimizer. +3. **The learning rate can be updated during training.** Learning rate schedulers are often used + during training and should be considered as a key aspect. + +### Solution + +In the following, the `Module` trait is defined in +[`crates/burn-core/src/module/base.rs`](https://github.com/tracel-ai/burn/blob/81a67b6a0992b9b5c33cda8b9784570143b67319/crates/burn-core/src/module/base.rs#L83) +and the `Optimizer` trait is defined in +[`crates/burn-core/src/optim/base.rs`](https://github.com/tracel-ai/burn/blob/81a67b6a0992b9b5c33cda8b9784570143b67319/crates/burn-core/src/optim/base.rs#L8) + +The solution to this problem comprises multiple parts. Firstly, the `Optimizer` trait is quite +similar to the `Module` trait, in terms of saving and loading the state. Please refer to the +[serialization](./serialization.md) section for more details. + +Secondly, two traits were created. The `Optimizer` trait is general and relatively unopinionated, +with a simple `step` method that takes a learning rate, a module, and the gradients. The other +trait, `SimpleOptimizer`, aims to provide an easier API for implementing new optimizers. The goal is +to allow implementations to avoid handling missing gradients, loading and exporting records, +navigating the module parameter structure, handling tracked and untracked tensors, and other such +tasks. + +Thirdly, each tensor that will be optimized needs to be wrapped into a `Param` struct, which gives +them an ID used for (de)serialization and to associate the state of the optimizer to each parameter. +The `Module` trait has two ways to navigate over parameters. The first one is the `map` function, +which returns `Self` and makes it easy to implement any transformation and mutate all parameters. +The second one is the `visit` function, which has a similar signature but does not mutate the +parameter tensors. + +#### SimpleOptimizer + +Located in +[`crates/burn-core/src/optim/simple/base.rs`](https://github.com/tracel-ai/burn/blob/81a67b6a0992b9b5c33cda8b9784570143b67319/crates/burn-core/src/optim/simple/base.rs#L9), +the `SimpleOptimizer` has two major assumptions: + +1. The state of the optimizer is linked to each parameter. In other words, each parameter has its + own optimizer state, decoupled from the other parameters. +2. The state of the optimizer implements `Record`, `Clone`, and has a `'static` lifetime. + +The benefits of those assumptions materialize in simplicity with little loss in flexibility. The +state associative type is also generic over the dimension, making it extremely easy to include +tensors in the state that share the same dimensionality as its parameter. + +To wrap a simple optimizer into the more general `Optimizer` trait, the `OptimizerAdaptor` struct is +used. + +#### OptimizerAdaptor + +Located in in +[`crates/burn-core/src/optim/simple/adaptor.rs`](https://github.com/tracel-ai/burn/blob/81a67b6a0992b9b5c33cda8b9784570143b67319/crates/burn-core/src/optim/simple/adaptor.rs#L14), +the `OptimizerAdaptor` is a simple struct composed of a `SimpleOptimizer` and a hashmap with all +records associated with each parameter ID. + +When performing an optimization step, the adaptor handles the following: + +1. Updates each parameter tensor in the given module using the `Module::map` function. +2. Checks if a gradient for the current tensor exists. +3. Makes sure that the gradient, the tensor, and the optimizer state associated with the current + parameter are on the same device. The device can be different if the state is loaded from disk to + restart training. +4. Performs the simple optimizer step using the inner tensor since the operations done by the + optimizer should not be tracked in the autodiff graph. +5. Updates the state for the current parameter and returns the updated tensor, making sure it's + properly registered into the autodiff graph if gradients are marked as required. + +Note that a parameter can still be updated by another process, as it is the case with running +metrics used in batch norm. These tensors are still wrapped using the `Param` struct so that they +are included in the module's state and given a proper parameter ID, but they are not registered in +the autodiff graph. diff --git a/contributor-book/src/project-architecture/serialization.md b/contributor-book/src/project-architecture/serialization.md new file mode 100644 index 0000000..97f0dfc --- /dev/null +++ b/contributor-book/src/project-architecture/serialization.md @@ -0,0 +1,124 @@ +# Serialization + +An important aspect of a deep learning framework is the ability to save and load training state to +and from disk. Burn serializes its records with the **burnpack** format, a compact binary container +implemented by the `burn-pack` crate. + +## Constraints + +1. **Users should be able to add any field to a module, even fields that are not serializable.** + + This can include constants, database connections, other module references, or any other + information. Only the parameters (tensors) should be serialized; the structure of the module + itself is encapsulated by its configuration (hyperparameters). + +2. **Records should be decoupled from the backend in use.** + + A record holds plain tensor data ([`TensorData`]), so weights saved with one backend can be + loaded on another. Parameter initialization is lazy, so loading a record does not require + eagerly materializing the module first. + +3. **The format should be fast to load and embeddable.** + + Tensor data is stored contiguously and aligned, so it can be read back with zero-copy / + memory-mapped loading, and a record can be saved straight to bytes for `no-std` environments. + +## The burnpack format + +The `burn-pack` crate is intentionally minimal and tensor-library-agnostic: it knows how to read and +write the container format but has no notion of Burn modules. A burnpack file has three parts: + +```text +┌────────────────────────────────────────────────────────────┐ +│ Header (fixed size) │ +│ magic "BURN", format version, metadata byte length │ +├────────────────────────────────────────────────────────────┤ +│ Metadata (CBOR) │ +│ tensors : map │ +│ dtype, shape, data_offsets, optional param_id │ +│ scalars : map │ +│ metadata : map user key/value pairs │ +├────────────────────────────────────────────────────────────┤ +│ Tensor data section │ +│ each tensor's bytes start on a 256-byte boundary so the │ +│ data can be sliced zero-copy / memory-mapped from a file │ +└────────────────────────────────────────────────────────────┘ +``` + +All multi-byte integers are little-endian. Tensor entries carry an optional `param_id` used to +preserve a parameter's identity across save/load. Besides tensors, a pack can store named **typed +scalars** (integers, floats, booleans), which the optimizer and learning rate scheduler records use +to persist their non-tensor state. + +A pack is written with `burn_pack::Writer` and read back with `burn_pack::Reader`, both operating on +`burn_pack::Tensor` entries plus a scalar map. + +## The three record types + +Higher layers bridge their state to and from burnpack through three record types. Each one can be +serialized to a file (`save` / `load`, appending the `.bpk` extension when the path has none) or to +an in-memory byte buffer (`into_bytes` / `from_bytes`). + +### `ModuleRecord` (`burn-core`, `burn::store`) + +Holds a module's parameters: a flat list of `(path, ParamId, TensorData)` entries keyed by module +path. It is produced and applied through the [`Module`] trait itself rather than a separate codegen +type: + +- `module.into_record()` walks the module with a `ModuleVisitor` (the `Collector`), recording each + float/int/bool parameter under its dotted path. +- `module.load_record(record)` (or the fallible `try_load_record`) walks the module with a + `ModuleMapper` that looks each parameter up by path and loads the matching tensor. + +Load-time behavior is configured with builder methods on the record, ignored when saving: + +- `allow_partial(bool)` — tolerate module parameters absent from the record. +- `validate(bool)` — toggle shape-mismatch / missing-tensor validation. +- `with_dtype_policy(..)` / `cast_to_module_dtype()` — choose whether a parameter adopts the + record's dtype (`DTypePolicy::FromRecord`, the default) or casts the data to the module + parameter's current dtype (`DTypePolicy::CastToModule`). + +The save-side dtype is not configurable: the record stores whatever dtype the module currently +holds. The dtype applied on load is controlled by the record's `DTypePolicy` +(`.cast_to_module_dtype()` / `.with_dtype_policy(..)`). + +This module in `burn-core` is intentionally tiny — no filtering, key remapping, or adapters. The +richer snapshot/import tooling (filtering, key remapping, PyTorch/SafeTensors cross-framework stores) +lives in the `burn-store` crate. + +### `OptimizerRecord` (`burn-optim`) + +Holds an optimizer's per-parameter state. Unlike a module record (keyed by module path), it is keyed +per parameter: each parameter's state is decomposed into tensors named `"{param_id}.{field}"` plus a +few typed scalar entries kept in the burnpack scalar map (including a `__rank` scalar so the state +can be reconstructed without inferring rank from tensor shapes). + +- `optimizer.to_record()` flattens each parameter's `DynState` into tensors and scalars. +- `optimizer.load_record(record)` reconstructs the states (no device argument: tensors load on the + default device and migrate to each parameter's device on the next step). + +### `LrSchedulerRecord` (`burn-optim`) + +Holds a learning rate scheduler's state, which is just a handful of scalars (step counters, current +learning rate) and no tensors. Produced/applied through the [`LrScheduler`] trait's `to_record()` / +`load_record()`. Composed schedulers nest their children's records under an index prefix +(`with_record` / `record`). + +## Checkpointing in `burn-train` + +During training, `burn-train` defines a `Checkpoint` trait (`save(path)` / `load(path)`) implemented +for all three record types — `ModuleRecord`, `OptimizerRecord`, `LrSchedulerRecord` — and for `()` +(a stateless no-op). The `Checkpointer` trait drives periodic saves; the +`FileCheckpointer` writes each record to `{name}-{epoch}.bpk` under the experiment directory. This is +how the model, optimizer, and scheduler are persisted and restored across epochs. + +## Notes + +- There is no `Recorder`, `PrecisionSettings`, `Module::Record` associated type, or `#[derive(Record)]` + any more. All of those were part of the previous serde-based record system, which has been removed. +- Cross-framework import/export (PyTorch `.pt`, SafeTensors) still lives in `burn-store` + (`PytorchStore`, `SafetensorsStore`, `BurnpackStore`). + +[`TensorData`]: https://burn.dev/docs/burn/tensor/struct.TensorData.html +[`Module`]: https://burn.dev/docs/burn/module/trait.Module.html +[`LrScheduler`]: https://burn.dev/docs/burn/lr_scheduler/trait.LrScheduler.html diff --git a/contributor-book/src/project-architecture/tensor.md b/contributor-book/src/project-architecture/tensor.md new file mode 100644 index 0000000..6bee357 --- /dev/null +++ b/contributor-book/src/project-architecture/tensor.md @@ -0,0 +1,65 @@ +# Tensor + +A proper deep learning framework should have a fast tensor implementation with autodiff support, and +Burn is no exception. The tensor API abstracts away backend implementation details and focuses on +usability without compromising performance. To make it as easy as possible to use, there is only one +tensor type, which is different from multiple tensor and deep learning crates in Rust. Generic +parameters are used instead to specialize the tensor type. + +- **B: Backend:** The first argument is the backend on which the tensor implementation lies. +- **const D: usize:** The second argument is the dimensionality of the tensor. +- **K: TensorKind:** The third argument is the tensor kind, which can be either Float, Int or Bool. + By default, the tensor kind is set to Float, so for most tensors, the kind argument is not + necessary. + +Having one struct for tensors reduces the complexity of the tensor API, which also means less +duplicated documentation to write and maintain. + +Tensors are thread-safe, which means that you can send a tensor to another thread, and everything +will work, including auto-differentiation. Note that there are no explicit in-place tensor +operations since all tensor operations take owned tensors as parameters, which make it possible to +mutate them. Tensors can be shared simply by cloning them, but if there is only one reference to a +tensor, the backend implementation is free to reuse the tensor's allocated data. For more +information about how it is done, you can have a look at this +[blog post](https://burn.dev/blog/burn-rusty-approach-to-tensor-handling). + +## Tensor Operations + +Operations on Tensors (sometimes shortened to Ops) are defined in traits (generally part of the +Backend Supertrait) and implemented for the Tensor struct. The appropriate parent trait of an +operation depends on the type of operation: + +- `base` => All tensor kinds should implement these operations (reshape, into_data, etc.). The + implementation is in + [crates/burn-tensor/src/tensor/api/base.rs](https://github.com/tracel-ai/burn/blob/6d96e8d8086d2309c425f2c8a43a8246f8c454d2/crates/burn-tensor/src/tensor/api/base.rs). +- `numeric` => All tensors that are numeric by nature should implement these operations (Add, Sub, + Div, etc.). The implementation is in + [crates/burn-tensor/src/tensor/api/numeric.rs](https://github.com/tracel-ai/burn/blob/6d96e8d8086d2309c425f2c8a43a8246f8c454d2/crates/burn-tensor/src/tensor/api/numeric.rs). +- `Float` => Tensor operations are only available for float tensors. The implementation is in + [burn-tensor/src/tensor/api/float.rs](https://github.com/tracel-ai/burn/blob/6d96e8d8086d2309c425f2c8a43a8246f8c454d2/crates/burn-tensor/src/tensor/api/float.rs). +- `Int` => Tensor operations are only available for int tensors. The implementation is in + [burn-tensor/src/tensor/api/int.rs](https://github.com/tracel-ai/burn/blob/6d96e8d8086d2309c425f2c8a43a8246f8c454d2/crates/burn-tensor/src/tensor/api/int.rs). +- `bool` => Tensor operations are only available for bool tensors. The implementation is in + [burn-tensor/src/tensor/api/bool.rs](https://github.com/tracel-ai/burn/blob/6d96e8d8086d2309c425f2c8a43a8246f8c454d2/crates/burn-tensor/src/tensor/api/bool.rs). + +`Numeric` is directly implemented for `Float` and `Int` tensors, and in general, The implementations +for these methods are calling the corresponding `{Int|Float}` method defined in the backend +supertrait. + +Anything that is implemented by numeric should have an implementation in the `{Int|Float}` traits, +though it may be avoidable if the operation for one type requires casting to the other type. To +provide an example, `powf` should be implemented for `Int` tensors, but it should not be an Int +Tensor Operation. The LHS should be converted to a float, and the output should be converted back to +an int. So it's possible to avoid implementing `IntTensorOp` altogether. + +Additionally there are some operations that should be defined as functions instead of tensor op +methods. These are: + +`module` => These should be exported as functions instead of methods on tensors. The implementation +is in +[crates/burn-tensor/src/tensor/ops/module.rs](https://github.com/tracel-ai/burn/tree/6d96e8d8086d2309c425f2c8a43a8246f8c454d2/crates/burn-tensor/src/tensor/ops/modules). +`activation` => These should also be exported as functions instead of methods on tensors. The +implementation is in +[crates/burn-tensor/src/tensor/ops/activation.rs](https://github.com/tracel-ai/burn/blob/6d96e8d8086d2309c425f2c8a43a8246f8c454d2/crates/burn-tensor/src/tensor/ops/activation.rs). +Note that some activations are just a combination of backend operations and are not declared in +there. diff --git a/crates/burn-autodiff/Cargo.toml b/crates/burn-autodiff/Cargo.toml new file mode 100644 index 0000000..ddc00e9 --- /dev/null +++ b/crates/burn-autodiff/Cargo.toml @@ -0,0 +1,47 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science"] +description = "Automatic differentiation backend for the Burn framework" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "data"] +license.workspace = true +name = "burn-autodiff" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-autodiff" +documentation = "https://docs.rs/burn-autodiff" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["std", "tracing"] +std = ["dep:parking_lot"] +export_tests = [] # check checkpointer is_empty in tests + +tracing = [ + "dep:tracing", + "burn-std/tracing", + "burn-backend/tracing", +] + +[dependencies] +burn-std = { workspace = true } +burn-backend = { workspace = true } + + +derive-new = { workspace = true } +spin = { workspace = true } +parking_lot = { workspace = true, optional = true } +log = { workspace = true } +hashbrown = { workspace = true } +num-traits = { workspace = true } +portable-atomic = { workspace = true } +tracing = { workspace = true, optional = true, features = ["default"] } + +[target.'cfg(not(target_has_atomic = "ptr"))'.dependencies] +portable-atomic-util = { workspace = true } + +[package.metadata.docs.rs] +features = ["default"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-autodiff/LICENSE-APACHE b/crates/burn-autodiff/LICENSE-APACHE new file mode 120000 index 0000000..1cd601d --- /dev/null +++ b/crates/burn-autodiff/LICENSE-APACHE @@ -0,0 +1 @@ +../../LICENSE-APACHE \ No newline at end of file diff --git a/crates/burn-autodiff/LICENSE-MIT b/crates/burn-autodiff/LICENSE-MIT new file mode 120000 index 0000000..b2cfbdc --- /dev/null +++ b/crates/burn-autodiff/LICENSE-MIT @@ -0,0 +1 @@ +../../LICENSE-MIT \ No newline at end of file diff --git a/crates/burn-autodiff/README.md b/crates/burn-autodiff/README.md new file mode 100644 index 0000000..0859a5c --- /dev/null +++ b/crates/burn-autodiff/README.md @@ -0,0 +1,8 @@ +# Burn Autodiff + +> [Burn](https://github.com/tracel-ai/burn) autodiff backend + +[![Current Crates.io Version](https://img.shields.io/crates/v/burn-autodiff.svg)](https://crates.io/crates/burn-autodiff) +[![license](https://shields.io/badge/license-MIT%2FApache--2.0-blue)](https://github.com/tracel-ai/burn-autodiff/blob/master/README.md) + +For now only first order reverse mode autodiff is supported. diff --git a/crates/burn-autodiff/src/backend.rs b/crates/burn-autodiff/src/backend.rs new file mode 100644 index 0000000..66efef1 --- /dev/null +++ b/crates/burn-autodiff/src/backend.rs @@ -0,0 +1,171 @@ +use crate::{ + checkpoint::strategy::{CheckpointStrategy, NoCheckpointing}, + grads::Gradients, + tensor::AutodiffTensor, +}; +use alloc::{format, string::String}; +use core::marker::PhantomData; + +use burn_backend::{ + backend::{AutodiffBackend, Backend, BackendTypes, ExecutionError}, + tensor::{BoolTensor, IntTensor, QuantizedTensor}, +}; + +use burn_backend::distributed::{DistributedParamId, DistributedParams}; + +/// Enable auto-differentiation on a backend. +/// +/// This works as a backend decorator, extending the functionality of any backend with +/// backpropagation. +#[derive(Clone, Copy, Debug, Default)] +pub struct Autodiff { + _b: PhantomData, + _checkpoint_strategy: PhantomData, +} + +impl BackendTypes for Autodiff { + type Device = B::Device; + + type FloatTensorPrimitive = AutodiffTensor; + + type IntTensorPrimitive = B::IntTensorPrimitive; + + type BoolTensorPrimitive = B::BoolTensorPrimitive; + + type QuantizedTensorPrimitive = B::QuantizedTensorPrimitive; + + // A replayed graph would skip re-recording the autodiff tape, so capture + // is not supported under autodiff. + type GraphPrimitive = burn_backend::GraphUnsupported; +} + +impl Backend for Autodiff { + fn ad_enabled(_device: &Self::Device) -> bool { + true + } + + fn name(device: &Self::Device) -> String { + format!("autodiff<{}>", B::name(device)) + } + + fn seed(device: &B::Device, seed: u64) { + B::seed(device, seed) + } + + fn sync(device: &B::Device) -> Result<(), ExecutionError> { + B::sync(device) + } + + fn memory_persistent_allocations< + Output: Send, + Input: Send, + Func: Fn(Input) -> Output + Send, + >( + device: &Self::Device, + input: Input, + func: Func, + ) -> Output { + B::memory_persistent_allocations(device, input, func) + } + + fn memory_cleanup(device: &Self::Device) { + B::memory_cleanup(device) + } + + fn staging<'a, Iter>(data: Iter, device: &Self::Device) + where + Iter: Iterator, + { + B::staging(data, device); + } + + fn supports_dtype(device: &Self::Device, dtype: burn_std::DType) -> bool { + B::supports_dtype(device, dtype) + } + + fn dtype_usage(device: &Self::Device, dtype: burn_std::DType) -> burn_backend::DTypeUsageSet { + B::dtype_usage(device, dtype) + } + + fn device_count(type_id: u16) -> usize { + B::device_count(type_id) + } + + fn flush(device: &Self::Device) { + B::flush(device) + } +} + +impl AutodiffBackend for Autodiff { + type InnerBackend = B; + type Gradients = Gradients; + + fn backward(tensor: AutodiffTensor) -> Gradients { + tensor.backward() + } + + fn grad(tensor: &AutodiffTensor, grads: &Gradients) -> Option { + tensor.grad(grads) + } + + fn grad_remove( + tensor: &AutodiffTensor, + grads: &mut Gradients, + ) -> Option { + tensor.grad_remove(grads) + } + fn inner(tensor: AutodiffTensor) -> B::FloatTensorPrimitive { + tensor.primitive + } + + fn from_inner(tensor: B::FloatTensorPrimitive) -> AutodiffTensor { + AutodiffTensor::new(tensor) + } + + fn grad_replace( + tensor: &AutodiffTensor, + grads: &mut Self::Gradients, + grad: B::FloatTensorPrimitive, + ) { + tensor.grad_replace(grads, grad); + } + + fn int_inner(tensor: IntTensor) -> IntTensor { + tensor + } + + fn bool_inner(tensor: BoolTensor) -> BoolTensor { + tensor + } + + fn int_from_inner(tensor: IntTensor) -> IntTensor { + tensor + } + + fn bool_from_inner(tensor: BoolTensor) -> BoolTensor { + tensor + } + + fn q_inner(tensor: QuantizedTensor) -> QuantizedTensor { + tensor + } + + fn q_from_inner(tensor: QuantizedTensor) -> QuantizedTensor { + tensor + } + + fn set_distributed_params( + tensor: AutodiffTensor, + param_id: DistributedParamId, + ) -> AutodiffTensor { + tensor.grad_distributed(param_id) + } + + fn distributed_params(tensor: &AutodiffTensor) -> Option { + tensor.node.distributed_params.clone() + } + + fn is_distributed(tensor: &AutodiffTensor) -> bool { + tensor.node.distributed_params.is_some() + } +} diff --git a/crates/burn-autodiff/src/checkpoint/base.rs b/crates/burn-autodiff/src/checkpoint/base.rs new file mode 100644 index 0000000..6663a87 --- /dev/null +++ b/crates/burn-autodiff/src/checkpoint/base.rs @@ -0,0 +1,92 @@ +use super::{ + retro_forward::RetroForwards, + state::{BackwardStates, State}, +}; +use crate::collections::HashMap; +use crate::graph::NodeId; + +use alloc::{format, vec, vec::Vec}; +use burn_std::config::{autodiff::AutodiffLogLevel, log_autodiff}; + +#[derive(new, Debug)] +/// Links a [NodeId] to its autodiff graph [NodeRef] +pub(crate) struct NodeTree { + map: HashMap>, +} + +impl NodeTree { + /// Gives the parents of the node in the autodiff graph + pub(crate) fn parents(&self, node_id: &NodeId) -> Option> { + self.map.get(node_id).cloned() + } +} + +#[derive(new, Debug)] +/// Struct responsible of fetching the output for a node in the autodiff graph during a backward pass +pub struct Checkpointer { + backward_states: BackwardStates, + retro_forwards: RetroForwards, + node_tree: NodeTree, +} + +impl Checkpointer { + /// Gives the output of the given node, by recursively asking parents to compute themselves + /// or give their pre-computed tensors. + pub fn retrieve_node_output(&mut self, node_id: NodeId) -> T + where + T: Clone + Send + 'static, + { + let sorted = self.topological_sort(node_id); + let num_nodes = sorted.len(); + log_autodiff(AutodiffLogLevel::Basic, move || { + format!("retrieve_node_output {node_id:?}: {num_nodes} node(s) to compute") + }); + + sorted.into_iter().for_each(|node| { + log_autodiff(AutodiffLogLevel::Full, move || { + format!("execute_retro_forward {node:?}") + }); + self.retro_forwards + .execute_retro_forward(node, &mut self.backward_states) + }); + + self.backward_states.get_state::(&node_id) + } + + /// Sorts the ancestors of NodeId in a way such that all parents come before their children + /// Useful to avoid recursivity later when mutating the states + /// + /// The sort on a compute bound state or a memory bound that is already computed is trivial. + /// The match on State::Computed also serves as a stopping criterion for the sort, + /// we don't need to look higher than that during recursivity. + fn topological_sort(&self, node_id: NodeId) -> Vec { + match self.backward_states.get_state_ref(&node_id) { + Some(state) => match state { + State::Recompute { n_required: _ } => { + let mut sorted = Vec::new(); + let parents = self.node_tree.parents(&node_id).unwrap(); + for parent_node in parents { + let parent_sorted = self.topological_sort(parent_node); + for ps in parent_sorted { + if !sorted.contains(&ps) { + sorted.push(ps) + } + } + } + sorted.push(node_id); + sorted + } + State::Computed { + state_content: _, + n_required: _, + } => vec![node_id], + }, + None => panic!("Node {node_id:?} is not in the backward_states. "), + } + } + + /// Checks if checkpointer has been drained adequately. Useful for testing + pub fn is_empty(&self) -> bool { + self.backward_states.is_empty() && self.retro_forwards.is_empty() + } +} diff --git a/crates/burn-autodiff/src/checkpoint/builder.rs b/crates/burn-autodiff/src/checkpoint/builder.rs new file mode 100644 index 0000000..aac9d3a --- /dev/null +++ b/crates/burn-autodiff/src/checkpoint/builder.rs @@ -0,0 +1,311 @@ +use crate::{ + collections::HashMap, + graph::{ComputingProperty, NodeId}, + tensor::AutodiffTensor, +}; +use alloc::{boxed::Box, vec::Vec}; + +#[cfg(target_has_atomic = "ptr")] +use alloc::sync::Arc; + +#[cfg(not(target_has_atomic = "ptr"))] +use portable_atomic_util::Arc; + +use burn_backend::Backend; +use core::any::Any; + +use super::{ + base::{Checkpointer, NodeTree}, + retro_forward::{RetroForward, RetroForwards}, + state::{BackwardStates, State}, +}; + +#[derive(Debug)] +/// Determines if a node should checkpoint its computed output or its retro_forward for recomputation +/// The action is normally created by the child of the node, once the node is determined to be needed +pub enum CheckpointingAction { + /// The node's already computed output should be saved + Computed { + /// The node + node_id: NodeId, + /// The node's output + state_content: Box, + }, + /// The node should recompute itself when asked + Recompute { + /// The node + node_id: NodeId, + /// How the node should recompute itself + retro_forward: Arc, + }, +} + +// TODO: Remove that when proper client server. +unsafe impl Send for CheckpointingAction {} + +impl CheckpointingAction { + /// Utility function to access the id of the node of the checkpointing action + pub fn id(&self) -> NodeId { + match self { + CheckpointingAction::Computed { + node_id: node_ref, + state_content: _, + } => *node_ref, + CheckpointingAction::Recompute { + node_id: node_ref, + retro_forward: _, + } => *node_ref, + } + } +} + +#[derive(new, Debug, Default)] +/// Accumulates checkpoints as checkpointing actions during the forward pass, +/// and builds a checkpointer right before the backward pass +pub struct CheckpointerBuilder { + explicit_actions: Vec, + backup_actions: Vec, +} + +/// Determines if a checkpoint should impact the n_required values (Main) +/// or if it should just keep the state in case it's required (Backup) +/// +pub(crate) enum ActionType { + /// Explicit actions have been explicitly requested by some operation to retrieve their state + Explicit, + /// Backup actions are not always needed. They exist to save the output of an operation + /// whose child is memory bound, in case the state is indirectly needed when computing + /// the child's retro_forward. If no explicit action ever asks for the child's output, then + /// the backup output will go out of scope when the checkpointer is built. + Backup, +} + +impl CheckpointerBuilder { + pub(crate) fn checkpoint( + &mut self, + tensor: &AutodiffTensor, + action_type: ActionType, + ) { + let action_list = match action_type { + ActionType::Explicit => &mut self.explicit_actions, + ActionType::Backup => &mut self.backup_actions, + }; + match &tensor.node.properties { + ComputingProperty::ComputeBound | ComputingProperty::Ambiguous => { + action_list.push(CheckpointingAction::Computed { + node_id: tensor.node.id, + state_content: Box::new(tensor.primitive.clone()), + }) + } + ComputingProperty::MemoryBound { retro_forward } => { + action_list.push(CheckpointingAction::Recompute { + node_id: tensor.node.id, + retro_forward: retro_forward.clone(), + }) + } + } + } + + pub(crate) fn extend(&mut self, other: CheckpointerBuilder) { + for other_action in other.explicit_actions { + self.explicit_actions.push(other_action) + } + for other_unsure in other.backup_actions { + self.backup_actions.push(other_unsure) + } + } + + pub(crate) fn build(self, node_tree: NodeTree) -> Checkpointer { + let mut backward_states_map = HashMap::new(); + let mut retro_forwards_map = HashMap::new(); + + // Find recursion stopping points + let stop_nodes: Vec = self.find_stop_nodes(); + + // We start by identifying how many times each node will be required. + let n_required_map = self.build_n_required_map(&node_tree, stop_nodes); + + // Then we checkpoint the nodes with the corresponding n_required value + self.insert_checkpoints( + &mut backward_states_map, + &mut retro_forwards_map, + n_required_map, + ); + + Checkpointer::new( + BackwardStates::new(backward_states_map), + RetroForwards::new(retro_forwards_map), + node_tree, + ) + } + + fn find_stop_nodes(&self) -> Vec { + let mut stop_nodes = Vec::default(); + for action in self + .explicit_actions + .iter() + .chain(self.backup_actions.iter()) + { + match action { + CheckpointingAction::Computed { + node_id: node_ref, + state_content: _, + } => stop_nodes.push(*node_ref), + CheckpointingAction::Recompute { + node_id: _, + retro_forward: _, + } => {} + } + } + stop_nodes + } + + fn build_n_required_map( + &self, + node_tree: &NodeTree, + stop_nodes: Vec, + ) -> HashMap { + let mut n_required_map = HashMap::::default(); + + for action in self.explicit_actions.iter() { + match action { + CheckpointingAction::Computed { + node_id: node_ref, + state_content: _, + } => { + let id = *node_ref; + match n_required_map.remove(&id) { + Some(n) => { + n_required_map.insert(id, n + 1); + } + None => { + n_required_map.insert(id, 1); + } + }; + } + CheckpointingAction::Recompute { + node_id: node_ref, + retro_forward: _, + } => { + let id = *node_ref; + Self::update_n_required_of_parents( + id, + &mut n_required_map, + node_tree, + &stop_nodes, + ); + } + } + } + + n_required_map + } + + fn insert_checkpoints( + mut self, + backward_states_map: &mut HashMap, + retro_forward_map: &mut HashMap>, + n_required_map: HashMap, + ) { + // We do not loop over checkpointing actions anymore because they can contain + // duplicates or miss some that are in backup. We loop over the n_required_map + // from which we use the ids to find them again in the checkpointing actions + for (node_id, n_required) in n_required_map { + // We find the checkpointing action for node_id. It's likely in checkpointing_actions + // so we check there first, otherwise it will be in backup. + // Technically it can be there several times but can never be of both types, so we can assume the first we find is fine + + let action = match self + .explicit_actions + .iter() + .position(|action| action.id() == node_id) + { + Some(pos) => self.explicit_actions.remove(pos), + None => { + let pos = self + .backup_actions + .iter() + .position(|action| action.id() == node_id); + self.backup_actions.remove(pos.unwrap_or_else(|| { + panic!("Node {:?} is needed but never checkpointed", node_id) + })) + } + }; + + match action { + CheckpointingAction::Computed { + node_id: _, + state_content, + } => { + self.checkpoint_compute(backward_states_map, node_id, state_content, n_required) + } + CheckpointingAction::Recompute { + node_id: _, + retro_forward, + } => self.checkpoint_lazy( + backward_states_map, + retro_forward_map, + node_id, + retro_forward, + n_required, + ), + }; + } + } + + fn update_n_required_of_parents( + id: NodeId, + n_required_map: &mut HashMap, + node_tree: &NodeTree, + stop_nodes: &Vec, + ) { + match n_required_map.remove(&id) { + Some(n) => { + n_required_map.insert(id, n + 1); + } + None => { + n_required_map.insert(id, 1); + if !stop_nodes.contains(&id) + && let Some(parents) = node_tree.parents(&id) + { + for p in parents { + Self::update_n_required_of_parents( + p, + n_required_map, + node_tree, + stop_nodes, + ); + } + } + } + } + } + + fn checkpoint_compute( + &self, + backward_states_map: &mut HashMap, + node_id: NodeId, + state_content: Box, + n_required: usize, + ) { + backward_states_map.insert( + node_id, + State::Computed { + state_content, + n_required, + }, + ); + } + + fn checkpoint_lazy( + &self, + backward_states_map: &mut HashMap, + retro_forward_map: &mut HashMap>, + node_id: NodeId, + retro_forward: Arc, + n_required: usize, + ) { + retro_forward_map.insert(node_id, retro_forward); + backward_states_map.insert(node_id, State::Recompute { n_required }); + } +} diff --git a/crates/burn-autodiff/src/checkpoint/mod.rs b/crates/burn-autodiff/src/checkpoint/mod.rs new file mode 100644 index 0000000..a67952b --- /dev/null +++ b/crates/burn-autodiff/src/checkpoint/mod.rs @@ -0,0 +1,9 @@ +/// Checkpointer module +pub mod base; +pub(crate) mod builder; +/// RetroForward module +pub mod retro_forward; +/// BackwardStates module +pub mod state; +/// CheckpointStrategy module +pub mod strategy; diff --git a/crates/burn-autodiff/src/checkpoint/retro_forward.rs b/crates/burn-autodiff/src/checkpoint/retro_forward.rs new file mode 100644 index 0000000..4b5e313 --- /dev/null +++ b/crates/burn-autodiff/src/checkpoint/retro_forward.rs @@ -0,0 +1,121 @@ +use crate::collections::HashMap; +use crate::graph::NodeId; + +#[cfg(target_has_atomic = "ptr")] +use alloc::sync::Arc; + +#[cfg(not(target_has_atomic = "ptr"))] +use portable_atomic_util::Arc; + +use core::fmt::Debug; + +use super::state::{BackwardStates, State}; + +/// Definition of the forward function of a node, called during retropropagation only. +/// This is different from the normal forward function because it reads and writes from +/// the [BackwardStates] map instead of having a clear function signature. +pub trait RetroForward: Debug + Send + 'static { + /// Applies the forward pass for retropropagation. + fn forward(&self, states: &mut BackwardStates, out_node: NodeId); +} + +#[derive(new, Debug)] +/// Links [NodeId]s to their corresponding [RetroForward] +pub(crate) struct RetroForwards { + map: HashMap>, +} + +impl RetroForwards { + /// Executes the [RetroForward] for a given [NodeId] if the node's + /// [State] is [State::Recompute], otherwise does nothing. + pub(crate) fn execute_retro_forward( + &mut self, + node_id: NodeId, + backward_states: &mut BackwardStates, + ) { + if let State::Recompute { n_required: _ } = backward_states + .get_state_ref(&node_id) + .unwrap_or_else(|| panic!("Should find node {node_id:?}")) + { + // Retro forwards are always used only once because afterwards their state is computed + let retro_forward = self.map.remove(&node_id).unwrap(); + retro_forward.forward(backward_states, node_id); + } + } + + pub(crate) fn is_empty(&self) -> bool { + self.map.is_empty() + } +} + +#[macro_export] +/// Creates a RetroForward struct for unary scalar operations +macro_rules! retro_unary_scalar { + ( + $name:ident, + $ops:expr + ) => { + #[derive(new, Debug, Clone)] + struct $name { + lhs_id: NodeId, + rhs: Scalar, + _backend: PhantomData, + } + + impl RetroForward for $name { + fn forward(&self, states: &mut BackwardStates, out_node: NodeId) { + let lhs = states.get_state::(&self.lhs_id); + let out = $ops(lhs, self.rhs); + states.save(out_node, out) + } + } + }; +} + +#[macro_export] +/// Creates a RetroForward struct for unary scalar operations +macro_rules! retro_unary { + ( + $name:ident, + $ops:expr + ) => { + #[derive(new, Debug, Clone)] + struct $name { + input_id: NodeId, + _backend: PhantomData, + } + + impl RetroForward for $name { + fn forward(&self, states: &mut BackwardStates, out_node: NodeId) { + let input = states.get_state::(&self.input_id); + let out = $ops(input); + states.save(out_node, out) + } + } + }; +} + +#[macro_export] +/// Creates a RetroForward struct for binary operations +macro_rules! retro_binary { + ( + $name:ident, + $ops:expr + ) => { + #[derive(new, Debug, Clone)] + struct $name { + lhs_id: NodeId, + rhs_id: NodeId, + _backend: PhantomData, + } + + impl RetroForward for $name { + fn forward(&self, states: &mut BackwardStates, out_node: NodeId) { + let lhs = states.get_state::(&self.lhs_id); + let rhs = states.get_state::(&self.rhs_id); + let out = $ops(lhs, rhs); + states.save(out_node, out) + } + } + }; +} diff --git a/crates/burn-autodiff/src/checkpoint/state.rs b/crates/burn-autodiff/src/checkpoint/state.rs new file mode 100644 index 0000000..28b9578 --- /dev/null +++ b/crates/burn-autodiff/src/checkpoint/state.rs @@ -0,0 +1,144 @@ +use core::any::Any; + +use crate::collections::HashMap; +use crate::graph::NodeId; +use alloc::boxed::Box; + +/// In order to accept arbitrary node output in the same hashmap, we need to upcast them to any. +pub(crate) type StateContent = Box; + +#[derive(Debug)] +/// The state contained at one node. Encapsulates the node output if precomputed, +/// or clearly asks that it needs to be recomputed from the parents. +/// Also keeps track of the number of times the state is required so it can be removed +/// from the map of states on its last use. +pub(crate) enum State { + /// The state was not checkpointed, will need to recompute it from the node's parents + Recompute { n_required: usize }, + /// The state was checkpointed or computed during retropropagation and can be directly accessed + Computed { + state_content: StateContent, + n_required: usize, + }, +} + +impl State { + /// Returns a reference to the (not yet) downcasted node output, if checkpointed + pub(crate) fn to_state_content(&self) -> &StateContent { + match self { + State::Recompute { n_required: _ } => { + unreachable!( + "Can't get state content of recompute state. A child has likely been accessed before its parents." + ) + } + State::Computed { + state_content, + n_required: _, + } => state_content, + } + } + + /// Returns a (not yet) downcasted node output, if checkpointed + pub(crate) fn into_state_content(self) -> StateContent { + match self { + State::Recompute { n_required: _ } => { + unreachable!( + "Can't get state content of recompute state. A child has likely been accessed before its parents." + ) + } + State::Computed { + state_content, + n_required: _, + } => state_content, + } + } + + /// Returns the number of time the state is required + pub(crate) fn n_required(&self) -> usize { + match self { + State::Recompute { n_required } => *n_required, + State::Computed { + state_content: _, + n_required, + } => *n_required, + } + } +} + +#[derive(new, Default, Debug)] +/// Links [NodeId]s to their current state +pub struct BackwardStates { + map: HashMap, +} + +impl BackwardStates { + /// Returns the output in the state of the given [NodeId], + /// and decrements the number of times this state is required. + /// This function always gives ownership of the output, but will clone it if needed for further uses. + pub fn get_state(&mut self, node_id: &NodeId) -> T + where + T: Clone + Send + 'static, + { + // Fetch the state and decrement its number of required + let state = self.map.remove(node_id).unwrap(); + let remaining_n_required = state.n_required() - 1; + + // Downcast the state to whatever it is supposed to be + // If still needed after giving ownership, we copy it back to the hashmap + if remaining_n_required > 0 { + let new_stored_state = match state { + State::Recompute { n_required: _ } => unreachable!(), + State::Computed { + state_content, + n_required: _, + } => State::Computed { + state_content, + n_required: remaining_n_required, + }, + }; + + let downcasted = new_stored_state + .to_state_content() + .downcast_ref::() + .unwrap() + .clone(); + + self.insert_state(*node_id, new_stored_state); + + downcasted + } else { + let downcasted = state.into_state_content().downcast::().unwrap(); + *downcasted + } + } + + /// Returns a reference to the [State] of the given node + /// Useful when we need [State] information without needing the underlying tensor + pub(crate) fn get_state_ref(&self, node_id: &NodeId) -> Option<&State> { + self.map.get(node_id) + } + + /// Associates a [State] to its [NodeId] + pub(crate) fn insert_state(&mut self, node_id: NodeId, state: State) { + self.map.insert(node_id, state); + } + + /// Saves the output to the state of the given [NodeId]. + pub fn save(&mut self, node_id: NodeId, saved_output: T) + where + T: Clone + Send + 'static, + { + let n_required = self.get_state_ref(&node_id).unwrap().n_required(); + self.insert_state( + node_id, + State::Computed { + state_content: Box::new(saved_output), + n_required, + }, + ); + } + + pub(crate) fn is_empty(&self) -> bool { + self.map.is_empty() + } +} diff --git a/crates/burn-autodiff/src/checkpoint/strategy.rs b/crates/burn-autodiff/src/checkpoint/strategy.rs new file mode 100644 index 0000000..afa6adf --- /dev/null +++ b/crates/burn-autodiff/src/checkpoint/strategy.rs @@ -0,0 +1,107 @@ +use core::fmt::Debug; + +use burn_backend::Backend; + +#[cfg(target_has_atomic = "ptr")] +use alloc::sync::Arc; + +#[cfg(not(target_has_atomic = "ptr"))] +use portable_atomic_util::Arc; + +use crate::{graph::ComputingProperty, tensor::AutodiffTensor}; + +use super::{ + builder::{ActionType, CheckpointerBuilder}, + retro_forward::RetroForward, +}; + +/// Strategy for the amount of checkpointing to do during autodiff +pub trait CheckpointStrategy: Clone + Copy + Debug + Default + Send + Sync + 'static { + /// May modify the compute property depending on the strategy + fn compute_property(retro_forward: R) -> ComputingProperty; + + /// Checkpoints parents if necessary in the strategy + fn checkpoint_parents<'a, B2, A>( + parents: A, + builder: &mut CheckpointerBuilder, + ) -> Result<(), CheckpointingError> + where + B2: Backend, + A: IntoIterator>; +} + +#[derive(Debug)] +/// Error that can happen when trying to checkpoint a tensor. +pub enum CheckpointingError { + /// When a parent is untracked, we can't easily checkpoint its state, since we don't know the + /// requirements in advanced. + UntrackedParent, +} + +#[derive(Clone, Copy, Debug, Default)] +/// All operations are considered compute bound, notwithstanding how they are marked +pub struct NoCheckpointing {} + +impl CheckpointStrategy for NoCheckpointing { + /// An operation marked as memory bound is actually compute bound. + fn compute_property(_retro_forward: R) -> ComputingProperty { + ComputingProperty::ComputeBound + } + + /// An operation marked as memory bound is actually compute bound. + /// It's therefore useless to checkpoint the parents + fn checkpoint_parents<'a, B2, A>( + _parents: A, + _builder: &mut CheckpointerBuilder, + ) -> Result<(), CheckpointingError> + where + B2: Backend, + A: IntoIterator>, + { + // Nothing to do here + Ok(()) + } +} + +#[derive(Clone, Copy, Debug, Default)] +/// Operation properties are as they are marked (compute or memory bound) +pub struct BalancedCheckpointing {} + +impl CheckpointStrategy for BalancedCheckpointing { + /// An operation marked as memory bound is memory bound. + /// When memory bound, an operation needs to save its RetroForward + fn compute_property(retro_forward: R) -> ComputingProperty { + ComputingProperty::MemoryBound { + retro_forward: Arc::new(retro_forward), + } + } + + /// An operation marked as memory bound is really memory bound. + /// Since the operation may not checkpoint its parents but may need them indirectly + /// if asked to recompute itself, the method needs to know the parent tensors to maybe checkpoint them + fn checkpoint_parents<'a, B2, A>( + parents: A, + builder: &mut CheckpointerBuilder, + ) -> Result<(), CheckpointingError> + where + B2: Backend, + A: IntoIterator>, + { + let mut can_checkpoint = true; + + for tensor in parents.into_iter() { + if let crate::graph::Requirement::None = tensor.node.requirement { + can_checkpoint = false; + } else { + builder.checkpoint(tensor, ActionType::Backup); + } + } + + if !can_checkpoint { + *builder = CheckpointerBuilder::default(); + return Err(CheckpointingError::UntrackedParent); + } + + Ok(()) + } +} diff --git a/crates/burn-autodiff/src/distributed.rs b/crates/burn-autodiff/src/distributed.rs new file mode 100644 index 0000000..7e55e20 --- /dev/null +++ b/crates/burn-autodiff/src/distributed.rs @@ -0,0 +1,59 @@ +use std::marker::PhantomData; + +use crate::{collections::HashMap, grads::DistributedRegistration}; +use burn_backend::{ + Backend, TensorPrimitive, + distributed::{DistributedParams, TensorRef}, +}; +use burn_std::container::TensorContainer; + +use crate::{NodeId, grads::GradID}; + +/// Submits sync operations on gradient registrations. +pub(crate) struct DistributedGradientRegistration { + n_required_map: HashMap, + sharded_parameters_map: HashMap, + _b: PhantomData, +} + +impl DistributedGradientRegistration { + /// Creates a new registration and immediately registers the distributed parameters + /// with the sync server so it can coordinate gradient reductions across devices. + pub(crate) fn new( + n_required_map: HashMap, + sharded_parameters_map: HashMap, + device: B::Device, + ) -> Self { + // For DDP, we register the distributed parameters of the tensors' nodes used in the graph and the number of times they + // appear as nodes to know when to launch gradients reducing. + if !sharded_parameters_map.is_empty() { + B::register_sync_parameters( + &device, + sharded_parameters_map.values().cloned().collect(), + ); + } + + Self { + n_required_map, + sharded_parameters_map, + _b: PhantomData, + } + } +} + +impl DistributedRegistration for DistributedGradientRegistration { + fn on_register(&mut self, id: &NodeId, container: &mut TensorContainer) { + if let Some(sharded_params) = self.sharded_parameters_map.get(id) { + let n_required = self.n_required_map.get_mut(id).unwrap(); + *n_required -= 1; + + if *n_required == 0 { + let tensor_ref = container + .get_mut_ref::>(&id.value) + .unwrap(); + let tensor_ref = TensorRef(tensor_ref.get_mut_ref()); + B::submit_gradient_sync(tensor_ref, sharded_params.clone()); + } + } + } +} diff --git a/crates/burn-autodiff/src/grads.rs b/crates/burn-autodiff/src/grads.rs new file mode 100644 index 0000000..5e332ab --- /dev/null +++ b/crates/burn-autodiff/src/grads.rs @@ -0,0 +1,146 @@ +use alloc::boxed::Box; +use burn_backend::{Backend, TensorMetadata, TensorPrimitive, tensor::FloatTensor}; +use burn_std::tensor::container::TensorContainer; + +use crate::{ + NodeId, + graph::{NodeRef, Requirement}, + tensor::AutodiffTensor, +}; + +#[cfg(feature = "std")] +use crate::collections::HashMap; +#[cfg(feature = "std")] +use burn_backend::distributed::DistributedParams; + +/// Gradient identifier. +pub type GradID = u64; + +#[cfg(feature = "std")] +#[derive(Clone)] +pub(crate) struct GradSyncContext { + pub n_required_map: HashMap, + pub distributed_params: HashMap, +} + +/// Hook type executed when a gradient is registered. +type OnRegisterHook = Box) + Send + Sync>; + +/// Trait for registering distributed gradients. +pub trait DistributedRegistration: Send + Sync { + /// Performs distributed registration operations on the tensor with the corresponding [`NodeId`]. + fn on_register(&mut self, node_id: &NodeId, container: &mut TensorContainer); +} + +#[derive(Default)] +pub(crate) enum BackwardMode { + #[default] + Standard, + // Distributed registration hook. + #[cfg(feature = "std")] + Distributed(Box Box>), +} + +/// Gradients container used during the backward pass. +pub struct Gradients { + container: TensorContainer, + /// Optional hook called after each gradient is registered, used to trigger + /// distributed gradient synchronization operations. + on_register: Option, +} + +impl Gradients { + /// Creates a new gradients container. + pub fn new(root_node: NodeRef, root_tensor: FloatTensor) -> Self { + Self::new_with_hook::(root_node, root_tensor, None) + } + + /// Creates a new gradients container. + fn new_with_hook( + root_node: NodeRef, + root_tensor: FloatTensor, + on_register: Option, + ) -> Self { + let mut gradients = Self { + container: TensorContainer::new(), + on_register, + }; + gradients.register::( + root_node.id, + B::float_ones( + root_tensor.shape(), + &root_tensor.device(), + root_tensor.dtype().into(), + ), + ); + gradients + } + + /// Creates a new gradients container with a registration hook for distributed gradients. + #[cfg(feature = "std")] + pub fn new_distributed( + root_node: NodeRef, + root_tensor: FloatTensor, + mut reg: Box, + ) -> Self { + let on_register: Option = Some(Box::new(move |id, container| { + reg.on_register(id, container); + })); + Self::new_with_hook::(root_node, root_tensor, on_register) + } + + /// Consumes the gradients for a given tensor. + /// + /// Each tensor should be consumed exactly 1 time if its gradients are only required during the + /// backward pass, otherwise, it may be consume multiple times. + pub fn consume(&mut self, node: &NodeRef) -> FloatTensor { + match node.requirement { + Requirement::Grad => self + .container + .get::>(&node.id.value) + .map(|tensor| tensor.tensor()) + .expect("Can't consume the gradients before they are registered at least once."), + Requirement::GradInBackward => self + .container + .remove::>(&node.id.value) + .map(|tensor| tensor.tensor()) + .expect("Can't consume the gradients before they are registered at least once."), + Requirement::None => panic!("Trying to consume the gradients for an untracked tensor"), + } + } + + /// Removes a grad tensor from the container. + pub fn remove(&mut self, tensor: &AutodiffTensor) -> Option> { + self.container + .remove::>(&tensor.node.id.value) + .map(|tensor| tensor.tensor()) + } + + /// Gets a grad tensor from the container. + pub fn get(&self, tensor: &AutodiffTensor) -> Option> { + self.container + .get::>(&tensor.node.id.value) + .map(|tensor| tensor.tensor()) + } + + /// Register a grad tensor in the container. + /// + /// If the tensor already exists, add both tensors together before saving the result. + /// + /// If the registered tensor is distributed, launches a syncing operation on the gradients. + pub fn register(&mut self, node_id: NodeId, value: FloatTensor) { + let out = + if let Some(tensor_old) = self.container.remove::>(&node_id.value) { + B::float_add(value, tensor_old.tensor()) + } else { + value + }; + + self.container + .register::>(node_id.value, TensorPrimitive::Float(out)); + + if let Some(hook) = &mut self.on_register { + hook(&node_id, &mut self.container); + } + } +} diff --git a/crates/burn-autodiff/src/graph/base.rs b/crates/burn-autodiff/src/graph/base.rs new file mode 100644 index 0000000..0f4d38e --- /dev/null +++ b/crates/burn-autodiff/src/graph/base.rs @@ -0,0 +1,26 @@ +use super::NodeId; +use crate::{checkpoint::base::Checkpointer, grads::Gradients, graph::Parent}; +use alloc::boxed::Box; + +use burn_backend::distributed::DistributedParams; + +/// Backward step for reverse mode autodiff. +pub trait Step: Send + core::fmt::Debug { + /// Executes the step and consumes it. + fn step(self: Box, grads: &mut Gradients, checkpointer: &mut Checkpointer); + /// Depth of the operation relative to the first node added to a graph. + fn depth(&self) -> usize; + /// The node associated to the step. + fn node(&self) -> NodeId; + /// The parents of the node associated to the step. + fn parents(&self) -> &[Parent]; + + /// Returns the [`DistributedParams`] of the node's tensor associated to the step. + /// + /// Defaults to `None`; steps that carry distributed parameters override this. + fn distributed_params(&self) -> Option { + None + } +} + +pub type StepBoxed = Box; diff --git a/crates/burn-autodiff/src/graph/mod.rs b/crates/burn-autodiff/src/graph/mod.rs new file mode 100644 index 0000000..7e78bed --- /dev/null +++ b/crates/burn-autodiff/src/graph/mod.rs @@ -0,0 +1,9 @@ +mod base; +mod node; +mod requirement; + +pub mod traversal; + +pub use base::*; +pub use node::*; +pub use requirement::*; diff --git a/crates/burn-autodiff/src/graph/node.rs b/crates/burn-autodiff/src/graph/node.rs new file mode 100644 index 0000000..3c271b7 --- /dev/null +++ b/crates/burn-autodiff/src/graph/node.rs @@ -0,0 +1,96 @@ +use alloc::vec::Vec; + +use burn_backend::distributed::DistributedParams; + +#[cfg(target_has_atomic = "64")] +use core::sync::atomic::{AtomicU64, Ordering}; +#[cfg(not(target_has_atomic = "64"))] +use portable_atomic::{AtomicU64, Ordering}; + +#[cfg(target_has_atomic = "ptr")] +use alloc::sync::Arc; + +#[cfg(not(target_has_atomic = "ptr"))] +use portable_atomic_util::Arc; + +use crate::checkpoint::retro_forward::RetroForward; +use crate::runtime::AutodiffClientImpl; + +use super::Requirement; + +#[derive(Debug, Clone)] +pub enum ComputingProperty { + ComputeBound, + MemoryBound { + retro_forward: Arc, + }, + Ambiguous, // Maybe autotune someday +} + +/// This is safe only because we only call RetroForward on the autodiff server. +/// Therefore, the trait will never be used by multiple threads at the same time. +/// +/// TODO: Find a way to avoid cloning the compute property, which will remove the need to add the +/// Arc, which will make (dyn RetroForward) safely implement Send. +unsafe impl Send for ComputingProperty {} +/// unsafe Sync is required because Send is only implemented for Arc, not Arc. +unsafe impl Sync for ComputingProperty {} + +/// A node contains graph metadata and should be used wrapped in an Arc for cheap cloning. +#[derive(new, Debug)] +pub struct Node { + pub parents: Vec, + pub order: usize, + pub id: NodeId, + pub requirement: Requirement, + pub properties: ComputingProperty, + pub client: AutodiffClientImpl, + pub distributed_params: Option, +} +pub type NodeRef = Arc; + +#[derive(new, Debug, Clone, PartialEq, Eq)] +pub struct Parent { + pub id: NodeId, +} + +impl Node { + /// Returns the [node](Node) only if gradients are required. + pub fn clone_if_require_grad(self: &Arc) -> Option { + match self.requirement.is_none() { + true => None, + false => Some(self.clone()), + } + } +} + +/// Unique identifier generated for each node. +#[derive(Clone, Hash, PartialEq, Eq, Debug, Copy, Ord, PartialOrd)] +pub struct NodeId { + /// The integer representation of the id + pub value: u64, +} + +impl core::fmt::Display for NodeId { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_fmt(format_args!("NodeId({})", self.value)) + } +} + +impl NodeId { + /// Create a unique [node id](NodeId). + pub fn new() -> Self { + static COUNTER: AtomicU64 = AtomicU64::new(0); + let value = COUNTER.fetch_add(1, Ordering::Relaxed); + if value == u64::MAX { + panic!("NodeId overflowed"); + } + Self { value } + } +} + +impl Default for NodeId { + fn default() -> Self { + Self::new() + } +} diff --git a/crates/burn-autodiff/src/graph/requirement.rs b/crates/burn-autodiff/src/graph/requirement.rs new file mode 100644 index 0000000..ce7291c --- /dev/null +++ b/crates/burn-autodiff/src/graph/requirement.rs @@ -0,0 +1,38 @@ +use super::NodeRef; + +/// Requirement for each tensor in the graph. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Requirement { + /// Operations that require gradients. + Grad, + /// Operations that require gradients only for backprop. + GradInBackward, + /// Operations that don't need gradients, therefore not to be included in the graph. + None, +} + +impl Requirement { + /// Returns true if gradients are not required. + pub fn is_none(&self) -> bool { + matches!(self, Self::None) + } + /// Returns the right requirement from a list of nodes. + pub fn from_nodes(nodes: &[NodeRef]) -> Self { + if nodes.len() == 1 { + return nodes[0].requirement.infer(&Requirement::None); + } + + nodes + .iter() + .map(|node| node.requirement) + .reduce(|acc, requirement| requirement.infer(&acc)) + .unwrap_or(Requirement::None) + } + + fn infer(&self, other: &Self) -> Self { + match self.is_none() && other.is_none() { + true => Self::None, + false => Self::GradInBackward, + } + } +} diff --git a/crates/burn-autodiff/src/graph/traversal.rs b/crates/burn-autodiff/src/graph/traversal.rs new file mode 100644 index 0000000..0a0afcb --- /dev/null +++ b/crates/burn-autodiff/src/graph/traversal.rs @@ -0,0 +1,74 @@ +use super::{Step, StepBoxed}; +use crate::{ + NodeId, + collections::{HashMap, HashSet}, + graph::Parent, +}; +use alloc::vec::Vec; + +/// Breadth for search algorithm. +pub struct BreadthFirstSearch; + +pub trait TraversalItem { + fn id(&self) -> NodeId; + fn parents(&self) -> &[Parent]; + fn parent_nodes(&self) -> Vec { + self.parents().iter().map(|p| p.id).collect() + } +} + +impl BreadthFirstSearch { + /// Traverse the graph of backward steps from a root node. + pub fn traverse( + &self, + root_id: NodeId, + root_step: I, + steps: &mut HashMap, + mut callback: F, + ) where + F: FnMut(NodeId, I), + I: TraversalItem, + { + let mut visited = HashSet::new(); + let mut parents = Vec::new(); + + visited.insert(root_id); + parents.append(&mut root_step.parent_nodes()); + + callback(root_id, root_step); + + while let Some(id) = parents.pop() { + let step = match steps.remove(&id) { + Some(step) => step, + None => continue, + }; + + let step_node = step.id(); + let step_parents = step.parent_nodes(); + + if visited.contains(&step_node) { + continue; + } + + visited.insert(step_node); + + for id in step_parents.iter() { + if !visited.contains(id) { + parents.push(*id); + } + } + + callback(step_node, step); + } + } +} + +impl TraversalItem for StepBoxed { + fn id(&self) -> NodeId { + Step::node(self.as_ref()) + } + + fn parents(&self) -> &[Parent] { + Step::parents(self.as_ref()) + } +} diff --git a/crates/burn-autodiff/src/lib.rs b/crates/burn-autodiff/src/lib.rs new file mode 100644 index 0000000..32b263c --- /dev/null +++ b/crates/burn-autodiff/src/lib.rs @@ -0,0 +1,46 @@ +#![cfg_attr(not(feature = "std"), no_std)] +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] + +//! # Burn Autodiff +//! +//! This autodiff library is a part of the Burn project. It is a standalone crate +//! that can be used to perform automatic differentiation on tensors. It is +//! designed to be used with the Burn Tensor crate, but it can be used with any +//! tensor library that implements the `Backend` trait. + +#[macro_use] +extern crate derive_new; + +extern crate alloc; + +/// Checkpoint module. +pub mod checkpoint; +#[cfg(feature = "std")] +/// Distributed utils. +pub mod distributed; +/// Gradients module. +pub mod grads; +/// Operation module. +pub mod ops; + +pub(crate) mod graph; +// Exported for backend extension +pub use graph::NodeId; +pub(crate) mod tensor; +pub(crate) mod utils; + +mod backend; + +pub(crate) mod runtime; + +pub use backend::*; + +/// A facade around for HashMap and HashSet. +/// This avoids elaborate import wrangling having to happen in every module. +mod collections { + #[cfg(not(feature = "std"))] + pub use hashbrown::{HashMap, HashSet}; + #[cfg(feature = "std")] + pub use std::collections::{HashMap, HashSet}; +} diff --git a/crates/burn-autodiff/src/ops/activation.rs b/crates/burn-autodiff/src/ops/activation.rs new file mode 100644 index 0000000..4be3912 --- /dev/null +++ b/crates/burn-autodiff/src/ops/activation.rs @@ -0,0 +1,167 @@ +use core::marker::PhantomData; + +use crate::{ + Autodiff, + checkpoint::{ + base::Checkpointer, retro_forward::RetroForward, state::BackwardStates, + strategy::CheckpointStrategy, + }, + grads::Gradients, + graph::NodeId, + ops::{Backward, Ops, OpsKind, unary}, + retro_unary, +}; +use burn_backend::{Backend, ops::ActivationOps, tensor::FloatTensor}; + +impl ActivationOps> for Autodiff { + fn gelu(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Gelu; + + retro_unary!(RetroGelu, B::gelu); + + impl Backward for Gelu { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + + unary::(ops.parents, ops.node, grads, |grad| { + B::gelu_backward(input, grad) + }); + } + } + + match Gelu + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroGelu::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::gelu(tensor.primitive.clone())) + } + OpsKind::UnTracked(prep) => prep.finish(B::gelu(tensor.primitive)), + } + } + + fn relu(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Relu; + + retro_unary!(RetroRelu, B::relu); + + impl Backward for Relu { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let state = checkpointer.retrieve_node_output(ops.state); + unary::(ops.parents, ops.node, grads, |grad| { + B::relu_backward(state, grad) + }); + } + } + + match Relu + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroRelu::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::relu(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::relu(tensor.primitive)), + } + } + + fn sigmoid(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Sigmoid; + + retro_unary!(RetroSigmoid, B::sigmoid); + + impl Backward for Sigmoid { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + let output = B::sigmoid(input); + unary::(ops.parents, ops.node, grads, |grad| { + B::sigmoid_backward(output, grad) + }); + } + } + + match Sigmoid + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroSigmoid::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::sigmoid(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::sigmoid(tensor.primitive)), + } + } + + fn log_sigmoid(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct LogSigmoid; + + retro_unary!(RetroLogSigmoid, B::log_sigmoid); + + impl Backward for LogSigmoid { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + + unary::(ops.parents, ops.node, grads, |grad| { + B::log_sigmoid_backward(input, grad) + }); + } + } + + match LogSigmoid + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroLogSigmoid::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::log_sigmoid(tensor.primitive.clone())) + } + OpsKind::UnTracked(prep) => prep.finish(B::log_sigmoid(tensor.primitive)), + } + } +} diff --git a/crates/burn-autodiff/src/ops/backward.rs b/crates/burn-autodiff/src/ops/backward.rs new file mode 100644 index 0000000..548acd7 --- /dev/null +++ b/crates/burn-autodiff/src/ops/backward.rs @@ -0,0 +1,88 @@ +use super::{Ops, OpsPrep}; +use crate::{ + checkpoint::{base::Checkpointer, builder::CheckpointerBuilder, strategy::CheckpointStrategy}, + grads::Gradients, + graph::{ComputingProperty, NodeRef, Requirement}, + utils::duplicate, +}; +use burn_backend::Backend; + +/// Trait for all operations. +/// +/// # Notes +/// +/// Concrete types implementing this trait should not have any state. +/// If a state is necessary during the backward pass, +/// they should be declared with the associated type 'State'. +pub trait Backward: Send + core::fmt::Debug +where + Self: Sized + 'static, + B: Backend, +{ + /// Associated type to compute the backward pass. + type State: Clone + Send + core::fmt::Debug + 'static; + + /// The backward pass. + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ); + + /// Prepare the backward ops. + fn prepare( + self, + nodes: [NodeRef; N], + ) -> OpsPrep { + let requirement = Requirement::from_nodes(&nodes); + OpsPrep::new( + nodes, + requirement, + self, + ComputingProperty::Ambiguous, // If not specified we start with ambiguous + CheckpointerBuilder::default(), + ) + } +} + +/// Execute a binary operation during the backward step. +pub fn binary( + parents: [Option; 2], + node: NodeRef, + grads: &mut Gradients, + func_lhs: FLhs, + func_rhs: FRhs, +) where + B: Backend, + FLhs: FnOnce(B::FloatTensorPrimitive) -> B::FloatTensorPrimitive, + FRhs: FnOnce(B::FloatTensorPrimitive) -> B::FloatTensorPrimitive, +{ + let [grad_4lhs, grad_4rhs] = duplicate(&parents, Some(grads.consume::(&node))); + let [node_lhs, node_rhs] = parents; + + if let Some(node) = node_lhs { + let grad = func_lhs(grad_4lhs.unwrap()); + grads.register::(node.id, grad) + } + + if let Some(node) = node_rhs { + let grad = func_rhs(grad_4rhs.unwrap()); + grads.register::(node.id, grad) + } +} + +/// Execute a unary operation during the backward step. +pub fn unary(parents: [Option; 1], node: NodeRef, grads: &mut Gradients, func: F) +where + B: Backend, + F: FnOnce(B::FloatTensorPrimitive) -> B::FloatTensorPrimitive, +{ + let [parent_node] = parents; + let grad = grads.consume::(&node); + + if let Some(node) = parent_node { + let grad = func(grad); + grads.register::(node.id, grad) + } +} diff --git a/crates/burn-autodiff/src/ops/base.rs b/crates/burn-autodiff/src/ops/base.rs new file mode 100644 index 0000000..b124544 --- /dev/null +++ b/crates/burn-autodiff/src/ops/base.rs @@ -0,0 +1,327 @@ +use super::Backward; +use crate::{ + checkpoint::{ + base::Checkpointer, + builder::{ActionType, CheckpointerBuilder}, + retro_forward::RetroForward, + strategy::CheckpointStrategy, + }, + grads::Gradients, + graph::{ComputingProperty, NodeId, NodeRef, Parent, Requirement, Step}, + tensor::AutodiffTensor, +}; +use alloc::boxed::Box; +use burn_backend::{Backend, TensorMetadata, tensor::FloatTensor}; +use burn_std::Shape; +use core::marker::PhantomData; + +use burn_backend::distributed::DistributedParams; + +/// Operation in preparation. +/// +/// Each mode has its own set of functions to minimize cloning for unused backward states. +#[derive(new)] +pub struct OpsPrep { + nodes: [NodeRef; N], + requirement: Requirement, + backward: Backward, + compute_property: ComputingProperty, + checkpointer_builder: CheckpointerBuilder, + checkpoint_strategy: PhantomData, + phantom_backend: PhantomData, + phantom_state: PhantomData, + marker: PhantomData, +} + +/// Operation is initialized +pub struct Init; +/// Operation has been tagged as memory bound +pub struct MemoryBound; +/// Memory bound operation has received its RetroForward +pub struct MemoryBoundRetroForward; +/// Operation's compute property is fixed +pub struct ComputePropertyDone; +/// Tracked operation tag. +pub struct Tracked; +/// Untracked operation tag. +pub struct UnTracked; + +impl OpsPrep +where + B: Backend, + BO: Backward, +{ + /// Indicates that the operation is compute bound, meaning its computation + /// is heavy and should not be recomputed + pub fn compute_bound(self) -> OpsPrep { + OpsPrep::new( + self.nodes, + self.requirement, + self.backward, + ComputingProperty::ComputeBound, + self.checkpointer_builder, + ) + } + + /// Indicates that the operation is memory bound, meaning its computation + /// is light and can be recomputed + pub fn memory_bound(self) -> OpsPrep { + OpsPrep::new( + self.nodes, + self.requirement, + self.backward, + self.compute_property, + self.checkpointer_builder, + ) + } +} + +impl OpsPrep +where + B: Backend, + BO: Backward, + C: CheckpointStrategy, +{ + /// Registers the retro forward, if needed + pub fn retro_forward( + self, + retro_forward: R, + ) -> OpsPrep { + OpsPrep::new( + self.nodes, + self.requirement, + self.backward, + C::compute_property(retro_forward), + self.checkpointer_builder, + ) + } +} + +impl OpsPrep +where + B: Backend, + BO: Backward, + C: CheckpointStrategy, +{ + /// Checkpoints the parents, if needed + pub fn parents<'a, B2, A>(mut self, parents: A) -> OpsPrep + where + B2: Backend, + A: IntoIterator>, + { + let compute_property = match C::checkpoint_parents(parents, &mut self.checkpointer_builder) + { + Ok(..) => self.compute_property, + Err(..) => ComputingProperty::ComputeBound, + }; + + OpsPrep::new( + self.nodes, + self.requirement, + self.backward, + compute_property, + self.checkpointer_builder, + ) + } +} + +impl OpsPrep +where + B: Backend, + BO: Backward, +{ + /// Prepare a stateless operation. + pub fn stateless(self, output: FloatTensor) -> AutodiffTensor { + match self.stateful() { + OpsKind::Tracked(prep) => prep.finish((), output), + OpsKind::UnTracked(prep) => prep.finish(output), + } + } +} + +impl OpsPrep +where + B: Backend, + S: Clone + Send + core::fmt::Debug + 'static, + BO: Backward, +{ + /// Prepare an operation that requires a state during the backward pass. + pub fn stateful(self) -> OpsKind { + match self.requirement.is_none() { + false => OpsKind::Tracked(OpsPrep::new( + self.nodes, + self.requirement, + self.backward, + self.compute_property, + self.checkpointer_builder, + )), + true => OpsKind::UnTracked(OpsPrep::new( + self.nodes, + self.requirement, + self.backward, + self.compute_property, + self.checkpointer_builder, + )), + } + } +} + +impl OpsPrep +where + B: Backend, + S: Clone + Send + core::fmt::Debug + 'static, + BO: Backward, +{ + /// Finish the preparation of an untracked operation and returns the output tensor. + pub fn finish(self, output: FloatTensor) -> AutodiffTensor { + let output = AutodiffTensor::from_parents( + output, + &self.nodes, + self.requirement, + self.compute_property, + ); + let parents = self.nodes.map(|node| node.clone_if_require_grad()); + let ops = Ops::new(parents, output.node.clone(), ()); + + // We register the ops in the graph even if untracked, otherwise memory bound operations + // that have an untracked parent would not be able to retrieve it + output.register_step(UntrackedOpsStep::new(ops), self.checkpointer_builder) + } +} + +impl OpsPrep +where + B: Backend, + S: Clone + Send + core::fmt::Debug + 'static, + BO: Backward, +{ + /// Finish the preparation of a tracked operation and returns the output tensor. + pub fn finish(self, state: S, output: FloatTensor) -> AutodiffTensor { + let output = AutodiffTensor::from_parents( + output, + &self.nodes, + self.requirement, + self.compute_property, + ); + let parents = self.nodes.map(|node| node.clone_if_require_grad()); + let ops = Ops::new(parents, output.node.clone(), state); + + output.register_step(OpsStep::new(ops, self.backward), self.checkpointer_builder) + } + + /// Checkpoints the tensor + pub fn checkpoint(&mut self, tensor: &AutodiffTensor) -> NodeId { + self.checkpointer_builder + .checkpoint(tensor, ActionType::Explicit); + + tensor.node.id + } +} + +/// Enum used before finishing tracked and untracked operations. +pub enum OpsKind { + /// Tracked operation preparation. + Tracked(OpsPrep), + /// Untracked operation preparation. + UnTracked(OpsPrep), +} + +/// Operation containing its parent nodes, its own node and the backward step state. +#[derive(new, Debug)] +pub struct Ops { + /// Parents nodes. + pub parents: [Option; N], + /// The node. + pub node: NodeRef, + /// The state. + pub state: S, +} + +/// Operation implementing backward [step](Step) with type erasing. +#[derive(new, Debug)] +struct OpsStep +where + B: Backend, + T: Backward, + SB: Clone + Send + core::fmt::Debug + 'static, +{ + ops: Ops, + backward: T, + phantom: PhantomData, +} + +impl Step for OpsStep +where + B: Backend, + T: Backward, + SB: Clone + Send + core::fmt::Debug + 'static, +{ + fn step(self: Box, grads: &mut Gradients, checkpointer: &mut Checkpointer) { + self.backward.backward(self.ops, grads, checkpointer); + } + + fn node(&self) -> NodeId { + self.ops.node.id + } + + fn parents(&self) -> &[Parent] { + &self.ops.node.parents + } + + fn depth(&self) -> usize { + self.ops.node.order + } + + fn distributed_params(&self) -> Option { + self.ops.node.distributed_params.clone() + } +} + +#[derive(new, Debug)] +struct UntrackedOpsStep { + ops: Ops<(), N>, +} + +impl Step for UntrackedOpsStep { + fn step(self: Box, _grads: &mut Gradients, _checkpointer: &mut Checkpointer) { + // Nothing to do + } + + fn node(&self) -> NodeId { + self.ops.node.id + } + + fn parents(&self) -> &[Parent] { + &self.ops.node.parents + } + fn depth(&self) -> usize { + self.ops.node.order + } + + fn distributed_params(&self) -> Option { + self.ops.node.distributed_params.clone() + } +} + +/// Make sure the grad tensor has the given shape. +/// +/// If broadcasting happened during the forward pass, the gradients will be sum along the +/// broadcasted dimension. +pub fn broadcast_shape(mut grad: FloatTensor, shape: &Shape) -> FloatTensor { + let shape_grad = grad.shape(); + let ndims = shape_grad.num_dims(); + + for i in 0..ndims { + if shape_grad[i] != shape[i] { + if shape[i] != 1 { + panic!( + "Invalid broadcast shapes: Next grad shape {:?}, Previous grad shape {:?}. {}", + shape, shape_grad, "Expected the shape of the next grad to be 1." + ); + } + grad = B::float_sum_dim(grad, i); + } + } + + grad +} diff --git a/crates/burn-autodiff/src/ops/bool_tensor.rs b/crates/burn-autodiff/src/ops/bool_tensor.rs new file mode 100644 index 0000000..2e3228b --- /dev/null +++ b/crates/burn-autodiff/src/ops/bool_tensor.rs @@ -0,0 +1,170 @@ +use crate::{Autodiff, checkpoint::strategy::CheckpointStrategy, tensor::AutodiffTensor}; +use alloc::vec::Vec; + +use burn_backend::{ + Backend, ExecutionError, Scalar, TensorData, + ops::BoolTensorOps, + tensor::{BoolTensor, Device, FloatTensor, IntTensor}, +}; +use burn_std::{BoolDType, FloatDType, IntDType, Shape}; + +impl BoolTensorOps for Autodiff { + fn bool_from_data(data: TensorData, device: &Device) -> BoolTensor { + B::bool_from_data(data, device) + } + + async fn bool_into_data(tensor: BoolTensor) -> Result { + B::bool_into_data(tensor).await + } + + fn bool_into_int(tensor: BoolTensor, out_dtype: IntDType) -> IntTensor { + B::bool_into_int(tensor, out_dtype) + } + + fn bool_to_device(tensor: BoolTensor, device: &Device) -> BoolTensor { + B::bool_to_device(tensor, device) + } + + fn bool_reshape(tensor: BoolTensor, shape: Shape) -> BoolTensor { + B::bool_reshape(tensor, shape) + } + + fn bool_slice(tensor: BoolTensor, slices: &[burn_std::Slice]) -> BoolTensor { + B::bool_slice(tensor, slices) + } + + fn bool_empty(shape: Shape, device: &Device, dtype: BoolDType) -> BoolTensor { + B::bool_empty(shape, device, dtype) + } + + fn bool_zeros(shape: Shape, device: &Device, dtype: BoolDType) -> BoolTensor { + B::bool_zeros(shape, device, dtype) + } + + fn bool_ones(shape: Shape, device: &Device, dtype: BoolDType) -> BoolTensor { + B::bool_ones(shape, device, dtype) + } + + fn bool_slice_assign( + tensor: BoolTensor, + slices: &[burn_std::Slice], + value: BoolTensor, + ) -> BoolTensor { + B::bool_slice_assign(tensor, slices, value) + } + + fn bool_cat(tensors: Vec>, dim: usize) -> BoolTensor { + B::bool_cat(tensors, dim) + } + + fn bool_equal(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + B::bool_equal(lhs, rhs) + } + + fn bool_not(tensor: BoolTensor) -> BoolTensor { + B::bool_not(tensor) + } + + fn bool_and(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + B::bool_and(lhs, rhs) + } + + fn bool_or(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + B::bool_or(lhs, rhs) + } + + fn bool_xor(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + B::bool_xor(lhs, rhs) + } + + fn bool_into_float(tensor: BoolTensor, out_dtype: FloatDType) -> FloatTensor { + AutodiffTensor::new(B::bool_into_float(tensor, out_dtype)) + } + + fn bool_swap_dims(tensor: BoolTensor, dim1: usize, dim2: usize) -> BoolTensor { + B::bool_swap_dims(tensor, dim1, dim2) + } + + fn bool_permute(tensor: BoolTensor, axes: &[usize]) -> BoolTensor { + B::bool_permute(tensor, axes) + } + + fn bool_flip(tensor: BoolTensor, axes: &[usize]) -> BoolTensor { + B::bool_flip(tensor, axes) + } + + async fn bool_argwhere(tensor: BoolTensor, out_dtype: burn_std::IntDType) -> IntTensor { + B::bool_argwhere(tensor, out_dtype).await + } + + fn bool_expand(tensor: BoolTensor, shape: Shape) -> BoolTensor { + B::bool_expand(tensor, shape) + } + + fn bool_repeat_dim(tensor: BoolTensor, dim: usize, times: usize) -> BoolTensor { + B::bool_repeat_dim(tensor, dim, times) + } + + fn bool_unfold( + tensor: BoolTensor, + dim: usize, + size: usize, + step: usize, + ) -> BoolTensor { + B::bool_unfold(tensor, dim, size, step) + } + + fn bool_mask_where( + tensor: BoolTensor, + mask: BoolTensor, + source: BoolTensor, + ) -> BoolTensor { + B::bool_mask_where(tensor, mask, source) + } + + fn bool_mask_fill( + tensor: BoolTensor, + mask: BoolTensor, + value: Scalar, + ) -> BoolTensor { + B::bool_mask_fill(tensor, mask, value) + } + + fn bool_gather( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + ) -> BoolTensor { + B::bool_gather(dim, tensor, indices) + } + + fn bool_scatter_or( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor { + B::bool_scatter_or(dim, tensor, indices, value) + } + + fn bool_equal_elem(lhs: BoolTensor, rhs: Scalar) -> BoolTensor { + B::bool_equal_elem(lhs, rhs) + } + + fn bool_select( + tensor: BoolTensor, + dim: usize, + indices: IntTensor, + ) -> BoolTensor { + B::bool_select(tensor, dim, indices) + } + + fn bool_select_or( + tensor: BoolTensor, + dim: usize, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor { + B::bool_select_or(tensor, dim, indices, value) + } +} diff --git a/crates/burn-autodiff/src/ops/distributed.rs b/crates/burn-autodiff/src/ops/distributed.rs new file mode 100644 index 0000000..785fd06 --- /dev/null +++ b/crates/burn-autodiff/src/ops/distributed.rs @@ -0,0 +1,90 @@ +use alloc::vec::Vec; +use burn_backend::{ + DeviceId, + distributed::{ + CollectiveTensor, DistributedConfig, DistributedOps, DistributedParams, ReduceOperation, + TensorRef, + }, + tensor::FloatTensor, +}; + +use burn_backend::Backend; + +use crate::{ + Autodiff, + checkpoint::strategy::CheckpointStrategy, + ops::{Backward, Ops, OpsKind, unary}, +}; + +impl DistributedOps for Autodiff { + fn start_communication_server(devices: &[B::Device], config: DistributedConfig) { + B::start_communication_server(devices, config); + } + + fn close_communication_server(device: &B::Device) { + B::close_communication_server(device); + } + + fn register_sync_parameters(device: &B::Device, distributed_params: Vec) { + B::register_sync_parameters(device, distributed_params); + } + + fn submit_sync_collective(device: &B::Device) { + B::submit_sync_collective(device); + } + + fn submit_gradient_sync(tensor: TensorRef, distributed_params: DistributedParams) { + let mut tensor = unsafe { (*tensor.0).clone() }; + B::submit_gradient_sync(TensorRef(&mut tensor.primitive), distributed_params); + } + + fn all_reduce( + tensor: FloatTensor, + op: ReduceOperation, + device_ids: Vec, + ) -> CollectiveTensor { + #[derive(Debug)] + struct AllReduce; + + impl Backward for AllReduce { + type State = (ReduceOperation, Vec); + + fn backward( + self, + ops: Ops, + grads: &mut crate::grads::Gradients, + _checkpointer: &mut crate::checkpoint::base::Checkpointer, + ) { + // Backward uses the same reduce op: local gradients are synchronized via the backend, which handles + // scaling (e.g., ncclAvg for mean). This works for the reduce ops that are currently supported, but we + // might need to rework it if we add other ops such as ncclMax. + unary::(ops.parents, ops.node, grads, |grad| { + B::all_reduce(grad, ops.state.0, ops.state.1).resolve() + }); + } + } + + let collective = B::all_reduce(tensor.primitive, op, device_ids.clone()); + // Safety: we call `assume_resolved` only to wrap it in a new `CollectiveTensor`. + let resolved = unsafe { collective.assume_resolved() }; + + match AllReduce + .prepare::([tensor.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(preps) => { + let output = preps.finish((op, device_ids), resolved); + CollectiveTensor::new(output) + } + OpsKind::UnTracked(preps) => { + let output = preps.finish(resolved); + CollectiveTensor::new(output) + } + } + } + + fn sync_collective(device: &B::Device) { + B::sync_collective(device); + } +} diff --git a/crates/burn-autodiff/src/ops/int_tensor.rs b/crates/burn-autodiff/src/ops/int_tensor.rs new file mode 100644 index 0000000..c4f94fa --- /dev/null +++ b/crates/burn-autodiff/src/ops/int_tensor.rs @@ -0,0 +1,416 @@ +use crate::{Autodiff, checkpoint::strategy::CheckpointStrategy, tensor::AutodiffTensor}; +use alloc::vec::Vec; + +use burn_backend::{ + Backend, Distribution, ExecutionError, Scalar, TensorData, + ops::IntTensorOps, + tensor::{BoolTensor, Device, FloatTensor, IntTensor}, +}; +use burn_std::{BoolDType, FloatDType, IndexingUpdateOp, IntDType, Shape}; + +impl IntTensorOps for Autodiff { + fn int_from_data(data: TensorData, device: &Device) -> IntTensor { + B::int_from_data(data, device) + } + + async fn int_into_data(tensor: IntTensor) -> Result { + B::int_into_data(tensor).await + } + + fn int_to_device(tensor: IntTensor, device: &Device) -> IntTensor { + B::int_to_device(tensor, device) + } + + fn int_reshape(tensor: IntTensor, shape: Shape) -> IntTensor { + B::int_reshape(tensor, shape) + } + + fn int_slice(tensor: IntTensor, slices: &[burn_std::Slice]) -> IntTensor { + B::int_slice(tensor, slices) + } + + fn int_empty(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + B::int_empty(shape, device, dtype) + } + + fn int_slice_assign( + tensor: IntTensor, + slices: &[burn_std::Slice], + value: IntTensor, + ) -> IntTensor { + B::int_slice_assign(tensor, slices, value) + } + + fn int_cat(tensors: Vec>, dim: usize) -> IntTensor { + B::int_cat(tensors, dim) + } + + fn int_equal(lhs: IntTensor, rhs: IntTensor, out_dtype: BoolDType) -> BoolTensor { + B::int_equal(lhs, rhs, out_dtype) + } + + fn int_equal_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor { + B::int_equal_elem(lhs, rhs, out_dtype) + } + + fn int_add(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + B::int_add(lhs, rhs) + } + + fn int_add_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + B::int_add_scalar(lhs, rhs) + } + + fn int_clamp_min(tensor: IntTensor, min: Scalar) -> IntTensor { + B::int_clamp_min(tensor, min) + } + + fn int_clamp_max(tensor: IntTensor, max: Scalar) -> IntTensor { + B::int_clamp_max(tensor, max) + } + + fn int_clamp(tensor: IntTensor, min: Scalar, max: Scalar) -> IntTensor { + B::int_clamp(tensor, min, max) + } + + fn int_sub(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + B::int_sub(lhs, rhs) + } + + fn int_sub_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + B::int_sub_scalar(lhs, rhs) + } + + fn int_mul(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + B::int_mul(lhs, rhs) + } + + fn int_mul_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + B::int_mul_scalar(lhs, rhs) + } + + fn int_div(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + B::int_div(lhs, rhs) + } + + fn int_div_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + B::int_div_scalar(lhs, rhs) + } + + fn int_remainder(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + B::int_remainder(lhs, rhs) + } + + fn int_remainder_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + B::int_remainder_scalar(lhs, rhs) + } + + fn int_matmul(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + B::int_matmul(lhs, rhs) + } + + fn int_neg(tensor: IntTensor) -> IntTensor { + B::int_neg(tensor) + } + + fn int_zeros(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + B::int_zeros(shape, device, dtype) + } + + fn int_ones(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + B::int_ones(shape, device, dtype) + } + + fn int_full( + shape: Shape, + fill_value: Scalar, + device: &Device, + dtype: IntDType, + ) -> IntTensor { + B::int_full(shape, fill_value, device, dtype) + } + + fn int_sum(tensor: IntTensor) -> IntTensor { + B::int_sum(tensor) + } + + fn int_sum_dim(tensor: IntTensor, dim: usize) -> IntTensor { + B::int_sum_dim(tensor, dim) + } + + fn int_mean(tensor: IntTensor) -> IntTensor { + B::int_mean(tensor) + } + + fn int_mean_dim(tensor: IntTensor, dim: usize) -> IntTensor { + B::int_mean_dim(tensor, dim) + } + + fn int_cumsum(tensor: IntTensor, dim: usize) -> IntTensor { + B::int_cumsum(tensor, dim) + } + + fn int_cumprod(tensor: IntTensor, dim: usize) -> IntTensor { + B::int_cumprod(tensor, dim) + } + + fn int_cummin(tensor: IntTensor, dim: usize) -> IntTensor { + B::int_cummin(tensor, dim) + } + + fn int_cummax(tensor: IntTensor, dim: usize) -> IntTensor { + B::int_cummax(tensor, dim) + } + + fn int_repeat_dim(tensor: IntTensor, dim: usize, times: usize) -> IntTensor { + B::int_repeat_dim(tensor, dim, times) + } + + fn int_greater(lhs: IntTensor, rhs: IntTensor, out_dtype: BoolDType) -> BoolTensor { + B::int_greater(lhs, rhs, out_dtype) + } + + fn int_greater_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor { + B::int_greater_elem(lhs, rhs, out_dtype) + } + + fn int_greater_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + B::int_greater_equal(lhs, rhs, out_dtype) + } + + fn int_greater_equal_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + B::int_greater_equal_elem(lhs, rhs, out_dtype) + } + + fn int_lower(lhs: IntTensor, rhs: IntTensor, out_dtype: BoolDType) -> BoolTensor { + B::int_lower(lhs, rhs, out_dtype) + } + + fn int_lower_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor { + B::int_lower_elem(lhs, rhs, out_dtype) + } + + fn int_lower_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + B::int_lower_equal(lhs, rhs, out_dtype) + } + + fn int_lower_equal_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor { + B::int_lower_equal_elem(lhs, rhs, out_dtype) + } + + fn int_gather(dim: usize, tensor: IntTensor, indices: IntTensor) -> IntTensor { + B::int_gather(dim, tensor, indices) + } + + fn int_scatter_add( + dim: usize, + tensor: IntTensor, + indices: IntTensor, + value: IntTensor, + ) -> IntTensor { + B::int_scatter_add(dim, tensor, indices, value) + } + + fn int_scatter_nd( + data: IntTensor, + indices: IntTensor, + values: IntTensor, + reduction: IndexingUpdateOp, + ) -> IntTensor { + B::int_scatter_nd(data, indices, values, reduction) + } + + fn int_select(tensor: IntTensor, dim: usize, indices: IntTensor) -> IntTensor { + B::int_select(tensor, dim, indices) + } + + fn int_select_add( + tensor: IntTensor, + dim: usize, + indices: IntTensor, + value: IntTensor, + ) -> IntTensor { + B::int_select_add(tensor, dim, indices, value) + } + + fn int_mask_where( + tensor: IntTensor, + mask: BoolTensor, + value: IntTensor, + ) -> IntTensor { + B::int_mask_where(tensor, mask, value) + } + + fn int_mask_fill(tensor: IntTensor, mask: BoolTensor, value: Scalar) -> IntTensor { + B::int_mask_fill(tensor, mask, value) + } + + fn int_argmax(tensor: IntTensor, dim: usize) -> IntTensor { + B::int_argmax(tensor, dim) + } + + fn int_argtopk(tensor: IntTensor, dim: usize, k: usize) -> IntTensor { + B::int_argtopk(tensor, dim, k) + } + + fn int_argmin(tensor: IntTensor, dim: usize) -> IntTensor { + B::int_argmin(tensor, dim) + } + fn int_max(tensor: IntTensor) -> IntTensor { + B::int_max(tensor) + } + fn int_max_dim(tensor: IntTensor, dim: usize) -> IntTensor { + B::int_max_dim(tensor, dim) + } + fn int_topk(tensor: IntTensor, dim: usize, k: usize) -> IntTensor { + B::int_topk(tensor, dim, k) + } + fn int_max_dim_with_indices(tensor: IntTensor, dim: usize) -> (IntTensor, IntTensor) { + B::int_max_dim_with_indices(tensor, dim) + } + fn int_min(tensor: IntTensor) -> IntTensor { + B::int_min(tensor) + } + fn int_min_dim(tensor: IntTensor, dim: usize) -> IntTensor { + B::int_min_dim(tensor, dim) + } + fn int_min_dim_with_indices(tensor: IntTensor, dim: usize) -> (IntTensor, IntTensor) { + B::int_min_dim_with_indices(tensor, dim) + } + fn int_abs(tensor: IntTensor) -> IntTensor { + B::int_abs(tensor) + } + fn int_into_float(tensor: IntTensor, out_dtype: FloatDType) -> FloatTensor { + AutodiffTensor::new(B::int_into_float(tensor, out_dtype)) + } + + fn int_swap_dims(tensor: IntTensor, dim1: usize, dim2: usize) -> IntTensor { + B::int_swap_dims(tensor, dim1, dim2) + } + + fn int_random( + shape: Shape, + distribution: Distribution, + device: &Device, + dtype: IntDType, + ) -> IntTensor { + B::int_random(shape, distribution, device, dtype) + } + + fn int_arange( + range: core::ops::Range, + device: &Device, + dtype: IntDType, + ) -> IntTensor { + B::int_arange(range, device, dtype) + } + + fn int_permute(tensor: IntTensor, axes: &[usize]) -> IntTensor { + B::int_permute(tensor, axes) + } + + fn int_flip(tensor: IntTensor, axes: &[usize]) -> IntTensor { + B::int_flip(tensor, axes) + } + + fn int_sign(tensor: IntTensor) -> IntTensor { + B::int_sign(tensor) + } + + fn int_prod(tensor: IntTensor) -> IntTensor { + B::int_prod(tensor) + } + + fn int_prod_dim(tensor: IntTensor, dim: usize) -> IntTensor { + B::int_prod_dim(tensor, dim) + } + + fn int_expand(tensor: IntTensor, shape: Shape) -> IntTensor { + B::int_expand(tensor, shape) + } + + fn int_sort(tensor: IntTensor, dim: usize, descending: bool) -> IntTensor { + B::int_sort(tensor, dim, descending) + } + + fn int_sort_with_indices( + tensor: IntTensor, + dim: usize, + descending: bool, + ) -> (IntTensor, IntTensor) { + B::int_sort_with_indices(tensor, dim, descending) + } + + fn int_argsort(tensor: IntTensor, dim: usize, descending: bool) -> IntTensor { + B::int_argsort(tensor, dim, descending) + } + + fn bitwise_and(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + B::bitwise_and(lhs, rhs) + } + + fn bitwise_and_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + B::bitwise_and_scalar(lhs, rhs) + } + + fn bitwise_or(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + B::bitwise_or(lhs, rhs) + } + + fn bitwise_or_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + B::bitwise_or_scalar(lhs, rhs) + } + + fn bitwise_xor(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + B::bitwise_xor(lhs, rhs) + } + + fn bitwise_xor_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + B::bitwise_xor_scalar(lhs, rhs) + } + + fn bitwise_not(tensor: IntTensor) -> IntTensor { + B::bitwise_not(tensor) + } + + fn bitwise_left_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + B::bitwise_left_shift(lhs, rhs) + } + + fn bitwise_left_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + B::bitwise_left_shift_scalar(lhs, rhs) + } + + fn bitwise_right_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + B::bitwise_right_shift(lhs, rhs) + } + + fn bitwise_right_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + B::bitwise_right_shift_scalar(lhs, rhs) + } + + fn int_cast(tensor: IntTensor, dtype: IntDType) -> IntTensor { + B::int_cast(tensor, dtype) + } + + fn int_unfold( + tensor: IntTensor, + dim: usize, + size: usize, + step: usize, + ) -> IntTensor { + B::int_unfold(tensor, dim, size, step) + } +} diff --git a/crates/burn-autodiff/src/ops/maxmin.rs b/crates/burn-autodiff/src/ops/maxmin.rs new file mode 100644 index 0000000..ecc25d8 --- /dev/null +++ b/crates/burn-autodiff/src/ops/maxmin.rs @@ -0,0 +1,27 @@ +use super::{Backward, Ops, unary}; +use crate::{checkpoint::base::Checkpointer, grads::Gradients}; +use burn_backend::{Backend, TensorMetadata}; +use burn_std::Shape; + +#[derive(Debug)] +pub(crate) struct MaxMinDim; + +impl Backward for MaxMinDim { + type State = (B::IntTensorPrimitive, Shape, usize); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + unary::(ops.parents, ops.node, grads, |grad| { + let (indices, shape, dim) = ops.state; + let device = grad.device(); + let dtype = grad.dtype(); + let zeros = B::float_zeros(shape, &device, dtype.into()); + + B::float_scatter_add(dim, zeros, indices, grad) + }); + } +} diff --git a/crates/burn-autodiff/src/ops/mod.rs b/crates/burn-autodiff/src/ops/mod.rs new file mode 100644 index 0000000..1a24067 --- /dev/null +++ b/crates/burn-autodiff/src/ops/mod.rs @@ -0,0 +1,16 @@ +mod activation; +mod backward; +mod base; +mod bool_tensor; +mod distributed; +mod int_tensor; +mod module; +mod qtensor; +mod tensor; +mod transaction; + +pub(crate) mod maxmin; +pub(crate) mod sort; + +pub use backward::*; +pub use base::*; diff --git a/crates/burn-autodiff/src/ops/module.rs b/crates/burn-autodiff/src/ops/module.rs new file mode 100644 index 0000000..ddc08d4 --- /dev/null +++ b/crates/burn-autodiff/src/ops/module.rs @@ -0,0 +1,2248 @@ +use alloc::{vec, vec::Vec}; + +use crate::Autodiff; +use crate::checkpoint::base::Checkpointer; +use crate::checkpoint::strategy::CheckpointStrategy; +use crate::grads::Gradients; +use crate::graph::NodeId; +use crate::ops::{Backward, Ops, unary}; +use crate::tensor::AutodiffTensor; + +use burn_backend::TensorMetadata; +use burn_backend::ops::attention::attention_fallback; +use burn_backend::ops::*; +use burn_backend::tensor::{FloatTensor, IntTensor}; +use burn_backend::{Backend, get_device_settings}; +use burn_std::{IntDType, Slice}; + +use super::OpsKind; + +impl ModuleOps> for Autodiff { + fn embedding(weights: AutodiffTensor, indices: IntTensor) -> AutodiffTensor { + #[derive(Debug)] + struct Embedding; + + impl Backward for Embedding { + type State = (B::FloatTensorPrimitive, IntTensor); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (weights, indices) = ops.state; + + unary::(ops.parents, ops.node, grads, |grad| { + B::embedding_backward(weights, grad, indices) + }); + } + } + + match Embedding + .prepare::([weights.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + (weights.primitive.clone(), indices.clone()), + B::embedding(weights.primitive, indices), + ), + OpsKind::UnTracked(prep) => prep.finish(B::embedding(weights.primitive, indices)), + } + } + + fn embedding_backward( + _weights: AutodiffTensor, + _output: AutodiffTensor, + _indices: IntTensor, + ) -> AutodiffTensor { + panic!("Can't differentiate embedding backward."); + } + + fn linear( + x: AutodiffTensor, + weight: AutodiffTensor, + bias: Option>, + ) -> AutodiffTensor { + #[derive(Debug)] + struct LinearWithBias; + #[derive(Debug)] + struct LinearNoBias; + + impl Backward for LinearWithBias { + type State = (Option, Option); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_weight, node_bias] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, weight_state) = ops.state; + let x = x_state + .map(|id| checkpointer.retrieve_node_output::(id)); + let weight = weight_state + .map(|id| checkpointer.retrieve_node_output::(id)); + + if let Some(node) = node_x { + let grad = B::linear_x_backward(weight.unwrap(), grad.clone()); + grads.register::(node.id, grad) + } + if let Some(node) = node_weight { + let grad = B::linear_weight_backward(x.unwrap(), grad.clone()); + grads.register::(node.id, grad) + } + if let Some(node) = node_bias { + let grad = B::linear_bias_backward(grad); + grads.register::(node.id, grad) + } + } + } + + impl Backward for LinearNoBias { + type State = (Option, Option); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_weight] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, weight_state) = ops.state; + let x = x_state + .map(|id| checkpointer.retrieve_node_output::(id)); + let weight = weight_state + .map(|id| checkpointer.retrieve_node_output::(id)); + + if let Some(node) = node_x { + let grad = B::linear_x_backward(weight.unwrap(), grad.clone()); + grads.register::(node.id, grad) + } + if let Some(node) = node_weight { + let grad = B::linear_weight_backward(x.unwrap(), grad); + grads.register::(node.id, grad) + } + } + } + + let x_tracked = x.is_tracked(); + let weight_tracked = weight.is_tracked(); + + match bias { + Some(bias) => match LinearWithBias + .prepare::([x.node.clone(), weight.node.clone(), bias.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + // x is only needed to compute the weight gradient, and vice versa. + let x_state = weight_tracked.then(|| prep.checkpoint(&x)); + let weight_state = x_tracked.then(|| prep.checkpoint(&weight)); + prep.finish( + (x_state, weight_state), + B::linear(x.primitive, weight.primitive, Some(bias.primitive)), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::linear( + x.primitive, + weight.primitive, + Some(bias.primitive), + )), + }, + None => match LinearNoBias + .prepare::([x.node.clone(), weight.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = weight_tracked.then(|| prep.checkpoint(&x)); + let weight_state = x_tracked.then(|| prep.checkpoint(&weight)); + prep.finish( + (x_state, weight_state), + B::linear(x.primitive, weight.primitive, None), + ) + } + OpsKind::UnTracked(prep) => { + prep.finish(B::linear(x.primitive, weight.primitive, None)) + } + }, + } + } + + fn linear_x_backward( + _weight: AutodiffTensor, + _output_grad: AutodiffTensor, + ) -> AutodiffTensor { + panic!("Can't differentiate linear_x_backward."); + } + + fn linear_weight_backward( + _x: AutodiffTensor, + _output_grad: AutodiffTensor, + ) -> AutodiffTensor { + panic!("Can't differentiate linear_weight_backward."); + } + + fn linear_bias_backward(_output_grad: AutodiffTensor) -> AutodiffTensor { + panic!("Can't differentiate linear_bias_backward."); + } + + fn conv1d( + x: AutodiffTensor, + weight: AutodiffTensor, + bias: Option>, + options: ConvOptions<1>, + ) -> AutodiffTensor { + #[derive(Debug)] + struct Conv1DWithBias; + #[derive(Debug)] + struct Conv1DNoBias; + + impl Backward for Conv1DWithBias { + type State = (NodeId, NodeId, NodeId, ConvOptions<1>); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_weight, node_bias] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, weight_state, bias_state, options) = ops.state; + let x = checkpointer.retrieve_node_output::(x_state); + let weight = + checkpointer.retrieve_node_output::(weight_state); + let bias = checkpointer.retrieve_node_output::(bias_state); + + if let Some(node) = node_x { + let grad = B::conv1d_x_backward( + x.clone(), + weight.clone(), + grad.clone(), + options.clone(), + ); + grads.register::(node.id, grad) + } + if let Some(node) = node_weight { + let grad = B::conv1d_weight_backward(x.clone(), weight, grad.clone(), options); + grads.register::(node.id, grad) + } + if let Some(node) = node_bias { + let grad = B::conv1d_bias_backward(x, bias, grad); + grads.register::(node.id, grad) + } + } + } + + impl Backward for Conv1DNoBias { + type State = (NodeId, NodeId, ConvOptions<1>); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_weight] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, weight_state, options) = ops.state; + let x = checkpointer.retrieve_node_output::(x_state); + let weight = + checkpointer.retrieve_node_output::(weight_state); + + if let Some(node) = node_x { + let grad = B::conv1d_x_backward( + x.clone(), + weight.clone(), + grad.clone(), + options.clone(), + ); + grads.register::(node.id, grad) + } + if let Some(node) = node_weight { + let grad = B::conv1d_weight_backward(x, weight, grad, options); + grads.register::(node.id, grad) + } + } + } + match bias { + Some(bias) => match Conv1DWithBias + .prepare::([x.node.clone(), weight.node.clone(), bias.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let weight_state = prep.checkpoint(&weight); + let bias_state = prep.checkpoint(&bias); + prep.finish( + (x_state, weight_state, bias_state, options.clone()), + B::conv1d(x.primitive, weight.primitive, Some(bias.primitive), options), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::conv1d( + x.primitive, + weight.primitive, + Some(bias.primitive), + options, + )), + }, + None => match Conv1DNoBias + .prepare::([x.node.clone(), weight.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let weight_state = prep.checkpoint(&weight); + prep.finish( + (x_state, weight_state, options.clone()), + B::conv1d(x.primitive, weight.primitive, None, options), + ) + } + OpsKind::UnTracked(prep) => { + prep.finish(B::conv1d(x.primitive, weight.primitive, None, options)) + } + }, + } + } + + fn conv_transpose1d( + x: AutodiffTensor, + weight: AutodiffTensor, + bias: Option>, + options: ConvTransposeOptions<1>, + ) -> AutodiffTensor { + #[derive(Debug)] + struct ConvTranspose1DWithBias; + #[derive(Debug)] + struct ConvTranspose1DNoBias; + + impl Backward for ConvTranspose1DWithBias { + type State = (NodeId, NodeId, NodeId, ConvTransposeOptions<1>); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_weight, node_bias] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, weight_state, bias_state, options) = ops.state; + let x = checkpointer.retrieve_node_output::(x_state); + let weight = + checkpointer.retrieve_node_output::(weight_state); + let bias = checkpointer.retrieve_node_output::(bias_state); + + if let Some(node) = node_x { + let grad = B::conv_transpose1d_x_backward( + weight.clone(), + grad.clone(), + options.clone(), + ); + grads.register::(node.id, grad) + } + if let Some(node) = node_weight { + let grad = B::conv_transpose1d_weight_backward( + x.clone(), + weight, + grad.clone(), + options, + ); + grads.register::(node.id, grad) + } + if let Some(node) = node_bias { + let grad = B::conv_transpose1d_bias_backward(x, bias, grad); + grads.register::(node.id, grad) + } + } + } + + impl Backward for ConvTranspose1DNoBias { + type State = (NodeId, NodeId, ConvTransposeOptions<1>); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_weight] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, weight_state, options) = ops.state; + let x = checkpointer.retrieve_node_output::(x_state); + let weight = + checkpointer.retrieve_node_output::(weight_state); + + if let Some(node) = node_x { + let grad = B::conv_transpose1d_x_backward( + weight.clone(), + grad.clone(), + options.clone(), + ); + grads.register::(node.id, grad) + } + if let Some(node) = node_weight { + let grad = B::conv_transpose1d_weight_backward(x, weight, grad, options); + grads.register::(node.id, grad) + } + } + } + + match bias { + Some(bias) => match ConvTranspose1DWithBias + .prepare::([x.node.clone(), weight.node.clone(), bias.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let weight_state = prep.checkpoint(&weight); + let bias_state = prep.checkpoint(&bias); + prep.finish( + (x_state, weight_state, bias_state, options.clone()), + B::conv_transpose1d( + x.primitive, + weight.primitive, + Some(bias.primitive), + options, + ), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::conv_transpose1d( + x.primitive, + weight.primitive, + Some(bias.primitive), + options, + )), + }, + None => match ConvTranspose1DNoBias + .prepare::([x.node.clone(), weight.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let weight_state = prep.checkpoint(&weight); + prep.finish( + (x_state, weight_state, options.clone()), + B::conv_transpose1d(x.primitive, weight.primitive, None, options), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::conv_transpose1d( + x.primitive, + weight.primitive, + None, + options, + )), + }, + } + } + + fn conv2d( + x: AutodiffTensor, + weight: AutodiffTensor, + bias: Option>, + options: ConvOptions<2>, + ) -> AutodiffTensor { + #[derive(Debug)] + struct Conv2DWithBias; + #[derive(Debug)] + struct Conv2DNoBias; + + impl Backward for Conv2DWithBias { + type State = (NodeId, NodeId, NodeId, ConvOptions<2>); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_weight, node_bias] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, weight_state, bias_state, options) = ops.state; + let x = checkpointer.retrieve_node_output::(x_state); + let weight = + checkpointer.retrieve_node_output::(weight_state); + let bias = checkpointer.retrieve_node_output::(bias_state); + + if let Some(node) = node_x { + let grad = B::conv2d_x_backward( + x.clone(), + weight.clone(), + grad.clone(), + options.clone(), + ); + grads.register::(node.id, grad) + } + if let Some(node) = node_weight { + let grad = + B::conv2d_weight_backward(x.clone(), weight.clone(), grad.clone(), options); + grads.register::(node.id, grad) + } + if let Some(node) = node_bias { + let grad = B::conv2d_bias_backward(x, bias, grad); + grads.register::(node.id, grad) + } + } + } + + impl Backward for Conv2DNoBias { + type State = (NodeId, NodeId, ConvOptions<2>); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_weight] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, weight_state, options) = ops.state; + let x = checkpointer.retrieve_node_output::(x_state); + let weight = + checkpointer.retrieve_node_output::(weight_state); + + if let Some(node) = node_x { + let grad = B::conv2d_x_backward( + x.clone(), + weight.clone(), + grad.clone(), + options.clone(), + ); + grads.register::(node.id, grad) + } + if let Some(node) = node_weight { + let grad = B::conv2d_weight_backward(x, weight, grad, options); + grads.register::(node.id, grad) + } + } + } + + match bias { + Some(bias) => match Conv2DWithBias + .prepare::([x.node.clone(), weight.node.clone(), bias.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let weight_state = prep.checkpoint(&weight); + let bias_state = prep.checkpoint(&bias); + prep.finish( + (x_state, weight_state, bias_state, options.clone()), + B::conv2d(x.primitive, weight.primitive, Some(bias.primitive), options), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::conv2d( + x.primitive, + weight.primitive, + Some(bias.primitive), + options, + )), + }, + None => match Conv2DNoBias + .prepare::([x.node.clone(), weight.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let weight_state = prep.checkpoint(&weight); + prep.finish( + (x_state, weight_state, options.clone()), + B::conv2d(x.primitive, weight.primitive, None, options), + ) + } + + OpsKind::UnTracked(prep) => { + prep.finish(B::conv2d(x.primitive, weight.primitive, None, options)) + } + }, + } + } + + fn deform_conv2d( + x: AutodiffTensor, + offset: AutodiffTensor, + weight: AutodiffTensor, + mask: Option>, + bias: Option>, + options: DeformConvOptions<2>, + ) -> AutodiffTensor { + #[derive(Debug)] + struct DeformConv2DWithMaskWithBias; + #[derive(Debug)] + struct DeformConv2DWithMaskNoBias; + #[derive(Debug)] + struct DeformConv2DNoMaskWithBias; + #[derive(Debug)] + struct DeformConv2DNoMaskNoBias; + + impl Backward for DeformConv2DWithMaskWithBias { + type State = (NodeId, NodeId, NodeId, NodeId, NodeId, DeformConvOptions<2>); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_offset, node_weight, node_mask, node_bias] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, offset_state, weight_state, mask_state, bias_state, options) = + ops.state; + let x = checkpointer.retrieve_node_output(x_state); + let offset = checkpointer.retrieve_node_output(offset_state); + let weight = checkpointer.retrieve_node_output(weight_state); + let mask = Some(checkpointer.retrieve_node_output(mask_state)); + let bias = Some(checkpointer.retrieve_node_output(bias_state)); + + let backward = + B::deform_conv2d_backward(x, offset, weight, mask, bias, grad, options); + + if let Some(node) = node_x { + grads.register::(node.id, backward.x_grad) + } + if let Some(node) = node_offset { + grads.register::(node.id, backward.offset_grad) + } + if let Some(node) = node_weight { + grads.register::(node.id, backward.weight_grad) + } + if let Some(node) = node_mask { + grads.register::(node.id, backward.mask_grad.unwrap()) + } + if let Some(node) = node_bias { + grads.register::(node.id, backward.bias_grad.unwrap()) + } + } + } + + impl Backward for DeformConv2DWithMaskNoBias { + type State = (NodeId, NodeId, NodeId, NodeId, DeformConvOptions<2>); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_offset, node_weight, node_mask] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, offset_state, weight_state, mask_state, options) = ops.state; + let x = checkpointer.retrieve_node_output(x_state); + let offset = checkpointer.retrieve_node_output(offset_state); + let weight = checkpointer.retrieve_node_output(weight_state); + let mask = Some(checkpointer.retrieve_node_output(mask_state)); + + let backward = + B::deform_conv2d_backward(x, offset, weight, mask, None, grad, options); + + if let Some(node) = node_x { + grads.register::(node.id, backward.x_grad) + } + if let Some(node) = node_offset { + grads.register::(node.id, backward.offset_grad) + } + if let Some(node) = node_weight { + grads.register::(node.id, backward.weight_grad) + } + if let Some(node) = node_mask { + grads.register::(node.id, backward.mask_grad.unwrap()) + } + } + } + + impl Backward for DeformConv2DNoMaskWithBias { + type State = (NodeId, NodeId, NodeId, NodeId, DeformConvOptions<2>); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_offset, node_weight, node_bias] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, offset_state, weight_state, bias_state, options) = ops.state; + let x = checkpointer.retrieve_node_output(x_state); + let offset = checkpointer.retrieve_node_output(offset_state); + let weight = checkpointer.retrieve_node_output(weight_state); + let bias = Some(checkpointer.retrieve_node_output(bias_state)); + + let backward = + B::deform_conv2d_backward(x, offset, weight, None, bias, grad, options); + + if let Some(node) = node_x { + grads.register::(node.id, backward.x_grad) + } + if let Some(node) = node_offset { + grads.register::(node.id, backward.offset_grad) + } + if let Some(node) = node_weight { + grads.register::(node.id, backward.weight_grad) + } + if let Some(node) = node_bias { + grads.register::(node.id, backward.bias_grad.unwrap()) + } + } + } + + impl Backward for DeformConv2DNoMaskNoBias { + type State = (NodeId, NodeId, NodeId, DeformConvOptions<2>); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_offset, node_weight] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, offset_state, weight_state, options) = ops.state; + let x = checkpointer.retrieve_node_output(x_state); + let offset = checkpointer.retrieve_node_output(offset_state); + let weight = checkpointer.retrieve_node_output(weight_state); + + let backward = + B::deform_conv2d_backward(x, offset, weight, None, None, grad, options); + + if let Some(node) = node_x { + grads.register::(node.id, backward.x_grad) + } + if let Some(node) = node_offset { + grads.register::(node.id, backward.offset_grad) + } + if let Some(node) = node_weight { + grads.register::(node.id, backward.weight_grad) + } + } + } + + match (mask, bias) { + (Some(mask), Some(bias)) => match DeformConv2DWithMaskWithBias + .prepare::([ + x.node.clone(), + offset.node.clone(), + weight.node.clone(), + mask.node.clone(), + bias.node.clone(), + ]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let offset_state = prep.checkpoint(&offset); + let weight_state = prep.checkpoint(&weight); + let mask_state = prep.checkpoint(&mask); + let bias_state = prep.checkpoint(&bias); + prep.finish( + ( + x_state, + offset_state, + weight_state, + mask_state, + bias_state, + options.clone(), + ), + B::deform_conv2d( + x.primitive, + offset.primitive, + weight.primitive, + Some(mask.primitive), + Some(bias.primitive), + options, + ), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::deform_conv2d( + x.primitive, + offset.primitive, + weight.primitive, + Some(mask.primitive), + Some(bias.primitive), + options, + )), + }, + (Some(mask), None) => match DeformConv2DWithMaskNoBias + .prepare::([ + x.node.clone(), + offset.node.clone(), + weight.node.clone(), + mask.node.clone(), + ]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let offset_state = prep.checkpoint(&offset); + let weight_state = prep.checkpoint(&weight); + let mask_state = prep.checkpoint(&mask); + prep.finish( + ( + x_state, + offset_state, + weight_state, + mask_state, + options.clone(), + ), + B::deform_conv2d( + x.primitive, + offset.primitive, + weight.primitive, + Some(mask.primitive), + None, + options, + ), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::deform_conv2d( + x.primitive, + offset.primitive, + weight.primitive, + Some(mask.primitive), + None, + options, + )), + }, + (None, Some(bias)) => match DeformConv2DNoMaskWithBias + .prepare::([ + x.node.clone(), + offset.node.clone(), + weight.node.clone(), + bias.node.clone(), + ]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let offset_state = prep.checkpoint(&offset); + let weight_state = prep.checkpoint(&weight); + let bias_state = prep.checkpoint(&bias); + prep.finish( + ( + x_state, + offset_state, + weight_state, + bias_state, + options.clone(), + ), + B::deform_conv2d( + x.primitive, + offset.primitive, + weight.primitive, + None, + Some(bias.primitive), + options, + ), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::deform_conv2d( + x.primitive, + offset.primitive, + weight.primitive, + None, + Some(bias.primitive), + options, + )), + }, + (None, None) => match DeformConv2DNoMaskNoBias + .prepare::([x.node.clone(), offset.node.clone(), weight.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let offset_state = prep.checkpoint(&offset); + let weight_state = prep.checkpoint(&weight); + prep.finish( + (x_state, offset_state, weight_state, options.clone()), + B::deform_conv2d( + x.primitive, + offset.primitive, + weight.primitive, + None, + None, + options, + ), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::deform_conv2d( + x.primitive, + offset.primitive, + weight.primitive, + None, + None, + options, + )), + }, + } + } + + fn deform_conv2d_backward( + _x: AutodiffTensor, + _offset: AutodiffTensor, + _weight: AutodiffTensor, + _mask: Option>, + _bias: Option>, + _output_grad: AutodiffTensor, + _options: DeformConvOptions<2>, + ) -> DeformConv2dBackward { + panic!("Can't differentiate deform conv 2d backward."); + } + + fn conv_transpose2d( + x: AutodiffTensor, + weight: AutodiffTensor, + bias: Option>, + options: ConvTransposeOptions<2>, + ) -> AutodiffTensor { + #[derive(Debug)] + struct ConvTranspose2DWithBias; + #[derive(Debug)] + struct ConvTranspose2DNoBias; + + impl Backward for ConvTranspose2DWithBias { + type State = (NodeId, NodeId, NodeId, ConvTransposeOptions<2>); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_weight, node_bias] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, weight_state, bias_state, options) = ops.state; + let x = checkpointer.retrieve_node_output::(x_state); + let weight = + checkpointer.retrieve_node_output::(weight_state); + let bias = checkpointer.retrieve_node_output::(bias_state); + + if let Some(node) = node_x { + let grad = B::conv_transpose2d_x_backward( + weight.clone(), + grad.clone(), + options.clone(), + ); + grads.register::(node.id, grad) + } + if let Some(node) = node_weight { + let grad = B::conv_transpose2d_weight_backward( + x.clone(), + weight, + grad.clone(), + options, + ); + grads.register::(node.id, grad) + } + if let Some(node) = node_bias { + let grad = B::conv_transpose2d_bias_backward(x, bias, grad); + grads.register::(node.id, grad) + } + } + } + + impl Backward for ConvTranspose2DNoBias { + type State = (NodeId, NodeId, ConvTransposeOptions<2>); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_weight] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, weight_state, options) = ops.state; + let x = checkpointer.retrieve_node_output::(x_state); + let weight = + checkpointer.retrieve_node_output::(weight_state); + + if let Some(node) = node_x { + let grad = B::conv_transpose2d_x_backward( + weight.clone(), + grad.clone(), + options.clone(), + ); + grads.register::(node.id, grad) + } + if let Some(node) = node_weight { + let grad = B::conv_transpose2d_weight_backward(x, weight, grad, options); + grads.register::(node.id, grad) + } + } + } + + match bias { + Some(bias) => match ConvTranspose2DWithBias + .prepare::([x.node.clone(), weight.node.clone(), bias.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let weight_state = prep.checkpoint(&weight); + let bias_state = prep.checkpoint(&bias); + + prep.finish( + (x_state, weight_state, bias_state, options.clone()), + B::conv_transpose2d( + x.primitive, + weight.primitive, + Some(bias.primitive), + options, + ), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::conv_transpose2d( + x.primitive, + weight.primitive, + Some(bias.primitive), + options, + )), + }, + None => match ConvTranspose2DNoBias + .prepare::([x.node.clone(), weight.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let weight_state = prep.checkpoint(&weight); + + prep.finish( + (x_state, weight_state, options.clone()), + B::conv_transpose2d(x.primitive, weight.primitive, None, options), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::conv_transpose2d( + x.primitive, + weight.primitive, + None, + options, + )), + }, + } + } + + fn conv3d( + x: AutodiffTensor, + weight: AutodiffTensor, + bias: Option>, + options: ConvOptions<3>, + ) -> AutodiffTensor { + #[derive(Debug)] + struct Conv3DWithBias; + #[derive(Debug)] + struct Conv3DNoBias; + + impl Backward for Conv3DWithBias { + type State = (NodeId, NodeId, NodeId, ConvOptions<3>); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_weight, node_bias] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, weight_state, bias_state, options) = ops.state; + let x = checkpointer.retrieve_node_output::(x_state); + let weight = + checkpointer.retrieve_node_output::(weight_state); + let bias = checkpointer.retrieve_node_output::(bias_state); + + if let Some(node) = node_x { + let grad = B::conv3d_x_backward( + x.clone(), + weight.clone(), + grad.clone(), + options.clone(), + ); + grads.register::(node.id, grad) + } + if let Some(node) = node_weight { + let grad = + B::conv3d_weight_backward(x.clone(), weight.clone(), grad.clone(), options); + grads.register::(node.id, grad) + } + if let Some(node) = node_bias { + let grad = B::conv3d_bias_backward(x, bias, grad); + grads.register::(node.id, grad) + } + } + } + + impl Backward for Conv3DNoBias { + type State = (NodeId, NodeId, ConvOptions<3>); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_weight] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, weight_state, options) = ops.state; + let x = checkpointer.retrieve_node_output::(x_state); + let weight = + checkpointer.retrieve_node_output::(weight_state); + + if let Some(node) = node_x { + let grad = B::conv3d_x_backward( + x.clone(), + weight.clone(), + grad.clone(), + options.clone(), + ); + grads.register::(node.id, grad) + } + if let Some(node) = node_weight { + let grad = B::conv3d_weight_backward(x, weight, grad, options); + grads.register::(node.id, grad) + } + } + } + + match bias { + Some(bias) => match Conv3DWithBias + .prepare::([x.node.clone(), weight.node.clone(), bias.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let weight_state = prep.checkpoint(&weight); + let bias_state = prep.checkpoint(&bias); + prep.finish( + (x_state, weight_state, bias_state, options.clone()), + B::conv3d(x.primitive, weight.primitive, Some(bias.primitive), options), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::conv3d( + x.primitive, + weight.primitive, + Some(bias.primitive), + options, + )), + }, + None => match Conv3DNoBias + .prepare::([x.node.clone(), weight.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let weight_state = prep.checkpoint(&weight); + prep.finish( + (x_state, weight_state, options.clone()), + B::conv3d(x.primitive, weight.primitive, None, options), + ) + } + + OpsKind::UnTracked(prep) => { + prep.finish(B::conv3d(x.primitive, weight.primitive, None, options)) + } + }, + } + } + + fn conv_transpose3d( + x: AutodiffTensor, + weight: AutodiffTensor, + bias: Option>, + options: ConvTransposeOptions<3>, + ) -> AutodiffTensor { + #[derive(Debug)] + struct ConvTranspose3DWithBias; + #[derive(Debug)] + struct ConvTranspose3DNoBias; + + impl Backward for ConvTranspose3DWithBias { + type State = (NodeId, NodeId, NodeId, ConvTransposeOptions<3>); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_weight, node_bias] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, weight_state, bias_state, options) = ops.state; + let x = checkpointer.retrieve_node_output::(x_state); + let weight = + checkpointer.retrieve_node_output::(weight_state); + let bias = checkpointer.retrieve_node_output::(bias_state); + + if let Some(node) = node_x { + let grad = B::conv_transpose3d_x_backward( + weight.clone(), + grad.clone(), + options.clone(), + ); + grads.register::(node.id, grad) + } + if let Some(node) = node_weight { + let grad = B::conv_transpose3d_weight_backward( + x.clone(), + weight, + grad.clone(), + options, + ); + grads.register::(node.id, grad) + } + if let Some(node) = node_bias { + let grad = B::conv_transpose3d_bias_backward(x, bias, grad); + grads.register::(node.id, grad) + } + } + } + + impl Backward for ConvTranspose3DNoBias { + type State = (NodeId, NodeId, ConvTransposeOptions<3>); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_x, node_weight] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, weight_state, options) = ops.state; + let x = checkpointer.retrieve_node_output::(x_state); + let weight = + checkpointer.retrieve_node_output::(weight_state); + + if let Some(node) = node_x { + let grad = B::conv_transpose3d_x_backward( + weight.clone(), + grad.clone(), + options.clone(), + ); + grads.register::(node.id, grad) + } + if let Some(node) = node_weight { + let grad = B::conv_transpose3d_weight_backward(x, weight, grad, options); + grads.register::(node.id, grad) + } + } + } + + match bias { + Some(bias) => match ConvTranspose3DWithBias + .prepare::([x.node.clone(), weight.node.clone(), bias.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let weight_state = prep.checkpoint(&weight); + let bias_state = prep.checkpoint(&bias); + + prep.finish( + (x_state, weight_state, bias_state, options.clone()), + B::conv_transpose3d( + x.primitive, + weight.primitive, + Some(bias.primitive), + options, + ), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::conv_transpose3d( + x.primitive, + weight.primitive, + Some(bias.primitive), + options, + )), + }, + None => match ConvTranspose3DNoBias + .prepare::([x.node.clone(), weight.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let weight_state = prep.checkpoint(&weight); + + prep.finish( + (x_state, weight_state, options.clone()), + B::conv_transpose3d(x.primitive, weight.primitive, None, options), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::conv_transpose3d( + x.primitive, + weight.primitive, + None, + options, + )), + }, + } + } + + // TODO: Support a custom unfold4d operation by overriding the default implementation. + // + // We don't override it now because the fold operation isn't available for the backward pass. + // This implies that when autodiff is enabled, custom unfold operations defined by backends + // won't be used. Instead, the conv2d operation with custom weights matrix will be used. + // Therefore, the conv2d backward pass will be used for the unfold4d backward pass. + // + // fn unfold4d( + // x:AutodiffTensor, + // kernel_size: [usize; 2], + // options: UnfoldOptions, + // ) -> AutodiffTensor { + // todo!() + // } + + fn avg_pool1d( + x: AutodiffTensor, + kernel_size: usize, + stride: usize, + padding: usize, + count_include_pad: bool, + ceil_mode: bool, + ) -> AutodiffTensor { + #[derive(Debug)] + struct AvgPool1D; + + impl Backward for AvgPool1D { + type State = (NodeId, usize, usize, usize, bool, bool); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_parent] = ops.parents; + let grad = grads.consume::(&ops.node); + let (x_state, kernel_size, stride, padding, count_include_pad, ceil_mode) = + ops.state; + let x = checkpointer.retrieve_node_output(x_state); + + if let Some(node) = node_parent { + let grad = B::avg_pool1d_backward( + x, + grad, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ); + grads.register::(node.id, grad); + } + } + } + + match AvgPool1D + .prepare::([x.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + prep.finish( + ( + x_state, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ), + B::avg_pool1d( + x.primitive.clone(), + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::avg_pool1d( + x.primitive, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + )), + } + } + + fn avg_pool2d( + x: AutodiffTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> AutodiffTensor { + #[derive(Debug)] + struct AvgPool2D; + + impl Backward for AvgPool2D { + type State = (NodeId, [usize; 2], [usize; 2], [usize; 2], bool, bool); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_parent] = ops.parents; + let grad = grads.consume::(&ops.node); + let (x_state, kernel_size, stride, padding, count_include_pad, ceil_mode) = + ops.state; + let x = checkpointer.retrieve_node_output(x_state); + + if let Some(node) = node_parent { + let grad = B::avg_pool2d_backward( + x, + grad, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ); + grads.register::(node.id, grad); + } + } + } + + match AvgPool2D + .prepare::([x.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + prep.finish( + ( + x_state, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ), + B::avg_pool2d( + x.primitive.clone(), + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::avg_pool2d( + x.primitive, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + )), + } + } + + fn avg_pool2d_backward( + _x: AutodiffTensor, + _grad: AutodiffTensor, + _kernel_size: [usize; 2], + _stride: [usize; 2], + _padding: [usize; 2], + _count_include_pad: bool, + _ceil_mode: bool, + ) -> AutodiffTensor { + panic!("Can't differentiate avg pool 2d backward."); + } + + fn max_pool1d( + x: AutodiffTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + ) -> AutodiffTensor { + match MaxPool1D + .prepare::([x.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let settings = get_device_settings::(&x.primitive.device()); + let output = B::max_pool1d_with_indices( + x.primitive, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + settings.int_dtype, + ); + prep.finish( + ( + x_state, + output.indices, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + ), + output.output, + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::max_pool1d( + x.primitive, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + )), + } + } + + fn max_pool1d_with_indices( + x: AutodiffTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + int_dtype: IntDType, + ) -> MaxPool1dWithIndices { + match MaxPool1D + .prepare::([x.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let output = B::max_pool1d_with_indices( + x.primitive, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + int_dtype, + ); + + let output_tensor = prep.finish( + ( + x_state, + output.indices.clone(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + ), + output.output, + ); + + MaxPool1dWithIndices::new(output_tensor, output.indices) + } + OpsKind::UnTracked(prep) => { + let output = B::max_pool1d_with_indices( + x.primitive, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + int_dtype, + ); + let output_tensor = prep.finish(output.output); + + MaxPool1dWithIndices::new(output_tensor, output.indices) + } + } + } + + fn max_pool1d_with_indices_backward( + x: AutodiffTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + output_grad: AutodiffTensor, + indices: IntTensor, + ) -> MaxPool1dBackward { + let output = B::max_pool1d_with_indices_backward( + x.primitive, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + output_grad.primitive, + indices, + ); + MaxPool1dBackward::new(AutodiffTensor::new(output.x_grad)) + } + + fn max_pool2d( + x: AutodiffTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + ) -> AutodiffTensor { + match MaxPool2D + .prepare::([x.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let settings = get_device_settings::(&x.primitive.device()); + let output = B::max_pool2d_with_indices( + x.primitive, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + settings.int_dtype, + ); + prep.finish( + ( + x_state, + output.indices, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + ), + output.output, + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::max_pool2d( + x.primitive, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + )), + } + } + + fn max_pool2d_with_indices( + x: AutodiffTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + int_dtype: IntDType, + ) -> MaxPool2dWithIndices { + match MaxPool2D + .prepare::([x.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + + let output = B::max_pool2d_with_indices( + x.primitive, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + int_dtype, + ); + + let output_tensor = prep.finish( + ( + x_state, + output.indices.clone(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + ), + output.output, + ); + + MaxPool2dWithIndices::new(output_tensor, output.indices) + } + OpsKind::UnTracked(prep) => { + let output = B::max_pool2d_with_indices( + x.primitive, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + int_dtype, + ); + let output_tensor = prep.finish(output.output); + + MaxPool2dWithIndices::new(output_tensor, output.indices) + } + } + } + + fn max_pool2d_with_indices_backward( + _x: AutodiffTensor, + _kernel_size: [usize; 2], + _stride: [usize; 2], + _padding: [usize; 2], + _dilation: [usize; 2], + _ceil_mode: bool, + _output_grad: AutodiffTensor, + _indices: IntTensor, + ) -> MaxPool2dBackward { + panic!("Can't differentiate max pool2d with indices backward."); + } + fn adaptive_avg_pool1d(x: AutodiffTensor, output_size: usize) -> AutodiffTensor { + #[derive(Debug)] + struct AdaptiveAvgPool1D; + + impl Backward for AdaptiveAvgPool1D { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_parent] = ops.parents; + let grad = grads.consume::(&ops.node); + let state = checkpointer.retrieve_node_output(ops.state); + + if let Some(node) = node_parent { + let grad = B::adaptive_avg_pool1d_backward(state, grad); + grads.register::(node.id, grad); + } + } + } + + match AdaptiveAvgPool1D + .prepare::([x.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + prep.finish(x_state, B::adaptive_avg_pool1d(x.primitive, output_size)) + } + OpsKind::UnTracked(prep) => { + prep.finish(B::adaptive_avg_pool1d(x.primitive, output_size)) + } + } + } + + fn adaptive_avg_pool2d(x: AutodiffTensor, output_size: [usize; 2]) -> AutodiffTensor { + #[derive(Debug)] + struct AdaptiveAvgPool2D; + + impl Backward for AdaptiveAvgPool2D { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_parent] = ops.parents; + let grad = grads.consume::(&ops.node); + let state = checkpointer.retrieve_node_output(ops.state); + + if let Some(node) = node_parent { + let grad = B::adaptive_avg_pool2d_backward(state, grad); + grads.register::(node.id, grad); + } + } + } + + match AdaptiveAvgPool2D + .prepare::([x.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + prep.finish(x_state, B::adaptive_avg_pool2d(x.primitive, output_size)) + } + OpsKind::UnTracked(prep) => { + prep.finish(B::adaptive_avg_pool2d(x.primitive, output_size)) + } + } + } + + fn adaptive_avg_pool2d_backward( + _x: AutodiffTensor, + _grad: AutodiffTensor, + ) -> AutodiffTensor { + panic!("Can't differentiate adaptive avg pool2d backward."); + } + + fn interpolate( + x: AutodiffTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> AutodiffTensor { + #[derive(Debug)] + struct Interpolate; + impl Backward for Interpolate { + type State = (NodeId, [usize; 2], InterpolateOptions); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_parent] = ops.parents; + let grad = grads.consume::(&ops.node); + + let (x_state, output_size, options) = ops.state; + let state = checkpointer.retrieve_node_output(x_state); + + if let Some(node) = node_parent { + let grad = B::interpolate_backward(state, grad, output_size, options); + grads.register::(node.id, grad); + } + } + } + + match Interpolate + .prepare::([x.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let x_state = prep.checkpoint(&x); + let output = B::interpolate(x.primitive.clone(), output_size, options.clone()); + prep.finish((x_state, output_size, options), output) + } + OpsKind::UnTracked(prep) => { + prep.finish(B::interpolate(x.primitive, output_size, options)) + } + } + } + + fn interpolate_backward( + _x: FloatTensor>, + _grad: FloatTensor>, + _output_size: [usize; 2], + _options: InterpolateOptions, + ) -> AutodiffTensor { + panic!("Can't differentiate interpolate backward."); + } + + fn attention( + query: FloatTensor>, + key: FloatTensor>, + value: FloatTensor>, + mask: Option>>, + attn_bias: Option>>, + options: AttentionModuleOptions, + ) -> FloatTensor> { + attention_fallback::(query, key, value, mask, attn_bias, options) + } + + fn ctc_loss( + log_probs: FloatTensor>, + targets: IntTensor>, + input_lengths: IntTensor>, + target_lengths: IntTensor>, + blank: usize, + ) -> FloatTensor> { + // Backends without a native ctc_loss_backward fall back to the default + // implementation, which is built from differentiable tensor ops so the + // autodiff layer derives the gradient automatically. + if !B::has_ctc_loss_backward() { + return burn_backend::ops::ctc::ctc_loss_default::( + log_probs, + targets, + input_lengths, + target_lengths, + blank, + ); + } + + #[derive(Debug)] + struct CtcLoss; + + impl Backward for CtcLoss { + type State = (NodeId, IntTensor, IntTensor, IntTensor, usize); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_parent] = ops.parents; + let grad_loss = grads.consume::(&ops.node); + + let (log_probs_state, targets, input_lengths, target_lengths, blank) = ops.state; + let log_probs: B::FloatTensorPrimitive = + checkpointer.retrieve_node_output(log_probs_state); + + if let Some(node) = node_parent { + let grad = B::ctc_loss_backward( + log_probs, + targets, + input_lengths, + target_lengths, + grad_loss, + blank, + ); + grads.register::(node.id, grad); + } + } + } + + match CtcLoss + .prepare::([log_probs.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let log_probs_state = prep.checkpoint(&log_probs); + let output = B::ctc_loss( + log_probs.primitive.clone(), + targets.clone(), + input_lengths.clone(), + target_lengths.clone(), + blank, + ); + prep.finish( + ( + log_probs_state, + targets, + input_lengths, + target_lengths, + blank, + ), + output, + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::ctc_loss( + log_probs.primitive, + targets, + input_lengths, + target_lengths, + blank, + )), + } + } + + fn rfft( + signal: FloatTensor>, + dim: usize, + n: Option, + ) -> (FloatTensor>, FloatTensor>) { + #[derive(Debug)] + struct Rfft; + + impl Backward for Rfft { + type State = (usize, Option, usize, usize, Vec, Vec); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + unary::(ops.parents, ops.node, grads, |grad| { + let (dim, n, input_len, n_fft, slices_re, slices_im) = ops.state; + + let grad_re = B::float_slice(grad.clone(), &slices_re); + let grad_im = B::float_slice(grad.clone(), &slices_im); + + let grad_re = mul_interior::(grad_re, dim, n_fft, 0.5); + let grad_im = mul_interior::(grad_im, dim, n_fft, 0.5); + + let grad = B::irfft(grad_re, grad_im, dim, n); + let grad = B::float_mul_scalar(grad, (n_fft as f64).into()); + + pad_to_length::(grad, dim, input_len) + }); + } + } + + let input_len = signal.shape()[dim]; + let n_fft = n.unwrap_or(input_len); + let (re, im) = B::rfft(signal.primitive, dim, n); + + // In order to perform only a single irfft in the backward pass, we have to temporarily + // bundle `re` and `im` into a single tensor. The following slice vecs are used to split + // them back up in both the forward and the backward pass. + let slices_re = re + .shape() + .iter() + .map(|&len| Slice::from(0..len)) + .collect::>(); + let slices_im = { + let mut slices = slices_re.clone(); + let len = slices[0].end.unwrap(); + slices[0].start = len; + slices[0].end = Some(2 * len); + slices + }; + let spectrum = B::float_cat(vec![re, im], 0); + let state = ( + dim, + n, + input_len, + n_fft, + slices_re.clone(), + slices_im.clone(), + ); + + let spectrum = match Rfft + .prepare::([signal.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish(state, spectrum), + OpsKind::UnTracked(prep) => prep.finish(spectrum), + }; + + let re = Self::float_slice(spectrum.clone(), &slices_re); + let im = Self::float_slice(spectrum.clone(), &slices_im); + + (re, im) + } + + fn irfft( + spectrum_re: FloatTensor>, + spectrum_im: FloatTensor>, + dim: usize, + n: Option, + ) -> FloatTensor> { + #[derive(Debug)] + struct Irfft; + + impl Backward for Irfft { + type State = (usize, Option, usize, usize); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (dim, n, input_len, n_fft) = ops.state; + let [node_re, node_im] = ops.parents; + let grad = grads.consume::(&ops.node); + + let grad = B::float_div_scalar(grad, (n_fft as f64).into()); + + let (grad_re, grad_im) = B::rfft(grad, dim, n); + + let grad_re = mul_interior::(grad_re, dim, n_fft, 2.0); + let grad_im = mul_interior::(grad_im, dim, n_fft, 2.0); + + let grad_re = pad_to_length::(grad_re, dim, input_len); + let grad_im = pad_to_length::(grad_im, dim, input_len); + + if let Some(node) = node_re { + grads.register::(node.id, grad_re); + } + if let Some(node) = node_im { + grads.register::(node.id, grad_im); + } + } + } + + let input_len = spectrum_re.shape()[dim]; + let signal = B::irfft(spectrum_re.primitive, spectrum_im.primitive, dim, n); + let n_fft = n.unwrap_or(signal.shape()[dim]); + let state = (dim, n, input_len, n_fft); + + match Irfft + .prepare::([spectrum_re.node.clone(), spectrum_im.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish(state, signal), + OpsKind::UnTracked(prep) => prep.finish(signal), + } + } +} + +// adapted from: burn_cubecl::kernel::fft::base::pad_to_length +fn pad_to_length(tensor: FloatTensor, dim: usize, target: usize) -> FloatTensor { + let shape = tensor.shape(); + let current = shape[dim]; + if current == target { + return tensor; + } + if current > target { + let slices: Vec<_> = shape + .iter() + .enumerate() + .map(|(i, &s)| Slice::from(if i == dim { 0..target } else { 0..s })) + .collect(); + return B::float_slice(tensor, &slices); + } + let mut padded_shape = shape.clone(); + padded_shape[dim] = target; + let padded = B::float_zeros(padded_shape, &tensor.device(), tensor.dtype().into()); + let slices: Vec = shape.iter().map(|&s| Slice::from(0..s)).collect(); + B::float_slice_assign(padded, &slices, tensor) +} + +fn mul_interior( + bins: FloatTensor, + dim: usize, + n_fft: usize, + factor: f64, +) -> FloatTensor { + // identify the interior bins (all bins except DC and Nyquist) + let slices_interior: Vec = { + let mut ranges = bins.shape().into_ranges(); + + // skip the DC bin + ranges[dim].start += 1; + + // if `n_fft` is even, we have a Nyquist bin to skip + if n_fft.is_multiple_of(2) { + ranges[dim].end -= 1; + } + + ranges.into_iter().map(Slice::from).collect() + }; + + // multiply only the interior bins by `factor` + let interior = B::float_slice(bins.clone(), &slices_interior); + let interior = B::float_mul_scalar(interior, factor.into()); + + B::float_slice_assign(bins, &slices_interior, interior) +} + +#[derive(Debug)] +struct MaxPool1D; + +impl Backward for MaxPool1D { + type State = (NodeId, IntTensor, usize, usize, usize, usize, bool); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_parent] = ops.parents; + let grad = grads.consume::(&ops.node); + let (x_state, indices, kernel_size, stride, padding, dilation, ceil_mode) = ops.state; + let x = checkpointer.retrieve_node_output(x_state); + + if let Some(node) = node_parent { + let grad = B::max_pool1d_with_indices_backward( + x, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + grad, + indices, + ); + + grads.register::(node.id, grad.x_grad); + } + } +} + +#[derive(Debug)] +struct MaxPool2D; + +impl Backward for MaxPool2D { + type State = ( + NodeId, + IntTensor, + [usize; 2], + [usize; 2], + [usize; 2], + [usize; 2], + bool, + ); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let [node_parent] = ops.parents; + let grad = grads.consume::(&ops.node); + let (x_state, indices, kernel_size, stride, padding, dilation, ceil_mode) = ops.state; + let x = checkpointer.retrieve_node_output(x_state); + + if let Some(node) = node_parent { + let grad = B::max_pool2d_with_indices_backward( + x, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + grad, + indices, + ); + + grads.register::(node.id, grad.x_grad); + } + } +} diff --git a/crates/burn-autodiff/src/ops/qtensor.rs b/crates/burn-autodiff/src/ops/qtensor.rs new file mode 100644 index 0000000..999b20d --- /dev/null +++ b/crates/burn-autodiff/src/ops/qtensor.rs @@ -0,0 +1,73 @@ +use burn_backend::{ + Backend, ExecutionError, TensorData, + ops::QTensorOps, + quantization::QuantizationParametersPrimitive, + tensor::{Device, FloatTensor, IntTensor, QuantizedTensor}, +}; +use burn_std::{FloatDType, IntDType, QuantScheme, Shape}; + +use crate::{Autodiff, checkpoint::strategy::CheckpointStrategy}; + +impl QTensorOps for Autodiff { + fn q_from_data(_data: TensorData, _device: &Device) -> QuantizedTensor { + todo!() + } + + fn quantize( + _tensor: FloatTensor, + _scheme: &QuantScheme, + _qparams: QuantizationParametersPrimitive, + ) -> QuantizedTensor { + todo!() // required for QAT + } + + fn quantize_dynamic( + _tensor: FloatTensor, + _scheme: &QuantScheme, + ) -> QuantizedTensor { + todo!() + } + + fn dequantize(_tensor: QuantizedTensor, _dtype: FloatDType) -> FloatTensor { + todo!() + } + + fn q_to_device( + _tensor: QuantizedTensor, + _device: &Device, + ) -> QuantizedTensor { + unimplemented!() + } + + fn q_reshape(tensor: QuantizedTensor, shape: Shape) -> QuantizedTensor { + B::q_reshape(tensor, shape) + } + + async fn q_into_data(tensor: QuantizedTensor) -> Result { + B::q_into_data(tensor).await + } + + fn q_swap_dims( + _tensor: QuantizedTensor, + _dim1: usize, + _dim2: usize, + ) -> QuantizedTensor { + unimplemented!() + } + + fn q_permute(_tensor: QuantizedTensor, _axes: &[usize]) -> QuantizedTensor { + unimplemented!() + } + + fn q_flip(_tensor: QuantizedTensor, _axes: &[usize]) -> QuantizedTensor { + unimplemented!() + } + + fn q_argmax(tensor: QuantizedTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + B::q_argmax(tensor, dim, out_dtype) + } + + fn q_argmin(tensor: QuantizedTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + B::q_argmin(tensor, dim, out_dtype) + } +} diff --git a/crates/burn-autodiff/src/ops/sort.rs b/crates/burn-autodiff/src/ops/sort.rs new file mode 100644 index 0000000..3a580a9 --- /dev/null +++ b/crates/burn-autodiff/src/ops/sort.rs @@ -0,0 +1,27 @@ +use super::{Backward, Ops, unary}; +use crate::{checkpoint::base::Checkpointer, grads::Gradients}; +use burn_backend::{Backend, TensorMetadata}; +use burn_std::Shape; + +#[derive(Debug)] +pub(crate) struct SortDim; + +impl Backward for SortDim { + type State = (B::IntTensorPrimitive, Shape, usize); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + unary::(ops.parents, ops.node, grads, |grad| { + let (indices, shape, dim) = ops.state; + let device = grad.device(); + let dtype = grad.dtype(); + let zeros = B::float_zeros(shape, &device, dtype.into()); + + B::float_scatter_add(dim, zeros, indices, grad) + }); + } +} diff --git a/crates/burn-autodiff/src/ops/tensor.rs b/crates/burn-autodiff/src/ops/tensor.rs new file mode 100644 index 0000000..38582d6 --- /dev/null +++ b/crates/burn-autodiff/src/ops/tensor.rs @@ -0,0 +1,4055 @@ +use alloc::{boxed::Box, vec, vec::Vec}; +use core::marker::PhantomData; + +#[cfg(not(feature = "std"))] +#[allow(unused_imports, reason = "required on aarch64, unused on x86_64")] +use num_traits::float::Float; + +use crate::{ + Autodiff, + checkpoint::{ + base::Checkpointer, builder::CheckpointerBuilder, retro_forward::RetroForward, + state::BackwardStates, strategy::CheckpointStrategy, + }, + grads::Gradients, + graph::{ComputingProperty, NodeId, NodeRef, Parent, Requirement, Step}, + ops::{Backward, Ops, OpsKind, binary, broadcast_shape, unary}, + retro_binary, retro_unary, retro_unary_scalar, + tensor::AutodiffTensor, + utils::duplicate, +}; + +use burn_backend::{ + Backend, ExecutionError, TensorData, TensorMetadata, get_device_settings, + ops::FloatTensorOps, + tensor::{BoolTensor, Device, FloatTensor, IntTensor}, +}; +use burn_backend::{Scalar, ops::unfold::calculate_unfold_windows}; +use burn_std::{BoolDType, FloatDType, IntDType, Shape, Slice}; + +use burn_backend::distributed::DistributedParams; + +use super::maxmin::MaxMinDim; + +// Unsqueeze op on primitive. +fn unsqueeze_like( + tensor: B::FloatTensorPrimitive, + shape: Shape, +) -> B::FloatTensorPrimitive { + let ndims_out = shape.num_dims(); + let shape = tensor.shape(); + let ndims_in = shape.num_dims(); + + let mut dims = vec![1; ndims_out]; + let num_ones = ndims_out - ndims_in; + dims[num_ones..(ndims_in + num_ones)].copy_from_slice(&shape[..ndims_in]); + + B::float_reshape(tensor, Shape::from(dims)) +} + +impl FloatTensorOps for Autodiff { + #[cfg_attr(feature = "tracing", tracing::instrument( + level="trace", + skip(data), + fields(?data.shape, ?data.dtype) + ))] + fn float_from_data(data: TensorData, device: &Device) -> FloatTensor { + AutodiffTensor::new(B::float_from_data(data, device)) + } + + fn float_random( + shape: Shape, + distribution: burn_backend::Distribution, + device: &Device, + dtype: FloatDType, + ) -> FloatTensor { + AutodiffTensor::new(B::float_random(shape, distribution, device, dtype)) + } + + fn float_zeros(shape: Shape, device: &Device, dtype: FloatDType) -> FloatTensor { + AutodiffTensor::new(B::float_zeros(shape, device, dtype)) + } + + fn float_ones(shape: Shape, device: &Device, dtype: FloatDType) -> FloatTensor { + AutodiffTensor::new(B::float_ones(shape, device, dtype)) + } + + #[cfg_attr(feature = "tracing", tracing::instrument( + level="trace", + skip(tensor), + fields( + from = ?tensor.node, + shape = ?tensor.shape(), + dtype = ?tensor.dtype(), + ) + ))] + async fn float_into_data(tensor: FloatTensor) -> Result { + B::float_into_data(tensor.primitive).await + } + + #[cfg_attr(feature = "tracing", tracing::instrument( + level="trace", + skip(tensor), + fields( + from = ?tensor.node, + shape = ?tensor.shape(), + dtype = ?tensor.dtype(), + ) + ))] + fn float_to_device(tensor: FloatTensor, device: &Device) -> FloatTensor { + #[derive(Debug)] + struct ToDevice; + + impl Backward for ToDevice { + type State = B::Device; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + unary::(ops.parents, ops.node, grads, |grad| { + B::float_to_device(grad, &ops.state) + }); + } + } + + match ToDevice + .prepare::([tensor.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => { + let device_old = tensor.primitive.device(); + prep.finish(device_old, B::float_to_device(tensor.primitive, device)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_to_device(tensor.primitive, device)), + } + } + + fn float_empty(shape: Shape, device: &Device, dtype: FloatDType) -> FloatTensor { + AutodiffTensor::new(B::float_empty(shape, device, dtype)) + } + + fn float_add(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Add; + + retro_binary!(RetroAdd, B::float_add); + + impl Backward for Add { + type State = (Shape, Shape); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (shape_lhs, shape_rhs) = ops.state; + + binary::( + ops.parents, + ops.node, + grads, + |grad| broadcast_shape::(grad, &shape_lhs), + |grad| broadcast_shape::(grad, &shape_rhs), + ); + } + } + + match Add + .prepare::([lhs.node.clone(), rhs.node.clone()]) + .memory_bound() + .retro_forward(RetroAdd::::new(lhs.node.id, rhs.node.id)) + .parents([&lhs, &rhs]) + .stateful() + { + OpsKind::Tracked(preps) => preps.finish( + (lhs.primitive.shape(), rhs.primitive.shape()), + B::float_add(lhs.primitive, rhs.primitive), + ), + OpsKind::UnTracked(preps) => preps.finish(B::float_add(lhs.primitive, rhs.primitive)), + } + } + + fn float_add_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + #[derive(Debug)] + struct AddScalar; + + retro_unary_scalar!(RetroAddScalar, B::float_add_scalar); + + impl Backward for AddScalar { + type State = (); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + unary::(ops.parents, ops.node, grads, |grad| grad); + } + } + + AddScalar + .prepare::([lhs.node.clone()]) + .memory_bound() + .retro_forward(RetroAddScalar::::new(lhs.node.id, rhs)) + .parents([&lhs]) + .stateless(B::float_add_scalar(lhs.primitive, rhs)) + } + + fn float_sub(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Sub; + + retro_binary!(RetroSub, B::float_sub); + + impl Backward for Sub { + type State = (Shape, Shape); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (shape_lhs, shape_rhs) = ops.state; + + binary::( + ops.parents, + ops.node, + grads, + |grad| broadcast_shape::(grad, &shape_lhs), + |grad| broadcast_shape::(B::float_neg(grad), &shape_rhs), + ); + } + } + + match Sub + .prepare::([lhs.node.clone(), rhs.node.clone()]) + .memory_bound() + .retro_forward(RetroSub::::new(lhs.node.id, rhs.node.id)) + .parents([&lhs, &rhs]) + .stateful() + { + OpsKind::Tracked(preps) => preps.finish( + (lhs.primitive.shape(), rhs.primitive.shape()), + B::float_sub(lhs.primitive, rhs.primitive), + ), + OpsKind::UnTracked(preps) => preps.finish(B::float_sub(lhs.primitive, rhs.primitive)), + } + } + + fn float_sub_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + #[derive(Debug)] + struct SubScalar; + + retro_unary_scalar!(RetroSubScalar, B::float_sub_scalar); + + impl Backward for SubScalar { + type State = (); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + unary::(ops.parents, ops.node, grads, |grad| grad); + } + } + + SubScalar + .prepare::([lhs.node.clone()]) + .memory_bound() + .retro_forward(RetroSubScalar::::new(lhs.node.id, rhs)) + .parents([&lhs]) + .stateless(B::float_sub_scalar(lhs.primitive, rhs)) + } + + fn float_mul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Mul; + + retro_binary!(RetroMul, B::float_mul); + + impl Backward for Mul { + type State = (Option, Option, BinaryOpsBroadcast); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let (lhs, rhs, broadcast) = ops.state; + let lhs = lhs.map(|lhs| checkpointer.retrieve_node_output(lhs)); + let rhs = rhs.map(|rhs| checkpointer.retrieve_node_output(rhs)); + + binary::( + ops.parents, + ops.node, + grads, + |grad| { + let grad = B::float_mul(grad, rhs.unwrap()); + broadcast.backward_lhs::(grad) + }, + |grad| { + let grad = B::float_mul(grad, lhs.unwrap()); + broadcast.backward_rhs::(grad) + }, + ); + } + } + + let lhs_tracked = lhs.is_tracked(); + let rhs_tracked = rhs.is_tracked(); + let broadcast = BinaryOpsBroadcast::new::(&lhs.primitive, &rhs.primitive); + + match Mul + .prepare::([lhs.node.clone(), rhs.node.clone()]) + .memory_bound() + .retro_forward(RetroMul::::new(lhs.node.id, rhs.node.id)) + .parents([&lhs, &rhs]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let lhs_state = rhs_tracked.then(|| prep.checkpoint(&lhs)); + let rhs_state = lhs_tracked.then(|| prep.checkpoint(&rhs)); + + prep.finish( + (lhs_state, rhs_state, broadcast), + B::float_mul(lhs.primitive, rhs.primitive), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_mul(lhs.primitive, rhs.primitive)), + } + } + + fn float_mul_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + #[derive(Debug)] + struct MulScalar; + + retro_unary_scalar!(RetroMulScalar, B::float_mul_scalar); + + impl Backward for MulScalar { + type State = Scalar; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + unary::(ops.parents, ops.node, grads, |grad| { + B::float_mul_scalar(grad, ops.state) + }); + } + } + + match MulScalar + .prepare::([lhs.node.clone()]) + .memory_bound() + .retro_forward(RetroMulScalar::::new(lhs.node.id, rhs)) + .parents([&lhs]) + .stateful() + { + OpsKind::Tracked(prep) => prep.finish(rhs, B::float_mul_scalar(lhs.primitive, rhs)), + OpsKind::UnTracked(prep) => prep.finish(B::float_mul_scalar(lhs.primitive, rhs)), + } + } + + fn float_div(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Div; + + retro_binary!(RetroDiv, B::float_div); + + impl Backward for Div { + type State = (Option, Option, BinaryOpsBroadcast); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let (lhs, rhs, broadcast) = ops.state; + let lhs = lhs.map(|lhs| checkpointer.retrieve_node_output(lhs)); + let rhs = rhs.map(|rhs| checkpointer.retrieve_node_output(rhs)); + let [rhs_4lhs, rhs_4rhs] = duplicate(&ops.parents, rhs); + + binary::( + ops.parents, + ops.node, + grads, + |grad| { + let rhs = rhs_4lhs.unwrap(); + let value = B::float_recip(rhs); + let grad = B::float_mul(grad, value); + + broadcast.backward_lhs::(grad) + }, + |grad| { + let rhs = rhs_4rhs.unwrap(); + let lhs = lhs.unwrap(); + let value = + B::float_div(B::float_neg(lhs), B::float_powi_scalar(rhs, 2.into())); + let grad = B::float_mul(grad, value); + + broadcast.backward_rhs::(grad) + }, + ); + } + } + + let lhs_tracked = lhs.is_tracked(); + let rhs_tracked = rhs.is_tracked(); + let broadcast = BinaryOpsBroadcast::new::(&lhs.primitive, &rhs.primitive); + + match Div + .prepare::([lhs.node.clone(), rhs.node.clone()]) + .memory_bound() + .retro_forward(RetroDiv::::new(lhs.node.id, rhs.node.id)) + .parents([&lhs, &rhs]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let lhs_state = rhs_tracked.then(|| prep.checkpoint(&lhs)); + let rhs_state = (lhs_tracked || rhs_tracked).then(|| prep.checkpoint(&rhs)); + + prep.finish( + (lhs_state, rhs_state, broadcast), + B::float_div(lhs.primitive, rhs.primitive), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_div(lhs.primitive, rhs.primitive)), + } + } + + fn float_div_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + #[derive(Debug)] + struct DivScalar; + + retro_unary_scalar!(RetroDivScalar, B::float_div_scalar); + + impl Backward for DivScalar { + type State = Scalar; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + unary::(ops.parents, ops.node, grads, |grad| { + let tmp = 1.0 / ops.state.elem::(); + B::float_mul_scalar(grad, tmp.into()) + }); + } + } + + match DivScalar + .prepare::([lhs.node.clone()]) + .memory_bound() + .retro_forward(RetroDivScalar::::new(lhs.node.id, rhs)) + .parents([&lhs]) + .stateful() + { + OpsKind::Tracked(prep) => prep.finish(rhs, B::float_div_scalar(lhs.primitive, rhs)), + OpsKind::UnTracked(prep) => prep.finish(B::float_div_scalar(lhs.primitive, rhs)), + } + } + + fn float_remainder(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Rem; + + retro_binary!(RetroRem, B::float_remainder); + + impl Backward for Rem { + type State = (Option, Option, BinaryOpsBroadcast); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let (lhs, rhs, broadcast) = ops.state; + let lhs = lhs.map(|lhs| checkpointer.retrieve_node_output(lhs)); + let rhs = rhs.map(|rhs| checkpointer.retrieve_node_output(rhs)); + + binary::( + ops.parents, + ops.node, + grads, + |grad| { + // remainder(x, y) = x - floor(x / y) * y + // partial(x - floor(x / y) * y, x) = 1 + broadcast.backward_lhs::(grad) + }, + |grad| { + // partial(x - floor(x / y) * y, y) = - floor(x / y) + let rhs = rhs.unwrap(); + let lhs = lhs.unwrap(); + let value = B::float_neg(B::float_floor(B::float_div(lhs, rhs))); + let grad = B::float_mul(grad, value); + broadcast.backward_rhs::(grad) + }, + ); + } + } + + let lhs_tracked = lhs.is_tracked(); + let rhs_tracked = rhs.is_tracked(); + let broadcast = BinaryOpsBroadcast::new::(&lhs.primitive, &rhs.primitive); + + match Rem + .prepare::([lhs.node.clone(), rhs.node.clone()]) + .memory_bound() + .retro_forward(RetroRem::::new(lhs.node.id, rhs.node.id)) + .parents([&lhs, &rhs]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let lhs_state = rhs_tracked.then(|| prep.checkpoint(&lhs)); + let rhs_state = (lhs_tracked || rhs_tracked).then(|| prep.checkpoint(&rhs)); + + prep.finish( + (lhs_state, rhs_state, broadcast), + B::float_remainder(lhs.primitive, rhs.primitive), + ) + } + OpsKind::UnTracked(prep) => { + prep.finish(B::float_remainder(lhs.primitive, rhs.primitive)) + } + } + } + + fn float_remainder_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + #[derive(Debug)] + struct RemainderScalar; + + retro_unary_scalar!(RetroRemainderScalar, B::float_remainder_scalar); + + impl Backward for RemainderScalar { + type State = (); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + unary::(ops.parents, ops.node, grads, |grad| grad); + } + } + + RemainderScalar + .prepare::([lhs.node.clone()]) + .memory_bound() + .retro_forward(RetroRemainderScalar::::new(lhs.node.id, rhs)) + .parents([&lhs]) + .stateless(B::float_remainder_scalar(lhs.primitive, rhs)) + } + + fn float_matmul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Matmul; + + impl Backward for Matmul { + type State = (Option, Option, BinaryOpsBroadcast); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let (lhs, rhs, broadcast) = ops.state; + let lhs = lhs.map(|lhs| checkpointer.retrieve_node_output(lhs)); + let rhs = rhs.map(|rhs| checkpointer.retrieve_node_output(rhs)); + + binary::( + ops.parents, + ops.node, + grads, + |grad| { + let rhs = B::float_transpose(rhs.unwrap()); + let grad = B::float_matmul(grad, rhs); + + broadcast.backward_lhs::(grad) + }, + |grad| { + let lhs = B::float_transpose(lhs.unwrap()); + let grad = B::float_matmul(lhs, grad); + + broadcast.backward_rhs::(grad) + }, + ); + } + } + + let lhs_tracked = lhs.is_tracked(); + let rhs_tracked = rhs.is_tracked(); + let broadcast = BinaryOpsBroadcast::new::(&lhs.primitive, &rhs.primitive); + + match Matmul + .prepare::([lhs.node.clone(), rhs.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let lhs_state = rhs_tracked.then(|| prep.checkpoint(&lhs)); + let rhs_state = lhs_tracked.then(|| prep.checkpoint(&rhs)); + prep.finish( + (lhs_state, rhs_state, broadcast), + B::float_matmul(lhs.primitive, rhs.primitive), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_matmul(lhs.primitive, rhs.primitive)), + } + } + + fn float_cross( + lhs: FloatTensor, + rhs: FloatTensor, + dim: usize, + ) -> FloatTensor { + #[derive(Debug)] + struct Cross; + + impl Backward for Cross { + type State = (Option, Option, usize); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let (lhs_id, rhs_id, dim) = ops.state; + let lhs = lhs_id.map(|id| checkpointer.retrieve_node_output(id)); + let rhs = rhs_id.map(|id| checkpointer.retrieve_node_output(id)); + + binary::( + ops.parents, + ops.node, + grads, + |grad| B::float_cross(rhs.unwrap(), grad, dim), + |grad| B::float_cross(grad, lhs.unwrap(), dim), + ); + } + } + + let lhs_tracked = lhs.is_tracked(); + let rhs_tracked = rhs.is_tracked(); + + match Cross + .prepare::([lhs.node.clone(), rhs.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let lhs_state = rhs_tracked.then(|| prep.checkpoint(&lhs)); + let rhs_state = lhs_tracked.then(|| prep.checkpoint(&rhs)); + prep.finish( + (lhs_state, rhs_state, dim), + B::float_cross(lhs.primitive, rhs.primitive, dim), + ) + } + OpsKind::UnTracked(prep) => { + prep.finish(B::float_cross(lhs.primitive, rhs.primitive, dim)) + } + } + } + + fn float_neg(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Neg; + + retro_unary!(RetroNeg, B::float_neg); + + impl Backward for Neg { + type State = (); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + unary::(ops.parents, ops.node, grads, |grad| B::float_neg(grad)); + } + } + + Neg.prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroNeg::::new(tensor.node.id)) + .parents([&tensor]) + .stateless(B::float_neg(tensor.primitive)) + } + + fn float_recip(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Recip; + + retro_unary!(RetroRecip, B::float_recip); + + impl Backward for Recip { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let tensor = checkpointer.retrieve_node_output(ops.state); + unary::(ops.parents, ops.node, grads, |grad| { + let tmp = B::float_powi_scalar(tensor, (-2).into()); + let value = B::float_neg(tmp); + + B::float_mul(grad, value) + }); + } + } + + match Recip + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroRecip::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_recip(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_recip(tensor.primitive)), + } + } + + fn float_swap_dims(tensor: FloatTensor, dim1: usize, dim2: usize) -> FloatTensor { + #[derive(Debug)] + struct SwapDim; + + #[derive(new, Debug)] + struct RetroSwapDims { + input_id: NodeId, + dim1: usize, + dim2: usize, + _backend: PhantomData, + } + + impl RetroForward for RetroSwapDims { + fn forward(&self, states: &mut BackwardStates, out_node: NodeId) { + let input = states.get_state::(&self.input_id); + let out = B::float_swap_dims(input, self.dim1, self.dim2); + states.save(out_node, out) + } + } + + impl Backward for SwapDim { + type State = (usize, usize); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (dim1, dim2) = ops.state; + + unary::(ops.parents, ops.node, grads, |grad| { + B::float_swap_dims(grad, dim2, dim1) + }); + } + } + + match SwapDim + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroSwapDims::::new(tensor.node.id, dim1, dim2)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + (dim1, dim2), + B::float_swap_dims(tensor.primitive, dim1, dim2), + ), + OpsKind::UnTracked(prep) => { + prep.finish(B::float_swap_dims(tensor.primitive, dim1, dim2)) + } + } + } + + fn float_permute(tensor: FloatTensor, axes: &[usize]) -> FloatTensor { + #[derive(Debug)] + struct PermuteDim; + + #[derive(new, Debug)] + struct RetroPermuteDims { + input_id: NodeId, + axes: Vec, + _backend: PhantomData, + } + + impl RetroForward for RetroPermuteDims { + fn forward(&self, states: &mut BackwardStates, out_node: NodeId) { + let input = states.get_state::(&self.input_id); + let out = B::float_permute(input, &self.axes); + states.save(out_node, out) + } + } + + impl Backward for PermuteDim { + type State = Vec; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let axes = ops.state; + + let mut inverse = vec![0usize; axes.len()]; + axes.iter() + .enumerate() + .for_each(|(i, &axis)| inverse[axis] = i); + + unary::(ops.parents, ops.node, grads, |grad| { + B::float_permute(grad, &inverse) + }); + } + } + + match PermuteDim + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroPermuteDims::::new(tensor.node.id, axes.to_vec())) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(prep) => { + prep.finish(axes.to_vec(), B::float_permute(tensor.primitive, axes)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_permute(tensor.primitive, axes)), + } + } + + fn float_flip(tensor: FloatTensor, axes: &[usize]) -> FloatTensor { + #[derive(Debug)] + struct FlipDim; + + #[derive(new, Debug)] + struct RetroFlipDims { + input_id: NodeId, + axes: Vec, + _backend: PhantomData, + } + + impl RetroForward for RetroFlipDims { + fn forward(&self, states: &mut BackwardStates, out_node: NodeId) { + let input = states.get_state::(&self.input_id); + let out = B::float_flip(input, &self.axes); + states.save(out_node, out) + } + } + + impl Backward for FlipDim { + type State = Vec; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let axes = ops.state; + + unary::(ops.parents, ops.node, grads, |grad| { + B::float_flip(grad, &axes) + }); + } + } + + match FlipDim + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroFlipDims::::new(tensor.node.id, axes.to_vec())) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(prep) => { + prep.finish(axes.to_vec(), B::float_flip(tensor.primitive, axes)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_flip(tensor.primitive, axes)), + } + } + + fn float_reshape(tensor: FloatTensor, shape: Shape) -> FloatTensor { + #[derive(Debug)] + struct ReshapeDim; + + #[derive(new, Debug)] + struct RetroReshape { + input_id: NodeId, + shape: Shape, + _backend: PhantomData, + } + + impl RetroForward for RetroReshape { + fn forward(&self, states: &mut BackwardStates, out_node: NodeId) { + let input = states.get_state::(&self.input_id); + let out = B::float_reshape(input, self.shape.clone()); + states.save(out_node, out) + } + } + + impl Backward for ReshapeDim { + type State = (Shape, Shape); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (shape_original, shape) = ops.state; + let ndims_out = shape.num_dims(); + + unary::(ops.parents, ops.node, grads, |grad| { + let shape_grad = grad.shape(); + let mut grad = grad; + + for i in 0..ndims_out { + if shape[i] == 1 && shape_grad[i] != 1 { + grad = B::float_sum_dim(grad, i); + } + } + + B::float_reshape(grad, shape_original) + }); + } + } + + match ReshapeDim + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroReshape::::new(tensor.node.id, shape.clone())) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + (tensor.primitive.shape(), shape.clone()), + B::float_reshape(tensor.primitive, shape), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_reshape(tensor.primitive, shape)), + } + } + + fn float_gather( + dim: usize, + tensor: FloatTensor, + indices: IntTensor, + ) -> FloatTensor { + #[derive(Debug)] + struct Gather; + + impl Backward for Gather { + type State = (usize, IntTensor, Shape, B::Device); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (dim, indices, shape, device) = ops.state; + + unary::(ops.parents, ops.node, grads, |grad| { + let zeros = B::float_zeros(shape, &device, grad.dtype().into()); + B::float_scatter_add(dim, zeros, indices, grad) + }); + } + } + + match Gather + .prepare::([tensor.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + ( + dim, + indices.clone(), + tensor.primitive.shape(), + tensor.primitive.device(), + ), + B::float_gather(dim, tensor.primitive, indices), + ), + OpsKind::UnTracked(prep) => { + prep.finish(B::float_gather(dim, tensor.primitive, indices)) + } + } + } + + fn float_scatter_add( + dim: usize, + tensor: FloatTensor, + indices: IntTensor, + value: FloatTensor, + ) -> FloatTensor { + #[derive(Debug)] + struct Scatter; + + impl Backward for Scatter { + type State = (usize, IntTensor); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (dim, indices) = ops.state; + let [_, indices_4rhs] = duplicate(&ops.parents, Some(indices)); + + binary::( + ops.parents, + ops.node, + grads, + |grad| grad, + |grad| B::float_gather(dim, grad, indices_4rhs.unwrap()), + ); + } + } + + match Scatter + .prepare::([tensor.node, value.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + (dim, indices.clone()), + B::float_scatter_add(dim, tensor.primitive, indices, value.primitive), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_scatter_add( + dim, + tensor.primitive, + indices, + value.primitive, + )), + } + } + + fn float_scatter_nd( + data: FloatTensor, + indices: IntTensor, + values: FloatTensor, + reduction: burn_backend::tensor::IndexingUpdateOp, + ) -> FloatTensor { + use burn_backend::tensor::IndexingUpdateOp; + + match reduction { + IndexingUpdateOp::Add => { + #[derive(Debug)] + struct ScatterNdAdd; + + impl Backward for ScatterNdAdd { + type State = IntTensor; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let indices = ops.state; + let [_, indices_4rhs] = duplicate(&ops.parents, Some(indices)); + + binary::( + ops.parents, + ops.node, + grads, + |grad| grad, + |grad| B::float_gather_nd(grad, indices_4rhs.unwrap()), + ); + } + } + + match ScatterNdAdd + .prepare::([data.node, values.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + indices.clone(), + B::float_scatter_nd( + data.primitive, + indices, + values.primitive, + IndexingUpdateOp::Add, + ), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_scatter_nd( + data.primitive, + indices, + values.primitive, + IndexingUpdateOp::Add, + )), + } + } + IndexingUpdateOp::Assign => { + #[derive(Debug)] + struct ScatterNdAssign; + + impl Backward for ScatterNdAssign { + type State = (IntTensor, Shape, B::Device); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (indices, values_shape, device) = ops.state; + let [indices_4lhs, indices_4rhs] = duplicate(&ops.parents, Some(indices)); + + binary::( + ops.parents, + ops.node, + grads, + |grad| { + // Zero out the scattered positions in grad + let zeros = + B::float_zeros(values_shape, &device, grad.dtype().into()); + B::float_scatter_nd( + grad, + indices_4lhs.unwrap(), + zeros, + IndexingUpdateOp::Assign, + ) + }, + |grad| B::float_gather_nd(grad, indices_4rhs.unwrap()), + ); + } + } + + match ScatterNdAssign + .prepare::([data.node, values.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + ( + indices.clone(), + values.primitive.shape(), + data.primitive.device(), + ), + B::float_scatter_nd( + data.primitive, + indices, + values.primitive, + IndexingUpdateOp::Assign, + ), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_scatter_nd( + data.primitive, + indices, + values.primitive, + IndexingUpdateOp::Assign, + )), + } + } + IndexingUpdateOp::Mul => { + // Forward: out[idx[i]] = data[idx[i]] * values[i] (unique indices assumed). + // Backward: + // grad_data = grad * scatter_nd(ones_like(data), idx, values, Assign) + // grad_values = gather_nd(grad, idx) * gather_nd(data, idx) + #[derive(Debug)] + struct ScatterNdMul; + + impl Backward for ScatterNdMul { + type State = (Option>, Option>, IntTensor); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (data_state, values_state, indices) = ops.state; + let [indices_4lhs, indices_4rhs] = duplicate(&ops.parents, Some(indices)); + + binary::( + ops.parents, + ops.node, + grads, + |grad| { + let device = grad.device(); + let dtype = grad.dtype(); + let shape = grad.shape(); + let ones = B::float_ones(shape, &device, dtype.into()); + let mult = B::float_scatter_nd( + ones, + indices_4lhs.unwrap(), + values_state.unwrap(), + IndexingUpdateOp::Assign, + ); + B::float_mul(grad, mult) + }, + |grad| { + let indices = indices_4rhs.unwrap(); + let g_idx = B::float_gather_nd(grad, indices.clone()); + let d_idx = B::float_gather_nd(data_state.unwrap(), indices); + B::float_mul(g_idx, d_idx) + }, + ); + } + } + + let data_tracked = data.is_tracked(); + let values_tracked = values.is_tracked(); + + match ScatterNdMul + .prepare::([data.node, values.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => { + let data_state = values_tracked.then(|| data.primitive.clone()); + let values_state = data_tracked.then(|| values.primitive.clone()); + prep.finish( + (data_state, values_state, indices.clone()), + B::float_scatter_nd( + data.primitive, + indices, + values.primitive, + IndexingUpdateOp::Mul, + ), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_scatter_nd( + data.primitive, + indices, + values.primitive, + IndexingUpdateOp::Mul, + )), + } + } + IndexingUpdateOp::Min | IndexingUpdateOp::Max => { + // Unique indices are required for this backward formula; duplicate indices have + // undefined behavior for scatter_nd Min/Max gradients. + // Forward (Max): out[idx] = max(data[idx], values); non-scattered: out = data. + // Backward, with ties contributing to both sides (matches cummin/cummax convention): + // data_at_idx = gather_nd(data, idx) + // data_won = data_at_idx >= values (Max) / <= values (Min) + // values_won = values >= data_at_idx (Max) / <= data_at_idx (Min) + // data_mask = scatter_nd(ones_like(data), idx, data_won, Assign) + // grad_data = grad * data_mask + // grad_values = gather_nd(grad, idx) * values_won + #[derive(Debug)] + struct ScatterNdMinMax; + + impl Backward for ScatterNdMinMax { + type State = (FloatTensor, FloatTensor, IntTensor, bool); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (data, values, indices, is_max) = ops.state; + let device = data.device(); + let data_dtype = data.dtype(); + let data_shape = data.shape(); + let settings = get_device_settings::(&device); + let bool_dtype = settings.bool_dtype; + + // data_at_idx — needed for both winner masks + let data_at_idx = B::float_gather_nd(data, indices.clone()); + + let (data_won_bool, values_won_bool) = if is_max { + ( + B::float_greater_equal( + data_at_idx.clone(), + values.clone(), + bool_dtype, + ), + B::float_greater_equal(values, data_at_idx, bool_dtype), + ) + } else { + ( + B::float_lower_equal( + data_at_idx.clone(), + values.clone(), + bool_dtype, + ), + B::float_lower_equal(values, data_at_idx, bool_dtype), + ) + }; + + let data_won_float = B::bool_into_float(data_won_bool, data_dtype.into()); + let values_won_float = + B::bool_into_float(values_won_bool, data_dtype.into()); + + // Build the full-shape data mask: 1 at non-scattered positions, + // data_won_float at scattered positions. + let ones = B::float_ones(data_shape, &device, data_dtype.into()); + let data_mask = B::float_scatter_nd( + ones, + indices.clone(), + data_won_float, + IndexingUpdateOp::Assign, + ); + + binary::( + ops.parents, + ops.node, + grads, + |grad| B::float_mul(grad, data_mask), + |grad| { + let g_idx = B::float_gather_nd(grad, indices); + B::float_mul(g_idx, values_won_float) + }, + ); + } + } + + let is_max = matches!(reduction, IndexingUpdateOp::Max); + + match ScatterNdMinMax + .prepare::([data.node, values.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + ( + data.primitive.clone(), + values.primitive.clone(), + indices.clone(), + is_max, + ), + B::float_scatter_nd(data.primitive, indices, values.primitive, reduction), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_scatter_nd( + data.primitive, + indices, + values.primitive, + reduction, + )), + } + } + } + } + + fn float_gather_nd(data: FloatTensor, indices: IntTensor) -> FloatTensor { + #[derive(Debug)] + struct GatherNd; + + impl Backward for GatherNd { + type State = (IntTensor, Shape, B::Device); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (indices, data_shape, device) = ops.state; + + unary::(ops.parents, ops.node, grads, |grad| { + let zeros = B::float_zeros(data_shape, &device, grad.dtype().into()); + B::float_scatter_nd( + zeros, + indices, + grad, + burn_backend::tensor::IndexingUpdateOp::Add, + ) + }); + } + } + + match GatherNd + .prepare::([data.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + ( + indices.clone(), + data.primitive.shape(), + data.primitive.device(), + ), + B::float_gather_nd(data.primitive, indices), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_gather_nd(data.primitive, indices)), + } + } + + fn float_select( + tensor: FloatTensor, + dim: usize, + indices: IntTensor, + ) -> FloatTensor { + #[derive(Debug)] + struct Select; + + #[derive(new, Debug)] + struct RetroSelect { + input_id: NodeId, + dim: usize, + indices: IntTensor, + } + + impl RetroForward for RetroSelect { + fn forward(&self, states: &mut BackwardStates, out_node: NodeId) { + let input = states.get_state::(&self.input_id); + let out = B::float_select(input, self.dim, self.indices.clone()); + states.save(out_node, out) + } + } + + impl Backward for Select { + type State = (usize, IntTensor, Shape, B::Device); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (dim, indices, shape, device) = ops.state; + + unary::(ops.parents, ops.node, grads, |grad| { + let zeros = B::float_zeros(shape, &device, grad.dtype().into()); + B::float_select_add(zeros, dim, indices, grad) + }); + } + } + + match Select + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroSelect::::new(tensor.node.id, dim, indices.clone())) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + ( + dim, + indices.clone(), + tensor.primitive.shape(), + tensor.primitive.device(), + ), + B::float_select(tensor.primitive, dim, indices), + ), + OpsKind::UnTracked(prep) => { + prep.finish(B::float_select(tensor.primitive, dim, indices)) + } + } + } + + fn float_select_add( + tensor: FloatTensor, + dim: usize, + indices: IntTensor, + value: FloatTensor, + ) -> FloatTensor { + #[derive(Debug)] + struct IndexSelectDimAssign; + + #[derive(new, Debug)] + struct RetroSelectAssign { + tensor_id: NodeId, + dim: usize, + indices: IntTensor, + value_id: NodeId, + } + + impl RetroForward for RetroSelectAssign { + fn forward(&self, states: &mut BackwardStates, out_node: NodeId) { + let tensor = states.get_state::(&self.tensor_id); + let value = states.get_state::(&self.value_id); + let out = B::float_select_add(tensor, self.dim, self.indices.clone(), value); + states.save(out_node, out) + } + } + + impl Backward for IndexSelectDimAssign { + type State = (usize, IntTensor); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (dim, indices) = ops.state; + + binary::( + ops.parents, + ops.node, + grads, + |grad| grad, + |grad| B::float_select(grad, dim, indices), + ); + } + } + + match IndexSelectDimAssign + .prepare::([tensor.node.clone(), value.node.clone()]) + .memory_bound() + .retro_forward(RetroSelectAssign::::new( + tensor.node.id, + dim, + indices.clone(), + value.node.id, + )) + .parents([&tensor, &value]) + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + (dim, indices.clone()), + B::float_select_add(tensor.primitive, dim, indices, value.primitive), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_select_add( + tensor.primitive, + dim, + indices, + value.primitive, + )), + } + } + + fn float_slice(tensor: FloatTensor, slices: &[Slice]) -> FloatTensor { + #[derive(Debug)] + struct Index; + + #[derive(new, Debug)] + struct RetroSlice { + tensor_id: NodeId, + slices: Vec, + _backend: PhantomData, + } + + impl RetroForward for RetroSlice { + fn forward(&self, states: &mut BackwardStates, out_node: NodeId) { + let tensor = states.get_state::(&self.tensor_id); + let out = B::float_slice(tensor, &self.slices); + states.save(out_node, out) + } + } + + impl Backward for Index { + type State = (Vec, Shape, B::Device); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (slices, shape, device) = ops.state; + + unary::(ops.parents, ops.node, grads, |grad| { + let zeros = B::float_zeros(shape, &device, grad.dtype().into()); + B::float_slice_assign(zeros, &slices, grad) + }); + } + } + + match Index + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroSlice::::new(tensor.node.id, slices.to_vec())) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + ( + slices.to_vec(), + tensor.primitive.shape(), + tensor.primitive.device(), + ), + B::float_slice(tensor.primitive, slices), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_slice(tensor.primitive, slices)), + } + } + + fn float_slice_assign( + tensor: FloatTensor, + slices: &[Slice], + value: FloatTensor, + ) -> FloatTensor { + #[derive(Debug)] + struct SliceAssign; + + #[derive(new, Debug)] + struct RetroSliceAssign { + tensor_id: NodeId, + slices: Vec, + value_id: NodeId, + _backend: PhantomData, + } + + impl RetroForward for RetroSliceAssign { + fn forward(&self, states: &mut BackwardStates, out_node: NodeId) { + let tensor = states.get_state::(&self.tensor_id); + let value = states.get_state::(&self.value_id); + let out = B::float_slice_assign(tensor, &self.slices, value); + states.save(out_node, out) + } + } + + impl Backward for SliceAssign { + type State = (Vec, Shape, B::Device); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (slices, shape_rhs, device) = ops.state; + let [slices_4lhs, slices_4rhs] = duplicate(&ops.parents, Some(slices)); + + binary::( + ops.parents, + ops.node, + grads, + |grad| { + let zeros = B::float_zeros(shape_rhs, &device, grad.dtype().into()); + B::float_slice_assign(grad, &slices_4lhs.unwrap(), zeros) + }, + |grad| B::float_slice(grad, &slices_4rhs.unwrap()), + ); + } + } + + match SliceAssign + .prepare::([tensor.node.clone(), value.node.clone()]) + .memory_bound() + .retro_forward(RetroSliceAssign::::new( + tensor.node.id, + slices.to_vec(), + value.node.id, + )) + .parents([&tensor, &value]) + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + ( + slices.to_vec(), + value.primitive.shape(), + value.primitive.device(), + ), + B::float_slice_assign(tensor.primitive, slices, value.primitive), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_slice_assign( + tensor.primitive, + slices, + value.primitive, + )), + } + } + + fn float_mask_where( + tensor: FloatTensor, + mask: BoolTensor, + source: FloatTensor, + ) -> FloatTensor { + #[derive(Debug)] + struct MaskWhere; + + impl Backward for MaskWhere { + type State = (BoolTensor, Shape, Shape, B::Device); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (mask, shape_lhs, shape_rhs, device) = ops.state; + let [mask_4lhs, mask_4rhs] = duplicate(&ops.parents, Some(mask)); + + binary::( + ops.parents, + ops.node, + grads, + |grad| { + let zeros = B::float_zeros(shape_lhs.clone(), &device, grad.dtype().into()); + let grad = B::float_mask_where(grad, mask_4lhs.unwrap(), zeros); + + broadcast_shape::(grad, &shape_lhs) + }, + |grad| { + let zeros = B::float_zeros(shape_rhs.clone(), &device, grad.dtype().into()); + let grad = B::float_mask_where(zeros, mask_4rhs.unwrap(), grad); + + broadcast_shape::(grad, &shape_rhs) + }, + ); + } + } + + match MaskWhere + .prepare::([tensor.node, source.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + ( + mask.clone(), + tensor.primitive.shape(), + source.primitive.shape(), + source.primitive.device(), + ), + B::float_mask_where(tensor.primitive, mask, source.primitive), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_mask_where( + tensor.primitive, + mask, + source.primitive, + )), + } + } + + fn float_mask_fill( + tensor: FloatTensor, + mask: BoolTensor, + value: Scalar, + ) -> FloatTensor { + #[derive(Debug)] + struct MaskFill; + + impl Backward for MaskFill { + type State = BoolTensor; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + unary::(ops.parents, ops.node, grads, |grad| { + B::float_mask_fill(grad, ops.state, 0f32.into()) + }); + } + } + + match MaskFill + .prepare::([tensor.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + mask.clone(), + B::float_mask_fill(tensor.primitive, mask, value), + ), + OpsKind::UnTracked(prep) => { + prep.finish(B::float_mask_fill(tensor.primitive, mask, value)) + } + } + } + + fn float_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + B::float_equal(lhs.primitive, rhs.primitive, out_dtype) + } + + fn float_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + B::float_equal_elem(lhs.primitive, rhs, out_dtype) + } + + fn float_greater( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + B::float_greater(lhs.primitive, rhs.primitive, out_dtype) + } + + fn float_greater_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + B::float_greater_elem(lhs.primitive, rhs, out_dtype) + } + + fn float_greater_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + B::float_greater_equal(lhs.primitive, rhs.primitive, out_dtype) + } + + fn float_greater_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + B::float_greater_equal_elem(lhs.primitive, rhs, out_dtype) + } + + fn float_lower( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + B::float_lower(lhs.primitive, rhs.primitive, out_dtype) + } + + fn float_lower_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + B::float_lower_elem(lhs.primitive, rhs, out_dtype) + } + + fn float_lower_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + B::float_lower_equal(lhs.primitive, rhs.primitive, out_dtype) + } + + fn float_lower_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + B::float_lower_equal_elem(lhs.primitive, rhs, out_dtype) + } + + fn float_is_nan(tensor: FloatTensor, out_dtype: BoolDType) -> BoolTensor { + B::float_is_nan(tensor.primitive, out_dtype) + } + + fn float_is_inf(tensor: FloatTensor, out_dtype: BoolDType) -> BoolTensor { + B::float_is_inf(tensor.primitive, out_dtype) + } + + fn float_detach(tensor: FloatTensor) -> FloatTensor { + // When we detach a tensor, we remove it from the graph, but we still want to keep the + // `require_grad` (and `distributed`) setting. + let is_require_grad = Self::float_is_require_grad(&tensor); + + let distributed_params = tensor.node.distributed_params.clone(); + + let mut tensor = AutodiffTensor::new(tensor.primitive); + + if is_require_grad { + tensor = tensor.require_grad(); + } + if let Some(params) = distributed_params { + tensor = tensor.grad_distributed(params.param_id); + } + tensor + } + + fn float_set_require_grad(tensor: FloatTensor, require_grad: bool) -> FloatTensor { + if require_grad { + return tensor.require_grad(); + } + + AutodiffTensor::new(tensor.primitive) + } + + fn float_is_require_grad(tensor: &FloatTensor) -> bool { + matches!(tensor.node.requirement, Requirement::Grad) + } + + fn float_mean(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Mean; + + impl Backward for Mean { + type State = Shape; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + unary::(ops.parents, ops.node, grads, |grad| { + let shape = ops.state; + let val = 1_f64 / shape.num_elements() as f64; + let ones = B::float_ones(shape, &grad.device(), grad.dtype().into()); + let val = B::float_mul_scalar(ones, val.into()); + + let grad = unsqueeze_like::(grad, val.shape()); + B::float_mul(val, grad) + }); + } + } + + match Mean.prepare::([tensor.node]).compute_bound().stateful() { + OpsKind::Tracked(prep) => { + prep.finish(tensor.primitive.shape(), B::float_mean(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_mean(tensor.primitive)), + } + } + + fn float_sum(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Sum; + + impl Backward for Sum { + type State = Shape; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + unary::(ops.parents, ops.node, grads, |grad| { + let val = B::float_ones(ops.state, &grad.device(), grad.dtype().into()); + + let grad = unsqueeze_like::(grad, val.shape()); + B::float_mul(val, grad) + }); + } + } + + match Sum.prepare::([tensor.node]).compute_bound().stateful() { + OpsKind::Tracked(prep) => { + prep.finish(tensor.primitive.shape(), B::float_sum(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_sum(tensor.primitive)), + } + } + + fn float_mean_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + #[derive(Debug)] + struct MeanDim; + + impl Backward for MeanDim { + type State = (Shape, usize); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (shape, dim) = ops.state; + + unary::(ops.parents, ops.node, grads, |grad| { + let val = 1_f64 / shape[dim] as f64; + let ones = B::float_ones(shape, &grad.device(), grad.dtype().into()); + let val = B::float_mul_scalar(ones, val.into()); + + let grad = B::float_sum_dim(grad, dim); + B::float_mul(val, grad) + }); + } + } + + match MeanDim + .prepare::([tensor.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + (tensor.primitive.shape(), dim), + B::float_mean_dim(tensor.primitive, dim), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_mean_dim(tensor.primitive, dim)), + } + } + + fn float_sum_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + #[derive(Debug)] + struct SumDim; + + impl Backward for SumDim { + type State = (Shape, usize); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (shape, dim) = ops.state; + + unary::(ops.parents, ops.node, grads, |grad| { + let ones = B::float_ones(shape, &grad.device(), grad.dtype().into()); + let grad = B::float_sum_dim(grad, dim); + + B::float_mul(ones, grad) + }); + } + } + + match SumDim + .prepare::([tensor.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + (tensor.primitive.shape(), dim), + B::float_sum_dim(tensor.primitive, dim), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_sum_dim(tensor.primitive, dim)), + } + } + + fn float_prod(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Prod; + + impl Backward for Prod { + // Saves the input and the output product so backward can compute + // `grad * prod(x) / x` without recomputing the reduction. + type State = (B::FloatTensorPrimitive, B::FloatTensorPrimitive); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (input, output) = ops.state; + + unary::(ops.parents, ops.node, grads, |grad| { + // d/dx_i prod(x) = prod(x) / x_i, so grad_input = grad * output / input, + // broadcast over the input shape (output is a single-element tensor). + // + // This divides by the input, so it produces NaN gradients when the + // input contains zeros. A zero-safe version requires the product of + // all other elements via exclusive cumulative products, same as the + // cumprod limitation tracked in https://github.com/tracel-ai/burn/issues/3864. + let ones = B::float_ones(input.shape(), &input.device(), input.dtype().into()); + let grad = B::float_mul(grad, output); + let grad = unsqueeze_like::(grad, ones.shape()); + let grad = B::float_mul(ones, grad); + + B::float_div(grad, input) + }); + } + } + + match Prod.prepare::([tensor.node]).compute_bound().stateful() { + OpsKind::Tracked(prep) => { + let output = B::float_prod(tensor.primitive.clone()); + prep.finish((tensor.primitive, output.clone()), output) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_prod(tensor.primitive)), + } + } + + fn float_prod_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + #[derive(Debug)] + struct ProdDim; + + impl Backward for ProdDim { + // Saves the input and the reduced product (size 1 along `dim`). + type State = (B::FloatTensorPrimitive, B::FloatTensorPrimitive); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (input, output) = ops.state; + + unary::(ops.parents, ops.node, grads, |grad| { + // grad_input = grad * prod_dim(x) / x. The grad and output both keep + // a size-1 reduced dim and broadcast back over the input along `dim`. + // + // Like `float_prod`, this divides by the input and produces NaN + // gradients when the input contains zeros (see + // https://github.com/tracel-ai/burn/issues/3864). + let ones = B::float_ones(input.shape(), &input.device(), input.dtype().into()); + let grad = B::float_mul(grad, output); + let grad = B::float_mul(ones, grad); + + B::float_div(grad, input) + }); + } + } + + match ProdDim + .prepare::([tensor.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => { + let output = B::float_prod_dim(tensor.primitive.clone(), dim); + prep.finish((tensor.primitive, output.clone()), output) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_prod_dim(tensor.primitive, dim)), + } + } + + fn float_cumsum(tensor: FloatTensor, dim: usize) -> FloatTensor { + #[derive(Debug)] + struct CumSum; + + impl Backward for CumSum { + type State = (Shape, usize); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (_shape, dim) = ops.state; + + unary::(ops.parents, ops.node, grads, |grad| { + // Gradient of cumsum is cumsum of gradient in reverse + let grad_reversed = B::float_flip(grad.clone(), &[dim]); + let grad_cumsum = B::float_cumsum(grad_reversed, dim); + B::float_flip(grad_cumsum, &[dim]) + }); + } + } + + match CumSum + .prepare::([tensor.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + (tensor.primitive.shape(), dim), + B::float_cumsum(tensor.primitive, dim), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_cumsum(tensor.primitive, dim)), + } + } + + fn float_cumprod(tensor: FloatTensor, dim: usize) -> FloatTensor { + #[derive(Debug)] + struct CumProd; + + impl Backward for CumProd { + type State = (B::FloatTensorPrimitive, usize); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (input, dim) = ops.state; + let output = B::float_cumprod(input.clone(), dim); + + unary::(ops.parents, ops.node, grads, |grad| { + // Gradient of cumprod using negative step slicing + // Formula: grad_input[i] = sum_{j>=i}(grad_output[j] * output[j] / input[i]) + // = (1 / input[i]) * sum_{j>=i}(grad_output[j] * output[j]) + // = (1 / input) * reverse_cumsum(grad * output) + // + // LIMITATION: This produces NaN when input contains zeros. + // A proper zero-safe implementation requires more sophisticated algorithms + // (see PyTorch's cumprod_backward or JAX's associative_scan approach). + // TODO: Implement zero-safe gradient computation. + // See: https://github.com/tracel-ai/burn/issues/3864 + + let grad_times_output = B::float_mul(grad, output.clone()); + + // Create slices to reverse along the specified dimension + let shape = grad_times_output.shape(); + let mut slices = vec![Slice::full(); shape.num_dims()]; + slices[dim] = Slice::with_step(0, None, -1); + + // Reverse, cumsum, reverse back using negative step slicing + let grad_reversed = B::float_slice(grad_times_output, &slices); + let grad_cumsum = B::float_cumsum(grad_reversed, dim); + let grad_result = B::float_slice(grad_cumsum, &slices); + + B::float_div(grad_result, input) + }); + } + } + + match CumProd + .prepare::([tensor.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + (tensor.primitive.clone(), dim), + B::float_cumprod(tensor.primitive, dim), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_cumprod(tensor.primitive, dim)), + } + } + + fn float_cummin(tensor: FloatTensor, dim: usize) -> FloatTensor { + #[derive(Debug)] + struct CumMin; + + impl Backward for CumMin { + type State = (B::FloatTensorPrimitive, usize); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (input, dim) = ops.state; + let output = B::float_cummin(input.clone(), dim); + + unary::(ops.parents, ops.node, grads, |grad| { + // Gradient flows to the input positions that produced each output + // Use scatter to accumulate gradients (scatter does sum reduction) + + let shape = input.shape(); + let device = input.device(); + let settings = get_device_settings::(&device); + let dim_size = shape[dim] as i64; + + // Create indices [0, 1, 2, ...] along the dimension + let arange_1d = B::int_arange(0..dim_size, &device, settings.int_dtype); + + // Reshape to broadcast along the specified dimension + let mut arange_shape = vec![1; shape.num_dims()]; + arange_shape[dim] = dim_size as usize; + let arange = B::int_reshape(arange_1d, Shape::from(arange_shape)); + + // Expand to match input shape + let arange = B::int_expand(arange, shape.clone()); + + // Find where cummin[i] == input[i] (these are source positions) + let is_source = + B::float_equal(output.clone(), input.clone(), settings.bool_dtype); + let is_source_int = B::bool_into_int(is_source, settings.int_dtype); + + // Mask: where is_source, use index; else 0 + let masked_indices = B::int_mul(arange, is_source_int); + + // Cummax propagates the last valid (non-zero) index forward + let source_indices = B::int_cummax(masked_indices, dim); + + // Scatter gradients to source positions (sum reduction) + let zeros = B::float_zeros(shape, &device, grad.dtype().into()); + B::float_scatter_add(dim, zeros, source_indices, grad) + }); + } + } + + match CumMin + .prepare::([tensor.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + (tensor.primitive.clone(), dim), + B::float_cummin(tensor.primitive, dim), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_cummin(tensor.primitive, dim)), + } + } + + fn float_cummax(tensor: FloatTensor, dim: usize) -> FloatTensor { + #[derive(Debug)] + struct CumMax; + + impl Backward for CumMax { + type State = (B::FloatTensorPrimitive, usize); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (input, dim) = ops.state; + let output = B::float_cummax(input.clone(), dim); + + unary::(ops.parents, ops.node, grads, |grad| { + // Gradient flows to the input positions that produced each output + // Use scatter to accumulate gradients (scatter does sum reduction) + + let shape = input.shape(); + let device = input.device(); + let settings = get_device_settings::(&device); + let dim_size = shape[dim] as i64; + + // Create indices [0, 1, 2, ...] along the dimension + let arange_1d = B::int_arange(0..dim_size, &device, settings.int_dtype); + + // Reshape to broadcast along the specified dimension + let mut arange_shape = vec![1; shape.num_dims()]; + arange_shape[dim] = dim_size as usize; + let arange = B::int_reshape(arange_1d, Shape::from(arange_shape)); + + // Expand to match input shape + let arange = B::int_expand(arange, shape.clone()); + + // Find where cummax[i] == input[i] (these are source positions) + let is_source = + B::float_equal(output.clone(), input.clone(), settings.bool_dtype); + let is_source_int = B::bool_into_int(is_source, settings.int_dtype); + + // Mask: where is_source, use index; else 0 + let masked_indices = B::int_mul(arange, is_source_int); + + // Cummax propagates the last valid (non-zero) index forward + let source_indices = B::int_cummax(masked_indices, dim); + + // Scatter gradients to source positions (sum reduction) + let zeros = B::float_zeros(shape, &device, grad.dtype().into()); + B::float_scatter_add(dim, zeros, source_indices, grad) + }); + } + } + + match CumMax + .prepare::([tensor.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + (tensor.primitive.clone(), dim), + B::float_cummax(tensor.primitive, dim), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_cummax(tensor.primitive, dim)), + } + } + + fn float_argmax(tensor: FloatTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + B::float_argmax(tensor.primitive, dim, out_dtype) + } + + fn float_argtopk( + tensor: FloatTensor, + dim: usize, + k: usize, + out_dtype: IntDType, + ) -> IntTensor { + B::float_argtopk(tensor.primitive, dim, k, out_dtype) + } + + fn float_topk(_tensor: FloatTensor, _dim: usize, _k: usize) -> FloatTensor { + unimplemented!("topk is not implemented for autodiff"); + } + + fn float_argmin(tensor: FloatTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + B::float_argmin(tensor.primitive, dim, out_dtype) + } + + fn float_exp(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Exp; + + retro_unary!(RetroExp, B::float_exp); + + impl Backward for Exp { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + let output = B::float_exp(input); + unary::(ops.parents, ops.node, grads, |grad| { + B::float_mul(grad, output) + }); + } + } + + match Exp + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroExp::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_exp(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_exp(tensor.primitive)), + } + } + + fn float_log(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Log; + + retro_unary!(RetroLog, B::float_log); + + impl Backward for Log { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + unary::(ops.parents, ops.node, grads, |grad| { + let value = B::float_recip(input); + B::float_mul(grad, value) + }); + } + } + + match Log + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroLog::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_log(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_log(tensor.primitive)), + } + } + + fn float_log1p(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Log1P; + + retro_unary!(RetroLog1P, B::float_log1p); + + impl Backward for Log1P { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + unary::(ops.parents, ops.node, grads, |grad| { + let value = B::float_add_scalar(input, 1f32.into()); + let value = B::float_recip(value); + + B::float_mul(grad, value) + }); + } + } + + match Log1P + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroLog1P::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_log1p(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_log1p(tensor.primitive)), + } + } + + fn float_powf_scalar_impl(tensor: FloatTensor, value: Scalar) -> FloatTensor { + #[derive(Debug)] + struct PowfScalar; + + #[derive(new, Debug)] + struct RetroPowfScalar { + lhs_id: NodeId, + rhs: f64, + _backend: PhantomData, + } + + impl RetroForward for RetroPowfScalar { + fn forward(&self, states: &mut BackwardStates, out_node: NodeId) { + let lhs = states.get_state::(&self.lhs_id); + let out = B::float_powf_scalar(lhs, self.rhs.into()); + states.save(out_node, out) + } + } + + impl Backward for PowfScalar { + type State = (NodeId, f64); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let (tensor_id, value) = ops.state; + let tensor = checkpointer.retrieve_node_output(tensor_id); + + unary::(ops.parents, ops.node, grads, |grad| { + let tmp = B::float_powf_scalar(tensor, (value - 1.).into()); + let value = B::float_mul_scalar(tmp, value.into()); + + B::float_mul(grad, value) + }); + } + } + + match PowfScalar + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroPowfScalar::::new(tensor.node.id, value.elem())) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = (prep.checkpoint(&tensor), value.elem()); + prep.finish(state, B::float_powf_scalar(tensor.primitive, value)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_powf_scalar(tensor.primitive, value)), + } + } + + fn float_sqrt(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Sqrt; + + retro_unary!(RetroSqrt, B::float_sqrt); + + impl Backward for Sqrt { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + unary::(ops.parents, ops.node, grads, |grad| { + let value = B::float_div_scalar( + B::float_powf_scalar(input, (-0.5).into()), + 2f32.into(), + ); + + B::float_mul(grad, value) + }); + } + } + + match Sqrt + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroSqrt::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_sqrt(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_sqrt(tensor.primitive)), + } + } + + fn float_abs(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Abs; + + retro_unary!(RetroAbs, B::float_abs); + + impl Backward for Abs { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let tensor: B::FloatTensorPrimitive = checkpointer.retrieve_node_output(ops.state); + let state = B::float_sign(tensor); + unary::(ops.parents, ops.node, grads, |grad| { + B::float_mul(grad, state) + }); + } + } + + match Abs + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroAbs::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_abs(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_abs(tensor.primitive)), + } + } + + fn float_cos(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Cos; + + retro_unary!(RetroCos, B::float_cos); + + impl Backward for Cos { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + unary::(ops.parents, ops.node, grads, |grad| { + let value = B::float_neg(B::float_sin(input)); + + B::float_mul(grad, value) + }); + } + } + + match Cos + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroCos::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_cos(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_cos(tensor.primitive)), + } + } + + fn float_sin(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Sin; + + retro_unary!(RetroSin, B::float_sin); + + impl Backward for Sin { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let state = checkpointer.retrieve_node_output(ops.state); + unary::(ops.parents, ops.node, grads, |grad| { + let value = B::float_cos(state); + B::float_mul(grad, value) + }); + } + } + + match Sin + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroSin::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_sin(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_sin(tensor.primitive)), + } + } + + fn float_tanh(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Tanh; + + retro_unary!(RetroTanh, B::float_tanh); + + impl Backward for Tanh { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + let state = B::float_tanh(input); + unary::(ops.parents, ops.node, grads, |grad| { + let value = B::float_add_scalar( + B::float_neg(B::float_powi_scalar(state, 2.into())), + 1f32.into(), + ); + B::float_mul(grad, value) + }); + } + } + + match Tanh + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroTanh::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_tanh(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_tanh(tensor.primitive)), + } + } + + fn float_cosh(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Cosh; + + retro_unary!(RetroCosh, B::float_cosh); + + impl Backward for Cosh { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + unary::(ops.parents, ops.node, grads, |grad| { + B::float_mul(grad, B::float_sinh(input)) + }); + } + } + + match Cosh + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroCosh::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_cosh(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_cosh(tensor.primitive)), + } + } + + fn float_sinh(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Sinh; + + retro_unary!(RetroSinh, B::float_sinh); + + impl Backward for Sinh { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + unary::(ops.parents, ops.node, grads, |grad| { + B::float_mul(grad, B::float_cosh(input)) + }); + } + } + + match Sinh + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroSinh::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_sinh(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_sinh(tensor.primitive)), + } + } + + fn float_tan(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Tan; + + retro_unary!(RetroTan, B::float_tan); + + impl Backward for Tan { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + let tan_x = B::float_tan(input); + unary::(ops.parents, ops.node, grads, |grad| { + // d/dx tan(x) = 1 + tan^2(x) + let tan_sq = B::float_powi_scalar(tan_x, 2.into()); + B::float_mul(grad, B::float_add_scalar(tan_sq, 1f32.into())) + }); + } + } + + match Tan + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroTan::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_tan(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_tan(tensor.primitive)), + } + } + + fn float_asin(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Asin; + + retro_unary!(RetroAsin, B::float_asin); + + impl Backward for Asin { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + unary::(ops.parents, ops.node, grads, |grad| { + // d/dx asin(x) = 1/sqrt(1 - x^2) + let x_sq = B::float_powi_scalar(input, 2.into()); + let denom = B::float_sqrt(B::float_add_scalar(B::float_neg(x_sq), 1f32.into())); + B::float_mul(grad, B::float_recip(denom)) + }); + } + } + + match Asin + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroAsin::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_asin(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_asin(tensor.primitive)), + } + } + + fn float_acos(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Acos; + + retro_unary!(RetroAcos, B::float_acos); + + impl Backward for Acos { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + unary::(ops.parents, ops.node, grads, |grad| { + // d/dx acos(x) = -1/sqrt(1 - x^2) + let x_sq = B::float_powi_scalar(input, 2.into()); + let denom = B::float_sqrt(B::float_add_scalar(B::float_neg(x_sq), 1f32.into())); + let value = B::float_neg(B::float_recip(denom)); + B::float_mul(grad, value) + }); + } + } + + match Acos + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroAcos::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_acos(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_acos(tensor.primitive)), + } + } + + fn float_atan(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Atan; + + retro_unary!(RetroAtan, B::float_atan); + + impl Backward for Atan { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + unary::(ops.parents, ops.node, grads, |grad| { + // d/dx atan(x) = 1/(1 + x^2) + let x_sq = B::float_powi_scalar(input, 2.into()); + let value = B::float_recip(B::float_add_scalar(x_sq, 1f32.into())); + B::float_mul(grad, value) + }); + } + } + + match Atan + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroAtan::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_atan(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_atan(tensor.primitive)), + } + } + + fn float_asinh(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Asinh; + + retro_unary!(RetroAsinh, B::float_asinh); + + impl Backward for Asinh { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + unary::(ops.parents, ops.node, grads, |grad| { + // d/dx asinh(x) = 1/sqrt(x^2 + 1) + let x_sq = B::float_powi_scalar(input, 2.into()); + let value = + B::float_recip(B::float_sqrt(B::float_add_scalar(x_sq, 1f32.into()))); + B::float_mul(grad, value) + }); + } + } + + match Asinh + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroAsinh::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_asinh(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_asinh(tensor.primitive)), + } + } + + fn float_acosh(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Acosh; + + retro_unary!(RetroAcosh, B::float_acosh); + + impl Backward for Acosh { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + unary::(ops.parents, ops.node, grads, |grad| { + // d/dx acosh(x) = 1/sqrt(x^2 - 1) + let x_sq = B::float_powi_scalar(input, 2.into()); + let value = + B::float_recip(B::float_sqrt(B::float_sub_scalar(x_sq, 1f32.into()))); + B::float_mul(grad, value) + }); + } + } + + match Acosh + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroAcosh::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_acosh(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_acosh(tensor.primitive)), + } + } + + fn float_atanh(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Atanh; + + retro_unary!(RetroAtanh, B::float_atanh); + + impl Backward for Atanh { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let input = checkpointer.retrieve_node_output(ops.state); + unary::(ops.parents, ops.node, grads, |grad| { + // d/dx atanh(x) = 1/(1 - x^2) + let x_sq = B::float_powi_scalar(input, 2.into()); + let value = + B::float_recip(B::float_add_scalar(B::float_neg(x_sq), 1f32.into())); + B::float_mul(grad, value) + }); + } + } + + match Atanh + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroAtanh::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_atanh(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_atanh(tensor.primitive)), + } + } + + fn float_atan2(y: FloatTensor, x: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Atan2; + + retro_binary!(RetroAtan2, B::float_atan2); + + impl Backward for Atan2 { + type State = (Option, Option, BinaryOpsBroadcast); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let (y_id, x_id, broadcast) = ops.state; + let y = y_id.map(|id| checkpointer.retrieve_node_output(id)); + let x = x_id.map(|id| checkpointer.retrieve_node_output(id)); + let [y_4y, y_4x] = duplicate(&ops.parents, y); + let [x_4y, x_4x]: [Option>; 2] = duplicate(&ops.parents, x); + + binary::( + ops.parents, + ops.node, + grads, + |grad| { + // d/dy atan2(y, x) = x/(x^2 + y^2) + let y = y_4y.unwrap(); + let x = x_4y.unwrap(); + let x_sq = B::float_powi_scalar(x.clone(), 2.into()); + let y_sq = B::float_powi_scalar(y, 2.into()); + let denom = B::float_add(x_sq, y_sq); + let value = B::float_div(x, denom); + let grad = B::float_mul(grad, value); + + broadcast.backward_lhs::(grad) + }, + |grad| { + // d/dx atan2(y, x) = -y/(x^2 + y^2) + let y = y_4x.unwrap(); + let x = x_4x.unwrap(); + let x_sq = B::float_powi_scalar(x, 2.into()); + let y_sq = B::float_powi_scalar(y.clone(), 2.into()); + let denom = B::float_add(x_sq, y_sq); + let value = B::float_neg(B::float_div(y, denom)); + let grad = B::float_mul(grad, value); + + broadcast.backward_rhs::(grad) + }, + ); + } + } + + let y_tracked = y.is_tracked(); + let x_tracked = x.is_tracked(); + let broadcast = BinaryOpsBroadcast::new::(&y.primitive, &x.primitive); + + match Atan2 + .prepare::([y.node.clone(), x.node.clone()]) + .memory_bound() + .retro_forward(RetroAtan2::::new(y.node.id, x.node.id)) + .parents([&y, &x]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let is_tracked = y_tracked || x_tracked; + let y_state = is_tracked.then(|| prep.checkpoint(&y)); + let x_state = is_tracked.then(|| prep.checkpoint(&x)); + + prep.finish( + (y_state, x_state, broadcast), + B::float_atan2(y.primitive, x.primitive), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_atan2(y.primitive, x.primitive)), + } + } + + fn float_round(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Round; + retro_unary!(RetroRound, B::float_round); + + impl Backward for Round { + type State = (Shape, B::Device); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (shape, device) = ops.state; + unary::(ops.parents, ops.node, grads, |grad| { + B::float_zeros(shape, &device, grad.dtype().into()) + }) + } + } + + match Round + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroRound::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(preps) => preps.finish( + (tensor.primitive.shape(), tensor.primitive.device()), + B::float_round(tensor.primitive), + ), + OpsKind::UnTracked(preps) => preps.finish(B::float_round(tensor.primitive)), + } + } + + fn float_floor(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Floor; + retro_unary!(RetroFloor, B::float_floor); + + impl Backward for Floor { + type State = (Shape, B::Device); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (shape, device) = ops.state; + unary::(ops.parents, ops.node, grads, |grad| { + B::float_zeros(shape, &device, grad.dtype().into()) + }) + } + } + + match Floor + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroFloor::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(preps) => preps.finish( + (tensor.primitive.shape(), tensor.primitive.device()), + B::float_floor(tensor.primitive), + ), + OpsKind::UnTracked(preps) => preps.finish(B::float_floor(tensor.primitive)), + } + } + + fn float_ceil(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Ceil; + retro_unary!(RetroCeil, B::float_ceil); + + impl Backward for Ceil { + type State = (Shape, B::Device); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (shape, device) = ops.state; + unary::(ops.parents, ops.node, grads, |grad| { + B::float_zeros(shape, &device, grad.dtype().into()) + }) + } + } + + match Ceil + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroCeil::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(preps) => preps.finish( + (tensor.primitive.shape(), tensor.primitive.device()), + B::float_ceil(tensor.primitive), + ), + OpsKind::UnTracked(preps) => preps.finish(B::float_ceil(tensor.primitive)), + } + } + + fn float_trunc(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Trunc; + retro_unary!(RetroTrunc, B::float_trunc); + + impl Backward for Trunc { + type State = (Shape, B::Device); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (shape, device) = ops.state; + unary::(ops.parents, ops.node, grads, |grad| { + B::float_zeros(shape, &device, grad.dtype().into()) + }) + } + } + + match Trunc + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroTrunc::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(preps) => preps.finish( + (tensor.primitive.shape(), tensor.primitive.device()), + B::float_trunc(tensor.primitive), + ), + OpsKind::UnTracked(preps) => preps.finish(B::float_trunc(tensor.primitive)), + } + } + + fn float_erf(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Erf; + + retro_unary!(RetroErf, B::float_erf); + + impl Backward for Erf { + type State = NodeId; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + unary::(ops.parents, ops.node, grads, |grad| { + let ops = checkpointer.retrieve_node_output(ops.state); + let exponent = B::float_neg(B::float_powi_scalar(ops, 2.into())); + let numerator = B::float_mul_scalar(B::float_exp(exponent), 2.0.into()); + let denominator = core::f64::consts::PI.sqrt().into(); + let value = B::float_div_scalar(numerator, denominator); + + B::float_mul(grad, value) + }); + } + } + + match Erf + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroErf::::new(tensor.node.id)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let state = prep.checkpoint(&tensor); + prep.finish(state, B::float_erf(tensor.primitive)) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_erf(tensor.primitive)), + } + } + + fn float_cat(tensors: Vec>, dim: usize) -> FloatTensor { + #[derive(new, Debug)] + struct CatStep { + nodes: Vec>, + // The dimension of each tensor along the dim dimension. + // This indicates the number of dimension concatenated for each tensor. + dim_sizes: Vec, + output: NodeRef, + phantom: PhantomData, + dim: usize, + parents: Vec, + } + + impl Step for CatStep { + fn step(self: Box, grads: &mut Gradients, _checkpointer: &mut Checkpointer) { + let grad = grads.consume::(&self.output); + let ranges_template: Vec<_> = grad.shape().iter().map(|&v| 0..v).collect(); + + self.nodes + .into_iter() + .zip(self.dim_sizes) + .scan(0, |offset, (node_opt, dim_size)| { + let start = *offset; + let end = start + dim_size; + *offset = end; + Some(node_opt.map(|node| (node, start, end))) + }) + .flatten() + .for_each(|(node, start, end)| { + let mut ranges = ranges_template.clone(); + ranges[self.dim] = start..end; + + let slices: Vec = ranges + .iter() + .map(|r| Slice::new(r.start as isize, Some(r.end as isize), 1)) + .collect(); + grads.register::(node.id, B::float_slice(grad.clone(), &slices)); + }); + } + + fn node(&self) -> NodeId { + self.output.id + } + + fn parents(&self) -> &[Parent] { + &self.parents + } + fn depth(&self) -> usize { + self.output.order + } + + fn distributed_params(&self) -> Option { + self.output.distributed_params.clone() + } + } + + let mut nodes = Vec::with_capacity(tensors.len()); + let mut primitives = Vec::with_capacity(tensors.len()); + let mut dim_sizes = Vec::with_capacity(tensors.len()); + + tensors.into_iter().for_each(|tensor| { + dim_sizes.push(tensor.primitive.shape()[dim]); + nodes.push(tensor.node); + primitives.push(tensor.primitive); + }); + + let requirement = Requirement::from_nodes(&nodes); + + // For simplicity, this operation does not checkpoint anything + let cat_computing_property = ComputingProperty::Ambiguous; + let checkpointer_builder = CheckpointerBuilder::default(); + + let output = B::float_cat(primitives, dim); + if requirement.is_none() { + return AutodiffTensor::from_parents( + output, + &nodes, + requirement, + cat_computing_property, + ); + } + + let output = + AutodiffTensor::from_parents(output, &nodes, requirement, cat_computing_property); + + let mut parents = Vec::new(); + + let nodes = nodes + .into_iter() + .map(|node| node.clone_if_require_grad()) + .collect::>(); + for node in nodes.iter().flatten() { + parents.push(Parent { id: node.id }); + } + let ops = CatStep::::new(nodes, dim_sizes, output.node.clone(), dim, parents); + output.register_step(ops, checkpointer_builder) + } + + fn float_max_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + match MaxMinDim + .prepare::([tensor.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => { + let shape = tensor.primitive.shape(); + let settings = get_device_settings::(&tensor.primitive.device()); + let (tensor, index) = + B::float_max_dim_with_indices(tensor.primitive, dim, settings.int_dtype); + prep.finish((index, shape, dim), tensor) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_max_dim(tensor.primitive, dim)), + } + } + fn float_max_dim_with_indices( + tensor: FloatTensor, + dim: usize, + indices_dtype: IntDType, + ) -> (FloatTensor, IntTensor) { + match MaxMinDim + .prepare::([tensor.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => { + let shape = tensor.primitive.shape(); + let (tensor, index) = + B::float_max_dim_with_indices(tensor.primitive, dim, indices_dtype); + let tensor = prep.finish((index.clone(), shape, dim), tensor); + + (tensor, index) + } + OpsKind::UnTracked(prep) => { + let (tensor, index) = + B::float_max_dim_with_indices(tensor.primitive, dim, indices_dtype); + let tensor = prep.finish(tensor); + + (tensor, index) + } + } + } + + fn float_min_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + match MaxMinDim + .prepare::([tensor.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => { + let shape = tensor.primitive.shape(); + let settings = get_device_settings::(&tensor.primitive.device()); + let (tensor, index) = + B::float_min_dim_with_indices(tensor.primitive, dim, settings.int_dtype); + prep.finish((index, shape, dim), tensor) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_min_dim(tensor.primitive, dim)), + } + } + fn float_min_dim_with_indices( + tensor: FloatTensor, + dim: usize, + indices_dtype: IntDType, + ) -> (FloatTensor, IntTensor) { + match MaxMinDim + .prepare::([tensor.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => { + let shape = tensor.primitive.shape(); + let (tensor, index) = + B::float_min_dim_with_indices(tensor.primitive, dim, indices_dtype); + let tensor = prep.finish((index.clone(), shape, dim), tensor); + + (tensor, index) + } + OpsKind::UnTracked(prep) => { + let (tensor, index) = + B::float_min_dim_with_indices(tensor.primitive, dim, indices_dtype); + let tensor = prep.finish(tensor); + + (tensor, index) + } + } + } + + fn float_into_int(tensor: FloatTensor, out_dtype: IntDType) -> IntTensor { + B::float_into_int(tensor.primitive, out_dtype) + } + + fn float_powf(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct PowF; + + retro_binary!(RetroPowf, B::float_powf); + + impl Backward for PowF { + type State = (NodeId, NodeId, BinaryOpsBroadcast); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let (lhs_id, rhs_id, broadcast) = ops.state; + let lhs: B::FloatTensorPrimitive = checkpointer.retrieve_node_output(lhs_id); + let rhs: B::FloatTensorPrimitive = checkpointer.retrieve_node_output(rhs_id); + + // Both lhs and rhs are needed for both lhs and rhs gradients, but we clone them + // the number of times required by the parents specification. + let [rhs_4lhs, rhs_4rhs] = duplicate(&ops.parents, Some(rhs)); + let [lhs_4lhs, lhs_4rhs] = duplicate(&ops.parents, Some(lhs)); + + binary::( + ops.parents, + ops.node, + grads, + |grad| { + //rhs*(lhs.val**(rhs-1))*grad + let rhs1 = rhs_4lhs.unwrap(); + let rhs2 = rhs1.clone(); + let lhs = lhs_4lhs.unwrap(); + + let tmp = B::float_powf(lhs, B::float_sub_scalar(rhs1, 1.0.into())); + let value = B::float_mul(tmp, rhs2); + let grad = B::float_mul(grad, value); + + broadcast.backward_lhs::(grad) + }, + |grad| { + //lhs**rhs * ln(lhs) * grad + let rhs = rhs_4rhs.unwrap(); + let lhs1 = lhs_4rhs.unwrap(); + let lhs2 = lhs1.clone(); + let tmp = B::float_powf(lhs1, rhs); + let value = B::float_mul(tmp, B::float_log(lhs2)); + let grad = B::float_mul(grad, value); + + broadcast.backward_rhs::(grad) + }, + ); + } + } + + let broadcast = BinaryOpsBroadcast::new::(&lhs.primitive, &rhs.primitive); + + match PowF + .prepare::([lhs.node.clone(), rhs.node.clone()]) + .memory_bound() + .retro_forward(RetroPowf::::new(lhs.node.id, rhs.node.id)) + .parents([&lhs, &rhs]) + .stateful() + { + OpsKind::Tracked(mut prep) => { + let lhs_state = prep.checkpoint(&lhs); + let rhs_state = prep.checkpoint(&rhs); + prep.finish( + (lhs_state, rhs_state, broadcast), + B::float_powf(lhs.primitive, rhs.primitive), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_powf(lhs.primitive, rhs.primitive)), + } + } + + fn float_hypot(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Hypot; + + impl Backward for Hypot { + type State = (NodeId, NodeId, BinaryOpsBroadcast); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + checkpointer: &mut Checkpointer, + ) { + let (lhs_id, rhs_id, broadcast) = ops.state; + let lhs: B::FloatTensorPrimitive = checkpointer.retrieve_node_output(lhs_id); + let rhs: B::FloatTensorPrimitive = checkpointer.retrieve_node_output(rhs_id); + + let [lhs_4lhs, lhs_4rhs] = duplicate(&ops.parents, Some(lhs.clone())); + let [rhs_4lhs, rhs_4rhs] = duplicate(&ops.parents, Some(rhs.clone())); + + binary::( + ops.parents, + ops.node, + grads, + |grad| { + // lhs / hypot(lhs, rhs) * grad + let value = + B::float_div(lhs, B::float_hypot(lhs_4lhs.unwrap(), rhs_4lhs.unwrap())); + broadcast.backward_lhs::(B::float_mul(grad, value)) + }, + |grad| { + // rhs / hypot(lhs, rhs) * grad + let value = + B::float_div(rhs, B::float_hypot(lhs_4rhs.unwrap(), rhs_4rhs.unwrap())); + broadcast.backward_rhs::(B::float_mul(grad, value)) + }, + ); + } + } + let broadcast = BinaryOpsBroadcast::new::(&lhs.primitive, &rhs.primitive); + match Hypot + .prepare::([lhs.node.clone(), rhs.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(mut prep) => { + let lhs_state = prep.checkpoint(&lhs); + let rhs_state = prep.checkpoint(&rhs); + prep.finish( + (lhs_state, rhs_state, broadcast), + B::float_hypot(lhs.primitive, rhs.primitive), + ) + } + OpsKind::UnTracked(prep) => prep.finish(B::float_hypot(lhs.primitive, rhs.primitive)), + } + } + + fn float_sign(tensor: FloatTensor) -> FloatTensor { + #[derive(Debug)] + struct Sign; + + retro_unary!(RetroSign, B::float_sign); + + impl Backward for Sign { + type State = (); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + unary::(ops.parents, ops.node, grads, |grad| + // Always return 0 because the derivative of the sign function + // does not contribute to gradient updates in a meaningful way. + B::float_mul_scalar(grad, 0f32.into())); + } + } + + Sign.prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroSign::::new(tensor.node.id)) + .parents([&tensor]) + .stateless(B::float_sign(tensor.primitive)) + } + + fn float_expand(tensor: FloatTensor, shape: Shape) -> FloatTensor { + // D1: tensor, D2: shape + #[derive(Debug)] + struct ExpandDim; + + #[derive(new, Debug)] + struct RetroExpand { + input_id: NodeId, + shape: Shape, + _backend: PhantomData, + } + + impl RetroForward for RetroExpand { + fn forward(&self, states: &mut BackwardStates, out_node: NodeId) { + let input = states.get_state::(&self.input_id); + let out = B::float_expand(input, self.shape.clone()); + states.save(out_node, out) + } + } + + impl Backward for ExpandDim { + type State = (Shape, Shape); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (shape_in, shape_out) = ops.state; + let ndims_in = shape_in.num_dims(); + let ndims_out = shape_out.num_dims(); + + let mut shape_expanded = vec![1; ndims_out]; + + debug_assert!(ndims_out >= ndims_in); + + for i in 0..ndims_in { + shape_expanded[i + (ndims_out - ndims_in)] = shape_in[i]; + } + + unary::(ops.parents, ops.node, grads, |grad| { + let shape_grad = grad.shape(); + let mut grad = grad; + + #[allow(clippy::needless_range_loop)] + for i in 0..ndims_out { + if shape_expanded[i] == 1 && shape_grad[i] != 1 { + grad = B::float_sum_dim(grad, i); + } + } + + B::float_reshape(grad, shape_in) + }); + } + } + + match ExpandDim + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroExpand::::new(tensor.node.id, shape.clone())) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + (tensor.primitive.shape(), shape.clone()), + B::float_expand(tensor.primitive, shape), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_expand(tensor.primitive, shape)), + } + } + + fn float_sort(tensor: FloatTensor, dim: usize, descending: bool) -> FloatTensor { + match super::sort::SortDim + .prepare::([tensor.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => { + let shape = tensor.primitive.shape(); + let settings = get_device_settings::(&tensor.primitive.device()); + let (tensor, indices) = B::float_sort_with_indices( + tensor.primitive, + dim, + descending, + settings.int_dtype, + ); + prep.finish((indices, shape, dim), tensor) + } + OpsKind::UnTracked(prep) => { + prep.finish(B::float_sort(tensor.primitive, dim, descending)) + } + } + } + + fn float_sort_with_indices( + tensor: FloatTensor, + dim: usize, + descending: bool, + indices_dtype: IntDType, + ) -> (FloatTensor, IntTensor) { + match super::sort::SortDim + .prepare::([tensor.node]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => { + let shape = tensor.primitive.shape(); + let (tensor, indices) = + B::float_sort_with_indices(tensor.primitive, dim, descending, indices_dtype); + let tensor = prep.finish((indices.clone(), shape, dim), tensor); + + (tensor, indices) + } + OpsKind::UnTracked(prep) => { + let (tensor, indices) = + B::float_sort_with_indices(tensor.primitive, dim, descending, indices_dtype); + let tensor = prep.finish(tensor); + + (tensor, indices) + } + } + } + + fn float_argsort( + tensor: FloatTensor, + dim: usize, + descending: bool, + out_dtype: IntDType, + ) -> IntTensor { + B::float_argsort(tensor.primitive, dim, descending, out_dtype) + } + + fn float_repeat_dim(tensor: FloatTensor, dim: usize, times: usize) -> FloatTensor { + #[derive(Debug)] + struct Repeat; + + #[derive(new, Debug)] + struct RetroRepeat { + tensor_id: NodeId, + dim: usize, + times: usize, + _backend: PhantomData, + } + + impl RetroForward for RetroRepeat { + fn forward(&self, states: &mut BackwardStates, out_node: NodeId) { + let tensor = states.get_state::(&self.tensor_id); + let out = B::float_repeat_dim(tensor, self.dim, self.times); + states.save(out_node, out) + } + } + + impl Backward for Repeat { + type State = (usize, usize); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (dim, times) = ops.state; + + unary::(ops.parents, ops.node, grads, |grad| { + let mut dims = grad.shape(); + let orig_dim_size = dims[dim] / times; + if orig_dim_size > 1 { + dims[dim] = orig_dim_size; + let orig_dims = dims.clone(); + dims.insert(dim + 1, times); // shape [..., orig_dim_size, times, ...] + let grad = B::float_reshape(grad, dims); + let grad = B::float_sum_dim(grad, dim + 1); // sum over repeat times + B::float_reshape(grad, orig_dims) + } else { + B::float_sum_dim(grad, dim) + } + }); + } + } + + match Repeat + .prepare::([tensor.node.clone()]) + .memory_bound() + .retro_forward(RetroRepeat::::new(tensor.node.id, dim, times)) + .parents([&tensor]) + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + (dim, times), + B::float_repeat_dim(tensor.primitive, dim, times), + ), + OpsKind::UnTracked(prep) => { + prep.finish(B::float_repeat_dim(tensor.primitive, dim, times)) + } + } + } + + fn float_cast(tensor: FloatTensor, dtype: burn_std::FloatDType) -> FloatTensor { + #[derive(Debug)] + struct Cast; + + impl Backward for Cast { + type State = FloatDType; + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let dtype = ops.state; + + unary::(ops.parents, ops.node, grads, |grad| { + B::float_cast(grad, dtype) + }); + } + } + + match Cast + .prepare::([tensor.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + tensor.dtype().into(), + B::float_cast(tensor.primitive, dtype), + ), + OpsKind::UnTracked(prep) => prep.finish(B::float_cast(tensor.primitive, dtype)), + } + } + + fn float_unfold( + tensor: FloatTensor, + dim: usize, + size: usize, + step: usize, + ) -> FloatTensor { + #[derive(Debug)] + struct Unfold; + + impl Backward for Unfold { + type State = (Shape, usize, usize, usize); + + fn backward( + self, + ops: Ops, + grads: &mut Gradients, + _checkpointer: &mut Checkpointer, + ) { + let (shape_in, dim, size, step) = ops.state; + let windows = calculate_unfold_windows(shape_in[dim], size, step); + + unary::(ops.parents, ops.node, grads, |grad| { + let device = grad.device(); + let mut grad_input = + B::float_zeros(shape_in.clone(), &device, grad.dtype().into()); + + if windows == 0 { + return grad_input; + } + + let ndims_in = shape_in.num_dims(); + let ndims_out = grad.shape().num_dims(); + + let mut target_shape = shape_in.clone(); + target_shape[dim] = size; + + for window_idx in 0..windows { + let mut slices_out = vec![Slice::new(0, None, 1); ndims_out]; + let start = window_idx * step; + let end = start + size; + slices_out[dim] = + Slice::new(window_idx as isize, Some((window_idx + 1) as isize), 1); + + let window_grad = B::float_slice(grad.clone(), &slices_out); + + let last_axis = ndims_out - 1; + let mut permutation: Vec = (0..dim).collect(); + permutation.push(last_axis); + permutation.extend(dim + 1..last_axis); + permutation.push(dim); + + let window_grad = B::float_permute(window_grad, &permutation); + let window_grad = B::float_reshape(window_grad, target_shape.clone()); + + let mut slices_in = vec![Slice::new(0, None, 1); ndims_in]; + slices_in[dim] = Slice::new(start as isize, Some(end as isize), 1); + + let current = B::float_slice(grad_input.clone(), &slices_in); + let updated = B::float_add(current, window_grad); + grad_input = B::float_slice_assign(grad_input, &slices_in, updated); + } + + grad_input + }); + } + } + + match Unfold + .prepare::([tensor.node.clone()]) + .compute_bound() + .stateful() + { + OpsKind::Tracked(prep) => prep.finish( + (tensor.primitive.shape(), dim, size, step), + B::float_unfold(tensor.primitive, dim, size, step), + ), + OpsKind::UnTracked(prep) => { + prep.finish(B::float_unfold(tensor.primitive, dim, size, step)) + } + } + } +} + +#[derive(Debug, Clone)] +enum BinaryOpsBroadcast { + Broadcasted(Shape, Shape), + None, +} + +impl BinaryOpsBroadcast { + fn new(lhs: &B::FloatTensorPrimitive, rhs: &B::FloatTensorPrimitive) -> Self { + let shape_lhs = lhs.shape(); + let shape_rhs = rhs.shape(); + let ndims = shape_lhs.num_dims(); + + for i in 0..ndims { + if shape_rhs[i] != shape_lhs[i] { + return Self::Broadcasted(shape_lhs, shape_rhs); + } + } + + Self::None + } + + fn backward_lhs(&self, grad: B::FloatTensorPrimitive) -> B::FloatTensorPrimitive { + match self { + BinaryOpsBroadcast::Broadcasted(lhs, _rhs) => broadcast_shape::(grad, lhs), + BinaryOpsBroadcast::None => grad, + } + } + + fn backward_rhs(&self, grad: B::FloatTensorPrimitive) -> B::FloatTensorPrimitive { + match self { + BinaryOpsBroadcast::Broadcasted(_lhs, rhs) => broadcast_shape::(grad, rhs), + BinaryOpsBroadcast::None => grad, + } + } +} diff --git a/crates/burn-autodiff/src/ops/transaction.rs b/crates/burn-autodiff/src/ops/transaction.rs new file mode 100644 index 0000000..17a07a4 --- /dev/null +++ b/crates/burn-autodiff/src/ops/transaction.rs @@ -0,0 +1,24 @@ +use burn_backend::{ + Backend, ExecutionError, + ops::{TransactionOps, TransactionPrimitive}, +}; + +use crate::{Autodiff, checkpoint::strategy::CheckpointStrategy}; + +impl TransactionOps for Autodiff { + async fn tr_execute( + transaction: TransactionPrimitive, + ) -> Result { + B::tr_execute(TransactionPrimitive::new( + transaction + .read_floats + .into_iter() + .map(|t| t.primitive) + .collect(), + transaction.read_qfloats, + transaction.read_ints, + transaction.read_bools, + )) + .await + } +} diff --git a/crates/burn-autodiff/src/runtime/client.rs b/crates/burn-autodiff/src/runtime/client.rs new file mode 100644 index 0000000..df9d88a --- /dev/null +++ b/crates/burn-autodiff/src/runtime/client.rs @@ -0,0 +1,18 @@ +use crate::{ + checkpoint::builder::CheckpointerBuilder, + grads::{BackwardMode, Gradients}, + graph::StepBoxed, + tensor::{AutodiffTensor, NodeRefCount}, +}; +use burn_backend::Backend; + +/// Client used to communicate with the autodiff server. +pub trait AutodiffClient: Send + Clone { + /// Register a new step. + fn register(&self, node_id: NodeRefCount, step: StepBoxed, actions: CheckpointerBuilder); + /// Call backpropagation from the given tensor. + fn backward(&self, tensor: AutodiffTensor, mode: BackwardMode) -> Gradients; +} + +/// Client implementation in used. +pub type AutodiffClientImpl = super::graph::GraphMutexClient; diff --git a/crates/burn-autodiff/src/runtime/graph.rs b/crates/burn-autodiff/src/runtime/graph.rs new file mode 100644 index 0000000..4b4a2f1 --- /dev/null +++ b/crates/burn-autodiff/src/runtime/graph.rs @@ -0,0 +1,344 @@ +use super::{AutodiffClient, server::AutodiffServer}; +use crate::{ + NodeId, + checkpoint::builder::CheckpointerBuilder, + grads::{BackwardMode, Gradients}, + graph::{Parent, StepBoxed}, + runtime::server::NodeCleaner, + tensor::{AutodiffTensor, NodeRefCount}, +}; +use alloc::vec::Vec; + +#[cfg(target_has_atomic = "ptr")] +use alloc::sync::Arc; + +#[cfg(not(target_has_atomic = "ptr"))] +use portable_atomic_util::Arc; + +use burn_backend::Backend; + +use alloc::collections::BTreeMap; +use hashbrown::HashSet; + +#[cfg(feature = "std")] +use parking_lot::{Mutex, MutexGuard}; + +#[cfg(not(feature = "std"))] +use spin::{Mutex, MutexGuard}; + +/// A client for managing multiple graphs using mutex-based synchronization. +/// +/// The biggest benefit of using this client implementation is that each graph can modify its own +/// data without blocking other graphs, which is essential for multi-device training. +/// +/// # Notes +/// +/// The [AutodiffServer] fully supports multiple graphs with sharing nodes, however those type of +/// graphs will be stored under a single mutex-protected graph by the client, limiting +/// parallelisation. +#[derive(Clone, new, Debug)] +pub struct GraphMutexClient; + +/// Manages a collection of graphs, mapping [node ids](NodeId) to their respective graph. +/// +/// The `GraphLocator` is responsible for selecting and merging graphs based on their IDs and parent +/// dependencies, ensuring proper synchronization and server allocation. +/// +/// # Notes +/// +/// Multiple node ids can point to the same graph, where the autodiff graph is stored. +#[derive(Default)] +pub struct GraphLocator { + graphs: BTreeMap>, + /// We keep a mapping of each original node id (graph id) => all nodes that point to that graph. + /// This is to ensure that when merging graphs, we correctly move all previous graphs to + /// the new merged one. + keys: BTreeMap>, +} + +/// Represents a single computation graph with a mutex-protected server. +/// +/// Each `Graph` contains an [AutodiffServer] and the original [NodeId] where the server was +/// first created. +pub(crate) struct Graph { + origin: NodeId, + state: Mutex, +} + +#[derive(Default)] +struct GraphState { + server: AutodiffServer, +} + +impl core::fmt::Debug for Graph { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("Graph") + .field("origin", &self.origin) + .finish() + } +} + +static STATE: Mutex> = Mutex::new(None); + +impl GraphMutexClient { + /// Retrieves or creates a graph for the given [NodeId] and parent dependencies. + /// + /// # Parameters + /// - `node`: The unique identifier for the stream. + /// - `parents`: A slice of parent nodes that the stream depends on. + /// + /// # Returns + /// An `Arc` representing the selected or newly created stream. + fn graph(node: NodeId, parents: &[Parent]) -> Arc { + let mut state = STATE.lock(); + + match state.as_mut() { + Some(locator) => locator.select(node, parents), + None => { + let mut locator = GraphLocator::default(); + let stream = locator.select(node, parents); + *state = Some(locator); + stream + } + } + } +} + +impl AutodiffClient for GraphMutexClient { + fn register(&self, node_id_ref: NodeRefCount, step: StepBoxed, actions: CheckpointerBuilder) { + let node_id = *node_id_ref; + let graph = GraphMutexClient::graph(node_id, step.parents()); + let mut state = graph.state.lock(); + + state.server.register(node_id_ref, step, actions); + } + + fn backward(&self, root: AutodiffTensor, mode: BackwardMode) -> Gradients { + let node_id = root.node.id; + let graph = GraphMutexClient::graph(root.node.id, &[]); + + let grads = { + let mut state = graph.state.lock(); + state + .server + .backward::(root.node, root.primitive, node_id, mode) + }; // lock released + + GraphCleaner::cleanup_orphaned_entries(); + + grads + } +} + +struct GraphCleaner<'a> { + guard: MutexGuard<'a, Option>, +} + +impl<'a> GraphCleaner<'a> { + fn cleanup_orphaned_entries() { + let graphs = { + // Get the available graphs and release the lock + match STATE.lock().as_ref() { + Some(state) => state.graphs.clone(), + None => return, + } + }; + + let mut should_remove = Vec::new(); + for graph in graphs.values() { + { + let mut guard = graph.state.lock(); + // Double safety: in case it was marked as no longer useful, but other + // nodes are still relevant, we only check which nodes can safely be removed. + if !guard.server.maybe_useful() { + guard + .server + .free_unused_roots(|node| should_remove.push(*node)); + } + } + } + + if !should_remove.is_empty() { + let mut state = STATE.lock(); + if let Some(state) = state.as_mut() { + for node in should_remove { + state.remove_entry(&node); + } + } + } + } +} + +impl<'a> NodeCleaner for GraphCleaner<'a> { + fn init() -> Self { + let guard = STATE.lock(); + Self { guard } + } + + fn clean(&mut self, node: &NodeId) { + if let Some(state) = self.guard.as_mut() { + state.remove_entry(node); + } + } +} + +impl GraphLocator { + /// Selects a single graph for the given [NodeId], considering parent dependencies. + /// + /// If multiple graphs are found, they are merged into a single one. + /// + /// # Parameters + /// - `node`: The node ID of the graph to select. + /// - `parents`: A slice of parent nodes that the graph depends on. + /// + /// # Returns + /// + /// An `Arc` representing the selected or merged graph. + pub(crate) fn select(&mut self, node: NodeId, parents: &[Parent]) -> Arc { + match self.analyse(node, parents) { + GraphAnalysis::NoCollision(graph) => { + if graph.origin != node { + self.graphs.insert(node, graph.clone()); + self.register_key(graph.origin, node); + } + + graph + } + GraphAnalysis::Collisions(graphs) => self.merge(node, graphs), + } + } + + /// Analyses the graph for a given node and its parents, returning the associated `GraphAnalysis`. + fn analyse(&mut self, node: NodeId, parents: &[Parent]) -> GraphAnalysis { + // If no parents, there is no collision, therefore a single graph is ok. + if parents.is_empty() { + let graph = match self.graphs.get(&node) { + Some(val) => val.clone(), + None => self.new_graph(node), + }; + return GraphAnalysis::NoCollision(graph); + }; + + // We collect all graphs of parents and of the current node based on their origin node id. + let mut graphs = BTreeMap::>::new(); + + if let Some(val) = self.graphs.get(&node) { + graphs.insert(val.origin, val.clone()); + } + + for parent in parents { + match self.graphs.get(&parent.id) { + Some(graph) => graphs.insert(graph.origin, graph.clone()), + None => continue, + }; + } + + if graphs.is_empty() { + return match self.graphs.get(&node) { + Some(old) => GraphAnalysis::NoCollision(old.clone()), + None => GraphAnalysis::NoCollision(self.new_graph(node)), + }; + } + + if graphs.len() == 1 { + return GraphAnalysis::NoCollision(graphs.into_values().next().unwrap()); + } + + GraphAnalysis::Collisions(graphs) + } + + /// Merges multiple graphs associated with a node into a single graph. + fn merge(&mut self, node: NodeId, graphs: BTreeMap>) -> Arc { + // `into_values()` on the `BTreeMap` yields graphs in ascending `NodeId` order, so the "main" + // graph and the merge order are deterministic (a `HashMap` here picked an arbitrary, + // run-to-run-varying main, which permuted the merged steps and the backward op stream). + let mut graphs = graphs.into_values(); + + let main = graphs.next().expect("At least one graph"); + self.register_key(main.origin, node); + + let mut state = main.state.lock(); + + for graph in graphs { + self.merge_two(&mut state, &main, graph); + } + + self.graphs.insert(main.origin, main.clone()); + self.graphs.insert(node, main.clone()); + + core::mem::drop(state); + + main + } + + /// Registers a key for a given origin node. + fn register_key(&mut self, origin: NodeId, key: NodeId) { + self.keys.entry(origin).or_default(); + + if origin != key { + // Register this node to point to the origin graph + self.keys.get_mut(&origin).unwrap().insert(key); + } + } + + /// Merges two graphs by combining their states and updating graph mappings. + fn merge_two(&mut self, main_state: &mut GraphState, main: &Arc, merged: Arc) { + let mut locked = merged.state.lock(); + let mut state_old = GraphState::default(); + core::mem::swap(&mut state_old, &mut locked); + main_state.server.extend(state_old.server); + + // Re-map merged origin to the main graph + self.graphs.insert(merged.origin, main.clone()); + + // Move all keys (node IDs) from the merged graph to the main graph + if let Some(locator_keys) = self.keys.remove(&merged.origin) { + for k in locator_keys.iter() { + self.graphs.insert(*k, main.clone()); + } + + let locator_keys_main = self + .keys + .get_mut(&main.origin) + .expect("Should be init before the merge."); + locator_keys_main.extend(locator_keys); + } + } + + /// Creates a new graph for a given node. + fn new_graph(&mut self, origin: NodeId) -> Arc { + let graph = Arc::new(Graph { + origin, + state: Mutex::new(GraphState::default()), + }); + self.graphs.insert(origin, graph.clone()); + self.keys.insert(origin, HashSet::new()); + graph + } + + fn remove_entry(&mut self, node: &NodeId) { + if let Some(graph) = self.graphs.remove(node) { + let mut remove = false; + + if let Some(entry) = self.keys.get_mut(&graph.origin) { + entry.remove(node); + if entry.is_empty() { + remove = true; + } + } + + if remove { + self.keys.remove(&graph.origin); + } + } + } +} + +/// Represents the analysis result of graph operations for a given node and its parents. +#[derive(Debug)] +enum GraphAnalysis { + /// No collision detected, contains the graph associated with the node. + NoCollision(Arc), + /// Collision detected, contains a map of node IDs to their associated graphs. + Collisions(BTreeMap>), +} diff --git a/crates/burn-autodiff/src/runtime/memory_management.rs b/crates/burn-autodiff/src/runtime/memory_management.rs new file mode 100644 index 0000000..0a1f4c6 --- /dev/null +++ b/crates/burn-autodiff/src/runtime/memory_management.rs @@ -0,0 +1,301 @@ +use crate::{ + NodeId, + collections::{HashMap, HashSet}, + graph::Parent, + tensor::NodeRefCount, +}; +use alloc::{borrow::ToOwned, vec, vec::Vec}; + +#[cfg(target_has_atomic = "ptr")] +use alloc::sync::Arc; + +#[cfg(not(target_has_atomic = "ptr"))] +use portable_atomic_util::Arc; + +use core::mem; + +#[derive(Default, Debug)] +pub struct GraphMemoryManagement { + nodes: HashMap>, + leaves: HashSet, + statuses: HashMap, +} + +#[derive(Debug, Clone, PartialEq)] +enum NodeMemoryStatus { + Useful, + Unavailable, + Unknown, +} + +impl GraphMemoryManagement { + pub fn extend(&mut self, other: Self) { + self.nodes.extend(other.nodes); + self.leaves.extend(other.leaves); + self.statuses.extend(other.statuses); + } + + /// Register a new node with its parent. + pub fn register(&mut self, node: NodeRefCount, parents: &[Parent]) { + let node_id = *node.as_ref(); + + for parent in parents.iter() { + self.leaves.remove(&parent.id); + } + + self.leaves.insert(node_id); + self.nodes + .insert(node, parents.iter().map(|p| p.id).collect()); + } + + /// Free the node from the state. + pub fn consume_node(&mut self, node_id: NodeId) { + if !self.is_referenced(node_id) { + self.leaves.remove(&node_id); + self.nodes.remove(&node_id); + } + } + + /// Free all nodes whose backward call has become impossible + /// + /// This function goes into three steps, which must happen for all leaves + /// before going into the next step. Then it deletes what can be safely deleted + pub(crate) fn free_unavailable_nodes(&mut self, mut on_free_graph: impl FnMut(&NodeId)) { + let leaves = self.leaves.clone(); + let mut new_leaves = HashSet::new(); + let mut deletables = Vec::new(); + + // When consuming nodes with a backward pass, some other backward passes become + // unavailable because some of their parents have been consumed. They are + // identified here. + for leaf in leaves.clone() { + self.unavailable_propagation(leaf); + } + + // Among the available nodes that remain, some may be useless if no + // available node with a tensor reference exist in their descendance. + // But some may seem useless from some leaf but be useful from another one, + // hence the need to iterate on all leaves. + self.useful_propagation(leaves.clone()); + + // New leaves are the roots of a useful backward sub-tree. + // Deletables are everything not marked as useful. + for leaf in leaves { + self.identify_leaves_and_deletables(leaf, &mut new_leaves, &mut deletables); + } + + // Replace leaves by the new ones and delete everything not useful anymore + mem::swap(&mut self.leaves, &mut new_leaves); + + self.clear_unused_roots(&mut deletables); + + self.statuses.clear(); + for node_to_delete in deletables { + self.nodes.remove(&node_to_delete); + on_free_graph(&node_to_delete) + } + } + + pub(crate) fn free_unused_roots(&mut self, mut on_free_graph: impl FnMut(&NodeId)) { + let mut deletables = Vec::new(); + self.clear_unused_roots(&mut deletables); + + for node_id in deletables { + self.nodes.remove(&node_id); + on_free_graph(&node_id); + } + } + + fn clear_unused_roots(&self, to_delete: &mut Vec) { + for (id, parents) in self.nodes.iter() { + let is_useful = matches!( + self.statuses.get(id.as_ref()), + Some(NodeMemoryStatus::Useful) + ); + + // Check if parents are either empty or absent from self.nodes + let parents_absent = parents.iter().all(|p| !self.nodes.contains_key(p)); + + if !is_useful && Arc::strong_count(id) == 1 && parents_absent { + to_delete.push(*id.as_ref()) + } + } + } + + fn unavailable_propagation(&mut self, node_id: NodeId) -> NodeMemoryStatus { + // If already visited + if let Some(status) = self.statuses.get(&node_id) { + return status.clone(); + } + + match self.nodes.get(&node_id).cloned() { + // If node exists and any of its parents is unavailable, it is unavailable as well + // If node exists but the parents vec is empty, it is a tensor that never had parents; + // the status remains unknown + Some(parents) => { + let mut node_status = NodeMemoryStatus::Unknown; + for parent in parents { + let parent_status = self.unavailable_propagation(parent); + if let NodeMemoryStatus::Unavailable = parent_status { + node_status = NodeMemoryStatus::Unavailable; + } + } + self.statuses.insert(node_id, node_status.clone()); + node_status + } + // If node does not exist, it was + // deleted, so this and all its descendants are unavailable + None => { + self.statuses.insert(node_id, NodeMemoryStatus::Unavailable); + NodeMemoryStatus::Unavailable + } + } + } + + fn useful_propagation(&mut self, leaves: HashSet) { + // Accumulate visited nodes + let mut explored = HashSet::new(); + let mut tagged_useful = HashSet::new(); + + // Queue of nodes to visit + let mut to_tag_useful = PopNodeSet::default(); + let mut to_explore = PopNodeSet::new(leaves); + + // Utility function to iterate over a node's parents + let parents = |node_id| { + self.nodes + .get(&node_id) + .cloned() + .unwrap_or_default() + .into_iter() + }; + + loop { + // Pop a node id, greedily looking at tag_useful ones first + let (node_id, status) = match to_tag_useful.pop() { + Some(node_id) => (node_id, NodeMemoryStatus::Useful), + None => match to_explore.pop() { + Some(node_id) => { + let node_status = self + .statuses + .get(&node_id) + .expect("All nodes should have received a status during unavailable_propagation") + .to_owned(); + + if let NodeMemoryStatus::Unknown = node_status { + match self.is_referenced(node_id) { + true => (node_id, NodeMemoryStatus::Useful), + false => (node_id, NodeMemoryStatus::Unknown), + } + } else { + (node_id, node_status) + } + } + None => { + // There are no nodes in the queues anymore + break; + } + }, + }; + + match status { + NodeMemoryStatus::Useful => { + tagged_useful.insert(node_id); + for parent in parents(node_id) { + // The node can be explored, as long as it's not already tagged useful + if !(tagged_useful.contains(&parent) || to_tag_useful.contains(&parent)) { + to_tag_useful.insert(parent); + } + } + } + _ => { + explored.insert(node_id); + for parent in parents(node_id) { + if !(explored.contains(&parent) || to_explore.contains(&parent)) { + to_explore.insert(parent); + } + } + } + } + + self.statuses.insert(node_id, status); + } + } + + fn identify_leaves_and_deletables( + &self, + leaf_id: NodeId, + new_leaves: &mut HashSet, + to_delete: &mut Vec, + ) { + let mut visited = HashSet::new(); + let mut to_visit = vec![leaf_id]; + + while let Some(node_id) = to_visit.pop() { + visited.insert(node_id); + + match self + .statuses + .get(&node_id) + .expect("Node should have status") + { + NodeMemoryStatus::Useful => { + new_leaves.insert(node_id); + } + _ => { + to_delete.push(node_id); + + for parent in self + .nodes + .get(&node_id) + .cloned() + .unwrap_or_default() + .into_iter() + { + if !visited.contains(&parent) { + to_visit.push(parent); + } + } + } + }; + } + } + + fn is_referenced(&self, node_id: NodeId) -> bool { + match self.nodes.get_key_value(&node_id) { + Some((key, _value)) => Arc::strong_count(key) > 1, + None => panic!("Node should be in the nodes map"), + } + } + + pub(crate) fn maybe_useful(&self) -> bool { + self.nodes.keys().any(|node| Arc::strong_count(node) > 1) + } +} + +/// Wrapper over hash set for fast popping of any node +#[derive(new, Default)] +struct PopNodeSet { + hash_set: HashSet, +} + +impl PopNodeSet { + #[inline(always)] + fn pop(&mut self) -> Option { + self.hash_set + .iter() + .next() + .copied() + .and_then(|node_id| self.hash_set.take(&node_id)) + } + + #[inline(always)] + fn contains(&self, node_id: &NodeId) -> bool { + self.hash_set.contains(node_id) + } + + #[inline(always)] + fn insert(&mut self, node_id: NodeId) { + self.hash_set.insert(node_id); + } +} diff --git a/crates/burn-autodiff/src/runtime/mod.rs b/crates/burn-autodiff/src/runtime/mod.rs new file mode 100644 index 0000000..fde3a5f --- /dev/null +++ b/crates/burn-autodiff/src/runtime/mod.rs @@ -0,0 +1,6 @@ +mod client; +mod memory_management; +mod server; + +pub mod graph; +pub use client::*; diff --git a/crates/burn-autodiff/src/runtime/server.rs b/crates/burn-autodiff/src/runtime/server.rs new file mode 100644 index 0000000..7c29c49 --- /dev/null +++ b/crates/burn-autodiff/src/runtime/server.rs @@ -0,0 +1,201 @@ +use super::memory_management::GraphMemoryManagement; +use crate::{ + NodeId, + checkpoint::{ + base::{Checkpointer, NodeTree}, + builder::CheckpointerBuilder, + }, + collections::HashMap, + grads::{BackwardMode, Gradients}, + graph::{ + NodeRef, StepBoxed, + traversal::{BreadthFirstSearch, TraversalItem}, + }, + tensor::NodeRefCount, +}; +use alloc::vec::Vec; +use burn_backend::Backend; +use burn_backend::tensor::FloatTensor; + +#[cfg(feature = "std")] +use crate::grads::GradSyncContext; + +struct TapeResult { + tape: Vec>, + checkpointer: Checkpointer, + #[cfg(feature = "std")] + distributed: Option, +} + +#[derive(Default)] +pub struct AutodiffServer { + steps: HashMap, + actions_builder: HashMap, + memory_management: GraphMemoryManagement, +} + +/// Defines how nodes are clean. +pub trait NodeCleaner { + /// Initialize a new cleaner. + fn init() -> Self; + /// Cleans a single [node](NodeId). + fn clean(&mut self, node: &NodeId); +} + +impl AutodiffServer { + pub fn extend(&mut self, other: AutodiffServer) { + self.steps.extend(other.steps); + self.actions_builder.extend(other.actions_builder); + self.memory_management.extend(other.memory_management); + } + + pub fn register(&mut self, rc: NodeRefCount, step: StepBoxed, actions: CheckpointerBuilder) { + let parents = step.parents(); + let node_id = *rc.as_ref(); + + self.memory_management.register(rc, parents); + + self.steps.insert(node_id, step); + self.actions_builder.insert(node_id, actions); + } + + pub fn backward( + &mut self, + root_node: NodeRef, + root_tensor: FloatTensor, + node_id: NodeId, + mode: BackwardMode, + ) -> Gradients { + let step = self.steps.remove(&node_id).expect( + "Node should have a step registered, did you forget to call \ + `Tensor::register_grad` on the tensor where you need gradients?", + ); + let builder = self.actions_builder.remove(&node_id).unwrap(); + + let mut consumed = Vec::new(); + let tape_result = self.build_tape(node_id, step, builder, &mut consumed); + + let grads = match mode { + #[cfg(feature = "std")] + BackwardMode::Distributed(factory) if tape_result.distributed.is_some() => { + let on_register = factory(tape_result.distributed.clone().unwrap()); + Gradients::new_distributed::(root_node, root_tensor, on_register) + } + _ => Gradients::new::(root_node, root_tensor), + }; + + let gradients = Self::execute_steps(tape_result.tape, grads, tape_result.checkpointer); + + self.cleanup::(&consumed); + + gradients + } + + fn cleanup(&mut self, consumed: &Vec) { + let mut cleaner = NC::init(); + self.memory_management + .free_unavailable_nodes(|node_id: &NodeId| { + self.steps.remove(node_id); + self.actions_builder.remove(node_id); + NC::clean(&mut cleaner, node_id); + }); + for node_id in consumed { + cleaner.clean(node_id) + } + } + + pub(crate) fn free_unused_roots(&mut self, mut on_free_graph: impl FnMut(&NodeId)) { + self.memory_management.free_unused_roots(|node_id| { + self.steps.remove(node_id); + self.actions_builder.remove(node_id); + on_free_graph(node_id); + }); + } + + fn build_tape( + &mut self, + node: NodeId, + node_step: StepBoxed, + mut builder: CheckpointerBuilder, + consumed: &mut Vec, + ) -> TapeResult { + let mut tape = (0..node_step.depth() + 1) + .map(|_| Vec::with_capacity(1)) + .collect::>(); + + let mut tree = HashMap::default(); + + #[cfg(feature = "std")] + let mut n_required_map = HashMap::default(); + #[cfg(feature = "std")] + let mut distributed_params = HashMap::default(); + + BreadthFirstSearch.traverse(node, node_step, &mut self.steps, |id, step| { + self.memory_management.consume_node(id); + // Clean up consumed node + consumed.push(id); + + let depth = step.depth(); + + #[cfg(feature = "std")] + step.distributed_params() + .and_then(|params| distributed_params.insert(id, params)); + + if let Some(steps) = tape.get_mut(depth) { + let parents = step + .parents() + .iter() + .map(|p| { + #[cfg(feature = "std")] + { + *n_required_map.entry(p.id).or_insert(0) += 1; + } + p.id + }) + .filter(|s| *s != id); + tree.insert(id, parents.collect()); + steps.push(step); + } + + if let Some(node_builder) = self.actions_builder.remove(&id) { + builder.extend(node_builder); + } + }); + + let checkpointer = builder.build(NodeTree::new(tree)); + #[cfg(feature = "std")] + let distributed = Some(GradSyncContext { + n_required_map, + distributed_params, + }); + + TapeResult { + tape, + checkpointer, + #[cfg(feature = "std")] + distributed, + } + } + + fn execute_steps( + tape: Vec>, + mut grads: Gradients, + mut checkpointer: Checkpointer, + ) -> Gradients { + tape.into_iter().rev().for_each(|steps| { + steps + .into_iter() + .for_each(|step| step.step(&mut grads, &mut checkpointer)) + }); + + // For checkpointing tests + #[cfg(feature = "export_tests")] + assert!(checkpointer.is_empty()); + + grads + } + + pub(crate) fn maybe_useful(&self) -> bool { + self.memory_management.maybe_useful() + } +} diff --git a/crates/burn-autodiff/src/tensor.rs b/crates/burn-autodiff/src/tensor.rs new file mode 100644 index 0000000..8b1db6d --- /dev/null +++ b/crates/burn-autodiff/src/tensor.rs @@ -0,0 +1,261 @@ +use crate::{ + checkpoint::{base::Checkpointer, builder::CheckpointerBuilder}, + grads::{BackwardMode, Gradients}, + graph::{ComputingProperty, Node, NodeId, NodeRef, Parent, Requirement, Step}, + runtime::{AutodiffClient, AutodiffClientImpl}, +}; +#[cfg(feature = "std")] +use crate::{distributed::DistributedGradientRegistration, grads::GradSyncContext}; +use alloc::{boxed::Box, vec}; +use burn_backend::{Backend, BackendTypes, TensorMetadata}; + +#[cfg(target_has_atomic = "ptr")] +use alloc::sync::Arc; + +#[cfg(not(target_has_atomic = "ptr"))] +use portable_atomic_util::Arc; + +use burn_backend::distributed::{DistributedParamId, DistributedParams}; + +#[derive(Debug, Clone)] +pub struct AutodiffTensor { + pub primitive: B::FloatTensorPrimitive, + pub node: NodeRef, + pub rc: NodeRefCount, +} + +impl TensorMetadata for AutodiffTensor { + type Device = B::Device; + fn dtype(&self) -> burn_std::DType { + self.primitive.dtype() + } + + fn shape(&self) -> burn_std::Shape { + self.primitive.shape() + } + + fn rank(&self) -> usize { + self.primitive.rank() + } + fn device(&self) -> B::Device { + self.primitive.device() + } + + fn can_mut(&self) -> bool { + // Precise: the inner handle's buffer refcount already accounts for any + // clone the autodiff graph retains (e.g. checkpointed states). + self.primitive.can_mut() + } +} + +pub type NodeRefCount = Arc; + +#[derive(new, Debug)] +pub(crate) struct RootStep { + node: NodeRef, +} + +impl Step for RootStep { + fn step(self: Box, _grads: &mut Gradients, _checkpointer: &mut Checkpointer) { + // Nothing to do + } + + fn node(&self) -> NodeId { + self.node.id + } + + fn parents(&self) -> &[Parent] { + &self.node.parents + } + + fn depth(&self) -> usize { + self.node.order + } + + fn distributed_params(&self) -> Option { + self.node.distributed_params.clone() + } +} + +impl AutodiffTensor { + /// Create a new leaf tensor. + pub fn new(primitive: B::FloatTensorPrimitive) -> Self { + let id = NodeId::new(); + let node: NodeRef = Node::new( + vec![], + 0, + id, + Requirement::None, + ComputingProperty::Ambiguous, + AutodiffClientImpl::new(), + None, + ) + .into(); + + Self { + rc: Arc::new(node.id), + primitive, + node: node.clone(), + } + } + + pub fn is_tracked(&self) -> bool { + !self.node.requirement.is_none() + } + + /// Mark the tensor as requiring gradients. + /// + /// # Panics + /// + /// It panics if the tensor is not a leaf. + pub fn require_grad(mut self) -> Self { + match self.node.requirement { + Requirement::Grad => self, + Requirement::GradInBackward => { + panic!("Can't convert a non leaf tensor into a tracked tensor") + } + Requirement::None => { + self.node = Node::new( + vec![], + 0, + self.node.id, + Requirement::Grad, + self.node.properties.clone(), + self.node.client.clone(), + self.node.distributed_params.clone(), + ) + .into(); + let step = RootStep::new(self.node.clone()); + + self.register_step(step, CheckpointerBuilder::default()) + } + } + } + + /// Create a tensor from parent infos. + pub fn from_parents( + primitive: B::FloatTensorPrimitive, + parent_nodes: &[NodeRef], + requirement: Requirement, + computing_properties: ComputingProperty, + ) -> Self { + let order = parent_nodes + .iter() + .map(|node| node.order) + .reduce(usize::max) + .unwrap_or(0) + + 1; + + let client = parent_nodes + .first() + .map(|node| node.client.clone()) + .unwrap_or_else(AutodiffClientImpl::new); + + let node: NodeRef = Node::new( + parent_nodes + .iter() + .filter_map(|node| node.clone_if_require_grad()) + .map(|node| Parent::new(node.id)) + .collect(), + order, + NodeId::new(), + requirement, + computing_properties, + client, + None, + ) + .into(); + + Self { + rc: Arc::new(node.id), + primitive, + node, + } + } + + /// Register a step into a graph for that tensor. + /// + /// # Warning + /// + /// This should be called only once per tensor. + pub fn register_step( + self, + step_that_created_the_tensor: S, + actions: CheckpointerBuilder, + ) -> Self { + self.node.client.register( + self.rc.clone(), + Box::new(step_that_created_the_tensor), + actions, + ); + self + } + + pub fn into_primitive(self) -> B::FloatTensorPrimitive { + self.primitive + } + + #[cfg(not(feature = "std"))] + pub fn backward(self) -> Gradients { + let client = self.node.client.clone(); + + AutodiffClient::backward::(&client, self, BackwardMode::default()) + } + + pub fn grad(&self, grads: &Gradients) -> Option { + grads.get::(self) + } + + pub fn grad_remove(&self, grads: &mut Gradients) -> Option { + grads.remove::(self) + } + + pub fn grad_replace(&self, grads: &mut Gradients, grad: B::FloatTensorPrimitive) { + grads.remove::(self); + grads.register::(self.node.id, grad); + } + + /// Mark the tensor as distributed across multiple devices. + /// Its gradients will be automatically aggregated from those devices after the backward pass. + /// + /// # Arguments + /// + /// * `param_id` - The module tensor's [`DistributedParamId`]. + pub fn grad_distributed(mut self, param_id: DistributedParamId) -> Self { + self.node = Node::new( + vec![], + 0, + self.node.id, + self.node.requirement, + self.node.properties.clone(), + self.node.client.clone(), + Some(DistributedParams { param_id }), + ) + .into(); + let step = RootStep::new(self.node.clone()); + + self.register_step(step, CheckpointerBuilder::default()) + } +} + +#[cfg(feature = "std")] +impl AutodiffTensor { + pub fn backward(self) -> Gradients { + let device = self.primitive.device(); + let device_cloned = device.clone(); + let client = self.node.client.clone(); + + let mode = BackwardMode::Distributed(Box::new(|ctx: GradSyncContext| { + let registration = DistributedGradientRegistration::::new( + ctx.n_required_map, + ctx.distributed_params, + device_cloned, + ); + Box::new(registration) + })); + + let grads = AutodiffClient::backward::(&client, self, mode); + B::submit_sync_collective(&device); + grads + } +} diff --git a/crates/burn-autodiff/src/utils.rs b/crates/burn-autodiff/src/utils.rs new file mode 100644 index 0000000..1f3549c --- /dev/null +++ b/crates/burn-autodiff/src/utils.rs @@ -0,0 +1,25 @@ +use alloc::vec::Vec; + +use crate::graph::NodeRef; +/// Duplicate the given object for each node that requires gradients. +/// +/// # Notes +/// +/// This is useful since you don't have to keep N cloned references alive event if just 1 node +/// will be updated. +/// +/// If the object is a tensor and if one reference exists, it can be updated inplace. +pub fn duplicate( + nodes: &[Option; N], + obj: Option, +) -> [Option; N] { + nodes + .iter() + .map(|node| match node { + Some(_) => obj.clone(), + None => None, + }) + .collect::>() + .try_into() + .unwrap() +} diff --git a/crates/burn-backend-extension/Cargo.toml b/crates/burn-backend-extension/Cargo.toml new file mode 100644 index 0000000..00d257f --- /dev/null +++ b/crates/burn-backend-extension/Cargo.toml @@ -0,0 +1,20 @@ +[package] +authors = ["nathanielsimard "] +description = "Burn backend extension proc macro" +edition.workspace = true +license.workspace = true +name = "burn-backend-extension" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-backend-extension" +version.workspace = true + +[lints] +workspace = true + +[lib] +proc-macro = true + +[dependencies] +proc-macro2 = { workspace = true } +quote = { workspace = true } +syn = { workspace = true } diff --git a/crates/burn-backend-extension/LICENSE-APACHE b/crates/burn-backend-extension/LICENSE-APACHE new file mode 120000 index 0000000..1cd601d --- /dev/null +++ b/crates/burn-backend-extension/LICENSE-APACHE @@ -0,0 +1 @@ +../../LICENSE-APACHE \ No newline at end of file diff --git a/crates/burn-backend-extension/LICENSE-MIT b/crates/burn-backend-extension/LICENSE-MIT new file mode 120000 index 0000000..b2cfbdc --- /dev/null +++ b/crates/burn-backend-extension/LICENSE-MIT @@ -0,0 +1 @@ +../../LICENSE-MIT \ No newline at end of file diff --git a/crates/burn-backend-extension/README.md b/crates/burn-backend-extension/README.md new file mode 100644 index 0000000..376a93d --- /dev/null +++ b/crates/burn-backend-extension/README.md @@ -0,0 +1,6 @@ +# Burn Backend Extension + +> [Burn](https://github.com/tracel-ai/burn) backend extension generation + +[![Current Crates.io Version](https://img.shields.io/crates/v/burn-backend-extension.svg)](https://crates.io/crates/burn-backend-extension) +[![license](https://shields.io/badge/license-MIT%2FApache--2.0-blue)](https://github.com/tracel-ai/burn-backend-extension/blob/master/README.md) diff --git a/crates/burn-backend-extension/src/lib.rs b/crates/burn-backend-extension/src/lib.rs new file mode 100644 index 0000000..4646490 --- /dev/null +++ b/crates/burn-backend-extension/src/lib.rs @@ -0,0 +1,875 @@ +use proc_macro::TokenStream; +use quote::{format_ident, quote}; + +use proc_macro2::TokenStream as TokenStream2; +use syn::parse::{Parse, ParseStream}; +use syn::punctuated::Punctuated; +use syn::{ + FnArg, GenericArgument, Ident, ItemStruct, ItemTrait, Meta, Pat, PathArguments, ReturnType, + Token, TraitItem, Type, TypeParamBound, parse_macro_input, +}; + +/// # `backend_extension` +/// +/// Attribute macro that generates dispatch glue for Burn backend extension traits. +/// +/// ## Usage +/// +/// ```rust,ignore +/// use burn_backend_extension::backend_extension; +/// +/// #[backend_extension(Wgpu, Cuda, Cpu, Autodiff)] +/// pub trait MyExtension: Backend { +/// fn fused_matmul_add_relu(lhs: FloatTensor, rhs: FloatTensor, bias: FloatTensor) -> FloatTensor; +/// fn custom_threshold(x: FloatTensor, threshold: f32) -> FloatTensor; +/// } +/// ``` +/// +/// ### What gets generated +/// +/// - An `impl Trait for Dispatch` is generated automatically. +/// - Each method dispatches to the corresponding implementation for the listed backends. +/// - If `Autodiff` is specified, autodiff variants are also handled automatically. +/// - All other backends are left as `unimplemented!()`. +/// +/// Supported tensor argument/return types: `FloatTensor`, `IntTensor`, +/// `BoolTensor`, and `QuantizedTensor` (a QFloat tensor passed to the +/// backend still quantized — the op reads the packed values/scales directly instead +/// of going through a dequantize). +#[proc_macro_attribute] +pub fn backend_extension(attr: TokenStream, item: TokenStream) -> TokenStream { + // Parse the backend list + let backends = parse_macro_input!(attr as Backends); + // Parse the trait definition + let trait_def = parse_macro_input!(item as ItemTrait); + + // Lower to extension representation, then expand codegen + let expanded = lower_extension(backends, &trait_def) + .map(|ir| expand_extension(ir, trait_def)) + .unwrap_or_else(|err| err.to_compile_error()); + + TokenStream::from(expanded) +} + +#[derive(Debug, Clone)] +struct Backend { + pub kind: BackendKind, + pub cfg: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BackendKind { + Cpu, + Cuda, + Rocm, + Metal, + Vulkan, + Wgpu, + WebGpu, + Flex, + NdArray, + LibTorch, + Remote, +} + +impl BackendKind { + fn try_from(ident: &Ident) -> syn::Result { + match ident.to_string().as_str() { + "Cpu" => Ok(BackendKind::Cpu), + "Cuda" => Ok(BackendKind::Cuda), + "Wgpu" => Ok(BackendKind::Wgpu), + "WebGpu" => Ok(BackendKind::WebGpu), + "Metal" => Ok(BackendKind::Metal), + "Rocm" => Ok(BackendKind::Rocm), + "Vulkan" => Ok(BackendKind::Vulkan), + "Flex" => Ok(BackendKind::Flex), + "NdArray" => Ok(BackendKind::NdArray), + "LibTorch" => Ok(BackendKind::LibTorch), + "Remote" => Ok(BackendKind::Remote), + other => Err(syn::Error::new_spanned( + ident, + format!("Unsupported backend `{}`", other), + )), + } + } +} + +struct Backends { + concrete: Vec, + autodiff: (bool, Option), +} + +// Helper to parse backend idents w/ optional cfg +struct BackendArg { + id: Ident, + cfg: Option, +} + +impl Parse for BackendArg { + fn parse(input: ParseStream) -> syn::Result { + let id: Ident = input.parse()?; + let cfg = if input.peek(Token![:]) { + input.parse::()?; + + // This parses cfg(feature = "...") or any other meta item + let meta: syn::Meta = input.parse()?; + Some(meta) + } else { + None + }; + + Ok(Self { id, cfg }) + } +} + +impl Parse for Backends { + fn parse(input: ParseStream) -> syn::Result { + let args = Punctuated::::parse_terminated(input)?; + + let mut concrete = vec![]; + let mut autodiff = (false, None); + + for arg in args { + if arg.id == "Autodiff" { + autodiff = (true, arg.cfg); + continue; + } + + concrete.push(Backend { + kind: BackendKind::try_from(&arg.id)?, + cfg: arg.cfg, + }); + } + + Ok(Backends { concrete, autodiff }) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TensorKind { + Float, + Int, + Bool, + Quantized, +} + +#[allow(clippy::large_enum_variant)] +enum ArgKind { + Tensor(TensorKind), + // Passthrough - unhandled by the macro + Other(Type), +} + +struct OperationArg { + name: Ident, + kind: ArgKind, +} + +#[allow(clippy::large_enum_variant)] +#[derive(Debug, Clone)] +enum OutputKind { + Tensor(TensorKind), + Custom(Type), +} + +#[allow(clippy::large_enum_variant)] +#[derive(Debug, Clone)] +enum OperationOutput { + Tensor(TensorKind), + Tuple(Vec), + Custom(Type), +} + +struct Operation { + name: Ident, + inputs: Vec, + output: OperationOutput, + asyncness: bool, +} + +struct Extension { + trait_name: Ident, + backends: Backends, + ops: Vec, +} + +impl TensorKind { + fn from_type(ty: &Type) -> Option { + match ty { + // Handle `::` + Type::Path(tp) if tp.qself.is_some() => { + let last = tp.path.segments.last()?.ident.to_string(); + + match last.as_str() { + "FloatTensorPrimitive" => Some(Self::Float), + "IntTensorPrimitive" => Some(Self::Int), + "BoolTensorPrimitive" => Some(Self::Bool), + "QuantizedTensorPrimitive" => Some(Self::Quantized), + _ => None, + } + } + // Handle simple paths: Float, FloatTensor, burn::...::FloatTensor + Type::Path(tp) => { + let last = tp.path.segments.last()?.ident.to_string(); + + match last.as_str() { + // Shorthand + "Float" => Some(Self::Float), + "Int" => Some(Self::Int), + "Bool" => Some(Self::Bool), + "Quantized" => Some(Self::Quantized), + + // Full tensor types + "FloatTensor" => Some(Self::Float), + "IntTensor" => Some(Self::Int), + "BoolTensor" => Some(Self::Bool), + "QuantizedTensor" => Some(Self::Quantized), + + // Associated primitive types + "FloatTensorPrimitive" => Some(Self::Float), + "IntTensorPrimitive" => Some(Self::Int), + "BoolTensorPrimitive" => Some(Self::Bool), + "QuantizedTensorPrimitive" => Some(Self::Quantized), + + _ => None, + } + } + + // Handle references like &FloatTensor + Type::Reference(r) => Self::from_type(&r.elem), + + // Handle parentheses `(FloatTensor)` + Type::Paren(p) => Self::from_type(&p.elem), + + // TODO: option/containers already handled? + _ => None, + } + } + + fn to_primitive_ty(self) -> TokenStream2 { + match self { + Self::Float => quote! { burn::backend::tensor::FloatTensor }, + Self::Int => quote! { burn::backend::tensor::IntTensor }, + Self::Bool => quote! { burn::backend::tensor::BoolTensor }, + Self::Quantized => quote! { burn::backend::tensor::QuantizedTensor }, + } + } + + fn variant(self) -> Ident { + match self { + Self::Float => format_ident!("Float"), + Self::Int => format_ident!("Int"), + Self::Bool => format_ident!("Bool"), + Self::Quantized => format_ident!("Quantized"), + } + } + + fn unwrap_method(self) -> Ident { + // e.g. tensor.float() (BackendTensor method) + format_ident!("{}", format!("{:?}", self).to_lowercase()) + } +} + +fn backend_to_ident(b: &Backend) -> Ident { + // Convert the enum variant to a string first, then to an Ident + format_ident!("{}", format!("{:?}", b.kind)) +} + +fn extract_future_output_type(ty: &Type) -> Option<&Type> { + if let Type::ImplTrait(impl_trait) = ty { + for bound in &impl_trait.bounds { + if let TypeParamBound::Trait(trait_bound) = bound { + let last_segment = trait_bound.path.segments.last()?; + if last_segment.ident == "Future" + && let PathArguments::AngleBracketed(args) = &last_segment.arguments + { + for arg in &args.args { + if let GenericArgument::AssocType(assoc) = arg + && assoc.ident == "Output" + { + return Some(&assoc.ty); + } + } + } + } + } + } + None +} + +fn lower_extension(attr: Backends, item: &ItemTrait) -> syn::Result { + let mut ops = Vec::new(); + + for trait_item in &item.items { + let TraitItem::Fn(f) = trait_item else { + continue; + }; + + // Parse Inputs + let mut inputs = Vec::new(); + for arg in &f.sig.inputs { + let FnArg::Typed(pt) = arg else { continue }; + let name = match pt.pat.as_ref() { + Pat::Ident(p) => p.ident.clone(), + _ => return Err(syn::Error::new_spanned(&pt.pat, "Unsupported pattern")), + }; + let kind = TensorKind::from_type(&pt.ty) + .map(ArgKind::Tensor) + .unwrap_or_else(|| ArgKind::Other((*pt.ty).clone())); + inputs.push(OperationArg { name, kind }); + } + + // Parse outputs + + // Parse outputs + let (actual_ty, is_async) = match &f.sig.output { + ReturnType::Default => { + return Err(syn::Error::new_spanned( + &f.sig.output, + "Operations must return a value", + )); + } + ReturnType::Type(_, ty) => { + // If it's `impl Future`, extract T and mark as async. + // Otherwise, use the type as-is and check for `async fn`. + if let Some(out_ty) = extract_future_output_type(ty) { + (out_ty, true) + } else { + (ty.as_ref(), f.sig.asyncness.is_some()) + } + } + }; + + let output = match actual_ty { + // TODO: expand support for vec and maybe nested containers + Type::Tuple(tup) => { + let elements = tup + .elems + .iter() + .map(|elem| { + if let Some(kind) = TensorKind::from_type(elem) { + Ok(OutputKind::Tensor(kind)) + } else { + Ok(OutputKind::Custom(elem.clone())) + } + }) + .collect::>>()?; + OperationOutput::Tuple(elements) + } + ty if TensorKind::from_type(ty).is_some() => { + OperationOutput::Tensor(TensorKind::from_type(ty).unwrap()) + } + ty => { + // ExtensionType + OperationOutput::Custom(ty.clone()) + } + }; + + ops.push(Operation { + name: f.sig.ident.clone(), + inputs, + output, + asyncness: is_async, + }); + } + + Ok(Extension { + trait_name: item.ident.clone(), + backends: attr, + ops, + }) +} + +fn expand_extension(ir: Extension, original_trait: ItemTrait) -> TokenStream2 { + let trait_name = &ir.trait_name; + + // Generate Dispatch Implementation + let dispatch_methods = ir.ops.iter().map(|op| gen_dispatch_method(&ir, op)); + + quote! { + #original_trait + + impl #trait_name for burn::backend::Dispatch { + #( #dispatch_methods )* + } + } +} + +fn gen_dispatch_method(ir: &Extension, op: &Operation) -> TokenStream2 { + let name = &op.name; + let has_ad = ir.backends.autodiff.0; + let ad_cfg_attr = ir + .backends + .autodiff + .1 + .as_ref() + .map(|meta| quote! { #[#meta] }); + + let maybe_async = if op.asyncness { + quote! { async } + } else { + quote! {} + }; + + let sig_args = op.inputs.iter().map(|arg| { + let name = &arg.name; + match &arg.kind { + ArgKind::Tensor(k) => { + let ty = k.to_primitive_ty(); + quote! { #name: #ty } + } + ArgKind::Other(ty) => quote! { #name: #ty }, + } + }); + + let ret_ty = match &op.output { + OperationOutput::Tensor(k) => k.to_primitive_ty(), + OperationOutput::Tuple(elems) => { + let types = elems.iter().map(|e| match e { + OutputKind::Tensor(k) => k.to_primitive_ty(), + OutputKind::Custom(ty) => quote! { #ty }, + }); + quote! { (#(#types),*) } + } + OperationOutput::Custom(ty) => quote! { #ty }, + }; + + let tensor_inputs: Vec<_> = op + .inputs + .iter() + .filter_map(|a| match &a.kind { + ArgKind::Tensor(_) => Some(&a.name), + _ => None, + }) + .collect(); + + let match_inputs = match tensor_inputs.len() { + 0 => quote! { () }, + 1 => { + let name = tensor_inputs[0]; + quote! { #name.kind } + } + _ => { + let kinds = tensor_inputs.iter().map(|n| quote! { #n.kind }); + quote! { (#(#kinds),*) } + } + }; + + let first_tensor = op.inputs.iter().find_map(|a| match &a.kind { + ArgKind::Tensor(_) => Some(&a.name), + _ => None, + }); + + let ckp_logic = if let Some(name) = first_tensor { + quote! { let checkpointing = #name.checkpointing.clone(); } + } else { + quote! { let checkpointing = None; } + }; + + let body = if tensor_inputs.is_empty() { + // No tensor input to select the backend from (e.g. `fn load_data(i: usize) -> FloatTensor`). + // There is nothing to match on, so this is only well-defined for a single backend — the + // remote backend is the motivating case (`#[backend_extension(Remote)]`), where the op is + // shipped to the server. Dispatch directly to that backend; reject the ambiguous cases. + if has_ad { + quote! { compile_error!("A backend extension operation with no tensor inputs can't be combined with `Autodiff` — there is no input tensor to carry the autodiff graph.") } + } else if ir.backends.concrete.len() == 1 { + let backend = &ir.backends.concrete[0]; + let call = gen_backend_call(ir, op, backend); + match &backend.cfg { + // Ungated backend: dispatch straight to it. + None => quote! { #ckp_logic #call }, + // The single backend is `cfg`-gated. Mirror the match path: gate the call on the + // backend's cfg and fall back to `unimplemented!` when it's compiled out, so the + // method still has a valid body (instead of referencing a backend that doesn't + // exist). `ckp_logic` lives inside the gated arm so its `checkpointing` binding + // isn't left dangling (and untypeable) when the backend is compiled out. + Some(meta) => quote! { + match () { + #[#meta] + () => { #ckp_logic #call } + #[allow(unreachable_patterns)] + _ => unimplemented!("Backend not supported for custom op `{}`", stringify!(#name)), + } + }, + } + } else { + quote! { compile_error!("A backend extension operation with no tensor inputs must list exactly one backend (e.g. `#[backend_extension(Remote)]`), since there is no input tensor to select the backend from.") } + } + } else { + let concrete_arms = ir + .backends + .concrete + .iter() + .map(|b| gen_backend_arm(ir, op, b)); + let ad_arm = if has_ad { + Some(gen_autodiff_arm(ir, op)) + } else { + None + }; + + quote! { + #ckp_logic + match #match_inputs { + #( #concrete_arms )* + #ad_cfg_attr + #ad_arm + _ => unimplemented!("Backend not supported for custom op `{}`", stringify!(#name)), + } + } + }; + + quote! { + #maybe_async fn #name(#(#sig_args),*) -> #ret_ty { + #body + } + } +} + +fn gen_backend_arm(ir: &Extension, op: &Operation, backend: &Backend) -> TokenStream2 { + let b_ident = backend_to_ident(backend); + + // If any cfg(..) was specified to gate the backend + let cfg_attr = backend.cfg.as_ref().map(|meta| quote! { #[#meta] }); + + // Filter for tensor arguments only + let tensor_args: Vec<_> = op + .inputs + .iter() + .filter_map(|a| match &a.kind { + ArgKind::Tensor(k) => Some((&a.name, k)), + _ => None, + }) + .collect(); + + // Build the match pattern for the same backend + // e.g., (DispatchTensorKind::Wgpu(lhs), DispatchTensorKind::Wgpu(rhs)) + let pattern = if tensor_args.len() == 1 { + let (name, _) = tensor_args[0]; + quote! { burn::backend::DispatchTensorKind::#b_ident(#name) } + } else { + let pats = tensor_args.iter().map(|(name, _)| { + quote! { burn::backend::DispatchTensorKind::#b_ident(#name) } + }); + quote! { (#(#pats),*) } + }; + + let call = gen_backend_call(ir, op, backend); + + quote! { + #cfg_attr + #pattern => { #call } + } +} + +/// Generate the body that unwraps the dispatch tensors, calls the backend's trait impl and wraps the +/// result back into a [`DispatchTensor`]. Shared by [`gen_backend_arm`] (inside a match) and the +/// no-tensor-input dispatch path (direct call). +fn gen_backend_call(ir: &Extension, op: &Operation, backend: &Backend) -> TokenStream2 { + let b_ident = backend_to_ident(backend); + let trait_name = &ir.trait_name; + let fn_name = &op.name; + + // Unwrap inner kind: lhs.float(), rhs.int(), etc. (no-op when there are no tensor inputs). + let unwraps = op.inputs.iter().filter_map(|a| match &a.kind { + ArgKind::Tensor(kind) => { + let name = &a.name; + let method = kind.unwrap_method(); + Some(quote! { let #name = #name.#method(); }) + } + _ => None, + }); + + let call_args = op.inputs.iter().map(|a| &a.name); + let maybe_await = if op.asyncness { + quote! { .await } + } else { + quote! {} + }; + + let wrap_out = match &op.output { + OperationOutput::Tensor(kind) => { + let wrapped = gen_tensor_wrap(kind, quote! { _out }, &b_ident, false); + quote! { burn::backend::DispatchTensor { kind: #wrapped, checkpointing } } + } + OperationOutput::Tuple(elems) => { + let elements = elems.iter().enumerate().map(|(i, elem)| { + let idx = syn::Index::from(i); + match elem { + OutputKind::Tensor(kind) => { + let wrapped = gen_tensor_wrap(kind, quote! { _out.#idx }, &b_ident, false); + quote! { burn::backend::DispatchTensor { kind: #wrapped, checkpointing } } + } + OutputKind::Custom(_) => { + quote! { + burn::backend::ExtensionType::map_type( + _out.#idx, + |tensor| burn::backend::DispatchTensorKind::#b_ident(tensor), + checkpointing, + ) + } + } + } + }); + quote! { (#(#elements),*) } + } + OperationOutput::Custom(_) => { + quote! { burn::backend::ExtensionType::map_type(_out, |tensor| burn::backend::DispatchTensorKind::#b_ident(tensor), checkpointing) } + } + }; + + quote! { + #(#unwraps)* + let _out = <#b_ident as #trait_name>::#fn_name(#(#call_args),*)#maybe_await; + #wrap_out + } +} + +fn gen_autodiff_arm(ir: &Extension, op: &Operation) -> TokenStream2 { + let trait_name = &ir.trait_name; + let fn_name = &op.name; + + // Filter for tensor arguments only + let tensor_args: Vec<_> = op + .inputs + .iter() + .filter_map(|a| match &a.kind { + ArgKind::Tensor(k) => Some((&a.name, k)), + _ => None, + }) + .collect(); + + // Build the match pattern for the same backend wrapped by autodiff + let inner_arms = ir.backends.concrete.iter().map(|backend| { + let cfg_attr = backend.cfg.as_ref().map(|meta| quote! { #[#meta] }); + let b_ident = backend_to_ident(backend); + + let pattern = if tensor_args.len() == 1 { + let (name, kind) = tensor_args[0]; + + if *kind == TensorKind::Float { + quote! { + burn::backend::DispatchTensorKind::Autodiff(#name) + } + } else { + quote! { + burn::backend::DispatchTensorKind::#b_ident(#name) + } + } + } else { + let pats = tensor_args.iter().map(|(name, kind)| { + if **kind == TensorKind::Float { + quote! { + burn::backend::DispatchTensorKind::Autodiff(#name) + } + } else { + quote! { + burn::backend::DispatchTensorKind::#b_ident(#name) + } + } + }); + + quote! { (#(#pats),*) } + }; + + let unwraps = tensor_args.iter().map(|(name, kind)| { + if **kind == TensorKind::Float { + quote! { + let #name = match *#name { + burn::backend::DispatchTensorKind::#b_ident(t) => t.autodiff(), + _ => unreachable!("Autodiff backend mismatch"), + }; + } + } else { + let method = kind.unwrap_method(); + quote! { + let #name = #name.#method(); + } + } + }); + + let call_args = op.inputs.iter().map(|a| &a.name); + let maybe_await = if op.asyncness { + quote! { .await } + } else { + quote! {} + }; + + let wrap_out = match &op.output { + OperationOutput::Tensor(kind) => { + let wrapped = gen_tensor_wrap(kind, quote! { _out }, &b_ident, true); + quote! { burn::backend::DispatchTensor { kind: #wrapped, checkpointing } } + } + OperationOutput::Tuple(elems) => { + let elements = elems.iter().enumerate().map(|(i, elem)| { + let idx = syn::Index::from(i); + match elem { + OutputKind::Tensor(kind) => { + let wrapped = gen_tensor_wrap(kind, quote! { _out.#idx }, &b_ident, true); + quote! { burn::backend::DispatchTensor { kind: #wrapped, checkpointing } } + } + OutputKind::Custom(_) => { + quote! { + burn::backend::ExtensionType::map_type( + _out.#idx, + |tensor| match tensor { + burn::backend::BackendTensor::Float(t) => { + burn::backend::DispatchTensorKind::Autodiff( + Box::new(burn::backend::DispatchTensorKind::#b_ident( + burn::backend::BackendTensor::Autodiff(t) + )) + ) + } + _ => burn::backend::DispatchTensorKind::#b_ident(tensor), + }, + checkpointing, + ) + } + } + } + }); + quote! { (#(#elements),*) } + } + OperationOutput::Custom(_) => { + quote! { burn::backend::ExtensionType::map_type(_out, |tensor| match tensor { + burn::backend::BackendTensor::Float(t) => { + burn::backend::DispatchTensorKind::Autodiff( + Box::new(burn::backend::DispatchTensorKind::#b_ident( + burn::backend::BackendTensor::Autodiff(t) + )) + ) + } + }, checkpointing) } + } + }; + + quote! { + #cfg_attr + #pattern => { + #(#unwraps)* + type _ADBackend = Autodiff<#b_ident>; + let _out = <_ADBackend as #trait_name>::#fn_name(#(#call_args),*)#maybe_await; + #wrap_out + } + } + }); + + quote! { + #( #inner_arms )* + } +} + +fn gen_tensor_wrap( + kind: &TensorKind, + val: TokenStream2, + b_ident: &Ident, + is_ad: bool, +) -> TokenStream2 { + let variant = kind.variant(); + if is_ad && *kind == TensorKind::Float { + quote! { + burn::backend::DispatchTensorKind::Autodiff( + Box::new(burn::backend::DispatchTensorKind::#b_ident( + burn::backend::BackendTensor::Autodiff(#val) + )) + ) + } + } else { + quote! { + burn::backend::DispatchTensorKind::#b_ident( + burn::backend::BackendTensor::#variant(#val) + ) + } + } +} + +/// Derive macro to implement `ExtensionType` for custom structures returned by backend extensions. +/// +/// When a custom backend extension operation needs to return multiple tensors or a mix of tensors +/// and metadata (instead of a single tensor primitive or container of primitives), this macro automates +/// the process of wrapping the tensor primitives so they can cross the boundary into the `Dispatch` backend. +/// +/// # Requirements +/// +/// - The struct must have named fields. +/// - The struct must be generic over a `Backend` type parameter (`B`). +/// +/// # Field Attributes +/// +/// By default, the macro inspects the type of each field: +/// - **Tensor Primitives**: Automatically mapped. +/// - **Other types**: Passed through unmodified. +/// +/// To nest another custom struct that also implements `ExtensionType`, you must annotate it with +/// `#[extension_type]` to tell the macro to traverse it recursively. +/// +/// # Example +/// +/// ```rust,ignore +/// #[derive(ExtensionType)] +/// pub struct OperationOutput { +/// pub bool: BoolTensor, +/// pub int: IntTensor, +/// pub float: FloaTensor, +/// pub count: usize, // Non-tensor field passes through automatically +/// } +/// ``` +#[proc_macro_derive(ExtensionType, attributes(extension_type))] +pub fn derive_extension_output(input: TokenStream) -> TokenStream { + let input = parse_macro_input!(input as ItemStruct); + let name = &input.ident; + + // Extract generics for the trait implementation block + let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl(); + let fields_wrap = match &input.fields { + syn::Fields::Named(fields) => { + let field_mappings = fields.named.iter().map(|f| { + let f_name = &f.ident; + let is_ext = f.attrs.iter().any(|attr| attr.path().is_ident("extension_type")); + + if is_ext { + // It's a nested extension type + quote! { + #f_name: self.#f_name.map_type( + &map_kind, + checkpointing + ), + } + } + else if let Some(tensor_kind) = TensorKind::from_type(&f.ty) { + // Tensor primitive + let variant_ident = tensor_kind.variant(); + quote! { + #f_name: burn::backend::DispatchTensor { + kind: map_kind(burn::backend::BackendTensor::#variant_ident(self.#f_name)), + checkpointing, + }, + } + } else { + // Passthrough + quote! { #f_name: self.#f_name, } + } + }); + + quote! { #name { #( #field_mappings )* } } + } + _ => panic!("ExtensionType derive only supports structs with named fields"), + }; + + TokenStream::from(quote! { + impl #impl_generics burn::backend::ExtensionType for #name #ty_generics #where_clause { + type Target = #name; + + fn map_type( + self, + map_kind: F, + checkpointing: Option, + ) -> Self::Target + where + F: Fn(burn::backend::BackendTensor) -> burn::backend::DispatchTensorKind, + { + #fields_wrap + } + } + }) +} diff --git a/crates/burn-backend-tests/.cargo/config.toml b/crates/burn-backend-tests/.cargo/config.toml new file mode 100644 index 0000000..6379268 --- /dev/null +++ b/crates/burn-backend-tests/.cargo/config.toml @@ -0,0 +1,31 @@ +[alias] +test-cpu = "test --release --no-default-features --features cpu,std,fusion" +test-cuda = "test --release --no-default-features --features cuda,std,fusion" +test-flex = "test --release --no-default-features --features flex,std" +test-ndarray = "test --release --no-default-features --features ndarray,std" +test-rocm = "test --release --no-default-features --features rocm,std,fusion" +test-remote = "test --release --no-default-features --features remote,std,fusion" +test-tch = "test --release --no-default-features --features tch,std" +test-wgpu = "test --release --no-default-features --features wgpu,std,fusion" +test-vulkan = "test --release --no-default-features --features vulkan,std,fusion" +test-metal = "test --release --no-default-features --features metal,std,fusion" + +# Fusion disabled tests +test-cpu-nofuse = "test --release --no-default-features --features cpu,std" +test-cuda-nofuse = "test --release --no-default-features --features cuda,std" +test-rocm-nofuse = "test --release --no-default-features --features rocm,std" +test-wgpu-nofuse = "test --release --no-default-features --features wgpu,std" +test-vulkan-nofuse = "test --release --no-default-features --features vulkan,std" +test-metal-nofuse = "test --release --no-default-features --features metal,std" +test-remote-nofuse = "test --release --no-default-features --features remote,std" + +# Benches +bench-cpu = "bench --no-default-features --features cpu,std,fusion" +bench-cuda = "bench --no-default-features --features cuda,std,fusion" +bench-flex = "bench --no-default-features --features flex,std" +bench-ndarray = "bench --no-default-features --features ndarray,std" +bench-rocm = "bench --no-default-features --features rocm,std,fusion" +bench-tch = "bench --no-default-features --features tch,std" +bench-wgpu = "bench --no-default-features --features wgpu,std,fusion" +bench-vulkan = "bench --no-default-features --features vulkan,std,fusion" +bench-metal = "bench --no-default-features --features metal,std,fusion" diff --git a/crates/burn-backend-tests/Cargo.toml b/crates/burn-backend-tests/Cargo.toml new file mode 100644 index 0000000..257a577 --- /dev/null +++ b/crates/burn-backend-tests/Cargo.toml @@ -0,0 +1,195 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science", "no-std", "embedded", "wasm"] +description = "Tensor tests for Burn backends" +documentation = "https://docs.rs/burn-backend-tests" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "tensor", "pytorch", "ndarray"] +license.workspace = true +name = "burn-backend-tests" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-backend-tests" +version.workspace = true + +[lints] +workspace = true + +[lib] +proc-macro = true + +[features] +default = [ + "burn-tensor/default", + # Default + # "ndarray", + "std", +] +std = ["burn-tensor/std"] + +# Backends +cuda = ["burn-tensor/cuda", "quantization", "cube"] +rocm = ["burn-tensor/rocm", "quantization", "cube"] +flex = ["burn-tensor/flex", "quantization"] +ndarray = ["burn-tensor/ndarray", "quantization"] +tch = ["burn-tensor/tch"] +vulkan = ["wgpu", "burn-tensor/vulkan"] +webgpu = ["wgpu", "burn-tensor/webgpu"] +metal = ["wgpu", "burn-tensor/metal"] +wgpu = ["burn-tensor/wgpu", "quantization", "cube"] +cpu = ["burn-tensor/cpu", "cube"] +# Runs the test suite against a remote `burn-remote` server. +# At runtime, set `BURN_DEVICE=remote` (and optionally `BURN_REMOTE_ADDRESS=ws://host:port`). +remote = ["burn-tensor/remote"] + +autotune = ["burn-tensor/autotune"] +autotune-checks = ["burn-tensor/autotune-checks"] +memory-checks = ["burn-fusion?/memory-checks"] +fusion = [ + "burn-tensor/fusion", + "burn-fusion/test-util", + "burn-cubecl-fusion?/test-util", +] + +# CubeCL backends +cube = [ + "cubek", + "autotune", + "burn-tensor/ndarray", + "burn-ndarray", + # Not used directly, but included to avoid recompiling with and without fusion + "burn-cubecl", + # Direct dependency so fusion tests can assert launch-level decisions (in-place aliasing). + "burn-cubecl-fusion", +] + +# Test configs +quantization = [] +# Enables the collective (all-reduce) tests, which require a CUDA backend built with NCCL. +# Kept separate from `cuda` so the rest of the suite can run on CUDA without NCCL. +distributed = ["cuda"] + +[dependencies] +burn-tensor = { workspace = true, features = ["autodiff"] } +burn-backend = { workspace = true } + +# For random tests +cubek = { workspace = true, features = ["random"], optional = true } + +# Explicitly enable test-only checks +burn-autodiff = { workspace = true, features = ["export_tests"] } +burn-ndarray = { workspace = true, optional = true, features = [ + "export_tests", +] } + +# Manually enable so running the tests again with fusion feature doesn't have to compile other crates +burn-cubecl = { workspace = true, optional = true, features = [ + "default", + "fusion", +] } +burn-cubecl-fusion = { workspace = true, optional = true } + +# Enabled alongside the `fusion` feature so spy-based tests can import the spy API. +burn-fusion = { workspace = true, optional = true } + +num-traits = { workspace = true } +serial_test = { workspace = true } +rand.workspace = true + +# might-panic macro +proc-macro2 = { workspace = true } +quote = { workspace = true } +syn = { workspace = true } + + +[dev-dependencies] +# Used to initialize device settings (test dtypes) +ctor = "0.10.1" +divan = "0.1" + +[[bench]] +name = "attention" +harness = false + +[[bench]] +name = "binary_ops" +harness = false + +[[bench]] +name = "cat_max_min_ops" +harness = false + +[[bench]] +name = "comparison_ops" +harness = false + +[[bench]] +name = "conv_ops" +harness = false + +[[bench]] +name = "conv_transpose_ops" +harness = false + +[[bench]] +name = "cross_unfold_ops" +harness = false + +[[bench]] +name = "cumulative_ops" +harness = false + +[[bench]] +name = "default_ops" +harness = false + +[[bench]] +name = "deform_conv_ops" +harness = false + +[[bench]] +name = "fft_ops" +harness = false + +[[bench]] +name = "gather_scatter_ops" +harness = false + +[[bench]] +name = "int_ops" +harness = false + +[[bench]] +name = "interpolate_ops" +harness = false + +[[bench]] +name = "matmul" +harness = false + +[[bench]] +name = "norm_ops" +harness = false + +[[bench]] +name = "pool_ops" +harness = false + +[[bench]] +name = "quantization_ops" +harness = false + +[[bench]] +name = "reduce_ops" +harness = false + +[[bench]] +name = "slice_ops" +harness = false + +[[bench]] +name = "softmax_ops" +harness = false + +[[bench]] +name = "unary_ops" +harness = false diff --git a/crates/burn-backend-tests/README.md b/crates/burn-backend-tests/README.md new file mode 100644 index 0000000..a9167cb --- /dev/null +++ b/crates/burn-backend-tests/README.md @@ -0,0 +1,187 @@ +# Burn Backend Tests + +This crate provides a comprehensive suite of tests for Burn backends, covering: + +- Tensor operations: [tests/tensor/](./tests/tensor/) +- Autodiff: [tests/autodiff/](./tests/autodiff/) +- (Optional) CubeCL kernels correctness: [tests/cubecl/](./tests/cubecl/) +- (Optional) Fusion correctness: [tests/fusion/](./tests/fusion/) +- Benchmarks: [benches/](./benches/) + +## Running Tests + +The test backend device is selected via feature flags. Use the provided shorthand commands for +convenience: + +```sh +# Cpu +cargo test-cpu +# Cuda +cargo test-cuda +# Rocm +cargo test-rocm +# Wgpu / WebGpu +cargo test-wgpu +# Vulkan +cargo test-vulkan +# Metal +cargo test-metal + +# NdArray +cargo test-ndarray +# LibTorch +cargo test-tch +``` + +By default, `cargo test` fail-fast across integration test binaries. When one integration test +binary fails, Cargo does not run the remaining test binaries. If you want to run all test binaries +regardless of failures, pass `--no-fail-fast`, for example: + +```sh +cargo test-cuda --no-fail-fast +``` + +> [!NOTE] +> CubeCL-based backends are tested with `fusion` by default. If you want to run the tests without +> fusion, just append `-nofuse` to the cargo command. For example: +> +> ```sh +> cargo test-cuda-nofuse +> ``` + +## Structure + +- `tests/tensor.rs`: Tensor tests +- `tests/tensor_f16.rs`: F16 tensor tests +- `tests/autodiff.rs`: Autodiff tests +- `tests/autodiff_f16.rs`: F16 autodiff tests +- `tests/cubecl.rs`: CubeCL kernel tests + +### Common Modules + +- `common/backend.rs`: Backend type definitions +- `common/tensor.rs`: Reusable tensor test suite, split across float, int and bool tensor kinds +- `common/autodiff.rs`: Reusable autodiff test suite, with and without checkpointing + +### Test Reusability + +This crate uses a pattern of parameterized test modules to run the same tests with different +configurations: + +1. **Element types**: Each test scope declares `FloatElem` and `IntElem`, which determine the + precision for the current test run +1. **`#[path = "..."]` references shared modules**: Points to test files outside the normal module + hierarchy, e.g. `"common/tensor.rs"` +1. **`include!()` imports test code**: Test modules are included multiple times with different type + configurations +1. **`use super::*;`** propagates types down the module tree: Each level re-exports parent types so + deeply nested tests have access to the configured types + +For example, all tensor tests are included in `tensor.rs` to execute with the default element types +(`FloatElem = f32`, `IntElem = i32`): + +```rust +#[path = "common/tensor.rs"] +mod tensor; +``` + +For f16 tests, only the float tensor tests are included to validate behavior with `FloatElem = f16`: + +```rust +#[path = "tensor/float/mod.rs"] +mod f16; +``` + +> [!WARNING] +> Tests for different data types (e.g., f32 vs f16) must be isolated, as each device's settings are +> global and can only be initialized once per process. + +## Running Benchmarks + +Benchmarks use [divan](https://docs.rs/divan) and run against the feature-selected backend via +`TestBackend = burn_dispatch::Dispatch`, the same way tests do. Shorthand aliases mirror the test +ones: + +```sh +# Cpu +cargo bench-cpu +# Cuda +cargo bench-cuda +# Rocm +cargo bench-rocm +# Wgpu / WebGpu +cargo bench-wgpu +# Vulkan +cargo bench-vulkan +# Metal +cargo bench-metal + +# Flex +cargo bench-flex +# NdArray +cargo bench-ndarray +# LibTorch +cargo bench-tch +``` + +Run a single bench file by passing `--bench `: + +```sh +cargo bench-flex --bench matmul +cargo bench-cuda --bench attention +``` + +Filter benches within a file by passing an argument after `--`: + +```sh +cargo bench-flex --bench matmul -- square +``` + +### Comparing backends + +Each run targets one backend. To compare (e.g. Flex vs NdArray), run twice and diff the output: + +```sh +cargo bench-flex --bench matmul > flex.txt +cargo bench-ndarray --bench matmul > ndarray.txt +diff flex.txt ndarray.txt +``` + +### GPU sync + +GPU backends (cuda, wgpu, rocm, metal, vulkan) dispatch ops asynchronously. Each bench wraps its op +with `bencher.bench_synced(...)` (defined in [benches/common/mod.rs](./benches/common/mod.rs)), +which inserts a `Backend::sync` call before returning - so the timed region covers actual execution +rather than dispatch latency. On CPU backends `sync` is a no-op via the `Backend::sync` default, so +this costs nothing there. + +When adding a new bench, use `bench_synced` rather than divan's raw `bench`: + +```rust +use common::BencherExt; + +#[divan::bench] +fn my_op(bencher: Bencher) { + let x = make_tensor(SIZE); + bencher.bench_synced(|| some_op(x.clone())); +} +``` + +## Adding New Tests + +Add test modules under `tests/tensor/`, `tests/autodiff/`, `tests/cubecl` and `tests/fusion` +respectively. They will automatically run for all required configurations. + +For tensor tests, make sure to add the test to each relevant tensor kind: + +- `tensor/bool`: boolean tensor tests +- `tensor/float`: float tensor tests +- `tensor/int`: integer tensor tests + +**Guidelines:** + +Import types with `use super::*;` at the top of each module to use the `FloatElem` and `IntElem` types. + +For autodiff tests, always use `AutodiffDevice::new()` to create the device. Autodiff is enabled on +the device itself when using the `Dispatch` test backend, ensuring the device supports automatic +differentiation without modifying the backend type. diff --git a/crates/burn-backend-tests/benches/attention.rs b/crates/burn-backend-tests/benches/attention.rs new file mode 100644 index 0000000..f8fc79d --- /dev/null +++ b/crates/burn-backend-tests/benches/attention.rs @@ -0,0 +1,271 @@ +//! Scaled dot-product attention benchmarks. +//! +//! Run with: +//! ```bash +//! cargo bench --bench attention +//! ``` + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::module::attention; +use burn_tensor::ops::AttentionModuleOptions; +use burn_tensor::{Tensor, TensorData}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Attention Benchmarks"); + println!("Memory allocation tracking enabled"); + println!(); + divan::main(); + common::report_failures(); +} + +/// Create Q, K, V tensors for attention benchmarking. +fn make_qkv( + batch: usize, + heads: usize, + seq_q: usize, + seq_k: usize, + head_dim: usize, +) -> (Tensor<4>, Tensor<4>, Tensor<4>) { + let dev = Default::default(); + let make = |rows: usize, cols: usize| -> Tensor<4> { + let n = batch * heads * rows * cols; + let data: Vec = (0..n).map(|i| ((i % 1000) as f32 / 1000.0) - 0.5).collect(); + Tensor::from_data(TensorData::new(data, [batch, heads, rows, cols]), &dev) + }; + ( + make(seq_q, head_dim), + make(seq_k, head_dim), + make(seq_k, head_dim), + ) +} + +/// Create an additive bias tensor [batch, heads, seq_q, seq_k]. +fn make_bias(batch: usize, heads: usize, seq_q: usize, seq_k: usize) -> Tensor<4> { + let n = batch * heads * seq_q * seq_k; + let data: Vec = (0..n).map(|i| ((i % 500) as f32 / 500.0) - 0.5).collect(); + Tensor::from_data( + TensorData::new(data, [batch, heads, seq_q, seq_k]), + &Default::default(), + ) +} + +macro_rules! bench_attention { + ($mod_name:ident, $name:literal) => { + #[divan::bench_group(name = $name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "self_attention")] + mod self_attention { + use super::*; + + #[divan::bench] + fn b1_h8_s64_d64(bencher: Bencher) { + let (q, k, v) = make_qkv(1, 8, 64, 64, 64); + bencher.bench_synced(|| { + attention( + q.clone(), + k.clone(), + v.clone(), + None, + None, + Default::default(), + ) + }); + } + + #[divan::bench] + fn b1_h12_s128_d64(bencher: Bencher) { + let (q, k, v) = make_qkv(1, 12, 128, 128, 64); + bencher.bench_synced(|| { + attention( + q.clone(), + k.clone(), + v.clone(), + None, + None, + Default::default(), + ) + }); + } + + #[divan::bench] + fn b1_h12_s256_d64(bencher: Bencher) { + let (q, k, v) = make_qkv(1, 12, 256, 256, 64); + bencher.bench_synced(|| { + attention( + q.clone(), + k.clone(), + v.clone(), + None, + None, + Default::default(), + ) + }); + } + + #[divan::bench] + fn b1_h12_s512_d64(bencher: Bencher) { + let (q, k, v) = make_qkv(1, 12, 512, 512, 64); + bencher.bench_synced(|| { + attention( + q.clone(), + k.clone(), + v.clone(), + None, + None, + Default::default(), + ) + }); + } + + #[divan::bench] + fn b1_h32_s256_d128(bencher: Bencher) { + let (q, k, v) = make_qkv(1, 32, 256, 256, 128); + bencher.bench_synced(|| { + attention( + q.clone(), + k.clone(), + v.clone(), + None, + None, + Default::default(), + ) + }); + } + + #[divan::bench] + fn b4_h12_s128_d64(bencher: Bencher) { + let (q, k, v) = make_qkv(4, 12, 128, 128, 64); + bencher.bench_synced(|| { + attention( + q.clone(), + k.clone(), + v.clone(), + None, + None, + Default::default(), + ) + }); + } + } + + #[divan::bench_group(name = "causal")] + mod causal { + use super::*; + + fn causal_opts() -> AttentionModuleOptions { + AttentionModuleOptions { + is_causal: true, + ..Default::default() + } + } + + #[divan::bench] + fn b1_h12_s128_d64(bencher: Bencher) { + let (q, k, v) = make_qkv(1, 12, 128, 128, 64); + bencher.bench_synced(|| { + attention(q.clone(), k.clone(), v.clone(), None, None, causal_opts()) + }); + } + + #[divan::bench] + fn b1_h12_s256_d64(bencher: Bencher) { + let (q, k, v) = make_qkv(1, 12, 256, 256, 64); + bencher.bench_synced(|| { + attention(q.clone(), k.clone(), v.clone(), None, None, causal_opts()) + }); + } + + #[divan::bench] + fn b1_h12_s512_d64(bencher: Bencher) { + let (q, k, v) = make_qkv(1, 12, 512, 512, 64); + bencher.bench_synced(|| { + attention(q.clone(), k.clone(), v.clone(), None, None, causal_opts()) + }); + } + } + + #[divan::bench_group(name = "with_bias")] + mod with_bias { + use super::*; + + #[divan::bench] + fn b1_h12_s128_d64(bencher: Bencher) { + let (q, k, v) = make_qkv(1, 12, 128, 128, 64); + let bias = make_bias(1, 12, 128, 128); + bencher.bench_synced(|| { + attention( + q.clone(), + k.clone(), + v.clone(), + None, + Some(bias.clone()), + Default::default(), + ) + }); + } + + #[divan::bench] + fn b1_h12_s256_d64(bencher: Bencher) { + let (q, k, v) = make_qkv(1, 12, 256, 256, 64); + let bias = make_bias(1, 12, 256, 256); + bencher.bench_synced(|| { + attention( + q.clone(), + k.clone(), + v.clone(), + None, + Some(bias.clone()), + Default::default(), + ) + }); + } + } + + #[divan::bench_group(name = "cross_attention")] + mod cross_attention { + use super::*; + + #[divan::bench] + fn b1_h12_sq128_sk512_d64(bencher: Bencher) { + let (q, k, v) = make_qkv(1, 12, 128, 512, 64); + bencher.bench_synced(|| { + attention( + q.clone(), + k.clone(), + v.clone(), + None, + None, + Default::default(), + ) + }); + } + + #[divan::bench] + fn b1_h12_sq32_sk1024_d64(bencher: Bencher) { + let (q, k, v) = make_qkv(1, 12, 32, 1024, 64); + bencher.bench_synced(|| { + attention( + q.clone(), + k.clone(), + v.clone(), + None, + None, + Default::default(), + ) + }); + } + } + } + }; +} + +bench_attention!(backend, "backend"); diff --git a/crates/burn-backend-tests/benches/binary_ops.rs b/crates/burn-backend-tests/benches/binary_ops.rs new file mode 100644 index 0000000..cc09fe2 --- /dev/null +++ b/crates/burn-backend-tests/benches/binary_ops.rs @@ -0,0 +1,390 @@ +//! Benchmarks for binary operations. +//! +//! Run with: +//! ```bash +//! cargo bench --bench binary_ops +//! ``` + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{Int, Tensor, TensorData}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Benchmarks"); + println!(); + divan::main(); + common::report_failures(); +} + +// Tensor sizes for benchmarking +const SMALL: usize = 64 * 64; // 4K elements +const MEDIUM: usize = 256 * 256; // 64K elements +const LARGE: usize = 1024 * 1024; // 1M elements + +fn make_tensor(size: usize) -> Tensor<1> { + let data: Vec = (0..size).map(|i| (i % 1000) as f32 / 1000.0).collect(); + Tensor::from_data(TensorData::new(data, [size]), &Default::default()) +} + +fn make_tensor_2d(rows: usize, cols: usize) -> Tensor<2> { + let data: Vec = (0..rows * cols) + .map(|i| (i % 1000) as f32 / 1000.0) + .collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) +} + +fn make_int_tensor(size: usize) -> Option> { + common::try_setup(|| { + let data: Vec = (0..size).map(|i| (i % 1000) as i32).collect(); + Tensor::from_data(TensorData::new(data, [size]), &Default::default()) + }) +} + +fn make_int_tensor_2d(rows: usize, cols: usize) -> Option> { + common::try_setup(|| { + let data: Vec = (0..rows * cols).map(|i| (i % 1000) as i32).collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) + }) +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + // Add operations + #[divan::bench_group(name = "add")] + mod add { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let a = make_tensor(SMALL); + let b = make_tensor(SMALL); + bencher.bench_synced(|| a.clone() + b.clone()); + } + + #[divan::bench] + fn medium(bencher: Bencher) { + let a = make_tensor(MEDIUM); + let b = make_tensor(MEDIUM); + bencher.bench_synced(|| a.clone() + b.clone()); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + let b = make_tensor(LARGE); + bencher.bench_synced(|| a.clone() + b.clone()); + } + } + + // Mul operations + #[divan::bench_group(name = "mul")] + mod mul { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let a = make_tensor(SMALL); + let b = make_tensor(SMALL); + bencher.bench_synced(|| a.clone() * b.clone()); + } + + #[divan::bench] + fn medium(bencher: Bencher) { + let a = make_tensor(MEDIUM); + let b = make_tensor(MEDIUM); + bencher.bench_synced(|| a.clone() * b.clone()); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + let b = make_tensor(LARGE); + bencher.bench_synced(|| a.clone() * b.clone()); + } + } + + // Div operations + #[divan::bench_group(name = "div")] + mod div { + use super::*; + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + let b = make_tensor(LARGE); + bencher.bench_synced(|| a.clone() / b.clone()); + } + } + + // Transposed (non-contiguous) + #[divan::bench_group(name = "add_transposed")] + mod add_transposed { + use super::*; + + #[divan::bench] + fn medium_256x256(bencher: Bencher) { + let a = make_tensor_2d(256, 256).transpose(); + let b = make_tensor_2d(256, 256); + bencher.bench_synced(|| a.clone() + b.clone()); + } + + #[divan::bench] + fn large_1024x1024(bencher: Bencher) { + let a = make_tensor_2d(1024, 1024).transpose(); + let b = make_tensor_2d(1024, 1024); + bencher.bench_synced(|| a.clone() + b.clone()); + } + } + + // Scalar operations + #[divan::bench_group(name = "scalar")] + mod scalar { + use super::*; + + #[divan::bench] + fn add_large(bencher: Bencher) { + let a = make_tensor(LARGE); + bencher.bench_synced(|| a.clone() + 1.5); + } + + #[divan::bench] + fn mul_large(bencher: Bencher) { + let a = make_tensor(LARGE); + bencher.bench_synced(|| a.clone() * 2.0); + } + } + + // Powf operations + #[divan::bench_group(name = "powf")] + mod powf { + use super::*; + + #[divan::bench] + fn medium(bencher: Bencher) { + let a = make_tensor(MEDIUM); + let b = make_tensor(MEDIUM); + bencher.bench_synced(|| a.clone().powf(b.clone())); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + let b = make_tensor(LARGE); + bencher.bench_synced(|| a.clone().powf(b.clone())); + } + + #[divan::bench] + fn scalar_large(bencher: Bencher) { + let a = make_tensor(LARGE); + bencher.bench_synced(|| a.clone().powf_scalar(2.5)); + } + } + + // Atan2 operations + #[divan::bench_group(name = "atan2")] + mod atan2 { + use super::*; + + #[divan::bench] + fn medium(bencher: Bencher) { + let a = make_tensor(MEDIUM); + let b = make_tensor(MEDIUM); + bencher.bench_synced(|| a.clone().atan2(b.clone())); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + let b = make_tensor(LARGE); + bencher.bench_synced(|| a.clone().atan2(b.clone())); + } + } + } + }; +} + +bench_backend!(backend, "backend"); + +// Int benchmarks +macro_rules! bench_int_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "int_add")] + mod int_add { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let Some(a) = make_int_tensor(SMALL) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_int_tensor(SMALL) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone() + b.clone()); + } + + #[divan::bench] + fn medium(bencher: Bencher) { + let Some(a) = make_int_tensor(MEDIUM) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_int_tensor(MEDIUM) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone() + b.clone()); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let Some(a) = make_int_tensor(LARGE) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_int_tensor(LARGE) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone() + b.clone()); + } + } + + #[divan::bench_group(name = "int_mul")] + mod int_mul { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let Some(a) = make_int_tensor(SMALL) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_int_tensor(SMALL) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone() * b.clone()); + } + + #[divan::bench] + fn medium(bencher: Bencher) { + let Some(a) = make_int_tensor(MEDIUM) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_int_tensor(MEDIUM) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone() * b.clone()); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let Some(a) = make_int_tensor(LARGE) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_int_tensor(LARGE) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone() * b.clone()); + } + } + + #[divan::bench_group(name = "int_div")] + mod int_div { + use super::*; + + #[divan::bench] + fn large(bencher: Bencher) { + let Some(a) = make_int_tensor(LARGE) else { + bencher.bench(|| ()); + return; + }; + // Avoid division by zero + let data: Vec = (0..LARGE).map(|i| (i % 999) as i32 + 1).collect(); + let b: Tensor<1, Int> = + Tensor::from_data(TensorData::new(data, [LARGE]), &Default::default()); + bencher.bench_synced(|| a.clone() / b.clone()); + } + } + + #[divan::bench_group(name = "int_add_transposed")] + mod int_add_transposed { + use super::*; + + #[divan::bench] + fn medium_256x256(bencher: Bencher) { + let Some(a) = make_int_tensor_2d(256, 256) else { + bencher.bench(|| ()); + return; + }; + let a = a.transpose(); + let Some(b) = make_int_tensor_2d(256, 256) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone() + b.clone()); + } + + #[divan::bench] + fn large_1024x1024(bencher: Bencher) { + let Some(a) = make_int_tensor_2d(1024, 1024) else { + bencher.bench(|| ()); + return; + }; + let a = a.transpose(); + let Some(b) = make_int_tensor_2d(1024, 1024) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone() + b.clone()); + } + } + + #[divan::bench_group(name = "int_scalar")] + mod int_scalar { + use super::*; + + #[divan::bench] + fn add_large(bencher: Bencher) { + let Some(a) = make_int_tensor(LARGE) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone() + 100); + } + + #[divan::bench] + fn mul_large(bencher: Bencher) { + let Some(a) = make_int_tensor(LARGE) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone() * 3); + } + } + } + }; +} + +bench_int_backend!(backend_int, "backend_int"); diff --git a/crates/burn-backend-tests/benches/cat_max_min_ops.rs b/crates/burn-backend-tests/benches/cat_max_min_ops.rs new file mode 100644 index 0000000..ff3f159 --- /dev/null +++ b/crates/burn-backend-tests/benches/cat_max_min_ops.rs @@ -0,0 +1,321 @@ +//! Benchmarks for cat, max, min, int power, bool select, and grid_sample_2d. +//! +//! Run with: +//! ```bash +//! cargo bench --bench cat_max_min_ops --features simd,rayon +//! ``` + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{Bool, Int, Tensor, TensorData, ops::GridSampleOptions}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Benchmarks for cat, max/min, int power, bool select, grid_sample_2d"); + println!("Memory allocation tracking enabled"); + println!(); + divan::main(); + common::report_failures(); +} + +// === Helpers === + +fn make_f32_1d(size: usize) -> Tensor<1> { + let data: Vec = (0..size).map(|i| (i % 1000) as f32 / 1000.0).collect(); + Tensor::from_data(TensorData::new(data, [size]), &Default::default()) +} + +fn make_f32_2d(rows: usize, cols: usize) -> Tensor<2> { + let data: Vec = (0..rows * cols) + .map(|i| (i % 1000) as f32 / 1000.0) + .collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) +} + +fn make_int_2d(rows: usize, cols: usize) -> Option> { + common::try_setup(|| { + let data: Vec = (0..rows * cols).map(|i| (i % 10) as i32 + 1).collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) + }) +} + +fn make_int_exp_2d(rows: usize, cols: usize) -> Option> { + common::try_setup(|| { + let data: Vec = (0..rows * cols).map(|i| (i % 4) as i32 + 1).collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) + }) +} + +fn make_bool_2d(rows: usize, cols: usize) -> Tensor<2, Bool> { + let data: Vec = (0..rows * cols).map(|i| i % 3 != 0).collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) +} + +fn make_indices_1d(size: usize, max_idx: usize) -> Option> { + common::try_setup(|| { + let data: Vec = (0..size).map(|i| (i % max_idx) as i32).collect(); + Tensor::from_data(TensorData::new(data, [size]), &Default::default()) + }) +} + +fn make_grid_4d(batch: usize, h_out: usize, w_out: usize) -> Tensor<4> { + let size = batch * h_out * w_out * 2; + // Grid values in [-1, 1] + let data: Vec = (0..size) + .map(|i| (i as f32 / size as f32) * 2.0 - 1.0) + .collect(); + Tensor::from_data( + TensorData::new(data, [batch, h_out, w_out, 2]), + &Default::default(), + ) +} + +fn make_image_4d(batch: usize, channels: usize, h: usize, w: usize) -> Tensor<4> { + let size = batch * channels * h * w; + let data: Vec = (0..size).map(|i| (i % 256) as f32 / 255.0).collect(); + Tensor::from_data( + TensorData::new(data, [batch, channels, h, w]), + &Default::default(), + ) +} + +// ============================================================================= +// Cat +// ============================================================================= + +macro_rules! bench_cat { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "cat")] + mod cat { + use super::*; + + // Cat along dim 0 (fast path: contiguous memcpy) + #[divan::bench] + fn dim0_4x_256x256(bencher: Bencher) { + let t = make_f32_2d(256, 256); + let tensors = vec![t.clone(), t.clone(), t.clone(), t.clone()]; + bencher.bench_synced(|| Tensor::cat(tensors.clone(), 0)); + } + + #[divan::bench] + fn dim0_4x_1024x256(bencher: Bencher) { + let t = make_f32_2d(1024, 256); + let tensors = vec![t.clone(), t.clone(), t.clone(), t.clone()]; + bencher.bench_synced(|| Tensor::cat(tensors.clone(), 0)); + } + + // Cat along dim 1 (general path) + #[divan::bench] + fn dim1_4x_256x64(bencher: Bencher) { + let t = make_f32_2d(256, 64); + let tensors = vec![t.clone(), t.clone(), t.clone(), t.clone()]; + bencher.bench_synced(|| Tensor::cat(tensors.clone(), 1)); + } + + #[divan::bench] + fn dim1_4x_1024x64(bencher: Bencher) { + let t = make_f32_2d(1024, 64); + let tensors = vec![t.clone(), t.clone(), t.clone(), t.clone()]; + bencher.bench_synced(|| Tensor::cat(tensors.clone(), 1)); + } + + // Many small tensors + #[divan::bench] + fn dim0_16x_64x64(bencher: Bencher) { + let t = make_f32_2d(64, 64); + let tensors: Vec<_> = (0..16).map(|_| t.clone()).collect(); + bencher.bench_synced(|| Tensor::cat(tensors.clone(), 0)); + } + + // 1D cat + #[divan::bench] + fn dim0_4x_16k_1d(bencher: Bencher) { + let t = make_f32_1d(16 * 1024); + let tensors = vec![t.clone(), t.clone(), t.clone(), t.clone()]; + bencher.bench_synced(|| Tensor::cat(tensors.clone(), 0)); + } + } + + #[divan::bench_group(name = "max")] + mod max { + use super::*; + + #[divan::bench] + fn s_1k(bencher: Bencher) { + let t = make_f32_1d(1024); + bencher.bench_synced(|| t.clone().max()); + } + + #[divan::bench] + fn s_64k(bencher: Bencher) { + let t = make_f32_1d(64 * 1024); + bencher.bench_synced(|| t.clone().max()); + } + + #[divan::bench] + fn s_1m(bencher: Bencher) { + let t = make_f32_1d(1024 * 1024); + bencher.bench_synced(|| t.clone().max()); + } + } + + #[divan::bench_group(name = "min")] + mod min { + use super::*; + + #[divan::bench] + fn s_1k(bencher: Bencher) { + let t = make_f32_1d(1024); + bencher.bench_synced(|| t.clone().min()); + } + + #[divan::bench] + fn s_64k(bencher: Bencher) { + let t = make_f32_1d(64 * 1024); + bencher.bench_synced(|| t.clone().min()); + } + + #[divan::bench] + fn s_1m(bencher: Bencher) { + let t = make_f32_1d(1024 * 1024); + bencher.bench_synced(|| t.clone().min()); + } + } + + #[divan::bench_group(name = "int_max")] + mod int_max { + use super::*; + + #[divan::bench] + fn s_256x256(bencher: Bencher) { + let Some(t) = make_int_2d(256, 256) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| t.clone().max()); + } + + #[divan::bench] + fn s_1024x1024(bencher: Bencher) { + let Some(t) = make_int_2d(1024, 1024) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| t.clone().max()); + } + } + + #[divan::bench_group(name = "int_powi")] + mod int_powi { + use super::*; + + #[divan::bench] + fn s_256x256(bencher: Bencher) { + let Some(base) = make_int_2d(256, 256) else { + bencher.bench(|| ()); + return; + }; + let Some(exp) = make_int_exp_2d(256, 256) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| base.clone().powi(exp.clone())); + } + + #[divan::bench] + fn s_1024x256(bencher: Bencher) { + let Some(base) = make_int_2d(1024, 256) else { + bencher.bench(|| ()); + return; + }; + let Some(exp) = make_int_exp_2d(1024, 256) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| base.clone().powi(exp.clone())); + } + } + + #[divan::bench_group(name = "bool_select")] + mod bool_select { + use super::*; + + #[divan::bench] + fn s_256x256_128idx(bencher: Bencher) { + let t = make_bool_2d(256, 256); + let Some(idx) = make_indices_1d(128, 256) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| t.clone().select(0, idx.clone())); + } + + #[divan::bench] + fn s_1024x256_512idx(bencher: Bencher) { + let t = make_bool_2d(1024, 256); + let Some(idx) = make_indices_1d(512, 1024) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| t.clone().select(0, idx.clone())); + } + } + + #[divan::bench_group(name = "grid_sample_2d")] + mod grid_sample { + use super::*; + + #[divan::bench] + fn s_b1_c3_32x32(bencher: Bencher) { + let img = make_image_4d(1, 3, 32, 32); + let grid = make_grid_4d(1, 32, 32); + bencher.bench_synced(|| { + img.clone() + .grid_sample_2d(grid.clone(), GridSampleOptions::default()) + }); + } + + #[divan::bench] + fn s_b1_c3_64x64(bencher: Bencher) { + let img = make_image_4d(1, 3, 64, 64); + let grid = make_grid_4d(1, 64, 64); + bencher.bench_synced(|| { + img.clone() + .grid_sample_2d(grid.clone(), GridSampleOptions::default()) + }); + } + + #[divan::bench] + fn s_b4_c3_32x32(bencher: Bencher) { + let img = make_image_4d(4, 3, 32, 32); + let grid = make_grid_4d(4, 32, 32); + bencher.bench_synced(|| { + img.clone() + .grid_sample_2d(grid.clone(), GridSampleOptions::default()) + }); + } + + #[divan::bench] + fn s_b1_c16_64x64(bencher: Bencher) { + let img = make_image_4d(1, 16, 64, 64); + let grid = make_grid_4d(1, 64, 64); + bencher.bench_synced(|| { + img.clone() + .grid_sample_2d(grid.clone(), GridSampleOptions::default()) + }); + } + } + } + }; +} + +bench_cat!(backend, "backend"); diff --git a/crates/burn-backend-tests/benches/common/mod.rs b/crates/burn-backend-tests/benches/common/mod.rs new file mode 100644 index 0000000..54b364a --- /dev/null +++ b/crates/burn-backend-tests/benches/common/mod.rs @@ -0,0 +1,162 @@ +//! Shared setup for benches. +//! +//! Mirrors `tests/common/backend.rs`: `Device::default()`, with a `#[ctor]` +//! that pins the default float/int dtypes to `f32`/`i32` so backends that advertise a different +//! default (e.g. bf16) don't silently convert bench inputs. + +use std::cell::Cell; +use std::panic::{self, AssertUnwindSafe, Location}; +use std::sync::Mutex; + +use burn_tensor::Element; +use ctor::ctor; + +pub type FloatElem = f32; +pub type IntElem = i32; + +#[ctor] +fn init_device_settings() { + let mut device = burn_tensor::Device::default(); + device + .configure( + burn_tensor::DeviceConfig::default() + .float_dtype(::dtype()) + .int_dtype(::dtype()), + ) + .unwrap(); +} + +/// Block until all outstanding ops on the default device complete. +/// +/// GPU backends (cuda, wgpu, rocm, metal, vulkan) dispatch ops asynchronously; without a sync +/// barrier inside the timed region a bench would measure dispatch latency, not execution time. +/// On CPU backends (flex, ndarray) this is a no-op via the `Backend::sync` default. +#[inline] +pub fn sync() { + burn_tensor::Device::default().sync().unwrap(); +} + +// --- Panic-tolerant bench execution ----------------------------------------- +// +// Some ops are not implemented on every backend (quantization on tch, deformable conv on metal, +// etc.). Without handling, the first panic in a bench binary aborts the whole run. We pre-flight +// each bench under `catch_unwind`; on panic we record the failure, skip the bench loop, and let +// subsequent benches in the same binary continue. `report_failures` at the end of `main()` +// prints a summary. + +thread_local! { + static SUPPRESS_PANIC_OUTPUT: Cell = const { Cell::new(false) }; +} + +struct BenchFailure { + location: String, + message: String, +} + +static FAILURES: Mutex> = Mutex::new(Vec::new()); + +#[ctor] +fn install_panic_hook() { + // Chain onto the existing hook rather than replacing it so truly unexpected panics (outside + // `bench_synced`'s catch_unwind window) still print normally. + let default_hook = panic::take_hook(); + panic::set_hook(Box::new(move |info| { + if !SUPPRESS_PANIC_OUTPUT.with(|s| s.get()) { + default_hook(info); + } + })); +} + +fn panic_message(payload: Box) -> String { + if let Some(s) = payload.downcast_ref::<&str>() { + s.to_string() + } else if let Some(s) = payload.downcast_ref::() { + s.clone() + } else { + "(non-string panic payload)".to_string() + } +} + +/// Run a setup closure tolerantly: if it panics (typically because the backend doesn't support +/// an op used during setup), returns `None` and records the failure location. Lets benches that +/// rely on op-heavy setup (e.g. `make_qtensor` calls `quantize_dynamic`, fft inverse benches call +/// `rfft` to produce the input) fall through to a no-op without taking down the whole binary. +#[allow(unused)] // it's used in benches +#[track_caller] +pub fn try_setup(f: impl FnOnce() -> T) -> Option { + let loc = Location::caller(); + SUPPRESS_PANIC_OUTPUT.with(|s| s.set(true)); + let r = panic::catch_unwind(AssertUnwindSafe(f)); + SUPPRESS_PANIC_OUTPUT.with(|s| s.set(false)); + match r { + Ok(v) => Some(v), + Err(payload) => { + FAILURES.lock().unwrap().push(BenchFailure { + location: format!("{}:{} (setup)", loc.file(), loc.line()), + message: panic_message(payload), + }); + None + } + } +} + +/// Print a summary of benches that panicked during this run, if any. Call from `main()` after +/// `divan::main()`. +pub fn report_failures() { + let failures = FAILURES.lock().unwrap(); + if failures.is_empty() { + return; + } + eprintln!(); + eprintln!("=== {} bench(es) skipped due to panic ===", failures.len()); + for f in failures.iter() { + eprintln!(" [{}] {}", f.location, f.message); + } +} + +/// Extension trait adding a synced, panic-tolerant variant of `Bencher::bench`. +/// +/// `bench_synced` runs the op then forces a device sync before returning, so the timed region +/// covers actual execution on async backends. If the op panics (typically "not implemented" on a +/// backend that doesn't support it), the failure is recorded and the bench is replaced with a +/// no-op so later benches in the same binary keep running. +pub trait BencherExt<'a, 'b> { + fn bench_synced(self, benched: F) + where + F: Fn() -> O + Sync; +} + +impl<'a, 'b> BencherExt<'a, 'b> for divan::Bencher<'a, 'b> { + #[track_caller] + fn bench_synced(self, benched: F) + where + F: Fn() -> O + Sync, + { + let loc = Location::caller(); + + SUPPRESS_PANIC_OUTPUT.with(|s| s.set(true)); + let result = panic::catch_unwind(AssertUnwindSafe(|| { + let r = benched(); + sync(); + drop(r); + })); + SUPPRESS_PANIC_OUTPUT.with(|s| s.set(false)); + + match result { + Ok(()) => { + self.bench(move || { + let r = benched(); + sync(); + r + }); + } + Err(payload) => { + FAILURES.lock().unwrap().push(BenchFailure { + location: format!("{}:{}", loc.file(), loc.line()), + message: panic_message(payload), + }); + self.bench(|| ()); + } + } + } +} diff --git a/crates/burn-backend-tests/benches/comparison_ops.rs b/crates/burn-backend-tests/benches/comparison_ops.rs new file mode 100644 index 0000000..97adb9e --- /dev/null +++ b/crates/burn-backend-tests/benches/comparison_ops.rs @@ -0,0 +1,214 @@ +//! Benchmarks for comparison and broadcasting operations. +//! +//! Run with: +//! ```bash +//! cargo bench --bench comparison_ops +//! ``` + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{Tensor, TensorData}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Benchmarking comparison and broadcast operations"); + println!(); + divan::main(); + common::report_failures(); +} + +const SMALL: usize = 64 * 64; // 4K elements +const MEDIUM: usize = 256 * 256; // 64K elements +const LARGE: usize = 1024 * 1024; // 1M elements + +fn make_tensor(size: usize) -> Tensor<1> { + let data: Vec = (0..size).map(|i| (i % 1000) as f32 / 1000.0).collect(); + Tensor::from_data(TensorData::new(data, [size]), &Default::default()) +} + +fn make_tensor_2d(rows: usize, cols: usize) -> Tensor<2> { + let data: Vec = (0..rows * cols) + .map(|i| (i % 1000) as f32 / 1000.0) + .collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + // Greater comparison + #[divan::bench_group(name = "greater")] + mod greater { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let a = make_tensor(SMALL); + let b = make_tensor(SMALL); + bencher.bench_synced(|| a.clone().greater(b.clone())); + } + + #[divan::bench] + fn medium(bencher: Bencher) { + let a = make_tensor(MEDIUM); + let b = make_tensor(MEDIUM); + bencher.bench_synced(|| a.clone().greater(b.clone())); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + let b = make_tensor(LARGE); + bencher.bench_synced(|| a.clone().greater(b.clone())); + } + } + + // Greater scalar comparison + #[divan::bench_group(name = "greater_elem")] + mod greater_elem { + use super::*; + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + bencher.bench_synced(|| a.clone().greater_elem(0.5)); + } + } + + // Equal comparison + #[divan::bench_group(name = "equal")] + mod equal { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let a = make_tensor(SMALL); + let b = make_tensor(SMALL); + bencher.bench_synced(|| a.clone().equal(b.clone())); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + let b = make_tensor(LARGE); + bencher.bench_synced(|| a.clone().equal(b.clone())); + } + } + + // Lower comparison + #[divan::bench_group(name = "lower")] + mod lower { + use super::*; + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + let b = make_tensor(LARGE); + bencher.bench_synced(|| a.clone().lower(b.clone())); + } + } + + // Non-contiguous (transposed) comparison + #[divan::bench_group(name = "greater_transposed")] + mod greater_transposed { + use super::*; + + #[divan::bench] + fn medium_256x256(bencher: Bencher) { + let a = make_tensor_2d(256, 256).transpose(); + let b = make_tensor_2d(256, 256); + bencher.bench_synced(|| a.clone().greater(b.clone())); + } + + #[divan::bench] + fn large_1024x1024(bencher: Bencher) { + let a = make_tensor_2d(1024, 1024).transpose(); + let b = make_tensor_2d(1024, 1024); + bencher.bench_synced(|| a.clone().greater(b.clone())); + } + } + + // Broadcast comparison: [N, 1] > [1, M] -> [N, M] + #[divan::bench_group(name = "greater_broadcast")] + mod greater_broadcast { + use super::*; + + #[divan::bench] + fn broadcast_256x256(bencher: Bencher) { + let a = make_tensor_2d(256, 1); + let b = make_tensor_2d(1, 256); + bencher.bench_synced(|| a.clone().greater(b.clone())); + } + + #[divan::bench] + fn broadcast_1024x1024(bencher: Bencher) { + let a = make_tensor_2d(1024, 1); + let b = make_tensor_2d(1, 1024); + bencher.bench_synced(|| a.clone().greater(b.clone())); + } + } + + // Expand operation + #[divan::bench_group(name = "expand")] + mod expand { + use super::*; + + #[divan::bench] + fn expand_1_to_1m(bencher: Bencher) { + let a = make_tensor_2d(1, 1); + bencher.bench_synced(|| a.clone().expand([1000, 1000])); + } + + #[divan::bench] + fn expand_row_1024x1_to_1024x1024(bencher: Bencher) { + let a = make_tensor_2d(1024, 1); + bencher.bench_synced(|| a.clone().expand([1024, 1024])); + } + + #[divan::bench] + fn expand_col_1x1024_to_1024x1024(bencher: Bencher) { + let a = make_tensor_2d(1, 1024); + bencher.bench_synced(|| a.clone().expand([1024, 1024])); + } + } + + // Bool operations + #[divan::bench_group(name = "bool_not")] + mod bool_not { + use super::*; + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + let b = make_tensor(LARGE); + let mask = a.greater(b); + bencher.bench_synced(|| mask.clone().bool_not()); + } + } + + #[divan::bench_group(name = "bool_and")] + mod bool_and { + use super::*; + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + let b = make_tensor(LARGE); + let mask1 = a.clone().greater(b.clone()); + let mask2 = a.lower(b); + bencher.bench_synced(|| mask1.clone().bool_and(mask2.clone())); + } + } + } + }; +} + +bench_backend!(backend, "backend"); diff --git a/crates/burn-backend-tests/benches/conv_ops.rs b/crates/burn-backend-tests/benches/conv_ops.rs new file mode 100644 index 0000000..9600676 --- /dev/null +++ b/crates/burn-backend-tests/benches/conv_ops.rs @@ -0,0 +1,498 @@ +//! Benchmarks for convolution operations. +//! +//! Run with: +//! ```bash +//! cargo bench --bench conv_ops --features simd,gemm +//! ``` +//! +//! Memory allocation tracking is enabled via divan's AllocProfiler. + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{Tensor, TensorData, module, ops::ConvOptions}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Convolution Benchmarks"); + println!("Memory allocation tracking enabled"); + println!(); + divan::main(); + common::report_failures(); +} + +fn make_input_2d(batch: usize, channels: usize, height: usize, width: usize) -> Tensor<4> { + let data: Vec = (0..batch * channels * height * width) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [batch, channels, height, width]), + &Default::default(), + ) +} + +fn make_kernel_2d(out_ch: usize, in_ch: usize, kh: usize, kw: usize) -> Tensor<4> { + let data: Vec = (0..out_ch * in_ch * kh * kw) + .map(|i| ((i % 100) as f32 / 100.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [out_ch, in_ch, kh, kw]), + &Default::default(), + ) +} + +fn make_input_1d(batch: usize, channels: usize, length: usize) -> Tensor<3> { + let data: Vec = (0..batch * channels * length) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [batch, channels, length]), + &Default::default(), + ) +} + +fn make_kernel_1d(out_ch: usize, in_ch: usize, kw: usize) -> Tensor<3> { + let data: Vec = (0..out_ch * in_ch * kw) + .map(|i| ((i % 100) as f32 / 100.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [out_ch, in_ch, kw]), + &Default::default(), + ) +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "conv2d_small")] + mod conv2d_small { + use super::*; + + #[divan::bench] + fn conv2d_1x3x32x32_k3x3(bencher: Bencher) { + let x = make_input_2d(1, 3, 32, 32); + let w = make_kernel_2d(16, 3, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn conv2d_1x16x32x32_k3x3(bencher: Bencher) { + let x = make_input_2d(1, 16, 32, 32); + let w = make_kernel_2d(32, 16, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn conv2d_1x32x16x16_k3x3(bencher: Bencher) { + let x = make_input_2d(1, 32, 16, 16); + let w = make_kernel_2d(64, 32, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + } + + #[divan::bench_group(name = "conv2d_medium")] + mod conv2d_medium { + use super::*; + + #[divan::bench] + fn conv2d_8x3x64x64_k3x3(bencher: Bencher) { + let x = make_input_2d(8, 3, 64, 64); + let w = make_kernel_2d(32, 3, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn conv2d_8x32x64x64_k3x3(bencher: Bencher) { + let x = make_input_2d(8, 32, 64, 64); + let w = make_kernel_2d(64, 32, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn conv2d_8x64x32x32_k3x3(bencher: Bencher) { + let x = make_input_2d(8, 64, 32, 32); + let w = make_kernel_2d(128, 64, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + } + + #[divan::bench_group(name = "conv2d_large")] + mod conv2d_large { + use super::*; + + #[divan::bench] + fn conv2d_16x64x128x128_k3x3(bencher: Bencher) { + let x = make_input_2d(16, 64, 128, 128); + let w = make_kernel_2d(128, 64, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn conv2d_16x128x64x64_k3x3(bencher: Bencher) { + let x = make_input_2d(16, 128, 64, 64); + let w = make_kernel_2d(256, 128, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + } + + #[divan::bench_group(name = "conv2d_resnet")] + mod conv2d_resnet { + use super::*; + + #[divan::bench] + fn resnet_conv1_1x3x224x224_k7x7_s2(bencher: Bencher) { + let x = make_input_2d(1, 3, 224, 224); + let w = make_kernel_2d(64, 3, 7, 7); + let opts = ConvOptions::new([2, 2], [3, 3], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn resnet_layer1_1x64x56x56_k3x3(bencher: Bencher) { + let x = make_input_2d(1, 64, 56, 56); + let w = make_kernel_2d(64, 64, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn resnet_layer2_1x128x28x28_k3x3(bencher: Bencher) { + let x = make_input_2d(1, 128, 28, 28); + let w = make_kernel_2d(128, 128, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn resnet_layer3_1x256x14x14_k3x3(bencher: Bencher) { + let x = make_input_2d(1, 256, 14, 14); + let w = make_kernel_2d(256, 256, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn resnet_layer4_1x512x7x7_k3x3(bencher: Bencher) { + let x = make_input_2d(1, 512, 7, 7); + let w = make_kernel_2d(512, 512, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + } + + #[divan::bench_group(name = "conv1d")] + mod conv1d { + use super::*; + + #[divan::bench] + fn conv1d_1x16x256_k3(bencher: Bencher) { + let x = make_input_1d(1, 16, 256); + let w = make_kernel_1d(32, 16, 3); + let opts = ConvOptions::new([1], [1], [1], 1); + bencher + .bench_synced(|| module::conv1d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn conv1d_8x32x512_k5(bencher: Bencher) { + let x = make_input_1d(8, 32, 512); + let w = make_kernel_1d(64, 32, 5); + let opts = ConvOptions::new([1], [2], [1], 1); + bencher + .bench_synced(|| module::conv1d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn conv1d_16x64x1024_k7(bencher: Bencher) { + let x = make_input_1d(16, 64, 1024); + let w = make_kernel_1d(128, 64, 7); + let opts = ConvOptions::new([1], [3], [1], 1); + bencher + .bench_synced(|| module::conv1d(x.clone(), w.clone(), None, opts.clone())); + } + } + + // Depthwise conv1d (groups == channels_in == channels_out). + // These flow through the same conv3d_depthwise_impl fast path as + // conv2d because conv1d expands to a trivial 3D shape (kd=kh=1). + #[divan::bench_group(name = "conv1d_depthwise")] + mod conv1d_depthwise { + use super::*; + + #[divan::bench] + fn depthwise_k3_8x32x512(bencher: Bencher) { + let x = make_input_1d(8, 32, 512); + let w = make_kernel_1d(32, 1, 3); + let opts = ConvOptions::new([1], [1], [1], 32); + bencher + .bench_synced(|| module::conv1d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn depthwise_k7_8x64x1024(bencher: Bencher) { + let x = make_input_1d(8, 64, 1024); + let w = make_kernel_1d(64, 1, 7); + let opts = ConvOptions::new([1], [3], [1], 64); + bencher + .bench_synced(|| module::conv1d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn depthwise_k15_4x128x2048(bencher: Bencher) { + // Larger receptive field, dilation 1. Tests the common + // ConvNeXt-1d / audio separable conv case. + let x = make_input_1d(4, 128, 2048); + let w = make_kernel_1d(128, 1, 15); + let opts = ConvOptions::new([1], [7], [1], 128); + bencher + .bench_synced(|| module::conv1d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn depthwise_k3_stride2_8x64x1024(bencher: Bencher) { + let x = make_input_1d(8, 64, 1024); + let w = make_kernel_1d(64, 1, 3); + let opts = ConvOptions::new([2], [1], [1], 64); + bencher + .bench_synced(|| module::conv1d(x.clone(), w.clone(), None, opts.clone())); + } + } + + // Small-channel groups=1 convs (Sobel-style edge filters and + // first-stage image preprocessors). These were reported as a + // 2.6-3.3x slowdown vs burn-ndarray before the small-channel + // fast path landed. + #[divan::bench_group(name = "conv2d_small_channel")] + mod conv2d_small_channel { + use super::*; + + #[divan::bench] + fn sobel_1x3x488x448_k3x3(bencher: Bencher) { + // 3-channel image, same-padding 3x3, 3 output channels. + let x = make_input_2d(1, 3, 488, 448); + let w = make_kernel_2d(3, 3, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn sobel_1x3x488x448_k1x3(bencher: Bencher) { + let x = make_input_2d(1, 3, 488, 448); + let w = make_kernel_2d(3, 3, 1, 3); + let opts = ConvOptions::new([1, 1], [0, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn sobel_1x3x488x448_k3x1(bencher: Bencher) { + let x = make_input_2d(1, 3, 488, 448); + let w = make_kernel_2d(3, 3, 3, 1); + let opts = ConvOptions::new([1, 1], [1, 0], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn sobel_1x3x488x448_k1x5(bencher: Bencher) { + // 1x5 Sobel pass. Separable filter, runs on RGB. + let x = make_input_2d(1, 3, 488, 448); + let w = make_kernel_2d(3, 3, 1, 5); + let opts = ConvOptions::new([1, 1], [0, 2], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn sobel_1x3x488x448_k5x1(bencher: Bencher) { + // 5x1 Sobel pass. Separable filter, runs on RGB. + let x = make_input_2d(1, 3, 488, 448); + let w = make_kernel_2d(3, 3, 5, 1); + let opts = ConvOptions::new([1, 1], [2, 0], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn downsample_1x4x128x128_k4x4_s2_to12(bencher: Bencher) { + // 4 in, 12 out, 4x4 kernel, stride 2. Small downsample stem. + let x = make_input_2d(1, 4, 128, 128); + let w = make_kernel_2d(12, 4, 4, 4); + let opts = ConvOptions::new([2, 2], [1, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn preproc_1x3x488x448_k5x5_to8(bencher: Bencher) { + // First-stage image preprocessor: 3 in, 8 out, 5x5. + let x = make_input_2d(1, 3, 488, 448); + let w = make_kernel_2d(8, 3, 5, 5); + let opts = ConvOptions::new([1, 1], [2, 2], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn mask_1x1x256x256_k3x3_to8(bencher: Bencher) { + // Single-channel mask input, early feature extractor. + let x = make_input_2d(1, 1, 256, 256); + let w = make_kernel_2d(8, 1, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn first_layer_4x3x224x224_k7x7_s2(bencher: Bencher) { + // Classic ImageNet first layer: 3 in, 7x7, stride 2. + // Output channels = 64 is large so gemm could compete; + // this bench shows whether our path is at worst on par. + let x = make_input_2d(4, 3, 224, 224); + let w = make_kernel_2d(64, 3, 7, 7); + let opts = ConvOptions::new([2, 2], [3, 3], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + } + + #[divan::bench_group(name = "conv2d_kernel_sizes")] + mod conv2d_kernel_sizes { + use super::*; + + #[divan::bench] + fn conv2d_k1x1(bencher: Bencher) { + let x = make_input_2d(4, 64, 56, 56); + let w = make_kernel_2d(128, 64, 1, 1); + let opts = ConvOptions::new([1, 1], [0, 0], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn conv2d_k3x3(bencher: Bencher) { + let x = make_input_2d(4, 64, 56, 56); + let w = make_kernel_2d(128, 64, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn conv2d_k5x5(bencher: Bencher) { + let x = make_input_2d(4, 64, 56, 56); + let w = make_kernel_2d(128, 64, 5, 5); + let opts = ConvOptions::new([1, 1], [2, 2], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn conv2d_k7x7(bencher: Bencher) { + let x = make_input_2d(4, 64, 56, 56); + let w = make_kernel_2d(128, 64, 7, 7); + let opts = ConvOptions::new([1, 1], [3, 3], [1, 1], 1); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + } + + // Depthwise (groups == channels_in == channels_out). Shapes cover + // MobileNet/ConvNeXt style depthwise blocks and a ConvNeXt-7x7 + // block reported in a user regression on burn-ndarray parity. + #[divan::bench_group(name = "conv2d_depthwise")] + mod conv2d_depthwise { + use super::*; + + #[divan::bench] + fn depthwise_3x3_4x32x56x56(bencher: Bencher) { + let x = make_input_2d(4, 32, 56, 56); + let w = make_kernel_2d(32, 1, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 32); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn depthwise_3x3_4x96x28x28(bencher: Bencher) { + let x = make_input_2d(4, 96, 28, 28); + let w = make_kernel_2d(96, 1, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 96); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn depthwise_3x3_4x192x14x14(bencher: Bencher) { + let x = make_input_2d(4, 192, 14, 14); + let w = make_kernel_2d(192, 1, 3, 3); + let opts = ConvOptions::new([1, 1], [1, 1], [1, 1], 192); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn depthwise_7x7_4x24x56x56(bencher: Bencher) { + // ConvNeXt-style 7x7 depthwise. This shape was 3x slower + // than burn-ndarray before the depthwise fast path landed. + let x = make_input_2d(4, 24, 56, 56); + let w = make_kernel_2d(24, 1, 7, 7); + let opts = ConvOptions::new([1, 1], [3, 3], [1, 1], 24); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn depthwise_7x7_4x48x56x56(bencher: Bencher) { + // Another regressed shape (channels_out = 48). + let x = make_input_2d(4, 48, 56, 56); + let w = make_kernel_2d(48, 1, 7, 7); + let opts = ConvOptions::new([1, 1], [3, 3], [1, 1], 48); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + + #[divan::bench] + fn depthwise_3x3_stride2_4x64x56x56(bencher: Bencher) { + // Downsampling depthwise 3x3. + let x = make_input_2d(4, 64, 56, 56); + let w = make_kernel_2d(64, 1, 3, 3); + let opts = ConvOptions::new([2, 2], [1, 1], [1, 1], 64); + bencher + .bench_synced(|| module::conv2d(x.clone(), w.clone(), None, opts.clone())); + } + } + } + }; +} + +bench_backend!(backend, "backend"); diff --git a/crates/burn-backend-tests/benches/conv_transpose_ops.rs b/crates/burn-backend-tests/benches/conv_transpose_ops.rs new file mode 100644 index 0000000..abe0118 --- /dev/null +++ b/crates/burn-backend-tests/benches/conv_transpose_ops.rs @@ -0,0 +1,284 @@ +//! Benchmarks for transposed convolution operations. +//! +//! Run with: +//! ```bash +//! cargo bench --bench conv_transpose_ops --features simd,rayon +//! ``` +//! +//! Memory allocation tracking is enabled via divan's AllocProfiler. + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::ops::ConvTransposeOptions; +use burn_tensor::{Tensor, TensorData, module}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Conv Transpose Benchmarks"); + println!("Memory allocation tracking enabled"); + println!(); + divan::main(); + common::report_failures(); +} + +fn make_input_2d(batch: usize, channels: usize, height: usize, width: usize) -> Tensor<4> { + let data: Vec = (0..batch * channels * height * width) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [batch, channels, height, width]), + &Default::default(), + ) +} + +fn make_weight_2d( + in_channels: usize, + out_channels: usize, + kernel_h: usize, + kernel_w: usize, +) -> Tensor<4> { + let data: Vec = (0..in_channels * out_channels * kernel_h * kernel_w) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [in_channels, out_channels, kernel_h, kernel_w]), + &Default::default(), + ) +} + +fn make_input_1d(batch: usize, channels: usize, length: usize) -> Tensor<3> { + let data: Vec = (0..batch * channels * length) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [batch, channels, length]), + &Default::default(), + ) +} + +fn make_weight_1d(in_channels: usize, out_channels: usize, kernel: usize) -> Tensor<3> { + let data: Vec = (0..in_channels * out_channels * kernel) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [in_channels, out_channels, kernel]), + &Default::default(), + ) +} + +fn make_input_3d( + batch: usize, + channels: usize, + depth: usize, + height: usize, + width: usize, +) -> Tensor<5> { + let data: Vec = (0..batch * channels * depth * height * width) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [batch, channels, depth, height, width]), + &Default::default(), + ) +} + +fn make_weight_3d( + in_channels: usize, + out_channels: usize, + kernel_d: usize, + kernel_h: usize, + kernel_w: usize, +) -> Tensor<5> { + let data: Vec = (0..in_channels * out_channels * kernel_d * kernel_h * kernel_w) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new( + data, + [in_channels, out_channels, kernel_d, kernel_h, kernel_w], + ), + &Default::default(), + ) +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "conv_transpose2d")] + mod conv_transpose2d { + use super::*; + + #[divan::bench] + fn conv_transpose2d_1x64x7x7_to_14x14(bencher: Bencher) { + // Upsample from 7x7 to 14x14 (common in decoder/generator) + let x = make_input_2d(1, 64, 7, 7); + let w = make_weight_2d(64, 64, 4, 4); + let opts = ConvTransposeOptions::new([2, 2], [1, 1], [0, 0], [1, 1], 1); + bencher.bench_synced(|| { + module::conv_transpose2d(x.clone(), w.clone(), None, opts.clone()) + }); + } + + #[divan::bench] + fn conv_transpose2d_1x128x14x14_to_28x28(bencher: Bencher) { + let x = make_input_2d(1, 128, 14, 14); + let w = make_weight_2d(128, 64, 4, 4); + let opts = ConvTransposeOptions::new([2, 2], [1, 1], [0, 0], [1, 1], 1); + bencher.bench_synced(|| { + module::conv_transpose2d(x.clone(), w.clone(), None, opts.clone()) + }); + } + + #[divan::bench] + fn conv_transpose2d_1x256x28x28_to_56x56(bencher: Bencher) { + let x = make_input_2d(1, 256, 28, 28); + let w = make_weight_2d(256, 128, 4, 4); + let opts = ConvTransposeOptions::new([2, 2], [1, 1], [0, 0], [1, 1], 1); + bencher.bench_synced(|| { + module::conv_transpose2d(x.clone(), w.clone(), None, opts.clone()) + }); + } + + #[divan::bench] + fn conv_transpose2d_8x64x14x14_to_28x28(bencher: Bencher) { + // Batch of 8 + let x = make_input_2d(8, 64, 14, 14); + let w = make_weight_2d(64, 64, 4, 4); + let opts = ConvTransposeOptions::new([2, 2], [1, 1], [0, 0], [1, 1], 1); + bencher.bench_synced(|| { + module::conv_transpose2d(x.clone(), w.clone(), None, opts.clone()) + }); + } + + #[divan::bench] + fn conv_transpose2d_1x512x7x7_k3x3_s1(bencher: Bencher) { + // No upsampling, just transpose conv + let x = make_input_2d(1, 512, 7, 7); + let w = make_weight_2d(512, 512, 3, 3); + let opts = ConvTransposeOptions::new([1, 1], [1, 1], [0, 0], [1, 1], 1); + bencher.bench_synced(|| { + module::conv_transpose2d(x.clone(), w.clone(), None, opts.clone()) + }); + } + } + + #[divan::bench_group(name = "conv_transpose2d_gan")] + mod conv_transpose2d_gan { + use super::*; + + #[divan::bench] + fn dcgan_layer1_1x512x1x1_to_4x4(bencher: Bencher) { + // DCGAN first layer: project and reshape + let x = make_input_2d(1, 512, 1, 1); + let w = make_weight_2d(512, 256, 4, 4); + let opts = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [1, 1], 1); + bencher.bench_synced(|| { + module::conv_transpose2d(x.clone(), w.clone(), None, opts.clone()) + }); + } + + #[divan::bench] + fn dcgan_layer2_1x256x4x4_to_8x8(bencher: Bencher) { + let x = make_input_2d(1, 256, 4, 4); + let w = make_weight_2d(256, 128, 4, 4); + let opts = ConvTransposeOptions::new([2, 2], [1, 1], [0, 0], [1, 1], 1); + bencher.bench_synced(|| { + module::conv_transpose2d(x.clone(), w.clone(), None, opts.clone()) + }); + } + + #[divan::bench] + fn dcgan_layer3_1x128x8x8_to_16x16(bencher: Bencher) { + let x = make_input_2d(1, 128, 8, 8); + let w = make_weight_2d(128, 64, 4, 4); + let opts = ConvTransposeOptions::new([2, 2], [1, 1], [0, 0], [1, 1], 1); + bencher.bench_synced(|| { + module::conv_transpose2d(x.clone(), w.clone(), None, opts.clone()) + }); + } + + #[divan::bench] + fn dcgan_layer4_1x64x16x16_to_32x32(bencher: Bencher) { + let x = make_input_2d(1, 64, 16, 16); + let w = make_weight_2d(64, 3, 4, 4); + let opts = ConvTransposeOptions::new([2, 2], [1, 1], [0, 0], [1, 1], 1); + bencher.bench_synced(|| { + module::conv_transpose2d(x.clone(), w.clone(), None, opts.clone()) + }); + } + } + + #[divan::bench_group(name = "conv_transpose1d")] + mod conv_transpose1d { + use super::*; + + #[divan::bench] + fn conv_transpose1d_1x64x32_to_64(bencher: Bencher) { + let x = make_input_1d(1, 64, 32); + let w = make_weight_1d(64, 64, 4); + let opts = ConvTransposeOptions::new([2], [1], [0], [1], 1); + bencher.bench_synced(|| { + module::conv_transpose1d(x.clone(), w.clone(), None, opts.clone()) + }); + } + + #[divan::bench] + fn conv_transpose1d_8x128x64_to_128(bencher: Bencher) { + let x = make_input_1d(8, 128, 64); + let w = make_weight_1d(128, 64, 4); + let opts = ConvTransposeOptions::new([2], [1], [0], [1], 1); + bencher.bench_synced(|| { + module::conv_transpose1d(x.clone(), w.clone(), None, opts.clone()) + }); + } + + #[divan::bench] + fn conv_transpose1d_1x256x128_to_256(bencher: Bencher) { + let x = make_input_1d(1, 256, 128); + let w = make_weight_1d(256, 128, 4); + let opts = ConvTransposeOptions::new([2], [1], [0], [1], 1); + bencher.bench_synced(|| { + module::conv_transpose1d(x.clone(), w.clone(), None, opts.clone()) + }); + } + } + + #[divan::bench_group(name = "conv_transpose3d")] + mod conv_transpose3d { + use super::*; + + #[divan::bench] + fn conv_transpose3d_1x32x4x4x4_to_8x8x8(bencher: Bencher) { + let x = make_input_3d(1, 32, 4, 4, 4); + let w = make_weight_3d(32, 32, 4, 4, 4); + let opts = + ConvTransposeOptions::new([2, 2, 2], [1, 1, 1], [0, 0, 0], [1, 1, 1], 1); + bencher.bench_synced(|| { + module::conv_transpose3d(x.clone(), w.clone(), None, opts.clone()) + }); + } + + #[divan::bench] + fn conv_transpose3d_1x64x8x8x8_to_16x16x16(bencher: Bencher) { + let x = make_input_3d(1, 64, 8, 8, 8); + let w = make_weight_3d(64, 32, 4, 4, 4); + let opts = + ConvTransposeOptions::new([2, 2, 2], [1, 1, 1], [0, 0, 0], [1, 1, 1], 1); + bencher.bench_synced(|| { + module::conv_transpose3d(x.clone(), w.clone(), None, opts.clone()) + }); + } + } + } + }; +} + +bench_backend!(backend, "backend"); diff --git a/crates/burn-backend-tests/benches/cross_unfold_ops.rs b/crates/burn-backend-tests/benches/cross_unfold_ops.rs new file mode 100644 index 0000000..e0a8adb --- /dev/null +++ b/crates/burn-backend-tests/benches/cross_unfold_ops.rs @@ -0,0 +1,154 @@ +//! Benchmarks for cross and unfold operations. +//! +//! Run with: +//! ```bash +//! cargo bench --bench cross_unfold_ops --features simd +//! ``` + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{Tensor, TensorData}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Cross and unfold ops Benchmarks"); + println!(); + divan::main(); + common::report_failures(); +} + +// Cross product requires 3 elements along the specified dimension +fn make_cross_tensor_1d(batch: usize) -> Tensor<2> { + let data: Vec = (0..batch * 3).map(|i| (i % 1000) as f32 / 100.0).collect(); + Tensor::from_data(TensorData::new(data, [batch, 3]), &Default::default()) +} + +fn make_cross_tensor_2d(d0: usize, d1: usize) -> Tensor<3> { + let data: Vec = (0..d0 * d1 * 3) + .map(|i| (i % 1000) as f32 / 100.0) + .collect(); + Tensor::from_data(TensorData::new(data, [d0, 3, d1]), &Default::default()) +} + +fn make_tensor_1d(size: usize) -> Tensor<1> { + let data: Vec = (0..size).map(|i| (i % 1000) as f32 / 1000.0).collect(); + Tensor::from_data(TensorData::new(data, [size]), &Default::default()) +} + +fn make_tensor_2d(rows: usize, cols: usize) -> Tensor<2> { + let data: Vec = (0..rows * cols) + .map(|i| (i % 1000) as f32 / 1000.0) + .collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) +} + +fn make_tensor_3d(d0: usize, d1: usize, d2: usize) -> Tensor<3> { + let size = d0 * d1 * d2; + let data: Vec = (0..size).map(|i| (i % 1000) as f32 / 1000.0).collect(); + Tensor::from_data(TensorData::new(data, [d0, d1, d2]), &Default::default()) +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "cross")] + mod cross { + use super::*; + + #[divan::bench] + fn s_1k_vectors(bencher: Bencher) { + let a = make_cross_tensor_1d(1024); + let b = make_cross_tensor_1d(1024); + bencher.bench_synced(|| a.clone().cross(b.clone(), 1)); + } + + #[divan::bench] + fn s_64k_vectors(bencher: Bencher) { + let a = make_cross_tensor_1d(64 * 1024); + let b = make_cross_tensor_1d(64 * 1024); + bencher.bench_synced(|| a.clone().cross(b.clone(), 1)); + } + + #[divan::bench] + fn s_256k_vectors(bencher: Bencher) { + let a = make_cross_tensor_1d(256 * 1024); + let b = make_cross_tensor_1d(256 * 1024); + bencher.bench_synced(|| a.clone().cross(b.clone(), 1)); + } + + #[divan::bench] + fn s_3d_64x64(bencher: Bencher) { + let a = make_cross_tensor_2d(64, 64); + let b = make_cross_tensor_2d(64, 64); + bencher.bench_synced(|| a.clone().cross(b.clone(), 1)); + } + } + + #[divan::bench_group(name = "unfold")] + mod unfold { + use super::*; + + // 1D unfold: sliding window extraction + #[divan::bench] + fn s_1d_1k_win8_step1(bencher: Bencher) { + let t: Tensor<1> = make_tensor_1d(1024); + bencher.bench_synced(|| -> Tensor<2> { t.clone().unfold(0, 8, 1) }); + } + + #[divan::bench] + fn s_1d_64k_win8_step1(bencher: Bencher) { + let t: Tensor<1> = make_tensor_1d(64 * 1024); + bencher.bench_synced(|| -> Tensor<2> { t.clone().unfold(0, 8, 1) }); + } + + #[divan::bench] + fn s_1d_64k_win64_step1(bencher: Bencher) { + let t: Tensor<1> = make_tensor_1d(64 * 1024); + bencher.bench_synced(|| -> Tensor<2> { t.clone().unfold(0, 64, 1) }); + } + + #[divan::bench] + fn s_1d_64k_win64_step32(bencher: Bencher) { + let t: Tensor<1> = make_tensor_1d(64 * 1024); + bencher.bench_synced(|| -> Tensor<2> { t.clone().unfold(0, 64, 32) }); + } + + // 2D unfold along dim 1 + #[divan::bench] + fn s_2d_256x256_dim1_win8_step1(bencher: Bencher) { + let t: Tensor<2> = make_tensor_2d(256, 256); + bencher.bench_synced(|| -> Tensor<3> { t.clone().unfold(1, 8, 1) }); + } + + #[divan::bench] + fn s_2d_256x256_dim1_win32_step16(bencher: Bencher) { + let t: Tensor<2> = make_tensor_2d(256, 256); + bencher.bench_synced(|| -> Tensor<3> { t.clone().unfold(1, 32, 16) }); + } + + #[divan::bench] + fn s_2d_1024x256_dim1_win8_step1(bencher: Bencher) { + let t: Tensor<2> = make_tensor_2d(1024, 256); + bencher.bench_synced(|| -> Tensor<3> { t.clone().unfold(1, 8, 1) }); + } + + // 3D unfold (batch scenarios) + #[divan::bench] + fn s_3d_32x64x64_dim2_win8_step4(bencher: Bencher) { + let t: Tensor<3> = make_tensor_3d(32, 64, 64); + bencher.bench_synced(|| -> Tensor<4> { t.clone().unfold(2, 8, 4) }); + } + } + } + }; +} + +bench_backend!(backend, "backend"); diff --git a/crates/burn-backend-tests/benches/cumulative_ops.rs b/crates/burn-backend-tests/benches/cumulative_ops.rs new file mode 100644 index 0000000..3068999 --- /dev/null +++ b/crates/burn-backend-tests/benches/cumulative_ops.rs @@ -0,0 +1,174 @@ +//! Benchmarks for cumulative operations. +//! +//! Run with: +//! ```bash +//! cargo bench --bench cumulative_ops --features simd +//! ``` + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{Tensor, TensorData}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Cumulative ops Benchmarks"); + println!(); + divan::main(); + common::report_failures(); +} + +fn make_tensor_1d(size: usize) -> Tensor<1> { + let data: Vec = (0..size).map(|i| (i % 1000) as f32 / 1000.0).collect(); + Tensor::from_data(TensorData::new(data, [size]), &Default::default()) +} + +fn make_tensor_2d(rows: usize, cols: usize) -> Tensor<2> { + let data: Vec = (0..rows * cols) + .map(|i| (i % 1000) as f32 / 1000.0) + .collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) +} + +fn make_tensor_3d(d0: usize, d1: usize, d2: usize) -> Tensor<3> { + let size = d0 * d1 * d2; + let data: Vec = (0..size).map(|i| (i % 1000) as f32 / 1000.0).collect(); + Tensor::from_data(TensorData::new(data, [d0, d1, d2]), &Default::default()) +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "cumsum")] + mod cumsum { + use super::*; + + #[divan::bench] + fn s_1k(bencher: Bencher) { + let t = make_tensor_1d(1024); + bencher.bench_synced(|| t.clone().cumsum(0)); + } + + #[divan::bench] + fn s_64k(bencher: Bencher) { + let t = make_tensor_1d(64 * 1024); + bencher.bench_synced(|| t.clone().cumsum(0)); + } + + #[divan::bench] + fn s_1m(bencher: Bencher) { + let t = make_tensor_1d(1024 * 1024); + bencher.bench_synced(|| t.clone().cumsum(0)); + } + + #[divan::bench] + fn s_256x256_dim0(bencher: Bencher) { + let t = make_tensor_2d(256, 256); + bencher.bench_synced(|| t.clone().cumsum(0)); + } + + #[divan::bench] + fn s_256x256_dim1(bencher: Bencher) { + let t = make_tensor_2d(256, 256); + bencher.bench_synced(|| t.clone().cumsum(1)); + } + + #[divan::bench] + fn s_1024x1024_dim1(bencher: Bencher) { + let t = make_tensor_2d(1024, 1024); + bencher.bench_synced(|| t.clone().cumsum(1)); + } + } + + #[divan::bench_group(name = "cumprod")] + mod cumprod { + use super::*; + + #[divan::bench] + fn s_1k(bencher: Bencher) { + let t = make_tensor_1d(1024); + bencher.bench_synced(|| t.clone().cumprod(0)); + } + + #[divan::bench] + fn s_256x256_dim1(bencher: Bencher) { + let t = make_tensor_2d(256, 256); + bencher.bench_synced(|| t.clone().cumprod(1)); + } + } + + #[divan::bench_group(name = "cummin")] + mod cummin { + use super::*; + + #[divan::bench] + fn s_1k(bencher: Bencher) { + let t = make_tensor_1d(1024); + bencher.bench_synced(|| t.clone().cummin(0)); + } + + #[divan::bench] + fn s_256x256_dim1(bencher: Bencher) { + let t = make_tensor_2d(256, 256); + bencher.bench_synced(|| t.clone().cummin(1)); + } + + #[divan::bench] + fn s_1024x1024_dim1(bencher: Bencher) { + let t = make_tensor_2d(1024, 1024); + bencher.bench_synced(|| t.clone().cummin(1)); + } + } + + #[divan::bench_group(name = "cummax")] + mod cummax { + use super::*; + + #[divan::bench] + fn s_1k(bencher: Bencher) { + let t = make_tensor_1d(1024); + bencher.bench_synced(|| t.clone().cummax(0)); + } + + #[divan::bench] + fn s_256x256_dim1(bencher: Bencher) { + let t = make_tensor_2d(256, 256); + bencher.bench_synced(|| t.clone().cummax(1)); + } + + #[divan::bench] + fn s_1024x1024_dim1(bencher: Bencher) { + let t = make_tensor_2d(1024, 1024); + bencher.bench_synced(|| t.clone().cummax(1)); + } + } + + // 3D cumulative (batch scenarios) + #[divan::bench_group(name = "cumsum_3d")] + mod cumsum_3d { + use super::*; + + #[divan::bench] + fn s_batch32_64x64_dim2(bencher: Bencher) { + let t = make_tensor_3d(32, 64, 64); + bencher.bench_synced(|| t.clone().cumsum(2)); + } + + #[divan::bench] + fn s_batch32_64x64_dim1(bencher: Bencher) { + let t = make_tensor_3d(32, 64, 64); + bencher.bench_synced(|| t.clone().cumsum(1)); + } + } + } + }; +} + +bench_backend!(backend, "backend"); diff --git a/crates/burn-backend-tests/benches/default_ops.rs b/crates/burn-backend-tests/benches/default_ops.rs new file mode 100644 index 0000000..6a4e847 --- /dev/null +++ b/crates/burn-backend-tests/benches/default_ops.rs @@ -0,0 +1,273 @@ +//! Benchmarks for default ops. +//! +//! These ops previously used burn's default trait implementations (TensorData +//! round-trips, slice_assign loops, etc.) and now have direct implementations. +//! +//! Run with: +//! ```bash +//! cargo bench --bench default_ops --features simd +//! ``` + +#![allow(clippy::approx_constant)] + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{Int, Tensor, TensorData}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Default ops Benchmarks"); + println!(); + divan::main(); + common::report_failures(); +} + +const SMALL: usize = 64 * 64; +const MEDIUM: usize = 256 * 256; +const LARGE: usize = 1024 * 1024; + +fn make_tensor_1d(size: usize) -> Tensor<1> { + let data: Vec = (0..size) + .map(|i| (i % 1000) as f32 / 1000.0 - 0.5) + .collect(); + Tensor::from_data(TensorData::new(data, [size]), &Default::default()) +} + +fn make_tensor_2d(rows: usize, cols: usize) -> Tensor<2> { + let data: Vec = (0..rows * cols) + .map(|i| (i % 1000) as f32 / 1000.0 - 0.5) + .collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + use burn_tensor::Shape; + + // --------------------------------------------------------------- + // sort (1D) + // --------------------------------------------------------------- + + #[divan::bench_group(name = "sort_1d")] + mod sort_1d { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let a = make_tensor_1d(SMALL); + bencher.bench_synced(|| a.clone().sort(0)); + } + + #[divan::bench] + fn medium(bencher: Bencher) { + let a = make_tensor_1d(MEDIUM); + bencher.bench_synced(|| a.clone().sort(0)); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor_1d(LARGE); + bencher.bench_synced(|| a.clone().sort(0)); + } + } + + // --------------------------------------------------------------- + // sort (2D along dim 1) + // --------------------------------------------------------------- + + #[divan::bench_group(name = "sort_2d_dim1")] + mod sort_2d { + use super::*; + + #[divan::bench] + fn rows64_cols64(bencher: Bencher) { + let a = make_tensor_2d(64, 64); + bencher.bench_synced(|| a.clone().sort(1)); + } + + #[divan::bench] + fn rows256_cols256(bencher: Bencher) { + let a = make_tensor_2d(256, 256); + bencher.bench_synced(|| a.clone().sort(1)); + } + + #[divan::bench] + fn rows1024_cols1024(bencher: Bencher) { + let a = make_tensor_2d(1024, 1024); + bencher.bench_synced(|| a.clone().sort(1)); + } + } + + // --------------------------------------------------------------- + // argsort (1D) + // --------------------------------------------------------------- + + #[divan::bench_group(name = "argsort_1d")] + mod argsort_1d { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let a = make_tensor_1d(SMALL); + bencher.bench_synced(|| a.clone().argsort(0)); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor_1d(LARGE); + bencher.bench_synced(|| a.clone().argsort(0)); + } + } + + // --------------------------------------------------------------- + // repeat_dim + // --------------------------------------------------------------- + + #[divan::bench_group(name = "repeat_dim")] + mod repeat_dim { + use super::*; + + #[divan::bench] + fn repeat_dim0_4x(bencher: Bencher) { + let a = make_tensor_2d(256, 256); + bencher.bench_synced(|| a.clone().repeat_dim(0, 4)); + } + + #[divan::bench] + fn repeat_dim1_4x(bencher: Bencher) { + let a = make_tensor_2d(256, 256); + bencher.bench_synced(|| a.clone().repeat_dim(1, 4)); + } + + #[divan::bench] + fn repeat_large_dim0_8x(bencher: Bencher) { + let a = make_tensor_2d(512, 512); + bencher.bench_synced(|| a.clone().repeat_dim(0, 8)); + } + } + + // --------------------------------------------------------------- + // zeros / ones / full + // --------------------------------------------------------------- + + #[divan::bench_group(name = "creation")] + mod creation { + use super::*; + + #[divan::bench] + fn zeros_large(bencher: Bencher) { + let dev = Default::default(); + bencher.bench_synced(|| Tensor::<1>::zeros(Shape::new([LARGE]), &dev)); + } + + #[divan::bench] + fn ones_large(bencher: Bencher) { + let dev = Default::default(); + bencher.bench_synced(|| Tensor::<1>::ones(Shape::new([LARGE]), &dev)); + } + + #[divan::bench] + fn full_large(bencher: Bencher) { + let dev = Default::default(); + bencher.bench_synced(|| Tensor::<1>::full(Shape::new([LARGE]), 3.14, &dev)); + } + } + + // --------------------------------------------------------------- + // arange + // --------------------------------------------------------------- + + #[divan::bench_group(name = "arange")] + mod arange { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let dev = Default::default(); + bencher.bench_synced(|| Tensor::<1, Int>::arange(0..SMALL as i64, &dev)); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let dev = Default::default(); + bencher.bench_synced(|| Tensor::<1, Int>::arange(0..LARGE as i64, &dev)); + } + } + + // --------------------------------------------------------------- + // embedding + // --------------------------------------------------------------- + + #[divan::bench_group(name = "embedding")] + mod embedding { + use super::*; + use burn_tensor::module; + + fn make_weights(vocab: usize, dim: usize) -> Tensor<2> { + let data: Vec = (0..vocab * dim) + .map(|i| (i % 1000) as f32 / 1000.0) + .collect(); + Tensor::from_data(TensorData::new(data, [vocab, dim]), &Default::default()) + } + + fn make_indices(batch: usize, seq: usize, vocab: usize) -> Option> { + common::try_setup(|| { + let data: Vec = (0..batch * seq).map(|i| (i % vocab) as i32).collect(); + Tensor::from_data(TensorData::new(data, [batch, seq]), &Default::default()) + }) + } + + #[divan::bench] + fn vocab30k_dim512_batch8_seq128(bencher: Bencher) { + let weights = make_weights(30000, 512); + let Some(indices) = make_indices(8, 128, 30000) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| module::embedding(weights.clone(), indices.clone())); + } + + #[divan::bench] + fn vocab50k_dim768_batch4_seq256(bencher: Bencher) { + let weights = make_weights(50000, 768); + let Some(indices) = make_indices(4, 256, 50000) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| module::embedding(weights.clone(), indices.clone())); + } + } + + // --------------------------------------------------------------- + // is_nan / is_inf + // --------------------------------------------------------------- + + #[divan::bench_group(name = "predicates")] + mod predicates { + use super::*; + + #[divan::bench] + fn is_nan_large(bencher: Bencher) { + let a = make_tensor_1d(LARGE); + bencher.bench_synced(|| a.clone().is_nan()); + } + + #[divan::bench] + fn is_inf_large(bencher: Bencher) { + let a = make_tensor_1d(LARGE); + bencher.bench_synced(|| a.clone().is_inf()); + } + } + } + }; +} + +bench_backend!(backend, "backend"); diff --git a/crates/burn-backend-tests/benches/deform_conv_ops.rs b/crates/burn-backend-tests/benches/deform_conv_ops.rs new file mode 100644 index 0000000..60c3c97 --- /dev/null +++ b/crates/burn-backend-tests/benches/deform_conv_ops.rs @@ -0,0 +1,415 @@ +//! Benchmarks for deformable convolution. +//! +//! Run with: +//! ```bash +//! cargo bench --bench deform_conv_ops --features simd,rayon,gemm +//! ``` +//! +//! Memory allocation tracking is enabled via divan's AllocProfiler. + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{Tensor, TensorData, module, ops::DeformConvOptions}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Deformable Convolution Benchmarks"); + println!("Memory allocation tracking enabled"); + println!(); + divan::main(); + common::report_failures(); +} + +/// Create input tensor [batch, channels, height, width] +fn make_input(batch: usize, channels: usize, height: usize, width: usize) -> Tensor<4> { + let data: Vec = (0..batch * channels * height * width) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [batch, channels, height, width]), + &Default::default(), + ) +} + +/// Create weight tensor [out_channels, in_channels/groups, kernel_h, kernel_w] +fn make_weight(out_ch: usize, in_ch_per_group: usize, kh: usize, kw: usize) -> Tensor<4> { + let data: Vec = (0..out_ch * in_ch_per_group * kh * kw) + .map(|i| ((i % 100) as f32 / 100.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [out_ch, in_ch_per_group, kh, kw]), + &Default::default(), + ) +} + +/// Create offset tensor [batch, offset_groups * kernel_h * kernel_w * 2, out_h, out_w] +fn make_offset( + batch: usize, + offset_groups: usize, + kh: usize, + kw: usize, + out_h: usize, + out_w: usize, +) -> Tensor<4> { + let channels = offset_groups * kh * kw * 2; + let data: Vec = (0..batch * channels * out_h * out_w) + .map(|i| ((i % 100) as f32 / 100.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [batch, channels, out_h, out_w]), + &Default::default(), + ) +} + +/// Create mask tensor [batch, offset_groups * kernel_h * kernel_w, out_h, out_w] +fn make_mask( + batch: usize, + offset_groups: usize, + kh: usize, + kw: usize, + out_h: usize, + out_w: usize, +) -> Tensor<4> { + let channels = offset_groups * kh * kw; + let data: Vec = (0..batch * channels * out_h * out_w) + .map(|i| (i % 100) as f32 / 100.0) + .collect(); + Tensor::from_data( + TensorData::new(data, [batch, channels, out_h, out_w]), + &Default::default(), + ) +} + +/// Create bias tensor [out_channels] +fn make_bias(out_ch: usize) -> Tensor<1> { + let data: Vec = (0..out_ch).map(|i| (i % 10) as f32 / 10.0).collect(); + Tensor::from_data(TensorData::new(data, [out_ch]), &Default::default()) +} + +/// Compute output dimensions for convolution +fn compute_output_size( + input_size: usize, + kernel_size: usize, + padding: usize, + stride: usize, + dilation: usize, +) -> usize { + (input_size + 2 * padding - dilation * (kernel_size - 1) - 1) / stride + 1 +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "deform_conv2d_tiny")] + mod deform_conv2d_tiny { + use super::*; + + // Tiny: 1x3x8x8 -> 8 output channels + #[divan::bench(sample_count = 20, min_time = 0)] + fn deform_conv2d_1x3x8x8_k3x3_c8(bencher: Bencher) { + let (batch, in_ch, h, w) = (1, 3, 8, 8); + let (out_ch, kh, kw) = (8, 3, 3); + let (stride, padding, dilation) = ([1, 1], [1, 1], [1, 1]); + let (weight_groups, offset_groups) = (1, 1); + + let out_h = compute_output_size(h, kh, padding[0], stride[0], dilation[0]); + let out_w = compute_output_size(w, kw, padding[1], stride[1], dilation[1]); + + let x = make_input(batch, in_ch, h, w); + let weight = make_weight(out_ch, in_ch / weight_groups, kh, kw); + let offset = make_offset(batch, offset_groups, kh, kw, out_h, out_w); + let mask = make_mask(batch, offset_groups, kh, kw, out_h, out_w); + let bias = make_bias(out_ch); + let opts = DeformConvOptions::new( + stride, + padding, + dilation, + weight_groups, + offset_groups, + ); + + bencher.bench_synced(|| { + module::deform_conv2d( + x.clone(), + offset.clone(), + weight.clone(), + Some(mask.clone()), + Some(bias.clone()), + opts.clone(), + ) + }); + } + + // Tiny without mask + #[divan::bench(sample_count = 20, min_time = 0)] + fn deform_conv2d_1x3x8x8_k3x3_no_mask(bencher: Bencher) { + let (batch, in_ch, h, w) = (1, 3, 8, 8); + let (out_ch, kh, kw) = (8, 3, 3); + let (stride, padding, dilation) = ([1, 1], [1, 1], [1, 1]); + let (weight_groups, offset_groups) = (1, 1); + + let out_h = compute_output_size(h, kh, padding[0], stride[0], dilation[0]); + let out_w = compute_output_size(w, kw, padding[1], stride[1], dilation[1]); + + let x = make_input(batch, in_ch, h, w); + let weight = make_weight(out_ch, in_ch / weight_groups, kh, kw); + let offset = make_offset(batch, offset_groups, kh, kw, out_h, out_w); + let opts = DeformConvOptions::new( + stride, + padding, + dilation, + weight_groups, + offset_groups, + ); + + bencher.bench_synced(|| { + module::deform_conv2d( + x.clone(), + offset.clone(), + weight.clone(), + None, + None, + opts.clone(), + ) + }); + } + } + + #[divan::bench_group(name = "deform_conv2d_small")] + mod deform_conv2d_small { + use super::*; + + // Small: 1x3x16x16 -> 16 output channels + #[divan::bench(sample_count = 20, min_time = 0)] + fn deform_conv2d_1x3x16x16_k3x3_c16(bencher: Bencher) { + let (batch, in_ch, h, w) = (1, 3, 16, 16); + let (out_ch, kh, kw) = (16, 3, 3); + let (stride, padding, dilation) = ([1, 1], [1, 1], [1, 1]); + let (weight_groups, offset_groups) = (1, 1); + + let out_h = compute_output_size(h, kh, padding[0], stride[0], dilation[0]); + let out_w = compute_output_size(w, kw, padding[1], stride[1], dilation[1]); + + let x = make_input(batch, in_ch, h, w); + let weight = make_weight(out_ch, in_ch / weight_groups, kh, kw); + let offset = make_offset(batch, offset_groups, kh, kw, out_h, out_w); + let mask = make_mask(batch, offset_groups, kh, kw, out_h, out_w); + let bias = make_bias(out_ch); + let opts = DeformConvOptions::new( + stride, + padding, + dilation, + weight_groups, + offset_groups, + ); + + bencher.bench_synced(|| { + module::deform_conv2d( + x.clone(), + offset.clone(), + weight.clone(), + Some(mask.clone()), + Some(bias.clone()), + opts.clone(), + ) + }); + } + + // Small with stride 2 + #[divan::bench(sample_count = 20, min_time = 0)] + fn deform_conv2d_1x3x16x16_k3x3_stride2(bencher: Bencher) { + let (batch, in_ch, h, w) = (1, 3, 16, 16); + let (out_ch, kh, kw) = (16, 3, 3); + let (stride, padding, dilation) = ([2, 2], [1, 1], [1, 1]); + let (weight_groups, offset_groups) = (1, 1); + + let out_h = compute_output_size(h, kh, padding[0], stride[0], dilation[0]); + let out_w = compute_output_size(w, kw, padding[1], stride[1], dilation[1]); + + let x = make_input(batch, in_ch, h, w); + let weight = make_weight(out_ch, in_ch / weight_groups, kh, kw); + let offset = make_offset(batch, offset_groups, kh, kw, out_h, out_w); + let mask = make_mask(batch, offset_groups, kh, kw, out_h, out_w); + let bias = make_bias(out_ch); + let opts = DeformConvOptions::new( + stride, + padding, + dilation, + weight_groups, + offset_groups, + ); + + bencher.bench_synced(|| { + module::deform_conv2d( + x.clone(), + offset.clone(), + weight.clone(), + Some(mask.clone()), + Some(bias.clone()), + opts.clone(), + ) + }); + } + + // Small batch of 2 + #[divan::bench(sample_count = 20, min_time = 0)] + fn deform_conv2d_2x8x16x16_k3x3_c16(bencher: Bencher) { + let (batch, in_ch, h, w) = (2, 8, 16, 16); + let (out_ch, kh, kw) = (16, 3, 3); + let (stride, padding, dilation) = ([1, 1], [1, 1], [1, 1]); + let (weight_groups, offset_groups) = (1, 1); + + let out_h = compute_output_size(h, kh, padding[0], stride[0], dilation[0]); + let out_w = compute_output_size(w, kw, padding[1], stride[1], dilation[1]); + + let x = make_input(batch, in_ch, h, w); + let weight = make_weight(out_ch, in_ch / weight_groups, kh, kw); + let offset = make_offset(batch, offset_groups, kh, kw, out_h, out_w); + let mask = make_mask(batch, offset_groups, kh, kw, out_h, out_w); + let bias = make_bias(out_ch); + let opts = DeformConvOptions::new( + stride, + padding, + dilation, + weight_groups, + offset_groups, + ); + + bencher.bench_synced(|| { + module::deform_conv2d( + x.clone(), + offset.clone(), + weight.clone(), + Some(mask.clone()), + Some(bias.clone()), + opts.clone(), + ) + }); + } + } + + #[divan::bench_group(name = "deform_conv2d_medium")] + mod deform_conv2d_medium { + use super::*; + + // Medium: 1x16x32x32 -> 32 output channels + #[divan::bench(sample_count = 20, min_time = 0)] + fn deform_conv2d_1x16x32x32_k3x3_c32(bencher: Bencher) { + let (batch, in_ch, h, w) = (1, 16, 32, 32); + let (out_ch, kh, kw) = (32, 3, 3); + let (stride, padding, dilation) = ([1, 1], [1, 1], [1, 1]); + let (weight_groups, offset_groups) = (1, 1); + + let out_h = compute_output_size(h, kh, padding[0], stride[0], dilation[0]); + let out_w = compute_output_size(w, kw, padding[1], stride[1], dilation[1]); + + let x = make_input(batch, in_ch, h, w); + let weight = make_weight(out_ch, in_ch / weight_groups, kh, kw); + let offset = make_offset(batch, offset_groups, kh, kw, out_h, out_w); + let mask = make_mask(batch, offset_groups, kh, kw, out_h, out_w); + let bias = make_bias(out_ch); + let opts = DeformConvOptions::new( + stride, + padding, + dilation, + weight_groups, + offset_groups, + ); + + bencher.bench_synced(|| { + module::deform_conv2d( + x.clone(), + offset.clone(), + weight.clone(), + Some(mask.clone()), + Some(bias.clone()), + opts.clone(), + ) + }); + } + + // Medium with weight groups + #[divan::bench(sample_count = 20, min_time = 0)] + fn deform_conv2d_1x16x32x32_k3x3_wg4(bencher: Bencher) { + let (batch, in_ch, h, w) = (1, 16, 32, 32); + let (out_ch, kh, kw) = (32, 3, 3); + let (stride, padding, dilation) = ([1, 1], [1, 1], [1, 1]); + let (weight_groups, offset_groups) = (4, 1); + + let out_h = compute_output_size(h, kh, padding[0], stride[0], dilation[0]); + let out_w = compute_output_size(w, kw, padding[1], stride[1], dilation[1]); + + let x = make_input(batch, in_ch, h, w); + let weight = make_weight(out_ch, in_ch / weight_groups, kh, kw); + let offset = make_offset(batch, offset_groups, kh, kw, out_h, out_w); + let mask = make_mask(batch, offset_groups, kh, kw, out_h, out_w); + let bias = make_bias(out_ch); + let opts = DeformConvOptions::new( + stride, + padding, + dilation, + weight_groups, + offset_groups, + ); + + bencher.bench_synced(|| { + module::deform_conv2d( + x.clone(), + offset.clone(), + weight.clone(), + Some(mask.clone()), + Some(bias.clone()), + opts.clone(), + ) + }); + } + + // Medium with offset groups + #[divan::bench(sample_count = 20, min_time = 0)] + fn deform_conv2d_1x16x32x32_k3x3_og4(bencher: Bencher) { + let (batch, in_ch, h, w) = (1, 16, 32, 32); + let (out_ch, kh, kw) = (32, 3, 3); + let (stride, padding, dilation) = ([1, 1], [1, 1], [1, 1]); + let (weight_groups, offset_groups) = (1, 4); + + let out_h = compute_output_size(h, kh, padding[0], stride[0], dilation[0]); + let out_w = compute_output_size(w, kw, padding[1], stride[1], dilation[1]); + + let x = make_input(batch, in_ch, h, w); + let weight = make_weight(out_ch, in_ch / weight_groups, kh, kw); + let offset = make_offset(batch, offset_groups, kh, kw, out_h, out_w); + let mask = make_mask(batch, offset_groups, kh, kw, out_h, out_w); + let bias = make_bias(out_ch); + let opts = DeformConvOptions::new( + stride, + padding, + dilation, + weight_groups, + offset_groups, + ); + + bencher.bench_synced(|| { + module::deform_conv2d( + x.clone(), + offset.clone(), + weight.clone(), + Some(mask.clone()), + Some(bias.clone()), + opts.clone(), + ) + }); + } + } + } + }; +} + +bench_backend!(backend, "backend"); diff --git a/crates/burn-backend-tests/benches/fft_ops.rs b/crates/burn-backend-tests/benches/fft_ops.rs new file mode 100644 index 0000000..5f6b201 --- /dev/null +++ b/crates/burn-backend-tests/benches/fft_ops.rs @@ -0,0 +1,138 @@ +//! Benchmarks for rfft / irfft. +//! +//! Run with: +//! ```bash +//! cargo bench --bench fft_ops +//! ``` + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{ + Tensor, TensorData, + signal::{irfft, rfft}, +}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("rfft / irfft benchmarks"); + println!(); + divan::main(); + common::report_failures(); +} + +fn make_signal_1d(n: usize) -> Tensor<1> { + let data: Vec = (0..n).map(|i| (i as f32 * 0.1).sin()).collect(); + Tensor::from_data(TensorData::new(data, [n]), &Default::default()) +} + +fn make_signal_2d(batch: usize, n: usize) -> Tensor<2> { + let data: Vec = (0..batch * n).map(|i| (i as f32 * 0.1).sin()).collect(); + Tensor::from_data(TensorData::new(data, [batch, n]), &Default::default()) +} + +#[divan::bench_group(name = "rfft_1d")] +mod rfft_1d { + use super::*; + + #[divan::bench] + fn n_256(bencher: Bencher) { + let s = make_signal_1d(256); + bencher.bench_synced(|| rfft(s.clone(), 0, None)); + } + + #[divan::bench] + fn n_1024(bencher: Bencher) { + let s = make_signal_1d(1024); + bencher.bench_synced(|| rfft(s.clone(), 0, None)); + } + + #[divan::bench] + fn n_4096(bencher: Bencher) { + let s = make_signal_1d(4096); + bencher.bench_synced(|| rfft(s.clone(), 0, None)); + } + + #[divan::bench] + fn n_16384(bencher: Bencher) { + let s = make_signal_1d(16384); + bencher.bench_synced(|| rfft(s.clone(), 0, None)); + } + + #[divan::bench] + fn n_65536(bencher: Bencher) { + let s = make_signal_1d(65536); + bencher.bench_synced(|| rfft(s.clone(), 0, None)); + } +} + +#[divan::bench_group(name = "rfft_2d_batch")] +mod rfft_2d_batch { + use super::*; + + #[divan::bench] + fn batch_16_n_1024(bencher: Bencher) { + let s = make_signal_2d(16, 1024); + bencher.bench_synced(|| rfft(s.clone(), 1, None)); + } + + #[divan::bench] + fn batch_64_n_1024(bencher: Bencher) { + let s = make_signal_2d(64, 1024); + bencher.bench_synced(|| rfft(s.clone(), 1, None)); + } + + #[divan::bench] + fn batch_256_n_256(bencher: Bencher) { + let s = make_signal_2d(256, 256); + bencher.bench_synced(|| rfft(s.clone(), 1, None)); + } +} + +#[divan::bench_group(name = "irfft_1d")] +mod irfft_1d { + use super::*; + + fn bench_irfft(bencher: Bencher, n: usize) { + // rfft is used to produce the input spectrum; if the backend doesn't implement it the + // setup panics before bench_synced's catch_unwind. Use try_setup to record and fall + // through to a no-op. + let Some((re, im)) = common::try_setup(|| { + let s = make_signal_1d(n); + rfft(s, 0, None) + }) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| irfft(re.clone(), im.clone(), 0, None)); + } + + #[divan::bench] + fn n_256(bencher: Bencher) { + bench_irfft(bencher, 256); + } + + #[divan::bench] + fn n_1024(bencher: Bencher) { + bench_irfft(bencher, 1024); + } + + #[divan::bench] + fn n_4096(bencher: Bencher) { + bench_irfft(bencher, 4096); + } + + #[divan::bench] + fn n_16384(bencher: Bencher) { + bench_irfft(bencher, 16384); + } + + #[divan::bench] + fn n_65536(bencher: Bencher) { + bench_irfft(bencher, 65536); + } +} diff --git a/crates/burn-backend-tests/benches/gather_scatter_ops.rs b/crates/burn-backend-tests/benches/gather_scatter_ops.rs new file mode 100644 index 0000000..579bd0a --- /dev/null +++ b/crates/burn-backend-tests/benches/gather_scatter_ops.rs @@ -0,0 +1,207 @@ +//! Benchmarks for gather/scatter operations. +//! +//! Run with: +//! ```bash +//! cargo bench --bench gather_scatter_ops --features simd +//! ``` + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{IndexingUpdateOp, Int, Tensor, TensorData}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Gather/scatter ops Benchmarks"); + println!(); + divan::main(); + common::report_failures(); +} + +fn make_tensor_2d(rows: usize, cols: usize) -> Tensor<2> { + let data: Vec = (0..rows * cols) + .map(|i| (i % 1000) as f32 / 1000.0) + .collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) +} + +fn make_indices_2d(rows: usize, cols: usize, max_idx: usize) -> Option> { + common::try_setup(|| { + let data: Vec = (0..rows * cols).map(|i| (i % max_idx) as i32).collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) + }) +} + +fn make_indices_1d(size: usize, max_idx: usize) -> Option> { + common::try_setup(|| { + let data: Vec = (0..size).map(|i| (i % max_idx) as i32).collect(); + Tensor::from_data(TensorData::new(data, [size]), &Default::default()) + }) +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "gather")] + mod gather { + use super::*; + + #[divan::bench] + fn s_256x256_dim1(bencher: Bencher) { + let tensor = make_tensor_2d(256, 256); + let Some(indices) = make_indices_2d(256, 128, 256) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| tensor.clone().gather(1, indices.clone())); + } + + #[divan::bench] + fn s_1024x1024_dim1(bencher: Bencher) { + let tensor = make_tensor_2d(1024, 1024); + let Some(indices) = make_indices_2d(1024, 512, 1024) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| tensor.clone().gather(1, indices.clone())); + } + + #[divan::bench] + fn s_256x256_dim0(bencher: Bencher) { + let tensor = make_tensor_2d(256, 256); + let Some(indices) = make_indices_2d(128, 256, 256) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| tensor.clone().gather(0, indices.clone())); + } + } + + #[divan::bench_group(name = "scatter_add")] + mod scatter_add { + use super::*; + + #[divan::bench] + fn s_256x256_dim1(bencher: Bencher) { + let tensor = make_tensor_2d(256, 256); + let Some(indices) = make_indices_2d(256, 128, 256) else { + bencher.bench(|| ()); + return; + }; + let values = make_tensor_2d(256, 128); + bencher.bench_synced(|| { + tensor.clone().scatter( + 1, + indices.clone(), + values.clone(), + IndexingUpdateOp::Add, + ) + }); + } + + #[divan::bench] + fn s_1024x1024_dim1(bencher: Bencher) { + let tensor = make_tensor_2d(1024, 1024); + let Some(indices) = make_indices_2d(1024, 512, 1024) else { + bencher.bench(|| ()); + return; + }; + let values = make_tensor_2d(1024, 512); + bencher.bench_synced(|| { + tensor.clone().scatter( + 1, + indices.clone(), + values.clone(), + IndexingUpdateOp::Add, + ) + }); + } + } + + #[divan::bench_group(name = "select")] + mod select { + use super::*; + + #[divan::bench] + fn s_256x256_dim0(bencher: Bencher) { + let tensor = make_tensor_2d(256, 256); + let Some(indices) = make_indices_1d(128, 256) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| tensor.clone().select(0, indices.clone())); + } + + #[divan::bench] + fn s_1024x1024_dim0(bencher: Bencher) { + let tensor = make_tensor_2d(1024, 1024); + let Some(indices) = make_indices_1d(512, 1024) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| tensor.clone().select(0, indices.clone())); + } + + #[divan::bench] + fn s_256x256_dim1(bencher: Bencher) { + let tensor = make_tensor_2d(256, 256); + let Some(indices) = make_indices_1d(128, 256) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| tensor.clone().select(1, indices.clone())); + } + } + + #[divan::bench_group(name = "select_add")] + mod select_add { + use super::*; + + #[divan::bench] + fn s_256x256_dim0(bencher: Bencher) { + let tensor = make_tensor_2d(256, 256); + let Some(indices) = make_indices_1d(128, 256) else { + bencher.bench(|| ()); + return; + }; + let values = make_tensor_2d(128, 256); + bencher.bench_synced(|| { + tensor.clone().select_assign( + 0, + indices.clone(), + values.clone(), + IndexingUpdateOp::Add, + ) + }); + } + + #[divan::bench] + fn s_1024x1024_dim0(bencher: Bencher) { + let tensor = make_tensor_2d(1024, 1024); + let Some(indices) = make_indices_1d(512, 1024) else { + bencher.bench(|| ()); + return; + }; + let values = make_tensor_2d(512, 1024); + bencher.bench_synced(|| { + tensor.clone().select_assign( + 0, + indices.clone(), + values.clone(), + IndexingUpdateOp::Add, + ) + }); + } + } + } + }; +} + +bench_backend!(backend, "backend"); diff --git a/crates/burn-backend-tests/benches/int_ops.rs b/crates/burn-backend-tests/benches/int_ops.rs new file mode 100644 index 0000000..bf2b23d --- /dev/null +++ b/crates/burn-backend-tests/benches/int_ops.rs @@ -0,0 +1,161 @@ +//! Benchmarks for integer operations: cast and random. +//! +//! Run with: +//! ```bash +//! cargo bench --bench int_ops --features simd,rayon +//! ``` + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{DType, Distribution, Int, Tensor, TensorData}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Integer Operations Benchmarks"); + println!("Memory allocation tracking enabled"); + println!(); + divan::main(); + common::report_failures(); +} + +fn make_int_tensor(shape: &[usize]) -> Option> { + common::try_setup(|| { + let size: usize = shape.iter().product(); + // Use values that fit in i8 (-128 to 127) so all casts work + let data: Vec = (0..size).map(|i| (i % 200) as i32 - 100).collect(); + Tensor::from_data(TensorData::new(data, shape.to_vec()), &Default::default()) + }) +} + +// ============================================================================= +// Int Cast Benchmarks +// ============================================================================= + +macro_rules! bench_cast_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "cast")] + mod cast { + use super::*; + + // Small tensor cast + #[divan::bench] + fn cast_i64_to_i32_64x64(bencher: Bencher) { + let Some(t) = make_int_tensor(&[64, 64]) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| t.clone().cast(DType::I32)); + } + + // Medium tensor cast + #[divan::bench] + fn cast_i64_to_i32_256x256(bencher: Bencher) { + let Some(t) = make_int_tensor(&[256, 256]) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| t.clone().cast(DType::I32)); + } + + // Large tensor cast + #[divan::bench] + fn cast_i64_to_i32_1024x1024(bencher: Bencher) { + let Some(t) = make_int_tensor(&[1024, 1024]) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| t.clone().cast(DType::I32)); + } + + // Cast to smaller type (i8) + #[divan::bench] + fn cast_i64_to_i8_256x256(bencher: Bencher) { + let Some(t) = make_int_tensor(&[256, 256]) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| t.clone().cast(DType::I8)); + } + } + } + }; +} + +bench_cast_backend!(backend_cast, "backend_cast"); + +// ============================================================================= +// Int Random Benchmarks +// ============================================================================= + +macro_rules! bench_random_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "random")] + mod random { + use super::*; + + // Small random tensor + #[divan::bench] + fn random_uniform_64x64(bencher: Bencher) { + bencher.bench_synced(|| { + Tensor::<2, Int>::random( + [64, 64], + Distribution::Uniform(0.0, 100.0), + &Default::default(), + ) + }); + } + + // Medium random tensor + #[divan::bench] + fn random_uniform_256x256(bencher: Bencher) { + bencher.bench_synced(|| { + Tensor::<2, Int>::random( + [256, 256], + Distribution::Uniform(0.0, 100.0), + &Default::default(), + ) + }); + } + + // Large random tensor + #[divan::bench] + fn random_uniform_1024x1024(bencher: Bencher) { + bencher.bench_synced(|| { + Tensor::<2, Int>::random( + [1024, 1024], + Distribution::Uniform(0.0, 100.0), + &Default::default(), + ) + }); + } + + // 3D random tensor (batch) + #[divan::bench] + fn random_uniform_batch_16x128x128(bencher: Bencher) { + bencher.bench_synced(|| { + Tensor::<3, Int>::random( + [16, 128, 128], + Distribution::Uniform(-1000.0, 1000.0), + &Default::default(), + ) + }); + } + } + } + }; +} + +bench_random_backend!(backend_random, "backend_random"); diff --git a/crates/burn-backend-tests/benches/interpolate_ops.rs b/crates/burn-backend-tests/benches/interpolate_ops.rs new file mode 100644 index 0000000..e605f67 --- /dev/null +++ b/crates/burn-backend-tests/benches/interpolate_ops.rs @@ -0,0 +1,221 @@ +//! Benchmarks for interpolation operations. +//! +//! Run with: +//! ```bash +//! cargo bench --bench interpolate_ops --features simd,rayon +//! ``` +//! +//! Memory allocation tracking is enabled via divan's AllocProfiler. + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{ + Tensor, TensorData, module, + ops::{InterpolateMode, InterpolateOptions}, +}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Interpolate Benchmarks"); + println!("Memory allocation tracking enabled"); + println!(); + divan::main(); + common::report_failures(); +} + +fn make_input(batch: usize, channels: usize, height: usize, width: usize) -> Tensor<4> { + let data: Vec = (0..batch * channels * height * width) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [batch, channels, height, width]), + &Default::default(), + ) +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "nearest")] + mod nearest { + use super::*; + + #[divan::bench] + fn upsample_2x_64x64_to_128x128(bencher: Bencher) { + let x = make_input(1, 3, 64, 64); + let opts = InterpolateOptions::new(InterpolateMode::Nearest); + bencher + .bench_synced(|| module::interpolate(x.clone(), [128, 128], opts.clone())); + } + + #[divan::bench] + fn upsample_4x_32x32_to_128x128(bencher: Bencher) { + let x = make_input(1, 3, 32, 32); + let opts = InterpolateOptions::new(InterpolateMode::Nearest); + bencher + .bench_synced(|| module::interpolate(x.clone(), [128, 128], opts.clone())); + } + + #[divan::bench] + fn downsample_2x_256x256_to_128x128(bencher: Bencher) { + let x = make_input(1, 3, 256, 256); + let opts = InterpolateOptions::new(InterpolateMode::Nearest); + bencher + .bench_synced(|| module::interpolate(x.clone(), [128, 128], opts.clone())); + } + + #[divan::bench] + fn batch8_upsample_2x_64x64_to_128x128(bencher: Bencher) { + let x = make_input(8, 3, 64, 64); + let opts = InterpolateOptions::new(InterpolateMode::Nearest); + bencher + .bench_synced(|| module::interpolate(x.clone(), [128, 128], opts.clone())); + } + + #[divan::bench] + fn channels64_upsample_2x_32x32_to_64x64(bencher: Bencher) { + let x = make_input(1, 64, 32, 32); + let opts = InterpolateOptions::new(InterpolateMode::Nearest); + bencher.bench_synced(|| module::interpolate(x.clone(), [64, 64], opts.clone())); + } + } + + #[divan::bench_group(name = "bilinear")] + mod bilinear { + use super::*; + + #[divan::bench] + fn upsample_2x_64x64_to_128x128(bencher: Bencher) { + let x = make_input(1, 3, 64, 64); + let opts = InterpolateOptions::new(InterpolateMode::Bilinear); + bencher + .bench_synced(|| module::interpolate(x.clone(), [128, 128], opts.clone())); + } + + #[divan::bench] + fn upsample_4x_32x32_to_128x128(bencher: Bencher) { + let x = make_input(1, 3, 32, 32); + let opts = InterpolateOptions::new(InterpolateMode::Bilinear); + bencher + .bench_synced(|| module::interpolate(x.clone(), [128, 128], opts.clone())); + } + + #[divan::bench] + fn downsample_2x_256x256_to_128x128(bencher: Bencher) { + let x = make_input(1, 3, 256, 256); + let opts = InterpolateOptions::new(InterpolateMode::Bilinear); + bencher + .bench_synced(|| module::interpolate(x.clone(), [128, 128], opts.clone())); + } + + #[divan::bench] + fn batch8_upsample_2x_64x64_to_128x128(bencher: Bencher) { + let x = make_input(8, 3, 64, 64); + let opts = InterpolateOptions::new(InterpolateMode::Bilinear); + bencher + .bench_synced(|| module::interpolate(x.clone(), [128, 128], opts.clone())); + } + + #[divan::bench] + fn channels64_upsample_2x_32x32_to_64x64(bencher: Bencher) { + let x = make_input(1, 64, 32, 32); + let opts = InterpolateOptions::new(InterpolateMode::Bilinear); + bencher.bench_synced(|| module::interpolate(x.clone(), [64, 64], opts.clone())); + } + + // Segmentation-model sized shapes mirroring the user's + // 488x448 input + UNet up/down pattern from issue #64 + // item 4. Covers the typical spatial dims across + // ConvNeXt-style downsample stages. + #[divan::bench] + fn model_downsample_488x448_to_244x224(bencher: Bencher) { + let x = make_input(1, 3, 488, 448); + let opts = InterpolateOptions::new(InterpolateMode::Bilinear); + bencher + .bench_synced(|| module::interpolate(x.clone(), [244, 224], opts.clone())); + } + + #[divan::bench] + fn model_upsample_244x224_to_488x448(bencher: Bencher) { + let x = make_input(1, 48, 244, 224); + let opts = InterpolateOptions::new(InterpolateMode::Bilinear); + bencher + .bench_synced(|| module::interpolate(x.clone(), [488, 448], opts.clone())); + } + + #[divan::bench] + fn model_upsample_61x56_to_122x112(bencher: Bencher) { + let x = make_input(1, 192, 61, 56); + let opts = InterpolateOptions::new(InterpolateMode::Bilinear); + bencher + .bench_synced(|| module::interpolate(x.clone(), [122, 112], opts.clone())); + } + + // Explicit align_corners=false variant at a model size, + // since the default is `true` everywhere else. + #[divan::bench] + fn model_upsample_61x56_to_122x112_halfpixel(bencher: Bencher) { + let x = make_input(1, 192, 61, 56); + let opts = InterpolateOptions::new(InterpolateMode::Bilinear) + .with_align_corners(false); + bencher + .bench_synced(|| module::interpolate(x.clone(), [122, 112], opts.clone())); + } + } + + #[divan::bench_group(name = "bicubic")] + mod bicubic { + use super::*; + + #[divan::bench] + fn upsample_2x_64x64_to_128x128(bencher: Bencher) { + let x = make_input(1, 3, 64, 64); + let opts = InterpolateOptions::new(InterpolateMode::Bicubic); + bencher + .bench_synced(|| module::interpolate(x.clone(), [128, 128], opts.clone())); + } + + #[divan::bench] + fn upsample_4x_32x32_to_128x128(bencher: Bencher) { + let x = make_input(1, 3, 32, 32); + let opts = InterpolateOptions::new(InterpolateMode::Bicubic); + bencher + .bench_synced(|| module::interpolate(x.clone(), [128, 128], opts.clone())); + } + + #[divan::bench] + fn downsample_2x_256x256_to_128x128(bencher: Bencher) { + let x = make_input(1, 3, 256, 256); + let opts = InterpolateOptions::new(InterpolateMode::Bicubic); + bencher + .bench_synced(|| module::interpolate(x.clone(), [128, 128], opts.clone())); + } + + #[divan::bench] + fn batch8_upsample_2x_64x64_to_128x128(bencher: Bencher) { + let x = make_input(8, 3, 64, 64); + let opts = InterpolateOptions::new(InterpolateMode::Bicubic); + bencher + .bench_synced(|| module::interpolate(x.clone(), [128, 128], opts.clone())); + } + + #[divan::bench] + fn channels64_upsample_2x_32x32_to_64x64(bencher: Bencher) { + let x = make_input(1, 64, 32, 32); + let opts = InterpolateOptions::new(InterpolateMode::Bicubic); + bencher.bench_synced(|| module::interpolate(x.clone(), [64, 64], opts.clone())); + } + } + } + }; +} + +bench_backend!(backend, "backend"); diff --git a/crates/burn-backend-tests/benches/matmul.rs b/crates/burn-backend-tests/benches/matmul.rs new file mode 100644 index 0000000..5fd3a54 --- /dev/null +++ b/crates/burn-backend-tests/benches/matmul.rs @@ -0,0 +1,364 @@ +//! Benchmarks for matrix multiplication. +//! +//! Run with: +//! ```bash +//! cargo bench --bench matmul --features simd,gemm +//! ``` +//! +//! Memory allocation tracking is enabled via divan's AllocProfiler. + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{Tensor, TensorData}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Matrix Multiplication Benchmarks"); + println!("Memory allocation tracking enabled"); + println!(); + divan::main(); + common::report_failures(); +} + +fn make_matrix(rows: usize, cols: usize) -> Tensor<2> { + let data: Vec = (0..rows * cols) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) +} + +fn make_batch_matrix(batch: usize, rows: usize, cols: usize) -> Tensor<3> { + let data: Vec = (0..batch * rows * cols) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [batch, rows, cols]), + &Default::default(), + ) +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + // Square matrices + #[divan::bench_group(name = "square")] + mod square { + use super::*; + + #[divan::bench] + fn matmul_64x64(bencher: Bencher) { + let a = make_matrix(64, 64); + let b = make_matrix(64, 64); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + #[divan::bench] + fn matmul_128x128(bencher: Bencher) { + let a = make_matrix(128, 128); + let b = make_matrix(128, 128); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + #[divan::bench] + fn matmul_256x256(bencher: Bencher) { + let a = make_matrix(256, 256); + let b = make_matrix(256, 256); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + #[divan::bench] + fn matmul_512x512(bencher: Bencher) { + let a = make_matrix(512, 512); + let b = make_matrix(512, 512); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + #[divan::bench] + fn matmul_1024x1024(bencher: Bencher) { + let a = make_matrix(1024, 1024); + let b = make_matrix(1024, 1024); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + } + + // Rectangular matrices (common in neural networks) + #[divan::bench_group(name = "rectangular")] + mod rectangular { + use super::*; + + // Linear layer: [batch, in] x [in, out] + #[divan::bench] + fn linear_256x512_512x256(bencher: Bencher) { + let a = make_matrix(256, 512); + let b = make_matrix(512, 256); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + // Attention: [seq, hidden] x [hidden, seq] + #[divan::bench] + fn attention_512x64_64x512(bencher: Bencher) { + let a = make_matrix(512, 64); + let b = make_matrix(64, 512); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + // Wide matrix (embedding lookup style) + #[divan::bench] + fn wide_128x1024_1024x128(bencher: Bencher) { + let a = make_matrix(128, 1024); + let b = make_matrix(1024, 128); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + } + + // Batched matmul (common in transformers) + #[divan::bench_group(name = "batched")] + mod batched { + use super::*; + + #[divan::bench] + fn batch8_64x64(bencher: Bencher) { + let a = make_batch_matrix(8, 64, 64); + let b = make_batch_matrix(8, 64, 64); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + #[divan::bench] + fn batch16_128x128(bencher: Bencher) { + let a = make_batch_matrix(16, 128, 128); + let b = make_batch_matrix(16, 128, 128); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + #[divan::bench] + fn batch32_64x64(bencher: Bencher) { + let a = make_batch_matrix(32, 64, 64); + let b = make_batch_matrix(32, 64, 64); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + // Transformer attention heads + #[divan::bench] + fn heads12_seq512_dim64(bencher: Bencher) { + let a = make_batch_matrix(12, 512, 64); + let b = make_batch_matrix(12, 64, 512); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + } + + // Transposed inputs (tests contiguous conversion overhead) + #[divan::bench_group(name = "transposed")] + mod transposed { + use super::*; + + #[divan::bench] + fn lhs_transposed_256x256(bencher: Bencher) { + let a = make_matrix(256, 256).transpose(); + let b = make_matrix(256, 256); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + #[divan::bench] + fn rhs_transposed_256x256(bencher: Bencher) { + let a = make_matrix(256, 256); + let b = make_matrix(256, 256).transpose(); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + #[divan::bench] + fn both_transposed_256x256(bencher: Bencher) { + let a = make_matrix(256, 256).transpose(); + let b = make_matrix(256, 256).transpose(); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + } + } + }; +} + +bench_backend!(backend, "backend"); + +// Integer matmul benchmarks +fn make_int_matrix(rows: usize, cols: usize) -> Option> { + common::try_setup(|| { + let data: Vec = (0..rows * cols).map(|i| (i % 1000) as i32 - 500).collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) + }) +} + +fn make_int_batch_matrix( + batch: usize, + rows: usize, + cols: usize, +) -> Option> { + common::try_setup(|| { + let data: Vec = (0..batch * rows * cols) + .map(|i| (i % 1000) as i32 - 500) + .collect(); + Tensor::from_data( + TensorData::new(data, [batch, rows, cols]), + &Default::default(), + ) + }) +} + +macro_rules! bench_int_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "int_square")] + mod square { + use super::*; + + #[divan::bench] + fn matmul_64x64(bencher: Bencher) { + let Some(a) = make_int_matrix(64, 64) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_int_matrix(64, 64) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + #[divan::bench] + fn matmul_128x128(bencher: Bencher) { + let Some(a) = make_int_matrix(128, 128) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_int_matrix(128, 128) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + #[divan::bench] + fn matmul_256x256(bencher: Bencher) { + let Some(a) = make_int_matrix(256, 256) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_int_matrix(256, 256) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + #[divan::bench] + fn matmul_512x512(bencher: Bencher) { + let Some(a) = make_int_matrix(512, 512) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_int_matrix(512, 512) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + } + + #[divan::bench_group(name = "int_batched")] + mod batched { + use super::*; + + #[divan::bench] + fn batch8_64x64(bencher: Bencher) { + let Some(a) = make_int_batch_matrix(8, 64, 64) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_int_batch_matrix(8, 64, 64) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + #[divan::bench] + fn batch16_128x128(bencher: Bencher) { + let Some(a) = make_int_batch_matrix(16, 128, 128) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_int_batch_matrix(16, 128, 128) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + } + } + }; +} + +bench_int_backend!(backend_int, "backend_int"); + +// Broadcast matmul benchmarks +fn make_batch_matrix_4d(b1: usize, b2: usize, rows: usize, cols: usize) -> Tensor<4> { + let data: Vec = (0..b1 * b2 * rows * cols) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [b1, b2, rows, cols]), + &Default::default(), + ) +} + +macro_rules! bench_broadcast_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + // Broadcast: [1, M, K] x [batch, K, N] -> [batch, M, N] + #[divan::bench] + fn broadcast_1_vs_8_64x64(bencher: Bencher) { + let a = make_batch_matrix(1, 64, 64); + let b = make_batch_matrix(8, 64, 64); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + // Broadcast: [batch, M, K] x [1, K, N] -> [batch, M, N] + #[divan::bench] + fn broadcast_8_vs_1_64x64(bencher: Bencher) { + let a = make_batch_matrix(8, 64, 64); + let b = make_batch_matrix(1, 64, 64); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + // 4D broadcast: [2, 1, M, K] x [1, 4, K, N] -> [2, 4, M, N] + #[divan::bench] + fn broadcast_4d_2x1_vs_1x4_32x32(bencher: Bencher) { + let a = make_batch_matrix_4d(2, 1, 32, 32); + let b = make_batch_matrix_4d(1, 4, 32, 32); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + // Larger 4D broadcast + #[divan::bench] + fn broadcast_4d_4x1_vs_1x4_64x64(bencher: Bencher) { + let a = make_batch_matrix_4d(4, 1, 64, 64); + let b = make_batch_matrix_4d(1, 4, 64, 64); + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + } + }; +} + +bench_broadcast_backend!(backend_broadcast, "backend_broadcast"); diff --git a/crates/burn-backend-tests/benches/norm_ops.rs b/crates/burn-backend-tests/benches/norm_ops.rs new file mode 100644 index 0000000..9736707 --- /dev/null +++ b/crates/burn-backend-tests/benches/norm_ops.rs @@ -0,0 +1,323 @@ +//! Benchmarks for normalization ops. +//! +//! Covers `burn::nn::LayerNorm::forward`'s decomposed primitive-ops path +//! (what you get by default from burn-nn), the backend `layer_norm` fast +//! path, and a `reshape -> layer_norm -> reshape-back` composite that +//! mirrors a ConvNeXt-style channels-last layer norm. See issue #64 item 2. +//! +//! Run with: +//! ```bash +//! cargo bench --bench norm_ops --features simd,rayon +//! ``` + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{Tensor, TensorData}; +use divan::{AllocProfiler, Bencher}; + +/// Route through the `B::layer_norm` backend hook. On Flex this hits the +/// fused override; on NdArray this hits the `ModuleOps` default decomposition. +fn trait_layer_norm( + input: Tensor, + gamma: Tensor<1>, + beta: Tensor<1>, + epsilon: f64, +) -> Tensor { + burn_tensor::module::layer_norm(input, gamma, Some(beta), epsilon) +} + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Normalization Benchmarks"); + println!(); + divan::main(); + common::report_failures(); +} + +fn make_tensor_3d(d0: usize, d1: usize, d2: usize) -> Tensor<3> { + let data: Vec = (0..d0 * d1 * d2) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data(TensorData::new(data, [d0, d1, d2]), &Default::default()) +} + +fn make_tensor_4d(d0: usize, d1: usize, d2: usize, d3: usize) -> Tensor<4> { + let data: Vec = (0..d0 * d1 * d2 * d3) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data(TensorData::new(data, [d0, d1, d2, d3]), &Default::default()) +} + +fn make_gamma_beta(d_model: usize) -> (Tensor<1>, Tensor<1>) { + let gamma: Vec = (0..d_model) + .map(|i| 0.5 + (i as f32 / d_model as f32)) + .collect(); + let beta: Vec = (0..d_model).map(|i| i as f32 / d_model as f32).collect(); + ( + Tensor::from_data(TensorData::new(gamma, [d_model]), &Default::default()), + Tensor::from_data(TensorData::new(beta, [d_model]), &Default::default()), + ) +} + +/// Decomposed layer_norm, mirroring what `burn::nn::LayerNorm::forward` +/// does via primitive tensor ops. Applies along the last dim. +fn decomposed_layer_norm( + input: Tensor, + gamma: Tensor<1>, + beta: Tensor<1>, + epsilon: f32, +) -> Tensor { + let dim = D - 1; + let mean = input.clone().mean_dim(dim); + let centered = input - mean; + let var = centered.clone().powi_scalar(2).mean_dim(dim); + let normalized = centered / (var + epsilon).sqrt(); + normalized * gamma.unsqueeze() + beta.unsqueeze() +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + // ConvNeXt-style: small last-dim C (48, 96, 192, 384), large + // N*H*W. This is the shape the user's segmentation model hits + // at the channels-last layer norm between convolution blocks. + #[divan::bench_group(name = "convnext_layer_norm_decomposed")] + mod convnext_layer_norm_decomposed { + use super::*; + + #[divan::bench] + fn c48_hw244x224(bencher: Bencher) { + let x = make_tensor_3d(1, 244 * 224, 48); + let (g, b) = make_gamma_beta(48); + bencher.bench_synced(|| { + decomposed_layer_norm(x.clone(), g.clone(), b.clone(), 1e-5) + }); + } + + #[divan::bench] + fn c96_hw122x112(bencher: Bencher) { + let x = make_tensor_3d(1, 122 * 112, 96); + let (g, b) = make_gamma_beta(96); + bencher.bench_synced(|| { + decomposed_layer_norm(x.clone(), g.clone(), b.clone(), 1e-5) + }); + } + + #[divan::bench] + fn c192_hw61x56(bencher: Bencher) { + let x = make_tensor_3d(1, 61 * 56, 192); + let (g, b) = make_gamma_beta(192); + bencher.bench_synced(|| { + decomposed_layer_norm(x.clone(), g.clone(), b.clone(), 1e-5) + }); + } + + #[divan::bench] + fn c384_hw30x28(bencher: Bencher) { + let x = make_tensor_3d(1, 30 * 28, 384); + let (g, b) = make_gamma_beta(384); + bencher.bench_synced(|| { + decomposed_layer_norm(x.clone(), g.clone(), b.clone(), 1e-5) + }); + } + + // Transformer-style: typical `[batch=2, seq=512, d=768]`. + #[divan::bench] + fn d768_seq512_b2(bencher: Bencher) { + let x = make_tensor_3d(2, 512, 768); + let (g, b) = make_gamma_beta(768); + bencher.bench_synced(|| { + decomposed_layer_norm(x.clone(), g.clone(), b.clone(), 1e-5) + }); + } + } + + // Routes through `B::layer_norm`; compare against the + // `convnext_layer_norm_decomposed` group above. + #[divan::bench_group(name = "convnext_layer_norm_fused_trait")] + mod convnext_layer_norm_fused_trait { + use super::*; + + #[divan::bench] + fn c48_hw244x224(bencher: Bencher) { + let x = make_tensor_3d(1, 244 * 224, 48); + let (g, b) = make_gamma_beta(48); + bencher + .bench_synced(|| trait_layer_norm(x.clone(), g.clone(), b.clone(), 1e-5)); + } + + #[divan::bench] + fn c96_hw122x112(bencher: Bencher) { + let x = make_tensor_3d(1, 122 * 112, 96); + let (g, b) = make_gamma_beta(96); + bencher + .bench_synced(|| trait_layer_norm(x.clone(), g.clone(), b.clone(), 1e-5)); + } + + #[divan::bench] + fn c192_hw61x56(bencher: Bencher) { + let x = make_tensor_3d(1, 61 * 56, 192); + let (g, b) = make_gamma_beta(192); + bencher + .bench_synced(|| trait_layer_norm(x.clone(), g.clone(), b.clone(), 1e-5)); + } + + #[divan::bench] + fn c384_hw30x28(bencher: Bencher) { + let x = make_tensor_3d(1, 30 * 28, 384); + let (g, b) = make_gamma_beta(384); + bencher + .bench_synced(|| trait_layer_norm(x.clone(), g.clone(), b.clone(), 1e-5)); + } + + #[divan::bench] + fn d768_seq512_b2(bencher: Bencher) { + let x = make_tensor_3d(2, 512, 768); + let (g, b) = make_gamma_beta(768); + bencher + .bench_synced(|| trait_layer_norm(x.clone(), g.clone(), b.clone(), 1e-5)); + } + } + + // Drill-down: time each primitive op used inside the + // decomposition, so we can see which one drives the gap. + // Shape is [1, 244*224, 48], matching the worst case above. + #[divan::bench_group(name = "layer_norm_primitive_drilldown")] + mod layer_norm_primitive_drilldown { + use super::*; + + fn shape() -> Tensor<3> { + make_tensor_3d(1, 244 * 224, 48) + } + + #[divan::bench] + fn op1_mean_dim_last(bencher: Bencher) { + let x = shape(); + bencher.bench_synced(|| x.clone().mean_dim(2)); + } + + #[divan::bench] + fn op2_broadcast_sub(bencher: Bencher) { + let x = shape(); + let m = x.clone().mean_dim(2); + bencher.bench_synced(|| x.clone() - m.clone()); + } + + #[divan::bench] + fn op3_powi_scalar_2(bencher: Bencher) { + let x = shape(); + bencher.bench_synced(|| x.clone().powi_scalar(2)); + } + + #[divan::bench] + fn op4_broadcast_div(bencher: Bencher) { + let x = shape(); + let m = x.clone().mean_dim(2); + bencher.bench_synced(|| x.clone() / m.clone()); + } + + #[divan::bench] + fn op5_broadcast_mul_1d(bencher: Bencher) { + let x = shape(); + let (g, _) = make_gamma_beta(48); + bencher.bench_synced(|| x.clone() * g.clone().unsqueeze()); + } + + #[divan::bench] + fn op6_broadcast_add_1d(bencher: Bencher) { + let x = shape(); + let (_, b) = make_gamma_beta(48); + bencher.bench_synced(|| x.clone() + b.clone().unsqueeze()); + } + } + + // Drill-down for the "non-contig-input layer_norm" case. + // Input is x.permute([0, 2, 3, 1]) from shape + // [1, 48, 244, 224], so the permuted last dim has an + // absurd stride (54656) and hurts every op that needs + // row-contig access. + #[divan::bench_group(name = "permuted_input_drilldown")] + mod permuted_input_drilldown { + use super::*; + + fn permuted() -> Tensor<4> { + let x = make_tensor_4d(1, 48, 244, 224); + x.permute([0, 2, 3, 1]) + } + + #[divan::bench] + fn p1_permute_only(bencher: Bencher) { + let x = make_tensor_4d(1, 48, 244, 224); + bencher.bench_synced(|| x.clone().permute([0, 2, 3, 1])); + } + + #[divan::bench] + fn p2_mean_dim_last_on_permuted(bencher: Bencher) { + let p = permuted(); + bencher.bench_synced(|| p.clone().mean_dim(3)); + } + + #[divan::bench] + fn p3_broadcast_sub_on_permuted(bencher: Bencher) { + let p = permuted(); + let m = p.clone().mean_dim(3); + bencher.bench_synced(|| p.clone() - m.clone()); + } + } + + // ConvNeXt pattern: input is (N, C, H, W), reshape/permute to + // (N, H, W, C), layer_norm over C, permute back. The full + // composite is what the user's benchmark bundle called + // "reshape + dims + layer_norm". + #[divan::bench_group(name = "convnext_reshape_layer_norm_composite")] + mod convnext_reshape_layer_norm_composite { + use super::*; + + fn run(x: Tensor<4>, g: Tensor<1>, b: Tensor<1>) -> Tensor<4> { + // (N, C, H, W) -> (N, H, W, C) + let y = x.permute([0, 2, 3, 1]); + let normed = decomposed_layer_norm(y, g, b, 1e-5); + // (N, H, W, C) -> (N, C, H, W) + normed.permute([0, 3, 1, 2]) + } + + #[divan::bench] + fn c48_488x448(bencher: Bencher) { + let x = make_tensor_4d(1, 48, 244, 224); + let (g, b) = make_gamma_beta(48); + bencher.bench_synced(|| run::<3>(x.clone(), g.clone(), b.clone())); + } + + #[divan::bench] + fn c96_122x112(bencher: Bencher) { + let x = make_tensor_4d(1, 96, 122, 112); + let (g, b) = make_gamma_beta(96); + bencher.bench_synced(|| run::<3>(x.clone(), g.clone(), b.clone())); + } + + #[divan::bench] + fn c192_61x56(bencher: Bencher) { + let x = make_tensor_4d(1, 192, 61, 56); + let (g, b) = make_gamma_beta(192); + bencher.bench_synced(|| run::<3>(x.clone(), g.clone(), b.clone())); + } + + #[divan::bench] + fn c384_30x28(bencher: Bencher) { + let x = make_tensor_4d(1, 384, 30, 28); + let (g, b) = make_gamma_beta(384); + bencher.bench_synced(|| run::<3>(x.clone(), g.clone(), b.clone())); + } + } + } + }; +} + +bench_backend!(backend, "backend"); diff --git a/crates/burn-backend-tests/benches/pool_ops.rs b/crates/burn-backend-tests/benches/pool_ops.rs new file mode 100644 index 0000000..d169b66 --- /dev/null +++ b/crates/burn-backend-tests/benches/pool_ops.rs @@ -0,0 +1,237 @@ +//! Benchmarks for pooling operations. +//! +//! Run with: +//! ```bash +//! cargo bench --bench pool_ops --features simd,rayon +//! ``` +//! +//! Memory allocation tracking is enabled via divan's AllocProfiler. + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{Tensor, TensorData, module}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Pooling Benchmarks"); + println!("Memory allocation tracking enabled"); + println!(); + divan::main(); + common::report_failures(); +} + +fn make_input_2d(batch: usize, channels: usize, height: usize, width: usize) -> Tensor<4> { + let data: Vec = (0..batch * channels * height * width) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [batch, channels, height, width]), + &Default::default(), + ) +} + +fn make_input_1d(batch: usize, channels: usize, length: usize) -> Tensor<3> { + let data: Vec = (0..batch * channels * length) + .map(|i| ((i % 1000) as f32 / 1000.0) - 0.5) + .collect(); + Tensor::from_data( + TensorData::new(data, [batch, channels, length]), + &Default::default(), + ) +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "max_pool2d")] + mod max_pool2d { + use super::*; + + #[divan::bench] + fn max_pool2d_1x64x56x56_k3x3_s2(bencher: Bencher) { + let x = make_input_2d(1, 64, 56, 56); + bencher.bench_synced(|| { + module::max_pool2d(x.clone(), [3, 3], [2, 2], [1, 1], [1, 1], false) + }); + } + + #[divan::bench] + fn max_pool2d_8x64x56x56_k3x3_s2(bencher: Bencher) { + let x = make_input_2d(8, 64, 56, 56); + bencher.bench_synced(|| { + module::max_pool2d(x.clone(), [3, 3], [2, 2], [1, 1], [1, 1], false) + }); + } + + #[divan::bench] + fn max_pool2d_16x128x28x28_k2x2_s2(bencher: Bencher) { + let x = make_input_2d(16, 128, 28, 28); + bencher.bench_synced(|| { + module::max_pool2d(x.clone(), [2, 2], [2, 2], [0, 0], [1, 1], false) + }); + } + + #[divan::bench] + fn max_pool2d_1x512x14x14_k2x2_s2(bencher: Bencher) { + let x = make_input_2d(1, 512, 14, 14); + bencher.bench_synced(|| { + module::max_pool2d(x.clone(), [2, 2], [2, 2], [0, 0], [1, 1], false) + }); + } + } + + #[divan::bench_group(name = "max_pool2d_resnet")] + mod max_pool2d_resnet { + use super::*; + + #[divan::bench] + fn resnet_maxpool_1x64x112x112_k3x3_s2(bencher: Bencher) { + // ResNet initial max pool after first conv + let x = make_input_2d(1, 64, 112, 112); + bencher.bench_synced(|| { + module::max_pool2d(x.clone(), [3, 3], [2, 2], [1, 1], [1, 1], false) + }); + } + + #[divan::bench] + fn resnet_maxpool_8x64x112x112_k3x3_s2(bencher: Bencher) { + let x = make_input_2d(8, 64, 112, 112); + bencher.bench_synced(|| { + module::max_pool2d(x.clone(), [3, 3], [2, 2], [1, 1], [1, 1], false) + }); + } + + #[divan::bench] + fn resnet_maxpool_16x64x112x112_k3x3_s2(bencher: Bencher) { + let x = make_input_2d(16, 64, 112, 112); + bencher.bench_synced(|| { + module::max_pool2d(x.clone(), [3, 3], [2, 2], [1, 1], [1, 1], false) + }); + } + } + + #[divan::bench_group(name = "avg_pool2d")] + mod avg_pool2d { + use super::*; + + #[divan::bench] + fn avg_pool2d_1x64x56x56_k3x3_s2(bencher: Bencher) { + let x = make_input_2d(1, 64, 56, 56); + bencher.bench_synced(|| { + module::avg_pool2d(x.clone(), [3, 3], [2, 2], [1, 1], false, false) + }); + } + + #[divan::bench] + fn avg_pool2d_8x64x56x56_k3x3_s2(bencher: Bencher) { + let x = make_input_2d(8, 64, 56, 56); + bencher.bench_synced(|| { + module::avg_pool2d(x.clone(), [3, 3], [2, 2], [1, 1], false, false) + }); + } + + #[divan::bench] + fn avg_pool2d_16x128x28x28_k2x2_s2(bencher: Bencher) { + let x = make_input_2d(16, 128, 28, 28); + bencher.bench_synced(|| { + module::avg_pool2d(x.clone(), [2, 2], [2, 2], [0, 0], false, false) + }); + } + } + + #[divan::bench_group(name = "adaptive_avg_pool2d")] + mod adaptive_avg_pool2d { + use super::*; + + #[divan::bench] + fn adaptive_avg_pool2d_1x512x7x7_to_1x1(bencher: Bencher) { + // Global average pooling (common in ResNet final layer) + let x = make_input_2d(1, 512, 7, 7); + bencher.bench_synced(|| module::adaptive_avg_pool2d(x.clone(), [1, 1])); + } + + #[divan::bench] + fn adaptive_avg_pool2d_8x512x7x7_to_1x1(bencher: Bencher) { + let x = make_input_2d(8, 512, 7, 7); + bencher.bench_synced(|| module::adaptive_avg_pool2d(x.clone(), [1, 1])); + } + + #[divan::bench] + fn adaptive_avg_pool2d_16x2048x7x7_to_1x1(bencher: Bencher) { + // ResNet-50/101/152 final layer + let x = make_input_2d(16, 2048, 7, 7); + bencher.bench_synced(|| module::adaptive_avg_pool2d(x.clone(), [1, 1])); + } + + #[divan::bench] + fn adaptive_avg_pool2d_1x256x56x56_to_7x7(bencher: Bencher) { + // Downsampling to fixed size + let x = make_input_2d(1, 256, 56, 56); + bencher.bench_synced(|| module::adaptive_avg_pool2d(x.clone(), [7, 7])); + } + } + + #[divan::bench_group(name = "max_pool1d")] + mod max_pool1d { + use super::*; + + #[divan::bench] + fn max_pool1d_1x64x256_k3_s2(bencher: Bencher) { + let x = make_input_1d(1, 64, 256); + bencher.bench_synced(|| module::max_pool1d(x.clone(), 3, 2, 1, 1, false)); + } + + #[divan::bench] + fn max_pool1d_8x128x512_k3_s2(bencher: Bencher) { + let x = make_input_1d(8, 128, 512); + bencher.bench_synced(|| module::max_pool1d(x.clone(), 3, 2, 1, 1, false)); + } + + #[divan::bench] + fn max_pool1d_16x256x1024_k3_s2(bencher: Bencher) { + let x = make_input_1d(16, 256, 1024); + bencher.bench_synced(|| module::max_pool1d(x.clone(), 3, 2, 1, 1, false)); + } + } + + #[divan::bench_group(name = "pool_kernel_sizes")] + mod pool_kernel_sizes { + use super::*; + + #[divan::bench] + fn max_pool2d_k2x2(bencher: Bencher) { + let x = make_input_2d(4, 64, 56, 56); + bencher.bench_synced(|| { + module::max_pool2d(x.clone(), [2, 2], [2, 2], [0, 0], [1, 1], false) + }); + } + + #[divan::bench] + fn max_pool2d_k3x3(bencher: Bencher) { + let x = make_input_2d(4, 64, 56, 56); + bencher.bench_synced(|| { + module::max_pool2d(x.clone(), [3, 3], [2, 2], [1, 1], [1, 1], false) + }); + } + + #[divan::bench] + fn max_pool2d_k5x5(bencher: Bencher) { + let x = make_input_2d(4, 64, 56, 56); + bencher.bench_synced(|| { + module::max_pool2d(x.clone(), [5, 5], [2, 2], [2, 2], [1, 1], false) + }); + } + } + } + }; +} + +bench_backend!(backend, "backend"); diff --git a/crates/burn-backend-tests/benches/quantization_ops.rs b/crates/burn-backend-tests/benches/quantization_ops.rs new file mode 100644 index 0000000..5f65186 --- /dev/null +++ b/crates/burn-backend-tests/benches/quantization_ops.rs @@ -0,0 +1,341 @@ +//! Benchmarks for quantized tensor operations. +//! +//! Run with: +//! ```bash +//! cargo bench --bench quantization_ops +//! ``` + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{Device, Tensor, TensorData}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Quantized ops Benchmarks"); + println!(); + divan::main(); + common::report_failures(); +} + +const SMALL: usize = 64 * 64; // 4K elements +const MEDIUM: usize = 256 * 256; // 64K elements +const LARGE: usize = 1024 * 1024; // 1M elements + +fn make_tensor(size: usize) -> Tensor<1> { + let data: Vec = (0..size) + .map(|i| (i % 1000) as f32 / 1000.0 - 0.5) + .collect(); + Tensor::from_data(TensorData::new(data, [size]), &Default::default()) +} + +fn make_matrix(rows: usize, cols: usize) -> Tensor<2> { + let data: Vec = (0..rows * cols) + .map(|i| (i % 1000) as f32 / 1000.0 - 0.5) + .collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) +} + +// Returns None (and records a setup failure) on backends without quantization support, so the +// caller can fall back to a no-op bench instead of taking down the whole binary. +fn make_qtensor(size: usize) -> Option> { + common::try_setup(|| { + make_tensor(size).quantize_dynamic(&Device::default().settings().quantization.scheme) + }) +} + +fn make_qmatrix(rows: usize, cols: usize) -> Option> { + common::try_setup(|| { + make_matrix(rows, cols).quantize_dynamic(&Device::default().settings().quantization.scheme) + }) +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "quantize")] + mod quantize { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let scheme = Device::default().settings().quantization.scheme; + bencher.bench_synced(|| make_tensor(SMALL).quantize_dynamic(&scheme)); + } + + #[divan::bench] + fn medium(bencher: Bencher) { + let scheme = Device::default().settings().quantization.scheme; + bencher.bench_synced(|| make_tensor(MEDIUM).quantize_dynamic(&scheme)); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let scheme = Device::default().settings().quantization.scheme; + bencher.bench_synced(|| make_tensor(LARGE).quantize_dynamic(&scheme)); + } + } + + #[divan::bench_group(name = "dequantize")] + mod dequantize { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let Some(qt) = make_qtensor(SMALL) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| qt.clone().dequantize()); + } + + #[divan::bench] + fn medium(bencher: Bencher) { + let Some(qt) = make_qtensor(MEDIUM) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| qt.clone().dequantize()); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let Some(qt) = make_qtensor(LARGE) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| qt.clone().dequantize()); + } + } + + #[divan::bench_group(name = "q_add")] + mod q_add { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let Some(a) = make_qtensor(SMALL) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_qtensor(SMALL) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone() + b.clone()); + } + + #[divan::bench] + fn medium(bencher: Bencher) { + let Some(a) = make_qtensor(MEDIUM) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_qtensor(MEDIUM) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone() + b.clone()); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let Some(a) = make_qtensor(LARGE) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_qtensor(LARGE) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone() + b.clone()); + } + } + + #[divan::bench_group(name = "q_matmul")] + mod q_matmul { + use super::*; + + #[divan::bench] + fn mat_64x64(bencher: Bencher) { + let Some(a) = make_qmatrix(64, 64) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_qmatrix(64, 64) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + #[divan::bench] + fn mat_256x256(bencher: Bencher) { + let Some(a) = make_qmatrix(256, 256) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_qmatrix(256, 256) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + + #[divan::bench] + fn mat_512x512(bencher: Bencher) { + let Some(a) = make_qmatrix(512, 512) else { + bencher.bench(|| ()); + return; + }; + let Some(b) = make_qmatrix(512, 512) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| a.clone().matmul(b.clone())); + } + } + + #[divan::bench_group(name = "q_sum")] + mod q_sum { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let Some(qt) = make_qtensor(SMALL) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| qt.clone().sum()); + } + + #[divan::bench] + fn medium(bencher: Bencher) { + let Some(qt) = make_qtensor(MEDIUM) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| qt.clone().sum()); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let Some(qt) = make_qtensor(LARGE) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| qt.clone().sum()); + } + } + + #[divan::bench_group(name = "q_permute")] + mod q_permute { + use super::*; + + #[divan::bench] + fn medium_256x256(bencher: Bencher) { + let Some(qt) = make_qmatrix(256, 256) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| qt.clone().permute([1, 0])); + } + + #[divan::bench] + fn large_1024x1024(bencher: Bencher) { + let Some(qt) = make_qmatrix(1024, 1024) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| qt.clone().permute([1, 0])); + } + } + + #[divan::bench_group(name = "q_argmax")] + mod q_argmax { + use super::*; + + #[divan::bench] + fn medium_256x256(bencher: Bencher) { + let Some(qt) = make_qmatrix(256, 256) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| qt.clone().argmax(1)); + } + + #[divan::bench] + fn large_1024x1024(bencher: Bencher) { + let Some(qt) = make_qmatrix(1024, 1024) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| qt.clone().argmax(1)); + } + } + + #[divan::bench_group(name = "q_argmin")] + mod q_argmin { + use super::*; + + #[divan::bench] + fn medium_256x256(bencher: Bencher) { + let Some(qt) = make_qmatrix(256, 256) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| qt.clone().argmin(1)); + } + + #[divan::bench] + fn large_1024x1024(bencher: Bencher) { + let Some(qt) = make_qmatrix(1024, 1024) else { + bencher.bench(|| ()); + return; + }; + bencher.bench_synced(|| qt.clone().argmin(1)); + } + } + + #[divan::bench_group(name = "q_gather")] + mod q_gather { + use super::*; + use burn_tensor::Tensor; + + fn make_indices(rows: usize, cols: usize) -> Tensor<2, burn_tensor::Int> { + let data: Vec = (0..rows * cols).map(|i| (i % cols) as i32).collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) + } + + #[divan::bench] + fn medium_256x256(bencher: Bencher) { + let Some(qt) = make_qmatrix(256, 256) else { + bencher.bench(|| ()); + return; + }; + let indices = make_indices(256, 256); + bencher.bench_synced(|| qt.clone().gather(1, indices.clone())); + } + + #[divan::bench] + fn large_1024x1024(bencher: Bencher) { + let Some(qt) = make_qmatrix(1024, 1024) else { + bencher.bench(|| ()); + return; + }; + let indices = make_indices(1024, 1024); + bencher.bench_synced(|| qt.clone().gather(1, indices.clone())); + } + } + } + }; +} + +bench_backend!(backend, "backend"); diff --git a/crates/burn-backend-tests/benches/reduce_ops.rs b/crates/burn-backend-tests/benches/reduce_ops.rs new file mode 100644 index 0000000..bdb0070 --- /dev/null +++ b/crates/burn-backend-tests/benches/reduce_ops.rs @@ -0,0 +1,279 @@ +//! Benchmarks for reduction operations. +//! +//! Run with: +//! ```bash +//! cargo bench --bench reduce_ops --features simd +//! ``` + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{FloatDType, Tensor, TensorData}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Reduction ops Benchmarks"); + println!(); + divan::main(); + common::report_failures(); +} + +fn make_tensor_1d(size: usize) -> Tensor<1> { + let data: Vec = (0..size).map(|i| (i % 1000) as f32 / 1000.0).collect(); + Tensor::from_data(TensorData::new(data, [size]), &Default::default()) +} + +fn make_tensor_2d(rows: usize, cols: usize) -> Tensor<2> { + let data: Vec = (0..rows * cols) + .map(|i| (i % 1000) as f32 / 1000.0) + .collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) +} + +fn make_tensor_3d(d0: usize, d1: usize, d2: usize) -> Tensor<3> { + let size = d0 * d1 * d2; + let data: Vec = (0..size).map(|i| (i % 1000) as f32 / 1000.0).collect(); + Tensor::from_data(TensorData::new(data, [d0, d1, d2]), &Default::default()) +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + // Sum all elements + #[divan::bench_group(name = "sum")] + mod sum { + use super::*; + + #[divan::bench] + fn s_1k(bencher: Bencher) { + let t = make_tensor_1d(1024); + bencher.bench_synced(|| t.clone().sum()); + } + + #[divan::bench] + fn s_64k(bencher: Bencher) { + let t = make_tensor_1d(64 * 1024); + bencher.bench_synced(|| t.clone().sum()); + } + + #[divan::bench] + fn s_1m(bencher: Bencher) { + let t = make_tensor_1d(1024 * 1024); + bencher.bench_synced(|| t.clone().sum()); + } + } + + // Sum along dimension + #[divan::bench_group(name = "sum_dim")] + mod sum_dim { + use super::*; + + #[divan::bench] + fn s_256x256_dim0(bencher: Bencher) { + let t = make_tensor_2d(256, 256); + bencher.bench_synced(|| t.clone().sum_dim(0)); + } + + #[divan::bench] + fn s_256x256_dim1(bencher: Bencher) { + let t = make_tensor_2d(256, 256); + bencher.bench_synced(|| t.clone().sum_dim(1)); + } + + #[divan::bench] + fn s_1024x1024_dim0(bencher: Bencher) { + let t = make_tensor_2d(1024, 1024); + bencher.bench_synced(|| t.clone().sum_dim(0)); + } + + #[divan::bench] + fn s_1024x1024_dim1(bencher: Bencher) { + let t = make_tensor_2d(1024, 1024); + bencher.bench_synced(|| t.clone().sum_dim(1)); + } + } + + // Mean along dimension + #[divan::bench_group(name = "mean_dim")] + mod mean_dim { + use super::*; + + #[divan::bench] + fn s_256x256_dim1(bencher: Bencher) { + let t = make_tensor_2d(256, 256); + bencher.bench_synced(|| t.clone().mean_dim(1)); + } + + #[divan::bench] + fn s_1024x1024_dim1(bencher: Bencher) { + let t = make_tensor_2d(1024, 1024); + bencher.bench_synced(|| t.clone().mean_dim(1)); + } + } + + // Argmax + #[divan::bench_group(name = "argmax")] + mod argmax { + use super::*; + + #[divan::bench] + fn s_1k(bencher: Bencher) { + let t = make_tensor_1d(1024); + bencher.bench_synced(|| t.clone().argmax(0)); + } + + #[divan::bench] + fn s_256x256_dim1(bencher: Bencher) { + let t = make_tensor_2d(256, 256); + bencher.bench_synced(|| t.clone().argmax(1)); + } + + #[divan::bench] + fn s_1024x1024_dim1(bencher: Bencher) { + let t = make_tensor_2d(1024, 1024); + bencher.bench_synced(|| t.clone().argmax(1)); + } + } + + // Sum on transposed (non-contiguous) + #[divan::bench_group(name = "sum_transposed")] + mod sum_transposed { + use super::*; + + #[divan::bench] + fn s_256x256(bencher: Bencher) { + let t = make_tensor_2d(256, 256).transpose(); + bencher.bench_synced(|| t.clone().sum()); + } + + #[divan::bench] + fn s_1024x1024(bencher: Bencher) { + let t = make_tensor_2d(1024, 1024).transpose(); + bencher.bench_synced(|| t.clone().sum()); + } + } + + // Sum_dim on transposed tensor (tests dim_stride=1 optimization) + #[divan::bench_group(name = "sum_dim_transposed")] + mod sum_dim_transposed { + use super::*; + + #[divan::bench] + fn s_256x256_dim0(bencher: Bencher) { + let t = make_tensor_2d(256, 256).transpose(); + bencher.bench_synced(|| t.clone().sum_dim(0)); + } + + #[divan::bench] + fn s_1024x1024_dim0(bencher: Bencher) { + let t = make_tensor_2d(1024, 1024).transpose(); + bencher.bench_synced(|| t.clone().sum_dim(0)); + } + } + + // 3D reductions (batch-like) + #[divan::bench_group(name = "sum_3d")] + mod sum_3d { + use super::*; + + #[divan::bench] + fn s_batch32_256x256_dim2(bencher: Bencher) { + let t = make_tensor_3d(32, 256, 256); + bencher.bench_synced(|| t.clone().sum_dim(2)); + } + + #[divan::bench] + fn s_batch32_256x256_dim1(bencher: Bencher) { + let t = make_tensor_3d(32, 256, 256); + bencher.bench_synced(|| t.clone().sum_dim(1)); + } + } + } + }; +} + +bench_backend!(backend, "backend"); + +// ============================================================================ +// f16 mean benches +// +// These exercise the half-precision mean / mean_dim fast path. Each tensor +// is cast to f16 once during setup so the bench measures only the reduction, +// not the cast. Sizes mirror the f32 sum_dim / mean_dim groups so before/after +// numbers can be compared apples-to-apples. +// ============================================================================ + +#[divan::bench_group(name = "f16")] +mod f16 { + use super::*; + + fn make_f16_2d(rows: usize, cols: usize) -> Tensor<2> { + make_tensor_2d(rows, cols).cast(FloatDType::F16) + } + + fn make_f16_3d(d0: usize, d1: usize, d2: usize) -> Tensor<3> { + make_tensor_3d(d0, d1, d2).cast(FloatDType::F16) + } + + #[divan::bench_group(name = "mean")] + mod mean { + use super::*; + + #[divan::bench] + fn s_64k(bencher: Bencher) { + let t = make_f16_2d(256, 256); + bencher.bench_synced(|| t.clone().mean()); + } + + #[divan::bench] + fn s_1m(bencher: Bencher) { + let t = make_f16_2d(1024, 1024); + bencher.bench_synced(|| t.clone().mean()); + } + } + + #[divan::bench_group(name = "mean_dim")] + mod mean_dim { + use super::*; + + // Last-dim path (sum_rows_f32) + #[divan::bench] + fn s_256x256_dim1(bencher: Bencher) { + let t = make_f16_2d(256, 256); + bencher.bench_synced(|| t.clone().mean_dim(1)); + } + + #[divan::bench] + fn s_1024x1024_dim1(bencher: Bencher) { + let t = make_f16_2d(1024, 1024); + bencher.bench_synced(|| t.clone().mean_dim(1)); + } + + // First-dim path (scatter_add_f32) + #[divan::bench] + fn s_256x256_dim0(bencher: Bencher) { + let t = make_f16_2d(256, 256); + bencher.bench_synced(|| t.clone().mean_dim(0)); + } + + #[divan::bench] + fn s_1024x1024_dim0(bencher: Bencher) { + let t = make_f16_2d(1024, 1024); + bencher.bench_synced(|| t.clone().mean_dim(0)); + } + + // Middle-dim path (scatter_add_batched) + #[divan::bench] + fn s_batch32_256x256_dim1(bencher: Bencher) { + let t = make_f16_3d(32, 256, 256); + bencher.bench_synced(|| t.clone().mean_dim(1)); + } + } +} diff --git a/crates/burn-backend-tests/benches/slice_ops.rs b/crates/burn-backend-tests/benches/slice_ops.rs new file mode 100644 index 0000000..63aa4ea --- /dev/null +++ b/crates/burn-backend-tests/benches/slice_ops.rs @@ -0,0 +1,240 @@ +//! Benchmarks for slice operations. +//! +//! Run with: +//! ```bash +//! cargo bench --bench slice_ops --features simd +//! ``` +//! +//! Memory allocation tracking is enabled via divan's AllocProfiler. + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{Tensor, TensorData, s}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Slice Operations Benchmarks"); + println!("Memory allocation tracking enabled"); + println!(); + divan::main(); + common::report_failures(); +} + +fn make_tensor_1d(size: usize) -> Tensor<1> { + let data: Vec = (0..size).map(|i| (i % 1000) as f32 / 1000.0).collect(); + Tensor::from_data(TensorData::new(data, [size]), &Default::default()) +} + +fn make_tensor_2d(rows: usize, cols: usize) -> Tensor<2> { + let data: Vec = (0..rows * cols) + .map(|i| (i % 1000) as f32 / 1000.0) + .collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) +} + +fn make_tensor_3d(d0: usize, d1: usize, d2: usize) -> Tensor<3> { + let data: Vec = (0..d0 * d1 * d2) + .map(|i| (i % 1000) as f32 / 1000.0) + .collect(); + Tensor::from_data(TensorData::new(data, [d0, d1, d2]), &Default::default()) +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + + + // Basic slice (contiguous, step=1) + #[divan::bench_group(name = "basic_slice")] + mod basic_slice { + use super::*; + + #[divan::bench] + fn slice_1d_1k(bencher: Bencher) { + let t = make_tensor_1d(1024); + bencher.bench_synced(|| t.clone().slice([256..768])); + } + + #[divan::bench] + fn slice_1d_1m(bencher: Bencher) { + let t = make_tensor_1d(1024 * 1024); + bencher.bench_synced(|| t.clone().slice([256 * 1024..768 * 1024])); + } + + #[divan::bench] + fn slice_2d_256x256(bencher: Bencher) { + let t = make_tensor_2d(256, 256); + bencher.bench_synced(|| t.clone().slice([64..192, 64..192])); + } + + #[divan::bench] + fn slice_2d_1024x1024(bencher: Bencher) { + let t = make_tensor_2d(1024, 1024); + bencher.bench_synced(|| t.clone().slice([256..768, 256..768])); + } + + #[divan::bench] + fn slice_3d_64x64x64(bencher: Bencher) { + let t = make_tensor_3d(64, 64, 64); + bencher.bench_synced(|| t.clone().slice([16..48, 16..48, 16..48])); + } + } + + // Slice with step > 1 + #[divan::bench_group(name = "slice_with_step")] + mod slice_with_step { + use super::*; + + #[divan::bench] + fn step2_1d_1k(bencher: Bencher) { + let t = make_tensor_1d(1024); + bencher.bench_synced(|| t.clone().slice(s![0..1024;2])); + } + + #[divan::bench] + fn step2_1d_1m(bencher: Bencher) { + let t = make_tensor_1d(1024 * 1024); + bencher.bench_synced(|| t.clone().slice(s![0..1024 * 1024;2])); + } + + #[divan::bench] + fn step4_2d_256x256(bencher: Bencher) { + let t = make_tensor_2d(256, 256); + bencher.bench_synced(|| t.clone().slice(s![0..256;4, 0..256;4])); + } + + #[divan::bench] + fn step2_2d_1024x1024(bencher: Bencher) { + let t = make_tensor_2d(1024, 1024); + bencher.bench_synced(|| t.clone().slice(s![0..1024;2, 0..1024;2])); + } + } + + // Slice on transposed tensor (non-contiguous) + #[divan::bench_group(name = "slice_transposed")] + mod slice_transposed { + use super::*; + + #[divan::bench] + fn transposed_256x256(bencher: Bencher) { + let t = make_tensor_2d(256, 256).transpose(); + bencher.bench_synced(|| t.clone().slice([64..192, 64..192])); + } + + #[divan::bench] + fn transposed_1024x1024(bencher: Bencher) { + let t = make_tensor_2d(1024, 1024).transpose(); + bencher.bench_synced(|| t.clone().slice([256..768, 256..768])); + } + } + + // Slice assign + #[divan::bench_group(name = "slice_assign")] + mod slice_assign { + use super::*; + + #[divan::bench] + fn assign_1d_1k(bencher: Bencher) { + let t = make_tensor_1d(1024); + let v = make_tensor_1d(512); + bencher.bench_synced(|| t.clone().slice_assign([256..768], v.clone())); + } + + #[divan::bench] + fn assign_2d_256x256(bencher: Bencher) { + let t = make_tensor_2d(256, 256); + let v = make_tensor_2d(128, 128); + bencher.bench_synced(|| t.clone().slice_assign([64..192, 64..192], v.clone())); + } + + #[divan::bench] + fn assign_2d_1024x1024(bencher: Bencher) { + let t = make_tensor_2d(1024, 1024); + let v = make_tensor_2d(512, 512); + bencher.bench_synced(|| t.clone().slice_assign([256..768, 256..768], v.clone())); + } + } + + // Slice fill (scalar fill over a slice region) + // + // This is the pattern the user hit in issue #64 item 3: + // `tensor.slice_fill(slice, scalar)`. burn-tensor's default impl + // allocates a 1-element tensor, calls expand() to broadcast it + // over the slice shape, and then hands that view to + // slice_assign. A naive backend materializes the broadcast (one + // full alloc + fill) before doing the actual strided write, so + // the slice fill ends up doing ~3x the memory traffic it should. + #[divan::bench_group(name = "slice_fill")] + mod slice_fill { + use super::*; + + #[divan::bench] + fn fill_1d_512_of_1k(bencher: Bencher) { + let t = make_tensor_1d(1024); + bencher.bench_synced(|| t.clone().slice_fill([256..768], 1.0f32)); + } + + #[divan::bench] + fn fill_2d_128x128_of_256x256(bencher: Bencher) { + let t = make_tensor_2d(256, 256); + bencher.bench_synced(|| t.clone().slice_fill([64..192, 64..192], 1.0f32)); + } + + #[divan::bench] + fn fill_2d_512x512_of_1024x1024(bencher: Bencher) { + let t = make_tensor_2d(1024, 1024); + bencher.bench_synced(|| t.clone().slice_fill([256..768, 256..768], 1.0f32)); + } + + // Segmentation-model-sized: 488x448 feature map, half the + // spatial extent. + #[divan::bench] + fn fill_2d_244x224_of_488x448(bencher: Bencher) { + let t = make_tensor_2d(488, 448); + bencher.bench_synced(|| t.clone().slice_fill([122..366, 112..336], 0.0f32)); + } + + // 3D feature map (B=1,C=64,H,W) half-spatial fill. + #[divan::bench] + fn fill_3d_64x64x64_of_64x128x128(bencher: Bencher) { + let t = make_tensor_3d(64, 128, 128); + bencher.bench_synced(|| t.clone().slice_fill([0..64, 32..96, 32..96], 0.0f32)); + } + } + + // Narrow (single dimension slice) + #[divan::bench_group(name = "narrow")] + mod narrow { + use super::*; + + #[divan::bench] + fn narrow_dim0_256x256(bencher: Bencher) { + let t = make_tensor_2d(256, 256); + bencher.bench_synced(|| t.clone().narrow(0, 64, 128)); + } + + #[divan::bench] + fn narrow_dim1_256x256(bencher: Bencher) { + let t = make_tensor_2d(256, 256); + bencher.bench_synced(|| t.clone().narrow(1, 64, 128)); + } + + #[divan::bench] + fn narrow_dim0_1024x1024(bencher: Bencher) { + let t = make_tensor_2d(1024, 1024); + bencher.bench_synced(|| t.clone().narrow(0, 256, 512)); + } + } + } + }; +} + +bench_backend!(backend, "backend"); diff --git a/crates/burn-backend-tests/benches/softmax_ops.rs b/crates/burn-backend-tests/benches/softmax_ops.rs new file mode 100644 index 0000000..ff266fc --- /dev/null +++ b/crates/burn-backend-tests/benches/softmax_ops.rs @@ -0,0 +1,98 @@ +//! Benchmarks comparing `B::softmax` against a manual decomposition baseline. +//! Covers last-axis and non-last-axis cases. +//! +//! Run with: +//! ```bash +//! cargo bench --bench softmax_ops --features simd,rayon +//! ``` + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{Element, Tensor, TensorData}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Softmax Benchmarks: fused (via B::softmax) vs decomposed"); + println!(); + divan::main(); + common::report_failures(); +} + +fn make_tensor_3d>(d0: usize, d1: usize, d2: usize) -> Tensor<3> { + let data: Vec = (0..d0 * d1 * d2) + .map(|i| E::from((((i % 997) as f32) / 997.0) - 0.5)) + .collect(); + Tensor::from_data(TensorData::new(data, [d0, d1, d2]), &Default::default()) +} + +/// `softmax(x, dim) = exp(x - max) / sum(exp(x - max))`, matching the +/// `ActivationOps::softmax` default in burn-backend. +fn decomposed_softmax(x: Tensor, dim: usize) -> Tensor { + let max = x.clone().detach().max_dim(dim); + let shifted = x - max; + let exp = shifted.exp(); + let sum = exp.clone().sum_dim(dim); + exp / sum +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "last_axis_f32")] + mod last_axis_f32 { + use super::*; + + #[divan::bench] + fn trait_hook_bert_seq512_h12(bencher: Bencher) { + let x = make_tensor_3d::(12, 512, 512); + bencher.bench_synced(|| burn_tensor::activation::softmax(x.clone(), 2)); + } + + #[divan::bench] + fn decomposed_bert_seq512_h12(bencher: Bencher) { + let x = make_tensor_3d::(12, 512, 512); + bencher.bench_synced(|| decomposed_softmax(x.clone(), 2)); + } + + #[divan::bench] + fn trait_hook_wide_d2048(bencher: Bencher) { + let x = make_tensor_3d::(2, 256, 2048); + bencher.bench_synced(|| burn_tensor::activation::softmax(x.clone(), 2)); + } + + #[divan::bench] + fn decomposed_wide_d2048(bencher: Bencher) { + let x = make_tensor_3d::(2, 256, 2048); + bencher.bench_synced(|| decomposed_softmax(x.clone(), 2)); + } + } + + #[divan::bench_group(name = "non_last_axis_f32")] + mod non_last_axis_f32 { + use super::*; + + #[divan::bench] + fn trait_hook_dim1_of_3(bencher: Bencher) { + let x = make_tensor_3d::(4, 1024, 512); + bencher.bench_synced(|| burn_tensor::activation::softmax(x.clone(), 1)); + } + + #[divan::bench] + fn decomposed_dim1_of_3(bencher: Bencher) { + let x = make_tensor_3d::(4, 1024, 512); + bencher.bench_synced(|| decomposed_softmax(x.clone(), 1)); + } + } + } + }; +} + +bench_backend!(backend, "backend"); diff --git a/crates/burn-backend-tests/benches/unary_ops.rs b/crates/burn-backend-tests/benches/unary_ops.rs new file mode 100644 index 0000000..7aa58a4 --- /dev/null +++ b/crates/burn-backend-tests/benches/unary_ops.rs @@ -0,0 +1,225 @@ +//! Benchmarks for unary operations. +//! +//! Run with: +//! ```bash +//! cargo bench --bench unary_ops --features simd +//! ``` + +#[path = "common/mod.rs"] +mod common; +use common::BencherExt; + +use burn_tensor::{Tensor, TensorData}; +use divan::{AllocProfiler, Bencher}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +fn main() { + println!("Unary ops Benchmarks"); + println!(); + divan::main(); + common::report_failures(); +} + +const SMALL: usize = 64 * 64; // 4K elements +const MEDIUM: usize = 256 * 256; // 64K elements +const LARGE: usize = 1024 * 1024; // 1M elements + +fn make_tensor(size: usize) -> Tensor<1> { + // Use values in range [0.1, 1.1] to avoid domain issues for log/sqrt + let data: Vec = (0..size) + .map(|i| 0.1 + (i % 1000) as f32 / 1000.0) + .collect(); + Tensor::from_data(TensorData::new(data, [size]), &Default::default()) +} + +fn make_tensor_2d(rows: usize, cols: usize) -> Tensor<2> { + let data: Vec = (0..rows * cols) + .map(|i| 0.1 + (i % 1000) as f32 / 1000.0) + .collect(); + Tensor::from_data(TensorData::new(data, [rows, cols]), &Default::default()) +} + +macro_rules! bench_backend { + ($mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name)] + mod $mod_name { + use super::*; + + #[divan::bench_group(name = "exp")] + mod exp { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let a = make_tensor(SMALL); + bencher.bench_synced(|| a.clone().exp()); + } + + #[divan::bench] + fn medium(bencher: Bencher) { + let a = make_tensor(MEDIUM); + bencher.bench_synced(|| a.clone().exp()); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + bencher.bench_synced(|| a.clone().exp()); + } + } + + #[divan::bench_group(name = "log")] + mod log { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let a = make_tensor(SMALL); + bencher.bench_synced(|| a.clone().log()); + } + + #[divan::bench] + fn medium(bencher: Bencher) { + let a = make_tensor(MEDIUM); + bencher.bench_synced(|| a.clone().log()); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + bencher.bench_synced(|| a.clone().log()); + } + } + + #[divan::bench_group(name = "sqrt")] + mod sqrt { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let a = make_tensor(SMALL); + bencher.bench_synced(|| a.clone().sqrt()); + } + + #[divan::bench] + fn medium(bencher: Bencher) { + let a = make_tensor(MEDIUM); + bencher.bench_synced(|| a.clone().sqrt()); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + bencher.bench_synced(|| a.clone().sqrt()); + } + } + + #[divan::bench_group(name = "sin")] + mod sin { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let a = make_tensor(SMALL); + bencher.bench_synced(|| a.clone().sin()); + } + + #[divan::bench] + fn medium(bencher: Bencher) { + let a = make_tensor(MEDIUM); + bencher.bench_synced(|| a.clone().sin()); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + bencher.bench_synced(|| a.clone().sin()); + } + } + + #[divan::bench_group(name = "cos")] + mod cos { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let a = make_tensor(SMALL); + bencher.bench_synced(|| a.clone().cos()); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + bencher.bench_synced(|| a.clone().cos()); + } + } + + #[divan::bench_group(name = "tanh")] + mod tanh { + use super::*; + + #[divan::bench] + fn small(bencher: Bencher) { + let a = make_tensor(SMALL); + bencher.bench_synced(|| a.clone().tanh()); + } + + #[divan::bench] + fn medium(bencher: Bencher) { + let a = make_tensor(MEDIUM); + bencher.bench_synced(|| a.clone().tanh()); + } + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + bencher.bench_synced(|| a.clone().tanh()); + } + } + + #[divan::bench_group(name = "abs")] + mod abs { + use super::*; + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + bencher.bench_synced(|| a.clone().abs()); + } + } + + #[divan::bench_group(name = "recip")] + mod recip { + use super::*; + + #[divan::bench] + fn large(bencher: Bencher) { + let a = make_tensor(LARGE); + bencher.bench_synced(|| a.clone().recip()); + } + } + + // Non-contiguous (transposed) + #[divan::bench_group(name = "exp_transposed")] + mod exp_transposed { + use super::*; + + #[divan::bench] + fn medium_256x256(bencher: Bencher) { + let a = make_tensor_2d(256, 256).transpose(); + bencher.bench_synced(|| a.clone().exp()); + } + + #[divan::bench] + fn large_1024x1024(bencher: Bencher) { + let a = make_tensor_2d(1024, 1024).transpose(); + bencher.bench_synced(|| a.clone().exp()); + } + } + } + }; +} + +bench_backend!(backend, "backend"); diff --git a/crates/burn-backend-tests/cubecl.toml b/crates/burn-backend-tests/cubecl.toml new file mode 100644 index 0000000..f1f12cc --- /dev/null +++ b/crates/burn-backend-tests/cubecl.toml @@ -0,0 +1,15 @@ +[profiling] +logger = { file = "target/profiling.log", level = "disabled" } + +[autotune] +logger = { file = "target/autotune.log", level = "disabled" } + +[compilation] +check_mode = "validate" +logger = { file = "target/compilation.log", level = "disabled" } + +[memory] +logger = { file = "target/memory.log", level = "disabled" } + +[streaming] +max_streams = 4 diff --git a/crates/burn-backend-tests/src/lib.rs b/crates/burn-backend-tests/src/lib.rs new file mode 100644 index 0000000..997f018 --- /dev/null +++ b/crates/burn-backend-tests/src/lib.rs @@ -0,0 +1,157 @@ +use proc_macro::TokenStream; +use quote::{format_ident, quote}; + +use syn::parse::{Parse, ParseStream}; +use syn::punctuated::Punctuated; +use syn::token::Comma; +use syn::{Attribute, Expr, ItemFn, Lit, Meta, MetaNameValue, parse_macro_input}; + +// Define a structure to parse the attribute arguments +struct AttributeArgs { + args: Punctuated, +} + +impl Parse for AttributeArgs { + fn parse(input: ParseStream) -> syn::Result { + Ok(AttributeArgs { + args: Punctuated::parse_terminated(input)?, + }) + } +} + +#[allow(clippy::test_attr_in_doctest)] +/// **This is only meaningful when the `reason` is specific and clear.** +/// +/// A proc macro attribute that adds panic handling to test functions. +/// +/// # Usage +/// ```rust, ignore +/// #[might_panic(reason = "expected panic message prefix")] +/// #[test] +/// fn test_that_might_panic() { +/// // test code that might panic (with acceptable reason) +/// } +/// ``` +/// +/// # Behavior +/// - If the test does not panic, it passes. +/// - If the test panics with a message starting with the expected prefix, the failure is ignored. +/// - If the test panics with a different message, the test fails. +/// +/// # Note +/// This proc macro uses [`std::panic::catch_unwind`]. As such, it does not work in a no-std environment. +/// Make sure it is feature gated when an `"std"` feature is available. +#[proc_macro_attribute] +pub fn might_panic(args: TokenStream, input: TokenStream) -> TokenStream { + // Parse the attribute arguments + let args = parse_macro_input!(args as AttributeArgs); + let input_fn = parse_macro_input!(input as ItemFn); + + // Extract the expected panic reason + let mut expected_reason = None; + for arg in args.args.iter() { + if let Meta::NameValue(MetaNameValue { path, value, .. }) = arg + && path.is_ident("reason") + && let Expr::Lit(lit) = value + && let Lit::Str(ref lit_str) = lit.lit + { + expected_reason = Some(lit_str.value()); + } + } + + let expected_reason = match expected_reason { + Some(reason) => reason, + None => { + return syn::Error::new( + proc_macro2::Span::call_site(), + "The #[might_panic] attribute requires a 'reason' parameter", + ) + .to_compile_error() + .into(); + } + }; + + let fn_name = &input_fn.sig.ident; + let fn_vis = &input_fn.vis; + let fn_generics = &input_fn.sig.generics; + let fn_block = &input_fn.block; + let fn_attrs = input_fn + .attrs + .iter() + .filter(|attr| !attr.path().is_ident("test")) + .collect::>(); + + // Create a wrapped test function + let wrapper_name = format_ident!("{}_might_panic", fn_name); + + quote! { + #(#fn_attrs)* + #fn_vis fn #fn_name #fn_generics() { #fn_block } + + #[test] + #fn_vis fn #wrapper_name #fn_generics() { + use std::panic::{self, AssertUnwindSafe}; + use std::sync::{Arc, Mutex, OnceLock}; + + let get_msg = |p: &(dyn std::any::Any + Send)| -> String { + p.downcast_ref::().cloned() + .or_else(|| p.downcast_ref::<&str>().map(|s| s.to_string())) + .unwrap_or_else(|| "Unknown panic".to_string()) + }; + + // An append-only list of all panic messages across the entire process. + // This is required because cubecl's `CallError` hides the original panic message + // occurring in the device threads. + // + // A global log also prevents parallel tests from overwriting each other's panic hooks. + static PANIC_LOG: OnceLock>> = OnceLock::new(); + let log = PANIC_LOG.get_or_init(|| Mutex::new(Vec::new())); + + static HOOK: OnceLock<()> = OnceLock::new(); + HOOK.get_or_init(|| { + let prev = panic::take_hook(); + panic::set_hook(Box::new(move |info| { + if let Ok(mut v) = log.lock() { + v.push(get_msg(info.payload())); + } + prev(info); + })); + }); + + // We only care about panics that occur during this test's execution window, so + // we start at the number of panics logged before this test starts. + let start_idx = log.lock().unwrap().len(); + let result = panic::catch_unwind(AssertUnwindSafe(|| #fn_name())); + + if let Err(e) = result { + let main_msg = get_msg(&*e); + let panic_logs = log.lock().unwrap(); + let window = &panic_logs[start_idx..]; + + let matched = window.iter().chain(std::iter::once(&main_msg)) + .any(|m| m.contains(#expected_reason)); + + let all = window.iter().chain(std::iter::once(&main_msg)) + .map(|m| format!("- {m}")).collect::>().join("\n"); + + if matched { + eprintln!( + "\n[SKIPPED - might_panic] Test '{}'\nReason: '{}'\nPanics:\n{}\n", + stringify!(#fn_name), + #expected_reason, + all + ); + return; + } else { + panic!( + "\nTest '{}' failed.\nExpected: '{}'\nFound:\n{}\n", + stringify!(#fn_name), + #expected_reason, + all + ); + } + } + } + } + .into() +} diff --git a/crates/burn-backend-tests/tests/autodiff.rs b/crates/burn-backend-tests/tests/autodiff.rs new file mode 100644 index 0000000..e4de92d --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff.rs @@ -0,0 +1,18 @@ +//! Burn autodiff tests. + +#![recursion_limit = "256"] +#![allow(clippy::single_range_in_vec_init, clippy::duplicate_mod)] + +extern crate alloc; + +pub type FloatElem = f32; +#[allow(unused)] +pub type IntElem = i32; + +#[path = "common/backend.rs"] +mod backend; +pub use backend::*; + +#[allow(clippy::module_inception)] +#[path = "common/autodiff.rs"] +mod autodiff; diff --git a/crates/burn-backend-tests/tests/autodiff/abs.rs b/crates/burn-backend-tests/tests/autodiff/abs.rs new file mode 100644 index 0000000..a1160d5 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/abs.rs @@ -0,0 +1,58 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn should_diff_abs() { + let data_1 = TensorData::from([[0.0, -1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[6.0, 7.0], [9.0, -10.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().abs()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[71.0, 107.0], [71.0, 107.0]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[84.0, 42.0], [90.0, 54.0]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_abs_no_nans() { + let data_1 = TensorData::from([[6.0, 7.0], [9.0, -10.0]]); + let data_2 = TensorData::from([[0.0, -1.0], [3.0, 4.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().abs()); + let grads = tensor_3.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[1.0, 7.0], [1.0, 7.0]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[0.0, -15.0], [-3.0, -3.0]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let contains_nan = grad_2.contains_nan(); + assert!(!contains_nan.into_scalar::()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/adaptive_avgpool1d.rs b/crates/burn-backend-tests/tests/autodiff/adaptive_avgpool1d.rs new file mode 100644 index 0000000..88b4f10 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/adaptive_avgpool1d.rs @@ -0,0 +1,50 @@ +use super::*; +use burn_tensor::module::adaptive_avg_pool1d; +use burn_tensor::{Shape, Tolerance}; + +#[test] +fn test_avg_pool1d_simple() { + let test = AdaptiveAvgPool1dTestCase { + batch_size: 1, + channels: 2, + length: 5, + output_size: 3, + }; + + test.assert_output(TestTensor::from_data( + [[ + [0.5000, 0.83333, 0.33333, 0.83333, 0.5000], + [0.5000, 0.83333, 0.33333, 0.83333, 0.5000], + ]], + &AutodiffDevice::new(), + )); +} + +struct AdaptiveAvgPool1dTestCase { + batch_size: usize, + channels: usize, + length: usize, + output_size: usize, +} + +impl AdaptiveAvgPool1dTestCase { + fn assert_output(self, x_grad: TestTensor<3>) { + let shape_x = Shape::new([self.batch_size, self.channels, self.length]); + let device = AutodiffDevice::new(); + let x = TestTensor::from_data( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<3, _>(shape_x) + .into_data(), + &device, + ) + .require_grad(); + let output = adaptive_avg_pool1d(x.clone(), self.output_size); + let grads = output.backward(); + let x_grad_actual = x.grad(&grads).unwrap(); + + x_grad.to_data().assert_approx_eq::( + &x_grad_actual.into_data(), + Tolerance::default().set_half_precision_relative(1e-3), + ); + } +} diff --git a/crates/burn-backend-tests/tests/autodiff/adaptive_avgpool2d.rs b/crates/burn-backend-tests/tests/autodiff/adaptive_avgpool2d.rs new file mode 100644 index 0000000..2bd3219 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/adaptive_avgpool2d.rs @@ -0,0 +1,96 @@ +use super::*; +use burn_tensor::module::adaptive_avg_pool2d; +use burn_tensor::{Shape, Tolerance}; + +#[test] +fn test_avg_pool2d_simple() { + let test = AdaptiveAvgPool2dTestCase { + batch_size: 1, + channels: 2, + height: 5, + width: 3, + output_size_1: 3, + output_size_2: 2, + }; + + test.assert_output(TestTensor::from_data( + [[ + [ + [0.2500, 0.5000, 0.2500], + [0.41667, 0.83333, 0.41667], + [0.16667, 0.33333, 0.16667], + [0.41667, 0.83333, 0.41667], + [0.2500, 0.5000, 0.2500], + ], + [ + [0.2500, 0.5000, 0.2500], + [0.41667, 0.83333, 0.41667], + [0.16667, 0.33333, 0.16667], + [0.41667, 0.83333, 0.41667], + [0.2500, 0.5000, 0.2500], + ], + ]], + &AutodiffDevice::new(), + )); +} + +#[test] +fn test_avg_pool2d_output_1() { + let test = AdaptiveAvgPool2dTestCase { + batch_size: 1, + channels: 1, + height: 4, + width: 8, + output_size_1: 1, + output_size_2: 1, + }; + + test.assert_output(TestTensor::from_data( + [[[ + [ + 0.03125, 0.03125, 0.03125, 0.03125, 0.03125, 0.03125, 0.03125, 0.03125, + ], + [ + 0.03125, 0.03125, 0.03125, 0.03125, 0.03125, 0.03125, 0.03125, 0.03125, + ], + [ + 0.03125, 0.03125, 0.03125, 0.03125, 0.03125, 0.03125, 0.03125, 0.03125, + ], + [ + 0.03125, 0.03125, 0.03125, 0.03125, 0.03125, 0.03125, 0.03125, 0.03125, + ], + ]]], + &AutodiffDevice::new(), + )); +} + +struct AdaptiveAvgPool2dTestCase { + batch_size: usize, + channels: usize, + height: usize, + width: usize, + output_size_1: usize, + output_size_2: usize, +} + +impl AdaptiveAvgPool2dTestCase { + fn assert_output(self, x_grad: TestTensor<4>) { + let shape_x = Shape::new([self.batch_size, self.channels, self.height, self.width]); + let device = AutodiffDevice::new(); + let x = TestTensor::from_data( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<4, _>(shape_x) + .into_data(), + &device, + ) + .require_grad(); + let output = adaptive_avg_pool2d(x.clone(), [self.output_size_1, self.output_size_2]); + let grads = output.backward(); + let x_grad_actual = x.grad(&grads).unwrap(); + + x_grad.to_data().assert_approx_eq::( + &x_grad_actual.into_data(), + Tolerance::default().set_half_precision_relative(1e-3), + ); + } +} diff --git a/crates/burn-backend-tests/tests/autodiff/add.rs b/crates/burn-backend-tests/tests/autodiff/add.rs new file mode 100644 index 0000000..60d3af8 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/add.rs @@ -0,0 +1,74 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_diff_add() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<1>::from_data([2.0, 5.0], &device).require_grad(); + let tensor_2 = TestTensor::from_data([4.0, 1.0], &device).require_grad(); + + let tensor_3 = tensor_1.clone() + tensor_2.clone(); + let grads = tensor_3.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([1.0, 1.0]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([1.0, 1.0]), false); + tensor_3 + .to_data() + .assert_eq(&TensorData::from([6.0, 6.0]), false); +} + +#[test] +fn should_diff_add_scalar() { + let data = TensorData::from([2.0, 10.0]); + + let tensor = TestTensor::<1>::from_data(data, &AutodiffDevice::new()).require_grad(); + let tensor_out = tensor.clone().add_scalar(5.0); + let grads = tensor_out.backward(); + + let grad = tensor.grad(&grads).unwrap(); + + grad.to_data() + .assert_eq(&TensorData::from([1.0, 1.0]), false); + tensor_out + .into_data() + .assert_eq(&TensorData::from([7.0, 15.0]), false); +} + +#[test] +fn test_add_complex_1() { + let data_1 = TensorData::from([[1.0, 7.0], [13.0, -3.0]]); + let data_2 = TensorData::from([[4.0, 7.0], [2.0, 3.0]]); + let data_3 = TensorData::from([[2.0, 2.0], [2.0, 2.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + let tensor_3 = TestTensor::from_data(data_3, &device).require_grad(); + + let tensor_4 = tensor_1.clone().add(tensor_2.clone()); + let tensor_5 = tensor_4 + .add(tensor_3) + .add_scalar(5.0) + .add(tensor_1.clone()) + .add(tensor_2.clone()); + let tensor_6 = tensor_1.clone().add(tensor_5); + + let grads = tensor_6.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[3.0, 3.0], [3.0, 3.0]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([[2.0, 2.0], [2.0, 2.0]]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/aggregation.rs b/crates/burn-backend-tests/tests/autodiff/aggregation.rs new file mode 100644 index 0000000..5ad1e65 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/aggregation.rs @@ -0,0 +1,261 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn should_diff_mean() { + let data_1 = TensorData::from([[1.0, 7.0], [-2.0, -3.0]]); + let data_2 = TensorData::from([[4.0, -7.0], [2.0, 3.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_1.clone().mul(tensor_3.mean().unsqueeze()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[3.5, 9.5], [3.5, 9.5]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[-0.75, -0.75], [3.0, 3.0]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_sum_1() { + let data_1 = TensorData::from([[1.0, 7.0], [-2.0, -3.0]]); + let data_2 = TensorData::from([[4.0, -7.0], [2.0, 3.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_1.clone().mul(tensor_3.sum().unsqueeze()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[14.0, 38.0], [14.0, 38.0]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[-3.0, -3.0], [12.0, 12.0]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_sum_2() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[6.0, 7.0], [9.0, 10.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_3.clone().sum_dim(1); + let tensor_5 = tensor_4.mul(tensor_3); + + let grads = tensor_5.sum().backward(); + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[494.0, 722.0], [2990.0, 4370.0]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[690.0, 690.0], [958.0, 958.0]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_mean_dim() { + let data_1 = TensorData::from([[1.0, 7.0], [-2.0, -3.0]]); + let data_2 = TensorData::from([[4.0, -7.0], [2.0, 3.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_1.clone().mul(tensor_3.mean_dim(1).unsqueeze()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[4.0, 36.0], [3.0, -17.0]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[9.0, 9.0], [35.5, 35.5]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_sum_dim() { + let data_1 = TensorData::from([[1.0, 7.0], [-2.0, -3.0]]); + let data_2 = TensorData::from([[4.0, -7.0], [2.0, 3.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_1.clone().mul(tensor_3.sum_dim(1).unsqueeze()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[8.0, 72.0], [6.0, -34.0]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[18.0, 18.0], [71.0, 71.0]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_prod() { + let device = AutodiffDevice::new(); + let tensor = + TestTensor::<1>::from_data(TensorData::from([2.0, 3.0, 4.0]), &device).require_grad(); + + let output = tensor.clone().prod(); + let grads = output.backward(); + let grad = tensor.grad(&grads).unwrap(); + + // grad_i = prod(x) / x_i = [24/2, 24/3, 24/4] + let expected = TensorData::from([12.0, 8.0, 6.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_prod_with_negatives() { + let device = AutodiffDevice::new(); + let tensor = + TestTensor::<1>::from_data(TensorData::from([2.0, -3.0, 4.0]), &device).require_grad(); + + let output = tensor.clone().prod(); + + // The forward value keeps its sign. The previous default impl computed + // exp(sum(log(x))), which returns NaN for any negative input. + output + .to_data() + .assert_approx_eq::(&TensorData::from([-24.0]), Tolerance::default()); + + let grads = output.backward(); + let grad = tensor.grad(&grads).unwrap(); + + // grad_i = prod(x) / x_i = [-24/2, -24/-3, -24/4] + let expected = TensorData::from([-12.0, 8.0, -6.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_prod_dim() { + let device = AutodiffDevice::new(); + let tensor = TestTensor::<2>::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), + &device, + ) + .require_grad(); + + let output = tensor.clone().prod_dim(1); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // Per row: grad_ij = prod(row_i) / x_ij + let expected = TensorData::from([[6.0, 3.0, 2.0], [30.0, 24.0, 20.0]]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_prod_dim_with_negatives() { + let device = AutodiffDevice::new(); + let tensor = TestTensor::<2>::from_data( + TensorData::from([[1.0, -2.0, 3.0], [-4.0, 5.0, 6.0]]), + &device, + ) + .require_grad(); + + let output = tensor.clone().prod_dim(0); + + output.clone().to_data().assert_approx_eq::( + &TensorData::from([[-4.0, -10.0, 18.0]]), + Tolerance::default(), + ); + + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // Per column: grad_ij = prod(col_j) / x_ij + let expected = TensorData::from([[-4.0, 5.0, 6.0], [1.0, -2.0, 3.0]]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +// The following tests are ignored due to the same limitation as cumprod: the +// gradient divides by the input, which produces NaN when the input contains +// zeros. The true gradient at a zero is finite (the product of the other +// elements), but recovering it needs the zero-safe exclusive-cumulative-product +// algorithm tracked in https://github.com/tracel-ai/burn/issues/3864. + +#[test] +#[ignore = "prod gradient with zeros not yet implemented - produces NaN due to division by zero"] +fn should_diff_prod_single_zero() { + let device = AutodiffDevice::new(); + let tensor = + TestTensor::<1>::from_data(TensorData::from([2.0, 0.0, 3.0, 4.0]), &device).require_grad(); + + let output = tensor.clone().prod(); + let grads = output.backward(); + let grad = tensor.grad(&grads).unwrap(); + + // Only the zero slot has a non-zero gradient (product of the others = 24). + let expected = TensorData::from([0.0, 24.0, 0.0, 0.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +#[ignore = "prod gradient with zeros not yet implemented - produces NaN due to division by zero"] +fn should_diff_prod_multiple_zeros() { + let device = AutodiffDevice::new(); + let tensor = TestTensor::<1>::from_data(TensorData::from([2.0, 0.0, 3.0, 0.0, 5.0]), &device) + .require_grad(); + + let output = tensor.clone().prod(); + let grads = output.backward(); + let grad = tensor.grad(&grads).unwrap(); + + // Every leave-one-out product still contains a zero, so the gradient is zero. + let expected = TensorData::from([0.0, 0.0, 0.0, 0.0, 0.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/all_reduce.rs b/crates/burn-backend-tests/tests/autodiff/all_reduce.rs new file mode 100644 index 0000000..17f09df --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/all_reduce.rs @@ -0,0 +1,166 @@ +use super::*; +use burn_tensor::{ + Device, DeviceType, TensorData, + distributed::{DistributedConfig, DistributedContext, ReduceOperation}, +}; +use serial_test::serial; + +// TODO + +#[test] +#[serial] +fn should_diff_all_reduce_sum() { + let devices = Device::enumerate(DeviceType::Cuda).autodiff().into_vec(); + if devices.len() < 2 { + return; + } + let (device_0, device_1) = (devices[0].clone(), devices[1].clone()); + + let in_tensor_0 = TestTensor::<1>::from_data([2.0, 5.0], &device_0).require_grad(); + let in_tensor_1 = TestTensor::<1>::from_data([4.0, 1.0], &device_1).require_grad(); + + let (output, grads) = compute_gradients( + vec![in_tensor_0, in_tensor_1], + ReduceOperation::Sum, + vec![device_0, device_1], + ); + compare_gradients(output, grads, &[6.0, 6.0], &[2.0, 2.0]); +} + +#[test] +#[serial] +fn should_diff_all_reduce_mean() { + let devices = Device::enumerate(DeviceType::Cuda).autodiff().into_vec(); + if devices.len() < 2 { + return; + } + let (device_0, device_1) = (devices[0].clone(), devices[1].clone()); + + let in_tensor_0 = TestTensor::<1>::from_data([2.0, 5.0], &device_0).require_grad(); + let in_tensor_1 = TestTensor::<1>::from_data([4.0, 1.0], &device_1).require_grad(); + + let (output, grads) = compute_gradients( + vec![in_tensor_0, in_tensor_1], + ReduceOperation::Mean, + vec![device_0, device_1], + ); + compare_gradients(output, grads, &[3.0, 3.0], &[1.0, 1.0]); +} + +#[test] +#[serial] +fn should_diff_all_reduce_complex_1() { + let devices = Device::enumerate(DeviceType::Cuda).autodiff().into_vec(); + if devices.len() < 2 { + return; + } + let (device_0, device_1) = (devices[0].clone(), devices[1].clone()); + + let in_tensor_0 = TestTensor::<1>::from_data([2.0, 5.0], &device_0).require_grad(); + let in_tensor_1 = TestTensor::<1>::from_data([4.0, 1.0], &device_1).require_grad(); + + let config = DistributedConfig { + all_reduce_op: ReduceOperation::Sum, + }; + let _context = DistributedContext::init(vec![device_0.clone(), device_1.clone()], config); + let [out_tensor_0, out_tensor_1] = &compute_all_reduce( + vec![in_tensor_0.clone(), in_tensor_1.clone()], + ReduceOperation::Sum, + vec![device_0, device_1], + )[..] else { + panic!("should have 2 tensors") + }; + + let out_tensor_1 = out_tensor_1.clone().mul_scalar(4.0); + + let grads_0 = out_tensor_0.backward(); + let grads_1 = out_tensor_1.backward(); + + let grad_0 = in_tensor_0.grad(&grads_0).unwrap(); + let grad_1 = in_tensor_1.grad(&grads_1).unwrap(); + + out_tensor_0.device().sync().unwrap(); + + out_tensor_0 + .to_data() + .assert_eq(&TensorData::from([6.0, 6.0]), false); + out_tensor_1 + .to_data() + .assert_eq(&TensorData::from([24.0, 24.0]), false); + grad_0 + .to_data() + .assert_eq(&TensorData::from([5.0, 5.0]), false); + grad_1 + .to_data() + .assert_eq(&TensorData::from([5.0, 5.0]), false); +} + +#[test] +#[serial] +fn should_diff_all_reduce_all_devices() { + let devices = Device::enumerate(DeviceType::Cuda).autodiff().into_vec(); + + let input = devices + .iter() + .enumerate() + .map(|(i, device)| { + let elem = i as f32; + TestTensor::<1>::from_data([elem, elem], device).require_grad() + }) + .collect(); + + let value: f32 = devices.iter().enumerate().map(|(i, _)| i as f32).sum(); + let grad_value = devices.len() as f32; + let (output, grads) = compute_gradients(input, ReduceOperation::Sum, devices); + compare_gradients(output, grads, &[value, value], &[grad_value, grad_value]); +} + +fn compare_gradients( + outputs: Vec>, + grads: Vec>, + expected_output: &[f32], + expected_grads: &[f32], +) { + for out in outputs { + out.to_data() + .assert_eq(&TensorData::from(expected_output), false); + } + for grad in grads { + grad.to_data() + .assert_eq(&TensorData::from(expected_grads), false); + } +} + +fn compute_gradients( + tensors: Vec>, + op: ReduceOperation, + devices: Vec, +) -> (Vec>, Vec>) { + let config = DistributedConfig { all_reduce_op: op }; + let _context = DistributedContext::init(devices.clone(), config); + + let out = compute_all_reduce(tensors.clone(), op, devices); + + let mut all_grads = vec![]; + for (in_tensor, out_tensor) in tensors.iter().zip(out.clone()) { + let grads = out_tensor.backward(); + all_grads.push(in_tensor.grad(&grads).unwrap()); + } + + (out, all_grads) +} + +fn compute_all_reduce( + tensors: Vec>, + op: ReduceOperation, + devices: Vec, +) -> Vec> { + let mut out = vec![]; + for tensor in tensors.clone() { + let out_tensor = burn_tensor::module::all_reduce(tensor, op, devices.clone()); + let out_tensor = out_tensor.resolve(); + out.push(out_tensor); + } + + out +} diff --git a/crates/burn-backend-tests/tests/autodiff/avgpool1d.rs b/crates/burn-backend-tests/tests/autodiff/avgpool1d.rs new file mode 100644 index 0000000..5ab9911 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/avgpool1d.rs @@ -0,0 +1,102 @@ +use super::*; +use burn_tensor::module::avg_pool1d; +use burn_tensor::{Shape, Tolerance}; + +#[test] +fn test_avg_pool1d_simple() { + let test = AvgPool1dTestCase { + batch_size: 1, + channels: 1, + kernel_size: 3, + padding: 0, + stride: 1, + length: 6, + count_include_pad: true, + }; + + test.assert_output(TestTensor::from_data( + [[[0.33333, 0.66667, 1.0000, 1.0000, 0.66667, 0.33333]]], + &AutodiffDevice::new(), + )); +} + +#[test] +fn test_avg_pool1d_complex() { + let test = AvgPool1dTestCase { + batch_size: 1, + channels: 2, + kernel_size: 3, + padding: 1, + stride: 2, + length: 6, + count_include_pad: true, + }; + + test.assert_output(TestTensor::from_data( + [[ + [0.33333, 0.66667, 0.33333, 0.66667, 0.33333, 0.33333], + [0.33333, 0.66667, 0.33333, 0.66667, 0.33333, 0.33333], + ]], + &AutodiffDevice::new(), + )); +} + +#[test] +fn test_avg_pool1d_complex_dont_count_pad() { + let test = AvgPool1dTestCase { + batch_size: 1, + channels: 2, + kernel_size: 3, + padding: 1, + stride: 2, + length: 6, + count_include_pad: false, + }; + + test.assert_output(TestTensor::from_data( + [[ + [0.5000, 0.83333, 0.33333, 0.66667, 0.33333, 0.33333], + [0.5000, 0.83333, 0.33333, 0.66667, 0.33333, 0.33333], + ]], + &AutodiffDevice::new(), + )); +} + +struct AvgPool1dTestCase { + batch_size: usize, + channels: usize, + kernel_size: usize, + padding: usize, + stride: usize, + length: usize, + count_include_pad: bool, +} + +impl AvgPool1dTestCase { + fn assert_output(self, x_grad: TestTensor<3>) { + let shape_x = Shape::new([self.batch_size, self.channels, self.length]); + let device = AutodiffDevice::new(); + let x = TestTensor::from_data( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<3, _>(shape_x) + .into_data(), + &device, + ) + .require_grad(); + let output = avg_pool1d( + x.clone(), + self.kernel_size, + self.stride, + self.padding, + self.count_include_pad, + false, + ); + let grads = output.backward(); + let x_grad_actual = x.grad(&grads).unwrap(); + + let tolerance = Tolerance::default().set_half_precision_relative(1e-3); + x_grad + .to_data() + .assert_approx_eq::(&x_grad_actual.into_data(), tolerance); + } +} diff --git a/crates/burn-backend-tests/tests/autodiff/avgpool2d.rs b/crates/burn-backend-tests/tests/autodiff/avgpool2d.rs new file mode 100644 index 0000000..bad5234 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/avgpool2d.rs @@ -0,0 +1,129 @@ +use super::*; +use burn_tensor::module::avg_pool2d; +use burn_tensor::{Shape, Tolerance}; + +#[test] +fn test_avg_pool2d_simple() { + let test = AvgPool2dTestCase { + batch_size: 1, + channels: 1, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + stride_1: 1, + stride_2: 1, + height: 6, + width: 6, + count_include_pad: true, + }; + + test.assert_output(TestTensor::from_data( + [[[ + [0.11111, 0.22222, 0.33333, 0.33333, 0.22222, 0.11111], + [0.22222, 0.44444, 0.66667, 0.66667, 0.44444, 0.22222], + [0.33333, 0.66667, 1.00000, 1.00000, 0.66667, 0.33333], + [0.33333, 0.66667, 1.00000, 1.00000, 0.66667, 0.33333], + [0.22222, 0.44444, 0.66667, 0.66667, 0.44444, 0.22222], + [0.11111, 0.22222, 0.33333, 0.33333, 0.22222, 0.11111], + ]]], + &AutodiffDevice::new(), + )); +} + +#[test] +fn test_avg_pool2d_complex() { + let test = AvgPool2dTestCase { + batch_size: 1, + channels: 1, + kernel_size_1: 3, + kernel_size_2: 4, + padding_1: 1, + padding_2: 2, + stride_1: 1, + stride_2: 2, + height: 4, + width: 6, + count_include_pad: true, + }; + + test.assert_output(TestTensor::from_data( + [[[ + [0.33333, 0.33333, 0.33333, 0.33333, 0.33333, 0.33333], + [0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000], + [0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000], + [0.33333, 0.33333, 0.33333, 0.33333, 0.33333, 0.33333], + ]]], + &AutodiffDevice::new(), + )); +} + +#[test] +fn test_avg_pool2d_complex_dont_include_pad() { + let test = AvgPool2dTestCase { + batch_size: 1, + channels: 1, + kernel_size_1: 3, + kernel_size_2: 4, + padding_1: 1, + padding_2: 2, + stride_1: 1, + stride_2: 2, + height: 4, + width: 6, + count_include_pad: false, + }; + + test.assert_output(TestTensor::from_data( + [[[ + [0.6250, 0.6250, 0.41667, 0.41667, 0.6250, 0.6250], + [0.8750, 0.8750, 0.58333, 0.58333, 0.8750, 0.8750], + [0.8750, 0.8750, 0.58333, 0.58333, 0.8750, 0.8750], + [0.6250, 0.6250, 0.41667, 0.41667, 0.6250, 0.6250], + ]]], + &AutodiffDevice::new(), + )); +} + +struct AvgPool2dTestCase { + batch_size: usize, + channels: usize, + kernel_size_1: usize, + kernel_size_2: usize, + padding_1: usize, + padding_2: usize, + stride_1: usize, + stride_2: usize, + height: usize, + width: usize, + count_include_pad: bool, +} + +impl AvgPool2dTestCase { + fn assert_output(self, x_grad: TestTensor<4>) { + let shape_x = Shape::new([self.batch_size, self.channels, self.height, self.width]); + let device = AutodiffDevice::new(); + let x = TestTensor::from_data( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<4, _>(shape_x) + .into_data(), + &device, + ) + .require_grad(); + let output = avg_pool2d( + x.clone(), + [self.kernel_size_1, self.kernel_size_2], + [self.stride_1, self.stride_2], + [self.padding_1, self.padding_2], + self.count_include_pad, + false, + ); + let grads = output.backward(); + let x_grad_actual = x.grad(&grads).unwrap(); + + x_grad.to_data().assert_approx_eq::( + &x_grad_actual.into_data(), + Tolerance::default().set_half_precision_relative(1e-3), + ); + } +} diff --git a/crates/burn-backend-tests/tests/autodiff/backward.rs b/crates/burn-backend-tests/tests/autodiff/backward.rs new file mode 100644 index 0000000..86cee98 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/backward.rs @@ -0,0 +1,24 @@ +use super::*; +use burn_tensor::{TensorData, module::embedding}; + +#[test] +fn test_embedding_backward() { + let weights = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let indices = TensorData::from([[0, 1], [1, 1]]); + let x = TensorData::from([ + [[1.0, 2.0], [4.0, 5.0], [3.0, 4.0]], + [[4.0, 5.0], [8.0, 5.0], [1.0, 9.0]], + ]); + let device = AutodiffDevice::new(); + let weights = TestTensor::<2>::from_data(weights, &device).require_grad(); + let indices = TestTensorInt::<2>::from_data(indices, &device); + let x = TestTensor::<3>::from_data(x, &device).require_grad(); + + let output = embedding(weights.clone(), indices); + let output = output.matmul(x); + let grads = output.backward(); + + let grad = weights.grad(&grads).unwrap(); + grad.to_data() + .assert_eq(&TensorData::from([[3., 9., 7.], [21., 35., 27.]]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/bridge.rs b/crates/burn-backend-tests/tests/autodiff/bridge.rs new file mode 100644 index 0000000..712f62f --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/bridge.rs @@ -0,0 +1,25 @@ +use super::*; +use burn_tensor::{DType, Distribution}; + +#[test] +fn test_full_precision() { + let device = AutodiffDevice::new(); + let x1 = TestTensor::<2>::random([32, 32], Distribution::Default, &device).require_grad(); + let x2 = TestTensor::<2>::random([32, 32], Distribution::Default, &device).require_grad(); + let dtype = x1.dtype(); + + let x3 = x1.clone().cast(DType::F32); + let x4 = x2.clone().cast(DType::F32); + + let x5 = x3.matmul(x4); + let x6 = x5.cast(dtype); + let x7 = x6 * x1.clone() / x2.clone(); + + let grads = x7.backward(); + + let x1_grad = x1.grad(&grads); + let x2_grad = x2.grad(&grads); + + assert!(x1_grad.is_some()); + assert!(x2_grad.is_some()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/broadcast.rs b/crates/burn-backend-tests/tests/autodiff/broadcast.rs new file mode 100644 index 0000000..3c90828 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/broadcast.rs @@ -0,0 +1,56 @@ +use super::*; + +#[test] +fn mul_broadcast() { + test_ops_broadcast_backward(|x, y| x * y); +} + +#[test] +fn div_broadcast() { + test_ops_broadcast_backward(|x, y| x / y); +} + +#[test] +fn sub_broadcast() { + test_ops_broadcast_backward(|x, y| x - y); +} + +#[test] +fn add_broadcast() { + test_ops_broadcast_backward(|x, y| x + y); +} + +#[test] +fn matmul_broadcast() { + test_ops_broadcast_backward(|x, y| x.matmul(y)); +} + +#[test] +fn mask_where_broadcast() { + test_ops_broadcast_backward(|x, y| { + let cond = y.clone().equal_elem(4); + x.mask_where(cond, y) + }); +} + +fn test_ops_broadcast_backward(func: F) +where + F: Fn(TestTensor<3>, TestTensor<3>) -> TestTensor<3>, +{ + let device = AutodiffDevice::new(); + let w = TestTensor::zeros([16, 5, 5], &device).require_grad(); + let x = TestTensor::zeros([4, 5, 5], &device).require_grad(); + + // Slice isn't a broadcastable operation, so it will fail when the previous backward pass + // of an operation that support broadcast doesn't support it during the backward pass. + let y = func(w.clone().slice([0..1]), x.clone()); + + // Will panic if broadcast isn't supported! + let grads = y.backward(); + + let w_grad = w.grad(&grads).unwrap(); + let x_grad = x.grad(&grads).unwrap(); + + assert_eq!(w_grad.shape(), w.shape()); + assert_eq!(x_grad.shape(), x.shape()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/cast.rs b/crates/burn-backend-tests/tests/autodiff/cast.rs new file mode 100644 index 0000000..503def8 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/cast.rs @@ -0,0 +1,25 @@ +// Skip on metal - F64 not supported +#![cfg(all(feature = "std", not(feature = "metal")))] + +use super::*; +use burn_backend_tests::might_panic; +use burn_tensor::{DType, TensorData}; + +#[might_panic(reason = "Unsupported precision for fusion")] +#[test] +fn cast_keeps_gradient_flow() { + let device = AutodiffDevice::new(); + + let x = TestTensor::<2>::from_data(TensorData::from([[1.0, 2.0], [3.0, 4.0]]), &device) + .require_grad(); + + let y = x.clone().cast(DType::F64); + let z = y.sum(); + + let grads = z.backward(); + let grad_x = x.grad(&grads).unwrap(); + + grad_x + .to_data() + .assert_eq(&TensorData::from([[1., 1.], [1., 1.]]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/cat.rs b/crates/burn-backend-tests/tests/autodiff/cat.rs new file mode 100644 index 0000000..f595d8e --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/cat.rs @@ -0,0 +1,105 @@ +use super::*; + +use burn_tensor::Tolerance; + +#[test] +fn should_diff_cat() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data([[2.0, -1.0], [5.0, 2.0]], &device).require_grad(); + let tensor_2 = TestTensor::<2>::from_data([[5.0, 4.0], [-1.0, 4.0]], &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let grads = tensor_3.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let mut tensor_1_list = Vec::new(); + let mut tensor_2_list = Vec::new(); + + for i in 0..2 { + tensor_1_list.push(tensor_1.clone().slice([i..i + 1])); + tensor_2_list.push(tensor_2.clone().slice([i..i + 1])); + } + + let tensor_1_cat = TestTensor::cat(tensor_1_list.clone(), 0); + let tensor_2_cat = TestTensor::cat(tensor_2_list.clone(), 0); + + let tensor_3_cat = tensor_1_cat.clone().matmul(tensor_2_cat.clone()); + let grads = tensor_3_cat.backward(); + + let grad_1_slice_1 = tensor_1.grad(&grads).unwrap().slice([0..1]); + let grad_1_slice_2 = tensor_1.grad(&grads).unwrap().slice([1..2]); + + let grad_2_slice_1 = tensor_2.grad(&grads).unwrap().slice([0..1]); + let grad_2_slice_2 = tensor_2.grad(&grads).unwrap().slice([1..2]); + + grad_1 + .clone() + .slice([0..1]) + .to_data() + .assert_approx_eq::(&grad_1_slice_1.to_data(), Tolerance::default()); + grad_1 + .slice([1..2]) + .to_data() + .assert_approx_eq::(&grad_1_slice_2.to_data(), Tolerance::default()); + + grad_2 + .clone() + .slice([0..1]) + .to_data() + .assert_approx_eq::(&grad_2_slice_1.to_data(), Tolerance::default()); + grad_2 + .slice([1..2]) + .to_data() + .assert_approx_eq::(&grad_2_slice_2.to_data(), Tolerance::default()); +} + +#[test] +fn should_diff_cat_more_than_1_dim() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data([[2.0, -1.0], [5.0, 2.0]], &device).require_grad(); + let tensor_2 = + TestTensor::<2>::from_data([[5.0, 4.0], [-1.0, 4.0], [4.0, 1.0]], &device).require_grad(); + + // Concat a tensor [2, 2] with another tensor [3, 2] along dim 0. + // The resulting tensor should be [5, 2] + let tensor_3 = TestTensor::cat(vec![tensor_1.clone(), tensor_2.clone()], 0); + assert_eq!(tensor_3.dims(), [5, 2]); + let grads = tensor_3.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + assert_eq!(tensor_1.dims(), grad_1.dims()); + assert_eq!(tensor_2.dims(), grad_2.dims()); +} + +#[test] +fn should_slice_grads_correctly_when_some_inputs_not_tracked() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data([[1.0]], &device).require_grad(); // tracked + let tensor_2 = TestTensor::<2>::from_data([[10.0, 20.0]], &device); // not tracked + let tensor_3 = TestTensor::<2>::from_data([[100.0, 200.0, 300.0]], &device).require_grad(); // tracked + + let cat = TestTensor::cat( + vec![tensor_1.clone(), tensor_2.clone(), tensor_3.clone()], + 1, + ); + + // Make gradient per column unique so wrong slicing shows up. + let weights = TestTensor::<2>::from_data([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]], &device); + let loss = (cat * weights).sum(); + + let grads = loss.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_3 = tensor_3.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&burn_tensor::TensorData::from([[1.0]]), false); + grad_3 + .to_data() + .assert_eq(&burn_tensor::TensorData::from([[4.0, 5.0, 6.0]]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/ceil.rs b/crates/burn-backend-tests/tests/autodiff/ceil.rs new file mode 100644 index 0000000..639d028 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/ceil.rs @@ -0,0 +1,21 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_diff_ceil() { + let data = TensorData::from([ + [-1.9751, 0.0714, 0.0643, 0.2406], + [-1.3172, 0.1252, -0.1119, -0.0127], + ]); + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data, &device).require_grad(); + let tensor_2 = tensor_1.clone().ceil(); + let grads = tensor_2.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + + grad_1.to_data().assert_eq( + &TensorData::from([[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]]), + false, + ); +} diff --git a/crates/burn-backend-tests/tests/autodiff/checkpoint.rs b/crates/burn-backend-tests/tests/autodiff/checkpoint.rs new file mode 100644 index 0000000..1d16f72 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/checkpoint.rs @@ -0,0 +1,209 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_autodiff_checkpoint_complicated_computation() { + let data_0 = TensorData::from([[0.0, 7.0], [7.0, 7.0]]); + let data_1 = TensorData::from([[0.1, 7.0], [7.0, 7.0]]); + let data_2 = TensorData::from([[0.2, 7.0], [7.0, 7.0]]); + let data_3 = TensorData::from([[0.3, 7.0], [7.0, 7.0]]); + let data_4 = TensorData::from([[0.4, 7.0], [7.0, 7.0]]); + + let device = AutodiffDevice::new(); + let tensor_0 = TestTensor::<2>::from_data(data_0, &device).require_grad(); + let tensor_1 = TestTensor::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + let tensor_3 = TestTensor::from_data(data_3, &device).require_grad(); + let tensor_4 = TestTensor::from_data(data_4, &device).require_grad(); + + let tensor_5 = compute_bound_eager(tensor_0, tensor_1); + let tensor_6 = compute_bound_lazy(tensor_2, tensor_3.clone()); + let tensor_7 = memory_bound_eager(tensor_3, tensor_4); + let tensor_8 = compute_bound_lazy(tensor_6, tensor_7.clone()); + let tensor_9 = memory_bound_eager_scalar(tensor_7, 11.); + let tensor_10 = memory_bound_lazy(tensor_5, tensor_8.clone()); + let tensor_11 = memory_bound_lazy(tensor_8, tensor_9); + let tensor_12 = compute_bound_lazy(tensor_10, tensor_11); + + assert_checkpoint(tensor_12); +} + +#[test] +fn test_autodiff_checkpoint_with_missing_requirement() { + let data_0 = TensorData::from([[0.0, 7.0], [7.0, 7.0]]); + let data_1 = TensorData::from([[0.1, 7.0], [7.0, 7.0]]); + + let device = AutodiffDevice::new(); + let tensor_0 = TestTensor::<2>::from_data(data_0, &device).require_grad(); + let tensor_1 = TestTensor::from_data(data_1, &device); // does not require_grad + + let tensor_2 = memory_bound_eager(tensor_0, tensor_1); + let tensor_3 = memory_bound_eager_scalar(tensor_2.clone(), 11.); + let tensor_4 = memory_bound_eager_scalar(tensor_2.clone(), 11.); + let tensor_5 = compute_bound_lazy(tensor_3, tensor_4); + let tensor_6 = compute_bound_eager_scalar(tensor_5.clone(), 11.); + let tensor_7 = memory_bound_eager(tensor_5, tensor_2); + let tensor_8 = memory_bound_eager(tensor_6, tensor_7); + + assert_checkpoint(tensor_8); +} + +#[test] +fn test_autodiff_checkpoint_with_many_duplicates() { + let data_0 = TensorData::from([[4.0, 7.0], [7.0, 7.0]]); + + let device = AutodiffDevice::new(); + let tensor_0 = TestTensor::<2>::from_data(data_0, &device).require_grad(); + + let tensor_1 = memory_bound_eager(tensor_0.clone(), tensor_0.clone()); + let tensor_2 = compute_bound_eager(tensor_0.clone(), tensor_0.clone()); + let tensor_3 = memory_bound_lazy(tensor_0.clone(), tensor_0.clone()); + let tensor_4 = compute_bound_lazy(tensor_0.clone(), tensor_0.clone()); + + let tensor_5 = memory_bound_eager(tensor_1.clone(), tensor_0.clone()); + let tensor_6 = memory_bound_eager(tensor_0.clone(), tensor_5.clone()); + let tensor_7 = compute_bound_lazy(tensor_3.clone(), tensor_5.clone()); + let tensor_8 = compute_bound_eager(tensor_4.clone(), tensor_2.clone()); + let tensor_9 = memory_bound_lazy(tensor_6, tensor_7); + let tensor_10 = memory_bound_eager(tensor_0, tensor_9); + let tensor_11 = memory_bound_eager_scalar(tensor_10, 9.); + let tensor_12 = compute_bound_lazy(tensor_8, tensor_11); + + assert_checkpoint(tensor_12); +} + +#[test] +fn test_autodiff_checkpoint_with_long_chain_of_eager_memory_bound() { + let data_0 = TensorData::from([[0.0, 7.0], [7.0, 7.0]]); + let data_1 = TensorData::from([[0.1, 7.0], [7.0, 7.0]]); + let data_2 = TensorData::from([[0.2, 7.0], [7.0, 7.0]]); + let data_3 = TensorData::from([[0.3, 7.0], [7.0, 7.0]]); + let data_4 = TensorData::from([[0.4, 7.0], [7.0, 7.0]]); + + let device = AutodiffDevice::new(); + let tensor_0 = TestTensor::<2>::from_data(data_0, &device).require_grad(); + let tensor_1 = TestTensor::from_data(data_1, &device); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + let tensor_3 = TestTensor::from_data(data_3, &device).require_grad(); + let tensor_4 = TestTensor::from_data(data_4, &device).require_grad(); + + let tensor_5 = memory_bound_eager(tensor_0, tensor_1.clone()); + let tensor_6 = memory_bound_eager(tensor_5, tensor_2); + let tensor_7 = memory_bound_eager(tensor_6, tensor_3); + let tensor_8 = memory_bound_eager(tensor_7, tensor_4); + let tensor_9 = memory_bound_lazy(tensor_8, tensor_1); + + assert_checkpoint(tensor_9) +} + +#[test] +fn test_autodiff_checkpoint_half_sub_graph_not_tracked() { + let data_0 = TensorData::from([[0.0, 7.0], [7.0, 7.0]]); + let data_1 = TensorData::from([[0.1, 7.0], [7.0, 7.0]]); + let data_2 = TensorData::from([[0.2, 7.0], [7.0, 7.0]]); + let data_3 = TensorData::from([[0.3, 7.0], [7.0, 7.0]]); + let data_4 = TensorData::from([[0.4, 7.0], [7.0, 7.0]]); + let data_5 = TensorData::from([[0.5, 7.0], [7.0, 7.0]]); + + let device = AutodiffDevice::new(); + let tensor_0 = TestTensor::<2>::from_data(data_0, &device); + let tensor_1 = TestTensor::from_data(data_1, &device); + let tensor_2 = TestTensor::from_data(data_2, &device); + let tensor_3 = TestTensor::from_data(data_3, &device).require_grad(); + let tensor_4 = TestTensor::from_data(data_4, &device).require_grad(); + let tensor_5 = TestTensor::from_data(data_5, &device).require_grad(); + + let tensor_6 = memory_bound_lazy(tensor_0, tensor_1); + let tensor_7 = compute_bound_eager(tensor_6, tensor_2); + + let tensor_8 = memory_bound_eager(tensor_3, tensor_4); + let tensor_9 = compute_bound_lazy(tensor_8, tensor_5); + + let tensor_10 = compute_bound_lazy(tensor_7, tensor_9); + + assert_checkpoint(tensor_10); +} + +#[test] +fn test_autodiff_checkpoint_very_complex() { + let data_0 = TensorData::from([[0.0, 7.0], [7.0, 7.0]]); + let data_1 = TensorData::from([[0.1, 7.0], [7.0, 7.0]]); + let data_2 = TensorData::from([[0.2, 7.0], [7.0, 7.0]]); + let data_3 = TensorData::from([[0.3, 7.0], [7.0, 7.0]]); + let data_4 = TensorData::from([[0.4, 7.0], [7.0, 7.0]]); + + let device = AutodiffDevice::new(); + let tensor_0 = TestTensor::<2>::from_data(data_0, &device).require_grad(); + let tensor_1 = TestTensor::from_data(data_1, &device); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + let tensor_3 = TestTensor::from_data(data_3, &device).require_grad(); + let tensor_4 = TestTensor::from_data(data_4, &device).require_grad(); + + let tensor_5 = memory_bound_eager_scalar(tensor_0, 8.); + let tensor_6 = memory_bound_lazy(tensor_5.clone(), tensor_1.clone()); + let tensor_7 = compute_bound_lazy(tensor_6.clone(), tensor_6); + let tensor_8 = memory_bound_lazy(tensor_1.clone(), tensor_5.clone()); + let tensor_9 = memory_bound_eager_scalar(tensor_7.clone(), 7.); + let tensor_10 = compute_bound_eager(tensor_5, tensor_8); + let tensor_11 = memory_bound_eager(tensor_2.clone(), tensor_9); + let tensor_12 = memory_bound_lazy(tensor_2.clone(), tensor_2); + let tensor_13 = compute_bound_eager(tensor_10.clone(), tensor_11); + let tensor_14 = compute_bound_eager_scalar(tensor_3, 8.); + let tensor_15 = compute_bound_lazy(tensor_4, tensor_12); + let tensor_16 = memory_bound_lazy(tensor_10, tensor_7); + let tensor_17 = compute_bound_lazy(tensor_13, tensor_1); + let tensor_18 = memory_bound_eager(tensor_15, tensor_16); + let tensor_19 = compute_bound_eager(tensor_14, tensor_17); + let tensor_20 = memory_bound_lazy(tensor_18, tensor_19); + let tensor_21 = memory_bound_eager_scalar(tensor_20, 8.); + + assert_checkpoint(tensor_21) +} + +fn assert_checkpoint(tensor: TestTensor) { + // Assert is not explicit here, but the test can fail + // - when a tensor is actually required more than n_required, it won't be found and will panic + // - when a tensor is actually required less than n_required, the backward states map won't be + // empty and will fail the assertion within the backward code, same for retro_forwards + tensor.backward(); +} + +// Does not save its state and does not need its parents +fn memory_bound_eager( + tensor_a: TestTensor, + tensor_b: TestTensor, +) -> TestTensor { + tensor_a.add(tensor_b) +} +fn memory_bound_eager_scalar(tensor_a: TestTensor, b: f32) -> TestTensor { + tensor_a.add_scalar(b) +} + +// Saves its own state and does not need its parents +fn compute_bound_eager( + tensor_a: TestTensor, + tensor_b: TestTensor, +) -> TestTensor { + let mask = TestTensorBool::::empty(tensor_a.shape(), &tensor_a.device()); + tensor_a.mask_where(mask, tensor_b) +} +fn compute_bound_eager_scalar(tensor_a: TestTensor, b: f32) -> TestTensor { + let mask = TestTensorBool::::empty(tensor_a.shape(), &tensor_a.device()); + tensor_a.mask_fill(mask, b) +} + +// Does not save its state and needs its parents +fn memory_bound_lazy( + tensor_a: TestTensor, + tensor_b: TestTensor, +) -> TestTensor { + tensor_a.mul(tensor_b) +} + +// Saves its own state and needs its parents +fn compute_bound_lazy( + tensor_a: TestTensor, + tensor_b: TestTensor, +) -> TestTensor { + tensor_a.matmul(tensor_b) +} diff --git a/crates/burn-backend-tests/tests/autodiff/complex.rs b/crates/burn-backend-tests/tests/autodiff/complex.rs new file mode 100644 index 0000000..12adec2 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/complex.rs @@ -0,0 +1,81 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_diff_full_complex_1() { + let data_1 = TensorData::from([[1.0, 7.0], [13.0, -3.0]]); + let data_2 = TensorData::from([[4.0, 7.0], [2.0, 3.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_3.matmul(tensor_1.clone()); + let tensor_5 = tensor_4.mul(tensor_2.clone()); + + let grads = tensor_5.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[593., 463.0], [487.0, 539.0]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([[734.0, 294.0], [1414.0, 242.0]]), false); +} + +#[test] +fn should_diff_full_complex_2() { + let data_1 = TensorData::from([[1.0, 7.0], [13.0, -3.0]]); + let data_2 = TensorData::from([[4.0, 7.0], [2.0, 3.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_3.matmul(tensor_1.clone()); + let tensor_5 = tensor_4.add_scalar(17.0).add(tensor_2.clone()); + + let grads = tensor_5.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[166.0, 110.0], [212.0, 156.0]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([[113.0, 141.0], [33.0, 41.0]]), false); +} + +#[test] +fn should_diff_full_complex_3() { + let data_1 = TensorData::from([[1.0, 7.0], [13.0, -3.0]]); + let data_2 = TensorData::from([[4.0, 7.0], [2.0, 3.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_3.matmul(tensor_1.clone()); + let tensor_5 = tensor_4.clone().sub(tensor_2.clone()); + let tensor_6 = tensor_5.add(tensor_4); + + let grads = tensor_6.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[332.0, 220.0], [424.0, 312.0]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([[223.0, 279.0], [63.0, 79.0]]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/conv1d.rs b/crates/burn-backend-tests/tests/autodiff/conv1d.rs new file mode 100644 index 0000000..b4a4795 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/conv1d.rs @@ -0,0 +1,347 @@ +#![allow(clippy::identity_op)] + +use super::*; +use burn_tensor::{ + Shape, Tolerance, + module::conv1d, + ops::{ConvOptions, PaddedConvOptions}, +}; + +#[test] +fn test_conv1d_basic() { + let test = Conv1dTestCase { + batch_size: 2, + channels_in: 2, + channels_out: 2, + kernel_size: 3, + padding: 1, + stride: 1, + dilation: 1, + groups: 1, + length: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [[14., 24., 24., 18.], [26., 42., 42., 30.]], + [[14., 24., 24., 18.], [26., 42., 42., 30.]], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [[30., 44., 36.], [54., 76., 60.]], + [[30., 44., 36.], [54., 76., 60.]], + ], + &device, + ), + bias: TestTensor::from_data([8., 8.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv1d_different_channels() { + let test = Conv1dTestCase { + batch_size: 2, + channels_in: 2, + channels_out: 3, + kernel_size: 3, + padding: 1, + stride: 1, + dilation: 1, + groups: 1, + length: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [[39., 63., 63., 45.], [57., 90., 90., 63.]], + [[39., 63., 63., 45.], [57., 90., 90., 63.]], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [[30., 44., 36.], [54., 76., 60.]], + [[30., 44., 36.], [54., 76., 60.]], + [[30., 44., 36.], [54., 76., 60.]], + ], + &device, + ), + bias: TestTensor::from_data([8., 8., 8.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv1d_with_padding() { + let test = Conv1dTestCase { + batch_size: 2, + channels_in: 2, + channels_out: 2, + kernel_size: 3, + padding: 2, + stride: 1, + dilation: 1, + groups: 1, + length: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [[24., 24., 24., 24.], [42., 42., 42., 42.]], + [[24., 24., 24., 24.], [42., 42., 42., 42.]], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [[44., 44., 44.], [76., 76., 76.]], + [[44., 44., 44.], [76., 76., 76.]], + ], + &device, + ), + bias: TestTensor::from_data([12., 12.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv1d_with_stride() { + let test = Conv1dTestCase { + batch_size: 2, + channels_in: 2, + channels_out: 2, + kernel_size: 3, + padding: 1, + stride: 2, + dilation: 1, + groups: 1, + length: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [[8., 16., 8., 10.], [14., 28., 14., 16.]], + [[8., 16., 8., 10.], [14., 28., 14., 16.]], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [[10., 20., 24.], [18., 36., 40.]], + [[10., 20., 24.], [18., 36., 40.]], + ], + &device, + ), + bias: TestTensor::from_data([4., 4.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv1d_dilation() { + let test = Conv1dTestCase { + batch_size: 2, + channels_in: 2, + channels_out: 2, + kernel_size: 3, + padding: 1, + stride: 1, + dilation: 2, + groups: 1, + length: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [[6., 8., 8., 10.], [12., 14., 14., 16.]], + [[6., 8., 8., 10.], [12., 14., 14., 16.]], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [[8., 22., 14.], [16., 38., 22.]], + [[8., 22., 14.], [16., 38., 22.]], + ], + &device, + ), + bias: TestTensor::from_data([4., 4.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv1d_groups() { + let test = Conv1dTestCase { + batch_size: 2, + channels_in: 2, + channels_out: 2, + kernel_size: 3, + padding: 1, + stride: 1, + dilation: 1, + groups: 2, + length: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [[1., 3., 3., 3.], [7., 12., 12., 9.]], + [[1., 3., 3., 3.], [7., 12., 12., 9.]], + ], + &device, + ), + weight: TestTensor::from_data([[[30., 44., 36.]], [[54., 76., 60.]]], &device), + bias: TestTensor::from_data([8., 8.], &device), + }; + test.assert_grads(grads); +} + +/// Regression test for https://github.com/tracel-ai/burn/issues/4799. +/// +/// When a conv drops more than one tail input (here length=11, kernel=4, +/// stride=4 → three dropped inputs), the transpose-conv used in the backward +/// pass must inject enough `padding_out` to reproduce the original input shape. +/// An earlier formula capped this at 1, which caused downstream ops (pad, +/// slice_assign) to panic on shape mismatch. +#[test] +fn test_conv1d_backward_shape_with_remainder() { + let device = AutodiffDevice::new(); + let length = 11; + let x = TestTensor::<3>::from_data( + TestTensorInt::arange(0..(1 * 1 * length) as i64, &device) + .reshape::<3, _>(Shape::new([1, 1, length])) + .into_data(), + &device, + ) + .require_grad(); + let weight = TestTensor::<3>::from_data( + TestTensorInt::arange(0..4, &device) + .reshape::<3, _>(Shape::new([1, 1, 4])) + .into_data(), + &device, + ) + .require_grad(); + + let output = conv1d( + x.clone(), + weight.clone(), + None, + ConvOptions::new([4], [0], [1], 1), + ); + let grads = output.sum().backward(); + + let x_grad = x.grad(&grads).unwrap(); + let weight_grad = weight.grad(&grads).unwrap(); + assert_eq!(x_grad.dims(), [1, 1, length]); + assert_eq!(weight_grad.dims(), [1, 1, 4]); +} + +/// Regression test for the asymmetric-padding backward-shape path referenced +/// in https://github.com/tracel-ai/burn/issues/4799. Reduced to a single +/// `conv1d`: the asymmetric path routes through `x.pad(...) -> B::conv1d(padding=0)`, +/// and the backward must restore the original input shape without triggering +/// shape-mismatch failures when the forward drops multiple tail inputs. +#[test] +fn test_conv1d_asymmetric_padding_backward_shape() { + let device = AutodiffDevice::new(); + let length = 1600; + let x = TestTensor::<3>::zeros([1, 32, length], &device).require_grad(); + let weight = TestTensor::<3>::zeros([64, 32, 4], &device).require_grad(); + + let output = conv1d( + x.clone(), + weight.clone(), + None, + PaddedConvOptions::asymmetric([4], [3], [0], [1], 1), + ); + let grads = output.sum().backward(); + + assert_eq!(x.grad(&grads).unwrap().dims(), [1, 32, length]); + assert_eq!(weight.grad(&grads).unwrap().dims(), [64, 32, 4]); +} + +struct Conv1dTestCase { + batch_size: usize, + channels_in: usize, + channels_out: usize, + kernel_size: usize, + padding: usize, + stride: usize, + dilation: usize, + groups: usize, + length: usize, +} + +struct Grads { + x: TestTensor<3>, + weight: TestTensor<3>, + bias: TestTensor<1>, +} + +impl Conv1dTestCase { + fn assert_grads(self, expected_grads: Grads) { + let shape_x = Shape::new([self.batch_size, self.channels_in, self.length]); + let shape_weight = Shape::new([ + self.channels_out, + self.channels_in / self.groups, + self.kernel_size, + ]); + let device = AutodiffDevice::new(); + let weight = TestTensor::from_data( + TestTensorInt::arange(0..shape_weight.num_elements() as i64, &device) + .reshape::<3, _>(shape_weight) + .into_data(), + &device, + ) + .require_grad(); + let bias = TestTensor::from_data( + TestTensorInt::arange(0..self.channels_out as i64, &device).into_data(), + &device, + ) + .require_grad(); + let x = TestTensor::from_data( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<3, _>(shape_x) + .into_data(), + &device, + ) + .require_grad(); + + let output = conv1d( + x.clone(), + weight.clone(), + Some(bias.clone()), + ConvOptions::new([self.stride], [self.padding], [self.dilation], self.groups), + ); + let grads = output.backward(); + + // Assert + let x_grad_actual = x.grad(&grads).unwrap(); + let weight_grad_actual = weight.grad(&grads).unwrap(); + let bias_grad_actual = bias.grad(&grads).unwrap(); + + let tolerance = Tolerance::default(); + expected_grads + .bias + .to_data() + .assert_approx_eq::(&bias_grad_actual.to_data(), tolerance); + expected_grads + .weight + .to_data() + .assert_approx_eq::(&weight_grad_actual.to_data(), tolerance); + expected_grads + .x + .to_data() + .assert_approx_eq::(&x_grad_actual.to_data(), tolerance); + } +} diff --git a/crates/burn-backend-tests/tests/autodiff/conv2d.rs b/crates/burn-backend-tests/tests/autodiff/conv2d.rs new file mode 100644 index 0000000..af640bc --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/conv2d.rs @@ -0,0 +1,962 @@ +use super::*; +use burn_tensor::{Shape, Tolerance, module::conv2d, ops::ConvOptions}; + +#[test] +fn test_conv2d_basic() { + let test = Conv2dTestCase { + batch_size: 2, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 1, + height: 4, + width: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [ + [ + [88., 138., 138., 96.], + [150., 234., 234., 162.], + [150., 234., 234., 162.], + [112., 174., 174., 120.], + ], + [ + [160., 246., 246., 168.], + [258., 396., 396., 270.], + [258., 396., 396., 270.], + [184., 282., 282., 192.], + ], + ], + [ + [ + [88., 138., 138., 96.], + [150., 234., 234., 162.], + [150., 234., 234., 162.], + [112., 174., 174., 120.], + ], + [ + [160., 246., 246., 168.], + [258., 396., 396., 270.], + [258., 396., 396., 270.], + [184., 282., 282., 192.], + ], + ], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [[378., 516., 396.], [552., 752., 576.], [450., 612., 468.]], + [[666., 900., 684.], [936., 1264., 960.], [738., 996., 756.]], + ], + [ + [[378., 516., 396.], [552., 752., 576.], [450., 612., 468.]], + [[666., 900., 684.], [936., 1264., 960.], [738., 996., 756.]], + ], + ], + &device, + ), + bias: TestTensor::from_data([32., 32.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv2d_different_channels() { + let test = Conv2dTestCase { + batch_size: 2, + channels_in: 2, + channels_out: 3, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 1, + height: 4, + width: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [ + [ + [240., 369., 369., 252.], + [387., 594., 594., 405.], + [387., 594., 594., 405.], + [276., 423., 423., 288.], + ], + [ + [348., 531., 531., 360.], + [549., 837., 837., 567.], + [549., 837., 837., 567.], + [384., 585., 585., 396.], + ], + ], + [ + [ + [240., 369., 369., 252.], + [387., 594., 594., 405.], + [387., 594., 594., 405.], + [276., 423., 423., 288.], + ], + [ + [348., 531., 531., 360.], + [549., 837., 837., 567.], + [549., 837., 837., 567.], + [384., 585., 585., 396.], + ], + ], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [[378., 516., 396.], [552., 752., 576.], [450., 612., 468.]], + [[666., 900., 684.], [936., 1264., 960.], [738., 996., 756.]], + ], + [ + [[378., 516., 396.], [552., 752., 576.], [450., 612., 468.]], + [[666., 900., 684.], [936., 1264., 960.], [738., 996., 756.]], + ], + [ + [[378., 516., 396.], [552., 752., 576.], [450., 612., 468.]], + [[666., 900., 684.], [936., 1264., 960.], [738., 996., 756.]], + ], + ], + &device, + ), + bias: TestTensor::from_data([32., 32., 32.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv2d_different_kernel_size() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 4, + padding_1: 1, + padding_2: 1, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 1, + height: 4, + width: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [116., 180., 192., 132.], + [198., 306., 324., 222.], + [198., 306., 324., 222.], + [148., 228., 240., 164.], + ], + [ + [212., 324., 336., 228.], + [342., 522., 540., 366.], + [342., 522., 540., 366.], + [244., 372., 384., 260.], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [ + [27., 45., 54., 39.], + [52., 84., 96., 68.], + [51., 81., 90., 63.], + ], + [ + [123., 189., 198., 135.], + [180., 276., 288., 196.], + [147., 225., 234., 159.], + ], + ], + [ + [ + [27., 45., 54., 39.], + [52., 84., 96., 68.], + [51., 81., 90., 63.], + ], + [ + [123., 189., 198., 135.], + [180., 276., 288., 196.], + [147., 225., 234., 159.], + ], + ], + ], + &device, + ), + bias: TestTensor::from_data([12., 12.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv2d_different_padding() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 2, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 1, + height: 4, + width: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [138., 138., 138., 138.], + [234., 234., 234., 234.], + [234., 234., 234., 234.], + [174., 174., 174., 174.], + ], + [ + [246., 246., 246., 246.], + [396., 396., 396., 396.], + [396., 396., 396., 396.], + [282., 282., 282., 282.], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [[66., 66., 66.], [120., 120., 120.], [114., 114., 114.]], + [[258., 258., 258.], [376., 376., 376.], [306., 306., 306.]], + ], + [ + [[66., 66., 66.], [120., 120., 120.], [114., 114., 114.]], + [[258., 258., 258.], [376., 376., 376.], [306., 306., 306.]], + ], + ], + &device, + ), + bias: TestTensor::from_data([24., 24.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv2d_different_width() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 1, + height: 4, + width: 5, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [88., 138., 138., 138., 96.], + [150., 234., 234., 234., 162.], + [150., 234., 234., 234., 162.], + [112., 174., 174., 174., 120.], + ], + [ + [160., 246., 246., 246., 168.], + [258., 396., 396., 396., 270.], + [258., 396., 396., 396., 270.], + [184., 282., 282., 282., 192.], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [[78., 105., 90.], [144., 190., 160.], [138., 180., 150.]], + [[318., 405., 330.], [464., 590., 480.], [378., 480., 390.]], + ], + [ + [[78., 105., 90.], [144., 190., 160.], [138., 180., 150.]], + [[318., 405., 330.], [464., 590., 480.], [378., 480., 390.]], + ], + ], + &device, + ), + bias: TestTensor::from_data([20., 20.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv2d_stride_2() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + stride_1: 2, + stride_2: 2, + dilation_1: 1, + dilation_2: 1, + groups: 1, + height: 6, + width: 6, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [26., 52., 26., 52., 26., 28.], + [52., 104., 52., 104., 52., 56.], + [26., 52., 26., 52., 26., 28.], + [52., 104., 52., 104., 52., 56.], + [26., 52., 26., 52., 26., 28.], + [32., 64., 32., 64., 32., 34.], + ], + [ + [44., 88., 44., 88., 44., 46.], + [88., 176., 88., 176., 88., 92.], + [44., 88., 44., 88., 44., 46.], + [88., 176., 88., 176., 88., 92.], + [44., 88., 44., 88., 44., 46.], + [50., 100., 50., 100., 50., 52.], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [[56., 84., 90.], [84., 126., 135.], [120., 180., 189.]], + [[200., 300., 306.], [300., 450., 459.], [336., 504., 513.]], + ], + [ + [[56., 84., 90.], [84., 126., 135.], [120., 180., 189.]], + [[200., 300., 306.], [300., 450., 459.], [336., 504., 513.]], + ], + ], + &device, + ), + bias: TestTensor::from_data([9., 9.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv2d_different_stride() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + stride_1: 3, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 1, + height: 8, + width: 8, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [50., 78., 78., 78., 78., 78., 78., 54.], + [62., 96., 96., 96., 96., 96., 96., 66.], + [38., 60., 60., 60., 60., 60., 60., 42.], + [50., 78., 78., 78., 78., 78., 78., 54.], + [62., 96., 96., 96., 96., 96., 96., 66.], + [38., 60., 60., 60., 60., 60., 60., 42.], + [50., 78., 78., 78., 78., 78., 78., 54.], + [62., 96., 96., 96., 96., 96., 96., 66.], + ], + [ + [86., 132., 132., 132., 132., 132., 132., 90.], + [98., 150., 150., 150., 150., 150., 150., 102.], + [74., 114., 114., 114., 114., 114., 114., 78.], + [86., 132., 132., 132., 132., 132., 132., 90.], + [98., 150., 150., 150., 150., 150., 150., 102.], + [74., 114., 114., 114., 114., 114., 114., 78.], + [86., 132., 132., 132., 132., 132., 132., 90.], + [98., 150., 150., 150., 150., 150., 150., 102.], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [[434., 504., 448.], [567., 660., 588.], [735., 852., 756.]], + [ + [1330., 1528., 1344.], + [1911., 2196., 1932.], + [2079., 2388., 2100.], + ], + ], + [ + [[434., 504., 448.], [567., 660., 588.], [735., 852., 756.]], + [ + [1330., 1528., 1344.], + [1911., 2196., 1932.], + [2079., 2388., 2100.], + ], + ], + ], + &device, + ), + bias: TestTensor::from_data([24., 24.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv2d_dilation_2() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + stride_1: 1, + stride_2: 1, + dilation_1: 2, + dilation_2: 2, + groups: 1, + height: 6, + width: 6, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [18., 38., 38., 42., 42., 22.], + [42., 88., 88., 96., 96., 50.], + [42., 88., 88., 96., 96., 50.], + [54., 112., 112., 120., 120., 62.], + [54., 112., 112., 120., 120., 62.], + [30., 62., 62., 66., 66., 34.], + ], + [ + [36., 74., 74., 78., 78., 40.], + [78., 160., 160., 168., 168., 86.], + [78., 160., 160., 168., 168., 86.], + [90., 184., 184., 192., 192., 98.], + [90., 184., 184., 192., 192., 98.], + [48., 98., 98., 102., 102., 52.], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [[63., 102., 90.], [192., 280., 228.], [225., 318., 252.]], + [[387., 534., 414.], [624., 856., 660.], [549., 750., 576.]], + ], + [ + [[63., 102., 90.], [192., 280., 228.], [225., 318., 252.]], + [[387., 534., 414.], [624., 856., 660.], [549., 750., 576.]], + ], + ], + &device, + ), + bias: TestTensor::from_data([16., 16.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv2d_different_dilation() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + stride_1: 1, + stride_2: 1, + dilation_1: 2, + dilation_2: 3, + groups: 1, + height: 6, + width: 6, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [18., 0., 20., 20., 0., 22.], + [42., 0., 46., 46., 0., 50.], + [42., 0., 46., 46., 0., 50.], + [54., 0., 58., 58., 0., 62.], + [54., 0., 58., 58., 0., 62.], + [30., 0., 32., 32., 0., 34.], + ], + [ + [36., 0., 38., 38., 0., 40.], + [78., 0., 82., 82., 0., 86.], + [78., 0., 82., 82., 0., 86.], + [90., 0., 94., 94., 0., 98.], + [90., 0., 94., 94., 0., 98.], + [48., 0., 50., 50., 0., 52.], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [[18., 51., 33.], [60., 140., 80.], [72., 159., 87.]], + [[126., 267., 141.], [204., 428., 224.], [180., 375., 195.]], + ], + [ + [[18., 51., 33.], [60., 140., 80.], [72., 159., 87.]], + [[126., 267., 141.], [204., 428., 224.], [180., 375., 195.]], + ], + ], + &device, + ), + bias: TestTensor::from_data([8., 8.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv2d_groups() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 2, + height: 5, + width: 5, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [0., 1., 3., 3., 2.], + [3., 8., 15., 12., 7.], + [9., 21., 36., 27., 15.], + [9., 20., 33., 24., 13.], + [6., 13., 21., 15., 8.], + ], + [ + [9., 19., 30., 21., 11.], + [21., 44., 69., 48., 25.], + [36., 75., 117., 81., 42.], + [27., 56., 87., 60., 31.], + [15., 31., 48., 33., 17.], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [[[54., 63., 72.], [99., 108., 117.], [144., 153., 162.]]], + [[[279., 288., 297.], [324., 333., 342.], [369., 378., 387.]]], + ], + &device, + ), + bias: TestTensor::from_data([9., 9.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv2d_groups_stride_2() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 4, + channels_out: 4, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + stride_1: 2, + stride_2: 2, + dilation_1: 1, + dilation_2: 1, + groups: 4, + height: 4, + width: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [4., 8., 4., 5.], + [8., 16., 8., 10.], + [4., 8., 4., 5.], + [7., 14., 7., 8.], + ], + [ + [13., 26., 13., 14.], + [26., 52., 26., 28.], + [13., 26., 13., 14.], + [16., 32., 16., 17.], + ], + [ + [22., 44., 22., 23.], + [44., 88., 44., 46.], + [22., 44., 22., 23.], + [25., 50., 25., 26.], + ], + [ + [31., 62., 31., 32.], + [62., 124., 62., 64.], + [31., 62., 31., 32.], + [34., 68., 34., 35.], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [[[5., 10., 12.], [10., 20., 24.], [18., 36., 40.]]], + [[[21., 42., 44.], [42., 84., 88.], [50., 100., 104.]]], + [[[37., 74., 76.], [74., 148., 152.], [82., 164., 168.]]], + [[[53., 106., 108.], [106., 212., 216.], [114., 228., 232.]]], + ], + &device, + ), + bias: TestTensor::from_data([4., 4., 4., 4.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv2d_groups_different_channels() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 3, + channels_out: 6, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 3, + height: 4, + width: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [9., 20., 24., 13.], + [24., 52., 60., 32.], + [36., 76., 84., 44.], + [21., 44., 48., 25.], + ], + [ + [45., 92., 96., 49.], + [96., 196., 204., 104.], + [108., 220., 228., 116.], + [57., 116., 120., 61.], + ], + [ + [81., 164., 168., 85.], + [168., 340., 348., 176.], + [180., 364., 372., 188.], + [93., 188., 192., 97.], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [[[10., 14., 18.], [26., 30., 34.], [42., 46., 50.]]], + [[[10., 14., 18.], [26., 30., 34.], [42., 46., 50.]]], + [[[74., 78., 82.], [90., 94., 98.], [106., 110., 114.]]], + [[[74., 78., 82.], [90., 94., 98.], [106., 110., 114.]]], + [[[138., 142., 146.], [154., 158., 162.], [170., 174., 178.]]], + [[[138., 142., 146.], [154., 158., 162.], [170., 174., 178.]]], + ], + &device, + ), + bias: TestTensor::from_data([4., 4., 4., 4., 4., 4.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv2d_complex() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 3, + kernel_size_1: 2, + kernel_size_2: 3, + padding_1: 1, + padding_2: 2, + stride_1: 1, + stride_2: 2, + dilation_1: 2, + dilation_2: 3, + groups: 1, + height: 4, + width: 5, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [36., 39., 0., 39., 42.], + [81., 87., 0., 87., 93.], + [81., 87., 0., 87., 93.], + [45., 48., 0., 48., 51.], + ], + [ + [54., 57., 0., 57., 60.], + [117., 123., 0., 123., 129.], + [117., 123., 0., 123., 129.], + [63., 66., 0., 66., 69.], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [[15., 42., 27.], [30., 72., 42.]], + [[75., 162., 87.], [90., 192., 102.]], + ], + [ + [[15., 42., 27.], [30., 72., 42.]], + [[75., 162., 87.], [90., 192., 102.]], + ], + [ + [[15., 42., 27.], [30., 72., 42.]], + [[75., 162., 87.], [90., 192., 102.]], + ], + ], + &device, + ), + bias: TestTensor::from_data([8., 8., 8.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv2d_groups_stride_2_no_pad() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 4, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + stride_1: 2, + stride_2: 2, + dilation_1: 1, + dilation_2: 1, + groups: 2, + height: 4, + width: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [0., 1., 2., 0.], + [3., 4., 5., 0.], + [6., 7., 8., 0.], + [0., 0., 0., 0.], + ], + [ + [9., 10., 11., 0.], + [12., 13., 14., 0.], + [15., 16., 17., 0.], + [0., 0., 0., 0.], + ], + [ + [18., 19., 20., 0.], + [21., 22., 23., 0.], + [24., 25., 26., 0.], + [0., 0., 0., 0.], + ], + [ + [27., 28., 29., 0.], + [30., 31., 32., 0.], + [33., 34., 35., 0.], + [0., 0., 0., 0.], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [[0., 1., 2.], [4., 5., 6.], [8., 9., 10.]], + [[16., 17., 18.], [20., 21., 22.], [24., 25., 26.]], + ], + [ + [[32., 33., 34.], [36., 37., 38.], [40., 41., 42.]], + [[48., 49., 50.], [52., 53., 54.], [56., 57., 58.]], + ], + ], + &device, + ), + bias: TestTensor::from_data([1., 1.], &device), + }; + test.assert_grads(grads); +} + +struct Conv2dTestCase { + batch_size: usize, + channels_in: usize, + channels_out: usize, + kernel_size_1: usize, + kernel_size_2: usize, + padding_1: usize, + padding_2: usize, + stride_1: usize, + stride_2: usize, + dilation_1: usize, + dilation_2: usize, + groups: usize, + height: usize, + width: usize, +} + +struct Grads { + x: TestTensor<4>, + weight: TestTensor<4>, + bias: TestTensor<1>, +} + +impl Conv2dTestCase { + fn assert_grads(self, expected_grads: Grads) { + let shape_x = Shape::new([self.batch_size, self.channels_in, self.height, self.width]); + let shape_weight = Shape::new([ + self.channels_out, + self.channels_in / self.groups, + self.kernel_size_1, + self.kernel_size_2, + ]); + let device = AutodiffDevice::new(); + let weight = TestTensor::from_data( + TestTensorInt::arange(0..shape_weight.num_elements() as i64, &device) + .reshape::<4, _>(shape_weight) + .into_data(), + &device, + ) + .require_grad(); + let bias = TestTensor::from_data( + TestTensorInt::arange(0..self.channels_out as i64, &device).into_data(), + &device, + ) + .require_grad(); + let x = TestTensor::from_data( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<4, _>(shape_x) + .into_data(), + &device, + ) + .require_grad(); + let output = conv2d( + x.clone(), + weight.clone(), + Some(bias.clone()), + ConvOptions::new( + [self.stride_1, self.stride_2], + [self.padding_1, self.padding_2], + [self.dilation_1, self.dilation_2], + self.groups, + ), + ); + let grads = output.backward(); + + // Assert + let x_grad_actual = x.grad(&grads).unwrap(); + let weight_grad_actual = weight.grad(&grads).unwrap(); + let bias_grad_actual = bias.grad(&grads).unwrap(); + + let tolerance = Tolerance::rel_abs(0.01, 0.01); + expected_grads + .bias + .to_data() + .assert_approx_eq::(&bias_grad_actual.to_data(), tolerance); + expected_grads + .x + .to_data() + .assert_approx_eq::(&x_grad_actual.to_data(), tolerance); + expected_grads + .weight + .to_data() + .assert_approx_eq::(&weight_grad_actual.to_data(), tolerance); + } +} diff --git a/crates/burn-backend-tests/tests/autodiff/conv3d.rs b/crates/burn-backend-tests/tests/autodiff/conv3d.rs new file mode 100644 index 0000000..da7d5ce --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/conv3d.rs @@ -0,0 +1,690 @@ +use super::*; +use burn_tensor::{Shape, Tolerance, module::conv3d, ops::ConvOptions}; + +#[test] +fn test_conv3d_basic() { + let test = Conv3dTestCase { + batch_size: 2, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + kernel_size_3: 3, + padding_1: 1, + padding_2: 1, + padding_3: 1, + stride_1: 1, + stride_2: 1, + stride_3: 1, + dilation_1: 1, + dilation_2: 1, + dilation_3: 1, + groups: 1, + depth: 4, + height: 4, + width: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [ + [ + [ + [536., 816., 816., 552.], + [840., 1278., 1278., 864.], + [840., 1278., 1278., 864.], + [584., 888., 888., 600.], + ], + [ + [912., 1386., 1386., 936.], + [1422., 2160., 2160., 1458.], + [1422., 2160., 2160., 1458.], + [984., 1494., 1494., 1008.], + ], + [ + [912., 1386., 1386., 936.], + [1422., 2160., 2160., 1458.], + [1422., 2160., 2160., 1458.], + [984., 1494., 1494., 1008.], + ], + [ + [680., 1032., 1032., 696.], + [1056., 1602., 1602., 1080.], + [1056., 1602., 1602., 1080.], + [728., 1104., 1104., 744.], + ], + ], + [ + [ + [968., 1464., 1464., 984.], + [1488., 2250., 2250., 1512.], + [1488., 2250., 2250., 1512.], + [1016., 1536., 1536., 1032.], + ], + [ + [1560., 2358., 2358., 1584.], + [2394., 3618., 3618., 2430.], + [2394., 3618., 3618., 2430.], + [1632., 2466., 2466., 1656.], + ], + [ + [1560., 2358., 2358., 1584.], + [2394., 3618., 3618., 2430.], + [2394., 3618., 3618., 2430.], + [1632., 2466., 2466., 1656.], + ], + [ + [1112., 1680., 1680., 1128.], + [1704., 2574., 2574., 1728.], + [1704., 2574., 2574., 1728.], + [1160., 1752., 1752., 1176.], + ], + ], + ], + [ + [ + [ + [536., 816., 816., 552.], + [840., 1278., 1278., 864.], + [840., 1278., 1278., 864.], + [584., 888., 888., 600.], + ], + [ + [912., 1386., 1386., 936.], + [1422., 2160., 2160., 1458.], + [1422., 2160., 2160., 1458.], + [984., 1494., 1494., 1008.], + ], + [ + [912., 1386., 1386., 936.], + [1422., 2160., 2160., 1458.], + [1422., 2160., 2160., 1458.], + [984., 1494., 1494., 1008.], + ], + [ + [680., 1032., 1032., 696.], + [1056., 1602., 1602., 1080.], + [1056., 1602., 1602., 1080.], + [728., 1104., 1104., 744.], + ], + ], + [ + [ + [968., 1464., 1464., 984.], + [1488., 2250., 2250., 1512.], + [1488., 2250., 2250., 1512.], + [1016., 1536., 1536., 1032.], + ], + [ + [1560., 2358., 2358., 1584.], + [2394., 3618., 3618., 2430.], + [2394., 3618., 3618., 2430.], + [1632., 2466., 2466., 1656.], + ], + [ + [1560., 2358., 2358., 1584.], + [2394., 3618., 3618., 2430.], + [2394., 3618., 3618., 2430.], + [1632., 2466., 2466., 1656.], + ], + [ + [1112., 1680., 1680., 1128.], + [1704., 2574., 2574., 1728.], + [1704., 2574., 2574., 1728.], + [1160., 1752., 1752., 1176.], + ], + ], + ], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [ + [ + [4590., 6156., 4644.], + [6264., 8400., 6336.], + [4806., 6444., 4860.], + ], + [ + [6696., 8976., 6768.], + [9120., 12224., 9216.], + [6984., 9360., 7056.], + ], + [ + [5454., 7308., 5508.], + [7416., 9936., 7488.], + [5670., 7596., 5724.], + ], + ], + [ + [ + [8046., 10764., 8100.], + [10872., 14544., 10944.], + [8262., 11052., 8316.], + ], + [ + [11304., 15120., 11376.], + [15264., 20416., 15360.], + [11592., 15504., 11664.], + ], + [ + [8910., 11916., 8964.], + [12024., 16080., 12096.], + [9126., 12204., 9180.], + ], + ], + ], + [ + [ + [ + [4590., 6156., 4644.], + [6264., 8400., 6336.], + [4806., 6444., 4860.], + ], + [ + [6696., 8976., 6768.], + [9120., 12224., 9216.], + [6984., 9360., 7056.], + ], + [ + [5454., 7308., 5508.], + [7416., 9936., 7488.], + [5670., 7596., 5724.], + ], + ], + [ + [ + [8046., 10764., 8100.], + [10872., 14544., 10944.], + [8262., 11052., 8316.], + ], + [ + [11304., 15120., 11376.], + [15264., 20416., 15360.], + [11592., 15504., 11664.], + ], + [ + [8910., 11916., 8964.], + [12024., 16080., 12096.], + [9126., 12204., 9180.], + ], + ], + ], + ], + &device, + ), + bias: TestTensor::from_data([128., 128.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv3d_complex() { + let test = Conv3dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 3, + kernel_size_1: 2, + kernel_size_2: 3, + kernel_size_3: 4, + padding_1: 1, + padding_2: 2, + padding_3: 3, + stride_1: 1, + stride_2: 2, + stride_3: 3, + dilation_1: 2, + dilation_2: 3, + dilation_3: 4, + groups: 1, + depth: 5, + height: 6, + width: 7, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [ + [0., 147., 0., 0., 0., 150., 0.], + [0., 159., 0., 0., 0., 162., 0.], + [0., 0., 0., 0., 0., 0., 0.], + [0., 159., 0., 0., 0., 162., 0.], + [0., 171., 0., 0., 0., 174., 0.], + [0., 0., 0., 0., 0., 0., 0.], + ], + [ + [0., 330., 0., 0., 0., 336., 0.], + [0., 354., 0., 0., 0., 360., 0.], + [0., 0., 0., 0., 0., 0., 0.], + [0., 354., 0., 0., 0., 360., 0.], + [0., 378., 0., 0., 0., 384., 0.], + [0., 0., 0., 0., 0., 0., 0.], + ], + [ + [0., 330., 0., 0., 0., 336., 0.], + [0., 354., 0., 0., 0., 360., 0.], + [0., 0., 0., 0., 0., 0., 0.], + [0., 354., 0., 0., 0., 360., 0.], + [0., 378., 0., 0., 0., 384., 0.], + [0., 0., 0., 0., 0., 0., 0.], + ], + [ + [0., 330., 0., 0., 0., 336., 0.], + [0., 354., 0., 0., 0., 360., 0.], + [0., 0., 0., 0., 0., 0., 0.], + [0., 354., 0., 0., 0., 360., 0.], + [0., 378., 0., 0., 0., 384., 0.], + [0., 0., 0., 0., 0., 0., 0.], + ], + [ + [0., 183., 0., 0., 0., 186., 0.], + [0., 195., 0., 0., 0., 198., 0.], + [0., 0., 0., 0., 0., 0., 0.], + [0., 195., 0., 0., 0., 198., 0.], + [0., 207., 0., 0., 0., 210., 0.], + [0., 0., 0., 0., 0., 0., 0.], + ], + ], + [ + [ + [0., 219., 0., 0., 0., 222., 0.], + [0., 231., 0., 0., 0., 234., 0.], + [0., 0., 0., 0., 0., 0., 0.], + [0., 231., 0., 0., 0., 234., 0.], + [0., 243., 0., 0., 0., 246., 0.], + [0., 0., 0., 0., 0., 0., 0.], + ], + [ + [0., 474., 0., 0., 0., 480., 0.], + [0., 498., 0., 0., 0., 504., 0.], + [0., 0., 0., 0., 0., 0., 0.], + [0., 498., 0., 0., 0., 504., 0.], + [0., 522., 0., 0., 0., 528., 0.], + [0., 0., 0., 0., 0., 0., 0.], + ], + [ + [0., 474., 0., 0., 0., 480., 0.], + [0., 498., 0., 0., 0., 504., 0.], + [0., 0., 0., 0., 0., 0., 0.], + [0., 498., 0., 0., 0., 504., 0.], + [0., 522., 0., 0., 0., 528., 0.], + [0., 0., 0., 0., 0., 0., 0.], + ], + [ + [0., 474., 0., 0., 0., 480., 0.], + [0., 498., 0., 0., 0., 504., 0.], + [0., 0., 0., 0., 0., 0., 0.], + [0., 498., 0., 0., 0., 504., 0.], + [0., 522., 0., 0., 0., 528., 0.], + [0., 0., 0., 0., 0., 0., 0.], + ], + [ + [0., 255., 0., 0., 0., 258., 0.], + [0., 267., 0., 0., 0., 270., 0.], + [0., 0., 0., 0., 0., 0., 0.], + [0., 267., 0., 0., 0., 270., 0.], + [0., 279., 0., 0., 0., 282., 0.], + [0., 0., 0., 0., 0., 0., 0.], + ], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [ + [ + [0., 256., 272., 0.], + [0., 624., 656., 0.], + [0., 368., 384., 0.], + ], + [ + [0., 424., 440., 0.], + [0., 960., 992., 0.], + [0., 536., 552., 0.], + ], + ], + [ + [ + [0., 1096., 1112., 0.], + [0., 2304., 2336., 0.], + [0., 1208., 1224., 0.], + ], + [ + [0., 1264., 1280., 0.], + [0., 2640., 2672., 0.], + [0., 1376., 1392., 0.], + ], + ], + ], + [ + [ + [ + [0., 256., 272., 0.], + [0., 624., 656., 0.], + [0., 368., 384., 0.], + ], + [ + [0., 424., 440., 0.], + [0., 960., 992., 0.], + [0., 536., 552., 0.], + ], + ], + [ + [ + [0., 1096., 1112., 0.], + [0., 2304., 2336., 0.], + [0., 1208., 1224., 0.], + ], + [ + [0., 1264., 1280., 0.], + [0., 2640., 2672., 0.], + [0., 1376., 1392., 0.], + ], + ], + ], + [ + [ + [ + [0., 256., 272., 0.], + [0., 624., 656., 0.], + [0., 368., 384., 0.], + ], + [ + [0., 424., 440., 0.], + [0., 960., 992., 0.], + [0., 536., 552., 0.], + ], + ], + [ + [ + [0., 1096., 1112., 0.], + [0., 2304., 2336., 0.], + [0., 1208., 1224., 0.], + ], + [ + [0., 1264., 1280., 0.], + [0., 2640., 2672., 0.], + [0., 1376., 1392., 0.], + ], + ], + ], + ], + &device, + ), + bias: TestTensor::from_data([10., 10., 10.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv3d_groups_stride_2_no_pad() { + let test = Conv3dTestCase { + batch_size: 1, + channels_in: 4, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + kernel_size_3: 3, + padding_1: 0, + padding_2: 0, + padding_3: 0, + stride_1: 2, + stride_2: 2, + stride_3: 2, + dilation_1: 1, + dilation_2: 1, + dilation_3: 1, + groups: 2, + depth: 4, + height: 4, + width: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [ + [0., 1., 2., 0.], + [3., 4., 5., 0.], + [6., 7., 8., 0.], + [0., 0., 0., 0.], + ], + [ + [9., 10., 11., 0.], + [12., 13., 14., 0.], + [15., 16., 17., 0.], + [0., 0., 0., 0.], + ], + [ + [18., 19., 20., 0.], + [21., 22., 23., 0.], + [24., 25., 26., 0.], + [0., 0., 0., 0.], + ], + [ + [0., 0., 0., 0.], + [0., 0., 0., 0.], + [0., 0., 0., 0.], + [0., 0., 0., 0.], + ], + ], + [ + [ + [27., 28., 29., 0.], + [30., 31., 32., 0.], + [33., 34., 35., 0.], + [0., 0., 0., 0.], + ], + [ + [36., 37., 38., 0.], + [39., 40., 41., 0.], + [42., 43., 44., 0.], + [0., 0., 0., 0.], + ], + [ + [45., 46., 47., 0.], + [48., 49., 50., 0.], + [51., 52., 53., 0.], + [0., 0., 0., 0.], + ], + [ + [0., 0., 0., 0.], + [0., 0., 0., 0.], + [0., 0., 0., 0.], + [0., 0., 0., 0.], + ], + ], + [ + [ + [54., 55., 56., 0.], + [57., 58., 59., 0.], + [60., 61., 62., 0.], + [0., 0., 0., 0.], + ], + [ + [63., 64., 65., 0.], + [66., 67., 68., 0.], + [69., 70., 71., 0.], + [0., 0., 0., 0.], + ], + [ + [72., 73., 74., 0.], + [75., 76., 77., 0.], + [78., 79., 80., 0.], + [0., 0., 0., 0.], + ], + [ + [0., 0., 0., 0.], + [0., 0., 0., 0.], + [0., 0., 0., 0.], + [0., 0., 0., 0.], + ], + ], + [ + [ + [81., 82., 83., 0.], + [84., 85., 86., 0.], + [87., 88., 89., 0.], + [0., 0., 0., 0.], + ], + [ + [90., 91., 92., 0.], + [93., 94., 95., 0.], + [96., 97., 98., 0.], + [0., 0., 0., 0.], + ], + [ + [99., 100., 101., 0.], + [102., 103., 104., 0.], + [105., 106., 107., 0.], + [0., 0., 0., 0.], + ], + [ + [0., 0., 0., 0.], + [0., 0., 0., 0.], + [0., 0., 0., 0.], + [0., 0., 0., 0.], + ], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [ + [[0., 1., 2.], [4., 5., 6.], [8., 9., 10.]], + [[16., 17., 18.], [20., 21., 22.], [24., 25., 26.]], + [[32., 33., 34.], [36., 37., 38.], [40., 41., 42.]], + ], + [ + [[64., 65., 66.], [68., 69., 70.], [72., 73., 74.]], + [[80., 81., 82.], [84., 85., 86.], [88., 89., 90.]], + [[96., 97., 98.], [100., 101., 102.], [104., 105., 106.]], + ], + ], + [ + [ + [[128., 129., 130.], [132., 133., 134.], [136., 137., 138.]], + [[144., 145., 146.], [148., 149., 150.], [152., 153., 154.]], + [[160., 161., 162.], [164., 165., 166.], [168., 169., 170.]], + ], + [ + [[192., 193., 194.], [196., 197., 198.], [200., 201., 202.]], + [[208., 209., 210.], [212., 213., 214.], [216., 217., 218.]], + [[224., 225., 226.], [228., 229., 230.], [232., 233., 234.]], + ], + ], + ], + &device, + ), + bias: TestTensor::from_data([1., 1.], &device), + }; + test.assert_grads(grads); +} + +struct Conv3dTestCase { + batch_size: usize, + channels_in: usize, + channels_out: usize, + kernel_size_1: usize, + kernel_size_2: usize, + kernel_size_3: usize, + padding_1: usize, + padding_2: usize, + padding_3: usize, + stride_1: usize, + stride_2: usize, + stride_3: usize, + dilation_1: usize, + dilation_2: usize, + dilation_3: usize, + groups: usize, + depth: usize, + height: usize, + width: usize, +} + +struct Grads { + x: TestTensor<5>, + weight: TestTensor<5>, + bias: TestTensor<1>, +} + +impl Conv3dTestCase { + fn assert_grads(self, expected_grads: Grads) { + let shape_x = Shape::new([ + self.batch_size, + self.channels_in, + self.depth, + self.height, + self.width, + ]); + let shape_weight = Shape::new([ + self.channels_out, + self.channels_in / self.groups, + self.kernel_size_1, + self.kernel_size_2, + self.kernel_size_3, + ]); + let device = AutodiffDevice::new(); + let weight = TestTensor::from_data( + TestTensorInt::arange(0..shape_weight.num_elements() as i64, &device) + .reshape::<5, _>(shape_weight) + .into_data(), + &device, + ) + .require_grad(); + let bias = TestTensor::from_data( + TestTensorInt::arange(0..self.channels_out as i64, &device).into_data(), + &device, + ) + .require_grad(); + let x = TestTensor::from_data( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<5, _>(shape_x) + .into_data(), + &device, + ) + .require_grad(); + let output = conv3d( + x.clone(), + weight.clone(), + Some(bias.clone()), + ConvOptions::new( + [self.stride_1, self.stride_2, self.stride_3], + [self.padding_1, self.padding_2, self.padding_3], + [self.dilation_1, self.dilation_2, self.dilation_3], + self.groups, + ), + ); + let grads = output.backward(); + + // Assert + let x_grad_actual = x.grad(&grads).unwrap(); + let weight_grad_actual = weight.grad(&grads).unwrap(); + let bias_grad_actual = bias.grad(&grads).unwrap(); + + let tolerance = Tolerance::default(); + expected_grads + .bias + .to_data() + .assert_approx_eq::(&bias_grad_actual.to_data(), tolerance); + expected_grads + .x + .to_data() + .assert_approx_eq::(&x_grad_actual.to_data(), tolerance); + expected_grads + .weight + .to_data() + .assert_approx_eq::(&weight_grad_actual.to_data(), tolerance); + } +} diff --git a/crates/burn-backend-tests/tests/autodiff/conv_transpose1d.rs b/crates/burn-backend-tests/tests/autodiff/conv_transpose1d.rs new file mode 100644 index 0000000..f66ec3b --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/conv_transpose1d.rs @@ -0,0 +1,338 @@ +use super::*; +use burn_tensor::{Shape, Tolerance, module::conv_transpose1d, ops::ConvTransposeOptions}; + +#[test] +fn test_conv_transpose1d_basic() { + let test = ConvTranspose1dTestCase { + batch_size: 2, + channels: [2, 2], + kernel_size: 3, + padding: 0, + padding_out: 0, + stride: 1, + dilation: 1, + groups: 1, + size: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [[15.0, 15.0, 15.0, 15.0], [51.0, 51.0, 51.0, 51.0]], + [[15.0, 15.0, 15.0, 15.0], [51.0, 51.0, 51.0, 51.0]], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [[44.0, 44.0, 44.0], [44.0, 44.0, 44.0]], + [[76.0, 76.0, 76.0], [76.0, 76.0, 76.0]], + ], + &device, + ), + bias: TestTensor::from_data([12., 12.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv_transpose1d_padding() { + let test = ConvTranspose1dTestCase { + batch_size: 2, + channels: [2, 2], + kernel_size: 3, + padding: 2, + padding_out: 0, + stride: 1, + dilation: 1, + groups: 1, + size: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [[7., 12., 8., 3.], [19., 36., 32., 15.]], + [[7., 12., 8., 3.], [19., 36., 32., 15.]], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [[26., 22., 18.], [26., 22., 18.]], + [[42., 38., 34.], [42., 38., 34.]], + ], + &device, + ), + bias: TestTensor::from_data([4., 4.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv_transpose1d_stride() { + let test = ConvTranspose1dTestCase { + batch_size: 2, + channels: [2, 2], + kernel_size: 3, + padding: 0, + padding_out: 0, + stride: 2, + dilation: 1, + groups: 1, + size: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [[15., 15., 15., 15.], [51., 51., 51., 51.]], + [[15., 15., 15., 15.], [51., 51., 51., 51.]], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [[44., 44., 44.], [44., 44., 44.]], + [[76., 76., 76.], [76., 76., 76.]], + ], + &device, + ), + bias: TestTensor::from_data([18., 18.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv_transpose1d_stride_padding_out() { + let test = ConvTranspose1dTestCase { + batch_size: 2, + channels: [2, 2], + kernel_size: 3, + padding: 0, + padding_out: 1, + stride: 2, + dilation: 1, + groups: 1, + size: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [[15., 15., 15., 15.], [51., 51., 51., 51.]], + [[15., 15., 15., 15.], [51., 51., 51., 51.]], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [[44., 44., 44.], [44., 44., 44.]], + [[76., 76., 76.], [76., 76., 76.]], + ], + &device, + ), + bias: TestTensor::from_data([20., 20.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv_transpose1d_dilation() { + let test = ConvTranspose1dTestCase { + batch_size: 2, + channels: [2, 2], + kernel_size: 3, + padding: 0, + padding_out: 0, + stride: 1, + dilation: 2, + groups: 1, + size: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [[15., 15., 15., 15.], [51., 51., 51., 51.]], + [[15., 15., 15., 15.], [51., 51., 51., 51.]], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [[44., 44., 44.], [44., 44., 44.]], + [[76., 76., 76.], [76., 76., 76.]], + ], + &device, + ), + bias: TestTensor::from_data([16., 16.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv_transpose1d_complex() { + let test = ConvTranspose1dTestCase { + batch_size: 2, + channels: [2, 4], + kernel_size: 3, + padding: 1, + padding_out: 1, + stride: 2, + dilation: 2, + groups: 2, + size: 8, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [ + [12.0, 15.0, 15.0, 15.0, 15.0, 15.0, 15.0, 15.0], + [36.0, 51.0, 51.0, 51.0, 51.0, 51.0, 51.0, 51.0], + ], + [ + [12.0, 15.0, 15.0, 15.0, 15.0, 15.0, 15.0, 15.0], + [36.0, 51.0, 51.0, 51.0, 51.0, 51.0, 51.0, 51.0], + ], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [[168.0, 184.0, 184.0], [168.0, 184.0, 184.0]], + [[280.0, 312.0, 312.0], [280.0, 312.0, 312.0]], + ], + &device, + ), + bias: TestTensor::from_data([36.0, 36.0, 36.0, 36.0], &device), + }; + test.assert_grads(grads); +} + +/// Regression test for #4845. +/// +/// `ConvTranspose1d` with `padding_out != 0` and `stride == 1` used to panic in +/// the backward pass because `conv_transpose1d_x_backward` did not account for +/// the trailing `padding_out` cells, producing a gradient longer than `x`. +#[test] +fn test_conv_transpose1d_padding_out_stride1_backward_shape() { + let device = AutodiffDevice::new(); + let batch_size = 2; + let channels_in = 2; + let channels_out = 2; + let kernel_size = 3; + let size_in = 4; + let padding_out = 1; + + let shape_x = Shape::new([batch_size, channels_in, size_in]); + let shape_weight = Shape::new([channels_in, channels_out, kernel_size]); + let weight = TestTensor::from_data( + TestTensorInt::arange(0..shape_weight.num_elements() as i64, &device) + .reshape::<3, _>(shape_weight.clone()) + .into_data(), + &device, + ) + .require_grad(); + let x = TestTensor::from_data( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<3, _>(shape_x.clone()) + .into_data(), + &device, + ) + .require_grad(); + + let output = conv_transpose1d( + x.clone(), + weight.clone(), + None, + ConvTransposeOptions::new([1], [0], [padding_out], [1], 1), + ); + let grads = output.backward(); + + let x_grad = x.grad(&grads).unwrap(); + let weight_grad = weight.grad(&grads).unwrap(); + assert_eq!(x_grad.shape(), shape_x); + assert_eq!(weight_grad.shape(), shape_weight); +} + +struct ConvTranspose1dTestCase { + batch_size: usize, + channels: [usize; 2], + kernel_size: usize, + padding: usize, + padding_out: usize, + stride: usize, + dilation: usize, + groups: usize, + size: usize, +} + +struct Grads { + x: TestTensor<3>, + weight: TestTensor<3>, + bias: TestTensor<1>, +} + +impl ConvTranspose1dTestCase { + fn assert_grads(self, expected_grads: Grads) { + let shape_x = Shape::new([self.batch_size, self.channels[0], self.size]); + let shape_weight = Shape::new([ + self.channels[0], + self.channels[1] / self.groups, + self.kernel_size, + ]); + let device = AutodiffDevice::new(); + let weight = TestTensor::from_data( + TestTensorInt::arange(0..shape_weight.num_elements() as i64, &device) + .reshape::<3, _>(shape_weight) + .into_data(), + &device, + ) + .require_grad(); + let bias = TestTensor::from_data( + TestTensorInt::arange(0..self.channels[1] as i64, &device).into_data(), + &device, + ) + .require_grad(); + let x = TestTensor::from_data( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<3, _>(shape_x) + .into_data(), + &device, + ) + .require_grad(); + let output = conv_transpose1d( + x.clone(), + weight.clone(), + Some(bias.clone()), + ConvTransposeOptions::new( + [self.stride], + [self.padding], + [self.padding_out], + [self.dilation], + self.groups, + ), + ); + let grads = output.backward(); + + // Assert + let x_grad_actual = x.grad(&grads).unwrap(); + let weight_grad_actual = weight.grad(&grads).unwrap(); + let bias_grad_actual = bias.grad(&grads).unwrap(); + + expected_grads + .bias + .to_data() + .assert_approx_eq::(&bias_grad_actual.to_data(), Tolerance::default()); + expected_grads + .x + .to_data() + .assert_approx_eq::(&x_grad_actual.to_data(), Tolerance::default()); + expected_grads + .weight + .to_data() + .assert_approx_eq::(&weight_grad_actual.to_data(), Tolerance::default()); + } +} diff --git a/crates/burn-backend-tests/tests/autodiff/conv_transpose2d.rs b/crates/burn-backend-tests/tests/autodiff/conv_transpose2d.rs new file mode 100644 index 0000000..a1a648a --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/conv_transpose2d.rs @@ -0,0 +1,706 @@ +use super::*; +use burn_tensor::{Shape, Tolerance, module::conv_transpose2d, ops::ConvTransposeOptions}; + +#[test] +fn test_conv_transpose2d_basic() { + let test = ConvTranspose2dTestCase { + batch_size: 2, + channels: [2, 2], + kernel_size: [3, 3], + padding: [0, 0], + padding_out: [0, 0], + stride: [1, 1], + dilation: [1, 1], + groups: 1, + size: [4, 4], + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [ + [ + [153., 153., 153., 153.], + [153., 153., 153., 153.], + [153., 153., 153., 153.], + [153., 153., 153., 153.], + ], + [ + [477., 477., 477., 477.], + [477., 477., 477., 477.], + [477., 477., 477., 477.], + [477., 477., 477., 477.], + ], + ], + [ + [ + [153., 153., 153., 153.], + [153., 153., 153., 153.], + [153., 153., 153., 153.], + [153., 153., 153., 153.], + ], + [ + [477., 477., 477., 477.], + [477., 477., 477., 477.], + [477., 477., 477., 477.], + [477., 477., 477., 477.], + ], + ], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [[752., 752., 752.], [752., 752., 752.], [752., 752., 752.]], + [[752., 752., 752.], [752., 752., 752.], [752., 752., 752.]], + ], + [ + [ + [1264., 1264., 1264.], + [1264., 1264., 1264.], + [1264., 1264., 1264.], + ], + [ + [1264., 1264., 1264.], + [1264., 1264., 1264.], + [1264., 1264., 1264.], + ], + ], + ], + &device, + ), + bias: TestTensor::from_data([72., 72.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv_transpose2d_padding() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels: [1, 1], + kernel_size: [3, 3], + padding: [1, 2], + padding_out: [0, 0], + stride: [1, 1], + dilation: [1, 1], + groups: 1, + size: [4, 4], + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[[ + [13., 24., 20., 9.], + [15., 27., 21., 9.], + [15., 27., 21., 9.], + [7., 12., 8., 3.], + ]]], + &device, + ), + weight: TestTensor::from_data( + [[[[63., 57., 51.], [68., 60., 52.], [39., 33., 27.]]]], + &device, + ), + bias: TestTensor::from_data([8.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv_transpose2d_stride() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels: [1, 1], + kernel_size: [3, 3], + padding: [0, 0], + padding_out: [0, 0], + stride: [2, 3], + dilation: [1, 1], + groups: 1, + size: [4, 4], + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[[ + [36., 36., 36., 36.], + [36., 36., 36., 36.], + [36., 36., 36., 36.], + [36., 36., 36., 36.], + ]]], + &device, + ), + weight: TestTensor::from_data( + [[[[120., 120., 120.], [120., 120., 120.], [120., 120., 120.]]]], + &device, + ), + bias: TestTensor::from_data([108.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv_transpose2d_stride_padding_out() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels: [1, 1], + kernel_size: [3, 3], + padding: [0, 0], + padding_out: [1, 2], + stride: [2, 3], + dilation: [1, 1], + groups: 1, + size: [4, 4], + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[[ + [36., 36., 36., 36.], + [36., 36., 36., 36.], + [36., 36., 36., 36.], + [36., 36., 36., 36.], + ]]], + &device, + ), + weight: TestTensor::from_data( + [[[[120., 120., 120.], [120., 120., 120.], [120., 120., 120.]]]], + &device, + ), + bias: TestTensor::from_data([140.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv_transpose2d_dilation() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels: [1, 1], + kernel_size: [3, 3], + padding: [0, 0], + padding_out: [0, 0], + stride: [1, 1], + dilation: [2, 3], + groups: 1, + size: [4, 4], + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[[ + [36., 36., 36., 36.], + [36., 36., 36., 36.], + [36., 36., 36., 36.], + [36., 36., 36., 36.], + ]]], + &device, + ), + weight: TestTensor::from_data( + [[[[120., 120., 120.], [120., 120., 120.], [120., 120., 120.]]]], + &device, + ), + bias: TestTensor::from_data([80.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv_transpose2d_channels() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels: [2, 3], + kernel_size: [3, 3], + padding: [0, 0], + padding_out: [0, 0], + stride: [1, 1], + dilation: [1, 1], + groups: 1, + size: [4, 4], + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [351., 351., 351., 351.], + [351., 351., 351., 351.], + [351., 351., 351., 351.], + [351., 351., 351., 351.], + ], + [ + [1080., 1080., 1080., 1080.], + [1080., 1080., 1080., 1080.], + [1080., 1080., 1080., 1080.], + [1080., 1080., 1080., 1080.], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [[120., 120., 120.], [120., 120., 120.], [120., 120., 120.]], + [[120., 120., 120.], [120., 120., 120.], [120., 120., 120.]], + [[120., 120., 120.], [120., 120., 120.], [120., 120., 120.]], + ], + [ + [[376., 376., 376.], [376., 376., 376.], [376., 376., 376.]], + [[376., 376., 376.], [376., 376., 376.], [376., 376., 376.]], + [[376., 376., 376.], [376., 376., 376.], [376., 376., 376.]], + ], + ], + &device, + ), + bias: TestTensor::from_data([36., 36., 36.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv_transpose2d_kernel_size() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels: [1, 1], + kernel_size: [3, 5], + padding: [0, 0], + padding_out: [0, 0], + stride: [1, 1], + dilation: [1, 1], + groups: 1, + size: [6, 6], + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[[ + [105., 105., 105., 105., 105., 105.], + [105., 105., 105., 105., 105., 105.], + [105., 105., 105., 105., 105., 105.], + [105., 105., 105., 105., 105., 105.], + [105., 105., 105., 105., 105., 105.], + [105., 105., 105., 105., 105., 105.], + ]]], + &device, + ), + weight: TestTensor::from_data( + [[[ + [630., 630., 630., 630., 630.], + [630., 630., 630., 630., 630.], + [630., 630., 630., 630., 630.], + ]]], + &device, + ), + bias: TestTensor::from_data([80.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv_transpose2d_groups() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels: [2, 2], + kernel_size: [3, 3], + padding: [0, 0], + padding_out: [0, 0], + stride: [1, 1], + dilation: [1, 1], + groups: 2, + size: [4, 4], + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [36., 36., 36., 36.], + [36., 36., 36., 36.], + [36., 36., 36., 36.], + [36., 36., 36., 36.], + ], + [ + [117., 117., 117., 117.], + [117., 117., 117., 117.], + [117., 117., 117., 117.], + [117., 117., 117., 117.], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [[[120., 120., 120.], [120., 120., 120.], [120., 120., 120.]]], + [[[376., 376., 376.], [376., 376., 376.], [376., 376., 376.]]], + ], + &device, + ), + bias: TestTensor::from_data([36., 36.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv_transpose2d_complex_no_groups() { + let test = ConvTranspose2dTestCase { + batch_size: 2, + channels: [2, 3], + kernel_size: [3, 5], + padding: [1, 2], + padding_out: [1, 2], + stride: [2, 3], + dilation: [2, 3], + groups: 1, + size: [6, 8], + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [ + [ + [600., 735., 735., 735., 735., 735., 735., 735.], + [810., 990., 990., 990., 990., 990., 990., 990.], + [810., 990., 990., 990., 990., 990., 990., 990.], + [810., 990., 990., 990., 990., 990., 990., 990.], + [810., 990., 990., 990., 990., 990., 990., 990.], + [810., 990., 990., 990., 990., 990., 990., 990.], + ], + [ + [1680., 2085., 2085., 2085., 2085., 2085., 2085., 2085.], + [2430., 3015., 3015., 3015., 3015., 3015., 3015., 3015.], + [2430., 3015., 3015., 3015., 3015., 3015., 3015., 3015.], + [2430., 3015., 3015., 3015., 3015., 3015., 3015., 3015.], + [2430., 3015., 3015., 3015., 3015., 3015., 3015., 3015.], + [2430., 3015., 3015., 3015., 3015., 3015., 3015., 3015.], + ], + ], + [ + [ + [600., 735., 735., 735., 735., 735., 735., 735.], + [810., 990., 990., 990., 990., 990., 990., 990.], + [810., 990., 990., 990., 990., 990., 990., 990.], + [810., 990., 990., 990., 990., 990., 990., 990.], + [810., 990., 990., 990., 990., 990., 990., 990.], + [810., 990., 990., 990., 990., 990., 990., 990.], + ], + [ + [1680., 2085., 2085., 2085., 2085., 2085., 2085., 2085.], + [2430., 3015., 3015., 3015., 3015., 3015., 3015., 3015.], + [2430., 3015., 3015., 3015., 3015., 3015., 3015., 3015.], + [2430., 3015., 3015., 3015., 3015., 3015., 3015., 3015.], + [2430., 3015., 3015., 3015., 3015., 3015., 3015., 3015.], + [2430., 3015., 3015., 3015., 3015., 3015., 3015., 3015.], + ], + ], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [ + [5320., 6040., 6040., 6040., 6040.], + [6048., 6864., 6864., 6864., 6864.], + [6048., 6864., 6864., 6864., 6864.], + ], + [ + [5320., 6040., 6040., 6040., 6040.], + [6048., 6864., 6864., 6864., 6864.], + [6048., 6864., 6864., 6864., 6864.], + ], + [ + [5320., 6040., 6040., 6040., 6040.], + [6048., 6864., 6864., 6864., 6864.], + [6048., 6864., 6864., 6864., 6864.], + ], + ], + [ + [ + [8680., 9880., 9880., 9880., 9880.], + [10080., 11472., 11472., 11472., 11472.], + [10080., 11472., 11472., 11472., 11472.], + ], + [ + [8680., 9880., 9880., 9880., 9880.], + [10080., 11472., 11472., 11472., 11472.], + [10080., 11472., 11472., 11472., 11472.], + ], + [ + [8680., 9880., 9880., 9880., 9880.], + [10080., 11472., 11472., 11472., 11472.], + [10080., 11472., 11472., 11472., 11472.], + ], + ], + ], + &device, + ), + bias: TestTensor::from_data([896., 896., 896.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv_transpose2d_complex_no_groups_2() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels: [4, 2], + kernel_size: [2, 3], + padding: [1, 2], + padding_out: [1, 2], + stride: [2, 3], + dilation: [1, 2], + groups: 1, + size: [10, 10], + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [30., 42., 42., 42., 42., 42., 42., 42., 42., 42.], + [48., 66., 66., 66., 66., 66., 66., 66., 66., 66.], + [48., 66., 66., 66., 66., 66., 66., 66., 66., 66.], + [48., 66., 66., 66., 66., 66., 66., 66., 66., 66.], + [48., 66., 66., 66., 66., 66., 66., 66., 66., 66.], + [48., 66., 66., 66., 66., 66., 66., 66., 66., 66.], + [48., 66., 66., 66., 66., 66., 66., 66., 66., 66.], + [48., 66., 66., 66., 66., 66., 66., 66., 66., 66.], + [48., 66., 66., 66., 66., 66., 66., 66., 66., 66.], + [48., 66., 66., 66., 66., 66., 66., 66., 66., 66.], + ], + [ + [78., 114., 114., 114., 114., 114., 114., 114., 114., 114.], + [144., 210., 210., 210., 210., 210., 210., 210., 210., 210.], + [144., 210., 210., 210., 210., 210., 210., 210., 210., 210.], + [144., 210., 210., 210., 210., 210., 210., 210., 210., 210.], + [144., 210., 210., 210., 210., 210., 210., 210., 210., 210.], + [144., 210., 210., 210., 210., 210., 210., 210., 210., 210.], + [144., 210., 210., 210., 210., 210., 210., 210., 210., 210.], + [144., 210., 210., 210., 210., 210., 210., 210., 210., 210.], + [144., 210., 210., 210., 210., 210., 210., 210., 210., 210.], + [144., 210., 210., 210., 210., 210., 210., 210., 210., 210.], + ], + [ + [126., 186., 186., 186., 186., 186., 186., 186., 186., 186.], + [240., 354., 354., 354., 354., 354., 354., 354., 354., 354.], + [240., 354., 354., 354., 354., 354., 354., 354., 354., 354.], + [240., 354., 354., 354., 354., 354., 354., 354., 354., 354.], + [240., 354., 354., 354., 354., 354., 354., 354., 354., 354.], + [240., 354., 354., 354., 354., 354., 354., 354., 354., 354.], + [240., 354., 354., 354., 354., 354., 354., 354., 354., 354.], + [240., 354., 354., 354., 354., 354., 354., 354., 354., 354.], + [240., 354., 354., 354., 354., 354., 354., 354., 354., 354.], + [240., 354., 354., 354., 354., 354., 354., 354., 354., 354.], + ], + [ + [174., 258., 258., 258., 258., 258., 258., 258., 258., 258.], + [336., 498., 498., 498., 498., 498., 498., 498., 498., 498.], + [336., 498., 498., 498., 498., 498., 498., 498., 498., 498.], + [336., 498., 498., 498., 498., 498., 498., 498., 498., 498.], + [336., 498., 498., 498., 498., 498., 498., 498., 498., 498.], + [336., 498., 498., 498., 498., 498., 498., 498., 498., 498.], + [336., 498., 498., 498., 498., 498., 498., 498., 498., 498.], + [336., 498., 498., 498., 498., 498., 498., 498., 498., 498.], + [336., 498., 498., 498., 498., 498., 498., 498., 498., 498.], + [336., 498., 498., 498., 498., 498., 498., 498., 498., 498.], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [[4455., 4905., 4905.], [4500., 4950., 4950.]], + [[4455., 4905., 4905.], [4500., 4950., 4950.]], + ], + [ + [[12555., 13905., 13905.], [13500., 14950., 14950.]], + [[12555., 13905., 13905.], [13500., 14950., 14950.]], + ], + [ + [[20655., 22905., 22905.], [22500., 24950., 24950.]], + [[20655., 22905., 22905.], [22500., 24950., 24950.]], + ], + [ + [[28755., 31905., 31905.], [31500., 34950., 34950.]], + [[28755., 31905., 31905.], [31500., 34950., 34950.]], + ], + ], + &device, + ), + bias: TestTensor::from_data([570., 570.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv_transpose2d_complex_groups() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels: [4, 2], + kernel_size: [2, 3], + padding: [1, 2], + padding_out: [1, 2], + stride: [2, 3], + dilation: [1, 2], + groups: 2, + size: [10, 10], + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [9., 12., 12., 12., 12., 12., 12., 12., 12., 12.], + [12., 15., 15., 15., 15., 15., 15., 15., 15., 15.], + [12., 15., 15., 15., 15., 15., 15., 15., 15., 15.], + [12., 15., 15., 15., 15., 15., 15., 15., 15., 15.], + [12., 15., 15., 15., 15., 15., 15., 15., 15., 15.], + [12., 15., 15., 15., 15., 15., 15., 15., 15., 15.], + [12., 15., 15., 15., 15., 15., 15., 15., 15., 15.], + [12., 15., 15., 15., 15., 15., 15., 15., 15., 15.], + [12., 15., 15., 15., 15., 15., 15., 15., 15., 15.], + [12., 15., 15., 15., 15., 15., 15., 15., 15., 15.], + ], + [ + [21., 30., 30., 30., 30., 30., 30., 30., 30., 30.], + [36., 51., 51., 51., 51., 51., 51., 51., 51., 51.], + [36., 51., 51., 51., 51., 51., 51., 51., 51., 51.], + [36., 51., 51., 51., 51., 51., 51., 51., 51., 51.], + [36., 51., 51., 51., 51., 51., 51., 51., 51., 51.], + [36., 51., 51., 51., 51., 51., 51., 51., 51., 51.], + [36., 51., 51., 51., 51., 51., 51., 51., 51., 51.], + [36., 51., 51., 51., 51., 51., 51., 51., 51., 51.], + [36., 51., 51., 51., 51., 51., 51., 51., 51., 51.], + [36., 51., 51., 51., 51., 51., 51., 51., 51., 51.], + ], + [ + [33., 48., 48., 48., 48., 48., 48., 48., 48., 48.], + [60., 87., 87., 87., 87., 87., 87., 87., 87., 87.], + [60., 87., 87., 87., 87., 87., 87., 87., 87., 87.], + [60., 87., 87., 87., 87., 87., 87., 87., 87., 87.], + [60., 87., 87., 87., 87., 87., 87., 87., 87., 87.], + [60., 87., 87., 87., 87., 87., 87., 87., 87., 87.], + [60., 87., 87., 87., 87., 87., 87., 87., 87., 87.], + [60., 87., 87., 87., 87., 87., 87., 87., 87., 87.], + [60., 87., 87., 87., 87., 87., 87., 87., 87., 87.], + [60., 87., 87., 87., 87., 87., 87., 87., 87., 87.], + ], + [ + [45., 66., 66., 66., 66., 66., 66., 66., 66., 66.], + [84., 123., 123., 123., 123., 123., 123., 123., 123., 123.], + [84., 123., 123., 123., 123., 123., 123., 123., 123., 123.], + [84., 123., 123., 123., 123., 123., 123., 123., 123., 123.], + [84., 123., 123., 123., 123., 123., 123., 123., 123., 123.], + [84., 123., 123., 123., 123., 123., 123., 123., 123., 123.], + [84., 123., 123., 123., 123., 123., 123., 123., 123., 123.], + [84., 123., 123., 123., 123., 123., 123., 123., 123., 123.], + [84., 123., 123., 123., 123., 123., 123., 123., 123., 123.], + [84., 123., 123., 123., 123., 123., 123., 123., 123., 123.], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [[[4455., 4905., 4905.], [4500., 4950., 4950.]]], + [[[12555., 13905., 13905.], [13500., 14950., 14950.]]], + [[[20655., 22905., 22905.], [22500., 24950., 24950.]]], + [[[28755., 31905., 31905.], [31500., 34950., 34950.]]], + ], + &device, + ), + bias: TestTensor::from_data([570., 570.], &device), + }; + test.assert_grads(grads); +} + +struct ConvTranspose2dTestCase { + batch_size: usize, + channels: [usize; 2], + kernel_size: [usize; 2], + padding: [usize; 2], + padding_out: [usize; 2], + stride: [usize; 2], + dilation: [usize; 2], + groups: usize, + size: [usize; 2], +} + +struct Grads { + x: TestTensor<4>, + weight: TestTensor<4>, + bias: TestTensor<1>, +} + +impl ConvTranspose2dTestCase { + fn assert_grads(self, expected_grads: Grads) { + let shape_x = Shape::new([ + self.batch_size, + self.channels[0], + self.size[0], + self.size[1], + ]); + let shape_weight = Shape::new([ + self.channels[0], + self.channels[1] / self.groups, + self.kernel_size[0], + self.kernel_size[1], + ]); + let device = AutodiffDevice::new(); + let weight = TestTensor::from_data( + TestTensorInt::arange(0..shape_weight.num_elements() as i64, &device) + .reshape::<4, _>(shape_weight) + .into_data(), + &device, + ) + .require_grad(); + let bias = TestTensor::from_data( + TestTensorInt::arange(0..self.channels[1] as i64, &device).into_data(), + &device, + ) + .require_grad(); + let x = TestTensor::from_data( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<4, _>(shape_x) + .into_data(), + &device, + ) + .require_grad(); + let output = conv_transpose2d( + x.clone(), + weight.clone(), + Some(bias.clone()), + ConvTransposeOptions::new( + self.stride, + self.padding, + self.padding_out, + self.dilation, + self.groups, + ), + ); + let grads = output.backward(); + + // Assert + let x_grad_actual = x.grad(&grads).unwrap(); + let weight_grad_actual = weight.grad(&grads).unwrap(); + let bias_grad_actual = bias.grad(&grads).unwrap(); + + let tolerance = Tolerance::permissive(); + expected_grads + .bias + .to_data() + .assert_approx_eq::(&bias_grad_actual.to_data(), tolerance); + expected_grads + .x + .to_data() + .assert_approx_eq::(&x_grad_actual.to_data(), tolerance); + expected_grads + .weight + .to_data() + .assert_approx_eq::(&weight_grad_actual.to_data(), tolerance); + } +} diff --git a/crates/burn-backend-tests/tests/autodiff/conv_transpose3d.rs b/crates/burn-backend-tests/tests/autodiff/conv_transpose3d.rs new file mode 100644 index 0000000..970cbdf --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/conv_transpose3d.rs @@ -0,0 +1,711 @@ +use super::*; +use burn_tensor::{Shape, Tolerance, module::conv_transpose3d, ops::ConvTransposeOptions}; + +#[test] +fn test_conv_transpose3d_basic() { + let test = ConvTranspose3dTestCase { + batch_size: 2, + channels: [2, 2], + kernel_size: [3, 3, 3], + padding: [0, 0, 0], + padding_out: [0, 0, 0], + stride: [1, 1, 1], + dilation: [1, 1, 1], + groups: 1, + size: [4, 4, 4], + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [ + [ + [ + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + ], + [ + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + ], + [ + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + ], + [ + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + ], + ], + [ + [ + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + ], + [ + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + ], + [ + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + ], + [ + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + ], + ], + ], + [ + [ + [ + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + ], + [ + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + ], + [ + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + ], + [ + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + [13.250001, 13.250001, 13.250001, 13.250001], + ], + ], + [ + [ + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + ], + [ + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + ], + [ + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + ], + [ + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + [40.249992, 40.249992, 40.249992, 40.249992], + ], + ], + ], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [ + [ + [47.750000, 47.750000, 47.750000], + [47.750000, 47.750000, 47.750000], + [47.750000, 47.750000, 47.750000], + ], + [ + [47.750000, 47.750000, 47.750000], + [47.750000, 47.750000, 47.750000], + [47.750000, 47.750000, 47.750000], + ], + [ + [47.750000, 47.750000, 47.750000], + [47.750000, 47.750000, 47.750000], + [47.750000, 47.750000, 47.750000], + ], + ], + [ + [ + [47.750000, 47.750000, 47.750000], + [47.750000, 47.750000, 47.750000], + [47.750000, 47.750000, 47.750000], + ], + [ + [47.750000, 47.750000, 47.750000], + [47.750000, 47.750000, 47.750000], + [47.750000, 47.750000, 47.750000], + ], + [ + [47.750000, 47.750000, 47.750000], + [47.750000, 47.750000, 47.750000], + [47.750000, 47.750000, 47.750000], + ], + ], + ], + [ + [ + [ + [79.750000, 79.750000, 79.750000], + [79.750000, 79.750000, 79.750000], + [79.750000, 79.750000, 79.750000], + ], + [ + [79.750000, 79.750000, 79.750000], + [79.750000, 79.750000, 79.750000], + [79.750000, 79.750000, 79.750000], + ], + [ + [79.750000, 79.750000, 79.750000], + [79.750000, 79.750000, 79.750000], + [79.750000, 79.750000, 79.750000], + ], + ], + [ + [ + [79.750000, 79.750000, 79.750000], + [79.750000, 79.750000, 79.750000], + [79.750000, 79.750000, 79.750000], + ], + [ + [79.750000, 79.750000, 79.750000], + [79.750000, 79.750000, 79.750000], + [79.750000, 79.750000, 79.750000], + ], + [ + [79.750000, 79.750000, 79.750000], + [79.750000, 79.750000, 79.750000], + [79.750000, 79.750000, 79.750000], + ], + ], + ], + ], + &device, + ), + bias: TestTensor::from_data([432., 432.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_conv_transpose3d_complex_groups() { + let test = ConvTranspose3dTestCase { + batch_size: 1, + channels: [4, 2], + kernel_size: [2, 3, 4], + padding: [1, 2, 3], + padding_out: [1, 2, 3], + stride: [2, 3, 4], + dilation: [1, 2, 3], + groups: 2, + size: [6, 6, 6], + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [ + [1.250000, 1.625000, 1.625000, 1.625000, 1.625000, 1.625000], + [1.687500, 2.187500, 2.187500, 2.187500, 2.187500, 2.187500], + [1.687500, 2.187500, 2.187500, 2.187500, 2.187500, 2.187500], + [1.687500, 2.187500, 2.187500, 2.187500, 2.187500, 2.187500], + [1.687500, 2.187500, 2.187500, 2.187500, 2.187500, 2.187500], + [1.687500, 2.187500, 2.187500, 2.187500, 2.187500, 2.187500], + ], + [ + [1.750000, 2.250000, 2.250000, 2.250000, 2.250000, 2.250000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + ], + [ + [1.750000, 2.250000, 2.250000, 2.250000, 2.250000, 2.250000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + ], + [ + [1.750000, 2.250000, 2.250000, 2.250000, 2.250000, 2.250000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + ], + [ + [1.750000, 2.250000, 2.250000, 2.250000, 2.250000, 2.250000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + ], + [ + [1.750000, 2.250000, 2.250000, 2.250000, 2.250000, 2.250000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + [2.250000, 2.875000, 2.875000, 2.875000, 2.875000, 2.875000], + ], + ], + [ + [ + [2.750000, 3.625000, 3.625000, 3.625000, 3.625000, 3.625000], + [3.937500, 5.187500, 5.187500, 5.187500, 5.187500, 5.187500], + [3.937500, 5.187500, 5.187500, 5.187500, 5.187500, 5.187500], + [3.937500, 5.187500, 5.187500, 5.187500, 5.187500, 5.187500], + [3.937500, 5.187500, 5.187500, 5.187500, 5.187500, 5.187500], + [3.937500, 5.187500, 5.187500, 5.187500, 5.187500, 5.187500], + ], + [ + [4.750000, 6.250000, 6.250000, 6.250000, 6.250000, 6.250000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + ], + [ + [4.750000, 6.250000, 6.250000, 6.250000, 6.250000, 6.250000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + ], + [ + [4.750000, 6.250000, 6.250000, 6.250000, 6.250000, 6.250000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + ], + [ + [4.750000, 6.250000, 6.250000, 6.250000, 6.250000, 6.250000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + ], + [ + [4.750000, 6.250000, 6.250000, 6.250000, 6.250000, 6.250000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + [6.750000, 8.875000, 8.875000, 8.875000, 8.875000, 8.875000], + ], + ], + [ + [ + [4.250000, 5.625000, 5.625000, 5.625000, 5.625000, 5.625000], + [6.187500, 8.187500, 8.187500, 8.187500, 8.187500, 8.187500], + [6.187500, 8.187500, 8.187500, 8.187500, 8.187500, 8.187500], + [6.187500, 8.187500, 8.187500, 8.187500, 8.187500, 8.187500], + [6.187500, 8.187500, 8.187500, 8.187500, 8.187500, 8.187500], + [6.187500, 8.187500, 8.187500, 8.187500, 8.187500, 8.187500], + ], + [ + [ + 7.750000, 10.250000, 10.250000, 10.250000, 10.250000, 10.250000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + ], + [ + [ + 7.750000, 10.250000, 10.250000, 10.250000, 10.250000, 10.250000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + ], + [ + [ + 7.750000, 10.250000, 10.250000, 10.250000, 10.250000, 10.250000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + ], + [ + [ + 7.750000, 10.250000, 10.250000, 10.250000, 10.250000, 10.250000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + ], + [ + [ + 7.750000, 10.250000, 10.250000, 10.250000, 10.250000, 10.250000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + [ + 11.250000, 14.875000, 14.875000, 14.875000, 14.875000, 14.875000, + ], + ], + ], + [ + [ + [5.750000, 7.625000, 7.625000, 7.625000, 7.625000, 7.625000], + [ + 8.437500, 11.187500, 11.187500, 11.187500, 11.187500, 11.187500, + ], + [ + 8.437500, 11.187500, 11.187500, 11.187500, 11.187500, 11.187500, + ], + [ + 8.437500, 11.187500, 11.187500, 11.187500, 11.187500, 11.187500, + ], + [ + 8.437500, 11.187500, 11.187500, 11.187500, 11.187500, 11.187500, + ], + [ + 8.437500, 11.187500, 11.187500, 11.187500, 11.187500, 11.187500, + ], + ], + [ + [ + 10.750000, 14.250000, 14.250000, 14.250000, 14.250000, 14.250000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + ], + [ + [ + 10.750000, 14.250000, 14.250000, 14.250000, 14.250000, 14.250000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + ], + [ + [ + 10.750000, 14.250000, 14.250000, 14.250000, 14.250000, 14.250000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + ], + [ + [ + 10.750000, 14.250000, 14.250000, 14.250000, 14.250000, 14.250000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + ], + [ + [ + 10.750000, 14.250000, 14.250000, 14.250000, 14.250000, 14.250000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + [ + 15.750000, 20.875000, 20.875000, 20.875000, 20.875000, 20.875000, + ], + ], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [[ + [ + [18.663193, 22.309027, 22.309027, 22.309027], + [21.875000, 26.145834, 26.145834, 26.145834], + [21.875000, 26.145834, 26.145834, 26.145834], + ], + [ + [19.270832, 23.020834, 23.020834, 23.020834], + [22.500000, 26.875002, 26.875002, 26.875002], + [22.500000, 26.875002, 26.875002, 26.875002], + ], + ]], + [[ + [ + [49.913193, 59.809029, 59.809029, 59.809029], + [59.375000, 71.145836, 71.145836, 71.145836], + [59.375000, 71.145836, 71.145836, 71.145836], + ], + [ + [56.770836, 68.020836, 68.020836, 68.020836], + [67.500000, 80.875000, 80.875000, 80.875000], + [67.500000, 80.875000, 80.875000, 80.875000], + ], + ]], + [[ + [ + [81.163193, 97.309029, 97.309029, 97.309029], + [96.875000, 116.145828, 116.145828, 116.145828], + [96.875000, 116.145828, 116.145828, 116.145828], + ], + [ + [94.270828, 113.020828, 113.020828, 113.020828], + [112.500000, 134.875000, 134.875000, 134.875000], + [112.500000, 134.875000, 134.875000, 134.875000], + ], + ]], + [[ + [ + [112.413200, 134.809021, 134.809021, 134.809021], + [134.375000, 161.145828, 161.145828, 161.145828], + [134.375000, 161.145828, 161.145828, 161.145828], + ], + [ + [131.770844, 158.020828, 158.020828, 158.020828], + [157.500000, 188.875000, 188.875000, 188.875000], + [157.500000, 188.875000, 188.875000, 188.875000], + ], + ]], + ], + &device, + ), + bias: TestTensor::from_data([5346., 5346.], &device), + }; + test.assert_grads(grads); +} + +struct ConvTranspose3dTestCase { + batch_size: usize, + channels: [usize; 2], + kernel_size: [usize; 3], + padding: [usize; 3], + padding_out: [usize; 3], + stride: [usize; 3], + dilation: [usize; 3], + groups: usize, + size: [usize; 3], +} + +struct Grads { + x: TestTensor<5>, + weight: TestTensor<5>, + bias: TestTensor<1>, +} + +impl ConvTranspose3dTestCase { + fn assert_grads(self, expected_grads: Grads) { + let shape_x = Shape::new([ + self.batch_size, + self.channels[0], + self.size[0], + self.size[1], + self.size[2], + ]); + let shape_weight = Shape::new([ + self.channels[0], + self.channels[1] / self.groups, + self.kernel_size[0], + self.kernel_size[1], + self.kernel_size[2], + ]); + let device = AutodiffDevice::new(); + let weight = TestTensor::from_data( + TestTensorInt::arange(0..shape_weight.num_elements() as i64, &device) + .reshape::<5, _>(shape_weight.clone()) + .into_data(), + &device, + ) + .div_scalar(shape_weight.num_elements() as f32) + .require_grad(); + let bias = TestTensor::from_data( + TestTensorInt::arange(0..self.channels[1] as i64, &device).into_data(), + &device, + ) + .require_grad(); + let x = TestTensor::from_data( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<5, _>(shape_x.clone()) + .into_data(), + &device, + ) + .div_scalar(shape_x.num_elements() as f32) + .require_grad(); + let output = conv_transpose3d( + x.clone(), + weight.clone(), + Some(bias.clone()), + ConvTransposeOptions::new( + self.stride, + self.padding, + self.padding_out, + self.dilation, + self.groups, + ), + ); + let grads = output.backward(); + + // Assert + let x_grad_actual = x.grad(&grads).unwrap(); + let weight_grad_actual = weight.grad(&grads).unwrap(); + let bias_grad_actual = bias.grad(&grads).unwrap(); + + let tolerance = Tolerance::permissive(); + expected_grads + .bias + .to_data() + .assert_approx_eq::(&bias_grad_actual.to_data(), tolerance); + expected_grads + .x + .to_data() + .assert_approx_eq::(&x_grad_actual.to_data(), tolerance); + expected_grads + .weight + .to_data() + .assert_approx_eq::(&weight_grad_actual.to_data(), tolerance); + } +} diff --git a/crates/burn-backend-tests/tests/autodiff/cross.rs b/crates/burn-backend-tests/tests/autodiff/cross.rs new file mode 100644 index 0000000..58e9dcd --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/cross.rs @@ -0,0 +1,130 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn backward_basic() { + let device = AutodiffDevice::new(); + let a = TestTensor::<2>::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), + &device, + ) + .require_grad(); + let b = TestTensor::<2>::from_data( + TensorData::from([[4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), + &device, + ) + .require_grad(); + + // Simple cross product; grad is a vector of ones. + let c = a.clone().cross(b.clone(), 1); + let grads = c.backward(); + + let a_grad = a.grad(&grads).unwrap().to_data(); + let b_grad = b.grad(&grads).unwrap().to_data(); + + // For a: b×grad_out, where grad_out = [1,1,1] + let expected_a = TensorData::from([[-1.0, 2.0, -1.0], [-1.0, 2.0, -1.0]]); + // For b: grad_out×a + let expected_b = TensorData::from([[1.0, -2.0, 1.0], [1.0, -2.0, 1.0]]); + + a_grad.assert_approx_eq::(&expected_a, Tolerance::default()); + b_grad.assert_approx_eq::(&expected_b, Tolerance::default()); +} + +#[test] +fn backward_after_sum() { + let device = AutodiffDevice::new(); + let a = TestTensor::<2>::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), + &device, + ) + .require_grad(); + let b = TestTensor::<2>::from_data( + TensorData::from([[4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), + &device, + ) + .require_grad(); + + // Sum reduces to scalar, but the gradient should be the same. + let c = a.clone().cross(b.clone(), 1).sum(); + let grads = c.backward(); + + let a_grad = a.grad(&grads).unwrap().to_data(); + let b_grad = b.grad(&grads).unwrap().to_data(); + + let expected_a = TensorData::from([[-1.0, 2.0, -1.0], [-1.0, 2.0, -1.0]]); + let expected_b = TensorData::from([[1.0, -2.0, 1.0], [1.0, -2.0, 1.0]]); + + a_grad.assert_approx_eq::(&expected_a, Tolerance::default()); + b_grad.assert_approx_eq::(&expected_b, Tolerance::default()); +} + +#[test] +fn different_dim() { + // Cross along a non-last dimension (dim 0) treats columns as 3-vectors. + let device = AutodiffDevice::new(); + let a_raw = [[1.0, 4.0, 7.0], [2.0, 5.0, 8.0], [3.0, 6.0, 9.0]]; + let b_raw = [[9.0, 6.0, 3.0], [8.0, 5.0, 2.0], [7.0, 4.0, 1.0]]; + + let a = TestTensor::<2>::from_data(TensorData::from(a_raw), &device); + let b = TestTensor::<2>::from_data(TensorData::from(b_raw), &device); + let out = a.cross(b.clone(), 0); + + // Manually compute cross of each column vector using raw arrays + let expected = [ + [ + a_raw[1][0] * b_raw[2][0] - a_raw[2][0] * b_raw[1][0], + a_raw[1][1] * b_raw[2][1] - a_raw[2][1] * b_raw[1][1], + a_raw[1][2] * b_raw[2][2] - a_raw[2][2] * b_raw[1][2], + ], + [ + a_raw[2][0] * b_raw[0][0] - a_raw[0][0] * b_raw[2][0], + a_raw[2][1] * b_raw[0][1] - a_raw[0][1] * b_raw[2][1], + a_raw[2][2] * b_raw[0][2] - a_raw[0][2] * b_raw[2][2], + ], + [ + a_raw[0][0] * b_raw[1][0] - a_raw[1][0] * b_raw[0][0], + a_raw[0][1] * b_raw[1][1] - a_raw[1][1] * b_raw[0][1], + a_raw[0][2] * b_raw[1][2] - a_raw[1][2] * b_raw[0][2], + ], + ]; + + out.to_data() + .assert_approx_eq::(&TensorData::from(expected), Tolerance::default()); +} + +#[test] +fn backward_non_last_dim() { + // Backward through a cross on dim 0. The autodiff rule recurses through + // float_cross with the same dim, so this also exercises the non-last-dim + // forward path on the gradient pass. + let device = AutodiffDevice::new(); + let a = TestTensor::<2>::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), + &device, + ) + .require_grad(); + let b = TestTensor::<2>::from_data( + TensorData::from([[9.0, 8.0, 7.0], [6.0, 5.0, 4.0], [3.0, 2.0, 1.0]]), + &device, + ) + .require_grad(); + + // Cross on the column dimension, then permute and cross on the last dim. + // Both paths share the same expected gradient magnitudes; we compare the + // dim-0 output against the reference last-dim output to keep the check + // backend-agnostic. + let c0 = a.clone().cross(b.clone(), 0); + let c_ref = a + .clone() + .permute([1, 0]) + .cross(b.clone().permute([1, 0]), 1) + .permute([1, 0]); + + c0.to_data() + .assert_approx_eq::(&c_ref.to_data(), Tolerance::default()); + + let grads = c0.sum().backward(); + assert!(a.grad(&grads).is_some()); + assert!(b.grad(&grads).is_some()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/cross_entropy.rs b/crates/burn-backend-tests/tests/autodiff/cross_entropy.rs new file mode 100644 index 0000000..2925039 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/cross_entropy.rs @@ -0,0 +1,32 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance, loss}; + +#[test] +fn test_cross_entropy_loss_grad() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[6.0, 7.0], [9.0, 10.0]]); + let data_targets = TensorData::from([[0.8, 0.2], [0.9, 0.1]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device).require_grad(); + let tensor_targets = TestTensor::<2>::from_data(data_targets, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = loss::cross_entropy_with_logits(tensor_3, tensor_targets); + + let grads = tensor_4.backward(); + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let tolerance = Tolerance::permissive(); + let expected = TensorData::from([[0.26553, 0.26553], [0.44954, 0.44954]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([[-1.34863, 1.34863], [-2.06371, 2.06371]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, tolerance); +} diff --git a/crates/burn-backend-tests/tests/autodiff/ctc.rs b/crates/burn-backend-tests/tests/autodiff/ctc.rs new file mode 100644 index 0000000..2d1b016 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/ctc.rs @@ -0,0 +1,754 @@ +#![allow(clippy::excessive_precision)] + +use super::*; +use burn_tensor::{Device, TensorData, Tolerance, activation::log_softmax, module::ctc_loss}; + +/// Deterministic logits: sin((t*7 + n*13 + c*3) * 0.1). Matches the +/// fixture generator in `burn-nn/src/loss/ctc.rs`'s PyTorch comparison tests. +fn generate_logits(t_size: usize, n_size: usize, c_size: usize, device: &Device) -> TestTensor<3> { + let mut data = Vec::with_capacity(t_size * n_size * c_size); + for t in 0..t_size { + for n in 0..n_size { + for c in 0..c_size { + data.push(((t * 7 + n * 13 + c * 3) as f32 * 0.1).sin()); + } + } + } + TestTensor::<3>::from_data(TensorData::new(data, [t_size, n_size, c_size]), device) +} + +#[allow(clippy::too_many_arguments)] +fn run_comparison( + t_size: usize, + n_size: usize, + c_size: usize, + targets_flat: Vec, + target_shape: [usize; 2], + input_lengths: Vec, + target_lengths: Vec, + blank: usize, + expected_grad_flat: &[f32], +) { + let device = AutodiffDevice::new(); + let logits = generate_logits(t_size, n_size, c_size, &device).require_grad(); + let log_probs = log_softmax(logits.clone(), 2); + + let targets = + TestTensorInt::<2>::from_data(TensorData::new(targets_flat, target_shape), &device); + let input_lengths = + TestTensorInt::<1>::from_data(TensorData::new(input_lengths, [n_size]), &device); + let target_lengths = + TestTensorInt::<1>::from_data(TensorData::new(target_lengths, [n_size]), &device); + + let loss = ctc_loss(log_probs, targets, input_lengths, target_lengths, blank).sum(); + let grads = loss.backward(); + let logits_grad = logits.grad(&grads).expect( + "logits should receive a gradient - if this fails on a backend with a fused ctc_loss \ + kernel, the autodiff override needs to skip the fused path", + ); + + logits_grad.into_data().assert_approx_eq::( + &TensorData::new(expected_grad_flat.to_vec(), [t_size, n_size, c_size]), + Tolerance::rel_abs(1e-3, 1e-3).set_half_precision_absolute(1e-2), + ); +} + +/// Original coverage: uniform input_lengths, no repeated labels, small T. +/// Keeps a first-row sanity check against the PyTorch reference so a bug in +/// the first batch element's first timestep is caught quickly. +#[test] +fn test_ctc_loss_grad() { + let t_size: usize = 5; + let n_size: usize = 3; + let c_size: usize = 4; + let device = AutodiffDevice::new(); + + let logits = generate_logits(t_size, n_size, c_size, &device).require_grad(); + let log_probs = log_softmax(logits.clone(), 2); + + let targets = + TestTensorInt::<2>::from_data(TensorData::from([[1, 2, 0], [1, 0, 0], [3, 2, 1]]), &device); + let input_lengths = TestTensorInt::<1>::from_data(TensorData::from([5, 5, 5]), &device); + let target_lengths = TestTensorInt::<1>::from_data(TensorData::from([2, 1, 3]), &device); + + let loss = ctc_loss(log_probs, targets, input_lengths, target_lengths, 0).sum(); + let grads = loss.backward(); + let logits_grad = logits.grad(&grads).expect( + "logits should receive a gradient - if this fails on a backend with a fused ctc_loss \ + kernel, the autodiff override needs to skip the fused path", + ); + + let expected_first_row = + TensorData::from([-0.1679008007_f32, -0.4595540464, 0.2795598209, 0.3478950262]); + let logits_grad_data = logits_grad.into_data(); + let first_row: Vec = logits_grad_data.iter::().take(4).collect(); + TensorData::from(first_row.as_slice()) + .assert_approx_eq::(&expected_first_row, Tolerance::rel_abs(1e-3, 1e-3)); +} + +/// Variable input_lengths exercise the per-sample `t_last = input_len - 1` +/// boundary init in the beta recursion and the OOB mask in the gradient +/// composition. Fixture lifted from burn-nn's pytorch_comparison_tests. +#[test] +fn test_ctc_loss_grad_mixed_input_lengths() { + // T=12, N=3, C=5, input_lengths=[12, 7, 10], target_lengths=[3, 2, 4]. + let expected_grad_flat = [ + -0.4790987670_f32, + -0.2554937005, + 0.1991624236, + 0.2478453964, + 0.2875846624, + -0.3495813310, + 0.2268397957, + 0.2150714993, + -0.2442178279, + 0.1518878639, + -0.2764556706, + 0.2474014312, + -0.2137086987, + 0.1371368915, + 0.1056260392, + -0.2729502618, + -0.3609606028, + 0.2159237266, + 0.2238420397, + 0.1941450834, + -0.2953839302, + 0.1920599341, + 0.1974952668, + -0.2054278404, + 0.1112565696, + -0.1719199270, + 0.2299505472, + -0.2864859998, + 0.1497263014, + 0.0787290633, + -0.2035763413, + -0.3042884767, + 0.2126964629, + 0.1810975969, + 0.1140707731, + -0.2759391963, + 0.0975771844, + 0.1823379993, + -0.1112988219, + 0.1073228419, + -0.1336459517, + 0.1869296581, + -0.1996247321, + 0.1846873760, + -0.0383463502, + -0.2254105806, + -0.1834360659, + 0.1925925612, + 0.1462381780, + 0.0700158924, + -0.2259973884, + -0.0393539183, + 0.1802661419, + -0.0571591072, + 0.1422442794, + -0.0609069727, + 0.1089282706, + -0.0313654318, + 0.2186669111, + -0.2353227735, + -0.2840364873, + -0.0632198900, + 0.1755636632, + 0.1377806067, + 0.0339120962, + -0.1904856712, + -0.2139032930, + 0.1827126741, + 0.0056131603, + 0.2160631120, + -0.0243270602, + -0.0070458520, + 0.1070247591, + 0.2239368409, + -0.2995886803, + -0.2955487072, + 0.0309870224, + 0.1654911339, + 0.1581364125, + -0.0590658709, + -0.2191396207, + -0.3791662455, + 0.1803640425, + 0.1225430891, + 0.2953987718, + -0.0436352938, + -0.1575258970, + 0.1785279512, + 0.1756918877, + -0.1530586481, + -0.1834939867, + 0.0909025446, + 0.1423641294, + 0.1959712654, + -0.2457439601, + -0.3619639874, + -0.3929221630, + 0.1820438206, + 0.2454170734, + 0.3274252713, + -0.0628800318, + -0.2567180395, + 0.2112283260, + 0.0507859327, + 0.0575838275, + -0.0587697029, + 0.1174769849, + 0.0783569664, + 0.2290501744, + -0.3661144078, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.0725664943, + -0.1532069892, + 0.2162397504, + -0.1248963475, + 0.1344300956, + -0.0362483934, + 0.1295878887, + -0.0502482466, + 0.2470482886, + -0.2901395261, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.1349253207, + 0.0867646411, + 0.1998746395, + -0.2658679783, + 0.1141540110, + -0.0705668628, + 0.1519546807, + -0.2509805560, + 0.2475892603, + -0.0779965296, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.2338010073, + 0.2471641302, + 0.1834627241, + -0.3026831448, + 0.1058573127, + -0.1155209392, + 0.1921830922, + -0.4129956067, + 0.2229512781, + 0.1133821756, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.2636392713, + 0.2323469073, + -0.2913427949, + 0.1800564528, + 0.1425786912, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + ]; + run_comparison( + 12, + 3, + 5, + vec![1, 4, 2, 0, 3, 1, 0, 0, 2, 4, 1, 3], + [3, 4], + vec![12, 7, 10], + vec![3, 2, 4], + 0, + &expected_grad_flat, + ); +} + +/// Consecutive repeated labels `[1, 1, ...]` force the alpha and beta skip +/// transitions to stay disabled (`l'[s] != l'[s-2]` fails). Exercises the +/// `skip_allowed` branch condition on the fused kernel. +#[test] +fn test_ctc_loss_grad_repeated_labels() { + // T=8, N=4, C=6, targets include the consecutive pair [1, 1, 2]. + let expected_grad_flat = [ + -0.2766432464_f32, + -0.5202965736, + 0.1523768753, + 0.1896236390, + 0.2200277001, + 0.2349116206, + -0.1854365915, + 0.2031330466, + -0.4260218740, + 0.1678018719, + 0.1360142529, + 0.1045092493, + -0.6603536606, + 0.2278252542, + 0.1691786796, + 0.1262856424, + 0.0972681716, + 0.0397959016, + -0.0894432291, + -0.5457318425, + 0.1490373611, + 0.1462858170, + 0.1569476575, + 0.1829041988, + -0.2842915654, + -0.4220107496, + 0.1822281033, + 0.1889107376, + 0.1791101843, + 0.1560532600, + -0.1155678406, + 0.2295538932, + -0.2645366490, + -0.0288553704, + 0.1027252972, + 0.0766806602, + -0.5448347330, + 0.2031028718, + 0.1589304954, + 0.1322451383, + 0.1189499870, + -0.0683937520, + -0.0873993114, + -0.3051757514, + -0.2355299890, + 0.1586059481, + 0.2018169016, + 0.2676822543, + -0.3225219846, + -0.2611543834, + 0.1922984123, + 0.1632783115, + 0.1297036558, + 0.0983960181, + -0.1507159024, + 0.2256962359, + -0.1040333956, + -0.1514528394, + 0.0985243544, + 0.0819815546, + -0.2940836251, + 0.1586865336, + 0.1468491107, + 0.1485087872, + 0.1639631987, + -0.3239239752, + -0.0767390430, + -0.0434846729, + -0.4023587406, + -0.0052628326, + 0.2273432612, + 0.3005020022, + -0.2598774135, + -0.2188862711, + 0.1678501070, + 0.1352078766, + 0.1002781317, + 0.0754275694, + -0.1502914876, + 0.1930875033, + -0.0709601715, + -0.2219523191, + 0.1243555173, + 0.1257609427, + -0.0574148744, + 0.1152269915, + 0.1307857931, + 0.1599020809, + 0.2068412602, + -0.5553412437, + -0.0536844917, + 0.0758557543, + -0.2106334567, + -0.2509877980, + 0.1757438034, + 0.2637061775, + -0.1759711355, + -0.2431350052, + 0.1071053818, + 0.1259848624, + 0.1004033238, + 0.0856125653, + -0.1173698306, + 0.1213828772, + -0.1768893301, + -0.2070008069, + 0.1709136516, + 0.2089634240, + 0.0153109450, + 0.0967332721, + 0.1268781722, + 0.1706230640, + 0.2291058898, + -0.6386513710, + -0.0536664203, + 0.1378114969, + 0.0360041447, + -0.2989685237, + -0.0084722806, + 0.1872915775, + -0.1523490399, + -0.2111770809, + -0.0390694551, + 0.1366800815, + 0.1302325875, + 0.1356829405, + -0.0982905105, + -0.0127884001, + -0.3586881459, + -0.0259541404, + 0.2114149332, + 0.2843062580, + -0.0324133746, + 0.1084750593, + 0.1447229236, + 0.1862253845, + 0.2259712219, + -0.6329812407, + -0.1173689738, + 0.1914442331, + 0.1654772907, + -0.1376858056, + -0.2194855511, + 0.1176188141, + -0.1529908478, + -0.0606661662, + -0.3384291232, + 0.1524862647, + 0.1777049750, + 0.2218948901, + -0.0923086405, + -0.2855934799, + -0.3215619624, + 0.1726681292, + 0.2303666323, + 0.2964293361, + -0.2508065701, + 0.1479703039, + 0.1753441393, + 0.1917535067, + 0.1919818372, + -0.4562432170, + -0.2350299209, + 0.2257601619, + 0.1863904297, + 0.0388212129, + -0.2966264784, + 0.0806845874, + -0.1992894858, + 0.1068909168, + -0.5761897564, + 0.1624972969, + 0.2155302167, + 0.2905607820, + -0.1168124676, + -0.6870660186, + 0.1488010883, + 0.1881926507, + 0.2230074406, + 0.2438773215, + -0.5771554708, + 0.1980127096, + 0.1924194694, + 0.1714663208, + 0.1415647417, + -0.1263078004, + -0.3408652246, + 0.2292248607, + 0.1707807332, + 0.1269564927, + -0.2634142637, + 0.0773174241, + ]; + run_comparison( + 8, + 4, + 6, + vec![1, 1, 2, 0, 2, 3, 2, 1, 5, 0, 0, 0, 1, 2, 3, 4], + [4, 4], + vec![8, 8, 8, 8], + vec![3, 4, 1, 4], + 0, + &expected_grad_flat, + ); +} + +/// Empty target (`target_length == 0`). Analytical gradient: with uniform +/// logits the only valid alignment is all-blank, so for each `(t, c)` the +/// gradient of `sum(loss)` w.r.t. logits is `-0.5` at the blank class and +/// `+0.5` at the other class. +#[test] +fn test_ctc_loss_grad_empty_target() { + let t_size: usize = 3; + let n_size: usize = 1; + let c_size: usize = 2; + let device = AutodiffDevice::new(); + + // Uniform logits (any value works; softmax normalizes). Use 0.0 for clarity. + let logits = TestTensor::<3>::from_data( + TensorData::new( + alloc::vec![0.0f32; t_size * n_size * c_size], + [t_size, n_size, c_size], + ), + &device, + ) + .require_grad(); + let log_probs = log_softmax(logits.clone(), 2); + + // Dummy target column (never read: target_length is 0). Using [1, 1] + // avoids backends that don't support zero-sized dims. + let targets = TestTensorInt::<2>::from_data(TensorData::from([[0_i64]]), &device); + let input_lengths = TestTensorInt::<1>::from_data(TensorData::from([3_i64]), &device); + let target_lengths = TestTensorInt::<1>::from_data(TensorData::from([0_i64]), &device); + + let loss = ctc_loss(log_probs, targets, input_lengths, target_lengths, 0).sum(); + let grads = loss.backward(); + let logits_grad = logits.grad(&grads).unwrap(); + + let expected = TensorData::new( + alloc::vec![-0.5_f32, 0.5, -0.5, 0.5, -0.5, 0.5], + [t_size, n_size, c_size], + ); + logits_grad + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(1e-4, 1e-4)); +} + +/// Regression guard for NaN gradient on unreachable targets, driven through +/// `.backward()` rather than a direct `ctc_loss_backward` call. The direct-path +/// version is covered by `test_ctc_loss_backward_unreachable_is_finite`; this +/// test exists because a lazy-backend regression in log_sum_exp or the +/// grad-composition mask can leave NaN in the autodiff trace even when the +/// direct backward call is finite. +/// +/// Note: the direct backward asserts grad == 0 (the mask inside +/// `ctc_grad_from_alpha_beta_default` zeroes entries for +inf nll). The autodiff +/// path currently produces a non-zero but finite gradient for this case because +/// `loss.sum().backward()` seeds through a +inf forward value. Finiteness is +/// what guards training stability; the zero-vs-finite-nonzero difference is a +/// separate concern. +#[test] +fn test_ctc_loss_grad_unreachable_is_finite() { + let device = AutodiffDevice::new(); + + // T=2, target=[1, 1] needs three steps, so nll = +inf for this sample. + let log_probs = TestTensor::<3>::full([2, 1, 3], (1.0f32 / 3.0).ln(), &device).require_grad(); + + let targets = TestTensorInt::<2>::from_data(TensorData::from([[1_i64, 1]]), &device); + let input_lengths = TestTensorInt::<1>::from_data(TensorData::from([2_i64]), &device); + let target_lengths = TestTensorInt::<1>::from_data(TensorData::from([2_i64]), &device); + + let loss = ctc_loss(log_probs.clone(), targets, input_lengths, target_lengths, 0).sum(); + let grads = loss.backward(); + let log_probs_grad = log_probs.grad(&grads).unwrap(); + for v in log_probs_grad.into_data().iter::() { + assert!( + v.is_finite(), + "gradient through .backward() must be finite for unreachable targets, got {v}", + ); + } +} + +/// Gradient through `.backward()` on a `swap_dims + narrow` input chain, a layout +/// produced by typical training code ([B, T, C] model output -> [T, B, C] -> +/// skip warmup). Forward-only tests cover `narrow` and `swap_dims+narrow`; this +/// also checks that the autodiff path preserves strides end-to-end rather than +/// materializing a contiguous copy with wrong offsets. +#[test] +fn test_ctc_loss_grad_swap_dims_then_narrow() { + let b: usize = 2; + let t_full: usize = 8; + let warmup: usize = 2; + let t_eff: usize = t_full - warmup; + let c: usize = 4; + + let device = AutodiffDevice::new(); + + let mut data = Vec::with_capacity(b * t_full * c); + for bi in 0..b { + for ti in 0..t_full { + for ci in 0..c { + data.push(((bi * 31 + ti * 7 + ci * 3) as f32 * 0.1).sin()); + } + } + } + + // Path A: [B, T, C] -> swap -> narrow, all on a single require_grad leaf. + let logits_btc_a = + TestTensor::<3>::from_data(TensorData::new(data.clone(), [b, t_full, c]), &device) + .require_grad(); + let log_probs_a = log_softmax(logits_btc_a.clone(), 2) + .swap_dims(0, 1) + .narrow(0, warmup, t_eff); + + let targets_a = TestTensorInt::<2>::from_data(TensorData::from([[1_i64, 2], [2, 3]]), &device); + let input_lengths_a = + TestTensorInt::<1>::from_data(TensorData::from([t_eff as i64, t_eff as i64]), &device); + let target_lengths_a = TestTensorInt::<1>::from_data(TensorData::from([2_i64, 2]), &device); + + let loss_a = ctc_loss(log_probs_a, targets_a, input_lengths_a, target_lengths_a, 0).sum(); + let grads_a = loss_a.backward(); + let grad_a = logits_btc_a.grad(&grads_a).unwrap(); + + // Path B: build the equivalent [T, B, C] contiguous layout up front so the + // log_softmax -> ctc_loss chain sees a plain contiguous tensor. Transpose + // the incoming data to [T, B, C] order. + let mut data_tbc = Vec::with_capacity(t_full * b * c); + for ti in 0..t_full { + for bi in 0..b { + for ci in 0..c { + data_tbc.push(((bi * 31 + ti * 7 + ci * 3) as f32 * 0.1).sin()); + } + } + } + let logits_tbc_b = + TestTensor::<3>::from_data(TensorData::new(data_tbc, [t_full, b, c]), &device) + .require_grad(); + let log_probs_b = log_softmax(logits_tbc_b.clone(), 2).narrow(0, warmup, t_eff); + + let targets_b = TestTensorInt::<2>::from_data(TensorData::from([[1_i64, 2], [2, 3]]), &device); + let input_lengths_b = + TestTensorInt::<1>::from_data(TensorData::from([t_eff as i64, t_eff as i64]), &device); + let target_lengths_b = TestTensorInt::<1>::from_data(TensorData::from([2_i64, 2]), &device); + + let loss_b = ctc_loss(log_probs_b, targets_b, input_lengths_b, target_lengths_b, 0).sum(); + let grads_b = loss_b.backward(); + let grad_b_tbc = logits_tbc_b.grad(&grads_b).unwrap(); + + // Compare path A's grad (laid out [B, T, C]) against path B's grad (laid + // out [T, B, C]) element by element. + let grad_a_data: Vec = grad_a.into_data().iter::().collect(); + let grad_b_data: Vec = grad_b_tbc.into_data().iter::().collect(); + for bi in 0..b { + for ti in 0..t_full { + for ci in 0..c { + let idx_a = (bi * t_full + ti) * c + ci; + let idx_b = (ti * b + bi) * c + ci; + let diff = (grad_a_data[idx_a] - grad_b_data[idx_b]).abs(); + assert!( + diff < 1e-3, + "grad mismatch at (b={bi}, t={ti}, c={ci}): A={} vs B={}", + grad_a_data[idx_a], + grad_b_data[idx_b], + ); + } + } + } +} + +/// Autodiff at production scale. The training repo (speechlet exp-11) reported +/// loss ~0.18 instead of ~1250 on a freshly-initialized model when log_probs +/// were routed through `log_softmax -> swap_dims(0,1) -> narrow(0, warmup, ...)` +/// on `Autodiff>`, while the contiguous-[T, B, C] reference +/// (same raw data) produced the correct value. Reproduces that exact shape +/// chain with a deterministic fixture and compares the two paths. +#[test] +fn test_ctc_loss_scaled_swap_dims_then_narrow_autodiff() { + let b: usize = 2; + let t_full: usize = 494; + let warmup: usize = 92; + let t_eff: usize = t_full - warmup; + let c: usize = 36; + + let device = AutodiffDevice::new(); + + let mut data_btc = Vec::with_capacity(b * t_full * c); + for bi in 0..b { + for ti in 0..t_full { + for ci in 0..c { + data_btc.push(((bi * 31 + ti * 7 + ci * 3) as f32 * 0.1).sin()); + } + } + } + + let logits_btc_a = + TestTensor::<3>::from_data(TensorData::new(data_btc.clone(), [b, t_full, c]), &device) + .require_grad(); + let log_probs_a = log_softmax(logits_btc_a, 2) + .swap_dims(0, 1) + .narrow(0, warmup, t_eff); + + let mut targets_data = alloc::vec![0_i64; b * 64]; + for bi in 0..b { + let max = if bi == 0 { 64 } else { 43 }; + for si in 0..max { + targets_data[bi * 64 + si] = ((si * 3 + bi * 7) % 35 + 1) as i64; + } + } + let targets_a = + TestTensorInt::<2>::from_data(TensorData::new(targets_data.clone(), [b, 64]), &device); + let input_lengths_a = + TestTensorInt::<1>::from_data(TensorData::from([t_eff as i64, t_eff as i64]), &device); + let target_lengths_a = TestTensorInt::<1>::from_data(TensorData::from([64_i64, 43]), &device); + + let loss_a_per_sample = ctc_loss(log_probs_a, targets_a, input_lengths_a, target_lengths_a, 0); + let loss_a_vec: Vec = loss_a_per_sample + .clone() + .into_data() + .iter::() + .collect(); + + // Path B: build the equivalent contiguous [T, B, C] layout up front. + let mut data_tbc = Vec::with_capacity(t_full * b * c); + for ti in 0..t_full { + for bi in 0..b { + for ci in 0..c { + data_tbc.push(((bi * 31 + ti * 7 + ci * 3) as f32 * 0.1).sin()); + } + } + } + let logits_tbc_b = + TestTensor::<3>::from_data(TensorData::new(data_tbc, [t_full, b, c]), &device); + let log_probs_b = log_softmax(logits_tbc_b, 2).narrow(0, warmup, t_eff); + + let targets_b = TestTensorInt::<2>::from_data(TensorData::new(targets_data, [b, 64]), &device); + let input_lengths_b = + TestTensorInt::<1>::from_data(TensorData::from([t_eff as i64, t_eff as i64]), &device); + let target_lengths_b = TestTensorInt::<1>::from_data(TensorData::from([64_i64, 43]), &device); + + let loss_b_per_sample = ctc_loss(log_probs_b, targets_b, input_lengths_b, target_lengths_b, 0); + let loss_b_vec: Vec = loss_b_per_sample.into_data().iter::().collect(); + + for i in 0..b { + assert!( + loss_a_vec[i] > 0.0, + "sample {i}: autodiff+swap_dims+narrow loss is not positive ({}); a valid CTC \ + loss must be >= 0 for log-probs <= 0", + loss_a_vec[i], + ); + assert!( + (loss_a_vec[i] - loss_b_vec[i]).abs() < 1.0, + "sample {i}: autodiff swap_dims+narrow loss ({}) diverges from contiguous \ + reference loss ({})", + loss_a_vec[i], + loss_b_vec[i], + ); + } +} diff --git a/crates/burn-backend-tests/tests/autodiff/cummax.rs b/crates/burn-backend-tests/tests/autodiff/cummax.rs new file mode 100644 index 0000000..17ef405 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/cummax.rs @@ -0,0 +1,115 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn should_diff_cummax() { + // Simple test to verify cummax gradients work + let device = AutodiffDevice::new(); + let tensor = + TestTensor::<1>::from_data(TensorData::from([1.0, 3.0, 2.0]), &device).require_grad(); + + let output = tensor.clone().cummax(0); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // PyTorch reference: [1.0, 2.0, 0.0] + let expected = TensorData::from([1.0, 2.0, 0.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_cummax_2d() { + // Test 2D cummax gradients + let device = AutodiffDevice::new(); + let tensor = TestTensor::<2>::from_data( + TensorData::from([[1.0, 3.0, 2.0], [2.0, 5.0, 4.0]]), + &device, + ) + .require_grad(); + + let output = tensor.clone().cummax(1); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // PyTorch reference: [[1.0, 2.0, 0.0], [1.0, 2.0, 0.0]] + let expected = TensorData::from([[1.0, 2.0, 0.0], [1.0, 2.0, 0.0]]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_cummax_duplicate_values() { + // Test with duplicate maximum values - critical edge case + let device = AutodiffDevice::new(); + let tensor = + TestTensor::<1>::from_data(TensorData::from([1.0, 3.0, 3.0, 2.0]), &device).require_grad(); + + let output = tensor.clone().cummax(0); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // input: [1.0, 3.0, 3.0, 2.0] + // cummax: [1.0, 3.0, 3.0, 3.0] + // PyTorch reference: [1.0, 1.0, 2.0, 0.0] + // Position 2 gets grad from itself + position 3 + let expected = TensorData::from([1.0, 1.0, 2.0, 0.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_cummax_all_same() { + // Test with all same values + let device = AutodiffDevice::new(); + let tensor = + TestTensor::<1>::from_data(TensorData::from([2.0, 2.0, 2.0]), &device).require_grad(); + + let output = tensor.clone().cummax(0); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // PyTorch reference: [1.0, 1.0, 1.0] + // Each position matches cummax, so each gets its own gradient + let expected = TensorData::from([1.0, 1.0, 1.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_cummax_increasing() { + // Test with increasing sequence + let device = AutodiffDevice::new(); + let tensor = + TestTensor::<1>::from_data(TensorData::from([1.0, 2.0, 3.0, 4.0]), &device).require_grad(); + + let output = tensor.clone().cummax(0); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // PyTorch reference: [1.0, 1.0, 1.0, 1.0] + // Each position is a new maximum + let expected = TensorData::from([1.0, 1.0, 1.0, 1.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_cummax_2d_duplicates() { + // Test 2D with duplicate values + let device = AutodiffDevice::new(); + let tensor = TestTensor::<2>::from_data( + TensorData::from([[1.0, 3.0, 3.0, 2.0], [2.0, 5.0, 5.0, 4.0]]), + &device, + ) + .require_grad(); + + let output = tensor.clone().cummax(1); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // PyTorch reference: [[1.0, 1.0, 2.0, 0.0], [1.0, 1.0, 2.0, 0.0]] + let expected = TensorData::from([[1.0, 1.0, 2.0, 0.0], [1.0, 1.0, 2.0, 0.0]]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/cummin.rs b/crates/burn-backend-tests/tests/autodiff/cummin.rs new file mode 100644 index 0000000..04147fe --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/cummin.rs @@ -0,0 +1,115 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn should_diff_cummin() { + // Simple test to verify cummin gradients work + let device = AutodiffDevice::new(); + let tensor = + TestTensor::<1>::from_data(TensorData::from([3.0, 2.0, 4.0]), &device).require_grad(); + + let output = tensor.clone().cummin(0); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // PyTorch reference: [1.0, 2.0, 0.0] + let expected = TensorData::from([1.0, 2.0, 0.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_cummin_2d() { + // Test 2D cummin gradients + let device = AutodiffDevice::new(); + let tensor = TestTensor::<2>::from_data( + TensorData::from([[3.0, 2.0, 4.0], [5.0, 1.0, 3.0]]), + &device, + ) + .require_grad(); + + let output = tensor.clone().cummin(1); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // PyTorch reference: [[1.0, 2.0, 0.0], [1.0, 2.0, 0.0]] + let expected = TensorData::from([[1.0, 2.0, 0.0], [1.0, 2.0, 0.0]]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_cummin_duplicate_values() { + // Test with duplicate minimum values - critical edge case + let device = AutodiffDevice::new(); + let tensor = + TestTensor::<1>::from_data(TensorData::from([3.0, 2.0, 2.0, 4.0]), &device).require_grad(); + + let output = tensor.clone().cummin(0); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // input: [3.0, 2.0, 2.0, 4.0] + // cummin: [3.0, 2.0, 2.0, 2.0] + // PyTorch reference: [1.0, 1.0, 2.0, 0.0] + // Position 2 gets grad from itself + position 3 + let expected = TensorData::from([1.0, 1.0, 2.0, 0.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_cummin_all_same() { + // Test with all same values + let device = AutodiffDevice::new(); + let tensor = + TestTensor::<1>::from_data(TensorData::from([2.0, 2.0, 2.0]), &device).require_grad(); + + let output = tensor.clone().cummin(0); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // PyTorch reference: [1.0, 1.0, 1.0] + // Each position matches cummin, so each gets its own gradient + let expected = TensorData::from([1.0, 1.0, 1.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_cummin_decreasing() { + // Test with decreasing sequence + let device = AutodiffDevice::new(); + let tensor = + TestTensor::<1>::from_data(TensorData::from([5.0, 4.0, 3.0, 2.0]), &device).require_grad(); + + let output = tensor.clone().cummin(0); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // PyTorch reference: [1.0, 1.0, 1.0, 1.0] + // Each position is a new minimum + let expected = TensorData::from([1.0, 1.0, 1.0, 1.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_cummin_2d_duplicates() { + // Test 2D with duplicate values + let device = AutodiffDevice::new(); + let tensor = TestTensor::<2>::from_data( + TensorData::from([[3.0, 2.0, 2.0, 4.0], [5.0, 1.0, 1.0, 3.0]]), + &device, + ) + .require_grad(); + + let output = tensor.clone().cummin(1); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // PyTorch reference: [[1.0, 1.0, 2.0, 0.0], [1.0, 1.0, 2.0, 0.0]] + let expected = TensorData::from([[1.0, 1.0, 2.0, 0.0], [1.0, 1.0, 2.0, 0.0]]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/cumprod.rs b/crates/burn-backend-tests/tests/autodiff/cumprod.rs new file mode 100644 index 0000000..1196001 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/cumprod.rs @@ -0,0 +1,128 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn should_diff_cumprod() { + // Simple test to verify cumprod gradients work + let device = AutodiffDevice::new(); + let tensor = + TestTensor::<1>::from_data(TensorData::from([2.0, 3.0, 4.0]), &device).require_grad(); + + let output = tensor.clone().cumprod(0); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // PyTorch reference: [16.0, 10.0, 6.0] + let expected = TensorData::from([16.0, 10.0, 6.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_cumprod_2d() { + // Test 2D cumprod gradients + let device = AutodiffDevice::new(); + let tensor = TestTensor::<2>::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), + &device, + ) + .require_grad(); + + let output = tensor.clone().cumprod(1); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // PyTorch reference: [[9.0, 4.0, 2.0], [36.0, 28.0, 20.0]] + let expected = TensorData::from([[9.0, 4.0, 2.0], [36.0, 28.0, 20.0]]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +// TODO: The following tests are currently ignored due to a known limitation +// in the cumprod gradient implementation. The current implementation uses +// division (grad / input), which produces NaN when the input contains zeros. +// +// A proper fix requires implementing a zero-safe algorithm using exclusive +// cumulative products (similar to PyTorch's cumprod_backward or JAX's +// associative_scan approach). This is a non-trivial implementation that +// requires careful handling of cumulative products in both forward and +// reverse directions. +// +// See: https://github.com/tracel-ai/burn/issues/3864 +// +// References: +// - PyTorch: https://github.com/pytorch/pytorch (cumprod_backward) +// - JAX PR #2596: Parallel prefix scan implementation +// - TensorFlow Issue #3862: tf.cumprod's gradient produces nans given zeros + +#[test] +#[ignore = "cumprod gradient with zeros not yet implemented - produces NaN due to division by zero"] +fn should_diff_cumprod_zero_in_middle() { + // Test cumprod with zero in the middle - edge case for division + let device = AutodiffDevice::new(); + let tensor = + TestTensor::<1>::from_data(TensorData::from([2.0, 0.0, 3.0, 4.0]), &device).require_grad(); + + let output = tensor.clone().cumprod(0); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // PyTorch reference: [1.0, 32.0, 0.0, 0.0] + let expected = TensorData::from([1.0, 32.0, 0.0, 0.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +#[ignore = "cumprod gradient with zeros not yet implemented - produces NaN due to division by zero"] +fn should_diff_cumprod_zero_at_start() { + // Test cumprod with zero at the beginning + let device = AutodiffDevice::new(); + let tensor = + TestTensor::<1>::from_data(TensorData::from([0.0, 2.0, 3.0, 4.0]), &device).require_grad(); + + let output = tensor.clone().cumprod(0); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // PyTorch reference: [33.0, 0.0, 0.0, 0.0] + let expected = TensorData::from([33.0, 0.0, 0.0, 0.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +#[ignore = "cumprod gradient with zeros not yet implemented - produces NaN due to division by zero"] +fn should_diff_cumprod_zero_at_end() { + // Test cumprod with zero at the end + let device = AutodiffDevice::new(); + let tensor = + TestTensor::<1>::from_data(TensorData::from([2.0, 3.0, 4.0, 0.0]), &device).require_grad(); + + let output = tensor.clone().cumprod(0); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // PyTorch reference: [16.0, 10.0, 6.0, 24.0] + let expected = TensorData::from([16.0, 10.0, 6.0, 24.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +#[ignore = "cumprod gradient with zeros not yet implemented - produces NaN due to division by zero"] +fn should_diff_cumprod_multiple_zeros() { + // Test cumprod with multiple zeros + let device = AutodiffDevice::new(); + let tensor = TestTensor::<1>::from_data(TensorData::from([2.0, 0.0, 3.0, 0.0, 5.0]), &device) + .require_grad(); + + let output = tensor.clone().cumprod(0); + let grads = output.sum().backward(); + let grad = tensor.grad(&grads).unwrap(); + + // PyTorch reference: [1.0, 8.0, 0.0, 0.0, 0.0] + let expected = TensorData::from([1.0, 8.0, 0.0, 0.0, 0.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/cumsum.rs b/crates/burn-backend-tests/tests/autodiff/cumsum.rs new file mode 100644 index 0000000..e0926d3 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/cumsum.rs @@ -0,0 +1,89 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn should_diff_cumsum_dim0() { + let data_1 = TensorData::from([[1.0, 7.0], [-2.0, -3.0]]); + let data_2 = TensorData::from([[4.0, -7.0], [2.0, 3.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_3.cumsum(0); + let tensor_5 = tensor_1.clone().mul(tensor_4); + let grads = tensor_5.sum().backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + // Expected gradients computed with PyTorch + let expected = TensorData::from([[-14.0, 24.0], [17.0, 6.0]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[3.0, 10.0], [-1.0, 37.0]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_cumsum_dim1() { + let data_1 = TensorData::from([[1.0, 7.0], [-2.0, -3.0]]); + let data_2 = TensorData::from([[4.0, -7.0], [2.0, 3.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_3.cumsum(1); + let tensor_5 = tensor_1.clone().mul(tensor_4); + let grads = tensor_5.sum().backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + // Expected gradients computed with PyTorch + let expected = TensorData::from([[1.0, 69.0], [-13.0, -28.0]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[18.0, 13.0], [71.0, 58.0]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_cumsum_complex() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[6.0, 7.0], [9.0, 10.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_3.clone().cumsum(1); + let tensor_5 = tensor_4.mul(tensor_3); + + let grads = tensor_5.sum().backward(); + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + // Expected gradients computed with PyTorch + let expected = TensorData::from([[371.0, 542.0], [2246.0, 3281.0]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[507.0, 528.0], [704.0, 733.0]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/deform_conv2d.rs b/crates/burn-backend-tests/tests/autodiff/deform_conv2d.rs new file mode 100644 index 0000000..2090491 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/deform_conv2d.rs @@ -0,0 +1,1804 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{Shape, module::deform_conv2d, ops::DeformConvOptions}; + +#[test] +fn test_deform_conv2d_basic() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 3, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 1, + offset_groups: 1, + height: 4, + width: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [0.000, 6.0678, 14.2071, 12.2477], + [11.2292, 33.7937, 50.1555, 44.0561], + [17.9294, 57.2174, 85.1505, 79.1840], + [18.0220, 73.6263, 126.8184, 151.6910], + ], + [ + [0.000, 8.9783, 20.7620, 17.7888], + [16.2326, 48.7386, 71.7961, 62.5845], + [25.3808, 80.5195, 119.0949, 110.0938], + [25.0567, 101.8461, 174.3329, 206.6013], + ], + ]], + &device, + ), + offset: TestTensor::from_data( + [[ + [[0.000, 15.0000], [30.000, 45.0000]], + [[0.000, 3.7500], [7.5000, 11.2500]], + [[62.6667, 78.3333], [94.0000, 109.6667]], + [[15.6667, 19.5833], [23.5000, 27.4167]], + [[130.6667, 104.1250], [163.3333, 122.2732]], + [[32.6667, -492.9583], [40.8333, -787.1620]], + [[204.0000, 221.0000], [238.0000, 255.0000]], + [[51.0000, 55.2500], [59.5000, 63.7500]], + [[282.6667, 300.3333], [318.0000, 335.6667]], + [[70.6667, 75.0833], [79.5000, 83.9167]], + [[366.6667, 144.3750], [403.3333, 146.4121]], + [[91.6667, -1788.9860], [100.8333, -2392.7456]], + [[456.0000, 475.0000], [-2718.6250, -2953.2188]], + [[114.0000, 118.7500], [37.7361, 37.4063]], + [[550.6667, 570.3334], [-3404.5139, -3672.5312]], + [[137.6667, 142.5833], [28.6806, 27.5197]], + [[650.6667, 27.9584], [-4174.3657, -59.7509]], + [[162.6667, -3991.0139], [14.4028, -298.7557]], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [ + [0.7029, 2.8356, 5.1067], + [12.7492, 19.4745, 17.8345], + [22.0687, 25.9156, 14.6394], + ], + [ + [3.3696, 12.6134, 19.2671], + [36.7492, 50.5856, 43.5506], + [50.8774, 56.3292, 30.7470], + ], + ], + [ + [ + [0.7029, 2.8356, 5.1067], + [12.7492, 19.4745, 17.8345], + [22.0687, 25.9156, 14.6394], + ], + [ + [3.3696, 12.6134, 19.2671], + [36.7492, 50.5856, 43.5506], + [50.8774, 56.3292, 30.7470], + ], + ], + [ + [ + [0.7029, 2.8356, 5.1067], + [12.7492, 19.4745, 17.8345], + [22.0687, 25.9156, 14.6394], + ], + [ + [3.3696, 12.6134, 19.2671], + [36.7492, 50.5856, 43.5506], + [50.8774, 56.3292, 30.7470], + ], + ], + ], + &device, + ), + mask: TestTensor::from_data( + [[ + [[1303.5000, 1447.8750], [1862.2500, 2006.6250]], + [[1571.1666, 1721.9581], [2154.7500, 2305.5417]], + [[1857.4999, 1396.7151], [2465.9167, 1753.2246]], + [[2315.5000, 2479.1250], [2948.7502, 3112.3750]], + [[2645.1665, 2815.2085], [3303.2500, 3473.2917]], + [[2993.5000, 1150.0625], [3676.4165, 1300.4055]], + [[3531.5000, 3714.3752], [1150.1876, 1148.4744]], + [[3923.1665, 4112.4585], [794.3865, 770.0470]], + [[4333.5000, 181.4101], [368.3260, 4.2679]], + ]], + &device, + ), + bias: TestTensor::from_data([4., 4., 4.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_deform_conv2d_batched() { + let test = Conv2dTestCase { + batch_size: 2, + channels_in: 2, + channels_out: 3, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 1, + offset_groups: 1, + height: 4, + width: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [ + [ + [ + [0.000, 3.4604, 8.7539, 6.8080], + [8.4661, 24.0784, 35.4610, 26.4276], + [19.5988, 51.0406, 68.4389, 53.4993], + [17.4698, 47.9106, 67.3808, 56.6063], + ], + [ + [0.000, 5.1185, 12.7803, 9.8796], + [12.1957, 34.5728, 50.4616, 37.3777], + [27.4521, 71.1227, 94.5778, 73.4724], + [24.1147, 65.8443, 91.8995, 76.7475], + ], + ], + [ + [ + [6.3750, 19.3553, 26.4935, 22.5650], + [17.0026, 57.8088, 85.5580, 78.0746], + [20.7334, 86.5793, 139.4667, 136.4133], + [16.8126, 103.0225, 186.4502, 206.9613], + ], + [ + [9.5625, 28.8786, 39.1137, 32.9178], + [25.1984, 85.0747, 124.6941, 112.5691], + [30.0242, 124.2863, 198.6056, 192.4489], + [23.5826, 143.4660, 257.8752, 283.2587], + ], + ], + ], + &device, + ), + + offset: TestTensor::from_data( + [ + [ + [[0.000, 7.5000], [15.0000, 22.5000]], + [[0.000, 1.8750], [3.7500, 5.6250]], + [[31.3333, 39.1667], [47.0000, 54.8333]], + [[7.8333, 9.7917], [11.7500, 13.7083]], + [[65.3333, 62.7813], [81.6667, 75.4849]], + [[16.3333, -237.8021], [20.4167, -381.7280]], + [[102.0000, 110.5000], [119.0000, 127.5000]], + [[25.5000, 27.6250], [29.7500, 31.8750]], + [[141.3333, 150.1667], [159.0000, 167.8333]], + [[35.3333, 37.5417], [39.7500, 41.9583]], + [[183.3333, 132.3438], [201.6667, 142.0197]], + [[45.8333, -839.6840], [50.4167, -1133.4155]], + [[228.0000, 237.5000], [-1336.1562, -1452.1173]], + [[57.0000, 59.3750], [40.3090, 41.4141]], + [[275.3333, 285.1667], [-1670.5034, -1802.9244]], + [[68.8333, 71.2917], [44.0451, 44.9841]], + [[325.3333, 174.7396], [-2045.1747, -1090.4585]], + [[81.3333, -1844.0659], [46.8090, -1150.2101]], + ], + [ + [[270.000, 277.5000], [285.0000, 292.5000]], + [[67.5000, 69.3750], [71.2500, 73.1250]], + [[313.3333, 321.1667], [329.0000, 336.8333]], + [[78.3333, 80.2917], [82.2500, 84.2083]], + [[359.3333, 130.1563], [375.6667, 130.6099]], + [[89.8333, -4312.7603], [93.9167, -4893.6035]], + [[408.0000, 416.5000], [425.0000, 433.5000]], + [[102.0000, 104.1250], [106.2500, 108.3750]], + [[459.3333, 468.1667], [477.0000, 485.8333]], + [[114.8333, 117.0417], [119.2500, 121.4583]], + [[513.3334, 97.9688], [531.6667, 93.8947]], + [[128.3333, -6720.3926], [132.9167, -7504.5405]], + [[570.000, 579.5000], [-7971.8438, -8251.0850]], + [[142.5000, 144.8750], [22.4965, 21.8203]], + [[629.3333, 639.1667], [-8948.2334, -9249.6641]], + [[157.3333, 159.7917], [15.7743, 14.8695]], + [[691.3333, 14.6145], [-9992.9453, -70.4040]], + [[172.8333, -9818.5234], [7.4132, -352.0222]], + ], + ], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [ + [77.7195, 89.8692, 69.0213], + [121.0760, 137.0775, 92.2989], + [100.0212, 106.5561, 61.1851], + ], + [ + [112.3862, 131.6470, 103.8793], + [177.0760, 200.1887, 138.2681], + [149.5922, 158.7074, 94.3991], + ], + ], + [ + [ + [77.7195, 89.8692, 69.0213], + [121.0760, 137.0775, 92.2989], + [100.0212, 106.5561, 61.1851], + ], + [ + [112.3862, 131.6470, 103.8793], + [177.0760, 200.1887, 138.2681], + [149.5922, 158.7074, 94.3991], + ], + ], + [ + [ + [77.7195, 89.8692, 69.0213], + [121.0760, 137.0775, 92.2989], + [100.0212, 106.5561, 61.1851], + ], + [ + [112.3862, 131.6470, 103.8793], + [177.0760, 200.1887, 138.2681], + [149.5922, 158.7074, 94.3991], + ], + ], + ], + &device, + ), + mask: TestTensor::from_data( + [ + [ + [[1299.7499, 1439.4375], [1849.1249, 1988.8125]], + [[1528.0834, 1673.9791], [2101.8750, 2247.7708]], + [[1771.7500, 1624.9811], [2369.9583, 2099.5039]], + [[2183.7500, 2342.0625], [2806.3750, 2964.6875]], + [[2464.0833, 2628.6042], [3111.1250, 3275.6458]], + [[2759.7500, 1979.2551], [3431.2085, 2390.0286]], + [[3241.7498, 3418.6873], [2415.3589, 2500.8682]], + [[3574.0835, 3757.2292], [2394.3889, 2471.7510]], + [[3921.7500, 2095.5293], [2345.9363, 1199.5048]], + ], + [ + [[5957.2500, 6096.9375], [6506.6250, 6646.3125]], + [[6392.5835, 6538.4790], [6966.3750, 7112.2705]], + [[6843.2500, 2443.8982], [7441.4585, 2550.9199]], + [[7462.2505, 7620.5625], [8084.8745, 8243.1875]], + [[7949.5835, 8114.1045], [8596.6250, 8761.1465]], + [[8452.2500, 1591.6719], [9123.7080, 1589.9454]], + [[9141.2500, 9318.1875], [1414.3584, 1375.1803]], + [[9680.5840, 9863.7285], [949.0560, 897.3544]], + [[10235.2500, 213.4454], [428.2699, 2.4790]], + ], + ], + &device, + ), + bias: TestTensor::from_data([8., 8., 8.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_deform_conv2d_different_kernel_size() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 4, + padding_1: 1, + padding_2: 1, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 1, + offset_groups: 1, + height: 4, + width: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [14.558521, 27.249609, 37.382030, 36.039406], + [33.151936, 60.480656, 81.264656, 78.618156], + [57.520061, 108.623283, 153.413559, 170.072998], + [54.706184, 102.596664, 144.367157, 162.643570], + ], + [ + [25.836353, 48.088451, 65.249161, 62.103317], + [56.805233, 102.995605, 136.983124, 131.120911], + [96.105408, 179.790192, 250.550934, 272.668793], + [90.210945, 167.567917, 232.847275, 257.934692], + ], + ]], + &device, + ), + offset: TestTensor::from_data( + [[ + [ + [0.0e+00, 5.355903e+00, 1.171528e+01], + [3.124999e-01, 8.000000e+00, 1.000000e+01], + [7.500000e-01, 1.400000e+01, 1.600000e+01], + [1.312500e+00, 2.000000e+01, 2.200000e+01], + ], + [ + [0.0e+00, 1.736104e-03, 6.944418e-03], + [1.606250e+01, 2.000000e+00, 2.500000e+00], + [4.425000e+01, 3.500000e+00, 4.000000e+00], + [8.456250e+01, 5.000000e+00, 5.500000e+00], + ], + [ + [6.745834e+01, 7.996479e+01, 9.353048e+01], + [3.166667e+01, 3.377778e+01, 3.588889e+01], + [3.800000e+01, 4.011111e+01, 4.222223e+01], + [4.433333e+01, 4.644444e+01, 4.855556e+01], + ], + [ + [5.277777e-01, 5.955827e-01, 6.670526e-01], + [7.916667e+00, 8.444445e+00, 8.972222e+00], + [9.500000e+00, 1.002778e+01, 1.055556e+01], + [1.108333e+01, 1.161111e+01, 1.213889e+01], + ], + [ + [1.547778e+02, 1.751640e+02, 1.518874e+02], + [6.000000e+01, 6.222223e+01, 4.989969e+01], + [6.666666e+01, 6.888889e+01, 5.432098e+01], + [7.333334e+01, 7.555556e+01, 5.860340e+01], + ], + [ + [2.222223e+00, 2.363040e+00, -3.360339e+01], + [1.500000e+01, 1.555556e+01, -2.277485e+02], + [1.666667e+01, 1.722222e+01, -3.231605e+02], + [1.833333e+01, 1.888889e+01, -4.320448e+02], + ], + [ + [2.641250e+02, 2.021189e+02, 0.0e+00], + [9.100000e+01, 6.481482e+01, 0.0e+00], + [9.800000e+01, 6.863078e+01, 0.0e+00], + [1.050000e+02, 7.230093e+01, 0.0e+00], + ], + [ + [5.250000e+00, -7.268316e+01, 0.0e+00], + [2.275000e+01, -3.346296e+02, 0.0e+00], + [2.450000e+01, -4.611053e+02, 0.0e+00], + [2.625000e+01, -6.017269e+02, 0.0e+00], + ], + [ + [4.400000e+01, 1.197778e+02, 1.222222e+02], + [4.804860e+01, 1.271111e+02, 1.295556e+02], + [5.225000e+01, 1.344444e+02, 1.368889e+02], + [-3.138958e+02, -8.007446e+02, -8.507313e+02], + ], + [ + [3.377778e+02, 2.994445e+01, 3.055556e+01], + [4.848542e+02, 3.177778e+01, 3.238889e+01], + [6.467500e+02, 3.361111e+01, 3.422222e+01], + [4.909653e+02, 2.239892e+01, 2.265992e+01], + ], + [ + [1.533333e+02, 1.558889e+02, 1.584444e+02], + [1.610000e+02, 1.635556e+02, 1.661111e+02], + [1.686667e+02, 1.712222e+02, 1.737778e+02], + [-9.952491e+02, -1.054551e+03, -1.115134e+03], + ], + [ + [3.833333e+01, 3.897222e+01, 3.961111e+01], + [4.025000e+01, 4.088889e+01, 4.152778e+01], + [4.216667e+01, 4.280556e+01, 4.344445e+01], + [2.433767e+01, 2.453511e+01, 2.472810e+01], + ], + [ + [1.920000e+02, 1.946667e+02, 8.907407e+01], + [2.000000e+02, 2.026667e+02, 9.054632e+01], + [2.080000e+02, 2.106667e+02, 9.185186e+01], + [-1.272938e+03, -1.343509e+03, -5.811921e+02], + ], + [ + [4.800000e+01, 4.866667e+01, -7.413704e+02], + [5.000000e+01, 5.066667e+01, -9.788981e+02], + [5.200000e+01, 5.266667e+01, -1.232593e+03], + [2.531250e+01, 2.543518e+01, -6.388311e+02], + ], + [ + [2.333333e+02, 8.772182e+01, 0.0e+00], + [2.416667e+02, 8.827161e+01, 0.0e+00], + [2.500000e+02, 8.864776e+01, 0.0e+00], + [-1.587216e+03, -5.535372e+02, 0.0e+00], + ], + [ + [5.833333e+01, -9.011902e+02, 0.0e+00], + [6.041667e+01, -1.179988e+03, 0.0e+00], + [6.250000e+01, -1.475625e+03, 0.0e+00], + [2.489150e+01, -6.213175e+02, 0.0e+00], + ], + [ + [1.964444e+02, 2.802222e+02, 2.831111e+02], + [2.055625e+02, 2.888889e+02, 2.917778e+02], + [-1.173472e+03, -1.679611e+03, -1.771290e+03], + [0.0e+00, 0.0e+00, 0.0e+00], + ], + [ + [1.144889e+03, 7.005556e+01, 7.077778e+01], + [1.469646e+03, 7.222223e+01, 7.294444e+01], + [5.029167e+02, 2.298823e+01, 2.295062e+01], + [0.0e+00, 0.0e+00, 0.0e+00], + ], + [ + [3.240000e+02, 3.270000e+02, 3.300000e+02], + [3.330000e+02, 3.360000e+02, 3.390000e+02], + [-1.931469e+03, -2.034961e+03, -2.139958e+03], + [0.0e+00, 0.0e+00, 0.0e+00], + ], + [ + [8.100000e+01, 8.175000e+01, 8.250000e+01], + [8.325000e+01, 8.400000e+01, 8.475000e+01], + [1.959376e+01, 1.946614e+01, 1.933334e+01], + [0.0e+00, 0.0e+00, 0.0e+00], + ], + [ + [3.733333e+02, 3.764445e+02, 4.480865e+01], + [3.826667e+02, 3.857778e+02, 4.185955e+01], + [-2.313792e+03, -2.431276e+03, -2.392101e+02], + [0.0e+00, 0.0e+00, 0.0e+00], + ], + [ + [9.333333e+01, 9.411111e+01, -1.904932e+03], + [9.566667e+01, 9.644444e+01, -2.344715e+03], + [1.429166e+01, 1.406212e+01, -3.417283e+02], + [0.0e+00, 0.0e+00, 0.0e+00], + ], + [ + [4.253333e+02, 1.636843e+01, 0.0e+00], + [4.350000e+02, 1.217279e+01, 0.0e+00], + [-2.738517e+03, -4.792887e+01, 0.0e+00], + [0.0e+00, 0.0e+00, 0.0e+00], + ], + [ + [1.063333e+02, -2.178747e+03, 0.0e+00], + [1.087500e+02, -2.670679e+03, 0.0e+00], + [6.947917e+00, -1.629574e+02, 0.0e+00], + [0.0e+00, 0.0e+00, 0.0e+00], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [ + [1.856041, 7.203409, 12.833395, 11.969448], + [24.236776, 40.125511, 41.396423, 27.642044], + [43.613083, 57.508926, 46.093338, 25.174383], + ], + [ + [6.989914, 26.580338, 42.618557, 37.501404], + [75.623192, 116.925674, 113.288368, 72.567764], + [112.724869, 139.826447, 107.653435, 56.799385], + ], + ], + [ + [ + [1.856041, 7.203409, 12.833395, 11.969448], + [24.236776, 40.125511, 41.396423, 27.642044], + [43.613083, 57.508926, 46.093338, 25.174383], + ], + [ + [6.989914, 26.580338, 42.618557, 37.501404], + [75.623192, 116.925674, 113.288368, 72.567764], + [112.724869, 139.826447, 107.653435, 56.799385], + ], + ], + ], + &device, + ), + mask: TestTensor::from_data( + [[ + [ + [0.0e+00, 2.677941e+00, 5.857617e+00], + [4.015623e+01, 7.759999e+02, 8.492499e+02], + [6.637500e+01, 1.067750e+03, 1.141000e+03], + [9.865628e+01, 1.359500e+03, 1.432750e+03], + ], + [ + [6.745831e+01, 7.688924e+01, 8.684974e+01], + [8.387916e+02, 9.161111e+02, 9.934306e+02], + [1.146750e+03, 1.224069e+03, 1.301389e+03], + [1.454708e+03, 1.532028e+03, 1.609347e+03], + ], + [ + [1.547778e+02, 1.716607e+02, 1.460455e+02], + [9.861667e+02, 1.067556e+03, 8.756536e+02], + [1.310333e+03, 1.391722e+03, 1.110864e+03], + [1.634500e+03, 1.715889e+03, 1.339339e+03], + ], + [ + [2.641250e+02, 1.993876e+02, 0.0e+00], + [1.144875e+03, 8.365740e+02, 0.0e+00], + [1.485250e+03, 1.056253e+03, 0.0e+00], + [1.825625e+03, 1.268859e+03, 0.0e+00], + ], + [ + [3.800000e+02, 1.047861e+03, 1.137389e+03], + [5.276354e+02, 1.404444e+03, 1.493972e+03], + [6.826807e+02, 1.761028e+03, 1.850555e+03], + [5.038855e+02, 1.256341e+03, 1.304936e+03], + ], + [ + [1.123500e+03, 1.217097e+03, 1.310694e+03], + [1.496292e+03, 1.589889e+03, 1.683486e+03], + [1.869083e+03, 1.962681e+03, 2.056278e+03], + [1.146700e+03, 1.190136e+03, 1.232930e+03], + ], + [ + [1.300000e+03, 1.397667e+03, 6.512036e+02], + [1.689000e+03, 1.786667e+03, 8.072734e+02], + [2.078000e+03, 2.175667e+03, 9.552593e+02], + [1.060781e+03, 1.097745e+03, 4.656539e+02], + ], + [ + [1.487833e+03, 5.672195e+02, 0.0e+00], + [1.893042e+03, 6.972655e+02, 0.0e+00], + [2.298250e+03, 8.188910e+02, 0.0e+00], + [9.472098e+02, 3.238781e+02, 0.0e+00], + ], + [ + [1.216444e+03, 1.792806e+03, 1.898611e+03], + [1.536448e+03, 2.214222e+03, 2.320028e+03], + [5.177084e+02, 7.256571e+02, 7.493920e+02], + [0.0e+00, 0.0e+00, 0.0e+00], + ], + [ + [1.897500e+03, 2.007375e+03, 2.117250e+03], + [2.335125e+03, 2.445000e+03, 2.554875e+03], + [5.591096e+02, 5.750975e+02, 5.903336e+02], + [0.0e+00, 0.0e+00, 0.0e+00], + ], + [ + [2.119333e+03, 2.233278e+03, 2.654414e+02], + [2.573167e+03, 2.687111e+03, 2.907444e+02], + [3.856317e+02, 3.924502e+02, 3.737657e+01], + [0.0e+00, 0.0e+00, 0.0e+00], + ], + [ + [2.352500e+03, 9.009851e+01, 0.0e+00], + [2.822542e+03, 7.854909e+01, 0.0e+00], + [1.785990e+02, 2.930897e+00, 0.0e+00], + [0.0e+00, 0.0e+00, 0.0e+00], + ], + ]], + &device, + ), + bias: TestTensor::from_data([12., 12.], &device), + }; + test.assert_grads(grads); +} + +#[test] +fn test_deform_conv2d_different_padding() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 2, + padding_2: 3, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 1, + offset_groups: 1, + height: 4, + width: 4, + }; + let device = AutodiffDevice::new(); + let grads = Grads { + x: TestTensor::from_data( + [[ + [ + [60.633026, 60.906506, 61.179493, 61.451954], + [122.557770, 123.088188, 123.618599, 124.149033], + [126.801132, 127.331535, 127.861938, 128.392365], + [131.044434, 131.574875, 132.105286, 132.635712], + ], + [ + [102.000595, 102.497604, 102.993835, 103.489281], + [198.932983, 199.830597, 200.728210, 201.625870], + [206.113968, 207.011627, 207.909256, 208.806870], + [213.294952, 214.192627, 215.090271, 215.987930], + ], + ]], + &device, + ), + // => Position 788: 10.421875 != 10.0546875 + // diff (rel = +1.79e-2, abs = +3.67e-1), tol (rel = +1.00e-2, abs = +9.77e-4) + offset: TestTensor::from_data( + [[ + [ + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [ + 0.0, 0.0, 0.895062, 14.760561, 17.604168, 20.698063, 22.200424, 0.0, + ], + [ + 0.0, 0.0, 0.687500, 9.500000, 10.0, 10.500000, 10.108797, 0.0, + ], + [ + 0.0, 0.0, 1.113426, 13.500000, 14.000000, 14.499999, 13.645835, 0.0, + ], + [ + 0.0, 0.0, 1.613426, 17.500000, 18.000000, 18.500000, 17.108795, 0.0, + ], + [ + 0.0, + 0.0, + -12.395836, + -122.399445, + -130.752319, + -139.355469, + -131.526810, + 0.0, + ], + ], + [ + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [ + 0.0, 0.0, 0.154321, 0.017506, 0.020833, 0.024450, -0.387539, 0.0, + ], + [ + 0.0, 0.0, 24.187502, 2.375000, 2.500000, 2.625000, -37.863422, 0.0, + ], + [ + 0.0, 0.0, 48.057869, 3.375000, 3.500000, 3.625000, -66.770836, 0.0, + ], + [ + 0.0, + 0.0, + 80.02312, + 4.375000, + 4.500000, + 4.625000, + -103.752319, + 0.0, + ], + [ + 0.0, + 0.0, + 113.215271, + 5.107495, + 5.219907, + 5.332031, + -139.725891, + 0.0, + ], + ], + [ + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [ + 0.0, 14.206017, 83.017586, 92.379395, 102.010040, 90.356323, 0.0, 0.0, + ], + [ + 0.0, 6.504737, 35.444443, 35.981483, 36.518517, 29.978970, 0.0, 0.0, + ], + [ + 0.0, 7.668316, 39.740742, 40.277779, 40.814816, 33.071907, 0.0, 0.0, + ], + [ + 0.0, 8.911458, 44.037037, 44.574074, 45.111111, 36.085281, 0.0, 0.0, + ], + [ + 0.0, + -57.523048, + -274.267914, + -289.547089, + -305.095093, + -248.578552, + 0.0, + 0.0, + ], + ], + [ + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [ + 0.0, 9.749230, 0.955354, 0.980994, 1.006945, -13.930464, 0.0, 0.0, + ], + [ + 0.0, + 96.046921, + 8.861111, + 8.995371, + 9.129629, + -129.920715, + 0.0, + 0.0, + ], + [ + 0.0, + 147.434769, + 9.935185, + 10.069445, + 10.203704, + -186.718735, + 0.0, + 0.0, + ], + [ + 0.0, + 207.494781, + 11.009259, + 11.143518, + 11.277778, + -252.188889, + 0.0, + 0.0, + ], + [ + 0.0, + 226.050003, + 10.153355, + 10.252030, + 10.350393, + -266.255280, + 0.0, + 0.0, + ], + ], + [ + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [ + 44.224964, 159.898483, 176.651901, 193.692688, 146.270813, 0.0, 0.0, 0.0, + ], + [ + 19.050755, 64.870377, 65.444443, 66.018517, 46.553150, 0.0, 0.0, 0.0, + ], + [ + 21.049385, 69.462967, 70.037033, 70.611115, 49.104595, 0.0, 0.0, 0.0, + ], + [ + 23.133059, 74.055557, 74.629631, 75.203705, 51.570988, 0.0, 0.0, 0.0, + ], + [ + -141.200272, + -445.302155, + -468.381012, + -491.747223, + -341.553131, + 0.0, + 0.0, + 0.0, + ], + ], + [ + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [ + 35.665298, 3.505739, 3.556735, 3.608062, -48.756947, 0.0, 0.0, 0.0, + ], + [ + 181.404663, + 16.217594, + 16.361111, + 16.504629, + -238.136124, + 0.0, + 0.0, + 0.0, + ], + [ + 263.888885, + 17.365742, + 17.509258, + 17.652779, + -326.403656, + 0.0, + 0.0, + 0.0, + ], + [ + 355.643341, + 18.513889, + 18.657408, + 18.800926, + -423.941345, + 0.0, + 0.0, + 0.0, + ], + [ + 318.709198, + 14.359658, + 14.441552, + 14.523109, + -369.819580, + 0.0, + 0.0, + 0.0, + ], + ], + [ + [ + 0.0, 0.0, 88.846703, 237.478439, 261.731201, 286.289917, 182.508713, 0.0, + ], + [ + 0.0, 0.0, 37.688015, 94.722221, 95.333328, 95.944450, 57.441605, 0.0, + ], + [ + 0.0, 0.0, 40.562500, 99.611107, 100.222229, 100.833336, 59.410744, 0.0, + ], + [ + 0.0, 0.0, 43.527519, 104.500000, 105.111107, 105.722221, 61.289349, 0.0, + ], + [ + 0.0, + 0.0, + -258.324371, + -618.353943, + -649.340271, + -680.632507, + -397.101013, + 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [ + 0.0, + 0.0, + 76.229431, + 7.564093, + 7.641718, + 7.719699, + -102.792252, + 0.0, + ], + [ + 0.0, + 0.0, + 272.015167, + 23.680555, + 23.833332, + 23.986113, + -351.944214, + 0.0, + ], + [ + 0.0, + 0.0, + 386.062500, + 24.902777, + 25.055557, + 25.208334, + -472.147888, + 0.0, + ], + [ + 0.0, + 0.0, + 509.978149, + 26.125000, + 26.277777, + 26.430555, + -602.219971, + 0.0, + ], + [ + 0.0, + 0.0, + 378.410248, + 17.123661, + 17.187500, + 17.250984, + -436.000732, + 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [ + 0.0, 157.623291, 331.938538, 365.283356, 398.952606, 205.988480, 0.0, 0.0, + ], + [ + 0.0, 66.495949, 130.925934, 131.574066, 132.222229, 64.435974, 0.0, 0.0, + ], + [ + 0.0, 70.396835, 136.111115, 136.759262, 137.407410, 65.672256, 0.0, 0.0, + ], + [ + 0.0, 74.393753, 141.296295, 141.944458, 142.592606, 66.812523, 0.0, 0.0, + ], + [ + 0.0, + -432.798035, + -827.492065, + -867.978455, + -908.789368, + -425.074158, + 0.0, + 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [ + 0.0, + 140.150024, + 14.043960, + 14.152921, + 14.262260, + -187.656906, + 0.0, + 0.0, + ], + [ + 0.0, + 386.813873, + 32.731483, + 32.893517, + 33.055557, + -494.779602, + 0.0, + 0.0, + ], + [ + 0.0, + 538.926697, + 34.027779, + 34.189816, + 34.351852, + -653.421875, + 0.0, + 0.0, + ], + [ + 0.0, + 701.505859, + 35.324074, + 35.486115, + 35.648151, + -822.530640, + 0.0, + 0.0, + ], + [ + 0.0, + 416.044586, + 18.903570, + 18.944647, + 18.985338, + -476.728790, + 0.0, + 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [ + 249.876541, 435.868500, 479.178772, 522.832031, 207.919815, 0.0, 0.0, 0.0, + ], + [ + 105.417015, 170.611115, 171.296295, 171.981476, 64.750000, 0.0, 0.0, 0.0, + ], + [ + 110.441696, 176.092590, 176.777771, 177.462952, 65.156044, 0.0, 0.0, 0.0, + ], + [ + 115.567902, 181.574066, 182.259247, 182.944458, 65.460571, 0.0, 0.0, 0.0, + ], + [ + -662.743530, + -1056.641846, + -1107.501953, + -1158.704712, + -409.510162, + 0.0, + 0.0, + 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [ + 227.160507, + 22.982454, + 23.125793, + 23.269531, + -303.112030, + 0.0, + 0.0, + 0.0, + ], + [ + 518.495178, + 42.652779, + 42.824074, + 42.995369, + -657.157410, + 0.0, + 0.0, + 0.0, + ], + [ + 712.252380, + 44.023148, + 44.194443, + 44.365738, + -857.817200, + 0.0, + 0.0, + 0.0, + ], + [ + 917.074036, + 45.393517, + 45.564812, + 45.736115, + -1069.541626, + 0.0, + 0.0, + 0.0, + ], + [ + 416.581482, + 18.997831, + 19.013102, + 19.027966, + -475.031525, + 0.0, + 0.0, + 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [ + 0.0, 0.0, 151.750259, 210.166672, 210.888885, 211.611099, 57.506927, 0.0, + ], + [ + 0.0, 0.0, 157.929276, 215.944443, 216.666672, 217.388901, 57.052204, 0.0, + ], + [ + 0.0, 0.0, 164.215271, 221.722229, 222.444458, 223.166672, 56.490482, 0.0, + ], + [ + 0.0, + 0.0, + -931.783752, + -1285.353760, + -1346.555908, + -1408.119385, + -346.739044, + 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [ + 0.0, + 0.0, + 655.669983, + 52.541668, + 52.722221, + 52.902775, + -824.946777, + 0.0, + ], + [ + 0.0, + 0.0, + 890.972473, + 53.986111, + 54.166668, + 54.347225, + -1067.525024, + 0.0, + ], + [ + 0.0, + 0.0, + 1137.937500, + 55.430557, + 55.611115, + 55.791668, + -1321.765625, + 0.0, + ], + [ + 0.0, + 0.0, + 375.580566, + 17.180984, + 17.169498, + 17.157579, + -425.993713, + 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [ + 0.0, 213.521454, 256.629639, 257.388885, 258.148132, 41.652927, 0.0, 0.0, + ], + [ + 0.0, 221.015625, 262.703705, 263.462982, 264.222229, 40.176598, 0.0, 0.0, + ], + [ + 0.0, 228.622284, 268.777802, 269.537048, 270.296295, 38.587788, 0.0, 0.0, + ], + [ + 0.0, + -1285.466797, + -1554.254517, + -1627.530640, + -1701.186646, + -228.291397, + 0.0, + 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [ + 0.0, + 823.380554, + 64.157410, + 64.347221, + 64.537033, + -1028.532715, + 0.0, + 0.0, + ], + [ + 0.0, + 1107.296509, + 65.675926, + 65.865746, + 66.055557, + -1320.097534, + 0.0, + 0.0, + ], + [ + 0.0, + 1403.473022, + 67.194450, + 67.384262, + 67.574074, + -1623.922974, + 0.0, + 0.0, + ], + [ + 0.0, + 288.151398, + 13.201796, + 13.158524, + 13.114797, + -323.577820, + 0.0, + 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [ + 288.790131, 306.574066, 307.370361, 308.166656, 15.734239, 0.0, 0.0, 0.0, + ], + [ + 297.696838, 312.944427, 313.740723, 314.537048, 13.138914, 0.0, 0.0, 0.0, + ], + [ + 306.721527, 319.314819, 320.111115, 320.907410, 10.425544, 0.0, 0.0, 0.0, + ], + [ + -1711.543213, + -1844.013062, + -1930.236572, + -2016.858643, + -46.846100, + 0.0, + 0.0, + 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [ + 1011.358093, + 76.643517, + 76.842590, + 77.041664, + -1255.045654, + 0.0, + 0.0, + 0.0, + ], + [ + 1347.466431, + 78.236107, + 78.435181, + 78.634262, + -1599.175903, + 0.0, + 0.0, + 0.0, + ], + [ + 1696.433350, + 79.828705, + 80.027779, + 80.226852, + -1956.164917, + 0.0, + 0.0, + 0.0, + ], + [ + 146.703568, + 6.690874, + 6.612756, + 6.534196, + -159.277222, + 0.0, + 0.0, + 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + ]], + &device, + ), + weight: TestTensor::from_data( + [ + [ + [ + [10.341997, 22.988085, 35.634174], + [46.920216, 59.566299, 72.212387], + [80.881615, 92.591522, 104.158524], + ], + [ + [29.213360, 68.837769, 108.462166], + [143.825104, 183.449509, 223.073944], + [228.029373, 256.751740, 283.807098], + ], + ], + [ + [ + [10.341997, 22.988085, 35.634174], + [46.920216, 59.566299, 72.212387], + [80.881615, 92.591522, 104.158524], + ], + [ + [29.213360, 68.837769, 108.462166], + [143.825104, 183.449509, 223.073944], + [228.029373, 256.751740, 283.807098], + ], + ], + ], + &device, + ), + mask: TestTensor::from_data( + [[ + [ + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [ + 0.0, 0.0, 0.447531, 7.380288, 8.802088, 10.349031, 11.100212, 0.0, + ], + [ + 0.0, 0.0, 44.343754, 584.937439, 639.250000, 693.562439, 683.262756, 0.0, + ], + [ + 0.0, 0.0, 68.390068, 803.437561, 857.750000, 912.062500, 874.698059, 0.0, + ], + [ + 0.0, + 0.0, + 96.473381, + 1021.937500, + 1076.250000, + 1130.562500, + 1062.095947, + 0.0, + ], + [ + 0.0, + 0.0, + 121.302101, + 1168.487915, + 1218.373779, + 1268.134888, + 1169.444702, + 0.0, + ], + ], + [ + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [ + 0.0, 13.084491, 75.860909, 83.767761, 91.809029, 80.728188, 0.0, 0.0, + ], + [ + 0.0, 118.950417, 649.486084, 707.821777, 766.157410, 658.076599, 0.0, 0.0, + ], + [ + 0.0, + 170.660782, + 884.171265, + 942.506958, + 1000.842651, + 837.809326, + 0.0, + 0.0, + ], + [ + 0.0, + 226.707260, + 1118.856445, + 1177.192261, + 1235.527710, + 1013.205933, + 0.0, + 0.0, + ], + [ + 0.0, + 234.939651, + 1106.213867, + 1153.415649, + 1200.482666, + 966.248901, + 0.0, + 0.0, + ], + ], + [ + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [ + 42.524002, 153.045700, 168.319275, 183.736511, 138.144653, 0.0, 0.0, 0.0, + ], + [ + 207.319611, 718.432800, 780.791626, 843.150391, 619.975037, 0.0, 0.0, 0.0, + ], + [ + 290.277802, + 969.303223, + 1031.661987, + 1094.020752, + 784.421631, + 0.0, + 0.0, + 0.0, + ], + [ + 377.871063, + 1220.173584, + 1282.532471, + 1344.891235, + 944.233032, + 0.0, + 0.0, + 0.0, + ], + [ + 328.083038, + 1025.494995, + 1069.130615, + 1112.622192, + 766.054932, + 0.0, + 0.0, + 0.0, + ], + ], + [ + [ + 0.0, 0.0, 88.238174, 235.055206, 258.194336, 281.486389, 178.858536, 0.0, + ], + [ + 0.0, 0.0, 305.575500, 789.868042, 856.250061, 922.631897, 572.466064, 0.0, + ], + [ + 0.0, + 0.0, + 421.809021, + 1056.923584, + 1123.305542, + 1189.687500, + 719.598816, + 0.0, + ], + [ + 0.0, + 0.0, + 542.976746, + 1323.979248, + 1390.361206, + 1456.743042, + 861.797302, + 0.0, + ], + [ + 0.0, + 0.0, + 393.291565, + 934.439697, + 974.010376, + 1013.428101, + 586.924011, + 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [ + 0.0, 157.214920, 330.227448, 362.473419, 394.881653, 203.374420, 0.0, 0.0, + ], + [ + 0.0, + 424.340576, + 867.495361, + 937.900452, + 1008.305542, + 505.640503, + 0.0, + 0.0, + ], + [ + 0.0, + 578.894897, + 1150.736084, + 1221.141235, + 1291.546265, + 630.414001, + 0.0, + 0.0, + ], + [ + 0.0, + 738.682495, + 1433.976929, + 1504.381958, + 1574.787109, + 749.954346, + 0.0, + 0.0, + ], + [ + 0.0, 429.912781, 816.507507, 850.771973, 884.873779, 411.152588, 0.0, 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [ + 249.876541, 434.964233, 477.198730, 519.604675, 206.215576, 0.0, 0.0, 0.0, + ], + [ + 560.309326, + 949.520813, + 1023.949097, + 1098.377319, + 422.458344, + 0.0, + 0.0, + 0.0, + ], + [ + 756.768127, + 1248.946777, + 1323.375000, + 1397.803223, + 521.289001, + 0.0, + 0.0, + 0.0, + ], + [ + 958.759216, + 1548.372803, + 1622.800903, + 1697.229248, + 614.587402, + 0.0, + 0.0, + 0.0, + ], + [ + 428.833923, 679.269775, 707.346252, 735.250916, 258.169373, 0.0, 0.0, 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [ + 0.0, + 0.0, + 707.671387, + 1033.687378, + 1112.138916, + 1190.590210, + 328.295044, + 0.0, + ], + [ + 0.0, + 0.0, + 947.779419, + 1349.298584, + 1427.750000, + 1506.201416, + 399.438080, + 0.0, + ], + [ + 0.0, + 0.0, + 1193.718872, + 1664.909668, + 1743.361084, + 1821.812500, + 464.749847, + 0.0, + ], + [ + 0.0, 0.0, 388.737854, 532.503540, 553.962891, 575.241089, 140.658310, 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [ + 0.0, + 880.797302, + 1124.393555, + 1206.868042, + 1289.342651, + 209.627625, + 0.0, + 0.0, + ], + [ + 0.0, + 1169.882812, + 1456.189819, + 1538.664429, + 1621.138916, + 247.754730, + 0.0, + 0.0, + ], + [ + 0.0, + 1465.098755, + 1787.986084, + 1870.460571, + 1952.935181, + 279.751526, + 0.0, + 0.0, + ], + [ + 0.0, 297.330719, 356.362152, 369.893524, 383.234344, 50.974621, 0.0, 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + [ + [ + 1074.567993, + 1219.497681, + 1305.995361, + 1392.493042, + 71.162437, + 0.0, + 0.0, + 0.0, + ], + [ + 1416.214722, + 1567.479126, + 1653.976929, + 1740.474609, + 72.689949, + 0.0, + 0.0, + 0.0, + ], + [ + 1764.290771, + 1915.460571, + 2001.958496, + 2088.456055, + 67.787628, + 0.0, + 0.0, + 0.0, + ], + [ + 151.018372, 160.055023, 164.776138, 169.298447, 3.865937, 0.0, 0.0, 0.0, + ], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ], + ]], + &device, + ), + bias: TestTensor::from_data([48., 48.], &device), + }; + test.assert_grads(grads); +} + +struct Conv2dTestCase { + batch_size: usize, + channels_in: usize, + channels_out: usize, + kernel_size_1: usize, + kernel_size_2: usize, + padding_1: usize, + padding_2: usize, + stride_1: usize, + stride_2: usize, + dilation_1: usize, + dilation_2: usize, + groups: usize, + offset_groups: usize, + height: usize, + width: usize, +} + +struct Grads { + x: TestTensor<4>, + offset: TestTensor<4>, + weight: TestTensor<4>, + mask: TestTensor<4>, + bias: TestTensor<1>, +} + +impl Conv2dTestCase { + fn assert_grads(self, expected_grads: Grads) { + let out_height = + (self.height + 2 * self.padding_1 - self.dilation_1 * (self.kernel_size_1 - 1) - 1) + / self.stride_1 + + 1; + let out_width = + (self.width + 2 * self.padding_2 - self.dilation_2 * (self.kernel_size_2 - 1) - 1) + / self.stride_2 + + 1; + + let shape_x = Shape::new([self.batch_size, self.channels_in, self.height, self.width]); + let shape_offset = Shape::new([ + self.batch_size, + 2 * self.offset_groups * self.kernel_size_1 * self.kernel_size_2, + out_height, + out_width, + ]); + let shape_weight = Shape::new([ + self.channels_out, + self.channels_in / self.groups, + self.kernel_size_1, + self.kernel_size_2, + ]); + let shape_mask = Shape::new([ + self.batch_size, + self.offset_groups * self.kernel_size_1 * self.kernel_size_2, + out_height, + out_width, + ]); + let device = AutodiffDevice::new(); + let weight = TestTensor::from_data( + TestTensorInt::arange(0..shape_weight.num_elements() as i64, &device) + .reshape::<4, _>(shape_weight) + .into_data(), + &device, + ) + .require_grad(); + let bias = TestTensor::from_data( + TestTensorInt::arange(0..self.channels_out as i64, &device).into_data(), + &device, + ) + .require_grad(); + let x = TestTensor::from_data( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<4, _>(shape_x) + .into_data(), + &device, + ) + .require_grad(); + let offset = TestTensor::from_data( + TestTensorInt::arange(0..shape_offset.num_elements() as i64, &device) + .reshape::<4, _>(shape_offset.clone()) + .into_data(), + &device, + ) + .div_scalar(shape_offset.num_elements() as f32) + .require_grad(); + + let mask = TestTensor::from_data( + TestTensorInt::arange(0..shape_mask.num_elements() as i64, &device) + .reshape::<4, _>(shape_mask.clone()) + .into_data(), + &device, + ) + .div_scalar(shape_mask.num_elements() as f32) + .require_grad(); + + let output = deform_conv2d( + x.clone(), + offset.clone(), + weight.clone(), + Some(mask.clone()), + Some(bias.clone()), + DeformConvOptions::new( + [self.stride_1, self.stride_2], + [self.padding_1, self.padding_2], + [self.dilation_1, self.dilation_2], + self.groups, + self.offset_groups, + ), + ); + let grads = output.backward(); + + // Assert + let x_grad_actual = x.grad(&grads).unwrap(); + let offset_grad_actual = offset.grad(&grads).unwrap(); + let weight_grad_actual = weight.grad(&grads).unwrap(); + let mask_grad_actual = mask.grad(&grads).unwrap(); + let bias_grad_actual = bias.grad(&grads).unwrap(); + + // Relative is set to 5%, which is much higher than typical numerical test tolerances. + // This is due to the complexity of the deformable convolution operation. + // Unlike regular conv2d, which samples from fixed integer grid positions, + // deformable conv2d samples input values at fractional offset locations (learned offsets). + // These non-integer positions require bilinear interpolation to estimate the input value. + // Gradients computed through all these floating-point operations can compound numerical differences. + let tolerance = Tolerance::relative(0.5); + + println!("Testing bias"); + expected_grads + .bias + .to_data() + .assert_approx_eq::(&bias_grad_actual.to_data(), tolerance); + println!("Testing input"); + expected_grads + .x + .to_data() + .assert_approx_eq::(&x_grad_actual.to_data(), tolerance); + println!("Testing offset"); + expected_grads + .offset + .to_data() + .assert_approx_eq::(&offset_grad_actual.to_data(), tolerance); + println!("Testing mask"); + expected_grads + .mask + .to_data() + .assert_approx_eq::(&mask_grad_actual.to_data(), tolerance); + println!("Testing weight"); + expected_grads + .weight + .to_data() + .assert_approx_eq::(&weight_grad_actual.to_data(), tolerance); + } +} diff --git a/crates/burn-backend-tests/tests/autodiff/div.rs b/crates/burn-backend-tests/tests/autodiff/div.rs new file mode 100644 index 0000000..04a480d --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/div.rs @@ -0,0 +1,105 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn should_diff_div() { + let data_1 = TensorData::from([1.0, 7.0]); + let data_2 = TensorData::from([4.0, 7.0]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<1>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().div(tensor_2.clone()); + let grads = tensor_3.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([0.25, 0.14285715]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([-0.0625, -0.14285715]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_div_scalar() { + let data = TensorData::from([1.0, 7.0]); + + let tensor = TestTensor::<1>::from_data(data, &AutodiffDevice::new()).require_grad(); + let tensor_out = tensor.clone().div_scalar(4.0); + + let grads = tensor_out.backward(); + let grad = tensor.grad(&grads).unwrap(); + + grad.to_data() + .assert_eq(&TensorData::from([0.25, 0.25]), false); +} + +#[test] +fn test_div_complex_1() { + let data_1 = TensorData::from([[1.0, 7.0], [13.0, -3.0]]); + let data_2 = TensorData::from([[4.0, 7.0], [2.0, 3.0]]); + let data_3 = TensorData::from([[2.0, 2.0], [2.0, 2.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + let tensor_3 = TestTensor::from_data(data_3, &device).require_grad(); + + let tensor_4 = tensor_1.clone().div(tensor_2.clone()); + let tensor_5 = tensor_4.div(tensor_3.clone()); + + let grads = tensor_5.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + let grad_3 = tensor_3.grad(&grads).unwrap(); + + let expected = TensorData::from([[0.1250, 0.07142857], [0.25, 0.16666667]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[-0.03125, -0.07142857], [-1.6250, 0.16666667]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + let expected = TensorData::from([[-0.0625, -0.25], [-1.6250, 0.25]]); + grad_3 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_div_complex_2() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[6.0, 7.0], [9.0, 10.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_3.div(tensor_2.clone()); + + let grads = tensor_4.backward(); + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let tolerance = Tolerance::default().set_half_precision_absolute(2e-3); + let expected = TensorData::from([[2.00, 2.92857146], [1.36666667, 2.0]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([[0.08333334, 0.09591837], [-0.05555558, -0.06714284]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, tolerance); +} diff --git a/crates/burn-backend-tests/tests/autodiff/erf.rs b/crates/burn-backend-tests/tests/autodiff/erf.rs new file mode 100644 index 0000000..f29132f --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/erf.rs @@ -0,0 +1,29 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn should_diff_erf() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[6.0, 7.0], [9.0, 10.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().erf()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[32.0, 32.0], [32.0, 32.0]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[8.0, 8.0], [8.0, 8.0]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/exp.rs b/crates/burn-backend-tests/tests/autodiff/exp.rs new file mode 100644 index 0000000..f6b9c4f --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/exp.rs @@ -0,0 +1,29 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn should_diff_exp() { + let data_1 = TensorData::from([[1.0, 7.0], [-2.0, -3.0]]); + let data_2 = TensorData::from([[4.0, -7.0], [2.0, 3.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().exp()); + let grads = tensor_3.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let tolerance = Tolerance::default(); + let expected = TensorData::from([[54.5991, 27.4746], [54.5991, 27.4746]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([[-5.4598e+01, -9.1188e-04], [2.9556e+01, 8.0342e+01]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, tolerance); +} diff --git a/crates/burn-backend-tests/tests/autodiff/expand.rs b/crates/burn-backend-tests/tests/autodiff/expand.rs new file mode 100644 index 0000000..5fff88c --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/expand.rs @@ -0,0 +1,39 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_diff_expand() { + // Python code to generate the test case values + // import torch + // x1 = torch.tensor([4.0, 7.0, 2.0, 3.0], requires_grad=True) + // x2 = torch.tensor([2.0, 4.5, 7.0, 3.0], requires_grad=True) + // y = x1.expand(4, 4) + // z = (x2 * y).sum() + // z.backward() + // print("x1", x1.grad) + // print("x2", x2.grad) + + let device = AutodiffDevice::new(); + + let data_1 = TensorData::from([4.0, 7.0, 2.0, 3.0]); + let tensor_1 = TestTensor::<1>::from_data(data_1, &device).require_grad(); + + let data_2 = TensorData::from([2.0, 4.5, 7.0, 3.0]); + let tensor_2 = TestTensor::<1>::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().expand([4, 4]); + + // Use unsqueeze to make tensor_2 have the same shape as tensor_3 + let tensor_4 = tensor_2.clone().unsqueeze().mul(tensor_3).sum(); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([8., 18., 28., 12.]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([16., 28., 8., 12.]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/flip.rs b/crates/burn-backend-tests/tests/autodiff/flip.rs new file mode 100644 index 0000000..9cea89e --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/flip.rs @@ -0,0 +1,29 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_diff_flip() { + let data_1 = TensorData::from([[[1.0, 7.0], [2.0, 3.0]]]); // 1x2x2 + let data_2 = TensorData::from([[[3.0, 2.0, 7.0], [3.0, 3.2, 1.0]]]); // 1x2x3 + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<3>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_2.clone().flip([1, 2]); + let tensor_4 = tensor_1.clone().matmul(tensor_3); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let tolerance = Tolerance::default().set_half_precision_relative(1e-3); + grad_1 + .into_data() + .assert_approx_eq::(&TensorData::from([[[7.2, 12.0], [7.2, 12.0]]]), tolerance); // 1x2x2 + grad_2.into_data().assert_approx_eq::( + &TensorData::from([[[10.0, 10.0, 10.0], [3.0, 3.0, 3.0]]]), + tolerance, + ); // 1x2x3 +} diff --git a/crates/burn-backend-tests/tests/autodiff/floor.rs b/crates/burn-backend-tests/tests/autodiff/floor.rs new file mode 100644 index 0000000..ec18491 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/floor.rs @@ -0,0 +1,21 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_diff_floor() { + let data = TensorData::from([ + [-1.9751, 0.0714, 0.0643, 0.2406], + [-1.3172, 0.1252, -0.1119, -0.0127], + ]); + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data, &device).require_grad(); + let tensor_2 = tensor_1.clone().floor(); + let grads = tensor_2.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + + grad_1.to_data().assert_eq( + &TensorData::from([[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]]), + false, + ); +} diff --git a/crates/burn-backend-tests/tests/autodiff/gather_scatter.rs b/crates/burn-backend-tests/tests/autodiff/gather_scatter.rs new file mode 100644 index 0000000..c03b48c --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/gather_scatter.rs @@ -0,0 +1,94 @@ +use super::*; +use burn_tensor::{IndexingUpdateOp, TensorData}; + +#[test] +fn test_gather_grad() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::from_data( + TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]), + &device, + ) + .require_grad(); + let indices = TestTensorInt::<2>::from_data( + TensorData::from([[2, 1, 0, 1, 2], [1, 0, 2, 1, 0]]), + &device, + ); + + let tensor_2 = tensor_1.clone().matmul(tensor_1.clone().transpose()); + let tensor_3 = tensor_1.clone().gather(1, indices); + let tensor_4 = tensor_2.matmul(tensor_3); + + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + + grad_1.to_data().assert_eq( + &TensorData::from([[94., 150., 187.], [242., 305., 304.]]), + false, + ); +} + +#[test] +fn test_scatter_grad() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::from_data( + TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]), + &device, + ) + .require_grad(); + let values = TestTensor::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), + &device, + ) + .require_grad(); + let indices = TestTensorInt::<2>::from_data(TensorData::from([[2, 1, 0], [2, 0, 1]]), &device); + + let tensor_2 = tensor_1.clone().matmul(tensor_1.clone().transpose()); + let tensor_3 = tensor_1 + .clone() + .scatter(1, indices, values.clone(), IndexingUpdateOp::Add); + let tensor_4 = tensor_2.matmul(tensor_3); + + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = values.grad(&grads).unwrap(); + + grad_1.to_data().assert_eq( + &TensorData::from([[127., 181., 235.], [226., 316., 406.]]), + false, + ); + grad_2 + .to_data() + .assert_eq(&TensorData::from([[19., 19., 19.], [64., 64., 64.]]), false); +} + +#[test] +fn test_scatter_add_grad_partial_indices() { + let device = AutodiffDevice::new(); + let tensor_1 = + TestTensor::from_data(TensorData::from([[0.0, 1.0, 2.0, 3.0, 4.0, 5.0]]), &device) + .require_grad(); + let tensor_2 = + TestTensor::from_data(TensorData::from([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]]), &device) + .require_grad(); + let values = TestTensor::from_data(TensorData::from([[4.0, 5.0, 6.0]]), &device).require_grad(); + let indices = TestTensorInt::<2>::from_data(TensorData::from([[2, 1, 0]]), &device); + + let tensor_3 = tensor_1.clone().mul(tensor_2); + let tensor_4 = tensor_3 + .clone() + .scatter(1, indices, values.clone(), IndexingUpdateOp::Add); + + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = values.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[1., 2., 3., 4., 5., 6.]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([[1., 1., 1.]]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/gather_scatter_nd.rs b/crates/burn-backend-tests/tests/autodiff/gather_scatter_nd.rs new file mode 100644 index 0000000..3bbf7a3 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/gather_scatter_nd.rs @@ -0,0 +1,453 @@ +use super::*; +use burn_tensor::{IndexingUpdateOp, TensorData}; + +#[test] +fn test_scatter_nd_add_grad() { + let device = AutodiffDevice::new(); + let data: TestTensor<2> = TestTensor::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), + &device, + ) + .require_grad(); + let values: TestTensor<2> = + TestTensor::from_data(TensorData::from([[10.0, 20.0, 30.0]]), &device).require_grad(); + let indices = TestTensorInt::<2>::from_data(TensorData::from([[1]]), &device); + + // scatter_nd_add: data[1, :] += values[0, :] + let result: TestTensor<2> = + data.clone() + .scatter_nd(indices, values.clone(), IndexingUpdateOp::Add); + let grads = result.sum().backward(); + + let grad_data = data.grad(&grads).unwrap(); + let grad_values = values.grad(&grads).unwrap(); + + // grad_data = all ones (identity for add) + grad_data.to_data().assert_eq( + &TensorData::from([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]), + false, + ); + // grad_values = gather_nd(ones, indices) = ones at row 1 + grad_values + .to_data() + .assert_eq(&TensorData::from([[1.0, 1.0, 1.0]]), false); +} + +#[test] +fn test_scatter_nd_assign_grad() { + let device = AutodiffDevice::new(); + let data: TestTensor<2> = TestTensor::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), + &device, + ) + .require_grad(); + let values: TestTensor<2> = + TestTensor::from_data(TensorData::from([[10.0, 20.0, 30.0]]), &device).require_grad(); + let indices = TestTensorInt::<2>::from_data(TensorData::from([[1]]), &device); + + // scatter_nd (assign): data[1, :] = values[0, :] + let result: TestTensor<2> = + data.clone() + .scatter_nd(indices, values.clone(), IndexingUpdateOp::Assign); + let grads = result.sum().backward(); + + let grad_data = data.grad(&grads).unwrap(); + let grad_values = values.grad(&grads).unwrap(); + + // grad_data: all ones except row 1 is zeroed out (overwritten positions get no gradient) + grad_data.to_data().assert_eq( + &TensorData::from([[1.0, 1.0, 1.0], [0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]), + false, + ); + // grad_values = gather_nd(ones, indices) = ones at row 1 + grad_values + .to_data() + .assert_eq(&TensorData::from([[1.0, 1.0, 1.0]]), false); +} + +#[test] +fn test_gather_nd_grad() { + let device = AutodiffDevice::new(); + let data: TestTensor<2> = TestTensor::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), + &device, + ) + .require_grad(); + let indices = TestTensorInt::<2>::from_data(TensorData::from([[0], [2]]), &device); + + // gather_nd: extract rows 0 and 2 + let result: TestTensor<2> = data.clone().gather_nd(indices); + let grads = result.sum().backward(); + + let grad_data = data.grad(&grads).unwrap(); + + // grad_data: scatter_nd(zeros, indices, ones, Add) -> row 0 and 2 get 1s, row 1 stays 0 + grad_data.to_data().assert_eq( + &TensorData::from([[1.0, 1.0, 1.0], [0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]), + false, + ); +} + +#[test] +fn test_scatter_nd_add_grad_k_equals_d() { + // K=D: scalar-level indexing (each index tuple selects a single element) + let device = AutodiffDevice::new(); + let data: TestTensor<2> = + TestTensor::from_data(TensorData::from([[1.0, 2.0], [3.0, 4.0]]), &device).require_grad(); + let values: TestTensor<1> = + TestTensor::from_data(TensorData::from([10.0, 20.0]), &device).require_grad(); + // indices shape: [2, 2] => K=2=D, M=2, DV = 1+2-2 = 1 + let indices = TestTensorInt::<2>::from_data(TensorData::from([[0, 1], [1, 0]]), &device); + + let result: TestTensor<2> = + data.clone() + .scatter_nd(indices, values.clone(), IndexingUpdateOp::Add); + let grads = result.sum().backward(); + + let grad_data = data.grad(&grads).unwrap(); + let grad_values = values.grad(&grads).unwrap(); + + // grad_data = all ones (identity for add) + grad_data + .to_data() + .assert_eq(&TensorData::from([[1.0, 1.0], [1.0, 1.0]]), false); + // grad_values = gather_nd(ones, indices) = [1.0, 1.0] + grad_values + .to_data() + .assert_eq(&TensorData::from([1.0, 1.0]), false); +} + +#[test] +fn test_gather_nd_grad_k_equals_d() { + // K=D: scalar-level gather + let device = AutodiffDevice::new(); + let data: TestTensor<2> = + TestTensor::from_data(TensorData::from([[1.0, 2.0], [3.0, 4.0]]), &device).require_grad(); + // indices shape: [2, 2] => K=2=D + let indices = TestTensorInt::<2>::from_data(TensorData::from([[0, 1], [1, 0]]), &device); + + let result: TestTensor<1> = data.clone().gather_nd(indices); + let grads = result.sum().backward(); + + let grad_data = data.grad(&grads).unwrap(); + + // grad = scatter_nd(zeros, indices, ones, Add) + // [0,1] => data[0][1] gets 1, [1,0] => data[1][0] gets 1 + grad_data + .to_data() + .assert_eq(&TensorData::from([[0.0, 1.0], [1.0, 0.0]]), false); +} + +#[test] +fn test_scatter_nd_mul_grad() { + let device = AutodiffDevice::new(); + let data: TestTensor<2> = TestTensor::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), + &device, + ) + .require_grad(); + let values: TestTensor<2> = + TestTensor::from_data(TensorData::from([[10.0, 20.0, 30.0]]), &device).require_grad(); + let indices = TestTensorInt::<2>::from_data(TensorData::from([[1]]), &device); + + // scatter_nd_mul: data[1, :] *= values[0, :] + let result: TestTensor<2> = + data.clone() + .scatter_nd(indices, values.clone(), IndexingUpdateOp::Mul); + let grads = result.sum().backward(); + + let grad_data = data.grad(&grads).unwrap(); + let grad_values = values.grad(&grads).unwrap(); + + // grad_data = grad * scatter_nd(ones, indices, values, Assign) + // row 1 gets the values, others stay at 1. + grad_data.to_data().assert_eq( + &TensorData::from([[1.0, 1.0, 1.0], [10.0, 20.0, 30.0], [1.0, 1.0, 1.0]]), + false, + ); + // grad_values = gather_nd(grad, indices) * gather_nd(data, indices) + // = ones * data[1, :] = [4.0, 5.0, 6.0] + grad_values + .to_data() + .assert_eq(&TensorData::from([[4.0, 5.0, 6.0]]), false); +} + +#[test] +fn test_scatter_nd_max_grad_data_wins() { + // values < data at scattered positions: gradient flows fully to data. + let device = AutodiffDevice::new(); + let data: TestTensor<2> = TestTensor::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), + &device, + ) + .require_grad(); + let values: TestTensor<2> = + TestTensor::from_data(TensorData::from([[1.0, 2.0, 3.0]]), &device).require_grad(); + let indices = TestTensorInt::<2>::from_data(TensorData::from([[1]]), &device); + + let result: TestTensor<2> = + data.clone() + .scatter_nd(indices, values.clone(), IndexingUpdateOp::Max); + let grads = result.sum().backward(); + + let grad_data = data.grad(&grads).unwrap(); + let grad_values = values.grad(&grads).unwrap(); + + // data wins everywhere on row 1, so grad_data = ones (full flow). + grad_data.to_data().assert_eq( + &TensorData::from([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]), + false, + ); + // values lose, so grad_values is zero. + grad_values + .to_data() + .assert_eq(&TensorData::from([[0.0, 0.0, 0.0]]), false); +} + +#[test] +fn test_scatter_nd_max_grad_values_win() { + // values > data at scattered positions: grad_data zeroed at scattered positions. + let device = AutodiffDevice::new(); + let data: TestTensor<2> = TestTensor::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), + &device, + ) + .require_grad(); + let values: TestTensor<2> = + TestTensor::from_data(TensorData::from([[10.0, 20.0, 30.0]]), &device).require_grad(); + let indices = TestTensorInt::<2>::from_data(TensorData::from([[1]]), &device); + + let result: TestTensor<2> = + data.clone() + .scatter_nd(indices, values.clone(), IndexingUpdateOp::Max); + let grads = result.sum().backward(); + + let grad_data = data.grad(&grads).unwrap(); + let grad_values = values.grad(&grads).unwrap(); + + // grad_data zeroed at row 1 (values won), ones elsewhere. + grad_data.to_data().assert_eq( + &TensorData::from([[1.0, 1.0, 1.0], [0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]), + false, + ); + grad_values + .to_data() + .assert_eq(&TensorData::from([[1.0, 1.0, 1.0]]), false); +} + +#[test] +fn test_scatter_nd_max_grad_tie() { + // Mixed: column 0 -> data wins, column 1 -> tie (both get grad), column 2 -> values wins. + let device = AutodiffDevice::new(); + let data: TestTensor<2> = TestTensor::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), + &device, + ) + .require_grad(); + let values: TestTensor<2> = + TestTensor::from_data(TensorData::from([[1.0, 5.0, 10.0]]), &device).require_grad(); + let indices = TestTensorInt::<2>::from_data(TensorData::from([[1]]), &device); + + let result: TestTensor<2> = + data.clone() + .scatter_nd(indices, values.clone(), IndexingUpdateOp::Max); + let grads = result.sum().backward(); + + let grad_data = data.grad(&grads).unwrap(); + let grad_values = values.grad(&grads).unwrap(); + + // Row 1: data_won = [T, T, F], values_won = [F, T, T] (tie at col 1 contributes to both). + grad_data.to_data().assert_eq( + &TensorData::from([[1.0, 1.0, 1.0], [1.0, 1.0, 0.0], [1.0, 1.0, 1.0]]), + false, + ); + grad_values + .to_data() + .assert_eq(&TensorData::from([[0.0, 1.0, 1.0]]), false); +} + +#[test] +fn test_scatter_nd_min_grad() { + // Min mirror of test_scatter_nd_max_grad_values_win: values < data => values win for Min. + let device = AutodiffDevice::new(); + let data: TestTensor<2> = TestTensor::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), + &device, + ) + .require_grad(); + let values: TestTensor<2> = + TestTensor::from_data(TensorData::from([[1.0, 2.0, 3.0]]), &device).require_grad(); + let indices = TestTensorInt::<2>::from_data(TensorData::from([[1]]), &device); + + let result: TestTensor<2> = + data.clone() + .scatter_nd(indices, values.clone(), IndexingUpdateOp::Min); + let grads = result.sum().backward(); + + let grad_data = data.grad(&grads).unwrap(); + let grad_values = values.grad(&grads).unwrap(); + + grad_data.to_data().assert_eq( + &TensorData::from([[1.0, 1.0, 1.0], [0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]), + false, + ); + grad_values + .to_data() + .assert_eq(&TensorData::from([[1.0, 1.0, 1.0]]), false); +} + +#[test] +fn test_scatter_nd_min_grad_tie() { + // Mixed: column 0 -> values win, column 1 -> tie, column 2 -> data wins. + let device = AutodiffDevice::new(); + let data: TestTensor<2> = TestTensor::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), + &device, + ) + .require_grad(); + let values: TestTensor<2> = + TestTensor::from_data(TensorData::from([[1.0, 5.0, 10.0]]), &device).require_grad(); + let indices = TestTensorInt::<2>::from_data(TensorData::from([[1]]), &device); + + let result: TestTensor<2> = + data.clone() + .scatter_nd(indices, values.clone(), IndexingUpdateOp::Min); + let grads = result.sum().backward(); + + let grad_data = data.grad(&grads).unwrap(); + let grad_values = values.grad(&grads).unwrap(); + + // Row 1: data_won = [F, T, T], values_won = [T, T, F]. + grad_data.to_data().assert_eq( + &TensorData::from([[1.0, 1.0, 1.0], [0.0, 1.0, 1.0], [1.0, 1.0, 1.0]]), + false, + ); + grad_values + .to_data() + .assert_eq(&TensorData::from([[1.0, 1.0, 0.0]]), false); +} + +#[test] +fn test_scatter_nd_max_grad_k_equals_d() { + // K=D scalar-level Max: each index targets one element. + let device = AutodiffDevice::new(); + let data: TestTensor<2> = + TestTensor::from_data(TensorData::from([[1.0, 2.0], [3.0, 4.0]]), &device).require_grad(); + let values: TestTensor<1> = + TestTensor::from_data(TensorData::from([10.0, 1.0]), &device).require_grad(); + let indices = TestTensorInt::<2>::from_data(TensorData::from([[0, 1], [1, 0]]), &device); + + let result: TestTensor<2> = + data.clone() + .scatter_nd(indices, values.clone(), IndexingUpdateOp::Max); + let grads = result.sum().backward(); + + let grad_data = data.grad(&grads).unwrap(); + let grad_values = values.grad(&grads).unwrap(); + + grad_data + .to_data() + .assert_eq(&TensorData::from([[1.0, 0.0], [1.0, 1.0]]), false); + grad_values + .to_data() + .assert_eq(&TensorData::from([1.0, 0.0]), false); +} + +#[test] +fn test_scatter_nd_mul_grad_k_equals_d() { + // K=D scalar-level Mul: each index targets one element. + let device = AutodiffDevice::new(); + let data: TestTensor<2> = + TestTensor::from_data(TensorData::from([[1.0, 2.0], [3.0, 4.0]]), &device).require_grad(); + let values: TestTensor<1> = + TestTensor::from_data(TensorData::from([10.0, 20.0]), &device).require_grad(); + // Multiply data[0,1] by 10 and data[1,0] by 20. + let indices = TestTensorInt::<2>::from_data(TensorData::from([[0, 1], [1, 0]]), &device); + + let result: TestTensor<2> = + data.clone() + .scatter_nd(indices, values.clone(), IndexingUpdateOp::Mul); + let grads = result.sum().backward(); + + let grad_data = data.grad(&grads).unwrap(); + let grad_values = values.grad(&grads).unwrap(); + + // grad_data[0,1] = 10 (multiplier), grad_data[1,0] = 20, others = 1. + grad_data + .to_data() + .assert_eq(&TensorData::from([[1.0, 10.0], [20.0, 1.0]]), false); + // grad_values = data at [0,1]=2, [1,0]=3. + grad_values + .to_data() + .assert_eq(&TensorData::from([2.0, 3.0]), false); +} + +#[test] +fn test_scatter_nd_mul_grad_only_data_tracked() { + let device = AutodiffDevice::new(); + let data: TestTensor<2> = + TestTensor::from_data(TensorData::from([[1.0, 2.0], [3.0, 4.0]]), &device).require_grad(); + let values: TestTensor<1> = TestTensor::from_data(TensorData::from([10.0]), &device); + let indices = TestTensorInt::<2>::from_data(TensorData::from([[0, 1]]), &device); + + let result: TestTensor<2> = data + .clone() + .scatter_nd(indices, values, IndexingUpdateOp::Mul); + let grads = result.sum().backward(); + + let grad_data = data.grad(&grads).unwrap(); + + grad_data + .to_data() + .assert_eq(&TensorData::from([[1.0, 10.0], [1.0, 1.0]]), false); +} + +#[test] +fn test_scatter_nd_mul_grad_only_values_tracked() { + let device = AutodiffDevice::new(); + let data: TestTensor<2> = + TestTensor::from_data(TensorData::from([[1.0, 2.0], [3.0, 4.0]]), &device); + let values: TestTensor<1> = + TestTensor::from_data(TensorData::from([10.0]), &device).require_grad(); + let indices = TestTensorInt::<2>::from_data(TensorData::from([[0, 1]]), &device); + + let result: TestTensor<2> = data.scatter_nd(indices, values.clone(), IndexingUpdateOp::Mul); + let grads = result.sum().backward(); + + let grad_values = values.grad(&grads).unwrap(); + + grad_values + .to_data() + .assert_eq(&TensorData::from([2.0]), false); +} + +#[test] +fn test_scatter_nd_assign_grad_k_equals_d() { + // K=D: scalar-level assign, each index overwrites a single element + let device = AutodiffDevice::new(); + let data: TestTensor<2> = + TestTensor::from_data(TensorData::from([[1.0, 2.0], [3.0, 4.0]]), &device).require_grad(); + let values: TestTensor<1> = + TestTensor::from_data(TensorData::from([10.0, 20.0]), &device).require_grad(); + // indices shape: [2, 2] => K=2=D + // Overwrite data[0,1] and data[1,0] + let indices = TestTensorInt::<2>::from_data(TensorData::from([[0, 1], [1, 0]]), &device); + + let result: TestTensor<2> = + data.clone() + .scatter_nd(indices, values.clone(), IndexingUpdateOp::Assign); + let grads = result.sum().backward(); + + let grad_data = data.grad(&grads).unwrap(); + let grad_values = values.grad(&grads).unwrap(); + + // grad_data: ones everywhere except the overwritten positions are zeroed + // [0,1] and [1,0] were overwritten, so those get 0 + grad_data + .to_data() + .assert_eq(&TensorData::from([[1.0, 0.0], [0.0, 1.0]]), false); + // grad_values = gather_nd(ones, indices) = [1.0, 1.0] + grad_values + .to_data() + .assert_eq(&TensorData::from([1.0, 1.0]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/gelu.rs b/crates/burn-backend-tests/tests/autodiff/gelu.rs new file mode 100644 index 0000000..a875535 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/gelu.rs @@ -0,0 +1,27 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance, activation}; + +#[test] +fn should_diff_gelu() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data([[0.0, 1.0], [-3.0, 4.0]], &device).require_grad(); + let tensor_2 = TestTensor::from_data([[6.0, -0.5], [9.0, 10.0]], &device).require_grad(); + + let x = tensor_1.clone().matmul(activation::gelu(tensor_2.clone())); + let x = tensor_1.clone().matmul(x); + let grads = x.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let tolerance = Tolerance::permissive(); + let expected = TensorData::from([[1.46281, 1.46281], [48.22866, 153.46280]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([[-15.0000, -1.98757], [17.0000, 17.0000]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, tolerance); +} diff --git a/crates/burn-backend-tests/tests/autodiff/gradients.rs b/crates/burn-backend-tests/tests/autodiff/gradients.rs new file mode 100644 index 0000000..f51400e --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/gradients.rs @@ -0,0 +1,23 @@ +use super::*; +use burn_tensor::{Distribution, activation}; + +#[test] +fn should_update_tensor_when_grad_replace() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::random([32, 32], Distribution::Default, &device).require_grad(); + let tensor_2 = TestTensor::random([32, 32], Distribution::Default, &device); + + let x = tensor_1.clone().matmul(activation::gelu(tensor_2)); + let mut grads = x.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + + let grad_1_updated = + TestTensor::random([32, 32], Distribution::Default, &device).require_grad(); + tensor_1.grad_replace(&mut grads, grad_1_updated.clone().inner()); + + let grad_1_new = tensor_1.grad(&grads).unwrap(); + + assert_ne!(grad_1_new.to_data(), grad_1.into_data()); + assert_eq!(grad_1_new.into_data(), grad_1_updated.into_data()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/hypot.rs b/crates/burn-backend-tests/tests/autodiff/hypot.rs new file mode 100644 index 0000000..e824335 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/hypot.rs @@ -0,0 +1,27 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn should_diff_hypot() { + let data_a = TensorData::from([[3.0, 4.0], [5.0, 12.0]]); + let data_b = TensorData::from([[4.0, 3.0], [12.0, 5.0]]); + + let device = AutodiffDevice::new(); + let tensor_a = TestTensor::<2>::from_data(data_a, &device).require_grad(); + let tensor_b = TestTensor::<2>::from_data(data_b, &device).require_grad(); + + let tensor_c = tensor_a.clone().hypot(tensor_b.clone()); + let grads = tensor_c.backward(); + + let grad_a = tensor_a.grad(&grads).unwrap(); + let grad_b = tensor_b.grad(&grads).unwrap(); + + grad_a.to_data().assert_approx_eq::( + &TensorData::from([[0.6, 0.8], [0.3846, 0.9231]]), + Tolerance::default(), + ); + grad_b.to_data().assert_approx_eq::( + &TensorData::from([[0.8, 0.6], [0.9231, 0.3846]]), + Tolerance::default(), + ); +} diff --git a/crates/burn-backend-tests/tests/autodiff/linear.rs b/crates/burn-backend-tests/tests/autodiff/linear.rs new file mode 100644 index 0000000..9949949 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/linear.rs @@ -0,0 +1,70 @@ +use super::*; +use burn_tensor::module::linear; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn should_diff_linear_with_bias() { + let device = AutodiffDevice::new(); + let x = TestTensor::<3>::from_data( + [[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]], + &device, + ) + .require_grad(); + let weight = + TestTensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device).require_grad(); + let bias = TestTensor::<1>::from_data([0.1, 0.2, 0.3], &device).require_grad(); + + let output = linear(x.clone(), weight.clone(), Some(bias.clone())); + let grads = output.backward(); + + let tolerance = Tolerance::default(); + x.grad(&grads) + .unwrap() + .to_data() + .assert_approx_eq::( + &TensorData::from([[[6.0, 15.0], [6.0, 15.0]], [[6.0, 15.0], [6.0, 15.0]]]), + tolerance, + ); + weight + .grad(&grads) + .unwrap() + .to_data() + .assert_approx_eq::( + &TensorData::from([[16.0, 16.0, 16.0], [20.0, 20.0, 20.0]]), + tolerance, + ); + bias.grad(&grads) + .unwrap() + .to_data() + .assert_approx_eq::(&TensorData::from([4.0, 4.0, 4.0]), tolerance); +} + +#[test] +fn should_diff_linear_matches_matmul_decomposition() { + let device = AutodiffDevice::new(); + let x_data = TensorData::from([[[1.0, -2.0], [3.0, 4.0]], [[-5.0, 6.0], [7.0, -8.0]]]); + let w_data = TensorData::from([[0.5, -1.0, 2.0], [1.5, 2.5, -0.5]]); + + // Gradients through the linear module op. + let x = TestTensor::<3>::from_data(x_data.clone(), &device).require_grad(); + let weight = TestTensor::<2>::from_data(w_data.clone(), &device).require_grad(); + let grads = linear(x.clone(), weight.clone(), None).backward(); + let x_grad = x.grad(&grads).unwrap(); + let weight_grad = weight.grad(&grads).unwrap(); + + // Gradients through the explicit broadcast matmul decomposition. + let x_ref = TestTensor::<3>::from_data(x_data, &device).require_grad(); + let weight_ref = TestTensor::<2>::from_data(w_data, &device).require_grad(); + let grads_ref = x_ref + .clone() + .matmul(weight_ref.clone().unsqueeze::<3>()) + .backward(); + + let tolerance = Tolerance::default(); + x_grad + .to_data() + .assert_approx_eq::(&x_ref.grad(&grads_ref).unwrap().to_data(), tolerance); + weight_grad + .to_data() + .assert_approx_eq::(&weight_ref.grad(&grads_ref).unwrap().to_data(), tolerance); +} diff --git a/crates/burn-backend-tests/tests/autodiff/log.rs b/crates/burn-backend-tests/tests/autodiff/log.rs new file mode 100644 index 0000000..4b4da5f --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/log.rs @@ -0,0 +1,30 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn should_diff_log() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[6.0, 7.0], [9.0, 10.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().log()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let tolerance = Tolerance::default().set_half_precision_relative(1e-3); + let expected = TensorData::from([[60.2652, 72.3130], [60.2652, 72.3130]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([[22.8614, 24.5043], [24.5729, 26.8507]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, tolerance); +} diff --git a/crates/burn-backend-tests/tests/autodiff/log1p.rs b/crates/burn-backend-tests/tests/autodiff/log1p.rs new file mode 100644 index 0000000..70fbf00 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/log1p.rs @@ -0,0 +1,29 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_diff_log1p() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data([[0.0, 1.0], [3.0, 4.0]], &device).require_grad(); + let tensor_2 = TestTensor::from_data([[6.0, 7.0], [9.0, 10.0]], &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().log1p()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let tolerance = Tolerance::default().set_half_precision_relative(1e-3); + let expected = TensorData::from([[64.80622101, 75.49362183], [64.80622101, 75.49362183]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([[22.92208481, 24.47565651], [24.72780228, 26.86416626]]); + + grad_2 + .to_data() + .assert_approx_eq::(&expected, tolerance); +} diff --git a/crates/burn-backend-tests/tests/autodiff/log_sigmoid.rs b/crates/burn-backend-tests/tests/autodiff/log_sigmoid.rs new file mode 100644 index 0000000..746a04b --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/log_sigmoid.rs @@ -0,0 +1,19 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn should_diff_log_sigmoid() { + let data = TensorData::from([[0.8762, -0.1423], [-300., 200.]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data, &device).require_grad(); + let tensor_2 = activation::log_sigmoid(tensor_1.clone()); + let grads = tensor_2.backward(); + + let grad = tensor_1.grad(&grads).unwrap(); + + let expected = TensorData::from([[0.293966, 0.535515], [1.000000, 0.000000]]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/mask.rs b/crates/burn-backend-tests/tests/autodiff/mask.rs new file mode 100644 index 0000000..c61758c --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/mask.rs @@ -0,0 +1,63 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_diff_mask_fill() { + let data_1 = TensorData::from([[1.0, 7.0], [2.0, 3.0]]); + let data_2 = TensorData::from([[4.0, 7.0], [2.0, 3.0]]); + let mask = TensorData::from([[true, false], [false, true]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + let mask = TestTensorBool::<2>::from_bool(mask, &device); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_3.mask_fill(mask, 2.0); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[7.0, 3.0], [4.0, 2.0]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([[2.0, 1.0], [3.0, 7.0]]), false); +} + +#[test] +fn should_diff_mask_where() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::from_data([[1.0, 7.0], [2.0, 3.0]], &device).require_grad(); + let tensor_2 = TestTensor::from_data([[4.0, 7.0], [2.0, 3.0]], &device).require_grad(); + let tensor_3 = TestTensor::from_data([[8.8, 9.8], [10.8, 11.8]], &device).require_grad(); + let mask = TestTensorBool::<2>::from_data([[true, false], [false, true]], &device); + + let tensor_4 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_5 = tensor_4.clone().matmul(tensor_3.clone()); + let tensor_6 = tensor_5.mask_where(mask, tensor_3.clone()); + let grads = tensor_6.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + let grad_3 = tensor_3.grad(&grads).unwrap(); + + let tolerance = Tolerance::default().set_half_precision_relative(1e-3); + let expected = TensorData::from([[121.8, 55.0], [110.8, 50.0]]); + grad_1 + .into_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([[27.4, 33.4], [95.0, 115.0]]); + grad_2 + .into_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([[15., 18.], [23., 29.]]); + grad_3 + .into_data() + .assert_approx_eq::(&expected, tolerance); +} diff --git a/crates/burn-backend-tests/tests/autodiff/matmul.rs b/crates/burn-backend-tests/tests/autodiff/matmul.rs new file mode 100644 index 0000000..39cb6db --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/matmul.rs @@ -0,0 +1,83 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_diff_matmul() { + let data_1 = TensorData::from([[1.0, 7.0], [2.0, 3.0]]); + let data_2 = TensorData::from([[4.0, 7.0], [2.0, 3.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let grads = tensor_3.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[11.0, 5.0], [11.0, 5.0]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([[3.0, 3.0], [10.0, 10.0]]), false); + tensor_3 + .to_data() + .assert_eq(&TensorData::from([[18.0, 28.0], [14.0, 23.0]]), false); +} + +#[test] +fn test_matmul_complex_1() { + let data_1 = TensorData::from([[1.0, 7.0], [13.0, -3.0]]); + let data_2 = TensorData::from([[4.0, 7.0], [2.0, 3.0]]); + let data_3 = TensorData::from([[2.0, 2.0], [2.0, 2.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + let tensor_3 = TestTensor::from_data(data_3, &device).require_grad(); + + let tensor_4 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_5 = tensor_4.matmul(tensor_3); + + let grads = tensor_5.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[44.0, 20.0], [44.0, 20.0]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([[56.0, 56.0], [16.0, 16.0]]), false); +} + +#[test] +fn test_matmul_complex_2() { + let data_1 = TensorData::from([[1.0, 7.0], [13.0, -3.0]]); + let data_2 = TensorData::from([[4.0, 7.0], [2.0, 3.0]]); + let data_3 = TensorData::from([[2.0, 2.0], [2.0, 2.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + let tensor_3 = TestTensor::from_data(data_3, &device).require_grad(); + + let tensor_4 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_5 = tensor_4.matmul(tensor_3.clone()); + let tensor_6 = tensor_1.clone().matmul(tensor_5); + + let grads = tensor_6.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[800.0, 792.0], [360.0, 592.0]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([[264., 264.0], [344.0, 344.0]]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/maxmin.rs b/crates/burn-backend-tests/tests/autodiff/maxmin.rs new file mode 100644 index 0000000..d808f83 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/maxmin.rs @@ -0,0 +1,76 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_diff_max_dim() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data([[1.0, 7.0], [-2.0, -3.0]], &device).require_grad(); + let tensor_2 = TestTensor::from_data([[4.0, -7.0], [2.0, 3.0]], &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_1.clone().mul(tensor_3.max_dim(1).unsqueeze()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[50.0, 34.0], [40.0, -10.0]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[8.0, 10.0], [56.0, 15.0]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_min_dim() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data([[1.0, 7.0], [-2.0, -3.0]], &device).require_grad(); + let tensor_2 = TestTensor::from_data([[4.0, -7.0], [2.0, 3.0]], &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_1.clone().mul(tensor_3.min_dim(1).unsqueeze()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[-42.0, 38.0], [-34.0, -24.0]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[10.0, 8.0], [15.0, 56.0]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_min_dim_3d_dim1() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<3>::from_data([[[1.0, 7.0], [-2.0, -3.0]]], &device).require_grad(); + let tensor_2 = TestTensor::<3>::from_data([[[4., -7.], [2., 3.]]], &device).require_grad(); + + let tensor_3 = tensor_1.clone().mul(tensor_2.clone()); + let tensor_4 = tensor_3.min_dim(1); + + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[[0., -7.], [2., 0.]]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[[0., 7.], [-2., -0.]]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/maxpool1d.rs b/crates/burn-backend-tests/tests/autodiff/maxpool1d.rs new file mode 100644 index 0000000..046ee15 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/maxpool1d.rs @@ -0,0 +1,133 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::module::max_pool1d; + +#[test] +fn test_max_pool1d_simple() { + let kernel_size = 4; + let padding = 0; + let stride = 1; + let dilation = 1; + + let device = AutodiffDevice::new(); + let x = TestTensor::from_data( + [[[0.9861, 0.5474, 0.4477, 0.0732, 0.3548, 0.8221]]], + &device, + ) + .require_grad(); + let x_grad_expected = TestTensor::<3>::from_data([[[1., 1., 0., 0., 0., 1.]]], &device); + + let output = max_pool1d(x.clone(), kernel_size, stride, padding, dilation, false); + let grads = output.backward(); + + // Asserts + let x_grad_actual = x.grad(&grads).unwrap(); + x_grad_expected + .to_data() + .assert_approx_eq::(&x_grad_actual.to_data(), Tolerance::default()); +} + +#[test] +fn test_max_pool1d_with_dilation() { + let kernel_size = 4; + let padding = 0; + let stride = 1; + let dilation = 2; + + let device = AutodiffDevice::new(); + let x = TestTensor::from_data( + [[[ + 0.5388, 0.0676, 0.7122, 0.8316, 0.0653, 0.9154, 0.1536, 0.9089, 0.8016, 0.7518, 0.2073, + 0.0501, 0.8811, 0.5604, 0.5075, 0.4384, 0.9963, 0.9698, 0.4988, 0.2609, 0.3391, 0.2230, + 0.4610, 0.5365, 0.6880, + ]]], + &device, + ) + .require_grad(); + let x_grad_expected = TestTensor::<3>::from_data( + [[[ + 0., 0., 1., 0., 0., 3., 0., 1., 2., 1., 0., 0., 2., 0., 0., 0., 4., 4., 0., 0., 0., 0., + 0., 0., 1., + ]]], + &device, + ); + + let output = max_pool1d(x.clone(), kernel_size, stride, padding, dilation, false); + let grads = output.backward(); + + // Asserts + let x_grad_actual = x.grad(&grads).unwrap(); + x_grad_expected + .to_data() + .assert_approx_eq::(&x_grad_actual.to_data(), Tolerance::default()); +} + +#[test] +fn test_max_pool1d_complex() { + let kernel_size = 4; + let padding = 0; + let stride = 1; + let dilation = 1; + + let device = AutodiffDevice::new(); + let x = TestTensor::from_data( + [[[ + 0.5388, 0.0676, 0.7122, 0.8316, 0.0653, 0.9154, 0.1536, 0.9089, 0.8016, 0.7518, 0.2073, + 0.0501, 0.8811, 0.5604, 0.5075, 0.4384, 0.9963, 0.9698, 0.4988, 0.2609, 0.3391, 0.2230, + 0.4610, 0.5365, 0.6880, + ]]], + &device, + ) + .require_grad(); + let x_grad_expected = TestTensor::<3>::from_data( + [[[ + 0., 0., 0., 2., 0., 4., 0., 2., 1., 0., 0., 0., 4., 0., 0., 0., 4., 1., 1., 0., 0., 0., + 1., 1., 1., + ]]], + &device, + ); + + let output = max_pool1d(x.clone(), kernel_size, stride, padding, dilation, false); + let grads = output.backward(); + + // Asserts + let x_grad_actual = x.grad(&grads).unwrap(); + x_grad_expected + .to_data() + .assert_approx_eq::(&x_grad_actual.to_data(), Tolerance::default()); +} + +#[test] +fn test_max_pool1d_complex_with_padding() { + let kernel_size = 4; + let padding = 2; + let stride = 1; + let dilation = 1; + + let device = AutodiffDevice::new(); + let x = TestTensor::from_data( + [[[ + 0.5388, 0.0676, 0.7122, 0.8316, 0.0653, 0.9154, 0.1536, 0.9089, 0.8016, 0.7518, 0.2073, + 0.0501, 0.8811, 0.5604, 0.5075, 0.4384, 0.9963, 0.9698, 0.4988, 0.2609, 0.3391, 0.2230, + 0.4610, 0.5365, 0.6880, + ]]], + &device, + ) + .require_grad(); + let x_grad_expected = TestTensor::<3>::from_data( + [[[ + 1., 0., 1., 2., 0., 4., 0., 2., 1., 0., 0., 0., 4., 0., 0., 0., 4., 1., 1., 0., 0., 0., + 1., 1., 3., + ]]], + &device, + ); + + let output = max_pool1d(x.clone(), kernel_size, stride, padding, dilation, false); + let grads = output.backward(); + + // Asserts + let x_grad_actual = x.grad(&grads).unwrap(); + x_grad_expected + .to_data() + .assert_approx_eq::(&x_grad_actual.to_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/maxpool2d.rs b/crates/burn-backend-tests/tests/autodiff/maxpool2d.rs new file mode 100644 index 0000000..87328de --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/maxpool2d.rs @@ -0,0 +1,271 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::module::max_pool2d; + +#[test] +fn test_max_pool2d_simple_1() { + let kernel_size_1 = 3; + let kernel_size_2 = 3; + let padding_1 = 0; + let padding_2 = 0; + let stride_1 = 1; + let stride_2 = 1; + let dilation_1 = 1; + let dilation_2 = 1; + + let device = AutodiffDevice::new(); + let x = TestTensor::from_data( + [[[ + [0.2479, 0.6386, 0.3166, 0.5742], + [0.7065, 0.1940, 0.6305, 0.8959], + [0.5416, 0.8602, 0.8129, 0.1662], + [0.3358, 0.3059, 0.8293, 0.0990], + ]]], + &device, + ) + .require_grad(); + let x_grad_expected = TestTensor::<4>::from_data( + [[[ + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 2.0], + [0.0, 2.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + ]]], + &device, + ); + + let output = max_pool2d( + x.clone(), + [kernel_size_1, kernel_size_2], + [stride_1, stride_2], + [padding_1, padding_2], + [dilation_1, dilation_2], + false, + ); + let grads = output.backward(); + + // Asserts + let x_grad_actual = x.grad(&grads).unwrap(); + x_grad_expected + .to_data() + .assert_approx_eq::(&x_grad_actual.to_data(), Tolerance::default()); +} + +#[test] +fn test_max_pool2d_simple_2() { + let kernel_size_1 = 2; + let kernel_size_2 = 2; + let padding_1 = 1; + let padding_2 = 1; + let stride_1 = 1; + let stride_2 = 1; + let dilation_1 = 1; + let dilation_2 = 1; + + let device = AutodiffDevice::new(); + let x = TestTensor::from_data( + [[[ + [0.2479, 0.6386, 0.3166, 0.5742], + [0.7065, 0.1940, 0.6305, 0.8959], + [0.5416, 0.8602, 0.8129, 0.1662], + [0.3358, 0.3059, 0.8293, 0.0990], + ]]], + &device, + ) + .require_grad(); + let x_grad_expected = TestTensor::<4>::from_data( + [[[ + [1., 3., 0., 2.], + [3., 0., 0., 4.], + [1., 4., 0., 1.], + [2., 0., 3., 1.], + ]]], + &device, + ); + + let output = max_pool2d( + x.clone(), + [kernel_size_1, kernel_size_2], + [stride_1, stride_2], + [padding_1, padding_2], + [dilation_1, dilation_2], + false, + ); + let grads = output.backward(); + + // Asserts + let x_grad_actual = x.grad(&grads).unwrap(); + x_grad_expected + .to_data() + .assert_approx_eq::(&x_grad_actual.to_data(), Tolerance::default()); +} + +#[test] +fn test_max_pool2d_with_dilation() { + let kernel_size_1 = 2; + let kernel_size_2 = 2; + let padding_1 = 1; + let padding_2 = 1; + let stride_1 = 1; + let stride_2 = 1; + let dilation_1 = 2; + let dilation_2 = 2; + + let device = AutodiffDevice::new(); + let x = TestTensor::from_data( + [[[ + [0.2479, 0.6386, 0.3166, 0.5742], + [0.7065, 0.1940, 0.6305, 0.8959], + [0.5416, 0.8602, 0.8129, 0.1662], + [0.3358, 0.3059, 0.8293, 0.0990], + ]]], + &device, + ) + .require_grad(); + let x_grad_expected = TestTensor::<4>::from_data( + [[[ + [0., 0., 0., 0.], + [1., 1., 1., 2.], + [0., 4., 4., 0.], + [0., 1., 2., 0.], + ]]], + &device, + ); + + let output = max_pool2d( + x.clone(), + [kernel_size_1, kernel_size_2], + [stride_1, stride_2], + [padding_1, padding_2], + [dilation_1, dilation_2], + false, + ); + let grads = output.backward(); + + // Asserts + let x_grad_actual = x.grad(&grads).unwrap(); + x_grad_expected + .to_data() + .assert_approx_eq::(&x_grad_actual.to_data(), Tolerance::default()); +} + +#[test] +fn test_max_pool2d_complex() { + let kernel_size_1 = 4; + let kernel_size_2 = 2; + let padding_1 = 2; + let padding_2 = 1; + let stride_1 = 1; + let stride_2 = 2; + let dilation_1 = 1; + let dilation_2 = 1; + + let device = AutodiffDevice::new(); + let x = TestTensor::from_data( + [[[ + [0.5388, 0.0676, 0.7122, 0.8316, 0.0653], + [0.9154, 0.1536, 0.9089, 0.8016, 0.7518], + [0.2073, 0.0501, 0.8811, 0.5604, 0.5075], + [0.4384, 0.9963, 0.9698, 0.4988, 0.2609], + [0.3391, 0.2230, 0.4610, 0.5365, 0.6880], + ]]], + &device, + ) + .require_grad(); + let x_grad_expected = TestTensor::<4>::from_data( + [[[ + [0., 0., 0., 3., 0.], + [4., 0., 2., 1., 0.], + [0., 0., 0., 0., 0.], + [2., 4., 0., 0., 0.], + [0., 0., 0., 0., 2.], + ]]], + &device, + ); + + let output = max_pool2d( + x.clone(), + [kernel_size_1, kernel_size_2], + [stride_1, stride_2], + [padding_1, padding_2], + [dilation_1, dilation_2], + false, + ); + let grads = output.backward(); + + // Asserts + let x_grad_actual = x.grad(&grads).unwrap(); + x_grad_expected + .to_data() + .assert_approx_eq::(&x_grad_actual.to_data(), Tolerance::default()); +} + +#[test] +fn test_max_pool2d_ceil_mode() { + // Test ceil_mode=true with gradient computation + // Using 1x1x6x6 input with kernel 3x3, stride 2x2, padding 0 + // Floor mode: output 2x2 + // Ceil mode: output 3x3 + let kernel_size_1 = 3; + let kernel_size_2 = 3; + let padding_1 = 0; + let padding_2 = 0; + let stride_1 = 2; + let stride_2 = 2; + let dilation_1 = 1; + let dilation_2 = 1; + + let device = AutodiffDevice::new(); + // Input (values 1-36): + let x = TestTensor::from_data( + [[[ + [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], + [7.0, 8.0, 9.0, 10.0, 11.0, 12.0], + [13.0, 14.0, 15.0, 16.0, 17.0, 18.0], + [19.0, 20.0, 21.0, 22.0, 23.0, 24.0], + [25.0, 26.0, 27.0, 28.0, 29.0, 30.0], + [31.0, 32.0, 33.0, 34.0, 35.0, 36.0], + ]]], + &device, + ) + .require_grad(); + + // Expected gradients for ceil_mode output 3x3: + // Output positions and their max value positions: + // (0,0): max at (2,2)=15 -> grad[2,2] += 1 + // (0,1): max at (2,4)=17 -> grad[2,4] += 1 + // (0,2): max at (2,5)=18 -> grad[2,5] += 1 + // (1,0): max at (4,2)=27 -> grad[4,2] += 1 + // (1,1): max at (4,4)=29 -> grad[4,4] += 1 + // (1,2): max at (4,5)=30 -> grad[4,5] += 1 + // (2,0): max at (5,2)=33 -> grad[5,2] += 1 + // (2,1): max at (5,4)=35 -> grad[5,4] += 1 + // (2,2): max at (5,5)=36 -> grad[5,5] += 1 + let x_grad_expected = TestTensor::<4>::from_data( + [[[ + [0., 0., 0., 0., 0., 0.], + [0., 0., 0., 0., 0., 0.], + [0., 0., 1., 0., 1., 1.], + [0., 0., 0., 0., 0., 0.], + [0., 0., 1., 0., 1., 1.], + [0., 0., 1., 0., 1., 1.], + ]]], + &device, + ); + + let output = max_pool2d( + x.clone(), + [kernel_size_1, kernel_size_2], + [stride_1, stride_2], + [padding_1, padding_2], + [dilation_1, dilation_2], + true, + ); + let grads = output.backward(); + + // Asserts + let x_grad_actual = x.grad(&grads).unwrap(); + x_grad_expected + .to_data() + .assert_approx_eq::(&x_grad_actual.to_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/memory_management.rs b/crates/burn-backend-tests/tests/autodiff/memory_management.rs new file mode 100644 index 0000000..9fc43fa --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/memory_management.rs @@ -0,0 +1,280 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_mm_independent_trees() { + let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + let device = AutodiffDevice::new(); + + // First tree + let tensor_0 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_1 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_2 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_3 = TestTensor::from_data(data.clone(), &device).require_grad(); + + let tensor_4 = tensor_0 * tensor_1; + let tensor_5 = tensor_2 * tensor_3; + let tensor_6 = tensor_4 * tensor_5; + + // Second tree + let tensor_7 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_8 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_9 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_10 = TestTensor::from_data(data.clone(), &device).require_grad(); + + let tensor_11 = tensor_7.clone() * tensor_8.clone(); + let tensor_12 = tensor_9.clone() * tensor_10.clone(); + let tensor_13 = tensor_11 * tensor_12; + + let _grads = tensor_6.backward(); + let grads = tensor_13.backward(); + + assert!(tensor_7.grad(&grads).is_some()); + assert!(tensor_8.grad(&grads).is_some()); + assert!(tensor_9.grad(&grads).is_some()); + assert!(tensor_10.grad(&grads).is_some()); +} + +#[test] +#[should_panic] +fn test_mm_crossover_trees_root_unavailable() { + let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + let device = AutodiffDevice::new(); + + // First tree + let tensor_0 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_1 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_2 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_3 = TestTensor::from_data(data.clone(), &device).require_grad(); + + let tensor_4 = tensor_0 * tensor_1; + let tensor_5 = tensor_2 * tensor_3; + let tensor_6 = tensor_4.clone() * tensor_5; + + // Second tree + let tensor_7 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_8 = TestTensor::from_data(data.clone(), &device).require_grad(); + + let tensor_9 = tensor_7.clone() * tensor_8.clone(); + let tensor_10 = tensor_4 * tensor_9; + + let _grads = tensor_6.backward(); + let _grads = tensor_10.backward(); +} + +#[test] +fn test_mm_crossover_trees_with_referred_subtree() { + let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + let device = AutodiffDevice::new(); + + // First tree + let tensor_0 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_1 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_2 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_3 = TestTensor::from_data(data.clone(), &device).require_grad(); + + let tensor_4 = tensor_0 * tensor_1; + let tensor_5 = tensor_2 * tensor_3; + let tensor_6 = tensor_4.clone() * tensor_5; + + // Second tree + let tensor_7 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_8 = TestTensor::from_data(data.clone(), &device).require_grad(); + + let tensor_9 = tensor_7.clone() * tensor_8.clone(); + let _tensor_10 = tensor_4 * tensor_9.clone(); + + let _grads = tensor_6.backward(); + let _grads = tensor_9.backward(); +} + +#[test] +fn test_mm_three_crossover_trees_last_still_usable() { + let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + let device = AutodiffDevice::new(); + + // First tree + let tensor_0 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_1 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_2 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_3 = TestTensor::from_data(data.clone(), &device).require_grad(); + + let tensor_4 = tensor_0 * tensor_1; + let tensor_5 = tensor_2 * tensor_3; + let tensor_6 = tensor_4 * tensor_5.clone(); + + // Third tree + let tensor_7 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_8 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_9 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_10 = TestTensor::from_data(data.clone(), &device).require_grad(); + + let tensor_11 = tensor_7 * tensor_8; + let tensor_12 = tensor_9 * tensor_10; + let tensor_13 = tensor_11 * tensor_12.clone(); + + // Second tree (in between) + let _tensor_14 = tensor_5 * tensor_12; + + let _grads = tensor_6.backward(); + let _grads = tensor_13.backward(); +} + +#[test] +#[should_panic] +fn test_mm_three_crossover_trees_middle_one_unavailable() { + let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + let device = AutodiffDevice::new(); + + // First tree + let tensor_0 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_1 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_2 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_3 = TestTensor::from_data(data.clone(), &device).require_grad(); + + let tensor_4 = tensor_0 * tensor_1; + let tensor_5 = tensor_2 * tensor_3; + let tensor_6 = tensor_4 * tensor_5.clone(); + + // Third tree + let tensor_7 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_8 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_9 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_10 = TestTensor::from_data(data.clone(), &device).require_grad(); + + let tensor_11 = tensor_7 * tensor_8; + let tensor_12 = tensor_9 * tensor_10; + let _tensor_13 = tensor_11 * tensor_12.clone(); + + // Second tree (in between) + let tensor_14 = tensor_5 * tensor_12; + + let _grads = tensor_6.backward(); + let _grads = tensor_14.backward(); +} + +#[test] +fn test_mm_self_referencing_tree() { + let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + let device = AutodiffDevice::new(); + + // First tree + let tensor_0 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_1 = TestTensor::from_data(data.clone(), &device).require_grad(); + let tensor_2 = TestTensor::from_data(data.clone(), &device).require_grad(); + + let tensor_3 = tensor_0 * tensor_1; + let tensor_5 = tensor_2 * tensor_3.clone(); + let tensor_6 = tensor_3 * tensor_5; + + let _grads = tensor_6.backward(); +} + +#[test] +fn test_mm_with_non_impacting_detach() { + let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_2 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_3 = TestTensor::<2>::from_data(data, &device).require_grad(); + + let tensor_4 = tensor_1.clone() * tensor_2.clone(); + let tensor_5 = tensor_4.detach() * tensor_3.clone(); + + let grads = tensor_5.backward(); + assert!(tensor_3.grad(&grads).is_some()); +} + +#[test] +fn test_mm_with_missing_require_grad_after_cleanup() { + let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + let device = AutodiffDevice::new(); + + let tensor_1 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_2 = TestTensor::<2>::from_data(data.clone(), &device); + let tensor_3 = TestTensor::<2>::from_data(data.clone(), &device); + + let tensor_4 = tensor_1.clone() * tensor_2.clone(); + let tensor_5 = tensor_4 * tensor_3.clone(); + + // Trivial backward, just to trigger cleanup + TestTensor::<2>::from_data(data, &device) + .require_grad() + .backward(); + + let grads = tensor_5.backward(); + assert!(tensor_1.grad(&grads).is_some()); + assert!(tensor_2.grad(&grads).is_none()); + assert!(tensor_3.grad(&grads).is_none()); +} + +#[test] +fn test_mm_with_detach_after_cleanup() { + let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + let device = AutodiffDevice::new(); + + let tensor_1 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_2 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_3 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + + let tensor_4 = tensor_1.clone() * tensor_2.clone(); + let tensor_5 = tensor_4 * tensor_3.clone().detach(); + + // Trivial backward, just to trigger cleanup + TestTensor::<2>::from_data(data, &device) + .require_grad() + .backward(); + + let grads = tensor_5.backward(); + assert!(tensor_1.grad(&grads).is_some()); + assert!(tensor_2.grad(&grads).is_some()); + assert!(tensor_3.grad(&grads).is_none()); +} + +#[test] +#[should_panic] +fn test_mm_deletables_propagate_well() { + let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + let device = AutodiffDevice::new(); + + let tensor_0 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_1 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + + let tensor_2 = tensor_0 * tensor_1; + let tensor_3 = tensor_2.clone().exp(); + let _tensor_4 = tensor_3.clone().log(); + + let _grads = tensor_2.backward(); + + // We are testing that after backward on tensor_2, not only the leaf tensor_4 is deleted, but + // the intermediate tensor_3 as well + let _grads = tensor_3.backward(); +} + +#[test] +fn test_mm_node_explored_once_can_still_be_tagged_as_useful_when_found_again_deeper() { + let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + let device = AutodiffDevice::new(); + + // The test has 50% chance of starting with leaf tensor_8 instead of tensor_4, which is not informative + // By repeating it many times it becomes almost impossible that it passes if it shouldn't + for _ in 0..12 { + let tensor_0 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_1 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + + let tensor_2 = tensor_1.clone().exp(); + let tensor_3 = tensor_0.exp(); + let _tensor_4 = tensor_3.clone() * tensor_2.clone(); + let tensor_5 = tensor_2.exp(); + let tensor_6 = tensor_5.exp(); + let tensor_7 = tensor_6.exp(); + let tensor_8 = tensor_7.exp(); + + // tensor_2 should be tagged unknown through the leaf tensor_4, then useful through the leaf tensor_8 + // which should happen after because tensor_2 is deeper from tensor_8 point of view and we're in breadth first search + tensor_3.backward(); + let grads = tensor_8.backward(); + + assert!(tensor_1.grad(&grads).is_some()); + } +} diff --git a/crates/burn-backend-tests/tests/autodiff/mod.rs b/crates/burn-backend-tests/tests/autodiff/mod.rs new file mode 100644 index 0000000..df025ee --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/mod.rs @@ -0,0 +1,81 @@ +#[allow(unused_imports)] // required for re-included modules +pub use super::*; + +mod abs; +mod adaptive_avgpool1d; +mod adaptive_avgpool2d; +mod add; +mod aggregation; +#[cfg(feature = "distributed")] +mod all_reduce; +mod avgpool1d; +mod avgpool2d; +mod backward; +mod bridge; +mod broadcast; +mod cast; +mod cat; +mod ceil; +mod checkpoint; +mod complex; +mod conv1d; +mod conv2d; +mod conv3d; +mod conv_transpose1d; +mod conv_transpose2d; +mod conv_transpose3d; +mod cross; +mod cross_entropy; +mod ctc; +mod cummax; +mod cummin; +mod cumprod; +mod cumsum; +mod deform_conv2d; +mod div; +mod erf; +mod exp; +mod expand; +mod flip; +mod floor; +mod gather_scatter; +mod gather_scatter_nd; +mod gelu; +mod gradients; +mod hypot; +mod linear; +mod log; +mod log1p; +mod log_sigmoid; +mod mask; +mod matmul; +mod maxmin; +mod maxpool1d; +mod maxpool2d; +mod memory_management; +mod mul; +mod multithread; +mod nearest_interpolate; +mod neg; +mod nonzero; +mod permute; +mod pow; +mod recip; +mod relu; +mod remainder; +mod repeat_dim; +mod reshape; +mod rfft; +mod round; +mod select; +mod sigmoid; +mod sign; +mod slice; +mod slice_assign; +mod softmax; +mod sort; +mod sqrt; +mod sub; +mod transpose; +mod trig; +mod unfold; diff --git a/crates/burn-backend-tests/tests/autodiff/mul.rs b/crates/burn-backend-tests/tests/autodiff/mul.rs new file mode 100644 index 0000000..7c43532 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/mul.rs @@ -0,0 +1,68 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_diff_mul() { + let data_1 = TensorData::from([1.0, 7.0]); + let data_2 = TensorData::from([4.0, 7.0]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<1>::from_data(data_1.clone(), &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2.clone(), &device).require_grad(); + + let tensor_3 = tensor_1.clone().mul(tensor_2.clone()); + let grads = tensor_3.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let _grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1.to_data().assert_eq(&data_2, false); + tensor_3 + .into_data() + .assert_eq(&TensorData::from([4.0, 49.0]), false); +} + +#[test] +fn should_diff_mul_scalar() { + let data = TensorData::from([2.0, 5.0]); + + let tensor = TestTensor::<1>::from_data(data, &AutodiffDevice::new()).require_grad(); + let tensor_out = tensor.clone().mul_scalar(4.0); + + let grads = tensor_out.backward(); + let grad = tensor.grad(&grads).unwrap(); + + tensor_out + .into_data() + .assert_eq(&TensorData::from([8.0, 20.0]), false); + grad.to_data() + .assert_eq(&TensorData::from([4.0, 4.0]), false); +} + +#[test] +fn test_mul_complex_1() { + let data_1 = TensorData::from([[1.0, 7.0], [13.0, -3.0]]); + let data_2 = TensorData::from([[4.0, 7.0], [2.0, 3.0]]); + let data_3 = TensorData::from([[2.0, 2.0], [2.0, 2.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + let tensor_3 = TestTensor::from_data(data_3, &device).require_grad(); + + let tensor_4 = tensor_1.clone().mul(tensor_2.clone()); + let tensor_5 = tensor_4.mul(tensor_3); + let tensor_6 = tensor_1.clone().mul(tensor_5); + + let grads = tensor_6.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[16.0, 196.0], [104.0, -36.0]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([[2.0, 98.0], [338.0, 18.0]]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/multithread.rs b/crates/burn-backend-tests/tests/autodiff/multithread.rs new file mode 100644 index 0000000..19f36a4 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/multithread.rs @@ -0,0 +1,90 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +// TODO: FIXME https://github.com/tracel-ai/burn/issues/4739 +#[ignore = "Flaky test"] +fn should_behave_the_same_with_multithread() { + let data_1 = TensorData::from([[1.0, 7.0], [13.0, -3.0]]); + let data_2 = TensorData::from([[4.0, 7.0], [2.0, 3.0]]); + + let with_move = || { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1.clone(), &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2.clone(), &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_3.clone().matmul(tensor_2.clone()); + let tensor_5 = tensor_4.matmul(tensor_3); + + // Task 1 + let tensor_1_cloned = tensor_1.clone(); + let tensor_2_cloned = tensor_2.clone(); + let tensor_5_cloned = tensor_5.clone(); + + let first_call = move || { + let tensor_6_1 = tensor_5_cloned.matmul(tensor_2_cloned); + tensor_6_1.matmul(tensor_1_cloned) + }; + + // Task 2 + let tensor_1_cloned = tensor_1.clone(); + let tensor_2_cloned = tensor_2.clone(); + let tensor_5_cloned = tensor_5; + + let second_call = move || { + let tensor_6_2 = tensor_5_cloned.matmul(tensor_1_cloned); + tensor_6_2.matmul(tensor_2_cloned) + }; + + let tensor_7_1_handle = std::thread::spawn(first_call); + let tensor_7_2_handle = std::thread::spawn(second_call); + + let tensor_7_1 = tensor_7_1_handle.join().unwrap(); + let tensor_7_2 = tensor_7_2_handle.join().unwrap(); + let tensor_8 = tensor_7_1.matmul(tensor_7_2); + + let grads = tensor_8.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + (grad_1, grad_2) + }; + let without_move = || { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1.clone(), &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2.clone(), &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_3.clone().matmul(tensor_2.clone()); + let tensor_5 = tensor_4.matmul(tensor_3); + + // Task 1 + let tensor_6_1 = tensor_5.clone().matmul(tensor_2.clone()); + let tensor_7_1 = tensor_6_1.matmul(tensor_1.clone()); + + // Task 2 + let tensor_6_2 = tensor_5.matmul(tensor_1.clone()); + let tensor_7_2 = tensor_6_2.matmul(tensor_2.clone()); + + let tensor_8 = tensor_7_1.matmul(tensor_7_2); + + let grads = tensor_8.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + (grad_1, grad_2) + }; + + let (grad_1, grad_2) = without_move(); + let (grad_1_moved, grad_2_moved) = with_move(); + + grad_1 + .into_data() + .assert_approx_eq::(&grad_1_moved.into_data(), Tolerance::default()); + grad_2 + .into_data() + .assert_approx_eq::(&grad_2_moved.into_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/nearest_interpolate.rs b/crates/burn-backend-tests/tests/autodiff/nearest_interpolate.rs new file mode 100644 index 0000000..11ce084 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/nearest_interpolate.rs @@ -0,0 +1,97 @@ +use super::*; +use burn_tensor::Shape; +use burn_tensor::Tolerance; +use burn_tensor::module::interpolate; +use burn_tensor::ops::{InterpolateMode, InterpolateOptions}; + +#[test] +fn test_upsample_interpolation() { + let test = InterpolateTestCase { + batch_size: 2, + channels: 1, + height: 7, + width: 5, + height_out: 8, + width_out: 7, + }; + + test.assert_output(TestTensor::from([ + [[ + [4., 2., 4., 2., 2.], + [2., 1., 2., 1., 1.], + [2., 1., 2., 1., 1.], + [2., 1., 2., 1., 1.], + [2., 1., 2., 1., 1.], + [2., 1., 2., 1., 1.], + [2., 1., 2., 1., 1.], + ]], + [[ + [4., 2., 4., 2., 2.], + [2., 1., 2., 1., 1.], + [2., 1., 2., 1., 1.], + [2., 1., 2., 1., 1.], + [2., 1., 2., 1., 1.], + [2., 1., 2., 1., 1.], + [2., 1., 2., 1., 1.], + ]], + ])); +} + +#[test] +fn test_downsample_interpolation() { + let test = InterpolateTestCase { + batch_size: 1, + channels: 1, + height: 8, + width: 8, + height_out: 4, + width_out: 6, + }; + + test.assert_output(TestTensor::from([[[ + [1., 1., 1., 0., 1., 1., 1., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.], + [1., 1., 1., 0., 1., 1., 1., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.], + [1., 1., 1., 0., 1., 1., 1., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.], + [1., 1., 1., 0., 1., 1., 1., 0.], + [0., 0., 0., 0., 0., 0., 0., 0.], + ]]])); +} + +struct InterpolateTestCase { + batch_size: usize, + channels: usize, + height: usize, + width: usize, + height_out: usize, + width_out: usize, +} + +impl InterpolateTestCase { + fn assert_output(self, x_grad: TestTensor<4>) { + let shape_x = Shape::new([self.batch_size, self.channels, self.height, self.width]); + let device = AutodiffDevice::new(); + let x = TestTensor::from_data( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &x_grad.device()) + .reshape::<4, _>(shape_x) + .into_data(), + &device, + ) + .require_grad(); + + let output = interpolate( + x.clone(), + [self.height_out, self.width_out], + InterpolateOptions::new(InterpolateMode::Nearest), + ); + + let grads = output.backward(); + let x_grad_actual = x.grad(&grads).unwrap(); + + x_grad + .to_data() + .assert_approx_eq::(&x_grad_actual.into_data(), Tolerance::permissive()); + } +} diff --git a/crates/burn-backend-tests/tests/autodiff/neg.rs b/crates/burn-backend-tests/tests/autodiff/neg.rs new file mode 100644 index 0000000..30cc317 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/neg.rs @@ -0,0 +1,26 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_diff_neg() { + let data_1 = TensorData::from([[1.0, 7.0], [2.0, 3.0]]); + let data_2 = TensorData::from([[4.0, 7.0], [2.0, 3.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().neg()); + let tensor_4 = tensor_3.neg(); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[11.0, 5.0], [11.0, 5.0]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([[3.0, 3.0], [10.0, 10.0]]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/nonzero.rs b/crates/burn-backend-tests/tests/autodiff/nonzero.rs new file mode 100644 index 0000000..5ac59c5 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/nonzero.rs @@ -0,0 +1,41 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_diff_nonzero() { + let data_1 = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + let data_2 = TensorData::from([-1.0, 1.0]); + let mask = TensorData::from([[false, true], [true, false]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::<1>::from_data(data_2, &device).require_grad(); + + // Multi-dimensional tensor indexing isn't really supported yet so the easiest way to do + // this is to flatten the mask and tensor to get proper indexing. Anyway the returned tensor would + // have dimensions different from the input, so this is somewhat equivalent. + let mask = TestTensorBool::<2>::from_bool(mask, &device).flatten::<1>(0, 1); + let indices = mask.nonzero(); + let tensor_3 = tensor_1 + .clone() + .flatten::<1>(0, 1) + .select(0, indices[0].clone()); + + // Vector dot product not supported (only 2D matmuls) so unsqueeze for test purposes + let tensor_4 = tensor_2 + .clone() + .unsqueeze_dim::<2>(0) + .matmul(tensor_3.unsqueeze_dim(1)); + + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[0.0, -1.0], [1.0, 0.0]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([2.0, 3.0]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/permute.rs b/crates/burn-backend-tests/tests/autodiff/permute.rs new file mode 100644 index 0000000..6a7510b --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/permute.rs @@ -0,0 +1,29 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_diff_permute() { + let data_1 = TensorData::from([[[1.0, 7.0], [2.0, 3.0]]]); // 1x2x2 + let data_2 = TensorData::from([[[1.0, 7.0], [3.2, 2.0], [3.0, 3.0]]]); // 1x3x2 + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_2.clone().permute([0, 2, 1]); + let tensor_4 = tensor_1.clone().matmul(tensor_3); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let tolerance = Tolerance::default().set_half_precision_relative(1e-3); + grad_1 + .into_data() + .assert_approx_eq::(&TensorData::from([[[7.2, 12.0], [7.2, 12.0]]]), tolerance); // 1x2x2 + grad_2.into_data().assert_approx_eq::( + &TensorData::from([[[3.0, 10.0], [3.0, 10.0], [3.0, 10.0]]]), + tolerance, + ); // 1x3x2 +} diff --git a/crates/burn-backend-tests/tests/autodiff/pow.rs b/crates/burn-backend-tests/tests/autodiff/pow.rs new file mode 100644 index 0000000..da55d04 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/pow.rs @@ -0,0 +1,93 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_diff_powf_scalar() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[6.0, 7.0], [9.0, 10.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().powf_scalar(0.4)); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let tolerance = Tolerance::default().set_half_precision_relative(2e-3); + let expected = TensorData::from([[68.0, 79.0328], [68.0, 79.0328]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([[23.5081, 25.2779], [26.0502, 28.6383]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn should_diff_powf() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<1>::from_data([2.0, 7.0], &device).require_grad(); + let tensor_2 = TestTensor::from_data([4.0, 2.0], &device).require_grad(); + + let tensor_3 = tensor_1.clone().powf(tensor_2.clone()); + let grads = tensor_3.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([32.0, 14.0]); + grad_1 + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([11.09035, 95.34960]); + grad_2 + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([16.0, 49.0]); + tensor_3 + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_powf_with_untracked_lhs() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<1>::from_data([2.0, 7.0], &device); + let tensor_2 = TestTensor::from_data([4.0, 2.0], &device).require_grad(); + + let tensor_3 = tensor_1.clone().powf(tensor_2.clone()); + let grads = tensor_3.backward(); + + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([11.09035, 95.34960]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_powf_with_untracked_rhs() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<1>::from_data([2.0, 7.0], &device).require_grad(); + let tensor_2 = TestTensor::from_data([4.0, 2.0], &device); + + let tensor_3 = tensor_1.clone().powf(tensor_2.clone()); + let grads = tensor_3.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + + let expected = TensorData::from([32.0, 14.0]); + grad_1 + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/recip.rs b/crates/burn-backend-tests/tests/autodiff/recip.rs new file mode 100644 index 0000000..f6fb667 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/recip.rs @@ -0,0 +1,22 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_diff_recip() { + let data = TensorData::from([2.0, 5.0, 0.4]); + + let tensor = TestTensor::<1>::from_data(data, &AutodiffDevice::new()).require_grad(); + let tensor_out = tensor.clone().recip(); + + let grads = tensor_out.backward(); + let grad = tensor.grad(&grads).unwrap(); + + tensor_out + .into_data() + .assert_eq(&TensorData::from([0.5, 0.2, 2.5]), false); + grad.to_data().assert_approx_eq::( + &TensorData::from([-0.25, -0.04, -6.25]), + Tolerance::default(), + ); +} diff --git a/crates/burn-backend-tests/tests/autodiff/relu.rs b/crates/burn-backend-tests/tests/autodiff/relu.rs new file mode 100644 index 0000000..118c4e9 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/relu.rs @@ -0,0 +1,27 @@ +use super::*; +use burn_tensor::{TensorData, activation}; + +#[test] +fn should_diff_relu() { + let data_1 = TensorData::from([[1.0, 7.0], [-2.0, -3.0]]); + let data_2 = TensorData::from([[4.0, -7.0], [2.0, 3.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = activation::relu(tensor_3); + let tensor_5 = tensor_4.matmul(tensor_2.clone()); + let grads = tensor_5.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[-47.0, 9.0], [-35.0, 15.0]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([[15.0, 13.0], [-2.0, 39.0]]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/remainder.rs b/crates/burn-backend-tests/tests/autodiff/remainder.rs new file mode 100644 index 0000000..e33dc09 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/remainder.rs @@ -0,0 +1,41 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_diff_remainder() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<1>::from_data( + TensorData::from([ + 0.9742, 0.3676, 0.0905, 0.8066, 0.7072, 0.7883, 0.6987, 0.1560, 0.7179, 0.7874, 0.9032, + 0.1845, + ]), + &device, + ) + .require_grad(); + let tensor_2 = TestTensor::<1>::from_data( + TensorData::from([ + 0.3357, 0.0285, 0.4115, 0.5511, 0.8637, 0.3593, 0.3885, 0.2569, 0.0936, 0.7172, 0.4792, + 0.4898, + ]), + &device, + ) + .require_grad(); + let tensor_3 = tensor_1.clone().remainder(tensor_2.clone()); + let grads = tensor_3.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([ + -2.0, -12.0, -0.0, -1.0, -0.0, -2.0, -1.0, -0.0, -7.0, -1.0, -1.0, -0.0, + ]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/repeat_dim.rs b/crates/burn-backend-tests/tests/autodiff/repeat_dim.rs new file mode 100644 index 0000000..42ead57 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/repeat_dim.rs @@ -0,0 +1,44 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_diff_repeat() { + let data_1 = TensorData::from([[1.0, 7.0], [-2.0, -3.0]]); + let data_2 = TensorData::from([[4.0], [2.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_2.clone().repeat_dim(1, 3); + + let tensor_3 = tensor_1.matmul(tensor_3); + let grads = tensor_3.backward(); + + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_2 + .to_data() + .assert_eq(&TensorData::from([[-3.0], [12.0]]), false); +} + +#[test] +fn should_diff_repeat_multi_dim() { + let data_1 = TensorData::from([[1.0, 7.0], [-2.0, -3.0]]); + let data_2 = TensorData::from([[4.0, 2.0], [2.0, 4.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_2.clone().repeat_dim(1, 3); + + let tensor_3 = tensor_1.matmul(tensor_3); + let grads = tensor_3.backward(); + + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_2 + .to_data() + .assert_eq(&TensorData::from([[-3.0, -3.0], [12.0, 12.0]]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/reshape.rs b/crates/burn-backend-tests/tests/autodiff/reshape.rs new file mode 100644 index 0000000..29086a7 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/reshape.rs @@ -0,0 +1,26 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_diff_reshape() { + let data_1 = TensorData::from([[1.0, 7.0], [2.0, 3.0]]); + let data_2 = TensorData::from([4.0, 7.0, 2.0, 3.0]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::<1>::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_2.clone().reshape([2, 2]); + let tensor_4 = tensor_1.clone().matmul(tensor_3); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[11.0, 5.0], [11.0, 5.0]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([3.0, 3.0, 10.0, 10.0]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/rfft.rs b/crates/burn-backend-tests/tests/autodiff/rfft.rs new file mode 100644 index 0000000..fae4acb --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/rfft.rs @@ -0,0 +1,201 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; +use burn_tensor::signal; + +#[cfg(not(feature = "ndarray"))] +use burn_tensor::{DType, Element}; + +#[test] +#[cfg(not(feature = "ndarray"))] +fn should_diff_rfft() { + // Lower precisions not supported + if !matches!(FloatElem::dtype(), DType::F32 | DType::F64) { + return; + } + + let device = AutodiffDevice::new(); + + let random1 = TensorData::from([0.26, 0.13, 0.36, 0.24, 0.40, 0.93, 0.12, 0.18]); + let random2 = TensorData::from([0.03, 0.74, 0.33, 0.70, 0.07, 0.61, 0.32, 0.66]); + + let x = TestTensor::<1>::from_data(random1, &device).require_grad(); + let y = TestTensor::<1>::from_data(random2, &device); + + let (x_re, x_im) = signal::rfft(x.clone(), 0, None); + let (y_re, y_im) = signal::rfft(y.clone(), 0, None); + + let loss = (x_re * y_re + x_im * y_im).sum(); + let grads = loss.backward(); + let x_grad = x.grad(&grads).unwrap(); + let prod = x_grad.mul(x.inner()).sum(); + + TensorData::assert_approx_eq::( + &prod.to_data(), + &loss.to_data(), + Tolerance::default(), + ); +} + +#[test] +#[cfg(not(feature = "ndarray"))] +fn round_trip() { + if !matches!(FloatElem::dtype(), DType::F32 | DType::F64) { + return; + } + + let device = AutodiffDevice::new(); + + let random = TensorData::from([ + 0.26, 0.13, 0.36, 0.24, 0.40, 0.93, 0.12, 0.18, 0.03, 0.74, 0.33, 0.70, 0.07, 0.61, 0.32, + 0.66, + ]); + + let tensor = TestTensor::<1>::from_data(random.clone(), &device).require_grad(); + + let y = signal::rfft(tensor.clone() * 3.0, 0, None); + let x = signal::irfft(y.0 * 2.0, y.1 * 2.0, 0, None) / 6.0; + + let loss = x.powi_scalar(2).sum() * 0.5; + let grads = loss.backward(); + let grad = tensor.grad(&grads).unwrap(); + + TensorData::assert_approx_eq::(&grad.to_data(), &random, Tolerance::default()); +} + +#[test] +#[cfg(not(feature = "ndarray"))] +fn round_trip_with_dim_nonzero() { + if !matches!(FloatElem::dtype(), DType::F32 | DType::F64) { + return; + } + + let device = AutodiffDevice::new(); + + let random = TensorData::from([ + 0.26, 0.13, 0.36, 0.24, 0.40, 0.93, 0.12, 0.18, 0.03, 0.74, 0.33, 0.70, 0.07, 0.61, 0.32, + 0.66, + ]); + + let tensor = TestTensor::<1>::from_data(random.clone(), &device); + let tensor = tensor.reshape([1, 1, -1, 1, 1]).require_grad(); + + let y = signal::rfft(tensor.clone() * 3.0, 2, None); + let x = signal::irfft(y.0 * 2.0, y.1 * 2.0, 2, None) / 6.0; + + let loss = x.powi_scalar(2).sum() * 0.5; + let grads = loss.backward(); + let grad = tensor.grad(&grads).unwrap(); + let grad = grad.reshape([-1]); + + TensorData::assert_approx_eq::(&grad.to_data(), &random, Tolerance::default()); +} + +#[test] +#[cfg(not(feature = "ndarray"))] +fn round_trip_with_some_n_greater() { + if !matches!(FloatElem::dtype(), DType::F32 | DType::F64) { + return; + } + + let device = AutodiffDevice::new(); + + let random = TensorData::from([ + 0.26, 0.13, 0.36, 0.24, 0.40, 0.93, 0.12, 0.18, 0.03, 0.74, 0.33, 0.70, 0.07, + ]); + let n = Some(16); + + let tensor = TestTensor::<1>::from_data(random.clone(), &device).require_grad(); + + let y = signal::rfft(tensor.clone() * 3.0, 0, n); + let x = signal::irfft(y.0 * 2.0, y.1 * 2.0, 0, n) / 6.0; + + let loss = x.powi_scalar(2).sum() * 0.5; + let grads = loss.backward(); + let grad = tensor.grad(&grads).unwrap(); + + TensorData::assert_approx_eq::(&grad.to_data(), &random, Tolerance::default()); +} + +#[test] +#[cfg(not(feature = "ndarray"))] +fn round_trip_with_some_n_less() { + if !matches!(FloatElem::dtype(), DType::F32 | DType::F64) { + return; + } + + let device = AutodiffDevice::new(); + + let random = TensorData::from([ + 0.26, 0.13, 0.36, 0.24, 0.40, 0.93, 0.12, 0.18, 0.03, 0.74, 0.33, 0.70, 0.07, 0.61, 0.32, + 0.66, 0.16, 0.05, 0.69, + ]); + let n = Some(16); + + let tensor = TestTensor::<1>::from_data(random.clone(), &device).require_grad(); + + let y = signal::rfft(tensor.clone() * 3.0, 0, n); + let x = signal::irfft(y.0 * 2.0, y.1 * 2.0, 0, n) / 6.0; + + let loss = x.powi_scalar(2).sum() * 0.5; + let grads = loss.backward(); + let grad = tensor.grad(&grads).unwrap(); + + let lhs = grad.slice_dim(0, 0..16); + let rhs = tensor.slice_dim(0, 0..16); + + TensorData::assert_approx_eq::(&lhs.to_data(), &rhs.to_data(), Tolerance::default()); +} + +#[test] +#[cfg(not(feature = "ndarray"))] +fn round_trip_inverse_with_some_n_greater() { + if !matches!(FloatElem::dtype(), DType::F32 | DType::F64) { + return; + } + + let device = AutodiffDevice::new(); + + let random = TensorData::from([0.26, 0.13, 0.36, 0.24, 0.40, 0.93, 0.12]); + let n = Some(16); + + let tensor = TestTensor::<1>::from_data(random.clone(), &device).require_grad(); + + let x = signal::irfft(tensor.clone() * 2.0, tensor.zeros_like(), 0, n) / 6.0; + let (y, _) = signal::rfft(x * 3.0, 0, n); + + let loss = y.powi_scalar(2).sum() * 0.5; + let grads = loss.backward(); + let grad = tensor.grad(&grads).unwrap(); + + TensorData::assert_approx_eq::(&grad.to_data(), &random, Tolerance::default()); +} + +#[test] +#[cfg(not(feature = "ndarray"))] +fn round_trip_inverse_with_some_n_less() { + if !matches!(FloatElem::dtype(), DType::F32 | DType::F64) { + return; + } + + let device = AutodiffDevice::new(); + + let random = TensorData::from([ + 0.26, 0.13, 0.36, 0.24, 0.40, 0.93, 0.12, 0.18, 0.03, 0.74, 0.33, 0.70, + ]); + let n = Some(16); + + let tensor = TestTensor::<1>::from_data(random.clone(), &device).require_grad(); + + let x = signal::irfft(tensor.clone() * 2.0, tensor.zeros_like(), 0, n) / 6.0; + let (y, _) = signal::rfft(x * 3.0, 0, n); + + let loss = y.powi_scalar(2).sum() * 0.5; + let grads = loss.backward(); + let grad = tensor.grad(&grads).unwrap(); + + let lhs = grad.slice_dim(0, 0..9); + let rhs = tensor.slice_dim(0, 0..9); + + TensorData::assert_approx_eq::(&lhs.to_data(), &rhs.to_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/round.rs b/crates/burn-backend-tests/tests/autodiff/round.rs new file mode 100644 index 0000000..b138697 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/round.rs @@ -0,0 +1,20 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_diff_round() { + let data = TensorData::from([ + [-1.9751, 0.0714, 0.0643, 0.2406], + [-1.3172, 0.1252, -0.1119, -0.0127], + ]); + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data.clone(), &device).require_grad(); + let tensor_2 = tensor_1.clone().round(); + let grads = tensor_2.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + grad_1.to_data().assert_eq( + &TensorData::from([[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]]), + false, + ); +} diff --git a/crates/burn-backend-tests/tests/autodiff/select.rs b/crates/burn-backend-tests/tests/autodiff/select.rs new file mode 100644 index 0000000..0ed9abe --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/select.rs @@ -0,0 +1,87 @@ +use super::*; +use burn_tensor::{IndexingUpdateOp, TensorData}; + +#[test] +fn test_select_grad() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data( + TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]), + &device, + ) + .require_grad(); + let indices = TestTensorInt::<1>::from_data(TensorData::from([1, 0]), &device); + + let tensor_2 = tensor_1.clone().matmul(tensor_1.clone().transpose()); + let tensor_3 = tensor_1.clone().select(0, indices); + let tensor_4 = tensor_2.matmul(tensor_3); + + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + + grad_1.into_data().assert_eq( + &TensorData::from([[109., 148., 187.], [37., 58., 79.]]), + false, + ); +} + +#[test] +fn test_select_add_grad() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data( + TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]), + &device, + ) + .require_grad(); + let values = TestTensor::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), + &device, + ) + .require_grad(); + let indices = TestTensorInt::<1>::from_data(TensorData::from([1, 0]), &device); + + let tensor_2 = tensor_1.clone().matmul(tensor_1.clone().transpose()); + let tensor_3 = + tensor_1 + .clone() + .select_assign(0, indices, values.clone(), IndexingUpdateOp::Add); + let tensor_4 = tensor_2.matmul(tensor_3); + + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = values.grad(&grads).unwrap(); + + grad_1.into_data().assert_eq( + &TensorData::from([[127., 199., 271.], [172., 244., 316.]]), + false, + ); + grad_2 + .into_data() + .assert_eq(&TensorData::from([[64., 64., 64.], [19., 19., 19.]]), false); +} + +#[test] +fn test_select_add_grad_different_shapes() { + let device = AutodiffDevice::new(); + + let indices = TestTensorInt::from_ints([1], &device); + let x = TestTensor::<2>::ones([1, 1], &device).require_grad(); + let y = TestTensor::ones([2, 1], &device).require_grad(); + + let w = y + .clone() + .select_assign(0, indices, x.clone(), IndexingUpdateOp::Add); + let w = w.matmul(y.clone().transpose()); + + let grads = w.backward(); + let x_grad = x.grad(&grads).unwrap(); + let y_grad = y.grad(&grads).unwrap(); + + x_grad + .into_data() + .assert_eq(&TensorData::from([[2.0]]), false); + y_grad + .into_data() + .assert_eq(&TensorData::from([[5.0], [5.0]]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/sigmoid.rs b/crates/burn-backend-tests/tests/autodiff/sigmoid.rs new file mode 100644 index 0000000..5ac1fc3 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/sigmoid.rs @@ -0,0 +1,35 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn should_diff_sigmoid() { + let data = TensorData::from([0.8762]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<1>::from_data(data, &device).require_grad(); + let tensor_2 = activation::sigmoid(tensor_1.clone()); + let grads = tensor_2.backward(); + + let grad = tensor_1.grad(&grads).unwrap(); + + let expected = TensorData::from([0.207549]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn small_neg_val_should_not_cause_grad_overflow() { + let data = TensorData::from([-90.0]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<1>::from_data(data, &device).require_grad(); + let tensor_2 = activation::sigmoid(tensor_1.clone()); + let grads = tensor_2.backward(); + + let grad = tensor_1.grad(&grads).unwrap(); + + let expected = TensorData::from([0.0]); + grad.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/sign.rs b/crates/burn-backend-tests/tests/autodiff/sign.rs new file mode 100644 index 0000000..48d7132 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/sign.rs @@ -0,0 +1,42 @@ +use super::*; +use burn_tensor::TensorData; + +/// Example using the sign function with PyTorch: +// >>> import torch +// >>> # Create a tensor with requires_grad=True +// >>> x = torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0], requires_grad=True) +// >>> # Forward pass: Apply the sign function +// >>> y = torch.sign(x) +// >>> print("Forward pass:") +// Forward pass: +// >>> print("x:", x) +// x: tensor([-2., -1., 0., 1., 2.], requires_grad=True) +// >>> print("y:", y) +// y: tensor([-1., -1., 0., 1., 1.], grad_fn=) +// >>> # Compute the loss (just an example) +// >>> loss = y.sum() +// >>> # Backward pass: Compute the gradients +// >>> loss.backward() +// >>> print("\nBackward pass:") +// Backward pass: +// >>> print("x.grad:", x.grad) +// x.grad: tensor([0., 0., 0., 0., 0.]) + +#[test] +fn should_diff_sign() { + let data = TensorData::from([-2.0, -1.0, 0.0, 1.0, 2.0]); + + let device = AutodiffDevice::new(); + let x = TestTensor::<1>::from_data(data, &device).require_grad(); + + let y = x.clone().sign(); + + let loss = y.clone().sum(); + let grads = loss.backward(); + let grad = x.grad(&grads).unwrap(); + + y.to_data() + .assert_eq(&TensorData::from([-1., -1., 0., 1., 1.]), false); + grad.to_data() + .assert_eq(&TensorData::from([0., 0., 0., 0., 0.]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/slice.rs b/crates/burn-backend-tests/tests/autodiff/slice.rs new file mode 100644 index 0000000..9e6141e --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/slice.rs @@ -0,0 +1,67 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_diff_matmul_with_slice() { + let data_1 = TensorData::from([[1.0, 7.0], [2.0, 3.0]]); + let data_2 = TensorData::from([[4.0, 7.0, 100.0], [2.0, 3.0, 15.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_2.clone().slice([0..2, 0..2]); + let tensor_4 = tensor_1.clone().matmul(tensor_3); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[11.0, 5.0], [11.0, 5.0]]), false); + grad_2.to_data().assert_eq( + &TensorData::from([[3.0, 3.0, 0.0], [10.0, 10.0, 0.0]]), + false, + ); +} + +#[test] +fn should_diff_matmul_with_slice_stepped() { + use burn_tensor::s; + let data_1 = TensorData::from([[1.0, 7.0], [100.0, 100.0], [2.0, 3.0], [100.0, 100.0]]); + let data_2 = TensorData::from([[4.0, 100.0, 7.0, 100.0], [2.0, 100.0, 3.0, 15.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().slice(s![0..;2, 0..2]); // [[1., 7.], [2., 3.]] + let tensor_4 = tensor_2.clone().slice(s![0..2, 0..;2]); // [[4., 7.], [2., 3.]] + let tensor_5 = tensor_3.clone().matmul(tensor_4); + let grads = tensor_5.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1.to_data().assert_eq( + &TensorData::from([[11., 5.], [0., 0.], [11., 5.], [0., 0.]]), + false, + ); + grad_2.to_data().assert_eq( + &TensorData::from([[3., 0., 3., 0.], [10., 0., 10., 0.]]), + false, + ); +} + +#[test] +fn should_panic_on_slice_with_step() { + use burn_tensor::s; + + let data = TensorData::from([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]); + let device = AutodiffDevice::new(); + let tensor = TestTensor::<2>::from_data(data, &device).require_grad(); + + // This should panic because step is 2 + let _sliced = tensor.slice(s![.., 0..4; 2]); +} diff --git a/crates/burn-backend-tests/tests/autodiff/slice_assign.rs b/crates/burn-backend-tests/tests/autodiff/slice_assign.rs new file mode 100644 index 0000000..bf47e7e --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/slice_assign.rs @@ -0,0 +1,163 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_diff_matmul_with_slice_assign() { + let data_1 = TensorData::from([[1.0, 7.0], [2.0, 3.0]]); + let data_2 = TensorData::from([[4.0, 7.0], [2.0, 3.0]]); + let data_assigned = TensorData::from([[9.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + let tensor_assigned = TestTensor::from_data(data_assigned, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_3.slice_assign([0..1, 0..1], tensor_assigned); + let tensor_5 = tensor_4.matmul(tensor_1.clone()); + + let grads = tensor_5.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[58.0, 38.0], [118.0, 82.0]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([[16.0, 15.0], [24.0, 50.0]]), false); +} + +#[test] +fn should_diff_matmul_with_slice_assign_complex() { + let data_1 = TensorData::from([[1.0, 7.0], [2.0, 3.0]]); + let data_2 = TensorData::from([[4.0, 7.0], [2.0, 3.0]]); + let data_3 = TensorData::from([[9.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + let tensor_3 = TestTensor::from_data(data_3, &device).require_grad(); + + let tensor_4 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_5 = tensor_2.clone().slice([0..1, 0..1]); + let tensor_6 = tensor_5.mul(tensor_3.clone()); + let tensor_7 = tensor_4.slice_assign([0..1, 0..1], tensor_6); + let tensor_8 = tensor_7.matmul(tensor_1.clone()); + + let grads = tensor_8.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + let grad_3 = tensor_3.grad(&grads).unwrap(); + + grad_3 + .to_data() + .assert_eq(&TensorData::from([[32.0]]), false); + grad_1 + .to_data() + .assert_eq(&TensorData::from([[85.0, 65.0], [118.0, 82.0]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([[88.0, 15.0], [24.0, 50.0]]), false); +} + +#[test] +fn slice_assign_diff_should_give_same_results_as_cat() { + let data_1 = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[5.0, 6.0], [7.0, 8.0]]); + let data_3 = TensorData::from([[14.0, 97.0, 100.0, 9.0], [2.0, 3.0, 15.0, 7.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + let tensor_3 = TestTensor::from_data(data_3, &device); + + let slice_assign_output = TestTensor::zeros([2, 4], &AutodiffDevice::new()); + let slice_assign_output = slice_assign_output.slice_assign([0..2, 0..2], tensor_1.clone()); + let slice_assign_output = slice_assign_output.slice_assign([0..2, 2..4], tensor_2.clone()); + let slice_assign_output = slice_assign_output / tensor_3.clone(); + + let cat_output = TestTensor::cat(vec![tensor_1.clone(), tensor_2.clone()], 1); + let cat_output = cat_output / tensor_3; + + slice_assign_output + .to_data() + .assert_approx_eq::(&cat_output.to_data(), Tolerance::default()); + + let slice_assign_grads = slice_assign_output.backward(); + let cat_grads = cat_output.backward(); + + let slice_assign_grad_1 = tensor_1.grad(&slice_assign_grads).unwrap(); + let slice_assign_grad_2 = tensor_2.grad(&slice_assign_grads).unwrap(); + let cat_grad_1 = tensor_1.grad(&cat_grads).unwrap(); + let cat_grad_2 = tensor_2.grad(&cat_grads).unwrap(); + + slice_assign_grad_1 + .to_data() + .assert_approx_eq::(&cat_grad_1.to_data(), Tolerance::default()); + slice_assign_grad_2 + .to_data() + .assert_approx_eq::(&cat_grad_2.to_data(), Tolerance::default()); +} + +#[test] +fn should_diff_slice_assign_with_step() { + use burn_tensor::s; + let data = TensorData::from([[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]]); + let value_data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + + let device = AutodiffDevice::new(); + let tensor = TestTensor::<2>::from_data(data, &device).require_grad(); + let value = TestTensor::<2>::from_data(value_data, &device).require_grad(); + + // Assign with step=2 + let result = tensor.clone().slice_assign(s![.., 0..4; 2], value.clone()); + let result = result * 2.0; // Scale to create gradients + let grads = result.backward(); + + let grad_tensor = tensor.grad(&grads).unwrap(); + let grad_value = value.grad(&grads).unwrap(); + + // The gradient for tensor should be 2.0 everywhere except the assigned positions + grad_tensor.to_data().assert_eq( + &TensorData::from([[0.0, 2.0, 0.0, 2.0], [0.0, 2.0, 0.0, 2.0]]), + false, + ); + // The gradient for value should be 2.0 at all positions + grad_value + .to_data() + .assert_eq(&TensorData::from([[2.0, 2.0], [2.0, 2.0]]), false); +} + +#[test] +fn should_diff_slice_assign_with_negative_step() { + use burn_tensor::s; + + let data = TensorData::from([[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]]); + let value_data = TensorData::from([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]); + let device = AutodiffDevice::new(); + let tensor = TestTensor::<2>::from_data(data, &device).require_grad(); + let value = TestTensor::<2>::from_data(value_data, &device).require_grad(); + + // Assign with step=-1 (reverse order, all elements) + let result = tensor.clone().slice_assign(s![.., ..;-1], value.clone()); + let result = result * 2.0; // Scale to create gradients + let grads = result.backward(); + + let grad_tensor = tensor.grad(&grads).unwrap(); + let grad_value = value.grad(&grads).unwrap(); + + // The gradient for tensor should be 0 since all values were replaced + grad_tensor.to_data().assert_eq( + &TensorData::from([[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]]), + false, + ); + // The gradient for value should be 2.0 at all positions + grad_value.to_data().assert_eq( + &TensorData::from([[2.0, 2.0, 2.0, 2.0], [2.0, 2.0, 2.0, 2.0]]), + false, + ); +} diff --git a/crates/burn-backend-tests/tests/autodiff/softmax.rs b/crates/burn-backend-tests/tests/autodiff/softmax.rs new file mode 100644 index 0000000..7900f8b --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/softmax.rs @@ -0,0 +1,90 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_softmax_grad() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[6.0, 7.0], [9.0, 10.0]]); + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = activation::softmax(tensor_3, 1).matmul(tensor_2.clone()); + + let grads = tensor_4.backward(); + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[1.179665, 1.179661], [0.005462, 0.005463]]); + + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(0.05, 0.5)); + + let expected = TensorData::from([[0.253469, 0.286237], [0.528630, 2.931664]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(0.05, 0.05)); +} + +#[test] +fn test_log_softmax_grad() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[6.0, 7.0], [9.0, 10.0]]); + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = activation::log_softmax(tensor_3, 1).matmul(tensor_2.clone()); + + let grads = tensor_4.backward(); + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[-4.3939, -4.3939], [-12.9709, -12.9709]]); + // f16 gradients from log-softmax + matmul amplify error, so we increase the tolerance + // to account for limited precision and large representable step sizes in this range. + let tolerance = Tolerance::permissive().set_half_precision_relative(6e-2); + + grad_1 + .to_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([[30.5984, -47.2267], [55.9631, -56.5914]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn test_quiet_softmax_grad() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[6.0, 7.0], [9.0, 10.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = activation::softmax(tensor_3, 1).matmul(tensor_2.clone()); + + let grads = tensor_4.backward(); + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[1.179665, 1.179661], [0.005462, 0.005463]]); + + // Precision is quite bad yet on softmax grad especially with half precision. + let tolerance = Tolerance::rel_abs(0.5, 0.2); + grad_1 + .to_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([[0.253469, 0.286237], [0.528630, 2.931664]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, tolerance); +} diff --git a/crates/burn-backend-tests/tests/autodiff/sort.rs b/crates/burn-backend-tests/tests/autodiff/sort.rs new file mode 100644 index 0000000..eb22eaf --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/sort.rs @@ -0,0 +1,76 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_diff_sort() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data([[1.0, 7.0], [-2.0, -3.0]], &device).require_grad(); + let tensor_2 = TestTensor::from_data([[4.0, -7.0], [2.0, 3.0]], &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_1.clone().mul(tensor_3.sort(1)); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[35.0, 35.0], [-1.0, -8.0]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[11.0, 7.0], [55.0, 16.0]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_sort_with_indices() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data([[1.0, 7.0], [-2.0, -3.0]], &device).require_grad(); + let tensor_2 = TestTensor::from_data([[4.0, -7.0], [2.0, 3.0]], &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let (values, _indices) = tensor_3.sort_with_indices(1); + let tensor_4 = tensor_1.clone().mul(values); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[35.0, 35.0], [-1.0, -8.0]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[11.0, 7.0], [55.0, 16.0]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_diff_sort_3d_dim1() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<3>::from_data([[[1.0, 7.0], [-2.0, -3.0]]], &device).require_grad(); + let tensor_2 = TestTensor::from_data([[[4.0, -7.0], [2.0, 3.0]]], &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone()); + let tensor_4 = tensor_1.clone().mul(tensor_3.sort(1)); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let expected = TensorData::from([[[-1., -8.], [-27., 37.]]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[[-4., -17.], [-17., -42.]]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/autodiff/sqrt.rs b/crates/burn-backend-tests/tests/autodiff/sqrt.rs new file mode 100644 index 0000000..9d5ecd7 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/sqrt.rs @@ -0,0 +1,31 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_diff_sqrt() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[6.0, 7.0], [9.0, 10.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().sqrt()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let tolerance = Tolerance::default().set_half_precision_relative(1e-3); + let expected = TensorData::from([[82.112640, 99.083275], [82.112640, 99.083275]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([[30.309311, 33.120457], [34.581974, 38.769463]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, tolerance); +} diff --git a/crates/burn-backend-tests/tests/autodiff/sub.rs b/crates/burn-backend-tests/tests/autodiff/sub.rs new file mode 100644 index 0000000..5ddddd7 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/sub.rs @@ -0,0 +1,73 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_diff_sub() { + let data_1 = TensorData::from([2.0, 5.0]); + let data_2 = TensorData::from([4.0, 1.0]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<1>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().sub(tensor_2.clone()); + let grads = tensor_3.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([1.0, 1.0]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([-1.0, -1.0]), false); + + tensor_3 + .into_data() + .assert_eq(&TensorData::from([-2.0, 4.0]), false); +} + +#[test] +fn should_diff_sub_scalar() { + let data = TensorData::from([2.0, 10.0]); + let tensor = TestTensor::<1>::from_data(data, &AutodiffDevice::new()).require_grad(); + let tensor_out = tensor.clone().sub_scalar(5.0); + let grads = tensor_out.backward(); + + let grad = tensor.grad(&grads).unwrap(); + + grad.to_data() + .assert_eq(&TensorData::from([1.0, 1.0]), false); + tensor_out + .into_data() + .assert_eq(&TensorData::from([-3.0, 5.0]), false); +} + +#[test] +fn test_sub_complex_1() { + let data_1 = TensorData::from([[1.0, 7.0], [13.0, -3.0]]); + let data_2 = TensorData::from([[4.0, 7.0], [2.0, 3.0]]); + let data_3 = TensorData::from([[2.0, 2.0], [2.0, 2.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + let tensor_3 = TestTensor::from_data(data_3, &device).require_grad(); + + let tensor_4 = tensor_1.clone().sub(tensor_2.clone()); + let tensor_5 = tensor_4.sub(tensor_3).sub_scalar(5.0); + let tensor_6 = tensor_1.clone().sub(tensor_5); + + let grads = tensor_6.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1 + .to_data() + .assert_eq(&TensorData::from([[0.0, 0.0], [0.0, 0.0]]), false); + grad_2 + .to_data() + .assert_eq(&TensorData::from([[1.0, 1.0], [1.0, 1.0]]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff/transpose.rs b/crates/burn-backend-tests/tests/autodiff/transpose.rs new file mode 100644 index 0000000..bedbbae --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/transpose.rs @@ -0,0 +1,60 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_diff_transpose() { + let data_1 = TensorData::from([[1.0, 7.0], [2.0, 3.0]]); + let data_2 = TensorData::from([[4.0, 7.0], [2.0, 3.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().transpose()); + let tensor_4 = tensor_3.transpose(); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1.to_data().assert_approx_eq::( + &TensorData::from([[6.0, 10.0], [6.0, 10.0]]), + Tolerance::default(), + ); + grad_2.to_data().assert_approx_eq::( + &TensorData::from([[3.0, 10.0], [3.0, 10.0]]), + Tolerance::default(), + ); +} + +#[test] +fn should_diff_swap_dims() { + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<3>::from_data( + [[[0.0, 1.0], [3.0, 4.0]], [[6.0, 7.0], [9.0, 10.0]]], + &device, + ) + .require_grad(); + let tensor_2 = TestTensor::from_data( + [[[1.0, 4.0], [2.0, 5.0]], [[7.0, 10.0], [8.0, 11.0]]], + &device, + ) + .require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().swap_dims(0, 2)); + let tensor_4 = tensor_3.matmul(tensor_2.clone().swap_dims(1, 2)); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1.to_data().assert_approx_eq::( + &TensorData::from([[[66., 78.], [66., 78.]], [[270., 306.], [270., 306.]]]), + Tolerance::default(), + ); + grad_2.to_data().assert_approx_eq::( + &TensorData::from([[[22., 286.], [28., 316.]], [[172., 652.], [190., 694.]]]), + Tolerance::default(), + ); +} diff --git a/crates/burn-backend-tests/tests/autodiff/trig.rs b/crates/burn-backend-tests/tests/autodiff/trig.rs new file mode 100644 index 0000000..20e3d28 --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/trig.rs @@ -0,0 +1,371 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_diff_cos() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[6.0, 7.0], [9.0, 10.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().cos()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + // Metal has less precise trigonometric functions + let tolerance = Tolerance::default().set_half_precision_relative(1e-2); + + grad_1.to_data().assert_approx_eq::( + &TensorData::from([[26.8063, -27.7870], [26.8063, -27.7870]]), + tolerance, + ); + + grad_2.to_data().assert_approx_eq::( + &TensorData::from([[9.222064, -39.123375], [-28.721354, 49.748356]]), + tolerance, + ); +} + +#[test] +fn should_diff_sin() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[6.0, 7.0], [9.0, 10.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().sin()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + // Metal has less precise trigonometric functions + let tolerance = Tolerance::default().set_half_precision_relative(1e-2); + + let expected = TensorData::from([[8.8500, -4.9790], [8.8500, -4.9790]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([[38.668987, 44.194775], [-59.97261, -80.46094]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn should_diff_tanh() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[6.0, 7.0], [9.0, 10.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().tanh()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + let tolerance = Tolerance::default().set_half_precision_relative(8e-3); + let expected = TensorData::from([[32.0, 32.0], [32.0, 32.0]]); + grad_1 + .to_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([[8.00092, 8.000153], [8.000003, 7.999995]]); + grad_2 + .to_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn should_diff_cosh() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[0.5, 1.0], [1.5, 2.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().cosh()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1.to_data().assert_approx_eq::( + &TensorData::from([[7.092221, 16.696301], [7.092221, 16.696301]]), + Tolerance::default(), + ); + + grad_2.to_data().assert_approx_eq::( + &TensorData::from([[17.489855, 27.484539], [39.409813, 86.910278]]), + Tolerance::default(), + ); +} + +#[test] +fn should_diff_sinh() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[0.5, 1.0], [1.5, 2.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().sinh()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1.to_data().assert_approx_eq::( + &TensorData::from([[4.894847, 15.887931], [4.894847, 15.887931]]), + Tolerance::default(), + ); + + grad_2.to_data().assert_approx_eq::( + &TensorData::from([[17.284000, 28.412029], [39.302979, 87.498329]]), + Tolerance::default(), + ); +} + +#[test] +fn should_diff_tan() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[0.5, 1.0], [0.3, 0.8]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().tan()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1.to_data().assert_approx_eq::( + &TensorData::from([[2.532602, 1.596607], [2.532602, 1.596607]]), + Tolerance::default(), + ); + + grad_2.to_data().assert_approx_eq::( + &TensorData::from([[9.028598, 14.489801], [18.038082, 21.151270]]), + Tolerance::default(), + ); +} + +#[test] +fn should_diff_asin() { + let data_1 = TensorData::from([[0.0, 0.1], [0.3, 0.4]]); + let data_2 = TensorData::from([[0.2, 0.3], [0.5, 0.6]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().asin()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1.to_data().assert_approx_eq::( + &TensorData::from([[0.435841, 0.969651], [0.435841, 0.969651]]), + Tolerance::default(), + ); + + grad_2.to_data().assert_approx_eq::( + &TensorData::from([[0.475300, 0.668141], [0.701834, 1.100658]]), + Tolerance::default(), + ); +} + +#[test] +fn should_diff_acos() { + let data_1 = TensorData::from([[0.0, 0.1], [0.3, 0.4]]); + let data_2 = TensorData::from([[0.2, 0.3], [0.5, 0.6]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().acos()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1.to_data().assert_approx_eq::( + &TensorData::from([[2.077433, 1.543624], [2.077433, 1.543624]]), + Tolerance::default(), + ); + + grad_2.to_data().assert_approx_eq::( + &TensorData::from([[0.781337, 0.588496], [0.554804, 0.155979]]), + Tolerance::default(), + ); +} + +#[test] +fn should_diff_atan() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[0.5, 1.0], [1.5, 2.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().atan()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1.to_data().assert_approx_eq::( + &TensorData::from([[3.444365, 5.349211], [3.444365, 5.349211]]), + Tolerance::default(), + ); + + grad_2.to_data().assert_approx_eq::( + &TensorData::from([[9.904911, 11.554912], [10.199631, 11.391938]]), + Tolerance::default(), + ); +} + +#[test] +fn should_diff_asinh() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[0.5, 1.0], [1.5, 2.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().asinh()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1.to_data().assert_approx_eq::( + &TensorData::from([[3.806625, 6.844869], [3.806625, 6.844869]]), + Tolerance::default(), + ); + + grad_2.to_data().assert_approx_eq::( + &TensorData::from([[11.442373, 14.842072], [14.022551, 17.688538]]), + Tolerance::default(), + ); +} + +#[test] +fn should_diff_acosh() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[1.5, 2.0], [2.5, 3.0]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().acosh()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1.to_data().assert_approx_eq::( + &TensorData::from([[10.611752, 15.178907], [10.611752, 15.178907]]), + Tolerance::default(), + ); + + grad_2.to_data().assert_approx_eq::( + &TensorData::from([[20.112753, 20.247547], [20.402235, 22.487328]]), + Tolerance::default(), + ); +} + +#[test] +fn should_diff_atanh() { + let data_1 = TensorData::from([[0.0, 0.1], [0.3, 0.4]]); + let data_2 = TensorData::from([[0.2, 0.3], [0.5, 0.6]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + + let tensor_3 = tensor_1.clone().matmul(tensor_2.clone().atanh()); + let tensor_4 = tensor_3.matmul(tensor_2.clone()); + let grads = tensor_4.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + + grad_1.to_data().assert_approx_eq::( + &TensorData::from([[0.441838, 1.037115], [0.441838, 1.037115]]), + Tolerance::default(), + ); + + grad_2.to_data().assert_approx_eq::( + &TensorData::from([[0.491723, 0.698110], [0.772763, 1.298805]]), + Tolerance::default(), + ); +} + +#[test] +fn should_diff_atan2() { + let data_1 = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + let data_2 = TensorData::from([[0.5, 1.0], [1.5, 2.0]]); + let data_3 = TensorData::from([[1.0, 0.5], [2.0, 1.5]]); + + let device = AutodiffDevice::new(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device).require_grad(); + let tensor_2 = TestTensor::from_data(data_2, &device).require_grad(); + let tensor_3 = TestTensor::from_data(data_3, &device).require_grad(); + + let tensor_4 = tensor_1 + .clone() + .matmul(tensor_2.clone().atan2(tensor_3.clone())); + let tensor_5 = tensor_4.matmul(tensor_2.clone()); + let grads = tensor_5.backward(); + + let grad_1 = tensor_1.grad(&grads).unwrap(); + let grad_2 = tensor_2.grad(&grads).unwrap(); + let grad_3 = tensor_3.grad(&grads).unwrap(); + + grad_1.to_data().assert_approx_eq::( + &TensorData::from([[4.570492, 4.210785], [4.570492, 4.210785]]), + Tolerance::default(), + ); + + grad_2.to_data().assert_approx_eq::( + &TensorData::from([[8.208448, 8.808449], [10.357923, 12.157923]]), + Tolerance::default(), + ); + + grad_3.to_data().assert_approx_eq::( + &TensorData::from([[-1.8, -8.4], [-1.8, -5.6]]), + Tolerance::default(), + ); +} diff --git a/crates/burn-backend-tests/tests/autodiff/unfold.rs b/crates/burn-backend-tests/tests/autodiff/unfold.rs new file mode 100644 index 0000000..674ad4a --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff/unfold.rs @@ -0,0 +1,18 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn unfold_backward_accumulates_overlaps() { + let device = AutodiffDevice::new(); + let x = TestTensor::<2>::from_data([[1.0, 2.0, 3.0, 4.0]], &device).require_grad(); + + let y = x.clone().unfold::<3, _>(1, 2, 1); + let loss = y.sum(); + + let grads = loss.backward(); + let grad_x = x.grad(&grads).unwrap(); + + grad_x + .to_data() + .assert_eq(&TensorData::from([[1., 2., 2., 1.]]), false); +} diff --git a/crates/burn-backend-tests/tests/autodiff_f16.rs b/crates/burn-backend-tests/tests/autodiff_f16.rs new file mode 100644 index 0000000..ade3aef --- /dev/null +++ b/crates/burn-backend-tests/tests/autodiff_f16.rs @@ -0,0 +1,23 @@ +//! Burn autodiff tests. + +#![recursion_limit = "256"] +#![cfg(any( + feature = "vulkan", + feature = "cuda", + feature = "rocm", + feature = "metal" +))] + +extern crate alloc; + +pub type FloatElem = burn_tensor::f16; +#[allow(unused)] +pub type IntElem = i32; + +#[path = "common/backend.rs"] +mod backend; +pub use backend::*; + +#[allow(clippy::module_inception)] +#[path = "common/autodiff.rs"] +mod autodiff; diff --git a/crates/burn-backend-tests/tests/common/autodiff.rs b/crates/burn-backend-tests/tests/common/autodiff.rs new file mode 100644 index 0000000..487243a --- /dev/null +++ b/crates/burn-backend-tests/tests/common/autodiff.rs @@ -0,0 +1,34 @@ +// Burn autodiff tests, reusable with element types. + +pub use super::*; + +// Autodiff-enabled device used for tests. +pub struct AutodiffDevice; + +impl AutodiffDevice { + #[allow(clippy::new_ret_no_self)] + pub fn new() -> burn_tensor::Device { + burn_tensor::Device::default().autodiff() + } +} + +#[path = "../autodiff/mod.rs"] +mod base; + +mod checkpointing { + pub use super::FloatElem; + + // Override autodiff device + pub struct AutodiffDevice; + + impl AutodiffDevice { + #[allow(clippy::new_ret_no_self)] + pub fn new() -> burn_tensor::Device { + burn_tensor::Device::default() + .autodiff() + .gradient_checkpointing() + } + } + + include!("../autodiff/mod.rs"); +} diff --git a/crates/burn-backend-tests/tests/common/backend.rs b/crates/burn-backend-tests/tests/common/backend.rs new file mode 100644 index 0000000..10f1a8b --- /dev/null +++ b/crates/burn-backend-tests/tests/common/backend.rs @@ -0,0 +1,31 @@ +use burn_tensor::Element; +use ctor::ctor; + +// Re-export +use super::{FloatElem, IntElem}; + +#[ctor] +fn init_device_settings() { + let mut device = burn_tensor::Device::default(); + device + .configure( + burn_tensor::DeviceConfig::default() + .float_dtype(::dtype()) + .int_dtype(::dtype()), + ) + .unwrap(); +} + +/// Collection of types used across tests +#[allow(unused)] +pub mod prelude { + pub use burn_tensor::Tensor; + + use super::*; + pub type TestTensor = Tensor; + pub type TestTensorInt = Tensor; + pub type TestTensorBool = Tensor; +} + +#[allow(unused)] +pub use prelude::*; diff --git a/crates/burn-backend-tests/tests/common/tensor.rs b/crates/burn-backend-tests/tests/common/tensor.rs new file mode 100644 index 0000000..858a15b --- /dev/null +++ b/crates/burn-backend-tests/tests/common/tensor.rs @@ -0,0 +1,26 @@ +// Burn backend tensor tests, reusable with element types. + +pub use super::*; + +#[path = "../tensor/clone_invariance.rs"] +mod clone_invariance; + +#[cfg(feature = "std")] +#[path = "../tensor/multi_threads.rs"] +mod multi_threads; + +#[cfg(feature = "distributed")] +#[path = "../tensor/distributed.rs"] +mod distributed; + +// Default float dtype +#[path = "../tensor/float/mod.rs"] +mod float; + +// Default integer dtype +#[path = "../tensor/int/mod.rs"] +mod int; + +// Default bool dtype +#[path = "../tensor/bool/mod.rs"] +mod bool; diff --git a/crates/burn-backend-tests/tests/cubecl.rs b/crates/burn-backend-tests/tests/cubecl.rs new file mode 100644 index 0000000..92112be --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl.rs @@ -0,0 +1,25 @@ +//! CubeCL kernel tests. +#![cfg(feature = "cube")] +#![recursion_limit = "256"] + +#[path = "."] +mod cube { + type FloatElem = f32; + type IntElem = i32; + + mod backend { + include!("common/backend.rs"); + + pub struct ReferenceDevice; + + impl ReferenceDevice { + pub fn new() -> burn_tensor::Device { + burn_ndarray::NdArrayDevice::Cpu.into() + } + } + } + pub use backend::*; + + #[path = "cubecl/mod.rs"] + mod kernel; +} diff --git a/crates/burn-backend-tests/tests/cubecl/avg_pool2d.rs b/crates/burn-backend-tests/tests/cubecl/avg_pool2d.rs new file mode 100644 index 0000000..bf3c7f3 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/avg_pool2d.rs @@ -0,0 +1,87 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{Device, Distribution, module}; + +#[test] +fn avg_pool2d_should_match_reference_backend() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::<4>::random([32, 32, 32, 32], Distribution::Default, &device); + let tensor_ref = TestTensor::<4>::from_data(tensor.to_data(), &ref_device); + let kernel_size = [3, 4]; + let stride = [1, 2]; + let padding = [1, 2]; + let count_include_pad = true; + + let pooled = module::avg_pool2d( + tensor, + kernel_size, + stride, + padding, + count_include_pad, + false, + ); + let pooled_ref = module::avg_pool2d( + tensor_ref, + kernel_size, + stride, + padding, + count_include_pad, + false, + ); + + pooled + .into_data() + .assert_approx_eq::(&pooled_ref.into_data(), Tolerance::default()); +} + +#[test] +fn avg_pool2d_backward_should_match_reference_backend() { + let device = Device::default(); + let ref_device = ReferenceDevice::new(); + + device.seed(0); + ref_device.seed(0); + + let tensor = TestTensor::<4>::random([32, 32, 32, 32], Distribution::Default, &device); + let tensor_ref = TestTensor::<4>::from_data(tensor.to_data(), &ref_device); + let kernel_size = [3, 3]; + let stride = [1, 1]; + let padding = [1, 1]; + let count_include_pad = true; + + let shape_out = module::avg_pool2d( + tensor.clone(), + kernel_size, + stride, + padding, + count_include_pad, + false, + ) + .shape(); + let grad_output = TestTensor::<4>::random(shape_out, Distribution::Default, &device); + let grad_output_ref = TestTensor::<4>::from_data(grad_output.to_data(), &ref_device); + + let grad = module::avg_pool2d_backward( + tensor, + grad_output, + kernel_size, + stride, + padding, + count_include_pad, + false, + ); + let grad_ref = module::avg_pool2d_backward( + tensor_ref, + grad_output_ref, + kernel_size, + stride, + padding, + count_include_pad, + false, + ); + + grad.into_data() + .assert_approx_eq::(&grad_ref.into_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/cubecl/bernoulli.rs b/crates/burn-backend-tests/tests/cubecl/bernoulli.rs new file mode 100644 index 0000000..d70ee2d --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/bernoulli.rs @@ -0,0 +1,43 @@ +use super::*; + +use serial_test::serial; + +use core::f32; + +use burn_tensor::{Device, Distribution, Shape}; + +use cubek::random::{assert_number_of_1_proportional_to_prob, assert_wald_wolfowitz_runs_test}; + +#[test] +#[serial] +fn number_of_1_proportional_to_prob() { + let device = Device::default(); + device.seed(0); + + let shape: Shape = [40, 40].into(); + let prob = 0.7; + + let tensor = + TestTensor::<2>::random(shape.clone(), Distribution::Bernoulli(prob), &device).into_data(); + + let numbers = tensor.as_slice::().unwrap(); + + assert_number_of_1_proportional_to_prob(numbers, prob as f32); +} + +#[test] +#[serial] +fn wald_wolfowitz_runs_test() { + let device = Device::default(); + device.seed(0); + + let shape = Shape::new([512, 512]); + let device = Device::default(); + let tensor = TestTensor::<2>::random(shape, Distribution::Bernoulli(0.5), &device); + + let data = tensor.into_data(); + let numbers = data.as_slice::().unwrap(); + + // High bound slightly over 1 so 1.0 is included in second bin + assert_wald_wolfowitz_runs_test(numbers, 0., 1.1); +} diff --git a/crates/burn-backend-tests/tests/cubecl/cast.rs b/crates/burn-backend-tests/tests/cubecl/cast.rs new file mode 100644 index 0000000..a36b0c0 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/cast.rs @@ -0,0 +1,45 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_cast_int_to_float() { + const START: usize = 0; + const END: usize = 100; + + let device = Default::default(); + let tensor = TestTensorInt::arange(START as i64..END as i64, &device); + + let data_int = tensor.to_data(); + let data_int = data_int.as_slice::().unwrap(); + let data_float = tensor.float().into_data(); + let data_float = data_float.as_slice::().unwrap(); + + for i in START..END { + assert_eq!(data_int[i], i as i32); + assert_eq!(data_float[i], i as f32); + } +} + +#[test] +fn should_cast_bool_to_int() { + let device = Default::default(); + + let tensor_1 = TestTensor::<2>::from_data([[1., 0., 3.], [0., 0., 900.]], &device); + let tensor_2: TestTensorInt<2> = tensor_1.clone().greater_elem(0.0).int(); + + tensor_2 + .to_data() + .assert_eq(&TensorData::from([[1, 0, 1], [0, 0, 1]]), false); +} + +#[test] +fn should_cast_bool_to_float() { + let device = Default::default(); + + let tensor_1 = TestTensor::<2>::from_data([[1., 0., 3.], [0., 0., 900.]], &device); + let tensor_2: TestTensor<2> = tensor_1.clone().greater_elem(0.0).float(); + + tensor_2 + .to_data() + .assert_eq(&TensorData::from([[1., 0., 1.], [0., 0., 1.]]), false); +} diff --git a/crates/burn-backend-tests/tests/cubecl/cat.rs b/crates/burn-backend-tests/tests/cubecl/cat.rs new file mode 100644 index 0000000..03897af --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/cat.rs @@ -0,0 +1,42 @@ +use super::*; +use burn_tensor::Device; +use burn_tensor::Distribution; +use burn_tensor::Tolerance; + +#[test] +fn cat_should_match_reference_backend_dim0() { + test_same_as_reference([6, 256], 2, 0); +} + +#[test] +fn cat_should_match_reference_backend_dim1() { + test_same_as_reference([6, 256], 2, 1); +} + +#[test] +fn cat_should_support_uneven_launch() { + test_same_as_reference([1, 137], 2, 0); +} + +fn test_same_as_reference(shape: [usize; 2], num_tensors: usize, dim: usize) { + let device = Device::default(); + let ref_device = ReferenceDevice::new(); + + device.seed(0); + ref_device.seed(0); + + let tensors = (0..num_tensors) + .map(|_| TestTensor::<2>::random(shape, Distribution::Default, &device)) + .collect::>(); + let tensors_ref = tensors + .iter() + .map(|tensor| TestTensor::<2>::from_data(tensor.to_data(), &ref_device)) + .collect::>(); + + let tensor = TestTensor::<2>::cat(tensors, dim); + let tensor_ref = TestTensor::<2>::cat(tensors_ref, dim); + + tensor + .into_data() + .assert_approx_eq::(&tensor_ref.into_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/cubecl/clamp.rs b/crates/burn-backend-tests/tests/cubecl/clamp.rs new file mode 100644 index 0000000..9033921 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/clamp.rs @@ -0,0 +1,19 @@ +use super::*; +use burn_tensor::Distribution; +use burn_tensor::Tolerance; + +#[test] +fn clamp_should_match_reference() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let input = TestTensor::<4>::random([1, 5, 32, 32], Distribution::Default, &device); + let input_ref = TestTensor::<4>::from_data(input.to_data(), &ref_device); + + let output = input.clamp(0.3, 0.7); + + output.into_data().assert_approx_eq::( + &input_ref.clamp(0.3, 0.7).into_data(), + Tolerance::default(), + ); +} diff --git a/crates/burn-backend-tests/tests/cubecl/contiguous.rs b/crates/burn-backend-tests/tests/cubecl/contiguous.rs new file mode 100644 index 0000000..5d4cd33 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/contiguous.rs @@ -0,0 +1,40 @@ +use super::*; +use burn_tensor::Tolerance; + +#[test] +pub fn into_contiguous_match_reference_backend_1() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + for shape in [ + [4, 4, 4, 4], + [32, 42, 24, 48], + [8, 3, 7, 4], + [1, 4, 1, 1], + [1, 32, 256, 128], + ] { + let num_elems = shape.iter().product::() as i64; + let tensor: TestTensor<4> = TestTensorInt::arange(0..num_elems, &device) + .reshape(shape) + .float(); + let tensor_ref = TestTensor::<4>::from_data(tensor.to_data(), &ref_device); + + for (i, j) in get_combinations(shape.len()) { + let view = tensor.clone().swap_dims(i, j); + let view_ref = tensor_ref.clone().swap_dims(i, j); + let data = view.into_data(); + let data_ref = view_ref.into_data(); + + data_ref.assert_approx_eq::(&data, Tolerance::default()); + } + } +} + +fn get_combinations(n: usize) -> impl Iterator { + // Iterate from 0 up to n + (0..n).flat_map(move |i| { + // For each i, iterate from i + 1 up to n + // This ensures no repeats (i == j) and no duplicates (j, i) + (i + 1..n).map(move |j| (i, j)) + }) +} diff --git a/crates/burn-backend-tests/tests/cubecl/conv2d.rs b/crates/burn-backend-tests/tests/cubecl/conv2d.rs new file mode 100644 index 0000000..71664af --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/conv2d.rs @@ -0,0 +1,116 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::ops::ConvOptions; +use burn_tensor::{Distribution, module}; + +#[test] +fn conv2d_should_match_reference_backend() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let input = TestTensor::<4>::random([6, 16, 32, 32], Distribution::Default, &device); + let weight = TestTensor::<4>::random([12, 8, 3, 3], Distribution::Default, &device); + let bias = TestTensor::<1>::random([12], Distribution::Default, &device); + + let input_ref = TestTensor::<4>::from_data(input.to_data(), &ref_device); + let weight_ref = TestTensor::<4>::from_data(weight.to_data(), &ref_device); + let bias_ref = TestTensor::<1>::from_data(bias.to_data(), &ref_device); + + let options = ConvOptions::new([2, 3], [2, 3], [2, 3], 2); + + let output = module::conv2d(input, weight, Some(bias), options.clone()); + let output_ref = module::conv2d(input_ref, weight_ref, Some(bias_ref), options); + + output + .into_data() + .assert_approx_eq::(&output_ref.into_data(), Tolerance::default()); +} + +#[test] +fn conv2d_should_match_reference_backend_implicit() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let input = TestTensor::<4>::random([4, 16, 6, 6], Distribution::Default, &device); + let weight = TestTensor::<4>::random([16, 16, 3, 3], Distribution::Default, &device); + let bias = TestTensor::<1>::random([16], Distribution::Default, &device); + + let input_ref = TestTensor::<4>::from_data(input.to_data(), &ref_device); + let weight_ref = TestTensor::<4>::from_data(weight.to_data(), &ref_device); + let bias_ref = TestTensor::<1>::from_data(bias.to_data(), &ref_device); + + let options = ConvOptions::new([1, 1], [2, 2], [1, 1], 1); + + let output = module::conv2d(input, weight, Some(bias), options.clone()); + let output_ref = module::conv2d(input_ref, weight_ref, Some(bias_ref), options); + + let tolerance = Tolerance::default(); + output + .into_data() + .assert_approx_eq::(&output_ref.into_data(), tolerance); +} + +/// Regression test for bias loader in new implicit GEMM +#[test] +fn conv2d_should_match_reference_backend_bias_regression() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let input = TestTensor::<4>::random([1, 1, 1, 1], Distribution::Default, &device); + let weight = TestTensor::<4>::random([32, 1, 3, 3], Distribution::Default, &device); + let bias = TestTensor::<1>::random([32], Distribution::Default, &device); + + let input_ref = TestTensor::<4>::from_data(input.to_data(), &ref_device); + let weight_ref = TestTensor::<4>::from_data(weight.to_data(), &ref_device); + let bias_ref = TestTensor::<1>::from_data(bias.to_data(), &ref_device); + + let options = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + + let output = module::conv2d(input, weight, Some(bias), options.clone()).permute([0, 2, 3, 1]); + let output_ref = + module::conv2d(input_ref, weight_ref, Some(bias_ref), options).permute([0, 2, 3, 1]); + + let tolerance = Tolerance::default(); + output + .into_data() + .assert_approx_eq::(&output_ref.into_data(), tolerance); +} + +#[test] +fn conv2d_weight_backward_should_run() { + // https://github.com/tracel-ai/burn/issues/4226#issuecomment-3911335769 + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let options = ConvOptions::new([1, 1], [0, 0], [1, 1], 1); + let x = TestTensor::<4>::random([1, 1, 1, 672], Distribution::Default, &device); + // let x = x.permute([0, 3, 1, 2]); + + let output_grad = TestTensor::<4>::random([1, 168, 1, 1], Distribution::Default, &device); + let weight = TestTensor::<4>::random([168, 672, 1, 1], Distribution::Default, &device); + + let x_ref = TestTensor::<4>::from_data(x.to_data(), &ref_device); + let output_grad_ref = TestTensor::<4>::from_data(output_grad.to_data(), &ref_device); + let weight_ref = TestTensor::<4>::from_data(weight.to_data(), &ref_device); + + // Input shape [672, 1] and strides [672, 672] should be valid + let output = module::conv2d_weight_backward( + x.permute([0, 3, 1, 2]), + weight, + output_grad, + options.clone(), + ); + + // Input shape [672, 1] and strides [672, 672] should be valid + let output_ref = module::conv2d_weight_backward( + x_ref.permute([0, 3, 1, 2]), + weight_ref, + output_grad_ref, + options, + ); + + let tolerance = Tolerance::default(); + output + .into_data() + .assert_approx_eq::(&output_ref.into_data(), tolerance); +} diff --git a/crates/burn-backend-tests/tests/cubecl/conv3d.rs b/crates/burn-backend-tests/tests/cubecl/conv3d.rs new file mode 100644 index 0000000..e432fc4 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/conv3d.rs @@ -0,0 +1,26 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{Distribution, module}; + +#[test] +fn conv3d_should_match_reference_backend() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let input = TestTensor::<5>::random([6, 16, 32, 32, 32], Distribution::Default, &device); + let weight = TestTensor::<5>::random([12, 8, 3, 3, 3], Distribution::Default, &device); + let bias = TestTensor::<1>::random([12], Distribution::Default, &device); + + let input_ref = TestTensor::<5>::from_data(input.to_data(), &ref_device); + let weight_ref = TestTensor::<5>::from_data(weight.to_data(), &ref_device); + let bias_ref = TestTensor::<1>::from_data(bias.to_data(), &ref_device); + + let options = burn_tensor::ops::ConvOptions::new([2, 3, 4], [2, 3, 4], [2, 3, 4], 2); + + let output = module::conv3d(input, weight, Some(bias), options.clone()); + let output_ref = module::conv3d(input_ref, weight_ref, Some(bias_ref), options); + + output + .into_data() + .assert_approx_eq::(&output_ref.into_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/cubecl/conv_transpose2d.rs b/crates/burn-backend-tests/tests/cubecl/conv_transpose2d.rs new file mode 100644 index 0000000..c8bdee4 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/conv_transpose2d.rs @@ -0,0 +1,48 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{Device, Distribution, module}; + +#[test] +fn conv_transpose2d_should_match_reference_backend() { + let device = Device::default(); + let ref_device = ReferenceDevice::new(); + + device.seed(0); + + let height = 8; + let width = 8; + let in_channels = 8; + let out_channels = 8; + let batch_size = 32; + let kernel_size_0 = 3; + let kernel_size_1 = 3; + let options = burn_tensor::ops::ConvTransposeOptions::new([1, 1], [1, 1], [0, 0], [1, 1], 1); + + let input = Tensor::<4>::random( + [batch_size, in_channels, height, width], + Distribution::Default, + &device, + ); + let weight = Tensor::<4>::random( + [ + in_channels, + out_channels / options.groups, + kernel_size_0, + kernel_size_1, + ], + Distribution::Default, + &device, + ); + let bias = TestTensor::<1>::random([out_channels], Distribution::Default, &device); + + let input_ref = TestTensor::<4>::from_data(input.to_data(), &ref_device); + let weight_ref = TestTensor::<4>::from_data(weight.to_data(), &ref_device); + let bias_ref = TestTensor::<1>::from_data(bias.to_data(), &ref_device); + + let output = module::conv_transpose2d(input, weight, Some(bias), options.clone()); + let output_ref = module::conv_transpose2d(input_ref, weight_ref, Some(bias_ref), options); + + output + .into_data() + .assert_approx_eq::(&output_ref.into_data(), Tolerance::rel_abs(0.01, 0.02)); +} diff --git a/crates/burn-backend-tests/tests/cubecl/conv_transpose3d.rs b/crates/burn-backend-tests/tests/cubecl/conv_transpose3d.rs new file mode 100644 index 0000000..d1abe57 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/conv_transpose3d.rs @@ -0,0 +1,52 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{Device, Distribution, module}; + +#[test] +fn conv_transpose3d_should_match_reference_backend() { + let device = Device::default(); + let ref_device = ReferenceDevice::new(); + + device.seed(0); + + let depth = 8; + let height = 8; + let width = 8; + let in_channels = 8; + let out_channels = 8; + let batch_size = 32; + let kernel_size_0 = 3; + let kernel_size_1 = 3; + let kernel_size_2 = 3; + let options = + burn_tensor::ops::ConvTransposeOptions::new([1, 1, 1], [1, 1, 1], [0, 0, 0], [1, 1, 1], 1); + + let input = TestTensor::<5>::random( + [batch_size, in_channels, depth, height, width], + Distribution::Default, + &device, + ); + let weight = TestTensor::<5>::random( + [ + in_channels, + out_channels / options.groups, + kernel_size_0, + kernel_size_1, + kernel_size_2, + ], + Distribution::Default, + &device, + ); + let bias = TestTensor::<1>::random([out_channels], Distribution::Default, &device); + + let input_ref = TestTensor::<5>::from_data(input.to_data(), &ref_device); + let weight_ref = TestTensor::<5>::from_data(weight.to_data(), &ref_device); + let bias_ref = TestTensor::<1>::from_data(bias.to_data(), &ref_device); + + let output = module::conv_transpose3d(input, weight, Some(bias), options.clone()); + let output_ref = module::conv_transpose3d(input_ref, weight_ref, Some(bias_ref), options); + + output + .into_data() + .assert_approx_eq::(&output_ref.into_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/cubecl/cross.rs b/crates/burn-backend-tests/tests/cubecl/cross.rs new file mode 100644 index 0000000..93e4da5 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/cross.rs @@ -0,0 +1,156 @@ +use super::*; +use burn_tensor::Tolerance; + +#[test] +fn test_cross_product() { + let device = Default::default(); + // Test with well-known orthogonal vectors for clearer validation + let a = TestTensor::<2>::from_data([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], &device); + let b = TestTensor::<2>::from_data([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], &device); + + let result = a.cross(b, 1); + // For orthogonal unit vectors: + // i × j = k + // j × k = i + let expected = TestTensor::<2>::from_data([[0.0, 0.0, 1.0], [1.0, 0.0, 0.0]], &device); + + // Use Tolerance for floating-point comparisons + let tolerance = Tolerance::::default(); + result + .to_data() + .assert_approx_eq(&expected.to_data(), tolerance); +} + +#[test] +fn test_cross_product_zeros() { + let device = Default::default(); + // Test cross product with zero vector - should always give zero vector + let a = TestTensor::<2>::from_data([[2.0, 3.0, 4.0]], &device); + let b = TestTensor::<2>::zeros([1, 3], &device); + + let result = a.cross(b, 1); + let expected = TestTensor::<2>::zeros([1, 3], &device); + + // For zeros, we can use exact equality or a very tight tolerance + let tolerance = Tolerance::::default(); + result + .to_data() + .assert_approx_eq(&expected.to_data(), tolerance); +} + +#[test] +fn test_cross_product_batch() { + let device = Default::default(); + // Test typical cross product computations in batch + let a = TestTensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + let b = TestTensor::<2>::from_data([[4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], &device); + + let result = a.cross(b, 1); + // Cross products: + // [1,2,3] × [4,5,6] = [-3,6,-3] + // [4,5,6] × [7,8,9] = [-3,6,-3] + let expected = TestTensor::<2>::from_data([[-3.0, 6.0, -3.0], [-3.0, 6.0, -3.0]], &device); + + let tolerance = Tolerance::::default(); + result + .to_data() + .assert_approx_eq(&expected.to_data(), tolerance); +} + +#[test] +#[should_panic] +fn test_cross_product_invalid_dimension() { + let device = Default::default(); + let a = TestTensor::<2>::zeros([1, 4], &device); + let b = TestTensor::<2>::zeros([1, 4], &device); + + let _ = a.cross(b, 1); +} + +#[test] +fn test_cross_product_parallel_vectors() { + let device = Default::default(); + // Test cross product of parallel vectors (should be zero) + let a = TestTensor::<2>::from_data([[1.0, 2.0, 3.0]], &device); + let b = TestTensor::<2>::from_data([[2.0, 4.0, 6.0]], &device); // b = 2 * a + + let result = a.cross(b, 1); + let expected = TestTensor::<2>::zeros([1, 3], &device); + + let tolerance = Tolerance::::default(); + result + .to_data() + .assert_approx_eq(&expected.to_data(), tolerance); +} + +#[test] +fn test_cross_product_3d_tensor() { + let device = Default::default(); + // Test with 3D tensor (batch of matrices) + let a = TestTensor::<3>::from_data( + [ + [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + ], + &device, + ); + + let b = TestTensor::<3>::from_data( + [ + [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], + [[4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], + ], + &device, + ); + + let result = a.cross(b, 2); // Cross on last dimension + let expected = TestTensor::<3>::from_data( + [ + [[0.0, 0.0, 1.0], [1.0, 0.0, 0.0]], + [[-3.0, 6.0, -3.0], [-3.0, 6.0, -3.0]], + ], + &device, + ); + + let tolerance = Tolerance::::default(); + result + .to_data() + .assert_approx_eq(&expected.to_data(), tolerance); +} + +// Test to verify that padding doesn't affect results +#[test] +fn test_cross_product_with_padding_awareness() { + let device = Default::default(); + // Create tensors that would span multiple 4-element blocks + // This tests that the padding doesn't corrupt adjacent data + let a = TestTensor::<2>::from_data( + [ + [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], // Two vectors: [1,2,3] and [4,5,6] + ], + &device, + ); + + let b = TestTensor::<2>::from_data( + [ + [7.0, 8.0, 9.0, 10.0, 11.0, 12.0], // Two vectors: [7,8,9] and [10,11,12] + ], + &device, + ); + + // Reshape to have proper 3-element vectors in last dimension + let a_reshaped = a.reshape([2, 3]); + let b_reshaped = b.reshape([2, 3]); + + let result = a_reshaped.cross(b_reshaped, 1); + + // Expected cross products: + // [1,2,3] × [7,8,9] = [-6,12,-6] + // [4,5,6] × [10,11,12] = [-6,12,-6] + let expected = TestTensor::<2>::from_data([[-6.0, 12.0, -6.0], [-6.0, 12.0, -6.0]], &device); + + let tolerance = Tolerance::::default(); + result + .to_data() + .assert_approx_eq(&expected.to_data(), tolerance); +} diff --git a/crates/burn-backend-tests/tests/cubecl/fft.rs b/crates/burn-backend-tests/tests/cubecl/fft.rs new file mode 100644 index 0000000..f509f1c --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/fft.rs @@ -0,0 +1,712 @@ +use super::*; +use burn_tensor::signal::{cfft, irfft, rfft}; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn rfft_zeros() { + let signal = TestTensor::<1>::from([0.0, 0.0, 0.0, 0.0]); + let (re, im) = rfft(signal, 0, None); + + let expected_re = TensorData::from([0.0, 0.0, 0.0]); + let expected_im = TensorData::from([0.0, 0.0, 0.0]); + + re.into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-4)); + im.into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-4)); +} + +#[test] +fn rfft_constant() { + let signal = TestTensor::<1>::from([1.0, 1.0, 1.0, 1.0]); + let (re, im) = rfft(signal, 0, None); + + let expected_re = TensorData::from([4.0, 0.0, 0.0]); + let expected_im = TensorData::from([0.0, 0.0, 0.0]); + + re.into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-4)); + im.into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-4)); +} + +#[test] +#[should_panic] // "RFFT requires n_fft >= 2" error is shadowed by the CallError +fn rfft_length1() { + let signal = TestTensor::<1>::from([5.0]); + let (re, im) = rfft(signal, 0, None); + + let expected_re = TensorData::from([5.0]); + let expected_im = TensorData::from([0.0]); + + re.into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-4)); + im.into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-4)); +} + +#[test] +fn rfft_length2() { + let signal = TestTensor::<1>::from([1.0, -1.0]); + let (re, im) = rfft(signal, 0, None); + + let expected_re = TensorData::from([0.0, 2.0]); + let expected_im = TensorData::from([0.0, 0.0]); + + re.into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-4)); + im.into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-4)); +} + +#[test] +fn rfft_dim1_sine_wave_produces_imaginary_spectrum() { + let signal = TestTensor::<2>::from([[0.0, 1.2071, 1.0, 0.2071, 0.0, -0.2071, -1.0, -1.2071]]); + let dim = 1; + let (spectrum_re, spectrum_im) = rfft(signal.clone(), dim, None); + let expected_re = TensorData::from([[0, 0, 0, 0, 0]]); + let expected_im = TensorData::from([[0, -4, -2, 0, 0]]); + + assert_eq!(spectrum_re.shape(), spectrum_im.shape()); + assert_eq!(spectrum_re.shape(), expected_re.shape); + + spectrum_re + .into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-4)); + + spectrum_im + .into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-4)); +} + +#[test] +fn rfft_dim1_cosine_wave_produces_real_spectrum() { + let signal = TestTensor::<2>::from([[1.0, 0.7071, 0.0, -0.7071, -1.0, -0.7071, 0.0, 0.7071]]); + + let (spectrum_re, spectrum_im) = rfft(signal, 1, None); + + let expected_re = TensorData::from([[0.0, 4.0, 0.0, 0.0, 0.0]]); + let expected_im = TensorData::from([[0.0, 0.0, 0.0, 0.0, 0.0]]); + + assert_eq!(spectrum_re.shape(), spectrum_im.shape()); + assert_eq!(spectrum_re.shape(), expected_re.shape); + + spectrum_re + .into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-3)); + + spectrum_im + .into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-3)); +} + +#[test] +fn rfft_dim1_2d_tensor_distinct_rows() { + let signal = TestTensor::<2>::from([ + [0.0, 0.7071, 1.0, 0.7071, 0.0, -0.7071, -1.0, -0.7071], + [0.0, 1.0, 0.0, -1.0, 0.0, 1.0, 0.0, -1.0], + ]); + + let (spectrum_re, spectrum_im) = rfft(signal, 1, None); + + let expected_re = TensorData::from([[0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]]); + + let expected_im = TensorData::from([[0.0, -4.0, 0.0, 0.0, 0.0], [0.0, 0.0, -4.0, 0.0, 0.0]]); + + assert_eq!(spectrum_re.shape(), spectrum_im.shape()); + assert_eq!(spectrum_re.shape(), expected_re.shape); + + spectrum_re + .into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-3)); + + spectrum_im + .into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-3)); +} + +#[test] +fn rfft_dim0_2d_tensor() { + let signal = TestTensor::<2>::from([ + [0.0, 0.0], + [0.7071, 1.0], + [1.0, 0.0], + [0.7071, -1.0], + [0.0, 0.0], + [-0.7071, 1.0], + [-1.0, 0.0], + [-0.7071, -1.0], + ]); + + let (spectrum_re, spectrum_im) = rfft(signal, 0, None); + + let expected_re = + TensorData::from([[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]); + + let expected_im = + TensorData::from([[0.0, 0.0], [-4.0, 0.0], [0.0, -4.0], [0.0, 0.0], [0.0, 0.0]]); + + assert_eq!(spectrum_re.shape(), spectrum_im.shape()); + assert_eq!(spectrum_re.shape(), expected_re.shape); + + spectrum_re + .into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-3)); + + spectrum_im + .into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-3)); +} + +#[test] +fn rfft_dim2_3d_tensor() { + let signal = TestTensor::<3>::from([ + [ + [0.0, 0.7071, 1.0, 0.7071, 0.0, -0.7071, -1.0, -0.7071], + [0.0, 1.0, 0.0, -1.0, 0.0, 1.0, 0.0, -1.0], + ], + [ + [0.0, 1.0, 0.0, -1.0, 0.0, 1.0, 0.0, -1.0], + [0.0, 0.7071, 1.0, 0.7071, 0.0, -0.7071, -1.0, -0.7071], + ], + ]); + + let (spectrum_re, spectrum_im) = rfft(signal, 2, None); + + let expected_re = TensorData::from([ + [[0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]], + [[0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]], + ]); + + let expected_im = TensorData::from([ + [[0.0, -4.0, 0.0, 0.0, 0.0], [0.0, 0.0, -4.0, 0.0, 0.0]], + [[0.0, 0.0, -4.0, 0.0, 0.0], [0.0, -4.0, 0.0, 0.0, 0.0]], + ]); + + assert_eq!(spectrum_re.shape(), spectrum_im.shape()); + assert_eq!(spectrum_re.shape(), expected_re.shape); + + spectrum_re + .into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-3)); + + spectrum_im + .into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-3)); +} + +#[test] +fn rfft_dim1_3d_tensor() { + let signal = TestTensor::<3>::from([ + [[1.0, 0.0], [0.0, 0.0], [-1.0, 0.0], [0.0, 0.0]], + [[0.0, 1.0], [0.0, 0.0], [0.0, -1.0], [0.0, 0.0]], + ]); + + let (spectrum_re, spectrum_im) = rfft(signal, 1, None); + + let expected_re = TensorData::from([ + [[0.0, 0.0], [2.0, 0.0], [0.0, 0.0]], + [[0.0, 0.0], [0.0, 2.0], [0.0, 0.0]], + ]); + + let expected_im = TensorData::from([ + [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], + [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], + ]); + + assert_eq!(spectrum_re.shape(), spectrum_im.shape()); + assert_eq!(spectrum_re.shape(), expected_re.shape); + + spectrum_re + .into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-3)); + + spectrum_im + .into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-3)); +} + +#[test] +fn irfft_dim1_imaginary_spectrum_produces_sine_wave() { + let spectrum_re = TestTensor::<2>::from([[0.0, 0.0, 0.0, 0.0, 0.0]]); + let spectrum_im = TestTensor::<2>::from([[0.0, -4.0, -2.0, 0.0, 0.0]]); + + let signal = irfft(spectrum_re, spectrum_im, 1, None); + + let expected = TensorData::from([[0.0, 1.2071, 1.0, 0.2071, 0.0, -0.2071, -1.0, -1.2071]]); + + signal + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-3)); +} + +#[test] +fn irfft_dim1_real_spectrum_produces_cosine_wave() { + let spectrum_re = TestTensor::<2>::from([[0.0, 4.0, 0.0, 0.0, 0.0]]); + let spectrum_im = TestTensor::<2>::from([[0.0, 0.0, 0.0, 0.0, 0.0]]); + + let signal = irfft(spectrum_re, spectrum_im, 1, None); + + let expected = TensorData::from([[1.0, 0.7071, 0.0, -0.7071, -1.0, -0.7071, 0.0, 0.7071]]); + + signal + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-3)); +} + +#[test] +fn irfft_dim1_2d_tensor_distinct_rows() { + let spectrum_re = TestTensor::<2>::from([[0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]]); + + let spectrum_im = + TestTensor::<2>::from([[0.0, -4.0, 0.0, 0.0, 0.0], [0.0, 0.0, -4.0, 0.0, 0.0]]); + + let signal = irfft(spectrum_re, spectrum_im, 1, None); + + let expected = TensorData::from([ + [0.0, 0.7071, 1.0, 0.7071, 0.0, -0.7071, -1.0, -0.7071], + [0.0, 1.0, 0.0, -1.0, 0.0, 1.0, 0.0, -1.0], + ]); + + signal + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-3)); +} + +#[test] +fn irfft_dim0_2d_tensor() { + let spectrum_re = + TestTensor::<2>::from([[0.0, 0.0], [4.0, 4.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]); + + let spectrum_im = + TestTensor::<2>::from([[0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]); + + let signal = irfft(spectrum_re, spectrum_im, 0, None); + + let expected = TensorData::from([ + [1.0, 1.0], + [0.7071, 0.7071], + [0.0, 0.0], + [-0.7071, -0.7071], + [-1.0, -1.0], + [-0.7071, -0.7071], + [0.0, 0.0], + [0.7071, 0.7071], + ]); + + signal + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-3)); +} + +#[test] +fn rfft_irfft_roundtrip_1d() { + let signal = TestTensor::<1>::from([0.0, 1.2071, 1.0, 0.2071, 0.0, -0.2071, -1.0, -1.2071]); + + let (re, im) = rfft(signal.clone(), 0, None); + let reconstructed = irfft(re, im, 0, None); + + reconstructed + .into_data() + .assert_approx_eq::(&signal.into_data(), Tolerance::absolute(1e-3)); +} + +#[test] +fn rfft_irfft_roundtrip_dim1_2d() { + let signal = TestTensor::<2>::from([ + [0.0, 0.7071, 1.0, 0.7071, 0.0, -0.7071, -1.0, -0.7071], + [0.0, 1.0, 0.0, -1.0, 0.0, 1.0, 0.0, -1.0], + ]); + + let (re, im) = rfft(signal.clone(), 1, None); + let reconstructed = irfft(re, im, 1, None); + + reconstructed + .into_data() + .assert_approx_eq::(&signal.into_data(), Tolerance::absolute(1e-3)); +} + +#[test] +fn rfft_irfft_roundtrip_dim0_2d() { + let signal = TestTensor::<2>::from([ + [1.0, 0.0], + [0.7071, 1.0], + [0.0, 0.0], + [-0.7071, -1.0], + [-1.0, 0.0], + [-0.7071, 1.0], + [0.0, 0.0], + [0.7071, -1.0], + ]); + + let (re, im) = rfft(signal.clone(), 0, None); + let reconstructed = irfft(re, im, 0, None); + + reconstructed + .into_data() + .assert_approx_eq::(&signal.into_data(), Tolerance::absolute(1e-3)); +} + +#[test] +fn rfft_irfft_roundtrip_dim1_3d() { + let signal = TestTensor::<3>::from([ + [ + [0.0, 1.0], + [1.0, 0.0], + [0.0, -1.0], + [-1.0, 0.0], + [0.0, 1.0], + [1.0, 0.0], + [0.0, -1.0], + [-1.0, 0.0], + ], + [ + [1.0, 0.0], + [0.7071, 1.0], + [0.0, 0.0], + [-0.7071, -1.0], + [-1.0, 0.0], + [-0.7071, 1.0], + [0.0, 0.0], + [0.7071, -1.0], + ], + ]); + + let (re, im) = rfft(signal.clone(), 1, None); + let reconstructed = irfft(re, im, 1, None); + + reconstructed + .into_data() + .assert_approx_eq::(&signal.into_data(), Tolerance::absolute(1e-3)); +} + +// ---- Padded input tests (n: Some(...)) ---- + +#[test] +fn rfft_with_n_larger_than_signal() { + // Signal of length 4, padded to n=8 + // DFT of [1,0,0,0, 0,0,0,0] = all-ones real, zero imag + let signal = TestTensor::<1>::from([1.0, 0.0, 0.0, 0.0]); + let (re, im) = rfft(signal, 0, Some(8)); + + let expected_re = TensorData::from([1.0, 1.0, 1.0, 1.0, 1.0]); + let expected_im = TensorData::from([0.0, 0.0, 0.0, 0.0, 0.0]); + + re.into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-3)); + im.into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-3)); +} + +#[test] +fn rfft_with_n_smaller_than_signal() { + // Signal of length 8, truncated to n=4 -> DFT of [1,0,0,0] + let signal = TestTensor::<1>::from([1.0, 0.0, 0.0, 0.0, 99.0, 99.0, 99.0, 99.0]); + let (re, im) = rfft(signal, 0, Some(4)); + + let expected_re = TensorData::from([1.0, 1.0, 1.0]); + let expected_im = TensorData::from([0.0, 0.0, 0.0]); + + re.into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-3)); + im.into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-3)); +} + +#[test] +#[should_panic(expected = "power of two")] +fn rfft_rejects_non_power_of_two_n() { + let signal = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0]); + let _ = rfft(signal, 0, Some(5)); +} + +#[test] +#[should_panic(expected = "power of two")] +fn irfft_rejects_non_power_of_two_n() { + let re = TestTensor::<1>::from([1.0, 2.0, 3.0]); + let im = TestTensor::<1>::from([0.0, 0.0, 0.0]); + let _ = irfft(re, im, 0, Some(5)); +} + +#[test] +fn rfft_irfft_roundtrip_with_n() { + let signal = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0]); + let (re, im) = rfft(signal.clone(), 0, Some(4)); + let reconstructed = irfft(re, im, 0, Some(4)); + + reconstructed + .into_data() + .assert_approx_eq::(&signal.into_data(), Tolerance::absolute(1e-3)); +} + +#[test] +fn irfft_with_n_different_from_natural() { + // Spectrum from length-4 signal (3 bins), reconstruct at length 8 + let signal = TestTensor::<1>::from([1.0, 0.0, 0.0, 0.0]); + let (re, im) = rfft(signal, 0, None); + let reconstructed = irfft(re, im, 0, Some(8)); + assert_eq!(reconstructed.dims(), [8]); +} + +#[test] +fn rfft_2d_with_n_padded() { + // 2D tensor, rfft along dim=1 with n=8 (signal is length 4) + let signal = TestTensor::<2>::from([[1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0]]); + let (re, im) = rfft(signal, 1, Some(8)); + + // Output: 8/2+1=5 frequency bins + assert_eq!(re.dims(), [2, 5]); + assert_eq!(im.dims(), [2, 5]); + + // Row 0: impulse zero-padded to 8 -> all-ones real + let re_data = re.into_data(); + let re_vals = re_data.to_vec::().unwrap(); + for k in 0..5 { + assert!( + (re_vals[k] - 1.0).abs() < 1e-3, + "row0 re[{k}] should be 1.0, got {}", + re_vals[k] + ); + } +} + +// ---- cfft tests ---- + +#[test] +fn cfft_output_has_n_bins() { + // cfft should return N bins, not N/2+1 + let re = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0]); + let im = TestTensor::<1>::from([0.0, 0.0, 0.0, 0.0]); + let (out_re, out_im) = cfft(re, im, 0, None); + + assert_eq!(out_re.dims(), [4]); + assert_eq!(out_im.dims(), [4]); +} + +#[test] +fn cfft_pure_real_input() { + // When imaginary part is zero, cfft should produce the same result + // as extending rfft to the full spectrum + let signal = [1.0f32, 2.0, 3.0, 4.0]; + let re = TestTensor::<1>::from(signal); + let im = TestTensor::<1>::from([0.0, 0.0, 0.0, 0.0]); + + let (cfft_re, cfft_im) = cfft(re, im, 0, None); + + // Expected: DFT of [1,2,3,4] + // X[0] = 10, X[1] = -2+2i, X[2] = -2, X[3] = -2-2i + let expected_re = TensorData::from([10.0, -2.0, -2.0, -2.0]); + let expected_im = TensorData::from([0.0, 2.0, 0.0, -2.0]); + + cfft_re + .into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-3)); + cfft_im + .into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-3)); +} + +#[test] +fn cfft_pure_imaginary_input() { + // Signal is purely imaginary: z[n] = i * [1, 2, 3, 4] + // FFT(i*x) = i*FFT(x), so result_re = -FFT(x)_im, result_im = FFT(x)_re + let re = TestTensor::<1>::from([0.0, 0.0, 0.0, 0.0]); + let im = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0]); + + let (cfft_re, cfft_im) = cfft(re, im, 0, None); + + // FFT([1,2,3,4]) = [10, -2+2i, -2, -2-2i] + // i * FFT(x) = i * [10, -2+2i, -2, -2-2i] + // = [-0, -2+(-2)i, 0, 2+(-2)i] → re = [0, -2, 0, 2], im = [10, -2, -2, -2] + let expected_re = TensorData::from([0.0, -2.0, 0.0, 2.0]); + let expected_im = TensorData::from([10.0, -2.0, -2.0, -2.0]); + + cfft_re + .into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-3)); + cfft_im + .into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-3)); +} + +#[test] +fn cfft_complex_exponential() { + // z[n] = exp(i * 2π * n / 4) for n=0..3, i.e. frequency bin 1 + // re = [cos(0), cos(π/2), cos(π), cos(3π/2)] = [1, 0, -1, 0] + // im = [sin(0), sin(π/2), sin(π), sin(3π/2)] = [0, 1, 0, -1] + // DFT should be: X[0]=0, X[1]=4, X[2]=0, X[3]=0 + let re = TestTensor::<1>::from([1.0, 0.0, -1.0, 0.0]); + let im = TestTensor::<1>::from([0.0, 1.0, 0.0, -1.0]); + + let (cfft_re, cfft_im) = cfft(re, im, 0, None); + + let expected_re = TensorData::from([0.0, 4.0, 0.0, 0.0]); + let expected_im = TensorData::from([0.0, 0.0, 0.0, 0.0]); + + cfft_re + .into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-3)); + cfft_im + .into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-3)); +} + +#[test] +fn cfft_zeros() { + let re = TestTensor::<1>::from([0.0, 0.0, 0.0, 0.0]); + let im = TestTensor::<1>::from([0.0, 0.0, 0.0, 0.0]); + + let (cfft_re, cfft_im) = cfft(re, im, 0, None); + + let expected = TensorData::from([0.0, 0.0, 0.0, 0.0]); + + cfft_re + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-4)); + cfft_im + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-4)); +} + +#[test] +fn cfft_dim1_2d_tensor() { + // Apply cfft along dim=1 on a 2D tensor + // Row 0: pure real [1, 2, 3, 4] → DFT = [10, -2+2i, -2, -2-2i] + // Row 1: complex exponential exp(i·2π·n/4) → DFT = [0, 4, 0, 0] + let re = TestTensor::<2>::from([[1.0, 2.0, 3.0, 4.0], [1.0, 0.0, -1.0, 0.0]]); + let im = TestTensor::<2>::from([[0.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, -1.0]]); + + let (cfft_re, cfft_im) = cfft(re, im, 1, None); + + // Output should be [2, 4] (N=4 bins per row) + assert_eq!(cfft_re.dims(), [2, 4]); + assert_eq!(cfft_im.dims(), [2, 4]); + + // Row 0: DFT of [1,2,3,4]+i*0 = [10, -2+2i, -2, -2-2i] + // Row 1: DFT of exp(i*2π*n/4) = [0, 4, 0, 0] + let expected_re = TensorData::from([[10.0, -2.0, -2.0, -2.0], [0.0, 4.0, 0.0, 0.0]]); + let expected_im = TensorData::from([[0.0, 2.0, 0.0, -2.0], [0.0, 0.0, 0.0, 0.0]]); + + cfft_re + .into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-3)); + cfft_im + .into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-3)); +} + +#[test] +fn cfft_with_n_padding() { + // Signal length 2, padded to N=4 + // z = [1+0i, 0+0i] padded to [1+0i, 0, 0, 0] + // DFT = [1, 1, 1, 1] (all real, zero imag) + let re = TestTensor::<1>::from([1.0, 0.0]); + let im = TestTensor::<1>::from([0.0, 0.0]); + + let (cfft_re, cfft_im) = cfft(re, im, 0, Some(4)); + + assert_eq!(cfft_re.dims(), [4]); + + let expected_re = TensorData::from([1.0, 1.0, 1.0, 1.0]); + let expected_im = TensorData::from([0.0, 0.0, 0.0, 0.0]); + + cfft_re + .into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-3)); + cfft_im + .into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-3)); +} + +#[test] +#[should_panic] // "RFFT requires n_fft >= 2" error is shadowed by the CallError +fn cfft_length_1() { + // N=1: DFT of a single complex value is itself + let re = TestTensor::<1>::from([3.0]); + let im = TestTensor::<1>::from([5.0]); + + let (cfft_re, cfft_im) = cfft(re, im, 0, None); + + assert_eq!(cfft_re.dims(), [1]); + cfft_re + .into_data() + .assert_approx_eq::(&TensorData::from([3.0]), Tolerance::absolute(1e-4)); + cfft_im + .into_data() + .assert_approx_eq::(&TensorData::from([5.0]), Tolerance::absolute(1e-4)); +} + +#[test] +fn cfft_length_2() { + // N=2: z = [a, b] → X[0] = a+b, X[1] = a-b + // z = [1+2i, 3+4i] + // X[0] = (1+3) + i(2+4) = 4+6i + // X[1] = (1-3) + i(2-4) = -2-2i + let re = TestTensor::<1>::from([1.0, 3.0]); + let im = TestTensor::<1>::from([2.0, 4.0]); + + let (cfft_re, cfft_im) = cfft(re, im, 0, None); + + assert_eq!(cfft_re.dims(), [2]); + cfft_re + .into_data() + .assert_approx_eq::(&TensorData::from([4.0, -2.0]), Tolerance::absolute(1e-4)); + cfft_im + .into_data() + .assert_approx_eq::(&TensorData::from([6.0, -2.0]), Tolerance::absolute(1e-4)); +} + +#[test] +#[should_panic(expected = "same shape")] +fn cfft_rejects_mismatched_shapes() { + let re = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0]); + let im = TestTensor::<1>::from([1.0, 2.0]); + let _ = cfft(re, im, 0, None); +} + +#[test] +fn cfft_dim0_2d_tensor() { + // Apply cfft along dim=0 on a 2D tensor (4 rows, 2 columns) + // Column 0: complex exponential exp(i·2π·n/4) → DFT = [0, 4, 0, 0] + // Column 1: pure real [1, 2, 3, 4] → DFT = [10, -2+2i, -2, -2-2i] + let re = TestTensor::<2>::from([[1.0, 1.0], [0.0, 2.0], [-1.0, 3.0], [0.0, 4.0]]); + let im = TestTensor::<2>::from([[0.0, 0.0], [1.0, 0.0], [0.0, 0.0], [-1.0, 0.0]]); + + let (cfft_re, cfft_im) = cfft(re, im, 0, None); + + assert_eq!(cfft_re.dims(), [4, 2]); + assert_eq!(cfft_im.dims(), [4, 2]); + + let expected_re = TensorData::from([[0.0, 10.0], [4.0, -2.0], [0.0, -2.0], [0.0, -2.0]]); + let expected_im = TensorData::from([[0.0, 0.0], [0.0, 2.0], [0.0, 0.0], [0.0, -2.0]]); + + cfft_re + .into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-3)); + cfft_im + .into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-3)); +} + +#[test] +fn cfft_with_n_truncation() { + // Signal length 8, truncated to n=4 → DFT of [1+0i, 2+0i, 3+0i, 4+0i] + // Trailing values are discarded, not included in the transform. + let re = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0, 99.0, 99.0, 99.0, 99.0]); + let im = TestTensor::<1>::from([0.0, 0.0, 0.0, 0.0, 99.0, 99.0, 99.0, 99.0]); + + let (cfft_re, cfft_im) = cfft(re, im, 0, Some(4)); + + assert_eq!(cfft_re.dims(), [4]); + + // DFT of [1,2,3,4] = [10, -2+2i, -2, -2-2i] + let expected_re = TensorData::from([10.0, -2.0, -2.0, -2.0]); + let expected_im = TensorData::from([0.0, 2.0, 0.0, -2.0]); + + cfft_re + .into_data() + .assert_approx_eq::(&expected_re, Tolerance::absolute(1e-3)); + cfft_im + .into_data() + .assert_approx_eq::(&expected_im, Tolerance::absolute(1e-3)); +} diff --git a/crates/burn-backend-tests/tests/cubecl/gather.rs b/crates/burn-backend-tests/tests/cubecl/gather.rs new file mode 100644 index 0000000..a33fc32 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/gather.rs @@ -0,0 +1,43 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{Device, Distribution, Shape}; + +#[test] +fn gather_should_work_with_multiple_workgroups_dim0() { + test_same_as_ref([6, 256], 0); +} + +#[test] +fn gather_should_work_with_multiple_workgroups_dim1() { + test_same_as_ref([6, 256], 1); +} + +fn test_same_as_ref(shape: [usize; D], dim: usize) { + let device = Device::default(); + let ref_device = ReferenceDevice::new(); + + device.seed(0); + + let max = shape[dim]; + let shape = Shape::new(shape); + let tensor = TestTensor::::random(shape.clone(), Distribution::Default, &device); + let indices = TestTensorInt::<1>::from_data( + TestTensor::<1>::random( + [shape.num_elements()], + Distribution::Uniform(0., max as f64), + &device, + ) + .into_data(), + &device, + ) + .reshape(shape); + let tensor_ref = TestTensor::::from_data(tensor.to_data(), &ref_device); + let indices_ref = TestTensorInt::::from_data(indices.to_data(), &ref_device); + + let actual = tensor.gather(dim, indices); + let expected = tensor_ref.gather(dim, indices_ref); + + expected + .into_data() + .assert_approx_eq::(&actual.into_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/cubecl/interpolate_nearest.rs b/crates/burn-backend-tests/tests/cubecl/interpolate_nearest.rs new file mode 100644 index 0000000..7f4ee00 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/interpolate_nearest.rs @@ -0,0 +1,31 @@ +use super::*; +use burn_tensor::ops::{InterpolateMode, InterpolateOptions}; +use burn_tensor::{Distribution, Tolerance, module}; + +/// Regression test for https://github.com/tracel-ai/burn/issues/4686 +/// +/// The nearest-neighbor interpolation kernel previously used f32 division +/// to map output coordinates to input coordinates. For certain spatial +/// dimensions, GPU f32 division produced values that truncated to the +/// wrong integer, selecting the wrong input pixel. +#[test] +pub fn nearest_interpolate_should_match_reference_backend() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + // These spatial sizes previously triggered the float precision bug. + for h in [214, 220, 227, 235, 255] { + let tensor = TestTensor::<4>::random([1, 64, h, h], Distribution::Default, &device); + let tensor_ref = TestTensor::<4>::from_data(tensor.to_data(), &ref_device); + + let opts = InterpolateOptions::new(InterpolateMode::Nearest); + let out_size = [h * 2, h * 2]; + + let output = module::interpolate(tensor, out_size, opts.clone()); + let output_ref = module::interpolate(tensor_ref, out_size, opts); + + output + .into_data() + .assert_approx_eq::(&output_ref.into_data(), Tolerance::default()); + } +} diff --git a/crates/burn-backend-tests/tests/cubecl/mask_fill.rs b/crates/burn-backend-tests/tests/cubecl/mask_fill.rs new file mode 100644 index 0000000..fbd6d8e --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/mask_fill.rs @@ -0,0 +1,51 @@ +use super::*; +use burn_tensor::Distribution; +use burn_tensor::Tolerance; + +#[test] +fn mask_fill_should_match_reference_backend() { + let (tensor, mask, tensor_ref, mask_ref) = inputs_mask_fill(); + + // MaskFillStrategy::Readonly + let _clone_for_readonly = tensor.clone(); + + let actual = tensor.mask_fill(mask, 4.0); + let expected = tensor_ref.mask_fill(mask_ref, 4.0); + + expected + .into_data() + .assert_approx_eq::(&actual.into_data(), Tolerance::default()); +} + +#[test] +fn mask_fill_inplace_should_match_reference_backend() { + let (tensor, mask, tensor_ref, mask_ref) = inputs_mask_fill(); + + // MaskFillStrategy::Inplace + let actual = tensor.mask_fill(mask, 4.0); + let expected = tensor_ref.mask_fill(mask_ref, 4.0); + + expected + .into_data() + .assert_approx_eq::(&actual.into_data(), Tolerance::default()); +} + +#[allow(clippy::type_complexity)] +fn inputs_mask_fill() -> ( + TestTensor<3>, + TestTensorBool<3>, + TestTensor<3>, + TestTensorBool<3>, +) { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::<3>::random([2, 6, 256], Distribution::Default, &device); + let mask = TestTensor::<3>::random([2, 6, 256], Distribution::Uniform(0., 1.), &device) + .lower_equal_elem(0.5); + + let tensor_ref = TestTensor::<3>::from_data(tensor.to_data(), &ref_device); + let mask_ref = TestTensorBool::<3>::from_data(mask.to_data(), &ref_device); + + (tensor, mask, tensor_ref, mask_ref) +} diff --git a/crates/burn-backend-tests/tests/cubecl/mask_where.rs b/crates/burn-backend-tests/tests/cubecl/mask_where.rs new file mode 100644 index 0000000..36b77e7 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/mask_where.rs @@ -0,0 +1,70 @@ +use super::*; +use burn_tensor::Device; +use burn_tensor::Distribution; +use burn_tensor::Tolerance; + +#[test] +fn mask_where_should_match_reference_backend() { + let (tensor, value, mask, tensor_ref, value_ref, mask_ref) = inputs_mask_where(); + + let actual = tensor.mask_where(mask, value); + let expected = tensor_ref.mask_where(mask_ref, value_ref); + + expected + .into_data() + .assert_approx_eq::(&actual.into_data(), Tolerance::default()); +} +#[test] +fn mask_where_inplace_lhs_should_match_reference_backend() { + let (tensor, value, mask, tensor_ref, value_ref, mask_ref) = inputs_mask_where(); + + // MaskWhereStrategy::InplaceLhs + let actual = tensor.mask_where(mask, value); + let expected = tensor_ref.mask_where(mask_ref, value_ref); + + expected + .into_data() + .assert_approx_eq::(&actual.into_data(), Tolerance::default()); +} + +#[test] +fn mask_where_inplace_rhs_should_match_reference_backend() { + let (tensor, value, mask, tensor_ref, value_ref, mask_ref) = inputs_mask_where(); + + // MaskWhereStrategy::InplaceRhs + let _clone_for_inplace_rhs = tensor.clone(); + + let actual = tensor.mask_where(mask, value); + let expected = tensor_ref.mask_where(mask_ref, value_ref); + + expected + .into_data() + .assert_approx_eq::(&actual.into_data(), Tolerance::default()); +} + +#[allow(clippy::type_complexity)] +fn inputs_mask_where() -> ( + TestTensor<3>, + TestTensor<3>, + TestTensorBool<3>, + TestTensor<3>, + TestTensor<3>, + TestTensorBool<3>, +) { + let device = Device::default(); + let ref_device = ReferenceDevice::new(); + + device.seed(0); + + let tensor = TestTensor::<3>::random([2, 6, 256], Distribution::Default, &device); + let value = TestTensor::<3>::random([2, 6, 256], Distribution::Default, &device); + let mask = TestTensor::<3>::random([2, 6, 256], Distribution::Uniform(0., 1.), &device) + .lower_equal_elem(0.5); + + let tensor_ref = TestTensor::<3>::from_data(tensor.to_data(), &ref_device); + let value_ref = TestTensor::<3>::from_data(value.to_data(), &ref_device); + let mask_ref = TestTensorBool::<3>::from_data(mask.to_data(), &ref_device); + mask.to_data().assert_eq(&mask_ref.to_data(), false); + + (tensor, value, mask, tensor_ref, value_ref, mask_ref) +} diff --git a/crates/burn-backend-tests/tests/cubecl/max_pool2d.rs b/crates/burn-backend-tests/tests/cubecl/max_pool2d.rs new file mode 100644 index 0000000..f0a2c0d --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/max_pool2d.rs @@ -0,0 +1,48 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{Distribution, module}; + +#[test] +pub fn max_pool2d_should_match_reference_backends() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::<4>::random([32, 32, 32, 32], Distribution::Default, &device); + let tensor_ref = TestTensor::<4>::from_data(tensor.to_data(), &ref_device); + let kernel_size = [3, 3]; + let stride = [2, 2]; + let padding = [1, 1]; + let dilation = [1, 1]; + + let pooled = module::max_pool2d(tensor, kernel_size, stride, padding, dilation, false); + let pooled_ref = module::max_pool2d(tensor_ref, kernel_size, stride, padding, dilation, false); + + pooled + .into_data() + .assert_approx_eq::(&pooled_ref.into_data(), Tolerance::default()); +} + +#[test] +pub fn max_pool2d_with_indices_should_match_reference_backend() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::<4>::random([32, 32, 32, 32], Distribution::Default, &device); + let tensor_ref = TestTensor::<4>::from_data(tensor.to_data(), &ref_device); + let kernel_size = [3, 3]; + let stride = [2, 2]; + let padding = [1, 1]; + let dilation = [1, 1]; + + let (pooled, indices) = + module::max_pool2d_with_indices(tensor, kernel_size, stride, padding, dilation, false); + let (pooled_ref, indices_ref) = + module::max_pool2d_with_indices(tensor_ref, kernel_size, stride, padding, dilation, false); + + pooled + .into_data() + .assert_approx_eq::(&pooled_ref.into_data(), Tolerance::default()); + indices + .into_data() + .assert_eq(&indices_ref.into_data(), false); +} diff --git a/crates/burn-backend-tests/tests/cubecl/max_pool2d_backward.rs b/crates/burn-backend-tests/tests/cubecl/max_pool2d_backward.rs new file mode 100644 index 0000000..1db27f4 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/max_pool2d_backward.rs @@ -0,0 +1,59 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{Distribution, module}; + +#[test] +pub fn max_pool2d_with_indices_backward_should_match_reference_backend() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::<4>::random([32, 32, 32, 32], Distribution::Default, &device); + let grad_output = TestTensor::<4>::random([32, 32, 16, 16], Distribution::Default, &device); + + let tensor_ref = TestTensor::<4>::from_data(tensor.to_data(), &ref_device); + let grad_output_ref = TestTensor::<4>::from_data(grad_output.to_data(), &ref_device); + let kernel_size = [3, 3]; + let stride = [2, 2]; + let padding = [1, 1]; + let dilation = [1, 1]; + + let (_, indices) = module::max_pool2d_with_indices( + tensor.clone(), + kernel_size, + stride, + padding, + dilation, + false, + ); + let (_, indices_ref) = module::max_pool2d_with_indices( + tensor_ref.clone(), + kernel_size, + stride, + padding, + dilation, + false, + ); + let grad = module::max_pool2d_with_indices_backward( + tensor, + kernel_size, + stride, + padding, + dilation, + false, + grad_output, + indices, + ); + let grad_ref = module::max_pool2d_with_indices_backward( + tensor_ref, + kernel_size, + stride, + padding, + dilation, + false, + grad_output_ref, + indices_ref, + ); + + grad.into_data() + .assert_approx_eq::(&grad_ref.into_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/cubecl/mod.rs b/crates/burn-backend-tests/tests/cubecl/mod.rs new file mode 100644 index 0000000..f7ab622 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/mod.rs @@ -0,0 +1,32 @@ +pub use super::*; + +mod avg_pool2d; +mod bernoulli; +mod cast; +mod cat; +mod clamp; +mod contiguous; +mod conv2d; +mod conv3d; +mod conv_transpose2d; +mod conv_transpose3d; +mod cross; +mod fft; +mod gather; +mod interpolate_nearest; +mod mask_fill; +mod mask_where; +mod max_pool2d; +mod max_pool2d_backward; +mod normal; +mod quantization; +mod reduce; +mod repeat_dim; +mod scatter; +mod select; +mod select_assign; +mod slice; +mod slice_assign; +mod stft; +mod unary; +mod uniform; diff --git a/crates/burn-backend-tests/tests/cubecl/normal.rs b/crates/burn-backend-tests/tests/cubecl/normal.rs new file mode 100644 index 0000000..ab6c300 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/normal.rs @@ -0,0 +1,36 @@ +use super::*; +use burn_tensor::{Device, Distribution, Shape}; +use cubek::random::{assert_mean_approx_equal, assert_normal_respects_68_95_99_rule}; +use serial_test::serial; + +#[test] +#[serial] +fn empirical_mean_close_to_expectation() { + let device = Device::default(); + device.seed(0); + + let shape = [100, 100]; + let mean = 10.; + let tensor = + TestTensor::<2>::random(shape, Distribution::Normal(mean, 2.), &device).into_data(); + let numbers = tensor.as_slice::().unwrap(); + + assert_mean_approx_equal(numbers, mean as f32); +} + +#[test] +#[serial] +fn normal_respects_68_95_99_rule() { + // https://en.wikipedia.org/wiki/68%E2%80%9395%E2%80%9399.7_rule + let shape: Shape = [1000, 1000].into(); + let device = Device::default(); + device.seed(0); + let mu = 0.; + let s = 1.; + let tensor = + TestTensor::<2>::random(shape.clone(), Distribution::Normal(mu, s), &device).into_data(); + + let numbers = tensor.as_slice::().unwrap(); + + assert_normal_respects_68_95_99_rule(numbers, mu as f32, s as f32); +} diff --git a/crates/burn-backend-tests/tests/cubecl/quantization/mod.rs b/crates/burn-backend-tests/tests/cubecl/quantization/mod.rs new file mode 100644 index 0000000..65ca126 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/quantization/mod.rs @@ -0,0 +1,14 @@ +pub use super::*; + +mod quantize_dequantize; +mod reshape; +mod slice; + +fn supports_native() -> bool { + let name = format!("{:?}", burn_tensor::Device::default()); + // TODO: Proper checks for i8 support. + name.contains("Cuda") + || name.contains("Rocm") + || name.contains("Vulkan") + || name.contains("Metal") +} diff --git a/crates/burn-backend-tests/tests/cubecl/quantization/quantize_dequantize.rs b/crates/burn-backend-tests/tests/cubecl/quantization/quantize_dequantize.rs new file mode 100644 index 0000000..d961934 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/quantization/quantize_dequantize.rs @@ -0,0 +1,235 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{ + Shape, + quantization::{QuantLevel, QuantScheme, QuantStore, QuantValue}, +}; + +fn should_quantize_dequantize_symmetric_arange>( + value: QuantValue, + store: QuantStore, + shape: S, +) { + let shape = shape.into(); + assert_eq!(shape.rank(), 2); // 2D tests + + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let scheme = QuantScheme::default().with_value(value).with_store(store); + let scheme_ref = scheme.clone().with_store(QuantStore::Native); + + let input: TestTensor<2> = TestTensorInt::arange(0..shape.num_elements() as i64, &device) + .float() + .reshape(shape); + let input_ref = TestTensor::<2>::from_data(input.to_data(), &ref_device); + + let output = input.quantize_dynamic(&scheme); + let output_ref = input_ref.quantize_dynamic(&scheme_ref); + + output.to_data().assert_eq(&output_ref.to_data(), false); + + let output = output.dequantize(); + let output_ref = output_ref.dequantize(); + + output + .into_data() + .assert_approx_eq::(&output_ref.to_data(), Tolerance::default()); +} + +fn should_quantize_dequantize_symmetric_per_block_arange>( + value: QuantValue, + block_size: usize, + store: QuantStore, + shape: S, +) { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let scheme = QuantScheme::default() + .with_value(value) + .with_level(QuantLevel::block([block_size as u8])) + .with_store(store); + let scheme_ref = scheme.clone().with_store(QuantStore::Native); + + let shape = shape.into(); + let input: TestTensor<2> = TestTensorInt::arange(0..shape.num_elements() as i64, &device) + .float() + .reshape(shape); + let input_ref = TestTensor::<2>::from_data(input.to_data(), &ref_device); + + let output = input.quantize_dynamic(&scheme); + let output_ref = input_ref.quantize_dynamic(&scheme_ref); + + output.to_data().assert_eq(&output_ref.to_data(), false); + + let output = output.dequantize(); + let output_ref = output_ref.dequantize(); + + output + .into_data() + .assert_approx_eq::(&output_ref.to_data(), Tolerance::default()); +} + +fn should_quantize_dequantize_symmetric_per_block( + value: QuantValue, + block_size: usize, + store: QuantStore, +) { + let scheme = QuantScheme::default() + .with_value(value) + .with_level(QuantLevel::block([block_size as u8])) + .with_store(store); + let scheme_ref = scheme.clone().with_store(QuantStore::Native); + + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let input = TestTensor::<2>::from_data( + [ + [ + -1.8, -1.0, 0.0, 0.5, -1.8, -1.0, 0.0, 0.5, 0.01, 0.025, 0.03, 0.04, 0.01, 0.025, + 0.03, 0.04, + ], + [ + 1.8, 1.0, 0.0, -0.5, 1.8, 1.0, 0.0, -0.5, -0.01, -0.025, -0.03, -0.04, -0.01, + -0.025, -0.03, -0.04, + ], + ], + &device, + ); + let input_ref = TestTensor::<2>::from_data(input.to_data(), &ref_device); + + let output = input.quantize_dynamic(&scheme); + let output_ref = input_ref.quantize_dynamic(&scheme_ref); + + output.to_data().assert_eq(&output_ref.to_data(), false); + + let output = output.dequantize(); + let output_ref = output_ref.dequantize(); + + output + .into_data() + .assert_approx_eq::(&output_ref.to_data(), Tolerance::default()); +} + +#[test] +fn should_quantize_dequantize_symmetric_arange_q8s_packed() { + should_quantize_dequantize_symmetric_arange(QuantValue::Q8S, QuantStore::PackedU32(0), [8, 16]) +} + +#[test] +fn should_quantize_dequantize_symmetric_arange_q8f_packed() { + should_quantize_dequantize_symmetric_arange(QuantValue::Q8F, QuantStore::PackedU32(0), [8, 16]) +} + +#[test] +fn should_quantize_dequantize_symmetric_arange_q4s_packed() { + should_quantize_dequantize_symmetric_arange(QuantValue::Q4S, QuantStore::PackedU32(0), [8, 16]) +} + +#[test] +fn should_quantize_dequantize_symmetric_arange_q4f_packed() { + should_quantize_dequantize_symmetric_arange(QuantValue::Q4F, QuantStore::PackedU32(0), [8, 16]) +} + +#[test] +fn should_quantize_dequantize_symmetric_arange_q2s_packed() { + should_quantize_dequantize_symmetric_arange(QuantValue::Q2S, QuantStore::PackedU32(0), [8, 16]) +} + +#[test] +fn should_quantize_dequantize_symmetric_arange_q2f_packed() { + should_quantize_dequantize_symmetric_arange(QuantValue::Q2F, QuantStore::PackedU32(0), [8, 16]) +} + +#[test] +fn should_quantize_dequantize_symmetric_per_block_q8s_packed() { + should_quantize_dequantize_symmetric_per_block(QuantValue::Q8S, 8, QuantStore::PackedU32(0)) +} + +#[test] +fn should_quantize_dequantize_symmetric_per_block_q4s_packed() { + should_quantize_dequantize_symmetric_per_block(QuantValue::Q4S, 8, QuantStore::PackedU32(0)) +} + +#[test] +#[should_panic] // "Block size must be divisible by 16" error is shadowed by the CallError +fn should_panic_when_block_size_cannot_store_num_quants() { + // num_quants in u32 = 32 bits / 2 bits = 16 + should_quantize_dequantize_symmetric_per_block(QuantValue::Q2S, 8, QuantStore::PackedU32(0)) +} + +#[test] +fn should_quantize_dequantize_symmetric_per_block_q2s_packed() { + should_quantize_dequantize_symmetric_per_block(QuantValue::Q2S, 16, QuantStore::PackedU32(0)) +} + +#[test] +fn should_quantize_dequantize_symmetric_arange_q8s_native() { + if supports_native() { + should_quantize_dequantize_symmetric_arange(QuantValue::Q8S, QuantStore::Native, [32, 32]) + } +} + +#[test] +fn should_quantize_dequantize_symmetric_per_block_q8s_native() { + if supports_native() { + should_quantize_dequantize_symmetric_per_block(QuantValue::Q8S, 8, QuantStore::Native) + } +} + +#[test] +fn should_quantize_dequantize_symmetric_per_block_arange_q8s_packed() { + should_quantize_dequantize_symmetric_per_block_arange( + QuantValue::Q8S, + 32, + QuantStore::PackedU32(0), + [32, 32], + ) +} + +#[test] +fn should_quantize_dequantize_symmetric_per_block_arange_q8s_native() { + if supports_native() { + should_quantize_dequantize_symmetric_per_block_arange( + QuantValue::Q8S, + 32, + QuantStore::Native, + [32, 32], + ) + } +} + +#[test] +fn should_quantize_dequantize_symmetric_arange_128x256_q8s_native() { + if supports_native() { + should_quantize_dequantize_symmetric_per_block_arange( + QuantValue::Q8S, + 32, + QuantStore::Native, + [128, 256], + ) + } +} + +#[test] +fn should_quantize_dequantize_symmetric_arange_128x256_q8s_packed() { + should_quantize_dequantize_symmetric_per_block_arange( + QuantValue::Q8S, + 32, + QuantStore::PackedU32(0), + [128, 256], + ) +} + +#[cfg(not(feature = "fusion"))] +#[test] +#[should_panic = "Can't store in u32"] +fn should_panic_when_shape_cannot_store_quants() { + let device = Default::default(); + let scheme = QuantScheme::default(); + + let _tensor_1 = TestTensor::<2>::from_data([[1.0, 6.35], [2.0, 3.0], [1.0, 3.0]], &device) + .quantize_dynamic(&scheme); +} diff --git a/crates/burn-backend-tests/tests/cubecl/quantization/reshape.rs b/crates/burn-backend-tests/tests/cubecl/quantization/reshape.rs new file mode 100644 index 0000000..e10168d --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/quantization/reshape.rs @@ -0,0 +1,151 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{ + Shape, + quantization::{BlockSize, QuantLevel, QuantScheme, QuantStore, QuantValue}, +}; + +fn should_quantize_dequantize_per_block_arange_reshaped( + level: QuantLevel, + value: QuantValue, + store: QuantStore, + shape: [usize; D1], + new_shape: [usize; D2], +) { + let numel = Shape::from(shape).num_elements() as i64; + + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let scheme = QuantScheme::default() + .with_level(level) + .with_value(value) + .with_store(store); + + let data = TestTensorInt::arange(0..numel, &ref_device) + .float() + .div_scalar(numel) + .reshape::(shape) + .into_data(); + + let input_ref = TestTensor::::from_data(data.clone(), &device).reshape::(new_shape); + let input = TestTensor::::from_data(data.clone(), &device) + .quantize_dynamic(&scheme) + .reshape::(new_shape); + + let output_ref = input_ref.into_data(); + let output = input.dequantize().into_data(); + + output.assert_approx_eq::(&output_ref, Tolerance::permissive()); +} + +#[test] +// https://github.com/tracel-ai/burn/issues/4659 +// Edge case where a single block is used, essentially like `QuantLevel::Tensor` +fn should_quantize_dequantize_per_block_reshaped_global_block_q8s_packed() { + should_quantize_dequantize_per_block_arange_reshaped( + QuantLevel::Block(BlockSize::new([16])), + QuantValue::Q8S, + QuantStore::PackedU32(0), + [32], + [2, 16], + ) +} + +#[test] +// FIXME: should work like tensor-level +#[should_panic] // "Reshape with sub-byte values is not supported"] error is shadowed by the CallError +fn should_quantize_dequantize_per_block_reshaped_global_block_q4s_packed() { + should_quantize_dequantize_per_block_arange_reshaped( + QuantLevel::Block(BlockSize::new([16])), + QuantValue::Q4S, + QuantStore::PackedU32(0), + [32], + [2, 16], + ) +} + +#[test] +// FIXME: should work +#[should_panic] // "Reshape with sub-byte values is not supported" error is shadowed by the CallError +fn should_quantize_dequantize_per_tensor_reshaped_q4s_packed() { + should_quantize_dequantize_per_block_arange_reshaped( + QuantLevel::Tensor, + QuantValue::Q4S, + QuantStore::PackedU32(0), + [32], + [2, 16], + ) +} + +#[test] +fn should_quantize_dequantize_per_block_reshaped_1d_q8s_native() { + if supports_native() { + should_quantize_dequantize_per_block_arange_reshaped( + QuantLevel::Block(BlockSize::new([16])), + QuantValue::Q8S, + QuantStore::Native, + [32], + [2, 16], + ) + } +} + +#[test] +fn should_quantize_dequantize_per_block_unsqueezed_q8s_packed() { + should_quantize_dequantize_per_block_arange_reshaped( + QuantLevel::Block(BlockSize::new([32])), + QuantValue::Q8S, + QuantStore::PackedU32(0), + [32], + [1, 1, 1, 32], + ) +} + +#[test] +#[should_panic] // "Reshape of ND block-quantized tensor is not yet supported" error is shadowed by the CallError +fn quantize_2d_block_reshape_should_panic() { + should_quantize_dequantize_per_block_arange_reshaped( + QuantLevel::Block(BlockSize::new([2, 4])), + QuantValue::Q8S, + QuantStore::PackedU32(0), + [4, 8], + [32], // invalid shape for 2D block boundaries + ) +} + +#[test] +#[should_panic] // "Reshape would split a block across multiple rows" error is shadowed by the CallError +fn quantize_per_block_reshaped_should_not_split_block() { + if supports_native() { + should_quantize_dequantize_per_block_arange_reshaped( + QuantLevel::Block(BlockSize::new([32])), + QuantValue::Q8S, + QuantStore::Native, + [2, 32], + [4, 16], + ) + } else { + // So it also panics with the same message when `QuantStore::Native` is not supported + panic!("Reshape would split a block across multiple rows") + } +} + +#[test] +#[should_panic] // "Reshape would split a block across multiple rows"] error is shadowed by the CallError +fn should_quantize_dequantize_per_block_reshaped_2d_q8s_packed() { + should_quantize_dequantize_per_block_arange_reshaped( + QuantLevel::Block(BlockSize::new([32])), + QuantValue::Q8S, + QuantStore::PackedU32(0), + [2, 32], + [4, 16], + ) +} + +// TODO: add tests for +// - ND reshape split (validation should panic) +// - broadcasted +// - multi-block successful reshape (all current success tests use exactly 1 block, e.g. [4, 32] -> [1, 4, 32] with block_size 32) +// - packed dimension alignment failure (invalid shape for packed num_quants) +// - ND-block unsqueeze (should succeed per the is_unsqueeze exemption, but nothing tests it) diff --git a/crates/burn-backend-tests/tests/cubecl/quantization/slice.rs b/crates/burn-backend-tests/tests/cubecl/quantization/slice.rs new file mode 100644 index 0000000..72b89aa --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/quantization/slice.rs @@ -0,0 +1,121 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{Shape, SliceArg, quantization::QuantScheme, s}; + +/// Slice a quantized tensor (exercising `q_slice`) and compare against slicing the same +/// quantized data in full precision (dequantized) - within a permissive tolerance. +fn should_quantize_slice_dequantize_arange(shape: [usize; D], slices: S) +where + S: SliceArg + Clone, +{ + let numel = Shape::from(shape).num_elements() as i64; + + let device = Default::default(); + + // Always use the default scheme when testing `q_slice`. + let scheme = QuantScheme::default(); + + let quantized = TestTensorInt::arange(0..numel, &device) + .float() + .div_scalar(numel) + .reshape::(shape) + .quantize_dynamic(&scheme); + + // Reference: dequantize, then slice in full precision. + let output_ref = quantized + .clone() + .dequantize() + .slice(slices.clone()) + .into_data(); + + // Slice the quantized tensor (exercises `q_slice`), then dequantize. + let output = quantized.slice(slices).dequantize().into_data(); + + output.assert_approx_eq::(&output_ref, Tolerance::permissive()); +} + +// Note: the default `QuantScheme` stores `Q8F` values packed into `u32` (4 values per element), +// so the last dimension of both the input shape and the sliced output must be a multiple of 4. + +#[test] +fn should_slice_1d_full() { + should_quantize_slice_dequantize_arange([8], s![..]); +} + +#[test] +fn should_slice_1d_range() { + should_quantize_slice_dequantize_arange([8], s![2..6]); +} + +#[test] +fn should_slice_2d_rows() { + should_quantize_slice_dequantize_arange([4, 8], s![1..3, ..]); +} + +#[test] +fn should_slice_2d_cols() { + should_quantize_slice_dequantize_arange([4, 8], s![.., 0..4]); +} + +#[test] +fn should_slice_2d_both_dims() { + should_quantize_slice_dequantize_arange([4, 8], s![1..3, 4..8]); +} + +#[test] +fn should_slice_3d() { + should_quantize_slice_dequantize_arange([2, 3, 8], s![0..1, 1..3, ..]); +} + +#[test] +fn should_slice_single_index_row() { + // A single index keeps the rank, producing a dimension of size 1. + should_quantize_slice_dequantize_arange([4, 8], s![1, ..]); +} + +#[test] +fn should_slice_inclusive_range() { + should_quantize_slice_dequantize_arange([4, 8], s![1..=2, ..]); +} + +#[test] +fn should_slice_negative_index() { + // Negative start counts from the end: last two rows. + should_quantize_slice_dequantize_arange([4, 8], s![-2.., ..]); +} + +#[test] +fn should_slice_negative_range_cols() { + // Last 4 columns via a negative range. + should_quantize_slice_dequantize_arange([4, 8], s![.., -4..]); +} + +#[test] +fn should_slice_with_step_rows() { + // Every other row. + should_quantize_slice_dequantize_arange([8, 8], s![0..8;2, ..]); +} + +#[test] +fn should_slice_with_step_cols() { + // Every other column: output last dim is ceil(8 / 2) = 4. + should_quantize_slice_dequantize_arange([4, 8], s![.., 0..8;2]); +} + +#[test] +fn should_slice_reversed_rows() { + // Reverse the first dimension (negative step). + should_quantize_slice_dequantize_arange([4, 8], s![..;-1, ..]); +} + +#[test] +fn should_slice_reversed_cols() { + // Reverse the last dimension; size stays a multiple of 4. + should_quantize_slice_dequantize_arange([4, 8], s![.., ..;-1]); +} + +#[test] +fn should_slice_partial_dims() { + // Fewer slices than dims: trailing dimension is sliced fully. + should_quantize_slice_dequantize_arange([4, 8], s![1..3]); +} diff --git a/crates/burn-backend-tests/tests/cubecl/reduce.rs b/crates/burn-backend-tests/tests/cubecl/reduce.rs new file mode 100644 index 0000000..66b9149 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/reduce.rs @@ -0,0 +1,367 @@ +use super::*; +use burn_tensor::Distribution; +use burn_tensor::Shape; +use burn_tensor::Tolerance; + +const RANK: usize = 4; +const SHAPE: [usize; RANK] = [2, 4, 8, 16]; + +#[test] +fn reduction_argmax_should_match_reference_backend() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::::random(SHAPE, Distribution::Default, &device); + let tensor_ref = TestTensor::::from_data(tensor.to_data(), &ref_device); + for dim in 0..RANK { + tensor + .clone() + .argmax(dim) + .into_data() + .assert_eq(&tensor_ref.clone().argmax(dim).into_data(), false); + } +} + +#[test] +fn reduction_argtopk_simple() { + let device = Default::default(); + + let tensor = TestTensor::<2>::from_data([[1, 7, 3], [8, 2, 8]], &device); + let actual = tensor.argtopk(2, 1); + let expected = TestTensor::<2>::from_data([[1, 2], [0, 2]], &device); + + let output_shape = Shape::new([2, 2]); + assert_eq!(actual.shape(), output_shape); + actual.into_data().assert_eq(&expected.into_data(), false); +} + +#[test] +fn reduction_argtopk_1d() { + let device = Default::default(); + + let tensor = TestTensor::<1>::from_data([10.0, 50.0, 20.0, 40.0, 30.0], &device); + let k = 3; + let actual = tensor.argtopk(k, 0); + + let expected = TestTensor::<1>::from_data([1, 3, 4], &device); + + assert_eq!(actual.shape(), Shape::new([k])); + actual.into_data().assert_eq(&expected.into_data(), false); +} + +#[test] +fn reduction_argtopk_3d_dim0() { + let device = Default::default(); + + // Shape [2, 2, 2] + let tensor = TestTensor::<3>::from_data( + [[[10.0, 1.0], [5.0, 20.0]], [[1.0, 10.0], [20.0, 5.0]]], + &device, + ); + let k = 1; + let actual = tensor.argtopk(k, 0); + + let expected = TestTensor::<3>::from_data([[[0, 1], [1, 0]]], &device); + + assert_eq!(actual.shape(), Shape::new([1, 2, 2])); + actual.into_data().assert_eq(&expected.into_data(), false); +} + +#[test] +fn reduction_argtopk_ties() { + let device = Default::default(); + + let tensor = TestTensor::<1>::from_data([5.0, 2.0, 5.0, 5.0], &device); + let k = 2; + let actual = tensor.argtopk(k, 0); + + let expected = TestTensor::<1>::from_data([0, 2], &device); + + assert_eq!(actual.shape(), Shape::new([k])); + actual.into_data().assert_eq(&expected.into_data(), false); +} + +#[test] +fn reduction_argtopk_3d_random_complex() { + let device = Default::default(); + + #[rustfmt::skip] + let tensor = TestTensor::<3>::from_data( + [ + [ + [0.5, 1.2, 0.8, 3.3], + [4.4, 2.1, 9.9, 0.1], + [7.7, 8.8, 6.6, 5.5], + ], + [ + [1.1, 0.2, 4.4, 2.2], + [6.0, 7.0, 5.0, 8.0], + [3.0, 3.0, 1.0, 2.0], + ], + ], + &device, + ); + + let k = 2; + let dim = 2; + let actual = tensor.argtopk(k, dim); + + #[rustfmt::skip] + let expected = TestTensor::<3>::from_data( + [ + [ + [3, 1], + [2, 0], + [1, 0], + ], + [ + [2, 3], + [3, 1], + [0, 1], + ], + ], + &device, + ); + + let output_shape = Shape::new([2, 3, 2]); + assert_eq!( + actual.shape(), + output_shape, + "Output shape should be [2, 3, 2]" + ); + actual.into_data().assert_eq(&expected.into_data(), false); +} + +#[test] +fn reduction_topk_simple() { + let device = Default::default(); + + let tensor = TestTensor::<2>::from_data([[1, 7, 3], [8, 2, 8]], &device); + let actual = tensor.topk(2, 1); + let expected = TestTensor::<2>::from_data([[7, 3], [8, 8]], &device); + + let output_shape = Shape::new([2, 2]); + assert_eq!(actual.shape(), output_shape); + actual.into_data().assert_eq(&expected.into_data(), false); +} + +#[test] +fn reduction_topk_3d_dim0() { + let device = Default::default(); + + // Shape [2, 2, 2] + let tensor = TestTensor::<3>::from_data( + [[[10.0, 1.0], [5.0, 20.0]], [[1.0, 10.0], [20.0, 5.0]]], + &device, + ); + let k = 1; + let actual = tensor.topk(k, 0); + + // Max values across dim 0 + let expected = TestTensor::<3>::from_data([[[10.0, 10.0], [20.0, 20.0]]], &device); + + assert_eq!(actual.shape(), Shape::new([1, 2, 2])); + actual.into_data().assert_eq(&expected.into_data(), false); +} + +#[test] +fn reduction_topk_ties() { + let device = Default::default(); + + let tensor = TestTensor::<1>::from_data([5.0, 2.0, 5.0, 5.0], &device); + let k = 2; + let actual = tensor.topk(k, 0); + + let expected = TestTensor::<1>::from_data([5.0, 5.0], &device); + + assert_eq!(actual.shape(), Shape::new([k])); + actual.into_data().assert_eq(&expected.into_data(), false); +} + +#[test] +fn reduction_topk_3d_random_complex() { + let device = Default::default(); + + #[rustfmt::skip] + let tensor = TestTensor::<3>::from_data( + [ + [ + [0.5, 1.2, 0.8, 3.3], + [4.4, 2.1, 9.9, 0.1], + [7.7, 8.8, 6.6, 5.5], + ], + [ + [1.1, 0.2, 4.4, 2.2], + [6.0, 7.0, 5.0, 8.0], + [3.0, 3.0, 1.0, 2.0], + ], + ], + &device, + ); + + let k = 2; + let dim = 2; + let actual = tensor.topk(k, dim); + + #[rustfmt::skip] + let expected = TestTensor::<3>::from_data( + [ + [ + [3.3, 1.2], + [9.9, 4.4], + [8.8, 7.7], + ], + [ + [4.4, 2.2], + [8.0, 7.0], + [3.0, 3.0], + ], + ], + &device, + ); + + let output_shape = Shape::new([2, 3, 2]); + assert_eq!( + actual.shape(), + output_shape, + "Output shape should be [2, 3, 2]" + ); + actual.into_data().assert_eq(&expected.into_data(), false); +} + +#[test] +fn reduction_argmin_should_match_reference_backend() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::::random(SHAPE, Distribution::Default, &device); + let tensor_ref = TestTensor::::from_data(tensor.to_data(), &ref_device); + for dim in 0..RANK { + tensor + .clone() + .argmin(dim) + .into_data() + .assert_eq(&tensor_ref.clone().argmin(dim).into_data(), false); + } +} + +#[test] +fn reduction_mean_dim_should_match_reference_backend() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::::random(SHAPE, Distribution::Default, &device); + let tensor_ref = TestTensor::::from_data(tensor.to_data(), &ref_device); + for dim in 0..RANK { + tensor + .clone() + .mean_dim(dim) + .into_data() + .assert_approx_eq::( + &tensor_ref.clone().mean_dim(dim).into_data(), + Tolerance::default(), + ); + } +} + +#[test] +fn reduction_mean_should_match_reference_backend() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::::random(SHAPE, Distribution::Default, &device); + let tensor_ref = TestTensor::::from_data(tensor.to_data(), &ref_device); + tensor + .clone() + .mean() + .into_data() + .assert_approx_eq::( + &tensor_ref.clone().mean().into_data(), + Tolerance::default(), + ); +} + +#[test] +fn reduction_prod_dim_should_match_reference_backend() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::::random(SHAPE, Distribution::Default, &device); + let tensor_ref = TestTensor::::from_data(tensor.to_data(), &ref_device); + for dim in 0..RANK { + tensor + .clone() + .prod_dim(dim) + .into_data() + .assert_approx_eq::( + &tensor_ref.clone().prod_dim(dim).into_data(), + Tolerance::default(), + ); + } +} + +#[test] +fn reduction_prod_should_match_reference_backend() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::::random(SHAPE, Distribution::Default, &device); + let tensor_ref = TestTensor::::from_data(tensor.to_data(), &ref_device); + tensor + .clone() + .prod() + .into_data() + .assert_approx_eq::( + &tensor_ref.clone().prod().into_data(), + Tolerance::default(), + ); +} + +#[test] +fn reduction_sum_dim_should_match_reference_backend() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::::random(SHAPE, Distribution::Default, &device); + let tensor_ref = TestTensor::::from_data(tensor.to_data(), &ref_device); + for dim in 0..RANK { + tensor + .clone() + .sum_dim(dim) + .into_data() + .assert_approx_eq::( + &tensor_ref.clone().sum_dim(dim).into_data(), + Tolerance::default(), + ); + } +} + +#[test] +fn reduction_sum_should_match_reference_backend() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::::random(SHAPE, Distribution::Default, &device); + let tensor_ref = TestTensor::::from_data(tensor.to_data(), &ref_device); + tensor + .clone() + .sum() + .into_data() + .assert_approx_eq::(&tensor_ref.clone().sum().into_data(), Tolerance::default()); +} + +#[test] +#[ignore = "Impossible to run unless you have tons of VRAM. Also reference backend is broken."] +fn reduction_sum_should_match_reference_backend_64bit() { + const SHAPE: [usize; RANK] = [33, 512, 512, 512]; + + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::::random(SHAPE, Distribution::Default, &device); + let tensor_ref = TestTensor::::from_data(tensor.to_data(), &ref_device); + let data = tensor.sum().into_data(); + let data_ref = tensor_ref.sum().into_data(); + println!("result: {:?}", data.as_slice::()); + data.assert_approx_eq::(&data_ref, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/cubecl/repeat_dim.rs b/crates/burn-backend-tests/tests/cubecl/repeat_dim.rs new file mode 100644 index 0000000..16c273f --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/repeat_dim.rs @@ -0,0 +1,75 @@ +use super::*; +use burn_tensor::Distribution; +use burn_tensor::Tolerance; + +#[test] +fn repeat_dim_0_few_times() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::<3>::random([1, 6, 6], Distribution::Default, &device); + let dim = 0; + let times = 4; + let tensor_ref = TestTensor::<3>::from_data(tensor.to_data(), &ref_device); + + let actual = tensor.repeat_dim(dim, times); + let expected = tensor_ref.repeat_dim(dim, times); + + expected + .into_data() + .assert_approx_eq::(&actual.into_data(), Tolerance::default()); +} + +#[test] +fn repeat_dim_1_few_times() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::<3>::random([6, 1, 6], Distribution::Default, &device); + let dim = 1; + let times = 4; + let tensor_ref = TestTensor::<3>::from_data(tensor.to_data(), &ref_device); + + let actual = tensor.repeat_dim(dim, times); + let expected = tensor_ref.repeat_dim(dim, times); + + expected + .into_data() + .assert_approx_eq::(&actual.into_data(), Tolerance::default()); +} + +#[test] +fn repeat_dim_2_few_times() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::<3>::random([6, 6, 1], Distribution::Default, &device); + let dim = 2; + let times = 4; + let tensor_ref = TestTensor::<3>::from_data(tensor.to_data(), &ref_device); + + let actual = tensor.repeat_dim(dim, times); + let expected = tensor_ref.repeat_dim(dim, times); + + expected + .into_data() + .assert_approx_eq::(&actual.into_data(), Tolerance::default()); +} + +#[test] +fn repeat_dim_2_many_times() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::<3>::random([10, 10, 1], Distribution::Default, &device); + let dim = 2; + let times = 200; + let tensor_ref = TestTensor::<3>::from_data(tensor.to_data(), &ref_device); + + let actual = tensor.repeat_dim(dim, times); + let expected = tensor_ref.repeat_dim(dim, times); + + expected + .into_data() + .assert_approx_eq::(&actual.into_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/cubecl/scatter.rs b/crates/burn-backend-tests/tests/cubecl/scatter.rs new file mode 100644 index 0000000..70670cf --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/scatter.rs @@ -0,0 +1,68 @@ +use super::*; +use burn_tensor::Distribution; +use burn_tensor::{Device, IndexingUpdateOp, Tolerance}; + +#[test] +fn scatter_should_work_with_multiple_workgroups_2d_dim0() { + same_as_reference_same_shape(0, [256, 32]); +} + +#[test] +fn scatter_should_work_with_multiple_workgroups_2d_dim1() { + same_as_reference_same_shape(1, [32, 256]); +} + +#[test] +fn scatter_should_work_with_multiple_workgroups_3d_dim0() { + same_as_reference_same_shape(0, [256, 6, 6]); +} + +#[test] +fn scatter_should_work_with_multiple_workgroups_3d_dim1() { + same_as_reference_same_shape(1, [6, 256, 6]); +} + +#[test] +fn scatter_should_work_with_multiple_workgroups_3d_dim2() { + same_as_reference_same_shape(2, [6, 6, 256]); +} + +#[test] +fn scatter_should_work_with_multiple_workgroups_diff_shapes() { + same_as_reference_diff_shape(1, [32, 128], [32, 1]); +} + +fn same_as_reference_diff_shape( + dim: usize, + shape1: [usize; D], + shape2: [usize; D], +) { + let device = Device::default(); + let ref_device = ReferenceDevice::new(); + + device.seed(0); + + let tensor = TestTensor::::random(shape1, Distribution::Default, &device); + let value = TestTensor::::random(shape2, Distribution::Default, &device); + let indices = TestTensorInt::<1>::random( + [shape2.iter().product::()], + Distribution::Uniform(0., shape2[dim] as f64), + &device, + ) + .reshape(shape2); + + let tensor_ref = TestTensor::::from_data(tensor.to_data(), &ref_device); + let value_ref = TestTensor::::from_data(value.to_data(), &ref_device); + let indices_ref = TestTensorInt::::from_data(indices.to_data(), &ref_device); + + let actual = tensor.scatter(dim, indices, value, IndexingUpdateOp::Add); + let expected = tensor_ref.scatter(dim, indices_ref, value_ref, IndexingUpdateOp::Add); + + expected + .into_data() + .assert_approx_eq::(&actual.into_data(), Tolerance::default()); +} + +fn same_as_reference_same_shape(dim: usize, shape: [usize; D]) { + same_as_reference_diff_shape(dim, shape, shape); +} diff --git a/crates/burn-backend-tests/tests/cubecl/select.rs b/crates/burn-backend-tests/tests/cubecl/select.rs new file mode 100644 index 0000000..ea3b3bc --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/select.rs @@ -0,0 +1,21 @@ +use super::*; +use burn_tensor::Distribution; +use burn_tensor::Tolerance; + +#[test] +fn select_should_work_with_multiple_workgroups() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::<2>::random([6, 256], Distribution::Default, &device); + let indices = TestTensorInt::<1>::arange(0..100, &device); + let tensor_ref = TestTensor::<2>::from_data(tensor.to_data(), &ref_device); + let indices_ref = TestTensorInt::<1>::from_data(indices.to_data(), &ref_device); + + let actual = tensor.select(1, indices); + let expected = tensor_ref.select(1, indices_ref); + + expected + .into_data() + .assert_approx_eq::(&actual.into_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/cubecl/select_assign.rs b/crates/burn-backend-tests/tests/cubecl/select_assign.rs new file mode 100644 index 0000000..1195507 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/select_assign.rs @@ -0,0 +1,53 @@ +use super::*; +use burn_tensor::Distribution; +use burn_tensor::{Device, IndexingUpdateOp, Tolerance}; + +#[test] +fn select_add_should_work_with_multiple_workgroups_2d_dim0() { + select_add_same_as_ref(0, [256, 6]); +} + +#[test] +fn select_add_should_work_with_multiple_workgroups_2d_dim1() { + select_add_same_as_ref(1, [6, 256]); +} + +#[test] +fn select_add_should_work_with_multiple_workgroups_3d_dim0() { + select_add_same_as_ref(0, [256, 6, 6]); +} + +#[test] +fn select_add_should_work_with_multiple_workgroups_3d_dim1() { + select_add_same_as_ref(1, [6, 256, 6]); +} + +#[test] +fn select_add_should_work_with_multiple_workgroups_3d_dim2() { + select_add_same_as_ref(2, [6, 6, 256]); +} + +fn select_add_same_as_ref(dim: usize, shape: [usize; D]) { + let device = Device::default(); + let ref_device = ReferenceDevice::new(); + + device.seed(0); + + let tensor = TestTensor::::random(shape, Distribution::Default, &device); + let value = TestTensor::::random(shape, Distribution::Default, &device); + let indices = TestTensorInt::<1>::random( + [shape[dim]], + Distribution::Uniform(0., shape[dim] as f64), + &device, + ); + let tensor_ref = TestTensor::::from_data(tensor.to_data(), &ref_device); + let value_ref = TestTensor::::from_data(value.to_data(), &ref_device); + let indices_ref = TestTensorInt::<1>::from_data(indices.to_data(), &ref_device); + + let actual = tensor.select_assign(dim, indices, value, IndexingUpdateOp::Add); + let expected = tensor_ref.select_assign(dim, indices_ref, value_ref, IndexingUpdateOp::Add); + + expected + .into_data() + .assert_approx_eq::(&actual.into_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/cubecl/slice.rs b/crates/burn-backend-tests/tests/cubecl/slice.rs new file mode 100644 index 0000000..c941a61 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/slice.rs @@ -0,0 +1,20 @@ +use super::*; +use burn_tensor::Distribution; +use burn_tensor::Tolerance; + +#[test] +fn slice_should_work_with_multiple_workgroups() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::<2>::random([6, 256], Distribution::Default, &device); + let indices = [3..5, 45..256]; + let tensor_ref = TestTensor::<2>::from_data(tensor.to_data(), &ref_device); + + let actual = tensor.slice(indices.clone()); + let expected = tensor_ref.slice(indices); + + expected + .into_data() + .assert_approx_eq::(&actual.into_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/cubecl/slice_assign.rs b/crates/burn-backend-tests/tests/cubecl/slice_assign.rs new file mode 100644 index 0000000..c65f524 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/slice_assign.rs @@ -0,0 +1,21 @@ +use super::*; +use burn_tensor::{Distribution, Tolerance}; + +#[test] +fn slice_assign_should_work_with_multiple_workgroups() { + let device = Default::default(); + let ref_device = ReferenceDevice::new(); + + let tensor = TestTensor::<2>::random([6, 256], Distribution::Default, &device); + let value = TestTensor::<2>::random([2, 211], Distribution::Default, &device); + let indices = [3..5, 45..256]; + let tensor_ref = TestTensor::<2>::from_data(tensor.to_data(), &ref_device); + let value_ref = TestTensor::<2>::from_data(value.to_data(), &ref_device); + + let actual = tensor.slice_assign(indices.clone(), value); + let expected = tensor_ref.slice_assign(indices, value_ref); + + expected + .into_data() + .assert_approx_eq::(&actual.into_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/cubecl/stft.rs b/crates/burn-backend-tests/tests/cubecl/stft.rs new file mode 100644 index 0000000..cb7df56 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/stft.rs @@ -0,0 +1,293 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::signal::{StftOptions, hann_window, istft, stft}; + +fn opts(n_fft: usize, hop_length: usize, center: bool, onesided: bool) -> StftOptions { + StftOptions { + n_fft, + hop_length, + win_length: None, + center, + onesided, + } +} + +#[test] +fn stft_constant_signal_rectangular_window() { + // Constant signal with rectangular window: only DC bin should be non-zero + let signal = TestTensor::<2>::from([[1.0, 1.0, 1.0, 1.0]]); + let result = stft(signal, None, opts(4, 1, false, true)); + + let [batch, n_frames, n_freqs, two] = result.dims(); + assert_eq!(batch, 1); + assert_eq!(n_frames, 1); // (4 - 4) / 1 + 1 = 1 + assert_eq!(n_freqs, 3); // 4/2 + 1 = 3 + assert_eq!(two, 2); + + // DC bin should be 4.0 + 0i (sum of ones) + let data = result.into_data(); + let values = data.to_vec::().unwrap(); + // [batch=0, frame=0, freq=0, re] = 4.0 + assert!( + (values[0] - 4.0).abs() < 1e-4, + "DC real should be 4.0, got {}", + values[0] + ); + // [batch=0, frame=0, freq=0, im] = 0.0 + assert!( + values[1].abs() < 1e-4, + "DC imag should be 0.0, got {}", + values[1] + ); +} + +#[test] +fn stft_output_shape_onesided() { + let signal = TestTensor::<2>::from([[1.0; 16]]); + let result = stft(signal, None, opts(8, 4, false, true)); + + let [batch, n_frames, n_freqs, two] = result.dims(); + assert_eq!(batch, 1); + assert_eq!(n_frames, 3); // (16 - 8) / 4 + 1 = 3 + assert_eq!(n_freqs, 5); // 8/2 + 1 = 5 + assert_eq!(two, 2); +} + +#[test] +fn stft_output_shape_twosided() { + let signal = TestTensor::<2>::from([[1.0; 16]]); + let result = stft(signal, None, opts(8, 4, false, false)); + + let [batch, n_frames, n_freqs, two] = result.dims(); + assert_eq!(batch, 1); + assert_eq!(n_frames, 3); + assert_eq!(n_freqs, 8); // full spectrum + assert_eq!(two, 2); +} + +#[test] +fn stft_center_padding() { + // With center=true, signal is padded by n_fft/2 on both sides + let signal = TestTensor::<2>::from([[1.0; 8]]); + let result = stft(signal, None, opts(4, 2, true, true)); + + // After padding: 2 + 8 + 2 = 12 samples + // n_frames = (12 - 4) / 2 + 1 = 5 + let [_, n_frames, _, _] = result.dims(); + assert_eq!(n_frames, 5); +} + +#[test] +fn stft_with_hann_window() { + let signal = TestTensor::<2>::from([[1.0; 8]]); + let window: TestTensor<1> = hann_window(4, true, &Default::default()); + let result = stft(signal, Some(window), opts(4, 2, false, true)); + + let [batch, n_frames, n_freqs, two] = result.dims(); + assert_eq!(batch, 1); + assert_eq!(n_frames, 3); + assert_eq!(n_freqs, 3); + assert_eq!(two, 2); +} + +#[test] +fn stft_batch_dimension() { + let signal = TestTensor::<2>::from([[1.0; 8], [2.0; 8]]); + let result = stft(signal, None, opts(4, 2, false, true)); + + let [batch, _, _, _] = result.dims(); + assert_eq!(batch, 2); +} + +#[test] +fn stft_istft_roundtrip_rectangular() { + let original = TestTensor::<2>::from([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]]); + let n_fft = 4; + let hop_length = 2; + + let o = opts(n_fft, hop_length, false, true); + let spectrum = stft(original.clone(), None, o); + let reconstructed = istft(spectrum, None, Some(8), o); + + reconstructed + .into_data() + .assert_approx_eq::(&original.into_data(), Tolerance::absolute(1e-3)); +} + +#[test] +fn stft_istft_roundtrip_centered() { + let original = TestTensor::<2>::from([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]]); + let n_fft = 4; + let hop_length = 2; + + let o = opts(n_fft, hop_length, true, true); + let spectrum = stft(original.clone(), None, o); + let reconstructed = istft(spectrum, None, Some(8), o); + + reconstructed + .into_data() + .assert_approx_eq::(&original.into_data(), Tolerance::absolute(1e-3)); +} + +#[test] +fn stft_istft_roundtrip_twosided() { + let original = TestTensor::<2>::from([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]]); + let n_fft = 4; + let hop_length = 2; + + let o = opts(n_fft, hop_length, false, false); + let spectrum = stft(original.clone(), None, o); + let reconstructed = istft(spectrum, None, Some(8), o); + + reconstructed + .into_data() + .assert_approx_eq::(&original.into_data(), Tolerance::absolute(1e-3)); +} + +#[test] +fn stft_istft_roundtrip_batch() { + let original = TestTensor::<2>::from([ + [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], + [8.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0], + ]); + let n_fft = 4; + let hop_length = 2; + + let o = opts(n_fft, hop_length, false, true); + let spectrum = stft(original.clone(), None, o); + let reconstructed = istft(spectrum, None, Some(8), o); + + reconstructed + .into_data() + .assert_approx_eq::(&original.into_data(), Tolerance::absolute(1e-3)); +} + +#[test] +fn stft_istft_roundtrip_hann_window() { + let original = TestTensor::<2>::from([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]]); + let n_fft = 4; + let hop_length = 1; + + let window: TestTensor<1> = hann_window(4, true, &Default::default()); + let o = opts(n_fft, hop_length, true, true); + let spectrum = stft(original.clone(), Some(window.clone()), o); + let reconstructed = istft(spectrum, Some(window), Some(8), o); + + reconstructed + .into_data() + .assert_approx_eq::(&original.into_data(), Tolerance::absolute(1e-2)); +} + +#[test] +fn stft_with_hamming_window() { + use burn_tensor::signal::hamming_window; + let signal = TestTensor::<2>::from([[1.0; 8]]); + let window: TestTensor<1> = hamming_window(4, true, &Default::default()); + let result = stft(signal, Some(window), opts(4, 2, false, true)); + + let [batch, n_frames, n_freqs, two] = result.dims(); + assert_eq!(batch, 1); + assert_eq!(n_frames, 3); + assert_eq!(n_freqs, 3); + assert_eq!(two, 2); +} + +#[test] +#[should_panic(expected = "hop_length")] +fn stft_rejects_hop_greater_than_window() { + // hop (5) > effective win_length (4) violates COLA/NOLA; must be rejected. + let signal = TestTensor::<2>::from([[1.0; 16]]); + let _ = stft(signal, None, opts(4, 5, false, true)); +} + +#[test] +#[should_panic(expected = "n_fft")] +fn stft_rejects_zero_nfft() { + let signal = TestTensor::<2>::from([[1.0; 4]]); + let _ = stft(signal, None, opts(0, 1, false, true)); +} + +#[test] +#[should_panic(expected = "power of two")] +fn stft_rejects_non_power_of_two_nfft() { + // n_fft=5 is not a power of two; should hard-fail in StftOptions::assert_valid. + let signal = TestTensor::<2>::from([[1.0; 8]]); + let _ = stft(signal, None, opts(5, 1, false, true)); +} + +#[test] +#[should_panic(expected = "hop_length")] +fn stft_rejects_zero_hop_length() { + let signal = TestTensor::<2>::from([[1.0; 4]]); + let _ = stft(signal, None, opts(4, 0, false, true)); +} + +#[test] +#[should_panic(expected = "win_length")] +fn stft_rejects_zero_win_length() { + let signal = TestTensor::<2>::from([[1.0; 4]]); + let o = StftOptions { + n_fft: 4, + hop_length: 1, + win_length: Some(0), + center: false, + onesided: true, + }; + let _ = stft(signal, None, o); +} + +#[test] +#[should_panic(expected = "reflect pad")] +fn stft_rejects_too_short_signal_with_center() { + // n_fft/2 = 2, signal length 2 is not > 2, so reflect pad would fail. + let signal = TestTensor::<2>::from([[1.0, 2.0]]); + let _ = stft(signal, None, opts(4, 1, true, true)); +} + +#[test] +#[should_panic(expected = "window length")] +fn istft_rejects_wrong_window_length() { + // Synthetic stft matrix: [batch=1, n_frames=3, n_freqs=3, 2 (re/im)]. + // Values are arbitrary; we only need istft to reach the window-length check. + let spectrum: TestTensor<4> = TestTensor::from([[ + [[1.0, 0.0], [0.0, 0.0], [0.0, 0.0]], + [[1.0, 0.0], [0.0, 0.0], [0.0, 0.0]], + [[1.0, 0.0], [0.0, 0.0], [0.0, 0.0]], + ]]); + + // n_fft=4, effective win_length=3 (per win_length=Some(3)); passed window length=4 mismatches. + let bad_window: TestTensor<1> = TestTensor::from([1.0, 1.0, 1.0, 1.0]); + let o_bad = StftOptions { + n_fft: 4, + hop_length: 2, + win_length: Some(3), + center: false, + onesided: true, + }; + let _ = istft(spectrum, Some(bad_window), Some(8), o_bad); +} + +#[test] +#[should_panic(expected = "n_freqs")] +fn istft_rejects_wrong_n_freqs() { + // n_fft=4, onesided=true: expected n_freqs = 4/2+1 = 3. + // Pass a spectrum with 4 bins to trigger the shape check. + let spectrum: TestTensor<4> = TestTensor::from([[ + [[1.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], + [[1.0, 0.0], [0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], + ]]); + let _ = istft(spectrum, None, Some(8), opts(4, 2, false, true)); +} + +#[test] +fn stft_options_default_and_new() { + // Spot-check the defaults match PyTorch (hop = n_fft/4, center, onesided). + let o = StftOptions::new(16); + assert_eq!(o.n_fft, 16); + assert_eq!(o.hop_length, 4); + assert_eq!(o.win_length, None); + assert!(o.center); + assert!(o.onesided); + let d = StftOptions::default(); + assert_eq!(d.n_fft, 400); +} diff --git a/crates/burn-backend-tests/tests/cubecl/unary.rs b/crates/burn-backend-tests/tests/cubecl/unary.rs new file mode 100644 index 0000000..739aff7 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/unary.rs @@ -0,0 +1,20 @@ +use super::*; + +#[test] +fn tanh_should_not_have_numerical_bugs_on_macos() { + fn tanh_one_value(input: f32) -> f32 { + let tensor = TestTensor::<1>::ones([1], &Default::default()) * input; + let output = tensor.tanh(); + output.into_data().as_slice().unwrap()[0] + } + + let ok = tanh_one_value(43.0); // metal tanh gives 1.0 which is the right answer + let zero = tanh_one_value(44.0); // metal tanh gives zero when within 43.67..44.36 + let nan = tanh_one_value(45.0); // metal tanh gives nan when over 44.36 + let neg = tanh_one_value(-45.0); // metal works correctly here + + assert!(!ok.is_nan() && ok == 1.0); + assert!(!zero.is_nan() && zero == 1.0); + assert!(!nan.is_nan() && nan == 1.0); + assert!(!neg.is_nan() && neg == -1.0); +} diff --git a/crates/burn-backend-tests/tests/cubecl/uniform.rs b/crates/burn-backend-tests/tests/cubecl/uniform.rs new file mode 100644 index 0000000..2736037 --- /dev/null +++ b/crates/burn-backend-tests/tests/cubecl/uniform.rs @@ -0,0 +1,111 @@ +use super::*; +use burn_tensor::{Device, Distribution, Shape}; +use burn_tensor::{ElementConversion, Tolerance}; + +use serial_test::serial; + +use cubek::random::{assert_at_least_one_value_per_bin, assert_wald_wolfowitz_runs_test}; + +#[test] +#[serial] +fn values_all_within_interval_default() { + let device = Device::default(); + device.seed(0); + let shape = [24, 24]; + + let tensor = TestTensor::<2>::random(shape, Distribution::Default, &device); + tensor + .to_data() + .assert_within_range::(0.elem()..1.elem()); +} + +#[test] +#[serial] +fn values_all_within_interval_uniform() { + let device = Device::default(); + device.seed(0); + let shape = [24, 24]; + + let tensor = TestTensor::<2>::random(shape, Distribution::Uniform(5., 17.), &device); + tensor + .to_data() + .assert_within_range::(5.elem()..17.elem()); +} + +#[test] +#[serial] +fn at_least_one_value_per_bin_uniform() { + let device = Device::default(); + device.seed(0); + let shape = [64, 64]; + + let tensor = + TestTensor::<2>::random(shape, Distribution::Uniform(-5., 10.), &device).into_data(); + let numbers = tensor.as_slice::().unwrap(); + + assert_at_least_one_value_per_bin(numbers, 3, -5., 10.); +} + +#[test] +#[serial] +fn runs_test() { + let device = Device::default(); + device.seed(0); + let shape = Shape::new([512, 512]); + let tensor = TestTensor::<2>::random(shape, Distribution::Default, &device).into_data(); + + let numbers = tensor.as_slice::().unwrap(); + + assert_wald_wolfowitz_runs_test(numbers, 0., 1.); +} + +#[test] +#[serial] +fn int_values_all_within_interval_uniform() { + let device = Device::default(); + device.seed(0); + let shape = Shape::new([20, 20]); + let tensor = TestTensorInt::<2>::random(shape, Distribution::Default, &device); + + let data_float = tensor.float().into_data(); + + data_float.assert_within_range(0..255); +} + +#[test] +#[serial] +fn at_least_one_value_per_bin_int_uniform() { + let device = Device::default(); + device.seed(0); + let shape = Shape::new([64, 64]); + + let tensor = TestTensorInt::<2>::random(shape, Distribution::Uniform(-10.0, 10.0), &device); + + let data_float = tensor.float().into_data(); + + let numbers = data_float.as_slice::().unwrap(); + + assert_at_least_one_value_per_bin(numbers, 10, -10., 10.); +} + +#[test] +fn should_not_fail_on_non_float_autotune() { + let device = Device::default(); + let tensor_1 = TestTensor::<2>::from_data([[1., 2., 3.], [3., 4., 5.]], &device); + + // Autotune of all (reduce) on lower_equal_elem's output calls uniform distribution + tensor_1.lower_equal_elem(1.0).all(); +} + +#[test] +#[serial] +fn test_seed_reproducibility() { + let device = Device::default(); + device.seed(42); + let t1 = TestTensor::<1>::random([5], Distribution::Default, &device); + device.seed(42); + let t2 = TestTensor::<1>::random([5], Distribution::Default, &device); + + t1.into_data() + .assert_approx_eq::(&t2.into_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/fusion/cat.rs b/crates/burn-backend-tests/tests/fusion/cat.rs new file mode 100644 index 0000000..d51f0d4 --- /dev/null +++ b/crates/burn-backend-tests/tests/fusion/cat.rs @@ -0,0 +1,341 @@ +//! Tests that `Cat` fuses with following element-wise operations into a single kernel +//! and computes correct values, including along the vectorization (last) axis. + +use super::*; +use burn_fusion::inspect::{BlockKind, FusionInspector, matchers}; +use burn_tensor::{TensorData, Tolerance}; + +/// `cat` on dim 0 followed by an element-wise op should collapse into a single +/// ElementWise fused kernel containing both operations. +#[test] +fn cat_then_elementwise_fuses_into_single_kernel() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + let a = TestTensor::<2>::from_data([[0.0, 1.0, 2.0, 3.0], [4.0, 5.0, 6.0, 7.0]], &device); + let b = + TestTensor::<2>::from_data([[8.0, 9.0, 10.0, 11.0], [12.0, 13.0, 14.0, 15.0]], &device); + let dtype = a.dtype(); + device.sync().unwrap(); + + let inspector = FusionInspector::install(stream); + let out = TestTensor::cat(vec![a, b], 0).mul_scalar(2.0); + out.into_data().assert_approx_eq::( + &TensorData::from([ + [0.0, 2.0, 4.0, 6.0], + [8.0, 10.0, 12.0, 14.0], + [16.0, 18.0, 20.0, 22.0], + [24.0, 26.0, 28.0, 30.0], + ]), + Tolerance::default(), + ); + device.sync().unwrap(); + + let reports = inspector.drain(); + let tables = reports + .iter() + .map(|report| report.format_table()) + .collect::>() + .join("\n\n"); + + let block = reports + .iter() + .flat_map(|report| report.blocks.iter()) + .find(|block| block.operations.iter().any(matchers::is_cat())) + .unwrap_or_else(|| panic!("no block containing cat found\n\n{tables}")); + + assert!( + matches!( + block.kind, + BlockKind::Fused { + name: "ElementWise", + .. + } + ), + "expected cat in an ElementWise fused block, got {:?}\n\n{tables}", + block.kind, + ); + assert!( + block + .operations + .iter() + .any(matchers::is_mul_scalar_float(dtype)), + "MulScalar should share the fused block with Cat\n\n{tables}", + ); + }); +} + +/// `cat` along the last (vectorization) axis with segments that aren't aligned to the +/// vectorization width: each element of a vector may come from a different input. +#[test] +fn cat_last_dim_with_unaligned_segments_fuses_and_computes_correctly() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + let a = TestTensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + let b = TestTensor::<2>::from_data( + [ + [10.0, 20.0, 30.0, 40.0, 50.0], + [60.0, 70.0, 80.0, 90.0, 100.0], + ], + &device, + ); + device.sync().unwrap(); + + let inspector = FusionInspector::install(stream); + let out = TestTensor::cat(vec![a, b], 1).add_scalar(1.0); + out.into_data().assert_approx_eq::( + &TensorData::from([ + [2.0, 3.0, 4.0, 11.0, 21.0, 31.0, 41.0, 51.0], + [5.0, 6.0, 7.0, 61.0, 71.0, 81.0, 91.0, 101.0], + ]), + Tolerance::default(), + ); + device.sync().unwrap(); + + let reports = inspector.drain(); + let tables = reports + .iter() + .map(|report| report.format_table()) + .collect::>() + .join("\n\n"); + + assert!( + reports + .iter() + .flat_map(|report| report.blocks.iter()) + .any(|block| matches!(block.kind, BlockKind::Fused { .. }) + && block.operations.iter().any(matchers::is_cat())), + "cat should be part of a fused block\n\n{tables}", + ); + }); +} + +/// `cat` of three inputs along a middle axis of rank-3 tensors. +#[test] +fn cat_three_inputs_middle_dim_computes_correctly() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + let a = TestTensor::<3>::from_data([[[1.0, 2.0]], [[3.0, 4.0]]], &device); + let b = TestTensor::<3>::from_data( + [[[5.0, 6.0], [7.0, 8.0]], [[9.0, 10.0], [11.0, 12.0]]], + &device, + ); + let c = TestTensor::<3>::from_data([[[13.0, 14.0]], [[15.0, 16.0]]], &device); + device.sync().unwrap(); + + let inspector = FusionInspector::install(stream); + let out = TestTensor::cat(vec![a, b, c], 1).mul_scalar(2.0); + out.into_data().assert_approx_eq::( + &TensorData::from([ + [[2.0, 4.0], [10.0, 12.0], [14.0, 16.0], [26.0, 28.0]], + [[6.0, 8.0], [18.0, 20.0], [22.0, 24.0], [30.0, 32.0]], + ]), + Tolerance::default(), + ); + device.sync().unwrap(); + + let reports = inspector.drain(); + let tables = reports + .iter() + .map(|report| report.format_table()) + .collect::>() + .join("\n\n"); + + assert!( + reports + .iter() + .flat_map(|report| report.blocks.iter()) + .any(|block| matches!(block.kind, BlockKind::Fused { .. }) + && block.operations.iter().any(matchers::is_cat())), + "cat should be part of a fused block\n\n{tables}", + ); + }); +} + +/// When a `cat` input is produced by a preceding element-wise op in the same stream +/// segment, the fuser must split (the input needs to be materialized) while still +/// producing correct values. +#[test] +fn cat_of_computed_input_splits_but_computes_correctly() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + let a = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + let b = TestTensor::<2>::from_data([[5.0, 6.0], [7.0, 8.0]], &device); + device.sync().unwrap(); + + // No sync between the mul and the cat: the fuser has to handle the split itself. + let out = TestTensor::cat(vec![a.mul_scalar(10.0), b], 0); + out.into_data().assert_approx_eq::( + &TensorData::from([[10.0, 20.0], [30.0, 40.0], [5.0, 6.0], [7.0, 8.0]]), + Tolerance::default(), + ); + device.sync().unwrap(); + }); +} + +/// `cat` on int tensors goes through the `BaseInt` path and should fuse as well. +#[test] +fn cat_int_then_elementwise_computes_correctly() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + let a = TestTensorInt::<1>::from_data([1, 2, 3], &device); + let b = TestTensorInt::<1>::from_data([4, 5], &device); + device.sync().unwrap(); + + let inspector = FusionInspector::install(stream); + let out = TestTensorInt::cat(vec![a, b], 0) + 10; + out.into_data() + .assert_eq(&TensorData::from([11, 12, 13, 14, 15]), false); + device.sync().unwrap(); + + let reports = inspector.drain(); + let tables = reports + .iter() + .map(|report| report.format_table()) + .collect::>() + .join("\n\n"); + + assert!( + reports + .iter() + .flat_map(|report| report.blocks.iter()) + .any(|block| matches!(block.kind, BlockKind::Fused { .. }) + && block.operations.iter().any(matchers::is_cat())), + "int cat should be part of a fused block\n\n{tables}", + ); + }); +} + +/// `cat` on bool tensors goes through the `BaseBool` path and should fuse with a following +/// element-wise op (`equal` is the base op available on bool). +#[test] +fn cat_bool_then_elementwise_computes_correctly() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + let a = TestTensorBool::<1>::from_data([true, false, true], &device); + let b = TestTensorBool::<1>::from_data([false, true], &device); + let c = TestTensorBool::<1>::from_data([true, true, false, false, true], &device); + device.sync().unwrap(); + + let inspector = FusionInspector::install(stream); + let out = TestTensorBool::cat(vec![a, b], 0).equal(c); + out.into_data() + .assert_eq(&TensorData::from([true, false, false, true, true]), false); + device.sync().unwrap(); + + let reports = inspector.drain(); + let tables = reports + .iter() + .map(|report| report.format_table()) + .collect::>() + .join("\n\n"); + + assert!( + reports + .iter() + .flat_map(|report| report.blocks.iter()) + .any(|block| matches!(block.kind, BlockKind::Fused { .. }) + && block.operations.iter().any(matchers::is_cat())), + "bool cat should be part of a fused block\n\n{tables}", + ); + }); +} + +/// `cat` of a non-contiguous (transposed) input: the fused kernel must follow the input's +/// actual strides. Concatenating on the last dim also exercises the per-element path. +#[test] +fn cat_transposed_input_last_dim_computes_correctly() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + let a = TestTensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + // Non-contiguous handle of shape [3, 2]. + let a = a.swap_dims(0, 1); + let b = TestTensor::<2>::from_data([[10.0, 20.0], [30.0, 40.0], [50.0, 60.0]], &device); + device.sync().unwrap(); + + let inspector = FusionInspector::install(stream); + let out = TestTensor::cat(vec![a, b], 1).mul_scalar(2.0); + out.into_data().assert_approx_eq::( + &TensorData::from([ + [2.0, 8.0, 20.0, 40.0], + [4.0, 10.0, 60.0, 80.0], + [6.0, 12.0, 100.0, 120.0], + ]), + Tolerance::default(), + ); + device.sync().unwrap(); + + let reports = inspector.drain(); + let tables = reports + .iter() + .map(|report| report.format_table()) + .collect::>() + .join("\n\n"); + + assert!( + reports + .iter() + .flat_map(|report| report.blocks.iter()) + .any(|block| matches!(block.kind, BlockKind::Fused { .. }) + && block.operations.iter().any(matchers::is_cat())), + "cat of a transposed input should be part of a fused block\n\n{tables}", + ); + }); +} + +/// The `cat` output is broadcast by the following element-wise op: the fused block's +/// reference shape is larger than the cat output along a non-cat axis, so cat input reads +/// must wrap around on that axis. +#[test] +fn cat_output_broadcast_by_elementwise_computes_correctly() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + let a = TestTensor::<2>::from_data([[1.0, 2.0, 3.0]], &device); + let b = TestTensor::<2>::from_data([[10.0, 20.0, 30.0, 40.0, 50.0]], &device); + let c = TestTensor::<2>::from_data([[100.0; 8], [200.0; 8]], &device); + device.sync().unwrap(); + + let inspector = FusionInspector::install(stream); + // cat([1, 3], [1, 5]) = [1, 8], broadcast against [2, 8]. + let out = TestTensor::cat(vec![a, b], 1) + c; + out.into_data().assert_approx_eq::( + &TensorData::from([ + [101.0, 102.0, 103.0, 110.0, 120.0, 130.0, 140.0, 150.0], + [201.0, 202.0, 203.0, 210.0, 220.0, 230.0, 240.0, 250.0], + ]), + Tolerance::default(), + ); + device.sync().unwrap(); + + let reports = inspector.drain(); + let tables = reports + .iter() + .map(|report| report.format_table()) + .collect::>() + .join("\n\n"); + + assert!( + reports + .iter() + .flat_map(|report| report.blocks.iter()) + .any(|block| matches!(block.kind, BlockKind::Fused { .. }) + && block.operations.iter().any(matchers::is_cat())), + "broadcast cat should be part of a fused block\n\n{tables}", + ); + }); +} diff --git a/crates/burn-backend-tests/tests/fusion/fusion_f16_broadcast.rs b/crates/burn-backend-tests/tests/fusion/fusion_f16_broadcast.rs new file mode 100644 index 0000000..189fe2e --- /dev/null +++ b/crates/burn-backend-tests/tests/fusion/fusion_f16_broadcast.rs @@ -0,0 +1,108 @@ +//! Regression test for fusion broadcast bug with conv + manual BN. +//! +//! Bug: When two parallel branches each do conv(64→256) followed by manual +//! batch-norm (sub, div-by-sqrt, mul, add with [1,C,1,1] broadcast params), +//! the second branch produces wrong values under Fusion. +//! +//! The sqrt(v + eps) computation on the [1,256,1,1] parameter tensor gets +//! fused into the same kernel as the broadcast ops on the [1,256,16,16] conv output + +use super::*; +use burn_tensor::{ + DType, Device, TensorCreationOptions, TensorData, Tolerance, module::conv2d, ops::ConvOptions, +}; + +const EPS: f64 = 1e-5; + +fn pseudo_random(n: usize, seed: f32, scale: f32) -> Vec { + (0..n) + .map(|i| (i as f32 * seed + 0.7).sin() * scale) + .collect() +} + +fn make_conv_weight( + in_ch: usize, + out_ch: usize, + seed: f32, + opts: TensorCreationOptions, +) -> TestTensor<4> { + // kernel_size=[1, 1]; stride=[1, 1]; dilation=[1, 1]; groups=1; + let n = out_ch * in_ch; + let scale = (2.0 / in_ch as f32).sqrt(); + let data = TensorData::new(pseudo_random(n, seed, scale), [out_ch, in_ch, 1, 1]); + TestTensor::from_data(data, opts) +} + +fn manual_bn( + x: TestTensor<4>, + ch: usize, + mean: &TestTensor<1>, + var: &TestTensor<1>, + gamma: &TestTensor<1>, + beta: &TestTensor<1>, +) -> TestTensor<4> { + let m = mean.clone().reshape([1, ch, 1, 1]); + let v = var.clone().reshape([1, ch, 1, 1]); + let g = gamma.clone().reshape([1, ch, 1, 1]); + let b = beta.clone().reshape([1, ch, 1, 1]); + let std = (v + EPS).sqrt(); + (x - m) / std * g + b +} + +fn two_conv_bn_branches(dev: Device, dtype: DType) -> TensorData { + let opts: TensorCreationOptions = (&dev, dtype).into(); + let weight_a = make_conv_weight(64, 256, 1.0, opts.clone()); + let weight_b = make_conv_weight(64, 256, 2.0, opts.clone()); + + let ch = 256; + let make_params = |seed: f32| { + let mean = TestTensor::<1>::from_data( + TensorData::new(pseudo_random(ch, seed + 300.0, 0.5), [ch]), + opts.clone(), + ); + let var_data: Vec = pseudo_random(ch, seed + 400.0, 0.3) + .iter() + .map(|v| 0.5 + v.abs()) + .collect(); + let var = TestTensor::<1>::from_data(TensorData::new(var_data, [ch]), opts.clone()); + let gamma_data: Vec = pseudo_random(ch, seed + 100.0, 0.5) + .iter() + .map(|v| 1.0 + v * 0.1) + .collect(); + let gamma = TestTensor::<1>::from_data(TensorData::new(gamma_data, [ch]), opts.clone()); + let beta = TestTensor::<1>::from_data( + TensorData::new(pseudo_random(ch, seed + 200.0, 0.3), [ch]), + opts.clone(), + ); + (mean, var, gamma, beta) + }; + let (mean_a, var_a, gamma_a, beta_a) = make_params(1.0); + let (mean_b, var_b, gamma_b, beta_b) = make_params(2.0); + dev.sync().unwrap(); + + let numel = 64 * 16 * 16; + let data: Vec = (0..numel).map(|i| (i as f32 * 7.13).sin() * 0.5).collect(); + let input = TestTensor::<4>::from_data(TensorData::new(data, [1, 64, 16, 16]), opts); + + let options = ConvOptions::new([1, 1], [0, 0], [1, 1], 1); + let conv_a_out = conv2d(input.clone(), weight_a, None, options.clone()); + let conv_b_out = conv2d(input, weight_b, None, options); + + let a = manual_bn(conv_a_out, ch, &mean_a, &var_a, &gamma_a, &beta_a); + let b = manual_bn(conv_b_out, ch, &mean_b, &var_b, &gamma_b, &beta_b); + let out = a + b; + + out.into_data() +} + +/// This test was failing only on Vulkan+Fusion+f16 +/// Reference: https://github.com/tracel-ai/burn/pull/4675 +#[test] +fn fusion_f16_two_branch_conv_manual_bn() { + let fused_f16 = two_conv_bn_branches(Default::default(), DType::F16); + let reference_f32 = two_conv_bn_branches(Default::default(), DType::F32); + + fused_f16 + .convert_dtype(DType::F32) + .assert_approx_eq::(&reference_f32, Tolerance::rel_abs(5e-3, 1e-2)); +} diff --git a/crates/burn-backend-tests/tests/fusion/fusion_f16_write_vectorization.rs b/crates/burn-backend-tests/tests/fusion/fusion_f16_write_vectorization.rs new file mode 100644 index 0000000..7ecdf58 --- /dev/null +++ b/crates/burn-backend-tests/tests/fusion/fusion_f16_write_vectorization.rs @@ -0,0 +1,107 @@ +//! Regression test for fusion kernel vector_size mismatch on output writes. +//! +//! Bug: When a fused kernel's computation width (config.width) differs from an +//! output tensor's registered vector_size, the SPIR-V store uses the wrong type +//! (e.g. vec4 into a scalar f16 slot), corrupting adjacent memory. +//! +//! This manifests as Inf/NaN values in downstream tensors. Only triggers with +//! f16 because f32 tensors tend to all receive the same vectorization on WGPU. +//! +//! The minimal trigger is two parallel conv+BatchNorm branches whose results +//! are added: the BatchNorm inplace running-stats outputs are alias tensors +//! registered with vector_size=1, while the fused block runs at width=4. + +use super::*; +use burn_tensor::{ + DType, Device, TensorCreationOptions, TensorData, Tolerance, module::conv2d, ops::ConvOptions, +}; +use std::sync::{Arc, Mutex}; + +const CH: usize = 8; +const INIT: f64 = 0.02; +const EPS: f64 = 1e-5; + +struct RunningState { + value: Arc>>, +} + +impl RunningState { + fn new(value: TestTensor<1>) -> Self { + Self { + value: Arc::new(Mutex::new(value)), + } + } + + fn value(&self) -> TestTensor<1> { + self.value.lock().unwrap().clone() + } +} + +/// Minimal clone of `BatchNorm` w/ `RunningState` for inference only. +struct BatchNorm { + gamma: TestTensor<1>, + beta: TestTensor<1>, + running_mean: RunningState, + running_var: RunningState, +} + +impl BatchNorm { + fn new(opts: TensorCreationOptions) -> Self { + Self { + gamma: TestTensor::<1>::ones([CH], opts.clone()), + beta: TestTensor::<1>::zeros([CH], opts.clone()), + running_mean: RunningState::new(TestTensor::<1>::zeros([CH], opts.clone())), + running_var: RunningState::new(TestTensor::<1>::ones([CH], opts)), + } + } + + fn forward(&self, input: TestTensor<4>) -> TestTensor<4> { + let device = input.device(); + let channels = input.dims()[1]; + let mean = self.running_mean.value().to_device(&device); + let var = self.running_var.value().to_device(&device); + + let mean = mean.reshape([1, channels, 1, 1]); + let var = var.reshape([1, channels, 1, 1]); + let std = (var + EPS).sqrt(); + + let x = (input - mean) / std; + let x = x * self.gamma.clone().reshape([1, channels, 1, 1]); + x + self.beta.clone().reshape([1, channels, 1, 1]) + } +} + +fn make_conv_weight(opts: TensorCreationOptions) -> TestTensor<4> { + TestTensor::full([CH, 3, 1, 1], INIT, opts) +} + +fn two_branch_conv_bn_add(dev: Device, dtype: DType) -> TensorData { + let opts: TensorCreationOptions = (&dev, dtype).into(); + let weight_a = make_conv_weight(opts.clone()); + let weight_b = make_conv_weight(opts.clone()); + let bn_a = BatchNorm::new(opts.clone()); + let bn_b = BatchNorm::new(opts.clone()); + + dev.sync().unwrap(); + + let input = TestTensor::<4>::ones([1, 3, 32, 32], opts) * 0.5; + let options = ConvOptions::new([1, 1], [0, 0], [1, 1], 1); + let a = bn_a.forward(conv2d(input.clone(), weight_a, None, options.clone())); + let b = bn_b.forward(conv2d(input, weight_b, None, options)); + (a + b).into_data() +} + +/// Two parallel conv+BN branches added together — the Fusion result must +/// match reference. +/// +/// This test was failing only on Vulkan+Fusion+f16 +/// Reference: https://github.com/tracel-ai/burn/pull/4675 +#[test] +fn fusion_f16_two_branch_conv_bn_add_matches_reference() { + let fused_f16 = two_branch_conv_bn_add(Default::default(), DType::F16); + let reference_f32 = two_branch_conv_bn_add(Default::default(), DType::F32); + + fused_f16 + .convert_dtype(DType::F32) + .assert_approx_eq::(&reference_f32, Tolerance::rel_abs(5e-3, 1e-2)); +} diff --git a/crates/burn-backend-tests/tests/fusion/fusion_shape.rs b/crates/burn-backend-tests/tests/fusion/fusion_shape.rs new file mode 100644 index 0000000..e8c0373 --- /dev/null +++ b/crates/burn-backend-tests/tests/fusion/fusion_shape.rs @@ -0,0 +1,277 @@ +//! Tests that assert on *which* operations were fused together, not just output values. +//! +//! These tests install a [`FusionInspector`], run a sequence of tensor ops, and then +//! inspect the captured [`FusionReport`]s to check that the expected operations ended +//! up in the same (or separate) fused kernels. +//! +//! `FusionInspector::install` captures operations on a specific stream, so two tests running +//! in parallel will not interfere when inspecting different streams. Use `test_stream()` +//! (which returns a fresh unique [`StreamId`] per call) to isolate each test. Installing +//! two inspectors for the *same* stream will panic. + +use super::*; +use burn_fusion::inspect::{BlockKind, FusionInspector, matchers}; + +/// `a + b` followed by `exp` should collapse into a single element-wise fused kernel. +#[test] +fn elementwise_add_then_exp_fuses_into_single_kernel() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + // Materialize the inputs first so that the `Ones` init ops aren't rolled into the + // fused block we're trying to observe. + let a = TestTensor::<2>::ones([4, 4], &device); + let b = TestTensor::<2>::ones([4, 4], &device); + let dtype = a.dtype(); + device.sync().unwrap(); + + let inspector = FusionInspector::install(stream); + let out = (a + b).exp(); + let _ = out.into_data(); + device.sync().unwrap(); + + let reports = inspector.drain(); + assert!(!reports.is_empty(), "expected at least one fusion report"); + + let target = reports + .iter() + .find(|r| r.total_operations() == 2) + .unwrap_or_else(|| panic!("no report with 2 ops; got {reports:#?}")); + + let block = target.assert_single_fused_block(); + assert_eq!(block.fuser_name(), Some("ElementWise")); + assert!( + block.ops_match(&[matchers::is_add_float(dtype), matchers::is_exp(dtype),]), + "block ops did not match add + exp: {:#?}", + block.operations, + ); + }); +} + +/// Forcing a sync between two sub-expressions should materialize the intermediate and +/// prevent fusion across the boundary. +#[test] +fn sync_between_ops_splits_into_separate_kernels() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + let inspector = FusionInspector::install(stream); + + let a = TestTensor::<2>::ones([4, 4], &device); + let b = TestTensor::<2>::ones([4, 4], &device); + let intermediate = a + b; + let dtype = intermediate.dtype(); + + // Force materialization of the intermediate. + device.sync().unwrap(); + + let out = intermediate.exp(); + let _ = out.into_data(); + device.sync().unwrap(); + + let reports = inspector.drain(); + + // Across all reports, we should have at least one fused add and one fused exp, and + // they should never appear together in the same block. + let add = matchers::is_add_float(dtype); + let exp = matchers::is_exp(dtype); + for report in &reports { + for block in &report.blocks { + let has_add = block.operations.iter().any(|op| add(op)); + let has_exp = block.operations.iter().any(|op| exp(op)); + assert!( + !(has_add && has_exp), + "add and exp should not share a fused block after sync: {:#?}", + block, + ); + } + } + + let saw_add = reports.iter().any(|r| { + r.blocks + .iter() + .any(|b| b.operations.iter().any(|op| add(op))) + }); + let saw_exp = reports.iter().any(|r| { + r.blocks + .iter() + .any(|b| b.operations.iter().any(|op| exp(op))) + }); + assert!( + saw_add && saw_exp, + "expected both ops to appear; got {reports:#?}" + ); + }); +} + +/// Multiple chained element-wise ops should still fuse into a single kernel. +#[test] +fn chained_elementwise_ops_fuse_together() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + let a = TestTensor::<2>::ones([8, 8], &device); + let b = TestTensor::<2>::ones([8, 8], &device); + let c = TestTensor::<2>::ones([8, 8], &device); + device.sync().unwrap(); + + let inspector = FusionInspector::install(stream); + // add → mul → exp, all element-wise on the same shape. + let out = ((a + b) * c).exp(); + let _ = out.into_data(); + device.sync().unwrap(); + + let reports = inspector.drain(); + + let target = reports + .iter() + .find(|r| r.total_operations() == 3) + .unwrap_or_else(|| panic!("no report with 3 ops; got {reports:#?}")); + + let block = target.assert_single_fused_block(); + assert!( + matches!( + block.kind, + BlockKind::Fused { + name: "ElementWise", + .. + } + ), + "expected a single ElementWise fused block, got {:?}", + block.kind, + ); + assert_eq!(block.operations.len(), 3); + }); +} + +/// A loop that materializes new tensors with `ones` and folds them into an ongoing +/// elementwise computation. Every op — Ones, MulScalar, Mul, Add across every +/// iteration — should collapse into a single fused ElementWise kernel. +/// +/// If this test fails, the panic message includes the fusion table so the split +/// points are visible directly in the test output. +#[test] +fn elementwise_and_creation_into_single_kernel() { + let stream = test_stream(); + stream.executes(|| { + const REPETITIONS: usize = 4; + let device = Default::default(); + + // Materialize the base tensor so it doesn't land in the inspector's first report. + let original = TestTensor::<2>::ones([8, 8], &device); + device.sync().unwrap(); + + let inspector = FusionInspector::install(stream); + + let mut tmp = original.clone(); + for i in 0..REPETITIONS { + let new = TestTensor::<2>::ones([8, 8], &device) * (i as f32); + tmp = tmp.clone().mul(original.clone()) + new; + } + + let _ = tmp.into_data(); + device.sync().unwrap(); + + let reports = inspector.drain(); + assert!(!reports.is_empty(), "expected at least one fusion report"); + + // Per iteration: ones, mul-by-scalar, mul, add => 4 ops. `Drop` ops (memory + // bookkeeping) also land in the reports but aren't counted here. + let expected_ops = REPETITIONS * 4; + let is_drop = matchers::is_drop(); + + let tables: String = reports + .iter() + .map(|r| r.format_table()) + .collect::>() + .join("\n\n"); + + // Everything should land in a single report... + assert_eq!( + reports.len(), + 1, + "expected 1 report, got {}\n\n{tables}", + reports.len(), + ); + let report = &reports[0]; + + // ...as a single fused block... + assert_eq!( + report.blocks.len(), + 1, + "expected 1 block, got {}\n\n{tables}", + report.blocks.len(), + ); + let block = &report.blocks[0]; + + // ...produced by the ElementWise fuser... + assert_eq!( + block.fuser_name(), + Some("ElementWise"), + "block was not ElementWise\n\n{tables}", + ); + + // ...containing exactly the expected non-Drop ops. + let non_drop_ops = block.operations.iter().filter(|op| !is_drop(op)).count(); + assert_eq!( + non_drop_ops, expected_ops, + "expected {expected_ops} non-Drop ops in the fused block, got {non_drop_ops}\n\n{tables}", + ); + }); +} + +/// A lone view (`swap_dims`) whose original tensor isn't otherwise read must NOT be fused: it runs +/// eagerly as a cheap shape/stride update, so the following element-wise math reads a clean, +/// vectorizable layout instead of recomputing the transposed index on every element. +/// +/// Regression guard for the "view fix": previously the swap-dims was fused into the element-wise +/// kernel (the original registered as a fresh input, read *through* the transpose), which defeats +/// vectorization for no benefit when the original is used nowhere else. +#[test] +fn lone_view_is_not_fused_into_kernel() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + // Materialize the inputs so their `ones` init ops don't land in the inspected reports. + let a = TestTensor::<2>::ones([4, 8], &device); + let b = TestTensor::<2>::ones([8, 4], &device); + device.sync().unwrap(); + + let inspector = FusionInspector::install(stream); + + // `swap_dims` is the only use of `a`. + let out = (a.swap_dims(0, 1) + b).exp(); + let _ = out.into_data(); + device.sync().unwrap(); + + let reports = inspector.drain(); + let tables: String = reports + .iter() + .map(|r| r.format_table()) + .collect::>() + .join("\n\n"); + + let is_swap = matchers::is_swap_dims(); + + // The view must never end up inside a fused kernel. + assert!( + reports + .iter() + .flat_map(|r| r.fused_blocks()) + .all(|blk| !blk.operations.iter().any(|op| is_swap(op))), + "swap_dims must not be fused (view fix):\n\n{tables}", + ); + + // ...and it must actually run, in an un-fused block. + assert!( + reports + .iter() + .flat_map(|r| r.unfused_blocks()) + .any(|blk| blk.operations.iter().any(|op| is_swap(op))), + "swap_dims should run eagerly in an un-fused block:\n\n{tables}", + ); + }); +} diff --git a/crates/burn-backend-tests/tests/fusion/inplace.rs b/crates/burn-backend-tests/tests/fusion/inplace.rs new file mode 100644 index 0000000..03cfd94 --- /dev/null +++ b/crates/burn-backend-tests/tests/fusion/inplace.rs @@ -0,0 +1,236 @@ +//! Tests that fused kernels actually write in-place (`HandleOutput::Alias`) when they can. +//! +//! Aliasing produces identical values to allocating a fresh output, so value-based tests +//! cannot catch a regression — the whole mechanism silently died for over a year without a +//! single test failing. These tests assert on the alias counter from +//! [`burn_cubecl_fusion::inspect`]. +//! +//! The counter is process-global, so assertions are deltas (`>=`): concurrent tests can +//! only inflate it. "Does NOT alias" is asserted through values instead (a wrong alias +//! corrupts the still-referenced input). + +use super::*; +use burn_cubecl_fusion::inspect::inplace_alias_count; +use burn_fusion::inspect::{BlockKind, FusionInspector}; +use burn_tensor::{TensorData, Tolerance}; + +fn assert_all_fused(reports: &[burn_fusion::inspect::FusionReport], what: &str) { + let tables = reports + .iter() + .map(|report| report.format_table()) + .collect::>() + .join("\n\n"); + + assert!( + !reports.is_empty() + && reports + .iter() + .flat_map(|report| report.blocks.iter()) + .all(|block| matches!(block.kind, BlockKind::Fused { .. })), + "{what}: every block should be fused, otherwise the aliasing assertion is \ + meaningless\n\n{tables}", + ); +} + +/// A consuming element-wise chain must reuse the input buffer for its output. +/// +/// This covers the two launch-side regressions that killed aliasing: the `can_mut()` +/// reference miscount (container clone kept during planning) and the reference-layout +/// gate that rejected the first output of every block. +#[test] +fn consuming_elemwise_chain_writes_inplace() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + let mut tensor = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + // Warmup with the same (relative) graph: first execution may go through autotune, + // where benchmark trials run on forked contexts that (correctly) never alias. + tensor = tensor.add_scalar(0.0).mul_scalar(1.0); + let _ = tensor.clone().into_data(); + device.sync().unwrap(); + + let inspector = FusionInspector::install(stream); + let before = inplace_alias_count(); + for _ in 0..4 { + tensor = tensor.add_scalar(1.0).mul_scalar(2.0); + // Sync so every iteration is its own execution (own launch plan). + let _ = tensor.clone().into_data(); + } + device.sync().unwrap(); + let after = inplace_alias_count(); + + tensor.into_data().assert_approx_eq::( + &TensorData::from([[46.0, 62.0], [78.0, 94.0]]), + Tolerance::default(), + ); + assert_all_fused(&inspector.drain(), "consuming elemwise chain"); + assert!( + after - before >= 4, + "each consuming iteration should alias its input buffer \ + (got {} aliases over 4 iterations)", + after - before, + ); + }); +} + +/// The consumed input keeps its buffer even when another input is broadcast: the output +/// (same shape as the consumed input) must alias it, and values must stay correct. This +/// is the shape that regressed with an invalid vector-size declaration on the alias. +#[test] +fn broadcast_elemwise_writes_inplace_and_computes_correctly() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + let bias = TestTensor::<2>::from_data([[10.0, 20.0, 30.0, 40.0]], &device); + let mut tensor = + TestTensor::<2>::from_data([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], &device); + tensor = tensor.add_scalar(0.0).mul_scalar(1.0); + let _ = tensor.clone().into_data(); + device.sync().unwrap(); + + let inspector = FusionInspector::install(stream); + let before = inplace_alias_count(); + let out = (tensor + bias).mul_scalar(2.0); + out.into_data().assert_approx_eq::( + &TensorData::from([[22.0, 44.0, 66.0, 88.0], [30.0, 52.0, 74.0, 96.0]]), + Tolerance::default(), + ); + device.sync().unwrap(); + let after = inplace_alias_count(); + + assert_all_fused(&inspector.drain(), "broadcast elemwise"); + assert!( + after > before, + "the broadcast chain should alias the consumed input buffer", + ); + }); +} + +/// A tensor that stays alive (read-only use) must NOT be written in-place: its values +/// must survive the op. Asserted through values — the alias counter is process-global, +/// so "did not move" can't be checked under parallel tests, but a wrong alias would +/// corrupt the still-referenced input. +#[test] +fn read_only_input_is_not_written_inplace() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + let mut tensor = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + tensor = tensor.add_scalar(0.0).mul_scalar(1.0); + let _ = tensor.clone().into_data(); + device.sync().unwrap(); + + // `tensor` stays alive across the chain: its buffer must not be reused. + let out = tensor.clone().add_scalar(1.0).mul_scalar(2.0); + out.into_data().assert_approx_eq::( + &TensorData::from([[4.0, 6.0], [8.0, 10.0]]), + Tolerance::default(), + ); + device.sync().unwrap(); + + // The original tensor must be untouched. + tensor.into_data().assert_approx_eq::( + &TensorData::from([[1.0, 2.0], [3.0, 4.0]]), + Tolerance::default(), + ); + }); +} + +/// A consumed elemwise input feeding a fused reduce must alias the elemwise output. +/// +/// The aliased output becomes the read block's reference layout, so the reduce runner +/// resolves the reference against an aliased output argument: this covers the +/// output-arg reference form and the `TensorArg::Alias` arms of +/// `GlobalArgsLaunch::shape/strides`, which elemwise-only chains don't exercise. +#[test] +fn reduce_fusion_elemwise_output_writes_inplace() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + let make_input = || { + let tensor = + TestTensor::<2>::from_data([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], &device); + // Materialize the init op before the fused chain, so the chain starts from + // an existing buffer and the init doesn't show up in the inspected reports. + let _ = tensor.clone().into_data(); + device.sync().unwrap(); + tensor + }; + + // `out` stays alive across the reduce, so it is a global output of the fused + // read block, while `tensor` is consumed and can be aliased. The elemwise ops + // on both sides of the reduce each eliminate an intermediate, giving the reduce + // fuser an I/O saving over a plain elemwise fusion (which must stop at + // `sum_dim`) so the reduce optimization wins the block. + let run = |tensor: TestTensor<2>, offset: f32| { + let out = tensor.add_scalar(offset).mul_scalar(2.0); + let sum = out.clone().sum_dim(1); + let res = sum.mul_scalar(3.0).add_scalar(offset); + (out.into_data(), res.into_data()) + }; + + // Warmup with the same (relative) graph: first execution may go through autotune, + // where benchmark trials run on forked contexts that (correctly) never alias. + let _ = run(make_input(), 0.0); + device.sync().unwrap(); + + let input = make_input(); + let inspector = FusionInspector::install(stream); + let before = inplace_alias_count(); + let (out, res) = run(input, 1.0); + device.sync().unwrap(); + let after = inplace_alias_count(); + + out.assert_approx_eq::( + &TensorData::from([[4.0, 6.0, 8.0, 10.0], [12.0, 14.0, 16.0, 18.0]]), + Tolerance::default(), + ); + res.assert_approx_eq::( + &TensorData::from([[85.0], [181.0]]), + Tolerance::default(), + ); + assert_all_fused(&inspector.drain(), "reduce with consumed elemwise input"); + assert!( + after > before, + "the elemwise output feeding the fused reduce should alias the consumed input buffer", + ); + }); +} + +/// Repeated consuming updates (the KV-cache pattern) must alias on every iteration while +/// the cached execution plan is reused, and values must stay correct throughout. +#[test] +fn repeated_consuming_updates_alias_with_cached_plan() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + let mut acc = TestTensor::<1>::zeros([32], &device); + acc = acc.add_scalar(0.0).mul_scalar(1.0); + let _ = acc.clone().into_data(); + device.sync().unwrap(); + + let inspector = FusionInspector::install(stream); + let before = inplace_alias_count(); + for step in 1..=3 { + let values = TestTensor::<1>::full([32], step as f32, &device); + acc = (acc + values).add_scalar(1.0); + let _ = acc.clone().into_data(); + } + device.sync().unwrap(); + let after = inplace_alias_count(); + + acc.into_data() + .assert_approx_eq::(&TensorData::from([9.0; 32]), Tolerance::default()); + assert_all_fused(&inspector.drain(), "repeated consuming updates"); + assert!( + after - before >= 3, + "each accumulation step should alias (got {} aliases over 3 steps)", + after - before, + ); + }); +} diff --git a/crates/burn-backend-tests/tests/fusion/int_bitwise.rs b/crates/burn-backend-tests/tests/fusion/int_bitwise.rs new file mode 100644 index 0000000..7111bdc --- /dev/null +++ b/crates/burn-backend-tests/tests/fusion/int_bitwise.rs @@ -0,0 +1,94 @@ +use super::*; +use burn_fusion::inspect::{BlockKind, FusionInspector, matchers}; +use burn_tensor::TensorData; + +#[test] +fn int_bitwise_chain_and_float_cast_fuse_into_single_kernel() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + + let lhs = TestTensorInt::<1>::from_data([13, 7, 3, 12], &device); + let rhs = TestTensorInt::<1>::from_data([11, 3, 5, 9], &device); + device.sync().unwrap(); + + let inspector = FusionInspector::install(stream); + + let bitwise = lhs + .clone() + .bitwise_xor(rhs) + .bitwise_and_scalar(7) + .bitwise_right_shift_scalar(1) + .bitwise_or(lhs.bitwise_not()); + + let output = bitwise.float().mul_scalar(0.5); + let dtype = output.dtype(); + output + .into_data() + .assert_eq(&TensorData::from([-6.5_f32, -3.0, -0.5, -6.5]), false); + device.sync().unwrap(); + + let reports = inspector.drain(); + let tables = reports + .iter() + .map(|report| report.format_table()) + .collect::>() + .join("\n\n"); + + let block = reports + .iter() + .flat_map(|report| report.blocks.iter()) + .find(|block| block.operations.iter().any(matchers::is_bitwise_xor_int())) + .unwrap_or_else(|| panic!("no bitwise fused block found\n\n{tables}")); + + assert!( + matches!( + block.kind, + BlockKind::Fused { + name: "ElementWise", + .. + } + ), + "expected ElementWise fused block, got {:?}\n\n{tables}", + block.kind, + ); + + assert!( + block.operations.iter().any(matchers::is_bitwise_xor_int()), + "BitwiseXor missing from fused block\n\n{tables}", + ); + assert!( + block + .operations + .iter() + .any(matchers::is_bitwise_and_scalar_int()), + "BitwiseAndScalar missing from fused block\n\n{tables}", + ); + assert!( + block + .operations + .iter() + .any(matchers::is_bitwise_right_shift_scalar_int()), + "BitwiseRightShiftScalar missing from fused block\n\n{tables}", + ); + assert!( + block.operations.iter().any(matchers::is_bitwise_not_int()), + "BitwiseNot missing from fused block\n\n{tables}", + ); + assert!( + block.operations.iter().any(matchers::is_bitwise_or_int()), + "BitwiseOr missing from fused block\n\n{tables}", + ); + assert!( + block.operations.iter().any(matchers::is_int_into_float()), + "IntoFloat missing from fused block\n\n{tables}", + ); + assert!( + block + .operations + .iter() + .any(matchers::is_mul_scalar_float(dtype)), + "MulScalar missing from fused block\n\n{tables}", + ); + }); +} diff --git a/crates/burn-backend-tests/tests/fusion/mod.rs b/crates/burn-backend-tests/tests/fusion/mod.rs new file mode 100644 index 0000000..2445787 --- /dev/null +++ b/crates/burn-backend-tests/tests/fusion/mod.rs @@ -0,0 +1,23 @@ +pub use super::*; + +mod cat; +mod fusion_f16_broadcast; +mod fusion_f16_write_vectorization; +mod fusion_shape; +// Asserts on launch-level decisions from `burn-cubecl-fusion`, which only the `cube` +// feature pulls in; the rest of this suite also runs on non-cube fusion backends (flex). +#[cfg(feature = "cube")] +mod inplace; +mod int_bitwise; +mod reduce_broadcasted; + +use burn_tensor::StreamId; +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Returns a unique `StreamId` for test isolation. +pub fn test_stream() -> StreamId { + static COUNTER: AtomicU64 = AtomicU64::new(1000); + StreamId { + value: COUNTER.fetch_add(1, Ordering::Relaxed), + } +} diff --git a/crates/burn-backend-tests/tests/fusion/reduce_broadcasted.rs b/crates/burn-backend-tests/tests/fusion/reduce_broadcasted.rs new file mode 100644 index 0000000..e083208 --- /dev/null +++ b/crates/burn-backend-tests/tests/fusion/reduce_broadcasted.rs @@ -0,0 +1,158 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn test_reduce_broadcasted_1() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..32, &device) + .reshape([4, 8]) + .float(); + let fused_on_read = TestTensorInt::<1>::arange(0..32, &device) + .reshape([4, 8]) + .float(); + let fused_on_write = TestTensorInt::<1>::arange(0..4, &device) + .reshape([4, 1]) + .float(); + + // Forces previous tensors to be materialized. + device.sync().unwrap(); + + let x = tensor + fused_on_read.clone(); + let x = x.sum_dim(1); + + let x = x + fused_on_write; + + // Broadcast + let end = x + fused_on_read; + let actual = end.into_data(); + let expected = TensorData::from([ + [56.0, 57.0, 58.0, 59.0, 60.0, 61.0, 62.0, 63.0], + [193.0, 194.0, 195.0, 196.0, 197.0, 198.0, 199.0, 200.0], + [330.0, 331.0, 332.0, 333.0, 334.0, 335.0, 336.0, 337.0], + [467.0, 468.0, 469.0, 470.0, 471.0, 472.0, 473.0, 474.0], + ]); + actual.assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_reduce_broadcasted_2() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..32, &device) + .reshape([4, 8]) + .float(); + let fused_on_read = TestTensorInt::<1>::arange(0..32, &device) + .reshape([4, 8]) + .float(); + let fused_on_write = TestTensorInt::<1>::arange(16..48, &device) + .reshape([4, 8]) + .float(); + // Second fuse on read + let y = TestTensorInt::<1>::arange(32..64, &device) + .reshape([4, 8]) + .float(); + + // Forces previous tensors to be materialized. + device.sync().unwrap(); + + let x = tensor + fused_on_read.clone(); + let x = x.sum_dim(1); + let x = x + fused_on_write; + let x = x.mean_dim(1); + + let end = x + y; + device.sync().unwrap(); + + let actual = end.into_data(); + let expected = TensorData::from([ + [107.5, 108.5, 109.5, 110.5, 111.5, 112.5, 113.5, 114.5], + [251.5, 252.5, 253.5, 254.5, 255.5, 256.5, 257.5, 258.5], + [395.5, 396.5, 397.5, 398.5, 399.5, 400.5, 401.5, 402.5], + [539.5, 540.5, 541.5, 542.5, 543.5, 544.5, 545.5, 546.5], + ]); + actual.assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_reduce_broadcasted_3() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..32, &device) + .reshape([4, 8]) + .float(); + let fused_on_read = TestTensorInt::<1>::arange(0..32, &device) + .reshape([4, 8]) + .float(); + let fused_on_write = TestTensorInt::<1>::arange(0..4, &device) + .reshape([4, 1]) + .float(); + // Second fuse on read + let y = TestTensorInt::<1>::arange(32..64, &device) + .reshape([4, 8]) + .float(); + + // Forces previous tensors to be materialized. + device.sync().unwrap(); + + let x = tensor + fused_on_read.clone(); + let x = x.sum_dim(1); + + let x = x + fused_on_write; + + // Broadcast + let x = x + fused_on_read; + // Second reduce + let x = x.mean_dim(1); + + let end = x + y; + let actual = end.into_data(); + let expected = TensorData::from([ + [91.5, 92.5, 93.5, 94.5, 95.5, 96.5, 97.5, 98.5], + [236.5, 237.5, 238.5, 239.5, 240.5, 241.5, 242.5, 243.5], + [381.5, 382.5, 383.5, 384.5, 385.5, 386.5, 387.5, 388.5], + [526.5, 527.5, 528.5, 529.5, 530.5, 531.5, 532.5, 533.5], + ]); + actual.assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_reduce_broadcasted_4_reused_partial() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..32, &device) + .reshape([4, 8]) + .float(); + let fused_on_read = TestTensorInt::<1>::arange(0..32, &device) + .reshape([4, 8]) + .float(); + let fused_on_write = TestTensorInt::<1>::arange(0..4, &device) + .reshape([4, 1]) + .float(); + let y = TestTensorInt::<1>::arange(32..64, &device) + .reshape([4, 8]) + .float(); + + // Forces previous tensors to be materialized. + device.sync().unwrap(); + + // In fusion we have to create a global buffer to keep the intermediate data for now. + let x_previous = tensor + fused_on_read; + let x = x_previous.clone().sum_dim(1); + + let x = x * fused_on_write; + + // Broadcast + let x = x + x_previous; + // Second reduce + let x = x.mean_dim(1); + + // Second fuse on read + let end = x + y; + let actual = end.into_data(); + let expected = TensorData::from([ + [39.0, 40.0, 41.0, 42.0, 43.0, 44.0, 45.0, 46.0], + [247.0, 248.0, 249.0, 250.0, 251.0, 252.0, 253.0, 254.0], + [711.0, 712.0, 713.0, 714.0, 715.0, 716.0, 717.0, 718.0], + [ + 1431.0, 1432.0, 1433.0, 1434.0, 1435.0, 1436.0, 1437.0, 1438.0, + ], + ]); + actual.assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/graph_capture.rs b/crates/burn-backend-tests/tests/graph_capture.rs new file mode 100644 index 0000000..cc1db64 --- /dev/null +++ b/crates/burn-backend-tests/tests/graph_capture.rs @@ -0,0 +1,68 @@ +//! Graph capture/replay integration tests for the closure-based +//! [`capture`](burn_tensor::capture) API. +//! +//! On a backend with hardware graph support (CUDA/HIP) the capture records the +//! closure's launches and every replay is a single dispatch against the +//! original buffers; elsewhere replay falls back to re-running the closure. +//! Both paths must produce the eager result. +//! +//! Isolated in this test binary: `capture` arms device-global allocation state +//! (the persistent pool) between `graph_prepare` and `stop_capture`, so it must +//! not interleave with unrelated tests allocating on the same device. Tests are +//! additionally `#[serial]` so two captures never overlap. + +#![cfg(feature = "cube")] + +extern crate alloc; + +pub type FloatElem = f32; +#[allow(unused)] +pub type IntElem = i32; + +#[path = "common/backend.rs"] +mod backend; +pub use backend::*; + +use burn_tensor::{Device, Tolerance}; +use serial_test::serial; + +/// Replaying a captured pure closure reproduces the eager result, repeatedly. +/// +/// Safety of the `replay` calls: the closure owns a clone of `input` (keeping +/// every captured buffer alive as long as the graph), and all replays and +/// output reads happen sequentially on this thread — nothing else touches the +/// graph's tensors. +#[test] +#[serial] +fn capture_replay_matches_eager() { + let device = Device::default(); + let input = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + let expected = input.clone().mul_scalar(2.0).add_scalar(1.0).into_data(); + + let mut graph = burn_tensor::capture(&device, || input.clone().mul_scalar(2.0).add_scalar(1.0)); + + for _ in 0..3 { + let out = unsafe { graph.replay() }.clone().into_data(); + out.assert_approx_eq::(&expected, Tolerance::default()); + } +} + +/// The output handle is stable across replays on the hardware path: `output()` +/// returns the same tensor whose buffer each replay overwrites. +#[test] +#[serial] +fn capture_output_is_stable() { + let device = Device::default(); + let input = TestTensor::<1>::from_data([1.0, 2.0, 3.0, 4.0], &device); + let expected = input.clone().add_scalar(10.0).into_data(); + + let mut graph = burn_tensor::capture(&device, || input.clone().add_scalar(10.0)); + + // Safety: see `capture_replay_matches_eager`. + unsafe { graph.replay() }; + graph + .output() + .clone() + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/memory_cleaning.rs b/crates/burn-backend-tests/tests/memory_cleaning.rs new file mode 100644 index 0000000..ffe4afe --- /dev/null +++ b/crates/burn-backend-tests/tests/memory_cleaning.rs @@ -0,0 +1,1030 @@ +//! Tests that the fusion backend releases every tensor handle once the corresponding +//! [`FusionTensor`](burn_fusion::FusionTensor) wrappers are dropped; both for tensors +//! that live on a single stream and for tensors that are shared across streams. +//! +//! # Assertion Strategy +//! The assertion target is [`FusionInspector::new_handles_since_baseline`], which +//! returns every [`TensorId`](burn_ir::TensorId) that appeared in the +//! [`HandleContainer`](burn_ir::HandleContainer) *after* the baseline was set. +//! +//! Because the [`HandleContainer`](burn_ir::HandleContainer) is a global, per-device +//! registry, these tests call [`FusionInspector::enable_leak_detection`] to acquire +//! an exclusive lock on the device's handle state. This prevents concurrent tests +//! from allocating tensors that would otherwise appear as "leaks" in the baseline diff. +//! +//! # Threading and Concurrency +//! While execution reports are stream-isolated, memory handles are not. Therefore: +//! 1. These tests are isolated in this test binary to avoid false "leaks". +//! 2. Every test is marked `#[serial]` to ensure they take turns capturing the +//! device-wide handle state. +//! +//! Cross-stream cases are simulated with [`StreamId::executes`], which swaps the +//! per-thread stream id for the duration of a closure. This is fast and deterministic; +//! the real OS-thread path is exercised by `tensor/multi_threads.rs`. + +#![cfg(all(feature = "cube", feature = "fusion"))] + +extern crate alloc; + +pub type FloatElem = f32; +#[allow(unused)] +pub type IntElem = i32; + +#[path = "common/backend.rs"] +mod backend; +pub use backend::*; + +#[path = "fusion/mod.rs"] +mod fusion { + + use super::*; + use burn_fusion::inspect::FusionInspector; + use burn_tensor::{Device, StreamId, Transaction}; + use serial_test::serial; + use std::sync::atomic::{AtomicU64, Ordering}; + + /// Returns a unique `StreamId` for test isolation. + pub fn test_stream() -> StreamId { + static COUNTER: AtomicU64 = AtomicU64::new(1000); + StreamId { + value: COUNTER.fetch_add(1, Ordering::Relaxed), + } + } + + /// Drain both the main stream and `peer` stream so the post-sync handle snapshot + /// reflects the full quiescent state. + fn sync_all(device: &Device, peer: StreamId) { + device.sync().unwrap(); + peer.executes(|| device.sync().unwrap()); + device.sync().unwrap(); + } + + /// Drive the inspector into a known state, then freeze that state as the baseline. + /// After this returns, [`FusionInspector::new_handles_since_baseline`] only reports + /// handles born during the test body — handles from other parallel tests that were + /// already live are excluded. + fn install_and_baseline(device: &Device, stream: StreamId) -> FusionInspector { + let inspector = FusionInspector::install(stream); + device.sync().unwrap(); + inspector.set_baseline(); + inspector + } + + /// Assert that no handles registered during the test's window are still alive. + #[track_caller] + fn assert_no_leaked_handles(inspector: &FusionInspector, context: &str) { + let leaked = inspector.new_handles_since_baseline(); + assert!( + leaked.is_empty(), + "{context}: {count} handle(s) registered during the test are still live: {leaked:?}", + count = leaked.len(), + ); + } + + /// Baseline: a tensor created and consumed on a single stream must not leave any + /// handles behind. + #[test] + #[serial] + fn single_stream_drop_frees_all_handles() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, stream); + + { + let a = TestTensor::<2>::ones([4, 4], &device); + let b = TestTensor::<2>::ones([4, 4], &device); + let _ = (a + b).into_data(); + } + device.sync().unwrap(); + + assert_no_leaked_handles(&inspector, "single-stream drop + sync"); + }); + } + + /// `into_data` consumes its source — the source handle must be released when the + /// stream is drained. + #[test] + #[serial] + fn into_data_releases_source_handle() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, stream); + + let a = TestTensor::<2>::ones([4, 4], &device); + let _ = a.into_data(); + device.sync().unwrap(); + + assert_no_leaked_handles(&inspector, "into_data + sync"); + }); + } + + /// Cloning a tensor does not allocate a new handle; dropping the last `Arc`-style + /// reference should release exactly one handle. + #[test] + #[serial] + fn clone_then_drop_frees_handle_once() { + let stream = test_stream(); + stream.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, stream); + + let a = TestTensor::<2>::ones([4, 4], &device); + let b = a.clone(); + device.sync().unwrap(); + + drop(a); + drop(b); + device.sync().unwrap(); + + assert_no_leaked_handles(&inspector, "clone + drop + sync"); + }); + } + + /// Cross-stream sharing happy path: tensor created on the main stream, used and + /// consumed on a peer stream, then dropped on the main stream. After every stream is + /// drained, no handles must remain. + #[test] + #[serial] + fn cross_stream_shared_tensor_released() { + let stream = test_stream(); + let peer = test_stream(); + + stream.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, stream); + + // Materialize `a` so its `Ones` init op doesn't get rolled into the cross-stream + // analysis we're trying to observe. + let a = TestTensor::<2>::ones([4, 4], &device); + device.sync().unwrap(); + + // Performing an op on `a` from peer stream is what flags `a` as shared inside + // `SharedTensors::analyse` (its creation stream is main, current is peer). + let a_for_peer = a.clone(); + peer.executes(|| { + let _ = (a_for_peer * 2.0).into_data(); + }); + + drop(a); + sync_all(&device, peer); + + assert_no_leaked_handles( + &inspector, + "shared tensor handle should be released once every sharing stream is drained", + ); + }); + } + + /// The owner stream releases its reference *before* the peer stream catches up. + /// The shared-tensor bookkeeping must defer the release until the peer drains; once + /// it does, the handle must be reclaimed. + #[test] + #[serial] + fn owner_drops_before_peer_finishes() { + let stream = test_stream(); + let peer = test_stream(); + + stream.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, stream); + + let a = TestTensor::<2>::ones([4, 4], &device); + device.sync().unwrap(); + + // Queue work on the peer stream that depends on `a`, but do not drain yet so the + // peer keeps a pending reference. + let a_for_peer = a.clone(); + let peer_pending = peer.executes(|| a_for_peer * 2.0); + + // Owner releases its reference first, while peer still holds a pending op. + drop(a); + + // Peer eventually catches up. + peer.executes(|| { + let _ = peer_pending.into_data(); + }); + + sync_all(&device, peer); + + assert_no_leaked_handles( + &inspector, + "deferred shared-tensor drop should fire once the peer stream catches up", + ); + }); + } + + /// The peer stream fully consumes its clone and closes *before* the owner drops its + /// reference. The owner's later drop must still free the handle and not get stuck + /// waiting on a stream that no longer exists. + #[test] + #[serial] + fn peer_closes_before_owner_drops() { + let stream = test_stream(); + let peer = test_stream(); + + stream.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, stream); + + let a = TestTensor::<2>::ones([4, 4], &device); + device.sync().unwrap(); + + let a_for_peer = a.clone(); + peer.executes(|| { + let _ = (a_for_peer * 2.0).into_data(); + }); + // Peer is now closed: every var consumed, queue empty. + + drop(a); + sync_all(&device, peer); + + assert_no_leaked_handles( + &inspector, + "owner-side drop after peer closure should still free the handle", + ); + }); + } + + /// Stress: many iterations of cross-stream sharing back-to-back. The handle count + /// must return to baseline at the end — catches regressions where each iteration + /// leaks one or more handles (the failure mode the `drop-triggered drain` warning at + /// `MultiStream::register` is meant to prevent). + #[test] + #[serial] + fn cross_stream_loop_no_leak() { + const REPS: usize = 8; + let stream = test_stream(); + let peer = test_stream(); + + stream.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, stream); + + for i in 0..REPS { + let a = TestTensor::<2>::ones([4, 4], &device) * (i as f32); + device.sync().unwrap(); + + let a_for_peer = a.clone(); + peer.executes(|| { + let _ = (a_for_peer + (i as f32)).into_data(); + }); + + let _ = (a * 3.0).into_data(); + } + + sync_all(&device, peer); + + assert_no_leaked_handles( + &inspector, + "handle set must return to baseline after repeated cross-stream cleanup", + ); + }); + } + + /// Drain `extra` peer streams in addition to the main stream. Used by tests that + /// share a single tensor across more than two streams. + fn sync_streams(device: &Device, extra: &[StreamId]) { + device.sync().unwrap(); + for s in extra { + s.executes(|| device.sync().unwrap()); + } + device.sync().unwrap(); + } + + /// A tensor created on the owner stream is consumed by three independent peer streams + /// concurrently (no peer-to-peer chain). Every peer must observe the same data and the + /// shared handle must be reclaimed after every stream has drained. + #[test] + #[serial] + fn shared_tensor_three_peers_released() { + let owner = test_stream(); + let p1 = test_stream(); + let p2 = test_stream(); + let p3 = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + let a = TestTensor::<2>::ones([4, 4], &device); + device.sync().unwrap(); + + // Each peer takes its own clone — each cross-stream `clone()` registers a + // shared view aliasing `a`'s handle under a fresh tensor id. + let c1 = a.clone(); + let c2 = a.clone(); + let c3 = a.clone(); + p1.executes(|| { + let _ = (c1 * 2.0).into_data(); + }); + p2.executes(|| { + let _ = (c2 + 1.0).into_data(); + }); + p3.executes(|| { + let _ = (c3 - 1.0).into_data(); + }); + + drop(a); + sync_streams(&device, &[p1, p2, p3]); + + assert_no_leaked_handles( + &inspector, + "handle should be freed once every peer stream drains its alias", + ); + }); + } + + /// Re-sharing path: tensor flows owner → peer → grandpeer. The grandpeer takes its + /// clone from a tensor whose home stream is already `peer`, so the share path is + /// triggered a second time from a different source stream. + #[test] + #[serial] + fn chained_share_across_three_streams() { + let owner = test_stream(); + let peer = test_stream(); + let grandpeer = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + let a = TestTensor::<2>::ones([4, 4], &device); + device.sync().unwrap(); + + // First hop: owner -> peer. The clone happens on `owner` (regular clone), + // then the multiplication on `peer` consumes the clone and triggers the + // cross-stream `into_ir` path. + let a_for_peer = a.clone(); + let on_peer = peer.executes(|| a_for_peer * 2.0); + drop(a); + + // Second hop: peer -> grandpeer. `on_peer.stream == peer`; cloning while + // current is grandpeer fires `shared_view` with `peer` as source. + let on_peer_for_grand = grandpeer.executes(|| on_peer.clone()); + grandpeer.executes(|| { + let _ = (on_peer_for_grand + 1.0).into_data(); + }); + + // The peer tensor stays alive until here so peer's queue still references it. + peer.executes(|| { + let _ = on_peer.into_data(); + }); + + sync_streams(&device, &[peer, grandpeer]); + + assert_no_leaked_handles( + &inspector, + "handle set must return to baseline after a chained owner→peer→grandpeer share", + ); + }); + } + + /// Share a tensor that was *computed* on the owner stream rather than freshly + /// initialised. This makes sure the drain-on-first-share path correctly flushes + /// pending ops before exposing the handle to the peer. + #[test] + #[serial] + fn share_computed_tensor_flushes_pending_ops() { + let owner = test_stream(); + let peer = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + // Build a pending computation on owner without draining it. + let a = TestTensor::<2>::ones([4, 4], &device); + let b = (a.clone() + 1.0) * 3.0; // pending: (1 + 1) * 3 = 6 + // Intentionally do NOT sync — peer's first-share must drain the source. + + let b_for_peer = b.clone(); + let peer_data = peer.executes(|| (b_for_peer + 0.0).into_data()); + + // 6 + 0 = 6 + peer_data.assert_approx_eq::( + &burn_tensor::TensorData::from([ + [6.0_f32, 6.0, 6.0, 6.0], + [6.0, 6.0, 6.0, 6.0], + [6.0, 6.0, 6.0, 6.0], + [6.0, 6.0, 6.0, 6.0], + ]), + burn_tensor::Tolerance::default(), + ); + + drop(a); + drop(b); + sync_streams(&device, &[peer]); + + assert_no_leaked_handles( + &inspector, + "computed-tensor share must flush pending ops and leave no handles behind", + ); + }); + } + + /// Peer takes a clone but never executes any op on it before it is dropped. + /// The aliased handle must still be reclaimed once the unused clone is freed. + #[test] + #[serial] + fn share_then_drop_unused_clone() { + let owner = test_stream(); + let peer = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + let a = TestTensor::<2>::ones([4, 4], &device); + device.sync().unwrap(); + + // Move the clone over to the peer stream but don't touch it. + let unused = peer.executes(|| a.clone()); + peer.executes(|| drop(unused)); + + drop(a); + sync_streams(&device, &[peer]); + + assert_no_leaked_handles( + &inspector, + "unused cross-stream clone must not leak its aliased handle", + ); + }); + } + + /// Owner keeps using the tensor on its own stream after exposing a clone to the + /// peer. Both sides must finish cleanly and the handle must be released only after + /// the last reader drains. + #[test] + #[serial] + fn owner_keeps_using_after_share() { + let owner = test_stream(); + let peer = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + let a = TestTensor::<2>::ones([4, 4], &device); + device.sync().unwrap(); + + // Hand a clone to peer. + let a_for_peer = a.clone(); + peer.executes(|| { + let _ = (a_for_peer * 5.0).into_data(); + }); + + // Owner continues doing work on `a` after the share point. + let owner_result = (a.clone() + 2.0) * 3.0; + let _ = owner_result.into_data(); + drop(a); + + sync_streams(&device, &[peer]); + + assert_no_leaked_handles( + &inspector, + "owner must be able to keep using its tensor after sharing a clone", + ); + }); + } + + /// Two distinct tensors are independently shared between the same pair of streams. + /// Each share path must be tracked separately — dropping one must not affect the + /// other. + #[test] + #[serial] + fn multiple_distinct_shared_tensors_same_pair() { + let owner = test_stream(); + let peer = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + let a = TestTensor::<2>::ones([4, 4], &device); + let b = TestTensor::<2>::ones([4, 4], &device) * 2.0; + device.sync().unwrap(); + + let a_for_peer = a.clone(); + let b_for_peer = b.clone(); + + peer.executes(|| { + let _ = (a_for_peer + b_for_peer).into_data(); + }); + + drop(a); + drop(b); + sync_streams(&device, &[peer]); + + assert_no_leaked_handles( + &inspector, + "two independently shared tensors must each be released", + ); + }); + } + + /// Int tensors travel through the same shared-view path as floats, but go through + /// a different dtype branch in the runner. Cover that branch explicitly. + #[test] + #[serial] + fn int_tensor_cross_stream_released() { + let owner = test_stream(); + let peer = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + let a = TestTensorInt::<1>::arange(0..16, &device); + device.sync().unwrap(); + + let a_for_peer = a.clone(); + peer.executes(|| { + let _ = (a_for_peer + 1).into_data(); + }); + + drop(a); + sync_streams(&device, &[peer]); + + assert_no_leaked_handles( + &inspector, + "int tensor sharing must release its aliased handle", + ); + }); + } + + /// The peer performs several chained ops on the shared tensor. The shared handle + /// is consumed (read-only) once on the peer side then dropped; we verify the full + /// chain executes without leaks. + #[test] + #[serial] + fn peer_runs_chain_on_shared_tensor() { + let owner = test_stream(); + let peer = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + let a = TestTensor::<2>::ones([4, 4], &device); + device.sync().unwrap(); + + let a_for_peer = a.clone(); + peer.executes(|| { + let x = a_for_peer * 2.0; + let x = x + 1.0; + let x = x * 4.0; + let _ = x.into_data(); + }); + + drop(a); + sync_streams(&device, &[peer]); + + assert_no_leaked_handles( + &inspector, + "multi-step chain on a shared tensor must not leak", + ); + }); + } + + /// Owner sends a clone to peer; peer immediately sends a clone back to owner. The + /// round trip exercises shared-view in both directions on a single underlying + /// allocation. + #[test] + #[serial] + fn share_round_trip_owner_peer_owner() { + let owner = test_stream(); + let peer = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + let a = TestTensor::<2>::ones([4, 4], &device); + device.sync().unwrap(); + + // owner -> peer + let a_for_peer = a.clone(); + let on_peer = peer.executes(|| a_for_peer + 1.0); + + // peer -> owner (cross-stream clone the other way) + let back_to_owner = on_peer.clone(); + + // Both sides consume their copy. + let _ = (back_to_owner * 2.0).into_data(); + peer.executes(|| { + let _ = on_peer.into_data(); + }); + + drop(a); + sync_streams(&device, &[peer]); + + assert_no_leaked_handles( + &inspector, + "round-trip share owner→peer→owner must leave no handles behind", + ); + }); + } + + /// `Tensor::into_data` consumes the tensor. When the consuming call happens on a + /// stream different from the tensor's home stream, `into_ir` (not `clone`) must + /// itself trigger the shared-view aliasing. + #[test] + #[serial] + fn into_data_cross_stream_no_explicit_clone() { + let owner = test_stream(); + let peer = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + let a = TestTensor::<2>::ones([4, 4], &device); + device.sync().unwrap(); + + // No explicit `.clone()` here — `a` is moved into the peer's closure and + // consumed there. The cross-stream `into_ir` path must do the right thing. + peer.executes(|| { + let _ = a.into_data(); + }); + + sync_streams(&device, &[peer]); + + assert_no_leaked_handles( + &inspector, + "cross-stream into_ir must alias the handle without an explicit clone", + ); + }); + } + + /// Same source tensor shared multiple times *to the same peer stream*. The + /// first share triggers a drain of the source; subsequent shares of the same + /// source must skip the drain (`MultiStream::resolved` short-circuit) but still + /// register a fresh aliased handle. + #[test] + #[serial] + fn repeated_share_same_source_same_peer() { + let owner = test_stream(); + let peer = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + let a = TestTensor::<2>::ones([4, 4], &device); + device.sync().unwrap(); + + // Take three clones of the same `a`, all consumed on the peer stream. + let c1 = a.clone(); + let c2 = a.clone(); + let c3 = a.clone(); + peer.executes(|| { + let _ = (c1 + 1.0).into_data(); + let _ = (c2 + 2.0).into_data(); + let _ = (c3 + 3.0).into_data(); + }); + + drop(a); + sync_streams(&device, &[peer]); + + assert_no_leaked_handles( + &inspector, + "repeated shares of the same source to the same peer must not leak", + ); + }); + } + + /// The peer must observe the values the owner produced. Most other tests only + /// check that handles are released; this one pins down data-side correctness across + /// the shared-view aliasing boundary. + #[test] + #[serial] + fn peer_observes_owner_values_after_share() { + let owner = test_stream(); + let peer = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + // Owner stages a non-trivial pending computation: (1 + 1) * 5 = 10, then + // shares the result with the peer *without* a manual sync. The first-share + // drain must flush the chain so the peer sees 10, not garbage. + let a = TestTensor::<2>::ones([2, 3], &device); + let b = (a.clone() + 1.0) * 5.0; + let b_for_peer = b.clone(); + + let peer_data = peer.executes(|| (b_for_peer + 0.5).into_data()); + peer_data.assert_approx_eq::( + &burn_tensor::TensorData::from([[10.5_f32, 10.5, 10.5], [10.5, 10.5, 10.5]]), + burn_tensor::Tolerance::default(), + ); + + // Owner still sees the same upstream values via its own copy of `b`. + let owner_data = b.into_data(); + owner_data.assert_approx_eq::( + &burn_tensor::TensorData::from([[10.0_f32, 10.0, 10.0], [10.0, 10.0, 10.0]]), + burn_tensor::Tolerance::default(), + ); + + drop(a); + sync_streams(&device, &[peer]); + + assert_no_leaked_handles( + &inspector, + "peer/owner views of the shared tensor must agree numerically and not leak", + ); + }); + } + + /// Bool tensors take a separate dtype branch through `register_bool_tensor` and the + /// boolean ops; sharing must route through the same path without leaks. + #[test] + #[serial] + fn bool_tensor_cross_stream_released() { + let owner = test_stream(); + let peer = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + let a = TestTensor::<2>::ones([4, 4], &device); + // Build a bool tensor on owner: `a == a` is all-true. + let mask = a.clone().equal(a.clone()); + device.sync().unwrap(); + + let mask_for_peer = mask.clone(); + peer.executes(|| { + let any_true = mask_for_peer.any(); + let data = any_true.into_data(); + let v: bool = data.iter::().next().expect("scalar bool"); + assert!(v, "shared bool tensor must remain all-true on peer side"); + }); + + drop(a); + drop(mask); + sync_streams(&device, &[peer]); + + assert_no_leaked_handles( + &inspector, + "bool tensor sharing must release its aliased handle", + ); + }); + } + + /// A sliced view is a separate `FusionTensor` with its own id that depends on the + /// parent's data. Sharing the view across streams must alias the *view's* handle + /// (which is materialised by the slice op on the source stream) — not the parent's. + #[test] + #[serial] + fn sliced_view_cross_stream_released() { + let owner = test_stream(); + let peer = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + let a = TestTensor::<2>::from_floats( + [ + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + [13.0, 14.0, 15.0, 16.0], + ], + &device, + ); + + // Slice on owner: 2x2 top-left block. The slice op produces a new tensor. + let slice = a.clone().slice([0..2, 0..2]); + let slice_for_peer = slice.clone(); + + // Peer doubles every element of the view and reads it back. + let peer_data = peer.executes(|| (slice_for_peer * 2.0).into_data()); + peer_data.assert_approx_eq::( + &burn_tensor::TensorData::from([[2.0_f32, 4.0], [10.0, 12.0]]), + burn_tensor::Tolerance::default(), + ); + + drop(a); + drop(slice); + sync_streams(&device, &[peer]); + + assert_no_leaked_handles( + &inspector, + "sliced-view sharing must release both parent and view aliases", + ); + }); + } + + /// High fan-out: a single source tensor is shared with eight peer streams. This + /// stresses the `MultiStream::resolved` short-circuit (one drain, many aliased + /// handles) and the handle-removal path when every alias eventually drops. + #[test] + #[serial] + fn high_fan_out_eight_peers_released() { + let owner = test_stream(); + let peers: Vec = (0..8).map(|_| test_stream()).collect(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + let a = TestTensor::<2>::ones([4, 4], &device); + device.sync().unwrap(); + + // Each peer takes its own clone and consumes it. + let clones: Vec<_> = (0..peers.len()).map(|_| a.clone()).collect(); + for (peer, c) in peers.iter().zip(clones) { + let peer = *peer; + peer.executes(|| { + let _ = (c + 1.0).into_data(); + }); + } + + drop(a); + sync_streams(&device, &peers); + + assert_no_leaked_handles( + &inspector, + "many fan-out peers must each release their aliased handle", + ); + }); + } + + /// Real OS-thread cross-stream sharing with leak detection. The other cross-stream + /// tests use [`StreamId::executes`] to swap the per-thread id in place; this one + /// actually spawns threads (the production code path) and checks that handles are + /// reclaimed after every thread joins. + #[test] + #[serial] + fn real_threads_cross_stream_no_leak() { + let owner = test_stream(); + owner.executes(|| { + let device: Device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + let a = TestTensor::<2>::ones([4, 4], &device); + device.sync().unwrap(); + + // Hand each worker thread its own clone. Cross-thread `clone()` is exactly + // the path `FusionTensor::clone` takes in user code that spawns threads. + let mut handles = Vec::new(); + for i in 0..4 { + let c = a.clone(); + let dev = device.clone(); + handles.push(std::thread::spawn(move || { + let _ = (c + i as f32).into_data(); + dev.sync().unwrap(); + })); + } + for h in handles { + h.join().unwrap(); + } + + drop(a); + device.sync().unwrap(); + + assert_no_leaked_handles( + &inspector, + "real-thread cross-stream sharing must not leak handles", + ); + }); + } + + /// Two peer streams *both* take a clone of the same source, and one of them + /// re-shares its clone onward to a third stream. Exercises the case where a + /// non-owner stream is itself the source of a `tag_shared_view` call. + #[test] + #[serial] + fn peer_reshares_to_third_stream() { + let owner = test_stream(); + let p1 = test_stream(); + let p2 = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + let a = TestTensor::<2>::ones([4, 4], &device); + device.sync().unwrap(); + + // owner -> p1 (consumed) + let c1 = a.clone(); + let on_p1 = p1.executes(|| c1 + 1.0); // on_p1.stream == p1, value = 2 + + // p1 -> p2 (re-share); the source stream for tag_shared_view is p1, not owner. + let on_p1_clone = p2.executes(|| on_p1.clone()); + let p2_data = p2.executes(|| (on_p1_clone * 3.0).into_data()); + p2_data.assert_approx_eq::( + &burn_tensor::TensorData::from([ + [6.0_f32, 6.0, 6.0, 6.0], + [6.0, 6.0, 6.0, 6.0], + [6.0, 6.0, 6.0, 6.0], + [6.0, 6.0, 6.0, 6.0], + ]), + burn_tensor::Tolerance::default(), + ); + + // Drain p1's pending tensor. + p1.executes(|| { + let _ = on_p1.into_data(); + }); + + drop(a); + sync_streams(&device, &[p1, p2]); + + assert_no_leaked_handles( + &inspector, + "peer-as-source share path must release every alias", + ); + }); + } + + /// Both streams have pending (un-drained) work referencing the shared tensor when + /// the top-level test syncs. The drain path on either stream must converge. + #[test] + #[serial] + fn concurrent_pending_work_on_both_streams() { + let owner = test_stream(); + let peer = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + let a = TestTensor::<2>::ones([4, 4], &device); + device.sync().unwrap(); + + // Build a pending op on peer, do not drain. + let a_for_peer = a.clone(); + let peer_pending = peer.executes(|| a_for_peer * 2.0); + + // Build a pending op on owner, do not drain. + let owner_pending = a.clone() + 5.0; + + // Now drop `a` and finalise both sides via a coordinated sync. + drop(a); + + peer.executes(|| { + let _ = peer_pending.into_data(); + }); + let _ = owner_pending.into_data(); + + sync_streams(&device, &[peer]); + + assert_no_leaked_handles( + &inspector, + "both streams holding pending work on a shared tensor must converge cleanly", + ); + }); + } + + #[test] + #[serial] + fn transaction_cross_stream_no_reentrant_panic() { + let owner = test_stream(); + let peer = test_stream(); + + owner.executes(|| { + let device = Default::default(); + let inspector = install_and_baseline(&device, owner); + + // Create and compute a tensor on the owner stream, leave it un-drained. + let a = TestTensor::<2>::ones([4, 4], &device); + let b = (a.clone() + 1.0) * 2.0; + + // Execute the transaction from the peer stream, `b.stream == owner` but + // `current == peer`, so `resolve_tensor_float` hits the cross-stream mismatch + // and previously called `into_ir()` inside the `submit_blocking` closure, + // triggering the re-entrancy panic. + let [data] = peer.executes(|| { + Transaction::default() + .register(b) + .execute() + .try_into() + .unwrap() + }); + + data.assert_approx_eq::( + &burn_tensor::TensorData::from([ + [4.0_f32, 4.0, 4.0, 4.0], + [4.0, 4.0, 4.0, 4.0], + [4.0, 4.0, 4.0, 4.0], + [4.0, 4.0, 4.0, 4.0], + ]), + burn_tensor::Tolerance::default(), + ); + + drop(a); + sync_streams(&device, &[peer]); + + assert_no_leaked_handles( + &inspector, + "transaction executed cross-stream must not panic or leak handles", + ); + }); + } +} diff --git a/crates/burn-backend-tests/tests/tensor.rs b/crates/burn-backend-tests/tests/tensor.rs new file mode 100644 index 0000000..5cbb8f3 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor.rs @@ -0,0 +1,21 @@ +//! Burn backend tensor tests. + +#![recursion_limit = "256"] +#![allow(clippy::single_range_in_vec_init)] + +extern crate alloc; + +pub type FloatElem = f32; +#[allow(unused)] +pub type IntElem = i32; + +#[path = "common/backend.rs"] +mod backend; +pub use backend::*; + +#[path = "common/tensor.rs"] +mod tensor; + +#[cfg(all(feature = "cube", feature = "fusion"))] +#[path = "fusion/mod.rs"] +mod fusion; diff --git a/crates/burn-backend-tests/tests/tensor/bool/mod.rs b/crates/burn-backend-tests/tests/tensor/bool/mod.rs new file mode 100644 index 0000000..5d0928d --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/mod.rs @@ -0,0 +1,3 @@ +pub use super::*; // re-export test types + +mod ops; diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/all.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/all.rs new file mode 100644 index 0000000..e7439d6 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/all.rs @@ -0,0 +1,34 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_all() { + let tensor = TestTensorBool::<2>::from([[false, true, false], [true, true, true]]); + let data_actual = tensor.all().into_data(); + let data_expected = TensorData::from([false]); + data_expected.assert_eq(&data_actual, false); + + let tensor = TestTensorBool::<2>::from([[true, true, true], [true, true, true]]); + let data_actual = tensor.all().into_data(); + let data_expected = TensorData::from([true]); + data_expected.assert_eq(&data_actual, false); +} + +#[test] +fn test_all_dim() { + let tensor = TestTensorBool::<2>::from([[false, true, false], [true, true, true]]); + let data_actual = tensor.all_dim(1).into_data(); + let data_expected = TensorData::from([[false], [true]]); + data_expected.assert_eq(&data_actual, false); +} + +#[test] +fn test_all_with_bool_from_lower_equal() { + let tensor1 = TestTensor::<2>::from([[0.0, 1.0, 0.0], [1.0, -1.0, 1.0]]) + 1e-6; + let tensor2 = TestTensor::from([[0.0, 1.0, 0.0], [1.0, -1.0, 1.0]]) + 1e-6; + + let ge = tensor1.lower_equal(tensor2); + let all = ge.clone().all(); + + TensorData::from([true]).assert_eq(&all.clone().into_data(), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/any.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/any.rs new file mode 100644 index 0000000..fd414fa --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/any.rs @@ -0,0 +1,23 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_any() { + let tensor = TestTensorBool::<2>::from([[false, false, false], [true, true, false]]); + let data_actual = tensor.any().into_data(); + let data_expected = TensorData::from([true]); + data_expected.assert_eq(&data_actual, false); + + let tensor = TestTensorBool::<2>::from([[false, false, false], [false, false, false]]); + let data_actual = tensor.any().into_data(); + let data_expected = TensorData::from([false]); + data_expected.assert_eq(&data_actual, false); +} + +#[test] +fn test_any_dim() { + let tensor = TestTensorBool::<2>::from([[false, false, false], [true, true, false]]); + let data_actual = tensor.any_dim(1).into_data(); + let data_expected = TensorData::from([[false], [true]]); + data_expected.assert_eq(&data_actual, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/argwhere_nonzero.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/argwhere_nonzero.rs new file mode 100644 index 0000000..2e2cb26 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/argwhere_nonzero.rs @@ -0,0 +1,107 @@ +use super::*; +use alloc::vec::Vec; +use burn_tensor::{Shape, TensorData}; + +#[test] +fn test_argwhere_1d() { + let tensor = TestTensorBool::<1>::from([false, true, false, true, true]); + let output = tensor.argwhere(); + + output + .into_data() + .assert_eq(&TensorData::from([[1], [3], [4]]), false); +} + +#[test] +fn test_argwhere_2d() { + let tensor = TestTensorBool::<2>::from([[false, false], [false, true], [true, true]]); + let output = tensor.argwhere(); + + output + .into_data() + .assert_eq(&TensorData::from([[1, 1], [2, 0], [2, 1]]), false); +} + +#[test] +fn test_argwhere_3d() { + let tensor = TestTensorBool::<3>::from([ + [[false, false, false], [false, true, false]], + [[true, false, true], [true, true, false]], + ]); + let output = tensor.argwhere(); + + output.into_data().assert_eq( + &TensorData::from([[0, 1, 1], [1, 0, 0], [1, 0, 2], [1, 1, 0], [1, 1, 1]]), + false, + ); +} + +#[test] +fn test_nonzero_1d() { + let tensor = TestTensorBool::<1>::from([false, true, false, true, true]); + let data_actual = tensor + .nonzero() + .into_iter() + .map(|t| t.into_data()) + .collect::>(); + + assert_eq!(data_actual.len(), 1); + data_actual[0].assert_eq(&TensorData::from([1, 3, 4]), false); +} + +#[test] +fn test_nonzero_2d() { + // 2-D tensor + let tensor = TestTensorBool::<2>::from([[false, false], [false, true], [true, true]]); + let data_actual = tensor + .nonzero() + .into_iter() + .map(|t| t.into_data()) + .collect::>(); + let data_expected = [TensorData::from([1, 2, 2]), TensorData::from([1, 0, 1])]; + + assert_eq!(data_actual.len(), 2); + for (idx, actual) in data_actual.iter().enumerate() { + actual.assert_eq(&data_expected[idx], false) + } +} + +#[test] +fn test_nonzero_3d() { + // 3-D tensor + let tensor = TestTensorBool::<3>::from([ + [[false, false, false], [false, true, false]], + [[true, false, true], [true, true, false]], + ]); + let data_actual = tensor + .nonzero() + .into_iter() + .map(|t| t.into_data()) + .collect::>(); + let data_expected = [ + TensorData::from([0, 1, 1, 1, 1]), + TensorData::from([1, 0, 0, 1, 1]), + TensorData::from([1, 0, 2, 0, 1]), + ]; + + assert_eq!(data_actual.len(), 3); + for (idx, actual) in data_actual.iter().enumerate() { + actual.assert_eq(&data_expected[idx], false) + } +} + +#[test] +fn test_nonzero_empty() { + let tensor = TestTensorBool::<1>::from([false, false, false, false, false]); + let output = tensor.nonzero(); + + assert_eq!(output.len(), 0); +} + +#[test] +fn test_argwhere_empty() { + let tensor = TestTensorBool::<1>::from([false, false, false, false, false]); + let output = tensor.argwhere(); + + assert_eq!(output.shape(), Shape::new([0, 1])); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/cast.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/cast.rs new file mode 100644 index 0000000..8070359 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/cast.rs @@ -0,0 +1,59 @@ +use super::*; +use burn_tensor::{DType, FloatDType, IntDType, TensorData}; + +#[test] +fn cast_bool_to_int_with_dtype() { + let tensor = TestTensorBool::<2>::from([[true, false, true], [false, false, true]]); + + let output = tensor.cast(IntDType::I32); + + assert_eq!(output.dtype(), DType::I32); + let expected = TensorData::from([[1i32, 0, 1], [0, 0, 1]]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn cast_bool_to_float_with_dtype() { + let tensor = TestTensorBool::<2>::from([[true, false, true], [false, false, true]]); + + let output = tensor.cast(FloatDType::F32); + + assert_eq!(output.dtype(), DType::F32); + let expected = TensorData::from([[1.0f32, 0.0, 1.0], [0.0, 0.0, 1.0]]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_bool_into_int_flipped() { + // [T, F, T, F] flipped -> [F, T, F, T] -> int [0, 1, 0, 1] + let t = TestTensorBool::<1>::from([true, false, true, false]).flip([0]); + + let output = t.int(); + + output + .into_data() + .assert_eq(&TensorData::from([0i64, 1, 0, 1]), false); +} + +#[test] +fn test_bool_into_float_flipped() { + let t = TestTensorBool::<1>::from([true, false, true, false]).flip([0]); + + let output = t.float(); + + output + .into_data() + .assert_eq(&TensorData::from([0.0f32, 1.0, 0.0, 1.0]), false); +} + +#[test] +fn test_bool_into_int_flipped_2d() { + // [[T, F], [F, T]] flipped on axis 0 -> [[F, T], [T, F]] -> int [[0, 1], [1, 0]] + let t = TestTensorBool::<2>::from([[true, false], [false, true]]).flip([0]); + + let output = t.int(); + + output + .into_data() + .assert_eq(&TensorData::from([[0i64, 1], [1, 0]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/cat.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/cat.rs new file mode 100644 index 0000000..3dcda65 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/cat.rs @@ -0,0 +1,29 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_cat_ops_bool() { + let device = Default::default(); + let tensor_1 = TestTensorBool::<2>::from_data([[false, true, true]], &device); + let tensor_2 = TestTensorBool::<2>::from_data([[true, true, false]], &device); + + let output = Tensor::cat(vec![tensor_1, tensor_2], 0); + + output.into_data().assert_eq( + &TensorData::from([[false, true, true], [true, true, false]]), + false, + ); +} + +#[test] +fn should_support_cat_with_empty_tensor_bool() { + let device = Default::default(); + let tensor_1 = TestTensorBool::<2>::from_data([[true, false, true]], &device); + let tensor_2: TestTensorBool<2> = TestTensorBool::empty([1, 0], &device); + + let output = Tensor::cat(vec![tensor_1, tensor_2], 1); + + output + .into_data() + .assert_eq(&TensorData::from([[true, false, true]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/comparison.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/comparison.rs new file mode 100644 index 0000000..1d3e49d --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/comparison.rs @@ -0,0 +1,112 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_bool_equal() { + let data_1 = TensorData::from([[false, true, true], [true, false, true]]); + let data_2 = TensorData::from([[false, false, true], [false, true, true]]); + let device = Default::default(); + let tensor_1 = TestTensorBool::<2>::from_data(data_1, &device); + let tensor_2 = TestTensorBool::<2>::from_data(data_2, &device); + + let data_actual_cloned = tensor_1.clone().equal(tensor_2.clone()); + let data_actual_inplace = tensor_1.equal(tensor_2); + + let data_expected = TensorData::from([[true, false, true], [false, false, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn should_support_bool_not_equal() { + let data_1 = TensorData::from([[false, true, true], [true, false, true]]); + let data_2 = TensorData::from([[false, false, true], [false, true, true]]); + let device = Default::default(); + let tensor_1 = TestTensorBool::<2>::from_data(data_1, &device); + let tensor_2 = TestTensorBool::<2>::from_data(data_2, &device); + + let data_actual_cloned = tensor_1.clone().not_equal(tensor_2.clone()); + let data_actual_inplace = tensor_1.not_equal(tensor_2); + + let data_expected = TensorData::from([[false, true, false], [true, true, false]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn should_support_bool_equal_broadcast() { + // [2, 2, 2] equal [1, 2, 2] broadcasts along the leading dim. + let device = Default::default(); + let tensor_1 = TestTensorBool::<3>::from_data( + TensorData::from([ + [[true, false], [false, true]], + [[true, true], [false, false]], + ]), + &device, + ); + let tensor_2 = + TestTensorBool::<2>::from_data(TensorData::from([[true, false], [false, true]]), &device); + + let actual = tensor_1.equal(tensor_2.unsqueeze::<3>()).into_data(); + let expected = TensorData::from([[[true, true], [true, true]], [[true, false], [true, false]]]); + expected.assert_eq(&actual, false); +} + +#[test] +fn should_support_bool_not_equal_broadcast() { + // [2, 2, 2] not_equal [1, 2, 2] broadcasts along the leading dim. + let device = Default::default(); + let tensor_1 = TestTensorBool::<3>::from_data( + TensorData::from([ + [[true, false], [false, true]], + [[true, true], [false, false]], + ]), + &device, + ); + let tensor_2 = + TestTensorBool::<2>::from_data(TensorData::from([[true, false], [false, true]]), &device); + + let actual = tensor_1.not_equal(tensor_2.unsqueeze::<3>()).into_data(); + let expected = TensorData::from([ + [[false, false], [false, false]], + [[false, true], [false, true]], + ]); + expected.assert_eq(&actual, false); +} + +#[test] +fn should_support_bool_not() { + let data_1 = TensorData::from([[false, true, true], [true, true, false]]); + let tensor_1 = TestTensorBool::<2>::from_data(data_1, &Default::default()); + + let data_actual_cloned = tensor_1.clone().bool_not(); + let data_actual_inplace = tensor_1.bool_not(); + + let data_expected = TensorData::from([[true, false, false], [false, false, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_bool_equal_elem() { + let tensor_1 = TestTensorBool::<2>::from([[true, false, true], [false, true, false]]); + + let data_actual_cloned = tensor_1.clone().equal_elem(false); + let data_actual_inplace = tensor_1.equal_elem(false); + + let data_expected = TensorData::from([[false, true, false], [true, false, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_bool_not_equal_elem() { + let tensor_1 = TestTensorBool::<2>::from([[true, false, true], [false, true, false]]); + + let data_actual_cloned = tensor_1.clone().not_equal_elem(true); + let data_actual_inplace = tensor_1.not_equal_elem(true); + + let data_expected = TensorData::from([[false, true, false], [true, false, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/create_like.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/create_like.rs new file mode 100644 index 0000000..accd268 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/create_like.rs @@ -0,0 +1,34 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_zeros_like() { + let tensor = TestTensorBool::<3>::from([ + [[false, true, false], [true, true, true]], + [[false, false, false], [true, true, false]], + ]); + + let tensor = tensor.zeros_like(); + let expected = TensorData::from([ + [[false, false, false], [false, false, false]], + [[false, false, false], [false, false, false]], + ]); + + tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_ones_like() { + let tensor = TestTensorBool::<3>::from([ + [[false, true, false], [true, true, true]], + [[false, false, false], [true, true, false]], + ]); + + let tensor = tensor.ones_like(); + let expected = TensorData::from([ + [[true, true, true], [true, true, true]], + [[true, true, true], [true, true, true]], + ]); + + tensor.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/expand.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/expand.rs new file mode 100644 index 0000000..c5fae07 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/expand.rs @@ -0,0 +1,30 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn expand_2d_bool() { + let tensor = TestTensorBool::<1>::from([false, true, false]); + let expanded_tensor = tensor.expand([3, 3]); + + let expected_data = TensorData::from([ + [false, true, false], + [false, true, false], + [false, true, false], + ]); + + expanded_tensor.into_data().assert_eq(&expected_data, false); +} + +#[test] +fn expand_bool_after_flip() { + // Non-palindromic input so the flip actually changes the data: + // [true, false, false] flipped -> [false, false, true]. + let tensor = TestTensorBool::<1>::from([true, false, false]).flip([0]); + + let output = tensor.expand([2, 3]); + + output.into_data().assert_eq( + &TensorData::from([[false, false, true], [false, false, true]]), + false, + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/flip.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/flip.rs new file mode 100644 index 0000000..bda4941 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/flip.rs @@ -0,0 +1,33 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn flip_bool() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device) + .reshape([2, 3, 4]) + .greater_elem(10); + + let flipped = tensor.clone().flip([0, 2]); + + // from pytorch: + // import torch; torch.arange(0, 24).reshape(2, 3, 4).flip((0, 2)).gt(10) + let data_expected = TensorData::from([ + [ + [true, true, true, true], + [true, true, true, true], + [true, true, true, true], + ], + [ + [false, false, false, false], + [false, false, false, false], + [true, false, false, false], + ], + ]); + + flipped.into_data().assert_eq(&data_expected, false); + + // Test with no flip + let flipped = tensor.clone().flip([]); + tensor.into_data().assert_eq(&flipped.into_data(), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/full.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/full.rs new file mode 100644 index 0000000..11c0b2e --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/full.rs @@ -0,0 +1,16 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_tensor_full() { + let device = Default::default(); + let bool_tensor = TestTensorBool::<2>::full([2, 2], true, &device); + bool_tensor + .into_data() + .assert_eq(&TensorData::from([[true, true], [true, true]]), false); + + let bool_tensor = TestTensorBool::<2>::full([2, 2], false, &device); + bool_tensor + .into_data() + .assert_eq(&TensorData::from([[false, false], [false, false]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/gather_scatter.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/gather_scatter.rs new file mode 100644 index 0000000..b7b6b70 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/gather_scatter.rs @@ -0,0 +1,29 @@ +use super::*; +use burn_tensor::{IndexingUpdateOp, TensorData}; + +#[test] +fn should_scatter_1d_bool() { + let device = Default::default(); + let tensor = TestTensorBool::<1>::from_data([true, false, false], &device); + let values = TestTensorBool::from_data([false, true, true], &device); + let indices = TestTensorInt::from_ints([1, 0, 2], &device); + + let output = tensor.scatter(0, indices, values, IndexingUpdateOp::Add); + + output + .into_data() + .assert_eq(&TensorData::from([true, false, true]), false); +} + +#[test] +fn should_gather_1d_dim0_bool() { + let device = Default::default(); + let tensor = TestTensorBool::<1>::from_data([true, false, false], &device); + let indices = TestTensorInt::from_ints([1, 1, 0, 1, 2], &device); + + let output = tensor.gather(0, indices); + + output + .into_data() + .assert_eq(&TensorData::from([false, false, true, false, false]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/init.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/init.rs new file mode 100644 index 0000000..d049ba9 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/init.rs @@ -0,0 +1,51 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_bool_empty() { + let shape = [2, 2]; + let tensor = TestTensorBool::<2>::empty(shape, &Default::default()); + assert_eq!(tensor.shape(), shape.into()) +} + +#[test] +fn should_support_bool_zeros() { + let shape = [2, 2]; + let tensor = TestTensorBool::<2>::zeros(shape, &Default::default()); + assert_eq!(tensor.shape(), shape.into()); + + tensor + .into_data() + .assert_eq(&TensorData::from([[false, false], [false, false]]), false); +} + +#[test] +fn should_support_bool_ones() { + let shape = [2, 2]; + let tensor = TestTensorBool::<2>::ones(shape, &Default::default()); + assert_eq!(tensor.shape(), shape.into()); + + tensor + .into_data() + .assert_eq(&TensorData::from([[true, true], [true, true]]), false); +} + +#[test] +fn should_load_bool_from_data_forcing_native_store() { + use burn_tensor::{BoolStore, DType}; + + // Model loaders (e.g. `burn-store`) reconstruct tensors with `Tensor::from_data`, + // forcing the serialized snapshot dtype. Bool tensors are serialized as + // `BoolStore::Native`, a storage layout that cubecl runtimes don't support. Forcing + // it used to panic in `bool_from_data`; `resolve_dtype` must now fall back to the + // device's default bool storage. Regression test for tracel-ai/burn#5094. + let device = Default::default(); + let data = TensorData::new(vec![true, false, true], [3]); + assert_eq!(data.dtype, DType::Bool(BoolStore::Native)); + + let tensor = TestTensorBool::<1>::from_data(data, (&device, DType::Bool(BoolStore::Native))); + + tensor + .into_data() + .assert_eq(&TensorData::from([true, false, true]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/logical.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/logical.rs new file mode 100644 index 0000000..174228f --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/logical.rs @@ -0,0 +1,336 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_bool_and() { + let tensor1 = TestTensorBool::<2>::from([[false, true, false], [true, false, true]]); + let tensor2 = TestTensorBool::<2>::from([[true, true, false], [false, false, true]]); + let data_actual = tensor1.bool_and(tensor2).into_data(); + let data_expected = TensorData::from([[false, true, false], [false, false, true]]); + data_expected.assert_eq(&data_actual, false); +} + +#[test] +fn test_bool_or() { + let tensor1 = TestTensorBool::<2>::from([[false, true, false], [true, false, true]]); + let tensor2 = TestTensorBool::<2>::from([[true, true, false], [false, false, true]]); + let data_actual = tensor1.bool_or(tensor2).into_data(); + let data_expected = TensorData::from([[true, true, false], [true, false, true]]); + data_expected.assert_eq(&data_actual, false); +} + +#[test] +fn test_bool_xor() { + let tensor1 = TestTensorBool::<2>::from([[false, true, false], [true, false, true]]); + let tensor2 = TestTensorBool::<2>::from([[true, true, false], [false, false, true]]); + let data_actual = tensor1.bool_xor(tensor2).into_data(); + let data_expected = TensorData::from([[true, false, false], [true, false, false]]); + data_expected.assert_eq(&data_actual, false); +} + +#[test] +fn test_bool_or_vec() { + let device = Default::default(); + let tensor1 = TestTensorBool::<1>::full([256], 0, &device); + let tensor2 = TestTensorBool::<1>::full([256], 1, &device); + let data_actual = tensor1.bool_or(tensor2).into_data(); + let data_expected = TensorData::from([true; 256]); + data_expected.assert_eq(&data_actual, false); +} + +#[test] +fn test_bool_and_vec() { + let device = Default::default(); + let tensor1 = TestTensorBool::<1>::full([256], 0, &device); + let tensor2 = TestTensorBool::<1>::full([256], 1, &device); + let data_actual = tensor1.bool_and(tensor2).into_data(); + let data_expected = TensorData::from([false; 256]); + data_expected.assert_eq(&data_actual, false); +} + +#[test] +fn test_bool_and_broadcast_scalar_lhs() { + // [1, 1] AND [2, 3] broadcasts to [2, 3]. + let device = Default::default(); + let tensor1 = TestTensorBool::<2>::from_data(TensorData::from([[true]]), &device); + let tensor2 = TestTensorBool::<2>::from_data( + TensorData::from([[true, false, true], [false, true, false]]), + &device, + ); + + let actual = tensor1.bool_and(tensor2).into_data(); + let expected = TensorData::from([[true, false, true], [false, true, false]]); + expected.assert_eq(&actual, false); +} + +#[test] +fn test_bool_or_broadcast_scalar_lhs() { + // [1, 1]=false OR [2, 3] broadcasts to [2, 3] and equals rhs. + let device = Default::default(); + let tensor1 = TestTensorBool::<2>::from_data(TensorData::from([[false]]), &device); + let tensor2 = TestTensorBool::<2>::from_data( + TensorData::from([[true, false, true], [false, true, false]]), + &device, + ); + + let actual = tensor1.bool_or(tensor2).into_data(); + let expected = TensorData::from([[true, false, true], [false, true, false]]); + expected.assert_eq(&actual, false); +} + +#[test] +fn test_bool_and_broadcast_rank_mismatch() { + // [2, 3, 4] AND [1, 3, 4] (unsqueezed from [3, 4]) broadcasts along dim 0. + let device = Default::default(); + let lhs = TestTensorBool::<3>::from_data( + TensorData::from([ + [ + [true, false, false, false], + [true, false, true, false], + [false, true, true, true], + ], + [ + [true, false, false, true], + [false, false, true, true], + [true, true, false, true], + ], + ]), + &device, + ); + let rhs = TestTensorBool::<2>::from_data( + TensorData::from([ + [false, false, true, true], + [false, true, true, false], + [false, false, false, true], + ]), + &device, + ); + + let actual = lhs.bool_and(rhs.unsqueeze::<3>()).into_data(); + let expected = TensorData::from([ + [ + [false, false, false, false], + [false, false, true, false], + [false, false, false, true], + ], + [ + [false, false, false, true], + [false, false, true, false], + [false, false, false, true], + ], + ]); + expected.assert_eq(&actual, false); +} + +#[test] +fn test_bool_or_broadcast_rank_mismatch() { + // [2, 3, 4] OR [1, 3, 4] (unsqueezed from [3, 4]) broadcasts along dim 0. + let device = Default::default(); + let lhs = TestTensorBool::<3>::from_data( + TensorData::from([ + [ + [true, false, false, false], + [true, false, true, false], + [false, true, true, true], + ], + [ + [true, false, false, true], + [false, false, true, true], + [true, true, false, true], + ], + ]), + &device, + ); + let rhs = TestTensorBool::<2>::from_data( + TensorData::from([ + [false, false, true, true], + [false, true, true, false], + [false, false, false, true], + ]), + &device, + ); + + let actual = lhs.bool_or(rhs.unsqueeze::<3>()).into_data(); + let expected = TensorData::from([ + [ + [true, false, true, true], + [true, true, true, false], + [false, true, true, true], + ], + [ + [true, false, true, true], + [false, true, true, true], + [true, true, false, true], + ], + ]); + expected.assert_eq(&actual, false); +} + +#[test] +fn test_bool_xor_broadcast_rank_mismatch() { + // [2, 2, 2] XOR [1, 2, 2] broadcasts along the leading dim. + let device = Default::default(); + let lhs = TestTensorBool::<3>::from_data( + TensorData::from([ + [[true, false], [false, true]], + [[true, true], [false, false]], + ]), + &device, + ); + let rhs = + TestTensorBool::<2>::from_data(TensorData::from([[true, false], [false, true]]), &device); + + let actual = lhs.bool_xor(rhs.unsqueeze::<3>()).into_data(); + let expected = TensorData::from([ + [[false, false], [false, false]], + [[false, true], [false, true]], + ]); + expected.assert_eq(&actual, false); +} + +#[test] +fn test_bool_and_broadcast_row_lhs() { + // [1, 3] AND [2, 3] broadcasts lhs along dim 0 only. Distinct from + // the scalar_lhs pattern, which broadcasts along both dims. + let device = Default::default(); + let lhs = TestTensorBool::<2>::from_data(TensorData::from([[true, false, true]]), &device); + let rhs = TestTensorBool::<2>::from_data( + TensorData::from([[true, true, false], [false, true, true]]), + &device, + ); + + let actual = lhs.bool_and(rhs).into_data(); + let expected = TensorData::from([[true, false, false], [false, false, true]]); + expected.assert_eq(&actual, false); +} + +#[test] +fn test_bool_and_broadcast_mutual() { + // [3, 1] AND [1, 3] - both operands need expansion along different + // axes, producing [3, 3]. Exercises the "neither operand matches the + // output shape" path through broadcast_binary. + let device = Default::default(); + let lhs = TestTensorBool::<2>::from_data(TensorData::from([[true], [false], [true]]), &device); + let rhs = TestTensorBool::<2>::from_data(TensorData::from([[true, false, true]]), &device); + + let actual = lhs.bool_and(rhs).into_data(); + let expected = TensorData::from([ + [true, false, true], + [false, false, false], + [true, false, true], + ]); + expected.assert_eq(&actual, false); +} + +#[test] +fn test_bool_and_broadcast_4d() { + // 4D broadcast: [2, 1, 2, 1] AND [1, 2, 1, 2] -> [2, 2, 2, 2]. + // Confirms the fix is rank-agnostic. Pre-fix, both operands had 4 + // storage elements so the inplace helper ran without panicking, but + // the output silently kept lhs's [2, 1, 2, 1] shape - shape assert + // catches that. + let device = Default::default(); + let lhs = TestTensorBool::<4>::from_data( + TensorData::from([[[[true], [false]]], [[[false], [true]]]]), + &device, + ); + let rhs = TestTensorBool::<4>::from_data( + TensorData::from([[[[true, false]], [[false, true]]]]), + &device, + ); + + let actual = lhs.bool_and(rhs).into_data(); + let expected = TensorData::from([ + [ + [[true, false], [false, false]], + [[false, true], [false, false]], + ], + [ + [[false, false], [true, false]], + [[false, false], [false, true]], + ], + ]); + expected.assert_eq(&actual, false); +} + +// Non-contiguous (negative-stride / flipped) bool inputs. Exercises stride +// handling through the logical ops on every backend. + +#[test] +fn test_bool_not_flipped() { + // [T, F, T, F] flipped -> [F, T, F, T], not -> [T, F, T, F] + let t = TestTensorBool::<1>::from([true, false, true, false]).flip([0]); + + let output = t.bool_not(); + + output + .into_data() + .assert_eq(&TensorData::from([true, false, true, false]), false); +} + +#[test] +fn test_bool_and_flipped() { + // a: [T, F, T, F] flipped -> [F, T, F, T]; b: [T, T, F, F] + // and: [F, T, F, F] + let a = TestTensorBool::<1>::from([true, false, true, false]).flip([0]); + let b = TestTensorBool::<1>::from([true, true, false, false]); + + let output = a.bool_and(b); + + output + .into_data() + .assert_eq(&TensorData::from([false, true, false, false]), false); +} + +#[test] +fn test_bool_or_flipped() { + // a: [T, F, T, F] flipped -> [F, T, F, T]; b: [T, F, F, F] + // or: [T, T, F, T] + let a = TestTensorBool::<1>::from([true, false, true, false]).flip([0]); + let b = TestTensorBool::<1>::from([true, false, false, false]); + + let output = a.bool_or(b); + + output + .into_data() + .assert_eq(&TensorData::from([true, true, false, true]), false); +} + +#[test] +fn test_bool_and_both_flipped() { + // a, b both flipped -> [F, T, F, T] & [F, F, T, T] = [F, F, F, T] + let a = TestTensorBool::<1>::from([true, false, true, false]).flip([0]); + let b = TestTensorBool::<1>::from([true, true, false, false]).flip([0]); + + let output = a.bool_and(b); + + output + .into_data() + .assert_eq(&TensorData::from([false, false, false, true]), false); +} + +#[test] +fn test_bool_xor_flipped() { + // a flipped -> [F, T, F, T]; b: [T, T, F, F]; xor: [T, F, F, T] + let a = TestTensorBool::<1>::from([true, false, true, false]).flip([0]); + let b = TestTensorBool::<1>::from([true, true, false, false]); + + let output = a.bool_xor(b); + + output + .into_data() + .assert_eq(&TensorData::from([true, false, false, true]), false); +} + +#[test] +fn test_bool_equal_flipped() { + // a flipped -> [F, T, F, T]; b: [F, T, F, T]; equal: [T, T, T, T] + let a = TestTensorBool::<1>::from([true, false, true, false]).flip([0]); + let b = TestTensorBool::<1>::from([false, true, false, true]); + + let output = a.equal(b); + + output + .into_data() + .assert_eq(&TensorData::from([true, true, true, true]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/mask.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/mask.rs new file mode 100644 index 0000000..243df0b --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/mask.rs @@ -0,0 +1,30 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_bool_mask_where_ops() { + let device = Default::default(); + let tensor = TestTensorBool::<2>::from_data([[true, false], [false, false]], &device); + let mask = + TestTensorBool::<2>::from_bool(TensorData::from([[true, false], [false, true]]), &device); + let value = + TestTensorBool::<2>::from_data(TensorData::from([[false, true], [true, false]]), &device); + + let output = tensor.mask_where(mask, value); + let expected = TensorData::from([[false, false], [false, false]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_bool_mask_fill_ops() { + let device = Default::default(); + let tensor = TestTensorBool::<2>::from_data([[false, true], [false, false]], &device); + let mask = + TestTensorBool::<2>::from_bool(TensorData::from([[true, false], [false, true]]), &device); + + let output = tensor.mask_fill(mask, true); + let expected = TensorData::from([[true, true], [false, true]]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/mod.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/mod.rs new file mode 100644 index 0000000..c33e226 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/mod.rs @@ -0,0 +1,27 @@ +pub use super::*; // re-export test types + +mod all; +mod any; +mod argwhere_nonzero; +mod cast; +mod cat; +mod comparison; +mod create_like; +mod expand; +mod flip; +mod full; +mod gather_scatter; +mod init; +mod logical; +mod mask; +mod movedim; +mod permute; +mod repeat; +mod repeat_dim; +mod reshape; +mod select; +mod stack; +mod take; +mod transpose; +mod tri_mask; +mod unfold; diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/movedim.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/movedim.rs new file mode 100644 index 0000000..fba3f4b --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/movedim.rs @@ -0,0 +1,56 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn movedim_bool() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device) + .reshape([2, 3, 4]) + .greater_elem(10); + + let permuted = tensor.clone().movedim(0, 2); + // from pytorch: + // import torch; torch.arange(0, 24).reshape(2, 3, 4).movedim(0, 2).gt(10) + let expected = TensorData::from([ + [[false, true], [false, true], [false, true], [false, true]], + [[false, true], [false, true], [false, true], [false, true]], + [[false, true], [false, true], [false, true], [true, true]], + ]); + + permuted.into_data().assert_eq(&expected, false); + + // Test with negative axis + let permuted = tensor.clone().movedim(0, -1); + permuted.into_data().assert_eq(&expected, false); + + // Test with the same axis + let permuted = tensor.clone().movedim(0, 0); + permuted.into_data().assert_eq(&tensor.into_data(), false); +} + +#[test] +fn vec_input_bool() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device) + .reshape([2, 3, 4]) + .greater_elem(10); + + let permuted = tensor.clone().movedim(vec![0, 1], vec![1, 0]); + // from pytorch + // import torch; torch.arange(0, 24).reshape(2, 3, 4).movedim([0, 1], [1, 0]).gt(10) + let expected = TensorData::from([ + [[false, false, false, false], [true, true, true, true]], + [[false, false, false, false], [true, true, true, true]], + [[false, false, false, true], [true, true, true, true]], + ]); + + permuted.into_data().assert_eq(&expected, false); + + // Test with negative axes + let permuted = tensor.clone().movedim(vec![-3, -2], vec![-2, -3]); + permuted.into_data().assert_eq(&expected, false); + + // Test with the same axes + let permuted = tensor.clone().movedim(vec![0, 1], vec![0, 1]); + permuted.into_data().assert_eq(&tensor.into_data(), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/permute.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/permute.rs new file mode 100644 index 0000000..ef095c0 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/permute.rs @@ -0,0 +1,31 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn permute_bool() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device) + .reshape([2, 3, 4]) + .greater_elem(10); + + let permuted = tensor.clone().permute([2, 1, 0]); + + // from pytorch: + // import torch; torch.arange(0, 24).reshape(2, 3, 4).permute(2, 1, 0).gt(10) + let expected = TensorData::from([ + [[false, true], [false, true], [false, true]], + [[false, true], [false, true], [false, true]], + [[false, true], [false, true], [false, true]], + [[false, true], [false, true], [true, true]], + ]); + + permuted.into_data().assert_eq(&expected, false); + + // Test with negative axis + let permuted = tensor.clone().permute([-1, 1, 0]); + permuted.into_data().assert_eq(&expected, false); + + // Test with the same axis + let permuted = tensor.clone().permute([0, 1, 2]); + permuted.into_data().assert_eq(&tensor.into_data(), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/repeat.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/repeat.rs new file mode 100644 index 0000000..ba0bd33 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/repeat.rs @@ -0,0 +1,64 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_bool_repeat_ops_one_dimension() { + let data = TensorData::from([[true, false, false]]); + let tensor = TestTensorBool::<2>::from_data(data, &Default::default()); + + let output = tensor.repeat(&[4, 1, 1]); + let expected = TensorData::from([ + [true, false, false], + [true, false, false], + [true, false, false], + [true, false, false], + ]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_bool_repeat_on_many_dimension() { + let data = TensorData::from([ + [[false, true], [true, false]], + [[true, true], [false, false]], + ]); + let tensor = TestTensorBool::<3>::from_data(data, &Default::default()); + + let output = tensor.repeat(&[2, 3, 2]); + let expected = TensorData::from([ + [ + [false, true, false, true], + [true, false, true, false], + [false, true, false, true], + [true, false, true, false], + [false, true, false, true], + [true, false, true, false], + ], + [ + [true, true, true, true], + [false, false, false, false], + [true, true, true, true], + [false, false, false, false], + [true, true, true, true], + [false, false, false, false], + ], + [ + [false, true, false, true], + [true, false, true, false], + [false, true, false, true], + [true, false, true, false], + [false, true, false, true], + [true, false, true, false], + ], + [ + [true, true, true, true], + [false, false, false, false], + [true, true, true, true], + [false, false, false, false], + [true, true, true, true], + [false, false, false, false], + ], + ]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/repeat_dim.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/repeat_dim.rs new file mode 100644 index 0000000..c3b33a8 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/repeat_dim.rs @@ -0,0 +1,34 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_bool_repeat_ops() { + let data = TensorData::from([[true, false, false]]); + let tensor = TestTensorBool::<2>::from_data(data, &Default::default()); + + let output = tensor.repeat_dim(0, 4); + let expected = TensorData::from([ + [true, false, false], + [true, false, false], + [true, false, false], + [true, false, false], + ]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_bool_repeat_on_dims_larger_than_1() { + let data = TensorData::from([ + [[false, true], [true, false]], + [[true, true], [false, false]], + ]); + let tensor = TestTensorBool::<3>::from_data(data, &Default::default()); + + let output = tensor.repeat_dim(1, 2); + let expected = TensorData::from([ + [[false, true], [true, false], [false, true], [true, false]], + [[true, true], [false, false], [true, true], [false, false]], + ]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/reshape.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/reshape.rs new file mode 100644 index 0000000..7652aef --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/reshape.rs @@ -0,0 +1,13 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_reshape_bool() { + let data = TensorData::from([false, true, false]); + let tensor = TestTensorBool::<1>::from_data(data, &Default::default()); + + let output = tensor.clone().reshape([1, 3]); + let expected = TensorData::from([[false, true, false]]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/select.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/select.rs new file mode 100644 index 0000000..05e30a0 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/select.rs @@ -0,0 +1,241 @@ +use super::*; +use burn_tensor::{IndexingUpdateOp, TensorData}; + +#[test] +fn should_select_bool_tensor_1d() { + // Test that select works for boolean tensors + let device = Default::default(); + let tensor = TestTensorBool::<1>::from_data([true, false, true], &device); + let indices = TestTensorInt::from_data([0, 2, 1, 0], &device); + + let output = tensor.select(0, indices); + let expected = TensorData::from([true, true, false, true]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_select_bool_tensor_2d() { + // Test that select works for boolean 2D tensors + let device = Default::default(); + let tensor = + TestTensorBool::<2>::from_data([[true, false, true], [false, true, false]], &device); + let indices = TestTensorInt::from_data([1, 0], &device); + + let output = tensor.select(0, indices); + let expected = TensorData::from([[false, true, false], [true, false, true]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_select_add_bool_tensor() { + // Test that select_add works for boolean tensors + let device = Default::default(); + let tensor = TestTensorBool::<1>::from_data([true, false, true], &device); + let values = TestTensorBool::<1>::from_data([false, true], &device); + let indices = TestTensorInt::from_data([0, 2], &device); + + let output = tensor.select_assign(0, indices, values, IndexingUpdateOp::Add); + // Note: select_add uses sum reduction, so: + // index 0: true OR false = true + // index 2: true OR true = true + // index 1: false (unchanged) + let expected = TensorData::from([true, false, true]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_select_add_bool_overlapping_indices() { + // Test accumulation behavior with overlapping indices + let device = Default::default(); + let tensor = TestTensorBool::<1>::from_data([false, true], &device); + let indices = TestTensorInt::from_data([0, 0], &device); + let values = TestTensorBool::<1>::from_data([true, false], &device); + + let output = tensor.select_assign(0, indices, values, IndexingUpdateOp::Add); + // Index 0: false OR true OR false = true + let expected = TensorData::from([true, true]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_select_add_bool_false_to_true_case() { + // Test false OR true = true + let device = Default::default(); + let tensor = TestTensorBool::<1>::from_data([false], &device); + let indices = TestTensorInt::from_data([0], &device); + let values = TestTensorBool::<1>::from_data([true], &device); + + let output = tensor.select_assign(0, indices, values, IndexingUpdateOp::Add); + let expected = TensorData::from([true]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_select_add_bool_true_or_true_accumulation() { + // Test multiple true accumulations + let device = Default::default(); + let tensor = TestTensorBool::<1>::from_data([true, false], &device); + let indices = TestTensorInt::from_data([0, 0, 0], &device); + let values = TestTensorBool::<1>::from_data([true, true, true], &device); + + let output = tensor.select_assign(0, indices, values, IndexingUpdateOp::Add); + let expected = TensorData::from([true, false]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_match_default_implementation_behavior() { + // Verify optimized implementation matches original default logic + let device = Default::default(); + let tensor = TestTensorBool::<1>::from_data([true, false, true], &device); + let indices = TestTensorInt::from_data([0, 1, 0], &device); + let values = TestTensorBool::<1>::from_data([false, true, true], &device); + + let optimized_result = + tensor + .clone() + .select_assign(0, indices.clone(), values.clone(), IndexingUpdateOp::Add); + + // Manual default implementation logic + let int_tensor = tensor.int(); + let int_values = values.int(); + let assigned = int_tensor.select_assign(0, indices, int_values, IndexingUpdateOp::Add); + let default_result = assigned.greater_elem(0); + + optimized_result + .into_data() + .assert_eq(&default_result.into_data(), false); +} + +#[test] +fn should_select_add_bool_overlapping_indices_vs_default() { + // Test overlapping indices against default implementation + let device = Default::default(); + let tensor = TestTensorBool::<1>::from_data([false, true], &device); + let indices = TestTensorInt::from_data([0, 0], &device); + let values = TestTensorBool::<1>::from_data([true, false], &device); + + let optimized_result = + tensor + .clone() + .select_assign(0, indices.clone(), values.clone(), IndexingUpdateOp::Add); + + let int_tensor = tensor.int(); + let int_values = values.int(); + let assigned = int_tensor.select_assign(0, indices, int_values, IndexingUpdateOp::Add); + let default_result = assigned.greater_elem(0); + + optimized_result + .into_data() + .assert_eq(&default_result.into_data(), false); +} + +#[test] +fn should_select_add_bool_true_or_true_accumulation_vs_default() { + // Test multiple true accumulations against default implementation + let device = Default::default(); + let tensor = TestTensorBool::<1>::from_data([true, false], &device); + let indices = TestTensorInt::from_data([0, 0, 0], &device); + let values = TestTensorBool::<1>::from_data([true, true, true], &device); + + let optimized_result = + tensor + .clone() + .select_assign(0, indices.clone(), values.clone(), IndexingUpdateOp::Add); + + let int_tensor = tensor.int(); + let int_values = values.int(); + let assigned = int_tensor.select_assign(0, indices, int_values, IndexingUpdateOp::Add); + let default_result = assigned.greater_elem(0); + + optimized_result + .into_data() + .assert_eq(&default_result.into_data(), false); +} + +#[test] +fn should_select_add_bool_false_to_true_case_vs_default() { + // Test false OR true case against default implementation + let device = Default::default(); + let tensor = TestTensorBool::<1>::from_data([false], &device); + let indices = TestTensorInt::from_data([0], &device); + let values = TestTensorBool::<1>::from_data([true], &device); + + let optimized_result = + tensor + .clone() + .select_assign(0, indices.clone(), values.clone(), IndexingUpdateOp::Add); + + let int_tensor = tensor.int(); + let int_values = values.int(); + let assigned = int_tensor.select_assign(0, indices, int_values, IndexingUpdateOp::Add); + let default_result = assigned.greater_elem(0); + + optimized_result + .into_data() + .assert_eq(&default_result.into_data(), false); +} + +#[test] +fn should_select_add_bool_tensor_vs_default() { + // Test existing basic case against default implementation + let device = Default::default(); + let tensor = TestTensorBool::<1>::from_data([true, false, true], &device); + let indices = TestTensorInt::from_data([0, 2], &device); + let values = TestTensorBool::<1>::from_data([false, false], &device); + + let optimized_result = + tensor + .clone() + .select_assign(0, indices.clone(), values.clone(), IndexingUpdateOp::Add); + + let int_tensor = tensor.int(); + let int_values = values.int(); + let assigned = int_tensor.select_assign(0, indices, int_values, IndexingUpdateOp::Add); + let default_result = assigned.greater_elem(0); + + optimized_result + .into_data() + .assert_eq(&default_result.into_data(), false); +} + +#[test] +#[should_panic(expected = "Tensors are not eq")] +fn should_fail_if_replacement_semantics_were_used() { + // Test that framework uses accumulation, not replacement + let device = Default::default(); + let tensor = TestTensorBool::<1>::from_data([true], &device); + let indices = TestTensorInt::from_data([0], &device); + let values = TestTensorBool::<1>::from_data([false], &device); + + let output = tensor.select_assign(0, indices, values, IndexingUpdateOp::Add); + let replacement_expected = TensorData::from([false]); + + output.into_data().assert_eq(&replacement_expected, false); +} + +#[test] +#[should_panic(expected = "Tensors are not eq")] +fn should_fail_if_replacement_semantics_were_used_vs_default() { + // Test that default implementation also uses accumulation, not replacement + let device = Default::default(); + let tensor = TestTensorBool::<1>::from_data([true], &device); + let indices = TestTensorInt::from_data([0], &device); + let values = TestTensorBool::<1>::from_data([false], &device); + + let int_tensor = tensor.int(); + let int_values = values.int(); + let assigned = int_tensor.select_assign(0, indices, int_values, IndexingUpdateOp::Add); + let default_result = assigned.greater_elem(0); + let replacement_expected = TensorData::from([false]); + + default_result + .into_data() + .assert_eq(&replacement_expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/stack.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/stack.rs new file mode 100644 index 0000000..aea179e --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/stack.rs @@ -0,0 +1,15 @@ +use super::*; +use alloc::vec; +use burn_tensor::{Tensor, TensorData}; + +#[test] +fn should_support_stack_ops_bool() { + let device = Default::default(); + let tensor_1 = TestTensorBool::<2>::from_data([[false, true, true]], &device); + let tensor_2 = TestTensorBool::<2>::from_data([[true, true, false]], &device); + + let output = Tensor::stack::<3>(vec![tensor_1, tensor_2], 0); + let expected = TensorData::from([[[false, true, true]], [[true, true, false]]]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/take.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/take.rs new file mode 100644 index 0000000..f11e29d --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/take.rs @@ -0,0 +1,41 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_take_bool_tensor() { + // Test take with boolean tensors + let device = Default::default(); + let tensor = TestTensorBool::<2>::from_data([[true, false], [false, true]], &device); + let indices = TestTensorInt::<1>::from_data([1, 0], &device); + + let output = tensor.take::<1, 2>(0, indices); + let expected = TensorData::from([[false, true], [true, false]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_take_bool_tensor_with_2d_indices() { + // Test take with boolean tensors - output will be 3D + let device = Default::default(); + let tensor = TestTensorBool::<2>::from_data( + [ + [true, false, true], + [false, true, false], + [true, true, false], + ], + &device, + ); + + // 2D indices - shape [2, 2] + let indices = TestTensorInt::<2>::from_data([[0, 2], [1, 0]], &device); + let output = tensor.take::<2, 3>(0, indices); + + // Expected: shape [2, 2, 3] + let expected = TensorData::from([ + [[true, false, true], [true, true, false]], + [[false, true, false], [true, false, true]], + ]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/transpose.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/transpose.rs new file mode 100644 index 0000000..51d44bc --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/transpose.rs @@ -0,0 +1,41 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_transpose_bool() { + let tensor = TestTensorBool::<3>::from_data( + [ + [[false, true, false], [false, false, false]], + [[false, false, true], [false, false, true]], + ], + &Default::default(), + ); + + let output = tensor.transpose(); + let expected = TensorData::from([ + [[false, false], [true, false], [false, false]], + [[false, false], [false, false], [true, true]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_swap_dims_bool() { + let tensor = TestTensorBool::<3>::from_data( + [ + [[false, true, false], [false, false, false]], + [[false, false, true], [false, false, true]], + ], + &Default::default(), + ); + + let output = tensor.swap_dims(0, 2); + let expected = TensorData::from([ + [[false, false], [false, false]], + [[true, false], [false, false]], + [[false, true], [false, true]], + ]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/tri_mask.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/tri_mask.rs new file mode 100644 index 0000000..dfe3fa0 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/tri_mask.rs @@ -0,0 +1,94 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn square_diag() { + let device = Default::default(); + let data_expected = TensorData::from([ + [false, true, true], + [true, false, true], + [true, true, false], + ]); + let tensor = TestTensorBool::<2>::diag_mask([3, 3], 0, &device); + tensor.into_data().assert_eq(&data_expected, false); +} + +#[test] +fn square_diag_offset() { + let device = Default::default(); + let data_expected = + TensorData::from([[true, false, true], [true, true, false], [true, true, true]]); + let tensor = TestTensorBool::<2>::diag_mask([3, 3], 1, &device); + tensor.into_data().assert_eq(&data_expected, false); +} + +#[test] +fn square_tri_upper() { + let device = Default::default(); + let data_expected = TensorData::from([ + [false, false, false], + [true, false, false], + [true, true, false], + ]); + let tensor = TestTensorBool::<2>::triu_mask([3, 3], 0, &device); + tensor.into_data().assert_eq(&data_expected, false); +} + +#[test] +fn square_tri_upper_offset() { + let device = Default::default(); + let data_expected = TensorData::from([ + [true, false, false], + [true, true, false], + [true, true, true], + ]); + let tensor = TestTensorBool::<2>::triu_mask([3, 3], 1, &device); + tensor.into_data().assert_eq(&data_expected, false); +} + +#[test] +fn square_tri_lower() { + let device = Default::default(); + + let data_expected = TensorData::from([ + [false, true, true], + [false, false, true], + [false, false, false], + ]); + let tensor = TestTensorBool::<2>::tril_mask([3, 3], 0, &device); + tensor.into_data().assert_eq(&data_expected, false); +} + +#[test] +fn square_tri_lower_offset() { + let device = Default::default(); + + let data_expected = TensorData::from([ + [true, true, true], + [false, true, true], + [false, false, true], + ]); + let tensor = TestTensorBool::<2>::tril_mask([3, 3], -1, &device); + tensor.into_data().assert_eq(&data_expected, false); +} + +#[test] +fn rect_diag() { + let device = Default::default(); + let data_expected = TensorData::from([ + [false, true, true, true], + [true, false, true, true], + [true, true, false, true], + ]); + let tensor = TestTensorBool::<2>::diag_mask([3, 4], 0, &device); + tensor.into_data().assert_eq(&data_expected, false); + + let data_expected = TensorData::from([ + [false, true, true], + [true, false, true], + [true, true, false], + [true, true, true], + ]); + let tensor = TestTensorBool::<2>::diag_mask([4, 3], 0, &device); + tensor.into_data().assert_eq(&data_expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/bool/ops/unfold.rs b/crates/burn-backend-tests/tests/tensor/bool/ops/unfold.rs new file mode 100644 index 0000000..fb02534 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/bool/ops/unfold.rs @@ -0,0 +1,36 @@ +use super::*; +use burn_tensor::Distribution; +use burn_tensor::s; + +#[test] +fn test_unfold_bool() { + let device = Default::default(); + + let input = + TestTensor::<3>::random([2, 6, 6], Distribution::Default, &device).greater_elem(0.5); + + let dim = 1; + let size = 3; + let step = 2; + let actual: TestTensorBool<4> = input.clone().unfold(dim, size, step); + + let expected = TestTensorBool::<4>::empty([2, 2, 6, 3], &device) + .slice_assign( + s![.., 0, .., ..], + input + .clone() + .slice(s![.., 0..3, ..]) + .swap_dims(1, 2) + .unsqueeze_dim::<4>(1), + ) + .slice_assign( + s![.., 1, .., ..], + input + .clone() + .slice(s![.., 2..5, ..]) + .swap_dims(1, 2) + .unsqueeze_dim::<4>(1), + ); + + actual.to_data().assert_eq(&expected.to_data(), true); +} diff --git a/crates/burn-backend-tests/tests/tensor/clone_invariance.rs b/crates/burn-backend-tests/tests/tensor/clone_invariance.rs new file mode 100644 index 0000000..806f960 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/clone_invariance.rs @@ -0,0 +1,761 @@ +/// This module tests whether basic tensor operations remain invariant when performed on clones, +/// meaning that cloning input tensors won't affect the results. +/// +/// Those are relevant tests because backends may employ unsafe optimizations to reuse tensor data +/// and use different kernels in such cases. We ensure that the results are consistent regardless +/// of the approach and that the input tensors are not modified when cloned. +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::activation::{ + gelu, log_sigmoid, log_softmax, mish, relu, sigmoid, silu, softmax, softplus, tanh, +}; +use burn_tensor::{Distribution, IndexingUpdateOp, TensorData}; + +pub trait CloneInvarianceTest { + type Args; + + fn args(&self) -> Self::Args; + + fn run(&self, args: &Self::Args, inplace: bool) -> TensorData; + + fn check(&self) { + let args = self.args(); + let out = self.run(&args, false); + let out_inplace = self.run(&args, true); + + out.assert_approx_eq::(&out_inplace, Tolerance::default()); + } +} + +macro_rules! clone_invariance_test { + (unary: $name:ident, ops_float: $ops:expr) => { + #[test] + #[allow(non_snake_case)] + fn $name() { + struct $name; + + impl CloneInvarianceTest<2> for $name { + type Args = TensorData; + + fn args(&self) -> Self::Args { + TestTensor::<2>::random([32, 32], Distribution::Default, &Default::default()) + .into_data() + .convert::() + } + + fn run(&self, args: &Self::Args, inplace: bool) -> TensorData { + let lhs = TestTensor::from_data(args.clone(), &Default::default()); + + if inplace { + $ops(lhs).into_data().convert::() + } else { + let out = $ops(lhs.clone()).into_data().convert::(); + lhs.into_data() + .assert_approx_eq::(args, Tolerance::default()); + out + } + } + } + + CloneInvarianceTest::<2>::check(&$name); + } + }; + + (binary: $name:ident, ops_float: $ops:expr) => { + #[test] + #[allow(non_snake_case)] + fn $name() { + struct $name; + + impl CloneInvarianceTest<2> for $name { + type Args = (TensorData, TensorData); + + fn args(&self) -> Self::Args { + let device = Default::default(); + ( + TestTensor::<2>::ones([32, 32], &device) + .into_data() + .convert::(), + // Avoid div by zero. + TestTensor::<2>::ones([32, 32], &device) + .into_data() + .convert::(), + ) + } + + fn run(&self, (lhs_arg, rhs_arg): &Self::Args, inplace: bool) -> TensorData { + let device = Default::default(); + let lhs = TestTensor::from_data(lhs_arg.clone(), &device); + let rhs = TestTensor::from_data(rhs_arg.clone(), &device); + + if inplace { + $ops(lhs, rhs).into_data().convert::() + } else { + let out = $ops(lhs.clone(), rhs.clone()).into_data().convert::(); + + lhs.into_data() + .assert_approx_eq::(lhs_arg, Tolerance::default()); + rhs.into_data() + .assert_approx_eq::(rhs_arg, Tolerance::default()); + + out + } + } + } + + CloneInvarianceTest::<2>::check(&$name); + } + }; + + (unary: $name:ident, ops_int: $ops:expr) => { + #[test] + #[allow(non_snake_case)] + fn $name() { + struct $name; + + impl CloneInvarianceTest<2> for $name { + type Args = TensorData; + + fn args(&self) -> Self::Args { + TestTensor::<2>::random( + [32, 32], + Distribution::Uniform(0.0, 50.0), + &Default::default(), + ) + .into_data() + .convert::() + } + + fn run(&self, args: &Self::Args, inplace: bool) -> TensorData { + let lhs = TestTensorInt::from_data(args.clone(), &Default::default()); + + if inplace { + $ops(lhs).into_data().convert::() + } else { + let out = $ops(lhs.clone()).into_data().convert::(); + lhs.into_data() + .convert::() + .assert_approx_eq::(args, Tolerance::default()); + out + } + } + } + + CloneInvarianceTest::<2>::check(&$name); + } + }; + + (binary: $name:ident, ops_int: $ops:expr) => { + #[test] + #[allow(non_snake_case)] + fn $name() { + struct $name; + + impl CloneInvarianceTest<2> for $name { + type Args = (TensorData, TensorData); + + fn args(&self) -> Self::Args { + let device = Default::default(); + ( + TestTensor::<2>::random([32, 32], Distribution::Uniform(0., 50.), &device) + .into_data() + .convert::(), + // Avoid div by zero. + TestTensor::<2>::random([32, 32], Distribution::Uniform(1., 51.), &device) + .into_data() + .convert::(), + ) + } + + fn run(&self, (lhs_arg, rhs_arg): &Self::Args, inplace: bool) -> TensorData { + let device = Default::default(); + let lhs = TestTensorInt::from_data(lhs_arg.clone(), &device); + let rhs = TestTensorInt::from_data(rhs_arg.clone(), &device); + + if inplace { + $ops(lhs, rhs).into_data().convert::() + } else { + let out = $ops(lhs.clone(), rhs.clone()).into_data().convert::(); + + lhs.into_data() + .convert::() + .assert_approx_eq::(lhs_arg, Tolerance::default()); + rhs.into_data() + .convert::() + .assert_approx_eq::(rhs_arg, Tolerance::default()); + + out + } + } + } + + CloneInvarianceTest::<2>::check(&$name); + } + }; +} + +mod float { + use super::*; + + // Unary ops + clone_invariance_test!( + unary: AddScalar, + ops_float: |tensor: TestTensor<2>| tensor.add_scalar(2.0) + ); + clone_invariance_test!( + unary: SubScalar, + ops_float: |tensor: TestTensor<2>| tensor.sub_scalar(2.0) + ); + clone_invariance_test!( + unary: DivScalar, + ops_float: |tensor: TestTensor<2>| tensor.div_scalar(2.0) + ); + clone_invariance_test!( + unary: MulScalar, + ops_float: |tensor: TestTensor<2>| tensor.mul_scalar(2.0) + ); + clone_invariance_test!( + unary: PowScalar, + ops_float: |tensor: TestTensor<2>| tensor.powf_scalar(2.0) + ); + clone_invariance_test!( + unary: Square, + ops_float: |tensor: TestTensor<2>| tensor.square() + ); + clone_invariance_test!( + unary: Sqrt, + ops_float: |tensor: TestTensor<2>| tensor.sqrt() + ); + clone_invariance_test!( + unary: Exp, + ops_float: |tensor: TestTensor<2>| tensor.exp() + ); + clone_invariance_test!( + unary: Neg, + ops_float: |tensor: TestTensor<2>| tensor.neg() + ); + clone_invariance_test!( + unary: MeanDim, + ops_float: |tensor: TestTensor<2>| tensor.mean_dim(1) + ); + clone_invariance_test!( + unary: SumDim, + ops_float: |tensor: TestTensor<2>| tensor.sum_dim(1) + ); + clone_invariance_test!( + unary: Sum, + ops_float: |tensor: TestTensor<2>| tensor.sum().unsqueeze::<2>() + ); + clone_invariance_test!( + unary: Mean, + ops_float: |tensor: TestTensor<2>| tensor.mean().unsqueeze::<2>() + ); + clone_invariance_test!( + unary: Clamp, + ops_float: |tensor: TestTensor<2>| tensor.clamp(-2., 2.) + ); + clone_invariance_test!( + unary: ClampMin, + ops_float: |tensor: TestTensor<2>| tensor.clamp_min(-2.) + ); + clone_invariance_test!( + unary: ClampMax, + ops_float: |tensor: TestTensor<2>| tensor.clamp_max(2.) + ); + clone_invariance_test!( + unary: Abs, + ops_float: |tensor: TestTensor<2>| tensor.abs() + ); + clone_invariance_test!( + unary: Cos, + ops_float: |tensor: TestTensor<2>| tensor.cos() + ); + clone_invariance_test!( + unary: Sin, + ops_float: |tensor: TestTensor<2>| tensor.sin() + ); + clone_invariance_test!( + unary: Tan, + ops_float: |tensor: TestTensor<2>| tensor.tan() + ); + clone_invariance_test!( + unary: Log, + ops_float: |tensor: TestTensor<2>| tensor.log() + ); + clone_invariance_test!( + unary: Log1P, + ops_float: |tensor: TestTensor<2>| tensor.log1p() + ); + clone_invariance_test!( + unary: SwapDims, + ops_float: |tensor: TestTensor<2>| tensor.swap_dims(0, 1) + ); + clone_invariance_test!( + unary: Transpose, + ops_float: |tensor: TestTensor<2>| tensor.transpose() + ); + clone_invariance_test!( + unary: Slice, + ops_float: |tensor: TestTensor<2>| tensor.slice([0..12, 12..24]) + ); + clone_invariance_test!( + unary: Erf, + ops_float: |tensor: TestTensor<2>| tensor.erf() + ); + clone_invariance_test!( + unary: EqualElem, + ops_float: |tensor: TestTensor<2>| tensor.equal_elem(0.5) + ); + clone_invariance_test!( + unary: NotEqualElem, + ops_float: |tensor: TestTensor<2>| tensor.not_equal_elem(0.5) + ); + clone_invariance_test!( + unary: GreaterElem, + ops_float: |tensor: TestTensor<2>| tensor.greater_elem(0.5) + ); + clone_invariance_test!( + unary: GreaterEqualElem, + ops_float: |tensor: TestTensor<2>| tensor.greater_equal_elem(0.5) + ); + clone_invariance_test!( + unary: LowerElem, + ops_float: |tensor: TestTensor<2>| tensor.lower_elem(0.5) + ); + clone_invariance_test!( + unary: LowerEqualElem, + ops_float: |tensor: TestTensor<2>| tensor.lower_equal_elem(0.5) + ); + clone_invariance_test!( + unary: Argmax, + ops_float: |tensor: TestTensor<2>| tensor.argmax(0) + ); + clone_invariance_test!( + unary: Argmin, + ops_float: |tensor: TestTensor<2>| tensor.argmin(0) + ); + clone_invariance_test!( + unary: Max, + ops_float: |tensor: TestTensor<2>| tensor.max().unsqueeze::<2>() + ); + clone_invariance_test!( + unary: Min, + ops_float: |tensor: TestTensor<2>| tensor.min().unsqueeze::<2>() + ); + clone_invariance_test!( + unary: MaxDim, + ops_float: |tensor: TestTensor<2>| tensor.max_dim(1) + ); + clone_invariance_test!( + unary: MaxDimWithIndices, + ops_float: |tensor: TestTensor<2>| tensor.max_dim_with_indices(1).0 + ); + clone_invariance_test!( + unary: MinDimWithIndices, + ops_float: |tensor: TestTensor<2>| tensor.min_dim_with_indices(1).0 + ); + clone_invariance_test!( + unary: MinDim, + ops_float: |tensor: TestTensor<2>| tensor.min_dim(1) + ); + clone_invariance_test!( + unary: Repeat, + ops_float: |tensor: TestTensor<2>| { + tensor.reshape([1, 32, 32]).repeat_dim(0, 4).reshape([4 * 32, 32]) + } + ); + clone_invariance_test!( + unary: Reshape, + ops_float: |tensor: TestTensor<2>| { + let shape = tensor.shape(); + let new_shape = [shape.num_elements(), 1]; + tensor.reshape(new_shape) + } + ); + clone_invariance_test!( + unary: Gatter, + ops_float: |tensor: TestTensor<2>| { + let shape = tensor.shape(); + let indices = TestTensorInt::ones(shape, &Default::default()); + tensor.gather(0, indices) + } + ); + clone_invariance_test!( + unary: Select, + ops_float: |tensor: TestTensor<2>| { + let indices = TestTensorInt::from_ints([1, 2, 0, 5], &Default::default()); + tensor.select(0, indices) + } + ); + clone_invariance_test!( + unary: MaskFill, + ops_float: |tensor: TestTensor<2>| { + let mask = tensor.clone().greater_elem(0.5); + tensor.mask_fill(mask, 77.0) + } + ); + + // Activation + clone_invariance_test!( + unary: Softmax, + ops_float: |tensor: TestTensor<2>| softmax(tensor, 1) + ); + clone_invariance_test!( + unary: LogSoftmax, + ops_float: |tensor: TestTensor<2>| log_softmax(tensor, 1) + ); + clone_invariance_test!( + unary: Sigmoid, + ops_float: |tensor: TestTensor<2>| sigmoid(tensor) + ); + clone_invariance_test!( + unary: LogSigmoid, + ops_float: |tensor: TestTensor<2>| log_sigmoid(tensor) + ); + clone_invariance_test!( + unary: Relu, + ops_float: |tensor: TestTensor<2>| relu(tensor) + ); + clone_invariance_test!( + unary: Gelu, + ops_float: |tensor: TestTensor<2>| gelu(tensor) + ); + clone_invariance_test!( + unary: Mish, + ops_float: |tensor: TestTensor<2>| mish(tensor) + ); + clone_invariance_test!( + unary: Silu, + ops_float: |tensor: TestTensor<2>| silu(tensor) + ); + clone_invariance_test!( + unary: Softplus, + ops_float: |tensor: TestTensor<2>| softplus(tensor, 1.0) + ); + clone_invariance_test!( + unary: Tanh, + ops_float: |tensor: TestTensor<2>| tanh(tensor) + ); + + // Binary ops + clone_invariance_test!( + binary: Add, + ops_float: |lhs: TestTensor<2>, rhs: TestTensor<2>| lhs.add(rhs) + ); + clone_invariance_test!( + binary: Sub, + ops_float: |lhs: TestTensor<2>, rhs: TestTensor<2>| lhs.sub(rhs) + ); + clone_invariance_test!( + binary: Div, + ops_float: |lhs: TestTensor<2>, rhs: TestTensor<2>| lhs.div(rhs) + ); + clone_invariance_test!( + binary: Mul, + ops_float: |lhs: TestTensor<2>, rhs: TestTensor<2>| lhs.mul(rhs) + ); + clone_invariance_test!( + binary: Matmul, + ops_float: |lhs: TestTensor<2>, rhs: TestTensor<2>| lhs.matmul(rhs) + ); + clone_invariance_test!( + binary: Equal, + ops_float: |lhs: TestTensor<2>, rhs: TestTensor<2>| lhs.equal(rhs) + ); + clone_invariance_test!( + binary: Greater, + ops_float: |lhs: TestTensor<2>, rhs: TestTensor<2>| lhs.greater(rhs) + ); + clone_invariance_test!( + binary: GreaterEqual, + ops_float: |lhs: TestTensor<2>, rhs: TestTensor<2>| lhs.greater_equal(rhs) + ); + clone_invariance_test!( + binary: Lower, + ops_float: |lhs: TestTensor<2>, rhs: TestTensor<2>| lhs.lower(rhs) + ); + clone_invariance_test!( + binary: LowerEqual, + ops_float: |lhs: TestTensor<2>, rhs: TestTensor<2>| lhs.lower_equal(rhs) + ); + clone_invariance_test!( + binary: Cat, + ops_float: |lhs: TestTensor<2>, rhs: TestTensor<2>| { + let lhs = lhs.reshape([1usize, 32, 32]); + let rhs = rhs.reshape([1usize, 32, 32]); + + TestTensor::cat(vec![lhs, rhs], 0).reshape([64, 32]) + } + ); + clone_invariance_test!( + binary: Scatter, + ops_float: |tensor: TestTensor<2>, values: TestTensor<2>| { + let shape = tensor.shape(); + let indices = TestTensorInt::ones(shape, &Default::default()); + tensor.scatter(0, indices, values, IndexingUpdateOp::Add) + } + ); + clone_invariance_test!( + binary: SliceAssign, + ops_float: |tensor: TestTensor<2>, values: TestTensor<2>| { + tensor.slice_assign([0..12, 12..24], values.slice([12..24, 0..12])) + } + ); + clone_invariance_test!( + binary: MaskWhere, + ops_float: |tensor: TestTensor<2>, values: TestTensor<2>| { + let mask = tensor.clone().greater_elem(0.5); + tensor.mask_where(mask, values) + } + ); + clone_invariance_test!( + binary: SelectAssign, + ops_float: |tensor: TestTensor<2>, values: TestTensor<2>| { + let indices = TestTensorInt::from_ints([1, 2, 0, 5], &Default::default()); + let values = values.select(0, indices.clone()); + tensor.select_assign(0, indices, values, IndexingUpdateOp::Add) + } + ); +} + +mod int { + use super::*; + + // Unary ops + clone_invariance_test!( + unary: AddScalar, + ops_int: |tensor: TestTensorInt<2>| tensor.add_scalar(2.0) + ); + clone_invariance_test!( + unary: SubScalar, + ops_int: |tensor: TestTensorInt<2>| tensor.sub_scalar(2.0) + ); + clone_invariance_test!( + unary: DivScalar, + ops_int: |tensor: TestTensorInt<2>| tensor.div_scalar(2.0) + ); + clone_invariance_test!( + unary: MulScalar, + ops_int: |tensor: TestTensorInt<2>| tensor.mul_scalar(2.0) + ); + clone_invariance_test!( + unary: Neg, + ops_int: |tensor: TestTensorInt<2>| tensor.neg() + ); + clone_invariance_test!( + unary: MeanDim, + ops_int: |tensor: TestTensorInt<2>| tensor.mean_dim(1) + ); + clone_invariance_test!( + unary: SumDim, + ops_int: |tensor: TestTensorInt<2>| tensor.sum_dim(1) + ); + clone_invariance_test!( + unary: Sum, + ops_int: |tensor: TestTensorInt<2>| tensor.sum().unsqueeze::<2>() + ); + clone_invariance_test!( + unary: Mean, + ops_int: |tensor: TestTensorInt<2>| tensor.mean().unsqueeze::<2>() + ); + clone_invariance_test!( + unary: Clamp, + ops_int: |tensor: TestTensorInt<2>| tensor.clamp(-2., 2.) + ); + clone_invariance_test!( + unary: ClampMin, + ops_int: |tensor: TestTensorInt<2>| tensor.clamp_min(-2.) + ); + clone_invariance_test!( + unary: ClampMax, + ops_int: |tensor: TestTensorInt<2>| tensor.clamp_max(2.) + ); + clone_invariance_test!( + unary: Abs, + ops_int: |tensor: TestTensorInt<2>| tensor.abs() + ); + clone_invariance_test!( + unary: SwapDims, + ops_int: |tensor: TestTensorInt<2>| tensor.swap_dims(0, 1) + ); + clone_invariance_test!( + unary: Transpose, + ops_int: |tensor: TestTensorInt<2>| tensor.transpose() + ); + clone_invariance_test!( + unary: Slice, + ops_int: |tensor: TestTensorInt<2>| tensor.slice([0..12, 12..24]) + ); + clone_invariance_test!( + unary: EqualElem, + ops_int: |tensor: TestTensorInt<2>| tensor.equal_elem(25) + ); + clone_invariance_test!( + unary: NotEqualElem, + ops_int: |tensor: TestTensorInt<2>| tensor.not_equal_elem(25) + ); + clone_invariance_test!( + unary: GreaterElem, + ops_int: |tensor: TestTensorInt<2>| tensor.greater_elem(25) + ); + clone_invariance_test!( + unary: GreaterEqualElem, + ops_int: |tensor: TestTensorInt<2>| tensor.greater_equal_elem(25) + ); + clone_invariance_test!( + unary: LowerElem, + ops_int: |tensor: TestTensorInt<2>| tensor.lower_elem(25) + ); + clone_invariance_test!( + unary: LowerEqualElem, + ops_int: |tensor: TestTensorInt<2>| tensor.lower_equal_elem(25) + ); + clone_invariance_test!( + unary: Argmax, + ops_int: |tensor: TestTensorInt<2>| tensor.argmax(0) + ); + clone_invariance_test!( + unary: Argmin, + ops_int: |tensor: TestTensorInt<2>| tensor.argmin(0) + ); + clone_invariance_test!( + unary: Max, + ops_int: |tensor: TestTensorInt<2>| tensor.max().unsqueeze::<2>() + ); + clone_invariance_test!( + unary: Min, + ops_int: |tensor: TestTensorInt<2>| tensor.min().unsqueeze::<2>() + ); + clone_invariance_test!( + unary: MaxDim, + ops_int: |tensor: TestTensorInt<2>| tensor.max_dim(1) + ); + clone_invariance_test!( + unary: MaxDimWithIndices, + ops_int: |tensor: TestTensorInt<2>| tensor.max_dim_with_indices(1).0 + ); + clone_invariance_test!( + unary: MinDimWithIndices, + ops_int: |tensor: TestTensorInt<2>| tensor.min_dim_with_indices(1).0 + ); + clone_invariance_test!( + unary: MinDim, + ops_int: |tensor: TestTensorInt<2>| tensor.min_dim(1) + ); + clone_invariance_test!( + unary: Repeat, + ops_int: |tensor: TestTensorInt<2>| { + tensor.reshape([1, 32, 32]).repeat_dim(0, 4).reshape([4 * 32, 32]) + } + ); + clone_invariance_test!( + unary: Reshape, + ops_int: |tensor: TestTensorInt<2>| { + let shape = tensor.shape(); + let new_shape = [shape.num_elements(), 1]; + tensor.reshape(new_shape) + } + ); + clone_invariance_test!( + unary: Gatter, + ops_int: |tensor: TestTensorInt<2>| { + let shape = tensor.shape(); + let indices = TestTensorInt::ones(shape, &Default::default()); + tensor.gather(0, indices) + } + ); + clone_invariance_test!( + unary: Select, + ops_int: |tensor: TestTensorInt<2>| { + let indices = TestTensorInt::from_ints([1, 2, 0, 5], &Default::default()); + tensor.select(0, indices) + } + ); + clone_invariance_test!( + unary: MaskFill, + ops_int: |tensor: TestTensorInt<2>| { + let mask = tensor.clone().greater_elem(0.5); + tensor.mask_fill(mask, 77.0) + } + ); + + // Binary ops + clone_invariance_test!( + binary: Add, + ops_int: |lhs: TestTensorInt<2>, rhs: TestTensorInt<2>| lhs.add(rhs) + ); + clone_invariance_test!( + binary: Sub, + ops_int: |lhs: TestTensorInt<2>, rhs: TestTensorInt<2>| lhs.sub(rhs) + ); + clone_invariance_test!( + binary: Div, + ops_int: |lhs: TestTensorInt<2>, rhs: TestTensorInt<2>| lhs.div(rhs) + ); + clone_invariance_test!( + binary: Mul, + ops_int: |lhs: TestTensorInt<2>, rhs: TestTensorInt<2>| lhs.mul(rhs) + ); + clone_invariance_test!( + binary: Equal, + ops_int: |lhs: TestTensorInt<2>, rhs: TestTensorInt<2>| lhs.equal(rhs) + ); + clone_invariance_test!( + binary: NotEqual, + ops_int: |lhs: TestTensorInt<2>, rhs: TestTensorInt<2>| lhs.not_equal(rhs) + ); + clone_invariance_test!( + binary: Greater, + ops_int: |lhs: TestTensorInt<2>, rhs: TestTensorInt<2>| lhs.greater(rhs) + ); + clone_invariance_test!( + binary: GreaterEqual, + ops_int: |lhs: TestTensorInt<2>, rhs: TestTensorInt<2>| lhs.greater_equal(rhs) + ); + clone_invariance_test!( + binary: Lower, + ops_int: |lhs: TestTensorInt<2>, rhs: TestTensorInt<2>| lhs.lower(rhs) + ); + clone_invariance_test!( + binary: LowerEqual, + ops_int: |lhs: TestTensorInt<2>, rhs: TestTensorInt<2>| lhs.lower_equal(rhs) + ); + clone_invariance_test!( + binary: Cat, + ops_int: |lhs: TestTensorInt<2>, rhs: TestTensorInt<2>| { + let lhs = lhs.reshape([1usize, 32, 32]); + let rhs = rhs.reshape([1usize, 32, 32]); + + TestTensorInt::cat(vec![lhs, rhs], 0).reshape([64, 32]) + } + ); + clone_invariance_test!( + binary: Scatter, + ops_int: |tensor: TestTensorInt<2>, values: TestTensorInt<2>| { + let shape = tensor.shape(); + let indices = TestTensorInt::ones(shape, &Default::default()); + tensor.scatter(0, indices, values, IndexingUpdateOp::Add) + } + ); + clone_invariance_test!( + binary: SliceAssign, + ops_int: |tensor: TestTensorInt<2>, values: TestTensorInt<2>| { + tensor.slice_assign([0..12, 12..24], values.slice([12..24, 0..12])) + } + ); + clone_invariance_test!( + binary: MaskWhere, + ops_int: |tensor: TestTensorInt<2>, values: TestTensorInt<2>| { + let mask = tensor.clone().greater_elem(0.5); + tensor.mask_where(mask, values) + } + ); + clone_invariance_test!( + binary: SelectAssign, + ops_int: |tensor: TestTensorInt<2>, values: TestTensorInt<2>| { + let indices = TestTensorInt::from_ints([1, 2, 0, 5], &Default::default()); + let values = values.select(0, indices.clone()); + tensor.select_assign(0, indices, values, IndexingUpdateOp::Add) + } + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/distributed.rs b/crates/burn-backend-tests/tests/tensor/distributed.rs new file mode 100644 index 0000000..be7318a --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/distributed.rs @@ -0,0 +1,118 @@ +use burn_tensor::Tolerance; +use burn_tensor::distributed::ReduceOperation; +use burn_tensor::module::all_reduce; +use burn_tensor::{Device, DeviceType, TensorData}; +use rand::RngExt; +use serial_test::serial; + +use super::*; + +#[test] +#[serial] +fn test_all_reduce() { + let devices = Device::enumerate(DeviceType::Cuda).autodiff().into_vec(); + let shape = [20, 20]; + run_all_reduce(devices, 100, shape); +} + +#[test] +#[serial] +fn test_all_reduce_multithread() { + let devices = Device::enumerate(DeviceType::Cuda).autodiff().into_vec(); + let shape = [20, 20]; + run_multithread(devices, 100, shape); +} + +fn run_all_reduce(devices: Vec, num_iterations: usize, shape: [usize; 2]) { + let mut rng = rand::rng(); + let size = shape[0] * shape[1]; + + for _ in 0..num_iterations { + let vec_data: Vec> = (0..devices.len()) + .map(|_| (0..size).map(|_| rng.random_range(0.0..10.0)).collect()) + .collect(); + let expected: Vec = (0..size) + .map(|i| vec_data.iter().map(|v| v[i]).sum::()) + .collect(); + let tensors: Vec> = vec_data + .iter() + .zip(devices.clone()) + .map(|(data, device)| { + let tensor_data = TensorData::from(data.as_slice()); + let tensor = Tensor::<1>::from_data(tensor_data, &device); + tensor.reshape(shape) + }) + .collect(); + + let mut out_tensors = vec![]; + for tensor in tensors.clone() { + let output = all_reduce(tensor, ReduceOperation::Sum, devices.clone()); + let output = output.resolve(); + out_tensors.push(output); + } + + for tensor in out_tensors { + let data = tensor.flatten::<1>(0, 1).to_data(); + data.assert_approx_eq::( + &TensorData::from(expected.as_slice()), + Tolerance::default(), + ); + } + } +} + +fn run_multithread(devices: Vec, num_iterations: usize, shape: [usize; 2]) { + let size = shape[0] * shape[1]; + let num_devices = devices.len(); + + let (expected_sender, expected_receiver) = std::sync::mpsc::channel(); + let (actual_sender, actual_receiver) = std::sync::mpsc::channel(); + let handles = devices + .iter() + .map(|device| { + let local_device = device.clone(); + let local_devices = devices.clone(); + let local_expected_sender = expected_sender.clone(); + let local_actual_sender = actual_sender.clone(); + std::thread::spawn(move || { + for _ in 0..num_iterations { + let mut rng = rand::rng(); + let vec_data: Vec = + (0..size).map(|_| rng.random_range(0.0..10.0)).collect(); + local_expected_sender.send(vec_data.clone()).unwrap(); + + let tensor_data = TensorData::from(vec_data.as_slice()); + let tensor = + Tensor::<1>::from_data(tensor_data, &local_device).reshape(shape.clone()); + let output = all_reduce(tensor, ReduceOperation::Sum, local_devices.clone()); + let output = output.resolve(); + + let data = output.flatten::<1>(0, 1).to_data(); + local_actual_sender.send(data).unwrap(); + } + }) + }) + .collect::>(); + + for _ in 0..num_iterations { + let mut expected_list = vec![]; + for _ in 0..num_devices { + expected_list.push(expected_receiver.recv().unwrap()); + } + let expected: Vec = (0..size) + .map(|i| expected_list.iter().map(|v| v[i]).sum::()) + .collect(); + + for _ in 0..num_devices { + let data = actual_receiver.recv().unwrap(); + data.assert_approx_eq::( + &TensorData::from(expected.as_slice()), + Tolerance::default(), + ); + } + } + + for h in handles { + h.join().unwrap(); + } +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/celu.rs b/crates/burn-backend-tests/tests/tensor/float/activation/celu.rs new file mode 100644 index 0000000..49b4812 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/celu.rs @@ -0,0 +1,34 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_celu_d2() { + let tensor = TestTensor::<2>::from([[1.0, 7.0], [-3.0, 0.5]]); + + let output = activation::celu(tensor, 1.0); + // celu(1, 1) = 1 + // celu(7, 1) = 7 + // celu(-3, 1) = 1 * (exp(-3) - 1) = -0.950213 + // celu(0.5, 1) = 0.5 + let expected = TensorData::from([[1.0, 7.0], [-0.950213, 0.5]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_celu_with_alpha() { + let tensor = TestTensor::<1>::from([0.0, -1.0, -2.0]); + + let output = activation::celu(tensor, 2.0); + // celu(0, 2) = 0 + // celu(-1, 2) = 2 * (exp(-0.5) - 1) = -0.786939 + // celu(-2, 2) = 2 * (exp(-1) - 1) = -1.264241 + let expected = TensorData::from([0.0, -0.786939, -1.264241]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/elu.rs b/crates/burn-backend-tests/tests/tensor/float/activation/elu.rs new file mode 100644 index 0000000..f38cc07 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/elu.rs @@ -0,0 +1,32 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_elu() { + let tensor = TestTensor::<2>::from([[1.0, 7.0], [13.0, -3.0]]); + + let output = activation::elu(tensor, 1.0); + // elu(1, 1) = 1, elu(7, 1) = 7, elu(13, 1) = 13 + // elu(-3, 1) = 1 * (exp(-3) - 1) = -0.950213 + let expected = TensorData::from([[1.0, 7.0], [13.0, -0.950213]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_elu_alpha() { + let tensor = TestTensor::<1>::from([0.0, -1.0, -2.0]); + + let output = activation::elu(tensor, 2.0); + // elu(0, 2) = 2*(exp(0)-1) = 0 + // elu(-1, 2) = 2*(exp(-1)-1) = 2*(-0.632121) = -1.264241 + // elu(-2, 2) = 2*(exp(-2)-1) = 2*(-0.864665) = -1.729329 + let expected = TensorData::from([0.0, -1.264241, -1.729329]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/gelu.rs b/crates/burn-backend-tests/tests/tensor/float/activation/gelu.rs new file mode 100644 index 0000000..3de0403 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/gelu.rs @@ -0,0 +1,32 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_gelu() { + let tensor = TestTensor::<2>::from([[ + 0.5447, 0.9809, 0.4114, 0.1398, 0.8045, 0.4103, 0.2388, 0.5262, 0.6677, 0.6737, + ]]); + let output = activation::gelu(tensor); + let expected = TensorData::from([[ + 0.3851, 0.8207, 0.2714, 0.0777, 0.6351, 0.2704, 0.1419, 0.3687, 0.4993, 0.5051, + ]]); + + // Low precision to allow approximation implementation using tanh + output.into_data().assert_approx_eq::( + &expected, + Tolerance::default().set_half_precision_absolute(2e-3), + ); +} + +#[test] +fn test_gelu_d1() { + let tensor = TestTensor::<1>::from([-3.0, 0.0, 3.0]); + + let output = activation::gelu(tensor); + + output.into_data().assert_approx_eq::( + &TensorData::from([0.0, 0.0, 3.0]), + Tolerance::absolute(0.01), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/glu.rs b/crates/burn-backend-tests/tests/tensor/float/activation/glu.rs new file mode 100644 index 0000000..f8e46fa --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/glu.rs @@ -0,0 +1,28 @@ +use super::*; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_glu_d3() { + let tensor = TestTensor::<3>::from([[ + [ + -0.5710, -1.3416, 1.9128, -0.8257, -0.1331, -1.4804, -0.6281, -0.6115, + ], + [ + 0.0267, -1.3834, 0.2752, 0.7844, -0.3549, -0.4274, 0.3290, -0.5459, + ], + [ + -1.6347, -2.0908, 1.8801, 0.3541, 0.2237, 1.0377, 2.4850, 0.3490, + ], + ]]); + + let output = activation::glu(tensor, 2); + + output.into_data().assert_approx_eq::( + &TensorData::from([[ + [-0.2665, -0.2487, 0.6656, -0.2904], + [0.0110, -0.5461, 0.1601, 0.2877], + [-0.9084, -1.5439, 1.7355, 0.2077], + ]]), + Default::default(), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/hard_sigmoid.rs b/crates/burn-backend-tests/tests/tensor/float/activation/hard_sigmoid.rs new file mode 100644 index 0000000..c1af47d --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/hard_sigmoid.rs @@ -0,0 +1,27 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_hard_sigmoid() { + let tensor = TestTensor::<2>::from([[1.0, 7.0], [13.0, -3.0]]); + + let output = activation::hard_sigmoid(tensor, 0.2, 0.5); + let expected = TensorData::from([[0.7, 1.0], [1.0, 0.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_hard_sigmoid_overflow() { + let tensor = TestTensor::<1>::from([FloatElem::MAX, FloatElem::MIN]); + + let output = activation::hard_sigmoid(tensor, 0.2, 0.5); + let expected = TensorData::from([1.0, 0.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/leaky_relu.rs b/crates/burn-backend-tests/tests/tensor/float/activation/leaky_relu.rs new file mode 100644 index 0000000..80c6e3b --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/leaky_relu.rs @@ -0,0 +1,28 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_leaky_relu_d2() { + let tensor = TestTensor::<2>::from([[0.0, -1.0, 2.0], [3.0, -4.0, 5.0]]); + + let output = activation::leaky_relu(tensor, 0.01); + + // Account for conversion errors if `FloatType != f32` + output.into_data().assert_approx_eq::( + &TensorData::from([[0.0, -0.01, 2.0], [3.0, -0.04, 5.0]]), + Tolerance::default(), + ); +} + +#[test] +fn test_leaky_relu_d1() { + let tensor = TestTensor::<1>::from([-2.0, -1.0, 0.0, 1.0, 2.0]); + + let output = activation::leaky_relu(tensor, 0.01); + + output.into_data().assert_approx_eq::( + &TensorData::from([-0.02, -0.01, 0.0, 1.0, 2.0]), + Tolerance::absolute(1e-6), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/log_sigmoid.rs b/crates/burn-backend-tests/tests/tensor/float/activation/log_sigmoid.rs new file mode 100644 index 0000000..a834f88 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/log_sigmoid.rs @@ -0,0 +1,51 @@ +#![allow(clippy::approx_constant)] + +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{ElementConversion, TensorData, activation}; + +#[test] +fn test_log_sigmoid() { + let tensor = TestTensor::<2>::from([[1.0, 7.0], [13.0, -3.0]]); + + let output = activation::log_sigmoid(tensor); + let expected = TensorData::from([[-3.132617e-1, -9.114665e-4], [-2.260327e-6, -3.0485873]]); + + let tolerance = Tolerance::rel_abs(0.01, 0.0001); + output + .into_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn test_log_sigmoid_d1() { + let tensor = TestTensor::<1>::from([-10.0, 0.0, 10.0]); + + let output = activation::log_sigmoid(tensor); + + output.into_data().assert_approx_eq::( + &TensorData::from([-10.0, -0.6931472, 0.0]), + Tolerance::absolute(1e-3), + ); +} + +#[test] +fn test_log_sigmoid_numerical_stability() { + let tensor = TestTensor::<1>::from([300.0, -300.0]); + + let output = activation::log_sigmoid(tensor); + + // For large negative values, the previous implementation −log(1 + exp(−x)) would give -inf + let expected = TensorData::from([0.0, -300.0]); + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let tensor = TestTensor::<1>::from([FloatElem::MAX, FloatElem::MIN]); + let output = activation::log_sigmoid(tensor); + let expected = TensorData::from([0.elem(), FloatElem::MIN]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/log_softmax.rs b/crates/burn-backend-tests/tests/tensor/float/activation/log_softmax.rs new file mode 100644 index 0000000..979a671 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/log_softmax.rs @@ -0,0 +1,16 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_log_softmax_d2() { + let tensor = TestTensor::<2>::from([[1.0, 0.0], [0.0, 1.0]]); + + let output = activation::log_softmax(tensor, 1); + let expected = TensorData::from([[-0.3132617, -1.3132617], [-1.3132617, -0.3132617]]); + + let tolerance = Tolerance::rel_abs(0.01, 0.0001); + output + .into_data() + .assert_approx_eq::(&expected, tolerance); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/mish.rs b/crates/burn-backend-tests/tests/tensor/float/activation/mish.rs new file mode 100644 index 0000000..2cb2f20 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/mish.rs @@ -0,0 +1,21 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_mish() { + let tensor = TestTensor::<2>::from([[-0.4240, -0.9574, -0.2215], [-0.5767, 0.7218, -0.1620]]); + + let output = activation::mish(tensor); + let expected = TensorData::from([ + [-0.19709, -0.30056, -0.11714], + [-0.24132, 0.58235, -0.08877], + ]); + + // Metal has less precise trigonometric functions (tanh inside mish) + let tolerance = Tolerance::default().set_half_precision_relative(1e-2); + + output + .into_data() + .assert_approx_eq::(&expected, tolerance); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/mod.rs b/crates/burn-backend-tests/tests/tensor/float/activation/mod.rs new file mode 100644 index 0000000..7ffadb8 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/mod.rs @@ -0,0 +1,23 @@ +use super::*; + +mod celu; +mod elu; +mod gelu; +mod glu; +mod hard_sigmoid; +mod leaky_relu; +mod log_sigmoid; +mod log_softmax; +mod mish; +mod prelu; +mod quiet_softmax; +mod relu; +mod selu; +mod sigmoid; +mod silu; +mod softmax; +mod softmin; +mod softplus; +mod softsign; +mod tanh_activation; +mod thresholded_relu; diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/prelu.rs b/crates/burn-backend-tests/tests/tensor/float/activation/prelu.rs new file mode 100644 index 0000000..f4d1c35 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/prelu.rs @@ -0,0 +1,101 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_prelu_2_dimension() { + let data = [ + [-1.1, 0.0, 1.2, 0.25, -5.4], + [-4.567, 0.56, -1.55, 99.9, 0.0], + ]; + let tensor = TestTensor::<2>::from(data); + let output = activation::prelu(tensor, TestTensor::from([0.5, 0.25, 0.0, -0.8, -0.4])); + let expected = TensorData::from([ + [-0.5500, 0.0000, 1.2000, 0.2500, 2.1600], + [-2.2835, 0.5600, -0.0000, 99.9000, -0.0000], + ]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} +#[test] +fn test_prelu_2_dimension_scalar_weight() { + let data = [ + [-1.1, 0.0, 1.2, 0.25, -5.4], + [-4.567, 0.56, -1.55, 99.9, 0.0], + ]; + let tensor = TestTensor::<2>::from(data); + let output = activation::prelu(tensor, TestTensor::from([-0.8])); + let expected = TensorData::from([ + [0.8800, -0.0000, 1.2000, 0.2500, 4.3200], + [3.6536, 0.5600, 1.2400, 99.9000, -0.0000], + ]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_prelu_positives() { + // Check that positives are untouched + let data = [[ + 0.5447, 0.9809, 0.4114, 0.1398, 0.8045, 0.4103, 0.2388, 0.5262, 0.6677, 0.6737, + ]]; + let tensor = TestTensor::<2>::from(data); + let output = activation::prelu(tensor, TestTensor::from([0.25])); + let expected = TensorData::from(data); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_prelu_zero_weight() { + // test that with weight 0 it behaves as relu + let data = [-1.1, 0.0, 1.2, 0.25, -5.4]; + let tensor = TestTensor::<1>::from(data); + let output = activation::prelu(tensor, TestTensor::from([0.0])); + let expected = TensorData::from([0.0, 0.0, 1.2, 0.25, 0.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_prelu_some_weight() { + // test that with some non zero weight it works like leaky relu + let data = [-1.1, 0.0, 1.2, 0.25, -5.4]; + let tensor = TestTensor::<1>::from(data); + let output = activation::prelu(tensor, TestTensor::from([0.5])); + let expected = TensorData::from([-0.550, 0.0, 1.20, 0.250, -2.70]); + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +#[should_panic] +fn test_prelu_single_dim_multi_weight() { + // should panic because the data has only 1 channel + let data = [-1.1, 2.0, 1.2, 0.25, -5.4]; + let tensor = TestTensor::<1>::from(data); + let data_actual = + activation::prelu(tensor, TestTensor::from([0.5, -0.25, 0.0, 0.5, -1.0])).into_data(); + let data_expected = TensorData::from([-0.550, 0.0, 1.20, 0.250, -2.70]); + data_expected.assert_approx_eq::(&data_actual, Tolerance::default()); +} + +#[test] +#[should_panic] +fn test_prelu_multi_dim_wrong_weights() { + let data = [ + [-1.1, 0.0, 1.2, 0.25, -5.4], + [-4.567, 0.56, -1.55, 99.9, 0.0], + ]; + let tensor = TestTensor::<2>::from(data); + let _ = activation::prelu(tensor, TestTensor::from([-0.8, 0.1])); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/quiet_softmax.rs b/crates/burn-backend-tests/tests/tensor/float/activation/quiet_softmax.rs new file mode 100644 index 0000000..32e0385 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/quiet_softmax.rs @@ -0,0 +1,15 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_quiet_softmax_d2() { + let tensor = TestTensor::<2>::from([[1.0, 7.0], [13.0, -3.0]]); + + let output = activation::quiet_softmax(tensor, 1); + let expected = TensorData::from([[2.47e-03, 9.975e-01], [1.0, 1.1254e-07]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/relu.rs b/crates/burn-backend-tests/tests/tensor/float/activation/relu.rs new file mode 100644 index 0000000..8cdb707 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/relu.rs @@ -0,0 +1,25 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance, activation}; + +#[test] +fn test_relu_d2() { + let tensor = TestTensor::<2>::from([[0.0, -1.0, 2.0], [3.0, -4.0, 5.0]]); + + let output = activation::relu(tensor); + + output + .into_data() + .assert_eq(&TensorData::from([[0.0, 0.0, 2.0], [3.0, 0.0, 5.0]]), false); +} + +#[test] +fn test_relu_d1() { + let tensor = TestTensor::<1>::from([-2.0, -1.0, 0.0, 1.0, 2.0]); + + let output = activation::relu(tensor); + + output.into_data().assert_approx_eq::( + &TensorData::from([0.0, 0.0, 0.0, 1.0, 2.0]), + Tolerance::absolute(1e-6), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/selu.rs b/crates/burn-backend-tests/tests/tensor/float/activation/selu.rs new file mode 100644 index 0000000..a17109a --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/selu.rs @@ -0,0 +1,37 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_selu() { + // selu(x) = gamma * x if x > 0, gamma * alpha * (exp(x) - 1) if x <= 0 + // alpha = 1.6733, gamma = 1.0507 + let tensor = TestTensor::<2>::from([[0.0, 1.0, -1.0], [2.0, -2.0, 0.5]]); + + let output = activation::selu(tensor); + + // Expected values computed from the formula: + // selu(0.0) = 1.0507 * 1.6733 * (exp(0) - 1) = 0.0 + // selu(1.0) = 1.0507 * 1.0 = 1.0507 + // selu(-1.0) = 1.0507 * 1.6733 * (exp(-1) - 1) = 1.7581 * (0.3679 - 1) = -1.1113 + // selu(2.0) = 1.0507 * 2.0 = 2.1014 + // selu(-2.0) = 1.0507 * 1.6733 * (exp(-2) - 1) = 1.7581 * (0.1353 - 1) = -1.5202 + // selu(0.5) = 1.0507 * 0.5 = 0.5254 + let expected = TensorData::from([[0.0, 1.0507, -1.1113], [2.1014, -1.5202, 0.5254]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_selu_zero() { + let tensor = TestTensor::<1>::from([0.0]); + + let output = activation::selu(tensor); + let expected = TensorData::from([0.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/sigmoid.rs b/crates/burn-backend-tests/tests/tensor/float/activation/sigmoid.rs new file mode 100644 index 0000000..aaf95ae --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/sigmoid.rs @@ -0,0 +1,39 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_sigmoid() { + let tensor = TestTensor::<2>::from([[1.0, 7.0], [13.0, -3.0]]); + + let output = activation::sigmoid(tensor); + let expected = TensorData::from([[0.731059, 0.999089], [0.999998, 0.047426]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_sigmoid_d1() { + let tensor = TestTensor::<1>::from([-10.0, 0.0, 10.0]); + + let output = activation::sigmoid(tensor); + + output.into_data().assert_approx_eq::( + &TensorData::from([0.0, 0.5, 1.0]), + Tolerance::absolute(1e-3), + ); +} + +#[test] +fn test_sigmoid_overflow() { + let tensor = TestTensor::<1>::from([FloatElem::MAX, FloatElem::MIN]); + + let output = activation::sigmoid(tensor); + let expected = TensorData::from([1.0, 0.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/silu.rs b/crates/burn-backend-tests/tests/tensor/float/activation/silu.rs new file mode 100644 index 0000000..da62121 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/silu.rs @@ -0,0 +1,15 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_silu() { + let tensor = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]); + + let output = activation::silu(tensor); + let expected = TensorData::from([[0.73106, 1.76159], [2.85772, 3.92806]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/softmax.rs b/crates/burn-backend-tests/tests/tensor/float/activation/softmax.rs new file mode 100644 index 0000000..6f57e09 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/softmax.rs @@ -0,0 +1,129 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_softmax_d2() { + let tensor = TestTensor::<2>::from([[1.0, 7.0], [13.0, -3.0]]); + + let output = activation::softmax(tensor, 1); + let expected = TensorData::from([[2.472623e-03, 9.975274e-01], [1.0, 1.125352e-07]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_softmax_d1() { + let tensor = TestTensor::<1>::from([1.0, 2.0, 3.0]); + + let output = activation::softmax(tensor, 0); + + output.into_data().assert_approx_eq::( + &TensorData::from([0.09003, 0.24473, 0.66524]), + Tolerance::default().set_half_precision_absolute(2e-3), + ); +} + +#[test] +fn test_softmax_d2_varied() { + let tensor = TestTensor::<2>::from([[-1.0, 0.0, 1.0, 2.0], [0.5, 0.5, 0.5, 0.5]]); + + let output = activation::softmax(tensor, 1); + + output.into_data().assert_approx_eq::( + &TensorData::from([ + [0.03205860, 0.08714432, 0.23688284, 0.64391422], + [0.25, 0.25, 0.25, 0.25], + ]), + Tolerance::default().set_half_precision_absolute(2e-3), + ); +} + +#[test] +fn test_softmax_d3_last_axis() { + let tensor = TestTensor::<3>::from([ + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + [[0.0, 0.0, 1.0], [1.0, 1.0, 1.0]], + ]); + + let output = activation::softmax(tensor, 2); + + output.into_data().assert_approx_eq::( + &TensorData::from([ + [ + [0.09003057, 0.24472848, 0.66524094], + [0.09003057, 0.24472848, 0.66524094], + ], + [ + [0.21194156, 0.21194156, 0.57611688], + [0.33333334, 0.33333334, 0.33333334], + ], + ]), + Tolerance::default().set_half_precision_absolute(2e-3), + ); +} + +#[test] +fn test_softmax_non_contiguous_input() { + // Softmax on a transposed (non-contiguous) input. The [3, 4] tensor is + // transposed to [4, 3] before softmax over the last axis, exercising + // stride-aware handling of the op in every backend. + // + // Every transposed row has the pattern [a, a+4, a+8] (columns of the + // original tensor); softmax is shift-invariant, so every row shares + // the same result. + let t = TestTensor::<2>::from([ + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + ]); + let t_transposed = t.transpose(); + + let output = activation::softmax(t_transposed, 1); + + let row = [3.2932044e-4, 1.7980287e-2, 0.98169035]; + let expected = TensorData::from([row, row, row, row]); + output.into_data().assert_approx_eq::( + &expected, + Tolerance::default().set_half_precision_absolute(2e-3), + ); +} + +#[test] +fn test_softmax_non_last_axis_d3() { + let tensor = TestTensor::<3>::from([ + [[1.0, -2.0, 0.5], [3.0, 0.0, -1.0]], + [[0.1, 2.5, -0.3], [1.2, -0.7, 2.1]], + ]); + + let output = activation::softmax(tensor.clone(), 1); + + // Reference: softmax along middle axis computed via exp / sum pattern. + let exp = tensor.exp(); + let sum = exp.clone().sum_dim(1); + let expected = exp / sum; + + output + .into_data() + .assert_approx_eq::(&expected.into_data(), Tolerance::default()); +} + +#[test] +fn test_softmax_non_last_axis_d4_dim0() { + let tensor = TestTensor::<4>::from([ + [[[1.0, -0.5], [0.3, 2.1]], [[-1.2, 0.0], [0.8, 1.5]]], + [[[0.4, -1.1], [2.0, 0.2]], [[-0.3, 0.9], [1.1, -0.7]]], + ]); + + let output = activation::softmax(tensor.clone(), 0); + + let exp = tensor.exp(); + let sum = exp.clone().sum_dim(0); + let expected = exp / sum; + + output + .into_data() + .assert_approx_eq::(&expected.into_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/softmin.rs b/crates/burn-backend-tests/tests/tensor/float/activation/softmin.rs new file mode 100644 index 0000000..e513c45 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/softmin.rs @@ -0,0 +1,15 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_softmin_d2() { + let tensor = TestTensor::<2>::from([[1.0, 7.0], [13.0, -3.0]]); + + let output = activation::softmin(tensor, 1); + let expected = TensorData::from([[9.975274e-01, 2.472623e-03], [1.125352e-07, 1.0000]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/softplus.rs b/crates/burn-backend-tests/tests/tensor/float/activation/softplus.rs new file mode 100644 index 0000000..a0911e2 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/softplus.rs @@ -0,0 +1,28 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_softplus_d2() { + let tensor = TestTensor::<2>::from([[-0.4240, -0.9574, -0.2215], [-0.5767, 0.7218, -0.1620]]); + + let output = activation::softplus(tensor.clone(), 1.0); + let expected = TensorData::from([ + [0.503453, 0.324898, 0.588517], + [0.445806, 1.117805, 0.615424], + ]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let output = activation::softplus(tensor, 2.0); + let expected = TensorData::from([ + [0.178232, 0.068737, 0.247990], + [0.137132, 0.827771, 0.272106], + ]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/softsign.rs b/crates/burn-backend-tests/tests/tensor/float/activation/softsign.rs new file mode 100644 index 0000000..a4c1334 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/softsign.rs @@ -0,0 +1,27 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_softsign() { + let tensor = TestTensor::<2>::from([[1.0, 7.0], [13.0, -3.0]]); + + let output = activation::softsign(tensor); + let expected = TensorData::from([[0.5, 0.875], [0.928571, -0.75]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_softsign_zero() { + let tensor = TestTensor::<1>::from([0.0]); + + let output = activation::softsign(tensor); + let expected = TensorData::from([0.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/tanh_activation.rs b/crates/burn-backend-tests/tests/tensor/float/activation/tanh_activation.rs new file mode 100644 index 0000000..2015672 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/tanh_activation.rs @@ -0,0 +1,15 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_tanh() { + let tensor = TestTensor::<2>::from([[1., 2.], [3., 4.]]); + + let output = activation::tanh(tensor); + let expected = TensorData::from([[0.761594, 0.964028], [0.995055, 0.999329]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/activation/thresholded_relu.rs b/crates/burn-backend-tests/tests/tensor/float/activation/thresholded_relu.rs new file mode 100644 index 0000000..d1be8eb --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/activation/thresholded_relu.rs @@ -0,0 +1,26 @@ +use super::*; +use burn_tensor::{TensorData, activation}; + +#[test] +fn test_thresholded_relu_d2() { + // alpha = 1.0 (ONNX default): x if x > 1.0, else 0 + let tensor = TestTensor::<2>::from([[0.0, -1.0, 2.0], [3.0, 1.0, 0.5]]); + + let output = activation::thresholded_relu(tensor, 1.0); + + output + .into_data() + .assert_eq(&TensorData::from([[0.0, 0.0, 2.0], [3.0, 0.0, 0.0]]), false); +} + +#[test] +fn test_thresholded_relu_d2_alpha() { + // alpha = 0.5: x if x > 0.5, else 0 + let tensor = TestTensor::<2>::from([[0.0, -1.0, 2.0], [3.0, 0.5, 0.6]]); + + let output = activation::thresholded_relu(tensor, 0.5); + + output + .into_data() + .assert_eq(&TensorData::from([[0.0, 0.0, 2.0], [3.0, 0.0, 0.6]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/grid/affine_grid.rs b/crates/burn-backend-tests/tests/tensor/float/grid/affine_grid.rs new file mode 100644 index 0000000..9a67b35 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/grid/affine_grid.rs @@ -0,0 +1,82 @@ +use super::*; +use burn_tensor::grid::affine_grid_2d; + +fn create_identity_transform(batch_size: usize) -> TestTensor<3> { + // Identity affine transform (batch_size, 2, 3) + TestTensor::<3>::from([[[1f32, 0., 0.], [0., 1., 0.]]]).expand([batch_size, 2, 3]) +} + +#[test] +fn test_affine_grid_identity() { + let batch_size = 1; + let channels = 1; + let height = 2; + let width = 2; + + let transform = create_identity_transform(batch_size); + + let output = affine_grid_2d(transform, [batch_size, channels, height, width]); + + // Expected normalized coords: + // [-1, -1], [ 1,-1] + // [-1, 1], [ 1, 1] + let expected = TestTensor::<4>::from([[ + [[-1f32, -1f32], [1f32, -1f32]], + [[-1f32, 1f32], [1f32, 1f32]], + ]]); + + output.into_data().assert_eq(&expected.into_data(), false); +} + +#[test] +fn test_affine_grid_scaling() { + let batch_size = 1; + let channels = 1; + let height = 2; + let width = 2; + + let scale = 2.0f32; + let transform = TestTensor::<3>::from([[[scale, 0., 0.], [0., scale, 0.]]]); + + let output = affine_grid_2d(transform, [batch_size, channels, height, width]); + + // Expect scaled coordinates from normalized grid, so coords * 2 + let expected = TestTensor::<4>::from([[ + [[-2f32, -2f32], [2f32, -2f32]], + [[-2f32, 2f32], [2f32, 2f32]], + ]]); + + output.into_data().assert_eq(&expected.into_data(), false); +} + +#[test] +fn test_affine_grid_translation() { + let batch_size = 1; + let channels = 1; + let height = 2; + let width = 2; + + // Translate by 0.5 in x and -0.5 in y (normalized coords) + let tx = 0.5f32; + let ty = -0.5f32; + + let transform = TestTensor::<3>::from([[[1.0, 0.0, tx], [0.0, 1.0, ty]]]); + + let output = affine_grid_2d(transform, [batch_size, channels, height, width]); + + // Expected coordinates: + // Original normalized coords are [-1,1] in x and y + // After translation, each coordinate shifts by tx and ty + // So points become: + // [-1 + 0.5, -1 - 0.5] = [-0.5, -1.5] + // [ 1 + 0.5, -1 - 0.5] = [1.5, -1.5] + // [-1 + 0.5, 1 - 0.5] = [-0.5, 0.5] + // [ 1 + 0.5, 1 - 0.5] = [1.5, 0.5] + + let expected = TestTensor::<4>::from([[ + [[-0.5f32, -1.5f32], [1.5f32, -1.5f32]], + [[-0.5f32, 0.5f32], [1.5f32, 0.5f32]], + ]]); + + output.into_data().assert_eq(&expected.into_data(), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/grid/meshgrid.rs b/crates/burn-backend-tests/tests/tensor/float/grid/meshgrid.rs new file mode 100644 index 0000000..65de990 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/grid/meshgrid.rs @@ -0,0 +1,150 @@ +use super::*; +use burn_tensor::Tensor; +use burn_tensor::TensorData; +use burn_tensor::grid::{ + GridIndexing, GridOptions, GridSparsity, IndexPos, meshgrid, meshgrid_stack, +}; +use burn_tensor::kind::Basic; + +fn assert_tensors_equal(actual: &[Tensor; N], expected: &[Tensor; N]) +where + K: Basic, +{ + for (a, e) in actual.iter().zip(expected.iter()) { + a.clone() + .into_data() + .assert_eq(&e.clone().into_data(), true); + } +} + +#[test] +fn test_meshgrid() { + let x = TestTensor::<1>::from([1, 2, 3, 4]); + let y = TestTensor::<1>::from([5, 6]); + let z = TestTensor::<1>::from([7, 8]); + + let grid_shape = [x.dims()[0], y.dims()[0], z.dims()[0]]; + + // 3D, Dense, Matrix + assert_tensors_equal( + &meshgrid(&[x.clone(), y.clone(), z.clone()], GridOptions::default()), + &[ + x.clone().reshape([4, 1, 1]).expand(grid_shape), + y.clone().reshape([1, 2, 1]).expand(grid_shape), + z.clone().reshape([1, 1, 2]).expand(grid_shape), + ], + ); + assert_tensors_equal( + &meshgrid(&[x.clone(), y.clone(), z.clone()], GridSparsity::Dense), + &[ + x.clone().reshape([4, 1, 1]).expand(grid_shape), + y.clone().reshape([1, 2, 1]).expand(grid_shape), + z.clone().reshape([1, 1, 2]).expand(grid_shape), + ], + ); + assert_tensors_equal( + &meshgrid(&[x.clone(), y.clone(), z.clone()], GridIndexing::Matrix), + &[ + x.clone().reshape([4, 1, 1]).expand(grid_shape), + y.clone().reshape([1, 2, 1]).expand(grid_shape), + z.clone().reshape([1, 1, 2]).expand(grid_shape), + ], + ); + + // 3D, Sparse, Matrix + assert_tensors_equal( + &meshgrid( + &[x.clone(), y.clone(), z.clone()], + GridOptions { + indexing: GridIndexing::Matrix, + sparsity: GridSparsity::Sparse, + }, + ), + &[ + x.clone().reshape([4, 1, 1]), + y.clone().reshape([1, 2, 1]), + z.clone().reshape([1, 1, 2]), + ], + ); + assert_tensors_equal( + &meshgrid(&[x.clone(), y.clone(), z.clone()], GridSparsity::Sparse), + &[ + x.clone().reshape([4, 1, 1]), + y.clone().reshape([1, 2, 1]), + z.clone().reshape([1, 1, 2]), + ], + ); + + // 3D, Dense, Cartesian + assert_tensors_equal( + &meshgrid(&[x.clone(), y.clone(), z.clone()], GridIndexing::Cartesian), + &[ + x.clone() + .reshape([4, 1, 1]) + .expand(grid_shape) + .swap_dims(0, 1), + y.clone() + .reshape([1, 2, 1]) + .expand(grid_shape) + .swap_dims(0, 1), + z.clone() + .reshape([1, 1, 2]) + .expand(grid_shape) + .swap_dims(0, 1), + ], + ); + + // 3D, Sparse, Cartesian + assert_tensors_equal( + &meshgrid( + &[x.clone(), y.clone(), z.clone()], + GridOptions::new(GridIndexing::Cartesian, GridSparsity::Sparse), + ), + &[ + x.clone().reshape([4, 1, 1]).swap_dims(0, 1), + y.clone().reshape([1, 2, 1]).swap_dims(0, 1), + z.clone().reshape([1, 1, 2]).swap_dims(0, 1), + ], + ); + assert_tensors_equal( + &meshgrid( + &[x.clone(), y.clone(), z.clone()], + GridOptions { + indexing: GridIndexing::Cartesian, + sparsity: GridSparsity::Sparse, + }, + ), + &[ + x.clone().reshape([4, 1, 1]).swap_dims(0, 1), + y.clone().reshape([1, 2, 1]).swap_dims(0, 1), + z.clone().reshape([1, 1, 2]).swap_dims(0, 1), + ], + ); +} + +#[test] +fn test_meshgrid_stack() { + let tensors = [ + TestTensor::from([0.5, 1.0, 2.5]), + TestTensor::from([0.5, 1.0]), + ]; + + let result: Tensor<3> = meshgrid_stack(&tensors, IndexPos::First); + result.to_data().assert_eq( + &TensorData::from([ + [[0.5, 0.5], [1.0, 1.0], [2.5, 2.5]], + [[0.5, 1.0], [0.5, 1.0], [0.5, 1.0]], + ]), + false, + ); + + let result: Tensor<3> = meshgrid_stack(&tensors, IndexPos::Last); + result.to_data().assert_eq( + &TensorData::from([ + [[0.5, 0.5], [0.5, 1.0]], + [[1.0, 0.5], [1.0, 1.0]], + [[2.5, 0.5], [2.5, 1.0]], + ]), + false, + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/grid/mod.rs b/crates/burn-backend-tests/tests/tensor/float/grid/mod.rs new file mode 100644 index 0000000..e2f874c --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/grid/mod.rs @@ -0,0 +1,4 @@ +use super::*; + +pub(crate) mod affine_grid; +pub(crate) mod meshgrid; diff --git a/crates/burn-backend-tests/tests/tensor/float/linalg/cosine_similarity.rs b/crates/burn-backend-tests/tests/tensor/float/linalg/cosine_similarity.rs new file mode 100644 index 0000000..644148a --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/linalg/cosine_similarity.rs @@ -0,0 +1,100 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{TensorData, linalg}; + +#[test] +fn test_cosine_similarity_basic() { + // Create test tensors + let x1 = TestTensor::<2>::from([[1.0, 2.0, 3.0], [0.5, 1.5, 2.5]]); + let x2 = TestTensor::<2>::from([[1.5, 2.5, 3.5], [0.7, 1.7, 2.7]]); + + // Test cosine similarity along dimension 1 + let expected = TensorData::from([[0.99983203], [0.99987257]]); + linalg::cosine_similarity(x1.clone(), x2.clone(), 1, None) + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + // Test with explicit epsilon + linalg::cosine_similarity(x1.clone(), x2.clone(), 1, Some(1e-8)) + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_cosine_similarity_orthogonal() { + // Create orthogonal vectors + let x1 = TestTensor::<2>::from([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]); + let x2 = TestTensor::<2>::from([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]); + + // Orthogonal vectors should have cosine similarity of 0 + let expected = TensorData::from([[0.0], [0.0]]); + linalg::cosine_similarity(x1, x2, 1, None) + .into_data() + .assert_eq(&expected, false); +} + +#[test] +fn test_cosine_similarity_parallel() { + // Create parallel vectors + let x1 = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let x2 = TestTensor::<2>::from([[2.0, 4.0, 6.0], [8.0, 10.0, 12.0]]); + + // Parallel vectors should have cosine similarity of 1 + let expected = TensorData::from([[1.0], [1.0]]); + linalg::cosine_similarity(x1, x2, 1, None) + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_cosine_similarity_opposite() { + // Create opposite direction vectors + let x1 = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let x2 = TestTensor::<2>::from([[-1.0, -2.0, -3.0], [-4.0, -5.0, -6.0]]); + + // Opposite vectors should have cosine similarity of -1 + let expected = TensorData::from([[-1.0], [-1.0]]); + linalg::cosine_similarity(x1, x2, 1, None) + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_cosine_similarity_different_dimension() { + // Test with a 3D tensor + let x1 = TestTensor::<3>::from([[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]); + let x2 = TestTensor::<3>::from([[[2.0, 3.0], [4.0, 5.0]], [[6.0, 7.0], [8.0, 9.0]]]); + + // Test along dimension 2 + let expected = TensorData::from([[[0.9959688], [0.9958376]], [[0.9955946], [0.9955169]]]); + + // sensitive to rounding in dot/norm; loosen f16 tolerance + let tolerance = Tolerance::default().set_half_precision_relative(7e-3); + + linalg::cosine_similarity(x1.clone(), x2.clone(), 2, None) + .into_data() + .assert_approx_eq::(&expected, tolerance); + + // Test with negative dimension (-1 is the last dimension, which is 2 in this case) + linalg::cosine_similarity(x1.clone(), x2.clone(), -1, None) + .into_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn test_cosine_similarity_near_zero() { + // Test with near-zero vectors + let x1 = TestTensor::<2>::from([[1e-10, 2e-10, 3e-10], [4e-10, 5e-10, 6e-10]]); + let x2 = TestTensor::<2>::from([[2e-10, 4e-10, 6e-10], [8e-10, 10e-10, 12e-10]]); + + // Update the expected values based on the actual implementation behavior + let expected = TensorData::from([[0.0028], [0.0154]]); + + // Smaller values result in NaN on metal f16 + let epsilon = Some(1e-2); + let tolerance = Tolerance::absolute(0.2); + + linalg::cosine_similarity(x1, x2, 1, epsilon) + .into_data() + .assert_approx_eq::(&expected, tolerance); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/linalg/det.rs b/crates/burn-backend-tests/tests/tensor/float/linalg/det.rs new file mode 100644 index 0000000..05b8de1 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/linalg/det.rs @@ -0,0 +1,537 @@ +use super::*; +use burn_tensor::{Distribution, Tolerance, linalg::det, s}; + +// --------------------------------------------------------------------- +// Small Matrices (single batch) +// --------------------------------------------------------------------- + +#[test] +fn test_det_1x1_batched() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data([[[5.0]]], &device); + let det_tensor = det::<3, 2, 1>(tensor); + let expected = TestTensor::<1>::from_data([5.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_2x2_batched() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data([[[4.0, 3.0], [6.0, 3.0]]], &device); + let det_tensor = det::<3, 2, 1>(tensor); + // det = 4*3 - 3*6 = 12 - 18 = -6 + let expected = TestTensor::<1>::from_data([-6.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_3x3_batched() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data( + [[[4.0, 7.0, 3.0], [6.0, 1.0, 3.0], [8.0, 3.0, 7.0]]], + &device, + ); + let det_tensor = det::<3, 2, 1>(tensor); + // det = 4*(1*7 - 3*3) - 7*(6*7 - 3*8) + 3*(6*3 - 1*8) + // = 4*(7-9) - 7*(42-24) + 3*(18-8) = 4*(-2) - 7*18 + 3*10 = -8 -126 +30 = -104 + let expected = TestTensor::<1>::from_data([-104.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +// --------------------------------------------------------------------- +// Special Matrices (single batch) +// --------------------------------------------------------------------- + +#[test] +fn test_det_3x3_identity_matrix() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data([[[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]], &device); + let det_tensor = det::<3, 2, 1>(tensor); + let expected = TestTensor::<1>::from_data([1.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_300x300_identity_matrix() { + let device = Default::default(); + let tensor = TestTensor::eye(300, &device).unsqueeze_dim(0); + let det_tensor = det::<3, 2, 1>(tensor); + let expected = TestTensor::<1>::from_data([1.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_3x3_singular_zero_col() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data( + [[[0.0, 4.0, 2.0], [0.0, 2.0, 1.0], [0.0, 3.0, 9.0]]], + &device, + ); + let det_tensor = det::<3, 2, 1>(tensor); + let expected = TestTensor::<1>::from_data([0.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_20x20_singular_zero_col() { + let device = Default::default(); + let mut tensor = TestTensor::random([1, 20, 20], Distribution::Default, &device); + tensor = tensor.slice_fill(s![.., .., 16], 0.0); + let det_tensor = det::<3, 2, 1>(tensor); + let expected = TestTensor::<1>::from_data([0.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_3x3_singular_linearly_dependent() { + let device = Default::default(); + // Rows: [1,2,3], [2,4,6], [3,6,9] -> linearly dependent + let tensor = TestTensor::<3>::from_data([[[1., 2., 3.], [2., 4., 6.], [3., 6., 9.]]], &device); + let det_tensor = det::<3, 2, 1>(tensor); + let expected = TestTensor::<1>::from_data([0.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_20x20_singular_linearly_dependent() { + let device = Default::default(); + let mut tensor = TestTensor::random([1, 20, 20], Distribution::Default, &device); + let lin_dep_row1 = tensor.clone().slice(s![.., 5, ..]); + let lin_dep_row2 = lin_dep_row1.mul_scalar(3); + tensor = tensor.slice_assign(s![.., 17, ..], lin_dep_row2); + let det_tensor = det::<3, 2, 1>(tensor); + let expected = TestTensor::<1>::from_data([0.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_3x3_upper_triangular() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data([[[1., 2., 3.], [0., 4., 5.], [0., 0., 6.]]], &device); + let det_tensor = det::<3, 2, 1>(tensor); + // det = 1 * 4 * 6 = 24 + let expected = TestTensor::<1>::from_data([24.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_3x3_lower_triangular() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data([[[2., 0., 0.], [3., 4., 0.], [5., 6., 7.]]], &device); + let det_tensor = det::<3, 2, 1>(tensor); + // det = 2 * 4 * 7 = 56 + let expected = TestTensor::<1>::from_data([56.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_3x3_diagonal_matrix() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data([[[3., 0., 0.], [0., 5., 0.], [0., 0., 2.]]], &device); + let det_tensor = det::<3, 2, 1>(tensor); + // det = 3 * 5 * 2 = 30 + let expected = TestTensor::<1>::from_data([30.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_10x10_all_zeros() { + let device = Default::default(); + let tensor = TestTensor::<3>::zeros([1, 10, 10], &device); + let det_tensor = det::<3, 2, 1>(tensor); + let expected = TestTensor::<1>::from_data([0.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_300x300_all_zeros() { + let device = Default::default(); + let tensor = TestTensor::<3>::zeros([1, 300, 300], &device); + let det_tensor = det::<3, 2, 1>(tensor); + let expected = TestTensor::<1>::from_data([0.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_10x10_all_ones() { + let device = Default::default(); + // matrix of all ones has rank 1 -> det = 0 + let tensor = TestTensor::<3>::ones([1, 10, 10], &device); + let det_tensor = det::<3, 2, 1>(tensor); + let expected = TestTensor::<1>::from_data([0.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_300x300_all_ones() { + let device = Default::default(); + // matrix of all ones has rank 1 -> det = 0 + let tensor = TestTensor::<3>::ones([1, 300, 300], &device); + let det_tensor = det::<3, 2, 1>(tensor); + let expected = TestTensor::<1>::from_data([0.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_large_diagonal() { + let device = Default::default(); + // 8x8 diagonal with known product: 8! = 40320 + let eye = Tensor::<2>::eye(8, &device); + let values = Tensor::arange(1..9, &device).float(); + let diag_2d = eye * values.unsqueeze::<2>(); + let tensor = diag_2d.unsqueeze_dim(0); + let det_tensor = det::<3, 2, 1>(tensor); + let expected = TestTensor::<1>::from_data([40320.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_mixed_singular_non_singular_batch() { + let device = Default::default(); + let mut singular_tensor = Tensor::<2>::random([10, 10], Distribution::Default, &device); + singular_tensor = singular_tensor.slice_fill(s![.., 8], 0.0); + let identity_tensor = Tensor::<2>::eye(10, &device); + let input_tensor = Tensor::stack(vec![singular_tensor, identity_tensor], 0); + let det_tensor = det::<3, 2, 1>(input_tensor); + let expected = TestTensor::<1>::from_data([0.0, 1.0], &device); + let tolerance = Tolerance::default(); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +// --------------------------------------------------------------------- +// Batched Tensors (3D, multiple matrices) +// --------------------------------------------------------------------- + +#[test] +fn test_det_batch_of_1x1() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data([[[2.0]], [[3.0]], [[5.0]]], &device); + let det_tensor = det::<3, 2, 1>(tensor); + let expected = TestTensor::<1>::from_data([2.0, 3.0, 5.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_batch_of_2x2() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data( + [ + [[1.0, 2.0], [3.0, 4.0]], // det = -2 + [[2.0, 0.0], [0.0, 3.0]], // det = 6 + [[5.0, 6.0], [7.0, 8.0]], // det = -2 + ], + &device, + ); + let det_tensor = det::<3, 2, 1>(tensor); + let expected = TestTensor::<1>::from_data([-2.0, 6.0, -2.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_batch_of_3x3() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data( + [ + [[4.0, 7.0, 3.0], [6.0, 1.0, 3.0], [8.0, 3.0, 7.0]], // det = -104 + [[1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 3.0]], // det = 6 + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], // det = 0 (singular) + ], + &device, + ); + let det_tensor = det::<3, 2, 1>(tensor); + let expected = TestTensor::<1>::from_data([-104.0, 6.0, 0.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +// --------------------------------------------------------------------- +// 4D Tensors (two batch dimensions) +// --------------------------------------------------------------------- + +#[test] +fn test_det_4d_square() { + let device = Default::default(); + // Shape [2, 2, 3, 3] + let tensor = TestTensor::<4>::from_data( + [ + [ + [[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]], + [[2., 0., 0.], [0., 2., 0.], [0., 0., 2.]], + ], + [ + [[3., 0., 0.], [0., 3., 0.], [0., 0., 3.]], + [[4., 0., 0.], [0., 4., 0.], [0., 0., 4.]], + ], + ], + &device, + ); + let det_tensor = det::<4, 3, 2>(tensor); + let expected = TestTensor::<2>::from_data([[1.0, 8.0], [27.0, 64.0]], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_4d_rectangular_batch() { + let device = Default::default(); + let tensor = TestTensor::<4>::random([3, 4, 5, 5], Distribution::Default, &device); + let det_tensor = det::<4, 3, 2>(tensor.clone()); + assert_eq!(det_tensor.dims(), [3, 4]); +} + +// --------------------------------------------------------------------- +// Edge Cases and Properties +// --------------------------------------------------------------------- + +#[test] +fn test_det_product_property() { + let device = Default::default(); + // det(A * B) = det(A) * det(B) + let a = TestTensor::<3>::random([1, 10, 10], Distribution::Default, &device); + let b = TestTensor::<3>::random([1, 10, 10], Distribution::Default, &device); + let c = a.clone().matmul(b.clone()); + let det_a = det::<3, 2, 1>(a); + let det_b = det::<3, 2, 1>(b); + let det_c = det::<3, 2, 1>(c); + let det_ab = det_a * det_b; + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_c + .into_data() + .assert_approx_eq::(&det_ab.into_data(), tolerance); +} + +#[test] +fn test_det_transpose_invariance() { + let device = Default::default(); + let tensor = TestTensor::<3>::random([1, 10, 10], Distribution::Default, &device); + let det_original = det::<3, 2, 1>(tensor.clone()); + let det_transpose = det::<3, 2, 1>(tensor.transpose()); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_original + .into_data() + .assert_approx_eq::(&det_transpose.into_data(), tolerance); +} + +#[test] +fn test_det_negative_elements() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data([[[-1.0, 2.0], [3.0, -4.0]]], &device); + let det_tensor = det::<3, 2, 1>(tensor); + // det = (-1)*(-4) - (2*3) = 4 - 6 = -2 + let expected = TestTensor::<1>::from_data([-2.0], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_f16_dtype_roundtrip_1x1() { + // Does not perform upcasting when the tests are run in normal mode (not f16) + let device = Default::default(); + let tensor = TestTensor::<3>::random([1, 1, 1], Distribution::Default, &device); + let det_tensor = det::<3, 2, 1>(tensor.clone()); + assert_eq!(tensor.dtype(), det_tensor.dtype()); +} + +#[test] +fn test_det_f16_dtype_roundtrip_2x2() { + // Does not perform upcasting when the tests are run in normal mode (not f16) + let device = Default::default(); + let tensor = TestTensor::<3>::random([1, 2, 2], Distribution::Default, &device); + let det_tensor = det::<3, 2, 1>(tensor.clone()); + assert_eq!(tensor.dtype(), det_tensor.dtype()); +} + +#[test] +fn test_det_f16_dtype_roundtrip_3x3() { + // Does not perform upcasting when the tests are run in normal mode (not f16) + let device = Default::default(); + let tensor = TestTensor::<3>::random([1, 3, 3], Distribution::Default, &device); + let det_tensor = det::<3, 2, 1>(tensor.clone()); + assert_eq!(tensor.dtype(), det_tensor.dtype()); +} + +#[test] +fn test_det_f16_dtype_roundtrip_10x10() { + // Does not perform upcasting when the tests are run in normal mode (not f16) + let device = Default::default(); + let tensor = TestTensor::<3>::random([1, 10, 10], Distribution::Default, &device); + let det_tensor = det::<3, 2, 1>(tensor.clone()); + assert_eq!(tensor.dtype(), det_tensor.dtype()); +} + +// --------------------------------------------------------------------- +// Panic Tests (invalid input) +// --------------------------------------------------------------------- + +#[test] +#[should_panic(expected = "The input tensor must have at least 3 dimensions")] +fn test_det_panic_rank_less_than_3() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + let _ = det::<2, 1, 0>(tensor); +} + +#[test] +#[should_panic(expected = "The last two dimensions of the input tensor must be equal")] +fn test_det_panic_non_square_last_dims() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]], &device); // shape [1,2,3] + let _ = det::<3, 2, 1>(tensor); +} + +#[test] +#[should_panic(expected = "D - 1 = D1 must hold for the generic parameters")] +fn test_det_panic_invalid_d1_generic() { + let device = Default::default(); + let tensor = TestTensor::<3>::ones([1, 2, 2], &device); + let _ = det::<3, 3, 2>(tensor); +} + +#[test] +#[should_panic(expected = "The output tensor rank must be less than input tensor rank by 2")] +fn test_det_panic_invalid_d2_generic() { + let device = Default::default(); + let tensor = TestTensor::<3>::ones([1, 2, 2], &device); + let _ = det::<3, 2, 2>(tensor); +} + +// --------------------------------------------------------------------- +// Random Tensors (expected determinants from PyTorch) +// --------------------------------------------------------------------- + +#[test] +fn test_det_random_2x3x3_pytorch_comparison() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data( + [ + [ + [0.8823, 0.9150, 0.3829], + [0.9593, 0.3904, 0.6009], + [0.2566, 0.7936, 0.9408], + ], + [ + [0.1332, 0.9346, 0.5936], + [0.8694, 0.5677, 0.7411], + [0.4294, 0.8854, 0.5739], + ], + ], + &device, + ); + let det_tensor = det::<3, 2, 1>(tensor); + // Expected determinants computed by PyTorch + let expected = TestTensor::<1>::from_data([-0.5283, 0.0993], &device); + let tolerance = Tolerance::default().set_half_precision_absolute(5e-3); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +#[test] +fn test_det_random_2x2x4x4_pytorch_comparison() { + let device = Default::default(); + let tensor = TestTensor::<4>::from_data( + [ + [ + [ + [0.7713, 0.0208, 0.6336, 0.7488], + [0.4985, 0.2248, 0.1981, 0.7605], + [0.1691, 0.0883, 0.6854, 0.9534], + [0.0039, 0.5122, 0.8126, 0.6125], + ], + [ + [0.7218, 0.2919, 0.9178, 0.7146], + [0.5425, 0.1422, 0.3733, 0.6741], + [0.4418, 0.4340, 0.6178, 0.5131], + [0.6503, 0.6011, 0.8052, 0.6985], + ], + ], + [ + [ + [0.2922, 0.0077, 0.5508, 0.5940], + [0.4512, 0.3444, 0.3599, 0.6226], + [0.8929, 0.8856, 0.1539, 0.8586], + [0.7683, 0.8819, 0.9567, 0.1534], + ], + [ + [0.4885, 0.6690, 0.9967, 0.7013], + [0.5040, 0.2374, 0.1123, 0.1505], + [0.3819, 0.4428, 0.0835, 0.7866], + [0.6515, 0.7998, 0.1242, 0.5865], + ], + ], + ], + &device, + ); + let det_tensor = det::<4, 3, 2>(tensor); + // Expected determinants computed by PyTorch + let expected = TestTensor::<2>::from_data([[-0.1527, -0.0053], [0.0307, -0.1039]], &device); + let tolerance = Tolerance::default().set_absolute(1e-4); + det_tensor + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/linalg/diag.rs b/crates/burn-backend-tests/tests/tensor/float/linalg/diag.rs new file mode 100644 index 0000000..c5930dc --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/linalg/diag.rs @@ -0,0 +1,259 @@ +use super::*; +use burn_tensor::{TensorData, linalg::diag}; + +#[test] +fn test_diag_2d_square() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + let result = diag::<2, 1, _>(tensor); + let expected = TensorData::from([1.0, 4.0]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_diag_2d_tall() { + let device = Default::default(); + // 4x2 matrix (tall) - min(4,2) = 2 diagonal elements + let tensor = + TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]], &device); + let result = diag::<2, 1, _>(tensor); + // Result should have shape [2] with values [1.0, 4.0] + let expected = TensorData::from([1.0, 4.0]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_diag_2d_wide() { + let device = Default::default(); + // 2x4 matrix (wide) - min(2,4) = 2 diagonal elements + let tensor = TestTensor::<2>::from_data([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], &device); + let result = diag::<2, 1, _>(tensor); + // Result should have shape [2] with values [1.0, 6.0] + let expected = TensorData::from([1.0, 6.0]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_diag_3d_batch_square() { + let device = Default::default(); + // Batch of 2 matrices, each 2x2 + let tensor = TestTensor::<3>::from_data( + [[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]], + &device, + ); + let result = diag::<3, 2, _>(tensor); + // Result should have shape [2, 2] + let expected = TensorData::from([[1.0, 4.0], [5.0, 8.0]]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_diag_3d_batch_tall() { + let device = Default::default(); + // Batch of 2 matrices, each 3x2 (tall) + let tensor = TestTensor::<3>::from_data( + [ + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], + [[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]], + ], + &device, + ); + let result = diag::<3, 2, _>(tensor); + // Result should have shape [2, 2] - min(3,2) = 2 diagonal elements each + let expected = TensorData::from([[1.0, 4.0], [7.0, 10.0]]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_diag_3d_batch_wide() { + let device = Default::default(); + // Batch of 2 matrices, each 2x3 (wide) + let tensor = TestTensor::<3>::from_data( + [ + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]], + ], + &device, + ); + let result = diag::<3, 2, _>(tensor); + // Result should have shape [2, 2] - min(2,3) = 2 diagonal elements each + let expected = TensorData::from([[1.0, 5.0], [7.0, 11.0]]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_diag_4d_batch_channel_square() { + let device = Default::default(); + // [batch=2, channel=2, rows=2, cols=2] + let tensor = TestTensor::<4>::from_data( + [ + [[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]], + [[[9.0, 10.0], [11.0, 12.0]], [[13.0, 14.0], [15.0, 16.0]]], + ], + &device, + ); + let result = diag::<4, 3, _>(tensor); + // Result should have shape [2, 2, 2] + let expected = TensorData::from([[[1.0, 4.0], [5.0, 8.0]], [[9.0, 12.0], [13.0, 16.0]]]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_diag_4d_batch_channel_tall() { + let device = Default::default(); + // [batch=2, channel=1, rows=3, cols=2] + let tensor = TestTensor::<4>::from_data( + [ + [[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]], + [[[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]]], + ], + &device, + ); + let result = diag::<4, 3, _>(tensor); + // Result should have shape [2, 1, 2] - min(3,2) = 2 diagonal elements each + let expected = TensorData::from([[[1.0, 4.0]], [[7.0, 10.0]]]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_diag_4d_batch_channel_wide() { + let device = Default::default(); + // [batch=1, channel=2, rows=2, cols=4] + let tensor = TestTensor::<4>::from_data( + [[ + [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], + [[9.0, 10.0, 11.0, 12.0], [13.0, 14.0, 15.0, 16.0]], + ]], + &device, + ); + let result = diag::<4, 3, _>(tensor); + // Result should have shape [1, 2, 2] - min(2,4) = 2 diagonal elements each + let expected = TensorData::from([[[1.0, 6.0], [9.0, 14.0]]]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_diag_1x1() { + let device = Default::default(); + // Single element matrix + let tensor = TestTensor::<2>::from_data([[5.0]], &device); + let result = diag::<2, 1, _>(tensor); + // Should return [5.0] with shape [1] + let expected = TensorData::from([5.0]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_diag_single_row() { + let device = Default::default(); + // Single row matrix + let tensor = TestTensor::<2>::from_data([[1.0, 2.0, 3.0]], &device); + let result = diag::<2, 1, _>(tensor); + // min(1,3) = 1, should return [1.0] with shape [1] + let expected = TensorData::from([1.0]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_diag_single_column() { + let device = Default::default(); + // Single column matrix + let tensor = TestTensor::<2>::from_data([[1.0], [2.0], [3.0]], &device); + let result = diag::<2, 1, _>(tensor); + // min(3,1) = 1, should return [1.0] with shape [1] + let expected = TensorData::from([1.0]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_diag_zeros() { + let device = Default::default(); + // Matrix with zeros on diagonal + let tensor = TestTensor::<2>::from_data([[0.0, 1.0], [2.0, 0.0]], &device); + let result = diag::<2, 1, _>(tensor); + // Should extract diagonal: [0.0, 0.0] + let expected = TensorData::from([0.0, 0.0]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_diag_batch_single_element() { + let device = Default::default(); + // Batch with single element matrices + let tensor = TestTensor::<3>::from_data([[[5.0]], [[7.0]]], &device); + let result = diag::<3, 2, _>(tensor); + // Should return [[5.0], [7.0]] with shape [2, 1] + let expected = TensorData::from([[5.0], [7.0]]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_diag_batch_mixed_zeros() { + let device = Default::default(); + // Batch with mixed zero and non-zero diagonal elements + let tensor = TestTensor::<3>::from_data( + [[[1.0, 2.0], [3.0, 0.0]], [[0.0, 5.0], [6.0, 7.0]]], + &device, + ); + let result = diag::<3, 2, _>(tensor); + // Should return [[1.0, 0.0], [0.0, 7.0]] with shape [2, 2] + let expected = TensorData::from([[1.0, 0.0], [0.0, 7.0]]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_diag_int_tensor() { + let device = Default::default(); + // Test with integer tensor + let tensor = TestTensorInt::<2>::from_data([[1, 2], [3, 4]], &device); + let result = diag::<2, 1, _>(tensor); + // Result should have shape [2] with values [1, 4] + let expected = TensorData::from([1, 4]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_diag_int_3x3() { + let device = Default::default(); + // Test with 3x3 integer matrix + let tensor = TestTensorInt::<2>::from_data([[1, 2, 3], [4, 5, 6], [7, 8, 9]], &device); + let result = diag::<2, 1, _>(tensor); + // Result should have shape [3] with values [1, 5, 9] + let expected = TensorData::from([1, 5, 9]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +#[should_panic] +fn test_diag_1d_should_panic() { + let device = Default::default(); + // 1D tensor should panic - diagonal requires at least 2 dimensions + let tensor = TestTensor::<1>::from_data([1.0, 2.0, 3.0], &device); + let _result = diag::<1, 0, _>(tensor); +} + +#[test] +#[should_panic] +fn test_diag_wrong_output_rank_should_panic() { + let device = Default::default(); + // Providing wrong output rank should panic + let tensor = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + let _result = diag::<2, 2, _>(tensor); // Should be 2,1 not 2,2 +} diff --git a/crates/burn-backend-tests/tests/tensor/float/linalg/lu.rs b/crates/burn-backend-tests/tests/tensor/float/linalg/lu.rs new file mode 100644 index 0000000..9f0affb --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/linalg/lu.rs @@ -0,0 +1,321 @@ +use super::*; +use burn_tensor::{Distribution, Tolerance, linalg::lu}; + +#[cfg(not(feature = "cuda"))] +const REL: f32 = 5e-3; +#[cfg(not(feature = "cuda"))] +const ABS: f32 = 1e-3; + +// CUDA may useTensor Core–accelerated matmul with reduced precision accumulation (TF32). +// In this block LU implementation (matmul-heavy), accumulation error and pivoting +// differences can amplify into large local elementwise deviations. +// TODO: for correctness it might be good to allow users to disable some of these configs in burn.toml. +#[cfg(feature = "cuda")] +const REL: f32 = 3e-2; +#[cfg(feature = "cuda")] +const ABS: f32 = 2e-2; + +// --------------------------------------------------------------------- +// Small Matrices +// --------------------------------------------------------------------- + +#[test] +fn test_lu_1x1() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[5.0]], &device); + let (p, l, u) = lu::<2, 1>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(5e-2); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +#[test] +fn test_lu_2x2() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[4.0, 3.0], [6.0, 3.0]], &device); + let (p, l, u) = lu::<2, 1>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(5e-2); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +#[test] +fn test_lu_3x3() { + let device = Default::default(); + let tensor = + TestTensor::<2>::from_data([[4.0, 7.0, 3.0], [6.0, 1.0, 3.0], [8.0, 3.0, 7.0]], &device); + let (p, l, u) = lu::<2, 1>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(5e-2); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +// --------------------------------------------------------------------- +// Special Matrices (using lu) +// --------------------------------------------------------------------- + +#[test] +fn test_lu_identity_matrix() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]], &device); + let (p, l, u) = lu::<2, 1>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(5e-2); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +#[test] +fn test_lu_singular_zero_pivot() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data( + [ + [0.0, 4.0, 2.0, 6.0], + [0.0, 2.0, 2.0, 11.0], + [0.0, 3.0, 9.0, 6.0], + [0.0, 7.0, 10.0, 9.0], + ], + &device, + ); + let (p, l, u) = lu::<2, 1>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(5e-2); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +// --------------------------------------------------------------------- +// 2D Tensors (no batch dimension) +// --------------------------------------------------------------------- + +#[test] +fn test_lu_2d_square() { + let device = Default::default(); + let tensor = TestTensor::<2>::random([6, 6], Distribution::Default, &device); + let (p, l, u) = lu::<2, 1>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(5e-2); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +#[test] +fn test_lu_2d_tall() { + let device = Default::default(); + let tensor = TestTensor::<2>::random([8, 5], Distribution::Default, &device); + let (p, l, u) = lu::<2, 1>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(5e-2); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +#[test] +fn test_lu_2d_wide() { + let device = Default::default(); + let tensor = TestTensor::<2>::random([5, 8], Distribution::Default, &device); + let (p, l, u) = lu::<2, 1>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(5e-2); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +#[test] +fn test_lu_medium_tall() { + let device = Default::default(); + let tensor = TestTensor::<2>::random([256, 128], Distribution::Default, &device); + let (p, l, u) = lu::<2, 1>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(5e-2); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +#[test] +fn test_lu_medium_wide() { + let device = Default::default(); + let tensor = TestTensor::<2>::random([128, 256], Distribution::Default, &device); + let (p, l, u) = lu::<2, 1>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(5e-2); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +// --------------------------------------------------------------------- +// 3D Tensors (1 batch dimension) +// --------------------------------------------------------------------- + +#[test] +fn test_lu_3d_square() { + let device = Default::default(); + let tensor = TestTensor::<3>::random([3, 6, 6], Distribution::Default, &device); + let (p, l, u) = lu::<3, 2>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(5e-2); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +#[test] +fn test_lu_3d_tall() { + let device = Default::default(); + let tensor = TestTensor::<3>::random([3, 8, 5], Distribution::Default, &device); + let (p, l, u) = lu::<3, 2>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(5e-2); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +#[test] +fn test_lu_3d_wide() { + let device = Default::default(); + let tensor = TestTensor::<3>::random([3, 5, 8], Distribution::Default, &device); + let (p, l, u) = lu::<3, 2>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(5e-2); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +// --------------------------------------------------------------------- +// 4D Tensors (2 batch dimensions) +// --------------------------------------------------------------------- + +#[test] +fn test_lu_4d_square() { + let device = Default::default(); + let tensor = TestTensor::<4>::random([2, 2, 6, 6], Distribution::Default, &device); + let (p, l, u) = lu::<4, 3>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(5e-2); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +#[test] +fn test_lu_4d_tall() { + let device = Default::default(); + let tensor = TestTensor::<4>::random([2, 2, 8, 5], Distribution::Default, &device); + let (p, l, u) = lu::<4, 3>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(5e-2); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +#[test] +fn test_lu_4d_wide() { + let device = Default::default(); + let tensor = TestTensor::<4>::random([2, 2, 5, 8], Distribution::Default, &device); + let (p, l, u) = lu::<4, 3>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(5e-2); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +// --------------------------------------------------------------------- +// Large Tensors (Triggers Block LU Dispatch) +// --------------------------------------------------------------------- + +// The block-dispatch tests below feed an unseeded 500-ish-element random +// matrix through LU + matmul reconstruction. On f16 the per-row error +// grows with n, so a borderline matrix can overshoot a tight absolute +// tolerance (seen: 6.1e-2 diff on a reconstructed value of 0.10 with +// 5e-2 tol). Using 1.5e-1 for f16 on these large sizes keeps the test +// deterministic across seeds while still catching real regressions +// (the f16 LU typo regression in #4738 produced errors of O(1)). + +#[test] +fn test_lu_500x500_block_dispatch() { + let device = Default::default(); + let tensor = TestTensor::<2>::random([500, 500], Distribution::Default, &device); + let (p, l, u) = lu::<2, 1>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(1.5e-1); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +#[test] +fn test_lu_500x300_block_dispatch() { + let device = Default::default(); + let tensor = TestTensor::<2>::random([500, 300], Distribution::Default, &device); + let (p, l, u) = lu::<2, 1>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(1.5e-1); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +#[test] +fn test_lu_300x500_block_dispatch() { + let device = Default::default(); + let tensor = TestTensor::<2>::random([300, 500], Distribution::Default, &device); + let (p, l, u) = lu::<2, 1>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(1.5e-1); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +#[test] +fn test_lu_5x300x300_block_dispatch() { + let device = Default::default(); + let tensor = TestTensor::<3>::random([5, 300, 300], Distribution::Default, &device); + let (p, l, u) = lu::<3, 2>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(1.5e-1); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +#[test] +fn test_lu_3x300x500_block_dispatch() { + let device = Default::default(); + let tensor = TestTensor::<3>::random([3, 300, 500], Distribution::Default, &device); + let (p, l, u) = lu::<3, 2>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(1.5e-1); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +#[test] +fn test_lu_3x500x300_block_dispatch() { + let device = Default::default(); + let tensor = TestTensor::<3>::random([3, 500, 300], Distribution::Default, &device); + let (p, l, u) = lu::<3, 2>(tensor.clone()); + let plu = p.matmul(l).matmul(u); + let tolerance = Tolerance::rel_abs(REL, ABS).set_half_precision_absolute(1.5e-1); + plu.into_data() + .assert_approx_eq::(&tensor.into_data(), tolerance); +} + +// --------------------------------------------------------------------- +// Tensor Check Panics +// --------------------------------------------------------------------- + +#[test] +#[should_panic] +fn test_lu_panic_rank_less_than_2() { + // Fails check: D >= 2 + let device = Default::default(); + let tensor = TestTensor::<1>::from_data([1.0, 2.0, 3.0], &device); + let _ = lu::<1, 0>(tensor); +} + +#[test] +#[should_panic] +fn test_lu_panic_invalid_d1() { + // Fails check: D - 1 == D1 (2 - 1 != 2) + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + let _ = lu::<2, 2>(tensor); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/linalg/matvec.rs b/crates/burn-backend-tests/tests/tensor/float/linalg/matvec.rs new file mode 100644 index 0000000..4324e7f --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/linalg/matvec.rs @@ -0,0 +1,102 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance, linalg}; + +#[test] +fn test_matvec_basic_float() { + let device = Default::default(); + let matrix = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + let vector = TestTensor::<1>::from_data([5.0, 6.0], &device); + + let result = linalg::matvec::<2, 1, _>(matrix, vector); + let expected = TensorData::from([17.0, 39.0]); + + result + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_matvec_basic_int() { + let device = Default::default(); + let matrix = TestTensorInt::<2>::from_ints([[2, 0, -1], [1, 3, 2]], &device); + let vector = TestTensorInt::<1>::from_ints([3, -2, 4], &device); + + let result = linalg::matvec::<2, 1, _>(matrix, vector); + let expected = TensorData::from([2, 5]); + + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_matvec_batched() { + let device = Default::default(); + let matrix = TestTensor::<3>::from_data( + [ + [[1.0, 0.0, 2.0], [3.0, 1.0, -1.0]], + [[-2.0, 1.0, 0.0], [0.5, -1.5, 2.0]], + ], + &device, + ); + let vector = TestTensor::<2>::from_data([[1.0, -1.0, 0.5], [2.0, 0.0, -1.0]], &device); + + let result = linalg::matvec::<3, 2, _>(matrix, vector); + let expected = TensorData::from([[2.0, 1.5], [-4.0, -1.0]]); + + result + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_matvec_vector_broadcasts_over_batches() { + let device = Default::default(); + let matrix = TestTensor::<3>::from_data( + [ + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + [[-1.0, 0.0, 2.0], [3.0, 1.0, -2.0]], + ], + &device, + ); + let vector = TestTensor::<2>::from_data([[1.0, 0.0, -1.0]], &device); + + let result = linalg::matvec::<3, 2, _>(matrix, vector); + let expected = TensorData::from([[-2.0, -2.0], [-3.0, 5.0]]); + + result + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_matvec_matrix_broadcasts_over_vector_batches() { + let device = Default::default(); + let matrix = TestTensor::<3>::from_data([[[1.0, 0.0, 2.0], [3.0, -1.0, 1.0]]], &device); + let vector = TestTensor::<2>::from_data([[2.0, 1.0, 0.0], [1.0, -1.0, 3.0]], &device); + + let result = linalg::matvec::<3, 2, _>(matrix, vector); + let expected = TensorData::from([[2.0, 5.0], [7.0, 7.0]]); + + result + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +#[should_panic] +fn test_matvec_invalid_inner_dim_panics() { + let device = Default::default(); + let matrix = TestTensor::<2>::zeros([2, 3], &device); + let vector = TestTensor::<1>::zeros([4], &device); + + let _ = linalg::matvec::<2, 1, _>(matrix, vector); +} + +#[test] +#[should_panic] +fn test_matvec_mismatched_batches_panics() { + let device = Default::default(); + let matrix = TestTensor::<3>::zeros([2, 3, 4], &device); + let vector = TestTensor::<2>::zeros([3, 4], &device); + + let _ = linalg::matvec::<3, 2, _>(matrix, vector); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/linalg/mod.rs b/crates/burn-backend-tests/tests/tensor/float/linalg/mod.rs new file mode 100644 index 0000000..65137dd --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/linalg/mod.rs @@ -0,0 +1,10 @@ +use super::*; + +pub(crate) mod cosine_similarity; +pub(crate) mod det; +pub(crate) mod diag; +pub(crate) mod lu; +pub(crate) mod matvec; +pub(crate) mod outer; +pub(crate) mod trace; +pub(crate) mod vector_norm; diff --git a/crates/burn-backend-tests/tests/tensor/float/linalg/outer.rs b/crates/burn-backend-tests/tests/tensor/float/linalg/outer.rs new file mode 100644 index 0000000..90e52d8 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/linalg/outer.rs @@ -0,0 +1,262 @@ +use super::*; +use burn_tensor::{ElementConversion, Tolerance}; +use burn_tensor::{TensorData, linalg}; + +// ---------- Vector (D=1, R=2) tests ---------- + +#[test] +fn test_outer_basic() { + let u = TestTensor::<1>::from([1.0, 2.0, 3.0]); + let v = TestTensor::<1>::from([4.0, 5.0]); + + let out = linalg::outer::<1, 2, _>(u, v).into_data(); + let expected = TensorData::from([[4.0, 5.0], [8.0, 10.0], [12.0, 15.0]]); + + out.assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_outer_shapes_only() { + let device = Default::default(); + let u = TestTensor::<1>::zeros([3], &device); + let v = TestTensor::<1>::zeros([5], &device); + let out = linalg::outer::<1, 2, _>(u, v); + assert_eq!(out.shape().dims(), [3, 5]); +} + +#[test] +fn test_outer_asymmetry_and_shapes() { + let u = TestTensor::<1>::from([1.0, 2.0]); + let v = TestTensor::<1>::from([3.0, 4.0, 5.0]); + + let uv = linalg::outer::<1, 2, _>(u.clone(), v.clone()); + let vu = linalg::outer::<1, 2, _>(v, u); + + assert_eq!(uv.shape().dims(), [2, 3]); + assert_eq!(vu.shape().dims(), [3, 2]); +} + +#[test] +fn test_outer_zero_left() { + let device = Default::default(); + let u = TestTensor::<1>::zeros([3], &device); + let v = TestTensor::<1>::from([7.0, 8.0]); + + let out = linalg::outer::<1, 2, _>(u, v).into_data(); + let expected = TensorData::zeros::([3, 2]); + + out.assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_outer_zero_right() { + let device = Default::default(); + let u = TestTensor::<1>::from([1.0, -2.0, 3.0]); + let v = TestTensor::<1>::zeros([4], &device); + + let out = linalg::outer::<1, 2, _>(u, v).into_data(); + let expected = TensorData::zeros::([3, 4]); + + out.assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_outer_signs() { + let u = TestTensor::<1>::from([-1.0, 2.0]); + let v = TestTensor::<1>::from([3.0, -4.0]); + + let out = linalg::outer::<1, 2, _>(u, v).into_data(); + let expected = TensorData::from([[-3.0, 4.0], [6.0, -8.0]]); + + out.assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_outer_integer_inputs() { + let u = TestTensorInt::<1>::from([1, 2, 3]); + let v = TestTensorInt::<1>::from([4, 5]); + + let out = linalg::outer::<1, 2, _>(u, v).into_data(); + let expected = TensorData::from([[4, 5], [8, 10], [12, 15]]); + + out.assert_eq(&expected, false); +} + +#[test] +fn test_outer_equivalence_to_matmul() { + let u = TestTensor::<1>::from([1.0, 2.0, 3.0]); + let v = TestTensor::<1>::from([4.0, 5.0]); + + let out = linalg::outer::<1, 2, _>(u.clone(), v.clone()).into_data(); + + let u2 = u.reshape([3, 1]); + let v2 = v.reshape([1, 2]); + let out_matmul = u2.matmul(v2).into_data(); + + out.assert_approx_eq::(&out_matmul, Tolerance::default()); +} + +#[test] +fn test_outer_vector_identity_right_mult() { + let u = TestTensor::<1>::from([2.0, -1.0]); + let v = TestTensor::<1>::from([3.0, 4.0]); + let w = TestTensor::<1>::from([5.0, 6.0]); + + let uv = linalg::outer::<1, 2, _>(u.clone(), v.clone()); + let left = uv.matmul(w.clone().reshape([2, 1])).reshape([2]); + + let v_dot_w = v.dot(w); + let right = u * v_dot_w; + + left.into_data() + .assert_approx_eq::(&right.into_data(), Tolerance::default()); +} + +#[test] +fn test_outer_length_one_vectors() { + let u = TestTensor::<1>::from([3.0]); + let v = TestTensor::<1>::from([4.0, 5.0, 6.0]); + + let out = linalg::outer::<1, 2, _>(u, v).into_data(); + let expected = TensorData::from([[12.0, 15.0, 18.0]]); + + out.assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_outer_large_values() { + let big = 1.0e10; + let u = TestTensor::<1>::from([big, -big]); + let v = TestTensor::<1>::from([big, big]); + + let out = linalg::outer::<1, 2, _>(u, v).into_data(); + let expected = TensorData::from([[big * big, big * big], [-big * big, -big * big]]); + + let tol = Tolerance::relative(1e-6).set_half_precision_relative(1e-3); + out.assert_approx_eq::(&expected, tol); +} + +#[test] +fn test_outer_nan_propagation() { + let u = TestTensor::<1>::from([f32::NAN, 2.0]); + let v = TestTensor::<1>::from([3.0, 4.0]); + + let out = linalg::outer::<1, 2, _>(u, v).into_data(); + + let s: &[FloatElem] = out + .as_slice::() + .expect("outer nan_propagation: as_slice failed"); + + assert!(s[0].is_nan()); + assert!(s[1].is_nan()); + assert_eq!(s[2], 6.0f32.elem::()); + assert_eq!(s[3], 8.0f32.elem::()); +} + +// ---------- Batched (D=2, R=3) tests ---------- + +#[test] +fn test_outer_batched_basic() { + let x = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]); + let y = TestTensor::<2>::from([[5.0, 6.0, 7.0], [8.0, 9.0, 10.0]]); + let out = linalg::outer::<2, 3, _>(x, y).into_data(); + + let expected = TensorData::from([ + [[5.0, 6.0, 7.0], [10.0, 12.0, 14.0]], + [[24.0, 27.0, 30.0], [32.0, 36.0, 40.0]], + ]); + + out.assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_outer_batched_shapes() { + let device = Default::default(); + let x = TestTensor::<2>::zeros([3, 4], &device); + let y = TestTensor::<2>::zeros([3, 5], &device); + let out = linalg::outer::<2, 3, _>(x, y); + assert_eq!(out.shape().dims(), [3, 4, 5]); +} + +#[test] +fn test_outer_batched_zero_left() { + let device = Default::default(); + let x = TestTensor::<2>::zeros([2, 3], &device); + let y = TestTensor::<2>::from([[7.0, 8.0], [9.0, 10.0]]); + let out = linalg::outer::<2, 3, _>(x, y).into_data(); + + let expected = TensorData::zeros::([2, 3, 2]); + out.assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_outer_batched_zero_right() { + let device = Default::default(); + let x = TestTensor::<2>::from([[1.0, -2.0, 3.0], [4.0, 5.0, -6.0]]); + let y = TestTensor::<2>::zeros([2, 4], &device); + let out = linalg::outer::<2, 3, _>(x, y).into_data(); + + let expected = TensorData::zeros::([2, 3, 4]); + out.assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_outer_batched_signs() { + let x = TestTensor::<2>::from([[-1.0, 2.0], [3.0, -4.0]]); + let y = TestTensor::<2>::from([[3.0, -4.0], [-5.0, 6.0]]); + let out = linalg::outer::<2, 3, _>(x, y).into_data(); + + let expected = TensorData::from([[[-3.0, 4.0], [6.0, -8.0]], [[-15.0, 18.0], [20.0, -24.0]]]); + + out.assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_outer_batched_equivalence_to_per_sample_outer() { + let x = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]); + let y = TestTensor::<2>::from([[5.0, 6.0, 7.0], [8.0, 9.0, 10.0]]); + let batched = linalg::outer::<2, 3, _>(x.clone(), y.clone()); + + for b in 0..2 { + let idx = TestTensorInt::<1>::from([b]); + + let xb2d = x.clone().select(0, idx.clone()); // (1, m) + let yb2d = y.clone().select(0, idx); // (1, n) + + let dims_x: [usize; 2] = xb2d.shape().dims(); + let dims_y: [usize; 2] = yb2d.shape().dims(); + let (m, n) = (dims_x[1], dims_y[1]); + + let per = linalg::outer::<1, 2, _>(xb2d.reshape([m]), yb2d.reshape([n])); + + let bat3d = batched.clone().select(0, TestTensorInt::<1>::from([b])); // (m, n) + + let per_len = per.shape().num_elements(); + let per_flat = per.reshape([per_len]).into_data(); + + let bat_len = bat3d.shape().num_elements(); + let bat_flat = bat3d.reshape([bat_len]).into_data(); + + bat_flat.assert_approx_eq::(&per_flat, Tolerance::default()); + } +} + +#[test] +#[should_panic] +fn test_outer_batched_mismatched_batches_panics() { + let device = Default::default(); + let x = TestTensor::<2>::zeros([2, 3], &device); + let y = TestTensor::<2>::zeros([3, 4], &device); + let _ = linalg::outer::<2, 3, _>(x, y); +} + +#[test] +fn test_outer_dim() { + let u = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]); + let v = TestTensor::<2>::from([[4.0, 5.0], [5.0, 6.0]]); + + let out = linalg::outer_dim::<2, 3, _, _>(u, v, 0).into_data(); + let expected = TensorData::from([[[4.0, 10.0], [5.0, 12.0]], [[12.0, 20.0], [15.0, 24.0]]]); + + out.assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/linalg/trace.rs b/crates/burn-backend-tests/tests/tensor/float/linalg/trace.rs new file mode 100644 index 0000000..6240309 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/linalg/trace.rs @@ -0,0 +1,114 @@ +use super::*; +use burn_tensor::linalg::trace; + +#[test] +fn test_trace_2d_square() { + let device = Default::default(); + let tensor = + TestTensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], &device); + let result = trace::<2, 1>(tensor); + let expected = TestTensor::<1>::from_data([15.0], &device); // 1 + 5 + 9 = 15 + + assert_eq!(result.to_data(), expected.to_data()); +} + +#[test] +fn test_trace_2d_rectangular_wide() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], &device); + let result = trace::<2, 1>(tensor); + let expected = TestTensor::<1>::from_data([7.0], &device); // 1 + 6 = 7 + + assert_eq!(result.to_data(), expected.to_data()); +} + +#[test] +fn test_trace_2d_rectangular_tall() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], &device); + let result = trace::<2, 1>(tensor); + let expected = TestTensor::<1>::from_data([5.0], &device); // 1 + 4 = 5 + + assert_eq!(result.to_data(), expected.to_data()); +} + +#[test] +fn test_trace_3d_batch() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data( + [[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]], + &device, + ); + + let result = trace::<3, 2>(tensor); + let expected = TestTensor::<2>::from_data([[5.0], [13.0]], &device); // [1+4=5, 5+8=13] + + assert_eq!(result.to_data(), expected.to_data()); +} + +#[test] +fn test_trace_4d_batch() { + let device = Default::default(); + let tensor = TestTensor::<4>::from_data( + [[ + // Batch 0, Channel 0 + [[1.0, 2.0], [3.0, 4.0]], + // Batch 0, Channel 1 + [[5.0, 6.0], [7.0, 8.0]], + ]], + &device, + ); + + let result = trace::<4, 3>(tensor); + let expected = TestTensor::<3>::from_data([[[5.0], [13.0]]], &device); + + assert_eq!(result.to_data(), expected.to_data()); +} + +#[test] +fn test_trace_single_element() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[42.0]], &device); + let result = trace::<2, 1>(tensor); + let expected = TestTensor::<1>::from_data([42.0], &device); + + assert_eq!(result.to_data(), expected.to_data()); +} + +#[test] +fn test_trace_zeros() { + let device = Default::default(); + let tensor = TestTensor::<2>::zeros([3, 3], &device); + let result = trace::<2, 1>(tensor); + let expected = TestTensor::<1>::from_data([0.0], &device); + + assert_eq!(result.to_data(), expected.to_data()); +} + +#[test] +fn test_trace_negative_values() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[-1.0, 2.0], [3.0, -4.0]], &device); + let result = trace::<2, 1>(tensor); + let expected = TestTensor::<1>::from_data([-5.0], &device); // -1 + (-4) = -5 + + assert_eq!(result.to_data(), expected.to_data()); +} + +#[test] +#[should_panic] +fn test_trace_1d_should_panic() { + let device = Default::default(); + // 1D tensor should panic - trace requires at least 2 dimensions + let tensor = TestTensor::<1>::from_data([1.0, 2.0, 3.0], &device); + let _result = trace::<1, 0>(tensor); +} + +#[test] +#[should_panic] +fn test_trace_wrong_output_rank_should_panic() { + let device = Default::default(); + // Providing wrong output rank should panic + let tensor = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + let _result = trace::<2, 2>(tensor); // Should be 2,1 not 2,2 +} diff --git a/crates/burn-backend-tests/tests/tensor/float/linalg/vector_norm.rs b/crates/burn-backend-tests/tests/tensor/float/linalg/vector_norm.rs new file mode 100644 index 0000000..09de479 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/linalg/vector_norm.rs @@ -0,0 +1,236 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; +use burn_tensor::linalg; + +#[test] +fn test_max_min_abs() { + let x = TestTensor::<2>::from([[1., 2.], [3., 4.]]); + + let expected = TestTensor::<2>::from([[3., 4.]]).into_data(); + linalg::vector_norm(x.clone(), linalg::Norm::LInf, 0) + .into_data() + .assert_eq(&expected, true); + linalg::max_abs_norm(x.clone(), 0) + .into_data() + .assert_eq(&expected, true); + + let expected = TestTensor::<2>::from([[1., 2.]]).into_data(); + linalg::vector_norm(x.clone(), -f64::INFINITY, 0) + .into_data() + .assert_eq(&expected, true); + linalg::vector_norm(x.clone(), f64::NEG_INFINITY, 0) + .into_data() + .assert_eq(&expected, true); + linalg::min_abs_norm(x.clone(), 0) + .into_data() + .assert_eq(&expected, true); + + let expected = TestTensor::<2>::from([[2.], [4.]]).into_data(); + linalg::vector_norm(x.clone(), f64::INFINITY, 1) + .into_data() + .assert_eq(&expected, true); + linalg::max_abs_norm(x.clone(), 1) + .into_data() + .assert_eq(&expected, true); + + let expected = TestTensor::<2>::from([[1.], [3.]]).into_data(); + linalg::vector_norm(x.clone(), -f64::INFINITY, 1) + .into_data() + .assert_eq(&expected, true); + linalg::vector_norm(x.clone(), f64::NEG_INFINITY, 1) + .into_data() + .assert_eq(&expected, true); + linalg::min_abs_norm(x, 1) + .into_data() + .assert_eq(&expected, true); + + // Test with integer tensor + let z = TestTensorInt::<2>::from([[1, 2], [3, 4]]); + + linalg::max_abs_norm(z.clone(), 0) + .into_data() + .assert_eq(&TestTensorInt::<2>::from([[3, 4]]).into_data(), true); + linalg::max_abs_norm(z.clone(), 1) + .into_data() + .assert_eq(&TestTensorInt::<2>::from([[2], [4]]).into_data(), true); + + linalg::min_abs_norm(z.clone(), 0) + .into_data() + .assert_eq(&TestTensorInt::<2>::from([[1, 2]]).into_data(), true); + linalg::min_abs_norm(z, 1) + .into_data() + .assert_eq(&TestTensorInt::<2>::from([[1], [3]]).into_data(), true); +} + +#[test] +fn test_l0_norm() { + let x = TestTensor::<2>::from([[1.0, -2.0, 0.], [0.0, 0., 4.]]); + + let expected = TestTensor::<2>::from([[1., 1., 1.]]).into_data(); + linalg::vector_norm(x.clone(), linalg::Norm::L0, 0) + .into_data() + .assert_eq(&expected, true); + linalg::l0_norm(x.clone(), 0) + .into_data() + .assert_eq(&expected, true); + + let expected = TestTensor::<2>::from([[2.], [1.]]).into_data(); + linalg::vector_norm(x.clone(), 0.0, 1) + .into_data() + .assert_eq(&expected, true); + linalg::l0_norm(x.clone(), 1) + .into_data() + .assert_eq(&expected, true); + + // Test with integer tensor + let z = TestTensorInt::<2>::from([[1, -2, 0], [0, 0, 4]]); + + linalg::l0_norm(z.clone(), 0) + .into_data() + .assert_eq(&TestTensor::<2>::from([[1, 1, 1]]).int().into_data(), true); + linalg::l0_norm(z.clone(), 1) + .into_data() + .assert_eq(&TestTensor::<2>::from([[2], [1]]).int().into_data(), true); +} + +#[test] +fn test_l1_norm() { + let x = TestTensor::<2>::from([[1., 2.], [3., 4.]]); + + let expected = TestTensor::<2>::from([[4.0, 6.0]]).into_data(); + linalg::vector_norm(x.clone(), linalg::Norm::L1, 0) + .into_data() + .assert_eq(&expected, true); + linalg::l1_norm(x.clone(), 0) + .into_data() + .assert_eq(&expected, true); + + let expected = TestTensor::<2>::from([[3.0], [7.0]]).into_data(); + linalg::vector_norm(x.clone(), 1.0, 1) + .into_data() + .assert_eq(&expected, true); + linalg::l1_norm(x.clone(), 1) + .into_data() + .assert_eq(&expected, true); +} + +#[test] +fn test_lp_norm() { + let x = TestTensor::<2>::from([[1., -2., 0.], [0., 3., 4.]]); + let tolerance = Tolerance::relative(1e-5).set_half_precision_relative(2e-3); + + fn lp_norm_naive(x: Tensor, p: f64, dim: usize) -> Tensor { + x.abs().powf_scalar(p).sum_dim(dim).powf_scalar(1. / p) + } + + // Arbitrary P + let expected = TestTensor::<2>::from([[1.0, 3.2710664, 4.0]]).into_data(); + linalg::vector_norm(x.clone(), 3, 0) + .into_data() + .assert_approx_eq::(&expected, tolerance); + linalg::lp_norm(x.clone(), 3., 0) + .into_data() + .assert_approx_eq::(&expected, tolerance); + + // L0 + let expected = TestTensor::<2>::from([[1., 2., 1.]]).into_data(); + linalg::vector_norm(x.clone(), linalg::Norm::L0, 0) + .into_data() + .assert_eq(&expected, true); + linalg::l0_norm(x.clone(), 0) + .into_data() + .assert_eq(&expected, true); + linalg::lp_norm(x.clone(), 0.0, 0) + .into_data() + .assert_eq(&expected, true); + + // L1 + let expected = TestTensor::<2>::from([[1.0, 5.0, 4.0]]).into_data(); + linalg::vector_norm(x.clone(), linalg::Norm::L1, 0) + .into_data() + .assert_eq(&expected, true); + linalg::l1_norm(x.clone(), 0) + .into_data() + .assert_eq(&expected, true); + lp_norm_naive(x.clone(), 1.0, 0) + .into_data() + .assert_eq(&expected, true); + linalg::lp_norm(x.clone(), 1.0, 0) + .into_data() + .assert_eq(&expected, true); + + // L2 + let expected = TestTensor::<2>::from([[1.0, 3.6055512, 4.0]]).into_data(); + linalg::vector_norm(x.clone(), linalg::Norm::L2, 0) + .into_data() + .assert_approx_eq::(&expected, tolerance); + linalg::l2_norm(x.clone(), 0) + .into_data() + .assert_approx_eq::(&expected, tolerance); + lp_norm_naive(x.clone(), 2.0, 0) + .into_data() + .assert_approx_eq::(&expected, tolerance); + linalg::lp_norm(x.clone(), 2.0, 0) + .into_data() + .assert_approx_eq::(&expected, tolerance); + + // LInf + let expected = TestTensor::<2>::from([[1.0, 3.0, 4.0]]).into_data(); + linalg::vector_norm(x.clone(), linalg::Norm::LInf, 0) + .into_data() + .assert_eq(&expected, true); + linalg::max_abs_norm(x.clone(), 0) + .into_data() + .assert_eq(&expected, true); + linalg::lp_norm(x.clone(), f64::INFINITY, 0) + .into_data() + .assert_approx_eq::(&expected, tolerance); + + // LNegInf + let expected = TestTensor::<2>::from([[0.0, 2.0, 0.0]]).into_data(); + linalg::vector_norm(x.clone(), linalg::Norm::LNegInf, 0) + .into_data() + .assert_approx_eq::(&expected, tolerance); + linalg::min_abs_norm(x.clone(), 0) + .into_data() + .assert_eq(&expected, true); + linalg::lp_norm(x.clone(), f64::NEG_INFINITY, 0) + .into_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn test_l2_norm() { + let x = TestTensor::<2>::from([[1., 2.], [3., 4.]]); + let tolerance = Tolerance::relative(1e-5).set_half_precision_relative(1e-3); + + let expected = TestTensor::<2>::from([[3.16227766, 4.47213595]]).into_data(); + linalg::vector_norm(x.clone(), linalg::Norm::L2, 0) + .into_data() + .assert_approx_eq::(&expected, tolerance); + linalg::l2_norm(x.clone(), 0) + .into_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TestTensor::<2>::from([[2.23606798], [5.0]]).into_data(); + linalg::vector_norm(x.clone(), 2.0, 1) + .into_data() + .assert_approx_eq::(&expected, tolerance); + linalg::l2_norm(x.clone(), 1) + .into_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn test_normalize() { + let x = TestTensor::<2>::from([[1., 2.], [3., 4.]]); + + let expected = TensorData::from([[1. / 4., 2. / 6.], [3. / 4., 4. / 6.]]); + let output = linalg::vector_normalize(x.clone(), 1.0, 0, 0.25).into_data(); + output.assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([[1. / 5., 2. / 6.], [3. / 5., 4. / 6.]]); + let output = linalg::vector_normalize(x.clone(), 1.0, 0, 5.0).into_data(); + output.assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/mod.rs b/crates/burn-backend-tests/tests/tensor/float/mod.rs new file mode 100644 index 0000000..b5d2650 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/mod.rs @@ -0,0 +1,13 @@ +#[allow(unused_imports)] +pub use super::*; // re-export test types + +mod activation; +mod grid; +mod linalg; +mod module; +mod ops; +mod primitive; +mod stats; + +#[cfg(feature = "quantization")] +mod quantization; diff --git a/crates/burn-backend-tests/tests/tensor/float/module/adaptive_avgpool1d.rs b/crates/burn-backend-tests/tests/tensor/float/module/adaptive_avgpool1d.rs new file mode 100644 index 0000000..4a9b284 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/adaptive_avgpool1d.rs @@ -0,0 +1,70 @@ +use super::*; +use burn_tensor::Shape; +use burn_tensor::Tolerance; +use burn_tensor::module::adaptive_avg_pool1d; + +#[test] +fn test_adaptive_avg_pool1d_simple() { + let test = AdaptiveAvgPool1dTestCase { + batch_size: 1, + channels: 2, + length: 8, + length_out: 4, + }; + + test.assert_output(TestTensor::from([[ + [0.5, 2.5, 4.5, 6.5], + [8.5, 10.5, 12.5, 14.5], + ]])); +} + +#[test] +fn test_adaptive_avg_pool1d_dyn_filter_size() { + let test = AdaptiveAvgPool1dTestCase { + batch_size: 1, + channels: 2, + length: 7, + length_out: 3, + }; + + test.assert_output(TestTensor::from([[[1.0, 3.0, 5.0], [8.0, 10.0, 12.0]]])); +} + +#[test] +fn test_adaptive_avg_pool1d_bigger_output() { + let test = AdaptiveAvgPool1dTestCase { + batch_size: 1, + channels: 2, + length: 4, + length_out: 8, + }; + + test.assert_output(TestTensor::from([[ + [0.0, 0.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0], + [4.0, 4.0, 5.0, 5.0, 6.0, 6.0, 7.0, 7.0], + ]])); +} + +struct AdaptiveAvgPool1dTestCase { + batch_size: usize, + channels: usize, + length: usize, + length_out: usize, +} + +impl AdaptiveAvgPool1dTestCase { + fn assert_output(self, y: TestTensor<3>) { + let shape_x = Shape::new([self.batch_size, self.channels, self.length]); + let device = Default::default(); + let x = TestTensor::from_data( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<3, _>(shape_x) + .into_data(), + &device, + ); + let output = adaptive_avg_pool1d(x, self.length_out); + + y.into_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); + } +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/adaptive_avgpool2d.rs b/crates/burn-backend-tests/tests/tensor/float/module/adaptive_avgpool2d.rs new file mode 100644 index 0000000..baeb034 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/adaptive_avgpool2d.rs @@ -0,0 +1,101 @@ +use super::*; +use burn_tensor::Shape; +use burn_tensor::Tolerance; +use burn_tensor::module::adaptive_avg_pool2d; + +#[test] +fn test_adaptive_avg_pool2d_simple() { + let test = AdaptiveAvgPool2dTestCase { + batch_size: 1, + channels: 2, + height: 8, + width: 6, + height_out: 4, + width_out: 4, + }; + + test.assert_output(TestTensor::from([[ + [ + [3.5000, 4.5000, 6.5000, 7.5000], + [15.5000, 16.5000, 18.5000, 19.5000], + [27.5000, 28.5000, 30.5000, 31.5000], + [39.5000, 40.5000, 42.5000, 43.5000], + ], + [ + [51.5000, 52.5000, 54.5000, 55.5000], + [63.5000, 64.5000, 66.5000, 67.5000], + [75.5000, 76.5000, 78.5000, 79.5000], + [87.5000, 88.5000, 90.5000, 91.5000], + ], + ]])); +} + +#[test] +fn test_adaptive_avg_pool2d_dyn_filter_size() { + let test = AdaptiveAvgPool2dTestCase { + batch_size: 1, + channels: 2, + height: 5, + width: 7, + height_out: 3, + width_out: 2, + }; + + test.assert_output(TestTensor::from([[ + [[5.0000, 8.0000], [15.5000, 18.5000], [26.0000, 29.0000]], + [[40.0000, 43.0000], [50.5000, 53.5000], [61.0000, 64.0000]], + ]])); +} + +#[test] +fn test_adaptive_avg_pool2d_bigger_output() { + let test = AdaptiveAvgPool2dTestCase { + batch_size: 1, + channels: 2, + height: 4, + width: 3, + height_out: 5, + width_out: 4, + }; + + test.assert_output(TestTensor::from([[ + [ + [0.0000, 0.5000, 1.5000, 2.0000], + [1.5000, 2.0000, 3.0000, 3.5000], + [4.5000, 5.0000, 6.0000, 6.5000], + [7.5000, 8.0000, 9.0000, 9.5000], + [9.0000, 9.5000, 10.5000, 11.0000], + ], + [ + [12.0000, 12.5000, 13.5000, 14.0000], + [13.5000, 14.0000, 15.0000, 15.5000], + [16.5000, 17.0000, 18.0000, 18.5000], + [19.5000, 20.0000, 21.0000, 21.5000], + [21.0000, 21.5000, 22.5000, 23.0000], + ], + ]])); +} + +struct AdaptiveAvgPool2dTestCase { + batch_size: usize, + channels: usize, + height: usize, + width: usize, + height_out: usize, + width_out: usize, +} + +impl AdaptiveAvgPool2dTestCase { + fn assert_output(self, y: TestTensor<4>) { + let shape_x = Shape::new([self.batch_size, self.channels, self.height, self.width]); + let x = TestTensor::from( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &y.device()) + .reshape::<4, _>(shape_x) + .into_data(), + ); + let output = adaptive_avg_pool2d(x, [self.height_out, self.width_out]); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); + } +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/attention.rs b/crates/burn-backend-tests/tests/tensor/float/module/attention.rs new file mode 100644 index 0000000..5314f5e --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/attention.rs @@ -0,0 +1,642 @@ +use super::*; +use burn_tensor::Distribution; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; +use burn_tensor::module::attention; +use burn_tensor::module::attention_fallback; +use burn_tensor::ops::AttentionModuleOptions; +use num_traits::cast::cast; + +#[allow(unused)] +use num_traits::Signed; // f16 + +#[test] +fn test_attention_no_mask() { + // Skip on metal with f16 - flash attention returns zeros + // Enable once this issue is fixed: https://github.com/tracel-ai/burn/issues/4325 + #[cfg(feature = "metal")] + if core::any::TypeId::of::() == core::any::TypeId::of::() { + return; + } + + let num_batches = 1; + let num_heads = 1; + let seq_q = 128; + let seq_kv = 128; + let head_dim = 64; + let val_dim = 64; + + let query = TestTensor::<4>::random( + [num_batches, num_heads, seq_q, head_dim], + Distribution::Uniform(0., 1.), + &Default::default(), + ); + let key = TestTensor::<4>::random( + [num_batches, num_heads, seq_kv, head_dim], + Distribution::Uniform(0., 1.), + &Default::default(), + ); + let value = TestTensor::<4>::random( + [num_batches, num_heads, seq_kv, val_dim], + Distribution::Uniform(0., 1.), + &Default::default(), + ); + + let output = attention( + query.clone(), + key.clone(), + value.clone(), + None, + None, + Default::default(), + ); + + let expected = attention_fallback(query, key, value, None, None, Default::default()); + + output.into_data().assert_approx_eq::( + &expected.into_data(), + Tolerance::rel_abs(1e-2, 1e-3).set_half_precision_relative(1e-1), + ); +} + +#[test] +fn test_attention_custom_scale() { + let [num_batches, num_heads, seq_len, head_dim] = [1, 2, 16, 32]; + + let query = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let key = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let value = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + + let options = AttentionModuleOptions { + scale: Some(0.1), + ..Default::default() + }; + + let output = attention( + query.clone(), + key.clone(), + value.clone(), + None, + None, + options, + ); + + let expected = attention_fallback(query, key, value, None, None, options); + + output.into_data().assert_approx_eq::( + &expected.into_data(), + Tolerance::rel_abs(1e-2, 1e-3).set_half_precision_relative(1e-1), + ); +} + +#[test] +fn test_attention_attn_bias() { + let [num_batches, num_heads, seq_len, head_dim] = [1, 2, 16, 32]; + + let query = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let key = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let value = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let bias = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, seq_len], + Distribution::Uniform(-0.5, 0.5), + &Default::default(), + ); + + let output = attention( + query.clone(), + key.clone(), + value.clone(), + None, + Some(bias.clone()), + Default::default(), + ); + + let expected = attention_fallback(query, key, value, None, Some(bias), Default::default()); + + output.into_data().assert_approx_eq::( + &expected.into_data(), + Tolerance::rel_abs(1e-2, 1e-3).set_half_precision_relative(1e-1), + ); +} + +#[test] +fn test_attention_softcap() { + let [num_batches, num_heads, seq_len, head_dim] = [1, 2, 16, 32]; + + let query = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let key = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let value = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + + let options = AttentionModuleOptions { + softcap: Some(50.0), + ..Default::default() + }; + + let output = attention( + query.clone(), + key.clone(), + value.clone(), + None, + None, + options, + ); + + let expected = attention_fallback(query, key, value, None, None, options); + + output.into_data().assert_approx_eq::( + &expected.into_data(), + Tolerance::rel_abs(1e-2, 1e-3).set_half_precision_relative(1e-1), + ); +} + +#[test] +fn test_attention_is_causal() { + let [num_batches, num_heads, seq_len, head_dim] = [2, 4, 16, 32]; + + let query = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let key = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let value = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + + let options = AttentionModuleOptions { + is_causal: true, + ..Default::default() + }; + + let output = attention( + query.clone(), + key.clone(), + value.clone(), + None, + None, + options, + ); + + let expected = attention_fallback(query, key, value, None, None, options); + + output.into_data().assert_approx_eq::( + &expected.into_data(), + Tolerance::rel_abs(1e-2, 1e-3).set_half_precision_relative(1e-1), + ); +} + +/// Cross-attention: seq_q != seq_k, with causal masking and additive bias. +#[test] +fn test_attention_cross_attention_with_bias() { + let [num_batches, num_heads, seq_q, seq_k, head_dim] = [2, 2, 8, 24, 32]; + + let query = TestTensor::<4>::random( + [num_batches, num_heads, seq_q, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let key = TestTensor::<4>::random( + [num_batches, num_heads, seq_k, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let value = TestTensor::<4>::random( + [num_batches, num_heads, seq_k, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let bias = TestTensor::<4>::random( + [num_batches, num_heads, seq_q, seq_k], + Distribution::Uniform(-0.5, 0.5), + &Default::default(), + ); + + let options = AttentionModuleOptions { + is_causal: true, + ..Default::default() + }; + + let output = attention( + query.clone(), + key.clone(), + value.clone(), + None, + Some(bias.clone()), + options, + ); + + let expected = attention_fallback(query, key, value, None, Some(bias), options); + + output.into_data().assert_approx_eq::( + &expected.into_data(), + Tolerance::rel_abs(1e-2, 1e-3).set_half_precision_relative(1e-1), + ); +} + +/// Regression: softcap must be applied before -inf masking. +/// With causal masking, position 0 can only attend to itself, so output[0] == value[0]. +/// If softcap were applied after masking, tanh(-inf/softcap) = -softcap (finite), +/// and the masked position would leak into the output. +#[test] +fn test_attention_softcap_preserves_causal_mask() { + let [num_batches, num_heads, seq_len, head_dim] = [1, 1, 4, 8]; + + let query = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let key = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let value = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + + let options = AttentionModuleOptions { + softcap: Some(20.0), + is_causal: true, + ..Default::default() + }; + + let output = attention_fallback(query, key, value.clone(), None, None, options); + + // With causal masking, position 0 can only attend to itself (softmax = [1, 0, 0, 0]). + // So output[..., 0, :] must equal value[..., 0, :]. + let output_row0 = output.slice([0..1, 0..1, 0..1, 0..head_dim]); + let value_row0 = value.slice([0..1, 0..1, 0..1, 0..head_dim]); + + output_row0 + .into_data() + .assert_approx_eq::(&value_row0.into_data(), Tolerance::rel_abs(1e-4, 1e-3)); +} + +/// Regression: fully-masked rows must produce 0, not NaN. +/// When a bool mask masks every key position for a query row, all attention +/// scores are -inf and naive softmax yields NaN. +#[test] +fn test_attention_fully_masked_rows_no_nan() { + // Skip test with f16 (fallback uses too big epsilon value) + if core::any::TypeId::of::() == core::any::TypeId::of::() { + return; + } + let [num_batches, num_heads, seq_len, head_dim] = [1, 1, 4, 8]; + + let query = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let key = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let value = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + + let mask = TestTensorBool::<4>::full( + [num_batches, num_heads, seq_len, seq_len], + true, + &Default::default(), + ); + + let output = attention_fallback(query, key, value, Some(mask), None, Default::default()); + + let output_data = output.into_data(); + let values = output_data.as_slice::().unwrap(); + let tol: FloatElem = cast(1e-4f64).unwrap(); + assert!( + !values.iter().any(|v| v.is_nan()), + "Fully-masked rows should produce 0, not NaN" + ); + assert!( + values.iter().all(|v| v.abs() < tol), + "Fully-masked rows should produce values near 0" + ); +} + +/// Causal + partial bool mask combine to fully mask early rows. +/// With seq_len=4, the causal mask allows row 0 to attend only to key 0. +/// The bool mask masks key 0, so row 0 is fully masked while rows 1-3 still +/// have valid positions. Row 0 output must be 0, not NaN. +#[test] +fn test_attention_fully_masked_rows_causal_no_nan() { + // Skip test with f16 (fallback uses too big epsilon value) + if core::any::TypeId::of::() == core::any::TypeId::of::() { + return; + } + + let [num_batches, num_heads, seq_len, head_dim] = [1, 1, 4, 8]; + + let query = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let key = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let value = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + + // Mask only column 0 (key position 0). Combined with causal mask, + // row 0 becomes fully masked since it can only attend to key 0. + #[rustfmt::skip] + let mask_data = TensorData::from([[[[ + true, false, false, false, + true, false, false, false, + true, false, false, false, + true, false, false, false, + ]]]]); + let mask = TestTensorBool::<4>::from_data(mask_data, &Default::default()).reshape([ + num_batches, + num_heads, + seq_len, + seq_len, + ]); + + let options = AttentionModuleOptions { + is_causal: true, + ..Default::default() + }; + + let output = attention_fallback(query, key, value, Some(mask), None, options); + + let output_data = output.into_data(); + let values = output_data.as_slice::().unwrap(); + let tol: FloatElem = cast(1e-4f64).unwrap(); + assert!( + !values.iter().any(|v| v.is_nan()), + "Fully-masked rows should produce 0, not NaN" + ); + // Row 0 (indices 0..head_dim) should be ~0 since it's fully masked + let row0 = &values[..head_dim]; + assert!( + row0.iter().all(|v| v.abs() < tol), + "Fully-masked row 0 should produce values near 0" + ); + // Rows 1-3 should have non-zero output (they have valid positions) + let rest = &values[head_dim..]; + assert!( + rest.iter().any(|v| v.abs() > tol), + "Non-masked rows should produce non-zero output" + ); +} + +/// Combined: mask + bias + custom scale + softcap together. +#[test] +fn test_attention_all_options() { + let [num_batches, num_heads, seq_len, head_dim] = [2, 2, 16, 32]; + + let query = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let key = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let value = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let bias = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, seq_len], + Distribution::Uniform(-0.5, 0.5), + &Default::default(), + ); + // Create a random bool mask by thresholding a uniform float tensor + let mask = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, seq_len], + Distribution::Uniform(0., 1.), + &Default::default(), + ) + .greater_elem(0.7); + + let options = AttentionModuleOptions { + scale: Some(0.05), + softcap: Some(30.0), + is_causal: true, + }; + + let output = attention( + query.clone(), + key.clone(), + value.clone(), + Some(mask.clone()), + Some(bias.clone()), + options, + ); + + let expected = attention_fallback(query, key, value, Some(mask), Some(bias), options); + + output.into_data().assert_approx_eq::( + &expected.into_data(), + Tolerance::rel_abs(1e-2, 1e-3).set_half_precision_relative(1e-1), + ); +} + +/// Regression for burn#4772: ONNX Attention-23 allows a `[1, 1, seq_q, seq_kv]` bias +/// that's shared across all batches and heads. The main attention path must accept this +/// shape and produce the same result as the fallback (whose elementwise `float_add` +/// broadcasts the bias naturally). +#[test] +fn test_attention_bias_broadcast_batch_and_heads() { + let [num_batches, num_heads, seq_len, head_dim] = [2, 3, 16, 32]; + + let query = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let key = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let value = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let bias = TestTensor::<4>::random( + [1, 1, seq_len, seq_len], + Distribution::Uniform(-0.5, 0.5), + &Default::default(), + ); + + let output = attention( + query.clone(), + key.clone(), + value.clone(), + None, + Some(bias.clone()), + Default::default(), + ); + + let expected = attention_fallback(query, key, value, None, Some(bias), Default::default()); + + output.into_data().assert_approx_eq::( + &expected.into_data(), + Tolerance::rel_abs(1e-2, 1e-3).set_half_precision_relative(1e-1), + ); +} + +/// Regression for burn#4772: `[batch, 1, seq_q, seq_kv]` bias is shared across heads +/// but distinct per batch (the ONNX `test_attention_4d_attn_mask_3d` pattern). +#[test] +fn test_attention_bias_broadcast_heads_only() { + let [num_batches, num_heads, seq_len, head_dim] = [2, 3, 16, 32]; + + let query = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let key = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let value = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let bias = TestTensor::<4>::random( + [num_batches, 1, seq_len, seq_len], + Distribution::Uniform(-0.5, 0.5), + &Default::default(), + ); + + let output = attention( + query.clone(), + key.clone(), + value.clone(), + None, + Some(bias.clone()), + Default::default(), + ); + + let expected = attention_fallback(query, key, value, None, Some(bias), Default::default()); + + output.into_data().assert_approx_eq::( + &expected.into_data(), + Tolerance::rel_abs(1e-2, 1e-3).set_half_precision_relative(1e-1), + ); +} + +/// Regression for burn#4772: a `[1, 1, seq_q, seq_kv]` bool mask must be shared across +/// all batches and heads (the ONNX `test_attention_4d_attn_mask_bool` pattern). +#[test] +fn test_attention_bool_mask_broadcast_batch_and_heads() { + // Skip on cubecl backends with f16: the main attention path diverges from + // attention_fallback starting at batch 0 head 1, the fingerprint of a broadcast + // mask being applied only to head 0. The equivalent full-shape mask case + // (test_attention_all_options) passes on the same backends, so this is a latent + // cubecl flash-attention broadcast issue (tracked in #4778, under the umbrella + // of #4325), not a burn-flex or fallback bug. Enable once #4778 is fixed. + #[cfg(feature = "cube")] + if core::any::TypeId::of::() == core::any::TypeId::of::() { + return; + } + + let [num_batches, num_heads, seq_len, head_dim] = [2, 3, 16, 32]; + + let query = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let key = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + let value = TestTensor::<4>::random( + [num_batches, num_heads, seq_len, head_dim], + Distribution::Uniform(-1., 1.), + &Default::default(), + ); + // Random bool mask via thresholding, matching test_attention_all_options. + let mask = TestTensor::<4>::random( + [1, 1, seq_len, seq_len], + Distribution::Uniform(0., 1.), + &Default::default(), + ) + .greater_elem(0.7); + + let output = attention( + query.clone(), + key.clone(), + value.clone(), + Some(mask.clone()), + None, + Default::default(), + ); + + let expected = attention_fallback(query, key, value, Some(mask), None, Default::default()); + + output.into_data().assert_approx_eq::( + &expected.into_data(), + Tolerance::rel_abs(1e-2, 1e-3).set_half_precision_relative(1e-1), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/avgpool1d.rs b/crates/burn-backend-tests/tests/tensor/float/module/avgpool1d.rs new file mode 100644 index 0000000..8dc2170 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/avgpool1d.rs @@ -0,0 +1,168 @@ +use super::*; +use burn_tensor::Shape; +use burn_tensor::Tolerance; +use burn_tensor::module::avg_pool1d; + +#[test] +fn test_avg_pool1d_simple() { + let test = AvgPool1dTestCase { + batch_size: 1, + channels: 1, + kernel_size: 3, + padding: 0, + stride: 1, + length: 6, + count_include_pad: true, + }; + + test.assert_output(TestTensor::from([[[1., 2., 3., 4.]]])); +} + +#[test] +fn test_avg_pool1d_complex() { + let test = AvgPool1dTestCase { + batch_size: 1, + channels: 2, + kernel_size: 3, + padding: 1, + stride: 2, + length: 6, + count_include_pad: true, + }; + + test.assert_output(TestTensor::from([[ + [0.33333, 2.0000, 4.0000], + [4.33333, 8.0000, 10.0000], + ]])); +} + +#[test] +fn test_avg_pool1d_complex_dont_count_pad() { + let test = AvgPool1dTestCase { + batch_size: 1, + channels: 2, + kernel_size: 3, + padding: 1, + stride: 2, + length: 6, + count_include_pad: false, + }; + + test.assert_output(TestTensor::from([[ + [0.5000, 2.0000, 4.0000], + [6.5000, 8.0000, 10.0000], + ]])); +} + +struct AvgPool1dTestCase { + batch_size: usize, + channels: usize, + kernel_size: usize, + padding: usize, + stride: usize, + length: usize, + count_include_pad: bool, +} + +impl AvgPool1dTestCase { + fn assert_output(self, y: TestTensor<3>) { + let shape_x = Shape::new([self.batch_size, self.channels, self.length]); + let x = TestTensor::from( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &y.device()) + .reshape::<3, _>(shape_x) + .into_data(), + ); + let output = avg_pool1d( + x, + self.kernel_size, + self.stride, + self.padding, + self.count_include_pad, + false, + ); + + y.to_data().assert_approx_eq::( + &output.into_data(), + Tolerance::default().set_half_precision_relative(1e-3), + ); + } +} + +#[test] +fn test_avg_pool1d_ceil_mode() { + // Test ceil_mode=true produces larger output when input doesn't divide evenly by stride + // Input: 1x1x6 (values 0-5), kernel: 3, stride: 2, padding: 0 + // Floor mode: output = (6-3)/2+1 = 2 elements + // Ceil mode: output = ceil((6-3)/2)+1 = ceil(1.5)+1 = 3 elements + let x = TestTensor::from([[[0.0, 1.0, 2.0, 3.0, 4.0, 5.0]]]); + + // With ceil_mode=false (floor): output is 2 elements + // Window 0: avg(0,1,2) = 1 + // Window 1: avg(2,3,4) = 3 + let y_floor = TestTensor::<3>::from([[[1.0, 3.0]]]); + + let output_floor = avg_pool1d( + x.clone(), + 3, // kernel_size + 2, // stride + 0, // padding + true, // count_include_pad + false, + ); + + y_floor.to_data().assert_approx_eq::( + &output_floor.into_data(), + Tolerance::default().set_half_precision_relative(1e-3), + ); + + // With ceil_mode=true: output is 3 elements + // Window 0: avg(0,1,2) = 1 + // Window 1: avg(2,3,4) = 3 + // Window 2: avg(4,5) = 4.5 (partial window, count_include_pad=false divides by 2) + let y_ceil = TestTensor::<3>::from([[[1.0, 3.0, 4.5]]]); + + let output_ceil = avg_pool1d( + x, 3, // kernel_size + 2, // stride + 0, // padding + false, // count_include_pad=false to get correct average for partial window + true, + ); + + y_ceil.to_data().assert_approx_eq::( + &output_ceil.into_data(), + Tolerance::default().set_half_precision_relative(1e-3), + ); +} + +#[test] +fn test_avg_pool1d_ceil_mode_count_include_pad() { + // Test count_include_pad=true + ceil_mode=true interaction for 1D + // When ceil_mode creates windows that extend beyond the padded input: + // - count_include_pad=true should count positions within padded bounds (not ceil_mode extensions) + // + // Input: 1x1x6, kernel 3, stride 2, padding 1, ceil_mode=true + // Output is 4 elements + let x = TestTensor::from([[[0.0, 1.0, 2.0, 3.0, 4.0, 5.0]]]); + + // Expected PyTorch output with padding=1, ceil_mode=true, count_include_pad=true: + // Window 0: positions -1,0,1 -> values 0,0,1 (0 is padding) / 3 = 0.333 + // Window 1: positions 1,2,3 -> values 1,2,3 / 3 = 2.0 + // Window 2: positions 3,4,5 -> values 3,4,5 / 3 = 4.0 + // Window 3: positions 5,6,7 -> only 5 is valid, 6 is padding, 7 is ceil_mode extension + // value 5 / 2 (only 2 positions within padded bounds) = 2.5 + let expected = TestTensor::<3>::from([[[0.3333, 2.0, 4.0, 2.5]]]); + + let output = avg_pool1d( + x, 3, // kernel_size + 2, // stride + 1, // padding + true, // count_include_pad=true + true, // ceil_mode=true + ); + + expected.to_data().assert_approx_eq::( + &output.into_data(), + Tolerance::default().set_half_precision_relative(1e-2), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/avgpool2d.rs b/crates/burn-backend-tests/tests/tensor/float/module/avgpool2d.rs new file mode 100644 index 0000000..233a5d6 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/avgpool2d.rs @@ -0,0 +1,220 @@ +use super::*; +use burn_tensor::Shape; +use burn_tensor::Tolerance; +use burn_tensor::module::avg_pool2d; + +#[test] +fn test_avg_pool2d_simple() { + let test = AvgPool2dTestCase { + batch_size: 1, + channels: 1, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + stride_1: 1, + stride_2: 1, + height: 6, + width: 6, + count_include_pad: true, + }; + + test.assert_output(TestTensor::from([[[ + [7., 8., 9., 10.], + [13., 14., 15., 16.], + [19., 20., 21., 22.], + [25., 26., 27., 28.], + ]]])); +} + +#[test] +fn test_avg_pool2d_complex() { + let test = AvgPool2dTestCase { + batch_size: 1, + channels: 1, + kernel_size_1: 3, + kernel_size_2: 4, + padding_1: 1, + padding_2: 2, + stride_1: 1, + stride_2: 2, + height: 4, + width: 6, + count_include_pad: true, + }; + + test.assert_output(TestTensor::from([[[ + [1.1667, 3.0000, 4.3333, 2.5000], + [3.2500, 7.5000, 9.5000, 5.2500], + [6.2500, 13.5000, 15.5000, 8.2500], + [5.1667, 11.0000, 12.3333, 6.5000], + ]]])); +} + +#[test] +fn test_avg_pool2d_complex_dont_include_pad() { + let test = AvgPool2dTestCase { + batch_size: 1, + channels: 1, + kernel_size_1: 3, + kernel_size_2: 4, + padding_1: 1, + padding_2: 2, + stride_1: 1, + stride_2: 2, + height: 4, + width: 6, + count_include_pad: false, + }; + + test.assert_output(TestTensor::from([[[ + [3.5000, 4.5000, 6.5000, 7.5000], + [6.5000, 7.5000, 9.5000, 10.5000], + [12.5000, 13.5000, 15.5000, 16.5000], + [15.5000, 16.5000, 18.5000, 19.5000], + ]]])); +} + +struct AvgPool2dTestCase { + batch_size: usize, + channels: usize, + kernel_size_1: usize, + kernel_size_2: usize, + padding_1: usize, + padding_2: usize, + stride_1: usize, + stride_2: usize, + height: usize, + width: usize, + count_include_pad: bool, +} + +impl AvgPool2dTestCase { + fn assert_output(self, y: TestTensor<4>) { + let shape_x = Shape::new([self.batch_size, self.channels, self.height, self.width]); + let x = TestTensor::from( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &y.device()) + .reshape::<4, _>(shape_x) + .into_data(), + ); + let output = avg_pool2d( + x, + [self.kernel_size_1, self.kernel_size_2], + [self.stride_1, self.stride_2], + [self.padding_1, self.padding_2], + self.count_include_pad, + false, + ); + + y.to_data().assert_approx_eq::( + &output.into_data(), + Tolerance::default().set_half_precision_relative(1e-3), + ); + } +} + +#[test] +fn test_avg_pool2d_ceil_mode() { + // Test ceil_mode=true produces larger output when input doesn't divide evenly by stride + // Input: 1x1x6x6 (values 0-35), kernel: 3x3, stride: 2x2, padding: 0x0 + // Floor mode: output = (6-3)/2+1 = 2 x 2 + // Ceil mode: output = ceil((6-3)/2)+1 = ceil(1.5)+1 = 3 x 3 + let x = TestTensor::from([[[ + [0.0, 1.0, 2.0, 3.0, 4.0, 5.0], + [6.0, 7.0, 8.0, 9.0, 10.0, 11.0], + [12.0, 13.0, 14.0, 15.0, 16.0, 17.0], + [18.0, 19.0, 20.0, 21.0, 22.0, 23.0], + [24.0, 25.0, 26.0, 27.0, 28.0, 29.0], + [30.0, 31.0, 32.0, 33.0, 34.0, 35.0], + ]]]); + + // With ceil_mode=false (floor): output is 2x2 + // Window (0,0): avg(0,1,2,6,7,8,12,13,14) = avg(63) = 7 + // Window (0,1): avg(2,3,4,8,9,10,14,15,16) = avg(81) = 9 + // Window (1,0): avg(12,13,14,18,19,20,24,25,26) = avg(171) = 19 + // Window (1,1): avg(14,15,16,20,21,22,26,27,28) = avg(189) = 21 + let y_floor = TestTensor::<4>::from([[[[7.0, 9.0], [19.0, 21.0]]]]); + + let output_floor = avg_pool2d( + x.clone(), + [3, 3], + [2, 2], + [0, 0], + true, // count_include_pad + false, + ); + + y_floor.to_data().assert_approx_eq::( + &output_floor.into_data(), + Tolerance::default().set_half_precision_relative(1e-3), + ); + + // With ceil_mode=true: output is 3x3 + // The extra windows at the edge include partial/padded regions + // When count_include_pad=false, only actual values are averaged + // Window (0,2): positions (0:3, 4:6) -> values 4,5,10,11,16,17 -> avg = 10.5 + // Window (1,2): positions (2:5, 4:6) -> values 16,17,22,23,28,29 -> avg = 22.5 + // Window (2,0): positions (4:6, 0:3) -> values 24,25,26,30,31,32 -> avg = 28 + // Window (2,1): positions (4:6, 2:5) -> values 26,27,28,32,33,34 -> avg = 30 + // Window (2,2): positions (4:6, 4:6) -> values 28,29,34,35 -> avg = 31.5 + let y_ceil = + TestTensor::<4>::from([[[[7.0, 9.0, 10.5], [19.0, 21.0, 22.5], [28.0, 30.0, 31.5]]]]); + + let output_ceil = avg_pool2d( + x, + [3, 3], + [2, 2], + [0, 0], + false, // count_include_pad=false to avoid dividing by full kernel size + true, + ); + + y_ceil.to_data().assert_approx_eq::( + &output_ceil.into_data(), + Tolerance::default().set_half_precision_relative(1e-3), + ); +} + +#[test] +fn test_avg_pool2d_ceil_mode_count_include_pad() { + // Test count_include_pad=true + ceil_mode=true interaction + // When ceil_mode creates windows that extend beyond the padded input: + // - count_include_pad=true should count positions within padded bounds (not ceil_mode extensions) + // + // For input 6x6, kernel 3, stride 2, padding 1, ceil_mode=true: + // - Output is 4x4 + // - Corner (3,3) window covers positions beyond even the user padding + // - Expected: 35/4 = 8.75 (divides by count of positions within padded bounds) + + let x = TestTensor::from([[[ + [0.0, 1.0, 2.0, 3.0, 4.0, 5.0], + [6.0, 7.0, 8.0, 9.0, 10.0, 11.0], + [12.0, 13.0, 14.0, 15.0, 16.0, 17.0], + [18.0, 19.0, 20.0, 21.0, 22.0, 23.0], + [24.0, 25.0, 26.0, 27.0, 28.0, 29.0], + [30.0, 31.0, 32.0, 33.0, 34.0, 35.0], + ]]]); + + // Expected PyTorch output with padding=1, ceil_mode=true, count_include_pad=true + // Note: corner (3,3) = 8.75 = 35/4, not 35/9 + let expected = TestTensor::<4>::from([[[ + [1.5556, 3.3333, 4.6667, 2.6667], + [8.3333, 14.0000, 16.0000, 8.5000], + [16.3333, 26.0000, 28.0000, 14.5000], + [10.1667, 16.0000, 17.0000, 8.7500], + ]]]); + + let output = avg_pool2d( + x, + [3, 3], + [2, 2], + [1, 1], + true, // count_include_pad=true + true, // ceil_mode=true + ); + + expected.to_data().assert_approx_eq::( + &output.into_data(), + Tolerance::default().set_half_precision_relative(1e-2), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/bicubic_interpolate.rs b/crates/burn-backend-tests/tests/tensor/float/module/bicubic_interpolate.rs new file mode 100644 index 0000000..f8225b6 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/bicubic_interpolate.rs @@ -0,0 +1,196 @@ +use super::*; +use burn_tensor::Shape; +use burn_tensor::Tolerance; +use burn_tensor::module::interpolate; +use burn_tensor::ops::{InterpolateMode, InterpolateOptions}; + +#[test] +fn test_upsample_interpolation() { + let test = InterpolateTestCase { + batch_size: 2, + channels: 1, + height: 7, + width: 5, + height_out: 8, + width_out: 7, + }; + + test.assert_output(TestTensor::from([ + [[ + [0.0000, 0.5741, 1.3704, 2.0000, 2.6296, 3.4259, 4.0000], + [4.0015, 4.5755, 5.3718, 6.0015, 6.6311, 7.4274, 8.0015], + [8.3528, 8.9268, 9.7231, 10.3528, 10.9824, 11.7787, 12.3528], + [ + 12.7697, 13.3438, 14.1400, 14.7697, 15.3993, 16.1956, 16.7697, + ], + [ + 17.2303, 17.8044, 18.6007, 19.2303, 19.8600, 20.6562, 21.2303, + ], + [ + 21.6472, 22.2213, 23.0176, 23.6472, 24.2769, 25.0731, 25.6472, + ], + [ + 25.9986, 26.5726, 27.3689, 27.9986, 28.6282, 29.4245, 29.9986, + ], + [ + 30.0000, 30.5741, 31.3704, 32.0000, 32.6296, 33.4259, 34.0000, + ], + ]], + [[ + [ + 35.0000, 35.5741, 36.3704, 37.0000, 37.6296, 38.4259, 39.0000, + ], + [ + 39.0015, 39.5755, 40.3718, 41.0015, 41.6311, 42.4274, 43.0015, + ], + [ + 43.3528, 43.9269, 44.7231, 45.3528, 45.9824, 46.7787, 47.3528, + ], + [ + 47.7697, 48.3438, 49.1400, 49.7697, 50.3993, 51.1956, 51.7697, + ], + [ + 52.2303, 52.8044, 53.6007, 54.2303, 54.8600, 55.6562, 56.2303, + ], + [ + 56.6472, 57.2213, 58.0176, 58.6472, 59.2769, 60.0731, 60.6472, + ], + [ + 60.9986, 61.5726, 62.3689, 62.9986, 63.6282, 64.4245, 64.9986, + ], + [ + 65.0000, 65.5741, 66.3704, 67.0000, 67.6296, 68.4259, 69.0000, + ], + ]], + ])); +} + +#[test] +fn test_downsample_interpolation() { + let test = InterpolateTestCase { + batch_size: 1, + channels: 1, + height: 45, + width: 14, + height_out: 4, + width_out: 6, + }; + + test.assert_output(TestTensor::from([[[ + [0.0000, 2.5760, 5.2480, 7.7520, 10.4240, 13.0000], + [204.8148, 207.3908, 210.0628, 212.5668, 215.2388, 217.8148], + [411.1852, 413.7612, 416.4331, 418.9371, 421.6091, 424.1852], + [616.0000, 618.576, 621.2479, 623.7519, 626.4239, 629.0000], + ]]])); +} + +#[test] +fn test_1d_bicubic() { + // Initialize the model without weights (because the exported file does not contain them) + let device = Default::default(); + + // Run the model + let input = TestTensor::<3>::from_data( + [[[1.5410, -0.2934, -2.1788, 0.5684, -1.0845, -1.3986]]], + &device, + ); + + let input = input.unsqueeze_dim(2); + + let output = interpolate( + input, + [1, 9], + InterpolateOptions::new(InterpolateMode::Bicubic), + ); + + assert_eq!(output.dims(), [1, 1, 1, 9]); + + // assert output data does not contain NaN + assert!( + !output + .clone() + .to_data() + .as_slice::() + .unwrap() + .iter() + .any(|&x| x.is_nan()), + "interpolate output contains NaN" + ); + + TestTensor::<4>::from([[[[ + 1.541, 0.5747652, -1.010614, -2.197787, -0.8269969, 0.59609234, -0.5803058, -1.3792794, + -1.3986, + ]]]]) + .to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); +} +struct InterpolateTestCase { + batch_size: usize, + channels: usize, + height: usize, + width: usize, + height_out: usize, + width_out: usize, +} + +impl InterpolateTestCase { + fn assert_output(self, y: TestTensor<4>) { + self.assert_output_with_align_corners(y, true); + } + + fn assert_output_with_align_corners(self, y: TestTensor<4>, align_corners: bool) { + let shape_x = Shape::new([self.batch_size, self.channels, self.height, self.width]); + let x = TestTensor::from( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &y.device()) + .reshape::<4, _>(shape_x) + .into_data(), + ); + let output = interpolate( + x, + [self.height_out, self.width_out], + InterpolateOptions::new(InterpolateMode::Bicubic).with_align_corners(align_corners), + ); + + let tolerance = Tolerance::permissive(); + y.to_data() + .assert_approx_eq::(&output.into_data(), tolerance); + } +} + +#[test] +fn test_upsample_half_pixel() { + let test = InterpolateTestCase { + batch_size: 1, + channels: 1, + height: 4, + width: 4, + height_out: 8, + width_out: 8, + }; + + test.assert_output_with_align_corners( + TestTensor::from([[[ + [ + -0.5273, -0.2305, 0.2461, 0.875, 1.2812, 1.9102, 2.3867, 2.6836, + ], + [ + 0.6602, 0.957, 1.4336, 2.0625, 2.4688, 3.0977, 3.5742, 3.8711, + ], + [ + 2.5664, 2.8633, 3.3398, 3.9688, 4.375, 5.0039, 5.4805, 5.7773, + ], + [5.082, 5.3789, 5.8555, 6.4844, 6.8906, 7.5195, 7.9961, 8.293], + [6.707, 7.0039, 7.4805, 8.1094, 8.5156, 9.1445, 9.6211, 9.918], + [ + 9.2227, 9.5195, 9.9961, 10.625, 11.0312, 11.6602, 12.1367, 12.4336, + ], + [ + 11.1289, 11.4258, 11.9023, 12.5312, 12.9375, 13.5664, 14.043, 14.3398, + ], + [ + 12.3164, 12.6133, 13.0898, 13.7188, 14.125, 14.7539, 15.2305, 15.5273, + ], + ]]]), + false, + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/bilinear_interpolate.rs b/crates/burn-backend-tests/tests/tensor/float/module/bilinear_interpolate.rs new file mode 100644 index 0000000..9050dc7 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/bilinear_interpolate.rs @@ -0,0 +1,270 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::module::interpolate; +use burn_tensor::ops::{InterpolateMode, InterpolateOptions}; +use burn_tensor::{DType, Shape}; + +#[test] +fn test_upsample_interpolation() { + let test = InterpolateTestCase { + batch_size: 2, + channels: 1, + height: 7, + width: 5, + height_out: 8, + width_out: 7, + }; + + test.assert_output(TestTensor::from([ + [[ + [0.0000, 0.6667, 1.3333, 2.0000, 2.6667, 3.3333, 4.0000], + [4.2857, 4.9524, 5.6190, 6.2857, 6.9524, 7.6190, 8.2857], + [8.5714, 9.2381, 9.9048, 10.5714, 11.2381, 11.9048, 12.5714], + [ + 12.8571, 13.5238, 14.1905, 14.8571, 15.5238, 16.1905, 16.8571, + ], + [ + 17.1429, 17.8095, 18.4762, 19.1429, 19.8095, 20.4762, 21.1429, + ], + [ + 21.4286, 22.0952, 22.7619, 23.4286, 24.0952, 24.7619, 25.4286, + ], + [ + 25.7143, 26.3810, 27.0476, 27.7143, 28.3810, 29.0476, 29.7143, + ], + [ + 30.0000, 30.6667, 31.3333, 32.0000, 32.6667, 33.3333, 34.0000, + ], + ]], + [[ + [ + 35.0000, 35.6667, 36.3333, 37.0000, 37.6667, 38.3333, 39.0000, + ], + [ + 39.2857, 39.9524, 40.6190, 41.2857, 41.9524, 42.6190, 43.2857, + ], + [ + 43.5714, 44.2381, 44.9048, 45.5714, 46.2381, 46.9048, 47.5714, + ], + [ + 47.8571, 48.5238, 49.1905, 49.8571, 50.5238, 51.1905, 51.8571, + ], + [ + 52.1429, 52.8095, 53.4762, 54.1429, 54.8095, 55.4762, 56.1429, + ], + [ + 56.4286, 57.0952, 57.7619, 58.4286, 59.0952, 59.7619, 60.4286, + ], + [ + 60.7143, 61.3810, 62.0476, 62.7143, 63.3810, 64.0476, 64.7143, + ], + [ + 65.0000, 65.6667, 66.3333, 67.0000, 67.6667, 68.3333, 69.0000, + ], + ]], + ])); +} + +#[test] +fn test_downsample_interpolation() { + let test = InterpolateTestCase { + batch_size: 1, + channels: 1, + height: 45, + width: 14, + height_out: 4, + width_out: 6, + }; + + test.assert_output(TestTensor::from([[[ + [0.0, 2.6, 5.2, 7.8, 10.4, 13.], + [205.3333, 207.9333, 210.5333, 213.1333, 215.7333, 218.3333], + [410.6667, 413.2667, 415.8667, 418.4667, 421.0667, 423.6667], + [616., 618.6, 621.2, 623.8, 626.4, 629.], + ]]])); +} + +#[test] +fn test_1d_bilinear() { + // Initialize the model without weights (because the exported file does not contain them) + let device = Default::default(); + + // Run the model + let input = TestTensor::<3>::from_data( + [[[1.5410, -0.2934, -2.1788, 0.5684, -1.0845, -1.3986]]], + &device, + ); + + let input = input.unsqueeze_dim(2); + + let output = interpolate( + input, + [1, 9], + InterpolateOptions::new(InterpolateMode::Bilinear), + ); + + assert_eq!(output.dims(), [1, 1, 1, 9]); + + // assert output data does not contain NaN + assert!( + !output + .clone() + .to_data() + .as_slice::() + .unwrap() + .iter() + .any(|&x| x.is_nan()), + "interpolate output contains NaN" + ); + + TestTensor::<4>::from([[[[ + 1.541f32, + 0.39450002, + -0.76475, + -1.943125, + -0.80520004, + 0.36178753, + -0.671275, + -1.2022874, + -1.3986, + ]]]]) + .to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); +} + +#[test] +fn test_interpolate_coord_float_precision_boundary() { + let test = InterpolateTestCase { + batch_size: 1, + channels: 1, + height: 28, + width: 4, + height_out: 24, + width_out: 2, + }; + + test.assert_output(TestTensor::from([[[ + [0.0, 3.0], + [4.6956, 7.6956], + [9.3913, 12.3913], + [14.0869, 17.0869], + [18.7826, 21.7826], + [23.4782, 26.4782], + [28.1739, 31.1739], + [32.8695, 35.8695], + [37.5652, 40.5652], + [42.2608, 45.2608], + [46.9565, 49.9565], + [51.6521, 54.6521], + [56.3478, 59.3478], + [61.0434, 64.0434], + [65.7391, 68.7391], + [70.4347, 73.4347], + [75.1304, 78.1304], + [79.8260, 82.8260], + [84.5217, 87.5217], + [89.2173, 92.2173], + [93.9130, 96.9130], + [98.6086, 101.6086], + [103.3043, 106.3043], + [108.0, 111.0], + ]]])); +} + +#[test] +fn should_interpolate_cast() { + let device = Default::default(); + let shape_x = Shape::new([1, 1, 4, 4]); + let x = TestTensor::from( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<4, _>(shape_x) + .into_data(), + ) + .cast(DType::F32); // ok for f32 backends, casts dtype for f16 tests + let output = interpolate( + x, + [8, 8], + InterpolateOptions::new(InterpolateMode::Bilinear), + ); + + let expected = TestTensor::<4>::from([[[ + [0.0, 0.42857, 0.8571, 1.2857, 1.7142, 2.1428, 2.5714, 3.0], + [1.7142, 2.1428, 2.5714, 3.0, 3.4285, 3.8571, 4.2857, 4.7142], + [3.4285, 3.8571, 4.2857, 4.7142, 5.1428, 5.5714, 6.0, 6.4285], + [5.1428, 5.5714, 6.0, 6.4285, 6.8571, 7.2857, 7.7142, 8.1428], + [6.8571, 7.2857, 7.7142, 8.1428, 8.5714, 9.0, 9.4285, 9.8571], + [ + 8.5714, 9.0, 9.4285, 9.8571, 10.2857, 10.7142, 11.1428, 11.5714, + ], + [ + 10.2857, 10.7142, 11.1428, 11.5714, 12.0, 12.4285, 12.8571, 13.2857, + ], + [ + 12.0, 12.4285, 12.8571, 13.2857, 13.7142, 14.1428, 14.5714, 15.0, + ], + ]]]); + + let tolerance = Tolerance::permissive(); + output + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); +} + +struct InterpolateTestCase { + batch_size: usize, + channels: usize, + height: usize, + width: usize, + height_out: usize, + width_out: usize, +} + +impl InterpolateTestCase { + fn assert_output(self, y: TestTensor<4>) { + self.assert_output_with_align_corners(y, true); + } + + fn assert_output_with_align_corners(self, y: TestTensor<4>, align_corners: bool) { + let shape_x = Shape::new([self.batch_size, self.channels, self.height, self.width]); + let x = TestTensor::from( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &y.device()) + .reshape::<4, _>(shape_x) + .into_data(), + ); + let output = interpolate( + x, + [self.height_out, self.width_out], + InterpolateOptions::new(InterpolateMode::Bilinear).with_align_corners(align_corners), + ); + + let tolerance = Tolerance::permissive(); + y.to_data() + .assert_approx_eq::(&output.into_data(), tolerance); + } +} + +#[test] +fn test_upsample_half_pixel() { + let test = InterpolateTestCase { + batch_size: 1, + channels: 1, + height: 4, + width: 4, + height_out: 8, + width_out: 8, + }; + + test.assert_output_with_align_corners( + TestTensor::from([[[ + [0.0, 0.25, 0.75, 1.25, 1.75, 2.25, 2.75, 3.0], + [1.0, 1.25, 1.75, 2.25, 2.75, 3.25, 3.75, 4.0], + [3.0, 3.25, 3.75, 4.25, 4.75, 5.25, 5.75, 6.0], + [5.0, 5.25, 5.75, 6.25, 6.75, 7.25, 7.75, 8.0], + [7.0, 7.25, 7.75, 8.25, 8.75, 9.25, 9.75, 10.0], + [9.0, 9.25, 9.75, 10.25, 10.75, 11.25, 11.75, 12.0], + [11.0, 11.25, 11.75, 12.25, 12.75, 13.25, 13.75, 14.0], + [12.0, 12.25, 12.75, 13.25, 13.75, 14.25, 14.75, 15.0], + ]]]), + false, + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/conv1d.rs b/crates/burn-backend-tests/tests/tensor/float/module/conv1d.rs new file mode 100644 index 0000000..456f182 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/conv1d.rs @@ -0,0 +1,138 @@ +use super::*; +use burn_tensor::Shape; +use burn_tensor::Tolerance; +use burn_tensor::module::conv1d; +use burn_tensor::ops::ConvOptions; + +#[test] +fn test_conv1d_simple() { + let test = Conv1dTestCase { + batch_size: 2, + channels_in: 2, + channels_out: 2, + kernel_size: 3, + padding: 1, + stride: 1, + dilation: 1, + groups: 1, + length: 4, + }; + + test.assert_output(TestTensor::from([ + [[43., 67., 82., 49.], [104., 176., 227., 158.]], + [[139., 187., 202., 113.], [392., 584., 635., 414.]], + ])); +} + +#[test] +fn test_conv1d_dilation() { + let test = Conv1dTestCase { + batch_size: 2, + channels_in: 2, + channels_out: 2, + kernel_size: 3, + padding: 1, + stride: 1, + dilation: 2, + groups: 1, + length: 4, + }; + + test.assert_output(TestTensor::from([ + [[62., 38.], [159., 111.]], + [[158., 102.], [447., 367.]], + ])); +} + +#[test] +fn test_conv1d_groups() { + let test = Conv1dTestCase { + batch_size: 2, + channels_in: 2, + channels_out: 2, + kernel_size: 3, + padding: 1, + stride: 1, + dilation: 1, + groups: 2, + length: 4, + }; + + test.assert_output(TestTensor::from([ + [[2., 5., 8., 3.], [42., 63., 75., 47.]], + [[26., 29., 32., 11.], [114., 159., 171., 103.]], + ])); +} + +#[test] +fn test_conv1d_complex() { + let test = Conv1dTestCase { + batch_size: 2, + channels_in: 3, + channels_out: 4, + kernel_size: 3, + padding: 1, + stride: 2, + dilation: 1, + groups: 1, + length: 4, + }; + + test.assert_output(TestTensor::from_data( + [ + [[171., 294.], [415., 781.], [659., 1268.], [903., 1755.]], + [[495., 726.], [1387., 2185.], [2279., 3644.], [3171., 5103.]], + ], + &Default::default(), + )); +} + +struct Conv1dTestCase { + batch_size: usize, + channels_in: usize, + channels_out: usize, + kernel_size: usize, + padding: usize, + stride: usize, + dilation: usize, + groups: usize, + length: usize, +} + +impl Conv1dTestCase { + fn assert_output(self, y: TestTensor<3>) { + let shape_x = Shape::new([self.batch_size, self.channels_in, self.length]); + let shape_weight = Shape::new([ + self.channels_out, + self.channels_in / self.groups, + self.kernel_size, + ]); + let device = Default::default(); + let weight = TestTensor::from_data( + TestTensorInt::arange(0..shape_weight.num_elements() as i64, &device) + .reshape::<3, _>(shape_weight) + .into_data(), + &device, + ); + let bias = TestTensor::from_data( + TestTensorInt::arange(0..self.channels_out as i64, &device).into_data(), + &device, + ); + let x = TestTensor::from_data( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<3, _>(shape_x) + .into_data(), + &device, + ); + let output = conv1d( + x, + weight, + Some(bias), + ConvOptions::new([self.stride], [self.padding], [self.dilation], self.groups), + ); + + let tolerance = Tolerance::relative(1e-5).set_half_precision_relative(1e-3); + y.to_data() + .assert_approx_eq::(&output.into_data(), tolerance); + } +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/conv2d.rs b/crates/burn-backend-tests/tests/tensor/float/module/conv2d.rs new file mode 100644 index 0000000..8cc527f --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/conv2d.rs @@ -0,0 +1,762 @@ +use super::*; +use alloc::{vec, vec::Vec}; +use burn_tensor::Shape; +use burn_tensor::activation::gelu; +use burn_tensor::module::conv2d; +use burn_tensor::ops::ConvOptions; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn test_conv2d_simple() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 1, + height: 4, + width: 4, + }; + + test.assert_output(TestTensor::from([[ + [ + [1196., 1796., 1916., 1264.], + [1881., 2793., 2946., 1923.], + [2313., 3405., 3558., 2307.], + [1424., 2072., 2156., 1380.], + ], + [ + [2709., 4173., 4509., 3065.], + [4582., 7006., 7483., 5056.], + [5878., 8914., 9391., 6304.], + [4089., 6177., 6477., 4333.], + ], + ]])); +} + +#[test] +fn test_conv2d_simple_implicit() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 1, + channels_out: 16, + kernel_size_1: 4, + kernel_size_2: 4, + padding_1: 1, + padding_2: 1, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 1, + height: 5, + width: 5, + }; + + test.assert_output(TestTensor::from([[ + [ + [666., 916., 1030., 774.], + [1124., 1500., 1620., 1190.], + [1604., 2100., 2220., 1610.], + [990., 1264., 1330., 936.], + ], + [ + [1531., 2165., 2471., 1927.], + [2757., 3805., 4181., 3207.], + [4197., 5685., 6061., 4587.], + [3295., 4433., 4691., 3529.], + ], + [ + [2396., 3414., 3912., 3080.], + [4390., 6110., 6742., 5224.], + [6790., 9270., 9902., 7564.], + [5600., 7602., 8052., 6122.], + ], + [ + [3261., 4663., 5353., 4233.], + [6023., 8415., 9303., 7241.], + [9383., 12855., 13743., 10541.], + [7905., 10771., 11413., 8715.], + ], + [ + [4126., 5912., 6794., 5386.], + [7656., 10720., 11864., 9258.], + [11976., 16440., 17584., 13518.], + [10210., 13940., 14774., 11308.], + ], + [ + [4991., 7161., 8235., 6539.], + [9289., 13025., 14425., 11275.], + [14569., 20025., 21425., 16495.], + [12515., 17109., 18135., 13901.], + ], + [ + [5856., 8410., 9676., 7692.], + [10922., 15330., 16986., 13292.], + [17162., 23610., 25266., 19472.], + [14820., 20278., 21496., 16494.], + ], + [ + [6721., 9659., 11117., 8845.], + [12555., 17635., 19547., 15309.], + [19755., 27195., 29107., 22449.], + [17125., 23447., 24857., 19087.], + ], + [ + [7586., 10908., 12558., 9998.], + [14188., 19940., 22108., 17326.], + [22348., 30780., 32948., 25426.], + [19430., 26616., 28218., 21680.], + ], + [ + [8451., 12157., 13999., 11151.], + [15821., 22245., 24669., 19343.], + [24941., 34365., 36789., 28403.], + [21735., 29785., 31579., 24273.], + ], + [ + [9316., 13406., 15440., 12304.], + [17454., 24550., 27230., 21360.], + [27534., 37950., 40630., 31380.], + [24040., 32954., 34940., 26866.], + ], + [ + [10181., 14655., 16881., 13457.], + [19087., 26855., 29791., 23377.], + [30127., 41535., 44471., 34357.], + [26345., 36123., 38301., 29459.], + ], + [ + [11046., 15904., 18322., 14610.], + [20720., 29160., 32352., 25394.], + [32720., 45120., 48312., 37334.], + [28650., 39292., 41662., 32052.], + ], + [ + [11911., 17153., 19763., 15763.], + [22353., 31465., 34913., 27411.], + [35313., 48705., 52153., 40311.], + [30955., 42461., 45023., 34645.], + ], + [ + [12776., 18402., 21204., 16916.], + [23986., 33770., 37474., 29428.], + [37906., 52290., 55994., 43288.], + [33260., 45630., 48384., 37238.], + ], + [ + [13641., 19651., 22645., 18069.], + [25619., 36075., 40035., 31445.], + [40499., 55875., 59835., 46265.], + [35565., 48799., 51745., 39831.], + ], + ]])); +} + +#[test] +fn test_conv2d_implicit_padded_in_channels() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 3, + channels_out: 16, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 1, + height: 4, + width: 4, + }; + + test.assert_output(TestTensor::from([[ + [ + [4521., 6753., 7014., 4635.], + [6858., 10197., 10548., 6939.], + [7830., 11601., 11952., 7839.], + [5007., 7383., 7590., 4953.], + ], + [ + [10516., 15988., 16735., 11278.], + [16822., 25507., 26587., 17875.], + [19738., 29827., 30907., 20719.], + [13594., 20506., 21199., 14188.], + ], + [ + [16511., 25223., 26456., 17921.], + [26786., 40817., 42626., 28811.], + [31646., 48053., 49862., 33599.], + [22181., 33629., 34808., 23423.], + ], + [ + [22506., 34458., 36177., 24564.], + [36750., 56127., 58665., 39747.], + [43554., 66279., 68817., 46479.], + [30768., 46752., 48417., 32658.], + ], + [ + [28501., 43693., 45898., 31207.], + [46714., 71437., 74704., 50683.], + [55462., 84505., 87772., 59359.], + [39355., 59875., 62026., 41893.], + ], + [ + [34496., 52928., 55619., 37850.], + [56678., 86747., 90743., 61619.], + [67370., 102731., 106727., 72239.], + [47942., 72998., 75635., 51128.], + ], + [ + [40491., 62163., 65340., 44493.], + [66642., 102057., 106782., 72555.], + [79278., 120957., 125682., 85119.], + [56529., 86121., 89244., 60363.], + ], + [ + [46486., 71398., 75061., 51136.], + [76606., 117367., 122821., 83491.], + [91186., 139183., 144637., 97999.], + [65116., 99244., 102853., 69598.], + ], + [ + [52481., 80633., 84782., 57779.], + [86570., 132677., 138860., 94427.], + [103094., 157409., 163592., 110879.], + [73703., 112367., 116462., 78833.], + ], + [ + [58476., 89868., 94503., 64422.], + [96534., 147987., 154899., 105363.], + [115002., 175635., 182547., 123759.], + [82290., 125490., 130071., 88068.], + ], + [ + [64471., 99103., 104224., 71065.], + [106498., 163297., 170938., 116299.], + [126910., 193861., 201502., 136639.], + [90877., 138613., 143680., 97303.], + ], + [ + [70466., 108338., 113945., 77708.], + [116462., 178607., 186977., 127235.], + [138818., 212087., 220457., 149519.], + [99464., 151736., 157289., 106538.], + ], + [ + [76461., 117573., 123666., 84351.], + [126426., 193917., 203016., 138171.], + [150726., 230313., 239412., 162399.], + [108051., 164859., 170898., 115773.], + ], + [ + [82456., 126808., 133387., 90994.], + [136390., 209227., 219055., 149107.], + [162634., 248539., 258367., 175279.], + [116638., 177982., 184507., 125008.], + ], + [ + [88451., 136043., 143108., 97637.], + [146354., 224537., 235094., 160043.], + [174542., 266765., 277322., 188159.], + [125225., 191105., 198116., 134243.], + ], + [ + [94446., 145278., 152829., 104280.], + [156318., 239847., 251133., 170979.], + [186450., 284991., 296277., 201039.], + [133812., 204228., 211725., 143478.], + ], + ]])); +} + +#[test] +fn test_conv2d_groups_channels_out() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 16, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 2, + height: 4, + width: 4, + }; + + test.assert_output(TestTensor::from([[ + [ + [73., 121., 154., 103.], + [171., 258., 294., 186.], + [279., 402., 438., 270.], + [139., 187., 202., 113.], + ], + [ + [164., 284., 371., 266.], + [415., 664., 781., 538.], + [739., 1132., 1249., 838.], + [518., 782., 851., 564.], + ], + [ + [255., 447., 588., 429.], + [659., 1070., 1268., 890.], + [1199., 1862., 2060., 1406.], + [897., 1377., 1500., 1015.], + ], + [ + [346., 610., 805., 592.], + [903., 1476., 1755., 1242.], + [1659., 2592., 2871., 1974.], + [1276., 1972., 2149., 1466.], + ], + [ + [437., 773., 1022., 755.], + [1147., 1882., 2242., 1594.], + [2119., 3322., 3682., 2542.], + [1655., 2567., 2798., 1917.], + ], + [ + [528., 936., 1239., 918.], + [1391., 2288., 2729., 1946.], + [2579., 4052., 4493., 3110.], + [2034., 3162., 3447., 2368.], + ], + [ + [619., 1099., 1456., 1081.], + [1635., 2694., 3216., 2298.], + [3039., 4782., 5304., 3678.], + [2413., 3757., 4096., 2819.], + ], + [ + [710., 1262., 1673., 1244.], + [1879., 3100., 3703., 2650.], + [3499., 5512., 6115., 4246.], + [2792., 4352., 4745., 3270.], + ], + [ + [5793., 8865., 9330., 6335.], + [9467., 14450., 15134., 10250.], + [11303., 17186., 17870., 12062.], + [7971., 12099., 12546., 8457.], + ], + [ + [6460., 9892., 10411., 7074.], + [10575., 16152., 16917., 11466.], + [12627., 19212., 19977., 13494.], + [8926., 13558., 14059., 9484.], + ], + [ + [7127., 10919., 11492., 7813.], + [11683., 17854., 18700., 12682.], + [13951., 21238., 22084., 14926.], + [9881., 15017., 15572., 10511.], + ], + [ + [7794., 11946., 12573., 8552.], + [12791., 19556., 20483., 13898.], + [15275., 23264., 24191., 16358.], + [10836., 16476., 17085., 11538.], + ], + [ + [8461., 12973., 13654., 9291.], + [13899., 21258., 22266., 15114.], + [16599., 25290., 26298., 17790.], + [11791., 17935., 18598., 12565.], + ], + [ + [9128., 14000., 14735., 10030.], + [15007., 22960., 24049., 16330.], + [17923., 27316., 28405., 19222.], + [12746., 19394., 20111., 13592.], + ], + [ + [9795., 15027., 15816., 10769.], + [16115., 24662., 25832., 17546.], + [19247., 29342., 30512., 20654.], + [13701., 20853., 21624., 14619.], + ], + [ + [10462., 16054., 16897., 11508.], + [17223., 26364., 27615., 18762.], + [20571., 31368., 32619., 22086.], + [14656., 22312., 23137., 15646.], + ], + ]])); +} + +#[test] +fn test_conv2d_groups_remainder_padding_regression() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 32, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 2, + height: 2, + width: 2, + }; + + test.assert_output(TestTensor::from([[ + [[43., 37.], [25., 19.]], + [[98., 92.], [80., 74.]], + [[153., 147.], [135., 129.]], + [[208., 202.], [190., 184.]], + [[263., 257.], [245., 239.]], + [[318., 312.], [300., 294.]], + [[373., 367.], [355., 349.]], + [[428., 422.], [410., 404.]], + [[483., 477.], [465., 459.]], + [[538., 532.], [520., 514.]], + [[593., 587.], [575., 569.]], + [[648., 642.], [630., 624.]], + [[703., 697.], [685., 679.]], + [[758., 752.], [740., 734.]], + [[813., 807.], [795., 789.]], + [[868., 862.], [850., 844.]], + [[3323., 3301.], [3257., 3235.]], + [[3522., 3500.], [3456., 3434.]], + [[3721., 3699.], [3655., 3633.]], + [[3920., 3898.], [3854., 3832.]], + [[4119., 4097.], [4053., 4031.]], + [[4318., 4296.], [4252., 4230.]], + [[4517., 4495.], [4451., 4429.]], + [[4716., 4694.], [4650., 4628.]], + [[4915., 4893.], [4849., 4827.]], + [[5114., 5092.], [5048., 5026.]], + [[5313., 5291.], [5247., 5225.]], + [[5512., 5490.], [5446., 5424.]], + [[5711., 5689.], [5645., 5623.]], + [[5910., 5888.], [5844., 5822.]], + [[6109., 6087.], [6043., 6021.]], + [[6308., 6286.], [6242., 6220.]], + ]])); +} + +#[test] +fn test_conv2d_groups_padding_larger_than_output_width_regression() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 32, + kernel_size_1: 1, + kernel_size_2: 8, + padding_1: 0, + padding_2: 4, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 2, + height: 1, + width: 2, + }; + + test.assert_output(TestTensor::from([[ + [[5., 4., 3.]], + [[14., 13., 12.]], + [[23., 22., 21.]], + [[32., 31., 30.]], + [[41., 40., 39.]], + [[50., 49., 48.]], + [[59., 58., 57.]], + [[68., 67., 66.]], + [[77., 76., 75.]], + [[86., 85., 84.]], + [[95., 94., 93.]], + [[104., 103., 102.]], + [[113., 112., 111.]], + [[122., 121., 120.]], + [[131., 130., 129.]], + [[140., 139., 138.]], + [[679., 674., 669.]], + [[720., 715., 710.]], + [[761., 756., 751.]], + [[802., 797., 792.]], + [[843., 838., 833.]], + [[884., 879., 874.]], + [[925., 920., 915.]], + [[966., 961., 956.]], + [[1007., 1002., 997.]], + [[1048., 1043., 1038.]], + [[1089., 1084., 1079.]], + [[1130., 1125., 1120.]], + [[1171., 1166., 1161.]], + [[1212., 1207., 1202.]], + [[1253., 1248., 1243.]], + [[1294., 1289., 1284.]], + ]])); +} + +#[test] +fn test_conv2d_groups() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 2, + height: 5, + width: 5, + }; + + test.assert_output(TestTensor::from([[ + [[312., 348., 384.], [492., 528., 564.], [672., 708., 744.]], + [ + [3724., 3841., 3958.], + [4309., 4426., 4543.], + [4894., 5011., 5128.], + ], + ]])); +} + +#[test] +fn test_conv2d_groups_multiple_channels() { + let test = Conv2dTestCase { + batch_size: 1, + channels_in: 4, + channels_out: 4, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 2, + height: 5, + width: 5, + }; + + test.assert_output(TestTensor::from([[ + [ + [4035., 4188., 4341.], + [4800., 4953., 5106.], + [5565., 5718., 5871.], + ], + [ + [10030., 10507., 10984.], + [12415., 12892., 13369.], + [14800., 15277., 15754.], + ], + [ + [56075., 56876., 57677.], + [60080., 60881., 61682.], + [64085., 64886., 65687.], + ], + [ + [78270., 79395., 80520.], + [83895., 85020., 86145.], + [89520., 90645., 91770.], + ], + ]])); +} + +#[test] +fn test_conv2d_complex() { + let test = Conv2dTestCase { + batch_size: 2, + channels_in: 3, + channels_out: 4, + kernel_size_1: 3, + kernel_size_2: 2, + padding_1: 1, + padding_2: 2, + stride_1: 2, + stride_2: 3, + dilation_1: 1, + dilation_2: 2, + groups: 1, + height: 4, + width: 5, + }; + + test.assert_output(TestTensor::from([ + [ + [[1845., 3789., 1926.], [3210., 6465., 3228.]], + [[4276., 9082., 4789.], [8071., 16834., 8737.]], + [[6707., 14375., 7652.], [12932., 27203., 14246.]], + [[9138., 19668., 10515.], [17793., 37572., 19755.]], + ], + [ + [[5445., 10629., 5166.], [8070., 15645., 7548.]], + [[14356., 28882., 14509.], [22651., 45454., 22777.]], + [[23267., 47135., 23852.], [37232., 75263., 38006.]], + [[32178., 65388., 33195.], [51813., 105072., 53235.]], + ], + ])); +} + +struct Conv2dTestCase { + batch_size: usize, + channels_in: usize, + channels_out: usize, + kernel_size_1: usize, + kernel_size_2: usize, + padding_1: usize, + padding_2: usize, + stride_1: usize, + stride_2: usize, + dilation_1: usize, + dilation_2: usize, + groups: usize, + height: usize, + width: usize, +} + +impl Conv2dTestCase { + fn assert_output(self, y: TestTensor<4>) { + let shape_x = Shape::new([self.batch_size, self.channels_in, self.height, self.width]); + let shape_weight = Shape::new([ + self.channels_out, + self.channels_in / self.groups, + self.kernel_size_1, + self.kernel_size_2, + ]); + let device = Default::default(); + let weight = TestTensor::from( + TestTensorInt::arange(0..shape_weight.num_elements() as i64, &device) + .reshape::<4, _>(shape_weight) + .into_data(), + ); + let bias = TestTensor::from( + TestTensorInt::arange(0..self.channels_out as i64, &device).into_data(), + ); + let x = TestTensor::from( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<4, _>(shape_x) + .into_data(), + ); + let output = conv2d( + x, + weight, + Some(bias), + ConvOptions::new( + [self.stride_1, self.stride_2], + [self.padding_1, self.padding_2], + [self.dilation_1, self.dilation_2], + self.groups, + ), + ); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); + } +} + +#[rustfmt::skip] // param values are too long + fn conv2d_weight() -> TensorData { + TensorData::new( + vec![0.048065186, -0.3059082, -0.10345459, -0.34643555, -0.20788574, -0.021072388, 0.13745117, -0.05102539, 0.024536133, -0.16479492, -0.19519043, 0.27270508, 0.17700195, -0.33764648, -0.08239746, -0.27929688, 0.17321777, -0.1315918, 0.04574585, -0.17980957, -0.33569336, 0.27612305, 0.30004883, -0.28979492, -0.17297363, -0.021759033, -0.27148438, 0.005657196, 0.29956055, -0.06958008, -0.29345703, -0.14440918, 0.10827637, -0.13305664, -0.20239258, 0.24890137, -0.1541748, -0.20019531, -0.2854004, 0.17016602, 0.07861328, -0.09075928, 0.30908203, -0.00013422966, 0.29589844, 0.15258789, -0.25708008, 0.20422363, -0.2529297, 0.07891846, -0.19506836, 0.23571777, 0.27124023, 0.17370605, -0.16992188, -0.23522949, 0.14648438, -0.09576416, -0.18310547, 0.21044922, -0.08911133, -0.2541504, -0.2775879, -0.2064209, -0.16271973, -0.048919678, -0.03555298, -0.11639404, 0.09661865, -0.10241699, 0.08929443, 0.2866211], + [8, 1, 3, 3], + ) + } + +#[test] +fn test_conv2d_binary_broadcasted() { + let device = Default::default(); + let x = TestTensor::<4>::full([1, 1, 28, 28], -0.42421296, &device); + + // conv2d -> batchnorm -> activation + let weight = TestTensor::from_data(conv2d_weight(), &device); + let bias = TestTensor::from([ + 0.082336426, + -0.049591064, + 0.0031795502, + 0.00095653534, + 0.02357483, + 0.005569458, + 0.07525635, + 0.056396484, + ]); + + // channels: [1, 8], kernel_size: [3, 3], stride: [1, 1], dilation: [1, 1], groups: 1, padding: [0, 0] + let opt = ConvOptions::new([1, 1], [0, 0], [1, 1], 1); + let x = conv2d(x, weight, Some(bias), opt); + + // simulate batchnorm binary ops with broadcasted params + let gamma = TestTensor::<1>::from([ + 1.0048828, 0.9902344, 1.0185547, 0.97558594, 1.0097656, 0.97802734, 1.0009766, 1.0146484, + ]); + let beta = TestTensor::<1>::from([ + 0.026290894, + 0.0007505417, + 0.006134033, + 0.02418518, + 0.07373047, + 0.020507813, + 0.01902771, + 0.02003479, + ]); + let mean = TestTensor::<1>::from([ + 0.029159546, + -0.08673096, + -0.03894043, + -0.01108551, + 0.032440186, + 0.03237915, + 0.013839722, + 0.04397583, + ]) + .reshape([1, 8, 1, 1]); + let var = TestTensor::<1>::from([ + 0.67089844, 0.29956055, 0.5209961, 0.1862793, 0.30419922, 0.21313477, 0.7504883, 0.26342773, + ]) + .reshape([1, 8, 1, 1]); + + let std = var.add_scalar(1e-5).sqrt(); + let x = x.sub(mean); + let x = x.div(std); + let x = x.mul(gamma.reshape([1, 8, 1, 1])); + let x = x.add(beta.reshape([1, 8, 1, 1])); + + let x = gelu(x); + + let expected: Vec = [ + 0.36432067f32, + 0.34909567, + 0.30684796, + 0.13217466, + -0.018471397, + -0.1389876, + 0.39402074, + 0.12394252, + ] + .iter() + .flat_map(|&v| core::iter::repeat_n(v, 676)) + .collect(); + let expected = TensorData::new(expected, [1, 8, 26, 26]); + + x.into_data().assert_approx_eq::( + &expected, + Tolerance::default().set_half_precision_absolute(1e-3), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/conv3d.rs b/crates/burn-backend-tests/tests/tensor/float/module/conv3d.rs new file mode 100644 index 0000000..1143edd --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/conv3d.rs @@ -0,0 +1,299 @@ +use super::*; +use burn_tensor::Shape; +use burn_tensor::Tolerance; +use burn_tensor::module::conv3d; +use burn_tensor::ops::ConvOptions; + +#[test] +fn test_conv3d_simple() { + let test = Conv3dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + kernel_size_3: 3, + padding_1: 1, + padding_2: 1, + padding_3: 1, + stride_1: 1, + stride_2: 1, + stride_3: 1, + dilation_1: 1, + dilation_2: 1, + dilation_3: 1, + groups: 1, + depth: 4, + height: 4, + width: 4, + }; + + test.assert_output(TestTensor::from([[ + [ + [ + [29980.0, 44860.0, 45640.0, 30324.0], + [45072.0, 67380.0, 68496.0, 45468.0], + [48096.0, 71844.0, 72960.0, 48396.0], + [31780.0, 47428.0, 48136.0, 31900.0], + ], + [ + [47292.0, 70548.0, 71556.0, 47400.0], + [70335.0, 104823.0, 106254.0, 70317.0], + [74223.0, 110547.0, 111978.0, 74061.0], + [48552.0, 72240.0, 73140.0, 48324.0], + ], + [ + [58236.0, 86676.0, 87684.0, 57960.0], + [85887.0, 127719.0, 129150.0, 85293.0], + [89775.0, 133443.0, 134874.0, 89037.0], + [58344.0, 86640.0, 87540.0, 57732.0], + ], + [ + [36148.0, 53620.0, 54184.0, 35692.0], + [52740.0, 78144.0, 78936.0, 51936.0], + [54900.0, 81312.0, 82104.0, 54000.0], + [35260.0, 52156.0, 52648.0, 34580.0], + ], + ], + [ + [ + [66701.0, 100589.0, 102665.0, 68773.0], + [102745.0, 154861.0, 157921.0, 105733.0], + [110953.0, 167101.0, 170161.0, 113845.0], + [75413.0, 113525.0, 115529.0, 77261.0], + ], + [ + [112741.0, 169693.0, 172645.0, 115441.0], + [172396.0, 259372.0, 263719.0, 176266.0], + [184060.0, 276760.0, 281107.0, 187786.0], + [124369.0, 186937.0, 189781.0, 126733.0], + ], + [ + [144421.0, 216925.0, 219877.0, 146737.0], + [219052.0, 328924.0, 333271.0, 222346.0], + [230716.0, 346312.0, 350659.0, 233866.0], + [154897.0, 232441.0, 235285.0, 156877.0], + ], + [ + [100517.0, 150821.0, 152681.0, 101789.0], + [151885.0, 227833.0, 230569.0, 153673.0], + [159229.0, 238777.0, 241513.0, 160921.0], + [106541.0, 159725.0, 161513.0, 107589.0], + ], + ], + ]])); +} + +#[test] +fn test_conv3d_groups() { + let test = Conv3dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + kernel_size_3: 3, + padding_1: 0, + padding_2: 0, + padding_3: 0, + stride_1: 1, + stride_2: 1, + stride_3: 1, + dilation_1: 1, + dilation_2: 1, + dilation_3: 1, + groups: 2, + depth: 5, + height: 5, + width: 5, + }; + + test.assert_output(TestTensor::from([[ + [ + [ + [15219., 15570., 15921.], + [16974., 17325., 17676.], + [18729., 19080., 19431.], + ], + [ + [23994., 24345., 24696.], + [25749., 26100., 26451.], + [27504., 27855., 28206.], + ], + [ + [32769., 33120., 33471.], + [34524., 34875., 35226.], + [36279., 36630., 36981.], + ], + ], + [ + [ + [172819., 173899., 174979.], + [178219., 179299., 180379.], + [183619., 184699., 185779.], + ], + [ + [199819., 200899., 201979.], + [205219., 206299., 207379.], + [210619., 211699., 212779.], + ], + [ + [226819., 227899., 228979.], + [232219., 233299., 234379.], + [237619., 238699., 239779.], + ], + ], + ]])); +} + +#[test] +fn test_conv3d_complex() { + let test = Conv3dTestCase { + batch_size: 2, + channels_in: 3, + channels_out: 4, + kernel_size_1: 4, + kernel_size_2: 3, + kernel_size_3: 2, + padding_1: 1, + padding_2: 2, + padding_3: 3, + stride_1: 2, + stride_2: 3, + stride_3: 4, + dilation_1: 1, + dilation_2: 2, + dilation_3: 3, + groups: 1, + depth: 4, + height: 5, + width: 6, + }; + + test.assert_output(TestTensor::from([ + [ + [ + [[149148., 299070., 149850.], [147636., 295758., 148050.]], + [[150660., 301014., 150282.], [147420., 294246., 146754.]], + ], + [ + [[351325., 709903., 358507.], [357589., 722143., 364483.]], + [[391717., 789607., 397819.], [396253., 798391., 402067.]], + ], + [ + [[553502., 1120736., 567164.], [567542., 1148528., 580916.]], + [[632774., 1278200., 645356.], [645086., 1302536., 657380.]], + ], + [ + [[755679., 1531569., 775821.], [777495., 1574913., 797349.]], + [[873831., 1766793., 892893.], [893919., 1806681., 912693.]], + ], + ], + [ + [ + [[408348., 810990., 402570.], [393876., 781758., 387810.]], + [[370980., 735174., 364122.], [354780., 702486., 347634.]], + ], + [ + [ + [1077085., 2154943., 1077787.], + [1070389., 2141263., 1070803.], + ], + [ + [1078597., 2156887., 1078219.], + [1070173., 2139751., 1069507.], + ], + ], + [ + [ + [1745822., 3498896., 1753004.], + [1746902., 3500768., 1753796.], + ], + [ + [1786214., 3578600., 1792316.], + [1785566., 3577016., 1791380.], + ], + ], + [ + [ + [2414559., 4842849., 2428221.], + [2423415., 4860273., 2436789.], + ], + [ + [2493831., 5000313., 2506413.], + [2500959., 5014281., 2513253.], + ], + ], + ], + ])); +} + +struct Conv3dTestCase { + batch_size: usize, + channels_in: usize, + channels_out: usize, + kernel_size_1: usize, + kernel_size_2: usize, + kernel_size_3: usize, + padding_1: usize, + padding_2: usize, + padding_3: usize, + stride_1: usize, + stride_2: usize, + stride_3: usize, + dilation_1: usize, + dilation_2: usize, + dilation_3: usize, + groups: usize, + depth: usize, + height: usize, + width: usize, +} + +impl Conv3dTestCase { + fn assert_output(self, y: TestTensor<5>) { + let shape_x = Shape::new([ + self.batch_size, + self.channels_in, + self.depth, + self.height, + self.width, + ]); + let shape_weight = Shape::new([ + self.channels_out, + self.channels_in / self.groups, + self.kernel_size_1, + self.kernel_size_2, + self.kernel_size_3, + ]); + let device = Default::default(); + let weight = TestTensor::from( + TestTensorInt::arange(0..shape_weight.num_elements() as i64, &device) + .reshape::<5, _>(shape_weight) + .into_data(), + ); + let bias = TestTensor::from( + TestTensorInt::arange(0..self.channels_out as i64, &device).into_data(), + ); + let x = TestTensor::from( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<5, _>(shape_x) + .into_data(), + ); + let output = conv3d( + x, + weight, + Some(bias), + ConvOptions::new( + [self.stride_1, self.stride_2, self.stride_3], + [self.padding_1, self.padding_2, self.padding_3], + [self.dilation_1, self.dilation_2, self.dilation_3], + self.groups, + ), + ); + + let tolerance = Tolerance::relative(1e-5).set_half_precision_relative(2e-3); + y.to_data() + .assert_approx_eq::(&output.into_data(), tolerance); + } +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/conv_transpose1d.rs b/crates/burn-backend-tests/tests/tensor/float/module/conv_transpose1d.rs new file mode 100644 index 0000000..2959e7f --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/conv_transpose1d.rs @@ -0,0 +1,145 @@ +use super::*; +use burn_tensor::Shape; +use burn_tensor::Tolerance; +use burn_tensor::module::conv_transpose1d; +use burn_tensor::ops::ConvTransposeOptions; + +#[test] +fn test_conv_transpose1d_diff_channels() { + let test = ConvTranspose1dTestCase { + batch_size: 1, + channels_in: 3, + channels_out: 2, + kernel_size: 3, + padding: 1, + padding_out: 0, + stride: 1, + dilation: 1, + groups: 1, + length: 4, + }; + + test.assert_output(TestTensor::from([[ + [270., 453., 516., 387.], + [352., 589., 679., 505.], + ]])); +} + +#[test] +fn test_conv_transpose1d_stride() { + let test = ConvTranspose1dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size: 3, + padding: 1, + padding_out: 1, + stride: 2, + dilation: 1, + groups: 1, + length: 4, + }; + + test.assert_output(TestTensor::from([[ + [28., 62., 36., 78., 44., 94., 52., 62.], + [41., 93., 55., 121., 69., 149., 83., 93.], + ]])); +} + +#[test] +fn test_conv_transpose1d_dilation() { + let test = ConvTranspose1dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size: 3, + padding: 1, + padding_out: 0, + stride: 1, + dilation: 2, + groups: 1, + length: 4, + }; + + test.assert_output(TestTensor::from([[ + [30., 64., 78., 76., 94., 52.], + [49., 101., 127., 113., 143., 77.], + ]])); +} + +#[test] +fn test_conv_transpose1d_groups() { + let test = ConvTranspose1dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size: 3, + padding: 1, + padding_out: 0, + stride: 1, + dilation: 1, + groups: 2, + length: 4, + }; + + test.assert_output(TestTensor::from_data( + [[[0., 1., 4., 7.], [32., 59., 71., 59.]]], + &Default::default(), + )); +} + +struct ConvTranspose1dTestCase { + batch_size: usize, + channels_in: usize, + channels_out: usize, + kernel_size: usize, + padding: usize, + padding_out: usize, + stride: usize, + dilation: usize, + groups: usize, + length: usize, +} + +impl ConvTranspose1dTestCase { + fn assert_output(self, y: TestTensor<3>) { + let shape_x = Shape::new([self.batch_size, self.channels_in, self.length]); + let shape_weights = Shape::new([ + self.channels_in, + self.channels_out / self.groups, + self.kernel_size, + ]); + let device = Default::default(); + let weights = TestTensor::from_data( + TestTensorInt::arange(0..shape_weights.num_elements() as i64, &device) + .reshape::<3, _>(shape_weights) + .into_data(), + &device, + ); + let bias = TestTensor::from_data( + TestTensorInt::arange(0..self.channels_out as i64, &device).into_data(), + &device, + ); + let x = TestTensor::from_data( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<3, _>(shape_x) + .into_data(), + &device, + ); + let output = conv_transpose1d( + x, + weights, + Some(bias), + ConvTransposeOptions::new( + [self.stride], + [self.padding], + [self.padding_out], + [self.dilation], + self.groups, + ), + ); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); + } +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/conv_transpose2d.rs b/crates/burn-backend-tests/tests/tensor/float/module/conv_transpose2d.rs new file mode 100644 index 0000000..ec8d069 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/conv_transpose2d.rs @@ -0,0 +1,361 @@ +use super::*; +use burn_tensor::Shape; +use burn_tensor::Tolerance; +use burn_tensor::module::conv_transpose2d; +use burn_tensor::ops::ConvTransposeOptions; + +#[test] +fn test_conv_transpose2d_simple_1() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels_in: 1, + channels_out: 1, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + padding_out_1: 0, + padding_out_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 1, + height: 2, + width: 2, + }; + + test.assert_output(TestTensor::from([[[[5.0, 11.0], [23.0, 29.0]]]])); +} + +#[test] +fn test_conv_transpose2d_simple_2() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels_in: 3, + channels_out: 3, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + padding_out_1: 0, + padding_out_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 1, + height: 4, + width: 4, + }; + + test.assert_output(TestTensor::from([[ + [ + [9855., 15207., 15738., 10797.], + [16290., 25119., 25956., 17793.], + [18486., 28467., 29304., 20061.], + [13593., 20913., 21498., 14703.], + ], + [ + [11854., 18286., 18979., 13012.], + [19612., 30223., 31303., 21439.], + [22456., 34543., 35623., 24355.], + [16456., 25288., 26035., 17782.], + ], + [ + [13853., 21365., 22220., 15227.], + [22934., 35327., 36650., 25085.], + [26426., 40619., 41942., 28649.], + [19319., 29663., 30572., 20861.], + ], + ]])); +} + +#[test] +fn test_conv_transpose2d_simple_3() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels_in: 1, + channels_out: 1, + kernel_size_1: 2, + kernel_size_2: 2, + padding_1: 0, + padding_2: 0, + padding_out_1: 0, + padding_out_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 1, + height: 2, + width: 2, + }; + + test.assert_output(TestTensor::from([[[ + [0.0, 0.0, 1.0], + [0.0, 4.0, 6.0], + [4.0, 12.0, 9.0], + ]]])); +} + +#[test] +fn test_conv_transpose2d_stride_2() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels_in: 1, + channels_out: 1, + kernel_size_1: 2, + kernel_size_2: 2, + padding_1: 0, + padding_2: 0, + padding_out_1: 0, + padding_out_2: 0, + stride_1: 2, + stride_2: 2, + dilation_1: 1, + dilation_2: 1, + groups: 1, + height: 2, + width: 2, + }; + + test.assert_output(TestTensor::from([[[ + [0.0, 0.0, 0.0, 1.0], + [0.0, 0.0, 2.0, 3.0], + [0.0, 2.0, 0.0, 3.0], + [4.0, 6.0, 6.0, 9.0], + ]]])); +} + +#[test] +fn test_conv_transpose2d_dilation_2() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + padding_out_1: 1, + padding_out_2: 1, + stride_1: 1, + stride_2: 1, + dilation_1: 2, + dilation_2: 2, + groups: 1, + height: 2, + width: 2, + }; + + test.assert_output(TestTensor::from([[ + [ + [126., 116., 136., 124., 146.], + [108., 88., 114., 92., 120.], + [156., 140., 166., 148., 176.], + [126., 100., 132., 104., 138.], + [186., 164., 196., 172., 206.], + ], + [ + [217., 189., 227., 197., 237.], + [163., 125., 169., 129., 175.], + [247., 213., 257., 221., 267.], + [181., 137., 187., 141., 193.], + [277., 237., 287., 245., 297.], + ], + ]])); +} + +#[test] +fn test_conv_transpose2d_stride2_out_padding() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + padding_out_1: 1, + padding_out_2: 1, + stride_1: 2, + stride_2: 2, + dilation_1: 1, + dilation_2: 1, + groups: 1, + height: 4, + width: 4, + }; + + test.assert_output(TestTensor::from([[ + [ + [352., 728., 378., 780., 404., 832., 430., 452.], + [784., 1616., 836., 1720., 888., 1824., 940., 992.], + [456., 936., 482., 988., 508., 1040., 534., 564.], + [992., 2032., 1044., 2136., 1096., 2240., 1148., 1216.], + [560., 1144., 586., 1196., 612., 1248., 638., 676.], + [1200., 2448., 1252., 2552., 1304., 2656., 1356., 1440.], + [664., 1352., 690., 1404., 716., 1456., 742., 788.], + [784., 1598., 816., 1662., 848., 1726., 880., 926.], + ], + [ + [497., 1035., 541., 1123., 585., 1211., 629., 651.], + [1145., 2373., 1233., 2549., 1321., 2725., 1409., 1461.], + [673., 1387., 717., 1475., 761., 1563., 805., 835.], + [1497., 3077., 1585., 3253., 1673., 3429., 1761., 1829.], + [849., 1739., 893., 1827., 937., 1915., 981., 1019.], + [1849., 3781., 1937., 3957., 2025., 4133., 2113., 2197.], + [1025., 2091., 1069., 2179., 1113., 2267., 1157., 1203.], + [1145., 2337., 1195., 2437., 1245., 2537., 1295., 1341.], + ], + ]])); +} + +#[test] +fn test_conv_transpose2d_groups_2() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 1, + padding_2: 1, + padding_out_1: 0, + padding_out_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 2, + height: 2, + width: 2, + }; + + test.assert_output(TestTensor::from([[ + [[5., 11.], [23., 29.]], + [[236., 258.], [302., 324.]], + ]])); +} + +#[test] +fn test_conv_transpose2d_groups_different_channels() { + let test = ConvTranspose2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 6, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + padding_out_1: 0, + padding_out_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + groups: 2, + height: 2, + width: 2, + }; + + test.assert_output(TestTensor::from([[ + [ + [0.0000e+00, 0.0000e+00, 1.0000e+00, 2.0000e+00], + [0.0000e+00, 5.0000e+00, 1.1000e+01, 1.1000e+01], + [6.0000e+00, 2.3000e+01, 2.9000e+01, 2.3000e+01], + [1.2000e+01, 3.2000e+01, 3.7000e+01, 2.4000e+01], + ], + [ + [1.0000e+00, 1.0000e+01, 1.1000e+01, 1.2000e+01], + [1.9000e+01, 6.0000e+01, 6.6000e+01, 4.8000e+01], + [2.5000e+01, 7.8000e+01, 8.4000e+01, 6.0000e+01], + [3.1000e+01, 7.8000e+01, 8.3000e+01, 5.2000e+01], + ], + [ + [2.0000e+00, 2.0000e+01, 2.1000e+01, 2.2000e+01], + [3.8000e+01, 1.1500e+02, 1.2100e+02, 8.5000e+01], + [4.4000e+01, 1.3300e+02, 1.3900e+02, 9.7000e+01], + [5.0000e+01, 1.2400e+02, 1.2900e+02, 8.0000e+01], + ], + [ + [1.1100e+02, 2.5000e+02, 2.5900e+02, 1.4800e+02], + [2.8500e+02, 6.3400e+02, 6.5600e+02, 3.6600e+02], + [3.1500e+02, 7.0000e+02, 7.2200e+02, 4.0200e+02], + [2.0100e+02, 4.3800e+02, 4.5100e+02, 2.4800e+02], + ], + [ + [1.4800e+02, 3.3200e+02, 3.4100e+02, 1.9400e+02], + [3.7600e+02, 8.3300e+02, 8.5500e+02, 4.7500e+02], + [4.0600e+02, 8.9900e+02, 9.2100e+02, 5.1100e+02], + [2.5600e+02, 5.5600e+02, 5.6900e+02, 3.1200e+02], + ], + [ + [1.8500e+02, 4.1400e+02, 4.2300e+02, 2.4000e+02], + [4.6700e+02, 1.0320e+03, 1.0540e+03, 5.8400e+02], + [4.9700e+02, 1.0980e+03, 1.1200e+03, 6.2000e+02], + [3.1100e+02, 6.7400e+02, 6.8700e+02, 3.7600e+02], + ], + ]])); +} + +struct ConvTranspose2dTestCase { + batch_size: usize, + channels_in: usize, + channels_out: usize, + kernel_size_1: usize, + kernel_size_2: usize, + padding_1: usize, + padding_2: usize, + padding_out_1: usize, + padding_out_2: usize, + stride_1: usize, + stride_2: usize, + dilation_1: usize, + dilation_2: usize, + groups: usize, + height: usize, + width: usize, +} + +impl ConvTranspose2dTestCase { + fn assert_output(self, y: TestTensor<4>) { + let shape_x = Shape::new([self.batch_size, self.channels_in, self.height, self.width]); + let shape_weights = Shape::new([ + self.channels_in, + self.channels_out / self.groups, + self.kernel_size_1, + self.kernel_size_2, + ]); + let device = Default::default(); + let weights = TestTensor::from( + TestTensorInt::arange(0..shape_weights.num_elements() as i64, &device) + .reshape::<4, _>(shape_weights) + .into_data(), + ); + let bias = TestTensor::from( + TestTensorInt::arange(0..self.channels_out as i64, &device).into_data(), + ); + let x = TestTensor::from( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<4, _>(shape_x) + .into_data(), + ); + let output = conv_transpose2d( + x, + weights, + Some(bias), + ConvTransposeOptions::new( + [self.stride_1, self.stride_2], + [self.padding_1, self.padding_2], + [self.padding_out_1, self.padding_out_2], + [self.dilation_1, self.dilation_2], + self.groups, + ), + ); + + y.into_data() + .assert_approx_eq::(&output.into_data(), Tolerance::rel_abs(1e-1, 0.01)); + } +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/conv_transpose3d.rs b/crates/burn-backend-tests/tests/tensor/float/module/conv_transpose3d.rs new file mode 100644 index 0000000..49911a6 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/conv_transpose3d.rs @@ -0,0 +1,749 @@ +use super::*; +use burn_tensor::Shape; +use burn_tensor::Tolerance; +use burn_tensor::module::conv_transpose3d; +use burn_tensor::ops::ConvTransposeOptions; + +#[test] +fn test_conv_transpose3d_simple_1() { + let test = ConvTranspose3dTestCase { + batch_size: 1, + channels_in: 1, + channels_out: 1, + kernel_size_1: 3, + kernel_size_2: 3, + kernel_size_3: 3, + padding_1: 1, + padding_2: 1, + padding_3: 1, + padding_out_1: 0, + padding_out_2: 0, + padding_out_3: 0, + stride_1: 1, + stride_2: 1, + stride_3: 1, + dilation_1: 1, + dilation_2: 1, + dilation_3: 1, + groups: 1, + depth: 2, + height: 2, + width: 2, + }; + + test.assert_output(TestTensor::from([[[ + [[96., 124.], [180., 208.]], + [[348., 376.], [432., 460.]], + ]]])); +} +#[test] +fn test_conv_transpose3d_simple_2() { + let test = ConvTranspose3dTestCase { + batch_size: 1, + channels_in: 3, + channels_out: 3, + kernel_size_1: 3, + kernel_size_2: 3, + kernel_size_3: 3, + padding_1: 1, + padding_2: 1, + padding_3: 1, + padding_out_1: 0, + padding_out_2: 0, + padding_out_3: 0, + stride_1: 1, + stride_2: 1, + stride_3: 1, + dilation_1: 1, + dilation_2: 1, + dilation_3: 1, + groups: 1, + depth: 4, + height: 4, + width: 4, + }; + + test.assert_output(TestTensor::from([[ + [ + [ + [238452., 360588., 363756., 244488.], + [367929., 556353., 561186., 377163.], + [380745., 575685., 580518., 390123.], + [261192., 394896., 398172., 267564.], + ], + [ + [394083., 595827., 600822., 403749.], + [607635., 918648., 926262., 622404.], + [627831., 949104., 956718., 642816.], + [430353., 650529., 655686., 440523.], + ], + [ + [447075., 675747., 680742., 457317.], + [688419., 1040472., 1048086., 704052.], + [708615., 1070928., 1078542., 724464.], + [485073., 733041., 738198., 495819.], + ], + [ + [328656., 496632., 500124., 335892.], + [505611., 763983., 769302., 516645.], + [519723., 785259., 790578., 530901.], + [355428., 536988., 540588., 363000.], + ], + ], + [ + [ + [286729., 433489., 437629., 294061.], + [442288., 668620., 674911., 453466.], + [458992., 693784., 700075., 470314.], + [314653., 475573., 479821., 322321.], + ], + [ + [474274., 716842., 723295., 485884.], + [730837., 1104544., 1114345., 748522.], + [756865., 1143748., 1153549., 774766.], + [518320., 783208., 789823., 530434.], + ], + [ + [542818., 820090., 826543., 555004.], + [834949., 1261360., 1271161., 853498.], + [860977., 1300564., 1310365., 879742.], + [588592., 889048., 895663., 601282.], + ], + [ + [397669., 600637., 605101., 406201.], + [611074., 922906., 929683., 624052.], + [629074., 950014., 956791., 642196.], + [429625., 648769., 653341., 438493.], + ], + ], + [ + [ + [335006., 506390., 511502., 343634.], + [516647., 780887., 788636., 529769.], + [537239., 811883., 819632., 550505.], + [368114., 556250., 561470., 377078.], + ], + [ + [554465., 837857., 845768., 568019.], + [854039., 1290440., 1302428., 874640.], + [885899., 1338392., 1350380., 906716.], + [606287., 915887., 923960., 620345.], + ], + [ + [638561., 964433., 972344., 652691.], + [981479., 1482248., 1494236., 1002944.], + [1013339., 1530200., 1542188., 1035020.], + [692111., 1045055., 1053128., 706745.], + ], + [ + [466682., 704642., 710078., 476510.], + [716537., 1081829., 1090064., 731459.], + [738425., 1114769., 1123004., 753491.], + [503822., 760550., 766094., 513986.], + ], + ], + ]])); +} + +#[test] +fn test_conv_transpose3d_stride_2() { + let test = ConvTranspose3dTestCase { + batch_size: 1, + channels_in: 1, + channels_out: 1, + kernel_size_1: 2, + kernel_size_2: 2, + kernel_size_3: 2, + padding_1: 0, + padding_2: 0, + padding_3: 0, + padding_out_1: 0, + padding_out_2: 0, + padding_out_3: 0, + stride_1: 2, + stride_2: 2, + stride_3: 2, + dilation_1: 1, + dilation_2: 1, + dilation_3: 1, + groups: 1, + depth: 2, + height: 2, + width: 2, + }; + + test.assert_output(TestTensor::from([[[ + [ + [0., 0., 0., 1.], + [0., 0., 2., 3.], + [0., 2., 0., 3.], + [4., 6., 6., 9.], + ], + [ + [0., 0., 4., 5.], + [0., 0., 6., 7.], + [8., 10., 12., 15.], + [12., 14., 18., 21.], + ], + [ + [0., 4., 0., 5.], + [8., 12., 10., 15.], + [0., 6., 0., 7.], + [12., 18., 14., 21.], + ], + [ + [16., 20., 20., 25.], + [24., 28., 30., 35.], + [24., 30., 28., 35.], + [36., 42., 42., 49.], + ], + ]]])); +} + +#[test] +fn test_conv_transpose3d_dilation_2() { + let test = ConvTranspose3dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + kernel_size_3: 3, + padding_1: 1, + padding_2: 1, + padding_3: 1, + padding_out_1: 1, + padding_out_2: 1, + padding_out_3: 1, + stride_1: 1, + stride_2: 1, + stride_3: 1, + dilation_1: 2, + dilation_2: 2, + dilation_3: 2, + groups: 1, + depth: 2, + height: 2, + width: 2, + }; + + test.assert_output(TestTensor::from([[ + [ + [ + [810., 776., 832., 796., 854.], + [756., 712., 774., 728., 792.], + [876., 836., 898., 856., 920.], + [810., 760., 828., 776., 846.], + [942., 896., 964., 916., 986.], + ], + [ + [720., 660., 734., 672., 748.], + [606., 536., 616., 544., 626.], + [762., 696., 776., 708., 790.], + [636., 560., 646., 568., 656.], + [804., 732., 818., 744., 832.], + ], + [ + [1008., 956., 1030., 976., 1052.], + [918., 856., 936., 872., 954.], + [1074., 1016., 1096., 1036., 1118.], + [972., 904., 990., 920., 1008.], + [1140., 1076., 1162., 1096., 1184.], + ], + [ + [846., 768., 860., 780., 874.], + [696., 608., 706., 616., 716.], + [888., 804., 902., 816., 916.], + [726., 632., 736., 640., 746.], + [930., 840., 944., 852., 958.], + ], + [ + [1206., 1136., 1228., 1156., 1250.], + [1080., 1000., 1098., 1016., 1116.], + [1272., 1196., 1294., 1216., 1316.], + [1134., 1048., 1152., 1064., 1170.], + [1338., 1256., 1360., 1276., 1382.], + ], + ], + [ + [ + [1405., 1317., 1427., 1337., 1449.], + [1243., 1145., 1261., 1161., 1279.], + [1471., 1377., 1493., 1397., 1515.], + [1297., 1193., 1315., 1209., 1333.], + [1537., 1437., 1559., 1457., 1581.], + ], + [ + [1099., 985., 1113., 997., 1127.], + [877., 753., 887., 761., 897.], + [1141., 1021., 1155., 1033., 1169.], + [907., 777., 917., 785., 927.], + [1183., 1057., 1197., 1069., 1211.], + ], + [ + [1603., 1497., 1625., 1517., 1647.], + [1405., 1289., 1423., 1305., 1441.], + [1669., 1557., 1691., 1577., 1713.], + [1459., 1337., 1477., 1353., 1495.], + [1735., 1617., 1757., 1637., 1779.], + ], + [ + [1225., 1093., 1239., 1105., 1253.], + [967., 825., 977., 833., 987.], + [1267., 1129., 1281., 1141., 1295.], + [997., 849., 1007., 857., 1017.], + [1309., 1165., 1323., 1177., 1337.], + ], + [ + [1801., 1677., 1823., 1697., 1845.], + [1567., 1433., 1585., 1449., 1603.], + [1867., 1737., 1889., 1757., 1911.], + [1621., 1481., 1639., 1497., 1657.], + [1933., 1797., 1955., 1817., 1977.], + ], + ], + ]])); +} + +#[test] +fn test_conv_transpose3d_stride2_out_padding() { + let test = ConvTranspose3dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + kernel_size_3: 3, + padding_1: 1, + padding_2: 1, + padding_3: 1, + padding_out_1: 1, + padding_out_2: 1, + padding_out_3: 1, + stride_1: 2, + stride_2: 2, + stride_3: 2, + dilation_1: 1, + dilation_2: 1, + dilation_3: 1, + groups: 1, + depth: 2, + height: 4, + width: 4, + }; + + test.assert_output(TestTensor::from([[ + [ + [ + [2144., 4366., 2224., 4526., 2304., 4686., 2384., 2422.], + [4584., 9324., 4744., 9644., 4904., 9964., 5064., 5148.], + [2464., 5006., 2544., 5166., 2624., 5326., 2704., 2750.], + [5224., 10604., 5384., 10924., 5544., 11244., 5704., 5804.], + [2784., 5646., 2864., 5806., 2944., 5966., 3024., 3078.], + [5864., 11884., 6024., 12204., 6184., 12524., 6344., 6460.], + [3104., 6286., 3184., 6446., 3264., 6606., 3344., 3406.], + [3272., 6628., 3358., 6800., 3444., 6972., 3530., 3592.], + ], + [ + [5280., 10716., 5440., 11036., 5600., 11356., 5760., 5868.], + [ + 11152., 22616., 11472., 23256., 11792., 23896., 12112., 12344., + ], + [5920., 11996., 6080., 12316., 6240., 12636., 6400., 6524.], + [ + 12432., 25176., 12752., 25816., 13072., 26456., 13392., 13656., + ], + [6560., 13276., 6720., 13596., 6880., 13916., 7040., 7180.], + [ + 13712., 27736., 14032., 28376., 14352., 29016., 14672., 14968., + ], + [7200., 14556., 7360., 14876., 7520., 15196., 7680., 7836.], + [7632., 15432., 7804., 15776., 7976., 16120., 8148., 8304.], + ], + [ + [3424., 6926., 3504., 7086., 3584., 7246., 3664., 3734.], + [7144., 14444., 7304., 14764., 7464., 15084., 7624., 7772.], + [3744., 7566., 3824., 7726., 3904., 7886., 3984., 4062.], + [7784., 15724., 7944., 16044., 8104., 16364., 8264., 8428.], + [4064., 8206., 4144., 8366., 4224., 8526., 4304., 4390.], + [8424., 17004., 8584., 17324., 8744., 17644., 8904., 9084.], + [4384., 8846., 4464., 9006., 4544., 9166., 4624., 4718.], + [4648., 9380., 4734., 9552., 4820., 9724., 4906., 5000.], + ], + [ + [4000., 8096., 4098., 8292., 4196., 8488., 4294., 4364.], + [8368., 16928., 8564., 17320., 8760., 17712., 8956., 9104.], + [4392., 8880., 4490., 9076., 4588., 9272., 4686., 4764.], + [9152., 18496., 9348., 18888., 9544., 19280., 9740., 9904.], + [4784., 9664., 4882., 9860., 4980., 10056., 5078., 5164.], + [ + 9936., 20064., 10132., 20456., 10328., 20848., 10524., 10704., + ], + [5176., 10448., 5274., 10644., 5372., 10840., 5470., 5564.], + [5440., 10982., 5544., 11190., 5648., 11398., 5752., 5846.], + ], + ], + [ + [ + [3009., 6149., 3143., 6417., 3277., 6685., 3411., 3449.], + [6529., 13321., 6797., 13857., 7065., 14393., 7333., 7417.], + [3545., 7221., 3679., 7489., 3813., 7757., 3947., 3993.], + [7601., 15465., 7869., 16001., 8137., 16537., 8405., 8505.], + [4081., 8293., 4215., 8561., 4349., 8829., 4483., 4537.], + [8673., 17609., 8941., 18145., 9209., 18681., 9477., 9593.], + [4617., 9365., 4751., 9633., 4885., 9901., 5019., 5081.], + [4785., 9707., 4925., 9987., 5065., 10267., 5205., 5267.], + ], + [ + [7873., 16009., 8141., 16545., 8409., 17081., 8677., 8785.], + [ + 16769., 34065., 17305., 35137., 17841., 36209., 18377., 18609., + ], + [8945., 18153., 9213., 18689., 9481., 19225., 9749., 9873.], + [ + 18913., 38353., 19449., 39425., 19985., 40497., 20521., 20785., + ], + [ + 10017., 20297., 10285., 20833., 10553., 21369., 10821., 10961., + ], + [ + 21057., 42641., 21593., 43713., 22129., 44785., 22665., 22961., + ], + [ + 11089., 22441., 11357., 22977., 11625., 23513., 11893., 12049., + ], + [ + 11521., 23317., 11801., 23877., 12081., 24437., 12361., 12517., + ], + ], + [ + [5153., 10437., 5287., 10705., 5421., 10973., 5555., 5625.], + [ + 10817., 21897., 11085., 22433., 11353., 22969., 11621., 11769., + ], + [5689., 11509., 5823., 11777., 5957., 12045., 6091., 6169.], + [ + 11889., 24041., 12157., 24577., 12425., 25113., 12693., 12857., + ], + [6225., 12581., 6359., 12849., 6493., 13117., 6627., 6713.], + [ + 12961., 26185., 13229., 26721., 13497., 27257., 13765., 13945., + ], + [6761., 13653., 6895., 13921., 7029., 14189., 7163., 7257.], + [7025., 14187., 7165., 14467., 7305., 14747., 7445., 7539.], + ], + [ + [5729., 11607., 5881., 11911., 6033., 12215., 6185., 6255.], + [ + 12041., 24381., 12345., 24989., 12649., 25597., 12953., 13101., + ], + [6337., 12823., 6489., 13127., 6641., 13431., 6793., 6871.], + [ + 13257., 26813., 13561., 27421., 13865., 28029., 14169., 14333., + ], + [6945., 14039., 7097., 14343., 7249., 14647., 7401., 7487.], + [ + 14473., 29245., 14777., 29853., 15081., 30461., 15385., 15565., + ], + [7553., 15255., 7705., 15559., 7857., 15863., 8009., 8103.], + [7817., 15789., 7975., 16105., 8133., 16421., 8291., 8385.], + ], + ], + ]])); +} + +#[test] +fn test_conv_transpose3d_groups_2() { + let test = ConvTranspose3dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 2, + kernel_size_1: 3, + kernel_size_2: 3, + kernel_size_3: 3, + padding_1: 1, + padding_2: 1, + padding_3: 1, + padding_out_1: 0, + padding_out_2: 0, + padding_out_3: 0, + stride_1: 1, + stride_2: 1, + stride_3: 1, + dilation_1: 1, + dilation_2: 1, + dilation_3: 1, + groups: 2, + depth: 2, + height: 2, + width: 2, + }; + + test.assert_output(TestTensor::from([[ + [[[96., 124.], [180., 208.]], [[348., 376.], [432., 460.]]], + [ + [[2997., 3089.], [3273., 3365.]], + [[3825., 3917.], [4101., 4193.]], + ], + ]])); +} + +#[test] +fn test_conv_transpose3d_groups_different_channels() { + let test = ConvTranspose3dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 6, + kernel_size_1: 3, + kernel_size_2: 3, + kernel_size_3: 3, + padding_1: 0, + padding_2: 0, + padding_3: 0, + padding_out_1: 0, + padding_out_2: 0, + padding_out_3: 0, + stride_1: 1, + stride_2: 1, + stride_3: 1, + dilation_1: 1, + dilation_2: 1, + dilation_3: 1, + groups: 2, + depth: 2, + height: 2, + width: 2, + }; + + test.assert_output(TestTensor::from([[ + [ + [ + [0., 0., 1., 2.], + [0., 5., 11., 11.], + [6., 23., 29., 23.], + [12., 32., 37., 24.], + ], + [ + [0., 13., 23., 21.], + [30., 96., 124., 86.], + [66., 180., 208., 134.], + [66., 161., 179., 107.], + ], + [ + [36., 103., 113., 75.], + [138., 348., 376., 230.], + [174., 432., 460., 278.], + [138., 323., 341., 197.], + ], + [ + [72., 166., 175., 100.], + [192., 433., 455., 255.], + [222., 499., 521., 291.], + [144., 318., 331., 182.], + ], + ], + [ + [ + [1., 28., 29., 30.], + [55., 168., 174., 120.], + [61., 186., 192., 132.], + [67., 168., 173., 106.], + ], + [ + [109., 284., 294., 184.], + [355., 853., 881., 519.], + [391., 937., 965., 567.], + [283., 648., 666., 378.], + ], + [ + [145., 374., 384., 238.], + [463., 1105., 1133., 663.], + [499., 1189., 1217., 711.], + [355., 810., 828., 468.], + ], + [ + [181., 410., 419., 236.], + [463., 1028., 1050., 580.], + [493., 1094., 1116., 616.], + [307., 670., 683., 372.], + ], + ], + [ + [ + [2., 56., 57., 58.], + [110., 331., 337., 229.], + [116., 349., 355., 241.], + [122., 304., 309., 188.], + ], + [ + [218., 555., 565., 347.], + [680., 1610., 1638., 952.], + [716., 1694., 1722., 1000.], + [500., 1135., 1153., 649.], + ], + [ + [254., 645., 655., 401.], + [788., 1862., 1890., 1096.], + [824., 1946., 1974., 1144.], + [572., 1297., 1315., 739.], + ], + [ + [290., 654., 663., 372.], + [734., 1623., 1645., 905.], + [764., 1689., 1711., 941.], + [470., 1022., 1035., 562.], + ], + ], + [ + [ + [651., 1388., 1405., 750.], + [1485., 3150., 3188., 1690.], + [1539., 3264., 3302., 1750.], + [873., 1840., 1861., 982.], + ], + [ + [1695., 3578., 3620., 1910.], + [3789., 7967., 8059., 4233.], + [3921., 8243., 8335., 4377.], + [2181., 4566., 4616., 2416.], + ], + [ + [1875., 3956., 3998., 2108.], + [4185., 8795., 8887., 4665.], + [4317., 9071., 9163., 4809.], + [2397., 5016., 5066., 2650.], + ], + [ + [1191., 2490., 2515., 1316.], + [2613., 5450., 5504., 2870.], + [2691., 5612., 5666., 2954.], + [1473., 3062., 3091., 1608.], + ], + ], + [ + [ + [868., 1848., 1865., 994.], + [1972., 4177., 4215., 2231.], + [2026., 4291., 4329., 2291.], + [1144., 2408., 2429., 1280.], + ], + [ + [2236., 4713., 4755., 2505.], + [4978., 10452., 10544., 5530.], + [5110., 10728., 10820., 5674.], + [2830., 5917., 5967., 3119.], + ], + [ + [2416., 5091., 5133., 2703.], + [5374., 11280., 11372., 5962.], + [5506., 11556., 11648., 6106.], + [3046., 6367., 6417., 3353.], + ], + [ + [1516., 3166., 3191., 1668.], + [3316., 6909., 6963., 3627.], + [3394., 7071., 7125., 3711.], + [1852., 3846., 3875., 2014.], + ], + ], + [ + [ + [1085., 2308., 2325., 1238.], + [2459., 5204., 5242., 2772.], + [2513., 5318., 5356., 2832.], + [1415., 2976., 2997., 1578.], + ], + [ + [2777., 5848., 5890., 3100.], + [6167., 12937., 13029., 6827.], + [6299., 13213., 13305., 6971.], + [3479., 7268., 7318., 3822.], + ], + [ + [2957., 6226., 6268., 3298.], + [6563., 13765., 13857., 7259.], + [6695., 14041., 14133., 7403.], + [3695., 7718., 7768., 4056.], + ], + [ + [1841., 3842., 3867., 2020.], + [4019., 8368., 8422., 4384.], + [4097., 8530., 8584., 4468.], + [2231., 4630., 4659., 2420.], + ], + ], + ]])); +} + +struct ConvTranspose3dTestCase { + batch_size: usize, + channels_in: usize, + channels_out: usize, + kernel_size_1: usize, + kernel_size_2: usize, + kernel_size_3: usize, + padding_1: usize, + padding_2: usize, + padding_3: usize, + padding_out_1: usize, + padding_out_2: usize, + padding_out_3: usize, + stride_1: usize, + stride_2: usize, + stride_3: usize, + dilation_1: usize, + dilation_2: usize, + dilation_3: usize, + groups: usize, + depth: usize, + height: usize, + width: usize, +} + +impl ConvTranspose3dTestCase { + fn assert_output(self, y: TestTensor<5>) { + let shape_x = Shape::new([ + self.batch_size, + self.channels_in, + self.depth, + self.height, + self.width, + ]); + let shape_weights = Shape::new([ + self.channels_in, + self.channels_out / self.groups, + self.kernel_size_1, + self.kernel_size_2, + self.kernel_size_3, + ]); + let device = Default::default(); + let weights = TestTensor::from( + TestTensorInt::arange(0..shape_weights.num_elements() as i64, &device) + .reshape::<5, _>(shape_weights) + .into_data(), + ); + let bias = TestTensor::from( + TestTensorInt::arange(0..self.channels_out as i64, &device).into_data(), + ); + let x = TestTensor::from( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<5, _>(shape_x) + .into_data(), + ); + let output = conv_transpose3d( + x, + weights, + Some(bias), + ConvTransposeOptions::new( + [self.stride_1, self.stride_2, self.stride_3], + [self.padding_1, self.padding_2, self.padding_3], + [self.padding_out_1, self.padding_out_2, self.padding_out_3], + [self.dilation_1, self.dilation_2, self.dilation_3], + self.groups, + ), + ); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); + } +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/ctc.rs b/crates/burn-backend-tests/tests/tensor/float/module/ctc.rs new file mode 100644 index 0000000..1b5aa73 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/ctc.rs @@ -0,0 +1,333 @@ +#![allow(clippy::excessive_precision)] + +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::activation::log_softmax; +use burn_tensor::module::ctc_loss; + +#[test] +fn test_ctc_loss_unreachable_target_is_inf() { + // T=2, N=1, C=3, target=[1, 1]. CTC requires at least `target_length + // + repeats = 3` steps to produce this target, but input_length is 2, + // so no valid alignment exists and the loss must be +inf. burn-nn's + // `zero_infinity` relies on `is_inf` to detect this case, so fused + // backend kernels must preserve the +inf (not collapse to a large + // finite number). + let device = Default::default(); + let log_probs = TestTensor::<3>::full([2, 1, 3], (1.0f32 / 3.0).ln(), &device); + let targets = TestTensorInt::<2>::from([[1_i64, 1]]); + let input_lengths = TestTensorInt::<1>::from([2_i64]); + let target_lengths = TestTensorInt::<1>::from([2_i64]); + + let loss = ctc_loss(log_probs, targets, input_lengths, target_lengths, 0); + let loss_data = loss.into_data(); + let value = loss_data.iter::().next().unwrap(); + assert!( + value.is_infinite() && value > 0.0, + "Expected +inf for an unreachable target, got {value}" + ); +} + +#[test] +fn test_ctc_loss_empty_target() { + // T=3, N=1, C=2, blank=0, target_length=0, uniform P = 1/2. + // With an empty target, the only valid alignment is all blanks. + // Loss = -ln((1/2)^3) = 3 * ln(2) ~ 2.0794. + // + // Uses a [1, 1] targets shape rather than [1, 0] since not every backend + // supports zero-sized dimensions. The target slot at [0, 0] is never read + // because target_lengths[0] == 0. + let log_probs = TestTensor::<3>::full([3, 1, 2], (0.5f32).ln(), &Default::default()); + let targets = TestTensorInt::<2>::from([[0]]); + let input_lengths = TestTensorInt::<1>::from([3]); + let target_lengths = TestTensorInt::<1>::from([0]); + + let loss = ctc_loss(log_probs, targets, input_lengths, target_lengths, 0); + + let expected = burn_tensor::TensorData::from([3.0f32 * 2.0f32.ln()]); + loss.into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(1e-3, 1e-3)); +} + +#[test] +fn test_ctc_loss_uniform() { + // T=3, N=1, C=2, blank=0, target=[1, 1], uniform P = 1/2. + // Only valid path is (1, 0, 1) with prob (1/2)^3. + // Loss = -ln(1/8) = 3 * ln(2) ~ 2.0794 + let log_probs = TestTensor::<3>::full([3, 1, 2], (0.5f32).ln(), &Default::default()); + let targets = TestTensorInt::<2>::from([[1, 1]]); + let input_lengths = TestTensorInt::<1>::from([3]); + let target_lengths = TestTensorInt::<1>::from([2]); + + let loss = ctc_loss(log_probs, targets, input_lengths, target_lengths, 0); + + let expected = burn_tensor::TensorData::from([3.0f32 * 2.0f32.ln()]); + loss.into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(1e-3, 1e-3)); +} + +#[test] +fn test_ctc_loss_long_sequence() { + // T=400, N=2, C=36, target_lengths=[60, 40]. Exercises the alpha + // recursion at training-realistic sequence lengths where numerical + // issues (shared-memory overflow, log-sum-exp accumulation drift) + // would surface but don't show in the short T=5..12 tests. + let t: usize = 400; + let n: usize = 2; + let c: usize = 36; + let mut data = Vec::with_capacity(t * n * c); + for ti in 0..t { + for ni in 0..n { + for ci in 0..c { + data.push(((ti * 7 + ni * 13 + ci * 3) as f32 * 0.1).sin()); + } + } + } + let logits = TestTensor::<3>::from(burn_tensor::TensorData::new(data, [t, n, c])); + let log_probs = log_softmax(logits, 2); + + // Targets: distinct labels for each batch element, no consecutive repeats. + let tgt_a: Vec = (0..60).map(|i| (i % 35 + 1) as i64).collect(); + let tgt_b: Vec = (0..60) + .map(|i| if i < 40 { (i % 35 + 1) as i64 } else { 0 }) + .collect(); + let mut tgt_flat = tgt_a; + tgt_flat.extend(tgt_b); + let targets = TestTensorInt::<2>::from(burn_tensor::TensorData::new(tgt_flat, [n, 60])); + let input_lengths = TestTensorInt::<1>::from([400, 400]); + let target_lengths = TestTensorInt::<1>::from([60, 40]); + + let loss = ctc_loss(log_probs, targets, input_lengths, target_lengths, 0); + // `iter::` converts from the backend's native dtype (f16, f32, ...) + // to f32 for comparison; `to_vec::` would require exact dtype match. + let loss_data: Vec = loss.into_data().iter::().collect(); + + // We don't have a PyTorch reference for this exact input, but the loss + // must be positive (it's -log P where 0 < P <= 1) and finite. + for (i, &v) in loss_data.iter().enumerate() { + assert!( + v.is_finite() && v > 0.0, + "sample {i}: expected positive finite loss, got {v}" + ); + } + // With 36 classes and random-ish logits, loss should be in the hundreds. + // An off-by-orders-of-magnitude result (< 1 or > 10000) indicates a bug. + for (i, &v) in loss_data.iter().enumerate() { + assert!( + v > 10.0 && v < 10000.0, + "sample {i}: loss {v} outside plausible range [10, 10000]" + ); + } +} + +#[test] +fn test_ctc_loss_long_mixed_input_lengths() { + // T=494, N=2, C=36, input_lengths=[494, 390], target_lengths=[64, 43]. + // The second sample has 104 frames of padding (input_length < T). + // Reproduces the exact dimensions of a real training batch that gave + // incorrect results on the cubecl kernel. + let t: usize = 494; + let n: usize = 2; + let c: usize = 36; + let mut data = Vec::with_capacity(t * n * c); + for ti in 0..t { + for ni in 0..n { + for ci in 0..c { + data.push(((ti * 7 + ni * 13 + ci * 3) as f32 * 0.1).sin()); + } + } + } + let logits = TestTensor::<3>::from(burn_tensor::TensorData::new(data, [t, n, c])); + let log_probs = log_softmax(logits, 2); + + let max_tgt = 64; + let tgt_a: Vec = (0..max_tgt).map(|i| (i % 35 + 1) as i64).collect(); + let tgt_b: Vec = (0..max_tgt) + .map(|i| if i < 43 { (i % 35 + 1) as i64 } else { 0 }) + .collect(); + let mut tgt_flat = tgt_a; + tgt_flat.extend(tgt_b); + let targets = TestTensorInt::<2>::from(burn_tensor::TensorData::new(tgt_flat, [n, max_tgt])); + let input_lengths = TestTensorInt::<1>::from([494, 390]); + let target_lengths = TestTensorInt::<1>::from([64, 43]); + + let loss = ctc_loss(log_probs, targets, input_lengths, target_lengths, 0); + let loss_data: Vec = loss.into_data().iter::().collect(); + + for (i, &v) in loss_data.iter().enumerate() { + assert!( + v.is_finite() && v > 0.0, + "sample {i}: expected positive finite loss, got {v}" + ); + assert!( + v > 10.0 && v < 10000.0, + "sample {i}: loss {v} outside plausible range [10, 10000]" + ); + } +} + +#[test] +fn test_ctc_loss_narrowed_input() { + // CTC on a narrowed (non-contiguous, offset != 0) log_probs tensor. + // Reproduces the pattern: model outputs [T_total, B, C], training + // code narrows to [T_total - warmup, B, C] before passing to CTC. + // If the kernel ignores the tensor's base offset, it reads the wrong + // data and produces garbage loss. + let t_full: usize = 10; + let warmup: usize = 3; + let t_eff: usize = t_full - warmup; + let n: usize = 1; + let c: usize = 4; + + let mut data = Vec::with_capacity(t_full * n * c); + for ti in 0..t_full { + for ni in 0..n { + for ci in 0..c { + data.push(((ti * 7 + ni * 13 + ci * 3) as f32 * 0.1).sin()); + } + } + } + let full_logits = TestTensor::<3>::from(burn_tensor::TensorData::new(data, [t_full, n, c])); + let full_log_probs = log_softmax(full_logits, 2); + + // Narrow: skip the first `warmup` frames (simulating warmup slice). + let narrowed = full_log_probs.clone().narrow(0, warmup, t_eff); + + // Reference: manually extract the same slice as a contiguous tensor. + // Collect in the backend's native dtype (f32 or f16) to preserve + // precision when rebuilding the tensor. + let reference_data: Vec = full_log_probs + .clone() + .slice([warmup..t_full, 0..n, 0..c]) + .into_data() + .iter::() + .collect(); + let reference = + TestTensor::<3>::from(burn_tensor::TensorData::new(reference_data, [t_eff, n, c])); + + let targets = TestTensorInt::<2>::from([[1, 2]]); + let input_lengths = TestTensorInt::<1>::from([t_eff as i64]); + let target_lengths = TestTensorInt::<1>::from([2]); + + let loss_narrowed = ctc_loss( + narrowed, + targets.clone(), + input_lengths.clone(), + target_lengths.clone(), + 0, + ); + let loss_reference = ctc_loss(reference, targets, input_lengths, target_lengths, 0); + + let v_narrowed: Vec = loss_narrowed.into_data().iter::().collect(); + let v_reference: Vec = loss_reference.into_data().iter::().collect(); + + assert!( + (v_narrowed[0] - v_reference[0]).abs() < 1e-3, + "narrowed loss ({}) diverges from contiguous reference ({})", + v_narrowed[0], + v_reference[0], + ); +} + +/// Training pipelines typically produce log_probs with a `[B, T, C] -> +/// swap_dims(0, 1) -> [T, B, C] -> narrow(0, warmup, ...)` chain before +/// handing the tensor to CTC. `swap_dims` makes strides non-monotonic, +/// `narrow` introduces a non-zero base offset, and on fused backends the +/// combination has historically produced garbage from manual stride +/// indexing. This guards against that regression with a shape-compatible +/// fixture (B=2 so swap_dims actually permutes strides non-trivially). +#[test] +fn test_ctc_loss_swap_dims_then_narrow() { + let b: usize = 2; + let t_full: usize = 8; + let warmup: usize = 2; + let t_eff: usize = t_full - warmup; + let c: usize = 4; + + // Build logits in [B, T, C] layout (as a typical model output would). + let mut data = Vec::with_capacity(b * t_full * c); + for bi in 0..b { + for ti in 0..t_full { + for ci in 0..c { + data.push(((bi * 31 + ti * 7 + ci * 3) as f32 * 0.1).sin()); + } + } + } + let logits_btc = TestTensor::<3>::from(burn_tensor::TensorData::new(data, [b, t_full, c])); + let log_probs_btc = log_softmax(logits_btc, 2); + + // Permute to [T, B, C] (non-monotonic strides) then narrow away warmup. + let log_probs_tbc = log_probs_btc.clone().swap_dims(0, 1); + let narrowed = log_probs_tbc.narrow(0, warmup, t_eff); + + // Reference: materialize the same slice as a contiguous [T, B, C] tensor + // in the backend's native dtype. + let reference_data: Vec = log_probs_btc + .swap_dims(0, 1) + .slice([warmup..t_full, 0..b, 0..c]) + .into_data() + .iter::() + .collect(); + let reference = + TestTensor::<3>::from(burn_tensor::TensorData::new(reference_data, [t_eff, b, c])); + + let targets = TestTensorInt::<2>::from([[1_i64, 2], [2, 3]]); + let input_lengths = TestTensorInt::<1>::from([t_eff as i64, t_eff as i64]); + let target_lengths = TestTensorInt::<1>::from([2_i64, 2]); + + let loss_narrowed = ctc_loss( + narrowed, + targets.clone(), + input_lengths.clone(), + target_lengths.clone(), + 0, + ); + let loss_reference = ctc_loss(reference, targets, input_lengths, target_lengths, 0); + + let v_narrowed: Vec = loss_narrowed.into_data().iter::().collect(); + let v_reference: Vec = loss_reference.into_data().iter::().collect(); + + for i in 0..b { + assert!( + (v_narrowed[i] - v_reference[i]).abs() < 1e-3, + "sample {i}: swap_dims+narrow loss ({}) diverges from contiguous reference ({})", + v_narrowed[i], + v_reference[i], + ); + } +} + +#[test] +fn test_ctc_loss_matches_pytorch() { + // T=5, N=3, C=4, deterministic logits via sin((t*7 + n*13 + c*3) * 0.1). + // Expected per-sample losses computed by PyTorch's nn.functional.ctc_loss. + let t_size: usize = 5; + let n_size: usize = 3; + let c_size: usize = 4; + + let mut data = Vec::with_capacity(t_size * n_size * c_size); + for t in 0..t_size { + for n in 0..n_size { + for c in 0..c_size { + data.push(((t * 7 + n * 13 + c * 3) as f32 * 0.1).sin()); + } + } + } + let logits = + TestTensor::<3>::from(burn_tensor::TensorData::new(data, [t_size, n_size, c_size])); + let log_probs = log_softmax(logits, 2); + + let targets = TestTensorInt::<2>::from([[1, 2, 0], [1, 0, 0], [3, 2, 1]]); + let input_lengths = TestTensorInt::<1>::from([5, 5, 5]); + let target_lengths = TestTensorInt::<1>::from([2, 1, 3]); + + let loss = ctc_loss(log_probs, targets, input_lengths, target_lengths, 0); + + let expected = burn_tensor::TensorData::from([ + 3.5236570835113525f32, + 3.495313882827759, + 4.262677192687988, + ]); + loss.into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(1e-3, 1e-3)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/deform_conv2d.rs b/crates/burn-backend-tests/tests/tensor/float/module/deform_conv2d.rs new file mode 100644 index 0000000..e32a1c9 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/deform_conv2d.rs @@ -0,0 +1,438 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::module::deform_conv2d; +use burn_tensor::ops::DeformConvOptions; +use burn_tensor::{Shape, Tensor}; + +#[test] +fn test_deform_conv2d_simple() { + let test = DeformConv2dTestCase { + batch_size: 1, + channels_in: 3, + channels_out: 5, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + weight_groups: 1, + offset_groups: 1, + height: 4, + width: 4, + }; + + test.assert_output(TestTensor::<4>::from([[ + [[0.9074, 0.6387], [0.5160, 0.4196]], + [[2.4259, 1.8008], [1.5449, 1.3112]], + [[3.9444, 2.9629], [2.5738, 2.2027]], + [[5.4629, 4.1250], [3.6027, 3.0943]], + [[6.9814, 5.2871], [4.6316, 3.9859]], + ]])); +} + +#[test] +fn test_deform_conv2d_batched() { + let test = DeformConv2dTestCase { + batch_size: 2, + channels_in: 3, + channels_out: 5, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + weight_groups: 1, + offset_groups: 1, + height: 4, + width: 4, + }; + + test.assert_output(TestTensor::<4>::from([ + [ + [[0.215466, 0.192846], [0.193407, 0.175496]], + [[0.725073, 0.675926], [0.687746, 0.648506]], + [[1.234679, 1.159006], [1.182085, 1.121516]], + [[1.744286, 1.642086], [1.676423, 1.594526]], + [[2.253892, 2.125167], [2.170762, 2.067536]], + ], + [ + [[1.652976, 1.136937], [0.984030, 0.718403]], + [[4.836801, 3.472453], [3.177263, 2.418021]], + [[8.020626, 5.807969], [5.370497, 4.117639]], + [[11.204453, 8.143486], [7.563731, 5.817256]], + [[14.388277, 10.479003], [9.756965, 7.516875]], + ], + ])) +} + +#[test] +fn test_deform_conv2d_weight_groups() { + let test = DeformConv2dTestCase { + batch_size: 1, + channels_in: 3, + channels_out: 6, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + weight_groups: 3, + offset_groups: 1, + height: 4, + width: 4, + }; + + test.assert_output(TestTensor::<4>::from([[ + [[0.101823, 0.065756], [0.046691, 0.036233]], + [[0.412523, 0.336674], [0.306863, 0.282386]], + [[1.307585, 1.024152], [0.902454, 0.800008]], + [[1.840507, 1.458072], [1.299371, 1.158781]], + [[3.402235, 2.634555], [2.305198, 2.014265]], + [[4.157379, 3.231476], [2.838861, 2.485659]], + ]])) +} + +#[test] +fn test_deform_conv2d_offset_groups() { + let test = DeformConv2dTestCase { + batch_size: 1, + channels_in: 3, + channels_out: 6, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + weight_groups: 1, + offset_groups: 3, + height: 4, + width: 4, + }; + + test.assert_output(TestTensor::<4>::from([[ + [[1.0794, 0.7676], [0.7209, 0.5337]], + [[2.7059, 2.0216], [1.9740, 1.5419]], + [[4.3325, 3.2755], [3.2271, 2.5501]], + [[5.9590, 4.5295], [4.4802, 3.5582]], + [[7.5855, 5.7835], [5.7333, 4.5664]], + [[9.2120, 7.0375], [6.9864, 5.5746]], + ]])) +} + +#[test] +fn test_deform_conv2d_different_kernel_size() { + let test = DeformConv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 3, + kernel_size_1: 3, + kernel_size_2: 4, + padding_1: 0, + padding_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + weight_groups: 1, + offset_groups: 1, + height: 4, + width: 4, + }; + + test.assert_output(TestTensor::<4>::from([[ + [[1.0669], [0.6329]], + [[2.9741], [2.0383]], + [[4.8812], [3.4437]], + ]])) +} + +#[test] +fn test_deform_conv2d_different_padding_size() { + let test = DeformConv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 3, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 2, + padding_2: 3, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + weight_groups: 1, + offset_groups: 1, + height: 4, + width: 4, + }; + + test.assert_output(TestTensor::<4>::from([[ + [ + [ + 0.199779, 0.376176, 0.528501, 0.605256, 0.384365, 0.198675, 0.048145, 0.000000, + ], + [ + 0.287923, 0.551719, 0.777562, 0.890479, 0.580469, 0.304325, 0.079554, 0.000000, + ], + [ + 0.372947, 0.721405, 1.013668, 1.151988, 0.756444, 0.393098, 0.101582, 0.000000, + ], + [ + 0.132138, 0.324872, 0.495372, 0.584617, 0.453122, 0.250084, 0.075703, 0.000000, + ], + [ + 0.059332, 0.160658, 0.244789, 0.297057, 0.239464, 0.132701, 0.047114, 0.000000, + ], + [ + 0.014338, 0.051338, 0.078303, 0.094190, 0.081278, 0.041954, 0.014506, 0.000000, + ], + ], + [ + [ + 0.766652, 1.164805, 1.521938, 1.711110, 1.230500, 0.807579, 0.450423, 0.333333, + ], + [ + 0.981162, 1.601005, 2.152534, 2.440920, 1.745547, 1.091843, 0.536749, 0.333333, + ], + [ + 1.196386, 2.044845, 2.785330, 3.152243, 2.242613, 1.351308, 0.604905, 0.333333, + ], + [ + 0.669465, 1.178133, 1.644096, 1.902188, 1.573183, 1.033924, 0.553577, 0.333333, + ], + [ + 0.495048, 0.786124, 1.039796, 1.204721, 1.052342, 0.743887, 0.483380, 0.333333, + ], + [ + 0.378767, 0.498209, 0.592867, 0.654230, 0.615487, 0.488202, 0.390890, 0.333333, + ], + ], + [ + [ + 1.333524, 1.953435, 2.515375, 2.816964, 2.076636, 1.416483, 0.852701, 0.666667, + ], + [ + 1.674402, 2.650291, 3.527507, 3.991360, 2.910625, 1.879361, 0.993943, 0.666667, + ], + [ + 2.019825, 3.368286, 4.556992, 5.152499, 3.728782, 2.309520, 1.108229, 0.666667, + ], + [ + 1.206791, 2.031395, 2.792820, 3.219759, 2.693245, 1.817763, 1.031452, 0.666667, + ], + [ + 0.930765, 1.411590, 1.834802, 2.112385, 1.865221, 1.355072, 0.919646, 0.666667, + ], + [ + 0.743195, 0.945081, 1.107431, 1.214270, 1.149695, 0.934451, 0.767274, 0.666667, + ], + ], + ]])) +} + +#[test] +fn test_deform_conv2d_different_stride() { + let test = DeformConv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 3, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + stride_1: 1, + stride_2: 2, + dilation_1: 1, + dilation_2: 1, + weight_groups: 1, + offset_groups: 1, + height: 4, + width: 4, + }; + + test.assert_output(TestTensor::<4>::from([[ + [[1.0647], [0.5783]], + [[2.9289], [1.8829]], + [[4.7931], [3.1875]], + ]])) +} + +#[test] +fn test_deform_conv2d_different_dilation() { + let test = DeformConv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 3, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 2, + weight_groups: 1, + offset_groups: 1, + height: 5, + width: 5, + }; + + test.assert_output(TestTensor::<4>::from([[ + [[0.6162], [0.7611], [0.4666]], + [[1.8578], [2.2684], [1.6208]], + [[3.0994], [3.7757], [2.7749]], + ]])) +} + +#[test] +fn test_deform_conv2d_different_width() { + let test = DeformConv2dTestCase { + batch_size: 1, + channels_in: 2, + channels_out: 3, + kernel_size_1: 3, + kernel_size_2: 3, + padding_1: 0, + padding_2: 0, + stride_1: 1, + stride_2: 1, + dilation_1: 1, + dilation_2: 1, + weight_groups: 1, + offset_groups: 1, + height: 6, + width: 4, + }; + + test.assert_output(TestTensor::<4>::from([[ + [ + [0.8909, 0.6016], + [1.0697, 0.7186], + [1.2618, 0.8433], + [0.6424, 0.5032], + ], + [ + [2.4670, 1.8168], + [2.9529, 2.1497], + [3.4805, 2.5090], + [2.0925, 1.7411], + ], + [ + [4.0432, 3.0321], + [4.8362, 3.5809], + [5.6992, 4.1746], + [3.5425, 2.9790], + ], + ]])) +} + +struct DeformConv2dTestCase { + batch_size: usize, + channels_in: usize, + channels_out: usize, + kernel_size_1: usize, + kernel_size_2: usize, + padding_1: usize, + padding_2: usize, + stride_1: usize, + stride_2: usize, + dilation_1: usize, + dilation_2: usize, + weight_groups: usize, + offset_groups: usize, + height: usize, + width: usize, +} + +impl DeformConv2dTestCase { + fn assert_output(self, y: Tensor<4>) { + let out_height = + (self.height + 2 * self.padding_1 - self.dilation_1 * (self.kernel_size_1 - 1) - 1) + / self.stride_1 + + 1; + let out_width = + (self.width + 2 * self.padding_2 - self.dilation_2 * (self.kernel_size_2 - 1) - 1) + / self.stride_2 + + 1; + + let shape_x = Shape::new([self.batch_size, self.channels_in, self.height, self.width]); + let shape_weight = Shape::new([ + self.channels_out, + self.channels_in / self.weight_groups, + self.kernel_size_1, + self.kernel_size_2, + ]); + let shape_offset = Shape::new([ + self.batch_size, + self.kernel_size_1 * self.kernel_size_2 * self.offset_groups * 2, + out_height, + out_width, + ]); + let shape_mask = Shape::new([ + self.batch_size, + self.kernel_size_1 * self.kernel_size_2 * self.offset_groups, + out_height, + out_width, + ]); + let device = Default::default(); + let weight = TestTensor::<4>::from( + TestTensorInt::arange(0..shape_weight.num_elements() as i64, &device) + .reshape::<4, _>(shape_weight.clone()) + .into_data(), + ) + .div_scalar(shape_weight.num_elements() as f32); + let bias = TestTensor::<1>::from( + TestTensorInt::arange(0..self.channels_out as i64, &device).into_data(), + ) + .div_scalar(self.channels_out as f32); + let x = TestTensor::<4>::from( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &device) + .reshape::<4, _>(shape_x.clone()) + .into_data(), + ) + .div_scalar(shape_x.num_elements() as f32); + let offset = TestTensor::<4>::from( + TestTensorInt::arange(0..shape_offset.num_elements() as i64, &device) + .reshape::<4, _>(shape_offset.clone()) + .into_data(), + ) + .div_scalar(shape_offset.num_elements() as f32); + let mask = TestTensor::<4>::from( + TestTensorInt::arange(0..shape_mask.num_elements() as i64, &device) + .reshape::<4, _>(shape_mask.clone()) + .into_data(), + ) + .div_scalar(shape_mask.num_elements() as f32); + + let output = deform_conv2d( + x, + offset, + weight, + Some(mask), + Some(bias), + DeformConvOptions::new( + [self.stride_1, self.stride_2], + [self.padding_1, self.padding_2], + [self.dilation_1, self.dilation_2], + self.weight_groups, + self.offset_groups, + ), + ); + + let tolerance = Tolerance::permissive(); + y.to_data() + .assert_approx_eq::(&output.into_data(), tolerance); + } +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/forward.rs b/crates/burn-backend-tests/tests/tensor/float/module/forward.rs new file mode 100644 index 0000000..5a9d097 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/forward.rs @@ -0,0 +1,18 @@ +use super::*; +use burn_tensor::{TensorData, module::embedding}; + +#[test] +fn test_embedding_forward() { + let weights = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let indices = TensorData::from([[0, 1], [1, 1]]); + let weights = TestTensor::<2>::from(weights); + let indices = TestTensorInt::<2>::from(indices); + + let output = embedding(weights, indices); + let expected = TensorData::from([ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[3.0, 4.0, 5.0], [3.0, 4.0, 5.0]], + ]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/lanczos3_interpolate.rs b/crates/burn-backend-tests/tests/tensor/float/module/lanczos3_interpolate.rs new file mode 100644 index 0000000..cd1503b --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/lanczos3_interpolate.rs @@ -0,0 +1,235 @@ +use super::*; +use burn_tensor::Shape; +use burn_tensor::Tolerance; +use burn_tensor::module::interpolate; +use burn_tensor::ops::{InterpolateMode, InterpolateOptions}; + +#[test] +fn test_upsample_interpolation() { + let test = InterpolateTestCase { + batch_size: 2, + channels: 1, + height: 7, + width: 5, + height_out: 8, + width_out: 7, + }; + + test.assert_output(TestTensor::from([ + [[ + [-0.0000, 0.5685, 1.3918, 2.0000, 2.6082, 3.4315, 4.0000], + [4.0822, 4.6507, 5.4740, 6.0822, 6.6904, 7.5137, 8.0822], + [8.7971, 9.3656, 10.1889, 10.7971, 11.4053, 12.2286, 12.7971], + [ + 12.8964, 13.4649, 14.2882, 14.8964, 15.5046, 16.3279, 16.8964, + ], + [ + 17.1036, 17.6721, 18.4954, 19.1036, 19.7118, 20.5351, 21.1036, + ], + [ + 21.2029, 21.7715, 22.5947, 23.2029, 23.8112, 24.6344, 25.2029, + ], + [ + 25.9178, 26.4863, 27.3096, 27.9178, 28.5260, 29.3493, 29.9178, + ], + [ + 30.0000, 30.5685, 31.3918, 32.0000, 32.6082, 33.4315, 34.0000, + ], + ]], + [[ + [ + 35.0000, 35.5685, 36.3918, 37.0000, 37.6082, 38.4315, 39.0000, + ], + [ + 39.0822, 39.6507, 40.4740, 41.0822, 41.6904, 42.5137, 43.0822, + ], + [ + 43.7971, 44.3656, 45.1888, 45.7971, 46.4053, 47.2286, 47.7971, + ], + [ + 47.8964, 48.4649, 49.2882, 49.8964, 50.5046, 51.3279, 51.8964, + ], + [ + 52.1036, 52.6721, 53.4954, 54.1036, 54.7118, 55.5351, 56.1036, + ], + [ + 56.2029, 56.7715, 57.5947, 58.2029, 58.8112, 59.6344, 60.2029, + ], + [ + 60.9178, 61.4863, 62.3096, 62.9178, 63.5260, 64.3493, 64.9178, + ], + [ + 65.0000, 65.5685, 66.3918, 67.0000, 67.6082, 68.4315, 69.0000, + ], + ]], + ])); +} + +#[test] +fn test_downsample_interpolation() { + let test = InterpolateTestCase { + batch_size: 1, + channels: 1, + height: 45, + width: 14, + height_out: 4, + width_out: 6, + }; + + test.assert_output(TestTensor::from([[[ + [-0.0000, 2.6107, 5.1803, 7.8197, 10.3893, 13.0000], + [205.5606, 208.1713, 210.7408, 213.3802, 215.9498, 218.5606], + [410.4395, 413.0502, 415.6198, 418.2592, 420.8287, 423.4395], + [616.0000, 618.6107, 621.1803, 623.8197, 626.3893, 629.0000], + ]]])); +} + +#[test] +fn test_upsample_2x() { + let test = InterpolateTestCase { + batch_size: 1, + channels: 1, + height: 4, + width: 4, + height_out: 8, + width_out: 8, + }; + + test.assert_output(TestTensor::from([[[ + [ + -0.0000, 0.2972, 0.8164, 1.3131, 1.6869, 2.1836, 2.7028, 3.0000, + ], + [ + 1.1889, 1.4861, 2.0053, 2.5020, 2.8758, 3.3725, 3.8917, 4.1889, + ], + [ + 3.2658, 3.5630, 4.0822, 4.5789, 4.9527, 5.4493, 5.9685, 6.2658, + ], + [ + 5.2524, 5.5496, 6.0689, 6.5655, 6.9393, 7.4360, 7.9552, 8.2524, + ], + [ + 6.7476, 7.0448, 7.5640, 8.0607, 8.4345, 8.9311, 9.4504, 9.7476, + ], + [ + 8.7342, 9.0315, 9.5507, 10.0473, 10.4211, 10.9178, 11.4370, 11.7342, + ], + [ + 10.8111, 11.1083, 11.6275, 12.1242, 12.4980, 12.9947, 13.5139, 13.8111, + ], + [ + 12.0000, 12.2972, 12.8164, 13.3131, 13.6869, 14.1836, 14.7028, 15.0000, + ], + ]]])); +} + +#[test] +fn test_upsample_half_pixel() { + let test = InterpolateTestCase { + batch_size: 1, + channels: 1, + height: 4, + width: 4, + height_out: 8, + width_out: 8, + }; + + test.assert_output_with_align_corners( + TestTensor::from([[[ + [ + -0.4626, -0.2276, 0.3055, 0.9087, 1.3512, 1.9543, 2.4875, 2.7225, + ], + [ + 0.4773, 0.7123, 1.2454, 1.8486, 2.2911, 2.8942, 3.4274, 3.6623, + ], + [ + 2.6099, 2.8449, 3.3780, 3.9812, 4.4237, 5.0268, 5.5600, 5.7949, + ], + [ + 5.0224, 5.2574, 5.7906, 6.3937, 6.8362, 7.4394, 7.9725, 8.2075, + ], + [ + 6.7925, 7.0275, 7.5606, 8.1638, 8.6063, 9.2094, 9.7426, 9.9776, + ], + [ + 9.2051, 9.4400, 9.9732, 10.5763, 11.0188, 11.6220, 12.1551, 12.3901, + ], + [ + 11.3377, 11.5726, 12.1058, 12.7089, 13.1514, 13.7546, 14.2877, 14.5227, + ], + [ + 12.2775, 12.5125, 13.0457, 13.6488, 14.0913, 14.6945, 15.2276, 15.4626, + ], + ]]]), + false, + ); +} + +#[test] +fn test_1d_lanczos3() { + let device = Default::default(); + + let input = TestTensor::<3>::from_data( + [[[1.5410, -0.2934, -2.1788, 0.5684, -1.0845, -1.3986]]], + &device, + ); + + let input = input.unsqueeze_dim(2); + + let output = interpolate( + input, + [1, 9], + InterpolateOptions::new(InterpolateMode::Lanczos3), + ); + assert_eq!(output.dims(), [1, 1, 1, 9]); + + assert!( + !output + .clone() + .to_data() + .as_slice::() + .unwrap() + .iter() + .any(|&x| x.is_nan()), + "interpolate output contains NaN" + ); + + TestTensor::<4>::from([[[[ + 1.5410, 0.7266, -1.1387, -2.2672, -0.7894, 0.6408, -0.4967, -1.4650, -1.3986, + ]]]]) + .to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::permissive()); +} + +struct InterpolateTestCase { + batch_size: usize, + channels: usize, + height: usize, + width: usize, + height_out: usize, + width_out: usize, +} + +impl InterpolateTestCase { + fn assert_output(self, y: TestTensor<4>) { + self.assert_output_with_align_corners(y, true); + } + + fn assert_output_with_align_corners(self, y: TestTensor<4>, align_corners: bool) { + let shape_x = Shape::new([self.batch_size, self.channels, self.height, self.width]); + let x = TestTensor::from( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &y.device()) + .reshape::<4, _>(shape_x) + .into_data(), + ); + let output = interpolate( + x, + [self.height_out, self.width_out], + InterpolateOptions::new(InterpolateMode::Lanczos3).with_align_corners(align_corners), + ); + + let tolerance = Tolerance::permissive(); + y.to_data() + .assert_approx_eq::(&output.into_data(), tolerance); + } +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/linear.rs b/crates/burn-backend-tests/tests/tensor/float/module/linear.rs new file mode 100644 index 0000000..07f06dc --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/linear.rs @@ -0,0 +1,59 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; +use burn_tensor::module::linear; + +#[test] +fn test_linear_1d() { + let weight = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]); + + let x = TestTensor::<1>::from([1.0, 2.0]); + let output = linear(x, weight, None); + + let expected = TensorData::from([7.0, 10.0]); + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::relative(1e-5)); +} + +#[test] +fn test_linear_1d_one_element_output() { + let weight = TestTensor::<2>::from([[3.0], [4.0]]); + + let x = TestTensor::<1>::from([1.0, 2.0]); + let output = linear(x, weight, None); + + let expected = TensorData::from([11.0]); + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::relative(1e-5)); +} + +#[test] +fn test_linear_forward_no_bias() { + let weight = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]); + + let x = TestTensor::<3>::from([[[1.0, 2.0], [3.0, 4.0]], [[-1.0, -2.0], [-3.0, -4.0]]]); + + let output = linear(x, weight, None); + + let expected = TensorData::from([[[7.0, 10.0], [15.0, 22.0]], [[-7.0, -10.0], [-15.0, -22.0]]]); + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::relative(1e-5)); +} + +#[test] +fn test_linear_forward_with_bias() { + let weight = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]); + let bias = Some(TestTensor::<1>::from([1.0, -1.0])); + + let x = TestTensor::<3>::from([[[1.0, 2.0], [3.0, 4.0]], [[-1.0, -2.0], [-3.0, -4.0]]]); + + let output = linear(x, weight, bias); + + let expected = TensorData::from([[[8.0, 9.0], [16.0, 21.0]], [[-6.0, -11.0], [-14.0, -23.0]]]); + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::relative(1e-5)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/maxpool1d.rs b/crates/burn-backend-tests/tests/tensor/float/module/maxpool1d.rs new file mode 100644 index 0000000..fdfee13 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/maxpool1d.rs @@ -0,0 +1,155 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; +use burn_tensor::module::{max_pool1d, max_pool1d_with_indices}; + +#[test] +fn test_max_pool1d_simple() { + let kernel_size = 3; + let padding = 0; + let stride = 1; + let dilation = 1; + + let x = TestTensor::from([[ + [0.9861, 0.5474, 0.4477, 0.0732, 0.3548, 0.8221], + [0.8148, 0.5474, 0.9490, 0.7890, 0.5537, 0.5689], + ]]); + let y = TestTensor::<3>::from([[ + [0.9861, 0.5474, 0.4477, 0.8221], + [0.949, 0.949, 0.949, 0.789], + ]]); + + let output = max_pool1d(x, kernel_size, stride, padding, dilation, false); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); +} + +#[test] +fn test_max_pool1d_different_padding_stride_kernel() { + let kernel_size = 3; + let padding = 1; + let stride = 2; + let dilation = 1; + + let x = TestTensor::from([[[0.6309, 0.6112, 0.6998, 0.4708]]]); + let y = TestTensor::<3>::from([[[0.6309, 0.6998]]]); + + let output = max_pool1d(x, kernel_size, stride, padding, dilation, false); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); +} + +#[test] +fn test_max_pool1d_with_neg() { + let kernel_size = 3; + let padding = 1; + let stride = 1; + let dilation = 1; + + let x = TestTensor::from([[[-0.6309, -0.6112, -0.6998, -0.4708]]]); + let y = TestTensor::<3>::from([[[-0.6112, -0.6112, -0.4708, -0.4708]]]); + + let output = max_pool1d(x, kernel_size, stride, padding, dilation, false); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); +} + +#[test] +fn test_max_pool1d_with_dilation() { + let kernel_size = 2; + let padding = 1; + let stride = 1; + let dilation = 2; + + let x = TestTensor::from([[ + [0.9861, 0.5474, 0.4477, 0.0732, 0.3548, 0.8221], + [0.8148, 0.5474, 0.9490, 0.7890, 0.5537, 0.5689], + ]]); + let y = TestTensor::<3>::from([[ + [0.5474, 0.9861, 0.5474, 0.4477, 0.8221, 0.3548], + [0.5474, 0.9490, 0.7890, 0.9490, 0.7890, 0.5537], + ]]); + + let output = max_pool1d(x, kernel_size, stride, padding, dilation, false); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); +} + +#[test] +fn test_max_pool1d_with_indices() { + let kernel_size = 2; + let padding = 0; + let stride = 1; + let dilation = 1; + + let x = TestTensor::from([[[0.2479, 0.6386, 0.3166, 0.5742]]]); + let indices = TensorData::from([[[1, 1, 3]]]); + let y = TestTensor::<3>::from([[[0.6386, 0.6386, 0.5742]]]); + + let (output, output_indices) = + max_pool1d_with_indices(x, kernel_size, stride, padding, dilation, false); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); + output_indices.into_data().assert_eq(&indices, false); +} + +#[test] +fn test_max_pool1d_complex() { + let kernel_size = 4; + let padding = 2; + let stride = 1; + let dilation = 1; + + let x = TestTensor::from([[[0.5388, 0.0676, 0.7122, 0.8316, 0.0653]]]); + let indices = TensorData::from([[[0, 2, 3, 3, 3, 3]]]); + let y = TestTensor::<3>::from([[[0.5388, 0.7122, 0.8316, 0.8316, 0.8316, 0.8316]]]); + + let (output, output_indices) = + max_pool1d_with_indices(x, kernel_size, stride, padding, dilation, false); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); + output_indices.into_data().assert_eq(&indices, false); +} + +#[test] +fn test_max_pool1d_ceil_mode() { + // Test ceil_mode=true produces larger output when input doesn't divide evenly by stride + // Input: 1x1x6, kernel: 3, stride: 2, padding: 0 + // Floor mode: output = (6-3)/2+1 = 2 elements + // Ceil mode: output = ceil((6-3)/2)+1 = ceil(1.5)+1 = 3 elements + let kernel_size = 3; + let padding = 0; + let stride = 2; + let dilation = 1; + + let x = TestTensor::from([[[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]]]); + + // With ceil_mode=false (floor): output is 2 elements + // Window 0: positions [0:3] -> max(1,2,3) = 3 + // Window 1: positions [2:5] -> max(3,4,5) = 5 + let y_floor = TestTensor::<3>::from([[[3.0, 5.0]]]); + + let output_floor = max_pool1d(x.clone(), kernel_size, stride, padding, dilation, false); + + y_floor + .to_data() + .assert_approx_eq::(&output_floor.into_data(), Tolerance::default()); + + // With ceil_mode=true: output is 3 elements + // Window 0: positions [0:3] -> max(1,2,3) = 3 + // Window 1: positions [2:5] -> max(3,4,5) = 5 + // Window 2: positions [4:7] -> max(5,6) = 6 (partial window) + let y_ceil = TestTensor::<3>::from([[[3.0, 5.0, 6.0]]]); + + let output_ceil = max_pool1d(x, kernel_size, stride, padding, dilation, true); + + y_ceil + .to_data() + .assert_approx_eq::(&output_ceil.into_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/maxpool2d.rs b/crates/burn-backend-tests/tests/tensor/float/module/maxpool2d.rs new file mode 100644 index 0000000..325d5bb --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/maxpool2d.rs @@ -0,0 +1,523 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; +use burn_tensor::module::{max_pool2d, max_pool2d_with_indices}; + +#[test] +fn test_max_pool2d_simple() { + let kernel_size_1 = 3; + let kernel_size_2 = 3; + let padding_1 = 1; + let padding_2 = 1; + let stride_1 = 1; + let stride_2 = 1; + let dilation_1 = 1; + let dilation_2 = 1; + + let x = TestTensor::from([ + [ + [ + [0.9861, 0.5474, 0.4477, 0.0732, 0.3548, 0.8221], + [0.8148, 0.5474, 0.9490, 0.7890, 0.5537, 0.5689], + [0.5986, 0.2059, 0.4897, 0.6136, 0.2965, 0.6182], + [0.1485, 0.9540, 0.4023, 0.6176, 0.7111, 0.3392], + [0.3703, 0.0472, 0.2771, 0.1868, 0.8855, 0.5605], + [0.5063, 0.1638, 0.9432, 0.7836, 0.8696, 0.1068], + ], + [ + [0.8872, 0.0137, 0.1652, 0.5505, 0.6127, 0.6473], + [0.1128, 0.0888, 0.1152, 0.5456, 0.6199, 0.7947], + [0.5911, 0.7781, 0.7256, 0.6578, 0.0989, 0.9149], + [0.5879, 0.5189, 0.6561, 0.0578, 0.7025, 0.6426], + [0.9590, 0.0325, 0.6455, 0.6248, 0.2009, 0.1544], + [0.7339, 0.1369, 0.6598, 0.5528, 0.6775, 0.1572], + ], + ], + [ + [ + [0.6853, 0.6439, 0.4639, 0.5573, 0.2723, 0.5910], + [0.5419, 0.7729, 0.6743, 0.8956, 0.2997, 0.9546], + [0.0334, 0.2178, 0.6917, 0.4958, 0.3357, 0.6584], + [0.7358, 0.9074, 0.2462, 0.5159, 0.6420, 0.2441], + [0.7602, 0.6297, 0.6073, 0.5937, 0.8037, 0.4881], + [0.8859, 0.0974, 0.3954, 0.6763, 0.1078, 0.7467], + ], + [ + [0.2991, 0.5012, 0.8024, 0.7653, 0.9378, 0.7952], + [0.7393, 0.2336, 0.9521, 0.2719, 0.8445, 0.0454], + [0.6479, 0.9822, 0.7905, 0.0318, 0.2474, 0.0628], + [0.9955, 0.7591, 0.4140, 0.3215, 0.4349, 0.1527], + [0.8064, 0.0164, 0.4002, 0.2024, 0.6128, 0.5827], + [0.5368, 0.7895, 0.8727, 0.7793, 0.0910, 0.3421], + ], + ], + ]); + let y = TestTensor::<4>::from([ + [ + [ + [0.9861, 0.9861, 0.9490, 0.9490, 0.8221, 0.8221], + [0.9861, 0.9861, 0.9490, 0.9490, 0.8221, 0.8221], + [0.9540, 0.9540, 0.9540, 0.9490, 0.7890, 0.7111], + [0.9540, 0.9540, 0.9540, 0.8855, 0.8855, 0.8855], + [0.9540, 0.9540, 0.9540, 0.9432, 0.8855, 0.8855], + [0.5063, 0.9432, 0.9432, 0.9432, 0.8855, 0.8855], + ], + [ + [0.8872, 0.8872, 0.5505, 0.6199, 0.7947, 0.7947], + [0.8872, 0.8872, 0.7781, 0.7256, 0.9149, 0.9149], + [0.7781, 0.7781, 0.7781, 0.7256, 0.9149, 0.9149], + [0.9590, 0.9590, 0.7781, 0.7256, 0.9149, 0.9149], + [0.9590, 0.9590, 0.6598, 0.7025, 0.7025, 0.7025], + [0.9590, 0.9590, 0.6598, 0.6775, 0.6775, 0.6775], + ], + ], + [ + [ + [0.7729, 0.7729, 0.8956, 0.8956, 0.9546, 0.9546], + [0.7729, 0.7729, 0.8956, 0.8956, 0.9546, 0.9546], + [0.9074, 0.9074, 0.9074, 0.8956, 0.9546, 0.9546], + [0.9074, 0.9074, 0.9074, 0.8037, 0.8037, 0.8037], + [0.9074, 0.9074, 0.9074, 0.8037, 0.8037, 0.8037], + [0.8859, 0.8859, 0.6763, 0.8037, 0.8037, 0.8037], + ], + [ + [0.7393, 0.9521, 0.9521, 0.9521, 0.9378, 0.9378], + [0.9822, 0.9822, 0.9822, 0.9521, 0.9378, 0.9378], + [0.9955, 0.9955, 0.9822, 0.9521, 0.8445, 0.8445], + [0.9955, 0.9955, 0.9822, 0.7905, 0.6128, 0.6128], + [0.9955, 0.9955, 0.8727, 0.8727, 0.7793, 0.6128], + [0.8064, 0.8727, 0.8727, 0.8727, 0.7793, 0.6128], + ], + ], + ]); + + let output = max_pool2d( + x, + [kernel_size_1, kernel_size_2], + [stride_1, stride_2], + [padding_1, padding_2], + [dilation_1, dilation_2], + false, + ); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); +} + +#[test] +fn test_max_pool2d_different_padding_stride_kernel() { + let kernel_size_1 = 3; + let kernel_size_2 = 1; + let padding_1 = 1; + let padding_2 = 0; + let stride_1 = 1; + let stride_2 = 2; + let dilation_1 = 1; + let dilation_2 = 1; + + let x = TestTensor::from([[[ + [0.6309, 0.6112, 0.6998], + [0.4708, 0.9161, 0.5402], + [0.4577, 0.7397, 0.9870], + [0.6380, 0.4352, 0.5884], + [0.6277, 0.5139, 0.4525], + [0.9333, 0.9846, 0.5006], + ]]]); + let y = TestTensor::<4>::from([[[ + [0.6309, 0.6998], + [0.6309, 0.9870], + [0.6380, 0.9870], + [0.6380, 0.9870], + [0.9333, 0.5884], + [0.9333, 0.5006], + ]]]); + + let output = max_pool2d( + x, + [kernel_size_1, kernel_size_2], + [stride_1, stride_2], + [padding_1, padding_2], + [dilation_1, dilation_2], + false, + ); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); +} + +#[test] +fn test_max_pool2d_with_neg() { + let kernel_size_1 = 3; + let kernel_size_2 = 3; + let padding_1 = 1; + let padding_2 = 1; + let stride_1 = 1; + let stride_2 = 1; + let dilation_1 = 1; + let dilation_2 = 1; + + let x = TestTensor::from([[[ + [0.6309, 0.6112, 0.6998], + [0.4708, 0.9161, 0.5402], + [0.4577, 0.7397, 0.9870], + [0.6380, 0.4352, 0.5884], + [0.6277, 0.5139, 0.4525], + [0.9333, 0.9846, 0.5006], + ]]]) + .neg(); + let y = TestTensor::<4>::from([[[ + [-0.4708, -0.4708, -0.5402], + [-0.4577, -0.4577, -0.5402], + [-0.4352, -0.4352, -0.4352], + [-0.4352, -0.4352, -0.4352], + [-0.4352, -0.4352, -0.4352], + [-0.5139, -0.4525, -0.4525], + ]]]); + + let output = max_pool2d( + x, + [kernel_size_1, kernel_size_2], + [stride_1, stride_2], + [padding_1, padding_2], + [dilation_1, dilation_2], + false, + ); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); +} + +#[test] +fn test_max_pool2d_with_dilation() { + let kernel_size_1 = 2; + let kernel_size_2 = 2; + let padding_1 = 0; + let padding_2 = 0; + let stride_1 = 1; + let stride_2 = 1; + let dilation_1 = 2; + let dilation_2 = 2; + + let x = TestTensor::from([[[ + [0.9861, 0.9861, 0.9490, 0.9490, 0.8221, 0.8221], + [0.9861, 0.9861, 0.9490, 0.9490, 0.8221, 0.8221], + [0.9540, 0.9540, 0.9540, 0.9490, 0.7890, 0.7111], + [0.9540, 0.9540, 0.9540, 0.8855, 0.8855, 0.8855], + [0.9540, 0.9540, 0.9540, 0.9432, 0.8855, 0.8855], + [0.5063, 0.9432, 0.9432, 0.9432, 0.8855, 0.8855], + ]]]); + let y = TestTensor::<4>::from([[[ + [0.9861, 0.9861, 0.9540, 0.9490], + [0.9861, 0.9861, 0.9540, 0.9490], + [0.9540, 0.9540, 0.9540, 0.9490], + [0.9540, 0.9540, 0.9540, 0.9432], + ]]]); + + let output = max_pool2d( + x, + [kernel_size_1, kernel_size_2], + [stride_1, stride_2], + [padding_1, padding_2], + [dilation_1, dilation_2], + false, + ); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); +} + +#[test] +fn test_max_pool2d_with_indices() { + let kernel_size_1 = 2; + let kernel_size_2 = 2; + let padding_1 = 1; + let padding_2 = 1; + let stride_1 = 1; + let stride_2 = 1; + let dilation_1 = 1; + let dilation_2 = 1; + + let x = TestTensor::from([[[ + [0.2479, 0.6386, 0.3166, 0.5742], + [0.7065, 0.1940, 0.6305, 0.8959], + [0.5416, 0.8602, 0.8129, 0.1662], + [0.3358, 0.3059, 0.8293, 0.0990], + ]]]); + let indices = TensorData::from([[[ + [0, 1, 1, 3, 3], + [4, 4, 1, 7, 7], + [4, 9, 9, 7, 7], + [8, 9, 9, 14, 11], + [12, 12, 14, 14, 15], + ]]]); + let y = TestTensor::<4>::from([[[ + [0.2479, 0.6386, 0.6386, 0.5742, 0.5742], + [0.7065, 0.7065, 0.6386, 0.8959, 0.8959], + [0.7065, 0.8602, 0.8602, 0.8959, 0.8959], + [0.5416, 0.8602, 0.8602, 0.8293, 0.1662], + [0.3358, 0.3358, 0.8293, 0.8293, 0.0990], + ]]]); + + let (output, output_indices) = max_pool2d_with_indices( + x, + [kernel_size_1, kernel_size_2], + [stride_1, stride_2], + [padding_1, padding_2], + [dilation_1, dilation_2], + false, + ); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); + output_indices.into_data().assert_eq(&indices, false); +} + +#[test] +fn test_max_pool2d_complex() { + let kernel_size_1 = 4; + let kernel_size_2 = 2; + let padding_1 = 2; + let padding_2 = 1; + let stride_1 = 1; + let stride_2 = 2; + let dilation_1 = 1; + let dilation_2 = 1; + + let x = TestTensor::from([[[ + [0.5388, 0.0676, 0.7122, 0.8316, 0.0653], + [0.9154, 0.1536, 0.9089, 0.8016, 0.7518], + [0.2073, 0.0501, 0.8811, 0.5604, 0.5075], + [0.4384, 0.9963, 0.9698, 0.4988, 0.2609], + [0.3391, 0.2230, 0.4610, 0.5365, 0.6880], + ]]]); + let indices = TensorData::from([[[ + [5, 7, 3], + [5, 7, 3], + [5, 16, 3], + [5, 16, 8], + [15, 16, 24], + [15, 16, 24], + ]]]); + let y = TestTensor::<4>::from([[[ + [0.9154, 0.9089, 0.8316], + [0.9154, 0.9089, 0.8316], + [0.9154, 0.9963, 0.8316], + [0.9154, 0.9963, 0.8016], + [0.4384, 0.9963, 0.688], + [0.4384, 0.9963, 0.688], + ]]]); + let (output, output_indices) = max_pool2d_with_indices( + x, + [kernel_size_1, kernel_size_2], + [stride_1, stride_2], + [padding_1, padding_2], + [dilation_1, dilation_2], + false, + ); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); + output_indices.into_data().assert_eq(&indices, false); +} + +#[test] +fn test_max_pool2d_ceil_mode() { + // Test ceil_mode=true which produces larger output when input doesn't divide evenly by stride + // Using 1x1x6x6 with kernel 3x3, stride 2x2, padding 0: + // Floor mode: output = (6+0-1*(3-1)-1)/2+1 = 3/2+1 = 2 x 2 + // Ceil mode: output = ceil(3/2)+1 = 2+1 = 3 x 3 + let kernel_size_1 = 3; + let kernel_size_2 = 3; + let padding_1 = 0; + let padding_2 = 0; + let stride_1 = 2; + let stride_2 = 2; + let dilation_1 = 1; + let dilation_2 = 1; + + // Input (values 1-36 arranged row by row): + // col: 0 1 2 3 4 5 + // row 0: 1 2 3 4 5 6 + // row 1: 7 8 9 10 11 12 + // row 2: 13 14 15 16 17 18 + // row 3: 19 20 21 22 23 24 + // row 4: 25 26 27 28 29 30 + // row 5: 31 32 33 34 35 36 + let x = TestTensor::from([[[ + [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], + [7.0, 8.0, 9.0, 10.0, 11.0, 12.0], + [13.0, 14.0, 15.0, 16.0, 17.0, 18.0], + [19.0, 20.0, 21.0, 22.0, 23.0, 24.0], + [25.0, 26.0, 27.0, 28.0, 29.0, 30.0], + [31.0, 32.0, 33.0, 34.0, 35.0, 36.0], + ]]]); + + // With ceil_mode=false (floor): output is 2x2 + // (0,0): rows 0-2, cols 0-2 -> max(1,2,3,7,8,9,13,14,15) = 15 + // (0,1): rows 0-2, cols 2-4 -> max(3,4,5,9,10,11,15,16,17) = 17 + // (1,0): rows 2-4, cols 0-2 -> max(13,14,15,19,20,21,25,26,27) = 27 + // (1,1): rows 2-4, cols 2-4 -> max(15,16,17,21,22,23,27,28,29) = 29 + let y_floor = TestTensor::<4>::from([[[[15.0, 17.0], [27.0, 29.0]]]]); + + let output_floor = max_pool2d( + x.clone(), + [kernel_size_1, kernel_size_2], + [stride_1, stride_2], + [padding_1, padding_2], + [dilation_1, dilation_2], + false, + ); + + y_floor + .to_data() + .assert_approx_eq::(&output_floor.into_data(), Tolerance::default()); + + // With ceil_mode=true: output is 3x3 + // Extra windows at edges use only available input values (padded with -inf for max pooling) + // (0,0): rows 0-2, cols 0-2 -> max = 15 + // (0,1): rows 0-2, cols 2-4 -> max = 17 + // (0,2): rows 0-2, cols 4-5 -> max(5,6,11,12,17,18) = 18 + // (1,0): rows 2-4, cols 0-2 -> max = 27 + // (1,1): rows 2-4, cols 2-4 -> max = 29 + // (1,2): rows 2-4, cols 4-5 -> max(17,18,23,24,29,30) = 30 + // (2,0): rows 4-5, cols 0-2 -> max(25,26,27,31,32,33) = 33 + // (2,1): rows 4-5, cols 2-4 -> max(27,28,29,33,34,35) = 35 + // (2,2): rows 4-5, cols 4-5 -> max(29,30,35,36) = 36 + let y_ceil = + TestTensor::<4>::from([[[[15.0, 17.0, 18.0], [27.0, 29.0, 30.0], [33.0, 35.0, 36.0]]]]); + + let output_ceil = max_pool2d( + x, + [kernel_size_1, kernel_size_2], + [stride_1, stride_2], + [padding_1, padding_2], + [dilation_1, dilation_2], + true, + ); + + y_ceil + .to_data() + .assert_approx_eq::(&output_ceil.into_data(), Tolerance::default()); +} + +#[test] +fn test_max_pool2d_ceil_mode_with_indices() { + // Test ceil_mode=true with indices to verify correct index calculation + // when pooling windows extend beyond original input bounds + let kernel_size_1 = 3; + let kernel_size_2 = 3; + let padding_1 = 0; + let padding_2 = 0; + let stride_1 = 2; + let stride_2 = 2; + let dilation_1 = 1; + let dilation_2 = 1; + + // Input 6x6 (indices 0-35 in row-major order): + // row 0: 0 1 2 3 4 5 + // row 1: 6 7 8 9 10 11 + // row 2: 12 13 14 15 16 17 + // row 3: 18 19 20 21 22 23 + // row 4: 24 25 26 27 28 29 + // row 5: 30 31 32 33 34 35 + let x = TestTensor::from([[[ + [0.0, 1.0, 2.0, 3.0, 4.0, 5.0], + [6.0, 7.0, 8.0, 9.0, 10.0, 11.0], + [12.0, 13.0, 14.0, 15.0, 16.0, 17.0], + [18.0, 19.0, 20.0, 21.0, 22.0, 23.0], + [24.0, 25.0, 26.0, 27.0, 28.0, 29.0], + [30.0, 31.0, 32.0, 33.0, 34.0, 35.0], + ]]]); + + // With ceil_mode=true: output is 3x3 + // (0,0): rows 0-2, cols 0-2 -> max at index 14 + // (0,1): rows 0-2, cols 2-4 -> max at index 16 + // (0,2): rows 0-2, cols 4-5 -> max at index 17 + // (1,0): rows 2-4, cols 0-2 -> max at index 26 + // (1,1): rows 2-4, cols 2-4 -> max at index 28 + // (1,2): rows 2-4, cols 4-5 -> max at index 29 + // (2,0): rows 4-5, cols 0-2 -> max at index 32 + // (2,1): rows 4-5, cols 2-4 -> max at index 34 + // (2,2): rows 4-5, cols 4-5 -> max at index 35 + let expected_values = + TestTensor::<4>::from([[[[14.0, 16.0, 17.0], [26.0, 28.0, 29.0], [32.0, 34.0, 35.0]]]]); + let expected_indices = TensorData::from([[[[14i64, 16, 17], [26, 28, 29], [32, 34, 35]]]]); + + let (output, output_indices) = max_pool2d_with_indices( + x, + [kernel_size_1, kernel_size_2], + [stride_1, stride_2], + [padding_1, padding_2], + [dilation_1, dilation_2], + true, + ); + + expected_values + .to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); + output_indices + .into_data() + .assert_eq(&expected_indices, false); +} + +#[test] +fn test_max_pool2d_ceil_mode_with_indices_and_padding() { + // Test ceil_mode=true with padding and indices to verify correct index calculation + // This exercises the case where both user padding and ceil_mode extra padding apply + let kernel_size_1 = 3; + let kernel_size_2 = 3; + let padding_1 = 1; + let padding_2 = 1; + let stride_1 = 2; + let stride_2 = 2; + let dilation_1 = 1; + let dilation_2 = 1; + + // Input 5x5 (indices 0-24 in row-major order): + // row 0: 0 1 2 3 4 + // row 1: 5 6 7 8 9 + // row 2: 10 11 12 13 14 + // row 3: 15 16 17 18 19 + // row 4: 20 21 22 23 24 + let x = TestTensor::from([[[ + [0.0, 1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0, 9.0], + [10.0, 11.0, 12.0, 13.0, 14.0], + [15.0, 16.0, 17.0, 18.0, 19.0], + [20.0, 21.0, 22.0, 23.0, 24.0], + ]]]); + + // With padding=1, ceil_mode=true: + // Effective input is 7x7 (5 + 2*1) + // Output size: ceil((5 + 2*1 - 3) / 2) + 1 = ceil(4/2) + 1 = 3 + // + // Windows (with -inf padding at boundaries): + // (0,0): rows -1 to 1, cols -1 to 1 -> valid: (0,0) to (1,1), max at (1,1)=6 + // (0,1): rows -1 to 1, cols 1 to 3 -> max at (1,3)=8 + // (0,2): rows -1 to 1, cols 3 to 5 -> max at (1,4)=9 + // (1,0): rows 1 to 3, cols -1 to 1 -> max at (3,1)=16 + // (1,1): rows 1 to 3, cols 1 to 3 -> max at (3,3)=18 + // (1,2): rows 1 to 3, cols 3 to 5 -> max at (3,4)=19 + // (2,0): rows 3 to 5, cols -1 to 1 -> max at (4,1)=21 + // (2,1): rows 3 to 5, cols 1 to 3 -> max at (4,3)=23 + // (2,2): rows 3 to 5, cols 3 to 5 -> max at (4,4)=24 + let expected_values = + TestTensor::<4>::from([[[[6.0, 8.0, 9.0], [16.0, 18.0, 19.0], [21.0, 23.0, 24.0]]]]); + let expected_indices = TensorData::from([[[[6i64, 8, 9], [16, 18, 19], [21, 23, 24]]]]); + + let (output, output_indices) = max_pool2d_with_indices( + x, + [kernel_size_1, kernel_size_2], + [stride_1, stride_2], + [padding_1, padding_2], + [dilation_1, dilation_2], + true, + ); + + expected_values + .to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); + output_indices + .into_data() + .assert_eq(&expected_indices, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/mod.rs b/crates/burn-backend-tests/tests/tensor/float/module/mod.rs new file mode 100644 index 0000000..e2be3ea --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/mod.rs @@ -0,0 +1,24 @@ +use super::*; + +mod adaptive_avgpool1d; +mod adaptive_avgpool2d; +mod attention; +mod avgpool1d; +mod avgpool2d; +mod bicubic_interpolate; +mod bilinear_interpolate; +mod conv1d; +mod conv2d; +mod conv3d; +mod conv_transpose1d; +mod conv_transpose2d; +mod conv_transpose3d; +mod ctc; +mod deform_conv2d; +mod forward; +mod lanczos3_interpolate; +mod linear; +mod maxpool1d; +mod maxpool2d; +mod nearest_interpolate; +mod unfold4d; diff --git a/crates/burn-backend-tests/tests/tensor/float/module/nearest_interpolate.rs b/crates/burn-backend-tests/tests/tensor/float/module/nearest_interpolate.rs new file mode 100644 index 0000000..8b6ef32 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/nearest_interpolate.rs @@ -0,0 +1,127 @@ +use super::*; +use burn_tensor::Shape; +use burn_tensor::Tolerance; +use burn_tensor::module::interpolate; +use burn_tensor::ops::{InterpolateMode, InterpolateOptions}; + +#[test] +fn test_upsample_interpolation() { + let test = InterpolateTestCase { + batch_size: 2, + channels: 1, + height: 7, + width: 5, + height_out: 8, + width_out: 7, + }; + + test.assert_output(TestTensor::from([ + [[ + [0., 0., 1., 2., 2., 3., 4.], + [0., 0., 1., 2., 2., 3., 4.], + [5., 5., 6., 7., 7., 8., 9.], + [10., 10., 11., 12., 12., 13., 14.], + [15., 15., 16., 17., 17., 18., 19.], + [20., 20., 21., 22., 22., 23., 24.], + [25., 25., 26., 27., 27., 28., 29.], + [30., 30., 31., 32., 32., 33., 34.], + ]], + [[ + [35., 35., 36., 37., 37., 38., 39.], + [35., 35., 36., 37., 37., 38., 39.], + [40., 40., 41., 42., 42., 43., 44.], + [45., 45., 46., 47., 47., 48., 49.], + [50., 50., 51., 52., 52., 53., 54.], + [55., 55., 56., 57., 57., 58., 59.], + [60., 60., 61., 62., 62., 63., 64.], + [65., 65., 66., 67., 67., 68., 69.], + ]], + ])); +} + +#[test] +fn test_downsample_interpolation() { + let test = InterpolateTestCase { + batch_size: 1, + channels: 1, + height: 45, + width: 14, + height_out: 4, + width_out: 6, + }; + + test.assert_output(TestTensor::from([[[ + [0., 2., 4., 7., 9., 11.], + [154., 156., 158., 161., 163., 165.], + [308., 310., 312., 315., 317., 319.], + [462., 464., 466., 469., 471., 473.], + ]]])); +} + +#[test] +fn test_1d_nearest() { + // Initialize the model without weights (because the exported file does not contain them) + let device = Default::default(); + + // Run the model + let input = TestTensor::<3>::from_data( + [[[1.5410, -0.2934, -2.1788, 0.5684, -1.0845, -1.3986]]], + &device, + ); + + let input = input.unsqueeze_dim(2); + + let output = interpolate( + input, + [1, 9], + InterpolateOptions::new(InterpolateMode::Nearest), + ); + assert_eq!(output.dims(), [1, 1, 1, 9]); + + // assert output data does not contain NaN + assert!( + !output + .clone() + .to_data() + .as_slice::() + .unwrap() + .iter() + .any(|&x| x.is_nan()), + "interpolate output contains NaN" + ); + + TestTensor::<4>::from([[[[ + 1.541, 1.541, -0.2934, -2.1788, -2.1788, 0.5684, -1.0845, -1.0845, -1.3986, + ]]]]) + .to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); +} + +struct InterpolateTestCase { + batch_size: usize, + channels: usize, + height: usize, + width: usize, + height_out: usize, + width_out: usize, +} + +impl InterpolateTestCase { + fn assert_output(self, y: TestTensor<4>) { + let shape_x = Shape::new([self.batch_size, self.channels, self.height, self.width]); + let x = TestTensor::from( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &y.device()) + .reshape::<4, _>(shape_x) + .into_data() + .convert::(), + ); + let output = interpolate( + x, + [self.height_out, self.width_out], + InterpolateOptions::new(InterpolateMode::Nearest), + ); + + y.to_data() + .assert_approx_eq::(&output.into_data(), Tolerance::default()); + } +} diff --git a/crates/burn-backend-tests/tests/tensor/float/module/unfold4d.rs b/crates/burn-backend-tests/tests/tensor/float/module/unfold4d.rs new file mode 100644 index 0000000..ed20837 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/module/unfold4d.rs @@ -0,0 +1,132 @@ +use super::*; +use burn_tensor::Shape; +use burn_tensor::Tolerance; +use burn_tensor::module::unfold4d; +use burn_tensor::ops::UnfoldOptions; + +#[test] +fn test_unfold4d_shape() { + let test = Unfold4dTestCase { + batch_size: 2, + channels_in: 5, + kernel_size: [2, 3], + padding: [0, 0], + stride: [1, 1], + dilation: [1, 1], + height: 3, + width: 4, + }; + + test.assert_shape([2, 30, 4]); +} + +#[test] +fn test_unfold4d_simple() { + let test = Unfold4dTestCase { + batch_size: 1, + channels_in: 2, + kernel_size: [2, 2], + padding: [0, 0], + stride: [1, 1], + dilation: [1, 1], + height: 4, + width: 4, + }; + + test.assert_output(TestTensor::from([[ + [0., 1., 2., 4., 5., 6., 8., 9., 10.], + [1., 2., 3., 5., 6., 7., 9., 10., 11.], + [4., 5., 6., 8., 9., 10., 12., 13., 14.], + [5., 6., 7., 9., 10., 11., 13., 14., 15.], + [16., 17., 18., 20., 21., 22., 24., 25., 26.], + [17., 18., 19., 21., 22., 23., 25., 26., 27.], + [20., 21., 22., 24., 25., 26., 28., 29., 30.], + [21., 22., 23., 25., 26., 27., 29., 30., 31.], + ]])); +} + +#[test] +fn test_unfold4d_complex() { + let test = Unfold4dTestCase { + batch_size: 1, + channels_in: 2, + kernel_size: [2, 3], + padding: [0, 1], + stride: [1, 2], + dilation: [1, 2], + height: 3, + width: 4, + }; + + test.assert_output(TestTensor::from([[ + [0., 0.], + [1., 5.], + [3., 7.], + [0., 0.], + [5., 9.], + [7., 11.], + [0., 0.], + [13., 17.], + [15., 19.], + [0., 0.], + [17., 21.], + [19., 23.], + ]])); +} + +struct Unfold4dTestCase { + batch_size: usize, + channels_in: usize, + kernel_size: [usize; 2], + padding: [usize; 2], + stride: [usize; 2], + dilation: [usize; 2], + height: usize, + width: usize, +} + +impl Unfold4dTestCase { + fn assert_shape(self, expected_shape: [usize; 3]) { + let shape_x = Shape::new([self.batch_size, self.channels_in, self.height, self.width]); + let x = TestTensor::from( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &Default::default()) + .reshape::<4, _>(shape_x) + .into_data() + .convert::(), + ); + + let output = unfold4d( + x, + self.kernel_size, + UnfoldOptions::new(self.stride, self.padding, self.dilation), + ); + + assert_eq!( + output.shape().as_slice(), + expected_shape, + "Expected shape doesn't match the actual shape" + ); + } + + fn assert_output(self, expected: TestTensor<3>) { + let shape_x = Shape::new([self.batch_size, self.channels_in, self.height, self.width]); + let x = TestTensor::from( + TestTensorInt::arange(0..shape_x.num_elements() as i64, &Default::default()) + .reshape::<4, _>(shape_x) + .into_data(), + ); + + let output = unfold4d( + x, + self.kernel_size, + UnfoldOptions::new(self.stride, self.padding, self.dilation), + ); + + let tolerance = Tolerance::default() + .set_half_precision_relative(2e-3) + .set_half_precision_absolute(2e-3); + output + .into_data() + .assert_approx_eq::(&expected.into_data(), tolerance); + } +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/abs.rs b/crates/burn-backend-tests/tests/tensor/float/ops/abs.rs new file mode 100644 index 0000000..ae72fa1 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/abs.rs @@ -0,0 +1,25 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_abs_flipped() { + // [1, -2, 3, -4] flipped -> [-4, 3, -2, 1]; abs -> [4, 3, 2, 1] + let tensor = TestTensor::<1>::from([1.0, -2.0, 3.0, -4.0]); + let flipped = tensor.flip([0]); + + let output = flipped.abs(); + let expected = TensorData::from([4.0, 3.0, 2.0, 1.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_abs_ops_float() { + let tensor = TestTensor::<2>::from([[0.0, -1.0, 2.0], [3.0, 4.0, -5.0]]); + + let output = tensor.abs(); + + output + .into_data() + .assert_eq(&TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/add.rs b/crates/burn-backend-tests/tests/tensor/float/ops/add.rs new file mode 100644 index 0000000..a44e079 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/add.rs @@ -0,0 +1,151 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_add_d2() { + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_2 = TestTensor::from([[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]]); + + let output = tensor_1 + tensor_2; + + output.into_data().assert_eq( + &TensorData::from([[6.0, 8.0, 10.0], [12.0, 14.0, 16.0]]), + false, + ); +} + +#[test] +fn test_add_broadcast() { + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0]]); + let tensor_2 = TestTensor::from([[3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]); + + let output = tensor_1 + tensor_2; + + output.into_data().assert_eq( + &TensorData::from([[3.0, 5.0, 7.0], [6.0, 8.0, 10.0]]), + false, + ); +} + +#[test] +fn test_add_different_strides_rhs() { + // We need to execute an operation after `from data` to trigger inplace in some backends. + // Which is the operation that might be problematic in this case. + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0], [2.0, 3.0]]) * 1; + let tensor_2 = TestTensor::from([[4.0, 5.0], [6.0, 7.0]]) * 1; + + let output = tensor_1 + tensor_2.transpose(); + + output + .into_data() + .assert_eq(&TensorData::from([[4.0, 7.0], [7.0, 10.0]]), false); +} + +#[test] +fn test_add_different_strides_lhs() { + // We need to execute an operation after `from data` to trigger inplace in some backends. + // Which is the operation that might be problematic in this case. + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0], [2.0, 3.0]]) * 1; + let tensor_2 = TestTensor::from([[4.0, 5.0], [6.0, 7.0]]) * 1; + + let output = tensor_1.transpose() + tensor_2; + + output + .into_data() + .assert_eq(&TensorData::from([[4.0, 7.0], [7.0, 10.0]]), false); +} + +#[test] +fn test_add_different_strides_broadcast() { + // We need to execute an operation after `from data` to trigger inplace in some backends. + // Which is the operation that might be problematic in this case. + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0], [2.0, 3.0]]) * 1; + let tensor_2 = TestTensor::from([[4.0, 5.0]]) * 1; + + let output = tensor_1.transpose() + tensor_2; + + output + .into_data() + .assert_eq(&TensorData::from([[4.0, 7.0], [5.0, 8.0]]), false); +} + +#[test] +fn should_support_add_scalar_ops() { + let scalar = 2.0; + let tensor = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor + scalar; + + output + .into_data() + .assert_eq(&TensorData::from([[2.0, 3.0, 4.0], [5.0, 6.0, 7.0]]), false); +} + +#[test] +fn add_maybe_fused_not_contiguous() { + let tensor1 = TestTensorInt::arange(0..8, &Default::default()).float(); + let tensor2 = TestTensorInt::arange(8..16, &Default::default()).float(); + let tensor1 = tensor1.reshape([2, 4]); + let tensor2 = tensor2.reshape([4, 2]); + let tensor2 = tensor2.swap_dims(0, 1); + + tensor2.device().sync().unwrap(); + + let output = tensor1 + tensor2; + + output.into_data().assert_eq( + &TensorData::from([[8.0, 11.0, 14.0, 17.0], [13.0, 16.0, 19.0, 22.0]]), + false, + ); +} + +#[test] +fn test_add_narrowed() { + // Non-contiguous via .narrow(): rows 1-2 of a [4, 4] tensor, then add + // to itself. Exercises stride handling on narrowed inputs. + let data: Vec = (0..16).map(|i| i as f32).collect(); + let device = Default::default(); + let a = TestTensor::<2>::from_data(TensorData::new(data.clone(), [4, 4]), &device); + let b = TestTensor::<2>::from_data(TensorData::new(data, [4, 4]), &device); + + let a_narrow = a.narrow(0, 1, 2); + let b_narrow = b.narrow(0, 1, 2); + + let output = a_narrow + b_narrow; + + output.into_data().assert_eq( + &TensorData::from([[8.0, 10.0, 12.0, 14.0], [16.0, 18.0, 20.0, 22.0]]), + false, + ); +} + +#[test] +fn test_add_scalar_transposed() { + // Scalar add on a non-contiguous (transposed) input. + let a = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]).transpose(); + + let output = a + 10.0; + + // a_t = [[1, 3], [2, 4]] + 10 = [[11, 13], [12, 14]] + output + .into_data() + .assert_eq(&TensorData::from([[11.0, 13.0], [12.0, 14.0]]), false); +} + +#[test] +fn add_maybe_fused_not_contiguous_broadcasted() { + let tensor1 = TestTensorInt::arange(0..8, &Default::default()).float(); + let tensor2 = TestTensorInt::arange(8..10, &Default::default()).float(); + let tensor1 = tensor1.reshape([2, 4]); + let tensor2 = tensor2.reshape([1, 2]); + let tensor2 = tensor2.swap_dims(0, 1); + + tensor2.device().sync().unwrap(); + + let output = tensor2 + tensor1; + + output.into_data().assert_eq( + &TensorData::from([[8.0, 9.0, 10.0, 11.0], [13.0, 14.0, 15.0, 16.0]]), + false, + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/aggregation.rs b/crates/burn-backend-tests/tests/tensor/float/ops/aggregation.rs new file mode 100644 index 0000000..f24484a --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/aggregation.rs @@ -0,0 +1,645 @@ +#![allow(clippy::identity_op)] + +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn test_should_mean() { + let tensor = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.mean(); + let expected = TensorData::from([15.0 / 6.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_should_sum() { + let tensor = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.sum(); + + output + .into_data() + .assert_eq(&TensorData::from([15.0]), false); +} + +#[test] +fn test_should_sum_dim_maybe_fused() { + let tensor = TestTensor::<2>::from([[5.0], [-12.0]]); + let tensor1 = TestTensor::<2>::from([[2.0, 3.0], [-1.0, -5.0]]); + let ones = TestTensor::<2>::ones([2, 2], &Default::default()); + let _x = ones.clone() * tensor; + let y = ones * tensor1; + + let output = y.clone().sum_dim(1); + output + .into_data() + .assert_eq(&TensorData::from([[5.0], [-6.0]]), false); + + // Negative Indexing. + let output = y.clone().sum_dim(-1); + output + .into_data() + .assert_eq(&TensorData::from([[5.0], [-6.0]]), false); +} + +#[test] +fn test_should_mean_last_dim() { + let tensor = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.clone().mean_dim(1); + let expected = TensorData::from([[3.0 / 3.0], [12.0 / 3.0]]); + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + // Negative Indexing. + let output = tensor.clone().mean_dim(-1); + let expected = TensorData::from([[3.0 / 3.0], [12.0 / 3.0]]); + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_should_sum_last_dim() { + let tensor = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.sum_dim(1); + + output + .into_data() + .assert_eq(&TensorData::from([[3.0], [12.0]]), false); +} + +#[test] +fn test_should_sum_first_dim() { + let tensor = TestTensor::<2>::from([[3.0, 1.0, 2.0], [4.0, 2.0, 3.0]]); + + let output = tensor.sum_dim(0); + + output + .into_data() + .assert_eq(&TensorData::from([[7.0, 3.0, 5.0]]), false); +} + +#[test] +fn test_should_mean_first_dim() { + let tensor = TestTensor::<2>::from([[3.0, 1.0, 2.0], [4.0, 2.0, 3.0]]); + + let output = tensor.mean_dim(0); + + output.into_data().assert_eq( + &TensorData::from([[7.0 / 2.0, 3.0 / 2.0, 5.0 / 2.0]]), + false, + ); +} + +#[test] +fn test_should_sum_mid_dim_3d_non_contiguous_1() { + let tensor = TestTensor::<3>::from([ + [[2.0, 4.0, 1.0], [7.0, -5.0, 3.0]], + [[3.0, 1.0, 2.0], [4.0, 2.0, 3.0]], + ]); + + let output = tensor.swap_dims(0, 2).sum_dim(1); + + output.into_data().assert_eq( + &TensorData::new(vec![9.0, 7.0, -1.0, 3.0, 4.0, 5.0], [3, 1, 2]), + false, + ); +} + +#[test] +fn test_should_sum_mid_dim_3d_non_contiguous_2() { + let tensor = TestTensor::<3>::from([ + [[2.0, 4.0, 1.0], [7.0, -5.0, 3.0]], + [[3.0, 1.0, 2.0], [4.0, 2.0, 3.0]], + ]); + + let output = tensor.swap_dims(0, 1).sum_dim(1); + + output.into_data().assert_eq( + &TensorData::new(vec![5.0, 5.0, 3.0, 11.0, -3.0, 6.0], [2, 1, 3]), + false, + ); +} + +#[test] +fn test_prod_float() { + let tensor = TestTensor::<2>::from([[2.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let output = tensor.prod(); + + // 2 * 1 * 2 * 3 * 4 * 5 = 240 but we need to check the precision because of the float + let expected = TensorData::from([240.0]); + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let tensor_with_zero = TestTensor::<2>::from([[2.0, 0.0, 2.0], [3.0, 4.0, 5.0]]); + let output = tensor_with_zero.prod(); + + output + .into_data() + .assert_eq(&TensorData::from([0.0]), false); +} + +#[test] +fn test_prod_dim_float() { + let tensor = TestTensor::<2>::from([[2.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let output = tensor.prod_dim(1); + let expected = TensorData::from([[4.0], [60.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let tensor_with_zero = TestTensor::<2>::from([[2.0, 0.0, 2.0], [3.0, 4.0, 5.0]]); + let output = tensor_with_zero.prod_dim(1); + let expected = TensorData::from([[0.0], [60.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_sum_dim_2d() { + let tensor = + TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &Default::default()); + + let output = tensor.clone().sum_dim(1); + let expected = TensorData::from([[3.], [12.]]); + + output.into_data().assert_eq(&expected, false); + + let output = tensor.sum_dim(0); + let expected = TensorData::from([[3., 5., 7.]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_sum_dims_2d() { + let tensor = + TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &Default::default()); + + tensor + .clone() + .sum_dims(&[1]) + .to_data() + .assert_eq(&TensorData::from([[3.], [12.]]), false); + + tensor + .clone() + .sum_dims(&[-1]) + .to_data() + .assert_eq(&TensorData::from([[3.], [12.]]), false); + + tensor + .clone() + .sum_dims(&[0, 1]) + .to_data() + .assert_eq(&TensorData::from([[15.]]), false); +} + +#[test] +fn test_sum_and_squeeze_dims() { + let tensor = TestTensor::<3>::from_data( + [ + [[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], + [[9.0, 2.0, 5.0], [5.0, 7.0, 7.0]], + ], + &Default::default(), + ); + + tensor + .sum_dims_squeeze::<1, _>(&[0, 1]) + .to_data() + .assert_eq(&TensorData::from([20., 16., 21.]), false); +} + +#[test] +fn test_sum_dim_1_reshape_maybe_fused() { + let tensor = TestTensorInt::arange(0..9, &Default::default()).float(); + tensor.device().sync().unwrap(); + + let output = tensor.reshape([3, 3]) + 2; + let output = output.sum_dim(1); + let expected = TensorData::from([[9.0], [18.0], [27.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_sum_dim_1_swap_dims_maybe_fused() { + // Note: the `+ 2` elementwise op materializes a contiguous intermediate, + // so the subsequent sum_dim sees contiguous strides - it does not + // exercise the transposed-reduce-last-dim path. That path is covered + // by `test_sum_transposed`. + let tensor = TestTensorInt::arange(0..9, &Default::default()).float(); + let tensor = tensor.reshape([3, 3]); + tensor.device().sync().unwrap(); + + let output = tensor.swap_dims(0, 1) + 2; + let output = output.sum_dim(1); + let expected = TensorData::from([[15.0], [18.0], [21.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_sum_dim_2_reshape_maybe_fused_broadcast() { + let tensor = TestTensorInt::arange(0..9, &Default::default()).float(); + tensor.device().sync().unwrap(); + + let output = tensor.reshape([1, 3, 3]) + 2; + let output = output.sum_dim(2); + let expected = TensorData::from([[[9.0], [18.0], [27.0]]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_sum_dim_2_maybe_fused_on_write() { + let tensor_1 = TestTensorInt::arange(0..8, &Default::default()).float(); + let tensor_2 = TestTensorInt::arange(10..12, &Default::default()).float(); + let tensor_1 = tensor_1.reshape([1, 2, 4]); + let tensor_2 = tensor_2.reshape([1, 2, 1]); + tensor_1.device().sync().unwrap(); + + let output = (tensor_1 + tensor_2.clone()).sum_dim(2) + tensor_2; + output.device().sync().unwrap(); + let expected = TensorData::from([[[56.0], [77.0]]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_sum_dim_3_maybe_fused_on_read_not_contiguous() { + let tensor_1 = TestTensorInt::arange(0..8, &Default::default()).float(); + let tensor_2 = TestTensorInt::arange(16..24, &Default::default()).float(); + + let tensor_1 = tensor_1.reshape([4, 2, 1]); + let tensor_1 = tensor_1.swap_dims(0, 2); + + let tensor_2 = tensor_2.reshape([1, 4, 2]); + let tensor_2 = tensor_2.swap_dims(1, 2); + tensor_1.device().sync().unwrap(); + + let output = (tensor_1 + tensor_2).sum_dim(2); + let expected = TensorData::from([[[88.0], [96.0]]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_sum_dim_4_maybe_fused_on_read_not_contiguous_mixed() { + let tensor_1 = TestTensorInt::arange(0..8, &Default::default()).float(); + let tensor_2 = TestTensorInt::arange(16..24, &Default::default()).float(); + let tensor_3 = TestTensorInt::arange(32..40, &Default::default()).float(); + + let tensor_1 = tensor_1.reshape([4, 2, 1]); + let tensor_3 = tensor_3.reshape([1, 2, 4]); + let tensor_1 = tensor_1.swap_dims(0, 2); + + let tensor_2 = tensor_2.reshape([1, 4, 2]); + let tensor_2 = tensor_2.swap_dims(1, 2); + tensor_1.device().sync().unwrap(); + + let output = (tensor_3 + tensor_1 + tensor_2).sum_dim(2); + let expected = TensorData::from([[[222.0], [246.0]]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_sum_dim_5_maybe_fused_on_read_not_contiguous_mixed() { + let tensor_1 = TestTensorInt::arange(0..8, &Default::default()).float(); + let tensor_2 = TestTensorInt::arange(16..24, &Default::default()).float(); + let tensor_3 = TestTensorInt::arange(32..40, &Default::default()).float(); + + let tensor_1 = tensor_1.reshape([4, 2, 1]); + let tensor_3 = tensor_3.reshape([1, 2, 4]); + let tensor_1 = tensor_1.swap_dims(0, 2); + + let tensor_2 = tensor_2.reshape([1, 4, 2]); + let tensor_2 = tensor_2.swap_dims(1, 2); + tensor_1.device().sync().unwrap(); + + let output = (tensor_3 + tensor_1 + tensor_2).sum_dim(1); + let expected = TensorData::from([[[102.0, 112.0, 122.0, 132.0]]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_sum_dim_6_maybe_fused_on_read_not_contiguous_broadcasted() { + let tensor_1 = TestTensorInt::arange(0..32, &Default::default()).float(); + let tensor_2 = TestTensorInt::arange(0..8, &Default::default()).float(); + + let tensor_1 = tensor_1.reshape([4, 2, 2, 2]); + let tensor_1 = tensor_1.swap_dims(3, 2); + let tensor_1 = tensor_1.swap_dims(1, 2); + + let tensor_2 = tensor_2.reshape([1, 2, 2, 2]); + + tensor_1.device().sync().unwrap(); + let sum = tensor_2.clone().sum_dim(0); + let sum = sum.sum_dim(1); + let sum = sum.sum_dim(2); + + tensor_1.device().sync().unwrap(); + + let _tmp = sum.clone() + 2; + let output = (tensor_1 + tensor_2 + sum).sum_dim(1); + let expected = TensorData::from([ + [[[29.0, 43.0], [41.0, 55.0]]], + [[[45.0, 59.0], [57.0, 71.0]]], + [[[61.0, 75.0], [73.0, 87.0]]], + [[[77.0, 91.0], [89.0, 103.0]]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_sum_dim_7_maybe_fused_on_read_reshaped() { + let tensor_1 = TestTensorInt::arange(0..16, &Default::default()).float(); + + let tensor_1 = tensor_1.reshape([4, 4]); + + tensor_1.device().sync().unwrap(); + + let reshaped = tensor_1.reshape([1, 4, 4]); + let tmp = reshaped + 5.0; + let output = tmp.sum_dim(2); + let expected = TensorData::from([[[26.0], [42.0], [58.0], [74.0]]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_reduce_dim_non_contiguous_input() { + let t = TestTensor::<2>::from([ + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + ]); + let t_transposed = t.swap_dims(0, 1) /*+ 0.*/; + + let dim = 1; + let max = t_transposed.clone().max_dim(dim); + let shifted = t_transposed.sub(max); + let exp = shifted.exp(); + let sum = exp.clone().sum_dim(dim); + let output = exp.div(sum); + + let row = [3.2932044e-4, 1.7980287e-2, 0.98169035]; + let expected = TensorData::from([row, row, row, row]); + output.into_data().assert_approx_eq::( + &expected, + Tolerance::default().set_half_precision_absolute(2e-3), + ); +} + +#[test] +fn test_mean_dim_fused_on_read_on_write() { + // https://github.com/tracel-ai/burn/issues/3987 + let device = Default::default(); + let x = TestTensor::ones([128, 32, 1], &device); + + let weight = TestTensor::ones([1, 32, 1], &device); + let options = burn_tensor::ops::ConvOptions::new([1], [0], [1], 1); + let x = burn_tensor::module::conv1d(x, weight, None, options); + let global = x.clone().powi_scalar(2).sum_dim(2).add_scalar(1e-5).sqrt(); + let norm = global.clone().div(global.mean_dim(1)); + let x = x.clone().mul(norm).add(x); + + let out = x.sum(); + + out.into_data() + .assert_eq(&TensorData::from([8192.0]), false); +} + +#[test] +fn test_mean_dim_2d() { + let tensor = + TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &Default::default()); + + let output = tensor.clone().mean_dim(1); + let expected = TensorData::from([[1.], [4.]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let output = tensor.mean_dim(0); + let expected = TensorData::from([[1.5, 2.5, 3.5]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_mean_dims_2d() { + let tensor = + TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &Default::default()); + + tensor + .clone() + .mean_dims(&[1]) + .to_data() + .assert_eq(&TensorData::from([[1.], [4.]]), false); + + tensor + .clone() + .mean_dims(&[-1]) + .to_data() + .assert_eq(&TensorData::from([[1.], [4.]]), false); + + tensor + .clone() + .mean_dims(&[0, 1]) + .to_data() + .assert_eq(&TensorData::from([[2.5]]), false); +} + +#[test] +fn test_multiple_reduce_dims_permuted() { + // Regression test for https://github.com/tracel-ai/burn/issues/4461. + // + // Also pins the f16/bf16 `mean_dim` overflow fix: after the first reduction + // the second `mean_dim` sums 256 values that peak near 1021, so the f32 + // intermediate reaches ~261k (well above f16::MAX = 65504). A naive + // sum-then-divide in f16 would overflow to +inf. See the sibling + // `test_should_mean_overflow_intermediate_sum` for the scalar `.mean()` + // code path. + let tensor = TestTensorInt::arange(0..2 * 2 * 256, &Default::default()) + .float() + .reshape([2, 2, 256]); + + let output = tensor + .permute([1, 2, 0]) + .mean_dim(0) + .mean_dim(1) + .squeeze_dims::<1>(&[0, 1]); + + output + .into_data() + .assert_approx_eq::(&TensorData::from([255.5, 767.5]), Tolerance::default()); +} + +#[test] +fn test_sum_transposed() { + // Stress the sum kernel on a non-contiguous (transposed) input. Reducing + // the (now last) dim with stride != 1 previously hit a fast path that + // read contiguous storage, producing sliding-pair sums instead of column + // sums. + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let output = tensor.transpose().sum_dim(1); + + output + .into_data() + .assert_eq(&TensorData::from([[5.0], [7.0], [9.0]]), false); +} + +#[test] +fn test_prod_dim_transposed() { + // Same fast-path gate as sum_dim: prod_dim on a transposed 2D input must + // honor the reduce-dim stride. + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let output = tensor.transpose().prod_dim(1); + + output.into_data().assert_approx_eq::( + &TensorData::from([[4.0], [10.0], [18.0]]), + Tolerance::default(), + ); +} + +#[test] +fn test_mean_dim_transposed() { + // mean_dim routes through sum_dim, so it inherits the same gate. + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let output = tensor.transpose().mean_dim(1); + + output.into_data().assert_approx_eq::( + &TensorData::from([[2.5], [3.5], [4.5]]), + Tolerance::default(), + ); +} + +#[test] +fn test_sum_flipped() { + // Flip axis 0, sum along axis 1: per-row sums appear in the flipped + // row order, so a no-op flip would give the unflipped order. + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]]); + let output = tensor.flip([0]).sum_dim(1); + + output + .into_data() + .assert_eq(&TensorData::from([[60.0], [6.0]]), false); +} + +#[test] +fn test_sum_dim_flipped() { + // Flip axis 0, reduce axis 1: rows are swapped so per-row sums appear + // reversed, which a flip-as-noop would not produce. + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let output = tensor.flip([0]).sum_dim(1); + + output + .into_data() + .assert_eq(&TensorData::from([[15.0], [6.0]]), false); +} + +#[test] +fn test_sum_dim_flipped_axis1() { + // Flip axis 1, reduce axis 0: columns are reversed, so column sums + // appear in reversed order. + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let output = tensor.flip([1]).sum_dim(0); + + output + .into_data() + .assert_eq(&TensorData::from([[9.0, 7.0, 5.0]]), false); +} + +#[test] +fn test_mean_dim_flipped() { + // Flip axis 0, mean axis 1: row means appear in reversed row order. + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let output = tensor.flip([0]).mean_dim(1); + + output + .into_data() + .assert_approx_eq::(&TensorData::from([[5.0], [2.0]]), Tolerance::default()); +} + +#[test] +fn test_prod_flipped() { + // Flip axis 0, prod along axis 1: per-row products appear in reversed + // row order. + let tensor = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]); + let output = tensor.flip([0]).prod_dim(1); + + output + .into_data() + .assert_approx_eq::(&TensorData::from([[12.0], [2.0]]), Tolerance::default()); +} + +#[test] +fn test_sum_narrowed() { + // Narrow to indices 1..4 of [0, 1, 2, 3, 4] -> [1, 2, 3]. Without the + // narrow the sum would be 10, not 6. + let tensor = TestTensor::<1>::from([0.0, 1.0, 2.0, 3.0, 4.0]); + let output = tensor.narrow(0, 1, 3).sum(); + + output + .into_data() + .assert_eq(&TensorData::from([6.0]), false); +} + +#[test] +fn test_sum_flipped_both_axes() { + // Flip both axes, sum along axis 0: column sums appear in reversed + // order because axis 1 was flipped; row pairing is also swapped so a + // missing axis-0 flip would give different pairings. + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let output = tensor.flip([0, 1]).sum_dim(0); + + output + .into_data() + .assert_eq(&TensorData::from([[9.0, 7.0, 5.0]]), false); +} + +#[test] +fn test_sum_dim_4d_middle_dim() { + // Regression: 4D tensor reducing a middle dim (shape [1, 84, 80, 80]). + // Fill with 1.0 so every output position should sum to 84.0. + let shape = [1, 84, 80, 80]; + let n: usize = shape.iter().product(); + let data: Vec = vec![1.0; n]; + let tensor = TestTensor::<4>::from_data(TensorData::new(data, shape), &Default::default()); + + let output = tensor.sum_dim(1); + + let expected = TensorData::new(vec![84.0f32; 1 * 1 * 80 * 80], [1, 1, 80, 80]); + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_should_mean_overflow_intermediate_sum() { + let tensor = TestTensorInt::arange(0..1024, &Default::default()).float(); + let output = tensor.mean(); + output + .into_data() + .assert_approx_eq::(&TensorData::from([511.5]), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/all.rs b/crates/burn-backend-tests/tests/tensor/float/ops/all.rs new file mode 100644 index 0000000..bd6ae0e --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/all.rs @@ -0,0 +1,47 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_all() { + let tensor = TestTensor::<2>::from([[0.0, 1.0, 0.0], [1.0, -1.0, 1.0]]); + let data_actual = tensor.all().into_data(); + let data_expected = TensorData::from([false]); + data_expected.assert_eq(&data_actual, false); +} + +#[test] +fn test_all_dim() { + let tensor = TestTensor::<2>::from([[0.0, 1.0, 0.0], [1.0, -1.0, 1.0]]); + let data_actual = tensor.all_dim(1).into_data(); + let data_expected = TensorData::from([[false], [true]]); + data_expected.assert_eq(&data_actual, false); +} + +/// Larger `all_dim` over a long, non-power-of-two axis so the reduction spans +/// multiple blocks/planes/cubes. Same mask as the large `any_dim` test; here a +/// row reduces to `true` only when every element is nonzero (so the single-zero +/// row 7 must reduce to `false`). Checked for float, int and bool inputs. +#[test] +fn test_all_dim_large() { + let device = Default::default(); + let (rows, cols) = (8usize, 513usize); + let mask = build_logical_mask(rows, cols); + + let expected: Vec = (0..rows) + .map(|r| (0..cols).all(|c| mask[r * cols + c])) + .collect(); + let expected = TensorData::new(expected, [rows, 1]); + + let float = TestTensor::<2>::from_data( + TensorData::new(mask_to_floats(&mask), [rows, cols]), + &device, + ); + expected.assert_eq(&float.all_dim(1).into_data(), false); + + let int = + TestTensorInt::<2>::from_data(TensorData::new(mask_to_ints(&mask), [rows, cols]), &device); + expected.assert_eq(&int.all_dim(1).into_data(), false); + + let boolean = TestTensorBool::<2>::from_data(TensorData::new(mask, [rows, cols]), &device); + expected.assert_eq(&boolean.all_dim(1).into_data(), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/any.rs b/crates/burn-backend-tests/tests/tensor/float/ops/any.rs new file mode 100644 index 0000000..b17839e --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/any.rs @@ -0,0 +1,87 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_any() { + // test float tensor + let tensor = TestTensor::<2>::from([[0.0, 0.0, 0.0], [1.0, -1.0, 0.0]]); + let data_actual = tensor.any().into_data(); + let data_expected = TensorData::from([true]); + data_expected.assert_eq(&data_actual, false); + + let tensor = TestTensor::<2>::from([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]); + let data_actual = tensor.any().into_data(); + let data_expected = TensorData::from([false]); + data_expected.assert_eq(&data_actual, false); + + // test int tensor + let tensor = TestTensorInt::<2>::from([[0, 0, 0], [1, -1, 0]]); + let data_actual = tensor.any().into_data(); + let data_expected = TensorData::from([true]); + data_expected.assert_eq(&data_actual, false); + + let tensor = TestTensorInt::<2>::from([[0, 0, 0], [0, 0, 0]]); + let data_actual = tensor.any().into_data(); + let data_expected = TensorData::from([false]); + data_expected.assert_eq(&data_actual, false); + + // test bool tensor + let tensor = TestTensorBool::<2>::from([[false, false, false], [true, true, false]]); + let data_actual = tensor.any().into_data(); + let data_expected = TensorData::from([true]); + data_expected.assert_eq(&data_actual, false); + + let tensor = TestTensorBool::<2>::from([[false, false, false], [false, false, false]]); + let data_actual = tensor.any().into_data(); + let data_expected = TensorData::from([false]); + data_expected.assert_eq(&data_actual, false); +} + +#[test] +fn test_any_dim() { + let tensor = TestTensor::<2>::from([[0.0, 0.0, 0.0], [1.0, -1.0, 0.0]]); + let data_actual = tensor.any_dim(1).into_data(); + let data_expected = TensorData::from([[false], [true]]); + data_expected.assert_eq(&data_actual, false); + + // test int tensor + let tensor = TestTensorInt::<2>::from([[0, 0, 0], [1, -1, 0]]); + let data_actual = tensor.any_dim(1).into_data(); + let data_expected = TensorData::from([[false], [true]]); + data_expected.assert_eq(&data_actual, false); + + // test bool tensor + let tensor = TestTensorBool::<2>::from([[false, false, false], [true, true, false]]); + let data_actual = tensor.any_dim(1).into_data(); + let data_expected = TensorData::from([[false], [true]]); + data_expected.assert_eq(&data_actual, false); +} + +/// Larger `any_dim` over a long, non-power-of-two axis so the reduction spans +/// multiple blocks/planes/cubes. Rows cover every merge case: all-zero, +/// all-nonzero, a single nonzero at the first/last/middle column, and an +/// all-nonzero row with a single zero. Checked for float, int and bool inputs. +#[test] +fn test_any_dim_large() { + let device = Default::default(); + let (rows, cols) = (8usize, 513usize); + let mask = build_logical_mask(rows, cols); + + let expected: Vec = (0..rows) + .map(|r| (0..cols).any(|c| mask[r * cols + c])) + .collect(); + let expected = TensorData::new(expected, [rows, 1]); + + let float = TestTensor::<2>::from_data( + TensorData::new(mask_to_floats(&mask), [rows, cols]), + &device, + ); + expected.assert_eq(&float.any_dim(1).into_data(), false); + + let int = + TestTensorInt::<2>::from_data(TensorData::new(mask_to_ints(&mask), [rows, cols]), &device); + expected.assert_eq(&int.any_dim(1).into_data(), false); + + let boolean = TestTensorBool::<2>::from_data(TensorData::new(mask, [rows, cols]), &device); + expected.assert_eq(&boolean.any_dim(1).into_data(), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/arg.rs b/crates/burn-backend-tests/tests/tensor/float/ops/arg.rs new file mode 100644 index 0000000..c86023a --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/arg.rs @@ -0,0 +1,212 @@ +#![allow(clippy::identity_op)] + +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_argmax_2d_dim0() { + let tensor = TestTensor::<2>::from([[10.0, 11.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.argmax(0); + + output + .into_data() + .assert_eq(&TensorData::from([[0, 0, 1]]), false); +} + +#[test] +fn test_argmin_2d_dim0() { + let tensor = TestTensor::<2>::from([[10.0, 11.0, 2.0], [30.0, 4.0, 5.0]]); + + let output = tensor.argmin(0); + + output + .into_data() + .assert_eq(&TensorData::from([[0, 1, 0]]), false); +} + +#[test] +fn test_argmax_2d_dim1() { + let tensor = TestTensor::<2>::from([[10.0, 11.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.argmax(1); + + output + .into_data() + .assert_eq(&TensorData::from([[1], [2]]), false); +} + +#[test] +fn test_argmin_2d_dim1() { + let tensor = TestTensor::<2>::from([[10.0, 11.0, 2.0], [30.0, 4.0, 5.0]]); + + let output = tensor.argmin(1); + + output + .into_data() + .assert_eq(&TensorData::from([[2], [1]]), false); +} + +#[test] +fn test_argmax_flipped() { + // Flip [1, 5, 3, 2, 4] -> [4, 2, 3, 5, 1]; max is at index 3. + let tensor = TestTensor::<1>::from([1.0, 5.0, 3.0, 2.0, 4.0]); + let output = tensor.flip([0]).argmax(0); + + output.into_data().assert_eq(&TensorData::from([3]), false); +} + +#[test] +fn test_argmax_2d_flipped() { + // [[1, 5, 3], [6, 2, 4]] axis-1-flipped -> [[3, 5, 1], [4, 2, 6]]; argmax dim 1 -> [[1], [2]]. + let tensor = TestTensor::<2>::from([[1.0, 5.0, 3.0], [6.0, 2.0, 4.0]]); + let output = tensor.flip([1]).argmax(1); + + output + .into_data() + .assert_eq(&TensorData::from([[1], [2]]), false); +} + +#[test] +fn test_argmin_flipped() { + // Flip [5, 1, 4, 2, 3] -> [3, 2, 4, 1, 5]; min is at index 3. + let tensor = TestTensor::<1>::from([5.0, 1.0, 4.0, 2.0, 3.0]); + let output = tensor.flip([0]).argmin(0); + + output.into_data().assert_eq(&TensorData::from([3]), false); +} + +#[test] +fn test_argmax_permuted_4d() { + // Regression: argmax on a permuted 4D tensor (was index OOB). + let n = 2 * 3 * 4 * 5; + let data: Vec = (0..n).map(|i| i as f32).collect(); + let tensor = + TestTensor::<4>::from_data(TensorData::new(data, [2, 3, 4, 5]), &Default::default()); + let permuted = tensor.permute([0, 2, 1, 3]); + + let result = permuted.clone().argmax(3); + assert_eq!(result.dims(), [2, 4, 3, 1]); + + let result = permuted.argmax(2); + assert_eq!(result.dims(), [2, 4, 1, 5]); +} + +#[test] +fn test_argmin_permuted_4d() { + let n = 2 * 3 * 4 * 5; + let data: Vec = (0..n).map(|i| (n - i) as f32).collect(); + let tensor = + TestTensor::<4>::from_data(TensorData::new(data, [2, 3, 4, 5]), &Default::default()); + let permuted = tensor.permute([0, 2, 1, 3]); + + let result = permuted.argmin(3); + assert_eq!(result.dims(), [2, 4, 3, 1]); +} + +#[test] +fn test_argmax_4d_middle_dim() { + // Regression (YOLOv8n): shape [1, 84, 80, 80], argmax dim=1. + let n = 1 * 84 * 80 * 80; + let data: Vec = (0..n).map(|i| (i % 84) as f32).collect(); + let tensor = + TestTensor::<4>::from_data(TensorData::new(data, [1, 84, 80, 80]), &Default::default()); + + let output = tensor.argmax(1); + assert_eq!(output.dims(), [1, 1, 80, 80]); +} + +#[test] +fn test_argmax_permuted_correctness() { + // Data [2, 2, 3] permuted [0, 2, 1] -> [2, 3, 2]; argmax dim 2 should be all 1s. + let data: Vec = (1..=12).map(|i| i as f32).collect(); + let tensor = TestTensor::<3>::from_data(TensorData::new(data, [2, 2, 3]), &Default::default()); + let output = tensor.permute([0, 2, 1]).argmax(2); + + output + .into_data() + .assert_eq(&TensorData::from([[[1], [1], [1]], [[1], [1], [1]]]), false); +} + +// NaN-propagation tests below. Only run when the `flex` backend feature +// is active, because flex is the only burn backend that currently +// propagates NaN from argmax/argmin (matching PyTorch/NumPy/JAX/TF). +// ndarray and the cubecl backends follow IEEE 754 min/max and drop NaN. +// The positive-gate form (rather than excluding specific backends) is +// used because the default-feature CI build selects a backend +// transitively without setting any of its identifying feature flags on +// burn-backend-tests, so a negative gate would still run the test on a +// NaN-dropping backend. See issue #4814. +#[cfg(feature = "flex")] +#[test] +fn test_argmax_nan_propagation() { + let tensor = TestTensor::<2>::from([[1.0, f32::NAN, 3.0]]); + let output = tensor.argmax(1); + + output + .into_data() + .assert_eq(&TensorData::from([[1]]), false); +} + +// First-NaN wins: when row[0] is NaN AND a later element is also NaN, +// argmax must return 0 (the earlier NaN index), not the later one. +#[cfg(feature = "flex")] +#[test] +fn test_argmax_nan_leading_with_trailing_nan() { + let tensor = TestTensor::<2>::from([[f32::NAN, f32::NAN, 3.0]]); + let output = tensor.argmax(1); + output + .into_data() + .assert_eq(&TensorData::from([[0]]), false); +} + +#[cfg(feature = "flex")] +#[test] +fn test_argmin_nan_leading_with_trailing_nan() { + let tensor = TestTensor::<2>::from([[f32::NAN, f32::NAN, 3.0]]); + let output = tensor.argmin(1); + output + .into_data() + .assert_eq(&TensorData::from([[0]]), false); +} + +// All-NaN row: argmax should return 0 (first NaN). +#[cfg(feature = "flex")] +#[test] +fn test_argmax_all_nan_row() { + let tensor = TestTensor::<2>::from([[f32::NAN, f32::NAN, f32::NAN]]); + let output = tensor.argmax(1); + output + .into_data() + .assert_eq(&TensorData::from([[0]]), false); +} + +// NaN propagation must also hold for non-last-dim argmax, which routes +// through a different kernel than the last-dim fast path. +#[cfg(feature = "flex")] +#[test] +fn test_argmax_nan_leading_non_last_dim() { + // Column 0: [NaN, NaN, 3.0] -> argmax along dim 0 = 0. + // Column 1: [1.0, 2.0, 4.0] -> argmax along dim 0 = 2. + let tensor = TestTensor::<2>::from([[f32::NAN, 1.0], [f32::NAN, 2.0], [3.0, 4.0]]); + let output = tensor.argmax(0); + output + .into_data() + .assert_eq(&TensorData::from([[0, 2]]), false); +} + +// max_dim_with_indices on a leading-NaN row should return (NaN, 0). +// Complements test_max_dim_with_indices_nan_propagation in maxmin.rs, +// which uses a single NaN in the middle of the row. +#[cfg(feature = "flex")] +#[test] +fn test_max_dim_with_indices_nan_leading() { + let tensor = TestTensor::<2>::from([[f32::NAN, f32::NAN, 3.0]]); + let (values, indices) = tensor.max_dim_with_indices(1); + let vdata = values.into_data(); + let slice = vdata.as_slice::().unwrap(); + assert!(slice[0].is_nan()); + indices + .into_data() + .assert_eq(&TensorData::from([[0]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/blackman_window.rs b/crates/burn-backend-tests/tests/tensor/float/ops/blackman_window.rs new file mode 100644 index 0000000..325a056 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/blackman_window.rs @@ -0,0 +1,65 @@ +use super::*; +use burn_tensor::signal::blackman_window; +use burn_tensor::{DType, TensorData, Tolerance}; + +#[test] +fn should_support_blackman_window_options_dtype() { + let tensor: TestTensor<1> = blackman_window(4, true, (&Default::default(), DType::F32)); + assert_eq!(tensor.dtype(), DType::F32); +} + +#[test] +fn should_support_blackman_window_size_0_symmetric() { + let tensor: TestTensor<1> = blackman_window(0, false, &Default::default()); + assert_eq!(tensor.dims(), [0]); +} + +#[test] +fn should_support_blackman_window_size_0_periodic() { + let tensor: TestTensor<1> = blackman_window(0, true, &Default::default()); + assert_eq!(tensor.dims(), [0]); +} + +#[test] +fn should_handle_blackman_window_size_1_symmetric() { + let tensor: TestTensor<1> = blackman_window(1, false, &Default::default()); + let expected = TensorData::from([1.0]); + + tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_handle_blackman_window_size_1_periodic() { + let tensor: TestTensor<1> = blackman_window(1, true, &Default::default()); + let expected = TensorData::from([1.0]); + + tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_blackman_window_size_8_periodic() { + let tensor: TestTensor<1> = blackman_window(8, true, &Default::default()); + let expected = TensorData::from([ + 0.0, 6.6447e-02, 3.4000e-01, 7.7355e-01, 1.0000e+00, 7.7355e-01, 3.4000e-01, 6.6447e-02, + ]); + + // Positions 0 and 7 have values close to 0, hence setting absolute tolerance + let tolerance = Tolerance::default().set_half_precision_absolute(1e-3); + tensor + .into_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn should_support_blackman_window_size_8_symmetric() { + let tensor: TestTensor<1> = blackman_window(8, false, &Default::default()); + let expected = TensorData::from([ + 0.0, 9.0453e-02, 4.5918e-01, 9.2036e-01, 9.2036e-01, 4.5918e-01, 9.0453e-02, 0.0, + ]); + + // Positions 0 and 7 have values close to 0, hence setting absolute tolerance + let tolerance = Tolerance::default().set_half_precision_absolute(1e-3); + tensor + .into_data() + .assert_approx_eq::(&expected, tolerance); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/cast.rs b/crates/burn-backend-tests/tests/tensor/float/ops/cast.rs new file mode 100644 index 0000000..e00bb35 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/cast.rs @@ -0,0 +1,91 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{DType, FloatDType, IntDType, TensorData}; + +#[test] +fn cast_float_to_bool() { + let tensor1 = TestTensor::<2>::from([[0.0, 43.0, 0.0], [2.0, -4.2, 31.33]]); + let data_actual = tensor1.bool().into_data(); + let data_expected = TensorData::from([[false, true, false], [true, true, true]]); + data_actual.assert_eq(&data_expected, false); +} + +#[test] +fn cast_float_to_int() { + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.4, 5.5, 6.6]]).int(); + let expected = TensorData::from([[1, 2, 3], [4, 5, 6]]); + + tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn cast_int_to_float_tensor() { + let tensor = TestTensorInt::<2>::from([[1, 2, 3], [4, 5, 6]]).float(); + + let expected = TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn cast_bool_to_float_tensor() { + let tensor = TestTensorBool::<2>::from([[true, false, true], [false, false, true]]).float(); + + let expected = TensorData::from([[1., 0., 1.], [0., 0., 1.]]); + + tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn cast_float_precision() { + let data = TensorData::from([[1.0, 2.0, 3.0], [4.4, 5.5, 6.6]]); + let tensor = TestTensor::<2>::from(data.clone()); + + let output = tensor.cast(DType::F32); + + assert_eq!(output.dtype(), DType::F32); + // Use precision 2 for parameterized tests in f16 and bf16 + output + .into_data() + .assert_approx_eq::(&data, Tolerance::default()); +} + +#[test] +fn cast_float_to_int_with_dtype() { + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.4, 5.5, 6.6]]); + + let output = tensor.cast(IntDType::I32); + + assert_eq!(output.dtype(), DType::I32); + let expected = TensorData::from([[1i32, 2, 3], [4, 5, 6]]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn cast_float_within_kind_with_float_dtype() { + let tensor = TestTensor::<1>::from([1.0, 2.5, 3.7]); + + let output = tensor.cast(FloatDType::F32); + + assert_eq!(output.dtype(), DType::F32); +} + +#[test] +fn cast_float_same_dtype_is_noop() { + let data = TensorData::from([1.0, 2.0, 3.0]); + let tensor = TestTensor::<1>::from(data.clone()); + let original_dtype = tensor.dtype(); + let float_dtype: FloatDType = original_dtype.into(); + + let output = tensor.cast(float_dtype); + + assert_eq!(output.dtype(), original_dtype); + output.into_data().assert_eq(&data, false); +} + +#[test] +#[should_panic] +fn cast_float_with_int_dtype_panics() { + let tensor = TestTensor::<1>::from([1.0, 2.0]); + let _ = tensor.cast(DType::I32); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/cat.rs b/crates/burn-backend-tests/tests/tensor/float/ops/cat.rs new file mode 100644 index 0000000..7c2d12b --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/cat.rs @@ -0,0 +1,148 @@ +use super::*; +use alloc::vec::Vec; +use burn_tensor::Tolerance; +use burn_tensor::{DType, TensorData}; + +#[test] +fn should_support_cat_ops_2d_dim0() { + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data([[1.0, 2.0, 3.0]], &device); + let tensor_2 = TestTensor::from_data([[4.0, 5.0, 6.0]], &device); + + let output = TestTensor::cat(vec![tensor_1, tensor_2], 0); + let expected = TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_cat_ops_2d_dim1() { + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data([[1.0, 2.0, 3.0]], &device); + let tensor_2 = TestTensor::from_data([[4.0, 5.0, 6.0]], &device); + + let output = TestTensor::cat(vec![tensor_1, tensor_2], 1); + let expected = TensorData::from([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_cat_ops_3d() { + let device = Default::default(); + let tensor_1 = TestTensor::<3>::from_data([[[1.0, 2.0, 3.0]], [[1.1, 2.1, 3.1]]], &device); + let tensor_2 = TestTensor::from_data([[[4.0, 5.0, 6.0]]], &device); + + let output = TestTensor::cat(vec![tensor_1, tensor_2], 0); + let expected = TensorData::from([[[1.0, 2.0, 3.0]], [[1.1, 2.1, 3.1]], [[4.0, 5.0, 6.0]]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +#[should_panic] +fn should_panic_when_dimensions_are_not_the_same() { + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]], &device); + let tensor_2 = TestTensor::from_data([[4.0, 5.0]], &device); + + TestTensor::cat(vec![tensor_1, tensor_2], 0).into_data(); +} + +#[test] +#[should_panic] +fn should_panic_when_list_of_vectors_is_empty() { + let tensor: Vec> = vec![]; + TestTensor::cat(tensor, 0).into_data(); +} + +#[test] +#[should_panic] +fn should_panic_when_cat_exceeds_dimension() { + let device = Default::default(); + let tensor_1 = TestTensor::<3>::from_data([[[1.0, 2.0, 3.0]], [[1.1, 2.1, 3.1]]], &device); + let tensor_2 = TestTensor::from_data([[[4.0, 5.0, 6.0]]], &device); + + TestTensor::cat(vec![tensor_1, tensor_2], 3).into_data(); +} + +#[test] +fn should_support_cat_ops_cast_dtype() { + let device = Default::default(); + // ok for f32 backends, casts dtype for f16 tests + let tensor_1 = TestTensor::<3>::from_data([[[1.0, 2.0, 3.0]], [[1.1, 2.1, 3.1]]], &device) + .cast(DType::F32); + let tensor_2 = TestTensor::from_data([[[4.0, 5.0, 6.0]]], &device).cast(DType::F32); + + let output = TestTensor::cat(vec![tensor_1, tensor_2], 0); + let expected = TensorData::from([[[1.0, 2.0, 3.0]], [[1.1, 2.1, 3.1]], [[4.0, 5.0, 6.0]]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_cat_with_empty_tensor() { + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data([[1.0, 2.0, 3.0]], &device); + let tensor_2: TestTensor<2> = TestTensor::empty([1, 0], &device); // Empty tensor with size 0 on dim 1 + + // Concatenating with an empty tensor should just return the non-empty tensor + let output = TestTensor::cat(vec![tensor_1.clone(), tensor_2], 1); + let expected = TensorData::from([[1.0, 2.0, 3.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_cat_with_empty_tensor_first() { + let device = Default::default(); + let tensor_1: TestTensor<2> = TestTensor::empty([1, 0], &device); // Empty tensor + let tensor_2 = TestTensor::<2>::from_data([[4.0, 5.0, 6.0]], &device); + + // Empty tensor first, then non-empty + let output = TestTensor::cat(vec![tensor_1, tensor_2.clone()], 1); + let expected = TensorData::from([[4.0, 5.0, 6.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_cat_with_multiple_empty_tensors() { + let device = Default::default(); + let tensor_1: TestTensor<2> = TestTensor::empty([2, 0], &device); + let tensor_2 = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + let tensor_3: TestTensor<2> = TestTensor::empty([2, 0], &device); + let tensor_4 = TestTensor::<2>::from_data([[5.0], [6.0]], &device); + + // Mix of empty and non-empty tensors + let output = TestTensor::cat(vec![tensor_1, tensor_2, tensor_3, tensor_4], 1); + let expected = TensorData::from([[1.0, 2.0, 5.0], [3.0, 4.0, 6.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_cat_all_empty_tensors() { + let device = Default::default(); + let tensor_1: TestTensor<2> = TestTensor::empty([2, 0], &device); + let tensor_2: TestTensor<2> = TestTensor::empty([2, 0], &device); + + // All empty tensors should produce an empty tensor + let output = TestTensor::cat(vec![tensor_1, tensor_2], 1); + + assert_eq!(output.shape().as_slice(), [2, 0]); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/categorical.rs b/crates/burn-backend-tests/tests/tensor/float/ops/categorical.rs new file mode 100644 index 0000000..c488f7b --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/categorical.rs @@ -0,0 +1,128 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn categorical_output_shape() { + let probs = TestTensor::<2>::from([[0.5, 0.5], [0.5, 0.5]]); + let samples = probs.categorical(5); + + assert_eq!(samples.dims(), [2, 5]); +} + +#[test] +fn categorical_1d() { + // 1D input: single distribution + let probs = TestTensor::<1>::from([0.0, 0.0, 1.0]); + let samples = probs.categorical(5); + + assert_eq!(samples.dims(), [5]); + let data = samples.into_data(); + let expected = TensorData::from([2, 2, 2, 2, 2i64]); + data.assert_eq(&expected, false); +} + +#[test] +fn categorical_3d() { + // 3D input: [2, 1, 3] — two batches, each with one sub-batch of 3 categories + let device = Default::default(); + let probs = TestTensor::<3>::from_data( + TensorData::from([[[1.0, 0.0, 0.0]], [[0.0, 0.0, 1.0]]]), + &device, + ); + let samples = probs.categorical(4); + + assert_eq!(samples.dims(), [2, 1, 4]); + let data = samples.into_data(); + let expected = TensorData::from([[[0, 0, 0, 0i64]], [[2, 2, 2, 2]]]); + data.assert_eq(&expected, false); +} + +#[test] +fn categorical_deterministic_single_category() { + let probs = TestTensor::<2>::from([[1.0, 0.0, 0.0]]); + let samples = probs.categorical(10); + + let data = samples.into_data(); + let expected = TensorData::from([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0i64]]); + data.assert_eq(&expected, false); +} + +#[test] +fn categorical_deterministic_last_category() { + let probs = TestTensor::<2>::from([[0.0, 0.0, 1.0]]); + let samples = probs.categorical(10); + + let data = samples.into_data(); + let expected = TensorData::from([[2, 2, 2, 2, 2, 2, 2, 2, 2, 2i64]]); + data.assert_eq(&expected, false); +} + +#[test] +fn categorical_indices_in_range() { + let num_categories = 5; + let probs = TestTensor::<2>::from([[0.2, 0.2, 0.2, 0.2, 0.2]]); + let samples = probs.categorical(100); + + let data = samples.into_data(); + data.assert_within_range::(0..num_categories); +} + +#[test] +fn categorical_batch_deterministic() { + // Row 0: always index 0, Row 1: always index 2 + let probs = TestTensor::<2>::from([[1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]); + let samples = probs.categorical(5); + + let data = samples.into_data(); + let expected = TensorData::from([[0, 0, 0, 0, 0i64], [2, 2, 2, 2, 2]]); + data.assert_eq(&expected, false); +} + +#[test] +fn categorical_single_category_always_zero() { + let probs = TestTensor::<2>::from([[1.0]]); + let samples = probs.categorical(5); + + let data = samples.into_data(); + let expected = TensorData::from([[0, 0, 0, 0, 0i64]]); + data.assert_eq(&expected, false); +} + +#[test] +fn categorical_unnormalized_weights() { + // Weights [0, 0, 10] — should always sample index 2 after normalization + let probs = TestTensor::<2>::from([[0.0, 0.0, 10.0]]); + let samples = probs.categorical(10); + + let data = samples.into_data(); + let expected = TensorData::from([[2, 2, 2, 2, 2, 2, 2, 2, 2, 2i64]]); + data.assert_eq(&expected, false); +} + +#[test] +fn categorical_statistical_distribution() { + // With equal probabilities [0.5, 0.5], samples should be roughly evenly split + let num_samples = 1000; + let probs = TestTensor::<2>::from([[0.5, 0.5]]); + let samples = probs.categorical(num_samples); + + let data: TensorData = samples.into_data(); + let values = data.to_vec::().unwrap(); + + let count_zero = values.iter().filter(|&&v| v == 0).count(); + let ratio = count_zero as f64 / num_samples as f64; + + // Expect ~50% zeros, allow wide tolerance for stochastic test + assert!( + ratio > 0.3 && ratio < 0.7, + "expected ~50% index-0 samples, got {:.1}%", + ratio * 100.0 + ); +} + +#[test] +#[should_panic] +fn categorical_zero_samples_panics() { + let probs = TestTensor::<2>::from([[0.5, 0.5]]); + let _ = probs.categorical(0); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/ceil.rs b/crates/burn-backend-tests/tests/tensor/float/ops/ceil.rs new file mode 100644 index 0000000..8164d3e --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/ceil.rs @@ -0,0 +1,16 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_ceil_ops() { + let data = TensorData::from([[24.0423, 87.9478, 76.1838], [59.6929, 43.8169, 94.8826]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.ceil(); + let expected = TensorData::from([[25., 88., 77.], [60., 44., 95.]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/chunk.rs b/crates/burn-backend-tests/tests/tensor/float/ops/chunk.rs new file mode 100644 index 0000000..1c93ee6 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/chunk.rs @@ -0,0 +1,85 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_chunk_evenly_divisible() { + let tensors = TestTensorInt::arange(0..12, &Default::default()) + .float() + .chunk(6, 0); + assert_eq!(tensors.len(), 6); + + let expected = [ + TensorData::from([0, 1]), + TensorData::from([2, 3]), + TensorData::from([4, 5]), + TensorData::from([6, 7]), + TensorData::from([8, 9]), + TensorData::from([10, 11]), + ]; + + for (index, tensor) in tensors.iter().enumerate() { + tensor.to_data().assert_eq(&expected[index], false); + } +} + +#[test] +fn test_chunk_not_evenly_divisible() { + let tensors = TestTensorInt::arange(0..11, &Default::default()) + .float() + .chunk(6, 0); + assert_eq!(tensors.len(), 6); + + let expected = [ + TensorData::from([0, 1]), + TensorData::from([2, 3]), + TensorData::from([4, 5]), + TensorData::from([6, 7]), + TensorData::from([8, 9]), + TensorData::from([10]), + ]; + + for (index, tensor) in tensors.iter().enumerate() { + tensor.to_data().assert_eq(&expected[index], false); + } +} + +#[test] +fn test_chunk_not_evenly_divisible_remains_several() { + let tensors = TestTensorInt::arange(0..100, &Default::default()) + .float() + .chunk(8, 0); + assert_eq!(tensors.len(), 8); + + let expected = [13, 13, 13, 13, 13, 13, 13, 9]; + + for (index, tensor) in tensors.iter().enumerate() { + assert_eq!(tensor.shape()[0], expected[index]); + } +} + +#[test] +fn test_chunk_not_divisible() { + let tensors = TestTensorInt::arange(0..6, &Default::default()) + .float() + .chunk(7, 0); + assert_eq!(tensors.len(), 6); + + let expected = [ + TensorData::from([0]), + TensorData::from([1]), + TensorData::from([2]), + TensorData::from([3]), + TensorData::from([4]), + TensorData::from([5]), + ]; + + for (index, tensor) in tensors.iter().enumerate() { + tensor.to_data().assert_eq(&expected[index], false); + } +} + +#[test] +#[should_panic] +fn test_invalid_dim() { + let _tensors = TestTensorInt::arange(0..12, &Default::default()).chunk(6, 1); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/clamp.rs b/crates/burn-backend-tests/tests/tensor/float/ops/clamp.rs new file mode 100644 index 0000000..9c15c00 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/clamp.rs @@ -0,0 +1,81 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn clamp_min() { + let device = Default::default(); + // test float tensor + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &device); + + let output = tensor.clamp_min(2.0); + + output + .into_data() + .assert_eq(&TensorData::from([[2.0, 2.0, 2.0], [3.0, 4.0, 5.0]]), false); + + // test int tensor + let data = TensorData::from([[0, 1, 2], [3, 4, 5]]); + let tensor = TestTensorInt::<2>::from_data(data, &device); + let output = tensor.clamp_min(2); + + output + .into_data() + .assert_eq(&TensorData::from([[2, 2, 2], [3, 4, 5]]), false); +} + +#[test] +fn clamp_max() { + let device = Default::default(); + // test float tensor + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &device); + + let output = tensor.clamp_max(2.0); + + output + .into_data() + .assert_eq(&TensorData::from([[0.0, 1.0, 2.0], [2.0, 2.0, 2.0]]), false); + + // test int tensor + let data = TensorData::from([[0, 1, 2], [3, 4, 5]]); + let tensor = TestTensorInt::<2>::from_data(data, &device); + let output = tensor.clamp_max(4); + + output + .into_data() + .assert_eq(&TensorData::from([[0, 1, 2], [3, 4, 4]]), false); +} + +#[test] +fn clamp_min_max() { + let device = Default::default(); + // test float tensor + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &device); + let output = tensor.clamp(1.0, 4.0); + + output + .into_data() + .assert_eq(&TensorData::from([[1.0, 1.0, 2.0], [3.0, 4.0, 4.0]]), false); + + // test int tensor + let data = TensorData::from([[0, 1, 2], [3, 4, 5]]); + let tensor = TestTensorInt::<2>::from_data(data, &device); + let output = tensor.clamp(1, 4); + + output + .into_data() + .assert_eq(&TensorData::from([[1, 1, 2], [3, 4, 4]]), false); +} + +#[test] +fn clamp_min_max_vec_should_compile() { + let input = TestTensor::<2>::ones([2, 4], &Default::default()); + let output = input.clamp(0., 0.5); + + output.into_data().assert_eq( + &TensorData::from([[0.5, 0.5, 0.5, 0.5], [0.5, 0.5, 0.5, 0.5]]), + false, + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/close.rs b/crates/burn-backend-tests/tests/tensor/float/ops/close.rs new file mode 100644 index 0000000..1f177ef --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/close.rs @@ -0,0 +1,40 @@ +use super::*; +use burn_tensor::{DEFAULT_ATOL, DEFAULT_RTOL, TensorData}; + +#[test] +fn test_is_close() { + let tensor1 = TestTensor::<2>::from([[0.0, 1.0, 0.0], [1.0, -1.0, 1.0]]); + let tensor2 = TestTensor::from([[0.0, 1.0, 0.0], [1.0, -1.0, 3.0]]) + 1e-9; + + let data_actual = tensor1 + .clone() + .is_close(tensor2.clone(), None, None) + .into_data(); + let defaults_expected = TensorData::from([[true, true, true], [true, true, false]]); + defaults_expected.assert_eq(&data_actual, false); + + // Using the defaults. + let data_actual = tensor1 + .is_close(tensor2, Some(DEFAULT_RTOL), Some(DEFAULT_ATOL)) + .into_data(); + defaults_expected.assert_eq(&data_actual, false); +} + +#[test] +fn test_all_close() { + let tensor1 = TestTensor::<2>::from([[0.0, 1.0, 0.0], [1.0, -1.0, 1.0]]); + let tensor2 = TestTensor::from([[0.0, 1.0, 0.0], [1.0, -1.0, 3.0]]) + 1e-9; + assert!(!tensor1.clone().all_close(tensor2.clone(), None, None)); + + let tensor2 = TestTensor::from([[0.0, 1.0, 0.0], [1.0, -1.0, 1.0]]) + 1e-9; + assert!(tensor1.all_close(tensor2, None, None)); + + // non finite values + let inf_plus = TestTensor::<2>::from([[f32::INFINITY]]); + let one = TestTensor::<2>::from([[1.]]); + let inf_minus = TestTensor::<2>::from([[-f32::INFINITY]]); + assert!(!inf_plus.clone().all_close(inf_minus.clone(), None, None)); + assert!(!one.clone().all_close(inf_minus.clone(), None, None)); + assert!(!one.all_close(inf_plus.clone(), None, None)); + assert!(inf_plus.clone().all_close(inf_plus, None, None)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/comparison.rs b/crates/burn-backend-tests/tests/tensor/float/ops/comparison.rs new file mode 100644 index 0000000..e022c2c --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/comparison.rs @@ -0,0 +1,393 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_equal_inf() { + let data_1 = TensorData::from([[0.0, 1.0, 2.0], [f32::INFINITY, 4.0, f32::NEG_INFINITY]]); + let data_2 = TensorData::from([[1.0, 1.0, 1.0], [f32::INFINITY, 3.0, f32::NEG_INFINITY]]); + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device); + + let data_actual_cloned = tensor_1.clone().equal(tensor_2.clone()); + let data_actual_inplace = tensor_1.equal(tensor_2); + + let data_expected = TensorData::from([[false, true, false], [true, false, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_not_equal_inf() { + let data_1 = TensorData::from([[0.0, 1.0, 2.0], [3.0, f32::INFINITY, 5.0]]); + let data_2 = TensorData::from([[1.0, 1.0, 1.0], [f32::INFINITY, 3.0, f32::NEG_INFINITY]]); + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device); + + let data_actual_cloned = tensor_1.clone().not_equal(tensor_2.clone()); + let data_actual_inplace = tensor_1.not_equal(tensor_2); + + let data_expected = TensorData::from([[true, false, true], [true, true, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_equal() { + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_2 = TestTensor::<2>::from([[1.0, 1.0, 1.0], [4.0, 3.0, 5.0]]); + + let data_actual_cloned = tensor_1.clone().equal(tensor_2.clone()); + let data_actual_inplace = tensor_1.equal(tensor_2); + + let data_expected = TensorData::from([[false, true, false], [false, false, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_not_equal() { + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_2 = TestTensor::<2>::from([[1.0, 1.0, 1.0], [4.0, 3.0, 5.0]]); + + let data_actual_cloned = tensor_1.clone().not_equal(tensor_2.clone()); + let data_actual_inplace = tensor_1.not_equal(tensor_2); + + let data_expected = TensorData::from([[true, false, true], [true, true, false]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_equal_elem() { + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 2.0, 5.0]]); + + let data_actual_cloned = tensor_1.clone().equal_elem(2); + let data_actual_inplace = tensor_1.equal_elem(2); + + let data_expected = TensorData::from([[false, false, true], [false, true, false]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_not_equal_elem() { + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 2.0, 5.0]]); + + let data_actual_cloned = tensor_1.clone().not_equal_elem(2); + let data_actual_inplace = tensor_1.not_equal_elem(2); + + let data_expected = TensorData::from([[true, true, false], [true, false, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn greater_elem() { + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let data_actual_cloned = tensor_1.clone().greater_elem(4); + let data_actual_inplace = tensor_1.greater_elem(4); + + let data_expected = TensorData::from([[false, false, false], [false, false, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_greater_equal_elem() { + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let data_actual_cloned = tensor_1.clone().greater_equal_elem(4.0); + let data_actual_inplace = tensor_1.greater_equal_elem(4.0); + + let data_expected = TensorData::from([[false, false, false], [false, true, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_greater() { + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_2 = TestTensor::<2>::from([[1.0, 1.0, 1.0], [4.0, 3.0, 50.0]]); + + let data_actual_cloned = tensor_1.clone().greater(tensor_2.clone()); + let data_actual_inplace = tensor_1.greater(tensor_2); + + let data_expected = TensorData::from([[false, false, true], [false, true, false]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_greater_equal() { + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_2 = TestTensor::<2>::from([[1.0, 1.0, 1.0], [4.0, 3.0, 50.0]]); + + let data_actual_cloned = tensor_1.clone().greater_equal(tensor_2.clone()); + let data_actual_inplace = tensor_1.greater_equal(tensor_2); + + let data_expected = TensorData::from([[false, true, true], [false, true, false]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_lower_elem() { + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let data_actual_cloned = tensor_1.clone().lower_elem(4.0); + let data_actual_inplace = tensor_1.lower_elem(4.0); + + let data_expected = TensorData::from([[true, true, true], [true, false, false]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_lower_equal_elem() { + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let data_actual_cloned = tensor_1.clone().lower_equal_elem(4.0); + let data_actual_inplace = tensor_1.lower_equal_elem(4.0); + + let data_expected = TensorData::from([[true, true, true], [true, true, false]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_lower() { + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_2 = TestTensor::<2>::from([[1.0, 1.0, 1.0], [4.0, 3.0, 50.0]]); + + let data_actual_cloned = tensor_1.clone().lower(tensor_2.clone()); + let data_actual_inplace = tensor_1.lower(tensor_2); + + let data_expected = TensorData::from([[true, false, false], [true, false, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_lower_equal() { + let tensor_1 = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_2 = TestTensor::<2>::from([[1.0, 1.0, 1.0], [4.0, 3.0, 50.0]]); + + let data_actual_cloned = tensor_1.clone().lower_equal(tensor_2.clone()); + let data_actual_inplace = tensor_1.lower_equal(tensor_2); + + let data_expected = TensorData::from([[true, true, false], [true, false, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_greater_broadcast() { + // Test broadcasting with shape [1, 4] vs [4, 4] + let device = Default::default(); + let data_1 = TensorData::from([[1.0, 2.0, 3.0, 4.0]]); + let data_2 = TensorData::from([ + [0.5, 1.5, 2.5, 3.5], + [1.5, 2.5, 3.5, 4.5], + [2.5, 3.5, 4.5, 5.5], + [3.5, 4.5, 5.5, 6.5], + ]); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device); + + let result = tensor_1.greater(tensor_2); + + let expected = TensorData::from([ + [true, true, true, true], + [false, false, false, false], + [false, false, false, false], + [false, false, false, false], + ]); + expected.assert_eq(&result.into_data(), false); +} + +#[test] +fn test_greater_equal_broadcast() { + // Test broadcasting with shape [4, 1] vs [1, 4] + let device = Default::default(); + let data_1 = TensorData::from([[1.0], [2.0], [3.0], [4.0]]); + let data_2 = TensorData::from([[1.0, 2.0, 3.0, 4.0]]); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device); + + let result = tensor_1.greater_equal(tensor_2); + + let expected = TensorData::from([ + [true, false, false, false], + [true, true, false, false], + [true, true, true, false], + [true, true, true, true], + ]); + expected.assert_eq(&result.into_data(), false); +} + +#[test] +fn test_lower_broadcast() { + // Test broadcasting mimicking CLIP pattern: [1, 5] vs [5, 1] + let device = Default::default(); + let data_1 = TensorData::from([[0.0, 1.0, -1.0, 2.0, -2.0]]); + let data_2 = TensorData::from([[0.5], [1.5], [-0.5], [-1.5], [2.5]]); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device); + + let result = tensor_1.lower(tensor_2); + + let expected = TensorData::from([ + [true, false, true, false, true], + [true, true, true, false, true], + [false, false, true, false, true], + [false, false, false, false, true], + [true, true, true, true, true], + ]); + expected.assert_eq(&result.into_data(), false); +} + +#[test] +fn test_lower_equal_broadcast() { + // Test broadcasting with shape [1, 1] vs [2, 4] + let device = Default::default(); + let data_1 = TensorData::from([[2.5]]); + let data_2 = TensorData::from([[1.0, 2.0, 3.0, 4.0], [2.0, 2.5, 3.0, 3.5]]); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device); + + let result = tensor_1.lower_equal(tensor_2); + + let expected = TensorData::from([[false, false, true, true], [false, true, true, true]]); + expected.assert_eq(&result.into_data(), false); +} + +#[test] +fn test_equal_broadcast() { + // Test broadcasting with different ranks + let device = Default::default(); + let data_1 = TensorData::from([[2.0], [3.0], [4.0]]); + let data_2 = TensorData::from([[2.0, 3.0, 4.0, 2.0]]); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device); + + let result = tensor_1.equal(tensor_2); + + let expected = TensorData::from([ + [true, false, false, true], + [false, true, false, false], + [false, false, true, false], + ]); + expected.assert_eq(&result.into_data(), false); +} + +#[test] +fn test_not_equal_broadcast() { + // Test broadcasting with shape [3, 1] vs [1, 3] + let device = Default::default(); + let data_1 = TensorData::from([[1.0], [2.0], [3.0]]); + let data_2 = TensorData::from([[1.0, 2.0, 3.0]]); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device); + + let result = tensor_1.not_equal(tensor_2); + + let expected = TensorData::from([ + [false, true, true], + [true, false, true], + [true, true, false], + ]); + expected.assert_eq(&result.into_data(), false); +} + +#[test] +fn test_greater_transposed() { + // [[1,2],[3,4]] transposed -> [[1,3],[2,4]]; > [[2,2],[2,2]] = [[F,T],[F,T]] + let lhs = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]).transpose(); + let rhs = TestTensor::<2>::from([[2.0, 2.0], [2.0, 2.0]]); + + let result = lhs.greater(rhs); + + result + .into_data() + .assert_eq(&TensorData::from([[false, true], [false, true]]), false); +} + +#[test] +fn test_equal_flipped_1d() { + // [1,2,3,4] flipped -> [4,3,2,1]; == [4,2,2,1] = [T,F,T,T] + let lhs = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0]).flip([0]); + let rhs = TestTensor::<1>::from([4.0, 2.0, 2.0, 1.0]); + + let result = lhs.equal(rhs); + + result + .into_data() + .assert_eq(&TensorData::from([true, false, true, true]), false); +} + +#[test] +fn test_lower_flipped_2d() { + // [[1,2],[3,4]] axis-0 flipped -> [[3,4],[1,2]]; < [[2,5],[2,1]] = [[F,T],[T,F]] + let lhs = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]).flip([0]); + let rhs = TestTensor::<2>::from([[2.0, 5.0], [2.0, 1.0]]); + + let result = lhs.lower(rhs); + + result + .into_data() + .assert_eq(&TensorData::from([[false, true], [true, false]]), false); +} + +#[test] +fn test_greater_elem_flipped() { + // [1,2,3,4] flipped -> [4,3,2,1]; > 2.5 = [T,T,F,F] + let lhs = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0]).flip([0]); + + let result = lhs.greater_elem(2.5); + + result + .into_data() + .assert_eq(&TensorData::from([true, true, false, false]), false); +} + +#[test] +fn test_equal_both_transposed() { + // Both transposed. [[1,2],[3,4]]^T == [[1,3],[2,4]]^T = ? + let lhs = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]).transpose(); + let rhs = TestTensor::<2>::from([[1.0, 3.0], [2.0, 4.0]]).transpose(); + + let result = lhs.equal(rhs); + + result + .into_data() + .assert_eq(&TensorData::from([[true, false], [false, true]]), false); +} + +#[test] +fn test_not_equal_narrowed() { + // [1,2,3,4,5,6] narrowed to [2,3,4,5]; != [2,2,4,4] = [F,T,F,T] + let lhs = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).narrow(0, 1, 4); + let rhs = TestTensor::<1>::from([2.0, 2.0, 4.0, 4.0]); + + let result = lhs.not_equal(rhs); + + result + .into_data() + .assert_eq(&TensorData::from([false, true, false, true]), false); +} + +#[test] +fn test_lower_flipped_both_axes() { + // [[1,2],[3,4]] flipped on both axes -> [[4,3],[2,1]]; < [[3,3],[3,3]] = [[F,F],[T,T]] + let lhs = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]).flip([0, 1]); + let rhs = TestTensor::<2>::from([[3.0, 3.0], [3.0, 3.0]]); + + let result = lhs.lower(rhs); + + result + .into_data() + .assert_eq(&TensorData::from([[false, false], [true, true]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/create_like.rs b/crates/burn-backend-tests/tests/tensor/float/ops/create_like.rs new file mode 100644 index 0000000..751302d --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/create_like.rs @@ -0,0 +1,57 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{Distribution, TensorData}; + +#[test] +fn should_support_zeros_like() { + let tensor = TestTensor::<3>::from_data( + [ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], + ], + &Default::default(), + ); + + let tensor = tensor.zeros_like(); + let expected = TensorData::from([[[0., 0., 0.], [0., 0., 0.]], [[0., 0., 0.], [0., 0., 0.]]]); + + tensor + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_ones_like() { + let tensor = TestTensor::<3>::from_data( + [ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], + ], + &Default::default(), + ); + + let tensor = tensor.ones_like(); + let expected = TensorData::from([[[1., 1., 1.], [1., 1., 1.]], [[1., 1., 1.], [1., 1., 1.]]]); + + tensor + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_randoms_like() { + let tensor = TestTensor::<3>::from_data( + [ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], + ], + &Default::default(), + ); + + let tensor = tensor.random_like(Distribution::Uniform(0.99999, 1.)); + let expected = TensorData::from([[[1., 1., 1.], [1., 1., 1.]], [[1., 1., 1.], [1., 1., 1.]]]); + + tensor + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/cross.rs b/crates/burn-backend-tests/tests/tensor/float/ops/cross.rs new file mode 100644 index 0000000..1a45315 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/cross.rs @@ -0,0 +1,148 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_cross_3d_last_dim() { + let tensor_1 = TestTensor::<2>::from([[1.0, 3.0, -5.0], [2.0, -1.0, 4.0]]); + let tensor_2 = TestTensor::from([[4.0, -2.0, 1.0], [3.0, 5.0, -2.0]]); + + let output = tensor_1.cross(tensor_2, -1); + + output.into_data().assert_eq( + &TensorData::from([[-7.0, -21.0, -14.0], [-18.0, 16.0, 13.0]]), + false, + ); +} + +#[test] +fn test_cross_3d_non_contiguous_last_dim() { + let tensor_1 = TestTensor::<2>::from([[1.0, 3.0, -5.0], [2.0, -1.0, 4.0]]); + let tensor_2 = TestTensor::from([[4.0, 3.0], [-2.0, 5.0], [1.0, -2.0]]); + + let output = tensor_1.cross(tensor_2.permute([1, 0]), -1); + + output.into_data().assert_eq( + &TensorData::from([[-7.0, -21.0, -14.0], [-18.0, 16.0, 13.0]]), + false, + ); +} + +#[test] +fn test_cross_3d_dim0() { + let tensor_1 = TestTensor::<2>::from([[1.0, 0.0], [0.0, 1.0], [0.0, 0.0]]); + let tensor_2 = TestTensor::from([[0.0, 1.0], [0.0, 0.0], [1.0, 0.0]]); + + let output = tensor_1.cross(tensor_2, 0); + + output.into_data().assert_eq( + &TensorData::from([[0.0, 0.0], [-1.0, 0.0], [0.0, -1.0]]), + false, + ); +} + +#[test] +fn test_cross_4d_middle_dim() { + // Shape [1, 3, 2]; cross along dim=1 should match the corresponding + // permuted last-dim cross. + let tensor_1 = TestTensor::<3>::from([[[1.0, 0.0], [0.0, 1.0], [0.0, 0.0]]]); + let tensor_2 = TestTensor::from([[[0.0, 1.0], [0.0, 0.0], [1.0, 0.0]]]); + + let output = tensor_1.cross(tensor_2, 1); + + output.into_data().assert_eq( + &TensorData::from([[[0.0, 0.0], [-1.0, 0.0], [0.0, -1.0]]]), + false, + ); +} + +#[test] +fn test_cross_non_last_dim_broadcast() { + // Broadcast on a non-last cross dim: lhs shape [3, 1] vs rhs shape [3, 4] + // crossed along dim=0 must equal the same op on the permuted [1, 3] vs + // [4, 3] along dim=1 (last) and then permuted back. + let lhs = TestTensor::<2>::from([[1.0], [2.0], [3.0]]); + let rhs = TestTensor::<2>::from([ + [4.0, 5.0, 6.0, 7.0], + [8.0, 9.0, 10.0, 11.0], + [12.0, 13.0, 14.0, 15.0], + ]); + + let non_last = lhs.clone().cross(rhs.clone(), 0); + let last_dim = lhs + .permute([1, 0]) + .cross(rhs.permute([1, 0]), 1) + .permute([1, 0]); + + non_last.into_data().assert_eq(&last_dim.into_data(), false); +} + +#[test] +fn test_cross_non_last_dim_matches_permuted_last_dim() { + // Cross on a non-last dim must equal: permute -> cross on last dim -> + // permute back. Here we cross [N, 3] along dim=1 (last) and along dim=0 + // after a transpose of the same data. + let a = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let b = TestTensor::<2>::from([[7.0, 8.0, 9.0], [1.0, 0.0, -1.0]]); + + let last_dim = a.clone().cross(b.clone(), 1); + let non_last = a.permute([1, 0]).cross(b.permute([1, 0]), 0); + + last_dim + .into_data() + .assert_eq(&non_last.permute([1, 0]).into_data(), false); +} + +#[test] +fn test_cross_3d_broadcast() { + let tensor_1 = TestTensor::<2>::from([[1.0, 3.0, -5.0]]); + let tensor_2 = TestTensor::from([[4.0, -2.0, 1.0], [3.0, 5.0, -2.0]]); + + let output = tensor_1.cross(tensor_2, -1); + + output.into_data().assert_eq( + &TensorData::from([[-7.0, -21.0, -14.0], [19.0, -13.0, -4.0]]), + false, + ); +} + +#[test] +fn test_cross_4d_last_dim() { + let tensor_1 = TestTensor::<3>::from([[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]]); + let tensor_2 = TestTensor::from([[[0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]]); + + let output = tensor_1.cross(tensor_2, -1); + + output.into_data().assert_eq( + &TensorData::from([[[0.0, 0.0, 1.0], [1.0, 0.0, 0.0]]]), + false, + ); +} + +// Helper to compute expected cross product for 2-D (N × 3) tensors. +fn manual_cross(a: &[[f32; 3]], b: &[[f32; 3]]) -> Vec<[f32; 3]> { + a.iter() + .zip(b.iter()) + .map(|(x, y)| { + [ + x[1] * y[2] - x[2] * y[1], + x[2] * y[0] - x[0] * y[2], + x[0] * y[1] - x[1] * y[0], + ] + }) + .collect() +} + +#[test] +fn forward_matches_manual_cross() { + let a_raw = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]; + let b_raw = [[7.0, 8.0, 9.0], [1.0, 0.0, -1.0]]; + let a = TestTensor::<2>::from(a_raw); + let b = TestTensor::<2>::from(b_raw); + + let out = a.cross(b.clone(), 1); + let expected_vec = manual_cross(&a_raw, &b_raw); + let expected: [[f32; 3]; 2] = [expected_vec[0], expected_vec[1]]; + + out.into_data() + .assert_eq(&TensorData::from(expected), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/cumulative.rs b/crates/burn-backend-tests/tests/tensor/float/ops/cumulative.rs new file mode 100644 index 0000000..cba4844 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/cumulative.rs @@ -0,0 +1,202 @@ +use super::*; +use burn_tensor::TensorData; + +#[cfg(feature = "flex")] +use burn_tensor::ElementConversion; + +#[test] +fn test_cumsum_float_dim_0() { + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + let output = tensor.cumsum(0); + + output + .into_data() + .assert_eq(&TensorData::from([[1.0, 2.0, 3.0], [5.0, 7.0, 9.0]]), false); +} + +#[test] +fn test_cumsum_float_dim_1() { + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + let output = tensor.cumsum(1); + + output.into_data().assert_eq( + &TensorData::from([[1.0, 3.0, 6.0], [4.0, 9.0, 15.0]]), + false, + ); +} + +#[test] +fn test_cumsum_non_contiguous() { + let tensor = TestTensor::<2>::from([[1., 2.], [3., 4.]]).swap_dims(0, 1); + + let output = tensor.cumsum(1); + + output + .into_data() + .assert_eq(&TensorData::from([[1., 4.], [2., 6.]]), false); +} + +#[test] +fn test_cumsum_float_3d() { + let tensor = TestTensor::<3>::from([[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]); + + let output = tensor.cumsum(2); + + output.into_data().assert_eq( + &TensorData::from([[[1.0, 3.0], [3.0, 7.0]], [[5.0, 11.0], [7.0, 15.0]]]), + false, + ); +} + +#[test] +fn test_cumprod_float_dim_0() { + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + let output = tensor.cumprod(0); + + output.into_data().assert_eq( + &TensorData::from([[1.0, 2.0, 3.0], [4.0, 10.0, 18.0]]), + false, + ); +} + +#[test] +fn test_cumprod_float_dim_1() { + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + let output = tensor.cumprod(1); + + output.into_data().assert_eq( + &TensorData::from([[1.0, 2.0, 6.0], [4.0, 20.0, 120.0]]), + false, + ); +} + +#[test] +fn test_cumprod_float_3d() { + let tensor = TestTensor::<3>::from([[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]); + + let output = tensor.cumprod(2); + + output.into_data().assert_eq( + &TensorData::from([[[1.0, 2.0], [3.0, 12.0]], [[5.0, 30.0], [7.0, 56.0]]]), + false, + ); +} + +#[test] +fn test_cummin_float_dim_0() { + let tensor = TestTensor::<2>::from([[3.0, 1.0, 4.0], [2.0, 5.0, 1.0]]); + + let output = tensor.cummin(0); + + output + .into_data() + .assert_eq(&TensorData::from([[3.0, 1.0, 4.0], [2.0, 1.0, 1.0]]), false); +} + +#[test] +fn test_cummin_float_dim_1() { + let tensor = TestTensor::<2>::from([[3.0, 1.0, 4.0], [2.0, 5.0, 1.0]]); + + let output = tensor.cummin(1); + + output + .into_data() + .assert_eq(&TensorData::from([[3.0, 1.0, 1.0], [2.0, 2.0, 1.0]]), false); +} + +#[test] +fn test_cummin_float_3d() { + let tensor = TestTensor::<3>::from([[[4.0, 2.0], [3.0, 1.0]], [[5.0, 6.0], [7.0, 8.0]]]); + + let output = tensor.cummin(2); + + output.into_data().assert_eq( + &TensorData::from([[[4.0, 2.0], [3.0, 1.0]], [[5.0, 5.0], [7.0, 7.0]]]), + false, + ); +} + +#[test] +fn test_cummax_float_dim_0() { + let tensor = TestTensor::<2>::from([[3.0, 1.0, 4.0], [1.0, 5.0, 2.0]]); + + let output = tensor.cummax(0); + + output + .into_data() + .assert_eq(&TensorData::from([[3.0, 1.0, 4.0], [3.0, 5.0, 4.0]]), false); +} + +#[test] +fn test_cummax_float_dim_1() { + let tensor = TestTensor::<2>::from([[3.0, 1.0, 4.0], [1.0, 5.0, 2.0]]); + + let output = tensor.cummax(1); + + output + .into_data() + .assert_eq(&TensorData::from([[3.0, 3.0, 4.0], [1.0, 5.0, 5.0]]), false); +} + +#[test] +fn test_cummax_float_3d() { + let tensor = TestTensor::<3>::from([[[1.0, 3.0], [2.0, 4.0]], [[5.0, 2.0], [6.0, 1.0]]]); + + let output = tensor.cummax(2); + + output.into_data().assert_eq( + &TensorData::from([[[1.0, 3.0], [2.0, 4.0]], [[5.0, 5.0], [6.0, 6.0]]]), + false, + ); +} + +// NaN-propagation tests below. Only run under the `flex` backend +// feature; other burn backends follow IEEE 754 min/max and drop NaN. +// Positive-gate form because the default CI build doesn't set +// identifying feature flags on burn-backend-tests. See issue #4814. +#[cfg(feature = "flex")] +#[test] +fn test_cummin_nan_propagation() { + // Once NaN appears, cummin propagates it forward. + let tensor = TestTensor::<1>::from([3.0, f32::NAN, 1.0, 2.0]); + + let output = tensor.cummin(0); + + let data: Vec = output.into_data().to_vec().unwrap(); + assert_eq!(data[0], 3.0.elem::()); + assert!(data[1].is_nan()); + assert!(data[2].is_nan()); + assert!(data[3].is_nan()); +} + +#[cfg(feature = "flex")] +#[test] +fn test_cummax_nan_propagation() { + let tensor = TestTensor::<1>::from([1.0, f32::NAN, 5.0, 2.0]); + + let output = tensor.cummax(0); + + let data: Vec = output.into_data().to_vec().unwrap(); + assert_eq!(data[0], 1.0.elem::()); + assert!(data[1].is_nan()); + assert!(data[2].is_nan()); + assert!(data[3].is_nan()); +} + +#[cfg(feature = "flex")] +#[test] +fn test_cummin_nan_at_start() { + // NaN on the first element should poison the entire output. + let tensor = TestTensor::<1>::from([f32::NAN, 1.0, 2.0]); + + let output = tensor.cummin(0); + + let data: Vec = output.into_data().to_vec().unwrap(); + assert!(data[0].is_nan()); + assert!(data[1].is_nan()); + assert!(data[2].is_nan()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/div.rs b/crates/burn-backend-tests/tests/tensor/float/ops/div.rs new file mode 100644 index 0000000..c19e6aa --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/div.rs @@ -0,0 +1,49 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_div_ops() { + let data_1 = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let data_2 = TensorData::from([[1.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device); + + let output = tensor_1 / tensor_2; + let expected = TensorData::from([[0.0, 1.0, 1.0], [1.0, 1.0, 1.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_div_broadcast() { + let data_1 = TensorData::from([[0.0, 1.0, 2.0]]); + let data_2 = TensorData::from([[1.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device); + + let output = tensor_1 / tensor_2; + + output.into_data().assert_eq( + &TensorData::from([[0.0, 1.0, 1.0], [0.0, 0.25, 0.4]]), + false, + ); +} + +#[test] +fn should_support_div_scalar_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let scalar = 2.0; + let device = Default::default(); + let tensor = TestTensor::<2>::from_data(data, &device); + + let output = tensor / scalar; + + output + .into_data() + .assert_eq(&TensorData::from([[0.0, 0.5, 1.0], [1.5, 2.0, 2.5]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/dot.rs b/crates/burn-backend-tests/tests/tensor/float/ops/dot.rs new file mode 100644 index 0000000..33c03e5 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/dot.rs @@ -0,0 +1,35 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_float() { + let device = Default::default(); + let tensor_1 = TestTensor::<1>::from_data([1.0, 2.0, 3.0], &device); + let tensor_2 = TestTensor::<1>::from_data([0.0, -1.0, 4.0], &device); + + let output = tensor_1.dot(tensor_2); + let expected = TensorData::from([10.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_int() { + let device = Default::default(); + let tensor_1 = TestTensor::<1>::from_data([1, 2, 3], &device); + let tensor_2 = TestTensor::<1>::from_data([0, -1, 4], &device); + + let output = tensor_1.dot(tensor_2); + let expected = TensorData::from([10]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +#[should_panic] +fn test_panics_for_different_sizes() { + let device = Default::default(); + let tensor_1 = TestTensor::<1>::from_data([1, 2], &device); + let tensor_2 = TestTensor::<1>::from_data([1, 2, 3], &device); + let _output = tensor_1.dot(tensor_2); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/erf.rs b/crates/burn-backend-tests/tests/tensor/float/ops/erf.rs new file mode 100644 index 0000000..76ab84e --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/erf.rs @@ -0,0 +1,34 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_erf_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.erf(); + let expected = TensorData::from([[0.0000, 0.8427, 0.99532], [0.99998, 1.0000, 1.0000]]); + + output.into_data().assert_approx_eq::( + &expected, + Tolerance::default().set_half_precision_absolute(2e-3), + ); +} + +#[test] +fn should_support_erf_ops_with_negative_number() { + let data = TensorData::from([[-0.056, -0.043, -0.089], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.erf(); + let expected = TensorData::from([ + [-0.06312324, -0.048490416, -0.10016122], + [0.99998, 1.0000, 1.0000], + ]); + + output.into_data().assert_approx_eq::( + &expected, + Tolerance::default().set_half_precision_absolute(3e-3), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/exp.rs b/crates/burn-backend-tests/tests/tensor/float/ops/exp.rs new file mode 100644 index 0000000..f881a99 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/exp.rs @@ -0,0 +1,31 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; +use core::f32::consts::E; + +#[test] +fn should_support_exp_transposed() { + // [[0, 1], [2, 3]] transposed -> [[0, 2], [1, 3]] + let tensor = TestTensor::<2>::from([[0.0, 1.0], [2.0, 3.0]]); + let transposed = tensor.transpose(); + + let output = transposed.exp(); + let expected = TensorData::from([[1.0, E * E], [E, E * E * E]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_exp_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.exp(); + let expected = TensorData::from([[1.0, 2.71830, 7.3891], [20.0855, 54.5981, 148.4132]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/expand.rs b/crates/burn-backend-tests/tests/tensor/float/ops/expand.rs new file mode 100644 index 0000000..6bc2f5c --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/expand.rs @@ -0,0 +1,138 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn expand_2d() { + let tensor = TestTensor::<1>::from_data([1.0, 2.0, 3.0], &Default::default()); + let output = tensor.expand([3, 3]); + + output.into_data().assert_eq( + &TensorData::from([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0], [1.0, 2.0, 3.0]]), + false, + ); + + let tensor = TestTensor::<1>::from_data([4.0, 7.0, 2.0, 3.0], &Default::default()); + let output = tensor.expand([2, 4]); + + output.into_data().assert_eq( + &TensorData::from([[4.0, 7.0, 2.0, 3.0], [4.0, 7.0, 2.0, 3.0]]), + false, + ); +} + +#[test] +fn expand_3d() { + let tensor = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &Default::default()); + let output = tensor.expand([3, 2, 2]); + let expected = TensorData::from([ + [[1.0, 2.0], [3.0, 4.0]], + [[1.0, 2.0], [3.0, 4.0]], + [[1.0, 2.0], [3.0, 4.0]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn expand_higher_dimensions() { + let tensor = TestTensor::<2>::from_data([[1.0, 2.0, 3.0, 4.0]], &Default::default()); + let output = tensor.expand([2, 3, 4]); + let expected = TensorData::from([ + [ + [1.0, 2.0, 3.0, 4.0], + [1.0, 2.0, 3.0, 4.0], + [1.0, 2.0, 3.0, 4.0], + ], + [ + [1.0, 2.0, 3.0, 4.0], + [1.0, 2.0, 3.0, 4.0], + [1.0, 2.0, 3.0, 4.0], + ], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn expand_sum_3d() { + let tensor = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &Default::default()); + let output = tensor.expand([3, 2, 2]).sum_dim(0); + let expected = TensorData::from([[[3.0, 6.0], [9.0, 12.0]]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn broadcast_single() { + let tensor = TestTensor::<1>::from_data([1.0], &Default::default()); + let output = tensor.expand([2, 3]); + + output + .into_data() + .assert_eq(&TensorData::from([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]), false); +} + +#[test] +#[should_panic] +fn should_fail_expand_incompatible_shapes() { + let tensor = TestTensor::<1>::from_data([1.0, 2.0, 3.0], &Default::default()); + let _expanded_tensor = tensor.expand([2, 2]); +} + +#[test] +fn expand_after_transpose() { + // [[1, 2], [3, 4]] transposed -> [[1, 3], [2, 4]] then expanded to [3, 2, 2]. + let tensor = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]).transpose(); + + let output = tensor.expand([3, 2, 2]); + + output.into_data().assert_eq( + &TensorData::from([ + [[1.0, 3.0], [2.0, 4.0]], + [[1.0, 3.0], [2.0, 4.0]], + [[1.0, 3.0], [2.0, 4.0]], + ]), + false, + ); +} + +#[test] +fn expand_after_flip_1d() { + // [1,2,3] flipped -> [3,2,1] then expanded to [2, 3]. + let tensor = TestTensor::<1>::from([1.0, 2.0, 3.0]).flip([0]); + + let output = tensor.expand([2, 3]); + + output + .into_data() + .assert_eq(&TensorData::from([[3.0, 2.0, 1.0], [3.0, 2.0, 1.0]]), false); +} + +#[test] +fn expand_after_flip_2d() { + // [[1,2],[3,4]] axis-0 flipped -> [[3,4],[1,2]] then expanded to [3,2,2]. + let tensor = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]).flip([0]); + + let output = tensor.expand([3, 2, 2]); + + output.into_data().assert_eq( + &TensorData::from([ + [[3.0, 4.0], [1.0, 2.0]], + [[3.0, 4.0], [1.0, 2.0]], + [[3.0, 4.0], [1.0, 2.0]], + ]), + false, + ); +} + +#[test] +fn expand_after_narrow() { + // [0,1,2,3,4] narrowed to [1,2,3] then expanded to [2, 3]. + let tensor = TestTensor::<1>::from([0.0, 1.0, 2.0, 3.0, 4.0]).narrow(0, 1, 3); + + let output = tensor.expand([2, 3]); + + output + .into_data() + .assert_eq(&TensorData::from([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/finite.rs b/crates/burn-backend-tests/tests/tensor/float/ops/finite.rs new file mode 100644 index 0000000..eb21c9a --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/finite.rs @@ -0,0 +1,23 @@ +use super::*; + +#[test] +fn is_finite() { + let all_finite = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let all_finite_expected = TestTensorBool::<2>::from([[true, true, true], [true, true, true]]); + + let with_inf_nan = TestTensor::<2>::from([ + [0.0, f32::INFINITY, f32::NAN], + [f32::NEG_INFINITY, f32::NAN, 5.0], + ]); + let with_inf_nan_expected = + TestTensorBool::<2>::from([[true, false, false], [false, false, true]]); + + all_finite_expected + .into_data() + .assert_eq(&all_finite.is_finite().into_data(), false); + + with_inf_nan + .is_finite() + .into_data() + .assert_eq(&with_inf_nan_expected.into_data(), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/flatten.rs b/crates/burn-backend-tests/tests/tensor/float/ops/flatten.rs new file mode 100644 index 0000000..f34a96c --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/flatten.rs @@ -0,0 +1,64 @@ +use super::*; +use burn_tensor::Shape; + +/// Test if the function can successfully flatten a 4D tensor to a 1D tensor. +#[test] +fn should_flatten_to_1d() { + let tensor = TestTensor::<4>::ones(Shape::new([2, 3, 4, 5]), &Default::default()); + let flattened_tensor: TestTensor<1> = tensor.flatten(0, 3); + let expected_shape = Shape::new([120]); + assert_eq!(flattened_tensor.shape(), expected_shape); +} + +/// Test if the function can successfully flatten the middle dimensions of a 4D tensor. +#[test] +fn should_flatten_middle() { + let tensor = TestTensor::<4>::ones(Shape::new([2, 3, 4, 5]), &Default::default()); + let flattened_tensor: TestTensor<3> = tensor.flatten(1, 2); + let expected_shape = Shape::new([2, 12, 5]); + assert_eq!(flattened_tensor.shape(), expected_shape); +} + +/// Test if the function can successfully flatten the first dimensions of a 4D tensor. +#[test] +fn should_flatten_begin() { + let tensor = TestTensor::<4>::ones(Shape::new([2, 3, 4, 5]), &Default::default()); + let flattened_tensor: TestTensor<2> = tensor.flatten(0, 2); + let expected_shape = Shape::new([24, 5]); + assert_eq!(flattened_tensor.shape(), expected_shape); +} + +/// Test if the function can successfully flatten the last dimensions of a 4D tensor. +#[test] +fn should_flatten_end() { + let tensor = TestTensor::<4>::ones(Shape::new([2, 3, 4, 5]), &Default::default()); + let flattened_tensor: TestTensor<2> = tensor.flatten(1, 3); + let expected_shape = Shape::new([2, 60]); + assert_eq!(flattened_tensor.shape(), expected_shape); +} + +/// Test if the function can flatten negative indices. +#[test] +fn should_flatten_end_negative_indices() { + let tensor = TestTensor::<4>::ones(Shape::new([2, 3, 4, 5]), &Default::default()); + let flattened_tensor: TestTensor<2> = tensor.flatten(-3, -1); + let expected_shape = Shape::new([2, 60]); + assert_eq!(flattened_tensor.shape(), expected_shape); +} + +/// Test if the function panics when the start dimension is greater than the end dimension. +#[test] +#[should_panic] +fn should_flatten_panic() { + let tensor = TestTensor::<4>::ones(Shape::new([2, 3, 4, 5]), &Default::default()); + let _flattened_tensor: TestTensor<2> = tensor.flatten(2, 0); +} + +#[test] +#[should_panic] +fn not_enough_destination_dimension() { + let tensor = TestTensor::<3>::ones(Shape::new([1, 5, 15]), &Default::default()); + let flattened_tensor: TestTensor<1> = tensor.flatten(1, 2); + let expected_shape = Shape::new([75]); + assert_eq!(flattened_tensor.shape(), expected_shape); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/flip.rs b/crates/burn-backend-tests/tests/tensor/float/ops/flip.rs new file mode 100644 index 0000000..4448ea6 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/flip.rs @@ -0,0 +1,48 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn flip_float() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device) + .reshape([2, 3, 4]) + .float(); + + let flipped = tensor.clone().flip([0, 2]); + // from pytorch: + // import torch; torch.arange(0, 24).reshape(2, 3, 4).flip((0, 2)).float() + let expected = TensorData::from([ + [ + [15., 14., 13., 12.], + [19., 18., 17., 16.], + [23., 22., 21., 20.], + ], + [[3., 2., 1., 0.], [7., 6., 5., 4.], [11., 10., 9., 8.]], + ]); + + flipped.into_data().assert_eq(&expected, false); + + // Test with no flip + let flipped = tensor.clone().flip([]); + tensor.into_data().assert_eq(&flipped.into_data(), false); +} + +#[test] +#[should_panic] +fn flip_duplicated_axes() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device).reshape([2, 3, 4]); + + // Test with a duplicated axis + let _ = tensor.clone().flip([0, 0, 1]); +} + +#[test] +#[should_panic] +fn flip_out_of_bound_axis() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device).reshape([2, 3, 4]); + + // Test with an out of bound axis + let _ = tensor.clone().flip([3, 0, 1]); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/floor.rs b/crates/burn-backend-tests/tests/tensor/float/ops/floor.rs new file mode 100644 index 0000000..1c29e9d --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/floor.rs @@ -0,0 +1,16 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_floor_ops() { + let data = TensorData::from([[24.0423, 87.9478, 76.1838], [59.6929, 43.8169, 94.8826]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.floor(); + let expected = TensorData::from([[24., 87., 76.], [59., 43., 94.]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/fmod.rs b/crates/burn-backend-tests/tests/tensor/float/ops/fmod.rs new file mode 100644 index 0000000..5962f31 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/fmod.rs @@ -0,0 +1,295 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{ElementConversion, TensorData}; + +#[allow(unused_imports)] // f16 +use num_traits::Float; + +#[test] +fn should_support_fmod_ops() { + let dividend = TensorData::from([[5.3, -5.3], [7.5, -7.5]]); + let divisor = TensorData::from([[2.0, 2.0], [3.0, 3.0]]); + + let dividend_tensor = TestTensor::<2>::from_data(dividend, &Default::default()); + let divisor_tensor = TestTensor::<2>::from_data(divisor, &Default::default()); + + let output = dividend_tensor.fmod(divisor_tensor); + let expected = TensorData::from([[1.3, -1.3], [1.5, -1.5]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_fmod_scalar() { + let data = TensorData::from([5.3, -5.3, 7.5, -7.5]); + let tensor = TestTensor::<1>::from_data(data, &Default::default()); + + let output = tensor.fmod_scalar(2.0); + let expected = TensorData::from([1.3, -1.3, 1.5, -1.5]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_handle_positive_dividend_positive_divisor() { + let dividend = TensorData::from([10.0, 7.5, 3.8, 1.2]); + let divisor = TensorData::from([3.0, 2.0, 1.5, 0.7]); + + let dividend_tensor = TestTensor::<1>::from_data(dividend, &Default::default()); + let divisor_tensor = TestTensor::<1>::from_data(divisor, &Default::default()); + + let output = dividend_tensor.fmod(divisor_tensor); + let expected = TensorData::from([1.0, 1.5, 0.8, 0.5]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_handle_negative_dividend() { + let dividend = TensorData::from([-10.0, -7.5, -3.8, -1.2]); + let divisor = TensorData::from([3.0, 2.0, 1.5, 0.7]); + + let dividend_tensor = TestTensor::<1>::from_data(dividend, &Default::default()); + let divisor_tensor = TestTensor::<1>::from_data(divisor, &Default::default()); + + let output = dividend_tensor.fmod(divisor_tensor); + let expected = TensorData::from([-1.0, -1.5, -0.8, -0.5]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_handle_mixed_signs() { + let dividend = TensorData::from([5.3, -5.3, 5.3, -5.3]); + let divisor = TensorData::from([2.0, 2.0, -2.0, -2.0]); + + let dividend_tensor = TestTensor::<1>::from_data(dividend, &Default::default()); + let divisor_tensor = TestTensor::<1>::from_data(divisor, &Default::default()); + + let output = dividend_tensor.fmod(divisor_tensor); + // fmod result has same sign as dividend + let expected = TensorData::from([1.3, -1.3, 1.3, -1.3]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_handle_infinity_dividend() { + // If x is ±∞ and y is not NaN, NaN is returned + let dividend = TensorData::from([ + f32::INFINITY, + f32::NEG_INFINITY, + f32::INFINITY, + f32::NEG_INFINITY, + ]); + let divisor = TensorData::from([2.0, 3.0, -2.0, -3.0]); + + let dividend_tensor = TestTensor::<1>::from_data(dividend, &Default::default()); + let divisor_tensor = TestTensor::<1>::from_data(divisor, &Default::default()); + + let output = dividend_tensor.fmod(divisor_tensor); + let data = output.into_data(); + let values = data.as_slice::().unwrap(); + + // All results should be NaN + assert!(values[0].is_nan(), "fmod(inf, 2.0) should be NaN"); + assert!(values[1].is_nan(), "fmod(-inf, 3.0) should be NaN"); + assert!(values[2].is_nan(), "fmod(inf, -2.0) should be NaN"); + assert!(values[3].is_nan(), "fmod(-inf, -3.0) should be NaN"); +} + +#[test] +fn should_handle_zero_divisor() { + // If y is ±0 and x is not NaN, NaN should be returned + let dividend = TensorData::from([5.3, -5.3, 0.0, 1.0]); + let divisor = TensorData::from([0.0, -0.0, 0.0, -0.0]); + + let dividend_tensor = TestTensor::<1>::from_data(dividend, &Default::default()); + let divisor_tensor = TestTensor::<1>::from_data(divisor, &Default::default()); + + let output = dividend_tensor.fmod(divisor_tensor); + let data = output.into_data(); + let values = data.as_slice::().unwrap(); + + // All results should be NaN + assert!(values[0].is_nan(), "fmod(5.3, 0.0) should be NaN"); + assert!(values[1].is_nan(), "fmod(-5.3, -0.0) should be NaN"); + assert!(values[2].is_nan(), "fmod(0.0, 0.0) should be NaN"); + assert!(values[3].is_nan(), "fmod(1.0, -0.0) should be NaN"); +} + +#[test] +fn should_handle_infinity_divisor() { + // If y is ±∞ and x is finite, x is returned + let dividend = TensorData::from([5.3, -5.3, 0.0, -0.0]); + let divisor = TensorData::from([ + f32::INFINITY, + f32::NEG_INFINITY, + f32::INFINITY, + f32::NEG_INFINITY, + ]); + + let dividend_tensor = TestTensor::<1>::from_data(dividend, &Default::default()); + let divisor_tensor = TestTensor::<1>::from_data(divisor, &Default::default()); + + let output = dividend_tensor.fmod(divisor_tensor); + let expected = TensorData::from([5.3, -5.3, 0.0, -0.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_handle_nan_arguments() { + // If either argument is NaN, NaN is returned + let dividend = TensorData::from([f32::NAN, 5.3, f32::NAN, 0.0]); + let divisor = TensorData::from([2.0, f32::NAN, f32::NAN, 3.0]); + + let dividend_tensor = TestTensor::<1>::from_data(dividend, &Default::default()); + let divisor_tensor = TestTensor::<1>::from_data(divisor, &Default::default()); + + let output = dividend_tensor.fmod(divisor_tensor); + let data = output.into_data(); + let values = data.as_slice::().unwrap(); + + assert!(values[0].is_nan(), "fmod(NaN, 2.0) should be NaN"); + assert!(values[1].is_nan(), "fmod(5.3, NaN) should be NaN"); + assert!(values[2].is_nan(), "fmod(NaN, NaN) should be NaN"); + assert!(!values[3].is_nan(), "fmod(0.0, 3.0) should be 0.0"); +} + +#[test] +fn should_handle_negative_zero() { + // If x is -0 and y is greater than zero, either +0 or -0 may be returned + let dividend = TensorData::from([-0.0_f32]); + let divisor = TensorData::from([2.0_f32]); + + let dividend_tensor = TestTensor::<1>::from_data(dividend, &Default::default()); + let divisor_tensor = TestTensor::<1>::from_data(divisor, &Default::default()); + + let output = dividend_tensor.fmod(divisor_tensor); + let data = output.into_data(); + let values = data.as_slice::().unwrap(); + + // Result should be zero (either +0 or -0 is acceptable) + assert_eq!( + values[0], + 0.0f32.elem::(), + "fmod(-0, 2.0) should be zero" + ); +} + +#[test] +fn should_support_fmod_broadcasting_2d() { + // Test broadcasting: 1x2 with 3x2 + let dividend = TensorData::from([[5.3, -5.3]]); // Shape: 1x2 + let divisor = TensorData::from([[2.0, 2.0], [3.0, 3.0], [1.5, 1.5]]); // Shape: 3x2 + + let dividend_tensor = TestTensor::<2>::from_data(dividend, &Default::default()); + let divisor_tensor = TestTensor::<2>::from_data(divisor, &Default::default()); + + let output = dividend_tensor.fmod(divisor_tensor); + let expected = TensorData::from([ + [1.3, -1.3], // 5.3 % 2.0, -5.3 % 2.0 + [2.3, -2.3], // 5.3 % 3.0, -5.3 % 3.0 + [0.8, -0.8], // 5.3 % 1.5, -5.3 % 1.5 + ]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_fmod_broadcasting_3d() { + // Test broadcasting: 1x1x3 with 2x1x3 + let dividend = TensorData::from([[[5.0, -7.0, 8.0]]]); // Shape: 1x1x3 + let divisor = TensorData::from([[[3.0, 3.0, 3.0]], [[4.0, 4.0, 4.0]]]); // Shape: 2x1x3 + + let dividend_tensor = TestTensor::<3>::from_data(dividend, &Default::default()); + let divisor_tensor = TestTensor::<3>::from_data(divisor, &Default::default()); + + let output = dividend_tensor.fmod(divisor_tensor); + let expected = TensorData::from([ + [[2.0, -1.0, 2.0]], // 5.0 % 3.0, -7.0 % 3.0, 8.0 % 3.0 + [[1.0, -3.0, 0.0]], // 5.0 % 4.0, -7.0 % 4.0, 8.0 % 4.0 + ]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_fmod_scalar_broadcasting() { + // Test scalar operation with different shapes + let data = TensorData::from([[5.3, -5.3, 7.5], [-7.5, 10.0, -10.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.fmod_scalar(3.0); + let expected = TensorData::from([[2.3, -2.3, 1.5], [-1.5, 1.0, -1.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_handle_edge_case_values() { + // Test various edge cases + let dividend = TensorData::from([0.0, -0.0, 1e-10, -1e-10, 10.0, -10.0]); + let divisor = TensorData::from([1.0, 1.0, 1.0, 1.0, 3.0, 3.0]); + + let dividend_tensor = TestTensor::<1>::from_data(dividend, &Default::default()); + let divisor_tensor = TestTensor::<1>::from_data(divisor, &Default::default()); + + let output = dividend_tensor.fmod(divisor_tensor); + let expected = TensorData::from([0.0, -0.0, 1e-10, -1e-10, 1.0, -1.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_handle_special_scalar_cases() { + // Test scalar operations with special values + let data = TensorData::from([5.3, -5.3, 0.0, -0.0]); + let tensor = TestTensor::<1>::from_data(data, &Default::default()); + + // Test with infinity divisor + let output_inf = tensor.clone().fmod_scalar(f32::INFINITY); + let expected_inf = TensorData::from([5.3, -5.3, 0.0, -0.0]); + output_inf + .into_data() + .assert_approx_eq::(&expected_inf, Tolerance::default()); + + // Test with very small divisor + // Doesn't work if the test divisor is subnormal + if FloatElem::MIN_POSITIVE > 1e-5f32.elem::() { + return; + } + + let output_small = tensor.clone().fmod_scalar(1e-5); + let data = output_small.into_data(); + let values = data.as_slice::().unwrap(); + + // let expected = TensorData::from([0.0, 0.0, 0.0, 0.0]); + + // Results should be very small remainders + assert!(values[0].abs() < 1e-5f32.elem::()); + assert!(values[1].abs() < 1e-5f32.elem::()); + assert_eq!(values[2], 0.0f32.elem::()); + assert_eq!(values[3], 0.0f32.elem::()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/full.rs b/crates/burn-backend-tests/tests/tensor/float/ops/full.rs new file mode 100644 index 0000000..87d824c --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/full.rs @@ -0,0 +1,28 @@ +use super::*; +use burn_tensor::{DType, TensorData}; + +#[test] +fn test_data_full() { + let tensor = TensorData::full([2, 3], 2.0); + + tensor.assert_eq(&TensorData::from([[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]]), false); +} + +#[test] +fn test_tensor_full() { + let device = Default::default(); + let tensor = TestTensor::<2>::full([2, 3], 2.1, &device); + tensor + .into_data() + .assert_eq(&TensorData::from([[2.1, 2.1, 2.1], [2.1, 2.1, 2.1]]), false); +} + +#[test] +fn test_tensor_full_options() { + let tensor = TestTensor::<2>::full([2, 3], 2.1, (&Default::default(), DType::F32)); + assert_eq!(tensor.dtype(), DType::F32); + + tensor + .into_data() + .assert_eq(&TensorData::from([[2.1, 2.1, 2.1], [2.1, 2.1, 2.1]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/gather_scatter.rs b/crates/burn-backend-tests/tests/tensor/float/ops/gather_scatter.rs new file mode 100644 index 0000000..248ffd0 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/gather_scatter.rs @@ -0,0 +1,195 @@ +use super::*; +use burn_tensor::{IndexingUpdateOp, TensorData}; + +#[test] +fn should_gather_1d_dim0() { + let device = Default::default(); + let tensor = TestTensor::<1>::from_data([0.0, 1.0, 2.0], &device); + let indices = TestTensorInt::from_ints([1, 1, 0, 1, 2], &device); + + let output = tensor.gather(0, indices); + + output + .into_data() + .assert_eq(&TensorData::from([1.0, 1.0, 0.0, 1.0, 2.0]), false); +} + +#[test] +fn should_gather_2d_dim0() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &device); + let indices = TestTensorInt::from_ints([[0, 1, 0], [1, 0, 1]], &device); + + let output = tensor.gather(0, indices); + + output + .into_data() + .assert_eq(&TensorData::from([[0.0, 4.0, 2.0], [3.0, 1.0, 5.0]]), false); +} + +#[test] +fn should_gather_2d_dim1() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &device); + let indices = TestTensorInt::from_ints([[2, 1, 0, 0], [2, 0, 1, 2]], &device); + + let output = tensor.gather(1, indices); + + output.into_data().assert_eq( + &TensorData::from([[2.0, 1.0, 0.0, 0.0], [5.0, 3.0, 4.0, 5.0]]), + false, + ); +} + +#[test] +fn should_gather_3d_dim1() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data( + [ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], + ], + &device, + ); + let indices = + TestTensorInt::from_ints([[[1, 0, 0], [0, 1, 0]], [[0, 0, 1], [0, 1, 1]]], &device); + + let output = tensor.gather(1, indices); + let expected = TensorData::from([ + [[3.0, 1.0, 2.0], [0.0, 4.0, 2.0]], + [[6.0, 7.0, 11.0], [6.0, 10.0, 11.0]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_gather_2d_only_1dim() { + let device = Default::default(); + let tensor = TestTensor::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &device); + let indices = TestTensorInt::<2>::from_ints([[1, 2]], &device).reshape([2, 1]); + + let output = tensor.gather(1, indices); + + output + .into_data() + .assert_eq(&TensorData::from([[1.0], [5.0]]), false); +} + +#[test] +fn should_scatter_add_1d() { + let device = Default::default(); + let tensor = TestTensor::<1>::from_data([0.0, 0.0, 0.0], &device); + let values = TestTensor::from_data([5.0, 4.0, 3.0], &device); + let indices = TestTensorInt::from_ints([1, 0, 2], &device); + + let output = tensor.scatter(0, indices, values, IndexingUpdateOp::Add); + + output + .into_data() + .assert_eq(&TensorData::from([4.0, 5.0, 3.0]), false); +} + +#[test] +fn should_scatter_add_2d_dim0() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], &device); + let values = TestTensor::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + let indices = TestTensorInt::from_ints([[1, 0, 1], [1, 1, 0]], &device); + + let output = tensor.scatter(0, indices, values, IndexingUpdateOp::Add); + + output + .into_data() + .assert_eq(&TensorData::from([[0.0, 2.0, 6.0], [5.0, 5.0, 3.0]]), false); +} + +#[test] +fn should_scatter_add_2d_dim1() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], &device); + let values = TestTensor::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + let indices = TestTensorInt::from_ints([[1, 0, 2], [1, 2, 0]], &device); + + let output = tensor.scatter(1, indices, values, IndexingUpdateOp::Add); + + output + .into_data() + .assert_eq(&TensorData::from([[2.0, 1.0, 3.0], [6.0, 4.0, 5.0]]), false); +} + +#[test] +fn should_scatter_add_3d_dim1() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data( + [ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], + ], + &device, + ); + let values = TestTensor::from_data( + [ + [[12.0, 13.0, 14.0], [15.0, 16.0, 17.0]], + [[18.0, 19.0, 20.0], [21.0, 22.0, 23.0]], + ], + &device, + ); + let indices = + TestTensorInt::from_ints([[[1, 0, 0], [0, 1, 0]], [[0, 0, 1], [0, 1, 1]]], &device); + + let output = tensor.scatter(1, indices, values, IndexingUpdateOp::Add); + let expected = TensorData::from([ + [[15.0, 14.0, 33.0], [15.0, 20.0, 5.0]], + [[45.0, 26.0, 8.0], [9.0, 32.0, 54.0]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_scatter_add_2d_dim1_diff_shape() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], &device); + let values = TestTensor::from_data([[1.0], [4.0]], &device); + let indices = TestTensorInt::from_ints([[1], [2]], &device); + + let output = tensor.scatter(1, indices, values, IndexingUpdateOp::Add); + + output + .into_data() + .assert_eq(&TensorData::from([[0.0, 1.0, 0.0], [0.0, 0.0, 4.0]]), false); +} + +#[test] +#[should_panic] +fn scatter_should_panic_on_mismatch_of_shapes() { + let device = Default::default(); + let tensor = TestTensor::<1>::from_data([0.0, 0.0, 0.0], &device); + let values = TestTensor::from_data([5.0, 4.0], &device); + let indices = TestTensorInt::from_ints([1, 0, 2], &device); + + tensor.scatter(0, indices, values, IndexingUpdateOp::Add); +} + +#[test] +fn should_scatter_add_3d_dim0_with_index_view() { + // Index tensor expanded from 1D via unsqueeze + repeat_dim — a common + // pattern to scatter whole rows along dim 0. The repeat of a unit dim + // is a broadcast view; reshaping it must keep stride 1 on dim 0, + // otherwise every iteration reads `indices[0]` and all updates pile + // onto one row while the other rows get none. + let device = Default::default(); + let tensor = TestTensor::<3>::from_data([[[1.0]], [[2.0]], [[3.0]], [[4.0]]], &device); + let values = TestTensor::<3>::from_data([[[10.0]], [[20.0]]], &device); + let indices = TestTensorInt::<1>::from_ints([2, 0], &device); + let indices: TestTensorInt<2> = indices.unsqueeze_dim(1).repeat_dim(1, 1); + let indices: TestTensorInt<3> = indices.unsqueeze_dim(2).repeat_dim(2, 1); + + let output = tensor.scatter(0, indices, values, IndexingUpdateOp::Add); + + output.into_data().assert_eq( + &TensorData::from([[[21.0]], [[2.0]], [[13.0]], [[4.0]]]), + false, + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/gather_scatter_nd.rs b/crates/burn-backend-tests/tests/tensor/float/ops/gather_scatter_nd.rs new file mode 100644 index 0000000..bab7274 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/gather_scatter_nd.rs @@ -0,0 +1,356 @@ +use super::*; +use burn_tensor::{IndexingUpdateOp, TensorData}; + +#[test] +fn test_gather_nd_2d_k_equals_d() { + // Gather single elements from a 2D tensor (K=D=2) + let device = Default::default(); + // data shape: [3, 3] + let data = + TestTensor::<2>::from_floats([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]], &device); + // indices shape: [3, 2] => K=2, M=2, DV = M-1+D-K = 1+2-2 = 1 + // Each row is [row, col] + let indices = TestTensorInt::<2>::from_ints([[0, 1], [1, 2], [2, 0]], &device); + + let output: TestTensor<1> = data.gather_nd(indices); + + output + .into_data() + .assert_eq(&TensorData::from([1.0, 5.0, 6.0]), false); +} + +#[test] +fn test_gather_nd_3d_k1_slices() { + // Gather 2D slices from a 3D tensor (K=1) + let device = Default::default(); + // data shape: [2, 3, 4] + let data = TestTensor::<3>::from_floats( + [ + [ + [0.0, 1.0, 2.0, 3.0], + [4.0, 5.0, 6.0, 7.0], + [8.0, 9.0, 10.0, 11.0], + ], + [ + [12.0, 13.0, 14.0, 15.0], + [16.0, 17.0, 18.0, 19.0], + [20.0, 21.0, 22.0, 23.0], + ], + ], + &device, + ); + // indices shape: [2, 1] => K=1, M=2, DV = 1 + 3 - 1 = 3 + // Each index selects a whole 2D slice along dim 0 + let indices = TestTensorInt::<2>::from_ints([[1], [0]], &device); + + let output: TestTensor<3> = data.gather_nd(indices); + + output.into_data().assert_eq( + &TensorData::from([ + [ + [12.0, 13.0, 14.0, 15.0], + [16.0, 17.0, 18.0, 19.0], + [20.0, 21.0, 22.0, 23.0], + ], + [ + [0.0, 1.0, 2.0, 3.0], + [4.0, 5.0, 6.0, 7.0], + [8.0, 9.0, 10.0, 11.0], + ], + ]), + false, + ); +} + +#[test] +fn test_gather_nd_3d_k_equals_d() { + // Gather scalar elements from a 3D tensor (K=D=3) + let device = Default::default(); + let data = TestTensor::<3>::from_floats( + [[[0.0, 1.0], [2.0, 3.0]], [[4.0, 5.0], [6.0, 7.0]]], + &device, + ); + // indices shape: [2, 3] => K=3, M=2, DV = 1+3-3 = 1 + let indices = TestTensorInt::<2>::from_ints([[0, 0, 1], [1, 1, 0]], &device); + + let output: TestTensor<1> = data.gather_nd(indices); + + output + .into_data() + .assert_eq(&TensorData::from([1.0, 6.0]), false); +} + +#[test] +fn test_gather_nd_batch() { + // Batch dimension in indices + let device = Default::default(); + // data shape: [2, 3] + let data = TestTensor::<2>::from_floats([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]], &device); + // indices shape: [3, 2] => K=2, M=2, DV = 1+2-2 = 1 + let indices = TestTensorInt::<2>::from_ints([[0, 0], [0, 2], [1, 1]], &device); + + let output: TestTensor<1> = data.gather_nd(indices); + + output + .into_data() + .assert_eq(&TensorData::from([10.0, 30.0, 50.0]), false); +} + +#[test] +fn test_scatter_nd_assign_2d() { + let device = Default::default(); + // data shape: [3, 3], initialized to zeros + let data = TestTensor::<2>::zeros([3, 3], &device); + // indices shape: [2, 2] => K=2, M=2, DV = 1+2-2 = 1 + let indices = TestTensorInt::<2>::from_ints([[0, 1], [2, 0]], &device); + // values shape: [2] (scalar per index tuple) + let values = TestTensor::<1>::from_floats([10.0, 20.0], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Assign); + + output.into_data().assert_eq( + &TensorData::from([[0.0, 10.0, 0.0], [0.0, 0.0, 0.0], [20.0, 0.0, 0.0]]), + false, + ); +} + +#[test] +fn test_scatter_nd_add_2d() { + let device = Default::default(); + let data = + TestTensor::<2>::from_floats([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], &device); + let indices = TestTensorInt::<2>::from_ints([[0, 1], [2, 0]], &device); + let values = TestTensor::<1>::from_floats([10.0, 20.0], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Add); + + output.into_data().assert_eq( + &TensorData::from([[1.0, 11.0, 1.0], [1.0, 1.0, 1.0], [21.0, 1.0, 1.0]]), + false, + ); +} + +// Duplicate indices are non-deterministic on GPU; only test on CPU backends. +#[cfg(feature = "ndarray")] +#[test] +fn test_scatter_nd_add_duplicate_indices() { + let device = Default::default(); + let data = TestTensor::<2>::zeros([2, 3], &device); + // Both index tuples point to the same location + let indices = TestTensorInt::<2>::from_ints([[0, 1], [0, 1]], &device); + let values = TestTensor::<1>::from_floats([5.0, 3.0], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Add); + + // Values should accumulate: 5.0 + 3.0 = 8.0 + output + .into_data() + .assert_eq(&TensorData::from([[0.0, 8.0, 0.0], [0.0, 0.0, 0.0]]), false); +} + +#[test] +fn test_scatter_nd_3d_slices() { + // Scatter 1D slices into a 3D tensor (K=1) + let device = Default::default(); + // data shape: [2, 3] + let data = TestTensor::<2>::zeros([2, 3], &device); + // indices shape: [2, 1] => K=1, M=2, DV = 1+2-1 = 2 + // Each index selects a row + let indices = TestTensorInt::<2>::from_ints([[0], [1]], &device); + // values shape: [2, 3] + let values = TestTensor::<2>::from_floats([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Assign); + + output.into_data().assert_eq( + &TensorData::from([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]]), + false, + ); +} + +#[test] +fn test_scatter_nd_mul() { + let device = Default::default(); + let data = TestTensor::<2>::from_floats([[2.0, 3.0, 4.0], [5.0, 6.0, 7.0]], &device); + let indices = TestTensorInt::<2>::from_ints([[0, 1], [1, 2]], &device); + let values = TestTensor::<1>::from_floats([10.0, 2.0], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Mul); + + output.into_data().assert_eq( + &TensorData::from([[2.0, 30.0, 4.0], [5.0, 6.0, 14.0]]), + false, + ); +} + +#[test] +fn test_scatter_nd_min() { + let device = Default::default(); + let data = TestTensor::<2>::from_floats([[5.0, 10.0], [15.0, 20.0]], &device); + let indices = TestTensorInt::<2>::from_ints([[0, 0], [1, 1]], &device); + let values = TestTensor::<1>::from_floats([3.0, 25.0], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Min); + + // min(5.0, 3.0) = 3.0; min(20.0, 25.0) = 20.0 + output + .into_data() + .assert_eq(&TensorData::from([[3.0, 10.0], [15.0, 20.0]]), false); +} + +#[test] +fn test_scatter_nd_max() { + let device = Default::default(); + let data = TestTensor::<2>::from_floats([[5.0, 10.0], [15.0, 20.0]], &device); + let indices = TestTensorInt::<2>::from_ints([[0, 0], [1, 1]], &device); + let values = TestTensor::<1>::from_floats([3.0, 25.0], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Max); + + // max(5.0, 3.0) = 5.0; max(20.0, 25.0) = 25.0 + output + .into_data() + .assert_eq(&TensorData::from([[5.0, 10.0], [15.0, 25.0]]), false); +} + +#[test] +fn test_scatter_nd_batch() { + // Scatter with batch dimensions + let device = Default::default(); + // data shape: [3, 3] + let data = TestTensor::<2>::zeros([3, 3], &device); + // indices shape: [3, 1] => K=1, M=2, DV = 1+2-1 = 2 + let indices = TestTensorInt::<2>::from_ints([[0], [2], [1]], &device); + // values shape: [3, 3] + let values = + TestTensor::<2>::from_floats([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Assign); + + output.into_data().assert_eq( + &TensorData::from([[1.0, 2.0, 3.0], [7.0, 8.0, 9.0], [4.0, 5.0, 6.0]]), + false, + ); +} + +#[test] +fn test_scatter_nd_single_element() { + // Single update to a single element in a 1D tensor + let device = Default::default(); + let data = TestTensor::<1>::from_floats([1.0, 2.0, 3.0], &device); + let indices = TestTensorInt::<2>::from_ints([[2]], &device); + let values = TestTensor::<1>::from_floats([99.0], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Assign); + + output + .into_data() + .assert_eq(&TensorData::from([1.0, 2.0, 99.0]), false); +} + +#[test] +fn test_scatter_nd_k1_high_rank() { + // K=1 on a 3D tensor: each index selects a full 2D slice + let device = Default::default(); + let data = TestTensor::<3>::ones([2, 2, 2], &device); + let indices = TestTensorInt::<2>::from_ints([[1]], &device); + let values = TestTensor::<3>::from_floats([[[10.0, 20.0], [30.0, 40.0]]], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Assign); + + output.into_data().assert_eq( + &TensorData::from([[[1.0, 1.0], [1.0, 1.0]], [[10.0, 20.0], [30.0, 40.0]]]), + false, + ); +} + +#[test] +fn test_gather_nd_single_element() { + // Gather a single scalar from a 1D tensor + let device = Default::default(); + let data = TestTensor::<1>::from_floats([10.0, 20.0, 30.0], &device); + let indices = TestTensorInt::<2>::from_ints([[1]], &device); + + let output: TestTensor<1> = data.gather_nd(indices); + + output + .into_data() + .assert_eq(&TensorData::from([20.0]), false); +} + +#[test] +fn test_gather_scatter_nd_roundtrip() { + // gather_nd then scatter_nd_add back should reconstruct selected rows + let device = Default::default(); + let data = + TestTensor::<2>::from_floats([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], &device); + let indices = TestTensorInt::<2>::from_ints([[0], [2]], &device); + + let gathered: TestTensor<2> = data.clone().gather_nd(indices.clone()); + let zeros = TestTensor::<2>::zeros([3, 3], &device); + let reconstructed: TestTensor<2> = zeros.scatter_nd(indices, gathered, IndexingUpdateOp::Add); + + reconstructed.into_data().assert_eq( + &TensorData::from([[1.0, 2.0, 3.0], [0.0, 0.0, 0.0], [7.0, 8.0, 9.0]]), + false, + ); +} + +#[test] +fn test_scatter_nd_3d_k2_partial_slices() { + // K=2 on 3D data: each index tuple selects a 1D row within a 2D slice + let device = Default::default(); + // data shape: [2, 3, 4] + let data = TestTensor::<3>::zeros([2, 3, 4], &device); + // indices shape: [2, 2] => K=2, M=2, DV = 1+3-2 = 2 + // [0, 1] => data[0, 1, :], [1, 2] => data[1, 2, :] + let indices = TestTensorInt::<2>::from_ints([[0, 1], [1, 2]], &device); + // values shape: [2, 4] + let values = TestTensor::<2>::from_floats( + [[10.0, 20.0, 30.0, 40.0], [50.0, 60.0, 70.0, 80.0]], + &device, + ); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Add); + + // Only data[0,1,:] and data[1,2,:] should be set + let output_data = output.into_data(); + output_data.assert_eq( + &TensorData::from([ + [ + [0.0, 0.0, 0.0, 0.0], + [10.0, 20.0, 30.0, 40.0], + [0.0, 0.0, 0.0, 0.0], + ], + [ + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [50.0, 60.0, 70.0, 80.0], + ], + ]), + false, + ); +} + +#[test] +fn test_gather_nd_3d_k2_partial_slices() { + // K=2 on 3D data: each index tuple gathers a 1D row + let device = Default::default(); + let data = TestTensor::<3>::from_floats( + [ + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], + [[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]], + ], + &device, + ); + // indices shape: [3, 2] => K=2, M=2, DV = 1+3-2 = 2 + let indices = TestTensorInt::<2>::from_ints([[0, 2], [1, 0], [0, 1]], &device); + + let output: TestTensor<2> = data.gather_nd(indices); + + // [0,2] => [5,6], [1,0] => [7,8], [0,1] => [3,4] + output.into_data().assert_eq( + &TensorData::from([[5.0, 6.0], [7.0, 8.0], [3.0, 4.0]]), + false, + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/grid_sample.rs b/crates/burn-backend-tests/tests/tensor/float/ops/grid_sample.rs new file mode 100644 index 0000000..fe21442 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/grid_sample.rs @@ -0,0 +1,126 @@ +use super::*; +use burn_tensor::{ + TensorData, Tolerance, + ops::{GridSampleOptions, GridSamplePaddingMode, InterpolateMode}, +}; + +/// Tests grid_sample_2d with default options (align_corners=false, zeros padding). +/// +/// For a 3x3 input with grid coordinates: +/// - (0.0, 0.0) maps to pixel (1.0, 1.0) -> center pixel = 4.0 +/// - (-1.0, 0.25) maps to pixel (-0.5, 1.375) -> partially out of bounds +/// - (1.0, 1.0) maps to pixel (2.5, 2.5) -> corner, partially out of bounds +/// - (0.2, -0.8) maps to pixel (1.3, 0.3) -> interpolates around center-top +#[test] +fn should_grid_sample_2d_default() { + let device = Default::default(); + let tensor = TestTensor::<4>::from_data( + [[[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]]], + &device, + ); + let grid = TestTensor::<4>::from_data( + [[[[0.0, 0.0], [-1.0, 0.25]], [[1.0, 1.0], [0.2, -0.8]]]], + &device, + ); + + let output = tensor.grid_sample_2d(grid, InterpolateMode::Bilinear); + + // Expected values computed with PyTorch grid_sample(align_corners=False, padding_mode='zeros') + let expected = TensorData::from([[[[4.0, 2.0625], [2.0, 1.04]]]]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +/// Tests grid_sample_2d with align_corners=true and border padding. +/// +/// This is the original Burn semantics before the API change. +#[test] +fn should_grid_sample_2d_align_corners_border() { + let device = Default::default(); + let tensor = TestTensor::<4>::from_data( + [[[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]]], + &device, + ); + let grid = TestTensor::<4>::from_data( + [[[[0.0, 0.0], [-1.0, 0.25]], [[1.0, 1.0], [0.2, -0.8]]]], + &device, + ); + + let options = GridSampleOptions::new(InterpolateMode::Bilinear) + .with_padding_mode(GridSamplePaddingMode::Border) + .with_align_corners(true); + let output = tensor.grid_sample_2d(grid, options); + + // Expected values computed with PyTorch grid_sample(align_corners=True, padding_mode='border') + let expected = TensorData::from([[[[4.0, 3.75], [8.0, 1.8]]]]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +/// Tests out-of-bounds grid coordinates with zeros padding. +/// Grid coordinate (0.0, -2.0) maps to pixel (1.0, -2.5) which is completely out of bounds. +#[test] +fn should_pad_zeros_grid_sample_2d() { + let device = Default::default(); + let tensor = TestTensor::<4>::from_data( + [[[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]]], + &device, + ); + let grid = TestTensor::<4>::from_data([[[[0.0, -2.0]]]], &device); + + let output = tensor.grid_sample_2d(grid, GridSampleOptions::default()); + + // With zeros padding, out-of-bounds samples return 0 + let expected = TensorData::from([[[[0.0]]]]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +/// Tests out-of-bounds grid coordinates with border padding. +#[test] +fn should_pad_border_grid_sample_2d() { + let device = Default::default(); + let tensor = TestTensor::<4>::from_data( + [[[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]]], + &device, + ); + let grid = TestTensor::<4>::from_data([[[[0.0, -2.0]]]], &device); + + let options = GridSampleOptions::new(InterpolateMode::Bilinear) + .with_padding_mode(GridSamplePaddingMode::Border); + let output = tensor.grid_sample_2d(grid, options); + + // With border padding, out-of-bounds coordinates are clamped to border + // Grid (0.0, -2.0) with align_corners=false: pixel (1.0, -2.5) -> clamped to (1.0, 0.0) = 1.0 + let expected = TensorData::from([[[[1.0]]]]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +/// Tests bilinear interpolation with reflection padding. +#[test] +fn should_pad_reflection_grid_sample_2d() { + let device = Default::default(); + let tensor = TestTensor::<4>::from_data( + [[[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]]], + &device, + ); + let grid = TestTensor::<4>::from_data( + [[[[0.0, 0.0], [-1.0, 0.25]], [[1.0, 1.0], [0.2, -0.8]]]], + &device, + ); + + let options = GridSampleOptions::new(InterpolateMode::Bilinear) + .with_padding_mode(GridSamplePaddingMode::Reflection); + let output = tensor.grid_sample_2d(grid, options); + + // Expected values computed with PyTorch F.grid_sample(mode='bilinear', padding_mode='reflection', align_corners=False) + let expected = TensorData::from([[[[4.0, 4.125], [8.0, 1.3]]]]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/hamming_window.rs b/crates/burn-backend-tests/tests/tensor/float/ops/hamming_window.rs new file mode 100644 index 0000000..0d75bb6 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/hamming_window.rs @@ -0,0 +1,65 @@ +use super::*; +use burn_tensor::signal::hamming_window; +use burn_tensor::{DType, TensorData, Tolerance}; + +#[test] +fn should_support_hamming_window_periodic() { + let tensor: TestTensor<1> = hamming_window(8, true, &Default::default()); + let expected = TensorData::from([ + 0.086957, 0.220669, 0.543478, 0.866288, 1.0, 0.866288, 0.543478, 0.220669, + ]); + + // Metal has less precise trigonometric functions. + let tolerance = Tolerance::default().set_half_precision_relative(1e-2); + + tensor + .into_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn should_support_hamming_window_symmetric() { + let tensor: TestTensor<1> = hamming_window(8, false, &Default::default()); + let expected = TensorData::from([ + 0.086957, 0.258842, 0.645064, 0.954790, 0.954790, 0.645064, 0.258842, 0.086957, + ]); + + // Metal has less precise trigonometric functions. + let tolerance = Tolerance::default().set_half_precision_relative(1e-2); + + tensor + .into_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn should_support_hamming_window_options_dtype() { + let tensor: TestTensor<1> = hamming_window(4, true, (&Default::default(), DType::F32)); + assert_eq!(tensor.dtype(), DType::F32); +} + +#[test] +fn should_support_hamming_window_empty() { + let tensor: TestTensor<1> = hamming_window(0, true, &Default::default()); + assert_eq!(tensor.shape().dims(), [0]); +} + +#[test] +fn should_handle_hamming_window_size_one_symmetric() { + let tensor: TestTensor<1> = hamming_window(1, false, &Default::default()); + let expected = TensorData::from([1.0]); + + tensor + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_handle_hamming_window_size_one_periodic() { + let tensor: TestTensor<1> = hamming_window(1, true, &Default::default()); + let expected = TensorData::from([1.0]); + + tensor + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/hann_window.rs b/crates/burn-backend-tests/tests/tensor/float/ops/hann_window.rs new file mode 100644 index 0000000..0604617 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/hann_window.rs @@ -0,0 +1,63 @@ +use super::*; +use burn_tensor::signal::hann_window; +use burn_tensor::{DType, TensorData, Tolerance}; + +#[test] +fn should_support_hann_window_periodic() { + let tensor: TestTensor<1> = hann_window(8, true, &Default::default()); + let expected = TensorData::from([0.0, 0.146447, 0.5, 0.853553, 1.0, 0.853553, 0.5, 0.146447]); + + // Metal has less precise trigonometric functions. + let tolerance = Tolerance::default().set_half_precision_relative(1e-2); + + tensor + .into_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn should_support_hann_window_symmetric() { + let tensor: TestTensor<1> = hann_window(8, false, &Default::default()); + let expected = TensorData::from([ + 0.0, 0.188255, 0.611260, 0.950484, 0.950484, 0.611260, 0.188255, 0.0, + ]); + + // Metal has less precise trigonometric functions. + let tolerance = Tolerance::default().set_half_precision_relative(1e-2); + + tensor + .into_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn should_support_hann_window_options_dtype() { + let tensor: TestTensor<1> = hann_window(4, true, (&Default::default(), DType::F32)); + assert_eq!(tensor.dtype(), DType::F32); +} + +#[test] +fn should_support_hann_window_empty() { + let tensor: TestTensor<1> = hann_window(0, true, &Default::default()); + assert_eq!(tensor.shape().dims(), [0]); +} + +#[test] +fn should_handle_hann_window_size_one_symmetric() { + let tensor: TestTensor<1> = hann_window(1, false, &Default::default()); + let expected = TensorData::from([1.0]); + + tensor + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_handle_hann_window_size_one_periodic() { + let tensor: TestTensor<1> = hann_window(1, true, &Default::default()); + let expected = TensorData::from([1.0]); + + tensor + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/hypot.rs b/crates/burn-backend-tests/tests/tensor/float/ops/hypot.rs new file mode 100644 index 0000000..258a3d2 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/hypot.rs @@ -0,0 +1,47 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn should_support_hypot_basic() { + let data_a = TensorData::from([[3.0, 4.0], [5.0, 12.0]]); + let data_b = TensorData::from([[4.0, 3.0], [12.0, 5.0]]); + let tensor_a = TestTensor::<2>::from_data(data_a, &Default::default()); + let tensor_b = TestTensor::<2>::from_data(data_b, &Default::default()); + + let result = tensor_a.hypot(tensor_b); + let expected = TensorData::from([[5.0, 5.0], [13.0, 13.0]]); + + result + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_hypot_broadcast() { + let data_a = TensorData::from([[3.0, 4.0, 5.0]]); + let data_b = TensorData::from([[4.0], [3.0]]); + let tensor_a = TestTensor::<2>::from_data(data_a, &Default::default()); + let tensor_b = TestTensor::<2>::from_data(data_b, &Default::default()); + + let result = tensor_a.hypot(tensor_b); + let expected = TensorData::from([[5.0, 5.656854, 6.404], [4.2426405, 5.0, 5.831]]); + + result + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_hypot_zero() { + let data_a = TensorData::from([[0.0, 3.0], [0.0, 0.0]]); + let data_b = TensorData::from([[4.0, 0.0], [0.0, 5.0]]); + let tensor_a = TestTensor::<2>::from_data(data_a, &Default::default()); + let tensor_b = TestTensor::<2>::from_data(data_b, &Default::default()); + + let result = tensor_a.hypot(tensor_b); + let expected = TensorData::from([[4.0, 3.0], [0.0, 5.0]]); + + result + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/inf.rs b/crates/burn-backend-tests/tests/tensor/float/ops/inf.rs new file mode 100644 index 0000000..f0c6cef --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/inf.rs @@ -0,0 +1,21 @@ +use super::*; + +#[test] +fn is_inf() { + let no_inf = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let no_inf_expected = TestTensorBool::<2>::from([[false, false, false], [false, false, false]]); + + let with_inf = + TestTensor::<2>::from([[0.0, f32::INFINITY, 2.0], [f32::NEG_INFINITY, 4.0, 5.0]]); + let with_inf_expected = TestTensorBool::<2>::from([[false, true, false], [true, false, false]]); + + no_inf + .is_inf() + .into_data() + .assert_eq(&no_inf_expected.into_data(), false); + + with_inf + .is_inf() + .into_data() + .assert_eq(&with_inf_expected.into_data(), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/init.rs b/crates/burn-backend-tests/tests/tensor/float/ops/init.rs new file mode 100644 index 0000000..eaa7b58 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/init.rs @@ -0,0 +1,62 @@ +use super::*; +use burn_tensor::{DType, TensorData}; + +#[test] +fn should_support_float_empty() { + let shape = [2, 2]; + let tensor = TestTensor::<2>::empty(shape, &Default::default()); + assert_eq!(tensor.shape(), shape.into()) +} + +#[test] +fn should_support_float_empty_options() { + let shape = [2, 2]; + let tensor = TestTensor::<2>::empty(shape, (&Default::default(), DType::F32)); + assert_eq!(tensor.shape(), shape.into()) +} + +#[test] +fn should_support_float_zeros() { + let shape = [2, 2]; + let tensor = TestTensor::<2>::zeros(shape, &Default::default()); + assert_eq!(tensor.shape(), shape.into()); + + tensor + .into_data() + .assert_eq(&TensorData::from([[0., 0.], [0., 0.]]), false); +} + +#[test] +fn should_support_float_zeros_options() { + let shape = [2, 2]; + let tensor = TestTensor::<2>::zeros(shape, (&Default::default(), DType::F32)); + assert_eq!(tensor.shape(), shape.into()); + assert_eq!(tensor.dtype(), DType::F32); + + tensor + .into_data() + .assert_eq(&TensorData::from([[0., 0.], [0., 0.]]), false); +} + +#[test] +fn should_support_float_ones() { + let shape = [2, 2]; + let tensor = TestTensor::<2>::ones(shape, &Default::default()); + assert_eq!(tensor.shape(), shape.into()); + + tensor + .into_data() + .assert_eq(&TensorData::from([[1., 1.], [1., 1.]]), false); +} + +#[test] +fn should_support_float_ones_options() { + let shape = [2, 2]; + let tensor = TestTensor::<2>::ones(shape, (&Default::default(), DType::F32)); + assert_eq!(tensor.shape(), shape.into()); + assert_eq!(tensor.dtype(), DType::F32); + + tensor + .into_data() + .assert_eq(&TensorData::from([[1., 1.], [1., 1.]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/iter_dim.rs b/crates/burn-backend-tests/tests/tensor/float/ops/iter_dim.rs new file mode 100644 index 0000000..6535566 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/iter_dim.rs @@ -0,0 +1,246 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_1d_iter_last_item() { + let data = [1, 2, 3, 4]; + let device = Default::default(); + let tensor = TestTensorInt::<1>::from_ints(data, &device); + tensor + .iter_dim(0) + .last() + .unwrap() + .into_data() + .assert_eq(&TensorData::from([4]), false); +} + +#[test] +#[should_panic] +fn test_too_high_dimension() { + TestTensor::<1>::zeros([10], &Default::default()).iter_dim(1); +} + +#[test] +fn test_transposed() { + let data = [ + [1., 2., 3., 1., 2.], + [4., 5., 6., 1., 2.], + [7., 8., 9., 1., 2.], + ]; + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + let lhs = tensor.clone().slice([1..2, 0..5]); + let rhs = tensor.transpose().iter_dim(1).nth(1).unwrap(); + assert_eq!( + lhs.into_data().as_slice::().unwrap(), + rhs.into_data().as_slice::().unwrap() + ); +} + +#[test] +fn test_2d_iter_dim() { + let tensor = + TestTensor::<2>::from_data([[3.0, 4.9, 2.0], [2.0, 1.9, 3.0]], &Default::default()); + + let mut iter = tensor.iter_dim(0); + + let iter1 = iter.next().unwrap(); + iter1 + .into_data() + .assert_eq(&TensorData::from([[3.0, 4.9, 2.0]]), false); + + let iter2 = iter.next().unwrap(); + iter2 + .into_data() + .assert_eq(&TensorData::from([[2.0, 1.9, 3.0]]), false); + + assert!(iter.next().is_none()); +} + +#[test] +fn test_2d_iter_dim1() { + let tensor = + TestTensor::<2>::from_data([[3.0, 4.9, 2.0], [2.0, 1.9, 3.0]], &Default::default()); + + let mut iter = tensor.iter_dim(1); + + let iter1 = iter.next().unwrap(); + iter1 + .into_data() + .assert_eq(&TensorData::from([[3.0], [2.0]]), false); + + let iter2 = iter.next().unwrap(); + iter2 + .into_data() + .assert_eq(&TensorData::from([[4.9], [1.9]]), false); + + let iter3 = iter.next().unwrap(); + iter3 + .into_data() + .assert_eq(&TensorData::from([[2.0], [3.0]]), false); + + assert!(iter.next().is_none()); +} + +#[test] +fn test_3d_iter_dim() { + let tensor = TestTensor::<3>::from([[ + [1., 2., 3., 1., 2.], + [4., 5., 6., 1., 2.], + [7., 8., 9., 1., 2.], + ]]); + + let mut iter = tensor.clone().iter_dim(0); + + let iter1 = iter.next().unwrap(); + iter1.into_data().assert_eq(&tensor.into_data(), true); + + assert!(iter.next().is_none()); +} + +#[test] +fn test_3d_iter_dim1() { + let tensor = TestTensor::<3>::from([[ + [1., 2., 3., 1., 2.], + [4., 5., 6., 1., 2.], + [7., 8., 9., 1., 2.], + ]]); + + let mut iter = tensor.iter_dim(1); + + let iter1 = iter.next().unwrap(); + iter1 + .into_data() + .assert_eq(&TensorData::from([[[1., 2., 3., 1., 2.]]]), false); + + let iter2 = iter.next().unwrap(); + iter2 + .into_data() + .assert_eq(&TensorData::from([[[4., 5., 6., 1., 2.]]]), false); + + let iter3 = iter.next().unwrap(); + iter3 + .into_data() + .assert_eq(&TensorData::from([[[7., 8., 9., 1., 2.]]]), false); + + assert!(iter.next().is_none()); +} + +#[test] +fn test_3d_iter_dim2() { + let tensor = TestTensor::<3>::from([[ + [1., 2., 3., 1., 2.], + [4., 5., 6., 1., 2.], + [7., 8., 9., 1., 2.], + ]]); + + let mut iter = tensor.iter_dim(2); + + let iter1 = iter.next().unwrap(); + iter1 + .into_data() + .assert_eq(&TensorData::from([[[1.], [4.], [7.]]]), false); + + let iter2 = iter.next().unwrap(); + iter2 + .into_data() + .assert_eq(&TensorData::from([[[2.], [5.], [8.]]]), false); + + let iter3 = iter.next().unwrap(); + iter3 + .into_data() + .assert_eq(&TensorData::from([[[3.], [6.], [9.]]]), false); + + let iter4 = iter.next().unwrap(); + iter4 + .into_data() + .assert_eq(&TensorData::from([[[1.], [1.], [1.]]]), false); + + let iter5 = iter.next().unwrap(); + iter5 + .into_data() + .assert_eq(&TensorData::from([[[2.], [2.], [2.]]]), false); + + assert!(iter.next().is_none()); +} + +#[test] +fn test_iteration_over_low_dim() { + let data = [[ + [1., 2., 3., 1., 2.], + [4., 5., 6., 1., 2.], + [7., 8., 9., 1., 2.], + ]]; + + let tensor = TestTensor::<3>::from_data(data, &Default::default()); + + let lhs = tensor.iter_dim(2).nth(1).unwrap(); + let rhs = TestTensor::<1>::from([2., 5., 8.]); + assert_eq!( + lhs.into_data().as_slice::().unwrap(), + rhs.into_data().as_slice::().unwrap() + ); +} + +#[test] +fn test_iter_dim_double_end() { + let input = TestTensorInt::<1>::arange(0..(4 * 6 * 3), &Default::default()).reshape([4, 6, 3]); + let mut iter = input.iter_dim(1); + + let ele0 = TensorData::from([[[0, 1, 2]], [[18, 19, 20]], [[36, 37, 38]], [[54, 55, 56]]]); + let ele1 = TensorData::from([[[3, 4, 5]], [[21, 22, 23]], [[39, 40, 41]], [[57, 58, 59]]]); + let ele2 = TensorData::from([[[6, 7, 8]], [[24, 25, 26]], [[42, 43, 44]], [[60, 61, 62]]]); + let ele3 = TensorData::from([ + [[9, 10, 11]], + [[27, 28, 29]], + [[45, 46, 47]], + [[63, 64, 65]], + ]); + let ele4 = TensorData::from([ + [[12, 13, 14]], + [[30, 31, 32]], + [[48, 49, 50]], + [[66, 67, 68]], + ]); + let ele5 = TensorData::from([ + [[15, 16, 17]], + [[33, 34, 35]], + [[51, 52, 53]], + [[69, 70, 71]], + ]); + + iter.next().unwrap().into_data().assert_eq(&ele0, false); + iter.next_back() + .unwrap() + .into_data() + .assert_eq(&ele5, false); + iter.next_back() + .unwrap() + .into_data() + .assert_eq(&ele4, false); + iter.next().unwrap().into_data().assert_eq(&ele1, false); + iter.next().unwrap().into_data().assert_eq(&ele2, false); + iter.next().unwrap().into_data().assert_eq(&ele3, false); + assert!(iter.next().is_none()); + assert!(iter.next_back().is_none()); +} + +#[test] +fn test_iter_dim_single_element() { + let input = TestTensorInt::<1>::arange(0..(4 * 3), &Default::default()).reshape([4, 1, 3]); + + let mut iter = input.clone().iter_dim(1); + iter.next() + .unwrap() + .into_data() + .assert_eq(&input.clone().into_data(), false); + assert!(iter.next_back().is_none()); + assert!(iter.next().is_none()); + + let mut iter = input.clone().iter_dim(1); + iter.next_back() + .unwrap() + .into_data() + .assert_eq(&input.clone().into_data(), false); + assert!(iter.next().is_none()); + assert!(iter.next_back().is_none()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/log.rs b/crates/burn-backend-tests/tests/tensor/float/ops/log.rs new file mode 100644 index 0000000..7aa64dc --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/log.rs @@ -0,0 +1,39 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; +use core::f32::consts::E; + +#[test] +fn should_support_log_3d_transposed() { + // 3D tensor with permuted dimensions; log should undo the e^k powers. + let data = TensorData::from([ + [[1.0, E], [E * E, E * E * E]], + [[1.0, E], [E * E, E * E * E]], + ]); + let tensor = TestTensor::<3>::from_data(data, &Default::default()); + let permuted = tensor.permute([2, 0, 1]); + + let output = permuted.log(); + let expected = TensorData::from([[[0.0, 2.0], [0.0, 2.0]], [[1.0, 3.0], [1.0, 3.0]]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_log_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.log(); + let expected = TensorData::from([ + [-f32::INFINITY, 0.0, core::f32::consts::LN_2], + [1.09861, 1.38629, 1.60944], + ]); + + output.into_data().assert_approx_eq::( + &expected, + Tolerance::default().set_half_precision_relative(1e-3), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/log1p.rs b/crates/burn-backend-tests/tests/tensor/float/ops/log1p.rs new file mode 100644 index 0000000..289b264 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/log1p.rs @@ -0,0 +1,20 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_exp_log1p() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.log1p(); + let expected = TensorData::from([ + [0.0, core::f32::consts::LN_2, 1.09861], + [1.38629, 1.60944, 1.79176], + ]); + + output.into_data().assert_approx_eq::( + &expected, + Tolerance::default().set_half_precision_relative(1e-3), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/mask.rs b/crates/burn-backend-tests/tests/tensor/float/ops/mask.rs new file mode 100644 index 0000000..9a06b57 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/mask.rs @@ -0,0 +1,246 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_mask_fill_swap_dims() { + let device = Default::default(); + let tensor_1 = TestTensorInt::arange(0..16, &device).float(); + let tensor_1 = tensor_1.reshape([2, 2, 4]); + let tensor_1 = tensor_1.swap_dims(0, 2); + + let mask = tensor_1.clone().lower_equal_elem(5.0); + let output = tensor_1.clone().mask_fill(mask, -5.0); + + let expected = TensorData::from([ + [[-5.0, 8.0], [-5.0, 12.0]], + [[-5.0, 9.0], [-5.0, 13.0]], + [[-5.0, 10.0], [6.0, 14.0]], + [[-5.0, 11.0], [7.0, 15.0]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_mask_where_ops() { + let device = Default::default(); + let tensor = TestTensor::from_data([[1.0, 7.0], [2.0, 3.0]], &device); + let mask = + TestTensorBool::<2>::from_bool(TensorData::from([[true, false], [false, true]]), &device); + let value = TestTensor::<2>::from_data(TensorData::from([[1.8, 2.8], [3.8, 4.8]]), &device); + + let output = tensor.mask_where(mask, value); + let expected = TensorData::from([[1.8, 7.0], [2.0, 4.8]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_mask_where_broadcast() { + let device = Default::default(); + // When broadcasted, the input [[2, 3], [4, 5]] is repeated 4 times + let tensor = TestTensorInt::<1>::arange(2..6, &device).reshape([1, 2, 2]); + let mask = TestTensorBool::<3>::from_bool( + TensorData::from([ + [[true, false], [false, true]], + [[false, true], [true, false]], + [[false, false], [false, false]], + [[true, true], [true, true]], + ]), + &device, + ); + let value = TestTensor::<3>::ones([4, 2, 2], &device); + + let output = tensor.float().mask_where(mask, value); + let expected = TensorData::from([ + [[1., 3.], [4., 1.]], + [[2., 1.], [1., 5.]], + [[2., 3.], [4., 5.]], + [[1., 1.], [1., 1.]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_mask_where_broadcast_value_small() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(2..4, &device).float(); + let mask = TestTensorBool::<1>::from_bool(TensorData::from([true, false]), &device); + let value = TestTensor::<1>::ones([1], &device); + + let output = tensor.mask_where(mask, value); + let expected = TensorData::from([1., 3.]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_handle_mask_where_nans() { + let device = Default::default(); + let tensor = TestTensor::from_data( + [ + [f32::NAN, f32::NAN, f32::NAN], + [f32::NAN, f32::NAN, f32::NAN], + [f32::NAN, f32::NAN, f32::NAN], + ], + &device, + ); + let mask = TestTensorBool::<2>::from_bool( + TensorData::from([ + [true, true, true], + [true, true, false], + [false, false, false], + ]), + &device, + ); + let value = TestTensor::<2>::from_data( + TensorData::from([[0.9, 0.8, 0.7], [0.6, 0.5, 0.4], [0.3, 0.2, 0.1]]), + &device, + ); + + let output = tensor.mask_where(mask, value); + let expected = TensorData::from([ + [0.9, 0.8, 0.7], + [0.6, 0.5, f32::NAN], + [f32::NAN, f32::NAN, f32::NAN], + ]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_mask_fill_ops() { + let device = Default::default(); + let tensor = TestTensor::from_data([[1.0, 7.0], [2.0, 3.0]], &device); + let mask = + TestTensorBool::<2>::from_bool(TensorData::from([[true, false], [false, true]]), &device); + + let output = tensor.mask_fill(mask, 2.0); + let expected = TensorData::from([[2.0, 7.0], [2.0, 2.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_mask_fill_broadcasted() { + let device = Default::default(); + let tensor = TestTensor::zeros([1, 4, 2, 2], &device); + let mask = TestTensorBool::<4>::from_bool( + TensorData::from([[[[true, false], [false, true]]]]), + &device, + ); + + let output = tensor.mask_fill(mask, 2.0); + let expected = TensorData::from([[ + [[2., 0.], [0., 2.]], + [[2., 0.], [0., 2.]], + [[2., 0.], [0., 2.]], + [[2., 0.], [0., 2.]], + ]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn float_mask_fill_infinite() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data( + [ + [f32::NEG_INFINITY, f32::NEG_INFINITY], + [f32::NEG_INFINITY, f32::NEG_INFINITY], + ], + &device, + ); + let mask = + TestTensorBool::<2>::from_bool(TensorData::from([[true, false], [false, true]]), &device); + + let output = tensor.mask_fill(mask, 10.0f32); + let expected = TensorData::from([[10f32, f32::NEG_INFINITY], [f32::NEG_INFINITY, 10f32]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_mask_fill_transposed_tensor() { + // [[1, 2], [3, 4]] transposed -> [[1, 3], [2, 4]]; mask [[T, F], [F, T]] -> [[0, 3], [2, 0]] + let tensor = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]).transpose(); + let mask = TestTensorBool::<2>::from([[true, false], [false, true]]); + + let output = tensor.mask_fill(mask, 0.0); + + output + .into_data() + .assert_eq(&TensorData::from([[0.0, 3.0], [2.0, 0.0]]), false); +} + +#[test] +fn should_support_mask_fill_flipped_tensor() { + // [1,2,3,4] flipped -> [4,3,2,1]; mask [T,T,F,F] -> [0,0,2,1] + let tensor = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0]).flip([0]); + let mask = TestTensorBool::<1>::from([true, true, false, false]); + + let output = tensor.mask_fill(mask, 0.0); + + output + .into_data() + .assert_eq(&TensorData::from([0.0, 0.0, 2.0, 1.0]), false); +} + +#[test] +fn should_support_mask_fill_flipped_mask() { + // tensor [1,2,3,4]; mask [F,F,T,T] flipped -> [T,T,F,F] -> [0,0,3,4] + let tensor = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0]); + let mask = TestTensorBool::<1>::from([false, false, true, true]).flip([0]); + + let output = tensor.mask_fill(mask, 0.0); + + output + .into_data() + .assert_eq(&TensorData::from([0.0, 0.0, 3.0, 4.0]), false); +} + +#[test] +fn should_support_mask_where_flipped_2d() { + // [[1,2],[3,4]] axis-0 flipped -> [[3,4],[1,2]]; mask [[T,F],[F,T]]; + // value [[10,20],[30,40]] -> [[10,4],[1,40]] + let tensor = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]).flip([0]); + let mask = TestTensorBool::<2>::from([[true, false], [false, true]]); + let value = TestTensor::<2>::from([[10.0, 20.0], [30.0, 40.0]]); + + let output = tensor.mask_where(mask, value); + + output + .into_data() + .assert_eq(&TensorData::from([[10.0, 4.0], [1.0, 40.0]]), false); +} + +#[test] +fn should_support_mask_fill_both_flipped() { + // tensor [1,2,3,4] flipped -> [4,3,2,1]; mask [T,F,T,F] flipped -> [F,T,F,T] + // result: [4, 0, 2, 0] + let tensor = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0]).flip([0]); + let mask = TestTensorBool::<1>::from([true, false, true, false]).flip([0]); + + let output = tensor.mask_fill(mask, 0.0); + + output + .into_data() + .assert_eq(&TensorData::from([4.0, 0.0, 2.0, 0.0]), false); +} + +#[test] +fn should_support_mask_fill_narrowed_tensor() { + // [1,2,3,4,5,6] narrowed to [2,3,4,5]; mask [T,F,F,T] -> [0,3,4,0] + let tensor = TestTensor::<1>::from([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).narrow(0, 1, 4); + let mask = TestTensorBool::<1>::from([true, false, false, true]); + + let output = tensor.mask_fill(mask, 0.0); + + output + .into_data() + .assert_eq(&TensorData::from([0.0, 3.0, 4.0, 0.0]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/matmul.rs b/crates/burn-backend-tests/tests/tensor/float/ops/matmul.rs new file mode 100644 index 0000000..a31f14c --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/matmul.rs @@ -0,0 +1,466 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::{ElementConversion, Tolerance}; + +#[test] +fn test_float_matmul_d2() { + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data([[1.0, 7.0], [2.0, 3.0], [1.0, 5.0]], &device); + let tensor_2 = TestTensor::from_data([[4.0, 7.0, 5.0], [2.0, 3.0, 5.0]], &device); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([[18.0, 28.0, 40.0], [14.0, 23.0, 25.0], [14.0, 22.0, 30.0]]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_float_matmul_d3() { + let device = Default::default(); + let tensor_1 = TestTensor::<3>::from_data([[[1.0, 7.0], [2.0, 3.0]]], &device); + let tensor_2 = TestTensor::from_data([[[4.0, 7.0], [2.0, 3.0]]], &device); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([[[18.0, 28.0], [14.0, 23.0]]]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_float_matmul_broadcast_1() { + let device = Default::default(); + let tensor_1 = TestTensor::<3>::from_data([[[1.0, 7.0], [2.0, 3.0]]], &device); + let tensor_2 = TestTensor::from_data( + [[[4.0, 7.0], [2.0, 3.0]], [[2.0, 5.0], [6.0, 3.0]]], + &device, + ); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([[[18.0, 28.0], [14.0, 23.0]], [[44.0, 26.0], [22.0, 19.0]]]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_float_matmul_broadcast_4d() { + let device = Default::default(); + // [2, 1, 2, 2] + let tensor_1 = TestTensor::<4>::from_data( + [[[[1.0, 7.0], [2.0, 3.0]]], [[[2.0, 5.0], [6.0, 3.0]]]], + &device, + ); + // [1, 2, 2, 2] + let tensor_2 = TestTensor::from_data( + [[[[9.0, 8.0], [1.0, 4.0]], [[2.0, 7.0], [3.0, 5.0]]]], + &device, + ); + + // [2, 1, 2, 2] @ [1, 2, 2, 2] -> [2, 2, 2, 2] + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([ + [[[16.0, 36.0], [21.0, 28.0]], [[23.0, 42.0], [13.0, 29.0]]], + [[[23.0, 36.0], [57.0, 60.0]], [[19.0, 39.0], [21.0, 57.0]]], + ]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_float_matmul_simple_1() { + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data([[5.0, 14.0], [14.0, 50.0]], &device); + let tensor_2 = TestTensor::from_data([[3.0, 4.0, 5.0], [0.0, 1.0, 2.0]], &device); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([[15.0, 34.0, 53.0], [42.0, 106.0, 170.0]]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_float_matmul_4_3() { + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data( + [[0., 1., 2., 3.], [4., 5., 6., 7.], [8., 9., 10., 11.]], + &device, + ); + let tensor_2 = TestTensor::from_data( + [[0., 1., 2.], [4., 5., 6.], [8., 9., 10.], [12., 13., 14.]], + &device, + ); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([[56., 62., 68.], [152., 174., 196.], [248., 286., 324.]]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_float_matmul_batch_vec_mat() { + let device = Default::default(); + + // [..., B, 1, K] = [3, 1, 2] + let tensor_1 = TestTensor::<3>::from_data([[[1.0, 7.0]], [[2.0, 3.0]], [[1.0, 5.0]]], &device); + + // [..., 1, K, N] = [1, 2, 3] + let tensor_2 = TestTensor::<3>::from_data([[[4.0, 7.0, 5.0], [2.0, 3.0, 5.0]]], &device); + + let tensor_3 = tensor_1.matmul(tensor_2); + + // [..., B, 1, N] = [3, 1, 3] + let expected = TensorData::from([ + [[18.0, 28.0, 40.0]], + [[14.0, 23.0, 25.0]], + [[14.0, 22.0, 30.0]], + ]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_float_matmul_vecmat() { + let device = Default::default(); + + // [..., B, 1, K] = [3, 1, 2] + let tensor_1 = TestTensor::<3>::from_data([[[1.0, 7.0]], [[2.0, 3.0]], [[1.0, 5.0]]], &device); + + // [..., B, K, N] = [3, 2, 3] + let tensor_2 = TestTensor::<3>::from_data( + [ + [[1.0, 4.0, 5.0], [3.0, 5.0, 6.0]], + [[8.0, 2.0, 3.0], [0.0, 2.0, 4.0]], + [[4.0, 7.0, 5.0], [2.0, 3.0, 5.0]], + ], + &device, + ); + + let tensor_3 = tensor_1.matmul(tensor_2); + + // [..., B, 1, N] = [3, 1, 3] + let expected = TensorData::from([ + [[22.0, 39.0, 47.0]], + [[16.0, 10.0, 18.0]], + [[14.0, 22.0, 30.0]], + ]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_float_matmul_trivial() { + let device = Default::default(); + + let tensor_1 = TestTensorInt::<1>::arange(0..16, &device) + .reshape([4, 4]) + .float(); + + let tensor_3 = tensor_1.clone().matmul(tensor_1); + + tensor_3.into_data().assert_approx_eq::( + &TensorData::from([ + [56., 62., 68., 74.], + [152., 174., 196., 218.], + [248., 286., 324., 362.], + [344., 398., 452., 506.], + ]), + Tolerance::default(), + ); +} + +#[test] +fn test_float_matmul_trivial_transposed() { + let device = Default::default(); + + let tensor_1 = TestTensorInt::<1>::arange(0..16, &device) + .reshape([4, 4]) + .float(); + + let tensor_3 = tensor_1.clone().matmul(tensor_1.transpose()); + + tensor_3.into_data().assert_approx_eq::( + &TensorData::from([ + [14., 38., 62., 86.], + [38., 126., 214., 302.], + [62., 214., 366., 518.], + [86., 302., 518., 734.], + ]), + Tolerance::default(), + ); +} + +/// Regression test for batch bug in fused matmul +#[test] +fn test_float_matmul_vecmat_transposed_fused() { + let device = Default::default(); + + let batch1 = 1; + let batch2 = 2; + let batch = batch1 * batch2; + let seq_length = 3; + let d_model = 32; + + // Guard int arange limits + #[allow(clippy::unnecessary_cast)] + if (IntElem::MAX as i64) < seq_length * d_model * batch { + return; + } + if FloatElem::MAX.elem::() < 269493.0 { + return; + } + + let weight: TestTensor<4> = TestTensorInt::arange(0..d_model * batch, &device) + .reshape([batch1, batch2, 1, d_model]) + .float(); + let signal: TestTensor<4> = TestTensorInt::arange(0..seq_length * d_model * batch, &device) + .reshape([batch1, batch2, seq_length, d_model]) + .float(); + + device.sync().unwrap(); + let weight = weight.transpose(); + let out = signal.matmul(weight) + 5; + let expected = TensorData::from([[ + [[10421.0], [26293.0], [42165.0]], + [[172213.0], [220853.0], [269493.0]], + ]]); + expected.assert_approx_eq(&out.into_data(), Tolerance::::strict()); +} + +#[test] +fn test_float_matmul_4_8() { + let device = Default::default(); + + let tensor_1 = TestTensorInt::<1>::arange(0..32, &device) + .reshape([4, 8]) + .float(); + + let tensor_3 = tensor_1.clone().matmul(tensor_1.transpose()); + + tensor_3.into_data().assert_approx_eq::( + &TensorData::from([ + [140., 364., 588., 812.], + [364., 1100., 1836., 2572.], + [588., 1836., 3084., 4332.], + [812., 2572., 4332., 6092.], + ]), + Tolerance::default(), + ); +} + +#[test] +fn test_float_matmul_simple_2() { + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data([[1.0, 2.0, 3.0, 4.0]], &device); + let tensor_2 = TestTensor::from_data([[3.0], [4.0], [5.0], [6.0]], &device); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([[50.0]]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_float_matmul_simple_3() { + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data( + [[3., 3., 3.], [4., 4., 4.], [5., 5., 5.], [6., 6., 6.]], + &device, + ); + let tensor_2 = TestTensor::from_data( + [[1., 2., 3., 4.], [1., 2., 3., 4.], [1., 2., 3., 4.]], + &device, + ); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([ + [9., 18., 27., 36.], + [12., 24., 36., 48.], + [15., 30., 45., 60.], + [18., 36., 54., 72.], + ]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_matmul_transposed_lhs() { + // [2, 3] -> transpose -> [3, 2], matmul with identity [2, 2] -> [3, 2]. + let device = Default::default(); + let lhs = TestTensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + let rhs = TestTensor::<2>::from_data([[1.0, 0.0], [0.0, 1.0]], &device); + + let output = lhs.transpose().matmul(rhs); + let expected = TensorData::from([[1.0, 4.0], [2.0, 5.0], [3.0, 6.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_matmul_transposed_rhs() { + // identity [2, 2] matmul [3, 2].transpose() -> [2, 3]. + let device = Default::default(); + let lhs = TestTensor::<2>::from_data([[1.0, 0.0], [0.0, 1.0]], &device); + let rhs = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], &device); + + let output = lhs.matmul(rhs.transpose()); + let expected = TensorData::from([[1.0, 3.0, 5.0], [2.0, 4.0, 6.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_matmul_both_transposed() { + // [2, 3].T [3, 2] matmul [3, 2].T [2, 3] -> [3, 3]. + let device = Default::default(); + let lhs = TestTensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + let rhs = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], &device); + + let output = lhs.transpose().matmul(rhs.transpose()); + // lhs.T = [[1,4],[2,5],[3,6]]; rhs.T = [[1,3,5],[2,4,6]] + // row0: 1+8,3+16,5+24 = 9,19,29 + // row1: 2+10,6+20,10+30 = 12,26,40 + // row2: 3+12,9+24,15+36 = 15,33,51 + let expected = TensorData::from([[9.0, 19.0, 29.0], [12.0, 26.0, 40.0], [15.0, 33.0, 51.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_matmul_batched_transposed_rhs() { + // QK^T attention pattern: q.matmul(k.swap_dims(1, 2)). + let device = Default::default(); + let q = TestTensor::<3>::from_data( + [ + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]], + ], + &device, + ); + let k = TestTensor::<3>::from_data( + [ + [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + [[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]], + ], + &device, + ); + + let output = q.matmul(k.swap_dims(1, 2)); + // batch 0: [[1,2,3],[4,5,6]] @ [[1,0],[0,1],[0,0]] = [[1,2],[4,5]] + // batch 1: [[7,8,9],[10,11,12]] @ [[1,2],[1,2],[1,2]] = [[24,48],[33,66]] + let expected = TensorData::from([[[1.0, 2.0], [4.0, 5.0]], [[24.0, 48.0], [33.0, 66.0]]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_matmul_batched_transposed_lhs() { + // [2, 2, 3].swap_dims(1, 2) -> [2, 3, 2]; rhs [2, 2, 2] -> [2, 3, 2]. + let device = Default::default(); + let a = TestTensor::<3>::from_data( + [ + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]], + ], + &device, + ); + let b = TestTensor::<3>::from_data( + [[[1.0, 0.0], [0.0, 1.0]], [[2.0, 0.0], [0.0, 2.0]]], + &device, + ); + + let output = a.swap_dims(1, 2).matmul(b); + // batch 0: a.T = [[1,4],[2,5],[3,6]] @ I = [[1,4],[2,5],[3,6]] + // batch 1: a.T = [[7,10],[8,11],[9,12]] @ 2I = [[14,20],[16,22],[18,24]] + let expected = TensorData::from([ + [[1.0, 4.0], [2.0, 5.0], [3.0, 6.0]], + [[14.0, 20.0], [16.0, 22.0], [18.0, 24.0]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_matmul_batched_both_transposed() { + // a: [2, 3, 2].swap_dims -> [2, 2, 3]; b: [2, 2, 3].swap_dims -> [2, 3, 2]; + // result: [2, 2, 2]. + let device = Default::default(); + let a = TestTensor::<3>::from_data( + [ + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], + [[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]], + ], + &device, + ); + let b = TestTensor::<3>::from_data( + [ + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]], + ], + &device, + ); + + let output = a.swap_dims(1, 2).matmul(b.swap_dims(1, 2)); + // batch 0: a.T = [[1,3,5],[2,4,6]]; b.T = [[1,4],[2,5],[3,6]] + // row0: 1+6+15=22, 4+15+30=49 + // row1: 2+8+18=28, 8+20+36=64 + // batch 1: a.T = [[7,9,11],[8,10,12]]; b.T = [[7,10],[8,11],[9,12]] + // row0: 49+72+99=220, 70+99+132=301 + // row1: 56+80+108=244, 80+110+144=334 + let expected = TensorData::from([ + [[22.0, 49.0], [28.0, 64.0]], + [[220.0, 301.0], [244.0, 334.0]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_matmul_batched_broadcast_transposed() { + // lhs [1, 2, 3].swap_dims(1, 2) -> [1, 3, 2] broadcasts against rhs [4, 2, 2]. + let device = Default::default(); + let lhs = TestTensor::<3>::from_data([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]], &device); + let rhs = TestTensor::<3>::from_data( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[2.0, 0.0], [0.0, 2.0]], + [[1.0, 1.0], [1.0, 1.0]], + [[0.0, 1.0], [1.0, 0.0]], + ], + &device, + ); + + let output = lhs.swap_dims(1, 2).matmul(rhs); + // lhs.T = [[1,4],[2,5],[3,6]] (shape [1, 3, 2]) broadcast to [4, 3, 2] + // batch 0 (I): [[1,4],[2,5],[3,6]] + // batch 1 (2I): [[2,8],[4,10],[6,12]] + // batch 2 (ones): each output row is [sum, sum] where sum is the row of lhs.T + // batch 3 (swap): [[4,1],[5,2],[6,3]] + let expected = TensorData::from([ + [[1.0, 4.0], [2.0, 5.0], [3.0, 6.0]], + [[2.0, 8.0], [4.0, 10.0], [6.0, 12.0]], + [[5.0, 5.0], [7.0, 7.0], [9.0, 9.0]], + [[4.0, 1.0], [5.0, 2.0], [6.0, 3.0]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +#[should_panic] +fn float_should_panic_when_inner_dimensions_are_not_equal() { + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data([[3., 3.], [4., 4.], [5., 5.], [6., 6.]], &device); + let tensor_2 = TestTensor::from_data( + [[1., 2., 3., 4.], [1., 2., 3., 4.], [1., 2., 3., 4.]], + &device, + ); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([ + [9., 18., 27., 36.], + [12., 24., 36., 48.], + [15., 30., 45., 60.], + [18., 36., 54., 72.], + ]); + + tensor_3.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/maxmin.rs b/crates/burn-backend-tests/tests/tensor/float/ops/maxmin.rs new file mode 100644 index 0000000..5e11eeb --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/maxmin.rs @@ -0,0 +1,309 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_max_dim_2d() { + let f = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &Default::default()); + + f.clone() + .max_dim(0) + .into_data() + .assert_eq(&TensorData::from([[3., 4., 5.]]), false); + + f.clone() + .max_dim(1) + .into_data() + .assert_eq(&TensorData::from([[2.], [5.]]), false); + + // Negative Index + f.clone() + .max_dim(-1) + .into_data() + .assert_eq(&TensorData::from([[2.], [5.]]), false); + + // Regression Test: https://github.com/tracel-ai/burn/issues/3139 + let z = f.clone().int(); + z.clone() + .max_dim(0) + .into_data() + .assert_eq(&TensorData::from([[3, 4, 5]]), false); + z.clone() + .max_dim(1) + .into_data() + .assert_eq(&TensorData::from([[2], [5]]), false); +} + +#[test] +fn test_max_dims_2d() { + let f = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &Default::default()); + + f.clone() + .max_dims(&[0]) + .into_data() + .assert_eq(&TensorData::from([[3., 4., 5.]]), false); + + f.clone() + .max_dims(&[-2]) + .into_data() + .assert_eq(&TensorData::from([[3., 4., 5.]]), false); + + f.clone() + .max_dims(&[0, 1]) + .into_data() + .assert_eq(&TensorData::from([[5.]]), false); +} + +#[test] +fn test_max_dim_with_indices_2d_with_dim_0th() { + let tensor = + TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &Default::default()); + + // Positive, Negative Index + for idx in [0, -2] { + let (output, index) = tensor.clone().max_dim_with_indices(idx); + + let output_expected = TensorData::from([[3., 4., 5.]]); + let index_expected = TensorData::from([[1, 1, 1]]); + + output.into_data().assert_eq(&output_expected, false); + index.into_data().assert_eq(&index_expected, false); + } +} + +#[test] +fn test_max_dim_with_indices_2d() { + let tensor = + TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &Default::default()); + + let (output, index) = tensor.max_dim_with_indices(1); + + let output_expected = TensorData::from([[2.], [5.]]); + let index_expected = TensorData::from([[2], [2]]); + + output.into_data().assert_eq(&output_expected, false); + index.into_data().assert_eq(&index_expected, false); +} + +#[test] +fn test_max_dim_2d_with_0th_dim() { + let tensor = + TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &Default::default()); + + let output = tensor.max_dim(0); + let expected = TensorData::from([[3., 4., 5.]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_max_pair() { + let a = TestTensor::<1>::from_data([1.0, 2.0, 3.0, 4.0], &Default::default()); + let b = TestTensor::from_data([2.0, 1.0, 4.0, 5.0], &Default::default()); + + let output = a.max_pair(b); + let expected = TensorData::from([2.0, 2.0, 4.0, 5.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_min_dim_2d() { + let f = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &Default::default()); + + f.clone() + .min_dim(0) + .into_data() + .assert_eq(&TensorData::from([[0., 1., 2.]]), false); + + f.clone() + .min_dim(1) + .into_data() + .assert_eq(&TensorData::from([[0.], [3.]]), false); + + // Negative Index + f.clone() + .min_dim(-1) + .into_data() + .assert_eq(&TensorData::from([[0.], [3.]]), false); + + // Regression Test: https://github.com/tracel-ai/burn/issues/3139 + let z = f.int(); + z.clone() + .min_dim(0) + .into_data() + .assert_eq(&TensorData::from([[0, 1, 2]]), false); + z.clone() + .min_dim(1) + .into_data() + .assert_eq(&TensorData::from([[0], [3]]), false); +} + +#[test] +fn test_min_dims_2d() { + let f = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &Default::default()); + + f.clone() + .min_dims(&[0]) + .into_data() + .assert_eq(&TensorData::from([[0., 1., 2.]]), false); + + f.clone() + .min_dims(&[-2]) + .into_data() + .assert_eq(&TensorData::from([[0., 1., 2.]]), false); + + f.clone() + .min_dims(&[0, 1]) + .into_data() + .assert_eq(&TensorData::from([[0.]]), false); +} + +#[test] +fn test_min_dim_with_indices_2d() { + let tensor = + TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &Default::default()); + + let (output, index) = tensor.min_dim_with_indices(1); + + let output_expected = TensorData::from([[0.], [3.]]); + let index_expected = TensorData::from([[0], [0]]); + + output.into_data().assert_eq(&output_expected, false); + index.into_data().assert_eq(&index_expected, false); +} + +#[test] +fn test_min_dim_2d_with_0th_dim() { + let tensor = + TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &Default::default()); + + let output = tensor.min_dim(0); + let expected = TensorData::from([[0., 1., 2.]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_min_dim_with_indices_2d_with_0th_dim() { + let tensor = + TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &Default::default()); + + // Positive, Negative Index + for idx in [0, -2] { + let (output, index) = tensor.clone().min_dim_with_indices(idx); + + let output_expected = TensorData::from([[0., 1., 2.]]); + let index_expected = TensorData::from([[0, 0, 0]]); + + output.into_data().assert_eq(&output_expected, false); + index.into_data().assert_eq(&index_expected, false); + } +} + +#[test] +fn test_min_pair() { + let a = TestTensor::<1>::from_data([1.0, 2.0, 3.0, 4.0], &Default::default()); + let b = TestTensor::from_data([2.0, 1.0, 4.0, 5.0], &Default::default()); + + let output = a.min_pair(b); + let expected = TensorData::from([1.0, 1.0, 3.0, 4.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_max_abs() { + let tensor = TestTensor::<2>::from_data([[0., 1., -2.], [-5., 6., 1.]], &Default::default()); + + let output = tensor.max_abs(); + let expected = TensorData::from([6.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_max_abs_dim_2d_dim_0() { + let tensor = TestTensor::<2>::from_data([[0., 1., -2.], [-5., 6., 1.]], &Default::default()); + + let output = tensor.clone().max_abs_dim(0); + let expected = TensorData::from([[5., 6., 2.]]); + output.into_data().assert_eq(&expected, false); + + // Negative Index + let output = tensor.clone().max_abs_dim(-2); + let expected = TensorData::from([[5., 6., 2.]]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_max_abs_dims_2d() { + let tensor = TestTensor::<2>::from_data([[0., 1., -2.], [-5., 6., 1.]], &Default::default()); + + tensor + .clone() + .max_abs_dims(&[0]) + .into_data() + .assert_eq(&TensorData::from([[5., 6., 2.]]), false); + + tensor + .clone() + .max_abs_dims(&[-2]) + .into_data() + .assert_eq(&TensorData::from([[5., 6., 2.]]), false); + + tensor + .clone() + .max_abs_dims(&[0, 1]) + .into_data() + .assert_eq(&TensorData::from([[6.]]), false); +} + +#[test] +fn test_max_abs_dim_2d_dim_1() { + let tensor = TestTensor::<2>::from_data([[0., 1., -2.], [-5., 6., 1.]], &Default::default()); + + let output = tensor.max_abs_dim(1); + let expected = TensorData::from([[2.], [6.]]); + + output.into_data().assert_eq(&expected, false); +} + +// NaN-propagation tests below. Only run when the `flex` backend feature +// is active, because flex is the only burn backend that currently +// propagates NaN from min/max (matching PyTorch/NumPy/JAX/TF). ndarray +// and the cubecl backends follow IEEE 754 min/max and drop NaN. The +// positive-gate form (rather than excluding specific backends) is used +// because the default-feature CI build selects a backend transitively +// without setting any of its identifying feature flags on +// burn-backend-tests, so a negative gate would still run the test on a +// NaN-dropping backend. See issue #4814. +#[cfg(feature = "flex")] +#[test] +fn test_max_dim_nan_propagation() { + let tensor = TestTensor::<2>::from([[1.0, f32::NAN, 3.0]]); + let data = tensor.max_dim(1).into_data(); + let values = data.as_slice::().unwrap(); + assert!(values[0].is_nan()); +} + +#[cfg(feature = "flex")] +#[test] +fn test_min_dim_nan_propagation() { + let tensor = TestTensor::<2>::from([[1.0, f32::NAN, 3.0]]); + let data = tensor.min_dim(1).into_data(); + let values = data.as_slice::().unwrap(); + assert!(values[0].is_nan()); +} + +#[cfg(feature = "flex")] +#[test] +fn test_max_dim_with_indices_nan_propagation() { + let tensor = TestTensor::<2>::from([[1.0, f32::NAN, 3.0]]); + let (values, indices) = tensor.max_dim_with_indices(1); + let vdata = values.into_data(); + let slice = vdata.as_slice::().unwrap(); + assert!(slice[0].is_nan()); + indices + .into_data() + .assert_eq(&TensorData::from([[1]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/mod.rs b/crates/burn-backend-tests/tests/tensor/float/ops/mod.rs new file mode 100644 index 0000000..61dd784 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/mod.rs @@ -0,0 +1,134 @@ +use super::*; + +/// Build a `[rows, cols]` row-major boolean mask whose rows exercise every +/// logical-reduction merge case along the last axis: all-zero, all-nonzero, a +/// single nonzero at the first/last/middle column, and an all-nonzero row with +/// a single zero. Used by the large `any_dim`/`all_dim` tests (`rows == 8`). +pub fn build_logical_mask(rows: usize, cols: usize) -> Vec { + assert!( + rows >= 8 && cols > 100, + "mask layout assumes rows>=8, cols>100" + ); + let mut mask = vec![false; rows * cols]; + for c in 0..cols { + mask[cols + c] = true; // row 1: all nonzero + mask[5 * cols + c] = true; // row 5: all nonzero + mask[7 * cols + c] = true; // row 7: all nonzero (one zero punched below) + } + mask[2 * cols] = true; // row 2: single nonzero at the first column + mask[3 * cols + (cols - 1)] = true; // row 3: single nonzero at the last column + mask[6 * cols + cols / 2] = true; // row 6: single nonzero in the middle + mask[7 * cols + 100] = false; // row 7: a single zero deep inside an all-nonzero row + // rows 0 and 4 stay all-zero + mask +} + +/// Map the mask to float input, using alternating nonzero magnitudes so the +/// reduction must normalize arbitrary nonzero values (not just `1.0`) to true. +pub fn mask_to_floats(mask: &[bool]) -> Vec { + mask.iter() + .enumerate() + .map(|(i, &m)| { + if m { + if i % 2 == 0 { 2.0 } else { -3.0 } + } else { + 0.0 + } + }) + .collect() +} + +/// Map the mask to int input with alternating nonzero magnitudes. +pub fn mask_to_ints(mask: &[bool]) -> Vec { + mask.iter() + .enumerate() + .map(|(i, &m)| { + if m { + if i % 2 == 0 { 2 } else { -3 } + } else { + 0 + } + }) + .collect() +} + +mod abs; +mod add; +mod aggregation; +mod all; +mod any; +mod arg; +mod blackman_window; +mod cast; +mod cat; +mod categorical; +mod ceil; +mod chunk; +mod clamp; +mod close; +mod comparison; +mod create_like; +mod cross; +mod cumulative; +mod div; +mod dot; +mod erf; +mod exp; +mod expand; +mod finite; +mod flatten; +mod flip; +mod floor; +mod fmod; +mod full; +mod gather_scatter; +mod gather_scatter_nd; +mod grid_sample; +mod hamming_window; +mod hann_window; +mod hypot; +mod inf; +mod init; +mod iter_dim; +mod log; +mod log1p; +mod mask; +mod matmul; +mod maxmin; +mod movedim; +mod mul; +mod nan; +mod narrow; +mod neg; +mod one_hot; +mod padding; +mod permute; +mod powf; +mod powf_scalar; +mod prod; +mod random; +mod recip; +mod remainder; +mod repeat; +mod repeat_dim; +mod reshape; +mod round; +mod select; +mod sign; +mod slice; +mod slice_assign; +mod sort_argsort; +mod split; +mod sqrt; +mod square; +mod squeeze; +mod stack; +mod sub; +mod take; +mod topk; +mod transaction; +mod transpose; +mod tri; +mod trig; +mod trunc; +mod unfold; diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/movedim.rs b/crates/burn-backend-tests/tests/tensor/float/ops/movedim.rs new file mode 100644 index 0000000..49ee98a --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/movedim.rs @@ -0,0 +1,123 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn movedim_float() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device) + .reshape([2, 3, 4]) + .float(); + + let permuted = tensor.clone().movedim(0, 2); + // from pytorch: + // import torch; torch.arange(0, 24).reshape(2, 3, 4).movedim(0, 2).float() + let expected = TensorData::from([ + [[0., 12.], [1., 13.], [2., 14.], [3., 15.]], + [[4., 16.], [5., 17.], [6., 18.], [7., 19.]], + [[8., 20.], [9., 21.], [10., 22.], [11., 23.]], + ]); + + permuted.into_data().assert_eq(&expected, false); + + // Test with negative axis + let permuted = tensor.clone().movedim(0, -1); + permuted.into_data().assert_eq(&expected, false); + + // Test with the same axis + let permuted = tensor.clone().movedim(0, 0); + permuted.into_data().assert_eq(&tensor.into_data(), true); +} + +#[test] +fn vec_input_float() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device) + .reshape([2, 3, 4]) + .float(); + + let permuted = tensor.clone().movedim(vec![0, 1], vec![1, 0]); + // from pytorch + // import torch; torch.arange(0, 24).reshape(2, 3, 4).movedim([0, 1], [1, 0]).float() + let expected = TensorData::from([ + [[0., 1., 2., 3.], [12., 13., 14., 15.]], + [[4., 5., 6., 7.], [16., 17., 18., 19.]], + [[8., 9., 10., 11.], [20., 21., 22., 23.]], + ]); + + permuted.into_data().assert_eq(&expected, false); + + // Test with negative axes + let permuted = tensor.clone().movedim(vec![-3, -2], vec![-2, -3]); + permuted.into_data().assert_eq(&expected, false); + + // Test with the same axes + let permuted = tensor.clone().movedim(vec![0, 1], vec![0, 1]); + permuted.into_data().assert_eq(&tensor.into_data(), true); +} + +#[test] +fn different_input_types() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device) + .reshape([2, 3, 4]) + .float(); + + let permuted = tensor.clone().movedim(0_usize, 2_i32); + // from pytorch: + // import torch; torch.arange(0, 24).reshape(2, 3, 4).movedim(0, 2).float() + let expected = TensorData::from([ + [[0., 12.], [1., 13.], [2., 14.], [3., 15.]], + [[4., 16.], [5., 17.], [6., 18.], [7., 19.]], + [[8., 20.], [9., 21.], [10., 22.], [11., 23.]], + ]); + + permuted.into_data().assert_eq(&expected, false); + + // Test with negative axis + let permuted = tensor.clone().movedim(0_usize, -1); + permuted.into_data().assert_eq(&expected, false); + + // Test with the same axis + let permuted = tensor.clone().movedim(0_i32, 0_usize); + permuted.into_data().assert_eq(&tensor.into_data(), true); +} + +#[test] +#[should_panic] +fn edge_different_sizes() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device).reshape([2, 3, 4]); + + // Test with a repeated axis + let _ = tensor.clone().movedim(vec![0, 1], vec![0]); +} + +#[test] +#[should_panic] +fn edge_out_of_bound_axis() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device).reshape([2, 3, 4]); + + // Test with an out of bound axis + let _ = tensor.clone().movedim(0, 100); +} + +#[test] +#[should_panic] +fn edge_vec_is_not_a_set() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device).reshape([2, 3, 4]); + + // Test with a repeated axis + let _ = tensor.clone().movedim(vec![0, 1, 1, 1, 1], vec![0, 0, 1]); +} + +#[test] +#[should_panic] +fn edge_out_of_bound_axis_vec() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device).reshape([2, 3, 4]); + + // Test with an out of bound axis + let _ = tensor.clone().movedim(vec![0, 100], vec![0, 1]); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/mul.rs b/crates/burn-backend-tests/tests/tensor/float/ops/mul.rs new file mode 100644 index 0000000..c40a33a --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/mul.rs @@ -0,0 +1,68 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_mul_ops() { + let data_1 = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let data_2 = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device); + + let output = tensor_1 * tensor_2; + let expected = TensorData::from([[0.0, 1.0, 4.0], [9.0, 16.0, 25.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_mul_broadcast() { + let data_1 = TensorData::from([[0.0, 1.0, 2.0]]); + let data_2 = TensorData::from([[3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]); + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device); + + let output = tensor_1 * tensor_2; + let expected = TensorData::from([[0.0, 4.0, 10.0], [0.0, 7.0, 16.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_mul_broadcast_2_dims() { + let device = Default::default(); + let tensor_1 = TestTensor::<1>::from_data([0.0, 1.0, 2.0], &device).reshape([3, 1]); + let tensor_2 = TestTensor::<1>::from_data([3.0, 4.0, 5.0], &device).reshape([1, 3]); + + let output = tensor_1 * tensor_2; + let expected = TensorData::from([[0.0, 0.0, 0.0], [3.0, 4.0, 5.0], [6.0, 8.0, 10.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_mul_both_transposed() { + // Non-contiguous on both sides: each operand is a transpose of a 2x2. + let a = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]).transpose(); + let b = TestTensor::<2>::from([[2.0, 3.0], [4.0, 5.0]]).transpose(); + + let output = a * b; + + // a_t = [[1, 3], [2, 4]], b_t = [[2, 4], [3, 5]] + output + .into_data() + .assert_eq(&TensorData::from([[2.0, 12.0], [6.0, 20.0]]), false); +} + +#[test] +fn should_support_mul_scalar_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let scalar = 2.0; + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor * scalar; + let expected = TensorData::from([[0.0, 2.0, 4.0], [6.0, 8.0, 10.0]]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/nan.rs b/crates/burn-backend-tests/tests/tensor/float/ops/nan.rs new file mode 100644 index 0000000..54a3b6f --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/nan.rs @@ -0,0 +1,34 @@ +use super::*; + +#[test] +fn is_nan() { + let no_nan = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let no_nan_expected = TestTensorBool::<2>::from([[false, false, false], [false, false, false]]); + + let with_nan = TestTensor::<2>::from([[0.0, f32::NAN, 2.0], [f32::NAN, 4.0, 5.0]]); + let with_nan_expected = TestTensorBool::<2>::from([[false, true, false], [true, false, false]]); + + assert_eq!(no_nan_expected.into_data(), no_nan.is_nan().into_data()); + + assert_eq!(with_nan_expected.into_data(), with_nan.is_nan().into_data()); +} + +#[test] +fn contains_nan() { + let no_nan = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + assert!(!no_nan.contains_nan().into_scalar::()); + + let with_nan = TestTensor::<2>::from([[0.0, f32::NAN, 2.0], [3.0, 4.0, 5.0]]); + assert!(with_nan.contains_nan().into_scalar::()); + + // Regression guard: a finite tensor must never be reported as containing NaN. + // The previous `sum().is_nan()` implementation could overflow the reduction + // on narrow-exponent floats (e.g. f16) and turn a finite tensor into a NaN sum, + // producing a false positive. Here, Inf + (-Inf) = NaN. + let device = Default::default(); + let pos = TestTensor::<1>::ones([2048], &device) * 60000.0; + let neg = TestTensor::<1>::ones([2048], &device) * -60000.0; + // Large positive and large negative numbers + let large_finite = TestTensor::<1>::cat(vec![pos, neg], 0); + assert!(!large_finite.contains_nan().into_scalar::()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/narrow.rs b/crates/burn-backend-tests/tests/tensor/float/ops/narrow.rs new file mode 100644 index 0000000..ec47e34 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/narrow.rs @@ -0,0 +1,98 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{Shape, TensorData}; + +#[test] +fn test_narrow_1() { + let tensor = TestTensor::<2>::from_data( + TensorData::from([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]), + &Default::default(), + ); + + let output = tensor.clone().narrow(0, 0, 2); + let expected = TensorData::from([[1., 2., 3.], [4., 5., 6.]]); + + assert_eq!(output.shape(), Shape::from([2, 3])); + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_narrow_2() { + let tensor = TestTensor::<2>::from_data( + TensorData::from([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]), + &Default::default(), + ); + + let output = tensor.clone().narrow(1, 1, 2); + let expected = TensorData::from([[2., 3.], [5., 6.], [8., 9.]]); + assert_eq!(output.shape(), Shape::from([3, 2])); + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_narrow_3() { + let device = &Default::default(); + let shape = Shape::new([8, 8]); + let tensor = TestTensorInt::arange(0..shape.num_elements() as i64, device) + .reshape::<2, _>(shape) + .float(); + + let output = tensor.clone().narrow(0, 3, 4); + let expected = TensorData::from([ + [24.0, 25.0, 26.0, 27.0, 28.0, 29.0, 30.0, 31.0], + [32.0, 33.0, 34.0, 35.0, 36.0, 37.0, 38.0, 39.0], + [40.0, 41.0, 42.0, 43.0, 44.0, 45.0, 46.0, 47.0], + [48.0, 49.0, 50.0, 51.0, 52.0, 53.0, 54.0, 55.0], + ]); + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +#[should_panic] +fn test_narrow_invalid_dim() { + let tensor = TestTensor::<2>::from_data( + TensorData::from([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]), + &Default::default(), + ); + + let _output = tensor.narrow(2, 0, 2); +} + +#[test] +#[should_panic] +fn test_narrow_invalid_start() { + let tensor = TestTensor::<2>::from_data( + TensorData::from([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]), + &Default::default(), + ); + + let _output = tensor.narrow(0, 3, 2); +} + +#[test] +#[should_panic] +fn test_narrow_invalid_zero_length() { + let tensor = TestTensor::<2>::from_data( + TensorData::from([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]), + &Default::default(), + ); + + let _output = tensor.narrow(0, 1, 0); +} + +#[test] +#[should_panic] +fn test_narrow_invalid_length() { + let tensor = TestTensor::<2>::from_data( + TensorData::from([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]]), + &Default::default(), + ); + + let _output = tensor.narrow(0, 0, 4); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/neg.rs b/crates/burn-backend-tests/tests/tensor/float/ops/neg.rs new file mode 100644 index 0000000..01c9aaa --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/neg.rs @@ -0,0 +1,21 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_neg_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.neg(); + let expected = TensorData::from([[-0.0, -1.0, -2.0], [-3.0, -4.0, -5.0]]).convert::(); + + // -0.0 is represented differently than 0.0 so we make sure the values are the same in f32 + assert_eq!( + output + .into_data() + .convert::() + .as_slice::() + .unwrap(), + expected.as_slice::().unwrap() + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/one_hot.rs b/crates/burn-backend-tests/tests/tensor/float/ops/one_hot.rs new file mode 100644 index 0000000..60430ad --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/one_hot.rs @@ -0,0 +1,71 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn float_should_support_one_hot() { + let tensor = TestTensor::<1>::from([0.0, 1.0, 4.0]); + let one_hot_tensor: TestTensor<2> = tensor.one_hot(5); + let expected = TensorData::from([ + [1.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 1.0], + ]); + one_hot_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn float_should_support_one_hot_index() { + let tensor = TestTensor::<1>::from([2.0]); + let one_hot_tensor: TestTensor<2> = tensor.one_hot::<2>(10); + let expected = TensorData::from([[0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]); + one_hot_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +#[should_panic] +fn float_one_hot_should_panic_when_index_exceeds_number_of_classes() { + let tensor = TestTensor::<1>::from([5.0]); + let _result: TestTensor<2> = tensor.one_hot(5); +} + +#[test] +#[should_panic] +fn float_one_hot_should_panic_when_number_of_classes_is_zero() { + let tensor = TestTensor::<1>::from([0.0]); + let _result: TestTensor<2> = tensor.one_hot(0); +} + +#[test] +fn one_hot_fill_with_negative_axis_and_indices() { + let tensor = TestTensor::<2>::from([[0, 2], [1, -1]]); + let expected = TensorData::from([ + [[5.0, 0.0, 0.0], [0.0, 0.0, 5.0]], + [[0.0, 5.0, 0.0], [0.0, 0.0, 5.0]], + ]); + + let one_hot_tensor: TestTensor<3> = tensor.one_hot_fill(3, 5.0, 0.0, -1); + + one_hot_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn one_hot_fill_with_negative_indices() { + let tensor = TestTensor::<1>::from([0.0, -7.0, -8.0]); + let expected = TensorData::from([ + [3.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + [1.0, 1.0, 1.0, 3.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + [1.0, 1.0, 3.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + ]); + + let one_hot_tensor: TestTensor<2> = tensor.one_hot_fill(10, 3.0, 1.0, 1); + + one_hot_tensor.into_data().assert_eq(&expected, false); +} + +#[should_panic] +#[test] +fn one_hot_fill_should_panic_when_axis_out_range_of_rank() { + let tensor = TestTensor::<2>::from([[0.0, 2.0], [1.0, -1.0]]); + + let _one_hot_tensor: TestTensor<3> = tensor.one_hot_fill(2, 5.0, 0.0, 3); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/padding.rs b/crates/burn-backend-tests/tests/tensor/float/ops/padding.rs new file mode 100644 index 0000000..fa2101d --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/padding.rs @@ -0,0 +1,470 @@ +use super::*; +use burn_tensor::{TensorData, ops::PadMode}; + +#[test] +fn padding_constant_2d_test() { + let unpadded_floats: [[f32; 3]; 2] = [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]; + let tensor = TestTensor::<2>::from(unpadded_floats); + + let padded_tensor = tensor.pad((2, 2, 2, 2), 1.1); + + let expected = TensorData::from([ + [1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1], + [1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1], + [1.1, 1.1, 0.0, 1.0, 2.0, 1.1, 1.1], + [1.1, 1.1, 3.0, 4.0, 5.0, 1.1, 1.1], + [1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1], + [1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1], + ]); + padded_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_constant_4d_test() { + let unpadded_floats = [[[[0.0, 1.0], [2.0, 3.0], [4.0, 5.0]]]]; + let tensor = TestTensor::<4>::from(unpadded_floats); + + let padded_tensor = tensor.pad((2, 2, 2, 2), 1.1); + + let expected = TensorData::from([[[ + [1.1, 1.1, 1.1, 1.1, 1.1, 1.1], + [1.1, 1.1, 1.1, 1.1, 1.1, 1.1], + [1.1, 1.1, 0.0, 1.0, 1.1, 1.1], + [1.1, 1.1, 2.0, 3.0, 1.1, 1.1], + [1.1, 1.1, 4.0, 5.0, 1.1, 1.1], + [1.1, 1.1, 1.1, 1.1, 1.1, 1.1], + [1.1, 1.1, 1.1, 1.1, 1.1, 1.1], + ]]]); + padded_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_constant_asymmetric_test() { + let unpadded_floats = [[[[0.0, 1.0], [2.0, 3.0], [4.0, 5.0]]]]; + let tensor = TestTensor::<4>::from(unpadded_floats); + + let padded_tensor = tensor.pad((2, 1, 4, 3), 1.1); + + let expected = TensorData::from([[[ + [1.1, 1.1, 1.1, 1.1, 1.1], + [1.1, 1.1, 1.1, 1.1, 1.1], + [1.1, 1.1, 1.1, 1.1, 1.1], + [1.1, 1.1, 1.1, 1.1, 1.1], + [1.1, 1.1, 0.0, 1.0, 1.1], + [1.1, 1.1, 2.0, 3.0, 1.1], + [1.1, 1.1, 4.0, 5.0, 1.1], + [1.1, 1.1, 1.1, 1.1, 1.1], + [1.1, 1.1, 1.1, 1.1, 1.1], + [1.1, 1.1, 1.1, 1.1, 1.1], + ]]]); + padded_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_reflect_2d_test() { + // Test reflect padding on a 2D tensor + // Input: [[1, 2, 3], [4, 5, 6]] + // With padding (1, 1, 1, 1): + // - Top: reflect row 1 -> [4, 5, 6] + // - Bottom: reflect row 0 -> [1, 2, 3] + // - Left: reflect col 1 + // - Right: reflect col 1 + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + let padded_tensor = tensor.pad((1, 1, 1, 1), PadMode::Reflect); + + // Expected: reflect excludes the edge value + // Before padding height: [[1,2,3], [4,5,6]] + // After top pad (reflect row at index 1): [[4,5,6], [1,2,3], [4,5,6]] + // After bottom pad (reflect row at index 1 from end): [[4,5,6], [1,2,3], [4,5,6], [1,2,3]] + // Then pad width similarly + let expected = TensorData::from([ + [5.0, 4.0, 5.0, 6.0, 5.0], + [2.0, 1.0, 2.0, 3.0, 2.0], + [5.0, 4.0, 5.0, 6.0, 5.0], + [2.0, 1.0, 2.0, 3.0, 2.0], + ]); + padded_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_reflect_width_only_test() { + // Test reflect padding on width dimension only + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0, 4.0]]); + + let padded_tensor = tensor.pad((2, 2, 0, 0), PadMode::Reflect); + + // Input: [1, 2, 3, 4] + // Reflect left 2: take indices [1, 2] = [2, 3], flip = [3, 2] + // Reflect right 2: take indices [1, 2] from end = [2, 3], flip = [3, 2] + // Result: [3, 2, 1, 2, 3, 4, 3, 2] + let expected = TensorData::from([[3.0, 2.0, 1.0, 2.0, 3.0, 4.0, 3.0, 2.0]]); + padded_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_reflect_4d_test() { + // Test reflect padding on 4D tensor (common for images: NCHW) + let tensor = TestTensor::<4>::from([[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]]]); + + let padded_tensor = tensor.pad((1, 1, 1, 1), PadMode::Reflect); + + let expected = TensorData::from([[[ + [5.0, 4.0, 5.0, 6.0, 5.0], + [2.0, 1.0, 2.0, 3.0, 2.0], + [5.0, 4.0, 5.0, 6.0, 5.0], + [8.0, 7.0, 8.0, 9.0, 8.0], + [5.0, 4.0, 5.0, 6.0, 5.0], + ]]]); + padded_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_edge_2d_test() { + // Test edge padding on a 2D tensor + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + let padded_tensor = tensor.pad((1, 1, 1, 1), PadMode::Edge); + + // Edge padding replicates the boundary values + let expected = TensorData::from([ + [1.0, 1.0, 2.0, 3.0, 3.0], + [1.0, 1.0, 2.0, 3.0, 3.0], + [4.0, 4.0, 5.0, 6.0, 6.0], + [4.0, 4.0, 5.0, 6.0, 6.0], + ]); + padded_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_edge_width_only_test() { + // Test edge padding on width dimension only + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0, 4.0]]); + + let padded_tensor = tensor.pad((2, 3, 0, 0), PadMode::Edge); + + // Input: [1, 2, 3, 4] + // Left 2: [1, 1] + // Right 3: [4, 4, 4] + // Result: [1, 1, 1, 2, 3, 4, 4, 4, 4] + let expected = TensorData::from([[1.0, 1.0, 1.0, 2.0, 3.0, 4.0, 4.0, 4.0, 4.0]]); + padded_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_edge_4d_test() { + // Test edge padding on 4D tensor + let tensor = TestTensor::<4>::from([[[[1.0, 2.0], [3.0, 4.0]]]]); + + let padded_tensor = tensor.pad((1, 1, 1, 1), PadMode::Edge); + + let expected = TensorData::from([[[ + [1.0, 1.0, 2.0, 2.0], + [1.0, 1.0, 2.0, 2.0], + [3.0, 3.0, 4.0, 4.0], + [3.0, 3.0, 4.0, 4.0], + ]]]); + padded_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_constant_default_test() { + // Test default PadMode (Constant with 0.0) + let tensor = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]); + + let padded_tensor = tensor.pad((1, 1, 1, 1), PadMode::default()); + + let expected = TensorData::from([ + [0.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 2.0, 0.0], + [0.0, 3.0, 4.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + ]); + padded_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_reflect_max_valid_test() { + // Test reflect padding at maximum valid size (dim_size - 1) + // For a 4-element dimension, max valid padding is 3 + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0, 4.0]]); + + // Padding of 3 on left is valid for width=4 (3 < 4) + let padded_tensor = tensor.pad((3, 3, 0, 0), PadMode::Reflect); + + // Input: [1, 2, 3, 4] + // Reflect left 3: take indices [1, 2, 3] = [2, 3, 4], flip = [4, 3, 2] + // Reflect right 3: take indices [0, 1, 2] = [1, 2, 3], flip = [3, 2, 1] + // Result: [4, 3, 2, 1, 2, 3, 4, 3, 2, 1] + let expected = TensorData::from([[4.0, 3.0, 2.0, 1.0, 2.0, 3.0, 4.0, 3.0, 2.0, 1.0]]); + padded_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_reflect_asymmetric_test() { + // Test asymmetric reflect padding + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]); + + // Asymmetric padding: left=2, right=1, top=1, bottom=2 + let padded_tensor = tensor.pad((2, 1, 1, 2), PadMode::Reflect); + + let expected = TensorData::from([ + [6.0, 5.0, 4.0, 5.0, 6.0, 5.0], + [3.0, 2.0, 1.0, 2.0, 3.0, 2.0], + [6.0, 5.0, 4.0, 5.0, 6.0, 5.0], + [9.0, 8.0, 7.0, 8.0, 9.0, 8.0], + [6.0, 5.0, 4.0, 5.0, 6.0, 5.0], + [3.0, 2.0, 1.0, 2.0, 3.0, 2.0], + ]); + padded_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +#[should_panic(expected = "Reflect padding")] +fn padding_reflect_exceeds_dimension_test() { + // Test that reflect padding panics when padding >= dim_size + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0]]); + + // Padding of 3 on width=3 should panic (3 >= 3, need padding < dim_size) + let _ = tensor.pad((3, 0, 0, 0), PadMode::Reflect); +} + +#[test] +fn padding_edge_asymmetric_test() { + // Test asymmetric edge padding + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + // Asymmetric padding: left=2, right=1, top=3, bottom=1 + let padded_tensor = tensor.pad((2, 1, 3, 1), PadMode::Edge); + + let expected = TensorData::from([ + [1.0, 1.0, 1.0, 2.0, 3.0, 3.0], + [1.0, 1.0, 1.0, 2.0, 3.0, 3.0], + [1.0, 1.0, 1.0, 2.0, 3.0, 3.0], + [1.0, 1.0, 1.0, 2.0, 3.0, 3.0], + [4.0, 4.0, 4.0, 5.0, 6.0, 6.0], + [4.0, 4.0, 4.0, 5.0, 6.0, 6.0], + ]); + padded_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_zero_padding_test() { + // Test that zero padding returns the original tensor unchanged + let tensor = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]); + + let padded_constant = tensor.clone().pad((0, 0, 0, 0), PadMode::Constant(5.0)); + let padded_reflect = tensor.clone().pad((0, 0, 0, 0), PadMode::Reflect); + let padded_edge = tensor.clone().pad((0, 0, 0, 0), PadMode::Edge); + + let expected = TensorData::from([[1.0, 2.0], [3.0, 4.0]]); + padded_constant.into_data().assert_eq(&expected, false); + padded_reflect.into_data().assert_eq(&expected, false); + padded_edge.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_empty_tensor_constant_test() { + // Test constant padding on an empty tensor (zero-sized dimension) + // This should work - creates a tensor filled with the constant value + let tensor: TestTensor<2> = TestTensor::empty([0, 3], &Default::default()); + + // Padding an empty height dimension with constant should create a tensor of just padding + let padded = tensor.pad((0, 0, 2, 2), 1.0); + + // Result should be 4x3 (0 + 2 + 2 = 4 rows) + assert_eq!(padded.dims(), [4, 3]); + + let expected = TensorData::from([ + [1.0, 1.0, 1.0], + [1.0, 1.0, 1.0], + [1.0, 1.0, 1.0], + [1.0, 1.0, 1.0], + ]); + padded.into_data().assert_eq(&expected, false); +} + +#[test] +#[should_panic(expected = "edge padding")] +fn padding_empty_tensor_edge_panics_test() { + // Test that edge padding panics on empty tensor + let tensor: TestTensor<2> = TestTensor::empty([0, 3], &Default::default()); + + // Edge padding on zero-sized dimension should panic + let _ = tensor.pad((0, 0, 1, 1), PadMode::Edge); +} + +#[test] +#[should_panic(expected = "Reflect padding")] +fn padding_empty_tensor_reflect_panics_test() { + // Test that reflect padding panics on empty tensor + let tensor: TestTensor<2> = TestTensor::empty([0, 3], &Default::default()); + + // Reflect padding on zero-sized dimension should panic + let _ = tensor.pad((0, 0, 1, 1), PadMode::Reflect); +} + +// --- Tests for N-dimensional padding using (before, after) pairs --- + +#[test] +fn padding_constant_pairs_2d_test() { + // Same as padding_constant_2d_test but using the new pairs API + let tensor = TestTensor::<2>::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + // [(row_before, row_after), (col_before, col_after)] + let padded_tensor = tensor.pad([(2, 2), (2, 2)], 1.1); + + let expected = TensorData::from([ + [1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1], + [1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1], + [1.1, 1.1, 0.0, 1.0, 2.0, 1.1, 1.1], + [1.1, 1.1, 3.0, 4.0, 5.0, 1.1, 1.1], + [1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1], + [1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1], + ]); + padded_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_constant_single_dim_test() { + // Pad only the last dimension + let tensor = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]); + + let padded_tensor = tensor.pad([(1, 1)], 0.0); + + let expected = TensorData::from([[0.0, 1.0, 2.0, 0.0], [0.0, 3.0, 4.0, 0.0]]); + padded_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_constant_all_dims_4d_test() { + // Pad all 4 dimensions of a 4D tensor (batch, channel, height, width) + // Input: shape [1, 1, 2, 2] + let tensor = TestTensor::<4>::from([[[[1.0, 2.0], [3.0, 4.0]]]]); + + // Pad: batch(1,1), channel(1,1), height(0,0), width(0,0) + let padded = tensor.pad([(1, 1), (1, 1), (0, 0), (0, 0)], 0.0); + + // Shape should be [3, 3, 2, 2] + assert_eq!(padded.dims(), [3, 3, 2, 2]); + + let expected = TensorData::from([ + [ + [[0.0, 0.0], [0.0, 0.0]], + [[0.0, 0.0], [0.0, 0.0]], + [[0.0, 0.0], [0.0, 0.0]], + ], + [ + [[0.0, 0.0], [0.0, 0.0]], + [[1.0, 2.0], [3.0, 4.0]], + [[0.0, 0.0], [0.0, 0.0]], + ], + [ + [[0.0, 0.0], [0.0, 0.0]], + [[0.0, 0.0], [0.0, 0.0]], + [[0.0, 0.0], [0.0, 0.0]], + ], + ]); + padded.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_constant_batch_dim_only_test() { + // Pad only the batch dimension of a 3D tensor [N, H, W] + let tensor = TestTensor::<3>::from([[[1.0, 2.0], [3.0, 4.0]]]); + + // 3 pairs for 3 dims: batch(1,1), height(0,0), width(0,0) + let padded = tensor.pad([(1, 1), (0, 0), (0, 0)], -1.0); + + assert_eq!(padded.dims(), [3, 2, 2]); + + let expected = TensorData::from([ + [[-1.0, -1.0], [-1.0, -1.0]], + [[1.0, 2.0], [3.0, 4.0]], + [[-1.0, -1.0], [-1.0, -1.0]], + ]); + padded.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_reflect_pairs_test() { + // Reflect padding using pairs API + let tensor = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]); + + let padded = tensor.pad([(1, 1), (1, 1)], PadMode::Reflect); + + let expected = TensorData::from([ + [5.0, 4.0, 5.0, 6.0, 5.0], + [2.0, 1.0, 2.0, 3.0, 2.0], + [5.0, 4.0, 5.0, 6.0, 5.0], + [8.0, 7.0, 8.0, 9.0, 8.0], + [5.0, 4.0, 5.0, 6.0, 5.0], + ]); + padded.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_edge_pairs_test() { + // Edge padding using pairs API + let tensor = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]); + + let padded = tensor.pad([(1, 1), (1, 1)], PadMode::Edge); + + let expected = TensorData::from([ + [1.0, 1.0, 2.0, 2.0], + [1.0, 1.0, 2.0, 2.0], + [3.0, 3.0, 4.0, 4.0], + [3.0, 3.0, 4.0, 4.0], + ]); + padded.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_reflect_batch_dim_3d_test() { + // Reflect pad the batch dimension of a 3D tensor [N, H, W] + // Input shape: [3, 1, 2] - 3 batches, 1 row, 2 cols + let tensor = TestTensor::<3>::from([[[1.0, 2.0]], [[3.0, 4.0]], [[5.0, 6.0]]]); + + // Pad batch dim with reflect(1, 1), no spatial padding + let padded = tensor.pad([(1, 1), (0, 0), (0, 0)], PadMode::Reflect); + + assert_eq!(padded.dims(), [5, 1, 2]); + + // Reflect on batch: [3,4] [1,2] [3,4] [5,6] [3,4] + let expected = TensorData::from([ + [[3.0, 4.0]], + [[1.0, 2.0]], + [[3.0, 4.0]], + [[5.0, 6.0]], + [[3.0, 4.0]], + ]); + padded.into_data().assert_eq(&expected, false); +} + +#[test] +fn padding_edge_batch_dim_3d_test() { + // Edge pad the batch dimension of a 3D tensor + let tensor = TestTensor::<3>::from([[[1.0, 2.0]], [[3.0, 4.0]]]); + + let padded = tensor.pad([(2, 1), (0, 0), (0, 0)], PadMode::Edge); + + assert_eq!(padded.dims(), [5, 1, 2]); + + let expected = TensorData::from([ + [[1.0, 2.0]], + [[1.0, 2.0]], + [[1.0, 2.0]], + [[3.0, 4.0]], + [[3.0, 4.0]], + ]); + padded.into_data().assert_eq(&expected, false); +} + +#[test] +#[should_panic(expected = "Padding has")] +fn padding_too_many_pairs_panics_test() { + let tensor = TestTensor::<2>::from([[1.0, 2.0]]); + + // 3 pairs for a 2D tensor should panic + let _ = tensor.pad([(1, 1), (1, 1), (1, 1)], 0.0); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/permute.rs b/crates/burn-backend-tests/tests/tensor/float/ops/permute.rs new file mode 100644 index 0000000..c9a2982 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/permute.rs @@ -0,0 +1,58 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn permute_float_a() { + let tensor = TestTensor::<1>::from([ + 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12., 13., 14., 15., + ]) + .reshape([2, 2, 4]); + + let permuted = tensor.clone().permute([2, 1, 0]); + + let expected = TensorData::from([ + [[0., 8.], [4., 12.]], + [[1., 9.], [5., 13.]], + [[2., 10.], [6., 14.]], + [[3., 11.], [7., 15.]], + ]); + + permuted.into_data().assert_eq(&expected, false); + + // Test with negative axis + let permuted = tensor.clone().permute([-1, 1, 0]); + permuted.into_data().assert_eq(&expected, false); + + // Test with the same axis + let permuted = tensor.clone().permute([0, 1, 2]); + permuted.into_data().assert_eq(&tensor.into_data(), false); +} + +#[test] +fn permute_float() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device) + .reshape([2, 3, 4]) + .float(); + + let permuted = tensor.clone().permute([2, 1, 0]); + + // from pytorch: + // import torch; torch.arange(0, 24).reshape(2, 3, 4).permute(2, 1, 0).float() + let expected = TensorData::from([ + [[0., 12.], [4., 16.], [8., 20.]], + [[1., 13.], [5., 17.], [9., 21.]], + [[2., 14.], [6., 18.], [10., 22.]], + [[3., 15.], [7., 19.], [11., 23.]], + ]); + + permuted.into_data().assert_eq(&expected, false); + + // Test with negative axis + let permuted = tensor.clone().permute([-1, 1, 0]); + permuted.into_data().assert_eq(&expected, false); + + // Test with the same axis + let permuted = tensor.clone().permute([0, 1, 2]); + permuted.into_data().assert_eq(&tensor.into_data(), true); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/powf.rs b/crates/burn-backend-tests/tests/tensor/float/ops/powf.rs new file mode 100644 index 0000000..ace2936 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/powf.rs @@ -0,0 +1,105 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_powf_ops() { + let data = TensorData::from([[1.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + let pow = TensorData::from([[1.0, 1.0, 2.0], [3.0, 4.0, 2.0]]); + let tensor_pow = TestTensor::<2>::from_data(pow, &Default::default()); + + let output = tensor.powf(tensor_pow); + let expected = TensorData::from([[1.0, 1.0, 4.0], [27.0, 256.0, 25.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_neg_power() { + let data = TensorData::from([[1.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + let pow = TensorData::from([[-0.95, -0.67, -0.45], [-0.24, -0.5, -0.6]]); + let tensor_pow = TestTensor::<2>::from_data(pow, &Default::default()); + + let output = tensor.powf(tensor_pow); + let expected = TensorData::from([[1., 1., 0.73204285], [0.76822936, 0.5, 0.38073079]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_neg_values_with_even_power() { + let data = TensorData::from([[1.0, -1.0, -2.0], [-3.0, -4.0, -5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + let pow = TensorData::from([[2.0, 2.0, 4.0], [4.0, 4.0, 2.0]]); + let tensor_pow = TestTensor::<2>::from_data(pow, &Default::default()); + + let output = tensor.powf(tensor_pow); + let expected = TensorData::from([[1.0, 1.0, 16.0], [81.0, 256.0, 25.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_neg_values_with_odd_power() { + let data = TensorData::from([[1.0, -1.0, -2.0], [-3.0, -4.0, -5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + let pow = TensorData::from([[3.0, 3.0, 3.0], [3.0, 3.0, 3.0]]); + let tensor_pow = TestTensor::<2>::from_data(pow, &Default::default()); + + let output = tensor.powf(tensor_pow); + let expected = TensorData::from([[1.0, -1.0, -8.0], [-27.0, -64.0, -125.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_powf_broadcasted() { + let device = Default::default(); + let tensor_1 = TestTensor::<1>::from_data([2.0, 3.0, 4.0], &device); + let tensor_2 = TestTensor::from_data([1.0], &device); + + // Broadcast rhs + let output = tensor_1.clone().powf(tensor_2.clone()); + output + .into_data() + .assert_approx_eq::(&tensor_1.to_data(), Tolerance::default()); + + // Broadcast lhs + let output = tensor_2.powf(tensor_1); + output + .into_data() + .assert_approx_eq::(&TensorData::from([1.0, 1.0, 1.0]), Tolerance::default()); +} + +fn outer(a: TestTensor<1>, b: TestTensor<1>) -> TestTensor<2> { + a.unsqueeze_dim::<2>(1) * b.unsqueeze_dim::<2>(0) +} + +#[test] +fn should_support_powf_scalar_tensor() { + let device = Default::default(); + let head_dim = 64; + let seq_len = 1024; + let base = 10000; + + let channel_range = TestTensorInt::arange_step(0..head_dim as i64, 2, &device).float(); + let base = TestTensor::<1>::from_data([base as f32], &device); + let inv_freq = base.powf(-channel_range / head_dim as f32); + + let t = TestTensorInt::arange(0..seq_len as i64, &device).float(); + + let freqs = outer(t, inv_freq); + + let _cos = freqs.clone().cos(); + let _sin = freqs.sin(); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/powf_scalar.rs b/crates/burn-backend-tests/tests/tensor/float/ops/powf_scalar.rs new file mode 100644 index 0000000..1b19589 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/powf_scalar.rs @@ -0,0 +1,55 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_powf_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.powf_scalar(0.71); + let expected = TensorData::from([[0.0, 1.0, 1.6358], [2.1815, 2.67586, 3.13522]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_neg_power() { + let data = TensorData::from([[1.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.powf_scalar(-0.33); + let expected = TensorData::from([[1.0, 1.0, 0.79553646], [0.695905, 0.6328783, 0.58794934]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_neg_values_with_even_power() { + let data = TensorData::from([[0.0, -1.0, -2.0], [-3.0, -4.0, -5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.powf_scalar(4.0); + let expected = TensorData::from([[0.0, 1.0, 16.0], [81.0, 256.0, 625.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_neg_values_with_odd_power() { + let data = TensorData::from([[0.0, -1.0, -2.0], [-3.0, -4.0, -5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.powf_scalar(3.0); + let expected = TensorData::from([[0.0, -1.0, -8.0], [-27.0, -64.0, -125.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/prod.rs b/crates/burn-backend-tests/tests/tensor/float/ops/prod.rs new file mode 100644 index 0000000..dc9156e --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/prod.rs @@ -0,0 +1,48 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_prod_float() { + let tensor_1 = TestTensor::<2>::from([[-5.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor_1.prod(); + + output + .into_data() + .assert_eq(&TensorData::from([-600.0]), false); +} + +#[test] +fn test_prod_dim_2d() { + let f = TestTensor::<2>::from([[-5.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + f.clone() + .prod_dim(1) + .into_data() + .assert_eq(&TensorData::from([[-10.0], [60.0]]), false); + + f.clone() + .prod_dim(-1) + .into_data() + .assert_eq(&TensorData::from([[-10.0], [60.0]]), false); +} + +#[test] +fn test_prod_dims_2d() { + let f = TestTensor::<2>::from([[-5.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + f.clone() + .prod_dims(&[1]) + .into_data() + .assert_eq(&TensorData::from([[-10.0], [60.0]]), false); + + f.clone() + .prod_dims(&[-1]) + .into_data() + .assert_eq(&TensorData::from([[-10.0], [60.0]]), false); + + f.clone() + .prod_dims(&[0, 1]) + .into_data() + .assert_eq(&TensorData::from([[-600.0]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/random.rs b/crates/burn-backend-tests/tests/tensor/float/ops/random.rs new file mode 100644 index 0000000..871a8ac --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/random.rs @@ -0,0 +1,53 @@ +use super::*; +use burn_tensor::{Device, Distribution, ElementConversion, TensorData, Tolerance}; + +#[test] +fn rand_default() { + let tensor = TestTensor::<1>::random([20], Distribution::Default, &Default::default()); + + // check that the tensor is within the range of [0..1) (1 is exclusive) + // the conversion can ceil the value if `FloatElem` is less precise than f32 + let low = 0.elem::(); + let high = 1.elem::(); + if FloatElem::EPSILON.elem::() > f32::EPSILON { + tensor.into_data().assert_within_range_inclusive(low..=high); + } else { + tensor.into_data().assert_within_range(low..high); + } +} + +#[test] +fn rand_uniform() { + let tensor = TestTensor::<1>::random([20], Distribution::Uniform(4., 5.), &Default::default()); + let low = 4.elem::(); + let high = 5.elem::(); + + if FloatElem::EPSILON.elem::() > f32::EPSILON { + tensor.into_data().assert_within_range_inclusive(low..=high); + } else { + tensor.into_data().assert_within_range(low..high); + } +} + +#[test] +fn rand_bernoulli() { + let tensor = TestTensor::<1>::random([20], Distribution::Bernoulli(1.), &Default::default()); + + tensor.into_data().assert_eq( + &TensorData::new::(vec![1.elem(); 20], [20]), + true, + ); +} + +#[test] +#[ignore] // TODO: mark serial for backends that handle the same devices (e.g. fusion)? +fn test_seed_reproducibility() { + let device = Device::default(); + device.seed(42); + let t1 = TestTensor::<1>::random([5], Distribution::Default, &device); + device.seed(42); + let t2 = TestTensor::<1>::random([5], Distribution::Default, &device); + + t1.into_data() + .assert_approx_eq::(&t2.into_data(), Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/recip.rs b/crates/burn-backend-tests/tests/tensor/float/ops/recip.rs new file mode 100644 index 0000000..689b023 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/recip.rs @@ -0,0 +1,16 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_recip_ops() { + let data = TensorData::from([[0.5, 1.0, 2.0], [3.0, -4.0, -5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.recip(); + let expected = TensorData::from([[2.0, 1.0, 0.5], [0.33333, -0.25, -0.2]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/remainder.rs b/crates/burn-backend-tests/tests/tensor/float/ops/remainder.rs new file mode 100644 index 0000000..e48eb92 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/remainder.rs @@ -0,0 +1,241 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +/// From https://pytorch.org/docs/stable/generated/torch.remainder.html +#[test] +fn should_support_remainder_basic() { + let device = Default::default(); + let lhs = + TestTensor::<1>::from_data(TensorData::from([-3.0, -2.0, -1.0, 1.0, 2.0, 3.0]), &device); + let rhs = TestTensor::<1>::from_data(TensorData::from([2.0, 3.0, 1.0, 2.0, 1.0, 3.0]), &device); + let output = lhs.remainder(rhs); + let expected = TensorData::from([1.0, 1.0, -0.0, 1.0, 0.0, 0.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_remainder_basic_scalar() { + let data = TensorData::from([-3.0, -2.0, -1.0, 1.0, 2.0, 3.0]); + let device = Default::default(); + let tensor = TestTensor::<1>::from_data(data, &device); + + let output = tensor.remainder_scalar(2.0); + let expected = TensorData::from([1.0, 0.0, 1.0, 1.0, 0.0, 1.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_remainder_float() { + let device = Default::default(); + let lhs = TestTensor::<1>::from_data(TensorData::from([1.0, 2.0, 3.0, 4.0, 5.0]), &device); + let rhs = TestTensor::<1>::from_data( + TensorData::from([1.4233, 2.7313, 0.2641, 1.9651, 0.5897]), + &device, + ); + let output = lhs.remainder(rhs); + let expected = TensorData::from([1.0, 2.0, 0.0949, 0.0698, 0.2824]); + + // Metal has less precise remainder function + let tolerance = Tolerance::default() + .set_half_precision_relative(1e-2) + .set_half_precision_absolute(2e-3); + + output + .into_data() + .assert_approx_eq::(&expected, tolerance); +} + +/// Also from https://pytorch.org/docs/stable/generated/torch.remainder.html +#[test] +fn should_support_remainder_float_scalar() { + let data = TensorData::from([1.0, 2.0, 3.0, 4.0, 5.0]); + let device = Default::default(); + let tensor = TestTensor::<1>::from_data(data, &device); + + let output = tensor.clone().remainder_scalar(-1.5); + let expected = TensorData::from([-0.5, -1.0, 0.0, -0.5, -1.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_be_zero() { + let device = Default::default(); + let lhs = TestTensor::<1>::from_data(TensorData::from([0.0, 0.0, 0.0]), &device); + let rhs = TestTensor::<1>::from_data(TensorData::from([3.5, -2.1, 1e-4]), &device); + + let output = lhs.remainder(rhs); + let expected = TensorData::from([0.0, 0.0, 0.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_be_zero_scalar() { + let data = TensorData::from([0.0, 0.0, 0.0]); + let device = Default::default(); + let tensor = TestTensor::<1>::from_data(data, &device); + + let output = tensor.clone().remainder_scalar(3.5); + let expected = TensorData::from([0.0, 0.0, 0.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_have_no_remainder() { + let device = Default::default(); + let lhs = TestTensor::<1>::from_data( + // Previous values failed on some vulkan backends (driver bug?) + // TensorData::from([-1.4843, 1.1350, -2.1563, 1.0862, 0.5, 3.6587]), + TensorData::from([-1.0, 1.5, -2.0, 2.5, 0.5, 4.0]), + &device, + ); + let rhs = TestTensor::<1>::from_data( + // TensorData::from([1.4843, 1.1350, 2.1563, 1.0862, 0.5, 3.6587]), + TensorData::from([1.0, 1.5, 2.0, 2.5, 0.5, 4.0]), + &device, + ); + + let output = lhs.remainder(rhs); + let expected = TensorData::from([-0., 0., -0., 0., 0., 0.]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_have_no_remainder_scalar() { + let data = TensorData::from([-4.0, 4.0]); + let device = Default::default(); + let tensor = TestTensor::<1>::from_data(data, &device); + + let output = tensor.remainder_scalar(4.0); + let expected = TensorData::from([-0.0, 0.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_be_negative() { + let device = Default::default(); + + let lhs = TestTensor::<1>::from_data(TensorData::from([-7.0, -3.0, 2.0, 6.0]), &device); + let rhs = TestTensor::<1>::from_data(TensorData::from([-2.5, -2.1, -1.5, -3.25]), &device); + + let output = lhs.remainder(rhs); + let expected = TensorData::from([-2.0, -0.9, -1.0, -0.5]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_be_negative_scalar() { + let data = TensorData::from([-7.0, -3.0, 2.0, 6.0]); + let device = Default::default(); + let tensor = TestTensor::<1>::from_data(data, &device); + + let output = tensor.clone().remainder_scalar(-2.5); + let expected = TensorData::from([-2.0, -0.50, -0.50, -1.5]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_fp_dividends() { + let data = TensorData::from([-7.5, -2.5, 2.5, 7.5]); + let device = Default::default(); + let tensor = TestTensor::<1>::from_data(data, &device); + + let output = tensor.remainder_scalar(3.0); + let expected = TensorData::from([1.5, 0.5, 2.5, 1.5]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + // for tensor.remainder case, tests above have already covered float point dividend cases +} + +#[test] +fn should_support_large_divisor() { + let device = Default::default(); + + let lhs = TestTensor::<1>::from_data( + TensorData::from([-1.0, 1.0, -1.5, 1.5, -1.0, 1.0, -1.5, 1.5]), + &device, + ); + let rhs = TestTensor::<1>::from_data( + TensorData::from([10.0, 10.0, 10.0, 10.0, -10.0, -10.0, -10.0, -10.0]), + &device, + ); + let output = lhs.remainder(rhs); + let expected = TensorData::from([9.0, 1.0, 8.5, 1.5, -1.0, -9.0, -1.5, -8.5]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_large_divisor_scalar() { + let data = TensorData::from([-1.0, 1.0]); + let device = Default::default(); + let tensor = TestTensor::<1>::from_data(data, &device); + + let output = tensor.remainder_scalar(10.0); + let expected = TensorData::from([9.0, 1.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_remainder_op() { + let device = Default::default(); + let lhs = + TestTensor::<1>::from_data(TensorData::from([-3.0, -2.0, -1.0, 1.0, 2.0, 3.0]), &device); + let rhs = TestTensor::<1>::from_data(TensorData::from([2.0, 3.0, 1.0, 2.0, 1.0, 3.0]), &device); + + let output = lhs % rhs; + let expected = TensorData::from([1.0, 1.0, -0.0, 1.0, 0.0, 0.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_remainder_scalar_op() { + let data = TensorData::from([-3.0, -2.0, -1.0, 1.0, 2.0, 3.0]); + let device = Default::default(); + let tensor = TestTensor::<1>::from_data(data, &device); + + let output = tensor % 2.0; + let expected = TensorData::from([1.0, 0.0, 1.0, 1.0, 0.0, 1.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/repeat.rs b/crates/burn-backend-tests/tests/tensor/float/ops/repeat.rs new file mode 100644 index 0000000..be843ef --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/repeat.rs @@ -0,0 +1,108 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_repeat_ops_one_dimension() { + let data = TensorData::from([[0.0f32, 1.0f32, 2.0f32]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.repeat(&[4, 1, 1]); + let expected = TensorData::from([ + [0.0f32, 1.0f32, 2.0f32], + [0.0f32, 1.0f32, 2.0f32], + [0.0f32, 1.0f32, 2.0f32], + [0.0f32, 1.0f32, 2.0f32], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_float_repeat_repeating_on_many_dimensions() { + let data = TensorData::from([ + [[1.0f32, 2.0f32], [3.0f32, 4.0f32]], + [[5.0f32, 6.0f32], [7.0f32, 8.0f32]], + [[9.0f32, 10.0f32], [11.0f32, 12.0f32]], + [[13.0f32, 14.0f32], [15.0f32, 16.0f32]], + ]); + let tensor = TestTensor::<3>::from_data(data, &Default::default()); + + let output = tensor.repeat(&[2, 3, 2]); + let expected = TensorData::from([ + [ + [1.0f32, 2.0f32, 1.0f32, 2.0f32], + [3.0f32, 4.0f32, 3.0f32, 4.0f32], + [1.0f32, 2.0f32, 1.0f32, 2.0f32], + [3.0f32, 4.0f32, 3.0f32, 4.0f32], + [1.0f32, 2.0f32, 1.0f32, 2.0f32], + [3.0f32, 4.0f32, 3.0f32, 4.0f32], + ], + [ + [5.0f32, 6.0f32, 5.0f32, 6.0f32], + [7.0f32, 8.0f32, 7.0f32, 8.0f32], + [5.0f32, 6.0f32, 5.0f32, 6.0f32], + [7.0f32, 8.0f32, 7.0f32, 8.0f32], + [5.0f32, 6.0f32, 5.0f32, 6.0f32], + [7.0f32, 8.0f32, 7.0f32, 8.0f32], + ], + [ + [9.0f32, 10.0f32, 9.0f32, 10.0f32], + [11.0f32, 12.0f32, 11.0f32, 12.0f32], + [9.0f32, 10.0f32, 9.0f32, 10.0f32], + [11.0f32, 12.0f32, 11.0f32, 12.0f32], + [9.0f32, 10.0f32, 9.0f32, 10.0f32], + [11.0f32, 12.0f32, 11.0f32, 12.0f32], + ], + [ + [13.0f32, 14.0f32, 13.0f32, 14.0f32], + [15.0f32, 16.0f32, 15.0f32, 16.0f32], + [13.0f32, 14.0f32, 13.0f32, 14.0f32], + [15.0f32, 16.0f32, 15.0f32, 16.0f32], + [13.0f32, 14.0f32, 13.0f32, 14.0f32], + [15.0f32, 16.0f32, 15.0f32, 16.0f32], + ], + [ + [1.0f32, 2.0f32, 1.0f32, 2.0f32], + [3.0f32, 4.0f32, 3.0f32, 4.0f32], + [1.0f32, 2.0f32, 1.0f32, 2.0f32], + [3.0f32, 4.0f32, 3.0f32, 4.0f32], + [1.0f32, 2.0f32, 1.0f32, 2.0f32], + [3.0f32, 4.0f32, 3.0f32, 4.0f32], + ], + [ + [5.0f32, 6.0f32, 5.0f32, 6.0f32], + [7.0f32, 8.0f32, 7.0f32, 8.0f32], + [5.0f32, 6.0f32, 5.0f32, 6.0f32], + [7.0f32, 8.0f32, 7.0f32, 8.0f32], + [5.0f32, 6.0f32, 5.0f32, 6.0f32], + [7.0f32, 8.0f32, 7.0f32, 8.0f32], + ], + [ + [9.0f32, 10.0f32, 9.0f32, 10.0f32], + [11.0f32, 12.0f32, 11.0f32, 12.0f32], + [9.0f32, 10.0f32, 9.0f32, 10.0f32], + [11.0f32, 12.0f32, 11.0f32, 12.0f32], + [9.0f32, 10.0f32, 9.0f32, 10.0f32], + [11.0f32, 12.0f32, 11.0f32, 12.0f32], + ], + [ + [13.0f32, 14.0f32, 13.0f32, 14.0f32], + [15.0f32, 16.0f32, 15.0f32, 16.0f32], + [13.0f32, 14.0f32, 13.0f32, 14.0f32], + [15.0f32, 16.0f32, 15.0f32, 16.0f32], + [13.0f32, 14.0f32, 13.0f32, 14.0f32], + [15.0f32, 16.0f32, 15.0f32, 16.0f32], + ], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_repeat_0_times_empty() { + let tensor = TestTensor::<3>::ones([2, 3, 4], &Default::default()); + + let output = tensor.repeat(&[1, 0, 2]); + + assert_eq!(output.shape(), [2, 0, 8].into()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/repeat_dim.rs b/crates/burn-backend-tests/tests/tensor/float/ops/repeat_dim.rs new file mode 100644 index 0000000..7504cad --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/repeat_dim.rs @@ -0,0 +1,182 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_repeat_ops() { + let data = TensorData::from([[0.0f64, 1.0f64, 2.0f64]]); + let tensor = TestTensor::<2>::from_data(data.clone(), &Default::default()); + + let output = tensor.repeat_dim(0, 4); + let expected = TensorData::from([ + [0.0f32, 1.0f32, 2.0f32], + [0.0f32, 1.0f32, 2.0f32], + [0.0f32, 1.0f32, 2.0f32], + [0.0f32, 1.0f32, 2.0f32], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_float_repeat_on_dims_larger_than_1() { + let data = TensorData::from([ + [[1.0f32, 2.0f32], [3.0f32, 4.0f32]], + [[5.0f32, 6.0f32], [7.0f32, 8.0f32]], + [[9.0f32, 10.0f32], [11.0f32, 12.0f32]], + [[13.0f32, 14.0f32], [15.0f32, 16.0f32]], + ]); + let tensor = TestTensor::<3>::from_data(data, &Default::default()); + + let output = tensor.repeat_dim(2, 2); + let expected = TensorData::from([ + [ + [1.0f32, 2.0f32, 1.0f32, 2.0f32], + [3.0f32, 4.0f32, 3.0f32, 4.0f32], + ], + [ + [5.0f32, 6.0f32, 5.0f32, 6.0f32], + [7.0f32, 8.0f32, 7.0f32, 8.0f32], + ], + [ + [9.0f32, 10.0f32, 9.0f32, 10.0f32], + [11.0f32, 12.0f32, 11.0f32, 12.0f32], + ], + [ + [13.0f32, 14.0f32, 13.0f32, 14.0f32], + [15.0f32, 16.0f32, 15.0f32, 16.0f32], + ], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn repeat_dim_swap_dims_1() { + let tensor = TestTensorInt::arange(0..16, &Default::default()).float(); + + let tensor = tensor.reshape([4, 1, 4]); + let tensor = tensor.swap_dims(0, 2); + let output = tensor.repeat_dim(1, 4); + + let expected = TensorData::from([ + [ + [0.0, 4.0, 8.0, 12.0], + [0.0, 4.0, 8.0, 12.0], + [0.0, 4.0, 8.0, 12.0], + [0.0, 4.0, 8.0, 12.0], + ], + [ + [1.0, 5.0, 9.0, 13.0], + [1.0, 5.0, 9.0, 13.0], + [1.0, 5.0, 9.0, 13.0], + [1.0, 5.0, 9.0, 13.0], + ], + [ + [2.0, 6.0, 10.0, 14.0], + [2.0, 6.0, 10.0, 14.0], + [2.0, 6.0, 10.0, 14.0], + [2.0, 6.0, 10.0, 14.0], + ], + [ + [3.0, 7.0, 11.0, 15.0], + [3.0, 7.0, 11.0, 15.0], + [3.0, 7.0, 11.0, 15.0], + [3.0, 7.0, 11.0, 15.0], + ], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn repeat_dim_swap_dims_2() { + let tensor = TestTensorInt::arange(0..16, &Default::default()).float(); + + let tensor = tensor.reshape([2, 2, 1, 4]); + let tensor = tensor.swap_dims(0, 1); + let output = tensor.repeat_dim(2, 4); + + let expected = TensorData::from([ + [ + [ + [0.0, 1.0, 2.0, 3.0], + [0.0, 1.0, 2.0, 3.0], + [0.0, 1.0, 2.0, 3.0], + [0.0, 1.0, 2.0, 3.0], + ], + [ + [8.0, 9.0, 10.0, 11.0], + [8.0, 9.0, 10.0, 11.0], + [8.0, 9.0, 10.0, 11.0], + [8.0, 9.0, 10.0, 11.0], + ], + ], + [ + [ + [4.0, 5.0, 6.0, 7.0], + [4.0, 5.0, 6.0, 7.0], + [4.0, 5.0, 6.0, 7.0], + [4.0, 5.0, 6.0, 7.0], + ], + [ + [12.0, 13.0, 14.0, 15.0], + [12.0, 13.0, 14.0, 15.0], + [12.0, 13.0, 14.0, 15.0], + [12.0, 13.0, 14.0, 15.0], + ], + ], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn repeat_dim_swap_dims_3() { + let tensor = TestTensorInt::arange(0..16, &Default::default()).float(); + + let tensor = tensor.reshape([1, 2, 2, 4]); + let tensor = tensor.swap_dims(0, 2); + let tensor = tensor.swap_dims(1, 3); + let output = tensor.repeat_dim(2, 4); + + let expected = TensorData::from([ + [ + [[0.0, 8.0], [0.0, 8.0], [0.0, 8.0], [0.0, 8.0]], + [[1.0, 9.0], [1.0, 9.0], [1.0, 9.0], [1.0, 9.0]], + [[2.0, 10.0], [2.0, 10.0], [2.0, 10.0], [2.0, 10.0]], + [[3.0, 11.0], [3.0, 11.0], [3.0, 11.0], [3.0, 11.0]], + ], + [ + [[4.0, 12.0], [4.0, 12.0], [4.0, 12.0], [4.0, 12.0]], + [[5.0, 13.0], [5.0, 13.0], [5.0, 13.0], [5.0, 13.0]], + [[6.0, 14.0], [6.0, 14.0], [6.0, 14.0], [6.0, 14.0]], + [[7.0, 15.0], [7.0, 15.0], [7.0, 15.0], [7.0, 15.0]], + ], + ]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_repeat_dim_0_times_empty() { + let tensor = TestTensor::<3>::ones([2, 3, 4], &Default::default()); + + let output = tensor.repeat_dim(2, 0); + + assert_eq!(output.shape(), [2, 3, 0].into()); +} + +#[test] +fn should_support_unsqueeze_of_repeat_dim_view() { + // `repeat_dim` on a size-1 dim returns a broadcast view (stride 0 on + // that dim). A following unsqueeze (reshape) must not let the unit + // dim's stride leak onto real dims: [3, 1] with strides [1, 0] + // reshaped to [3, 1, 1] has to keep stride 1 on dim 0, otherwise + // every row reads element 0. + let device = Default::default(); + let tensor = TestTensor::<1>::from_data([1.0, 2.0, 3.0], &device); + let view: TestTensor<2> = tensor.unsqueeze_dim(1).repeat_dim(1, 1); + let view: TestTensor<3> = view.unsqueeze_dim(2); + + view.into_data() + .assert_eq(&TensorData::from([[[1.0]], [[2.0]], [[3.0]]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/reshape.rs b/crates/burn-backend-tests/tests/tensor/float/ops/reshape.rs new file mode 100644 index 0000000..4b1a772 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/reshape.rs @@ -0,0 +1,90 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_rank() { + let data = TensorData::from([0.0, 1.0, 2.0]); + let tensor = TestTensor::<1>::from_data(data, &Default::default()); + assert_eq!(tensor.rank(), 1); + + let data = TensorData::from([[0.0, 1.0, 2.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + assert_eq!(tensor.rank(), 2); +} + +#[test] +fn should_support_reshape_1d() { + let data = TensorData::from([0.0, 1.0, 2.0]); + let tensor = TestTensor::<1>::from_data(data, &Default::default()); + + let output = tensor.clone().reshape([1, 3]); + let expected = TensorData::from([[0.0, 1.0, 2.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_reshape_2d() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.clone().reshape([6]); + let expected = TensorData::from([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_dim_infererence() { + let data = TensorData::from([ + [0.0, 1.0, 2.0], + [3.0, 4.0, 5.0], + [6.0, 7.0, 8.0], + [9.0, 10.0, 11.0], + ]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + // Infer the dimension via -1 + let reshaped = tensor.clone().reshape([2, -1]); + assert_eq!(reshaped.shape(), [2, 6].into()); + + // Infer the dimension via 0 (keep from the source) and -1 (infer) + let reshaped = reshaped.reshape([0, 2, -1]); + assert_eq!(reshaped.shape(), [2, 2, 3].into()); + + // This is effectively as if we did a flatten + let reshaped = tensor.clone().reshape([-1]); + assert_eq!(reshaped.shape(), [12].into()); + + // Keeping the first dimension the same (using 0) + let reshaped = tensor.clone().reshape([0, 3]); + assert_eq!(reshaped.shape(), [4, 3].into()); +} + +#[test] +fn should_not_corrupt_after_slice() { + let zeros = TestTensor::<1>::zeros([2], &Default::default()); + zeros.clone().slice([1..2]).reshape([1]).exp(); + + // May lead to zeroes being equal to [0.0, 1.0] + zeros.into_data().assert_eq( + &TestTensor::<1>::zeros([2], &Default::default()).to_data(), + true, + ); +} + +#[test] +#[should_panic] +fn multiple_neg_ones() { + let data = TensorData::from([0.0, 1.0, 2.0]); + let tensor = TestTensor::<1>::from_data(data, &Default::default()); + let _data_actual = tensor.reshape([-1, -1]).into_data(); +} + +#[test] +#[should_panic] +fn neg_value() { + let data = TensorData::from([0.0, 1.0, 2.0]); + let tensor = TestTensor::<1>::from_data(data, &Default::default()); + let _data_actual = tensor.reshape([-2, -1]).into_data(); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/round.rs b/crates/burn-backend-tests/tests/tensor/float/ops/round.rs new file mode 100644 index 0000000..f6161da --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/round.rs @@ -0,0 +1,29 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_round_ops() { + let data = TensorData::from([[24.0423, 87.9478, 76.1838], [59.6929, 43.8169, 94.8826]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.round(); + let expected = TensorData::from([[24., 88., 76.], [60., 44., 95.]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_round_ties_even() { + let data = TensorData::from([1.5, 2.5, 3.5, 4.5, 5.5, 6.5]); + let tensor = TestTensor::<1>::from_data(data, &Default::default()); + + let output = tensor.round(); + let expected = TensorData::from([2., 2., 4., 4., 6., 6.]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/select.rs b/crates/burn-backend-tests/tests/tensor/float/ops/select.rs new file mode 100644 index 0000000..3231c85 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/select.rs @@ -0,0 +1,251 @@ +use super::*; +use burn_tensor::{IndexingUpdateOp, TensorData}; + +#[test] +fn should_select_1d() { + let device = Default::default(); + let tensor = TestTensor::<1>::from_data([0.0, 1.0, 2.0], &device); + let indices = TestTensorInt::from_data([1, 1, 0, 1, 2], &device); + + let output = tensor.select(0, indices); + let expected = TensorData::from([1.0, 1.0, 0.0, 1.0, 2.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_select_2d_dim0_same_num_dim() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &device); + let indices = TestTensorInt::from_data([1, 0], &device); + + let output = tensor.select(0, indices); + let expected = TensorData::from([[3.0, 4.0, 5.0], [0.0, 1.0, 2.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_select_2d_dim0_more_num_dim() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &device); + let indices = TestTensorInt::from_data([1, 0, 1, 1], &device); + + let output = tensor.select(0, indices); + let expected = TensorData::from([ + [3.0, 4.0, 5.0], + [0.0, 1.0, 2.0], + [3.0, 4.0, 5.0], + [3.0, 4.0, 5.0], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_select_2d_dim0_vec() { + let device = Default::default(); + let tensor = + TestTensor::<2>::from_data([[0.0, 1.0], [2.0, 3.0], [4.0, 5.0], [6.0, 7.0]], &device); + let indices = TestTensorInt::from_data([1, 0, 3, 2], &device); + + let output = tensor.select(0, indices); + let expected = TensorData::from([[2.0, 3.0], [0.0, 1.0], [6.0, 7.0], [4.0, 5.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_select_2d_dim1() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &device); + let indices = TestTensorInt::from_data([1, 1, 0, 1, 2], &device); + + let output = tensor.select(1, indices); + let expected = TensorData::from([[1.0, 1.0, 0.0, 1.0, 2.0], [4.0, 4.0, 3.0, 4.0, 5.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_select_add_1d() { + let device = Default::default(); + let tensor = TestTensor::<1>::from_data([0.0, 1.0, 2.0], &device); + let values = TestTensor::from_data([5.0, 4.0, 3.0, 2.0, 1.0], &device); + let indices = TestTensorInt::from_data(TensorData::from([1, 1, 0, 1, 2]), &device); + + let output = tensor.select_assign(0, indices, values, IndexingUpdateOp::Add); + let expected = TensorData::from([3.0, 12.0, 3.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_select_add_1d_int() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::from_data([7, 8, 9], &device); + let values = TestTensorInt::from_data([5, 4, 3, 2, 1], &device); + let indices = TestTensorInt::from_data(TensorData::from([1, 1, 0, 1, 2]), &device); + + let output = tensor.select_assign(0, indices, values, IndexingUpdateOp::Add); + let expected = TensorData::from([10, 19, 10]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_select_add_2d_dim0() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &device); + let values = TestTensor::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + let indices = TestTensorInt::from_data(TensorData::from([1, 0]), &device); + + let output = tensor.select_assign(0, indices, values, IndexingUpdateOp::Add); + let expected = TensorData::from([[4.0, 6.0, 8.0], [4.0, 6.0, 8.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_select_add_2d_dim1() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &device); + let values = TestTensor::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + let indices = TestTensorInt::from_data(TensorData::from([1, 0, 2]), &device); + + let output = tensor.select_assign(1, indices, values, IndexingUpdateOp::Add); + let expected = TensorData::from([[2.0, 2.0, 5.0], [8.0, 8.0, 11.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_select_3d_dim1_vec() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data( + [ + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]], + [[-1.0, -2.0], [-3.0, -4.0], [-5.0, -6.0], [-7.0, -8.0]], + ], + &device, + ); + let indices = TestTensorInt::from_data([1, 0, 3, 2], &device); + + let output = tensor.select(1, indices); + let expected = TensorData::from([ + [[3.0, 4.0], [1.0, 2.0], [7.0, 8.0], [5.0, 6.0]], + [[-3.0, -4.0], [-1.0, -2.0], [-7.0, -8.0], [-5.0, -6.0]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +#[should_panic] +fn should_select_panic_invalid_dimension() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &device); + let indices = TestTensorInt::from_data([1, 1, 0, 1, 2], &device); + + tensor.select(10, indices); +} + +#[test] +fn should_match_default_implementation_behavior() { + // Verify optimized implementation matches original default logic + let device = Default::default(); + let tensor = TestTensorBool::<1>::from_data([true, false, true], &device); + let indices = TestTensorInt::from_data([0, 1, 0], &device); + let values = TestTensorBool::<1>::from_data([false, true, true], &device); + + let optimized_result = + tensor + .clone() + .select_assign(0, indices.clone(), values.clone(), IndexingUpdateOp::Add); + + // Manual default implementation logic + let int_tensor = tensor.int(); + let int_values = values.int(); + let assigned = int_tensor.select_assign(0, indices, int_values, IndexingUpdateOp::Add); + let default_result = assigned.greater_elem(0); + + optimized_result + .into_data() + .assert_eq(&default_result.into_data(), false); +} + +#[test] +fn should_select_with_negative_dim_2d() { + // Test using negative dimension indexing on 2D tensor + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &device); + let indices = TestTensorInt::from_data([1, 0, 2], &device); + + // Using -1 should refer to the last dimension (dim 1) + let output_neg = tensor.clone().select(-1, indices.clone()); + let output_pos = tensor.select(1, indices); + + // Both should produce the same result + output_neg + .into_data() + .assert_eq(&output_pos.into_data(), false); +} + +#[test] +fn should_select_add_with_negative_dim_2d() { + // Test select_add with negative dimension on 2D tensor + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &device); + let values = TestTensor::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + let indices = TestTensorInt::from_data([0, 2], &device); + + // Using -1 should refer to the last dimension (dim 1) + let output_neg = + tensor + .clone() + .select_assign(-1, indices.clone(), values.clone(), IndexingUpdateOp::Add); + let output_pos = tensor.select_assign(1, indices, values, IndexingUpdateOp::Add); + + output_neg + .into_data() + .assert_eq(&output_pos.into_data(), false); +} + +#[test] +#[should_panic] +fn should_panic_select_negative_dim_out_of_bounds() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + let indices = TestTensorInt::from_data([0, 1], &device); + + // This should panic because -3 is out of bounds for a 2D tensor + tensor.select(-3, indices); +} + +#[test] +#[should_panic] +fn should_panic_select_add_negative_dim_out_of_bounds() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + let values = TestTensor::from_data([[5.0], [6.0]], &device); + let indices = TestTensorInt::from_data([0], &device); + + // This should panic because -3 is out of bounds for a 2D tensor + tensor.select_assign(-3, indices, values, IndexingUpdateOp::Add); +} + +#[test] +#[ignore = "Empty dims support is sparse"] +fn should_select_2d_dim0_empty_indices() { + // Zero-length index tensor: output should keep shape [0, cols] without + // triggering OOB reads. + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], &device); + let indices = TestTensorInt::<1>::from_data(TensorData::new(Vec::::new(), [0]), &device); + + let output = tensor.select(0, indices); + + assert_eq!(output.dims(), [0, 2]); + let out: Vec = output.into_data().to_vec().unwrap(); + assert!(out.is_empty()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/sign.rs b/crates/burn-backend-tests/tests/tensor/float/ops/sign.rs new file mode 100644 index 0000000..6b71edc --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/sign.rs @@ -0,0 +1,12 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_sign_ops_float() { + let tensor = TestTensor::<2>::from([[-0.2, -1.0, 2.0], [3.0, 0.0, -5.0]]); + + let output = tensor.sign(); + let expected = TensorData::from([[-1.0, -1.0, 1.0], [1.0, 0.0, -1.0]]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/slice.rs b/crates/burn-backend-tests/tests/tensor/float/ops/slice.rs new file mode 100644 index 0000000..5080fc1 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/slice.rs @@ -0,0 +1,591 @@ +use super::*; +use burn_tensor::{ElementConversion, Slice, TensorData, s}; + +#[test] +fn should_support_slice_dim_1d() { + let data = TensorData::from([0.0, 1.0, 2.0]); + let tensor = TestTensor::<1>::from_data(data.clone(), &Default::default()); + + // Test with range (negative index) + let output = tensor.clone().slice_dim(0, -2..); + output + .into_data() + .assert_eq(&TensorData::from([1.0, 2.0]), false); + + // Test with Slice directly + let slice = Slice::new(1, None, 1); // equivalent to 1.. + let output = tensor.slice_dim(0, slice); + output + .into_data() + .assert_eq(&TensorData::from([1.0, 2.0]), false); +} + +#[test] +#[should_panic(expected = "The provided dimension exceeds the tensor dimensions")] +fn should_panic_when_slice_dim_1d_bad_dim() { + let data = TensorData::from([0.0, 1.0, 2.0]); + let tensor = TestTensor::<1>::from_data(data.clone(), &Default::default()); + + let _output = tensor.slice_dim(1, 1..); +} + +#[test] +fn should_support_slice_dim_2d() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data.clone(), &Default::default()); + + let output = tensor.slice_dim(1, 1..); + output + .into_data() + .assert_eq(&TensorData::from([[1.0, 2.0], [4.0, 5.0]]), false); +} + +#[test] +fn should_support_slice_dim_with_step() { + let data = TensorData::from([[0.0, 1.0, 2.0, 3.0], [4.0, 5.0, 6.0, 7.0]]); + let tensor = TestTensor::<2>::from_data(data.clone(), &Default::default()); + + // Test 1: Slice dimension 1 with step=2 using s! macro + let output = tensor.clone().slice_dim(1, s![0..4;2]); + output + .into_data() + .assert_eq(&TensorData::from([[0.0, 2.0], [4.0, 6.0]]), false); + + // Test 2: Slice dimension 1 with step=2 using Slice directly + let slice = Slice::new(0, Some(4), 2); + let output = tensor.slice_dim(1, slice); + output + .into_data() + .assert_eq(&TensorData::from([[0.0, 2.0], [4.0, 6.0]]), false); +} + +#[test] +fn should_support_slice_dim_with_negative_step() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data.clone(), &Default::default()); + + // Slice dimension 1 with negative step (reverse columns) + let output = tensor.slice_dim(1, s![..;-1]); + output + .into_data() + .assert_eq(&TensorData::from([[2.0, 1.0, 0.0], [5.0, 4.0, 3.0]]), false); +} + +#[test] +fn should_support_full_sliceing_1d() { + let data = TensorData::from([0.0, 1.0, 2.0]); + let tensor = TestTensor::<1>::from_data(data.clone(), &Default::default()); + + let output = tensor.slice([0..3]); + + output.into_data().assert_eq(&data, false); +} + +#[test] +fn should_support_full_sliceing_vec() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data.clone(), &Default::default()); + + let slices: Vec = vec![(0..2).into()]; + + let output = tensor.clone().slice(&slices); + output.into_data().assert_eq(&data, false); + + let output = tensor.slice([0..2, 0..3]); + output.into_data().assert_eq(&data, false); +} + +#[test] +fn should_support_partial_sliceing_1d() { + let data = TensorData::from([0.0, 1.0, 2.0]); + let tensor = TestTensor::<1>::from_data(data, &Default::default()); + + let output = tensor.slice([1..3]); + let expected = TensorData::from([1.0, 2.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_full_sliceing_2d() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data.clone(), &Default::default()); + + let output = tensor.clone().slice([0..2]); + output.into_data().assert_eq(&data, false); + + let output = tensor.slice([0..2, 0..3]); + output.into_data().assert_eq(&data, false); +} + +#[test] +fn should_support_partial_sliceing_2d() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.slice([0..2, 0..2]); + let expected = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_slice_range_first_dim() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.slice(0..1); + let expected = TensorData::from([[0.0, 1.0, 2.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_partial_sliceing_3d() { + let tensor = TestTensor::<3>::from_data( + [ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], + ], + &Default::default(), + ); + + let output = tensor.slice([1..2, 1..2, 0..2]); + let expected = TensorData::from([[[9.0, 10.0]]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_partial_sliceing_3d_non_contiguous() { + let tensor = TestTensor::<3>::from_data( + [ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], + ], + &Default::default(), + ); + + let output = tensor.transpose().slice([1..2, 1..2, 0..2]); + let expected = TensorData::from([[[7.0, 10.0]]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_slice_fill_1d() { + let data = TensorData::from([0.0, 1.0, 2.0]); + + let device = Default::default(); + let tensor = TestTensor::<1>::from_data(data, &device); + + let output = tensor.slice_fill([0..2], -1.0); + let expected = TensorData::from([-1.0, -1.0, 2.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_slice_fill_vec() { + let data = TensorData::from([0.0, 1.0, 2.0]); + + let device = Default::default(); + let tensor = TestTensor::<1>::from_data(data, &device); + + let slices: Vec = vec![(0..2).into()]; + + let output = tensor.slice_fill(&slices, -1.0); + let expected = TensorData::from([-1.0, -1.0, 2.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_slice_fill_cast_f32() { + let data = TensorData::from([0.0, 1.0, 2.0]); + let device = Default::default(); + let tensor = TestTensor::<1>::from_data(data, &device).cast(burn_tensor::DType::F32); + + tensor + .slice_fill(s![0..2], 1.0) + .into_data() + .assert_eq(&TensorData::from([1.0, 1.0, 2.0]), false); +} + +// Skip on metal - F64 not supported +#[cfg(not(feature = "metal"))] +#[test] +fn should_support_slice_fill_cast_f64() { + let data = TensorData::from([0.0, 1.0, 2.0]); + let device = Default::default(); + let tensor = TestTensor::<1>::from_data(data, &device).cast(burn_tensor::DType::F64); + + tensor + .slice_fill(s![0..2], 1.0) + .into_data() + .assert_eq(&TensorData::from([1.0, 1.0, 2.0]), false); +} + +#[test] +fn should_support_slice_fill_1d_neg() { + let data = TensorData::from([0.0, 1.0, 2.0]); + + let device = Default::default(); + let tensor = TestTensor::<1>::from_data(data, &device); + + let output = tensor.slice_fill([-1..], -1.0); + let expected = TensorData::from([0.0, 1.0, -1.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_slice_fill_2d() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let device = Default::default(); + let tensor = TestTensor::<2>::from_data(data, &device); + + let output = tensor.slice_fill([1..2, 0..2], -1.0); + let expected = TensorData::from([[0.0, 1.0, 2.0], [-1.0, -1.0, 5.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_slice_fill_with_positive_step() { + let device = Default::default(); + + // Test 1D tensor with step + let tensor = TestTensor::<1>::zeros([10], &device); + let output = tensor.slice_fill(s![0..10;2], 5.0); + let expected = TensorData::from([5.0, 0.0, 5.0, 0.0, 5.0, 0.0, 5.0, 0.0, 5.0, 0.0]); + output.into_data().assert_eq(&expected, false); + + // Test 2D tensor with step on first dimension + let tensor = TestTensor::<2>::zeros([4, 4], &device); + let output = tensor.slice_fill(s![0..4;2, ..], 3.0); + let expected = TensorData::from([ + [3.0, 3.0, 3.0, 3.0], + [0.0, 0.0, 0.0, 0.0], + [3.0, 3.0, 3.0, 3.0], + [0.0, 0.0, 0.0, 0.0], + ]); + output.into_data().assert_eq(&expected, false); + + // Test 2D tensor with step on second dimension + let tensor = TestTensor::<2>::zeros([3, 6], &device); + let output = tensor.slice_fill(s![.., 0..6;3], 2.0); + let expected = TensorData::from([ + [2.0, 0.0, 0.0, 2.0, 0.0, 0.0], + [2.0, 0.0, 0.0, 2.0, 0.0, 0.0], + [2.0, 0.0, 0.0, 2.0, 0.0, 0.0], + ]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_slice_fill_with_negative_step() { + let device = Default::default(); + + // Test 1D tensor with negative step (reverse fill) + let tensor = TestTensor::<1>::from_data([1.0, 2.0, 3.0, 4.0, 5.0], &device); + let output = tensor.slice_fill(s![0..5;-1], 10.0); + // Should reverse the indices [4,3,2,1,0] and fill them with 10.0 + let expected = TensorData::from([10.0, 10.0, 10.0, 10.0, 10.0]); + output.into_data().assert_eq(&expected, false); + + // Test 2D tensor with negative step + let tensor = + TestTensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], &device); + let output = tensor.slice_fill(s![.., 0..3;-2], -1.0); + // Should fill columns in reverse order with step 2: indices 2, 0 + let expected = TensorData::from([[-1.0, 2.0, -1.0], [-1.0, 5.0, -1.0], [-1.0, 8.0, -1.0]]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_slice_fill_with_mixed_steps() { + let device = Default::default(); + + // Test 2D tensor with mixed positive and negative steps + let tensor = TestTensor::<2>::zeros([4, 6], &device); + let output = tensor.slice_fill(s![0..4;2, 0..6;-3], 7.0); + // Step 2 on dim 0 selects rows 0, 2 + // Step -3 on dim 1 with range 0..6 reverses and takes every 3rd: indices [5, 2] + let expected = TensorData::from([ + [0.0, 0.0, 7.0, 0.0, 0.0, 7.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 7.0, 0.0, 0.0, 7.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ]); + output.into_data().assert_eq(&expected, false); + + // Test 3D tensor with steps + let tensor = TestTensor::<3>::zeros([2, 4, 4], &device); + let output = tensor.slice_fill(s![.., 0..4;2, 0..4;-2], 1.0); + // Step 2 on dim 1 selects rows 0, 2 + // Step -2 on dim 2 with range 0..4 reverses and takes every 2nd: indices [3, 1] + let expected_slice = [ + [0.0, 1.0, 0.0, 1.0], + [0.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 1.0], + [0.0, 0.0, 0.0, 0.0], + ]; + let expected = TensorData::from([expected_slice, expected_slice]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn clamp_when_slice_exceeds_dimension() { + let tensor = TestTensor::<1>::from([0.0, 1.0, 2.0]); + let data = tensor.to_data(); + + let output = tensor.slice([0..4]); + output.into_data().assert_eq(&data, true); +} + +#[test] +fn negative_dimensions() { + let tensor = TestTensor::<2>::from([[0.0f32, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let data = tensor.to_data(); + + // Clamping to the tensor dimensions + let output = tensor.clone().slice([0..4, 0..4]); + output.into_data().assert_eq(&data, true); + + // Negative dimensions + let output = tensor.clone().slice([0..1, 0..1]); + let data = TensorData::from([[0.elem::()]]); + output.into_data().assert_eq(&data, true); + + let output = tensor.slice(s![0..-1, 0..-2]); + output.into_data().assert_eq(&data, true); +} + +#[test] +fn missing_dimensions() { + let tensor = TestTensor::<2>::from([[0.0f32, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let data = tensor.to_data(); + + // Clamping to the tensor dimensions + let output = tensor.clone().slice([0..4, 0..4]); + output.into_data().assert_eq(&data, true); + + // Negative dimensions + let data = TensorData::from([[0.elem::()]]); + let output = tensor.clone().slice(s![0..-1, 0..-2]); + output.into_data().assert_eq(&data, true); + + // Missing dimensions + let output = tensor.clone().slice(s![0..1, ..]); + let data = TensorData::from([[0.0f32, 1.0, 2.0]]); + output.into_data().assert_eq(&data, false); + + let output = tensor.clone().slice(s![.., 0..2]); + let data = TensorData::from([[0.0f32, 1.0], [3.0, 4.0]]); + output.into_data().assert_eq(&data, false); + + let output = tensor.clone().slice([.., ..]); + let data = TensorData::from([[0.0f32, 1.0, 2.0], [3.0, 4.0, 5.0]]); + output.into_data().assert_eq(&data, false); +} + +#[test] +fn should_slice_aggregation_result() { + let tensor = TestTensor::<1>::from([0.0, 1.0, 2.0]).mean(); + + let output = tensor.clone().slice([(0..1)]); + output.into_data().assert_eq(&tensor.into_data(), true); +} + +#[test] +#[should_panic] +fn should_panic_when_slice_with_too_many_dimensions() { + let tensor = TestTensor::<1>::from([0.0, 1.0, 2.0]); + + let _output = tensor.slice([0..1, 0..1]); +} + +#[test] +fn should_support_descending_slice_as_empty() { + // Like PyTorch, x[3:1] should return an empty tensor, not panic + let data = TensorData::from([0.0, 1.0, 2.0]); + let tensor = TestTensor::<1>::from_data(data, &Default::default()); + + let output = tensor.slice(s![2..1]); + + // Should produce an empty tensor with shape [0] + assert_eq!(output.dims(), [0]); +} + +#[test] +fn should_support_empty_slice() { + // ONNX models can have empty slices where start == end + // This should produce a tensor with size 0 in that dimension + let data = TensorData::from([0.0, 1.0, 2.0]); + let tensor = TestTensor::<1>::from_data(data, &Default::default()); + + let output = tensor.slice([1..1]); + + // Should produce an empty tensor with shape [0] + assert_eq!(output.dims(), [0]); +} + +#[test] +fn should_support_empty_slice_2d() { + // Test empty slice on 2D tensor + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + // Empty slice on first dimension + let output = tensor.clone().slice([1..1, 0..3]); + assert_eq!(output.dims(), [0, 3]); + + // Empty slice on second dimension + let output = tensor.slice([0..2, 2..2]); + assert_eq!(output.dims(), [2, 0]); +} + +#[test] +fn test_slice_with_positive_step() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data( + [ + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + ], + &device, + ); + + // Test step=2 along first dimension + let sliced = tensor.clone().slice([s![0..3;2]]); + let expected = TensorData::from([[1.0, 2.0, 3.0, 4.0], [9.0, 10.0, 11.0, 12.0]]); + sliced.into_data().assert_eq(&expected, false); + + // Test step=2 along second dimension + let sliced = tensor.clone().slice(s![.., 0..4;2]); + let expected = TensorData::from([[1.0, 3.0], [5.0, 7.0], [9.0, 11.0]]); + sliced.into_data().assert_eq(&expected, false); + + // Test step=2 along both dimensions + let sliced = tensor.clone().slice(s![0..3;2, 0..4;2]); + let expected = TensorData::from([[1.0, 3.0], [9.0, 11.0]]); + sliced.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_slice_with_negative_step() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data( + [ + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + ], + &device, + ); + + // Test step=-1 along first dimension (reverse rows) + let sliced = tensor.clone().slice([s![0..3;-1]]); + let expected = TensorData::from([ + [9.0, 10.0, 11.0, 12.0], + [5.0, 6.0, 7.0, 8.0], + [1.0, 2.0, 3.0, 4.0], + ]); + sliced.into_data().assert_eq(&expected, false); + + // Test step=-1 along second dimension (reverse columns) + let sliced = tensor.clone().slice(s![.., 0..4;-1]); + let expected = TensorData::from([ + [4.0, 3.0, 2.0, 1.0], + [8.0, 7.0, 6.0, 5.0], + [12.0, 11.0, 10.0, 9.0], + ]); + sliced.into_data().assert_eq(&expected, false); + + // Test step=-2 along first dimension + let sliced = tensor.clone().slice([s![0..3;-2]]); + let expected = TensorData::from([[9.0, 10.0, 11.0, 12.0], [1.0, 2.0, 3.0, 4.0]]); + sliced.into_data().assert_eq(&expected, false); + + // Test step=-2 along second dimension + let sliced = tensor.clone().slice(s![.., 0..4;-2]); + let expected = TensorData::from([[4.0, 2.0], [8.0, 6.0], [12.0, 10.0]]); + sliced.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_slice_with_mixed_steps() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data( + [ + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + ], + &device, + ); + + // Test positive step along first dimension, negative along second + let sliced = tensor.clone().slice(s![0..3;2, 0..4;-1]); + let expected = TensorData::from([[4.0, 3.0, 2.0, 1.0], [12.0, 11.0, 10.0, 9.0]]); + sliced.into_data().assert_eq(&expected, false); + + // Test negative step along first dimension, positive along second + let sliced = tensor.clone().slice(s![0..3;-1, 0..4;2]); + let expected = TensorData::from([[9.0, 11.0], [5.0, 7.0], [1.0, 3.0]]); + sliced.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_slice_with_steps_1d() { + let device = Default::default(); + let tensor = + TestTensor::<1>::from_data([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], &device); + + // Test positive step + let sliced = tensor.clone().slice([s![0..10;2]]); + let expected = TensorData::from([1.0, 3.0, 5.0, 7.0, 9.0]); + sliced.into_data().assert_eq(&expected, false); + + // Test negative step + let sliced = tensor.clone().slice([s![0..10;-1]]); + let expected = TensorData::from([10.0, 9.0, 8.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0]); + sliced.into_data().assert_eq(&expected, false); + + // Test negative step with partial range + let sliced = tensor.clone().slice([s![2..8;-2]]); + let expected = TensorData::from([8.0, 6.0, 4.0]); + sliced.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_slice_with_steps_3d() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data( + [ + [[1.0, 2.0], [3.0, 4.0]], + [[5.0, 6.0], [7.0, 8.0]], + [[9.0, 10.0], [11.0, 12.0]], + [[13.0, 14.0], [15.0, 16.0]], + ], + &device, + ); + + // Test step=2 along first dimension + let sliced = tensor.clone().slice(s![0..4;2, .., ..]); + let expected = TensorData::from([[[1.0, 2.0], [3.0, 4.0]], [[9.0, 10.0], [11.0, 12.0]]]); + sliced.into_data().assert_eq(&expected, false); + + // Test step=-1 along all dimensions + let sliced = tensor.clone().slice(s![0..4;-1, 0..2;-1, 0..2;-1]); + let expected = TensorData::from([ + [[16.0, 15.0], [14.0, 13.0]], + [[12.0, 11.0], [10.0, 9.0]], + [[8.0, 7.0], [6.0, 5.0]], + [[4.0, 3.0], [2.0, 1.0]], + ]); + sliced.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/slice_assign.rs b/crates/burn-backend-tests/tests/tensor/float/ops/slice_assign.rs new file mode 100644 index 0000000..8799f5d --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/slice_assign.rs @@ -0,0 +1,366 @@ +use super::*; +use burn_tensor::{Slice, TensorData, s}; + +#[test] +fn should_support_slice_assign_1d() { + let data = TensorData::from([0.0, 1.0, 2.0]); + let data_assigned = TensorData::from([10.0, 5.0]); + + let device = Default::default(); + let tensor = TestTensor::<1>::from_data(data, &device); + let tensor_assigned = TestTensor::<1>::from_data(data_assigned, &device); + + let output = tensor.slice_assign([0..2], tensor_assigned); + let expected = TensorData::from([10.0, 5.0, 2.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_slice_assign_2d() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let data_assigned = TensorData::from([[10.0, 5.0]]); + + let device = Default::default(); + let tensor = TestTensor::<2>::from_data(data, &device); + let tensor_assigned = TestTensor::<2>::from_data(data_assigned, &device); + + let output = tensor.slice_assign([1..2, 0..2], tensor_assigned); + let expected = TensorData::from([[0.0, 1.0, 2.0], [10.0, 5.0, 5.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_slice_assign_vec() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let data_assigned = TensorData::from([[10.0, 5.0]]); + + let device = Default::default(); + let tensor = TestTensor::<2>::from_data(data, &device); + let tensor_assigned = TestTensor::<2>::from_data(data_assigned, &device); + + let slices: Vec = vec![1..2, 0..2].into_iter().map(Slice::from).collect(); + + let output = tensor.slice_assign(&slices, tensor_assigned); + let expected = TensorData::from([[0.0, 1.0, 2.0], [10.0, 5.0, 5.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn slice_assign_now_supports_non_unit_step() { + let device = Default::default(); + // Create tensors where the shapes match for stepped slicing + let tensor = TestTensor::<2>::ones([4, 4], &device); + // With step=2 on first dim, we select indices 0 and 2, so we need a [2, 4] values tensor + let values = TestTensor::<2>::zeros([2, 4], &device); + + // This now works because slice_assign supports steps != 1 + // We use s! macro to create a slice with step=2 + let result = tensor.slice_assign(s![0..3;2, ..], values); + + // Verify the result: rows 0 and 2 should be zeros, rows 1 and 3 should be ones + let expected = TensorData::from([ + [0.0, 0.0, 0.0, 0.0], + [1.0, 1.0, 1.0, 1.0], + [0.0, 0.0, 0.0, 0.0], + [1.0, 1.0, 1.0, 1.0], + ]); + result.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_slice_assign_with_positive_step_1d() { + let device = Default::default(); + let tensor = TestTensor::<1>::from_data([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &device); + let values = TestTensor::<1>::from_data([10.0, 20.0, 30.0], &device); + + // Assign to indices 0, 2, 4 (step=2) + let output = tensor.slice_assign([s![0..6;2]], values); + let expected = TensorData::from([10.0, 2.0, 20.0, 4.0, 30.0, 6.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_slice_assign_with_positive_step_2d() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data( + [ + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + [13.0, 14.0, 15.0, 16.0], + ], + &device, + ); + + // Assign to rows 0, 2 (step=2) + let values = TestTensor::<2>::from_data( + [[100.0, 101.0, 102.0, 103.0], [200.0, 201.0, 202.0, 203.0]], + &device, + ); + let output = tensor.clone().slice_assign([s![0..4;2]], values); + let expected = TensorData::from([ + [100.0, 101.0, 102.0, 103.0], + [5.0, 6.0, 7.0, 8.0], + [200.0, 201.0, 202.0, 203.0], + [13.0, 14.0, 15.0, 16.0], + ]); + output.into_data().assert_eq(&expected, false); + + // Assign to columns 0, 2 (step=2) + let values = TestTensor::<2>::from_data( + [ + [100.0, 200.0], + [101.0, 201.0], + [102.0, 202.0], + [103.0, 203.0], + ], + &device, + ); + let output = tensor.clone().slice_assign(s![.., 0..4;2], values); + let expected = TensorData::from([ + [100.0, 2.0, 200.0, 4.0], + [101.0, 6.0, 201.0, 8.0], + [102.0, 10.0, 202.0, 12.0], + [103.0, 14.0, 203.0, 16.0], + ]); + output.into_data().assert_eq(&expected, false); + + // Assign with step=2 on both dimensions + let values = TestTensor::<2>::from_data([[100.0, 200.0], [300.0, 400.0]], &device); + let output = tensor.slice_assign(s![0..4;2, 0..4;2], values); + let expected = TensorData::from([ + [100.0, 2.0, 200.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [300.0, 10.0, 400.0, 12.0], + [13.0, 14.0, 15.0, 16.0], + ]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_slice_assign_with_negative_step_1d() { + let device = Default::default(); + let tensor = TestTensor::<1>::from_data([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &device); + let values = TestTensor::<1>::from_data([60.0, 50.0, 40.0, 30.0, 20.0, 10.0], &device); + + // Assign in reverse order (step=-1) + let output = tensor.slice_assign([s![0..6;-1]], values); + let expected = TensorData::from([10.0, 20.0, 30.0, 40.0, 50.0, 60.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_slice_assign_with_negative_step_2d() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data( + [ + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + ], + &device, + ); + + // Assign to rows in reverse order (step=-1) + let values = TestTensor::<2>::from_data( + [ + [30.0, 31.0, 32.0, 33.0], + [20.0, 21.0, 22.0, 23.0], + [10.0, 11.0, 12.0, 13.0], + ], + &device, + ); + let output = tensor.clone().slice_assign([s![0..3;-1]], values); + let expected = TensorData::from([ + [10.0, 11.0, 12.0, 13.0], + [20.0, 21.0, 22.0, 23.0], + [30.0, 31.0, 32.0, 33.0], + ]); + output.into_data().assert_eq(&expected, false); + + // Assign to columns in reverse order (step=-1) + let values = TestTensor::<2>::from_data( + [ + [40.0, 30.0, 20.0, 10.0], + [80.0, 70.0, 60.0, 50.0], + [120.0, 110.0, 100.0, 90.0], + ], + &device, + ); + let output = tensor.clone().slice_assign(s![.., 0..4;-1], values); + let expected = TensorData::from([ + [10.0, 20.0, 30.0, 40.0], + [50.0, 60.0, 70.0, 80.0], + [90.0, 100.0, 110.0, 120.0], + ]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_slice_assign_with_mixed_steps() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data( + [ + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + [13.0, 14.0, 15.0, 16.0], + ], + &device, + ); + + // Positive step along rows, negative along columns + let values = TestTensor::<2>::from_data( + [[100.0, 101.0, 102.0, 103.0], [200.0, 201.0, 202.0, 203.0]], + &device, + ); + let output = tensor.clone().slice_assign(s![0..4;2, 0..4;-1], values); + let expected = TensorData::from([ + [103.0, 102.0, 101.0, 100.0], + [5.0, 6.0, 7.0, 8.0], + [203.0, 202.0, 201.0, 200.0], + [13.0, 14.0, 15.0, 16.0], + ]); + output.into_data().assert_eq(&expected, false); + + // Negative step along rows, positive along columns + let values = TestTensor::<2>::from_data( + [ + [100.0, 200.0], + [101.0, 201.0], + [102.0, 202.0], + [103.0, 203.0], + ], + &device, + ); + let output = tensor.slice_assign(s![0..4;-1, 0..4;2], values); + let expected = TensorData::from([ + [103.0, 2.0, 203.0, 4.0], + [102.0, 6.0, 202.0, 8.0], + [101.0, 10.0, 201.0, 12.0], + [100.0, 14.0, 200.0, 16.0], + ]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_slice_assign_3d_with_steps() { + let device = Default::default(); + let tensor = TestTensor::<3>::from_data( + [ + [[1.0, 2.0], [3.0, 4.0]], + [[5.0, 6.0], [7.0, 8.0]], + [[9.0, 10.0], [11.0, 12.0]], + [[13.0, 14.0], [15.0, 16.0]], + ], + &device, + ); + + // Test step=2 along first dimension + let values = TestTensor::<3>::from_data( + [ + [[100.0, 101.0], [102.0, 103.0]], + [[200.0, 201.0], [202.0, 203.0]], + ], + &device, + ); + let output = tensor.clone().slice_assign(s![0..4;2, .., ..], values); + let expected = TensorData::from([ + [[100.0, 101.0], [102.0, 103.0]], + [[5.0, 6.0], [7.0, 8.0]], + [[200.0, 201.0], [202.0, 203.0]], + [[13.0, 14.0], [15.0, 16.0]], + ]); + output.into_data().assert_eq(&expected, false); + + // Test step=-1 along all dimensions + let values = TestTensor::<3>::from_data( + [ + [[400.0, 399.0], [398.0, 397.0]], + [[396.0, 395.0], [394.0, 393.0]], + [[392.0, 391.0], [390.0, 389.0]], + [[388.0, 387.0], [386.0, 385.0]], + ], + &device, + ); + let output = tensor.slice_assign(s![0..4;-1, 0..2;-1, 0..2;-1], values); + let expected = TensorData::from([ + [[385.0, 386.0], [387.0, 388.0]], + [[389.0, 390.0], [391.0, 392.0]], + [[393.0, 394.0], [395.0, 396.0]], + [[397.0, 398.0], [399.0, 400.0]], + ]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_slice_assign_partial_with_steps() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data( + [ + [1.0, 2.0, 3.0, 4.0, 5.0], + [6.0, 7.0, 8.0, 9.0, 10.0], + [11.0, 12.0, 13.0, 14.0, 15.0], + [16.0, 17.0, 18.0, 19.0, 20.0], + [21.0, 22.0, 23.0, 24.0, 25.0], + ], + &device, + ); + + // Assign to a subset with step=2 + let values = TestTensor::<2>::from_data([[100.0, 200.0], [300.0, 400.0]], &device); + let output = tensor.slice_assign(s![1..4;2, 1..4;2], values); + let expected = TensorData::from([ + [1.0, 2.0, 3.0, 4.0, 5.0], + [6.0, 100.0, 8.0, 200.0, 10.0], + [11.0, 12.0, 13.0, 14.0, 15.0], + [16.0, 300.0, 18.0, 400.0, 20.0], + [21.0, 22.0, 23.0, 24.0, 25.0], + ]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_slice_assign_empty_range() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + let values: TestTensor<2> = TestTensor::empty([2, 0], &device); + + // Empty slice assignment (start == end) should be a no-op + let output = tensor.clone().slice_assign([0..2, 1..1], values); + let expected = TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_slice_assign_empty_range_1d() { + let device = Default::default(); + let tensor = TestTensor::<1>::from_data([1.0, 2.0, 3.0, 4.0, 5.0], &device); + let values: TestTensor<1> = TestTensor::empty([0], &device); + + // Empty slice assignment should return tensor unchanged + let output = tensor.clone().slice_assign([2..2], values); + let expected = TensorData::from([1.0, 2.0, 3.0, 4.0, 5.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_slice_assign_single_dim_slice() { + let device = Default::default(); + let x = TestTensor::<3>::ones([2, 3, 1], &device); + let values = TestTensor::<3>::zeros([1, 3, 1], &device); + + let output = x.slice_assign(s![1], values); + + output.into_data().assert_eq( + &TensorData::from([[[1.0], [1.0], [1.0]], [[0.0], [0.0], [0.0]]]), + false, + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/sort_argsort.rs b/crates/burn-backend-tests/tests/tensor/float/ops/sort_argsort.rs new file mode 100644 index 0000000..7da93cd --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/sort_argsort.rs @@ -0,0 +1,216 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn test_sort_1d_float() { + let tensor = TestTensor::<1>::from([ + 0.5, 1.2, -0.21, 0., 2.1, 0.94, -0.3, 2.3, 199.412, 4., 0.99, 3., -8.1, + ]); + + // Sort along dim=0 + let values = tensor.sort(0); + + let values_expected = TensorData::from([ + -8.1, -0.3, -0.21, 0., 0.5, 0.94, 0.99, 1.2, 2.1, 2.3, 3., 4., 199.412, + ]); + values + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::default()); +} + +#[test] +fn test_argsort_1d_float() { + let tensor = TestTensor::<1>::from([ + 0.5, 1.2, -0.21, 0., 2.1, 0.94, -0.3, 2.3, 199.412, 4., 0.99, 3., -8.1, + ]); + + // Sort along dim=0 + let indices = tensor.argsort(0); + + let indices_expected = TensorData::from([12, 6, 2, 3, 0, 5, 10, 1, 4, 7, 11, 9, 8]); + indices.into_data().assert_eq(&indices_expected, false); +} + +#[test] +fn test_sort_with_indices_descending_float() { + // 1D + let tensor = TestTensor::<1>::from([ + 0.5, 1.2, -0.21, 0., 2.1, 0.94, -0.3, 2.3, 199.412, 4., 0.99, 3., -8.1, + ]); + + // Sort along dim=0 + let (values, indices) = tensor.sort_descending_with_indices(0); + + let values_expected = TensorData::from([ + 199.412, 4., 3., 2.3, 2.1, 1.2, 0.99, 0.94, 0.5, 0., -0.21, -0.3, -8.1, + ]); + values + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::default()); + + let indices_expected = TensorData::from([8, 9, 11, 7, 4, 1, 10, 5, 0, 3, 2, 6, 12]); + indices.into_data().assert_eq(&indices_expected, false); + + // 2D + let tensor = TestTensor::<3>::from([ + [[-0.5, 1.2, -0.21], [0., 2.1, 0.94]], + [[-0.3, 2.3, 4.], [0.99, 3., -8.1]], + ]); + + // Sort along dim=1 + let (values, indices) = tensor.sort_descending_with_indices(1); + + let values_expected = TensorData::from([ + [[0., 2.1, 0.94], [-0.5, 1.2, -0.21]], + [[0.99, 3., 4.], [-0.3, 2.3, -8.1]], + ]); + values + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::default()); + + let indices_expected = TensorData::from([[[1, 1, 1], [0, 0, 0]], [[1, 1, 0], [0, 0, 1]]]); + indices.into_data().assert_eq(&indices_expected, false); +} + +#[test] +fn test_sort_float() { + let tensor = TestTensor::<3>::from([ + [[-0.5, 1.2, -0.21], [0., 2.1, 0.94]], + [[-0.3, 2.3, 4.], [0.99, 3., -8.1]], + ]); + + // Sort along dim=0 + let values = tensor.clone().sort(0); + + let values_expected = TensorData::from([ + [[-0.5, 1.2, -0.21], [0., 2.1, -8.1]], + [[-0.3, 2.3, 4.], [0.99, 3., 0.94]], + ]); + values + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::default()); + + // Sort along dim=1 + let values = tensor.clone().sort(1); + + let values_expected = TensorData::from([ + [[-0.5, 1.2, -0.21], [0., 2.1, 0.94]], + [[-0.3, 2.3, -8.1], [0.99, 3., 4.]], + ]); + values + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::default()); + + // Sort along dim=2 + let values = tensor.sort(2); + + let values_expected = TensorData::from([ + [[-0.5, -0.21, 1.2], [0., 0.94, 2.1]], + [[-0.3, 2.3, 4.], [-8.1, 0.99, 3.]], + ]); + values + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::default()); +} + +#[test] +fn test_sort_with_indices_float() { + let tensor = TestTensor::<3>::from([ + [[-0.5, 1.2, -0.21], [0., 2.1, 0.94]], + [[-0.3, 2.3, 4.], [0.99, 3., -8.1]], + ]); + + // Sort along dim=0 + let (values, indices) = tensor.clone().sort_with_indices(0); + let values_expected = TensorData::from([ + [[-0.5, 1.2, -0.21], [0., 2.1, -8.1]], + [[-0.3, 2.3, 4.], [0.99, 3., 0.94]], + ]); + values + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::default()); + + let indices_expected = TensorData::from([[[0, 0, 0], [0, 0, 1]], [[1, 1, 1], [1, 1, 0]]]); + indices.into_data().assert_eq(&indices_expected, false); + + // Sort along dim=1 + let (values, indices) = tensor.clone().sort_with_indices(1); + + let values_expected = TensorData::from([ + [[-0.5, 1.2, -0.21], [0., 2.1, 0.94]], + [[-0.3, 2.3, -8.1], [0.99, 3., 4.]], + ]); + values + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::default()); + + let indices_expected = TensorData::from([[[0, 0, 0], [1, 1, 1]], [[0, 0, 1], [1, 1, 0]]]); + indices.into_data().assert_eq(&indices_expected, false); + + // Sort along dim=2 + let (values, indices) = tensor.sort_with_indices(2); + + let values_expected = TensorData::from([ + [[-0.5, -0.21, 1.2], [0., 0.94, 2.1]], + [[-0.3, 2.3, 4.], [-8.1, 0.99, 3.]], + ]); + values + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::default()); + + let indices_expected = TensorData::from([[[0, 2, 1], [0, 2, 1]], [[0, 1, 2], [2, 0, 1]]]); + indices.into_data().assert_eq(&indices_expected, false); +} + +#[test] +fn test_argsort_float() { + let tensor = TestTensor::<3>::from([ + [[-0.5, 1.2, -0.21], [0., 2.1, 0.94]], + [[-0.3, 2.3, 4.], [0.99, 3., -8.1]], + ]); + + // Sort along dim=0 + let indices = tensor.clone().argsort(0); + + let indices_expected = TensorData::from([[[0, 0, 0], [0, 0, 1]], [[1, 1, 1], [1, 1, 0]]]); + indices.into_data().assert_eq(&indices_expected, false); + + // Sort along dim=1 + let indices = tensor.clone().argsort(1); + + let indices_expected = TensorData::from([[[0, 0, 0], [1, 1, 1]], [[0, 0, 1], [1, 1, 0]]]); + indices.into_data().assert_eq(&indices_expected, false); + + // Sort along dim=2 + let indices = tensor.argsort(2); + + let indices_expected = TensorData::from([[[0, 2, 1], [0, 2, 1]], [[0, 1, 2], [2, 0, 1]]]); + indices.into_data().assert_eq(&indices_expected, false); +} + +#[test] +fn test_sort_float_nan() { + let tensor = TestTensor::<2>::from([[-0.5, f32::NAN], [0., 0.94], [-0.3, f32::NAN]]); + + // Sort along dim=0 + let values = tensor.sort(0); + + let values_expected = TensorData::from([[-0.5, 0.94], [-0.3, f32::NAN], [0., f32::NAN]]); + values + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::default()); +} + +#[test] +fn test_sort_descending_1d() { + let tensor = TestTensor::<1>::from([1., 2., 3., 4., 5.]); + + // Sort along dim=0 + let values = tensor.sort_descending(0); + + let values_expected = TensorData::from([5., 4., 3., 2., 1.]); + values + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/split.rs b/crates/burn-backend-tests/tests/tensor/float/ops/split.rs new file mode 100644 index 0000000..7607b75 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/split.rs @@ -0,0 +1,206 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_split_evenly_divisible() { + let device = Default::default(); + let tensors = + TestTensor::<2>::from_data([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11]], &device); + + let split_tensors = tensors.split(2, 0); + assert_eq!(split_tensors.len(), 3); + + let expected = [ + TensorData::from([[0, 1], [2, 3]]), + TensorData::from([[4, 5], [6, 7]]), + TensorData::from([[8, 9], [10, 11]]), + ]; + + for (index, tensor) in split_tensors.iter().enumerate() { + tensor.to_data().assert_eq(&expected[index], false); + } +} + +#[test] +fn test_split_not_evenly_divisible() { + let device = Default::default(); + let tensors = TestTensor::<2>::from_data([[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]], &device); + + let split_tensors = tensors.split(2, 0); + assert_eq!(split_tensors.len(), 3); + + let expected = [ + TensorData::from([[0, 1], [2, 3]]), + TensorData::from([[4, 5], [6, 7]]), + TensorData::from([[8, 9]]), + ]; + + for (index, tensor) in split_tensors.iter().enumerate() { + tensor.to_data().assert_eq(&expected[index], false); + } +} + +#[test] +fn test_split_along_dim1() { + let device = Default::default(); + let tensors = TestTensor::<2>::from_data([[0, 1, 2], [3, 4, 5]], &device); + + let split_tensors = tensors.split(2, 1); + assert_eq!(split_tensors.len(), 2); + + let expected = [ + TensorData::from([[0, 1], [3, 4]]), + TensorData::from([[2], [5]]), + ]; + + for (index, tensor) in split_tensors.iter().enumerate() { + tensor.to_data().assert_eq(&expected[index], false); + } +} + +#[test] +fn test_split_split_size_larger_than_tensor_size() { + let device = Default::default(); + let tensors = TestTensor::<1>::from_data([0, 1, 2, 3, 4], &device); + + let split_tensors = tensors.split(10, 0); + assert_eq!(split_tensors.len(), 1); + + let expected = [TensorData::from([0, 1, 2, 3, 4])]; + + for (index, tensor) in split_tensors.iter().enumerate() { + tensor.to_data().assert_eq(&expected[index], false); + } +} + +#[test] +fn test_split_with_zero_split_size_zero_tensor_size() { + let device = Default::default(); + let empty_array: [i32; 0] = []; + let tensors = TestTensor::<1>::from_data(empty_array, &device); + + let split_tensors = tensors.split(0, 0); + assert_eq!(split_tensors.len(), 0); +} + +#[test] +fn test_split_zero_sized_tensor() { + let device = Default::default(); + let empty_array: [i32; 0] = []; + let tensors = TestTensor::<1>::from_data(empty_array, &device); + + let split_tensors = tensors.split(1, 0); + assert_eq!(split_tensors.len(), 0); +} + +#[test] +#[should_panic( + expected = "split_size must be greater than 0 unless the tensor size along the dimension is 0." +)] +fn test_split_with_zero_split_size_non_zero_tensor() { + let device = Default::default(); + let tensors = TestTensor::<1>::from_data([0, 1, 2, 3, 4], &device); + + let _split_tensors = tensors.split(0, 0); +} + +#[test] +#[should_panic(expected = "Given dimension is greater than or equal to the tensor rank.")] +fn test_split_invalid_dim() { + let device = Default::default(); + let tensors = TestTensor::<1>::from_data([0, 1, 2], &device); + + let _split_tensors = tensors.split(1, 2); +} + +#[test] +fn test_split_3d_tensor_along_dim0() { + let device = Default::default(); + let tensors = TestTensor::<3>::from_data( + [ + [[0, 1], [2, 3]], + [[4, 5], [6, 7]], + [[8, 9], [10, 11]], + [[12, 13], [14, 15]], + ], + &device, + ); + + let split_tensors = tensors.split(2, 0); + assert_eq!(split_tensors.len(), 2); + + let expected = [ + TensorData::from([[[0, 1], [2, 3]], [[4, 5], [6, 7]]]), + TensorData::from([[[8, 9], [10, 11]], [[12, 13], [14, 15]]]), + ]; + + for (index, tensor) in split_tensors.iter().enumerate() { + tensor.to_data().assert_eq(&expected[index], false); + } +} + +#[test] +fn test_split_3d_tensor_along_dim1() { + let device = Default::default(); + let tensors = TestTensor::<3>::from_data( + [[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]]], + &device, + ); + + let split_tensors = tensors.split(2, 1); + assert_eq!(split_tensors.len(), 2); + + let expected = [ + TensorData::from([[[0, 1], [2, 3]], [[6, 7], [8, 9]]]), + TensorData::from([[[4, 5]], [[10, 11]]]), + ]; + + for (index, tensor) in split_tensors.iter().enumerate() { + tensor.to_data().assert_eq(&expected[index], false); + } +} + +#[test] +fn test_split_with_sizes() { + let device = Default::default(); + let tensors = TestTensor::<1>::from_data([0, 1, 2, 3, 4, 5], &device); + + let split_tensors = tensors.split_with_sizes(vec![2, 3, 1], 0); + assert_eq!(split_tensors.len(), 3); + + let expected = [ + TensorData::from([0, 1]), + TensorData::from([2, 3, 4]), + TensorData::from([5]), + ]; + + for (index, tensor) in split_tensors.iter().enumerate() { + tensor.to_data().assert_eq(&expected[index], false); + } +} + +#[test] +#[should_panic( + expected = "The sum of split_sizes must equal the tensor size along the specified dimension." +)] +fn test_split_with_sizes_invalid_sum() { + let device = Default::default(); + let tensors = TestTensor::<1>::from_data([0, 1, 2, 3, 4, 5], &device); + + let _split_tensors = tensors.split_with_sizes(vec![2, 2, 1], 0); +} + +#[test] +fn test_split_with_sizes_zero_length() { + let device = Default::default(); + let tensors = TestTensor::<1>::from_data([0, 1, 2], &device); + + let split_tensors = tensors.split_with_sizes(vec![0, 1, 2], 0); + assert_eq!(split_tensors.len(), 2); + + let expected = [TensorData::from([0]), TensorData::from([1, 2])]; + + for (index, tensor) in split_tensors.iter().enumerate() { + tensor.to_data().assert_eq(&expected[index], false); + } +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/sqrt.rs b/crates/burn-backend-tests/tests/tensor/float/ops/sqrt.rs new file mode 100644 index 0000000..05f0153 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/sqrt.rs @@ -0,0 +1,46 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; +use core::f32::consts::SQRT_2; + +#[test] +fn should_support_sqrt_narrowed() { + // [1, 4, 9, 16, 25, 36] narrowed to [4, 9, 16, 25] + let tensor = TestTensor::<1>::from([1.0, 4.0, 9.0, 16.0, 25.0, 36.0]); + let narrowed = tensor.narrow(0, 1, 4); + + let output = narrowed.sqrt(); + let expected = TensorData::from([2.0, 3.0, 4.0, 5.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_sqrt_flipped_2d() { + // [[1, 4], [9, 16]] flipped on axis 0 -> [[9, 16], [1, 4]] + let tensor = TestTensor::<2>::from([[1.0, 4.0], [9.0, 16.0]]); + let flipped = tensor.flip([0]); + + let output = flipped.sqrt(); + let expected = TensorData::from([[3.0, 4.0], [1.0, 2.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_sqrt_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.sqrt(); + let expected = TensorData::from([[0.0, 1.0, SQRT_2], [1.73205, 2.0, 2.2360]]); + + output.into_data().assert_approx_eq::( + &expected, + Tolerance::relative(1e-4).set_half_precision_relative(1e-3), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/square.rs b/crates/burn-backend-tests/tests/tensor/float/ops/square.rs new file mode 100644 index 0000000..994282b --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/square.rs @@ -0,0 +1,17 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_sqrt_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.square(); + let expected = TensorData::from([[0.0, 1.0, 4.0], [9.0, 16.0, 25.0]]); + + output.into_data().assert_approx_eq::( + &expected, + Tolerance::relative(1e-4).set_half_precision_relative(1e-3), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/squeeze.rs b/crates/burn-backend-tests/tests/tensor/float/ops/squeeze.rs new file mode 100644 index 0000000..f432310 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/squeeze.rs @@ -0,0 +1,263 @@ +use super::*; +use burn_tensor::Shape; + +/// Test if the function can successfully squeeze the size 1 dimension of a 3D tensor. +#[test] +fn should_squeeze_dim() { + let tensor = TestTensor::<3>::ones(Shape::new([2, 1, 4]), &Default::default()); + let squeezed_tensor: TestTensor<2> = tensor.squeeze_dim(1); + let expected_shape = Shape::new([2, 4]); + assert_eq!(squeezed_tensor.shape(), expected_shape); +} + +#[test] +fn should_squeeze() { + let tensor = TestTensor::<3>::ones(Shape::new([2, 1, 4]), &Default::default()); + let squeezed_tensor: TestTensor<2> = tensor.squeeze(); + let expected_shape = Shape::new([2, 4]); + assert_eq!(squeezed_tensor.shape(), expected_shape); +} + +/// Test if the function can successfully squeeze the first size 1 dimension of a 4D tensor. +#[test] +fn should_squeeze_first() { + let tensor = TestTensor::<4>::ones(Shape::new([1, 3, 4, 5]), &Default::default()); + let squeezed_tensor: TestTensor<3> = tensor.squeeze_dim(0); + let expected_shape = Shape::new([3, 4, 5]); + assert_eq!(squeezed_tensor.shape(), expected_shape); +} +/// Test if the function can successfully squeeze the last size 1 dimension of a 4D tensor. +#[test] +fn should_squeeze_last() { + let tensor = TestTensor::<4>::ones(Shape::new([2, 3, 4, 1]), &Default::default()); + let squeezed_tensor: TestTensor<3> = tensor.squeeze_dim(3); + let expected_shape = Shape::new([2, 3, 4]); + assert_eq!(squeezed_tensor.shape(), expected_shape); +} +/// Test if the function panics when the squeezed dimension is not of size 1. +#[test] +#[should_panic] +fn should_squeeze_panic() { + let tensor = TestTensor::<4>::ones(Shape::new([2, 3, 4, 5]), &Default::default()); + let _squeezed_tensor: TestTensor<3> = tensor.squeeze_dim(2); +} + +/// Test if the function works with an empty slice +#[test] +fn should_squeeze_dims_with_empty_slice() { + let tensor = TestTensor::<3>::ones(Shape::new([1, 1, 3]), &Default::default()); + let squeezed_tensor: TestTensor<1> = tensor.squeeze_dims(&[]); + let expected_shape = Shape::new([3]); + assert_eq!(squeezed_tensor.shape(), expected_shape); +} + +#[test] +fn should_squeeze_all_dims() { + let tensor = TestTensor::<3>::ones(Shape::new([1, 3, 1]), &Default::default()); + let squeezed_tensor: TestTensor<1> = tensor.squeeze(); + let expected_shape = Shape::new([3]); + assert_eq!(squeezed_tensor.shape(), expected_shape); +} + +/// Test if the function works with positive indices +#[test] +fn should_squeeze_dims_with_positive_indices() { + let tensor = TestTensor::<4>::ones(Shape::new([1, 3, 1, 5]), &Default::default()); + let squeezed_tensor: TestTensor<2> = tensor.squeeze_dims(&[0, 2]); + let expected_shape = Shape::new([3, 5]); + assert_eq!(squeezed_tensor.shape(), expected_shape); +} + +/// Test if the function works with negative indices +#[test] +fn should_squeeze_dims_with_negative_indices() { + let tensor = TestTensor::<4>::ones(Shape::new([2, 1, 3, 1]), &Default::default()); + let squeezed_tensor: TestTensor<2> = tensor.squeeze_dims(&[-3, -1]); + let expected_shape = Shape::new([2, 3]); + assert_eq!(squeezed_tensor.shape(), expected_shape); +} + +/// Test to make sure the function panics if a non-singleton dimension is squeezed +#[test] +#[should_panic] +fn should_squeeze_dims_work_if_non_singleton() { + let tensor = TestTensor::<3>::ones(Shape::new([2, 3, 4]), &Default::default()); + let squeezed_tensor: TestTensor<3> = tensor.squeeze_dims(&[1]); + let expected_shape = Shape::new([2, 3, 4]); + assert_eq!(squeezed_tensor.shape(), expected_shape); +} + +#[test] +#[should_panic] +fn should_panic_squeeze_consumes_all_singleton() { + let tensor = TestTensor::<3>::ones(Shape::new([1, 3, 1]), &Default::default()); + let _squeezed_tensor: TestTensor<2> = tensor.squeeze(); // output rank should be 1 +} + +/// Test to make sure the function panics if too many dimensions are requested to be squeezed +#[test] +#[should_panic] +fn should_squeeze_dims_panic_on_too_many_dimensions() { + let tensor = TestTensor::<3>::ones(Shape::new([1, 1, 1]), &Default::default()); + let _: TestTensor<1> = tensor.squeeze_dims(&[0, 1, 2]); +} + +/// Test to make sure function panics if dimensions are mismatched +#[test] +#[should_panic] +fn should_squeeze_dims_dimension_mismatch_panic() { + let tensor = TestTensor::<4>::ones(Shape::new([1, 3, 1, 5]), &Default::default()); + let _: TestTensor<3> = tensor.squeeze_dims(&[0, 2]); +} + +/// Test if the function can successfully unsqueeze the size 1 dimension at the specified position of a 3D tensor. +#[test] +fn should_unsqueeze_dim() { + let tensor = TestTensor::<3>::ones(Shape::new([2, 4, 1]), &Default::default()); + let unsqueezed_tensor: TestTensor<4> = tensor.unsqueeze_dim(1); + let expected_shape = Shape::new([2, 1, 4, 1]); + assert_eq!(unsqueezed_tensor.shape(), expected_shape); +} + +/// Test if the function can successfully unsqueeze the first size 1 dimension of a 4D tensor. +#[test] +fn should_unsqueeze_dim_first() { + let tensor = TestTensor::<4>::ones(Shape::new([2, 3, 4, 5]), &Default::default()); + let unsqueezed_tensor: TestTensor<5> = tensor.unsqueeze_dim(0); + let expected_shape = Shape::new([1, 2, 3, 4, 5]); + assert_eq!(unsqueezed_tensor.shape(), expected_shape); +} + +/// Test if the function can successfully unsqueeze the last size 1 dimension of a 4D tensor. +#[test] +fn should_unsqueeze_dim_last() { + let tensor = TestTensor::<4>::ones(Shape::new([5, 4, 3, 2]), &Default::default()); + let unsqueezed_tensor: TestTensor<5> = tensor.unsqueeze_dim(4); + let expected_shape = Shape::new([5, 4, 3, 2, 1]); + assert_eq!(unsqueezed_tensor.shape(), expected_shape); +} + +/// Test if the function panics when the unsqueezed dimension is out of bounds. +#[test] +#[should_panic] +fn should_unsqueeze_dim_panic() { + let tensor = TestTensor::<4>::ones(Shape::new([2, 3, 4, 5]), &Default::default()); + let _unsqueezed_tensor: TestTensor<5> = tensor.unsqueeze_dim(5); +} + +#[test] +fn should_unsqueeze_dims_support_dim_inference() { + let input_tensor = TestTensor::<3>::ones(Shape::new([3, 4, 5]), &Default::default()); + let output_tensor = input_tensor.unsqueeze_dims::<5>(&[1, -2]); + let expected_shape = Shape::new([3, 1, 4, 1, 5]); + assert_eq!(output_tensor.shape(), expected_shape); +} + +#[test] +fn should_unsqueeze_dims_handle_first_last() { + let input_tensor = TestTensor::<3>::ones(Shape::new([3, 4, 5]), &Default::default()); + let output_tensor = input_tensor.unsqueeze_dims::<5>(&[0, 4]); + let expected_shape = Shape::new([1, 3, 4, 5, 1]); + assert_eq!(output_tensor.shape(), expected_shape); +} + +#[test] +fn should_unsqueeze_dims_work_with_single_dim() { + //bruh, just call unsqueeze_dim + let input_tensor = TestTensor::<3>::ones(Shape::new([3, 4, 5]), &Default::default()); + let output_tensor: TestTensor<4> = input_tensor.unsqueeze_dims(&[1]); + let expected_shape = Shape::new([3, 1, 4, 5]); + assert_eq!(output_tensor.shape(), expected_shape); +} + +#[test] +fn should_unsqueeze_dims_multiple_trailing_negatives() { + let input_tensor = TestTensor::<3>::ones(Shape::new([3, 4, 5]), &Default::default()); + let output_tensor: TestTensor<6> = input_tensor.unsqueeze_dims(&[0, -1, -1]); + let expected_shape = Shape::new([1, 3, 4, 5, 1, 1]); + assert_eq!(output_tensor.shape(), expected_shape); +} + +#[test] +fn should_unsqueeze_dims_consecutive_first_then_gap() { + let input_tensor = TestTensor::<3>::ones(Shape::new([3, 4, 5]), &Default::default()); + let output_tensor: TestTensor<6> = input_tensor.unsqueeze_dims(&[0, 1, 3]); + let expected_shape = Shape::new([1, 1, 3, 1, 4, 5]); + assert_eq!(output_tensor.shape(), expected_shape); +} + +/// Regression test: `unsqueeze_dims` previously panicked with "index out of bounds" +/// when the axes produced duplicate values after sorting. The documented semantic +/// is that N duplicate axes insert N dims at that index; the fix normalizes sorted +/// duplicates so the N insertions occupy N consecutive output positions. +/// +/// This pattern comes up in ONNX matmul broadcasting when a 1D vector is unsqueezed +/// to match a higher-rank matrix (e.g. `[4]` -> `[1, 1, 4, 1]` for `4D @ 1D` matmul). +#[test] +fn should_unsqueeze_dims_handle_duplicate_axes_after_sort() { + // Matches the codegen pattern from burn-onnx for 4D @ 1D matmul broadcasting: + // axes `[-1, 0, 0]` convert to `[3, 0, 0]`, sort to `[0, 0, 3]`, and should + // produce `[1, 1, 4, 1]` after normalization. + let input_tensor = TestTensor::<1>::ones(Shape::new([4]), &Default::default()); + let output_tensor: TestTensor<4> = input_tensor.unsqueeze_dims(&[-1, 0, 0]); + let expected_shape = Shape::new([1, 1, 4, 1]); + assert_eq!(output_tensor.shape(), expected_shape); +} + +#[test] +fn should_unsqueeze_dims_handle_repeated_zero_axes() { + // Pure positive duplicates: `[0, 0]` should insert two leading 1s. + let input_tensor = TestTensor::<1>::ones(Shape::new([5]), &Default::default()); + let output_tensor: TestTensor<3> = input_tensor.unsqueeze_dims(&[0, 0]); + let expected_shape = Shape::new([1, 1, 5]); + assert_eq!(output_tensor.shape(), expected_shape); +} + +#[test] +fn should_unsqueeze_dims_handle_repeated_middle_axes() { + // Duplicates at an interior position: `[1, 1, 1]` should place three 1s + // starting at index 1, keeping the original dim at index 0. + let input_tensor = TestTensor::<1>::ones(Shape::new([5]), &Default::default()); + let output_tensor: TestTensor<4> = input_tensor.unsqueeze_dims(&[1, 1, 1]); + let expected_shape = Shape::new([5, 1, 1, 1]); + assert_eq!(output_tensor.shape(), expected_shape); +} + +#[test] +#[should_panic] +fn should_unsqueeze_dims_panic() { + let input_tensor = TestTensor::<3>::ones(Shape::new([3, 4, 5]), &Default::default()); + let _output_tensor: TestTensor<5> = input_tensor.unsqueeze_dims(&[0, -6]); +} + +/// Duplicate-axis normalization can push the last index past the valid output +/// range (e.g. `[2, 2]` targeting rank 3 normalizes to `[2, 3]`). This must be +/// rejected with a clear error instead of triggering an out-of-bounds read in +/// the copy loop. +#[test] +#[should_panic] +fn should_unsqueeze_dims_panic_duplicate_pushed_out_of_range() { + let input_tensor = TestTensor::<1>::ones(Shape::new([4]), &Default::default()); + let _: TestTensor<3> = input_tensor.unsqueeze_dims(&[2, 2]); +} + +#[test] +#[should_panic] +fn squeeze_all_singleton_not_supported() { + let tensor = TestTensor::<3>::ones(Shape::new([1, 1, 1]), &Default::default()); + let _ = tensor.squeeze::<0>(); +} + +#[test] +#[should_panic] +fn squeeze_dim_singleton_not_supported() { + let tensor = TestTensor::<1>::ones(Shape::new([1]), &Default::default()); + let _ = tensor.squeeze_dim::<0>(0); +} + +#[test] +#[should_panic] +fn squeeze_dims_all_singleton_not_supported() { + let tensor = TestTensor::<3>::ones(Shape::new([1, 1, 1]), &Default::default()); + let _ = tensor.squeeze_dims::<0>(&[0, 1, 2]); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/stack.rs b/crates/burn-backend-tests/tests/tensor/float/ops/stack.rs new file mode 100644 index 0000000..63d8c85 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/stack.rs @@ -0,0 +1,69 @@ +use super::*; +use alloc::{vec, vec::Vec}; +use burn_tensor::{Tensor, TensorData}; + +#[test] +fn should_support_stack_ops_2d_dim0() { + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data([[1.0, 2.0, 3.0]], &device); + let tensor_2 = TestTensor::from_data([[4.0, 5.0, 6.0]], &device); + + let output = Tensor::stack::<3>(vec![tensor_1, tensor_2], 0); + let expected = TensorData::from([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_stack_ops_2d_dim1() { + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data([[1.0, 2.0, 3.0]], &device); + let tensor_2 = TestTensor::from_data([[4.0, 5.0, 6.0]], &device); + + let output = Tensor::stack::<3>(vec![tensor_1, tensor_2], 1); + let expected = TensorData::from([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_stack_ops_3d() { + let device = Default::default(); + let tensor_1 = TestTensor::<3>::from_data([[[1.0, 2.0, 3.0]], [[1.1, 2.1, 3.1]]], &device); + let tensor_2 = TestTensor::from_data([[[4.0, 5.0, 6.0]], [[4.1, 5.1, 6.1]]], &device); + + let output = Tensor::stack::<4>(vec![tensor_1, tensor_2], 0); + let expected = TensorData::from([ + [[[1.0000, 2.0000, 3.0000]], [[1.1000, 2.1000, 3.1000]]], + [[[4.0000, 5.0000, 6.0000]], [[4.1000, 5.1000, 6.1000]]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +#[should_panic] +fn should_panic_when_dimensions_are_not_the_same() { + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]], &device); + let tensor_2 = TestTensor::from_data([[4.0, 5.0]], &device); + + let _output = Tensor::stack::<3>(vec![tensor_1, tensor_2], 0); +} + +#[test] +#[should_panic] +fn should_panic_when_list_of_vectors_is_empty() { + let tensors: Vec> = vec![]; + let _output = Tensor::stack::<3>(tensors, 0); +} + +#[test] +#[should_panic] +fn should_panic_when_stack_exceeds_dimension() { + let device = Default::default(); + let tensor_1 = TestTensor::<3>::from_data([[[1.0, 2.0, 3.0]], [[1.1, 2.1, 3.1]]], &device); + let tensor_2 = TestTensor::from_data([[[4.0, 5.0, 6.0]]], &device); + + let _output = Tensor::stack::<4>(vec![tensor_1, tensor_2], 3); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/sub.rs b/crates/burn-backend-tests/tests/tensor/float/ops/sub.rs new file mode 100644 index 0000000..d390caf --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/sub.rs @@ -0,0 +1,42 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_sub_ops() { + let data_1 = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let data_2 = TensorData::from([[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]]); + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device); + + let output = tensor_1 - tensor_2; + let expected = TensorData::from([[-6.0, -6.0, -6.0], [-6.0, -6.0, -6.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_sub_broadcast() { + let data_1 = TensorData::from([[0.0, 1.0, 2.0]]); + let data_2 = TensorData::from([[3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]); + let device = Default::default(); + let tensor_1 = TestTensor::<2>::from_data(data_1, &device); + let tensor_2 = TestTensor::<2>::from_data(data_2, &device); + + let output = tensor_1 - tensor_2; + let expected = TensorData::from([[-3.0, -3.0, -3.0], [-6.0, -6.0, -6.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_sub_scalar_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let scalar = 2.0; + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor - scalar; + let expected = TensorData::from([[-2.0, -1.0, 0.0], [1.0, 2.0, 3.0]]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/take.rs b/crates/burn-backend-tests/tests/tensor/float/ops/take.rs new file mode 100644 index 0000000..db27d77 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/take.rs @@ -0,0 +1,206 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_take_1d() { + // Test that take works with 1D indices + let device = Default::default(); + let tensor = TestTensor::<1>::from_data([0.0, 1.0, 2.0], &device); + let indices = TestTensorInt::<1>::from_data([1, 1, 0, 1, 2], &device); + + let output = tensor.take::<1, 1>(0, indices); + let expected = TensorData::from([1.0, 1.0, 0.0, 1.0, 2.0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_take_2d_dim0() { + // Test take on 2D tensor along dimension 0 + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &device); + let indices = TestTensorInt::<1>::from_data([1, 0, 1, 1], &device); + + let output = tensor.take::<1, 2>(0, indices); + let expected = TensorData::from([ + [3.0, 4.0, 5.0], + [0.0, 1.0, 2.0], + [3.0, 4.0, 5.0], + [3.0, 4.0, 5.0], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_take_2d_dim1() { + // Test take on 2D tensor along dimension 1 + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &device); + let indices = TestTensorInt::<1>::from_data([2, 0, 1], &device); + + let output = tensor.take::<1, 2>(1, indices); + let expected = TensorData::from([[2.0, 0.0, 1.0], [5.0, 3.0, 4.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn take_and_select_should_be_equivalent() { + // Verify that take and select produce identical results + let device = Default::default(); + let tensor = TestTensor::<2>::from_data( + [ + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + ], + &device, + ); + let indices = TestTensorInt::<1>::from_data([2, 0, 1, 1], &device); + + let result_take = tensor.clone().take::<1, 2>(0, indices.clone()); + let result_select = tensor.select(0, indices); + + let take_data = result_take.into_data(); + let select_data = result_select.into_data(); + + take_data.assert_eq(&select_data, false); +} + +#[test] +fn should_take_with_2d_indices() { + // Test take with 2D indices - output will be 3D with shape [2, 2, 4] + let device = Default::default(); + let tensor = TestTensor::<2>::from_data( + [ + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + ], + &device, + ); + + // 2D indices to select along dimension 0 - shape [2, 2] + let indices = TestTensorInt::<2>::from_data([[0, 2], [1, 0]], &device); + let output = tensor.take::<2, 3>(0, indices); + + // Expected: shape [2, 2, 4] - indices shape replaces dim 0 + let expected = TensorData::from([ + [[1.0, 2.0, 3.0, 4.0], [9.0, 10.0, 11.0, 12.0]], + [[5.0, 6.0, 7.0, 8.0], [1.0, 2.0, 3.0, 4.0]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_take_with_2d_indices_dim1() { + // Test take with 2D indices along dimension 1 - output will be 3D with shape [2, 2, 2] + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], &device); + + // 2D indices to select along dimension 1 - shape [2, 2] + let indices = TestTensorInt::<2>::from_data([[0, 3], [2, 1]], &device); + let output = tensor.take::<2, 3>(1, indices); + + // Expected: shape [2, 2, 2] - indices shape replaces dim 1 + let expected = TensorData::from([[[1.0, 4.0], [3.0, 2.0]], [[5.0, 8.0], [7.0, 6.0]]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_take_3d_tensor() { + // Test take with 3D tensor - output will be 4D with shape [2, 2, 2, 2] + let device = Default::default(); + let tensor = TestTensor::<3>::from_data( + [ + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], + [[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]], + ], + &device, + ); + + // 2D indices to select along dimension 1 - shape [2, 2] + let indices = TestTensorInt::<2>::from_data([[0, 2], [1, 0]], &device); + let output = tensor.take::<2, 4>(1, indices); + + // Expected: shape [2, 2, 2, 2] - indices shape replaces dim 1 + let expected = TensorData::from([ + [[[1.0, 2.0], [5.0, 6.0]], [[3.0, 4.0], [1.0, 2.0]]], + [[[7.0, 8.0], [11.0, 12.0]], [[9.0, 10.0], [7.0, 8.0]]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_take_with_3d_indices() { + // Test take with 3D indices - output will be 4D + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + + // 3D indices to select along dimension 1 - shape [2, 2, 2] + let indices = TestTensorInt::<3>::from_data([[[0, 2], [1, 0]], [[2, 1], [0, 2]]], &device); + let output = tensor.take::<3, 4>(1, indices); + + // Expected: shape [2, 2, 2, 2] - indices shape replaces dim 1 + let expected = TensorData::from([ + [[[1.0, 3.0], [2.0, 1.0]], [[3.0, 2.0], [1.0, 3.0]]], + [[[4.0, 6.0], [5.0, 4.0]], [[6.0, 5.0], [4.0, 6.0]]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +#[should_panic] +fn should_panic_take_invalid_dimension() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], &device); + let indices = TestTensorInt::<1>::from_data([1, 0], &device); + + // This should panic because dimension 10 is out of bounds + tensor.take::<1, 2>(10, indices); +} + +#[test] +fn should_take_with_single_index() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + let indices = TestTensorInt::<1>::from_data([1], &device); + + let output = tensor.take::<1, 2>(0, indices); + let expected = TensorData::from([[4.0, 5.0, 6.0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_take_with_negative_dim_2d() { + // Test using negative dimension indexing on 2D tensor + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + let indices = TestTensorInt::<1>::from_data([2, 0, 1], &device); + + // Using -1 should refer to the last dimension (dim 1) + let output_neg = tensor.clone().take::<1, 2>(-1, indices.clone()); + let output_pos = tensor.take::<1, 2>(1, indices); + + // Both should produce the same result + let neg_data = output_neg.into_data(); + let pos_data = output_pos.into_data(); + neg_data.assert_eq(&pos_data, false); +} + +#[test] +#[should_panic] +fn should_panic_take_negative_dim_out_of_bounds() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + let indices = TestTensorInt::<1>::from_data([0, 1], &device); + + // This should panic because -3 is out of bounds for a 2D tensor + tensor.take::<1, 2>(-3, indices); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/topk.rs b/crates/burn-backend-tests/tests/tensor/float/ops/topk.rs new file mode 100644 index 0000000..3672586 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/topk.rs @@ -0,0 +1,21 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn test_topk_with_indices_3d() { + let tensor = + TestTensor::<3>::from([[[1., 4., 7.], [2., 5., 6.]], [[3., 0., 9.], [8., 2., 7.]]]); + + let (values, indices) = tensor.topk_with_indices(2, /*dim*/ 2); + + let values_expected = TensorData::from([[[7., 4.], [6., 5.]], [[9., 3.], [8., 7.]]]); + + values + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::default()); + + let indices_expected = TensorData::from([[[2, 1], [2, 1]], [[2, 0], [0, 2]]]); + + indices.into_data().assert_eq(&indices_expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/transaction.rs b/crates/burn-backend-tests/tests/tensor/float/ops/transaction.rs new file mode 100644 index 0000000..ac9dbbf --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/transaction.rs @@ -0,0 +1,31 @@ +use super::*; +use burn_tensor::Transaction; + +// https://github.com/tracel-ai/burn/issues/4021 +#[test] +fn should_support_transaction() { + let rows = 261120; + let cols = 408; + + let device = Default::default(); + + let j = TestTensor::<2>::zeros([rows, cols], &device); + let jt = j.clone().transpose(); + + let g = jt.matmul(j); + + let g = g.transpose(); + let expected = g.to_data(); + + assert_eq!(g.shape().dims(), [cols, cols]); + + // Fails + let [data] = Transaction::default() + .register(g) + .execute() + .try_into() + .unwrap(); + + // check byte equality + assert_eq!(data, expected); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/transpose.rs b/crates/burn-backend-tests/tests/tensor/float/ops/transpose.rs new file mode 100644 index 0000000..0146e8b --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/transpose.rs @@ -0,0 +1,116 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_transpose_ops() { + let tensor = TestTensor::<3>::from_data( + [ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], + ], + &Default::default(), + ); + + // Check the .t() alias. + let output = tensor.t(); + + let expected = TensorData::from([ + [[0.0, 3.0], [1.0, 4.0], [2.0, 5.0]], + [[6.0, 9.0], [7.0, 10.0], [8.0, 11.0]], + ]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_transpose_maybe_fused_with_one() { + let tensor = TestTensor::<3>::from_data( + [ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], + ], + &Default::default(), + ); + let ones = TestTensor::<3>::ones([1, 1, 1], &Default::default()); + + let output = tensor.transpose(); + let expected = TensorData::from([ + [[0.0, 3.0], [1.0, 4.0], [2.0, 5.0]], + [[6.0, 9.0], [7.0, 10.0], [8.0, 11.0]], + ]); + let expected_ones = TensorData::from([[[1.0]]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + ones.into_data() + .assert_approx_eq::(&expected_ones, Tolerance::default()); +} + +#[test] +fn should_support_swap_dims_no_op() { + let tensor = TestTensor::<3>::from_data( + [ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], + ], + &Default::default(), + ); + + let output = tensor.swap_dims(0, 0); + let expected = TensorData::from([ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], + ]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_swap_dims() { + let tensor = TestTensor::<3>::from_data( + [ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], + ], + &Default::default(), + ); + + let output = tensor.swap_dims(0, 2); + let expected = TensorData::from([ + [[0.0, 6.0], [3.0, 9.0]], + [[1.0, 7.0], [4.0, 10.0]], + [[2.0, 8.0], [5.0, 11.0]], + ]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_swap_dims_neg_index() { + let tensor = TestTensor::<3>::from_data( + [ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], + ], + &Default::default(), + ); + + let output = tensor.swap_dims(-3, -1); + let expected = TensorData::from([ + [[0.0, 6.0], [3.0, 9.0]], + [[1.0, 7.0], [4.0, 10.0]], + [[2.0, 8.0], [5.0, 11.0]], + ]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/tri.rs b/crates/burn-backend-tests/tests/tensor/float/ops/tri.rs new file mode 100644 index 0000000..3325e10 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/tri.rs @@ -0,0 +1,21 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_triu() { + let tensor = TestTensor::<2>::from([[1., 1., 1.], [1., 1., 1.], [1., 1., 1.]]); + let output = tensor.triu(0); + let expected = TensorData::from([[1., 1., 1.], [0., 1., 1.], [0., 0., 1.]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_triu_positive_diagonal() { + let tensor = TestTensor::<2>::from([[1, 1, 1], [1, 1, 1], [1, 1, 1]]); + + let output = tensor.triu(1); + let expected = TensorData::from([[0, 1, 1], [0, 0, 1], [0, 0, 0]]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/trig.rs b/crates/burn-backend-tests/tests/tensor/float/ops/trig.rs new file mode 100644 index 0000000..61f065b --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/trig.rs @@ -0,0 +1,292 @@ +#![allow(clippy::approx_constant)] + +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; +use burn_tensor::s; +use core::f32::consts::{FRAC_PI_2, FRAC_PI_3, FRAC_PI_4, FRAC_PI_6, FRAC_PI_8, PI}; + +#[test] +fn should_support_cos_flipped_axis1() { + // [[0, pi], [pi/2, 3pi/2]] flipped on axis 1 -> [[pi, 0], [3pi/2, pi/2]] + let tensor = TestTensor::<2>::from([[0.0, PI], [FRAC_PI_2, 3.0 * FRAC_PI_2]]); + let flipped = tensor.flip([1]); + + let output = flipped.cos(); + let expected = TensorData::from([[-1.0, 1.0], [0.0, 0.0]]); + + output.into_data().assert_approx_eq::( + &expected, + Tolerance::default().set_half_precision_absolute(2e-3), + ); +} + +#[test] +fn should_support_sin_step_sliced() { + // Step-2 slice exercises stride=2 path in unary ops. + let tensor = TestTensor::<2>::from([[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]]); + let sliced = tensor.slice(s![.., 0..8;2]); + + let output = sliced.sin(); + let expected = TensorData::from([[0.0f32.sin(), 2.0f32.sin(), 4.0f32.sin(), 6.0f32.sin()]]); + + output.into_data().assert_approx_eq::( + &expected, + Tolerance::default().set_half_precision_absolute(2e-3), + ); +} + +#[test] +fn should_support_cos_step_sliced_3d() { + // 3D tensor with step-2 slice on last dim (the RF-DETR pattern). + let data: Vec = (0..12).map(|i| i as f32 * 0.5).collect(); + let tensor = TestTensor::<3>::from_data(TensorData::new(data, [1, 2, 6]), &Default::default()); + let sliced = tensor.slice(s![.., .., 0..6;2]); + + let output = sliced.cos(); + let expected = TensorData::from([[ + [0.0f32.cos(), 1.0f32.cos(), 2.0f32.cos()], + [3.0f32.cos(), 4.0f32.cos(), 5.0f32.cos()], + ]]); + + output.into_data().assert_approx_eq::( + &expected, + Tolerance::default().set_half_precision_absolute(2e-3), + ); +} + +#[test] +fn should_support_cos_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.cos(); + let expected = TensorData::from([[1.0, 0.54030, -0.41615], [-0.98999, -0.65364, 0.28366]]); + + // Metal has less precise trigonometric functions + let tolerance = Tolerance::default().set_half_precision_relative(1e-2); + + output + .into_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn should_support_cosh_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.cosh(); + let expected = TensorData::from([[1.0000, 1.5431, 3.7622], [10.0677, 27.3082, 74.2099]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_sin_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.sin(); + let expected = TensorData::from([[0.0, 0.841471, 0.909297], [0.141120, -0.756802, -0.958924]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_sinh_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.sinh(); + let expected = TensorData::from([[0.0000, 1.1752, 3.6269], [10.0179, 27.2899, 74.2032]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_tan_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.tan(); + let expected = TensorData::from([[0.0, 1.557408, -2.185040], [-0.142547, 1.157821, -3.380515]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_tanh_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.tanh(); + let expected = TensorData::from([[0.0, 0.761594, 0.964028], [0.995055, 0.999329, 0.999909]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_asin_ops() { + let data = TensorData::from([[0.0, 0.5, 0.707107], [-0.5, -0.707107, -1.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.asin(); + let expected = TensorData::from([[0.0, 0.523599, 0.785398], [-0.523599, -0.785398, -1.570796]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_acos_ops() { + let data = TensorData::from([[0.0, 0.5, 0.707107], [-0.5, -0.707107, -1.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.acos(); + let expected = TensorData::from([ + [1.570796, 1.047198, 0.785398], + [2.094395, 2.356194, 3.141593], + ]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_atan_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.atan(); + let expected = TensorData::from([[0.0, 0.785398, 1.107149], [1.249046, 1.325818, 1.373401]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_asinh_ops() { + let data = TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.asinh(); + let expected = TensorData::from([[0.0, 0.881374, 1.443635], [1.818446, 2.094713, 2.312438]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_acosh_ops() { + let data = TensorData::from([[1.0, 1.5, 2.0], [3.0, 4.0, 5.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.acosh(); + let expected = TensorData::from([[0.0, 0.962424, 1.316958], [1.762747, 2.063437, 2.292432]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_atanh_ops() { + let data = TensorData::from([[0.0, 0.5, 0.707107], [-0.5, -0.707107, -0.9]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.atanh(); + let expected = TensorData::from([[0.0, 0.549306, 0.881374], [-0.549306, -0.881374, -1.472219]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_atan2_ops() { + let y = TensorData::from([[0.0, 1.0, 1.0], [-1.0, -1.0, 0.0]]); + let x = TensorData::from([[1.0, 1.0, 0.0], [1.0, 0.0, -1.0]]); + + let y_tensor = TestTensor::<2>::from_data(y, &Default::default()); + let x_tensor = TestTensor::<2>::from_data(x, &Default::default()); + + let output = y_tensor.atan2(x_tensor); + let expected = TensorData::from([[0.0, 0.785398, 1.570796], [-0.785398, -1.570796, 3.141593]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_deg2rad_ops() { + let device = Default::default(); + let tensor = TestTensor::<1>::from_data( + [ + 0.0, 22.5, 30.0, 45.0, 60.0, 90.0, 135.0, 180.0, 270.0, 360.0, + ], + &device, + ); + + let output = tensor.deg2rad(); + let expected = TensorData::from([ + 0.0f32, + FRAC_PI_8, + FRAC_PI_6, + FRAC_PI_4, + FRAC_PI_3, + FRAC_PI_2, + 0.75 * PI, + PI, + 1.5 * PI, + 2.0 * PI, + ]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_support_rad2deg_ops() { + let device = Default::default(); + let tensor = TestTensor::<1>::from_data( + [ + 0.0, + FRAC_PI_8, + FRAC_PI_6, + FRAC_PI_4, + FRAC_PI_3, + FRAC_PI_2, + PI, + 1.5 * PI, + 2.0 * PI, + -FRAC_PI_3, + ], + &device, + ); + + let output = tensor.rad2deg(); + let expected = TensorData::from([ + 0.0f32, 22.5, 30.0, 45.0, 60.0, 90.0, 180.0, 270.0, 360.0, -60.0, + ]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/trunc.rs b/crates/burn-backend-tests/tests/tensor/float/ops/trunc.rs new file mode 100644 index 0000000..d29692c --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/trunc.rs @@ -0,0 +1,67 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{ElementConversion, TensorData}; + +#[test] +fn should_support_trunc_ops() { + let data = TensorData::from([[2.3, -1.7, 0.5], [-0.5, 3.9, -4.2]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.trunc(); + let expected = TensorData::from([[2.0, -1.0, 0.0], [0.0, 3.0, -4.0]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_truncate_positive_values_like_floor() { + let data = TensorData::from([1.7, 2.9, 3.1, 4.5]); + let tensor = TestTensor::<1>::from_data(data, &Default::default()); + + let output = tensor.trunc(); + let expected = TensorData::from([1.0, 2.0, 3.0, 4.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_truncate_negative_values_like_ceil() { + let data = TensorData::from([-1.7, -2.9, -3.1, -4.5]); + let tensor = TestTensor::<1>::from_data(data, &Default::default()); + + let output = tensor.trunc(); + let expected = TensorData::from([-1.0, -2.0, -3.0, -4.0]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn should_handle_special_cases() { + // Test special IEEE 754 cases + let data = TensorData::from([0.0, -0.0, f32::INFINITY, f32::NEG_INFINITY, f32::NAN]); + let tensor = TestTensor::<1>::from_data(data, &Default::default()); + + let output = tensor.trunc(); + let values = output.into_data().as_slice::().unwrap().to_vec(); + + // Check positive zero + assert_eq!(values[0], 0.0f32.elem::()); + assert!(values[0].is_sign_positive()); + + // Check negative zero is preserved + assert_eq!(values[1], 0.0f32.elem::()); + assert!(values[1].is_sign_negative()); + + // Check infinity is preserved + assert!(values[2].is_infinite() && values[2].is_sign_positive()); + assert!(values[3].is_infinite() && values[3].is_sign_negative()); + + // Check NaN is preserved + assert!(values[4].is_nan()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/ops/unfold.rs b/crates/burn-backend-tests/tests/tensor/float/ops/unfold.rs new file mode 100644 index 0000000..385cf0b --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/ops/unfold.rs @@ -0,0 +1,68 @@ +use super::*; +use burn_tensor::Distribution; +use burn_tensor::TensorData; +use burn_tensor::s; + +#[test] +fn test_unfold_float() { + let device = Default::default(); + + let input = TestTensor::<3>::random([2, 6, 6], Distribution::Default, &device); + + let dim = 1; + let size = 3; + let step = 2; + let actual: TestTensor<4> = input.clone().unfold(dim, size, step); + + let expected = TestTensor::<4>::empty([2, 2, 6, 3], &device) + .slice_assign( + s![.., 0, .., ..], + input + .clone() + .slice(s![.., 0..3, ..]) + .swap_dims(1, 2) + .unsqueeze_dim::<4>(1), + ) + .slice_assign( + s![.., 1, .., ..], + input + .clone() + .slice(s![.., 2..5, ..]) + .swap_dims(1, 2) + .unsqueeze_dim::<4>(1), + ); + + actual.to_data().assert_eq(&expected.to_data(), true); +} + +#[test] +fn test_unfold_reshape_overlap() { + let device = Default::default(); + let signal = TestTensor::<2>::from_data( + [[1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12.]], + &device, + ); + + let size = 4; + let step = 1; + + let frames = signal.unfold::<3, _>(1, size, step); + + let flat = frames.reshape([9, 4]); + + let data = flat.into_data(); + + let expected = TensorData::from([ + [1., 2., 3., 4.], + [2., 3., 4., 5.], + [3., 4., 5., 6.], + [4., 5., 6., 7.], + [5., 6., 7., 8.], + [6., 7., 8., 9.], + [7., 8., 9., 10.], + [8., 9., 10., 11.], + [9., 10., 11., 12.], + ]); + + data.assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/primitive.rs b/crates/burn-backend-tests/tests/tensor/float/primitive.rs new file mode 100644 index 0000000..fc64fb4 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/primitive.rs @@ -0,0 +1,27 @@ +use super::*; +use burn_tensor::{Element, Shape, TensorData}; + +#[test] +fn should_support_float_dtype() { + let tensor = TestTensor::<2>::from([[0.0, -1.0, 2.0], [3.0, 4.0, -5.0]])/*.into_primitive()*/; + + assert_eq!(tensor.shape(), Shape::new([2, 3])); + assert_eq!( + tensor.dtype(), + FloatElem::dtype() // default float elem type + ); +} + +#[test] +fn should_support_into_data_from_data() { + let device = Default::default(); + let data = + TestTensor::<2>::from_data([[0.0, -1.0, 2.0], [3.0, 4.0, -5.0]], &device).into_data(); + let tensor = TestTensor::<2>::from_data(data, &device).slice(0); + + // Regression test for `LazyDeviceController` from_data(tensor.into_data()) roundtrips + // These unnecessary round-trips should be avoided, but should not panic + tensor + .into_data() + .assert_eq(&TensorData::from([[0.0, -1.0, 2.0]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/calibration.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/calibration.rs new file mode 100644 index 0000000..5e4f801 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/calibration.rs @@ -0,0 +1,110 @@ +use super::*; +use burn_tensor::{ + TensorData, Tolerance, + quantization::{Calibration, QuantLevel, QuantValue, compute_range}, +}; + +// NOTE: The scheme variant fields are not important for calibration, only the "main" variant (e.g., per-tensor) +#[test] +fn min_max_calibration_range_per_tensor() { + let device = Default::default(); + let tensor = TestTensor::<1>::from_data([-1.8, -1.0, 0.0, 0.5], &device); + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S); + + let range = compute_range(&scheme, &tensor, &Calibration::MinMax); + + range + .min + .into_data() + .assert_eq(&TensorData::from([-1.8]), false); + range + .max + .into_data() + .assert_eq(&TensorData::from([0.5]), false); +} + +#[test] +fn min_max_calibration_range_per_block() { + let device = Default::default(); + let tensor = TestTensor::<2>::from_data( + [ + [-1.8, -1.0, 0.0, 0.5], + [1.8, 1.0, 0.0, -0.5], + [0.01, 0.02, 0.03, 0.04], + [-0.01, -0.02, -0.03, -0.04], + ], + &device, + ); + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::block([4])); + + let range = compute_range(&scheme, &tensor, &Calibration::MinMax); + + range + .min + .into_data() + .assert_eq(&TensorData::from([[-1.8], [-0.5], [0.01], [-0.04]]), false); + range + .max + .into_data() + .assert_eq(&TensorData::from([[0.5], [1.8], [0.04], [-0.01]]), false); +} + +// abs-mean calibration: gamma = mean(|W|), range = [-gamma, +gamma] +// weights: [-0.9, -0.3, 0.0, 0.6] => mean(|w|) = (0.9+0.3+0.0+0.6)/4 = 0.45 +#[test] +fn abs_mean_calibration_range_per_tensor() { + let device = Default::default(); + let tensor = TestTensor::<1>::from_data([-0.9_f32, -0.3, 0.0, 0.6], &device); + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q2S); + + let range = compute_range(&scheme, &tensor, &Calibration::AbsMean); + + range + .min + .into_data() + .assert_approx_eq::(&TensorData::from([-0.45_f32]), Tolerance::default()); + range + .max + .into_data() + .assert_approx_eq::(&TensorData::from([0.45_f32]), Tolerance::default()); +} + +// block abs-mean: 2 blocks of 4 weights each +// block 0: [-0.9, -0.3, 0.0, 0.6] gamma = 0.45 +// block 1: [0.1, 0.2, 0.3, 0.4] gamma = 0.25 +#[test] +fn abs_mean_calibration_range_per_block() { + let device = Default::default(); + let tensor = + TestTensor::<2>::from_data([[-0.9_f32, -0.3, 0.0, 0.6], [0.1, 0.2, 0.3, 0.4]], &device); + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q2S) + .with_level(QuantLevel::block([4])); + + let range = compute_range(&scheme, &tensor, &Calibration::AbsMean); + + range.min.into_data().assert_approx_eq::( + &TensorData::from([[-0.45_f32], [-0.25]]), + Tolerance::default(), + ); + range.max.into_data().assert_approx_eq::( + &TensorData::from([[0.45_f32], [0.25]]), + Tolerance::default(), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/data.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/data.rs new file mode 100644 index 0000000..71669cc --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/data.rs @@ -0,0 +1,48 @@ +use super::*; +use alloc::vec; +use burn_tensor::quantization::{QuantLevel, QuantValue}; +use burn_tensor::{Device, TensorData}; + +#[test] +fn should_support_per_tensor_symmetric_int8() { + let device = Device::default(); + let data = TensorData::quantized( + vec![-127i8, -71, 0, 35], + [4], + device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S), + &[0.014_173_228], + ); + let tensor = TestTensor::<1>::from_data(data.clone(), &device); + + let q_data = tensor.into_data(); + q_data.assert_eq(&data, true); + + let tensor = TestTensor::<1>::from_data(q_data.clone(), &device); + + tensor.into_data().assert_eq(&q_data, true); +} + +#[test] +fn should_support_per_block_symmetric_int8() { + let device = Device::default(); + let data = TensorData::quantized( + vec![ + -127i8, -71, 0, 35, -127i8, -71, 0, 35, -32, -63, -95, -127, -32, -63, -95, -127, + ], + [16], + device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::block([8])), + &[0.014_173_228, 0.000_314_96], + ); + let tensor = TestTensor::<1>::from_data(data.clone(), &device); + + tensor.into_data().assert_eq(&data, true); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/mod.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/mod.rs new file mode 100644 index 0000000..8e92afd --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/mod.rs @@ -0,0 +1,49 @@ +pub use super::*; // re-export test types + +mod calibration; +mod data; +mod ops; +mod scheme; + +/// Quantized tensor utilities +pub mod qtensor { + use super::TestTensor; + + use burn_tensor::quantization::QuantLevel; + use burn_tensor::{TensorData, quantization::QuantValue}; + + pub struct QTensor; + + impl QTensor { + /// Creates a quantized int8 tensor from the floating point data using the default quantization scheme + /// (i.e., per-tensor symmetric quantization). + pub fn int8>(floats: F) -> TestTensor { + Self::int8_symmetric(floats) + } + + /// Creates a quantized int8 tensor from the floating point data using blocks of size 16 + pub fn int8_block>(floats: F) -> TestTensor { + let device = Default::default(); + TestTensor::from_data(floats, &device).quantize_dynamic( + &device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::block([16])), + ) + } + + /// Creates a quantized int8 tensor from the floating point data using per-tensor symmetric quantization. + pub fn int8_symmetric>(floats: F) -> TestTensor { + let device = Default::default(); + TestTensor::from_data(floats, &device).quantize_dynamic( + &device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S), + ) + } + } +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/abs.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/abs.rs new file mode 100644 index 0000000..33e1dd9 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/abs.rs @@ -0,0 +1,19 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_abs_ops() { + let tensor = QTensor::<2>::int8([[0.0, -1.0, 2.0], [3.0, 4.0, -5.0]]); + + let output = tensor.abs(); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]), + Tolerance::absolute(1e-1), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/add.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/add.rs new file mode 100644 index 0000000..a7afe09 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/add.rs @@ -0,0 +1,106 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn test_add_d2() { + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_2 = QTensor::<2>::int8([[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]]); + + let output = tensor_1 + tensor_2; + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[6.0, 8.0, 10.0], [12.0, 14.0, 16.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn test_add_broadcast() { + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0, 2.0]]); + let tensor_2 = QTensor::<2>::int8([[3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]); + + let output = tensor_1 + tensor_2; + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[3.0, 5.0, 7.0], [6.0, 8.0, 10.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn test_add_different_strides_rhs() { + // We need to execute an operation after `from data` to trigger inplace in some backends. + // Which is the operation that might be problematic in this case. + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0], [2.0, 3.0]]) * 1; + let tensor_2 = QTensor::<2>::int8([[4.0, 5.0], [6.0, 7.0]]) * 1; + + let output = tensor_1 + tensor_2.transpose(); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[4.0, 7.0], [7.0, 10.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn test_add_different_strides_lhs() { + // We need to execute an operation after `from data` to trigger inplace in some backends. + // Which is the operation that might be problematic in this case. + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0], [2.0, 3.0]]) * 1; + let tensor_2 = QTensor::<2>::int8([[4.0, 5.0], [6.0, 7.0]]) * 1; + + let output = tensor_1.transpose() + tensor_2; + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[4.0, 7.0], [7.0, 10.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn test_add_different_strides_broadcast() { + // We need to execute an operation after `from data` to trigger inplace in some backends. + // Which is the operation that might be problematic in this case. + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0], [2.0, 3.0]]) * 1; + let tensor_2 = QTensor::<2>::int8([[4.0, 5.0]]) * 1; + + let output = tensor_1.transpose() + tensor_2; + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[4.0, 7.0], [5.0, 8.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn should_support_add_scalar_ops() { + let scalar = 2.0; + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor + scalar; + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[2.0, 3.0, 4.0], [5.0, 6.0, 7.0]]), + Tolerance::absolute(1e-1), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/aggregation.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/aggregation.rs new file mode 100644 index 0000000..a684c9b --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/aggregation.rs @@ -0,0 +1,166 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn test_should_mean() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.mean(); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&TensorData::from([15.0 / 6.0]), Tolerance::absolute(1e-1)); +} + +#[test] +fn test_should_sum() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.sum(); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&TensorData::from([15.0]), Tolerance::absolute(1e-1)); +} + +#[test] +fn test_should_mean_last_dim() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.mean_dim(1); + let expected = TensorData::from([[3.0 / 3.0], [12.0 / 3.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn test_should_sum_last_dim() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.sum_dim(1); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[3.0], [12.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn test_should_sum_first_dim() { + let tensor = QTensor::<2>::int8([[3.0, 1.0, 2.0], [4.0, 2.0, 3.0]]); + + let output = tensor.sum_dim(0); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[7.0, 3.0, 5.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn test_should_mean_first_dim() { + let tensor = QTensor::<2>::int8([[3.0, 1.0, 2.0], [4.0, 2.0, 3.0]]); + + let output = tensor.mean_dim(0); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[7.0 / 2.0, 3.0 / 2.0, 5.0 / 2.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn test_should_sum_mid_dim_3d_non_contiguous_1() { + let tensor = QTensor::<3>::int8([ + [[2.0, 4.0, 1.0], [7.0, -5.0, 3.0]], + [[3.0, 1.0, 2.0], [4.0, 2.0, 3.0]], + ]); + + let output = tensor.swap_dims(0, 2).sum_dim(1); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::new(vec![9.0, 7.0, -1.0, 3.0, 4.0, 5.0], [3, 1, 2]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn test_should_sum_mid_dim_3d_non_contiguous_2() { + let tensor = QTensor::<3>::int8([ + [[2.0, 4.0, 1.0], [7.0, -5.0, 3.0]], + [[3.0, 1.0, 2.0], [4.0, 2.0, 3.0]], + ]); + + let output = tensor.swap_dims(0, 1).sum_dim(1); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::new(vec![5.0, 5.0, 3.0, 11.0, -3.0, 6.0], [2, 1, 3]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn test_prod_float() { + let tensor = QTensor::<2>::int8([[2.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.prod(); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&TensorData::from([240.0]), Tolerance::rel_abs(1e-1, 1e-1)); + + let tensor_with_zero = QTensor::<2>::int8([[2.0, 0.0, 2.0], [3.0, 4.0, 5.0]]); + let output = tensor_with_zero.prod(); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&TensorData::from([0.0]), Tolerance::rel_abs(1e-1, 1e-1)); +} + +#[test] +fn test_prod_dim_float() { + let tensor = QTensor::<2>::int8([[2.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.prod_dim(1); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[4.0], [60.0]]), + Tolerance::absolute(1e-1), + ); + + let tensor_with_zero = QTensor::<2>::int8([[2.0, 0.0, 2.0], [3.0, 4.0, 5.0]]); + let output = tensor_with_zero.prod_dim(1); + let expected = TensorData::from([[0.0], [60.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/all.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/all.rs new file mode 100644 index 0000000..30fb967 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/all.rs @@ -0,0 +1,24 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_all() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 0.0], [1.0, -1.0, 1.0]]); + let data_actual = tensor.all().into_data(); + let data_expected = TensorData::from([false]); + assert_eq!(data_expected, data_actual); + + let tensor = QTensor::<2>::int8([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]); + let data_actual = tensor.all().into_data(); + let data_expected = TensorData::from([true]); + assert_eq!(data_expected, data_actual); +} + +#[test] +fn test_all_dim() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 0.0], [1.0, -1.0, 1.0]]); + let data_actual = tensor.all_dim(1).into_data(); + let data_expected = TensorData::from([[false], [true]]); + assert_eq!(data_expected, data_actual); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/any.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/any.rs new file mode 100644 index 0000000..6f0671a --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/any.rs @@ -0,0 +1,25 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_any() { + let tensor = QTensor::<2>::int8([[0.0, 0.0, 0.0], [1.0, -1.0, 0.0]]); + let data_actual = tensor.any().into_data(); + let data_expected = TensorData::from([true]); + assert_eq!(data_expected, data_actual); + + let tensor = QTensor::<2>::int8([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]); + let data_actual = tensor.any().into_data(); + let data_expected = TensorData::from([false]); + assert_eq!(data_expected, data_actual); +} + +#[test] +fn test_any_dim() { + let tensor = QTensor::<2>::int8([[0.0, 0.0, 0.0], [1.0, -1.0, 0.0]]); + + let data_actual = tensor.any_dim(1).into_data(); + let data_expected = TensorData::from([[false], [true]]); + assert_eq!(data_expected, data_actual); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/arg.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/arg.rs new file mode 100644 index 0000000..f4138ac --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/arg.rs @@ -0,0 +1,47 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_argmax_2d_dim0() { + let tensor = QTensor::<2>::int8([[10.0, 11.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.argmax(0); + + output + .into_data() + .assert_eq(&TensorData::from([[0, 0, 1]]), false); +} + +#[test] +fn test_argmin_2d_dim0() { + let tensor = QTensor::<2>::int8([[10.0, 11.0, 2.0], [30.0, 4.0, 5.0]]); + + let output = tensor.argmin(0); + + output + .into_data() + .assert_eq(&TensorData::from([[0, 1, 0]]), false); +} + +#[test] +fn test_argmax_2d_dim1() { + let tensor = QTensor::<2>::int8([[10.0, 11.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.argmax(1); + + output + .into_data() + .assert_eq(&TensorData::from([[1], [2]]), false); +} + +#[test] +fn test_argmin_2d_dim1() { + let tensor = QTensor::<2>::int8([[10.0, 11.0, 2.0], [30.0, 4.0, 5.0]]); + + let output = tensor.argmin(1); + + output + .into_data() + .assert_eq(&TensorData::from([[2], [1]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/cat.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/cat.rs new file mode 100644 index 0000000..e856f49 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/cat.rs @@ -0,0 +1,65 @@ +use super::qtensor::*; +use super::*; +use alloc::vec; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_cat_ops_2d_dim0() { + let tensor_1 = QTensor::<2>::int8([[1.0, 2.0, 3.0]]); + let tensor_2 = QTensor::<2>::int8([[4.0, 5.0, 6.0]]); + + let output = TestTensor::cat(vec![tensor_1, tensor_2], 0); + let expected = TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_support_cat_ops_2d_dim1() { + let tensor_1 = QTensor::<2>::int8([[1.0, 2.0, 3.0]]); + let tensor_2 = QTensor::<2>::int8([[4.0, 5.0, 6.0]]); + + let output = TestTensor::cat(vec![tensor_1, tensor_2], 1); + let expected = TensorData::from([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_support_cat_ops_3d() { + let tensor_1 = QTensor::<3>::int8([[[1.0, 2.0, 3.0]], [[1.1, 2.1, 3.1]]]); + let tensor_2 = QTensor::<3>::int8([[[4.0, 5.0, 6.0]]]); + + let output = TestTensor::cat(vec![tensor_1, tensor_2], 0); + let expected = TensorData::from([[[1.0, 2.0, 3.0]], [[1.1, 2.1, 3.1]], [[4.0, 5.0, 6.0]]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +#[should_panic] +fn should_panic_when_dimensions_are_not_the_same() { + let tensor_1 = QTensor::<2>::int8([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]]); + let tensor_2 = QTensor::<2>::int8([[4.0, 5.0]]); + + let _output = TestTensor::cat(vec![tensor_1, tensor_2], 0); +} + +#[test] +#[should_panic] +fn should_panic_when_cat_exceeds_dimension() { + let tensor_1 = QTensor::<2>::int8([[1.0, 2.0, 3.0]]); + let tensor_2 = QTensor::<2>::int8([[4.0, 5.0, 6.0]]); + + let _output = TestTensor::cat(vec![tensor_1, tensor_2], 3); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/ceil.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/ceil.rs new file mode 100644 index 0000000..168fdd4 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/ceil.rs @@ -0,0 +1,16 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_ceil_ops() { + let tensor = QTensor::<2>::int8([[24.0423, 87.9478, 76.1838], [59.6929, 43.8169, 94.8826]]); + + let output = tensor.ceil(); + let expected = TensorData::from([[25., 88., 77.], [60., 44., 96.]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(1e-1, 1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/chunk.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/chunk.rs new file mode 100644 index 0000000..790478a --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/chunk.rs @@ -0,0 +1,98 @@ +use super::qtensor::*; +use super::*; +use alloc::vec::Vec; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn test_chunk_evenly_divisible() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]); + + let tensors: Vec> = tensor.chunk(3, 0); + assert_eq!(tensors.len(), 3); + + let expected = [ + TensorData::from([0., 1.]), + TensorData::from([2., 3.]), + TensorData::from([4., 5.]), + ]; + + for (index, tensor) in tensors.into_iter().enumerate() { + tensor + .dequantize() + .to_data() + .assert_approx_eq::(&expected[index], Tolerance::absolute(1e-1)); + } +} + +#[test] +fn test_chunk_not_evenly_divisible() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + + let tensors: Vec> = tensor.chunk(4, 0); + assert_eq!(tensors.len(), 4); + + let expected = [ + TensorData::from([0., 1.]), + TensorData::from([2., 3.]), + TensorData::from([4., 5.]), + TensorData::from([6.]), + ]; + + for (index, tensor) in tensors.into_iter().enumerate() { + tensor + .dequantize() + .to_data() + .assert_approx_eq::(&expected[index], Tolerance::absolute(1e-1)); + } +} + +#[test] +fn test_chunk_not_divisible() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]); + + let tensors: Vec> = tensor.chunk(7, 0); + assert_eq!(tensors.len(), 6); + + let expected = [ + TensorData::from([0.]), + TensorData::from([1.]), + TensorData::from([2.]), + TensorData::from([3.]), + TensorData::from([4.]), + TensorData::from([5.]), + ]; + + for (index, tensor) in tensors.into_iter().enumerate() { + tensor + .dequantize() + .to_data() + .assert_approx_eq::(&expected[index], Tolerance::absolute(1e-1)); + } +} + +#[test] +fn test_chunk_multi_dimension() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0, 3.0, 4.0, 5.0]]); + + let tensors: Vec> = tensor.chunk(2, 1); + assert_eq!(tensors.len(), 2); + + let expected = [ + TensorData::from([[0., 1., 2.]]), + TensorData::from([[3., 4., 5.]]), + ]; + + for (index, tensor) in tensors.into_iter().enumerate() { + tensor + .dequantize() + .to_data() + .assert_approx_eq::(&expected[index], Tolerance::absolute(1e-1)); + } +} + +#[test] +#[should_panic] +fn test_invalid_dim() { + let _tensors = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]).chunk(6, 1); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/clamp.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/clamp.rs new file mode 100644 index 0000000..b571a83 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/clamp.rs @@ -0,0 +1,49 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn clamp_min() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.clamp_min(2.0); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[2.0, 2.0, 2.0], [3.0, 4.0, 5.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn clamp_max() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.clamp_max(2.0); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[0.0, 1.0, 2.0], [2.0, 2.0, 2.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn clamp_min_max() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.clamp(1.0, 4.0); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[1.0, 1.0, 2.0], [3.0, 4.0, 4.0]]), + Tolerance::absolute(1e-1), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/cos.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/cos.rs new file mode 100644 index 0000000..457d4e9 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/cos.rs @@ -0,0 +1,17 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_cos_ops() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.cos(); + let expected = TensorData::from([[1.0, 0.5403, -0.4161], [-0.9899, -0.6536, 0.2836]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/cosh.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/cosh.rs new file mode 100644 index 0000000..deca1ad --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/cosh.rs @@ -0,0 +1,17 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_cosh_ops() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.cosh(); + let expected = TensorData::from([[1.0000, 1.5431, 3.7622], [10.0677, 27.3082, 74.2100]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/div.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/div.rs new file mode 100644 index 0000000..47aaee3 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/div.rs @@ -0,0 +1,50 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_div_ops() { + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_2 = QTensor::<2>::int8([[1.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor_1 / tensor_2; + let expected = TensorData::from([[0.0, 1.0, 1.0], [1.0, 1.0, 1.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn test_div_broadcast() { + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0, 2.0]]); + let tensor_2 = QTensor::<2>::int8([[1.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor_1 / tensor_2; + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[0.0, 1.0, 1.0], [0.0, 0.25, 0.4]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn should_support_div_scalar_ops() { + let scalar = 2.0; + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor / scalar; + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[0.0, 0.5, 1.0], [1.5, 2.0, 2.5]]), + Tolerance::absolute(1e-1), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/erf.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/erf.rs new file mode 100644 index 0000000..96776e9 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/erf.rs @@ -0,0 +1,33 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_erf_ops() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.erf(); + let expected = TensorData::from([[0.0000, 0.8427, 0.9953], [1.0000, 1.0000, 1.0000]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_support_erf_ops_with_negative_number() { + let tensor = QTensor::<2>::int8([[-0.056, -0.043, -0.089], [3.0, 4.0, 5.0]]); + + let output = tensor.erf(); + let expected = TensorData::from([ + [-0.06312324, -0.048490416, -0.10016122], + [1.0000, 1.0000, 1.0000], + ]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/exp.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/exp.rs new file mode 100644 index 0000000..7800697 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/exp.rs @@ -0,0 +1,17 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_exp_ops() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.exp(); + let expected = TensorData::from([[1.0, 2.71830, 7.3891], [20.0855, 54.5981, 148.4132]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/expand.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/expand.rs new file mode 100644 index 0000000..41f1fee --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/expand.rs @@ -0,0 +1,112 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn expand_2d() { + let tensor = QTensor::<1>::int8([1.0, 2.0, 3.0]); + let output = tensor.expand([3, 3]); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0], [1.0, 2.0, 3.0]]), + Tolerance::absolute(1e-1), + ); + + // Quantized [4.0, 7.0, 2.0, 3.0] + let tensor = QTensor::<1>::int8([4.0, 7.0, 2.0, 3.0]); + let output = tensor.expand([2, 4]); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[4.0, 7.0, 2.0, 3.0], [4.0, 7.0, 2.0, 3.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn expand_3d() { + let tensor = QTensor::<2>::int8([[1.0, 2.0], [3.0, 4.0]]); + + let output = tensor.expand([3, 2, 2]); + let expected = TensorData::from([ + [[1.0, 2.0], [3.0, 4.0]], + [[1.0, 2.0], [3.0, 4.0]], + [[1.0, 2.0], [3.0, 4.0]], + ]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn expand_higher_dimensions() { + let tensor = QTensor::<2>::int8([[1.0, 2.0, 3.0, 4.0]]); + + let output = tensor.expand([2, 3, 4]); + let expected = TensorData::from([ + [ + [1.0, 2.0, 3.0, 4.0], + [1.0, 2.0, 3.0, 4.0], + [1.0, 2.0, 3.0, 4.0], + ], + [ + [1.0, 2.0, 3.0, 4.0], + [1.0, 2.0, 3.0, 4.0], + [1.0, 2.0, 3.0, 4.0], + ], + ]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn broadcast_single() { + let tensor = QTensor::<1>::int8([1.0]); + + let output = tensor.expand([2, 3]); + + output + .dequantize() + .into_data() + .assert_eq(&TensorData::from([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]), false); +} + +#[test] +#[should_panic] +fn should_fail_expand_incompatible_shapes() { + let tensor = QTensor::<1>::int8([1.0, 2.0, 3.0]); + let _expanded_tensor = tensor.expand([2, 2]); +} + +#[test] +fn should_all_negative_one() { + let tensor = QTensor::<1>::int8([1.0, 2.0, 3.0]); + + let output = tensor.expand([2, -1]); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[1., 2., 3.], [1., 2., 3.]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +#[should_panic] +fn should_panic_negative_one_on_non_existing_dim() { + let tensor = QTensor::<1>::int8([1.0, 2.0, 3.0]); + let _expanded_tensor = tensor.expand([-1, 3]); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/flip.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/flip.rs new file mode 100644 index 0000000..a1930aa --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/flip.rs @@ -0,0 +1,39 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn flip_float() { + let tensor = QTensor::<3>::int8([[[0.0, 1.0, 2.0]], [[3.0, 4.0, 5.0]]]); + + let flipped = tensor.clone().flip([0, 2]); + let expected = TensorData::from([[[5., 4., 3.]], [[2., 1., 0.]]]); + + flipped + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); + + // Test with no flip + let flipped = tensor.clone().flip([]); + tensor.into_data().assert_eq(&flipped.into_data(), true); +} + +#[test] +#[should_panic] +fn flip_duplicated_axes() { + let tensor = QTensor::<3>::int8([[[0.0, 1.0, 2.0]], [[3.0, 4.0, 5.0]]]); + + // Test with a duplicated axis + let _ = tensor.flip([0, 0, 1]); +} + +#[test] +#[should_panic] +fn flip_out_of_bound_axis() { + let tensor = QTensor::<3>::int8([[[0.0, 1.0, 2.0]], [[3.0, 4.0, 5.0]]]); + + // Test with an out of bound axis + let _ = tensor.clone().flip([3, 0, 1]); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/floor.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/floor.rs new file mode 100644 index 0000000..438c6f0 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/floor.rs @@ -0,0 +1,16 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_floor_ops() { + let tensor = QTensor::<2>::int8([[24.0423, 87.9478, 76.1838], [59.6929, 43.8169, 94.8826]]); + + let output = tensor.floor(); + let expected = TensorData::from([[24., 87., 76.], [59., 43., 95.]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(1e-1, 1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/gather_scatter.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/gather_scatter.rs new file mode 100644 index 0000000..3ca2395 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/gather_scatter.rs @@ -0,0 +1,198 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::IndexingUpdateOp; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_gather_1d_dim0() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0]); + let indices = TestTensorInt::from_ints([1, 1, 0, 1, 2], &Default::default()); + + let output = tensor.gather(0, indices); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([1.0, 1.0, 0.0, 1.0, 2.0]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn should_gather_2d_dim0() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let indices = TestTensorInt::from_ints([[0, 1, 0], [1, 0, 1]], &Default::default()); + + let output = tensor.gather(0, indices); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[0.0, 4.0, 2.0], [3.0, 1.0, 5.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn should_gather_2d_dim1() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let indices = TestTensorInt::from_ints([[2, 1, 0, 0], [2, 0, 1, 2]], &Default::default()); + + let output = tensor.gather(1, indices); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[2.0, 1.0, 0.0, 0.0], [5.0, 3.0, 4.0, 5.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn should_gather_3d_dim1() { + let tensor = QTensor::<3>::int8([ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], + ]); + let indices = TestTensorInt::from_ints( + [[[1, 0, 0], [0, 1, 0]], [[0, 0, 1], [0, 1, 1]]], + &Default::default(), + ); + + let output = tensor.gather(1, indices); + let expected = TensorData::from([ + [[3.0, 1.0, 2.0], [0.0, 4.0, 2.0]], + [[6.0, 7.0, 11.0], [6.0, 10.0, 11.0]], + ]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_gather_2d_only_1dim() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let indices = TestTensorInt::<2>::from_ints([[1, 2]], &Default::default()).reshape([2, 1]); + + let output = tensor.gather(1, indices); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[1.0], [5.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn should_scatter_1d() { + let tensor = QTensor::<1>::int8([0.0, 0.0, 0.0]); + let values = QTensor::<1>::int8([5.0, 4.0, 3.0]); + let indices = TestTensorInt::from_ints([1, 0, 2], &Default::default()); + + let output = tensor.scatter(0, indices, values, IndexingUpdateOp::Add); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([4.0, 5.0, 3.0]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn should_scatter_2d_dim0() { + let tensor = QTensor::<2>::int8([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]); + let values = QTensor::<2>::int8([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let indices = TestTensorInt::from_ints([[1, 0, 1], [1, 1, 0]], &Default::default()); + + let output = tensor.scatter(0, indices, values, IndexingUpdateOp::Add); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[0.0, 2.0, 6.0], [5.0, 5.0, 3.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn should_scatter_2d_dim1() { + let tensor = QTensor::<2>::int8([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]); + let values = QTensor::<2>::int8([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let indices = TestTensorInt::from_ints([[1, 0, 2], [1, 2, 0]], &Default::default()); + + let output = tensor.scatter(1, indices, values, IndexingUpdateOp::Add); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[2.0, 1.0, 3.0], [6.0, 4.0, 5.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +fn should_scatter_3d_dim1() { + let tensor = QTensor::<3>::int8([ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], + ]); + let values = QTensor::<3>::int8([ + [[12.0, 13.0, 14.0], [15.0, 16.0, 17.0]], + [[18.0, 19.0, 20.0], [21.0, 22.0, 23.0]], + ]); + let indices = TestTensorInt::from_ints( + [[[1, 0, 0], [0, 1, 0]], [[0, 0, 1], [0, 1, 1]]], + &Default::default(), + ); + + let output = tensor.scatter(1, indices, values, IndexingUpdateOp::Add); + let expected = TensorData::from([ + [[15.0, 14.0, 33.0], [15.0, 20.0, 5.0]], + [[45.0, 26.0, 8.0], [9.0, 32.0, 54.0]], + ]); + + // Set higher tolerance (0.2) due to larger de/quantization errors + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(2e-1)); +} + +#[test] +fn should_scatter_2d_dim1_diff_shape() { + let tensor = QTensor::<2>::int8([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]); + let values = QTensor::<2>::int8([[1.0], [4.0]]); + let indices = TestTensorInt::from_ints([[1], [2]], &Default::default()); + + let output = tensor.scatter(1, indices, values, IndexingUpdateOp::Add); + + output + .dequantize() + .into_data() + .assert_approx_eq::( + &TensorData::from([[0.0, 1.0, 0.0], [0.0, 0.0, 4.0]]), + Tolerance::absolute(1e-1), + ); +} + +#[test] +#[should_panic] +fn scatter_should_panic_on_mismatch_of_shapes() { + let tensor = QTensor::<1>::int8([0.0, 0.0, 0.0]); + let values = QTensor::<1>::int8([1.0, 4.0]); + let indices = TestTensorInt::from_ints([1, 0, 2], &Default::default()); + + tensor.scatter(0, indices, values, IndexingUpdateOp::Add); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/log.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/log.rs new file mode 100644 index 0000000..ed31e2d --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/log.rs @@ -0,0 +1,20 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_log_ops() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.log(); + let expected = TensorData::from([ + [-f32::INFINITY, 0.0, core::f32::consts::LN_2], + [1.0986, 1.3862, 1.6094], + ]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/log1p.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/log1p.rs new file mode 100644 index 0000000..1861de7 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/log1p.rs @@ -0,0 +1,20 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_exp_log1p() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.log1p(); + let expected = TensorData::from([ + [0.0, core::f32::consts::LN_2, 1.0986], + [1.3862, 1.6094, 1.7917], + ]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/map_comparison.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/map_comparison.rs new file mode 100644 index 0000000..cbf2f2c --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/map_comparison.rs @@ -0,0 +1,157 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_equal() { + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_2 = QTensor::<2>::int8([[0.0, 1.0, 1.0], [3.0, 5.0, 4.0]]); + + let data_actual_cloned = tensor_1.clone().equal(tensor_2.clone()); + let data_actual_inplace = tensor_1.equal(tensor_2); + + let data_expected = TensorData::from([[true, true, false], [true, false, false]]); + assert_eq!(data_expected, data_actual_cloned.into_data()); + assert_eq!(data_expected, data_actual_inplace.into_data()); +} + +#[test] +fn test_not_equal() { + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_2 = QTensor::<2>::int8([[0.0, 1.0, 1.0], [3.0, 5.0, 4.0]]); + + let data_actual_cloned = tensor_1.clone().not_equal(tensor_2.clone()); + let data_actual_inplace = tensor_1.not_equal(tensor_2); + + let data_expected = TensorData::from([[false, false, true], [false, true, true]]); + assert_eq!(data_expected, data_actual_cloned.into_data()); + assert_eq!(data_expected, data_actual_inplace.into_data()); +} + +#[test] +#[ignore = "quantization equality with float element is undefined"] +fn test_equal_elem() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 2.0, 5.0]]); + + let data_actual_cloned = tensor.clone().equal_elem(2); + let data_actual_inplace = tensor.equal_elem(2); + + let data_expected = TensorData::from([[false, false, true], [false, true, false]]); + assert_eq!(data_expected, data_actual_cloned.into_data()); + assert_eq!(data_expected, data_actual_inplace.into_data()); +} + +#[test] +#[ignore = "quantization equality with float element is undefined"] +fn test_not_equal_elem() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 2.0, 5.0]]); + + let data_actual_cloned = tensor.clone().not_equal_elem(2); + let data_actual_inplace = tensor.not_equal_elem(2); + + let data_expected = TensorData::from([[true, true, false], [true, false, true]]); + assert_eq!(data_expected, data_actual_cloned.into_data()); + assert_eq!(data_expected, data_actual_inplace.into_data()); +} + +#[test] +#[ignore = "quantization equality with float element is undefined"] +fn test_greater_elem() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let data_actual_cloned = tensor.clone().greater_elem(4); + let data_actual_inplace = tensor.greater_elem(4); + + let data_expected = TensorData::from([[false, false, false], [false, false, true]]); + assert_eq!(data_expected, data_actual_cloned.into_data()); + assert_eq!(data_expected, data_actual_inplace.into_data()); +} + +#[test] +fn test_greater_equal_elem() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let data_actual_cloned = tensor.clone().greater_equal_elem(4.0); + let data_actual_inplace = tensor.greater_equal_elem(4.0); + + let data_expected = TensorData::from([[false, false, false], [false, true, true]]); + assert_eq!(data_expected, data_actual_cloned.into_data()); + assert_eq!(data_expected, data_actual_inplace.into_data()); +} + +#[test] +fn test_greater() { + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_2 = QTensor::<2>::int8([[0.0, 1.0, 1.0], [3.0, 5.0, 4.0]]); + + let data_actual_cloned = tensor_1.clone().greater(tensor_2.clone()); + let data_actual_inplace = tensor_1.greater(tensor_2); + + let data_expected = TensorData::from([[false, false, true], [false, false, true]]); + assert_eq!(data_expected, data_actual_cloned.into_data()); + assert_eq!(data_expected, data_actual_inplace.into_data()); +} + +#[test] +fn test_greater_equal() { + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0, 1.0], [3.0, 4.0, 5.0]]); + let tensor_2 = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 5.0, 4.0]]); + + let data_actual_cloned = tensor_1.clone().greater_equal(tensor_2.clone()); + let data_actual_inplace = tensor_1.greater_equal(tensor_2); + + let data_expected = TensorData::from([[true, true, false], [true, false, true]]); + assert_eq!(data_expected, data_actual_cloned.into_data()); + assert_eq!(data_expected, data_actual_inplace.into_data()); +} + +#[test] +fn test_lower_elem() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let data_actual_cloned = tensor.clone().lower_elem(4.0); + let data_actual_inplace = tensor.lower_elem(4.0); + + let data_expected = TensorData::from([[true, true, true], [true, false, false]]); + assert_eq!(data_expected, data_actual_cloned.into_data()); + assert_eq!(data_expected, data_actual_inplace.into_data()); +} + +#[test] +#[ignore = "quantization equality with float element is undefined"] +fn test_lower_equal_elem() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let data_actual_cloned = tensor.clone().lower_equal_elem(4.0); + let data_actual_inplace = tensor.lower_equal_elem(4.0); + + let data_expected = TensorData::from([[true, true, true], [true, true, false]]); + assert_eq!(data_expected, data_actual_cloned.into_data()); + assert_eq!(data_expected, data_actual_inplace.into_data()); +} + +#[test] +fn test_lower() { + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0, 1.0], [3.0, 4.0, 5.0]]); + let tensor_2 = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 5.0, 4.0]]); + + let data_actual_cloned = tensor_1.clone().lower(tensor_2.clone()); + let data_actual_inplace = tensor_1.lower(tensor_2); + + let data_expected = TensorData::from([[false, false, true], [false, true, false]]); + assert_eq!(data_expected, data_actual_cloned.into_data()); + assert_eq!(data_expected, data_actual_inplace.into_data()); +} + +#[test] +fn test_lower_equal() { + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_2 = QTensor::<2>::int8([[0.0, 1.0, 1.0], [3.0, 5.0, 4.0]]); + + let data_actual_cloned = tensor_1.clone().lower_equal(tensor_2.clone()); + let data_actual_inplace = tensor_1.lower_equal(tensor_2); + + let data_expected = TensorData::from([[true, true, false], [true, true, false]]); + assert_eq!(data_expected, data_actual_cloned.into_data()); + assert_eq!(data_expected, data_actual_inplace.into_data()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/mask.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/mask.rs new file mode 100644 index 0000000..25871d8 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/mask.rs @@ -0,0 +1,39 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_mask_where_ops() { + let tensor = QTensor::<2>::int8([[1.0, 7.0], [2.0, 3.0]]); + let mask = TestTensorBool::<2>::from_bool( + TensorData::from([[true, false], [false, true]]), + &Default::default(), + ); + let value = QTensor::<2>::int8([[1.8, 2.8], [3.8, 4.8]]); + + let output = tensor.mask_where(mask, value); + let expected = TensorData::from([[1.8, 7.0], [2.0, 4.8]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_support_mask_fill_ops() { + let tensor = QTensor::<2>::int8([[1.0, 7.0], [2.0, 3.0]]); + let mask = TestTensorBool::<2>::from_bool( + TensorData::from([[true, false], [false, true]]), + &Default::default(), + ); + + let output = tensor.mask_fill(mask, 2.0); + let expected = TensorData::from([[2.0, 7.0], [2.0, 2.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/maxmin.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/maxmin.rs new file mode 100644 index 0000000..d77920d --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/maxmin.rs @@ -0,0 +1,149 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn test_max_dim_2d() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.max_dim(1); + let expected = TensorData::from([[2.], [5.]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn test_max_dim_with_indices_2d_with_dim_0th() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let (output, index) = tensor.max_dim_with_indices(0); + + let output_expected = TensorData::from([[3., 4., 5.]]); + let index_expected = TensorData::from([[1, 1, 1]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&output_expected, Tolerance::rel_abs(2e-2, 1e-2)); + index.into_data().assert_eq(&index_expected, false); +} + +#[test] +fn test_max_dim_with_indices_2d() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let (output, index) = tensor.max_dim_with_indices(1); + + let output_expected = TensorData::from([[2.], [5.]]); + let index_expected = TensorData::from([[2], [2]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&output_expected, Tolerance::rel_abs(2e-2, 1e-2)); + index.into_data().assert_eq(&index_expected, false); +} + +#[test] +fn test_min_dim_2d() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.min_dim(1); + + let expected = TensorData::from([[0.], [3.]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn test_min_dim_with_indices_2d() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let (output, index) = tensor.min_dim_with_indices(1); + + let output_expected = TensorData::from([[0.], [3.]]); + let index_expected = TensorData::from([[0], [0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&output_expected, Tolerance::rel_abs(2e-2, 1e-2)); + index.into_data().assert_eq(&index_expected, false); +} + +#[test] +fn test_min_dim_2d_with_0th_dim() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.min_dim(0); + let expected = TensorData::from([[0., 1., 2.]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 3e-2)); +} + +#[test] +fn test_max_dim_2d_with_0th_dim() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.max_dim(0); + let expected = TensorData::from([[3., 4., 5.]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn test_min_dim_with_indices_2d_with_0th_dim() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let (output, index) = tensor.min_dim_with_indices(0); + + let output_expected = TensorData::from([[0., 1., 2.]]); + let index_expected = TensorData::from([[0, 0, 0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&output_expected, Tolerance::rel_abs(2e-2, 3e-2)); + index.into_data().assert_eq(&index_expected, false); +} + +#[test] +fn test_maximum_pair() { + let a = QTensor::<1>::int8([1.0, 5.0, 3.0, 4.0]); + let b = QTensor::<1>::int8([2.0, 1.0, 4.0, 5.0]); + + let output = a.max_pair(b); + let expected = TensorData::from([2.0, 5.0, 4.0, 5.0]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn test_minimum_pair() { + let a = QTensor::<1>::int8([1.0, 5.0, 3.0, 4.0]); + let b = QTensor::<1>::int8([2.0, 1.0, 4.0, 5.0]); + + let output = a.min_pair(b); + let expected = TensorData::from([1.0, 1.0, 3.0, 4.0]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/mod.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/mod.rs new file mode 100644 index 0000000..a3cc89f --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/mod.rs @@ -0,0 +1,50 @@ +pub use super::*; + +mod abs; +mod add; +mod aggregation; +mod all; +mod any; +mod arg; +mod cat; +mod ceil; +mod chunk; +mod clamp; +mod cos; +mod cosh; +mod div; +mod erf; +mod exp; +mod expand; +mod flip; +mod floor; +mod gather_scatter; +mod log; +mod log1p; +mod map_comparison; +mod mask; +mod maxmin; +mod mul; +mod narrow; +mod neg; +mod permute; +mod powf; +mod powf_scalar; +mod recip; +mod remainder; +mod repeat_dim; +mod reshape; +mod round; +mod select; +mod sin; +mod sinh; +mod slice; +mod sort_argsort; +mod split; +mod sqrt; +mod stack; +mod sub; +mod tan; +mod tanh; +mod topk; +mod transpose; diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/mul.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/mul.rs new file mode 100644 index 0000000..f7706d1 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/mul.rs @@ -0,0 +1,60 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_mul_ops() { + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_2 = tensor_1.clone(); + + let output = tensor_1 * tensor_2; + let expected = TensorData::from([[0.0, 1.0, 4.0], [9.0, 16.0, 25.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(5e-2, 1e-2)); +} + +#[test] +fn test_mul_broadcast() { + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0, 2.0]]); + let tensor_2 = QTensor::<2>::int8([[3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]); + + let output = tensor_1 * tensor_2; + let expected = TensorData::from([[0.0, 4.0, 10.0], [0.0, 7.0, 16.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn test_mul_broadcast_2_dims() { + let tensor_1 = QTensor::<2>::int8([[0.0], [1.0], [2.0]]); + let tensor_2 = QTensor::<2>::int8([[3.0, 4.0, 5.0]]); + + let output = tensor_1 * tensor_2; + let expected = TensorData::from([[0.0, 0.0, 0.0], [3.0, 4.0, 5.0], [6.0, 8.0, 10.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn should_support_mul_scalar_ops() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let scalar = 2.0; + + let output = tensor * scalar; + let expected = TensorData::from([[0.0, 2.0, 4.0], [6.0, 8.0, 10.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/narrow.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/narrow.rs new file mode 100644 index 0000000..cf92235 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/narrow.rs @@ -0,0 +1,58 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{Shape, TensorData}; + +#[test] +fn test_narrow() { + let tensor = QTensor::<2>::int8([[1., 2., 3.], [7., 8., 9.], [13., 14., 15.]]); + + let output = tensor.clone().narrow(0, 0, 2); + let expected = TensorData::from([[1., 2., 3.], [7., 8., 9.]]); + + assert_eq!(output.shape(), Shape::from([2, 3])); + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); + + let output = tensor.narrow(1, 1, 2); + let expected = TensorData::from([[2., 3.], [8., 9.], [14., 15.]]); + assert_eq!(output.shape(), Shape::from([3, 2])); + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +#[should_panic] +fn test_narrow_invalid_dim() { + let tensor = QTensor::<2>::int8([[1., 2., 3.], [7., 8., 9.], [13., 14., 15.]]); + + let _output = tensor.narrow(2, 0, 2); +} + +#[test] +#[should_panic] +fn test_narrow_invalid_start() { + let tensor = QTensor::<2>::int8([[1., 2., 3.], [7., 8., 9.], [13., 14., 15.]]); + + let _output = tensor.narrow(0, 3, 2); +} + +#[test] +#[should_panic] +fn test_narrow_invalid_zero_length() { + let tensor = QTensor::<2>::int8([[1., 2., 3.], [7., 8., 9.], [13., 14., 15.]]); + + let _output = tensor.narrow(0, 1, 0); +} + +#[test] +#[should_panic] +fn test_narrow_invalid_length() { + let tensor = QTensor::<2>::int8([[1., 2., 3.], [7., 8., 9.], [13., 14., 15.]]); + + let _output = tensor.narrow(0, 0, 4); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/neg.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/neg.rs new file mode 100644 index 0000000..b061c38 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/neg.rs @@ -0,0 +1,19 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_neg_ops() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.neg(); + let expected = TensorData::from([[-0.0, -1.0, -2.0], [-3.0, -4.0, -5.0]]).convert::(); + + // -0.0 is represented differently than 0.0 so we make sure the values are the same in f32 + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/permute.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/permute.rs new file mode 100644 index 0000000..fb5018c --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/permute.rs @@ -0,0 +1,66 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn permute_float() { + let tensor = QTensor::<1>::int8([ + 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12., 13., 14., 15., + ]) + .reshape([2, 2, 4]); + + let permuted = tensor.clone().permute([2, 1, 0]); + + let expected = TensorData::from([ + [[0., 8.], [4., 12.]], + [[1., 9.], [5., 13.]], + [[2., 10.], [6., 14.]], + [[3., 11.], [7., 15.]], + ]); + + permuted + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(1e-1, 1e-1)); + + // Test with negative axis + let permuted = tensor.clone().permute([-1, 1, 0]); + permuted + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(1e-1, 1e-1)); + + // Test with the same axis + let permuted = tensor.clone().permute([0, 1, 2]); + permuted + .dequantize() + .into_data() + .assert_approx_eq::( + &tensor.dequantize().into_data(), + Tolerance::rel_abs(1e-4, 1e-4), // dequant error should be the same + ); +} + +#[test] +#[should_panic] +fn edge_repeated_axes() { + let tensor = QTensor::<1>::int8([ + 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12., 13., 14., 15., + ]) + .reshape([2, 2, 4]); + + // Test with a repeated axis + let _ = tensor.permute([0, 0, 1]); +} + +#[test] +#[should_panic] +fn edge_out_of_bound_axis() { + let tensor = QTensor::<1>::int8([ + 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12., 13., 14., 15., + ]) + .reshape([2, 2, 4]); + + // Test with an invalid axis + let _ = tensor.permute([3, 0, 1]); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/powf.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/powf.rs new file mode 100644 index 0000000..0030ce9 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/powf.rs @@ -0,0 +1,60 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_powf_ops() { + let tensor = QTensor::<2>::int8([[1.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_pow = QTensor::<2>::int8([[1.0, 1.0, 2.0], [3.0, 4.0, 2.0]]); + + let output = tensor.powf(tensor_pow); + let expected = TensorData::from([[1.0, 1.0, 4.0], [27.0, 256.0, 25.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(4e-2, 1e-2)); +} + +#[test] +fn should_support_neg_power() { + let tensor = QTensor::<2>::int8([[1.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_pow = QTensor::<2>::int8([[-0.95, -0.67, -0.45], [-0.24, -0.5, -0.6]]); + + let output = tensor.powf(tensor_pow); + let expected = TensorData::from([[1., 1., 0.73204285], [0.76822936, 0.5, 0.38073079]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(4e-2, 1e-2)); +} + +#[test] +fn should_support_neg_values_with_even_power() { + let tensor = QTensor::<2>::int8([[0.0, -1.0, -2.0], [-3.0, -4.0, -5.0]]); + let tensor_pow = QTensor::<2>::int8([[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]]); + + let output = tensor.powf(tensor_pow); + let expected = TensorData::from([[0.0, 1.0, 4.0], [9.0, 16.0, 25.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(4e-2, 1e-2)); +} + +#[test] +fn should_support_neg_values_with_odd_power() { + let tensor = QTensor::<2>::int8([[0.0, -1.0, -2.0], [-3.0, -4.0, -4.0]]); + let tensor_pow = QTensor::<2>::int8([[3.0, 3.0, 3.0], [3.0, 3.0, 3.0]]); + + let output = tensor.powf(tensor_pow); + let expected = TensorData::from([[0.0, -1.0, -8.0], [-27.0, -64.0, -64.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(4e-2, 1e-2)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/powf_scalar.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/powf_scalar.rs new file mode 100644 index 0000000..2a2b665 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/powf_scalar.rs @@ -0,0 +1,56 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_powf_ops() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.powf_scalar(0.71); + let expected = TensorData::from([[0.0, 1.0, 1.6358], [2.182, 2.6759, 3.1352]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(4e-2, 1e-2)); +} + +#[test] +fn should_support_neg_power() { + let tensor = QTensor::<2>::int8([[1.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.powf_scalar(-0.33); + let expected = TensorData::from([[1.0, 1.0, 0.79553646], [0.695905, 0.6328783, 0.58794934]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(4e-2, 1e-2)); +} + +#[test] +fn should_support_neg_values_with_even_power() { + let tensor = QTensor::<2>::int8([[0.0, -1.0, -2.0], [-3.0, -4.0, -5.0]]); + + let output = tensor.powf_scalar(2.0); + let expected = TensorData::from([[0., 1., 4.], [9., 16., 25.]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(4e-2, 1e-2)); +} + +#[test] +fn should_support_neg_values_with_odd_power() { + let tensor = QTensor::<2>::int8([[0.0, -1.0, -2.0], [-3.0, -4.0, -4.0]]); + + let output = tensor.powf_scalar(3.0); + let expected = TensorData::from([[0.0, -1.0, -8.0], [-27.0, -64.0, -64.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(4e-2, 1e-2)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/recip.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/recip.rs new file mode 100644 index 0000000..b0f6fd2 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/recip.rs @@ -0,0 +1,17 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_recip_ops() { + let tensor = QTensor::<2>::int8([[0.5, 1.0, 2.0], [3.0, -4.0, -5.0]]); + + let output = tensor.recip(); + let expected = TensorData::from([[2.0, 1.0, 0.5], [0.33333, -0.25, -0.2]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/remainder.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/remainder.rs new file mode 100644 index 0000000..741aa64 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/remainder.rs @@ -0,0 +1,208 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_remainder_basic() { + let lhs = QTensor::<1>::int8([-3.0, -2.0, -1.0, 1.0, 2.0, 2.0]); + let rhs = QTensor::<1>::int8([2.0, 3.0, 1.0, 2.0, 1.0, 2.0]); + + let output = lhs.remainder(rhs); + let expected = TensorData::from([1., 1., 0., 1., 0., 0.]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +#[ignore = "quantization remainder with float element is undefined"] +fn should_support_remainder_basic_scalar() { + let tensor = QTensor::<1>::int8([-3.0, -2.0, -1.0, 1.0, 2.0, 3.0]); + + let output = tensor.remainder_scalar(2.0); + let expected = TensorData::from([1.0, 0.0, 1.0, 1.0, 0.0, 1.0]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_support_remainder_float() { + let lhs = QTensor::<1>::int8([1.0, 2.0, 3.0, 4.0, 5.0]); + let rhs = QTensor::<1>::int8([1.4233, 2.7313, 0.2641, 1.9651, 0.5897]); + + let output = lhs.remainder(rhs); + let expected = TensorData::from([1., 2., 0.0949, 0.0698, 0.2824]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_support_remainder_float_scalar() { + let tensor = QTensor::<1>::int8([1.0, 2.0, 3.0, 4.0, 5.0]); + + let output = tensor.remainder_scalar(-1.5); + let expected = TensorData::from([-0.5, -1.0, 0.0, -0.5, -1.0]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_be_zero() { + let lhs = QTensor::<1>::int8([0.0, 0.0, 0.0]); + let rhs = QTensor::<1>::int8([3.5, -2.1, 1.5]); + + let output = lhs.remainder(rhs); + let expected = TensorData::from([0.0, 0.0, 0.0]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_be_zero_scalar() { + let tensor = QTensor::<1>::int8([0.0, 0.0, 0.0]); + + let output = tensor.remainder_scalar(3.5); + let expected = TensorData::from([0.0, 0.0, 0.0]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_have_no_remainder() { + let lhs = QTensor::<1>::int8([1.0, 2.0, 3.0, 4.0, 5.0]); + let rhs = QTensor::<1>::int8([1.0, 2.0, 3.0, 4.0, 5.0]); + + let output = lhs.remainder(rhs); + let expected = TensorData::from([0.0, 0.0, 0.0, 0.0, 0.0]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_have_no_remainder_scalar() { + let tensor = QTensor::<1>::int8([4.0, 4.0]); + + let output = tensor.remainder_scalar(4.0); + let expected = TensorData::from([0.0, 0.0]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_be_negative() { + let lhs = QTensor::<1>::int8([-7.0, -3.0, 2.0, 6.0]); + let rhs = QTensor::<1>::int8([-2.5, -2.1, -1.5, -3.25]); + + let output = lhs.remainder(rhs); + let expected = TensorData::from([-2., -0.9, -1., -0.5]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_be_negative_scalar() { + let tensor = QTensor::<1>::int8([-7.0, -3.0, 2.0, 6.0]); + + let output = tensor.remainder_scalar(-2.5); + let expected = TensorData::from([-2.0, -0.50, -0.50, -1.5]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_support_fp_dividends() { + let tensor = QTensor::<1>::int8([-7.5, -2.5, 2.5, 7.5]); + + let output = tensor.remainder_scalar(3.0); + let expected = TensorData::from([1.5, 0.5, 2.5, 1.5]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_support_large_divisor() { + let lhs = QTensor::<1>::int8([-1.0, 1.0, -1.5, 1.5, -1.0, 1.0, -1.5, 1.5]); + let rhs = QTensor::<1>::int8([10.0, 10.0, 10.0, 10.0, -10.0, -10.0, -10.0, -10.0]); + + let output = lhs.remainder(rhs); + let expected = TensorData::from([9., 1., 8.5, 1.5, -1., -9., -1.5, -8.5]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_support_large_divisor_scalar() { + let tensor = QTensor::<1>::int8([-1.0, 1.0]); + + let output = tensor.remainder_scalar(10.0); + let expected = TensorData::from([9.0, 1.0]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_support_remainder_op() { + let lhs = QTensor::<1>::int8([-3.0, -2.0, -1.0, 1.0, 2.0, 2.0]); + let rhs = QTensor::<1>::int8([2.0, 3.0, 1.0, 2.0, 1.0, 2.0]); + + let output = lhs % rhs; + let expected = TensorData::from([1., 1., 0., 1., 0., 0.]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +#[ignore = "quantization remainder with float element is undefined"] +fn should_support_remainder_scalar_op() { + let tensor = QTensor::<1>::int8([-3.0, -2.0, -1.0, 1.0, 2.0, 3.0]); + + let output = tensor % 2.0; + let expected = TensorData::from([1.0, 0.0, 1.0, 1.0, 0.0, 1.0]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/repeat_dim.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/repeat_dim.rs new file mode 100644 index 0000000..e3e3db5 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/repeat_dim.rs @@ -0,0 +1,42 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn should_support_repeat_ops() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0, 3.0]]); + + let output = tensor.repeat_dim(0, 4); + let expected = TensorData::from([ + [0.0, 1.0, 2.0, 3.0], + [0.0, 1.0, 2.0, 3.0], + [0.0, 1.0, 2.0, 3.0], + [0.0, 1.0, 2.0, 3.0], + ]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::permissive()); +} + +#[test] +fn should_support_repeat_on_dims_larger_than_1() { + let tensor = QTensor::<1>::int8([ + 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12., 13., 14., 15., + ]) + .reshape([4, 2, 2]); + + let output = tensor.repeat_dim(2, 2); + let expected = TensorData::from([ + [[0., 1., 0., 1.], [2., 3., 2., 3.]], + [[4., 5., 4., 5.], [6., 7., 6., 7.]], + [[8., 9., 8., 9.], [10., 11., 10., 11.]], + [[12., 13., 12., 13.], [14., 15., 14., 15.]], + ]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(1e-1, 1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/reshape.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/reshape.rs new file mode 100644 index 0000000..528a378 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/reshape.rs @@ -0,0 +1,77 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn should_support_reshape_1d() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0, 3.0]]); + + let output = tensor.clone().reshape([1, 4]); + let expected = TensorData::from([[0.0, 1.0, 2.0, 3.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn should_support_reshape_2d() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.clone().reshape([6]); + let expected = TensorData::from([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn should_support_dim_infererence() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0]) + .reshape([4, 3]); + + // Infer the dimension via -1 + let reshaped = tensor.clone().reshape([2, -1]); + assert_eq!(reshaped.shape(), [2, 6].into()); + + // Infer the dimension via 0 (keep from the source) and -1 (infer) + let reshaped = reshaped.reshape([0, 2, -1]); + assert_eq!(reshaped.shape(), [2, 2, 3].into()); + + // This is effectively as if we did a flatten + let reshaped = tensor.clone().reshape([-1]); + assert_eq!(reshaped.shape(), [12].into()); + + // Keeping the first dimension the same (using 0) + let reshaped = tensor.clone().reshape([0, 3]); + assert_eq!(reshaped.shape(), [4, 3].into()); +} + +#[test] +fn should_not_corrupt_after_slice() { + let zeros = QTensor::<1>::int8([0.0, 0.0]); + zeros.clone().slice([1..2]).reshape([1]).exp(); + + // May lead to zeroes being equal to [0.0, 1.0] + zeros.dequantize().into_data().assert_eq( + &TestTensor::<1>::zeros([2], &Default::default()).to_data(), + true, + ); +} + +#[test] +#[should_panic] +fn multiple_neg_ones() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0]); + let _ = tensor.reshape([-1, -1]); +} + +#[test] +#[should_panic] +fn neg_value() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0]); + let _ = tensor.reshape([-2, -1]); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/round.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/round.rs new file mode 100644 index 0000000..25f8539 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/round.rs @@ -0,0 +1,31 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_round_ops() { + let tensor = QTensor::<2>::int8([[24.0423, 87.9478, 76.1838], [59.6929, 43.8169, 94.8826]]); + + let output = tensor.round(); + let expected = TensorData::from([[24., 88., 76.], [60., 44., 95.]]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_round_ties_even() { + // NOTE: round ties to even only affects values that are exact halfway from ceil/floor, so quantization + // errors can impact this. This basically only guarantees the values for the max value in the range since + // it is always represented correctly. + let tensor = QTensor::<1>::int8([5.5]); + + let output = tensor.round(); + let expected = TensorData::from([6.]); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/select.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/select.rs new file mode 100644 index 0000000..92c87ea --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/select.rs @@ -0,0 +1,120 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::IndexingUpdateOp; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_select_1d() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0]); + let indices = TestTensorInt::from_data([1, 1, 0, 1, 2], &Default::default()); + + let output = tensor.select(0, indices); + let expected = TensorData::from([1.0, 1.0, 0.0, 1.0, 2.0]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn should_select_2d_dim0_same_num_dim() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let indices = TestTensorInt::from_data([1, 0], &Default::default()); + + let output = tensor.select(0, indices); + let expected = TensorData::from([[3.0, 4.0, 5.0], [0.0, 1.0, 2.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn should_select_2d_dim0_more_num_dim() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let indices = TestTensorInt::from_data([1, 0, 1, 1], &Default::default()); + + let output = tensor.select(0, indices); + let expected = TensorData::from([ + [3.0, 4.0, 5.0], + [0.0, 1.0, 2.0], + [3.0, 4.0, 5.0], + [3.0, 4.0, 5.0], + ]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn should_select_2d_dim1() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let indices = TestTensorInt::from_data([1, 1, 0, 1, 2], &Default::default()); + + let output = tensor.select(1, indices); + let expected = TensorData::from([[1.0, 1.0, 0.0, 1.0, 2.0], [4.0, 4.0, 3.0, 4.0, 5.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn should_select_assign_1d() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0]); + let values = QTensor::<1>::int8([5.0, 4.0, 3.0, 2.0, 1.0]); + let indices = TestTensorInt::from_data(TensorData::from([1, 1, 0, 1, 2]), &Default::default()); + + let output = tensor.select_assign(0, indices, values, IndexingUpdateOp::Add); + let expected = TensorData::from([3.0, 12.0, 3.0]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_select_assign_2d_dim0() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let values = tensor.clone(); + let indices = TestTensorInt::from_data(TensorData::from([1, 0]), &Default::default()); + + let output = tensor.select_assign(0, indices, values, IndexingUpdateOp::Add); + let expected = TensorData::from([[3.0, 5.0, 7.0], [3.0, 5.0, 7.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_select_assign_2d_dim1() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let values = tensor.clone(); + let indices = TestTensorInt::from_data(TensorData::from([1, 0, 2]), &Default::default()); + + let output = tensor.select_assign(1, indices, values, IndexingUpdateOp::Add); + let expected = TensorData::from([[1.0, 1.0, 4.0], [7.0, 7.0, 10.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +#[should_panic] +fn should_select_panic_invalid_dimension() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let indices = TestTensorInt::from_data([1, 1, 0, 1, 2], &Default::default()); + + tensor.select(10, indices); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/sin.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/sin.rs new file mode 100644 index 0000000..5bae654 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/sin.rs @@ -0,0 +1,17 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_sin_ops() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.sin(); + let expected = TensorData::from([[0.0, 0.8414, 0.9092], [0.1411, -0.7568, -0.9589]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/sinh.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/sinh.rs new file mode 100644 index 0000000..da7e6c1 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/sinh.rs @@ -0,0 +1,17 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_sinh_ops() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.sinh(); + let expected = TensorData::from([[0.0000, 1.1752, 3.6269], [10.0179, 27.2899, 74.2032]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(3e-2, 1e-2)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/slice.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/slice.rs new file mode 100644 index 0000000..ecbe14c --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/slice.rs @@ -0,0 +1,229 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::{Tolerance, s}; + +#[test] +fn should_support_full_sliceing_1d() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0]); + let data = tensor.to_data(); + + let output = tensor.slice([0..4]); + + output.into_data().assert_eq(&data, false); +} + +#[test] +fn should_support_partial_sliceing_1d() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0]); + + let output = tensor.slice([1..3]); + let expected = TensorData::from([1.0, 2.0]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn should_support_full_sliceing_2d() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let data = tensor.to_data(); + + let output = tensor.clone().slice([0..2]); + output.into_data().assert_eq(&data, true); + + let output = tensor.slice([0..2, 0..3]); + output.into_data().assert_eq(&data, true); +} + +#[test] +fn should_support_partial_sliceing_2d() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.slice([0..2, 0..2]); + let expected = TensorData::from([[0.0, 1.0], [3.0, 4.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn should_support_partial_sliceing_3d() { + let tensor = QTensor::<3>::int8([ + [[0., 1., 2., 3.], [4., 5., 6., 7.]], + [[8., 9., 10., 11.], [12., 13., 14., 15.]], + ]); + + let output = tensor.slice([1..2, 1..2, 0..2]); + let expected = TensorData::from([[[12.0, 13.0]]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn should_support_partial_sliceing_3d_non_contiguous() { + let tensor = QTensor::<3>::int8([ + [[0., 1., 2., 3.], [4., 5., 6., 7.]], + [[8., 9., 10., 11.], [12., 13., 14., 15.]], + ]); + + let output = tensor.transpose().slice([1..2, 1..2, 0..2]); + let expected = TensorData::from([[[9.0, 13.0]]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn should_support_slice_assign_1d() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0]); + let tensor_assigned = QTensor::<1>::int8([10.0, 5.0]); + + let output = tensor.slice_assign([0..2], tensor_assigned); + let expected = TensorData::from([10.0, 5.0, 2.0]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_support_slice_assign_2d() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_assigned = QTensor::<2>::int8([[10.0, 5.0]]); + + let output = tensor.slice_assign([1..2, 0..2], tensor_assigned); + let expected = TensorData::from([[0.0, 1.0, 2.0], [10.0, 5.0, 5.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn slice_should_not_corrupt_potentially_inplace_operations() { + let tensor = QTensor::<1>::int8([1.0, 2.0, 3.0, 4.0, 5.0]); + let tensor = tensor.clone().slice([0..3]) + tensor.clone().slice([2..5]); + + let expected = TensorData::from([4., 6., 8.]); + + tensor + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn slice_assign_should_not_corrupt_potentially_inplace_operations() { + let tensor = QTensor::<1>::int8([1.0, 2.0, 3.0, 4.0, 5.0]); + let values = QTensor::<1>::int8([10., 20., 30.]); + + let tensor_1 = tensor.clone().slice_assign([0..3], values); + let tensor_2 = tensor + 2; + + let expected = TensorData::from([10., 20., 30., 4., 5.]); + + tensor_1 + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); + + let expected = TensorData::from([3., 4., 5., 6., 7.]); + + tensor_2 + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn clamp_when_slice_exceeds_dimension() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0]); + let data = tensor.to_data(); + + let output = tensor.slice([0..4]); + output.into_data().assert_eq(&data, true); +} + +#[test] +fn negative_dimensions() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let data = tensor.to_data(); + + // Clamping to the tensor dimensions + let output = tensor.clone().slice([0..4, 0..4]); + output.into_data().assert_eq(&data, true); + + // Negative dimensions + let output = tensor.clone().slice([0..1, 0..1]); + let data = TensorData::from([[0.0f32]]); + output + .dequantize() + .into_data() + .assert_approx_eq::(&data, Tolerance::rel_abs(2e-2, 1e-2)); + + let output = tensor.slice(s![0..-1, 0..-2]); + output + .dequantize() + .into_data() + .assert_approx_eq::(&data, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +fn missing_dimensions() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let data = tensor.to_data(); + + // Clamping to the tensor dimensions + let output = tensor.clone().slice([0..4, 0..4]); + output.into_data().assert_eq(&data, true); + + // Negative dimensions + let data = TensorData::from([[0.0f32]]); + let output = tensor.clone().slice(s![0..-1, 0..-2]); + output + .dequantize() + .into_data() + .assert_approx_eq::(&data, Tolerance::rel_abs(2e-2, 1e-2)); + + // Missing dimensions + let output = tensor.clone().slice(s![0..1, ..]); + let data = TensorData::from([[0.0f32, 1.0, 2.0]]); + output + .dequantize() + .into_data() + .assert_approx_eq::(&data, Tolerance::rel_abs(2e-2, 1e-2)); + + let output = tensor.clone().slice(s![.., 0..2]); + let data = TensorData::from([[0.0f32, 1.0], [3.0, 4.0]]); + output + .dequantize() + .into_data() + .assert_approx_eq::(&data, Tolerance::rel_abs(2e-2, 1e-2)); + + let output = tensor.clone().slice([.., ..]); + let data = TensorData::from([[0.0f32, 1.0, 2.0], [3.0, 4.0, 5.0]]); + output + .dequantize() + .into_data() + .assert_approx_eq::(&data, Tolerance::rel_abs(2e-2, 1e-2)); +} + +#[test] +#[should_panic] +fn should_panic_when_slice_with_too_many_dimensions() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0]); + + let _output = tensor.slice([0..1, 0..1]); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/sort_argsort.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/sort_argsort.rs new file mode 100644 index 0000000..421fa8e --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/sort_argsort.rs @@ -0,0 +1,225 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn test_sort_1d_float() { + // Quantized [0.5, 1.2, -0.21, 0., 2.1, 0.94, -0.3, 2.3, 5.2, 4., 0.99, 3., -8.1] + let tensor = QTensor::<1>::int8([ + 0.5, 1.2, -0.21, 0., 2.1, 0.94, -0.3, 2.3, 5.2, 4., 0.99, 3., -8.1, + ]); + + // Sort along dim=0 + let values = tensor.sort(0); + + let values_expected = TensorData::from([ + -8.1, -0.3, -0.21, 0., 0.5, 0.94, 0.99, 1.2, 2.1, 2.3, 3., 4., 5.2, + ]); + + values + .dequantize() + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn test_argsort_1d_float() { + let tensor = QTensor::<1>::int8([ + 0.5, 1.2, -0.21, 0., 2.1, 0.94, -0.3, 2.3, 5.2, 4., 0.99, 3., -8.1, + ]); + + // Sort along dim=0 + let indices = tensor.argsort(0); + + let indices_expected = TensorData::from([12, 6, 2, 3, 0, 5, 10, 1, 4, 7, 11, 9, 8]); + indices.into_data().assert_eq(&indices_expected, false); +} + +#[test] +fn test_sort_with_indices_descending_float() { + // 1D + let tensor = QTensor::<1>::int8([ + 0.5, 1.2, -0.21, 0., 2.1, 0.94, -0.3, 2.3, 5.2, 4., 0.99, 3., -8.1, + ]); + + // Sort along dim=0 + let (values, indices) = tensor.sort_descending_with_indices(0); + + let values_expected = TensorData::from([ + 5.2, 4., 3., 2.3, 2.1, 1.2, 0.99, 0.94, 0.5, 0., -0.21, -0.3, -8.1, + ]); + + values + .dequantize() + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::absolute(1e-1)); + + let indices_expected = TensorData::from([8, 9, 11, 7, 4, 1, 10, 5, 0, 3, 2, 6, 12]); + indices.into_data().assert_eq(&indices_expected, false); + + // 3D + // Quantized [-0.5, 1.2, -0.21, 0., 2.1, 0.94, -0.3, 2.3, 4., 0.99, 3., -8.1] + let tensor = QTensor::<1>::int8([ + -0.5, 1.2, -0.21, 0., 2.1, 0.94, -0.3, 2.3, 4., 0.99, 3., -8.1, + ]) + .reshape([2, 2, 3]); + + // Sort along dim=1 + let (values, indices) = tensor.sort_descending_with_indices(1); + + let values_expected = TensorData::from([ + [[0., 2.1, 0.94], [-0.5, 1.2, -0.21]], + [[0.99, 3., 4.], [-0.3, 2.3, -8.1]], + ]); + + values + .dequantize() + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::absolute(1e-1)); + + let indices_expected = TensorData::from([[[1, 1, 1], [0, 0, 0]], [[1, 1, 0], [0, 0, 1]]]); + indices.into_data().assert_eq(&indices_expected, false); +} + +#[test] +fn test_sort_float() { + let tensor = QTensor::<1>::int8([ + -0.5, 1.2, -0.21, 0., 2.1, 0.94, -0.3, 2.3, 4., 0.99, 3., -8.1, + ]) + .reshape([2, 2, 3]); + + // Sort along dim=0 + let values = tensor.clone().sort(0); + + let values_expected = TensorData::from([ + [[-0.5, 1.2, -0.21], [0., 2.1, -8.1]], + [[-0.3, 2.3, 4.], [0.99, 3., 0.94]], + ]); + + values + .dequantize() + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::absolute(1e-1)); + + // Sort along dim=1 + let values = tensor.clone().sort(1); + + let values_expected = TensorData::from([ + [[-0.5, 1.2, -0.21], [0., 2.1, 0.94]], + [[-0.3, 2.3, -8.1], [0.99, 3., 4.]], + ]); + + values + .dequantize() + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::absolute(1e-1)); + + // Sort along dim=2 + let values = tensor.sort(2); + + let values_expected = TensorData::from([ + [[-0.5, -0.21, 1.2], [0., 0.94, 2.1]], + [[-0.3, 2.3, 4.], [-8.1, 0.99, 3.]], + ]); + + values + .dequantize() + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn test_sort_with_indices_float() { + let tensor = QTensor::<1>::int8([ + -0.5, 1.2, -0.21, 0., 2.1, 0.94, -0.3, 2.3, 4., 0.99, 3., -8.1, + ]) + .reshape([2, 2, 3]); + + // Sort along dim=0 + let (values, indices) = tensor.clone().sort_with_indices(0); + let values_expected = TensorData::from([ + [[-0.5, 1.2, -0.21], [0., 2.1, -8.1]], + [[-0.3, 2.3, 4.], [0.99, 3., 0.94]], + ]); + + values + .dequantize() + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::absolute(1e-1)); + + let indices_expected = TensorData::from([[[0, 0, 0], [0, 0, 1]], [[1, 1, 1], [1, 1, 0]]]); + indices.into_data().assert_eq(&indices_expected, false); + + // Sort along dim=1 + let (values, indices) = tensor.clone().sort_with_indices(1); + + let values_expected = TensorData::from([ + [[-0.5, 1.2, -0.21], [0., 2.1, 0.94]], + [[-0.3, 2.3, -8.1], [0.99, 3., 4.]], + ]); + + values + .dequantize() + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::absolute(1e-1)); + + let indices_expected = TensorData::from([[[0, 0, 0], [1, 1, 1]], [[0, 0, 1], [1, 1, 0]]]); + indices.into_data().assert_eq(&indices_expected, false); + + // Sort along dim=2 + let (values, indices) = tensor.sort_with_indices(2); + + let values_expected = TensorData::from([ + [[-0.5, -0.21, 1.2], [0., 0.94, 2.1]], + [[-0.3, 2.3, 4.], [-8.1, 0.99, 3.]], + ]); + + values + .dequantize() + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::absolute(1e-1)); + + let indices_expected = TensorData::from([[[0, 2, 1], [0, 2, 1]], [[0, 1, 2], [2, 0, 1]]]); + indices.into_data().assert_eq(&indices_expected, false); +} + +#[test] +fn test_argsort_float() { + let tensor = QTensor::<1>::int8([ + -0.5, 1.2, -0.21, 0., 2.1, 0.94, -0.3, 2.3, 4., 0.99, 3., -8.1, + ]) + .reshape([2, 2, 3]); + + // Sort along dim=0 + let indices = tensor.clone().argsort(0); + + let indices_expected = TensorData::from([[[0, 0, 0], [0, 0, 1]], [[1, 1, 1], [1, 1, 0]]]); + indices.into_data().assert_eq(&indices_expected, false); + + // Sort along dim=1 + let indices = tensor.clone().argsort(1); + + let indices_expected = TensorData::from([[[0, 0, 0], [1, 1, 1]], [[0, 0, 1], [1, 1, 0]]]); + indices.into_data().assert_eq(&indices_expected, false); + + // Sort along dim=2 + let indices = tensor.argsort(2); + + let indices_expected = TensorData::from([[[0, 2, 1], [0, 2, 1]], [[0, 1, 2], [2, 0, 1]]]); + indices.into_data().assert_eq(&indices_expected, false); +} + +#[test] +fn test_sort_descending_1d() { + let tensor = QTensor::<1>::int8([1.0, 2.0, 3.0, 4.0, 5.0]); + + // Sort along dim=0 + let values = tensor.sort_descending(0); + + let values_expected = TensorData::from([5., 4., 3., 2., 1.]); + values + .dequantize() + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::absolute(1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/split.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/split.rs new file mode 100644 index 0000000..033990b --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/split.rs @@ -0,0 +1,151 @@ +use super::qtensor::*; +use super::*; +use alloc::vec; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn test_split_evenly_divisible() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]); + + let tensors = tensor.split(2, 0); + assert_eq!(tensors.len(), 3); + + let expected = [ + TensorData::from([0., 1.]), + TensorData::from([2., 3.]), + TensorData::from([4., 5.]), + ]; + + for (index, tensor) in tensors.into_iter().enumerate() { + tensor + .dequantize() + .to_data() + .assert_approx_eq::(&expected[index], Tolerance::absolute(1e-1)); + } +} + +#[test] +fn test_split_not_evenly_divisible() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + + let tensors = tensor.split(2, 0); + assert_eq!(tensors.len(), 4); + + let expected = [ + TensorData::from([0., 1.]), + TensorData::from([2., 3.]), + TensorData::from([4., 5.]), + TensorData::from([6.]), + ]; + + for (index, tensor) in tensors.into_iter().enumerate() { + tensor + .dequantize() + .to_data() + .assert_approx_eq::(&expected[index], Tolerance::absolute(1e-1)); + } +} + +#[test] +fn test_split_along_dim1() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let tensors = tensor.split(2, 1); + assert_eq!(tensors.len(), 2); + + let expected = [ + TensorData::from([[0., 1.], [3., 4.]]), + TensorData::from([[2.], [5.]]), + ]; + + for (index, tensor) in tensors.into_iter().enumerate() { + tensor + .dequantize() + .to_data() + .assert_approx_eq::(&expected[index], Tolerance::absolute(1e-1)); + } +} + +#[test] +fn test_split_split_size_larger_than_tensor_size() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]); + + let tensors = tensor.split(10, 0); + assert_eq!(tensors.len(), 1); + + let expected = [TensorData::from([0.0, 1.0, 2.0, 3.0, 4.0, 5.0])]; + + for (index, tensor) in tensors.into_iter().enumerate() { + tensor + .dequantize() + .to_data() + .assert_approx_eq::(&expected[index], Tolerance::absolute(1e-1)); + } +} + +#[test] +#[should_panic( + expected = "split_size must be greater than 0 unless the tensor size along the dimension is 0." +)] +fn test_split_with_zero_split_size_non_zero_tensor() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]); + + let _ = tensor.split(0, 0); +} + +#[test] +#[should_panic(expected = "Given dimension is greater than or equal to the tensor rank.")] +fn test_split_invalid_dim() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]); + + let _ = tensor.split(1, 2); +} + +#[test] +fn test_split_with_sizes() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]); + + let tensors = tensor.split_with_sizes(vec![2, 3, 1], 0); + assert_eq!(tensors.len(), 3); + + let expected = [ + TensorData::from([0., 1.]), + TensorData::from([2., 3., 4.]), + TensorData::from([5.]), + ]; + + for (index, tensor) in tensors.into_iter().enumerate() { + tensor + .dequantize() + .to_data() + .assert_approx_eq::(&expected[index], Tolerance::absolute(1e-1)); + } +} + +#[test] +#[should_panic( + expected = "The sum of split_sizes must equal the tensor size along the specified dimension." +)] +fn test_split_with_sizes_invalid_sum() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0, 4.0, 5.0]); + + let _ = tensor.split_with_sizes(vec![2, 2, 1], 0); +} + +#[test] +fn test_split_with_sizes_zero_length() { + let tensor = QTensor::<1>::int8([0.0, 2.0, 5.0]); + + let tensors = tensor.split_with_sizes(vec![0, 1, 2], 0); + assert_eq!(tensors.len(), 2); + + let expected = [TensorData::from([0.]), TensorData::from([2., 5.])]; + + for (index, tensor) in tensors.into_iter().enumerate() { + tensor + .dequantize() + .to_data() + .assert_approx_eq::(&expected[index], Tolerance::absolute(1e-1)); + } +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/sqrt.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/sqrt.rs new file mode 100644 index 0000000..9e6fc32 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/sqrt.rs @@ -0,0 +1,18 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; +use core::f32::consts::SQRT_2; + +#[test] +fn should_support_sqrt_ops() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.sqrt(); + let expected = TensorData::from([[0.0, 1.0, SQRT_2], [1.73205, 2.0, 2.2360]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/stack.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/stack.rs new file mode 100644 index 0000000..ee96087 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/stack.rs @@ -0,0 +1,69 @@ +use super::qtensor::*; +use super::*; +use alloc::vec; +use burn_tensor::Tensor; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_stack_ops_2d_dim0() { + let tensor_1 = QTensor::<2>::int8([[1.0, 2.0, 3.0]]); + let tensor_2 = QTensor::<2>::int8([[4.0, 5.0, 6.0]]); + + let output = Tensor::stack::<3>(vec![tensor_1, tensor_2], 0); + let expected = TensorData::from([[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_support_stack_ops_2d_dim1() { + let tensor_1 = QTensor::<2>::int8([[1.0, 2.0, 3.0]]); + let tensor_2 = QTensor::<2>::int8([[4.0, 5.0, 6.0]]); + + let output = Tensor::stack::<3>(vec![tensor_1, tensor_2], 1); + let expected = TensorData::from([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_support_stack_ops_3d() { + let tensor_1 = QTensor::<3>::int8([[[1.0, 2.0, 3.0]], [[3.0, 2.0, 1.0]]]); + let tensor_2 = QTensor::<3>::int8([[[4.0, 5.0, 6.0]], [[6.0, 5.0, 4.0]]]); + + let output = Tensor::stack::<4>(vec![tensor_1, tensor_2], 0); + let expected = TensorData::from([ + [[[1.0, 2.0, 3.0]], [[3.0, 2.0, 1.0]]], + [[[4.0, 5.0, 6.0]], [[6.0, 5.0, 4.0]]], + ]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +#[should_panic] +fn should_panic_when_dimensions_are_not_the_same() { + let tensor_1 = QTensor::<2>::int8([[1.0, 2.0, 3.0]]); + let tensor_2 = QTensor::<2>::int8([[4.0, 5.0]]); + + let _output = Tensor::stack::<3>(vec![tensor_1, tensor_2], 0); +} + +#[test] +#[should_panic] +fn should_panic_when_stack_exceeds_dimension() { + let tensor_1 = QTensor::<3>::int8([[[1.0, 2.0, 3.0]], [[3.0, 2.0, 1.0]]]); + let tensor_2 = QTensor::<3>::int8([[[4.0, 5.0, 6.0]]]); + + let _output = Tensor::stack::<4>(vec![tensor_1, tensor_2], 3); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/sub.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/sub.rs new file mode 100644 index 0000000..1e516ca --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/sub.rs @@ -0,0 +1,46 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_sub_ops() { + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let tensor_2 = QTensor::<2>::int8([[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]]); + + let output = tensor_1 - tensor_2; + let expected = TensorData::from([[-6.0, -6.0, -6.0], [-6.0, -6.0, -6.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn test_sub_broadcast() { + let tensor_1 = QTensor::<2>::int8([[0.0, 1.0, 2.0]]); + let tensor_2 = QTensor::<2>::int8([[3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]); + + let output = tensor_1 - tensor_2; + let expected = TensorData::from([[-3.0, -3.0, -3.0], [-6.0, -6.0, -6.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_support_sub_scalar_ops() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + let scalar = 2.0; + + let output = tensor - scalar; + let expected = TensorData::from([[-2.0, -1.0, 0.0], [1.0, 2.0, 3.0]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(2e-2, 1e-2)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/tan.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/tan.rs new file mode 100644 index 0000000..7b0ebc0 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/tan.rs @@ -0,0 +1,17 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_tan_ops() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.tan(); + let expected = TensorData::from([[0.0, 1.5574, -2.1850], [-0.1425, 1.1578, -3.3805]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/tanh.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/tanh.rs new file mode 100644 index 0000000..b3f727d --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/tanh.rs @@ -0,0 +1,17 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_tanh_ops() { + let tensor = QTensor::<2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]); + + let output = tensor.tanh(); + let expected = TensorData::from([[0.0, 0.7615, 0.9640], [0.9950, 0.9993, 0.9999]]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/topk.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/topk.rs new file mode 100644 index 0000000..0435baf --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/topk.rs @@ -0,0 +1,68 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +#[ignore] +fn test_topk_1d() { + let tensor = QTensor::<1>::int8([1.0, 2.0, 3.0, 4.0, 5.0]); + + let values = tensor.topk(3, /*dim*/ 0); + let expected = TensorData::from([5., 4., 3.]); + + values + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +#[ignore] +fn test_topk() { + let tensor = QTensor::<3>::int8([[[1., 4., 7.], [2., 5., 6.]], [[3., 0., 9.], [8., 2., 7.]]]); + + let values = tensor.topk(2, /*dim*/ 2); + let expected = TensorData::from([[[7., 4.], [6., 5.]], [[9., 3.], [8., 7.]]]); + + values + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn test_topk_with_indices() { + // 1D + let tensor = QTensor::<1>::int8([1.0, 2.0, 3.0, 4.0, 5.0]); + + let (values, indices) = tensor.topk_with_indices(3, /*dim*/ 0); + + let values_expected = TensorData::from([5., 4., 3.]); + values + .dequantize() + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::permissive()); + + let indices_expected = TensorData::from([4, 3, 2]); + indices.into_data().assert_eq(&indices_expected, false); +} + +#[test] +fn test_topk_with_indices_3d() { + // 3D + let tensor = QTensor::<3>::int8([[[1., 4., 7.], [2., 5., 6.]], [[3., 0., 9.], [8., 2., 7.]]]); + + let (values, indices) = tensor.topk_with_indices(2, /*dim*/ 2); + + let values_expected = TensorData::from([[[7., 4.], [6., 5.]], [[9., 3.], [8., 7.]]]); + + values + .dequantize() + .into_data() + .assert_approx_eq::(&values_expected, Tolerance::absolute(1e-1)); + + let indices_expected = TensorData::from([[[2, 1], [2, 1]], [[2, 0], [0, 2]]]); + + indices.into_data().assert_eq(&indices_expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/transpose.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/transpose.rs new file mode 100644 index 0000000..4ede156 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/transpose.rs @@ -0,0 +1,39 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn should_support_transpose_ops() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0]) + .reshape([2, 2, 3]); + + let output = tensor.transpose(); + let expected = TensorData::from([ + [[0.0, 3.0], [1.0, 4.0], [2.0, 5.0]], + [[6.0, 9.0], [7.0, 10.0], [8.0, 11.0]], + ]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} + +#[test] +fn should_support_swap_dims() { + let tensor = QTensor::<1>::int8([0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0]) + .reshape([2, 2, 3]); + + let output = tensor.swap_dims(0, 2); + let expected = TensorData::from([ + [[0.0, 6.0], [3.0, 9.0]], + [[1.0, 7.0], [4.0, 10.0]], + [[2.0, 8.0], [5.0, 11.0]], + ]); + + output + .dequantize() + .into_data() + .assert_approx_eq::(&expected, Tolerance::absolute(1e-1)); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/matmul.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/matmul.rs new file mode 100644 index 0000000..ba15b4d --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/matmul.rs @@ -0,0 +1,248 @@ +use super::qtensor::*; +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +#[ignore] +fn test_matmul_vectors() { + let tensor_1 = QTensor::<2>::int8([[1.0, 2.0, 3.0, 6.35]]); + let tensor_2 = QTensor::<2>::int8([[12.7], [4.0], [5.0], [1.0]]); + + let tensor_3 = tensor_1.matmul(tensor_2); + + let expected = TensorData::from([[42.05]]); + tensor_3 + .into_data() + .assert_approx_eq::(&expected, Tolerance::relative(2e-2)); +} + +#[test] +#[ignore] +fn test_matmul_2d() { + let tensor_1 = QTensor::<2>::int8([[1.0, 6.35], [2.0, 3.0], [1.0, 3.0]]); + let tensor_2 = QTensor::<2>::int8([[4.0, 8.0, 12.7], [2.0, 3.0, 6.0]]); + let tensor_3 = tensor_1.matmul(tensor_2); + + let expected = TensorData::from([[16.7, 27.05, 50.8], [14., 25., 43.4], [10., 17., 30.7]]); + tensor_3 + .into_data() + .assert_approx_eq::(&expected, Tolerance::relative(2e-2)); +} + +#[test] +fn test_matmul_2d_aligned() { + let tensor_1 = QTensor::<2>::int8([ + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + ]); + let tensor_2 = QTensor::<2>::int8([ + [2.0, 0.0, 1.0, 0.0], + [1.0, 2.0, 0.0, 0.0], + [0.0, 1.0, 2.0, 0.0], + [1.0, 0.0, 0.0, 1.0], + ]); + let tensor_3 = tensor_1.matmul(tensor_2); + + let expected = TensorData::from([ + [8.0, 7.0, 7.0, 4.0], + [24.0, 19.0, 19.0, 8.0], + [40.0, 31.0, 31.0, 12.0], + ]); + tensor_3 + .into_data() + .assert_approx_eq::(&expected, Tolerance::relative(2e-2)); +} + +#[test] +fn test_matmul_2d_aligned_fused() { + let tensor_1 = QTensor::<2>::int8([ + [1.0, 2.0, 3.0, 4.0], + [5.0, 6.0, 7.0, 8.0], + [9.0, 10.0, 11.0, 12.0], + ]); + let tensor_2 = QTensor::<2>::int8([ + [2.0, 0.0, 1.0, 0.0], + [1.0, 2.0, 0.0, 0.0], + [0.0, 1.0, 2.0, 0.0], + [1.0, 0.0, 0.0, 1.0], + ]); + let tensor_3 = tensor_1.matmul(tensor_2); + let tensor_4 = tensor_3 / 2.0; + + let expected = TensorData::from([ + [4.0, 3.5, 3.5, 2.0], + [12.0, 9.5, 9.5, 4.0], + [20.0, 15.5, 15.5, 6.0], + ]); + tensor_4 + .into_data() + .assert_approx_eq::(&expected, Tolerance::relative(2e-2)); +} + +#[test] +#[ignore] +fn test_matmul_3d() { + let tensor_1 = QTensor::<3>::int8([[[1.0, 6.35], [2.0, 3.0]]]); + let tensor_2 = QTensor::<3>::int8([[[12.7, 4.0], [2.0, 3.0]]]); + + let tensor_3 = tensor_1.matmul(tensor_2); + + let expected = TensorData::from([[[25.4, 23.05], [31.4, 17.0]]]); + tensor_3 + .into_data() + .assert_approx_eq::(&expected, Tolerance::relative(2e-2)); +} + +#[test] +#[ignore] +fn test_matmul_broadcast_4d() { + let tensor_1 = QTensor::<4>::int8([[[[1.0, 7.0], [2.0, 3.0]]], [[[2.0, 5.0], [6.0, 3.0]]]]); + let tensor_2 = QTensor::<4>::int8([[[[9.0, 8.0], [1.0, 4.0]], [[2.0, 7.0], [3.0, 5.0]]]]); + + // [2, 1, 2, 2] @ [1, 2, 2, 2] -> [2, 2, 2, 2] + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([ + [[[16.0, 36.0], [21.0, 28.0]], [[23.0, 42.0], [13.0, 29.0]]], + [[[23.0, 36.0], [57.0, 60.0]], [[19.0, 39.0], [21.0, 57.0]]], + ]); + + tensor_3 + .into_data() + .assert_approx_eq::(&expected, Tolerance::relative(2e-2)); +} + +#[test] +#[ignore] +fn test_matmul_broadcast() { + let tensor_1 = QTensor::<3>::int8([[[1.0, 7.0], [2.0, 3.0]]]); + let tensor_2 = QTensor::<3>::int8([[[4.0, 7.0], [2.0, 3.0]], [[2.0, 5.0], [6.0, 3.0]]]); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([[[18.0, 28.0], [14.0, 23.0]], [[44.0, 26.0], [22.0, 19.0]]]); + + tensor_3 + .into_data() + .assert_approx_eq::(&expected, Tolerance::relative(2e-2)); +} + +#[test] +#[should_panic] +fn should_panic_when_inner_dimensions_are_not_equal() { + let tensor_1 = QTensor::<2>::int8([[3., 3.], [4., 4.], [5., 5.], [6., 6.]]); + let tensor_2 = QTensor::<2>::int8([[1., 2., 3., 4.], [1., 2., 3., 4.], [1., 2., 3., 4.]]); + + let _ = tensor_1.matmul(tensor_2); +} + +#[test] +fn test_matmul_lhs_float_rhs_quantized() { + // Simulates a typical workflow with linear layers (e.g., transformers), where the rhs + // represents the weights. The lhs might be a float if a previous operation did not propagate + // the quantization. We still want to perform an efficient matmul with quantized weights. + let tensor_1 = TestTensor::<2>::from([ + [1.0, 6.35, 2.0, 3.0], + [2.0, 3.0, 4.0, 5.0], + [1.0, 3.0, 5.0, 7.0], + ]); + let tensor_2 = QTensor::<2>::int8([ + [4.0, 8.0, 12.7, 1.6], + [2.0, 3.0, 6.0, 4.0], + [1.0, 5.0, 9.0, 2.5], + [3.0, 7.0, 11.0, 0.5], + ]); + let tensor_3 = tensor_1.matmul(tensor_2); + + let expected = TensorData::from([ + [27.7, 58.05, 101.8, 33.5], + [33., 80., 134.4, 27.7], + [36., 91., 152.7, 29.6], + ]); + let output = tensor_3.into_data(); + output.assert_approx_eq::(&expected, Tolerance::default()); + + // Default quantization scheme does not propagate quantization with matmul + assert!(output.dtype.is_float()); +} + +#[test] +fn test_matmul_lhs_quantized_rhs_float() { + // Mirror of test_matmul_lhs_float_rhs_quantized: pins the symmetric + // (QFloat, Float) arm of the q_matmul dispatch so a future split of the + // collapsed match arm cannot regress one direction silently. Uses the + // relative tolerance shared by the other quantized matmul tests in this + // file because the lhs dynamic range here is smaller than the rhs in the + // sibling test, so int8 noise dominates the result more. + let tensor_1 = QTensor::<2>::int8([ + [1.0, 6.35, 2.0, 3.0], + [2.0, 3.0, 4.0, 5.0], + [1.0, 3.0, 5.0, 7.0], + ]); + let tensor_2 = TestTensor::<2>::from([ + [4.0, 8.0, 12.7, 1.6], + [2.0, 3.0, 6.0, 4.0], + [1.0, 5.0, 9.0, 2.5], + [3.0, 7.0, 11.0, 0.5], + ]); + let tensor_3 = tensor_1.matmul(tensor_2); + + let expected = TensorData::from([ + [27.7, 58.05, 101.8, 33.5], + [33., 80., 134.4, 27.7], + [36., 91., 152.7, 29.6], + ]); + let output = tensor_3.into_data(); + output.assert_approx_eq::(&expected, Tolerance::relative(2e-2)); + + // Default quantization scheme does not propagate quantization with matmul + assert!(output.dtype.is_float()); +} + +#[test] +fn test_matmul_mixed_block_scale() { + let tensor_1 = TestTensor::<2>::from([ + [1.0, 6.35, 2.0, 3.0], + [2.0, 3.0, 4.0, 5.0], + [1.0, 3.0, 5.0, 7.0], + ]); + let tensor_2 = QTensor::<2>::int8_block([ + [ + 6.110, 4.0, 9.360, 7.850, 0.630, 1.770, 0.430, 7.550, 9.690, 3.560, 2.920, 9.130, + 3.390, 0.510, 1.620, 1.460, + ], + [ + 6.140, 8.260, 5.660, 5.610, 7.070, 3.050, 9.890, 5.520, 1.350, 3.810, 5.630, 0.250, + 0.350, 8.860, 3.610, 6.240, + ], + [ + 8.810, 4.620, 7.420, 8.110, 2.560, 4.710, 5.730, 8.980, 1.170, 6.090, 4.140, 3.610, + 4.960, 9.720, 5.710, 1.470, + ], + [ + 2.260, 9.640, 6.320, 6.980, 9.860, 1.030, 8.340, 1.570, 4.140, 4.760, 4.590, 6.400, + 5.350, 1.430, 4.960, 1.180, + ], + ]); + let tensor_3 = tensor_1.matmul(tensor_2); + + let expected = TensorData::from([ + [ + 69.499, 94.611, 79.101, 80.633, 80.225, 33.647, 99.711, 65.272, 33.022, 54.213, 60.721, + 37.138, 31.582, 80.501, 50.843, 47.564, + ], + [ + 77.180, 99.460, 96.980, 99.870, 82.010, 36.680, 95.150, 75.430, 48.810, 66.710, 62.240, + 65.450, 54.420, 73.630, 61.710, 33.420, + ], + [ + 84.400, 119.360, 107.680, 114.090, 103.660, 41.680, 117.130, 80.0, 48.570, 78.760, + 72.640, 72.730, 66.690, 85.700, 75.720, 35.790, + ], + ]); + let output = tensor_3.into_data(); + output.assert_approx_eq::(&expected, Tolerance::permissive()); + + // Default quantization scheme does not propagate quantization with matmul + assert!(output.dtype.is_float()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/mod.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/mod.rs new file mode 100644 index 0000000..596c0b6 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/mod.rs @@ -0,0 +1,31 @@ +pub use super::*; + +mod matmul; +mod quantize; + +// The `extended` suite is only enabled for backends with native (non-packed) quantized +// storage AND a complete set of quantized ops. Today that means ndarray and flex. +// +// - cube backends are excluded: PackedU32 storage requires the last dim to be a multiple of +// the pack factor (4 int8s per u32), which most of these test shapes violate, so the +// quantized tensors can't even be constructed (`q_from_data` panics with "Can't store in u32"). +// +// - candle, router, tch and autodiff are excluded too. We dropped their `unimplemented!()` +// overrides for `q_gather`/`q_select`/`q_slice`/`q_expand` so they fall back to the default +// `dequantize -> float op -> quantize` path, but that does NOT make them functional: their +// other quantized methods are largely unimplemented. The quantization primitives themselves +// (`q_from_data`, `quantize`, `dequantize`, ...) are still `unimplemented!()`/`todo!()`, so the +// fallback simply moves the panic into `dequantize`. Running `extended` against any of these +// backends would fail. (They also don't enable the `quantization` feature, so they aren't +// selected here in the first place.) +// +// Enabling `flex` here means the `extended` suite now also runs under the `tensor_f16` target. +// That f16 path is why a couple of `maxmin` tests use a slightly +// looser tolerance: reductions like `min_dim` re-quantize their output, so a value is rounded to +// int8 twice (input quantization, then re-quantization of the reduced result). For small-magnitude +// values this accumulated rounding lands a hair over the tight `rel_abs(2e-2, 1e-2)` bound at f16 +// (e.g. `1.0` -> ~0.97998, rel error 2.00e-2), whereas f32's finer scale representation keeps the +// same value just inside it (~0.9802, rel error 1.98e-2). It is quantization noise, not a logic +// error, so those cases use `rel_abs(2e-2, 3e-2)`. +#[cfg(any(feature = "ndarray", feature = "flex"))] +mod extended; diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/ops/quantize.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/quantize.rs new file mode 100644 index 0000000..d930c07 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/ops/quantize.rs @@ -0,0 +1,255 @@ +use super::*; +use alloc::{vec, vec::Vec}; +use burn_tensor::Tolerance; +use burn_tensor::quantization::{ + QParams, QuantLevel, QuantScheme, QuantStore, QuantValue, QuantizationParameters, + QuantizedBytes, +}; +use burn_tensor::{DType, Element, TensorData}; + +fn get_q_params(data: TensorData) -> QParams> { + let num_elements = data.num_elements(); + let scheme = if let DType::QFloat(scheme) = data.dtype { + scheme + } else { + unreachable!() + }; + let q_bytes = QuantizedBytes { + bytes: data.into_bytes(), + scheme, + num_elements, + }; + q_bytes.into_vec_i8().1 +} + +#[test] +fn should_support_quantize_symmetric_int8() { + // Strict equality was based on full precision + if !matches!(FloatElem::dtype(), DType::F32) { + return; + } + let device = Default::default(); + let tensor = TestTensor::<1>::from_data([-1.8, -1.0, 0.0, 0.5], &device); + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S); + let qparams = QuantizationParameters { + scales: TestTensor::from_data([0.014_173_228], &device), + }; + + let x_q = tensor.clone().quantize(&scheme, qparams); + + let x_q_data = x_q.to_data(); + let expected = TensorData::quantized( + vec![-127i8, -71, 0, 35], + [4], + scheme.with_store(QuantStore::Native), + &[0.014_173_228], // scale + ); + + // Values equality + x_q_data.assert_eq(&expected, false); + + // Quantization parameters check + let qparams = get_q_params(x_q_data); + let expected = get_q_params(expected); + assert_eq!(qparams.scales.len(), 1); + // TODO: check scales + assert_eq!(qparams.scales, expected.scales); + + // Dequantize + let x = x_q.dequantize(); + + x.into_data() + .assert_approx_eq::(&tensor.into_data(), Tolerance::rel_abs(1e-1, 1e-2)); +} + +#[test] +fn should_support_quantize_dynamic_int8() { + let device = Default::default(); + // NOTE: we use fully representable values since different backend implementations could differ slightly + // due to rounding discrepancies + let tensor = TestTensor::<1>::from_data([5., 0., 4., -12.7], &device); + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S); + + let x_q = tensor.quantize_dynamic(&scheme); + + let expected = TensorData::quantized( + vec![50i8, 0, 40, -127], + [4], + scheme.with_store(QuantStore::Native), + &[0.1], // scale + ); + + x_q.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_quantize_dequantize_symmetric_single_with_transform() { + let device = Default::default(); + let input = TestTensorInt::<1>::arange(0..32, &device).float(); + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S); + + let quant = input.quantize_dynamic(&scheme); + let result = quant * 10; + + let data = result.into_data(); + let expected = [ + 0.0, 9.76378, 19.52756, 29.29134, 39.05512, 48.818897, 61.02362, 70.7874, 80.551186, + 90.31496, 100.07874, 109.84252, 119.60631, 129.37009, 139.13387, 148.89764, 161.10237, + 170.86615, 180.62991, 190.39369, 200.15749, 209.92126, 219.68504, 229.44882, 239.21262, + 248.97638, 261.1811, 270.9449, 280.70865, 290.47244, 300.23624, 310.0, + ]; + data.assert_approx_eq::(&TensorData::from(expected), Tolerance::permissive()); +} + +#[test] +fn should_quantize_dequantize_symmetric_arange_16x16() { + let device = Default::default(); + + let input: TestTensor<2> = TestTensorInt::arange(0..256, &device) + .float() + .div_scalar(256.) + .reshape([16, 16]); + + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S); + let output = input.clone().quantize_dynamic(&scheme); + let output = output.dequantize(); + + output.into_data().assert_approx_eq::( + &input.into_data(), + Tolerance::absolute(1e-1).set_relative(1e-2), + ); +} + +#[test] +fn should_quantize_dequantize_symmetric_per_block_arange_16x16() { + let device = Default::default(); + + let input: TestTensor<2> = TestTensorInt::arange(0..256, &device) + .float() + .div_scalar(256.) + .reshape([16, 16]); + + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::block([2, 16])); + + let output = input.clone().quantize_dynamic(&scheme); + let output = output.dequantize(); + + output.into_data().assert_approx_eq::( + &input.into_data(), + Tolerance::absolute(1e-1).set_relative(1e-2), + ); +} + +fn should_quantize_transposed(tensor: Tensor, scheme: QuantScheme) { + let tensor_t = tensor.clone().transpose(); + + let output = tensor_t.quantize_dynamic(&scheme).dequantize().transpose(); + + tensor.into_data().assert_approx_eq::( + &output.into_data(), + Tolerance::absolute(1e-1).set_relative(1e-2), + ); +} + +#[test] +fn should_quantize_symmetric_int8_transposed_8x32() { + let device = Default::default(); + + let tensor = TestTensorInt::arange(0..256, &device) + .float() + .div_scalar(256.) + .reshape([8, 32]); + + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S); + should_quantize_transposed(tensor, scheme); +} + +#[test] +fn should_quantize_symmetric_int8_transposed_48x64() { + let device = Default::default(); + + let tensor = TestTensorInt::arange(0..3072, &device) + .float() + .div_scalar(3072.) + .reshape([48, 64]); + + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S); + should_quantize_transposed(tensor, scheme); +} + +#[test] +fn should_quantize_symmetric_per_block_int8_transposed_32x64() { + let device = Default::default(); + + let tensor = TestTensorInt::arange(0..2048, &device) + .float() + .div_scalar(2048.) + .reshape([32, 64]); + + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::block([32])); + should_quantize_transposed(tensor, scheme); +} + +#[test] +fn should_quantize_symmetric_int8_permuted_batch_dims() { + let device = Default::default(); + + let tensor = TestTensorInt::arange(0..2048, &device) + .float() + .div_scalar(2048.) + .reshape([2, 4, 8, 32]); + + // Permute [0,1,2,3] -> [1,2,0,3] + // This rearranges batch dims but keeps packed dim in place + let tensor_permuted = tensor.clone().permute([1, 2, 0, 3]); + + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S); + + let output = tensor_permuted + .quantize_dynamic(&scheme) + .dequantize() + .permute([2, 0, 1, 3]); // reverse permutation + + tensor.into_data().assert_approx_eq::( + &output.into_data(), + Tolerance::absolute(1e-1).set_relative(1e-2), + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/quantization/scheme.rs b/crates/burn-backend-tests/tests/tensor/float/quantization/scheme.rs new file mode 100644 index 0000000..f975afa --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/quantization/scheme.rs @@ -0,0 +1,81 @@ +use super::*; +use burn_tensor::Tolerance; +use burn_tensor::{ + Device, Element, TensorData, + quantization::{CalibrationRange, QuantLevel, QuantValue, compute_q_params}, +}; + +#[test] +fn per_tensor_symmetric_int8() { + let device = Default::default(); + let range = CalibrationRange { + min: TestTensor::<1>::from_data([0.5], &device), + max: TestTensor::<1>::from_data([1.8], &device), + }; + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S); + + let qparams = compute_q_params(&scheme, range); + + qparams + .scales + .into_data() + .assert_approx_eq::(&TensorData::from([0.014_173_23]), Tolerance::default()); +} + +#[test] +fn per_block_symmetric_int8() { + let device = Default::default(); + let range = CalibrationRange { + min: TestTensor::<1>::from_data([-1.8, -0.5, 0.01, -0.04], &device), + max: TestTensor::<1>::from_data([0.5, 1.8, 0.04, -0.01], &device), + }; + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::block([4])); + + let qparams = compute_q_params(&scheme, range); + + qparams.scales.into_data().assert_approx_eq::( + &TensorData::from([0.014_173_23, 0.014_173_23, 0.000_314_96, 0.000_314_96]), + Tolerance::default(), + ); +} + +#[test] +fn quant_scheme_should_inhibit_by_default() { + let device = Device::default(); + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S); + + let tensor_1 = TestTensor::<2>::from_data( + [[1.0, 6.35, 0., 0.], [2.0, 3.0, 0., 0.], [1.0, 3.0, 0., 0.]], + &device, + ) + .quantize_dynamic(&scheme); + let _tensor_2 = TestTensor::<2>::from_data( + [ + [4.0, 8.0, 12.7, 0.], + [2.0, 3.0, 6.0, 0.], + [0., 0., 0., 0.], + [0., 0., 0., 0.], + ], + &device, + ) + .quantize_dynamic(&scheme); + + // let tensor_3 = tensor_1.clone().matmul(tensor_2); + // assert_eq!(tensor_3.to_data().dtype, FloatElem::dtype()); + + let tensor_4 = tensor_1.add_scalar(1.); + assert_eq!(tensor_4.to_data().dtype, FloatElem::dtype()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/stats/cov.rs b/crates/burn-backend-tests/tests/tensor/float/stats/cov.rs new file mode 100644 index 0000000..e0b4075 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/stats/cov.rs @@ -0,0 +1,66 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn test_cov_1() { + let data = TensorData::from([[0.5, 1.8, 0.2, -2.0], [3.0, -4.0, 5.0, 0.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.cov(1, 1); + let expected = + TensorData::from([[2.48917, -1.73333], [-1.73333, 15.33333]]).convert::(); + + let tolerance = Tolerance::default().set_half_precision_relative(1e-3); + output + .into_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn test_cov_4() { + let data = TensorData::from([[0.5, 1.8, 0.2, -2.0], [3.0, -4.0, 5.0, 0.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.cov(1, 0); + let expected = TensorData::from([[1.86687, -1.30000], [-1.30000, 11.5]]).convert::(); + + let tolerance = Tolerance::default().set_half_precision_relative(1e-3); + output + .into_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn test_cov_2() { + let data = TensorData::from([[0.5, 1.8], [0.2, -2.0], [3.0, -4.0], [5.0, 0.0]]); + let tensor = TestTensor::<2>::from_data(data, &Default::default()); + + let output = tensor.cov(1, 1); + let expected = TensorData::from([ + [0.845, -1.43, -4.55, -3.25], + [-1.43, 2.42, 7.7, 5.5], + [-4.55, 7.7, 24.5, 17.5], + [-3.25, 5.5, 17.5, 12.5], + ]) + .convert::(); + + let tolerance = Tolerance::default().set_half_precision_relative(1e-3); + output + .into_data() + .assert_approx_eq::(&expected, tolerance); +} + +#[test] +fn test_cov_3() { + let data = TensorData::from([ + [[0.5, 1.8, 0.2, -2.0], [3.0, -4.0, 5.0, 0.0]], + [[0.5, 1.8, 0.2, -2.0], [3.0, -4.0, 5.0, 0.0]], + [[0.5, 1.8, 0.2, -2.0], [3.0, -4.0, 5.0, 0.0]], + [[0.5, 1.8, 0.2, -2.0], [3.0, -4.0, 5.0, 0.0]], + ]); + let device = Default::default(); + let tensor = TestTensor::<3>::from_data(data, &device); + let data_actual = tensor.cov(0, 1).into_data(); + let data_expected = TestTensor::<3>::zeros([4, 4, 4], &device).to_data(); + data_expected.assert_approx_eq::(&data_actual, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/stats/display.rs b/crates/burn-backend-tests/tests/tensor/float/stats/display.rs new file mode 100644 index 0000000..fc21a48 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/stats/display.rs @@ -0,0 +1,321 @@ +use super::*; +use burn_tensor::{DType, Element, Shape, TensorData}; + +// Floating point values might not match for other precisions +fn skip_precision_not_f32() -> bool { + core::any::TypeId::of::() != core::any::TypeId::of::() +} + +#[test] +fn test_display_2d_int_tensor() { + let int_data = TensorData::from([[1, 2, 3], [4, 5, 6], [7, 8, 9]]); + let tensor_int = TestTensorInt::<2>::from_data(int_data, &Default::default()); + + let output = format!("{}", tensor_int); + let expected = format!( + r#"Tensor {{ + data: +[[1, 2, 3], + [4, 5, 6], + [7, 8, 9]], + shape: [3, 3], + device: {:?}, + kind: "Int", + dtype: "{dtype}", +}}"#, + tensor_int.device(), + dtype = core::any::type_name::(), + ); + assert_eq!(output, expected); +} + +#[test] +fn test_display_2d_float_tensor() { + if skip_precision_not_f32() { + return; + } + + let float_data = TensorData::from([[1.1, 2.2, 3.3], [4.4, 5.5, 6.6], [7.7, 8.8, 9.9]]); + let tensor_float = TestTensor::<2>::from_data(float_data, &Default::default()); + + let output = format!("{}", tensor_float); + let expected = format!( + r#"Tensor {{ + data: +[[1.1, 2.2, 3.3], + [4.4, 5.5, 6.6], + [7.7, 8.8, 9.9]], + shape: [3, 3], + device: {:?}, + kind: "Float", + dtype: "f32", +}}"#, + tensor_float.device(), + ); + assert_eq!(output, expected); +} + +#[test] +fn test_display_2d_bool_tensor() { + let device = Default::default(); + let bool_data = TensorData::from([ + [true, false, true], + [false, true, false], + [false, true, true], + ]); + let tensor_bool = TestTensorBool::<2>::from_data(bool_data, &device); + + let output = format!("{}", tensor_bool); + // TODO: remove once backends no longer rely on generics for default elem types + let expected_dtype_name = match device.settings().bool_dtype { + burn_tensor::BoolDType::Native => DType::Bool(burn_tensor::BoolStore::Native).name(), + burn_tensor::BoolDType::U8 => DType::Bool(burn_tensor::BoolStore::U8).name(), + burn_tensor::BoolDType::U32 => DType::Bool(burn_tensor::BoolStore::U32).name(), + }; + let expected = format!( + r#"Tensor {{ + data: +[[true, false, true], + [false, true, false], + [false, true, true]], + shape: [3, 3], + device: {:?}, + kind: "Bool", + dtype: {:?}, +}}"#, + tensor_bool.device(), + expected_dtype_name, + ); + assert_eq!(output, expected); +} + +#[test] +fn test_display_3d_tensor() { + let data = TensorData::from([ + [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], + [[13, 14, 15, 16], [17, 18, 19, 20], [21, 22, 23, 24]], + ]); + let tensor = TestTensorInt::<3>::from_data(data, &Default::default()); + + let output = format!("{}", tensor); + let expected = format!( + r#"Tensor {{ + data: +[[[1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12]], + [[13, 14, 15, 16], + [17, 18, 19, 20], + [21, 22, 23, 24]]], + shape: [2, 3, 4], + device: {:?}, + kind: "Int", + dtype: "{dtype}", +}}"#, + tensor.device(), + dtype = core::any::type_name::(), + ); + assert_eq!(output, expected); +} + +#[test] +fn test_display_4d_tensor() { + let data = TensorData::from([ + [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], + [[[13, 14, 15], [16, 17, 18]], [[19, 20, 21], [22, 23, 24]]], + ]); + + let tensor = TestTensorInt::<4>::from_data(data, &Default::default()); + + let output = format!("{}", tensor); + let expected = format!( + r#"Tensor {{ + data: +[[[[1, 2, 3], + [4, 5, 6]], + [[7, 8, 9], + [10, 11, 12]]], + [[[13, 14, 15], + [16, 17, 18]], + [[19, 20, 21], + [22, 23, 24]]]], + shape: [2, 2, 2, 3], + device: {:?}, + kind: "Int", + dtype: "{dtype}", +}}"#, + tensor.device(), + dtype = core::any::type_name::(), + ); + assert_eq!(output, expected); +} + +#[test] +fn test_display_tensor_summarize_1() { + let tensor = TestTensor::<4>::zeros(Shape::new([2, 2, 2, 1000]), &Default::default()); + + let output = format!("{}", tensor); + let expected = format!( + r#"Tensor {{ + data: +[[[[0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0]], + [[0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0]]], + [[[0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0]], + [[0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0]]]], + shape: [2, 2, 2, 1000], + device: {:?}, + kind: "Float", + dtype: "{dtype}", +}}"#, + tensor.device(), + dtype = FloatElem::dtype().name(), + ); + assert_eq!(output, expected); +} + +#[test] +fn test_display_tensor_summarize_2() { + let tensor = TestTensor::<4>::zeros(Shape::new([2, 2, 20, 100]), &Default::default()); + + let output = format!("{}", tensor); + let expected = format!( + r#"Tensor {{ + data: +[[[[0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + ... + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0]], + [[0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + ... + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0]]], + [[[0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + ... + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0]], + [[0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + ... + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, ..., 0.0, 0.0, 0.0]]]], + shape: [2, 2, 20, 100], + device: {:?}, + kind: "Float", + dtype: "{dtype}", +}}"#, + tensor.device(), + dtype = FloatElem::dtype().name(), + ); + assert_eq!(output, expected); +} + +#[test] +fn test_display_tensor_summarize_3() { + let tensor = TestTensor::<4>::zeros(Shape::new([2, 2, 200, 6]), &Default::default()); + + let output = format!("{}", tensor); + let expected = format!( + r#"Tensor {{ + data: +[[[[0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ... + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]], + [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ... + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]], + [[[0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ... + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]], + [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ... + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]]], + shape: [2, 2, 200, 6], + device: {:?}, + kind: "Float", + dtype: "{dtype}", +}}"#, + tensor.device(), + dtype = FloatElem::dtype().name(), + ); + assert_eq!(output, expected); +} +#[test] +fn test_display_precision() { + if skip_precision_not_f32() { + return; + } + + let tensor = TestTensor::<2>::full([1, 1], 0.123456789, &Default::default()); + + let output = format!("{}", tensor); + let expected = format!( + r#"Tensor {{ + data: +[[0.12345679]], + shape: [1, 1], + device: {:?}, + kind: "Float", + dtype: "f32", +}}"#, + tensor.device(), + ); + assert_eq!(output, expected); + + // CAN'T DO THIS BECAUSE OF GLOBAL STATE + // let print_options = PrintOptions { + // precision: Some(3), + // ..Default::default() + // }; + // set_print_options(print_options); + + let tensor = TestTensor::<2>::full([3, 2], 0.123456789, &Default::default()); + + // Set precision to 3 + let output = format!("{:.3}", tensor); + + let expected = format!( + r#"Tensor {{ + data: +[[0.123, 0.123], + [0.123, 0.123], + [0.123, 0.123]], + shape: [3, 2], + device: {:?}, + kind: "Float", + dtype: "f32", +}}"#, + tensor.device(), + ); + assert_eq!(output, expected); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/stats/eye.rs b/crates/burn-backend-tests/tests/tensor/float/stats/eye.rs new file mode 100644 index 0000000..13ac86f --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/stats/eye.rs @@ -0,0 +1,17 @@ +use super::*; + +#[test] +fn test_eye_float() { + let device = Default::default(); + let tensor = TestTensor::<2>::from([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]); + let rhs = TestTensor::<2>::eye(3, &device); + assert_eq!(tensor.to_data(), rhs.to_data()); +} + +#[test] +fn test_eye_int() { + let device = Default::default(); + let tensor = TestTensorInt::<2>::from([[1, 0, 0], [0, 1, 0], [0, 0, 1]]); + let rhs = TestTensorInt::<2>::eye(3, &device); + assert_eq!(tensor.to_data(), rhs.to_data()); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/stats/median.rs b/crates/burn-backend-tests/tests/tensor/float/stats/median.rs new file mode 100644 index 0000000..b8d21fe --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/stats/median.rs @@ -0,0 +1,92 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_median_even() { + let tensor = TestTensor::<2>::from_data( + [[0.5, 1.8, 0.2, -2.0], [3.0, -4.0, 5.0, 0.0]], + &Default::default(), + ); + + let median_actual_1 = tensor.clone().median(1); + let median_expected_1 = TensorData::from([[0.2], [0.0]]).convert::(); + median_actual_1 + .into_data() + .assert_eq(&median_expected_1, false); + + let median_actual_0 = tensor.median(0); + let median_expected_0 = TensorData::from([[0.5, -4.0, 0.2, -2.0]]).convert::(); + median_actual_0 + .into_data() + .assert_eq(&median_expected_0, false); +} + +#[test] +fn test_median_odd() { + let tensor = TestTensor::<2>::from_data( + [ + [0.5, 1.8, 0.2, -2.0, 1.0], + [3.0, -4.0, 5.0, 0.0, -1.0], + [5.0, -5.0, 1.0, 3.0, -2.0], + ], + &Default::default(), + ); + + let median_actual_1 = tensor.clone().median(1); + let median_expected_1 = TensorData::from([[0.5], [0.0], [1.0]]).convert::(); + median_actual_1 + .into_data() + .assert_eq(&median_expected_1, false); + + let median_actual_0 = tensor.median(0); + let median_expected_0 = TensorData::from([[3.0, -4.0, 1.0, 0.0, -1.0]]).convert::(); + median_actual_0 + .into_data() + .assert_eq(&median_expected_0, false); +} + +#[test] +fn test_median_with_indices() { + let device = Default::default(); + let tensor = TestTensor::<1>::from_data([3.0, 1.0, 2.0], &device); + // median = 2, original index = 2 + let (values, indices) = tensor.median_with_indices(0); + values + .into_data() + .assert_eq(&TensorData::from([2.0]), false); + indices + .into_data() + .assert_eq(&TensorData::from([2i64]), false); + + let tensor = TestTensor::<2>::from_data([[5.0, 1.0, 3.0], [2.0, 8.0, 4.0]], &device); + // Along dim 1: + // Row 0: median = 3, original index = 2 + // Row 1: median = 4, original index = 2 + let (values, indices) = tensor.median_with_indices(1); + values + .into_data() + .assert_eq(&TensorData::from([[3.0], [4.0]]), false); + indices + .into_data() + .assert_eq(&TensorData::from([[2i64], [2i64]]), false); +} + +#[test] +fn test_median_all_elements() { + let tensor = TestTensor::<2>::from_data( + [ + [0.5, 1.8, 0.2, -2.0, 1.0], + [3.0, -4.0, 5.0, 0.0, -1.0], + [5.0, -5.0, 1.0, 3.0, -2.0], + ], + &Default::default(), + ); + + // Sorted: [-5, -4, -2, -2, -1, 0, 0.2, 0.5, 1, 1, 1.8, 3, 3, 5, 5] + let dims = tensor.dims().len(); + let flattened_tensor: Tensor<1> = tensor.flatten(0, dims - 1); + let result = flattened_tensor.median(0); + result + .into_data() + .assert_eq(&TensorData::from([0.5]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/float/stats/mod.rs b/crates/burn-backend-tests/tests/tensor/float/stats/mod.rs new file mode 100644 index 0000000..dd5ee4c --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/stats/mod.rs @@ -0,0 +1,7 @@ +pub use super::*; // re-export test types + +mod cov; +mod display; +mod eye; +mod median; +mod var; diff --git a/crates/burn-backend-tests/tests/tensor/float/stats/var.rs b/crates/burn-backend-tests/tests/tensor/float/stats/var.rs new file mode 100644 index 0000000..a0a901a --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/float/stats/var.rs @@ -0,0 +1,69 @@ +use super::*; +use burn_tensor::TensorData; +use burn_tensor::Tolerance; + +#[test] +fn test_var() { + let tensor = TestTensor::<2>::from_data( + [[0.5, 1.8, 0.2, -2.0], [3.0, -4.0, 5.0, 0.0]], + &Default::default(), + ); + + let output = tensor.var(1); + let expected = TensorData::from([[2.4892], [15.3333]]).convert::(); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_var_mean() { + let tensor = TestTensor::<2>::from_data( + [[0.5, 1.8, 0.2, -2.0], [3.0, -4.0, 5.0, 0.0]], + &Default::default(), + ); + + let (var, mean) = tensor.var_mean(1); + + let var_expected = TensorData::from([[2.4892], [15.3333]]).convert::(); + let mean_expected = TensorData::from([[0.125], [1.]]).convert::(); + + var.into_data() + .assert_approx_eq::(&var_expected, Tolerance::default()); + mean.into_data() + .assert_approx_eq::(&mean_expected, Tolerance::default()); +} + +#[test] +fn test_var_bias() { + let tensor = TestTensor::<2>::from_data( + [[0.5, 1.8, 0.2, -2.0], [3.0, -4.0, 5.0, 0.0]], + &Default::default(), + ); + + let output = tensor.var_bias(1); + let expected = TensorData::from([[1.86688], [11.5]]).convert::(); + + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_var_mean_bias() { + let tensor = TestTensor::<2>::from_data( + [[0.5, 1.8, 0.2, -2.0], [3.0, -4.0, 5.0, 0.0]], + &Default::default(), + ); + + let (var, mean) = tensor.var_mean_bias(1); + + let var_expected = TensorData::from([[1.86688], [11.5]]).convert::(); + let mean_expected = TensorData::from([[0.125], [1.]]).convert::(); + + var.into_data() + .assert_approx_eq::(&var_expected, Tolerance::default()); + mean.into_data() + .assert_approx_eq::(&mean_expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/mod.rs b/crates/burn-backend-tests/tests/tensor/int/mod.rs new file mode 100644 index 0000000..a071d89 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/mod.rs @@ -0,0 +1,4 @@ +pub use super::*; // re-export test types + +mod ops; +mod primitive; diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/abs.rs b/crates/burn-backend-tests/tests/tensor/int/ops/abs.rs new file mode 100644 index 0000000..56b8138 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/abs.rs @@ -0,0 +1,24 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_abs_ops_int() { + let tensor = TestTensorInt::<2>::from([[0, -1, 2], [3, 4, -5]]); + + let output = tensor.abs(); + + output + .into_data() + .assert_eq(&TensorData::from([[0, 1, 2], [3, 4, 5]]), false); +} + +#[test] +fn should_support_abs_ops_int_signed_min() { + let tensor = TestTensorInt::<2>::from([[IntElem::MIN]]); + + let output = tensor.abs(); + + output + .into_data() + .assert_eq(&TensorData::from([[IntElem::MIN]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/add.rs b/crates/burn-backend-tests/tests/tensor/int/ops/add.rs new file mode 100644 index 0000000..da73ef5 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/add.rs @@ -0,0 +1,104 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_add_d2_int() { + let tensor_1 = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + let tensor_2 = TestTensorInt::from([[6, 7, 8], [9, 10, 11]]); + + let output = tensor_1 + tensor_2; + + output + .into_data() + .assert_eq(&TensorData::from([[6, 8, 10], [12, 14, 16]]), false); +} + +#[test] +fn test_add_broadcast_int() { + let tensor_1 = TestTensorInt::<2>::from([[0, 1, 2]]); + let tensor_2 = TestTensorInt::from([[3, 4, 5], [6, 7, 8]]); + + let output = tensor_1 + tensor_2; + + output + .into_data() + .assert_eq(&TensorData::from([[3, 5, 7], [6, 8, 10]]), false); +} + +#[test] +fn should_support_add_scalar_ops_int() { + let scalar = 2; + let tensor = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + + let output = tensor + scalar; + + output + .into_data() + .assert_eq(&TensorData::from([[2, 3, 4], [5, 6, 7]]), false); +} + +#[test] +fn scalar_add_not_contiguous() { + let tensor = TestTensorInt::<1>::arange(0..32, &Default::default()).float(); + let tensor = tensor.reshape([1, 4, 4, 2]).permute([0, 3, 1, 2]); + + let tensor = tensor.slice([0..1, 0..2, 0..4, 0..4]); + let before = tensor.clone(); + + let after = tensor.add_scalar(0.0); + + before + .into_data() + .assert_approx_eq::(&after.into_data(), Default::default()); +} + +#[test] +fn scalar_add_not_contiguous_int() { + let tensor = TestTensorInt::<1>::arange(0..32, &Default::default()); + let tensor = tensor.reshape([1, 4, 4, 2]).permute([0, 3, 1, 2]); + + let tensor = tensor.slice([0..1, 0..2, 0..4, 0..4]); + let before = tensor.clone(); + + let after = tensor.add_scalar(0); + + before.into_data().assert_eq(&after.into_data(), true); +} + +#[test] +fn test_int_add_transposed() { + // Non-contig on both sides via transpose. + let a = TestTensorInt::<2>::from([[1, 2], [3, 4]]).transpose(); + let b = TestTensorInt::<2>::from([[10, 20], [30, 40]]).transpose(); + + let output = a + b; + + output + .into_data() + .assert_eq(&TensorData::from([[11, 33], [22, 44]]), false); +} + +#[test] +fn test_int_add_flipped() { + // [1, 2, 3, 4] flipped + [10, 20, 30, 40] flipped = [44, 33, 22, 11] + let a = TestTensorInt::<1>::from([1, 2, 3, 4]).flip([0]); + let b = TestTensorInt::<1>::from([10, 20, 30, 40]).flip([0]); + + let output = a + b; + + output + .into_data() + .assert_eq(&TensorData::from([44, 33, 22, 11]), false); +} + +#[test] +fn test_int_add_scalar_flipped() { + // Scalar add on a flipped input. [1,2,3,4] flipped + 10 = [14,13,12,11] + let a = TestTensorInt::<1>::from([1, 2, 3, 4]).flip([0]); + + let output = a + 10; + + output + .into_data() + .assert_eq(&TensorData::from([14, 13, 12, 11]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/aggregation.rs b/crates/burn-backend-tests/tests/tensor/int/ops/aggregation.rs new file mode 100644 index 0000000..0d1d37e --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/aggregation.rs @@ -0,0 +1,81 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_should_mean_int() { + let tensor = TestTensorInt::<2>::from([[2, 2, 2], [3, 4, 5]]); + + let output = tensor.mean(); + + output.into_data().assert_eq(&TensorData::from([3]), false); +} + +#[test] +fn test_should_mean_last_dim_int() { + let tensor = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + + let output = tensor.mean_dim(1); + + output + .into_data() + .assert_eq(&TensorData::from([[1], [4]]), false); +} + +#[test] +fn test_should_sum_last_dim_int() { + let tensor = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + + let output = tensor.sum_dim(1); + + output + .into_data() + .assert_eq(&TensorData::from([[3], [12]]), false); +} + +#[test] +fn test_should_sum_int() { + let tensor = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + + let output = tensor.sum(); + + output.into_data().assert_eq(&TensorData::from([15]), false); +} + +#[test] +#[ignore = "Not implemented for all backends yet"] +fn test_prod_int() { + let tensor = TestTensorInt::<2>::from([[2, 1, 2], [3, 4, 5]]); + let output = tensor.prod(); + + output + .into_data() + .assert_eq(&TensorData::from([240]), false); + + let tensor_with_zero = TestTensorInt::<2>::from([[2, 0, 2], [3, 4, 5]]); + let output = tensor_with_zero.prod(); + + output.into_data().assert_eq(&TensorData::from([0]), false); +} + +#[test] +#[ignore = "Not implemented for all backends yet"] +fn test_prod_dim_int() { + let tensor = TestTensorInt::<2>::from([[2, 1, 2], [3, 4, 5]]); + let output = tensor.prod_dim(1); + output + .into_data() + .assert_eq(&TensorData::from([[4], [60]]), false); + + let tensor_with_zero = TestTensorInt::<2>::from([[2, 0, 2], [3, 4, 5]]); + let output = tensor_with_zero.prod_dim(1); + output + .into_data() + .assert_eq(&TensorData::from([[0], [60]]), false); + + // Negative Indexing. + let tensor_with_zero = TestTensorInt::<2>::from([[2, 0, 2], [3, 4, 5]]); + let output = tensor_with_zero.prod_dim(-1); + output + .into_data() + .assert_eq(&TensorData::from([[0], [60]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/all.rs b/crates/burn-backend-tests/tests/tensor/int/ops/all.rs new file mode 100644 index 0000000..490874c --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/all.rs @@ -0,0 +1,23 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_all() { + let tensor = TestTensorInt::<2>::from([[0, 1, 0], [1, -1, 1]]); + let data_actual = tensor.all().into_data(); + let data_expected = TensorData::from([false]); + data_expected.assert_eq(&data_actual, false); + + let tensor = TestTensorInt::<2>::from([[1, 1, 1], [1, 1, 1]]); + let data_actual = tensor.all().into_data(); + let data_expected = TensorData::from([true]); + data_expected.assert_eq(&data_actual, false); +} + +#[test] +fn test_all_dim() { + let tensor = TestTensorInt::<2>::from([[0, 1, 0], [1, -1, 1]]); + let data_actual = tensor.all_dim(1).into_data(); + let data_expected = TensorData::from([[false], [true]]); + data_expected.assert_eq(&data_actual, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/any.rs b/crates/burn-backend-tests/tests/tensor/int/ops/any.rs new file mode 100644 index 0000000..967129f --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/any.rs @@ -0,0 +1,23 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_any() { + let tensor = TestTensorInt::<2>::from([[0, 0, 0], [1, -1, 0]]); + let data_actual = tensor.any().into_data(); + let data_expected = TensorData::from([true]); + data_expected.assert_eq(&data_actual, false); + + let tensor = TestTensorInt::<2>::from([[0, 0, 0], [0, 0, 0]]); + let data_actual = tensor.any().into_data(); + let data_expected = TensorData::from([false]); + data_expected.assert_eq(&data_actual, false); +} + +#[test] +fn test_any_dim() { + let tensor = TestTensorInt::<2>::from([[0, 0, 0], [1, -1, 0]]); + let data_actual = tensor.any_dim(1).into_data(); + let data_expected = TensorData::from([[false], [true]]); + data_expected.assert_eq(&data_actual, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/arange.rs b/crates/burn-backend-tests/tests/tensor/int/ops/arange.rs new file mode 100644 index 0000000..3399fa1 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/arange.rs @@ -0,0 +1,31 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_arange() { + let device = Default::default(); + + let tensor = TestTensorInt::<1>::arange(2..5, &device); + tensor + .into_data() + .assert_eq(&TensorData::from([2, 3, 4]), false); + + // Test arange with negative numbers + let tensor = TestTensorInt::<1>::arange(-10..-5, &device); + tensor + .into_data() + .assert_eq(&TensorData::from([-10, -9, -8, -7, -6]), false); + + let tensor = TestTensorInt::<1>::arange(-3..0, &device); + tensor + .into_data() + .assert_eq(&TensorData::from([-3, -2, -1]), false); + + // Test arange with a mix of positive and negative numbers + let tensor = TestTensorInt::<1>::arange(-2..3, &device); + tensor + .clone() + .into_data() + .assert_eq(&TensorData::from([-2, -1, 0, 1, 2]), false); + assert_eq!(tensor.device(), device); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/arange_step.rs b/crates/burn-backend-tests/tests/tensor/int/ops/arange_step.rs new file mode 100644 index 0000000..7015c29 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/arange_step.rs @@ -0,0 +1,44 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_arange_step() { + let device = Default::default(); + + // Test correct sequence of numbers when the range is 0..9 and the step is 1 + let tensor = TestTensorInt::<1>::arange_step(0..9, 1, &device); + tensor + .into_data() + .assert_eq(&TensorData::from([0, 1, 2, 3, 4, 5, 6, 7, 8]), false); + + // Test correct sequence of numbers when the range is 0..3 and the step is 2 + let tensor = TestTensorInt::<1>::arange_step(0..3, 2, &device); + tensor + .into_data() + .assert_eq(&TensorData::from([0, 2]), false); + + // Test correct sequence of numbers when the range is 0..2 and the step is 5 + let tensor = TestTensorInt::<1>::arange_step(0..2, 5, &device); + tensor.into_data().assert_eq(&TensorData::from([0]), false); + + // Test correct sequence of numbers when the range includes negative numbers + let tensor = TestTensorInt::<1>::arange_step(-3..3, 2, &device); + tensor + .into_data() + .assert_eq(&TensorData::from([-3, -1, 1]), false); + + let tensor = TestTensorInt::<1>::arange_step(-5..1, 5, &device); + tensor + .clone() + .into_data() + .assert_eq(&TensorData::from([-5, 0]), false); + assert_eq!(tensor.device(), device); +} + +#[test] +#[should_panic] +fn should_panic_when_step_is_zero() { + let device = Default::default(); + // Test that arange_step panics when the step is 0 + let _tensor = TestTensorInt::<1>::arange_step(0..3, 0, &device); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/arg.rs b/crates/burn-backend-tests/tests/tensor/int/ops/arg.rs new file mode 100644 index 0000000..9ba36af --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/arg.rs @@ -0,0 +1,24 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_argmax_2d_dim0_int() { + let tensor = TestTensorInt::<2>::from([[10, 11, 2], [3, 4, 5]]); + + let output = tensor.argmax(0); + + output + .into_data() + .assert_eq(&TensorData::from([[0, 0, 1]]), false); +} + +#[test] +fn test_argmin_2d_dim0_int() { + let tensor = TestTensorInt::<2>::from([[10, 11, 2], [30, 4, 5]]); + + let output = tensor.argmin(0); + + output + .into_data() + .assert_eq(&TensorData::from([[0, 1, 0]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/bitwise.rs b/crates/burn-backend-tests/tests/tensor/int/ops/bitwise.rs new file mode 100644 index 0000000..cfdfc15 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/bitwise.rs @@ -0,0 +1,173 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_apply_bitwise_and_2d() { + let tensor_1 = TestTensorInt::<2>::from([[3, 4, 5], [9, 3, 8]]); + let tensor_2 = TestTensorInt::from([[6, 7, 8], [9, 10, 15]]); + + let output = tensor_1.bitwise_and(tensor_2); + + output + .into_data() + .assert_eq(&TensorData::from([[2, 4, 0], [9, 2, 8]]), false); +} + +#[test] +fn should_apply_bitwise_and_1d() { + let tensor_1 = TestTensorInt::<1>::from([13, 7]); + let tensor_2 = TestTensorInt::from([11, 3]); + + let output = tensor_1.bitwise_and(tensor_2); + + output + .into_data() + .assert_eq(&TensorData::from([9, 3]), false); +} + +#[test] +fn should_apply_bitwise_and_scalar_2d() { + let tensor_1 = TestTensorInt::<2>::from([[3, 4, 5], [9, 3, 8]]); + let scalar = 5; + + let output = tensor_1.bitwise_and_scalar(scalar); + + output + .into_data() + .assert_eq(&TensorData::from([[1, 4, 5], [1, 1, 0]]), false); +} + +#[test] +fn should_apply_bitwise_not_2d() { + let tensor_1 = TestTensorInt::<2>::from([[3, 4, 5], [9, 3, 8]]); + + let output = tensor_1.bitwise_not(); + + output + .into_data() + .assert_eq(&TensorData::from([[-4, -5, -6], [-10, -4, -9]]), false); +} + +#[test] +fn should_apply_bitwise_or_scalar_2d() { + let tensor_1 = TestTensorInt::<2>::from([[3, 4, 5], [9, 3, 8]]); + let scalar = 5; + + let output = tensor_1.bitwise_or_scalar(scalar); + + output + .into_data() + .assert_eq(&TensorData::from([[7, 5, 5], [13, 7, 13]]), false); +} + +#[test] +fn should_apply_bitwise_or_2d() { + let tensor_1 = TestTensorInt::<2>::from([[3, 4, 5], [9, 3, 8]]); + let tensor_2 = TestTensorInt::from([[6, 7, 8], [9, 10, 15]]); + + let output = tensor_1.bitwise_or(tensor_2); + + output + .into_data() + .assert_eq(&TensorData::from([[7, 7, 13], [9, 11, 15]]), false); +} + +#[test] +fn should_apply_bitwise_or_1d() { + let tensor_1 = TestTensorInt::<1>::from([13, 7]); + let tensor_2 = TestTensorInt::from([11, 3]); + + let output = tensor_1.bitwise_or(tensor_2); + + output + .into_data() + .assert_eq(&TensorData::from([15, 7]), false); +} + +#[test] +fn should_apply_bitwise_xor_scalar_2d() { + let tensor_1 = TestTensorInt::<2>::from([[3, 4, 5], [9, 3, 8]]); + let scalar = 5; + + let output = tensor_1.bitwise_xor_scalar(scalar); + + output + .into_data() + .assert_eq(&TensorData::from([[6, 1, 0], [12, 6, 13]]), false); +} + +#[test] +fn should_apply_bitwise_xor_2d() { + let tensor_1 = TestTensorInt::<2>::from([[3, 4, 5], [9, 3, 8]]); + let tensor_2 = TestTensorInt::from([[6, 7, 8], [9, 10, 15]]); + + let output = tensor_1.bitwise_xor(tensor_2); + + output + .into_data() + .assert_eq(&TensorData::from([[5, 3, 13], [0, 9, 7]]), false); +} + +#[test] +fn should_apply_bitwise_xor_1d() { + let tensor_1 = TestTensorInt::<1>::from([13, 7]); + let tensor_2 = TestTensorInt::from([11, 3]); + + let output = tensor_1.bitwise_xor(tensor_2); + + output + .into_data() + .assert_eq(&TensorData::from([6, 4]), false); +} + +#[test] +fn should_apply_bitwise_left_shift_2d() { + if (IntElem::MAX as u32) < 512 { + return; + } + + let tensor_1 = TestTensorInt::<2>::from([[3, 4, 5], [9, 3, 8]]); + let tensor_2 = TestTensorInt::from([[1, 2, 3], [4, 5, 6]]); + + let output = tensor_1.bitwise_left_shift(tensor_2); + + output + .into_data() + .assert_eq(&TensorData::from([[6, 16, 40], [144, 96, 512]]), false); +} + +#[test] +fn should_apply_bitwise_left_shift_scalar_2d() { + let tensor_1 = TestTensorInt::<2>::from([[3, 4, 5], [9, 3, 8]]); + let scalar = 2; + + let output = tensor_1.bitwise_left_shift_scalar(scalar); + + output + .into_data() + .assert_eq(&TensorData::from([[12, 16, 20], [36, 12, 32]]), false); +} + +#[test] +fn should_apply_bitwise_right_shift_2d() { + let tensor_1 = TestTensorInt::<2>::from([[3, 4, 5], [9, 3, 8]]); + let tensor_2 = TestTensorInt::from([[1, 2, 3], [4, 5, 6]]); + + let output = tensor_1.bitwise_right_shift(tensor_2); + + output + .into_data() + .assert_eq(&TensorData::from([[1, 1, 0], [0, 0, 0]]), false); +} + +#[test] +fn should_apply_bitwise_right_shift_scalar_2d() { + let tensor_1 = TestTensorInt::<2>::from([[3, 4, 5], [9, 3, 8]]); + let scalar = 2; + + let output = tensor_1.bitwise_right_shift_scalar(scalar); + + output + .into_data() + .assert_eq(&TensorData::from([[0, 1, 1], [2, 0, 2]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/cartesian_grid.rs b/crates/burn-backend-tests/tests/tensor/int/ops/cartesian_grid.rs new file mode 100644 index 0000000..87db450 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/cartesian_grid.rs @@ -0,0 +1,20 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_cartesian_grid() { + let device = Default::default(); + + // Test a single element tensor + let tensor: TestTensorInt<2> = TestTensorInt::<1>::cartesian_grid([1], &device); + tensor + .into_data() + .assert_eq(&TensorData::from([[0]]), false); + + // Test for a 2x2 tensor + let tensor: TestTensorInt<3> = TestTensorInt::<2>::cartesian_grid([2, 2], &device); + tensor.into_data().assert_eq( + &TensorData::from([[[0, 0], [0, 1]], [[1, 0], [1, 1]]]), + false, + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/cast.rs b/crates/burn-backend-tests/tests/tensor/int/ops/cast.rs new file mode 100644 index 0000000..f366d70 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/cast.rs @@ -0,0 +1,84 @@ +use super::*; +use burn_tensor::{DType, FloatDType, IntDType, TensorData}; + +#[test] +fn cast_int_to_bool() { + let tensor1 = TestTensorInt::<2>::from([[0, 43, 0], [2, -4, 31]]); + let data_actual = tensor1.bool().into_data(); + let data_expected = TensorData::from([[false, true, false], [true, true, true]]); + data_actual.assert_eq(&data_expected, false); +} + +#[test] +fn cast_bool_to_int_tensor() { + let tensor = TestTensorBool::<2>::from([[true, false, true], [false, false, true]]).int(); + + let expected = TensorData::from([[1, 0, 1], [0, 0, 1]]); + + tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn cast_int_precision() { + let data = TensorData::from([[1, 2, 3], [4, 5, 6]]); + let tensor = TestTensorInt::<2>::from(data.clone()); + + let output = tensor.cast(DType::I32); + + assert_eq!(output.dtype(), DType::I32); + output.into_data().assert_eq(&data, false); +} + +#[test] +fn cast_int_to_float_with_dtype() { + let tensor = TestTensorInt::<2>::from([[1, 2, 3], [4, 5, 6]]); + + let output = tensor.cast(FloatDType::F32); + + assert_eq!(output.dtype(), DType::F32); + let expected = TensorData::from([[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn cast_int_within_kind() { + let tensor = TestTensorInt::<1>::from([1, 2, 3]); + + let output = tensor.cast(IntDType::I32); + + assert_eq!(output.dtype(), DType::I32); + let expected = TensorData::from([1i32, 2, 3]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn cast_int_same_dtype_is_noop() { + let data = TensorData::from([1, 2, 3]); + let tensor = TestTensorInt::<1>::from(data.clone()); + let original_dtype = tensor.dtype(); + let int_dtype: IntDType = original_dtype.into(); + + let output = tensor.cast(int_dtype); + + assert_eq!(output.dtype(), original_dtype); + output.into_data().assert_eq(&data, false); +} + +#[test] +#[should_panic] +fn cast_int_with_float_dtype_panics() { + let tensor = TestTensorInt::<1>::from([1, 2]); + let _ = tensor.cast(DType::F32); +} + +#[test] +fn test_int_into_float_flipped() { + // [1, 2, 3, 4] flipped -> [4, 3, 2, 1] -> float [4.0, 3.0, 2.0, 1.0] + let t = TestTensorInt::<1>::from([1, 2, 3, 4]).flip([0]); + + let output = t.float(); + + output + .into_data() + .assert_eq(&TensorData::from([4.0f32, 3.0, 2.0, 1.0]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/cat.rs b/crates/burn-backend-tests/tests/tensor/int/ops/cat.rs new file mode 100644 index 0000000..640ed06 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/cat.rs @@ -0,0 +1,28 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_cat_ops_int() { + let device = Default::default(); + let tensor_1 = TestTensorInt::<2>::from_data([[1, 2, 3]], &device); + let tensor_2 = TestTensorInt::<2>::from_data([[4, 5, 6]], &device); + + let output = Tensor::cat(vec![tensor_1, tensor_2], 0); + + output + .into_data() + .assert_eq(&TensorData::from([[1, 2, 3], [4, 5, 6]]), false); +} + +#[test] +fn should_support_cat_with_empty_tensor_int() { + let device = Default::default(); + let tensor_1 = TestTensorInt::<2>::from_data([[1, 2, 3]], &device); + let tensor_2: TestTensorInt<2> = TestTensorInt::empty([1, 0], &device); + + let output = Tensor::cat(vec![tensor_1, tensor_2], 1); + + output + .into_data() + .assert_eq(&TensorData::from([[1, 2, 3]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/chunk.rs b/crates/burn-backend-tests/tests/tensor/int/ops/chunk.rs new file mode 100644 index 0000000..e040fad --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/chunk.rs @@ -0,0 +1,16 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_chunk_multi_dimension() { + let tensors = + TestTensorInt::<2>::from_data(TensorData::from([[0, 1, 2, 3]]), &Default::default()) + .chunk(2, 1); + assert_eq!(tensors.len(), 2); + + let expected = [TensorData::from([[0, 1]]), TensorData::from([[2, 3]])]; + + for (index, tensor) in tensors.iter().enumerate() { + tensor.to_data().assert_eq(&expected[index], false); + } +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/comparison.rs b/crates/burn-backend-tests/tests/tensor/int/ops/comparison.rs new file mode 100644 index 0000000..522ad44 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/comparison.rs @@ -0,0 +1,280 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_equal() { + let tensor_1 = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + let tensor_2 = TestTensorInt::<2>::from([[1, 1, 1], [4, 3, 5]]); + + let data_actual_cloned = tensor_1.clone().equal(tensor_2.clone()); + let data_actual_inplace = tensor_1.equal(tensor_2); + + let data_expected = TensorData::from([[false, true, false], [false, false, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_not_equal() { + let tensor_1 = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + let tensor_2 = TestTensorInt::<2>::from([[1, 1, 1], [4, 3, 5]]); + + let data_actual_cloned = tensor_1.clone().not_equal(tensor_2.clone()); + let data_actual_inplace = tensor_1.not_equal(tensor_2); + + let data_expected = TensorData::from([[true, false, true], [true, true, false]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_equal_elem() { + let tensor_1 = TestTensorInt::<2>::from([[0, 1, 2], [3, 2, 5]]); + + let data_actual_cloned = tensor_1.clone().equal_elem(2); + let data_actual_inplace = tensor_1.equal_elem(2); + + let data_expected = TensorData::from([[false, false, true], [false, true, false]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_not_equal_elem() { + let tensor_1 = TestTensorInt::<2>::from([[0, 1, 2], [3, 2, 5]]); + + let data_actual_cloned = tensor_1.clone().not_equal_elem(2); + let data_actual_inplace = tensor_1.not_equal_elem(2); + + let data_expected = TensorData::from([[true, true, false], [true, false, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn greater_elem() { + let tensor_1 = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + + let data_actual_cloned = tensor_1.clone().greater_elem(4); + let data_actual_inplace = tensor_1.greater_elem(4); + + let data_expected = TensorData::from([[false, false, false], [false, false, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_greater_equal_elem() { + let tensor_1 = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + + let data_actual_cloned = tensor_1.clone().greater_equal_elem(4); + let data_actual_inplace = tensor_1.greater_equal_elem(4); + + let data_expected = TensorData::from([[false, false, false], [false, true, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_greater() { + let tensor_1 = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + let tensor_2 = TestTensorInt::<2>::from([[1, 1, 1], [4, 3, 50]]); + + let data_actual_cloned = tensor_1.clone().greater(tensor_2.clone()); + let data_actual_inplace = tensor_1.greater(tensor_2); + + let data_expected = TensorData::from([[false, false, true], [false, true, false]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_greater_equal() { + let tensor_1 = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + let tensor_2 = TestTensorInt::<2>::from([[1, 1, 1], [4, 3, 50]]); + + let data_actual_cloned = tensor_1.clone().greater_equal(tensor_2.clone()); + let data_actual_inplace = tensor_1.greater_equal(tensor_2); + + let data_expected = TensorData::from([[false, true, true], [false, true, false]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_lower_elem() { + let tensor_1 = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + + let data_actual_cloned = tensor_1.clone().lower_elem(4); + let data_actual_inplace = tensor_1.lower_elem(4); + + let data_expected = TensorData::from([[true, true, true], [true, false, false]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_lower_equal_elem() { + let tensor_1 = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + + let data_actual_cloned = tensor_1.clone().lower_equal_elem(4); + let data_actual_inplace = tensor_1.lower_equal_elem(4); + + let data_expected = TensorData::from([[true, true, true], [true, true, false]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_lower() { + let tensor_1 = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + let tensor_2 = TestTensorInt::<2>::from([[1, 1, 1], [4, 3, 50]]); + + let data_actual_cloned = tensor_1.clone().lower(tensor_2.clone()); + let data_actual_inplace = tensor_1.lower(tensor_2); + + let data_expected = TensorData::from([[true, false, false], [true, false, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_lower_equal() { + let tensor_1 = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + let tensor_2 = TestTensorInt::<2>::from([[1, 1, 1], [4, 3, 50]]); + + let data_actual_cloned = tensor_1.clone().lower_equal(tensor_2.clone()); + let data_actual_inplace = tensor_1.lower_equal(tensor_2); + + let data_expected = TensorData::from([[true, true, false], [true, false, true]]); + data_expected.assert_eq(&data_actual_cloned.into_data(), false); + data_expected.assert_eq(&data_actual_inplace.into_data(), false); +} + +#[test] +fn test_greater_broadcast() { + // Test broadcasting with shape [1, 4] vs [4, 4] + let device = Default::default(); + let data_1 = TensorData::from([[1, 2, 3, 4]]); + let data_2 = TensorData::from([ + [0.5, 1.5, 2.5, 3.5], + [1.5, 2.5, 3.5, 4.5], + [2.5, 3.5, 4.5, 5.5], + [3.5, 4.5, 5.5, 6.5], + ]); + let tensor_1 = TestTensorInt::<2>::from_data(data_1, &device); + let tensor_2 = TestTensorInt::<2>::from_data(data_2, &device); + + let result = tensor_1.greater(tensor_2); + + let expected = TensorData::from([ + [true, true, true, true], + [false, false, false, false], + [false, false, false, false], + [false, false, false, false], + ]); + expected.assert_eq(&result.into_data(), false); +} + +#[test] +fn test_greater_equal_broadcast() { + // Test broadcasting with shape [4, 1] vs [1, 4] + let device = Default::default(); + let data_1 = TensorData::from([[1], [2], [3], [4]]); + let data_2 = TensorData::from([[1, 2, 3, 4]]); + let tensor_1 = TestTensorInt::<2>::from_data(data_1, &device); + let tensor_2 = TestTensorInt::<2>::from_data(data_2, &device); + + let result = tensor_1.greater_equal(tensor_2); + + let expected = TensorData::from([ + [true, false, false, false], + [true, true, false, false], + [true, true, true, false], + [true, true, true, true], + ]); + expected.assert_eq(&result.into_data(), false); +} + +#[test] +fn test_equal_broadcast() { + // Test broadcasting with different ranks + let device = Default::default(); + let data_1 = TensorData::from([[2], [3], [4]]); + let data_2 = TensorData::from([[2, 3, 4, 2]]); + let tensor_1 = TestTensorInt::<2>::from_data(data_1, &device); + let tensor_2 = TestTensorInt::<2>::from_data(data_2, &device); + + let result = tensor_1.equal(tensor_2); + + let expected = TensorData::from([ + [true, false, false, true], + [false, true, false, false], + [false, false, true, false], + ]); + expected.assert_eq(&result.into_data(), false); +} + +#[test] +fn test_not_equal_broadcast() { + // Test broadcasting with shape [3, 1] vs [1, 3] + let device = Default::default(); + let data_1 = TensorData::from([[1], [2], [3]]); + let data_2 = TensorData::from([[1, 2, 3]]); + let tensor_1 = TestTensorInt::<2>::from_data(data_1, &device); + let tensor_2 = TestTensorInt::<2>::from_data(data_2, &device); + + let result = tensor_1.not_equal(tensor_2); + + let expected = TensorData::from([ + [false, true, true], + [true, false, true], + [true, true, false], + ]); + expected.assert_eq(&result.into_data(), false); +} + +#[test] +fn test_int_greater_broadcast() { + let device = Default::default(); + let data_1 = TensorData::from([[1i32, 2, 3]]); + let data_2 = TensorData::from([[0i32], [2], [4]]); + let tensor_1 = TestTensorInt::<2>::from_data(data_1, &device); + let tensor_2 = TestTensorInt::<2>::from_data(data_2, &device); + + let result = tensor_1.greater(tensor_2); + + let expected = TensorData::from([ + [true, true, true], + [false, false, true], + [false, false, false], + ]); + expected.assert_eq(&result.into_data(), false); +} + +#[test] +fn test_int_lower_equal_broadcast() { + let device = Default::default(); + let data_1 = TensorData::from([[2i32], [4]]); + let data_2 = TensorData::from([[1i32, 2, 3]]); + let tensor_1 = TestTensorInt::<2>::from_data(data_1, &device); + let tensor_2 = TestTensorInt::<2>::from_data(data_2, &device); + + let result = tensor_1.lower_equal(tensor_2); + + let expected = TensorData::from([[false, true, true], [false, false, false]]); + expected.assert_eq(&result.into_data(), false); +} + +#[test] +fn test_int_greater_flipped() { + // [1,2,3,4] flipped -> [4,3,2,1]; > [3,3,3,3] = [T,F,F,F] + let lhs = TestTensorInt::<1>::from([1, 2, 3, 4]).flip([0]); + let rhs = TestTensorInt::<1>::from([3, 3, 3, 3]); + + let result = lhs.greater(rhs); + + result + .into_data() + .assert_eq(&TensorData::from([true, false, false, false]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/create_like.rs b/crates/burn-backend-tests/tests/tensor/int/ops/create_like.rs new file mode 100644 index 0000000..c674971 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/create_like.rs @@ -0,0 +1,22 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_zeros_like() { + let tensor = TestTensorInt::<3>::from([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]]); + + let tensor = tensor.zeros_like(); + let expected = TensorData::from([[[0, 0, 0], [0, 0, 0]], [[0, 0, 0], [0, 0, 0]]]); + + tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_ones_like() { + let tensor = TestTensorInt::<3>::from([[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]]); + + let tensor = tensor.ones_like(); + let expected = TensorData::from([[[1, 1, 1], [1, 1, 1]], [[1, 1, 1], [1, 1, 1]]]); + + tensor.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/cumulative.rs b/crates/burn-backend-tests/tests/tensor/int/ops/cumulative.rs new file mode 100644 index 0000000..6af39a9 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/cumulative.rs @@ -0,0 +1,90 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_cumsum_int_dim_0() { + let tensor = TestTensorInt::<2>::from([[1, 2, 3], [4, 5, 6]]); + + let output = tensor.cumsum(0); + + output + .into_data() + .assert_eq(&TensorData::from([[1, 2, 3], [5, 7, 9]]), false); +} + +#[test] +fn test_cumsum_int_dim_1() { + let tensor = TestTensorInt::<2>::from([[1, 2, 3], [4, 5, 6]]); + + let output = tensor.cumsum(1); + + output + .into_data() + .assert_eq(&TensorData::from([[1, 3, 6], [4, 9, 15]]), false); +} + +#[test] +fn test_cumprod_int_dim_0() { + let tensor = TestTensorInt::<2>::from([[1, 2, 3], [4, 5, 6]]); + + let output = tensor.cumprod(0); + + output + .into_data() + .assert_eq(&TensorData::from([[1, 2, 3], [4, 10, 18]]), false); +} + +#[test] +fn test_cumprod_int_dim_1() { + let tensor = TestTensorInt::<2>::from([[1, 2, 3], [4, 5, 6]]); + + let output = tensor.cumprod(1); + + output + .into_data() + .assert_eq(&TensorData::from([[1, 2, 6], [4, 20, 120]]), false); +} + +#[test] +fn test_cummin_int_dim_0() { + let tensor = TestTensorInt::<2>::from([[3, 1, 4], [2, 5, 1]]); + + let output = tensor.cummin(0); + + output + .into_data() + .assert_eq(&TensorData::from([[3, 1, 4], [2, 1, 1]]), false); +} + +#[test] +fn test_cummin_int_dim_1() { + let tensor = TestTensorInt::<2>::from([[3, 1, 4], [2, 5, 1]]); + + let output = tensor.cummin(1); + + output + .into_data() + .assert_eq(&TensorData::from([[3, 1, 1], [2, 2, 1]]), false); +} + +#[test] +fn test_cummax_int_dim_0() { + let tensor = TestTensorInt::<2>::from([[3, 1, 4], [1, 5, 2]]); + + let output = tensor.cummax(0); + + output + .into_data() + .assert_eq(&TensorData::from([[3, 1, 4], [3, 5, 4]]), false); +} + +#[test] +fn test_cummax_int_dim_1() { + let tensor = TestTensorInt::<2>::from([[3, 1, 4], [1, 5, 2]]); + + let output = tensor.cummax(1); + + output + .into_data() + .assert_eq(&TensorData::from([[3, 3, 4], [1, 5, 5]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/div.rs b/crates/burn-backend-tests/tests/tensor/int/ops/div.rs new file mode 100644 index 0000000..167b3a3 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/div.rs @@ -0,0 +1,45 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_div_ops_int() { + let data_1 = TensorData::from([[0, 1, 2], [3, 4, 5]]); + let data_2 = TensorData::from([[1, 1, 2], [1, 1, 2]]); + let device = Default::default(); + let tensor_1 = TestTensorInt::<2>::from_data(data_1, &device); + let tensor_2 = TestTensorInt::<2>::from_data(data_2, &device); + + let output = tensor_1 / tensor_2; + + output + .into_data() + .assert_eq(&TensorData::from([[0, 1, 1], [3, 4, 2]]), false); +} + +#[test] +fn test_div_broadcast_int() { + let data_1 = TensorData::from([[0, 1, 2]]); + let data_2 = TensorData::from([[1, 1, 2], [3, 4, 5]]); + let device = Default::default(); + let tensor_1 = TestTensorInt::<2>::from_data(data_1, &device); + let tensor_2 = TestTensorInt::<2>::from_data(data_2, &device); + + let output = tensor_1 / tensor_2; + + output + .into_data() + .assert_eq(&TensorData::from([[0, 1, 1], [0, 0, 0]]), false); +} + +#[test] +fn should_support_div_scalar_ops_int() { + let data = TensorData::from([[0, 1, 2], [3, 4, 5]]); + let scalar = 2; + let tensor = TestTensorInt::<2>::from_data(data, &Default::default()); + + let output = tensor / scalar; + + output + .into_data() + .assert_eq(&TensorData::from([[0, 0, 1], [1, 2, 2]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/expand.rs b/crates/burn-backend-tests/tests/tensor/int/ops/expand.rs new file mode 100644 index 0000000..c34ee7b --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/expand.rs @@ -0,0 +1,75 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn expand_2d_int() { + let tensor = TestTensorInt::<1>::from([1, 2, 3]); + let output = tensor.expand([3, 3]); + + output + .into_data() + .assert_eq(&TensorData::from([[1, 2, 3], [1, 2, 3], [1, 2, 3]]), false); +} + +#[test] +fn should_all_negative_one() { + let tensor = TestTensorInt::<1>::from([1, 2, 3]); + let output = tensor.expand([2, -1]); + + output + .into_data() + .assert_eq(&TensorData::from([[1, 2, 3], [1, 2, 3]]), false); +} + +#[test] +#[should_panic] +fn should_panic_negative_one_on_non_existing_dim() { + let tensor = TestTensorInt::<1>::from([1, 2, 3]); + let _expanded_tensor = tensor.expand([-1, 3]); +} + +/// Regression test for https://github.com/tracel-ai/burn/issues/2091 +#[test] +fn inplace_op_after_expand() { + let tensor = TestTensorInt::<1>::from([1, 2, 3]); + let mut output = tensor.expand([2, 3]); + output = output + 1; + + output + .into_data() + .assert_eq(&TensorData::from([[2, 3, 4], [2, 3, 4]]), false); +} + +#[test] +fn expand_int_after_transpose() { + let tensor = TestTensorInt::<2>::from([[1, 2], [3, 4]]).transpose(); + + let output = tensor.expand([3, 2, 2]); + + output.into_data().assert_eq( + &TensorData::from([[[1, 3], [2, 4]], [[1, 3], [2, 4]], [[1, 3], [2, 4]]]), + false, + ); +} + +#[test] +fn expand_int_after_flip() { + let tensor = TestTensorInt::<1>::from([1, 2, 3]).flip([0]); + + let output = tensor.expand([2, 3]); + + output + .into_data() + .assert_eq(&TensorData::from([[3, 2, 1], [3, 2, 1]]), false); +} + +#[test] +fn expand_int_after_narrow() { + let tensor = TestTensorInt::<1>::from([0, 1, 2, 3, 4]).narrow(0, 1, 3); + + let output = tensor.expand([2, 3]); + + output + .into_data() + .assert_eq(&TensorData::from([[1, 2, 3], [1, 2, 3]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/flip.rs b/crates/burn-backend-tests/tests/tensor/int/ops/flip.rs new file mode 100644 index 0000000..c302322 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/flip.rs @@ -0,0 +1,22 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn flip_int() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device).reshape([2, 3, 4]); + + let flipped = tensor.clone().flip([0, 2]); + // from pytorch: + // import torch; torch.arange(0, 24).reshape(2, 3, 4).flip((0, 2)) + let expected = TensorData::from([ + [[15, 14, 13, 12], [19, 18, 17, 16], [23, 22, 21, 20]], + [[3, 2, 1, 0], [7, 6, 5, 4], [11, 10, 9, 8]], + ]); + + flipped.into_data().assert_eq(&expected, false); + + // Test with no flip + let flipped = tensor.clone().flip([]); + assert_eq!(tensor.into_data(), flipped.into_data()); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/full.rs b/crates/burn-backend-tests/tests/tensor/int/ops/full.rs new file mode 100644 index 0000000..4541795 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/full.rs @@ -0,0 +1,11 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_tensor_full() { + let device = Default::default(); + let int_tensor = TestTensorInt::<2>::full([2, 2], 2, &device); + int_tensor + .into_data() + .assert_eq(&TensorData::from([[2, 2], [2, 2]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/gather_scatter.rs b/crates/burn-backend-tests/tests/tensor/int/ops/gather_scatter.rs new file mode 100644 index 0000000..ca8d281 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/gather_scatter.rs @@ -0,0 +1,69 @@ +use super::*; +use burn_tensor::{IndexingUpdateOp, TensorData}; + +#[test] +fn should_gather_1d_dim0_int() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::from_ints([5, 6, 7], &device); + let indices = TestTensorInt::from_ints([1, 1, 0, 1, 2], &device); + + let output = tensor.gather(0, indices); + + output + .into_data() + .assert_eq(&TensorData::from([6, 6, 5, 6, 7]), false); +} + +#[test] +fn should_gather_indices_broadcasted() { + let device = Default::default(); + + let batch_size = 3; + let fft_size = 4; + let shape = [batch_size, fft_size, 2]; + let x = TestTensorInt::arange( + 0..shape.iter().product::() as i64, + &Default::default(), + ) + .reshape(shape); + let idx = TestTensorInt::<1>::from_ints([0, 2, 1, 3], &device); + + let expected = TestTensorInt::<3>::from([ + [[0, 1], [4, 5], [2, 3], [6, 7]], + [[8, 9], [12, 13], [10, 11], [14, 15]], + [[16, 17], [20, 21], [18, 19], [22, 23]], + ]) + .into_data(); + + // Case 1: gather dim 2 + let perm = idx + .clone() + .reshape([1, 1, fft_size]) + .repeat_dim(0, batch_size) + .repeat_dim(1, 2); + + let input = x.clone().permute([0, 2, 1]); + let out = input.gather(2, perm).permute([0, 2, 1]); + + out.into_data().assert_eq(&expected, true); + + // Case 2: gather directly on dim 1 + let perm = idx.reshape([1, fft_size, 1]).repeat_dim(0, batch_size); + let out2 = x.gather(1, perm.repeat_dim(2, 2)); + + out2.into_data().assert_eq(&expected, true); +} + +#[test] +fn should_scatter_add_1d_int() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::from_ints([0, 0, 0], &device); + let values = TestTensorInt::from_ints([5, 4, 3], &device); + let indices = TestTensorInt::from_ints([1, 0, 2], &device); + + let output = tensor.scatter(0, indices, values, IndexingUpdateOp::Add); + + output + .into_data() + .assert_eq(&TensorData::from([4, 5, 3]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/gather_scatter_nd.rs b/crates/burn-backend-tests/tests/tensor/int/ops/gather_scatter_nd.rs new file mode 100644 index 0000000..6f3dfe7 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/gather_scatter_nd.rs @@ -0,0 +1,131 @@ +use super::*; +use burn_tensor::{IndexingUpdateOp, TensorData}; + +#[test] +fn test_int_gather_nd_2d() { + let device = Default::default(); + let data = TestTensorInt::<2>::from_ints([[10, 20, 30], [40, 50, 60]], &device); + // indices shape: [2, 2] => K=2, M=2, DV=1 + let indices = TestTensorInt::<2>::from_ints([[0, 1], [1, 2]], &device); + + let output: TestTensorInt<1> = data.gather_nd(indices); + + output + .into_data() + .assert_eq(&TensorData::from([20, 60]), false); +} + +#[test] +fn test_int_scatter_nd_add() { + let device = Default::default(); + let data = TestTensorInt::<2>::from_ints([[1, 2, 3], [4, 5, 6]], &device); + let indices = TestTensorInt::<2>::from_ints([[0, 0], [1, 2]], &device); + let values = TestTensorInt::<1>::from_ints([10, 20], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Add); + + output + .into_data() + .assert_eq(&TensorData::from([[11, 2, 3], [4, 5, 26]]), false); +} + +#[test] +fn test_int_scatter_nd_mul() { + let device = Default::default(); + let data = TestTensorInt::<2>::from_ints([[2, 3], [4, 5]], &device); + let indices = TestTensorInt::<2>::from_ints([[0, 1], [1, 0]], &device); + let values = TestTensorInt::<1>::from_ints([10, 3], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Mul); + + output + .into_data() + .assert_eq(&TensorData::from([[2, 30], [12, 5]]), false); +} + +#[test] +fn test_int_scatter_nd_assign() { + let device = Default::default(); + let data = TestTensorInt::<2>::from_ints([[1, 2, 3], [4, 5, 6]], &device); + let indices = TestTensorInt::<2>::from_ints([[0, 1], [1, 2]], &device); + let values = TestTensorInt::<1>::from_ints([99, 88], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Assign); + + output + .into_data() + .assert_eq(&TensorData::from([[1, 99, 3], [4, 5, 88]]), false); +} + +#[test] +fn test_int_scatter_nd_min() { + let device = Default::default(); + let data = TestTensorInt::<2>::from_ints([[5, 10], [15, 20]], &device); + let indices = TestTensorInt::<2>::from_ints([[0, 0], [1, 1]], &device); + let values = TestTensorInt::<1>::from_ints([3, 25], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Min); + + // min(5, 3) = 3; min(20, 25) = 20 + output + .into_data() + .assert_eq(&TensorData::from([[3, 10], [15, 20]]), false); +} + +#[test] +fn test_int_scatter_nd_max() { + let device = Default::default(); + let data = TestTensorInt::<2>::from_ints([[5, 10], [15, 20]], &device); + let indices = TestTensorInt::<2>::from_ints([[0, 0], [1, 1]], &device); + let values = TestTensorInt::<1>::from_ints([3, 25], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Max); + + // max(5, 3) = 5; max(20, 25) = 25 + output + .into_data() + .assert_eq(&TensorData::from([[5, 10], [15, 25]]), false); +} + +#[test] +fn test_int_scatter_nd_slices() { + // K=1: each index selects a full row + let device = Default::default(); + let data = TestTensorInt::<2>::zeros([2, 3], &device); + let indices = TestTensorInt::<2>::from_ints([[0], [1]], &device); + let values = TestTensorInt::<2>::from_ints([[10, 20, 30], [40, 50, 60]], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Assign); + + output + .into_data() + .assert_eq(&TensorData::from([[10, 20, 30], [40, 50, 60]]), false); +} + +#[test] +fn test_int_gather_nd_slices() { + // K=1: each index gathers a full row + let device = Default::default(); + let data = TestTensorInt::<2>::from_ints([[10, 20, 30], [40, 50, 60]], &device); + let indices = TestTensorInt::<2>::from_ints([[1], [0]], &device); + + let output: TestTensorInt<2> = data.gather_nd(indices); + + output + .into_data() + .assert_eq(&TensorData::from([[40, 50, 60], [10, 20, 30]]), false); +} + +#[test] +fn test_int_scatter_nd_single_element() { + let device = Default::default(); + let data = TestTensorInt::<1>::from_ints([1, 2, 3], &device); + let indices = TestTensorInt::<2>::from_ints([[0]], &device); + let values = TestTensorInt::<1>::from_ints([99], &device); + + let output = data.scatter_nd(indices, values, IndexingUpdateOp::Assign); + + output + .into_data() + .assert_eq(&TensorData::from([99, 2, 3]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/init.rs b/crates/burn-backend-tests/tests/tensor/int/ops/init.rs new file mode 100644 index 0000000..b95fa38 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/init.rs @@ -0,0 +1,31 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_int_empty() { + let shape = [2, 2]; + let tensor = TestTensorInt::<2>::empty(shape, &Default::default()); + assert_eq!(tensor.shape(), shape.into()) +} + +#[test] +fn should_support_int_zeros() { + let shape = [2, 2]; + let tensor = TestTensorInt::<2>::zeros(shape, &Default::default()); + assert_eq!(tensor.shape(), shape.into()); + + tensor + .into_data() + .assert_eq(&TensorData::from([[0, 0], [0, 0]]), false); +} + +#[test] +fn should_support_int_ones() { + let shape = [2, 2]; + let tensor = TestTensorInt::<2>::ones(shape, &Default::default()); + assert_eq!(tensor.shape(), shape.into()); + + tensor + .into_data() + .assert_eq(&TensorData::from([[1, 1], [1, 1]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/mask.rs b/crates/burn-backend-tests/tests/tensor/int/ops/mask.rs new file mode 100644 index 0000000..90c027d --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/mask.rs @@ -0,0 +1,69 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_mask_where_broadcast_int() { + let device = Default::default(); + // When broadcasted, the input [[2, 3], [4, 5]] is repeated 4 times + let tensor = TestTensorInt::<1>::arange(2..6, &device).reshape([1, 2, 2]); + let mask = TestTensorBool::<3>::from_bool( + TensorData::from([ + [[true, false], [false, true]], + [[false, true], [true, false]], + [[false, false], [false, false]], + [[true, true], [true, true]], + ]), + &device, + ); + let value = TestTensorInt::<3>::ones([4, 2, 2], &device); + + let output = tensor.mask_where(mask, value); + let expected = TensorData::from([ + [[1, 3], [4, 1]], + [[2, 1], [1, 5]], + [[2, 3], [4, 5]], + [[1, 1], [1, 1]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_int_mask_where_ops() { + let device = Default::default(); + let tensor = TestTensorInt::<2>::from_data([[1, 7], [2, 3]], &device); + let mask = + TestTensorBool::<2>::from_bool(TensorData::from([[true, false], [false, true]]), &device); + let value = TestTensorInt::<2>::from_data(TensorData::from([[8, 9], [10, 11]]), &device); + + let output = tensor.mask_where(mask, value); + let expected = TensorData::from([[8, 7], [2, 11]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_int_mask_fill_ops() { + let device = Default::default(); + let tensor = TestTensorInt::<2>::from_data([[1, 7], [2, 3]], &device); + let mask = + TestTensorBool::<2>::from_bool(TensorData::from([[true, false], [false, true]]), &device); + + let output = tensor.mask_fill(mask, 9); + let expected = TensorData::from([[9, 7], [2, 9]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_int_mask_fill_flipped() { + // [10, 20, 30, 40] flipped -> [40, 30, 20, 10]; mask [T, F, T, F] -> [-1, 30, -1, 10] + let tensor = TestTensorInt::<1>::from([10, 20, 30, 40]).flip([0]); + let mask = TestTensorBool::<1>::from([true, false, true, false]); + + let output = tensor.mask_fill(mask, -1); + + output + .into_data() + .assert_eq(&TensorData::from([-1, 30, -1, 10]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/matmul.rs b/crates/burn-backend-tests/tests/tensor/int/ops/matmul.rs new file mode 100644 index 0000000..c01af14 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/matmul.rs @@ -0,0 +1,203 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_int_matmul_d2() { + let device = Default::default(); + let tensor_1 = TestTensorInt::<2>::from_ints([[1, 7], [2, 3], [1, 5]], &device); + let tensor_2 = TestTensorInt::<2>::from_ints([[4, 7, 5], [2, 3, 5]], &device); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([[18, 28, 40], [14, 23, 25], [14, 22, 30]]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_int_matmul_d3() { + let device = Default::default(); + let tensor_1 = TestTensorInt::<3>::from_ints([[[1, 7], [2, 3]]], &device); + let tensor_2 = TestTensorInt::<3>::from_ints([[[4, 7], [2, 3]]], &device); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([[[18, 28], [14, 23]]]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_int_matmul_broadcast_1() { + let device = Default::default(); + let tensor_1 = TestTensorInt::<3>::from_ints([[[1, 7], [2, 3]]], &device); + let tensor_2 = TestTensorInt::from_ints([[[4, 7], [2, 3]], [[2, 5], [6, 3]]], &device); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([[[18, 28], [14, 23]], [[44, 26], [22, 19]]]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_int_matmul_broadcast_4d() { + let device = Default::default(); + // [2, 1, 2, 2] + let tensor_1 = TestTensorInt::<4>::from_ints([[[[1, 7], [2, 3]]], [[[2, 5], [6, 3]]]], &device); + // [1, 2, 2, 2] + let tensor_2 = TestTensorInt::from_ints([[[[9, 8], [1, 4]], [[2, 7], [3, 5]]]], &device); + + // [2, 1, 2, 2] @ [1, 2, 2, 2] -> [2, 2, 2, 2] + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([ + [[[16, 36], [21, 28]], [[23, 42], [13, 29]]], + [[[23, 36], [57, 60]], [[19, 39], [21, 57]]], + ]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_int_matmul_simple_1() { + let device = Default::default(); + let tensor_1 = TestTensorInt::<2>::from_ints([[5, 14], [14, 25]], &device); + let tensor_2 = TestTensorInt::from_ints([[3, 4, 5], [0, 1, 2]], &device); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([[15, 34, 53], [42, 81, 120]]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_int_matmul_4_3() { + if (IntElem::MAX as u32) < 324 { + return; + } + + let device = Default::default(); + let tensor_1 = + TestTensorInt::<2>::from_ints([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], &device); + let tensor_2 = + TestTensorInt::from_ints([[0, 1, 2], [4, 5, 6], [8, 9, 10], [12, 13, 14]], &device); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([[56, 62, 68], [152, 174, 196], [248, 286, 324]]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_int_matmul_trivial() { + if (IntElem::MAX as u32) < 506 { + return; + } + + let device = Default::default(); + + let tensor_1 = TestTensorInt::<1>::arange(0..16, &device).reshape([4, 4]); + + let tensor_3 = tensor_1.clone().matmul(tensor_1); + + tensor_3.into_data().assert_eq( + &TensorData::from([ + [56, 62, 68, 74], + [152, 174, 196, 218], + [248, 286, 324, 362], + [344, 398, 452, 506], + ]), + false, + ); +} + +#[test] +fn test_int_matmul_trivial_transposed() { + if (IntElem::MAX as u32) < 734 { + return; + } + + let device = Default::default(); + + let tensor_1 = TestTensorInt::<1>::arange(0..16, &device).reshape([4, 4]); + + let tensor_3 = tensor_1.clone().matmul(tensor_1.transpose()); + + tensor_3.into_data().assert_eq( + &TensorData::from([ + [14, 38, 62, 86], + [38, 126, 214, 302], + [62, 214, 366, 518], + [86, 302, 518, 734], + ]), + false, + ); +} + +#[test] +fn test_int_matmul_4_8() { + if (IntElem::MAX as u32) < 6092 { + return; + } + + let device = Default::default(); + + let tensor_1 = TestTensorInt::<1>::arange(0..32, &device).reshape([4, 8]); + + let tensor_3 = tensor_1.clone().matmul(tensor_1.transpose()); + + tensor_3.into_data().assert_eq( + &TensorData::from([ + [140, 364, 588, 812], + [364, 1100, 1836, 2572], + [588, 1836, 3084, 4332], + [812, 2572, 4332, 6092], + ]), + false, + ); +} + +#[test] +fn test_int_matmul_simple_2() { + let device = Default::default(); + let tensor_1 = TestTensorInt::<2>::from_ints([[1, 2, 3, 4]], &device); + let tensor_2 = TestTensorInt::from_ints([[3], [4], [5], [6]], &device); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([[50]]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_int_matmul_simple_3() { + let device = Default::default(); + let tensor_1 = + TestTensorInt::<2>::from_ints([[3, 3, 3], [4, 4, 4], [5, 5, 5], [6, 6, 6]], &device); + let tensor_2 = TestTensorInt::from_ints([[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]], &device); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([ + [9, 18, 27, 36], + [12, 24, 36, 48], + [15, 30, 45, 60], + [18, 36, 54, 72], + ]); + + tensor_3.into_data().assert_eq(&expected, false); +} + +#[test] +#[should_panic] +fn int_should_panic_when_inner_dimensions_are_not_equal() { + let device = Default::default(); + let tensor_1 = TestTensorInt::<2>::from_ints([[3, 3], [4, 4], [5, 5], [6, 6]], &device); + let tensor_2 = TestTensorInt::from_ints([[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]], &device); + + let tensor_3 = tensor_1.matmul(tensor_2); + let expected = TensorData::from([ + [9, 18, 27, 36], + [12, 24, 36, 48], + [15, 30, 45, 60], + [18, 36, 54, 72], + ]); + + tensor_3.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/mod.rs b/crates/burn-backend-tests/tests/tensor/int/ops/mod.rs new file mode 100644 index 0000000..c6d6d83 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/mod.rs @@ -0,0 +1,48 @@ +pub use super::*; // re-export test types + +mod abs; +mod add; +mod aggregation; +mod all; +mod any; +mod arange; +mod arange_step; +mod bitwise; +mod cartesian_grid; +mod cast; +mod cat; +mod chunk; +mod comparison; +mod create_like; +mod cumulative; +mod div; +mod expand; +mod flip; +mod full; +mod gather_scatter; +mod gather_scatter_nd; +mod init; +mod mask; +mod matmul; +mod movedim; +mod mul; +mod one_hot; +mod permute; +mod random; +mod remainder; +mod repeat; +mod repeat_dim; +mod reshape; +mod roll; +mod select; +mod sign; +mod slice; +mod slice_assign; +mod sort_argsort; +mod stack; +mod sub; +mod take; +mod topk; +mod transpose; +mod tri; +mod unfold; diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/movedim.rs b/crates/burn-backend-tests/tests/tensor/int/ops/movedim.rs new file mode 100644 index 0000000..613e37e --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/movedim.rs @@ -0,0 +1,52 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn movedim_int() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device).reshape([2, 3, 4]); + + let permuted = tensor.clone().movedim(0, 2); + // from pytorch: + // import torch; torch.arange(0, 24).reshape(2, 3, 4).movedim(0, 2) + let expected = TensorData::from([ + [[0, 12], [1, 13], [2, 14], [3, 15]], + [[4, 16], [5, 17], [6, 18], [7, 19]], + [[8, 20], [9, 21], [10, 22], [11, 23]], + ]); + + permuted.into_data().assert_eq(&expected, false); + + // Test with negative axis + let permuted = tensor.clone().movedim(0, -1); + permuted.into_data().assert_eq(&expected, false); + + // Test with the same axis + let permuted = tensor.clone().movedim(0, 0); + permuted.into_data().assert_eq(&tensor.into_data(), true); +} + +#[test] +fn vec_input_int() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device).reshape([2, 3, 4]); + + let permuted = tensor.clone().movedim(vec![0, 1], vec![1, 0]); + // from pytorch + // import torch; torch.arange(0, 24).reshape(2, 3, 4).movedim([0, 1], [1, 0]) + let expected = TensorData::from([ + [[0, 1, 2, 3], [12, 13, 14, 15]], + [[4, 5, 6, 7], [16, 17, 18, 19]], + [[8, 9, 10, 11], [20, 21, 22, 23]], + ]); + + permuted.into_data().assert_eq(&expected, false); + + // Test with negative axes + let permuted = tensor.clone().movedim(vec![-3, -2], vec![-2, -3]); + permuted.into_data().assert_eq(&expected, false); + + // Test with the same axes + let permuted = tensor.clone().movedim(vec![0, 1], vec![0, 1]); + permuted.into_data().assert_eq(&tensor.into_data(), true); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/mul.rs b/crates/burn-backend-tests/tests/tensor/int/ops/mul.rs new file mode 100644 index 0000000..1ab881a --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/mul.rs @@ -0,0 +1,70 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_mul_ops_int() { + let data_1 = TensorData::from([[0, 1, 2], [3, 4, 5]]); + let data_2 = TensorData::from([[0, 1, 2], [3, 4, 5]]); + let device = Default::default(); + let tensor_1 = TestTensorInt::<2>::from_data(data_1, &device); + let tensor_2 = TestTensorInt::<2>::from_data(data_2, &device); + + let output = tensor_1 * tensor_2; + let expected = TensorData::from([[0, 1, 4], [9, 16, 25]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_mul_broadcast_int() { + let data_1 = TensorData::from([[0, 1, 2]]); + let data_2 = TensorData::from([[3, 4, 5], [6, 7, 8]]); + let device = Default::default(); + let tensor_1 = TestTensorInt::<2>::from_data(data_1, &device); + let tensor_2 = TestTensorInt::<2>::from_data(data_2, &device); + + let output = tensor_1 * tensor_2; + let expected = TensorData::from([[0, 4, 10], [0, 7, 16]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_mul_scalar_ops_int() { + let data = TensorData::from([[0, 1, 2], [3, 4, 5]]); + let scalar = 2; + let tensor = TestTensorInt::<2>::from_data(data, &Default::default()); + + let output = tensor * scalar; + let expected = TensorData::from([[0, 2, 4], [6, 8, 10]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_int_mul_flipped_2d() { + // [[1, 2], [3, 4]] axis 0 flipped -> [[3, 4], [1, 2]] + // * [[10, 20], [30, 40]] = [[30, 80], [30, 80]] + let a = TestTensorInt::<2>::from([[1, 2], [3, 4]]).flip([0]); + let b = TestTensorInt::<2>::from([[10, 20], [30, 40]]); + + let output = a * b; + + output + .into_data() + .assert_eq(&TensorData::from([[30, 80], [30, 80]]), false); +} + +#[test] +fn test_int_mul_flipped_both_axes() { + // [[1, 2], [3, 4]] flipped on both axes -> [[4, 3], [2, 1]] + // * [[5, 5], [5, 5]] = [[20, 15], [10, 5]] + let a = TestTensorInt::<2>::from([[1, 2], [3, 4]]).flip([0, 1]); + let b = TestTensorInt::<2>::from([[5, 5], [5, 5]]); + + let output = a * b; + + output + .into_data() + .assert_eq(&TensorData::from([[20, 15], [10, 5]]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/one_hot.rs b/crates/burn-backend-tests/tests/tensor/int/ops/one_hot.rs new file mode 100644 index 0000000..cd6c0aa --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/one_hot.rs @@ -0,0 +1,59 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn int_should_support_one_hot() { + let tensor = TestTensorInt::<1>::from([0, 1, 4]); + let one_hot_tensor: TestTensorInt<2> = tensor.one_hot(5); + let expected = TensorData::from([[1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 0, 0, 1]]); + one_hot_tensor.into_data().assert_eq(&expected, false); +} + +#[test] +#[should_panic] +fn int_one_hot_should_panic_when_index_exceeds_number_of_classes() { + let tensor = TestTensorInt::<1>::from([5]); + let _result: TestTensorInt<2> = tensor.one_hot(5); +} + +#[test] +#[should_panic] +fn int_one_hot_should_panic_when_number_of_classes_is_zero() { + let tensor = TestTensorInt::<1>::from([2]); + let _result: TestTensorInt<2> = tensor.one_hot(0); +} + +#[test] +fn one_hot_fill_with_positive_axis_and_indices() { + let tensor = TestTensorInt::<2>::from([[1, 9], [2, 4]]); + let expected = TensorData::from([ + [ + [1, 1], + [3, 1], + [1, 1], + [1, 1], + [1, 1], + [1, 1], + [1, 1], + [1, 1], + [1, 1], + [1, 3], + ], + [ + [1, 1], + [1, 1], + [3, 1], + [1, 1], + [1, 3], + [1, 1], + [1, 1], + [1, 1], + [1, 1], + [1, 1], + ], + ]); + + let one_hot_tensor: TestTensorInt<3> = tensor.one_hot_fill(10, 3.0, 1.0, 1); + + one_hot_tensor.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/permute.rs b/crates/burn-backend-tests/tests/tensor/int/ops/permute.rs new file mode 100644 index 0000000..3f008f8 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/permute.rs @@ -0,0 +1,49 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn permute_int() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device).reshape([2, 3, 4]); + + let permuted = tensor.clone().permute([2, 1, 0]); + + // from pytorch: + // import torch; torch.arange(0, 24).reshape(2, 3, 4).permute(2, 1, 0) + let expected = TensorData::from([ + [[0, 12], [4, 16], [8, 20]], + [[1, 13], [5, 17], [9, 21]], + [[2, 14], [6, 18], [10, 22]], + [[3, 15], [7, 19], [11, 23]], + ]); + + permuted.into_data().assert_eq(&expected, false); + + // Test with negative axis + let permuted = tensor.clone().permute([-1, 1, 0]); + permuted.into_data().assert_eq(&expected, false); + + // Test with the same axis + let permuted = tensor.clone().permute([0, 1, 2]); + permuted.into_data().assert_eq(&tensor.into_data(), true); +} + +#[test] +#[should_panic] +fn edge_repeated_axes() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device).reshape([2, 3, 4]); + + // Test with a repeated axis + let _ = tensor.clone().permute([0, 0, 1]); +} + +#[test] +#[should_panic] +fn edge_out_of_bound_axis() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(0..24, &device).reshape([2, 3, 4]); + + // Test with a repeated axis + let _ = tensor.clone().permute([3, 0, 1]); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/random.rs b/crates/burn-backend-tests/tests/tensor/int/ops/random.rs new file mode 100644 index 0000000..83138c1 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/random.rs @@ -0,0 +1,18 @@ +use super::*; +use burn_tensor::{Distribution, ElementConversion}; + +#[test] +fn rand_uniform_int() { + let low = 0.; + let high = 5.; + + let tensor = TestTensorInt::<1>::random( + [100_000], + Distribution::Uniform(low, high), + &Default::default(), + ); + + tensor + .into_data() + .assert_within_range::(low.elem()..high.elem()); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/remainder.rs b/crates/burn-backend-tests/tests/tensor/int/ops/remainder.rs new file mode 100644 index 0000000..c0a216e --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/remainder.rs @@ -0,0 +1,39 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_int_remainder_basic() { + let data = TensorData::from([-3, -2, -1, 1, 2, 3]); + let device = Default::default(); + let lhs = TestTensorInt::<1>::from_data(data, &device); + + let rhs = TestTensorInt::from_data(TensorData::from([2, 3, 1, 2, 1, 3]), &device); + let output = lhs.remainder(rhs); + let expected = TensorData::from([1, 1, -0, 1, 0, 0]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_int_remainder_broadcast() { + let device = Default::default(); + let lhs = TestTensorInt::<2>::from_data(TensorData::from([[10, 20, 30]]), &device); + let rhs = TestTensorInt::<2>::from_data(TensorData::from([[7]]), &device); + + let output = lhs.remainder(rhs); + let expected = TensorData::from([[3, 6, 2]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_int_remainder_basic_scalar() { + let data = TensorData::from([-3, -2, -1, 1, 2, 3]); + let device = Default::default(); + let tensor = TestTensorInt::<1>::from_data(data, &device); + + let output = tensor.remainder_scalar(2); + let expected = TensorData::from([1, 0, 1, 1, 0, 1]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/repeat.rs b/crates/burn-backend-tests/tests/tensor/int/ops/repeat.rs new file mode 100644 index 0000000..44e488a --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/repeat.rs @@ -0,0 +1,100 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_int_repeat_ops_one_dimension() { + let data = TensorData::from([[0i32, 1i32, 2i32]]); + let tensor = TestTensorInt::<2>::from_data(data, &Default::default()); + + let output = tensor.repeat(&[4, 1, 1]); + let expected = TensorData::from([ + [0i32, 1i32, 2i32], + [0i32, 1i32, 2i32], + [0i32, 1i32, 2i32], + [0i32, 1i32, 2i32], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_int_repeat_on_many_dims() { + let data = TensorData::from([ + [[1i32, 2i32], [3i32, 4i32]], + [[5i32, 6i32], [7i32, 8i32]], + [[9i32, 10i32], [11i32, 12i32]], + [[13i32, 14i32], [15i32, 16i32]], + ]); + let tensor = TestTensorInt::<3>::from_data(data, &Default::default()); + + let output = tensor.repeat(&[2, 3, 2]); + + let expected = TensorData::from([ + [ + [1i32, 2i32, 1i32, 2i32], + [3i32, 4i32, 3i32, 4i32], + [1i32, 2i32, 1i32, 2i32], + [3i32, 4i32, 3i32, 4i32], + [1i32, 2i32, 1i32, 2i32], + [3i32, 4i32, 3i32, 4i32], + ], + [ + [5i32, 6i32, 5i32, 6i32], + [7i32, 8i32, 7i32, 8i32], + [5i32, 6i32, 5i32, 6i32], + [7i32, 8i32, 7i32, 8i32], + [5i32, 6i32, 5i32, 6i32], + [7i32, 8i32, 7i32, 8i32], + ], + [ + [9i32, 10i32, 9i32, 10i32], + [11i32, 12i32, 11i32, 12i32], + [9i32, 10i32, 9i32, 10i32], + [11i32, 12i32, 11i32, 12i32], + [9i32, 10i32, 9i32, 10i32], + [11i32, 12i32, 11i32, 12i32], + ], + [ + [13i32, 14i32, 13i32, 14i32], + [15i32, 16i32, 15i32, 16i32], + [13i32, 14i32, 13i32, 14i32], + [15i32, 16i32, 15i32, 16i32], + [13i32, 14i32, 13i32, 14i32], + [15i32, 16i32, 15i32, 16i32], + ], + [ + [1i32, 2i32, 1i32, 2i32], + [3i32, 4i32, 3i32, 4i32], + [1i32, 2i32, 1i32, 2i32], + [3i32, 4i32, 3i32, 4i32], + [1i32, 2i32, 1i32, 2i32], + [3i32, 4i32, 3i32, 4i32], + ], + [ + [5i32, 6i32, 5i32, 6i32], + [7i32, 8i32, 7i32, 8i32], + [5i32, 6i32, 5i32, 6i32], + [7i32, 8i32, 7i32, 8i32], + [5i32, 6i32, 5i32, 6i32], + [7i32, 8i32, 7i32, 8i32], + ], + [ + [9i32, 10i32, 9i32, 10i32], + [11i32, 12i32, 11i32, 12i32], + [9i32, 10i32, 9i32, 10i32], + [11i32, 12i32, 11i32, 12i32], + [9i32, 10i32, 9i32, 10i32], + [11i32, 12i32, 11i32, 12i32], + ], + [ + [13i32, 14i32, 13i32, 14i32], + [15i32, 16i32, 15i32, 16i32], + [13i32, 14i32, 13i32, 14i32], + [15i32, 16i32, 15i32, 16i32], + [13i32, 14i32, 13i32, 14i32], + [15i32, 16i32, 15i32, 16i32], + ], + ]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/repeat_dim.rs b/crates/burn-backend-tests/tests/tensor/int/ops/repeat_dim.rs new file mode 100644 index 0000000..fd29a32 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/repeat_dim.rs @@ -0,0 +1,46 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_int_repeat_ops() { + let data = TensorData::from([[0, 1, 2]]); + let tensor = TestTensorInt::<2>::from_data(data, &Default::default()); + + let output = tensor.repeat_dim(0, 4); + let expected = TensorData::from([[0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_int_repeat_on_dims_larger_than_1() { + let data = TensorData::from([ + [[1i32, 2i32], [3i32, 4i32]], + [[5i32, 6i32], [7i32, 8i32]], + [[9i32, 10i32], [11i32, 12i32]], + [[13i32, 14i32], [15i32, 16i32]], + ]); + let tensor = TestTensorInt::<3>::from_data(data, &Default::default()); + + let output = tensor.repeat_dim(2, 3); + let expected = TensorData::from([ + [ + [1i32, 2i32, 1i32, 2i32, 1i32, 2i32], + [3i32, 4i32, 3i32, 4i32, 3i32, 4i32], + ], + [ + [5i32, 6i32, 5i32, 6i32, 5i32, 6i32], + [7i32, 8i32, 7i32, 8i32, 7i32, 8i32], + ], + [ + [9i32, 10i32, 9i32, 10i32, 9i32, 10i32], + [11i32, 12i32, 11i32, 12i32, 11i32, 12i32], + ], + [ + [13i32, 14i32, 13i32, 14i32, 13i32, 14i32], + [15i32, 16i32, 15i32, 16i32, 15i32, 16i32], + ], + ]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/reshape.rs b/crates/burn-backend-tests/tests/tensor/int/ops/reshape.rs new file mode 100644 index 0000000..ba9dee1 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/reshape.rs @@ -0,0 +1,179 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_reshape_maybe_fused_1() { + let tensor = TestTensorInt::arange(0..32, &Default::default()); + let tensor0 = TestTensorInt::zeros([8, 4, 8], &Default::default()); + let tensor1 = tensor.clone().reshape([1, 4, 8]); + let output = tensor0 + tensor1; + + let expected = TensorData::from([ + [ + [0, 1, 2, 3, 4, 5, 6, 7], + [8, 9, 10, 11, 12, 13, 14, 15], + [16, 17, 18, 19, 20, 21, 22, 23], + [24, 25, 26, 27, 28, 29, 30, 31], + ], + [ + [0, 1, 2, 3, 4, 5, 6, 7], + [8, 9, 10, 11, 12, 13, 14, 15], + [16, 17, 18, 19, 20, 21, 22, 23], + [24, 25, 26, 27, 28, 29, 30, 31], + ], + [ + [0, 1, 2, 3, 4, 5, 6, 7], + [8, 9, 10, 11, 12, 13, 14, 15], + [16, 17, 18, 19, 20, 21, 22, 23], + [24, 25, 26, 27, 28, 29, 30, 31], + ], + [ + [0, 1, 2, 3, 4, 5, 6, 7], + [8, 9, 10, 11, 12, 13, 14, 15], + [16, 17, 18, 19, 20, 21, 22, 23], + [24, 25, 26, 27, 28, 29, 30, 31], + ], + [ + [0, 1, 2, 3, 4, 5, 6, 7], + [8, 9, 10, 11, 12, 13, 14, 15], + [16, 17, 18, 19, 20, 21, 22, 23], + [24, 25, 26, 27, 28, 29, 30, 31], + ], + [ + [0, 1, 2, 3, 4, 5, 6, 7], + [8, 9, 10, 11, 12, 13, 14, 15], + [16, 17, 18, 19, 20, 21, 22, 23], + [24, 25, 26, 27, 28, 29, 30, 31], + ], + [ + [0, 1, 2, 3, 4, 5, 6, 7], + [8, 9, 10, 11, 12, 13, 14, 15], + [16, 17, 18, 19, 20, 21, 22, 23], + [24, 25, 26, 27, 28, 29, 30, 31], + ], + [ + [0, 1, 2, 3, 4, 5, 6, 7], + [8, 9, 10, 11, 12, 13, 14, 15], + [16, 17, 18, 19, 20, 21, 22, 23], + [24, 25, 26, 27, 28, 29, 30, 31], + ], + ]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_reshape_maybe_fused_2() { + let tensor = TestTensorInt::<3>::from_data([[[0, 2], [1, 2]]], &Default::default()); + let tensor1 = tensor.reshape([2, 2, 1]); + let tensor2 = TestTensorInt::<3>::full([2, 2, 4], 4, &Default::default()); + let output = tensor2 + tensor1; + + let expected_tensor1 = + TensorData::from([[[4, 4, 4, 4], [6, 6, 6, 6]], [[5, 5, 5, 5], [6, 6, 6, 6]]]); + output.into_data().assert_eq(&expected_tensor1, false); +} + +#[test] +fn should_support_reshape_maybe_fused_3() { + let tensor = TestTensorInt::<3>::from_data([[[0, 2], [1, 2]]], &Default::default()); + let tensor1 = tensor.reshape([2, 2, 1]); + let _tensor2 = TestTensorInt::<3>::full([2, 2, 3], 5, &Default::default()); + + let expected_tensor1 = TensorData::from([[[0], [2]], [[1], [2]]]); + tensor1.into_data().assert_eq(&expected_tensor1, false); +} + +#[test] +fn should_support_reshape_maybe_fused_4() { + let tensor = TestTensorInt::<3>::from_data([[[0, 2], [1, 2]]], &Default::default()); + let tensor2 = TestTensorInt::<3>::full([2, 2, 4], 4, &Default::default()); + let tensor2 = tensor2.swap_dims(0, 1); + let tensor1 = tensor.reshape([2, 2, 1]); + let output = tensor2 + tensor1; + + let expected_tensor1 = + TensorData::from([[[4, 4, 4, 4], [6, 6, 6, 6]], [[5, 5, 5, 5], [6, 6, 6, 6]]]); + output.into_data().assert_eq(&expected_tensor1, false); +} + +#[test] +fn should_support_reshape_maybe_fused_5() { + let tensor = TestTensorInt::<3>::from_data([[[0], [1], [2], [3]]], &Default::default()); + let tensor1 = tensor.clone().reshape([2, 1, 2]); + let tensor2 = TestTensorInt::<3>::full([2, 4, 2], 0, &Default::default()); + let output = tensor2.clone() + tensor1 + tensor.clone(); + + let expected_tensor1 = TensorData::from([ + [[0, 1], [1, 2], [2, 3], [3, 4]], + [[2, 3], [3, 4], [4, 5], [5, 6]], + ]); + output.into_data().assert_eq(&expected_tensor1, false); +} + +#[test] +fn should_support_reshape_maybe_fused_6() { + let device = Default::default(); + + let tensor1 = TestTensorInt::arange(0..32, &device); + let tensor1 = tensor1.reshape([2, 4, 4]); + + let tensor2 = TestTensorInt::arange(0..16, &device); + let tensor2 = tensor2.reshape([1, 4, 4]); + + let tensor3 = TestTensorInt::arange(0..8, &device); + let tensor3 = tensor3.reshape([4, 1, 2]); + let tensor3 = tensor3.swap_dims(0, 2); + + let out = tensor1 + tensor2 + tensor3; + + let expected = TensorData::from([ + [ + [0, 4, 8, 12], + [8, 12, 16, 20], + [16, 20, 24, 28], + [24, 28, 32, 36], + ], + [ + [17, 21, 25, 29], + [25, 29, 33, 37], + [33, 37, 41, 45], + [41, 45, 49, 53], + ], + ]); + out.to_data().assert_eq(&expected, false); +} + +// Skip on metal - cubecl autotune error +// Enable once this issue is fixed: https://github.com/tracel-ai/burn/issues/4327 +#[cfg(not(feature = "metal"))] +#[test] +fn should_support_multiple_reshapes_cloned_tensor() { + let device = Default::default(); + + let lhs = TestTensorInt::<1>::arange(0..4, &device).reshape([2, 2]); + // fusion should preserve correct strides when operating on the same tensor + let rhs = lhs.clone(); + + let lhs = lhs.reshape([2, 2, 1]); + let rhs = rhs.reshape([1, 2, 2]); + + let p = lhs.mul(rhs); + + let s = p.sum_dim(1); + + let out = s.reshape([2, 2]); + + out.into_data() + .assert_eq(&TensorData::from([[2, 3], [6, 11]]), false); +} + +#[test] +fn should_support_reshape_int() { + let data = TensorData::from([0, 1, 2]); + let tensor = TestTensorInt::<1>::from_data(data, &Default::default()); + + let output = tensor.clone().reshape([1, 3]); + let expected = TensorData::from([[0, 1, 2]]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/roll.rs b/crates/burn-backend-tests/tests/tensor/int/ops/roll.rs new file mode 100644 index 0000000..e53e928 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/roll.rs @@ -0,0 +1,108 @@ +use super::*; +use burn_tensor::TensorData; + +#[ignore = "0 size resources are not yet supported"] +#[test] +fn test_roll_empty() { + let device = Default::default(); + let input = TestTensorInt::<2>::zeros([12, 0], &device); + + let result = input.clone().roll(&[1, 2], &[0, 1]); + + assert_eq!(&*result.shape(), &[12, 0]); + + // TODO: Rolling an empty tensor should return the same empty tensor; + // but we have no way to compare tensor references yet. +} + +#[test] +fn test_roll() { + let input = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + + // No-op shift: + input + .clone() + .roll(&[0, 0], &[0, 1]) + .to_data() + .assert_eq(&input.clone().to_data(), false); + + input + .clone() + .roll(&[1, -1], &[0, 1]) + .to_data() + .assert_eq(&TensorData::from([[5, 3, 4], [2, 0, 1]]), false); + + input + .clone() + .roll(&[-1, 1], &[1, 0]) + .to_data() + .assert_eq(&TensorData::from([[5, 3, 4], [2, 0, 1]]), false); + + input + .clone() + .roll(&[2 * 32 + 1, 3 * (-400) - 1], &[0, 1]) + .to_data() + .assert_eq(&TensorData::from([[5, 3, 4], [2, 0, 1]]), false); +} + +#[should_panic] +#[test] +fn test_roll_dim_too_big() { + let input = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + + // Attempting to roll on a dimension that doesn't exist should panic + let _d = input.roll(&[1], &[2]); +} + +#[should_panic] +#[test] +fn test_roll_dim_too_small() { + let input = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + + // Attempting to roll on a dimension that doesn't exist should panic + let _d = input.roll(&[1], &[-3]); +} + +#[should_panic] +#[test] +fn test_roll_shift_size_mismatch() { + let input = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + + // Attempting to roll with a shift size that doesn't match the number of dimensions should panic + let _d = input.roll(&[1, 2], &[0]); +} + +#[test] +fn test_roll_dim() { + let input = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + + input + .clone() + .roll_dim(1, 0) + .to_data() + .assert_eq(&TensorData::from([[3, 4, 5], [0, 1, 2]]), false); + + input + .clone() + .roll_dim(-1, 1) + .to_data() + .assert_eq(&TensorData::from([[2, 0, 1], [5, 3, 4]]), false); +} + +#[should_panic] +#[test] +fn test_roll_dim_dim_too_big() { + let input = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + + // Attempting to roll on a dimension that doesn't exist should panic + let _d = input.roll_dim(1, 2); +} + +#[should_panic] +#[test] +fn test_roll_dim_dim_too_small() { + let input = TestTensorInt::<2>::from([[0, 1, 2], [3, 4, 5]]); + + // Attempting to roll on a dimension that doesn't exist should panic + let _d = input.roll_dim(1, -3); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/select.rs b/crates/burn-backend-tests/tests/tensor/int/ops/select.rs new file mode 100644 index 0000000..ef1d5c4 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/select.rs @@ -0,0 +1,38 @@ +use super::*; +use burn_tensor::{IndexingUpdateOp, TensorData}; + +#[test] +fn should_select_1d_int() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::from_data([5, 6, 7], &device); + let indices = TestTensorInt::from_data([1, 1, 0, 1, 2], &device); + + let output = tensor.select(0, indices); + let expected = TensorData::from([6, 6, 5, 6, 7]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_select_add_1d_int() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::from_data([7, 8, 9], &device); + let values = TestTensorInt::from_data([5, 4, 3, 2, 1], &device); + let indices = TestTensorInt::from_data(TensorData::from([1, 1, 0, 1, 2]), &device); + + let output = tensor.select_assign(0, indices, values, IndexingUpdateOp::Add); + let expected = TensorData::from([10, 19, 10]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +#[should_panic] +fn should_panic_select_add_invalid_num_indices() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::from_data([0; 12], &device); + let values = TestTensorInt::from_data([1; 12], &device); + let indices = TestTensorInt::from_data(TensorData::from([1]), &device); + + tensor.select_assign(0, indices, values, IndexingUpdateOp::Add); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/sign.rs b/crates/burn-backend-tests/tests/tensor/int/ops/sign.rs new file mode 100644 index 0000000..db4f6f6 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/sign.rs @@ -0,0 +1,12 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_sign_ops_int() { + let tensor = TestTensorInt::<2>::from([[-2, -1, 2], [3, 0, -5]]); + + let output = tensor.sign(); + let expected = TensorData::from([[-1, -1, 1], [1, 0, -1]]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/slice.rs b/crates/burn-backend-tests/tests/tensor/int/ops/slice.rs new file mode 100644 index 0000000..2e725bf --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/slice.rs @@ -0,0 +1,29 @@ +use super::*; +use burn_tensor::{TensorData, s}; + +#[test] +fn slice_should_not_corrupt_potentially_inplace_operations() { + let tensor = TestTensorInt::<1>::from([1, 2, 3, 4, 5]); + let tensor = tensor.clone().slice([0..3]) + tensor.clone().slice([2..5]); + + let expected = TensorData::from([4, 6, 8]); + + tensor.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_slice_int_tensor_with_steps() { + let device = Default::default(); + let tensor = + TestTensorInt::<2>::from_data([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], &device); + + // Test step=2 along first dimension + let sliced = tensor.clone().slice([s![0..3;2]]); + let expected = TensorData::from([[1i32, 2, 3, 4], [9, 10, 11, 12]]); + sliced.into_data().assert_eq(&expected, false); + + // Test step=-1 along second dimension + let sliced = tensor.clone().slice(s![.., 0..4;-1]); + let expected = TensorData::from([[4i32, 3, 2, 1], [8, 7, 6, 5], [12, 11, 10, 9]]); + sliced.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/slice_assign.rs b/crates/burn-backend-tests/tests/tensor/int/ops/slice_assign.rs new file mode 100644 index 0000000..5e5b1f6 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/slice_assign.rs @@ -0,0 +1,55 @@ +use super::*; +use burn_tensor::{TensorData, s}; + +#[test] +fn slice_assign_should_not_corrupt_potentially_inplace_operations() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::from_data([1, 2, 3, 4, 5], &device); + let values = TestTensorInt::<1>::from_data([10, 20, 30], &device); + let tensor_1 = tensor.clone().slice_assign([0..3], values); + let tensor_2 = tensor + 2; + + let expected = TensorData::from([10, 20, 30, 4, 5]); + + tensor_1.into_data().assert_eq(&expected, false); + + let expected = TensorData::from([3, 4, 5, 6, 7]); + + tensor_2.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_slice_assign_int_tensor_with_steps() { + let device = Default::default(); + let tensor = + TestTensorInt::<2>::from_data([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], &device); + + // Test step=2 along first dimension + let values = + TestTensorInt::<2>::from_data([[100, 101, 102, 103], [200, 201, 202, 203]], &device); + let output = tensor.clone().slice_assign([s![0..3;2]], values); + let expected = TensorData::from([[100i32, 101, 102, 103], [5, 6, 7, 8], [200, 201, 202, 203]]); + output.into_data().assert_eq(&expected, false); + + // Test step=-1 along second dimension + let values = TestTensorInt::<2>::from_data( + [[40, 30, 20, 10], [80, 70, 60, 50], [120, 110, 100, 90]], + &device, + ); + let output = tensor.slice_assign(s![.., 0..4;-1], values); + let expected = TensorData::from([[10i32, 20, 30, 40], [50, 60, 70, 80], [90, 100, 110, 120]]); + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_slice_assign_empty_range_int() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::from_data([1, 2, 3, 4, 5], &device); + let values: TestTensorInt<1> = TestTensorInt::empty([0], &device); + + // Empty slice assignment for int tensor + let output = tensor.clone().slice_assign([3..3], values); + let expected = TensorData::from([1i32, 2, 3, 4, 5]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/sort_argsort.rs b/crates/burn-backend-tests/tests/tensor/int/ops/sort_argsort.rs new file mode 100644 index 0000000..21322db --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/sort_argsort.rs @@ -0,0 +1,153 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_sort_1d_int() { + // Skip with u8 + if (IntElem::MAX as u32) < 1000u32 { + return; + } + + let tensor = TestTensorInt::<1>::from([1, 4, 7, 2, 5, 6, 3, 0, 9, 8, 2, 8, -10, 42, 1000]); + + // Sort along dim=0 + let values = tensor.sort(0); + let values_expected = TensorData::from([-10, 0, 1, 2, 2, 3, 4, 5, 6, 7, 8, 8, 9, 42, 1000]); + + values.into_data().assert_eq(&values_expected, false); +} + +#[test] +fn test_argsort_1d_int() { + // Skip with u8 + if (IntElem::MAX as u32) < 1000u32 { + return; + } + + let tensor = TestTensorInt::<1>::from([1, 4, 7, 2, 5, 6, 3, 0, 9, 8, -10, 42, 1000]); + + // Sort along dim=0 + let indices = tensor.argsort(0); + let indices_expected = TensorData::from([10, 7, 0, 3, 6, 1, 4, 5, 2, 9, 8, 11, 12]); + + indices.into_data().assert_eq(&indices_expected, false); +} + +#[test] +fn test_sort_with_indices_descending_int() { + // Skip with u8 + if (IntElem::MAX as u32) >= 1000u32 { + // 1D + let tensor = TestTensorInt::<1>::from([1, 4, 7, 2, 5, 6, 3, 0, 9, 8, -10, 42, 1000]); + + // Sort along dim=0 + let (values, indices) = tensor.sort_descending_with_indices(0); + + let values_expected = TensorData::from([1000, 42, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0, -10]); + values.into_data().assert_eq(&values_expected, false); + + let indices_expected = TensorData::from([12, 11, 8, 9, 2, 5, 4, 1, 6, 3, 0, 7, 10]); + indices.into_data().assert_eq(&indices_expected, false); + } + + // 2D + let tensor = TestTensorInt::<3>::from([[[1, 4, 7], [2, 5, 6]], [[3, 0, 9], [8, 2, 8]]]); + + // Sort along dim=1 + let (values, indices) = tensor.sort_descending_with_indices(1); + + let values_expected = TensorData::from([[[2, 5, 7], [1, 4, 6]], [[8, 2, 9], [3, 0, 8]]]); + values.into_data().assert_eq(&values_expected, false); + + let indices_expected = TensorData::from([[[1, 1, 0], [0, 0, 1]], [[1, 1, 0], [0, 0, 1]]]); + indices.into_data().assert_eq(&indices_expected, false); +} + +#[test] +fn test_sort_int() { + let tensor = TestTensorInt::<3>::from([[[1, 4, 7], [2, 5, 6]], [[3, 0, 9], [8, 2, 8]]]); + + // Sort along dim=0 + let values = tensor.clone().sort(0); + + let values_expected = TensorData::from([[[1, 0, 7], [2, 2, 6]], [[3, 4, 9], [8, 5, 8]]]); + values.into_data().assert_eq(&values_expected, false); + + // Sort along dim=1 + let values = tensor.clone().sort(1); + + let values_expected = TensorData::from([[[1, 4, 6], [2, 5, 7]], [[3, 0, 8], [8, 2, 9]]]); + values.into_data().assert_eq(&values_expected, false); + + // Sort along dim=2 + let values = tensor.sort(2); + + let values_expected = TensorData::from([[[1, 4, 7], [2, 5, 6]], [[0, 3, 9], [2, 8, 8]]]); + values.into_data().assert_eq(&values_expected, false); +} + +#[test] +fn test_sort_with_indices_int() { + let tensor = TestTensorInt::<3>::from([[[1, 4, 7], [2, 5, 6]], [[3, 0, 9], [7, 2, 8]]]); + + // Sort along dim=0 + let (values, indices) = tensor.clone().sort_with_indices(0); + + let values_expected = TensorData::from([[[1, 0, 7], [2, 2, 6]], [[3, 4, 9], [7, 5, 8]]]); + values.into_data().assert_eq(&values_expected, false); + + let indices_expected = TensorData::from([[[0, 1, 0], [0, 1, 0]], [[1, 0, 1], [1, 0, 1]]]); + indices.into_data().assert_eq(&indices_expected, false); + + // Sort along dim=1 + let (values, indices) = tensor.clone().sort_with_indices(1); + + let values_expected = TensorData::from([[[1, 4, 6], [2, 5, 7]], [[3, 0, 8], [7, 2, 9]]]); + values.into_data().assert_eq(&values_expected, false); + + let indices_expected = TensorData::from([[[0, 0, 1], [1, 1, 0]], [[0, 0, 1], [1, 1, 0]]]); + indices.into_data().assert_eq(&indices_expected, false); + + // Sort along dim=2 + let (values, indices) = tensor.sort_with_indices(2); + + let values_expected = TensorData::from([[[1, 4, 7], [2, 5, 6]], [[0, 3, 9], [2, 7, 8]]]); + values.into_data().assert_eq(&values_expected, false); + + let indices_expected = TensorData::from([[[0, 1, 2], [0, 1, 2]], [[1, 0, 2], [1, 0, 2]]]); + indices.into_data().assert_eq(&indices_expected, false); +} + +#[test] +fn test_argsort_int() { + let tensor = TestTensorInt::<3>::from([[[1, 4, 7], [2, 5, 6]], [[3, 0, 9], [7, 2, 8]]]); + + // Sort along dim=0 + let indices = tensor.clone().argsort(0); + + let indices_expected = TensorData::from([[[0, 1, 0], [0, 1, 0]], [[1, 0, 1], [1, 0, 1]]]); + indices.into_data().assert_eq(&indices_expected, false); + + // Sort along dim=1 + let indices = tensor.clone().argsort(1); + + let indices_expected = TensorData::from([[[0, 0, 1], [1, 1, 0]], [[0, 0, 1], [1, 1, 0]]]); + indices.into_data().assert_eq(&indices_expected, false); + + // Sort along dim=2 + let indices = tensor.argsort(2); + + let indices_expected = TensorData::from([[[0, 1, 2], [0, 1, 2]], [[1, 0, 2], [1, 0, 2]]]); + indices.into_data().assert_eq(&indices_expected, false); +} + +#[test] +fn test_sort_descending_1d() { + let tensor = TestTensorInt::<1>::from([1, 2, 3, 4, 5]); + + // Sort along dim=0 + let values = tensor.sort_descending(0); + + let values_expected = TensorData::from([5, 4, 3, 2, 1]); + values.into_data().assert_eq(&values_expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/stack.rs b/crates/burn-backend-tests/tests/tensor/int/ops/stack.rs new file mode 100644 index 0000000..ebda5ce --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/stack.rs @@ -0,0 +1,33 @@ +use super::*; +use alloc::vec; +use burn_tensor::{Tensor, TensorData}; + +#[test] +fn should_support_stack_ops_int() { + let device = Default::default(); + let tensor_1 = TestTensorInt::<2>::from_data([[1, 2, 3]], &device); + let tensor_2 = TestTensorInt::<2>::from_data([[4, 5, 6]], &device); + + let output = Tensor::stack::<3>(vec![tensor_1, tensor_2], 0); + let expected = TensorData::from([[[1, 2, 3]], [[4, 5, 6]]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_generate_row_major_layout() { + let device = Default::default(); + let tensor = TestTensorInt::<1>::arange(1..25, &device).reshape([4, 6]); + let zeros = TestTensorInt::zeros([4, 6], &device); + let intersperse = + Tensor::stack::<3>([tensor.clone(), zeros.clone()].to_vec(), 2).reshape([4, 12]); + + let expected = TensorData::from([ + [1, 0, 2, 0, 3, 0, 4, 0, 5, 0, 6, 0], + [7, 0, 8, 0, 9, 0, 10, 0, 11, 0, 12, 0], + [13, 0, 14, 0, 15, 0, 16, 0, 17, 0, 18, 0], + [19, 0, 20, 0, 21, 0, 22, 0, 23, 0, 24, 0], + ]); + + intersperse.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/sub.rs b/crates/burn-backend-tests/tests/tensor/int/ops/sub.rs new file mode 100644 index 0000000..cd644cd --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/sub.rs @@ -0,0 +1,55 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_sub_ops_int() { + let data_1 = TensorData::from([[0, 1, 2], [3, 4, 5]]); + let data_2 = TensorData::from([[6, 7, 8], [9, 10, 11]]); + let device = Default::default(); + let tensor_1 = TestTensorInt::<2>::from_data(data_1, &device); + let tensor_2 = TestTensorInt::<2>::from_data(data_2, &device); + + let output = tensor_1 - tensor_2; + let expected = TensorData::from([[-6, -6, -6], [-6, -6, -6]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_sub_broadcast_int() { + let data_1 = TensorData::from([[0, 1, 2]]); + let data_2 = TensorData::from([[3, 4, 5], [6, 7, 8]]); + let device = Default::default(); + let tensor_1 = TestTensorInt::<2>::from_data(data_1, &device); + let tensor_2 = TestTensorInt::<2>::from_data(data_2, &device); + + let output = tensor_1 - tensor_2; + let expected = TensorData::from([[-3, -3, -3], [-6, -6, -6]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_sub_scalar_ops_int() { + let data = TensorData::from([[0, 1, 2], [3, 4, 5]]); + let scalar = 2; + let tensor = TestTensorInt::<2>::from_data(data, &Default::default()); + + let output = tensor - scalar; + let expected = TensorData::from([[-2, -1, 0], [1, 2, 3]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_int_sub_flipped() { + // [10, 20, 30, 40] flipped - [1, 2, 3, 4] = [39, 28, 17, 6] + let a = TestTensorInt::<1>::from([10, 20, 30, 40]).flip([0]); + let b = TestTensorInt::<1>::from([1, 2, 3, 4]); + + let output = a - b; + + output + .into_data() + .assert_eq(&TensorData::from([39, 28, 17, 6]), false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/take.rs b/crates/burn-backend-tests/tests/tensor/int/ops/take.rs new file mode 100644 index 0000000..d9c28ef --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/take.rs @@ -0,0 +1,31 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_take_int_tensor() { + // Test take with integer tensors + let device = Default::default(); + let tensor = TestTensorInt::<2>::from_data([[10, 20, 30], [40, 50, 60]], &device); + let indices = TestTensorInt::<1>::from_data([1, 0], &device); + + let output = tensor.take::<1, 2>(0, indices); + let expected = TensorData::from([[40, 50, 60], [10, 20, 30]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_take_int_tensor_with_2d_indices() { + // Test take with integer tensors - output will be 3D + let device = Default::default(); + let tensor = TestTensorInt::<2>::from_data([[10, 20, 30], [40, 50, 60], [70, 80, 90]], &device); + + // 2D indices - shape [2, 2] + let indices = TestTensorInt::<2>::from_data([[0, 2], [2, 1]], &device); + let output = tensor.take::<2, 3>(0, indices); + + // Expected: shape [2, 2, 3] + let expected = TensorData::from([[[10, 20, 30], [70, 80, 90]], [[70, 80, 90], [40, 50, 60]]]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/topk.rs b/crates/burn-backend-tests/tests/tensor/int/ops/topk.rs new file mode 100644 index 0000000..16b6367 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/topk.rs @@ -0,0 +1,58 @@ +use super::*; +use burn_tensor::{TensorData, Tolerance}; + +#[test] +fn test_topk_with_indices_1d() { + let tensor = TestTensorInt::<1>::from([1, 2, 3, 4, 5]); + + let (values, indices) = tensor.topk_with_indices(3, /*dim*/ 0); + + let values_expected = TensorData::from([5, 4, 3]); + values.into_data().assert_eq(&values_expected, false); + + let indices_expected = TensorData::from([4, 3, 2]); + indices.into_data().assert_eq(&indices_expected, false); +} + +#[test] +fn test_topk_1d() { + // Int + let tensor = TestTensorInt::<1>::from([1, 2, 3, 4, 5]); + + let values = tensor.topk(3, /*dim*/ 0); + let expected = TensorData::from([5, 4, 3]); + + values.into_data().assert_eq(&expected, false); + + // Float + let tensor = TestTensor::<1>::from([1., 2., 3., 4., 5.]); + + let values = tensor.topk(3, /*dim*/ 0); + let expected = TensorData::from([5., 4., 3.]); + + values + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} + +#[test] +fn test_topk() { + // 3D Int + let tensor = TestTensorInt::<3>::from([[[1, 4, 7], [2, 5, 6]], [[3, 0, 9], [8, 2, 8]]]); + + let values = tensor.topk(2, /*dim*/ 2); + let expected = TensorData::from([[[7, 4], [6, 5]], [[9, 3], [8, 8]]]); + + values.into_data().assert_eq(&expected, false); + + // 3D Float + let tensor = + TestTensor::<3>::from([[[1., 4., 7.], [2., 5., 6.]], [[3., 0., 9.], [8., 2., 8.]]]); + + let values = tensor.topk(2, /*dim*/ 2); + let expected = TensorData::from([[[7., 4.], [6., 5.]], [[9., 3.], [8., 8.]]]); + + values + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/transpose.rs b/crates/burn-backend-tests/tests/tensor/int/ops/transpose.rs new file mode 100644 index 0000000..7de0b0d --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/transpose.rs @@ -0,0 +1,28 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn should_support_transpose_ops_int() { + let tensor = TestTensorInt::<3>::from_data( + [[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]], + &Default::default(), + ); + + let output = tensor.transpose(); + let expected = TensorData::from([[[0, 3], [1, 4], [2, 5]], [[6, 9], [7, 10], [8, 11]]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn should_support_swap_dims_int() { + let tensor = TestTensorInt::<3>::from_data( + [[[0, 1, 2], [3, 4, 5]], [[6, 7, 8], [9, 10, 11]]], + &Default::default(), + ); + + let output = tensor.swap_dims(0, 2); + let expected = TensorData::from([[[0, 6], [3, 9]], [[1, 7], [4, 10]], [[2, 8], [5, 11]]]); + + output.into_data().assert_eq(&expected, false); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/tri.rs b/crates/burn-backend-tests/tests/tensor/int/ops/tri.rs new file mode 100644 index 0000000..c62de88 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/tri.rs @@ -0,0 +1,85 @@ +use super::*; +use burn_tensor::TensorData; + +#[test] +fn test_triu_negative_diagonal() { + let tensor = TestTensorInt::<2>::from([[1, 1, 1], [1, 1, 1], [1, 1, 1]]); + + let output = tensor.triu(-1); + let expected = TensorData::from([[1, 1, 1], [1, 1, 1], [0, 1, 1]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_triu_batch_tensors() { + let tensor = TestTensorInt::<4>::from([ + [[[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]]], + [[[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]]], + ]); + let output = tensor.triu(1); + let expected = TensorData::from([ + [[[0, 1, 1, 1], [0, 0, 1, 1], [0, 0, 0, 1], [0, 0, 0, 0]]], + [[[0, 1, 1, 1], [0, 0, 1, 1], [0, 0, 0, 1], [0, 0, 0, 0]]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +#[should_panic] +fn test_triu_too_few_dims() { + let tensor = TestTensorInt::<1>::from([1, 2, 3]); + let _output = tensor.triu(0); +} + +#[test] +fn test_tril() { + let tensor = TestTensor::<2>::from([[1., 1., 1.], [1., 1., 1.], [1., 1., 1.]]); + let output = tensor.tril(0); + let expected = TensorData::from([[1., 0., 0.], [1., 1., 0.], [1., 1., 1.]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_tril_positive_diagonal() { + let tensor = TestTensorInt::<2>::from([[1, 1, 1], [1, 1, 1], [1, 1, 1]]); + + let output = tensor.tril(1); + let expected = TensorData::from([[1, 1, 0], [1, 1, 1], [1, 1, 1]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_tril_negative_diagonal() { + let tensor = TestTensorInt::<2>::from([[1, 1, 1], [1, 1, 1], [1, 1, 1]]); + + let output = tensor.tril(-1); + let expected = TensorData::from([[0, 0, 0], [1, 0, 0], [1, 1, 0]]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +fn test_tril_batch_tensors() { + let tensor = TestTensorInt::<4>::from([ + [[[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]]], + [[[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]]], + ]); + let output = tensor.tril(1); + let expected = TensorData::from([ + [[[1, 1, 0, 0], [1, 1, 1, 0], [1, 1, 1, 1], [1, 1, 1, 1]]], + [[[1, 1, 0, 0], [1, 1, 1, 0], [1, 1, 1, 1], [1, 1, 1, 1]]], + ]); + + output.into_data().assert_eq(&expected, false); +} + +#[test] +#[should_panic] +fn test_tril_too_few_dims() { + let tensor = TestTensorInt::<1>::from([1, 2, 3]); + let _output = tensor.tril(0); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/ops/unfold.rs b/crates/burn-backend-tests/tests/tensor/int/ops/unfold.rs new file mode 100644 index 0000000..779d37e --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/ops/unfold.rs @@ -0,0 +1,39 @@ +use super::*; +use burn_tensor::Distribution; +use burn_tensor::s; + +#[test] +fn test_unfold_int() { + // Distribution::Default samples from [0, 255) + if (IntElem::MAX as u32) < 255 - 1 { + return; + } + let device = Default::default(); + + let input = TestTensorInt::<3>::random([2, 6, 6], Distribution::Default, &device); + + let dim = 1; + let size = 3; + let step = 2; + let actual: TestTensorInt<4> = input.clone().unfold(dim, size, step); + + let expected = TestTensorInt::<4>::empty([2, 2, 6, 3], &device) + .slice_assign( + s![.., 0, .., ..], + input + .clone() + .slice(s![.., 0..3, ..]) + .swap_dims(1, 2) + .unsqueeze_dim::<4>(1), + ) + .slice_assign( + s![.., 1, .., ..], + input + .clone() + .slice(s![.., 2..5, ..]) + .swap_dims(1, 2) + .unsqueeze_dim::<4>(1), + ); + + actual.to_data().assert_eq(&expected.to_data(), true); +} diff --git a/crates/burn-backend-tests/tests/tensor/int/primitive.rs b/crates/burn-backend-tests/tests/tensor/int/primitive.rs new file mode 100644 index 0000000..99ce4f8 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/int/primitive.rs @@ -0,0 +1,13 @@ +use super::*; +use burn_tensor::{Element, Shape}; + +#[test] +fn should_support_int_dtype() { + let tensor = TestTensorInt::<2>::from([[0, -1, 2], [3, 4, -5]])/*.into_primitive()*/; + + assert_eq!(tensor.shape(), Shape::new([2, 3])); + assert_eq!( + tensor.dtype(), + IntElem::dtype() // default int elem type + ); +} diff --git a/crates/burn-backend-tests/tests/tensor/mod.rs b/crates/burn-backend-tests/tests/tensor/mod.rs new file mode 100644 index 0000000..7cb2ee0 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/mod.rs @@ -0,0 +1,12 @@ +pub use super::*; // re-export test types + +mod clone_invariance; +#[cfg(feature = "distributed")] +mod distributed; +#[cfg(feature = "std")] +mod multi_threads; + +// Data types +mod bool; +mod float; +mod int; diff --git a/crates/burn-backend-tests/tests/tensor/multi_threads.rs b/crates/burn-backend-tests/tests/tensor/multi_threads.rs new file mode 100644 index 0000000..8710947 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor/multi_threads.rs @@ -0,0 +1,171 @@ +use super::*; +use core::time::Duration; +use std::sync::{ + Arc, + atomic::{AtomicU32, Ordering}, +}; + +struct MultiThreadTestSettings { + num_threads: usize, + // The number of operations that are applied while the tensor is still alive and has a + // reference count > 1 on the new thread. + num_ops_alive: usize, + // The number of operations that are applied after the tensor is consumed for the last time. + num_ops_consumed: usize, + // Number of operations that needs to execute before continuing execution on the main thread. + sleep_before: Duration, + sleep_alive: Duration, + sleep_consumed: Duration, + // If the output is dropped, otherwise it will be consumed by an operation. + dropped: bool, +} + +#[test] +fn should_handle_multi_threads_dropped() { + run_multi_thread_test(MultiThreadTestSettings { + num_threads: 3, + num_ops_alive: 5, + num_ops_consumed: 5, + sleep_before: Duration::from_millis(100), + sleep_alive: Duration::from_millis(100), + sleep_consumed: Duration::from_millis(100), + dropped: true, + }) +} + +#[test] +fn should_handle_multi_threads_consumed() { + run_multi_thread_test(MultiThreadTestSettings { + num_threads: 3, + num_ops_alive: 5, + num_ops_consumed: 5, + sleep_before: Duration::from_millis(100), + sleep_alive: Duration::from_millis(100), + sleep_consumed: Duration::from_millis(100), + dropped: false, + }) +} + +#[test] +fn should_handle_multi_threads_drop_no_wait() { + run_multi_thread_test(MultiThreadTestSettings { + num_threads: 3, + num_ops_alive: 5, + num_ops_consumed: 5, + sleep_before: Duration::from_millis(100), + sleep_alive: Duration::from_millis(100), + sleep_consumed: Duration::from_millis(100), + dropped: true, + }) +} + +#[test] +fn should_handle_multi_threads_consumed_no_wait() { + run_multi_thread_test(MultiThreadTestSettings { + num_threads: 3, + num_ops_alive: 5, + num_ops_consumed: 5, + sleep_before: Duration::from_millis(100), + sleep_alive: Duration::from_millis(100), + sleep_consumed: Duration::from_millis(100), + dropped: false, + }) +} + +#[test] +fn should_handle_multi_threads_no_async_op() { + run_multi_thread_test(MultiThreadTestSettings { + num_threads: 3, + num_ops_alive: 0, + num_ops_consumed: 0, + sleep_before: Duration::from_millis(100), + sleep_alive: Duration::from_millis(100), + sleep_consumed: Duration::from_millis(100), + dropped: false, + }) +} + +// Skip test - flaky (works when ran alone) +// Enable once this issue is fixed: https://github.com/tracel-ai/burn/issues/4328 +#[ignore = "Flaky test"] +#[test] +fn should_handle_multi_threads_no_async_op_no_wait() { + run_multi_thread_test(MultiThreadTestSettings { + num_threads: 3, + num_ops_alive: 0, + num_ops_consumed: 0, + sleep_before: Duration::from_millis(0), + sleep_alive: Duration::from_millis(100), + sleep_consumed: Duration::from_millis(100), + dropped: false, + }) +} + +fn run_multi_thread_test(settings: MultiThreadTestSettings) { + let tensor = TestTensor::<2>::from([[0.0, -1.0, 2.0], [3.0, 4.0, -5.0]]); + + let mut joined = Vec::with_capacity(settings.num_threads); + + let counter_alive = Arc::new(AtomicU32::new(0)); + let counter_consumed = Arc::new(AtomicU32::new(0)); + + for i in 0..settings.num_threads { + let tensor_moved = tensor.clone(); + let ca_moved = counter_alive.clone(); + let cc_moved = counter_consumed.clone(); + + let handle = std::thread::spawn(move || { + let mut base = tensor_moved.clone(); + std::thread::sleep(settings.sleep_before); + + if settings.num_ops_alive == 0 && settings.num_ops_consumed == 0 { + core::mem::drop(tensor_moved); + core::mem::drop(base); + } else { + if settings.num_ops_alive > 1 { + for j in 0..(settings.num_ops_alive - 1) { + base = tensor_moved.clone() + j as u32; + ca_moved.fetch_add(1, Ordering::Relaxed); + std::thread::sleep(settings.sleep_alive); + } + } + + base = base * tensor_moved + i as u32; + ca_moved.fetch_add(1, Ordering::Relaxed); + + for n in 0..settings.num_ops_consumed { + base = base + n as i32; + cc_moved.fetch_add(1, Ordering::Relaxed); + std::thread::sleep(settings.sleep_consumed); + } + let _data = base.into_data(); + } + }); + joined.push(handle); + } + + fn wait(counter: Arc, limit: usize) { + loop { + let counter_curr = counter.load(Ordering::Relaxed); + if counter_curr as usize >= limit { + break; + } else { + std::thread::sleep(Duration::from_millis(10)); + } + } + } + + wait(counter_alive, settings.num_ops_alive); + wait(counter_consumed, settings.num_ops_consumed); + + if settings.dropped { + core::mem::drop(tensor); + } else { + let t = tensor * 2.0; + let _t = t.into_data(); + } + + for j in joined { + j.join().unwrap(); + } +} diff --git a/crates/burn-backend-tests/tests/tensor_f16.rs b/crates/burn-backend-tests/tests/tensor_f16.rs new file mode 100644 index 0000000..1d3d4e1 --- /dev/null +++ b/crates/burn-backend-tests/tests/tensor_f16.rs @@ -0,0 +1,29 @@ +// //! Burn backend tensor tests. + +#![recursion_limit = "256"] +#![cfg(any( + feature = "vulkan", + feature = "cuda", + feature = "rocm", + feature = "metal", + feature = "flex" +))] + +extern crate alloc; + +pub type FloatElem = burn_tensor::f16; +#[allow(unused)] +pub type IntElem = i32; + +#[path = "common/backend.rs"] +mod backend; +pub use backend::*; + +#[path = "tensor/float/mod.rs"] +mod f16; + +#[cfg(feature = "fusion")] +#[path = "fusion/mod.rs"] +mod fusion; + +// TODO: bf16 (vulkan only supports bf16 for matmul, metal/wgpu doesn't support bf16) diff --git a/crates/burn-backend/Cargo.toml b/crates/burn-backend/Cargo.toml new file mode 100644 index 0000000..b474036 --- /dev/null +++ b/crates/burn-backend/Cargo.toml @@ -0,0 +1,58 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science", "no-std", "embedded", "wasm"] +description = "Core backend interfaces and data structures for executing tensor operations in Burn." +documentation = "https://docs.rs/burn-backend" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "tensor", "pytorch", "ndarray"] +license.workspace = true +name = "burn-backend" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-backend" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["std"] +doc = ["default"] +std = ["rand/std", "num-traits/std", "burn-std/std", "cubecl?/std", "dep:oneshot"] + +tracing = ["burn-std/tracing", "cubecl/tracing"] + +# For DTypeUsage de/serialization +serde = ["enumset/serde"] + +cubecl = ["dep:cubecl"] +cubecl-cuda = ["cubecl", "cubecl/cuda"] +cubecl-hip = ["cubecl", "cubecl/hip"] +cubecl-wgpu = ["cubecl", "cubecl/wgpu"] +cubecl-metal = ["cubecl", "cubecl/metal"] +cubecl-vulkan = ["cubecl", "cubecl/vulkan"] +cubecl-webgpu = ["cubecl", "cubecl/webgpu"] +cubecl-cpu = ["cubecl", "cubecl/cpu"] + +[dependencies] +burn-std = { workspace = true } +cubecl = { workspace = true, optional = true, default-features = false } + +bytemuck = { workspace = true, features = ["extern_crate_alloc"] } +derive-new = { workspace = true } +enumset = { workspace = true } +hashbrown = { workspace = true } +num-traits = { workspace = true } +rand = { workspace = true, default-features = false } +rand_distr = { workspace = true } +serde = { workspace = true } +oneshot = { version = "0.2.1", optional = true } +spin = { workspace = true } + +[target.'cfg(not(target_has_atomic = "ptr"))'.dependencies] +portable-atomic-util = { workspace = true } + +[dev-dependencies] +rand = { workspace = true, features = ["thread_rng"] } +paste = { workspace = true } +serde_json = { workspace = true, features = ["alloc"] } +serial_test = { workspace = true } diff --git a/crates/burn-backend/README.md b/crates/burn-backend/README.md new file mode 100644 index 0000000..e79c1b5 --- /dev/null +++ b/crates/burn-backend/README.md @@ -0,0 +1,4 @@ +# Burn Backend + +This crate includes the core backend interfaces and data structures for executing tensor operations +in Burn. diff --git a/crates/burn-backend/src/alias.rs b/crates/burn-backend/src/alias.rs new file mode 100644 index 0000000..54223f7 --- /dev/null +++ b/crates/burn-backend/src/alias.rs @@ -0,0 +1,20 @@ +/// Backend type aliases. +pub mod tensor { + use crate::backend::BackendTypes; + // We provide some type aliases to improve the readability of using associated types without + // having to use the disambiguation syntax. + + pub use burn_std::ops::IndexingUpdateOp; + + /// Device type used by the backend. + pub type Device = ::Device; + + /// Float tensor primitive type used by the backend. + pub type FloatTensor = ::FloatTensorPrimitive; + /// Integer tensor primitive type used by the backend. + pub type IntTensor = ::IntTensorPrimitive; + /// Boolean tensor primitive type used by the backend. + pub type BoolTensor = ::BoolTensorPrimitive; + /// Quantized tensor primitive type used by the backend. + pub type QuantizedTensor = ::QuantizedTensorPrimitive; +} diff --git a/crates/burn-backend/src/backend/base.rs b/crates/burn-backend/src/backend/base.rs new file mode 100644 index 0000000..c3feb4a --- /dev/null +++ b/crates/burn-backend/src/backend/base.rs @@ -0,0 +1,469 @@ +use burn_std::DType; +pub use burn_std::{ExecutionError, backtrace::BackTrace}; + +use crate::distributed::DistributedOps; +pub use crate::element::Element; +use crate::ops::*; +use crate::tensor::{BoolTensor, FloatTensor, IntTensor, QuantizedTensor}; +use crate::{TensorData, TensorMetadata}; +use alloc::string::String; +use enumset::{EnumSet, EnumSetType}; + +use crate::distributed::{DistributedParamId, DistributedParams}; + +use super::DeviceOps; + +/// The mapping of types used by Backend and traits. +pub trait BackendTypes: Clone + Send + Sync + core::fmt::Debug + 'static { + /// Device type. + type Device: DeviceOps; + + /// Tensor primitive to be used for all float operations. + type FloatTensorPrimitive: TensorMetadata + 'static; + + /// Tensor primitive to be used for all int operations. + type IntTensorPrimitive: TensorMetadata + 'static; + + /// Tensor primitive to be used for all bool operations. + type BoolTensorPrimitive: TensorMetadata + 'static; + + /// Tensor primitive to be used for all quantized operations. + type QuantizedTensorPrimitive: TensorMetadata + 'static; + + /// Captured graph primitive returned by [`Backend::graph_stop_capture`] and + /// consumed by [`Backend::graph_replay`]: a backend-owned recording of a + /// launch sequence that replays as a single dispatch. + /// + /// Backends without graph-capture support use [`GraphUnsupported`], an + /// uninhabited type — their capture methods only ever error, so no value of + /// it can exist. + type GraphPrimitive: Clone + Send + Sync + core::fmt::Debug + 'static; +} + +/// Captured graph primitive type used by the backend (see +/// [`BackendTypes::GraphPrimitive`]). +pub type BackendGraph = ::GraphPrimitive; + +/// Placeholder [graph primitive](BackendTypes::GraphPrimitive) for backends +/// without graph-capture support. +/// +/// Uninhabited: `graph_stop_capture` on such backends always errors, so a value +/// of this type can never be constructed (and `graph_replay` can never be called). +#[derive(Debug, Clone, Copy)] +pub enum GraphUnsupported {} + +/// The error returned by the default (unsupported) graph-capture methods. +fn graph_unsupported() -> ExecutionError { + ExecutionError::Generic { + reason: alloc::string::String::from("graph capture is not supported by this backend"), + backtrace: BackTrace::capture(), + } +} + +/// This trait defines all types and functions needed for a backend to be used with burn. +/// +/// ## Design +/// +/// This trait aims to be as unopinionated as possible and allows implementations to define +/// their own types and patterns. Therefore, there are few pre-defined abstractions baked +/// into this trait. +/// +/// Backends must define their own tensor types for each data type: `float`, `int`, and `bool`. +/// Since we minimize assumptions, we chose to separate these types, as they are used in +/// different contexts. However, some backends may have a generic tensor type that is used +/// for all data types. +/// +/// ### Eager Mode +/// +/// Because burn supports dynamic graphs, the backend trait is designed around kernel +/// implementations that can be called without any mutable context or graph. This may not be +/// ideal for backends that want to configure their computational graphs and execute them +/// multiple times. +/// +/// To implement this kind of backend, channels could be used to communicate with a backend +/// server thread to build the computation graphs and re-execute the ones that are repeated, +/// with some form of cache. Once that pattern has matured, a graph mode backend trait could +/// be extracted from it, allowing other backends of the same kind to be quickly integrated +/// with burn. This pattern could also be used to create an operation fusion trait, which +/// allows backends to define what kind of graph structures can be fused into one operation. +/// +/// ### Multi-Threaded +/// +/// Backend tensor types are all `Clone` + `Send`, which allows them to be safely +/// sent between threads. It is recommended to wrap tensors with [Arc](alloc::sync::Arc), +/// which avoids copying the tensor's buffer. Note that it is still possible to mutate and +/// reuse tensors' buffer without locking; see the next section on the Mutable API. +/// +/// ### Mutable API +/// +/// There is no mutable or inplace operation API to implement, but that does not mean that +/// backends cannot support them. Using [try_unwrap](alloc::sync::Arc::try_unwrap) and +/// [get_mut](alloc::sync::Arc::get_mut) allows backends to have access to an owned or mutable +/// reference to their tensor buffer data structure if the tensor is not shared. In that case, +/// backends can dispatch to their owned inplace operations for better performance. +/// +/// ## Documentation +/// +/// Most of the documentation for each function can be found on the user API +#[cfg_attr(doc, doc = crate::doc_tensor!())] +#[cfg_attr(not(doc), doc = "`Tensor`")] +/// struct in the `burn-tensor` crate. +/// For modules, public functions are often created, which can be used by `burn-core` modules. +pub trait Backend: + BackendTypes + + FloatTensorOps + + BoolTensorOps + + IntTensorOps + + ModuleOps + + ActivationOps + + QTensorOps + + TransactionOps + + DistributedOps + + Clone + + Default + + Sized + + Send + + Sync + + core::fmt::Debug + + 'static +{ + /// If autodiff is enabled. + fn ad_enabled(_device: &Self::Device) -> bool { + false + } + + /// Sets the current allocation mode to persistent. + #[allow(unused_variables)] + fn memory_persistent_allocations< + Output: Send, + Input: Send, + Func: Fn(Input) -> Output + Send, + >( + device: &Self::Device, + input: Input, + func: Func, + ) -> Output { + func(input) + } + + /// Manually triggers a memory cleanup on the given device. + #[allow(unused_variables)] + fn memory_cleanup(device: &Self::Device) {} + + /// Name of the backend. + fn name(device: &Self::Device) -> String; + + /// Seeds the backend on the specified device. + /// + /// There is no guarantee that only the specified device will be seeded, but it is guaranteed + /// that at least the specified device will be seeded. + /// + /// In all cases, this should ensure deterministic execution for a single-threaded program. + fn seed(device: &Self::Device, seed: u64); + + /// Sync the backend, ensure that all computation are finished. + fn sync(_device: &Self::Device) -> Result<(), ExecutionError> { + Ok(()) + } + + /// Prepare `device` for an upcoming graph capture: route allocations into a + /// stable pool so every buffer allocated before graph_stop_capture can + /// be pinned. Call before the warmup run. No-op by default. + /// + /// See [`burn_graph`](crate) — the closure-based `capture` helper drives + /// this whole sequence. + fn graph_prepare(_device: &Self::Device) -> Result<(), ExecutionError> { + Ok(()) + } + + /// Begin recording launches on `device` into a graph (see + /// [`graph_stop_capture`](Backend::graph_stop_capture)). Errors on backends + /// without hardware graph support, so callers fall back to re-running. + fn graph_start_capture(_device: &Self::Device) -> Result<(), ExecutionError> { + Err(graph_unsupported()) + } + + /// Stop recording and return the captured [graph](BackendTypes::GraphPrimitive), + /// ready to [`graph_replay`](Backend::graph_replay). + fn graph_stop_capture(_device: &Self::Device) -> Result, ExecutionError> { + Err(graph_unsupported()) + } + + /// Replay a captured [graph](BackendTypes::GraphPrimitive) — one dispatch + /// re-running the recorded launches against their original buffers. + /// + /// # Safety + /// + /// The replay dispatches raw device work against the exact buffers recorded + /// at capture time, with nothing tracking whether those buffers are still + /// valid. The caller must guarantee, for every tensor the captured closure + /// read or wrote: + /// + /// - its buffer is still alive — no tensor referenced by the graph has been + /// freed (and its memory possibly reallocated) since capture; + /// - it is not concurrently read or written by work on another stream or + /// thread while the replay executes; + /// - input refreshes and output reads are issued on the stream the graph + /// was captured on, so they order correctly against the replay. + unsafe fn graph_replay( + _device: &Self::Device, + _graph: &BackendGraph, + ) -> Result<(), ExecutionError> { + Err(graph_unsupported()) + } + + /// Flush any pending operation of the backend. + fn flush(_device: &Self::Device); + + /// Marks the given data as being used as a staging buffer for transfer between CPU and + /// accelerators like GPUs. + /// + /// The given data might be transferred to pinned memory or another format to improve data transfer + /// speed. + fn staging<'a, Iter>(_data: Iter, _device: &Self::Device) + where + Iter: Iterator, + { + } + + /// Whether the type is fully supported by the specified device for general operations. + /// + /// A type is considered supported if it can be used for the full suite of tensor + /// operations, including storage, conversion, and basic arithmetic. + /// + /// Returning `false` does not necessarily mean the device cannot handle the type at all. + /// For instance, a device might support a type only for specialized hardware + /// acceleration (e.g., matrix multiplication) but lack general arithmetic support. Such + /// types should return `false` here as they are not globally supported. + fn supports_dtype(device: &Self::Device, dtype: DType) -> bool { + Self::dtype_usage(device, dtype).is_superset(DTypeUsage::general()) + } + + /// Returns the [DTypeUsageSet] for the given [DType] on the specified device. + fn dtype_usage(device: &Self::Device, dtype: DType) -> DTypeUsageSet; + + /// Returns the number of devices available on this backend. + /// `device` is a reference device used to determine the underlying backend that should be queried. + /// A CUDA device will return all devices available to CUDA, a Vulkan device will return all + /// devices available to Vulkan, etc. + fn device_count(type_id: u16) -> usize; +} + +/// Trait that allows a backend to support autodiff. +pub trait AutodiffBackend: Backend { + /// The inner backend type. + type InnerBackend: Backend; + + /// Gradients type. + type Gradients: Send; + + /// Backward pass. + /// + /// # Arguments + /// + /// * `tensor` - The tensor is the last node of computational graph where the gradients are computed. + /// + /// # Returns + /// + /// The gradients. + fn backward(tensor: FloatTensor) -> Self::Gradients; + + /// Returns the gradients of a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to extract the gradients from. + /// + /// # Returns + /// + /// An optional tensor containing the gradient. + fn grad( + tensor: &FloatTensor, + grads: &Self::Gradients, + ) -> Option>; + + /// Pops the gradients of a tensor and returns them. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to pop the gradients from. + /// * `grads` - The gradients. + /// + /// # Returns + /// + /// An optional tensor containing the given gradients. + fn grad_remove( + tensor: &FloatTensor, + grads: &mut Self::Gradients, + ) -> Option>; + + /// Replace the gradients of a tensor with the one provided. + /// + /// If no gradient existed for the provided tensor, register it. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to pop the gradients from. + /// * `grads` - The gradients. + /// * `grad` - The updated grad tensor. + fn grad_replace( + tensor: &FloatTensor, + grads: &mut Self::Gradients, + grad: FloatTensor, + ); + + /// Returns the tensor with inner backend type. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the inner backend tensor for. + /// + /// # Returns + /// + /// The inner backend tensor. + fn inner(tensor: FloatTensor) -> FloatTensor; + + /// Returns the tensor with inner backend type. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the inner backend tensor for. + /// + /// # Returns + /// + /// The inner backend tensor. + fn int_inner(tensor: IntTensor) -> IntTensor; + + /// Returns the tensor with inner backend type. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the inner backend tensor for. + /// + /// # Returns + /// + /// The inner backend tensor. + fn bool_inner(tensor: BoolTensor) -> BoolTensor; + + /// Returns the tensor with inner backend type. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the inner backend tensor for. + /// + /// # Returns + /// + /// The inner backend tensor. + fn q_inner(tensor: QuantizedTensor) -> QuantizedTensor; + + /// Converts the inner backend tensor to the autodiff backend tensor. + /// + /// # Arguments + /// + /// * `tensor` - The inner backend tensor to convert. + /// + /// + /// # Returns + /// + /// The autodiff backend tensor. + fn from_inner(tensor: FloatTensor) -> FloatTensor; + + /// Converts the inner backend tensor to the autodiff backend tensor. + /// + /// # Arguments + /// + /// * `tensor` - The inner backend tensor to convert. + /// + /// + /// # Returns + /// + /// The autodiff backend tensor. + fn int_from_inner(tensor: IntTensor) -> IntTensor; + + /// Converts the inner backend tensor to the autodiff backend tensor. + /// + /// # Arguments + /// + /// * `tensor` - The inner backend tensor to convert. + /// + /// + /// # Returns + /// + /// The autodiff backend tensor. + fn bool_from_inner(tensor: BoolTensor) -> BoolTensor; + + /// Converts the inner backend tensor to the autodiff backend tensor. + /// + /// # Arguments + /// + /// * `tensor` - The inner backend tensor to convert. + /// + /// + /// # Returns + /// + /// The autodiff backend tensor. + fn q_from_inner(tensor: QuantizedTensor) -> QuantizedTensor; + + /// Mark the tensor as distributed across multiple devices. + /// The gradients will be aggregated during the backward pass. + /// + /// This function does nothing when distributed training is not available. + fn set_distributed_params( + tensor: FloatTensor, + _param_id: DistributedParamId, + ) -> FloatTensor { + tensor + } + + /// Returns the distributed parameters if the tensor was marked as distributed. + fn distributed_params(_tensor: &FloatTensor) -> Option { + None + } + + /// Returns true if the tensor was marked as distributed. + fn is_distributed(_tensor: &FloatTensor) -> bool { + false + } +} + +/// Describes how a data type can be used on a given device. +/// +/// A data type may be supported for different classes of operations. Not all +/// data types that appear in hardware or kernel implementations are suitable +/// for general-purpose tensor operations. +#[derive(Debug, EnumSetType)] +pub enum DTypeUsage { + /// The type can be stored in device memory and converted to and from + /// other supported data types. + Storage, + /// The type supports general-purpose arithmetic and common tensor + /// operations (e.g. elementwise ops, reductions, etc.). + Arithmetic, + /// The type is supported by hardware-accelerated execution paths. + /// + /// This typically indicates support for accelerator-backed compute units (e.g., tensor + /// cores executing MMA instructions) for high-performance operations such as matrix + /// multiplication and operations that lower to it. + /// + /// # Notes + /// - A type can be both [`Arithmetic`](DTypeUsage::Arithmetic) and + /// [`Accelerated`](DTypeUsage::Accelerated) if it supports general-purpose operations + /// *and* accelerated paths. + /// - If a type is marked as `Accelerated` but not `Arithmetic`, it is not + /// suitable for general-purpose tensor operations and may only be used + /// in specific accelerated operations. + /// + /// `Accelerated` is a **flag**, not a detailed descriptor. It does not enumerate which + /// operations are accelerated or which accelerator features are available. + Accelerated, +} + +/// A set of [DTypeUsage] representing the total capabilities of a data type on a device. +pub type DTypeUsageSet = EnumSet; + +impl DTypeUsage { + /// Returns the usage set required for general-purpose tensor support. + pub fn general() -> DTypeUsageSet { + DTypeUsage::Storage | DTypeUsage::Arithmetic + } +} diff --git a/crates/burn-backend/src/backend/device.rs b/crates/burn-backend/src/backend/device.rs new file mode 100644 index 0000000..ec96807 --- /dev/null +++ b/crates/burn-backend/src/backend/device.rs @@ -0,0 +1,403 @@ +pub use burn_std::device::*; +use burn_std::{BoolDType, DType, FloatDType, IntDType}; +pub use burn_std::{DeviceError, DeviceSettings}; + +use burn_std::stub::RwLock; + +#[cfg(target_has_atomic = "ptr")] +use alloc::sync::Arc; + +#[cfg(not(target_has_atomic = "ptr"))] +use portable_atomic_util::Arc; + +use core::any::TypeId; + +#[cfg(feature = "std")] +pub use std::collections::HashMap; +#[cfg(feature = "std")] +use std::sync::{LazyLock, OnceLock}; + +#[cfg(not(feature = "std"))] +pub use hashbrown::HashMap; +#[cfg(not(feature = "std"))] +use spin::{Lazy as LazyLock, Once as OnceLock}; + +use crate::Backend; + +/// Device trait for all burn backend devices. +pub trait DeviceOps: Clone + Default + PartialEq + Send + Sync + core::fmt::Debug + Device { + /// Returns the [device id](DeviceId). + fn id(&self) -> DeviceId { + self.to_id() + } + + /// Returns the default [settings](DeviceSettings) for this device. + fn defaults(&self) -> DeviceSettings; +} + +/// Key for the registry: physical device type + device id +type RegistryKey = (DeviceId, TypeId); + +/// Global registry mapping devices to their settings. +/// +/// Each value is wrapped in a `OnceLock` to enforce that settings are initialized only once +/// per device. +static REGISTRY: LazyLock>>>> = + LazyLock::new(|| RwLock::new(HashMap::new())); + +struct DeviceSettingsRegistry; + +impl DeviceSettingsRegistry { + /// Returns the settings for the given device, inserting the default if absent. + fn get_or_insert( + device: &D, + default_fn: impl FnOnce() -> DeviceSettings, + ) -> DeviceSettings { + let key = Self::key(device); + #[cfg(feature = "std")] + { + let cached = LOCAL_CACHE.with(|cache| cache.borrow().get(&key).copied()); + if let Some(settings) = cached { + return settings; + } + + // Entry does not exist in cache + let settings = { + let read = REGISTRY.read().unwrap(); + read.get(&key).cloned() + } + .unwrap_or_else(|| { + let mut map = REGISTRY.write().unwrap(); + Arc::clone(map.entry(key).or_default()) + }); + + let settings = *settings.get_or_init(default_fn); + + LOCAL_CACHE.with(|cache| { + cache.borrow_mut().insert(key, settings); + }); + + settings + } + #[cfg(not(feature = "std"))] + { + let settings = { + let read = REGISTRY.read().unwrap(); + read.get(&key).cloned() + } + .unwrap_or_else(|| { + let mut map = REGISTRY.write().unwrap(); + Arc::clone(map.entry(key).or_default()) + }); + + settings.call_once(default_fn); + *settings.get().unwrap() + } + } + + /// Initializes the settings for the given device. + /// + /// Returns `Err` with the existing settings if already initialized. + fn init(device: &D, settings: DeviceSettings) -> Result<(), DeviceError> { + let key = Self::key(device); + let mut map = REGISTRY.write().unwrap(); + let cell = map.entry(key).or_insert_with(|| Arc::new(OnceLock::new())); + + #[cfg(feature = "std")] + return cell + .set(settings) + .map_err(|_| DeviceError::already_initialized(device)); + + #[cfg(not(feature = "std"))] + if cell.get().is_some() { + Err(DeviceError::already_initialized(device)) + } else { + cell.call_once(|| settings); + Ok(()) + } + } + + /// Returns the device registry key. + fn key(device: &D) -> RegistryKey { + (device.to_id(), TypeId::of::()) + } +} + +#[cfg(feature = "std")] +thread_local! { + /// Thread-local cache access to initialized device settings is lock-free. + static LOCAL_CACHE: core::cell::RefCell> = + core::cell::RefCell::new(HashMap::new()); +} + +/// Get the [`device`'s settings](DeviceSettings). +pub fn get_device_settings(device: &B::Device) -> DeviceSettings { + DeviceSettingsRegistry::get_or_insert(device, || device.defaults()) +} + +fn check_dtype_support( + device: &B::Device, + dtype: impl Into, +) -> Result<(), DeviceError> { + let dtype = dtype.into(); + // Default dtypes should have `DTypeUsage::general()`. Types restricted to specialized + // operations should not be used as default. + if B::supports_dtype(device, dtype) { + Ok(()) + } else { + Err(DeviceError::unsupported_dtype(device, dtype)) + } +} + +/// Sets the default data types for the device. +/// +/// This updates the device's default data types used for tensor creation. +/// +/// Settings can only be initialized once per device. Subsequent calls for +/// the same device return [`DeviceError::AlreadyInitialized`]. +/// +/// # Note +/// +/// Initialization must happen before any tensor creation on the device. +/// The first tensor operation will lock the device to its defaults, causing +/// any subsequent initialization attempt to return [`DeviceError::AlreadyInitialized`]. +/// +/// # Example +/// +/// ```rust, ignore +/// fn example() { +/// let device = B::Device::default(); +/// +/// // Update the device settings +/// set_default_dtypes::(&device, DType::F16, DType::I32); +/// +/// // All float tensors created after this will use F16 by default +/// let tensor = Tensor::::zeros([2, 3], &device); +/// // All int tensors created after this will use I32 default +/// let tensor = Tensor::::zeros([2, 3], &device); +/// } +/// ``` +pub fn set_default_dtypes( + device: &B::Device, + float_dtype: impl Into, + int_dtype: impl Into, + bool_dtype: impl Into, +) -> Result<(), DeviceError> { + let float_dtype = float_dtype.into(); + let int_dtype = int_dtype.into(); + let bool_dtype = bool_dtype.into(); + check_dtype_support::(device, float_dtype)?; + check_dtype_support::(device, int_dtype)?; + check_dtype_support::(device, bool_dtype)?; + + let q_config = device.defaults().quantization; + let settings = DeviceSettings::new(float_dtype, int_dtype, bool_dtype, q_config); + + initialize_unchecked(device, settings)?; + Ok(()) +} + +// Unchecked dtypes +fn initialize_unchecked( + device: &D, + settings: DeviceSettings, +) -> Result<(), DeviceError> { + DeviceSettingsRegistry::init(device, settings) +} + +#[cfg(all(test, feature = "std"))] +mod tests { + use serial_test::serial; + + use super::*; + + fn clear_registry() { + REGISTRY.write().unwrap().clear(); + } + + #[derive(Clone, Debug, Default, PartialEq, new)] + pub struct TestDeviceA { + index: u16, + } + + impl Device for TestDeviceA { + fn from_id(device_id: DeviceId) -> Self { + Self { + index: device_id.index_id, + } + } + + fn to_id(&self) -> DeviceId { + DeviceId { + type_id: 0, + index_id: self.index, + } + } + } + + impl DeviceOps for TestDeviceA { + fn defaults(&self) -> DeviceSettings { + DeviceSettings::with_dtypes(FloatDType::F32, IntDType::I32, BoolDType::Native) + } + } + + #[derive(Clone, Debug, Default, PartialEq, new)] + pub struct TestDeviceB { + index: u16, + } + + impl Device for TestDeviceB { + fn from_id(device_id: DeviceId) -> Self { + Self { + index: device_id.index_id, + } + } + + fn to_id(&self) -> DeviceId { + DeviceId { + type_id: 0, + index_id: self.index, + } + } + } + + impl DeviceOps for TestDeviceB { + fn defaults(&self) -> DeviceSettings { + DeviceSettings::with_dtypes(FloatDType::F32, IntDType::I32, BoolDType::Native) + } + } + + fn get_test_device_settings(device: &D) -> DeviceSettings { + DeviceSettingsRegistry::get_or_insert(device, || device.defaults()) + } + + #[test] + #[serial] + fn default_settings_returned_when_uninitialized() { + clear_registry(); // reset registry for each test + + let device = TestDeviceA::new(0); + + let s1 = get_test_device_settings(&device); + let s2 = get_test_device_settings(&device); + + assert_eq!(s1, s2); + assert_eq!(s1, device.defaults()); + } + + #[test] + #[serial] + fn initialized_settings_are_returned() { + clear_registry(); // reset registry for each test + + let device = TestDeviceA::new(0); + let settings = + DeviceSettings::with_dtypes(FloatDType::BF16, IntDType::I32, BoolDType::Native); + + initialize_unchecked(&device, settings).unwrap(); + let s1 = get_test_device_settings(&device); + let s2 = get_test_device_settings(&device); + + assert_eq!(s1, s2); + assert_eq!(s1, settings); + assert_eq!(s2, settings); + } + + #[test] + #[serial] + fn settings_are_device_id_specific() { + clear_registry(); // reset registry for each test + + let d1 = TestDeviceA::new(0); + let d2 = TestDeviceA::new(1); + let settings = + DeviceSettings::with_dtypes(FloatDType::F16, IntDType::I64, BoolDType::Native); + + initialize_unchecked(&d1, settings).unwrap(); + + let s1 = get_test_device_settings(&d1); + let s2 = get_test_device_settings(&d2); + + assert_ne!(s1, s2); + assert_eq!(s1, settings); + assert_eq!(s2, d2.defaults()); + } + + #[test] + #[serial] + fn settings_are_device_type_specific() { + clear_registry(); // reset registry for each test + + let d1 = TestDeviceA::new(0); + let d2 = TestDeviceB::new(0); + let settings = + DeviceSettings::with_dtypes(FloatDType::F16, IntDType::I64, BoolDType::Native); + + initialize_unchecked(&d2, settings).unwrap(); + + let s1 = get_test_device_settings(&d1); + let s2 = get_test_device_settings(&d2); + + assert_ne!(s1, s2); + assert_eq!(s1, d1.defaults()); + assert_eq!(s2, settings); + } + + #[test] + #[serial] + fn initialization_after_default_returns_error() { + clear_registry(); // reset registry for each test + + let device = TestDeviceA::new(0); + // Settings are set to default on first access, which forces consistency + let _before = get_test_device_settings(&device); + + let settings = + DeviceSettings::with_dtypes(FloatDType::BF16, IntDType::I64, BoolDType::Native); + let result = initialize_unchecked(&device, settings); + + assert!(matches!( + result, + Err(DeviceError::AlreadyInitialized { .. }) + )); + } + + #[test] + #[serial] + fn second_initialization_returns_error() { + clear_registry(); // reset registry for each test + + let device = TestDeviceA::new(0); + let settings = + DeviceSettings::with_dtypes(FloatDType::F16, IntDType::I32, BoolDType::Native); + initialize_unchecked(&device, settings).unwrap(); + + let result = initialize_unchecked(&device, device.defaults()); + assert!(matches!( + result, + Err(DeviceError::AlreadyInitialized { .. }) + )); + } + + #[cfg(feature = "std")] + #[test] + #[serial] + fn initialized_settings_are_global() { + clear_registry(); + + let device = TestDeviceA::new(0); + let settings = + DeviceSettings::with_dtypes(FloatDType::F16, IntDType::I32, BoolDType::Native); + + initialize_unchecked(&device, settings).unwrap(); + let settings_actual = get_test_device_settings(&device); + assert_eq!(settings_actual, settings); + + // The other thread will see the initialized settings + let seen_by_new_thread = + std::thread::spawn(move || get_test_device_settings(&TestDeviceA::new(0))) + .join() + .unwrap(); + assert_eq!(seen_by_new_thread, settings_actual); + } +} diff --git a/crates/burn-backend/src/backend/distributed/api.rs b/crates/burn-backend/src/backend/distributed/api.rs new file mode 100644 index 0000000..4ba0c89 --- /dev/null +++ b/crates/burn-backend/src/backend/distributed/api.rs @@ -0,0 +1,59 @@ +use std::{ + any::{Any, TypeId}, + sync::{Mutex, MutexGuard, OnceLock}, +}; + +use std::collections::HashMap; + +use crate::Backend; + +use super::{DistributedConfig, client::DistributedSyncClient}; + +/// The type-erased box type for [`DistributedSyncClient`]. +type ClientBox = Box; + +/// Global state map from [Backend] to boxed [`DistributedSyncClient`]. +static BACKEND_CLIENT_MAP: OnceLock>> = OnceLock::new(); + +// TODO: Replace TypeId with DeviceId, the index being i32::MAX, a.k.a. communication index. +/// Gets a locked mutable view of the `STATE_MAP`. +pub(crate) fn get_backend_client_map() -> MutexGuard<'static, HashMap> { + BACKEND_CLIENT_MAP + .get_or_init(Default::default) + .lock() + .unwrap() +} + +/// Get the distributed sync client for the given [Backend]. +pub fn get_distributed_sync_client() -> Option> { + let typeid = TypeId::of::(); + let state_map = get_backend_client_map(); + state_map + .get(&typeid) + .map(|val| val.downcast_ref().cloned().unwrap()) +} + +/// Remove the client form the map for the given [Backend]. +pub(crate) fn remove_distributed_sync_client() { + let typeid = TypeId::of::(); + let mut state_map = get_backend_client_map(); + state_map.remove(&typeid); +} + +/// Starts the server used to sync the gradients of parameters sharded across multiple devices. +pub fn start_distributed_sync_server(devices: &[B::Device], config: DistributedConfig) { + if get_distributed_sync_client::().is_none() { + let typeid = TypeId::of::(); + let mut state_map = get_backend_client_map(); + let client = DistributedSyncClient::::new(devices.len(), config); + state_map.insert(typeid, Box::new(client.clone())); + } +} + +/// Close the gradient syncing server. +pub fn close_distributed_sync_server() { + if let Some(client) = get_distributed_sync_client::() { + client.close(); + remove_distributed_sync_client::(); + } +} diff --git a/crates/burn-backend/src/backend/distributed/client.rs b/crates/burn-backend/src/backend/distributed/client.rs new file mode 100644 index 0000000..94e80fd --- /dev/null +++ b/crates/burn-backend/src/backend/distributed/client.rs @@ -0,0 +1,75 @@ +use std::{sync::mpsc::Sender, thread::spawn}; + +use crate::Backend; +use crate::tensor::Device; + +use crate::distributed::{ + DistributedConfig, DistributedParams, TensorRef, server::DistributedSyncServer, +}; + +pub(crate) enum ActionMessage { + Message(DistributedSyncMessage), + Close(), +} + +pub(crate) enum DistributedSyncMessage { + RegisterSyncParameters(Vec), + TensorSync((TensorRef, DistributedParams)), + #[allow(clippy::type_complexity)] + CollectiveSync((Device, oneshot::Sender>)), +} + +#[derive(Clone)] +pub struct DistributedSyncClient { + sender: Sender>, +} + +impl DistributedSyncClient { + pub(crate) fn new(num_devices: usize, config: DistributedConfig) -> Self { + let (tx, rx) = std::sync::mpsc::channel(); + + let mut server = DistributedSyncServer::new(num_devices, config); + spawn(move || { + while let ActionMessage::Message(msg) = + rx.recv().expect("Gradient sync server disconnected.") + { + server.process_message(msg) + } + }); + Self { sender: tx } + } + + pub fn register_sync_parameters(&self, sharded_params: Vec) { + self.sender + .send(ActionMessage::Message( + DistributedSyncMessage::RegisterSyncParameters(sharded_params), + )) + .unwrap(); + } + + pub fn submit_gradient_sync(&self, tensor: TensorRef, params: DistributedParams) { + self.sender + .send(ActionMessage::Message(DistributedSyncMessage::TensorSync( + (tensor, params), + ))) + .unwrap(); + } + + pub fn submit_sync_collective(&self, device: Device) { + let (tx, rx) = oneshot::channel(); + + self.sender + .send(ActionMessage::Message( + DistributedSyncMessage::CollectiveSync((device.clone(), tx)), + )) + .unwrap(); + + let sync = rx.recv().expect("Can receive callback"); + + sync(); + } + + pub(crate) fn close(&self) { + self.sender.send(ActionMessage::Close()).unwrap(); + } +} diff --git a/crates/burn-backend/src/backend/distributed/mod.rs b/crates/burn-backend/src/backend/distributed/mod.rs new file mode 100644 index 0000000..d2f7197 --- /dev/null +++ b/crates/burn-backend/src/backend/distributed/mod.rs @@ -0,0 +1,51 @@ +//! Types and helpers for inter-device operations. + +#[cfg(feature = "std")] +pub(crate) mod api; +#[cfg(feature = "std")] +pub(crate) mod client; +mod ops; +#[cfg(feature = "std")] +pub(crate) mod server; + +#[cfg(feature = "std")] +pub use api::*; +pub use ops::*; + +pub use burn_std::distributed::*; +/// A unique identifier for a parameter distributed across multiple devices. +pub type DistributedParamId = burn_std::id::ParamId; + +use crate::{Backend, TensorMetadata, tensor::FloatTensor}; + +/// Parameters for a tensor that is sharded across multiple devices. +#[derive(Debug, Clone)] +pub struct DistributedParams { + /// The tensor's [DistributedParamId]. + pub param_id: DistributedParamId, +} + +/// A tensor handle used for a collective operation, that is not yet valid for use. +/// We must ensure collective operations are completed before accessing the underlying data. +#[derive(new, Clone)] +pub struct CollectiveTensor { + handle: FloatTensor, +} + +impl CollectiveTensor { + /// Synchronizes the collective operation and returns a valid tensor handle. + pub fn resolve(self) -> FloatTensor { + B::sync_collective(&self.handle.device()); + self.handle + } + + /// Returns the tensor handle without synchronizing. + /// + /// # Safety + /// + /// The caller must ensure that `sync_collective()` is called before + /// the returned handle is used in any computation. + pub unsafe fn assume_resolved(self) -> FloatTensor { + self.handle + } +} diff --git a/crates/burn-backend/src/backend/distributed/ops.rs b/crates/burn-backend/src/backend/distributed/ops.rs new file mode 100644 index 0000000..12e7c3a --- /dev/null +++ b/crates/burn-backend/src/backend/distributed/ops.rs @@ -0,0 +1,166 @@ +use alloc::vec::Vec; + +use crate::{ + Backend, DeviceId, TensorMetadata, + distributed::CollectiveTensor, + tensor::{Device, FloatTensor}, +}; + +use crate::distributed::{DistributedConfig, DistributedParams, ReduceOperation}; + +#[cfg(feature = "std")] +use crate::distributed::{ + close_distributed_sync_server, get_distributed_sync_client, start_distributed_sync_server, +}; + +/// Mutable reference to a float tensor. +#[derive(Clone)] +pub struct TensorRef(pub *mut FloatTensor); +unsafe impl Sync for TensorRef where B: Backend {} +unsafe impl Send for TensorRef where B: Backend {} + +// TODO : The following functions should be moved in the `burn-autodiff` crate. The difficulty is in not discriminating between +// the dispatch backend and its inner dispatched backend when calling the communication server API. This implementation makes +// it easy by implementing `DistributedOps` for `DispatchBackend`. +// The functions in question : +// * `start_communication_server` +// * `close_communication_server` +// * `register_sync_parameters` +// * `submit_sync_collective` +// * `submit_gradient_sync` + +/// Operations on communication tensors. +pub trait DistributedOps { + /// Start the communication server used to orchestrate tensor syncing between devices. + /// + /// # Arguments + /// + /// * `devices` - The devices to orchestrate. + fn start_communication_server(devices: &[B::Device], config: DistributedConfig) { + #[cfg(feature = "std")] + start_distributed_sync_server::(devices, config); + #[cfg(not(feature = "std"))] + let _ = (devices, config); + } + + /// Close the communication server used to orchestrate syncing between devices. + /// + /// # Arguments + /// + /// * `device` - A device on the backend. + fn close_communication_server(_device: &B::Device) { + #[cfg(feature = "std")] + close_distributed_sync_server::(); + } + + /// Register the parameters that will require gradient synchronization for the upcoming backward pass. + /// + /// This must be called before the backward pass on each device so the gradient sync server + /// can coordinate collective operations across all devices. + /// + /// # Arguments + /// + /// * `device` - The device calling the initialization. + /// * `distributed_params` - A list of [`DistributedParams`] of the tensors to sync. + fn register_sync_parameters(_device: &B::Device, distributed_params: Vec) { + #[cfg(feature = "std")] + if let Some(sync_client) = get_distributed_sync_client::() { + sync_client.register_sync_parameters(distributed_params); + }; + #[cfg(not(feature = "std"))] + let _ = distributed_params; + } + + /// Tell the gradient sync server that this device has submitted all its sync operations and is ready to be synchronized. + /// + /// # Arguments + /// + /// * `device` - The device on which to sync. + fn submit_sync_collective(device: &B::Device) { + #[cfg(feature = "std")] + if let Some(sync_client) = get_distributed_sync_client::() { + sync_client.submit_sync_collective(device.clone()); + }; + #[cfg(not(feature = "std"))] + let _ = device; + } + + /// Submit a gradient tensor for synchronization across all devices. + /// + /// The gradient is sent to the gradient sync server, which will launch the all-reduce + /// operation once all devices have submitted their gradient for this parameter. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to synchronize. + /// * `distributed_params` - The [`DistributedParams`] for the parameter. + fn submit_gradient_sync(tensor: TensorRef, distributed_params: DistributedParams) { + #[cfg(feature = "std")] + if let Some(sync_client) = get_distributed_sync_client::() { + sync_client.submit_gradient_sync(tensor, distributed_params); + }; + #[cfg(not(feature = "std"))] + let _ = (tensor, distributed_params); + } + + /// all_reduce operation. + /// + /// # Arguments + /// + /// * `tensors` - The tensors on which to perform all_reduce. + /// * `op` - The [`ReduceOperation`]. + /// + /// # Returns + /// + /// The corresponding [CollectiveTensor]. + fn all_reduce( + _tensor: FloatTensor, + _op: ReduceOperation, + _device_ids: Vec, + ) -> CollectiveTensor { + unimplemented!() + } + + /// Sync the collective operations. + /// + /// # Arguments + /// + /// * `device` - The device to sync. + fn sync_collective(_device: &B::Device) { + unimplemented!() + } + + /// Get the device of the tensor reference. + /// + /// # Arguments + /// + /// * `tensor` - The tensor reference. + /// + /// # Returns + /// + /// The device of the tensor. + /// + /// # Safety + /// + /// Ensure that the tensors are not accessed/modified when calling. + unsafe fn comm_device(tensor: &TensorRef) -> Device { + unsafe { &(*tensor.0) }.device() + } + + /// Returns a clone of the float tensor from the tensor reference. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// + /// # Returns + /// + /// A float tensor containing a copy of the data of the given tensor. + /// + /// # Safety + /// + /// Ensure that the tensors are not accessed/modified when calling. + unsafe fn float_from_ref(tensor: &TensorRef) -> FloatTensor { + unsafe { (*tensor.0).clone() } + } +} diff --git a/crates/burn-backend/src/backend/distributed/server.rs b/crates/burn-backend/src/backend/distributed/server.rs new file mode 100644 index 0000000..1e52fa7 --- /dev/null +++ b/crates/burn-backend/src/backend/distributed/server.rs @@ -0,0 +1,140 @@ +use std::collections::HashMap; + +use crate::{Backend, TensorMetadata}; +use crate::{DeviceId, DeviceOps, tensor::Device}; + +use crate::distributed::{ + DistributedConfig, DistributedParamId, DistributedParams, TensorRef, + client::DistributedSyncMessage, +}; + +pub(crate) struct DistributedSyncServer { + config: DistributedConfig, + all_reduce_ops_queue: HashMap>>, + param_required_map: HashMap, + num_devices: usize, + devices_registered: usize, + syncing_devices: Vec>, + devices_synced: usize, + callbacks: HashMap>>, +} + +impl DistributedSyncServer { + /// Create a new gradient sync server instance. + pub(crate) fn new(num_devices: usize, config: DistributedConfig) -> Self { + Self { + config, + all_reduce_ops_queue: HashMap::default(), + param_required_map: HashMap::default(), + num_devices, + devices_registered: 0, + syncing_devices: vec![], + devices_synced: 0, + callbacks: HashMap::default(), + } + } + + /// Process message from client. + pub(crate) fn process_message(&mut self, msg: DistributedSyncMessage) { + match msg { + DistributedSyncMessage::RegisterSyncParameters(params) => { + self.register_sync_params(params) + } + DistributedSyncMessage::TensorSync((tensor, params)) => { + self.register_tensor(tensor, params) + } + DistributedSyncMessage::CollectiveSync((device, callback)) => { + self.collective_sync(device, callback) + } + } + } + + /// Called at the start of the backward process. Lets the device announce what parameters are nodes in the autodiff graph and how many times they are required. + fn register_sync_params(&mut self, sharded_params: Vec) { + sharded_params.iter().for_each(|params| { + *self.param_required_map.entry(params.param_id).or_insert(0) += 1; + }); + self.devices_registered += 1; + } + + /// Called on registration of a gradient. Calls the all_reduce operation for any parameter that is no longer required in the autodiff graph. + fn register_tensor(&mut self, tensor: TensorRef, sharded_params: DistributedParams) { + let op_queue = self + .all_reduce_ops_queue + .entry(sharded_params.param_id) + .or_insert(vec![]); + op_queue.push(tensor.clone()); + self.launch_ops(); + } + + fn collective_sync( + &mut self, + device: Device, + callback: oneshot::Sender>, + ) { + self.callbacks.insert(device.id(), callback); + self.syncing_devices.push(device.clone()); + self.try_launch_sync(); + } + + fn try_launch_sync(&mut self) { + if self.all_reduce_ops_queue.is_empty() { + for d in self.syncing_devices.clone() { + let callback = self.callbacks.remove(&d.id()).unwrap(); + let closure = Box::new(move || B::sync_collective(&d)); + callback.send(closure).expect("Can send callback"); + self.devices_synced += 1; + } + self.syncing_devices.clear(); + } + + if self.devices_synced == self.num_devices { + self.devices_registered = 0; + self.devices_synced = 0; + self.param_required_map.clear(); + self.callbacks.clear(); + } + } + + fn launch_ops(&mut self) { + if self.devices_registered == self.num_devices { + for (param_id, num_tensors) in self.param_required_map.clone() { + let queued_tensors = self.all_reduce_ops_queue.entry(param_id).or_insert(vec![]); + + if num_tensors == queued_tensors.len() { + // Safety: Tensors sent to the `DistributedSyncServer` should not be accessed or modified until the end of the backward pass. + let device_ids = queued_tensors + .iter() + .map(|t| unsafe { &*t.0 }.device().id()) + .collect::>(); + let reduced_tensors: Vec = queued_tensors + .iter() + .map(|tensor| + // Safety: we can call `assume_resolved` on these tensors since we know `B::sync_collective` is called + // at the end of the backward pass. + unsafe { + B::all_reduce( + (*tensor.0).clone(), + self.config.all_reduce_op, + device_ids.clone(), + ) + .assume_resolved() + }) + .collect(); + + // Make the tensor reference point to the reduced tensor to perform an in-place all_reduce. + // Safety: `B::sync_collective` should be automatically called after the backward pass. + unsafe { + queued_tensors.iter().zip(reduced_tensors).for_each( + |(tensor_ref, reduced_tensor)| *tensor_ref.0 = reduced_tensor, + ); + } + + self.all_reduce_ops_queue.remove(¶m_id).unwrap(); + self.param_required_map.remove(¶m_id).unwrap(); + self.try_launch_sync(); + } + } + } + } +} diff --git a/crates/burn-backend/src/backend/mod.rs b/crates/burn-backend/src/backend/mod.rs new file mode 100644 index 0000000..e707635 --- /dev/null +++ b/crates/burn-backend/src/backend/mod.rs @@ -0,0 +1,13 @@ +mod base; +mod device; +mod primitive; + +pub use base::*; +pub use device::*; +pub use primitive::*; + +/// Backend operations on tensors. +pub mod ops; + +/// Distributed backend extension. +pub mod distributed; diff --git a/crates/burn-backend/src/backend/ops/activation.rs b/crates/burn-backend/src/backend/ops/activation.rs new file mode 100644 index 0000000..3190b57 --- /dev/null +++ b/crates/burn-backend/src/backend/ops/activation.rs @@ -0,0 +1,343 @@ +use crate::tensor::FloatTensor; +use crate::{Backend, Scalar, TensorMetadata, get_device_settings}; +use core::f64::consts::SQRT_2; + +/// Activation function operations. +/// +/// This trait let backend implementations override activation functions for better performance. +pub trait ActivationOps { + /// Applies the LeakyReLU activation function. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `negative_slope` - The negative_slope value that values smaller than 0 are multiplied with. + /// + /// # Returns + /// + /// The output tensor. + fn leaky_relu(tensor: FloatTensor, negative_slope: Scalar) -> FloatTensor { + let bool_dtype = get_device_settings::(&tensor.device()).bool_dtype; + let mask = B::float_lower_elem(tensor.clone(), 0f32.into(), bool_dtype); + let scaled_tensor = B::float_mul_scalar(tensor.clone(), negative_slope); + + // Update the tensor where the values are `< 0` by `tensor * negative_slope`. + B::float_mask_where(tensor, mask, scaled_tensor) + } + + /// Applies the ReLU activation function. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// + /// # Returns + /// + /// The output tensor. + fn relu(tensor: FloatTensor) -> FloatTensor { + let bool_dtype = get_device_settings::(&tensor.device()).bool_dtype; + let mask = B::float_lower_equal_elem(tensor.clone(), 0f32.into(), bool_dtype); + + B::float_mask_fill(tensor, mask, 0f32.into()) + } + + /// Applies the ReLU activation function backward. + /// + /// # Arguments + /// + /// * `output` - The output tensor. + /// + /// # Returns + /// + /// The gradient. + fn relu_backward(output: FloatTensor, grad: FloatTensor) -> FloatTensor { + let bool_dtype = get_device_settings::(&output.device()).bool_dtype; + let mask = B::float_lower_equal_elem(output, 0f32.into(), bool_dtype); + + B::float_mask_fill(grad, mask, 0.into()) + } + + /// Applies the Gelu activation function. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// + /// # Returns + /// + /// The output tensor. + fn gelu(tensor: FloatTensor) -> FloatTensor { + let x = B::float_div_scalar(tensor.clone(), SQRT_2.into()); + let x = B::float_erf(x); + let x = B::float_add_scalar(x, 1f32.into()); + let x = B::float_mul(tensor, x); + + B::float_div_scalar(x, 2f32.into()) + } + /// Applies the PReLu activation function. + /// # Arguments + /// * `tensor` - The input tensor + /// * `alpha` - The weight tensor + fn prelu(tensor: FloatTensor, alpha: FloatTensor) -> FloatTensor { + let bool_dtype = get_device_settings::(&tensor.device()).bool_dtype; + let mask = B::float_lower_elem(tensor.clone(), 0f32.into(), bool_dtype); + let scaled_tensor = B::float_mul(tensor.clone(), alpha); + B::float_mask_where(tensor, mask, scaled_tensor) + } + + /// Applies the Gelu activation function backward. + /// + /// # Arguments + /// + /// * `x` - The tensor. + /// * `grad` - The gradient. + /// + /// # Returns + /// + /// The output tensor. + fn gelu_backward(x: FloatTensor, grad: FloatTensor) -> FloatTensor { + // Derivative of the approximate gelu implementation based on tanh. + + let constant_1 = 0.0356774; + let constant_2 = 0.797885; + let constant_3 = 0.0535161; + let constant_4 = 0.398942; + + let x3 = B::float_powi_scalar(x.clone(), 3.into()); + + let c1 = B::float_mul_scalar(x3.clone(), constant_1.into()); + let c2 = B::float_mul_scalar(x.clone(), constant_2.into()); + let c3 = B::float_mul_scalar(x3, constant_3.into()); + let c4 = B::float_mul_scalar(x, constant_4.into()); + + let inner1 = B::float_add(c1, c2); + let inner2 = B::float_add(c3, c4); + + let tanh = B::float_tanh(inner1); + + let sech = B::float_powi_scalar(tanh.clone(), 2.into()); + let sech = B::float_neg(sech); + let sech = B::float_add_scalar(sech, 1.into()); + + let y1 = B::float_mul_scalar(tanh, 0.5.into()); + let y2 = B::float_mul(inner2, sech); + let y2 = B::float_add_scalar(y2, 0.5.into()); + let y = B::float_add(y1, y2); + + B::float_mul(y, grad) + } + + /// Applies the Sigmoid activation function. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// + /// # Returns + /// + /// The output tensor. + fn sigmoid(tensor: FloatTensor) -> FloatTensor { + let dtype = tensor.dtype(); + let tensor_full = B::float_cast(tensor, burn_std::FloatDType::F32); + let tensor_tmp = B::float_exp(B::float_neg(B::float_log(B::float_add_scalar( + B::float_exp(B::float_neg(tensor_full)), + 1.0.into(), + )))); + + B::float_cast(tensor_tmp, dtype.into()) + } + + /// Applies the Sigmoid activation function backward. + /// + /// # Arguments + /// + /// * `output` - The output tensor of the sigmoid function. + /// * `grad` - The gradient. + /// + /// # Returns + /// + /// The output tensor. + fn sigmoid_backward(output: FloatTensor, grad: FloatTensor) -> FloatTensor { + let value = B::float_mul( + output.clone(), + B::float_add_scalar(B::float_neg(output), 1.0.into()), + ); + B::float_mul(value, grad) + } + + /// Applies the hard Sigmoid activation function. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `alpha` - The alpha value that the tensor is multiplied with. + /// * `beta` - The beta value that is added to the tensor + /// + /// # Returns + /// + /// The output tensor. + fn hard_sigmoid(tensor: FloatTensor, alpha: Scalar, beta: Scalar) -> FloatTensor { + let dtype = tensor.dtype(); + let tensor_full = B::float_cast(tensor, burn_std::FloatDType::F32); + + let tensor_tmp = B::float_clamp( + B::float_add_scalar(B::float_mul_scalar(tensor_full, alpha), beta), + 0.0.into(), + 1.0.into(), + ); + + B::float_cast(tensor_tmp, dtype.into()) + } + + /// Applies the LogSigmoid activation function. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// + /// # Returns + /// + /// The output tensor. + fn log_sigmoid(tensor: FloatTensor) -> FloatTensor { + // To avoid overflow, we use the log-sum-exp trick. + // + // ```ignore + // log(sigmoid(x)) = log(1/(1 + exp(-x))) + // = log(1) - log(1 + exp(-x)) + // = -log(1 + exp(-x)) + // = -log(exp(0) + exp(-x)) + // ``` + // The `exp(t)` of even a moderate-magnitude positive number can be astronomically huge, so we + // subtract the `max(t, 0)` of each value (where `t = -x` in this case). This results in the + // following equivalence: + // ```ignore + // log(sigmoid(x)) = -(max(-x, 0) + log(exp(-max(-x, 0)) + exp(-x - max(-x, 0)))) + // ``` + // + // This extends the range of values for which we obtain accurate results. + + // max(-x, 0) + let bool_dtype = get_device_settings::(&tensor.device()).bool_dtype; + let tensor_neg = B::float_neg(tensor); + let mask = B::float_lower_elem(tensor_neg.clone(), 0f32.into(), bool_dtype); + let max_elem = B::float_mask_fill(tensor_neg.clone(), mask, 0f32.into()); + let max_elem_neg = B::float_neg(max_elem.clone()); + + // z = exp(-max(-x, 0)) + exp(-x - max(-x, 0)) + let z = B::float_add( + B::float_exp(max_elem_neg.clone()), + B::float_exp(B::float_sub(tensor_neg, max_elem.clone())), + ); + + // -max(-x, 0) - log(-z) + B::float_sub(max_elem_neg, B::float_log(z)) + } + + /// Applies the softmax function along the given dimension. + /// + /// Uses the max-shift trick for numerical stability: the per-row `max` is detached + /// so no gradient flows back through it (the shift is a numerical-stability + /// transformation, not part of the function). + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `dim` - The dimension along which softmax is computed. + /// + /// # Returns + /// + /// The output tensor. + fn softmax(tensor: FloatTensor, dim: usize) -> FloatTensor { + let max = B::float_max_dim(B::float_detach(tensor.clone()), dim); + let shifted = B::float_sub(tensor, max); + let exp = B::float_exp(shifted); + let sum = B::float_sum_dim(exp.clone(), dim); + B::float_div(exp, sum) + } + + /// Applies the log-softmax function along the given dimension. + /// + /// Computed via the log-sum-exp trick with a detached max-shift for numerical + /// stability. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `dim` - The dimension along which log-softmax is computed. + /// + /// # Returns + /// + /// The output tensor. + fn log_softmax(tensor: FloatTensor, dim: usize) -> FloatTensor { + let max = B::float_max_dim(B::float_detach(tensor.clone()), dim); + let shifted = B::float_sub(tensor, max); + let log_sum_exp = B::float_log(B::float_sum_dim(B::float_exp(shifted.clone()), dim)); + B::float_sub(shifted, log_sum_exp) + } + + /// Applies the softmin function along the given dimension. + /// + /// Equivalent to `softmax(-tensor, dim)`. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `dim` - The dimension along which softmin is computed. + /// + /// # Returns + /// + /// The output tensor. + fn softmin(tensor: FloatTensor, dim: usize) -> FloatTensor { + Self::softmax(B::float_neg(tensor), dim) + } + + /// Applies the LogSigmoid activation function backward. + /// + /// # Arguments + /// + /// * `x` - The input tensor. + /// * `grad` - The gradient. + /// + /// # Returns + /// + /// The output gradient. + fn log_sigmoid_backward(x: FloatTensor, grad: FloatTensor) -> FloatTensor { + // Derivative of -max(-x, 0) - log(exp(-max(-x, 0)) - exp(-x - max(-x, 0)))) is + // -max_derive - (-max_derive * exp(-max(-x, 0)) + (-1 - max_derive) * exp(-x - max(-x, 0))) / z + // where z = exp(-max(-x, 0)) + exp(-x - max(-x, 0)) + // + // This simplifies to: + // -max_derive - (z-1)/z if x is >= 0 + // -max_derive + (z-1)/z if x is < 0 + + let shape = x.shape(); + let dtype = x.dtype(); + let device = x.device(); + let bool_dtype = get_device_settings::(&device).bool_dtype; + + // max(-x, 0) + let x_neg = B::float_neg(x); + let mask = B::float_lower_elem(x_neg.clone(), 0f32.into(), bool_dtype); // -x < 0 or x >= 0 + let max_elem = B::float_mask_fill(x_neg.clone(), mask.clone(), 0f32.into()); + + // z = exp(-max(-x, 0)) + exp(-x - max(-x, 0)) + let z = B::float_add( + B::float_exp(B::float_neg(max_elem.clone())), + B::float_exp(B::float_sub(x_neg, max_elem)), + ); + + // Derivative of max(-x, 0) is 1 if x < 0 or 0 if x >= 0 + let ones = B::float_ones(shape, &device, dtype.into()); + let max_derive = B::float_mask_fill(ones.clone(), mask.clone(), 0f32.into()); + let sign = B::float_mask_fill(ones.clone(), mask, (-1f32).into()); + + // grad * (max_derive - sign * (1 - (1 / z))) + B::float_mul( + grad, + B::float_sub( + max_derive, + B::float_mul(sign, B::float_sub(ones, B::float_recip(z))), + ), + ) + } +} diff --git a/crates/burn-backend/src/backend/ops/argwhere.rs b/crates/burn-backend/src/backend/ops/argwhere.rs new file mode 100644 index 0000000..afad83b --- /dev/null +++ b/crates/burn-backend/src/backend/ops/argwhere.rs @@ -0,0 +1,72 @@ +use crate::tensor::{Device, IntTensor}; +use crate::{Backend, TensorData}; +use alloc::vec::Vec; +use burn_std::{Element, ElementConversion, IntDType}; + +/// Compute the indices of the elements that are non-zero, grouped by element. +/// +/// # Arguments +/// +/// * `data` - The input tensor data. +/// +/// # Returns +/// +/// A 2D tensor containing the indices of all non-zero elements of the given tensor. +/// Each row contains the indices of a non-zero element. +/// +/// # Remarks +/// +/// This is a fallback solution that used only when the backend doesn't have the corresponding implementation. +/// Ideally, it is supposed to be implemented by the backend and the backend implementation will be resolved +/// by static dispatch. It is not designed for direct usage by users, and not recommended to import +/// or use this function directly. +pub fn argwhere_data( + data: TensorData, + device: &Device, + out_dtype: IntDType, +) -> IntTensor { + let out = match out_dtype { + IntDType::I64 => argwhere_data_impl::(data), + IntDType::I32 => argwhere_data_impl::(data), + IntDType::I16 => argwhere_data_impl::(data), + IntDType::I8 => argwhere_data_impl::(data), + IntDType::U64 => argwhere_data_impl::(data), + IntDType::U32 => argwhere_data_impl::(data), + IntDType::U16 => argwhere_data_impl::(data), + IntDType::U8 => argwhere_data_impl::(data), + }; + + B::int_from_data(out, device) +} + +fn argwhere_data_impl(data: TensorData) -> TensorData { + let dims = &data.shape; + let ndims = dims.len(); + let count_nonzero = data.iter::().filter(|&v| v).count(); + + /// Converts a flat index into a vector of indices for the specified tensor shape + fn unravel_index(index: usize, shape: &[usize]) -> Vec { + shape + .iter() + .rev() + .scan(index, |i, size| { + let dim_idx = *i % size; + *i /= size; + Some((dim_idx as i64).elem()) + }) + .collect::>() + .into_iter() + .rev() + .collect() + } + + let indices = data + .iter::() + .enumerate() + .filter_map(|(index, v)| if v { Some(index) } else { None }) + .map(|index| unravel_index::(index, dims)) + .collect::>() + .concat(); + + TensorData::new(indices, [count_nonzero, ndims]) +} diff --git a/crates/burn-backend/src/backend/ops/bool_tensor.rs b/crates/burn-backend/src/backend/ops/bool_tensor.rs new file mode 100644 index 0000000..4ecb3f2 --- /dev/null +++ b/crates/burn-backend/src/backend/ops/bool_tensor.rs @@ -0,0 +1,573 @@ +use super::{ + argwhere::argwhere_data, cat::cat_with_slice_assign, repeat_dim::repeat_with_slice_assign, +}; +use crate::tensor::{BoolTensor, Device, FloatTensor, IntTensor}; +use crate::{Backend, TensorData, TensorMetadata, get_device_settings}; +use crate::{ExecutionError, Scalar}; +use alloc::vec::Vec; +use burn_std::{BoolDType, FloatDType, IntDType, Shape, Slice}; +use core::future::Future; + +/// Bool Tensor API for basic operations, see +#[cfg_attr(doc, doc = crate::doc_tensor!())] +#[cfg_attr(not(doc), doc = "`Tensor`")] +/// for documentation on each function. +pub trait BoolTensorOps { + /// Creates a new bool tensor. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `device` - The device to create the tensor on. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The boolean tensor with the given shape. + fn bool_empty(shape: Shape, device: &Device, dtype: BoolDType) -> BoolTensor; + + /// Creates a new bool tensor filled false. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `device` - The device to create the tensor on. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The boolean tensor filled with false. + fn bool_zeros(shape: Shape, device: &Device, dtype: BoolDType) -> BoolTensor; + + /// Creates a new bool tensor filled true. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `device` - The device to create the tensor on. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The boolean tensor filled with true. + fn bool_ones(shape: Shape, device: &Device, dtype: BoolDType) -> BoolTensor; + + /// Converts the tensor to a data structure. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// + /// # Returns + /// + /// The data structure with the tensor's data. + fn bool_into_data( + tensor: BoolTensor, + ) -> impl Future> + Send; + + /// Creates a tensor from the data structure. + /// + /// # Arguments + /// + /// * `data` - The data structure. + /// * `device` - The device to create the tensor on. + /// + /// # Returns + /// + /// The tensor with the data. + fn bool_from_data(data: TensorData, device: &Device) -> BoolTensor; + + /// Converts bool tensor to int tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// The int tensor with the same data as the bool tensor. + fn bool_into_int(tensor: BoolTensor, out_dtype: IntDType) -> IntTensor; + + /// Converts bool tensor to float tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// The float tensor with the same data as the bool tensor. + fn bool_into_float(tensor: BoolTensor, out_dtype: FloatDType) -> FloatTensor; + + /// Moves the tensor to the device. + fn bool_to_device(tensor: BoolTensor, device: &Device) -> BoolTensor; + + /// Reshapes the tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `shape` - The new shape. + /// + /// # Returns + /// + /// The tensor with the new shape. + fn bool_reshape(tensor: BoolTensor, shape: Shape) -> BoolTensor; + + /// Gets the values from the tensor for the given ranges. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `slices` - The slices specifying ranges and steps for each dimension. + /// + /// # Returns + /// + /// The tensor with the values for the given slices. + /// + /// # Note + /// + /// Empty slices (where start >= end) are handled at the high-level tensor API and will not + /// be passed to this method. Backend implementations do not need to handle empty slices. + fn bool_slice(tensor: BoolTensor, slices: &[Slice]) -> BoolTensor; + + /// Sets the values in the tensor for the given ranges. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `ranges` - The ranges to set the values for. + /// * `value` - The values to set. + /// + /// # Returns + /// + /// The tensor with the values set for the given ranges. + /// + /// # Note + /// + /// Empty slice assignments (where any slice range produces 0 elements) are handled at the + /// high-level tensor API and will not be passed to this method. Backend implementations do + /// not need to handle empty slice assignments. + fn bool_slice_assign( + tensor: BoolTensor, + slices: &[Slice], + value: BoolTensor, + ) -> BoolTensor; + + /// Fills the tensor with values from the value tensor if the mask is true at the given + /// indices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `mask` - The mask. + /// * `value` - The value tensor. + /// + /// # Returns + /// + /// The tensor with the values filled. + fn bool_mask_where( + tensor: BoolTensor, + mask: BoolTensor, + value: BoolTensor, + ) -> BoolTensor; + + /// Fills the tensor with the given value if the mask is true at the given indices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `mask` - The mask. + /// * `value` - The value. + /// + /// # Returns + /// + /// The tensor with the values filled. + fn bool_mask_fill(tensor: BoolTensor, mask: BoolTensor, value: Scalar) -> BoolTensor; + + /// Gather elements from the tensor at the given indices. + /// + /// # Arguments + /// + /// * `dim` - The dimension to gather from. + /// * `tensor` - The tensor. + /// * `indices` - The indices. + fn bool_gather(dim: usize, tensor: BoolTensor, indices: IntTensor) -> BoolTensor; + + /// Scatter a given value to the tensor at the given indices using boolean or reduction. + /// + /// # Arguments + /// + /// * `dim` - The dimension to scatter to. + /// * `tensor` - The tensor. + /// * `indices` - The indices. + /// * `value` - The value. + /// + /// # Returns + /// + /// The tensor with the values scattered. + fn bool_scatter_or( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor; + + /// Select tensor elements along the given dimension corresponding to the given indices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to select from. + /// * `dim` - The dimension to select from. + /// * `indices` - The indices of the elements to select. + /// + /// # Returns + /// + /// The tensor with the selected elements. + fn bool_select(tensor: BoolTensor, dim: usize, indices: IntTensor) -> BoolTensor; + + /// Assign the selected elements along the given dimension corresponding to the given indices + /// to the given value using sum reduction. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to assign the values to. + /// * `dim` - The dimension to select from. + /// * `indices` - The indices of the elements to assign. + /// * `value` - The values to assign. + /// + /// # Returns + /// + /// The tensor with the assigned values. + fn bool_select_or( + tensor: BoolTensor, + dim: usize, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor; + + /// Repeats one dimension of the tensor a given number of times along that dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `dim` - The dimension to repeat. + /// * `times` - The number of times to repeat the dimension. + /// + /// # Returns + /// + /// The tensor with the dimension repeated. + fn bool_repeat_dim(tensor: BoolTensor, dim: usize, times: usize) -> BoolTensor { + let device = tensor.device(); + repeat_with_slice_assign::( + tensor, + dim, + times, + device, + |shape, device, dtype| B::bool_empty(shape, device, dtype.into()), + B::bool_slice_assign, + ) + } + + /// Concatenates the tensors along the given dimension. + /// + /// # Arguments + /// + /// * `tensors` - The tensors to concatenate. + /// * `dim` - The dimension to concatenate along. + /// + /// # Returns + /// + /// The tensor with the tensors concatenated along the given dimension. + /// + /// # Note + /// + /// Empty tensors (where the concatenation dimension has size 0) are filtered out at the + /// high-level tensor API and will not be passed to this method. Backend implementations do + /// not need to handle empty tensors. + fn bool_cat(tensors: Vec>, dim: usize) -> BoolTensor { + let first_tensor = tensors.first().expect("Tensors should not be empty"); + let device = first_tensor.device(); + cat_with_slice_assign::( + tensors, + dim, + device, + |shape, device, dtype| B::bool_empty(shape, device, dtype.into()), + B::bool_slice_assign, + ) + } + + /// Equates the two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// The tensor with the result of the equate. + fn bool_equal(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor; + + /// Element-wise non-equality comparison. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// The tensor with the result of the comparison. + fn bool_not_equal(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + let equal_tensor = B::bool_equal(lhs, rhs); + B::bool_not(equal_tensor) + } + + /// Element-wise equality comparison with a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The boolean tensor with the result of the comparison. + fn bool_equal_elem(lhs: BoolTensor, rhs: Scalar) -> BoolTensor; + + /// Element-wise non-equality comparison with a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The boolean tensor with the result of the comparison. + fn bool_not_equal_elem(lhs: BoolTensor, rhs: Scalar) -> BoolTensor { + let equal_tensor = B::bool_equal_elem(lhs, rhs); + B::bool_not(equal_tensor) + } + + /// Inverses boolean values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// + /// # Returns + /// + /// The tensor with the result of the negation. + fn bool_not(tensor: BoolTensor) -> BoolTensor; + + /// Executes the logical and (`&&`) operation on two boolean tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// The tensor with the result of the logical and. + fn bool_and(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor; + + /// Executes the logical or (`||`) operation on two boolean tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// The tensor with the result of the logical or. + fn bool_or(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor; + + /// Element-wise exclusive or. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// The tensor with the result of the comparison. + fn bool_xor(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + Self::bool_not_equal(lhs, rhs) + } + + /// Transposes a bool tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to transpose. + /// + /// # Returns + /// + /// The transposed tensor. + fn bool_transpose(tensor: BoolTensor) -> BoolTensor { + let ndims = tensor.shape().num_dims(); + Self::bool_swap_dims(tensor, ndims - 2, ndims - 1) + } + + /// Swaps two dimensions of a bool tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to swap the dimensions of. + /// * `dim1` - The first dimension to swap. + /// * `dim2` - The second dimension to swap. + /// + /// # Returns + /// + /// The tensor with the dimensions swapped. + fn bool_swap_dims(tensor: BoolTensor, dim1: usize, dim2: usize) -> BoolTensor; + + /// Permutes the dimensions of a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to permute the dimensions of. + /// * `axes` - The new order of the dimensions. + /// # Returns + /// + /// The tensor with the dimensions permuted. + fn bool_permute(tensor: BoolTensor, axes: &[usize]) -> BoolTensor; + + /// Reverse the order of elements in a tensor along the given axes. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to reverse. + /// * `axes` - The axes to reverse. + /// + /// The tensor with the elements reversed. + fn bool_flip(tensor: BoolTensor, axes: &[usize]) -> BoolTensor; + + /// Tests if any element in the boolean `tensor` evaluates to True. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// + /// # Returns + /// + /// A boolean tensor with a single element, True if any element in the tensor is True, False otherwise. + fn bool_any(tensor: BoolTensor) -> BoolTensor { + let dtype = tensor.dtype(); + let int_dtype = get_device_settings::(&tensor.device()).int_dtype; + let sum = B::int_sum(B::bool_into_int(tensor, int_dtype)); + B::int_greater_elem(sum, 0.into(), dtype.into()) + } + + /// Tests if any element in the boolean `tensor` evaluates to True along a given dimension `dim`. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// * `dim` - The axis along which to test. + /// + /// # Returns + /// + /// A boolean tensor `Tensor` with the same size as input `tensor`, except in the `dim` axis + /// where the size is 1. The elem in the `dim` axis is True if any element along this dim in the input + /// evaluates to True, False otherwise. + fn bool_any_dim(tensor: BoolTensor, dim: usize) -> BoolTensor { + let dtype = tensor.dtype(); + let int_dtype = get_device_settings::(&tensor.device()).int_dtype; + let sum = B::int_sum_dim(B::bool_into_int(tensor, int_dtype), dim); + B::int_greater_elem(sum, 0.into(), dtype.into()) + } + + /// Tests if all elements in the boolean `tensor` evaluate to True. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// + /// # Returns + /// + /// A boolean tensor `Tensor` with a single element, True if all elements in the input tensor + /// evaluate to True, False otherwise. + fn bool_all(tensor: BoolTensor) -> BoolTensor { + let dtype = tensor.dtype(); + let int_dtype = get_device_settings::(&tensor.device()).int_dtype; + let num_elems = tensor.shape().num_elements() as i64; + let sum = B::int_sum(B::bool_into_int(tensor, int_dtype)); + B::int_equal_elem(sum, num_elems.into(), dtype.into()) + } + + /// Tests if all elements in the boolean `tensor` evaluate to True along a given dimension `dim`. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// * `dim` - The axis along which to test. + /// + /// # Returns + /// + /// A boolean tensor `Tensor` with the same size as input `tensor`, except in the `dim` axis + /// where the size is 1. The elem in the `dim` axis is True if all elements along this dim in the input + /// evaluates to True, False otherwise. + fn bool_all_dim(tensor: BoolTensor, dim: usize) -> BoolTensor { + let dtype = tensor.dtype(); + let int_dtype = get_device_settings::(&tensor.device()).int_dtype; + let num_elems = tensor.shape()[dim] as i64; + let sum = B::int_sum_dim(B::bool_into_int(tensor, int_dtype), dim); + B::int_equal_elem(sum, num_elems.into(), dtype.into()) + } + + /// Compute the indices of the elements that are non-zero, grouped by element. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A 2D tensor containing the indices of all non-zero elements of the given tensor. + /// Each row contains the indices of a non-zero element. + fn bool_argwhere( + tensor: BoolTensor, + out_dtype: IntDType, + ) -> impl Future> + 'static + Send { + async move { + // Size of each output tensor is variable (= number of nonzero elements in the tensor). + // Reading the data to count the number of truth values might cause sync but is required. + let device = &tensor.device(); + let data = B::bool_into_data(tensor) + .await + .expect("Can read the data without error"); + argwhere_data::(data, device, out_dtype) + } + } + + /// Broadcasts the bool `tensor` to the given `shape`. + fn bool_expand(tensor: BoolTensor, shape: Shape) -> BoolTensor; + + /// Unfold windows along a dimension. + /// + /// Returns a view of the tensor with all complete windows of size `size` in dimension `dim`; + /// where windows are advanced by `step` at each index. + /// + /// The number of windows is `max(0, (shape[dim] - size).ceil_div(step))`. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor to unfold; of shape ``[pre=..., dim shape, post=...]`` + /// * `dim` - the selected dim. + /// * `size` - the size of each unfolded window. + /// * `step` - the step between each window. + /// + /// # Returns + /// + /// A tensor view with shape ``[pre=..., windows, size, post=...]``. + fn bool_unfold(tensor: BoolTensor, dim: usize, size: usize, step: usize) -> BoolTensor; +} diff --git a/crates/burn-backend/src/backend/ops/cat.rs b/crates/burn-backend/src/backend/ops/cat.rs new file mode 100644 index 0000000..6334fce --- /dev/null +++ b/crates/burn-backend/src/backend/ops/cat.rs @@ -0,0 +1,45 @@ +use crate::{Backend, TensorMetadata, tensor::Device}; +use alloc::vec::Vec; +use burn_std::{DType, Shape, Slice}; + +pub(crate) fn cat_with_slice_assign( + tensors: Vec, + dim: usize, + device: Device, + empty: E, + slice_assign: SA, +) -> T +where + T: TensorMetadata, + B: Backend, + E: Fn(Shape, &Device, DType) -> T, + SA: Fn(T, &[Slice], T) -> T, +{ + let first_tensor = tensors.first().expect("Tensors should not be empty"); + let mut shape = first_tensor.shape(); + let dtype = first_tensor.dtype(); + + let output_dim_length: usize = tensors.iter().map(|tensor| tensor.shape()[dim]).sum(); + shape[dim] = output_dim_length; + + let mut tensor_output = empty(shape.clone(), &device, dtype); + + let indices_select_all = shape.iter().map(|d| 0..*d).collect::>(); + + let mut output_index = 0; + for tensor in tensors { + let mut indices = indices_select_all.clone(); + let tensor_dim_length = tensor.shape()[dim]; + indices[dim] = output_index..output_index + tensor_dim_length; + output_index += tensor_dim_length; + + // Convert ranges to Slice + let slices: Vec = indices + .iter() + .map(|r| Slice::new(r.start as isize, Some(r.end as isize), 1)) + .collect(); + tensor_output = slice_assign(tensor_output, &slices, tensor); + } + + tensor_output +} diff --git a/crates/burn-backend/src/backend/ops/int_tensor.rs b/crates/burn-backend/src/backend/ops/int_tensor.rs new file mode 100644 index 0000000..6f0d890 --- /dev/null +++ b/crates/burn-backend/src/backend/ops/int_tensor.rs @@ -0,0 +1,1460 @@ +use super::cat::cat_with_slice_assign; +use super::repeat_dim::repeat_with_slice_assign; +use super::sort::{argsort, sort, sort_with_indices}; +use crate::tensor::{BoolTensor, Device, FloatTensor, IntTensor}; +use crate::{Backend, Distribution, TensorData, TensorMetadata}; +use crate::{ExecutionError, Scalar, get_device_settings}; +use alloc::vec::Vec; +use burn_std::reader::try_read_sync; +use burn_std::{BoolDType, FloatDType, IntDType, Shape, Slice}; +use core::ops::Range; + +/// Int Tensor API for basic and numeric operations, see +#[cfg_attr(doc, doc = crate::doc_tensor!())] +#[cfg_attr(not(doc), doc = "`Tensor`")] +/// for documentation on each function. +pub trait IntTensorOps { + /// Creates a new int tensor. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `device` - The device to create the tensor on. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The integer tensor with the given shape. + fn int_empty(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor; + + /// Converts the tensor to a data structure. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// + /// # Returns + /// + /// The data structure with the tensor's data. + fn int_into_data( + tensor: IntTensor, + ) -> impl Future> + Send; + + /// Creates a tensor from the data structure. + /// + /// # Arguments + /// + /// * `data` - The data structure. + /// * `device` - The device to create the tensor on. + /// + /// # Returns + /// + /// The tensor with the data. + fn int_from_data(data: TensorData, device: &Device) -> IntTensor; + + /// Moves the tensor to the given device. + fn int_to_device(tensor: IntTensor, device: &Device) -> IntTensor; + + /// Reshapes the tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `shape` - The new shape. + /// + /// # Returns + /// + /// The tensor with the new shape. + fn int_reshape(tensor: IntTensor, shape: Shape) -> IntTensor; + + /// Gets the element at the given indices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `slices` - The slices specifying ranges and steps for each dimension. + /// + /// # Returns + /// + /// The elements at the given indices. + /// + /// # Note + /// + /// Empty slices (where start >= end) are handled at the high-level tensor API and will not + /// be passed to this method. Backend implementations do not need to handle empty slices. + fn int_slice(tensor: IntTensor, slices: &[Slice]) -> IntTensor; + + /// Sets the values in the tensor for the given ranges. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `ranges` - The ranges to set the values for. + /// + /// # Returns + /// + /// The tensor with the values set for the given ranges. + /// + /// # Note + /// + /// Empty slice assignments (where any slice range produces 0 elements) are handled at the + /// high-level tensor API and will not be passed to this method. Backend implementations do + /// not need to handle empty slice assignments. + fn int_slice_assign( + tensor: IntTensor, + slices: &[Slice], + value: IntTensor, + ) -> IntTensor; + + /// Converts int tensor to float tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// The int tensor with the same data as the float tensor. + fn int_into_float(tensor: IntTensor, out_dtype: FloatDType) -> FloatTensor; + + /// Fills the tensor with values from the value tensor if the mask is true at the given + /// indices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `mask` - The mask. + /// * `value` - The value tensor. + /// + /// # Returns + /// + /// The tensor with the values filled. + fn int_mask_where( + tensor: IntTensor, + mask: BoolTensor, + value: IntTensor, + ) -> IntTensor; + + /// Fills the tensor with the given value if the mask is true at the given indices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `mask` - The mask. + /// * `value` - The value. + /// + /// # Returns + /// + /// The tensor with the values filled. + fn int_mask_fill(tensor: IntTensor, mask: BoolTensor, value: Scalar) -> IntTensor; + + /// Gather elements from the tensor at the given indices. + /// + /// # Arguments + /// + /// * `dim` - The dimension to gather from. + /// * `tensor` - The tensor. + /// * `indices` - The indices. + fn int_gather(dim: usize, tensor: IntTensor, indices: IntTensor) -> IntTensor; + + /// Scatter a given value to the tensor at the given indices using sum reduction. + /// + /// # Arguments + /// + /// * `dim` - The dimension to scatter to. + /// * `tensor` - The tensor. + /// * `indices` - The indices. + /// * `value` - The value. + /// + /// # Returns + /// + /// The tensor with the values scattered. + fn int_scatter_add( + dim: usize, + tensor: IntTensor, + indices: IntTensor, + value: IntTensor, + ) -> IntTensor; + + /// Multi-dimensional scatter for int tensors. + fn int_scatter_nd( + _data: IntTensor, + _indices: IntTensor, + _values: IntTensor, + _reduction: crate::tensor::IndexingUpdateOp, + ) -> IntTensor { + unimplemented!("int_scatter_nd is not implemented for this backend") + } + + /// Multi-dimensional gather for int tensors. + fn int_gather_nd(_data: IntTensor, _indices: IntTensor) -> IntTensor { + unimplemented!("int_gather_nd is not implemented for this backend") + } + + /// Select tensor elements along the given dimension corresponding to the given indices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `dim` - The dimension to select from. + /// * `indices` - The indices. + /// + /// # Returns + /// + /// The tensor with the selected elements. + fn int_select(tensor: IntTensor, dim: usize, indices: IntTensor) -> IntTensor; + + /// Assign the selected elements along the given dimension corresponding to the given indices + /// to the given value using sum reduction. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `dim` - The dimension to select from. + /// * `indices` - The indices. + /// * `value` - The value. + /// + /// # Returns + /// + /// The tensor with the selected elements assigned to the given value. + fn int_select_add( + tensor: IntTensor, + dim: usize, + indices: IntTensor, + value: IntTensor, + ) -> IntTensor; + + /// Repeats the tensor along the given dimension the given number of times. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `dim` - The dimension to repeat. + /// * `times` - The number of times to repeat. + /// + /// # Returns + /// + /// The tensor with the given dimension repeated the given number of times. + fn int_repeat_dim(tensor: IntTensor, dim: usize, times: usize) -> IntTensor { + let device = tensor.device(); + repeat_with_slice_assign::( + tensor, + dim, + times, + device, + |shape, device, dtype| B::int_empty(shape, device, dtype.into()), + B::int_slice_assign, + ) + } + + /// Concatenates the given tensors along the given dimension. + /// + /// # Arguments + /// + /// * `tensors` - The tensors. + /// * `dim` - The dimension to concatenate along. + /// + /// # Returns + /// + /// The concatenated tensor. + /// + /// # Note + /// + /// Empty tensors (where the concatenation dimension has size 0) are filtered out at the + /// high-level tensor API and will not be passed to this method. Backend implementations do + /// not need to handle empty tensors. + fn int_cat(tensors: Vec>, dim: usize) -> IntTensor { + let first_tensor = tensors.first().expect("Tensors should not be empty"); + let device = first_tensor.device(); + cat_with_slice_assign::( + tensors, + dim, + device, + |shape, device, dtype| B::int_empty(shape, device, dtype.into()), + B::int_slice_assign, + ) + } + + /// Element-wise equality comparison. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// The boolean tensor with the result of the comparison. + fn int_equal(lhs: IntTensor, rhs: IntTensor, out_dtype: BoolDType) -> BoolTensor; + + /// Element-wise non-equality comparison. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// The boolean tensor with the result of the comparison. + fn int_not_equal(lhs: IntTensor, rhs: IntTensor, out_dtype: BoolDType) -> BoolTensor { + let equal_tensor = B::int_equal(lhs, rhs, out_dtype); + B::bool_not(equal_tensor) + } + + /// Element-wise equality comparison with a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// The boolean tensor with the result of the comparison. + fn int_equal_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor; + + /// Element-wise non-equality comparison with a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// The boolean tensor with the result of the comparison. + fn int_not_equal_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor { + let equal_tensor = B::int_equal_elem(lhs, rhs, out_dtype); + B::bool_not(equal_tensor) + } + + /// Element-wise greater than comparison. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// The boolean tensor with the result of the comparison. + fn int_greater(lhs: IntTensor, rhs: IntTensor, out_dtype: BoolDType) -> BoolTensor; + + /// Element-wise greater than comparison with a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// The boolean tensor with the result of the comparison. + fn int_greater_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor; + + /// Element-wise greater than or equal comparison. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// The boolean tensor with the result of the comparison. + fn int_greater_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor; + + /// Element-wise greater than or equal comparison with a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// The boolean tensor with the result of the comparison. + fn int_greater_equal_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor; + + /// Element-wise less than comparison. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// The boolean tensor with the result of the comparison. + fn int_lower(lhs: IntTensor, rhs: IntTensor, out_dtype: BoolDType) -> BoolTensor; + + /// Element-wise less than comparison with a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// The boolean tensor with the result of the comparison. + fn int_lower_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor; + + /// Element-wise less than or equal comparison. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// The boolean tensor with the result of the comparison. + fn int_lower_equal(lhs: IntTensor, rhs: IntTensor, out_dtype: BoolDType) + -> BoolTensor; + + /// Element-wise less than or equal comparison with a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// The boolean tensor with the result of the comparison. + fn int_lower_equal_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor; + + // ==== NUMERIC ==== // + + /// Element-wise addition. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// + /// # Returns + /// + /// The result of the addition. + fn int_add(lhs: IntTensor, rhs: IntTensor) -> IntTensor; + + /// Element-wise addition with a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The result of the addition. + fn int_add_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor; + + /// Element-wise power with a IntTensor. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side IntTensor. + /// * `rhs` - The right-hand side IntTensor. + /// + /// # Returns + /// + /// The elements of `lhs` raised to the power of the elements of `rhs`. + fn int_powi(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + let dtype = lhs.dtype(); + let float_dtype = get_device_settings::(&lhs.device()).float_dtype; + B::float_into_int( + B::float_powi(B::int_into_float(lhs, float_dtype), rhs), + dtype.into(), + ) + } + + /// Element-wise power with a scalar. + /// + /// # Backend Implementors Note + /// + /// A number of common exponent cases can be implemented with operations + /// which are much cheaper than generic exponentiation. + /// + /// This (`Backend` impl overridable) operation handles generic optimizations + /// for several common integer exponent cases; and then dispatches to + /// the (`Backend` impl overridable) [`Self::int_powi_scalar_impl`] + /// operation to handle the generic case. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The elements of `lhs` raised to the value of `rhs`. + fn int_powi_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let exp = rhs.elem::(); + match exp { + 0 => Self::int_ones(lhs.shape(), &lhs.device(), lhs.dtype().into()), + 1 => lhs, + 2 => Self::int_mul(lhs.clone(), lhs), + _ => Self::int_powi_scalar_impl(lhs, rhs), + } + } + + /// Element-wise power with a scalar. + /// + /// # Backend Implementors Note + /// + /// This is the generic implementation of integer exponentiation + /// called by [`Self::int_powi_scalar`] in the fallback case. + /// + /// By default, this performs a relatively expensive conversion to float, + /// exponentiation in float, and conversion back to int. + /// This reduces the minimal operation set for `Backend`s, + /// at the cost of performance. + /// + /// This is a good target for specialized optimizations in `Backend` implementations. + /// + /// As a general rule, this should not be called directly. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The elements of `lhs` raised to the value of `rhs`. + fn int_powi_scalar_impl(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let dtype = lhs.dtype(); + let float_dtype = get_device_settings::(&lhs.device()).float_dtype; + B::float_into_int( + B::float_powi_scalar_impl(B::int_into_float(lhs, float_dtype), rhs), + dtype.into(), + ) + } + + /// Clamps a tensor under a minimum value. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to clamp. + /// * `min` - The minimum value. + /// + /// # Returns + /// + /// The clamped tensor. + fn int_clamp_min(tensor: IntTensor, min: Scalar) -> IntTensor { + let dtype = get_device_settings::(&tensor.device()).bool_dtype; + let mask = Self::int_lower_elem(tensor.clone(), min, dtype); + Self::int_mask_fill(tensor, mask, min) + } + + /// Clamps a tensor over a maximum value. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to clamp. + /// * `max` - The maximum value. + /// + /// # Returns + /// + /// The clamped tensor. + fn int_clamp_max(tensor: IntTensor, max: Scalar) -> IntTensor { + let dtype = get_device_settings::(&tensor.device()).bool_dtype; + let mask = Self::int_greater_elem(tensor.clone(), max, dtype); + Self::int_mask_fill(tensor, mask, max) + } + + /// Clamps a tensor between a minimum and maximum value. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to clamp. + /// * `min` - The minimum value. + /// * `max` - The maximum value. + /// + /// # Returns + /// + /// The clamped tensor. + fn int_clamp(tensor: IntTensor, min: Scalar, max: Scalar) -> IntTensor { + Self::int_clamp_min(Self::int_clamp_max(tensor, max), min) + } + + /// Element-wise subtraction. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// + /// # Returns + /// + /// The result of the subtraction. + fn int_sub(lhs: IntTensor, rhs: IntTensor) -> IntTensor; + + /// Element-wise subtraction with a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The result of the subtraction. + fn int_sub_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor; + + /// Element-wise multiplication. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// + /// # Returns + /// + /// The result of the multiplication. + fn int_mul(lhs: IntTensor, rhs: IntTensor) -> IntTensor; + + /// Element-wise multiplication with a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The result of the multiplication. + fn int_mul_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor; + + /// Element-wise division. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// + /// # Returns + /// + /// The result of the division. + fn int_div(lhs: IntTensor, rhs: IntTensor) -> IntTensor; + + /// Element-wise division with a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The result of the division. + fn int_div_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor; + + /// Element-wise modulus. + /// + /// # Arguments + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The result of applying the modulus of the scalar to the tensor. + fn int_remainder(lhs: IntTensor, rhs: IntTensor) -> IntTensor; + + /// Element-wise modulus with a scalar. + /// + /// # Arguments + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The result of applying the modulus of the scalar to the tensor. + fn int_remainder_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor; + + /// Multiplies two tensors together using matrix multiplication. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// + /// # Returns + /// + /// The result of multiplying the two tensors together using matrix multiplication. + fn int_matmul(lhs: IntTensor, rhs: IntTensor) -> IntTensor; + + /// Element-wise negation. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to negate. + /// + /// # Returns + /// + /// The negated tensor. + fn int_neg(tensor: IntTensor) -> IntTensor { + Self::int_mul_scalar(tensor, (-1).into()) + } + + /// Creates a tensor of zeros. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `device` - The device to create the tensor on. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The tensor of zeros. + fn int_zeros(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + Self::int_from_data(TensorData::full_dtype(shape, 0, dtype.into()), device) + } + + /// Creates a tensor of ones. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `device` - The device to create the tensor on. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The tensor of ones. + fn int_ones(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + Self::int_from_data(TensorData::full_dtype(shape, 1, dtype.into()), device) + } + + /// Creates a tensor filled with given value. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `fill_value` - The value with which to fill the tensor. + /// * `device` - The device to create the tensor on. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The tensor filled with given value + fn int_full( + shape: Shape, + fill_value: Scalar, + device: &Device, + dtype: IntDType, + ) -> IntTensor { + Self::int_from_data( + TensorData::full_dtype(shape, fill_value, dtype.into()), + device, + ) + } + + /// Sums all elements in the tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to sum. + /// + /// # Returns + /// + /// The sum of all elements in the tensor. + fn int_sum(tensor: IntTensor) -> IntTensor; + + /// Sums all elements in the tensor along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to sum. + /// * `dim` - The dimension to sum along. + /// + /// # Returns + /// + /// The sum of all elements in the tensor along the dimension. + fn int_sum_dim(tensor: IntTensor, dim: usize) -> IntTensor; + + /// Computes the product of all elements in the tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the product of. + /// + /// # Returns + /// + /// The product of all elements in the tensor. + fn int_prod(tensor: IntTensor) -> IntTensor; + + /// Computes the product of all elements in the tensor along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the product of. + /// * `dim` - The dimension to compute the product along. + /// + /// # Returns + /// + /// The product of all elements in the tensor along the dimension. + fn int_prod_dim(tensor: IntTensor, dim: usize) -> IntTensor; + + /// Computes the mean of all elements in the tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the mean of. + /// + /// # Returns + /// + /// The mean of all elements in the tensor. + fn int_mean(tensor: IntTensor) -> IntTensor { + let num_elems = tensor.shape().num_elements() as i64; + B::int_div_scalar(B::int_sum(tensor), num_elems.into()) + } + + /// Computes the mean of all elements in the tensor along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the mean of. + /// + /// # Returns + /// + /// The mean of all elements in the tensor along the dimension. + fn int_mean_dim(tensor: IntTensor, dim: usize) -> IntTensor; + + /// Computes the cumulative sum of elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the cumulative sum of. + /// * `dim` - The dimension along which to compute the cumulative sum. + /// + /// # Returns + /// + /// A tensor with the same shape where each element is the cumulative sum + /// of all elements up to and including that position along the dimension. + fn int_cumsum(tensor: IntTensor, dim: usize) -> IntTensor; + + /// Computes the cumulative product of elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the cumulative product of. + /// * `dim` - The dimension along which to compute the cumulative product. + /// + /// # Returns + /// + /// A tensor with the same shape where each element is the cumulative product + /// of all elements up to and including that position along the dimension. + fn int_cumprod(tensor: IntTensor, dim: usize) -> IntTensor; + + /// Computes the cumulative minimum of elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the cumulative minimum of. + /// * `dim` - The dimension along which to compute the cumulative minimum. + /// + /// # Returns + /// + /// A tensor with the same shape where each element is the minimum + /// of all elements up to and including that position along the dimension. + fn int_cummin(tensor: IntTensor, dim: usize) -> IntTensor; + + /// Computes the cumulative maximum of elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the cumulative maximum of. + /// * `dim` - The dimension along which to compute the cumulative maximum. + /// + /// # Returns + /// + /// A tensor with the same shape where each element is the maximum + /// of all elements up to and including that position along the dimension. + fn int_cummax(tensor: IntTensor, dim: usize) -> IntTensor; + + /// Gets the indices of the maximum elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum indices of. + /// * `dim` - The dimension to get the maximum indices along. + /// + /// # Returns + /// + /// The indices of the maximum elements along the dimension. + fn int_argmax(tensor: IntTensor, dim: usize) -> IntTensor; + + /// Gets the indices of the k maximum elements along a dimension. + /// If two elements share the same value, it will be ordered by the lowest + /// coordinate + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum indices of. + /// * `dim` - The dimension to get the maximum indices along. + /// * `k` - number of maximum elements. + /// + /// # Returns + /// + /// The indices of the maximum elements along the dimension. + fn int_argtopk(tensor: IntTensor, dim: usize, k: usize) -> IntTensor; + + /// Gets the values of the k maximum elements along a dimension. + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum values of. + /// * `dim` - The dimension to get the maximum values along. + /// * `k` - number of maximum elements. + /// + /// # Returns + /// + /// The values of the maximum elements along the dimension. + fn int_topk(tensor: IntTensor, dim: usize, k: usize) -> IntTensor { + let device = &tensor.device(); + let dtype = get_device_settings::(device).int_dtype; + let k_indices = Self::int_arange(0..k as i64, device, dtype); + Self::int_select(Self::int_sort(tensor, dim, true), dim, k_indices) + } + + /// Gets the indices of the minimum elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the minimum indices of. + /// * `dim` - The dimension to get the minimum indices along. + /// + /// # Returns + /// + /// The indices of the minimum elements along the dimension. + fn int_argmin(tensor: IntTensor, dim: usize) -> IntTensor; + + /// Gets the maximum element in the tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum element of. + /// + /// # Returns + /// + /// The maximum element in the tensor. + fn int_max(tensor: IntTensor) -> IntTensor { + let shape = tensor.shape(); + let tensor = B::int_reshape(tensor, Shape::new([shape.num_elements()])); + + B::int_max_dim(tensor, 0) + } + + /// Gets the maximum element in the tensor along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum element of. + /// * `dim` - The dimension to get the maximum element along. + /// + /// # Returns + /// + /// The maximum element in the tensor along the dimension. + fn int_max_dim(tensor: IntTensor, dim: usize) -> IntTensor { + let index = B::int_argmax(tensor.clone(), dim); + B::int_gather(dim, tensor, index) + } + + /// Gets the maximum elements and corresponding indices along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements and indices of. + /// * `dim` - The dimension to get the maximum elements and indices along. + /// + /// # Returns + /// + /// The maximum elements and corresponding indices along the dimension. + fn int_max_dim_with_indices(tensor: IntTensor, dim: usize) -> (IntTensor, IntTensor) { + let index = B::int_argmax(tensor.clone(), dim); + let values = B::int_gather(dim, tensor, index.clone()); + + (values, index) + } + + /// Gets the maximum absolute element in the tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum element of. + /// + /// # Returns + /// + /// The maximum element in the tensor. + fn int_max_abs(tensor: IntTensor) -> IntTensor { + let shape = tensor.shape(); + let tensor = B::int_reshape(tensor, Shape::new([shape.num_elements()])); + + B::int_max_abs_dim(tensor, 0) + } + + /// Gets the maximum absolute element in the tensor along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum element of. + /// * `dim` - The dimension to get the maximum element along. + /// + /// # Returns + /// + /// The maximum element in the tensor along the dimension. + fn int_max_abs_dim(tensor: IntTensor, dim: usize) -> IntTensor { + B::int_max_dim(B::int_abs(tensor), dim) + } + + /// Gets the minimum element in the tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the minimum element of. + /// + /// # Returns + /// + /// The minimum element in the tensor. + fn int_min(tensor: IntTensor) -> IntTensor { + let shape = tensor.shape(); + let tensor = B::int_reshape(tensor, Shape::new([shape.num_elements()])); + + B::int_min_dim(tensor, 0) + } + + /// Gets the minimum elements in the tensor along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the minimum element of. + /// * `dim` - The dimension to get the minimum element along. + /// + /// # Returns + /// + /// The minimum element in the tensor along the dimension. + fn int_min_dim(tensor: IntTensor, dim: usize) -> IntTensor { + let index = B::int_argmin(tensor.clone(), dim); + B::int_gather(dim, tensor, index) + } + + /// Gets the minimum elements and corresponding indices along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the minimum elements and indices of. + /// * `dim` - The dimension to get the minimum elements and indices along. + /// + /// # Returns + /// + /// The minimum elements and corresponding indices along the dimension. + fn int_min_dim_with_indices(tensor: IntTensor, dim: usize) -> (IntTensor, IntTensor) { + let indices = B::int_argmin(tensor.clone(), dim); + let values = B::int_gather(dim, tensor, indices.clone()); + + (values, indices) + } + + /// Returns a new tensor with absolute values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take absolute value of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with absolute values. + fn int_abs(tensor: IntTensor) -> IntTensor; + + /// Transposes an int tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to transpose. + /// + /// # Returns + /// + /// The transposed tensor. + fn int_transpose(tensor: IntTensor) -> IntTensor { + let ndims = tensor.shape().num_dims(); + Self::int_swap_dims(tensor, ndims - 2, ndims - 1) + } + + /// Swaps two dimensions of an int tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to swap the dimensions of. + /// * `dim1` - The first dimension to swap. + /// * `dim2` - The second dimension to swap. + /// + /// # Returns + /// + /// The tensor with the dimensions swapped. + fn int_swap_dims(tensor: IntTensor, dim1: usize, dim2: usize) -> IntTensor; + + /// Permutes the dimensions of a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to permute the dimensions of. + /// * `axes` - The new order of the dimensions. + /// # Returns + /// + /// The tensor with the dimensions permuted. + fn int_permute(tensor: IntTensor, axes: &[usize]) -> IntTensor; + + /// Reverse the order of elements in a tensor along the given axes. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to reverse. + /// * `axes` - The axes to reverse. + /// + /// The tensor with the elements reversed. + fn int_flip(tensor: IntTensor, axes: &[usize]) -> IntTensor; + + /// Creates a new int tensor with random values. + /// + /// # Arguments + /// * `shape` - The shape of the tensor. + /// * `distribution` - The distribution to sample from. + /// * `device` - The device to create the tensor on. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The tensor with the given shape and random values. + fn int_random( + shape: Shape, + distribution: Distribution, + device: &Device, + dtype: IntDType, + ) -> IntTensor; + + /// Creates a new tensor with values from the given range with the given step size. + /// + /// # Arguments + /// + /// * `range` - The range of values. + /// * `step` - The step size. + /// * `device` - The device to create the tensor on. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The tensor with the given values. + fn int_arange_step( + range: Range, + step: usize, + device: &Device, + dtype: IntDType, + ) -> IntTensor { + let value = range.step_by(step).collect::>(); + let shape = Shape::new([value.len()]); + let data = TensorData::new(value, shape).convert_dtype(dtype.into()); + B::int_from_data(data, device) + } + + /// Creates a new tensor with values from the given range. + /// + /// # Arguments + /// + /// * `range` - The range of values. + /// * `device` - The device to create the tensor on. + /// + /// # Returns + /// + /// The tensor with the given values. + /// + /// # Remarks + /// + /// Uses `arange_step` with a step size of 1 under the hood. + fn int_arange(range: Range, device: &Device, dtype: IntDType) -> IntTensor { + Self::int_arange_step(range, 1, device, dtype) + } + + /// Tests if any element in the int `tensor` evaluates to True. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// + /// # Returns + /// + /// A boolean tensor with a single element, True if any element in the tensor is True, False otherwise. + fn int_any(tensor: IntTensor, out_dtype: BoolDType) -> BoolTensor { + let int_dtype = tensor.dtype(); + let bool_tensor = B::int_equal_elem(tensor, 0.into(), out_dtype); + let bool_tensor = B::bool_not(bool_tensor); + let sum = B::int_sum(B::bool_into_int(bool_tensor, int_dtype.into())); + B::int_greater_elem(sum, 0.into(), out_dtype) + } + + /// Tests if any element in the int `tensor` evaluates to True along a given dimension `dim`. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// * `dim` - The axis along which to test. + /// + /// # Returns + /// + /// A boolean tensor `Tensor` with the same size as input `tensor`, except in the `dim` axis + /// where the size is 1. The elem in the `dim` axis is True if any element along this dim in the input + /// evaluates to True, False otherwise. + fn int_any_dim(tensor: IntTensor, dim: usize, out_dtype: BoolDType) -> BoolTensor { + let int_dtype = tensor.dtype(); + let bool_tensor = B::int_equal_elem(tensor, 0.into(), out_dtype); + let bool_tensor = B::bool_not(bool_tensor); + let sum = B::int_sum_dim(B::bool_into_int(bool_tensor, int_dtype.into()), dim); + B::int_greater_elem(sum, 0.into(), out_dtype) + } + + /// Tests if all elements in the int `tensor` evaluate to True. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor `Tensor` with a single element, True if all elements in the input tensor + /// evaluate to True, False otherwise. + fn int_all(tensor: IntTensor, out_dtype: BoolDType) -> BoolTensor { + let int_dtype = tensor.dtype(); + let num_elems = tensor.shape().num_elements() as i64; + let bool_tensor = B::int_equal_elem(tensor, 0.into(), out_dtype); + let bool_tensor = B::bool_not(bool_tensor); + let sum = B::int_sum(B::bool_into_int(bool_tensor, int_dtype.into())); + B::int_equal_elem(sum, num_elems.into(), out_dtype) + } + + /// Tests if all elements in the int `tensor` evaluate to True along a given dimension `dim`. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// * `dim` - The axis along which to test. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor `Tensor` with the same size as input `tensor`, except in the `dim` axis + /// where the size is 1. The elem in the `dim` axis is True if all elements along this dim in the input + /// evaluates to True, False otherwise. + fn int_all_dim(tensor: IntTensor, dim: usize, out_dtype: BoolDType) -> BoolTensor { + let int_dtype = tensor.dtype(); + let num_elems = tensor.shape()[dim] as i64; + let bool_tensor = B::int_equal_elem(tensor, 0.into(), out_dtype); + let bool_tensor = B::bool_not(bool_tensor); + let sum = B::int_sum_dim(B::bool_into_int(bool_tensor, int_dtype.into()), dim); + B::int_equal_elem(sum, num_elems.into(), out_dtype) + } + + /// Returns the signs of the int `tensor`. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to extract the signs from. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` containing the signs of the elements of `tensor`. + fn int_sign(tensor: IntTensor) -> IntTensor { + let dtype = tensor.dtype(); + let device = &tensor.device(); + let bool_dtype = get_device_settings::(&tensor.device()).bool_dtype; + let zeros = B::int_zeros(tensor.shape(), device, dtype.into()); + let less_than_zero = B::int_lower_elem(tensor.clone(), 0.into(), bool_dtype); + let greater_than_zero = B::int_greater_elem(tensor, 0.into(), bool_dtype); + + let mut result = B::int_mask_fill(zeros, less_than_zero, (-1).into()); + result = B::int_mask_fill(result, greater_than_zero, 1.into()); + result + } + + /// Broadcasts the int `tensor` to the given `shape`. + fn int_expand(tensor: IntTensor, shape: Shape) -> IntTensor; + + /// Sort the elements of the input `tensor` by value along a given dimension. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// * `dim` - The axis along which to sort. + /// * `descending` - The sorting order. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor, where the elements are sorted by value. + fn int_sort(tensor: IntTensor, dim: usize, descending: bool) -> IntTensor { + let device = tensor.device(); + sort::( + tensor, + dim, + descending, + device, + |tensor| { + let msg = "Failed to synchronously read tensor data. This operation is not supported until this backend has a GPU sorting implementation."; + try_read_sync(B::int_into_data(tensor)) + .expect(msg) + .expect(msg) + }, + |data, device, _dtype| B::int_from_data(data, device), + ) + } + + /// Sort the elements of the input `tensor` by value along a given dimension. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// * `dim` - The axis along which to sort. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor and corresponding indices, where + /// the elements are sorted by value and the indices map back to the original input tensor. + fn int_sort_with_indices( + tensor: IntTensor, + dim: usize, + descending: bool, + ) -> (IntTensor, IntTensor) { + let dtype = tensor.dtype(); + let device = tensor.device(); + sort_with_indices::( + tensor, + dim, + descending, + dtype.into(), + device, + |tensor| { + let msg = "Failed to synchronously read tensor data. This operation is not supported until this backend has a GPU sorting implementation."; + try_read_sync(B::int_into_data(tensor)) + .expect(msg) + .expect(msg) + }, + |data, device, _dtype| B::int_from_data(data, device), + ) + } + + /// Returns the indices that sort the elements of the input `tensor` by value + /// along a given dimension. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// * `dim` - The axis along which to sort. + /// * `descending` - The sorting order. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor the indices map back to the original input tensor. + fn int_argsort(tensor: IntTensor, dim: usize, descending: bool) -> IntTensor { + let dtype = tensor.dtype(); + let device = tensor.device(); + argsort::(tensor, dim, descending, dtype.into(), device, |tensor| { + let msg = "Failed to synchronously read tensor data. This operation is not supported until this backend has a GPU sorting implementation."; + try_read_sync(B::int_into_data(tensor)) + .expect(msg) + .expect(msg) + }) + } + + /// Bitwise AND operation for Int Tensors + fn bitwise_and(lhs: IntTensor, rhs: IntTensor) -> IntTensor; + + /// Bitwise AND operation for Int Tensors with a scalar + fn bitwise_and_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor; + + /// Bitwise OR operation for Int Tensors + fn bitwise_or(lhs: IntTensor, rhs: IntTensor) -> IntTensor; + + /// Bitwise OR operation for Int Tensors with a scalar + fn bitwise_or_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor; + + /// Bitwise XOR operation for Int Tensors + fn bitwise_xor(lhs: IntTensor, rhs: IntTensor) -> IntTensor; + + /// Bitwise XOR operation for Int Tensors with a scalar + fn bitwise_xor_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor; + + /// Bitwise NOT operation for Int Tensors + fn bitwise_not(tensor: IntTensor) -> IntTensor; + + /// Bitwise left shift operation for Int Tensors + fn bitwise_left_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor; + + /// Bitwise left shift operation for Int Tensors with a scalar + fn bitwise_left_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor; + + /// Bitwise right shift operation for Int Tensors + fn bitwise_right_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor; + + /// Bitwise right shift operation for Int Tensors with a scalar + fn bitwise_right_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor; + + /// Converts a tensor to another integer data type. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to convert. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// A tensor with the same values as `tensor` but in the target integer data type. + fn int_cast(tensor: IntTensor, dtype: IntDType) -> IntTensor; + + /// Unfold windows along a dimension. + /// + /// Returns a view of the tensor with all complete windows of size `size` in dimension `dim`; + /// where windows are advanced by `step` at each index. + /// + /// The number of windows is `max(0, (shape[dim] - size).ceil_div(step))`. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor to unfold; of shape ``[pre=..., dim shape, post=...]`` + /// * `dim` - the selected dim. + /// * `size` - the size of each unfolded window. + /// * `step` - the step between each window. + /// + /// # Returns + /// + /// A tensor view with shape ``[pre=..., windows, size, post=...]``. + fn int_unfold(tensor: IntTensor, dim: usize, size: usize, step: usize) -> IntTensor; +} diff --git a/crates/burn-backend/src/backend/ops/mod.rs b/crates/burn-backend/src/backend/ops/mod.rs new file mode 100644 index 0000000..0485608 --- /dev/null +++ b/crates/burn-backend/src/backend/ops/mod.rs @@ -0,0 +1,20 @@ +mod activation; +mod bool_tensor; +mod int_tensor; +mod modules; +mod qtensor; +mod tensor; +mod transaction; + +pub(crate) mod argwhere; +pub(crate) mod cat; +pub(crate) mod repeat_dim; +pub(crate) mod sort; + +pub use activation::*; +pub use bool_tensor::*; +pub use int_tensor::*; +pub use modules::*; +pub use qtensor::*; +pub use tensor::*; +pub use transaction::*; diff --git a/crates/burn-backend/src/backend/ops/modules/attention.rs b/crates/burn-backend/src/backend/ops/modules/attention.rs new file mode 100644 index 0000000..1740af0 --- /dev/null +++ b/crates/burn-backend/src/backend/ops/modules/attention.rs @@ -0,0 +1,116 @@ +use core::f32; +#[allow(unused_imports)] +use num_traits::Float as _; + +use burn_std::Shape; + +use crate::{ + Backend, TensorMetadata, get_device_settings, + ops::AttentionModuleOptions, + tensor::{BoolTensor, FloatTensor}, +}; + +/// Computes softmax(QKᵗ * scale) · V using separate kernels. +/// Serves as a fallback when FlashAttention is not used. +pub fn attention_fallback( + query: FloatTensor, + key: FloatTensor, + value: FloatTensor, + mask: Option>, + attn_bias: Option>, + options: AttentionModuleOptions, +) -> FloatTensor { + if let Some(softcap) = options.softcap { + assert!(softcap > 0.0, "softcap must be positive, got {softcap}"); + } + + // Attention scores: A = QKᵗ * scale + let query_shape = query.shape().dims::<4>(); + let scale = options + .scale + .unwrap_or_else(|| 1.0 / (*query_shape.last().unwrap() as f64).sqrt()); + let transposed_key = B::float_transpose(key); + let qk = B::float_matmul(query, transposed_key); + let attention_scores = B::float_mul_scalar(qk, scale.into()); + + // Softcap: softcap * tanh(scores / softcap) + // Applied to raw logits before any -inf masking, so that tanh does not + // map -inf to a finite value (which would break masking semantics). + let attention_scores = if let Some(softcap) = options.softcap { + let scaled = B::float_div_scalar(attention_scores, softcap.into()); + let tanh = B::float_tanh(scaled); + B::float_mul_scalar(tanh, softcap.into()) + } else { + attention_scores + }; + + // Bool masking + let attention_scores = if let Some(mask) = mask { + B::float_mask_fill(attention_scores, mask, f32::NEG_INFINITY.into()) + } else { + attention_scores + }; + + // Causal masking: mask positions where col > row (future positions) + let attention_scores = if options.is_causal { + let causal_mask = build_causal_mask::(&attention_scores); + B::float_mask_fill(attention_scores, causal_mask, f32::NEG_INFINITY.into()) + } else { + attention_scores + }; + + // Additive bias (ALiBi, relative position biases, etc.) + let attention_scores = if let Some(bias) = attn_bias { + B::float_add(attention_scores, bias) + } else { + attention_scores + }; + + // NaN-safe softmax: S = softmax(A) + // When all positions in a row are masked (-inf), naive softmax has two NaN paths: + // (1) max is -inf, so the shift -inf - (-inf) = NaN; + // (2) after fixing (1), all exp values are 0, so sum is 0 and 0/0 = NaN. + // Clamping max to finfo.min (most negative finite value) and sum to finfo.min_positive + // (smallest positive normal) avoids both, yielding 0 for fully-masked rows. + let finfo = attention_scores.dtype().finfo().expect("float tensor"); + let max_per_dim = B::float_max_dim(attention_scores.clone(), 3); + let max_per_dim = B::float_clamp_min(max_per_dim, finfo.min.into()); + let minus_max = B::float_sub(attention_scores, max_per_dim); + let numerator = B::float_exp(minus_max); + let sum_exp = B::float_sum_dim(numerator.clone(), 3); + let sum_exp = B::float_clamp_min(sum_exp, finfo.min_positive.into()); + let softmax = B::float_div(numerator, sum_exp); + + // Context: S · V + B::float_matmul(softmax, value) +} + +/// Builds a causal (upper-triangular) bool mask where `true` means "mask this position". +/// Shape: [batch_size, num_heads, seq_q, seq_k], masking positions where col > row. +fn build_causal_mask(attention_scores: &FloatTensor) -> BoolTensor { + let device = attention_scores.device(); + let scores_shape = attention_scores.shape().dims::<4>(); + let [batch_size, num_heads, seq_q, seq_k] = scores_shape; + let settings = get_device_settings::(&device); + + // row indices [seq_q, 1] and col indices [1, seq_k] + // Offset col indices so that the causal boundary aligns at the bottom-right corner, + // which handles cross-attention (seq_k > seq_q) correctly. + let offset = seq_k as i64 - seq_q as i64; + let rows = B::int_reshape( + B::int_arange(0..seq_q as i64, &device, settings.int_dtype), + Shape::new([seq_q, 1]), + ); + let cols = B::int_reshape( + B::int_arange(0..seq_k as i64, &device, settings.int_dtype), + Shape::new([1, seq_k]), + ); + + // mask where col > row + offset (upper triangle) + let rows_shifted = B::int_add_scalar(rows, offset.into()); + let mask_2d = B::int_lower(rows_shifted, cols, settings.bool_dtype); + + // Reshape to [1, 1, seq_q, seq_k] then expand to [batch_size, num_heads, seq_q, seq_k] + let mask_4d = B::bool_reshape(mask_2d, Shape::new([1, 1, seq_q, seq_k])); + B::bool_expand(mask_4d, Shape::new([batch_size, num_heads, seq_q, seq_k])) +} diff --git a/crates/burn-backend/src/backend/ops/modules/base.rs b/crates/burn-backend/src/backend/ops/modules/base.rs new file mode 100644 index 0000000..7f6d51d --- /dev/null +++ b/crates/burn-backend/src/backend/ops/modules/base.rs @@ -0,0 +1,866 @@ +use super::{conv, ctc, linear, pool}; +use crate::ops::unfold::unfold4d_using_conv2d; +use crate::tensor::{BoolTensor, FloatTensor, IntTensor}; +use crate::{Backend, TensorMetadata}; +pub use burn_std::ops::{ + AttentionModuleOptions, ConvOptions, ConvTransposeOptions, DeformConvOptions, + GridSampleOptions, GridSamplePaddingMode, InterpolateMode, InterpolateOptions, PadMode, + PaddedConvOptions, UnfoldOptions, +}; +use burn_std::{IntDType, Shape}; + +/// Gradient computed during the backward pass for each tensor used by [conv2d](ModuleOps::conv2d). +#[derive(new)] +pub struct Conv2dBackward { + /// Gradient. + pub x_grad: FloatTensor, + + /// Weights gradient. + pub weights_grad: FloatTensor, + + /// Bias gradient. + pub bias_grad: Option>, +} + +/// Gradient computed during the backward pass for each tensor used by [deform_conv2d](ModuleOps::deform_conv2d). +#[derive(new)] +pub struct DeformConv2dBackward { + /// Gradient. + pub x_grad: FloatTensor, + + /// Offset gradient. + pub offset_grad: FloatTensor, + + /// Weights gradient. + pub weight_grad: FloatTensor, + + /// Mask gradient. + pub mask_grad: Option>, + + /// Bias gradient. + pub bias_grad: Option>, +} + +/// Gradient computed during the backward pass for each tensor used by [conv3d](ModuleOps::conv3d). +#[derive(new)] +pub struct Conv3dBackward { + /// Gradient. + pub x_grad: FloatTensor, + + /// Weights gradient. + pub weights_grad: FloatTensor, + + /// Bias gradient. + pub bias_grad: Option>, +} + +/// Gradient computed during the backward pass for each tensor used by [max_pool1d](ModuleOps::max_pool1d). +#[derive(new)] +pub struct MaxPool1dBackward { + /// Gradient. + pub x_grad: FloatTensor, +} + +/// Results from [max_pool1d](ModuleOps::max_pool1d_with_indices). +#[derive(new)] +pub struct MaxPool1dWithIndices { + /// The output tensor. + pub output: FloatTensor, + + /// The indices tensor. + pub indices: IntTensor, +} + +/// Gradient computed during the backward pass for each tensor used by [max_pool2d](ModuleOps::max_pool2d). +#[derive(new)] +pub struct MaxPool2dBackward { + /// Gradient. + pub x_grad: FloatTensor, +} + +/// Results from [max_pool2d](ModuleOps::max_pool2d_with_indices). +#[derive(new)] +pub struct MaxPool2dWithIndices { + /// The output tensor. + pub output: FloatTensor, + + /// The indices tensor. + pub indices: IntTensor, +} + +/// Gradient computed during the backward pass for each tensor used by [interpolate](ModuleOps::interpolate). +#[derive(new)] +pub struct InterpolateBackward { + /// Gradient. + pub x_grad: FloatTensor, +} + +/// Module operations trait. +pub trait ModuleOps { + /// Embedding operation. + /// + /// # Arguments + /// + /// * `weights` - The embedding weights. + /// * `indices` - The indices tensor. + /// + /// # Returns + /// + /// The output tensor. + fn embedding(weights: FloatTensor, indices: IntTensor) -> FloatTensor { + let [batch_size, seq_length] = indices.shape().dims(); + let [_, d_model] = weights.shape().dims(); + + let indices = B::int_reshape(indices, Shape::new([batch_size * seq_length])); + let output = B::float_select(weights, 0, indices); + + B::float_reshape(output, Shape::new([batch_size, seq_length, d_model])) + } + + /// Embedding backward operation. + /// + /// # Arguments + /// + /// * `weights` - The embedding weights. + /// * `output_grad` - The output gradient. + /// * `indices` - The indices tensor. + /// + /// # Returns + /// + /// The gradient. + fn embedding_backward( + weights: FloatTensor, + output_grad: FloatTensor, + indices: IntTensor, + ) -> FloatTensor { + let [batch_size, seq_length] = indices.shape().dims(); + let [n_embeddings, d_model] = weights.shape().dims(); + let device = weights.device(); + let dtype = output_grad.dtype(); + + let indices = B::int_reshape(indices, Shape::new([batch_size * seq_length])); + let output_grad = + B::float_reshape(output_grad, Shape::new([batch_size * seq_length, d_model])); + let grad = B::float_zeros(Shape::new([n_embeddings, d_model]), &device, dtype.into()); + + B::float_select_add(grad, 0, indices, output_grad) + } + + /// Linear transformation. + /// + /// # Shapes + /// + /// x: `[..., d_input]`, + /// weight: `[d_input, d_output]`, + /// bias: `[d_output]`, + fn linear( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + ) -> FloatTensor { + linear::linear::(x, weight, bias) + } + /// Backward pass for [linear](ModuleOps::linear), returning the gradient for `x`. + fn linear_x_backward(weight: FloatTensor, output_grad: FloatTensor) -> FloatTensor { + linear::linear_x_backward::(weight, output_grad) + } + /// Backward pass for [linear](ModuleOps::linear), returning the gradient for `weight`. + fn linear_weight_backward(x: FloatTensor, output_grad: FloatTensor) -> FloatTensor { + linear::linear_weight_backward::(x, output_grad) + } + /// Backward pass for [linear](ModuleOps::linear), returning the gradient for `bias`. + fn linear_bias_backward(output_grad: FloatTensor) -> FloatTensor { + linear::linear_bias_backward::(output_grad) + } + + /// One dimensional convolution. + /// + /// # Shapes + /// + /// x: `[batch_size, channels_in, length]`, + /// weight: `[channels_out, channels_in, kernel_size]`, + /// bias: `[channels_out]`, + fn conv1d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<1>, + ) -> FloatTensor { + conv::conv1d_from_conv2d::(x, weight, bias, options) + } + /// Backward pass for the [conv1d](ModuleOps::conv1d) operation, returning the gradient for `x`. + fn conv1d_x_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<1>, + ) -> FloatTensor { + conv::conv1d_x_backward::(x, weight, output_grad, options) + } + /// Backward pass for the [conv1d](ModuleOps::conv1d) operation, returning the gradient for `weight`. + fn conv1d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<1>, + ) -> FloatTensor { + conv::conv1d_weight_backward::(x, weight, output_grad, options) + } + /// Backward pass for the [conv1d](ModuleOps::conv1d) operation, returning the gradient for `bias`. + fn conv1d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + conv::conv1d_bias_backward::(x, bias, output_grad) + } + /// Two dimensional convolution. + /// + /// # Shapes + /// + /// x: `[batch_size, channels_in, height, width]`, + /// weight: `[channels_out, channels_in, kernel_size_1, kernel_size_2]`, + /// bias: `[channels_out]`, + fn conv2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<2>, + ) -> FloatTensor; + /// Backward pass for the [conv2d](ModuleOps::conv2d) operation, returning the gradient for `x`. + fn conv2d_x_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<2>, + ) -> FloatTensor { + conv::conv2d_x_backward::(x, weight, output_grad, options) + } + /// Backward pass for the [conv2d](ModuleOps::conv2d) operation, returning the gradient for `weight`. + fn conv2d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<2>, + ) -> FloatTensor { + conv::conv2d_weight_backward::(x, weight, output_grad, options) + } + /// Backward pass for the [conv2d](ModuleOps::conv2d) operation, returning the gradient for `bias`. + fn conv2d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + conv::conv2d_bias_backward::(x, bias, output_grad) + } + + /// Two dimensional deformable convolution. + /// + /// # Shapes + /// + /// x: `[batch_size, channels_in, height, width]`, + /// weight: `[channels_out, channels_in, kernel_size_1, kernel_size_2]`, + /// bias: `[channels_out]`, + fn deform_conv2d( + x: FloatTensor, + offset: FloatTensor, + weight: FloatTensor, + mask: Option>, + bias: Option>, + options: DeformConvOptions<2>, + ) -> FloatTensor; + /// Backward pass for the [deform_conv2d](ModuleOps::deform_conv2d) operation. + fn deform_conv2d_backward( + x: FloatTensor, + offset: FloatTensor, + weight: FloatTensor, + mask: Option>, + bias: Option>, + output_grad: FloatTensor, + options: DeformConvOptions<2>, + ) -> DeformConv2dBackward; + + /// Three dimensional convolution. + /// + /// # Shapes + /// + /// x: `[batch_size, channels_in, depth, height, width]`, + /// weight: `[channels_out, channels_in, kernel_size_1, kernel_size_2, kernel_size_3]`, + /// bias: `[channels_out]`, + fn conv3d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<3>, + ) -> FloatTensor; + /// Backward pass for the [conv3d](ModuleOps::conv3d) operation, returning the gradient for `x`. + fn conv3d_x_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<3>, + ) -> FloatTensor { + conv::conv3d_x_backward::(x, weight, output_grad, options) + } + /// Backward pass for the [conv3d](ModuleOps::conv3d) operation, returning the gradient for `weight`. + fn conv3d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<3>, + ) -> FloatTensor { + conv::conv3d_weight_backward::(x, weight, output_grad, options) + } + /// Backward pass for the [conv3d](ModuleOps::conv3d) operation, returning the gradient for `bias`. + fn conv3d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + conv::conv3d_bias_backward::(x, bias, output_grad) + } + /// One dimensional transposed convolution. + /// + /// # Shapes + /// + /// x: `[batch_size, channels_in, length]`, + /// weight: `[channels_in, channels_out, length]`, + /// bias: `[channels_out]`, + fn conv_transpose1d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<1>, + ) -> FloatTensor { + conv::conv_transpose1d_from_conv_transpose2d::(x, weight, bias, options) + } + /// Backward pass for the [conv transpose 1d](ModuleOps::conv_transpose1d) operation, returning the gradient for `x`. + fn conv_transpose1d_x_backward( + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<1>, + ) -> FloatTensor { + conv::conv_transpose1d_x_backward::(weight, output_grad, options) + } + /// Backward pass for the [conv transpose 1d](ModuleOps::conv_transpose1d) operation, returning the gradient for `weight`. + fn conv_transpose1d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<1>, + ) -> FloatTensor { + conv::conv_transpose1d_weight_backward::(x, weight, output_grad, options) + } + /// Backward pass for the [conv transpose 1d](ModuleOps::conv_transpose1d) operation, returning the gradient for `bias`. + fn conv_transpose1d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + conv::conv_transpose1d_bias_backward::(x, bias, output_grad) + } + + /// Two dimensional transposed convolution. + /// + /// # Shapes + /// + /// x: `[batch_size, channels_in, height, width]`, + /// weight: `[channels_in, channels_out, kernel_size_1, kernel_size_2]`, + /// bias: `[channels_out]`, + fn conv_transpose2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<2>, + ) -> FloatTensor; + /// Backward pass for the [conv transpose 2d](ModuleOps::conv_transpose2d) operation, returning the gradient for `x`. + fn conv_transpose2d_x_backward( + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<2>, + ) -> FloatTensor { + conv::conv_transpose2d_x_backward::(weight, output_grad, options) + } + /// Backward pass for the [conv transpose 2d](ModuleOps::conv_transpose2d) operation, returning the gradient for `weight`. + fn conv_transpose2d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<2>, + ) -> FloatTensor { + conv::conv_transpose2d_weight_backward::(x, weight, output_grad, options) + } + /// Backward pass for the [conv transpose 2d](ModuleOps::conv_transpose2d) operation, returning the gradient for `bias`. + fn conv_transpose2d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + conv::conv_transpose2d_bias_backward::(x, bias, output_grad) + } + + /// Three dimensional transposed convolution. + /// + /// # Shapes + /// + /// x: `[batch_size, channels_in, height, width]`, + /// weight: `[channels_in, channels_out, kernel_size_1, kernel_size_2, kernel_size_3]`, + /// bias: `[channels_out]`, + fn conv_transpose3d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<3>, + ) -> FloatTensor; + /// Backward pass for the [conv transpose 3d](ModuleOps::conv_transpose3d) operation, returning the gradient for `x`. + fn conv_transpose3d_x_backward( + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<3>, + ) -> FloatTensor { + conv::conv_transpose3d_x_backward::(weight, output_grad, options) + } + /// Backward pass for the [conv transpose 3d](ModuleOps::conv_transpose3d) operation, returning the gradient for `weight`. + fn conv_transpose3d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<3>, + ) -> FloatTensor { + conv::conv_transpose3d_weight_backward::(x, weight, output_grad, options) + } + /// Backward pass for the [conv transpose 3d](ModuleOps::conv_transpose3d) operation, returning the gradient for `bias`. + fn conv_transpose3d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + conv::conv_transpose3d_bias_backward::(x, bias, output_grad) + } + + /// Four-dimensional unfolding. + /// + /// # Shapes + /// + /// * x: ``[batch_size, channels_in, height, width]``, + /// * returns: ``[batch_size, channels_in * kernel_size_1 * kernel_size_2, number of blocks]``, + fn unfold4d( + x: FloatTensor, + kernel_size: [usize; 2], + options: UnfoldOptions, + ) -> FloatTensor { + if options.padding == [0, 0] && options.dilation == [1, 1] { + let blocks = B::float_unfold(x, 2, kernel_size[0], options.stride[0]); + let blocks = B::float_unfold(blocks, 3, kernel_size[1], options.stride[1]); + + // batch, channels, h_blocks, w_blocks, h_kern, w_kern + + let blocks = B::float_permute(blocks, &[0, 1, 4, 5, 2, 3]); + let shape = blocks.shape(); + + // batch, channels, h_kern, w_kern, h_blocks, w_blocks + + B::float_reshape( + blocks, + [ + shape[0], + shape[1] * shape[2] * shape[3], + shape[4] * shape[5], + ] + .into(), + ) + } else { + unfold4d_using_conv2d::(x, kernel_size, options) + } + } + + /// One dimensional avg pooling. + /// + /// # Shapes + /// + /// x: [batch_size, channels, length], + fn avg_pool1d( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + pool::avg_pool1d_from_2d::( + x, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ) + } + /// Backward pass for the [avg pooling 1d](ModuleOps::avg_pool1d) operation. + fn avg_pool1d_backward( + x: FloatTensor, + grad: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + pool::avg_pool1d_backward_from_2d::( + x, + grad, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ) + } + /// Two dimensional avg pooling. + /// + /// # Shapes + /// + /// x: [batch_size, channels, height, width], + fn avg_pool2d( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor; + /// Backward pass for the [avg pooling 2d](ModuleOps::avg_pool2d) operation. + fn avg_pool2d_backward( + x: FloatTensor, + grad: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor; + /// Two dimensional adaptive avg pooling. + /// + /// # Shapes + /// + /// x: [batch_size, channels, height, width], + fn adaptive_avg_pool2d(x: FloatTensor, output_size: [usize; 2]) -> FloatTensor; + /// Backward pass for the [adaptive avg pooling 2d](ModuleOps::adaptive_avg_pool2d) operation. + fn adaptive_avg_pool2d_backward(x: FloatTensor, grad: FloatTensor) -> FloatTensor; + /// One dimensional adaptive avg pooling. + /// + /// # Shapes + /// + /// x: [batch_size, channels, length], + fn adaptive_avg_pool1d(x: FloatTensor, output_size: usize) -> FloatTensor { + pool::adaptive_avg_pool1d_from_2d::(x, output_size) + } + /// Backward pass for the [adaptive avg pooling 1d](ModuleOps::adaptive_avg_pool1d) operation. + fn adaptive_avg_pool1d_backward(x: FloatTensor, grad: FloatTensor) -> FloatTensor { + pool::adaptive_avg_pool1d_backward_from_2d::(x, grad) + } + /// One dimensional max pooling. + /// + /// # Shapes + /// + /// x: [batch_size, channels, length], + fn max_pool1d( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + ) -> FloatTensor { + pool::max_pool1d_from_2d::(x, kernel_size, stride, padding, dilation, ceil_mode) + } + + /// One dimensional max pooling with indices. + /// + /// # Shapes + /// + /// x: [batch_size, channels, height, width], + fn max_pool1d_with_indices( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + indices_dtype: IntDType, + ) -> MaxPool1dWithIndices { + pool::max_pool1d_with_indices_from_2d::( + x, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + indices_dtype, + ) + } + /// Backward pass for the [max pooling 1d](ModuleOps::max_pool1d_with_indices) operation. + #[allow(clippy::too_many_arguments)] + fn max_pool1d_with_indices_backward( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + output_grad: FloatTensor, + indices: IntTensor, + ) -> MaxPool1dBackward { + pool::max_pool1d_with_indices_backward_from_2d::( + x, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + output_grad, + indices, + ) + } + + /// Two dimensional max pooling. + /// + /// # Shapes + /// + /// x: [batch_size, channels, height, width], + fn max_pool2d( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + ) -> FloatTensor; + + /// Two dimensional max pooling with indices. + /// + /// # Shapes + /// + /// x: [batch_size, channels, height, width], + fn max_pool2d_with_indices( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + indices_dtype: IntDType, + ) -> MaxPool2dWithIndices; + /// Backward pass for the [max pooling 2d](ModuleOps::max_pool2d_with_indices) operation. + #[allow(clippy::too_many_arguments)] + fn max_pool2d_with_indices_backward( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + output_grad: FloatTensor, + indices: IntTensor, + ) -> MaxPool2dBackward; + + /// Down/up samples the input. + /// + /// # Shapes + /// + /// x: `[batch_size, channels, height, width]`, + fn interpolate( + x: FloatTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> FloatTensor; + + /// Backward pass for the [interpolate](ModuleOps::interpolate) operation. + fn interpolate_backward( + x: FloatTensor, + grad: FloatTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> FloatTensor; + + /// Computes scaled dot-product attention: softmax(QKᵗ * scale) · V, + /// where scale defaults to 1/sqrt(head_dim). Optionally applies masking, + /// additive bias, causal masking, and softcap to the attention scores. + /// + /// # Arguments + /// - `query`: Query tensor of shape `[batch_size, num_heads, seq_len_q, head_dim]` + /// - `key`: Key tensor of shape `[batch_size, num_heads, seq_len_k, head_dim]` + /// - `value`: Value tensor of shape `[batch_size, num_heads, seq_len_k, val_dim]` + /// - `mask`: Optional boolean mask of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`, + /// where `true` indicates positions to mask (i.e. set to -inf before softmax). + /// - `attn_bias`: Optional float tensor of shape `[batch_size, num_heads, seq_len_q, seq_len_k]` + /// added to the attention scores before softmax (e.g. ALiBi, relative position biases). + /// - `options`: Additional attention options (custom scale, softcap, causal masking). + /// + /// # Returns + /// A tensor of shape `[batch_size, num_heads, seq_len_q, val_dim]` + /// representing the attended context per head. + /// + /// # Note + /// This implementation does not support dropout and is intended for inference or + /// use cases where dropout is not needed. + fn attention( + query: FloatTensor, + key: FloatTensor, + value: FloatTensor, + mask: Option>, + attn_bias: Option>, + options: AttentionModuleOptions, + ) -> FloatTensor; + + /// Applies Layer Normalization over the last dimension of the input tensor. + /// + /// Computes `(x - mean) / sqrt(var + epsilon) * gamma + beta`, where `mean` and + /// (biased) `var` are reduced over the last axis. + /// + /// # Arguments + /// + /// * `tensor` - Input tensor of shape `[..., d_model]`. + /// * `gamma` - Scale tensor of shape `[d_model]`. + /// * `beta` - Optional bias tensor of shape `[d_model]`. + /// * `epsilon` - Numerical stability term added to the variance before the square root. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor`. + fn layer_norm( + tensor: FloatTensor, + gamma: FloatTensor, + beta: Option>, + epsilon: f64, + ) -> FloatTensor { + let shape = tensor.shape(); + let rank = shape.num_dims(); + let last_dim = rank - 1; + let d_model = shape[last_dim]; + + let mean = B::float_mean_dim(tensor.clone(), last_dim); + let centered = B::float_sub(tensor, mean); + let var = B::float_mean_dim(B::float_mul(centered.clone(), centered.clone()), last_dim); + let denom = B::float_sqrt(B::float_add_scalar(var, epsilon.into())); + let normalized = B::float_div(centered, denom); + + let broadcast_dims: alloc::vec::Vec = (0..rank) + .map(|i| if i == last_dim { d_model } else { 1 }) + .collect(); + let gamma_b = B::float_reshape(gamma, Shape::from(broadcast_dims.clone())); + let scaled = B::float_mul(normalized, gamma_b); + + match beta { + Some(beta) => { + let beta_b = B::float_reshape(beta, Shape::from(broadcast_dims)); + B::float_add(scaled, beta_b) + } + None => scaled, + } + } + + /// Computes the Connectionist Temporal Classification (CTC) loss. + /// + /// Sums over all valid alignments between the input and target sequences + /// using the forward (alpha) algorithm. + /// + /// # Arguments + /// + /// * `log_probs` - Log-probabilities of shape `[T, N, C]` + /// * `targets` - Target label indices of shape `[N, S]` + /// * `input_lengths` - Actual input sequence lengths per batch element `[N]` + /// * `target_lengths` - Actual target lengths per batch element `[N]` + /// * `blank` - Index of the blank label + /// + /// # Returns + /// + /// Per-sample loss of shape `[N]` + fn ctc_loss( + log_probs: FloatTensor, + targets: IntTensor, + input_lengths: IntTensor, + target_lengths: IntTensor, + blank: usize, + ) -> FloatTensor { + ctc::ctc_loss_default::(log_probs, targets, input_lengths, target_lengths, blank) + } + + /// Returns `true` if this backend implements [ctc_loss_backward](ModuleOps::ctc_loss_backward) + /// natively. + /// + /// Autodiff queries this flag to decide between two paths: + /// - `true`: use the backend's [ctc_loss](ModuleOps::ctc_loss) and + /// [ctc_loss_backward](ModuleOps::ctc_loss_backward) directly. + /// - `false`: call [ctc::ctc_loss_default] for the forward pass; autodiff + /// then differentiates through the decomposed tensor ops. + /// + /// Backends that override `ctc_loss_backward` must also override this to + /// return `true`. + fn has_ctc_loss_backward() -> bool { + false + } + + /// Backward pass for [ctc_loss](ModuleOps::ctc_loss): gradient w.r.t. `log_probs`. + /// + /// Only called when [has_ctc_loss_backward](ModuleOps::has_ctc_loss_backward) + /// returns `true`. Backends without a native implementation should leave + /// both methods at their defaults; the gradient is computed automatically by + /// autodiff against the decomposed [ctc::ctc_loss_default] forward. + /// + /// # Arguments + /// + /// * `log_probs` - Log-probabilities of shape `[T, N, C]` + /// * `targets` - Target label indices of shape `[N, S]` + /// * `input_lengths` - Actual input sequence lengths per batch element `[N]` + /// * `target_lengths` - Actual target lengths per batch element `[N]` + /// * `grad_loss` - Upstream gradient w.r.t. the per-sample loss `[N]` + /// * `blank` - Index of the blank label + /// + /// # Returns + /// + /// Gradient w.r.t. `log_probs` of shape `[T, N, C]` + fn ctc_loss_backward( + _log_probs: FloatTensor, + _targets: IntTensor, + _input_lengths: IntTensor, + _target_lengths: IntTensor, + _grad_loss: FloatTensor, + _blank: usize, + ) -> FloatTensor { + unreachable!( + "ctc_loss_backward called on a backend whose has_ctc_loss_backward() returns false" + ) + } + + /// Real-valued FFT with optional size parameter. + /// + /// When `n` is `None`, the signal must be a power of two along `dim`, and the output has + /// `signal_len / 2 + 1` frequency bins. + /// + /// When `n` is `Some(size)`, `size` must also be a power of two. The signal is truncated + /// or zero-padded to `size` and the output has `size / 2 + 1` frequency bins. Non-power- + /// of-two sizes are currently rejected at the public API boundary; true arbitrary-`n` DFT + /// support (Bluestein's algorithm) is tracked as a follow-up. + /// + /// Returns two tensors: the real part and the imaginary part. + fn rfft( + signal: FloatTensor, + dim: usize, + n: Option, + ) -> (FloatTensor, FloatTensor); + + /// Inverse real-valued FFT with optional output size. + /// + /// When `n` is `None`, the reconstructed signal length `2 * (spectrum_size - 1)` must be + /// a power of two. + /// + /// When `n` is `Some(size)`, `size` must also be a power of two. Output has exactly + /// `size` samples. + fn irfft( + spectrum_re: FloatTensor, + spectrum_im: FloatTensor, + dim: usize, + n: Option, + ) -> FloatTensor; +} diff --git a/crates/burn-backend/src/backend/ops/modules/conv.rs b/crates/burn-backend/src/backend/ops/modules/conv.rs new file mode 100644 index 0000000..e1b52a0 --- /dev/null +++ b/crates/burn-backend/src/backend/ops/modules/conv.rs @@ -0,0 +1,1533 @@ +#![allow(clippy::single_range_in_vec_init)] +use super::{ConvOptions, ConvTransposeOptions}; +use crate::{Backend, TensorMetadata, tensor::FloatTensor}; +use burn_std::{MetadataError, Shape, Slice}; + +use alloc::{vec, vec::Vec}; +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float as _; + +/// Calculate the expected output shape `[batch_size, channels_out, spatial_dims, ..]` for a pooling operation. +pub fn calculate_pool_output_shape( + in_shape: &Shape, + kernel_size: &[usize; N], + stride: &[usize; N], + padding: &[usize; N], + dilation: &[usize; N], + ceil_mode: bool, +) -> Result { + if in_shape.rank() != N + 2 { + return Err(MetadataError::RankMismatch { + left: in_shape.rank(), + right: N + 2, + }); + } + + let mut out_shape = in_shape.clone(); + // Spatial dims + for (i, size_i) in out_shape[2..].iter_mut().enumerate() { + *size_i = calculate_pool_output_size( + kernel_size[i], + stride[i], + padding[i], + dilation[i], + *size_i, + ceil_mode, + ); + } + + Ok(out_shape) +} + +/// Calculate the expected output shape `[batch_size, channels_out, spatial_dims, ..]` for a convolution. +pub fn calculate_conv_output_shape( + in_shape: &Shape, + weight_shape: &Shape, + stride: &[usize; N], + padding: &[usize; N], + dilation: &[usize; N], +) -> Result { + if weight_shape.rank() != N + 2 { + return Err(MetadataError::RankMismatch { + left: weight_shape.rank(), + right: N + 2, + }); + } + + if in_shape.rank() != N + 2 { + return Err(MetadataError::RankMismatch { + left: in_shape.rank(), + right: N + 2, + }); + } + + let kernel_size = &weight_shape[2..]; + + let mut out_shape = in_shape.clone(); + // Spatial dims + for (i, size_i) in out_shape[2..].iter_mut().enumerate() { + *size_i = + calculate_conv_output_size(kernel_size[i], stride[i], padding[i], dilation[i], *size_i); + } + // Output channels + out_shape[1] = weight_shape[0]; + + Ok(out_shape) +} + +/// Calculate the expected output shape `[batch_size, channels_out, spatial_dims, ..]` for a transposed convolution. +pub fn calculate_conv_transpose_output_shape( + in_shape: &Shape, + weight_shape: &Shape, + stride: &[usize; N], + padding: &[usize; N], + padding_out: &[usize; N], + dilation: &[usize; N], + groups: usize, +) -> Result { + if weight_shape.rank() != N + 2 { + return Err(MetadataError::RankMismatch { + left: weight_shape.rank(), + right: N + 2, + }); + } + + if in_shape.rank() != N + 2 { + return Err(MetadataError::RankMismatch { + left: in_shape.rank(), + right: N + 2, + }); + } + + let kernel_size = &weight_shape[2..]; + + let mut out_shape = in_shape.clone(); + // Spatial dims + for (i, size_i) in out_shape[2..].iter_mut().enumerate() { + *size_i = calculate_conv_transpose_output_size( + kernel_size[i], + stride[i], + padding[i], + padding_out[i], + dilation[i], + *size_i, + ); + } + // Output channels + out_shape[1] = weight_shape[1] * groups; + + Ok(out_shape) +} + +/// Calculate the expected padding size required when applying a convolution. +pub fn calculate_conv_padding( + kernel_size: usize, + stride: usize, + size_in: usize, + size_out: usize, +) -> usize { + let kernel_size = kernel_size as f32; + let stride = stride as f32; + let size_in = size_in as f32; + let size_out = size_out as f32; + + let padding = stride * (size_out - 1.) - size_in + kernel_size; + let padding = (padding / 2.).ceil(); + + padding as usize +} + +/// Calculate the expected output size when doing a convolution operation. +pub fn calculate_conv_output_size( + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + size_in: usize, +) -> usize { + (size_in + 2 * padding - dilation * (kernel_size - 1) - 1) / stride + 1 +} + +/// Calculate the expected output sizes when doing a convolution operation. +pub fn calculate_conv_output_sizes( + kernel_size: &[usize], + stride: &[usize], + padding: &[usize], + dilation: &[usize], + size_in: &[usize], +) -> Vec { + size_in + .iter() + .enumerate() + .map(|(i, size_in)| { + calculate_conv_output_size(kernel_size[i], stride[i], padding[i], dilation[i], *size_in) + }) + .collect() +} + +/// Calculate the expected output size when doing a pooling operation. +/// +/// # Arguments +/// +/// * `kernel_size` - Size of the pooling kernel +/// * `stride` - Stride of the pooling operation +/// * `padding` - Padding applied to input +/// * `dilation` - Dilation of the pooling kernel +/// * `size_in` - Input size (height or width) +/// * `ceil_mode` - If true, use ceiling instead of floor for output size calculation. +/// This allows the last pooling window to go out-of-bounds if needed. +pub fn calculate_pool_output_size( + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + size_in: usize, + ceil_mode: bool, +) -> usize { + let numerator = size_in + 2 * padding - dilation * (kernel_size - 1) - 1; + if ceil_mode { + // Ceiling division: (a + b - 1) / b + numerator.div_ceil(stride) + 1 + } else { + // Floor division (default) + numerator / stride + 1 + } +} + +/// Calculate the expected output size when doing a transposed convolution operation. +pub fn calculate_conv_transpose_output_size( + kernel_size: usize, + stride: usize, + padding: usize, + padding_out: usize, + dilation: usize, + size_in: usize, +) -> usize { + (size_in - 1) * stride + (dilation * (kernel_size - 1) + 1) + padding_out - 2 * padding +} + +/// Calculate the original input size that was used for a transposed convolution. +/// This is used during the backward pass to recover the correct gradient shape. +fn calculate_conv_transpose_input_size( + kernel_size: usize, + stride: usize, + padding: usize, + padding_out: usize, + dilation: usize, + size_out: usize, +) -> usize { + // We solve the forward formula for size_in: + // size_out = (size_in - 1) * stride + (dilation * (kernel_size - 1) + 1) + padding_out - 2 * padding + (size_out + 2 * padding - dilation * (kernel_size - 1) - padding_out - 1) / stride + 1 +} + +/// Calculate the original input sizes that were used for a transposed convolution. +fn calculate_conv_transpose_input_sizes( + kernel_size: [usize; D], + stride: [usize; D], + padding: [usize; D], + padding_out: [usize; D], + dilation: [usize; D], + size_out: [usize; D], +) -> [usize; D] { + let mut res = [0; D]; + for i in 0..D { + res[i] = calculate_conv_transpose_input_size( + kernel_size[i], + stride[i], + padding[i], + padding_out[i], + dilation[i], + size_out[i], + ); + } + res +} + +/// Calculate the [1D convolution](crate::ops::ModuleOps::conv1d) backward pass, returning the gradient for `x`. +pub(crate) fn conv1d_x_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<1>, +) -> FloatTensor { + let weight_shape = weight.shape(); + + let [_batch_size, _, length_in] = x.shape().dims(); + let [_batch_size, _channels_out, length_out] = output_grad.shape().dims(); + let [_, _, kernel_size] = weight_shape.dims(); + + let padding_out = calculate_padding_out( + kernel_size, + options.stride[0], + options.padding[0], + options.dilation[0], + length_in, + length_out, + ); + + B::conv_transpose1d( + output_grad, + weight, + None, + ConvTransposeOptions::new( + options.stride, + options.padding, + [padding_out], + options.dilation, + options.groups, + ), + ) +} + +/// Calculate the [1D convolution](crate::ops::ModuleOps::conv1d) backward pass, returning the gradient for `weight`. +pub(crate) fn conv1d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<1>, +) -> FloatTensor { + let weight_dtype = weight.dtype(); + let weight_shape = weight.shape(); + let weight_device = weight.device(); + + match options.groups == 1 { + true => conv1d_weight_grad_no_groups::(x, output_grad, weight_shape, options), + false => conv1d_weight_grad_groups::( + x, + B::float_zeros(weight_shape, &weight_device, weight_dtype.into()), + output_grad, + options, + ), + } +} + +/// Calculate the [1D convolution](crate::ops::ModuleOps::conv1d) backward pass, returning the gradient for `bias`. +pub(crate) fn conv1d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, +) -> FloatTensor { + let [batch_size, _, _length_in] = x.shape().dims(); + let [_batch_size, channels_out, length_out] = output_grad.shape().dims(); + + let grad = B::float_swap_dims(output_grad, 0, 1); + let grad = B::float_reshape(grad, Shape::new([channels_out, batch_size * length_out])); + let grad = B::float_sum_dim(grad, 1); + + B::float_reshape(grad, bias.shape()) +} + +/// Calculate the [2D convolution](crate::ops::ModuleOps::conv2d) backward pass, returning the gradient for `x`. +pub(crate) fn conv2d_x_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<2>, +) -> FloatTensor { + let weight_shape = weight.shape(); + + let [_batch_size, _channels_in, height_in, width_in] = x.shape().dims(); + let [_, _, height_out, width_out] = output_grad.shape().dims(); + let [_channels_out, _, kernel_size_1, kernel_size_2] = weight_shape.dims(); + + let padding_1_out = calculate_padding_out( + kernel_size_1, + options.stride[0], + options.padding[0], + options.dilation[0], + height_in, + height_out, + ); + let padding_2_out = calculate_padding_out( + kernel_size_2, + options.stride[1], + options.padding[1], + options.dilation[1], + width_in, + width_out, + ); + + B::conv_transpose2d( + output_grad, + weight, + None, + ConvTransposeOptions::new( + options.stride, + options.padding, + [padding_1_out, padding_2_out], + options.dilation, + options.groups, + ), + ) +} + +/// Calculate the [2D convolution](crate::ops::ModuleOps::conv2d) backward pass, returning the gradient for `weight`. +pub(crate) fn conv2d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<2>, +) -> FloatTensor { + let weight_dtype = weight.dtype(); + let weight_shape = weight.shape(); + let weight_device = weight.device(); + + match options.groups == 1 { + true => conv2d_weight_grad_no_groups::(x, output_grad, weight_shape, options), + false => conv2d_weight_grad_groups::( + x, + B::float_zeros(weight_shape, &weight_device, weight_dtype.into()), + output_grad, + options, + ), + } +} + +/// Calculate the [2D convolution](crate::ops::ModuleOps::conv2d) backward pass, returning the gradient for `bias`. +pub(crate) fn conv2d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, +) -> FloatTensor { + let [batch_size, _, _, _] = x.shape().dims(); + let [_, channels_out, height_out, width_out] = output_grad.shape().dims(); + + let grad = B::float_swap_dims(output_grad, 0, 1); + let grad = B::float_reshape( + grad, + Shape::new([channels_out, batch_size * height_out * width_out]), + ); + let grad = B::float_sum_dim(grad, 1); + + B::float_reshape(grad, bias.shape()) +} + +/// Calculate the [3D convolution](crate::ops::ModuleOps::conv3d) backward pass, returning the gradient for `x`. +pub(crate) fn conv3d_x_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<3>, +) -> FloatTensor { + let weight_shape = weight.shape(); + + let [_batch_size, _channels_in, depth_in, height_in, width_in] = x.shape().dims(); + let [_, _, depth_out, height_out, width_out] = output_grad.shape().dims(); + let [ + _channels_out, + _, + kernel_size_1, + kernel_size_2, + kernel_size_3, + ] = weight_shape.dims(); + + let padding_1_out = calculate_padding_out( + kernel_size_1, + options.stride[0], + options.padding[0], + options.dilation[0], + depth_in, + depth_out, + ); + let padding_2_out = calculate_padding_out( + kernel_size_2, + options.stride[1], + options.padding[1], + options.dilation[1], + height_in, + height_out, + ); + let padding_3_out = calculate_padding_out( + kernel_size_3, + options.stride[2], + options.padding[2], + options.dilation[2], + width_in, + width_out, + ); + + B::conv_transpose3d( + output_grad, + weight, + None, + ConvTransposeOptions::new( + options.stride, + options.padding, + [padding_1_out, padding_2_out, padding_3_out], + options.dilation, + options.groups, + ), + ) +} + +/// Calculate the [3D convolution](crate::ops::ModuleOps::conv3d) backward pass, returning the gradient for `weight`. +pub(crate) fn conv3d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<3>, +) -> FloatTensor { + let weight_dtype = weight.dtype(); + let weight_shape = weight.shape(); + let weight_device = weight.device(); + + match options.groups == 1 { + true => conv3d_weight_grad_no_groups::(x, output_grad, weight_shape, options), + false => conv3d_weight_grad_groups::( + x, + B::float_zeros(weight_shape, &weight_device, weight_dtype.into()), + output_grad, + options, + ), + } +} + +/// Calculate the [3D convolution](crate::ops::ModuleOps::conv3d) backward pass, returning the gradient for `bias`. +pub(crate) fn conv3d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, +) -> FloatTensor { + let [batch_size, _channels_in, _depth_in, _height_in, _width_in] = x.shape().dims(); + let [_, channels_out, depth_out, height_out, width_out] = output_grad.shape().dims(); + + let grad = B::float_swap_dims(output_grad, 0, 1); + let grad = B::float_reshape( + grad, + Shape::new([ + channels_out, + batch_size * depth_out * height_out * width_out, + ]), + ); + let grad = B::float_sum_dim(grad, 1); + + B::float_reshape(grad, bias.shape()) +} + +/// Calculate the [1D convolution transpose](crate::ops::ModuleOps::conv_transpose1d) backward pass, returning the gradient for `x`. +pub(crate) fn conv_transpose1d_x_backward( + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<1>, +) -> FloatTensor { + let [batch_size, _c_out, out_length] = output_grad.shape().dims(); + let [c_in, _c_out_groups, kernel_size] = weight.shape().dims(); + + let grad = B::conv1d( + output_grad, + weight, + None, + ConvOptions::new( + options.stride, + options.padding, + options.dilation, + options.groups, + ), + ); + + if options.padding_out[0] == 0 { + return grad; + } + + let exp_length = calculate_conv_transpose_input_size( + kernel_size, + options.stride[0], + options.padding[0], + options.padding_out[0], + options.dilation[0], + out_length, + ); + + B::float_slice( + grad, + &[ + Slice::from(0..batch_size), + Slice::from(0..c_in), + Slice::from(0..exp_length), + ], + ) +} + +/// Calculate the [1D convolution transpose](crate::ops::ModuleOps::conv_transpose1d) backward pass, returning the gradient for `weight`. +pub(crate) fn conv_transpose1d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<1>, +) -> FloatTensor { + let weight_dtype = weight.dtype(); + let weight_shape = weight.shape(); + let weight_device = weight.device(); + + match options.groups == 1 { + true => conv_transpose1d_weight_grad_no_groups::(x, output_grad, weight_shape, options), + false => conv_transpose1d_weight_grad_groups::( + x, + B::float_zeros(weight_shape, &weight_device, weight_dtype.into()), + output_grad, + options, + ), + } +} + +/// Calculate the [1D convolution transpose](crate::ops::ModuleOps::conv_transpose1d) backward pass, returning the gradient for `bias`. +pub(crate) fn conv_transpose1d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, +) -> FloatTensor { + let [batch_size, _channels_in, _] = x.shape().dims(); + let [_, channels_out, length_out] = output_grad.shape().dims(); + + let grad = B::float_swap_dims(output_grad, 0, 1); + let grad = B::float_reshape(grad, Shape::new([channels_out, batch_size * length_out])); + let grad = B::float_sum_dim(grad, 1); + + B::float_reshape(grad, bias.shape()) +} + +/// Calculate the [2D convolution transpose](crate::ops::ModuleOps::conv_transpose2d) backward pass, returning the gradient for `x`. +pub(crate) fn conv_transpose2d_x_backward( + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<2>, +) -> FloatTensor { + let [batch_size, _c_out, out_h, out_w] = output_grad.shape().dims(); + let [c_in, _c_out_groups, k_h, k_w] = weight.shape().dims(); + + let grad = B::conv2d( + output_grad, + weight, + None, + ConvOptions::new( + options.stride, + options.padding, + options.dilation, + options.groups, + ), + ); + + if options.padding_out[0] == 0 && options.padding_out[1] == 0 { + return grad; + } + + let [exp_h, exp_w] = calculate_conv_transpose_input_sizes( + [k_h, k_w], + options.stride, + options.padding, + options.padding_out, + options.dilation, + [out_h, out_w], + ); + + B::float_slice( + grad, + &[ + Slice::from(0..batch_size), + Slice::from(0..c_in), + Slice::from(0..exp_h), + Slice::from(0..exp_w), + ], + ) +} + +/// Calculate the [2D convolution transpose](crate::ops::ModuleOps::conv_transpose2d) backward pass, returning the gradient for `weight`. +pub(crate) fn conv_transpose2d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<2>, +) -> FloatTensor { + let weight_dtype = weight.dtype(); + let weight_shape = weight.shape(); + let weight_device = weight.device(); + + match options.groups == 1 { + true => conv_transpose2d_weight_grad_no_groups::(x, output_grad, weight_shape, options), + false => conv_transpose2d_weight_grad_groups::( + x, + B::float_zeros(weight_shape, &weight_device, weight_dtype.into()), + output_grad, + options, + ), + } +} + +/// Calculate the [2D convolution transpose](crate::ops::ModuleOps::conv_transpose2d) backward pass, returning the gradient for `bias`. +pub(crate) fn conv_transpose2d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, +) -> FloatTensor { + let [batch_size, _channels_in, _, _] = x.shape().dims(); + let [_, channels_out, height_out, width_out] = output_grad.shape().dims(); + + let grad = B::float_swap_dims(output_grad, 0, 1); + let grad = B::float_reshape( + grad, + Shape::new([channels_out, batch_size * height_out * width_out]), + ); + let grad = B::float_sum_dim(grad, 1); + + B::float_reshape(grad, bias.shape()) +} + +/// Calculate the [3D convolution transpose](crate::ops::ModuleOps::conv_transpose3d) backward pass, returning the gradient for `x`. +pub(crate) fn conv_transpose3d_x_backward( + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<3>, +) -> FloatTensor { + let [batch_size, _c_out, out_d, out_h, out_w] = output_grad.shape().dims(); + let [c_in, _c_out_groups, k_d, k_h, k_w] = weight.shape().dims(); + + let grad = B::conv3d( + output_grad, + weight, + None, + ConvOptions::new( + options.stride, + options.padding, + options.dilation, + options.groups, + ), + ); + + if options.padding_out[0] == 0 && options.padding_out[1] == 0 && options.padding_out[2] == 0 { + return grad; + } + + let [exp_d, exp_h, exp_w] = calculate_conv_transpose_input_sizes( + [k_d, k_h, k_w], + options.stride, + options.padding, + options.padding_out, + options.dilation, + [out_d, out_h, out_w], + ); + + B::float_slice( + grad, + &[ + Slice::from(0..batch_size), + Slice::from(0..c_in), + Slice::from(0..exp_d), + Slice::from(0..exp_h), + Slice::from(0..exp_w), + ], + ) +} + +/// Calculate the [3D convolution transpose](crate::ops::ModuleOps::conv_transpose3d) backward pass, returning the gradient for `weight`. +pub(crate) fn conv_transpose3d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<3>, +) -> FloatTensor { + let weight_dtype = weight.dtype(); + let weight_shape = weight.shape(); + let weight_device = weight.device(); + + match options.groups == 1 { + true => conv_transpose3d_weight_grad_no_groups::(x, output_grad, weight_shape, options), + false => conv_transpose3d_weight_grad_groups::( + x, + B::float_zeros(weight_shape, &weight_device, weight_dtype.into()), + output_grad, + options, + ), + } +} + +/// Calculate the [3D convolution transpose](crate::ops::ModuleOps::conv_transpose3d) backward pass, returning the gradient for `bias`. +pub(crate) fn conv_transpose3d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, +) -> FloatTensor { + let [batch_size, _channels_in, _, _, _] = x.shape().dims(); + let [_, channels_out, depth_out, height_out, width_out] = output_grad.shape().dims(); + + let grad = B::float_swap_dims(output_grad, 0, 1); + let grad = B::float_reshape( + grad, + Shape::new([ + channels_out, + batch_size * depth_out * height_out * width_out, + ]), + ); + let grad = B::float_sum_dim(grad, 1); + + B::float_reshape(grad, bias.shape()) +} + +/// Execute a 1D convolution using a 2D convolution. +pub(crate) fn conv1d_from_conv2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<1>, +) -> FloatTensor { + let [channels_out, _channels_in, kernel_size] = weight.shape().dims(); + let [batch_size, channels_in, length_in] = x.shape().dims(); + + let weight = B::float_reshape( + weight, + Shape::new([channels_out, channels_in / options.groups, kernel_size, 1]), + ); + let x = B::float_reshape(x, Shape::new([batch_size, channels_in, length_in, 1])); + + let tensor = B::conv2d( + x, + weight, + bias, + ConvOptions::new( + [options.stride[0], 1], + [options.padding[0], 0], + [options.dilation[0], 1], + options.groups, + ), + ); + let [batch_size, channels_out, height_out, _weight_out] = tensor.shape().dims(); + B::float_reshape(tensor, Shape::from([batch_size, channels_out, height_out])) +} + +/// Execute a 1D transposed convolution using a 2D transposed convolution. +pub(crate) fn conv_transpose1d_from_conv_transpose2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<1>, +) -> FloatTensor { + let [channels_in, channels_out, kernel_size] = weight.shape().dims(); + let [batch_size, _channels_in, length_in] = x.shape().dims(); + + let weight = B::float_reshape( + weight, + Shape::new([channels_in, channels_out, kernel_size, 1]), + ); + let x = B::float_reshape(x, Shape::new([batch_size, channels_in, length_in, 1])); + + let tensor = B::conv_transpose2d( + x, + weight, + bias, + ConvTransposeOptions::new( + [options.stride[0], 1], + [options.padding[0], 0], + [options.padding_out[0], 0], + [options.dilation[0], 1], + options.groups, + ), + ); + let [batch_size, channels_out, height_out, _weight_out] = tensor.shape().dims(); + B::float_reshape(tensor, Shape::from([batch_size, channels_out, height_out])) +} + +fn conv1d_weight_grad_no_groups( + x: FloatTensor, + output_grad: FloatTensor, + weight_shape: Shape, + options: ConvOptions<1>, +) -> FloatTensor { + let x_swapped = B::float_swap_dims(x, 0, 1); + let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1); + let weight_grad_swapped = B::conv1d( + x_swapped, + output_grad_swapped, + None, + ConvOptions::new(options.dilation, options.padding, options.stride, 1), + ); + let mut weight_grad = B::float_swap_dims(weight_grad_swapped, 0, 1); + + if weight_grad.shape() != weight_shape { + let slices = vec![ + Slice::from(0..weight_shape[0]), + Slice::from(0..weight_shape[1]), + Slice::from(0..weight_shape[2]), + ]; + weight_grad = B::float_slice(weight_grad, &slices); + } + weight_grad +} + +fn conv2d_weight_grad_no_groups( + x: FloatTensor, + output_grad: FloatTensor, + weight_shape: Shape, + options: ConvOptions<2>, +) -> FloatTensor { + let x_swapped = B::float_swap_dims(x, 0, 1); + let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1); + let weight_grad_swapped = B::conv2d( + x_swapped, + output_grad_swapped, + None, + ConvOptions::new(options.dilation, options.padding, options.stride, 1), + ); + let mut weight_grad = B::float_swap_dims(weight_grad_swapped, 0, 1); + + if weight_grad.shape() != weight_shape { + let slices = vec![ + Slice::from(0..weight_shape[0]), + Slice::from(0..weight_shape[1]), + Slice::from(0..weight_shape[2]), + Slice::from(0..weight_shape[3]), + ]; + weight_grad = B::float_slice(weight_grad, &slices); + } + weight_grad +} + +fn conv3d_weight_grad_no_groups( + x: FloatTensor, + output_grad: FloatTensor, + weight_shape: Shape, + options: ConvOptions<3>, +) -> FloatTensor { + let x_swapped = B::float_swap_dims(x, 0, 1); + let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1); + let weight_grad_swapped = B::conv3d( + x_swapped, + output_grad_swapped, + None, + ConvOptions::new(options.dilation, options.padding, options.stride, 1), + ); + let mut weight_grad = B::float_swap_dims(weight_grad_swapped, 0, 1); + + if weight_grad.shape() != weight_shape { + let slices = vec![ + Slice::from(0..weight_shape[0]), + Slice::from(0..weight_shape[1]), + Slice::from(0..weight_shape[2]), + Slice::from(0..weight_shape[3]), + Slice::from(0..weight_shape[4]), + ]; + weight_grad = B::float_slice(weight_grad, &slices); + } + weight_grad +} + +fn conv1d_weight_grad_groups( + x: FloatTensor, + mut weight_grad: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<1>, +) -> FloatTensor { + let [channels_out, increment_ci, kernel_size] = weight_grad.shape().dims(); + let increment_co = channels_out / options.groups; + + let x_swapped = B::float_swap_dims(x, 0, 1); + let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1); + + for g in 0..options.groups { + let start_idx_ci = g * increment_ci; + let end_idx_ci = (g + 1) * increment_ci; + let start_idx_co = g * increment_co; + let end_idx_co = (g + 1) * increment_co; + + let x_slice = vec![Slice::new( + start_idx_ci as isize, + Some(end_idx_ci as isize), + 1, + )]; + let x = B::float_slice(x_swapped.clone(), &x_slice); + let grad_slice = vec![Slice::new( + start_idx_co as isize, + Some(end_idx_co as isize), + 1, + )]; + let grad = B::float_slice(output_grad_swapped.clone(), &grad_slice); + let mut weight_grad_tmp = B::conv1d( + x, + grad, + None, + ConvOptions::new(options.dilation, options.padding, options.stride, 1), + ); + weight_grad_tmp = B::float_swap_dims(weight_grad_tmp, 0, 1); + weight_grad = B::float_slice_assign( + weight_grad, + &[ + Slice::from(start_idx_co..end_idx_co), + Slice::from(0..increment_ci), + Slice::from(0..kernel_size), + ], + weight_grad_tmp, + ); + } + + weight_grad +} + +fn conv2d_weight_grad_groups( + x: FloatTensor, + mut weight_grad: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<2>, +) -> FloatTensor { + let [channels_out, increment_ci, kernel_size_1, kernel_size_2] = weight_grad.shape().dims(); + let increment_co = channels_out / options.groups; + + let x_swapped = B::float_swap_dims(x, 0, 1); + let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1); + + for g in 0..options.groups { + let start_idx_ci = g * increment_ci; + let end_idx_ci = (g + 1) * increment_ci; + let start_idx_co = g * increment_co; + let end_idx_co = (g + 1) * increment_co; + + let x_slice = vec![Slice::new( + start_idx_ci as isize, + Some(end_idx_ci as isize), + 1, + )]; + let x = B::float_slice(x_swapped.clone(), &x_slice); + let grad_slice = vec![Slice::new( + start_idx_co as isize, + Some(end_idx_co as isize), + 1, + )]; + let grad = B::float_slice(output_grad_swapped.clone(), &grad_slice); + let mut weight_grad_tmp = B::conv2d( + x, + grad, + None, + ConvOptions::new(options.dilation, options.padding, options.stride, 1), + ); + weight_grad_tmp = B::float_swap_dims(weight_grad_tmp, 0, 1); + let [_, _, kernel_size_1_tmp, kernel_size_2_tmp] = weight_grad_tmp.shape().dims(); + + if kernel_size_1_tmp != kernel_size_1 || kernel_size_2_tmp != kernel_size_2 { + let slices = vec![ + Slice::from(0..increment_co), + Slice::from(0..increment_ci), + Slice::from(0..kernel_size_1), + Slice::from(0..kernel_size_2), + ]; + weight_grad_tmp = B::float_slice(weight_grad_tmp, &slices); + } + + weight_grad = B::float_slice_assign( + weight_grad, + &[ + Slice::from(start_idx_co..end_idx_co), + Slice::from(0..increment_ci), + Slice::from(0..kernel_size_1), + Slice::from(0..kernel_size_2), + ], + weight_grad_tmp, + ); + } + + weight_grad +} + +fn conv3d_weight_grad_groups( + x: FloatTensor, + mut weight_grad: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<3>, +) -> FloatTensor { + let [ + channels_out, + increment_ci, + kernel_size_1, + kernel_size_2, + kernel_size_3, + ] = weight_grad.shape().dims(); + let increment_co = channels_out / options.groups; + + let x_swapped = B::float_swap_dims(x, 0, 1); + let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1); + + for g in 0..options.groups { + let start_idx_ci = g * increment_ci; + let end_idx_ci = (g + 1) * increment_ci; + let start_idx_co = g * increment_co; + let end_idx_co = (g + 1) * increment_co; + + let x_slice = vec![Slice::new( + start_idx_ci as isize, + Some(end_idx_ci as isize), + 1, + )]; + let x = B::float_slice(x_swapped.clone(), &x_slice); + let grad_slice = vec![Slice::new( + start_idx_co as isize, + Some(end_idx_co as isize), + 1, + )]; + let grad = B::float_slice(output_grad_swapped.clone(), &grad_slice); + let mut weight_grad_tmp = B::conv3d( + x, + grad, + None, + ConvOptions::new(options.dilation, options.padding, options.stride, 1), + ); + weight_grad_tmp = B::float_swap_dims(weight_grad_tmp, 0, 1); + let [ + _, + _, + kernel_size_1_tmp, + kernel_size_2_tmp, + kernel_size_3_tmp, + ] = weight_grad_tmp.shape().dims(); + + if kernel_size_1_tmp != kernel_size_1 + || kernel_size_2_tmp != kernel_size_2 + || kernel_size_3_tmp != kernel_size_3 + { + let slices = vec![ + Slice::from(0..increment_co), + Slice::from(0..increment_ci), + Slice::from(0..kernel_size_1), + Slice::from(0..kernel_size_2), + Slice::from(0..kernel_size_3), + ]; + weight_grad_tmp = B::float_slice(weight_grad_tmp, &slices); + } + + weight_grad = B::float_slice_assign( + weight_grad, + &[ + Slice::from(start_idx_co..end_idx_co), + Slice::from(0..increment_ci), + Slice::from(0..kernel_size_1), + Slice::from(0..kernel_size_2), + Slice::from(0..kernel_size_3), + ], + weight_grad_tmp, + ); + } + + weight_grad +} + +fn conv_transpose1d_weight_grad_no_groups( + x: FloatTensor, + output_grad: FloatTensor, + weight_shape: Shape, + options: ConvTransposeOptions<1>, +) -> FloatTensor { + let x_swapped = B::float_swap_dims(x, 0, 1); + let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1); + let weight_grad_swapped = B::conv1d( + output_grad_swapped, + x_swapped, + None, + ConvOptions::new(options.dilation, options.padding, options.stride, 1), + ); + let mut weight_grad = B::float_swap_dims(weight_grad_swapped, 0, 1); + + let grad_shape = weight_grad.shape(); + if grad_shape != weight_shape { + let slices = vec![ + Slice::from(0..weight_shape[0]), + Slice::from(0..weight_shape[1]), + Slice::from(0..weight_shape[2]), + ]; + weight_grad = B::float_slice(weight_grad, &slices); + } + weight_grad +} + +fn conv_transpose2d_weight_grad_no_groups( + x: FloatTensor, + output_grad: FloatTensor, + weight_shape: Shape, + options: ConvTransposeOptions<2>, +) -> FloatTensor { + let x_swapped = B::float_swap_dims(x, 0, 1); + let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1); + let weight_grad_swapped = B::conv2d( + output_grad_swapped, + x_swapped, + None, + ConvOptions::new(options.dilation, options.padding, options.stride, 1), + ); + let mut weight_grad = B::float_swap_dims(weight_grad_swapped, 0, 1); + + let grad_shape = weight_grad.shape(); + if grad_shape != weight_shape { + let slices = vec![ + Slice::from(0..weight_shape[0]), + Slice::from(0..weight_shape[1]), + Slice::from(0..weight_shape[2]), + Slice::from(0..weight_shape[3]), + ]; + weight_grad = B::float_slice(weight_grad, &slices); + } + weight_grad +} + +fn conv_transpose3d_weight_grad_no_groups( + x: FloatTensor, + output_grad: FloatTensor, + weight_shape: Shape, + options: ConvTransposeOptions<3>, +) -> FloatTensor { + let x_swapped = B::float_swap_dims(x, 0, 1); + let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1); + let weight_grad_swapped = B::conv3d( + output_grad_swapped, + x_swapped, + None, + ConvOptions::new(options.dilation, options.padding, options.stride, 1), + ); + let mut weight_grad = B::float_swap_dims(weight_grad_swapped, 0, 1); + + let grad_shape = weight_grad.shape(); + if grad_shape != weight_shape { + let slices = vec![ + Slice::from(0..weight_shape[0]), + Slice::from(0..weight_shape[1]), + Slice::from(0..weight_shape[2]), + Slice::from(0..weight_shape[3]), + Slice::from(0..weight_shape[4]), + ]; + weight_grad = B::float_slice(weight_grad, &slices); + } + weight_grad +} + +fn conv_transpose1d_weight_grad_groups( + x: FloatTensor, + mut weight_grad: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<1>, +) -> FloatTensor { + let [channels_in, increment_co, kernel_size] = weight_grad.shape().dims(); + let increment_ci = channels_in / options.groups; + + let x_swapped = B::float_swap_dims(x, 0, 1); + let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1); + + for g in 0..options.groups { + let start_idx_ci = g * increment_ci; + let end_idx_ci = (g + 1) * increment_ci; + let start_idx_co = g * increment_co; + let end_idx_co = (g + 1) * increment_co; + + let x_slice = vec![Slice::new( + start_idx_ci as isize, + Some(end_idx_ci as isize), + 1, + )]; + let x = B::float_slice(x_swapped.clone(), &x_slice); + let grad_slice = vec![Slice::new( + start_idx_co as isize, + Some(end_idx_co as isize), + 1, + )]; + let grad = B::float_slice(output_grad_swapped.clone(), &grad_slice); + let mut weight_grad_tmp = B::conv1d( + grad, + x, + None, + ConvOptions::new(options.dilation, options.padding, options.stride, 1), + ); + weight_grad_tmp = B::float_swap_dims(weight_grad_tmp, 0, 1); + let [_, _, kernel_size_tmp] = weight_grad_tmp.shape().dims(); + + if kernel_size_tmp != kernel_size { + let slices = vec![ + Slice::from(0..increment_ci), + Slice::from(0..increment_co), + Slice::from(0..kernel_size), + ]; + weight_grad_tmp = B::float_slice(weight_grad_tmp, &slices); + } + + weight_grad = B::float_slice_assign( + weight_grad, + &[ + Slice::from(start_idx_ci..end_idx_ci), + Slice::from(0..increment_co), + Slice::from(0..kernel_size), + ], + weight_grad_tmp, + ); + } + + weight_grad +} + +fn conv_transpose2d_weight_grad_groups( + x: FloatTensor, + mut weight_grad: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<2>, +) -> FloatTensor { + let [channels_in, increment_co, kernel_size_1, kernel_size_2] = weight_grad.shape().dims(); + let increment_ci = channels_in / options.groups; + + let x_swapped = B::float_swap_dims(x, 0, 1); + let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1); + + for g in 0..options.groups { + let start_idx_ci = g * increment_ci; + let end_idx_ci = (g + 1) * increment_ci; + let start_idx_co = g * increment_co; + let end_idx_co = (g + 1) * increment_co; + + let x_slice = vec![Slice::new( + start_idx_ci as isize, + Some(end_idx_ci as isize), + 1, + )]; + let x = B::float_slice(x_swapped.clone(), &x_slice); + let grad_slice = vec![Slice::new( + start_idx_co as isize, + Some(end_idx_co as isize), + 1, + )]; + let grad = B::float_slice(output_grad_swapped.clone(), &grad_slice); + let mut weight_grad_tmp = B::conv2d( + grad, + x, + None, + ConvOptions::new(options.dilation, options.padding, options.stride, 1), + ); + weight_grad_tmp = B::float_swap_dims(weight_grad_tmp, 0, 1); + let [_, _, kernel_size_1_tmp, kernel_size_2_tmp] = weight_grad_tmp.shape().dims(); + + if kernel_size_1_tmp != kernel_size_1 || kernel_size_2_tmp != kernel_size_2 { + let slices = vec![ + Slice::from(0..increment_ci), + Slice::from(0..increment_co), + Slice::from(0..kernel_size_1), + Slice::from(0..kernel_size_2), + ]; + weight_grad_tmp = B::float_slice(weight_grad_tmp, &slices); + } + + weight_grad = B::float_slice_assign( + weight_grad, + &[ + Slice::from(start_idx_ci..end_idx_ci), + Slice::from(0..increment_co), + Slice::from(0..kernel_size_1), + Slice::from(0..kernel_size_2), + ], + weight_grad_tmp, + ); + } + + weight_grad +} + +fn conv_transpose3d_weight_grad_groups( + x: FloatTensor, + mut weight_grad: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<3>, +) -> FloatTensor { + let [ + channels_in, + increment_co, + kernel_size_1, + kernel_size_2, + kernel_size_3, + ] = weight_grad.shape().dims(); + let increment_ci = channels_in / options.groups; + + let x_swapped = B::float_swap_dims(x, 0, 1); + let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1); + + for g in 0..options.groups { + let start_idx_ci = g * increment_ci; + let end_idx_ci = (g + 1) * increment_ci; + let start_idx_co = g * increment_co; + let end_idx_co = (g + 1) * increment_co; + + let x_slice = vec![Slice::new( + start_idx_ci as isize, + Some(end_idx_ci as isize), + 1, + )]; + let x = B::float_slice(x_swapped.clone(), &x_slice); + let grad_slice = vec![Slice::new( + start_idx_co as isize, + Some(end_idx_co as isize), + 1, + )]; + let grad = B::float_slice(output_grad_swapped.clone(), &grad_slice); + let mut weight_grad_tmp = B::conv3d( + grad, + x, + None, + ConvOptions::new(options.dilation, options.padding, options.stride, 1), + ); + weight_grad_tmp = B::float_swap_dims(weight_grad_tmp, 0, 1); + let [ + _, + _, + kernel_size_1_tmp, + kernel_size_2_tmp, + kernel_size_3_tmp, + ] = weight_grad_tmp.shape().dims(); + + if kernel_size_1_tmp != kernel_size_1 + || kernel_size_2_tmp != kernel_size_2 + || kernel_size_3_tmp != kernel_size_3 + { + let slices = vec![ + Slice::from(0..increment_ci), + Slice::from(0..increment_co), + Slice::from(0..kernel_size_1), + Slice::from(0..kernel_size_2), + Slice::from(0..kernel_size_3), + ]; + weight_grad_tmp = B::float_slice(weight_grad_tmp, &slices); + } + weight_grad = B::float_slice_assign( + weight_grad, + &[ + Slice::from(start_idx_ci..end_idx_ci), + Slice::from(0..increment_co), + Slice::from(0..kernel_size_1), + Slice::from(0..kernel_size_2), + Slice::from(0..kernel_size_3), + ], + weight_grad_tmp, + ); + } + + weight_grad +} + +/// Compute the `padding_out` for a transpose conv that exactly recovers the +/// original `size_in` from `size_out`, accounting for any input elements the +/// forward conv dropped. Shared by `conv{1,2,3}d_x_backward` and the CubeCL +/// dgrad fallback so the two paths can't drift. +pub fn calculate_padding_out( + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + size_in: usize, + size_out: usize, +) -> usize { + if stride <= 1 { + return 0; + } + + // Invert the transpose conv output formula to recover the exact number of + // input elements that a forward conv would drop for this (size_in, size_out). + // + // Forward: size_out = floor((size_in + 2*padding - dilated_kernel) / stride) + 1 + // Transpose: trans_out = (size_out - 1)*stride + dilated_kernel + padding_out - 2*padding + // Setting trans_out == size_in and solving for padding_out: + let dilated_kernel = dilation * (kernel_size - 1) + 1; + let base = (size_out as i64 - 1) * stride as i64 + dilated_kernel as i64 - 2 * padding as i64; + i64::max(0, size_in as i64 - base) as usize +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_calculate_output_size_1() { + let kernel_size = 3; + let stride = 1; + let padding = 1; + let size_in = 3; + let dilation = 1; + + let size_out = calculate_conv_output_size(kernel_size, stride, padding, dilation, size_in); + + assert_eq!(size_out, 3); + } + + #[test] + fn test_calculate_output_size_2() { + let kernel_size = 5; + let stride = 2; + let padding = 3; + let size_in = 27; + let dilation = 1; + + let size_out = calculate_conv_output_size(kernel_size, stride, padding, dilation, size_in); + + assert_eq!(size_out, 15); + } + + #[test] + fn test_calculate_output_size_3() { + let kernel_size = 5; + let stride = 2; + let padding = 3; + let size_in = 27; + let dilation = 2; + + let size_out = calculate_conv_output_size(kernel_size, stride, padding, dilation, size_in); + + assert_eq!(size_out, 13); + } + + #[test] + fn test_calculate_same_padding_1() { + let kernel_size = 3; + let stride = 1; + let size_in = 3; + let dilation = 1; + + let padding = calculate_conv_padding(kernel_size, stride, size_in, size_in); + let size_out = calculate_conv_output_size(kernel_size, stride, padding, dilation, size_in); + + assert_eq!(size_in, size_out, "Expected size"); + } + + #[test] + fn test_calculate_same_padding_2() { + let kernel_size = 3; + let stride = 2; + let size_in = 7; + let dilation = 1; + + let padding = calculate_conv_padding(kernel_size, stride, size_in, size_in); + let size_out = calculate_conv_output_size(kernel_size, stride, padding, dilation, size_in); + + assert_eq!(size_in, size_out, "Expected size"); + } + + #[test] + fn test_calculate_output_padding_1() { + let kernel_size = 3; + let stride = 2; + let size_in = 7; + let size_out = 10; + let dilation = 1; + + let padding = calculate_conv_padding(kernel_size, stride, size_in, size_out); + let size_out_expected = + calculate_conv_output_size(kernel_size, stride, padding, dilation, size_in); + + assert_eq!(size_out, size_out_expected, "Expected size"); + } + + #[test] + fn test_expect_conv2d_output_shape() { + // in channels: 3 + // out channels: 8 + // size in: [27, 3] + // kernel size: [5, 3] + let stride = [2, 1]; + let padding = [3, 1]; + let dilation = [2, 1]; + let shape = calculate_conv_output_shape( + &Shape::new([12, 3, 27, 3]), + &Shape::new([8, 3, 5, 3]), + &stride, + &padding, + &dilation, + ) + .unwrap(); + assert_eq!(shape, Shape::new([12, 8, 13, 3])) + } +} diff --git a/crates/burn-backend/src/backend/ops/modules/ctc.rs b/crates/burn-backend/src/backend/ops/modules/ctc.rs new file mode 100644 index 0000000..383efbe --- /dev/null +++ b/crates/burn-backend/src/backend/ops/modules/ctc.rs @@ -0,0 +1,613 @@ +use burn_std::{Shape, Slice}; + +use crate::{ + Backend, TensorMetadata, get_device_settings, + tensor::{BoolTensor, FloatTensor, IntTensor}, +}; + +/// Default CTC loss implementation using the forward (alpha) algorithm. +/// +/// Computes the Connectionist Temporal Classification loss by summing over +/// all valid alignments between the input and target sequences. +/// +/// # Arguments +/// +/// * `log_probs` - Log-probabilities of shape `[T, N, C]` +/// * `targets` - Target indices of shape `[N, S]` +/// * `input_lengths` - Actual input sequence lengths per batch element `[N]` +/// * `target_lengths` - Actual target lengths per batch element `[N]` +/// * `blank` - Index of the blank label +/// +/// # Returns +/// +/// Per-sample loss of shape `[N]` +pub fn ctc_loss_default( + log_probs: FloatTensor, + targets: IntTensor, + input_lengths: IntTensor, + target_lengths: IntTensor, + blank: usize, +) -> FloatTensor { + let alpha = AlphaCtx::::compute( + log_probs, + &targets, + input_lengths, + target_lengths.clone(), + blank, + ); + extract_loss::(&alpha, target_lengths) +} + +/// Compose the CTC gradient w.r.t. `log_probs` from pre-computed alpha, beta, and nll. +/// +/// The T-iteration alpha and beta recursions are the dominant cost of the backward +/// pass. Backends that fuse those recursions into a single kernel launch can call +/// this helper to reuse the gradient composition. +/// +/// # Arguments +/// +/// * `log_probs` - Log-probabilities `[T, N, C]` +/// * `targets` - Target label indices `[N, S]` +/// * `input_lengths` - Actual input sequence lengths per batch element `[N]` +/// * `grad_loss` - Upstream gradient w.r.t. the per-sample loss `[N]` +/// * `log_alpha_full` - Alpha recursion output `[T, N, 2S+1]` +/// * `log_beta_full` - Beta recursion output `[T, N, 2S+1]` +/// * `nll` - Per-sample negative log-likelihood (forward loss) `[N]` +/// * `blank` - Index of the blank label +#[allow(clippy::too_many_arguments)] +pub fn ctc_grad_from_alpha_beta_default( + log_probs: FloatTensor, + targets: IntTensor, + input_lengths: IntTensor, + grad_loss: FloatTensor, + log_alpha_full: FloatTensor, + log_beta_full: FloatTensor, + nll: FloatTensor, + blank: usize, +) -> FloatTensor { + let log_probs_shape = log_probs.shape(); + let [max_input_length, batch_size, num_classes] = log_probs_shape.dims::<3>(); + let target_shape = targets.shape(); + let max_target_len = target_shape.dims::<2>()[1]; + let max_l_prime_len = 2 * max_target_len + 1; + let device = log_probs.device(); + let int_dtype: burn_std::IntDType = targets.dtype().into(); + let settings = get_device_settings::(&device); + + let blank_inserted_targets = insert_blanks::( + &targets, + batch_size, + max_target_len, + max_l_prime_len, + blank, + &device, + int_dtype, + ); + + // Both log_alpha[t, n, s] and log_beta[t, n, s] include a factor of + // log_probs[t, n, l'[s]] (added on every recursion step). The CTC paper's + // alpha_hat * beta_hat product divides one of those factors out, so we + // subtract log_probs[t, n, l'[s]] when forming log_post. + // + // We then divide by total_prob = exp(-nll) to obtain the alignment + // posterior, which in log space means *adding* nll (since nll = -log P, + // dividing by P is adding nll). Per PyTorch's CTC backward kernel: + // log_post[t, n, s] = log_alpha + log_beta - log_probs[t, n, l'[s]] - log P + // = log_alpha + log_beta - log_probs[t, n, l'[s]] + nll + let indices_3d = B::int_reshape( + blank_inserted_targets, + Shape::new([1, batch_size, max_l_prime_len]), + ); + let indices_3d = B::int_expand( + indices_3d, + Shape::new([max_input_length, batch_size, max_l_prime_len]), + ); + let log_probs_at_l = B::float_gather(2, log_probs.clone(), indices_3d.clone()); + + // Samples with an unreachable target yield nll = +inf. For those, log_alpha + // stays at -inf at many (t, s) while log_beta is finite at the boundary, so + // log_post = (-inf) + finite - finite + (+inf) = NaN and -exp(NaN) = NaN + // contaminates the gradient. `NaN * 0 = NaN` under IEEE 754, so zero_infinity + // masking on the outer grad_loss can't clear it. Capture the mask now and + // zero the gradient for those samples at the end. + let nll_is_inf = B::float_is_inf(nll.clone(), settings.bool_dtype); + + let nll_b = B::float_reshape(nll, Shape::new([1, batch_size, 1])); + let nll_b = B::float_expand( + nll_b, + Shape::new([max_input_length, batch_size, max_l_prime_len]), + ); + let log_post = B::float_add( + B::float_sub(B::float_add(log_alpha_full, log_beta_full), log_probs_at_l), + nll_b, + ); + + // grad starts as exp(log_probs) * grad_loss[None, :, None]. + let grad_loss_3d = B::float_reshape(grad_loss, Shape::new([1, batch_size, 1])); + let grad_loss_b = B::float_expand( + grad_loss_3d.clone(), + Shape::new([max_input_length, batch_size, num_classes]), + ); + let mut grad = B::float_mul(B::float_exp(log_probs), grad_loss_b); + + // Subtract sum over s of grad_loss[n] * exp(log_post[t, n, s]) at index l'[n, s]. + let grad_loss_post = B::float_expand( + grad_loss_3d, + Shape::new([max_input_length, batch_size, max_l_prime_len]), + ); + let scatter_value = B::float_neg(B::float_mul(B::float_exp(log_post), grad_loss_post)); + + grad = B::float_scatter_add(2, grad, indices_3d, scatter_value); + + // Mask out timesteps where t >= input_lengths[n]. + let t_indices = B::int_arange(0..max_input_length as i64, &device, int_dtype); + let t_indices = B::int_reshape(t_indices, Shape::new([max_input_length, 1, 1])); + let t_indices = B::int_expand( + t_indices, + Shape::new([max_input_length, batch_size, num_classes]), + ); + let il_b = B::int_reshape(input_lengths, Shape::new([1, batch_size, 1])); + let il_b = B::int_expand( + il_b, + Shape::new([max_input_length, batch_size, num_classes]), + ); + let oob_mask = B::int_greater_equal(t_indices, il_b, settings.bool_dtype); + + // Broadcast the nll-is-inf mask across [T, N, C] and OR with oob_mask so a + // single mask_fill zeros both unreachable samples and out-of-bound timesteps. + let nll_inf_b = B::bool_reshape(nll_is_inf, Shape::new([1, batch_size, 1])); + let nll_inf_b = B::bool_expand( + nll_inf_b, + Shape::new([max_input_length, batch_size, num_classes]), + ); + let mask = B::bool_or(oob_mask, nll_inf_b); + B::float_mask_fill(grad, mask, 0.0.into()) +} + +/// Cached state from the alpha recursion. Only `last` is consumed by +/// `ctc_loss_default` (via `extract_loss`); the other fields hold intermediate +/// products that backends with a native backward kernel could reuse if wired +/// up. They are kept here to document the recursion's outputs. +#[allow(dead_code)] +struct AlphaCtx { + /// `log_alpha[T, N, 2S+1]` (full history). + full: FloatTensor, + /// `log_alpha[T-1, :, :]` (last timestep; used to read out the loss). + last: FloatTensor, + /// `l'` after blank insertion `[N, 2S+1]`. + blank_inserted_targets: IntTensor, + /// `log_probs[t, n, l'[n, s]]` pre-gathered as `[T, N, 2S+1]`. + log_probs_at_l_full: FloatTensor, + max_l_prime_len: usize, +} + +impl AlphaCtx { + fn compute( + log_probs: FloatTensor, + targets: &IntTensor, + input_lengths: IntTensor, + target_lengths: IntTensor, + blank: usize, + ) -> Self { + let log_probs_shape = log_probs.shape(); + let [max_input_length, batch_size, num_classes] = log_probs_shape.dims::<3>(); + let target_shape = targets.shape(); + let max_target_len = target_shape.dims::<2>()[1]; + let device = log_probs.device(); + let float_dtype: burn_std::FloatDType = log_probs.dtype().into(); + let int_dtype: burn_std::IntDType = targets.dtype().into(); + let settings = get_device_settings::(&device); + + let max_l_prime_len = 2 * max_target_len + 1; + let blank_inserted_targets = insert_blanks::( + targets, + batch_size, + max_target_len, + max_l_prime_len, + blank, + &device, + int_dtype, + ); + + // Pre-allocate the full alpha tensor [T, N, 2S+1] filled with -inf. + let mut alpha_full = B::float_full( + Shape::new([max_input_length, batch_size, max_l_prime_len]), + f32::NEG_INFINITY.into(), + &device, + float_dtype, + ); + + // Initialize alpha[0, :, 0] = log_probs[0, :, blank] + // and alpha[0, :, 1] = log_probs[0, :, l'[1]]. + let log_probs_t0 = B::float_slice( + log_probs.clone(), + &[Slice::new(0, Some(1), 1), Slice::full(), Slice::full()], + ); + let log_probs_t0 = B::float_reshape(log_probs_t0, Shape::new([batch_size, num_classes])); + + let first_blank = B::int_slice( + blank_inserted_targets.clone(), + &[Slice::full(), Slice::new(0, Some(1), 1)], + ); + let log_prob_blank = B::float_gather(1, log_probs_t0.clone(), first_blank); + // Broadcast to [1, N, 1] for slice_assign into alpha_full. + let log_prob_blank_3d = B::float_reshape(log_prob_blank, Shape::new([1, batch_size, 1])); + alpha_full = B::float_slice_assign( + alpha_full, + &[ + Slice::new(0, Some(1), 1), + Slice::full(), + Slice::new(0, Some(1), 1), + ], + log_prob_blank_3d, + ); + + if max_l_prime_len > 1 { + let first_label = B::int_slice( + blank_inserted_targets.clone(), + &[Slice::full(), Slice::new(1, Some(2), 1)], + ); + let log_prob_first = B::float_gather(1, log_probs_t0, first_label); + let log_prob_first_3d = + B::float_reshape(log_prob_first, Shape::new([1, batch_size, 1])); + alpha_full = B::float_slice_assign( + alpha_full, + &[ + Slice::new(0, Some(1), 1), + Slice::full(), + Slice::new(1, Some(2), 1), + ], + log_prob_first_3d, + ); + } + + // Track the latest row separately for the recursion (cheaper than + // re-slicing alpha_full each iteration). + let mut log_alpha = B::float_slice( + alpha_full.clone(), + &[Slice::new(0, Some(1), 1), Slice::full(), Slice::full()], + ); + log_alpha = B::float_reshape(log_alpha, Shape::new([batch_size, max_l_prime_len])); + + let l_prime_mask = create_l_prime_mask::( + &blank_inserted_targets, + batch_size, + max_l_prime_len, + blank, + &device, + int_dtype, + settings.bool_dtype, + ); + let s_mask = create_s_mask::( + &target_lengths, + batch_size, + max_l_prime_len, + &device, + int_dtype, + settings.bool_dtype, + ); + + // Hoist out of the T-loop: padding tensors for right_shift (same + // value/shape at every iteration) and the full `[T, N, 2S+1]` + // gather of log_probs at l' (one T-sized gather replaces T small + // gathers). + let pad_1 = B::float_full( + Shape::new([batch_size, 1]), + f32::NEG_INFINITY.into(), + &device, + float_dtype, + ); + let pad_2 = B::float_full( + Shape::new([batch_size, 2]), + f32::NEG_INFINITY.into(), + &device, + float_dtype, + ); + let indices_3d = B::int_expand( + B::int_reshape( + blank_inserted_targets.clone(), + Shape::new([1, batch_size, max_l_prime_len]), + ), + Shape::new([max_input_length, batch_size, max_l_prime_len]), + ); + let log_probs_at_l_full = B::float_gather(2, log_probs.clone(), indices_3d); + + // Precompute `combined_mask_all[t, n, s] = (input_lengths[n] > t) AND + // s_mask[n, s]` for every t in one shot. The T-loop reads its row via + // a metadata-only slice instead of recomputing the `int_greater_elem` + // + bool_and per iteration. + let t_indices_2d = B::int_expand( + B::int_reshape( + B::int_arange(0..max_input_length as i64, &device, int_dtype), + Shape::new([max_input_length, 1]), + ), + Shape::new([max_input_length, batch_size]), + ); + let il_tn = B::int_expand( + B::int_reshape(input_lengths.clone(), Shape::new([1, batch_size])), + Shape::new([max_input_length, batch_size]), + ); + let t_mask_all = B::bool_expand( + B::bool_reshape( + B::int_greater(il_tn, t_indices_2d, settings.bool_dtype), + Shape::new([max_input_length, batch_size, 1]), + ), + Shape::new([max_input_length, batch_size, max_l_prime_len]), + ); + let s_mask_bcast = B::bool_expand( + B::bool_reshape(s_mask.clone(), Shape::new([1, batch_size, max_l_prime_len])), + Shape::new([max_input_length, batch_size, max_l_prime_len]), + ); + let combined_mask_all = B::bool_and(t_mask_all, s_mask_bcast); + + for t in 1..max_input_length { + let combined_mask = B::bool_reshape( + B::bool_slice( + combined_mask_all.clone(), + &[ + Slice::new(t as isize, Some(t as isize + 1), 1), + Slice::full(), + Slice::full(), + ], + ), + Shape::new([batch_size, max_l_prime_len]), + ); + + let log_alpha_s = log_alpha.clone(); + let log_alpha_s_m1 = right_shift::(&log_alpha, &pad_1, max_l_prime_len, 1); + let log_alpha_s_m2 = right_shift::(&log_alpha, &pad_2, max_l_prime_len, 2); + + let bar = log_sum_exp::(log_alpha_s, log_alpha_s_m1, settings.bool_dtype); + let bar_with_skip = log_sum_exp::(bar.clone(), log_alpha_s_m2, settings.bool_dtype); + let log_alpha_combined = B::float_mask_where(bar, l_prime_mask.clone(), bar_with_skip); + + // Slice row t from the pre-gathered `[T, N, 2S+1]` tensor. + let log_probs_at_l = B::float_reshape( + B::float_slice( + log_probs_at_l_full.clone(), + &[ + Slice::new(t as isize, Some(t as isize + 1), 1), + Slice::full(), + Slice::full(), + ], + ), + Shape::new([batch_size, max_l_prime_len]), + ); + let new_alpha = B::float_add(log_alpha_combined, log_probs_at_l); + log_alpha = B::float_mask_where(log_alpha, combined_mask, new_alpha); + + let log_alpha_3d = B::float_reshape( + log_alpha.clone(), + Shape::new([1, batch_size, max_l_prime_len]), + ); + alpha_full = B::float_slice_assign( + alpha_full, + &[ + Slice::new(t as isize, Some(t as isize + 1), 1), + Slice::full(), + Slice::full(), + ], + log_alpha_3d, + ); + } + + Self { + full: alpha_full, + last: log_alpha, + blank_inserted_targets, + log_probs_at_l_full, + max_l_prime_len, + } + } +} + +/// Extract the per-sample loss from the last alpha row. +fn extract_loss(alpha: &AlphaCtx, target_lengths: IntTensor) -> FloatTensor { + let log_alpha_shape = alpha.last.shape(); + let [batch_size, _] = log_alpha_shape.dims::<2>(); + let settings = get_device_settings::(&alpha.last.device()); + + let last_blank_idx = B::int_mul_scalar(target_lengths.clone(), 2.into()); + let last_blank_idx = B::int_reshape(last_blank_idx, Shape::new([batch_size, 1])); + let last_label_idx = B::int_clamp_min( + B::int_sub_scalar(last_blank_idx.clone(), 1.into()), + 0.into(), + ); + + let log_alpha_last_blank = B::float_gather(1, alpha.last.clone(), last_blank_idx); + let log_alpha_last_blank = B::float_reshape(log_alpha_last_blank, Shape::new([batch_size])); + + let log_alpha_last_label = B::float_gather(1, alpha.last.clone(), last_label_idx); + let log_alpha_last_label = B::float_reshape(log_alpha_last_label, Shape::new([batch_size])); + + // For target_lengths == 0, last_label is meaningless: substitute -inf. + let target_len_zero = B::int_equal_elem(target_lengths, 0.into(), settings.bool_dtype); + let log_alpha_last_label = B::float_mask_fill( + log_alpha_last_label, + target_len_zero, + f32::NEG_INFINITY.into(), + ); + + let log_likelihood = log_sum_exp::( + log_alpha_last_blank, + log_alpha_last_label, + settings.bool_dtype, + ); + B::float_neg(log_likelihood) +} + +/// Insert blank labels between each target label: [b, l1, b, l2, ..., b] +fn insert_blanks( + targets: &IntTensor, + batch_size: usize, + max_target_len: usize, + max_l_prime_len: usize, + blank: usize, + device: &B::Device, + int_dtype: burn_std::IntDType, +) -> IntTensor { + let result = B::int_full( + Shape::new([batch_size, max_l_prime_len]), + (blank as i64).into(), + device, + int_dtype, + ); + + if max_target_len == 0 { + return result; + } + + // Place every target label at odd columns {1, 3, 5, ...} in one + // strided slice_assign, equivalent to `result[:, 1::2] = targets`. + B::int_slice_assign( + result, + &[Slice::full(), Slice::new(1, None, 2)], + targets.clone(), + ) +} + +/// Right-shift a 2D float tensor by `shift` positions, prepending the +/// pre-allocated `padding` tensor (shape `[batch_size, shift]`, value +/// `-inf`) instead of materializing it each call. +/// +/// Called inside the T-loop of the alpha recursion; hoisting the padding +/// out of the loop eliminates `O(T)` `float_full` allocations. +fn right_shift( + tensor: &FloatTensor, + padding: &FloatTensor, + cols: usize, + shift: usize, +) -> FloatTensor { + // Shifting by more than the column count pushes every data slot off + // the right. Avoid the `cols - shift` usize underflow when + // `max_target_len == 0` (so `max_l_prime_len == 1`) by narrowing the + // all-`-inf` padding down to `cols`. + if cols < shift { + return B::float_slice( + padding.clone(), + &[Slice::full(), Slice::new(0, Some(cols as isize), 1)], + ); + } + let shortened = B::float_slice( + tensor.clone(), + &[ + Slice::full(), + Slice::new(0, Some((cols - shift) as isize), 1), + ], + ); + B::float_cat(alloc::vec![padding.clone(), shortened], 1) +} + +/// Compute `log(exp(a) + exp(b))` in a numerically stable way. +/// +/// `log_sum_exp(a, b) = max(a, b) + log1p(exp(-|a - b|))`. The edge case is +/// `a = b = -inf`, where `-|(-inf) - (-inf)| = NaN`; we detect `max == -inf` +/// and substitute a `-inf` diff so the final sum stays `-inf` (both via the +/// mask and because `log1p(exp(-inf)) = 0`). Gradient-safe: no `NaN` flows +/// through the forward intermediates when inputs are `-inf`. +/// +/// Precondition: inputs must be `<= 0` (log-probabilities). `+inf` inputs are +/// not guarded and produce `NaN`; callers outside the CTC recursion should +/// validate this themselves. +fn log_sum_exp( + a: FloatTensor, + b: FloatTensor, + bool_dtype: burn_std::BoolDType, +) -> FloatTensor { + // `-inf` values in `a` or `b` would make `a - b` evaluate to `NaN` + // (when both are `-inf`) and the backward pass through that `NaN` + // intermediate propagates `NaN` into the gradient even when the + // forward mask discards it (`0 * NaN = NaN` in IEEE). Clamp `-inf` + // to `0` on safe copies used only for the diff computation; compute + // `max` on the original values so its output is correct in the + // `-inf` cases. + let a_is_neg_inf = B::float_equal_elem(a.clone(), f32::NEG_INFINITY.into(), bool_dtype); + let b_is_neg_inf = B::float_equal_elem(b.clone(), f32::NEG_INFINITY.into(), bool_dtype); + let either_neg_inf = B::bool_or(a_is_neg_inf.clone(), b_is_neg_inf.clone()); + + let a_safe = B::float_mask_fill(a.clone(), a_is_neg_inf, 0.0.into()); + let b_safe = B::float_mask_fill(b.clone(), b_is_neg_inf, 0.0.into()); + + let lt_mask = B::float_lower(a.clone(), b.clone(), bool_dtype); + let mx = B::float_mask_where(a, lt_mask, b); + + // diff_safe = -|a_safe - b_safe|. Finite by construction. When either + // input was `-inf`, force it to `-inf` so `exp(diff) == 0` and the + // `log1p` term contributes nothing (`result = mx`). When both were + // `-inf`, `mx = -inf` so `result = -inf + 0 = -inf`. + let diff_safe = B::float_neg(B::float_abs(B::float_sub(a_safe, b_safe))); + let diff_final = B::float_mask_fill(diff_safe, either_neg_inf, f32::NEG_INFINITY.into()); + + B::float_add(mx, B::float_log1p(B::float_exp(diff_final))) +} + +/// Mask for the alpha skip transition: `l'[s] != blank AND l'[s] != l'[s-2] AND s >= 2`. +fn create_l_prime_mask( + blank_inserted_targets: &IntTensor, + batch_size: usize, + max_l_prime_len: usize, + blank: usize, + device: &B::Device, + int_dtype: burn_std::IntDType, + bool_dtype: burn_std::BoolDType, +) -> BoolTensor { + // The mask requires `s >= 2`, which is unsatisfiable when max_l_prime_len < 2 + // (i.e. targets have shape [N, 0]). Bail out before the `max_l_prime_len - 2` + // usize subtraction underflows. + if max_l_prime_len < 2 { + return B::bool_zeros( + Shape::new([batch_size, max_l_prime_len]), + device, + bool_dtype, + ); + } + let l_prime = blank_inserted_targets.clone(); + + let not_blank = B::int_not_equal_elem(l_prime.clone(), (blank as i64).into(), bool_dtype); + + let l_prime_shifted = { + let padding = B::int_full( + Shape::new([batch_size, 2]), + (blank as i64).into(), + device, + int_dtype, + ); + let shortened = B::int_slice( + l_prime.clone(), + &[ + Slice::full(), + Slice::new(0, Some((max_l_prime_len - 2) as isize), 1), + ], + ); + B::int_cat(alloc::vec![padding, shortened], 1) + }; + let not_equal_s_m2 = B::int_not_equal(l_prime, l_prime_shifted, bool_dtype); + + let col_indices = B::int_arange(0..max_l_prime_len as i64, device, int_dtype); + let col_indices = B::int_reshape(col_indices, Shape::new([1, max_l_prime_len])); + let col_indices = B::int_expand(col_indices, Shape::new([batch_size, max_l_prime_len])); + let s_ge_2 = B::int_greater_equal_elem(col_indices, 2.into(), bool_dtype); + + B::bool_and(B::bool_and(not_blank, not_equal_s_m2), s_ge_2) +} + +/// Create a mask for valid s positions: s < 2 * target_length + 1 +fn create_s_mask( + target_lengths: &IntTensor, + batch_size: usize, + max_l_prime_len: usize, + device: &B::Device, + int_dtype: burn_std::IntDType, + bool_dtype: burn_std::BoolDType, +) -> BoolTensor { + let col_indices = B::int_arange(0..max_l_prime_len as i64, device, int_dtype); + let col_indices = B::int_reshape(col_indices, Shape::new([1, max_l_prime_len])); + let col_indices = B::int_expand(col_indices, Shape::new([batch_size, max_l_prime_len])); + + let lengths = B::int_mul_scalar(target_lengths.clone(), 2.into()); + let lengths = B::int_add_scalar(lengths, 1.into()); + let lengths = B::int_reshape(lengths, Shape::new([batch_size, 1])); + let lengths = B::int_expand(lengths, Shape::new([batch_size, max_l_prime_len])); + + B::int_lower(col_indices, lengths, bool_dtype) +} diff --git a/crates/burn-backend/src/backend/ops/modules/grid_sample.rs b/crates/burn-backend/src/backend/ops/modules/grid_sample.rs new file mode 100644 index 0000000..37ccc54 --- /dev/null +++ b/crates/burn-backend/src/backend/ops/modules/grid_sample.rs @@ -0,0 +1,320 @@ +use crate::{ + Backend, TensorMetadata, get_device_settings, + ops::{GridSampleOptions, GridSamplePaddingMode, InterpolateMode}, + tensor::FloatTensor, +}; +use alloc::vec; +use burn_std::{Shape, Slice}; + +/// Reference implementation of grid_sample_2d that supports all options. +/// +/// # Arguments +/// +/// * `tensor` - The tensor being sampled from, must be contiguous with shape (N, C, H_in, W_in) +/// * `grid` - A tensor of locations, with shape (N, H_out, W_out, 2). Values are [-1, 1]. +/// A [x = -1, y = -1] means top-left, and [x = 1, y = 1] means bottom-right +/// * `options` - Grid sampling options +/// +/// # Returns +/// +/// A tensor with shape (N, C, H_out, W_out) +pub fn float_grid_sample_2d_ref( + tensor: FloatTensor, + grid: FloatTensor, + options: GridSampleOptions, +) -> FloatTensor { + match options.mode { + InterpolateMode::Bilinear => float_grid_sample_2d_bilinear::( + tensor, + grid, + options.padding_mode, + options.align_corners, + ), + _ => todo!( + "Default implementation for grid_sample_2d with {:?} unimplemented", + options.mode + ), + } +} + +/// Bilinear grid sampling implementation. +fn float_grid_sample_2d_bilinear( + tensor: FloatTensor, + grid: FloatTensor, + padding_mode: GridSamplePaddingMode, + align_corners: bool, +) -> FloatTensor { + let n = tensor.shape()[0]; + let c = tensor.shape()[1]; + let h_in = tensor.shape()[2]; + let w_in = tensor.shape()[3]; + let h_out = grid.shape()[1]; + let w_out = grid.shape()[2]; + let spatial_in = h_in * w_in; + let spatial_out = h_out * w_out; + let device = tensor.device(); + + // Separate x and y coordinates from grid + // shape: (N, H_out, W_out, 1) + let grid_x_slice = vec![ + Slice::new(0, Some(n as isize), 1), + Slice::new(0, Some(h_out as isize), 1), + Slice::new(0, Some(w_out as isize), 1), + Slice::new(0, Some(1), 1), + ]; + let grid_y_slice = vec![ + Slice::new(0, Some(n as isize), 1), + Slice::new(0, Some(h_out as isize), 1), + Slice::new(0, Some(w_out as isize), 1), + Slice::new(1, Some(2), 1), + ]; + + let grid_x = B::float_slice(grid.clone(), &grid_x_slice); + let grid_x = B::float_reshape(grid_x, Shape::new([n, 1, h_out, w_out])); + let grid_y = B::float_slice(grid.clone(), &grid_y_slice); + let grid_y = B::float_reshape(grid_y, Shape::new([n, 1, h_out, w_out])); + + // Convert normalized grid coordinates [-1, 1] to pixel coordinates + let w_in_f = w_in as f64; + let h_in_f = h_in as f64; + + let (grid_x, grid_y) = if align_corners { + // align_corners=true: x_pixel = (x_norm + 1) * (width - 1) / 2 + // Maps -1 to 0 and 1 to width - 1 + let grid_x = B::float_add_scalar(grid_x, 1f32.into()); + let grid_x = B::float_mul_scalar(grid_x, ((w_in_f - 1.0) / 2.0).into()); + + let grid_y = B::float_add_scalar(grid_y, 1f32.into()); + let grid_y = B::float_mul_scalar(grid_y, ((h_in_f - 1.0) / 2.0).into()); + + (grid_x, grid_y) + } else { + // align_corners=false: x_pixel = (x_norm + 1) * width / 2 - 0.5 + // Maps -1 to -0.5 and 1 to width - 0.5 + let grid_x = B::float_add_scalar(grid_x, 1f32.into()); + let grid_x = B::float_mul_scalar(grid_x, (w_in_f / 2.0).into()); + let grid_x = B::float_sub_scalar(grid_x, 0.5f32.into()); + + let grid_y = B::float_add_scalar(grid_y, 1f32.into()); + let grid_y = B::float_mul_scalar(grid_y, (h_in_f / 2.0).into()); + let grid_y = B::float_sub_scalar(grid_y, 0.5f32.into()); + + (grid_x, grid_y) + }; + + // Apply padding mode to coordinates + let (grid_x, grid_y) = match padding_mode { + GridSamplePaddingMode::Border => { + // Clamp coordinates to valid range [0, size-1] + let grid_x = B::float_clamp(grid_x, 0f32.into(), ((w_in - 1) as f32).into()); + let grid_y = B::float_clamp(grid_y, 0f32.into(), ((h_in - 1) as f32).into()); + (grid_x, grid_y) + } + GridSamplePaddingMode::Reflection => { + // Reflect coordinates at boundaries + let grid_x = reflect_coordinates::(grid_x, w_in_f, align_corners); + let grid_y = reflect_coordinates::(grid_y, h_in_f, align_corners); + (grid_x, grid_y) + } + GridSamplePaddingMode::Zeros => { + // Keep coordinates as-is, we'll mask out-of-bounds later + (grid_x, grid_y) + } + }; + + // Get floor indices for the four corners + let grid_x_floored = B::float_floor(grid_x.clone()); + let grid_y_floored = B::float_floor(grid_y.clone()); + + // Compute interpolation weights (fractional part) + let x_frac = B::float_sub(grid_x.clone(), grid_x_floored.clone()); + let y_frac = B::float_sub(grid_y.clone(), grid_y_floored.clone()); + + // Convert to integer indices + let settings = get_device_settings::(&device); + let x0 = B::float_into_int(grid_x_floored.clone(), settings.int_dtype); + let y0 = B::float_into_int(grid_y_floored.clone(), settings.int_dtype); + let x1 = B::float_into_int( + B::float_add_scalar(grid_x_floored, 1f32.into()), + settings.int_dtype, + ); + let y1 = B::float_into_int( + B::float_add_scalar(grid_y_floored, 1f32.into()), + settings.int_dtype, + ); + + // Create masks for out-of-bounds coordinates (only used for zeros padding) + let (mask_00, mask_01, mask_10, mask_11) = if padding_mode == GridSamplePaddingMode::Zeros { + let x0_valid = B::int_greater_equal_elem(x0.clone(), 0.into(), settings.bool_dtype); + let x0_valid = B::bool_and( + x0_valid, + B::int_lower_elem(x0.clone(), (w_in as i32).into(), settings.bool_dtype), + ); + let x1_valid = B::int_greater_equal_elem(x1.clone(), 0.into(), settings.bool_dtype); + let x1_valid = B::bool_and( + x1_valid, + B::int_lower_elem(x1.clone(), (w_in as i32).into(), settings.bool_dtype), + ); + let y0_valid = B::int_greater_equal_elem(y0.clone(), 0.into(), settings.bool_dtype); + let y0_valid = B::bool_and( + y0_valid, + B::int_lower_elem(y0.clone(), (h_in as i32).into(), settings.bool_dtype), + ); + let y1_valid = B::int_greater_equal_elem(y1.clone(), 0.into(), settings.bool_dtype); + let y1_valid = B::bool_and( + y1_valid, + B::int_lower_elem(y1.clone(), (h_in as i32).into(), settings.bool_dtype), + ); + + ( + Some(B::bool_and(x0_valid.clone(), y0_valid.clone())), + Some(B::bool_and(x0_valid.clone(), y1_valid.clone())), + Some(B::bool_and(x1_valid.clone(), y0_valid)), + Some(B::bool_and(x1_valid, y1_valid)), + ) + } else { + (None, None, None, None) + }; + + // Clamp indices to valid range for gather + let x0_clamped = B::int_clamp(x0, 0.into(), ((w_in - 1) as i32).into()); + let x1_clamped = B::int_clamp(x1, 0.into(), ((w_in - 1) as i32).into()); + let y0_clamped = B::int_clamp(y0, 0.into(), ((h_in - 1) as i32).into()); + let y1_clamped = B::int_clamp(y1, 0.into(), ((h_in - 1) as i32).into()); + + // Linear indices: idx = y * W_in + x + let w_in_scalar: i32 = w_in as i32; + let idx_00 = B::int_add( + B::int_mul_scalar(y0_clamped.clone(), w_in_scalar.into()), + x0_clamped.clone(), + ); + let idx_01 = B::int_add( + B::int_mul_scalar(y1_clamped.clone(), w_in_scalar.into()), + x0_clamped, + ); + let idx_10 = B::int_add( + B::int_mul_scalar(y0_clamped, w_in_scalar.into()), + x1_clamped.clone(), + ); + let idx_11 = B::int_add( + B::int_mul_scalar(y1_clamped, w_in_scalar.into()), + x1_clamped, + ); + + // [N, 1, H_out, W_out] -> [N, 1, H_out * W_out] + let idx_00 = B::int_reshape(idx_00, Shape::new([n, 1, spatial_out])); + let idx_01 = B::int_reshape(idx_01, Shape::new([n, 1, spatial_out])); + let idx_10 = B::int_reshape(idx_10, Shape::new([n, 1, spatial_out])); + let idx_11 = B::int_reshape(idx_11, Shape::new([n, 1, spatial_out])); + + // [N, 1, spatial] -> [N, C, spatial] + let idx_00 = B::int_expand(idx_00, Shape::new([n, c, spatial_out])); + let idx_01 = B::int_expand(idx_01, Shape::new([n, c, spatial_out])); + let idx_10 = B::int_expand(idx_10, Shape::new([n, c, spatial_out])); + let idx_11 = B::int_expand(idx_11, Shape::new([n, c, spatial_out])); + + let tensor_flat = B::float_reshape(tensor, Shape::new([n, c, spatial_in])); + + let sample_00 = B::float_gather(2, tensor_flat.clone(), idx_00); + let sample_01 = B::float_gather(2, tensor_flat.clone(), idx_01); + let sample_10 = B::float_gather(2, tensor_flat.clone(), idx_10); + let sample_11 = B::float_gather(2, tensor_flat, idx_11); + + // Reshape samples to (N, C, H_out, W_out) + let sample_00 = B::float_reshape(sample_00, Shape::new([n, c, h_out, w_out])); + let sample_01 = B::float_reshape(sample_01, Shape::new([n, c, h_out, w_out])); + let sample_10 = B::float_reshape(sample_10, Shape::new([n, c, h_out, w_out])); + let sample_11 = B::float_reshape(sample_11, Shape::new([n, c, h_out, w_out])); + + // Apply masks for zeros padding (set out-of-bounds samples to 0) + let (sample_00, sample_01, sample_10, sample_11) = + if padding_mode == GridSamplePaddingMode::Zeros { + let mask_00 = mask_00.unwrap(); + let mask_01 = mask_01.unwrap(); + let mask_10 = mask_10.unwrap(); + let mask_11 = mask_11.unwrap(); + + let mask_00_inv = B::bool_not(mask_00); + let mask_00_inv = B::bool_reshape(mask_00_inv, Shape::new([n, 1, h_out, w_out])); + let mask_00_inv = B::bool_expand(mask_00_inv, Shape::new([n, c, h_out, w_out])); + let mask_01_inv = B::bool_not(mask_01); + let mask_01_inv = B::bool_reshape(mask_01_inv, Shape::new([n, 1, h_out, w_out])); + let mask_01_inv = B::bool_expand(mask_01_inv, Shape::new([n, c, h_out, w_out])); + let mask_10_inv = B::bool_not(mask_10); + let mask_10_inv = B::bool_reshape(mask_10_inv, Shape::new([n, 1, h_out, w_out])); + let mask_10_inv = B::bool_expand(mask_10_inv, Shape::new([n, c, h_out, w_out])); + let mask_11_inv = B::bool_not(mask_11); + let mask_11_inv = B::bool_reshape(mask_11_inv, Shape::new([n, 1, h_out, w_out])); + let mask_11_inv = B::bool_expand(mask_11_inv, Shape::new([n, c, h_out, w_out])); + + ( + B::float_mask_fill(sample_00, mask_00_inv, 0f32.into()), + B::float_mask_fill(sample_01, mask_01_inv, 0f32.into()), + B::float_mask_fill(sample_10, mask_10_inv, 0f32.into()), + B::float_mask_fill(sample_11, mask_11_inv, 0f32.into()), + ) + } else { + (sample_00, sample_01, sample_10, sample_11) + }; + + // Compute bilinear interpolation weights + let one_minus_x = B::float_neg(x_frac.clone()); + let one_minus_x = B::float_add_scalar(one_minus_x, 1f32.into()); + + let one_minus_y = B::float_neg(y_frac.clone()); + let one_minus_y = B::float_add_scalar(one_minus_y, 1f32.into()); + + let weight_00 = B::float_mul(one_minus_x.clone(), one_minus_y.clone()); + let weight_01 = B::float_mul(one_minus_x.clone(), y_frac.clone()); + let weight_10 = B::float_mul(x_frac.clone(), one_minus_y); + let weight_11 = B::float_mul(x_frac, y_frac); + + // Bilinear interpolation + let result = B::float_mul(sample_00, weight_00); + let result = B::float_add(result, B::float_mul(sample_01, weight_01)); + let result = B::float_add(result, B::float_mul(sample_10, weight_10)); + + B::float_add(result, B::float_mul(sample_11, weight_11)) +} + +/// Reflect coordinates at boundaries using a triangle wave pattern. +/// +/// For align_corners=true: reflects within [0, size-1] +/// For align_corners=false: reflects within [-0.5, size-0.5] +fn reflect_coordinates( + coords: FloatTensor, + size: f64, + align_corners: bool, +) -> FloatTensor { + let (min_val, max_val) = if align_corners { + (0.0f32, (size - 1.0) as f32) + } else { + (-0.5f32, (size - 0.5) as f32) + }; + + let span = max_val - min_val; + if span <= 0.0 { + // Edge case: size is 1, just return min_val everywhere + let zeros = B::float_mul_scalar(coords, 0f32.into()); + return B::float_add_scalar(zeros, min_val.into()); + } + + // Triangle wave formula: span - |((x mod 2*span) - span)| + min_val + let period = 2.0 * span; + + // x = abs(coord - min_val) + let x = B::float_sub_scalar(coords, min_val.into()); + let x = B::float_abs(x); + + // x_mod = x - floor(x / period) * period + let x_div = B::float_div_scalar(x.clone(), period.into()); + let x_div_floor = B::float_floor(x_div); + let x_mod = B::float_sub(x, B::float_mul_scalar(x_div_floor, period.into())); + + // result = span - abs(x_mod - span) + min_val + let diff = B::float_sub_scalar(x_mod, span.into()); + let abs_diff = B::float_abs(diff); + let reflected = B::float_sub_scalar(abs_diff, span.into()); + let reflected = B::float_neg(reflected); + B::float_add_scalar(reflected, min_val.into()) +} diff --git a/crates/burn-backend/src/backend/ops/modules/linear.rs b/crates/burn-backend/src/backend/ops/modules/linear.rs new file mode 100644 index 0000000..e43b5f1 --- /dev/null +++ b/crates/burn-backend/src/backend/ops/modules/linear.rs @@ -0,0 +1,128 @@ +use alloc::vec; +use alloc::vec::Vec; + +use crate::tensor::FloatTensor; +use crate::{Backend, TensorMetadata}; +use burn_std::{MatmulTransformAction, MatmulTransformAnalysis, MatmulTransformPolicy, Shape}; + +/// Default [linear](crate::ops::ModuleOps::linear) forward implementation. +/// +/// Computes `y = x @ weight [+ bias]`. +/// +/// The weight is shared by every leading dim of x, so when the batched matmul +/// would tile poorly, the [transform policy](MatmulTransformPolicy) folds the +/// leading dims into the rows and the product runs as one `[rows, d_input] @ +/// [d_input, d_output]` matmul. Otherwise weight and bias broadcast to x's rank +/// for a batched matmul. +pub(crate) fn linear( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, +) -> FloatTensor { + let shape = x.shape(); + let ndims = shape.num_dims(); + + let analysis = MatmulTransformAnalysis::from_shapes(&shape, &weight.shape()); + + match MatmulTransformPolicy::default().action(&analysis) { + MatmulTransformAction::MergeBatches { rows } => { + let d_input = shape[ndims - 1]; + let d_output = weight.shape()[1]; + + let x = B::float_reshape(x, Shape::new([rows, d_input])); + let mut output = B::float_matmul(x, weight); + + if let Some(bias) = bias { + let bias = B::float_reshape(bias, Shape::new([1, d_output])); + output = B::float_add(output, bias); + } + + let out_dims: Vec = (0..ndims - 1).map(|i| shape[i]).chain([d_output]).collect(); + B::float_reshape(output, Shape::from(out_dims)) + } + MatmulTransformAction::Keep => { + // Reshape weight [d_input, d_output] -> [1, ..., 1, d_input, d_output] + // for batch matmul. + let weight = unsqueeze_leading::(weight, ndims); + let output = B::float_matmul(x, weight); + + match bias { + Some(bias) => { + // Reshape bias [d_output] -> [1, ..., 1, d_output] to match + // output rank. + let bias = unsqueeze_leading::(bias, ndims); + B::float_add(output, bias) + } + None => output, + } + } + } +} + +/// Reshape a tensor by prepending size-1 dimensions until it has `target_ndims` dimensions. +fn unsqueeze_leading(tensor: FloatTensor, target_ndims: usize) -> FloatTensor { + let shape = tensor.shape(); + let ndims = shape.num_dims(); + if ndims >= target_ndims { + return tensor; + } + let mut new_dims = vec![1usize; target_ndims - ndims]; + for i in 0..ndims { + new_dims.push(shape[i]); + } + B::float_reshape(tensor, Shape::from(new_dims)) +} + +/// Reshape `tensor` `[..., d_last]` to `[num_rows, d_last]`, flattening every +/// leading dim into one. +fn flatten_leading(tensor: FloatTensor) -> FloatTensor { + let shape = tensor.shape(); + let ndims = shape.num_dims(); + if ndims == 2 { + return tensor; + } + let d_last = shape[ndims - 1]; + let num_rows = shape.num_elements() / d_last; + B::float_reshape(tensor, Shape::new([num_rows, d_last])) +} + +/// Default [linear_x_backward](crate::ops::ModuleOps::linear_x_backward) implementation. +/// +/// Computes `dx = output_grad @ weight^T` — a [linear] forward against the +/// transposed weight, so it shares the transform policy. +pub(crate) fn linear_x_backward( + weight: FloatTensor, + output_grad: FloatTensor, +) -> FloatTensor { + // weight is [d_input, d_output], transpose to [d_output, d_input] + let weight = B::float_swap_dims(weight, 0, 1); + linear::(output_grad, weight, None) +} + +/// Default [linear_weight_backward](crate::ops::ModuleOps::linear_weight_backward) implementation. +/// +/// Computes `dW = x^T @ output_grad` summed over batch dimensions, as a single +/// general matmul on the flattened operands: `[d_input, num_rows] @ [num_rows, +/// d_output]` — the flattening performs the batch reduction. +pub(crate) fn linear_weight_backward( + x: FloatTensor, + output_grad: FloatTensor, +) -> FloatTensor { + let x = B::float_swap_dims(flatten_leading::(x), 0, 1); + let output_grad = flatten_leading::(output_grad); + B::float_matmul(x, output_grad) +} + +/// Default [linear_bias_backward](crate::ops::ModuleOps::linear_bias_backward) implementation. +/// +/// Computes `db = sum(output_grad)` over all dimensions except the last, as a +/// single reduction on the flattened gradient. +pub(crate) fn linear_bias_backward(output_grad: FloatTensor) -> FloatTensor { + let shape = output_grad.shape(); + let d_output = shape[shape.num_dims() - 1]; + + let grad = flatten_leading::(output_grad); + // float_sum_dim preserves rank (keepdim): [num_rows, d_output] -> [1, d_output]. + let grad = B::float_sum_dim(grad, 0); + B::float_reshape(grad, Shape::new([d_output])) +} diff --git a/crates/burn-backend/src/backend/ops/modules/mod.rs b/crates/burn-backend/src/backend/ops/modules/mod.rs new file mode 100644 index 0000000..160ef4b --- /dev/null +++ b/crates/burn-backend/src/backend/ops/modules/mod.rs @@ -0,0 +1,24 @@ +/// Module with convolution operations. +pub mod conv; + +/// Module with linear operations. +pub mod linear; + +/// Module with attention operations. +pub mod attention; + +/// Module with CTC loss operations. +pub mod ctc; + +/// Module with unfold operations. +pub mod unfold; + +/// Module with pooling operations. +pub mod pool; + +/// Module for grid_sample operations +pub mod grid_sample; + +mod base; + +pub use base::*; diff --git a/crates/burn-backend/src/backend/ops/modules/pool.rs b/crates/burn-backend/src/backend/ops/modules/pool.rs new file mode 100644 index 0000000..9883168 --- /dev/null +++ b/crates/burn-backend/src/backend/ops/modules/pool.rs @@ -0,0 +1,178 @@ +use crate::tensor::{FloatTensor, IntTensor}; +use crate::{Backend, TensorMetadata}; +use burn_std::{IntDType, Shape}; + +use super::{MaxPool1dBackward, MaxPool1dWithIndices}; + +pub(crate) fn avg_pool1d_from_2d( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + count_include_pad: bool, + ceil_mode: bool, +) -> FloatTensor { + let [batch_size, channels, length] = x.shape().dims(); + + let x = B::float_reshape(x, Shape::from([batch_size, channels, length, 1])); + let x = B::avg_pool2d( + x, + [kernel_size, 1], + [stride, 1], + [padding, 0], + count_include_pad, + ceil_mode, + ); + + let [batch_size, channels, length, _] = x.shape().dims(); + + B::float_reshape(x, Shape::from([batch_size, channels, length])) +} + +pub(crate) fn avg_pool1d_backward_from_2d( + x: FloatTensor, + grad: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + count_include_pad: bool, + ceil_mode: bool, +) -> FloatTensor { + let [batch_size, channels, length_in] = x.shape().dims(); + let [_, _, length_out] = grad.shape().dims(); + + let x = B::float_reshape(x, Shape::from([batch_size, channels, length_in, 1])); + let grad_x = B::float_reshape(grad, Shape::from([batch_size, channels, length_out, 1])); + + let grad_x = B::avg_pool2d_backward( + x, + grad_x, + [kernel_size, 1], + [stride, 1], + [padding, 0], + count_include_pad, + ceil_mode, + ); + + B::float_reshape(grad_x, Shape::from([batch_size, channels, length_in])) +} + +pub(crate) fn adaptive_avg_pool1d_from_2d( + x: FloatTensor, + output_size: usize, +) -> FloatTensor { + let [batch_size, channels, length] = x.shape().dims(); + + let x = B::float_reshape(x, Shape::from([batch_size, channels, length, 1])); + let x = B::adaptive_avg_pool2d(x, [output_size, 1]); + + let [batch_size, channels, length, _] = x.shape().dims(); + + B::float_reshape(x, Shape::from([batch_size, channels, length])) +} + +pub(crate) fn adaptive_avg_pool1d_backward_from_2d( + x: FloatTensor, + grad: FloatTensor, +) -> FloatTensor { + let [batch_size, channels, length_in] = x.shape().dims(); + let [_, _, length_out] = grad.shape().dims(); + + let x = B::float_reshape(x, Shape::from([batch_size, channels, length_in, 1])); + let grad_x = B::float_reshape(grad, Shape::from([batch_size, channels, length_out, 1])); + + let grad_x = B::adaptive_avg_pool2d_backward(x, grad_x); + + B::float_reshape(grad_x, Shape::from([batch_size, channels, length_in])) +} + +pub(crate) fn max_pool1d_from_2d( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, +) -> FloatTensor { + let [batch_size, channels, length] = x.shape().dims(); + + let x = B::float_reshape(x, Shape::from([batch_size, channels, length, 1])); + let x = B::max_pool2d( + x, + [kernel_size, 1], + [stride, 1], + [padding, 0], + [dilation, 1], + ceil_mode, + ); + + let [batch_size, channels, length, _] = x.shape().dims(); + + B::float_reshape(x, Shape::from([batch_size, channels, length])) +} + +pub(crate) fn max_pool1d_with_indices_from_2d( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + indices_dtype: IntDType, +) -> MaxPool1dWithIndices { + let [batch_size, channels, length] = x.shape().dims(); + + let x = B::float_reshape(x, Shape::from([batch_size, channels, 1, length])); + let x = B::max_pool2d_with_indices( + x, + [1, kernel_size], + [1, stride], + [0, padding], + [1, dilation], + ceil_mode, + indices_dtype, + ); + let [batch_size, channels, _, length] = x.output.shape().dims(); + let output = B::float_reshape(x.output, Shape::from([batch_size, channels, length])); + let indices = B::int_reshape(x.indices, Shape::from([batch_size, channels, length])); + MaxPool1dWithIndices::new(output, indices) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn max_pool1d_with_indices_backward_from_2d( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + output_grad: FloatTensor, + indices: IntTensor, +) -> MaxPool1dBackward { + let [batch_size, channels, length_in] = x.shape().dims(); + let [_, _, length_out] = output_grad.shape().dims(); + + let x = B::float_reshape(x, Shape::from([batch_size, channels, length_in, 1])); + let grad_x = B::float_reshape( + output_grad, + Shape::from([batch_size, channels, length_out, 1]), + ); + let indices = B::int_reshape(indices, Shape::from([batch_size, channels, length_out, 1])); + + let grad_x = B::max_pool2d_with_indices_backward( + x, + [kernel_size, 1], + [stride, 1], + [padding, 0], + [dilation, 1], + ceil_mode, + grad_x, + indices, + ) + .x_grad; + + MaxPool1dBackward::new(B::float_reshape( + grad_x, + Shape::from([batch_size, channels, length_in]), + )) +} diff --git a/crates/burn-backend/src/backend/ops/modules/unfold.rs b/crates/burn-backend/src/backend/ops/modules/unfold.rs new file mode 100644 index 0000000..b25cfc9 --- /dev/null +++ b/crates/burn-backend/src/backend/ops/modules/unfold.rs @@ -0,0 +1,146 @@ +use super::{ConvOptions, UnfoldOptions}; +use crate::tensor::FloatTensor; +use crate::{Backend, TensorData, TensorMetadata}; +use alloc::vec; +use burn_std::{DType, Shape}; + +/// Constructs a special weight tensor used for unfolding. +/// +/// # Notes +/// +/// The idea behind using convolution for unfolding is to leverage the sliding window mechanism of +/// convolution. By creating a weight tensor with ones in a particular pattern, we are able to borrow +/// the convolution operation's mechanism as it moves across the input tensor, picking up the desired +/// values in the pattern of the unfolding operation. +pub(crate) fn create_unfolding_weight( + in_channels: usize, + kernel_size: [usize; 2], + device: &B::Device, + dtype: DType, +) -> FloatTensor { + let shape = Shape::new([ + in_channels * kernel_size[0] * kernel_size[1], + in_channels, + kernel_size[0], + kernel_size[1], + ]); + + let mut strides = [0; 4]; + let mut current = 1; + shape.iter().enumerate().rev().for_each(|(index, val)| { + strides[index] = current; + current *= val; + }); + + let num_elements = shape.num_elements(); + + let mut weight = vec![0f32; num_elements]; + + for k in 0..in_channels { + for i in 0..kernel_size[0] { + for j in 0..kernel_size[1] { + let output_channel = k * kernel_size[0] * kernel_size[1] + i * kernel_size[1] + j; + let index = + output_channel * strides[0] + k * strides[1] + i * strides[2] + j * strides[3]; + + weight[index] = 1.; + } + } + } + + B::float_from_data(TensorData::new(weight, shape).convert_dtype(dtype), device) +} + +/// Compute the unfold4d operation using the conv2d operations. +pub(crate) fn unfold4d_using_conv2d( + x: FloatTensor, + kernel_size: [usize; 2], + options: UnfoldOptions, +) -> FloatTensor { + let [_batch_size, in_channels, _in_height, _in_width] = x.shape().dims(); + let weight = create_unfolding_weight::(in_channels, kernel_size, &x.device(), x.dtype()); + let unfolded = B::conv2d( + x, + weight, + None, + ConvOptions::new(options.stride, options.padding, options.dilation, 1), + ); + + let [batch_size, channels_out, out_height, out_width] = unfolded.shape().dims(); + + B::float_reshape( + unfolded, + Shape::new([batch_size, channels_out, out_height * out_width]), + ) +} + +/// Calculate the number of unfolding windows that can be extracted from a dimension of given size. +pub fn calculate_unfold_windows(dim_size: usize, window_size: usize, step_size: usize) -> usize { + assert!(step_size > 0); + let x = dim_size + step_size; + if x < window_size { + 0 + } else { + (x - window_size) / step_size + } +} + +/// Calculate the output shape for an unfold operation. +/// +/// The operation yields a view with all complete windows of size `size` in dimension `dim`; +/// where windows are advanced by `step` at each index. +/// +/// The number of windows is `max(0, (shape[dim] - size).ceil_div(step))`. +/// +/// # Arguments +/// +/// * `shape` - The input shape to unfold; of shape ``[pre=..., dim shape, post=...]`` +/// * `dim` - the dimension to unfold. +/// * `size` - the size of each unfolded window. +/// * `step` - the step between each window. +/// +/// # Returns +/// +/// A shape with ``[pre=..., windows, post=..., size]``. +pub fn calculate_unfold_shape>( + shape: S, + dim: usize, + size: usize, + step: usize, +) -> Shape { + let mut shape = shape.into(); + let d_shape = shape[dim]; + let windows = calculate_unfold_windows(d_shape, size, step); + shape[dim] = windows; + shape.push(size); + + shape +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_calculate_unfold_windows() { + assert_eq!(calculate_unfold_windows(2, 5, 1), 0); + + assert_eq!(calculate_unfold_windows(2, 3, 1), 0); + assert_eq!(calculate_unfold_windows(3, 3, 1), 1); + assert_eq!(calculate_unfold_windows(4, 3, 1), 2); + assert_eq!(calculate_unfold_windows(5, 3, 1), 3); + + assert_eq!(calculate_unfold_windows(2, 3, 2), 0); + assert_eq!(calculate_unfold_windows(3, 3, 2), 1); + assert_eq!(calculate_unfold_windows(4, 3, 2), 1); + assert_eq!(calculate_unfold_windows(5, 3, 2), 2); + } + + #[test] + fn test_calculate_unfold_shape() { + assert_eq!( + calculate_unfold_shape([2, 6, 6], 1, 3, 2), + Shape::new([2, 2, 6, 3]) + ); + } +} diff --git a/crates/burn-backend/src/backend/ops/qtensor.rs b/crates/burn-backend/src/backend/ops/qtensor.rs new file mode 100644 index 0000000..69539e3 --- /dev/null +++ b/crates/burn-backend/src/backend/ops/qtensor.rs @@ -0,0 +1,1292 @@ +use alloc::vec::Vec; +use burn_std::{ + BoolDType, FloatDType, IntDType, Shape, Slice, + quantization::{QuantPropagation, QuantScheme}, +}; + +use crate::{ + Backend, ExecutionError, TensorData, TensorMetadata, TensorPrimitive, get_device_settings, +}; +use crate::{ + Scalar, + quantization::{Calibration, QuantizationParametersPrimitive, compute_q_params, compute_range}, + tensor::{BoolTensor, Device, FloatTensor, IntTensor, QuantizedTensor}, +}; + +/// Automatically applies `dequantization -> float operation -> quantization`. +/// +/// Used for tensor ops that should always return a quantized output. +#[macro_export] +macro_rules! dequant_op_quant { + // Binary tensor float op w/ lhs & rhs + ( + float_op $float_op:expr, $t1:expr, $t2:expr + ) => {{ + // Heuristic: prioritize lhs scheme + let scheme = $t1.scheme().clone(); + + let t1_f = Self::dequantize($t1); + let t2_f = Self::dequantize($t2); + #[allow(clippy::redundant_closure_call)] + let out_f = $float_op(t1_f, t2_f); + + Self::quantize_dynamic(out_f, &scheme) + }}; + // Unary tensor float op + ( + float_op $float_op:expr, $tensor:expr + ) => {{ + let scheme = $tensor.scheme().clone(); + let dtype = get_device_settings::(&$tensor.device()).float_dtype; + + let tensor_f = Self::dequantize($tensor, dtype); + #[allow(clippy::redundant_closure_call)] + let out_f = $float_op(tensor_f); + + Self::quantize_dynamic(out_f, &scheme) + }}; +} + +/// Automatically applies `dequantization -> float operation [-> quantization]`. +/// +/// The output quantization step is optional. +/// It is only performed when the input quantization scheme is propagated. +#[macro_export] +macro_rules! dequant_op_flow { + // Binary tensor float op w/ lhs & rhs + ( + float_op $float_op:expr, $t1:expr, $t2:expr + ) => {{ + // Heuristic: prioritize lhs scheme + let scheme = $t1.scheme().clone(); + let settings = get_device_settings::(&$t1.device()); + let dtype = settings.float_dtype; + let propagation = settings.quantization.propagation; + + let t1_f = Self::dequantize($t1, dtype); + let t2_f = Self::dequantize($t2, dtype); + #[allow(clippy::redundant_closure_call)] + let out_f = $float_op(t1_f, t2_f); + + match propagation { + QuantPropagation::Propagate => { + TensorPrimitive::QFloat(Self::quantize_dynamic(out_f, &scheme)) + } + QuantPropagation::Inhibit => TensorPrimitive::Float(out_f), + } + }}; + // Unary tensor float op + ( + float_op $float_op:expr, $tensor:expr + ) => {{ + let scheme = $tensor.scheme().clone(); + let settings = get_device_settings::(&$tensor.device()); + let dtype = settings.float_dtype; + let propagation = settings.quantization.propagation; + + let tensor_f = Self::dequantize($tensor, dtype); + #[allow(clippy::redundant_closure_call)] + let out_f = $float_op(tensor_f); + + match propagation { + QuantPropagation::Propagate => { + TensorPrimitive::QFloat(Self::quantize_dynamic(out_f, &scheme)) + } + QuantPropagation::Inhibit => TensorPrimitive::Float(out_f), + } + }}; +} + +/// Operations on quantized tensors. +/// +/// # Return Type Semantics +/// +/// The return type of each operation indicates how quantization is handled: +/// +/// ## [`QuantizedTensor`] +/// If the method returns a `QuantizedTensor`, the operation is expected to preserve the quantized +/// representation. Implementations should avoid dequantizing when possible to maintain performance. +/// For example, shape or layout changes such as expand or transpose preserve quantization. +/// +/// *Note: while this currently doesn't affect the quantized tensor parameters (only per-tensor is +/// supported at the time of writing), other quantization levels (e.g., per-block) may require re-ordering +/// the quantization parameters to match the new layout.* +/// +/// +/// ## [`TensorPrimitive`] +/// If the method returns a `TensorPrimitive` enum, the return type should align with propagation +/// strategy specified in the quantization scheme. The output should remain quantized ([`TensorPrimitive::QFloat`]) +/// returned in floating-point form ([`TensorPrimitive::Float`]). +/// +/// This distinction allows for fine-grained control over mixed-precision flows while still operating +/// through a unified API. +pub trait QTensorOps { + /// Creates a new tensor from the data structure. + /// + /// # Arguments + /// + /// * `data` - The data structure. + /// * `device` - The device to create the tensor on. + /// + /// # Returns + /// + /// The tensor with the given data. + fn q_from_data(data: TensorData, device: &Device) -> QuantizedTensor; + + /// Convert the tensor to a lower precision data type based on the quantization scheme and parameters. + fn quantize( + tensor: FloatTensor, + scheme: &QuantScheme, + qparams: QuantizationParametersPrimitive, + ) -> QuantizedTensor; + + /// Dynamically convert the tensor to a lower precision data type based on the quantization scheme. + fn quantize_dynamic(tensor: FloatTensor, scheme: &QuantScheme) -> QuantizedTensor { + // Dynamically compute min/max tensor range and qparams before quantizing + let (min, max) = compute_range::(scheme, tensor.clone(), &Calibration::MinMax); + let qparams = compute_q_params(scheme, min, max); + Self::quantize(tensor, scheme, qparams) + } + + /// Convert the tensor back to a higher precision data type. + fn dequantize(tensor: QuantizedTensor, dtype: FloatDType) -> FloatTensor; + + /// Moves the tensor to the given device. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `device` - The device to move the tensor to. + /// + /// # Returns + /// + /// The tensor on the given device. + fn q_to_device(tensor: QuantizedTensor, device: &Device) -> QuantizedTensor; + + /// Reshapes a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to reshape. + /// * `shape` - The new shape of the tensor. + /// + /// # Returns + /// + /// The tensor with the new shape. + fn q_reshape(tensor: QuantizedTensor, shape: Shape) -> QuantizedTensor; + + /// Converts the tensor to a data structure. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// + /// # Returns + /// + /// The data structure with the tensor's data. + fn q_into_data( + tensor: QuantizedTensor, + ) -> impl Future> + Send; + + /// Detaches a tensor from the computation graph. + fn q_detach(tensor: QuantizedTensor) -> QuantizedTensor { + // Should only be overridden by autodiff backends. + tensor + } + + /// Sets the `require_grad` flag of a tensor. + fn q_set_require_grad(tensor: QuantizedTensor, _require_grad: bool) -> QuantizedTensor { + // Should only be overridden by autodiff backends. + tensor + } + + /// Returns the `require_grad` flag of a tensor. + fn q_is_require_grad(_tensor: &QuantizedTensor) -> bool { + // Should only be overridden by autodiff backends. + false + } + + /// Broadcasts the `tensor` to the given `shape`. + fn q_expand(tensor: QuantizedTensor, shape: Shape) -> QuantizedTensor { + // Default implementation. Backends can expand on the quantized values when supported. + dequant_op_quant!(float_op | tensor | B::float_expand(tensor, shape), tensor) + } + + /// Transposes a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to transpose. + /// + /// # Returns + /// + /// The transposed tensor. + fn q_transpose(tensor: QuantizedTensor) -> QuantizedTensor { + let ndims = tensor.shape().num_dims(); + Self::q_swap_dims(tensor, ndims - 2, ndims - 1) + } + + /// Swaps two dimensions of a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to swap the dimensions of. + /// * `dim1` - The first dimension to swap. + /// * `dim2` - The second dimension to swap. + /// + /// # Returns + /// + /// The tensor with the dimensions swapped. + fn q_swap_dims(tensor: QuantizedTensor, dim1: usize, dim2: usize) -> QuantizedTensor; + + /// Permutes the dimensions of a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to permute the dimensions of. + /// * `axes` - The new order of the dimensions. + /// # Returns + /// + /// The tensor with the dimensions permuted. + fn q_permute(tensor: QuantizedTensor, axes: &[usize]) -> QuantizedTensor; + + /// Reverse the order of elements in a tensor along the given axes. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to reverse. + /// * `axes` - The axes to reverse. + /// + /// The tensor with the elements reversed. + fn q_flip(tensor: QuantizedTensor, axes: &[usize]) -> QuantizedTensor; + + /// Select tensor elements along the given dimension corresponding for the given indices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to select from. + /// * `dim` - The dimension to select from. + /// * `indices` - The indices to select. + /// + /// # Returns + /// + /// The selected elements. + fn q_select( + tensor: QuantizedTensor, + dim: usize, + indices: IntTensor, + ) -> QuantizedTensor { + // Default implementation. Backends can select on the quantized values when supported. + dequant_op_quant!( + float_op | tensor | B::float_select(tensor, dim, indices), + tensor + ) + } + + /// Select tensor elements corresponding to the given slices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to select from. + /// * `slices` - The slices specifying ranges and steps for each dimension. + /// + /// # Returns + /// + /// The selected elements in a new tensor. + fn q_slice(tensor: QuantizedTensor, slices: &[Slice]) -> QuantizedTensor { + // Default implementation. Backends can slice on the quantized values when supported. + dequant_op_quant!(float_op | tensor | B::float_slice(tensor, slices), tensor) + } + + /// Gather elements from a tensor. + /// + /// # Arguments + /// + /// * `dim` - The dimension to gather from. + /// * `tensor` - The tensor to gather from. + /// * `indices` - The indices to gather. + /// + /// # Returns + /// + /// The gathered elements. + fn q_gather( + dim: usize, + tensor: QuantizedTensor, + indices: IntTensor, + ) -> QuantizedTensor { + // Default implementation. Backends can gather on the quantized values when supported. + dequant_op_quant!( + float_op | tensor | B::float_gather(dim, tensor, indices), + tensor + ) + } + + /// Repeat the tensor along the given dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `dim` - The dimension to repeat. + /// * `times` - The number of times to repeat the dimension. + /// + /// # Returns + /// + /// The tensor with the given dimension repeated. + fn q_repeat_dim(tensor: QuantizedTensor, dim: usize, times: usize) -> QuantizedTensor { + dequant_op_quant!( + float_op | tensor | B::float_repeat_dim(tensor, dim, times), + tensor + ) + } + + /// Adds two tensors together. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// The result of adding the two tensors together. + fn q_add(lhs: QuantizedTensor, rhs: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | lhs, rhs | B::float_add(lhs, rhs), lhs, rhs) + } + + /// Adds a scalar to a tensor. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side scalar. + /// + /// # Returns + /// + /// The result of adding the scalar to the tensor. + fn q_add_scalar(lhs: QuantizedTensor, rhs: Scalar) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_add_scalar(tensor, rhs), lhs) + } + + /// Clamps a tensor under a minimum value. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to clamp. + /// * `min` - The minimum value. + /// + /// # Returns + /// + /// The clamped tensor. + fn q_clamp_min(tensor: QuantizedTensor, min: Scalar) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_clamp_min(tensor, min), tensor) + } + + /// Clamps a tensor over a maximum value. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to clamp. + /// * `max` - The maximum value. + /// + /// # Returns + /// + /// The clamped tensor. + fn q_clamp_max(tensor: QuantizedTensor, max: Scalar) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_clamp_max(tensor, max), tensor) + } + + /// Clamps a tensor between a minimum and maximum value. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to clamp. + /// * `min` - The minimum value. + /// * `max` - The maximum value. + /// + /// # Returns + /// + /// The clamped tensor. + fn q_clamp(tensor: QuantizedTensor, min: Scalar, max: Scalar) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_clamp(tensor, min, max), tensor) + } + + /// Subtracts two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// The result of subtracting the two tensors. + fn q_sub(lhs: QuantizedTensor, rhs: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | lhs, rhs | B::float_sub(lhs, rhs), lhs, rhs) + } + + /// Subtracts a scalar from a tensor. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side scalar. + /// + /// # Returns + /// + /// The result of subtracting the scalar from the tensor. + fn q_sub_scalar(lhs: QuantizedTensor, rhs: Scalar) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_sub_scalar(tensor, rhs), lhs) + } + + /// Multiplies two tensors together element-wise. + fn q_mul(lhs: QuantizedTensor, rhs: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | lhs, rhs | B::float_mul(lhs, rhs), lhs, rhs) + } + + /// Multiplies a tensor by a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side scalar. + /// + /// # Returns + /// + /// The result of multiplying the tensor by the scalar. + fn q_mul_scalar(lhs: QuantizedTensor, rhs: Scalar) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_mul_scalar(tensor, rhs), lhs) + } + + /// Divides two tensors element-wise. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// The result of dividing the two tensors. + fn q_div(lhs: QuantizedTensor, rhs: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | lhs, rhs | B::float_div(lhs, rhs), lhs, rhs) + } + + /// Divides a tensor by a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side scalar. + /// + /// # Returns + /// + /// The result of dividing the tensor by the scalar. + fn q_div_scalar(lhs: QuantizedTensor, rhs: Scalar) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_div_scalar(tensor, rhs), lhs) + } + + /// Multiplies two tensors together using matrix multiplication. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// The result of multiplying the two tensors together using matrix multiplication. + fn q_matmul(lhs: TensorPrimitive, rhs: TensorPrimitive) -> TensorPrimitive { + let mut propagation = QuantPropagation::Inhibit; + let mut scheme = QuantScheme::default(); + + // Pick a target dtype for any dequantization. If either operand is already + // a Float tensor, take its dtype so a Float-QFloat (or QFloat-Float) pair + // ends up matching after dequantize and `float_matmul` doesn't see a + // dtype mismatch. Only when both operands are QFloat do we fall back to + // the device default. + let target_dtype: Option = match (&lhs, &rhs) { + (TensorPrimitive::Float(t), _) | (_, TensorPrimitive::Float(t)) => { + Some(t.dtype().into()) + } + _ => None, + }; + + let lhs = match lhs { + TensorPrimitive::Float(lhs) => lhs, + TensorPrimitive::QFloat(lhs) => { + let settings = get_device_settings::(&lhs.device()); + propagation = settings.quantization.propagation; + scheme = lhs.scheme(); + let float_dtype = target_dtype.unwrap_or(settings.float_dtype); + + Self::dequantize(lhs, float_dtype) + } + }; + let rhs = match rhs { + TensorPrimitive::Float(rhs) => rhs, + TensorPrimitive::QFloat(rhs) => { + let settings = get_device_settings::(&rhs.device()); + propagation = settings.quantization.propagation; + scheme = rhs.scheme(); + let float_dtype = target_dtype.unwrap_or(settings.float_dtype); + + Self::dequantize(rhs, float_dtype) + } + }; + + let out_f = B::float_matmul(lhs, rhs); + match propagation { + QuantPropagation::Propagate => { + TensorPrimitive::QFloat(::quantize_dynamic(out_f, &scheme)) + } + QuantPropagation::Inhibit => TensorPrimitive::Float(out_f), + } + } + + /// Negates a tensor element-wise. + fn q_neg(tensor: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_neg(tensor), tensor) + } + + /// Calculates the reciprocals element-wise + fn q_recip(tensor: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_recip(tensor), tensor) + } + + /// Sum of all elements in a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to sum. + /// + /// # Returns + /// + /// A scalar tensor with the sum of all elements in `tensor`. + fn q_sum(tensor: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_sum(tensor), tensor) + } + + /// Sum of all elements in a tensor along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to sum. + /// * `dim` - The dimension along which to sum. + /// + /// # Returns + /// + /// A tensor with the sum of all elements in `tensor` along `dim`. + fn q_sum_dim(tensor: QuantizedTensor, dim: usize) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_sum_dim(tensor, dim), tensor) + } + + /// Product of all elements in a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to product. + /// + /// # Returns + /// + /// A scalar tensor with the product of all elements in `tensor`. + fn q_prod(tensor: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_prod(tensor), tensor) + } + + /// Product of all elements in a tensor along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to product. + /// + /// # Returns + /// + /// A tensor with the product of all elements in `tensor` along `dim`. + fn q_prod_dim(tensor: QuantizedTensor, dim: usize) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_prod_dim(tensor, dim), tensor) + } + + /// Mean of all elements in a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to mean. + /// + /// # Returns + /// + /// A scalar tensor with the mean of all elements in `tensor`. + fn q_mean(tensor: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_mean(tensor), tensor) + } + + /// Mean of all elements in a tensor along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to mean. + /// * `dim` - The dimension along which to mean. + /// + /// # Returns + /// + /// A tensor with the mean of all elements in `tensor` along `dim`. + fn q_mean_dim(tensor: QuantizedTensor, dim: usize) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_mean_dim(tensor, dim), tensor) + } + + /// Computes the cumulative sum of elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the cumulative sum of. + /// * `dim` - The dimension along which to compute the cumulative sum. + /// + /// # Returns + /// + /// A tensor with the same shape where each element is the cumulative sum + /// of all elements up to and including that position along the dimension. + fn q_cumsum(tensor: QuantizedTensor, dim: usize) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_cumsum(tensor, dim), tensor) + } + + /// Computes the cumulative product of elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the cumulative product of. + /// * `dim` - The dimension along which to compute the cumulative product. + /// + /// # Returns + /// + /// A tensor with the same shape where each element is the cumulative product + /// of all elements up to and including that position along the dimension. + fn q_cumprod(tensor: QuantizedTensor, dim: usize) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_cumprod(tensor, dim), tensor) + } + + /// Computes the cumulative minimum of elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the cumulative minimum of. + /// * `dim` - The dimension along which to compute the cumulative minimum. + /// + /// # Returns + /// + /// A tensor with the same shape where each element is the minimum + /// of all elements up to and including that position along the dimension. + fn q_cummin(tensor: QuantizedTensor, dim: usize) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_cummin(tensor, dim), tensor) + } + + /// Computes the cumulative maximum of elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the cumulative maximum of. + /// * `dim` - The dimension along which to compute the cumulative maximum. + /// + /// # Returns + /// + /// A tensor with the same shape where each element is the maximum + /// of all elements up to and including that position along the dimension. + fn q_cummax(tensor: QuantizedTensor, dim: usize) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_cummax(tensor, dim), tensor) + } + + /// Returns a new tensor with exponential values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to exponentiate. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with exponential values. + fn q_exp(tensor: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_exp(tensor), tensor) + } + + /// Returns a new tensor with natural logarithm values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the logarithm of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with natural logarithm values. + fn q_log(tensor: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_log(tensor), tensor) + } + + /// Returns a new tensor with logarithm values of (1 + Xi). + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the logarithm of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with logarithm values of (1 + Xi). + fn q_log1p(tensor: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_log1p(tensor), tensor) + } + + /// Element-wise power with another tensor. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// The elements of `lhs` raised to the power of the elements of `rhs`. + fn q_powf(lhs: QuantizedTensor, rhs: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | lhs, rhs | B::float_powf(lhs, rhs), lhs, rhs) + } + + /// Element-wise power with an IntTensor. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side floatTensor. + /// + /// # Returns + /// + /// The elements of `lhs` raised to the value of `rhs`. Result is an IntTensor. + fn q_powi(lhs: QuantizedTensor, rhs: IntTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_powi(tensor, rhs), lhs) + } + + /// Element-wise power with an int scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side scalar. + /// + /// # Returns + /// + /// The elements of `lhs` raised to the value of `rhs`. + fn q_powi_scalar(lhs: QuantizedTensor, rhs: Scalar) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_powi_scalar(tensor, rhs), lhs) + } + + /// Element-wise power with a float scalar. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to exponentiate. + /// * `value` - The exponent. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with values raised to the power of `value`. + fn q_powf_scalar(tensor: QuantizedTensor, value: Scalar) -> TensorPrimitive { + dequant_op_flow!( + float_op | tensor | B::float_powf_scalar(tensor, value), + tensor + ) + } + + /// Returns a new tensor with square root values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the square root of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with square root values. + fn q_sqrt(tensor: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_sqrt(tensor), tensor) + } + + /// Returns a new tensor with absolute values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take absolute value of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with absolute values. + fn q_abs(tensor: QuantizedTensor) -> QuantizedTensor { + dequant_op_quant!(float_op | tensor | B::float_abs(tensor), tensor) + } + + /// Returns a new tensor with cosine values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the cosine of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with cosine values. + fn q_cos(tensor: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_cos(tensor), tensor) + } + + /// Returns a new tensor with sine values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the sine of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with sine values. + fn q_sin(tensor: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_sin(tensor), tensor) + } + + /// Returns a new tensor with tangent values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the tangent of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with tangent values. + fn q_tan(tensor: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_tan(tensor), tensor) + } + + /// Returns a new tensor with hyperbolic cosine values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the hyperbolic cosine of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with hyperbolic cosine values. + fn q_cosh(tensor: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_cosh(tensor), tensor) + } + + /// Returns a new tensor with hyperbolic sine values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the hyperbolic sine of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with hyperbolic sine values. + fn q_sinh(tensor: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_sinh(tensor), tensor) + } + + /// Returns a new tensor with hyperbolic tangent values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the hyperbolic tangent of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with hyperbolic tangent values. + fn q_tanh(tensor: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_tanh(tensor), tensor) + } + + /// Returns a new tensor with the error function values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the error function of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with error function values. + fn q_erf(tensor: QuantizedTensor) -> TensorPrimitive { + dequant_op_flow!(float_op | tensor | B::float_erf(tensor), tensor) + } + + /// Concatenates tensors along a dimension. + /// + /// # Arguments + /// + /// * `tensors` - The tensors to concatenate. + /// * `dim` - The dimension along which to concatenate. + /// + /// # Returns + /// + /// A tensor with the concatenated tensors along `dim`. + fn q_cat(tensors: Vec>, dim: usize) -> QuantizedTensor { + // Heuristic: prioritize first tensor scheme + let first = tensors.first().unwrap(); + let scheme = first.scheme(); + let dtype = get_device_settings::(&first.device()).float_dtype; + + let tensor_f = tensors + .into_iter() + .map(|tensor| Self::dequantize(tensor, dtype)) + .collect(); + + let out_f = B::float_cat(tensor_f, dim); + + Self::quantize_dynamic(out_f, &scheme) + } + + /// Gets the indices of the maximum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements of. + /// * `dim` - The dimension along which to get the maximum elements. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A tensor with the indices of the maximum elements of `tensor` along `dim`. + fn q_argmax(tensor: QuantizedTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + let dtype = get_device_settings::(&tensor.device()).float_dtype; + let tensor_f = Self::dequantize(tensor, dtype); + B::float_argmax(tensor_f, dim, out_dtype) + } + + /// Gets the indices of the k maximum elements of a tensor along an axis. + /// If two elements are equals, order them by the lowest indices + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the k maximum elements of. + /// * `dim` - The dimension along which to get the maximum elements. + /// * `k` - number of k maximums to keep + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A tensor with the indices of the `k` maximum elements of `tensor` along `dim`. + fn q_argtopk( + tensor: QuantizedTensor, + dim: usize, + k: usize, + out_dtype: IntDType, + ) -> IntTensor { + let dtype = get_device_settings::(&tensor.device()).float_dtype; + let tensor_f = Self::dequantize(tensor, dtype); + B::float_argtopk(tensor_f, dim, k, out_dtype) + } + + /// Gets the values of the k maximum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the k maximum elements of. + /// * `dim` - The dimension along which to get the maximum elements. + /// * `k` - number of k maximums to keep + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A tensor with the values of the `k` maximum elements of `tensor` along `dim`. + fn q_topk(tensor: QuantizedTensor, dim: usize, k: usize) -> QuantizedTensor { + dequant_op_quant!(float_op | tensor | B::float_topk(tensor, dim, k), tensor) + } + + /// Gets the indices of the minimum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the minimum elements of. + /// * `dim` - The dimension along which to get the minimum elements. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A tensor with the indices of the minimum elements of `tensor` along `dim`. + fn q_argmin(tensor: QuantizedTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + let dtype = get_device_settings::(&tensor.device()).float_dtype; + let tensor_f = Self::dequantize(tensor, dtype); + B::float_argmin(tensor_f, dim, out_dtype) + } + + /// Gets the maximum element of a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements of. + /// + /// # Returns + /// + /// A tensor with the maximum element of `tensor`. + fn q_max(tensor: QuantizedTensor) -> QuantizedTensor { + let shape = tensor.shape(); + let tensor = B::q_reshape(tensor, Shape::new([shape.num_elements()])); + + B::q_max_dim(tensor, 0) + } + + /// Gets the maximum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements of. + /// * `dim` - The dimension along which to get the maximum elements. + /// + /// # Returns + /// + /// A tensor with the maximum elements of `tensor` along `dim`. + fn q_max_dim(tensor: QuantizedTensor, dim: usize) -> QuantizedTensor { + let int_dtype = get_device_settings::(&tensor.device()).int_dtype; + let index = B::q_argmax(tensor.clone(), dim, int_dtype); + + B::q_gather(dim, tensor, index) + } + + /// Gets the maximum elements of a tensor along an axis and their indices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements of. + /// * `dim` - The dimension along which to get the maximum elements. + /// + /// # Returns + /// + /// A tuple with the maximum elements of `tensor` along `dim` and their indices. + fn q_max_dim_with_indices( + tensor: QuantizedTensor, + dim: usize, + out_dtype: IntDType, + ) -> (QuantizedTensor, IntTensor) { + let index = B::q_argmax(tensor.clone(), dim, out_dtype); + let values = B::q_gather(dim, tensor, index.clone()); + + (values, index) + } + + /// Gets the minimum element of a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the minimum elements of. + /// + /// # Returns + /// + /// A tensor with the minimum element of `tensor`. + fn q_min(tensor: QuantizedTensor) -> QuantizedTensor { + let shape = tensor.shape(); + let tensor = B::q_reshape(tensor, Shape::new([shape.num_elements()])); + + B::q_min_dim(tensor, 0) + } + + /// Gets the minimum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the minimum elements of. + /// * `dim` - The dimension along which to get the minimum elements. + /// + /// # Returns + /// + /// A tensor with the minimum elements of `tensor` along `dim`. + fn q_min_dim(tensor: QuantizedTensor, dim: usize) -> QuantizedTensor { + let int_dtype = get_device_settings::(&tensor.device()).int_dtype; + let index = B::q_argmin(tensor.clone(), dim, int_dtype); + + B::q_gather(dim, tensor, index) + } + + /// Gets the minimum elements of a tensor along an axis and their indices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the minimum elements of. + /// * `dim` - The dimension along which to get the minimum elements. + /// + /// # Returns + /// + /// A tuple with the minimum elements of `tensor` along `dim` and their indices. + fn q_min_dim_with_indices( + tensor: QuantizedTensor, + dim: usize, + out_dtype: IntDType, + ) -> (QuantizedTensor, IntTensor) { + let index = B::q_argmin(tensor.clone(), dim, out_dtype); + let values = B::q_gather(dim, tensor, index.clone()); + + (values, index) + } + + /// Gets the maximum element of a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements of. + /// + /// # Returns + /// + /// A tensor with the maximum element of `tensor`. + fn q_max_abs(tensor: QuantizedTensor) -> QuantizedTensor { + let shape = tensor.shape(); + let tensor = B::q_reshape(tensor, Shape::new([shape.num_elements()])); + + B::q_max_abs_dim(tensor, 0) + } + + /// Gets the maximum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements of. + /// * `dim` - The dimension along which to get the maximum elements. + /// + /// # Returns + /// + /// A tensor with the maximum elements of `tensor` along `dim`. + fn q_max_abs_dim(tensor: QuantizedTensor, dim: usize) -> QuantizedTensor { + let int_dtype = get_device_settings::(&tensor.device()).int_dtype; + let index = B::q_argmax(B::q_abs(tensor.clone()), dim, int_dtype); + + B::q_gather(dim, tensor, index) + } + + /// Tests if any element in the `tensor` evaluates to True. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// + /// # Returns + /// + /// A boolean tensor with a single element, True if any element in the tensor is True, False otherwise. + fn q_any(tensor: QuantizedTensor, out_dtype: BoolDType) -> BoolTensor { + let dtype = get_device_settings::(&tensor.device()).float_dtype; + let tensor_f = Self::dequantize(tensor, dtype); + B::float_any(tensor_f, out_dtype) + } + + /// Tests if any element in the float `tensor` evaluates to True along a given dimension `dim`. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// * `dim` - The axis along which to test. + /// + /// # Returns + /// + /// A boolean tensor `Tensor` with the same size as input `tensor`, except in the `dim` axis + /// where the size is 1. The elem in the `dim` axis is True if any element along this dim in the + /// input evaluates to True, False otherwise. + fn q_any_dim(tensor: QuantizedTensor, dim: usize, out_dtype: BoolDType) -> BoolTensor { + let dtype = get_device_settings::(&tensor.device()).float_dtype; + let tensor_f = Self::dequantize(tensor, dtype); + B::float_any_dim(tensor_f, dim, out_dtype) + } + + /// Tests if all elements in the `tensor` evaluate to True. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// + /// # Returns + /// + /// A boolean tensor `Tensor` with a single element, True if all elements in the input tensor + /// evaluate to True, False otherwise. + fn q_all(tensor: QuantizedTensor, out_dtype: BoolDType) -> BoolTensor { + let dtype = get_device_settings::(&tensor.device()).float_dtype; + let tensor_f = Self::dequantize(tensor, dtype); + B::float_all(tensor_f, out_dtype) + } + + /// Tests if all elements in the `tensor` evaluate to True along a given dimension `dim`. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// * `dim` - The axis along which to test. + /// + /// # Returns + /// + /// A boolean tensor `Tensor` with the same size as input `tensor`, except in the `dim` axis + /// where the size is 1. The elem in the `dim` axis is True if all elements along this dim in the input + /// evaluates to True, False otherwise. + fn q_all_dim(tensor: QuantizedTensor, dim: usize, out_dtype: BoolDType) -> BoolTensor { + let dtype = get_device_settings::(&tensor.device()).float_dtype; + let tensor_f = Self::dequantize(tensor, dtype); + B::float_all_dim(tensor_f, dim, out_dtype) + } + + /// Sort the elements of the input `tensor` by value in along a given dimension. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// * `dim` - The axis along which to sort. + /// * `descending` - The sorting order. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor, where the elements are sorted by value. + fn q_sort(tensor: QuantizedTensor, dim: usize, descending: bool) -> QuantizedTensor { + // Default implementation. Backends can sort on the int values since qparams remain the same. + dequant_op_quant!( + float_op | tensor | B::float_sort(tensor, dim, descending), + tensor + ) + } + + /// Sort the elements of the input `tensor` by value in along a given dimension. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// * `dim` - The axis along which to sort. + /// * `descending` - The sorting order. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor and corresponding indices, where + /// the elements are sorted by value and the indices map back to the original input tensor. + fn q_sort_with_indices( + tensor: QuantizedTensor, + dim: usize, + descending: bool, + out_dtype: IntDType, + ) -> (QuantizedTensor, IntTensor) { + let scheme = tensor.scheme(); + let dtype = get_device_settings::(&tensor.device()).float_dtype; + + let tensor_f = Self::dequantize(tensor, dtype); + let (out_f, indices) = B::float_sort_with_indices(tensor_f, dim, descending, out_dtype); + + (Self::quantize_dynamic(out_f, &scheme), indices) + } + + /// Returns the indices that sort the elements of the input `tensor` by value along a given dimension. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// * `dim` - The axis along which to sort. + /// * `descending` - The sorting order. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor the indices map back to the original input tensor. + fn q_argsort( + tensor: QuantizedTensor, + dim: usize, + descending: bool, + out_dtype: IntDType, + ) -> IntTensor { + let dtype = get_device_settings::(&tensor.device()).float_dtype; + let tensor_f = Self::dequantize(tensor, dtype); + B::float_argsort(tensor_f, dim, descending, out_dtype) + } +} diff --git a/crates/burn-backend/src/backend/ops/repeat_dim.rs b/crates/burn-backend/src/backend/ops/repeat_dim.rs new file mode 100644 index 0000000..5747cfe --- /dev/null +++ b/crates/burn-backend/src/backend/ops/repeat_dim.rs @@ -0,0 +1,44 @@ +use crate::{Backend, TensorMetadata, tensor::Device}; +use alloc::vec::Vec; +use burn_std::{DType, Shape, Slice}; + +pub(crate) fn repeat_with_slice_assign( + tensor: T, + dim: usize, + times: usize, + device: Device, + empty: E, + slice_assign: SA, +) -> T +where + T: TensorMetadata, + B: Backend, + E: Fn(Shape, &Device, DType) -> T, + SA: Fn(T, &[Slice], T) -> T, +{ + let shape = tensor.shape(); + let dtype = tensor.dtype(); + + let original_dim_length = shape[dim]; + let shape = shape.repeat(dim, times).unwrap(); + + let mut tensor_output = empty(shape.clone(), &device, dtype); + + let indices_select_all = shape.iter().map(|d| 0..*d).collect::>(); + + let mut output_index = 0; + for _ in 0..times { + let mut indices = indices_select_all.clone(); + indices[dim] = output_index..output_index + original_dim_length; + output_index += original_dim_length; + + // Convert ranges to Slice + let slices: Vec = indices + .iter() + .map(|r| Slice::new(r.start as isize, Some(r.end as isize), 1)) + .collect(); + tensor_output = slice_assign(tensor_output, &slices, tensor.clone()); + } + + tensor_output +} diff --git a/crates/burn-backend/src/backend/ops/sort.rs b/crates/burn-backend/src/backend/ops/sort.rs new file mode 100644 index 0000000..ad67b66 --- /dev/null +++ b/crates/burn-backend/src/backend/ops/sort.rs @@ -0,0 +1,433 @@ +use core::cmp::Ordering; + +use crate::{ + Backend, DType, TensorData, + element::{ElementConversion, ElementOrdered}, + tensor::{Device, IntTensor}, +}; +use alloc::{vec, vec::Vec}; +use burn_std::{Element, IntDType}; +use burn_std::{bf16, f16}; + +/// Macro used to dispatch sort operations based on dtype. +macro_rules! sort_dispatch_dtype { + // Dispatch both element dtype and index dtype. + ($fn:ident, |$index_dtype:ident|, $data:ident, $($args:expr),*) => {{ + macro_rules! dispatch_index { + ($index_ty:ty) => { + match $data.dtype { + DType::F64 => $fn::($data, $($args),*), + DType::F32 | DType::Flex32 => { + $fn::($data, $($args),*) + } + DType::F16 => $fn::($data, $($args),*), + DType::BF16 => $fn::($data, $($args),*), + DType::I64 => $fn::($data, $($args),*), + DType::I32 => $fn::($data, $($args),*), + DType::I16 => $fn::($data, $($args),*), + DType::I8 => $fn::($data, $($args),*), + DType::U64 => $fn::($data, $($args),*), + DType::U32 => $fn::($data, $($args),*), + DType::U16 => $fn::($data, $($args),*), + DType::U8 => $fn::($data, $($args),*), + DType::Bool(_) | DType::QFloat(_) => { + unimplemented!("not supported for sorting operations") + } + } + }; + } + + match $index_dtype { + IntDType::I64 => dispatch_index!(i64), + IntDType::I32 => dispatch_index!(i32), + IntDType::I16 => dispatch_index!(i16), + IntDType::I8 => dispatch_index!(i8), + IntDType::U64 => dispatch_index!(u64), + IntDType::U32 => dispatch_index!(u32), + IntDType::U16 => dispatch_index!(u16), + IntDType::U8 => dispatch_index!(u8), + } + }}; + + // Dispatch only element dtype. + ($fn:ident, $data:ident, $($args:expr),*) => { + match $data.dtype { + DType::F64 => $fn::($data, $($args),*), + DType::F32 | DType::Flex32 => $fn::($data, $($args),*), + DType::F16 => $fn::($data, $($args),*), + DType::BF16 => $fn::($data, $($args),*), + DType::I64 => $fn::($data, $($args),*), + DType::I32 => $fn::($data, $($args),*), + DType::I16 => $fn::($data, $($args),*), + DType::I8 => $fn::($data, $($args),*), + DType::U64 => $fn::($data, $($args),*), + DType::U32 => $fn::($data, $($args),*), + DType::U16 => $fn::($data, $($args),*), + DType::U8 => $fn::($data, $($args),*), + DType::Bool(_) | DType::QFloat(_) => { + unimplemented!("not supported for sorting operations") + } + } + }; +} + +/// Sort the elements of the input `tensor` by value along a given dimension. +/// +/// This sort is unstable (i.e., may reorder equal elements). +/// +/// # Arguments +/// +/// * `tensor` - The input tensor. +/// * `dim` - The axis along which to sort. +/// * `descending` - The sorting order. +/// +/// # Returns +/// +/// A tensor with the same shape as the input tensor, where the elements are sorted by value. +/// +/// # Remarks +/// +/// This is a fallback solution that used only when the backend doesn't have the corresponding implementation. +/// Ideally, it is supposed to be implemented by the backend and the backend implementation will be resolved +/// by static dispatch. It is not designed for direct usage by users, and not recommended to import +/// or use this function directly. +pub fn sort( + tensor: T, + dim: usize, + descending: bool, + device: Device, + into_data: ID, + from_data: FD, +) -> T +where + B: Backend, + ID: Fn(T) -> TensorData, + FD: Fn(TensorData, &Device, DType) -> T, +{ + let data = into_data(tensor); + let dtype = data.dtype; + let data = sort_dispatch_dtype!(sort_data, data, dim, descending); + from_data(data, &device, dtype) +} + +pub fn sort_data( + mut data: TensorData, + dim: usize, + descending: bool, +) -> TensorData { + let dims = data.shape.clone(); + let data_slice = data.as_mut_slice().unwrap(); + if dims.len() == 1 { + // 1D sort + data_slice.sort_unstable_by(|&a, &b| compare(&a, &b, descending)); + } else { + sort_slice::(data_slice, &dims, dim, None, false, descending); + } + + data +} + +/// Sort the elements of the input `tensor` by value along a given dimension. +/// +/// This sort is unstable (i.e., may reorder equal elements). +/// +/// # Arguments +/// +/// * `tensor` - The input tensor. +/// * `dim` - The axis along which to sort. +/// * `descending` - The sorting order. +/// * `indices_dtype` - The indices tensor dtype. +/// +/// # Returns +/// +/// A tensor with the same shape as the input tensor and corresponding indices, where +/// the elements are sorted by value and the indices map back to the original input tensor. +/// +/// # Remarks +/// +/// This is a fallback solution that used only when the backend doesn't have the corresponding implementation. +/// Ideally, it is supposed to be implemented by the backend and the backend implementation will be resolved +/// by static dispatch. It is not designed for direct usage by users, and not recommended to import +/// or use this function directly. +pub fn sort_with_indices( + tensor: T, + dim: usize, + descending: bool, + indices_dtype: IntDType, + device: Device, + into_data: ID, + from_data: FD, +) -> (T, IntTensor) +where + B: Backend, + ID: Fn(T) -> TensorData, + FD: Fn(TensorData, &Device, DType) -> T, +{ + let data = into_data(tensor); + let dtype = data.dtype; + let (values, indices) = + sort_dispatch_dtype!(sort_data_with_indices, |indices_dtype|, data, dim, descending); + + ( + from_data(values, &device, dtype), + B::int_from_data(indices.convert_dtype(indices_dtype.into()), &device), + ) +} + +fn sort_data_with_indices( + mut data: TensorData, + dim: usize, + descending: bool, +) -> (TensorData, TensorData) { + let dims = data.shape.clone(); + let mut indices_data = dim_indices::(&dims, dim); + let data_slice = data.as_mut_slice().unwrap(); + if dims.len() == 1 { + // 1D sort + indices_data.sort_unstable_by(|&a, &b| { + compare( + &data_slice[a.elem::() as usize], + &data_slice[b.elem::() as usize], + descending, + ) + }); + + // Permute data in-place by the sorted indices + let mut indices = indices_data + .clone() + .iter() + .map(|i| i.elem::() as usize) + .collect::>(); + for idx in 0..indices.len() { + if indices[idx] != idx { + let mut current_idx = idx; + loop { + let target_idx = indices[current_idx]; + indices[current_idx] = current_idx; + if indices[target_idx] == target_idx { + // correct position + break; + } + + // Permute data by indices + data_slice.swap(current_idx, target_idx); + current_idx = target_idx; + } + } + } + } else { + sort_slice::( + data_slice, + &dims, + dim, + Some(&mut indices_data), + true, + descending, + ); + } + + (data, TensorData::new(indices_data, dims)) +} + +/// Returns the indices that sort the elements of the input `tensor` along a given dimension. +/// +/// This sort is unstable (i.e., may reorder equal elements). +/// +/// # Arguments +/// +/// * `tensor` - The input tensor. +/// * `dim` - The axis along which to sort. +/// * `descending` - The sorting order. +/// * `out_dtype` - The output tensor dtype. +/// +/// # Returns +/// +/// A tensor with the same shape as the input tensor the indices map back to the original input tensor. +/// +/// # Remarks +/// +/// This is a fallback solution that used only when the backend doesn't have the corresponding implementation. +/// Ideally, it is supposed to be implemented by the backend and the backend implementation will be resolved +/// by static dispatch. It is not designed for direct usage by users, and not recommended to import +/// or use this function directly. +pub fn argsort( + tensor: T, + dim: usize, + descending: bool, + out_dtype: IntDType, + device: Device, + into_data: ID, +) -> IntTensor +where + B: Backend, + ID: Fn(T) -> TensorData, +{ + let data = into_data(tensor); + let data = sort_dispatch_dtype!(argsort_data, |out_dtype|, data, dim, descending); + + B::int_from_data(data, &device) +} + +fn argsort_data( + mut data: TensorData, + dim: usize, + descending: bool, +) -> TensorData { + let dims = data.shape.clone(); + let mut indices_data = dim_indices::(&dims, dim); + if dims.len() == 1 { + // 1D sort + let slice = data.as_slice::().unwrap(); + indices_data.sort_unstable_by(|&a, &b| { + compare( + &slice[a.elem::() as usize], + &slice[b.elem::() as usize], + descending, + ) + }); + } else { + sort_slice::( + data.as_mut_slice().unwrap(), + &dims, + dim, + Some(&mut indices_data), + false, + descending, + ); + } + + TensorData::new(indices_data, dims) +} + +/// Sort the elements by value along a given dimension. +/// +/// When `indices` are not provided, the `data` is sorted. +/// Otherwise, the `indices` are sorted based on the value of the elements in `data`, +/// and if `permute_both` is enabled then the data is also sorted. +/// +/// This sort is unstable (i.e., may reorder equal elements). +fn sort_slice( + data: &mut [E], + dims: &[usize], + dim: usize, + mut indices: Option<&mut [I]>, + permute_both: bool, + descending: bool, +) { + let ndims = dims.len(); + let strides = compute_strides(dims); + // Dimensions to access elements to sort + let mut sort_dims = dims.to_vec(); + sort_dims[dim] = 1; + let strides_out = compute_strides(&sort_dims); + + // Number of groups to sort + let num_sorts: usize = dims + .iter() + .enumerate() + .filter(|&(i, _)| i != dim) + .map(|(_, d)| d) + .product(); + + // TODO: run each sort in parallel + // run_par!(|| { + // iter_range_par!(0, num_sorts).for_each(|id| {...}) + for id in 0..num_sorts { + let mut index_offset = 0; + let mut stride_dim = 0; + let mut shape_dim = 0; + for d in 0..ndims { + let stride_input = strides[d]; + let stride_output = strides_out[d]; + let shape_output = sort_dims[d]; + + let num_block = id / stride_output % shape_output; + + if d != dim { + index_offset += num_block * stride_input; + } else { + let shape_input = dims[d]; + stride_dim = stride_input; + shape_dim = shape_input; + index_offset += num_block; + } + } + + // For each group, sort the indices based on the element values + // NOTE: Sorting methods like `sort_unstable_by` are in-place but we need to sort + // different views/groups of the underlying data, so the swap is performed on the elements + // of the (flat index, element value) collection. + let mut elements = (0..shape_dim) + .map(|d| { + let flat_index = d * stride_dim + index_offset; + let elem = data[flat_index]; + (d, flat_index, elem) + }) + .collect::>(); + + elements.sort_unstable_by(|&(_, _, a), &(_, _, b)| compare(&a, &b, descending)); + + // Permute data in-place by the sorted indices + for idx in 0..elements.len() { + if elements[idx].0 != idx { + let mut current_idx = idx; + loop { + let target_idx = elements[current_idx].0; + elements[current_idx].0 = current_idx; + if elements[target_idx].0 == target_idx { + // correct position + break; + } + + if indices.is_none() || permute_both { + // Permute data by indices + data.swap(elements[current_idx].1, elements[target_idx].1); + } + + if let Some(ref mut indices_data) = indices { + // Permute data element indices + indices_data.swap(elements[current_idx].1, elements[target_idx].1); + } + + current_idx = target_idx; + } + } + } + } +} + +/// Computes the steps for each dimension when traversing an array. +fn compute_strides(dims: &[usize]) -> Vec { + let mut strides = vec![0; dims.len()]; + let mut current = 1; + + dims.iter().enumerate().rev().for_each(|(index, val)| { + strides[index] = current; + current *= val; + }); + + strides +} + +/// Generates the indices for each element along the specified dimension. +fn dim_indices(dims: &[usize], dim: usize) -> Vec { + if dims.len() == 1 { + (0..dims[dim]) + .map(|i| (i as i64).elem::()) + .collect::>() + } else { + // Dimension indices tensor + let numel_leading_dims: usize = dims[..dim].iter().product(); + let numel_trailing_dims: usize = dims[dim + 1..].iter().product(); + (0..dims[dim]) + .map(|i| [(i as i64).elem::()].repeat(numel_trailing_dims)) + .collect::>() + .concat() + .repeat(numel_leading_dims) + } +} + +/// Compare two elements +fn compare(a: &E, b: &E, descending: bool) -> Ordering { + if descending { b.cmp(a) } else { a.cmp(b) } +} diff --git a/crates/burn-backend/src/backend/ops/tensor.rs b/crates/burn-backend/src/backend/ops/tensor.rs new file mode 100644 index 0000000..9286b94 --- /dev/null +++ b/crates/burn-backend/src/backend/ops/tensor.rs @@ -0,0 +1,1875 @@ +use super::cat::cat_with_slice_assign; +use super::grid_sample::float_grid_sample_2d_ref; +use super::repeat_dim::repeat_with_slice_assign; +use super::sort::{argsort, sort, sort_with_indices}; +use crate::ops::GridSampleOptions; +use crate::tensor::{BoolTensor, Device, FloatTensor, IntTensor}; +use crate::{Backend, Distribution, TensorData, get_device_settings}; +use crate::{ExecutionError, Scalar, TensorMetadata}; +use alloc::vec::Vec; +use burn_std::reader::try_read_sync; +use burn_std::{BoolDType, FloatDType, IntDType, Shape, Slice}; + +/// Operations on float tensors. +pub trait FloatTensorOps { + /// Creates a new tensor from the data structure. + /// + /// # Arguments + /// + /// * `data` - The data structure. + /// * `device` - The device to create the tensor on. + /// + /// # Returns + /// + /// The tensor with the given data. + fn float_from_data(data: TensorData, device: &Device) -> FloatTensor; + + /// Creates a new tensor with random values. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `distribution` - The distribution to sample from. + /// * `device` - The device to create the tensor on. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The tensor with the given shape and random values. + fn float_random( + shape: Shape, + distribution: Distribution, + device: &Device, + dtype: FloatDType, + ) -> FloatTensor; + + /// Creates a new tensor with zeros. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `device` - The device to create the tensor on. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The tensor with the given shape and zeros. + fn float_zeros(shape: Shape, device: &Device, dtype: FloatDType) -> FloatTensor { + Self::float_from_data(TensorData::full_dtype(shape, 0., dtype.into()), device) + } + + /// Creates a new tensor with ones. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `device` - The device to create the tensor on. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The tensor with the given shape and ones. + fn float_ones(shape: Shape, device: &Device, dtype: FloatDType) -> FloatTensor { + Self::float_from_data(TensorData::full_dtype(shape, 1., dtype.into()), device) + } + + /// Creates a tensor filled with given value. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `fill_value` - The value with which to fill the tensor. + /// * `device` - The device to create the tensor on. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The tensor filled with given value + fn float_full( + shape: Shape, + fill_value: Scalar, + device: &Device, + dtype: FloatDType, + ) -> FloatTensor { + Self::float_from_data( + TensorData::full_dtype(shape, fill_value, dtype.into()), + device, + ) + } + + /// Converts the tensor to a data structure. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// + /// # Returns + /// + /// The data structure with the tensor's data. + fn float_into_data( + tensor: FloatTensor, + ) -> impl Future> + Send; + + /// Moves the tensor to the given device. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `device` - The device to move the tensor to. + /// + /// # Returns + /// + /// The tensor on the given device. + fn float_to_device(tensor: FloatTensor, device: &Device) -> FloatTensor; + + /// Converts float tensor to int tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// The int tensor with the same data as the float tensor. + fn float_into_int(tensor: FloatTensor, out_dtype: IntDType) -> IntTensor; + + /// Creates an empty tensor with the given shape. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `device` - The device to create the tensor on. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The empty tensor with the given shape. + fn float_empty(shape: Shape, device: &Device, dtype: FloatDType) -> FloatTensor; + + /// Repeat the tensor along the given dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `dim` - The dimension to repeat. + /// * `times` - The number of times to repeat the dimension. + /// + /// # Returns + /// + /// The tensor with the given dimension repeated. + fn float_repeat_dim(tensor: FloatTensor, dim: usize, times: usize) -> FloatTensor { + let device = tensor.device(); + repeat_with_slice_assign::( + tensor, + dim, + times, + device, + |shape, device, dtype| B::float_empty(shape, device, dtype.into()), + B::float_slice_assign, + ) + } + + /// Adds two tensors together. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// + /// # Returns + /// + /// The result of adding the two tensors together. + fn float_add(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor; + + /// Adds a scalar to a tensor. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The result of adding the scalar to the tensor. + fn float_add_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor; + + /// Clamps a tensor under a minimum value. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to clamp. + /// * `min` - The minimum value. + /// + /// # Returns + /// + /// The clamped tensor. + fn float_clamp_min(tensor: FloatTensor, min: Scalar) -> FloatTensor { + let dtype = get_device_settings::(&tensor.device()).bool_dtype; + let mask = Self::float_lower_elem(tensor.clone(), min, dtype); + B::float_mask_fill(tensor, mask, min) + } + + /// Clamps a tensor over a maximum value. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to clamp. + /// * `max` - The maximum value. + /// + /// # Returns + /// + /// The clamped tensor. + fn float_clamp_max(tensor: FloatTensor, max: Scalar) -> FloatTensor { + let dtype = get_device_settings::(&tensor.device()).bool_dtype; + let mask = Self::float_greater_elem(tensor.clone(), max, dtype); + B::float_mask_fill(tensor, mask, max) + } + + /// Clamps a tensor between a minimum and maximum value. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to clamp. + /// * `min` - The minimum value. + /// * `max` - The maximum value. + /// + /// # Returns + /// + /// The clamped tensor. + fn float_clamp(tensor: FloatTensor, min: Scalar, max: Scalar) -> FloatTensor { + // Default implementation + Self::float_clamp_min(Self::float_clamp_max(tensor, max), min) + } + + /// Subtracts two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// + /// # Returns + /// + /// The result of subtracting the two tensors. + fn float_sub(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor; + + /// Subtracts a scalar from a tensor. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The result of subtracting the scalar from the tensor. + fn float_sub_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor; + + /// Multiplies two tensors together element-wise. + fn float_mul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor; + + /// Multiplies a tensor by a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The result of multiplying the tensor by the scalar. + fn float_mul_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor; + + /// Divides two tensors element-wise. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// + /// # Returns + /// + /// The result of dividing the two tensors. + fn float_div(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor; + + /// Divides a tensor by a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The result of dividing the tensor by the scalar. + fn float_div_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor; + + /// Computes the remainder of division between two tensors element-wise. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// + /// # Returns + /// + /// The element-wise remainder when dividing `lhs` by `rhs`. + fn float_remainder(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor; + + /// Computes the modulus of a tensor given a scalar. + /// + /// # Arguments + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The result of applying the modulus of the scalar to the tensor. + fn float_remainder_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor; + + /// Multiplies two tensors together using matrix multiplication. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// + /// # Returns + /// + /// The result of multiplying the two tensors together using matrix multiplication. + fn float_matmul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor; + + /// Computes the cross product of two tensors along a given dimension. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// * `dim` - The dimension to compute the cross product along. + /// + /// # Returns + /// + /// The cross product of the two tensors. + fn float_cross(lhs: FloatTensor, rhs: FloatTensor, dim: usize) -> FloatTensor; + + /// Negates a tensor element-wise. + fn float_neg(tensor: FloatTensor) -> FloatTensor { + Self::float_mul_scalar(tensor, (-1f32).into()) + } + + /// Calculates the reciprocals element-wise + fn float_recip(tensor: FloatTensor) -> FloatTensor; + + /// Transposes a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to transpose. + /// + /// # Returns + /// + /// The transposed tensor. + fn float_transpose(tensor: FloatTensor) -> FloatTensor { + let ndims = tensor.shape().num_dims(); + Self::float_swap_dims(tensor, ndims - 2, ndims - 1) + } + + /// Swaps two dimensions of a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to swap the dimensions of. + /// * `dim1` - The first dimension to swap. + /// * `dim2` - The second dimension to swap. + /// + /// # Returns + /// + /// The tensor with the dimensions swapped. + fn float_swap_dims(tensor: FloatTensor, dim1: usize, dim2: usize) -> FloatTensor; + + /// Permutes the dimensions of a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to permute the dimensions of. + /// * `axes` - The new order of the dimensions. + /// # Returns + /// + /// The tensor with the dimensions permuted. + fn float_permute(tensor: FloatTensor, axes: &[usize]) -> FloatTensor; + + /// Reverse the order of elements in a tensor along the given axes. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to reverse. + /// * `axes` - The axes to reverse. + /// + /// The tensor with the elements reversed. + fn float_flip(tensor: FloatTensor, axes: &[usize]) -> FloatTensor; + + /// Reshapes a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to reshape. + /// * `shape` - The new shape of the tensor. + /// + /// # Returns + /// + /// The tensor with the new shape. + fn float_reshape(tensor: FloatTensor, shape: Shape) -> FloatTensor; + + /// Gather elements from a tensor. + /// + /// # Arguments + /// + /// * `dim` - The dimension to gather from. + /// * `tensor` - The tensor to gather from. + /// * `indices` - The indices to gather. + /// + /// # Returns + /// + /// The gathered elements. + fn float_gather(dim: usize, tensor: FloatTensor, indices: IntTensor) -> FloatTensor; + + /// Scatter elements into a tensor using sum reduction. + /// + /// # Arguments + /// + /// * `dim` - The dimension to scatter into. + /// * `tensor` - The tensor to scatter into. + /// * `indices` - The indices to scatter into. + /// * `value` - The value to scatter. + /// + /// # Returns + /// + /// The tensor with the scattered elements. + fn float_scatter_add( + dim: usize, + tensor: FloatTensor, + indices: IntTensor, + value: FloatTensor, + ) -> FloatTensor; + + /// Multi-dimensional scatter: update `data` at locations specified by `indices` with `values`. + /// + /// # Arguments + /// + /// * `data` - The tensor to scatter into. + /// * `indices` - An M-dimensional integer tensor whose last dimension indexes into `data`. + /// * `values` - The values to scatter. + /// * `reduction` - How to combine with existing values. + /// + /// # Returns + /// + /// The tensor with scattered values. + fn float_scatter_nd( + _data: FloatTensor, + _indices: IntTensor, + _values: FloatTensor, + _reduction: crate::tensor::IndexingUpdateOp, + ) -> FloatTensor { + unimplemented!("float_scatter_nd is not implemented for this backend") + } + + /// Multi-dimensional gather: collect slices from `data` at locations specified by `indices`. + /// + /// # Arguments + /// + /// * `data` - The tensor to gather from. + /// * `indices` - An M-dimensional integer tensor whose last dimension indexes into `data`. + /// + /// # Returns + /// + /// The gathered tensor. + fn float_gather_nd(_data: FloatTensor, _indices: IntTensor) -> FloatTensor { + unimplemented!("float_gather_nd is not implemented for this backend") + } + + /// Select tensor elements along the given dimension corresponding for the given indices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to select from. + /// * `dim` - The dimension to select from. + /// * `indices` - The indices to select. + /// + /// # Returns + /// + /// The selected elements. + fn float_select(tensor: FloatTensor, dim: usize, indices: IntTensor) -> FloatTensor; + + /// Assign the selected elements along the given dimension corresponding for the given indices + /// to the given value using sum reduction. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to select from. + /// * `dim` - The dimension to select from. + /// * `indices` - The indices to select. + /// * `value` - The value to assign. + /// + /// # Returns + /// + /// The tensor with the selected elements assigned to the given value. + fn float_select_add( + tensor: FloatTensor, + dim: usize, + indices: IntTensor, + value: FloatTensor, + ) -> FloatTensor; + + /// Select tensor elements corresponding to the given slices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to select from. + /// * `slices` - The slices specifying ranges and steps for each dimension. + /// + /// # Returns + /// + /// The selected elements in a new tensor. + /// + /// # Note + /// + /// Empty slices (where start >= end) are handled at the high-level tensor API and will not + /// be passed to this method. Backend implementations do not need to handle empty slices. + fn float_slice(tensor: FloatTensor, slices: &[Slice]) -> FloatTensor; + + /// Assign the selected elements corresponding to the given slices to the given value. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to select from. + /// * `ranges` - The ranges to select. + /// * `value` - The value to assign. + /// + /// # Returns + /// + /// The tensor with the selected elements assigned to the given value. + /// + /// # Note + /// + /// Empty slice assignments (where any slice range produces 0 elements) are handled at the + /// high-level tensor API and will not be passed to this method. Backend implementations do + /// not need to handle empty slice assignments. + fn float_slice_assign( + tensor: FloatTensor, + slices: &[Slice], + value: FloatTensor, + ) -> FloatTensor; + + /// Update the given tensor with the value tensor where the mask is true. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to select from. + /// * `mask` - The boolean mask to select with. + /// * `value` - The value to assign to the selected elements from the value tensor. + /// + /// # Returns + /// + /// The tensor with the selected elements assigned to the given value. + fn float_mask_where( + tensor: FloatTensor, + mask: BoolTensor, + value: FloatTensor, + ) -> FloatTensor; + + /// Update the given tensor with the value where the mask is true. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to select from. + /// * `mask` - The boolean mask to select with. + /// * `value` - The value to assign to the selected elements. + /// + /// # Returns + /// + /// The tensor with the selected elements assigned to the given value. + fn float_mask_fill( + tensor: FloatTensor, + mask: BoolTensor, + value: Scalar, + ) -> FloatTensor; + + /// Equal comparison of two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor with the result of the comparison. + fn float_equal(lhs: FloatTensor, rhs: FloatTensor, out_dtype: BoolDType) + -> BoolTensor; + + /// Element-wise non-equality comparison. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor with the result of the comparison. + fn float_not_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + let equal_tensor = B::float_equal(lhs, rhs, out_dtype); + B::bool_not(equal_tensor) + } + + /// Equal comparison of a tensor and a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor with the result of the comparison. + fn float_equal_elem(lhs: FloatTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor; + + /// Element-wise non-equality comparison with a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor with the result of the comparison. + fn float_not_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let equal_tensor = B::float_equal_elem(lhs, rhs, out_dtype); + B::bool_not(equal_tensor) + } + + /// Greater than comparison of two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor with the result of the comparison. + fn float_greater( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor; + + /// Greater than comparison of a tensor and a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor with the result of the comparison. + fn float_greater_elem(lhs: FloatTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor; + + /// Greater than or equal comparison of two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor with the result of the comparison. + fn float_greater_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor; + + /// Greater than or equal comparison of a tensor and a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor with the result of the comparison. + fn float_greater_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor; + + /// Less than comparison of two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor with the result of the comparison. + fn float_lower(lhs: FloatTensor, rhs: FloatTensor, out_dtype: BoolDType) + -> BoolTensor; + + /// Less than comparison of a tensor and a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor with the result of the comparison. + fn float_lower_elem(lhs: FloatTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor; + + /// Less than or equal comparison of two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor with the result of the comparison. + fn float_lower_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor; + + /// Less than or equal comparison of a tensor and a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor with the result of the comparison. + fn float_lower_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor; + + /// Detaches a tensor from the computation graph. + fn float_detach(tensor: FloatTensor) -> FloatTensor { + // Should only be overridden by autodiff backends. + tensor + } + + /// Sets the `require_grad` flag of a tensor. + fn float_set_require_grad(tensor: FloatTensor, _require_grad: bool) -> FloatTensor { + // Should only be overridden by autodiff backends. + tensor + } + + /// Returns the `require_grad` flag of a tensor. + fn float_is_require_grad(_tensor: &FloatTensor) -> bool { + // Should only be overridden by autodiff backends. + false + } + + /// Sum of all elements in a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to sum. + /// + /// # Returns + /// + /// A scalar tensor with the sum of all elements in `tensor`. + fn float_sum(tensor: FloatTensor) -> FloatTensor; + + /// Sum of all elements in a tensor along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to sum. + /// * `dim` - The dimension along which to sum. + /// + /// # Returns + /// + /// A tensor with the sum of all elements in `tensor` along `dim`. + fn float_sum_dim(tensor: FloatTensor, dim: usize) -> FloatTensor; + + /// Product of all elements in a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to product. + /// + /// # Returns + /// + /// A scalar tensor with the product of all elements in `tensor`. + fn float_prod(tensor: FloatTensor) -> FloatTensor { + // Product of all elements in a tensor + B::float_exp(B::float_sum(B::float_log(tensor))) + } + + /// Product of all elements in a tensor along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to product. + /// + /// # Returns + /// + /// A tensor with the product of all elements in `tensor` along `dim`. + fn float_prod_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + // Product of all elements in a tensor along a dimension + B::float_exp(B::float_sum_dim(B::float_log(tensor), dim)) + } + + /// Mean of all elements in a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to mean. + /// + /// # Returns + /// + /// A scalar tensor with the mean of all elements in `tensor`. + fn float_mean(tensor: FloatTensor) -> FloatTensor { + let num_elems = tensor.shape().num_elements() as f32; + B::float_div_scalar(B::float_sum(tensor), num_elems.into()) + } + + /// Mean of all elements in a tensor along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to mean. + /// * `dim` - The dimension along which to mean. + /// + /// # Returns + /// + /// A tensor with the mean of all elements in `tensor` along `dim`. + fn float_mean_dim(tensor: FloatTensor, dim: usize) -> FloatTensor; + + /// Computes the cumulative sum of elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the cumulative sum of. + /// * `dim` - The dimension along which to compute the cumulative sum. + /// + /// # Returns + /// + /// A tensor with the same shape where each element is the cumulative sum + /// of all elements up to and including that position along the dimension. + fn float_cumsum(tensor: FloatTensor, dim: usize) -> FloatTensor; + + /// Computes the cumulative product of elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the cumulative product of. + /// * `dim` - The dimension along which to compute the cumulative product. + /// + /// # Returns + /// + /// A tensor with the same shape where each element is the cumulative product + /// of all elements up to and including that position along the dimension. + fn float_cumprod(tensor: FloatTensor, dim: usize) -> FloatTensor; + + /// Computes the cumulative minimum of elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the cumulative minimum of. + /// * `dim` - The dimension along which to compute the cumulative minimum. + /// + /// # Returns + /// + /// A tensor with the same shape where each element is the minimum + /// of all elements up to and including that position along the dimension. + fn float_cummin(tensor: FloatTensor, dim: usize) -> FloatTensor; + + /// Computes the cumulative maximum of elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the cumulative maximum of. + /// * `dim` - The dimension along which to compute the cumulative maximum. + /// + /// # Returns + /// + /// A tensor with the same shape where each element is the maximum + /// of all elements up to and including that position along the dimension. + fn float_cummax(tensor: FloatTensor, dim: usize) -> FloatTensor; + + /// Converts a tensor to another floating point data type. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to convert. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// A tensor with the same values as `tensor` but in the target floating point data type. + fn float_cast(tensor: FloatTensor, dtype: FloatDType) -> FloatTensor; + + /// Returns a new tensor with exponential values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to exponentiate. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with exponential values. + fn float_exp(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with natural logarithm values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the logarithm of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with natural logarithm values. + fn float_log(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with logarithm values of (1 + Xi). + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the logarithm of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with logarithm values of (1 + Xi). + fn float_log1p(tensor: FloatTensor) -> FloatTensor; + + /// Element-wise power with a FloatTensor. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// + /// # Returns + /// + /// The elements of `lhs` raised to the power of the elements of `rhs`. + fn float_powf(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor; + + /// Element-wise power with an IntTensor. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side floatTensor. + /// + /// # Returns + /// + /// The elements of `lhs` raised to the value of `rhs`. Result is an IntTensor. + fn float_powi(lhs: FloatTensor, rhs: IntTensor) -> FloatTensor { + let dtype = lhs.dtype(); + Self::float_powf(lhs, B::int_into_float(rhs, dtype.into())) + } + + /// Raises a tensor to the power of an int scalar. + /// + /// # Backend Implementors Note + /// + /// A number of common exponent cases can be implemented with operations + /// which are much cheaper than generic exponentiation. + /// + /// This (`Backend` impl overridable) operation handles generic optimizations + /// for several common integer exponent cases; and then dispatches to + /// the (`Backend` impl overridable) [`Self::float_powi_scalar_impl`] + /// operation to handle the generic case. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The elements of `lhs` raised to the value of `rhs`. + fn float_powi_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + match rhs.elem::() { + 0 => Self::float_ones(lhs.shape(), &lhs.device(), lhs.dtype().into()), + 1 => lhs, + 2 => B::float_mul(lhs.clone(), lhs), + -1 => Self::float_recip(lhs), + -2 => Self::float_recip(B::float_mul(lhs.clone(), lhs)), + _ => Self::float_powi_scalar_impl(lhs, rhs), + } + } + + /// Raises a tensor to the power of an int scalar. + /// + /// # Backend Implementors Note + /// + /// This is the generic implementation of integer exponentiation + /// called by [`Self::float_powi_scalar`] in the fallback case. + /// + /// As a general rule, this should not be called directly. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side scalar. + /// + /// # Returns + /// + /// The elements of `lhs` raised to the value of `rhs`. + fn float_powi_scalar_impl(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + // Avoid a recursive loop by deferring directly to float_powf_scalar_impl. + Self::float_powf_scalar_impl(lhs, rhs) + } + + /// Returns a new tensor with values raised to the power of float `value`. + /// + /// # Backend Implementors Note + /// + /// This (`Backend` impl overridable) operation dispatches integer exponentiation + /// to [`Self::float_powi_scalar`], and the remaining non-integer exponent cases to + /// the (`Backend` impl overridable) [`Self::float_powf_scalar_impl`] + /// operation to handle the generic case. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to exponentiate. + /// * `value` - The exponent. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with values raised to the power of `value`. + fn float_powf_scalar(tensor: FloatTensor, value: Scalar) -> FloatTensor { + if let Some(exp) = value.try_as_integer() { + Self::float_powi_scalar(tensor, exp) + } else { + Self::float_powf_scalar_impl(tensor, value) + } + } + + /// Returns a new tensor with values raised to the power of float `value`. + /// + /// # Backend Implementors Note + /// + /// This is the generic implementation of integer exponentiation + /// called by [`Self::float_powf_scalar`] in the fallback case. + /// + /// This is the minimal required support a `Backend` must implement + /// for exponentiation. + /// + /// As a general rule, this should not be called directly. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to exponentiate. + /// * `value` - The exponent. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with values raised to the power of `value`. + fn float_powf_scalar_impl(tensor: FloatTensor, value: Scalar) -> FloatTensor; + + /// Returns a new tensor with square root values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the square root of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with square root values. + fn float_sqrt(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with the Euclidean distance values. + /// + /// # Arguments + /// + /// * `lhs` - The left-hand side tensor. + /// * `rhs` - The right-hand side tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `lhs` and `rhs` with hypotenuse values. + fn float_hypot(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + // default implementation for any backend that can't either iterator over elements or doesn't have + // a native hypot implementation + + // Mirrors glibc's approach: scale by max(|lhs|, |rhs|) to avoid + // overflow/underflow in the intermediate squaring step. + // + // hypot(x, y) = |max| * sqrt(1 + (min/max)^2) + // + // Edge cases: + // - If max == 0, both inputs are 0, result is 0 (division guarded by clamp) + // - If max is inf, result is inf (propagates naturally through sqrt) + // - NaN propagates naturally + let abs_lhs = B::float_abs(lhs); + let abs_rhs = B::float_abs(rhs); + + let diff = B::float_clamp_min(B::float_sub(abs_rhs.clone(), abs_lhs.clone()), 0.0.into()); + let max = B::float_add(abs_lhs.clone(), diff.clone()); + let min = B::float_sub(abs_rhs, diff); + + // Clamp max to at least epsilon to avoid 0/0; result will be 0 anyway + // since min <= max, so (min/clamped_max)^2 won't blow up meaningfully. + let max_safe = B::float_clamp_min( + max.clone(), + max.dtype().finfo().unwrap().min_positive.into(), + ); + + let ratio = B::float_div(min, max_safe); + let ratio_sq = B::float_mul(ratio.clone(), ratio); + + let inner = B::float_add_scalar(ratio_sq, 1.0.into()); + + B::float_mul(max, B::float_sqrt(inner)) + } + + /// Returns a new tensor with absolute values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take absolute value of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with absolute values. + fn float_abs(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with cosine values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the cosine of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with cosine values. + fn float_cos(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with sine values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the sine of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with sine values. + fn float_sin(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with tangent values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the tangent of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with tangent values. + fn float_tan(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with hyperbolic cosine values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the hyperbolic cosine of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with hyperbolic cosine values. + fn float_cosh(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with hyperbolic sine values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the hyperbolic sine of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with hyperbolic sine values. + fn float_sinh(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with hyperbolic tangent values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the hyperbolic tangent of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with hyperbolic tangent values. + fn float_tanh(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with inverse cosine values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with inverse cosine values. + fn float_acos(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with inverse hyperbolic cosine values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with inverse hyperbolic cosine values. + fn float_acosh(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with inverse sine values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with inverse sine values. + fn float_asin(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with inverse hyperbolic sine values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with inverse hyperbolic sine values. + fn float_asinh(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with the inverse tangent values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with the inverse tangent values. + fn float_atan(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with the inverse hyperbolic tangent values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with the inverse hyperbolic tangent values. + fn float_atanh(tensor: FloatTensor) -> FloatTensor; + + /// Returns a tensor with the four-quadrant inverse tangent values of `y` and `x`. + /// + /// # Arguments + /// + /// * `lhs` - The tensor with y coordinates. + /// * `rhs` - The tensor with x coordinates. + /// + /// # Returns + /// + /// A tensor with the four-quadrant inverse tangent values. + fn float_atan2(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with rounded values. + /// + /// This function should implement the [round half to even](https://en.wikipedia.org/wiki/Rounding#Rounding_half_to_even) + /// strategy, with halfway cases rounded to the nearest even integer value. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to be rounded. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with rounded values. + fn float_round(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with floored values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to be floored. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with floored values. + fn float_floor(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with ceiled values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to be ceiled. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with ceiled values. + fn float_ceil(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with truncated values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to be truncated. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with truncated values. + fn float_trunc(tensor: FloatTensor) -> FloatTensor; + + /// Returns a new tensor with the error function values. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to take the error function of. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with error function values. + fn float_erf(tensor: FloatTensor) -> FloatTensor; + + /// Concatenates tensors along a dimension. + /// + /// # Arguments + /// + /// * `tensors` - The tensors to concatenate. + /// * `dim` - The dimension along which to concatenate. + /// + /// # Returns + /// + /// A tensor with the concatenated tensors along `dim`. + /// + /// # Note + /// + /// Empty tensors (where the concatenation dimension has size 0) are filtered out at the + /// high-level tensor API and will not be passed to this method. Backend implementations do + /// not need to handle empty tensors. + fn float_cat(tensors: Vec>, dim: usize) -> FloatTensor { + let first_tensor = tensors.first().expect("Tensors should not be empty"); + let device = first_tensor.device(); + + cat_with_slice_assign::( + tensors, + dim, + device, + |shape, device, dtype| B::float_empty(shape, device, dtype.into()), + B::float_slice_assign, + ) + } + + /// Gets the indices of the maximum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements of. + /// * `dim` - The dimension along which to get the maximum elements. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A tensor with the indices of the maximum elements of `tensor` along `dim`. + fn float_argmax(tensor: FloatTensor, dim: usize, out_dtype: IntDType) -> IntTensor; + + /// Gets the indices of the k maximum elements of a tensor along an axis. + /// if two elements are equals, it will be ordered by lowest indices + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements of. + /// * `dim` - The dimension along which to get the maximum elements. + /// * `k` - number of maximum elements + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A tensor with the indices of the maximum elements of `tensor` along `dim`. + fn float_argtopk( + tensor: FloatTensor, + dim: usize, + k: usize, + out_dtype: IntDType, + ) -> IntTensor; + + /// Gets the values of the k maximum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements of. + /// * `dim` - The dimension along which to get the maximum elements. + /// * `k` - number of maximum elements + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A tensor with the values of the maximum elements of `tensor` along `dim`. + fn float_topk(tensor: FloatTensor, dim: usize, k: usize) -> FloatTensor { + let device = tensor.device(); + let dtype = get_device_settings::(&device).int_dtype; + let k_indices = B::int_arange(0..k as i64, &device, dtype); + Self::float_select(Self::float_sort(tensor, dim, true), dim, k_indices) + } + + /// Gets the indices of the minimum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the minimum elements of. + /// * `dim` - The dimension along which to get the minimum elements. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A tensor with the indices of the minimum elements of `tensor` along `dim`. + fn float_argmin(tensor: FloatTensor, dim: usize, out_dtype: IntDType) -> IntTensor; + + /// Gets the maximum element of a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements of. + /// + /// # Returns + /// + /// A tensor with the maximum element of `tensor`. + fn float_max(tensor: FloatTensor) -> FloatTensor { + let shape = tensor.shape(); + let tensor = B::float_reshape(tensor, Shape::new([shape.num_elements()])); + + B::float_max_dim(tensor, 0) + } + + /// Gets the maximum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements of. + /// * `dim` - The dimension along which to get the maximum elements. + /// + /// # Returns + /// + /// A tensor with the maximum elements of `tensor` along `dim`. + fn float_max_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + let dtype = get_device_settings::(&tensor.device()).int_dtype; + let index = B::float_argmax(tensor.clone(), dim, dtype); + + B::float_gather(dim, tensor, index) + } + + /// Gets the maximum elements of a tensor along an axis and their indices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements of. + /// * `dim` - The dimension along which to get the maximum elements. + /// * `indices_dtype` - The indices tensor dtype. + /// + /// # Returns + /// + /// A tuple with the maximum elements of `tensor` along `dim` and their indices. + fn float_max_dim_with_indices( + tensor: FloatTensor, + dim: usize, + indices_dtype: IntDType, + ) -> (FloatTensor, IntTensor) { + let index = B::float_argmax(tensor.clone(), dim, indices_dtype); + let values = B::float_gather(dim, tensor, index.clone()); + + (values, index) + } + + /// Gets the minimum element of a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the minimum elements of. + /// + /// # Returns + /// + /// A tensor with the minimum element of `tensor`. + fn float_min(tensor: FloatTensor) -> FloatTensor { + let shape = tensor.shape(); + let tensor = B::float_reshape(tensor, Shape::new([shape.num_elements()])); + + B::float_min_dim(tensor, 0) + } + + /// Gets the minimum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the minimum elements of. + /// * `dim` - The dimension along which to get the minimum elements. + /// + /// # Returns + /// + /// A tensor with the minimum elements of `tensor` along `dim`. + fn float_min_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + let dtype = get_device_settings::(&tensor.device()).int_dtype; + let index = B::float_argmin(tensor.clone(), dim, dtype); + + B::float_gather(dim, tensor, index) + } + + /// Gets the minimum elements of a tensor along an axis and their indices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the minimum elements of. + /// * `dim` - The dimension along which to get the minimum elements. + /// * `indices_dtype` - The indices tensor dtype. + /// + /// # Returns + /// + /// A tuple with the minimum elements of `tensor` along `dim` and their indices. + fn float_min_dim_with_indices( + tensor: FloatTensor, + dim: usize, + indices_dtype: IntDType, + ) -> (FloatTensor, IntTensor) { + let index = B::float_argmin(tensor.clone(), dim, indices_dtype); + let values = B::float_gather(dim, tensor, index.clone()); + + (values, index) + } + + /// Gets the maximum absolute element of a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements of. + /// + /// # Returns + /// + /// A tensor with the maximum element of `tensor`. + fn float_max_abs(tensor: FloatTensor) -> FloatTensor { + let shape = tensor.shape(); + let tensor = B::float_reshape(tensor, Shape::new([shape.num_elements()])); + + B::float_max_abs_dim(tensor, 0) + } + + /// Gets the maximum absolute elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements of. + /// * `dim` - The dimension along which to get the maximum elements. + /// + /// # Returns + /// + /// A tensor with the maximum elements of `tensor` along `dim`. + fn float_max_abs_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + B::float_max_dim(B::float_abs(tensor), dim) + } + + /// Tests if any element in the float `tensor` evaluates to True. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor with a single element, True if any element in the tensor is True, False otherwise. + fn float_any(tensor: FloatTensor, out_dtype: BoolDType) -> BoolTensor { + let float_dtype = tensor.dtype(); + let bool_tensor = B::float_equal_elem(tensor, 0f32.into(), out_dtype); + let bool_tensor = B::bool_not(bool_tensor); + let sum = B::float_sum(B::bool_into_float(bool_tensor, float_dtype.into())); + B::float_greater_elem(sum, 0f32.into(), out_dtype) + } + + /// Tests if any element in the float `tensor` evaluates to True along a given dimension `dim`. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// * `dim` - The axis along which to test. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor `Tensor` with the same size as input `tensor`, except in the `dim` axis + /// where the size is 1. The elem in the `dim` axis is True if any element along this dim in the + /// input evaluates to True, False otherwise. + fn float_any_dim(tensor: FloatTensor, dim: usize, out_dtype: BoolDType) -> BoolTensor { + let float_dtype = tensor.dtype(); + let bool_tensor = B::float_equal_elem(tensor, 0f32.into(), out_dtype); + let bool_tensor = B::bool_not(bool_tensor); + let sum = B::float_sum_dim(B::bool_into_float(bool_tensor, float_dtype.into()), dim); + B::float_greater_elem(sum, 0f32.into(), out_dtype) + } + + /// Tests if all elements in the float `tensor` evaluate to True. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor `Tensor` with a single element, True if all elements in the input tensor + /// evaluate to True, False otherwise. + fn float_all(tensor: FloatTensor, out_dtype: BoolDType) -> BoolTensor { + let float_dtype = tensor.dtype(); + let num_elems = tensor.shape().num_elements() as f32; + let bool_tensor = B::float_equal_elem(tensor, 0f32.into(), out_dtype); + let bool_tensor = B::bool_not(bool_tensor); + let sum = B::float_sum(B::bool_into_float(bool_tensor, float_dtype.into())); + B::float_equal_elem(sum, num_elems.into(), out_dtype) + } + + /// Tests if all elements in the float `tensor` evaluate to True along a given dimension `dim`. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// * `dim` - The axis along which to test. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A boolean tensor `Tensor` with the same size as input `tensor`, except in the `dim` axis + /// where the size is 1. The elem in the `dim` axis is True if all elements along this dim in the input + /// evaluates to True, False otherwise. + fn float_all_dim(tensor: FloatTensor, dim: usize, out_dtype: BoolDType) -> BoolTensor { + let float_dtype = tensor.dtype(); + let num_elems = tensor.shape()[dim] as f32; + let bool_tensor = B::float_equal_elem(tensor, 0f32.into(), out_dtype); + let bool_tensor = B::bool_not(bool_tensor); + let sum = B::float_sum_dim(B::bool_into_float(bool_tensor, float_dtype.into()), dim); + B::float_equal_elem(sum, num_elems.into(), out_dtype) + } + + /// Returns the signs of the float `tensor`. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to extract the signs from. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` containing the signs of the elements of `tensor`. + fn float_sign(tensor: FloatTensor) -> FloatTensor { + let device = tensor.device(); + let bool_dtype = get_device_settings::(&tensor.device()).bool_dtype; + let zeros = B::float_zeros(tensor.shape(), &device, tensor.dtype().into()); + let less_than_zero = B::float_lower_elem(tensor.clone(), 0f32.into(), bool_dtype); + let greater_than_zero = B::float_greater_elem(tensor, 0f32.into(), bool_dtype); + + let mut result = B::float_mask_fill(zeros, less_than_zero, (-1f32).into()); + result = B::float_mask_fill(result, greater_than_zero, 1f32.into()); + result + } + + /// Broadcasts the float `tensor` to the given `shape`. + fn float_expand(tensor: FloatTensor, shape: Shape) -> FloatTensor; + + /// Sort the elements of the input `tensor` by value in along a given dimension. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// * `dim` - The axis along which to sort. + /// * `descending` - The sorting order. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor, where the elements are sorted by value. + fn float_sort(tensor: FloatTensor, dim: usize, descending: bool) -> FloatTensor { + let device = tensor.device(); + sort::( + tensor, + dim, + descending, + device, + |tensor| { + let msg = "Failed to synchronously read tensor data. This operation is not supported until this backend has a GPU sorting implementation."; + try_read_sync(B::float_into_data(tensor)) + .expect(msg) + .expect(msg) + }, + |data, device, _dtype| B::float_from_data(data, device), + ) + } + + /// Sort the elements of the input `tensor` by value in along a given dimension. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// * `dim` - The axis along which to sort. + /// * `descending` - The sorting order. + /// * `indices_dtype` - The indices tensor dtype. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor and corresponding indices, where + /// the elements are sorted by value and the indices map back to the original input tensor. + fn float_sort_with_indices( + tensor: FloatTensor, + dim: usize, + descending: bool, + indices_dtype: IntDType, + ) -> (FloatTensor, IntTensor) { + let device = tensor.device(); + sort_with_indices::( + tensor, + dim, + descending, + indices_dtype, + device, + |tensor| { + let msg = "Failed to synchronously read tensor data. This operation is not supported until this backend has a GPU sorting implementation."; + try_read_sync(B::float_into_data(tensor)) + .expect(msg) + .expect(msg) + }, + |data, device, _dtype| B::float_from_data(data, device), + ) + } + + /// Returns the indices that sort the elements of the input `tensor` by value along a given dimension. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// * `dim` - The axis along which to sort. + /// * `descending` - The sorting order. + /// * `out_dtype` - The output tensor dtype. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor the indices map back to the original input tensor. + fn float_argsort( + tensor: FloatTensor, + dim: usize, + descending: bool, + out_dtype: IntDType, + ) -> IntTensor { + let device = tensor.device(); + argsort::(tensor, dim, descending, out_dtype, device, |tensor| { + let msg = "Failed to synchronously read tensor data. This operation is not supported until this backend has a GPU sorting implementation."; + try_read_sync(B::float_into_data(tensor)) + .expect(msg) + .expect(msg) + }) + } + + /// Samples tensor as a two-dimensional spatial grid of (possibly multi-channel) values, + /// using the given locations in [-1, 1]. + /// + /// # Arguments + /// + /// * `tensor` - The tensor being sampled from, must be contiguous with shape (N, C, H_in, W_in) + /// * `grid` - A tensor of locations, with shape (N, H_out, W_out, 2). Values are [-1, 1]. + /// A [x = -1, y = -1] means top-left, and [x = 1, y = 1] means bottom-right + /// * `options` - Grid sampling options (mode, padding_mode, align_corners) + /// + /// # Returns + /// + /// A tensor with shape (N, C, H_out, W_out) + fn float_grid_sample_2d( + tensor: FloatTensor, + grid: FloatTensor, + options: GridSampleOptions, + ) -> FloatTensor { + // TODO: default impl should get int default dtype + float_grid_sample_2d_ref::(tensor, grid, options) + } + + /// Unfold windows along a dimension. + /// + /// Returns a view of the tensor with all complete windows of size `size` in dimension `dim`; + /// where windows are advanced by `step` at each index. + /// + /// The number of windows is `max(0, (shape[dim] - size).ceil_div(step))`. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor to unfold; of shape ``[pre=..., dim shape, post=...]`` + /// * `dim` - the selected dim. + /// * `size` - the size of each unfolded window. + /// * `step` - the step between each window. + /// + /// # Returns + /// + /// A tensor view with shape ``[pre=..., windows, size, post=...]``. + fn float_unfold(tensor: FloatTensor, dim: usize, size: usize, step: usize) + -> FloatTensor; + + /// Returns a new tensor with boolean elements indicating whether each element of the input is NaN. + /// + /// # Returns + /// + /// A boolean tensor where `true` indicates NaN and `false` indicates a non-NaN value. + fn float_is_nan(tensor: FloatTensor, out_dtype: BoolDType) -> BoolTensor { + // Check if the input tensor is NaN by comparing it to itself + // NaN is the only value that is not equal to itself + B::float_not_equal(tensor.clone(), tensor, out_dtype) + } + + /// Returns a new tensor with boolean elements indicating whether each element of the input is infinite (either +INF or -INF). + /// + /// # Returns + /// + /// A boolean tensor where `true` indicates that the value is infinite + fn float_is_inf(tensor: FloatTensor, out_dtype: BoolDType) -> BoolTensor { + B::float_equal_elem(B::float_abs(tensor), f64::INFINITY.into(), out_dtype) + } +} diff --git a/crates/burn-backend/src/backend/ops/transaction.rs b/crates/burn-backend/src/backend/ops/transaction.rs new file mode 100644 index 0000000..59440a9 --- /dev/null +++ b/crates/burn-backend/src/backend/ops/transaction.rs @@ -0,0 +1,142 @@ +use alloc::vec::Vec; +use core::future::Future; + +use crate::tensor::{BoolTensor, FloatTensor, IntTensor, QuantizedTensor}; +use crate::{Backend, ExecutionError, TensorData, TensorPrimitive}; + +enum Order { + Float(usize), + QFloat(usize), + Int(usize), + Bool(usize), +} + +#[derive(Default)] +/// Contains all tensor primitives that are going to be read. +pub struct TransactionPrimitive { + /// Float tensors. + pub read_floats: Vec>, + /// Quantized tensors. + pub read_qfloats: Vec>, + /// Int tensors. + pub read_ints: Vec>, + /// Bool tensors. + pub read_bools: Vec>, + orders: Vec, +} + +#[derive(Default)] +/// Contains all [data](TensorData) related to a [transaction](TransactionPrimitive). +pub struct TransactionPrimitiveData { + /// Float tensor data. + pub read_floats: Vec, + /// Quantized tensor data. + pub read_qfloats: Vec, + /// Int tensor data. + pub read_ints: Vec, + /// Bool tensor data. + pub read_bools: Vec, +} + +/// Operations that are sync by nature and that can be batch together in transactions to improve +/// compute utilization with efficient laziness. +pub trait TransactionOps { + /// Executes a [transaction](TransactionPrimitive) and return its + /// [data](TransactionPrimitiveData). + fn tr_execute( + transaction: TransactionPrimitive, + ) -> impl Future> + Send { + async move { + let mut floats = Vec::new(); + let mut qfloats = Vec::new(); + let mut ints = Vec::new(); + let mut bools = Vec::new(); + + for t in transaction.read_floats { + floats.push(B::float_into_data(t).await?); + } + for t in transaction.read_qfloats { + qfloats.push(B::q_into_data(t).await?); + } + for t in transaction.read_ints { + ints.push(B::int_into_data(t).await?); + } + for t in transaction.read_bools { + bools.push(B::bool_into_data(t).await?); + } + + Ok(TransactionPrimitiveData { + read_floats: floats, + read_qfloats: qfloats, + read_ints: ints, + read_bools: bools, + }) + } + } +} + +impl TransactionPrimitive { + /// Creates a new transaction. + pub fn new( + read_floats: Vec>, + read_qfloats: Vec>, + read_ints: Vec>, + read_bools: Vec>, + ) -> Self { + Self { + read_floats, + read_qfloats, + read_ints, + read_bools, + orders: Vec::default(), + } + } + /// Executes the transaction asynchronously and returns the [data](TensorData) in the same order + /// in which they were registered. + pub async fn execute_async(mut self) -> Result, ExecutionError> { + let mut orders = Vec::new(); + core::mem::swap(&mut orders, &mut self.orders); + let result = B::tr_execute(self).await?; + + let mut floats: Vec<_> = result.read_floats.into_iter().map(Some).collect(); + let mut qfloats: Vec<_> = result.read_qfloats.into_iter().map(Some).collect(); + let mut ints: Vec<_> = result.read_ints.into_iter().map(Some).collect(); + let mut bools: Vec<_> = result.read_bools.into_iter().map(Some).collect(); + + Ok(orders + .into_iter() + .map(|order| match order { + Order::Float(index) => floats.get_mut(index).unwrap().take().unwrap(), + Order::QFloat(index) => qfloats.get_mut(index).unwrap().take().unwrap(), + Order::Int(index) => ints.get_mut(index).unwrap().take().unwrap(), + Order::Bool(index) => bools.get_mut(index).unwrap().take().unwrap(), + }) + .collect::>()) + } + + /// Register a float (or quantized-float) tensor to read in this transaction. + pub fn register_float(&mut self, tensor: TensorPrimitive) { + match tensor { + TensorPrimitive::Float(tensor) => { + self.orders.push(Order::Float(self.read_floats.len())); + self.read_floats.push(tensor); + } + TensorPrimitive::QFloat(tensor) => { + self.orders.push(Order::QFloat(self.read_qfloats.len())); + self.read_qfloats.push(tensor); + } + } + } + + /// Register an int tensor to read in this transaction. + pub fn register_int(&mut self, tensor: IntTensor) { + self.orders.push(Order::Int(self.read_ints.len())); + self.read_ints.push(tensor); + } + + /// Register a bool tensor to read in this transaction. + pub fn register_bool(&mut self, tensor: BoolTensor) { + self.orders.push(Order::Bool(self.read_bools.len())); + self.read_bools.push(tensor); + } +} diff --git a/crates/burn-backend/src/backend/primitive.rs b/crates/burn-backend/src/backend/primitive.rs new file mode 100644 index 0000000..33ae236 --- /dev/null +++ b/crates/burn-backend/src/backend/primitive.rs @@ -0,0 +1,107 @@ +use crate::{Backend, BackendTypes, DeviceOps, get_device_settings}; +use burn_std::{DType, QuantScheme, Shape}; + +#[derive(Debug, Clone)] +/// A primitive tensor representation. +pub enum TensorPrimitive { + /// Float tensor primitive. + Float(B::FloatTensorPrimitive), + /// Quantized float tensor primitive. + QFloat(B::QuantizedTensorPrimitive), +} + +impl TensorPrimitive { + /// Returns the full tensor representation. + pub fn tensor(self) -> B::FloatTensorPrimitive { + match self { + Self::QFloat(tensor) => { + let dtype = get_device_settings::(&tensor.device()).float_dtype; + B::dequantize(tensor, dtype) + } + Self::Float(tensor) => tensor, + } + } + + /// Returns a mutable reference to the full tensor representation. + pub fn get_mut_ref(&mut self) -> &mut B::FloatTensorPrimitive { + match self { + Self::QFloat(_tensor) => todo!(), + Self::Float(tensor) => tensor, + } + } +} + +impl TensorMetadata for TensorPrimitive { + type Device = B::Device; + + fn dtype(&self) -> DType { + match self { + TensorPrimitive::Float(tensor) => tensor.dtype(), + TensorPrimitive::QFloat(tensor) => tensor.dtype(), + } + } + + fn shape(&self) -> Shape { + match self { + TensorPrimitive::Float(tensor) => tensor.shape(), + TensorPrimitive::QFloat(tensor) => tensor.shape(), + } + } + + fn rank(&self) -> usize { + match self { + TensorPrimitive::Float(tensor) => tensor.rank(), + TensorPrimitive::QFloat(tensor) => tensor.rank(), + } + } + fn device(&self) -> Self::Device { + match self { + TensorPrimitive::Float(tensor) => tensor.device(), + TensorPrimitive::QFloat(tensor) => tensor.device(), + } + } + + fn can_mut(&self) -> bool { + match self { + TensorPrimitive::Float(tensor) => tensor.can_mut(), + TensorPrimitive::QFloat(tensor) => tensor.can_mut(), + } + } +} + +/// Tensor metadata trait for tensor primitive. +pub trait TensorMetadata: Clone + Send + Sync + core::fmt::Debug { + /// The device type associated with the tensor. + type Device: DeviceOps; + /// Get the dtype of the tensor. + fn dtype(&self) -> DType; + /// Get the shape of the tensor. + fn shape(&self) -> Shape; + + /// Get the number of dimensions of the tensor. + fn rank(&self) -> usize { + self.shape().num_dims() + } + /// Get the device associated with the tensor. + fn device(&self) -> Self::Device; + + /// Whether the tensor's buffer can be mutated in place — i.e. this handle + /// uniquely owns it, so an in-place op (`slice_assign`, an inplace kernel) + /// writes the existing allocation instead of copying it first. + /// + /// Backends that track buffer ownership (cubecl, fusion, tch) answer + /// precisely; a backend that can't must return a conservative `false` — + /// the buffer may be aliased, so an in-place write can't be assumed safe. + fn can_mut(&self) -> bool; + + /// Get the [quantization scheme](QuantScheme) for a quantized float tensor. + /// + /// # Panics + /// Panics if the tensor is not quantized. + fn scheme(&self) -> QuantScheme { + match self.dtype() { + DType::QFloat(scheme) => scheme, + other => panic!("Quantization scheme is not valid for dtype {other:?}"), + } + } +} diff --git a/crates/burn-backend/src/cubecl.rs b/crates/burn-backend/src/cubecl.rs new file mode 100644 index 0000000..d337410 --- /dev/null +++ b/crates/burn-backend/src/cubecl.rs @@ -0,0 +1,111 @@ +//! Conversion helpers between burn's [`DType`] and cubecl's +//! `ElemType` / `StorageType`. +//! +//! These are exposed as plain named functions (rather than `From`/`Into` +//! trait impls in `burn-std`) so that `burn-std` doesn't need to depend on +//! `cubecl` and the cubecl type tree doesn't leak into the shared standard +//! library. Backend implementations that want these conversions depend on +//! `burn-backend` (with the `cubecl` feature) and call these functions +//! explicitly. + +use burn_std::{BoolStore, DType, QuantScheme, QuantStore, QuantValue}; +use cubecl::ir::{ElemType, FloatKind, IntKind, StorageType, UIntKind}; + +/// Convert a cubecl [`ElemType`] into the corresponding burn [`DType`]. +/// +/// Panics if the cubecl type has no direct burn equivalent (e.g. `TF32`). +pub fn elem_type_to_dtype(value: ElemType) -> DType { + match value { + ElemType::Float(float_kind) => match float_kind { + FloatKind::F16 => DType::F16, + FloatKind::BF16 => DType::BF16, + FloatKind::Flex32 => DType::Flex32, + FloatKind::F32 => DType::F32, + FloatKind::F64 => DType::F64, + FloatKind::TF32 => panic!("Not a valid DType for tensors."), + FloatKind::E2M1 + | FloatKind::E2M3 + | FloatKind::E3M2 + | FloatKind::E4M3 + | FloatKind::E5M2 + | FloatKind::UE8M0 => { + unimplemented!("Not yet supported, will be used for quantization") + } + }, + ElemType::Int(int_kind) => match int_kind { + IntKind::I8 => DType::I8, + IntKind::I16 => DType::I16, + IntKind::I32 => DType::I32, + IntKind::I64 => DType::I64, + }, + ElemType::UInt(uint_kind) => match uint_kind { + UIntKind::U8 => DType::U8, + UIntKind::U16 => DType::U16, + UIntKind::U32 => DType::U32, + UIntKind::U64 => DType::U64, + }, + _ => panic!("Not a valid DType for tensors."), + } +} + +/// Convert a burn [`DType`] into the corresponding cubecl [`ElemType`]. +/// +/// Sub-byte quantization variants whose storage is `PackedNative` may panic +/// (see inline `panic!`s) — those configurations are not representable as a +/// single `ElemType`; use [`dtype_to_storage_type`] instead. +pub fn dtype_to_elem_type(dtype: DType) -> ElemType { + match dtype { + DType::F64 => ElemType::Float(FloatKind::F64), + DType::F32 => ElemType::Float(FloatKind::F32), + DType::Flex32 => ElemType::Float(FloatKind::Flex32), + DType::F16 => ElemType::Float(FloatKind::F16), + DType::BF16 => ElemType::Float(FloatKind::BF16), + DType::I64 => ElemType::Int(IntKind::I64), + DType::I32 => ElemType::Int(IntKind::I32), + DType::I16 => ElemType::Int(IntKind::I16), + DType::I8 => ElemType::Int(IntKind::I8), + DType::U64 => ElemType::UInt(UIntKind::U64), + DType::U32 => ElemType::UInt(UIntKind::U32), + DType::U16 => ElemType::UInt(UIntKind::U16), + DType::U8 => ElemType::UInt(UIntKind::U8), + DType::Bool(store) => match store { + BoolStore::Native => ElemType::Bool, + BoolStore::U8 => ElemType::UInt(UIntKind::U8), + BoolStore::U32 => ElemType::UInt(UIntKind::U32), + }, + DType::QFloat(scheme) => match scheme.store { + QuantStore::Native => match scheme.value { + QuantValue::Q8F | QuantValue::Q8S => ElemType::Int(IntKind::I8), + QuantValue::E4M3 => ElemType::Float(FloatKind::E4M3), + QuantValue::E5M2 => ElemType::Float(FloatKind::E5M2), + QuantValue::Q4F + | QuantValue::Q4S + | QuantValue::Q2F + | QuantValue::Q2S + | QuantValue::E2M1 => { + panic!("Can't store native sub-byte values") + } + }, + QuantStore::PackedU32(_) => ElemType::UInt(UIntKind::U32), + QuantStore::PackedNative(_) => match scheme.value { + QuantValue::E2M1 => panic!("Can't store native sub-byte values"), + other => panic!("{other:?} doesn't support native packing"), + }, + }, + } +} + +/// Convert a burn [`DType`] into the corresponding cubecl [`StorageType`]. +/// +/// Handles sub-byte packed quantization configurations that cannot be expressed +/// as a plain [`ElemType`] by emitting a `StorageType::Packed(...)`. +pub fn dtype_to_storage_type(dtype: DType) -> StorageType { + match dtype { + DType::QFloat(QuantScheme { + store: QuantStore::PackedNative(_), + value: QuantValue::E2M1, + .. + }) => StorageType::Packed(ElemType::Float(FloatKind::E2M1), 2), + _ => dtype_to_elem_type(dtype).into(), + } +} diff --git a/crates/burn-backend/src/lib.rs b/crates/burn-backend/src/lib.rs new file mode 100644 index 0000000..c543d37 --- /dev/null +++ b/crates/burn-backend/src/lib.rs @@ -0,0 +1,176 @@ +#![cfg_attr(not(feature = "std"), no_std)] +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] + +//! This library provides the core types that define how Burn tensor data is represented, stored, and interpreted. + +#[macro_use] +extern crate derive_new; + +extern crate alloc; + +/// [`Backend`] trait and required types. +pub mod backend; +pub use backend::*; + +// Re-exported types +pub use burn_std::reader::*; // Useful so that backends don't have to add `burn_std` as a dependency. +pub use burn_std::{ + AllocationProperty, BoolDType, BoolStore, Bytes, DType, DataError, DeviceHandle, Distribution, + DistributionSampler, DistributionSamplerKind, Element, ElementConversion, ElementEq, + ElementOrdered, ElementRandom, FloatDType, IntDType, Scalar, SplitPolicy, TensorData, + Tolerance, bf16, distribution, element, f16, stream_id::StreamId, +}; + +/// Shape definition. +pub mod shape { + pub use burn_std::shape::*; +} +pub use shape::*; + +/// Slice utilities. +pub mod slice { + pub use burn_std::{s, slice::*}; +} +pub use slice::*; + +/// Indexing utilities. +pub mod indexing { + pub use burn_std::indexing::*; +} +pub use indexing::*; + +mod alias; +pub use alias::*; + +/// Quantization data representation. +pub mod quantization; + +/// CubeCL inter-operation helpers (gated by the `cubecl` feature). +/// +/// Provides plain conversion functions between burn's [`DType`] and cubecl's +/// `ElemType` / `StorageType`. They are intentionally exposed as named +/// functions rather than `From`/`Into` impls so the cubecl type tree does not +/// leak into `burn-std`'s public surface. +#[cfg(feature = "cubecl")] +pub mod cubecl; + +#[cfg(any( + feature = "cubecl-wgpu", + feature = "cubecl-metal", + feature = "cubecl-vulkan", + feature = "cubecl-webgpu" +))] +mod cube_wgpu { + use crate::backend::DeviceOps; + use burn_std::{BoolStore, DType, DeviceSettings}; + use cubecl::wgpu::WgpuDevice; + + impl DeviceOps for WgpuDevice { + #[cfg(not(any(feature = "cubecl-metal", feature = "cubecl-vulkan")))] + fn defaults(&self) -> DeviceSettings { + DeviceSettings::new( + DType::F32, + DType::I32, + DType::Bool(BoolStore::U32), + Default::default(), + ) + } + + #[cfg(any(feature = "cubecl-metal", feature = "cubecl-vulkan"))] + fn defaults(&self) -> DeviceSettings { + DeviceSettings::new( + DType::F32, + DType::I32, + DType::Bool(BoolStore::U8), + Default::default(), + ) + } + } +} + +#[cfg(feature = "cubecl-cuda")] +mod cube_cuda { + use crate::backend::DeviceOps; + use burn_std::{BoolStore, DType, DeviceSettings}; + use cubecl::cuda::CudaDevice; + + impl DeviceOps for CudaDevice { + fn defaults(&self) -> DeviceSettings { + DeviceSettings::new( + DType::F32, + DType::I32, + DType::Bool(BoolStore::U8), + Default::default(), + ) + } + } +} + +#[cfg(feature = "cubecl-cpu")] +mod cube_cpu { + use crate::backend::DeviceOps; + use burn_std::{BoolStore, DType, DeviceSettings}; + use cubecl::cpu::CpuDevice; + + impl DeviceOps for CpuDevice { + fn defaults(&self) -> DeviceSettings { + DeviceSettings::new( + DType::F32, + DType::I32, + DType::Bool(BoolStore::U8), + Default::default(), + ) + } + } +} + +#[cfg(feature = "cubecl-hip")] +mod cube_hip { + use crate::backend::DeviceOps; + use burn_std::{BoolStore, DType, DeviceSettings}; + use cubecl::hip::AmdDevice; + + impl DeviceOps for AmdDevice { + fn defaults(&self) -> DeviceSettings { + DeviceSettings::new( + DType::F32, + DType::I32, + DType::Bool(BoolStore::U8), + Default::default(), + ) + } + } +} + +/// Convenience macro to link to the `burn-tensor` docs for this crate version. +/// +/// Usage: +/// ```rust,ignore +/// # use burn_backend::doc_tensor; +/// doc_tensor!(); // Links to `Tensor` struct +/// doc_tensor!("zeros"); // Links to `Tensor::zeros` method +/// ``` +#[macro_export] +macro_rules! doc_tensor { + () => { + concat!( + "[`Tensor`](https://docs.rs/burn-tensor/", + env!("CARGO_PKG_VERSION"), + "/burn_tensor/struct.Tensor.html)" + ) + }; + + ($method:literal) => { + concat!( + "[`Tensor::", + $method, + "`](", + "https://docs.rs/burn-tensor/", + env!("CARGO_PKG_VERSION"), + "/burn_tensor/struct.Tensor.html#method.", + $method, + ")" + ) + }; +} diff --git a/crates/burn-backend/src/quantization/mod.rs b/crates/burn-backend/src/quantization/mod.rs new file mode 100644 index 0000000..401a88c --- /dev/null +++ b/crates/burn-backend/src/quantization/mod.rs @@ -0,0 +1,10 @@ +mod parameters; +mod scheme; + +pub use parameters::*; +pub use scheme::*; + +pub use burn_std::quantization::{ + BlockSize, Calibration, QuantLevel, QuantMode, QuantParam, QuantPropagation, QuantScheme, + QuantStore, QuantValue, QuantizedBytes, +}; diff --git a/crates/burn-backend/src/quantization/parameters.rs b/crates/burn-backend/src/quantization/parameters.rs new file mode 100644 index 0000000..5b50882 --- /dev/null +++ b/crates/burn-backend/src/quantization/parameters.rs @@ -0,0 +1,15 @@ +use crate::Backend; + +pub use burn_std::quantization::{QParamTensor, QParams}; + +/// The quantization parameters primitive. +/// +/// # Remarks +/// +/// This is a low-level struct used internally by the library to provide the quantization parameters +/// to the backends. It is not designed for direct usage by users, and not recommended to import +/// or use this struct directly. +pub struct QuantizationParametersPrimitive { + /// The scaling factor. + pub scales: B::FloatTensorPrimitive, +} diff --git a/crates/burn-backend/src/quantization/scheme.rs b/crates/burn-backend/src/quantization/scheme.rs new file mode 100644 index 0000000..7e22d20 --- /dev/null +++ b/crates/burn-backend/src/quantization/scheme.rs @@ -0,0 +1,98 @@ +pub use burn_std::{QPARAM_ALIGN, params_shape}; +use burn_std::{QuantLevel, QuantMode, QuantScheme, Shape}; + +use super::{Calibration, QuantizationParametersPrimitive}; +use crate::{Backend, TensorMetadata, get_device_settings}; + +/// Compute the quantization range mapping. +pub fn compute_range( + scheme: &QuantScheme, + tensor: B::FloatTensorPrimitive, + calibration: &Calibration, +) -> (B::FloatTensorPrimitive, B::FloatTensorPrimitive) { + match calibration { + Calibration::MinMax => match scheme.level { + QuantLevel::Tensor => (B::float_min(tensor.clone()), B::float_max(tensor)), + QuantLevel::Block(block_size) => { + let block_elems = block_size.num_elements(); + let shape = tensor.shape(); + let numel = shape.num_elements(); + + assert_eq!( + numel % block_elems, + 0, + "Tensor {shape:?} must be evenly divisible by block size {block_elems}" + ); + + let num_blocks = numel / block_elems; + + let params_shape = params_shape(&shape, scheme.level); + + let blocks = B::float_reshape(tensor, Shape::new([num_blocks, block_elems])); + let blocks_min = + B::float_reshape(B::float_min_dim(blocks.clone(), 1), params_shape.clone()); + let blocks_max = B::float_reshape(B::float_max_dim(blocks, 1), params_shape); + (blocks_min, blocks_max) + } + }, + Calibration::AbsMean => { + // gamma = mean(|W|) per tensor or block — symmetric range [-gamma, +gamma] + let gamma = match scheme.level { + QuantLevel::Tensor => B::float_mean(B::float_abs(tensor)), + QuantLevel::Block(block_size) => { + let block_elems = block_size.num_elements(); + let shape = tensor.shape(); + let numel = shape.num_elements(); + + assert_eq!( + numel % block_elems, + 0, + "Tensor {shape:?} must be evenly divisible by block size {block_elems}" + ); + + let num_blocks = numel / block_elems; + let params_shape = params_shape(&shape, scheme.level); + let blocks = B::float_reshape( + B::float_abs(tensor), + Shape::new([num_blocks, block_elems]), + ); + B::float_reshape(B::float_mean_dim(blocks, 1), params_shape) + } + }; + let neg_gamma = B::float_neg(gamma.clone()); + (neg_gamma, gamma) + } + } +} + +/// Compute the quantization parameters. +pub fn compute_q_params( + scheme: &QuantScheme, + min: B::FloatTensorPrimitive, + max: B::FloatTensorPrimitive, +) -> QuantizationParametersPrimitive { + match scheme { + QuantScheme { + level: QuantLevel::Tensor | QuantLevel::Block(_), + mode: QuantMode::Symmetric, + .. + } => { + let bool_dtype = get_device_settings::(&min.device()).bool_dtype; + // Quantized range `[a, b]` + let (a, b) = scheme.value.range(); + + // Compute scale to convert an input value in range `[-alpha, alpha]` + let min_abs = B::float_abs(min); + let max_abs = B::float_abs(max); + + // `min_abs.max_pair(max_abs)` + let mask = B::float_lower(min_abs.clone(), max_abs.clone(), bool_dtype); + let values_range = + B::float_mul_scalar(B::float_mask_where(min_abs, mask, max_abs), 2f32.into()); + + QuantizationParametersPrimitive { + scales: B::float_div_scalar(values_range, (b - a).into()), + } + } + } +} diff --git a/crates/burn-candle/Cargo.toml b/crates/burn-candle/Cargo.toml new file mode 100644 index 0000000..bedfd3c --- /dev/null +++ b/crates/burn-candle/Cargo.toml @@ -0,0 +1,43 @@ +[package] +authors = ["louisfd "] +categories = ["science"] +description = "[Deprecated] Candle backend for the Burn framework - use burn-cubecl, burn-ndarray, or burn-tch instead" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "data"] +license.workspace = true +name = "burn-candle" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-candle" +documentation = "https://docs.rs/burn-candle" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["std"] +std = [] +doc = ["default"] +tracing = [ + "burn-backend/tracing", + "burn-std/tracing", +] + +cuda = ["candle-core/cuda"] +metal = ["candle-core/metal"] +accelerate = ["candle-core/accelerate"] + +[dependencies] +burn-backend = { workspace = true } +# For rand utils and stub mutex +burn-std = { workspace = true } + +candle-core = { workspace = true } +derive-new = { workspace = true } + +[dev-dependencies] +burn-tch = { workspace = true } + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-candle/LICENSE-APACHE b/crates/burn-candle/LICENSE-APACHE new file mode 120000 index 0000000..1cd601d --- /dev/null +++ b/crates/burn-candle/LICENSE-APACHE @@ -0,0 +1 @@ +../../LICENSE-APACHE \ No newline at end of file diff --git a/crates/burn-candle/LICENSE-MIT b/crates/burn-candle/LICENSE-MIT new file mode 120000 index 0000000..b2cfbdc --- /dev/null +++ b/crates/burn-candle/LICENSE-MIT @@ -0,0 +1 @@ +../../LICENSE-MIT \ No newline at end of file diff --git a/crates/burn-candle/README.md b/crates/burn-candle/README.md new file mode 100644 index 0000000..d535f6f --- /dev/null +++ b/crates/burn-candle/README.md @@ -0,0 +1,14 @@ +# Burn Candle Backend + +> **Deprecated:** This crate is deprecated as of `0.21.0-pre.2` and will be removed in a future release. +> Please migrate to one of the actively maintained backends: +> - **CubeCL backends** (CUDA, ROCm, Vulkan, Metal, WebGPU) for GPU acceleration +> - **[Flex](../burn-flex)** (`burn-flex`) for portable pure-Rust CPU execution (std, no_std, WASM) +> - **LibTorch** (`burn-tch`) for a mature CPU/GPU backend + +This crate provides a backend for [Burn](https://github.com/tracel-ai/burn) based on the [Candle](https://github.com/huggingface/candle) framework. + +## Feature Flags + +- `cuda` - Cuda GPU device (NVIDIA only) +- `accelerate` - Accelerate framework (macOS only) diff --git a/crates/burn-candle/src/backend.rs b/crates/burn-candle/src/backend.rs new file mode 100644 index 0000000..ce76a60 --- /dev/null +++ b/crates/burn-candle/src/backend.rs @@ -0,0 +1,304 @@ +use burn_backend::{ + BackTrace, Backend, BackendTypes, DType, DTypeUsage, DeviceId, DeviceOps, ExecutionError, + tensor::Device, +}; +use burn_std::{ + BoolStore, DeviceSettings, + rand::{SeedableRng, StdRng}, + stub::Mutex, +}; +use candle_core::{DeviceLocation, backend::BackendDevice}; + +use crate::{ + CandleTensor, IntoDType, + element::{CandleElement, FloatCandleElement, IntCandleElement}, +}; + +/// Tensor backend that uses the [candle](candle_core) crate for executing tensor operations. +/// +/// It is compatible with a wide range of hardware configurations, including CPUs and GPUs +/// that support CUDA or Metal. Additionally, the backend can be compiled to `wasm` when using the CPU. +#[derive(Clone, Default, Debug)] +pub struct Candle {} + +// Seed for CPU device +pub(crate) static SEED: Mutex> = Mutex::new(None); + +pub(crate) fn get_seeded_rng() -> StdRng { + let mut seed = SEED.lock().unwrap(); + seed.take().unwrap_or_else(burn_std::rand::get_seeded_rng) +} + +pub(crate) fn set_seeded_rng(rng_seeded: StdRng) { + let mut seed = SEED.lock().unwrap(); + *seed = Some(rng_seeded); +} + +/// The device type for the candle backend. +#[derive(Clone, Debug, PartialEq, Eq)] +/// The device struct when using the `candle` backend. +/// +/// To create a Cuda or Metal device from the index, use the associated methods to create the variant: +/// ```no_run +/// use burn_candle::CandleDevice; +/// +/// // Create a Cuda device from its index +/// let device = CandleDevice::cuda(0); +/// // Create a Metal device from its index +/// let device = CandleDevice::metal(0); +/// ``` +#[derive(Default)] +pub enum CandleDevice { + /// CPU device. + #[default] + Cpu, + + /// Cuda device with the given index. The index is the index of the Cuda device in the list of + /// all Cuda devices found on the system. + Cuda(CudaDevice), + + /// Metal device with the given index. The index is the index of the Metal device in the list of + /// all Metal devices found on the system. + Metal(MetalDevice), +} + +impl CandleDevice { + /// Create a Cuda device with the given index. + /// The index is the index of the Cuda device in the list of all Cuda devices found on the system. + pub fn cuda(index: usize) -> Self { + CandleDevice::Cuda(CudaDevice { + device: candle_core::CudaDevice::new(index).unwrap(), + index, + }) + } + + /// Create a Metal device with the given index. + /// The index is the index of the Metal device in the list of all Metal devices found on the system. + pub fn metal(index: usize) -> Self { + CandleDevice::Metal(MetalDevice { + device: candle_core::MetalDevice::new(index).unwrap(), + index, + }) + } + + pub(crate) fn set_seed(&self, seed: u64) { + match self { + CandleDevice::Cpu => { + // candle_core::cpu_backend::CpuDevice.set_seed(seed).unwrap(); + // Candle does not support seeding the CPU rng so we use a global seed + let rng = StdRng::seed_from_u64(seed); + set_seeded_rng(rng); + } + CandleDevice::Cuda(cuda_device) => cuda_device.device.set_seed(seed).unwrap(), + CandleDevice::Metal(metal_device) => metal_device.device.set_seed(seed).unwrap(), + } + } +} + +#[derive(Clone, Debug)] +/// A Cuda device for the `candle` backend. +pub struct CudaDevice { + pub(crate) device: candle_core::CudaDevice, + /// The index of the Cuda device in the list of all devices on the system. + pub index: usize, +} + +impl PartialEq for CudaDevice { + fn eq(&self, other: &Self) -> bool { + self.device.same_device(&other.device) && self.index == other.index + } +} + +impl Eq for CudaDevice {} + +#[derive(Clone, Debug)] +/// A Metal device for the `candle` backend. +pub struct MetalDevice { + pub(crate) device: candle_core::MetalDevice, + /// The index of the Metal device in the list of all devices on the system. + pub index: usize, +} + +impl PartialEq for MetalDevice { + fn eq(&self, other: &Self) -> bool { + self.device.same_device(&other.device) && self.index == other.index + } +} + +impl Eq for MetalDevice {} + +impl From for candle_core::Device { + fn from(device: CandleDevice) -> Self { + match device { + CandleDevice::Cpu => candle_core::Device::Cpu, + CandleDevice::Cuda(device) => candle_core::Device::Cuda(device.device), + CandleDevice::Metal(device) => candle_core::Device::Metal(device.device), + } + } +} + +impl From for CandleDevice { + fn from(device: candle_core::Device) -> Self { + match device.location() { + DeviceLocation::Cpu => CandleDevice::Cpu, + DeviceLocation::Cuda { gpu_id } => { + if let candle_core::Device::Cuda(device) = device { + CandleDevice::Cuda(CudaDevice { + device, + index: gpu_id, + }) + } else { + panic!("Expected CUDA device."); + } + } + DeviceLocation::Metal { gpu_id } => { + if let candle_core::Device::Metal(device) = device { + CandleDevice::Metal(MetalDevice { + device, + index: gpu_id, + }) + } else { + panic!("Expected Metal device."); + } + } + } + } +} + +impl burn_backend::Device for CandleDevice { + fn to_id(&self) -> burn_backend::DeviceId { + match self { + CandleDevice::Cuda(device) => DeviceId::new(0, device.index as u16), + CandleDevice::Metal(device) => DeviceId::new(1, device.index as u16), + CandleDevice::Cpu => DeviceId::new(2, 0), + } + } + + fn from_id(device_id: DeviceId) -> Self { + match device_id.type_id { + 0 => CandleDevice::cuda(device_id.index_id as usize), + 1 => CandleDevice::metal(device_id.index_id as usize), + _ => CandleDevice::Cpu, + } + } +} +impl DeviceOps for CandleDevice { + fn defaults(&self) -> DeviceSettings { + DeviceSettings::new( + DType::F32, + DType::I64, + DType::Bool(BoolStore::U8), + Default::default(), + ) + } +} + +impl BackendTypes for Candle { + type Device = CandleDevice; + + type FloatTensorPrimitive = CandleTensor; + + type IntTensorPrimitive = CandleTensor; + + type BoolTensorPrimitive = CandleTensor; + + type QuantizedTensorPrimitive = CandleTensor; + + type GraphPrimitive = burn_backend::GraphUnsupported; +} + +impl Backend for Candle { + fn ad_enabled(_device: &Self::Device) -> bool { + false + } + + fn name(device: &Self::Device) -> String { + match device { + CandleDevice::Cpu => "candle", + CandleDevice::Cuda(..) => "candle", + CandleDevice::Metal(..) => "candle", + } + .to_string() + } + + fn seed(device: &CandleDevice, seed: u64) { + device.set_seed(seed); + } + + fn sync(device: &Device) -> Result<(), ExecutionError> { + let device: candle_core::Device = (device.clone()).into(); + + match device { + candle_core::Device::Cpu => (), + candle_core::Device::Cuda(device) => { + #[cfg(feature = "cuda")] + device + .synchronize() + .map_err(|err| ExecutionError::Generic { + reason: format!("Can't sync the cuda device: {err}"), + backtrace: BackTrace::capture(), + })?; + } + candle_core::Device::Metal(device) => { + // For some reason, device.wait_until_completed() does not seem to work, + // and neither does writing and reading a value with into_data + return Err(ExecutionError::Generic { + reason: + "Device synchronization unavailable with Metal device on Candle backend" + .into(), + backtrace: BackTrace::capture(), + }); + } + } + + Ok(()) + } + + fn dtype_usage(device: &Self::Device, dtype: DType) -> burn_backend::DTypeUsageSet { + if dtype.try_into_dtype().is_ok() { + burn_backend::DTypeUsage::general() + } else { + burn_backend::DTypeUsageSet::empty() + } + } + + fn device_count(_: u16) -> usize { + 1 + } + + fn flush(_device: &Self::Device) {} +} + +#[cfg(test)] +mod tests { + use burn_std::{BoolStore, QuantScheme}; + + use super::*; + + #[test] + fn should_support_dtypes() { + type B = Candle; + let device = Default::default(); + + assert!(B::supports_dtype(&device, DType::F64)); + assert!(B::supports_dtype(&device, DType::F32)); + assert!(B::supports_dtype(&device, DType::Flex32)); + assert!(B::supports_dtype(&device, DType::F16)); + assert!(B::supports_dtype(&device, DType::BF16)); + assert!(B::supports_dtype(&device, DType::I64)); + assert!(B::supports_dtype(&device, DType::U32)); + assert!(B::supports_dtype(&device, DType::U8)); + assert!(B::supports_dtype(&device, DType::I32)); + assert!(B::supports_dtype(&device, DType::I16)); + assert!(B::supports_dtype(&device, DType::Bool(BoolStore::U8))); + + assert!(!B::supports_dtype(&device, DType::U64)); + assert!(!B::supports_dtype(&device, DType::U16)); + assert!(!B::supports_dtype(&device, DType::I8)); + assert!(!B::supports_dtype(&device, DType::Bool(BoolStore::Native))); + assert!(!B::supports_dtype( + &device, + DType::QFloat(QuantScheme::default()) + )); + } +} diff --git a/crates/burn-candle/src/element.rs b/crates/burn-candle/src/element.rs new file mode 100644 index 0000000..bc78c13 --- /dev/null +++ b/crates/burn-candle/src/element.rs @@ -0,0 +1,32 @@ +use std::borrow::Borrow; + +use burn_backend::{Element, bf16, f16}; +use candle_core::{FloatDType, Tensor, WithDType}; + +/// Candle element +pub trait CandleElement: Element + WithDType {} +/// Candle float element +pub trait FloatCandleElement: CandleElement + FloatDType {} +/// Candle int element +pub trait IntCandleElement: CandleElement {} + +impl CandleElement for f64 {} +impl FloatCandleElement for f64 {} + +impl CandleElement for f32 {} +impl FloatCandleElement for f32 {} + +impl CandleElement for f16 {} +impl FloatCandleElement for f16 {} + +impl CandleElement for bf16 {} +impl FloatCandleElement for bf16 {} + +impl CandleElement for u8 {} +impl IntCandleElement for u8 {} + +impl CandleElement for u32 {} +impl IntCandleElement for u32 {} + +impl CandleElement for i64 {} +impl IntCandleElement for i64 {} diff --git a/crates/burn-candle/src/lib.rs b/crates/burn-candle/src/lib.rs new file mode 100644 index 0000000..3138378 --- /dev/null +++ b/crates/burn-candle/src/lib.rs @@ -0,0 +1,27 @@ +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] +#![allow(unused)] // TODO remove when backend filled +#![deprecated( + since = "0.21.0", + note = "burn-candle is deprecated and will be removed in a future release. Use burn-cubecl (CUDA/ROCm/Vulkan/Metal/WebGPU), burn-flex, or burn-tch instead." +)] + +//! Burn Candle Backend +//! +//! **Deprecated:** This backend is deprecated and will be removed in a future release. +//! Please migrate to one of the actively maintained backends: +//! - CubeCL backends (CUDA, ROCm, Vulkan, Metal, WebGPU) for GPU acceleration +//! - Flex (`burn-flex`) for portable pure-Rust CPU execution (std, no_std, WASM) +//! - LibTorch (`burn-tch`) for a mature CPU/GPU backend + +#[macro_use] +extern crate derive_new; + +mod backend; +mod element; +mod ops; +mod tensor; + +pub use backend::*; +pub use element::*; +pub use tensor::*; diff --git a/crates/burn-candle/src/ops/activation.rs b/crates/burn-candle/src/ops/activation.rs new file mode 100644 index 0000000..b7ffe33 --- /dev/null +++ b/crates/burn-candle/src/ops/activation.rs @@ -0,0 +1,17 @@ +use burn_backend::{ops::ActivationOps, tensor::FloatTensor}; + +use crate::{ + Candle, CandleTensor, + element::{CandleElement, FloatCandleElement, IntCandleElement}, + tensor, +}; + +impl ActivationOps for Candle { + fn gelu(tensor: FloatTensor) -> FloatTensor { + CandleTensor::new(tensor.tensor.gelu().unwrap()) + } + + fn relu(tensor: FloatTensor) -> FloatTensor { + CandleTensor::new(tensor.tensor.relu().unwrap()) + } +} diff --git a/crates/burn-candle/src/ops/base.rs b/crates/burn-candle/src/ops/base.rs new file mode 100644 index 0000000..0f66548 --- /dev/null +++ b/crates/burn-candle/src/ops/base.rs @@ -0,0 +1,691 @@ +use std::cmp::max; +use std::marker::PhantomData; + +use crate::{ + Candle, CandleDevice, CandleTensor, + element::{CandleElement, FloatCandleElement, IntCandleElement}, +}; +use burn_backend::{ + BackTrace, Backend, Distribution, ExecutionError, Slice, bf16, f16, + ops::unfold::{calculate_unfold_shape, calculate_unfold_windows}, +}; +use burn_backend::{DType, Element, Shape, TensorData, TensorMetadata}; +use candle_core::{Layout, WithDType}; + +use super::tensor; + +pub fn cpu_random(shape: Shape, distribution: Distribution, dtype: DType) -> TensorData { + let mut rng = crate::get_seeded_rng(); + let data = match dtype { + DType::F64 => TensorData::random::(shape, distribution, &mut rng), + DType::F32 => TensorData::random::(shape, distribution, &mut rng), + DType::Flex32 => TensorData::random::(shape, distribution, &mut rng), + DType::F16 => TensorData::random::(shape, distribution, &mut rng), + DType::BF16 => { + TensorData::random::(shape, distribution, &mut rng) + } + DType::I64 => TensorData::random::(shape, distribution, &mut rng), + DType::U32 => TensorData::random::(shape, distribution, &mut rng), + DType::U8 => TensorData::random::(shape, distribution, &mut rng), + other => unimplemented!("DType not supported {dtype:?}"), + }; + crate::set_seeded_rng(rng); + data +} + +pub fn cat(tensors: Vec, dim: usize) -> CandleTensor { + let tensors: Vec = tensors.into_iter().map(|t| t.tensor).collect(); + CandleTensor::new(candle_core::Tensor::cat(&tensors, dim).unwrap()) +} + +pub fn from_data(data: TensorData, device: &CandleDevice) -> CandleTensor { + CandleTensor::from_data::(data, device.clone()) +} +pub fn into_data(tensor: CandleTensor) -> Result { + fn tensor_data_from_dtype( + tensor: &CandleTensor, + ) -> Result { + let data = tensor + .tensor + .flatten_all() + .map_err(|err| ExecutionError::Generic { + reason: format!("{err}"), + backtrace: BackTrace::capture(), + })? + .to_vec1::() + .map_err(|err| ExecutionError::Generic { + reason: format!("{err}"), + backtrace: BackTrace::capture(), + })?; + Ok(TensorData::new(data, tensor.shape())) + } + + match tensor.tensor.dtype() { + candle_core::DType::BF16 => tensor_data_from_dtype::(&tensor), + candle_core::DType::F16 => tensor_data_from_dtype::(&tensor), + candle_core::DType::F32 => tensor_data_from_dtype::(&tensor), + candle_core::DType::F64 => tensor_data_from_dtype::(&tensor), + candle_core::DType::U8 => tensor_data_from_dtype::(&tensor), + candle_core::DType::U32 => tensor_data_from_dtype::(&tensor), + candle_core::DType::I16 => tensor_data_from_dtype::(&tensor), + candle_core::DType::I32 => tensor_data_from_dtype::(&tensor), + candle_core::DType::I64 => tensor_data_from_dtype::(&tensor), + other => todo!("{other:?} not yet supported"), + } +} + +pub fn to_device(tensor: CandleTensor, device: &CandleDevice) -> CandleTensor { + CandleTensor::new(tensor.tensor.to_device(&(device.clone()).into()).unwrap()) +} + +pub fn empty(shape: Shape, device: &CandleDevice, dtype: candle_core::DType) -> CandleTensor { + zeros(shape, device, dtype) +} + +/// Flatten K-dimensional index tuples into 1D linear offsets. +/// +/// `indices` has shape `[num_updates, k]` (i64). +/// Returns a 1D tensor of shape `[num_updates * slice_size]` (u32, for candle index ops). +fn flatten_nd_indices( + indices: &candle_core::Tensor, + data_shape: &[usize], + k: usize, + slice_size: usize, +) -> candle_core::Result { + let device = indices.device(); + let num_updates = indices.dim(0)?; + + // Compute per-dimension strides for the first K dims + let mut strides = vec![0i64; k]; + if k > 0 { + strides[k - 1] = slice_size as i64; + for i in (0..k - 1).rev() { + strides[i] = strides[i + 1] * data_shape[i + 1] as i64; + } + } + + let strides_tensor = candle_core::Tensor::from_slice(&strides, (k,), device)?; + // base_offsets: [num_updates] = indices @ strides + let base_offsets = indices + .to_dtype(candle_core::DType::I64)? + .matmul(&strides_tensor.unsqueeze(1)?)? + .squeeze(1)?; + + // Expand to [num_updates, slice_size] + intra-slice offsets + let slice_offsets = candle_core::Tensor::arange(0i64, slice_size as i64, device)?; + base_offsets + .unsqueeze(1)? + .broadcast_add(&slice_offsets.unsqueeze(0)?)? + .reshape((num_updates * slice_size,))? + .to_dtype(candle_core::DType::U32) +} + +pub fn scatter_nd( + data: CandleTensor, + indices: CandleTensor, + values: CandleTensor, + reduction: burn_backend::tensor::IndexingUpdateOp, +) -> CandleTensor { + try_scatter_nd(data, indices, values, reduction) + .unwrap_or_else(|e| panic!("candle scatter_nd: {e}")) +} + +fn try_scatter_nd( + data: CandleTensor, + indices: CandleTensor, + values: CandleTensor, + reduction: burn_backend::tensor::IndexingUpdateOp, +) -> candle_core::Result { + use burn_backend::tensor::IndexingUpdateOp; + + let data_shape: Vec = data.tensor.dims().to_vec(); + let idx_shape: Vec = indices.tensor.dims().to_vec(); + let m = idx_shape.len(); + let k = idx_shape[m - 1]; + let num_updates: usize = idx_shape[..m - 1].iter().product(); + let total_elems: usize = data_shape.iter().product(); + let slice_size: usize = data_shape[k..].iter().product(); + + let flat_indices = indices.tensor.reshape((num_updates, k))?; + let linear_idx = flatten_nd_indices(&flat_indices, &data_shape, k, slice_size)?; + + let flat_data = data.tensor.contiguous()?.reshape(total_elems)?; + let flat_values = values + .tensor + .contiguous()? + .reshape(num_updates * slice_size)?; + + let result = match reduction { + IndexingUpdateOp::Assign => flat_data.scatter(&linear_idx, &flat_values, 0)?, + IndexingUpdateOp::Add => flat_data.scatter_add(&linear_idx, &flat_values, 0)?, + _ => panic!( + "scatter_nd with {:?} reduction is not supported by the candle backend", + reduction + ), + }; + + Ok(CandleTensor::new(result.reshape(data_shape)?)) +} + +pub fn gather_nd(data: CandleTensor, indices: CandleTensor) -> CandleTensor { + try_gather_nd(data, indices).unwrap_or_else(|e| panic!("candle gather_nd: {e}")) +} + +fn try_gather_nd(data: CandleTensor, indices: CandleTensor) -> candle_core::Result { + let data_shape: Vec = data.tensor.dims().to_vec(); + let idx_shape: Vec = indices.tensor.dims().to_vec(); + let m = idx_shape.len(); + let k = idx_shape[m - 1]; + let num_indices: usize = idx_shape[..m - 1].iter().product(); + let total_elems: usize = data_shape.iter().product(); + let slice_size: usize = data_shape[k..].iter().product(); + + let flat_indices = indices.tensor.reshape((num_indices, k))?; + let linear_idx = flatten_nd_indices(&flat_indices, &data_shape, k, slice_size)?; + + let flat_data = data.tensor.contiguous()?.reshape(total_elems)?; + let gathered = flat_data.index_select(&linear_idx, 0)?; + + // Output shape: idx_shape[..m-1] ++ data_shape[k..] + let mut out_shape: Vec = idx_shape[..m - 1].to_vec(); + out_shape.extend_from_slice(&data_shape[k..]); + + Ok(CandleTensor::new(gathered.reshape(out_shape)?)) +} + +pub fn zeros(shape: Shape, device: &CandleDevice, dtype: candle_core::DType) -> CandleTensor { + CandleTensor::new( + candle_core::Tensor::zeros(shape.to_vec(), dtype, &(device.clone()).into()).unwrap(), + ) +} + +pub fn ones(shape: Shape, device: &CandleDevice, dtype: candle_core::DType) -> CandleTensor { + CandleTensor::new( + candle_core::Tensor::ones(shape.to_vec(), dtype, &(device.clone()).into()).unwrap(), + ) +} + +pub fn swap_dims(mut tensor: CandleTensor, dim1: usize, dim2: usize) -> CandleTensor { + CandleTensor::new(tensor.tensor.transpose(dim1, dim2).unwrap()) +} + +pub fn permute(tensor: CandleTensor, axes: &[usize]) -> CandleTensor { + CandleTensor::new(tensor.tensor.permute(axes).unwrap()) +} + +pub fn flip(tensor: CandleTensor, axes: &[usize]) -> CandleTensor { + // FIXME: Replace with an appropriate method when Candle provides one. + let mut tensor = tensor.tensor; + for &axis in axes { + // Ensure tensor is contiguous before index_select (required by Candle) + tensor = tensor.contiguous().unwrap(); + + let indexes = candle_core::Tensor::arange_step( + tensor.dim(axis).unwrap() as i64 - 1, + -1, + -1, + tensor.device(), + ) + .unwrap(); + tensor = tensor.index_select(&indexes, axis).unwrap(); + } + + CandleTensor::new(tensor) +} + +pub fn reshape(tensor: CandleTensor, shape: Shape) -> CandleTensor { + CandleTensor::new(tensor.tensor.reshape(shape.to_vec()).unwrap()) +} + +pub fn shape(tensor: &CandleTensor) -> Shape { + tensor.shape() +} + +pub fn slice(tensor: CandleTensor, ranges: &[std::ops::Range]) -> CandleTensor { + let mut narrow_tensor = tensor.tensor; + for (i, range) in ranges.iter().enumerate().take(ranges.len()) { + narrow_tensor = narrow_tensor + .narrow(i, range.start, range.end - range.start) + .unwrap() + } + CandleTensor::new(narrow_tensor) +} + +pub fn slice_with_steps(tensor: CandleTensor, slices: &[Slice]) -> CandleTensor { + let mut result_tensor = tensor.tensor; + + for (dim, slice) in slices.iter().enumerate() { + if slice.step == 1 { + // Use narrow for step=1 (more efficient) + // Convert slice to range using tensor shape + let dim_size = result_tensor.dim(dim).unwrap(); + let range = slice.to_range(dim_size); + let start = range.start; + let length = range.end - range.start; + result_tensor = result_tensor.narrow(dim, start, length).unwrap(); + } else { + // Use index_select for step != 1 + let dim_size = result_tensor.dim(dim).unwrap(); + let range = slice.to_range(dim_size); + let start = range.start; + let end = range.end; + let step = slice.step; + + // Generate indices based on step direction + let indices_vec = if step > 0 { + // Forward stepping + let step_usize = step as usize; + (start..end).step_by(step_usize).collect::>() + } else { + // Backward stepping (negative step) + let step_usize = step.unsigned_abs(); + // Start from end-1 and go backwards + let mut indices = Vec::new(); + let mut idx = end - 1; + while idx >= start && idx < end { + // Check for underflow + indices.push(idx); + if idx >= step_usize { + idx -= step_usize; + } else { + break; + } + } + indices + }; + + // Convert indices to tensor and use index_select + let indices_len = indices_vec.len(); + let device = result_tensor.device(); + let indices = candle_core::Tensor::from_vec( + indices_vec.iter().map(|&x| x as u32).collect::>(), + indices_len, + device, + ) + .unwrap(); + + result_tensor = result_tensor.index_select(&indices, dim).unwrap(); + } + } + + CandleTensor::new(result_tensor) +} + +pub fn slice_assign(tensor: CandleTensor, slices: &[Slice], value: CandleTensor) -> CandleTensor { + // Check if all slices have step=1 (candle's native slice_assign requirement) + let all_unit_steps = slices.iter().all(|s| s.step == 1); + + if all_unit_steps { + // Convert Slice to Range for candle's native slice_assign + let ranges: Vec> = slices + .iter() + .enumerate() + .map(|(dim, slice)| { + let dim_size = tensor.tensor.dim(dim).unwrap_or(usize::MAX); + slice.to_range(dim_size) + }) + .collect(); + + CandleTensor::new(tensor.tensor.slice_assign(&ranges, &value.tensor).unwrap()) + } else { + // Implement slice_assign with steps using scatter operations + slice_assign_with_steps_workaround(tensor, slices, value) + } +} + +/// Implements slice_assign for non-unit steps using index operations +fn slice_assign_with_steps_workaround( + tensor: CandleTensor, + slices: &[Slice], + value: CandleTensor, +) -> CandleTensor { + let shape = tensor.shape(); + let ndims = shape.num_dims(); + let device = tensor.tensor.device(); + + // Generate indices for each dimension based on slice specifications + let indices_per_dim = generate_slice_indices(slices, &shape); + + // Early return if no elements to assign + let total_elements: usize = indices_per_dim.iter().map(|v| v.len()).product(); + if total_elements == 0 { + return tensor; + } + + // Flatten tensors and get metadata + let value_flat = value.tensor.flatten_all().unwrap(); + let strides = tensor.tensor.stride(); + let tensor_shape = tensor.tensor.dims(); + + // Use a macro to handle different dtypes without code duplication + macro_rules! apply_slice_assign { + ($dtype:ty, $to_vec_fn:ident) => {{ + let mut tensor_vec: Vec<$dtype> = + tensor.tensor.flatten_all().unwrap().$to_vec_fn().unwrap(); + let value_vec: Vec<$dtype> = value_flat.$to_vec_fn().unwrap(); + + // Apply assignments using cartesian product of indices + for (value_idx, &value) in value_vec.iter().enumerate() { + let flat_idx = compute_flat_index(value_idx, &indices_per_dim, &strides); + if flat_idx < tensor_vec.len() { + tensor_vec[flat_idx] = value; + } + } + + candle_core::Tensor::from_vec(tensor_vec, tensor_shape, device).unwrap() + }}; + } + + use candle_core::DType; + let result = match tensor.tensor.dtype() { + DType::F32 => apply_slice_assign!(f32, to_vec1), + DType::F64 => apply_slice_assign!(f64, to_vec1), + DType::I64 => apply_slice_assign!(i64, to_vec1), + DType::U32 => apply_slice_assign!(u32, to_vec1), + DType::U8 => apply_slice_assign!(u8, to_vec1), + _ => panic!( + "Unsupported dtype {:?} for slice_assign with steps", + tensor.tensor.dtype() + ), + }; + + CandleTensor::new(result) +} + +/// Generate indices for each dimension based on slice specifications +fn generate_slice_indices(slices: &[Slice], tensor_dims: &[usize]) -> Vec> { + let ndims = tensor_dims.len(); + let mut indices_per_dim = Vec::with_capacity(ndims); + + // Process provided slices + for (dim_idx, slice) in slices.iter().enumerate() { + let dim_size = tensor_dims[dim_idx]; + let range = slice.to_range(dim_size); + let indices = generate_stepped_indices(range.start, range.end, slice.step); + indices_per_dim.push(indices); + } + + // Fill remaining dimensions with full ranges + for &dim_size in tensor_dims.iter().skip(slices.len()) { + indices_per_dim.push((0..dim_size).collect()); + } + + indices_per_dim +} + +/// Generate indices for a single dimension with stepping +fn generate_stepped_indices(start: usize, end: usize, step: isize) -> Vec { + if step > 0 { + // Forward stepping + (start..end).step_by(step as usize).collect() + } else if step < 0 { + // Backward stepping: start from end-1 and go backwards + let step_size = step.unsigned_abs(); + let mut indices = Vec::new(); + let mut idx = end.saturating_sub(1); + + while idx >= start && idx < end { + indices.push(idx); + if idx >= step_size { + idx -= step_size; + } else { + break; + } + } + indices + } else { + // This branch should never be reached since step is validated to be non-zero + panic!("Step cannot be zero") + } +} + +/// Compute flat index from multi-dimensional indices using cartesian product logic +fn compute_flat_index( + value_idx: usize, + indices_per_dim: &[Vec], + strides: &[usize], +) -> usize { + let mut flat_idx = 0; + let mut remainder = value_idx; + + // Convert value_idx to multi-dimensional indices and compute flat tensor index + for dim in (0..indices_per_dim.len()).rev() { + let dim_size = indices_per_dim[dim].len(); + let idx_in_dim = remainder % dim_size; + remainder /= dim_size; + + let actual_idx = indices_per_dim[dim][idx_in_dim]; + flat_idx += actual_idx * strides[dim]; + } + + flat_idx +} + +pub fn narrow(tensor: CandleTensor, dim: usize, start: usize, length: usize) -> CandleTensor { + let tensor = tensor.tensor.narrow(dim, start, length); + match tensor { + Ok(tensor) => CandleTensor::new(tensor), + Err(e) => panic!("error narrow from Candle"), + } +} + +pub fn chunk(tensor: CandleTensor, chunks: usize, dim: usize) -> Vec { + let tensors = tensor.tensor.chunk(chunks, dim); + match tensors { + Ok(tensors) => tensors.into_iter().map(CandleTensor::new).collect(), + Err(e) => panic!("error chunk from Candle"), + } +} + +pub fn expand(tensor: CandleTensor, shape: Shape) -> CandleTensor { + CandleTensor::new(tensor.tensor.broadcast_as(shape.to_vec()).unwrap()) +} + +pub fn unfold(tensor: CandleTensor, dim: usize, size: usize, step: usize) -> CandleTensor { + let result_shape = calculate_unfold_shape(tensor.shape(), dim, size, step); + let windows = result_shape[dim]; + + let mut select_ranges = tensor.shape().into_ranges(); + let new_axis = select_ranges.len(); + + let mut stack = Vec::with_capacity(windows); + for widx in 0..windows { + let start = widx * step; + let end = start + size; + select_ranges[dim] = start..end; + + let mut window_slice = slice(tensor.clone(), &select_ranges); + + window_slice = swap_dims(window_slice, dim, new_axis); + let window_slice = CandleTensor::new(window_slice.tensor.unsqueeze(new_axis).unwrap()); + + stack.push(window_slice); + } + cat(stack, dim) +} + +pub fn sign(tensor: CandleTensor) -> CandleTensor { + CandleTensor::new(tensor.tensor.sign().unwrap()) +} + +pub fn mask_where_broadcasted( + tensor: CandleTensor, + mask: CandleTensor, + value: CandleTensor, +) -> CandleTensor { + let shape = tensor + .tensor + .shape() + .broadcast_shape_binary_op(mask.tensor.shape(), "where_cond") + .unwrap(); + + let mut tensor = tensor.tensor; + let mut mask = mask.tensor; + let mut value = value.tensor; + + if shape != *tensor.shape() { + tensor = tensor.broadcast_as(shape.clone()).unwrap(); + } + if shape != *mask.shape() { + mask = mask.broadcast_as(shape.clone()).unwrap(); + } + if shape != *value.shape() { + value = value.broadcast_as(shape).unwrap(); + } + + CandleTensor::new(mask.where_cond(&value, &tensor).unwrap()) +} + +pub fn cross(lhs: CandleTensor, rhs: CandleTensor, dim: usize) -> CandleTensor { + let shape_lhs = lhs.shape(); + let shape_rhs = rhs.shape(); + let ndims = shape_lhs.num_dims(); + + // Broadcast the shapes except along dim + let mut broadcast_shape = vec![0; ndims]; + for (i, item) in broadcast_shape.iter_mut().enumerate().take(ndims) { + if i == dim { + *item = shape_lhs[i]; + } else { + let l = shape_lhs[i]; + let r = shape_rhs[i]; + if l == r { + *item = l; + } else if l == 1 { + *item = r; + } else if r == 1 { + *item = l; + } else { + panic!("Tensors are not broadcastable along dimension {}", i); + } + } + } + + // Broadcast lhs and rhs + let lhs_broadcast = if shape_lhs == Shape::from(broadcast_shape.clone()) { + lhs + } else { + expand(lhs, Shape::from(broadcast_shape.clone())) + }; + let rhs_broadcast = if shape_rhs == Shape::from(broadcast_shape.clone()) { + rhs + } else { + expand(rhs, Shape::from(broadcast_shape.clone())) + }; + + // Now, move dim to the last dimension + let mut perm = (0..ndims).collect::>(); + perm.remove(dim); + perm.push(dim); + + let lhs_permuted = permute(lhs_broadcast, &perm); + let rhs_permuted = permute(rhs_broadcast, &perm); + + // Reshape to (*, 3) + let total_elements = lhs_permuted.shape().num_elements(); + let batch_size = total_elements / 3; + let lhs_reshaped = reshape(lhs_permuted, Shape::new([batch_size, 3])); + let rhs_reshaped = reshape(rhs_permuted, Shape::new([batch_size, 3])); + + // Extract components using narrow and squeeze + let lhs_0 = CandleTensor::new( + lhs_reshaped + .tensor + .narrow(1, 0, 1) + .unwrap() + .squeeze(1) + .unwrap(), + ); + let lhs_1 = CandleTensor::new( + lhs_reshaped + .tensor + .narrow(1, 1, 1) + .unwrap() + .squeeze(1) + .unwrap(), + ); + let lhs_2 = CandleTensor::new( + lhs_reshaped + .tensor + .narrow(1, 2, 1) + .unwrap() + .squeeze(1) + .unwrap(), + ); + let rhs_0 = CandleTensor::new( + rhs_reshaped + .tensor + .narrow(1, 0, 1) + .unwrap() + .squeeze(1) + .unwrap(), + ); + let rhs_1 = CandleTensor::new( + rhs_reshaped + .tensor + .narrow(1, 1, 1) + .unwrap() + .squeeze(1) + .unwrap(), + ); + let rhs_2 = CandleTensor::new( + rhs_reshaped + .tensor + .narrow(1, 2, 1) + .unwrap() + .squeeze(1) + .unwrap(), + ); + + // Compute cross product components + let result_0 = CandleTensor::new( + lhs_1 + .tensor + .mul(&rhs_2.tensor) + .unwrap() + .sub(&lhs_2.tensor.mul(&rhs_1.tensor).unwrap()) + .unwrap(), + ); + let result_1 = CandleTensor::new( + lhs_2 + .tensor + .mul(&rhs_0.tensor) + .unwrap() + .sub(&lhs_0.tensor.mul(&rhs_2.tensor).unwrap()) + .unwrap(), + ); + let result_2 = CandleTensor::new( + lhs_0 + .tensor + .mul(&rhs_1.tensor) + .unwrap() + .sub(&lhs_1.tensor.mul(&rhs_0.tensor).unwrap()) + .unwrap(), + ); + + // Stack the components + let result_0_unsqueezed = CandleTensor::new(result_0.tensor.unsqueeze(1).unwrap()); + let result_1_unsqueezed = CandleTensor::new(result_1.tensor.unsqueeze(1).unwrap()); + let result_2_unsqueezed = CandleTensor::new(result_2.tensor.unsqueeze(1).unwrap()); + let result = cat( + vec![ + result_0_unsqueezed, + result_1_unsqueezed, + result_2_unsqueezed, + ], + 1, + ); + + // Reshape back to the broadcast shape with dim at the end + let mut result_shape = broadcast_shape; + result_shape.remove(dim); + result_shape.push(3); + let result_reshaped = reshape(result, Shape::from(result_shape)); + + // Permute back + let mut inv_perm = vec![0; ndims]; + for (i, &p) in perm.iter().enumerate() { + inv_perm[p] = i; + } + permute(result_reshaped, &inv_perm) +} diff --git a/crates/burn-candle/src/ops/bool_tensor.rs b/crates/burn-candle/src/ops/bool_tensor.rs new file mode 100644 index 0000000..27c2a3e --- /dev/null +++ b/crates/burn-candle/src/ops/bool_tensor.rs @@ -0,0 +1,209 @@ +use burn_backend::{ + BackTrace, DType, ExecutionError, Scalar, Shape, Slice, TensorData, TensorMetadata, + ops::BoolTensorOps, + tensor::{BoolTensor, Device, FloatTensor, IntTensor}, +}; +use burn_std::{BoolDType, FloatDType, IntDType}; + +use crate::{ + Candle, CandleTensor, IntoDType, + element::{CandleElement, FloatCandleElement, IntCandleElement}, +}; + +use super::base::{expand, permute, unfold}; + +impl BoolTensorOps for Candle { + fn bool_empty(shape: Shape, device: &Device, _dtype: BoolDType) -> BoolTensor { + super::base::empty(shape, device, candle_core::DType::U8) + } + + fn bool_zeros(shape: Shape, device: &Device, _dtype: BoolDType) -> BoolTensor { + super::base::zeros(shape, device, candle_core::DType::U8) + } + + fn bool_ones(shape: Shape, device: &Device, _dtype: BoolDType) -> BoolTensor { + super::base::ones(shape, device, candle_core::DType::U8) + } + + async fn bool_into_data(tensor: BoolTensor) -> Result { + let x: Vec = tensor + .tensor + .flatten_all() + .map_err(|err| ExecutionError::Generic { + reason: format!("{err}"), + backtrace: BackTrace::capture(), + })? + .to_vec1() + .map_err(|err| ExecutionError::Generic { + reason: format!("{err}"), + backtrace: BackTrace::capture(), + })?; + + let y = x.iter().map(|b| !matches!(b, 0)).collect(); + + Ok(TensorData::new(y, tensor.shape())) + } + + fn bool_from_data(data: TensorData, device: &Device) -> BoolTensor { + match data.dtype { + DType::U8 => super::base::from_data::(data, device), + _ => unimplemented!("Unsupported dtype for `bool_from_data`"), + } + } + + fn bool_into_int(tensor: BoolTensor, out_dtype: IntDType) -> IntTensor { + CandleTensor::new(tensor.tensor.to_dtype(out_dtype.into_dtype()).unwrap()) + } + + fn bool_into_float(tensor: BoolTensor, out_dtype: FloatDType) -> FloatTensor { + CandleTensor::new(tensor.tensor.to_dtype(out_dtype.into_dtype()).unwrap()) + } + + fn bool_to_device(tensor: BoolTensor, device: &Device) -> BoolTensor { + super::base::to_device(tensor, device) + } + + fn bool_reshape(tensor: BoolTensor, shape: Shape) -> BoolTensor { + super::base::reshape(tensor, shape) + } + + fn bool_slice(tensor: BoolTensor, slices: &[Slice]) -> BoolTensor { + super::base::slice_with_steps(tensor, slices) + } + + fn bool_slice_assign( + tensor: BoolTensor, + slices: &[Slice], + value: BoolTensor, + ) -> BoolTensor { + super::base::slice_assign(tensor, slices, value) + } + + fn bool_cat(tensors: Vec>, dim: usize) -> BoolTensor { + super::base::cat(tensors, dim) + } + + fn bool_equal(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + let (lhs_broadcast, rhs_broadcast) = + super::candle_utils::broadcast_for_comparison(&lhs.tensor, &rhs.tensor).unwrap(); + CandleTensor::new(lhs_broadcast.eq(&rhs_broadcast).unwrap()) + } + + fn bool_not(tensor: BoolTensor) -> BoolTensor { + let x = (candle_core::Tensor::zeros_like(&tensor.tensor).unwrap()); + CandleTensor::new(tensor.tensor.eq(&x).unwrap()) + } + + fn bool_and(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + let x = candle_core::Tensor::ones_like(&lhs.tensor).unwrap(); + CandleTensor::new(lhs.tensor.add(&rhs.tensor).unwrap().gt(&x).unwrap()) + } + + fn bool_or(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + CandleTensor::new( + lhs.tensor + .add(&rhs.tensor) + .unwrap() + .clamp(0u32, 1u32) + .unwrap(), + ) + } + + fn bool_swap_dims(tensor: BoolTensor, dim1: usize, dim2: usize) -> BoolTensor { + super::base::swap_dims(tensor, dim1, dim2) + } + + fn bool_permute(tensor: BoolTensor, axes: &[usize]) -> BoolTensor { + super::base::permute(tensor, axes) + } + + fn bool_flip(tensor: BoolTensor, axes: &[usize]) -> BoolTensor { + super::base::flip(tensor, axes) + } + + fn bool_select( + tensor: BoolTensor, + dim: usize, + indices: IntTensor, + ) -> BoolTensor { + CandleTensor::new(tensor.tensor.index_select(&indices.tensor, dim).unwrap()) + } + + fn bool_select_or( + tensor: BoolTensor, + dim: usize, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor { + CandleTensor::new( + tensor + .tensor + .index_add(&indices.tensor, &value.tensor, dim) + .unwrap(), + ) + } + + fn bool_expand(tensor: BoolTensor, shape: Shape) -> BoolTensor { + expand(tensor, shape) + } + + fn bool_unfold( + tensor: BoolTensor, + dim: usize, + size: usize, + step: usize, + ) -> BoolTensor { + unfold(tensor, dim, size, step) + } + + fn bool_mask_where( + tensor: BoolTensor, + mask: BoolTensor, + value: BoolTensor, + ) -> BoolTensor { + super::base::mask_where_broadcasted(tensor, mask, value) + } + + fn bool_mask_fill( + tensor: BoolTensor, + mask: BoolTensor, + value: Scalar, + ) -> BoolTensor { + CandleTensor::new( + mask.tensor + .where_cond( + &super::candle_utils::fill_like::(value.elem(), &tensor.tensor), + &tensor.tensor, + ) + .unwrap(), + ) + } + + fn bool_gather( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + ) -> BoolTensor { + let tensor = tensor.tensor.contiguous().unwrap(); + let indices = indices.tensor.contiguous().unwrap(); + CandleTensor::new(tensor.gather(&indices, dim).unwrap()) + } + + fn bool_scatter_or( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor { + CandleTensor::new( + tensor + .tensor + .scatter_add(&indices.tensor, &value.tensor, dim) + .unwrap(), + ) + } + + fn bool_equal_elem(lhs: BoolTensor, rhs: Scalar) -> BoolTensor { + CandleTensor::new(lhs.tensor.eq(rhs.elem::()).unwrap()) + } +} diff --git a/crates/burn-candle/src/ops/candle_utils.rs b/crates/burn-candle/src/ops/candle_utils.rs new file mode 100644 index 0000000..957da6e --- /dev/null +++ b/crates/burn-candle/src/ops/candle_utils.rs @@ -0,0 +1,46 @@ +use candle_core::{DType, Device, Shape, Tensor}; + +use crate::element::CandleElement; + +pub(crate) fn fill>( + value: E, + shape: S, + dtype: DType, + device: &Device, +) -> Tensor { + let values = (Tensor::ones((1), dtype, device).unwrap() * value.elem::()).unwrap(); + values.expand(shape).unwrap() +} + +pub(crate) fn fill_like(value: E, reference_tensor: &Tensor) -> Tensor { + fill( + value, + reference_tensor.shape(), + reference_tensor.dtype(), + reference_tensor.device(), + ) +} + +/// Broadcasts two tensors to a common shape for comparison operations +pub(crate) fn broadcast_for_comparison( + lhs: &Tensor, + rhs: &Tensor, +) -> Result<(Tensor, Tensor), candle_core::Error> { + let broadcast_shape = lhs + .shape() + .broadcast_shape_binary_op(rhs.shape(), "comparison")?; + + let lhs = if broadcast_shape != *lhs.shape() { + lhs.broadcast_as(&broadcast_shape)? + } else { + lhs.clone() + }; + + let rhs = if broadcast_shape != *rhs.shape() { + rhs.broadcast_as(&broadcast_shape)? + } else { + rhs.clone() + }; + + Ok((lhs, rhs)) +} diff --git a/crates/burn-candle/src/ops/int_tensor.rs b/crates/burn-candle/src/ops/int_tensor.rs new file mode 100644 index 0000000..fdb0502 --- /dev/null +++ b/crates/burn-candle/src/ops/int_tensor.rs @@ -0,0 +1,570 @@ +use burn_backend::{ + DType, Distribution, ElementConversion, ExecutionError, IntDType, Scalar, Shape, Slice, + TensorData, TensorMetadata, + ops::{FloatTensorOps, IntTensorOps}, + tensor::{BoolTensor, Device, FloatTensor, IntTensor}, +}; +use burn_std::{BoolDType, FloatDType}; + +use crate::{ + Candle, CandleDevice, CandleTensor, IntoDType, + element::{CandleElement, FloatCandleElement, IntCandleElement}, +}; + +use super::base::{cpu_random, expand, permute, sign, unfold}; + +impl IntTensorOps for Candle { + fn int_empty(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + super::base::empty(shape, device, dtype.into_dtype()) + } + + async fn int_into_data(tensor: IntTensor) -> Result { + super::base::into_data(tensor) + } + + fn int_from_data(data: TensorData, device: &Device) -> IntTensor { + match data.dtype { + DType::I64 => super::base::from_data::(data, device), + DType::U32 => super::base::from_data::(data, device), + DType::U8 => super::base::from_data::(data, device), + _ => unimplemented!("Unsupported dtype for `int_from_data`"), + } + } + + fn int_to_device(tensor: IntTensor, device: &Device) -> IntTensor { + super::base::to_device(tensor, device) + } + + fn int_reshape(tensor: IntTensor, shape: Shape) -> IntTensor { + super::base::reshape(tensor, shape) + } + + fn int_slice(tensor: IntTensor, slices: &[Slice]) -> IntTensor { + super::base::slice_with_steps(tensor, slices) + } + + fn int_slice_assign( + tensor: IntTensor, + slices: &[Slice], + value: IntTensor, + ) -> IntTensor { + super::base::slice_assign(tensor, slices, value) + } + + fn int_into_float(tensor: IntTensor, out_dtype: FloatDType) -> FloatTensor { + CandleTensor::new(tensor.tensor.to_dtype(out_dtype.into_dtype()).unwrap()) + } + + fn int_mask_where( + tensor: IntTensor, + mask: BoolTensor, + source: IntTensor, + ) -> IntTensor { + super::base::mask_where_broadcasted(tensor, mask, source) + } + + fn int_mask_fill( + tensor: IntTensor, + mask: BoolTensor, + value: Scalar, + ) -> IntTensor { + CandleTensor::new( + mask.tensor + .where_cond( + &super::candle_utils::fill_like::(value.elem(), &tensor.tensor), + &tensor.tensor, + ) + .unwrap(), + ) + } + + fn int_gather( + dim: usize, + tensor: IntTensor, + indices: IntTensor, + ) -> IntTensor { + let tensor = tensor.tensor.contiguous().unwrap(); + let indices = indices.tensor.contiguous().unwrap(); + CandleTensor::new(tensor.gather(&indices, dim).unwrap()) + } + + fn int_scatter_add( + dim: usize, + tensor: IntTensor, + indices: IntTensor, + value: IntTensor, + ) -> IntTensor { + CandleTensor::new( + tensor + .tensor + .scatter_add(&indices.tensor, &value.tensor, dim) + .unwrap(), + ) + } + + fn int_scatter_nd( + data: IntTensor, + indices: IntTensor, + values: IntTensor, + reduction: burn_backend::tensor::IndexingUpdateOp, + ) -> IntTensor { + super::base::scatter_nd(data, indices, values, reduction) + } + + fn int_gather_nd(data: IntTensor, indices: IntTensor) -> IntTensor { + super::base::gather_nd(data, indices) + } + + fn int_select( + tensor: IntTensor, + dim: usize, + indices: IntTensor, + ) -> IntTensor { + CandleTensor::new(tensor.tensor.index_select(&indices.tensor, dim).unwrap()) + } + + fn int_select_add( + tensor: IntTensor, + dim: usize, + indices: IntTensor, + value: IntTensor, + ) -> IntTensor { + CandleTensor::new( + tensor + .tensor + .index_add(&indices.tensor, &value.tensor, dim) + .unwrap(), + ) + } + + fn int_cat(tensors: Vec>, dim: usize) -> IntTensor { + super::base::cat(tensors, dim) + } + + fn int_equal( + lhs: IntTensor, + rhs: IntTensor, + _out_dtype: BoolDType, + ) -> BoolTensor { + let (lhs_broadcast, rhs_broadcast) = + super::candle_utils::broadcast_for_comparison(&lhs.tensor, &rhs.tensor).unwrap(); + CandleTensor::new(lhs_broadcast.eq(&rhs_broadcast).unwrap()) + } + + fn int_equal_elem( + lhs: IntTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> BoolTensor { + CandleTensor::new(lhs.tensor.eq(rhs.elem::()).unwrap()) + } + + fn int_greater( + lhs: IntTensor, + rhs: IntTensor, + _out_dtype: BoolDType, + ) -> BoolTensor { + let (lhs_broadcast, rhs_broadcast) = + super::candle_utils::broadcast_for_comparison(&lhs.tensor, &rhs.tensor).unwrap(); + CandleTensor::new(lhs_broadcast.gt(&rhs_broadcast).unwrap()) + } + + fn int_greater_elem( + lhs: IntTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> BoolTensor { + CandleTensor::new( + lhs.tensor + .gt(&super::candle_utils::fill_like::( + rhs.elem(), + &lhs.tensor, + )) + .unwrap(), + ) + } + + fn int_greater_equal( + lhs: IntTensor, + rhs: IntTensor, + _out_dtype: BoolDType, + ) -> BoolTensor { + let (lhs_broadcast, rhs_broadcast) = + super::candle_utils::broadcast_for_comparison(&lhs.tensor, &rhs.tensor).unwrap(); + CandleTensor::new(lhs_broadcast.ge(&rhs_broadcast).unwrap()) + } + + fn int_greater_equal_elem( + lhs: IntTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> BoolTensor { + CandleTensor::new( + lhs.tensor + .ge(&super::candle_utils::fill_like::( + rhs.elem(), + &lhs.tensor, + )) + .unwrap(), + ) + } + + fn int_lower( + lhs: IntTensor, + rhs: IntTensor, + _out_dtype: BoolDType, + ) -> BoolTensor { + let (lhs_broadcast, rhs_broadcast) = + super::candle_utils::broadcast_for_comparison(&lhs.tensor, &rhs.tensor).unwrap(); + CandleTensor::new(lhs_broadcast.lt(&rhs_broadcast).unwrap()) + } + + fn int_lower_elem( + lhs: IntTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> BoolTensor { + CandleTensor::new( + lhs.tensor + .lt(&super::candle_utils::fill_like::( + rhs.elem(), + &lhs.tensor, + )) + .unwrap(), + ) + } + + fn int_lower_equal( + lhs: IntTensor, + rhs: IntTensor, + _out_dtype: BoolDType, + ) -> BoolTensor { + let (lhs_broadcast, rhs_broadcast) = + super::candle_utils::broadcast_for_comparison(&lhs.tensor, &rhs.tensor).unwrap(); + CandleTensor::new(lhs_broadcast.le(&rhs_broadcast).unwrap()) + } + + fn int_lower_equal_elem( + lhs: IntTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> BoolTensor { + CandleTensor::new( + lhs.tensor + .le(&super::candle_utils::fill_like::( + rhs.elem(), + &lhs.tensor, + )) + .unwrap(), + ) + } + + fn int_add(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + CandleTensor::new(lhs.tensor.broadcast_add(&rhs.tensor).unwrap()) + } + + fn int_add_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + CandleTensor::new((lhs.tensor + rhs.elem::()).unwrap()) + } + + fn int_sub(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + CandleTensor::new(lhs.tensor.broadcast_sub(&rhs.tensor).unwrap()) + } + + fn int_sub_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + CandleTensor::new((lhs.tensor - rhs.elem::()).unwrap()) + } + + fn int_mul(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + CandleTensor::new(lhs.tensor.broadcast_mul(&rhs.tensor).unwrap()) + } + + fn int_mul_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + CandleTensor::new((lhs.tensor * rhs.elem::()).unwrap()) + } + + fn int_div(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + CandleTensor::new(lhs.tensor.broadcast_div(&rhs.tensor).unwrap()) + } + + fn int_div_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + // Candle implements scalar a/b as a * (1/b). With ints 1/b is rounded to 0 so we always obtain 0. + panic!("Not supported by Candle") + } + + fn int_remainder(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + CandleTensor::new( + (lhs.tensor.clone() + - lhs + .tensor + .broadcast_div(&rhs.tensor) + .unwrap() + .broadcast_mul(&rhs.tensor) + .unwrap()) + .unwrap(), + ) + } + + fn int_remainder_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + // Same problem as int_div_scalar. + panic!("Not supported by Candle") + } + + fn int_zeros(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + CandleTensor::new( + candle_core::Tensor::zeros( + shape.to_vec(), + dtype.into_dtype(), + &(device.clone()).into(), + ) + .unwrap(), + ) + } + + fn int_ones(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + CandleTensor::new( + candle_core::Tensor::ones(shape.to_vec(), dtype.into_dtype(), &(device.clone()).into()) + .unwrap(), + ) + } + + fn int_sum(tensor: IntTensor) -> IntTensor { + let sum = tensor.tensor.sum_all().unwrap().reshape((1,)).unwrap(); + CandleTensor::new(sum) + } + + fn int_sum_dim(tensor: IntTensor, dim: usize) -> IntTensor { + CandleTensor::new(tensor.tensor.sum_keepdim(dim).unwrap()) + } + + fn int_prod(tensor: IntTensor) -> IntTensor { + todo!( + "prod is not implemented for Candle IntTensor (see https://github.com/tracel-ai/burn/issues/1454)" + ) + } + + fn int_prod_dim(tensor: IntTensor, dim: usize) -> IntTensor { + todo!( + "prod_int is not implemented for Candle IntTensor (see https://github.com/tracel-ai/burn/issues/1454)" + ) + } + + fn int_mean_dim(tensor: IntTensor, dim: usize) -> IntTensor { + // Candle implements scalar a/b as a * (1/b). With ints 1/b is rounded to 0 so we always obtain 0. + panic!("Not supported by Candle") + } + + fn int_cumsum(tensor: IntTensor, dim: usize) -> IntTensor { + // Candle's cumsum doesn't support integer types, so we convert to float, + // compute cumsum, and convert back to int + let dtype = tensor.tensor.dtype(); + let tensor_float = tensor.tensor.to_dtype(candle_core::DType::F32).unwrap(); + let result_float = tensor_float.cumsum(dim).unwrap(); + CandleTensor::new(result_float.to_dtype(dtype).unwrap()) + } + + fn int_cumprod(tensor: IntTensor, dim: usize) -> IntTensor { + // Convert to float for computation, then convert back + let dtype = tensor.tensor.dtype(); + let tensor_float = tensor.tensor.to_dtype(candle_core::DType::F32).unwrap(); + + let result_float = super::utils::cumulative_with_op(&tensor_float, dim, |prev, curr| { + prev.broadcast_mul(curr) + }); + CandleTensor::new(result_float.to_dtype(dtype).unwrap()) + } + + fn int_cummin(tensor: IntTensor, dim: usize) -> IntTensor { + // Convert to float for computation, then convert back + let dtype = tensor.tensor.dtype(); + let tensor_float = tensor.tensor.to_dtype(candle_core::DType::F32).unwrap(); + + let result_float = super::utils::cumulative_with_op(&tensor_float, dim, |prev, curr| { + prev.broadcast_minimum(curr) + }); + CandleTensor::new(result_float.to_dtype(dtype).unwrap()) + } + + fn int_cummax(tensor: IntTensor, dim: usize) -> IntTensor { + let result = super::utils::cumulative_with_op(&tensor.tensor, dim, |prev, curr| { + prev.broadcast_maximum(curr) + }); + CandleTensor::new(result) + } + + fn int_argmax(tensor: IntTensor, dim: usize) -> IntTensor { + CandleTensor::new(tensor.tensor.argmax_keepdim(dim).unwrap()) + } + + fn int_argtopk(tensor: IntTensor, dim: usize, k: usize) -> IntTensor { + panic!("argtopk not implemented for candle backend") + } + + fn int_argmin(tensor: IntTensor, dim: usize) -> IntTensor { + CandleTensor::new(tensor.tensor.argmin_keepdim(dim).unwrap()) + } + + fn int_abs(tensor: IntTensor) -> IntTensor { + // Ugly type conversion here as Candle does not support unary ops on ints + match tensor.tensor.dtype() { + candle_core::DType::U8 | candle_core::DType::U32 => tensor, + candle_core::DType::I64 => CandleTensor::new( + tensor + .tensor + .to_dtype(candle_core::DType::F32) + .unwrap() + .abs() + .unwrap() + .to_dtype(candle_core::DType::I64) + .unwrap(), + ), + _ => unreachable!(), + } + } + + fn int_swap_dims(tensor: IntTensor, dim1: usize, dim2: usize) -> IntTensor { + super::base::swap_dims(tensor, dim1, dim2) + } + + fn int_random( + shape: Shape, + distribution: Distribution, + device: &Device, + dtype: IntDType, + ) -> IntTensor { + if let CandleDevice::Cpu = device { + let distribution = if distribution == Distribution::Default { + Distribution::Uniform(0.0, 255.0) + } else { + distribution + }; + // Use our own seed since candle doesn't support it on CPU + return Self::int_from_data(cpu_random(shape, distribution, dtype.into()), device); + } + + let shape = shape.to_vec(); + let device = &(device.clone()).into(); + match distribution { + Distribution::Default => CandleTensor::new( + candle_core::Tensor::rand(0f32, 255f32, shape, device) + .unwrap() + .to_dtype(dtype.into_dtype()) + .unwrap(), + ), + Distribution::Bernoulli(prob) => CandleTensor::new( + candle_core::Tensor::rand(0f32, 1f32, shape.clone(), device) + .unwrap() + .to_dtype(dtype.into_dtype()) + .unwrap() + .lt(&super::candle_utils::fill( + prob, + shape, + dtype.into_dtype(), + device, + )) + .unwrap() + .to_dtype(dtype.into_dtype()) + .unwrap(), + ), + Distribution::Uniform(from, to) => CandleTensor::new( + candle_core::Tensor::rand(from, to, shape, device) + .unwrap() + .to_dtype(dtype.into_dtype()) + .unwrap(), + ), + Distribution::Normal(mean, std) => CandleTensor::new( + candle_core::Tensor::randn(mean, std, shape, device) + .unwrap() + .to_dtype(dtype.into_dtype()) + .unwrap(), + ), + } + } + + fn int_permute(tensor: IntTensor, axes: &[usize]) -> IntTensor { + super::base::permute(tensor, axes) + } + + fn int_flip(tensor: IntTensor, axes: &[usize]) -> IntTensor { + super::base::flip(tensor, axes) + } + + fn int_expand(tensor: IntTensor, shape: Shape) -> IntTensor { + expand(tensor, shape) + } + + fn int_unfold( + tensor: IntTensor, + dim: usize, + size: usize, + step: usize, + ) -> IntTensor { + unfold(tensor, dim, size, step) + } + + fn int_sign(tensor: IntTensor) -> IntTensor { + sign(tensor) + } + fn bitwise_and(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + unimplemented!("bitwise_and is not implemented for Candle IntTensor"); + } + + fn bitwise_and_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + unimplemented!("bitwise_and_scalar is not implemented for Candle IntTensor"); + } + + fn bitwise_or(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + unimplemented!("bitwise_or is not implemented for Candle IntTensor"); + } + + fn bitwise_or_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + unimplemented!("bitwise_or_scalar is not implemented for Candle IntTensor"); + } + + fn bitwise_xor(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + unimplemented!("bitwise_xor is not implemented for Candle IntTensor"); + } + + fn bitwise_xor_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + unimplemented!("bitwise_xor_scalar is not implemented for Candle IntTensor"); + } + + fn bitwise_not(tensor: IntTensor) -> IntTensor { + unimplemented!("bitwise_not is not implemented for Candle IntTensor"); + } + + fn bitwise_left_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + unimplemented!("bitwise_left_shift is not implemented for Candle IntTensor"); + } + + fn bitwise_right_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + unimplemented!("bitwise_right_shift is not implemented for Candle IntTensor"); + } + + fn bitwise_left_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + unimplemented!("bitwise_left_shift_scalar is not implemented for Candle IntTensor"); + } + + fn bitwise_right_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + unimplemented!("bitwise_right_shift_scalar is not implemented for Candle IntTensor"); + } + + fn int_matmul(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + let int_dtype = lhs.dtype(); + let lhs = Self::int_into_float(lhs, FloatDType::F32); + let rhs = Self::int_into_float(rhs, FloatDType::F32); + + let out = Self::float_matmul(lhs, rhs); + Self::float_into_int(out, int_dtype.into()) + } + + fn int_cast(tensor: IntTensor, dtype: IntDType) -> IntTensor { + let dtype = dtype.into_dtype(); + + if tensor.tensor.dtype() == dtype { + tensor + } else { + CandleTensor::new(tensor.tensor.to_dtype(dtype).unwrap()) + } + } +} diff --git a/crates/burn-candle/src/ops/mod.rs b/crates/burn-candle/src/ops/mod.rs new file mode 100644 index 0000000..49de10c --- /dev/null +++ b/crates/burn-candle/src/ops/mod.rs @@ -0,0 +1,10 @@ +mod activation; +mod base; +mod bool_tensor; +mod candle_utils; +mod int_tensor; +mod module; +mod qtensor; +mod tensor; +mod transaction; +mod utils; diff --git a/crates/burn-candle/src/ops/module.rs b/crates/burn-candle/src/ops/module.rs new file mode 100644 index 0000000..d397cd4 --- /dev/null +++ b/crates/burn-candle/src/ops/module.rs @@ -0,0 +1,352 @@ +use burn_backend::{ + Shape, + ops::{ + ConvOptions, ConvTransposeOptions, DeformConv2dBackward, DeformConvOptions, + InterpolateMode, InterpolateOptions, MaxPool2dBackward, MaxPool2dWithIndices, ModuleOps, + UnfoldOptions, attention::attention_fallback, + }, + tensor::{FloatTensor, IntTensor}, +}; +use burn_std::IntDType; +use candle_core::ToUsize2; + +use crate::{ + Candle, CandleTensor, + element::{CandleElement, FloatCandleElement, IntCandleElement}, + ops::base::reshape, +}; + +impl ModuleOps for Candle { + fn conv1d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<1>, + ) -> FloatTensor { + let conv = x + .tensor + .conv1d( + &weight.tensor, + options.padding[0], + options.stride[0], + options.dilation[0], + options.groups, + ) + .unwrap(); + CandleTensor::new(match bias { + Some(bias) => conv + .broadcast_add(&bias.tensor.unsqueeze(1).unwrap()) + .unwrap(), + None => conv, + }) + } + + fn conv2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<2>, + ) -> FloatTensor { + assert!( + options.dilation[0] == options.dilation[1] + && options.padding[0] == options.padding[1] + && options.stride[0] == options.stride[1], + "Candle does not support per dimension options in convolutions" + ); + let conv = x + .tensor + .conv2d( + &weight.tensor, + options.padding[0], + options.stride[0], + options.dilation[0], + options.groups, + ) + .unwrap(); + CandleTensor::new(match bias { + Some(bias) => conv + .broadcast_add( + &bias + .tensor + .unsqueeze(0) + .unwrap() + .unsqueeze(2) + .unwrap() + .unsqueeze(3) + .unwrap(), + ) + .unwrap(), + None => conv, + }) + } + + fn deform_conv2d( + x: FloatTensor, + offset: FloatTensor, + weight: FloatTensor, + mask: Option>, + bias: Option>, + options: DeformConvOptions<2>, + ) -> FloatTensor { + unimplemented!("Candle does not support deformable convolutions") + } + + fn deform_conv2d_backward( + x: FloatTensor, + offset: FloatTensor, + weight: FloatTensor, + mask: Option>, + bias: Option>, + output_grad: FloatTensor, + options: DeformConvOptions<2>, + ) -> DeformConv2dBackward { + unimplemented!("Candle does not support deformable convolutions") + } + + fn conv3d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<3>, + ) -> FloatTensor { + panic!("Candle does not support 3D convolutions"); + } + + fn conv_transpose1d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<1>, + ) -> FloatTensor { + let conv_transpose = x + .tensor + .conv_transpose1d( + &weight.tensor, + options.padding[0], + options.padding_out[0], + options.stride[0], + options.dilation[0], + options.groups, + ) + .unwrap(); + CandleTensor::new(match bias { + Some(bias) => conv_transpose + .broadcast_add(&bias.tensor.unsqueeze(0).unwrap().unsqueeze(2).unwrap()) + .unwrap(), + None => conv_transpose, + }) + } + + fn conv_transpose2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<2>, + ) -> FloatTensor { + assert!( + options.dilation[0] == options.dilation[1] + && options.padding[0] == options.padding[1] + && options.padding_out[0] == options.padding_out[1] + && options.stride[0] == options.stride[1], + "Candle does not support per dimension options in transposed convolutions" + ); + assert!( + options.groups == 1, + "Candle does not support groups in transposed convolutions" + ); + let conv_transpose = x + .tensor + .conv_transpose2d( + &weight.tensor, + options.padding[0], + options.padding_out[0], + options.stride[0], + options.dilation[0], + ) + .unwrap(); + CandleTensor::new(match bias { + Some(bias) => conv_transpose + .broadcast_add( + &bias + .tensor + .unsqueeze(0) + .unwrap() + .unsqueeze(2) + .unwrap() + .unsqueeze(3) + .unwrap(), + ) + .unwrap(), + None => conv_transpose, + }) + } + + fn conv_transpose3d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<3>, + ) -> FloatTensor { + panic!("Candle does not support 3D transposed convolutions"); + } + + fn avg_pool2d( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + assert!( + padding[0] == 0 && padding[1] == 0, + "Candle does not support padding in pooling" + ); + assert!( + count_include_pad, + "Candle does not support excluding pad count in pooling" + ); + assert!(!ceil_mode, "Candle does not support ceil_mode in pooling"); + CandleTensor::new( + x.tensor + .avg_pool2d_with_stride((kernel_size[0], kernel_size[1]), (stride[0], stride[1])) + .unwrap(), + ) + } + + fn avg_pool2d_backward( + x: FloatTensor, + grad: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + _ceil_mode: bool, + ) -> FloatTensor { + panic!("avg_pool2d_backward is not supported by Candle") + } + + fn max_pool2d( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + ) -> FloatTensor { + assert!( + padding[0] == 0 && padding[1] == 0, + "Candle does not support padding in pooling" + ); + assert!( + dilation[0] == 1 && dilation[1] == 1, + "Candle does not support dilation in pooling" + ); + assert!(!ceil_mode, "Candle does not support ceil_mode in pooling"); + CandleTensor::new( + x.tensor + .max_pool2d_with_stride((kernel_size[0], kernel_size[1]), (stride[0], stride[1])) + .unwrap(), + ) + } + + fn max_pool2d_with_indices( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + _ceil_mode: bool, + indices_dtype: IntDType, + ) -> MaxPool2dWithIndices { + panic!("max_pool2d_with_indices is not supported by Candle") + } + + fn max_pool2d_with_indices_backward( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + _ceil_mode: bool, + output_grad: FloatTensor, + indices: IntTensor, + ) -> MaxPool2dBackward { + panic!("max_pool2d_with_indices_backward is not supported by Candle") + } + + fn adaptive_avg_pool2d(x: FloatTensor, output_size: [usize; 2]) -> FloatTensor { + panic!("adaptive_avg_pool2 is not supported by Candle") + } + + fn adaptive_avg_pool2d_backward( + x: FloatTensor, + grad: FloatTensor, + ) -> FloatTensor { + panic!("adaptive_avg_pool2d_backward is not supported by Candle") + } + + fn interpolate( + x: FloatTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> FloatTensor { + let tensor = match options.mode { + InterpolateMode::Nearest => x + .tensor + .upsample_nearest2d(output_size[0], output_size[1]) + .unwrap(), + InterpolateMode::NearestExact => { + panic!("nearest exact interpolation is not implemented in Candle backend") + } + InterpolateMode::Bilinear => { + panic!("bilinear interpolation is not supported by Candle") + } + InterpolateMode::Bicubic => { + panic!("bicubic interpolation is not supported by Candle") + } + InterpolateMode::Lanczos3 => { + panic!("lanczos3 interpolation is not supported by Candle") + } + }; + + CandleTensor::new(tensor) + } + + fn interpolate_backward( + x: FloatTensor, + grad: FloatTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> FloatTensor { + panic!("interpolate_backward is not supported by Candle") + } + + fn attention( + query: FloatTensor, + key: FloatTensor, + value: FloatTensor, + mask: Option>, + attn_bias: Option>, + options: burn_backend::ops::AttentionModuleOptions, + ) -> FloatTensor { + attention_fallback::(query, key, value, mask, attn_bias, options) + } + + fn rfft( + signal: FloatTensor, + dim: usize, + _n: Option, + ) -> (FloatTensor, FloatTensor) { + todo!("rfft is unsupported in Candle") + } + + fn irfft( + spectrum_re: FloatTensor, + spectrum_im: FloatTensor, + dim: usize, + _n: Option, + ) -> FloatTensor { + todo!("irfft is unsupported in Candle") + } +} diff --git a/crates/burn-candle/src/ops/qtensor.rs b/crates/burn-candle/src/ops/qtensor.rs new file mode 100644 index 0000000..1ee2e60 --- /dev/null +++ b/crates/burn-candle/src/ops/qtensor.rs @@ -0,0 +1,60 @@ +use burn_backend::{ + Backend, DType, ExecutionError, FloatDType, Shape, TensorData, + ops::QTensorOps, + quantization::{QuantScheme, QuantizationParametersPrimitive}, + tensor::{Device, FloatTensor, QuantizedTensor}, +}; + +use crate::{ + Candle, + element::{FloatCandleElement, IntCandleElement}, +}; + +impl QTensorOps for Candle { + fn q_from_data(data: TensorData, device: &Device) -> QuantizedTensor { + unimplemented!() + } + + fn quantize( + _tensor: FloatTensor, + _scheme: &QuantScheme, + _qparams: QuantizationParametersPrimitive, + ) -> QuantizedTensor { + unimplemented!() + } + + fn dequantize(_tensor: QuantizedTensor, _dtype: FloatDType) -> FloatTensor { + unimplemented!() + } + + fn q_to_device( + _tensor: QuantizedTensor, + _device: &Device, + ) -> QuantizedTensor { + unimplemented!() + } + + fn q_reshape(_tensor: QuantizedTensor, _shape: Shape) -> QuantizedTensor { + unimplemented!() + } + + async fn q_into_data(tensor: QuantizedTensor) -> Result { + unimplemented!() + } + + fn q_swap_dims( + _tensor: QuantizedTensor, + _dim1: usize, + _dim2: usize, + ) -> QuantizedTensor { + unimplemented!() + } + + fn q_permute(_tensor: QuantizedTensor, _axes: &[usize]) -> QuantizedTensor { + unimplemented!() + } + + fn q_flip(_tensor: QuantizedTensor, _axes: &[usize]) -> QuantizedTensor { + unimplemented!() + } +} diff --git a/crates/burn-candle/src/ops/tensor.rs b/crates/burn-candle/src/ops/tensor.rs new file mode 100644 index 0000000..ad622ac --- /dev/null +++ b/crates/burn-candle/src/ops/tensor.rs @@ -0,0 +1,679 @@ +use std::borrow::Borrow; + +use burn_backend::{ + DType, Distribution, ElementConversion, ExecutionError, FloatDType, Scalar, Shape, Slice, + TensorData, bf16, f16, + ops::FloatTensorOps, + tensor::{BoolTensor, Device, FloatTensor, IntTensor}, +}; +use burn_std::{BoolDType, IntDType}; +use candle_core::{Tensor, backend::BackendStorage, shape}; + +use crate::{ + Candle, CandleDevice, CandleTensor, IntoDType, + element::{CandleElement, FloatCandleElement, IntCandleElement}, +}; + +use super::base::{cpu_random, expand, permute, sign, unfold}; + +impl FloatTensorOps for Candle { + fn float_from_data(data: TensorData, device: &Device) -> CandleTensor { + match data.dtype { + DType::F64 => super::base::from_data::(data, device), + DType::F32 => super::base::from_data::(data, device), + DType::F16 => super::base::from_data::(data, device), + DType::BF16 => super::base::from_data::(data, device), + _ => unimplemented!("Unsupported dtype for `float_from_data`"), + } + } + + fn float_random( + shape: Shape, + distribution: Distribution, + device: &Device, + dtype: FloatDType, + ) -> FloatTensor { + if let CandleDevice::Cpu = device { + // Use our own seed since candle doesn't support it on CPU + return Self::float_from_data(cpu_random(shape, distribution, dtype.into()), device); + } + + let shape = shape.to_vec(); + let device = &(device.clone()).into(); + match distribution { + Distribution::Default => CandleTensor::new( + candle_core::Tensor::rand(0f32, 1f32, shape, device) + .unwrap() + .to_dtype(dtype.into_dtype()) + .unwrap(), + ), + Distribution::Bernoulli(prob) => CandleTensor::new( + candle_core::Tensor::rand(0f32, 1f32, shape.clone(), device) + .unwrap() + .to_dtype(dtype.into_dtype()) + .unwrap() + .lt(&super::candle_utils::fill( + prob, + shape, + dtype.into_dtype(), + device, + )) + .unwrap() + .to_dtype(dtype.into_dtype()) + .unwrap(), + ), + Distribution::Uniform(from, to) => CandleTensor::new( + candle_core::Tensor::rand(from, to, shape, device) + .unwrap() + .to_dtype(dtype.into_dtype()) + .unwrap(), + ), + Distribution::Normal(mean, std) => CandleTensor::new( + candle_core::Tensor::randn(mean, std, shape, device) + .unwrap() + .to_dtype(dtype.into_dtype()) + .unwrap(), + ), + } + } + + async fn float_into_data(tensor: CandleTensor) -> Result { + super::base::into_data(tensor) + } + + fn float_to_device(tensor: CandleTensor, device: &Device) -> CandleTensor { + super::base::to_device(tensor, device) + } + + fn float_into_int(tensor: CandleTensor, out_dtype: IntDType) -> IntTensor { + CandleTensor::new(tensor.tensor.to_dtype(out_dtype.into_dtype()).unwrap()) + } + + fn float_empty(shape: Shape, device: &Device, dtype: FloatDType) -> FloatTensor { + super::base::empty(shape, device, dtype.into_dtype()) + } + + fn float_add(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + CandleTensor::new(lhs.tensor.broadcast_add(&rhs.tensor).unwrap()) + } + + fn float_add_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + CandleTensor::new((lhs.tensor + rhs.elem::()).unwrap()) + } + + fn float_sub(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + CandleTensor::new(lhs.tensor.broadcast_sub(&rhs.tensor).unwrap()) + } + + fn float_sub_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + CandleTensor::new((lhs.tensor - rhs.elem::()).unwrap()) + } + + fn float_mul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + CandleTensor::new(lhs.tensor.broadcast_mul(&rhs.tensor).unwrap()) + } + + fn float_mul_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + CandleTensor::new((lhs.tensor * rhs.elem::()).unwrap()) + } + + fn float_div(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + CandleTensor::new(lhs.tensor.broadcast_div(&rhs.tensor).unwrap()) + } + + fn float_div_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + CandleTensor::new((lhs.tensor / rhs.elem::()).unwrap()) + } + + fn float_remainder(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + CandleTensor::new( + (lhs.tensor.clone() + - lhs + .tensor + .broadcast_div(&rhs.tensor) + .unwrap() + .floor() + .unwrap() + .broadcast_mul(&rhs.tensor) + .unwrap()) + .unwrap(), + ) + } + + fn float_remainder_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + // In PyTorch, remainder can also be defined as torch.remainder(a, b) == a - a.div(b, rounding_mode="floor") * b + let rhs_val = rhs.elem::(); + let division_result = (lhs.tensor.clone() / rhs_val).unwrap().floor().unwrap(); + let product = division_result * rhs_val; + + CandleTensor::new((lhs.tensor - product).unwrap()) + } + + fn float_matmul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + let lhs_contiguous = if !lhs.tensor.is_contiguous() { + lhs.tensor.contiguous().unwrap() + } else { + lhs.tensor + }; + let rhs_contiguous = if !rhs.tensor.is_contiguous() { + rhs.tensor.contiguous().unwrap() + } else { + rhs.tensor + }; + CandleTensor::new(lhs_contiguous.broadcast_matmul(&rhs_contiguous).unwrap()) + } + + fn float_cross( + lhs: FloatTensor, + rhs: FloatTensor, + dim: usize, + ) -> FloatTensor { + super::base::cross(lhs, rhs, dim) + } + + fn float_swap_dims(tensor: FloatTensor, dim1: usize, dim2: usize) -> FloatTensor { + super::base::swap_dims(tensor, dim1, dim2) + } + + fn float_reshape(tensor: FloatTensor, shape: Shape) -> FloatTensor { + super::base::reshape(tensor, shape) + } + + fn float_gather( + dim: usize, + tensor: FloatTensor, + indices: IntTensor, + ) -> FloatTensor { + let tensor = tensor.tensor.contiguous().unwrap(); + let indices = indices.tensor.contiguous().unwrap(); + CandleTensor::new(tensor.gather(&indices, dim).unwrap()) + } + + fn float_scatter_add( + dim: usize, + tensor: FloatTensor, + indices: IntTensor, + value: FloatTensor, + ) -> FloatTensor { + CandleTensor::new( + tensor + .tensor + .scatter_add(&indices.tensor, &value.tensor, dim) + .unwrap(), + ) + } + + fn float_scatter_nd( + data: FloatTensor, + indices: IntTensor, + values: FloatTensor, + reduction: burn_backend::tensor::IndexingUpdateOp, + ) -> FloatTensor { + super::base::scatter_nd(data, indices, values, reduction) + } + + fn float_gather_nd(data: FloatTensor, indices: IntTensor) -> FloatTensor { + super::base::gather_nd(data, indices) + } + + fn float_select( + tensor: FloatTensor, + dim: usize, + indices: IntTensor, + ) -> FloatTensor { + CandleTensor::new(tensor.tensor.index_select(&indices.tensor, dim).unwrap()) + } + + fn float_select_add( + tensor: FloatTensor, + dim: usize, + indices: IntTensor, + value: FloatTensor, + ) -> FloatTensor { + CandleTensor::new( + tensor + .tensor + .index_add(&indices.tensor, &value.tensor, dim) + .unwrap(), + ) + } + + fn float_slice(tensor: FloatTensor, slices: &[Slice]) -> FloatTensor { + super::base::slice_with_steps(tensor, slices) + } + + fn float_slice_assign( + tensor: FloatTensor, + slices: &[Slice], + value: FloatTensor, + ) -> FloatTensor { + super::base::slice_assign(tensor, slices, value) + } + + fn float_mask_where( + tensor: FloatTensor, + mask: BoolTensor, + value: FloatTensor, + ) -> FloatTensor { + super::base::mask_where_broadcasted(tensor, mask, value) + } + + fn float_mask_fill( + tensor: FloatTensor, + mask: BoolTensor, + value: Scalar, + ) -> FloatTensor { + let value = super::candle_utils::fill_like::(value.elem(), &tensor.tensor); + super::base::mask_where_broadcasted(tensor, mask, CandleTensor::new(value)) + } + + fn float_equal( + lhs: FloatTensor, + rhs: FloatTensor, + _out_dtype: BoolDType, + ) -> BoolTensor { + let (lhs_broadcast, rhs_broadcast) = + super::candle_utils::broadcast_for_comparison(&lhs.tensor, &rhs.tensor).unwrap(); + CandleTensor::new(lhs_broadcast.eq(&rhs_broadcast).unwrap()) + } + + fn float_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> BoolTensor { + CandleTensor::new(lhs.tensor.eq(rhs.elem::()).unwrap()) + } + + fn float_greater( + lhs: FloatTensor, + rhs: FloatTensor, + _out_dtype: BoolDType, + ) -> BoolTensor { + let (lhs_broadcast, rhs_broadcast) = + super::candle_utils::broadcast_for_comparison(&lhs.tensor, &rhs.tensor).unwrap(); + CandleTensor::new(lhs_broadcast.gt(&rhs_broadcast).unwrap()) + } + + fn float_greater_elem( + lhs: FloatTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> BoolTensor { + CandleTensor::new( + lhs.tensor + .gt(&super::candle_utils::fill_like::( + rhs.elem(), + &lhs.tensor, + )) + .unwrap(), + ) + } + + fn float_greater_equal( + lhs: FloatTensor, + rhs: FloatTensor, + _out_dtype: BoolDType, + ) -> BoolTensor { + let (lhs_broadcast, rhs_broadcast) = + super::candle_utils::broadcast_for_comparison(&lhs.tensor, &rhs.tensor).unwrap(); + CandleTensor::new(lhs_broadcast.ge(&rhs_broadcast).unwrap()) + } + + fn float_greater_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> BoolTensor { + CandleTensor::new( + lhs.tensor + .ge(&super::candle_utils::fill_like::( + rhs.elem(), + &lhs.tensor, + )) + .unwrap(), + ) + } + + fn float_lower( + lhs: FloatTensor, + rhs: FloatTensor, + _out_dtype: BoolDType, + ) -> BoolTensor { + let (lhs_broadcast, rhs_broadcast) = + super::candle_utils::broadcast_for_comparison(&lhs.tensor, &rhs.tensor).unwrap(); + CandleTensor::new(lhs_broadcast.lt(&rhs_broadcast).unwrap()) + } + + fn float_lower_elem( + lhs: FloatTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> BoolTensor { + CandleTensor::new( + lhs.tensor + .lt(&super::candle_utils::fill_like::( + rhs.elem(), + &lhs.tensor, + )) + .unwrap(), + ) + } + + fn float_lower_equal( + lhs: FloatTensor, + rhs: FloatTensor, + _out_dtype: BoolDType, + ) -> BoolTensor { + let (lhs_broadcast, rhs_broadcast) = + super::candle_utils::broadcast_for_comparison(&lhs.tensor, &rhs.tensor).unwrap(); + CandleTensor::new(lhs_broadcast.le(&rhs_broadcast).unwrap()) + } + + fn float_lower_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> BoolTensor { + CandleTensor::new( + lhs.tensor + .le(&super::candle_utils::fill_like::( + rhs.elem(), + &lhs.tensor, + )) + .unwrap(), + ) + } + + fn float_sum(tensor: FloatTensor) -> FloatTensor { + let sum = tensor.tensor.sum_all().unwrap().reshape((1,)).unwrap(); + CandleTensor::new(sum) + } + + fn float_sum_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + CandleTensor::new(tensor.tensor.sum_keepdim(dim).unwrap()) + } + + fn float_mean_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + CandleTensor::new(tensor.tensor.mean_keepdim(dim).unwrap()) + } + + fn float_cumsum(tensor: FloatTensor, dim: usize) -> FloatTensor { + CandleTensor::new(tensor.tensor.cumsum(dim).unwrap()) + } + + fn float_cumprod(tensor: FloatTensor, dim: usize) -> FloatTensor { + let result = super::utils::cumulative_with_op(&tensor.tensor, dim, |prev, curr| { + prev.broadcast_mul(curr) + }); + CandleTensor::new(result) + } + + fn float_cummin(tensor: FloatTensor, dim: usize) -> FloatTensor { + let result = super::utils::cumulative_with_op(&tensor.tensor, dim, |prev, curr| { + prev.broadcast_minimum(curr) + }); + CandleTensor::new(result) + } + + fn float_cummax(tensor: FloatTensor, dim: usize) -> FloatTensor { + let result = super::utils::cumulative_with_op(&tensor.tensor, dim, |prev, curr| { + prev.broadcast_maximum(curr) + }); + CandleTensor::new(result) + } + + fn float_exp(tensor: FloatTensor) -> FloatTensor { + CandleTensor::new(tensor.tensor.exp().unwrap()) + } + + fn float_log(tensor: FloatTensor) -> FloatTensor { + CandleTensor::new(tensor.tensor.log().unwrap()) + } + + fn float_log1p(tensor: FloatTensor) -> FloatTensor { + CandleTensor::new((tensor.tensor + 1.).unwrap().log().unwrap()) + } + + fn float_powf_scalar_impl(tensor: FloatTensor, value: Scalar) -> FloatTensor { + CandleTensor::new(tensor.tensor.powf(value.elem()).unwrap()) + } + + fn float_sqrt(tensor: FloatTensor) -> FloatTensor { + CandleTensor::new(tensor.tensor.sqrt().unwrap()) + } + + fn float_abs(tensor: FloatTensor) -> FloatTensor { + CandleTensor::new(tensor.tensor.abs().unwrap()) + } + + fn float_cos(tensor: FloatTensor) -> FloatTensor { + CandleTensor::new(tensor.tensor.cos().unwrap()) + } + + fn float_cosh(tensor: FloatTensor) -> FloatTensor { + // cosh(x) = (e^x + e^(-x)) / 2 + let exp_x = tensor.tensor.exp().unwrap(); + CandleTensor::new(((exp_x.clone() + exp_x.recip().unwrap()).unwrap() / 2.0).unwrap()) + } + + fn float_sin(tensor: FloatTensor) -> FloatTensor { + CandleTensor::new(tensor.tensor.sin().unwrap()) + } + + fn float_sinh(tensor: FloatTensor) -> FloatTensor { + // sinh(x) = (e^x - e^(-x)) / 2 + let exp_x = tensor.tensor.exp().unwrap(); + CandleTensor::new(((exp_x.clone() - exp_x.recip().unwrap()).unwrap() / 2.0).unwrap()) + } + + fn float_tan(tensor: FloatTensor) -> FloatTensor { + CandleTensor::new((tensor.tensor.sin().unwrap() / tensor.tensor.cos().unwrap()).unwrap()) + } + + fn float_tanh(tensor: FloatTensor) -> FloatTensor { + CandleTensor::new(tensor.tensor.tanh().unwrap()) + } + + fn float_acos(tensor: FloatTensor) -> FloatTensor { + // acos(x) = PI/2 - asin(x) + let neg_asin_x = Self::float_neg(Self::float_asin(tensor)); + Self::float_add_scalar(neg_asin_x, core::f64::consts::FRAC_PI_2.into()) + } + + fn float_acosh(tensor: FloatTensor) -> FloatTensor { + // acosh(x) = ln(x + sqrt(x^2 - 1)) + let x_squared = Self::float_powi_scalar(tensor.clone(), 2.into()); + let x_sq_minus_one = Self::float_sub_scalar(x_squared, 1f64.into()); + let sqrt_term = Self::float_sqrt(x_sq_minus_one); + Self::float_log(Self::float_add(tensor, sqrt_term)) + } + + fn float_asin(tensor: FloatTensor) -> FloatTensor { + // asin(x) = atan(x / sqrt(1 - x^2)) + let x_squared = Self::float_powi_scalar(tensor.clone(), 2.into()); + let one_minus_x_sq = Self::float_add_scalar(Self::float_neg(x_squared), 1f64.into()); + let sqrt_term = Self::float_sqrt(one_minus_x_sq); + Self::float_atan(Self::float_div(tensor, sqrt_term)) + } + + fn float_asinh(tensor: FloatTensor) -> FloatTensor { + // asinh(x) = ln(x + sqrt(x^2 + 1)) + let x_squared = Self::float_powi_scalar(tensor.clone(), 2.into()); + let x_sq_plus_one = Self::float_add_scalar(x_squared, 1f64.into()); + let sqrt_term = Self::float_sqrt(x_sq_plus_one); + Self::float_log(Self::float_add(tensor, sqrt_term)) + } + + fn float_atan(tensor: FloatTensor) -> FloatTensor { + // atan(x) = asin(x / sqrt(1 + x^2)) + let x_squared = Self::float_powi_scalar(tensor.clone(), 2.into()); + let one_plus_x_sq = Self::float_add_scalar(x_squared, 1f64.into()); + let sqrt_term = Self::float_sqrt(one_plus_x_sq); + Self::float_asin(Self::float_div(tensor, sqrt_term)) + } + + fn float_atanh(tensor: FloatTensor) -> FloatTensor { + // atanh(x) = ln((1 + x) / (1 - x)) / 2 + let num = (1.0 + tensor.tensor.clone()).unwrap(); + let denom = (1.0 - tensor.tensor).unwrap(); + CandleTensor::new(((num / denom).unwrap().log().unwrap() / 2.0).unwrap()) + } + + fn float_atan2(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + // atan2(y, x) = 2 * atan(y / (sqrt(x^2 + y^2) + x)) + let x_squared = Self::float_powi_scalar(rhs.clone(), 2.into()); + let y_squared = Self::float_powi_scalar(lhs.clone(), 2.into()); + let r = Self::float_sqrt(Self::float_add(x_squared, y_squared)); + let ratio = Self::float_div(lhs, Self::float_add(r, rhs)); + Self::float_mul_scalar(Self::float_atan(ratio), 2f64.into()) + } + + fn float_round(tensor: FloatTensor) -> FloatTensor { + let inner = |tensor: FloatTensor| -> candle_core::Result> { + // implements round_to_even for consistent behavior vs libtorch + // https://github.com/pytorch/pytorch/blob/main/torch/csrc/jit/runtime/register_ops_utils.h#L65-L67 + + let floor_a = tensor.tensor.floor()?; + let frac_part = tensor.tensor.sub(&floor_a)?; + + let half = (candle_core::Tensor::ones_like(&tensor.tensor)? * 0.5)?; + let mask_half = frac_part.eq(&half)?; + let half_tensor = tensor.tensor.mul(&half)?; + let rounded_half = half_tensor.round()?; + let doubled = + rounded_half.mul(&(candle_core::Tensor::ones_like(&tensor.tensor)? * 2.0)?)?; + let standard_round = tensor.tensor.round()?; + Ok(CandleTensor::new( + mask_half.where_cond(&doubled, &standard_round)?, + )) + }; + inner(tensor).unwrap() + } + + fn float_floor(tensor: FloatTensor) -> FloatTensor { + CandleTensor::new(tensor.tensor.floor().unwrap()) + } + + fn float_ceil(tensor: FloatTensor) -> FloatTensor { + CandleTensor::new(tensor.tensor.ceil().unwrap()) + } + + fn float_trunc(tensor: FloatTensor) -> FloatTensor { + // truncate(x) = ⌊x⌋ if x ≥ 0, and ⌈x⌉ if x < 0 + // This preserves the sign of zero and handles all special cases correctly + let is_negative = tensor.tensor.lt(0.0).unwrap(); + let floored = tensor.tensor.floor().unwrap(); + let ceiled = tensor.tensor.ceil().unwrap(); + CandleTensor::new(is_negative.where_cond(&ceiled, &floored).unwrap()) + } + + fn float_erf(tensor: FloatTensor) -> FloatTensor { + CandleTensor::new(tensor.tensor.erf().unwrap()) + } + + fn float_cat(tensors: Vec>, dim: usize) -> FloatTensor { + super::base::cat(tensors, dim) + } + + fn float_argmax(tensor: FloatTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + CandleTensor::new( + tensor + .tensor + .argmax_keepdim(dim) + .unwrap() + .to_dtype(out_dtype.into_dtype()) + .unwrap(), + ) + } + + fn float_argtopk( + tensor: FloatTensor, + dim: usize, + k: usize, + out_dtype: IntDType, + ) -> IntTensor { + panic!("argtopk not implemented for candle backend") + } + + fn float_argmin(tensor: FloatTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + CandleTensor::new( + tensor + .tensor + .argmin_keepdim(dim) + .unwrap() + .to_dtype(out_dtype.into_dtype()) + .unwrap(), + ) + } + + fn float_clamp_max(tensor: FloatTensor, max: Scalar) -> FloatTensor { + CandleTensor::new(tensor.tensor.minimum(max.elem::()).unwrap()) + } + + fn float_clamp_min(tensor: FloatTensor, min: Scalar) -> FloatTensor { + CandleTensor::new(tensor.tensor.maximum(min.elem::()).unwrap()) + } + + fn float_clamp(tensor: FloatTensor, min: Scalar, max: Scalar) -> FloatTensor { + CandleTensor::new( + tensor + .tensor + .clamp(min.elem::(), max.elem::()) + .unwrap(), + ) + } + + fn float_recip(tensor: FloatTensor) -> FloatTensor { + CandleTensor::new(tensor.tensor.recip().unwrap()) + } + + fn float_powf(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + //broadcast_pow is in main but not yet published + //note: probably replace once pow once 0.3.3 is out + //see: https://github.com/huggingface/candle/pull/1583/files#diff-6319fa1e16dadc4c7b4e25698139703d93b70f30a1f8e2ac0999978e39efaa81R2594 + + CandleTensor::new( + rhs.tensor + .broadcast_mul(&lhs.tensor.log().unwrap()) + .unwrap() + .exp() + .unwrap(), + ) + } + + fn float_permute(tensor: FloatTensor, axes: &[usize]) -> FloatTensor { + super::base::permute(tensor, axes) + } + + fn float_flip(tensor: FloatTensor, axes: &[usize]) -> FloatTensor { + super::base::flip(tensor, axes) + } + + fn float_expand(tensor: FloatTensor, shape: Shape) -> FloatTensor { + expand(tensor, shape) + } + + fn float_unfold( + tensor: FloatTensor, + dim: usize, + size: usize, + step: usize, + ) -> FloatTensor { + unfold(tensor, dim, size, step) + } + + fn float_sign(tensor: FloatTensor) -> FloatTensor { + sign(tensor) + } + + fn float_cast(tensor: FloatTensor, dtype: FloatDType) -> FloatTensor { + let dtype = dtype.into_dtype(); + + if tensor.tensor.dtype() == dtype { + tensor + } else { + CandleTensor::new(tensor.tensor.to_dtype(dtype).unwrap()) + } + } +} diff --git a/crates/burn-candle/src/ops/transaction.rs b/crates/burn-candle/src/ops/transaction.rs new file mode 100644 index 0000000..708d4df --- /dev/null +++ b/crates/burn-candle/src/ops/transaction.rs @@ -0,0 +1,15 @@ +use burn_backend::{ + Backend, + distributed::DistributedOps, + ops::{TransactionOps, TransactionPrimitive}, +}; + +use crate::{ + Candle, + element::{FloatCandleElement, IntCandleElement}, +}; + +impl TransactionOps for Candle {} + +// DistributedOps has default implementations; Candle does not support collective operations. +impl DistributedOps for Candle {} diff --git a/crates/burn-candle/src/ops/utils.rs b/crates/burn-candle/src/ops/utils.rs new file mode 100644 index 0000000..7686018 --- /dev/null +++ b/crates/burn-candle/src/ops/utils.rs @@ -0,0 +1,29 @@ +/// Helper function for cumulative operations in Candle backend +/// +/// This function reduces code duplication for cumulative operations (cumprod, cummin, cummax) +/// which all follow the same pattern of slicing, applying an operation, and concatenating. +/// +/// # Arguments +/// +/// * `tensor` - The input tensor +/// * `dim` - The dimension along which to apply the cumulative operation +/// * `op` - A closure that takes two tensor references and produces a result tensor +pub fn cumulative_with_op(tensor: &candle_core::Tensor, dim: usize, op: F) -> candle_core::Tensor +where + F: Fn(&candle_core::Tensor, &candle_core::Tensor) -> candle_core::Result, +{ + let dim_size = tensor.dims()[dim]; + let mut slices = Vec::with_capacity(dim_size); + + // First slice is the initial value + slices.push(tensor.narrow(dim, 0, 1).unwrap()); + + // Apply cumulative operation + for i in 1..dim_size { + let curr = tensor.narrow(dim, i, 1).unwrap(); + let result = op(&slices[i - 1], &curr).unwrap(); + slices.push(result); + } + + candle_core::Tensor::cat(&slices, dim).unwrap() +} diff --git a/crates/burn-candle/src/tensor.rs b/crates/burn-candle/src/tensor.rs new file mode 100644 index 0000000..4de5b92 --- /dev/null +++ b/crates/burn-candle/src/tensor.rs @@ -0,0 +1,120 @@ +use burn_backend::{DType, FloatDType, IntDType, Shape, quantization::QuantScheme}; +use burn_backend::{Element, TensorData, TensorMetadata}; +use burn_std::BoolStore; + +use crate::{CandleDevice, element::CandleElement}; + +/// A tensor that uses the candle backend. +#[derive(Debug, Clone)] +pub struct CandleTensor { + pub(crate) tensor: candle_core::Tensor, +} + +impl TensorMetadata for CandleTensor { + type Device = CandleDevice; + fn dtype(&self) -> DType { + match self.tensor.dtype() { + candle_core::DType::U8 => DType::U8, + candle_core::DType::U32 => DType::U32, + candle_core::DType::I64 => DType::I64, + candle_core::DType::BF16 => DType::BF16, + candle_core::DType::F16 => DType::F16, + candle_core::DType::F32 => DType::F32, + candle_core::DType::F64 => DType::F64, + candle_core::DType::I16 => DType::I16, + candle_core::DType::I32 => DType::I32, + other => todo!("{other:?} not yet supported"), + } + } + + fn shape(&self) -> Shape { + Shape::from(self.tensor.dims().to_vec()) + } + + fn rank(&self) -> usize { + self.tensor.dims().len() + } + + fn device(&self) -> CandleDevice { + self.tensor.device().clone().into() + } + + fn can_mut(&self) -> bool { + // Candle tensors share storage behind an `Arc` with no public + // uniqueness check, so in-place mutation can never be assumed safe. + false + } +} + +impl CandleTensor { + /// Create a new tensor. + pub fn new(tensor: candle_core::Tensor) -> Self { + Self { tensor } + } + + /// Creates a new tensor from data and a device. + /// + /// # Arguments + /// + /// * `data` - The tensor's data. + /// * `device` - The device on which the tensor will be allocated. + /// + /// # Returns + /// + /// A new tensor. + pub fn from_data(data: TensorData, device: CandleDevice) -> Self { + let candle_shape: candle_core::Shape = data.shape.to_vec().into(); + let tensor = candle_core::Tensor::from_slice( + data.as_slice::().unwrap(), + candle_shape, + &device.into(), + ); + Self::new(tensor.unwrap()) + } +} + +pub(crate) trait IntoDType { + fn try_into_dtype(self) -> Result; + + fn into_dtype(self) -> candle_core::DType + where + Self: Sized, + { + self.try_into_dtype().unwrap() + } +} + +impl IntoDType for IntDType { + fn try_into_dtype(self) -> Result { + let dtype: DType = self.into(); + dtype.try_into_dtype() + } +} + +impl IntoDType for FloatDType { + fn try_into_dtype(self) -> Result { + let dtype: DType = self.into(); + dtype.try_into_dtype() + } +} + +impl IntoDType for DType { + fn try_into_dtype(self) -> Result { + match self { + DType::F64 => Ok(candle_core::DType::F64), + DType::F32 => Ok(candle_core::DType::F32), + DType::Flex32 => Ok(candle_core::DType::F32), + DType::F16 => Ok(candle_core::DType::F16), + DType::BF16 => Ok(candle_core::DType::BF16), + DType::I64 => Ok(candle_core::DType::I64), + DType::U32 => Ok(candle_core::DType::U32), + DType::U8 => Ok(candle_core::DType::U8), + DType::I16 => Ok(candle_core::DType::I16), + DType::I32 => Ok(candle_core::DType::I32), + DType::Bool(BoolStore::U8) => Ok(candle_core::DType::U8), + _ => Err(candle_core::Error::Msg(format!( + "Unsupported dtype {self:?}" + ))), + } + } +} diff --git a/crates/burn-communication/Cargo.toml b/crates/burn-communication/Cargo.toml new file mode 100644 index 0000000..678e075 --- /dev/null +++ b/crates/burn-communication/Cargo.toml @@ -0,0 +1,49 @@ +[package] +authors = ["Guilhem Ané (@Cielbird)", "Nathaniel Simard (@nathanielsimard)"] +description = "Abstractions for network communication for Burn" +edition.workspace = true +license.workspace = true +name = "burn-communication" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-communication" +version.workspace = true + +[lints] +workspace = true + +[features] +tracing = [ + "burn-std/tracing", + "burn-backend?/tracing", +] + +data-service = ["burn-backend"] +websocket = ["axum", "tokio-tungstenite", "futures"] + +[dependencies] +burn-std = { workspace = true, features = ["default"] } +bytes = { workspace = true } +derive-new = { workspace = true } +futures-util = { workspace = true } +log = { workspace = true } +rmp-serde = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_bytes = { workspace = true } +tokio = { workspace = true, features = ["rt-multi-thread", "sync", "signal", "tracing"] } +tokio-util = { workspace = true } +tracing = { workspace = true, features = ["default"] } +tracing-core = { workspace = true, features = ["default"] } +tracing-subscriber = { workspace = true, features = ["default", "fmt", "env-filter"] } + +# Tensor Data Service +burn-backend = { workspace = true, optional = true, features = ["default"] } + +# Websocket +axum = { workspace = true, features = ["ws"], optional = true } +tokio-tungstenite = { workspace = true, optional = true } +futures = { workspace = true, optional = true } + +[dev-dependencies] +# Integration tests for the websocket server (run with `--features websocket`). +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time", "net"] } +tokio-tungstenite = { workspace = true } diff --git a/crates/burn-communication/README.md b/crates/burn-communication/README.md new file mode 100644 index 0000000..d5f1490 --- /dev/null +++ b/crates/burn-communication/README.md @@ -0,0 +1,15 @@ +# Burn Communication + +Abstractions for network communication + +The Protocol trait defines how to communicate in a server/client style. +The server can set up routes with callbacks upon connection. + +## WebSocket + +Communication with WebSockets is implemented with the `websocket` feature. + +## Tensor Data Service + +The tensor data service provides easy utilities to share tensors peer-to-peer. +One peer can expose a tensor, and another can download it. Each peer is both a client and a server. diff --git a/crates/burn-communication/src/base.rs b/crates/burn-communication/src/base.rs new file mode 100644 index 0000000..123b24d --- /dev/null +++ b/crates/burn-communication/src/base.rs @@ -0,0 +1,221 @@ +use burn_std::future::DynFut; +use serde::{Deserialize, Serialize}; +use std::fmt::{Debug, Display}; +use std::hash::Hash; +use std::str::FromStr; + +/// A parsed network endpoint used by nodes to find each other. +/// +/// Construction normalizes the input so that equivalent spellings of the same endpoint +/// compare equal: a missing scheme defaults to `ws`, and a trailing path/slash is dropped. +/// As a result `"localhost:3000"`, `"ws://localhost:3000"`, and `"ws://localhost:3000/"` +/// all parse to the same `Address`. Equality is over `(scheme, host, port)`, so two devices +/// that point at the same server (e.g. different device indices) share one `Address` — that +/// is the signal the remote backend uses to detect a same-host transfer. +/// +/// Note: host names are compared textually, not resolved — `localhost` and `127.0.0.1` are +/// *not* considered equal. Aliases simply fall back to the network transfer path. +#[derive(Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Debug)] +pub struct Address { + scheme: String, + host: String, + port: Option, +} + +impl Address { + /// The scheme component (e.g. `ws`), lowercased. Defaults to `ws` when none was given. + pub fn scheme(&self) -> &str { + &self.scheme + } + + /// The host component (host name or IP literal), exactly as written. + pub fn host(&self) -> &str { + &self.host + } + + /// The port component, if one was specified. + pub fn port(&self) -> Option { + self.port + } + + /// Parse an endpoint string into its components, applying canonicalization. + fn parse(s: &str) -> Self { + let s = s.trim(); + let (scheme, rest) = match s.split_once("://") { + Some((scheme, rest)) => (scheme.to_ascii_lowercase(), rest), + None => ("ws".to_string(), s), + }; + + // Drop any path/query/fragment — only the authority identifies the endpoint. + let authority = rest + .split(['/', '?', '#']) + .next() + .unwrap_or(rest) + .trim_end_matches('/'); + + // Split host:port from the right so IPv6 literals like `[::1]:3000` keep their + // colons; only the final `:port` (when it parses as a port) is treated as the port. + let (host, port) = match authority.rsplit_once(':') { + Some((host, port)) if !host.is_empty() => match port.parse::() { + Ok(port) => (host.to_string(), Some(port)), + Err(_) => (authority.to_string(), None), + }, + _ => (authority.to_string(), None), + }; + + Self { scheme, host, port } + } +} + +impl From<&str> for Address { + fn from(s: &str) -> Self { + Address::parse(s) + } +} + +impl FromStr for Address { + type Err = String; + + fn from_str(s: &str) -> Result { + Ok(Address::parse(s)) + } +} + +impl Display for Address { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}://{}", self.scheme, self.host)?; + if let Some(port) = self.port { + write!(f, ":{port}")?; + } + Ok(()) + } +} + +/// The protocol used for the communications. +pub trait Protocol: Clone + Send + Sync + 'static { + /// The client implementation for the current protocol. + type Client: ProtocolClient; + /// The server implementation for the current protocol. + type Server: ProtocolServer; +} + +/// Error that happens during a communication. +pub trait CommunicationError: Debug + Send + 'static {} + +/// The client is only used to create a [channel](CommunicationChannel), which should be use to +/// transmit information with the [server](ProtocolServer). +pub trait ProtocolClient: Send + Sync + 'static { + /// Channel used by this protocol. + type Channel: CommunicationChannel; + /// The error type. + type Error: CommunicationError; + + /// Opens a new [channel](CommunicationChannel) with the current protocol at the given + /// [address](Address) and route. + /// + /// * `address` - Address to connect to + /// * `route` - The name of the route (no slashes) + /// + /// Returns `Err(Self::Error)` if the connection couldn't be established (address parse + /// failure, server unreachable, handshake failure, …). The error carries enough context + /// to identify the cause; callers should surface it rather than swallowing it. + fn connect(address: Address, route: &str) -> DynFut>; +} + +/// Data sent and received by the client and server. +#[derive(new)] +pub struct Message { + /// The data is always encoded as bytes. + pub data: bytes::Bytes, +} + +/// Defines how to create a server that respond to a [channel](CommunicationChannel). +pub trait ProtocolServer: Sized + Send + Sync + 'static { + /// Channel used by this protocol. + type Channel: CommunicationChannel; + /// The error type. + type Error: CommunicationError; + + /// Defines an endpoint with the function that responds. + /// TODO Docs: does it need a slash? + fn route(self, path: &str, callback: C) -> Self + where + C: FnOnce(Self::Channel) -> Fut + Clone + Send + Sync + 'static, + Fut: Future + Send + 'static; + + /// Start the server. + fn serve( + self, + shutdown: F, + ) -> impl Future> + Send + 'static + where + F: Future + Send + 'static; +} + +/// Handles communications. +pub trait CommunicationChannel: Send + 'static { + type Error: CommunicationError; + + /// Send a [message](Message) on the channel. + fn send( + &mut self, + message: Message, + ) -> impl std::future::Future> + Send; + + /// Receive a [message](Message) on the channel and returns a new [response message](Message). + fn recv( + &mut self, + ) -> impl std::future::Future, Self::Error>> + Send; + + fn close(&mut self) -> impl std::future::Future> + Send; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn address_defaults_scheme_and_canonicalizes() { + let bare = Address::from("localhost:3000"); + let with_scheme = Address::from("ws://localhost:3000"); + let trailing = Address::from("ws://localhost:3000/"); + + // Equivalent spellings canonicalize to the same value. + assert_eq!(bare, with_scheme); + assert_eq!(with_scheme, trailing); + assert_eq!(bare.scheme(), "ws"); + assert_eq!(bare.host(), "localhost"); + assert_eq!(bare.port(), Some(3000)); + assert_eq!(with_scheme.to_string(), "ws://localhost:3000"); + } + + #[test] + fn address_distinguishes_port_host_and_scheme() { + assert_ne!( + Address::from("ws://host:3000"), + Address::from("ws://host:3001") + ); + assert_ne!(Address::from("ws://a:3000"), Address::from("ws://b:3000")); + assert_ne!( + Address::from("ws://host:3000"), + Address::from("wss://host:3000") + ); + // Host names are textual, not resolved. + assert_ne!( + Address::from("ws://localhost:3000"), + Address::from("ws://127.0.0.1:3000") + ); + } + + #[test] + fn address_handles_ipv6_and_missing_port() { + let v6 = Address::from("ws://[::1]:3000"); + assert_eq!(v6.host(), "[::1]"); + assert_eq!(v6.port(), Some(3000)); + + let no_port = Address::from("ws://example.com"); + assert_eq!(no_port.host(), "example.com"); + assert_eq!(no_port.port(), None); + assert_eq!(no_port.to_string(), "ws://example.com"); + } +} diff --git a/crates/burn-communication/src/external_comm.rs b/crates/burn-communication/src/external_comm.rs new file mode 100644 index 0000000..0fc187a --- /dev/null +++ b/crates/burn-communication/src/external_comm.rs @@ -0,0 +1,260 @@ +//! This module enables direct data transfer between servers without blocking the client or any server. +//! +//! It eliminates the need for intermediate data transfer through the client, avoiding the process of downloading data from one server and reuploading it to another. +//! +//! The module provides an optimized mechanism for servers to communicate directly, streamlining data movement between them without involving the client. + +use crate::Message; +use crate::base::Protocol; +use crate::base::{Address, CommunicationChannel, ProtocolClient, ProtocolServer}; +use burn_backend::{TensorData, backend::Backend}; +use serde::{Deserialize, Serialize}; +use std::{collections::HashMap, marker::PhantomData, sync::Arc}; +use tokio::sync::Mutex; +use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct TensorTransferId(u64); + +impl From for TensorTransferId { + fn from(value: u64) -> Self { + Self(value) + } +} + +impl TensorTransferId { + pub fn next(&mut self) { + self.0 += 1; + } +} + +#[derive(Debug, Serialize, Deserialize)] +enum ExternalCommMessage { + TensorRequest(TensorTransferId), + Tensor(TensorData), +} + +type ClientChannelRef = Arc::Channel>>; + +pub struct ExternalCommService> { + /// Maps tensor transfer IDs to their exposed state. + pub exposed_tensors: Mutex>, + /// Maps node addresses to their channels. + pub channels: Mutex>>, + /// Notify when a new tensor is exposed. + pub new_tensor_notify: Arc, + + cancel_token: CancellationToken, + + _phantom_data: PhantomData, +} + +pub struct TensorExposeState { + /// The bytes of the tensor data message. Message::Data(...) serialized with rmp_serde + pub bytes: bytes::Bytes, + /// How many times the tensor will be downloaded + pub max_downloads: u32, + /// How man times the tensor has been downloaded + pub cur_download_count: u32, +} + +/// Provides a routing function for a tensor data service for a communications server +pub trait ExternalCommServer { + /// Routes the tensor data service to the "/data" route + fn route_external_comm(self, state: Arc>) -> Self; +} + +impl + 'static> + ExternalCommServer for S +{ + fn route_external_comm(self, state: Arc>) -> Self { + self.route("/data", async move |stream: S::Channel| { + state.handle_data_channel(stream).await; + }) + } +} + +impl ExternalCommService { + pub fn new(cancel_token: CancellationToken) -> Self { + Self { + exposed_tensors: Mutex::new(HashMap::new()), + channels: Mutex::new(HashMap::new()), + new_tensor_notify: Arc::new(Notify::new()), + cancel_token, + _phantom_data: PhantomData::, + } + } + + /// Exposes a tensor to the data server, allowing it to be downloaded by other nodes. + pub async fn expose( + &self, + tensor: B::FloatTensorPrimitive, + max_downloads: u32, + transfer_id: TensorTransferId, + ) { + let data = B::float_into_data(tensor).await.unwrap(); + self.expose_data(data, max_downloads, transfer_id).await + } + + /// Exposes a tensor data to the data server, allowing it to be downloaded by other nodes. + pub async fn expose_data( + &self, + tensor_data: TensorData, + max_downloads: u32, + transfer_id: TensorTransferId, + ) { + let bytes: bytes::Bytes = rmp_serde::to_vec(&ExternalCommMessage::Tensor(tensor_data)) + .unwrap() + .into(); + let mut exposed_tensors = self.exposed_tensors.lock().await; + exposed_tensors.insert( + transfer_id, + TensorExposeState { + bytes, + max_downloads, + cur_download_count: 0, + }, + ); + core::mem::drop(exposed_tensors); + self.new_tensor_notify.notify_waiters(); + } + + pub async fn close(&self) { + // Send a closing message to every open WebSocket stream + + let mut streams = self.channels.lock().await; + for (_, stream) in streams.drain() { + let mut stream = stream.lock().await; + + stream + .close() + .await + .expect("Failed to close WebSocket stream"); + } + } + + /// Downloads a tensor that is exposed on another server. Requires a Tokio 1.x runtime + /// + /// Returns None if the peer closes the connection + pub async fn download_tensor( + &self, + remote: Address, + transfer_id: TensorTransferId, + ) -> Option { + log::info!("Downloading tensor from {remote:?}"); + + let stream = self.get_data_stream(remote).await; + let mut stream = stream.lock().await; + + // Send the download request with the download id + let bytes: bytes::Bytes = + rmp_serde::to_vec(&ExternalCommMessage::TensorRequest(transfer_id)) + .unwrap() + .into(); + stream + .send(Message::new(bytes)) + .await + .expect("Failed to send download id"); + + if let Ok(msg) = stream.recv().await { + let Some(msg) = msg else { + log::warn!("Received None message from the websocket, closing connection."); + return None; + }; + + let ExternalCommMessage::Tensor(data) = rmp_serde::from_slice(&msg.data) + .expect("Can deserialize messages from the websocket.") + else { + panic!("Message should have been TensorData") + }; + return Some(data); + } + log::warn!("Closed connection"); + None + } + + /// Get the WebSocket stream for the given address, or create a new one if it doesn't exist. + async fn get_data_stream( + &self, + address: Address, + ) -> Arc::Channel>> { + let mut streams = self.channels.lock().await; + match streams.get(&address) { + Some(stream) => stream.clone(), + None => { + // Open a new WebSocket connection to the address + let stream = match P::Client::connect(address.clone(), "data").await { + Ok(stream) => stream, + Err(err) => panic!( + "Failed to open remote 'data' channel to {address}: {err:?}. \ + Is a `burn-remote` server running at that address?" + ), + }; + + let stream = Arc::new(Mutex::new(stream)); + streams.insert(address.clone(), stream.clone()); + + stream + } + } + } + + /// Get the requested exposed tensor data, and update download counter + async fn get_exposed_tensor_bytes( + &self, + transfer_id: TensorTransferId, + ) -> Option { + loop { + { + let mut exposed_tensors = self.exposed_tensors.lock().await; + // take the tensor out of the hashmap while we download + if let Some(mut exposed_state) = exposed_tensors.remove(&transfer_id) { + exposed_state.cur_download_count += 1; + let bytes = if exposed_state.cur_download_count == exposed_state.max_downloads { + exposed_state.bytes + } else { + let bytes = exposed_state.bytes.clone(); + exposed_tensors.insert(transfer_id, exposed_state); + bytes + }; + return Some(bytes); + } + } + // No matching tensor, wait for a new one to come in. + self.new_tensor_notify.notified().await; + } + } + + /// Handle incoming connections for downloading tensors. + pub(crate) async fn handle_data_channel( + &self, + mut channel: ::Channel, + ) { + log::info!("[Data Handler] New connection for download."); + + while !self.cancel_token.is_cancelled() { + match channel.recv().await { + Ok(message) => { + if let Some(msg) = message { + let bytes = msg.data; + let msg: ExternalCommMessage = rmp_serde::from_slice(&bytes) + .expect("Can deserialize messages from the websocket."); + let ExternalCommMessage::TensorRequest(transfer_id) = msg else { + panic!("Received a message that wasn't a tensor request! {msg:?}"); + }; + + let bytes = self.get_exposed_tensor_bytes(transfer_id).await.unwrap(); + + channel.send(Message::new(bytes)).await.unwrap(); + } else { + log::info!("Closed connection"); + return; + } + } + Err(err) => panic!("Failed to receive message from websocket: {err:?}"), + }; + } + log::info!("[Data Service] Closing connection for download."); + } +} diff --git a/crates/burn-communication/src/lib.rs b/crates/burn-communication/src/lib.rs new file mode 100644 index 0000000..49acfbb --- /dev/null +++ b/crates/burn-communication/src/lib.rs @@ -0,0 +1,13 @@ +#[macro_use] +extern crate derive_new; + +mod base; +pub use base::*; + +pub mod util; + +#[cfg(feature = "websocket")] +pub mod websocket; + +#[cfg(feature = "data-service")] +pub mod external_comm; diff --git a/crates/burn-communication/src/util.rs b/crates/burn-communication/src/util.rs new file mode 100644 index 0000000..7b2b909 --- /dev/null +++ b/crates/burn-communication/src/util.rs @@ -0,0 +1,46 @@ +use tracing_core::{Level, LevelFilter}; +use tracing_subscriber::{ + Layer, filter::filter_fn, layer::SubscriberExt, registry, util::SubscriberInitExt, +}; + +/// Utilities to help handle communication termination. +pub async fn os_shutdown_signal() { + let ctrl_c = async { + tokio::signal::ctrl_c() + .await + .expect("failed to install Ctrl+C handler"); + }; + + #[cfg(unix)] + let terminate = async { + tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .expect("failed to install signal handler") + .recv() + .await; + }; + + #[cfg(not(unix))] + let terminate = std::future::pending::<()>(); + + tokio::select! { + _ = ctrl_c => {}, + _ = terminate => {}, + } +} + +pub(crate) fn init_logging() { + let layer = tracing_subscriber::fmt::layer() + .with_filter(LevelFilter::INFO) + .with_filter(filter_fn(|m| { + if let Some(path) = m.module_path() { + // The wgpu crate is logging too much, so we skip `info` level. + if path.starts_with("wgpu") && *m.level() >= Level::INFO { + return false; + } + } + true + })); + + // If we start multiple servers in the same process, this will fail, it's ok + let _ = registry().with(layer).try_init(); +} diff --git a/crates/burn-communication/src/websocket/base.rs b/crates/burn-communication/src/websocket/base.rs new file mode 100644 index 0000000..af7981f --- /dev/null +++ b/crates/burn-communication/src/websocket/base.rs @@ -0,0 +1,26 @@ +use crate::{ + base::{Address, Protocol}, + websocket::{client::WsClient, server::WsServer}, +}; + +#[derive(Clone)] +/// A websocket implements a [communication protocol](Protocol) that can be used to communicate +/// over the internet. +pub struct WebSocket {} + +impl Protocol for WebSocket { + type Client = WsClient; + type Server = WsServer; +} + +/// Validate that an [`Address`] uses the websocket scheme. +/// +/// The [`Address`] is already canonicalized at construction (scheme defaults to `ws`, path +/// stripped), so this only has to reject a non-`ws` scheme. The address is returned +/// unchanged on success and its [`Display`](std::fmt::Display) form is the connection url. +pub(crate) fn parse_ws_address(address: Address) -> Result { + match address.scheme() { + "ws" | "wss" => Ok(address), + other => Err(format!("Invalid scheme: {other}")), + } +} diff --git a/crates/burn-communication/src/websocket/client.rs b/crates/burn-communication/src/websocket/client.rs new file mode 100644 index 0000000..436d25d --- /dev/null +++ b/crates/burn-communication/src/websocket/client.rs @@ -0,0 +1,207 @@ +use crate::{ + base::{Address, CommunicationChannel, CommunicationError, Message, ProtocolClient}, + websocket::base::parse_ws_address, +}; +use burn_std::future::DynFut; +use futures::{ + SinkExt, StreamExt, + stream::{SplitSink, SplitStream}, +}; +use tokio::net::TcpStream; +use tokio_tungstenite::{ + MaybeTlsStream, WebSocketStream, connect_async_with_config, + tungstenite::{self, protocol::WebSocketConfig}, +}; + +#[derive(Clone)] +pub struct WsClient; + +impl ProtocolClient for WsClient { + type Channel = WsClientChannel; + type Error = WsClientError; + + fn connect(address: Address, route: &str) -> DynFut> { + Box::pin(connect_ws(address, route.to_owned())) + } +} + +/// Open a new WebSocket connection to the address. +async fn connect_ws(address: Address, route: String) -> Result { + let address = parse_ws_address(address).map_err(WsClientError::Address)?; + let url = format!("{address}/{route}"); + const MB: usize = 1024 * 1024; + let (stream, _) = connect_async_with_config( + url, + Some( + WebSocketConfig::default() + .write_buffer_size(0) + .max_message_size(None) + .max_frame_size(Some(MB * 512)) + .accept_unmasked_frames(true) + .read_buffer_size(64 * 1024), // 64 KiB (previous default) + ), + true, + ) + .await?; + + Ok(WsClientChannel { inner: stream }) +} +pub struct WsClientChannel { + inner: WebSocketStream>, +} + +impl CommunicationChannel for WsClientChannel { + type Error = WsClientError; + + async fn send(&mut self, msg: Message) -> Result<(), WsClientError> { + self.inner + .send(tungstenite::Message::Binary(msg.data)) + .await?; + + Ok(()) + } + + async fn recv(&mut self) -> Result, WsClientError> { + // Keep reading until we get a data frame, a close, an error, or the stream ends. + // Ping/Pong (keepalives) and Text frames must be skipped, not turned into errors: a + // single keepalive ping would otherwise look like a fatal error to the consuming loop + // (e.g. the client's response demux) and kill an otherwise-healthy connection. + // tungstenite already answers pings with pongs at the protocol layer, so ignoring them + // here is safe. Mirrors the server channel's `recv`. + loop { + match self.inner.next().await { + Some(Ok(tungstenite::Message::Binary(data))) => return Ok(Some(Message { data })), + // A close frame is an orderly end-of-stream. + Some(Ok(tungstenite::Message::Close(_close_frame))) => return Ok(None), + // Control/text frames carry no protocol payload: skip and keep reading. + Some(Ok( + tungstenite::Message::Ping(_) + | tungstenite::Message::Pong(_) + | tungstenite::Message::Text(_) + | tungstenite::Message::Frame(_), + )) => continue, + Some(Err(err)) => return Err(WsClientError::Tungstenite(err)), + // The stream is exhausted: the peer went away without a close frame (server + // shut the socket, network drop). A common disconnect path, not an error — + // treat it as a clean end-of-stream rather than panicking with `todo!()`. + None => return Ok(None), + } + } + } + + async fn close(&mut self) -> Result<(), WsClientError> { + let reason = "Peer is closing".to_string(); + + self.inner + .send(tungstenite::Message::Close(Some( + tungstenite::protocol::CloseFrame { + code: tungstenite::protocol::frame::coding::CloseCode::Normal, + reason: reason.clone().into(), + }, + ))) + .await?; + + Ok(()) + } +} + +impl WsClientChannel { + /// Split into independently-owned send and receive halves, so a writer task and a reader loop + /// can run concurrently over one full-duplex socket. + pub fn split(self) -> (WsClientSink, WsClientStream) { + let (sink, stream) = self.inner.split(); + ( + WsClientSink { inner: sink }, + WsClientStream { inner: stream }, + ) + } +} + +/// Send half of a split [`WsClientChannel`]. +pub struct WsClientSink { + inner: SplitSink>, tungstenite::Message>, +} + +impl WsClientSink { + pub async fn send(&mut self, msg: Message) -> Result<(), WsClientError> { + self.inner + .send(tungstenite::Message::Binary(msg.data)) + .await?; + Ok(()) + } + + pub async fn close(&mut self) -> Result<(), WsClientError> { + self.inner + .send(tungstenite::Message::Close(Some( + tungstenite::protocol::CloseFrame { + code: tungstenite::protocol::frame::coding::CloseCode::Normal, + reason: "Peer is closing".to_string().into(), + }, + ))) + .await?; + Ok(()) + } +} + +/// Receive half of a split [`WsClientChannel`]. +pub struct WsClientStream { + inner: SplitStream>>, +} + +impl WsClientStream { + /// Mirrors [`WsClientChannel::recv`]: surface binary frames, skip keepalive/text frames, and + /// treat a close frame or an exhausted stream as a clean end-of-stream. + pub async fn recv(&mut self) -> Result, WsClientError> { + loop { + match self.inner.next().await { + Some(Ok(tungstenite::Message::Binary(data))) => return Ok(Some(Message { data })), + Some(Ok(tungstenite::Message::Close(_close_frame))) => return Ok(None), + Some(Ok( + tungstenite::Message::Ping(_) + | tungstenite::Message::Pong(_) + | tungstenite::Message::Text(_) + | tungstenite::Message::Frame(_), + )) => continue, + Some(Err(err)) => return Err(WsClientError::Tungstenite(err)), + None => return Ok(None), + } + } + } +} + +#[derive(Debug)] +pub enum WsClientError { + /// The address couldn't be parsed (e.g. missing scheme, unsupported scheme). + Address(String), + Io(std::io::Error), + Tungstenite(tungstenite::Error), + UnknownMessage(String), + Other(String), +} +impl CommunicationError for WsClientError {} + +impl core::fmt::Display for WsClientError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + Self::Address(msg) => write!(f, "invalid address: {msg}"), + Self::Io(err) => write!(f, "io error: {err}"), + Self::Tungstenite(err) => write!(f, "websocket error: {err}"), + Self::UnknownMessage(msg) => write!(f, "unknown message: {msg}"), + Self::Other(msg) => write!(f, "{msg}"), + } + } +} + +impl std::error::Error for WsClientError {} + +impl From for WsClientError { + fn from(err: std::io::Error) -> Self { + Self::Io(err) + } +} + +impl From for WsClientError { + fn from(err: tungstenite::Error) -> Self { + Self::Tungstenite(err) + } +} diff --git a/crates/burn-communication/src/websocket/mod.rs b/crates/burn-communication/src/websocket/mod.rs new file mode 100644 index 0000000..ca64da1 --- /dev/null +++ b/crates/burn-communication/src/websocket/mod.rs @@ -0,0 +1,7 @@ +mod base; +mod client; +mod server; + +pub use base::*; +pub use client::*; +pub use server::*; diff --git a/crates/burn-communication/src/websocket/server.rs b/crates/burn-communication/src/websocket/server.rs new file mode 100644 index 0000000..fe7d3f4 --- /dev/null +++ b/crates/burn-communication/src/websocket/server.rs @@ -0,0 +1,250 @@ +use std::net::SocketAddr; + +use crate::{ + base::{CommunicationChannel, CommunicationError, Message, ProtocolServer}, + util::init_logging, +}; +use axum::{ + Router, + extract::{ + State, WebSocketUpgrade, + ws::{self, WebSocket}, + }, + routing::get, +}; +use futures::{ + SinkExt, StreamExt, + stream::{SplitSink, SplitStream}, +}; + +#[derive(Clone, Debug)] +pub struct WsServer { + port: u16, + router: Router<()>, +} + +pub struct WsServerChannel { + inner: WebSocket, +} + +impl WsServer { + pub fn new(port: u16) -> Self { + Self { + port, + router: Router::new(), + } + } + + /// Serve on an already-bound listener instead of binding `0.0.0.0:port` ourselves. + /// + /// This is what [`serve`](ProtocolServer::serve) delegates to; it is exposed so a caller + /// that needs the actual bound address — e.g. a test binding an ephemeral `:0` port and + /// reading the port back from `listener.local_addr()` — can do so without a fixed port. + pub async fn serve_on( + self, + listener: tokio::net::TcpListener, + shutdown: F, + ) -> Result<(), WsServerError> + where + F: Future + Send + 'static, + { + init_logging(); + + // Report the address the listener actually bound to (resolves an ephemeral `:0` port to + // the real one), so the log confirms the server is accepting connections and tells the + // operator where. + match listener.local_addr() { + Ok(addr) => log::info!("Server started, listening on {addr}"), + Err(err) => log::info!("Server started (could not resolve bound address: {err})"), + } + + axum::serve( + listener, + self.router + .into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(shutdown) + .await?; + + Ok(()) + } +} + +impl ProtocolServer for WsServer { + type Channel = WsServerChannel; + type Error = WsServerError; + + async fn serve(self, shutdown: F) -> Result<(), Self::Error> + where + F: Future + Send + 'static, + { + let address = format!("0.0.0.0:{}", self.port); + log::info!("Starting server {address}"); + + let listener = tokio::net::TcpListener::bind(address).await?; + + self.serve_on(listener, shutdown).await + } + + fn route(mut self, path: &str, callback: C) -> Self + where + C: FnOnce(WsServerChannel) -> Fut + Clone + Send + Sync + 'static, + Fut: Future + Send + 'static, + { + // Format path: should start with a / + let path = if path.starts_with("/") { + path.to_owned() + } else { + format!("/{path}") + }; + + let method = get(|ws: WebSocketUpgrade, _: State<()>| async { + ws.on_upgrade(async move |socket| { + callback(WsServerChannel { inner: socket }).await; + }) + }); + + self.router = self.router.route(&path, method); + + self + } +} + +impl CommunicationChannel for WsServerChannel { + type Error = WsServerError; + + async fn send(&mut self, message: Message) -> Result<(), WsServerError> { + self.inner.send(ws::Message::Binary(message.data)).await?; + + Ok(()) + } + + async fn recv(&mut self) -> Result, WsServerError> { + // Keep reading until we get a data frame, a close, an error, or the stream ends. + // Ping/Pong (keepalives) and Text frames must be skipped, not turned into errors: a + // single keepalive ping would otherwise look like a fatal error to the consuming loop + // and kill an otherwise-healthy connection. tungstenite already answers pings with + // pongs at the protocol layer, so ignoring them here is safe. + loop { + match self.inner.next().await { + Some(Ok(ws::Message::Binary(data))) => return Ok(Some(Message { data })), + // A close frame is an orderly end-of-stream. + Some(Ok(ws::Message::Close(_close_frame))) => return Ok(None), + // Control/text frames carry no protocol payload: skip and keep reading. + Some(Ok(ws::Message::Ping(_) | ws::Message::Pong(_) | ws::Message::Text(_))) => { + continue; + } + Some(Err(err)) => return Err(WsServerError::Axum(err)), + // The stream is exhausted: the peer went away without a close frame (closed + // tab, network drop, killed client). This is a common disconnect path, not an + // error — treat it as a clean end-of-stream rather than panicking. + None => return Ok(None), + } + } + } + + async fn close(&mut self) -> Result<(), WsServerError> { + let reason = "Peer is closing".to_string(); + + self.inner + .send(ws::Message::Close(Some(ws::CloseFrame { + code: 1000, // code: Normal + reason: reason.clone().into(), + }))) + .await?; + + Ok(()) + } +} + +impl WsServerChannel { + /// Split into independently-owned send and receive halves, so a writer task and a reader loop + /// can run concurrently over one full-duplex socket. + pub fn split(self) -> (WsServerSink, WsServerStream) { + let (sink, stream) = self.inner.split(); + ( + WsServerSink { inner: sink }, + WsServerStream { inner: stream }, + ) + } +} + +/// Send half of a split [`WsServerChannel`]. +pub struct WsServerSink { + inner: SplitSink, +} + +impl WsServerSink { + pub async fn send(&mut self, message: Message) -> Result<(), WsServerError> { + self.inner.send(ws::Message::Binary(message.data)).await?; + Ok(()) + } + + pub async fn close(&mut self) -> Result<(), WsServerError> { + self.inner + .send(ws::Message::Close(Some(ws::CloseFrame { + code: 1000, // Normal + reason: "Peer is closing".to_string().into(), + }))) + .await?; + Ok(()) + } +} + +/// Receive half of a split [`WsServerChannel`]. +pub struct WsServerStream { + inner: SplitStream, +} + +impl WsServerStream { + /// Mirrors [`WsServerChannel::recv`]: surface binary frames, skip keepalive/text frames, and + /// treat a close frame or an exhausted stream as a clean end-of-stream. + pub async fn recv(&mut self) -> Result, WsServerError> { + loop { + match self.inner.next().await { + Some(Ok(ws::Message::Binary(data))) => return Ok(Some(Message { data })), + Some(Ok(ws::Message::Close(_close_frame))) => return Ok(None), + Some(Ok(ws::Message::Ping(_) | ws::Message::Pong(_) | ws::Message::Text(_))) => { + continue; + } + Some(Err(err)) => return Err(WsServerError::Axum(err)), + None => return Ok(None), + } + } + } +} + +#[derive(Debug)] +pub enum WsServerError { + Io(std::io::Error), + Axum(axum::Error), + UnknownMessage(String), + Other(String), +} + +impl CommunicationError for WsServerError {} + +impl From for WsServerError { + fn from(err: std::io::Error) -> Self { + Self::Io(err) + } +} + +impl From for WsServerError { + fn from(err: axum::Error) -> Self { + Self::Axum(err) + } +} + +impl core::fmt::Display for WsServerError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + Self::Io(err) => write!(f, "io error: {err}"), + Self::Axum(err) => write!(f, "websocket error: {err}"), + Self::UnknownMessage(msg) => write!(f, "unknown message: {msg}"), + Self::Other(msg) => write!(f, "{msg}"), + } + } +} + +impl std::error::Error for WsServerError {} diff --git a/crates/burn-communication/tests/websocket_server.rs b/crates/burn-communication/tests/websocket_server.rs new file mode 100644 index 0000000..7df5537 --- /dev/null +++ b/crates/burn-communication/tests/websocket_server.rs @@ -0,0 +1,367 @@ +//! Integration tests for the websocket [`WsServer`] channel implementation. +//! +//! These drive a real `WsServer` (bound to an ephemeral `127.0.0.1:0` port read back from the +//! listener) with a real `tokio-tungstenite` client, exercising the concurrency and +//! disconnect behaviour of the channel. Every await that could hang is wrapped in a timeout so +//! a regression fails fast instead of blocking CI. +//! +//! Run with: `cargo test -p burn-communication --features websocket`. +#![cfg(feature = "websocket")] + +use std::time::Duration; + +use burn_communication::websocket::{WsServer, WsServerChannel}; +use burn_communication::{CommunicationChannel, Message, ProtocolServer}; +use futures_util::{SinkExt, StreamExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::{mpsc, oneshot}; +use tokio::time::timeout; +use tokio_tungstenite::tungstenite::Message as WsMessage; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; + +/// Generous cap on every awaited operation: a passing test finishes in milliseconds; a +/// regression (panic, deadlock, dropped connection) trips the timeout instead of hanging. +const TIMEOUT: Duration = Duration::from_secs(10); + +type Client = WebSocketStream>; + +/// How a single server-side `recv` resolved — used by handlers that report back to the test. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Recv { + Message, + End, + Error, +} + +fn classify(result: Result, E>) -> Recv { + match result { + Ok(Some(_)) => Recv::Message, + Ok(None) => Recv::End, + Err(_) => Recv::Error, + } +} + +/// A running test server bound to an ephemeral port, with a shutdown the test controls. +struct TestServer { + port: u16, + shutdown: Option>, + handle: tokio::task::JoinHandle<()>, +} + +impl TestServer { + /// Bind an ephemeral port, let `build` register the routes, and serve on it. Binding + /// happens before the serve task is spawned, so the port is in the listen backlog by the + /// time this returns and clients can connect without a readiness race. + async fn start(build: impl FnOnce(WsServer) -> WsServer) -> Self { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind ephemeral port"); + let port = listener.local_addr().expect("local_addr").port(); + + // The port passed to `new` is unused — we serve on our own pre-bound listener. + let server = build(WsServer::new(0)); + + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let handle = tokio::spawn(async move { + let shutdown = async move { + let _ = shutdown_rx.await; + }; + if let Err(err) = server.serve_on(listener, shutdown).await { + eprintln!("test server error: {err:?}"); + } + }); + + Self { + port, + shutdown: Some(shutdown_tx), + handle, + } + } + + fn url(&self, path: &str) -> String { + format!( + "ws://127.0.0.1:{}/{}", + self.port, + path.trim_start_matches('/') + ) + } + + async fn shutdown(mut self) { + if let Some(tx) = self.shutdown.take() { + let _ = tx.send(()); + } + let _ = timeout(TIMEOUT, self.handle).await; + } +} + +/// A route handler that echoes every binary message back, looping until the stream ends. +async fn echo(mut channel: WsServerChannel) { + while let Ok(Some(msg)) = channel.recv().await { + if channel.send(msg).await.is_err() { + break; + } + } +} + +// --- client helpers ----------------------------------------------------------------------- + +async fn connect(url: &str) -> Client { + let (ws, _resp) = timeout(TIMEOUT, connect_async(url)) + .await + .expect("connect timed out") + .expect("connect failed"); + ws +} + +async fn send_binary(ws: &mut Client, data: &[u8]) { + timeout(TIMEOUT, ws.send(WsMessage::Binary(data.to_vec().into()))) + .await + .expect("send timed out") + .expect("send failed"); +} + +/// Read until the next binary frame, skipping pongs/pings/text the server may interleave. +async fn recv_binary(ws: &mut Client) -> Vec { + let fut = async { + loop { + match ws.next().await { + Some(Ok(WsMessage::Binary(data))) => return data.to_vec(), + Some(Ok(_)) => continue, + Some(Err(err)) => panic!("client recv error: {err:?}"), + None => panic!("client stream ended before a binary frame"), + } + } + }; + timeout(TIMEOUT, fut).await.expect("recv_binary timed out") +} + +/// Drive a clean websocket close handshake to completion from the client side. +async fn close_cleanly(mut ws: Client) { + let _ = ws.close(None).await; + // Drain so the server's close echo is read and the handshake completes on both sides. + while let Some(Ok(_)) = ws.next().await {} +} + +// --- tests -------------------------------------------------------------------------------- + +/// Many clients at once each get their own response — proves connections aren't serialized. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_connections_are_not_serialized() { + let server = TestServer::start(|s| s.route("/echo", echo)).await; + let url = server.url("echo"); + + // Spawn all clients eagerly (collect forces the spawns) so they run concurrently, then + // await them — rather than spawn-then-await one at a time, which would serialize them. + let clients: Vec<_> = (0..50) + .map(|i| { + let url = url.clone(); + tokio::spawn(async move { + let mut ws = connect(&url).await; + let payload = format!("hello-{i}").into_bytes(); + send_binary(&mut ws, &payload).await; + assert_eq!( + recv_binary(&mut ws).await, + payload, + "client {i} echo mismatch" + ); + close_cleanly(ws).await; + }) + }) + .collect(); + + for (i, client) in clients.into_iter().enumerate() { + timeout(TIMEOUT, client) + .await + .unwrap_or_else(|_| panic!("client {i} timed out")) + .unwrap_or_else(|_| panic!("client {i} panicked")); + } + + server.shutdown().await; +} + +/// One peer vanishing abruptly (TCP drop, no close frame) must not take down the others. +/// Regression for bug #1: the old `None => todo!()` would panic the handler on disconnect. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn abrupt_disconnect_does_not_affect_other_connections() { + // Echo handler that also reports how its stream ended, so the test can assert the + // disconnected peer's handler terminated rather than panicked or hung. + let (report_tx, mut report_rx) = mpsc::unbounded_channel::(); + let server = TestServer::start(move |s| { + s.route("/echo", move |mut channel: WsServerChannel| { + let report = report_tx.clone(); + async move { + loop { + match channel.recv().await { + Ok(Some(msg)) => { + if channel.send(msg).await.is_err() { + let _ = report.send(Recv::Error); + return; + } + } + Ok(None) => { + let _ = report.send(Recv::End); + return; + } + Err(_) => { + let _ = report.send(Recv::Error); + return; + } + } + } + } + }) + }) + .await; + let url = server.url("echo"); + + let mut a = connect(&url).await; + let mut b = connect(&url).await; + + // Both work to start with. + send_binary(&mut a, b"a1").await; + assert_eq!(recv_binary(&mut a).await, b"a1"); + send_binary(&mut b, b"b1").await; + assert_eq!(recv_binary(&mut b).await, b"b1"); + + // Abruptly drop A's TCP connection — no websocket close handshake. + drop(a); + + // A's handler must terminate (without panicking); with the `todo!()` bug it would panic on + // the bare-None path and never report, tripping this timeout. + let ended = timeout(TIMEOUT, report_rx.recv()) + .await + .expect("A's handler did not terminate — it likely panicked") + .expect("report channel closed"); + assert!( + matches!(ended, Recv::End | Recv::Error), + "A's handler ended in an unexpected state: {ended:?}" + ); + + // B must still be alive and responsive after A vanished. + send_binary(&mut b, b"b2").await; + assert_eq!( + recv_binary(&mut b).await, + b"b2", + "B died after A disconnected" + ); + + close_cleanly(b).await; + server.shutdown().await; +} + +/// Hit the bare end-of-stream arm directly. Regression for bug #1: a second `recv` after the +/// stream has terminated yields `None`, which the old `todo!()` turned into a panic. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn stream_end_after_close_does_not_panic() { + let (report_tx, mut report_rx) = mpsc::unbounded_channel::<(Recv, Recv)>(); + let server = TestServer::start(move |s| { + s.route("/drain", move |mut channel: WsServerChannel| { + let report = report_tx.clone(); + async move { + // First recv sees the client's Close -> Ok(None). The stream is then + // terminated, so the second recv hits the bare-None arm -> Ok(None) with the + // fix, panic with the old `todo!()`. + let first = classify(channel.recv().await); + let second = classify(channel.recv().await); + let _ = report.send((first, second)); + } + }) + }) + .await; + + let ws = connect(&server.url("drain")).await; + close_cleanly(ws).await; + + let (first, second) = timeout(TIMEOUT, report_rx.recv()) + .await + .expect("/drain handler did not finish — it likely panicked on the bare-None path") + .expect("report channel closed"); + assert_eq!( + first, + Recv::End, + "close frame should make the first recv return Ok(None)" + ); + assert_eq!( + second, + Recv::End, + "a terminated stream should make recv return Ok(None), not panic" + ); + + server.shutdown().await; +} + +/// A keepalive ping must be skipped, not treated as a fatal error. Regression for bug #2: +/// the old catch-all turned Ping/Pong/Text into `Err(UnknownMessage)`, killing the connection. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn keepalive_ping_does_not_kill_connection() { + let server = TestServer::start(|s| s.route("/echo", echo)).await; + let mut ws = connect(&server.url("echo")).await; + + // Send a keepalive ping, then a real message; the echo must still come back. + timeout( + TIMEOUT, + ws.send(WsMessage::Ping(b"keepalive".to_vec().into())), + ) + .await + .expect("ping send timed out") + .expect("ping send failed"); + + send_binary(&mut ws, b"after-ping").await; + assert_eq!( + recv_binary(&mut ws).await, + b"after-ping", + "connection died after a keepalive ping" + ); + + close_cleanly(ws).await; + server.shutdown().await; +} + +/// A client-initiated close makes the server's `recv` return `Ok(None)` and the handler end. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn normal_close_ends_recv_cleanly() { + let (report_tx, mut report_rx) = mpsc::unbounded_channel::(); + let server = TestServer::start(move |s| { + s.route("/echo", move |mut channel: WsServerChannel| { + let report = report_tx.clone(); + async move { + loop { + match channel.recv().await { + Ok(Some(msg)) => { + if channel.send(msg).await.is_err() { + let _ = report.send(Recv::Error); + return; + } + } + Ok(None) => { + let _ = report.send(Recv::End); + return; + } + Err(_) => { + let _ = report.send(Recv::Error); + return; + } + } + } + } + }) + }) + .await; + + let mut ws = connect(&server.url("echo")).await; + send_binary(&mut ws, b"ping").await; + assert_eq!(recv_binary(&mut ws).await, b"ping"); + close_cleanly(ws).await; + + let ended = timeout(TIMEOUT, report_rx.recv()) + .await + .expect("handler did not finish after close") + .expect("report channel closed"); + assert_eq!( + ended, + Recv::End, + "a client Close should make recv return Ok(None)" + ); + + server.shutdown().await; +} diff --git a/crates/burn-core/Cargo.toml b/crates/burn-core/Cargo.toml new file mode 100644 index 0000000..62dd13b --- /dev/null +++ b/crates/burn-core/Cargo.toml @@ -0,0 +1,133 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science", "no-std", "embedded", "wasm"] +description = "Flexible and Comprehensive Deep Learning Framework in Rust" +documentation = "https://docs.rs/burn-core" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "tensor", "pytorch", "ndarray"] +license.workspace = true +name = "burn-core" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-core" +version.workspace = true + +[lints] +workspace = true + +[features] +default = [ + "std", + "burn-std/default", + "burn-dataset?/default", + "burn-tensor/default", +] +doc = [ + "std", + "dataset", + "audio", + # Doc features + "burn-std/doc", + "burn-dataset/doc", + "burn-tensor/doc", + "dep:regex" +] +tracing = ["burn-std/tracing", "burn-tensor/tracing", "burn-dataset?/tracing"] + +dataset = ["burn-dataset"] + +network = ["burn-std/network"] +sqlite = ["burn-dataset?/sqlite"] +sqlite-bundled = ["burn-dataset?/sqlite-bundled"] +std = [ + "burn-std/std", + "burn-pack/std", + "burn-tensor/std", + "half/std", + "log", + "rand/std", + "serde/std", + "serde_json/std", + "num-traits/std", +] +vision = ["burn-dataset?/vision"] +audio = ["burn-dataset?/audio"] + +# Backends +remote = ["burn-tensor/remote"] +remote-server = ["burn-tensor/remote-server"] +remote-websocket = ["burn-tensor/remote-websocket"] +cuda = ["burn-tensor/cuda"] +rocm = ["burn-tensor/rocm"] +flex = ["burn-tensor/flex"] +ndarray = ["burn-tensor/ndarray"] +tch = ["burn-tensor/tch"] +vulkan = ["burn-tensor/vulkan"] +webgpu = ["burn-tensor/webgpu"] +wgpu = ["burn-tensor/wgpu"] +metal = ["burn-tensor/metal"] +cpu = ["burn-tensor/cpu"] +autodiff = ["burn-tensor/autodiff"] +ir = ["burn-ir"] + +# Backend features +autotune = ["burn-tensor/autotune"] +autotune-checks = ["burn-tensor/autotune-checks"] +fusion = ["burn-tensor/fusion"] + +# Backend extensions +# burn-dispatch features are enabled from burn-tensor, we just re-export the public types +extension = ["burn-tensor/extension", "burn-backend-extension", "burn-dispatch", "burn-backend"] + +[dependencies] + +# ** Please make sure all dependencies support no_std when std is disabled ** + +burn-std = { workspace = true } +burn-dataset = { workspace = true, optional = true } +burn-derive = { workspace = true } +burn-pack = { workspace = true, default-features = false } +burn-tensor = { workspace = true } + +# Backend extensions +burn-backend = { workspace = true, optional = true } +burn-dispatch = { workspace = true, optional = true } +burn-backend-extension = { workspace = true, optional = true } +burn-ir = { workspace = true, optional = true } + +derive-new = { workspace = true } +log = { workspace = true, optional = true } +rand = { workspace = true } + +# The same implementation of HashMap in std but with no_std support (only alloc crate is needed) +hashbrown = { workspace = true, features = ["serde"] } # no_std compatible + +# Serialize Deserialize +serde = { workspace = true, features = ["derive"] } + +ahash = { workspace = true } +half = { workspace = true } +num-traits = { workspace = true } +serde_json = { workspace = true, features = ["alloc"] } #Default enables std +spin = { workspace = true } # Using in place of use std::sync::Mutex when std is disabled +thiserror = { workspace = true, optional = true } +regex = { workspace = true, optional = true } + +[target.'cfg(target_has_atomic = "ptr")'.dependencies] +regex = { workspace = true } + +[target.'cfg(not(target_has_atomic = "ptr"))'.dependencies] +portable-atomic-util = { workspace = true } +portable-atomic = { workspace = true } + +[dev-dependencies] +burn-dataset = { workspace = true, features = ["fake"] } +rstest = { workspace = true } +serial_test = { workspace = true } +tempfile = { workspace = true } + +# Default device for tests +burn-tensor = { workspace = true, features = ["flex"] } + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-core/LICENSE-APACHE b/crates/burn-core/LICENSE-APACHE new file mode 120000 index 0000000..1cd601d --- /dev/null +++ b/crates/burn-core/LICENSE-APACHE @@ -0,0 +1 @@ +../../LICENSE-APACHE \ No newline at end of file diff --git a/crates/burn-core/LICENSE-MIT b/crates/burn-core/LICENSE-MIT new file mode 120000 index 0000000..b2cfbdc --- /dev/null +++ b/crates/burn-core/LICENSE-MIT @@ -0,0 +1 @@ +../../LICENSE-MIT \ No newline at end of file diff --git a/crates/burn-core/README.md b/crates/burn-core/README.md new file mode 100644 index 0000000..5790128 --- /dev/null +++ b/crates/burn-core/README.md @@ -0,0 +1,14 @@ +# Burn Core + +This crate should be used with [burn](https://github.com/tracel-ai/burn). It contains the core +traits and components for building and training deep learning models with Burn. + +[![Current Crates.io Version](https://img.shields.io/crates/v/burn-core.svg)](https://crates.io/crates/burn-core) +[![license](https://shields.io/badge/license-MIT%2FApache--2.0-blue)](https://github.com/tracel-ai/burn-core/blob/master/README.md) + +## Feature Flags + +This crate can be used without the standard library (`#![no_std]`) with `alloc` by disabling the +default `std` feature. + +- `std` - enables the standard library. Enabled by default. diff --git a/crates/burn-core/src/backend.rs b/crates/burn-core/src/backend.rs new file mode 100644 index 0000000..1103c38 --- /dev/null +++ b/crates/burn-core/src/backend.rs @@ -0,0 +1,39 @@ +#[cfg(feature = "ir")] +pub use burn_ir as ir; + +pub use burn_backend::*; +pub use burn_backend_extension::*; + +// Dispatch backend extension types +pub use burn_dispatch::{backend::*, device::*, tensor::*}; +// Re-export backends (e.g., Cuda) +pub use burn_dispatch::backends::*; + +/// A trait to allow mapping custom backend-specific structures into their generic dispatch equivalents. +/// +/// This trait is designed to cooperate with the [`#[backend_extension]`](backend_extension) macro. When an +/// extension operation returns a custom struct containing multiple tensor primitives rather than a single +/// output tensor primitive, this trait provides the mechanism to traverse that struct and recursively wrap +/// each internal field into a [`DispatchTensor`]. +/// +/// Implementations of this trait are generated automatically using `#[derive(ExtensionType)]`. +pub trait ExtensionType { + /// The target struct layout where all internal concrete backend tensors are transformed + /// into [`DispatchTensor`]s. + type Target; + + /// Transforms the internal fields of the struct by applying a backend-specific wrapping closure. + /// + /// # Arguments + /// + /// * `map_kind` - A closure provided by the dispatch macro that knows how to map a backend-agnostic + /// [`BackendTensor`] variant into the correct [`DispatchTensorKind`] variant (e.g., `Wgpu`, `Cuda`, `Cpu`). + /// * `checkpointing` - The active tensor recording/checkpointing strategy to attach to the [`DispatchTensor`]. + /// + /// # Returns + /// + /// A new instance of the struct mapped to the [`Dispatch`] backend. + fn map_type(self, map_kind: F, checkpointing: Option) -> Self::Target + where + F: Fn(BackendTensor) -> DispatchTensorKind; +} diff --git a/crates/burn-core/src/config.rs b/crates/burn-core/src/config.rs new file mode 100644 index 0000000..0b9d5c9 --- /dev/null +++ b/crates/burn-core/src/config.rs @@ -0,0 +1,98 @@ +use alloc::{format, string::String, string::ToString}; +pub use burn_derive::Config; +use core::fmt::Debug; + +/// Configuration IO error. +#[derive(Debug)] +pub enum ConfigError { + /// Invalid format. + InvalidFormat(String), + + /// File not found. + FileNotFound(String), +} + +impl core::fmt::Display for ConfigError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let mut message = "Config error => ".to_string(); + + match self { + Self::InvalidFormat(err) => { + message += format!("Invalid format: {err}").as_str(); + } + Self::FileNotFound(err) => { + message += format!("File not found: {err}").as_str(); + } + }; + + f.write_str(message.as_str()) + } +} + +impl core::error::Error for ConfigError {} + +/// Configuration trait. +pub trait Config: Debug + serde::Serialize + serde::de::DeserializeOwned { + /// Saves the configuration to a file. + /// + /// # Arguments + /// + /// * `file` - File to save the configuration to. + /// + /// # Returns + /// + /// The output of the save operation. + #[cfg(feature = "std")] + fn save>(&self, file: P) -> std::io::Result<()> { + std::fs::write(file, config_to_json(self)) + } + + /// Loads the configuration from a file. + /// + /// # Arguments + /// + /// * `file` - File to load the configuration from. + /// + /// # Returns + /// + /// The loaded configuration. + #[cfg(feature = "std")] + fn load>(file: P) -> Result { + let content = std::fs::read_to_string(file.as_ref()) + .map_err(|_| ConfigError::FileNotFound(file.as_ref().to_string_lossy().to_string()))?; + config_from_str(&content) + } + + /// Loads the configuration from a binary buffer. + /// + /// # Arguments + /// + /// * `data` - Binary buffer to load the configuration from. + /// + /// # Returns + /// + /// The loaded configuration. + fn load_binary(data: &[u8]) -> Result { + let content = core::str::from_utf8(data).map_err(|_| { + ConfigError::InvalidFormat("Could not parse data as utf-8.".to_string()) + })?; + config_from_str(content) + } +} + +/// Converts a configuration to a JSON string. +/// +/// # Arguments +/// +/// * `config` - Configuration to convert. +/// +/// # Returns +/// +/// The JSON string. +pub fn config_to_json(config: &C) -> String { + serde_json::to_string_pretty(config).unwrap() +} + +fn config_from_str(content: &str) -> Result { + serde_json::from_str(content).map_err(|err| ConfigError::InvalidFormat(format!("{err}"))) +} diff --git a/crates/burn-core/src/data/dataloader/base.rs b/crates/burn-core/src/data/dataloader/base.rs new file mode 100644 index 0000000..53b32f9 --- /dev/null +++ b/crates/burn-core/src/data/dataloader/base.rs @@ -0,0 +1,52 @@ +use burn_tensor::Device; + +pub use crate::data::dataset::{Dataset, DatasetIterator}; +use core::iter::Iterator; +use std::sync::Arc; + +/// A progress struct that can be used to track the progress of a data loader. +#[derive(new, Clone, Debug)] +pub struct Progress { + /// The number of items that have been processed. + pub items_processed: usize, + + /// The total number of items that need to be processed. + pub items_total: usize, + + /// The unit of measurement for the progress (e.g., "items", "batches"). + pub unit: Option, +} + +/// A data loader iterator that can be used to iterate over a data loader. +pub trait DataLoaderIterator: Iterator { + /// Returns the progress of the data loader. + fn progress(&self) -> Progress; +} + +/// A data loader that can be used to iterate over a dataset. +pub trait DataLoader: Send + Sync { + /// Returns a boxed [iterator](DataLoaderIterator) to iterate over the data loader. + fn iter<'a>(&'a self) -> Box + 'a>; + + /// The number of items (not the number of batches nor the number of iterations), + /// corresponding to the items_total of the progress returned by the iterator. + fn num_items(&self) -> usize; + + /// Move the data loader to the given device, ensuring the batches are assigned to the correct device. + fn to_device(&self, device: &Device) -> Arc>; + + /// Returns a new data loader containing a subset of the data. + /// + /// The subset includes items from `start` (inclusive) to `end` (exclusive), + /// preserving the batch size and ordering of the original data loader. + /// + /// # Arguments + /// + /// * `start` - The starting index of the subset (inclusive). + /// * `end` - The ending index of the subset (exclusive). + /// + /// # Returns + /// + /// A boxed [`DataLoader`] instance containing only the specified range. + fn slice(&self, start: usize, end: usize) -> Arc>; +} diff --git a/crates/burn-core/src/data/dataloader/batch.rs b/crates/burn-core/src/data/dataloader/batch.rs new file mode 100644 index 0000000..b4b5501 --- /dev/null +++ b/crates/burn-core/src/data/dataloader/batch.rs @@ -0,0 +1,260 @@ +use super::{BatchStrategy, DataLoader, DataLoaderIterator, Progress, batcher::Batcher}; +use burn_dataset::{ + Dataset, + transform::{PartialDataset, ShuffledDataset}, +}; +use burn_tensor::Device; +use rand::SeedableRng; +use std::ops::DerefMut; +use std::sync::Arc; + +/// A data loader that can be used to iterate over a dataset in batches. +pub struct BatchDataLoader { + strategy: Box>, + dataset: Arc>, + batcher: Arc>, + device: Device, + rng: Option>>, +} + +impl Clone for BatchDataLoader { + fn clone(&self) -> Self { + Self { + strategy: self.strategy.clone_dyn(), + dataset: self.dataset.clone(), + batcher: self.batcher.clone(), + device: self.device.clone(), + rng: self.rng.clone(), + } + } +} + +impl BatchDataLoader { + /// Creates a new batch data loader. + /// + /// # Arguments + /// + /// * `strategy` - The batch strategy. + /// * `dataset` - The dataset. + /// * `batcher` - The batcher. + /// * `device` - The device to use when loading a batch. + /// * `rng` - The rng determining if the dataset is shuffled each time a dataloader + /// iterator is created. + /// + /// # Returns + /// + /// The batch data loader. + pub fn new( + strategy: Box>, + dataset: Arc>, + batcher: Arc>, + device: Device, + rng: Option, + ) -> Self { + Self { + strategy, + dataset, + batcher, + device, + rng: rng.map(|rng| Arc::new(spin::Mutex::new(rng))), + } + } +} + +/// A data loader iterator that can be used to iterate over a data loader. +struct BatchDataloaderIterator { + current_index: usize, + strategy: Box>, + dataset: Arc>, + batcher: Arc>, + device: Device, +} + +impl DataLoader for BatchDataLoader +where + I: Send + Sync + Clone + 'static, + O: Send + 'static, +{ + fn iter<'a>(&'a self) -> Box + 'a> { + // When starting a new iteration, we first check if the dataloader was created with an rng, + // implying that we should shuffle the dataset beforehand, while advancing the current + // rng to ensure that each new iteration shuffles the dataset differently. + let dataset = match &self.rng { + Some(rng) => Arc::new(ShuffledDataset::new( + self.dataset.clone(), + rng.lock().deref_mut(), + )), + None => self.dataset.clone(), + }; + Box::new(BatchDataloaderIterator::new( + self.strategy.clone_dyn(), + dataset, + self.batcher.clone(), + self.device.clone(), + )) + } + + fn num_items(&self) -> usize { + self.dataset.len() + } + + fn to_device(&self, device: &Device) -> Arc> { + let rng = self.rng.as_ref().map(|rng| { + let mut rng = rng.lock(); + rng.fork() + }); + Arc::new(Self::new( + self.strategy.clone_dyn(), + self.dataset.clone(), + self.batcher.clone(), + device.clone(), + rng, + )) + } + + fn slice(&self, start: usize, end: usize) -> Arc> { + let rng = self.rng.as_ref().map(|rng| { + let mut rng = rng.lock(); + rng.fork() + }); + let dataloader = Self::new( + self.strategy.clone_dyn(), + Arc::new(PartialDataset::new(self.dataset.clone(), start, end)), + self.batcher.clone(), + self.device.clone(), + rng, + ); + Arc::new(dataloader) + } +} + +impl BatchDataloaderIterator { + /// Creates a new batch data loader iterator. + /// + /// # Arguments + /// + /// * `strategy` - The batch strategy. + /// * `dataset` - The dataset. + /// * `batcher` - The batcher. + /// * `device` - The device to use when loading a batch. + /// + /// # Returns + /// + /// The batch data loader iterator. + pub fn new( + strategy: Box>, + dataset: Arc>, + batcher: Arc>, + device: Device, + ) -> Self { + BatchDataloaderIterator { + current_index: 0, + strategy, + dataset, + batcher, + device, + } + } +} + +impl Iterator for BatchDataloaderIterator { + type Item = O; + + fn next(&mut self) -> Option { + while let Some(item) = self.dataset.get(self.current_index) { + self.current_index += 1; + self.strategy.add(item); + + if let Some(items) = self.strategy.batch(false) { + return Some(self.batcher.batch(items, &self.device)); + } + } + + if let Some(items) = self.strategy.batch(true) { + return Some(self.batcher.batch(items, &self.device)); + } + + None + } +} + +impl DataLoaderIterator for BatchDataloaderIterator { + fn progress(&self) -> Progress { + let unit: Option = Some("items".to_string()); + + Progress::new(self.current_index, self.dataset.len(), unit) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use super::*; + use crate::data::dataloader::FixBatchStrategy; + use crate::data::dataloader::batcher::TestBatcher; + use crate::data::dataset::FakeDataset; + + #[test] + fn test_batch_dataloader() { + let batcher = Arc::new(TestBatcher::new()); + let dataset = Arc::new(FakeDataset::::new(27)); + let dataloader = BatchDataLoader::new( + Box::new(FixBatchStrategy::new(5)), + dataset.clone(), + batcher, + Default::default(), + None, + ); + + let mut items_dataset = HashSet::new(); + let mut items_dataloader = HashSet::new(); + + for item in dataset.iter() { + items_dataset.insert(item); + } + + for items in dataloader.iter() { + for item in items { + items_dataloader.insert(item); + } + } + + assert_eq!(items_dataset, items_dataloader); + } + + #[test] + fn test_batch_dataloader_slice() { + let batcher = Arc::new(TestBatcher::new()); + let dataset = Arc::new(FakeDataset::::new(27)); + let dataloader = BatchDataLoader::new( + Box::new(FixBatchStrategy::new(5)), + dataset.clone(), + batcher, + Default::default(), + None, + ); + let dataloader_slice = dataloader.slice(5, 15); + + let mut items_dataloader = HashSet::new(); + let mut items_dataloader_slice = HashSet::new(); + + let mut idx = 0; + for items in dataloader.iter() { + for item in items { + if (5..15).contains(&idx) { + items_dataloader.insert(item); + } + idx += 1; + } + } + + for items in dataloader_slice.iter() { + for item in items { + items_dataloader_slice.insert(item); + } + } + + assert_eq!(items_dataloader, items_dataloader_slice); + } +} diff --git a/crates/burn-core/src/data/dataloader/batcher.rs b/crates/burn-core/src/data/dataloader/batcher.rs new file mode 100644 index 0000000..0dcd0c2 --- /dev/null +++ b/crates/burn-core/src/data/dataloader/batcher.rs @@ -0,0 +1,28 @@ +use burn_tensor::Device; + +/// A trait for batching items of type `I` into items of type `O`. +pub trait Batcher: Send + Sync { + /// Batches the given items on the specified device. + /// + /// # Arguments + /// + /// * `items` - The items to batch. + /// * `device` - The backend device to use. + /// + /// # Returns + /// + /// The batched items. + fn batch(&self, items: Vec, device: &Device) -> O; +} + +/// Test batcher +#[cfg(test)] +#[derive(new, Clone)] +pub struct TestBatcher; + +#[cfg(test)] +impl Batcher> for TestBatcher { + fn batch(&self, items: Vec, _device: &Device) -> Vec { + items + } +} diff --git a/crates/burn-core/src/data/dataloader/builder.rs b/crates/burn-core/src/data/dataloader/builder.rs new file mode 100644 index 0000000..ce74c72 --- /dev/null +++ b/crates/burn-core/src/data/dataloader/builder.rs @@ -0,0 +1,240 @@ +use super::{ + BatchDataLoader, BatchStrategy, DataLoader, FixBatchStrategy, MultiThreadDataLoader, + batcher::Batcher, +}; +use burn_dataset::Dataset; +use burn_tensor::Device; +use rand::{SeedableRng, rngs::StdRng}; +use std::sync::Arc; + +/// A builder for data loaders. +pub struct DataLoaderBuilder { + strategy: Option>>, + batcher: Arc>, + num_threads: Option, + shuffle: Option, + device: Option, +} + +impl DataLoaderBuilder +where + I: Send + Sync + Clone + std::fmt::Debug + 'static, + O: Send + Clone + std::fmt::Debug + 'static, +{ + /// Creates a new data loader builder. + /// + /// # Arguments + /// + /// * `batcher` - The batcher. + /// + /// # Returns + /// + /// The data loader builder. + pub fn new(batcher: Bt) -> Self + where + Bt: Batcher + 'static, + { + Self { + batcher: Arc::new(batcher), + strategy: None, + num_threads: None, + shuffle: None, + device: None, + } + } + + /// Sets the batch size to a fix number. + /// + /// The [fix batch strategy](FixBatchStrategy) will be used. + /// + /// # Arguments + /// + /// * `batch_size` - The batch size. + /// + /// # Returns + /// + /// The data loader builder. + pub fn batch_size(mut self, batch_size: usize) -> Self { + self.strategy = Some(Box::new(FixBatchStrategy::new(batch_size))); + self + } + + /// Sets the seed for shuffling. + /// + /// Each time the dataloader starts a new iteration, the dataset will be shuffled. + /// + /// # Arguments + /// + /// * `seed` - The seed. + /// + /// # Returns + /// + /// The data loader builder. + pub fn shuffle(mut self, seed: u64) -> Self { + self.shuffle = Some(seed); + self + } + + /// Sets the number of workers. + /// + /// - `Some(0)` or `None`: the dataloader will run without work threads. + /// - `Some(n); n > 0`: the dataloader will run with `n` background threads. + /// + /// A 1-worker threaded dataloader will run loads in a background thread, + /// while a 0-worker threaded dataloader will run loads in the main thread. + /// + /// # Arguments + /// + /// * `num_workers` - The number of workers. + /// + /// # Returns + /// + /// The data loader builder. + pub fn num_workers(mut self, num_workers: usize) -> Self { + self.num_threads = Some(num_workers); + self + } + + /// Sets the data loader device. + /// + /// # Arguments + /// + /// * `device` - The device to use when loading a batch. + /// + /// # Returns + /// + /// The data loader builder. + pub fn set_device(mut self, device: Device) -> Self { + self.device = Some(device); + self + } + + /// Builds the data loader. + /// + /// # Arguments + /// + /// * `dataset` - The dataset. + /// + /// # Returns + /// + /// The data loader. + pub fn build(self, dataset: D) -> Arc> + where + D: Dataset + 'static, + { + let dataset = Arc::new(dataset); + + let device = self.device.unwrap_or_default(); + let rng = self.shuffle.map(StdRng::seed_from_u64); + let strategy = match self.strategy { + Some(strategy) => strategy, + None => Box::new(FixBatchStrategy::new(1)), + }; + + if let Some(num_threads) = self.num_threads + && num_threads > 0 + { + return Arc::new(MultiThreadDataLoader::new( + strategy, + dataset, + self.batcher, + num_threads, + device, + rng, + )); + } + + Arc::new(BatchDataLoader::new( + strategy, + dataset, + self.batcher, + device, + rng, + )) + } +} + +#[cfg(test)] +mod tests { + #[cfg(test)] + use burn_tensor::Device; + + use super::*; + use crate::data::dataset::FakeDataset; + + #[derive(new, Clone)] + struct TestBatcherDevice; + + #[cfg(test)] + impl Batcher for TestBatcherDevice { + fn batch(&self, _items: Vec, device: &Device) -> Device { + device.clone() + } + } + + #[test] + fn test_dataloader_no_workers() { + let default_device = Device::default(); + let dataloader = DataLoaderBuilder::new(TestBatcherDevice::new()) + .batch_size(1) + .build(FakeDataset::::new(9)); + + assert_eq!(dataloader.num_items(), 9); + + for device in dataloader.iter() { + assert_eq!(device, default_device) + } + } + + #[test] + fn test_dataloader_default_device() { + let default_device = Device::default(); + let dataloader = DataLoaderBuilder::new(TestBatcherDevice::new()) + .batch_size(1) + .num_workers(1) + .build(FakeDataset::::new(9)); + + assert_eq!(dataloader.num_items(), 9); + + for device in dataloader.iter() { + assert_eq!(device, default_device) + } + } + + #[test] + fn test_dataloader_slice_multi_device() { + let dataloader = DataLoaderBuilder::new(TestBatcherDevice::new()) + .batch_size(1) + .num_workers(1) + .build(FakeDataset::::new(11)); + + #[cfg(all(test, not(feature = "tch"), not(feature = "cuda")))] + // Only one device exists... + let (device1, device2) = (Device::flex(), Device::flex()); + + #[cfg(all(test, feature = "tch"))] + let (device1, device2) = (Device::libtorch_cuda(0), Device::libtorch_cuda(1)); + + #[cfg(all(test, feature = "cuda"))] + let (device1, device2) = (Device::cuda(0), Device::cuda(1)); + + assert_eq!(dataloader.num_items(), 11); + let dataloader_1 = dataloader.slice(0, 5).to_device(&device1); + let dataloader_2 = dataloader.slice(5, 11).to_device(&device2); + + assert_eq!(dataloader_1.num_items(), 5); + assert_eq!(dataloader_2.num_items(), 6); + + let (mut iterator_1, mut iterator_2) = (dataloader_1.iter(), dataloader_2.iter()); + + for _ in 0..5 { + assert_eq!(iterator_1.next().as_ref(), Some(&device1)); + assert_eq!(iterator_2.next().as_ref(), Some(&device2)); + } + + assert_eq!(iterator_1.next(), None); + // For uneven split, the last dataloader (partial dataset) will have the remaining item + assert_eq!(iterator_2.next(), Some(device2)); + assert_eq!(iterator_2.next(), None); + } +} diff --git a/crates/burn-core/src/data/dataloader/mod.rs b/crates/burn-core/src/data/dataloader/mod.rs new file mode 100644 index 0000000..52556af --- /dev/null +++ b/crates/burn-core/src/data/dataloader/mod.rs @@ -0,0 +1,16 @@ +mod base; +mod batch; +mod builder; +mod multithread; +mod strategy; + +/// Module for batching items. +pub mod batcher; +/// Module to split a dataloader. +pub mod split; + +pub use base::*; +pub use batch::*; +pub use builder::*; +pub use multithread::*; +pub use strategy::*; diff --git a/crates/burn-core/src/data/dataloader/multithread.rs b/crates/burn-core/src/data/dataloader/multithread.rs new file mode 100644 index 0000000..70ca69d --- /dev/null +++ b/crates/burn-core/src/data/dataloader/multithread.rs @@ -0,0 +1,553 @@ +use burn_dataset::Dataset; +use burn_dataset::transform::PartialDataset; +use burn_tensor::Device; +use rand::distr::{Distribution, StandardUniform}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; + +use super::batcher::Batcher; +use super::{BatchDataLoader, BatchStrategy, DataLoader, DataLoaderIterator, Progress}; +use std::sync::{Arc, OnceLock, mpsc, mpsc::SyncSender}; +use std::thread; + +const MAX_QUEUED_ITEMS: usize = 100; + +type RngSeed = ::Seed; + +/// A multi-threaded data loader that can be used to iterate over a dataset. +pub struct MultiThreadDataLoader { + // Configuration parameters needed for initialization + strategy: Box>, + dataset: Arc>, + batcher: Arc>, + device: Device, + seed: Option, + num_threads: usize, + + // The lazily initialized data loaders + dataloaders: OnceLock>>, + + // Spawned once and reused across every `iter()` call so each worker keeps a + // stable CubeCL stream (and its memory pool) instead of leaking one per epoch (#4792). + workers: OnceLock>, +} + +/// A message that can be sent between threads. +#[derive(Debug)] +pub enum Message { + /// A batch of items. + Batch(usize, O, Progress), + + /// The thread is done. + Done, +} + +struct MultiThreadsDataloaderIterator { + num_done: usize, + num_workers: usize, + receiver: mpsc::Receiver>, + progresses: Vec, +} + +/// Per-epoch channel a worker streams its batches into; handed to the worker to start a pass. +type WorkerCommand = SyncSender>; + +struct WorkerPool { + /// One command channel per worker; sending a per-epoch sender starts a pass. + senders: Vec>>, + handles: Vec>, + item_counts: Vec, +} + +impl Drop for WorkerPool { + fn drop(&mut self) { + // Dropping the senders makes each worker's `recv()` return Err, ending its loop. + self.senders.clear(); + for handle in self.handles.drain(..) { + let _ = handle.join(); + } + } +} + +impl MultiThreadDataLoader +where + I: Send + Sync + Clone + 'static, + O: Send + 'static, +{ + /// Creates a new multi-threaded batch data loader. + /// + /// # Arguments + /// + /// * `strategy` - The batch strategy. + /// * `dataset` - The dataset. + /// * `batcher` - The batcher. + /// * `num_threads` - The number of threads. + /// * `device` - The device to use when loading a batch. + /// * `rng` - The rng determining if the dataset is shuffled each time a dataloader + /// iterator is created. + /// + /// # Returns + /// + /// The multi-threaded batch data loader. + pub fn new( + strategy: Box>, + dataset: Arc>, + batcher: Arc>, + num_threads: usize, + device: Device, + rng: Option, + ) -> Self { + let mut seed = None; + if let Some(mut rng) = rng { + // RNG stream splitting (not state cloning): derive a new seed from the RNG's output. + // This is exactly what `rng.fork()` does. + let mut s = RngSeed::default(); + rng.fill_bytes(&mut s); + + seed = Some(s); + } + Self::from_seed(strategy, dataset, batcher, num_threads, device, seed) + } + + fn from_seed( + strategy: Box>, + dataset: Arc>, + batcher: Arc>, + num_threads: usize, + device: Device, + seed: Option, + ) -> Self { + Self { + strategy, + dataset, + batcher, + num_threads, + device, + seed, + dataloaders: OnceLock::new(), + workers: OnceLock::new(), + } + } + + /// Force initialization if needed. + fn initialize(&self) -> &[BatchDataLoader] { + self.dataloaders + .get_or_init(|| { + let mut dataset = self.dataset.clone(); + if let Some(seed) = self.seed.as_ref() { + // Pre-shuffle the dataset before split if shuffle is enabled. + // This ensures that each thread gets a uniform random sample of the dataset. + let mut rng = StdRng::from_seed(*seed); + dataset = Arc::new(burn_dataset::transform::ShuffledDataset::new( + dataset, &mut rng, + )); + } + + let datasets = match self.strategy.batch_size() { + Some(batch_size) => { + PartialDataset::split_chunks(dataset, self.num_threads, batch_size) + } + None => PartialDataset::split(dataset, self.num_threads), + }; + + // Create more rngs from the first one, one for each new dataloader. + let mut rng = self.seed.map(StdRng::from_seed); + let rngs = (0..self.num_threads).map(|_| { + rng.as_mut().map(|rng| { + StdRng::seed_from_u64(Distribution::sample(&StandardUniform, rng)) + }) + }); + + datasets + .into_iter() + .zip(rngs) + .map(|(dataset, rng)| { + let strategy = self.strategy.clone_dyn(); + BatchDataLoader::new( + strategy, + Arc::new(dataset), + self.batcher.clone(), + self.device.clone(), + rng, + ) + }) + .collect() + }) + .as_ref() + } + + /// Lazily spawns the persistent worker pool (once) and returns it. + fn workers(&self) -> &WorkerPool { + self.workers.get_or_init(|| { + let dataloaders = self.initialize(); + let item_counts: Vec = dataloaders.iter().map(|d| d.num_items()).collect(); + + let mut senders = Vec::with_capacity(dataloaders.len()); + let mut handles = Vec::with_capacity(dataloaders.len()); + + for (index, dataloader) in dataloaders.iter().enumerate() { + let dataloader = dataloader.clone(); + let (command_sender, command_receiver) = mpsc::channel::>(); + + let handle = thread::Builder::new() + .name(std::format!("dataloader-{index}")) + .spawn(move || { + while let Ok(sender) = command_receiver.recv() { + let mut iterator = dataloader.iter(); + while let Some(item) = iterator.next() { + let progress = iterator.progress(); + + // Consumer dropped the receiver: abandon this pass, + // don't terminate the worker. + if sender.send(Message::Batch(index, item, progress)).is_err() { + break; + } + } + sender.send(Message::Done).ok(); + } + }) + .unwrap(); + + senders.push(command_sender); + handles.push(handle); + } + + WorkerPool { + senders, + handles, + item_counts, + } + }) + } +} + +impl DataLoader for MultiThreadDataLoader +where + I: Send + Sync + Clone + 'static, + O: Send + 'static + std::fmt::Debug, +{ + fn iter<'a>(&'a self) -> Box + 'a> { + let workers = self.workers(); + + let (sender, receiver) = mpsc::sync_channel::>(MAX_QUEUED_ITEMS); + let unit: Option = Some("items".to_string()); + + let mut progresses = Vec::with_capacity(workers.senders.len()); + for (command_sender, &num_items) in workers.senders.iter().zip(workers.item_counts.iter()) { + progresses.push(Progress::new(0, num_items, unit.clone())); + command_sender + .send(sender.clone()) + .expect("Dataloader worker thread should be alive"); + } + let num_workers = workers.senders.len(); + + // Drop our sender so the channel disconnects once every worker is done. + drop(sender); + + Box::new(MultiThreadsDataloaderIterator::new( + receiver, + num_workers, + progresses, + )) + } + + fn num_items(&self) -> usize { + // For num_items, we can directly use the dataset size without + // necessarily initializing the full loader + self.dataset.len() + } + + fn to_device(&self, device: &Device) -> Arc> { + Arc::new(Self::from_seed( + self.strategy.clone_dyn(), + self.dataset.clone(), + self.batcher.clone(), + self.num_threads, + device.clone(), + self.seed, + )) + } + + fn slice(&self, start: usize, end: usize) -> Arc> { + let dataloader = Self::from_seed( + self.strategy.clone_dyn(), + Arc::new(PartialDataset::new(self.dataset.clone(), start, end)), + self.batcher.clone(), + self.num_threads, + self.device.clone(), + self.seed, + ); + Arc::new(dataloader) + } +} + +impl MultiThreadsDataloaderIterator { + pub fn new( + receiver: mpsc::Receiver>, + num_workers: usize, + progresses: Vec, + ) -> Self { + MultiThreadsDataloaderIterator { + num_done: 0, + num_workers, + receiver, + progresses, + } + } +} +impl DataLoaderIterator for MultiThreadsDataloaderIterator { + fn progress(&self) -> Progress { + let mut items_total = 0; + let mut items_processed = 0; + let unit: Option = Some("items".to_string()); + + for progress in self.progresses.iter() { + items_total += progress.items_total; + items_processed += progress.items_processed; + } + + Progress::new(items_processed, items_total, unit) + } +} + +impl Iterator for MultiThreadsDataloaderIterator { + type Item = O; + + fn next(&mut self) -> Option { + if self.num_workers == 0 { + return None; + } + + loop { + match self.receiver.recv() { + Ok(Message::Batch(index, item, progress)) => { + if let Some(current) = self.progresses.get_mut(index) { + *current = progress; + } + return Some(item); + } + Ok(Message::Done) => { + self.num_done += 1; + if self.num_done == self.num_workers { + // Workers stay alive for the next epoch; nothing to join. + return None; + } + } + // Every worker dropped its sender (the pool is shutting down). + Err(_) => return None, + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::data::dataloader::FixBatchStrategy; + use crate::data::dataloader::batcher::TestBatcher; + use crate::data::dataset::FakeDataset; + use burn_dataset::InMemDataset; + use std::collections::HashSet; + + #[test] + fn test_multi_thread_batch_dataloader() { + let batcher = Arc::new(TestBatcher::new()); + let dataset = Arc::new(FakeDataset::::new(27)); + let dataloader_single_thread = BatchDataLoader::new( + Box::new(FixBatchStrategy::new(5)), + dataset.clone(), + batcher.clone(), + Default::default(), + None, + ); + let dataloader_multi_thread = MultiThreadDataLoader::new( + Box::new(FixBatchStrategy::new(5)), + dataset, + batcher, + 4, + Default::default(), + None, + ); + + let mut items_single_thread = HashSet::new(); + let mut items_multi_thread = HashSet::new(); + + for items in dataloader_single_thread.iter() { + for item in items { + items_single_thread.insert(item); + } + } + + for items in dataloader_multi_thread.iter() { + for item in items { + items_multi_thread.insert(item); + } + } + + assert_eq!(items_single_thread, items_multi_thread); + } + + #[test] + fn test_multi_thread_batch_dataloader_shuffle() { + let num_classes = 2; + let class_size = 100; + let batch_size = 10; + + // Items is a deliberately ordered dataset. + let mut items = Vec::new(); + for class in 0..num_classes { + items.extend(vec![class; class_size]); + } + + { + // Unshuffled multithreaded loader + let dataset = Arc::new(InMemDataset::new(items.clone())); + let batcher = Arc::new(TestBatcher::new()); + + let loader = MultiThreadDataLoader::new( + Box::new(FixBatchStrategy::new(batch_size)), + dataset, + batcher, + num_classes, + Default::default(), + // No rng means no shuffling. + None, + ); + + for batch in loader.iter() { + let mut batch_items = HashSet::new(); + for item in batch { + batch_items.insert(item); + } + + // Since the dataset is not shuffled, we expect each batch to contain the same item. + assert_eq!(batch_items.len(), 1); + } + } + + { + // Shuffled multithreaded loader + let dataset = Arc::new(InMemDataset::new(items.clone())); + let batcher = Arc::new(TestBatcher::new()); + + let loader = MultiThreadDataLoader::new( + Box::new(FixBatchStrategy::new(batch_size)), + dataset.clone(), + batcher.clone(), + num_classes, + Default::default(), + // The rng enables shuffling. + Some(StdRng::seed_from_u64(42)), + ); + + for batch in loader.iter() { + let mut batch_items = HashSet::new(); + for item in batch { + batch_items.insert(item); + } + + // Since the dataset is shuffled, we expect to see all items. + assert_eq!(batch_items.len(), num_classes); + } + } + } + + #[test] + fn test_multi_thread_batch_dataloader_incomplete_batches() { + let batcher = Arc::new(TestBatcher::new()); + let dataset = Arc::new(FakeDataset::::new(27)); + let dataloader_single_thread = BatchDataLoader::new( + Box::new(FixBatchStrategy::new(5)), + dataset.clone(), + batcher.clone(), + Default::default(), + None, + ); + let dataloader_multi_thread = MultiThreadDataLoader::new( + Box::new(FixBatchStrategy::new(5)), + dataset, + batcher, + 4, + Default::default(), + None, + ); + + let mut items_single_thread = HashSet::new(); + let mut items_multi_thread = HashSet::new(); + + let mut single_thread_cnt = 0; + let mut multi_thread_cnt = 0; + for items in dataloader_single_thread.iter() { + items_single_thread.insert(items); + single_thread_cnt += 1; + } + + for items in dataloader_multi_thread.iter() { + items_multi_thread.insert(items); + multi_thread_cnt += 1; + } + + assert_eq!(single_thread_cnt, multi_thread_cnt); + assert_eq!(items_single_thread, items_multi_thread); + } + + // Iterating the same loader over several epochs must keep yielding the full dataset (#4792). + #[test] + fn test_multi_thread_batch_dataloader_multiple_epochs() { + let batcher = Arc::new(TestBatcher::new()); + let dataset = Arc::new(FakeDataset::::new(27)); + + let expected: HashSet<_> = dataset.iter().collect(); + + let dataloader = MultiThreadDataLoader::new( + Box::new(FixBatchStrategy::new(5)), + dataset, + batcher, + 4, + Default::default(), + None, + ); + + for _epoch in 0..3 { + let mut items = HashSet::new(); + for batch in dataloader.iter() { + for item in batch { + items.insert(item); + } + } + assert_eq!(items, expected); + } + } + + // Dropping an iterator early must not kill the workers; the next pass still yields everything. + #[test] + fn test_multi_thread_batch_dataloader_resumes_after_early_drop() { + let batcher = Arc::new(TestBatcher::new()); + let dataset = Arc::new(FakeDataset::::new(27)); + + let expected: HashSet<_> = dataset.iter().collect(); + + let dataloader = MultiThreadDataLoader::new( + Box::new(FixBatchStrategy::new(5)), + dataset, + batcher, + 4, + Default::default(), + None, + ); + + // Consume a single batch then drop the iterator early. + { + let mut iterator = dataloader.iter(); + let _ = iterator.next(); + } + + let mut items = HashSet::new(); + for batch in dataloader.iter() { + for item in batch { + items.insert(item); + } + } + assert_eq!(items, expected); + } +} diff --git a/crates/burn-core/src/data/dataloader/split.rs b/crates/burn-core/src/data/dataloader/split.rs new file mode 100644 index 0000000..eda994c --- /dev/null +++ b/crates/burn-core/src/data/dataloader/split.rs @@ -0,0 +1,115 @@ +use std::sync::Arc; + +use burn_tensor::Device; + +use super::DataLoader; + +/// Splits a dataloader into multiple partial dataloaders (one per device). +pub fn split_dataloader( + dataloader: Arc>, + devices: &[Device], +) -> Vec>> { + let num_splits = devices.len(); + if num_splits > 1 { + let num_items = dataloader.num_items(); + let mut dataloaders = Vec::with_capacity(num_splits); + + let mut start = 0; + let step = num_items / num_splits; + for (i, device) in devices.iter().enumerate() { + let end = if i == (num_splits - 1) { + num_items + } else { + start + step + }; + let dataloader = dataloader.slice(start, end).to_device(device); + dataloaders.push(dataloader); + start = end; + } + dataloaders + } else { + vec![dataloader] + } +} + +#[cfg(test)] +mod tests { + use burn_tensor::Device; + use std::collections::HashSet; + + use super::*; + use crate::data::dataloader::batcher::Batcher; + use crate::data::dataloader::{BatchDataLoader, FixBatchStrategy}; + use crate::data::dataset::FakeDataset; + + #[test] + fn test_split_batch_dataloader() { + #[derive(new, Clone)] + pub struct TestBatcher; + + #[cfg(test)] + impl Batcher, Device)> for TestBatcher { + fn batch(&self, items: Vec, device: &Device) -> (Vec, Device) { + (items, device.clone()) + } + } + + let batcher = Arc::new(TestBatcher::new()); + let dataset = Arc::new(FakeDataset::::new(11)); + + #[allow(clippy::arc_with_non_send_sync)] + let dataloader = Arc::new(BatchDataLoader::new( + Box::new(FixBatchStrategy::new(5)), + dataset.clone(), + batcher, + Default::default(), + None, + )); + + #[cfg(all(test, not(feature = "tch"), not(feature = "cuda")))] + // Only one device exists... + let (device1, device2) = (Device::flex(), Device::flex()); + + #[cfg(all(test, feature = "tch"))] + let (device1, device2) = (Device::libtorch_cuda(0), Device::libtorch_cuda(1)); + + #[cfg(all(test, feature = "cuda"))] + let (device1, device2) = (Device::cuda(0), Device::cuda(1)); + + let dataloaders = split_dataloader(dataloader.clone(), &[device1.clone(), device2.clone()]); + + assert_eq!(dataloaders.len(), 2); + + let [dataloader_1, dataloader_2] = match dataloaders.try_into() { + Ok(arr) => arr, + Err(_) => unreachable!(), + }; + assert_eq!(dataloader_1.num_items(), 5); + assert_eq!(dataloader_2.num_items(), 6); + + let mut items_dataloader = HashSet::new(); + let mut items_dataloader_split = HashSet::new(); + + for (items, _device) in dataloader.iter() { + for item in items { + items_dataloader.insert(item); + } + } + + for (items, device) in dataloader_1.iter() { + assert_eq!(&device, &device1); + for item in items { + items_dataloader_split.insert(item); + } + } + + for (items, device) in dataloader_2.iter() { + assert_eq!(&device, &device2); + for item in items { + items_dataloader_split.insert(item); + } + } + + assert_eq!(items_dataloader, items_dataloader_split); + } +} diff --git a/crates/burn-core/src/data/dataloader/strategy.rs b/crates/burn-core/src/data/dataloader/strategy.rs new file mode 100644 index 0000000..d6e3413 --- /dev/null +++ b/crates/burn-core/src/data/dataloader/strategy.rs @@ -0,0 +1,87 @@ +/// A strategy to batch items. +pub trait BatchStrategy: Send + Sync { + /// Adds an item to the strategy. + /// + /// # Arguments + /// + /// * `item` - The item to add. + fn add(&mut self, item: I); + + /// Batches the items. + /// + /// # Arguments + /// + /// * `force` - Whether to force batching. + /// + /// # Returns + /// + /// The batched items. + fn batch(&mut self, force: bool) -> Option>; + + /// Creates a new strategy of the same type. + /// + /// # Returns + /// + /// The new strategy. + fn clone_dyn(&self) -> Box>; + + /// Returns the expected batch size for this strategy. + /// + /// # Returns + /// + /// The batch size, or None if the strategy doesn't have a fixed batch size. + fn batch_size(&self) -> Option; +} + +/// A strategy to batch items with a fixed batch size. +pub struct FixBatchStrategy { + items: Vec, + batch_size: usize, +} + +impl FixBatchStrategy { + /// Creates a new strategy to batch items with a fixed batch size. + /// + /// # Arguments + /// + /// * `batch_size` - The batch size. + /// + /// # Returns + /// + /// The strategy. + pub fn new(batch_size: usize) -> Self { + FixBatchStrategy { + items: Vec::with_capacity(batch_size), + batch_size, + } + } +} + +impl BatchStrategy for FixBatchStrategy { + fn add(&mut self, item: I) { + self.items.push(item); + } + + fn batch(&mut self, force: bool) -> Option> { + if self.items.len() < self.batch_size && !force { + return None; + } + + let mut items = Vec::with_capacity(self.batch_size); + std::mem::swap(&mut items, &mut self.items); + + if items.is_empty() { + return None; + } + + Some(items) + } + + fn clone_dyn(&self) -> Box> { + Box::new(Self::new(self.batch_size)) + } + + fn batch_size(&self) -> Option { + Some(self.batch_size) + } +} diff --git a/crates/burn-core/src/data/mod.rs b/crates/burn-core/src/data/mod.rs new file mode 100644 index 0000000..8881032 --- /dev/null +++ b/crates/burn-core/src/data/mod.rs @@ -0,0 +1,15 @@ +/// Dataloader module. +#[cfg(feature = "dataset")] +pub mod dataloader; + +/// Dataset module. +#[cfg(feature = "dataset")] +pub mod dataset { + pub use burn_dataset::*; +} + +/// Network module. +#[cfg(feature = "network")] +pub mod network { + pub use burn_std::network::*; +} diff --git a/crates/burn-core/src/lib.rs b/crates/burn-core/src/lib.rs new file mode 100644 index 0000000..79f2a06 --- /dev/null +++ b/crates/burn-core/src/lib.rs @@ -0,0 +1,93 @@ +#![cfg_attr(not(feature = "std"), no_std)] +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] +#![recursion_limit = "135"] + +//! The core crate of Burn. + +// `derive_new` provides the `#[derive(new)]` macro used across the crate; the lint mistakenly +// reports the `#[macro_use]` as unused. +#[allow(unused_imports)] +#[macro_use] +extern crate derive_new; + +/// Re-export serde for proc macros. +pub use serde; + +/// The configuration module. +pub mod config; + +/// Data module. +#[cfg(feature = "std")] +pub mod data; + +/// Module for the neural network module. +pub mod module; + +/// Module for saving and loading module/optimizer state in the burnpack format. +pub mod store; + +/// Module for the tensor. +pub mod tensor; +// Tensor at root: `burn::Tensor` +pub use tensor::Tensor; + +#[cfg(feature = "extension")] +/// Backend module. +pub mod backend; + +extern crate alloc; + +// TODO: configurable device priority +#[cfg(test)] +#[allow(missing_docs)] +pub fn test_device() -> burn_tensor::Device { + burn_tensor::Device::flex() +} + +#[cfg(test)] +mod test_utils { + use crate as burn; + use crate::module::Module; + use crate::module::Param; + use burn_tensor::Device; + use burn_tensor::Tensor; + + /// Simple linear module. + #[derive(Module, Debug)] + pub struct SimpleLinear { + pub weight: Param>, + pub bias: Option>>, + } + + impl SimpleLinear { + pub fn new(in_features: usize, out_features: usize, device: &Device) -> Self { + let weight = Tensor::random( + [out_features, in_features], + burn_tensor::Distribution::Default, + device, + ); + let bias = Tensor::random([out_features], burn_tensor::Distribution::Default, device); + + Self { + weight: Param::from_tensor(weight), + bias: Some(Param::from_tensor(bias)), + } + } + } +} + +pub mod prelude { + //! Structs and macros used by most projects. Add `use + //! burn::prelude::*` to your code to quickly get started with + //! Burn. + pub use crate::{ + config::Config, + module::Module, + tensor::{ + Bool, Device, DeviceIndex, DeviceKind, ElementConversion, Float, Int, Shape, SliceArg, + Tensor, TensorData, cast::ToElement, s, + }, + }; + pub use burn_std::device::Device as DeviceOps; +} diff --git a/crates/burn-core/src/module/base.rs b/crates/burn-core/src/module/base.rs new file mode 100644 index 0000000..9a441dd --- /dev/null +++ b/crates/burn-core/src/module/base.rs @@ -0,0 +1,459 @@ +use super::{Param, ParamId, Quantizer}; +use alloc::{string::String, vec::Vec}; +pub use burn_derive::Module; +use burn_tensor::{Bool, Device, Int, Tensor}; + +/// Type alias to `Vec` which supports `no_std` environments, but automatically using +/// the `alloc` crate. +pub type Devices = Vec; + +// At the moment, our plan is to continue experimenting with the macro internally and monitor its development. +// We may consider making it public in the future. +macro_rules! module { + (map=$module:ident, ops=$item:expr) => {{ + struct Mapper; + impl ModuleMapper for Mapper { + fn map_float(&mut self, param: Param>) -> Param> { + let (id, tensor, mapper) = param.consume(); + let func = $item; + let tensor = func(tensor); + Param::from_mapped_value(id, tensor, mapper) + } + } + let mut mapper = Mapper; + $module.map(&mut mapper) + }}; + (visit_float=$module:ident, ops=$item:expr, state=$state_ty:ty, init=$init:expr) => {{ + struct Visitor<'a> { + state: &'a mut $state_ty, + } + impl<'a> ModuleVisitor for Visitor<'a> { + fn visit_float(&mut self, param: &Param>) { + let func = $item; + func(¶m.val(), &mut self.state) + } + } + #[allow(clippy::redundant_closure_call)] + let mut state = $init(); + let mut visitor = Visitor { state: &mut state }; + $module.visit(&mut visitor); + state + }}; +} + +/// Trait for all neural network modules. +/// +/// Modules should be created using the [derive](burn_derive::Module) attribute. +/// This will make your module trainable, savable and loadable via +/// `state` and `load`. +/// +/// # Example +/// +/// ```rust, ignore +/// // Not necessary when using the burn crate directly. +/// use burn_core as burn; +/// +/// use burn::{ +/// module::Module, +/// nn::Linear, +/// tensor::Tensor, +/// }; +/// +/// #[derive(Module, Debug)] +/// struct MyModule { +/// my_param: Linear, +/// my_other_field: usize, +/// } +/// ``` +pub trait Module: Clone + Send + core::fmt::Debug { + /// Return all the devices found in the underneath module tree added to the given vector + /// without duplicates. + fn collect_devices(&self, devices: Devices) -> Devices; + + /// Return all the devices found in the underneath module tree without duplicates. + fn devices(&self) -> Devices { + self.collect_devices(Devices::new()) + } + + /// Fork the module and all of its sub-modules to the given device. + /// + /// # Notes + /// + /// This is similar to [to_device](Module::to_device), but it ensures the output module on the + /// new device will have its own autodiff graph. + fn fork(self, device: &Device) -> Self; + + /// Move the module and all of its sub-modules to the given device. + /// + /// # Warnings + /// + /// The operation supports autodiff and it will be registered when activated. However, this may + /// not be what you want. The output model will be an intermediary model, meaning that you + /// can't optimize it with gradient descent. If you want to optimize the output network on the + /// target device, use [fork](Module::fork) instead. + fn to_device(self, device: &Device) -> Self; + + /// Each tensor in the module tree will not require grad. + /// + /// # Warnings + /// + /// This should not be used for inference, use [valid](AutodiffModule::valid) when using + /// AD modules. This is mostly useful when performing partial finetuning, which is updating only + /// a small fraction of the parameters instead of finetuning all of them. + fn no_grad(self) -> Self { + module!( + map = self, + ops = |tensor: Tensor| tensor.set_require_grad(false) + ) + } + + /// Move the module and all of its sub-modules to the autodiff backend. + /// + /// # Notes + /// + /// * Only plain modules (not already on an autodiff backend) can be moved. + /// * Calling `train()` on a module that is already on an autodiff backend + /// will result in a type error, because the module's inner backend does not match. + fn train(self) -> Self + where + Self: AutodiffModule, + { + // ::TrainModule::from_inner(self) + AutodiffModule::from_inner(self) + } + + /// Get the number of parameters the module has, including all of its sub-modules. + fn num_params(&self) -> usize { + module!( + visit_float = self, + ops = |tensor: &Tensor, state: &mut usize| { + *state += tensor.shape().num_elements(); + }, + state = usize, + init = || 0 + ) + } + /// Visit each tensor parameter in the module with a [visitor](ModuleVisitor). + fn visit(&self, visitor: &mut Visitor); + + /// Map each tensor parameter in the module with a [mapper](ModuleMapper). + fn map(self, mapper: &mut Mapper) -> Self; + + /// Quantize the weights of the module. + fn quantize_weights(self, quantizer: &mut Quantizer) -> Self { + self.map(quantizer) + } + + /// Collect this module's parameters into a [`ModuleRecord`](crate::store::ModuleRecord). + /// + /// The record can be saved to a burnpack file or byte buffer and applied back with + /// [`load_record`](Module::load_record). + fn into_record(self) -> crate::store::ModuleRecord + where + Self: Sized, + { + crate::store::ModuleRecord::from_module(self) + } + + /// Apply a [`ModuleRecord`](crate::store::ModuleRecord) to this module, returning the loaded + /// module. + /// + /// Honors the record's [`DTypePolicy`](crate::store::DTypePolicy), `validate`, and + /// `allow_partial` settings. + fn try_load_record( + self, + record: crate::store::ModuleRecord, + ) -> Result + where + Self: Sized, + { + record.apply(self) + } + + /// Apply a [`ModuleRecord`](crate::store::ModuleRecord) to this module, consuming and returning + /// it. + /// + /// Panics if validation fails; use [`try_load_record`](Module::try_load_record) for the + /// fallible variant. + fn load_record(self, record: crate::store::ModuleRecord) -> Self + where + Self: Sized, + { + self.try_load_record(record).expect("Failed to load record") + } + + /// Save this module's parameters to a burnpack file on disk. + /// + /// Convenience for [`into_record`](Module::into_record) followed by + /// [`ModuleRecord::save`](crate::store::ModuleRecord::save). For non-default load behavior + /// (dtype policy, partial loading, validation), go through the record directly. + #[cfg(feature = "std")] + fn save_file>(self, path: P) -> Result<(), crate::store::RecordError> + where + Self: Sized, + { + self.into_record().save(path) + } + + /// Load this module's parameters from a burnpack file on disk, returning the loaded module. + /// + /// Uses the default load behavior. Panics on I/O or validation errors; use + /// [`try_load_file`](Module::try_load_file) for the fallible variant, or go through + /// [`ModuleRecord`](crate::store::ModuleRecord) to configure dtype policy, partial loading or + /// validation. + #[cfg(feature = "std")] + fn load_file>(self, path: P) -> Self + where + Self: Sized, + { + self.try_load_file(path) + .expect("Failed to load module from file") + } + + /// Fallible variant of [`load_file`](Module::load_file). + /// + /// Reads the record from `path` with [`ModuleRecord::load`](crate::store::ModuleRecord::load) + /// and applies it through [`try_load_record`](Module::try_load_record). + #[cfg(feature = "std")] + fn try_load_file>( + self, + path: P, + ) -> Result + where + Self: Sized, + { + let record = crate::store::ModuleRecord::load(path)?; + self.try_load_record(record) + } +} + +/// Module visitor trait for traversing and inspecting module parameters. +pub trait ModuleVisitor { + /// Visit a float parameter in the module. + /// + /// # Parameters + /// - `param`: The float parameter to visit + #[allow(unused_variables)] + fn visit_float(&mut self, param: &Param>) {} + + /// Visit an int parameter in the module. + /// + /// # Parameters + /// - `param`: The integer parameter to visit + #[allow(unused_variables)] + fn visit_int(&mut self, param: &Param>) {} + + /// Visit a bool parameter in the module. + /// + /// # Parameters + /// - `param`: The boolean parameter to visit + #[allow(unused_variables)] + fn visit_bool(&mut self, param: &Param>) {} + + /// Called when entering a submodule. + /// + /// # Parameters + /// - `name`: The name of the submodule being entered + /// - `container_type`: The type of the container with format: + /// - For user-defined structs: "Struct:TypeName" (e.g., "Struct:Linear") + /// - For user-defined enums: "Enum:TypeName" (e.g., "Enum:MyEnum") + /// - For Vec containers: "Vec" (name is the index) + /// - For Tuple containers: "Tuple" (name is the index) + /// - For Array containers: "Array" (name is the index) + /// + /// Note: Option containers do not call enter_module/exit_module to preserve + /// the field name in the path (e.g., "bias" instead of "bias.Some") + #[allow(unused_variables)] + fn enter_module(&mut self, name: &str, container_type: &str) {} + + /// Called when exiting a submodule. + /// + /// # Parameters + /// - `name`: The name of the submodule being exited + /// - `container_type`: The type of the container with format: + /// - For user-defined structs: "Struct:TypeName" (e.g., "Struct:Linear") + /// - For user-defined enums: "Enum:TypeName" (e.g., "Enum:MyEnum") + /// - For Vec containers: "Vec" (name is the index) + /// - For Tuple containers: "Tuple" (name is the index) + /// - For Array containers: "Array" (name is the index) + /// + /// Note: Option containers do not call enter_module/exit_module to preserve + /// the field name in the path (e.g., "bias" instead of "bias.Some") + #[allow(unused_variables)] + fn exit_module(&mut self, name: &str, container_type: &str) {} + + /// Visit a float tensor with its full module path. + /// + /// # Parameters + /// - `path`: The path components to the tensor as a slice (e.g., &["encoder", "layer1", "weight"]). + /// Each element represents a module name in the hierarchy, with the final element + /// being the parameter name. This allows efficient reuse of the path stack. + /// - `id`: The unique identifier of the parameter + /// - `tensor`: The float tensor to visit + #[allow(unused_variables)] + fn visit_float_with_path( + &mut self, + path: &[String], + id: ParamId, + tensor: &Tensor, + ) { + } + + /// Visit an int tensor with its full module path. + /// + /// # Parameters + /// - `path`: The path components to the tensor as a slice (e.g., &["encoder", "layer1", "weight"]). + /// Each element represents a module name in the hierarchy, with the final element + /// being the parameter name. This allows efficient reuse of the path stack. + /// - `id`: The unique identifier of the parameter + /// - `tensor`: The integer tensor to visit + #[allow(unused_variables)] + fn visit_int_with_path( + &mut self, + path: &[String], + id: ParamId, + tensor: &Tensor, + ) { + } + + /// Visit a bool tensor with its full module path. + /// + /// # Parameters + /// - `path`: The path components to the tensor as a slice (e.g., &["encoder", "layer1", "weight"]). + /// Each element represents a module name in the hierarchy, with the final element + /// being the parameter name. This allows efficient reuse of the path stack. + /// - `id`: The unique identifier of the parameter + /// - `tensor`: The boolean tensor to visit + #[allow(unused_variables)] + fn visit_bool_with_path( + &mut self, + path: &[String], + id: ParamId, + tensor: &Tensor, + ) { + } +} + +/// Module mapper trait for transforming module parameters. +pub trait ModuleMapper { + /// Called when entering a submodule. + /// + /// # Parameters + /// - `name`: The name of the submodule being entered + /// - `container_type`: The type of the container with format: + /// - For user-defined structs: "Struct:TypeName" (e.g., "Struct:Linear") + /// - For user-defined enums: "Enum:TypeName" (e.g., "Enum:MyEnum") + /// - For Vec containers: "Vec" (name is the index) + /// - For Tuple containers: "Tuple" (name is the index) + /// - For Array containers: "Array" (name is the index) + /// + /// Note: Option containers do not call enter_module/exit_module to preserve + /// the field name in the path (e.g., "bias" instead of "bias.Some") + #[allow(unused_variables)] + fn enter_module(&mut self, name: &str, container_type: &str) {} + + /// Called when exiting a submodule. + /// + /// # Parameters + /// - `name`: The name of the submodule being exited + /// - `container_type`: The type of the container with format: + /// - For user-defined structs: "Struct:TypeName" (e.g., "Struct:Linear") + /// - For user-defined enums: "Enum:TypeName" (e.g., "Enum:MyEnum") + /// - For Vec containers: "Vec" (name is the index) + /// - For Tuple containers: "Tuple" (name is the index) + /// - For Array containers: "Array" (name is the index) + /// + /// Note: Option containers do not call enter_module/exit_module to preserve + /// the field name in the path (e.g., "bias" instead of "bias.Some") + #[allow(unused_variables)] + fn exit_module(&mut self, name: &str, container_type: &str) {} + + /// Map a float parameter in the module. + /// + /// # Parameters + /// - `param`: The float parameter to transform + /// + /// # Returns + /// The transformed parameter + #[allow(unused_variables)] + fn map_float(&mut self, param: Param>) -> Param> { + let (id, tensor, mapper) = param.consume(); + Param::from_mapped_value(id, tensor, mapper) + } + + /// Map an int parameter in the module. + /// + /// # Parameters + /// - `param`: The integer parameter to transform + /// + /// # Returns + /// The transformed parameter + #[allow(unused_variables)] + fn map_int(&mut self, param: Param>) -> Param> { + let (id, tensor, mapper) = param.consume(); + Param::from_mapped_value(id, tensor, mapper) + } + + /// Map a bool parameter in the module. + /// + /// # Parameters + /// - `param`: The boolean parameter to transform + /// + /// # Returns + /// The transformed parameter + #[allow(unused_variables)] + fn map_bool( + &mut self, + param: Param>, + ) -> Param> { + let (id, tensor, mapper) = param.consume(); + Param::from_mapped_value(id, tensor, mapper) + } +} + +/// Module with auto-differentiation backend. +pub trait AutodiffModule: Module + Send + core::fmt::Debug { + /// Returns the same module, but on the inner backend without auto-differentiation. + fn valid(&self) -> Self; + + /// Wraps an inner module back into an auto-diff module. + fn from_inner(module: Self) -> Self; +} + +#[cfg(all(test, feature = "autodiff"))] +mod tests { + use super::*; + + use crate::{test_device, test_utils::SimpleLinear}; + + #[test] + fn test_module_val_train_stateful() { + let device = test_device().autodiff(); + let module = SimpleLinear::new(4, 4, &device); + + assert!(module.weight.is_require_grad()); + assert!(module.weight.require_grad); + + let module = module.valid(); + assert!(!module.weight.is_require_grad()); + assert!(module.weight.require_grad); // stateful + + // Without `HasAutodiffModule`, we would need to specify the module type as well, which would be annoying + // let module: SimpleLinear = module.train(); + let module = module.train(); + assert!(module.weight.is_require_grad()); + assert!(module.weight.require_grad); // stateful + + let module = module.no_grad(); + assert!(!module.weight.is_require_grad()); + assert!(!module.weight.require_grad); // stateful + + let module = module.valid(); + assert!(!module.weight.is_require_grad()); // always + assert!(!module.weight.require_grad); // stateful + + let module = module.train(); + assert!(!module.weight.is_require_grad()); + assert!(!module.weight.require_grad); // stateful + } +} diff --git a/crates/burn-core/src/module/display.rs b/crates/burn-core/src/module/display.rs new file mode 100644 index 0000000..40d8eb6 --- /dev/null +++ b/crates/burn-core/src/module/display.rs @@ -0,0 +1,578 @@ +use alloc::{ + borrow::ToOwned, + format, + string::{String, ToString}, + vec::Vec, +}; +use core::any; +use core::fmt::{Debug, Display, Write}; + +/// Default display settings for a module. +pub trait ModuleDisplayDefault { + /// Attributes of the module used for display purposes. + /// + /// # Arguments + /// + /// * `_content` - The content object that contains display settings and attributes. + /// + /// # Returns + /// + /// An optional content object containing the display attributes. + fn content(&self, _content: Content) -> Option; + + /// Gets the number of the parameters of the module. + fn num_params(&self) -> usize { + 0 + } +} + +/// Trait to implement custom display settings for a module. +/// +/// In order to implement custom display settings for a module, +/// 1. Add #[module(custom_display)] attribute to the module struct after #[derive(Module)] +/// 2. Implement ModuleDisplay trait for the module +pub trait ModuleDisplay: ModuleDisplayDefault { + /// Formats the module with provided display settings. + /// + /// # Arguments + /// + /// * `passed_settings` - Display settings passed to the module. + /// + /// # Returns + /// + /// A string representation of the formatted module. + fn format(&self, passed_settings: DisplaySettings) -> String { + let settings = if let Some(custom_settings) = self.custom_settings() { + custom_settings.inherit(passed_settings) + } else { + passed_settings + }; + + let indent = " ".repeat(settings.level * settings.indentation_size()); + let indent_close_braces = " ".repeat((settings.level - 1) * settings.indentation_size()); + + let settings = settings.level_up(); + + let self_type = extract_type_name::(); + + // Use custom content if it is implemented and show_all_attributes is false, + // otherwise use default content + let content = if !settings.show_all_attributes() { + self.custom_content(Content::new(settings.clone())) + .unwrap_or_else(|| { + self.content(Content::new(settings.clone())) + .unwrap_or_else(|| { + panic!("Default content should be implemented for {self_type}.") + }) + }) + } else { + self.content(Content::new(settings.clone())) + .unwrap_or_else(|| panic!("Default content should be implemented for {self_type}.")) + }; + + let top_level_type = if let Some(top_level_type) = content.top_level_type { + top_level_type.to_owned() + } else { + self_type.to_owned() + }; + + // If there is only one item in the content, return it or no attributes + if let Some(item) = content.single_item { + return item; + } else if content.attributes.is_empty() { + return top_level_type.to_string(); + } + + let mut result = String::new(); + + // Print the struct name + if settings.new_line_after_attribute() { + writeln!(result, "{top_level_type} {{").unwrap(); + } else { + write!(result, "{top_level_type} {{").unwrap(); + } + + for (i, attribute) in content.attributes.iter().enumerate() { + if settings.new_line_after_attribute() { + writeln!(result, "{indent}{}: {}", attribute.name, attribute.value).unwrap(); + } else if i == 0 { + write!(result, "{}: {}", attribute.name, attribute.value).unwrap(); + } else { + write!(result, ", {}: {}", attribute.name, attribute.value).unwrap(); + } + } + + if settings.show_num_parameters() { + let num_params = self.num_params(); + if num_params > 0 { + if settings.new_line_after_attribute() { + writeln!(result, "{indent}params: {num_params}").unwrap(); + } else { + write!(result, ", params: {num_params}").unwrap(); + } + } + } + + if settings.new_line_after_attribute() { + write!(result, "{indent_close_braces}}}").unwrap(); + } else { + write!(result, "}}").unwrap(); + } + + result + } + + /// Custom display settings for the module. + /// + /// # Returns + /// + /// An optional display settings object. + fn custom_settings(&self) -> Option { + None + } + + /// Custom attributes for the module. + /// + /// # Arguments + /// + /// * `_content` - The content object that contains display settings and attributes. + /// + /// # Returns + /// + /// An optional content object containing the custom attributes. + fn custom_content(&self, _content: Content) -> Option { + None + } +} + +/// Custom module display settings. +#[derive(Debug, Clone)] +pub struct DisplaySettings { + /// Whether to print the module parameter ids. + show_param_id: Option, + + /// Whether to print the module attributes. + show_all_attributes: Option, + + /// Whether to print the module number of parameters. + show_num_parameters: Option, + + /// Print new line after an attribute. + new_line_after_attribute: Option, + + /// Indentation size. + indentation_size: Option, + + /// Level of indentation. + level: usize, +} + +impl Default for DisplaySettings { + fn default() -> Self { + DisplaySettings { + show_param_id: None, + show_all_attributes: None, + show_num_parameters: None, + new_line_after_attribute: None, + indentation_size: None, + level: 1, + } + } +} + +impl DisplaySettings { + /// Create a new format settings. + /// + /// # Returns + /// + /// A new instance of `DisplaySettings`. + pub fn new() -> Self { + Default::default() + } + + /// Sets a flag to show module parameters. + /// + /// # Arguments + /// + /// * `flag` - Boolean flag to show module parameters. + /// + /// # Returns + /// + /// Updated `DisplaySettings` instance. + pub fn with_show_param_id(mut self, flag: bool) -> Self { + self.show_param_id = Some(flag); + self + } + + /// Sets a flag to show module attributes. + /// + /// # Arguments + /// + /// * `flag` - Boolean flag to show all module attributes. + /// + /// # Returns + /// + /// Updated `DisplaySettings` instance. + pub fn with_show_all_attributes(mut self, flag: bool) -> Self { + self.show_all_attributes = Some(flag); + self + } + + /// Sets a flag to show the number of module parameters. + /// + /// # Arguments + /// + /// * `flag` - Boolean flag to show the number of module parameters. + /// + /// # Returns + /// + /// Updated `DisplaySettings` instance. + pub fn with_show_num_parameters(mut self, flag: bool) -> Self { + self.show_num_parameters = Some(flag); + self + } + + /// Sets a flag to print a new line after an attribute. + /// + /// # Arguments + /// + /// * `flag` - Boolean flag to print a new line after an attribute. + /// + /// # Returns + /// + /// Updated `DisplaySettings` instance. + pub fn with_new_line_after_attribute(mut self, flag: bool) -> Self { + self.new_line_after_attribute = Some(flag); + self + } + + /// Sets the indentation size. + /// + /// # Arguments + /// + /// * `size` - The size of the indentation. + /// + /// # Returns + /// + /// Updated `DisplaySettings` instance. + pub fn with_indentation_size(mut self, size: usize) -> Self { + self.indentation_size = Some(size); + self + } + + /// Inherits settings from the provided settings and return a new settings object. + /// + /// # Arguments + /// + /// * `top` - The top level `DisplaySettings` to inherit from. + /// + /// # Returns + /// + /// Updated `DisplaySettings` instance. + pub fn inherit(self, top: Self) -> Self { + let mut updated = self.clone(); + + if let Some(show_param_id) = top.show_param_id { + updated.show_param_id = Some(show_param_id); + }; + + if let Some(show_all_attributes) = top.show_all_attributes { + updated.show_all_attributes = Some(show_all_attributes); + } + + if let Some(show_num_parameters) = top.show_num_parameters { + updated.show_num_parameters = Some(show_num_parameters); + } + + if let Some(new_line_after_attribute) = top.new_line_after_attribute { + updated.new_line_after_attribute = Some(new_line_after_attribute); + } + + if let Some(indentation_size) = top.indentation_size { + updated.indentation_size = Some(indentation_size); + } + + updated.level = top.level; + + updated + } + + /// A convenience method to wrap the DisplaySettings struct in an option. + /// + /// # Returns + /// + /// An optional `DisplaySettings`. + pub fn optional(self) -> Option { + Some(self) + } + + /// Increases the level of indentation. + /// + /// # Returns + /// + /// Updated `DisplaySettings` instance with increased indentation level. + pub fn level_up(mut self) -> Self { + self.level += 1; + self + } + + /// Gets `show_param_id` flag, substitutes false if not set. + /// + /// This flag is used to print the module parameter ids. + /// + /// # Returns + /// + /// A boolean value indicating whether to show parameter ids. + pub fn show_param_id(&self) -> bool { + self.show_param_id.unwrap_or(false) + } + + /// Gets `show_all_attributes`, substitutes false if not set. + /// + /// This flag is used to force to print all module attributes, overriding custom attributes. + /// + /// # Returns + /// + /// A boolean value indicating whether to show all attributes. + pub fn show_all_attributes(&self) -> bool { + self.show_all_attributes.unwrap_or(false) + } + + /// Gets `show_num_parameters`, substitutes true if not set. + /// + /// This flag is used to print the number of module parameters. + /// + /// # Returns + /// + /// A boolean value indicating whether to show the number of parameters. + pub fn show_num_parameters(&self) -> bool { + self.show_num_parameters.unwrap_or(true) + } + + /// Gets `new_line_after_attribute`, substitutes true if not set. + /// + /// This flag is used to print a new line after an attribute. + /// + /// # Returns + /// + /// A boolean value indicating whether to print a new line after an attribute. + pub fn new_line_after_attribute(&self) -> bool { + self.new_line_after_attribute.unwrap_or(true) + } + + /// Gets `indentation_size`, substitutes 2 if not set. + /// + /// This flag is used to set the size of indentation. + /// + /// # Returns + /// + /// An integer value indicating the size of indentation. + pub fn indentation_size(&self) -> usize { + self.indentation_size.unwrap_or(2) + } +} + +/// Struct to store the attributes of a module for formatting. +#[derive(Clone, Debug)] +pub struct Content { + /// List of attributes. + pub attributes: Vec, + + /// Single item content. + pub single_item: Option, + + /// Display settings. + pub display_settings: DisplaySettings, + + /// Top level type name. + pub top_level_type: Option, +} + +impl Content { + /// Creates a new attributes struct. + /// + /// # Arguments + /// + /// * `display_settings` - Display settings for the content. + /// + /// # Returns + /// + /// A new instance of `Content`. + pub fn new(display_settings: DisplaySettings) -> Self { + Content { + attributes: Vec::new(), + single_item: None, + display_settings, + top_level_type: None, + } + } + + /// Adds an attribute to the format settings. The value will be formatted and stored as a string. + /// + /// # Arguments + /// + /// * `name` - Name of the attribute. + /// * `value` - Value of the attribute. + /// + /// # Returns + /// + /// Updated `Content` instance with the new attribute added. + pub fn add(mut self, name: &str, value: &T) -> Self { + if self.single_item.is_some() { + panic!("Cannot add multiple attributes when single item is set."); + } + + let attribute = Attribute { + name: name.to_owned(), + value: value.format(self.display_settings.clone()), // TODO level + 1 + ty: any::type_name::().to_string(), + }; + self.attributes.push(attribute); + self + } + + /// Adds an attribute using its `Debug` representation. + /// + /// This is intended for fields that do not implement [`ModuleDisplay`]. + /// + /// # Arguments + /// + /// * `name` - Name of the attribute. + /// * `value` - Value of the attribute. + /// + /// # Returns + /// + /// Updated `Content` instance with the new attribute added. + pub fn add_debug_attribute(mut self, name: &str, value: &T) -> Self { + if self.single_item.is_some() { + panic!("Cannot add multiple attributes when single item is set."); + } + self.attributes.push(Attribute { + name: name.to_owned(), + value: DisplayAdapter(value).format(self.display_settings.clone()), + ty: any::type_name::().to_string(), + }); + self + } + + /// Adds a single item. + /// + /// # Arguments + /// + /// * `value` - Rendered string of the single item. + /// + /// # Returns + /// + /// Updated `Content` instance with the single item added. + pub fn add_single(mut self, value: &T) -> Self { + if !self.attributes.is_empty() { + panic!("Cannot add single item when attributes are set."); + } + + self.single_item = Some(value.format(self.display_settings.clone())); + + self + } + + /// Adds a single item. + /// + /// # Arguments + /// + /// * `value` - Formatted display value. + /// + /// # Returns + /// + /// Updated `Content` instance with the formatted single item added. + pub fn add_formatted(mut self, value: &T) -> Self { + if !self.attributes.is_empty() { + panic!("Cannot add single item when attributes are set."); + } + + self.single_item = Some(format!("{value}")); + self + } + + /// A convenience method to wrap the Attributes struct in an option + /// because it is often used as an optional field. + /// + /// # Returns + /// + /// An optional `Content`. + pub fn optional(self) -> Option { + if self.attributes.is_empty() && self.single_item.is_none() && self.top_level_type.is_none() + { + None + } else { + Some(self) + } + } + + /// Sets the top level type name. + /// + /// # Arguments + /// + /// * `ty` - The type name to set. + /// + /// # Returns + /// + /// Updated `Content` instance with the top level type name set. + pub fn set_top_level_type(mut self, ty: &str) -> Self { + self.top_level_type = Some(ty.to_owned()); + self + } +} + +/// Minimal display adapter for non-module types. +struct DisplayAdapter<'a, T: Debug>(&'a T); + +impl<'a, T: Debug> ModuleDisplayDefault for DisplayAdapter<'a, T> { + fn content(&self, content: Content) -> Option { + content.add_single(&format!("{:?}", self.0)).optional() + } +} + +impl<'a, T: Debug> ModuleDisplay for DisplayAdapter<'a, T> {} + +/// Attribute to print in the display method. +#[derive(Clone, Debug)] +pub struct Attribute { + /// Name of the attribute. + pub name: String, + + /// Value of the attribute. + pub value: String, + + /// Type of the attribute. + pub ty: String, +} + +/// Extracts the short name of a type T +/// +/// # Returns +/// +/// A string slice representing the short name of the type. +pub fn extract_type_name() -> &'static str { + // Get the full type name of T, including module path and generic parameters + let ty = any::type_name::(); + + // Find the first occurrence of '<' in the full type name + // If not found, use the length of the type name + let end = ty.find('<').unwrap_or(ty.len()); + + // Slice the type name up to the first '<' or the end + let ty = &ty[0..end]; + + // Find the last occurrence of "::" in the sliced type name + // If found, add 2 to skip the "::" itself + // If not found, start from the beginning of the type name + let start = ty.rfind("::").map(|i| i + 2).unwrap_or(0); + + // Find the last occurrence of '<' in the sliced type name + // If not found, use the length of the type name + let end = ty.rfind('<').unwrap_or(ty.len()); + + // If the start index is less than the end index, + // return the slice of the type name from start to end + // Otherwise, return the entire sliced type name + if start < end { &ty[start..end] } else { ty } +} diff --git a/crates/burn-core/src/module/initializer.rs b/crates/burn-core/src/module/initializer.rs new file mode 100644 index 0000000..42a1517 --- /dev/null +++ b/crates/burn-core/src/module/initializer.rs @@ -0,0 +1,617 @@ +use crate::tensor::Shape; + +use crate::config::Config; +use crate::module::{Param, ParamId}; +use crate::tensor::{Distribution, Tensor, s}; + +use crate as burn; + +use burn_tensor::Device; +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float as _; + +/// Enum specifying with what values a tensor should be initialized +#[derive(Config, Debug, PartialEq)] +pub enum Initializer { + /// Fills tensor with specified value everywhere + Constant { + /// The value to fill the tensor with + value: f64, + }, + /// Fills tensor with 1s everywhere + Ones, + /// Fills tensor with 0s everywhere + Zeros, + /// Fills tensor with values drawn uniformly between specified values + Uniform { + /// The minimum value to draw from + min: f64, + + /// The maximum value to draw from + max: f64, + }, + /// Fills tensor with values drawn from normal distribution with specified mean and std + Normal { + /// The mean of the normal distribution + mean: f64, + + /// The standard deviation of the normal distribution + std: f64, + }, + /// Fills tensor with values according to the uniform version of Kaiming initialization + KaimingUniform { + /// The gain to use in initialization formula + gain: f64, + + /// Whether to use fan out only in initialization formula + fan_out_only: bool, + }, + /// Fills tensor with values according to the uniform version of Kaiming initialization + KaimingNormal { + /// The gain to use in initialization formula + gain: f64, + + /// Whether to use fan out only in initialization formula + fan_out_only: bool, + }, + /// Fills tensor with values according to the uniform version of Xavier Glorot initialization + /// described in [Understanding the difficulty of training deep feedforward neural networks + /// ](https://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf) + XavierUniform { + /// The gain to use in initialization formula + gain: f64, + }, + /// Fills tensor with values according to the normal version of Xavier Glorot initialization + /// described in [Understanding the difficulty of training deep feedforward neural networks + /// ](https://proceedings.mlr.press/v9/glorot10a/glorot10a.pdf) + XavierNormal { + /// The gain to use in initialization formula + gain: f64, + }, + /// Fills tensor with values according to the (semi) orthogonal initialization + /// described in [Exact solutions to the nonlinear dynamics of learning in deep linear neural networks` + /// - [Saxe, A. et al. (2013)](https://arxiv.org/abs/1312.6120) + Orthogonal { + /// The gain to use in initialization formula + gain: f64, + }, +} + +impl Initializer { + /// Inits a tensor parameter of given shape with values depending on initializer kind. + /// + /// # Params + /// + /// - shape: Shape of the initiated tensor. + pub fn init>( + &self, + shape: S, + device: &Device, + ) -> Param> { + self.init_with(shape, None, None, device) + } + + /// Inits a tensor parameter of given shape with values depending on initializer kind. + /// + /// # Params + /// + /// - shape: Shape of the initiated tensor. + pub fn init_with>( + &self, + shape: S, + fan_in: Option, + fan_out: Option, + device: &Device, + ) -> Param> { + let device = device.clone(); + let shape: Shape = shape.into(); + let config = self.clone(); + let shape_for_closure = shape.clone(); + + Param::uninitialized( + ParamId::new(), + move |device, require_grad| { + let config = config.clone(); + let shape = shape.clone(); + device.memory_persistent_allocations((), move |_| { + let mut tensor = config.init_tensor(shape.clone(), fan_in, fan_out, device); + + if require_grad { + tensor = tensor.require_grad(); + } + + tensor + }) + }, + device, + true, + shape_for_closure, + ) + } + + fn init_tensor>( + &self, + shape: S, + fan_in: Option, + fan_out: Option, + device: &Device, + ) -> Tensor { + let shape = shape.into(); + match self { + Initializer::Constant { value } => Tensor::::full(shape, *value, device), + Initializer::Ones => Tensor::::ones(shape, device), + Initializer::Zeros => Tensor::::zeros(shape, device), + Initializer::Uniform { min, max } => uniform_draw(shape, *min, *max, device), + Initializer::Normal { mean, std } => normal_draw(shape, *mean, *std, device), + Initializer::KaimingUniform { gain, fan_out_only } => { + let a = 3.0f64.sqrt() * *gain * self.kaiming_std(*fan_out_only, fan_in, fan_out); + uniform_draw(shape, -a, a, device) + } + Initializer::KaimingNormal { gain, fan_out_only } => { + let std = *gain * self.kaiming_std(*fan_out_only, fan_in, fan_out); + normal_draw(shape, 0.0, std, device) + } + Initializer::XavierUniform { gain } => { + let a = 3.0f64.sqrt() * *gain * self.xavier_std(fan_in, fan_out); + uniform_draw(shape, -a, a, device) + } + Initializer::XavierNormal { gain } => { + let std = *gain * self.xavier_std(fan_in, fan_out); + normal_draw(shape, 0.0, std, device) + } + Initializer::Orthogonal { gain } => { + // following the implementation in pytorch: + // https://github.com/pytorch/pytorch/blob/v2.7.0/torch/nn/init.py#L574 + + assert!( + D >= 2, + "Expected D (in Tensor) to be greater or equal 2; (D >= 2)" + ); + + let rows: usize = shape.dims::()[0]; + let cols: usize = shape.num_elements() / rows; + + let mut t: Tensor<2> = normal_draw([rows, cols], 0.0, 1.0, device); + + if rows < cols { + t = t.transpose(); + } + + let (q, r) = qr_decomposition(t, device); + let [r_rows, r_cols] = r.clone().dims(); + + let diag_r = Tensor::<2>::ones([1, r_rows], device) + .matmul(Tensor::<2>::eye(r_cols, device).mul(r.clone())); + + let ph = diag_r.clone().sign(); + + let mut q = q.mul(ph); + + if rows < cols { + q = q.transpose(); + } + + q.reshape(shape).mul_scalar(*gain) + } + } + } + + fn kaiming_std( + &self, + fan_out_only: bool, + fan_in: Option, + fan_out: Option, + ) -> f64 { + let fan = if fan_out_only { fan_out } else { fan_in }; + let fan = fan.expect( + "Can't use Kaiming initialization without specifying fan. Use init_with method.", + ); + + 1.0 / (fan as f64).sqrt() + } + + fn xavier_std(&self, fan_in: Option, fan_out: Option) -> f64 { + let fan_in = fan_in.expect( + "Can't use Xavier initialization without specifying fan in. Use init_with method and \ + provide fan_in.", + ); + let fan_out = fan_out.expect( + "Can't use Xavier initialization without specifying fan out. Use init_with method and \ + provide fan_out.", + ); + (2.0 / (fan_in + fan_out) as f64).sqrt() + } +} + +fn uniform_draw>( + shape: S, + low: f64, + high: f64, + device: &Device, +) -> Tensor { + let distribution = Distribution::Uniform(low, high); + Tensor::::random(shape, distribution, device) +} + +fn normal_draw>( + shape: S, + mean: f64, + std: f64, + device: &Device, +) -> Tensor { + let distribution = Distribution::Normal(mean, std); + Tensor::::random(shape, distribution, device) +} + +fn qr_decomposition(a: Tensor<2>, device: &Device) -> (Tensor<2>, Tensor<2>) { + // Calculate the QR decomposition using Gram-Schmidt-process: https://en.wikipedia.org/wiki/Gram%E2%80%93Schmidt_process + + let [m, n] = a.clone().dims(); + let mut q = Tensor::<2>::zeros([m, n], device); + let mut r = Tensor::<2>::zeros([n, n], device); + + for j in 0..n { + let mut v: Tensor<1> = a.clone().slice(s![.., j..=j]).squeeze_dim(1); + + for i in 0..j { + let q_i: Tensor<1> = q.clone().slice(s![.., i..=i]).squeeze_dim(1); + let r_ij = q_i.clone().mul(v.clone()).sum(); + + r = r + .clone() + .slice_assign([i..i + 1, j..j + 1], r_ij.clone().unsqueeze()); + + v = v - q_i.mul(r_ij); + } + + // norm of v + let r_jj = v + .clone() + .powf(Tensor::from_floats([2.0], device)) + .sum() + .sqrt(); + + r = r + .clone() + .slice_assign([j..j + 1, j..j + 1], r_jj.clone().unsqueeze()); + + let q_j = v / r_jj; + + q = q + .clone() + .slice_assign([0..m, j..j + 1], q_j.unsqueeze_dim(1)); + } + + (q, r) +} + +#[cfg(test)] +mod tests { + use crate::test_device; + + use super::*; + + use burn_tensor::{ElementConversion, TensorData}; + use num_traits::Pow; + + use burn_tensor::Tolerance; + type FT = f32; + + fn assert_normal_init(expected_mean: f64, expected_var: f64, tensor: &Tensor<2>) { + let (actual_vars, actual_means) = tensor.clone().var_mean(0); + let actual_vars = actual_vars.to_data(); + let actual_vars = actual_vars.as_slice::().unwrap(); + let actual_means = actual_means.to_data(); + let actual_means = actual_means.as_slice::().unwrap(); + + for i in 0..tensor.shape()[0] { + let actual_var = actual_vars[i] as f64; + let actual_mean = actual_means[i] as f64; + + assert!( + (expected_var - actual_var).abs() <= 0.1, + "Expected variance to be between {expected_var} += 0.1, but got {actual_var}" + ); + assert!( + (expected_mean - actual_mean).abs() <= 0.1, + "Expected mean to be between {expected_mean} += 0.1, but got {actual_mean}" + ); + } + } + + #[test] + fn initializer_uniform_init() { + let device = test_device(); + device.seed(0); + + let (min, max) = (0.0, 1.0); + let uniform = Initializer::Uniform { min, max }; + let tensor: Tensor<4> = uniform.init([2, 2, 2, 2], &device).into_value(); + + tensor + .into_data() + .assert_within_range::(min.elem()..max.elem()); + } + + #[test] + fn initializer_normal_init() { + // seed random generator + let device = test_device(); + device.seed(0); + + let (mean, std) = (0.0, 1.0); + let normal: Tensor<1> = Initializer::Normal { mean, std } + .init([10000], &device) + .into_value(); + let (var_act, mean_act) = normal.var_mean(0); + + let var_act: f32 = var_act.into_scalar(); + let mean_act: f32 = mean_act.into_scalar(); + + assert!( + var_act >= 0.9 && var_act <= 1.1, + "Expected variance to be between 1.0 += 0.1, but got {var_act}" + ); + assert!( + mean_act >= -0.1 && mean_act <= 0.1, + "Expected mean to be between 0.0 += 0.1, but got {mean_act}" + ); + } + + #[test] + fn initializer_constant_init() { + let value = 5.0; + let constants: Tensor<4> = Initializer::Constant { value } + .init([2, 2, 2, 2], &test_device()) + .into_value(); + constants.sum().to_data().assert_approx_eq::( + &TensorData::from([value as f32 * 16.0]), + Tolerance::default(), + ); + } + + #[test] + fn initializer_zeros_init() { + let zeros: Tensor<4> = Initializer::Zeros + .init([2, 2, 2, 2], &test_device()) + .into_value(); + zeros + .sum() + .to_data() + .assert_approx_eq::(&TensorData::from([0.0]), Tolerance::default()); + } + + #[test] + fn initializer_ones_init() { + let ones: Tensor<4> = Initializer::Ones + .init([2, 2, 2, 2], &test_device()) + .into_value(); + ones.sum() + .to_data() + .assert_approx_eq::(&TensorData::from([16.0]), Tolerance::default()); + } + + #[test] + fn initializer_kaiming_uniform_init() { + let device = test_device(); + device.seed(0); + + let gain = 2_f64; + let (fan_in, fan_out) = (5, 6); + let k = (gain * (3.0 / fan_in as f64).sqrt()).elem::(); + + let tensor: Tensor<2> = Initializer::KaimingUniform { + gain, + fan_out_only: false, + } + .init_with([fan_out, fan_in], Some(fan_in), None, &device) + .into_value(); + tensor.into_data().assert_within_range(-k..k); + } + + #[test] + fn initializer_kaiming_normal_init() { + let device = test_device(); + device.seed(0); + + let gain = 2.; + let (fan_in, fan_out) = (1000, 10); + let expected_mean = 0_f64; + + let expected_var = (gain * (1. / (fan_in as f64)).sqrt()).pow(2.); + let tensor: Tensor<2> = Initializer::KaimingNormal { + gain, + fan_out_only: false, + } + .init_with([fan_out, fan_in], Some(fan_in), None, &device) + .into_value(); + assert_normal_init(expected_mean, expected_var, &tensor) + } + + #[test] + fn initializer_kaiming_uniform_init_bias() { + let device = test_device(); + device.seed(0); + + let gain = 2_f64; + let shape = [3]; + let fan_in = 5; + let k = (gain * (3.0 / fan_in as f64).sqrt()).elem::(); + + let tensor: Tensor<1> = Initializer::KaimingUniform { + gain, + fan_out_only: false, + } + .init_with(shape, Some(fan_in), None, &device) + .into_value(); + tensor.into_data().assert_within_range(-k..k); + } + + #[test] + fn initializer_kaiming_uniform_init_fan_out() { + let device = test_device(); + device.seed(0); + + let gain = 2_f64; + let (fan_in, fan_out) = (5, 6); + let k = (gain * (3.0 / fan_out as f64).sqrt()).elem::(); + + let tensor: Tensor<2> = Initializer::KaimingUniform { + gain, + fan_out_only: true, + } + .init_with([fan_out, fan_in], None, Some(fan_out), &device) + .into_value(); + tensor.into_data().assert_within_range(-k..k); + } + + #[test] + #[should_panic] + fn initializer_kaiming_uniform_no_fan() { + let device = test_device(); + device.seed(0); + + let gain = 2_f64; + let (fan_in, fan_out) = (5, 6); + + let _: Tensor<2> = Initializer::KaimingUniform { + gain, + fan_out_only: false, + } + .init([fan_out, fan_in], &device) + .into_value(); + } + + #[test] + fn initializer_xavier_uniform_init() { + let device = test_device(); + device.seed(0); + + let gain = 2.; + let (fan_in, fan_out) = (5, 6); + let bound = (gain * (6. / (fan_in + fan_out) as f64).sqrt()).elem::(); + let tensor: Tensor<2> = Initializer::XavierUniform { gain } + .init_with([fan_out, fan_in], Some(fan_in), Some(fan_out), &device) + .into_value(); + + tensor.into_data().assert_within_range(-bound..bound); + } + + #[test] + fn initializer_xavier_normal_init() { + let device = test_device(); + device.seed(0); + + let gain = 2.; + let (fan_in, fan_out) = (1000, 10); + let expected_mean = 0_f64; + + let expected_var = (gain * (2. / (fan_in as f64 + fan_out as f64)).sqrt()).powf(2.); + let tensor: Tensor<2> = Initializer::XavierNormal { gain } + .init_with([fan_out, fan_in], Some(fan_in), Some(fan_out), &device) + .into_value(); + assert_normal_init(expected_mean, expected_var, &tensor) + } + + #[test] + #[should_panic] + fn initializer_xavier_uniform_no_fan() { + let device = test_device(); + device.seed(0); + + let gain = 2.; + let (fan_in, fan_out) = (5, 6); + let _: Tensor<2> = Initializer::XavierUniform { gain } + .init([fan_out, fan_in], &device) + .into_value(); + } + + #[test] + fn test_qr_decomposition() { + let device = test_device(); + device.seed(0); + + // test values follow the example from https://pytorch.org/docs/stable/generated/torch.linalg.qr.html#torch.linalg.qr + let a = Tensor::<2>::from_floats( + [[12., -51., 4.], [6., 167., -68.], [-4., 24., -41.]], + &device, + ); + let qr = qr_decomposition(a.clone(), &device); + + // Q @ R should reconstruct input `a` + let q_matmul_r = qr.0.clone().matmul(qr.1.clone()); + + // assert that the difference between input (`a`) and Q @ R is (almost) zero + q_matmul_r + .into_data() + .assert_approx_eq::(&a.into_data(), Tolerance::rel_abs(0.1, 0.1)); + } + + #[test] + fn initializer_orthogonal_correct() { + let device = test_device(); + device.seed(0); + + let gain = 1.; + + // test 2D tensor + let size = 10; + let q: Tensor<2> = Initializer::Orthogonal { gain } + .init([size, size], &device) + .into_value(); + let eye = Tensor::<2>::eye(size, &device); + + // Q.T @ Q should be close to identity matrix + q.clone() + .transpose() + .matmul(q) + .into_data() + .assert_approx_eq::(&eye.into_data(), Tolerance::rel_abs(0.1, 0.1)); + } + + #[test] + fn initializer_orthogonal_init() { + let device = test_device(); + device.seed(0); + + let gain = 1.; + + // test 2D tensor + let shape = [25, 30]; + let t: Tensor<2> = Initializer::Orthogonal { gain } + .init(shape, &device) + .into_value(); + let dims = t.dims(); + assert_eq!( + shape, dims, + "Expected the shape of the input tensor to match the shape of the output. ({shape:?}, {dims:?})" + ); + + // test 3D tensor + let shape = [24, 6, 85]; + let t: Tensor<3> = Initializer::Orthogonal { gain } + .init(shape, &device) + .into_value(); + let dims = t.dims(); + assert_eq!( + shape, dims, + "Expected the shape of the input tensor to match the shape of the output. ({shape:?}, {dims:?})" + ); + } + + #[test] + #[should_panic] + fn initializer_orthogonal_init_1d() { + let device = test_device(); + device.seed(0); + + let gain = 1.; + + // test 1D tensor + let shape = [3]; + let _: Tensor<1> = Initializer::Orthogonal { gain } + .init(shape, &device) + .into_value(); + } +} diff --git a/crates/burn-core/src/module/mod.rs b/crates/burn-core/src/module/mod.rs new file mode 100644 index 0000000..a7d01ef --- /dev/null +++ b/crates/burn-core/src/module/mod.rs @@ -0,0 +1,11 @@ +mod base; +mod display; +mod initializer; +mod param; +mod quantize; + +pub use base::*; +pub use display::*; +pub use initializer::*; +pub use param::*; +pub use quantize::*; diff --git a/crates/burn-core/src/module/param/base.rs b/crates/burn-core/src/module/param/base.rs new file mode 100644 index 0000000..a00c8e8 --- /dev/null +++ b/crates/burn-core/src/module/param/base.rs @@ -0,0 +1,671 @@ +use super::ParamId; +use super::sync_once_cell::SyncOnceCell; +use alloc::format; + +use alloc::boxed::Box; +use burn_std::stub::RwLock; +use burn_tensor::{Device, Shape}; +use core::ops::Deref; + +#[cfg(target_has_atomic = "ptr")] +use alloc::sync::Arc; + +#[cfg(not(target_has_atomic = "ptr"))] +use portable_atomic_util::Arc; + +#[cfg(target_has_atomic = "ptr")] +type Mapper = Arc T + Send + Sync>; + +#[cfg(not(target_has_atomic = "ptr"))] +type Mapper = Arc T + Send + Sync>>; + +#[cfg(target_has_atomic = "ptr")] +fn new_mapper T + Send + Sync + 'static>(func: F) -> Mapper { + Arc::new(func) +} + +#[cfg(not(target_has_atomic = "ptr"))] +fn new_mapper T + Send + Sync + 'static>(func: F) -> Mapper { + Arc::new(Box::new(func)) +} + +type InitFn

= Box P + Send + Sync>; + +fn new_init_fn P + Send + Sync + 'static>( + func: F, +) -> InitFn

{ + Box::new(func) +} + +/// Coordinates lazy initialization across all clones of a [`Param`]. +/// +/// The sole purpose of this shared state is to ensure the initialization function runs at most +/// once: whichever clone first calls [`val`](Param::val) initializes the value, and all other +/// clones observe the same result. +/// +/// # State Management +/// +/// **Two logical states:** +/// +/// 1. **Initialized**: `value` contains the parameter value and `initialization` is `None`. +/// 2. **Lazily Managed**: `initialization` contains `Some(RwLock<...>)`. +/// - *Before initialization*: `value` is empty, inner option is `Some(Uninitialized)`. +/// - *After initialization*: `value` contains the parameter value, inner option is `None`. +/// +/// The transition from uninitialized to initialized happens exactly once and is synchronized +/// across all clones. +pub(crate) struct LazyInitState { + /// The SyncOnceCell holding the initialized parameter value. + /// Empty for uninitialized parameters, populated after first access or explicit initialization. + pub value: SyncOnceCell, + /// The deferred initialization state for lazy parameters. + /// + /// **State Transitions:** + /// - Initialized params: `None` + /// - Uninitialized params: `Some(RwLock)>)` + /// - After lazy init triggers: `Some(RwLock)` (inner Option is taken) + pub initialization: Option>>>, +} + +impl LazyInitState { + /// Create a new parameter state that is already initialized. + fn initialized(value: T) -> Arc { + Arc::new(Self { + value: SyncOnceCell::initialized(value), + initialization: None, + }) + } + + /// Create a new parameter state that is not already initialized. + fn uninitialized(uninit: Uninitialized) -> Arc { + Arc::new(Self { + value: SyncOnceCell::new(), + initialization: Some(RwLock::new(Some(uninit))), + }) + } + + /// Gets the parameter value, initializing it lazily if needed. + fn val(&self) -> &T { + self.value.get_or_init(|| { + let mut init = self + .initialization + .as_ref() + .expect("Should have an initialization when no state provided.") + .write() + .unwrap(); + let state = init.take().expect("Should exist when not initialized"); + state.initialize() + }) + } +} + +/// Parameters are the fundamental building blocks of [modules](crate::module::Module) where they +/// serve as containers for [tensors](crate::tensor::Tensor) that can be updated during +/// training, and loaded during inference. If you don't want to save the tensors +/// and/or don't want to update it during training, you don't need this type to wrap your tensor. +/// +/// # Cloning +/// +/// Cloning a parameter is always cheap; it never allocates or initializes tensors. +/// Clones share the same lazy initialization state, so initialization happens at most once and +/// all clones resolve to the same value regardless of which one triggered it. +/// +/// This sharing is strictly scoped to lazy initialization. It only guarantees that all clones +/// observe the same initialization result. Subsequent transformations operate on independent +/// parameter values and never propagate across clones. +pub struct Param { + /// The unique ID of this parameter. This is used by eg. optimizers to associate a gradient with a specific parameter. + pub id: ParamId, + /// Shared lazy initialization state across all clones of this parameter. + /// The `Arc` exists solely to coordinate lazy initialization. It is not a general + /// shared-ownership mechanism. Any mutation forks into a new `LazyInitState`. + pub(crate) state: Arc>, + pub(crate) param_mapper: ParamMapper, + // For stateful `module.valid()` <> `module.train()` + pub(crate) require_grad: bool, +} + +#[derive(Clone)] +/// Applies transformations when loading and saving parameters. +/// +/// # Mapper System +/// +/// `ParamMapper` allows applying transformations during serialization and deserialization: +/// - `load: Option>` - transformation during deserialization (applied in `transform_for_load()`) +/// - `save: Option>` - transformation during serialization (applied in `transform_for_save()`) +/// +/// These are commonly used for: +/// - Quantization/dequantization +/// - Precision conversion (e.g., FP32 ↔ FP16) +/// - Custom parameter transformations +pub struct ParamMapper { + load: Option>, + save: Option>, +} + +impl core::fmt::Debug for ParamMapper { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_fmt(format_args!( + "ParamMapper {{ load: {}, save: {} }}", + self.load.is_some(), + self.save.is_some(), + )) + } +} + +impl ParamMapper { + /// Applies the transformation when loading the given parameter. + pub fn on_load(&self, param: T) -> T { + match &self.load { + Some(mapper) => mapper(param), + None => param, + } + } + /// Applies the transformation when saving the given parameter. + pub fn on_save(&self, param: T) -> T { + match &self.save { + Some(mapper) => mapper(param), + None => param, + } + } +} + +impl Default for ParamMapper { + fn default() -> Self { + Self { + load: None, + save: None, + } + } +} + +impl core::fmt::Display for Param { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_str(format!("Param: {}", self.id).as_str()) + } +} + +impl core::fmt::Debug for Param { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_str(format!("Param: {} - {:?}", self.id, self.param_mapper).as_str()) + } +} + +pub(crate) mod sealed { + pub trait Sealed {} +} + +/// Trait that defines what is necessary for a type to be a parameter. +/// +/// # Notes +/// This trait is intentionally sealed to keep the set of parameters closed. +/// +/// Although exposed publicly, parameter types are not meant to be extensible: +/// the parameter loading/saving, module system and optimizers assume a fixed, +/// closed set of parameter types represented exclusively by [`Tensor`](crate::Tensor) instances. +pub trait Parameter: sealed::Sealed + Clone + core::fmt::Debug + Send { + /// Fetch the device. + fn device(&self) -> Device; + + /// Fetch the gradient requirement. + fn is_require_grad(&self) -> bool; + + /// Set the gradient requirement. + fn set_require_grad(self, require_grad: bool) -> Self; + + /// Fetch the shape of the parameter. + fn shape(&self) -> Shape; + + /// Moves the parameter to the target device if it is not already on it, + /// applying any kind-specific preparation required for the loading lifecycle (e.g. detach). + fn load_to_device(self, device: &Device) -> Self; +} + +/// The deferred initialization state for lazy parameters. +#[allow(clippy::type_complexity)] +pub(crate) struct Uninitialized { + /// The initialization function. Called with `(device, is_require_grad) -> Parameter`. + init: InitFn

, + /// The target device on which the parameter should be initialized. + /// Used by `lazy_device()` to provide device information without triggering initialization. + pub(crate) device: Device, + /// The gradient requirement for the parameter. + /// Used by `lazy_is_require_grad()` to provide gradient settings without triggering initialization. + pub(crate) is_require_grad: bool, + /// The shape of the tensor parameter. + /// Used by `lazy_shape()` to provide shape information without triggering initialization. + pub(crate) shape: Shape, +} + +impl Uninitialized

{ + /// Runs the initialization function. + /// + /// This is called by [Param::val] when accessing an uninitialized parameter for the first time. + /// The function is given the stored device and gradient requirement, and returns the initialized parameter. + fn initialize(self) -> P { + (self.init)(&self.device, self.is_require_grad) + } +} + +impl Param { + /// Create a new parameter that is already initialized. + pub fn initialized(id: ParamId, value: T) -> Self { + let require_grad = value.is_require_grad(); + Self { + id, + state: LazyInitState::initialized(value), + param_mapper: Default::default(), + require_grad, + } + } + + /// Create a new parameter that is not already initialized. + pub fn uninitialized( + id: ParamId, + init: F, + device: Device, + is_require_grad: bool, + shape: Shape, + ) -> Self + where + F: FnOnce(&Device, bool) -> T + Send + Sync + 'static, + { + Self { + id, + state: LazyInitState::uninitialized(Uninitialized { + init: new_init_fn(init), + device, + is_require_grad, + shape, + }), + param_mapper: Default::default(), + require_grad: is_require_grad, + } + } + + /// Gets the parameter value, initializing it lazily if needed. + /// + /// For initialized parameters, this returns a clone of the cached value. + /// For uninitialized parameters, this triggers initialization: + pub fn val(&self) -> T { + self.deref().clone() + } + + /// Check if the parameter has been initialized. + /// + /// Returns `true` if the parameter's value has been computed and cached, + /// `false` if it's still lazy and will be initialized on first access. + pub fn is_initialized(&self) -> bool { + self.state.value.get().is_some() + } + + /// Gets the parameter's value while consuming the parameter. + pub fn into_value(self) -> T { + self.consume().1 + } + + /// Gets the parameter id and value while consuming the parameter. + pub fn consume(self) -> (ParamId, T, ParamMapper) { + let tensor = self.val(); + + core::mem::drop(self.state); + + (self.id, tensor, self.param_mapper) + } + + /// Execute the given function on the inner value. + pub fn map T>(self, func: F) -> Self { + let (id, tensor, param_mapper) = self.consume(); + let tensor = func(tensor); + let require_grad = tensor.is_require_grad(); + + Self { + id, + state: LazyInitState::initialized(tensor), + param_mapper, + require_grad, + } + } + + /// Create an initialized parameter with the given id, value, and param mapper. + /// + /// This is a helper method for creating parameters while preserving the param mapper, + /// typically used in ModuleMapper implementations. + pub fn from_mapped_value(id: ParamId, value: T, param_mapper: ParamMapper) -> Self { + let require_grad = value.is_require_grad(); + Self { + id, + state: LazyInitState::initialized(value), + param_mapper, + require_grad, + } + } + + /// Runs a transformation on the parameter when loading. + pub fn load_mapper T + Send + Sync + 'static>(mut self, func: F) -> Self { + self.param_mapper.load = Some(new_mapper(func)); + + self + } + + /// Runs a transformation on the parameter when saving. + pub fn save_mapper T + Send + Sync + 'static>(mut self, func: F) -> Self { + self.param_mapper.save = Some(new_mapper(func)); + + self + } + + /// Returns a new parameter whose initialization value is transformed by the given function. + /// + /// If the parameter is still uninitialized (lazy), the transformation is chained onto the + /// existing initialization without triggering evaluation. If the parameter is already + /// initialized, it immediately applies the transformation to the current value. + pub fn init_mapper T + Send + Sync + 'static>(self, func: F) -> Self + where + T: Sync + 'static, + { + let initialization = match &self.state.initialization { + Some(init) => init, + None => return self.map(func), + }; + + let mut init = initialization.write().unwrap(); + + match init.as_mut() { + Some(value) => { + let device = value.device.clone(); + let is_require_grad = value.is_require_grad; + let shape = value.shape.clone(); + core::mem::drop(init); + + let base = self; + Self { + id: base.id, + param_mapper: base.param_mapper.clone(), + require_grad: base.require_grad, + state: LazyInitState::uninitialized(Uninitialized { + // (device, require_grad) are already encoded in `Uninitialized` state and + // applied when `base.val()` triggers initialization. The transformed tensor + // inherits those settings automatically, but since the mapper function + // `F: Fn(T) -> T` is applied on the tensor, we need to ensure the require + // grad setting is preserved. + init: new_init_fn(move |_a, b| func(base.val()).set_require_grad(b)), + device, + is_require_grad, + shape, + }), + } + } + None => { + core::mem::drop(init); + self.map(func) + } + } + } + + /// The device on which the parameter is or will be initialized, **without triggering initialization**. + /// + /// This is critical for the load optimization: when loading tensors into an uninitialized parameter, + /// we need to know the target device to move the loaded tensor appropriately, but we don't want to + /// trigger the initialization function (which would allocate an unnecessary tensor). + /// + /// Use this instead of [crate::tensor::Tensor::device] when you need the device but want to + /// preserve lazy initialization. + pub fn lazy_device(&self) -> Device { + let initialization = match &self.state.initialization { + Some(init) => init, + None => return self.device(), + }; + + let init = initialization.read().unwrap(); + + match init.as_ref() { + Some(value) => value.device.clone(), + None => self.device(), + } + } + + /// The gradient requirement on which the parameter is or will be initialized, **without triggering initialization**. + /// + /// Similar to [lazy_device](Self::lazy_device), this is critical for the load optimization. + /// When loading tensors into an uninitialized parameter, we need to apply the correct gradient + /// setting to the loaded tensor without triggering the initialization function. + /// + /// # Notes + /// + /// This is a crate-private function, since users are not expected to use `is_require_grad` of an + /// uninitialized module to then override its value. All low-level functions should be provided + /// by `burn` and should handle those details. + pub(crate) fn lazy_is_require_grad(&self) -> bool { + let initialization = match &self.state.initialization { + Some(init) => init, + None => return self.is_require_grad(), + }; + + let init = initialization.read().unwrap(); + + match init.as_ref() { + Some(value) => value.is_require_grad, + None => self.is_require_grad(), + } + } + + /// Override the gradient requirement for the current parameter. + pub fn set_require_grad(self, require_grad: bool) -> Self { + let initialization = match &self.state.initialization { + Some(init) => init, + None => return self.map(|tensor| tensor.set_require_grad(require_grad)), + }; + + let mut init = initialization.write().unwrap(); + let mut is_lazy = false; + + if let Some(value) = init.as_mut() { + is_lazy = true; + value.is_require_grad = require_grad; + }; + + core::mem::drop(init); + + if is_lazy { + return self; + } + + self.map(|tensor| tensor.set_require_grad(require_grad)) + } + + /// The shape of the parameter, **without triggering initialization**. + /// + /// This is critical for shape validation during loading: when applying tensors to an + /// uninitialized parameter, we need to validate the shape without triggering the + /// initialization function (which would allocate an unnecessary tensor). + /// + /// Use this instead of [crate::tensor::Tensor::shape] when you need the shape but want to + /// preserve lazy initialization. + pub fn lazy_shape(&self) -> burn_tensor::Shape { + let initialization = match &self.state.initialization { + Some(init) => init, + None => return self.shape(), + }; + + let init = initialization.read().unwrap(); + + match init.as_ref() { + Some(value) => value.shape.clone(), + None => self.shape(), + } + } + + /// Transform a parameter for loading by applying load transformations. + /// + /// This method is used to restore a parameter from a tensor (typically during deserialization). + /// It ensures the tensor is moved to the expected device, applies the param mapper's + /// `on_load` transformation, and preserves the autodiff settings (require_grad). + pub fn transform_for_load(self, tensor: T, param_id: ParamId) -> Self { + let mut new_tensor = tensor; + + let mapper = self.param_mapper.clone(); + + let expected_device = self.lazy_device(); + let expected_require_grad = self.lazy_is_require_grad(); + + // Make sure we load the tensor into the same module device. + new_tensor = new_tensor.load_to_device(&expected_device); + + new_tensor = mapper.on_load(new_tensor); + + // Make sure we load the tensor with the same autodiff setting. + new_tensor = new_tensor.set_require_grad(expected_require_grad); + + let mut loaded = Self::initialized(param_id, new_tensor); + loaded.param_mapper = mapper; + loaded + } + + /// Transform a parameter for saving by applying save transformations. + /// + /// This method is used to prepare a parameter for saving (typically during serialization). + /// It applies the param mapper's `on_save` transformation, which can be used + /// to modify the tensor before serialization (e.g., quantization, precision conversion). + pub fn transform_for_save(&self) -> Self { + let mut tensor = self.val(); + let mapper = self.param_mapper.clone(); + + tensor = mapper.on_save(tensor); + + Self::initialized(self.id, tensor) + } +} + +impl Clone for Param { + fn clone(&self) -> Self { + Self { + id: self.id, + state: self.state.clone(), + param_mapper: self.param_mapper.clone(), + require_grad: self.require_grad, + } + } +} + +impl Deref for Param { + type Target = T; + + fn deref(&self) -> &Self::Target { + self.state.val() + } +} +#[cfg(test)] +mod tests { + use super::*; + use burn_tensor::Tensor; + + // Param should be Sync so that models can be shared across threads + // (e.g. parallel inference with rayon). + fn _assert_sync() {} + + #[test] + fn param_is_sync() { + fn check() { + _assert_sync::>>(); + } + check(); + } + + /// Concurrent lazy initialization must not panic. + /// + /// Multiple threads call `val()` on an uninitialized `Param` simultaneously. + /// `SyncOnceCell::get_or_init` guarantees only one thread runs the initializer; + /// the others block and receive the same value. + #[cfg(feature = "std")] + #[test] + fn param_concurrent_lazy_init() { + use alloc::vec::Vec; + + let device = Default::default(); + + let param: Param> = Param::uninitialized( + ParamId::new(), + |device, _require_grad| Tensor::random([2, 3], Default::default(), device), + device, + false, + [2, 3].into(), + ); + + // Share across threads via ¶m (requires Sync). + std::thread::scope(|s| { + let handles: Vec<_> = (0..4).map(|_| s.spawn(|| param.val())).collect(); + + let results: Vec<_> = handles.into_iter().map(|h| h.join().unwrap()).collect(); + + // All threads must get the same value. + let expected = results[0].to_data(); + for result in &results[1..] { + assert_eq!(result.to_data(), expected); + } + }); + } + + #[test] + fn param_clones_share_lazy_initialization() { + let device = Default::default(); + + // We use random values so that if it initializes twice, the data will mismatch. + let param_original: Param> = Param::uninitialized( + ParamId::new(), + |device, _require_grad| Tensor::random([2, 3], Default::default(), device), + device, + false, + [2, 3].into(), + ); + + // Regression: https://github.com/tracel-ai/burn/issues/5040 + // Clone the parameter while it is still uninitialized. + // Previously, this would clone the init function only, leading to different parameter states. + let param_clone = param_original.clone(); + + let tensor_original = param_original.val(); + assert!(param_original.is_initialized()); + assert!(param_clone.is_initialized()); + + let tensor_clone = param_clone.val(); + + tensor_original + .into_data() + .assert_eq(&tensor_clone.into_data(), true); + } + + #[test] + fn param_set_require_grad_forks_from_shared_state() { + let device = Default::default(); + + let param1: Param> = Param::uninitialized( + ParamId::new(), + |device, require_grad| Tensor::ones([2, 3], device).set_require_grad(require_grad), + device, + true, + [2, 3].into(), + ); + + // Clone param; both now point to the exact same Arc + let param2 = param1.clone(); + + // Force initialization via the first clone. + let _tensor1 = param1.val(); + assert!(param1.is_initialized()); + assert!(param2.is_initialized()); + + // set_require_grad intentionally forks: param2 gets a new Arc with the mutated tensor. + let param2 = param2.set_require_grad(false); + + // The fork produced the correct require_grad state. + assert_eq!(param2.require_grad, false); + assert_eq!(param1.require_grad, true); // param1 is unaffected + + // Values are still identical (same tensor data, different grad setting). + param1 + .val() + .into_data() + .assert_eq(¶m2.val().into_data(), true); + } +} diff --git a/crates/burn-core/src/module/param/constant.rs b/crates/burn-core/src/module/param/constant.rs new file mode 100644 index 0000000..9335126 --- /dev/null +++ b/crates/burn-core/src/module/param/constant.rs @@ -0,0 +1,299 @@ +use alloc::format; +use burn_tensor::kind::{Autodiff, Basic}; +use core::fmt::Display; + +use crate as burn; +use crate::module::{ + AutodiffModule, Content, Devices, Module, ModuleDisplay, ModuleDisplayDefault, ModuleMapper, + ModuleVisitor, +}; +use burn_tensor::{Device, Tensor}; + +/// Constant macro. +#[macro_export] +macro_rules! empty { + (module) => { + fn visit(&self, _visitor: &mut V) { + // Nothing to do + } + + fn map(self, _mapper: &mut M) -> Self { + self + } + + fn to_device(self, _: &burn::tensor::Device) -> Self { + self + } + + fn fork(self, _: &burn::tensor::Device) -> Self { + self + } + + fn collect_devices(&self, devices: burn::module::Devices) -> burn::module::Devices { + devices + } + }; + + (ad_module, $type:ty) => { + fn valid(&self) -> Self { + self.clone() + } + + fn from_inner(module: Self) -> Self { + module + } + }; + + ($type:ty) => { + impl burn::module::Module for $type { + empty!(module); + } + + impl burn::module::AutodiffModule for $type { + empty!(ad_module, $type); + } + + impl burn::module::ModuleDisplayDefault for $type { + fn content(&self, content: burn::module::Content) -> Option { + let string = format!("{}", self); + content.add_formatted(&string).optional() + } + } + + impl burn::module::ModuleDisplay for $type {} + }; +} + +// TODO: breaking change for these constant types (currently empty record, non-persistent)? + +// General Types +empty!(alloc::string::String); +empty!(bool); + +// Float Types +empty!(f64); +empty!(f32); +empty!(half::bf16); +empty!(half::f16); + +// Unsigned Integer Types +empty!(usize); +empty!(u64); +empty!(u32); +empty!(u16); +empty!(u8); + +// Signed Integer Types +empty!(isize); +empty!(i64); +empty!(i32); +empty!(i16); +empty!(i8); + +impl burn::module::ModuleDisplay for str {} +impl burn::module::ModuleDisplayDefault for str { + fn content(&self, content: burn::module::Content) -> Option { + content.add_formatted(&self).optional() + } +} + +// TODO: tensor record should persist +impl Module for Tensor { + fn visit(&self, _visitor: &mut V) {} + + fn map(self, _mapper: &mut M) -> Self { + self + } + + fn to_device(self, device: &Device) -> Self { + self.to_device(device) + } + + fn fork(self, device: &Device) -> Self { + self.to_device(device) + } + + fn collect_devices(&self, mut devices: Devices) -> Devices { + let device = self.device(); + + if !devices.contains(&device) { + devices.push(device) + } + + devices + } +} + +impl ModuleDisplayDefault for Tensor { + fn content(&self, content: Content) -> Option { + let string = format!("Tensor {{rank: {D}, shape: {:?}}}", self.shape().as_slice()); + content.add_single(&string).optional() + } +} + +impl ModuleDisplay for Tensor {} + +impl AutodiffModule for Tensor { + fn valid(&self) -> Self { + self.clone().inner() + } + + fn from_inner(tensor: Self) -> Self { + Tensor::from_inner(tensor) + } +} + +// TODO: no longer necessary? +// impl Module for PhantomData { +// type Record = EmptyRecord; + +// fn visit(&self, _visitor: &mut V) { +// // Nothing to do +// } + +// fn map(self, _mapper: &mut M) -> Self { +// self +// } + +// fn load_record(self, _record: Self::Record) -> Self { +// self +// } + +// fn into_record(self) -> Self::Record { +// EmptyRecord::new() +// } + +// fn to_device(self, _: &Device) -> Self { +// self +// } + +// fn fork(self, _: &Device) -> Self { +// self +// } + +// fn collect_devices(&self, devices: Devices) -> Devices { +// devices +// } +// } + +// impl ModuleDisplayDefault for PhantomData { +// fn content(&self, content: Content) -> Option { +// content.add_single(&"PhantomData".to_string()).optional() +// } +// } + +// impl ModuleDisplay for PhantomData {} + +// impl AutodiffModule for PhantomData { +// fn valid(&self) -> Self { +// PhantomData +// } + +// fn from_inner(_module: Self) -> Self { +// Self +// } +// } + +/// Container to satisfy the Module trait for types that are not modules. +#[derive(Clone, Debug)] +#[deprecated( + since = "0.21.0", + note = "Ignored is deprecated. Use #[module(skip)] for non-persistent fields (same behavior)." +)] +pub struct Ignored(pub T); + +#[allow(deprecated)] +impl Module for Ignored +where + T: Sync + Send + core::fmt::Debug + Clone, +{ + fn visit(&self, _visitor: &mut V) { + // Nothing to do + } + + fn map(self, _mapper: &mut M) -> Self { + self + } + + fn to_device(self, _: &Device) -> Self { + self + } + + fn fork(self, _: &Device) -> Self { + self + } + + fn collect_devices(&self, devices: Devices) -> Devices { + devices + } +} + +#[allow(deprecated)] +impl ModuleDisplayDefault for Ignored +where + T: Sync + Send + core::fmt::Debug + Clone, +{ + fn content(&self, content: Content) -> Option { + // For now, just print the debug representation of the ignored value + content.add_single(&format!("{:?}", self.0)).optional() + } +} + +#[allow(deprecated)] +impl ModuleDisplay for Ignored where T: Sync + Send + core::fmt::Debug + Clone {} + +#[allow(deprecated)] +impl Display for Ignored +where + T: Sync + Send + core::fmt::Debug + Clone, +{ + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + write!(f, "{:?}", self.0) + } +} + +#[allow(deprecated)] +impl AutodiffModule for Ignored +where + T: Sync + Send + core::fmt::Debug + Clone, +{ + fn valid(&self) -> Self { + self.clone() + } + + fn from_inner(module: Self) -> Self { + module + } +} + +#[allow(deprecated)] +// Implement deref for Ignored +impl core::ops::Deref for Ignored { + type Target = T; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +#[cfg(all(test, feature = "std"))] +mod tests { + use core::marker::PhantomData; + + use burn::module::Module; + + use crate as burn; + + #[test] + fn empty_module_with_phantom() { + #[derive(Module, Debug, new)] + struct EmptyModule { + #[module(skip)] + _phantom: PhantomData, + } + + let _module = EmptyModule::::new(); + + assert_eq!(core::mem::size_of::>(), 0); + } +} diff --git a/crates/burn-core/src/module/param/group.rs b/crates/burn-core/src/module/param/group.rs new file mode 100644 index 0000000..09b9a74 --- /dev/null +++ b/crates/burn-core/src/module/param/group.rs @@ -0,0 +1,545 @@ +use crate::module::{Module, ModuleVisitor, Param}; + +use alloc::string::String; +#[cfg(target_has_atomic = "ptr")] +use alloc::sync::Arc; +use alloc::vec; +use alloc::vec::Vec; +#[cfg(not(target_has_atomic = "ptr"))] +use portable_atomic_util::Arc; +#[cfg(feature = "std")] +use regex::Regex; + +use burn_std::id::ParamId; +use burn_tensor::{Bool, Int, Tensor}; + +/// Errors tied to [ParamGroup]'s. +#[derive(Debug)] +pub enum ParamGroupError { + /// Use of an invalid Regex pattern. + InvalidPatternError(String), +} + +#[derive(Default)] +struct ParamIdCollector { + ids: Vec, +} + +impl ParamIdCollector { + pub fn ids(&self) -> Vec { + self.ids.clone() + } +} + +impl ModuleVisitor for ParamIdCollector { + fn visit_float(&mut self, param: &Param>) { + self.ids.push(param.id); + } + + fn visit_int(&mut self, param: &Param>) { + self.ids.push(param.id); + } + + fn visit_bool(&mut self, param: &Param>) { + self.ids.push(param.id); + } +} + +/// A way to represent a group of parameter for a Burn module. +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +pub struct ParamGroup { + matcher: ParamGroupMatcher, + excludes: Option, +} + +impl ParamGroup { + /// Evaluates whether a given parameter ID and its module path match this group. + pub fn matches(&self, id: &ParamId, path: Option<&str>) -> bool { + let matched = self.matcher.matches(id, path); + + let excluded = if let Some(exclude_matcher) = &self.excludes { + exclude_matcher.matches(id, path) + } else { + false + }; + + matched && !excluded + } + + /// Returns a parameter group with all of the module's parameters. + pub fn ids_from_module(module: M) -> Self { + let mut collector = ParamIdCollector::default(); + module.visit(&mut collector); + Self { + matcher: ParamGroupMatcher::from_ids(collector.ids()), + excludes: None, + } + } + + /// Matches parameters by exact text path (e.g., "model.backbone.linear.weight") + pub fn from_path(path: impl Into) -> Self { + ParamGroup::from_paths(vec![path]) + } + + /// Matches parameters by exact text paths (e.g., "model.backbone.linear.weight", etc.) + pub fn from_paths(paths: Vec>) -> Self { + Self { + matcher: ParamGroupMatcher::Path(Arc::new(PathMatcher::Exact( + paths.into_iter().map(|p| p.into()).collect(), + ))), + excludes: None, + } + } + + /// Matches parameters that include the predicate in their paths (e.g., "backbone") + pub fn from_predicate(path: impl Into) -> Self { + ParamGroup::from_predicates(vec![path]) + } + + /// Matches parameters that include all the predicates in their path (AND logic). + /// (e.g., parameter path contains "backbone" and "linear") + pub fn from_predicates(paths: Vec>) -> Self { + Self { + matcher: ParamGroupMatcher::Path(Arc::new(PathMatcher::Include( + paths.into_iter().map(|p| p.into()).collect(), + ))), + excludes: None, + } + } + + /// Matches parameters that include any of the predicates in their path (OR logic). + /// (e.g., parameter path contains "backbone" or "linear") + pub fn from_any_predicates(paths: Vec>) -> Self { + let mut matchers: Vec = paths + .into_iter() + .map(|p| ParamGroupMatcher::Path(Arc::new(PathMatcher::Include(vec![p.into()])))) + .collect(); + let mut main_matcher = if let Some(value) = matchers.pop() { + value + } else { + return Self { + matcher: ParamGroupMatcher::Path(Arc::new(PathMatcher::Include(vec![]))), + excludes: None, + }; + }; + + matchers + .iter() + .for_each(|m| main_matcher = main_matcher.clone().fuse(m)); + + Self { + matcher: main_matcher, + excludes: None, + } + } + + #[cfg(feature = "std")] + /// Matches parameters by regex pattern (e.g., "^model\.layer\.\d+$") + /// + /// # Errors + /// Returns a [ParamGroupError::InvalidPatternError] if the string cannot be compiled into a valid regex. + pub fn from_regex>(pattern: S) -> Result { + ParamGroup::from_regexes(vec![pattern]) + } + + #[cfg(feature = "std")] + /// Matches parameters for all the regex patterns (AND logic). + /// (e.g., "^encoder\.layer\.\d+", and "bias$" ) + /// + /// # Errors + /// Returns a [ParamGroupError::InvalidPatternError] if the strings cannot be compiled into a valid regex. + pub fn from_regexes>(patterns: Vec) -> Result { + let mut new_patterns = vec![]; + for pattern in patterns { + match Regex::new(pattern.as_ref()) { + Ok(re) => new_patterns.push(re), + Err(e) => { + return Err(ParamGroupError::InvalidPatternError(format!( + "Invalid regex pattern: {e}" + ))); + } + } + } + Ok(Self { + matcher: ParamGroupMatcher::Path(Arc::new(PathMatcher::Regex(new_patterns))), + excludes: None, + }) + } + + #[cfg(feature = "std")] + /// Matches parameters for any the regex patterns (OR logic). + /// (e.g., "^encoder\.layer\.\d+$", or "^decoder\.layer\.\d+$" ) + /// + /// # Errors + /// Returns a [ParamGroupError::InvalidPatternError] if the strings cannot be compiled into a valid regex. + pub fn from_any_regexes>(patterns: Vec) -> Result { + let mut matchers = vec![]; + for pattern in patterns { + match Regex::new(pattern.as_ref()) { + Ok(re) => { + matchers.push(ParamGroupMatcher::Path(Arc::new(PathMatcher::Regex(vec![ + re, + ])))) + } + Err(e) => { + return Err(ParamGroupError::InvalidPatternError(format!( + "Invalid regex pattern: {e}" + ))); + } + } + } + + let mut main_matcher = if let Some(value) = matchers.pop() { + value + } else { + return Ok(Self { + matcher: ParamGroupMatcher::Path(Arc::new(PathMatcher::Include(vec![]))), + excludes: None, + }); + }; + + matchers + .iter() + .for_each(|m| main_matcher = main_matcher.clone().fuse(m)); + + Ok(Self { + matcher: main_matcher, + excludes: None, + }) + } + + /// Matches any parameter. + pub fn all() -> Self { + Self { + matcher: ParamGroupMatcher::All, + excludes: None, + } + } + + /// Matches a specific slice of predefined parameter IDs. + pub fn from_ids(ids: Vec) -> Self { + Self { + matcher: ParamGroupMatcher::Explicit(Arc::new(ids)), + excludes: None, + } + } + + /// Fuse two parameter group. + pub fn fuse(self, other: &Self) -> Self { + Self { + matcher: self.matcher.fuse(&other.matcher), + excludes: None, + } + } + + /// Exclude the given group from the current group + pub fn exclude(mut self, group: Self) -> Self { + self.excludes = match &self.excludes { + Some(excluded) => Some(excluded.clone().fuse(&group.matcher)), + None => Some(group.matcher.clone()), + }; + self + } +} + +mod arc_serde { + use serde::{Deserialize, Deserializer, Serialize, Serializer}; + + use super::*; + + pub fn serialize(val: &Arc, serializer: S) -> Result + where + S: Serializer, + T: Serialize, + { + val.as_ref().serialize(serializer) + } + + pub fn deserialize<'de, D, T>(deserializer: D) -> Result, D::Error> + where + D: Deserializer<'de>, + T: Deserialize<'de>, + { + let v = T::deserialize(deserializer)?; + Ok(Arc::new(v)) + } +} + +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +enum ParamGroupMatcher { + All, + #[serde(with = "arc_serde")] + Explicit(Arc>), + #[serde(with = "arc_serde")] + Path(Arc), + #[serde(with = "arc_serde")] + Combined(Arc>), +} + +impl ParamGroupMatcher { + pub fn from_ids(ids: Vec) -> Self { + Self::Explicit(Arc::new(ids)) + } + + pub(crate) fn matches(&self, id: &ParamId, path: Option<&str>) -> bool { + match self { + Self::All => true, + Self::Explicit(ids) => ids.contains(id), + Self::Path(matcher) => path.is_some_and(|p| matcher.matches(p)), + Self::Combined(matchers) => matchers.iter().any(|m| m.matches(id, path)), + } + } + + fn push_combined(self, other: Self) -> Self { + match self { + ParamGroupMatcher::Combined(param_group_matchers) => match other.clone() { + ParamGroupMatcher::All => Self::All, + ParamGroupMatcher::Explicit(_) => { + let mut matchers = (*param_group_matchers).clone(); + matchers.push(other); + Self::Combined(Arc::new(matchers)) + } + ParamGroupMatcher::Path(_) => { + let mut matchers = (*param_group_matchers).clone(); + matchers.push(other); + Self::Combined(Arc::new(matchers)) + } + ParamGroupMatcher::Combined(other_matchers) => { + let mut matchers = (*param_group_matchers).clone(); + matchers.append(&mut (*other_matchers).clone()); + Self::Combined(Arc::new(matchers)) + } + }, + _ => panic!( + "`push_combined` should only be called on a ParamGroupMatcher::Combined variant." + ), + } + } + + pub(crate) fn fuse(self, other: &Self) -> Self { + match (self.clone(), other.clone()) { + (ParamGroupMatcher::All, _) => Self::All, + (_, ParamGroupMatcher::All) => Self::All, + (ParamGroupMatcher::Explicit(_), ParamGroupMatcher::Combined(_)) => { + other.clone().push_combined(self) + } + (ParamGroupMatcher::Path(_), ParamGroupMatcher::Combined(_)) => { + other.clone().push_combined(self) + } + (ParamGroupMatcher::Combined(_), ParamGroupMatcher::Explicit(_)) => { + self.push_combined(other.clone()) + } + (ParamGroupMatcher::Combined(_), ParamGroupMatcher::Path(_)) => { + self.push_combined(other.clone()) + } + (ParamGroupMatcher::Combined(_), ParamGroupMatcher::Combined(_)) => { + self.push_combined(other.clone()) + } + _ => ParamGroupMatcher::Combined(Arc::new(vec![self, other.clone()])), + } + } +} + +#[cfg(feature = "std")] +mod regex_serde { + use serde::{Deserialize, Deserializer, Serialize, Serializer}; + + use super::*; + + pub fn serialize(regexes: &[Regex], serializer: S) -> Result + where + S: Serializer, + { + let v: Vec<&str> = regexes.iter().map(|r| r.as_str()).collect(); + v.serialize(serializer) + } + + pub fn deserialize<'de, D>(deserializer: D) -> Result, D::Error> + where + D: Deserializer<'de>, + { + let strings: Vec = Vec::deserialize(deserializer)?; + strings + .into_iter() + .map(|s| Regex::new(&s).map_err(serde::de::Error::custom)) + .collect() + } +} + +#[derive(Clone, serde::Serialize, serde::Deserialize)] +enum PathMatcher { + Exact(Vec), + #[cfg(feature = "std")] + #[serde(with = "regex_serde")] + Regex(Vec), + Include(Vec), +} + +impl PathMatcher { + pub(crate) fn matches(&self, path: &str) -> bool { + match self { + PathMatcher::Exact(paths) => paths.iter().any(|p| p == path), + #[cfg(feature = "std")] + PathMatcher::Regex(patterns) => patterns.iter().all(|r| r.is_match(path)), + PathMatcher::Include(includes) => includes.iter().all(|inc| path.contains(inc)), + } + } +} + +impl alloc::fmt::Debug for PathMatcher { + fn fmt(&self, f: &mut alloc::fmt::Formatter<'_>) -> alloc::fmt::Result { + match self { + Self::Exact(arg0) => f.debug_tuple("Exact").field(arg0).finish(), + #[cfg(feature = "std")] + Self::Regex(arg0) => f.debug_tuple("Regex").field(arg0).finish(), + Self::Include(arg0) => f.debug_tuple("Include").field(arg0).finish(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::SimpleLinear; + + #[test] + fn all_matches_any_parameter() { + let group = ParamGroup::all(); + let id = ParamId::new(); + + assert!(group.matches(&id, None)); + } + + #[test] + fn explicit_matches_only_selected_ids() { + let id = ParamId::new(); + let other_id = ParamId::new(); + let group = ParamGroup::from_ids(vec![id.clone()]); + + assert!(group.matches(&id, None)); + assert!(!group.matches(&other_id, None)); + } + + #[test] + fn path_matcher_requires_path_and_matches_exactly() { + let group = ParamGroup::from_path("model.backbone.weight"); + let id = ParamId::new(); + + assert!(group.matches(&id, Some("model.backbone.weight"))); + assert!(!group.matches(&id, Some("model.backbone.bias"))); + assert!(!group.matches(&id, None)); + } + + #[test] + fn predicate_matcher_matches_substrings() { + let group = ParamGroup::from_predicate("backbone"); + let id = ParamId::new(); + + assert!(group.matches(&id, Some("model.backbone.weight"))); + assert!(!group.matches(&id, Some("model.other.weight"))); + } + + #[cfg(feature = "std")] + #[test] + fn regex_matcher_matches_pattern() { + let group = ParamGroup::from_regex(r"^model\.layer\.[0-9]+\.weight$").unwrap(); + let id = ParamId::new(); + + assert!(group.matches(&id, Some("model.layer.3.weight"))); + assert!(!group.matches(&id, Some("model.layer.weight"))); + } + + #[test] + fn ids_from_module_collects_all_param_ids() { + let device = crate::test_device(); + let module = SimpleLinear::new(4, 8, &device); + let weight_id = module.weight.id; + let bias_id = module.bias.as_ref().unwrap().id; + let group = ParamGroup::ids_from_module(module); + + assert!(group.matches(&weight_id, Some("weight"))); + assert!(group.matches(&bias_id, Some("bias"))); + } + + #[test] + fn fuse_combines_multiple_groups() { + let id = ParamId::new(); + let group1 = ParamGroup::from_ids(vec![id.clone()]); + let group2 = ParamGroup::from_path("model.layer.weight"); + let fused = group1.fuse(&group2); + + assert!(fused.matches(&id, Some("model.other.bias"))); + assert!(fused.matches(&ParamId::new(), Some("model.layer.weight"))); + assert!(!fused.matches(&ParamId::new(), Some("model.layer.bias"))); + } + + #[test] + fn exclude_removes_matching_ids_from_a_group() { + let id = ParamId::new(); + let excluded_id = ParamId::new(); + let exclude_group = ParamGroup::from_ids(vec![excluded_id.clone()]); + let group = + ParamGroup::from_ids(vec![id.clone(), excluded_id.clone()]).exclude(exclude_group); + + assert!(group.matches(&id, None)); + assert!(!group.matches(&excluded_id, None)); + } + + #[test] + fn exclude_filters_path_matches() { + let id = ParamId::new(); + let group = ParamGroup::from_predicate("backbone") + .exclude(ParamGroup::from_path("model.backbone.bias")); + + assert!(group.matches(&id, Some("model.backbone.weight"))); + assert!(!group.matches(&id, Some("model.backbone.bias"))); + } + + #[test] + fn exclude_ignores_excludes_on_the_provided_group() { + let id = ParamId::new(); + let excluded = ParamGroup::from_path("model.backbone.bias") + .exclude(ParamGroup::from_path("model.backbone.weight")); + let group = ParamGroup::from_path("model.backbone.weight").exclude(excluded); + + assert!(group.matches(&id, Some("model.backbone.weight"))); + assert!(!group.matches(&id, Some("model.backbone.bias"))); + } + + #[test] + fn from_any_predicates_matches_any_predicate() { + let group = ParamGroup::from_any_predicates(vec!["backbone", "encoder"]); + let id = ParamId::new(); + + assert!(group.matches(&id, Some("model.backbone.weight"))); + assert!(group.matches(&id, Some("model.encoder.weight"))); + assert!(!group.matches(&id, Some("model.decoder.weight"))); + } + + #[cfg(feature = "std")] + #[test] + fn from_any_regexes_matches_any_pattern() { + let group = + ParamGroup::from_any_regexes(vec![r"^model\.layer\.[0-9]+\.weight$", r"^model\.bias$"]) + .unwrap(); + let id = ParamId::new(); + + assert!(group.matches(&id, Some("model.layer.3.weight"))); + assert!(group.matches(&id, Some("model.bias"))); + assert!(!group.matches(&id, Some("model.layer.weight"))); + assert!(!group.matches(&id, Some("model.other.weight"))); + } + + #[cfg(feature = "std")] + #[test] + fn from_regexes_with_invalid_pattern() { + let result = ParamGroup::from_regex(r"[invalid("); + assert!(result.is_err()); + + let result = ParamGroup::from_regexes(vec![r"[invalid("]); + assert!(result.is_err()); + + let result = ParamGroup::from_any_regexes(vec![r"[invalid("]); + assert!(result.is_err()); + } +} diff --git a/crates/burn-core/src/module/param/id.rs b/crates/burn-core/src/module/param/id.rs new file mode 100644 index 0000000..e6223ce --- /dev/null +++ b/crates/burn-core/src/module/param/id.rs @@ -0,0 +1 @@ +pub use burn_std::id::ParamId; diff --git a/crates/burn-core/src/module/param/mod.rs b/crates/burn-core/src/module/param/mod.rs new file mode 100644 index 0000000..52d6abe --- /dev/null +++ b/crates/burn-core/src/module/param/mod.rs @@ -0,0 +1,16 @@ +mod base; +mod constant; +mod group; +mod id; +mod primitive; +mod running; +mod sync_once_cell; +mod tensor; +mod visitor; + +pub use base::*; +pub use constant::*; +pub use group::*; +pub use id::*; +pub use running::*; +pub use visitor::*; diff --git a/crates/burn-core/src/module/param/primitive.rs b/crates/burn-core/src/module/param/primitive.rs new file mode 100644 index 0000000..1dcd23a --- /dev/null +++ b/crates/burn-core/src/module/param/primitive.rs @@ -0,0 +1,319 @@ +use crate::module::{ + AutodiffModule, Content, Module, ModuleDisplay, ModuleDisplayDefault, ModuleMapper, + ModuleVisitor, +}; + +use alloc::{format, string::ToString, vec::Vec}; + +use burn_tensor::Device; +use core::fmt::Debug; + +impl Module for Option +where + T: Module + Debug + Send + Clone, +{ + fn visit(&self, visitor: &mut V) { + if let Some(module) = self { + module.visit(visitor) + } + } + + fn map(self, mapper: &mut M) -> Self { + self.map(|module| module.map(mapper)) + } + + fn to_device(self, device: &Device) -> Self { + self.map(|module| module.to_device(device)) + } + + fn fork(self, device: &Device) -> Self { + self.map(|module| module.fork(device)) + } + + fn collect_devices(&self, mut devices: Vec) -> Vec { + if let Some(module) = self.as_ref() { + devices = module.collect_devices(devices); + } + + devices + } +} + +impl ModuleDisplayDefault for Option { + fn content(&self, content: Content) -> Option { + match self { + Some(module) => content.add_single(module).optional(), + None => content.add_single("None").optional(), + } + } +} + +impl ModuleDisplay for Option {} + +impl AutodiffModule for Option +where + T: AutodiffModule + Debug + Send + Clone, +{ + fn valid(&self) -> Self { + self.as_ref().map(|module| module.valid()) + } + + fn from_inner(module: Self) -> Self { + module.map(|module| T::from_inner(module)) + } +} + +impl Module for Vec +where + T: Module + Debug + Send + Clone, +{ + fn num_params(&self) -> usize { + let mut num_params = 0; + for module in self.iter() { + num_params += module.num_params(); + } + + num_params + } + + fn visit(&self, visitor: &mut V) { + for (i, module) in self.iter().enumerate() { + let index_str = alloc::format!("{}", i); + visitor.enter_module(&index_str, "Vec"); + module.visit(visitor); + visitor.exit_module(&index_str, "Vec"); + } + } + + fn map(self, mapper: &mut M) -> Self { + self.into_iter() + .enumerate() + .map(|(i, module)| { + let index_str = alloc::format!("{}", i); + mapper.enter_module(&index_str, "Vec"); + let mapped = module.map(mapper); + mapper.exit_module(&index_str, "Vec"); + mapped + }) + .collect() + } + + fn to_device(self, device: &Device) -> Self { + self.into_iter() + .map(|module| module.to_device(device)) + .collect() + } + + fn fork(self, device: &Device) -> Self { + self.into_iter().map(|module| module.fork(device)).collect() + } + + fn collect_devices(&self, mut devices: Vec) -> Vec { + for module in self.iter() { + devices = module.collect_devices(devices); + } + + devices + } +} + +impl ModuleDisplayDefault for Vec { + fn content(&self, content: Content) -> Option { + self.iter() + .enumerate() + .fold(content, |acc, (i, module)| { + let index = format!("{i}"); + acc.add(&index, module) + }) + .set_top_level_type(format!("Vec<0..{}>", self.len()).as_str()) + .optional() + } +} + +impl ModuleDisplay for Vec {} + +impl AutodiffModule for Vec +where + T: AutodiffModule + Debug + Send + Clone, +{ + fn valid(&self) -> Self { + self.iter().map(|module| module.valid()).collect() + } + + fn from_inner(module: Self) -> Self { + module + .into_iter() + .map(|module| T::from_inner(module)) + .collect() + } +} + +impl Module for [T; N] +where + T: Module + Debug + Send + Clone, +{ + fn collect_devices(&self, mut devices: Vec) -> Vec { + for module in self.iter() { + devices = module.collect_devices(devices); + } + + devices + } + + fn num_params(&self) -> usize { + let mut num_params = 0; + for module in self.iter() { + num_params += module.num_params(); + } + + num_params + } + + fn visit(&self, visitor: &mut V) { + for (i, module) in self.iter().enumerate() { + let index_str = alloc::format!("{}", i); + visitor.enter_module(&index_str, "Array"); + module.visit(visitor); + visitor.exit_module(&index_str, "Array"); + } + } + + fn map(self, mapper: &mut M) -> Self { + let mut result = Vec::with_capacity(N); + for (i, module) in IntoIterator::into_iter(self).enumerate() { + let index_str = alloc::format!("{}", i); + mapper.enter_module(&index_str, "Array"); + let mapped = module.map(mapper); + mapper.exit_module(&index_str, "Array"); + result.push(mapped); + } + result + .try_into() + .unwrap_or_else(|v: Vec| panic!("Expected array of length {}, got {}", N, v.len())) + } + + fn to_device(self, device: &Device) -> Self { + self.map(|module| module.to_device(device)) + } + + fn fork(self, device: &Device) -> Self { + self.map(|module| module.fork(device)) + } +} + +impl ModuleDisplayDefault for [T; N] { + fn content(&self, content: Content) -> Option { + self.iter() + .enumerate() + .fold(content, |acc, (i, module)| { + let index = format!("{i}"); + acc.add(&index, module) + }) + .set_top_level_type(format!("[0..{}]", self.len()).as_str()) + .optional() + } +} + +impl ModuleDisplay for [T; N] {} + +impl AutodiffModule for [T; N] +where + T: AutodiffModule + Debug + Send + Clone, +{ + fn valid(&self) -> Self { + self.clone().map(|module| module.valid()) + } + + fn from_inner(module: Self) -> Self { + module.map(|module| T::from_inner(module)) + } +} + +/// A macro for generating implementations for tuple modules of different sizes. +/// For example: `impl_module_tuple!([L0, L1][0, 1])`. +/// Would generate an implementation for a tuple of size 2. +/// For this macro to work properly, please adhere to the convention: +/// `impl_module_tuple!([L0, L1, ..., Ln][0, 1, ..., n])`. +macro_rules! impl_module_tuple { + // `$l` represents the generic modules. + // `$i` represents the indices of the modules in the tuple. + ([$($l:ident),*][$($i:tt),*]) => { + impl<$($l,)*> Module for ($($l,)*) + where + $($l: Module + Debug + Send + Clone,)* + { + fn collect_devices(&self, mut devices: Vec) -> Vec { + $(devices = self.$i.collect_devices(devices);)* + devices + } + + fn fork(self, device: &Device) -> Self { + ($(self.$i.fork(device),)*) + } + + fn to_device(self, device: &Device) -> Self { + ($(self.$i.to_device(device),)*) + } + + fn visit(&self, visitor: &mut V) { + $( + let index_str = $i.to_string(); + visitor.enter_module(&index_str, "Tuple"); + self.$i.visit(visitor); + visitor.exit_module(&index_str, "Tuple"); + )* + } + + fn map(self, mapper: &mut M) -> Self { + ($( + { + let index_str = $i.to_string(); + mapper.enter_module(&index_str, "Tuple"); + let mapped = self.$i.map(mapper); + mapper.exit_module(&index_str, "Tuple"); + mapped + } + ,)*) + } + + } + + impl<$($l,)*> AutodiffModule for ($($l,)*) + where + $($l: AutodiffModule + Debug + Send + Clone,)* + { + fn valid(&self) -> Self { + ($(self.$i.valid(),)*) + } + + fn from_inner(module: Self) -> Self { + ($($l::from_inner(module.$i),)*) + } + } + + impl<$($l,)*> ModuleDisplayDefault for ($($l,)*) + where + $($l: ModuleDisplay,)* + { + fn content(&self, content: Content) -> Option { + let content = content + $(.add(&format!("{}", $i), &self.$i))* + .set_top_level_type(format!("({})", stringify!($($l),*)).as_str()); + content.optional() + } + } + + impl<$($l,)*> ModuleDisplay for ($($l,)*) where $($l: ModuleDisplay,)* {} + + }; +} + +impl_module_tuple!([L0, L1][0, 1]); +impl_module_tuple!([L0, L1, L2][0, 1, 2]); +impl_module_tuple!([L0, L1, L2, L3][0, 1, 2, 3]); +impl_module_tuple!([L0, L1, L2, L3, L4][0, 1, 2, 3, 4]); +impl_module_tuple!([L0, L1, L2, L3, L4, L5][0, 1, 2, 3, 4, 5]); +impl_module_tuple!([L0, L1, L2, L3, L4, L5, L6][0, 1, 2, 3, 4, 5, 6]); +impl_module_tuple!([L0, L1, L2, L3, L4, L5, L6, L7][0, 1, 2, 3, 4, 5, 6, 7]); +impl_module_tuple!([L0, L1, L2, L3, L4, L5, L6, L7, L8][0, 1, 2, 3, 4, 5, 6, 7, 8]); +impl_module_tuple!([L0, L1, L2, L3, L4, L5, L6, L7, L8, L9][0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); diff --git a/crates/burn-core/src/module/param/running.rs b/crates/burn-core/src/module/param/running.rs new file mode 100644 index 0000000..856a0e3 --- /dev/null +++ b/crates/burn-core/src/module/param/running.rs @@ -0,0 +1,233 @@ +use super::ParamId; +use crate::module::{ + AutodiffModule, Content, Module, ModuleDisplay, ModuleDisplayDefault, ModuleMapper, + ModuleVisitor, Param, +}; + +use alloc::string::ToString; +use alloc::vec::Vec; + +#[cfg(target_has_atomic = "ptr")] +use alloc::sync::Arc; + +#[cfg(not(target_has_atomic = "ptr"))] +use portable_atomic_util::Arc; + +use burn_std::stub::Mutex; +use burn_tensor::{Device, Tensor}; + +#[cfg(feature = "std")] +mod threading { + pub(super) use std::collections::HashMap; + pub(super) use std::thread::ThreadId; + + #[inline(always)] + pub(super) fn get_thread_current_id() -> ThreadId { + std::thread::current().id() + } +} + +#[cfg(not(feature = "std"))] +mod threading { + pub(super) use burn_std::stub::ThreadId; + pub(super) use hashbrown::HashMap; + + #[inline(always)] + pub(super) fn get_thread_current_id() -> ThreadId { + panic!("Current thread id is not available") + } +} + +// Re-export items from the disabled/enabled blocks +use threading::*; + +/// A state that can be updated during the forward pass while being thread safe. +/// +/// # Note +/// +/// The state value is the average of all updates on all threads. +#[derive(Clone, Debug)] +pub struct RunningState { + id: ParamId, + values: Arc>>, + value: Arc>, +} + +// Implement display for the module + +impl core::fmt::Display for RunningState { + fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result { + write!(f, "RunningState(id={})", self.id) + } +} + +impl ModuleDisplayDefault for RunningState { + fn content(&self, content: Content) -> Option { + content + .add_formatted(&"RunningState".to_string()) + .optional() + } +} + +impl ModuleDisplay for RunningState {} + +impl Module for RunningState> { + fn visit(&self, visitor: &mut V) { + let tensor = self.value.lock().unwrap(); + let param = Param::initialized(self.id, tensor.clone()); + visitor.visit_float(¶m) + } + + fn map(self, mapper: &mut M) -> Self { + let mut tensor = self.value.lock().unwrap(); + let param = Param::initialized(self.id, tensor.clone()); + let param_out = mapper.map_float(param); + let (_, tensor_out, _) = param_out.consume(); + + *tensor = tensor_out; + core::mem::drop(tensor); + + self + } + + fn to_device(self, device: &Device) -> Self { + let mut tensor = self.value.lock().unwrap(); + let tensor_out = tensor.clone().to_device(device); + + *tensor = tensor_out; + core::mem::drop(tensor); + + self + } + + fn fork(self, device: &Device) -> Self { + self.to_device(device) // Same thing here since no grad. + } + + fn collect_devices(&self, mut devices: Vec) -> Vec { + let device = self.value.lock().unwrap().device(); + + if !devices.contains(&device) { + devices.push(device) + } + + devices + } +} + +impl RunningState> { + /// Create a new running state. + pub fn new(value: Tensor) -> Self { + Self { + id: ParamId::new(), + values: Arc::new(Mutex::new(HashMap::new())), + value: Arc::new(Mutex::new(value)), + } + } + + /// Create a new running state. + pub fn with_id(id: ParamId, value: Tensor) -> Self { + Self { + id, + values: Arc::new(Mutex::new(HashMap::new())), + value: Arc::new(Mutex::new(value)), + } + } + + /// Create a new running state from a record. + pub fn from_record(record: Param>) -> Self { + let tensor = record.val(); + Self { + id: record.id, + values: Arc::new(Mutex::new(HashMap::new())), + value: Arc::new(Mutex::new(tensor)), + } + } + + /// Update the value on the current thread. + pub fn update(&self, value: Tensor) { + let thread_id = get_thread_current_id(); + let mut map = self.values.lock().unwrap(); + + if map.contains_key(&thread_id) { + self.update_value(&mut map); + } + + map.insert(thread_id, value); + } + + /// Get the current value, + /// + /// # Note + /// + /// The current value might be outdated by one update. + pub fn value(&self) -> Tensor { + let value = self.value.lock().unwrap(); + value.clone() + } + + /// Get the current value and make sure it is sync. + /// + /// # Note + /// + /// Don't use this function after an update on the same thread where other threads might have to + /// register their update before the actual synchronization needs to happen. + pub fn value_sync(&self) -> Tensor { + let thread_id = get_thread_current_id(); + let mut map = self.values.lock().unwrap(); + + if map.contains_key(&thread_id) { + self.update_value(&mut map); + } + + let value = self.value.lock().unwrap(); + value.clone() + } + + fn sync(&self) { + let mut map = self.values.lock().unwrap(); + + if !map.is_empty() { + self.update_value(&mut map); + } + } + + fn update_value(&self, map: &mut HashMap>) { + let mut value_updated: Option> = None; + let mut counter = 0; + + for (_key, tensor) in map.drain() { + counter += 1; + + value_updated = match value_updated { + Some(current) => { + let device = current.device(); + Some(tensor.to_device(&device).add(current)) + } + None => Some(tensor), + }; + } + + if let Some(value) = value_updated { + let value = value.div_scalar(counter); + let mut value_old = self.value.lock().unwrap(); + *value_old = value; + } + } +} + +impl AutodiffModule for RunningState> { + fn valid(&self) -> Self { + self.sync(); + let value = self.value(); + + RunningState::with_id(self.id, value.inner()) + } + + fn from_inner(module: Self) -> Self { + module.sync(); + let value = module.value(); + + RunningState::with_id(module.id, Tensor::from_inner(value)) + } +} diff --git a/crates/burn-core/src/module/param/sync_once_cell.rs b/crates/burn-core/src/module/param/sync_once_cell.rs new file mode 100644 index 0000000..68cc133 --- /dev/null +++ b/crates/burn-core/src/module/param/sync_once_cell.rs @@ -0,0 +1,60 @@ +//! A `Sync`-compatible single-initialization cell for both `std` and `no_std`. +//! +//! Wraps `std::sync::OnceLock` (with `std`) or `spin::Once` (without `std`) behind +//! a unified API. This makes `Param` `Sync` so models can be shared across threads +//! for parallel inference. +//! +//! We define our own wrapper instead of reusing `burn_std::stub::SyncOnceCell` because +//! that version requires `T: Debug` on all methods and lacks `get()`. + +#[cfg(feature = "std")] +use std::sync::OnceLock as Inner; + +#[cfg(not(feature = "std"))] +use spin::Once as Inner; + +pub(crate) struct SyncOnceCell(Inner); + +impl core::fmt::Debug for SyncOnceCell { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_tuple("SyncOnceCell").field(&self.get()).finish() + } +} + +impl SyncOnceCell { + /// Create a new empty cell. + pub fn new() -> Self { + Self(Inner::new()) + } + + /// Create a cell pre-populated with `value`. + pub fn initialized(value: T) -> Self { + #[cfg(feature = "std")] + { + let cell = Inner::new(); + cell.set(value).ok().expect("cell was just created"); + Self(cell) + } + #[cfg(not(feature = "std"))] + { + Self(Inner::initialized(value)) + } + } + + /// Returns `Some(&T)` if initialized, `None` otherwise. + pub fn get(&self) -> Option<&T> { + self.0.get() + } + + /// Returns `&T`, initializing with `f` on first call. + pub fn get_or_init T>(&self, f: F) -> &T { + #[cfg(feature = "std")] + { + self.0.get_or_init(f) + } + #[cfg(not(feature = "std"))] + { + self.0.call_once(f) + } + } +} diff --git a/crates/burn-core/src/module/param/tensor.rs b/crates/burn-core/src/module/param/tensor.rs new file mode 100644 index 0000000..cab5973 --- /dev/null +++ b/crates/burn-core/src/module/param/tensor.rs @@ -0,0 +1,341 @@ +use super::{Param, ParamId, Parameter}; +use crate::module::{ + AutodiffModule, Content, Module, ModuleDisplay, ModuleDisplayDefault, ModuleMapper, + ModuleVisitor, +}; +use alloc::{format, string::ToString, vec::Vec}; +use burn_tensor::{Bool, Device, Float, Int, Tensor, TensorData}; + +impl super::sealed::Sealed for Tensor {} +impl super::sealed::Sealed for Tensor {} +impl super::sealed::Sealed for Tensor {} + +impl Parameter for Tensor { + fn device(&self) -> Device { + Tensor::device(self) + } + + fn is_require_grad(&self) -> bool { + Tensor::is_require_grad(self) + } + + fn set_require_grad(self, require_grad: bool) -> Self { + Tensor::set_require_grad(self, require_grad) + } + + fn shape(&self) -> burn_std::Shape { + Tensor::shape(self) + } + + fn load_to_device(self, device: &Device) -> Self { + if self.device() != *device { + Tensor::to_device(self, device).detach() + } else { + self + } + } +} + +impl Parameter for Tensor { + fn device(&self) -> Device { + Tensor::device(self) + } + + fn is_require_grad(&self) -> bool { + false + } + + fn set_require_grad(self, _require_grad: bool) -> Self { + self + } + + fn shape(&self) -> burn_std::Shape { + Tensor::shape(self) + } + + fn load_to_device(self, device: &Device) -> Self { + if self.device() != *device { + Tensor::to_device(self, device) + } else { + self + } + } +} + +impl Parameter for Tensor { + fn device(&self) -> Device { + Tensor::device(self) + } + + fn is_require_grad(&self) -> bool { + false + } + + fn set_require_grad(self, _require_grad: bool) -> Self { + self + } + + fn shape(&self) -> burn_std::Shape { + Tensor::shape(self) + } + + fn load_to_device(self, device: &Device) -> Self { + if self.device() != *device { + Tensor::to_device(self, device) + } else { + self + } + } +} + +impl Param> { + /// Create a new parameter from a float tensor. + /// + /// # Warnings + /// + /// We strongly recommend using [Param::uninitialized] if you are using this method to + /// initialize parameters inside a module, since the tensor initialization will be lazy, + /// making the loading of weights more performant. + pub fn from_tensor(value: Tensor) -> Self { + // When creating a parameter from a float tensor, we automatically mark it as requiring + // gradients, so that it can be updated by an optimizer. + Param::initialized(ParamId::new(), value.require_grad()) + } + + /// Create a new parameter from data. + pub fn from_data(data: T, device: &Device) -> Self + where + T: Into, + { + let data: TensorData = data.into(); + // When creating a parameter from a float tensor, we automatically mark it as requiring + // gradients, so that it can be updated by an optimizer. + device.memory_persistent_allocations(data, |data| { + let value = Tensor::from_data(data, device); + Param::initialized(ParamId::new(), value.require_grad()) + }) + } +} + +impl Module for Param> { + fn visit(&self, visitor: &mut V) { + visitor.visit_float(self) + } + + fn map(self, mapper: &mut M) -> Self { + mapper.map_float(self) + } + + fn to_device(self, device: &Device) -> Self { + self.map(|tensor| tensor.to_device(device)) + } + + fn fork(self, device: &Device) -> Self { + self.map(|tensor| { + let is_require_grad = tensor.is_require_grad(); + let mut tensor = tensor.to_device(device).detach(); + + if is_require_grad { + tensor = tensor.require_grad(); + } + + tensor + }) + } + + fn collect_devices(&self, mut devices: Vec) -> Vec { + let device = self.val().device(); + + if !devices.contains(&device) { + devices.push(device) + } + + devices + } +} + +impl ModuleDisplayDefault for Param> { + fn content(&self, content: Content) -> Option { + let id = if content.display_settings.show_param_id() { + format!(", id: {}", self.id) + } else { + "".to_string() + }; + let string = format!( + "ParamTensor {{rank: {D}, shape: {:?}, kind: float{id}}}", + self.shape().as_slice() + ); + content.add_formatted(&string).optional() + } +} +impl ModuleDisplay for Param> {} + +impl Module for Param> { + fn visit(&self, visitor: &mut V) { + visitor.visit_int(self) + } + + fn map(self, mapper: &mut M) -> Self { + mapper.map_int(self) + } + + fn to_device(self, device: &Device) -> Self { + self.map(|tensor| tensor.to_device(device)) + } + + fn fork(self, device: &Device) -> Self { + self.to_device(device) // Don't support autodiff. + } + + fn collect_devices(&self, mut devices: Vec) -> Vec { + let device = self.val().device(); + + if !devices.contains(&device) { + devices.push(device) + } + + devices + } +} + +impl ModuleDisplayDefault for Param> { + fn content(&self, content: Content) -> Option { + let id = if content.display_settings.show_param_id() { + format!(", id: {}", self.id) + } else { + "".to_string() + }; + let string = format!( + "ParamTensor {{rank: {D}, shape: {:?}, kind: int{id}}}", + self.shape().as_slice() + ); + content.add_formatted(&string).optional() + } +} +impl ModuleDisplay for Param> {} + +impl Module for Param> { + fn visit(&self, visitor: &mut V) { + visitor.visit_bool(self) + } + + fn map(self, mapper: &mut M) -> Self { + mapper.map_bool(self) + } + + fn to_device(self, device: &Device) -> Self { + self.map(|tensor| tensor.to_device(device)) + } + + fn fork(self, device: &Device) -> Self { + self.to_device(device) // Don't support autodiff. + } + + fn collect_devices(&self, mut devices: Vec) -> Vec { + let device = self.val().device(); + + if !devices.contains(&device) { + devices.push(device) + } + + devices + } +} + +impl ModuleDisplayDefault for Param> { + fn content(&self, content: Content) -> Option { + let id = if content.display_settings.show_param_id() { + format!(", id: {}", self.id) + } else { + "".to_string() + }; + + let string = format!( + "ParamTensor {{rank: {D}, shape: {:?}, kind: bool{id}}}", + self.shape().as_slice() + ); + content.add_formatted(&string).optional() + } +} + +impl ModuleDisplay for Param> {} + +impl AutodiffModule for Param> { + fn valid(&self) -> Self { + // Preserve initialized param `require_grad` state, but reset the inner value's + let require_grad = self.require_grad; + let mut param = Param::initialized(self.id, self.val().inner().set_require_grad(false)); + param.require_grad = require_grad; + param + } + + fn from_inner(module: Self) -> Self { + // Reinstate the param's `require_grad` state + let tensor = Tensor::from_inner(module.val()).set_require_grad(module.require_grad); + Param::initialized(module.id, tensor) + } +} + +// impl HasAutodiffModule +// for Param> +// { +// type TrainModule = Param>; +// } + +impl AutodiffModule for Param> { + fn valid(&self) -> Self { + Param::initialized(self.id, self.val().inner()) + } + + fn from_inner(module: Self) -> Self { + Param::initialized(module.id, Tensor::from_inner(module.val())) + } +} + +impl AutodiffModule for Param> { + fn valid(&self) -> Self { + Param::initialized(self.id, self.val().inner()) + } + + fn from_inner(module: Self) -> Self { + Param::initialized(module.id, Tensor::from_inner(module.val())) + } +} + +#[cfg(all(test, feature = "std", feature = "autodiff"))] +mod tests { + use super::*; + use crate::{module::Module, test_device}; + + #[test] + fn test_param_require_grad_stateful() { + let device = test_device().autodiff(); + let tensor = Tensor::<2>::ones([3, 3], &device).require_grad(); + + let param = Param::initialized(ParamId::new(), tensor); + assert!(param.is_require_grad()); + assert!(param.require_grad); + + let param = param.valid(); + assert!(!param.is_require_grad()); + assert!(param.require_grad); // stateful + + // Without `HasAutodiffModule`, we would need to specify the param type as well, which would be annoying: + // let param: Param> = param.train(); + let param = param.train(); + assert!(param.is_require_grad()); + assert!(param.require_grad); // stateful + + let param = param.no_grad(); + assert!(!param.is_require_grad()); + assert!(!param.require_grad); // stateful + + let param = param.valid(); + assert!(!param.is_require_grad()); // always + assert!(!param.require_grad); // stateful + + let param = param.train(); + assert!(!param.is_require_grad()); + assert!(!param.require_grad); // stateful + } +} diff --git a/crates/burn-core/src/module/param/visitor.rs b/crates/burn-core/src/module/param/visitor.rs new file mode 100644 index 0000000..e099993 --- /dev/null +++ b/crates/burn-core/src/module/param/visitor.rs @@ -0,0 +1,37 @@ +use super::{Param, ParamId}; +use crate::module::{Module, ModuleVisitor}; +use alloc::vec::Vec; +use burn_tensor::{Bool, Int, Tensor}; +use core::marker::PhantomData; + +struct ParamIdCollector<'a, M> { + param_ids: &'a mut Vec, + phantom: PhantomData, +} + +impl ModuleVisitor for ParamIdCollector<'_, M> +where + M: Module, +{ + fn visit_float(&mut self, param: &Param>) { + self.param_ids.push(param.id); + } + fn visit_int(&mut self, param: &Param>) { + self.param_ids.push(param.id); + } + fn visit_bool(&mut self, param: &Param>) { + self.param_ids.push(param.id); + } +} + +/// List all the parameter ids in a module. +pub fn list_param_ids(module: &M) -> Vec { + let mut params_ids = Vec::new(); + let mut visitor = ParamIdCollector { + param_ids: &mut params_ids, + phantom: PhantomData::, + }; + module.visit(&mut visitor); + + params_ids +} diff --git a/crates/burn-core/src/module/quantize.rs b/crates/burn-core/src/module/quantize.rs new file mode 100644 index 0000000..283c7a7 --- /dev/null +++ b/crates/burn-core/src/module/quantize.rs @@ -0,0 +1,62 @@ +use burn_tensor::{ + Tensor, + quantization::{Calibration, QuantScheme, compute_q_params, compute_range}, +}; + +use crate::module::{ModuleMapper, Param}; + +/// Describes how to quantize a module. +pub struct Quantizer { + /// The calibration method used in quantization. + pub calibration: Calibration, + /// The quantization scheme. + pub scheme: QuantScheme, +} + +impl ModuleMapper for Quantizer { + fn map_float(&mut self, param: Param>) -> Param> { + let (id, tensor, mapper) = param.consume(); + let range = compute_range(&self.scheme, &tensor, &self.calibration); + let qparams = compute_q_params(&self.scheme, range); + let tensor = tensor.quantize(&self.scheme, qparams); + Param::from_mapped_value(id, tensor, mapper) + } +} + +#[cfg(all(test, not(feature = "tch")))] +mod tests { + use crate::module::{Module, Quantizer}; + use crate::test_device; + use crate::test_utils::SimpleLinear; + use burn_tensor::{ + Tolerance, + quantization::{Calibration, QuantLevel, QuantParam, QuantValue}, + }; + + #[test] + fn should_quantize_module() { + let device = test_device(); + let module = SimpleLinear::new(32, 32, &device); + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::Tensor) + .with_param(QuantParam::F32); + + let result = module.weight.val(); + + let calibration = Calibration::MinMax; + let mut quantizer = Quantizer { + calibration, + scheme, + }; + let q_module = module.quantize_weights(&mut quantizer); + let q_result = q_module.weight.val().dequantize(); + + result + .into_data() + .assert_approx_eq::(&q_result.into_data(), Tolerance::permissive()); + } +} diff --git a/crates/burn-core/src/store/mod.rs b/crates/burn-core/src/store/mod.rs new file mode 100644 index 0000000..a494c02 --- /dev/null +++ b/crates/burn-core/src/store/mod.rs @@ -0,0 +1,554 @@ +//! Minimal, non-generic record system for saving and loading module parameters. +//! +//! A [`ModuleRecord`](crate::store::ModuleRecord) holds a module's parameters (path + +//! [`ParamId`](crate::module::ParamId) + [`TensorData`](crate::tensor::TensorData)) and +//! serializes them with the [burnpack](burn_pack) format. It is produced and applied through the +//! [`Module`](crate::module::Module) trait itself ([`Module::into_record`](crate::module::Module::into_record) / +//! [`Module::load_record`](crate::module::Module::load_record)). +//! +//! This module is intentionally tiny: traversal is a straightforward +//! [`ModuleVisitor`](crate::module::ModuleVisitor) / [`ModuleMapper`](crate::module::ModuleMapper) +//! keyed by parameter path, with no filtering, adapters, or lazy snapshots. +//! The richer snapshot/import tooling (filtering, key remapping, PyTorch/SafeTensors adapters, +//! cross-framework stores) lives in the `burn-store` crate. + +use alloc::format; +use alloc::string::{String, ToString}; +use alloc::vec::Vec; + +use hashbrown::HashMap; + +use crate::module::{Module, ModuleMapper, ModuleVisitor, Param, ParamId}; +use crate::tensor::{Bool, DType, Device, Float, Int, Shape, Tensor, TensorData, kind::Basic}; + +use burn_pack::{Reader, Writer}; + +/// Controls how a parameter's dtype is resolved when loading a [`ModuleRecord`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum DTypePolicy { + /// The module parameter adopts the record's dtype (data is loaded verbatim). Default. + #[default] + FromRecord, + /// The record's data is cast to the module parameter's current dtype on load. + /// + /// Note this materializes each target parameter to read its dtype. + CastToModule, +} + +/// Error returned by [`ModuleRecord`] save/load and [`Module`] apply operations. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RecordError { + /// An I/O or format error occurred while reading or writing the record. + Io(String), + /// Validation failed while applying the record (shape mismatch, or missing tensors + /// when partial loading is not allowed). + Validation(String), +} + +impl core::fmt::Display for RecordError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + RecordError::Io(msg) => write!(f, "Record I/O error: {msg}"), + RecordError::Validation(msg) => write!(f, "Record validation error: {msg}"), + } + } +} + +#[cfg(feature = "std")] +impl std::error::Error for RecordError {} + +impl From for RecordError { + fn from(err: burn_pack::Error) -> Self { + RecordError::Io(err.to_string()) + } +} + +/// A single recorded tensor: its module path, parameter id, and data. +#[derive(Clone)] +struct RecordTensor { + path: String, + id: ParamId, + data: TensorData, +} + +/// A non-generic record holding a module's parameters. +/// +/// Obtain one from a module with [`Module::into_record`], then either save it +/// ([`save`](ModuleRecord::save) / [`into_bytes`](ModuleRecord::into_bytes)) or apply it back with +/// [`Module::load_record`]. Load-time behavior is +/// configured with the builder methods; they are ignored when saving. +/// +/// The save-side dtype is intentionally not configurable: use `module.cast(dtype)` before +/// taking the record. The record stores whatever dtype the module currently holds. +#[derive(Clone)] +pub struct ModuleRecord { + tensors: Vec, + dtype_policy: DTypePolicy, + allow_partial: bool, + validate: bool, +} + +impl core::fmt::Debug for ModuleRecord { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("ModuleRecord") + .field("num_tensors", &self.tensors.len()) + .field("dtype_policy", &self.dtype_policy) + .field("allow_partial", &self.allow_partial) + .field("validate", &self.validate) + .finish() + } +} + +impl ModuleRecord { + fn from_tensors(tensors: Vec) -> Self { + Self { + tensors, + dtype_policy: DTypePolicy::default(), + allow_partial: false, + validate: true, + } + } + + /// The number of tensors in the record. + pub fn len(&self) -> usize { + self.tensors.len() + } + + /// Whether the record holds no tensors. + pub fn is_empty(&self) -> bool { + self.tensors.is_empty() + } + + /// Set the dtype policy used when loading into a module. + pub fn with_dtype_policy(mut self, policy: DTypePolicy) -> Self { + self.dtype_policy = policy; + self + } + + /// Cast the record's data to the module parameter dtypes on load. + /// + /// Sugar for [`with_dtype_policy(DTypePolicy::CastToModule)`](ModuleRecord::with_dtype_policy). + pub fn cast_to_module_dtype(self) -> Self { + self.with_dtype_policy(DTypePolicy::CastToModule) + } + + /// Allow loading even when some module parameters are absent from the record. + pub fn allow_partial(mut self, allow: bool) -> Self { + self.allow_partial = allow; + self + } + + /// Enable or disable validation while loading. + pub fn validate(mut self, validate: bool) -> Self { + self.validate = validate; + self + } + + /// Serialize the record to an in-memory burnpack byte buffer. + pub fn into_bytes(self) -> Result { + Ok(Writer::new(self.pack_tensors()).into_bytes()?) + } + + /// Reconstruct a record from an in-memory burnpack byte buffer. + pub fn from_bytes(bytes: crate::tensor::Bytes) -> Result { + Self::from_reader(Reader::from_bytes(bytes)?) + } + + /// Save the record to a burnpack file on disk. + #[cfg(feature = "std")] + pub fn save>(self, path: P) -> Result<(), RecordError> { + Writer::new(self.pack_tensors()).write_to_file(path)?; + Ok(()) + } + + /// Load a record from a burnpack file on disk. + #[cfg(feature = "std")] + pub fn load>(path: P) -> Result { + Self::from_reader(Reader::from_file(path)?) + } + + fn pack_tensors(self) -> Vec { + self.tensors + .into_iter() + .map(|t| { + burn_pack::Tensor::new( + t.path, + t.data.dtype, + t.data.shape, + Some(t.id.val()), + t.data.bytes, + ) + }) + .collect() + } + + fn from_reader(reader: Reader) -> Result { + let tensors = reader + .into_tensors()? + .into_iter() + .map(|t| { + let id = t.param_id.map(ParamId::from).unwrap_or_else(ParamId::new); + let data = TensorData::from_bytes(t.bytes, t.shape, t.dtype); + RecordTensor { + path: t.name, + id, + data, + } + }) + .collect(); + Ok(Self::from_tensors(tensors)) + } + + /// Collect a module's parameters into a [`ModuleRecord`]. + /// + /// Backs [`Module::into_record`](crate::module::Module::into_record). + pub(crate) fn from_module(module: M) -> Self { + let mut collector = Collector::default(); + module.visit(&mut collector); + ModuleRecord::from_tensors(collector.tensors) + } + + /// Apply this record to a module, returning the loaded module. + /// + /// Honors the record's [`DTypePolicy`], `validate`, and `allow_partial` settings. Backs + /// [`Module::try_load_record`](crate::module::Module::try_load_record). + pub(crate) fn apply(self, module: M) -> Result { + let validate = self.validate; + let allow_partial = self.allow_partial; + + let mut mapper = ModuleRecordMapper::new(self); + let module = module.map(&mut mapper); + + if validate && !mapper.errors.is_empty() { + return Err(RecordError::Validation(format!( + "Apply errors: {:?}", + mapper.errors + ))); + } + if !allow_partial && !mapper.missing.is_empty() { + return Err(RecordError::Validation(format!( + "Missing tensors: {:?}", + mapper.missing + ))); + } + + Ok(module) + } +} + +/// Visitor that collects every parameter as a [`RecordTensor`], keyed by its module path. +#[derive(Default)] +struct Collector { + path: Vec, + tensors: Vec, +} + +impl Collector { + fn record(&mut self, id: ParamId, data: TensorData) { + self.tensors.push(RecordTensor { + path: self.path.join("."), + id, + data, + }); + } +} + +impl ModuleVisitor for Collector { + fn enter_module(&mut self, name: &str, _container_type: &str) { + self.path.push(name.to_string()); + } + + fn exit_module(&mut self, _name: &str, _container_type: &str) { + self.path.pop(); + } + + fn visit_float(&mut self, param: &Param>) { + self.record(param.id, param.val().into_data()); + } + + fn visit_int(&mut self, param: &Param>) { + self.record(param.id, param.val().into_data()); + } + + fn visit_bool(&mut self, param: &Param>) { + self.record(param.id, param.val().into_data()); + } +} + +/// Mapper that loads recorded tensors back onto matching parameters by module path. +struct ModuleRecordMapper { + path: Vec, + tensors: HashMap, + dtype_policy: DTypePolicy, + missing: Vec, + errors: Vec, +} + +impl ModuleRecordMapper { + fn new(record: ModuleRecord) -> Self { + let tensors = record + .tensors + .into_iter() + .map(|t| (t.path, t.data)) + .collect(); + Self { + path: Vec::new(), + tensors, + dtype_policy: record.dtype_policy, + missing: Vec::new(), + errors: Vec::new(), + } + } + + /// Look up the recorded tensor for the current path and build the tensor to load, + /// or `None` (recording it as missing / errored) to leave the parameter unchanged. + /// + /// `module_dtype` is only evaluated on a hit under [`DTypePolicy::CastToModule`], so a + /// missing parameter never materializes its current value just to read a dtype. + fn take( + &mut self, + device: &Device, + target_shape: Shape, + module_dtype: impl FnOnce() -> DType, + ) -> Option> { + let path = self.path.join("."); + let data = match self.tensors.remove_entry(&path) { + Some(data) => data.1, + None => { + self.missing.push(path); + return None; + } + }; + + // Resolve the dtype to load with (CastToModule casts to the module's dtype). + let dtype = match self.dtype_policy { + DTypePolicy::FromRecord => data.dtype, + DTypePolicy::CastToModule => module_dtype(), + }; + + if data.shape != target_shape { + self.errors.push(format!( + "{path}: shape mismatch, expected {:?} but record has {:?}", + target_shape, data.shape + )); + return None; + } + + Some(Tensor::from_data(data, (device, dtype))) + } +} + +/// Generate a `ModuleMapper::map_*` method for one tensor kind. The three kinds differ only +/// in the tensor type, so the body — collect identity, look the tensor up, load on a hit — is +/// shared here. +macro_rules! map_kind { + ($method:ident, $kind:ty) => { + fn $method( + &mut self, + param: Param>, + ) -> Param> { + let id = param.id; + let device = param.lazy_device(); + let shape = param.lazy_shape(); + match self.take(&device, shape, || param.val().dtype()) { + Some(tensor) => param.transform_for_load(tensor, id), + None => param, + } + } + }; +} + +impl ModuleMapper for ModuleRecordMapper { + fn enter_module(&mut self, name: &str, _container_type: &str) { + self.path.push(name.to_string()); + } + + fn exit_module(&mut self, _name: &str, _container_type: &str) { + self.path.pop(); + } + + map_kind!(map_float, Float); + map_kind!(map_int, Int); + map_kind!(map_bool, Bool); +} + +#[cfg(all(test, feature = "std"))] +mod tests { + use super::*; + use crate as burn; + use crate::module::{Module, Param}; + use crate::tensor::Tensor; + use burn_tensor::Device; + + #[derive(Module, Debug)] + struct Tiny { + weight: Param>, + bias: Param>, + } + + impl Tiny { + fn new(weight: [[f32; 2]; 2], bias: [f32; 2], device: &Device) -> Self { + Self { + weight: Param::from_data(weight, device), + bias: Param::from_data(bias, device), + } + } + } + + #[derive(Module, Debug)] + struct TinyWide { + weight: Param>, + bias: Param>, + gamma: Param>, + } + + impl TinyWide { + fn zeros(device: &Device) -> Self { + Self { + weight: Param::from_data([[0.0, 0.0], [0.0, 0.0]], device), + bias: Param::from_data([0.0, 0.0], device), + gamma: Param::from_data([0.0, 0.0], device), + } + } + } + + fn weights(model: &Tiny) -> (Vec, Vec) { + ( + model.weight.val().to_data().to_vec().unwrap(), + model.bias.val().to_data().to_vec().unwrap(), + ) + } + + #[test] + fn round_trip_in_memory() { + let device = Default::default(); + let model = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device); + + let bytes = model.into_record().into_bytes().unwrap(); + let record = ModuleRecord::from_bytes(bytes).unwrap(); + assert_eq!(record.len(), 2); + + let loaded = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device).load_record(record); + let (w, b) = weights(&loaded); + assert_eq!(w, vec![1.0, 2.0, 3.0, 4.0]); + assert_eq!(b, vec![5.0, 6.0]); + } + + #[test] + fn round_trip_file() { + let device = Default::default(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("tiny.bpk"); + + Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device) + .into_record() + .save(&path) + .unwrap(); + + let record = ModuleRecord::load(&path).unwrap(); + let loaded = Tiny::new([[0.0; 2]; 2], [0.0; 2], &device).load_record(record); + let (w, b) = weights(&loaded); + assert_eq!(w, vec![1.0, 2.0, 3.0, 4.0]); + assert_eq!(b, vec![5.0, 6.0]); + } + + #[test] + fn missing_tensor_requires_allow_partial() { + let device = Default::default(); + // Tiny has weight+bias; TinyWide also expects `gamma`, which the record lacks. + let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record(); + + let strict = TinyWide::zeros(&device).try_load_record(record.clone()); + assert!(matches!(strict, Err(RecordError::Validation(_)))); + + let partial = TinyWide::zeros(&device).try_load_record(record.allow_partial(true)); + assert!(partial.is_ok()); + let loaded = partial.unwrap(); + // weight/bias were loaded; gamma kept its (zero) initialization. + assert_eq!( + loaded.weight.val().to_data().to_vec::().unwrap(), + vec![1.0, 2.0, 3.0, 4.0] + ); + assert_eq!( + loaded.gamma.val().to_data().to_vec::().unwrap(), + vec![0.0, 0.0] + ); + } + + /// Build a `Tiny` whose parameters are zero-valued and carry the given dtype. + fn tiny_with_dtype(device: &Device, dtype: DType) -> Tiny { + Tiny { + weight: Param::from_tensor( + Tensor::<2>::from_data([[0.0, 0.0], [0.0, 0.0]], device).cast(dtype), + ), + bias: Param::from_tensor(Tensor::<1>::from_data([0.0, 0.0], device).cast(dtype)), + } + } + + #[test] + fn dtype_policy_from_record_keeps_record_dtype() { + let device = Default::default(); + // Record holds f32 data. + let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record(); + + // Target params are f64, but the default policy (FromRecord) loads the data verbatim (f32). + let loaded = tiny_with_dtype(&device, DType::F64).load_record(record); + assert_eq!(loaded.weight.val().dtype(), DType::F32); + assert_eq!(loaded.bias.val().dtype(), DType::F32); + } + + #[test] + fn dtype_policy_cast_to_module_uses_module_dtype() { + let device = Default::default(); + let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record(); + + // CastToModule casts the record's f32 data to the module's f64 params on load. + let loaded = + tiny_with_dtype(&device, DType::F64).load_record(record.cast_to_module_dtype()); + assert_eq!(loaded.weight.val().dtype(), DType::F64); + assert_eq!(loaded.bias.val().dtype(), DType::F64); + // Values survive the cast. + assert_eq!( + loaded.weight.val().to_data().to_vec::().unwrap(), + vec![1.0, 2.0, 3.0, 4.0] + ); + } + + /// Build a `Tiny` whose `bias` shape (len 3) does not match the record's (len 2). + fn tiny_wrong_bias_shape(device: &Device) -> Tiny { + Tiny { + weight: Param::from_data([[0.0, 0.0], [0.0, 0.0]], device), + bias: Param::from_data([0.0, 0.0, 0.0], device), + } + } + + #[test] + fn shape_mismatch_fails_validation() { + let device = Default::default(); + let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record(); + + // `bias` shape mismatch is reported as a validation error by default. + let result = tiny_wrong_bias_shape(&device).try_load_record(record); + assert!(matches!(result, Err(RecordError::Validation(_)))); + } + + #[test] + fn validate_false_ignores_shape_mismatch() { + let device = Default::default(); + let record = Tiny::new([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device).into_record(); + + // With validation disabled, the shape-mismatched `bias` is skipped (keeps its value) + // while the matching `weight` still loads. + let loaded = tiny_wrong_bias_shape(&device) + .try_load_record(record.validate(false)) + .unwrap(); + assert_eq!( + loaded.weight.val().to_data().to_vec::().unwrap(), + vec![1.0, 2.0, 3.0, 4.0] + ); + assert_eq!( + loaded.bias.val().to_data().to_vec::().unwrap(), + vec![0.0, 0.0, 0.0] + ); + } +} diff --git a/crates/burn-core/src/tensor.rs b/crates/burn-core/src/tensor.rs new file mode 100644 index 0000000..074606b --- /dev/null +++ b/crates/burn-core/src/tensor.rs @@ -0,0 +1 @@ +pub use burn_tensor::*; diff --git a/crates/burn-core/tests/test_derive_config.rs b/crates/burn-core/tests/test_derive_config.rs new file mode 100644 index 0000000..65c4376 --- /dev/null +++ b/crates/burn-core/tests/test_derive_config.rs @@ -0,0 +1,113 @@ +use burn::config::{Config, config_to_json}; +use burn_core as burn; + +#[derive(Config, Debug, PartialEq, Eq)] +pub struct TestEmptyStructConfig {} + +#[derive(Config, Debug, PartialEq)] +pub struct TestStructConfig { + int: i32, + #[config(default = 2)] + int_default: i32, + float: f32, + #[config(default = 2.0)] + float_default: f32, + string: String, + other_config: TestEmptyStructConfig, +} + +#[derive(Config, Debug, PartialEq)] +pub enum TestEnumConfig { + None, + Single(f32), + Multiple(f32, String), + Named { first: f32, second: String }, +} + +#[cfg(feature = "std")] +#[inline(always)] +fn file_path(file_name: &str) -> std::path::PathBuf { + std::env::temp_dir().join(file_name) +} + +#[cfg(feature = "std")] +#[test] +fn struct_config_should_impl_serde() { + let config = TestStructConfig::new(2, 3.0, "Allow".to_string(), TestEmptyStructConfig::new()); + let file_path = file_path("test_struct_config.json"); + + config.save(&file_path).unwrap(); + + let config_loaded = TestStructConfig::load(&file_path).unwrap(); + assert_eq!(config, config_loaded); +} + +#[test] +fn struct_config_should_impl_clone() { + let config = TestStructConfig::new(2, 3.0, "Allow".to_string(), TestEmptyStructConfig::new()); + assert_eq!(config, config.clone()); +} + +#[test] +fn struct_config_should_impl_display() { + let config = TestStructConfig::new(2, 3.0, "Allow".to_string(), TestEmptyStructConfig::new()); + assert_eq!(burn::config::config_to_json(&config), config.to_string()); +} + +#[cfg(feature = "std")] +#[test] +fn enum_config_no_value_should_impl_serde() { + let config = TestEnumConfig::None; + let file_path = file_path("test_enum_no_value_config.json"); + + config.save(&file_path).unwrap(); + + let config_loaded = TestEnumConfig::load(&file_path).unwrap(); + assert_eq!(config, config_loaded); +} + +#[cfg(feature = "std")] +#[test] +fn enum_config_one_value_should_impl_serde() { + let config = TestEnumConfig::Single(42.0); + let file_path = file_path("test_enum_one_value_config.json"); + + config.save(&file_path).unwrap(); + + let config_loaded = TestEnumConfig::load(&file_path).unwrap(); + assert_eq!(config, config_loaded); +} + +#[cfg(feature = "std")] +#[test] +fn enum_config_multiple_values_should_impl_serde() { + let config = TestEnumConfig::Multiple(42.0, "Allow".to_string()); + let file_path = file_path("test_enum_multiple_values_config.json"); + + config.save(&file_path).unwrap(); + + let config_loaded = TestEnumConfig::load(&file_path).unwrap(); + assert_eq!(config, config_loaded); +} + +#[test] +fn enum_config_should_impl_clone() { + let config = TestEnumConfig::Multiple(42.0, "Allow".to_string()); + assert_eq!(config, config.clone()); +} + +#[test] +fn enum_config_should_impl_display() { + let config = TestEnumConfig::Multiple(42.0, "Allow".to_string()); + assert_eq!(burn::config::config_to_json(&config), config.to_string()); +} + +#[test] +fn struct_config_can_load_binary() { + let config = TestStructConfig::new(2, 3.0, "Allow".to_string(), TestEmptyStructConfig::new()); + + let binary = config_to_json(&config).as_bytes().to_vec(); + + let config_loaded = TestStructConfig::load_binary(&binary).unwrap(); + assert_eq!(config, config_loaded); +} diff --git a/crates/burn-core/tests/test_derive_module.rs b/crates/burn-core/tests/test_derive_module.rs new file mode 100644 index 0000000..096bd3e --- /dev/null +++ b/crates/burn-core/tests/test_derive_module.rs @@ -0,0 +1,772 @@ +use burn::module::Initializer; +use burn::module::{Module, Param}; +use burn::tensor::{Int, Tensor}; +use burn_core as burn; +use burn_tensor::Device; + +#[derive(Module, Debug)] +pub struct ModuleBasic { + weight_basic: Param>, +} + +#[derive(Module, Debug)] +#[allow(unused)] +struct ModuleTensorConstInt { + weight_basic: Tensor<2, Int>, +} + +impl ModuleBasic { + fn new(device: &Device) -> Self { + Self { + weight_basic: Initializer::Normal { + std: 1.0, + mean: 0.0, + } + .init([20, 20], device), + } + } +} + +#[derive(Module, Debug)] +struct ModuleWithConstGeneric { + modules: [ModuleBasic; N], +} + +#[derive(Module, Debug)] +struct ModuleWithGenericModule { + module: M, +} + +#[derive(Module, Debug)] +#[allow(clippy::large_enum_variant)] +enum ModuleEnum { + Basic(ModuleBasic), + Composed(ModuleComposed), +} + +#[derive(Module, Debug)] +#[allow(unused)] +enum ModuleEnumNested { + AnotherEnum(ModuleEnum), +} + +#[derive(Module, Debug)] +#[allow(clippy::large_enum_variant)] +enum ModuleEnumWithGenericModule { + Basic(ModuleBasic), + Generic(ModuleWithGenericModule), +} + +#[derive(Module, Debug)] +pub struct ModuleComposed { + weight: Param>, + basic: ModuleBasic, + tuple: (ModuleBasic, ModuleBasic), +} + +impl ModuleComposed { + fn new(device: &Device) -> Self { + let weight = Initializer::Normal { + std: 1.0, + mean: 0.0, + } + .init([20, 20], device); + + Self { + weight, + basic: ModuleBasic::new(device), + tuple: (ModuleBasic::new(device), ModuleBasic::new(device)), + } + } +} + +#[derive(Debug, Clone)] +pub enum PaddingConfig { + Default, + Other, +} + +#[derive(Module, Debug)] +pub struct ModuleWithAttributes { + /// A normal parameter. + weight: Param>, + /// A nested module. + nested: ModuleEnumWithGenericModule, + /// By default, primitives were not persistent (same as `#[module(skip)]`). + other_prob: f64, + /// By default, tensors were not persistent and not visited/mapped (same as `#[module(skip)]`). + tensor: Tensor<1>, + /// A field that is recomputed at runtime. + #[module(skip)] + cached_mask: Option>, + /// A field that contains some debug state. + debug_state: String, + /// Hint required: this generic is NOT a module. + #[module(skip)] + config: N, +} + +impl ModuleWithAttributes { + fn new(device: &Device) -> Self { + let basic = ModuleBasic::new(device); + let weight = basic.weight_basic.clone(); + + Self { + weight, + nested: ModuleEnumWithGenericModule::Basic(basic), + other_prob: 1., + tensor: Tensor::ones([2], device), + cached_mask: Some(Tensor::ones([2, 2], device)), + debug_state: "Hello World".into(), + config: PaddingConfig::Default, + } + } +} + +pub fn test_device() -> Device { + burn_tensor::Device::flex() +} + +mod state { + use super::*; + use burn::store::RecordError; + + #[test] + fn should_load_from_record_basic() { + let device = test_device(); + let module_1 = ModuleBasic::new(&device); + let module_2 = ModuleBasic::new(&device); + + // Access module_1 to trigger initialization before cloning. + // Cloning an uninitialized module preserves lazy state (no memory allocation), + // so we need to initialize first if we want the clone to have the same values. + assert_ne!( + module_1.weight_basic.to_data(), + module_2.weight_basic.to_data() + ); + + let record = module_1.clone().into_record(); + let module_2 = module_2.load_record(record); + + assert_eq!( + module_1.weight_basic.to_data(), + module_2.weight_basic.to_data() + ); + } + + #[test] + fn should_load_from_record_compose() { + let device = test_device(); + let module_1 = ModuleComposed::new(&device); + let module_2 = ModuleComposed::new(&device); + assert_ne!(module_1.weight.to_data(), module_2.weight.to_data()); + assert_ne!( + module_1.basic.weight_basic.to_data(), + module_2.basic.weight_basic.to_data() + ); + + let record = module_1.clone().into_record(); + let module_2 = module_2.load_record(record); + + assert_eq!(module_1.weight.to_data(), module_2.weight.to_data()); + assert_eq!( + module_1.basic.weight_basic.to_data(), + module_2.basic.weight_basic.to_data() + ); + // The tuple of nested modules round-trips too. + assert_eq!( + module_1.tuple.0.weight_basic.to_data(), + module_2.tuple.0.weight_basic.to_data() + ); + assert_eq!( + module_1.tuple.1.weight_basic.to_data(), + module_2.tuple.1.weight_basic.to_data() + ); + } + + #[test] + fn should_load_from_record_enum() { + let device = test_device(); + let module_1 = ModuleEnum::Basic(ModuleBasic::new(&device)); + let module_2 = ModuleEnum::Basic(ModuleBasic::new(&device)); + + // Trigger initialization before cloning so clone has the same values. + let ModuleEnum::Basic(ref module_1_basic) = module_1 else { + panic!("Invalid module type") + }; + let ModuleEnum::Basic(ref module_2_basic) = module_2 else { + panic!("Invalid module type") + }; + assert_ne!( + module_1_basic.weight_basic.to_data(), + module_2_basic.weight_basic.to_data() + ); + + let record = module_1.clone().into_record(); + let module_2 = module_2.load_record(record); + + let ModuleEnum::Basic(module_2_basic) = module_2 else { + panic!("Invalid module type") + }; + assert_eq!( + module_1_basic.weight_basic.to_data(), + module_2_basic.weight_basic.to_data() + ); + } + + #[test] + fn should_load_from_record_const_generic() { + let device = test_device(); + let module_1 = ModuleWithConstGeneric { + modules: [ModuleBasic::new(&device), ModuleBasic::new(&device)], + }; + let module_2 = ModuleWithConstGeneric { + modules: [ModuleBasic::new(&device), ModuleBasic::new(&device)], + }; + + // Trigger initialization before cloning so clone has the same values. + assert_ne!( + module_1.modules[0].weight_basic.to_data(), + module_2.modules[0].weight_basic.to_data(), + ); + assert_ne!( + module_1.modules[1].weight_basic.to_data(), + module_2.modules[1].weight_basic.to_data(), + ); + + let record = module_1.clone().into_record(); + let module_2 = module_2.load_record(record); + + assert_eq!( + module_1.modules[0].weight_basic.to_data(), + module_2.modules[0].weight_basic.to_data(), + ); + assert_eq!( + module_1.modules[1].weight_basic.to_data(), + module_2.modules[1].weight_basic.to_data(), + ); + } + + #[test] + fn should_load_from_record_based_on_attributes() { + let device = test_device(); + let mut module_1 = ModuleWithAttributes::new(&device); + let module_2 = ModuleWithAttributes::new(&device); + + assert_ne!(module_1.weight.to_data(), module_2.weight.to_data()); + + let ModuleEnumWithGenericModule::Basic(ref m1_basic) = module_1.nested else { + panic!("Invalid module type") + }; + let ModuleEnumWithGenericModule::Basic(ref m2_basic) = module_2.nested else { + panic!("Invalid module type") + }; + assert_ne!( + m1_basic.weight_basic.to_data(), + m2_basic.weight_basic.to_data(), + ); + + // Alter the non-persistent fields (skipped fields and plain tensors) to validate that they + // are not captured by the module record. + module_1.cached_mask = Some(module_1.cached_mask.unwrap() * 2); + module_1.tensor = module_1.tensor * 2; + module_1.other_prob = 0.; + module_1.debug_state = "Hello World!".into(); + module_1.config = PaddingConfig::Other; + + let record = module_1.clone().into_record(); + let module_2 = module_2.load_record(record); + + // Modules & params are loaded. + assert_eq!(module_1.weight.to_data(), module_2.weight.to_data()); + let ModuleEnumWithGenericModule::Basic(m2_basic) = module_2.nested else { + panic!("Invalid module type") + }; + assert_eq!( + m1_basic.weight_basic.to_data(), + m2_basic.weight_basic.to_data(), + ); + + // `#[module(skip)]` fields and plain (non-param) tensors are not persisted: module_2 keeps + // its own initialization. + assert_ne!(module_1.other_prob, module_2.other_prob); + assert_ne!(module_1.debug_state, module_2.debug_state); + assert!(matches!(module_1.config, PaddingConfig::Other)); + assert!(matches!(module_2.config, PaddingConfig::Default)); + assert_ne!(module_1.tensor.to_data(), module_2.tensor.to_data()); + assert_ne!( + module_1.cached_mask.as_ref().unwrap().to_data(), + module_2.cached_mask.as_ref().unwrap().to_data() + ); + } + + #[test] + fn incompatible_enum_variant_record_fails_to_load() { + let device = test_device(); + let module_1 = ModuleEnum::Basic(ModuleBasic::new(&device)); + let module_2 = ModuleEnum::Composed(ModuleComposed::new(&device)); + + // The Basic variant's record does not provide the Composed variant's params, so a strict + // (default) load reports the missing tensors as a validation error. + let record = module_1.into_record(); + let result = module_2.try_load_record(record); + assert!(matches!(result, Err(RecordError::Validation(_)))); + } +} + +mod lazy_clone { + use burn_tensor::Device; + + use super::*; + + #[test] + fn clone_uninitialized_param_should_not_trigger_init() { + let device = test_device(); + let module = ModuleBasic::new(&device); + + // Module starts uninitialized (lazy). + assert!(!module.weight_basic.is_initialized()); + + // Cloning should preserve the lazy state, not trigger initialization. + let cloned = module.clone(); + assert!(!module.weight_basic.is_initialized()); + assert!(!cloned.weight_basic.is_initialized()); + } + + #[test] + fn clone_initialized_param_should_share_values() { + let device = test_device(); + let module = ModuleBasic::new(&device); + + // Force initialization by accessing the tensor. + let _ = module.weight_basic.to_data(); + assert!(module.weight_basic.is_initialized()); + + // Clone of an initialized param should have the same values. + let cloned = module.clone(); + assert_eq!(module.weight_basic.to_data(), cloned.weight_basic.to_data()); + } + + #[test] + fn lazy_clone_should_produce_valid_tensor_on_access() { + let device = test_device(); + let module = ModuleBasic::new(&device); + let cloned = module.clone(); + + // Both are uninitialized. + assert!(!module.weight_basic.is_initialized()); + assert!(!cloned.weight_basic.is_initialized()); + + // Accessing the clone should produce a valid tensor with the right shape. + let data = cloned.weight_basic.to_data(); + assert_eq!(data.shape, [20, 20].into()); + + data.assert_eq(&module.weight_basic.val().into_data(), true); + } + + #[test] + fn lazy_clone_and_original_have_same_init() { + let device = test_device(); + let module = ModuleBasic::new(&device); + let cloned = module.clone(); + + let clone_data = cloned.weight_basic.to_data(); + let orig_data = module.weight_basic.to_data(); + + assert_eq!(clone_data.shape, [20, 20].into()); + assert_eq!(orig_data.shape, [20, 20].into()); + assert_eq!(clone_data, orig_data); + } + + #[test] + fn lazy_clone_deref_should_trigger_init() { + let device = test_device(); + let module = ModuleBasic::new(&device); + let cloned = module.clone(); + + // Access via Deref (shape() uses Deref, not val()) on the clone. + let shape = cloned.weight_basic.shape(); + assert_eq!(shape, [20, 20].into()); + assert!(cloned.weight_basic.is_initialized()); + assert!(module.weight_basic.is_initialized()); + } + + #[test] + fn init_mapper_on_lazy_clone_should_not_affect_original() { + use burn::module::ParamId; + use burn::tensor::Shape; + + let device: Device = test_device(); + + // Create two uninitialized params from the same init function. + let param: Param> = Param::uninitialized( + ParamId::new(), + move |d, _| Tensor::random([4, 4], Default::default(), d), + device, + false, + Shape::from([4, 4]), + ); + + let mut cloned = param.clone(); + + // Apply init_mapper on the clone to double all values. + cloned = cloned.init_mapper(|t| t.mul_scalar(2.0)); + + // Random tensor should still have the same initialization point, but * 2.0 + cloned + .val() + .div_scalar(2.0) + .into_data() + .assert_approx_eq::(¶m.val().into_data(), Default::default()); + } + + #[test] + fn load_record_into_uninitialized_module_should_work() { + let device = test_device(); + let module_1 = ModuleBasic::new(&device); + + // Initialize module_1 so we have a record to load. + let _ = module_1.weight_basic.to_data(); + let record = module_1.clone().into_record(); + + // Create a fresh uninitialized module and load weights into it. + let module_2 = ModuleBasic::new(&device); + assert!(!module_2.weight_basic.is_initialized()); + + let module_2 = module_2.load_record(record); + + // After loading, the param should be initialized with the loaded values. + assert_eq!( + module_1.weight_basic.to_data(), + module_2.weight_basic.to_data() + ); + } +} + +mod num_params { + use super::*; + + #[test] + fn should_calculate_num_params_basic() { + let device = test_device(); + let module = ModuleBasic::new(&device); + assert_eq!(20 * 20, module.num_params()); + } + + #[test] + fn should_output_state_composed() { + let device = test_device(); + let module = ModuleComposed::new(&device); + assert_eq!(4 * 20 * 20, module.num_params()); + } + + #[test] + fn should_calculate_num_params_enum() { + let device = test_device(); + let module = ModuleEnum::Basic(ModuleBasic::new(&device)); + assert_eq!(20 * 20, module.num_params()); + + let module = ModuleEnum::Composed(ModuleComposed::new(&device)); + assert_eq!(4 * 20 * 20, module.num_params()); + } + + #[test] + fn should_calculate_num_params_based_on_attributes() { + let device = test_device(); + let module = ModuleWithAttributes::new(&device); + assert_eq!(20 * 20 * 2, module.num_params()); + } +} + +#[cfg(all(feature = "std", feature = "autodiff"))] +mod require_grad { + use burn_tensor::TensorData; + use rand::{ + SeedableRng, + rngs::{StdRng, SysRng}, + }; + + use super::*; + + #[test] + fn should_have_grad_by_default() { + let device = test_device().autodiff(); + let module = ModuleBasic::new(&device); + let grad_x = calculate_grads(&module, |weights, x| weights.matmul(x)); + + assert!(grad_x.is_some()); + } + + #[test] + fn should_have_no_grad_after_no_grad() { + let device = test_device().autodiff(); + let module = ModuleBasic::new(&device).no_grad(); + let grad_x = calculate_grads(&module, |weights, x| weights.matmul(x)); + + assert!(grad_x.is_none()); + } + + #[test] + fn should_have_grad_when_from_record() { + let device = test_device().autodiff(); + let module = ModuleBasic::new(&device); + + // Round-tripping through a record must preserve the params' require_grad state. + let record = module.clone().into_record(); + let module = module.load_record(record); + + let grad_x = calculate_grads(&module, |weights, x| weights.matmul(x)); + + assert!(grad_x.is_some()); + } + + fn calculate_grads( + module: &ModuleBasic, + transformation: fn(Tensor<2>, Tensor<2>) -> Tensor<2>, + ) -> Option> { + let device = module.weight_basic.device(); + let data = TensorData::random::( + module.weight_basic.shape(), + burn_tensor::Distribution::Default, + &mut StdRng::try_from_rng(&mut SysRng).unwrap(), + ); + let x = Tensor::from_data(data, &device).require_grad(); + let t = module.weight_basic.val(); + let y = transformation(t, x); + + let mut grads = y.backward(); + module.weight_basic.grad_remove(&mut grads) + } +} + +#[cfg(feature = "cuda")] +mod grad_distributed { + use burn_tensor::TensorData; + use burn_tensor::distributed::{DistributedContext, ReduceOperation}; + use burn_tensor::{Device, DeviceType, Tolerance}; + use rand::{ + SeedableRng, + rngs::{StdRng, SysRng}, + }; + use serial_test::serial; + use std::sync::mpsc::{Receiver, Sender}; + + use super::*; + + #[test] + #[serial] + fn sharded_module_should_sync_gradients_sum() { + compare_sync_gradient(ReduceOperation::Sum, |weights, x| weights.matmul(x)); + } + + #[test] + #[serial] + fn sharded_module_should_sync_gradients_mean() { + compare_sync_gradient(ReduceOperation::Mean, |weights, x| weights.matmul(x)); + } + + #[test] + #[serial] + fn sharded_module_should_sync_gradients_sum_residual() { + compare_sync_gradient(ReduceOperation::Sum, |weights, x| { + let y = weights.clone().matmul(x); + y.add(weights) + }); + } + + #[test] + #[serial] + fn sharded_module_should_sync_gradients_mean_residual() { + compare_sync_gradient(ReduceOperation::Mean, |weights, x| { + let y = weights.clone().matmul(x); + y.add(weights) + }); + } + + #[test] + #[serial] + fn sharded_module_should_sync_gradients_sum_activation() { + compare_sync_gradient(ReduceOperation::Sum, |weights, x| { + let y = weights.clone().matmul(x); + let y = y.add(weights); + burn_tensor::activation::relu(y) + }); + } + + #[test] + #[serial] + fn sharded_module_should_sync_gradients_mean_activation() { + compare_sync_gradient(ReduceOperation::Mean, |weights, x| { + let y = weights.clone().matmul(x); + let y = y.add(weights); + burn_tensor::activation::relu(y) + }); + } + + #[test] + #[serial] + fn sharded_module_should_sync_gradients_sum_diamond_graph() { + compare_sync_gradient(ReduceOperation::Sum, |weights, x| { + let left = weights.clone().matmul(x.clone().mul_scalar(2)); + let right = weights.clone().matmul(x.clone().exp()); + Tensor::cat(vec![left, right], 0) + }); + } + + #[test] + #[serial] + fn sharded_module_should_sync_gradients_mean_diamond_graph() { + compare_sync_gradient(ReduceOperation::Mean, |weights, x| { + let left = weights.clone().matmul(x.clone().mul_scalar(2)); + let right = weights.clone().matmul(x.clone().exp()); + Tensor::cat(vec![left, right], 0) + }); + } + + fn compare_sync_gradient( + op: ReduceOperation, + transformation: fn(Tensor<2>, Tensor<2>) -> Tensor<2>, + ) { + use burn_tensor::distributed::DistributedConfig; + + const NUM_ITERATIONS: usize = 100; + let devices = Device::enumerate(DeviceType::Cuda).autodiff().into_vec(); + + let module = ModuleBasic::new(&devices[0]); + let (synced_senders, synced_receivers): ( + Vec>, + Vec>, + ) = (0..devices.len()) + .map(|_| std::sync::mpsc::channel()) + .unzip(); + let (original_senders, original_receivers): ( + Vec>, + Vec>, + ) = (0..devices.len()) + .map(|_| std::sync::mpsc::channel()) + .unzip(); + + let config = DistributedConfig { all_reduce_op: op }; + let _context = DistributedContext::init(devices.clone(), config); + + let join_handles = spawn_peer_threads( + &module, + &devices, + synced_senders, + original_senders, + transformation, + NUM_ITERATIONS, + ); + + for _ in 0..NUM_ITERATIONS { + let device = devices.first().unwrap(); + let mut expected: Tensor<2> = + Tensor::from_data(original_receivers.first().unwrap().recv().unwrap(), device); + for r in original_receivers[1..].iter().by_ref() { + let data = r.recv().unwrap(); + expected = expected.add(Tensor::from_data(data, device)); + } + if op == ReduceOperation::Mean { + expected = expected.div_scalar(original_receivers.len() as f32); + } + + for r in synced_receivers.iter().by_ref() { + let data = r.recv().unwrap(); + data.assert_approx_eq::(&expected.to_data(), Tolerance::default()); + } + } + + for handle in join_handles { + handle.join().unwrap(); + } + + // DistributedContext goes out of scope -> close_communication_server + } + + fn spawn_peer_threads( + module: &ModuleBasic, + devices: &[Device], + synced_senders: Vec>, + original_senders: Vec>, + transformation: fn(Tensor<2>, Tensor<2>) -> Tensor<2>, + num_iter: usize, + ) -> Vec> { + let mut handles = vec![]; + + for i in 0..devices.len() { + let module_clone = module.clone(); + let device = devices[i].clone(); + let synced_sender = synced_senders[i].clone(); + let original_sender = original_senders[i].clone(); + handles.push(std::thread::spawn(move || { + run_peer_sharded( + &module_clone, + synced_sender, + original_sender, + transformation, + device, + num_iter, + ) + })); + } + + handles + } + + pub fn run_peer_sharded( + module: &ModuleBasic, + synced_sender: Sender, + original_sender: Sender, + transformation: fn(Tensor<2>, Tensor<2>) -> Tensor<2>, + device: Device, + num_iter: usize, + ) { + let mut module = module.clone().fork(&device); + + for _ in 0..num_iter { + module = set_distributed(&module, &device); + let (grads_synced, grads_original) = calculate_grads(&module, transformation); + + let data = grads_original.unwrap().to_data(); + original_sender.clone().send(data).unwrap(); + + let data = grads_synced.unwrap().to_data(); + synced_sender.clone().send(data).unwrap(); + } + } + + fn set_distributed(module: &ModuleBasic, device: &Device) -> ModuleBasic { + let mut module = module.clone().fork(&device); + let (id, tensor, mapper) = module.weight_basic.consume(); + let tensor = tensor.set_distributed(id); + module.weight_basic = Param::from_mapped_value(id, tensor, mapper); + module + } + + fn calculate_grads( + module: &ModuleBasic, + transformation: fn(Tensor<2>, Tensor<2>) -> Tensor<2>, + ) -> (Option>, Option>) { + let device = module.weight_basic.device(); + let data = TensorData::random::( + module.weight_basic.shape(), + burn_tensor::Distribution::Default, + &mut StdRng::try_from_rng(&mut SysRng).unwrap(), + ); + let x = Tensor::from_data(data.clone(), &device).require_grad(); + let t = module.weight_basic.val(); + let y = transformation(t, x); + + let mut grads = y.backward(); + let grads_synced = module.weight_basic.grad_remove(&mut grads); + + // Kind of hacky, but running the backward pass again without marking the tensor as distributed will not sync gradients. + // We can use this to compute the expected sum. + let x = Tensor::from_data(data, &device).require_grad(); + let t = module.weight_basic.val(); + let y = transformation(t, x); + let mut grads = y.backward(); + let grads_original = module.weight_basic.grad_remove(&mut grads); + (grads_synced, grads_original) + } +} diff --git a/crates/burn-core/tests/test_record.rs b/crates/burn-core/tests/test_record.rs new file mode 100644 index 0000000..6d937ac --- /dev/null +++ b/crates/burn-core/tests/test_record.rs @@ -0,0 +1,134 @@ +//! Integration test for the non-generic record system (`burn_core::store`). +//! +//! Saves a (nested) module to a burnpack record, loads it back into a freshly initialized +//! module, and checks the parameters round-trip — exercising the public API exactly as a user +//! would (`into_record` / `save` / `load` / `load_record`). + +use burn::module::{Module, Param}; +use burn::store::ModuleRecord; +use burn::tensor::Tensor; +use burn_core as burn; +use burn_tensor::Device; + +#[derive(Module, Debug)] +struct Layer { + weight: Param>, + bias: Param>, +} + +impl Layer { + fn from_values(weight: [[f32; 2]; 2], bias: [f32; 2], device: &Device) -> Self { + Self { + weight: Param::from_data(weight, device), + bias: Param::from_data(bias, device), + } + } + + fn zeros(device: &Device) -> Self { + Self::from_values([[0.0; 2]; 2], [0.0; 2], device) + } + + fn values(&self) -> (Vec, Vec) { + ( + self.weight.val().to_data().to_vec().unwrap(), + self.bias.val().to_data().to_vec().unwrap(), + ) + } +} + +/// A module with only the `first` layer — its record holds `first.weight` / `first.bias`, +/// a strict subset of [`Mlp`]'s parameters. +#[derive(Module, Debug)] +struct FirstOnly { + first: Layer, +} + +/// Nested module so the record exercises path building (`first.weight`, `second.bias`, ...). +#[derive(Module, Debug)] +struct Mlp { + first: Layer, + second: Layer, +} + +impl Mlp { + fn sample(device: &Device) -> Self { + Self { + first: Layer::from_values([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], device), + second: Layer::from_values([[7.0, 8.0], [9.0, 10.0]], [11.0, 12.0], device), + } + } + + fn zeros(device: &Device) -> Self { + Self { + first: Layer::zeros(device), + second: Layer::zeros(device), + } + } +} + +fn assert_matches_sample(model: &Mlp) { + let (w1, b1) = model.first.values(); + let (w2, b2) = model.second.values(); + assert_eq!(w1, vec![1.0, 2.0, 3.0, 4.0]); + assert_eq!(b1, vec![5.0, 6.0]); + assert_eq!(w2, vec![7.0, 8.0, 9.0, 10.0]); + assert_eq!(b2, vec![11.0, 12.0]); +} + +#[cfg(feature = "std")] +#[test] +fn save_and_load_module_via_file() { + let device = Default::default(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("mlp.bpk"); + + // Save a trained module to a record on disk. + let record = Mlp::sample(&device).into_record(); + assert_eq!(record.len(), 4); // first.weight, first.bias, second.weight, second.bias + record.save(&path).unwrap(); + + // Load the record and apply it onto a freshly initialized (zeroed) module. + let record = ModuleRecord::load(&path).unwrap(); + let loaded = Mlp::zeros(&device).load_record(record); + + assert_matches_sample(&loaded); +} + +#[test] +fn save_and_load_module_via_bytes() { + let device = Default::default(); + + let bytes = Mlp::sample(&device).into_record().into_bytes().unwrap(); + + let record = ModuleRecord::from_bytes(bytes).unwrap(); + let loaded = Mlp::zeros(&device).load_record(record); + + assert_matches_sample(&loaded); +} + +#[test] +fn missing_parameters_require_allow_partial() { + let device = Default::default(); + + // A record holding only `first.weight` / `first.bias` (a subset of the full Mlp). + let partial = FirstOnly { + first: Layer::from_values([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device), + } + .into_record(); + let bytes = partial.into_bytes().unwrap(); + + // Strict load into the full Mlp fails: `second.*` is missing from the record. + let strict = ModuleRecord::from_bytes(bytes.clone()).unwrap(); + assert!(Mlp::zeros(&device).try_load_record(strict).is_err()); + + // With `allow_partial`, the present params load and the rest keep their init. + let lenient = ModuleRecord::from_bytes(bytes).unwrap().allow_partial(true); + let loaded = Mlp::zeros(&device).load_record(lenient); + + let (w1, b1) = loaded.first.values(); + assert_eq!(w1, vec![1.0, 2.0, 3.0, 4.0]); + assert_eq!(b1, vec![5.0, 6.0]); + let (w2, b2) = loaded.second.values(); + assert_eq!(w2, vec![0.0, 0.0, 0.0, 0.0]); + assert_eq!(b2, vec![0.0, 0.0]); +} diff --git a/crates/burn-cpu/Cargo.toml b/crates/burn-cpu/Cargo.toml new file mode 100644 index 0000000..cb22c8e --- /dev/null +++ b/crates/burn-cpu/Cargo.toml @@ -0,0 +1,44 @@ +[package] +authors = ["marcantoinem "] +categories = ["science"] +description = "MLIR based CPU backend for the Burn framework" +documentation = "https://docs.rs/burn-cpu" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "cpu"] +license.workspace = true +name = "burn-cpu" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-cpu" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["std", "fusion", "autotune", "burn-cubecl/default", "cubecl/default"] +doc = ["burn-cubecl/doc"] +std = ["burn-cubecl/std", "cubecl/std"] + +tracing = [ + "burn-backend/tracing", + "burn-cubecl/tracing", + "burn-fusion?/tracing", + "cubecl/tracing", +] + +fusion = ["burn-fusion", "burn-cubecl/fusion"] +autotune = ["burn-cubecl/autotune"] +autotune-checks = ["burn-cubecl/autotune-checks"] + +[dependencies] +burn-fusion = { workspace = true, optional = true, features = ["default"] } +burn-cubecl = { workspace = true } +burn-backend = { workspace = true, features = [ + "default", + "cubecl-cpu", +] } +cubecl = { workspace = true, features = ["cpu"] } + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-cpu/README.md b/crates/burn-cpu/README.md new file mode 100644 index 0000000..e1f1639 --- /dev/null +++ b/crates/burn-cpu/README.md @@ -0,0 +1,12 @@ +# Burn CPU Backend + +[Burn](https://github.com/tracel-ai/burn) CubeCL CPU backend + +[![Current Crates.io Version](https://img.shields.io/crates/v/burn-cuda.svg)](https://crates.io/crates/burn-cuda) + +This crate provides a MLIR based CPU backend for [Burn](https://github.com/tracel-ai/burn) using the +[cubecl](https://github.com/tracel-ai/cubecl.git) crates. + +## Usage Example + +Example coming soon diff --git a/crates/burn-cpu/src/lib.rs b/crates/burn-cpu/src/lib.rs new file mode 100644 index 0000000..a6e5d65 --- /dev/null +++ b/crates/burn-cpu/src/lib.rs @@ -0,0 +1,44 @@ +#![cfg_attr(docsrs, feature(doc_cfg))] + +extern crate alloc; + +use burn_cubecl::CubeBackend; +pub use cubecl::cpu::CpuDevice; +use cubecl::cpu::CpuRuntime; + +#[cfg(not(feature = "fusion"))] +pub type Cpu = CubeBackend; + +#[cfg(feature = "fusion")] +pub type Cpu = burn_fusion::Fusion>; + +#[cfg(test)] +mod tests { + use super::*; + use burn_backend::{Backend, BoolStore, DType, DeviceOps}; + + #[test] + fn should_support_dtypes() { + type B = Cpu; + let device = CpuDevice; + let scheme = device.defaults().quantization.scheme; + + assert!(B::supports_dtype(&device, DType::F64)); + assert!(B::supports_dtype(&device, DType::F32)); + assert!(B::supports_dtype(&device, DType::F16)); + assert!(B::supports_dtype(&device, DType::BF16)); + assert!(B::supports_dtype(&device, DType::I64)); + assert!(B::supports_dtype(&device, DType::I32)); + assert!(B::supports_dtype(&device, DType::I16)); + assert!(B::supports_dtype(&device, DType::I8)); + assert!(B::supports_dtype(&device, DType::U64)); + assert!(B::supports_dtype(&device, DType::U32)); + assert!(B::supports_dtype(&device, DType::U16)); + assert!(B::supports_dtype(&device, DType::U8)); + assert!(B::supports_dtype(&device, DType::QFloat(scheme))); + + // Currently not registered in supported types + assert!(!B::supports_dtype(&device, DType::Flex32)); + assert!(!B::supports_dtype(&device, DType::Bool(BoolStore::Native))); + } +} diff --git a/crates/burn-cubecl-fusion/Cargo.toml b/crates/burn-cubecl-fusion/Cargo.toml new file mode 100644 index 0000000..5d499a1 --- /dev/null +++ b/crates/burn-cubecl-fusion/Cargo.toml @@ -0,0 +1,59 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science"] +description = "Provide optimizations that can be used with cubecl based backends." +documentation = "https://docs.rs/burn-cubecl-fusion" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "gpu"] +license.workspace = true +name = "burn-cubecl-fusion" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-cubecl-fusion" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["autotune", "std", "cubecl/default", "burn-fusion/default"] + +autotune = [] +autotune-checks = ["cubecl/autotune-checks", "half"] +doc = ["default"] +test-util = [] +std = ["cubecl/std", "burn-backend/std", "burn-fusion/std"] +tracing = [ + "cubecl/tracing", + "burn-std/tracing", + "burn-backend/tracing", + "burn-fusion/tracing", +] + +[dependencies] +burn-fusion = { workspace = true } +burn-ir = { workspace = true } +burn-std = { workspace = true, features = ["default"] } +cubecl = { workspace = true } +cubek = { workspace = true, features = [ + "matmul", + "reduce", + "quantization", + "stdlib", +] } +half = { workspace = true, optional = true } + +# Needed for the cubecl dtype helpers (`burn_backend::cubecl::*`). +burn-backend = { workspace = true, features = ["cubecl"] } + +derive-new = { workspace = true } +hashbrown = { workspace = true } +serde = { workspace = true } +spin = { workspace = true } + +[dev-dependencies] +cubecl = { workspace = true, features = ["test-runtime"] } +half = { workspace = true } + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-cubecl-fusion/README.md b/crates/burn-cubecl-fusion/README.md new file mode 100644 index 0000000..88c5797 --- /dev/null +++ b/crates/burn-cubecl-fusion/README.md @@ -0,0 +1,3 @@ +# Burn CubeCl Fusion + +Provide optimizations that can be used with [cubecl](../burn-cubecl) based backends. diff --git a/crates/burn-cubecl-fusion/src/base.rs b/crates/burn-cubecl-fusion/src/base.rs new file mode 100644 index 0000000..9c26cea --- /dev/null +++ b/crates/burn-cubecl-fusion/src/base.rs @@ -0,0 +1,124 @@ +use burn_fusion::stream::Context; +use burn_std::{DType, Shape, Strides, quantization::QParamTensor, strides}; +use cubecl::quant::scheme::{QuantParam, QuantScheme}; +use cubecl::{ + Runtime, + client::ComputeClient, + ir::AddressType, + prelude::{TensorArg, TensorBinding}, +}; +use std::marker::PhantomData; + +/// Defines a fallback operation when fusion isn't possible. +pub trait FallbackOperation: Send + Sync { + /// Executes the fallback procedure. + fn run(&self, context: &mut Context>); +} + +/// Runtime parameters for quantization. Can be used to construct a scales handle from the base +/// tensor handle. +pub type QParams = burn_std::quantization::QParams; + +/// Handle to be used when fusing operations. +pub struct CubeFusionHandle { + /// Compute client for jit. + pub client: ComputeClient, + /// The buffer where the data are stored. + pub handle: cubecl::server::Handle, + /// The device of the current tensor. + pub device: R::Device, + /// The element type of the tensor. + pub dtype: DType, + /// The strides of the tensor. + pub strides: Strides, + /// Quantization runtime parameters, if applicable + pub qparams: Option, +} + +impl core::fmt::Debug for CubeFusionHandle { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_fmt(format_args!( + "CubeFusionHandle {{ device: {:?}, runtime: {}}}", + self.device, + R::name(&self.client), + )) + } +} + +impl Clone for CubeFusionHandle { + fn clone(&self) -> Self { + Self { + client: self.client.clone(), + handle: self.handle.clone(), + device: self.device.clone(), + strides: self.strides.clone(), + dtype: self.dtype, + qparams: self.qparams.clone(), + } + } +} + +unsafe impl Send for CubeFusionHandle {} +unsafe impl Sync for CubeFusionHandle {} + +impl CubeFusionHandle { + /// Return the reference to a tensor handle. + pub fn binding(self, shape: Shape) -> TensorBinding { + TensorBinding { + handle: self.handle.binding(), + strides: self.strides.clone(), + shape, + runtime: PhantomData, + } + } + + pub fn required_address_type(&self) -> AddressType { + match self.dtype { + DType::QFloat(scheme) => { + let len = self.handle.size() as usize * 8 / scheme.size_bits_value(); + AddressType::from_len(len) + } + _ => AddressType::from_len(self.handle.size() as usize / self.dtype.size()), + } + } + + /// Return the reference to a tensor argument. + pub fn into_tensor_arg(self, shape: Shape) -> TensorArg { + let handle = self.binding(shape); + handle.into_tensor_arg() + } + + /// Construct a separate tensor for the quantization scales, if present + pub fn params(&self, scheme: QuantScheme) -> Option { + let qparams = self.qparams.as_ref()?; + let mut handle = self.handle.clone(); + handle.offset_start = Some(qparams.scales.offset_start as u64); + handle.offset_end = Some(qparams.scales.offset_end as u64); + + Some(Self { + client: self.client.clone(), + handle, + device: self.device.clone(), + dtype: match scheme.param { + QuantParam::F32 => DType::F32, + QuantParam::F16 => DType::F16, + QuantParam::BF16 => DType::BF16, + QuantParam::UE8M0 | QuantParam::UE4M3 => unimplemented!("Not yet supported"), + }, + strides: qparams.scales.metadata.strides().clone(), + qparams: None, + }) + } +} + +pub(crate) fn strides_dyn_rank(shape: &[usize]) -> Strides { + let mut strides = strides![0; shape.len()]; + + let mut current = 1; + shape.iter().enumerate().rev().for_each(|(index, val)| { + strides[index] = current; + current *= val; + }); + + strides +} diff --git a/crates/burn-cubecl-fusion/src/engine/codegen/base.rs b/crates/burn-cubecl-fusion/src/engine/codegen/base.rs new file mode 100644 index 0000000..c31d028 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/codegen/base.rs @@ -0,0 +1,5 @@ +use cubecl::{define_scalar, define_size, prelude::Vector}; + +define_scalar!(pub DynElem); +define_size!(pub DynSize); +pub type DynVector = Vector; diff --git a/crates/burn-cubecl-fusion/src/engine/codegen/io.rs b/crates/burn-cubecl-fusion/src/engine/codegen/io.rs new file mode 100644 index 0000000..9e108bd --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/codegen/io.rs @@ -0,0 +1,875 @@ +//! This module declares input-output primitives to read and write values during kernel expansion. +use crate::engine::codegen::{DynElem, DynSize}; + +use super::{ir::*, tensor::GlobalTensor}; +use burn_std::quantization::QuantScheme; +use cubecl::quant::scheme::QuantLevel; +use cubecl::{ + intrinsic, + ir::Value, + prelude::*, + std::{FastDivmod, tensor::View}, +}; +use cubek::quantization::layout::{BlockScaledLayout, PerTensorLayout, ScalesLayout}; +use serde::{Deserialize, Serialize}; + +/// Define how a tensor might be transformed at runtime. +#[derive(Clone, Debug, Hash, PartialEq, Eq, Serialize, Deserialize, PartialOrd, Ord)] +pub enum Transform { + /// A reshape operation has been registered on a tensor. + /// + /// This enum entry contains a sequence of [arguments](FuseArg) that points to global scalars representing the + /// new shape for the current tensor. + Reshape(Vec), + /// Two axes have been swapped on a tensor. + /// + /// The enum entry contains those two axes. + SwapDims(usize, usize), +} + +/// Reads the value from the [arg](FuseArg) and cast it to the generic cube primitive. +/// +/// # Notes +/// +/// The [global arguments](GlobalArgs) for both inputs and outputs as well as the +/// [local arguments](LocalArgs) need to be passed to this function. +/// +/// This is because the [argument](FuseArg) might point to a global input, output or local variable +/// created during kernel expansion. +#[cube] +pub fn read( + inputs: &GlobalArgs, + outputs: &GlobalArgs, + locals: &LocalArgs, + ref_pos: usize, + #[comptime] arg: FuseArg, + #[comptime] config: &FuseBlockConfig, +) -> Vector { + set_polyfill_typed::, DynElem, DynSize>(); + match arg { + FuseArg::Input(pos, _precision, layout) => { + let global = inputs.tensors.index(pos); + let vector_size = global.tensor.vector_size(); + + if comptime![!global.broadcasted && vector_size != config.width] { + read_input_aligned(inputs, locals, pos, ref_pos, layout, config, None) + } else { + read_input(inputs, locals, pos, ref_pos, layout, config, None) + } + } + FuseArg::MultiBlockLocal(key, _) | FuseArg::MultiBlockGlobal(key, _) => { + Vector::cast_from(outputs.variables.read(key)) + } + FuseArg::Output(pos, _precision, layout) => { + read_output(inputs, outputs, locals, pos, ref_pos, layout, config) + } + FuseArg::BlockLocal { pos, ty } => match comptime![ty] { + FuseType::F64 => Vector::cast_from(locals.l_f64.find(pos)), + FuseType::F32 | FuseType::Flex32 => Vector::cast_from(locals.l_f32.find(pos)), + FuseType::F16 => Vector::cast_from(locals.l_f16.find(pos)), + FuseType::BF16 => Vector::cast_from(locals.l_bf16.find(pos)), + FuseType::U64 => Vector::cast_from(locals.l_u64.find(pos)), + FuseType::U32 => Vector::cast_from(locals.l_u32.find(pos)), + FuseType::U16 => Vector::cast_from(locals.l_u16.find(pos)), + FuseType::U8 => Vector::cast_from(locals.l_u8.find(pos)), + FuseType::I64 => Vector::cast_from(locals.l_i64.find(pos)), + FuseType::I32 => Vector::cast_from(locals.l_i32.find(pos)), + FuseType::I16 => Vector::cast_from(locals.l_i16.find(pos)), + FuseType::I8 => Vector::cast_from(locals.l_i8.find(pos)), + }, + FuseArg::Scalar(..) => { + let scalar = read_scalar::(inputs, arg); + Vector::new(scalar) + } + FuseArg::ScalarShape(_) => { + let scalar = read_scalar_shape(inputs, arg); + Vector::cast_from(scalar) + } + FuseArg::Literal(val, _precision) => Vector::new(from_const_int::(val)), + FuseArg::InputReshaped { + original, + shape, + broadcasted, + } => match comptime![original.as_ref().clone()] { + FuseArg::Input(pos, _precision, layout) => { + let global = inputs.tensors.index(pos); + let vector_size = global.tensor.vector_size(); + + if comptime![!broadcasted && vector_size != config.width] { + read_input_aligned( + inputs, + locals, + pos, + ref_pos, + layout, + config, + comptime![Some(Transform::Reshape(shape))], + ) + } else { + read_input( + inputs, + locals, + pos, + ref_pos, + layout, + config, + comptime![Some(Transform::Reshape(shape))], + ) + } + } + _ => comptime![panic!("Only input can be reshaped")], + }, + FuseArg::InputSwapDims { + original, + dims, + broadcasted, + } => match comptime![original.as_ref().clone()] { + FuseArg::Input(pos, _precision, layout) => { + let global = inputs.tensors.index(pos); + let vector_size = global.tensor.vector_size(); + + if comptime![!broadcasted && vector_size != config.width] { + read_input_aligned( + inputs, + locals, + pos, + ref_pos, + layout, + config, + comptime![Some(Transform::SwapDims(dims.0, dims.1))], + ) + } else { + read_input( + inputs, + locals, + pos, + ref_pos, + layout, + config, + comptime![Some(Transform::SwapDims(dims.0, dims.1))], + ) + } + } + _ => comptime![panic!("Only input can be swapped dims")], + }, + } +} + +/// Computes the offset for the current global tensor with a quantized layout. +/// +/// The offset can be used to fetch the correct data from the quantized tensor as if it was in a +/// linear contiguous format. +#[cube] +fn index_offset_with_quant_layout( + tensor: &GlobalTensor, + locals: &LocalArgs, + index: usize, + #[comptime] rank: usize, + #[comptime] scheme: QuantScheme, +) -> usize { + let (start, end) = (0, rank - 1); + let num_quants = scheme.num_quants(); + + let offset_ref = index * locals.ref_vector_size; + let mut offset = 0; + + #[unroll] + for i in start..end { + let ogwl = offset_ref / locals.ref_strides[i]; + offset += ogwl % tensor.tensor.shape(i) * tensor.tensor.stride(i); + } + + // Handle packed representation in last dim + let ogwl = offset_ref / locals.ref_strides[end]; + let shape_last = tensor.tensor.shape(end).div_ceil(num_quants); + let stride_last = tensor.tensor.stride(end); + offset += (ogwl.div_ceil(num_quants)) % shape_last * stride_last; + + offset / tensor.tensor.vector_size() +} + +/// Reads a global quantized tensor at the given position. +/// +/// # Notes +/// +/// The values returned in the [Vector] are not dequantized. +#[cube] +pub fn read_quantized( + inputs: &GlobalArgs, + locals: &LocalArgs, + ref_pos: usize, + #[comptime] arg: FuseArg, + #[comptime] config: &FuseBlockConfig, + #[comptime] scheme: QuantScheme, +) -> Vector { + match arg { + FuseArg::Input(pos, _precision, _layout) => { + let global = inputs.tensors.index(pos); + + let offset = + index_offset_with_quant_layout(global, locals, ref_pos, config.rank, scheme); + let val = global.tensor[offset]; + Vector::cast_from(val) + } + _ => panic!("Not supported"), + } +} + +/// Reads a global scalar. +#[cube] +pub fn read_scalar(inputs: &GlobalArgs, #[comptime] arg: FuseArg) -> C { + match arg { + FuseArg::Scalar(pos, _precision) => { + let scalar = inputs.scalars.index(pos); + scalar.get::() + } + _ => comptime![panic!("Not a scalar")], + } +} + +/// Reads a global scalar that is used as a reshape position. +#[cube] +pub fn read_scalar_shape(inputs: &GlobalArgs, #[comptime] arg: FuseArg) -> usize { + match arg { + FuseArg::ScalarShape(pos) => inputs.reshapes[pos], + _ => comptime![panic!("Not a scalar shape")], + } +} + +/// Reads an input tensor. +#[cube] +pub fn read_input( + inputs: &GlobalArgs, + locals: &LocalArgs, + #[comptime] pos: usize, + ref_pos: usize, + #[comptime] layout: LayoutInfo, + #[comptime] config: &FuseBlockConfig, + #[comptime] transform: Option, +) -> Vector { + set_polyfill_typed::, DynElem, DynSize>(); + let tensor = inputs.tensors.index(pos); + let offset = match layout { + LayoutInfo::SameAsRef => ref_pos, + LayoutInfo::IsRef => ref_pos, + LayoutInfo::Unknown => get_offset(inputs, locals, tensor, ref_pos, None, config, transform), + }; + Vector::cast_from(tensor.tensor[offset]) +} + +/// Returns a slice of data in the asked precision of the input tensor at the given position. +#[cube] +pub fn read_input_window( + inputs: &GlobalArgs, + #[comptime] pos: usize, + start: usize, + end: usize, +) -> &[C] { + set_polyfill_typed::(); + let tensor = inputs.tensors.index(pos); + let slice = &tensor.tensor[start..end]; + slice.downcast() +} + +/// Returns the input as a slice. +#[cube] +pub fn input_as_slice(inputs: &GlobalArgs, #[comptime] pos: usize) -> &[C] { + set_polyfill_typed::(); + let tensor = inputs.tensors.index(pos); + let slice = tensor.tensor.as_slice(); + slice.downcast() +} + +/// Returns the input tensor as a quantized scale view. +#[cube] +pub fn input_as_scales_view( + inputs: &GlobalArgs, + #[comptime] pos: usize, + #[comptime] tensor_pos: usize, + #[comptime] level: QuantLevel, + #[comptime] config: &FuseBlockConfig, +) -> View<'static, C, usize> { + set_polyfill_typed::, DynElem, DynSize>(); + let tensor = inputs.tensors.index(tensor_pos); + let scales = inputs.tensors.index(pos); + let tensor_len = tensor.tensor.len(); + let rank = config.rank; + let layout = match level { + QuantLevel::Tensor => ScalesLayout::new_PerTensor(PerTensorLayout::new(tensor_len)), + QuantLevel::Block(block_size) => { + let block_size = comptime![block_size.to_dim_vec(rank)]; + let mut tensor_shape = Sequence::new(); + let mut scales_strides = Sequence::new(); + #[unroll] + for i in 0..rank { + tensor_shape.push(FastDivmod::new_Fallback(tensor.tensor.shape(i))); + scales_strides.push(scales.tensor.stride(i)); + } + let vector_size = scales.tensor.vector_size(); + let layout = BlockScaledLayout::new( + tensor_shape, + tensor_len, + scales_strides, + block_size, + vector_size, + ); + ScalesLayout::new_BlockScaled(layout) + } + }; + let slice = scales.tensor.as_slice().downcast(); + View::new::, usize>(unsafe { slice.as_boxed_unchecked() }, layout) +} + +/// Reads the input tensor aligned. +#[cube] +pub fn read_input_aligned( + inputs: &GlobalArgs, + locals: &LocalArgs, + #[comptime] pos: usize, + ref_pos: usize, + #[comptime] layout: LayoutInfo, + #[comptime] config: &FuseBlockConfig, + #[comptime] transform: Option, +) -> Vector { + let mut result = Vector::::empty(); + let tensor = inputs.tensors.index(pos); + + match transform.clone() { + Some(Transform::Reshape(shape)) => { + // Very brute force, not really efficient, but not easy to optimize and not a very + // frequent workflow. + let ref_pos = ref_pos * config.width; + #[unroll] + for i in 0..config.width { + let index = reshaped_index( + inputs, + locals, + ref_pos + i, + config.rank, + comptime![shape.clone()], + ); + let index = reshaped_index_to_original_index(&tensor.tensor, index, config.rank); + result.insert(i, C::cast_from(tensor.tensor[index].extract(0))) + } + } + Some(Transform::SwapDims(dim1, dim2)) => { + let offset = + get_offset_aligned(inputs, locals, tensor, ref_pos, layout, config, transform); + let i = comptime![swap_dims_transform(config.rank - 1, (dim1, dim2))]; + let stride = tensor.tensor.stride(i); + + #[unroll] + for i in 0..config.width { + let index = offset + i * stride; + result.insert(i, C::cast_from(tensor.tensor[index].extract(0))) + } + } + None => { + let offset = + get_offset_aligned(inputs, locals, tensor, ref_pos, layout, config, transform); + let stride = tensor.tensor.stride(config.rank - 1); + #[unroll] + for i in 0..config.width { + let index = offset + i * stride; + result.insert(i, C::cast_from(tensor.tensor[index].extract(0))) + } + } + } + + result +} + +/// Computes the offset of the given [GlobalTensor] at on the reference position with a linear +/// layout. +#[cube] +pub fn get_offset_aligned( + inputs: &GlobalArgs, + locals: &LocalArgs, + tensor: &GlobalTensor, + ref_pos: usize, + #[comptime] layout: LayoutInfo, + #[comptime] config: &FuseBlockConfig, + #[comptime] transform: Option, +) -> usize { + match layout { + LayoutInfo::SameAsRef | LayoutInfo::IsRef => { + (ref_pos * locals.ref_vector_size) / tensor.tensor.vector_size() + } + LayoutInfo::Unknown => get_offset( + inputs, + locals, + tensor, + ref_pos, + None, + config, + comptime!(transform.clone()), + ), + } +} + +/// Reads an output tensor. +#[cube] +pub fn read_output( + inputs: &GlobalArgs, + outputs: &GlobalArgs, + locals: &LocalArgs, + #[comptime] pos: usize, + ref_pos: usize, + #[comptime] layout: LayoutInfo, + #[comptime] config: &FuseBlockConfig, +) -> Vector { + let tensor = outputs.tensors.index(pos); + let offset = match layout { + LayoutInfo::SameAsRef => ref_pos, + LayoutInfo::IsRef => ref_pos, + LayoutInfo::Unknown => get_offset(inputs, locals, tensor, ref_pos, None, config, None), + }; + Vector::cast_from(tensor.tensor[offset]) +} + +#[cube] +/// Write the given value at the [arg](Arg) position. +pub fn write( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + ref_pos: usize, + value: Vector, + #[comptime] arg: FuseArg, + #[comptime] config: &FuseBlockConfig, +) { + set_polyfill_typed::, DynElem, DynSize>(); + + match arg { + FuseArg::Output(pos, _, layout) => { + let tensor = outputs.tensors.index(pos); + let output_vs = tensor.tensor.vector_size(); + if comptime![output_vs != config.width] { + // Output tensor has a different vector_size than the computation width. + // Write element-by-element to avoid SPIR-V type mismatch. + write_output_aligned( + inputs, outputs, &*locals, pos, ref_pos, value, layout, config, + ); + } else { + // Vector sizes match - set polyfill to output type and write directly. + set_polyfill::(comptime![tensor.ty]); + let offset = match layout { + LayoutInfo::SameAsRef => ref_pos, + LayoutInfo::IsRef => ref_pos, + LayoutInfo::Unknown => { + get_offset(inputs, &*locals, tensor, ref_pos, None, config, None) + } + }; + let tensor = outputs.tensors.index_mut(pos); + + let value = Vector::cast_from(value); + + tensor.tensor[offset] = value; + } + } + FuseArg::BlockLocal { .. } => write_scalar::(locals, value, arg), + FuseArg::MultiBlockLocal(key, _) | FuseArg::MultiBlockGlobal(key, _) => { + outputs.variables.write(key, Vector::cast_from(value)) + } + _ => comptime![panic!("Can't write into inputs and scalars")], + } +} + +/// Writes a [Vector] value element-by-element to an output tensor whose vector_size +/// differs from the computation width. Mirrors [read_input_aligned] for the write path. +#[cube] +fn write_output_aligned( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &LocalArgs, + #[comptime] pos: usize, + ref_pos: usize, + value: Vector, + #[comptime] layout: LayoutInfo, + #[comptime] config: &FuseBlockConfig, +) { + let tensor = outputs.tensors.index(pos); + set_polyfill::(comptime![tensor.ty]); + + match layout { + LayoutInfo::SameAsRef | LayoutInfo::IsRef => { + let offset = (ref_pos * config.width) / tensor.tensor.vector_size(); + let output = outputs.tensors.index_mut(pos); + let stride = output.tensor.stride(config.rank - 1); + + #[unroll] + for i in 0..config.width { + let idx = offset + i * stride; + let val = if comptime![value.vector_size() == config.width] { + Vector::cast_from(value.extract(i)) + } else { + Vector::cast_from(value.extract(i % value.vector_size())) + }; + output.tensor[idx] = val; + } + } + LayoutInfo::Unknown => { + // When layout differs from ref, each of the config.width elements may map + // to a different offset in the output (or to the same offset when the output + // has broadcast dimensions). Compute each element's offset individually to + // avoid cross-contamination of neighboring elements. + let base_flat = ref_pos * config.width; + + #[unroll] + for i in 0..config.width { + let tensor_ro = outputs.tensors.index(pos); + let offset = index_offset_with_layout( + inputs, + tensor_ro, + locals, + base_flat + i, + None, + config.rank, + None, + ); + let output = outputs.tensors.index_mut(pos); + + let val = if comptime![value.vector_size() == config.width] { + Vector::cast_from(value.extract(i)) + } else { + Vector::cast_from(value.extract(i % value.vector_size())) + }; + output.tensor[offset] = val; + } + } + } +} + +#[cube] +/// Write the given value at the [arg](Arg) position. +pub fn write_scalar( + locals: &mut LocalArgs, + value: Vector, + #[comptime] arg: FuseArg, +) { + match arg { + FuseArg::BlockLocal { pos, ty } => match comptime![ty] { + FuseType::F64 => locals.l_f64.insert(pos, Vector::cast_from(value)), + FuseType::F32 | FuseType::Flex32 => locals.l_f32.insert(pos, Vector::cast_from(value)), + FuseType::F16 => locals.l_f16.insert(pos, Vector::cast_from(value)), + FuseType::BF16 => locals.l_bf16.insert(pos, Vector::cast_from(value)), + FuseType::U64 => locals.l_u64.insert(pos, Vector::cast_from(value)), + FuseType::U32 => locals.l_u32.insert(pos, Vector::cast_from(value)), + FuseType::U16 => locals.l_u16.insert(pos, Vector::cast_from(value)), + FuseType::U8 => locals.l_u8.insert(pos, Vector::cast_from(value)), + FuseType::I64 => locals.l_i64.insert(pos, Vector::cast_from(value)), + FuseType::I32 => locals.l_i32.insert(pos, Vector::cast_from(value)), + FuseType::I16 => locals.l_i16.insert(pos, Vector::cast_from(value)), + FuseType::I8 => locals.l_i8.insert(pos, Vector::cast_from(value)), + }, + _ => comptime![panic!("Can't write into something else than scalars")], + } +} + +#[cube] +pub(crate) fn global_offset( + inputs: &GlobalArgs, + outputs: &GlobalArgs, + locals: &LocalArgs, + index: usize, + #[comptime] arg: FuseArg, + #[comptime] range: Option<(usize, usize)>, + #[comptime] config: &FuseBlockConfig, +) -> usize { + match arg { + FuseArg::Input(pos, _precision, _layout) => { + let tensor = inputs.tensors.index(pos); + get_offset(inputs, locals, tensor, index, range, config, None) + } + FuseArg::Output(pos, _precision, _layout) => { + let tensor = outputs.tensors.index(pos); + get_offset(inputs, locals, tensor, index, range, config, None) + } + _ => panic!("Only input and output tensors have global offset."), + } +} + +#[cube] +fn get_offset( + inputs: &GlobalArgs, + locals: &LocalArgs, + tensor: &GlobalTensor, + ref_pos: usize, + #[comptime] range: Option<(usize, usize)>, + #[comptime] config: &FuseBlockConfig, + #[comptime] transform: Option, +) -> usize { + let offset_ref = ref_pos * locals.ref_vector_size; + index_offset_with_layout( + inputs, + tensor, + locals, + offset_ref, + range, + config.rank, + transform, + ) +} + +#[cube] +/// Gets the vector size for a global tensor. +pub fn global_vector_size( + global: &GlobalArgs, + #[comptime] pos: usize, +) -> comptime_type!(VectorSize) { + let tensor = global.tensors.index(pos); + tensor.tensor.vector_size() +} + +#[cube] +/// Gets the rank for a global tensor. +pub fn global_rank(global: &GlobalArgs, #[comptime] pos: usize) -> usize { + let tensor = global.tensors.index(pos); + tensor.tensor.rank() +} + +#[cube] +/// Gets the length for a global tensor. +pub fn global_len(global: &GlobalArgs, #[comptime] pos: usize) -> usize { + let tensor = global.tensors.index(pos); + tensor.tensor.len() +} + +#[cube] +/// Gets the buffer length for a global tensor. +pub fn global_buffer_len(global: &GlobalArgs, #[comptime] pos: usize) -> usize { + let tensor = global.tensors.index(pos); + tensor.tensor.buffer_len() +} + +#[cube] +/// Gets the reference tensor length. +pub fn ref_len( + inputs: &GlobalArgs, + outputs: &GlobalArgs, + locals: &LocalArgs, + #[comptime] config: &FuseBlockConfig, +) -> usize { + match config.ref_layout.clone() { + RefLayout::Concrete(arg) => match comptime![arg] { + FuseArg::Input(index, _, _) => global_len(inputs, index), + FuseArg::Output(index, _, _) => global_len(outputs, index), + _ => panic!("Invalid concrete ref layout."), + }, + RefLayout::Virtual(..) => num_elements(locals, config), + } +} + +#[cube] +/// Gets the reference buffer tensor length. +pub fn ref_buffer_len( + inputs: &GlobalArgs, + outputs: &GlobalArgs, + locals: &LocalArgs, + #[comptime] config: &FuseBlockConfig, +) -> usize { + match config.ref_layout.clone() { + RefLayout::Concrete(arg) => match comptime![arg] { + FuseArg::Input(index, _, _) => global_buffer_len(inputs, index), + FuseArg::Output(index, _, _) => global_buffer_len(outputs, index), + _ => panic!("Invalid concrete ref layout."), + }, + RefLayout::Virtual(VirtualLayout::SwapDims(arg, ..)) => match arg { + FuseArg::Input(index, _, _) => global_buffer_len(inputs, index), + FuseArg::Output(index, _, _) => global_buffer_len(outputs, index), + _ => panic!("Invalid concrete ref layout."), + }, + RefLayout::Virtual(VirtualLayout::Reshaped { .. }) => num_elements(locals, config), + RefLayout::Virtual(VirtualLayout::Shape(..)) => num_elements(locals, config), + RefLayout::Virtual(VirtualLayout::Runtime { .. }) => num_elements(locals, config), + } +} + +#[cube] +/// Gets the reference number of elements. +pub fn num_elements(locals: &LocalArgs, #[comptime] config: &FuseBlockConfig) -> usize { + let mut length = 1; + + for i in 0..config.rank { + length *= locals.ref_shape[i]; + } + + length +} + +#[cube] +/// Gets the reference axis shape. +pub fn ref_shape(locals: &LocalArgs, axis: usize) -> usize { + locals.ref_shape[axis] +} + +#[cube] +/// Gets the reference axis stride. +pub fn ref_stride(locals: &LocalArgs, axis: usize) -> usize { + locals.ref_strides[axis] +} + +#[cube] +/// Gets the reference vector size. +pub fn ref_vector_size(locals: &LocalArgs) -> comptime_type!(VectorSize) { + comptime![locals.ref_vector_size] +} + +#[cube] +/// Gets the given tensor axis shape. +pub fn global_shape(global: &GlobalArgs, axis: usize, #[comptime] pos: usize) -> usize { + let tensor = global.tensors.index(pos); + tensor.tensor.shape(axis) +} + +#[cube] +/// Gets the given tensor axis stride. +pub fn global_stride(global: &GlobalArgs, dim: usize, #[comptime] pos: usize) -> usize { + let tensor = global.tensors.index(pos); + tensor.tensor.stride(dim) +} + +#[cube] +fn index_offset_with_layout( + inputs: &GlobalArgs, + tensor: &GlobalTensor, + locals: &LocalArgs, + offset_ref: usize, + #[comptime] range: Option<(usize, usize)>, + #[comptime] rank: usize, + #[comptime] transform: Option, +) -> usize { + match comptime![transform.clone()] { + Some(Transform::Reshape(shape)) => { + comptime![assert!( + range.is_none(), + "Can't get a range on a reshaped tensor." + )]; + + let index = reshaped_index(inputs, locals, offset_ref, rank, shape); + reshaped_index_to_original_index(&tensor.tensor, index, rank) + } + Some(Transform::SwapDims(dim1, dim2)) => { + let (start, end) = comptime! {match range { + Some(range) => range, + None => (0, rank), + }}; + + let mut offset = 0; + + #[unroll] + for i in start..end { + let index = comptime![swap_dims_transform(i, (dim1, dim2))]; + let ogwl = offset_ref / locals.ref_strides[i]; + offset += ogwl % tensor.tensor.shape(index) * tensor.tensor.stride(index); + } + + offset / tensor.tensor.vector_size() + } + None => { + let (start, end) = comptime! {match range { + Some(range) => range, + None => (0, rank), + }}; + + let mut offset = 0; + + #[unroll] + for i in start..end { + let ogwl = offset_ref / locals.ref_strides[i]; + offset += ogwl % tensor.tensor.shape(i) * tensor.tensor.stride(i); + } + + offset / tensor.tensor.vector_size() + } + } +} + +pub(crate) fn swap_dims_transform(i: usize, dims: (usize, usize)) -> usize { + if i == dims.0 { + dims.1 + } else if i == dims.1 { + dims.0 + } else { + i + } +} + +#[cube] +#[allow(clippy::clone_on_copy)] +/// The index the input tensor would be at if it was contiguous. +fn reshaped_index( + inputs: &GlobalArgs, + locals: &LocalArgs, + index: usize, + #[comptime] rank: usize, + #[comptime] shape: Vec, +) -> usize { + let mut offset = 0; + let mut stride_curr = 1; + + #[unroll] + for r in 0..rank { + let i = reverse_index(rank, r).comptime(); + let arg = shape[i].clone(); + let shape_i = read_scalar_shape(inputs, arg); + let ogwl = index / locals.ref_strides[i]; + + offset += ogwl % shape_i * stride_curr; + + stride_curr *= shape_i; + } + + offset +} + +#[allow(unreachable_code)] +#[cube] +#[allow(clippy::clone_on_copy)] +fn reshaped_index_to_original_index( + original: &Tensor>, + index_reshaped: usize, + #[comptime] rank: usize, +) -> usize { + let mut remaining = index_reshaped; + let mut offset = 0; + + #[unroll] + for r in 0..rank { + let i = reverse_index(rank, r); + let shape = original.shape(i); + let stride = original.stride(i); + + let coordinate = remaining % shape; + + remaining /= shape; + offset += coordinate * stride; + } + + offset / original.vector_size() +} + +#[cube] +#[allow(unused_variables)] +pub(crate) fn reverse_index( + #[comptime] rank: usize, + #[comptime] iter: usize, +) -> comptime_type!(usize) { + rank - iter - 1 +} + +/// Generic way to construct any [`CubePrimitive`] from an int. Used for fusion. +#[allow(unused_variables)] +#[cube] +fn from_const_int(#[comptime] value: usize) -> C { + intrinsic!(|scope| { Value::constant(value.into(), C::__expand_as_type(scope)).into() }) +} + +#[cube] +#[allow(clippy::extra_unused_type_parameters)] +pub(crate) fn set_polyfill_typed() { + intrinsic!(|scope| { + let elem_type = C::__expand_as_type(scope); + set_polyfill::expand::(scope, elem_type); + }) +} diff --git a/crates/burn-cubecl-fusion/src/engine/codegen/ir.rs b/crates/burn-cubecl-fusion/src/engine/codegen/ir.rs new file mode 100644 index 0000000..f114f42 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/codegen/ir.rs @@ -0,0 +1,1037 @@ +use super::tensor::GlobalTensor; +use crate::engine::codegen::{DynElem, DynSize, DynVector}; +use burn_std::{ + BoolStore, DType, Shape, Strides, bf16, f16, + quantization::{QuantScheme, QuantStore, QuantValue}, + strides, +}; +use core::fmt::Display; +use cubecl::{ + ir::{ElemType, FloatKind, IntKind, StorageType, UIntKind}, + prelude::*, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Hash, PartialEq, Eq, Serialize, Deserialize, PartialOrd, Ord)] +/// Argument to a [fuse operation](FuseOp). +pub enum FuseArg { + /// A readonly input tensor. + Input(usize, FuseType, LayoutInfo), + /// A readwrite output tensor. + Output(usize, FuseType, LayoutInfo), + /// A temporary local variable within a single [block](FuseBlockConfig). + BlockLocal { + /// The position of the current variable relative to all local variables within a single block. + pos: usize, + /// The type of the current variable. + ty: FuseType, + }, + /// A variable shared between multiple [block](FuseBlockConfig) that must have a compatible + /// scope. + MultiBlockLocal(MultiBlockPos, FuseType), + /// A variable shared between multiple [blocks](FuseBlockConfig) within a global accessible + /// scope. + MultiBlockGlobal(MultiBlockPos, FuseType), + /// A global scalar. + Scalar(usize, FuseType), + /// A global scalar used in a reshape operation. + /// + /// This is not a scalar defined by a user for computation, but a scalar defined as part of + /// a reshape operation. + ScalarShape(usize), + /// Only constant that can be encoded into an u32 can be used as literal. + Literal(usize, FuseType), + /// A readonly input tensor that is reshaped. + InputReshaped { + original: Box, + shape: Vec, + broadcasted: bool, + }, + /// A readonly input tensor with swapped dimensions. + InputSwapDims { + original: Box, + dims: (usize, usize), + broadcasted: bool, + }, +} + +/// Metadata of a variable shared between blocks. +#[derive(Clone, Debug, Hash, PartialEq, Eq, Serialize, Deserialize, PartialOrd, Ord)] +pub struct MultiBlockPos { + /// The block position in all blocks included in a fused trace. + pub block_pos: usize, + /// The [FuseArg::BlockLocal] position in the block where the variable is first initialized. + pub block_local_pos: usize, +} + +#[derive( + CubeType, Clone, Copy, Debug, Hash, PartialEq, Eq, Serialize, Deserialize, PartialOrd, Ord, +)] +/// Layout information. +pub enum LayoutInfo { + /// The layout if the same as the reference. + SameAsRef, + /// The reference layout. + IsRef, + /// The layout if unknown. + Unknown, +} + +impl FuseArg { + pub fn precision(&self) -> FuseType { + *match self { + FuseArg::Input(_, p, _) => p, + FuseArg::BlockLocal { ty, .. } => ty, + FuseArg::MultiBlockLocal(_, p) => p, + FuseArg::MultiBlockGlobal(_, p) => p, + FuseArg::Output(_, p, _) => p, + FuseArg::Scalar(_, p) => p, + FuseArg::Literal(_, p) => p, + FuseArg::ScalarShape(_) => return FuseType::U32, + FuseArg::InputReshaped { original, .. } => return original.precision(), + FuseArg::InputSwapDims { original, .. } => return original.precision(), + } + } +} + +impl CubeType for FuseArg { + type ExpandType = Self; +} + +impl IntoExpand for FuseArg { + type Expand = Self; + + fn into_expand(self, _: &Scope) -> Self::Expand { + self + } +} + +impl ExpandTypeClone for FuseArg { + fn clone_unchecked(&self) -> Self { + self.clone() + } +} + +impl AsRefExpand for FuseArg { + fn __expand_ref_method(&self, _: &Scope) -> &Self { + self + } +} + +impl AsMutExpand for FuseArg { + fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self { + self + } +} + +impl IntoMut for FuseArg { + fn into_mut(self, _context: &Scope) -> Self { + self + } +} + +impl IntoRuntime for FuseArg { + fn __expand_runtime_method(self, _context: &Scope) -> Self::ExpandType { + self + } +} + +impl CubeDebug for FuseArg {} + +#[derive(CubeType, Clone, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)] +/// Operations that can be executed and fused automatically using a fuse-on-read and/or +/// fuse-on-write strategy. +pub enum FuseOp { + Add(BinaryFuseArgs), + Sub(BinaryFuseArgs), + Mul(BinaryFuseArgs), + Div(BinaryFuseArgs), + Powf(BinaryFuseArgs), + Abs(UnaryFuseArgs), + Exp(UnaryFuseArgs), + Log(UnaryFuseArgs), + Log1p(UnaryFuseArgs), + Cos(UnaryFuseArgs), + Sin(UnaryFuseArgs), + Tanh(UnaryFuseArgs), + Erf(UnaryFuseArgs), + Sqrt(UnaryFuseArgs), + Recip(UnaryFuseArgs), + Assign(UnaryFuseArgs), + Equal(BinaryFuseArgs), + Lower(BinaryFuseArgs), + Greater(BinaryFuseArgs), + LowerEqual(BinaryFuseArgs), + Rem(BinaryFuseArgs), + GreaterEqual(BinaryFuseArgs), + BitwiseAnd(BinaryFuseArgs), + BitwiseOr(BinaryFuseArgs), + BitwiseXor(BinaryFuseArgs), + BitwiseLeftShift(BinaryFuseArgs), + BitwiseRightShift(BinaryFuseArgs), + BitwiseNot(UnaryFuseArgs), + Clamp { + input: FuseArg, + min: FuseArg, + max: FuseArg, + out: FuseArg, + }, + ConditionalAssign { + cond: FuseArg, + lhs: FuseArg, + rhs: FuseArg, + out: FuseArg, + }, + Gather { + input: FuseArg, + indices: FuseArg, + output: FuseArg, + dim: usize, + }, + Select { + input: FuseArg, + indices: FuseArg, + output: FuseArg, + dim: usize, + }, + Cat { + inputs: Vec, + output: FuseArg, + dim: usize, + }, + Dequantize { + values: FuseArg, + params: FuseArg, + output: FuseArg, + scheme: QuantSchemeFuse, + }, +} + +impl Display for FuseOp { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + FuseOp::Add(args) => write!(f, "{} = {} + {}", args.out, args.lhs, args.rhs), + FuseOp::Sub(args) => write!(f, "{} = {} - {}", args.out, args.lhs, args.rhs), + FuseOp::Mul(args) => write!(f, "{} = {} * {}", args.out, args.lhs, args.rhs), + FuseOp::Div(args) => write!(f, "{} = {} / {}", args.out, args.lhs, args.rhs), + FuseOp::Powf(args) => write!(f, "{} = powf({}, {})", args.out, args.lhs, args.rhs), + FuseOp::Abs(args) => write!(f, "{} = abs({})", args.out, args.input), + FuseOp::Exp(args) => write!(f, "{} = exp({})", args.out, args.input), + FuseOp::Log(args) => write!(f, "{} = log({})", args.out, args.input), + FuseOp::Log1p(args) => write!(f, "{} = log1p({})", args.out, args.input), + FuseOp::Cos(args) => write!(f, "{} = cos({})", args.out, args.input), + FuseOp::Sin(args) => write!(f, "{} = sin({})", args.out, args.input), + FuseOp::Tanh(args) => write!(f, "{} = tanh({})", args.out, args.input), + FuseOp::Erf(args) => write!(f, "{} = erf({})", args.out, args.input), + FuseOp::Sqrt(args) => write!(f, "{} = sqrt({})", args.out, args.input), + FuseOp::Recip(args) => write!(f, "{} = recip({})", args.out, args.input), + FuseOp::Assign(args) => write!(f, "{} = {}", args.out, args.input), + FuseOp::Equal(args) => write!(f, "{} = {} == {}", args.out, args.lhs, args.rhs), + FuseOp::Lower(args) => write!(f, "{} = {} < {}", args.out, args.lhs, args.rhs), + FuseOp::Greater(args) => write!(f, "{} = {} > {}", args.out, args.lhs, args.rhs), + FuseOp::LowerEqual(args) => write!(f, "{} = {} <= {}", args.out, args.lhs, args.rhs), + FuseOp::Rem(args) => write!(f, "{} = {} % {}", args.out, args.lhs, args.rhs), + FuseOp::GreaterEqual(args) => write!(f, "{} = {} >= {}", args.out, args.lhs, args.rhs), + FuseOp::BitwiseAnd(args) => write!(f, "{} = {} & {}", args.out, args.lhs, args.rhs), + FuseOp::BitwiseOr(args) => write!(f, "{} = {} | {}", args.out, args.lhs, args.rhs), + FuseOp::BitwiseXor(args) => write!(f, "{} = {} ^ {}", args.out, args.lhs, args.rhs), + FuseOp::BitwiseLeftShift(args) => { + write!(f, "{} = {} << {}", args.out, args.lhs, args.rhs) + } + FuseOp::BitwiseRightShift(args) => { + write!(f, "{} = {} >> {}", args.out, args.lhs, args.rhs) + } + FuseOp::BitwiseNot(args) => write!(f, "{} = !{}", args.out, args.input), + FuseOp::Clamp { + input, + min, + max, + out, + } => write!(f, "{} = clamp({}, min={}, max={})", out, input, min, max), + FuseOp::ConditionalAssign { + cond, + lhs, + rhs, + out, + } => write!( + f, + "{} = select(cond={}, lhs={}, rhs={})", + out, cond, lhs, rhs + ), + FuseOp::Gather { + input, + indices, + output, + dim, + } => write!( + f, + "{} = gather(input={}, indices={}, dim={})", + output, input, indices, dim + ), + FuseOp::Select { + input, + indices, + output, + dim, + } => write!( + f, + "{} = select(input={}, indices={}, dim={})", + output, input, indices, dim + ), + FuseOp::Cat { + inputs, + output, + dim, + } => { + write!(f, "{output} = cat(inputs=[")?; + for (i, input) in inputs.iter().enumerate() { + if i > 0 { + write!(f, ", ")?; + } + write!(f, "{input}")?; + } + write!(f, "], dim={dim})") + } + FuseOp::Dequantize { + values, + params, + output, + scheme: _, + } => write!( + f, + "{} = dequantize(values={}, params={})", + output, values, params + ), + } + } +} + +#[derive( + CubeType, CubeLaunch, Clone, Debug, Hash, PartialEq, Eq, Serialize, Deserialize, PartialOrd, Ord, +)] +pub struct QuantSchemeFuse { + #[cube(comptime)] + pub(crate) scheme: QuantScheme, +} + +impl FuseOp { + /// Element type used for the computation. + pub(crate) fn cmp_elem(&self) -> ElemType { + match self { + FuseOp::Add(op) => op.lhs.precision().into_elem(), + FuseOp::Sub(op) => op.lhs.precision().into_elem(), + FuseOp::Mul(op) => op.lhs.precision().into_elem(), + FuseOp::Div(op) => op.lhs.precision().into_elem(), + FuseOp::Powf(op) => op.lhs.precision().into_elem(), + FuseOp::Abs(op) => op.out.precision().into_elem(), + FuseOp::Exp(op) => op.out.precision().into_elem(), + FuseOp::Log(op) => op.out.precision().into_elem(), + FuseOp::Log1p(op) => op.out.precision().into_elem(), + FuseOp::Cos(op) => op.out.precision().into_elem(), + FuseOp::Sin(op) => op.out.precision().into_elem(), + FuseOp::Tanh(op) => op.out.precision().into_elem(), + FuseOp::Erf(op) => op.out.precision().into_elem(), + FuseOp::Recip(op) => op.out.precision().into_elem(), + FuseOp::Sqrt(op) => op.out.precision().into_elem(), + FuseOp::Assign(op) => op.out.precision().into_elem(), + FuseOp::Equal(op) => op.lhs.precision().into_elem(), + FuseOp::Lower(op) => op.lhs.precision().into_elem(), + FuseOp::Greater(op) => op.lhs.precision().into_elem(), + FuseOp::LowerEqual(op) => op.lhs.precision().into_elem(), + FuseOp::GreaterEqual(op) => op.lhs.precision().into_elem(), + FuseOp::BitwiseAnd(op) => op.lhs.precision().into_elem(), + FuseOp::BitwiseOr(op) => op.lhs.precision().into_elem(), + FuseOp::BitwiseXor(op) => op.lhs.precision().into_elem(), + FuseOp::BitwiseLeftShift(op) => op.lhs.precision().into_elem(), + FuseOp::BitwiseRightShift(op) => op.lhs.precision().into_elem(), + FuseOp::BitwiseNot(op) => op.out.precision().into_elem(), + FuseOp::ConditionalAssign { out, .. } => out.precision().into_elem(), + FuseOp::Gather { output, .. } => output.precision().into_elem(), + FuseOp::Select { output, .. } => output.precision().into_elem(), + FuseOp::Cat { output, .. } => output.precision().into_elem(), + FuseOp::Dequantize { output, .. } => output.precision().into_elem(), + FuseOp::Rem(op) => op.out.precision().into_elem(), + FuseOp::Clamp { out, .. } => out.precision().into_elem(), + } + } + + pub(crate) fn cmp_storage_ty(&self) -> StorageType { + self.cmp_elem().into() + } +} + +#[derive(CubeType, CubeLaunch, Default, Clone)] +#[expand(derive(Clone))] +/// Global arguments that are used for fusing [element wise operations](ElemTypewiseOp). +pub struct GlobalArgs { + /// Tensors that are stored in global memory. + pub tensors: Sequence, + /// Scalars that are stored in global memory. + pub scalars: Sequence, + /// To be used to perform reshape inside a fused kernel. + pub reshapes: Sequence, + /// When there are no metadata as a reference layout, we provide runtime shape/strides in this + /// sequence instead. + pub runtime_layouts: Sequence, + /// Variables shared between blocks. + pub variables: MultiBlockVariables, +} + +impl GlobalArgsLaunch { + pub fn required_address_type(&self) -> AddressType { + self.tensors + .values + .iter() + .map(|it| it.address_type) + .max() + .unwrap_or_default() + } +} + +/// Variables shared between blocks. +#[derive(CubeType, Default, Clone)] +#[expand(derive(Clone))] +pub struct MultiBlockVariables { + variables: Registry>>, +} + +#[cube] +impl MultiBlockVariables { + /// Initializes the variable with the given key and vector size. + /// + /// # Notes + /// + /// The type of [`NumericExpand`] must be set before calling this function. + pub fn init(&mut self, #[comptime] key: MultiBlockPos) { + let mut registers = + Registry::>>::find_or_default::( + &mut self.variables, + key.block_pos, + ); + let cell = RuntimeCell::new(Vector::empty()); + registers.insert(key.block_local_pos, cell); + } + + /// Read the variable using the provided key. + /// + /// # Notes + /// + /// The variable must be initialized. + pub fn read(&self, #[comptime] key: MultiBlockPos) -> DynVector { + let registers = self.variables.find(key.block_pos); + let cell = registers.find(key.block_local_pos); + cell.read() + } + + /// Write to the variable using the provided key and value. + /// + /// # Notes + /// + /// The variable must be initialized. + pub fn write(&mut self, #[comptime] key: MultiBlockPos, value: DynVector) { + let registers = self.variables.find(key.block_pos); + // Try find for local(visibility) registers. + let cell = registers.find(key.block_local_pos); + cell.store(value); + } +} + +// Because we only create it DURING compilation, not as a real launch arg. +unsafe impl Send for MultiBlockVariables {} +unsafe impl Sync for MultiBlockVariables {} + +impl LaunchArg for MultiBlockVariables { + type RuntimeArg = (); + type CompilationArg = (); + + fn register(_arg: Self::RuntimeArg, _launcher: &mut KernelLauncher) {} + + fn expand( + _arg: &Self::CompilationArg, + _builder: &mut KernelBuilder, + ) -> ::ExpandType { + MultiBlockVariablesExpand { + variables: Default::default(), + } + } +} + +impl Default for GlobalArgsLaunch { + fn default() -> Self { + Self { + tensors: Default::default(), + scalars: Default::default(), + reshapes: Default::default(), + variables: Default::default(), + runtime_layouts: Default::default(), + _phantom_runtime: std::marker::PhantomData, + } + } +} + +impl core::fmt::Debug for GlobalArgsLaunch { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "({:?})", self.tensors.values) + } +} + +impl GlobalArgsLaunch { + /// Get the shape of the given [argument](Arg). + /// + /// # Panics + /// + /// If the argument doesn't have an handle. + pub fn shape(&self, arg: &FuseArg) -> Shape { + match self.resolve_arg(arg) { + TensorArg::Handle { handle, .. } => handle.shape.clone(), + TensorArg::Alias { shape, .. } => shape.clone(), + } + } + + /// Shape used by the reference tensor. + pub fn shape_ref(&self, ref_layout: &RefLayout, rank: usize) -> Shape { + match ref_layout { + RefLayout::Concrete(arg) => self.shape(arg), + RefLayout::Virtual(layout) => match layout { + VirtualLayout::SwapDims(original, dims) => { + let mut shape = self.shape(original); + shape.swap(dims.0, dims.1); + shape + } + VirtualLayout::Reshaped { reshape_pos, .. } => { + let start = *reshape_pos * rank; + let end = start + rank; + self.reshapes.values[start..end].iter().copied().collect() + } + VirtualLayout::Shape(original, _) => self.shape(original), + VirtualLayout::Runtime { pos } => { + let start = (*pos * 2) * rank; + let end = start + rank; + self.runtime_layouts.values[start..end] + .iter() + .copied() + .collect() + } + }, + } + } + + /// Get the strides of the given [argument](Arg). + /// + /// # Panics + /// + /// If the argument doesn't have an handle. + pub fn strides(&self, arg: &FuseArg) -> Strides { + match self.resolve_arg(arg) { + TensorArg::Handle { handle, .. } => handle.strides.clone(), + TensorArg::Alias { strides, .. } => strides.clone(), + } + } + + pub fn strides_ref(&self, ref_layout: &RefLayout, rank: usize) -> Strides { + match ref_layout { + RefLayout::Concrete(arg) => self.strides(arg), + // When not concrete, we operate on the contiguous layout. + _ => { + let shape = self.shape_ref(ref_layout, rank); + let mut strides = strides![0; shape.len()]; + + let mut current = 1; + shape.iter().enumerate().rev().for_each(|(index, val)| { + strides[index] = current; + current *= val; + }); + + strides + } + } + } + + /// Get the vector size of the given [argument](Arg). + /// + /// # Panics + /// + /// If the argument doesn't have an handle. + pub fn vector_size(&self, arg: &FuseArg) -> VectorSize { + match arg { + FuseArg::Input(pos, _, _) => self.tensors.values[*pos].ty.vector_size(), + FuseArg::Output(pos, _, _) => self.tensors.values[*pos].ty.vector_size(), + other => panic!("Arg not found: {other:?}"), + } + } + + /// Resolve the [argument](Arg) to a [tensor argument](TensorArg). + /// + /// # Panics + /// + /// If the argument isn't a global input or output tensor. + pub fn resolve_arg(&self, arg: &FuseArg) -> &TensorArg { + match arg { + FuseArg::Input(pos, _, _) => &self.tensors.values[*pos].tensor, + FuseArg::Output(pos, _, _) => &self.tensors.values[*pos].tensor, + other => panic!("Arg not found: {other:?}"), + } + } +} + +#[derive(CubeType, Clone)] +#[expand(derive(Clone))] +/// Keep track of all local variables that are used as argument in fused +/// [element wise operations](ElemwiseOp). +pub struct LocalArgs { + pub l_f64: Registry>, + pub l_f32: Registry>, + pub l_f16: Registry>, + pub l_bf16: Registry>, + pub l_i64: Registry>, + pub l_i32: Registry>, + pub l_i16: Registry>, + pub l_i8: Registry>, + pub l_u64: Registry>, + pub l_u32: Registry>, + pub l_u16: Registry>, + pub l_u8: Registry>, + pub ref_shape: Box<[usize]>, + pub ref_strides: Box<[usize]>, + #[cube(comptime)] + pub ref_vector_size: VectorSize, +} + +#[cube] +impl LocalArgs { + /// Creates a new [LocalArgs] container. + pub fn new( + ref_shape: &[usize], + ref_strides: &[usize], + #[comptime] ref_vector_size: VectorSize, + ) -> LocalArgs { + LocalArgs { + l_f64: Registry::>::new(), + l_f32: Registry::>::new(), + l_f16: Registry::>::new(), + l_bf16: Registry::>::new(), + l_i64: Registry::>::new(), + l_i32: Registry::>::new(), + l_i16: Registry::>::new(), + l_i8: Registry::>::new(), + l_u64: Registry::>::new(), + l_u32: Registry::>::new(), + l_u16: Registry::>::new(), + l_u8: Registry::>::new(), + ref_shape: unsafe { ref_shape.as_boxed_unchecked() }, + ref_strides: unsafe { ref_strides.as_boxed_unchecked() }, + ref_vector_size, + } + } +} + +#[derive(CubeType, Clone, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)] +/// Unary [element wise operation](ElemwiseOp) arguments. +pub struct UnaryFuseArgs { + pub input: FuseArg, + pub out: FuseArg, +} + +#[derive(CubeType, Clone, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)] +/// Binary [element wise operation](ElemwiseOp) arguments. +pub struct BinaryFuseArgs { + pub lhs: FuseArg, + pub rhs: FuseArg, + pub out: FuseArg, +} + +#[derive( + CubeType, Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, +)] +/// Precisions supported by [element wise operations](ElemwiseOp). +/// +/// This is a custom type instead of [ElemType] so it can implement [CubeType] +/// and restricts the supported types for fusion. +pub enum FuseType { + F64, + F32, + Flex32, + F16, + BF16, + I64, + I32, + I16, + I8, + U64, + U32, + U16, + U8, +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)] +/// Configuration that encapsulates all comptime information necessary for element wise fusion. +pub struct FuseBlockConfig { + pub rank: usize, + pub ref_layout: RefLayout, + pub ops: Vec, + pub width: VectorSize, +} + +impl AsRefExpand for FuseBlockConfig { + fn __expand_ref_method(&self, _: &Scope) -> &Self { + self + } +} + +impl FuseBlockConfig { + pub fn multi_block_variables(&self, registers: &mut Vec<(MultiBlockPos, StorageType)>) { + for op in self.ops.iter() { + op.multi_block_variables(registers); + } + } +} + +impl FuseArg { + pub fn multi_block_variable(&self, registers: &mut Vec<(MultiBlockPos, StorageType)>) { + match self { + FuseArg::MultiBlockGlobal(arg, fuse_type) + // TODO: we need to init the multi-block local, but at some point we could avoid + // that for performance (easier for the underlying compiler). + | FuseArg::MultiBlockLocal(arg, fuse_type) => { + registers.push((arg.clone(), fuse_type.into_storage_type())) + } + _ => {} + }; + } +} + +impl FuseOp { + pub fn multi_block_variables(&self, registers: &mut Vec<(MultiBlockPos, StorageType)>) { + match self { + FuseOp::Add(binary_fuse_args) + | FuseOp::Sub(binary_fuse_args) + | FuseOp::Mul(binary_fuse_args) + | FuseOp::Div(binary_fuse_args) + | FuseOp::Powf(binary_fuse_args) + | FuseOp::Equal(binary_fuse_args) + | FuseOp::Lower(binary_fuse_args) + | FuseOp::Greater(binary_fuse_args) + | FuseOp::LowerEqual(binary_fuse_args) + | FuseOp::Rem(binary_fuse_args) + | FuseOp::GreaterEqual(binary_fuse_args) + | FuseOp::BitwiseAnd(binary_fuse_args) + | FuseOp::BitwiseOr(binary_fuse_args) + | FuseOp::BitwiseXor(binary_fuse_args) + | FuseOp::BitwiseLeftShift(binary_fuse_args) + | FuseOp::BitwiseRightShift(binary_fuse_args) => { + binary_fuse_args.lhs.multi_block_variable(registers); + binary_fuse_args.rhs.multi_block_variable(registers); + binary_fuse_args.out.multi_block_variable(registers); + } + FuseOp::Abs(unary_fuse_args) + | FuseOp::Exp(unary_fuse_args) + | FuseOp::Log(unary_fuse_args) + | FuseOp::Log1p(unary_fuse_args) + | FuseOp::Cos(unary_fuse_args) + | FuseOp::Sin(unary_fuse_args) + | FuseOp::Tanh(unary_fuse_args) + | FuseOp::Erf(unary_fuse_args) + | FuseOp::Sqrt(unary_fuse_args) + | FuseOp::Recip(unary_fuse_args) + | FuseOp::Assign(unary_fuse_args) + | FuseOp::BitwiseNot(unary_fuse_args) => { + unary_fuse_args.input.multi_block_variable(registers); + unary_fuse_args.out.multi_block_variable(registers); + } + FuseOp::Clamp { + input, + min, + max, + out, + } => { + input.multi_block_variable(registers); + min.multi_block_variable(registers); + max.multi_block_variable(registers); + out.multi_block_variable(registers); + } + FuseOp::ConditionalAssign { + cond, + lhs, + rhs, + out, + } => { + cond.multi_block_variable(registers); + lhs.multi_block_variable(registers); + rhs.multi_block_variable(registers); + out.multi_block_variable(registers); + } + FuseOp::Gather { + input, + indices, + output, + dim: _, + } => { + input.multi_block_variable(registers); + indices.multi_block_variable(registers); + output.multi_block_variable(registers); + } + FuseOp::Select { + input, + indices, + output, + dim: _, + } => { + input.multi_block_variable(registers); + indices.multi_block_variable(registers); + output.multi_block_variable(registers); + } + FuseOp::Cat { + inputs, + output, + dim: _, + } => { + for input in inputs { + input.multi_block_variable(registers); + } + output.multi_block_variable(registers); + } + FuseOp::Dequantize { + values, + params, + output, + scheme: _, + } => { + values.multi_block_variable(registers); + params.multi_block_variable(registers); + output.multi_block_variable(registers); + } + } + } +} + +#[cube] +/// Initializes block variables, both globals and locals. +pub fn multi_block_variables_init( + #[comptime] block: &FuseBlockConfig, + variables: &mut MultiBlockVariables, +) { + let output = comptime! { + let mut output = Vec::<(MultiBlockPos, StorageType)>::new(); + block.multi_block_variables(&mut output); + output + }; + + #[unroll] + for i in 0..comptime!(output.len()) { + let (key, dtype) = comptime!(output.get(i).unwrap().clone()); + set_polyfill::(comptime![Type::new(dtype).with_vector_size(block.width)]); + variables.init(key); + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)] +/// A reference layout determines how a fuse execution will access elements in tensors. +/// +/// It can either follow the same layout as a concrete tensor, or follow a virtual layout. +pub enum RefLayout { + Concrete(FuseArg), + Virtual(VirtualLayout), +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)] +/// A virtual layout is always contiguous and retrieves its shape from either a reshaped tensor or a +/// tensor with swap dimensions. +pub enum VirtualLayout { + /// Virtual tensor with the provided shape id and contiguous strides. + Reshaped { + reshape_pos: usize, + vector_size: VectorSize, + }, + /// Virtual tensor with the same shape as the given input, but with swap dims and contiguous + /// strides. + SwapDims(FuseArg, (usize, usize)), + /// Virtual tensor with the same shape as the given input, but with contiguous strides. + Shape(FuseArg, usize), + /// We don't have access to global metadata, they are passed as runtime values. + Runtime { pos: usize }, +} + +impl FuseArg { + /// Adds layout information. + /// + /// It's going to impact how the input or output is read and written to. + pub fn add_layout_info(&mut self, layout: LayoutInfo) { + match self { + FuseArg::Input(_, _, old) => { + *old = layout; + } + FuseArg::Output(_, _, old) => { + *old = layout; + } + _ => {} + } + } +} + +impl RegistryQuery for FuseArg {} + +impl From for FuseType { + fn from(value: ElemType) -> Self { + match value { + ElemType::Float(kind) => match kind { + FloatKind::F16 => Self::F16, + FloatKind::BF16 => Self::BF16, + FloatKind::F32 => Self::F32, + FloatKind::F64 => Self::F64, + FloatKind::Flex32 => Self::Flex32, + _ => panic!("Unsupported precision for fusion: {value}"), + }, + ElemType::Int(kind) => match kind { + IntKind::I64 => Self::I64, + IntKind::I32 => Self::I32, + IntKind::I16 => Self::I16, + IntKind::I8 => Self::I8, + }, + ElemType::UInt(kind) => match kind { + UIntKind::U64 => Self::U64, + UIntKind::U32 => Self::U32, + UIntKind::U16 => Self::U16, + UIntKind::U8 => Self::U8, + }, + ElemType::Bool => Self::U32, + } + } +} + +impl From for FuseType { + fn from(value: StorageType) -> Self { + value.elem_type().into() + } +} + +impl FuseType { + /// Converts the [fused element type](FuseType) into the [cubecl element type](ElemType). + pub fn into_elem(self) -> ElemType { + match self { + FuseType::F32 => ElemType::Float(FloatKind::F32), + FuseType::Flex32 => ElemType::Float(FloatKind::Flex32), + FuseType::F16 => ElemType::Float(FloatKind::F16), + FuseType::BF16 => ElemType::Float(FloatKind::BF16), + FuseType::I64 => ElemType::Int(IntKind::I64), + FuseType::I32 => ElemType::Int(IntKind::I32), + FuseType::I16 => ElemType::Int(IntKind::I16), + FuseType::I8 => ElemType::Int(IntKind::I8), + FuseType::U64 => ElemType::UInt(UIntKind::U64), + FuseType::U32 => ElemType::UInt(UIntKind::U32), + FuseType::U16 => ElemType::UInt(UIntKind::U16), + FuseType::U8 => ElemType::UInt(UIntKind::U8), + FuseType::F64 => ElemType::Float(FloatKind::F64), + } + } + + /// Convert the [fused element type](FuseType) into the [cubecl storage type](StorageType). + pub fn into_storage_type(self) -> StorageType { + self.into_elem().into() + } + + /// Convert the [fused element type](FuseType) into the [cubecl type](Type) + pub fn into_type(self, vector_size: VectorSize) -> Type { + Type::new(self.into_storage_type()).with_vector_size(vector_size) + } +} + +impl From for FuseType { + fn from(value: DType) -> Self { + match value { + DType::F32 => Self::F32, + DType::Flex32 => Self::Flex32, + DType::F16 => Self::F16, + DType::BF16 => Self::BF16, + DType::I64 => Self::I64, + DType::I32 => Self::I32, + DType::I16 => Self::I16, + DType::I8 => Self::I8, + DType::U64 => Self::U64, + DType::U32 => Self::U32, + DType::U16 => Self::U16, + DType::U8 => Self::U8, + DType::Bool(BoolStore::Native) => Self::U32, + DType::Bool(BoolStore::U8) => Self::U8, + DType::Bool(BoolStore::U32) => Self::U32, + DType::F64 => Self::F64, + DType::QFloat(scheme) => match scheme.store { + QuantStore::Native => match scheme.value { + QuantValue::Q8F | QuantValue::Q8S => Self::I8, + QuantValue::E4M3 | QuantValue::E5M2 => { + unimplemented!("Unsupported precision for fusion") + } + QuantValue::Q4F + | QuantValue::Q4S + | QuantValue::Q2F + | QuantValue::Q2S + | QuantValue::E2M1 => { + panic!("Can't store native sub-byte values") + } + }, + QuantStore::PackedU32(_) => Self::U32, + QuantStore::PackedNative(_) => match scheme.value { + QuantValue::E2M1 => unimplemented!("Unsupported precision for fusion"), + other => panic!("{other:?} doesn't support native packing"), + }, + }, + } + } +} + +impl Display for FuseArg { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + FuseArg::Input(pos, ..) => write!(f, "input({pos})"), + FuseArg::Output(pos, ..) => write!(f, "output({pos})"), + FuseArg::BlockLocal { pos, ty } => write!(f, "local({pos}, {ty:?})"), + FuseArg::MultiBlockLocal(mbp, ..) => write!(f, "{mbp}"), + FuseArg::MultiBlockGlobal(mbp, ..) => write!(f, "global_{mbp}"), + FuseArg::Scalar(pos, ..) => write!(f, "scalar({pos})"), + FuseArg::ScalarShape(pos) => write!(f, "scalar_shape({pos})"), + FuseArg::Literal(val, ..) => write!(f, "literal_{val}"), + FuseArg::InputReshaped { original, .. } => write!(f, "input_reshaped_{original}"), + FuseArg::InputSwapDims { original, .. } => write!(f, "input_swap_dims_{original}"), + } + } +} + +impl Display for MultiBlockPos { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "block_local({}-{})", + self.block_pos, self.block_local_pos + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn fuse_type_roundtrips_through_elem_type() { + let all = [ + FuseType::F64, + FuseType::F32, + FuseType::Flex32, + FuseType::F16, + FuseType::BF16, + FuseType::I64, + FuseType::I32, + FuseType::I16, + FuseType::I8, + FuseType::U64, + FuseType::U32, + FuseType::U16, + FuseType::U8, + ]; + for ft in all { + assert_eq!( + FuseType::from(ft.into_elem()), + ft, + "round-trip failed for {ft:?}", + ); + } + } +} diff --git a/crates/burn-cubecl-fusion/src/engine/codegen/kernel.rs b/crates/burn-cubecl-fusion/src/engine/codegen/kernel.rs new file mode 100644 index 0000000..b609e64 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/codegen/kernel.rs @@ -0,0 +1,1097 @@ +use crate::engine::codegen::{DynElem, DynSize}; + +use super::{io::*, ir::*}; +use burn_std::quantization::{QuantScheme, QuantStore, QuantValue}; +use cubecl::{ + ir::{ElemType, FloatKind, StorageType, UIntKind}, + prelude::*, +}; +use cubek::quantization::{dequantize::dequantize_symmetric_packed_value_at, scheme::QuantMode}; + +#[cube] +/// Fuse element-wise operations at the given write position. +/// +/// # Arguments +/// +/// - `inputs`: Contains all readonly global kernel arguments. +/// - `outputs`: Contains all readwrite global kernel arguments. +/// - `locals`: Contains all local variables defined during kernel expansion. +/// - `write_pos`: The logical position the values are written to. +/// - `write_values`: The explicit values to write at the given position. +/// - `write_args`: The arguments associated to the `writes_values`. +/// - `config`: The current [fuse block configuration](FuseBlockConfig). +/// +/// # Notes +/// +/// The function will start by writing `write_values`. +pub fn fuse_on_write( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + write_pos: usize, + write_values: Registry>, + #[comptime] write_args: Vec, + #[comptime] config: &FuseBlockConfig, +) { + comment!("Fuse on write begin"); + // Write the values given as arguments. + #[unroll] + for i in 0..write_args.len() { + let arg = comptime![write_args.get(i).unwrap().clone()]; + let val = write_values.find(arg.clone()); + + write::(inputs, outputs, locals, write_pos, val, arg, config); + } + + fuse(inputs, outputs, locals, write_pos, config); + comment!("Fuse on write end"); +} + +#[cube] +/// Fuse element-wise operations at the given read position. +/// +/// # Arguments +/// +/// - `inputs`: Contains all readonly global kernel arguments. +/// - `outputs`: Contains all readwrite global kernel arguments. +/// - `locals`: Contains all local variables defined during kernel expansion. +/// - `read_pos`: The logical position the values are read from. +/// - `read_args`: The arguments associated to the `read_pos`. +/// - `config`: The current [fuse block configuration](FuseBlockConfig). +/// +/// # Returns +/// +/// - A sequence of values associated to the given `read_args`. +pub fn fuse_on_read( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + read_pos: usize, + #[comptime] read_args: Sequence, + #[comptime] config: &FuseBlockConfig, +) -> Sequence> { + comment!("Fuse on read begin"); + fuse(inputs, outputs, locals, read_pos, config); + + let mut output = Sequence::new(); + + #[unroll] + for i in 0..read_args.len() { + let arg = comptime![read_args.index(i).clone()]; + let value = read::(inputs, &*outputs, &*locals, read_pos, arg, config); + + output.push(value); + } + + comment!("Fuse on read end"); + output +} + +#[cube] +/// Initializes [LocalArgs] given the input and output [arguments](GlobalArgs) with the [FuseBlockConfig]. +/// +/// # Notes +/// +/// The goal is to resolve and cache the reference shape and strides, as it is used in many +/// different function during kernel expansion. +pub fn init_locals( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + #[comptime] config: &FuseBlockConfig, +) -> LocalArgs { + comment!("Init locals begin"); + let mut ref_shape = Array::new(config.rank); + let mut ref_strides = Array::new(config.rank); + + let locals = match config.ref_layout.clone() { + RefLayout::Concrete(arg) => match comptime![arg] { + FuseArg::Input(index, ..) => { + let layout = inputs.tensors.index(index); + + #[unroll] + for i in 0..config.rank { + ref_shape[i] = layout.tensor.shape(i); + ref_strides[i] = layout.tensor.stride(i); + } + + LocalArgs::new( + ref_shape.as_slice(), + ref_strides.as_slice(), + layout.tensor.vector_size(), + ) + } + FuseArg::Output(index, ..) => { + let layout = outputs.tensors.index(index); + + #[unroll] + for i in 0..config.rank { + ref_shape[i] = layout.tensor.shape(i); + ref_strides[i] = layout.tensor.stride(i); + } + + LocalArgs::new( + ref_shape.as_slice(), + ref_strides.as_slice(), + layout.tensor.vector_size(), + ) + } + _ => comptime![panic!("Invalid concrete ref layout.")], + }, + RefLayout::Virtual(layout) => match layout { + VirtualLayout::SwapDims(original, dims) => { + let layout = match original.clone() { + FuseArg::Input(pos, ..) => inputs.tensors.index(pos), + FuseArg::Output(pos, ..) => outputs.tensors.index(pos), + _ => comptime![panic!("Unsupported")], + }; + + let mut stride_curr = 1; + + #[unroll] + #[allow(clippy::clone_on_copy)] + for i in 0..config.rank { + let reverse = reverse_index(config.rank, i); + let swap = comptime![swap_dims_transform(reverse, dims)]; + let shape = layout.tensor.shape(swap.clone()); + + ref_shape[reverse] = shape; + ref_strides[reverse] = stride_curr; + + stride_curr *= ref_shape[comptime![reverse]]; + } + + LocalArgs::new( + ref_shape.as_slice(), + ref_strides.as_slice(), + layout.tensor.vector_size(), + ) + } + VirtualLayout::Reshaped { + reshape_pos, + vector_size, + } => { + let mut stride_curr = 1; + let start = reshape_pos * config.rank; + + #[unroll] + #[allow(clippy::clone_on_copy)] + for i in 0..config.rank { + let reverse = reverse_index(config.rank, i); + let arg = comptime![FuseArg::ScalarShape(start + reverse)]; + let shape = read_scalar_shape(inputs, arg.clone()); + + ref_shape[comptime![reverse]] = shape; + ref_strides[comptime![reverse]] = stride_curr; + + stride_curr *= ref_shape[comptime![reverse]]; + } + + LocalArgs::new(ref_shape.as_slice(), ref_strides.as_slice(), vector_size) + } + VirtualLayout::Runtime { pos } => { + let start_shape = (pos * 2) * config.rank; + let start_strides = start_shape + config.rank; + + #[unroll] + for i in 0..config.rank { + let shape_index = start_shape + i; + let strides_index = start_strides + i; + + ref_shape[i] = *inputs.runtime_layouts.index(shape_index); + ref_strides[i] = *inputs.runtime_layouts.index(strides_index); + } + + LocalArgs::new(ref_shape.as_slice(), ref_strides.as_slice(), config.width) + } + VirtualLayout::Shape(original, vector_size) => { + let layout = match original.clone() { + FuseArg::Input(pos, ..) => inputs.tensors.index(pos), + FuseArg::Output(pos, ..) => outputs.tensors.index(pos), + _ => comptime![panic!("Unsupported")], + }; + let mut stride_curr = 1; + + #[unroll] + #[allow(clippy::clone_on_copy)] + for i in 0..config.rank { + let reverse = reverse_index(config.rank, i); + let shape = layout.tensor.shape(reverse); + + ref_shape[comptime![reverse]] = shape; + ref_strides[comptime![reverse]] = stride_curr; + + stride_curr *= ref_shape[comptime![reverse]]; + } + + LocalArgs::new(ref_shape.as_slice(), ref_strides.as_slice(), vector_size) + } + }, + }; + comment!("Init locals end"); + locals +} + +#[cube] +/// Expands all [operations](FuseOp) registered in the [block config](FuseBlockConfig]. +fn fuse( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + pos: usize, + #[comptime] config: &FuseBlockConfig, +) { + #[unroll] + for index in 0..config.ops.len() { + let op = config.ops[index].clone(); + let define!(E) = op.cmp_storage_ty(); + let size!(N) = config.width; + + match op { + FuseOp::Add(op) => add::(inputs, outputs, locals, pos, op, config), + FuseOp::Div(op) => div::(inputs, outputs, locals, pos, op, config), + FuseOp::Sub(op) => sub::(inputs, outputs, locals, pos, op, config), + FuseOp::Mul(op) => mul::(inputs, outputs, locals, pos, op, config), + FuseOp::Powf(op) => powf::(inputs, outputs, locals, pos, op, config), + FuseOp::Erf(op) => erf::(inputs, outputs, locals, pos, op, config), + FuseOp::Sqrt(op) => sqrt::(inputs, outputs, locals, pos, op, config), + FuseOp::Abs(op) => abs::(inputs, outputs, locals, pos, op, config), + FuseOp::Log(op) => log::(inputs, outputs, locals, pos, op, config), + FuseOp::Log1p(op) => log1p::(inputs, outputs, locals, pos, op, config), + FuseOp::Recip(op) => recip::(inputs, outputs, locals, pos, op, config), + FuseOp::Assign(op) => assign::(inputs, outputs, locals, pos, op, config), + FuseOp::Exp(op) => exp::(inputs, outputs, locals, pos, op, config), + FuseOp::Cos(op) => cos::(inputs, outputs, locals, pos, op, config), + FuseOp::Sin(op) => sin::(inputs, outputs, locals, pos, op, config), + FuseOp::Tanh(op) => tanh::(inputs, outputs, locals, pos, op, config), + FuseOp::Equal(op) => equal::(inputs, outputs, locals, pos, op, config), + FuseOp::Greater(op) => greater::(inputs, outputs, locals, pos, op, config), + FuseOp::GreaterEqual(op) => { + greater_equal::(inputs, outputs, locals, pos, op, config) + } + FuseOp::Lower(op) => lower::(inputs, outputs, locals, pos, op, config), + FuseOp::LowerEqual(op) => lower_equal::(inputs, outputs, locals, pos, op, config), + FuseOp::BitwiseAnd(op) => bitwise_and::(inputs, outputs, locals, pos, op, config), + FuseOp::BitwiseOr(op) => bitwise_or::(inputs, outputs, locals, pos, op, config), + FuseOp::BitwiseXor(op) => bitwise_xor::(inputs, outputs, locals, pos, op, config), + FuseOp::BitwiseLeftShift(op) => { + bitwise_left_shift::(inputs, outputs, locals, pos, op, config) + } + FuseOp::BitwiseRightShift(op) => { + bitwise_right_shift::(inputs, outputs, locals, pos, op, config) + } + FuseOp::BitwiseNot(op) => bitwise_not::(inputs, outputs, locals, pos, op, config), + FuseOp::ConditionalAssign { + cond, + lhs, + rhs, + out, + } => conditional_assign::( + inputs, outputs, locals, pos, cond, lhs, rhs, out, config, + ), + FuseOp::Gather { + input, + indices, + output, + dim, + } => gather::( + inputs, outputs, locals, pos, dim, input, indices, output, config, + ), + FuseOp::Select { + input, + indices, + output, + dim, + } => select_indices::( + inputs, outputs, locals, pos, dim, input, indices, output, config, + ), + FuseOp::Cat { + inputs: tensors, + output, + dim, + } => concat::(inputs, outputs, locals, pos, dim, tensors, output, config), + FuseOp::Dequantize { + values, + params, + output, + scheme, + } => dequantize::( + inputs, + outputs, + locals, + pos, + values, + params, + output, + scheme.scheme, + config, + ), + FuseOp::Rem(op) => rem::(inputs, outputs, locals, pos, op, config), + FuseOp::Clamp { + input, + min, + max, + out, + } => clamp::(inputs, outputs, locals, pos, input, min, max, out, config), + } + } +} + +macro_rules! binary_op { + ($ident:ident, $op:tt) => { + #[cube] + fn $ident( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + write_pos: usize, + #[comptime] op: BinaryFuseArgs, + #[comptime] config: &FuseBlockConfig, + ) { + let lhs = read::(inputs, &*outputs, &*locals, write_pos, op.lhs, config); + let rhs = read::(inputs, &*outputs, &*locals, write_pos, op.rhs, config); + let result = lhs $op rhs; + + write::(inputs, outputs, locals, write_pos, result, op.out, config); + } + }; +} + +macro_rules! binary_int_op { + ($ident:ident, $op:tt) => { + #[cube] + fn $ident( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + write_pos: usize, + #[comptime] op: BinaryFuseArgs, + #[comptime] config: &FuseBlockConfig, + ) { + let lhs = read::(inputs, &*outputs, &*locals, write_pos, op.lhs, config); + let rhs = read::(inputs, &*outputs, &*locals, write_pos, op.rhs, config); + let result = lhs $op rhs; + + write::(inputs, outputs, locals, write_pos, result, op.out, config); + } + }; +} + +macro_rules! unary_int_op { + ($ident:ident, $op:tt) => { + #[cube] + fn $ident( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + write_pos: usize, + #[comptime] op: UnaryFuseArgs, + #[comptime] config: &FuseBlockConfig, + ) { + let input = read::(inputs, &*outputs, &*locals, write_pos, op.input, config); + let result = $op input; + + write::(inputs, outputs, locals, write_pos, result, op.out, config); + } + }; +} + +macro_rules! binary_func { + ($ident:ident, $func:expr, $c:tt) => { + #[cube] + fn $ident( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + write_pos: usize, + #[comptime] op: BinaryFuseArgs, + #[comptime] config: &FuseBlockConfig, + ) { + let lhs = read::(inputs, &*outputs, &*locals, write_pos, op.lhs, config); + let rhs = read::(inputs, &*outputs, &*locals, write_pos, op.rhs, config); + let result = $func(lhs, rhs); + + write::(inputs, outputs, locals, write_pos, result, op.out, config); + } + }; +} + +macro_rules! comparison_op { + ($ident:ident, $op:tt) => { + #[cube] + fn $ident( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + write_pos: usize, + #[comptime] op: BinaryFuseArgs, + #[comptime] config: &FuseBlockConfig, + ) { + let lhs = read::(inputs, &*outputs, &*locals, write_pos, op.lhs, config); + let rhs = read::(inputs, &*outputs, &*locals, write_pos, op.rhs, config); + let result = Vector::new(lhs $op rhs); + + write::(inputs, outputs, locals, write_pos, result, op.out, config); + } + }; +} + +macro_rules! unary_func { + ($ident:ident, $func:expr, $c:tt) => { + #[cube] + fn $ident( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + write_pos: usize, + #[comptime] op: UnaryFuseArgs, + #[comptime] config: &FuseBlockConfig, + ) { + let input = read::(inputs, &*outputs, &*locals, write_pos, op.input, config); + let result = $func(input); + + write::(inputs, outputs, locals, write_pos, result, op.out, config); + } + }; +} + +#[cube] +fn assign( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + write_pos: usize, + #[comptime] op: UnaryFuseArgs, + #[comptime] config: &FuseBlockConfig, +) { + let input = read::(inputs, &*outputs, &*locals, write_pos, op.input, config); + + write::(inputs, outputs, locals, write_pos, input, op.out, config); +} + +#[cube] +fn gather( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + write_pos: usize, + #[comptime] dim: usize, + #[comptime] input: FuseArg, + #[comptime] indices: FuseArg, + #[comptime] output: FuseArg, + #[comptime] config: &FuseBlockConfig, +) { + let vector_size = locals.ref_vector_size; + + let pos_input = comptime! { + match input { + FuseArg::Input(pos, ..) => pos, + _ => panic!("Input tensor isn't an input"), + } + }; + let pos_indices = comptime! { + match indices { + FuseArg::Input(pos, ..) => pos, + _ => panic!("Indices tensor isn't an input"), + } + }; + + let stride_input_dim = global_stride(inputs, dim, pos_input); + + let mut index = 0; + let mut result = Vector::::empty(); + + if comptime![dim > 0] { + let index_before = global_offset( + inputs, + &*outputs, + &*locals, + write_pos, + input.clone(), + comptime![Some((0, dim))], + config, + ); + index += index_before; + } + + if comptime![dim + 1 < config.rank] { + let index_after = global_offset( + inputs, + &*outputs, + &*locals, + write_pos, + input, + comptime![Some((dim + 1, config.rank))], + config, + ); + index += index_after; + } + + let index_offset = global_offset( + inputs, + &*outputs, + &*locals, + write_pos, + indices, + comptime![Some((0, config.rank))], + config, + ); + + // TODO: new IR to differentiate between Gather and GatherBroadcasted at comptime? + let stride_indices_vector = global_stride(inputs, config.rank - 1, pos_indices); + + if comptime![dim == config.rank - 1] { + // Per-element indexing (along the dimension) + #[unroll] + for i in 0..vector_size { + let offset = read_input::>( + inputs, + &*locals, + pos_indices, + index_offset + i * stride_indices_vector, + LayoutInfo::IsRef, + config, + None, + ); + + let input = read_input::>( + inputs, + &*locals, + pos_input, + index + (offset.extract(0) as usize * stride_input_dim), + LayoutInfo::IsRef, + config, + None, + ); + result.insert(i, input.extract(0)); + } + } else { + let stride_input_vector = global_stride(inputs, config.rank - 1, pos_input); + + #[unroll] + for i in 0..vector_size { + let offset = read_input::>( + inputs, + &*locals, + pos_indices, + index_offset + i * stride_indices_vector, + LayoutInfo::IsRef, + config, + None, + ); + + let current_index = + index + (offset.extract(0) as usize * stride_input_dim) + (i * stride_input_vector); + + let input = read_input::>( + inputs, + &*locals, + pos_input, + current_index, + LayoutInfo::IsRef, + config, + None, + ); + + result.insert(i, input.extract(0)); + } + } + + write::(inputs, outputs, locals, write_pos, result, output, config); +} + +#[cube] +fn select_indices( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + write_pos: usize, + #[comptime] dim: usize, + #[comptime] input: FuseArg, + #[comptime] indices: FuseArg, + #[comptime] output: FuseArg, + #[comptime] config: &FuseBlockConfig, +) { + let (vector_size_ref, stride_dim_ref, shape_dim_ref) = ( + locals.ref_vector_size, + locals.ref_strides[dim], + locals.ref_shape[dim], + ); + + let pos_input = comptime! { + match input { + FuseArg::Input(pos, ..) => pos, + _ => panic!("Input tensor isn't an input"), + } + }; + let pos_indices = match indices { + FuseArg::Input(pos, ..) => pos, + _ => panic!("Indices tensor isn't an input"), + }; + + let stride_input_dim = global_stride(inputs, dim, pos_input); + + let mut index = 0; + let mut result = Vector::empty(); + + if comptime![dim != config.rank - 1] { + // In this scenario the select is actually broadcasted along the axis we're working on. + // + // Therefore the same indices are used to fetch multiple entries in the input tensor. + + if comptime![dim > 0] { + let index_before = global_offset( + inputs, + &*outputs, + &*locals, + write_pos, + input.clone(), + comptime![Some((0, dim))], + config, + ); + index += index_before; + } + + if comptime![dim + 1 < config.rank] { + let index_after = global_offset( + inputs, + &*outputs, + &*locals, + write_pos, + input.clone(), + comptime![Some((dim + 1, config.rank))], + config, + ); + index += index_after; + } + + let stride_input_vector = global_stride(inputs, comptime![config.rank - 1], pos_input); + let write_pos_input = write_pos * vector_size_ref; + let coordinate_dim = write_pos_input / stride_dim_ref % shape_dim_ref; + let offset_dim = read_input::>( + inputs, + &*locals, + pos_indices, + coordinate_dim, + LayoutInfo::IsRef, + config, + None, + ); + + index += offset_dim.extract(0) as usize * stride_input_dim; + + #[unroll] + for i in 0..vector_size_ref { + let input = read_input::>( + inputs, + &*locals, + pos_input, + index + i * stride_input_vector, + LayoutInfo::IsRef, + config, + None, + ); + result.insert(i, input.extract(0)); + } + } else { + // In this scenario the select is actually performed on the last dimension we're working on. + // + // Therefore we need to fetch multiple indices that correspond to different entries in the + // input tensor. + + if comptime![dim > 0] { + let index_before = global_offset( + inputs, + &*outputs, + &*locals, + write_pos, + input.clone(), + comptime![Some((0, dim))], + config, + ); + index += index_before; + } + + if comptime![dim + 1 < config.rank] { + let index_after = global_offset( + inputs, + &*outputs, + &*locals, + write_pos, + input, + comptime![Some((dim + 1, config.rank))], + config, + ); + index += index_after; + } + + let write_pos_indices = write_pos * vector_size_ref; + + #[unroll] + for i in 0..vector_size_ref { + let coordinate_dim = (write_pos_indices + i) / stride_dim_ref % shape_dim_ref; + let offset_dim = read_input::>( + inputs, + &*locals, + pos_indices, + coordinate_dim, + LayoutInfo::IsRef, + config, + None, + ); + + let input = read_input::>( + inputs, + &*locals, + pos_input, + index + (offset_dim.extract(0) as usize * stride_input_dim), + LayoutInfo::IsRef, + config, + None, + ); + result.insert(i, input.extract(0)); + } + } + + write::(inputs, outputs, locals, write_pos, result, output, config); +} + +#[cube] +fn concat( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + write_pos: usize, + #[comptime] dim: usize, + #[comptime] tensors: Vec, + #[comptime] output: FuseArg, + #[comptime] config: &FuseBlockConfig, +) { + let (vector_size_ref, stride_dim_ref, shape_dim_ref) = ( + locals.ref_vector_size, + locals.ref_strides[dim], + locals.ref_shape[dim], + ); + + let mut result = Vector::::empty(); + let write_pos_elem = write_pos * vector_size_ref; + + if comptime![dim != config.rank - 1] { + // The concatenation axis isn't the vectorization axis, therefore the whole vector is + // contained in a single input tensor. + let coordinate_dim = write_pos_elem / stride_dim_ref % shape_dim_ref; + + let mut offset_dim_start = 0; + + #[unroll] + for t in 0..tensors.len() { + let tensor = comptime![tensors.get(t).unwrap().clone()]; + let pos_tensor = comptime![concat_input_pos(&tensor)]; + + let shape_tensor_dim = global_shape(inputs, dim, pos_tensor); + let offset_dim_end = offset_dim_start + shape_tensor_dim; + + if coordinate_dim >= offset_dim_start && coordinate_dim < offset_dim_end { + let index = concat_input_offset( + inputs, + &*outputs, + &*locals, + write_pos, + coordinate_dim - offset_dim_start, + tensor, + pos_tensor, + dim, + config, + ); + + let stride_tensor_vector = + global_stride(inputs, comptime![config.rank - 1], pos_tensor); + + #[unroll] + for i in 0..vector_size_ref { + let input = read_input::>( + inputs, + &*locals, + pos_tensor, + index + i * stride_tensor_vector, + LayoutInfo::IsRef, + config, + None, + ); + result.insert(i, input.extract(0)); + } + } + + offset_dim_start = offset_dim_end; + } + } else { + // The concatenation happens along the vectorization axis, therefore each element of the + // vector may come from a different input tensor. + #[unroll] + for i in 0..vector_size_ref { + let coordinate_dim = (write_pos_elem + i) / stride_dim_ref % shape_dim_ref; + let mut offset_dim_start = 0; + + #[unroll] + for t in 0..tensors.len() { + let tensor = comptime![tensors.get(t).unwrap().clone()]; + let pos_tensor = comptime![concat_input_pos(&tensor)]; + + let shape_tensor_dim = global_shape(inputs, dim, pos_tensor); + let offset_dim_end = offset_dim_start + shape_tensor_dim; + + if coordinate_dim >= offset_dim_start && coordinate_dim < offset_dim_end { + let index = concat_input_offset( + inputs, + &*outputs, + &*locals, + write_pos, + coordinate_dim - offset_dim_start, + tensor, + pos_tensor, + dim, + config, + ); + + let input = read_input::>( + inputs, + &*locals, + pos_tensor, + index, + LayoutInfo::IsRef, + config, + None, + ); + result.insert(i, input.extract(0)); + } + + offset_dim_start = offset_dim_end; + } + } + } + + write::(inputs, outputs, locals, write_pos, result, output, config); +} + +fn concat_input_pos(arg: &FuseArg) -> usize { + match arg { + FuseArg::Input(pos, ..) => *pos, + _ => panic!("Cat input isn't a global input"), + } +} + +#[cube] +/// Element index inside a cat input for the vector at `write_pos`, where +/// `coordinate_dim_tensor` is the coordinate along the concat axis relative to the start of +/// this input's segment. +fn concat_input_offset( + inputs: &GlobalArgs, + outputs: &GlobalArgs, + locals: &LocalArgs, + write_pos: usize, + coordinate_dim_tensor: usize, + #[comptime] tensor: FuseArg, + #[comptime] pos_tensor: usize, + #[comptime] dim: usize, + #[comptime] config: &FuseBlockConfig, +) -> usize { + let stride_tensor_dim = global_stride(inputs, dim, pos_tensor); + let mut index = coordinate_dim_tensor * stride_tensor_dim; + + if comptime![dim > 0] { + let index_before = global_offset( + inputs, + outputs, + locals, + write_pos, + tensor.clone(), + comptime![Some((0, dim))], + config, + ); + index += index_before; + } + + if comptime![dim + 1 < config.rank] { + let index_after = global_offset( + inputs, + outputs, + locals, + write_pos, + tensor, + comptime![Some((dim + 1, config.rank))], + config, + ); + index += index_after; + } + + index +} + +#[cube] +fn conditional_assign( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + write_pos: usize, + #[comptime] cond: FuseArg, + #[comptime] lhs: FuseArg, + #[comptime] rhs: FuseArg, + #[comptime] out: FuseArg, + #[comptime] config: &FuseBlockConfig, +) { + let cond = read::(inputs, &*outputs, &*locals, write_pos, cond, config); + let lhs = read::(inputs, &*outputs, &*locals, write_pos, lhs, config); + let rhs = read::(inputs, &*outputs, &*locals, write_pos, rhs, config); + let result = select_many(cond, lhs, rhs); + + write::(inputs, outputs, locals, write_pos, result, out, config); +} + +#[cube] +fn clamp( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + write_pos: usize, + #[comptime] input: FuseArg, + #[comptime] min: FuseArg, + #[comptime] max: FuseArg, + #[comptime] out: FuseArg, + #[comptime] config: &FuseBlockConfig, +) { + let input = read::(inputs, &*outputs, &*locals, write_pos, input, config); + let min = read::(inputs, &*outputs, &*locals, write_pos, min, config); + let max = read::(inputs, &*outputs, &*locals, write_pos, max, config); + let result = cubecl::prelude::clamp(input, min, max); + + write::(inputs, outputs, locals, write_pos, result, out, config); +} + +#[cube] +#[allow(clippy::explicit_counter_loop)] +fn dequantize( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + write_pos: usize, + #[comptime] input: FuseArg, + #[comptime] scales: FuseArg, + #[comptime] output: FuseArg, + #[comptime] scheme: QuantScheme, + #[comptime] config: &FuseBlockConfig, +) { + comptime!(assert_eq!( + scheme.mode, + QuantMode::Symmetric, + "Only symmetric quantization mode is supported." + )); + + let quant_ty = comptime![match scheme.store { + QuantStore::Native => match scheme.value { + QuantValue::Q8F | QuantValue::Q8S => StorageType::Scalar(ElemType::UInt(UIntKind::U8)), + QuantValue::E4M3 => StorageType::Scalar(ElemType::Float(FloatKind::E4M3)), + QuantValue::E5M2 => StorageType::Scalar(ElemType::Float(FloatKind::E5M2)), + QuantValue::Q4F + | QuantValue::Q4S + | QuantValue::Q2F + | QuantValue::Q2S + | QuantValue::E2M1 => unreachable!("Can't store native sub-byte values"), + }, + QuantStore::PackedU32(_) => ElemType::UInt(UIntKind::U32).into(), + QuantStore::PackedNative(_) => match scheme.value { + QuantValue::E2M1 => StorageType::Packed(ElemType::Float(FloatKind::E4M3), 2), + other => panic!("{other:?} doesn't support native packing"), + }, + }]; + let param_ty = comptime![match scheme.param { + cubecl::quant::scheme::QuantParam::F32 => + StorageType::Scalar(ElemType::Float(FloatKind::F32)), + cubecl::quant::scheme::QuantParam::F16 => + StorageType::Scalar(ElemType::Float(FloatKind::F16)), + cubecl::quant::scheme::QuantParam::BF16 => + StorageType::Scalar(ElemType::Float(FloatKind::BF16)), + cubecl::quant::scheme::QuantParam::UE8M0 => + StorageType::Scalar(ElemType::Float(FloatKind::UE8M0)), + cubecl::quant::scheme::QuantParam::UE4M3 => + StorageType::Scalar(ElemType::Float(FloatKind::E4M3)), + }]; + let q_vector_size = N::value().comptime() / scheme.num_quants(); + + let define!(QStoreType) = quant_ty; + let size!(QStoreSize) = q_vector_size; + let define!(QParamType) = param_ty; + let size!(NumQuant) = scheme.num_quants(); + + let tensor_pos = comptime!(match input { + FuseArg::Input(pos, _, _) => pos, + _ => panic!("Not supported"), + }); + let pos = comptime!(match scales { + FuseArg::Input(pos, ..) => pos, + _ => unreachable!(""), + }); + + let num_quants = scheme.num_quants(); + + set_polyfill_typed::, DynElem, DynSize>(); + let input = read_quantized::( + inputs, &*locals, write_pos, input, config, scheme, + ); + + let scales = + input_as_scales_view::>(inputs, pos, tensor_pos, scheme.level, config); + + let result = dequantize_symmetric_packed_value_at::< + C, + NumQuant, + QParamType, + QStoreType, + QStoreSize, + >(write_pos * num_quants, input, &scales, scheme); + + let mut vector = Vector::empty(); + + #[unroll] + for i in 0..q_vector_size { + let value = result[i]; + + #[unroll] + for j in 0..num_quants { + let index = i * num_quants + j; + vector.insert(index, value.extract(j)); + } + } + + write::(inputs, outputs, locals, write_pos, vector, output, config); +} + +binary_op!(add, +); +binary_op!(mul, *); +binary_op!(div, /); +binary_op!(sub, -); + +binary_int_op!(bitwise_and, &); +binary_int_op!(bitwise_or, |); +binary_int_op!(bitwise_xor, ^); +binary_int_op!(bitwise_left_shift, <<); +binary_int_op!(bitwise_right_shift, >>); +unary_int_op!(bitwise_not, !); + +comparison_op!(equal, ==); +comparison_op!(greater, >); +comparison_op!(greater_equal, >=); +comparison_op!(lower, <); +comparison_op!(lower_equal, <=); + +binary_func!(powf, Vector::::powf, Float); +binary_func!(rem, Vector::::mod_floor, Float); + +unary_func!(exp, Vector::::exp, Float); +unary_func!(log, Vector::::ln, Float); +unary_func!(log1p, Vector::::log1p, Float); +unary_func!(sqrt, Vector::::sqrt, Float); +unary_func!(cos, Vector::::cos, Float); +unary_func!(sin, Vector::::sin, Float); +unary_func!(tanh, Vector::::tanh, Float); +unary_func!(erf, Vector::::erf, Float); +unary_func!(recip, Vector::::recip, Float); +unary_func!(abs, Vector::::abs, Numeric); diff --git a/crates/burn-cubecl-fusion/src/engine/codegen/mod.rs b/crates/burn-cubecl-fusion/src/engine/codegen/mod.rs new file mode 100644 index 0000000..00951a5 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/codegen/mod.rs @@ -0,0 +1,8 @@ +pub(crate) mod io; +pub(crate) mod ir; +pub(crate) mod kernel; +pub(crate) mod tensor; +pub(crate) mod view; + +mod base; +pub(crate) use base::*; diff --git a/crates/burn-cubecl-fusion/src/engine/codegen/tensor.rs b/crates/burn-cubecl-fusion/src/engine/codegen/tensor.rs new file mode 100644 index 0000000..c013fad --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/codegen/tensor.rs @@ -0,0 +1,68 @@ +use crate::engine::codegen::{DynElem, DynSize, DynVector}; + +use cubecl::{ir::Type, prelude::*}; +use std::hash::Hash; + +/// Represents a global tensor with the given [element type](ElemType). +/// +/// # Warning +/// +/// The `tensor` field type [Vector>] must be set using polyfill before +/// use. +#[derive(CubeType, Clone)] +#[expand(derive(Clone))] +pub struct GlobalTensor { + /// The global tensor type. + pub tensor: OwnedTensor, + /// The element type of the tensor. + #[cube(comptime)] + pub ty: Type, + /// Whether the current tensor is logically broadcasted. + #[cube(comptime)] + pub broadcasted: bool, +} + +// Everything below is to implement [LaunchArg]. + +#[derive(Clone, PartialEq, Eq, Hash, Debug)] +pub struct GlobalTensorCompilationArg { + tensor: TensorCompilationArg, + ty: Type, + broadcasted: bool, +} + +#[derive(new, Debug)] +pub struct GlobalTensorArg { + pub tensor: as LaunchArg>::RuntimeArg, + pub ty: Type, + pub broadcasted: bool, + pub address_type: AddressType, +} + +impl LaunchArg for GlobalTensor { + type RuntimeArg = GlobalTensorArg; + type CompilationArg = GlobalTensorCompilationArg; + + fn register( + arg: Self::RuntimeArg, + launcher: &mut KernelLauncher, + ) -> Self::CompilationArg { + launcher.with_scope(|scope| set_polyfill::expand::(scope, arg.ty)); + let tensor = OwnedTensor::::register(arg.tensor, launcher); + GlobalTensorCompilationArg { + tensor, + ty: arg.ty, + broadcasted: arg.broadcasted, + } + } + + fn expand(arg: &Self::CompilationArg, builder: &mut KernelBuilder) -> GlobalTensorExpand { + set_polyfill::expand::(&builder.scope, arg.ty); + let tensor = OwnedTensor::expand(&arg.tensor, builder); + GlobalTensorExpand { + tensor, + ty: arg.ty, + broadcasted: arg.broadcasted, + } + } +} diff --git a/crates/burn-cubecl-fusion/src/engine/codegen/view.rs b/crates/burn-cubecl-fusion/src/engine/codegen/view.rs new file mode 100644 index 0000000..41b85f0 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/codegen/view.rs @@ -0,0 +1,358 @@ +use crate::engine::codegen::{DynElem, DynSize, io::set_polyfill_typed}; + +use super::{ + io::{ + Transform, global_buffer_len, global_vector_size, input_as_slice, read_input, + read_input_window, ref_buffer_len, ref_len, + }, + ir::{FuseArg, FuseBlockConfig, GlobalArgs, LayoutInfo, LocalArgs}, + kernel::fuse_on_write, +}; +use cubecl::{ + CubeType, + io::read_masked, + ir::StorageType, + prelude::{barrier::BarrierExpand, *}, + std::tensor::{ + ViewOperations, ViewOperationsExpand, ViewOperationsMut, ViewOperationsMutExpand, + layout::Coords1d, + }, +}; + +#[allow(dead_code, reason = "only used in expand")] +#[derive(CubeType)] +pub struct GlobalInput { + inputs: GlobalArgs, + locals: LocalArgs, + #[cube(comptime)] + pos: usize, + #[cube(comptime)] + ty: StorageType, + #[cube(comptime)] + layout: LayoutInfo, + #[cube(comptime)] + config: FuseBlockConfig, + #[cube(comptime)] + transform: Option, +} + +#[cube] +impl GlobalInput { + pub fn new( + inputs: &GlobalArgs, + locals: &LocalArgs, + #[comptime] arg: FuseArg, + #[comptime] config: FuseBlockConfig, + #[comptime] transform: Option, + ) -> GlobalInput { + let (pos, ty, layout) = comptime![match arg { + FuseArg::Input(pos, prec, layout) => (pos, prec.into_storage_type(), layout), + _ => unreachable!("Must be concrete input"), + }]; + + GlobalInput { + inputs: inputs.clone(), + locals: locals.clone(), + pos, + ty, + layout, + config, + transform, + } + } +} + +impl ViewOperations for GlobalInput {} +impl ViewOperationsExpand for GlobalInputExpand { + #[allow(clippy::too_many_arguments)] + fn __expand_read_method( + &self, + scope: &Scope, + pos: NativeExpand, + ) -> ::ExpandType { + ViewOperationsExpand::::__expand_read_unchecked_method(self, scope, pos) + } + + #[allow(clippy::too_many_arguments)] + fn __expand_read_checked_method( + &self, + scope: &Scope, + pos: NativeExpand, + ) -> ::ExpandType { + let zero = E::__expand_cast_from(scope, 0.into()); + ViewOperationsExpand::::__expand_read_masked_method(self, scope, pos, zero) + } + + #[allow(clippy::too_many_arguments)] + fn __expand_read_masked_method( + &self, + scope: &Scope, + pos: NativeExpand, + value: ::ExpandType, + ) -> ::ExpandType { + let in_bounds = + ViewOperationsExpand::::__expand_is_in_bounds_method(self, scope, pos); + set_polyfill_typed::expand::(scope); + let slice = input_as_slice::expand(scope, &self.inputs, self.pos); + read_masked::expand::(scope, in_bounds, slice, pos, value) + } + + #[allow(clippy::too_many_arguments)] + fn __expand_read_unchecked_method( + &self, + scope: &Scope, + pos: NativeExpand, + ) -> ::ExpandType { + set_polyfill_typed::expand::(scope); + let value = read_input::expand::( + scope, + &self.inputs, + &self.locals, + self.pos, + pos, + self.layout, + &self.config, + self.transform.clone(), + ); + E::__expand_cast_from(scope, value) + } + + #[allow(clippy::too_many_arguments)] + fn __expand_as_linear_slice_method( + &self, + scope: &Scope, + pos: NativeExpand, + end: NativeExpand, + ) -> &SliceExpand { + set_polyfill_typed::expand::(scope); + let end = end.__expand_add_method(scope, 1usize.into_expand(scope)); + read_input_window::expand(scope, &self.inputs, self.pos, pos, end) + } + + #[allow(clippy::too_many_arguments)] + fn __expand_tensor_map_load_method( + &self, + _scope: &Scope, + _barrier: &BarrierExpand, + _shared_memory: &mut SliceExpand, + _pos: NativeExpand, + ) { + panic!("Not a tensor map") + } + + #[allow(clippy::too_many_arguments)] + fn __expand_shape_method(&self, scope: &Scope) -> NativeExpand { + global_buffer_len::expand(scope, &self.inputs, self.pos) + } + + #[allow(clippy::too_many_arguments)] + fn __expand_is_in_bounds_method( + &self, + scope: &Scope, + pos: NativeExpand, + ) -> NativeExpand { + let buffer_len = global_buffer_len::expand(scope, &self.inputs, self.pos); + pos.__expand_lt_method(scope, &buffer_len) + } +} + +impl Vectorized for GlobalInput {} +impl VectorizedExpand for GlobalInputExpand { + fn vector_size(&self) -> VectorSize { + let temp_scope = Scope::root(false); + global_vector_size::expand(&temp_scope, &self.inputs, self.pos) + } +} + +#[allow(dead_code, reason = "only used in expand")] +#[derive(CubeType)] +pub struct FusedOutput { + inputs: GlobalArgs, + outputs: GlobalArgs, + locals: LocalArgs, + arg: FuseArg, + #[cube(comptime)] + config: FuseBlockConfig, +} + +#[cube] +impl FusedOutput { + pub fn new( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + arg: FuseArg, + #[comptime] config: FuseBlockConfig, + ) -> Self { + FusedOutput { + inputs: inputs.clone(), + outputs: outputs.clone(), + locals: locals.clone(), + arg, + config, + } + } +} + +impl ViewOperations for FusedOutput {} +impl ViewOperationsExpand for FusedOutputExpand { + #[allow(clippy::too_many_arguments)] + fn __expand_read_method( + &self, + _scope: &Scope, + _pos: NativeExpand, + ) -> ::ExpandType { + todo!() + } + + #[allow(clippy::too_many_arguments)] + fn __expand_read_checked_method( + &self, + _scope: &Scope, + _pos: NativeExpand, + ) -> ::ExpandType { + todo!() + } + + #[allow(clippy::too_many_arguments)] + fn __expand_read_masked_method( + &self, + _scope: &Scope, + _pos: NativeExpand, + _value: ::ExpandType, + ) -> ::ExpandType { + todo!() + } + + #[allow(clippy::too_many_arguments)] + fn __expand_read_unchecked_method( + &self, + _scope: &Scope, + _pos: NativeExpand, + ) -> ::ExpandType { + todo!() + } + + #[allow(clippy::too_many_arguments)] + fn __expand_as_linear_slice_method( + &self, + _scope: &Scope, + _pos: NativeExpand, + _size: NativeExpand, + ) -> &SliceExpand { + todo!() + } + + #[allow(clippy::too_many_arguments)] + fn __expand_tensor_map_load_method( + &self, + _scope: &Scope, + _barrier: &BarrierExpand, + _shared_memory: &mut SliceExpand, + _pos: NativeExpand, + ) { + panic!("Not a tensor map") + } + + #[allow(clippy::too_many_arguments)] + fn __expand_shape_method(&self, scope: &Scope) -> NativeExpand { + ref_len::expand( + scope, + &self.inputs, + &self.outputs, + &self.locals, + &self.config, + ) + } + + #[allow(clippy::too_many_arguments)] + fn __expand_is_in_bounds_method( + &self, + scope: &Scope, + pos: NativeExpand, + ) -> NativeExpand { + let buffer_len = ref_buffer_len::expand( + scope, + &self.inputs, + &self.outputs, + &self.locals, + &self.config, + ); + pos.__expand_lt_method(scope, &buffer_len) + } +} + +impl ViewOperationsMut for FusedOutput {} +impl ViewOperationsMutExpand for FusedOutputExpand { + #[allow(clippy::too_many_arguments)] + fn __expand_write_method( + &self, + scope: &Scope, + pos: NativeExpand, + value: ::ExpandType, + ) { + let values = Registry::>::__expand_new(scope); + let mut args = comptime![Vec::::new()]; + + let value = Vector::__expand_cast_from(scope, value); + values + .clone() + .__expand_insert_method(scope, comptime![self.arg.clone()], value); + comptime![args.push(self.arg.clone())]; + + let mut outputs = self.outputs.clone(); + let mut locals = self.locals.clone(); + + fuse_on_write::expand( + scope, + &self.inputs, + &mut outputs, + &mut locals, + pos, + values, + args, + &self.config, + ); + } + + #[allow(clippy::too_many_arguments)] + fn __expand_write_checked_method( + &self, + scope: &Scope, + pos: NativeExpand, + value: ::ExpandType, + ) { + let in_bounds = + ViewOperationsExpand::::__expand_is_in_bounds_method(self, scope, pos); + if_expand(scope, in_bounds, |scope| { + ViewOperationsMutExpand::::__expand_write_method(self, scope, pos, value); + }) + } + + #[allow(clippy::too_many_arguments)] + fn __expand_as_linear_slice_mut_method( + &self, + _scope: &Scope, + _pos: NativeExpand, + _size: NativeExpand, + ) -> &mut SliceExpand { + todo!("Not yet supported") + } + + #[allow(clippy::too_many_arguments)] + fn __expand_tensor_map_store_method( + &self, + _scope: &Scope, + _shared_memory: &SliceExpand, + _pos: NativeExpand, + ) { + panic!("Not a tensor map") + } +} + +impl Vectorized for FusedOutput {} +impl VectorizedExpand for FusedOutputExpand { + fn vector_size(&self) -> VectorSize { + self.locals.ref_vector_size + } +} diff --git a/crates/burn-cubecl-fusion/src/engine/fuser.rs b/crates/burn-cubecl-fusion/src/engine/fuser.rs new file mode 100644 index 0000000..9b3dee4 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/fuser.rs @@ -0,0 +1,884 @@ +use super::{ + codegen::ir::{BinaryFuseArgs, FuseArg, FuseOp, UnaryFuseArgs}, + settings::FuseSettings, + trace::{FuseTrace, TraceFuser, block::QuantInput}, +}; +use crate::engine::{codegen::ir::QuantSchemeFuse, scoring::Scoring}; +use burn_backend::cubecl::dtype_to_elem_type; +use burn_fusion::{FuserProperties, FuserStatus, OperationFuser}; +use burn_ir::{ + BaseOperationIr, BinaryOpIr, FloatOperationIr, IntOperationIr, NumericOperationIr, OperationIr, + ScalarOpIr, TensorIr, UnaryOpIr, +}; +use burn_std::{ + DType, Shape, + config::{fusion::FusionLogLevel, log_fusion}, +}; +use cubecl::ir::ElemType; + +/// The base operation fuser that can be used to fuse [all supported fuse operations](FuseOp). +/// +/// +/// This fuser doesn't create a ready-to-execute kernel, but rather generates a +/// [trace](FuseTrace) that be used with a [runner](super::trace::TraceRunner). +/// +/// Since this fuser supports fusing multiple blocks, you can fuse any compute-bound operations +/// with the combination of fuse-on-read and fuse-on-write strategy. +/// +/// # Notes +/// +/// It is responsible to translate [OperationIr] into [FuseOp] and it uses the [TraceFuser] +/// to actually fuse the [FuseOp] when possible. +#[derive(Debug, Clone)] +pub(crate) struct TraceOperationFuser { + pub(crate) fuser: TryTraceFuser, + scoring: Scoring, + pub(crate) settings: FuseSettings, + pub(crate) current_output_shape: Shape, + status: FuserStatus, + pub(crate) num_ops: usize, + pub(crate) num_views: usize, + pub(crate) max_bindings: u32, +} + +impl TraceOperationFuser { + /// Emits a `Full`-level fusion log explaining why the fuser transitioned to `Closed`. + /// `prev_num_ops` is the number of ops already fused before the rejected op. + fn log_closed(&self, op: &OperationIr, prev_num_ops: usize, reason: &'static str) { + let max = self.max_bindings; + log_fusion(FusionLogLevel::Full, || { + // Debug-format the op and estimate bindings lazily: both are expensive + // enough to skip when logging is off. + let op_dbg = format!("{op:?}"); + let op_short = op_dbg + .split_once('(') + .map(|(head, _)| head) + .unwrap_or(&op_dbg); + let est = self.fuser.fuser.estimate_bindings(); + format!( + "[fuser] closed on {op_short} ({reason}); had {prev_num_ops} ops, est_bindings={est}/{max}" + ) + }); + } + + /// Checks if the [operation](OperationIr) can be fused with the current fuser. + pub(crate) fn can_fuse(&self, op: &OperationIr) -> bool { + let len_previous = self.len(); + let mut fuser_cloned = self.clone(); + + fuser_cloned.fuse(op); + let len_after = fuser_cloned.len(); + + len_after > len_previous + } +} + +impl OperationFuser for TraceOperationFuser { + fn fuse(&mut self, op: &OperationIr) { + if let FuserStatus::Closed = self.status { + return; + } + + // Capture state before the fuse attempt so we can log a useful reason + // if the fuser closes on this op. + let prev_num_ops = self.num_ops; + + match op { + OperationIr::Drop(tensor) => { + if self.num_ops == 0 { + self.status = FuserStatus::Closed; + self.log_closed(op, prev_num_ops, "drop on empty fuser"); + return; + } + + self.fuser.fuser.fuse_dropped(tensor); + } + OperationIr::BaseFloat(ops) => { + if !self.fuse_base(ops) { + self.status = FuserStatus::Closed; + self.log_closed(op, prev_num_ops, "base fuse rejected"); + return; + } + } + OperationIr::BaseInt(ops) => { + if !self.fuse_base(ops) { + self.status = FuserStatus::Closed; + self.log_closed(op, prev_num_ops, "base fuse rejected"); + return; + } + } + OperationIr::Float(_dtype, ops) => { + if !self.fuse_float(ops) { + self.status = FuserStatus::Closed; + self.log_closed(op, prev_num_ops, "float fuse rejected"); + return; + } + } + OperationIr::Int(ops) => { + if !self.fuse_int(ops) { + self.status = FuserStatus::Closed; + self.log_closed(op, prev_num_ops, "int fuse rejected"); + return; + } + } + OperationIr::NumericFloat(_dtype, ops) => { + if !self.fuse_numeric(ops) { + self.status = FuserStatus::Closed; + self.log_closed(op, prev_num_ops, "numeric fuse rejected"); + return; + } + } + OperationIr::NumericInt(_dtype, ops) => { + if !self.fuse_numeric(ops) { + self.status = FuserStatus::Closed; + self.log_closed(op, prev_num_ops, "numeric fuse rejected"); + return; + } + } + OperationIr::BaseBool(ops) => { + if !self.fuse_base(ops) { + self.status = FuserStatus::Closed; + self.log_closed(op, prev_num_ops, "base fuse rejected"); + return; + } + } + _ => { + self.status = FuserStatus::Closed; + self.log_closed(op, prev_num_ops, "unsupported op variant"); + return; + } + }; + + self.status = FuserStatus::Open; + self.scoring.register(op); + self.num_ops += 1; + } + + fn finish(&mut self) -> FuseTrace { + self.fuser.finish(self.current_output_shape.clone()) + } + + fn len(&self) -> usize { + self.num_ops + } + + fn reset(&mut self) { + self.num_ops = 0; + self.scoring.reset(); + self.num_views = 0; + self.status = FuserStatus::Open; + self.fuser = TryTraceFuser::new(self.max_bindings, self.settings); + self.current_output_shape = Shape::new([]); + } + + fn status(&self) -> FuserStatus { + self.status + } + + fn properties(&self) -> FuserProperties { + let ready = self.num_ops > 0; + let score = self + .scoring + .evaluate(&self.fuser.clone().finish(self.current_output_shape.clone())); + + FuserProperties { ready, score } + } + + fn clone_dyn(&self) -> Box> { + Box::new(self.clone()) + } +} + +impl TraceOperationFuser { + /// Creates a new fuser. + pub fn new(max_bindings: u32, settings: FuseSettings) -> Self { + Self { + fuser: TryTraceFuser::new(max_bindings, settings), + settings, + scoring: Scoring::default(), + num_ops: 0, + num_views: 0, + max_bindings, + current_output_shape: Shape::new([]), + status: FuserStatus::Open, + } + } + + /// Closes the fuser. + pub fn close(&mut self) { + self.status = FuserStatus::Closed; + } + + /// Declares an input tensor argument where the kernel is responsible to load. + /// + /// # Returns + /// + /// - The argument that maps to the tensor to be used during kernel expansion. + pub fn input_unhandled(&mut self, tensor: &TensorIr) -> FuseArg { + self.fuser.fuser.input_unhandled(tensor) + } + + /// Declares an input quantized tensor argument where the kernel is responsible to load. + /// + /// # Returns + /// + /// None if it's not possible to fuse a quantized tensor. Otherwise: + /// + /// - The argument that maps to the tensor values to be used during kernel expansion. + /// - The argument that maps to the tensor params to be used during kernel expansion. + pub fn input_quantized_unhandled(&mut self, tensor: &TensorIr) -> Option<(FuseArg, FuseArg)> { + self.fuser.fuser.input_quantized_unhandled(tensor) + } + + /// Declares an output tensor argument where the kernel is responsible to write values. + /// + /// # Notes + /// + /// Normally you don't have to declare outputs explicitly before they are going to be + /// fused based on the operations [fused](Self::fuse). + /// + /// # Returns + /// + /// - The argument that maps to the tensor to be used during kernel expansion. + pub fn output_unhandled(&mut self, tensor: &TensorIr) -> FuseArg { + if self.current_output_shape.is_empty() { + self.current_output_shape = tensor.shape.clone(); + } else if self.current_output_shape.iter().sum::() < tensor.shape.iter().sum() { + // The larguest shape win. + self.current_output_shape = tensor.shape.clone(); + } + + self.fuser.fuser.output_unhandled(tensor) + } + + /// Closes the previous block and declares a new one. + /// + /// # Arguments + /// + /// - arguments: Tensors that are logical outputs of the current block and inputs of the following blocks. + /// - settings: [FuseSettings] to be used by the next block. + /// + /// # Returns + /// + /// None if it's impossible to create a next block with the given arguments. Otherwise, the + /// corresponding [arguments](Arg) to the given tensors are returned. + pub fn next_block( + &mut self, + arguments: [&TensorIr; N], + settings: FuseSettings, + global: bool, + ) -> [FuseArg; N] { + let block_pos = self.fuser.fuser.num_previous_blocks(); + let current_output_shape = + core::mem::replace(&mut self.current_output_shape, Shape::new([])); + + self.fuser.fuser.next_block(current_output_shape, settings); + + self.settings = settings; + self.status = FuserStatus::Open; + + arguments.map(|arg| self.fuser.fuser.block_local_input(arg, block_pos, global)) + } + + /// Tag the [tensor](TensorIr) as received from a previous block. + /// + /// This will avoid reading the input again and instead use le local version when possible. + pub fn block_local_input(&mut self, tensor: &TensorIr, block_pos: usize, global: bool) { + self.fuser + .fuser + .block_local_input(tensor, block_pos, global); + } + + fn fuse_base(&mut self, ops: &BaseOperationIr) -> bool { + match ops { + BaseOperationIr::Equal(desc) => self.fuse_binary_ops(desc, |lhs, rhs, out| { + FuseOp::Equal(BinaryFuseArgs { lhs, rhs, out }) + }), + BaseOperationIr::EqualElem(desc) => self.fuse_scalar_ops(desc, |lhs, rhs, out| { + FuseOp::Equal(BinaryFuseArgs { lhs, rhs, out }) + }), + BaseOperationIr::Cast(desc) => { + self.fuse_unary_op(&desc.input, &desc.out, |input, out| { + FuseOp::Assign(UnaryFuseArgs { input, out }) + }) + } + BaseOperationIr::SwapDims(desc) => { + if !self.output_is_compatible(&desc.out) { + return false; + } + + if self.fuser.fuse(|fuser| { + fuser.input_swap_dims(&desc.input, &desc.out, (desc.dim1, desc.dim2))?; + + Some(()) + }) { + self.num_views += 1; + true + } else { + false + } + } + BaseOperationIr::Reshape(desc) => { + if desc.input.shape == desc.out.shape { + return self.fuse_unary_op(&desc.input, &desc.out, |input, out| { + FuseOp::Assign(UnaryFuseArgs { input, out }) + }); + } + + if desc.input.shape.rank() > desc.out.shape.rank() { + // Not yet supported. + return false; + } + + if !self.output_is_compatible(&desc.out) { + return false; + } + + if self.fuser.fuse(|fuser| { + fuser.input_reshaped(&desc.input, &desc.out)?; + Some(()) + }) { + self.num_views += 1; + true + } else { + false + } + } + BaseOperationIr::Ones(desc) => { + if !self.output_is_compatible(&desc.out) { + return false; + } + + let elem: ElemType = dtype_to_elem_type(desc.out.dtype); + let precision = elem.into(); + let input = FuseArg::Literal(1, precision); + + self.fuser.fuse(|fuser| { + let out = fuser.output(&desc.out)?; + + fuser.fuse_operation(FuseOp::Assign(UnaryFuseArgs { input, out })); + + Some(()) + }) + } + BaseOperationIr::Zeros(desc) => { + if !self.output_is_compatible(&desc.out) { + return false; + } + + let elem: ElemType = dtype_to_elem_type(desc.out.dtype); + let precision = elem.into(); + let input = FuseArg::Literal(0, precision); + + self.fuser.fuse(|fuser| { + let out = fuser.output(&desc.out)?; + + fuser.fuse_operation(FuseOp::Assign(UnaryFuseArgs { input, out })); + + Some(()) + }) + } + BaseOperationIr::Gather(desc) => { + if !self.output_is_compatible(&desc.out) { + return false; + } + + self.fuser.fuse(|build| { + let input = build.input_indexed(&desc.tensor)?; + let indices = build.input_indexed(&desc.indices)?; + let output = build.output(&desc.out)?; + + build.fuse_operation(FuseOp::Gather { + input, + indices, + output, + dim: desc.dim, + }); + + Some(()) + }) + } + BaseOperationIr::Select(desc) => { + if !self.output_is_compatible(&desc.out) { + return false; + } + + self.fuser.fuse(|build| { + let input = build.input_indexed(&desc.tensor)?; + let indices = build.input_indexed(&desc.indices)?; + let output = build.output(&desc.out)?; + + build.fuse_operation(FuseOp::Select { + input, + indices, + output, + dim: desc.dim, + }); + + Some(()) + }) + } + BaseOperationIr::Cat(desc) => { + if !self.output_is_compatible(&desc.out) { + return false; + } + + self.fuser.fuse(|build| { + let mut tensors = Vec::with_capacity(desc.tensors.len()); + + for tensor in desc.tensors.iter() { + tensors.push(build.input_indexed(tensor)?); + } + + let output = build.output(&desc.out)?; + + build.fuse_operation(FuseOp::Cat { + inputs: tensors, + output, + dim: desc.dim, + }); + + Some(()) + }) + } + BaseOperationIr::MaskWhere(desc) => { + if !self.output_is_compatible(&desc.out) { + return false; + } + + self.fuser.fuse(|build| { + let cond = build.input(&desc.mask)?; + let rhs = build.input(&desc.tensor)?; + let lhs = build.input(&desc.value)?; + let out = build.output(&desc.out)?; + + build.fuse_operation(FuseOp::ConditionalAssign { + cond, + lhs, + rhs, + out, + }); + + Some(()) + }) + } + BaseOperationIr::MaskFill(desc) => { + if !self.output_is_compatible(&desc.out) { + return false; + } + + self.fuser.fuse(|build| { + let cond = build.input(&desc.mask)?; + let lhs = build.scalar(&desc.value, desc.out.dtype); + let rhs = build.input(&desc.tensor)?; + let out = build.output(&desc.out)?; + + build.fuse_operation(FuseOp::ConditionalAssign { + cond, + lhs, + rhs, + out, + }); + + Some(()) + }) + } + _ => false, + } + } + + fn fuse_float(&mut self, ops: &FloatOperationIr) -> bool { + match ops { + FloatOperationIr::Exp(desc) => { + self.fuse_unary_ops(desc, |input, out| FuseOp::Exp(UnaryFuseArgs { input, out })) + } + FloatOperationIr::Log(desc) => { + self.fuse_unary_ops(desc, |input, out| FuseOp::Log(UnaryFuseArgs { input, out })) + } + FloatOperationIr::Powf(desc) => self.fuse_binary_ops(desc, |lhs, rhs, out| { + FuseOp::Powf(BinaryFuseArgs { lhs, rhs, out }) + }), + FloatOperationIr::Log1p(desc) => self.fuse_unary_ops(desc, |input, out| { + FuseOp::Log1p(UnaryFuseArgs { input, out }) + }), + FloatOperationIr::Cos(desc) => { + self.fuse_unary_ops(desc, |input, out| FuseOp::Cos(UnaryFuseArgs { input, out })) + } + FloatOperationIr::Sin(desc) => { + self.fuse_unary_ops(desc, |input, out| FuseOp::Sin(UnaryFuseArgs { input, out })) + } + FloatOperationIr::PowfScalar(desc) => self.fuse_scalar_ops(desc, |lhs, rhs, out| { + FuseOp::Powf(BinaryFuseArgs { lhs, rhs, out }) + }), + FloatOperationIr::Tanh(desc) => self.fuse_unary_ops(desc, |input, out| { + FuseOp::Tanh(UnaryFuseArgs { input, out }) + }), + FloatOperationIr::Erf(desc) => { + self.fuse_unary_ops(desc, |input, out| FuseOp::Erf(UnaryFuseArgs { input, out })) + } + FloatOperationIr::Sqrt(desc) => self.fuse_unary_ops(desc, |input, out| { + FuseOp::Sqrt(UnaryFuseArgs { input, out }) + }), + FloatOperationIr::Recip(desc) => self.fuse_unary_ops(desc, |input, out| { + FuseOp::Recip(UnaryFuseArgs { input, out }) + }), + FloatOperationIr::Dequantize(desc) => { + if !self.output_is_compatible(&desc.out) { + return false; + } + + self.fuser.fuse(|build| { + let qinput = build.input_quantized(&desc.input)?; + let out = build.output(&desc.out)?; + + match qinput { + QuantInput::AlreadyDequantized { local } => { + build.fuse_operation(FuseOp::Assign(UnaryFuseArgs { + input: local, + out, + })); + } + QuantInput::Quantized { values, params } => { + build.fuse_operation(FuseOp::Dequantize { + values, + params, + output: out, + scheme: match desc.input.dtype { + DType::QFloat(scheme) => QuantSchemeFuse { scheme }, + _ => unreachable!("Should be a quant tensor."), + }, + }); + } + } + + Some(()) + }) + } + _ => false, + } + } + + fn fuse_int(&mut self, ops: &IntOperationIr) -> bool { + match ops { + IntOperationIr::IntoFloat(desc) => { + self.fuse_unary_op(&desc.input, &desc.out, |input, out| { + FuseOp::Assign(UnaryFuseArgs { input, out }) + }) + } + IntOperationIr::BitwiseAnd(desc) => self.fuse_binary_ops(desc, |lhs, rhs, out| { + FuseOp::BitwiseAnd(BinaryFuseArgs { lhs, rhs, out }) + }), + IntOperationIr::BitwiseAndScalar(desc) => self + .fuse_scalar_ops(desc, |lhs, rhs, out| { + FuseOp::BitwiseAnd(BinaryFuseArgs { lhs, rhs, out }) + }), + IntOperationIr::BitwiseOr(desc) => self.fuse_binary_ops(desc, |lhs, rhs, out| { + FuseOp::BitwiseOr(BinaryFuseArgs { lhs, rhs, out }) + }), + IntOperationIr::BitwiseOrScalar(desc) => self.fuse_scalar_ops(desc, |lhs, rhs, out| { + FuseOp::BitwiseOr(BinaryFuseArgs { lhs, rhs, out }) + }), + IntOperationIr::BitwiseXor(desc) => self.fuse_binary_ops(desc, |lhs, rhs, out| { + FuseOp::BitwiseXor(BinaryFuseArgs { lhs, rhs, out }) + }), + IntOperationIr::BitwiseXorScalar(desc) => self + .fuse_scalar_ops(desc, |lhs, rhs, out| { + FuseOp::BitwiseXor(BinaryFuseArgs { lhs, rhs, out }) + }), + IntOperationIr::BitwiseNot(desc) => self.fuse_unary_ops(desc, |input, out| { + FuseOp::BitwiseNot(UnaryFuseArgs { input, out }) + }), + IntOperationIr::BitwiseLeftShift(desc) => self + .fuse_binary_ops(desc, |lhs, rhs, out| { + FuseOp::BitwiseLeftShift(BinaryFuseArgs { lhs, rhs, out }) + }), + IntOperationIr::BitwiseLeftShiftScalar(desc) => self + .fuse_scalar_ops(desc, |lhs, rhs, out| { + FuseOp::BitwiseLeftShift(BinaryFuseArgs { lhs, rhs, out }) + }), + IntOperationIr::BitwiseRightShift(desc) => self + .fuse_binary_ops(desc, |lhs, rhs, out| { + FuseOp::BitwiseRightShift(BinaryFuseArgs { lhs, rhs, out }) + }), + IntOperationIr::BitwiseRightShiftScalar(desc) => self + .fuse_scalar_ops(desc, |lhs, rhs, out| { + FuseOp::BitwiseRightShift(BinaryFuseArgs { lhs, rhs, out }) + }), + _ => false, + } + } + + fn fuse_numeric(&mut self, op: &NumericOperationIr) -> bool { + match op { + NumericOperationIr::Add(desc) => self.fuse_binary_ops(desc, |lhs, rhs, out| { + FuseOp::Add(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::AddScalar(desc) => self.fuse_scalar_ops(desc, |lhs, rhs, out| { + FuseOp::Add(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::Sub(desc) => self.fuse_binary_ops(desc, |lhs, rhs, out| { + FuseOp::Sub(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::SubScalar(desc) => self.fuse_scalar_ops(desc, |lhs, rhs, out| { + FuseOp::Sub(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::Mul(desc) => self.fuse_binary_ops(desc, |lhs, rhs, out| { + FuseOp::Mul(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::MulScalar(desc) => self.fuse_scalar_ops(desc, |lhs, rhs, out| { + FuseOp::Mul(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::Div(desc) => self.fuse_binary_ops(desc, |lhs, rhs, out| { + FuseOp::Div(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::DivScalar(desc) => self.fuse_scalar_ops(desc, |lhs, rhs, out| { + FuseOp::Div(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::Abs(desc) => { + self.fuse_unary_ops(desc, |input, out| FuseOp::Abs(UnaryFuseArgs { input, out })) + } + NumericOperationIr::Lower(desc) => self.fuse_binary_ops(desc, |lhs, rhs, out| { + FuseOp::Lower(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::LowerElem(desc) => self.fuse_scalar_ops(desc, |lhs, rhs, out| { + FuseOp::Lower(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::Greater(desc) => self.fuse_binary_ops(desc, |lhs, rhs, out| { + FuseOp::Greater(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::GreaterElem(desc) => self.fuse_scalar_ops(desc, |lhs, rhs, out| { + FuseOp::Greater(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::LowerEqual(desc) => self.fuse_binary_ops(desc, |lhs, rhs, out| { + FuseOp::LowerEqual(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::LowerEqualElem(desc) => self + .fuse_scalar_ops(desc, |lhs, rhs, out| { + FuseOp::LowerEqual(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::GreaterEqual(desc) => self + .fuse_binary_ops(desc, |lhs, rhs, out| { + FuseOp::GreaterEqual(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::GreaterEqualElem(desc) => self + .fuse_scalar_ops(desc, |lhs, rhs, out| { + FuseOp::GreaterEqual(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::Full(desc) => { + if !self.output_is_compatible(&desc.out) { + return false; + } + + self.fuser.fuse(|build| { + let input = build.scalar(&desc.value, desc.out.dtype); + let out = build.output(&desc.out)?; + + build.fuse_operation(FuseOp::Assign(UnaryFuseArgs { input, out })); + + Some(()) + }) + } + NumericOperationIr::Rem(desc) => self.fuse_binary_ops(desc, |lhs, rhs, out| { + FuseOp::Rem(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::RemScalar(desc) => self.fuse_scalar_ops(desc, |lhs, rhs, out| { + FuseOp::Rem(BinaryFuseArgs { lhs, rhs, out }) + }), + NumericOperationIr::Clamp(desc) => { + if !self.output_is_compatible(&desc.out) { + return false; + } + + self.fuser.fuse(|build| { + let input = build.input(&desc.tensor)?; + let min = build.scalar(&desc.min, desc.out.dtype); + let max = build.scalar(&desc.max, desc.out.dtype); + let out = build.output(&desc.out)?; + + build.fuse_operation(FuseOp::Clamp { + input, + min, + max, + out, + }); + + Some(()) + }) + } + _ => false, + } + } + + fn fuse_binary_ops(&mut self, desc: &BinaryOpIr, func: Func) -> bool + where + Func: Fn(FuseArg, FuseArg, FuseArg) -> FuseOp, + { + if !self.output_is_compatible(&desc.out) { + return false; + } + + self.fuser.fuse(|build| { + let lhs = build.input(&desc.lhs)?; + let rhs = build.input(&desc.rhs)?; + let out = build.output(&desc.out)?; + + build.fuse_operation(func(lhs, rhs, out)); + + Some(()) + }) + } + + fn fuse_unary_ops(&mut self, desc: &UnaryOpIr, func: Func) -> bool + where + Func: Fn(FuseArg, FuseArg) -> FuseOp, + { + self.fuse_unary_op(&desc.input, &desc.out, func) + } + + fn fuse_unary_op(&mut self, input: &TensorIr, out: &TensorIr, func: Func) -> bool + where + Func: Fn(FuseArg, FuseArg) -> FuseOp, + { + if !self.output_is_compatible(out) { + return false; + } + + self.fuser.fuse(|build| { + let input = build.input(input)?; + let out = build.output(out)?; + build.fuse_operation(func(input, out)); + Some(()) + }) + } + + fn fuse_scalar_ops(&mut self, desc: &ScalarOpIr, func: Func) -> bool + where + Func: Fn(FuseArg, FuseArg, FuseArg) -> FuseOp, + { + if !self.output_is_compatible(&desc.out) { + return false; + } + + self.fuser.fuse(|build| { + let elem = desc.lhs.dtype; + let lhs = build.input(&desc.lhs)?; + let rhs = build.scalar(&desc.rhs, elem); + let out = build.output(&desc.out)?; + + build.fuse_operation(func(lhs, rhs, out)); + + Some(()) + }) + } + + fn output_is_compatible(&mut self, out: &TensorIr) -> bool { + if self.current_output_shape.is_empty() { + self.current_output_shape.clone_from(&out.shape); + return true; + } + + let rank = self.current_output_shape.len(); + + // Rank should be equal. + if rank != out.shape.num_dims() { + return false; + } + + let mut updated = self.current_output_shape.clone(); + let mut should_update = false; + + #[allow(clippy::needless_range_loop)] + for i in 0..rank { + let curr = self.current_output_shape[i]; + let new = out.shape[i]; + + if curr == new { + continue; + } + + // Broadcast not enabled. + if !self.settings.broadcast { + return false; + } + + // Broadcasted on new dim. + if new == 0 { + continue; + } + + // Broadcasted on curr dim - update reference output shape. + if curr == 0 && self.settings.output_shape_updates { + should_update = true; + updated[i] = new; + continue; + } + + return false; + } + + if should_update { + // For now forced to have exact shape. + if updated != out.shape { + return false; + } + + self.current_output_shape.clone_from_slice(&out.shape); + } + + true + } +} + +#[derive(Debug, Clone)] +/// Builder wrapper to limit the number of bindings in generated kernels. +pub(crate) struct TryTraceFuser { + pub(crate) fuser: TraceFuser, + max_bindings: u32, + max_ops: u32, + added_ops: bool, +} + +impl TryTraceFuser { + fn new(max_bindings: u32, settings: FuseSettings) -> Self { + Self { + fuser: TraceFuser::new(settings), + max_bindings, + // A good default, avoid errors with for loops over only memory + // bound operations. + max_ops: 64, + added_ops: false, + } + } + + fn fuse(&mut self, add_ops: impl FnOnce(&mut TraceFuser) -> Option<()>) -> bool { + if self.fuser.num_ops_fused() > self.max_ops { + return false; + } + + // Always allow the first operation to be added. + if !self.added_ops { + self.added_ops = true; + + if add_ops(&mut self.fuser).is_none() { + return false; + } + return true; + } + + let mut cloned = self.fuser.clone(); + if add_ops(&mut cloned).is_none() { + return false; + } + + if cloned.estimate_bindings() > self.max_bindings { + return false; + } + + self.fuser = cloned; + true + } + + fn finish(&mut self, shape: Shape) -> FuseTrace { + self.fuser.finish(shape) + } +} diff --git a/crates/burn-cubecl-fusion/src/engine/launch/base.rs b/crates/burn-cubecl-fusion/src/engine/launch/base.rs new file mode 100644 index 0000000..9e5aed8 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/launch/base.rs @@ -0,0 +1,102 @@ +use crate::{ + CubeFusionHandle, + engine::{ + launch::{ + HandleInput, HandleOutput, LaunchPlan, executor::LaunchPlanExecutor, + input::InputPlanner, output::OutputPlanner, runner::TraceRunner, + vectorization::VectorizationPlanner, + }, + trace::{FuseTrace, TraceError, TuneOutput}, + }, +}; +use burn_fusion::stream::Context; +use cubecl::{Runtime, client::ComputeClient}; +use std::marker::PhantomData; + +/// The launcher is responsible to launch a fused kernel using the [TraceRunner] and a [FuseTrace]. +/// +/// TODO: We can reuse the same launcher between runs and avoid a lot of allocation, by simply +/// resetting the state. +pub struct FuseTraceLauncher<'a, R: Runtime, Runner: TraceRunner> { + trace: &'a FuseTrace, + runner: &'a Runner, + _runtime: PhantomData, +} + +impl<'a, R: Runtime, Runner: TraceRunner> FuseTraceLauncher<'a, R, Runner> { + /// Creates a new launcher. + pub fn new(trace: &'a FuseTrace, runner: &'a Runner) -> Self { + Self { + trace, + runner, + _runtime: PhantomData, + } + } + /// Launches the fuse kernel on the given device modifying the context. + pub fn launch( + &self, + client: &ComputeClient, + device: &R::Device, + context: &mut Context>, + ) -> Result, TraceError> { + let mut plan = LaunchPlan::new(&self.trace.blocks); + + InputPlanner::new(&self.trace.resources, &self.trace.blocks).run(context, &mut plan); + + OutputPlanner::new(&self.trace.resources, &self.trace.blocks) + .run(client, device, context, &mut plan); + + VectorizationPlanner::new(&self.trace.resources, &self.trace.blocks).run( + client, + self.runner, + context, + &mut plan, + ); + + match LaunchPlanExecutor::new(&self.trace.resources, &self.trace.blocks).execute::<_>( + client, + self.runner, + context, + plan, + ) { + Err(err) => { + self.rollback(context, err.handles_input, err.handles_output); + Err(err.error) + } + Ok(val) => Ok(val), + } + } + + fn rollback( + &self, + context: &mut Context>, + handle_inputs: Vec>, + handle_outputs: Vec>, + ) { + for input in handle_inputs { + match input { + HandleInput::Normal(input) => { + context + .handles + .register_handle(input.global_ir.id, input.handle_rollback()); + } + HandleInput::QuantValues(input) => { + context + .handles + .register_handle(input.global_ir.id, input.handle); + } + HandleInput::QuantParams(_) => { + // The scales are part of the quant data handle. + } + }; + } + for output in handle_outputs { + if let HandleOutput::Owned { + global_id, handle, .. + } = output + { + context.handles.register_handle(global_id, handle); + } + } + } +} diff --git a/crates/burn-cubecl-fusion/src/engine/launch/executor.rs b/crates/burn-cubecl-fusion/src/engine/launch/executor.rs new file mode 100644 index 0000000..40ebfe6 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/launch/executor.rs @@ -0,0 +1,316 @@ +use super::{HandleInput, HandleOutput, LaunchPlan, ReferenceSelection}; +use crate::engine::launch::runner::TraceRunner; +use crate::engine::trace::{FuseResources, TensorView, TraceError, TuneOutput, block::FuseBlock}; +use crate::{ + CubeFusionHandle, + engine::{ + codegen::ir::{ + FuseBlockConfig, FuseOp, FuseType, GlobalArgsLaunch, RefLayout, VirtualLayout, + }, + codegen::tensor::GlobalTensorArg, + }, +}; +use burn_fusion::stream::{Context, ScalarId}; +use burn_ir::ScalarIr; +use cubecl::{ + Runtime, + client::ComputeClient, + ir::{AddressType, Type}, + prelude::{InputScalar, TensorArg}, +}; +use std::marker::PhantomData; + +/// Execute a [plan](LaunchPlan) using a [runner](TraceRunner) modifying the [context](Context). +pub struct LaunchPlanExecutor<'a, R: Runtime> { + resources: &'a FuseResources, + blocks: &'a Vec, + _r: PhantomData, +} + +#[derive(new, Debug)] +pub struct ExecutionError> { + pub error: TraceError, + pub handles_input: Vec>, + pub handles_output: Vec>, +} + +impl<'a, R: Runtime> LaunchPlanExecutor<'a, R> { + pub fn new(resources: &'a FuseResources, blocks: &'a Vec) -> Self { + Self { + resources, + blocks, + _r: PhantomData, + } + } + + pub fn execute>( + self, + client: &ComputeClient, + runner: &Runner, + context: &mut Context>, + plan: LaunchPlan<'a, R>, + ) -> Result, ExecutionError> { + let mut num_writes = 0; + for b in plan.blocks.iter() { + for writes in b.writes.values() { + num_writes += writes.len(); + } + } + + #[cfg(feature = "autotune-checks")] + let mut tune_output = TuneOutput::Checked { + handles: std::collections::HashMap::new(), + }; + + #[cfg(not(feature = "autotune-checks"))] + let mut tune_output = TuneOutput::UnChecked(PhantomData); + + if num_writes == 0 { + // Nothing to write, can skip execution. + return Ok(tune_output); + } + + let mut inputs = GlobalArgsLaunch::default(); + let mut outputs = GlobalArgsLaunch::default(); + + // An aliased output reuses the input's buffer variable in the compiled kernel + // (`builder.inplace`), so it must be declared with the exact type the input was + // launched with — including its vector size. + let input_vector_sizes: Vec = plan + .handle_inputs + .iter() + .map(|hi| match hi { + HandleInput::Normal(hi) => hi.vector_size, + HandleInput::QuantValues(hi) => hi.vector_size, + HandleInput::QuantParams(_) => 1, + }) + .collect(); + + register_inputs(plan.handle_inputs.clone(), &mut inputs); + register_scalars( + self.resources.scalars.iter(), + self.resources.views.iter(), + context, + &mut inputs, + ); + register_outputs::( + plan.handle_outputs.clone(), + &input_vector_sizes, + &mut outputs, + &mut tune_output, + ); + + for layout in plan.runtime_layouts { + for s in layout.shape.iter() { + inputs.runtime_layouts.push(*s); + } + for s in layout.strides.iter() { + inputs.runtime_layouts.push(*s); + } + } + + let mut configs = Vec::with_capacity(plan.blocks.len()); + + for (block_plan, block) in plan.blocks.into_iter().zip(self.blocks) { + let reference = match block_plan.reference { + ReferenceSelection::Concrete { layout, .. } => RefLayout::Concrete(layout), + ReferenceSelection::VirtualShape { original, .. } => { + RefLayout::Virtual(VirtualLayout::Shape(original, block_plan.width)) + } + ReferenceSelection::SwapDims { original, dims } => { + RefLayout::Virtual(VirtualLayout::SwapDims(original, dims)) + } + ReferenceSelection::Reshaped { reshape_pos } => { + RefLayout::Virtual(VirtualLayout::Reshaped { + reshape_pos, + vector_size: block_plan.width, + }) + } + ReferenceSelection::Runtime { pos } => { + RefLayout::Virtual(VirtualLayout::Runtime { pos }) + } + ReferenceSelection::Searching => { + return Err(ExecutionError::new( + TraceError::ReferenceNotFound, + plan.handle_inputs, + plan.handle_outputs, + )); + } + }; + + let mut ops = Vec::::new(); + + for read_ops in block_plan.reads.into_values() { + for op in read_ops { + ops.push(op); + } + } + + for op in block.ops.iter() { + ops.push(op.clone()); + } + + for opsw in block_plan.writes.into_values() { + for op in opsw { + ops.push(op); + } + } + + let config = FuseBlockConfig { + rank: plan.rank, + ref_layout: reference, + ops, + width: block_plan.width, + }; + configs.push(config); + } + + Runner::run(runner, client, inputs, outputs, &configs).map_err(|err| { + ExecutionError::new( + TraceError::RunnerError(err), + plan.handle_inputs, + plan.handle_outputs, + ) + })?; + + Ok(tune_output) + } +} + +fn register_inputs( + handle_inputs: Vec>, + inputs: &mut GlobalArgsLaunch, +) { + for hi in handle_inputs { + match hi { + HandleInput::Normal(hi) => { + let at = hi.handle.required_address_type(); + let arg = hi.handle.into_tensor_arg(hi.global_ir.shape.clone()); + inputs.tensors.push(GlobalTensorArg::new( + arg, + hi.precision.into_type(hi.vector_size), + hi.broadcated, + at, + )); + } + HandleInput::QuantValues(hi) => { + let at = hi.handle.required_address_type(); + let arg = hi.handle.into_tensor_arg(hi.global_ir.shape.clone()); + inputs.tensors.push(GlobalTensorArg::new( + arg, + hi.precision.into_type(hi.vector_size), + false, + at, + )); + } + HandleInput::QuantParams(hi) => { + let at = hi.handle.required_address_type(); + let arg = hi.handle.into_tensor_arg(hi.shape.clone()); + inputs.tensors.push(GlobalTensorArg::new( + arg, + hi.precision.into_type(1), + false, + at, + )); + } + } + } +} + +fn register_outputs( + handle_outputs: Vec>, + input_vector_sizes: &[usize], + outputs: &mut GlobalArgsLaunch, + #[allow(unused_variables)] tune_output: &mut TuneOutput, +) { + for item in handle_outputs { + match item { + HandleOutput::Alias { + input_pos, + precision, + global_shape, + strides, + #[cfg(feature = "autotune-checks")] + debug_info, + } => { + outputs.tensors.push(GlobalTensorArg::new( + TensorArg::Alias { + input_pos, + strides, + shape: global_shape, + }, + // Must match the aliased input's launched type: the compiled kernel + // reuses the input's buffer variable, so declaring another vector + // size here would make the comptime type disagree with the actual + // variable and produce invalid casts on the write path. + precision.into_type(input_vector_sizes[input_pos]), + false, + AddressType::default(), + )); + + #[cfg(feature = "autotune-checks")] + if let TuneOutput::Checked { handles, .. } = tune_output { + handles.insert( + debug_info.relative_id, + (debug_info.global_shape.clone(), debug_info.handle.clone()), + ); + } + } + HandleOutput::Owned { + precision, + handle, + global_shape, + vectorization: vector_size, + #[cfg(feature = "autotune-checks")] + relative_id, + .. + } => { + let at = handle.required_address_type(); + + #[cfg(feature = "autotune-checks")] + if let TuneOutput::Checked { handles, .. } = tune_output { + handles.insert(relative_id, (global_shape.clone(), handle.clone())); + } + + let arg = handle.into_tensor_arg(global_shape.clone()); + + let elem = precision.into_elem(); + let ty = Type::new(elem.into()).with_vector_size(vector_size); + + outputs + .tensors + .push(GlobalTensorArg::new(arg, ty, false, at)); + } + } + } +} + +fn register_scalars<'h, R: Runtime>( + scalars: impl Iterator, + views: impl DoubleEndedIterator, + context: &mut Context>, + inputs: &mut GlobalArgsLaunch, +) { + for (precision, id) in scalars { + let dtype = precision.into_storage_type(); + match context.scalars.get(&ScalarId { value: *id }) { + Some(scalar) => match scalar { + ScalarIr::Float(val) => inputs.scalars.push(InputScalar::new(*val, dtype)), + ScalarIr::Int(val) => inputs.scalars.push(InputScalar::new(*val, dtype)), + ScalarIr::UInt(val) => inputs.scalars.push(InputScalar::new(*val, dtype)), + ScalarIr::Bool(val) => inputs.scalars.push(InputScalar::new(*val as u8, dtype)), + }, + None => panic!("Scalar ID not found"), + } + } + + for relative in views { + if let TensorView::Reshape { reshaped, .. } = relative { + let global = context.tensors.get(reshaped).unwrap(); + + for shape in global.shape.iter() { + inputs.reshapes.push(*shape); + } + } + } +} diff --git a/crates/burn-cubecl-fusion/src/engine/launch/input.rs b/crates/burn-cubecl-fusion/src/engine/launch/input.rs new file mode 100644 index 0000000..1126492 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/launch/input.rs @@ -0,0 +1,252 @@ +use super::{BlockPlan, HandleInput, InputReference}; +use super::{LaunchPlan, NormalHandleInput, PotentialInplace}; +use crate::CubeFusionHandle; +use crate::engine::launch::{QuantParamsHandleInput, QuantValuesHandleInput}; +use crate::engine::trace::block::FuseBlock; +use crate::engine::trace::{FuseResources, RegisterTensor, TensorView}; +use burn_fusion::stream::Context; +use burn_ir::{TensorIr, TensorStatus}; +use burn_std::quantization::params_shape; +use cubecl::Runtime; +use std::marker::PhantomData; + +/// Fetch and register [input handles](HandleInput). Also identifies potential inputs that +/// can be used inplace and/or as the [reference layout](super::super::ir::RefLayout). +pub struct InputPlanner<'a, R: Runtime> { + resources: &'a FuseResources, + blocks: &'a Vec, + _r: PhantomData, +} + +impl<'a, R: Runtime> InputPlanner<'a, R> { + pub fn new(resources: &'a FuseResources, blocks: &'a Vec) -> Self { + Self { + resources, + blocks, + _r: PhantomData, + } + } + + pub fn run(self, context: &mut Context>, plan: &mut LaunchPlan<'a, R>) { + for (pos, input) in self.resources.inputs.iter().enumerate() { + match input { + RegisterTensor::Normal(tensor_relative, precision) => { + let mut tensor_global = + context.tensors.get(&tensor_relative.id).unwrap().clone(); + // Fetching with the real status pops consumed (`ReadWrite`) tensors + // out of the container so that `can_mut()` sees the true reference + // count (memory pool + this handle). Fetching `ReadOnly` and keeping + // a clone in the container during planning made `can_mut()` always + // false, disabling every in-place alias. On execution failure the + // handle is restored by [FuseTraceLauncher::rollback](super::base). + let handle = context + .handles + .get_handle(&tensor_global.id, &tensor_relative.status); + + let mut new_strides = handle.strides.clone(); + + self.analyze(plan, pos, tensor_relative, &handle); + + if tensor_global.shape.rank() < plan.rank { + let num_elem: usize = tensor_global.shape.iter().product(); + for _ in 0..(plan.rank - tensor_global.shape.rank()) { + tensor_global.shape.insert(0, 1); + new_strides.insert(0, num_elem); + } + } + + plan.handle_inputs + .push(HandleInput::Normal(NormalHandleInput::new( + tensor_global, + tensor_relative, + *precision, + handle, + new_strides, + ))); + } + RegisterTensor::QuantValues(tensor_relative) => { + let tensor_global = context.tensors.get(&tensor_relative.id).unwrap().clone(); + let handle = context + .handles + .get_handle(&tensor_global.id, &TensorStatus::ReadOnly); + + let scheme = match tensor_relative.dtype { + burn_std::DType::QFloat(scheme) => scheme, + _ => unreachable!("Can't have quant data without QFloat"), + }; + let params = handle.params(scheme).unwrap(); + let precision = tensor_relative.dtype.into(); + let precision_scales = params.dtype.into(); + + let global_shape = tensor_global.shape.clone(); + let shape_params = params_shape(&global_shape, scheme.level); + plan.handle_inputs + .push(HandleInput::QuantValues(QuantValuesHandleInput { + relative_id: tensor_relative.id, + global_ir: tensor_global, + precision, + handle, + vector_size: 1, + })); + + plan.handle_inputs + .push(HandleInput::QuantParams(QuantParamsHandleInput { + precision: precision_scales, + handle: params, + shape: shape_params, + })); + } + RegisterTensor::QuantParams(_) => { + // It is registered at the same time as quant data. + // The order is important and the index in the vector as well, so that's why we + // have QuantParams. + } + } + } + } + + fn analyze( + &self, + plan: &mut LaunchPlan<'a, R>, + pos: usize, + tensor_relative: &'a TensorIr, + handle: &CubeFusionHandle, + ) { + if !self + .resources + .inputs_unhandled + .contains(&tensor_relative.id) + { + let mut is_a_view = false; + // For each view we try to see if it's not possible to set it as a reference input. + for view in self.resources.views.iter() { + for (block_plan, block) in plan.blocks.iter_mut().zip(self.blocks) { + is_a_view = is_a_view + || Self::analyze_view(pos, tensor_relative, block, block_plan, view); + } + } + + if !is_a_view { + self.analyze_normal(plan, pos, tensor_relative, handle); + } + } + } + + /// Analyzes if the given tensor can be used inplace in one of the block. + fn analyze_normal( + &self, + plan: &mut LaunchPlan<'a, R>, + pos: usize, + tensor_relative: &'a TensorIr, + handle: &CubeFusionHandle, + ) { + enum BlockInplaceSelection { + Notinit, + /// The block reads the input, and therefore can use it for inplace. + Selected(usize), + /// The same input is used in multiple blocks. + Unavailable, + } + + let mut block_inplace_selection = BlockInplaceSelection::Notinit; + + for (idx, block) in plan.blocks.iter().enumerate() { + if block.reads.contains_key(&tensor_relative.id) { + match block_inplace_selection { + BlockInplaceSelection::Notinit => { + block_inplace_selection = BlockInplaceSelection::Selected(idx); + } + BlockInplaceSelection::Selected(_) => { + block_inplace_selection = BlockInplaceSelection::Unavailable; + } + BlockInplaceSelection::Unavailable => {} + } + } + } + + if let BlockInplaceSelection::Selected(idx) = block_inplace_selection { + if self.blocks[idx].shape_ref != tensor_relative.shape { + return; + } + + let block_plan = &mut plan.blocks[idx]; + if tensor_relative.status == TensorStatus::ReadWrite { + // TODO: Autotune trials run on forked contexts whose extra handle clone + // makes `can_mut()` false, so trials never alias: the tuner measures the + // non-aliased kernel while the winner then executes the aliased one. + if self.blocks[idx].settings.inplace && handle.handle.can_mut() { + block_plan.potential_inplaces.push(PotentialInplace { + input_pos: pos, + tensor_relative, + strides: handle.strides.clone(), + }); + } + // Inplace tensors are normally really good as the reference layout, since + // it's normally better to be based on writes rather than on reads. + block_plan.potential_reference_input = + Some(InputReference::Normal { input_pos: pos }); + } else { + block_plan.potential_reference_input = + Some(InputReference::Normal { input_pos: pos }); + } + } + } + + /// Analyzes if the given tensor is also the view provided, and check if it can be used as the reference layout + /// for the given block. + fn analyze_view( + pos: usize, + tensor_relative: &'a TensorIr, + block: &FuseBlock, + block_plan: &mut BlockPlan<'a>, + view: &TensorView, + ) -> bool { + match view { + TensorView::Reshape { + reshaped, + original, + reshape_pos, + shape_relative, + } => { + if original == &tensor_relative.id || reshaped == &tensor_relative.id { + if block_plan.potential_reference_input.is_none() + && shape_relative == &block.shape_ref + { + block_plan.potential_reference_input = Some(InputReference::Reshaped { + reshape_pos: *reshape_pos, + }); + } + return true; + } + } + TensorView::SwapDims { + swapped, + original, + dims, + .. + } => { + if swapped == &tensor_relative.id { + return true; + } + + if original == &tensor_relative.id { + let shape = tensor_relative + .shape + .clone() + .swapped(dims.0, dims.1) + .unwrap(); + + if block_plan.potential_reference_input.is_none() && shape == block.shape_ref { + block_plan.potential_reference_input = Some(InputReference::SwapDims { + original_pos: pos, + dims: *dims, + }); + } + return true; + } + } + }; + + false + } +} diff --git a/crates/burn-cubecl-fusion/src/engine/launch/mod.rs b/crates/burn-cubecl-fusion/src/engine/launch/mod.rs new file mode 100644 index 0000000..94c9405 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/launch/mod.rs @@ -0,0 +1,11 @@ +pub(crate) mod executor; +pub(crate) mod input; +pub(crate) mod output; +pub(crate) mod runner; +pub(crate) mod vectorization; + +pub(crate) mod plan; +pub use plan::*; + +mod base; +pub use base::*; diff --git a/crates/burn-cubecl-fusion/src/engine/launch/output.rs b/crates/burn-cubecl-fusion/src/engine/launch/output.rs new file mode 100644 index 0000000..a81d859 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/launch/output.rs @@ -0,0 +1,720 @@ +use super::{ + super::codegen::ir::FuseType, BlockPlan, HandleOutput, InputReference, LaunchPlan, + NormalHandleInput, ReferenceSelection, +}; +use crate::{ + CubeFusionHandle, + engine::{ + codegen::ir::{FuseArg, FuseOp, LayoutInfo}, + launch::HandleInput, + settings::RefLayoutSetting, + trace::{FuseResources, RegisterTensor, RuntimeLayout, TensorView, block::FuseBlock}, + }, + strides_dyn_rank, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_fusion::stream::Context; +use burn_ir::{TensorId, TensorIr}; +use burn_std::Shape; +use burn_std::{ + Strides, + tensor::{ReshapeAction, contiguous_strides, is_contiguous, reshape_action}, +}; +use cubecl::{Runtime, client::ComputeClient}; + +/// Create or reuse handles for the outputs. +/// +/// It is also responsible to select the reference tensor. +pub struct OutputPlanner<'a, R: Runtime> { + resources: &'a FuseResources, + outputs_sorted: Vec>, + handles: Vec>>, + globals: Vec>, + blocks: &'a Vec, +} + +#[derive(Debug)] +struct OutputSorted<'a> { + pos_original: usize, + precision: FuseType, + tensor_relative: &'a TensorIr, +} + +#[derive(Debug)] +enum OutputKind { + Normal, + Inplace { + /// The position in the potential inplace vector + input_pos: usize, + }, + Transform(TensorView), +} + +impl<'a, R: Runtime> OutputPlanner<'a, R> { + pub fn new(resources: &'a FuseResources, blocks: &'a Vec) -> Self { + let mut outputs_sorted: Vec<_> = resources + .outputs + .iter() + .enumerate() + .filter_map(|(pos, entry)| match entry { + RegisterTensor::Normal(ir, p) => Some((pos, ir, p)), + RegisterTensor::QuantValues(_) => None, + RegisterTensor::QuantParams(_) => None, + }) + .map(|(pos, tensor, precision)| OutputSorted { + pos_original: pos, + precision: *precision, + tensor_relative: tensor, + }) + .collect(); + + outputs_sorted.sort_by(|a, b| { + let a_val: usize = a.tensor_relative.shape.iter().sum(); + let b_val: usize = b.tensor_relative.shape.iter().sum(); + + b_val.cmp(&a_val) + }); + + let mut handles = Vec::with_capacity(resources.outputs.len()); + let mut globals = Vec::with_capacity(resources.outputs.len()); + + for _ in 0..resources.outputs.len() { + handles.push(None); + globals.push(None); + } + + Self { + resources, + outputs_sorted, + handles, + globals, + blocks, + } + } + + pub fn run( + mut self, + client: &ComputeClient, + device: &R::Device, + context: &mut Context>, + plan: &mut LaunchPlan<'a, R>, + ) { + // So that we can borrow self during the iteration. + let mut outputs = Vec::new(); + core::mem::swap(&mut outputs, &mut self.outputs_sorted); + + for output in outputs.into_iter() { + let tensor_global = context + .tensors + .get(&output.tensor_relative.id) + .unwrap() + .clone(); + let strides = strides_dyn_rank(&tensor_global.shape); + let (kind, block_idx) = self.output_kind(plan, &tensor_global, &output, &strides); + + match kind { + OutputKind::Inplace { input_pos } => { + self.inplace_output( + context, + plan, + output, + tensor_global, + strides, + input_pos, + block_idx, + ); + } + OutputKind::Normal => { + self.normal_output( + client, + device, + context, + plan, + output, + tensor_global, + strides, + block_idx, + ); + } + OutputKind::Transform(TensorView::Reshape { original, .. }) => { + self.reshaped_output( + client, + device, + context, + plan, + output, + tensor_global, + strides, + original, + block_idx, + ); + } + OutputKind::Transform(TensorView::SwapDims { original, dims, .. }) => { + self.swapped_dims_output( + client, + device, + context, + plan, + output, + tensor_global, + original, + dims, + block_idx, + ); + } + } + } + + for (handle, global) in self.handles.into_iter().zip(self.globals) { + plan.handle_outputs.push(handle.unwrap()); + plan.global_outputs.push(global.unwrap()); + } + + for i in 0..plan.blocks.len() { + if !plan.blocks[i].reference.is_found() { + match self.blocks[i].settings.ref_layout { + RefLayoutSetting::SameAsBlock { block_pos } => { + plan.blocks[i].reference = + plan.blocks[block_pos as usize].reference.clone(); + } + _ => { + let new_runtime = Self::select_reference_from_inputs( + &self.blocks[i], + &mut plan.blocks[i], + &plan.handle_inputs, + ); + + if let Some(shape) = new_runtime { + let pos = plan.runtime_layouts.len(); + let mut shape_global = shape.clone(); + for (i, s) in shape.iter().enumerate() { + shape_global[i] = *context.shapes_relative2global.get(s).unwrap(); + } + + let strides = strides_dyn_rank(&shape_global); + + plan.blocks[i].reference = ReferenceSelection::Runtime { pos }; + plan.runtime_layouts.push(RuntimeLayout { + shape: shape_global, + strides, + }); + } + } + }; + } else { + Self::add_layout_info_inputs(&mut plan.blocks[i], &plan.handle_inputs); + } + } + + // Make sure dropped are correctly executed. + for id in self.resources.dropped.iter() { + if let Some(tensor_global) = context.tensors.get(id) { + context.handles.remove_handle(tensor_global.id); + } + } + } + + fn select_reference_from_inputs( + block: &FuseBlock, + block_plan: &mut BlockPlan<'_>, + handle_inputs: &[HandleInput], + ) -> Option { + if let Some(input_ref) = block_plan.potential_reference_input.take() { + match input_ref { + InputReference::Normal { input_pos } => { + let reference = handle_inputs + .get(input_pos) + .unwrap() + .as_normal() + .expect("Quant can't be used as inplace"); + + let set_ref_as_concrete = |block: &mut BlockPlan<'_>| { + block.reference = ReferenceSelection::Concrete { + layout: FuseArg::Input( + input_pos, + reference.precision, + LayoutInfo::IsRef, + ), + shape: reference.global_ir.shape.clone(), + strides: reference.handle.strides.clone(), + }; + }; + + let set_ref_as_virtual = |block: &mut BlockPlan<'_>| { + block.reference = ReferenceSelection::VirtualShape { + original: FuseArg::Input( + input_pos, + reference.precision, + LayoutInfo::Unknown, + ), + shape: reference.global_ir.shape.clone(), + strides: contiguous_strides(&reference.global_ir.shape), + }; + }; + + match block.settings.ref_layout { + RefLayoutSetting::Any => set_ref_as_concrete(block_plan), + RefLayoutSetting::SameAsBlock { .. } => { + // Skip set ref. + } + RefLayoutSetting::OnlyContiguous => { + if is_contiguous(&reference.global_ir.shape, &reference.handle.strides) + { + set_ref_as_concrete(block_plan) + } else { + set_ref_as_virtual(block_plan) + } + } + } + + Self::add_layout_info_inputs(block_plan, handle_inputs); + } + InputReference::SwapDims { original_pos, dims } => { + let reference = handle_inputs + .get(original_pos) + .unwrap() + .as_normal() + .expect("Quant can't be used in swap dims operation"); + block_plan.reference = ReferenceSelection::SwapDims { + original: FuseArg::Input( + original_pos, + reference.precision, + LayoutInfo::Unknown, + ), + dims, + }; + } + InputReference::Reshaped { reshape_pos } => { + block_plan.reference = ReferenceSelection::Reshaped { reshape_pos }; + } + }; + None + } else { + Some(block.shape_ref.clone()) + } + } + + fn add_layout_info_inputs(block: &mut BlockPlan<'_>, handle_inputs: &[HandleInput]) { + for hi in handle_inputs.iter().filter_map(|h| match h { + HandleInput::Normal(input) => Some(input), + _ => None, + }) { + let (strides, shape) = match &block.reference { + ReferenceSelection::Concrete { strides, shape, .. } + | ReferenceSelection::VirtualShape { strides, shape, .. } => (strides, shape), + _ => continue, + }; + + if strides == &hi.handle.strides + && shape == &hi.global_ir.shape + && let Some(ops) = block.reads.get_mut(&hi.relative_id) + { + for op in ops.iter_mut() { + if let FuseOp::Assign(op) = op { + op.input.add_layout_info(LayoutInfo::SameAsRef); + } + } + } + } + } + + fn output_kind( + &self, + plan: &mut LaunchPlan<'a, R>, + tensor_global: &TensorIr, + output: &OutputSorted, + strides: &[usize], + ) -> (OutputKind, usize) { + let mut block_idx = None; + for (i, block) in plan.blocks.iter().enumerate() { + if block.writes.contains_key(&output.tensor_relative.id) { + block_idx = Some(i); + break; + } + } + let block_idx = block_idx.unwrap(); + + if let Some(transform) = self.resources.views.iter().find(|v| match v { + TensorView::Reshape { reshaped, .. } => reshaped == &output.tensor_relative.id, + TensorView::SwapDims { swapped, .. } => swapped == &output.tensor_relative.id, + }) { + return (OutputKind::Transform(transform.clone()), block_idx); + } + + let block = &plan.blocks[block_idx]; + let ref_layout_setting = &self.blocks[block_idx].settings.ref_layout; + let kind = block + .potential_inplaces + .iter() + .enumerate() + .find(|(_pos, pi)| { + pi.tensor_relative.dtype == tensor_global.dtype + && pi.tensor_relative.shape == output.tensor_relative.shape + && &*pi.strides == strides + && if block.reference.is_found() { + // An already-selected reference must have compatible strides. + block.reference.compatible_strides_for_inplace(strides) + } else { + // When no reference has been selected yet, this output becomes + // the reference (see [Self::inplace_output]); requiring an + // existing reference here made the first output of every block + // ineligible, since the reference is only selected while + // processing outputs. Blocks that inherit their reference from + // another block are excluded: the inherited layout is unknown + // at this point, so it cannot be validated against the + // candidate's strides. + !matches!(ref_layout_setting, RefLayoutSetting::SameAsBlock { .. }) + } + }) + .map(|(pos, _)| OutputKind::Inplace { input_pos: pos }) + .unwrap_or(OutputKind::Normal); + + (kind, block_idx) + } + + #[allow(clippy::too_many_arguments)] + fn inplace_output( + &mut self, + context: &mut Context>, + plan: &mut LaunchPlan<'a, R>, + output: OutputSorted, + tensor_global: TensorIr, + strides: Strides, + input_index: usize, + block_idx: usize, + ) { + let block = &mut plan.blocks[block_idx]; + #[cfg(feature = "test-util")] + crate::inspect::record_inplace_alias(); + let potential_inplace = block.potential_inplaces.remove(input_index); + let handle_input = match plan.handle_inputs.get(potential_inplace.input_pos).unwrap() { + HandleInput::Normal(handle) => handle, + _ => { + unreachable!("Quant tensor handle can't be used inplace yet.") + } + }; + + // [Self::output_kind] only selects inplace when the reference is already + // validated or this output can become the reference, which blocks inheriting + // their reference from another block (`SameAsBlock`) never can. + debug_assert!( + block.reference.is_found() + || !matches!( + self.blocks[block_idx].settings.ref_layout, + RefLayoutSetting::SameAsBlock { .. } + ), + "inplace alias selected for a `SameAsBlock` block without a validated reference", + ); + + if !block.reference.is_found() { + // The aliased output shares the input's buffer and layout, so the reference + // is expressed with the output argument. Runners assume references are + // either output-concrete or virtual; an input-concrete reference here would + // be resolved against the wrong argument list. + block.reference = ReferenceSelection::Concrete { + layout: FuseArg::Output(output.pos_original, output.precision, LayoutInfo::IsRef), + shape: tensor_global.shape.clone(), + strides: handle_input.handle.strides.clone(), + }; + + if let Some(ops) = block.reads.get_mut(&handle_input.relative_id) { + for op in ops.iter_mut() { + if let FuseOp::Assign(op) = op { + op.input.add_layout_info(LayoutInfo::IsRef); + break; + }; + } + } + + if let Some(ops) = block.writes.get_mut(&output.tensor_relative.id) { + for op in ops { + if let FuseOp::Assign(op) = op { + op.out.add_layout_info(LayoutInfo::IsRef); + break; + } + } + }; + } else { + // Already validated, necessary for correctness. + if let Some(ops) = block.writes.get_mut(&output.tensor_relative.id) { + for op in ops { + if let FuseOp::Assign(op) = op { + op.out.add_layout_info(LayoutInfo::SameAsRef); + break; + } + } + }; + } + + context + .handles + .register_handle(tensor_global.id, handle_input.handle.clone()); + + self.handles[output.pos_original] = Some(HandleOutput::Alias { + input_pos: potential_inplace.input_pos, + precision: output.precision, + global_shape: tensor_global.shape.clone(), + strides, + #[cfg(feature = "autotune-checks")] + debug_info: super::HandleOutputAliasDebugInfo { + relative_id: output.tensor_relative.id, + handle: handle_input.handle.clone(), + global_shape: tensor_global.shape.clone(), + }, + }); + self.globals[output.pos_original] = Some(tensor_global); + } + + #[allow(clippy::too_many_arguments)] + fn normal_output( + &mut self, + client: &ComputeClient, + device: &R::Device, + context: &mut Context>, + plan: &mut LaunchPlan<'a, R>, + output: OutputSorted, + tensor_global: TensorIr, + strides: Strides, + block_idx: usize, + ) { + let block = &mut plan.blocks[block_idx]; + + if !block.reference.is_found() + && self.blocks[block_idx].shape_ref == output.tensor_relative.shape + && !matches!( + self.blocks[block_idx].settings.ref_layout, + RefLayoutSetting::SameAsBlock { .. } + ) + { + block.reference = ReferenceSelection::Concrete { + layout: FuseArg::Output(output.pos_original, output.precision, LayoutInfo::IsRef), + shape: tensor_global.shape.clone(), + strides: strides.clone(), + }; + + // Sometimes outputs that are manually handled don't have any write registered. + if let Some(ops) = block.writes.get_mut(&output.tensor_relative.id) { + for op in ops { + if let FuseOp::Assign(op) = op { + op.out.add_layout_info(LayoutInfo::IsRef); + break; + } + } + }; + } else if let ReferenceSelection::Concrete { + shape: ref_shape, + strides: ref_strides, + .. + } = &block.reference + && ref_strides == &strides + && ref_shape == &tensor_global.shape + && let Some(ops) = block.writes.get_mut(&output.tensor_relative.id) + { + for op in ops { + if let FuseOp::Assign(op) = op { + op.out.add_layout_info(LayoutInfo::SameAsRef); + break; + } + } + }; + + let dtype = tensor_global.dtype; + let size = + tensor_global.shape.iter().product::() * dtype_to_storage_type(dtype).size(); + + let handle = CubeFusionHandle { + client: client.clone(), + handle: client.empty(size), + device: device.clone(), + strides, + dtype, + qparams: None, + }; + + plan.rank = usize::max(tensor_global.shape.rank(), plan.rank); + context + .handles + .register_handle(tensor_global.id, handle.clone()); + + self.handles[output.pos_original] = Some(HandleOutput::Owned { + precision: output.precision, + handle, + global_shape: tensor_global.shape.clone(), + global_id: tensor_global.id, + relative_id: output.tensor_relative.id, + vectorization: 1, + }); + self.globals[output.pos_original] = Some(tensor_global); + } + + #[allow(clippy::too_many_arguments)] + fn reshaped_output( + &mut self, + client: &ComputeClient, + device: &R::Device, + context: &mut Context>, + plan: &mut LaunchPlan<'a, R>, + output: OutputSorted, + tensor_global: TensorIr, + strides: Strides, + original: TensorId, + block_idx: usize, + ) { + let block = &mut plan.blocks[block_idx]; + + let (pos_input, original_handle) = Self::find_child_input(&plan.handle_inputs, original); + + let dtype = tensor_global.dtype; + + let action = reshape_action( + &original_handle.global_ir.shape, + &original_handle.handle.strides, + &tensor_global.shape, + ); + + let update = match action { + ReshapeAction::UpdateStrides { strides } => Some(strides), + ReshapeAction::NoChange => Some(original_handle.handle.strides.clone()), + ReshapeAction::Recompute => None, + }; + + match update { + Some(strides) => { + // We modify the metadata instead. + remove_concrete_write(block, output.tensor_relative.id, output.pos_original); + + let handle = CubeFusionHandle { + client: client.clone(), + handle: original_handle.handle.handle.clone(), + device: device.clone(), + strides, + dtype, + qparams: original_handle.handle.qparams.clone(), + }; + context + .handles + .register_handle(tensor_global.id, handle.clone()); + + // IT will never be access, just a way to keep the original position working. + self.handles[output.pos_original] = Some(HandleOutput::Alias { + input_pos: pos_input, + precision: output.precision, + global_shape: tensor_global.shape.clone(), + strides: handle.strides.clone(), + #[cfg(feature = "autotune-checks")] + debug_info: super::HandleOutputAliasDebugInfo { + relative_id: output.tensor_relative.id, + handle: handle.clone(), + global_shape: tensor_global.shape.clone(), + }, + }); + self.globals[output.pos_original] = Some(tensor_global); + } + None => { + self.normal_output( + client, + device, + context, + plan, + output, + tensor_global, + strides, + block_idx, + ); + } + } + } + + #[allow(clippy::too_many_arguments)] + fn swapped_dims_output( + &mut self, + client: &ComputeClient, + device: &R::Device, + context: &mut Context>, + plan: &mut LaunchPlan<'a, R>, + output: OutputSorted, + tensor_global: TensorIr, + original: TensorId, + dims: (usize, usize), + block_idx: usize, + ) { + let block = &mut plan.blocks[block_idx]; + let (pos_input, original_handle) = Self::find_child_input(&plan.handle_inputs, original); + + let dtype = tensor_global.dtype; + + // TODO: Check if we can also remove the read, if we have a dead partial graph. + // + // We modify the metadata instead. + remove_concrete_write(block, output.tensor_relative.id, output.pos_original); + + let strides = original_handle.handle.strides.clone(); + + let mut handle = CubeFusionHandle { + client: client.clone(), + handle: original_handle.handle.handle.clone(), + device: device.clone(), + strides, + dtype, + qparams: original_handle.handle.qparams.clone(), + }; + handle.strides.swap(dims.0, dims.1); + + context + .handles + .register_handle(tensor_global.id, handle.clone()); + + // IT will never be access, just a way to keep the original position working. + self.handles[output.pos_original] = Some(HandleOutput::Alias { + input_pos: pos_input, + precision: output.precision, + global_shape: tensor_global.shape.clone(), + strides: handle.strides.clone(), + #[cfg(feature = "autotune-checks")] + debug_info: super::HandleOutputAliasDebugInfo { + relative_id: output.tensor_relative.id, + handle: handle.clone(), + global_shape: tensor_global.shape.clone(), + }, + }); + self.globals[output.pos_original] = Some(tensor_global); + } + + fn find_child_input( + handle_inputs: &[HandleInput], + original: TensorId, + ) -> (usize, &NormalHandleInput) { + handle_inputs + .iter() + .enumerate() + .find_map(|(pi, handle)| match handle { + HandleInput::Normal(handle) => match handle.relative_id == original { + true => Some((pi, handle)), + false => None, + }, + _ => None, // Quant tensor can't be reshaped. + }) + .unwrap() + } +} + +fn remove_concrete_write(block: &mut BlockPlan, id: TensorId, output_pos: usize) { + let ops = block.writes.remove(&id); + + if let Some(ops) = ops { + let mut keep = Vec::with_capacity(ops.len()); + + for op in ops { + if let FuseOp::Assign(args) = &op { + if let FuseArg::Output(pos, ..) = args.out { + if pos != output_pos { + keep.push(op); + } + } else { + keep.push(op); + } + } + } + block.writes.insert(id, keep); + } +} diff --git a/crates/burn-cubecl-fusion/src/engine/launch/plan.rs b/crates/burn-cubecl-fusion/src/engine/launch/plan.rs new file mode 100644 index 0000000..4cdd549 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/launch/plan.rs @@ -0,0 +1,272 @@ +use crate::{ + CubeFusionHandle, + engine::{ + codegen::ir::{FuseArg, FuseOp, FuseType}, + launch::vectorization::Vect, + trace::{RuntimeLayout, block::FuseBlock}, + }, +}; +use burn_ir::{TensorId, TensorIr}; +use burn_std::{Shape, Strides}; +use cubecl::{Runtime, ir::VectorSize}; +use std::collections::BTreeMap; + +/// The `LaunchPlan` is responsible for aggregating all runtime information required +/// to dispatch a fused kernel. +/// +/// It maps abstract IR tensors to memory handles, manages vectorization +/// strategies, and tracks layout transformations. +#[derive(Debug)] +pub struct LaunchPlan<'a, R: Runtime> { + /// The IR representation of tensors that are results of the fusion. + pub global_outputs: Vec, + /// Memory handles and metadata for all input tensors. + pub handle_inputs: Vec>, + /// Memory handles and metadata for all output tensors, including aliased inputs. + pub handle_outputs: Vec>, + /// The rank across all tensors in the plan. + /// + /// Smaller tensors are unsqueezed during launch. + pub rank: usize, + /// Detailed planning for each individual computation block within the fusion. + pub blocks: Vec>, + /// Mapping of tensor IDs to their specific vectorization factors. + pub vectorizations: BTreeMap, + /// Metadata for shapes and strides passed from the host when they cannot be + /// inferred from input tensors (e.g., complex deep fusions). + pub runtime_layouts: Vec, +} + +/// Information regarding the execution of a specific block of operations within a fusion. +#[derive(Debug)] +pub struct BlockPlan<'a> { + /// List of inputs that are candidates for in-place memory reuse within this block. + pub potential_inplaces: Vec>, + /// The input tensor chosen to define the iteration space, if any. + pub potential_reference_input: Option, + /// How the master layout is determined for this block. + pub reference: ReferenceSelection, + /// Mapping of tensor IDs to the read operations performed on them. + pub reads: BTreeMap>, + /// Mapping of tensor IDs to the write operations performed on them. + pub writes: BTreeMap>, + /// The width for the operations in this block. + pub width: VectorSize, +} + +/// Metadata for an input tensor being used as a reference for a block's layout. +#[derive(Debug)] +pub enum InputReference { + /// Standard input at the specified position. + Normal { input_pos: usize }, + /// Input that has an axis swapped. + SwapDims { + original_pos: usize, + dims: (usize, usize), + }, + /// Input that has been reshaped. + Reshaped { reshape_pos: usize }, +} + +/// Strategies for selecting the reference layout of a fused block. +/// +/// The reference layout determines how global indices are mapped to tensor coordinates. +#[derive(Clone, Debug)] +pub enum ReferenceSelection { + /// The engine is still calculating the optimal reference. + Searching, + /// Layout from a normal tensor. + Concrete { + layout: FuseArg, + shape: Shape, + strides: Strides, + }, + /// Layout from a swapped dim tensor. + SwapDims { + original: FuseArg, + dims: (usize, usize), + }, + /// Layout from a reshaped tensor. + Reshaped { reshape_pos: usize }, + /// Layout that has the shape of an input, but not its strides. + VirtualShape { + original: FuseArg, + shape: Shape, + strides: Strides, + }, + /// The layout is provided dynamically by the host at runtime. + Runtime { pos: usize }, +} + +impl LaunchPlan<'_, R> { + /// Creates a new `LaunchPlan` from a slice of fusion blocks. + /// + /// Initializes blocks with default "Searching" references and calculates + /// the initial max rank. + pub fn new(fuse_blocks: &[FuseBlock]) -> Self { + let mut rank = 0; + let mut blocks = Vec::with_capacity(fuse_blocks.len()); + + for b in fuse_blocks.iter() { + rank = usize::max(b.shape_ref.len(), rank); + let block = BlockPlan { + reference: ReferenceSelection::Searching, + reads: b.reads.clone(), + writes: b.writes.clone(), + width: 0, + potential_inplaces: Vec::new(), + potential_reference_input: None, + }; + blocks.push(block); + } + + LaunchPlan { + global_outputs: Vec::new(), + handle_inputs: Vec::new(), + handle_outputs: Vec::new(), + rank, + blocks, + vectorizations: Default::default(), + runtime_layouts: Default::default(), + } + } +} + +/// Debugging information for aliased handles when `autotune-checks` is enabled. +#[cfg(feature = "autotune-checks")] +#[derive(Debug, Clone)] +pub struct HandleOutputAliasDebugInfo { + pub handle: CubeFusionHandle, + pub relative_id: TensorId, + pub global_shape: Shape, +} + +/// Represents the output of a fused kernel execution. +#[derive(Debug, Clone)] +#[allow(clippy::large_enum_variant)] +pub enum HandleOutput { + /// An output that reuses the memory of an input tensor (In-place). + Alias { + /// Index of the input handle being aliased. + input_pos: usize, + /// Data type precision. + precision: FuseType, + global_shape: Shape, + strides: Strides, + #[cfg(feature = "autotune-checks")] + debug_info: HandleOutputAliasDebugInfo, + }, + /// An output that requires a newly allocated memory buffer. + Owned { + global_id: TensorId, + relative_id: TensorId, + precision: FuseType, + handle: CubeFusionHandle, + global_shape: Shape, + vectorization: VectorSize, + }, +} + +/// A standard input handle with associated layout and vectorization metadata. +#[derive(Debug, Clone)] +pub struct NormalHandleInput { + pub relative_id: TensorId, + pub global_ir: TensorIr, + pub precision: FuseType, + pub handle: CubeFusionHandle, + pub vector_size: VectorSize, + pub broadcated: bool, + /// Stores the original strides of the handle for restoration during plan rollback. + pub orig_strides: Strides, +} + +/// An input handle containing values for a quantized tensor. +#[derive(Debug, Clone)] +pub struct QuantValuesHandleInput { + pub relative_id: TensorId, + pub global_ir: TensorIr, + pub precision: FuseType, + pub handle: CubeFusionHandle, + pub vector_size: VectorSize, +} + +/// An input handle containing parameters (scales/offsets) for quantization. +#[derive(Debug, Clone)] +pub struct QuantParamsHandleInput { + pub precision: FuseType, + pub handle: CubeFusionHandle, + pub shape: Shape, +} + +/// Different types of inputs that can be passed to a fused kernel. +#[derive(Debug, Clone)] +pub enum HandleInput { + Normal(NormalHandleInput), + QuantValues(QuantValuesHandleInput), + QuantParams(QuantParamsHandleInput), +} + +impl HandleInput { + /// Returns a reference to the inner `NormalHandleInput` if the variant matches. + pub fn as_normal(&self) -> Option<&NormalHandleInput> { + match self { + HandleInput::Normal(normal) => Some(normal), + _ => None, + } + } +} + +impl NormalHandleInput { + /// Creates a new `NormalHandleInput` tracking original strides. + pub fn new( + tensor_global: TensorIr, + tensor_relative: &TensorIr, + precision: FuseType, + mut handle: CubeFusionHandle, + mut strides: Strides, + ) -> Self { + // Swap current handle strides with provided strides to track the original state for rollback. + core::mem::swap(&mut handle.strides, &mut strides); + Self { + precision, + handle, + relative_id: tensor_relative.id, + global_ir: tensor_global, + vector_size: 1, + broadcated: false, + orig_strides: strides, + } + } + + /// Restores the handle's original strides and returns the handle. + /// + /// Used when a plan is invalidated or needs to be rolled back. + pub fn handle_rollback(mut self) -> CubeFusionHandle { + core::mem::swap(&mut self.handle.strides, &mut self.orig_strides); + self.handle + } +} + +/// A candidate for in-place optimization. +#[derive(Debug)] +pub struct PotentialInplace<'a> { + /// Position of the input handle in the `handle_inputs` vector. + pub input_pos: usize, + /// Reference to the IR of the relative tensor. + pub tensor_relative: &'a TensorIr, + /// Current strides of the potential in-place candidate. + pub strides: Strides, +} + +impl ReferenceSelection { + pub fn is_found(&self) -> bool { + !matches!(self, Self::Searching) + } + + pub fn compatible_strides_for_inplace(&self, strides_inplace: &[usize]) -> bool { + match self { + ReferenceSelection::Concrete { strides, .. } => &**strides == strides_inplace, + _ => false, + } + } +} diff --git a/crates/burn-cubecl-fusion/src/engine/launch/runner.rs b/crates/burn-cubecl-fusion/src/engine/launch/runner.rs new file mode 100644 index 0000000..248019b --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/launch/runner.rs @@ -0,0 +1,96 @@ +use super::super::codegen::ir::{FuseBlockConfig, GlobalArgsLaunch}; +use crate::{ + CubeFusionHandle, + engine::launch::{ + LaunchPlan, + vectorization::{Vect, vectorization_default}, + }, +}; +use burn_fusion::stream::Context; +use burn_ir::{TensorId, TensorIr}; +use cubecl::prelude::*; +use std::collections::{BTreeMap, HashMap}; + +/// A trace runner is responsible for determining the vectorization factor as well as launching +/// a kernel based on global [inputs](GlobalArgsLaunch) and [outputs](GlobalArgsLaunch) +/// with provided [fuse block configs](FuseBlockConfig). +pub trait TraceRunner: Vectorization { + /// The error that might happen while running the trace. + type Error; + + /// Run the trace with the given inputs and outputs. + /// + /// There is one [fuse config](FuseBlockConfig) for each [block](super::block::FuseBlock) registered + /// in the [optimization builder](burn_fusion::OptimizationBuilder). + fn run<'a>( + &'a self, + client: &'a ComputeClient, + inputs: GlobalArgsLaunch, + outputs: GlobalArgsLaunch, + configs: &'a [FuseBlockConfig], + ) -> Result<(), Self::Error>; +} + +pub enum VectorizationHandle<'a, R: Runtime> { + NormalInput(&'a CubeFusionHandle, &'a TensorIr), + QuantValues(&'a CubeFusionHandle, &'a TensorIr), + QuantParams, +} + +impl<'a, R: Runtime> VectorizationHandle<'a, R> { + /// Returns if the current vectorization handle is from the given tensor id. + pub fn is_from_tensor(&self, id: TensorId) -> bool { + match self { + VectorizationHandle::NormalInput(_, tensor_ir) => tensor_ir.id == id, + VectorizationHandle::QuantValues(_, tensor_ir) => tensor_ir.id == id, + VectorizationHandle::QuantParams => false, + } + } +} + +#[derive(Default)] +pub struct VectorizationAxis { + axis: HashMap, +} + +impl VectorizationAxis { + pub fn get usize>(&self, id: TensorId, default: F) -> usize { + self.axis.get(&id).copied().unwrap_or_else(default) + } + pub fn insert(&mut self, id: TensorId, axis: usize) { + self.axis.insert(id, axis); + } +} + +pub trait Vectorization { + /// Returns the vectorization options. + fn axis(&self, _plan: &LaunchPlan<'_, R>) -> VectorizationAxis { + VectorizationAxis::default() + } + /// The vectorization factor for all inputs and outputs. + #[allow(clippy::too_many_arguments)] + fn vectorization<'a>( + &self, + _context: &Context>, + vectorizations: &mut BTreeMap, + inputs: impl Iterator>, + outputs: impl Iterator, + reshaped: impl Iterator, + swapped: impl Iterator, + vector_sizes: &[VectorSize], + max: VectorSize, + axis: VectorizationAxis, + ) { + vectorization_default( + vectorizations, + inputs, + outputs, + reshaped, + swapped, + vector_sizes, + &Default::default(), + max, + &axis, + ) + } +} diff --git a/crates/burn-cubecl-fusion/src/engine/launch/vectorization/base.rs b/crates/burn-cubecl-fusion/src/engine/launch/vectorization/base.rs new file mode 100644 index 0000000..10810e0 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/launch/vectorization/base.rs @@ -0,0 +1,474 @@ +use crate::{ + CubeFusionHandle, + engine::launch::runner::{VectorizationAxis, VectorizationHandle}, +}; +use burn_fusion::stream::Context; +use burn_ir::{TensorId, TensorIr}; +use cubecl::{Runtime, ir::VectorSize}; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; + +#[derive(Debug, Clone, Copy)] +pub enum Vect { + Broadcasted, + Aligned(VectorSize), +} + +impl Vect { + pub fn vector_size(&self) -> VectorSize { + match self { + Vect::Broadcasted => 1, + Vect::Aligned(val) => *val, + } + } + + pub fn is_broadcast(&self) -> bool { + matches!(self, Vect::Broadcasted) + } +} + +#[derive(Default, Clone, Serialize, Deserialize, Debug)] +pub struct VectorSizeOverrides { + state: Option>>, + default: Option>, +} + +#[allow(unused)] +impl VectorSizeOverrides { + pub fn overrides(&mut self, tensor_id: &TensorId, vector_sizes: Vec) { + let map = match &mut self.state { + Some(val) => val, + None => { + self.state = Some(BTreeMap::new()); + self.state.as_mut().unwrap() + } + }; + + map.insert(*tensor_id, vector_sizes); + } + pub fn overrides_default(&mut self, vector_sizes: Vec) { + self.default = Some(vector_sizes); + } + + pub fn mapping(&self, context: &Context>) -> Self { + match &self.state { + Some(state) => { + let mut state_new = BTreeMap::new(); + + for (k, v) in state.iter() { + let global = context.tensors.get(k).unwrap(); + state_new.insert(global.id, v.clone()); + } + + Self { + state: Some(state_new), + default: self.default.clone(), + } + } + None => Self { + state: None, + default: self.default.clone(), + }, + } + } + + pub fn tensor(&self, tensor_id: &TensorId) -> Option<&Vec> { + let map = match &self.state { + Some(val) => val, + None => return self.default.as_ref(), + }; + + match map.get(tensor_id) { + Some(val) => Some(val), + None => self.default.as_ref(), + } + } +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn vectorization_default<'a, R: Runtime>( + vectorizations: &mut BTreeMap, + inputs: impl Iterator>, + outputs: impl Iterator, + reshaped: impl Iterator, + swapped: impl Iterator, + vector_sizes: &[VectorSize], + overrides: &VectorSizeOverrides, + max: VectorSize, + axis: &VectorizationAxis, +) { + let swapped: Vec<_> = swapped.collect(); + + for input in inputs { + if let Some((s, o, mr, dims)) = swapped + .iter() + .find(|(_s, o, _mr, _dims)| input.is_from_tensor(o.id)) + { + let (handle, id) = match input { + VectorizationHandle::NormalInput(handle, tensor_ir) => (handle, &tensor_ir.id), + VectorizationHandle::QuantValues(..) => panic!("Can't be swapped"), + VectorizationHandle::QuantParams => panic!("Can't be swapped"), + }; + let val = vectorization_swapped( + handle, + s, + o, + *mr, + dims, + max, + axis, + vector_sizes, + overrides.tensor(id), + ); + multi_reads_vectorization_update(vectorizations, o.id, val); + } else { + match input { + VectorizationHandle::NormalInput(handle, tensor_ir) => { + let val = vectorization_input( + handle, + tensor_ir, + axis, + vector_sizes, + overrides.tensor(&tensor_ir.id), + ); + vectorizations.insert(tensor_ir.id, val); + } + VectorizationHandle::QuantValues(handle, tensor_ir) => { + let val = vectorization_input( + handle, + tensor_ir, + axis, + vector_sizes, + overrides.tensor(&tensor_ir.id), + ); + let num_quants = match tensor_ir.dtype { + burn_std::DType::QFloat(quant_scheme) => quant_scheme.num_quants(), + _ => panic!(""), + }; + let val = match val { + Vect::Broadcasted => Vect::Aligned(1), + Vect::Aligned(val) => Vect::Aligned(val.div_ceil(num_quants)), + }; + vectorizations.insert(tensor_ir.id, val); + } + VectorizationHandle::QuantParams => { + // Doesn't have vectorization for now. + } + }; + } + } + + for (reshaped, original, multi_reads) in reshaped { + let val = vectorization_reshape( + reshaped, + original, + multi_reads, + axis, + vector_sizes, + max, + overrides.tensor(&original.id), + ); + multi_reads_vectorization_update(vectorizations, original.id, val); + } + + for tensor in outputs { + let val = vectorization_output( + tensor, + axis, + vector_sizes, + max, + overrides.tensor(&tensor.id), + ); + vectorizations.insert(tensor.id, val); + } +} + +fn multi_reads_vectorization_update( + vectorizations: &mut BTreeMap, + original: TensorId, + vect: Vect, +) { + if let Some(ori_vect) = vectorizations.get(&original).cloned() { + match ori_vect { + Vect::Broadcasted => { + // keep the original as is. + } + Vect::Aligned(ori) => match vect { + Vect::Broadcasted => { + vectorizations.insert(original, Vect::Aligned(1)); + } + Vect::Aligned(new) => { + let val = if new != ori { 1 } else { new }; + vectorizations.insert(original, Vect::Aligned(val)); + } + }, + }; + } else { + vectorizations.insert(original, vect); + } +} + +// The default version uses the last dimension as vectorization axis and assumes a +// perpendicular contiguous vector. +fn vectorization_input( + handle: &CubeFusionHandle, + desc: &TensorIr, + axis: &VectorizationAxis, + vector_sizes: &[VectorSize], + overrides: Option<&Vec>, +) -> Vect { + let axis = axis.get(desc.id, || handle.strides.len() - 1); + let shape_axis = desc.shape[axis]; + + if shape_axis == 1 { + return Vect::Broadcasted; + } + + // Last dimension strides should be 1, otherwise vecX won't be contiguous. + if handle.strides[axis] != 1 { + return Vect::Aligned(1); + } + + // Reflect changes from https://github.com/tracel-ai/cubecl/pull/1312 + // Calculate the smallest non-zero stride among other dimensions + let mut next_stride = handle + .strides + .iter() + .enumerate() + .filter_map(|(i, &s)| (i != axis && s != 0).then_some(s)) + .min() + .unwrap_or(0); + + // Since when we calculate the vectorization of an input value, we will divide the + // vectorization by the num_quants, we need to adapt the stride check accordingly. + if let burn_std::DType::QFloat(quant_scheme) = handle.dtype { + next_stride *= quant_scheme.num_quants(); + }; + + let inner = |s: VectorSize| { + // The last dimension should be a multiple of the vector size or broadcated. + if shape_axis.is_multiple_of(s) && next_stride % s == 0 { + return Some(Vect::Aligned(s)); + } + None + }; + + match overrides { + Some(vals) => { + for s in vals { + if let Some(val) = inner(*s) { + return val; + } + } + } + None => { + for s in vector_sizes { + if let Some(val) = inner(*s) { + return val; + } + } + } + } + + Vect::Aligned(1) +} + +fn vectorization_output( + desc: &TensorIr, + axis: &VectorizationAxis, + vector_sizes: &[VectorSize], + max: VectorSize, + overrides: Option<&Vec>, +) -> Vect { + let axis = axis.get(desc.id, || desc.shape.rank() - 1); + + let inner = |s: VectorSize| { + // The dimension should be a multiple of the vector size. + if desc.shape[axis].is_multiple_of(s) && s <= max { + return Some(Vect::Aligned(s)); + } + + None + }; + match overrides { + Some(val) => { + for s in val { + if let Some(val) = inner(*s) { + return val; + } + } + } + None => { + for s in vector_sizes { + if let Some(val) = inner(*s) { + return val; + } + } + } + } + + Vect::Aligned(1) +} + +fn vectorization_reshape( + reshaped: &TensorIr, + original: &TensorIr, + multi_reads: bool, + axis: &VectorizationAxis, + vector_sizes: &[VectorSize], + max: VectorSize, + overrides: Option<&Vec>, +) -> Vect { + let axis = axis.get(reshaped.id, || reshaped.shape.rank() - 1); + let reshape_shape_axis = reshaped.shape[axis]; + + if !multi_reads && reshape_shape_axis == 1 { + return Vect::Broadcasted; + } + + // If the axis is not the last dim, didn't think of it, return Aligned(1) to be sure. + if axis != reshaped.shape.rank() - 1 { + return Vect::Aligned(1); + } + + let original_shape_axis = original.shape[original.shape.rank() - 1]; + + if original_shape_axis != reshape_shape_axis { + return Vect::Aligned(1); + } + + let inner = |s: VectorSize| { + if !multi_reads { + // The last dimension should be a multiple of the vector size or broadcated. + if reshape_shape_axis.is_multiple_of(s) && s <= max { + Some(Vect::Aligned(s)) + } else { + None + } + } else { + // Since the original tensor must share the same vectorization factor as the + // reshaped tensor, they must have compatible shapes when both are access + // independently. + if reshape_shape_axis.is_multiple_of(s) + && original_shape_axis.is_multiple_of(s) + && s <= max + { + Some(Vect::Aligned(s)) + } else { + None + } + } + }; + + match overrides { + Some(val) => { + for i in val { + if let Some(vect) = inner(*i) { + return vect; + } + } + } + None => { + for s in vector_sizes { + if let Some(vect) = inner(*s) { + return vect; + } + } + } + } + + Vect::Aligned(1) +} + +#[allow(clippy::too_many_arguments)] +fn vectorization_swapped( + handle: &CubeFusionHandle, + swapped: &TensorIr, + original: &TensorIr, + multi_reads: bool, + dims: &(usize, usize), + max: VectorSize, + axis: &VectorizationAxis, + vector_sizes: &[VectorSize], + overrides: Option<&Vec>, +) -> Vect { + let axis = axis.get(swapped.id, || swapped.shape.rank() - 1); + + let swapped_axis = swapped.shape[axis]; + let shape_axis = original.shape[axis]; + + let axis_index = axis; + let dim_index = if dims.0 == axis_index { + dims.1 + } else if dims.1 == axis_index { + dims.0 + } else { + axis_index + }; + + // Last dimension strides should be 1, otherwise vecX won't be contiguous. + if multi_reads { + if handle.strides[axis_index] != 1 { + return Vect::Aligned(1); + } + if handle.strides[dim_index] != 1 { + return Vect::Aligned(1); + } + } else if handle.strides[dim_index] != 1 { + return Vect::Aligned(1); + } + + if !multi_reads && swapped_axis == 1 { + return Vect::Broadcasted; + } + + // Reflect changes from https://github.com/tracel-ai/cubecl/pull/1312 + let mut next_stride = handle + .strides + .iter() + .enumerate() + .filter_map(|(i, &s)| (i != axis_index && s != 0).then_some(s)) + .min() + .unwrap_or(0); + + // Since when we calculate the vectorization of an input value, we will divide the + // vectorization by the num_quants, we need to adapt the stride check accordingly. + if let burn_std::DType::QFloat(quant_scheme) = handle.dtype { + next_stride *= quant_scheme.num_quants(); + }; + + let inner = |s: VectorSize| { + // The last dimension should be a multiple of the vector size or broadcated. + if multi_reads { + if swapped_axis.is_multiple_of(s) && next_stride % s == 0 && s <= max { + return Some(Vect::Aligned(s)); + } + } else if swapped_axis.is_multiple_of(s) + && shape_axis.is_multiple_of(s) + && next_stride % s == 0 + && s <= max + { + return Some(Vect::Aligned(s)); + } + None + }; + + match overrides { + Some(val) => { + for s in val { + if let Some(val) = inner(*s) { + return val; + } + } + } + None => { + for s in vector_sizes { + if let Some(val) = inner(*s) { + return val; + } + } + } + } + + Vect::Aligned(1) +} diff --git a/crates/burn-cubecl-fusion/src/engine/launch/vectorization/mod.rs b/crates/burn-cubecl-fusion/src/engine/launch/vectorization/mod.rs new file mode 100644 index 0000000..e7de0b2 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/launch/vectorization/mod.rs @@ -0,0 +1,5 @@ +mod base; +mod planner; + +pub use base::*; +pub use planner::*; diff --git a/crates/burn-cubecl-fusion/src/engine/launch/vectorization/planner.rs b/crates/burn-cubecl-fusion/src/engine/launch/vectorization/planner.rs new file mode 100644 index 0000000..7d6d2be --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/launch/vectorization/planner.rs @@ -0,0 +1,450 @@ +use super::{ + super::{BlockPlan, HandleOutput, LaunchPlan}, + Vect, +}; +use crate::{ + CubeFusionHandle, + engine::{ + launch::{ + HandleInput, + runner::{Vectorization, VectorizationHandle}, + }, + settings::VectorizationSetting, + trace::{FuseResources, TensorView, block::FuseBlock}, + }, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_fusion::stream::Context; +use burn_ir::TensorId; +use cubecl::{ + Runtime, + client::ComputeClient, + ir::{ElemType, StorageType, UIntKind}, +}; +use cubecl::{ + ir::VectorSize, + quant::scheme::{QuantScheme, QuantStore, QuantValue}, +}; +use std::marker::PhantomData; + +/// Select the best vectorization factor for each tensor handle. +pub struct VectorizationPlanner<'a, R: Runtime> { + resources: &'a FuseResources, + blocks: &'a Vec, + _r: PhantomData, +} + +impl<'a, R: Runtime> VectorizationPlanner<'a, R> { + pub fn new(resources: &'a FuseResources, blocks: &'a Vec) -> Self { + Self { + resources, + blocks, + _r: PhantomData, + } + } + pub fn run>( + self, + client: &ComputeClient, + runner: &Runner, + context: &Context>, + plan: &mut LaunchPlan<'a, R>, + ) { + let has_multiple_read = |tensor: &TensorId| { + let mut read_count = 0; + for block in plan.blocks.iter() { + read_count += block.reads.get(tensor).map(|a| a.len()).unwrap_or(0); + } + read_count > 1 + }; + let tensors_reshaped = self.resources.views.iter().filter_map(|view| match view { + TensorView::Reshape { + reshaped, original, .. + } => Some(( + context.tensors.get(reshaped).unwrap(), + context.tensors.get(original).unwrap(), + has_multiple_read(original), + )), + TensorView::SwapDims { .. } => None, + }); + let tensors_swapped = self.resources.views.iter().filter_map(|view| match view { + TensorView::SwapDims { + swapped, + original, + dims, + .. + } => Some(( + context.tensors.get(swapped).unwrap(), + context.tensors.get(original).unwrap(), + has_multiple_read(original), + dims, + )), + TensorView::Reshape { .. } => None, + }); + + let mut ref_elem = (ElemType::UInt(UIntKind::U64).into(), 8); + let mut quants_vector_sizes: Option> = None; + + for input in plan.handle_inputs.iter() { + let elem: StorageType = match input { + HandleInput::Normal(h) => dtype_to_storage_type(h.global_ir.dtype), + HandleInput::QuantValues(handle) => match handle.global_ir.dtype { + burn_std::DType::QFloat(scheme) => { + vector_sizes_quants(client, &mut quants_vector_sizes, scheme); + continue; + } + _ => panic!("Unable to retrieve the scheme for quantized values."), + }, + HandleInput::QuantParams(..) => continue, + }; + let elem_size = elem.size(); + + if ref_elem.1 >= elem_size { + ref_elem = (elem, elem_size); + } + } + for r in plan.global_outputs.iter() { + let elem: StorageType = dtype_to_storage_type(r.dtype); + let elem_size = elem.size(); + + if ref_elem.1 >= elem_size { + ref_elem = (elem, elem_size); + } + } + + let filtered = plan + .handle_inputs + .iter() + .map(|item| { + item.as_normal() + // Filter out indexed resources. + .map(|item| !self.resources.indexed.contains_key(&item.relative_id)) + .unwrap_or(true) + }) + .collect::>(); + + let vector_sizes = match quants_vector_sizes { + // Quantization normally triggers higher vectorization than anything else, no need to + // compare to ref elem. + Some(vector_sizes) => vector_sizes, + None => client + .io_optimized_vector_sizes(ref_elem.0.size()) + .collect::>(), + }; + let vectorization_axis = runner.axis(plan); + + runner.vectorization( + context, + &mut plan.vectorizations, + plan.handle_inputs + .iter() + .enumerate() + .filter_map(|(i, item)| { + if filtered[i] { + Some(match item { + HandleInput::Normal(h) => { + VectorizationHandle::NormalInput(&h.handle, &h.global_ir) + } + HandleInput::QuantValues(h) => { + VectorizationHandle::QuantValues(&h.handle, &h.global_ir) + } + HandleInput::QuantParams(_) => VectorizationHandle::QuantParams, + }) + } else { + None + } + }), + plan.global_outputs.iter(), + tensors_reshaped, + tensors_swapped, + &vector_sizes, + u8::MAX as usize, + vectorization_axis, + ); + + for tensor in self.resources.indexed.keys() { + let global = context.tensors.get(tensor).unwrap(); + plan.vectorizations.insert(global.id, Vect::Aligned(1)); + } + + let mut block_vectorization = Vec::with_capacity(self.blocks.len()); + for _ in 0..self.blocks.len() { + block_vectorization.push(Vec::new()); + } + + for (input_pos, handle) in plan.handle_inputs.iter_mut().enumerate() { + let (global_ir, relative_id) = match handle { + HandleInput::Normal(h) => (&h.global_ir, &h.relative_id), + HandleInput::QuantValues(h) => (&h.global_ir, &h.relative_id), + HandleInput::QuantParams(_) => continue, + }; + let (vect, br) = match plan.vectorizations.get(&global_ir.id) { + Some(v) => (v.vector_size(), v.is_broadcast()), + None => panic!("No vectorization factor found for {:?}", global_ir.id), + }; + + for (block_pos, block_plan) in plan.blocks.iter().enumerate() { + if block_plan.reads.contains_key(relative_id) { + block_vectorization[block_pos].push(BlockVectorization { + action: VectorizationAction::Input(input_pos), + potential: vect, + broadcasted: br, + }); + } + } + } + + for (output_pos, handle) in plan.handle_outputs.iter().enumerate() { + if let HandleOutput::Owned { + global_id, + relative_id, + .. + } = handle + { + for (block_pos, block_plan) in plan.blocks.iter().enumerate() { + if block_plan.writes.contains_key(relative_id) { + let vectorization = + plan.vectorizations.get(global_id).unwrap().vector_size(); + block_vectorization[block_pos].push(BlockVectorization { + action: VectorizationAction::Output(output_pos), + potential: vectorization, + broadcasted: false, + }); + } + } + } + } + + let mut previous_widths = Vec::with_capacity(block_vectorization.len()); + + // Unhandled inputs might not get included in any fused blocks for now. + // + // So we ensure they are vectorized by setting their vectorization before we set the + // vectorizations in blocks. + // + // Unhandled Outputs are correctly vectorized, so this is only necessary for inputs. + for input in self.resources.inputs_unhandled.iter() { + let pos = self + .resources + .inputs + .get_index(*input) + .unwrap_or_else(|| self.resources.inputs.get_index_quant(*input).unwrap()); + let input_global = context.tensors.get(input).unwrap(); + + match plan.vectorizations.get(&input_global.id).unwrap() { + Vect::Aligned(vect) => { + let handle = &mut plan.handle_inputs[pos]; + match handle { + HandleInput::Normal(handle) => { + handle.vector_size = *vect; + } + HandleInput::QuantValues(handle) => { + handle.vector_size = *vect; + } + HandleInput::QuantParams(_) => {} + } + } + Vect::Broadcasted => {} + } + } + + for ((tmp, block_plan), block) in block_vectorization + .into_iter() + .zip(plan.blocks.iter_mut()) + .zip(self.blocks) + { + match block.settings.vectorization { + VectorizationSetting::Activated => { + apply_vectorization_block( + tmp, + &mut plan.handle_inputs, + &mut plan.handle_outputs, + block_plan, + u8::MAX as usize, + ); + } + VectorizationSetting::SmallerOrEqualThanPreviousBlock { block_pos } => { + apply_vectorization_block( + tmp, + &mut plan.handle_inputs, + &mut plan.handle_outputs, + block_plan, + previous_widths[block_pos], + ); + if block_plan.width == 0 { + block_plan.width = previous_widths[block_pos]; + } + } + VectorizationSetting::EqualThanPreviousBlock { block_pos } => { + apply_vectorization_block( + tmp, + &mut plan.handle_inputs, + &mut plan.handle_outputs, + block_plan, + previous_widths[block_pos], + ); + // Enforces the width. + block_plan.width = previous_widths[block_pos]; + } + VectorizationSetting::Deactivated => { + apply_vectorization_block( + tmp, + &mut plan.handle_inputs, + &mut plan.handle_outputs, + block_plan, + 1, + ); + block_plan.width = 1; + } + } + + // When only virtual inputs/outputs are present for a block, we need to set a width. + if block_plan.width == 0 { + if let Some(w) = previous_widths.last() { + block_plan.width = *w; + } else { + block_plan.width = 1; + } + } + + previous_widths.push(block_plan.width); + } + } +} + +#[derive(Debug)] +enum VectorizationAction { + Input(usize), + Output(usize), +} + +#[derive(Debug)] +struct BlockVectorization { + action: VectorizationAction, + potential: VectorSize, + broadcasted: bool, +} + +fn apply_vectorization_block( + block_vectorization: Vec, + inputs: &mut [HandleInput], + outputs: &mut [HandleOutput], + block_plan: &mut BlockPlan, + max: VectorSize, +) { + for item in block_vectorization { + match item.action { + VectorizationAction::Input(pos) => { + let (vect, br) = if item.potential <= max { + (item.potential, item.broadcasted) + } else { + (1, false) + }; + + match &mut inputs[pos] { + HandleInput::Normal(input) => { + input.vector_size = vect; + input.broadcated = br; + } + HandleInput::QuantValues(input) => { + input.vector_size = vect; + } + HandleInput::QuantParams(_) => { + // Not vectorized + } + } + + if block_plan.width < vect { + block_plan.width = vect; + } + } + VectorizationAction::Output(pos) => { + if let HandleOutput::Owned { vectorization, .. } = &mut outputs[pos] { + let vect = if item.potential <= max { + item.potential + } else { + 1 + }; + *vectorization = vect; + + if block_plan.width < vect { + block_plan.width = vect; + } + } + } + } + } +} + +fn vector_sizes_quants( + client: &ComputeClient, + quants_vector_sizes: &mut Option>, + scheme: QuantScheme, +) { + match scheme.store { + QuantStore::Native => match scheme.value { + // Type sizes are the same so just treat fp8/fp4x2 as i8 + QuantValue::Q8F + | QuantValue::Q8S + | QuantValue::E4M3 + | QuantValue::E5M2 + | QuantValue::E2M1 => { + let vector_sizes = client + .io_optimized_vector_sizes(size_of::()) + .collect::>(); + + match &quants_vector_sizes { + Some(sizes) => { + if sizes[0] < vector_sizes[0] { + *quants_vector_sizes = Some(vector_sizes); + } + } + None => { + *quants_vector_sizes = Some(vector_sizes); + } + } + } + QuantValue::Q4F | QuantValue::Q4S | QuantValue::Q2F | QuantValue::Q2S => { + unreachable!("Can't store native sub-byte values") + } + }, + QuantStore::PackedU32(_) => { + let mut vector_sizes = client + .io_optimized_vector_sizes(size_of::()) + .collect::>(); + + for val in vector_sizes.iter_mut() { + *val *= scheme.num_quants(); + } + + let min = *vector_sizes.last().unwrap(); + + // We need to put back values that are not multiple of num_quants, but may be good + // vectorization factor for other handles in a fused trace. + for val in client.io_optimized_vector_sizes(size_of::()) { + if val < min { + vector_sizes.push(val); + } + } + + match &quants_vector_sizes { + Some(sizes) => { + if sizes[0] < vector_sizes[0] { + let mut min = *vector_sizes.last().unwrap(); + + while min > 1 { + min /= 2; + vector_sizes.push(min); + } + *quants_vector_sizes = Some(vector_sizes); + } + } + None => { + *quants_vector_sizes = Some(vector_sizes); + } + } + } + QuantStore::PackedNative(_) => { + panic!("Not yet supported") + } + }; +} diff --git a/crates/burn-cubecl-fusion/src/engine/mod.rs b/crates/burn-cubecl-fusion/src/engine/mod.rs new file mode 100644 index 0000000..304adef --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/mod.rs @@ -0,0 +1,7 @@ +pub(crate) mod codegen; +pub(crate) mod fuser; +pub(crate) mod launch; +pub(crate) mod scoring; +pub(crate) mod settings; + +pub mod trace; diff --git a/crates/burn-cubecl-fusion/src/engine/scoring.rs b/crates/burn-cubecl-fusion/src/engine/scoring.rs new file mode 100644 index 0000000..8270d27 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/scoring.rs @@ -0,0 +1,162 @@ +use crate::engine::{ + codegen::ir::{FuseArg, FuseOp, UnaryFuseArgs}, + trace::FuseTrace, +}; +use burn_ir::OperationIr; + +#[derive(Debug, Clone, Default)] +/// Tracks and evaluates the efficiency of operation fusion. +pub struct Scoring { + num_writes: usize, + num_reads: usize, + num_ops: usize, +} + +impl Scoring { + /// Resets the internal O counters. + pub fn reset(&mut self) { + self.num_writes = 0; + self.num_reads = 0; + self.num_ops = 0; + } + + /// Registers an unfused operation to the score, counting its total potential I/O. + pub fn register(&mut self, op: &OperationIr) { + self.num_writes += op.outputs().count(); + self.num_reads += op.inputs().count(); + self.num_ops += 1; + } + + /// Evaluates the efficiency of a fused trace by comparing its actual I/O + /// against the registered unfused I/O. Returns the number of saved I/O operations. + pub fn evaluate(&self, trace: &FuseTrace) -> u64 { + let mut num_reads_fused = 0; + let mut num_writes_fused = 0; + let mut num_penalty = 0; + + for b in trace.blocks.iter() { + // Count reads in block + for ops in b.reads.values() { + let result = self.count_fused_io(ops, |args| &args.input); + num_reads_fused += result.0; + num_penalty += result.1; + } + // Count writes in block + for ops in b.writes.values() { + let result = self.count_fused_io(ops, |args| &args.out); + num_writes_fused += result.0; + num_penalty += result.1; + } + } + + self.calculate_score(num_reads_fused, num_writes_fused, num_penalty) + } + + fn calculate_score(&self, reads_fused: usize, writes_fused: usize, num_penalty: usize) -> u64 { + // Those could be tweaked eventually. + + const FACTOR_IO: u64 = 100; + const FACTOR_LAUNCH: u64 = 10; + const FACTOR_PENALTY: u64 = 50; + + let num_fused = reads_fused + writes_fused; + let num_unfused = self.num_reads + self.num_writes; + + let score_io = match num_fused >= num_unfused { + true => 0, + false => (num_unfused - num_fused) as u64 * FACTOR_IO, + }; + + // We minus 1 since at least one kernel launch is necessary. + let score_launch = self.num_ops.saturating_sub(1) as u64 * FACTOR_LAUNCH; + + let score_penalty = num_penalty as u64 * FACTOR_PENALTY; + + (score_io + score_launch).saturating_sub(score_penalty) + } + + fn count_fused_io(&self, ops: &[FuseOp], arg_extractor: F) -> (usize, usize) + where + F: Fn(&UnaryFuseArgs) -> &FuseArg, + { + let mut num_io = 0; + let mut penalty = 0; + + for op in ops.iter() { + let FuseOp::Assign(args) = op else { + unreachable!() + }; + let count_normal = matches!( + arg_extractor(args), + FuseArg::Input(..) | FuseArg::Output(..) + ) as usize; + let count_view = matches!( + arg_extractor(args), + FuseArg::InputReshaped { .. } | FuseArg::InputSwapDims { .. } + ) as usize; + num_io += count_normal + count_view; + penalty += count_view; + } + + (num_io, penalty) + } +} + +#[cfg(test)] +#[allow(clippy::field_reassign_with_default)] +mod tests { + use super::*; + + #[test] + fn test_scoring_io_savings() { + let mut scoring = Scoring::default(); + scoring.num_reads = 2; + scoring.num_writes = 2; + scoring.num_ops = 2; + + let score = scoring.calculate_score(1, 1, 0); + assert_eq!(score, 210); + } + + #[test] + fn test_scoring_with_penalties() { + let mut scoring = Scoring::default(); + scoring.num_reads = 2; + scoring.num_writes = 2; + scoring.num_ops = 2; + + let score = scoring.calculate_score(1, 1, 1); + assert_eq!(score, 160); + } + + #[test] + fn test_penalty_outweighs_benefit() { + let mut scoring = Scoring::default(); + scoring.num_reads = 1; + scoring.num_writes = 1; + scoring.num_ops = 2; + + let score = scoring.calculate_score(1, 1, 1); + assert_eq!(score, 0); + } + + #[test] + fn test_scoring_no_ops() { + let scoring = Scoring::default(); + let score = scoring.calculate_score(0, 0, 0); + assert_eq!(score, 0); + } + + #[test] + fn test_reset() { + let mut scoring = Scoring { + num_writes: 10, + num_reads: 10, + num_ops: 10, + }; + scoring.reset(); + assert_eq!(scoring.num_writes, 0); + assert_eq!(scoring.num_reads, 0); + assert_eq!(scoring.num_ops, 0); + } +} diff --git a/crates/burn-cubecl-fusion/src/engine/settings.rs b/crates/burn-cubecl-fusion/src/engine/settings.rs new file mode 100644 index 0000000..8d0bb25 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/settings.rs @@ -0,0 +1,59 @@ +use serde::{Deserialize, Serialize}; + +/// Controls which operations can be fused. +#[derive(Clone, Copy, Debug, Serialize, Deserialize)] +pub struct FuseSettings { + /// Enables broadcasting of shapes. + pub broadcast: bool, + /// Enables output shape updates. + /// + /// When broadcast is enabled, the output shape can become bigger after a fusion, + /// therefore an update is needed. + pub output_shape_updates: bool, + /// Enables the reuse of input buffers. + pub inplace: bool, + /// Whether vectorization is enabled. + pub vectorization: VectorizationSetting, + /// How [reference layout](super::ir::RefLayout) selection is done. + pub ref_layout: RefLayoutSetting, +} + +impl Default for FuseSettings { + fn default() -> Self { + Self { + broadcast: true, + output_shape_updates: true, + inplace: true, + vectorization: VectorizationSetting::Activated, + ref_layout: RefLayoutSetting::Any, + } + } +} + +#[derive(Clone, Copy, Debug, Serialize, Deserialize)] +/// How vectorization is handled during fusion. +pub enum VectorizationSetting { + /// The biggest vector_size possible will be used. + Activated, + /// Equivalent to using vector_size of one. + Deactivated, + /// This is a good setting when a block processes values calculated from a previous block. + SmallerOrEqualThanPreviousBlock { block_pos: usize }, + /// This is a good setting when a block processes values calculated from a previous block. + EqualThanPreviousBlock { block_pos: usize }, +} + +#[derive(Clone, Copy, Debug, Serialize, Deserialize)] +/// Influence how the [reference layout](super::ir::RefLayout) selection is done. +pub enum RefLayoutSetting { + /// Any reference layout is allowed. + Any, + /// Only contiguous reference layout is allowed. + /// + /// Note that forcing a contiguous reference layout might reduce the opportunity of inplace + /// fusion. + OnlyContiguous, + SameAsBlock { + block_pos: u32, + }, +} diff --git a/crates/burn-cubecl-fusion/src/engine/trace/base.rs b/crates/burn-cubecl-fusion/src/engine/trace/base.rs new file mode 100644 index 0000000..d9fced8 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/trace/base.rs @@ -0,0 +1,377 @@ +use crate::engine::{ + codegen::ir::{FuseArg, FuseType}, + trace::block::FuseBlock, +}; +use burn_ir::{TensorId, TensorIr}; +use burn_std::{Shape, Strides}; +use cubecl::prelude::*; +use serde::{Deserialize, Serialize}; +use std::{ + collections::{BTreeMap, HashSet}, + marker::PhantomData, +}; + +#[cfg(feature = "autotune-checks")] +use crate::CubeFusionHandle; +#[cfg(feature = "autotune-checks")] +use burn_backend::TensorData; +#[cfg(feature = "autotune-checks")] +use burn_backend::cubecl::dtype_to_storage_type; +#[cfg(feature = "autotune-checks")] +use std::collections::HashMap; + +#[derive(Clone, Serialize, Deserialize, Debug)] +/// A trace contains all [blocks](FuseBlock) and the [resources](FuseResources) used by the +/// kernel. +pub struct FuseTrace { + pub blocks: Vec, + pub resources: FuseResources, +} + +impl core::fmt::Display for FuseTrace { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!(f, "FuseTrace")?; + for b in self.blocks.iter() { + writeln!(f, " - Block shape={:?}", b.shape_ref)?; + for (tensor, ops) in b.reads.iter() { + for op in ops.iter() { + writeln!(f, " - {op} <== {tensor}")?; + } + } + for op in b.ops.iter() { + writeln!(f, " - {op}")?; + } + for (tensor, ops) in b.writes.iter() { + for op in ops.iter() { + writeln!(f, " - {op} <== {tensor}")?; + } + } + } + + Ok(()) + } +} + +pub enum TuneOutput { + UnChecked(PhantomData), + #[cfg(feature = "autotune-checks")] + Checked { + handles: HashMap)>, + }, +} + +impl TuneOutput { + #[allow(unused_variables)] + pub fn merge(self, other: Self) -> Self { + let mut result = self; + + match &mut result { + TuneOutput::UnChecked(..) => {} + #[cfg(feature = "autotune-checks")] + TuneOutput::Checked { handles } => match other { + TuneOutput::UnChecked(..) => {} + TuneOutput::Checked { handles: o } => { + for (k, v) in o.into_iter() { + handles.insert(k, v); + } + } + }, + } + + result + } +} + +impl cubecl::tune::AutotuneOutput for TuneOutput { + #[cfg(feature = "autotune-checks")] + fn check_equivalence(&self, other: Self) { + use burn_backend::Tolerance; + use burn_std::DType; + + if let ( + TuneOutput::Checked { + handles: handles_ref, + }, + TuneOutput::Checked { handles }, + ) = (self, &other) + { + let mut num_checked = 0; + let mut num_handles = 0; + for (id, (shape, handle)) in handles_ref.iter() { + num_handles += 1; + if let Some((shape_other, other)) = handles.get(id) { + use burn_std::is_contiguous; + use cubecl::std::tensor::into_contiguous; + + let current_handle = if !is_contiguous(shape, &handle.strides) { + into_contiguous::( + &handle.client, + handle.clone().binding(shape.clone()), + dtype_to_storage_type(handle.dtype), + ) + .handle + } else { + handle.handle.clone() + }; + let other_handle = if !is_contiguous(shape, &other.strides) { + into_contiguous::( + &other.client, + other.clone().binding(shape.clone()), + dtype_to_storage_type(other.dtype), + ) + .handle + } else { + other.handle.clone() + }; + + let data_ref = handle.client.read_one(current_handle).unwrap(); + let data_other = other.client.read_one(other_handle).unwrap(); + let data_ref = TensorData::from_bytes(data_ref, shape.clone(), handle.dtype); + let data_other = + TensorData::from_bytes(data_other, shape_other.clone(), handle.dtype); + + match handle.dtype { + DType::F64 => { + data_ref.assert_approx_eq::(&data_other, Tolerance::permissive()) + } + DType::F32 => { + data_ref.assert_approx_eq::(&data_other, Tolerance::permissive()) + } + DType::F16 => data_ref + .assert_approx_eq::(&data_other, Tolerance::permissive()), + DType::BF16 => data_ref + .assert_approx_eq::(&data_other, Tolerance::permissive()), + _ => data_ref.assert_eq(&data_other, true), + } + num_checked += 1; + } else { + // Debug info for the tests. + println!("No tensor found for {id:?}=>{shape:?}"); + } + } + + // At least one check is needed per output when there is an output. + // + // Some optimizations might write more outputs than needed, so it might be fined if + // the number of handles is different, but at least one is required. + // + // An optimization might not create outputs if its dead code detection is triggered, + // therefore avoiding useless computation. + if num_handles > 0 { + assert!(num_checked >= 1); + } + } + } +} + +#[derive(Clone, Serialize, Deserialize, Debug, Default)] +/// Declare all resources used by the kernel, and potentially multiple [blocks](FuseBlock). +/// +/// # Notes +/// +/// Each block can't contain their own resources, since they are shared between blocks. The +/// vectorization factor of one input tensor must be the same for all blocks. +pub struct FuseResources { + pub outputs: RegisteredTensors, + pub inputs: RegisteredTensors, + pub scalars: Vec<(FuseType, u64)>, + // TODO: Making put a map of global registers. + pub views: Vec, + pub indexed: BTreeMap, + pub inputs_unhandled: Vec, + pub outputs_unhandled: Vec, + pub num_reshaped: usize, + /// Necessary to remove some entries from the context. + pub dropped: HashSet, + /// We know during fusion that we have to have those buffers has global. + /// The pos here can be interpreted as GLOBAL pos where the output pos are locals. + pub buffers: RegisteredTensors, + /// Global registers available everywhere. + /// + /// TODO: Not all registers should be globals. + pub registers: BTreeMap, +} + +#[derive(Clone, Serialize, Deserialize, Debug)] +pub struct RuntimeLayout { + pub shape: Shape, + pub strides: Strides, +} + +impl Default for RuntimeLayout { + fn default() -> Self { + Self { + shape: Shape::new([]), + strides: Strides::new(&[]), + } + } +} + +#[derive(Debug)] +pub enum TraceError { + ReferenceNotFound, + RunnerError(Err), +} + +#[derive(Clone, Serialize, Deserialize, Debug)] +pub enum TensorView { + Reshape { + reshaped: TensorId, + original: TensorId, + reshape_pos: usize, + shape_relative: Shape, + }, + SwapDims { + swapped: TensorId, + original: TensorId, + dims: (usize, usize), + }, +} + +#[derive(Default, Clone, Serialize, Deserialize, Debug)] +pub struct RegisteredTensors { + tensors: Vec, +} + +#[derive(Clone, Serialize, Deserialize, Debug)] +pub enum RegisterTensor { + Normal(TensorIr, FuseType), + QuantValues(TensorIr), + QuantParams(TensorId), +} + +impl RegisterTensor { + pub fn as_normal_tensor(&self) -> Option<(&TensorIr, &FuseType)> { + match self { + RegisterTensor::Normal(tensor_ir, precision) => Some((tensor_ir, precision)), + RegisterTensor::QuantValues(_) => None, + RegisterTensor::QuantParams(_) => None, + } + } +} + +impl RegisteredTensors { + /// Iterate over all the registered tensors. + pub fn iter(&self) -> impl Iterator { + self.tensors.iter() + } + + /// Consumes and iterate over all the registered tensors. + pub fn into_iter(self) -> impl Iterator { + self.tensors.into_iter() + } + + /// Returns the number of tensors registered. + pub fn len(&self) -> usize { + self.tensors.len() + } + + /// Retrieve the [tensor id](TensorId) at the given index. + pub fn get_id(&self, index: usize) -> Option { + self.tensors.get(index).map(|entry| match entry { + RegisterTensor::Normal(tensor_ir, _) => tensor_ir.id, + RegisterTensor::QuantValues(tensor_ir) => tensor_ir.id, + RegisterTensor::QuantParams(tensor_id) => *tensor_id, + }) + } + + /// Doesn't return quantized tensor. + pub fn get_index(&self, tensor_id: TensorId) -> Option { + self.tensors + .iter() + .enumerate() + .find(|(_pos, entry)| match entry { + RegisterTensor::Normal(tensor_ir, _) => tensor_ir.id == tensor_id, + RegisterTensor::QuantValues(_) => false, + RegisterTensor::QuantParams(_) => false, + }) + .map(|(pos, _)| pos) + } + + /// Get the index of a quantized tensor. + pub fn get_index_quant(&self, tensor_id: TensorId) -> Option { + self.tensors + .iter() + .enumerate() + .find(|(_pos, entry)| match entry { + RegisterTensor::Normal(..) => false, + RegisterTensor::QuantValues(tensor_ir) => tensor_ir.id == tensor_id, + RegisterTensor::QuantParams(_) => false, + }) + .map(|(pos, _)| pos) + } + + /// Doesn't return quantized tensor. + pub fn get(&self, tensor_id: TensorId) -> Option<(&TensorIr, &FuseType)> { + self.tensors + .iter() + .find(|entry| match entry { + RegisterTensor::Normal(tensor_ir, _) => tensor_ir.id == tensor_id, + RegisterTensor::QuantValues(_) => false, + RegisterTensor::QuantParams(_) => false, + }) + .and_then(|entry| match entry { + RegisterTensor::Normal(tensor_ir, fuse_precision) => { + Some((tensor_ir, fuse_precision)) + } + RegisterTensor::QuantValues(_) => None, + RegisterTensor::QuantParams(_) => None, + }) + } + + /// Insert a quantized tensor. + /// + /// It will return the positions for both the value tensor and param tensor. + pub fn insert_quant(&mut self, tensor: TensorIr) -> (usize, usize) { + if let Some(old) = self.tensors.iter().enumerate().find(|(_, val)| match &val { + RegisterTensor::QuantValues(tensor_ir) => tensor_ir == &tensor, + _ => false, + }) { + let values = old.0; + let params = values + 1; + return (values, params); + } + + let params = RegisterTensor::QuantParams(tensor.id); + let values = RegisterTensor::QuantValues(tensor); + let pos_values = self.len(); + self.tensors.push(values); + + let pos_params = self.len(); + self.tensors.push(params); + + (pos_values, pos_params) + } + + /// Insert a normal tensor with the given [precision](FusePrecision) in the current block. + pub fn insert(&mut self, precision: FuseType, tensor: TensorIr) -> usize { + if let Some(old) = self.tensors.iter().enumerate().find(|(_, val)| match &val { + RegisterTensor::Normal(tensor_ir, _) => tensor_ir.id == tensor.id, + _ => false, + }) { + return old.0; + } + + let value = RegisterTensor::Normal(tensor, precision); + let pos = self.len(); + + self.tensors.push(value); + + pos + } + + /// Update the already registered tensor with the given [tensor ir](TensorIr). + /// + /// # Notes + /// + /// This function only works with normal tensors, not quantized tensors. + pub fn update(&mut self, tensor: &TensorIr) { + if let Some(entry) = self.tensors.iter_mut().find(|entry| match entry { + RegisterTensor::Normal(tensor_ir, _) => tensor_ir.id == tensor.id, + _ => false, + }) && let RegisterTensor::Normal(tensor_ir, _) = entry + { + tensor_ir.status = tensor.status + } + } +} diff --git a/crates/burn-cubecl-fusion/src/engine/trace/block.rs b/crates/burn-cubecl-fusion/src/engine/trace/block.rs new file mode 100644 index 0000000..663aa13 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/trace/block.rs @@ -0,0 +1,541 @@ +use super::{FuseResources, RegisteredTensors, TensorView}; +use crate::engine::{ + codegen::ir::{FuseArg, FuseOp, FuseType, LayoutInfo, MultiBlockPos, UnaryFuseArgs}, + settings::FuseSettings, +}; +use burn_ir::{TensorId, TensorIr, TensorStatus}; +use burn_std::{DType, Shape, quantization::QuantParam}; +use serde::{Deserialize, Serialize}; +use std::collections::{BTreeMap, btree_map::Entry}; + +#[derive(Clone, Serialize, Deserialize, Debug)] +/// A block containing all [operations](FuseOp) as well as reads and writes for each tensor along +/// with the [fusion settings](FuseSettings). +pub struct FuseBlock { + /// Contains the [fusion settings](FuseSettings) associated to the current block. + pub settings: FuseSettings, + /// Contains all the [operations](FuseOp) registered in the current block. + pub ops: Vec, + /// The reference shape of the current block. + pub shape_ref: Shape, + /// Contains all tensor inputs of the current block except for manually handled tensors. + /// + /// # Notes + /// + /// Some reads might not have read operations registered, such as dequantization, but it's + /// important to be registered here for vectorization. Input tensors that are not + /// registered here must be vectorized manually. + pub reads: BTreeMap>, + /// Contains all tensor outputs of the current block except for manually handled tensors. + /// We can have multiple writes when the same variable is reused after in another block. + pub writes: BTreeMap>, +} + +#[derive(Clone, Debug)] +/// It is responsible to build a [trace](FuseBlock). +pub struct FuseBlockBuilder { + pub settings: FuseSettings, + locals: LocalVariablePool, + pub ops: Vec, + reads: BTreeMap>, + // Only for global registers. + writes: BTreeMap>, + // Output declared in this block alone. + outputs: RegisteredTensors, + pub outputs_unhandled: Vec, + pub local_inputs: BTreeMap, + /// The reference shape used by this block. + pub shape_ref: Shape, +} + +#[derive(Debug)] +/// How a quantized input can be read. +pub enum QuantInput { + /// If already dequantized, we cache the dequantization and returns the local variable + /// corresponding to the float value. + AlreadyDequantized { local: FuseArg }, + /// Otherwise we return the information necessary to dequantize the tensor. + Quantized { values: FuseArg, params: FuseArg }, +} + +impl FuseBlockBuilder { + pub fn new(settings: FuseSettings) -> Self { + Self { + settings, + locals: Default::default(), + ops: Default::default(), + reads: Default::default(), + writes: Default::default(), + outputs: Default::default(), + outputs_unhandled: Default::default(), + local_inputs: Default::default(), + shape_ref: Shape::new([]), + } + } + + /// Register an output tensor. + pub fn output(&mut self, tensor: &TensorIr, resources: &mut FuseResources) -> Option { + if resources.indexed.contains_key(&tensor.id) { + return None; + } + if matches!(tensor.dtype, DType::QFloat(..)) { + return None; + } + let precision = tensor.dtype.into(); + + let out = match self.locals.get(precision, tensor.id) { + Some(local) => local, + None => { + let out = self.locals.create(precision, tensor.id); + + self.outputs.insert(precision, tensor.clone()); + resources.outputs.insert(precision, tensor.clone()); + + out + } + }; + + Some(out) + } + + /// Register an input tensor. + pub fn multi_block_variable( + &mut self, + block_pos: usize, + tensor: &TensorIr, + global: bool, + ) -> Option { + let precision = tensor.dtype.into(); + + if let Some(val) = self.local_inputs.get(&tensor.id) { + return Some(val.clone()); + } + + let val = match self.locals.get(precision, tensor.id) { + Some(val) => val, + None => { + return None; + } + }; + + let arg = if global { + FuseArg::MultiBlockGlobal( + MultiBlockPos { + block_pos, + block_local_pos: self.writes.len(), + }, + val.precision(), + ) + } else { + FuseArg::MultiBlockLocal( + MultiBlockPos { + block_pos, + block_local_pos: self.writes.len(), + }, + val.precision(), + ) + }; + + let ops = match self.writes.get_mut(&tensor.id) { + Some(ops) => ops, + None => { + self.writes.insert(tensor.id, Vec::new()); + self.writes.get_mut(&tensor.id).unwrap() + } + }; + ops.push(FuseOp::Assign(UnaryFuseArgs { + input: val, + out: arg.clone(), + })); + + Some(arg) + } + + /// Register an input tensor. + pub fn input(&mut self, tensor: &TensorIr, resources: &mut FuseResources) -> Option { + if resources.indexed.contains_key(&tensor.id) { + return None; + } + + if matches!(tensor.dtype, DType::QFloat(..)) { + return None; + } + let precision = tensor.dtype.into(); + + if let Some(val) = self.local_inputs.get(&tensor.id) { + return Some(val.clone()); + } + + let arg = match self.locals.get(precision, tensor.id) { + Some(local) => { + resources.inputs.update(tensor); + + local + } + None => { + let input = if resources.outputs.get_index(tensor.id).is_some() { + if let Some(val) = resources.registers.get(&tensor.id) { + return Some(val.clone()); + }; + + let pos = resources.buffers.insert(precision, tensor.clone()); + FuseArg::Output(pos, precision, LayoutInfo::Unknown) + } else { + let pos = resources.inputs.insert(precision, tensor.clone()); + FuseArg::Input(pos, precision, LayoutInfo::Unknown) + }; + + let out = self.locals.create(precision, tensor.id); + + let reads = if let Entry::Vacant(e) = self.reads.entry(tensor.id) { + e.insert(Vec::with_capacity(1)); + self.reads.get_mut(&tensor.id).unwrap() + } else { + self.reads.get_mut(&tensor.id).unwrap() + }; + + reads.push(FuseOp::Assign(UnaryFuseArgs { + input, + out: out.clone(), + })); + + out + } + }; + + Some(arg) + } + + /// Register an input quantized tensor. + pub fn input_quant( + &mut self, + tensor: &TensorIr, + resources: &mut FuseResources, + ) -> Option { + if resources.indexed.contains_key(&tensor.id) { + return None; + } + + let precision = tensor.dtype.into(); + let precision_scales = match tensor.dtype { + DType::QFloat(scheme) => match scheme.param { + QuantParam::F32 => FuseType::F32, + QuantParam::F16 => FuseType::F16, + QuantParam::BF16 => FuseType::BF16, + QuantParam::UE8M0 | QuantParam::UE4M3 => { + unimplemented!("Unsupported fuse precision"); + } + }, + _ => return None, + }; + + let arg = match self.locals.get(precision, tensor.id) { + Some(local) => { + resources.inputs.update(tensor); + QuantInput::AlreadyDequantized { local } + } + None => { + let (new_input, q_index) = resources.inputs.insert_quant(tensor.clone()); + let input = FuseArg::Input(new_input, precision, LayoutInfo::Unknown); + let scales = FuseArg::Input(q_index, precision_scales, LayoutInfo::Unknown); + + // Important to flag that there is a read, even if no operation is registered. + if let Entry::Vacant(e) = self.reads.entry(tensor.id) { + e.insert(Vec::new()); + }; + + QuantInput::Quantized { + values: input, + params: scales, + } + } + }; + + Some(arg) + } + + /// Register an input with swapped dims. + pub fn input_swap_dims( + &mut self, + tensor: &TensorIr, + output: &TensorIr, + dims: (usize, usize), + resources: &mut FuseResources, + ) -> Option { + if matches!(tensor.dtype, DType::QFloat(..)) { + return None; + } + let precision = tensor.dtype.into(); + + let input_index = match self.locals.get(precision, tensor.id) { + Some(_) => { + // Can't fused an already fused input. + if resources.outputs.get(tensor.id).is_some() { + return None; + } + + match resources.inputs.get_index(tensor.id) { + Some(index) => { + resources.inputs.update(tensor); + index + } + None => { + return None; + } + } + } + // First appearance of the original: fusing a single-use view isn't worth it. The + // reshape/slice/swap_dims executes eagerly as a cheap shape/stride update and the result + // is then read with a clean, vectorizable layout. Fusing would instead recompute the view + // index on every read and can break vectorization. We only fuse a view when the original + // is already used elsewhere (the `Some` arm), i.e. when we genuinely need it in multiple + // layouts. + None => return None, + }; + + let out = self.output(output, resources)?; + let original = FuseArg::Input(input_index, precision, LayoutInfo::Unknown); + + let broadcasted = output.shape[output.shape.rank() - 1] == 0; + + resources.views.push(TensorView::SwapDims { + swapped: output.id, + original: tensor.id, + dims, + }); + + let input = FuseArg::InputSwapDims { + original: Box::new(original), + dims, + broadcasted, + }; + + let reads = if let Entry::Vacant(e) = self.reads.entry(tensor.id) { + e.insert(Vec::with_capacity(1)); + self.reads.get_mut(&tensor.id).unwrap() + } else { + self.reads.get_mut(&tensor.id).unwrap() + }; + + reads.push(FuseOp::Assign(UnaryFuseArgs { + input, + out: out.clone(), + })); + + Some(out) + } + + /// Register an input that is reshaped. + pub fn input_reshaped( + &mut self, + tensor: &TensorIr, + output: &TensorIr, + resources: &mut FuseResources, + ) -> Option { + if matches!(tensor.dtype, DType::QFloat(..)) { + return None; + } + let precision = tensor.dtype.into(); + + let input_index = match self.locals.get(precision, tensor.id) { + Some(_) => { + // Can't fused an already fused input. + if resources.outputs.get(tensor.id).is_some() { + return None; + } + + match resources.inputs.get_index(tensor.id) { + Some(index) => { + resources.inputs.update(tensor); + index + } + None => { + return None; + } + } + } + // First appearance of the original: fusing a single-use view isn't worth it. The + // reshape/slice/swap_dims executes eagerly as a cheap shape/stride update and the result + // is then read with a clean, vectorizable layout. Fusing would instead recompute the view + // index on every read and can break vectorization. We only fuse a view when the original + // is already used elsewhere (the `Some` arm), i.e. when we genuinely need it in multiple + // layouts. + None => return None, + }; + + let out = self.output(output, resources)?; + let original = FuseArg::Input(input_index, precision, LayoutInfo::Unknown); + + let mut shape = Vec::new(); + + let index = resources.num_reshaped; + resources.num_reshaped += 1; + + let rank = output.shape.rank(); + + for i in 0..output.shape.rank() { + let id = index * rank + i; + shape.push(FuseArg::ScalarShape(id)); + } + + resources.views.push(TensorView::Reshape { + reshaped: output.id, + original: tensor.id, + reshape_pos: index, + shape_relative: output.shape.clone(), + }); + + let input = FuseArg::InputReshaped { + original: Box::new(original), + shape, + broadcasted: output.shape[rank - 1] == 0, + }; + + let reads = if let Entry::Vacant(e) = self.reads.entry(tensor.id) { + e.insert(Vec::with_capacity(1)); + self.reads.get_mut(&tensor.id).unwrap() + } else { + self.reads.get_mut(&tensor.id).unwrap() + }; + + reads.push(FuseOp::Assign(UnaryFuseArgs { + input, + out: out.clone(), + })); + + Some(out) + } + + /// Build into a fuse block. + pub fn build( + &self, + resources: &FuseResources, + outputs: &mut RegisteredTensors, + buffers: &mut Vec, + ) -> FuseBlock { + let ops = self.ops.clone(); + let reads = self.reads.clone(); + let tensor_writes = self.tensor_writes(resources, buffers); + + let mut writes = self.writes.clone(); + + for (tensor, precision) in tensor_writes + .iter() + .filter_map(|entry| entry.as_normal_tensor()) + { + if let Some(local) = self.locals.get_any_precision(tensor.id) { + let out_index = outputs.insert(*precision, tensor.clone()); + + let ops = match writes.get_mut(&tensor.id) { + Some(ops) => ops, + None => { + writes.insert(tensor.id, Vec::new()); + writes.get_mut(&tensor.id).unwrap() + } + }; + + ops.push(FuseOp::Assign(UnaryFuseArgs { + input: local, + out: FuseArg::Output(out_index, *precision, LayoutInfo::Unknown), + })); + } + } + + FuseBlock { + settings: self.settings, + ops, + shape_ref: self.shape_ref.clone(), + reads, + writes, + } + } + + /// Return the tensor that needs to be written to. + /// + /// # Notes + /// + /// The buffers vector passed as input is only to track the intermediary buffer writes needed + /// during execution. + pub fn tensor_writes( + &self, + resources: &FuseResources, + buffers: &mut Vec, + ) -> RegisteredTensors { + let mut result = RegisteredTensors::default(); + + // All tensors where their latest representation is not read write should be written to since they + // are going to be used after the fused kernel by other operations. + for output in self.outputs.iter() { + if let Some((tensor, _precision)) = output.as_normal_tensor() { + // We get the latest representation from the resources, not just this block. + if let Some((tensor, precision)) = resources.outputs.get(tensor.id) { + if !matches!(tensor.status, TensorStatus::ReadWrite) { + result.insert(*precision, tensor.clone()); + } else if resources.buffers.get(tensor.id).is_some() + && !buffers.contains(&tensor.id) + { + result.insert(*precision, tensor.clone()); + // We make sure we don't write multiple time in the same buffer, only the + // earliest possible. + buffers.push(tensor.id); + } + } + } + } + + result + } +} + +#[derive(Default, Clone, Debug)] +pub struct LocalVariablePool { + values: BTreeMap>, +} + +impl LocalVariablePool { + fn get(&self, precision: FuseType, tensor_id: TensorId) -> Option { + if let Some(indexes) = self.values.get(&precision) + && let Some(index) = indexes.get(&tensor_id) + { + return Some(FuseArg::BlockLocal { + pos: *index, + ty: precision, + }); + } + + None + } + + fn get_any_precision(&self, tensor_id: TensorId) -> Option { + for (precision, indexes) in self.values.iter() { + if let Some(index) = indexes.get(&tensor_id) { + return Some(FuseArg::BlockLocal { + pos: *index, + ty: *precision, + }); + } + } + + None + } + + fn create(&mut self, precision: FuseType, tensor_id: TensorId) -> FuseArg { + if let Some(indexes) = self.values.get_mut(&precision) { + let new_index = indexes.len(); + indexes.insert(tensor_id, new_index); + return FuseArg::BlockLocal { + pos: new_index, + ty: precision, + }; + } + + let new_index = 0; + self.values + .insert(precision, BTreeMap::from_iter([(tensor_id, new_index)])); + + FuseArg::BlockLocal { + pos: new_index, + ty: precision, + } + } +} diff --git a/crates/burn-cubecl-fusion/src/engine/trace/fuser.rs b/crates/burn-cubecl-fusion/src/engine/trace/fuser.rs new file mode 100644 index 0000000..fab4c8d --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/trace/fuser.rs @@ -0,0 +1,321 @@ +use super::{ + super::{ + codegen::ir::{FuseArg, FuseOp, FuseType, LayoutInfo}, + settings::FuseSettings, + }, + FuseResources, + block::FuseBlockBuilder, +}; +use super::{FuseTrace, RegisteredTensors}; +use crate::engine::trace::block::QuantInput; +use burn_fusion::stream::ScalarId; +use burn_ir::{ScalarIr, TensorIr}; +use burn_std::{DType, Shape}; +use cubecl::quant::scheme::QuantParam; + +#[derive(Clone, Debug)] +/// It is responsible to create a [trace](FuseTrace) composed of multiple [blocks](super::block::FuseBlock). +/// +/// It mostly handles the [resources](KernelResources) needed by the generated fused kernel, and +/// delegates most of the work to the [block builder](FuseBlockBuilder). +pub struct TraceFuser { + settings: FuseSettings, + // The tensors returned by the block that don't need to be written to global memory. + block_current: FuseBlockBuilder, + blocks_previous: Vec, + resources: FuseResources, +} + +impl TraceFuser { + /// Create a new trace builder with the given bool precision and [fuse settings](FuseSettings). + pub fn new(settings: FuseSettings) -> Self { + Self { + settings, + block_current: FuseBlockBuilder::new(settings), + blocks_previous: Default::default(), + resources: Default::default(), + } + } + + /// Get the number of blocks that are closed. + pub fn num_previous_blocks(&self) -> usize { + self.blocks_previous.len() + } + + /// Tag a tensor as dropped. + pub fn fuse_dropped(&mut self, tensor: &TensorIr) { + self.resources.outputs.update(tensor); + self.resources.inputs.update(tensor); + self.resources.dropped.insert(tensor.id); + } + + /// Register an operation. + pub fn fuse_operation(&mut self, op: FuseOp) { + self.block_current.ops.push(op); + } + + /// The number of operations fused. + pub fn num_ops_fused(&self) -> u32 { + let mut num_ops_fused = 0; + + for block in self.blocks_previous.iter() { + num_ops_fused += block.ops.len(); + } + + num_ops_fused += self.block_current.ops.len(); + num_ops_fused as u32 + } + + /// Close the current block with the given reference shape and creates a new one with new [fusion settings](FuseSettings). + pub fn next_block(&mut self, shape_ref: Shape, settings: FuseSettings) { + let mut block_new = FuseBlockBuilder::new(settings); + core::mem::swap(&mut self.block_current, &mut block_new); + block_new.shape_ref = shape_ref; + self.blocks_previous.push(block_new); + self.settings = settings; + } + + // Estimate how many bindings are in use right now. This can return more than the actual number + // but should never return less. + pub fn estimate_bindings(&self) -> u32 { + let mut buffers = Vec::new(); + let mut estimation = 1; // Metadata takes one. + + // We assume we are not going to write multiple times in the same output buffer. + for b in self.blocks_previous.iter() { + estimation += b.tensor_writes(&self.resources, &mut buffers).len() as u32; + } + + estimation += self + .block_current + .tensor_writes(&self.resources, &mut buffers) + .len() as u32; + estimation += self.resources.inputs.len() as u32; + + estimation + } + + pub(crate) fn num_multi_block_local_inputs(&self) -> usize { + self.block_current + .local_inputs + .iter() + .filter(|(_, x)| matches!(x, FuseArg::MultiBlockLocal(..))) + .count() + } + + /// Tag the [tensor](TensorIr) as received from a previous block. + /// + /// This will avoid reading the input again and instead use le local version when possible. + pub fn block_local_input( + &mut self, + tensor: &TensorIr, + block_pos: usize, + global: bool, + ) -> FuseArg { + let block = &mut self.blocks_previous[block_pos]; + + let src_arg = match block.multi_block_variable(block_pos, tensor, global) { + Some(val) => val, + None => { + // We try to read the input if not present. + block.input(tensor, &mut self.resources); + block + .multi_block_variable(block_pos, tensor, global) + .unwrap() + } + }; + + self.resources.outputs.update(tensor); + + if global { + self.resources.registers.insert(tensor.id, src_arg.clone()); + } + + self.block_current + .local_inputs + .insert(tensor.id, src_arg.clone()); + src_arg + } + + /// Register an output tensor that won't be automatically synced into global memory. + /// + /// It is therefore the responsibility of the operation to write the result to given tensor. + pub fn output_unhandled(&mut self, tensor: &TensorIr) -> FuseArg { + let arg = self + .output(tensor) + .expect("Can't add a new output that is already used in an index operation"); + + self.resources.outputs_unhandled.push(arg.clone()); + self.block_current.outputs_unhandled.push(arg.clone()); + arg + } + + /// Register an input tensor that won't be automatically read into a local variable. + /// + /// It is therefore the responsibility of the operation to read the given tensor. + pub fn input_unhandled(&mut self, tensor: &TensorIr) -> FuseArg { + if self.resources.indexed.contains_key(&tensor.id) { + panic!("Can't add a new input that is already used in an index operation"); + } + + self.resources.outputs.update(tensor); + + let precision = tensor.dtype.into(); + let new_input = self.resources.inputs.insert(precision, tensor.clone()); + let arg = FuseArg::Input(new_input, precision, LayoutInfo::Unknown); + + self.resources.inputs_unhandled.push(tensor.id); + arg + } + + /// Register an input tensor. + pub fn input_quantized_unhandled(&mut self, tensor: &TensorIr) -> Option<(FuseArg, FuseArg)> { + if self.resources.indexed.contains_key(&tensor.id) { + panic!("Can't add a new input that is already used in an index operation"); + } + self.resources.outputs.update(tensor); + + let precision = tensor.dtype.into(); + let precision_scales = match tensor.dtype { + DType::QFloat(scheme) => match scheme.param { + QuantParam::F32 => FuseType::F32, + QuantParam::F16 => FuseType::F16, + QuantParam::BF16 => FuseType::BF16, + QuantParam::UE8M0 | QuantParam::UE4M3 => { + unimplemented!("Unsupported fuse precision"); + } + }, + _ => return None, + }; + + let (new_input, q_index) = self.resources.inputs.insert_quant(tensor.clone()); + let input = FuseArg::Input(new_input, precision, LayoutInfo::Unknown); + let scales = FuseArg::Input(q_index, precision_scales, LayoutInfo::Unknown); + + self.resources.inputs_unhandled.push(tensor.id); + Some((input, scales)) + } + + /// Register an input tensor. + pub fn input(&mut self, tensor: &TensorIr) -> Option { + if matches!(tensor.dtype, DType::QFloat(_)) { + return None; + } + + self.resources.outputs.update(tensor); + + self.block_current.input(tensor, &mut self.resources) + } + + /// Register an input tensor. + pub fn input_quantized(&mut self, tensor: &TensorIr) -> Option { + self.resources.outputs.update(tensor); + self.block_current.input_quant(tensor, &mut self.resources) + } + + /// Register an output tensor. + pub fn output(&mut self, tensor: &TensorIr) -> Option { + if matches!(tensor.dtype, DType::QFloat(_)) { + return None; + } + self.block_current.output(tensor, &mut self.resources) + } + + /// Register an input that will be accessed using custom indexing with no vectorization. + pub fn input_indexed(&mut self, tensor: &TensorIr) -> Option { + if matches!(tensor.dtype, DType::QFloat(_)) { + return None; + } + + if let Some(val) = self.resources.indexed.get(&tensor.id) { + self.resources.outputs.update(tensor); + return Some(val.clone()); + }; + + if self.resources.inputs.get(tensor.id).is_some() { + return None; + } + + if self.resources.outputs.get(tensor.id).is_some() { + return None; + } + + let input = self.input_unhandled(tensor); + self.resources.indexed.insert(tensor.id, input.clone()); + + Some(input) + } + + /// Register an input with swapped dims. + pub fn input_swap_dims( + &mut self, + tensor: &TensorIr, + output: &TensorIr, + dims: (usize, usize), + ) -> Option { + if matches!(tensor.dtype, DType::QFloat(_)) { + return None; + } + + self.resources.outputs.update(tensor); + self.block_current + .input_swap_dims(tensor, output, dims, &mut self.resources) + } + + /// Register an input that is reshaped. + pub fn input_reshaped(&mut self, tensor: &TensorIr, output: &TensorIr) -> Option { + if matches!(tensor.dtype, DType::QFloat(_)) { + return None; + } + + self.resources.outputs.update(tensor); + self.block_current + .input_reshaped(tensor, output, &mut self.resources) + } + + /// Register a scalar value. + pub fn scalar(&mut self, elem: &ScalarIr, dtype: DType) -> FuseArg { + let precision = dtype.into(); + let id = if let ScalarIr::UInt(value) = elem { + ScalarId { value: *value } + } else { + unreachable!() // should always be u64 + }; + + let new_index = self.resources.scalars.len(); + + self.resources.scalars.push((precision, id.value)); + FuseArg::Scalar(new_index, precision) + } + + /// Finish fusing and returns the created trace. + pub fn finish(&mut self, shape_ref: Shape) -> FuseTrace { + let mut resources = self.resources.clone(); + let mut outputs = RegisteredTensors::default(); + let mut buffers = Vec::new(); + + for tensor in resources.buffers.iter() { + let (tensor, ty) = tensor.as_normal_tensor().unwrap(); + outputs.insert(*ty, tensor.clone()); + } + + let mut blocks = Vec::new(); + + let mut register_block = |block: &FuseBlockBuilder| { + let block = block.build(&self.resources, &mut outputs, &mut buffers); + blocks.push(block); + }; + + for block in self.blocks_previous.iter() { + register_block(block); + } + self.block_current.shape_ref = shape_ref; + register_block(&self.block_current); + + // We update the output tensors registered to be the ones that are written to in global + // memory. + resources.outputs = outputs; + + FuseTrace { blocks, resources } + } +} diff --git a/crates/burn-cubecl-fusion/src/engine/trace/mod.rs b/crates/burn-cubecl-fusion/src/engine/trace/mod.rs new file mode 100644 index 0000000..286c232 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/engine/trace/mod.rs @@ -0,0 +1,7 @@ +pub(crate) mod block; + +mod base; +mod fuser; + +pub use base::*; +pub use fuser::*; diff --git a/crates/burn-cubecl-fusion/src/inspect.rs b/crates/burn-cubecl-fusion/src/inspect.rs new file mode 100644 index 0000000..f0792df --- /dev/null +++ b/crates/burn-cubecl-fusion/src/inspect.rs @@ -0,0 +1,26 @@ +//! Test-only introspection into launch-level decisions. +//! +//! Some launch decisions — like whether a fused kernel wrote its output in-place into an +//! input buffer — produce identical results either way, so value-based tests cannot catch +//! a silent regression. The counters here make those decisions observable. +//! +//! Counters are process-global: concurrent tests on the same process all increment them. +//! Assert on deltas (`>= n`), or serialize the test (e.g. `#[serial]`) when asserting that +//! a counter did *not* move. + +use core::sync::atomic::{AtomicU64, Ordering}; + +static INPLACE_ALIAS_COUNT: AtomicU64 = AtomicU64::new(0); + +/// The number of fused-kernel outputs planned to alias (reuse) an input buffer instead +/// of allocating their own, since process start. +/// +/// Counted at plan time: a plan whose execution later fails and rolls back still +/// increments the counter, so treat it as an upper bound on executed aliases. +pub fn inplace_alias_count() -> u64 { + INPLACE_ALIAS_COUNT.load(Ordering::Relaxed) +} + +pub(crate) fn record_inplace_alias() { + INPLACE_ALIAS_COUNT.fetch_add(1, Ordering::Relaxed); +} diff --git a/crates/burn-cubecl-fusion/src/lib.rs b/crates/burn-cubecl-fusion/src/lib.rs new file mode 100644 index 0000000..597c23b --- /dev/null +++ b/crates/burn-cubecl-fusion/src/lib.rs @@ -0,0 +1,14 @@ +#[macro_use] +extern crate derive_new; + +pub mod optim; + +#[cfg(feature = "test-util")] +pub mod inspect; + +mod base; + +pub(crate) mod engine; +pub(crate) mod tune; + +pub use base::*; diff --git a/crates/burn-cubecl-fusion/src/optim/base.rs b/crates/burn-cubecl-fusion/src/optim/base.rs new file mode 100644 index 0000000..76abf3e --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/base.rs @@ -0,0 +1,72 @@ +use crate::optim::{ + elemwise::{ElemwiseOptimization, ElemwiseOptimizationState}, + matmul::{MatmulOptimization, MatmulOptimizationState}, + reduce::{ReduceOptimization, ReduceOptimizationState}, + reduce_broadcasted::{ReduceBroadcastedOptimization, ReduceBroadcastedOptimizationState}, +}; +use cubecl::Runtime; +use serde::{Deserialize, Serialize}; + +/// Fusion optimization type for cubecl. +/// +/// More optimization variants should be added here. +#[allow(clippy::large_enum_variant)] +pub enum CubeOptimization { + ElementWise(ElemwiseOptimization), + Matmul(MatmulOptimization), + Reduce(ReduceOptimization), + ReduceBroadcasted(ReduceBroadcastedOptimization), +} + +impl core::fmt::Debug for CubeOptimization { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let value = self.to_opt_state(); + f.write_fmt(format_args!("{value:?}")) + } +} + +impl CubeOptimization { + /// Serializes the current optimization to its state. + pub fn to_opt_state(&self) -> CubeOptimizationState { + match self { + Self::ElementWise(value) => CubeOptimizationState::ElementWise(value.to_state()), + Self::Matmul(value) => CubeOptimizationState::Matmul(value.to_state()), + Self::Reduce(value) => CubeOptimizationState::Reduce(value.to_state()), + Self::ReduceBroadcasted(value) => { + CubeOptimizationState::ReduceBroadcasted(value.to_state()) + } + } + } +} + +impl burn_fusion::NumOperations for CubeOptimization { + fn len(&self) -> usize { + match self { + Self::ElementWise(op) => op.num_ops_fused(), + Self::Matmul(op) => op.num_ops_fused(), + Self::Reduce(op) => op.num_ops_fused(), + Self::ReduceBroadcasted(op) => op.num_ops_fused(), + } + } + + fn name(&self) -> &'static str { + match self { + CubeOptimization::ElementWise(..) => "ElementWise", + CubeOptimization::Matmul(..) => "Matmul", + CubeOptimization::Reduce(..) => "Reduce", + CubeOptimization::ReduceBroadcasted(..) => "ReduceBroadcasted", + } + } +} + +/// Fusion optimization state type for cubecl. +/// +/// More optimization variants should be added here. +#[allow(clippy::large_enum_variant)] +#[derive(Serialize, Deserialize, Debug)] +pub enum CubeOptimizationState { + ElementWise(ElemwiseOptimizationState), + Matmul(MatmulOptimizationState), + Reduce(ReduceOptimizationState), + ReduceBroadcasted(ReduceBroadcastedOptimizationState), +} diff --git a/crates/burn-cubecl-fusion/src/optim/elemwise/fuser.rs b/crates/burn-cubecl-fusion/src/optim/elemwise/fuser.rs new file mode 100644 index 0000000..90ed6a1 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/elemwise/fuser.rs @@ -0,0 +1,85 @@ +use super::optimization::ElemwiseOptimization; +use crate::{ + engine::{ + fuser::TraceOperationFuser, + settings::{FuseSettings, RefLayoutSetting, VectorizationSetting}, + }, + optim::CubeOptimization, +}; +use burn_fusion::OperationFuser; +use burn_std::Shape; +use cubecl::Runtime; + +/// Fuses element wise operations. +pub struct ElementWiseFuser { + fuser: TraceOperationFuser, + device: R::Device, +} + +impl Clone for ElementWiseFuser { + fn clone(&self) -> Self { + Self { + fuser: self.fuser.clone(), + device: self.device.clone(), + } + } +} + +impl ElementWiseFuser { + pub fn shape_id(&self) -> Shape { + self.fuser.current_output_shape.clone() + } + pub fn new(device: R::Device) -> Self { + let client = R::client(&device); + let props = client.properties(); + let max_bindings = props.hardware.max_bindings; + + Self { + fuser: TraceOperationFuser::new( + max_bindings, + FuseSettings { + broadcast: true, + output_shape_updates: true, + inplace: true, + vectorization: VectorizationSetting::Activated, + ref_layout: RefLayoutSetting::Any, + }, + ), + device, + } + } +} + +impl OperationFuser> for ElementWiseFuser { + fn fuse(&mut self, operation: &burn_ir::OperationIr) { + self.fuser.fuse(operation); + } + + fn finish(&mut self) -> CubeOptimization { + let client = R::client(&self.device); + let trace = self.fuser.finish(); + let elementwise = ElemwiseOptimization::new(trace, client, self.device.clone(), self.len()); + + CubeOptimization::ElementWise(elementwise) + } + + fn reset(&mut self) { + self.fuser.reset() + } + + fn status(&self) -> burn_fusion::FuserStatus { + self.fuser.status() + } + + fn properties(&self) -> burn_fusion::FuserProperties { + self.fuser.properties() + } + + fn len(&self) -> usize { + self.fuser.len() + } + + fn clone_dyn(&self) -> Box>> { + Box::new(self.clone()) + } +} diff --git a/crates/burn-cubecl-fusion/src/optim/elemwise/mod.rs b/crates/burn-cubecl-fusion/src/optim/elemwise/mod.rs new file mode 100644 index 0000000..7c4308c --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/elemwise/mod.rs @@ -0,0 +1,5 @@ +mod fuser; +mod optimization; + +pub use fuser::*; +pub use optimization::*; diff --git a/crates/burn-cubecl-fusion/src/optim/elemwise/optimization.rs b/crates/burn-cubecl-fusion/src/optim/elemwise/optimization.rs new file mode 100644 index 0000000..93d094e --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/elemwise/optimization.rs @@ -0,0 +1,141 @@ +use crate::{ + CubeFusionHandle, + engine::{ + codegen::{ + DynSize, + io::ref_len, + ir::{ + FuseArg, FuseBlockConfig, GlobalArgs, GlobalArgsLaunch, RefLayout, + multi_block_variables_init, + }, + kernel::{fuse_on_write, init_locals}, + }, + launch::{ + FuseTraceLauncher, + runner::{TraceRunner, Vectorization}, + }, + trace::FuseTrace, + }, +}; +use burn_fusion::stream::Context; +use cubecl::{CubeDim, calculate_cube_count_elemwise, client::ComputeClient, prelude::*}; +use serde::{Deserialize, Serialize}; + +#[derive(new)] +/// Fuse element wise operations into a single kernel. +pub struct ElemwiseOptimization { + pub(crate) trace: FuseTrace, + client: ComputeClient, + device: R::Device, + len: usize, +} + +#[derive(Serialize, Deserialize, Debug)] +/// State for the [elemwise optimization](ElemwiseOptimization). +pub struct ElemwiseOptimizationState { + trace: FuseTrace, + len: usize, +} + +impl ElemwiseOptimization { + /// Execute the optimization. + pub fn execute(&self, context: &mut Context>) { + let launcher = FuseTraceLauncher::new(&self.trace, &ElemwiseRunner); + + match launcher.launch(&self.client, &self.device, context) { + Ok(_) => (), + Err(err) => { + panic!("{err:?} - {:?}", self.trace); + } + } + } + + /// Number of element wise operations fused. + pub fn num_ops_fused(&self) -> usize { + self.len + } + + /// Create an optimization from its [state](ElemwiseOptimizationState). + pub fn from_state(device: &R::Device, state: ElemwiseOptimizationState) -> Self { + Self { + trace: state.trace, + len: state.len, + client: R::client(device), + device: device.clone(), + } + } + + /// Convert the optimization to its [state](ElemwiseOptimizationState). + pub fn to_state(&self) -> ElemwiseOptimizationState { + ElemwiseOptimizationState { + trace: self.trace.clone(), + len: self.len, + } + } +} + +pub struct ElemwiseRunner; + +impl Vectorization for ElemwiseRunner {} +impl TraceRunner for ElemwiseRunner { + type Error = LaunchError; // No error possible + + fn run<'a>( + &'a self, + client: &'a ComputeClient, + inputs: GlobalArgsLaunch, + outputs: GlobalArgsLaunch, + configs: &[FuseBlockConfig], + ) -> Result<(), Self::Error> { + let config = &configs[0]; + let shape = match &config.ref_layout { + RefLayout::Concrete(arg) => match arg { + FuseArg::Input(..) => inputs.shape_ref(&config.ref_layout, config.rank), + FuseArg::Output(..) => outputs.shape_ref(&config.ref_layout, config.rank), + _ => panic!("Invalid concreate ref layout"), + }, + RefLayout::Virtual(_) => inputs.shape_ref(&config.ref_layout, config.rank), + }; + let working_units = shape.iter().product::() / config.width; + let cube_dim = CubeDim::new(client, working_units); + let cube_count = calculate_cube_count_elemwise(client, working_units, cube_dim); + let address_type = inputs + .required_address_type() + .max(outputs.required_address_type()); + + unsafe { + elemwise_fuse::launch_unchecked( + client, + cube_count, + cube_dim, + address_type, + inputs, + outputs, + config.clone(), + ); + }; + + Ok(()) + } +} + +#[cube(launch_unchecked, address_type = "dynamic")] +fn elemwise_fuse( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + #[comptime] config: FuseBlockConfig, +) { + // We write no values for this fusion. + let values = Registry::>::new(); + let args = comptime![Vec::::new()]; + let pos = ABSOLUTE_POS; + + multi_block_variables_init(&config, &mut outputs.variables); + + let mut locals = init_locals(inputs, outputs, &config); + let length = ref_len(inputs, &*outputs, &locals, &config); + + if pos < length { + fuse_on_write::(inputs, outputs, &mut locals, pos, values, args, &config) + } +} diff --git a/crates/burn-cubecl-fusion/src/optim/matmul/args.rs b/crates/burn-cubecl-fusion/src/optim/matmul/args.rs new file mode 100644 index 0000000..d8fa80a --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/matmul/args.rs @@ -0,0 +1,571 @@ +use crate::engine::codegen::{ + io::ref_vector_size, + ir::{FuseArg, FuseBlockConfig, FuseType, GlobalArgs, LocalArgs, multi_block_variables_init}, + kernel::init_locals, + view::{FusedOutput, GlobalInput, GlobalInputExpand}, +}; +use cubecl::{ + intrinsic, + prelude::*, + quant::scheme::{QuantLevel, QuantScheme}, + std::{ + FastDivmod, + quant::{ + RunWithQuantType, + view::{QuantizedView, run_with_quant_type}, + }, + tensor::{ + View, ViewExpand, ViewMut, + layout::{Coords1d, Coords2d, VirtualLayout}, + }, + }, +}; +use cubek::{ + matmul::{ + args::{BatchedCoords, MatmulArgs}, + components::global::memory::{ + BatchLayout, BlockScaledLayout, GlobalLayout, GlobalLayoutConfig, GlobalLayoutExpand, + GlobalScaleLayout, GlobalScaleLayoutExpand, NoopLayout, + }, + }, + std::MatrixLayout, +}; +use serde::{Deserialize, Serialize}; +use std::marker::PhantomData; + +#[derive(Clone)] +pub struct FusedMatmulArgs; + +#[derive(CubeLaunch, CubeType)] +pub struct FusedMatmulInput { + global: GlobalArgs, + #[cube(comptime)] + config: FuseBlockConfig, + #[cube(comptime)] + a: MatmulArg, + #[cube(comptime)] + b: MatmulArg, + #[cube(comptime)] + c: Option, + #[cube(comptime)] + out: FuseArg, +} + +#[cube] +impl MatmulArgs for FusedMatmulArgs { + type Output = GlobalArgs; + type Input = FusedMatmulInput; + type State = FusedMatmulState; + type Config = (); + + fn init_state( + inputs: &Self::Input, + outputs: &mut Self::Output, + _config: (), + #[comptime] lhs_layout_config: GlobalLayoutConfig, + #[comptime] rhs_layout_config: GlobalLayoutConfig, + #[comptime] out_layout_config: GlobalLayoutConfig, + ) -> Self::State { + multi_block_variables_init(&inputs.config, &mut outputs.variables); + + let mut locals = init_locals(&inputs.global, outputs, &inputs.config); + let rank = comptime![inputs.config.rank]; + + let mut batch_shape = Sequence::new(); + let mut batch_strides_out = Sequence::new(); + + #[unroll] + for i in 0..rank - 2 { + batch_shape.push(FastDivmod::new_Fallback(locals.ref_shape[i] as u32)); + batch_strides_out.push(locals.ref_strides[i]); + } + + let batch_lhs = input_batch_layout( + &inputs.global, + &batch_shape, + comptime![inputs.a.clone()], + comptime![inputs.config.clone()], + ); + let batch_rhs = input_batch_layout( + &inputs.global, + &batch_shape, + comptime![inputs.b.clone()], + comptime![inputs.config.clone()], + ); + let batch_acc = match comptime![inputs.c.clone()] { + Some(c) => ComptimeOption::Some(input_batch_layout( + &inputs.global, + &batch_shape, + comptime![c], + comptime![inputs.config.clone()], + )), + None => ComptimeOption::new_None(), + }; + let batch_out = BatchLayout::new(batch_strides_out, batch_shape.clone()); + + FusedMatmulState::new( + inputs, + outputs, + &mut locals, + batch_lhs, + batch_rhs, + batch_acc, + VirtualLayout::new::(batch_out), + batch_shape, + &inputs.config, + lhs_layout_config, + rhs_layout_config, + out_layout_config, + ) + } + + fn view_lhs( + state: &Self::State, + ) -> View<'_, Lhs, BatchedCoords> { + global_view( + &state.inputs, + &state.locals, + &state.batch_shape, + comptime![state.a.clone()], + comptime![state.config.clone()], + state.lhs_layout_config, + ) + } + + fn batch_lhs( + state: &Self::State, + batch: usize, + ) -> usize { + state.a_batch.to_source_pos(batch) + } + + fn view_rhs( + state: &Self::State, + ) -> View<'_, Rhs, BatchedCoords> { + global_view( + &state.inputs, + &state.locals, + &state.batch_shape, + comptime![state.b.clone()], + comptime![state.config.clone()], + comptime![state.rhs_layout_config], + ) + } + + fn batch_rhs( + state: &Self::State, + batch: usize, + ) -> usize { + state.b_batch.to_source_pos(batch) + } + + fn view_acc( + state: &Self::State, + ) -> ComptimeOption> { + match comptime![state.c.clone()] { + Some(c) => { + let view = global_view( + &state.inputs, + &state.locals, + &state.batch_shape, + c, + comptime![state.config.clone()], + comptime![state.out_layout_config], + ); + ComptimeOption::Some(view) + } + None => ComptimeOption::new_None(), + } + } + + fn batch_acc( + state: &Self::State, + batch: usize, + ) -> usize { + #[comptime] + match &state.c_batch { + ComptimeOption::Some(c_batch) => c_batch.to_source_pos(batch), + ComptimeOption::None => batch, + } + } + + fn view_out( + state: &Self::State, + ) -> ViewMut<'_, EO, BatchedCoords> { + let rank = comptime![state.config.rank]; + + let shape_row = state.locals.ref_shape[rank - 2] as u32; + let shape_col = state.locals.ref_shape[rank - 1] as u32; + + let stride_row = state.locals.ref_strides[rank - 2]; + let stride_col = state.locals.ref_strides[rank - 1]; + + let mut outputs = state.outputs.clone(); + let mut locals = state.locals.clone(); + + let layout = GlobalLayout::new( + VirtualLayout::new::(NoopLayout::new()), + shape_row, + shape_col, + stride_row, + stride_col, + ref_vector_size(&state.locals), + 1u32, + state.out_layout_config, + ); + let buffer = FusedOutput::new( + &state.inputs, + &mut outputs, + &mut locals, + comptime![state.out.clone()], + comptime![state.config.clone()], + ); + ViewMut::new::(buffer, layout) + } + + fn batch_out( + state: &Self::State, + batch: usize, + ) -> usize { + state.out_batch.to_source_pos(batch) + } + + fn runtime_config( + _state: &Self::State, + ) { + } +} + +#[cube] +#[allow(clippy::missing_transmute_annotations)] +fn global_view( + inputs: &GlobalArgs, + locals: &LocalArgs, + batch_shape: &Sequence>, + #[comptime] arg: MatmulArg, + #[comptime] config: FuseBlockConfig, + #[comptime] layout_config: GlobalLayoutConfig, +) -> View<'static, E, BatchedCoords> { + let rank = comptime![config.rank]; + let data = comptime![arg.data().clone()]; + let data_tensor = match comptime![data.clone()] { + FuseArg::Input(pos, ..) => inputs.tensors.index(pos), + _ => panic!("Input must be concrete"), + }; + + let mut shape_row = data_tensor.tensor.shape(rank - 2) as u32; + let mut shape_col = data_tensor.tensor.shape(rank - 1) as u32; + let mut packing = comptime![1]; + + if arg.scheme().is_some() { + let scheme = arg.scheme().unwrap(); + let num_quants = scheme.num_quants() as u32; + comptime![packing = num_quants]; + match comptime![layout_config.matrix_layout] { + MatrixLayout::RowMajor => shape_col *= num_quants, + MatrixLayout::ColMajor => shape_row *= num_quants, + }; + } + + let shape = (shape_row, shape_col); + + // Noop for normal inputs because batch offset is cached, quantized uses logical batches + let batch_layout = match comptime![arg.clone()] { + MatmulArg::Normal(_) => VirtualLayout::new::(NoopLayout::new()), + MatmulArg::Quantized { data, .. } => { + let data_arg = comptime![MatmulArg::Normal(data)]; + input_batch_layout(inputs, batch_shape, data_arg, comptime![config.clone()]) + } + }; + + let data_layout = global_layout( + inputs, + shape, + batch_layout, + arg.data().clone(), + config.clone(), + data_tensor.tensor.vector_size(), + layout_config, + packing, + ); + let data_buf = GlobalInput::new(inputs, locals, data, comptime![config.clone()], None); + + match comptime![arg.clone()] { + MatmulArg::Normal(_) => View::new::(data_buf, data_layout), + MatmulArg::Quantized { scales, scheme, .. } => { + let scales_layout = match comptime![scheme.level] { + QuantLevel::Tensor => GlobalScaleLayout::new_PerTensor(shape), + QuantLevel::Block(block_size) => { + let block_size = comptime![block_size.as_dim::<2>()]; + + let scales_arg = comptime![MatmulArg::Normal(scales.clone())]; + let batch_layout = input_batch_layout( + inputs, + batch_shape, + scales_arg, + comptime![config.clone()], + ); + + let scales_layout = global_layout( + inputs, + shape, + batch_layout, + comptime![scales.clone()], + comptime![config.clone()], + 1usize, + layout_config, + 1u32, + ); + GlobalScaleLayout::new_BlockScaled(BlockScaledLayout::new( + shape, + scales_layout, + comptime![(block_size[0] as u32, block_size[1] as u32)], + )) + } + }; + let scales_buf = GlobalInput::new(inputs, locals, scales, config, None); + + // Redefine because of `Numeric` bound, kinda hacky but I can't figure out a way to + // assert `Vector::Scalar: Numeric` + let define!(T) = storage_type_of::(); + let view = create_quant_view_dynamic::( + data_buf, + data_layout, + scales_buf, + scales_layout, + scheme, + ); + // Safety: should be fine since `Vector` is guaranteed equal to `E` + comptime![unsafe { core::mem::transmute(view) }] + } + } +} + +#[cube] +fn input_batch_layout( + inputs: &GlobalArgs, + batch_shape: &Sequence>, + #[comptime] arg: MatmulArg, + #[comptime] config: FuseBlockConfig, +) -> VirtualLayout { + let rank = comptime![config.rank]; + match comptime![arg.clone()] { + MatmulArg::Normal(arg) => { + let data_tensor = match comptime![arg.clone()] { + FuseArg::Input(pos, ..) => inputs.tensors.index(pos), + _ => panic!("Input must be concrete"), + }; + + let mut batch_strides = Sequence::new(); + #[unroll] + for i in 0..rank - 2 { + let shape = data_tensor.tensor.shape(i); + let stride = select(shape == 1, 0, data_tensor.tensor.stride(i)); + batch_strides.push(stride); + } + + VirtualLayout::new::(BatchLayout::new(batch_strides, batch_shape.clone())) + } + MatmulArg::Quantized { .. } => VirtualLayout::new::(NoopLayout::new()), + } +} + +#[cube] +fn global_layout( + inputs: &GlobalArgs, + shape: Coords2d, + batch_layout: VirtualLayout, + #[comptime] arg: FuseArg, + #[comptime] config: FuseBlockConfig, + #[comptime] vector_size: VectorSize, + #[comptime] layout_config: GlobalLayoutConfig, + #[comptime] packing: u32, +) -> GlobalLayout { + let rank = comptime![config.rank]; + let data_tensor = match comptime![arg.clone()] { + FuseArg::Input(pos, ..) => inputs.tensors.index(pos), + _ => panic!("Input must be concrete"), + }; + + let (shape_row, shape_col) = shape; + + let stride_row = data_tensor.tensor.stride(rank - 2); + let stride_col = data_tensor.tensor.stride(rank - 1); + + GlobalLayout::new( + batch_layout, + shape_row, + shape_col, + stride_row, + stride_col, + vector_size, + packing, + layout_config, + ) +} + +struct CreateQuantView<'a, E: Numeric, N: Size> { + scope: &'a Scope, + data_buf: GlobalInputExpand, + data_layout: GlobalLayoutExpand, + scales_buf: GlobalInputExpand, + scales_layout: GlobalScaleLayoutExpand, + scheme: QuantScheme, + _ty: PhantomData<(E, N)>, +} + +impl<'a, E: Numeric, N: Size> RunWithQuantType for CreateQuantView<'a, E, N> { + type Output = ViewExpand<'static, Vector, BatchedCoords>; + + fn execute(self) -> Self::Output { + create_quant_view::expand::( + self.scope, + self.data_buf, + self.data_layout, + self.scales_buf, + self.scales_layout, + self.scheme, + ) + } +} + +#[cube] +#[allow(unused)] +fn create_quant_view_dynamic( + data_buf: GlobalInput, + data_layout: GlobalLayout, + scales_buf: GlobalInput, + scales_layout: GlobalScaleLayout, + #[comptime] scheme: QuantScheme, +) -> View<'static, Vector, BatchedCoords> { + intrinsic!(|scope| { + let func = CreateQuantView { + scope, + data_buf, + data_layout, + scales_buf, + scales_layout, + scheme, + _ty: PhantomData, + }; + run_with_quant_type(func, scheme) + }) +} + +#[cube] +fn create_quant_view( + data_buf: GlobalInput, + data_layout: GlobalLayout, + scales_buf: GlobalInput, + scales_layout: GlobalScaleLayout, + #[comptime] scheme: QuantScheme, +) -> View<'static, Vector, BatchedCoords> { + let size!(NQ) = N::value().comptime() / scheme.num_quants(); + + let data_view: View, BatchedCoords> = + View::new::(data_buf, data_layout); + let scales_view: View = + View::new::(scales_buf, scales_layout); + QuantizedView::new(data_view, scales_view, scheme).view() +} + +#[derive(CubeType)] +pub struct FusedMatmulState { + inputs: GlobalArgs, + outputs: GlobalArgs, + locals: LocalArgs, + a_batch: VirtualLayout, + b_batch: VirtualLayout, + c_batch: ComptimeOption>, + out_batch: VirtualLayout, + #[cube(comptime)] + config: FuseBlockConfig, + #[cube(comptime)] + a: MatmulArg, + #[cube(comptime)] + b: MatmulArg, + #[cube(comptime)] + c: Option, + #[cube(comptime)] + out: FuseArg, + #[cube(comptime)] + lhs_layout_config: GlobalLayoutConfig, + #[cube(comptime)] + rhs_layout_config: GlobalLayoutConfig, + #[cube(comptime)] + out_layout_config: GlobalLayoutConfig, + batch_shape: Sequence>, +} + +#[cube] +impl FusedMatmulState { + #[allow(clippy::too_many_arguments)] + pub fn new( + inputs: &FusedMatmulInput, + outputs: &mut GlobalArgs, + locals: &mut LocalArgs, + a_batch: VirtualLayout, + b_batch: VirtualLayout, + c_batch: ComptimeOption>, + out_batch: VirtualLayout, + batch_shape: Sequence>, + #[comptime] config: &FuseBlockConfig, + #[comptime] lhs_layout_config: GlobalLayoutConfig, + #[comptime] rhs_layout_config: GlobalLayoutConfig, + #[comptime] out_layout_config: GlobalLayoutConfig, + ) -> FusedMatmulState { + FusedMatmulState { + inputs: inputs.global.clone(), + outputs: outputs.clone(), + config: comptime![config.clone()], + locals: locals.clone(), + a_batch, + b_batch, + c_batch, + out_batch, + a: comptime![inputs.a.clone()], + b: comptime![inputs.b.clone()], + c: comptime![inputs.c.clone()], + out: comptime![inputs.out.clone()], + lhs_layout_config, + rhs_layout_config, + out_layout_config, + batch_shape, + } + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq, Serialize, Deserialize, PartialOrd, Ord)] +/// Argument to a matmul operation. +pub enum MatmulArg { + Normal(FuseArg), + Quantized { + data: FuseArg, + scales: FuseArg, + precision: FuseType, + scheme: QuantScheme, + }, +} + +impl MatmulArg { + pub fn data(&self) -> &FuseArg { + match self { + MatmulArg::Normal(arg) => arg, + MatmulArg::Quantized { data, .. } => data, + } + } + + pub fn scheme(&self) -> Option<&QuantScheme> { + match self { + MatmulArg::Normal(_) => None, + MatmulArg::Quantized { scheme, .. } => Some(scheme), + } + } + + pub fn precision(&self) -> FuseType { + match self { + MatmulArg::Normal(arg) => arg.precision(), + MatmulArg::Quantized { precision, .. } => *precision, + } + } +} diff --git a/crates/burn-cubecl-fusion/src/optim/matmul/fuser.rs b/crates/burn-cubecl-fusion/src/optim/matmul/fuser.rs new file mode 100644 index 0000000..42e4f3f --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/matmul/fuser.rs @@ -0,0 +1,154 @@ +use super::optimization::{FusedMatmul, MatmulOptimization}; +use crate::{ + engine::{fuser::TraceOperationFuser, settings::FuseSettings}, + optim::CubeOptimization, + optim::matmul::args::MatmulArg, +}; +use burn_fusion::{FuserStatus, OperationFuser}; +use burn_ir::{FloatOperationIr, OperationIr}; +use burn_std::DType; +use cubecl::Runtime; + +/// Fused element wise operations that are normally memory bound. +pub struct MatmulFuser { + fuser: TraceOperationFuser, + fuser_fallback: TraceOperationFuser, + device: R::Device, + matmul: Option, +} + +impl Clone for MatmulFuser { + fn clone(&self) -> Self { + Self { + fuser: self.fuser.clone(), + fuser_fallback: self.fuser_fallback.clone(), + device: self.device.clone(), + matmul: self.matmul.clone(), + } + } +} + +impl MatmulFuser { + pub fn new(device: R::Device) -> Self { + let client = R::client(&device); + let props = client.properties(); + let max_bindings = props.hardware.max_bindings; + let settings_matmul = FuseSettings { + output_shape_updates: false, + ..Default::default() + }; + let settings_fallback = FuseSettings::default(); + + Self { + fuser: TraceOperationFuser::new(max_bindings, settings_matmul), + fuser_fallback: TraceOperationFuser::new(max_bindings, settings_fallback), + device, + matmul: None, + } + } +} + +impl OperationFuser> for MatmulFuser { + fn fuse(&mut self, operation: &OperationIr) { + if let FuserStatus::Closed = self.fuser.status() { + return; + } + + if self.matmul.is_none() { + if let OperationIr::Float(_, FloatOperationIr::Matmul(op)) = operation { + // Precision shouldn't be hardcoded but I don't know how to get float precision of the backend + let lhs = match op.lhs.dtype { + DType::QFloat(scheme) => { + let (data, scales) = self.fuser.input_quantized_unhandled(&op.lhs).unwrap(); + MatmulArg::Quantized { + data, + scales, + precision: op.out.dtype.into(), + scheme, + } + } + _ => MatmulArg::Normal(self.fuser.input_unhandled(&op.lhs)), + }; + let rhs = match op.rhs.dtype { + DType::QFloat(scheme) => { + let (data, scales) = self.fuser.input_quantized_unhandled(&op.rhs).unwrap(); + MatmulArg::Quantized { + data, + scales, + precision: op.out.dtype.into(), + scheme, + } + } + _ => MatmulArg::Normal(self.fuser.input_unhandled(&op.rhs)), + }; + + let out = self.fuser.output_unhandled(&op.out); + + self.matmul = Some(FusedMatmul::new( + lhs, + rhs, + out, + op.clone().into(), + Default::default(), + )); + } else { + self.fuser.close(); + self.fuser_fallback.close(); + } + } else { + let can_register = + self.fuser.can_fuse(operation) && self.fuser_fallback.can_fuse(operation); + + match can_register { + true => { + self.fuser.fuse(operation); + self.fuser_fallback.fuse(operation); + } + false => { + self.fuser.close(); + self.fuser_fallback.close(); + } + }; + } + } + + fn finish(&mut self) -> CubeOptimization { + let client = R::client(&self.device); + let trace = self.fuser.finish(); + let trace_fallback = self.fuser_fallback.finish(); + + let matmul = MatmulOptimization::new( + trace, + trace_fallback, + client, + self.device.clone(), + self.len(), + self.matmul.as_ref().unwrap().clone(), + ); + + CubeOptimization::Matmul(matmul) + } + + fn reset(&mut self) { + self.fuser.reset(); + self.fuser_fallback.reset(); + self.matmul = None; + } + + fn status(&self) -> burn_fusion::FuserStatus { + self.fuser.status() + } + + fn properties(&self) -> burn_fusion::FuserProperties { + self.fuser.properties() + } + + fn len(&self) -> usize { + // Matmul operation isn't registered in the fuser + self.fuser.len() + 1 + } + + fn clone_dyn(&self) -> Box>> { + Box::new(self.clone()) + } +} diff --git a/crates/burn-cubecl-fusion/src/optim/matmul/mod.rs b/crates/burn-cubecl-fusion/src/optim/matmul/mod.rs new file mode 100644 index 0000000..75a082a --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/matmul/mod.rs @@ -0,0 +1,8 @@ +mod fuser; +mod optimization; + +pub(crate) mod args; +pub(crate) mod tune; + +pub use fuser::*; +pub use optimization::*; diff --git a/crates/burn-cubecl-fusion/src/optim/matmul/optimization.rs b/crates/burn-cubecl-fusion/src/optim/matmul/optimization.rs new file mode 100644 index 0000000..31dd82a --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/matmul/optimization.rs @@ -0,0 +1,708 @@ +use super::args::FusedMatmulInputLaunch; +#[cfg(feature = "autotune")] +use super::tune::fused_matmul_autotune; +use crate::{ + CubeFusionHandle, FallbackOperation, + engine::{ + codegen::ir::{FuseArg, FuseBlockConfig, FuseType, GlobalArgsLaunch, RefLayout}, + launch::{ + FuseTraceLauncher, HandleInput, LaunchPlan, + runner::{TraceRunner, Vectorization, VectorizationAxis}, + }, + trace::{FuseTrace, TraceError, TuneOutput}, + }, + optim::{ + elemwise::ElemwiseRunner, + matmul::args::{FusedMatmulArgs, MatmulArg}, + }, +}; +use burn_fusion::stream::Context; +use burn_ir::BinaryOpIr; +use cubecl::{ + client::ComputeClient, + prelude::*, + std::tensor::{MatrixBatchLayout, matrix_batch_layout}, +}; +use cubek::{ + matmul::{ + components::tile::TileMatmulKind, + definition::{ + MatmulElems, MatmulGlobalElems, MatmulProblem, MatmulSetupError, MatmulVectorSizes, + }, + routines::{ + BatchMatmulRoutine, BlueprintStrategy, + batch::{ + double_buffering::{CyclicDoubleBufferingAlgorithm, DoubleBufferingArgs}, + double_unit::DoubleUnitAlgorithm, + gemv_innerproduct::{ + DoubleVecMatInnerProductAlgorithm, VecMatInnerProductAlgorithm, + }, + ordered_double_buffering::{OrderedDoubleBufferingAlgorithm, OrderedSelectionArgs}, + simple::{SimpleAlgorithm, SimpleArgs}, + simple_unit::SimpleUnitAlgorithm, + }, + gemm::GemmRoutine, + gemv_unit_perpendicular::GemvUnitPerpendicularRoutine, + }, + strategy::launch_kernel_virtual, + }, + std::MatrixLayout, +}; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; + +/// Fuse matmul operation followed by elemwise operations into a single kernel. +pub struct MatmulOptimization { + pub(crate) info: Arc>, +} + +pub struct MatmulOptimizationTuneArg { + pub(crate) info: Arc>, + pub(crate) fallback: Box>, +} + +pub(crate) struct MatmulOptimizationInfo { + trace: FuseTrace, + trace_fallback: FuseTrace, + pub(crate) client: ComputeClient, + pub(crate) device: R::Device, + pub(crate) len: usize, + pub(crate) matmul: FusedMatmul, +} + +#[derive(Serialize, Deserialize, Debug)] +/// State for the [matrix optimization](MatmulOptimizationState). +pub struct MatmulOptimizationState { + trace: FuseTrace, + trace_fallback: FuseTrace, + matmul: FusedMatmul, + len: usize, +} + +impl MatmulOptimizationInfo { + /// Returns the number of output buffers added by fusion. + pub fn num_output_buffers(&self) -> usize { + self.trace_fallback.resources.outputs.len() + } + + /// Number of operations fused. + pub fn num_ops_fused(&self) -> usize { + self.len + } +} + +impl MatmulOptimizationTuneArg { + pub(crate) fn execute_fused( + &self, + context: &mut Context>, + selector: FusedMatmulSelector, + ) -> Result, TraceError> { + let launch = FusedMatmulLaunch::new(&self.info.matmul, selector); + let launcher = FuseTraceLauncher::new(&self.info.trace, &launch); + + launcher.launch(&self.info.client, &self.info.device, context) + } + + pub fn execute_fallback(&self, context: &mut Context>) -> TuneOutput { + self.fallback.run(context); + + #[cfg(feature = "autotune-checks")] + let mut output = TuneOutput::Checked { + handles: Default::default(), + }; + #[cfg(not(feature = "autotune-checks"))] + let output = TuneOutput::UnChecked(core::marker::PhantomData); + + #[cfg(feature = "autotune-checks")] + if let TuneOutput::Checked { handles } = &mut output { + let out_desc = context.tensors.get(&self.info.matmul.op.out.id).unwrap(); + let handle_out = context + .handles + .get_handle(&out_desc.id, &burn_ir::TensorStatus::ReadOnly); + + handles.insert( + self.info.matmul.op.out.id, + (out_desc.shape.clone(), handle_out.clone()), + ); + } + + let launcher = FuseTraceLauncher::new(&self.info.trace_fallback, &ElemwiseRunner); + let output_write = launcher + .launch(&self.info.client, &self.info.device, context) + .unwrap(); + + output.merge(output_write) + } +} + +impl MatmulOptimization { + pub fn new( + trace: FuseTrace, + trace_fallback: FuseTrace, + client: ComputeClient, + device: R::Device, + len: usize, + matmul: FusedMatmul, + ) -> Self { + let info = MatmulOptimizationInfo { + trace, + trace_fallback, + client, + device, + len, + matmul, + }; + + Self { + info: Arc::new(info), + } + } + /// Execute the optimization. + pub fn execute( + &mut self, + context: &mut Context>, + fallback: impl FnOnce(usize) -> Box>, + ) { + // The index of the fallback matmul is always 0. + let fallback = fallback(0); + let arg = MatmulOptimizationTuneArg { + info: self.info.clone(), + fallback, + }; + + #[cfg(feature = "autotune")] + fused_matmul_autotune::(arg, context); + + #[cfg(not(feature = "autotune"))] + if arg + .execute_fused(context, FusedMatmulSelector::default()) + .is_err() + { + arg.execute_fallback(context); + } + } + + /// Number of operations fused. + pub fn num_ops_fused(&self) -> usize { + self.info.num_ops_fused() + } + + /// Create an optimization from its [state](MatmulOptimizationState). + pub fn from_state(device: &R::Device, state: MatmulOptimizationState) -> Self { + let info = MatmulOptimizationInfo { + trace: state.trace, + trace_fallback: state.trace_fallback, + len: state.len, + client: R::client(device), + device: device.clone(), + matmul: state.matmul.clone(), + }; + + Self { + info: Arc::new(info), + } + } + + /// Convert the optimization to its [state](MatmulOptimizationState). + pub fn to_state(&self) -> MatmulOptimizationState { + MatmulOptimizationState { + trace: self.info.trace.clone(), + trace_fallback: self.info.trace_fallback.clone(), + matmul: self.info.matmul.clone(), + len: self.info.len, + } + } +} + +#[derive(Clone, Copy, Serialize, Deserialize, Debug)] +pub enum FusedMatmulSelector { + Simple { + multi_rows: bool, + tile_matmul: AcceleratedTileKind, + }, + DoubleBuffering { + specialized: bool, + tile_matmul: AcceleratedTileKind, + }, + OrderedDoubleBuffering { + tile_matmul: AcceleratedTileKind, + }, + SimpleVecMat, + DoubleVecMat, + GemmNoStage, + GemvUnitPerpendicular, + SimpleUnit, + DoubleUnit, +} + +impl FusedMatmulSelector { + /// Not efficient, but only called once when initializing the tunables. + pub fn name(&self) -> String { + let name = match self { + FusedMatmulSelector::Simple { + multi_rows, + tile_matmul, + } => match multi_rows { + false => format!("simple_{tile_matmul:?}"), + true => format!("simple_multirows_{tile_matmul:?}"), + }, + FusedMatmulSelector::DoubleBuffering { + specialized, + tile_matmul, + } => match specialized { + false => format!("double_buffering_{tile_matmul:?}"), + true => format!("double_buffering_specialized_{tile_matmul:?}"), + }, + FusedMatmulSelector::OrderedDoubleBuffering { tile_matmul } => { + format!("double_buffering_ordered_{tile_matmul:?}").to_lowercase() + } + FusedMatmulSelector::SimpleVecMat => "simple_vec_mat".into(), + FusedMatmulSelector::DoubleVecMat => "double_buffering_vec_mat".into(), + FusedMatmulSelector::GemmNoStage => "gemm".into(), + FusedMatmulSelector::GemvUnitPerpendicular => "gemv_unit_perpendicular".into(), + FusedMatmulSelector::SimpleUnit => "simple_unit".into(), + FusedMatmulSelector::DoubleUnit => "double_buffering_unit".into(), + }; + + format!("fused_{name}") + } +} + +impl Default for FusedMatmulSelector { + fn default() -> Self { + FusedMatmulSelector::Simple { + multi_rows: false, + tile_matmul: AcceleratedTileKind::Cmma, + } + } +} + +#[derive(new, Clone, Serialize, Deserialize, Debug)] +pub struct FusedMatmul { + pub(crate) lhs: MatmulArg, + pub(crate) rhs: MatmulArg, + out: FuseArg, + pub(crate) op: BinaryOpIr, + pub(crate) selector: FusedMatmulSelector, +} + +#[derive(new)] +pub struct FusedMatmulLaunch<'a> { + pub(crate) matmul: &'a FusedMatmul, + pub(crate) selector: FusedMatmulSelector, +} + +#[derive(Debug)] +pub enum FusedMatmulError { + LaunchError(MatmulSetupError), + InvalidInput(&'static str), +} + +impl From for FusedMatmulError { + fn from(value: MatmulSetupError) -> Self { + Self::LaunchError(value) + } +} + +impl<'a, R: Runtime> Vectorization for FusedMatmulLaunch<'a> { + fn axis(&self, plan: &LaunchPlan<'_, R>) -> VectorizationAxis { + let lhs_id = self.matmul.op.lhs.id; + let rhs_id = self.matmul.op.rhs.id; + + let mut tensor_lhs = None; + let mut tensor_rhs = None; + + for input in plan.handle_inputs.iter() { + match input { + HandleInput::Normal(input) => { + if input.relative_id == lhs_id { + tensor_lhs = Some((input.global_ir.id, &input.handle.strides)); + } + if input.relative_id == rhs_id { + tensor_rhs = Some((input.global_ir.id, &input.handle.strides)); + } + } + HandleInput::QuantValues(input) => { + if input.relative_id == lhs_id { + tensor_lhs = Some((input.global_ir.id, &input.handle.strides)); + } + if input.relative_id == rhs_id { + tensor_rhs = Some((input.global_ir.id, &input.handle.strides)); + } + } + HandleInput::QuantParams(_) => {} + } + } + + let (lhs_id_global, lhs_strides) = tensor_lhs.unwrap(); + let (rhs_id_global, rhs_strides) = tensor_rhs.unwrap(); + + let mut axis = VectorizationAxis::default(); + + if let MatrixBatchLayout::MildlyPermuted { transposed, .. } = + matrix_batch_layout(lhs_strides, self.matmul.lhs.scheme()) + && transposed + { + axis.insert(lhs_id_global, lhs_strides.len() - 2); + } + + if let MatrixBatchLayout::MildlyPermuted { transposed, .. } = + matrix_batch_layout(rhs_strides, self.matmul.rhs.scheme()) + && transposed + { + axis.insert(rhs_id_global, rhs_strides.len() - 2); + } + + axis + } +} + +impl TraceRunner for FusedMatmulLaunch<'_> { + type Error = FusedMatmulError; + + fn run<'a>( + &'a self, + client: &'a ComputeClient, + inputs: GlobalArgsLaunch, + outputs: GlobalArgsLaunch, + configs: &'a [FuseBlockConfig], + ) -> Result<(), FusedMatmulError> { + let global_elems = MatmulGlobalElems { + lhs: self.matmul.lhs.precision().into_storage_type(), + rhs: self.matmul.rhs.precision().into_storage_type(), + out: self.matmul.out.precision().into_storage_type(), + }; + let dtypes = MatmulElems::from_globals(&global_elems); + self.matmul_fused(client, inputs, outputs, &configs[0], dtypes) + } +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +/// Which tile matmul to use for accelerated algorithms +pub enum AcceleratedTileKind { + #[default] + Cmma, + Mma, +} + +impl FusedMatmulLaunch<'_> { + fn matmul_fused<'a, R: Runtime>( + &'a self, + client: &'a ComputeClient, + inputs: GlobalArgsLaunch, + outputs: GlobalArgsLaunch, + config: &'a FuseBlockConfig, + dtypes: MatmulElems, + ) -> Result<(), FusedMatmulError> { + let lhs_shape = inputs.shape(self.matmul.lhs.data()); + let rhs_shape = inputs.shape(self.matmul.rhs.data()); + let out_shape = outputs.shape_ref(&config.ref_layout, config.rank); + + let lhs_strides = inputs.strides(self.matmul.lhs.data()); + let lhs_scheme = self.matmul.lhs.scheme(); + let rhs_strides = inputs.strides(self.matmul.rhs.data()); + let rhs_scheme = self.matmul.rhs.scheme(); + + if matrix_batch_layout(&lhs_strides, lhs_scheme) == MatrixBatchLayout::HighlyPermuted { + return Err(FusedMatmulError::InvalidInput( + "Lhs needs to be contiguous, but can't when fusing.", + )); + } + if matrix_batch_layout(&rhs_strides, rhs_scheme) == MatrixBatchLayout::HighlyPermuted { + return Err(FusedMatmulError::InvalidInput( + "Rhs needs to be contiguous, but can't when fusing.", + )); + } + + let mut vector_sizes = MatmulVectorSizes { + lhs: inputs.vector_size(self.matmul.lhs.data()), + rhs: inputs.vector_size(self.matmul.rhs.data()), + out: match &config.ref_layout { + RefLayout::Concrete(arg) => match arg { + FuseArg::Input(..) => inputs.vector_size(arg), + FuseArg::Output(..) => outputs.vector_size(arg), + _ => panic!("Invalid ref layout"), + }, + RefLayout::Virtual(_) => 1, + }, + }; + + let address_type = inputs + .required_address_type() + .max(outputs.required_address_type()); + + if vector_sizes.out == 1 && (vector_sizes.lhs > 1 || vector_sizes.rhs > 1) { + return Err(FusedMatmulError::InvalidInput( + "Output vector size of 1 removes the gain from fusion", + )); + } + + if let MatmulArg::Quantized { scheme, .. } = self.matmul.lhs { + vector_sizes.lhs *= scheme.num_quants(); + } + if let MatmulArg::Quantized { scheme, .. } = self.matmul.rhs { + vector_sizes.rhs *= scheme.num_quants(); + } + + let out_strides = MatrixLayout::RowMajor.to_strides(&out_shape); + let problem = MatmulProblem::from_shapes_and_strides( + lhs_shape, + rhs_shape, + out_shape, + lhs_strides, + rhs_strides, + out_strides, + dtypes.as_global_elems(), + address_type, + self.matmul.lhs.scheme(), + self.matmul.rhs.scheme(), + )?; + + match self.selector { + FusedMatmulSelector::Simple { + multi_rows, + tile_matmul, + } => { + let args = SimpleArgs { + multi_rows, + tile_matmul: match tile_matmul { + AcceleratedTileKind::Cmma => TileMatmulKind::Cmma, + AcceleratedTileKind::Mma => TileMatmulKind::Mma, + }, + }; + + match launch_inner_fix_dtype::( + client, + FusedMatmulInputLaunch::new( + inputs, + config.clone(), + self.matmul.lhs.clone(), + self.matmul.rhs.clone(), + None, + self.matmul.out.clone(), + ), + outputs, + problem, + vector_sizes, + &BlueprintStrategy::Inferred(args), + ) { + Ok(_) => Ok(()), + Err(err) => Err(FusedMatmulError::LaunchError(err)), + } + } + + FusedMatmulSelector::DoubleBuffering { + specialized, + tile_matmul, + } => { + let args = DoubleBufferingArgs { + specialized, + tile_matmul: match tile_matmul { + AcceleratedTileKind::Cmma => TileMatmulKind::Cmma, + AcceleratedTileKind::Mma => TileMatmulKind::Mma, + }, + }; + + match launch_inner_fix_dtype::( + client, + FusedMatmulInputLaunch::new( + inputs, + config.clone(), + self.matmul.lhs.clone(), + self.matmul.rhs.clone(), + None, + self.matmul.out.clone(), + ), + outputs, + problem, + vector_sizes, + &BlueprintStrategy::Inferred(args), + ) { + Ok(_) => Ok(()), + Err(err) => Err(FusedMatmulError::LaunchError(err)), + } + } + + FusedMatmulSelector::OrderedDoubleBuffering { tile_matmul } => { + let row_count = match self.matmul.lhs.precision() { + FuseType::F16 | FuseType::BF16 => 8, + _ => 4, + }; + + let args = OrderedSelectionArgs { + row_count: Some(row_count), + rows_per_plane: Some(2), + partition_k: Some(2), + tile_matmul: match tile_matmul { + AcceleratedTileKind::Cmma => TileMatmulKind::Cmma, + AcceleratedTileKind::Mma => TileMatmulKind::Mma, + }, + }; + + match launch_inner_fix_dtype::( + client, + FusedMatmulInputLaunch::new( + inputs, + config.clone(), + self.matmul.lhs.clone(), + self.matmul.rhs.clone(), + None, + self.matmul.out.clone(), + ), + outputs, + problem, + vector_sizes, + &BlueprintStrategy::Inferred(args), + ) { + Ok(_) => Ok(()), + Err(err) => Err(FusedMatmulError::LaunchError(err)), + } + } + + FusedMatmulSelector::SimpleUnit => { + match launch_inner_fix_dtype::( + client, + FusedMatmulInputLaunch::new( + inputs, + config.clone(), + self.matmul.lhs.clone(), + self.matmul.rhs.clone(), + None, + self.matmul.out.clone(), + ), + outputs, + problem, + vector_sizes, + &Default::default(), + ) { + Ok(_) => Ok(()), + Err(err) => Err(FusedMatmulError::LaunchError(err)), + } + } + + FusedMatmulSelector::DoubleUnit => { + match launch_inner_fix_dtype::( + client, + FusedMatmulInputLaunch::new( + inputs, + config.clone(), + self.matmul.lhs.clone(), + self.matmul.rhs.clone(), + None, + self.matmul.out.clone(), + ), + outputs, + problem, + vector_sizes, + &Default::default(), + ) { + Ok(_) => Ok(()), + Err(err) => Err(FusedMatmulError::LaunchError(err)), + } + } + + FusedMatmulSelector::SimpleVecMat => { + match launch_inner_fix_dtype::( + client, + FusedMatmulInputLaunch::new( + inputs, + config.clone(), + self.matmul.lhs.clone(), + self.matmul.rhs.clone(), + None, + self.matmul.out.clone(), + ), + outputs, + problem, + vector_sizes, + &Default::default(), + ) { + Ok(_) => Ok(()), + Err(err) => Err(FusedMatmulError::LaunchError(err)), + } + } + + FusedMatmulSelector::DoubleVecMat => { + match launch_inner_fix_dtype::( + client, + FusedMatmulInputLaunch::new( + inputs, + config.clone(), + self.matmul.lhs.clone(), + self.matmul.rhs.clone(), + None, + self.matmul.out.clone(), + ), + outputs, + problem, + vector_sizes, + &Default::default(), + ) { + Ok(_) => Ok(()), + Err(err) => Err(FusedMatmulError::LaunchError(err)), + } + } + + FusedMatmulSelector::GemmNoStage => { + match launch_inner_fix_dtype::( + client, + FusedMatmulInputLaunch::new( + inputs, + config.clone(), + self.matmul.lhs.clone(), + self.matmul.rhs.clone(), + None, + self.matmul.out.clone(), + ), + outputs, + problem, + vector_sizes, + &Default::default(), + ) { + Ok(_) => Ok(()), + Err(err) => Err(FusedMatmulError::LaunchError(err)), + } + } + + FusedMatmulSelector::GemvUnitPerpendicular => { + match launch_inner_fix_dtype::( + client, + FusedMatmulInputLaunch::new( + inputs, + config.clone(), + self.matmul.lhs.clone(), + self.matmul.rhs.clone(), + None, + self.matmul.out.clone(), + ), + outputs, + problem, + vector_sizes, + &Default::default(), + ) { + Ok(_) => Ok(()), + Err(err) => Err(FusedMatmulError::LaunchError(err)), + } + } + } + } +} + +fn launch_inner_fix_dtype>( + client: &ComputeClient, + input: FusedMatmulInputLaunch, + output: GlobalArgsLaunch, + problem: MatmulProblem, + vector_sizes: MatmulVectorSizes, + blueprint_strategy: &BlueprintStrategy<(), A>, +) -> Result<(), MatmulSetupError> { + launch_kernel_virtual::( + client, + input, + output, + (), + problem, + vector_sizes, + blueprint_strategy, + ) +} diff --git a/crates/burn-cubecl-fusion/src/optim/matmul/tune.rs b/crates/burn-cubecl-fusion/src/optim/matmul/tune.rs new file mode 100644 index 0000000..fb35f99 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/matmul/tune.rs @@ -0,0 +1,299 @@ +use super::optimization::MatmulOptimizationTuneArg; +use crate::{ + CubeFusionHandle, + engine::trace::TuneOutput, + optim::matmul::{AcceleratedTileKind, FusedMatmulSelector}, + tune::{FusionInputGen, TuneInput}, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_fusion::stream::Context; +use cubecl::{ + AutotuneKey, CubeTuneId, Runtime, + std::tensor::MatrixBatchLayout, + tune::{LocalTuner, Tunable, TunableSet, TuneGroup, local_tuner}, +}; +use cubek::matmul::{ + definition::MatmulKind, + strategy::{MatmulAutotuneKey, MatmulGlobalScale, should_tune_double_buffering}, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Hash, Eq, PartialEq, Debug, Clone, Serialize, Deserialize, AutotuneKey)] +pub struct FusedMatmulAutotuneKey { + matmul_key: MatmulAutotuneKey, + #[autotune(anchor)] + num_out_buffers: usize, + #[autotune(anchor)] + num_ops: usize, +} + +/// Executes autotune on matmul operations +pub fn fused_matmul_autotune( + optimization: MatmulOptimizationTuneArg, + context: &mut Context>, +) { + static TUNER: LocalTuner = local_tuner!(); + + let tunables = TUNER.init(|| { + const PRIORITY_MAX: i8 = 3; + const PRIORITY_HIGH: i8 = 2; + const PRIORITY_MEDIUM: i8 = 1; + const PRIORITY_MIN: i8 = 0; + const PRIORITY_NEVER: i8 = -1; + + let accelerated = TuneGroup::::new("accelerated", |key| { + if matches!(key.matmul_key.analysis.kind, MatmulKind::General) { + match key.matmul_key.analysis.scale_global { + MatmulGlobalScale::Large => PRIORITY_MAX, + _ => PRIORITY_HIGH, + } + } else if matches!( + key.matmul_key.analysis.kind, + MatmulKind::MatVec | MatmulKind::VecMat + ) { + PRIORITY_MAX + } else { + PRIORITY_MEDIUM + } + }); + + let unit = TuneGroup::::new("unit", |key| { + if !matches!(key.matmul_key.analysis.kind, MatmulKind::General) + || matches!( + key.matmul_key.analysis.scale_global, + MatmulGlobalScale::Small + ) + { + PRIORITY_HIGH + } else { + PRIORITY_MEDIUM + } + }); + + let gemv = TuneGroup::::new("gemv", |key| { + if matches!(key.matmul_key.analysis.kind, MatmulKind::MatVec) { + // LHS is the matrix. + match key.matmul_key.definition.matrix_layout_lhs { + MatrixBatchLayout::Contiguous => PRIORITY_MAX, + MatrixBatchLayout::MildlyPermuted { transposed, .. } => { + if transposed { + PRIORITY_HIGH + } else { + PRIORITY_MAX + } + } + MatrixBatchLayout::HighlyPermuted => PRIORITY_MAX, + } + } else if matches!(key.matmul_key.analysis.kind, MatmulKind::VecMat) { + // RHS is the matrix. + match key.matmul_key.definition.matrix_layout_rhs { + MatrixBatchLayout::Contiguous => PRIORITY_HIGH, + MatrixBatchLayout::MildlyPermuted { transposed, .. } => { + if transposed { + PRIORITY_MAX + } else { + PRIORITY_HIGH + } + } + MatrixBatchLayout::HighlyPermuted => PRIORITY_HIGH, + } + } else { + PRIORITY_NEVER + } + }); + + let odd = TuneGroup::::new("odd", |key| { + if key.matmul_key.definition.lhs_pow2_factor == 0 + || key.matmul_key.definition.rhs_pow2_factor == 0 + { + PRIORITY_MAX + } else { + PRIORITY_MIN + } + }); + + fn double_buffering_priority(key: &FusedMatmulAutotuneKey, max: i8, min: i8) -> i8 { + if should_tune_double_buffering(key.num_out_buffers > 1, &key.matmul_key) { + max + } else { + min + } + } + + // First entry should always work, since it is considered the fallback. + let mut set = TunableSet::new(create_key::, FusionInputGen).with( + Tunable::new("fused_matmul_fallback", tune_fallback::).group(&unit, |key| { + if matches!(key.matmul_key.analysis.kind, MatmulKind::InnerProduct) { + PRIORITY_MAX + } else if matches!( + key.matmul_key.analysis.scale_global, + MatmulGlobalScale::Small + ) { + PRIORITY_HIGH + } else { + PRIORITY_MIN + } + }), + ); + + // Vector-matrix kernels. + for (selector, double_buf) in [ + (FusedMatmulSelector::SimpleVecMat, false), + (FusedMatmulSelector::DoubleVecMat, true), + (FusedMatmulSelector::GemmNoStage, false), + (FusedMatmulSelector::GemvUnitPerpendicular, false), + ] { + set = set.with( + Tunable::new(&selector.name(), move |input| { + tune_fused::(input, selector) + }) + .group(&gemv, move |key| match double_buf { + false => PRIORITY_MAX, + true => double_buffering_priority(key, PRIORITY_MAX, PRIORITY_HIGH), + }), + ); + } + + // Unit matmuls + for (selector, double_buf) in [ + (FusedMatmulSelector::SimpleUnit, false), + (FusedMatmulSelector::DoubleUnit, true), + ] { + set = set.with( + Tunable::new(&selector.name(), move |input| { + tune_fused::(input, selector) + }) + .group(&unit, move |key| match double_buf { + false => PRIORITY_MAX, + true => double_buffering_priority(key, PRIORITY_MAX, PRIORITY_HIGH), + }), + ); + } + + // Accelerated matmuls + for tile_matmul in [AcceleratedTileKind::Cmma, AcceleratedTileKind::Mma] { + for (selector, double_buf, extra_group) in [ + ( + FusedMatmulSelector::Simple { + multi_rows: false, + tile_matmul, + }, + false, + None, + ), + ( + FusedMatmulSelector::Simple { + multi_rows: true, + tile_matmul, + }, + false, + None, + ), + ( + FusedMatmulSelector::OrderedDoubleBuffering { tile_matmul }, + true, + None, + ), + ( + FusedMatmulSelector::DoubleBuffering { + specialized: false, + tile_matmul, + }, + true, + None, + ), + ( + FusedMatmulSelector::DoubleBuffering { + specialized: true, + tile_matmul, + }, + true, + Some(&odd), + ), + ] { + let priority_within_group = + |key: &FusedMatmulAutotuneKey, double_buf: bool| match double_buf { + false => PRIORITY_MAX, + true => double_buffering_priority(key, PRIORITY_MAX, PRIORITY_HIGH), + }; + let mut tunable = Tunable::new(&selector.name(), move |input| { + tune_fused::(input, selector) + }) + .group(&accelerated, move |key| { + priority_within_group(key, double_buf) + }); + + if let Some(group) = extra_group { + tunable = + tunable.group(group, move |key| priority_within_group(key, double_buf)); + } + set = set.with(tunable); + } + } + + set + }); + + TUNER.execute( + &CubeTuneId::new(&optimization.info.client, &optimization.info.device), + &optimization.info.client.clone(), + tunables, + TuneInput::new(context, optimization), + ); +} + +pub(crate) fn create_key( + input: &TuneInput>, +) -> FusedMatmulAutotuneKey { + let opt = input.optimization(); + assert!(input.is_original(), "Not supported when generating key"); + let tensors = input.tensors(); + let handles = input.handles(); + + let lhs = tensors.get(&opt.info.matmul.op.lhs.id).unwrap(); + let rhs = tensors.get(&opt.info.matmul.op.rhs.id).unwrap(); + let out = tensors.get(&opt.info.matmul.op.out.id).unwrap(); + + let lhs_strides = handles + .get_handle_ref(&lhs.id) + .expect("lhs handle") + .strides + .clone(); + let rhs_strides = handles + .get_handle_ref(&rhs.id) + .expect("rhs handle") + .strides + .clone(); + + let key = MatmulAutotuneKey::generate( + &opt.info.client, + &lhs.shape, + &rhs.shape, + &lhs_strides, + &rhs_strides, + dtype_to_storage_type(lhs.dtype), + dtype_to_storage_type(rhs.dtype), + dtype_to_storage_type(out.dtype), + opt.info.matmul.lhs.scheme(), + opt.info.matmul.rhs.scheme(), + ); + FusedMatmulAutotuneKey::new(key, opt.info.num_output_buffers(), opt.info.num_ops_fused()) +} + +fn tune_fused( + input: TuneInput>, + selector: FusedMatmulSelector, +) -> Result, String> { + let is_original = input.is_original(); + input.execute(|ctx, opt| match opt.execute_fused(ctx, selector) { + Ok(out) => Ok(out), + Err(_) if is_original => Ok(opt.execute_fallback(ctx)), + Err(e) => Err(format!("{e:?}")), + }) +} + +fn tune_fallback( + input: TuneInput>, +) -> Result, String> { + Ok(input.execute(|ctx, opt| opt.execute_fallback(ctx))) +} diff --git a/crates/burn-cubecl-fusion/src/optim/mod.rs b/crates/burn-cubecl-fusion/src/optim/mod.rs new file mode 100644 index 0000000..ff3f4d8 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/mod.rs @@ -0,0 +1,8 @@ +pub mod elemwise; +pub mod matmul; +pub mod reduce; +pub mod reduce_broadcasted; + +mod base; + +pub use base::*; diff --git a/crates/burn-cubecl-fusion/src/optim/reduce/args.rs b/crates/burn-cubecl-fusion/src/optim/reduce/args.rs new file mode 100644 index 0000000..0ac5f79 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/reduce/args.rs @@ -0,0 +1,246 @@ +use crate::engine::codegen::{ + io::{ref_buffer_len, ref_len, ref_shape, ref_stride, ref_vector_size}, + ir::{FuseArg, FuseBlockConfig, GlobalArgs, GlobalArgsExpand, LocalArgs, LocalArgsExpand}, + kernel::{fuse_on_read, fuse_on_write, init_locals}, +}; +use cubecl::prelude::*; +use cubek::reduce::components::args::{ReduceArgs, ReduceDType}; + +#[derive(Clone)] +pub struct FusedReduceArgs; + +#[derive(CubeType, CubeLaunch)] +pub struct FusedReduceInput { + pub global: GlobalArgs, + #[cube(comptime)] + pub config: FuseBlockConfig, + #[cube(comptime)] + pub arg: FuseArg, +} + +#[derive(CubeType, CubeLaunch)] +pub struct FusedReduceOutput { + pub global: GlobalArgs, + #[cube(comptime)] + pub config: FuseBlockConfig, + #[cube(comptime)] + pub arg: FuseArg, +} + +#[derive(Clone)] +pub struct FusedReduceState { + inputs: GlobalArgs, + outputs: GlobalArgs, + locals_on_read: LocalArgs, + locals_on_write: LocalArgs, + config_on_read: FuseBlockConfig, + config_on_write: FuseBlockConfig, + // TODO: Should be a list when multiple blocks are there. + input: FuseArg, + out: FuseArg, +} + +#[derive(Clone)] +pub struct FusedReduceStateExpand { + inputs: GlobalArgsExpand, + outputs: GlobalArgsExpand, + locals_on_read: LocalArgsExpand, + locals_on_write: LocalArgsExpand, + config_on_read: FuseBlockConfig, + config_on_write: FuseBlockConfig, + input: FuseArg, + out: FuseArg, +} + +#[cube] +impl ReduceArgs for FusedReduceArgs { + type Input = FusedReduceInput; + type Output = FusedReduceOutput; + type State = FusedReduceState; + + fn init_state( + input: &Self::Input, + output: &mut Self::Output, + ) -> Self::State

{ + let mut locals_read = init_locals(&input.global, &mut output.global, &input.config); + let mut locals_write = init_locals(&input.global, &mut output.global, &output.config); + // TODO Add stuff from previous blocks to the local of each block. + FusedReduceState::new(input, output, &mut locals_read, &mut locals_write) + } + + fn read_input( + state: &Self::State

, + index: usize, + ) -> Vector { + let mut state = state.clone(); + let value = fuse_on_read::( + &state.inputs, + &mut state.outputs, + &mut state.locals_on_read, + index, + comptime! { + let mut sequence = Sequence::new(); + // TODO: Register local arguments from previous blocks. + sequence.push(state.input.clone()); + sequence + }, + &state.config_on_read, + ); + value[0] + } + + fn read_output( + _state: &Self::State

, + _index: usize, + ) -> Vector { + Vector::empty() + } + + fn write_output( + state: &mut Self::State

, + index: usize, + value: Vector, + ) { + let mut values = Registry::>::new(); + let mut args = comptime![Vec::::new()]; + + values.insert(comptime![state.out.clone()], value); + comptime![args.push(state.out.clone())]; + fuse_on_write( + &state.inputs, + &mut state.outputs, + &mut state.locals_on_write, + index, + values, + args, + &state.config_on_write, + ); + } + + fn len_input(state: &Self::State

) -> usize { + ref_len( + &state.inputs, + &state.outputs, + &state.locals_on_read, + &state.config_on_read, + ) + } + + fn len_output(state: &Self::State

) -> usize { + ref_len( + &state.inputs, + &state.outputs, + &state.locals_on_write, + &state.config_on_write, + ) + } + + fn buffer_len_input(state: &Self::State

) -> usize { + ref_buffer_len( + &state.inputs, + &state.outputs, + &state.locals_on_read, + &state.config_on_read, + ) + } + + fn buffer_len_output(state: &Self::State

) -> usize { + ref_buffer_len( + &state.inputs, + &state.outputs, + &state.locals_on_write, + &state.config_on_write, + ) + } + + fn rank_input(state: &Self::State

) -> usize { + state.config_on_read.rank.runtime() + } + + fn rank_output(state: &Self::State

) -> usize { + state.config_on_write.rank.runtime() + } + + fn shape_input(state: &Self::State

, dim: usize) -> usize { + ref_shape(&state.locals_on_read, dim) + } + + fn shape_output(state: &Self::State

, dim: usize) -> usize { + ref_shape(&state.locals_on_write, dim) + } + + fn stride_input(state: &Self::State

, dim: usize) -> usize { + ref_stride(&state.locals_on_read, dim) + } + + fn stride_output(state: &Self::State

, dim: usize) -> usize { + ref_stride(&state.locals_on_write, dim) + } + + fn vector_size_input(state: &Self::State

) -> comptime_type!(VectorSize) { + ref_vector_size(&state.locals_on_read) + } + + fn vector_size_output(state: &Self::State

) -> comptime_type!(VectorSize) { + ref_vector_size(&state.locals_on_write) + } +} + +#[cube] +impl FusedReduceState { + pub fn new( + inputs: &FusedReduceInput, + outputs: &mut FusedReduceOutput, + locals_on_read: &mut LocalArgs, + locals_on_write: &mut LocalArgs, + ) -> FusedReduceState { + FusedReduceState { + inputs: inputs.global.clone(), + outputs: outputs.global.clone(), + locals_on_read: locals_on_read.clone(), + locals_on_write: locals_on_write.clone(), + config_on_read: comptime![inputs.config.clone()], + config_on_write: comptime![outputs.config.clone()], + input: comptime![inputs.arg.clone()], + out: comptime![outputs.arg.clone()], + } + } +} + +impl CubeType for FusedReduceState { + type ExpandType = FusedReduceStateExpand; +} + +impl IntoExpand for FusedReduceStateExpand { + type Expand = Self; + + fn into_expand(self, _: &Scope) -> Self::Expand { + self + } +} + +impl ExpandTypeClone for FusedReduceStateExpand { + fn clone_unchecked(&self) -> Self { + self.clone() + } +} + +impl AsRefExpand for FusedReduceStateExpand { + fn __expand_ref_method(&self, _: &Scope) -> &Self { + self + } +} + +impl AsMutExpand for FusedReduceStateExpand { + fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self { + self + } +} + +impl IntoMut for FusedReduceStateExpand { + fn into_mut(self, _context: &Scope) -> Self { + self + } +} + +impl CubeDebug for FusedReduceStateExpand {} diff --git a/crates/burn-cubecl-fusion/src/optim/reduce/fuser.rs b/crates/burn-cubecl-fusion/src/optim/reduce/fuser.rs new file mode 100644 index 0000000..2aa0c59 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/reduce/fuser.rs @@ -0,0 +1,314 @@ +use super::{ + ReduceSettings, + optimization::{FusedReduce, ReduceInstruction, ReduceOptimization}, +}; +use crate::{ + engine::{ + codegen::ir::FuseType, + fuser::TraceOperationFuser, + settings::{FuseSettings, RefLayoutSetting, VectorizationSetting}, + }, + optim::CubeOptimization, +}; +use burn_fusion::{FuserStatus, OperationFuser}; +use burn_ir::{NumericOperationIr, OperationIr, ReduceDimOpIr}; +use burn_std::Shape; +use cubecl::Runtime; + +/// Fuses element wise operations around a reduce operation. +pub struct ReduceFuser { + pub(crate) fuser: TraceOperationFuser, + pub(crate) fuser_read_fallback: TraceOperationFuser, + fuser_write_fallback: TraceOperationFuser, + settings_write: FuseSettings, + pub(crate) device: R::Device, + pub(crate) reduce: Option, + settings: ReduceSettings, +} + +impl Clone for ReduceFuser { + fn clone(&self) -> Self { + Self { + fuser: self.fuser.clone(), + fuser_read_fallback: self.fuser_read_fallback.clone(), + fuser_write_fallback: self.fuser_write_fallback.clone(), + settings_write: self.settings_write, + device: self.device.clone(), + reduce: self.reduce.clone(), + settings: self.settings, + } + } +} + +#[derive(Debug)] +pub enum ReduceFuserInfo { + FusedReduce { shape_input_id: Shape, axis: usize }, + FusedElemwise { shape_id: Shape }, +} + +impl ReduceFuser { + pub fn new(device: R::Device, settings: ReduceSettings) -> Self { + let client = R::client(&device); + let props = client.properties(); + let max_bindings = props.hardware.max_bindings; + let settings_read = FuseSettings { + // Inplace would work, but not when we have a concrete output to write too. + inplace: true, + ref_layout: RefLayoutSetting::OnlyContiguous, + broadcast: false, + output_shape_updates: true, + vectorization: VectorizationSetting::Activated, + }; + let settings_write = FuseSettings { + inplace: false, + output_shape_updates: false, + vectorization: VectorizationSetting::SmallerOrEqualThanPreviousBlock { block_pos: 0 }, + broadcast: false, + ref_layout: RefLayoutSetting::OnlyContiguous, + }; + let settings_fallback = FuseSettings::default(); + + Self { + fuser: TraceOperationFuser::new(max_bindings, settings_read), + fuser_read_fallback: TraceOperationFuser::new(max_bindings, settings_fallback), + fuser_write_fallback: TraceOperationFuser::new(max_bindings, settings_fallback), + settings_write, + device, + reduce: None, + settings, + } + } + + pub fn reduce_info(&self) -> ReduceFuserInfo { + match &self.reduce { + Some(reduce) => { + let shape_input_id = reduce.op.input.shape.clone(); + let axis = reduce.axis; + + ReduceFuserInfo::FusedReduce { + shape_input_id, + axis, + } + } + None => { + let shape_id = self.fuser_read_fallback.current_output_shape.clone(); + ReduceFuserInfo::FusedElemwise { shape_id } + } + } + } + + fn on_reduce(&mut self, op: &ReduceDimOpIr, inst: ReduceInstruction) { + // TODO: Fix: we need to have fuse-on-read with an identity block. + // + // if self.fuser.num_ops == 0 && false { + // self.fuser.current_output_shape = op.input.shape.dims.clone(); + // } else if self.fuser.current_output_shape != op.input.shape.dims { + + if self.fuser.current_output_shape != op.input.shape { + self.fuser.close(); + self.fuser_read_fallback.close(); + return; + } + + let [input] = self + .fuser + .next_block([&op.input], self.settings_write, false); + + let output = self.fuser.output_unhandled(&op.out); + let axis = op.axis; + + let fuse_on_write_activated = match self.settings { + ReduceSettings::Always => true, + // We only activate fuse-on-write when the reduction isn't on the last dimension, otherwise + // vectorization is impossible. Only [VectorizationMode::Perpendicular] supports vectorization. + // + // We could still fuse some output operations, but it would probably lead to worse performance. + ReduceSettings::OnlyParallel => axis != op.input.shape.rank() - 1, + ReduceSettings::Never => false, + }; + + if !fuse_on_write_activated { + self.fuser.close(); + } + + let acc = match inst { + ReduceInstruction::Mean | ReduceInstruction::Prod | ReduceInstruction::Sum => { + match input.precision() { + FuseType::F16 | FuseType::BF16 => FuseType::F32, + FuseType::I16 | FuseType::I8 => FuseType::I32, + FuseType::U16 | FuseType::U8 => FuseType::U32, + _ => input.precision(), + } + } + _ => input.precision(), + }; + + self.reduce = Some(FusedReduce { + input, + output, + acc, + axis, + op: op.clone(), + use_planes: false, + shared: false, + inst, + }); + + self.fuser_read_fallback.close(); + } + + fn on_elemwise_read(&mut self, operation: &OperationIr) { + let can_register = + self.fuser.can_fuse(operation) && self.fuser_read_fallback.can_fuse(operation); + + match can_register { + true => { + self.fuser.fuse(operation); + self.fuser_read_fallback.fuse(operation); + } + false => { + self.fuser.close(); + self.fuser_read_fallback.close(); + } + }; + } + + fn on_elemwise_write(&mut self, operation: &OperationIr) { + let can_register = + self.fuser.can_fuse(operation) && self.fuser_write_fallback.can_fuse(operation); + + match can_register { + true => { + self.fuser.fuse(operation); + self.fuser_write_fallback.fuse(operation); + } + false => { + self.fuser.close(); + self.fuser_write_fallback.close(); + } + }; + } +} + +impl OperationFuser> for ReduceFuser { + fn fuse(&mut self, operation: &OperationIr) { + if let FuserStatus::Closed = self.fuser.status() { + return; + } + + if self.reduce.is_none() { + if let OperationIr::NumericFloat(_, op) = operation { + match op { + NumericOperationIr::SumDim(op) => { + self.on_reduce(op, ReduceInstruction::Sum); + } + NumericOperationIr::MeanDim(op) => { + self.on_reduce(op, ReduceInstruction::Mean); + } + NumericOperationIr::ProdDim(op) => { + self.on_reduce(op, ReduceInstruction::Prod); + } + NumericOperationIr::ArgMax(op) => { + self.on_reduce(op, ReduceInstruction::ArgMax); + } + NumericOperationIr::ArgMin(op) => { + self.on_reduce(op, ReduceInstruction::ArgMin); + } + NumericOperationIr::MinDim(op) => { + self.on_reduce(op, ReduceInstruction::Min); + } + NumericOperationIr::MaxDim(op) => { + self.on_reduce(op, ReduceInstruction::Max); + } + NumericOperationIr::MaxAbsDim(op) => { + self.on_reduce(op, ReduceInstruction::MaxAbs); + } + _ => { + self.on_elemwise_read(operation); + } + }; + } else if let OperationIr::NumericInt(_, op) = operation { + match op { + NumericOperationIr::SumDim(op) => { + self.on_reduce(op, ReduceInstruction::Sum); + } + NumericOperationIr::MeanDim(op) => { + self.on_reduce(op, ReduceInstruction::Mean); + } + NumericOperationIr::ProdDim(op) => { + self.on_reduce(op, ReduceInstruction::Prod); + } + NumericOperationIr::ArgMax(op) => { + self.on_reduce(op, ReduceInstruction::ArgMax); + } + NumericOperationIr::ArgMin(op) => { + self.on_reduce(op, ReduceInstruction::ArgMin); + } + NumericOperationIr::MinDim(op) => { + self.on_reduce(op, ReduceInstruction::Min); + } + NumericOperationIr::MaxDim(op) => { + self.on_reduce(op, ReduceInstruction::Max); + } + NumericOperationIr::MaxAbsDim(op) => { + self.on_reduce(op, ReduceInstruction::MaxAbs); + } + _ => { + self.on_elemwise_read(operation); + } + }; + } else { + self.on_elemwise_read(operation); + } + } else { + self.on_elemwise_write(operation); + } + } + + fn finish(&mut self) -> CubeOptimization { + let client = R::client(&self.device); + let trace = self.fuser.finish(); + let trace_read_fallback = self.fuser_read_fallback.finish(); + let trace_write_fallback = self.fuser_write_fallback.finish(); + let fuse_reduce = self.reduce.as_ref().unwrap(); + + let reduce = ReduceOptimization::new( + trace, + trace_read_fallback, + trace_write_fallback, + client, + self.device.clone(), + self.len(), + self.fuser_read_fallback.len(), + fuse_reduce.clone(), + self.settings, + ); + + CubeOptimization::Reduce(reduce) + } + + fn reset(&mut self) { + self.fuser.reset(); + self.fuser_read_fallback.reset(); + self.fuser_write_fallback.reset(); + self.reduce = None; + } + + fn status(&self) -> burn_fusion::FuserStatus { + self.fuser.status() + } + + fn properties(&self) -> burn_fusion::FuserProperties { + let mut properties = self.fuser.properties(); + properties.ready = self.reduce.is_some(); + properties + } + + fn len(&self) -> usize { + self.fuser.len() + if self.reduce.is_some() { 1 } else { 0 } + } + + fn clone_dyn(&self) -> Box>> { + Box::new(self.clone()) + } +} diff --git a/crates/burn-cubecl-fusion/src/optim/reduce/mod.rs b/crates/burn-cubecl-fusion/src/optim/reduce/mod.rs new file mode 100644 index 0000000..75a082a --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/reduce/mod.rs @@ -0,0 +1,8 @@ +mod fuser; +mod optimization; + +pub(crate) mod args; +pub(crate) mod tune; + +pub use fuser::*; +pub use optimization::*; diff --git a/crates/burn-cubecl-fusion/src/optim/reduce/optimization.rs b/crates/burn-cubecl-fusion/src/optim/reduce/optimization.rs new file mode 100644 index 0000000..3836bd5 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/reduce/optimization.rs @@ -0,0 +1,518 @@ +use super::args::{ + FusedReduceInput, FusedReduceInputLaunch, FusedReduceOutput, FusedReduceOutputLaunch, +}; +#[cfg(feature = "autotune")] +use super::tune::fused_reduce_autotune; +use crate::{ + CubeFusionHandle, FallbackOperation, + engine::{ + codegen::ir::{ + FuseArg, FuseBlockConfig, FuseType, GlobalArgsLaunch, RefLayout, + multi_block_variables_init, + }, + launch::{ + FuseTraceLauncher, + runner::{TraceRunner, Vectorization}, + }, + trace::{FuseTrace, TraceError, TuneOutput}, + }, + optim::{elemwise::ElemwiseRunner, reduce::args::FusedReduceArgs}, +}; +use burn_backend::cubecl::{dtype_to_storage_type, elem_type_to_dtype}; +use burn_fusion::stream::Context; +use burn_ir::ReduceDimOpIr; +use burn_std::DType; +use cubecl::{Runtime, client::ComputeClient, ir::StorageType, prelude::*}; +use cubek::reduce::{ + ReduceDtypes, ReduceError, VectorizationMode, + components::instructions::ReduceOperationConfig, + init_tensors, + launch::{RoutineStrategy, reduce_kernel_virtual}, + output_vectorization_axis, + routines::{ + ReduceBlueprint, ReduceLaunchSettings, ReduceProblem, ReduceVectorSettings, Routine, + cube::CubeRoutine, plane::PlaneRoutine, unit::UnitRoutine, + }, +}; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; + +#[cfg(not(feature = "autotune"))] +use cubek::reduce::routines::{BlueprintStrategy, unit::UnitStrategy}; + +pub struct ReduceOptimization { + pub(crate) info: Arc>, +} + +pub(crate) struct ReduceOptimizationInfo { + pub(crate) trace: FuseTrace, + trace_read_fallback: FuseTrace, + trace_write_fallback: FuseTrace, + pub(crate) client: ComputeClient, + pub(crate) device: R::Device, + pub(crate) len: usize, + pub(crate) len_read: usize, + pub(crate) reduce: FusedReduce, + settings: ReduceSettings, +} + +impl ReduceOptimizationInfo { + pub fn from_state(device: &R::Device, state: ReduceOptimizationState) -> Self { + let client = R::client(device); + + Self { + trace: state.trace, + trace_read_fallback: state.trace_read_fallback, + trace_write_fallback: state.trace_write_fallback, + client, + device: device.clone(), + len: state.len, + len_read: state.len_read, + reduce: state.reduce, + settings: state.settings, + } + } + pub fn to_state(&self) -> ReduceOptimizationState { + ReduceOptimizationState { + trace: self.trace.clone(), + trace_read_fallback: self.trace_read_fallback.clone(), + trace_write_fallback: self.trace_write_fallback.clone(), + len: self.len, + len_read: self.len_read, + reduce: self.reduce.clone(), + settings: self.settings, + } + } +} + +#[derive(Serialize, Deserialize, Copy, Clone)] +pub enum ReduceSettings { + Always, + /// We only activate fuse-on-write when the reduction isn't on the last dimension, otherwise + /// vectorization is impossible. Only [VectorizationMode::Perpendicular] supports vectorization. + /// + /// We could still fuse some output operations, but it would probably lead to worse performance. + OnlyParallel, + Never, +} + +pub(crate) struct ReduceOptimizationTuneArg { + pub(crate) info: Arc>, + pub(crate) fallback: Arc>>, +} + +impl Clone for ReduceOptimizationTuneArg { + fn clone(&self) -> Self { + Self { + info: self.info.clone(), + fallback: self.fallback.clone(), + } + } +} + +#[derive(Clone, Copy, Serialize, Deserialize, Debug)] +pub enum ReduceInstruction { + ArgMax, + ArgMin, + Mean, + Prod, + Sum, + Max, + Min, + MaxAbs, +} + +pub trait ReduceFallbackFn: Send + Sync { + fn run(&self, context: &mut Context>); +} + +#[derive(Serialize, Deserialize)] +pub struct ReduceOptimizationState { + pub(crate) trace: FuseTrace, + pub(crate) trace_read_fallback: FuseTrace, + pub(crate) trace_write_fallback: FuseTrace, + pub(crate) reduce: FusedReduce, + pub(crate) len: usize, + pub(crate) len_read: usize, + pub(crate) settings: ReduceSettings, +} + +impl core::fmt::Debug for ReduceOptimizationState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_fmt(format_args!( + "{{ len_read: {}, len_total: {} }}", + self.len_read, self.len + )) + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct FusedReduce { + pub(crate) input: FuseArg, + pub(crate) output: FuseArg, + pub(crate) acc: FuseType, + pub(crate) axis: usize, + pub(crate) op: ReduceDimOpIr, + pub(crate) use_planes: bool, + pub(crate) shared: bool, + pub(crate) inst: ReduceInstruction, +} + +#[derive(new)] +pub struct FusedReduceLaunch<'a> { + reduce: &'a FusedReduce, + strategy: RoutineStrategy, +} + +#[derive(Debug)] +pub enum FusedReduceError { + Reduce(ReduceError), + InvalidSelection(Box<&'static str>), + InvalidInput, +} + +impl From for FusedReduceError { + fn from(value: ReduceError) -> Self { + Self::Reduce(value) + } +} + +impl ReduceOptimizationTuneArg { + pub fn execute_fused( + &self, + context: &mut Context>, + strategy: RoutineStrategy, + ) -> Result, TraceError> { + let launch = FusedReduceLaunch::new(&self.info.reduce, strategy); + let launcher = FuseTraceLauncher::new(&self.info.trace, &launch); + launcher.launch(&self.info.client, &self.info.device, context) + } + + pub fn execute_fallback(&self, context: &mut Context>) -> TuneOutput { + let launcher = FuseTraceLauncher::new(&self.info.trace_read_fallback, &ElemwiseRunner); + + #[allow(unused_mut)] // It is used when `autotune-checks` is activated. + let mut output_read = launcher + .launch(&self.info.client, &self.info.device, context) + .unwrap(); + + self.fallback.run(context); + + #[cfg(feature = "autotune-checks")] + if let TuneOutput::Checked { handles } = &mut output_read { + let out_desc = context.tensors.get(&self.info.reduce.op.out.id).unwrap(); + let handle_out = context + .handles + .get_handle(&out_desc.id, &burn_ir::TensorStatus::ReadOnly); + + handles.insert( + self.info.reduce.op.out.id, + (out_desc.shape.clone(), handle_out.clone()), + ); + } + + let launcher = FuseTraceLauncher::new(&self.info.trace_write_fallback, &ElemwiseRunner); + + let output_write = launcher + .launch(&self.info.client, &self.info.device, context) + .unwrap(); + + output_read.merge(output_write) + } +} + +#[allow(clippy::too_many_arguments)] +impl ReduceOptimization { + pub fn new( + trace: FuseTrace, + trace_read_fallback: FuseTrace, + trace_write_fallback: FuseTrace, + client: ComputeClient, + device: R::Device, + len: usize, + len_read: usize, + reduce: FusedReduce, + settings: ReduceSettings, + ) -> Self { + let info = ReduceOptimizationInfo { + trace, + trace_read_fallback, + trace_write_fallback, + client, + device, + len, + len_read, + reduce, + settings, + }; + + Self { + info: Arc::new(info), + } + } + /// Execute the optimization. + pub fn execute( + &mut self, + context: &mut Context>, + fallback: impl FnOnce(usize) -> Box>, + ) { + // The index of the fallback reduce is the number of ops fused as read. + let fallback = fallback(self.info.len_read); + let arg = ReduceOptimizationTuneArg { + info: self.info.clone(), + fallback: Arc::new(fallback), + }; + + #[cfg(feature = "autotune")] + fused_reduce_autotune::(arg, context); + + #[cfg(not(feature = "autotune"))] + if arg + .execute_fused( + context, + RoutineStrategy::Unit(BlueprintStrategy::Inferred(UnitStrategy)), + ) + .is_err() + { + arg.execute_fallback(context); + } + } + + pub fn num_output_buffers(&self) -> usize { + self.info.trace_read_fallback.resources.outputs.len() + } + + pub fn to_state(&self) -> ReduceOptimizationState { + ReduceOptimizationState { + trace: self.info.trace.clone(), + trace_read_fallback: self.info.trace_read_fallback.clone(), + trace_write_fallback: self.info.trace_write_fallback.clone(), + reduce: self.info.reduce.clone(), + len: self.info.len, + len_read: self.info.len_read, + settings: self.info.settings, + } + } + + pub fn from_state(device: &R::Device, state: ReduceOptimizationState) -> Self { + let client = R::client(device); + + let info = ReduceOptimizationInfo { + trace: state.trace, + trace_read_fallback: state.trace_read_fallback, + trace_write_fallback: state.trace_write_fallback, + reduce: state.reduce, + len: state.len, + len_read: state.len_read, + client, + device: device.clone(), + settings: state.settings, + }; + + Self { + info: Arc::new(info), + } + } + + /// Returns the number of output buffers added by fusion. + pub fn num_ops_fused(&self) -> usize { + self.info.len + } +} + +// TODO: Implement better vectorization here. +impl Vectorization for FusedReduceLaunch<'_> {} + +impl TraceRunner for FusedReduceLaunch<'_> { + type Error = FusedReduceError; + + fn run<'a>( + &'a self, + client: &'a ComputeClient, + inputs: GlobalArgsLaunch, + outputs: GlobalArgsLaunch, + configs: &'a [FuseBlockConfig], + ) -> Result<(), FusedReduceError> { + let [config_read, config_write] = [&configs[0], &configs[1]]; + let shape = match &config_read.ref_layout { + RefLayout::Concrete(FuseArg::Output(..)) => { + outputs.shape_ref(&config_read.ref_layout, config_read.rank) + } + _ => inputs.shape_ref(&config_read.ref_layout, config_read.rank), + }; + let reduce_count: usize = shape + .iter() + .enumerate() + .map(|(i, s)| if i == self.reduce.axis { 1 } else { *s }) + .product(); + + let vectorization_mode = match self.reduce.axis == config_read.rank - 1 { + true => VectorizationMode::Parallel, + false => VectorizationMode::Perpendicular, + }; + let address_type = inputs + .required_address_type() + .max(outputs.required_address_type()); + + let settings = ReduceVectorSettings { + vectorization_mode, + vector_size_input: config_read.width, + vector_size_output: config_write.width, + }; + let problem = ReduceProblem { + reduce_len: shape[self.reduce.axis], + reduce_count, + axis: self.reduce.axis, + dtypes: ReduceDtypes { + input: dtype_to_storage_type(self.reduce.op.input.dtype), + output: dtype_to_storage_type(self.reduce.op.out.dtype), + accumulation: self.reduce.acc.into_elem().into(), + }, + address_type, + instruction: reduce_instruction2config(&self.reduce.inst), + }; + + let (blueprint, settings) = match self.strategy.clone() { + RoutineStrategy::Unit(strategy) => { + let routine = UnitRoutine; + routine.prepare(client, problem, settings, strategy)? + } + RoutineStrategy::Plane(strategy) => { + let routine = PlaneRoutine; + routine.prepare(client, problem, settings, strategy)? + } + RoutineStrategy::Cube(strategy) => { + let routine = CubeRoutine; + routine.prepare(client, problem, settings, strategy)? + } + }; + + let out_vec_axis = output_vectorization_axis( + &inputs.strides_ref(&config_read.ref_layout, config_read.rank), + self.reduce.axis, + vectorization_mode, + ); + + let kwargs = ReduceKwArgs { + client, + inputs, + outputs, + reduce_axis: self.reduce.axis, + out_vec_axis, + config_fuse_read: config_read.clone(), + config_fuse_write: config_write.clone(), + input: self.reduce.input.clone(), + output: self.reduce.output.clone(), + blueprint, + settings, + }; + let result = launch_reduce_mixed_precision( + kwargs, + self.reduce.inst, + self.reduce.op.input.dtype, + self.reduce.op.out.dtype, + elem_type_to_dtype(self.reduce.acc.into_elem()), + ); + + match result { + Ok(_) => Ok(()), + Err(err) => Err(FusedReduceError::Reduce(ReduceError::Launch(err))), + } + } +} + +struct ReduceKwArgs<'b, Run: Runtime> { + client: &'b ComputeClient, + inputs: GlobalArgsLaunch, + outputs: GlobalArgsLaunch, + reduce_axis: usize, + out_vec_axis: usize, + blueprint: ReduceBlueprint, + settings: ReduceLaunchSettings, + config_fuse_read: FuseBlockConfig, + config_fuse_write: FuseBlockConfig, + input: FuseArg, + output: FuseArg, +} + +fn launch_reduce_mixed_precision( + kwargs: ReduceKwArgs<'_, Run>, + instruction: ReduceInstruction, + dtype_input: DType, + dtype_output: DType, + dtype_acc: DType, +) -> Result<(), LaunchError> { + let config = reduce_instruction2config(&instruction); + launch_reduce::(kwargs, config, dtype_input, dtype_output, dtype_acc) +} + +pub(crate) fn reduce_instruction2config(instruction: &ReduceInstruction) -> ReduceOperationConfig { + match instruction { + ReduceInstruction::ArgMax => ReduceOperationConfig::ArgMax, + ReduceInstruction::ArgMin => ReduceOperationConfig::ArgMin, + ReduceInstruction::Prod => ReduceOperationConfig::Prod, + ReduceInstruction::Mean => ReduceOperationConfig::Mean, + ReduceInstruction::Sum => ReduceOperationConfig::Sum, + ReduceInstruction::Max => ReduceOperationConfig::Max, + ReduceInstruction::Min => ReduceOperationConfig::Min, + ReduceInstruction::MaxAbs => ReduceOperationConfig::MaxAbs, + } +} + +fn launch_reduce( + kwargs: ReduceKwArgs<'_, Run>, + inst: ReduceOperationConfig, + dtype_input: DType, + dtype_output: DType, + dtype_acc: DType, +) -> Result<(), LaunchError> { + unsafe { + reduce_kernel_fused::launch_unchecked::( + kwargs.client, + kwargs.settings.cube_count, + kwargs.settings.cube_dim, + kwargs.settings.address_type, + kwargs.config_fuse_read.width, + kwargs.config_fuse_write.width, + FusedReduceInputLaunch::new(kwargs.inputs, kwargs.config_fuse_read, kwargs.input), + FusedReduceOutputLaunch::new(kwargs.outputs, kwargs.config_fuse_write, kwargs.output), + kwargs.reduce_axis, + kwargs.out_vec_axis, + kwargs.blueprint, + inst, + dtype_to_storage_type(dtype_input), + dtype_to_storage_type(dtype_output), + dtype_to_storage_type(dtype_acc), + ) + }; + + Ok(()) +} + +#[cube(launch_unchecked, address_type = "dynamic")] +pub fn reduce_kernel_fused( + input: &FusedReduceInput, + output: &mut FusedReduceOutput, + reduce_axis: usize, + out_vec_axis: usize, + #[comptime] blueprint: ReduceBlueprint, + #[comptime] config: ReduceOperationConfig, + #[define(In)] _input_dtype: StorageType, + #[define(Out)] _output_dtype: StorageType, + #[define(Acc)] _acc_dtype: StorageType, +) { + multi_block_variables_init(&input.config, &mut output.global.variables); + multi_block_variables_init(&output.config, &mut output.global.variables); + + let (input, mut output) = + init_tensors::(input, output); + + reduce_kernel_virtual::( + &input, + &mut output, + reduce_axis, + out_vec_axis, + blueprint, + config, + ); +} diff --git a/crates/burn-cubecl-fusion/src/optim/reduce/tune.rs b/crates/burn-cubecl-fusion/src/optim/reduce/tune.rs new file mode 100644 index 0000000..5500156 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/reduce/tune.rs @@ -0,0 +1,173 @@ +use super::optimization::ReduceOptimizationTuneArg; +use crate::{ + CubeFusionHandle, + engine::trace::TuneOutput, + tune::{FusionInputGen, TuneInput}, +}; +use burn_backend::cubecl::dtype_to_elem_type; +use burn_fusion::stream::Context; +use cubecl::{ + AutotuneKey, CubeTuneId, Runtime, + tune::{LocalTuner, Tunable, TunableSet, TuneGroup, local_tuner}, +}; +use cubek::reduce::{ + launch::{RoutineStrategy, tune_key::ReduceAutotuneKey}, + routines::{BlueprintStrategy, cube::CubeStrategy, plane::PlaneStrategy, unit::UnitStrategy}, +}; +use serde::{Deserialize, Serialize}; + +/// Autotune key for standard fused reduction operations. +/// +/// Records metadata about the fusion graph (IO and ops) alongside +/// the core reduction parameters to ensure stable kernel selection. +#[derive(Hash, Eq, PartialEq, Debug, Clone, Serialize, Deserialize, AutotuneKey)] +pub struct FusedReduceAutotuneKey { + reduce_key: ReduceAutotuneKey, + #[autotune(anchor)] + fuse_num_reads: usize, + #[autotune(anchor)] + fuse_num_writes: usize, + #[autotune(anchor)] + fuse_num_ops: usize, +} + +/// Executes autotuning for fused reduction operations. +/// +/// This tuner evaluates different hardware-specific strategies (Plane, Cube, Unit) +/// and assigns priorities based on the `vector_count` of the reduction. +pub fn fused_reduce_autotune( + arg: ReduceOptimizationTuneArg, + context: &mut Context>, +) { + static TUNER: LocalTuner = local_tuner!(); + + let tunables = TUNER.init(|| { + const PRIORITY_MAX: i8 = 2; + const PRIORITY_MIN: i8 = 1; + + let mut set = TunableSet::new(create_key::, FusionInputGen); + let group = TuneGroup::::new("fused_reduce", |_key| PRIORITY_MAX); + + // Fallback implementation for robustness. + set = set.with(Tunable::new("fused_reduce_fallback", tune_fallback::)); + + // Define properties to categorize hardware strategies. + enum ReduceProps { + GreatWithLowReduceCount, + GreatWithHighReduceCount, + Balanced, + } + + let strategies = [ + ( + "fused_unit", + RoutineStrategy::Unit(BlueprintStrategy::Inferred(UnitStrategy)), + ReduceProps::GreatWithHighReduceCount, + ), + ( + "fused_plane", + RoutineStrategy::Plane(BlueprintStrategy::Inferred(PlaneStrategy { + independent: true, + })), + ReduceProps::Balanced, + ), + ( + "fused_cube", + RoutineStrategy::Cube(BlueprintStrategy::Inferred(CubeStrategy { + // Two steps reduction doesn't work with fuse-on-write, we can't activate plane + // when using the cube algo. + use_planes: false, + })), + ReduceProps::GreatWithLowReduceCount, + ), + ]; + + for (name, strategy, props) in strategies { + let tunable = Tunable::new(name, move |input| tune_reduce::(input, &strategy)) + .group(&group, move |key| match props { + ReduceProps::GreatWithLowReduceCount => { + if key.reduce_key.vector_count < 128 { + PRIORITY_MAX + } else { + PRIORITY_MIN + } + } + ReduceProps::GreatWithHighReduceCount => { + if key.reduce_key.vector_count > 64 { + PRIORITY_MAX + } else { + PRIORITY_MIN + } + } + ReduceProps::Balanced => PRIORITY_MAX, + }); + + set = set.with(tunable); + } + + set + }); + + TUNER.execute( + &CubeTuneId::new(&arg.info.client, &arg.info.device), + &arg.info.client.clone(), + tunables, + TuneInput::new(context, arg), + ); +} + +/// Creates the autotune key by extracting tensor metadata and fusion block statistics. +pub(crate) fn create_key( + input: &TuneInput>, +) -> FusedReduceAutotuneKey { + let opt = input.optimization(); + assert!( + input.is_original(), + "Forked context not supported for key generation" + ); + let tensors = input.tensors(); + + let input_tensor = tensors.get(&opt.info.reduce.op.input.id).unwrap(); + let out_tensor = tensors.get(&opt.info.reduce.op.out.id).unwrap(); + let acc = opt.info.reduce.acc.into_elem(); + + let key = ReduceAutotuneKey::generate( + dtype_to_elem_type(input_tensor.dtype), + dtype_to_elem_type(out_tensor.dtype), + acc, + &input_tensor.shape, + opt.info.reduce.axis == input_tensor.shape.rank() - 1, + opt.info.reduce.axis, + ); + + // Assume the fusion contains at least a read and a write block. + let read_block = &opt.info.trace.blocks[0]; + let write_block = &opt.info.trace.blocks[1]; + + FusedReduceAutotuneKey::new( + key, + read_block.reads.len() + write_block.reads.len(), + read_block.writes.len() + write_block.writes.len(), + read_block.ops.len() + write_block.ops.len(), + ) +} + +/// Executes a fused reduction optimization. +fn tune_reduce( + input: TuneInput>, + strategy: &RoutineStrategy, +) -> Result, String> { + input + .execute(|ctx, opt| opt.execute_fused(ctx, strategy.clone())) + .map_err(|e| format!("{e:?}")) +} + +/// Executes the fallback path for a reduction optimization. +fn tune_fallback( + input: TuneInput>, +) -> Result, String> { + input.execute(|ctx, opt| { + opt.execute_fallback(ctx); + }); + Ok(TuneOutput::UnChecked(std::marker::PhantomData)) +} diff --git a/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/fuser/base.rs b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/fuser/base.rs new file mode 100644 index 0000000..9fb906b --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/fuser/base.rs @@ -0,0 +1,388 @@ +use crate::optim::{ + CubeOptimization, + reduce::{ReduceFuser, ReduceFuserInfo, ReduceSettings}, + reduce_broadcasted::{ + ReduceBlockOptimInfo, ReduceBroadcastedOptimization, ReduceBroadcastedOptimizationInfo, + fuser::{ + block::{ReduceBlockFuser, ReduceBlockFusionAnalysis, ReduceBroadcastedStatus}, + full::ReduceBroadcastedFullFuser, + full_analyzer::FullFuserAnalyzer, + }, + }, +}; +use burn_fusion::{FuserProperties, FuserStatus, OperationFuser}; +use burn_ir::OperationIr; +use cubecl::Runtime; +use std::sync::Arc; + +/// Fuses element wise operations around a reduce operation. +pub struct ReduceBroadcastedFuser { + blocks: Vec>, + fuser_default: ReduceFuser, + num_ops: usize, + state: ReduceBroadcastedStatus, + max_bindings: u32, +} + +impl Clone for ReduceBroadcastedFuser { + fn clone(&self) -> Self { + Self { + blocks: self.blocks.clone(), + fuser_default: self.fuser_default.clone(), + num_ops: self.num_ops, + state: self.state.clone(), + max_bindings: self.max_bindings, + } + } +} + +impl ReduceBroadcastedFuser { + pub fn new(device: R::Device) -> Self { + let fuser = ReduceFuser::new(device, ReduceSettings::Always); + let max_bindings = fuser.fuser.max_bindings; + let block = ReduceBlockFuser::new(fuser.clone()); + + Self { + blocks: vec![block], + fuser_default: fuser, + num_ops: 0, + state: ReduceBroadcastedStatus::Starting, + max_bindings, + } + } + + /// Checks whether the full fuser and all fallback blocks together account for every + /// operation that has been registered across all blocks. + /// + /// This is a dry-run consistency check: it simulates finishing all blocks without + /// mutating any state, then verifies two invariants: + /// + /// 1. **Full fuser coverage** — the number of operations absorbed by the + /// [`ReduceBroadcastedFullFuser`] matches the total operation count. + /// 2. **Fallback coverage** — the sum of operations across all fallback + /// [`ReduceBlockOptimInfo`] entries also matches the total operation count. + /// + /// If either invariant fails, the fuser's state should be marked as + /// [`ReduceBroadcastedStatus::Abort`] because the optimization would produce + /// incorrect results. + fn is_consistent(&self) -> bool { + let analyzer = FullFuserAnalyzer::new(&self.blocks); + let mut full = ReduceBroadcastedFullFuser::new(self.max_bindings, analyzer); + let mut num_ops = 0; + let fallbacks = self + .blocks + .clone() + .iter_mut() + .map(|block| block.finish(&mut num_ops, &mut full)) + .collect::>(); + + let mut num_ops_fallback = 0; + + for f in fallbacks.iter() { + num_ops_fallback += match f { + ReduceBlockOptimInfo::Reduce(info) => info.len, + ReduceBlockOptimInfo::Elemwise(info) => info.num_ops_fused(), + }; + } + + let full_fuser_covers_all = full.num_ops_fused() == num_ops; + let fallbacks_cover_all = num_ops_fallback == num_ops; + + full_fuser_covers_all && fallbacks_cover_all + } + + /// Fuses without checking consistency and the current state. + fn fuse_no_check(&mut self, operation: &OperationIr) { + let block = self.blocks.last_mut().unwrap(); + let analyze = block.analyze(operation, &self.state, &self.fuser_default); + + let info = match analyze { + ReduceBlockFusionAnalysis::Accept => { + block.fuse(operation); + self.num_ops += 1; + block.fuser.reduce_info() + } + ReduceBlockFusionAnalysis::Refuse => { + self.state = ReduceBroadcastedStatus::Closed; + return; + } + ReduceBlockFusionAnalysis::NewBlockRequired => { + let info = block.fuser.reduce_info(); + let mut block = ReduceBlockFuser::new(self.fuser_default.clone()); + block.fuse(operation); + self.num_ops += 1; + self.blocks.push(block); + info + } + }; + + match info { + ReduceFuserInfo::FusedReduce { + shape_input_id, + axis, + } => { + // Only support last axis for now. + if axis != shape_input_id.len() - 1 { + self.state = ReduceBroadcastedStatus::Abort; + } else { + self.state = ReduceBroadcastedStatus::Init { + shape_id: shape_input_id, + axis, + }; + } + } + ReduceFuserInfo::FusedElemwise { .. } => {} + } + } +} + +impl OperationFuser> for ReduceBroadcastedFuser { + fn fuse(&mut self, operation: &OperationIr) { + if matches!( + &self.state, + ReduceBroadcastedStatus::Closed | ReduceBroadcastedStatus::Abort + ) { + return; + } + + // We first need to simulate the fusion to check the consistency, then we perform the + // fusion. + let mut next = self.clone(); + next.fuse_no_check(operation); + + // We can only check for consistency if the optimization is ready. + if next.properties().ready && !next.is_consistent() { + // Fusions that lead to inconsistent trace are closed. + self.state = ReduceBroadcastedStatus::Closed; + } else { + // Normal path. + self.fuse_no_check(operation); + } + } + + fn finish(&mut self) -> CubeOptimization { + let analyzer = FullFuserAnalyzer::new(&self.blocks); + let mut full = ReduceBroadcastedFullFuser::new(self.max_bindings, analyzer); + let mut num_ops = 0; + let fallbacks = self + .blocks + .iter_mut() + .map(|block| block.finish(&mut num_ops, &mut full)) + .collect::>(); + + let broadcasted = Arc::new(full.finish()); + let info = Arc::new(ReduceBroadcastedOptimizationInfo { + fallbacks, + broadcasted, + }); + CubeOptimization::ReduceBroadcasted(ReduceBroadcastedOptimization { info, num_ops }) + } + + fn reset(&mut self) { + let block = ReduceBlockFuser::new(self.fuser_default.clone()); + self.blocks = vec![block]; + self.num_ops = 0; + self.state = ReduceBroadcastedStatus::Starting; + } + + fn status(&self) -> FuserStatus { + match self.state { + ReduceBroadcastedStatus::Closed | ReduceBroadcastedStatus::Abort => { + return FuserStatus::Closed; + } + _ => {} + }; + + let fuser = self.blocks.last().unwrap(); + fuser.fuser.status() + } + + fn properties(&self) -> FuserProperties { + let ready = match self.state { + ReduceBroadcastedStatus::Starting | ReduceBroadcastedStatus::Abort => false, + ReduceBroadcastedStatus::Closed if self.blocks.len() == 1 => { + !self.blocks[0].is_elemwise() + } + _ => true, + }; + let mut props = FuserProperties { score: 0, ready }; + for block in self.blocks.iter() { + let p = block.properties(); + props.score += p.score; + props.ready = p.ready && props.ready; + } + props + } + + fn len(&self) -> usize { + self.num_ops + } + + fn clone_dyn(&self) -> Box>> { + Box::new(self.clone()) + } +} + +#[cfg(test)] +mod tests { + use burn_ir::{ + BaseOperationIr, BinaryOpIr, CreationOpIr, ReduceDimOpIr, TensorId, TensorIr, TensorStatus, + }; + use burn_std::{DType, Shape}; + + use super::*; + + type Run = cubecl::TestRuntime; + + #[test] + fn reduce_broadcast_workflow_1() { + let device: ::Device = Default::default(); + let mut fuser = ReduceBroadcastedFuser::::new(device); + let (tensor1_out, tensor1) = tensor(0, &[1, 2], TensorStatus::ReadWrite); + let (tensor2_out, tensor2) = tensor(1, &[1, 0], TensorStatus::ReadWrite); + + fuser.fuse(&OperationIr::BaseFloat(BaseOperationIr::Ones( + CreationOpIr { out: tensor1_out }, + ))); + fuser.fuse(&OperationIr::NumericFloat( + DType::F32, + burn_ir::NumericOperationIr::SumDim(ReduceDimOpIr { + input: tensor1, + out: tensor2_out, + axis: 1, + accumulator_len: 1, + }), + )); + + let status = fuser.status(); + assert_eq!(2, fuser.len()); + assert_eq!(status, FuserStatus::Open); + assert!(fuser.properties().ready,); + + // An existing tensor + let (_tensor3_out, tensor3) = tensor(2, &[1, 0], TensorStatus::ReadWrite); + // A new tensor + let (tensor4_out, tensor4) = tensor(3, &[1, 0], TensorStatus::ReadWrite); + fuser.fuse(&OperationIr::NumericFloat( + DType::F32, + burn_ir::NumericOperationIr::Add(BinaryOpIr { + lhs: tensor2, + rhs: tensor3, + out: tensor4_out, + }), + )); + + let status = fuser.status(); + assert_eq!(3, fuser.len()); + assert_eq!(status, FuserStatus::Open); + assert!(fuser.properties().ready,); + + // An existing tensor + let (_tensor5_out, tensor5) = tensor(4, &[1, 2], TensorStatus::ReadWrite); + // A new tensor + let (tensor6_out, tensor6) = tensor(5, &[1, 2], TensorStatus::ReadWrite); + fuser.fuse(&OperationIr::NumericFloat( + DType::F32, + burn_ir::NumericOperationIr::Add(BinaryOpIr { + lhs: tensor4, + rhs: tensor5, + out: tensor6_out, + }), + )); + + let status = fuser.status(); + assert_eq!(4, fuser.len()); + assert_eq!(status, FuserStatus::Open); + assert!(fuser.properties().ready,); + + let (tensor7_out, _tensor7) = tensor(6, &[1, 0], TensorStatus::ReadWrite); + fuser.fuse(&OperationIr::NumericFloat( + DType::F32, + burn_ir::NumericOperationIr::SumDim(ReduceDimOpIr { + input: tensor6, + out: tensor7_out, + axis: 1, + accumulator_len: 1, + }), + )); + assert_eq!(5, fuser.len()); + assert_eq!(status, FuserStatus::Open); + assert!(fuser.properties().ready,); + + let _optimization = fuser.finish(); + } + + #[test] + fn reduce_broadcast_workflow_2() { + let device: ::Device = Default::default(); + let mut fuser = ReduceBroadcastedFuser::::new(device); + let (tensor1_out, tensor1) = tensor(0, &[1, 2], TensorStatus::ReadWrite); + // An existing tensor + let (_tensor2_out, mut tensor2) = tensor(2, &[1, 2], TensorStatus::ReadOnly); + let (tensor3_out, tensor3) = tensor(3, &[1, 2], TensorStatus::ReadWrite); + + // First reduce output + let (tensor4_out, tensor4) = tensor(1, &[1, 0], TensorStatus::ReadWrite); + + fuser.fuse(&OperationIr::BaseFloat(BaseOperationIr::Ones( + CreationOpIr { out: tensor1_out }, + ))); + + fuser.fuse(&OperationIr::NumericFloat( + DType::F32, + burn_ir::NumericOperationIr::Add(BinaryOpIr { + lhs: tensor1, + rhs: tensor2.clone(), + out: tensor3_out, + }), + )); + + fuser.fuse(&OperationIr::NumericFloat( + DType::F32, + burn_ir::NumericOperationIr::SumDim(ReduceDimOpIr { + input: tensor3, + out: tensor4_out, + axis: 1, + accumulator_len: 1, + }), + )); + + let status = fuser.status(); + assert_eq!(3, fuser.len()); + assert_eq!(status, FuserStatus::Open); + assert!(fuser.properties().ready,); + + // A new tensor + let (tensor5_out, _tensor5) = tensor(5, &[1, 2], TensorStatus::ReadWrite); + // Last time we use tensor2. + tensor2.status = TensorStatus::ReadWrite; + fuser.fuse(&OperationIr::NumericFloat( + DType::F32, + burn_ir::NumericOperationIr::Add(BinaryOpIr { + lhs: tensor4, + rhs: tensor2, + out: tensor5_out, + }), + )); + + let status = fuser.status(); + assert_eq!(4, fuser.len()); + assert_eq!(status, FuserStatus::Open); + assert!(fuser.properties().ready,); + + let _optimization = fuser.finish(); + } + + fn tensor(id: u64, shape: &[usize], status: TensorStatus) -> (TensorIr, TensorIr) { + let tensor = TensorIr { + id: TensorId::new(id), + shape: Shape::from(shape), + status: TensorStatus::NotInit, + dtype: DType::F32, + }; + let mut tensor_init = tensor.clone(); + tensor_init.status = status; + + (tensor, tensor_init) + } +} diff --git a/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/fuser/block.rs b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/fuser/block.rs new file mode 100644 index 0000000..6496c88 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/fuser/block.rs @@ -0,0 +1,227 @@ +use crate::optim::{ + CubeOptimization, + elemwise::ElemwiseOptimization, + reduce::{FusedReduce, ReduceFuser, ReduceFuserInfo}, + reduce_broadcasted::{ReduceBlockOptimInfo, fuser::full::ReduceBroadcastedFullFuser}, +}; +use burn_fusion::{FuserProperties, OperationFuser}; +use burn_ir::OperationIr; +use burn_std::Shape; +use cubecl::Runtime; +use std::sync::Arc; + +/// Responsible for fusing a single reduce block or elementwise block. +/// +/// When the block kind is reduce, it supports fuse-on-read and fuse-on-write fusion. +/// Broadcasting isn't supported; another block should handle it instead. +pub struct ReduceBlockFuser { + /// We use [ReduceFuser] for both elementwise and reduce blocks, keeping only the + /// fuse-on-read trace if the block is tagged as elementwise. + /// + /// # Notes + /// + /// A single elementwise block can only exist at the end of a full [ReduceBlockFuser], + /// otherwise the optimization will be included in the reduce fusion block. + pub fuser: ReduceFuser, + pub(crate) ops: Vec, + pub(crate) kind: ReduceBlockKind, +} + +/// The current state of the fusion process. +#[derive(Debug, Clone)] +pub enum ReduceBroadcastedStatus { + /// Fusion is starting; no reduction has been fused yet. + Starting, + /// Fusion is initialized with at least one reduce operation. + /// + /// # Notes + /// + /// Subsequent reduce operations must be compatible with the previous reduction to fuse. + Init { shape_id: Shape, axis: usize }, + /// No more operations can be fused. + Closed, + /// Invalid axis. + Abort, +} + +/// The [ReduceBlockFuser] capacity to accept an [OperationIr]. +#[derive(Clone, Copy, Debug)] +pub enum ReduceBlockFusionAnalysis { + /// The operation can be fused; call [ReduceBlockFuser::fuse()]. + Accept, + /// The operation cannot be fused; the optimization should close. + Refuse, + /// The operation can be fused, but requires a new block. + NewBlockRequired, +} + +impl ReduceBlockFuser { + /// Creates a new block. + pub fn new(fuser: ReduceFuser) -> Self { + Self { + fuser: fuser.clone(), + ops: Vec::new(), + kind: ReduceBlockKind::Elemwise, + } + } + + /// Returns true if this is an elementwise fuser. + pub fn is_elemwise(&self) -> bool { + matches!(self.kind, ReduceBlockKind::Elemwise) + } + + /// Analyzes if fusion is possible within this block. + pub fn analyze( + &self, + op: &OperationIr, + status: &ReduceBroadcastedStatus, + default_node: &ReduceFuser, + ) -> ReduceBlockFusionAnalysis { + let mut fuser_try = self.fuser.clone(); + let before = fuser_try.len(); + fuser_try.fuse(op); + let after = fuser_try.len(); + + if after > before { + return ReduceBlockFusionAnalysis::Accept; + } + + // Can't create a new block if the previous one was not a reduction. + if self.fuser.reduce.is_none() { + return ReduceBlockFusionAnalysis::Refuse; + } + + // TODO: `RefLayoutSetting::SameAsBlock` might point to a `SwapDims` layout even though reduce + // requires `OnlyContiguous`. For now we just refuse / close the block. + + // Make sure the reference layout of this block is not a tensor view. + // If a block is only composed of a view, it's not accepted. + let num_views = self.fuser.fuser.num_views; + let input_ref_not_concrete = + self.fuser.fuser.fuser.fuser.num_multi_block_local_inputs() > 0; + + if num_views > 0 && input_ref_not_concrete { + return ReduceBlockFusionAnalysis::Refuse; + } + + let mut fuser_try = default_node.clone(); + let before = fuser_try.len(); + fuser_try.fuse(op); + let after = fuser_try.len(); + + if after > before { + let info = fuser_try.reduce_info(); + + return match (info, status) { + ( + ReduceFuserInfo::FusedReduce { + shape_input_id, + axis, + }, + ReduceBroadcastedStatus::Init { + shape_id, + axis: axis_init, + }, + ) => { + if shape_id == &shape_input_id && axis_init == &axis { + ReduceBlockFusionAnalysis::NewBlockRequired + } else { + ReduceBlockFusionAnalysis::Refuse + } + } + ( + ReduceFuserInfo::FusedElemwise { shape_id }, + ReduceBroadcastedStatus::Init { + shape_id: shape_init, + .. + }, + ) => { + if &shape_id == shape_init { + ReduceBlockFusionAnalysis::NewBlockRequired + } else { + ReduceBlockFusionAnalysis::Refuse + } + } + _ => ReduceBlockFusionAnalysis::Refuse, + }; + } + + ReduceBlockFusionAnalysis::Refuse + } + + /// Fuses an operation within this block. + /// + /// # Warning + /// + /// Ensure [Self::analyze()] is called before this function to confirm the operation is accepted. + pub fn fuse(&mut self, op: &OperationIr) { + self.fuser.fuse(op); + self.ops.push(op.clone()); + + // Update the kind if a reduction is introduced to an elementwise block. + if let (Some(reduce), ReduceBlockKind::Elemwise) = (&self.fuser.reduce, &self.kind) { + self.kind = ReduceBlockKind::Reduce { + ops_index: self.ops.len() - 1, + reduce: Box::new(reduce.clone()), + }; + } + } + + /// Computes the fuser properties. + pub fn properties(&self) -> FuserProperties { + let mut properties = self.fuser.properties(); + if let ReduceBlockKind::Elemwise = &self.kind { + // Elementwise traces are always ready to run. + properties.ready = true; + } + properties + } + + pub fn finish( + &mut self, + num_ops: &mut usize, + full: &mut ReduceBroadcastedFullFuser, + ) -> ReduceBlockOptimInfo { + full.register(self); + + match &self.kind { + ReduceBlockKind::Elemwise => { + let len = self.fuser.fuser_read_fallback.len(); + let device = self.fuser.device.clone(); + *num_ops += len; + let trace = self.fuser.fuser_read_fallback.finish(); + let client = R::client(&device); + let elementwise = ElemwiseOptimization::new(trace, client, device, len); + ReduceBlockOptimInfo::Elemwise(Arc::new(elementwise)) + } + ReduceBlockKind::Reduce { .. } => { + *num_ops += self.fuser.len(); + let optim = self.fuser.finish(); + let info = match optim { + CubeOptimization::Reduce(optim) => optim.info, + _ => unreachable!("Expected Reduce optimization"), + }; + ReduceBlockOptimInfo::Reduce(info) + } + } + } +} + +#[derive(Clone, Debug)] +pub enum ReduceBlockKind { + Elemwise, + Reduce { + ops_index: usize, + reduce: Box, + }, +} + +impl Clone for ReduceBlockFuser { + fn clone(&self) -> Self { + Self { + fuser: self.fuser.clone(), + ops: self.ops.clone(), + kind: self.kind.clone(), + } + } +} diff --git a/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/fuser/full.rs b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/fuser/full.rs new file mode 100644 index 0000000..d1d3238 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/fuser/full.rs @@ -0,0 +1,177 @@ +use crate::{ + engine::{ + fuser::TraceOperationFuser, + settings::{FuseSettings, RefLayoutSetting, VectorizationSetting}, + }, + optim::{ + reduce::{FusedReduce, ReduceInstruction}, + reduce_broadcasted::{ + ReduceBroadcastedInfo, + fuser::{ + block::{ReduceBlockFuser, ReduceBlockKind}, + full_analyzer::FullFuserAnalyzer, + }, + launch::ReduceBroadcastedFuseBlock, + }, + }, +}; +use burn_fusion::OperationFuser; +use cubecl::Runtime; +use cubek::reduce::components::instructions::ReduceOperationConfig; + +/// Responsible for fusing a single trace for all operations involved in this optimization. +pub struct ReduceBroadcastedFullFuser { + pub(crate) fuser: TraceOperationFuser, + analyzer: FullFuserAnalyzer, + blocks: Vec, + settings_read: FuseSettings, + settings_write: FuseSettings, +} + +impl ReduceBroadcastedFullFuser { + /// Creates a new fuser with the given settings. + pub fn new(max_bindings: u32, analyzer: FullFuserAnalyzer) -> Self { + let settings_read = FuseSettings { + output_shape_updates: true, + broadcast: true, + inplace: false, + ref_layout: RefLayoutSetting::OnlyContiguous, + vectorization: VectorizationSetting::Activated, + }; + let settings_write = FuseSettings { + output_shape_updates: false, + inplace: false, + broadcast: false, + ref_layout: RefLayoutSetting::OnlyContiguous, + // Deactivated for now, but would be cool to support vectorization of the output. + vectorization: VectorizationSetting::Deactivated, + }; + let fuser = TraceOperationFuser::new(max_bindings, settings_read); + + Self { + fuser, + blocks: Vec::new(), + settings_write, + settings_read, + analyzer, + } + } + + /// Returns the amount of operations fused. + pub fn num_ops_fused(&self) -> usize { + let mut fused = self.fuser.num_ops; + + for block in self.blocks.iter() { + match block { + ReduceBlockKind::Elemwise => {} + // The base fuser doesn't hold the reduce ops. + ReduceBlockKind::Reduce { .. } => fused += 1, + } + } + + fused + } + + /// Finishes fusing all blocks. + pub fn finish(mut self) -> ReduceBroadcastedInfo { + let mut reduce_axis = 0; + let mut blocks = Vec::new(); + + for block in self.blocks.iter() { + match block { + ReduceBlockKind::Elemwise => {} + ReduceBlockKind::Reduce { reduce, .. } => { + let config = match reduce.inst { + ReduceInstruction::ArgMax => ReduceOperationConfig::ArgMax, + ReduceInstruction::ArgMin => ReduceOperationConfig::ArgMin, + ReduceInstruction::Prod => ReduceOperationConfig::Prod, + ReduceInstruction::Mean => ReduceOperationConfig::Mean, + ReduceInstruction::Sum => ReduceOperationConfig::Sum, + ReduceInstruction::Max => ReduceOperationConfig::Max, + ReduceInstruction::Min => ReduceOperationConfig::Min, + ReduceInstruction::MaxAbs => ReduceOperationConfig::MaxAbs, + }; + + let block = ReduceBroadcastedFuseBlock { + op: config, + input: reduce.input.clone(), + output: reduce.output.clone(), + }; + reduce_axis = reduce.axis; + blocks.push(block); + } + } + } + + let trace = self.fuser.finish(); + + ReduceBroadcastedInfo { + blocks, + trace, + reduce_axis, + } + } + + /// Registers a [ReduceBlockFuser] to build the trace. + pub fn register(&mut self, block: &ReduceBlockFuser) { + // Helper to close previous blocks if necessary + if !self.fuser.is_empty() { + let mut settings = self.settings_read; + settings.vectorization = VectorizationSetting::EqualThanPreviousBlock { block_pos: 0 }; + settings.ref_layout = RefLayoutSetting::SameAsBlock { block_pos: 0 }; + self.fuser.next_block([], settings, false); + + let analysis = self.analyzer.retrieve_next(); + + for (tensor, block_pos) in analysis.inputs { + self.fuser.block_local_input(&tensor, block_pos, false); + } + } + + match &block.kind { + ReduceBlockKind::Elemwise => { + for op in &block.ops { + self.fuser.fuse(op); + } + self.blocks.push(ReduceBlockKind::Elemwise); + } + ReduceBlockKind::Reduce { ops_index, reduce } => { + for op in &block.ops[0..*ops_index] { + self.fuser.fuse(op); + } + + let [input] = self + .fuser + .next_block([&reduce.op.input], self.settings_write, false); + + let output = self.fuser.output_unhandled(&reduce.op.out); + let analysis = self.analyzer.retrieve_next(); + + // Can be broadcasted so the generated buffer can be global. + for (tensor, block_pos) in analysis.inputs { + self.fuser.block_local_input(&tensor, block_pos, false); + } + + let fused_reduce = FusedReduce { + input, + output, + acc: reduce.acc, + axis: reduce.axis, + op: reduce.op.clone(), + use_planes: reduce.use_planes, + shared: reduce.shared, + inst: reduce.inst, + }; + + self.blocks.push(ReduceBlockKind::Reduce { + ops_index: *ops_index, + reduce: Box::new(fused_reduce), + }); + + for op in &block.ops[*ops_index + 1..block.ops.len()] { + self.fuser.fuse(op); + } + } + } + } +} diff --git a/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/fuser/full_analyzer.rs b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/fuser/full_analyzer.rs new file mode 100644 index 0000000..b845c2e --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/fuser/full_analyzer.rs @@ -0,0 +1,159 @@ +use super::block::ReduceBlockKind; +use crate::optim::reduce_broadcasted::fuser::block::ReduceBlockFuser; +use burn_ir::{TensorId, TensorIr}; +use cubecl::Runtime; +use std::collections::BTreeMap; + +#[derive(Debug)] +pub struct FullFuserAnalyzer { + // We need to know the block id of which we can reuse the read local input. + analyses: Vec>, +} + +impl FullFuserAnalyzer { + pub fn new(blocks: &[ReduceBlockFuser]) -> Self { + let mut state = AnalysisState::default(); + + for block in blocks.iter() { + for (pos, op) in block.ops.iter().enumerate() { + let potential_from_previous_blocks = op.inputs(); + let potential_to_next_blocks = op.outputs(); + + match &block.kind { + ReduceBlockKind::Elemwise => { + state.register( + potential_from_previous_blocks, + potential_to_next_blocks, + BlockKind::Full, + ); + } + ReduceBlockKind::Reduce { ops_index, .. } => { + if pos < *ops_index { + state.register( + potential_from_previous_blocks, + potential_to_next_blocks, + BlockKind::Full, + ); + } else if pos > *ops_index { + state.register( + potential_from_previous_blocks, + potential_to_next_blocks, + BlockKind::Single, + ); + } else { + state.next_block(); + } + } + } + } + state.next_block(); + } + + // First one is never called. + state.analyses.remove(0); + + Self { + analyses: state.analyses, + } + } + + pub fn retrieve_next(&mut self) -> FullFuserAnalysis { + let inputs = self.analyses.remove(0); + FullFuserAnalysis { inputs } + } +} + +#[derive(Debug)] +pub struct FullFuserAnalysis { + /// The tensor received from a previous block. + pub inputs: Vec<(TensorIr, usize)>, +} + +#[derive(Default)] +struct AnalysisState { + /// That pool contains tensors that are available in the fuse-on-write part of a reduce, not + /// broadcasted. + available_from_previous_single: BTreeMap, + /// That pool contains tensors that are available in the fuse-on-read of a reduce and the + /// element-wise broadcasted part + available_from_previous_full: BTreeMap, + block_data: Vec<(TensorIr, usize)>, + analyses: Vec>, + current_full: Vec, + current_single: Vec, +} + +enum BlockKind { + Full, + Single, +} + +impl AnalysisState { + fn next_block(&mut self) { + let block_pos = self.analyses.len(); + let data = core::mem::take(&mut self.block_data); + self.analyses.push(data); + + // Makes the current tensor reads available for the next block. + for p in self.current_single.drain(..) { + // We need to keep the earliest block position. + self.available_from_previous_single + .entry(p.id) + .or_insert(block_pos); + } + for p in self.current_full.drain(..) { + // We need to keep the earliest block position. + self.available_from_previous_full + .entry(p.id) + .or_insert(block_pos); + } + } + + fn register<'a>( + &mut self, + potential_from_previous_blocks: impl Iterator, + potential_to_next_blocks: impl Iterator, + kind: BlockKind, + ) { + match kind { + BlockKind::Full => { + for potential in potential_from_previous_blocks { + // We can't since it's not in the same scope. + // + // TODO: Find a way to merge multiple reduce loops. + // + // if let Some(block_pos) = self.available_from_previous_full.get(&potential.id) { + // self.block_data.push((potential.clone(), *block_pos)); + // } + + // We can since it's a broadcast. + if let Some(block_pos) = self.available_from_previous_single.get(&potential.id) + { + self.block_data.push((potential.clone(), *block_pos)); + } + + // Can reuse the read. + self.current_full.push(potential.clone()); + } + + for p in potential_to_next_blocks { + self.current_full.push(p.clone()); + } + } + BlockKind::Single => { + for potential in potential_from_previous_blocks { + if let Some(block_pos) = self.available_from_previous_single.get(&potential.id) + { + self.block_data.push((potential.clone(), *block_pos)); + } + // Can reuse the read. + self.current_single.push(potential.clone()); + } + + for p in potential_to_next_blocks { + self.current_single.push(p.clone()); + } + } + } + } +} diff --git a/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/fuser/mod.rs b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/fuser/mod.rs new file mode 100644 index 0000000..22bb521 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/fuser/mod.rs @@ -0,0 +1,6 @@ +mod base; +mod block; +mod full; +mod full_analyzer; + +pub use base::*; diff --git a/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/launch.rs b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/launch.rs new file mode 100644 index 0000000..334a50e --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/launch.rs @@ -0,0 +1,149 @@ +use crate::{ + engine::{ + codegen::ir::{FuseArg, FuseBlockConfig, GlobalArgsLaunch, RefLayout}, + launch::runner::{TraceRunner, Vectorization}, + }, + optim::reduce_broadcasted::unit::{ + ElemwiseFuseBlockLaunch, ReduceFuseBlockLaunch, reduce_kernel_broadcasted, + }, +}; +use cubecl::{ + Runtime, + ir::{ElemType, FloatKind, StorageType}, + prelude::*, + server::LaunchError, +}; +use cubek::reduce::{ + ReduceDtypes, VectorizationMode, + components::instructions::ReduceOperationConfig, + launch::RoutineStrategy, + output_vectorization_axis, + routines::{ + BlueprintStrategy, GlobalReduceBlueprint, ReduceProblem, ReduceVectorSettings, Routine, + unit::{UnitRoutine, UnitStrategy}, + }, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct ReduceBroadcastedFuseBlock { + pub(crate) op: ReduceOperationConfig, + pub(crate) input: FuseArg, + pub(crate) output: FuseArg, +} + +#[derive(new)] +pub struct FusedReduceBroadcastedLaunch<'a> { + blocks: &'a Vec, + reduce_axis: usize, + // TODO: Support multiple strategies. + _strategy: RoutineStrategy, +} + +impl Vectorization for FusedReduceBroadcastedLaunch<'_> {} + +impl TraceRunner for FusedReduceBroadcastedLaunch<'_> { + type Error = LaunchError; + + fn run<'a>( + &'a self, + client: &'a ComputeClient, + inputs: GlobalArgsLaunch, + outputs: GlobalArgsLaunch, + configs: &'a [FuseBlockConfig], + ) -> Result<(), Self::Error> { + let routine = UnitRoutine; + let first_config = &configs[0]; + + let shape = match &first_config.ref_layout { + RefLayout::Concrete(FuseArg::Output(..)) => { + outputs.shape_ref(&first_config.ref_layout, first_config.rank) + } + _ => inputs.shape_ref(&first_config.ref_layout, first_config.rank), + }; + + let reduce_len = shape[self.reduce_axis]; + let reduce_count = shape.iter().product::() / reduce_len; + let address_type = inputs + .required_address_type() + .max(outputs.required_address_type()); + + let (blueprint, settings) = routine + .prepare::( + client, + ReduceProblem { + reduce_len, + reduce_count, + axis: self.reduce_axis, + dtypes: ReduceDtypes { + input: StorageType::Scalar(ElemType::Float(FloatKind::F32)), + output: StorageType::Scalar(ElemType::Float(FloatKind::F32)), + accumulation: StorageType::Scalar(ElemType::Float(FloatKind::F32)), + }, + address_type, + // We assume at least one block. + instruction: self.blocks.first().unwrap().op, + }, + ReduceVectorSettings { + vectorization_mode: VectorizationMode::Parallel, + vector_size_input: first_config.width, + vector_size_output: 1, + }, + BlueprintStrategy::Inferred(UnitStrategy), + ) + .unwrap(); + + assert_eq!(blueprint.vectorization_mode, VectorizationMode::Parallel); + + let mut blocks = SequenceArg::new(); + let mut index = 0; + + for block in self.blocks { + let arg = ReduceFuseBlockLaunch::new( + block.op, + configs[index].clone(), + configs[index + 1].clone(), + block.input.clone(), + block.output.clone(), + match blueprint.global { + GlobalReduceBlueprint::Unit(bpt) => bpt, + _ => panic!(), + }, + ); + index += 2; + blocks.push(arg); + } + + let block_end = match configs.len() > index { + true => ComptimeOptionArgs::Some(ElemwiseFuseBlockLaunch::new( + configs.last().cloned().unwrap(), + )), + false => ComptimeOptionArgs::None, + }; + + let out_vec_axis = output_vectorization_axis( + &inputs.strides_ref(&first_config.ref_layout, first_config.rank), + self.reduce_axis, + VectorizationMode::Parallel, + ); + + // TODO: Ensure parallel is selected. + + unsafe { + reduce_kernel_broadcasted::launch_unchecked::( + client, + settings.cube_count, + settings.cube_dim, + settings.address_type, + inputs, + outputs, + self.reduce_axis, + out_vec_axis, + blocks, + block_end, + ); + } + + Ok(()) + } +} diff --git a/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/mod.rs b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/mod.rs new file mode 100644 index 0000000..b29f224 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/mod.rs @@ -0,0 +1,9 @@ +mod fuser; +mod optimization; + +pub(crate) mod launch; +pub(crate) mod tune; +pub(crate) mod unit; + +pub use fuser::*; +pub use optimization::*; diff --git a/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/optimization.rs b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/optimization.rs new file mode 100644 index 0000000..ee7cf88 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/optimization.rs @@ -0,0 +1,219 @@ +#[cfg(feature = "autotune")] +use crate::optim::reduce::tune::fused_reduce_autotune; +use crate::{ + CubeFusionHandle, FallbackOperation, + engine::{ + launch::FuseTraceLauncher, + trace::{FuseTrace, TraceError, TuneOutput}, + }, + optim::{ + elemwise::{ElemwiseOptimization, ElemwiseOptimizationState}, + reduce::{ReduceOptimizationInfo, ReduceOptimizationState, ReduceOptimizationTuneArg}, + reduce_broadcasted::{ + launch::{FusedReduceBroadcastedLaunch, ReduceBroadcastedFuseBlock}, + tune::fused_broadcasted_reduce_autotune, + }, + }, +}; +use burn_fusion::stream::Context; +use cubecl::{Runtime, prelude::*}; +use cubek::reduce::launch::RoutineStrategy; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; + +pub struct ReduceBroadcastedOptimization { + pub(crate) info: Arc>, + pub(crate) num_ops: usize, +} + +pub(crate) struct ReduceBroadcastedOptimizationInfo { + pub(crate) fallbacks: Vec>, + pub(crate) broadcasted: Arc, +} + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub(crate) struct ReduceBroadcastedInfo { + pub(crate) blocks: Vec, + pub(crate) trace: FuseTrace, + pub(crate) reduce_axis: usize, +} + +pub(crate) enum ReduceBlockOptimInfo { + Reduce(Arc>), + Elemwise(Arc>), +} + +impl ReduceBlockOptimInfo { + pub fn from_state(device: &R::Device, state: ReduceBlockState) -> Self { + match state { + ReduceBlockState::Reduce(state) => { + Self::Reduce(Arc::new(ReduceOptimizationInfo::from_state(device, state))) + } + ReduceBlockState::Elemwise(state) => { + Self::Elemwise(Arc::new(ElemwiseOptimization::from_state(device, state))) + } + } + } + pub fn to_state(&self) -> ReduceBlockState { + match self { + Self::Reduce(info) => ReduceBlockState::Reduce(info.to_state()), + Self::Elemwise(info) => ReduceBlockState::Elemwise(info.to_state()), + } + } +} + +pub(crate) struct ReduceBroadcastedOptimizationTuneArg { + pub(crate) fallbacks: Vec>, + pub(crate) broadcasted: Arc, + pub(crate) client: ComputeClient, + pub(crate) device: R::Device, +} + +pub(crate) enum ReduceBlockOptimArg { + Reduce(ReduceOptimizationTuneArg), + Elemwise(Arc>), +} + +impl ReduceBlockOptimArg { + pub fn execute_fallback( + &self, + context: &mut Context>, + ) -> Option> { + match self { + ReduceBlockOptimArg::Reduce(reduce) => { + #[cfg(feature = "autotune")] + { + fused_reduce_autotune::(reduce.clone(), context); + None + } + #[cfg(not(feature = "autotune"))] + Some(reduce.execute_fallback(context)) + } + ReduceBlockOptimArg::Elemwise(elem) => { + elem.execute(context); + None + } + } + } +} + +#[derive(Serialize, Deserialize, Debug)] +pub struct ReduceBroadcastedOptimizationState { + fallbacks: Vec, + broadcasted: ReduceBroadcastedInfo, + num_ops: usize, +} + +#[derive(Serialize, Deserialize, Debug)] +#[allow(clippy::large_enum_variant)] // Only for serialization. +pub enum ReduceBlockState { + Reduce(ReduceOptimizationState), + Elemwise(ElemwiseOptimizationState), +} + +impl ReduceBroadcastedOptimizationTuneArg { + pub fn execute_fused( + &self, + context: &mut Context>, + strategy: RoutineStrategy, + ) -> Result, TraceError> { + let launch = FusedReduceBroadcastedLaunch::new( + &self.broadcasted.blocks, + self.broadcasted.reduce_axis, + strategy, + ); + let launcher = FuseTraceLauncher::new(&self.broadcasted.trace, &launch); + + launcher + .launch(&self.client, &self.device, context) + .map_err(|err| TraceError::RunnerError(format!("{:?}", err))) + } + + pub fn execute_fallback(&self, context: &mut Context>) { + for fallback in self.fallbacks.iter() { + fallback.execute_fallback(context); + } + } +} + +#[allow(clippy::too_many_arguments)] +impl ReduceBroadcastedOptimization { + /// Execute the optimization. + pub fn execute( + &mut self, + context: &mut Context>, + fallback: impl Fn(usize) -> Box>, + ) { + let mut current_index = 0; + let mut client = None; + let mut device = None; + + let fallbacks = self + .info + .fallbacks + .iter() + .map(|info| { + match info { + ReduceBlockOptimInfo::Reduce(info) => { + // The index of the fallback reduce is the number of ops fused as read. + let fallback = fallback(current_index + info.len_read); + client = Some(info.client.clone()); + device = Some(info.device.clone()); + let arg = ReduceOptimizationTuneArg { + info: info.clone(), + fallback: Arc::new(fallback), + }; + current_index += info.len; + ReduceBlockOptimArg::Reduce(arg) + } + ReduceBlockOptimInfo::Elemwise(op) => ReduceBlockOptimArg::Elemwise(op.clone()), + } + }) + .collect(); + + let arg = ReduceBroadcastedOptimizationTuneArg { + fallbacks, + client: client.unwrap(), + device: device.unwrap(), + broadcasted: self.info.broadcasted.clone(), + }; + + #[cfg(feature = "autotune")] + fused_broadcasted_reduce_autotune::(arg, context); + + #[cfg(not(feature = "autotune"))] + arg.execute_fallback(context); + } + + pub fn to_state(&self) -> ReduceBroadcastedOptimizationState { + ReduceBroadcastedOptimizationState { + fallbacks: self + .info + .fallbacks + .iter() + .map(|info| info.to_state()) + .collect(), + broadcasted: self.info.broadcasted.as_ref().clone(), + num_ops: self.num_ops, + } + } + + pub fn from_state(device: &R::Device, state: ReduceBroadcastedOptimizationState) -> Self { + Self { + info: Arc::new(ReduceBroadcastedOptimizationInfo { + fallbacks: state + .fallbacks + .into_iter() + .map(|state| ReduceBlockOptimInfo::from_state(device, state)) + .collect(), + broadcasted: Arc::new(state.broadcasted), + }), + num_ops: state.num_ops, + } + } + + /// Returns the number of output buffers added by fusion. + pub fn num_ops_fused(&self) -> usize { + self.num_ops + } +} diff --git a/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/tune.rs b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/tune.rs new file mode 100644 index 0000000..4a61ae7 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/tune.rs @@ -0,0 +1,160 @@ +use super::optimization::ReduceBroadcastedOptimizationTuneArg; +use crate::{ + CubeFusionHandle, + engine::trace::TuneOutput, + optim::{reduce::ReduceOptimizationInfo, reduce_broadcasted::ReduceBlockOptimArg}, + tune::{FusionInputGen, TuneInput}, +}; +use burn_backend::cubecl::dtype_to_elem_type; +use burn_fusion::stream::Context; +use cubecl::{ + AutotuneKey, CubeTuneId, Runtime, + tune::{LocalTuner, Tunable, TunableSet, TuneGroup, local_tuner}, +}; +use cubek::reduce::{ + launch::{RoutineStrategy, tune_key::ReduceAutotuneKey}, + routines::{BlueprintStrategy, unit::UnitStrategy}, +}; +use serde::{Deserialize, Serialize}; + +/// Autotune key for fused broadcasted reduction operations. +/// +/// Captures the characteristics of the fusion (reads, writes, ops) to ensure +/// the best kernel is selected for specific fused graph shapes. +#[derive(Hash, Eq, PartialEq, Debug, Clone, Serialize, Deserialize, AutotuneKey)] +pub struct FusedBroadcastedReduceAutotuneKey { + reduce_key: ReduceAutotuneKey, + #[autotune(anchor)] + fuse_num_reads: usize, + #[autotune(anchor)] + fuse_num_writes: usize, + #[autotune(anchor)] + fuse_num_ops: usize, + fuse_num_blocks: usize, +} + +/// Executes the autotuning process for fused reduction operations. +/// +/// This function initializes a local tuner and attempts multiple strategies +/// (fallback vs. unit strategy) to find the most efficient execution path. +pub fn fused_broadcasted_reduce_autotune( + arg: ReduceBroadcastedOptimizationTuneArg, + context: &mut Context>, +) { + static TUNER: LocalTuner = local_tuner!(); + + let tunables = TUNER.init(|| { + const PRIORITY_MAX: i8 = 2; + let mut set = TunableSet::new(create_key::, FusionInputGen); + + let group = TuneGroup::::new( + "fused_reduce_broadcasted", + |_key| PRIORITY_MAX, + ); + + // Standard fallback implementation - guaranteed to work. + set = set.with(Tunable::new( + "fused_reduce_broadcasted_fallback", + tune_fallback::, + )); + + // Specialized unit strategy for fused reductions. + set = set.with( + Tunable::new("fused_reduce_broadcasted_unit", move |input| { + tune_reduce::( + input, + &RoutineStrategy::Unit(BlueprintStrategy::Inferred(UnitStrategy)), + ) + }) + .group(&group, |_| PRIORITY_MAX), + ); + + set + }); + + TUNER.execute( + &CubeTuneId::new(&arg.client, &arg.device), + &arg.client.clone(), + tunables, + TuneInput::new(context, arg), + ); +} + +/// Generates the autotune key based on the current optimization context and trace blocks. +pub(crate) fn create_key( + input: &TuneInput>, +) -> FusedBroadcastedReduceAutotuneKey { + let opt = input.optimization(); + assert!( + input.is_original(), + "Forked context not supported for key generation" + ); + let context = input.context(); + + // The fusion must start with a reduction block to be valid here. + let info = match &opt.fallbacks[0] { + ReduceBlockOptimArg::Reduce(reduce) => &reduce.info, + ReduceBlockOptimArg::Elemwise(_) => { + unreachable!("Fusion must start with a reduction block") + } + }; + + let key = generate_reduce_autotune_key(info, context); + + // Sum up complexity metrics across all blocks in the fused trace. + let (mut num_reads, mut num_writes, mut num_ops) = (0, 0, 0); + + for block in opt.broadcasted.trace.blocks.iter() { + num_reads += block.reads.len(); + num_writes += block.writes.len(); + num_ops += block.ops.len(); + } + + FusedBroadcastedReduceAutotuneKey::new( + key, + num_reads, + num_writes, + num_ops, + info.trace.blocks.len(), + ) +} + +/// Helper to generate the base reduction key (shapes, types, axes). +fn generate_reduce_autotune_key( + info: &ReduceOptimizationInfo, + context: &Context>, +) -> ReduceAutotuneKey { + let input = context.tensors.get(&info.reduce.op.input.id).unwrap(); + let out = context.tensors.get(&info.reduce.op.out.id).unwrap(); + let acc = info.reduce.acc.into_elem(); + + ReduceAutotuneKey::generate( + dtype_to_elem_type(input.dtype), + dtype_to_elem_type(out.dtype), + acc, + &input.shape, + info.reduce.axis == input.shape.rank() - 1, // Is it the last dimension? + info.reduce.axis, + ) +} + +/// Executes a fused reduction using a specific routine strategy. +fn tune_reduce( + input: TuneInput>, + strategy: &RoutineStrategy, +) -> Result, String> { + input + .execute(|ctx, opt| opt.execute_fused(ctx, strategy.clone())) + .map_err(|e| format!("{e:?}")) +} + +/// Executes the fallback implementation for the reduction. +fn tune_fallback( + input: TuneInput>, +) -> Result, String> { + input.execute(|ctx, opt| { + opt.execute_fallback(ctx); + }); + // Fallback is often used as a baseline, returning unchecked output. + Ok(TuneOutput::UnChecked(std::marker::PhantomData)) +} diff --git a/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/unit.rs b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/unit.rs new file mode 100644 index 0000000..6fb31e8 --- /dev/null +++ b/crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/unit.rs @@ -0,0 +1,220 @@ +use crate::{ + engine::codegen::{ + ir::{FuseArg, FuseBlockConfig, FuseType, GlobalArgs, multi_block_variables_init}, + kernel::{fuse_on_write, init_locals}, + }, + optim::reduce::args::{FusedReduceArgs, FusedReduceInput, FusedReduceOutput}, +}; +use cubecl::{Runtime, define_size, prelude::*, std::tensor::r#virtual::VirtualTensor}; +use cubek::reduce::{ + ReduceInstruction, ReducePrecision, VectorizationMode, + components::{ + args::NumericVector, + global::unit::GlobalFullUnitReduce, + instructions::{ReduceOperation, ReduceOperationConfig}, + }, + init_tensors, + routines::UnitReduceBlueprint, +}; + +/// A configuration block for a reduction operation within a fused kernel. +/// +/// This struct holds all the compile-time information needed to perform a +/// reduction, including the operation type (Sum, Max, etc.) and the layout +/// configuration for both input and output. +#[derive(CubeType, CubeLaunch, Clone)] +pub struct ReduceFuseBlock { + #[cube(comptime)] + op: ReduceOperationConfig, + #[cube(comptime)] + config_input: FuseBlockConfig, + #[cube(comptime)] + config_output: FuseBlockConfig, + #[cube(comptime)] + input: FuseArg, + #[cube(comptime)] + output: FuseArg, + #[cube(comptime)] + blueprint: UnitReduceBlueprint, +} + +/// A configuration block for an elementwise operation that follows a reduction. +#[derive(CubeType, CubeLaunch, Clone)] +pub struct ElemwiseFuseBlock { + #[cube(comptime)] + config: FuseBlockConfig, +} + +/// The entry point for a broadcasted reduction kernel. +/// +/// This kernel initializes local variables for multiple reduction blocks and then +/// executes the reduction sequence. +/// +/// # Arguments +/// +/// * `inputs` - Global arguments containing input tensor handles. +/// * `outputs` - Global arguments containing output tensor handles. +/// * `reduce_axis` - The dimension along which the reduction is performed. +/// * `blocks` - A sequence of reduction operations to execute. +/// * `block_end` - An optional elementwise block to execute after reductions are complete. +#[cube(launch_unchecked, address_type = "dynamic")] +pub fn reduce_kernel_broadcasted( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + reduce_axis: usize, + out_vec_axis: usize, + blocks: Sequence, + block_end: ComptimeOption, +) { + #[unroll] + for i in 0..blocks.len() { + let block = blocks.index(i); + multi_block_variables_init(&block.config_input, &mut outputs.variables); + multi_block_variables_init(&block.config_output, &mut outputs.variables); + } + + reduce_many( + inputs, + outputs, + reduce_axis, + out_vec_axis, + blocks, + block_end, + ); +} + +define_scalar!(In); +define_scalar!(Acc); +define_scalar!(Out); + +define_size!(InSize); +define_size!(OutSize); + +/// Configures the precision polyfills for the reduction based on the block's `FuseType`. +#[cube] +fn set_polyfill_block(block: &ReduceFuseBlock) { + let input_precision = comptime!(block.input.precision()); + let output_precision = comptime!(block.output.precision()); + let acc_precision = comptime!(match input_precision { + FuseType::F64 => FuseType::F64, + FuseType::F32 => FuseType::F32, + FuseType::Flex32 => FuseType::F32, + FuseType::F16 => FuseType::F32, + FuseType::BF16 => FuseType::F32, + FuseType::I64 => FuseType::I64, + FuseType::I32 => FuseType::I32, + FuseType::I16 => FuseType::I32, + FuseType::I8 => FuseType::I32, + FuseType::U64 => FuseType::U64, + FuseType::U32 => FuseType::U32, + FuseType::U16 => FuseType::U32, + FuseType::U8 => FuseType::U32, + }); + + set_polyfill::(comptime!( + input_precision.into_type(block.config_input.width) + )); + set_polyfill::(comptime!( + output_precision.into_type(block.config_output.width) + )); + set_polyfill::(comptime!(acc_precision.into_type(block.config_input.width))); +} + +/// Internal logic for executing a sequence of reduction blocks followed by an optional +/// trailing elementwise block. +#[cube] +#[allow(clippy::clone_on_copy)] +fn reduce_many( + inputs: &GlobalArgs, + outputs: &mut GlobalArgs, + reduce_axis: usize, + out_vec_axis: usize, + blocks: Sequence, + block_end: ComptimeOption, +) { + let mut axis_size = 0; + + #[unroll] + for i in 0..blocks.len() { + let block = blocks.index(i); + let input = FusedReduceInput { + global: inputs.clone(), + config: comptime!(block.config_input.clone()), + arg: comptime!(block.input.clone()), + }; + let global = outputs.clone(); + let config = comptime!(block.config_output.clone()); + let arg = comptime!(block.output.clone()); + let mut output = FusedReduceOutput { + global, + config, + arg, + }; + + set_polyfill_block(block); + let (input, mut output) = + init_tensors::(&input, &mut output); + + axis_size = reduce_step::<(In, InSize, Acc), (Out, OutSize), ReduceOperation>( + &input, + &mut output, + reduce_axis, + out_vec_axis, + block.op, + comptime!(block.blueprint.clone()), + ); + } + + #[comptime] + if let ComptimeOption::Some(block) = block_end { + let global_index = ABSOLUTE_POS; + let width = block.config.width; + let num_iter = axis_size / width; + let size!(N) = width; + + for i in 0..num_iter { + // Register block local inputs. + let values = Registry::>::new(); + let args = comptime![Vec::::new()]; + let index = global_index * num_iter + i; + let mut locals = init_locals(inputs, outputs, &block.config); + + fuse_on_write::( + inputs, + outputs, + &mut locals, + index, + values, + args, + &block.config.clone(), + ) + } + } +} + +#[cube] +/// Executes a single reduction step using a specified instruction and blueprint. +/// +/// Returns the size of the axis that was reduced. +fn reduce_step>( + input: &VirtualTensor, + output: &mut VirtualTensor, + reduce_axis: usize, + out_vec_axis: usize, + #[comptime] config: I::Config, + #[comptime] blueprint: UnitReduceBlueprint, +) -> usize { + let inst = I::from_config(config); + let axis_size = input.shape(reduce_axis); + + GlobalFullUnitReduce::execute::( + input, + output, + reduce_axis, + out_vec_axis, + &inst, + VectorizationMode::Parallel, + comptime!(blueprint), + ); + axis_size +} diff --git a/crates/burn-cubecl-fusion/src/tune.rs b/crates/burn-cubecl-fusion/src/tune.rs new file mode 100644 index 0000000..8cd4baa --- /dev/null +++ b/crates/burn-cubecl-fusion/src/tune.rs @@ -0,0 +1,215 @@ +use crate::CubeFusionHandle; +use burn_fusion::stream::Context; +use burn_ir::{HandleContainer, TensorId, TensorIr}; +use cubecl::Runtime; +use cubecl::tune::{InputGenerator, TuneInputs}; +use hashbrown::HashMap; +use std::marker::PhantomData; +use std::sync::Arc; + +/// [`TuneInputs`] marker for [`TuneInput`]. This is the indirection that lets a +/// `TunableSet<_, FusionTuneInputs, _>` live `'static` inside `LocalTuner::init`'s +/// cache while still accepting a borrowing `TuneInput<'a, …>` at `execute` time, via HRTB +/// over `'a`. +#[allow(clippy::type_complexity)] +pub(crate) struct FusionTuneInputs(PhantomData<(fn() -> R, fn() -> O)>); + +impl TuneInputs for FusionTuneInputs { + type At<'a> = TuneInput<'a, R, O>; +} + +/// [`InputGenerator`] for [`TuneInput`]: produces a benchmark-only [`TuneState::Fork`] +/// (i.e. with `new_handles = None`). Benchmarks discard their outputs, so the returned +/// input skips the handle-tracking machinery that a wasm-fallback fork carries. +pub(crate) struct FusionInputGen; + +impl InputGenerator> for FusionInputGen +where + K: 'static, + R: Runtime, + O: Send + Sync + 'static, +{ + fn generate<'a>( + &self, + _key: &K, + inputs: & as TuneInputs>::At<'a>, + ) -> as TuneInputs>::At<'a> { + inputs.for_benchmark() + } +} + +/// Shared staging area for handles produced by a wasm try-all fork. +/// +/// [`Fork`](TuneState::Fork)'s `Drop` dumps all of the fork's handles here (replacing any +/// previous contents), and [`Original`](TuneState::Original)'s `Drop` drains this and +/// filters each entry against the real context's current handles — anything the real +/// context doesn't already have is a genuinely new output and gets registered. +/// +/// The pipeline is strictly serial, so +/// the lock is always uncontended; `spin::Mutex` is there purely to give +/// `Arc>` `Send + Sync`. +pub(crate) struct HandleCollector(spin::Mutex>>); + +impl HandleCollector { + fn new() -> Self { + Self(spin::Mutex::new(HashMap::new())) + } + + fn capture(&self, handles: &HandleContainer>) { + let mut bag = self.0.lock(); + bag.clear(); + for id in handles.handle_ids() { + if let Some(h) = handles.get_handle_ref(id) { + bag.insert(*id, h.clone()); + } + } + } + + fn take(&self) -> HashMap> { + core::mem::take(&mut *self.0.lock()) + } +} + +/// Fusion input for autotuning. Thread the caller's `&mut Context` through the tuning +/// pipeline without `unsafe` by riding cubecl's `'a` lifetime parameter. +pub(crate) struct TuneInput<'a, R: Runtime, O> { + optimization: Arc, + state: TuneState<'a, R>, +} + +enum TuneState<'a, R: Runtime> { + Original { + context: &'a mut Context>, + /// Shared with any `Fork` spawned from this `Original`. `Drop` drains from here + /// if the cache-hit path never ran (e.g. wasm fallback succeeded on a fork). + new_handles: Arc>, + /// Set by [`TuneInput::execute`] when the winner ran on this context, which + /// suppresses the drain above. + executed: bool, + }, + /// Owned fork. `new_handles = None` is a benchmark-only sandbox (outputs discarded); + /// `Some(…)` is the wasm try-all path that promotes outputs back into the paired + /// `Original` on drop. + Fork { + context: Box>>, + new_handles: Option>>, + }, +} + +impl<'a, R: Runtime, O> TuneInput<'a, R, O> { + pub(crate) fn new(context: &'a mut Context>, optimization: O) -> Self { + Self { + optimization: Arc::new(optimization), + state: TuneState::Original { + context, + new_handles: Arc::new(HandleCollector::new()), + executed: false, + }, + } + } + + /// Fork the context into a non-tracking `Fork` for a benchmark trial. See + /// [`FusionInputGen`]. + fn for_benchmark(&self) -> Self { + Self { + optimization: self.optimization.clone(), + state: TuneState::Fork { + context: Box::new(self.context().fork()), + new_handles: None, + }, + } + } + + pub(crate) fn is_original(&self) -> bool { + matches!(self.state, TuneState::Original { .. }) + } + + /// Read-only access to the wrapped context. + pub(crate) fn context(&self) -> &Context> { + match &self.state { + TuneState::Original { context, .. } => context, + TuneState::Fork { context, .. } => context, + } + } + + pub(crate) fn tensors(&self) -> &HashMap { + &self.context().tensors + } + + pub(crate) fn handles(&self) -> &HandleContainer> { + &self.context().handles + } + + pub(crate) fn optimization(&self) -> &O { + &self.optimization + } + + /// Consume the input and run `f` with mutable access to the context and + /// optimization. + pub(crate) fn execute(mut self, f: F) -> T + where + F: FnOnce(&mut Context>, &O) -> T, + { + match &mut self.state { + TuneState::Original { + context, executed, .. + } => { + // Suppresses drop-time persistence; the real context was written directly. + *executed = true; + f(context, &self.optimization) + } + TuneState::Fork { context, .. } => f(context, &self.optimization), + } + } +} + +impl<'a, R: Runtime, O> Clone for TuneInput<'a, R, O> { + fn clone(&self) -> Self { + // `Original` clones come from the wasm fallback path (`operations.fastest(i) + // .execute(inputs.clone())`) and must track outputs. `Fork` clones inherit the + // source's tracking (benchmark forks stay non-tracking). + let new_handles = match &self.state { + TuneState::Original { new_handles, .. } => Some(new_handles.clone()), + TuneState::Fork { new_handles, .. } => new_handles.clone(), + }; + Self { + optimization: self.optimization.clone(), + state: TuneState::Fork { + context: Box::new(self.context().fork()), + new_handles, + }, + } + } +} + +impl<'a, R: Runtime, O> Drop for TuneInput<'a, R, O> { + fn drop(&mut self) { + match &mut self.state { + TuneState::Original { + context, + new_handles, + executed, + } => { + if *executed { + return; + } + // Cache-hit path never ran. Drain anything the most recent `Fork` + // deposited and register entries the real context doesn't already have + // (those are genuine new outputs, not inputs inherited via `fork`). + for (id, handle) in new_handles.take() { + if context.handles.get_handle_ref(&id).is_none() { + context.handles.register_handle(id, handle); + } + } + } + TuneState::Fork { + context, + new_handles, + } => { + if let Some(collector) = new_handles { + collector.capture(&context.handles); + } + } + } + } +} diff --git a/crates/burn-cubecl/Cargo.toml b/crates/burn-cubecl/Cargo.toml new file mode 100644 index 0000000..b9938c9 --- /dev/null +++ b/crates/burn-cubecl/Cargo.toml @@ -0,0 +1,92 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science"] +description = "Generic backend that can be compiled just-in-time to any shader language target" +documentation = "https://docs.rs/burn-cubecl" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "gpu"] +license.workspace = true +name = "burn-cubecl" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-cubecl" +version.workspace = true + +[lints] +workspace = true + +[features] +default = [ + "autotune", + "std", + "fusion", + "cubecl/default", + "burn-backend/default", + "burn-fusion?/default", + "burn-cubecl-fusion?/default", +] +std = [ + "cubecl/std", + "burn-backend/std", + "burn-fusion?/std", + "burn-cubecl-fusion?/std", +] +doc = ["default"] + +tracing = [ + "dep:tracing", + "cubecl/tracing", + "burn-std/tracing", + "burn-backend/tracing", + "burn-fusion?/tracing", + "burn-cubecl-fusion?/tracing", +] + +autotune = ["burn-cubecl-fusion?/autotune"] +autotune-checks = [ + "autotune", + "cubecl/autotune-checks", + "burn-cubecl-fusion?/autotune-checks", +] + +fusion = ["burn-fusion", "burn-cubecl-fusion"] +fusion-experimental = ["fusion"] + +template = [] + + +[dependencies] +burn-cubecl-fusion = { workspace = true, optional = true } +burn-fusion = { workspace = true, optional = true } +burn-ir = { workspace = true } +burn-std = { workspace = true } +burn-backend = { workspace = true, features = [ + "cubecl", +] } +cubecl = { workspace = true, features = ["stdlib"] } +cubek = { workspace = true, features = [ + "attention", + "matmul", + "convolution", + "interpolate", + "pool", + "reduce", + "random", + "quantization", + "stdlib", + "fft", +] } +tracing = { workspace = true, features = ["attributes"], optional = true } + +derive-new = { workspace = true } +log = { workspace = true } + +# Async +futures-lite = { workspace = true, features = ["std"] } + +# Template +serde = { workspace = true } +text_placeholder = { workspace = true, features = ["struct_context"] } + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-cubecl/LICENSE-APACHE b/crates/burn-cubecl/LICENSE-APACHE new file mode 120000 index 0000000..1cd601d --- /dev/null +++ b/crates/burn-cubecl/LICENSE-APACHE @@ -0,0 +1 @@ +../../LICENSE-APACHE \ No newline at end of file diff --git a/crates/burn-cubecl/LICENSE-MIT b/crates/burn-cubecl/LICENSE-MIT new file mode 120000 index 0000000..b2cfbdc --- /dev/null +++ b/crates/burn-cubecl/LICENSE-MIT @@ -0,0 +1 @@ +../../LICENSE-MIT \ No newline at end of file diff --git a/crates/burn-cubecl/README.md b/crates/burn-cubecl/README.md new file mode 100644 index 0000000..d1811f3 --- /dev/null +++ b/crates/burn-cubecl/README.md @@ -0,0 +1,3 @@ +# Burn CubeCL Backend + +Generic backend that can be compiled just-in-time (JIT) to any shader language target. diff --git a/crates/burn-cubecl/src/backend.rs b/crates/burn-cubecl/src/backend.rs new file mode 100644 index 0000000..238da91 --- /dev/null +++ b/crates/burn-cubecl/src/backend.rs @@ -0,0 +1,255 @@ +use crate::{CubeRuntime, tensor::CubeTensor}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{ + Backend, BackendGraph, BackendTypes, DTypeUsage, DTypeUsageSet, DeviceOps, ExecutionError, + TensorData, +}; +use burn_std::{BoolStore, DType}; +use cubecl::{ + features::{MmaConfig, TypeUsage}, + server::ComputeServer, +}; +use std::marker::PhantomData; + +#[cfg(not(feature = "fusion"))] +use burn_backend::tensor::{BoolTensor, FloatTensor, IntTensor, QuantizedTensor}; +#[cfg(not(feature = "fusion"))] +use burn_ir::{BackendIr, TensorHandle}; + +/// Turn a cubecl graph-capture error into a backend [`ExecutionError`]. +fn graph_err(err: impl core::fmt::Display) -> ExecutionError { + ExecutionError::WithContext { + reason: format!("{err}"), + } +} + +/// Generic tensor backend that can be compiled just-in-time to any shader runtime +#[derive(new)] +pub struct CubeBackend { + _runtime: PhantomData, +} + +impl BackendTypes for CubeBackend +where + R: CubeRuntime, + R::Server: ComputeServer, + R::Device: DeviceOps, +{ + type Device = R::Device; + + type FloatTensorPrimitive = CubeTensor; + type IntTensorPrimitive = CubeTensor; + type BoolTensorPrimitive = CubeTensor; + type QuantizedTensorPrimitive = CubeTensor; + + type GraphPrimitive = cubecl::client::Graph; +} + +impl Backend for CubeBackend +where + R: CubeRuntime, + R::Server: ComputeServer, + R::Device: DeviceOps, +{ + fn name(device: &Self::Device) -> String { + let client = R::client(device); + format!("cubecl<{}>", R::name(&client)) + } + + fn seed(_device: &Self::Device, seed: u64) { + cubek::random::seed(seed); + } + + fn ad_enabled(_device: &Self::Device) -> bool { + false + } + + fn sync(device: &Self::Device) -> Result<(), ExecutionError> { + let client = R::client(device); + futures_lite::future::block_on(client.sync()).map_err(|err| ExecutionError::WithContext { + reason: format!("{err}"), + }) + } + + fn graph_prepare(device: &Self::Device) -> Result<(), ExecutionError> { + let client = R::client(device); + client.graph_prepare().map_err(graph_err) + } + + fn graph_start_capture(device: &Self::Device) -> Result<(), ExecutionError> { + let client = R::client(device); + client.start_capture().map_err(graph_err) + } + + fn graph_stop_capture(device: &Self::Device) -> Result, ExecutionError> { + let client = R::client(device); + client.stop_capture().map_err(graph_err) + } + + unsafe fn graph_replay( + _device: &Self::Device, + graph: &BackendGraph, + ) -> Result<(), ExecutionError> { + // cubecl's `Graph::replay` is fire-and-forget: it enqueues the dispatch + // and returns immediately, so a replay failure is not reported here — it + // lands in the stream's error queue and surfaces on the next sync/flush. + // + // Safety: the buffer-liveness and stream-ordering obligations are the + // caller's, forwarded verbatim from this method's own contract. + unsafe { graph.replay() }; + Ok(()) + } + + fn memory_persistent_allocations< + Output: Send, + Input: Send, + Func: Fn(Input) -> Output + Send, + >( + device: &Self::Device, + input: Input, + func: Func, + ) -> Output { + let client = R::client(device); + client.memory_persistent_allocation(input, func).unwrap() + } + + fn memory_cleanup(device: &Self::Device) { + let client = R::client(device); + client.memory_cleanup(); + } + + fn staging<'a, Iter>(data: Iter, device: &Self::Device) + where + Iter: Iterator, + { + let client = R::client(device); + client.staging(data.map(|td| &mut td.bytes), false); + } + + fn supports_dtype(device: &Self::Device, dtype: DType) -> bool { + // Right now no cubecl backend actually works with native bool, even if + // the `TypeUsage` might indicate otherwise. + if let DType::Bool(BoolStore::Native) = dtype { + return false; + } + + let client = R::client(device); + + let type_usage = client.properties().type_usage(dtype_to_storage_type(dtype)); + // Same as `TypeUsage::all_scalar()`, but we make the usage explicit here + type_usage.is_superset( + TypeUsage::Buffer + | TypeUsage::Conversion + | TypeUsage::Arithmetic + | TypeUsage::DotProduct, + ) + } + + fn dtype_usage(device: &Self::Device, dtype: DType) -> DTypeUsageSet { + // Right now no cubecl backend actually works with native bool, even if + // the `TypeUsage` might indicate otherwise. + if let DType::Bool(BoolStore::Native) = dtype { + return DTypeUsageSet::empty(); + } + + let client = R::client(device); + + let props = client.properties(); + let storage = dtype_to_storage_type(dtype); + let usage = props.type_usage(storage); + + let mut out = DTypeUsageSet::new(); + + if usage.is_superset(TypeUsage::Buffer | TypeUsage::Conversion) { + out |= DTypeUsage::Storage; + } + + if usage.contains(TypeUsage::Arithmetic) { + out |= DTypeUsage::Arithmetic; + } + + let has_mma = |cfg: &MmaConfig| { + cfg.a_type == storage || cfg.b_type == storage || cfg.cd_type == storage + }; + if props.features.matmul.cmma.iter().any(has_mma) + || props.features.matmul.mma.iter().any(has_mma) + { + out |= DTypeUsage::Accelerated; + } + + out + } + + fn device_count(type_id: u16) -> usize { + let client = R::client(&Default::default()); + client.device_count(type_id) + } + + fn flush(device: &Self::Device) { + let client = R::client(device); + client.flush().unwrap(); + } +} + +impl core::fmt::Debug for CubeBackend { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("CubeCLBackend") + } +} + +impl Clone for CubeBackend { + fn clone(&self) -> Self { + Self::new() + } +} + +impl Default for CubeBackend { + fn default() -> Self { + Self::new() + } +} + +impl CubeRuntime for R +where + R::Device: DeviceOps, +{ + type CubeDevice = R::Device; + type CubeServer = R::Server; +} + +#[cfg(not(feature = "fusion"))] +impl BackendIr for CubeBackend { + type Handle = CubeTensor; + + fn float_tensor(handle: TensorHandle) -> FloatTensor { + handle.handle + } + + fn int_tensor(handle: TensorHandle) -> IntTensor { + handle.handle + } + + fn bool_tensor(handle: TensorHandle) -> BoolTensor { + handle.handle + } + + fn quantized_tensor(handle: TensorHandle) -> QuantizedTensor { + handle.handle + } + + fn float_tensor_handle(tensor: FloatTensor) -> Self::Handle { + tensor + } + + fn int_tensor_handle(tensor: IntTensor) -> Self::Handle { + tensor + } + + fn bool_tensor_handle(tensor: BoolTensor) -> Self::Handle { + tensor + } + + fn quantized_tensor_handle(tensor: QuantizedTensor) -> Self::Handle { + tensor + } +} diff --git a/crates/burn-cubecl/src/element.rs b/crates/burn-cubecl/src/element.rs new file mode 100644 index 0000000..c9a02b7 --- /dev/null +++ b/crates/burn-cubecl/src/element.rs @@ -0,0 +1,94 @@ +use burn_backend::{Element, bf16, f16}; +use cubecl::{ + CubeElement as CubeElem, flex32, + prelude::{Float, Int, Numeric}, +}; +use cubek::{ + matmul::definition::{MatmulPrecision, MatrixPrecision}, + reduce::ReducePrecision, +}; + +/// The base element trait for the jit backend. +pub trait CubeElement: Element + CubeElem + PartialEq + Numeric {} + +/// Element that can be used for matrix multiplication. Includes ints and floats. +pub trait MatmulElement: + CubeElement + MatmulPrecision> +{ +} + +/// The float element type for the jit backend. +pub trait FloatElement: MatmulElement + Float {} + +/// The int element type for the jit backend. +pub trait IntElement: + MatmulElement + Int + ReducePrecision +{ +} + +/// The element type for booleans for the jit backend. +pub trait BoolElement: CubeElement + Int { + /// The true value for the boolean element. + fn true_val() -> Self { + Self::from_int(1) + } + + /// The false value for the boolean element. + fn false_val() -> Self { + Self::from_int(0) + } + + /// New bool element from Rust bool. + fn new_bool(val: bool) -> Self { + match val { + true => Self::true_val(), + false => Self::false_val(), + } + } +} + +impl CubeElement for u64 {} +impl CubeElement for u32 {} +impl CubeElement for u16 {} +impl CubeElement for u8 {} +impl CubeElement for i64 {} +impl CubeElement for i32 {} +impl CubeElement for i16 {} +impl CubeElement for i8 {} +impl CubeElement for f64 {} +impl CubeElement for f32 {} +impl CubeElement for flex32 {} +impl CubeElement for f16 {} +impl CubeElement for bf16 {} + +impl FloatElement for f64 {} +impl FloatElement for f32 {} +impl FloatElement for flex32 {} +impl FloatElement for bf16 {} +impl FloatElement for f16 {} +impl IntElement for i64 {} +impl IntElement for i32 {} +impl IntElement for i16 {} +impl IntElement for i8 {} +impl IntElement for u64 {} +impl IntElement for u32 {} +impl IntElement for u16 {} +impl IntElement for u8 {} + +impl BoolElement for u8 {} +impl BoolElement for u32 {} + +impl MatmulElement for f64 {} +impl MatmulElement for f32 {} +impl MatmulElement for flex32 {} +impl MatmulElement for bf16 {} +impl MatmulElement for f16 {} + +impl MatmulElement for i64 {} +impl MatmulElement for i32 {} +impl MatmulElement for i16 {} +impl MatmulElement for i8 {} +impl MatmulElement for u64 {} +impl MatmulElement for u32 {} +impl MatmulElement for u16 {} +impl MatmulElement for u8 {} diff --git a/crates/burn-cubecl/src/fusion.rs b/crates/burn-cubecl/src/fusion.rs new file mode 100644 index 0000000..99ff13e --- /dev/null +++ b/crates/burn-cubecl/src/fusion.rs @@ -0,0 +1,192 @@ +use crate::{CubeBackend, CubeRuntime, kernel, tensor::CubeTensor}; +use burn_backend::tensor::{BoolTensor, FloatTensor, IntTensor, QuantizedTensor}; +use burn_backend::{DType, Shape}; +use burn_cubecl_fusion::optim::reduce::ReduceSettings; +use burn_cubecl_fusion::optim::reduce_broadcasted::ReduceBroadcastedFuser; +use burn_cubecl_fusion::{ + CubeFusionHandle, FallbackOperation, + optim::{ + CubeOptimization, CubeOptimizationState, + elemwise::{ElementWiseFuser, ElemwiseOptimization}, + matmul::{MatmulFuser, MatmulOptimization}, + reduce::{ReduceFuser, ReduceOptimization}, + reduce_broadcasted::ReduceBroadcastedOptimization, + }, +}; +use burn_fusion::UnfusedOp; +use burn_fusion::{ + FusionBackend, FusionRuntime, + stream::{Operation, OrderedExecution}, +}; +use burn_ir::{BackendIr, TensorHandle}; +use burn_std::Metadata; +use core::marker::PhantomData; +use std::sync::Arc; + +impl burn_fusion::Optimization> for CubeOptimization +where + R: CubeRuntime, +{ + fn execute( + &mut self, + context: &mut burn_fusion::stream::Context< + as FusionRuntime>::FusionHandle, + >, + execution: &OrderedExecution>, + ) { + match self { + Self::ElementWise(op) => op.execute(context), + Self::Matmul(op) => op.execute(context, |index| { + let operation = execution.operation_within_optimization(index); + Box::new(FallbackOperationWrapper::new(operation)) + }), + Self::Reduce(op) => op.execute(context, |index| { + let operation = execution.operation_within_optimization(index); + Box::new(FallbackOperationWrapper::new(operation)) + }), + Self::ReduceBroadcasted(op) => op.execute(context, |index| { + let operation = execution.operation_within_optimization(index); + Box::new(FallbackOperationWrapper::new(operation)) + }), + } + } + + fn to_state(&self) -> CubeOptimizationState { + self.to_opt_state() + } + + fn from_state(device: &R::Device, state: CubeOptimizationState) -> Self { + match state { + CubeOptimizationState::ElementWise(state) => { + Self::ElementWise(ElemwiseOptimization::from_state(device, state)) + } + CubeOptimizationState::Matmul(state) => { + Self::Matmul(MatmulOptimization::from_state(device, state)) + } + CubeOptimizationState::Reduce(state) => { + Self::Reduce(ReduceOptimization::from_state(device, state)) + } + CubeOptimizationState::ReduceBroadcasted(state) => { + Self::ReduceBroadcasted(ReduceBroadcastedOptimization::from_state(device, state)) + } + } + } +} + +struct FallbackOperationWrapper { + operation: O, +} + +impl FallbackOperationWrapper { + fn new(op: O) -> Self { + Self { operation: op } + } +} + +impl FallbackOperation + for FallbackOperationWrapper>>> +{ + fn run(&self, context: &mut burn_fusion::stream::Context>) { + self.operation.as_ref().execute(&mut context.handles); + } +} + +impl FallbackOperation + for FallbackOperationWrapper>> +{ + fn run(&self, context: &mut burn_fusion::stream::Context>) { + self.operation.execute(&mut context.handles); + } +} + +impl BackendIr for CubeBackend { + type Handle = CubeFusionHandle; + + fn float_tensor(handle: TensorHandle) -> FloatTensor { + into_tensor(handle.handle, handle.shape) + } + + fn int_tensor(handle: TensorHandle) -> IntTensor { + into_tensor(handle.handle, handle.shape) + } + + fn bool_tensor(handle: TensorHandle) -> BoolTensor { + into_tensor(handle.handle, handle.shape) + } + + fn quantized_tensor(handle: TensorHandle) -> QuantizedTensor { + into_tensor(handle.handle, handle.shape) + } + + fn float_tensor_handle(tensor: FloatTensor) -> Self::Handle { + tensor.into() + } + + fn int_tensor_handle(tensor: IntTensor) -> Self::Handle { + tensor.into() + } + + fn bool_tensor_handle(tensor: BoolTensor) -> Self::Handle { + tensor.into() + } + + fn quantized_tensor_handle(tensor: QuantizedTensor) -> Self::Handle { + tensor.into() + } +} + +impl FusionRuntime for FusionCubeRuntime { + type OptimizationState = CubeOptimizationState; + type Optimization = CubeOptimization; + type FusionHandle = CubeFusionHandle; + type FusionDevice = R::CubeDevice; + + fn fusers(device: R::Device) -> Vec>> { + vec![ + Box::new(ElementWiseFuser::new(device.clone())), + Box::new(MatmulFuser::new(device.clone())), + Box::new(ReduceFuser::new(device.clone(), ReduceSettings::Always)), + Box::new(ReduceBroadcastedFuser::new(device.clone())), + ] + } +} + +/// Fusion runtime for JIT runtimes. +#[derive(Debug)] +pub struct FusionCubeRuntime { + _b: PhantomData, +} + +impl FusionBackend for CubeBackend { + type FusionRuntime = FusionCubeRuntime; + + type FullPrecisionBackend = CubeBackend; + + fn cast_float(tensor: FloatTensor, dtype: DType) -> Self::Handle { + kernel::cast(tensor, dtype).into() + } +} + +fn into_tensor(handle: CubeFusionHandle, shape: Shape) -> CubeTensor { + CubeTensor { + client: handle.client.clone(), + handle: handle.handle.clone(), + device: handle.device.clone(), + meta: Box::new(Metadata::new(shape, handle.strides.clone())), + dtype: handle.dtype, + qparams: handle.qparams.clone(), + } +} + +impl From> for CubeFusionHandle { + fn from(value: CubeTensor) -> Self { + Self { + client: value.client.clone(), + handle: value.handle.clone(), + device: value.device.clone(), + strides: value.meta.strides.clone(), + dtype: value.dtype, + qparams: value.qparams.clone(), + } + } +} diff --git a/crates/burn-cubecl/src/kernel/attention/base.rs b/crates/burn-cubecl/src/kernel/attention/base.rs new file mode 100644 index 0000000..627f56d --- /dev/null +++ b/crates/burn-cubecl/src/kernel/attention/base.rs @@ -0,0 +1,145 @@ +use crate::{ + CubeBackend, CubeRuntime, kernel::attention::attention_autotune, + ops::numeric::empty_device_dtype, tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{ + DType, Shape, + ops::{AttentionModuleOptions, attention::attention_fallback}, +}; +use cubek::attention::forward::{ + definition::{ + AccumulatorPrecision, AttentionGlobalTypes, AttentionOptions, AttentionSetupError, + }, + launch, + routines::blackbox_accelerated::BlackboxAcceleratedStrategy, +}; + +#[derive(Debug)] +/// Strategy used to select which attention implementation to run. +pub enum AttentionStrategy { + /// Flash Attention using accelerated inner matmuls. + FlashBlackboxAccelerated(BlackboxAcceleratedStrategy), + + /// Flash Attention using unit inner matmuls. + FlashUnit, + + /// Fallback implementation using multiple separate kernels. + Fallback, + + /// Automatically benchmark and select the best strategy at runtime. + #[cfg(feature = "autotune")] + Autotune, +} + +impl Default for AttentionStrategy { + fn default() -> Self { + // if autotune is enabled, default to autotune + #[cfg(feature = "autotune")] + return AttentionStrategy::Autotune; + + // if autotune is disabled, default to fallback to make sure it runs + #[cfg(not(feature = "autotune"))] + AttentionStrategy::Fallback + } +} + +#[allow(clippy::too_many_arguments)] +/// Launch an attention kernel with given strategy +pub fn attention( + query: CubeTensor, + key: CubeTensor, + value: CubeTensor, + mask: Option>, + attn_bias: Option>, + options: AttentionModuleOptions, + strategy: AttentionStrategy, +) -> Result, AttentionSetupError> { + match strategy { + AttentionStrategy::FlashBlackboxAccelerated(strategy) => flash_attention( + query, + key, + value, + mask, + attn_bias, + options, + launch::Strategy::BlackboxAccelerated(launch::BlueprintStrategy::Inferred(strategy)), + ), + AttentionStrategy::FlashUnit => flash_attention( + query, + key, + value, + mask, + attn_bias, + options, + launch::Strategy::Unit(launch::BlueprintStrategy::Inferred(())), + ), + AttentionStrategy::Fallback => Ok(attention_fallback::>( + query, key, value, mask, attn_bias, options, + )), + #[cfg(feature = "autotune")] + AttentionStrategy::Autotune => Ok(attention_autotune( + query, key, value, mask, attn_bias, options, + )), + } +} + +#[allow(clippy::too_many_arguments)] +/// Launch a flash attention kernel +pub fn flash_attention( + query: CubeTensor, + key: CubeTensor, + value: CubeTensor, + mask: Option>, + _attn_bias: Option>, + options: AttentionModuleOptions, + strategy: launch::Strategy, +) -> Result, AttentionSetupError> { + let client = query.client.clone(); + let out = init_attention_output(&query, &value); + + let dtypes = AttentionGlobalTypes { + query: dtype_to_storage_type(query.dtype), + key: dtype_to_storage_type(key.dtype), + value: dtype_to_storage_type(value.dtype), + mask: dtype_to_storage_type(mask.as_ref().map(|m| m.dtype).unwrap_or(DType::U8)), + out: dtype_to_storage_type(out.dtype), + }; + + launch::launch_ref::( + strategy, + &client, + query.binding(), + key.binding(), + value.binding(), + mask.map(|mask| mask.binding()), + out.clone().binding(), + &dtypes, + AttentionOptions { + causal: options.is_causal, + accumulator_precision: AccumulatorPrecision::Strict(cubecl::ir::StorageType::Scalar( + cubecl::ir::ElemType::Float(cubecl::ir::FloatKind::F32), + )), + }, + )?; + + Ok(out) +} + +pub(crate) fn init_attention_output( + query: &CubeTensor, + value: &CubeTensor, +) -> CubeTensor { + let num_batches = query.meta.shape[0]; + let num_heads = query.meta.shape[1]; + let seq_q = query.meta.shape[2]; + let val_dim = value.meta.shape[3]; + let out_shape = Shape::new([num_batches, num_heads, seq_q, val_dim]); + + empty_device_dtype::( + query.client.clone(), + query.device.clone(), + out_shape, + query.dtype, + ) +} diff --git a/crates/burn-cubecl/src/kernel/attention/mod.rs b/crates/burn-cubecl/src/kernel/attention/mod.rs new file mode 100644 index 0000000..8ff38a9 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/attention/mod.rs @@ -0,0 +1,5 @@ +mod base; +mod tune; + +pub use base::*; +pub use tune::*; diff --git a/crates/burn-cubecl/src/kernel/attention/tune.rs b/crates/burn-cubecl/src/kernel/attention/tune.rs new file mode 100644 index 0000000..10fbfc8 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/attention/tune.rs @@ -0,0 +1,178 @@ +use crate::{ + CubeRuntime, CubeTuneId, + kernel::attention::{AttentionStrategy, attention}, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_elem_type; +use burn_backend::ops::AttentionModuleOptions; +use cubecl::tune::{LocalTuner, Tunable, TunableSet, TuneGroup, local_tuner}; +use cubek::attention::forward::{ + launch::AttentionAutotuneKey, routines::blackbox_accelerated::BlackboxAcceleratedStrategy, +}; + +/// Executes autotune on attention operations +pub fn attention_autotune( + query: CubeTensor, + key: CubeTensor, + value: CubeTensor, + mask: Option>, + attn_bias: Option>, + options: AttentionModuleOptions, +) -> CubeTensor { + let client = query.client.clone(); + + static TUNER: LocalTuner = local_tuner!(); + + let tunables = TUNER.init(|| { + const PRIORITY_MAX: i8 = 3; + const PRIORITY_MIN: i8 = 0; + + let flash_attention = + TuneGroup::::new("flash_attention", |_key| PRIORITY_MAX); + + let fallback = TuneGroup::::new("fallback", |key| { + if key.seq_q > 4096 { + PRIORITY_MIN + } else { + PRIORITY_MAX + } + }); + + let mut set = TunableSet::new(create_key::, input_gen::); + + // First entry should always work, since it is considered the fallback. + set = set.with( + Tunable::new( + "fallback", + |(query, key, value, mask, attn_bias, options)| { + attention::( + query, + key, + value, + mask, + attn_bias, + options, + AttentionStrategy::Fallback, + ) + .map_err(|err| std::format!("{err:?}")) + }, + ) + .group(&fallback, |_key| PRIORITY_MAX), + ); + + let seq_q = 1; + let seq_kv = 1; + for num_planes in [2, 4, 8] { + let name = format!("blackbox_accelerated_{num_planes}_planes_p_{seq_q}-{seq_kv}"); + set = set.with( + Tunable::new( + &name, + move |(query, key, value, mask, attn_bias, options)| { + attention::( + query, + key, + value, + mask, + attn_bias, + options, + AttentionStrategy::FlashBlackboxAccelerated( + BlackboxAcceleratedStrategy { + num_planes, + seq_q, + seq_kv, + }, + ), + ) + .map_err(|err| std::format!("{err:?}")) + }, + ) + .group(&flash_attention, |_key| PRIORITY_MAX), + ); + } + + set = set.with( + Tunable::new("unit", |(query, key, value, mask, attn_bias, options)| { + attention::( + query, + key, + value, + mask, + attn_bias, + options, + AttentionStrategy::FlashUnit, + ) + .map_err(|err| std::format!("{err:?}")) + }) + .group(&flash_attention, |_key| PRIORITY_MIN), + ); + + set + }); + + TUNER.execute( + &CubeTuneId::new(&client, &query.device), + &client, + tunables, + (query, key, value, mask, attn_bias, options), + ) +} + +#[allow(clippy::type_complexity)] +fn create_key( + (query, key, value, mask, _attn_bias, _options): &( + CubeTensor, + CubeTensor, + CubeTensor, + Option>, + Option>, + AttentionModuleOptions, + ), +) -> AttentionAutotuneKey { + let total_batches = query.meta.shape[0] * query.meta.shape[1]; + let seq_q = query.meta.shape[2]; + let head_dim = query.meta.shape[3]; + let seq_kv = value.meta.shape[2]; + let val_dim = value.meta.shape[3]; + + AttentionAutotuneKey::generate( + dtype_to_elem_type(query.dtype), + dtype_to_elem_type(key.dtype), + dtype_to_elem_type(value.dtype), + dtype_to_elem_type(query.dtype), + total_batches, + seq_q, + head_dim, + seq_kv, + val_dim, + mask.is_some(), + ) +} + +#[allow(clippy::type_complexity)] +fn input_gen( + _key: &AttentionAutotuneKey, + (query, key, value, mask, attn_bias, options): &( + CubeTensor, + CubeTensor, + CubeTensor, + Option>, + Option>, + AttentionModuleOptions, + ), +) -> ( + CubeTensor, + CubeTensor, + CubeTensor, + Option>, + Option>, + AttentionModuleOptions, +) { + ( + query.clone(), + key.clone(), + value.clone(), + mask.clone(), + attn_bias.clone(), + *options, + ) +} diff --git a/crates/burn-cubecl/src/kernel/binary.rs b/crates/burn-cubecl/src/kernel/binary.rs new file mode 100644 index 0000000..71c746b --- /dev/null +++ b/crates/burn-cubecl/src/kernel/binary.rs @@ -0,0 +1,346 @@ +use crate::{ + CubeRuntime, + kernel::utils::{address_type, broadcast_shape}, + ops::{max_vector_size, numeric::empty_device_dtype}, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{TensorMetadata, bf16, f16}; +use cubecl::{ + calculate_cube_count_elemwise, intrinsic, + prelude::*, + std::tensor::layout::linear::{LinearView, LinearViewMut}, +}; + +pub(crate) trait BinaryOpFamily: Send + Sync + 'static { + type BinaryOp: BinaryOp; +} + +#[cube] +pub(crate) trait BinaryOp: 'static + Send + Sync { + /// Execute a binary operation. + fn execute(lhs: Vector, rhs: Vector) -> Vector; +} + +pub(crate) struct AddOp; +pub(crate) struct SubOp; +pub(crate) struct MulOp; +pub(crate) struct DivOp; +pub(crate) struct RemainderOp; +pub(crate) struct AndOp; +pub(crate) struct OrOp; +pub(crate) struct PowOp; +pub(crate) struct AssignOp; +pub(crate) struct BinaryMinOp; +pub(crate) struct BinaryMaxOp; + +impl BinaryOpFamily for AddOp { + type BinaryOp = Self; +} + +impl BinaryOpFamily for SubOp { + type BinaryOp = Self; +} + +impl BinaryOpFamily for MulOp { + type BinaryOp = Self; +} + +impl BinaryOpFamily for DivOp { + type BinaryOp = Self; +} + +impl BinaryOpFamily for RemainderOp { + type BinaryOp = Self; +} + +impl BinaryOpFamily for PowOp { + type BinaryOp = Self; +} + +impl BinaryOpFamily for AndOp { + type BinaryOp = Self; +} + +impl BinaryOpFamily for OrOp { + type BinaryOp = Self; +} + +impl BinaryOpFamily for AssignOp { + type BinaryOp = Self; +} + +impl BinaryOpFamily for BinaryMinOp { + type BinaryOp = Self; +} + +impl BinaryOpFamily for BinaryMaxOp { + type BinaryOp = Self; +} + +#[cube] +impl BinaryOp for AddOp { + fn execute(lhs: Vector, rhs: Vector) -> Vector { + lhs + rhs + } +} + +#[cube] +impl BinaryOp for SubOp { + fn execute(lhs: Vector, rhs: Vector) -> Vector { + lhs - rhs + } +} + +#[cube] +impl BinaryOp for MulOp { + fn execute(lhs: Vector, rhs: Vector) -> Vector { + lhs * rhs + } +} + +#[cube] +impl BinaryOp for DivOp { + fn execute(lhs: Vector, rhs: Vector) -> Vector { + lhs / rhs + } +} + +#[cube] +impl BinaryOp for RemainderOp { + fn execute(lhs: Vector, rhs: Vector) -> Vector { + Vector::mod_floor(lhs, rhs) + } +} + +#[cube] +impl BinaryOp for PowOp { + #[allow(unused)] + fn execute(lhs: Vector, rhs: Vector) -> Vector { + intrinsic!(|scope| { + let elem = T::__expand_as_type(scope).elem_type(); + + if let cubecl::ir::ElemType::Float(kind) = elem { + match kind { + cubecl::ir::FloatKind::F16 => { + let lhs = as Cast>::__expand_cast_from(scope, lhs); + let rhs = as Cast>::__expand_cast_from(scope, rhs); + let out = Vector::__expand_powf(scope, lhs, rhs); + return as Cast>::__expand_cast_from(scope, out); + } + cubecl::ir::FloatKind::BF16 => { + let lhs = as Cast>::__expand_cast_from(scope, lhs); + let rhs = as Cast>::__expand_cast_from(scope, rhs); + let out = Vector::__expand_powf(scope, lhs, rhs); + return as Cast>::__expand_cast_from(scope, out); + } + cubecl::ir::FloatKind::F64 => { + let lhs = as Cast>::__expand_cast_from(scope, lhs); + let rhs = as Cast>::__expand_cast_from(scope, rhs); + let out = Vector::__expand_powf(scope, lhs, rhs); + return as Cast>::__expand_cast_from(scope, out); + } + _ => {} + } + }; + + let lhs = as Cast>::__expand_cast_from(scope, lhs); + let rhs = as Cast>::__expand_cast_from(scope, rhs); + let out = Vector::__expand_powf(scope, lhs, rhs); + return as Cast>::__expand_cast_from(scope, out); + }) + } +} + +#[cube] +impl BinaryOp for AndOp { + fn execute(lhs: Vector, rhs: Vector) -> Vector { + Vector::cast_from( + Vector::::cast_from(lhs).vec_and(Vector::::cast_from(rhs)), + ) + } +} + +#[cube] +impl BinaryOp for OrOp { + fn execute(lhs: Vector, rhs: Vector) -> Vector { + Vector::cast_from(Vector::::cast_from(lhs).or(Vector::::cast_from(rhs))) + } +} + +#[cube] +impl BinaryOp for AssignOp { + fn execute(_lhs: Vector, rhs: Vector) -> Vector { + rhs + } +} + +#[cube] +impl BinaryOp for BinaryMinOp { + fn execute(lhs: Vector, rhs: Vector) -> Vector { + clamp_max(lhs, rhs) + } +} + +#[cube] +impl BinaryOp for BinaryMaxOp { + fn execute(lhs: Vector, rhs: Vector) -> Vector { + clamp_min(lhs, rhs) + } +} + +#[cube(launch_unchecked, address_type = "dynamic")] +pub(crate) fn kernel_scalar_binop( + input: LinearView<'_, Vector>, + scalar: InputScalar, + mut output: LinearViewMut<'_, Vector>, + #[define(C)] _dtype: StorageType, +) { + if !output.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + output.write( + ABSOLUTE_POS, + O::BinaryOp::::execute(input.read(ABSOLUTE_POS), Vector::new(scalar.get::())), + ); +} + +#[cube(launch_unchecked, address_type = "dynamic")] +pub(crate) fn kernel_binop( + lhs: LinearView<'_, Vector>, + rhs: LinearView<'_, Vector>, + mut out: LinearViewMut<'_, Vector>, + #[define(C)] _dtype: StorageType, +) { + if !out.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + out.write( + ABSOLUTE_POS, + O::BinaryOp::::execute(lhs.read(ABSOLUTE_POS), rhs.read(ABSOLUTE_POS)), + ); +} + +pub(crate) fn launch_binop( + lhs: CubeTensor, + rhs: CubeTensor, +) -> CubeTensor { + let vector_size_lhs = max_vector_size(&lhs); + let vector_size_rhs = max_vector_size(&rhs); + let vector_size = Ord::min(vector_size_lhs, vector_size_rhs); + + let shape_out = broadcast_shape(&[&lhs, &rhs]); + let dtype = lhs.dtype; + + let client = lhs.client.clone(); + let num_elems = shape_out.num_elements(); + let working_units = num_elems / vector_size as usize; + + let cube_dim = CubeDim::new(&lhs.client, working_units); + let cube_count = calculate_cube_count_elemwise(&lhs.client, working_units, cube_dim); + + unsafe { + if lhs.can_mut_broadcast(&rhs) { + kernel_binop::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(lhs, rhs), + vector_size, + lhs.clone().into_linear_view(), + rhs.into_linear_view_like(&lhs), + lhs.as_linear_view_alias(0), + dtype_to_storage_type(dtype), + ); + + lhs + } else if rhs.can_mut_broadcast(&lhs) { + kernel_binop::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(lhs, rhs), + vector_size, + lhs.into_linear_view_like(&rhs), + rhs.clone().into_linear_view(), + rhs.as_linear_view_alias(1), + dtype_to_storage_type(dtype), + ); + + rhs + } else { + let output = + empty_device_dtype(lhs.client.clone(), lhs.device.clone(), shape_out, dtype); + + kernel_binop::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(lhs, rhs, output), + vector_size, + lhs.into_linear_view_like(&output), + rhs.into_linear_view_like(&output), + output.clone().into_linear_view(), + dtype_to_storage_type(dtype), + ); + + output + } + } +} + +pub(crate) fn launch_scalar_binop( + tensor: CubeTensor, + scalar: InputScalar, +) -> CubeTensor { + // Vectorization is only enabled when the last dimension is contiguous. + let vector_size = max_vector_size(&tensor); + let client = tensor.client.clone(); + let num_elems = tensor.meta.num_elements(); + let dtype = tensor.dtype; + + let working_units = num_elems / vector_size as usize; + let cube_dim = CubeDim::new(&tensor.client, working_units); + let cube_count = calculate_cube_count_elemwise(&tensor.client, working_units, cube_dim); + + unsafe { + if tensor.can_mut() && tensor.is_nonoverlapping() { + kernel_scalar_binop::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(tensor), + vector_size, + tensor.clone().into_linear_view(), + scalar, + tensor.as_linear_view_alias(0), + dtype_to_storage_type(dtype), + ); + + tensor + } else { + let output = empty_device_dtype( + tensor.client.clone(), + tensor.device.clone(), + tensor.shape(), + dtype, + ); + + kernel_scalar_binop::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(tensor, output), + vector_size, + tensor.into_linear_view(), + scalar, + output.clone().into_linear_view(), + dtype_to_storage_type(dtype), + ); + + output + } + } +} diff --git a/crates/burn-cubecl/src/kernel/binary_float.rs b/crates/burn-cubecl/src/kernel/binary_float.rs new file mode 100644 index 0000000..813d118 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/binary_float.rs @@ -0,0 +1,124 @@ +use crate::{ + CubeRuntime, + kernel::utils::{address_type, broadcast_shape}, + ops::{max_vector_size, numeric::empty_device_dtype}, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::{ + calculate_cube_count_elemwise, + prelude::*, + std::tensor::layout::linear::{LinearView, LinearViewMut}, +}; + +pub(crate) trait BinaryOpFloatFamily: Send + Sync + 'static { + type BinaryOp: BinaryOpFloat; +} + +#[cube] +pub(crate) trait BinaryOpFloat: 'static + Send + Sync { + /// Execute a binary operation. + fn execute(lhs: Vector, rhs: Vector) -> Vector; +} + +pub(crate) struct ArcTan2Op; + +impl BinaryOpFloatFamily for ArcTan2Op { + type BinaryOp = Self; +} + +#[cube] +impl BinaryOpFloat for ArcTan2Op { + fn execute(lhs: Vector, rhs: Vector) -> Vector { + Vector::atan2(lhs, rhs) + } +} + +#[cube(launch_unchecked, address_type = "dynamic")] +pub(crate) fn kernel_binop( + lhs: LinearView<'_, Vector>, + rhs: LinearView<'_, Vector>, + mut out: LinearViewMut<'_, Vector>, + #[define(C)] _dtype: StorageType, +) { + if !out.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + let res = O::BinaryOp::::execute(lhs.read(ABSOLUTE_POS), rhs.read(ABSOLUTE_POS)); + + out.write(ABSOLUTE_POS, res) +} + +pub(crate) fn launch_binop_float( + lhs: CubeTensor, + rhs: CubeTensor, +) -> CubeTensor { + let vector_size_lhs = max_vector_size(&lhs); + let vector_size_rhs = max_vector_size(&rhs); + let vector_size = Ord::min(vector_size_lhs, vector_size_rhs); + + let shape_out = broadcast_shape(&[&lhs, &rhs]); + let dtype = lhs.dtype; + + let client = lhs.client.clone(); + let num_elems = shape_out.num_elements(); + let working_units = num_elems / vector_size as usize; + + let cube_dim = CubeDim::new(&lhs.client, working_units); + let cube_count = calculate_cube_count_elemwise(&lhs.client, working_units, cube_dim); + + unsafe { + if lhs.can_mut_broadcast(&rhs) { + kernel_binop::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(lhs, rhs), + vector_size, + lhs.clone().into_linear_view(), + rhs.clone().into_linear_view_like(&lhs), + lhs.as_linear_view_alias(0), + dtype_to_storage_type(dtype), + ); + + lhs + } else if rhs.can_mut_broadcast(&lhs) { + kernel_binop::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(lhs, rhs), + vector_size, + lhs.into_linear_view_like(&rhs), + rhs.clone().into_linear_view(), + rhs.as_linear_view_alias(1), + dtype_to_storage_type(dtype), + ); + + rhs + } else { + let output = + empty_device_dtype(lhs.client.clone(), lhs.device.clone(), shape_out, dtype); + + kernel_binop::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(lhs, rhs, output), + vector_size, + lhs.into_linear_view_like(&output), + rhs.into_linear_view_like(&output), + output.clone().into_linear_view(), + dtype_to_storage_type(dtype), + ); + + output + } + } +} + +/// Calculate the four-quadrant inverse tangent of `lhs / rhs`. +pub fn atan2(lhs: CubeTensor, rhs: CubeTensor) -> CubeTensor { + launch_binop_float::(lhs, rhs) +} diff --git a/crates/burn-cubecl/src/kernel/binary_int.rs b/crates/burn-cubecl/src/kernel/binary_int.rs new file mode 100644 index 0000000..2fed09e --- /dev/null +++ b/crates/burn-cubecl/src/kernel/binary_int.rs @@ -0,0 +1,238 @@ +use crate::{ + CubeRuntime, + kernel::utils::{address_type, broadcast_shape}, + ops::{max_vector_size, numeric::empty_device_dtype}, + tensor::CubeTensor, +}; +use burn_backend::TensorMetadata; +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::{ + calculate_cube_count_elemwise, + prelude::*, + std::tensor::layout::linear::{LinearView, LinearViewMut}, +}; + +pub(crate) trait BinaryOpIntFamily: Send + Sync + 'static { + type BinaryOp: BinaryOpInt; +} + +#[cube] +pub(crate) trait BinaryOpInt: 'static + Send + Sync { + /// Execute a binary operation. + fn execute(lhs: Vector, rhs: Vector) -> Vector; +} + +pub(crate) struct BitwiseAndOp; +pub(crate) struct BitwiseOrOp; +pub(crate) struct BitwiseXorOp; +pub(crate) struct BitwiseShrOp; +pub(crate) struct BitwiseShlOp; + +impl BinaryOpIntFamily for BitwiseAndOp { + type BinaryOp = Self; +} + +impl BinaryOpIntFamily for BitwiseOrOp { + type BinaryOp = Self; +} + +impl BinaryOpIntFamily for BitwiseXorOp { + type BinaryOp = Self; +} + +impl BinaryOpIntFamily for BitwiseShrOp { + type BinaryOp = Self; +} + +impl BinaryOpIntFamily for BitwiseShlOp { + type BinaryOp = Self; +} + +#[cube] +impl BinaryOpInt for BitwiseAndOp { + fn execute(lhs: Vector, rhs: Vector) -> Vector { + lhs & rhs + } +} + +#[cube] +impl BinaryOpInt for BitwiseOrOp { + fn execute(lhs: Vector, rhs: Vector) -> Vector { + lhs | rhs + } +} + +#[cube] +impl BinaryOpInt for BitwiseXorOp { + fn execute(lhs: Vector, rhs: Vector) -> Vector { + lhs ^ rhs + } +} + +#[cube] +impl BinaryOpInt for BitwiseShrOp { + fn execute(lhs: Vector, rhs: Vector) -> Vector { + lhs >> rhs + } +} + +#[cube] +impl BinaryOpInt for BitwiseShlOp { + fn execute(lhs: Vector, rhs: Vector) -> Vector { + lhs << rhs + } +} + +#[cube(launch_unchecked, address_type = "dynamic")] +pub(crate) fn kernel_scalar_binop_int( + input: LinearView<'_, Vector>, + scalar: InputScalar, + mut output: LinearViewMut<'_, Vector>, + #[define(C)] _dtype: StorageType, +) { + if !output.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + output.write( + ABSOLUTE_POS, + O::BinaryOp::::execute(input.read(ABSOLUTE_POS), Vector::new(scalar.get::())), + ); +} + +#[cube(launch_unchecked, address_type = "dynamic")] +pub(crate) fn kernel_binop_int( + lhs: LinearView<'_, Vector>, + rhs: LinearView<'_, Vector>, + mut out: LinearViewMut<'_, Vector>, + #[define(C)] _dtype: StorageType, +) { + if !out.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + out.write( + ABSOLUTE_POS, + O::BinaryOp::::execute(lhs.read(ABSOLUTE_POS), rhs.read(ABSOLUTE_POS)), + ); +} + +pub(crate) fn launch_binop_int( + lhs: CubeTensor, + rhs: CubeTensor, +) -> CubeTensor { + let vector_size_lhs = max_vector_size(&lhs); + let vector_size_rhs = max_vector_size(&rhs); + let vector_size = Ord::min(vector_size_lhs, vector_size_rhs); + + let shape_out = broadcast_shape(&[&lhs, &rhs]); + + let client = lhs.client.clone(); + let num_elems = shape_out.num_elements(); + + let working_units = num_elems / vector_size as usize; + let cube_dim = CubeDim::new(&lhs.client, working_units); + let cube_count = calculate_cube_count_elemwise(&lhs.client, working_units, cube_dim); + let dtype = lhs.dtype; + + unsafe { + if lhs.can_mut_broadcast(&rhs) { + kernel_binop_int::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(lhs, rhs), + vector_size, + lhs.clone().into_linear_view(), + rhs.into_linear_view_like(&lhs), + lhs.as_linear_view_alias(0), + dtype_to_storage_type(dtype), + ); + + lhs + } else if rhs.can_mut_broadcast(&lhs) { + kernel_binop_int::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(lhs, rhs), + vector_size, + lhs.into_linear_view_like(&rhs), + rhs.clone().into_linear_view(), + rhs.as_linear_view_alias(1), + dtype_to_storage_type(dtype), + ); + + rhs + } else { + let output = + empty_device_dtype(lhs.client.clone(), lhs.device.clone(), shape_out, lhs.dtype); + + kernel_binop_int::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(lhs, rhs, output), + vector_size, + lhs.into_linear_view_like(&output), + rhs.into_linear_view_like(&output), + output.clone().into_linear_view(), + dtype_to_storage_type(dtype), + ); + + output + } + } +} + +pub(crate) fn launch_scalar_binop_int( + tensor: CubeTensor, + scalar: InputScalar, +) -> CubeTensor { + let vector_size = max_vector_size(&tensor); + let client = tensor.client.clone(); + let num_elems = tensor.meta.shape.num_elements(); + + let working_units = num_elems / vector_size as usize; + let cube_dim = CubeDim::new(&tensor.client, working_units); + let cube_count = calculate_cube_count_elemwise(&tensor.client, working_units, cube_dim); + + unsafe { + if tensor.can_mut() && tensor.is_nonoverlapping() { + kernel_scalar_binop_int::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(tensor), + vector_size, + tensor.clone().into_linear_view(), + scalar, + tensor.as_linear_view_alias(0), + dtype_to_storage_type(tensor.dtype), + ); + + tensor + } else { + let output = empty_device_dtype( + tensor.client.clone(), + tensor.device.clone(), + tensor.shape(), + tensor.dtype, + ); + + kernel_scalar_binop_int::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(tensor, output), + vector_size, + tensor.into_linear_view(), + scalar, + output.clone().into_linear_view(), + dtype_to_storage_type(output.dtype), + ); + + output + } + } +} diff --git a/crates/burn-cubecl/src/kernel/cast/base.rs b/crates/burn-cubecl/src/kernel/cast/base.rs new file mode 100644 index 0000000..3cb9176 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/cast/base.rs @@ -0,0 +1,74 @@ +use crate::{ + CubeRuntime, + kernel::utils::address_type, + ops::{max_vector_size, numeric::empty_device_dtype}, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{DType, TensorMetadata}; +use cubecl::std::tensor::layout::linear::{LinearView, LinearViewMut}; +use cubecl::{calculate_cube_count_elemwise, prelude::*}; + +#[cube(launch, address_type = "dynamic")] +pub(crate) fn cast_element( + input: LinearView<'_, Vector>, + mut output: LinearViewMut<'_, Vector>, + #[define(I, O)] _dtypes: [StorageType; 2], +) { + if !output.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + output.write(ABSOLUTE_POS, Vector::cast_from(input.read(ABSOLUTE_POS))); +} + +/// Cast a tensor to the given element type. +/// +/// Note: When input element is semantically a boolean, prefer bool_cast function. +pub fn cast(input: CubeTensor, dtype: DType) -> CubeTensor { + let dtype_output = match dtype { + DType::Flex32 => DType::F32, + _ => dtype, + }; + let dtype_input = match input.dtype { + DType::Flex32 => DType::F32, + _ => input.dtype, + }; + + if dtype_input == dtype_output { + return input; + } + + let client = input.client.clone(); + + let vector_size = max_vector_size(&input); + + let num_elems: usize = input.meta.num_elements(); + + let working_units = num_elems / vector_size as usize; + let cube_dim = CubeDim::new(&client, working_units); + let cube_count = calculate_cube_count_elemwise(&client, working_units, cube_dim); + + let output = empty_device_dtype( + client.clone(), + input.device.clone(), + input.shape(), + dtype, // We take the same dtype as passed as input (Flex32 not F32) + ); + + cast_element::launch( + &client, + cube_count, + cube_dim, + address_type!(input, output), + vector_size, + input.into_linear_view(), + output.clone().into_linear_view(), + [ + dtype_to_storage_type(dtype_input), + dtype_to_storage_type(dtype_output), + ], + ); + + output +} diff --git a/crates/burn-cubecl/src/kernel/cast/bool_cast.rs b/crates/burn-cubecl/src/kernel/cast/bool_cast.rs new file mode 100644 index 0000000..6b2e000 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/cast/bool_cast.rs @@ -0,0 +1,72 @@ +use crate::{ + CubeRuntime, + kernel::utils::address_type, + ops::{max_vector_size, numeric::empty_device_dtype}, + tensor::CubeTensor, +}; +use burn_backend::TensorMetadata; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_std::DType; +use cubecl::{ + CubeDim, calculate_cube_count_elemwise, + num_traits::One, + prelude::*, + std::tensor::layout::linear::{LinearView, LinearViewMut}, +}; + +#[cube(launch_unchecked, address_type = "dynamic")] +fn bool_cast_kernel( + input: LinearView<'_, Vector>, + mut output: LinearViewMut<'_, Vector>, + #[define(B, T)] _dtypes: [StorageType; 2], +) { + if !output.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + output.write( + ABSOLUTE_POS, + Vector::cast_from(input.read(ABSOLUTE_POS) & Vector::one()), + ); +} + +/// Cast a bool tensor to the given element type. +/// +/// This alternative to cast is necessary because bool are represented as u32 or u8 +/// where any non-zero value means true. Depending how it was created +/// it may hold an uncanny bit combination. Naively casting it would not +/// necessarily yield 0 or 1. +pub fn bool_cast(tensor: CubeTensor, out_dtype: DType) -> CubeTensor { + let output = empty_device_dtype( + tensor.client.clone(), + tensor.device.clone(), + tensor.shape(), + out_dtype, + ); + + let vector_size = max_vector_size(&tensor); + let num_elems = tensor.meta.num_elements(); + let working_units = num_elems / vector_size as usize; + let cube_dim = CubeDim::new(&tensor.client, working_units); + let cube_count = calculate_cube_count_elemwise(&tensor.client, working_units, cube_dim); + + let dtype = tensor.dtype; + + unsafe { + bool_cast_kernel::launch_unchecked( + &output.client, + cube_count, + cube_dim, + address_type!(tensor, output), + vector_size, + tensor.into_linear_view(), + output.clone().into_linear_view(), + [ + dtype_to_storage_type(dtype), + dtype_to_storage_type(out_dtype), + ], + ) + }; + + output +} diff --git a/crates/burn-cubecl/src/kernel/cast/mod.rs b/crates/burn-cubecl/src/kernel/cast/mod.rs new file mode 100644 index 0000000..52d9d7b --- /dev/null +++ b/crates/burn-cubecl/src/kernel/cast/mod.rs @@ -0,0 +1,5 @@ +mod base; +mod bool_cast; + +pub use base::*; +pub use bool_cast::*; diff --git a/crates/burn-cubecl/src/kernel/clamp.rs b/crates/burn-cubecl/src/kernel/clamp.rs new file mode 100644 index 0000000..9a705e1 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/clamp.rs @@ -0,0 +1,41 @@ +use cubecl::prelude::*; + +use crate::{ + CubeRuntime, + kernel::{NumericUnaryOp, NumericUnaryOpFamily, launch_unary_numeric}, + tensor::CubeTensor, +}; + +#[derive(CubeLaunch, CubeType)] +struct Options { + min_value: InputScalar, + max_value: InputScalar, +} + +pub(crate) fn clamp( + input: CubeTensor, + min_value: InputScalar, + max_value: InputScalar, +) -> CubeTensor { + struct ClampOp; + + #[cube] + impl NumericUnaryOp for ClampOp { + type Options = Options; + + fn execute(input: Vector, options: &Self::Options) -> Vector { + cubecl::prelude::clamp( + input, + Vector::new(options.min_value.get::()), + Vector::new(options.max_value.get::()), + ) + } + } + + impl NumericUnaryOpFamily for ClampOp { + type Options = Options; + type Unary = Self; + } + + launch_unary_numeric::(input, |_| OptionsLaunch::new(min_value, max_value)) +} diff --git a/crates/burn-cubecl/src/kernel/comparison.rs b/crates/burn-cubecl/src/kernel/comparison.rs new file mode 100644 index 0000000..b839939 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/comparison.rs @@ -0,0 +1,465 @@ +use crate::{ + CubeRuntime, + kernel::utils::{address_type, broadcast_shape}, + ops::{max_vector_size, numeric::empty_device_dtype}, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{DType, TensorMetadata}; +use cubecl::{ + calculate_cube_count_elemwise, + prelude::*, + std::tensor::layout::linear::{LinearView, LinearViewMut}, +}; + +#[cube] +pub(crate) trait ComparisonOpFamily: 'static + Send + Sync { + type Operation: ComparisonOp; +} + +#[cube] +pub(crate) trait ComparisonOp: 'static + Send + Sync { + /// Execute a comparison operation. + fn execute(lhs: Vector, rhs: Vector) -> bool; +} + +struct EqualOp; +struct GreaterEqualOp; +struct LowerEqualOp; +struct GreaterOp; +struct LowerOp; + +impl ComparisonOpFamily for EqualOp { + type Operation = Self; +} + +#[cube] +impl ComparisonOp for EqualOp { + fn execute(lhs: Vector, rhs: Vector) -> bool { + lhs == rhs + } +} + +impl ComparisonOpFamily for GreaterEqualOp { + type Operation = Self; +} + +#[cube] +impl ComparisonOp for GreaterEqualOp { + fn execute(lhs: Vector, rhs: Vector) -> bool { + lhs >= rhs + } +} + +impl ComparisonOpFamily for LowerEqualOp { + type Operation = Self; +} + +#[cube] +impl ComparisonOp for LowerEqualOp { + fn execute(lhs: Vector, rhs: Vector) -> bool { + lhs <= rhs + } +} + +impl ComparisonOpFamily for GreaterOp { + type Operation = Self; +} + +#[cube] +impl ComparisonOp for GreaterOp { + fn execute(lhs: Vector, rhs: Vector) -> bool { + lhs > rhs + } +} + +impl ComparisonOpFamily for LowerOp { + type Operation = Self; +} + +#[cube] +impl ComparisonOp for LowerOp { + fn execute(lhs: Vector, rhs: Vector) -> bool { + lhs < rhs + } +} + +#[cube(launch_unchecked, address_type = "dynamic")] +pub(crate) fn kernel_scalar_cmp( + input: LinearView<'_, Vector>, + scalar: InputScalar, + mut output: LinearViewMut<'_, Vector>, + #[define(T, Bool)] _dtypes: [StorageType; 2], +) { + if !output.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + output.write( + ABSOLUTE_POS, + Vector::cast_from(O::Operation::::execute( + input.read(ABSOLUTE_POS), + Vector::new(scalar.get::()), + )), + ); +} + +#[cube(launch_unchecked, address_type = "dynamic")] +pub(crate) fn kernel_cmp( + lhs: LinearView<'_, Vector>, + rhs: LinearView<'_, Vector>, + mut out: LinearViewMut<'_, Vector>, + #[define(T, Bool)] _dtype: [StorageType; 2], +) { + if !out.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + out.write( + ABSOLUTE_POS, + Vector::cast_from(O::Operation::::execute( + lhs.read(ABSOLUTE_POS), + rhs.read(ABSOLUTE_POS), + )), + ); +} + +pub(crate) fn launch_cmp( + lhs: CubeTensor, + rhs: CubeTensor, + dtype_bool: DType, +) -> CubeTensor { + let vector_size_lhs = max_vector_size(&lhs); + let vector_size_rhs = max_vector_size(&rhs); + + let vector_size = Ord::min(vector_size_lhs, vector_size_rhs); + + let shape_out = broadcast_shape(&[&lhs, &rhs]); + let client = lhs.client.clone(); + let num_elems = shape_out.num_elements(); + + let working_units = num_elems / vector_size as usize; + let cube_dim = CubeDim::new(&lhs.client, working_units); + let cube_count = calculate_cube_count_elemwise(&lhs.client, working_units, cube_dim); + + let dtypes = [ + dtype_to_storage_type(lhs.dtype), + dtype_to_storage_type(dtype_bool), + ]; + let same_tensor_type = dtypes[0] == dtypes[1]; + if same_tensor_type && lhs.can_mut_broadcast(&rhs) { + unsafe { + kernel_cmp::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(lhs, rhs), + vector_size, + lhs.clone().into_linear_view(), + rhs.into_linear_view_like(&lhs), + lhs.as_linear_view_alias(0), + dtypes, + ); + } + + CubeTensor::new( + lhs.client.clone(), + lhs.handle.clone(), + *lhs.meta.clone(), + lhs.device.clone(), + dtype_bool, + ) + } else if same_tensor_type && rhs.can_mut_broadcast(&lhs) { + unsafe { + kernel_cmp::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(lhs, rhs), + vector_size, + lhs.into_linear_view_like(&rhs), + rhs.clone().into_linear_view(), + rhs.as_linear_view_alias(1), + dtypes, + ); + }; + + CubeTensor::new( + rhs.client.clone(), + rhs.handle.clone(), + *rhs.meta.clone(), + rhs.device.clone(), + dtype_bool, + ) + } else { + let output = empty_device_dtype( + lhs.client.clone(), + lhs.device.clone(), + shape_out, + dtype_bool, + ); + + unsafe { + kernel_cmp::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(lhs, rhs, output), + vector_size, + lhs.into_linear_view_like(&output), + rhs.into_linear_view_like(&output), + output.clone().into_linear_view(), + dtypes, + ); + }; + + output + } +} + +pub(crate) fn launch_scalar_cmp( + tensor: CubeTensor, + scalar: InputScalar, + dtype_bool: DType, +) -> CubeTensor { + let vector_size = max_vector_size(&tensor); + let client = tensor.client.clone(); + let num_elems = tensor.meta.num_elements(); + + let working_units = num_elems / vector_size as usize; + let cube_dim = CubeDim::new(&tensor.client, working_units); + let cube_count = calculate_cube_count_elemwise(&tensor.client, working_units, cube_dim); + + let dtypes = [ + dtype_to_storage_type(tensor.dtype), + dtype_to_storage_type(dtype_bool), + ]; + let same_tensor_type = dtypes[0] == dtypes[1]; + + if same_tensor_type && tensor.can_mut() && tensor.is_nonoverlapping() { + unsafe { + kernel_scalar_cmp::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(tensor), + vector_size, + tensor.clone().into_linear_view(), + scalar, + tensor.as_linear_view_alias(0), + dtypes, + ); + } + + CubeTensor::new( + tensor.client.clone(), + tensor.handle.clone(), + *tensor.meta.clone(), + tensor.device.clone(), + dtype_bool, + ) + } else { + let output = empty_device_dtype( + tensor.client.clone(), + tensor.device.clone(), + tensor.shape(), + dtype_bool, + ); + + unsafe { + kernel_scalar_cmp::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(tensor, output), + vector_size, + tensor.into_linear_view(), + scalar, + output.clone().into_linear_view(), + dtypes, + ); + } + + output + } +} + +pub fn equal( + lhs: CubeTensor, + rhs: CubeTensor, + dtype_bool: DType, +) -> CubeTensor { + launch_cmp::(lhs, rhs, dtype_bool) +} + +pub fn greater( + lhs: CubeTensor, + rhs: CubeTensor, + dtype_bool: DType, +) -> CubeTensor { + launch_cmp::(lhs, rhs, dtype_bool) +} + +pub fn greater_equal( + lhs: CubeTensor, + rhs: CubeTensor, + dtype_bool: DType, +) -> CubeTensor { + launch_cmp::(lhs, rhs, dtype_bool) +} + +pub fn lower( + lhs: CubeTensor, + rhs: CubeTensor, + dtype_bool: DType, +) -> CubeTensor { + launch_cmp::(lhs, rhs, dtype_bool) +} + +pub fn lower_equal( + lhs: CubeTensor, + rhs: CubeTensor, + dtype_bool: DType, +) -> CubeTensor { + launch_cmp::(lhs, rhs, dtype_bool) +} + +pub fn equal_elem( + lhs: CubeTensor, + rhs: InputScalar, + dtype_bool: DType, +) -> CubeTensor { + launch_scalar_cmp::(lhs, rhs, dtype_bool) +} + +pub fn greater_elem( + lhs: CubeTensor, + rhs: InputScalar, + dtype_bool: DType, +) -> CubeTensor { + launch_scalar_cmp::(lhs, rhs, dtype_bool) +} + +pub fn lower_elem( + lhs: CubeTensor, + rhs: InputScalar, + dtype_bool: DType, +) -> CubeTensor { + launch_scalar_cmp::(lhs, rhs, dtype_bool) +} + +pub fn greater_equal_elem( + lhs: CubeTensor, + rhs: InputScalar, + dtype_bool: DType, +) -> CubeTensor { + launch_scalar_cmp::(lhs, rhs, dtype_bool) +} + +pub fn lower_equal_elem( + lhs: CubeTensor, + rhs: InputScalar, + dtype_bool: DType, +) -> CubeTensor { + launch_scalar_cmp::(lhs, rhs, dtype_bool) +} + +// Unary comparison / predicate / relational ops + +#[cube] +pub(crate) trait PredicateOp: 'static + Send + Sync { + /// Execute a predicate operation. + fn execute(input: Vector) -> Vector; +} + +pub(crate) trait PredicateOpFamily: 'static + Send + Sync { + type Operation: PredicateOp; +} + +struct IsNanOp; +struct IsInfOp; + +impl PredicateOpFamily for IsNanOp { + type Operation = Self; +} + +#[cube] +impl PredicateOp for IsNanOp { + fn execute(input: Vector) -> Vector { + Vector::is_nan(input) + } +} + +impl PredicateOpFamily for IsInfOp { + type Operation = Self; +} +#[cube] +impl PredicateOp for IsInfOp { + fn execute(input: Vector) -> Vector { + Vector::is_inf(input) + } +} + +#[cube(launch_unchecked, address_type = "dynamic")] +pub(crate) fn kernel_predicate( + input: LinearView<'_, Vector>, + mut output: LinearViewMut<'_, Vector>, + #[define(F, Bool)] _dtypes: [StorageType; 2], +) { + if !output.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + output.write( + ABSOLUTE_POS, + Vector::cast_from(O::Operation::::execute(input.read(ABSOLUTE_POS))), + ); +} + +pub(crate) fn launch_predicate( + tensor: CubeTensor, + dtype_bool: DType, +) -> CubeTensor { + let vector_size = max_vector_size(&tensor); + + let client = tensor.client.clone(); + let num_elems = tensor.meta.num_elements(); + + let dtypes = [ + dtype_to_storage_type(tensor.dtype), + dtype_to_storage_type(dtype_bool), + ]; + let working_units = num_elems / vector_size as usize; + let cube_dim = CubeDim::new(&tensor.client, working_units); + let cube_count = calculate_cube_count_elemwise(&tensor.client, working_units, cube_dim); + + let output = empty_device_dtype( + tensor.client.clone(), + tensor.device.clone(), + tensor.shape(), + dtype_bool, + ); + + unsafe { + kernel_predicate::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(tensor, output), + vector_size, + tensor.into_linear_view_like(&output), + output.clone().into_linear_view(), + dtypes, + ); + } + + output +} + +pub fn is_nan(tensor: CubeTensor, dtype_bool: DType) -> CubeTensor { + launch_predicate::(tensor, dtype_bool) +} + +pub fn is_inf(tensor: CubeTensor, dtype_bool: DType) -> CubeTensor { + launch_predicate::(tensor, dtype_bool) +} diff --git a/crates/burn-cubecl/src/kernel/contiguous.rs b/crates/burn-cubecl/src/kernel/contiguous.rs new file mode 100644 index 0000000..4b90738 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/contiguous.rs @@ -0,0 +1,125 @@ +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{DType, TensorMetadata}; +use cubecl::quant::scheme::{QuantStore, QuantValue}; +use cubecl::server::MemoryLayoutStrategy; + +use crate::{CubeRuntime, ops::empty_qtensor, tensor::CubeTensor}; + +/// Make a jit tensor contiguous. +pub fn into_contiguous(tensor: CubeTensor) -> CubeTensor { + if tensor.is_contiguous() { + return tensor; + } + + if tensor.qparams.is_some() { + return into_contiguous_quantized(tensor, MemoryLayoutStrategy::Contiguous); + } + + let (client, device, dtype) = (tensor.client.clone(), tensor.device.clone(), tensor.dtype); + + let output = cubecl::std::tensor::into_contiguous( + &client, + tensor.binding(), + dtype_to_storage_type(dtype), + ); + + CubeTensor::new( + client.clone(), + output.handle, + *output.metadata, + device, + dtype, + ) +} + +/// Make a jit tensor contiguous with an aligned last stride. Tensor is considered already contiguous +/// if runtime can read it as is. This is equivalent in practice. +#[cfg_attr( + feature = "tracing", + tracing::instrument(level = "trace", skip(tensor)) +)] +pub fn into_contiguous_aligned(tensor: CubeTensor) -> CubeTensor { + if R::can_read_tensor(tensor.meta.shape(), tensor.meta.strides()) { + return tensor; + } + + if tensor.qparams.is_some() { + return into_contiguous_quantized(tensor, MemoryLayoutStrategy::Optimized); + } + + let (client, device, dtype) = (tensor.client.clone(), tensor.device.clone(), tensor.dtype); + + let output = cubecl::std::tensor::into_contiguous_pitched( + &client, + tensor.binding(), + dtype_to_storage_type(dtype), + ); + + CubeTensor::new( + client.clone(), + output.handle, + *output.metadata, + device, + dtype, + ) +} + +#[cfg_attr( + feature = "tracing", + tracing::instrument(level = "trace", skip(tensor)) +)] +fn into_contiguous_quantized( + tensor: CubeTensor, + strategy: MemoryLayoutStrategy, +) -> CubeTensor { + let scheme = tensor.scheme(); + let output = empty_qtensor(tensor.shape(), tensor.scheme(), &tensor.device, strategy); + let (values, scales) = tensor.quantized_handles().unwrap(); + let (out_values, out_scales) = output.quantized_handles().unwrap(); + + let (client, dtype_scales, dtype_value) = (scales.client.clone(), scales.dtype, values.dtype); + + match scheme.store { + QuantStore::PackedU32(packed_dim) => { + cubecl::std::tensor::into_contiguous_packed_ref( + &client, + values.binding(), + out_values.binding(), + packed_dim, + tensor.meta.shape(), + scheme.num_quants(), + dtype_to_storage_type(DType::U32), + ); + } + // e2m1 is special because it has a native packed representation, `e2m1x2`. + // It's internally stored as `u8` with a packing factor of 2. + QuantStore::PackedNative(packed_dim) if scheme.value == QuantValue::E2M1 => { + cubecl::std::tensor::into_contiguous_packed_ref( + &client, + values.binding(), + out_values.binding(), + packed_dim, + tensor.meta.shape(), + scheme.num_quants(), + dtype_to_storage_type(DType::U8), + ); + } + _ => { + cubecl::std::tensor::copy_into( + &client, + values.binding(), + out_values.binding(), + dtype_to_storage_type(dtype_value), + ); + } + } + + cubecl::std::tensor::copy_into( + &client, + scales.binding(), + out_scales.binding(), + dtype_to_storage_type(dtype_scales), + ); + + output +} diff --git a/crates/burn-cubecl/src/kernel/conv/backward_data/fallback.rs b/crates/burn-cubecl/src/kernel/conv/backward_data/fallback.rs new file mode 100644 index 0000000..698bd5f --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/backward_data/fallback.rs @@ -0,0 +1,124 @@ +use burn_backend::{ + TensorMetadata, + ops::{ConvOptions, ConvTransposeOptions, conv::calculate_padding_out}, +}; +use burn_std::Shape; +use cubek::convolution::components::ConvSetupError; + +use crate::{ + CubeRuntime, + kernel::conv::{conv_transpose2d, conv_transpose3d}, + ops::{permute_nchw_to_nhwc, permute_nhwc_to_nchw, reshape}, + tensor::CubeTensor, +}; + +pub(crate) fn conv_data_backward_fallback( + out_grad: CubeTensor, + weights: CubeTensor, + in_shape: Shape, + options: ConvOptions, +) -> Result, ConvSetupError> { + let dim_c = out_grad.rank(); + + let kernel_size = &weights.meta.shape()[1..dim_c]; + let in_shape = &in_shape[1..dim_c]; + let out_shape = &out_grad.meta.shape()[1..dim_c]; + + let mut padding_out = [0; N_DIM]; + + for i in 0..N_DIM { + padding_out[i] = calculate_padding_out( + kernel_size[i], + options.stride[i], + options.padding[i], + options.dilation[i], + in_shape[i], + out_shape[i], + ); + } + + // We don't yet have NHWC kernels for conv_transpose so need to do this. + // Should eventually use NHWC kernels instead + let out_grad = permute_nhwc_to_nchw(out_grad); + let weights = permute_nhwc_to_nchw(weights); + + let in_grad = match N_DIM { + 1 => conv_transpose1d_from_conv_transpose2d( + out_grad, + weights, + ConvTransposeOptions::new( + [options.stride[0]], + [options.padding[0]], + [padding_out[0]], + [options.dilation[0]], + options.groups, + ), + ), + 2 => conv_transpose2d( + out_grad, + weights, + None, + ConvTransposeOptions::new( + [options.stride[0], options.stride[1]], + [options.padding[0], options.padding[1]], + [padding_out[0], padding_out[1]], + [options.dilation[0], options.dilation[1]], + options.groups, + ), + Default::default(), + ), + 3 => Ok(conv_transpose3d( + out_grad, + weights, + None, + ConvTransposeOptions::new( + [options.stride[0], options.stride[1], options.stride[2]], + [options.padding[0], options.padding[1], options.padding[2]], + [padding_out[0], padding_out[1], padding_out[2]], + [ + options.dilation[0], + options.dilation[1], + options.dilation[2], + ], + options.groups, + ), + ) + .unwrap()), + _ => unimplemented!("Invalid dimensionality"), + }?; + Ok(permute_nchw_to_nhwc(in_grad)) +} + +fn conv_transpose1d_from_conv_transpose2d( + x: CubeTensor, + weight: CubeTensor, + options: ConvTransposeOptions<1>, +) -> Result, ConvSetupError> { + let [channels_in, channels_out, kernel_size] = weight.shape().dims(); + let [batch_size, _channels_in, length_in] = x.shape().dims(); + + let weight = reshape( + weight, + Shape::new([channels_in, channels_out, kernel_size, 1]), + ); + let x = reshape(x, Shape::new([batch_size, channels_in, length_in, 1])); + + let tensor = conv_transpose2d( + x, + weight, + None, + ConvTransposeOptions::new( + [options.stride[0], 1], + [options.padding[0], 0], + [options.padding_out[0], 0], + [options.dilation[0], 1], + options.groups, + ), + Default::default(), + )?; + let [batch_size, channels_out, height_out, _weight_out] = tensor.shape().dims(); + Ok(reshape( + tensor, + Shape::from([batch_size, channels_out, height_out]), + )) +} diff --git a/crates/burn-cubecl/src/kernel/conv/backward_data/implicit_gemm/launch.rs b/crates/burn-cubecl/src/kernel/conv/backward_data/implicit_gemm/launch.rs new file mode 100644 index 0000000..a11c5d8 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/backward_data/implicit_gemm/launch.rs @@ -0,0 +1,135 @@ +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::ops::ConvOptions; +use burn_std::Shape; +use cubek::{ + convolution::{ + AcceleratedTileKind, ConvAlgorithm, ConvolutionArgs, ConvolutionInputs, Strategy, + components::ConvSetupError, launch_ref, + }, + matmul::definition::{MatmulElems, MatmulGlobalElems}, + std::InputBinding, +}; + +use crate::{CubeRuntime, ops::numeric::empty_device_dtype, tensor::CubeTensor}; + +pub fn dgrad_gemm_simple_sync( + out_grad: CubeTensor, + weights: CubeTensor, + input_shape: Shape, + options: ConvOptions, + tile_kind: AcceleratedTileKind, +) -> Result, ConvSetupError> { + let algorithm = match tile_kind { + AcceleratedTileKind::Cmma => ConvAlgorithm::SimpleSyncCyclic, + AcceleratedTileKind::Mma => ConvAlgorithm::SimpleSyncStrided, + }; + launch_backwards_data::( + &Strategy::Inferred { + algorithm, + tile_kind, + }, + out_grad, + weights, + input_shape, + options, + ) +} + +pub fn dgrad_gemm_simple_async( + out_grad: CubeTensor, + weights: CubeTensor, + input_shape: Shape, + options: ConvOptions, + tile_kind: AcceleratedTileKind, +) -> Result, ConvSetupError> { + let algorithm = match tile_kind { + AcceleratedTileKind::Cmma => ConvAlgorithm::SimpleAsyncCyclic, + AcceleratedTileKind::Mma => ConvAlgorithm::SimpleAsyncStrided, + }; + launch_backwards_data::( + &Strategy::Inferred { + algorithm, + tile_kind, + }, + out_grad, + weights, + input_shape, + options, + ) +} + +pub fn dgrad_gemm_simple_tma( + out_grad: CubeTensor, + weights: CubeTensor, + input_shape: Shape, + options: ConvOptions, + tile_kind: AcceleratedTileKind, +) -> Result, ConvSetupError> { + launch_backwards_data::( + &Strategy::Inferred { + algorithm: ConvAlgorithm::SimpleAsyncTma, + tile_kind, + }, + out_grad, + weights, + input_shape, + options, + ) +} + +/// Perform a convolution backwards data pass using the implicit GEMM (im2col) algorithm, using +/// cubecl tiling matmul components. +/// +/// * `input` - The input feature map +/// * `out_grad` - The output gradients +/// * `weight_shape` - The shape of the weights/weight gradients +/// * `options` - The options to use for the convolution +pub fn launch_backwards_data( + strategy: &Strategy, + out_grad: CubeTensor, + weights: CubeTensor, + input_shape: Shape, + options: ConvOptions, +) -> Result, ConvSetupError> { + if options.groups != 1 || options.stride.iter().any(|&s| s != 1) { + return Err(ConvSetupError::Groups(options.groups)); + } + + let out_dtype = out_grad.dtype; + + let in_grad = empty_device_dtype( + out_grad.client.clone(), + out_grad.device.clone(), + input_shape, + out_dtype, + ); + + let client = out_grad.client.clone(); + let dtypes = MatmulElems::from_globals(&MatmulGlobalElems { + lhs: dtype_to_storage_type(out_grad.dtype), + rhs: dtype_to_storage_type(weights.dtype), + out: dtype_to_storage_type(out_dtype), + }); + let out_grad_dtype = out_grad.dtype; + let weights_dtype = weights.dtype; + let out_grad = InputBinding::new(out_grad.binding(), dtype_to_storage_type(out_grad_dtype)); + let weights = InputBinding::new(weights.binding(), dtype_to_storage_type(weights_dtype)); + + launch_ref::( + strategy, + &client, + ConvolutionInputs::BackwardData { + out_grad, + weights, + in_grad: in_grad.clone().binding(), + }, + ConvolutionArgs { + stride: options.stride, + padding: options.padding, + dilation: options.dilation, + }, + dtypes, + )?; + + Ok(in_grad) +} diff --git a/crates/burn-cubecl/src/kernel/conv/backward_data/implicit_gemm/mod.rs b/crates/burn-cubecl/src/kernel/conv/backward_data/implicit_gemm/mod.rs new file mode 100644 index 0000000..0df8c51 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/backward_data/implicit_gemm/mod.rs @@ -0,0 +1,2 @@ +pub mod launch; +pub use launch::*; diff --git a/crates/burn-cubecl/src/kernel/conv/backward_data/mod.rs b/crates/burn-cubecl/src/kernel/conv/backward_data/mod.rs new file mode 100644 index 0000000..80e94c4 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/backward_data/mod.rs @@ -0,0 +1,8 @@ +pub mod fallback; +pub mod implicit_gemm; + +#[cfg(feature = "autotune")] +pub mod tune; + +#[cfg(feature = "autotune")] +pub(crate) use tune::*; diff --git a/crates/burn-cubecl/src/kernel/conv/backward_data/tune.rs b/crates/burn-cubecl/src/kernel/conv/backward_data/tune.rs new file mode 100644 index 0000000..d626d8f --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/backward_data/tune.rs @@ -0,0 +1,182 @@ +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::ops::ConvOptions; +use burn_std::Shape; +use cubecl::{ + ir::StorageType, + tune::{LocalTuner, Tunable, TunableSet, anchor, local_tuner}, +}; +use cubek::convolution::AcceleratedTileKind; + +use crate::{ + CubeAutotuneKey, CubeRuntime, CubeTuneId, + kernel::conv::{ + ConvAutotuneKey, + backward_data::{fallback::conv_data_backward_fallback, implicit_gemm::*}, + }, + tensor::CubeTensor, +}; + +/// Executes autotune on conv2d operations +pub fn dgrad_autotune( + out_grad: CubeTensor, + weights: CubeTensor, + input_shape: Shape, + options: ConvOptions, +) -> CubeTensor { + let client = out_grad.client.clone(); + + static TUNER: LocalTuner = local_tuner!(); + + // Note: TMA isn't currently implemented properly, and will always error. + // It's kept here so it gets automatically enabled as soon as cubek updates. + // No CMMA for TMA because swizzling will be mandatory for good performance on dgrad. + let tunables = TUNER.init(|| { + TunableSet::new(create_key::, create_wgrad_input::) + .with(Tunable::new( + "wgrad_fallback", + |(out_grad, weights, input_shape, options)| { + conv_data_backward_fallback::(out_grad, weights, input_shape, options) + }, + )) + .with(Tunable::new( + "simple_sync_cmma", + |(input, grad, shape, options)| { + dgrad_gemm_simple_sync(input, grad, shape, options, AcceleratedTileKind::Cmma) + }, + )) + .with(Tunable::new( + "simple_sync_mma", + |(input, grad, shape, options)| { + dgrad_gemm_simple_sync(input, grad, shape, options, AcceleratedTileKind::Mma) + }, + )) + .with(Tunable::new( + "simple_async_cmma", + |(input, grad, shape, options)| { + dgrad_gemm_simple_async(input, grad, shape, options, AcceleratedTileKind::Cmma) + }, + )) + .with(Tunable::new( + "simple_async_mma", + |(input, grad, shape, options)| { + dgrad_gemm_simple_async(input, grad, shape, options, AcceleratedTileKind::Mma) + }, + )) + .with(Tunable::new( + "simple_tma_mma", + |(input, grad, shape, options)| { + dgrad_gemm_simple_tma(input, grad, shape, options, AcceleratedTileKind::Mma) + }, + )) + }); + + TUNER.execute( + &CubeTuneId::new(&out_grad.client, &out_grad.device), + &client, + tunables, + (out_grad, weights, input_shape, options), + ) +} + +pub fn create_wgrad_input( + _key: &CubeAutotuneKey, + (out_grad, weights, input_shape, options): &( + CubeTensor, + CubeTensor, + Shape, + ConvOptions, + ), +) -> (CubeTensor, CubeTensor, Shape, ConvOptions) { + ( + out_grad.clone(), + weights.clone(), + input_shape.clone(), + options.clone(), + ) +} + +fn create_key( + (out_grad, weights, input_shape, options): &( + CubeTensor, + CubeTensor, + Shape, + ConvOptions, + ), +) -> CubeAutotuneKey { + let dtype = out_grad.dtype; + let rank = out_grad.meta.num_dims(); + let dim_c = rank - 1; + + let batch_size = out_grad.meta.shape()[0]; + let in_channels = input_shape[dim_c]; + let out_channels = out_grad.meta.shape()[dim_c]; + + let kernel_size = weights.meta.shape()[1..dim_c].to_vec(); + let in_shape = input_shape[1..dim_c] + .iter() + .map(|shape| anchor(*shape, None, None, None)) + .collect(); + + let ConvOptions { + stride, + padding, + dilation, + groups, + } = options.clone(); + + let lhs_stride_align = if out_grad.meta.strides()[dim_c] == 1 { + stride_align( + out_grad.meta.strides(), + dtype_to_storage_type(out_grad.dtype), + ) + } else { + 0 + }; + let lhs_shape_align = pow2_factor(out_channels).min(lhs_stride_align); + let rhs_stride_align = if weights.meta.strides()[dim_c] == 1 { + stride_align(weights.meta.strides(), dtype_to_storage_type(weights.dtype)) + } else { + 0 + }; + let rhs_shape_align = pow2_factor(in_channels).min(rhs_stride_align); + + CubeAutotuneKey::Conv(ConvAutotuneKey::new( + kernel_size, + stride.to_vec(), + padding.to_vec(), + dilation.to_vec(), + groups, + in_channels, + out_channels, + in_shape, + batch_size, + false, + dtype, + lhs_shape_align, + lhs_stride_align, + rhs_shape_align, + rhs_stride_align, + )) +} + +/// Maximum factor relevant for strides. Currently set to 2^10 because that's 128-byte swizzle's +/// repeat number, so it's the largest align that can have performance impacts. +const MAX_STRIDE_FACTOR: u32 = 10; + +/// Defines the non-contiguous stride alignment in terms of powers of two +fn stride_align(strides: &[usize], elem: StorageType) -> u8 { + let max = MAX_STRIDE_FACTOR; + let dim_c = strides.len() - 1; + let factor = strides[..dim_c] + .iter() + .map(|it| (*it * elem.size_bits()) / 8) + .map(|it| it.trailing_zeros()) + .min() + .unwrap_or(max); + factor.min(max) as u8 +} + +/// Defines the potential vectorization. +fn pow2_factor(axis: usize) -> u8 { + axis.trailing_zeros().min(4) as u8 +} diff --git a/crates/burn-cubecl/src/kernel/conv/backward_weight/fallback.rs b/crates/burn-cubecl/src/kernel/conv/backward_weight/fallback.rs new file mode 100644 index 0000000..6ddd9f2 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/backward_weight/fallback.rs @@ -0,0 +1,112 @@ +use burn_backend::{TensorMetadata, ops::ConvOptions}; +use burn_std::{Shape, Slice}; +use cubek::convolution::components::ConvSetupError; + +use crate::{ + CubeRuntime, + kernel::{conv::base::conv_forward_nhwc, slice, slice_assign}, + ops::{numeric::empty_device_dtype, swap_dims}, + tensor::CubeTensor, +}; + +/// Calculate the convolution backward pass with regard to the weight gradients. +pub fn conv_weight_backward_fallback( + input: CubeTensor, + output_grad: CubeTensor, + weight_shape: Shape, + options: ConvOptions, +) -> Result, ConvSetupError> { + match options.groups == 1 { + true => conv_weight_grad_no_groups::(input, output_grad, weight_shape, options), + false => conv_weight_grad_groups::(input, output_grad, weight_shape, options), + } +} + +fn conv_weight_grad_no_groups( + input: CubeTensor, + output_grad: CubeTensor, + weight_shape: Shape, + options: ConvOptions, +) -> Result, ConvSetupError> { + let dim_c = input.rank() - 1; + + let input_swapped = swap_dims(input, 0, dim_c); + let out_grad_swapped = swap_dims(output_grad, 0, dim_c); + let weight_grad_swapped = conv_forward_nhwc( + input_swapped, + out_grad_swapped, + None, + ConvOptions::new(options.dilation, options.padding, options.stride, 1), + Default::default(), + )?; + let mut weight_grad = swap_dims(weight_grad_swapped, 0, dim_c); + if weight_grad.shape() != weight_shape { + let ranges = weight_shape.iter().map(|&s| 0..s).collect::>(); + weight_grad = slice(weight_grad, &ranges); + } + + Ok(weight_grad) +} + +#[allow(clippy::single_range_in_vec_init, reason = "False positive")] +fn conv_weight_grad_groups( + input: CubeTensor, + output_grad: CubeTensor, + weight_shape: Shape, + options: ConvOptions, +) -> Result, ConvSetupError> { + let mut weight_grad = empty_device_dtype( + input.client.clone(), + input.device.clone(), + weight_shape.clone(), + input.dtype, + ); + + let dim_c = input.rank() - 1; + + let channels_out = weight_shape[0]; + let increment_co = channels_out / options.groups; + + let input_swapped = swap_dims(input, 0, dim_c); + let output_grad_swapped = swap_dims(output_grad, 0, dim_c); + + let kernel_size = &weight_shape[1..dim_c]; + let kernel_size_slice = kernel_size.iter().map(|&s| 0..s).collect::>(); + let increment_ci = weight_grad.meta.shape()[dim_c]; + + for g in 0..options.groups { + let start_idx_ci = g * increment_ci; + let end_idx_ci = (g + 1) * increment_ci; + let start_idx_co = g * increment_co; + let end_idx_co = (g + 1) * increment_co; + + let input = slice(input_swapped.clone(), &[start_idx_ci..end_idx_ci]); + let grad = slice(output_grad_swapped.clone(), &[start_idx_co..end_idx_co]); + + let weight_grad_tmp = conv_forward_nhwc( + input, + grad, + None, + ConvOptions::new(options.dilation, options.padding, options.stride, 1), + Default::default(), + )?; + let mut weight_grad_tmp = swap_dims(weight_grad_tmp, 0, dim_c); + let kernel_size_tmp = &weight_grad_tmp.meta.shape()[1..dim_c]; + + if kernel_size != kernel_size_tmp { + let mut slices = vec![0..increment_co]; + slices.extend(kernel_size_slice.clone()); + slices.push(0..increment_ci); + weight_grad_tmp = slice(weight_grad_tmp, &slices); + } + + let mut slices = vec![start_idx_co..end_idx_co]; + slices.extend(kernel_size_slice.clone()); + slices.push(0..increment_ci); + let slices = slices.into_iter().map(Slice::from).collect::>(); + + weight_grad = slice_assign(weight_grad, &slices, weight_grad_tmp); + } + + Ok(weight_grad) +} diff --git a/crates/burn-cubecl/src/kernel/conv/backward_weight/implicit_gemm/launch.rs b/crates/burn-cubecl/src/kernel/conv/backward_weight/implicit_gemm/launch.rs new file mode 100644 index 0000000..af9215e --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/backward_weight/implicit_gemm/launch.rs @@ -0,0 +1,135 @@ +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::ops::ConvOptions; +use burn_std::Shape; +use cubek::{ + convolution::{ + AcceleratedTileKind, ConvAlgorithm, ConvolutionArgs, ConvolutionInputs, Strategy, + components::ConvSetupError, launch_ref, + }, + matmul::definition::{MatmulElems, MatmulGlobalElems}, + std::InputBinding, +}; + +use crate::{CubeRuntime, ops::numeric::empty_device_dtype, tensor::CubeTensor}; + +pub(crate) fn wgrad_gemm_simple_sync( + input: CubeTensor, + out_grad: CubeTensor, + weight_shape: Shape, + options: ConvOptions, + tile_kind: AcceleratedTileKind, +) -> Result, ConvSetupError> { + let algorithm = match tile_kind { + AcceleratedTileKind::Cmma => ConvAlgorithm::SimpleSyncCyclic, + AcceleratedTileKind::Mma => ConvAlgorithm::SimpleSyncStrided, + }; + launch_backwards_weight::( + &Strategy::Inferred { + algorithm, + tile_kind, + }, + input, + out_grad, + weight_shape, + options, + ) +} + +pub(crate) fn wgrad_gemm_simple_async( + input: CubeTensor, + out_grad: CubeTensor, + weight_shape: Shape, + options: ConvOptions, + tile_kind: AcceleratedTileKind, +) -> Result, ConvSetupError> { + let algorithm = match tile_kind { + AcceleratedTileKind::Cmma => ConvAlgorithm::SimpleAsyncCyclic, + AcceleratedTileKind::Mma => ConvAlgorithm::SimpleAsyncStrided, + }; + launch_backwards_weight::( + &Strategy::Inferred { + algorithm, + tile_kind, + }, + input, + out_grad, + weight_shape, + options, + ) +} + +pub(crate) fn wgrad_gemm_simple_tma( + input: CubeTensor, + out_grad: CubeTensor, + weight_shape: Shape, + options: ConvOptions, + tile_kind: AcceleratedTileKind, +) -> Result, ConvSetupError> { + launch_backwards_weight::( + &Strategy::Inferred { + algorithm: ConvAlgorithm::SimpleAsyncTma, + tile_kind, + }, + input, + out_grad, + weight_shape, + options, + ) +} + +/// Perform a convolution backwards weight pass using the implicit GEMM (im2col) algorithm, using +/// cubecl tiling matmul components. +/// +/// * `input` - The input feature map +/// * `out_grad` - The output gradients +/// * `weight_shape` - The shape of the weights/weight gradients +/// * `options` - The options to use for the convolution +pub fn launch_backwards_weight( + strategy: &Strategy, + input: CubeTensor, + out_grad: CubeTensor, + weight_shape: Shape, + options: ConvOptions, +) -> Result, ConvSetupError> { + if options.groups != 1 { + return Err(ConvSetupError::Groups(options.groups)); + } + + let out_dtype = out_grad.dtype; + + let weight_grad = empty_device_dtype( + input.client.clone(), + input.device.clone(), + weight_shape, + out_dtype, + ); + + let client = input.client.clone(); + let dtypes = MatmulElems::from_globals(&MatmulGlobalElems { + lhs: dtype_to_storage_type(input.dtype), + rhs: dtype_to_storage_type(out_grad.dtype), + out: dtype_to_storage_type(out_dtype), + }); + let input_dtype = input.dtype; + let out_grad_dtype = out_grad.dtype; + let input = InputBinding::new(input.binding(), dtype_to_storage_type(input_dtype)); + let out_grad = InputBinding::new(out_grad.binding(), dtype_to_storage_type(out_grad_dtype)); + + launch_ref::( + strategy, + &client, + ConvolutionInputs::BackwardWeight { + input, + out_grad, + weight_grad: weight_grad.clone().binding(), + }, + ConvolutionArgs { + stride: options.stride, + padding: options.padding, + dilation: options.dilation, + }, + dtypes, + )?; + + Ok(weight_grad) +} diff --git a/crates/burn-cubecl/src/kernel/conv/backward_weight/implicit_gemm/mod.rs b/crates/burn-cubecl/src/kernel/conv/backward_weight/implicit_gemm/mod.rs new file mode 100644 index 0000000..0df8c51 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/backward_weight/implicit_gemm/mod.rs @@ -0,0 +1,2 @@ +pub mod launch; +pub use launch::*; diff --git a/crates/burn-cubecl/src/kernel/conv/backward_weight/mod.rs b/crates/burn-cubecl/src/kernel/conv/backward_weight/mod.rs new file mode 100644 index 0000000..80e94c4 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/backward_weight/mod.rs @@ -0,0 +1,8 @@ +pub mod fallback; +pub mod implicit_gemm; + +#[cfg(feature = "autotune")] +pub mod tune; + +#[cfg(feature = "autotune")] +pub(crate) use tune::*; diff --git a/crates/burn-cubecl/src/kernel/conv/backward_weight/tune.rs b/crates/burn-cubecl/src/kernel/conv/backward_weight/tune.rs new file mode 100644 index 0000000..f957e1d --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/backward_weight/tune.rs @@ -0,0 +1,185 @@ +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::ops::ConvOptions; +use burn_std::Shape; +use cubecl::{ + ir::StorageType, + tune::{LocalTuner, Tunable, TunableSet, anchor, local_tuner}, +}; +use cubek::convolution::AcceleratedTileKind; + +use crate::{ + CubeAutotuneKey, CubeRuntime, CubeTuneId, + kernel::conv::{ + ConvAutotuneKey, + backward_weight::{fallback::conv_weight_backward_fallback, implicit_gemm::*}, + }, + tensor::CubeTensor, +}; + +/// Executes autotune on the weight gradients pass for convolution +pub fn wgrad_autotune( + input: CubeTensor, + out_grad: CubeTensor, + weight_shape: Shape, + options: ConvOptions, +) -> CubeTensor { + let client = input.client.clone(); + + static TUNER: LocalTuner = local_tuner!(); + + let tunables = TUNER.init(|| { + TunableSet::new(create_key::, create_wgrad_input::) + .with(Tunable::new( + "wgrad_fallback", + |(input, grad, shape, options)| { + conv_weight_backward_fallback::(input, grad, shape, options) + }, + )) + .with(Tunable::new( + "simple_sync_cmma", + |(input, grad, shape, options)| { + wgrad_gemm_simple_sync(input, grad, shape, options, AcceleratedTileKind::Cmma) + }, + )) + .with(Tunable::new( + "simple_sync_mma", + |(input, grad, shape, options)| { + wgrad_gemm_simple_sync(input, grad, shape, options, AcceleratedTileKind::Mma) + }, + )) + .with(Tunable::new( + "simple_async_cmma", + |(input, grad, shape, options)| { + wgrad_gemm_simple_async(input, grad, shape, options, AcceleratedTileKind::Cmma) + }, + )) + .with(Tunable::new( + "simple_async_mma", + |(input, grad, shape, options)| { + wgrad_gemm_simple_async(input, grad, shape, options, AcceleratedTileKind::Mma) + }, + )) + .with(Tunable::new( + "simple_tma_cmma", + |(input, grad, shape, options)| { + wgrad_gemm_simple_tma(input, grad, shape, options, AcceleratedTileKind::Cmma) + }, + )) + .with(Tunable::new( + "simple_tma_mma", + |(input, grad, shape, options)| { + wgrad_gemm_simple_tma(input, grad, shape, options, AcceleratedTileKind::Mma) + }, + )) + }); + + TUNER.execute( + &CubeTuneId::new(&input.client, &input.device), + &client, + tunables, + (input, out_grad, weight_shape, options), + ) +} + +pub fn create_wgrad_input( + _key: &CubeAutotuneKey, + (input, out_grad, weight_shape, options): &( + CubeTensor, + CubeTensor, + Shape, + ConvOptions, + ), +) -> (CubeTensor, CubeTensor, Shape, ConvOptions) { + ( + input.clone(), + out_grad.clone(), + weight_shape.clone(), + options.clone(), + ) +} + +fn create_key( + (input, out_grad, weight_shape, options): &( + CubeTensor, + CubeTensor, + Shape, + ConvOptions, + ), +) -> CubeAutotuneKey { + let dtype = input.dtype; + let rank = input.meta.num_dims(); + let dim_c = rank - 1; + + let batch_size = input.meta.shape()[0]; + let in_channels = input.meta.shape()[dim_c]; + let out_channels = weight_shape[0]; + + let kernel_size = weight_shape[1..dim_c].to_vec(); + let in_shape = input.meta.shape()[1..dim_c] + .iter() + .map(|shape| anchor(*shape, None, None, None)) + .collect(); + + let ConvOptions { + stride, + padding, + dilation, + groups, + } = options.clone(); + + let lhs_stride_align = if out_grad.meta.strides()[dim_c] == 1 { + stride_align( + out_grad.meta.strides(), + dtype_to_storage_type(out_grad.dtype), + ) + } else { + 0 + }; + let lhs_shape_align = pow2_factor(out_channels).min(lhs_stride_align); + let rhs_stride_align = if input.meta.strides()[dim_c] == 1 { + stride_align(input.meta.strides(), dtype_to_storage_type(input.dtype)) + } else { + 0 + }; + let rhs_shape_align = pow2_factor(in_channels).min(rhs_stride_align); + + CubeAutotuneKey::Conv(ConvAutotuneKey::new( + kernel_size, + stride.to_vec(), + padding.to_vec(), + dilation.to_vec(), + groups, + in_channels, + out_channels, + in_shape, + batch_size, + false, + dtype, + lhs_shape_align, + lhs_stride_align, + rhs_shape_align, + rhs_stride_align, + )) +} + +/// Maximum factor relevant for strides. Currently set to 2^10 because that's 128-byte swizzle's +/// repeat number, so it's the largest align that can have performance impacts. +const MAX_STRIDE_FACTOR: u32 = 10; + +/// Defines the non-contiguous stride alignment in terms of powers of two +fn stride_align(strides: &[usize], elem: StorageType) -> u8 { + let max = MAX_STRIDE_FACTOR; + let dim_c = strides.len() - 1; + let factor = strides[..dim_c] + .iter() + .map(|it| (*it * elem.size_bits()) / 8) + .map(|it| it.trailing_zeros()) + .min() + .unwrap_or(max); + factor.min(max) as u8 +} + +/// Defines the potential vectorization. +fn pow2_factor(axis: usize) -> u8 { + axis.trailing_zeros().min(4) as u8 +} diff --git a/crates/burn-cubecl/src/kernel/conv/base.rs b/crates/burn-cubecl/src/kernel/conv/base.rs new file mode 100644 index 0000000..7d81082 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/base.rs @@ -0,0 +1,189 @@ +use burn_backend::ops::ConvOptions; +use burn_std::Shape; +use cubek::convolution::{AcceleratedTileKind, components::ConvSetupError}; + +#[cfg(feature = "autotune")] +use crate::kernel::conv::{backward_weight::wgrad_autotune, dgrad_autotune}; +use crate::{ + CubeRuntime, + kernel::conv::{ + backward_data::{fallback::conv_data_backward_fallback, implicit_gemm::*}, + backward_weight::{fallback::conv_weight_backward_fallback, implicit_gemm::*}, + forward::implicit_gemm::conv_gemm_simple_sync, + }, + ops::{permute_nchw_to_nhwc, permute_nchw_to_nhwc_shape, permute_nhwc_to_nchw}, + tensor::CubeTensor, +}; + +use super::conv_direct; +#[cfg(feature = "autotune")] +use super::forward::conv_autotune; + +/// The strategy to be used when launching a convolution kernel. +pub enum ConvStrategy { + /// A simple direct convolution. + Direct, + #[cfg(feature = "autotune")] + /// Using autotune to choose the best kernel based on runtime information. + Autotune, + /// Implicit GEMM implementation of convolution. Lower memory usage but requires CMMA and + /// has constraints on tensor shape. + ImplicitGemm, +} + +impl Default for ConvStrategy { + fn default() -> Self { + // if autotune is enabled, default to autotune + #[cfg(feature = "autotune")] + return ConvStrategy::Autotune; + + // if autotune is disabled, default to the more memory-conservative algorithm + #[cfg(not(feature = "autotune"))] + ConvStrategy::Direct + } +} + +/// Performs an N-dimensional convolution with the given strategy +/// +/// * `input` - The input feature map +/// * `weight` - The weights (filter) applied to each kernel +/// * `bias` - The bias added to each channel +/// * `options` - The options to use for the convolution +/// * `strategy` - The convolution algorithm to use. Autotune will pick the fastest available option. +pub fn conv_forward( + input: CubeTensor, + weight: CubeTensor, + bias: Option>, + options: ConvOptions, + strategy: ConvStrategy, +) -> Result, ConvSetupError> { + let input = permute_nchw_to_nhwc(input); + let weight = permute_nchw_to_nhwc(weight); + + let out = conv_forward_nhwc(input, weight, bias, options, strategy)?; + + Ok(permute_nhwc_to_nchw(out)) +} + +/// Performs an N-dimensional convolution with the given strategy on NHWC inputs/outputs +/// +/// * `input` - The input feature map +/// * `weight` - The weights (filter) applied to each kernel +/// * `bias` - The bias added to each channel +/// * `options` - The options to use for the convolution +/// * `strategy` - The convolution algorithm to use. Autotune will pick the fastest available option. +pub fn conv_forward_nhwc( + input: CubeTensor, + weight: CubeTensor, + bias: Option>, + options: ConvOptions, + strategy: ConvStrategy, +) -> Result, ConvSetupError> { + match strategy { + ConvStrategy::Direct => conv_direct::(input, weight, bias, options), + #[cfg(feature = "autotune")] + ConvStrategy::Autotune => Ok(conv_autotune::(input, weight, bias, options)), + ConvStrategy::ImplicitGemm => { + if options.groups != 1 { + conv_direct::(input, weight, bias, options) + } else { + conv_gemm_simple_sync::( + input, + weight, + bias, + options, + AcceleratedTileKind::Cmma, + ) + } + } + } +} + +/// Performs an N-dimensional convolution backwards pass with regard to weight, with the given strategy +/// +/// * `input` - The input feature map +/// * `out_grad` - The output gradients +/// * `weight_shape` - The shape of the weights/weight gradients +/// * `options` - The options used for the convolution +/// * `strategy` - The convolution algorithm to use. Autotune will pick the fastest available option. +pub fn conv_weight_backward( + input: CubeTensor, + out_grad: CubeTensor, + weight_shape: Shape, + options: ConvOptions, + strategy: ConvStrategy, +) -> Result, ConvSetupError> { + let input = permute_nchw_to_nhwc(input); + let out_grad = permute_nchw_to_nhwc(out_grad); + let weight_shape = permute_nchw_to_nhwc_shape(weight_shape); + + let weight_grad = match strategy { + ConvStrategy::Direct => { + conv_weight_backward_fallback::(input, out_grad, weight_shape, options) + } + #[cfg(feature = "autotune")] + ConvStrategy::Autotune => Ok(wgrad_autotune::( + input, + out_grad, + weight_shape, + options, + )), + ConvStrategy::ImplicitGemm => { + if options.groups != 1 { + conv_weight_backward_fallback::(input, out_grad, weight_shape, options) + } else { + wgrad_gemm_simple_sync::( + input, + out_grad, + weight_shape, + options, + AcceleratedTileKind::Cmma, + ) + } + } + }?; + + Ok(permute_nhwc_to_nchw(weight_grad)) +} + +/// Performs an N-dimensional convolution backwards data pass with the given strategy +/// +/// * `input` - The input feature map +/// * `weight` - The weights (filter) applied to each kernel +/// * `in_shape` - The shape of the input to the layer +/// * `options` - The options to use for the convolution +/// * `strategy` - The convolution algorithm to use. Autotune will pick the fastest available option. +pub fn conv_data_backward( + out_grad: CubeTensor, + weights: CubeTensor, + in_shape: Shape, + options: ConvOptions, + strategy: ConvStrategy, +) -> Result, ConvSetupError> { + let out_grad = permute_nchw_to_nhwc(out_grad); + let weights = permute_nchw_to_nhwc(weights); + let in_shape = permute_nchw_to_nhwc_shape(in_shape); + + let weight_grad = match strategy { + ConvStrategy::Direct => { + conv_data_backward_fallback::(out_grad, weights, in_shape, options)? + } + #[cfg(feature = "autotune")] + ConvStrategy::Autotune => dgrad_autotune::(out_grad, weights, in_shape, options), + ConvStrategy::ImplicitGemm => { + if options.groups != 1 || options.stride.iter().any(|&s| s != 1) { + conv_data_backward_fallback::(out_grad, weights, in_shape, options)? + } else { + dgrad_gemm_simple_sync::( + out_grad, + weights, + in_shape, + options, + AcceleratedTileKind::Cmma, + )? + } + } + }; + + Ok(permute_nhwc_to_nchw(weight_grad)) +} diff --git a/crates/burn-cubecl/src/kernel/conv/conv_transpose2d/base.rs b/crates/burn-cubecl/src/kernel/conv/conv_transpose2d/base.rs new file mode 100644 index 0000000..5fa9ed2 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/conv_transpose2d/base.rs @@ -0,0 +1,54 @@ +use crate::{CubeRuntime, tensor::CubeTensor}; +use burn_backend::ops::ConvTransposeOptions; +use cubek::convolution::components::ConvSetupError; + +#[cfg(feature = "autotune")] +use super::conv_transpose2d_autotune; +use super::{conv_transpose2d_col2im, conv_transpose2d_direct}; + +/// The strategy to be used when launching a conv_transpose kernel. +pub enum ConvTranspose2dStrategy { + /// A simple direct convolution. + Direct, + #[cfg(feature = "autotune")] + /// Using autotune to choose the best kernel based on runtime information. + Autotune, + /// GEMM (im2col) based implementation of convolution. Significantly increased memory usage. + Gemm, +} + +impl Default for ConvTranspose2dStrategy { + fn default() -> Self { + // if autotune is enabled, default to autotune + #[cfg(feature = "autotune")] + return ConvTranspose2dStrategy::Autotune; + + // if autotune is disabled, default to the more memory-conservative algorithm + #[cfg(not(feature = "autotune"))] + ConvTranspose2dStrategy::Direct + } +} + +/// Performs a 2D convolution with the given strategy +/// +/// * `input` - The input feature map +/// * `weight` - The weights (filter) applied to each kernel +/// * `bias` - The bias added to each channel +/// * `options` - The options to use for the convolution +/// * `strategy` - The convolution algorithm to use. Autotune will pick the fastest available option. +pub fn conv_transpose2d( + input: CubeTensor, + weight: CubeTensor, + bias: Option>, + options: ConvTransposeOptions<2>, + strategy: ConvTranspose2dStrategy, +) -> Result, ConvSetupError> { + match strategy { + ConvTranspose2dStrategy::Direct => conv_transpose2d_direct(input, weight, bias, options), + #[cfg(feature = "autotune")] + ConvTranspose2dStrategy::Autotune => { + Ok(conv_transpose2d_autotune(input, weight, bias, options)) + } + ConvTranspose2dStrategy::Gemm => conv_transpose2d_col2im(input, weight, bias, options), + } +} diff --git a/crates/burn-cubecl/src/kernel/conv/conv_transpose2d/col2im.rs b/crates/burn-cubecl/src/kernel/conv/conv_transpose2d/col2im.rs new file mode 100644 index 0000000..cec7846 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/conv_transpose2d/col2im.rs @@ -0,0 +1,307 @@ +use crate::{ + CubeRuntime, + kernel::{ + conv::batches_per_run, + into_contiguous_aligned, + matmul::{MatmulStrategy, matmul}, + slice, + utils::{address_type, decompose_linear, shape_divmod}, + }, + ops::{numeric::empty_device_dtype, reshape, swap_dims}, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{ + Shape, + ops::{ConvTransposeOptions, conv::calculate_conv_transpose_output_size}, +}; +use cubecl::{ + calculate_cube_count_elemwise, + prelude::*, + std::{FastDivmod, tensor::layout::linear::LinearViewMut}, +}; +use cubek::convolution::components::ConvSetupError; + +/// Perform a 2D convolution transposition using the GEMM (col2im) algorithm. +/// +/// * `input` - The input feature map +/// * `weight` - The weights (filter) applied to each kernel +/// * `bias` - The bias added to each channel +/// * `options` - The options to use for the convolution +pub fn conv_transpose2d_col2im( + input: CubeTensor, + weight: CubeTensor, + bias: Option>, + options: ConvTransposeOptions<2>, +) -> Result, ConvSetupError> { + let [input_channels, im_ch_per_group, kernel_h, kernel_w] = weight.meta.shape().dims(); + let [batch_size, _, input_h, input_w] = input.meta.shape().dims(); + let groups = options.groups; + let input_ch_per_group = input_channels / groups; + let ConvTransposeOptions { + padding: [padding_h, padding_w], + padding_out: [padding_out_h, padding_out_w], + dilation: [dilation_h, dilation_w], + stride: [stride_h, stride_w], + .. + } = options.clone(); + + let im_h = calculate_conv_transpose_output_size( + kernel_h, + stride_h, + padding_h, + padding_out_h, + dilation_h, + input_h, + ); + let im_w = calculate_conv_transpose_output_size( + kernel_w, + stride_w, + padding_w, + padding_out_w, + dilation_w, + input_w, + ); + let im_channels = im_ch_per_group * groups; + + let batches_per_run = batches_per_run( + batch_size, + input_h * input_w, + input.client.properties().hardware.plane_size_max as usize, + )?; + let col_shape_0 = im_ch_per_group * kernel_h * kernel_w; + + let weight = reshape( + weight.clone(), + Shape::new([groups, input_ch_per_group, col_shape_0]), + ); + let weight = into_contiguous_aligned(swap_dims(weight, 1, 2)); + + if batches_per_run != batch_size { + let runs = batch_size / batches_per_run; + + let im_shape = Shape::new([runs, batches_per_run, im_channels, im_h, im_w]); + let image = empty_device_dtype( + input.client.clone(), + input.device.clone(), + im_shape, + input.dtype, + ); + + let input_shape = Shape::new([runs, batches_per_run, input_channels, input_h, input_w]); + let input = reshape(input, input_shape); + let input_shape_run = Shape::new([batches_per_run, input_channels, input_h, input_w]); + + for run in 0..runs { + let input = index(input.clone(), run); + let input = reshape(input, input_shape_run.clone()); + let im_shape = Shape::new([batches_per_run, im_channels, im_h, im_w]); + let image_slice = index(image.clone(), run); + let image_slice = reshape(image_slice, im_shape); + execute( + input, + weight.clone(), + bias.clone(), + image_slice, + options.clone(), + kernel_h, + kernel_w, + )?; + } + Ok(reshape( + image, + Shape::new([batch_size, im_channels, im_h, im_w]), + )) + } else { + let im_shape = Shape::new([batches_per_run, im_channels, im_h, im_w]); + let image = empty_device_dtype( + input.client.clone(), + input.device.clone(), + im_shape, + input.dtype, + ); + execute( + input, + weight, + bias, + image.clone(), + options, + kernel_h, + kernel_w, + )?; + Ok(image) + } +} + +pub(crate) fn index(tensor: CubeTensor, i: usize) -> CubeTensor { + #[allow(clippy::single_range_in_vec_init)] + let mut indices = vec![i..i + 1]; + for dim in tensor.meta.shape()[1..].iter() { + indices.push(0..*dim); + } + let mut tensor = slice(tensor, &indices); + tensor.meta.remove(0); + tensor +} + +#[allow(clippy::too_many_arguments)] +fn execute( + input: CubeTensor, + weight: CubeTensor, + bias: Option>, + image: CubeTensor, + options: ConvTransposeOptions<2>, + kernel_h: usize, + kernel_w: usize, +) -> Result<(), ConvSetupError> { + let [batch_size, _, input_h, input_w] = input.meta.shape().dims(); + let [groups, col_shape_0, input_ch_per_group] = weight.meta.shape().dims(); + + let col_shape_1 = batch_size * input_h * input_w; + + let input = swap_dims(input, 0, 1); + let input_shape = Shape::new([groups, input_ch_per_group, col_shape_1]); + let input = reshape(input, input_shape); + + let dtype = input.dtype; + let columns = matmul(weight, input, None, MatmulStrategy::default(), dtype)?; + let columns = reshape(columns, Shape::new([col_shape_0 * groups, col_shape_1])); + + col2im( + columns, bias, image, kernel_h, kernel_w, input_h, input_w, options, + )?; + + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +fn col2im( + columns: CubeTensor, + bias: Option>, + out: CubeTensor, + kernel_h: usize, + kernel_w: usize, + out_h: usize, + out_w: usize, + options: ConvTransposeOptions<2>, +) -> Result<(), LaunchError> { + let dtype = columns.dtype; + + let columns = into_contiguous_aligned(columns); + let bias = bias.map(into_contiguous_aligned); + + let num_elems = out.meta.num_elements(); + + let cube_dim = CubeDim::new(&columns.client, num_elems); + let cube_count = calculate_cube_count_elemwise(&columns.client, num_elems, cube_dim); + + let shape = shape_divmod(&out); + unsafe { + col2im_kernel::launch_unchecked( + &columns.client.clone(), + cube_count, + cube_dim, + address_type!(columns, bias, out), + columns.into_tensor_arg(), + bias.map(|bias| bias.into_buffer_arg()).into(), + out.into_linear_view(), + shape, + Col2ImArgsLaunch::new( + out_h, + out_w, + kernel_h, + kernel_w, + options.padding[0], + options.padding[1], + options.dilation[0], + options.dilation[1], + options.stride[0], + options.stride[1], + ), + dtype_to_storage_type(dtype), + ) + }; + + Ok(()) +} + +#[derive(CubeLaunch, CubeType)] +struct Col2ImArgs { + out_h: usize, + out_w: usize, + + kernel_h: usize, + kernel_w: usize, + + pad_h: usize, + pad_w: usize, + dilation_h: usize, + dilation_w: usize, + stride_h: usize, + stride_w: usize, +} + +#[cube(launch_unchecked, address_type = "dynamic")] +fn col2im_kernel( + columns: &Tensor, + bias: ComptimeOption<&[E]>, + mut image: LinearViewMut<'_, E>, + image_shape: Sequence>, + args: &Col2ImArgs, + #[define(E)] _dtype: StorageType, +) { + if ABSOLUTE_POS >= image.shape() { + terminate!(); + } + + let (_, pos) = decompose_linear(ABSOLUTE_POS, &image_shape); + let [batch, ch_im, im_y, im_x] = *pos else { + unreachable!() + }; + + let im_x = im_x + args.pad_w; + let im_y = im_y + args.pad_h; + + let kernel_extent_w = (args.kernel_w - 1) * args.dilation_w + 1; + let kernel_extent_h = (args.kernel_h - 1) * args.dilation_h + 1; + + let mut val = E::zero(); + + let x_col_start = if im_x >= kernel_extent_w { + (im_x - kernel_extent_w) / args.stride_w + 1 + } else { + 0usize.runtime() + }; + let x_col_end = clamp_max(im_x / args.stride_w + 1, args.out_w); + let y_col_start = if im_y >= kernel_extent_h { + (im_y - kernel_extent_h) / args.stride_h + 1 + } else { + 0usize.runtime() + }; + let y_col_end = clamp_max(im_y / args.stride_h + 1, args.out_h); + + for col_y in y_col_start..y_col_end { + let kernel_y = im_y - col_y * args.stride_h; + for col_x in x_col_start..x_col_end { + let kernel_x = im_x - col_x * args.stride_w; + + if kernel_y.is_multiple_of(args.dilation_h) && kernel_x.is_multiple_of(args.dilation_w) + { + let kernel_y = kernel_y / args.dilation_h; + let kernel_x = kernel_x / args.dilation_w; + + let col_k = + ch_im * args.kernel_h * args.kernel_w + kernel_y * args.kernel_w + kernel_x; + let col_n = batch * args.out_h * args.out_w + col_y * args.out_w + col_x; + let col_pos = col_k * columns.stride(0) + col_n * columns.stride(1); + val += columns[col_pos]; + } + } + } + + #[comptime] + match bias { + ComptimeOption::Some(bias) => image.write(ABSOLUTE_POS, val + bias[ch_im]), + ComptimeOption::None => image.write(ABSOLUTE_POS, val), + } +} diff --git a/crates/burn-cubecl/src/kernel/conv/conv_transpose2d/mod.rs b/crates/burn-cubecl/src/kernel/conv/conv_transpose2d/mod.rs new file mode 100644 index 0000000..214333f --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/conv_transpose2d/mod.rs @@ -0,0 +1,15 @@ +mod base; +mod col2im; + +mod transpose_direct; + +#[cfg(feature = "autotune")] +mod tune; + +pub use base::*; +pub use col2im::*; + +pub use transpose_direct::*; + +#[cfg(feature = "autotune")] +pub use tune::*; diff --git a/crates/burn-cubecl/src/kernel/conv/conv_transpose2d/transpose_direct.rs b/crates/burn-cubecl/src/kernel/conv/conv_transpose2d/transpose_direct.rs new file mode 100644 index 0000000..7bd33c2 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/conv_transpose2d/transpose_direct.rs @@ -0,0 +1,187 @@ +use crate::{ + CubeRuntime, + kernel::utils::{address_type, decompose_linear, shape_divmod}, + ops::numeric::empty_device_dtype, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{Shape, ops::ConvTransposeOptions}; +use cubecl::{ + calculate_cube_count_elemwise, + prelude::*, + std::{FastDivmod, tensor::layout::linear::LinearViewMut}, +}; +use cubek::convolution::components::ConvSetupError; + +#[derive(CubeLaunch, CubeType)] +struct ConvArgs { + conv_stride_0: usize, + conv_stride_1: usize, + dilation_0: usize, + dilation_1: usize, + padding_0: usize, + padding_1: usize, + groups: usize, +} + +#[cube(launch, address_type = "dynamic")] +fn conv_transpose2d_direct_kernel( + input: &Tensor, + weight: &Tensor, + bias: ComptimeOption<&[E]>, + mut output: LinearViewMut<'_, E>, + out_shape: Sequence>, + args: ConvArgs, + #[define(E)] _dtype: StorageType, +) { + if ABSOLUTE_POS >= output.shape() { + terminate!(); + } + + let in_c_per_group = weight.shape(0) / args.groups; + let out_c_per_group = weight.shape(1); + let kernel_h = weight.shape(2); + let kernel_w = weight.shape(3); + + let (_, pos) = decompose_linear(ABSOLUTE_POS, &out_shape); + let [batch, oc_out, out_y, out_x] = *pos else { + unreachable!() + }; + + let k = oc_out / out_c_per_group; + let group = k % args.groups; + let out_c = oc_out - out_c_per_group * group; + + let in_c_start = group * in_c_per_group; + let in_c_end = in_c_start + in_c_per_group; + + let stride_0_i = args.conv_stride_0 as i32; + let stride_1_i = args.conv_stride_1 as i32; + + let kms_h = (kernel_h * args.dilation_0) as i32 - stride_0_i; + let kms_w = (kernel_w * args.dilation_1) as i32 - stride_1_i; + + let y_start = ((out_y + args.padding_0) as i32 - kms_h) / stride_0_i; + let x_start = ((out_x + args.padding_1) as i32 - kms_w) / stride_1_i; + + let y_end = clamp(kms_h + y_start + 1, 0, input.shape(2) as i32) as usize; + let x_end = clamp(kms_w + x_start + 1, 0, input.shape(3) as i32) as usize; + let y_start = clamp_min(y_start, 0) as usize; + let x_start = clamp_min(x_start, 0) as usize; + + let idx_input_batch = batch * input.stride(0); + let idx_weight_oc = out_c * weight.stride(1); + + let bias: ComptimeOption = bias.as_ref().map(|bias| bias[oc_out]); + let mut sum = bias.unwrap_or_default(); + + let numerator_h_base = out_y + args.padding_0; + let numerator_w_base = out_x + args.padding_1; + + for in_c in in_c_start..in_c_end { + let idx_input_ic = in_c * input.stride(1); + let idx_weight_ic = in_c * weight.stride(0); + + for in_y in y_start..y_end { + let numerator_tmp = in_y * args.conv_stride_0; + let numerator_h = numerator_h_base - numerator_tmp; + + if numerator_h_base >= numerator_tmp && numerator_h.is_multiple_of(args.dilation_0) { + let kernel_y = numerator_h / args.dilation_0; + let idx_input_y = in_y * input.stride(2); + let idx_weight_ky = kernel_y * weight.stride(2); + + for in_x in x_start..x_end { + let numerator_tmp = in_x * args.conv_stride_1; + let numerator_w = numerator_w_base - numerator_tmp; + + if numerator_w_base >= numerator_tmp + && numerator_w.is_multiple_of(args.dilation_1) + { + let kernel_x = numerator_w / args.dilation_1; + let idx_input_x = in_x * input.stride(3); + let idx_weight_kx = kernel_x * weight.stride(3); + + let index_input = + idx_input_batch + idx_input_ic + idx_input_y + idx_input_x; + let index_weight = + idx_weight_ic + idx_weight_oc + idx_weight_ky + idx_weight_kx; + + let value = input[index_input]; + let weight = weight[index_weight]; + + sum += value * weight; + } + } + } + } + } + + output.write(ABSOLUTE_POS, sum); +} + +/// Perform a 2D convolution transposition using the direct algorithm. +/// +/// * `input` - The input feature map +/// * `weight` - The weights (filter) applied to each kernel +/// * `bias` - The bias added to each channel +/// * `options` - The options to use for the convolution +/// +pub fn conv_transpose2d_direct( + input: CubeTensor, + weight: CubeTensor, + bias: Option>, + options: ConvTransposeOptions<2>, +) -> Result, ConvSetupError> { + let [batch_size, _, in_height, in_width] = input.meta.shape().dims(); + let [_, out_channels, kernel_0, kernel_1] = weight.meta.shape().dims(); + + let out_0 = (in_height - 1) * options.stride[0] + + options.dilation[0] * (kernel_0 - 1) + + options.padding_out[0] + - 2 * options.padding[0] + + 1; + let out_1 = (in_width - 1) * options.stride[1] + + options.dilation[1] * (kernel_1 - 1) + + options.padding_out[1] + - 2 * options.padding[1] + + 1; + + let shape_out = Shape::new([batch_size, out_channels * options.groups, out_0, out_1]); + + let output = empty_device_dtype( + input.client.clone(), + input.device.clone(), + shape_out.clone(), + input.dtype, + ); + + let num_elems = output.meta.num_elements(); + let cube_dim = CubeDim::new(&input.client, num_elems); + let cube_count = calculate_cube_count_elemwise(&input.client, num_elems, cube_dim); + let dtype = input.dtype; + + conv_transpose2d_direct_kernel::launch( + &output.client, + cube_count, + cube_dim, + address_type!(input, weight, bias, output), + input.into_tensor_arg(), + weight.into_tensor_arg(), + bias.map(|bias| bias.into_buffer_arg()).into(), + output.clone().into_linear_view(), + shape_divmod(&output), + ConvArgsLaunch::new( + options.stride[0], + options.stride[1], + options.dilation[0], + options.dilation[1], + options.padding[0], + options.padding[1], + options.groups, + ), + dtype_to_storage_type(dtype), + ); + + Ok(output) +} diff --git a/crates/burn-cubecl/src/kernel/conv/conv_transpose2d/tune.rs b/crates/burn-cubecl/src/kernel/conv/conv_transpose2d/tune.rs new file mode 100644 index 0000000..e799c09 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/conv_transpose2d/tune.rs @@ -0,0 +1,99 @@ +use burn_backend::ops::ConvTransposeOptions; +use cubecl::tune::{LocalTuner, Tunable, TunableSet, local_tuner}; + +use crate::{ + CubeAutotuneKey, CubeRuntime, CubeTuneId, + kernel::conv::{ConvTranspose2dAutotuneKey, conv_transpose2d_col2im, conv_transpose2d_direct}, + tensor::CubeTensor, +}; + +/// Executes autotune on conv2d operations +pub fn conv_transpose2d_autotune( + input: CubeTensor, + weights: CubeTensor, + bias: Option>, + options: ConvTransposeOptions<2>, +) -> CubeTensor { + let client = input.client.clone(); + + static TUNER: LocalTuner = local_tuner!(); + + let tune_set = TUNER.init(|| { + TunableSet::new(create_key::, create_transpose2d_input::) + .with(Tunable::new( + "conv_transpose2d_direct", + |(input, weights, bias, options)| { + conv_transpose2d_direct::(input, weights, bias, options) + }, + )) + .with(Tunable::new( + "conv_transpose2d_col2im", + |(input, weights, bias, options)| { + conv_transpose2d_col2im::(input, weights, bias, options) + }, + )) + }); + + TUNER.execute( + &CubeTuneId::new(&input.client, &input.device), + &client, + tune_set, + (input, weights, bias, options), + ) +} + +pub fn create_transpose2d_input( + _key: &CubeAutotuneKey, + (input, weights, bias, options): &( + CubeTensor, + CubeTensor, + Option>, + ConvTransposeOptions<2>, + ), +) -> ( + CubeTensor, + CubeTensor, + Option>, + ConvTransposeOptions<2>, +) { + ( + input.clone(), + weights.clone(), + bias.clone(), + options.clone(), + ) +} + +fn create_key( + (input, weights, bias, options): &( + CubeTensor, + CubeTensor, + Option>, + ConvTransposeOptions<2>, + ), +) -> CubeAutotuneKey { + let [batch_size, in_channels, height, width] = input.meta.shape().dims(); + let [out_channels, _, kernel_h, kernel_w] = weights.meta.shape().dims(); + let ConvTransposeOptions { + stride, + padding, + dilation, + groups, + padding_out, + } = options.clone(); + CubeAutotuneKey::ConvTranspose(ConvTranspose2dAutotuneKey::new( + [kernel_h, kernel_w], + stride, + padding, + padding_out, + dilation, + groups, + in_channels, + out_channels, + height, + width, + batch_size, + bias.is_some(), + input.dtype, + )) +} diff --git a/crates/burn-cubecl/src/kernel/conv/conv_transpose3d.rs b/crates/burn-cubecl/src/kernel/conv/conv_transpose3d.rs new file mode 100644 index 0000000..79eadd5 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/conv_transpose3d.rs @@ -0,0 +1,228 @@ +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::{ + calculate_cube_count_elemwise, + prelude::*, + std::{FastDivmod, tensor::layout::linear::LinearViewMut}, +}; + +use crate::{ + CubeRuntime, + kernel::utils::{address_type, decompose_linear, shape_divmod}, + ops::numeric::empty_device_dtype, + tensor::CubeTensor, +}; +use burn_backend::{Shape, ops::ConvTransposeOptions}; + +#[derive(CubeLaunch, CubeType)] +struct ConvArgs { + conv_stride_0: usize, + conv_stride_1: usize, + conv_stride_2: usize, + dilation_0: usize, + dilation_1: usize, + dilation_2: usize, + padding_0: usize, + padding_1: usize, + padding_2: usize, + groups: usize, +} + +#[cube(launch, address_type = "dynamic")] +fn conv_transpose3d_kernel( + input: &Tensor, + weight: &Tensor, + bias: ComptimeOption<&[E]>, + mut output: LinearViewMut<'_, E>, + out_shape: Sequence>, + args: ConvArgs, + #[define(E)] _dtype: StorageType, +) { + if !output.is_in_bounds(ABSOLUTE_POS) { + terminate!() + } + + let in_channels = weight.shape(0); + let out_c_per_group = weight.shape(1); + let kernel_size_0 = weight.shape(2); + let kernel_size_1 = weight.shape(3); + let kernel_size_2 = weight.shape(4); + + let stride_0_i = args.conv_stride_0 as i32; + let stride_1_i = args.conv_stride_1 as i32; + let stride_2_i = args.conv_stride_2 as i32; + + let (_, pos) = decompose_linear(ABSOLUTE_POS, &out_shape); + let [batch, out_c_out, out_z, out_y, out_x] = *pos else { + unreachable!() + }; + + let groups = args.groups; + let in_c_per_group = in_channels / groups; + + let k = out_c_out / out_c_per_group; + let group = k % groups; + let out_channel = out_c_out - out_c_per_group * group; + + let in_c_start = group * in_c_per_group; + let in_c_end = in_c_start + in_c_per_group; + + let kernel_d = (kernel_size_0 * args.dilation_0 - args.conv_stride_0) as i32; + let kernel_h = (kernel_size_1 * args.dilation_1 - args.conv_stride_1) as i32; + let kernel_w = (kernel_size_2 * args.dilation_2 - args.conv_stride_2) as i32; + + let z_start = ((out_z + args.padding_0) as i32 - kernel_d) / stride_0_i; + let y_start = ((out_y + args.padding_1) as i32 - kernel_h) / stride_1_i; + let x_start = ((out_x + args.padding_2) as i32 - kernel_w) / stride_2_i; + + let z_end = clamp(kernel_d + z_start + 1, 0, input.shape(2) as i32) as usize; + let y_end = clamp(kernel_h + y_start + 1, 0, input.shape(3) as i32) as usize; + let x_end = clamp(kernel_w + x_start + 1, 0, input.shape(4) as i32) as usize; + + let z_start = clamp_min(z_start, 0) as usize; + let y_start = clamp_min(y_start, 0) as usize; + let x_start = clamp_min(x_start, 0) as usize; + + let index_input_batch = batch * input.stride(0); + let index_weight_out_c = out_channel * weight.stride(1); + + let bias: ComptimeOption = bias.as_ref().map(|bias| bias[out_c_out]); + let mut sum = bias.unwrap_or_default(); + + let numerator_d_base = out_z + args.padding_0; + let numerator_h_base = out_y + args.padding_1; + let numerator_w_base = out_x + args.padding_2; + + for in_c in in_c_start..in_c_end { + let index_input_in_c = in_c * input.stride(1); + let index_weight_in_c = in_c * weight.stride(0); + + for in_z in z_start..z_end { + let numerator_tmp = in_z * args.conv_stride_0; + let numerator_d = numerator_d_base - numerator_tmp; + + if numerator_d_base >= numerator_tmp && numerator_d.is_multiple_of(args.dilation_0) { + let kernel_z = numerator_d / args.dilation_0; + let index_input_z = in_z * input.stride(2); + let index_weight_kz = kernel_z * weight.stride(2); + + for in_y in y_start..y_end { + let numerator_tmp = in_y * args.conv_stride_1; + let numerator_h = numerator_h_base - numerator_tmp; + + if numerator_h_base >= numerator_tmp + && numerator_h.is_multiple_of(args.dilation_1) + { + let kernel_y = numerator_h / args.dilation_1; + let index_input_y = in_y * input.stride(3); + let index_weight_ky = kernel_y * weight.stride(3); + + for in_x in x_start..x_end { + let numerator_tmp = in_x * args.conv_stride_2; + let numerator_w = numerator_w_base - numerator_tmp; + + if numerator_w_base >= numerator_tmp + && numerator_w.is_multiple_of(args.dilation_2) + { + let kernel_x = numerator_w / args.dilation_2; + let index_input_x = in_x * input.stride(4); + let index_weight_kx = kernel_x * weight.stride(4); + + let index_input = index_input_batch + + index_input_in_c + + index_input_z + + index_input_y + + index_input_x; + + let index_weight = index_weight_in_c + + index_weight_out_c + + index_weight_kz + + index_weight_ky + + index_weight_kx; + + let value = input[index_input]; + let weight = weight[index_weight]; + + sum += value * weight; + } + } + } + } + } + } + } + + output.write(ABSOLUTE_POS, sum); +} + +pub(crate) fn conv_transpose3d( + input: CubeTensor, + weight: CubeTensor, + bias: Option>, + options: ConvTransposeOptions<3>, +) -> Result, LaunchError> { + let [batch_size, _, in_depth, in_height, in_width] = input.meta.shape().dims(); + let [_, out_channels, kernel_0, kernel_1, kernel_2] = weight.meta.shape().dims(); + + let out_0 = (in_depth - 1) * options.stride[0] + + options.dilation[0] * (kernel_0 - 1) + + options.padding_out[0] + - 2 * options.padding[0] + + 1; + let out_1 = (in_height - 1) * options.stride[1] + + options.dilation[1] * (kernel_1 - 1) + + options.padding_out[1] + - 2 * options.padding[1] + + 1; + let out_2 = (in_width - 1) * options.stride[2] + + options.dilation[2] * (kernel_2 - 1) + + options.padding_out[2] + - 2 * options.padding[2] + + 1; + + let shape_out = Shape::new([ + batch_size, + out_channels * options.groups, + out_0, + out_1, + out_2, + ]); + + let output = empty_device_dtype( + input.client.clone(), + input.device.clone(), + shape_out.clone(), + input.dtype, + ); + + let num_elems = output.meta.num_elements(); + let cube_dim = CubeDim::new(&input.client, num_elems); + let cube_count = calculate_cube_count_elemwise(&input.client, num_elems, cube_dim); + + let dtype = input.dtype; + conv_transpose3d_kernel::launch( + &output.client, + cube_count, + cube_dim, + address_type!(input, weight, bias, output), + input.into_tensor_arg(), + weight.into_tensor_arg(), + bias.map(|bias| bias.into_buffer_arg()).into(), + output.clone().into_linear_view(), + shape_divmod(&output), + ConvArgsLaunch::new( + options.stride[0], + options.stride[1], + options.stride[2], + options.dilation[0], + options.dilation[1], + options.dilation[2], + options.padding[0], + options.padding[1], + options.padding[2], + options.groups, + ), + dtype_to_storage_type(dtype), + ); + + Ok(output) +} diff --git a/crates/burn-cubecl/src/kernel/conv/deform_conv2d.rs b/crates/burn-cubecl/src/kernel/conv/deform_conv2d.rs new file mode 100644 index 0000000..79e0ce5 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/deform_conv2d.rs @@ -0,0 +1,311 @@ +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::{calculate_cube_count_elemwise, prelude::*, std::FastDivmod}; +use cubek::convolution::components::ConvSetupError; + +use burn_backend::{ + Shape, + ops::{DeformConvOptions, conv::calculate_conv_output_size}, +}; + +use crate::{ + CubeRuntime, + kernel::{ + AddOp, into_contiguous_aligned, launch_binop, + matmul::{MatmulStrategy, matmul}, + utils::address_type, + }, + ops::{numeric::zeros_client, reshape, swap_dims}, + tensor::CubeTensor, +}; + +#[derive(CubeLaunch, CubeType)] +struct DeformConv2dArgs { + conv_stride_h: usize, + conv_stride_w: usize, + dilation_h: usize, + dilation_w: usize, + padding_h: InputScalar, + padding_w: InputScalar, + offset_groups: usize, + + kernel_height: usize, + kernel_width: usize, + out_h: usize, + out_w: usize, +} + +#[cube(launch, address_type = "dynamic")] +fn deform_im2col_kernel( + input: &Tensor, + offset: &Tensor, + mask: ComptimeOption<&Tensor>, + columns: &mut Tensor, + pos_shape: Sequence>, + args: &DeformConv2dArgs, + #[comptime] kernel_h_unroll: Option, + #[comptime] kernel_w_unroll: Option, + #[define(F)] _dtype: StorageType, +) { + // position shape: [in_channels, batch_size, out_h, out_w] + // columns shape: [[in_channels, kernel_h, kernel_w], [batch_size, out_h, out_w]] + + let kernel_height = kernel_h_unroll.unwrap_or(args.kernel_height); + let unroll_h = kernel_h_unroll.is_some(); + let kernel_width = kernel_w_unroll.unwrap_or(args.kernel_width); + let unroll_w = kernel_w_unroll.is_some(); + + let out_h = args.out_h; + let out_w = args.out_w; + let in_channels = input.shape(1); + let height = input.shape(2); + let width = input.shape(3); + let col_stride_0 = columns.stride(0); + + let (rem, out_x) = pos_shape[3].div_mod(ABSOLUTE_POS); + let (rem, out_y) = pos_shape[2].div_mod(rem); + let (in_channel, batch) = pos_shape[1].div_mod(rem); + + if in_channel >= in_channels { + terminate!() + } + + let out_k_base = in_channel * kernel_height * kernel_width; + let out_n = batch * out_h * out_w + out_y * out_w + out_x; + + let channels_per_offset_group = in_channels / args.offset_groups; + let group_index = in_channel / channels_per_offset_group; + + let mut col_base_idx = out_k_base * columns.stride(0) + out_n * columns.stride(1); + + let input_base_idx = batch * input.stride(0) + in_channel * input.stride(1); + + let offset_base_idx = batch * offset.stride(0) + + group_index * kernel_height * kernel_width * 2 * offset.stride(1); + + let mask_base_idx = mask.as_ref().map(|mask| { + batch * mask.stride(0) + group_index * kernel_height * kernel_width * mask.stride(1) + }); + + #[unroll(unroll_h)] + for kernel_y in 0..kernel_height { + #[unroll(unroll_w)] + for kernel_x in 0..kernel_width { + let mask_index = kernel_y * kernel_width + kernel_x; + let offset_index = mask_index * 2; + + let offset_y = offset[offset_base_idx + + offset_index * offset.stride(1) + + out_y * offset.stride(2) + + out_x * offset.stride(3)]; + let offset_x = offset[offset_base_idx + + (offset_index + 1) * offset.stride(1) + + out_y * offset.stride(2) + + out_x * offset.stride(3)]; + let y = F::cast_from(out_y * args.conv_stride_h + kernel_y * args.dilation_h) + - args.padding_h.get::() + + offset_y; + let x = F::cast_from(out_x * args.conv_stride_w + kernel_x * args.dilation_w) + - args.padding_w.get::() + + offset_x; + + let interpolated = bilinear_interpolate(input, height, width, y, x, input_base_idx); + #[comptime] + let value = match mask.zip::(mask_base_idx) { + ComptimeOption::Some((mask, base_idx)) => { + let mask_value = mask[base_idx + + mask_index * mask.stride(1) + + out_y * mask.stride(2) + + out_x * mask.stride(3)]; + mask_value * interpolated + } + ComptimeOption::None => interpolated, + }; + + columns[col_base_idx] = value; + col_base_idx += col_stride_0; + } + } +} + +#[cube] +pub(crate) fn bilinear_interpolate( + input: &Tensor, + height: usize, + width: usize, + y: F, + x: F, + offset: usize, +) -> F { + // To simplify code + let y = f32::cast_from(y); + let x = f32::cast_from(x); + let stride_y = input.stride(2); + let stride_x = input.stride(3); + + let mut result = F::new(0.0_f32); + if y > -1.0f32 && height as f32 > y && x > -1.0f32 && width as f32 > x { + let y_low = y.floor(); + let x_low = x.floor(); + let y_high = (y_low + 1.) as usize; + let x_high = (x_low + 1.) as usize; + + let zero = F::new(0.0_f32); + let v1: F = if y_low >= 0. && x_low >= 0. { + input[offset + y_low as usize * stride_y + x_low as usize * stride_x] + } else { + zero + }; + let v2: F = if y_low >= 0. && x_high < width { + input[offset + y_low as usize * stride_y + x_high * stride_x] + } else { + zero + }; + let v3: F = if y_high < height && x_low >= 0. { + input[offset + y_high * stride_y + x_low as usize * stride_x] + } else { + zero + }; + let v4: F = if y_high < height && x_high < width { + input[offset + y_high * stride_y + x_high * stride_x] + } else { + zero + }; + + let l_y = y - y_low; + let l_x = x - x_low; + let h_y = 1.0 - l_y; + let h_x = 1.0 - l_x; + + let w1 = F::cast_from(h_y * h_x); + let w2 = F::cast_from(h_y * l_x); + let w3 = F::cast_from(l_y * h_x); + let w4 = F::cast_from(l_y * l_x); + + result = w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4; + } + result +} + +pub(crate) fn deform_im2col( + input: CubeTensor, + offset: CubeTensor, + mask: Option>, + options: DeformConvOptions<2>, + out_dims: (usize, usize), + kernel_dims: (usize, usize), +) -> Result, LaunchError> { + let client = input.client.clone(); + let device = input.device.clone(); + let dtype = input.dtype; + + let [batch_size, in_channels, _, _] = input.meta.shape().dims(); + let (out_height, out_width) = out_dims; + let (kernel_height, kernel_width) = kernel_dims; + + let shape_out = Shape::new([ + in_channels * kernel_height * kernel_width, + batch_size * out_height * out_width, + ]); + + let pos_shape = [in_channels, batch_size, out_height, out_width] + .into_iter() + .collect(); + + let output = zeros_client(client.clone(), device.clone(), shape_out.clone(), dtype); + + let num_kernels = in_channels * batch_size * out_height * out_width; + let cube_dim = CubeDim::new(&input.client, num_kernels); + let cube_count = calculate_cube_count_elemwise(&input.client, num_kernels, cube_dim); + + deform_im2col_kernel::launch( + &output.client, + cube_count, + cube_dim, + address_type!(input, offset, mask, output), + input.into_tensor_arg(), + offset.into_tensor_arg(), + mask.map(|mask| mask.into_tensor_arg()).into(), + output.clone().binding().into_tensor_arg(), + pos_shape, + DeformConv2dArgsLaunch::new( + options.stride[0], + options.stride[1], + options.dilation[0], + options.dilation[1], + { + let val = options.padding[0] as f32; + InputScalar::new(val, dtype_to_storage_type(dtype)) + }, + { + let val = options.padding[1] as f32; + InputScalar::new(val, dtype_to_storage_type(dtype)) + }, + options.offset_groups, + kernel_height, + kernel_width, + out_height, + out_width, + ), + Some(kernel_height), + Some(kernel_width), + dtype_to_storage_type(dtype), + ); + + Ok(output) +} + +pub(crate) fn deform_conv2d( + input: CubeTensor, + offset: CubeTensor, + weight: CubeTensor, + mask: Option>, + bias: Option>, + options: DeformConvOptions<2>, +) -> Result, ConvSetupError> { + let input = into_contiguous_aligned(input); + let offset = into_contiguous_aligned(offset); + let weight = into_contiguous_aligned(weight); + let mask = mask.map(|it| into_contiguous_aligned(it)); + let bias = bias.map(|it| into_contiguous_aligned(it)); + + let [batch_size, _, in_height, in_width] = input.meta.shape().dims(); + let [out_channels, _, kernel_h, kernel_w] = weight.meta.shape().dims(); + let groups = options.weight_groups; + + let out_h = calculate_conv_output_size( + kernel_h, + options.stride[0], + options.padding[0], + options.dilation[0], + in_height, + ); + let out_w = calculate_conv_output_size( + kernel_w, + options.stride[1], + options.padding[1], + options.dilation[1], + in_width, + ); + let out_dims = (out_h, out_w); + + let columns = deform_im2col(input, offset, mask, options, out_dims, (kernel_h, kernel_w))?; + + let [col_size_0, col_size_1] = columns.meta.shape().dims(); + let col_size_0 = col_size_0 / groups; + let out_c_per_group = out_channels / groups; + + let dtype = weight.dtype; + let weight = reshape(weight, Shape::new([groups, out_c_per_group, col_size_0])); + let columns = reshape(columns, Shape::new([groups, col_size_0, col_size_1])); + let out = matmul(weight, columns, None, MatmulStrategy::default(), dtype)?; + + let out = reshape(out, Shape::new([out_channels, batch_size, out_h, out_w])); + let out = swap_dims(out, 0, 1); + + if let Some(bias) = bias { + let bias = reshape(bias, Shape::new([1, out_channels, 1, 1])); + Ok(launch_binop::(out, bias)) + } else { + Ok(out) + } +} diff --git a/crates/burn-cubecl/src/kernel/conv/deform_conv_transpose2d.rs b/crates/burn-cubecl/src/kernel/conv/deform_conv_transpose2d.rs new file mode 100644 index 0000000..ad8e180 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/deform_conv_transpose2d.rs @@ -0,0 +1,727 @@ +use super::{bilinear_interpolate, deform_im2col, index}; +use crate::{ + CubeRuntime, + kernel::{ + cast, into_contiguous_aligned, + matmul::{MatmulStrategy, matmul}, + reduce::reduce_dim, + slice_assign, + utils::{address_type, decompose_linear}, + }, + ops::{ + numeric::{empty_device_dtype, zeros_client}, + reshape, swap_dims, + }, + tensor::CubeTensor, +}; +use burn_backend::cubecl::{dtype_to_elem_type, dtype_to_storage_type}; +use burn_backend::{DType, Shape, TensorMetadata, ops::DeformConvOptions}; +use cubecl::{ + CubeDim, CubeLaunch, calculate_cube_count_elemwise, cube, + features::AtomicUsage, + ir::FloatKind, + prelude::*, + std::{ + FastDivmod, + tensor::layout::linear::{LinearView, LinearViewMut}, + }, +}; +use cubek::{ + convolution::components::ConvSetupError, + reduce::components::instructions::ReduceOperationConfig, +}; +use std::marker::PhantomData; + +/// Calculate the [deformable 2D convolution](crate::ops::ModuleOps::deform_conv2d) backward pass using convolutions. +#[allow( + clippy::single_range_in_vec_init, + clippy::type_complexity, + clippy::too_many_arguments +)] +pub(crate) fn deform_conv2d_backward( + input: CubeTensor, + offset: CubeTensor, + weight: CubeTensor, + mask: Option>, + bias: Option>, + out_grad: CubeTensor, + options: DeformConvOptions<2>, +) -> Result< + ( + CubeTensor, + CubeTensor, + CubeTensor, + Option>, + Option>, + ), + ConvSetupError, +> { + let [_, _, out_h, out_w] = out_grad.meta.shape().dims(); + let [_, _, kernel_h, kernel_w] = weight.meta.shape().dims(); + + let gradient_bias = bias.map(|bias| { + let grad = reduce_dim( + out_grad.clone(), + None, + 0, + Default::default(), + ReduceOperationConfig::Sum, + ) + .unwrap(); + let grad = reduce_dim( + grad, + None, + 2, + Default::default(), + ReduceOperationConfig::Sum, + ) + .unwrap(); + let grad = reduce_dim( + grad, + None, + 3, + Default::default(), + ReduceOperationConfig::Sum, + ) + .unwrap(); + + reshape(grad, bias.meta.shape.clone()) + }); + + let input = into_contiguous_aligned(input); + let offset = into_contiguous_aligned(offset); + let weight = into_contiguous_aligned(weight); + let mask = mask.map(|it| into_contiguous_aligned(it)); + + let (input_gradient, offset_gradient, mask_gradient) = backward_gradient_inputs( + input.clone(), + weight.clone(), + offset.clone(), + mask.clone(), + out_grad.clone(), + &options, + (kernel_h, kernel_w), + )?; + + let weight_grad = compute_weight_grad( + input, + offset, + mask, + out_grad, + options, + (kernel_h, kernel_w), + (out_h, out_w), + )?; + + Ok(( + input_gradient, + offset_gradient, + weight_grad, + mask_gradient, + gradient_bias, + )) +} + +fn compute_weight_grad( + input: CubeTensor, + offset: CubeTensor, + mask: Option>, + out_grad: CubeTensor, + options: DeformConvOptions<2>, + kernel_dims: (usize, usize), + out_dims: (usize, usize), +) -> Result, ConvSetupError> { + let [_, in_channels, _, _] = input.meta.shape().dims(); + let [_, out_channels, _, _] = out_grad.meta.shape().dims(); + let (kernel_h, kernel_w) = kernel_dims; + let groups = options.weight_groups; + let dtype = input.dtype; + + let in_c_per_group = in_channels / groups; + let out_c_per_group = out_channels / groups; + + let columns = deform_im2col(input, offset, mask, options, out_dims, kernel_dims)?; + let [col_size_0, col_size_1] = columns.meta.shape().dims(); + let col_size_0 = col_size_0 / groups; + + let out_grad = swap_dims(out_grad, 0, 1); + let out_grad = reshape(out_grad, Shape::new([groups, out_c_per_group, col_size_1])); + + let columns = reshape(columns, Shape::new([groups, col_size_0, col_size_1])); + let columns = swap_dims(columns, 1, 2); + + let grad_weight = matmul(out_grad, columns, None, MatmulStrategy::default(), dtype)?; + + Ok(reshape( + grad_weight, + Shape::new([out_channels, in_c_per_group, kernel_h, kernel_w]), + )) +} + +type InputGradients = (CubeTensor, CubeTensor, Option>); + +fn backward_gradient_inputs( + image: CubeTensor, + weight: CubeTensor, + offset: CubeTensor, + mask: Option>, + out_grad: CubeTensor, + options: &DeformConvOptions<2>, + kernel_dims: (usize, usize), +) -> Result, ConvSetupError> { + let client = out_grad.client.clone(); + let device = out_grad.device.clone(); + + let [out_channels, in_c_per_group, kernel_h, kernel_w] = weight.meta.shape().dims(); + let [batch_size, _, out_h, out_w] = out_grad.meta.shape().dims(); + + let groups = options.weight_groups; + let out_c_per_group = out_channels / groups; + + let col_shape_0 = in_c_per_group * kernel_h * kernel_w; + let col_shape_1 = batch_size * out_h * out_w; + let col_shape = Shape::new([groups, col_shape_0, col_shape_1]); + let mut columns = empty_device_dtype(client, device, col_shape, weight.dtype); + + let weight = reshape(weight, Shape::new([groups, out_c_per_group, col_shape_0])); + + let out_grad = swap_dims(out_grad, 0, 1); + let out_grad_shape = Shape::new([groups, out_c_per_group, col_shape_1]); + let out_grad = reshape(out_grad, out_grad_shape); + + for group in 0..groups { + let dtype = weight.dtype; + let weight = swap_dims(index(weight.clone(), group), 0, 1); + let out_grad = index(out_grad.clone(), group); + let values = matmul(weight, out_grad, None, MatmulStrategy::default(), dtype)?; + let values = reshape(values, Shape::new([1, col_shape_0, col_shape_1])); + columns = slice_assign( + columns, + &[ + burn_backend::Slice::from(group..group + 1), + burn_backend::Slice::from(0..col_shape_0), + burn_backend::Slice::from(0..col_shape_1), + ], + values, + ); + } + + let columns = reshape(columns, Shape::new([col_shape_0 * groups, col_shape_1])); + + let input_shape = image.shape(); + let (offset_gradient, mask_gradient) = compute_offset_and_mask_gradient( + columns.clone(), + image, + offset.clone(), + mask.clone(), + options, + kernel_dims, + )?; + + let input_gradient = + compute_input_grad(columns, offset, mask, options, kernel_dims, input_shape)?; + + Ok((input_gradient, offset_gradient, mask_gradient)) +} + +fn compute_offset_and_mask_gradient( + columns: CubeTensor, + image: CubeTensor, + offset: CubeTensor, + mask: Option>, + options: &DeformConvOptions<2>, + kernel_dims: (usize, usize), +) -> Result<(CubeTensor, Option>), ConvSetupError> { + let client = offset.client.clone(); + let device = offset.device.clone(); + let (kernel_h, kernel_w) = kernel_dims; + + let [batches, _, out_h, out_w] = offset.meta.shape().dims(); + let offset_groups = options.offset_groups; + + let pos_shape = [batches, offset_groups, kernel_h, kernel_w, 2, out_h, out_w]; + let pos_shape = pos_shape.into_iter().collect(); + + let grad_offset = + empty_device_dtype(client.clone(), device.clone(), offset.shape(), offset.dtype); + let grad_mask = mask + .as_ref() + .map(|mask| empty_device_dtype(client.clone(), device.clone(), mask.shape(), mask.dtype)); + + let num_elements_offset = offset.meta.num_elements(); + let cube_dim = CubeDim::new(&image.client, num_elements_offset); + let cube_count = calculate_cube_count_elemwise(&image.client, num_elements_offset, cube_dim); + + let dtype: StorageType = dtype_to_storage_type(image.dtype); + unsafe { + deform_col2img_coord_kernel::launch_unchecked( + &grad_offset.client, + cube_count, + cube_dim, + address_type!(image, offset, mask, grad_offset, grad_mask), + image.into_tensor_arg(), + offset.into_tensor_arg(), + mask.map(|mask| mask.into_tensor_arg()).into(), + columns.into_tensor_arg(), + grad_offset.clone().into_linear_view(), + grad_mask + .clone() + .map(|grad_mask| grad_mask.into_tensor_arg()) + .into(), + pos_shape, + DeformConv2dCol2ImgCoordArgsLaunch::new( + options.stride[0], + options.stride[1], + options.dilation[0], + options.dilation[1], + InputScalar::new(options.padding[0] as f32, dtype.elem_type()), + InputScalar::new(options.padding[1] as f32, dtype.elem_type()), + offset_groups, + kernel_h, + kernel_w, + ), + dtype, + ) + }; + + Ok((grad_offset, grad_mask)) +} + +#[derive(CubeLaunch, CubeType)] +struct DeformConv2dCol2ImgCoordArgs { + stride_h: usize, + stride_w: usize, + dilation_h: usize, + dilation_w: usize, + pad_h: InputScalar, + pad_w: InputScalar, + offset_groups: usize, + kernel_height: usize, + kernel_width: usize, +} + +#[allow(clippy::collapsible_if)] +#[cube(launch_unchecked, address_type = "dynamic")] +fn deform_col2img_coord_kernel( + image: &Tensor, + offset: &Tensor, + mask: ComptimeOption<&Tensor>, + columns: &Tensor, + mut grad_offset: LinearViewMut<'_, F>, + grad_mask: ComptimeOption<&mut Tensor>, + pos_shape: Sequence>, + args: &DeformConv2dCol2ImgCoordArgs, + #[define(F)] _dtype: StorageType, +) { + // Position format: [batch, [offset_groups, kernel_h, kernel_w, 2], out_h, out_w] + // Columns format: [[in_channel, kernel_h, kernel_w], [batch, out_h, out_w]] + // Alternatively : [batch, offset_channels, out_h, out_w] + + if ABSOLUTE_POS >= grad_offset.shape() { + terminate!(); + } + + let out_h = offset.shape(2); + let out_w = offset.shape(3); + let in_channels = image.shape(1); + let height = image.shape(2); + let width = image.shape(3); + let kernel_w = args.kernel_width; + let kernel_h = args.kernel_height; + + let mut grad_offset_val = F::new(0.0_f32); + let mut grad_mask_val = F::new(0.0_f32); + + let (_, pos) = decompose_linear(ABSOLUTE_POS, &pos_shape); + let [batch, offset_group, kernel_y, kernel_x, dir, out_y, out_x] = *pos else { + unreachable!() + }; + + let channels_per_offset_group = in_channels / args.offset_groups; + + let col_n = batch * out_h * out_w + out_y * out_w + out_x; + + let col_base_idx = + offset_group * channels_per_offset_group * kernel_h * kernel_w * columns.stride(0) + + col_n * columns.stride(1); + let mut image_base_idx = + batch * image.stride(0) + offset_group * channels_per_offset_group * image.stride(1); + + let offset_pos_1 = + offset_group * kernel_h * kernel_w * 2 + kernel_y * kernel_w * 2 + kernel_x * 2; + let offset_base_idx = batch * offset.stride(0) + + offset_pos_1 * offset.stride(1) + + out_y * offset.stride(2) + + out_x * offset.stride(3); + + let offset_y_idx = offset_base_idx; + let offset_x_idx = offset_base_idx + offset.stride(1); + + let offset_y = offset[offset_y_idx]; + let offset_x = offset[offset_x_idx]; + + let mask_pos_1 = offset_group * kernel_h * kernel_w + kernel_y * kernel_w + kernel_x; + #[comptime] + let mask_value = match &mask { + ComptimeOption::Some(mask) => { + let mask_idx = batch * mask.stride(0) + + mask_pos_1 * mask.stride(1) + + out_y * mask.stride(2) + + out_x * mask.stride(3); + mask[mask_idx] + } + ComptimeOption::None => F::new(1.0_f32), + }; + + let is_y_direction = dir == 0; + + for col_c in 0..channels_per_offset_group { + let col_pos = col_base_idx + col_c * kernel_h * kernel_w * columns.stride(0); + + let y = F::cast_from(out_y * args.stride_h + kernel_y * args.dilation_h) + - args.pad_h.get::() + + offset_y; + let x = F::cast_from(out_x * args.stride_w + kernel_x * args.dilation_w) + - args.pad_w.get::() + + offset_x; + + let weight = + get_coordinate_weight(image, image_base_idx, height, width, y, x, is_y_direction); + let columns_value = columns[col_pos]; + + grad_offset_val += mask_value * weight * columns_value; + + if grad_mask.is_some() && is_y_direction { + grad_mask_val += + columns_value * bilinear_interpolate(image, height, width, y, x, image_base_idx); + } + + image_base_idx += image.stride(1); + } + + grad_offset.write(ABSOLUTE_POS, grad_offset_val); + + #[comptime] + if let ComptimeOption::Some(grad_mask) = grad_mask { + if is_y_direction { + let idx = batch * grad_mask.stride(0) + + mask_pos_1 * grad_mask.stride(1) + + out_y * grad_mask.stride(2) + + out_x * grad_mask.stride(3); + + grad_mask[idx] = grad_mask_val + } + } +} + +#[cube] +fn get_coordinate_weight( + input: &Tensor, + offset: usize, + height: usize, + width: usize, + y: F, + x: F, + is_y_direction: bool, +) -> F { + let stride_y = input.stride(2); + let stride_x = input.stride(3); + + let y = f32::cast_from(y); + let x = f32::cast_from(x); + + let y_low = f32::floor(y); + let x_low = f32::floor(x); + let y_high = y_low + 1.; + let x_high = x_low + 1.; + + let valid_y_low = y_low >= 0. && y_low < height as f32; + let valid_y_high = y_high >= 0. && y_high < height as f32; + let valid_x_low = x_low >= 0. && x_low < width as f32; + let valid_x_high = x_high >= 0. && x_high < width as f32; + + let bottom_left = if valid_y_low && valid_x_low { + input[offset + y_low as usize * stride_y + x_low as usize * stride_x] + } else { + F::new(0.0_f32) + }; + let bottom_right = if valid_y_low && valid_x_high { + input[offset + y_low as usize * stride_y + x_high as usize * stride_x] + } else { + F::new(0.0_f32) + }; + let top_left = if valid_y_high && valid_x_low { + input[offset + y_high as usize * stride_y + x_low as usize * stride_x] + } else { + F::new(0.0_f32) + }; + let top_right = if valid_y_high && valid_x_high { + input[offset + y_high as usize * stride_y + x_high as usize * stride_x] + } else { + F::new(0.0_f32) + }; + + if is_y_direction { + let delta_x = F::cast_from(x - x_low); + delta_x * (top_right - bottom_right) + + (F::new(1.0_f32) - delta_x) * (top_left - bottom_left) + } else { + let delta_y = F::cast_from(y - y_low); + delta_y * (top_right - top_left) + + (F::new(1.0_f32) - delta_y) * (bottom_right - bottom_left) + } +} + +fn compute_input_grad( + columns: CubeTensor, + offset: CubeTensor, + mask: Option>, + options: &DeformConvOptions<2>, + kernel_dims: (usize, usize), + input_shape: Shape, +) -> Result, LaunchError> { + let client = offset.client.clone(); + let device = offset.device.clone(); + + let supports_fadd = client + .properties() + .atomic_type_usage(Type::atomic(FloatKind::F32)) + .contains(AtomicUsage::Add); + let supports_same_type = client + .properties() + .atomic_type_usage(Type::atomic(dtype_to_elem_type(columns.dtype))) + .contains(AtomicUsage::Add); + + let [batches, in_channels, height, width] = input_shape.dims(); + let [_, _, out_h, out_w] = offset.meta.shape().dims(); + let (kernel_h, kernel_w) = kernel_dims; + + let pos_shape = [in_channels, kernel_h, kernel_w, batches, out_h, out_w]; + let pos_shape = pos_shape.into_iter().collect(); + + let shape = Shape::new([batches, in_channels, height, width]); + let grad_in = match supports_fadd && supports_same_type { + // Use type as is to save a cast + true => zeros_client(client.clone(), device.clone(), shape, columns.dtype), + // Force `f32` to enable bitcasting as `u32`, or use intrinsic when supported + false => zeros_client(client.clone(), device.clone(), shape, DType::F32), + }; + let grad_arg = grad_in.clone().into_tensor_arg(); + + let num_elements = columns.meta.num_elements(); + let cube_dim = CubeDim::new(&offset.client, num_elements); + let cube_count = calculate_cube_count_elemwise(&offset.client, num_elements, cube_dim); + + let launch = match supports_fadd { + true => deform_col2img_kernel::launch_unchecked::, + false => deform_col2img_kernel::launch_unchecked::, + }; + let dtype = offset.dtype; + let dtypes: [StorageType; 2] = match supports_same_type { + true => [dtype_to_storage_type(dtype), dtype_to_storage_type(dtype)], + false => [ + dtype_to_storage_type(dtype), + dtype_to_storage_type(DType::F32), + ], + }; + + unsafe { + launch( + &grad_in.client, + cube_count, + cube_dim, + address_type!(offset, mask, columns, grad_in), + offset.into_tensor_arg(), + mask.map(|mask| mask.into_tensor_arg()).into(), + reshape(columns, Shape::new([num_elements])).into_linear_view(), + grad_arg, + pos_shape, + DeformConv2dCol2ImgArgsLaunch::new( + options.stride[0], + options.stride[1], + options.dilation[0], + options.dilation[1], + InputScalar::new(options.padding[0] as f32, dtypes[0].elem_type()), + InputScalar::new(options.padding[1] as f32, dtypes[0].elem_type()), + options.offset_groups, + kernel_h, + kernel_w, + ), + dtypes, + ) + }; + + Ok(if !supports_same_type || !supports_fadd { + cast(grad_in, dtype) + } else { + grad_in + }) +} + +#[derive(CubeLaunch, CubeType)] +struct DeformConv2dCol2ImgArgs { + stride_h: usize, + stride_w: usize, + dilation_h: usize, + dilation_w: usize, + pad_h: InputScalar, + pad_w: InputScalar, + offset_groups: usize, + kernel_height: usize, + kernel_width: usize, +} + +#[cube(launch_unchecked, address_type = "dynamic")] +fn deform_col2img_kernel( + offset: &Tensor, + mask: ComptimeOption<&Tensor>, + columns: LinearView<'_, F>, + grad_input: &mut Tensor>>, + pos_shape: Sequence>, + args: &DeformConv2dCol2ImgArgs, + #[define(F, FP)] _dtype: [StorageType; 2], +) { + // Position format: [[in_channels, kernel_h, kernel_w], [batch_size, out_h, out_w]] + if ABSOLUTE_POS >= columns.shape() { + terminate!(); + } + + let n_in_channels = grad_input.shape(1); + let height = grad_input.shape(2); + let width = grad_input.shape(3); + let kernel_h = args.kernel_height; + let kernel_w = args.kernel_width; + let n_offset_groups = args.offset_groups; + + let (_, pos) = decompose_linear(ABSOLUTE_POS, &pos_shape); + let [in_channel, kernel_y, kernel_x, batch, out_y, out_x] = *pos else { + unreachable!() + }; + + let channels_per_offset_group = n_in_channels / n_offset_groups; + let offset_group = in_channel / channels_per_offset_group; + + let offset_pos_1 = + offset_group * kernel_h * kernel_w * 2 + kernel_y * kernel_w * 2 + kernel_x * 2; + let offset_base_idx = batch * offset.stride(0) + + offset_pos_1 * offset.stride(1) + + out_y * offset.stride(2) + + out_x * offset.stride(3); + + let offset_y_idx = offset_base_idx; + let offset_x_idx = offset_base_idx + offset.stride(1); + + let offset_y = offset[offset_y_idx]; + let offset_x = offset[offset_x_idx]; + + #[comptime] + let mask_value = match mask { + ComptimeOption::Some(mask) => { + let mask_pos_1 = offset_group * kernel_h * kernel_w + kernel_y * kernel_w + kernel_x; + mask[batch * mask.stride(0) + + mask_pos_1 * mask.stride(1) + + out_y * mask.stride(2) + + out_x * mask.stride(3)] + } + ComptimeOption::None => F::new(1.0_f32), + }; + + let y = F::cast_from(out_y * args.stride_h + kernel_y * args.dilation_h) + - args.pad_h.get::() + + offset_y; + let x = F::cast_from(out_x * args.stride_w + kernel_x * args.dilation_w) + - args.pad_w.get::() + + offset_x; + + for dy in -1..=1i32 { + #[unroll] + for dx in -1..=1i32 { + let yp = y.floor() + F::cast_from(dy); + let xp = x.floor() + F::cast_from(dx); + + if yp >= F::new(0.0_f32) + && yp < F::cast_from(height) + && xp >= F::new(0.0_f32) + && xp < F::cast_from(width) + && F::abs(y - yp) < F::new(1.0_f32) + && F::abs(x - xp) < F::new(1.0_f32) + { + let gradient_pos = batch * grad_input.stride(0) + + in_channel * grad_input.stride(1) + + usize::cast_from(yp) * grad_input.stride(2) + + usize::cast_from(xp) * grad_input.stride(3); + + let weight = + (F::new(1.0_f32) - F::abs(y - yp)) * (F::new(1.0_f32) - F::abs(x - xp)); + + let value = mask_value * F::cast_from(weight) * columns.read(ABSOLUTE_POS); + + FAdd::Op::::float_atomic_add::(&mut grad_input[gradient_pos], value); + } + } + } +} + +type ProxyType = <::Op as FloatAtomicAdd>::ProxyType; + +#[cube] +trait FloatAtomicAddFamily: Send + Sync + 'static { + type Op: FloatAtomicAdd; +} + +#[cube] +trait FloatAtomicAdd: Send + Sync + 'static { + type ProxyType: Numeric; + + fn float_atomic_add(ptr: &mut Atomic, value: F); +} + +#[derive(CubeType)] +struct IntrinsicFloatAtomicAdd { + #[cube(comptime)] + _ty: PhantomData, +} + +#[derive(CubeType)] +struct CASFloatAtomicAdd; + +struct IntrinsicFloatAtomicAddFamily; + +impl FloatAtomicAddFamily for IntrinsicFloatAtomicAddFamily { + type Op = IntrinsicFloatAtomicAdd; +} + +impl FloatAtomicAddFamily for CASFloatAtomicAdd { + type Op = Self; +} + +#[cube] +impl FloatAtomicAdd for IntrinsicFloatAtomicAdd { + type ProxyType = FAdd; + + fn float_atomic_add(ptr: &mut Atomic, value: F) { + let value = FAdd::cast_from(value); + ptr.fetch_add(value); + } +} + +#[cube] +impl FloatAtomicAdd for CASFloatAtomicAdd { + type ProxyType = u32; + + fn float_atomic_add(ptr: &mut Atomic, value: F) { + let value = f32::cast_from(value); + if value != 0.0 { + let mut v = ptr.load(); + loop { + let prev = v; + let v_float = f32::from_bits(v); + let new = (v_float + value).to_bits(); + v = ptr.compare_exchange_weak(v, new); + if prev == v { + break; + } + } + } + } +} diff --git a/crates/burn-cubecl/src/kernel/conv/direct.rs b/crates/burn-cubecl/src/kernel/conv/direct.rs new file mode 100644 index 0000000..a74c98b --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/direct.rs @@ -0,0 +1,321 @@ +use crate::{ + CubeRuntime, + kernel::{into_contiguous_aligned, utils::address_type}, + ops::max_vector_size, + tensor::CubeTensor, +}; +use crate::{kernel::utils::decompose_linear, ops::numeric::empty_device_dtype}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{ + TensorMetadata, + ops::{ConvOptions, conv::calculate_conv_output_sizes}, +}; +use cubecl::{ + calculate_cube_count_elemwise, prelude::*, std::tensor::layout::linear::LinearViewMut, + tensor_vector_size_parallel, +}; +use cubecl::{num_traits::Zero, std::FastDivmod}; +use cubek::convolution::components::ConvSetupError; + +#[derive(CubeLaunch, CubeType, Clone)] +pub(crate) struct ConvParam { + pub stride: u32, + pub dilation: u32, + pub padding: i32, +} + +#[derive(CubeLaunch, CubeType)] +struct Conv2dArgs { + conv_params: Sequence, + channels_per_group: u32, +} + +#[cube(launch_unchecked, address_type = "dynamic")] +#[allow(clippy::redundant_closure)] +fn direct_conv2d_kernel( + input: &Tensor>, + weight: &Tensor>, + bias: ComptimeOption<&[Vector]>, + mut output: LinearViewMut<'_, Vector>, + args: Conv2dArgs, + shape_out: Sequence>, + shape_out_c: FastDivmod, + #[comptime] has_padding: bool, + #[define(E)] _dtype: StorageType, +) { + if !output.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + let n_spatial = comptime![shape_out.len()]; + + let vector_size_out = output.vector_size(); + let pos = ABSOLUTE_POS * vector_size_out; + + let in_c_per_group = weight.shape(weight.rank() - 1) as u32; + + let (rem, out_c) = shape_out_c.div_mod(pos as u32); + let (b, spatial_pos) = decompose_linear(rem, &shape_out); + + let g = out_c / args.channels_per_group; + let ic_start = in_c_per_group * g; + + let bias: ComptimeOption> = + bias.map(|bias| bias[out_c as usize / vector_size_out]); + let mut sum = bias.unwrap_or_else(|| Vector::zero()); + + let in_offs = b as usize * input.stride(0) + ic_start as usize; + + let stride_oc = weight.stride(0); + + let mut in_shape = Sequence::new(); + let mut in_strides = Sequence::new(); + let mut kernel_shape = Sequence::new(); + let mut kernel_strides = Sequence::new(); + + #[unroll] + for i in 0..n_spatial { + in_shape.push(input.shape(i + 1) as u32); + in_strides.push(input.stride(i + 1)); + kernel_shape.push(weight.shape(i + 1) as u32); + kernel_strides.push(weight.stride(i + 1)); + } + + let weight_offs = out_c as usize * stride_oc; + + let loop_params = LoopParams { + out_pos: spatial_pos, + in_shape, + in_strides, + kernel_shape, + kernel_strides, + conv_params: args.conv_params, + in_c_per_group, + stride_oc, + }; + + kernel_loop( + input, + weight, + &mut sum, + in_offs, + true, + weight_offs, + &loop_params, + 0usize, + has_padding, + ); + + output.write(ABSOLUTE_POS, sum); +} + +#[derive(CubeType, Clone)] +struct LoopParams { + out_pos: Sequence, + in_shape: Sequence, + in_strides: Sequence, + kernel_shape: Sequence, + kernel_strides: Sequence, + conv_params: Sequence, + + in_c_per_group: u32, + stride_oc: usize, +} + +#[cube] +fn kernel_loop( + input: &Tensor>, + weight: &Tensor>, + sum: &mut Vector, + in_offs: usize, + in_bounds: bool, + weight_offs: usize, + params: &LoopParams, + #[comptime] kernel_dim: usize, + #[comptime] has_padding: bool, +) { + if comptime![kernel_dim < params.kernel_shape.len()] { + let out_idx = *params.out_pos.index(kernel_dim); + let conv = params.conv_params.index(kernel_dim); + let shape = *params.in_shape.index(kernel_dim); + let stride = *params.in_strides.index(kernel_dim); + let k_stride = *params.kernel_strides.index(kernel_dim); + + for pos in 0..*params.kernel_shape.index(kernel_dim) { + let in_pos = (out_idx * conv.stride + pos * conv.dilation) as i32 - conv.padding; + let in_offs = in_offs + in_pos as usize * stride; + let weight_offs = weight_offs + pos as usize * k_stride; + let mut in_bounds = in_bounds; + + if has_padding { + in_bounds &= in_pos >= 0 && (in_pos as u32) < shape; + } + + kernel_loop( + input, + weight, + sum, + in_offs, + in_bounds, + weight_offs, + params, + comptime![kernel_dim + 1], + has_padding, + ); + } + } else { + kernel_loop_inner( + input, + weight, + sum, + in_offs, + in_bounds, + weight_offs, + params.in_c_per_group, + params.stride_oc, + ); + } +} + +#[cube] +fn kernel_loop_inner( + input: &Tensor>, + weight: &Tensor>, + sum: &mut Vector, + in_offs: usize, + in_bounds: bool, + weight_offs: usize, + in_c_per_group: u32, + stride_oc: usize, +) { + let vector_size_in = input.vector_size(); + let vector_size_out = sum.size(); + + if in_bounds { + for in_c in range_stepped(0, in_c_per_group, vector_size_in as u32) { + let in_pos = in_offs + in_c as usize; + let mut weight_pos = weight_offs + in_c as usize; + + let val = input[in_pos / vector_size_in]; + + #[unroll] + for v in 0..vector_size_out { + let weight = weight[weight_pos / vector_size_in]; + let val = val * weight; + + #[unroll] + for i in 0..vector_size_in { + sum.insert(v, sum.extract(v) + val.extract(i)); + } + weight_pos += stride_oc; + } + } + } +} + +/// Perform a 2D convolution using the direct convolution algorithm. +/// +/// * `input` - The input feature map +/// * `weight` - The weights (filter) applied to each kernel +/// * `bias` - The bias added to each channel +/// * `options` - The options to use for the convolution +/// +pub fn conv_direct( + mut input: CubeTensor, + mut weight: CubeTensor, + bias: Option>, + options: ConvOptions, +) -> Result, ConvSetupError> { + let out_dtype = input.dtype; + let rank = input.meta.shape().num_dims(); + let dim_c = rank - 1; + + // We only care about the channels here, everything else can be permuted + if input.meta.strides()[dim_c] != 1 { + input = into_contiguous_aligned(input); + } + if weight.meta.strides()[dim_c] != 1 { + weight = into_contiguous_aligned(weight); + } + + let batch_size = input.meta.shape()[0]; + let in_shape = &input.meta.shape()[1..dim_c]; + let out_channels = weight.meta.shape()[0]; + let kernel_shape = &weight.meta.shape()[1..dim_c]; + + let channels_per_group = out_channels / options.groups; + + let out_size = calculate_conv_output_sizes( + kernel_shape, + &options.stride, + &options.padding, + &options.dilation, + in_shape, + ); + + let mut shape_out = vec![batch_size]; + shape_out.extend(out_size.iter().copied()); + shape_out.push(out_channels); + + let output = empty_device_dtype( + input.client.clone(), + input.device.clone(), + shape_out.into(), + out_dtype, + ); + + // Need custom vector size calculation here to account for the groups division. Need to vectorize + // over `channels_per_group` instead. + let mut grouped_out_shape = output.shape(); + grouped_out_shape[dim_c] = channels_per_group; + let vector_size_out = tensor_vector_size_parallel( + input.client.io_optimized_vector_sizes(input.dtype.size()), + &grouped_out_shape, + output.meta.strides(), + dim_c, + ); + // Use channels_per_group instead of in_channels to avoid issues here + let vector_size_in = max_vector_size(&weight); + + let shape_out = output.meta.shape()[1..dim_c] + .iter() + .map(|s| *s as u32) + .collect(); + let shape_out_c = out_channels as u32; + + let mut conv_params = SequenceArg::new(); + + for i in 0..kernel_shape.len() { + conv_params.push(ConvParamLaunch::new( + options.stride[i] as u32, + options.dilation[i] as u32, + options.padding[i] as i32, + )); + } + + let working_units = output.meta.num_elements() / vector_size_out; + let cube_dim = CubeDim::new(&input.client, working_units); + let cube_count = calculate_cube_count_elemwise(&input.client, working_units, cube_dim); + + unsafe { + direct_conv2d_kernel::launch_unchecked( + &output.client, + cube_count, + cube_dim, + address_type!(input, weight, bias, output), + vector_size_in, + vector_size_out, + input.into_tensor_arg(), + weight.into_tensor_arg(), + bias.map(|b| b.into_buffer_arg()).into(), + output.clone().into_linear_view(), + Conv2dArgsLaunch::new(conv_params, channels_per_group as u32), + shape_out, + shape_out_c, + options.padding.iter().any(|it| *it != 0), + dtype_to_storage_type(out_dtype), + ) + }; + + Ok(output) +} diff --git a/crates/burn-cubecl/src/kernel/conv/forward/implicit_gemm/launch.rs b/crates/burn-cubecl/src/kernel/conv/forward/implicit_gemm/launch.rs new file mode 100644 index 0000000..ac2c710 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/forward/implicit_gemm/launch.rs @@ -0,0 +1,171 @@ +use crate::{CubeRuntime, ops::numeric::empty_device_dtype, tensor::CubeTensor}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::ops::{ConvOptions, conv::calculate_conv_output_sizes}; +use cubek::{ + convolution::{ + AcceleratedTileKind, ConvAlgorithm, ConvolutionArgs, ConvolutionInputs, Strategy, + components::ConvSetupError, launch_ref, + }, + matmul::definition::{MatmulElems, MatmulGlobalElems}, + std::InputBinding, +}; + +/// Perform a 2D convolution using the implicit GEMM (im2col) algorithm, using cubecl tiling matmul +/// components. Uses [`CmmaLargeMAlgorithm`] for the stage size +/// +/// * `input` - The input feature map +/// * `weight` - The weights (filter) applied to each kernel +/// * `bias` - The bias added to each channel +/// * `options` - The options to use for the convolution +pub fn conv_gemm_simple_sync( + input: CubeTensor, + weight: CubeTensor, + bias: Option>, + options: ConvOptions, + tile_kind: AcceleratedTileKind, +) -> Result, ConvSetupError> { + let algorithm = match tile_kind { + AcceleratedTileKind::Cmma => ConvAlgorithm::SimpleSyncCyclic, + AcceleratedTileKind::Mma => ConvAlgorithm::SimpleSyncStrided, + }; + launch_convolution_forward::( + &Strategy::Inferred { + algorithm, + tile_kind, + }, + input, + weight, + bias, + options, + ) +} + +pub fn conv_gemm_simple_async( + input: CubeTensor, + weight: CubeTensor, + bias: Option>, + options: ConvOptions, + tile_kind: AcceleratedTileKind, +) -> Result, ConvSetupError> { + let algorithm = match tile_kind { + AcceleratedTileKind::Cmma => ConvAlgorithm::SimpleAsyncCyclic, + AcceleratedTileKind::Mma => ConvAlgorithm::SimpleAsyncStrided, + }; + launch_convolution_forward::( + &Strategy::Inferred { + algorithm, + tile_kind, + }, + input, + weight, + bias, + options, + ) +} + +/// Perform a 2D convolution using the implicit GEMM (im2col) algorithm, using cubecl tiling matmul +/// components. Uses [`CmmaLargeMAlgorithm`] for the stage size +/// +/// * `input` - The input feature map +/// * `weight` - The weights (filter) applied to each kernel +/// * `bias` - The bias added to each channel +/// * `options` - The options to use for the convolution +pub fn conv_gemm_simple_tma( + input: CubeTensor, + weight: CubeTensor, + bias: Option>, + options: ConvOptions, + tile_kind: AcceleratedTileKind, +) -> Result, ConvSetupError> { + launch_convolution_forward::( + &Strategy::Inferred { + algorithm: ConvAlgorithm::SimpleAsyncTma, + tile_kind, + }, + input, + weight, + bias, + options, + ) +} + +/// Perform a 2D convolution using the implicit GEMM (im2col) algorithm, using cubecl tiling matmul +/// components, using the specified algorithm. +/// +/// * `input` - The input feature map +/// * `weight` - The weights (filter) applied to each kernel +/// * `bias` - The bias added to each channel +/// * `options` - The options to use for the convolution +pub fn launch_convolution_forward( + strategy: &Strategy, + input: CubeTensor, + weight: CubeTensor, + bias: Option>, + options: ConvOptions, +) -> Result, ConvSetupError> { + if options.groups != 1 { + return Err(ConvSetupError::Groups(options.groups)); + } + + let out_dtype = input.dtype; + let rank = input.meta.shape().num_dims(); + let batch_size = input.meta.shape()[0]; + let dim_c = rank - 1; + let shape = &input.meta.shape()[1..dim_c]; + + let out_channels = weight.meta.shape()[0]; + let weight_shape = &weight.meta.shape()[1..dim_c]; + + let mut out_shape = calculate_conv_output_sizes( + weight_shape, + &options.stride, + &options.padding, + &options.dilation, + shape, + ); + + out_shape.insert(0, batch_size); + out_shape.push(out_channels); + + let out = empty_device_dtype( + input.client.clone(), + input.device.clone(), + out_shape.into(), + out_dtype, + ); + + let bias = bias.map(|bias| { + let dtype = bias.dtype; + InputBinding::Normal(bias.binding(), dtype_to_storage_type(dtype)) + }); + + let client = input.client.clone(); + let dtypes = MatmulElems::from_globals(&MatmulGlobalElems { + lhs: dtype_to_storage_type(input.dtype), + rhs: dtype_to_storage_type(weight.dtype), + out: dtype_to_storage_type(out_dtype), + }); + let input_dtype = input.dtype; + let weight_dtype = weight.dtype; + let input = InputBinding::new(input.binding(), dtype_to_storage_type(input_dtype)); + let weight = InputBinding::new(weight.binding(), dtype_to_storage_type(weight_dtype)); + + launch_ref::( + strategy, + &client, + ConvolutionInputs::Forward { + input, + weight, + bias, + out: out.clone().binding(), + }, + ConvolutionArgs { + stride: options.stride, + padding: options.padding, + dilation: options.dilation, + }, + dtypes, + )?; + + Ok(out) +} diff --git a/crates/burn-cubecl/src/kernel/conv/forward/implicit_gemm/mod.rs b/crates/burn-cubecl/src/kernel/conv/forward/implicit_gemm/mod.rs new file mode 100644 index 0000000..0df8c51 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/forward/implicit_gemm/mod.rs @@ -0,0 +1,2 @@ +pub mod launch; +pub use launch::*; diff --git a/crates/burn-cubecl/src/kernel/conv/forward/mod.rs b/crates/burn-cubecl/src/kernel/conv/forward/mod.rs new file mode 100644 index 0000000..d5091ae --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/forward/mod.rs @@ -0,0 +1,7 @@ +pub mod implicit_gemm; + +#[cfg(feature = "autotune")] +pub mod tune; + +#[cfg(feature = "autotune")] +pub(crate) use tune::*; diff --git a/crates/burn-cubecl/src/kernel/conv/forward/tune.rs b/crates/burn-cubecl/src/kernel/conv/forward/tune.rs new file mode 100644 index 0000000..d9d8f76 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/forward/tune.rs @@ -0,0 +1,187 @@ +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::ops::ConvOptions; +use cubecl::{ + ir::StorageType, + tune::{LocalTuner, Tunable, TunableSet, anchor, local_tuner}, +}; +use cubek::convolution::AcceleratedTileKind; + +use crate::{ + CubeAutotuneKey, CubeRuntime, CubeTuneId, + kernel::conv::{ConvAutotuneKey, conv_direct, conv_im2col_1x1, forward::implicit_gemm::*}, + tensor::CubeTensor, +}; + +/// Executes autotune on convolution operations +pub fn conv_autotune( + input: CubeTensor, + weight: CubeTensor, + bias: Option>, + options: ConvOptions, +) -> CubeTensor { + let client = input.client.clone(); + + static TUNER: LocalTuner = local_tuner!(); + + let tunables = TUNER.init(|| { + TunableSet::new(create_key::, create_conv_input::) + .with(Tunable::new( + "conv_direct", + |(input, weight, bias, options)| conv_direct::(input, weight, bias, options), + )) + .with(Tunable::new( + "conv_im2col_1x1", + |(input, weight, bias, options)| { + conv_im2col_1x1::(input, weight, bias, options) + }, + )) + .with(Tunable::new( + "simple_sync_cmma", + |(input, weight, bias, options)| { + conv_gemm_simple_sync(input, weight, bias, options, AcceleratedTileKind::Cmma) + }, + )) + .with(Tunable::new( + "simple_sync_mma", + |(input, weight, bias, options)| { + conv_gemm_simple_sync(input, weight, bias, options, AcceleratedTileKind::Mma) + }, + )) + .with(Tunable::new( + "simple_async_cmma", + |(input, weight, bias, options)| { + conv_gemm_simple_async(input, weight, bias, options, AcceleratedTileKind::Cmma) + }, + )) + .with(Tunable::new( + "simple_async_mma", + |(input, weight, bias, options)| { + conv_gemm_simple_async(input, weight, bias, options, AcceleratedTileKind::Mma) + }, + )) + .with(Tunable::new( + "simple_tma_cmma", + |(input, weight, bias, options)| { + conv_gemm_simple_tma(input, weight, bias, options, AcceleratedTileKind::Cmma) + }, + )) + .with(Tunable::new( + "simple_tma_mma", + |(input, weight, bias, options)| { + conv_gemm_simple_tma(input, weight, bias, options, AcceleratedTileKind::Mma) + }, + )) + }); + + TUNER.execute( + &CubeTuneId::new(&input.client, &input.device), + &client, + tunables, + (input, weight, bias, options), + ) +} + +pub fn create_conv_input( + _key: &CubeAutotuneKey, + (input, weights, bias, options): &( + CubeTensor, + CubeTensor, + Option>, + ConvOptions, + ), +) -> ( + CubeTensor, + CubeTensor, + Option>, + ConvOptions, +) { + ( + input.clone(), + weights.clone(), + bias.clone(), + options.clone(), + ) +} + +fn create_key( + (input, weights, bias, options): &( + CubeTensor, + CubeTensor, + Option>, + ConvOptions, + ), +) -> CubeAutotuneKey { + let dtype = input.dtype; + let rank = input.meta.shape().num_dims(); + let dim_c = rank - 1; + + let batch_size = input.meta.shape()[0]; + let in_channels = input.meta.shape()[dim_c]; + let out_channels = weights.meta.shape()[0]; + + let kernel_size = weights.meta.shape()[1..dim_c].to_vec(); + let in_shape = input.meta.shape()[1..dim_c] + .iter() + .map(|shape| anchor(*shape, None, None, None)) + .collect(); + + let ConvOptions { + stride, + padding, + dilation, + groups, + } = options.clone(); + + let lhs_stride_align = if input.meta.strides()[dim_c] == 1 { + stride_align(input.meta.strides(), dtype_to_storage_type(input.dtype)) + } else { + 0 + }; + let lhs_shape_align = pow2_factor(in_channels).min(lhs_stride_align); + let rhs_stride_align = if weights.meta.strides()[dim_c] == 1 { + stride_align(weights.meta.strides(), dtype_to_storage_type(weights.dtype)) + } else { + 0 + }; + let rhs_shape_align = pow2_factor(in_channels).min(rhs_stride_align); + + CubeAutotuneKey::Conv(ConvAutotuneKey::new( + kernel_size, + stride.to_vec(), + padding.to_vec(), + dilation.to_vec(), + groups, + in_channels, + out_channels, + in_shape, + batch_size, + bias.is_some(), + dtype, + lhs_shape_align, + lhs_stride_align, + rhs_shape_align, + rhs_stride_align, + )) +} + +/// Maximum factor relevant for strides. Currently set to 2^10 because that's 128-byte swizzle's +/// repeat number, so it's the largest align that can have performance impacts. +const MAX_STRIDE_FACTOR: u32 = 10; + +/// Defines the non-contiguous stride alignment in terms of powers of two +fn stride_align(strides: &[usize], elem: StorageType) -> u8 { + let max = MAX_STRIDE_FACTOR; + let dim_c = strides.len() - 1; + let factor = strides[..dim_c] + .iter() + .map(|it| (*it * elem.size_bits()) / 8) + .map(|it| it.trailing_zeros()) + .min() + .unwrap_or(max); + factor.min(max) as u8 +} + +/// Defines the potential vectorization. +fn pow2_factor(axis: usize) -> u8 { + axis.trailing_zeros().min(4) as u8 +} diff --git a/crates/burn-cubecl/src/kernel/conv/im2col.rs b/crates/burn-cubecl/src/kernel/conv/im2col.rs new file mode 100644 index 0000000..6a73e30 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/im2col.rs @@ -0,0 +1,191 @@ +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{ + DType, + ops::{ConvOptions, conv::calculate_conv_output_sizes}, +}; +use burn_std::{Metadata, Shape}; +use core::iter; +use cubecl::{ + prelude::*, + std::tensor::{TensorHandle, into_contiguous_pitched}, +}; +use cubek::convolution::components::ConvSetupError; + +use crate::{ + CubeRuntime, + kernel::{ + AddOp, into_contiguous_aligned, launch_binop, + matmul::{MatmulStrategy, matmul}, + utils::split_dim, + }, + ops::{reshape, swap_dims}, + tensor::CubeTensor, +}; + +#[cfg(not(test))] +pub(crate) fn batches_per_run( + batch_size: usize, + out_shape: usize, + plane_size: usize, +) -> Result { + use cubek::matmul::definition::MatmulAvailabilityError; + + let cube_count_per_batch = out_shape.div_ceil(plane_size); + let max_cube_count = u16::MAX as usize; + let max_simultaneous = Ord::min(max_cube_count / cube_count_per_batch, batch_size); + if max_simultaneous == 0 { + return Err(MatmulAvailabilityError::CubeCountTooBig(CubeCount::Static( + cube_count_per_batch as u32, + 1, + 1, + )) + .into()); + } + Ok((0..=max_simultaneous) + .rev() + .find(|per_run| batch_size.is_multiple_of(*per_run)) + .expect("Logically not possible")) +} + +#[cfg(test)] +#[allow(unused)] +pub(crate) fn batches_per_run( + batch_size: usize, + out_shape: usize, + plane_size: usize, +) -> Result { + Ok(1) +} + +pub fn conv_im2col_1x1( + input: CubeTensor, + mut weight: CubeTensor, + bias: Option>, + options: ConvOptions, +) -> Result, ConvSetupError> { + if options.groups != 1 { + return Err(ConvSetupError::Groups(options.groups)); + } + + let rank = input.meta.num_dims(); + let dim_c = rank - 1; + + let batch_size = input.meta.shape()[0]; + let in_channels = input.meta.shape()[dim_c]; + let in_shape = &input.meta.shape()[1..dim_c]; + let out_channels = weight.meta.shape()[0]; + let kernel_shape = &weight.meta.shape()[1..dim_c]; + + if kernel_shape.iter().any(|s| *s != 1) { + return Err(ConvSetupError::Unknown); + } + + let out_shape = calculate_conv_output_sizes( + kernel_shape, + &options.stride, + &options.padding, + &options.dilation, + in_shape, + ); + + let mut split_m = vec![batch_size]; + split_m.extend(out_shape.iter().copied()); + + if kernel_shape.iter().any(|it| *it != 1) || in_shape != out_shape { + return Err(ConvSetupError::Unknown); + } + + let input = reshape_input(input); // [(NHW), C] : [M, K] + let dtype = input.dtype; + + // Efficient permutation that takes the stride required for TMA into account + let weight = if weight.meta.strides()[dim_c] != 1 { + // Remove kernel dims so padded dim is channels + *weight.meta = Metadata::new( + [out_channels, in_channels], // [N, K] + [weight.meta.strides()[0], weight.meta.strides()[dim_c]], + ); + // Pitched contiguous to skip running another kernel for TMA + into_contiguous_aligned(weight) + } else { + // Already compatible, skip initial reshape + *weight.meta = Metadata::new([out_channels, in_channels], [weight.meta.strides()[0], 1]); + weight + }; + + // Permute to N-major, while keeping memory layout K-major. K-major for both sides is the most + // efficient for matmul, and allows skipping a contiguous kernel + let weight = swap_dims(weight, 0, 1); // [K, N] + + let out = matmul(input, weight, None, MatmulStrategy::default(), dtype)?; // [M, N] + + // Skip reshape to avoid potential `into_contiguous`. We're only splitting dims so it's safe. + let mut out = split_dim(out, 0, &split_m); // [N, H, W, C] + + if let Some(bias) = bias { + let mut bias_shape = iter::repeat_n(1, rank - 1).collect::>(); + bias_shape.push(out_channels); + let bias = reshape(bias, bias_shape.into()); + out = launch_binop::(out, bias); + } + + Ok(out) +} + +/// Reshapes NHWC input to [(N, H, W), C] +fn reshape_input(input: CubeTensor) -> CubeTensor { + let rank = input.meta.num_dims(); + let dim_c = rank - 1; + let dtype = input.dtype; + + let batch_size = input.meta.shape()[0]; + let in_c: usize = input.meta.shape()[dim_c]; + let in_shape: Shape = input.meta.shape()[1..dim_c].into(); + + let mut input = if !is_spatial_contiguous(input.meta.shape(), input.meta.strides()) { + let (client, device) = (input.client.clone(), input.device.clone()); + let contiguous = + into_contiguous_pitched(&client, input.binding(), dtype_to_storage_type(dtype)); + from_handle(client, device, contiguous, dtype) + } else { + input + }; + + *input.meta = Metadata::new( + [batch_size * in_shape.num_elements(), in_c], // [M, K] + [input.meta.strides()[dim_c - 1], input.meta.strides()[dim_c]], + ); + input +} + +fn is_spatial_contiguous(shape: &[usize], strides: &[usize]) -> bool { + let rank = shape.len(); + let dim_c = rank - 1; + + // Channel must be contiguous for the [(N, H, W), C] reshape to be valid + if strides[dim_c] != 1 { + return false; + } + + for i in (1..dim_c).rev() { + if strides[i + 1] * shape[i + 1] != strides[i] { + return false; + } + } + true +} + +fn from_handle( + client: ComputeClient, + device: R::Device, + handle: TensorHandle, + dtype: DType, +) -> CubeTensor { + CubeTensor::new( + client.clone(), + handle.handle, + *handle.metadata, + device.clone(), + dtype, + ) +} diff --git a/crates/burn-cubecl/src/kernel/conv/mod.rs b/crates/burn-cubecl/src/kernel/conv/mod.rs new file mode 100644 index 0000000..03dde01 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/mod.rs @@ -0,0 +1,25 @@ +mod backward_data; +mod backward_weight; +mod base; +mod conv_transpose2d; +mod conv_transpose3d; +mod deform_conv2d; +mod deform_conv_transpose2d; +mod direct; +mod forward; +mod im2col; + +mod tune_key; + +pub(crate) use backward_data::*; +pub(crate) use conv_transpose2d::*; +pub(crate) use conv_transpose3d::*; +pub(crate) use deform_conv_transpose2d::*; +pub(crate) use deform_conv2d::*; +pub(crate) use direct::*; +pub(crate) use im2col::*; + +pub use base::*; +pub use conv_transpose2d::{ConvTranspose2dStrategy, conv_transpose2d}; + +pub(crate) use tune_key::*; diff --git a/crates/burn-cubecl/src/kernel/conv/tune_key.rs b/crates/burn-cubecl/src/kernel/conv/tune_key.rs new file mode 100644 index 0000000..1df268d --- /dev/null +++ b/crates/burn-cubecl/src/kernel/conv/tune_key.rs @@ -0,0 +1,50 @@ +use burn_backend::DType; +use cubecl::AutotuneKey; +use serde::{Deserialize, Serialize}; + +#[derive(Hash, Eq, PartialEq, Debug, Clone, Serialize, Deserialize, AutotuneKey)] +/// Autotune key representative of matmul versions +pub struct ConvAutotuneKey { + pub kernel_size: Vec, + pub stride: Vec, + pub padding: Vec, + pub dilation: Vec, + pub groups: usize, + #[autotune(anchor)] + pub in_channels: usize, + #[autotune(anchor)] + pub out_channels: usize, + pub shape: Vec, + #[autotune(anchor)] + pub batch_size: usize, + pub has_bias: bool, + pub dtype: DType, + + pub lhs_shape_align: u8, + pub lhs_stride_align: u8, + pub rhs_shape_align: u8, + pub rhs_stride_align: u8, +} + +#[derive(Hash, Eq, PartialEq, Debug, Clone, Serialize, Deserialize, AutotuneKey)] +/// Autotune key representative of matmul versions +pub struct ConvTranspose2dAutotuneKey { + pub kernel_size: [usize; 2], + pub stride: [usize; 2], + pub padding: [usize; 2], + pub padding_out: [usize; 2], + pub dilation: [usize; 2], + pub groups: usize, + #[autotune(anchor)] + pub in_channels: usize, + #[autotune(anchor)] + pub out_channels: usize, + #[autotune(anchor)] + pub height: usize, + #[autotune(anchor)] + pub width: usize, + #[autotune(anchor)] + pub batch_size: usize, + pub has_bias: bool, + pub dtype: DType, +} diff --git a/crates/burn-cubecl/src/kernel/cross.rs b/crates/burn-cubecl/src/kernel/cross.rs new file mode 100644 index 0000000..2dfc2fd --- /dev/null +++ b/crates/burn-cubecl/src/kernel/cross.rs @@ -0,0 +1,108 @@ +use crate::{ + CubeRuntime, + kernel::{ + into_contiguous, + utils::{address_type, broadcast_shape}, + }, + ops::{numeric::empty_device_dtype, swap_dims}, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::std::tensor::layout::linear::{LinearView, LinearViewMut}; +use cubecl::{calculate_cube_count_elemwise, prelude::*}; + +#[cube(launch_unchecked, address_type = "dynamic")] +fn cross_kernel( + lhs: LinearView<'_, E>, + rhs: LinearView<'_, E>, + mut output: LinearViewMut<'_, E>, + #[define(E)] _dtype: StorageType, +) { + // Each thread processes one 3-element vector + let vector_idx = ABSOLUTE_POS; + let base_pos = vector_idx * 3; + + if !output.is_in_bounds(base_pos) { + terminate!(); + } + + // Extract vectors + let a0 = lhs.read(base_pos); + let a1 = lhs.read(base_pos + 1); + let a2 = lhs.read(base_pos + 2); + let b0 = rhs.read(base_pos); + let b1 = rhs.read(base_pos + 1); + let b2 = rhs.read(base_pos + 2); + + // Compute cross product: a × b + let x = a1 * b2 - a2 * b1; + let y = a2 * b0 - a0 * b2; + let z = a0 * b1 - a1 * b0; + + // Store result + output.write(base_pos, x); + output.write(base_pos + 1, y); + output.write(base_pos + 2, z); +} + +pub(crate) fn cross( + lhs: CubeTensor, + rhs: CubeTensor, + dim: usize, +) -> CubeTensor { + let ndims = lhs.meta.num_dims(); + + // Validate that the cross dimension has size 3 + if lhs.meta.shape()[dim] != 3 || rhs.meta.shape()[dim] != 3 { + panic!( + "Cross product requires dimension {} to have size 3, but got {} and {}", + dim, + lhs.meta.shape()[dim], + rhs.meta.shape()[dim] + ); + } + + // The kernel reads each 3-vector from contiguous memory, so it expects the + // cross dimension to be the last (innermost) and physically contiguous. + // For non-last dims we permute the cross dim to the last position, run the + // kernel, then permute the result back. swap_dims only updates strides, so + // make the permuted operands contiguous before launch. + if dim != ndims - 1 { + let last = ndims - 1; + let lhs = into_contiguous(swap_dims(lhs, dim, last)); + let rhs = into_contiguous(swap_dims(rhs, dim, last)); + let result = cross(lhs, rhs, last); + return swap_dims(result, dim, last); + } + + let output_shape = broadcast_shape(&[&lhs, &rhs]); + + let output = empty_device_dtype( + lhs.client.clone(), + lhs.device.clone(), + output_shape.clone(), + lhs.dtype, + ); + + // Number of vectors to process + let num_vectors = output_shape.num_elements() / 3; + + let cube_dim = CubeDim::new(&lhs.client, num_vectors); + let cube_count = calculate_cube_count_elemwise(&lhs.client, num_vectors, cube_dim); + let dtype = lhs.dtype; + + unsafe { + cross_kernel::launch_unchecked( + &output.client, + cube_count, + cube_dim, + address_type!(lhs, rhs, output), + lhs.into_linear_view_like(&output), + rhs.into_linear_view_like(&output), + output.clone().into_linear_view(), + dtype_to_storage_type(dtype), + ); + }; + + output +} diff --git a/crates/burn-cubecl/src/kernel/ctc.rs b/crates/burn-cubecl/src/kernel/ctc.rs new file mode 100644 index 0000000..dc39fd9 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/ctc.rs @@ -0,0 +1,670 @@ +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::prelude::*; + +use crate::{ + CubeRuntime, kernel::into_contiguous, ops::numeric::empty_device_dtype, tensor::CubeTensor, +}; +use burn_backend::{Shape, TensorMetadata}; + +/// Maximum `2 * max_target_len + 1` the kernel supports. The alpha/beta state is +/// held in shared memory as two f32 buffers of this size (active row + scratch), +/// so peak shared use at full capacity is `2 * 8192 * 4 = 64 KB`. Apple Metal +/// caps shared memory at 32 KB per block, so the launch site sizes the buffer to +/// the actual per-batch `max_l_prime`; this constant is only the kernel-side +/// upper bound. Inputs exceeding it panic rather than silently degrade. +const SHARED_ALPHA_CAPACITY: u32 = 8192; + +/// Class label at position `s` of the blank-inserted label sequence `l'`. +/// Odd `s` reads the underlying target at index `(s-1)/2`; even `s` is a blank. +#[cube] +fn l_prime_class( + s: usize, + targets: &Tensor, + n: usize, + tgt_n: usize, + tgt_s: usize, + blank: usize, +) -> usize { + if s % 2 == 1 { + u32::cast_from(targets[n * tgt_n + ((s - 1) / 2) * tgt_s]) as usize + } else { + blank + } +} + +/// Numerically stable `log(exp(a) + exp(b))` with a sentinel short-circuit. +/// When `max(a, b) < unreachable_threshold`, returns `max(a, b)` directly so +/// the sentinel value doesn't drift upward each recursion step when both +/// inputs sit at the `-6e4` floor. +/// +/// The threshold's magnitude is forced by f16: the sentinel can't go below +/// `-65504` (f16 max magnitude), so it's `-6e4`, and the threshold has to sit +/// above the sentinel but below any plausible legit alpha value, leaving a +/// narrow band around `-1e4`. On sufficiently long sequences where legit +/// alpha values naturally drop below `-1e4` (roughly `T * log(1/C) < -1e4`), +/// reachable states get misclassified as unreachable. Mitigation is a +/// WGSL-only path with a smaller sentinel; WGSL spec 8.7 lets implementations +/// replace runtime `1/0` with zero, so `-inf` can't be synthesized reliably. +#[cube] +fn log_sum_exp2(a: F, b: F, unreachable_threshold: F, one: F) -> F { + let mut mx = a; + let mut mn = b; + if b > a { + mx = b; + mn = a; + } + if mx < unreachable_threshold { + mx + } else { + mx + (one + (mn - mx).exp()).ln() + } +} + +/// Single alpha (or beta) recurrence step. `near`, `near_m1`, `near_m2` are +/// the three values from the previous time row (alpha: `t-1`; beta: `t+1`). +/// `log_p` is the emission log-prob at the current `(t, l'[s])` and +/// `skip_allowed` toggles the 2-position skip transition. +#[cube] +fn recurrence_step( + near: F, + near_m1: F, + near_m2: F, + log_p: F, + skip_allowed: bool, + unreachable_threshold: F, + one: F, +) -> F { + let lse_01 = log_sum_exp2::(near, near_m1, unreachable_threshold, one); + let combined = if skip_allowed { + log_sum_exp2::(lse_01, near_m2, unreachable_threshold, one) + } else { + lse_01 + }; + log_p + combined +} + +/// Final `-log(alpha_last_blank + alpha_last_label)` reduction. Synthesizes a +/// true `+inf` via `exp()` overflow when both final alphas are at the sentinel +/// (the target is unreachable), so downstream `zero_infinity` logic can detect +/// it via `is_inf`. Builds the overflow arithmetically from a runtime-dependent +/// value (`target_len`, guaranteed >= 1 here) to keep WGSL's comptime-overflow +/// validator quiet. +#[cube] +fn finalize_nll( + last_blank: F, + last_label: F, + target_len: usize, + unreachable_threshold: F, + one: F, +) -> F { + let mut mx = last_blank; + let mut mn = last_label; + if last_label > last_blank { + mx = last_label; + mn = last_blank; + } + if mx < unreachable_threshold { + (F::new(1000.0_f32) * F::cast_from(target_len as u32)).exp() + } else { + F::new(0.0_f32) - (mx + (one + (mn - mx).exp()).ln()) + } +} + +/// Value to emit when `input_len == 0`. `target_len == 0` is the only case +/// with a valid alignment (P(empty | empty) = 1, nll = 0); otherwise the +/// target is unreachable and the output is `+inf` synthesized via overflow. +#[cube] +fn empty_input_nll(target_len: usize) -> F { + if target_len == 0 { + F::new(0.0_f32) + } else { + (F::new(1000.0_f32) * F::cast_from(target_len as u32)).exp() + } +} + +/// CTC alpha-recursion kernel. +/// +/// Each cube handles one batch element. `cube_dim.x` is fixed at launch time +/// (capped to the runtime's hardware limit); each thread strides over the `s` +/// positions of the modified label sequence `l'` (length `2 * target_len + 1`), +/// covering arbitrary target lengths up to `SHARED_ALPHA_CAPACITY`. `alpha` is +/// kept in shared memory and the time loop runs sequentially inside the kernel, +/// using two `sync_cube()` barriers per iteration: one to fence reads of +/// `alpha[t-1]` before any thread writes `alpha[t]`, one to publish the new row +/// before the next iteration. This collapses what would otherwise be roughly +/// `40 * T` host-side dispatches into a single kernel launch. +/// +/// Impossible alignments use a large finite negative sentinel (`-6.0e4`) +/// rather than true `-inf`, because WGSL rejects `f32(-inf)` as an identifier +/// and f16's range caps at ~65504. The recurrence treats values below a +/// threshold (`-1.0e4`) as unreachable. If an entire sequence has no valid +/// alignment (e.g. `target_length > input_length`), the kernel synthesizes +/// `+inf` in the output so downstream `zero_infinity` masking in `burn-nn` +/// can detect it via `is_inf`. +#[cube(launch)] +fn ctc_loss_kernel( + log_probs: &Tensor, // [T, N, C] + targets: &Tensor, // [N, S_max] + input_lengths: &Tensor, // [N] + target_lengths: &Tensor, // [N] + output: &mut Tensor, // [N] + blank: u32, + #[comptime] alpha_capacity: u32, + #[define(F, I)] _dtypes: [StorageType; 2], +) { + let n = CUBE_POS_X as usize; + let cube_dim = CUBE_DIM_X as usize; + let alpha_cap = alpha_capacity as usize; + let blank_u = blank as usize; + + let target_len = u32::cast_from(target_lengths[n]) as usize; + let input_len = u32::cast_from(input_lengths[n]) as usize; + let l_prime_len = 2 * target_len + 1; + + // Empty-input edge case: handled identically in both kernels to keep the + // forward loss and the backward nll agreeing for this sample. + if input_len == 0 { + if UNIT_POS_X == 0 { + output[n] = empty_input_nll::(target_len); + } + terminate!(); + } + + let lp_t = log_probs.stride(0); + let lp_n = log_probs.stride(1); + let lp_c = log_probs.stride(2); + let tgt_n = targets.stride(0); + let tgt_s = targets.stride(1); + + // Two adjacent regions: alpha[0..alpha_cap] is the active row, the second + // half [alpha_cap..2*alpha_cap] is a write scratch buffer that we copy back + // to the active region after a sync. This avoids RAW hazards across stride + // batches in the t-loop (a thread writing alpha[s] races with another + // thread still reading alpha[s-1] or alpha[s-2] for its own s). + let mut alpha = Shared::new_slice(2 * alpha_cap); + // Sentinel for unreachable states. f16 caps at ~65504 magnitude, so we + // can't go lower than `-6e4` without blowing past that range; WGSL also + // rejects `f32(-inf)` as an identifier, so a real -inf literal isn't an + // option anyway. On f32 the sentinel drifts slightly each recursion step + // (log(2) per step when both log_sum_exp inputs sit at the sentinel), + // which is why the recurrence compares against a threshold instead of + // checking `== neg_inf`. See `log_sum_exp2` for the long-sequence caveat. + let neg_inf = F::new(-6.0e4_f32); + let unreachable_threshold = F::new(-1.0e4_f32); + let one = F::new(1.0_f32); + + // Initialize alpha at t = 0 for s < l_prime_len; positions beyond that + // are never read by the recurrence (s < l_prime_len in every read) so + // they don't need to be touched. + let mut s = UNIT_POS_X as usize; + while s < l_prime_len { + let mut init = neg_inf; + if s == 0 { + init = log_probs[n * lp_n + blank_u * lp_c]; + } else if s == 1 { + let l1 = u32::cast_from(targets[n * tgt_n]) as usize; + init = log_probs[n * lp_n + l1 * lp_c]; + } + alpha[s] = init; + s += cube_dim; + } + sync_cube(); + + // Sequential time loop. Each iteration re-strides over s positions to + // compute alpha[t, s] from alpha[t-1, *] and writes back to the same + // shared memory after a full read fence. + for t in 1..input_len { + let mut s = UNIT_POS_X as usize; + while s < l_prime_len { + let l_class = l_prime_class::(s, targets, n, tgt_n, tgt_s, blank_u); + let log_p = log_probs[t * lp_t + n * lp_n + l_class * lp_c]; + + let l_class_m2 = if s >= 2 { + l_prime_class::(s - 2, targets, n, tgt_n, tgt_s, blank_u) + } else { + blank_u + }; + let skip_allowed = s >= 2 && l_class != blank_u && l_class != l_class_m2; + + let a_s = alpha[s]; + let mut a_s_m1 = neg_inf; + if s >= 1 { + a_s_m1 = alpha[s - 1]; + } + let mut a_s_m2 = neg_inf; + if s >= 2 { + a_s_m2 = alpha[s - 2]; + } + + alpha[alpha_cap + s] = recurrence_step::( + a_s, + a_s_m1, + a_s_m2, + log_p, + skip_allowed, + unreachable_threshold, + one, + ); + s += cube_dim; + } + sync_cube(); + + // Second pass: copy scratch back into the active alpha slots. + let mut s = UNIT_POS_X as usize; + while s < l_prime_len { + alpha[s] = alpha[alpha_cap + s]; + s += cube_dim; + } + sync_cube(); + } + + // Reduce: only thread 0 writes the output for this batch element. + if UNIT_POS_X == 0 { + let last_blank = alpha[2 * target_len]; + // Guard target_len = 0: index 2*0 - 1 underflows. Use -inf so + // log_sum_exp(last_blank, -inf) = last_blank (log_sum_exp(x, x) = x+ln2 + // would be wrong here). + let mut last_label = neg_inf; + if target_len > 0 { + last_label = alpha[2 * target_len - 1]; + } + output[n] = finalize_nll::( + last_blank, + last_label, + target_len, + unreachable_threshold, + one, + ); + } +} + +/// Fused CTC loss for burn-cubecl. Single kernel launch covers the entire +/// alpha recursion across all timesteps. +/// +/// Panics if `2 * max_target_len + 1` exceeds `SHARED_ALPHA_CAPACITY` (8192). +pub fn ctc_loss( + log_probs: CubeTensor, + targets: CubeTensor, + input_lengths: CubeTensor, + target_lengths: CubeTensor, + blank: usize, +) -> CubeTensor { + // Manual stride indexing below requires a contiguous physical layout; + // fusion-produced tensors may arrive with layouts that break that + // assumption. No-op when already contiguous. + let log_probs = into_contiguous(log_probs); + let targets = into_contiguous(targets); + let input_lengths = into_contiguous(input_lengths); + let target_lengths = into_contiguous(target_lengths); + + let log_probs_shape = log_probs.shape(); + let [_t, batch_size, _c] = log_probs_shape.dims::<3>(); + let target_shape = targets.shape(); + let max_target_len = target_shape.dims::<2>()[1]; + let max_l_prime = 2 * max_target_len + 1; + + assert!( + max_l_prime as u32 <= SHARED_ALPHA_CAPACITY, + "ctc_loss: 2 * max_target_len + 1 = {} exceeds the kernel's shared-memory \ + alpha capacity ({}). Reduce target length or raise SHARED_ALPHA_CAPACITY.", + max_l_prime, + SHARED_ALPHA_CAPACITY, + ); + + // Pick a thread count that fits the runtime's per-cube limit. We don't + // need one thread per s position - threads stride over s. + let hw_max = log_probs.client.properties().hardware.max_cube_dim.0; + let cube_dim_x = (max_l_prime as u32).min(hw_max).min(256); + + let client = log_probs.client.clone(); + let device = log_probs.device.clone(); + let f_dtype = log_probs.dtype; + let i_dtype = targets.dtype; + let output = empty_device_dtype::(client.clone(), device, Shape::new([batch_size]), f_dtype); + + let cube_count = CubeCount::Static(batch_size as u32, 1, 1); + let cube_dim = CubeDim::new_1d(cube_dim_x); + + // Pass the actual max_l_prime (not the static capacity) so shared memory + // is sized to what we need. Metal limits threadgroup memory to 32 KB; + // allocating 2 * 8192 * sizeof(f32) = 64 KB would silently corrupt on + // Apple GPUs. Different max_l_prime values trigger separate kernel + // compilations (it's a comptime param), but that's fine: target lengths + // are stable within a dataset. + ctc_loss_kernel::launch::( + &client, + cube_count, + cube_dim, + log_probs.into_tensor_arg(), + targets.into_tensor_arg(), + input_lengths.into_tensor_arg(), + target_lengths.into_tensor_arg(), + output.clone().into_tensor_arg(), + blank as u32, + max_l_prime as u32, + [ + dtype_to_storage_type(f_dtype), + dtype_to_storage_type(i_dtype), + ], + ); + + output +} + +/// Fused CTC alpha + beta recursion kernel. +/// +/// Runs the full forward alpha recursion and reverse beta recursion for one +/// batch element per cube, reusing the same shared-memory layout twice. +/// Writes `alpha_out[T, N, 2S+1]`, `beta_out[T, N, 2S+1]` and the per-sample +/// negative log-likelihood `nll_out[N]`. The three outputs are everything the +/// default CTC gradient-composition helper needs, so the caller can finish the +/// backward pass with a handful of element-wise tensor ops. +/// +/// The alpha phase is identical to `ctc_loss_kernel` except it additionally +/// publishes each row to global memory. The beta phase mirrors it in reverse: +/// initialize at `t = input_len - 1` from `log_probs[t, l'[s]]` at the two +/// boundary `s` positions, then step backward reading `beta[t+1, s]`, +/// `beta[t+1, s+1]`, and (when the skip transition is allowed) `beta[t+1, s+2]`. +#[cube(launch)] +fn ctc_alpha_beta_kernel( + log_probs: &Tensor, // [T, N, C] + targets: &Tensor, // [N, S_max] + input_lengths: &Tensor, // [N] + target_lengths: &Tensor, // [N] + alpha_out: &mut Tensor, // [T, N, 2S+1] + beta_out: &mut Tensor, // [T, N, 2S+1] + nll_out: &mut Tensor, // [N] + blank: u32, + #[comptime] alpha_capacity: u32, + #[define(F, I)] _dtypes: [StorageType; 2], +) { + let n = CUBE_POS_X as usize; + let cube_dim = CUBE_DIM_X as usize; + let alpha_cap = alpha_capacity as usize; + let blank_u = blank as usize; + + let target_len = u32::cast_from(target_lengths[n]) as usize; + let input_len = u32::cast_from(input_lengths[n]) as usize; + let l_prime_len = 2 * target_len + 1; + + // Empty input: alpha_out and beta_out stay at the host-side -inf pre-fill. + // Emit the semantically correct nll (0 for target_len=0, +inf otherwise). + if input_len == 0 { + if UNIT_POS_X == 0 { + nll_out[n] = empty_input_nll::(target_len); + } + terminate!(); + } + + let lp_t = log_probs.stride(0); + let lp_n = log_probs.stride(1); + let lp_c = log_probs.stride(2); + let tgt_n = targets.stride(0); + let tgt_s = targets.stride(1); + let ao_t = alpha_out.stride(0); + let ao_n = alpha_out.stride(1); + let ao_s = alpha_out.stride(2); + let bo_t = beta_out.stride(0); + let bo_n = beta_out.stride(1); + let bo_s = beta_out.stride(2); + + // Shared memory layout: [0..alpha_cap] is the active row; [alpha_cap..2*alpha_cap] + // is scratch for the next row. Same layout is reused for alpha and beta. Beta + // reads are guarded by `s + 1 < l_prime_len` / `s + 2 < l_prime_len`, so the + // residual alpha values sitting in the active row between phases are never + // observed by beta (its boundary init overwrites every slot it reads). + let mut state = Shared::new_slice(2 * alpha_cap); + // Sentinel for unreachable states. See ctc_loss_kernel for the full + // rationale: f16's 65504 magnitude cap forces the -6e4 floor, WGSL + // rejects f32(-inf) literals, and the threshold catches sentinel drift. + let neg_inf = F::new(-6.0e4_f32); + let unreachable_threshold = F::new(-1.0e4_f32); + let one = F::new(1.0_f32); + + // Alpha phase (forward). + // + // Initialize alpha at t = 0 for s < l_prime_len. Positions beyond + // l_prime_len are never read by the recurrence, so they don't need + // to be touched in shared memory; and they stay at the host-side -inf + // pre-fill in alpha_out. + let mut s = UNIT_POS_X as usize; + while s < l_prime_len { + let mut init = neg_inf; + if s == 0 { + init = log_probs[n * lp_n + blank_u * lp_c]; + } else if s == 1 { + let l1 = u32::cast_from(targets[n * tgt_n]) as usize; + init = log_probs[n * lp_n + l1 * lp_c]; + } + state[s] = init; + alpha_out[n * ao_n + s * ao_s] = init; + s += cube_dim; + } + sync_cube(); + + for t in 1..input_len { + let mut s = UNIT_POS_X as usize; + while s < l_prime_len { + let l_class = l_prime_class::(s, targets, n, tgt_n, tgt_s, blank_u); + let log_p = log_probs[t * lp_t + n * lp_n + l_class * lp_c]; + + let l_class_m2 = if s >= 2 { + l_prime_class::(s - 2, targets, n, tgt_n, tgt_s, blank_u) + } else { + blank_u + }; + let skip_allowed = s >= 2 && l_class != blank_u && l_class != l_class_m2; + + let a_s = state[s]; + let mut a_s_m1 = neg_inf; + if s >= 1 { + a_s_m1 = state[s - 1]; + } + let mut a_s_m2 = neg_inf; + if s >= 2 { + a_s_m2 = state[s - 2]; + } + + state[alpha_cap + s] = recurrence_step::( + a_s, + a_s_m1, + a_s_m2, + log_p, + skip_allowed, + unreachable_threshold, + one, + ); + s += cube_dim; + } + sync_cube(); + + let mut s = UNIT_POS_X as usize; + while s < l_prime_len { + state[s] = state[alpha_cap + s]; + alpha_out[t * ao_t + n * ao_n + s * ao_s] = state[s]; + s += cube_dim; + } + sync_cube(); + } + + if UNIT_POS_X == 0 { + let last_blank = state[2 * target_len]; + // See ctc_loss_kernel: -inf sentinel keeps log_sum_exp correct for target_len = 0. + let mut last_label = neg_inf; + if target_len > 0 { + last_label = state[2 * target_len - 1]; + } + nll_out[n] = finalize_nll::( + last_blank, + last_label, + target_len, + unreachable_threshold, + one, + ); + } + + // Fence thread 0's read of state[2*target_len] / state[2*target_len - 1] + // against the beta boundary init, which writes those same positions. + sync_cube(); + + // Beta phase (reverse). + // + // Boundary initialization at t = input_len - 1: set beta[s] = log_probs[t, l'[s]] + // at s = 2*target_len, and when target_len > 0 also at s = 2*target_len - 1. + // All other s positions in range get -inf. + let t_last = input_len - 1; + let mut s = UNIT_POS_X as usize; + while s < l_prime_len { + let is_last_blank = s == 2 * target_len; + let is_last_label = target_len > 0 && s == 2 * target_len - 1; + let mut init = neg_inf; + if is_last_blank || is_last_label { + let l_class = l_prime_class::(s, targets, n, tgt_n, tgt_s, blank_u); + init = log_probs[t_last * lp_t + n * lp_n + l_class * lp_c]; + } + state[s] = init; + beta_out[t_last * bo_t + n * bo_n + s * bo_s] = init; + s += cube_dim; + } + sync_cube(); + + // Step back from t = input_len - 2 down to t = 0. + for t_rev in 1..input_len { + let t = input_len - 1 - t_rev; + + let mut s = UNIT_POS_X as usize; + while s < l_prime_len { + let l_class = l_prime_class::(s, targets, n, tgt_n, tgt_s, blank_u); + let log_p = log_probs[t * lp_t + n * lp_n + l_class * lp_c]; + + let l_class_p2 = if s + 2 < l_prime_len { + l_prime_class::(s + 2, targets, n, tgt_n, tgt_s, blank_u) + } else { + blank_u + }; + let skip_allowed = s + 2 < l_prime_len && l_class != blank_u && l_class != l_class_p2; + + let b_s = state[s]; + let mut b_s_p1 = neg_inf; + if s + 1 < l_prime_len { + b_s_p1 = state[s + 1]; + } + let mut b_s_p2 = neg_inf; + if s + 2 < l_prime_len { + b_s_p2 = state[s + 2]; + } + + state[alpha_cap + s] = recurrence_step::( + b_s, + b_s_p1, + b_s_p2, + log_p, + skip_allowed, + unreachable_threshold, + one, + ); + s += cube_dim; + } + sync_cube(); + + let mut s = UNIT_POS_X as usize; + while s < l_prime_len { + state[s] = state[alpha_cap + s]; + beta_out[t * bo_t + n * bo_n + s * bo_s] = state[s]; + s += cube_dim; + } + sync_cube(); + } +} + +/// Host entry point for the fused alpha + beta + nll kernel. +/// +/// Returns `(log_alpha_full, log_beta_full, nll)` with shapes +/// `([T, N, 2S+1], [T, N, 2S+1], [N])`. Positions outside the valid +/// `(t < input_length, s < 2*target_length+1)` rectangle hold the +/// pre-fill value `-inf`, matching the default backend's convention. +/// +/// Panics if `2 * max_target_len + 1` exceeds `SHARED_ALPHA_CAPACITY`. +pub fn ctc_alpha_beta( + log_probs: CubeTensor, + targets: CubeTensor, + input_lengths: CubeTensor, + target_lengths: CubeTensor, + blank: usize, +) -> (CubeTensor, CubeTensor, CubeTensor) { + // Manual stride indexing below requires a contiguous physical layout; + // fusion-produced tensors may arrive with layouts that break that + // assumption. No-op when already contiguous. + let log_probs = into_contiguous(log_probs); + let targets = into_contiguous(targets); + let input_lengths = into_contiguous(input_lengths); + let target_lengths = into_contiguous(target_lengths); + + let log_probs_shape = log_probs.shape(); + let [max_input_length, batch_size, _c] = log_probs_shape.dims::<3>(); + let target_shape = targets.shape(); + let max_target_len = target_shape.dims::<2>()[1]; + let max_l_prime = 2 * max_target_len + 1; + + assert!( + max_l_prime as u32 <= SHARED_ALPHA_CAPACITY, + "ctc_loss_backward: 2 * max_target_len + 1 = {} exceeds the kernel's shared-memory \ + alpha capacity ({}). Reduce target length or raise SHARED_ALPHA_CAPACITY.", + max_l_prime, + SHARED_ALPHA_CAPACITY, + ); + + let hw_max = log_probs.client.properties().hardware.max_cube_dim.0; + let cube_dim_x = (max_l_prime as u32).min(hw_max).min(256); + + let client = log_probs.client.clone(); + let device = log_probs.device.clone(); + let f_dtype = log_probs.dtype; + let i_dtype = targets.dtype; + + // Pre-fill alpha/beta with -inf so positions the kernel doesn't touch + // (s >= 2U+1, or t outside the valid range for an individual batch + // element) are not read as stale zeros by the gradient composition. + let shape_abt = Shape::new([max_input_length, batch_size, max_l_prime]); + let neg_inf = InputScalar::new(f32::NEG_INFINITY, dtype_to_storage_type(f_dtype)); + let alpha_out = crate::ops::numeric::full_device_dtype::( + client.clone(), + shape_abt.clone(), + device.clone(), + neg_inf, + f_dtype, + ); + let beta_out = crate::ops::numeric::full_device_dtype::( + client.clone(), + shape_abt, + device.clone(), + neg_inf, + f_dtype, + ); + let nll_out = + empty_device_dtype::(client.clone(), device, Shape::new([batch_size]), f_dtype); + + let cube_count = CubeCount::Static(batch_size as u32, 1, 1); + let cube_dim = CubeDim::new_1d(cube_dim_x); + + ctc_alpha_beta_kernel::launch::( + &client, + cube_count, + cube_dim, + log_probs.into_tensor_arg(), + targets.into_tensor_arg(), + input_lengths.into_tensor_arg(), + target_lengths.into_tensor_arg(), + alpha_out.clone().into_tensor_arg(), + beta_out.clone().into_tensor_arg(), + nll_out.clone().into_tensor_arg(), + blank as u32, + max_l_prime as u32, + [ + dtype_to_storage_type(f_dtype), + dtype_to_storage_type(i_dtype), + ], + ); + + (alpha_out, beta_out, nll_out) +} diff --git a/crates/burn-cubecl/src/kernel/fft/base.rs b/crates/burn-cubecl/src/kernel/fft/base.rs new file mode 100644 index 0000000..3aded0c --- /dev/null +++ b/crates/burn-cubecl/src/kernel/fft/base.rs @@ -0,0 +1,164 @@ +use crate::kernel::index::slice; +use crate::ops::numeric::{empty_device_dtype, zeros}; +use crate::{CubeRuntime, tensor::CubeTensor}; +use burn_backend::{DType, TensorMetadata}; +use burn_std::Slice; +use cubecl::prelude::*; +use cubek::fft::{irfft_launch, rfft_launch}; + +// Materializes a padded tensor (allocate + copy) because rfft_launch/irfft_launch +// in the external cubek crate don't support virtual padding via a length parameter. +// See: https://github.com/tracel-ai/cubek/issues/194 +fn pad_to_length( + tensor: CubeTensor, + dim: usize, + target: usize, +) -> CubeTensor { + let shape = tensor.shape(); + let current = shape[dim]; + if current == target { + return tensor; + } + if current > target { + let ranges: Vec<_> = shape + .iter() + .enumerate() + .map(|(i, &s)| if i == dim { 0..target } else { 0..s }) + .collect(); + return slice(tensor, &ranges); + } + let mut padded_shape = shape.clone(); + padded_shape[dim] = target; + let padded = zeros::(tensor.device.clone(), padded_shape, tensor.dtype); + let slices: Vec = shape.iter().map(|&s| Slice::from(0..s)).collect(); + crate::kernel::index::slice_assign::(padded, &slices, tensor) +} + +/// Launch the rfft kernel with optional padding for non-power-of-two sizes. +/// +/// Signal is first truncated or zero-padded to `n` (when provided), then internally +/// padded to the next power of two so the kernel operates on a pow2 length. +/// Output bin count is `fft_size / 2 + 1` where `fft_size = next_pow2(n)`. +pub fn rfft( + signal: CubeTensor, + dim: usize, + n: Option, +) -> (CubeTensor, CubeTensor) { + let dtype = match signal.dtype { + DType::F64 => f64::as_type_native_unchecked().storage_type(), + DType::F32 => f32::as_type_native_unchecked().storage_type(), + _ => panic!("Unsupported type {:?}", signal.dtype), + }; + + let input_device = signal.device.clone(); + let input_dtype = signal.dtype; + let input_shape = signal.shape(); + let requested_n = n.unwrap_or(input_shape[dim]); + let fft_size = requested_n.next_power_of_two(); + + // Truncate/pad to requested_n, THEN pad to fft_size; otherwise for + // requested_n < input_len < fft_size we would keep bogus samples in [n, fft_size). + let signal = pad_to_length(signal, dim, requested_n); + let signal = pad_to_length(signal, dim, fft_size); + + let signal_shape = signal.shape(); + let mut output_shape = signal_shape.clone(); + output_shape[dim] = fft_size / 2 + 1; + + let output_re = empty_device_dtype( + signal.client.clone(), + signal.device.clone(), + output_shape.clone(), + signal.dtype, + ); + let output_im = empty_device_dtype( + signal.client.clone(), + signal.device.clone(), + output_shape.clone(), + signal.dtype, + ); + + rfft_launch( + &signal.client.clone(), + signal.binding(), + output_re.clone().binding(), + output_im.clone().binding(), + dim, + dtype, + ) + .unwrap_or_else(|e| { + panic!( + "rfft kernel launch failed (device={input_device:?}, dtype={input_dtype:?}, \ + dim={dim}, requested_n={requested_n}, fft_size={fft_size}): {e}" + ) + }); + + (output_re, output_im) +} + +/// Launch the irfft kernel with optional padding for non-power-of-two sizes. +pub fn irfft( + spectrum_re: CubeTensor, + spectrum_im: CubeTensor, + dim: usize, + n: Option, +) -> CubeTensor { + assert!( + spectrum_re.shape() == spectrum_im.shape(), + "irfft: spectrum_re and spectrum_im shapes must match" + ); + assert!( + spectrum_re.shape()[dim] >= 1, + "irfft: spectrum dimension cannot be empty" + ); + assert!( + !matches!(n, Some(0)), + "irfft: n must be >= 1 when specified, got Some(0)" + ); + + let dtype = match spectrum_re.dtype { + DType::F64 => f64::as_type_native_unchecked().storage_type(), + DType::F32 => f32::as_type_native_unchecked().storage_type(), + _ => panic!("Unsupported type {:?}", spectrum_re.dtype), + }; + + let input_device = spectrum_re.device.clone(); + let input_dtype = spectrum_re.dtype; + let requested_n = n.unwrap_or((spectrum_re.shape()[dim] - 1) * 2); + let fft_size = requested_n.next_power_of_two().max(1); + let half_fft = fft_size / 2 + 1; + + let spectrum_re = pad_to_length(spectrum_re, dim, half_fft); + let spectrum_im = pad_to_length(spectrum_im, dim, half_fft); + + let mut signal_shape = spectrum_re.shape().clone(); + signal_shape[dim] = fft_size; + + let signal = empty_device_dtype( + spectrum_re.client.clone(), + spectrum_re.device.clone(), + signal_shape, + spectrum_re.dtype, + ); + + irfft_launch( + &spectrum_re.client.clone(), + spectrum_re.binding(), + spectrum_im.binding(), + signal.clone().binding(), + dim, + dtype, + ) + .unwrap_or_else(|e| { + panic!( + "irfft kernel launch failed (device={input_device:?}, dtype={input_dtype:?}, \ + dim={dim}, requested_n={requested_n}, fft_size={fft_size}): {e}" + ) + }); + + if fft_size > requested_n { + pad_to_length(signal, dim, requested_n) + } else { + signal + } +} diff --git a/crates/burn-cubecl/src/kernel/fft/mod.rs b/crates/burn-cubecl/src/kernel/fft/mod.rs new file mode 100644 index 0000000..cbcb6ac --- /dev/null +++ b/crates/burn-cubecl/src/kernel/fft/mod.rs @@ -0,0 +1,3 @@ +mod base; + +pub use base::*; diff --git a/crates/burn-cubecl/src/kernel/grid_sample/base.rs b/crates/burn-cubecl/src/kernel/grid_sample/base.rs new file mode 100644 index 0000000..ef8d36a --- /dev/null +++ b/crates/burn-cubecl/src/kernel/grid_sample/base.rs @@ -0,0 +1,163 @@ +use cubecl::prelude::*; + +use crate::{CubeRuntime, tensor::CubeTensor}; +use burn_backend::ops::{GridSampleOptions, GridSamplePaddingMode, InterpolateMode}; + +use super::bilinear::grid_sample_bilinear_launch; + +/// Grid sample operation supporting bilinear interpolation +pub fn grid_sample( + input: CubeTensor, + grid: CubeTensor, + options: GridSampleOptions, +) -> CubeTensor { + match options.mode { + InterpolateMode::Bilinear => grid_sample_bilinear_launch(input, grid, options), + _ => panic!( + "Unsupported grid_sample interpolation mode: {:?}", + options.mode + ), + } +} + +/// Compile-time padding mode for kernel specialization +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum PaddingMode { + /// Fill with zeros for out-of-bounds coordinates. + Zeros, + /// Clamp coordinates to the border (use nearest edge value). + Border, + /// Reflect coordinates at the boundary. + Reflection, +} + +impl From for PaddingMode { + fn from(mode: GridSamplePaddingMode) -> Self { + match mode { + GridSamplePaddingMode::Zeros => PaddingMode::Zeros, + GridSamplePaddingMode::Border => PaddingMode::Border, + GridSamplePaddingMode::Reflection => PaddingMode::Reflection, + } + } +} + +/// Fetch value based on padding mode (dispatch to appropriate handler) +#[cube] +pub(crate) fn fetch_value( + input: &Tensor, + base: usize, + stride_h: usize, + stride_w: usize, + y: i32, + x: i32, + h: i32, + w: i32, + #[comptime] padding_mode: PaddingMode, +) -> F { + match padding_mode { + PaddingMode::Zeros => fetch_with_zeros(input, base, stride_h, stride_w, y, x, h, w), + PaddingMode::Border => fetch_with_border(input, base, stride_h, stride_w, y, x, h, w), + PaddingMode::Reflection => { + fetch_with_reflection(input, base, stride_h, stride_w, y, x, h, w) + } + } +} + +/// Fetch value with zeros padding (return 0 for out-of-bounds). +#[cube] +pub(crate) fn fetch_with_zeros( + input: &Tensor, + base: usize, + stride_h: usize, + stride_w: usize, + y: i32, + x: i32, + h: i32, + w: i32, +) -> F { + let in_bounds = x >= 0 && x < w && y >= 0 && y < h; + let x_clamped = clamp(x, 0, w - 1) as usize; + let y_clamped = clamp(y, 0, h - 1) as usize; + let idx = base + y_clamped * stride_h + x_clamped * stride_w; + select(in_bounds, input[idx], F::new(0.0_f32)) +} + +/// Fetch value with border padding (clamp to edge). +#[cube] +pub(crate) fn fetch_with_border( + input: &Tensor, + base: usize, + stride_h: usize, + stride_w: usize, + y: i32, + x: i32, + h: i32, + w: i32, +) -> F { + let x_clamped = clamp(x, 0, w - 1) as usize; + let y_clamped = clamp(y, 0, h - 1) as usize; + let idx = base + y_clamped * stride_h + x_clamped * stride_w; + input[idx] +} + +/// Fetch value with reflection padding. +/// Assumes float reflection was applied to center, so indices are at most 2 steps out of bounds. +#[cube] +pub(crate) fn fetch_with_reflection( + input: &Tensor, + base: usize, + stride_h: usize, + stride_w: usize, + y: i32, + x: i32, + h: i32, + w: i32, +) -> F { + let x_reflected = reflect_coord_bounded(x, w); + let y_reflected = reflect_coord_bounded(y, h); + let idx = base + y_reflected * stride_h + x_reflected * stride_w; + input[idx] +} + +/// Reflect an integer index that may be out of bounds. +/// After float reflection, indices can be up to 2 steps out for bicubic (1 step for bilinear). +#[cube] +fn reflect_coord_bounded(idx: i32, size: i32) -> usize { + let max_idx = size - 1; + let neg_reflected = -idx - 1; + let pos_reflected = 2 * max_idx + 1 - idx; + let result = select( + idx < 0, + neg_reflected, + select(idx > max_idx, pos_reflected, idx), + ); + clamp(result, 0, max_idx) as usize +} + +/// Reflect a float coordinate into the valid sampling range. +#[cube] +pub(crate) fn reflect_coord(coord: F, size: u32, #[comptime] align_corners: bool) -> F { + let size_f = F::cast_from(size); + if align_corners { + reflect_float_impl::(coord, F::new(0.0_f32), size_f - F::new(1.0_f32)) + } else { + reflect_float_impl::(coord, F::new(-0.5_f32), size_f - F::new(0.5_f32)) + } +} + +/// Reflect a float coordinate into [min_val, max_val] using a triangle wave pattern. +#[cube] +fn reflect_float_impl(coord: F, min_val: F, max_val: F) -> F { + let span = max_val - min_val; + + let is_valid = span > F::new(0.0_f32); + let safe_span = select(is_valid, span, F::new(1.0_f32)); + + // Triangle wave formula: span - |((x mod 2*span) - span)| + min_val + let period = safe_span * F::new(2.0_f32); + let x = (coord - min_val).abs(); + let x_mod = x - (x / period).floor() * period; + let reflected = safe_span - (x_mod - safe_span).abs() + min_val; + + select(is_valid, reflected, min_val) +} diff --git a/crates/burn-cubecl/src/kernel/grid_sample/bilinear.rs b/crates/burn-cubecl/src/kernel/grid_sample/bilinear.rs new file mode 100644 index 0000000..801f20e --- /dev/null +++ b/crates/burn-cubecl/src/kernel/grid_sample/bilinear.rs @@ -0,0 +1,181 @@ +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::std::FastDivmod; +use cubecl::{calculate_cube_count_elemwise, prelude::*}; + +use crate::{ + CubeRuntime, kernel::utils::address_type, ops::numeric::empty_device_dtype, tensor::CubeTensor, +}; +use burn_backend::{Shape, ops::GridSampleOptions}; + +use super::base::{PaddingMode, fetch_value, reflect_coord}; + +/// Grid sample with bilinear interpolation. +/// +/// Each thread processes all channels for one spatial output position: +/// 1. Reading (x, y) coordinates from the grid tensor (once per spatial position) +/// 2. Converting normalized [-1, 1] coords to pixel coordinates (once) +/// 3. For each channel: fetch 4 corner values, interpolate, and write output +#[cube(launch, address_type = "dynamic")] +fn grid_sample_bilinear_kernel( + input: &Tensor, // [N, C, H_in, W_in] + grid: &Tensor, // [N, H_out, W_out, 2] + output: &mut Tensor, // [N, C, H_out, W_out] + shape_spatial: Sequence>, // [N, H_out, W_out] for thread decomposition + #[comptime] align_corners: bool, + #[comptime] pad_mode: PaddingMode, + #[define(F)] _dtype: StorageType, +) { + // Thread index maps to spatial position (n, h_out, w_out) only + let spatial_idx = ABSOLUTE_POS; + let num_spatial = output.shape(0) * output.shape(2) * output.shape(3); + if spatial_idx >= num_spatial { + terminate!(); + } + + // Decompose spatial index into (n, h_out, w_out) + let (rem, w_out) = shape_spatial[2].div_mod(spatial_idx); + let (n, h_out) = shape_spatial[1].div_mod(rem); + + let channels = input.shape(1) as u32; + let h_in = input.shape(2) as u32; + let w_in = input.shape(3) as u32; + + // Read grid coordinates once per spatial position + let grid_offset = n * grid.stride(0) + h_out * grid.stride(1) + w_out * grid.stride(2); + let gx = grid[grid_offset]; // x coordinate in [-1, 1] + let gy = grid[grid_offset + 1]; // y coordinate in [-1, 1] + + // Convert normalized coordinates to pixel coordinates + let (px, py) = if align_corners { + let px = (gx + F::new(1.0_f32)) * F::cast_from((w_in - 1) as f32) / F::new(2.0_f32); + let py = (gy + F::new(1.0_f32)) * F::cast_from((h_in - 1) as f32) / F::new(2.0_f32); + (px, py) + } else { + let px = + (gx + F::new(1.0_f32)) * F::cast_from(w_in as f32) / F::new(2.0_f32) - F::new(0.5_f32); + let py = + (gy + F::new(1.0_f32)) * F::cast_from(h_in as f32) / F::new(2.0_f32) - F::new(0.5_f32); + (px, py) + }; + + // For reflection padding, reflect the coordinate into the valid sampling range. + // This ensures integer indices are at most 1 step out of bounds. + let (px, py) = if comptime!(pad_mode == PaddingMode::Reflection) { + let px = reflect_coord::(px, w_in, align_corners); + let py = reflect_coord::(py, h_in, align_corners); + (px, py) + } else { + (px, py) + }; + + // Compute floor and ceil indices + let x0_f = px.floor(); + let y0_f = py.floor(); + let x1_f = x0_f + F::new(1.0_f32); + let y1_f = y0_f + F::new(1.0_f32); + + // Compute interpolation weights + let wx = px - x0_f; + let wy = py - y0_f; + let wx_ = F::new(1.0_f32) - wx; + let wy_ = F::new(1.0_f32) - wy; + + // Convert to integers for indexing + let x0 = i32::cast_from(x0_f); + let y0 = i32::cast_from(y0_f); + let x1 = i32::cast_from(x1_f); + let y1 = i32::cast_from(y1_f); + + let w_in = w_in as i32; + let h_in = h_in as i32; + + // Pre-compute strides + let stride_n = input.stride(0); + let stride_c = input.stride(1); + let stride_h = input.stride(2); + let stride_w = input.stride(3); + let out_stride_n = output.stride(0); + let out_stride_c = output.stride(1); + let out_stride_h = output.stride(2); + let out_stride_w = output.stride(3); + + // Base offsets for this spatial position + let in_base_n = n * stride_n; + let out_base_spatial = n * out_stride_n + h_out * out_stride_h + w_out * out_stride_w; + + // Loop over all channels - grid coords and weights are reused + for c in 0..channels { + let in_base = in_base_n + c as usize * stride_c; + + let v00 = fetch_value( + input, in_base, stride_h, stride_w, y0, x0, h_in, w_in, pad_mode, + ); + let v01 = fetch_value( + input, in_base, stride_h, stride_w, y1, x0, h_in, w_in, pad_mode, + ); + let v10 = fetch_value( + input, in_base, stride_h, stride_w, y0, x1, h_in, w_in, pad_mode, + ); + let v11 = fetch_value( + input, in_base, stride_h, stride_w, y1, x1, h_in, w_in, pad_mode, + ); + + // Bilinear interpolation + let result = wx_ * wy_ * v00 + wx_ * wy * v01 + wx * wy_ * v10 + wx * wy * v11; + + let out_idx = out_base_spatial + c as usize * out_stride_c; + output[out_idx] = result; + } +} + +/// Launch the grid sample bilinear kernel +pub(crate) fn grid_sample_bilinear_launch( + input: CubeTensor, + grid: CubeTensor, + options: GridSampleOptions, +) -> CubeTensor { + let [batch_size, channels, _h_in, _w_in] = input.meta.shape().dims(); + let [_n, h_out, w_out, two] = grid.meta.shape().dims(); + assert_eq!(two, 2, "Grid last dimension must be 2"); + + // Create output tensor [N, C, H_out, W_out] + let output_shape = Shape::new([batch_size, channels, h_out, w_out]); + let output = empty_device_dtype( + input.client.clone(), + input.device.clone(), + output_shape, + input.dtype, + ); + + // Spatial threading: one thread per (n, h_out, w_out) + let spatial_shape = Shape::new([batch_size, h_out, w_out]); + let num_spatial = spatial_shape.num_elements(); + + let mut shape_spatial = SequenceArg::new(); + for dim in spatial_shape.iter() { + shape_spatial.push(*dim); + } + + let cube_dim = CubeDim::new(&input.client, num_spatial); + let cube_count = calculate_cube_count_elemwise(&input.client, num_spatial, cube_dim); + + let padding_mode: PaddingMode = options.padding_mode.into(); + + let dtype = input.dtype; + + grid_sample_bilinear_kernel::launch( + &output.client, + cube_count, + cube_dim, + address_type!(input, grid, output), + input.into_tensor_arg(), + grid.into_tensor_arg(), + output.clone().into_tensor_arg(), + shape_spatial, + options.align_corners, + padding_mode, + dtype_to_storage_type(dtype), + ); + + output +} diff --git a/crates/burn-cubecl/src/kernel/grid_sample/mod.rs b/crates/burn-cubecl/src/kernel/grid_sample/mod.rs new file mode 100644 index 0000000..7bd74df --- /dev/null +++ b/crates/burn-cubecl/src/kernel/grid_sample/mod.rs @@ -0,0 +1,4 @@ +mod base; +mod bilinear; + +pub use base::*; diff --git a/crates/burn-cubecl/src/kernel/index/flip.rs b/crates/burn-cubecl/src/kernel/index/flip.rs new file mode 100644 index 0000000..f8daf6e --- /dev/null +++ b/crates/burn-cubecl/src/kernel/index/flip.rs @@ -0,0 +1,103 @@ +use crate::{ + CubeRuntime, + kernel::utils::{address_type, shape_divmod}, + ops::numeric::empty_device_dtype, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{DType, TensorMetadata}; +use cubecl::{ + calculate_cube_count_elemwise, + prelude::*, + std::{FastDivmod, tensor::layout::linear::LinearViewMut}, +}; + +#[cube(launch_unchecked, address_type = "dynamic")] +fn flip_kernel( + input: &Tensor, + mut output: LinearViewMut<'_, E>, + in_shape: Sequence>, + indices: Sequence, + #[define(E, Bool)] _dtypes: [StorageType; 2], +) { + if !output.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + let rank = in_shape.len().comptime(); + + let mut offset = ABSOLUTE_POS; + let mut offset_input = 0; + + #[unroll] + for i in 0..rank { + let dim = rank - i - 1; + let shape = input.shape(dim); + + let (rem, offset_local) = in_shape[dim].div_mod(offset); + offset = rem; + + let flip = indices.index(dim).get::() == Bool::from_int(1); + let offset_local = select(flip, shape - offset_local - 1, offset_local); + + offset_input += offset_local * input.stride(dim); + } + + output.write(ABSOLUTE_POS, input[offset_input]); +} + +pub(crate) fn flip( + tensor: CubeTensor, + indices: &[usize], + dtype_bool: DType, +) -> CubeTensor { + let output = empty_device_dtype( + tensor.client.clone(), + tensor.device.clone(), + tensor.shape(), + tensor.dtype, + ); + flip_on_output(tensor, output, indices, dtype_bool) +} + +pub(crate) fn flip_on_output( + tensor: CubeTensor, + output: CubeTensor, + indices: &[usize], + dtype_bool: DType, +) -> CubeTensor { + let dtype_input = tensor.dtype; + let ndims = tensor.meta.num_dims(); + let mut indices_sequence = SequenceArg::::new(); + + for i in 0..ndims { + indices_sequence.push({ + let val = indices.contains(&i) as u8; + InputScalar::new(val, dtype_to_storage_type(dtype_bool)) + }); + } + + let num_elements = output.meta.num_elements(); + let cube_dim = CubeDim::new(&tensor.client, num_elements); + let cube_count = calculate_cube_count_elemwise(&tensor.client, num_elements, cube_dim); + + let shape = shape_divmod(&tensor); + unsafe { + flip_kernel::launch_unchecked( + &output.client, + cube_count, + cube_dim, + address_type!(tensor, output), + tensor.into_tensor_arg(), + output.clone().into_linear_view(), + shape, + indices_sequence, + [ + dtype_to_storage_type(dtype_input), + dtype_to_storage_type(dtype_bool), + ], + ) + } + + output +} diff --git a/crates/burn-cubecl/src/kernel/index/gather.rs b/crates/burn-cubecl/src/kernel/index/gather.rs new file mode 100644 index 0000000..f30aaab --- /dev/null +++ b/crates/burn-cubecl/src/kernel/index/gather.rs @@ -0,0 +1,84 @@ +use crate::{ + CubeRuntime, + kernel::utils::{address_type, broadcast_strides, shape_divmod}, + ops::numeric::empty_device_dtype, + tensor::CubeTensor, +}; +use burn_backend::TensorMetadata; +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::std::{FastDivmod, tensor::index_offset_contiguous_fastdivmod}; +use cubecl::{CubeDim, std::tensor::layout::linear::LinearView}; +use cubecl::{calculate_cube_count_elemwise, prelude::*}; +use cubecl::{ + frontend::{ABSOLUTE_POS, Numeric, Tensor}, + std::tensor::layout::linear::LinearViewMut, +}; + +#[cube(launch_unchecked, address_type = "dynamic")] +fn gather_kernel( + input: &Tensor, + indices: LinearView<'_, I>, + mut output: LinearViewMut<'_, T>, + in_strides: Sequence, // zeroed out for broadcast dims and `dim` + out_shape: Sequence>, + dim: usize, + #[define(T, I)] _dtypes: [StorageType; 2], +) { + if !indices.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + let mut offset = index_offset_contiguous_fastdivmod( + ABSOLUTE_POS, + &out_shape, + &in_strides, + input.vector_size(), + ); + + offset += usize::cast_from(indices.read(ABSOLUTE_POS)) * input.stride(dim); + + output.write(ABSOLUTE_POS, input[offset]); +} + +pub(crate) fn gather( + dim: usize, + tensor: CubeTensor, + indices: CubeTensor, +) -> CubeTensor { + let shape_output = indices.shape(); + let total_elem = shape_output.num_elements(); + let output = empty_device_dtype( + tensor.client.clone(), + tensor.device.clone(), + shape_output, + tensor.dtype, + ); + + let cube_dim = CubeDim::new(&tensor.client, total_elem); + let cube_count = calculate_cube_count_elemwise(&tensor.client, total_elem, cube_dim); + let mut in_strides = broadcast_strides(&output, &tensor); + in_strides.values[dim] = 0; // Zero `dim` to exclude it from the indexing + + let (dtype, indices_dtype) = (tensor.dtype, indices.dtype); + + unsafe { + gather_kernel::launch_unchecked( + &output.client, + cube_count, + cube_dim, + address_type!(tensor, indices, output), + tensor.into_tensor_arg(), + indices.into_linear_view(), + output.clone().into_linear_view(), + in_strides, + shape_divmod(&output), + dim, + [ + dtype_to_storage_type(dtype), + dtype_to_storage_type(indices_dtype), + ], + ) + } + + output +} diff --git a/crates/burn-cubecl/src/kernel/index/gather_nd.rs b/crates/burn-cubecl/src/kernel/index/gather_nd.rs new file mode 100644 index 0000000..7dc5772 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/index/gather_nd.rs @@ -0,0 +1,124 @@ +use crate::kernel::utils::{shape_divmod, shape_divmod_range}; +use crate::{ + CubeRuntime, kernel::utils::address_type, ops::numeric::empty_device_dtype, tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::prelude::*; +use cubecl::std::FastDivmod; +use cubecl::std::tensor::layout::linear::LinearView; +use cubecl::{CubeDim, calculate_cube_count_elemwise}; + +/// gather_nd GPU kernel. +/// +/// Each thread handles one element of the output. +/// Work items = num_indices * slice_size. +#[cube(launch_unchecked, address_type = "dynamic")] +fn gather_nd_kernel( + data: &Tensor, + indices: LinearView<'_, I>, + output: &mut Tensor, + output_shape: Sequence>, + data_slice_shape: Sequence>, + slice_size: usize, + k: usize, + working_units: usize, + #[define(T, I)] _dtypes: [StorageType; 2], +) { + if ABSOLUTE_POS >= working_units { + terminate!(); + } + + let slice_offset = ABSOLUTE_POS % slice_size; + let index_idx = ABSOLUTE_POS / slice_size; + + // Compute flat offset into data from the K-dimensional index tuple + let idx_base = index_idx * k; + let mut base_offset = 0usize; + for j in 0..k { + let idx_val = usize::cast_from(indices.read(idx_base + j)); + base_offset += idx_val * data.stride(j); + } + + let slice_rank = data_slice_shape.len().comptime(); + let mut data_slice_offset = 0usize; + let mut remainder = slice_offset; + #[unroll] + for i in 0..slice_rank { + let dim = slice_rank - i - 1; + let (rem, coord) = data_slice_shape[dim].div_mod(remainder); + remainder = rem; + data_slice_offset += coord * data.stride(k + dim); + } + + let out_rank = output_shape.len().comptime(); + let mut out_offset = 0usize; + let mut remainder_o = ABSOLUTE_POS; + #[unroll] + for i in 0..out_rank { + let dim = out_rank - i - 1; + let (rem, coord) = output_shape[dim].div_mod(remainder_o); + remainder_o = rem; + out_offset += coord * output.stride(dim); + } + + output[out_offset] = data[base_offset + data_slice_offset]; +} + +pub(crate) fn gather_nd( + tensor: CubeTensor, + indices: CubeTensor, +) -> CubeTensor { + let data_shape = &tensor.meta.shape; + let idx_shape = &indices.meta.shape; + let m = idx_shape.num_dims(); + let k = idx_shape[m - 1]; + + // num_indices = product of first M-1 dims of indices + let num_indices: usize = idx_shape.as_slice()[..m - 1].iter().product(); + // slice_size = product of data.shape[K..] + let slice_size: usize = data_shape.as_slice()[k..].iter().product(); + let total_elem = num_indices * slice_size; + + // Output shape: idx_shape[..m-1] ++ data_shape[k..] + let mut out_dims: Vec = idx_shape.as_slice()[..m - 1].to_vec(); + out_dims.extend_from_slice(&data_shape.as_slice()[k..]); + let out_shape = burn_backend::Shape::from(out_dims); + + let output = empty_device_dtype( + tensor.client.clone(), + tensor.device.clone(), + out_shape, + tensor.dtype, + ); + + let cube_dim = CubeDim::new(&tensor.client, total_elem); + let cube_count = calculate_cube_count_elemwise(&tensor.client, total_elem, cube_dim); + + let (dtype, indices_dtype) = (tensor.dtype, indices.dtype); + + let data_slice_shape = shape_divmod_range(&tensor, k..data_shape.num_dims()); + let output_shape_divmod = shape_divmod(&output); + + unsafe { + gather_nd_kernel::launch_unchecked( + &output.client, + cube_count, + cube_dim, + address_type!(tensor, indices, output), + tensor.into_tensor_arg(), + indices.into_linear_view(), + output.clone().into_tensor_arg(), + output_shape_divmod, + data_slice_shape, + slice_size, + k, + total_elem, + [ + dtype_to_storage_type(dtype), + dtype_to_storage_type(indices_dtype), + ], + ) + } + + output +} diff --git a/crates/burn-cubecl/src/kernel/index/mod.rs b/crates/burn-cubecl/src/kernel/index/mod.rs new file mode 100644 index 0000000..86ee25b --- /dev/null +++ b/crates/burn-cubecl/src/kernel/index/mod.rs @@ -0,0 +1,22 @@ +mod flip; +mod gather; +mod gather_nd; +mod repeat_dim; +mod scatter; +mod scatter_nd; +mod select; +mod select_assign; +mod slice; +mod slice_assign; + +pub(crate) use flip::*; +pub(crate) use gather_nd::*; +pub(crate) use repeat_dim::*; +pub(crate) use scatter_nd::*; +pub(crate) use select::*; +pub(crate) use select_assign::*; +pub use slice::*; +pub(crate) use slice_assign::*; + +pub(crate) use gather::*; +pub(crate) use scatter::*; diff --git a/crates/burn-cubecl/src/kernel/index/repeat_dim.rs b/crates/burn-cubecl/src/kernel/index/repeat_dim.rs new file mode 100644 index 0000000..ee9956c --- /dev/null +++ b/crates/burn-cubecl/src/kernel/index/repeat_dim.rs @@ -0,0 +1,91 @@ +use crate::{ + CubeRuntime, + kernel::utils::{address_type, shape_divmod}, + ops::numeric::empty_device_dtype, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::{calculate_cube_count_elemwise, prelude::*, std::FastDivmod}; + +#[cube(launch_unchecked, address_type = "dynamic")] +fn repeat_dim_kernel( + input: &Tensor, + output: &mut Tensor, + out_shape: Sequence>, + in_shape: FastDivmod, + #[comptime] dim: usize, + #[define(E)] _dtype: StorageType, +) { + if ABSOLUTE_POS >= output.len() { + terminate!(); + } + + let rank = out_shape.len().comptime(); + + let mut pos = ABSOLUTE_POS; + let mut offset_input = 0; + let mut offset_output = 0; + + #[unroll] + for i in 0..rank { + let i = rank - i - 1; + + let (rem, mut local_pos) = out_shape[i].div_mod(pos); + pos = rem; + + offset_output += local_pos * output.stride(i); + + if i == dim { + local_pos = in_shape.modulo(local_pos); + } + + offset_input += local_pos * input.stride(i); + } + + output[offset_output] = input[offset_input]; +} + +pub(crate) fn repeat_dim( + mut input: CubeTensor, + dim: usize, + times: usize, +) -> CubeTensor { + if input.meta.shape()[dim] == 1 { + input.meta.strides[dim] = 0; + input.meta.shape = input.meta.shape.clone().repeat(dim, times).unwrap(); + return input; + } + + let shape = input.meta.shape.clone().repeat(dim, times).unwrap(); + + // Create output handle + let output = empty_device_dtype( + input.client.clone(), + input.device.clone(), + shape, + input.dtype, + ); + + let working_units = output.meta.num_elements(); + let cube_dim = CubeDim::new(&input.client, working_units); + let cube_count = calculate_cube_count_elemwise(&input.client, working_units, cube_dim); + + let shape_arg = input.meta.shape()[dim]; + + unsafe { + repeat_dim_kernel::launch_unchecked( + &output.client, + cube_count, + cube_dim, + address_type!(input, output), + input.into_tensor_arg(), + output.clone().into_tensor_arg(), + shape_divmod(&output), + shape_arg, + dim, + dtype_to_storage_type(output.dtype), + ) + }; + + output +} diff --git a/crates/burn-cubecl/src/kernel/index/scatter.rs b/crates/burn-cubecl/src/kernel/index/scatter.rs new file mode 100644 index 0000000..2e36c60 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/index/scatter.rs @@ -0,0 +1,116 @@ +use crate::{ + CubeRuntime, + kernel::{ + AddOp, BinaryOp, BinaryOpFamily, OrOp, + utils::{address_type, shape_divmod}, + }, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::{CubeDim, calculate_cube_count_elemwise}; +use cubecl::{prelude::*, std::FastDivmod}; + +#[cube(launch_unchecked, address_type = "dynamic")] +fn scatter_kernel( + input: &mut Tensor, + indices: &Tensor, + value: &Tensor, + in_shape: Sequence>, + #[comptime] dim: usize, + #[define(T, I)] _dtypes: [StorageType; 2], +) { + let rank = in_shape.len().comptime(); + let stride_input = input.stride(dim); + let stride_value = value.stride(dim); + let stride_indices = indices.stride(dim); + let shape_value = value.shape(dim); + + let mut offset = ABSOLUTE_POS; + let mut offset_input = 0; + let mut offset_indices = 0; + let mut offset_value = 0; + let mut num_elems = 1; + + #[unroll] + for i in 0..rank { + let i = rank - i - 1; + if i != dim { + let shape_input_loop = input.shape(i); + + let (rem, local_pos) = in_shape[i].div_mod(offset); + offset = rem; + + offset_input += local_pos * input.stride(i); + offset_indices += local_pos * indices.stride(i); + offset_value += local_pos * value.stride(i); + + num_elems *= shape_input_loop; + } + } + + let should_stop = ABSOLUTE_POS >= num_elems; + if should_stop { + terminate!(); + } + + for i in 0..shape_value { + let value_idx = (stride_value * i) + offset_value; + let index_idx = (stride_indices * i) + offset_indices; + + let value = value[value_idx]; + let index = usize::cast_from(indices[index_idx]); + + let input_idx = (stride_input * index) + offset_input; + + let value = Op::BinaryOp::>::execute( + Vector::cast_from(input[input_idx]), + Vector::cast_from(value), + ); + input[input_idx] = value.extract(0); + } +} + +pub(crate) fn scatter( + dim: usize, + tensor: CubeTensor, + indices: CubeTensor, + value: CubeTensor, + is_bool: bool, +) -> CubeTensor { + let tensor = match tensor.can_mut() && tensor.is_nonoverlapping() { + true => tensor, + false => tensor.copy(), + }; + + let num_elems = tensor.meta.num_elements() / tensor.meta.shape()[dim]; + + let working_units = num_elems; + let cube_dim = CubeDim::new(&indices.client, working_units); + let cube_count = calculate_cube_count_elemwise(&indices.client, working_units, cube_dim); + + let launch = match is_bool { + true => scatter_kernel::launch_unchecked::, + false => scatter_kernel::launch_unchecked::, + }; + + let (tensor_dtype, indices_dtype) = (tensor.dtype, indices.dtype); + + unsafe { + launch( + &tensor.client.clone(), + cube_count, + cube_dim, + address_type!(tensor, indices, value), + tensor.clone().into_tensor_arg(), + indices.into_tensor_arg(), + value.into_tensor_arg(), + shape_divmod(&tensor), + dim, + [ + dtype_to_storage_type(tensor_dtype), + dtype_to_storage_type(indices_dtype), + ], + ) + } + tensor +} diff --git a/crates/burn-cubecl/src/kernel/index/scatter_nd.rs b/crates/burn-cubecl/src/kernel/index/scatter_nd.rs new file mode 100644 index 0000000..3f3fc7c --- /dev/null +++ b/crates/burn-cubecl/src/kernel/index/scatter_nd.rs @@ -0,0 +1,139 @@ +use crate::{ + CubeRuntime, + kernel::{ + AddOp, AssignOp, BinaryMaxOp, BinaryMinOp, BinaryOp, BinaryOpFamily, MulOp, + utils::{address_type, shape_divmod_range}, + }, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::tensor::IndexingUpdateOp; +use cubecl::std::tensor::layout::linear::LinearView; +use cubecl::{CubeDim, calculate_cube_count_elemwise}; +use cubecl::{prelude::*, std::FastDivmod}; + +/// scatter_nd GPU kernel. +/// +/// Each thread handles one element across all update slices. +/// Work items = num_updates * slice_size. +#[cube(launch_unchecked, address_type = "dynamic")] +fn scatter_nd_kernel( + data: &mut Tensor, + indices: LinearView<'_, I>, + values: &Tensor, + data_slice_shape: Sequence>, + values_shape: Sequence>, + slice_size: usize, + k: usize, + working_units: usize, + #[define(T, I)] _dtypes: [StorageType; 2], +) { + if ABSOLUTE_POS >= working_units { + terminate!(); + } + + let slice_offset = ABSOLUTE_POS % slice_size; + let update_idx = ABSOLUTE_POS / slice_size; + + let idx_base = update_idx * k; + let mut base_offset = 0usize; + for j in 0..k { + let idx_val = usize::cast_from(indices.read(idx_base + j)); + base_offset += idx_val * data.stride(j); + } + + // Decompose slice_offset over data's trailing dims (k..n) + let slice_rank = data_slice_shape.len().comptime(); + let mut data_slice_offset = 0usize; + let mut remainder = slice_offset; + #[unroll] + for i in 0..slice_rank { + let dim = slice_rank - i - 1; + let (rem, coord) = data_slice_shape[dim].div_mod(remainder); + remainder = rem; + data_slice_offset += coord * data.stride(k + dim); + } + + // Decompose slice_offset over values' dims (1..n_v), with update_idx for dim 0 + let val_rank = values_shape.len().comptime(); + let mut val_offset = update_idx * values.stride(0); + let mut remainder_v = slice_offset; + #[unroll] + for i in 0..val_rank { + let dim = val_rank - i - 1; + let (rem, coord) = values_shape[dim].div_mod(remainder_v); + remainder_v = rem; + val_offset += coord * values.stride(1 + dim); + } + + let data_idx = base_offset + data_slice_offset; + let result = Op::BinaryOp::>::execute( + Vector::cast_from(data[data_idx]), + Vector::cast_from(values[val_offset]), + ); + data[data_idx] = result.extract(0); +} + +pub(crate) fn scatter_nd( + tensor: CubeTensor, + indices: CubeTensor, + values: CubeTensor, + reduction: IndexingUpdateOp, +) -> CubeTensor { + // Ensure we can write in-place + let tensor = match tensor.can_mut() && tensor.is_nonoverlapping() { + true => tensor, + false => tensor.copy(), + }; + + let data_shape = &tensor.meta.shape; + let idx_shape = &indices.meta.shape; + let m = idx_shape.num_dims(); + let k = idx_shape[m - 1]; + + // num_updates = product of first M-1 dims of indices + let num_updates: usize = idx_shape.as_slice()[..m - 1].iter().product(); + // slice_size = product of data.shape[K..] + let slice_size: usize = data_shape.as_slice()[k..].iter().product(); + let working_units = num_updates * slice_size; + + let cube_dim = CubeDim::new(&indices.client, working_units); + let cube_count = calculate_cube_count_elemwise(&indices.client, working_units, cube_dim); + + let (tensor_dtype, indices_dtype) = (tensor.dtype, indices.dtype); + + let launch = match reduction { + IndexingUpdateOp::Assign => scatter_nd_kernel::launch_unchecked::, + IndexingUpdateOp::Add => scatter_nd_kernel::launch_unchecked::, + IndexingUpdateOp::Mul => scatter_nd_kernel::launch_unchecked::, + IndexingUpdateOp::Min => scatter_nd_kernel::launch_unchecked::, + IndexingUpdateOp::Max => scatter_nd_kernel::launch_unchecked::, + }; + + let data_slice_shape = shape_divmod_range(&tensor, k..data_shape.num_dims()); + // values dims 1.. (skip the num_updates leading dim) + let values_slice_shape = shape_divmod_range(&values, 1..values.meta.shape.num_dims()); + + unsafe { + launch( + &tensor.client.clone(), + cube_count, + cube_dim, + address_type!(tensor, indices, values), + tensor.clone().into_tensor_arg(), + indices.into_linear_view(), + values.into_tensor_arg(), + data_slice_shape, + values_slice_shape, + slice_size, + k, + working_units, + [ + dtype_to_storage_type(tensor_dtype), + dtype_to_storage_type(indices_dtype), + ], + ) + } + + tensor +} diff --git a/crates/burn-cubecl/src/kernel/index/select.rs b/crates/burn-cubecl/src/kernel/index/select.rs new file mode 100644 index 0000000..f5c0bad --- /dev/null +++ b/crates/burn-cubecl/src/kernel/index/select.rs @@ -0,0 +1,85 @@ +use crate::{CubeRuntime, kernel::utils::address_type, tensor::CubeTensor}; +use crate::{kernel::utils::shape_divmod, ops::numeric::empty_device_dtype}; +use burn_backend::TensorMetadata; +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::{ + CubeDim, calculate_cube_count_elemwise, + std::tensor::layout::linear::{LinearView, LinearViewMut}, +}; +use cubecl::{prelude::*, std::FastDivmod}; + +#[cube(launch_unchecked, address_type = "dynamic")] +fn select_kernel( + input: &Tensor, + indices: LinearView<'_, I>, + mut output: LinearViewMut<'_, T>, + out_shape: Sequence>, + dim: usize, + #[define(T, I)] _dtypes: [StorageType; 2], +) { + if ABSOLUTE_POS >= output.shape() { + terminate!(); + } + + let rank = out_shape.len().comptime(); + + let mut offset = ABSOLUTE_POS; + let mut offset_input = 0; + + #[unroll] + for i in 0..rank { + let i = rank - i - 1; + let (rem, mut offset_local) = out_shape[i].div_mod(offset); + offset = rem; + + if i == dim { + offset_local = usize::cast_from(indices.read(offset_local)); + } + + offset_input += offset_local * input.stride(i); + } + + output.write(ABSOLUTE_POS, input[offset_input]); +} + +pub(crate) fn select( + tensor: CubeTensor, + dim: usize, + indices: CubeTensor, +) -> CubeTensor { + let mut shape_output = tensor.shape(); + shape_output[dim] = indices.meta.shape()[0]; + let total_elem = shape_output.num_elements(); + + let output = empty_device_dtype( + tensor.client.clone(), + tensor.device.clone(), + shape_output, + tensor.dtype, + ); + + let working_units = total_elem; + let cube_dim = CubeDim::new(&indices.client, working_units); + let cube_count = calculate_cube_count_elemwise(&indices.client, working_units, cube_dim); + + let (tensor_dtype, indices_dtype) = (tensor.dtype, indices.dtype); + + unsafe { + select_kernel::launch_unchecked( + &output.client, + cube_count, + cube_dim, + address_type!(tensor, indices, output), + tensor.into_tensor_arg(), + indices.into_linear_view(), + output.clone().into_linear_view(), + shape_divmod(&output), + dim, + [ + dtype_to_storage_type(tensor_dtype), + dtype_to_storage_type(indices_dtype), + ], + ) + }; + output +} diff --git a/crates/burn-cubecl/src/kernel/index/select_assign.rs b/crates/burn-cubecl/src/kernel/index/select_assign.rs new file mode 100644 index 0000000..0c8ee22 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/index/select_assign.rs @@ -0,0 +1,104 @@ +use crate::kernel::{ + AddOp, BinaryOp, BinaryOpFamily, OrOp, + utils::{address_type, shape_divmod}, +}; +use crate::{CubeRuntime, tensor::CubeTensor}; +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::{CubeDim, calculate_cube_count_elemwise, std::tensor::layout::linear::LinearView}; +use cubecl::{prelude::*, std::FastDivmod}; + +/// Uses checked launch mode because user-provided `indices` may contain out-of-bounds values +/// that would cause invalid writes into `tensor`. Checked mode clamps these accesses rather +/// than producing undefined behavior. +#[cube(launch, address_type = "dynamic")] +fn select_assign_kernel( + tensor: &mut Tensor, + indices: LinearView<'_, I>, + value: &Tensor, + value_shape: Sequence>, + working_units: usize, + #[comptime] axis: usize, + #[define(F, I)] _dtypes: [StorageType; 2], +) { + if ABSOLUTE_POS >= working_units { + terminate!(); + } + + let rank = value_shape.len().comptime(); + + let mut offset = ABSOLUTE_POS; + let mut offset_tensor = 0; + let mut offset_value = 0; + + // Calculate offsets and num_elems + #[unroll] + for i in 0..rank { + let i = rank - i - 1; + if i != axis { + let (rem, local_pos) = value_shape[i].div_mod(offset); + offset = rem; + + offset_tensor += local_pos * tensor.stride(i); + offset_value += local_pos * value.stride(i); + } + } + + let strides_tensor_dim = tensor.stride(axis); + let strides_value_dim = value.stride(axis); + + // Main operation + for i in 0..value.shape(axis) { + let index_tensor = usize::cast_from(indices.read(i)) * strides_tensor_dim + offset_tensor; + let index_value = i * strides_value_dim + offset_value; + + let value = Op::BinaryOp::>::execute( + Vector::cast_from(tensor[index_tensor]), + Vector::cast_from(value[index_value]), + ); + write_checked(tensor.as_mut_slice(), index_tensor, F::cast_from(value)); + } +} + +pub(crate) fn select_assign( + tensor: CubeTensor, + dim: usize, + indices: CubeTensor, + value: CubeTensor, + is_bool: bool, +) -> CubeTensor { + let tensor = match tensor.can_mut() && tensor.is_nonoverlapping() { + true => tensor, + false => tensor.copy(), + }; + + let working_units = tensor.meta.num_elements() / tensor.meta.shape()[dim]; + let cube_dim = CubeDim::new(&indices.client, working_units); + let cube_count = calculate_cube_count_elemwise(&indices.client, working_units, cube_dim); + + let launch = match is_bool { + true => select_assign_kernel::launch::, + false => select_assign_kernel::launch::, + }; + + let (tensor_dtype, indices_dtype) = (tensor.dtype, indices.dtype); + + let shape = shape_divmod(&value); + launch( + &tensor.client, + cube_count, + cube_dim, + address_type!(tensor, indices, value), + tensor.clone().into_tensor_arg(), + indices.into_linear_view(), + value.into_tensor_arg(), + shape, + working_units, + dim, + [ + dtype_to_storage_type(tensor_dtype), + dtype_to_storage_type(indices_dtype), + ], + ); + + tensor +} diff --git a/crates/burn-cubecl/src/kernel/index/slice.rs b/crates/burn-cubecl/src/kernel/index/slice.rs new file mode 100644 index 0000000..d4b1ed1 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/index/slice.rs @@ -0,0 +1,249 @@ +use crate::{ + CubeRuntime, + kernel::utils::{address_type, shape_divmod}, + ops::numeric::empty_device_dtype, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{Slice, TensorMetadata}; +use burn_std::{Metadata, SliceOps}; +use cubecl::{ + calculate_cube_count_elemwise, intrinsic, + prelude::*, + std::{FastDivmod, tensor::layout::linear::LinearViewMut}, +}; +use std::ops::Range; + +/// Slice a jit tensor with a set of ranges +pub fn slice(tensor: CubeTensor, indices: &[Range]) -> CubeTensor { + let mut dims = tensor.shape(); + let mut offset_start = 0u64; + let mut offset_end = 0u64; + + for i in 0..indices.len() { + offset_start += (tensor.meta.strides()[i] * indices[i].start) as u64; + offset_end += (tensor.meta.strides()[i] * (dims[i] - indices[i].end)) as u64; + dims[i] = indices[i].end - indices[i].start; + } + + let offset_start = offset_start * tensor.dtype.size() as u64; + let offset_end = offset_end * tensor.dtype.size() as u64; + + let memory_offset_alignment = tensor.client.properties().memory.alignment; + + if offset_start.is_multiple_of(memory_offset_alignment) + && offset_end.is_multiple_of(memory_offset_alignment) + { + CubeTensor::new( + tensor.client.clone(), + tensor + .handle + .clone() + .offset_start(offset_start) + .offset_end(offset_end), + Metadata::new(dims, tensor.meta.strides.clone()), + tensor.device.clone(), + tensor.dtype, + ) + } else { + let output = empty_device_dtype( + tensor.client.clone(), + tensor.device.clone(), + dims, + tensor.dtype, + ); + slice_on_output(tensor, output, indices) + } +} + +#[cube(launch_unchecked, address_type = "dynamic")] +fn slice_kernel( + input: &Tensor, + mut output: LinearViewMut<'_, E>, + out_shape: Sequence>, + indices: Sequence, + #[define(E)] _dtype: StorageType, +) { + if !output.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + let rank = comptime![out_shape.len()]; + let mut offset_output = ABSOLUTE_POS; + let mut offset_input = 0; + + #[unroll] + for i in 0..rank { + // Iterate in reverse to use divmod + let dim = rank - i - 1; + + let range_start = indices[dim]; + let (rem, offset_local) = out_shape[dim].div_mod(offset_output); + offset_output = rem; + + let offset_local = offset_local + range_start; + + offset_input += offset_local * input.stride(dim); + } + + output.write(ABSOLUTE_POS, input[offset_input]); +} + +pub(crate) fn slice_on_output( + tensor: CubeTensor, + output: CubeTensor, + indices: &[Range], +) -> CubeTensor { + let ndims = tensor.meta.num_dims(); + let mut indices_sequence = SequenceArg::::new(); + + for i in 0..ndims { + let start = indices.get(i).map(|index| index.start).unwrap_or(0); + indices_sequence.push(start); + } + + let working_units = output.meta.num_elements(); + let cube_dim = CubeDim::new(&tensor.client, working_units); + let cube_count = calculate_cube_count_elemwise(&tensor.client, working_units, cube_dim); + let dtype = tensor.dtype; + + unsafe { + slice_kernel::launch_unchecked( + &output.client, + cube_count, + cube_dim, + address_type!(tensor, output), + tensor.into_tensor_arg(), + output.clone().into_linear_view(), + shape_divmod(&output), + indices_sequence, + dtype_to_storage_type(dtype), + ) + }; + + output +} + +/// Kernel for slicing with steps +#[cube(launch_unchecked, address_type = "dynamic")] +fn slice_with_steps_kernel( + input: &Tensor, + mut output: LinearViewMut<'_, E>, + out_shape: Sequence>, + starts: Sequence, + ends: Sequence, + steps: Sequence, + #[define(E)] _dtype: StorageType, +) { + if !output.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + let rank = comptime![out_shape.len()]; + let mut output_offset = ABSOLUTE_POS; + let mut input_offset = 0; + + // Calculate the input offset based on output position and slice info + #[unroll] + for i in 0..rank { + // Iterate in reverse to use divmod + let dim = rank - i - 1; + let start = starts[dim]; + let end = ends[dim]; + let step = steps[dim]; + + let (rem, output_idx) = out_shape[dim].div_mod(output_offset); + output_offset = rem; + + let input_idx = if step > 0 { + // Forward stepping + start + output_idx * (step as usize) + } else { + // Backward stepping - start from end-1 + let abs_step = (-step) as usize; + let end_minus_1 = end - 1; + end_minus_1 - output_idx * abs_step + }; + + input_offset += input_idx * input.stride(dim); + } + + output.write(ABSOLUTE_POS, input[input_offset]); +} + +/// Slice a tensor with steps +pub fn slice_with_steps(tensor: CubeTensor, slices: &[Slice]) -> CubeTensor { + // Check if all steps are 1 - if so, use the optimized regular slice + let all_steps_one = slices.iter().all(|info| info.step == 1); + + if all_steps_one { + // Convert Slice to Range for step=1 + let simple_ranges: Vec> = slices + .iter() + .enumerate() + .map(|(i, slice)| slice.to_range(tensor.meta.shape()[i])) + .collect(); + return slice(tensor, &simple_ranges); + } + + // Calculate output shape + let shape_output = tensor.shape().slice(slices).unwrap(); + + // Create output tensor + let output = empty_device_dtype( + tensor.client.clone(), + tensor.device.clone(), + shape_output.clone(), + tensor.dtype, + ); + + // Prepare three separate sequences for kernel + let mut starts = SequenceArg::::new(); + let mut ends = SequenceArg::::new(); + let mut steps = SequenceArg::::new(); + + for (dim, slice) in slices.iter().enumerate() { + let range = slice.to_range(tensor.meta.shape()[dim]); + starts.push(range.start); + ends.push(range.end); + steps.push(slice.step as i32); + } + + // Pad with default values if needed to match tensor dimensions + for dim in slices.len()..tensor.meta.num_dims() { + starts.push(0); + ends.push(tensor.meta.shape[dim]); + steps.push(1); + } + + // Launch kernel + let working_units = shape_output.num_elements(); + let cube_dim = CubeDim::new(&tensor.client, working_units); + let cube_count = calculate_cube_count_elemwise(&tensor.client, working_units, cube_dim); + let dtype = tensor.dtype; + + unsafe { + slice_with_steps_kernel::launch_unchecked( + &output.client, + cube_count, + cube_dim, + address_type!(tensor, output), + tensor.into_tensor_arg(), + output.clone().into_linear_view(), + shape_divmod(&output), + starts, + ends, + steps, + dtype_to_storage_type(dtype), + ); + } + + output +} + +/// This is annoying and we need to find a way to do this automatically at some point +#[allow(unused)] +#[cube] +fn unwrap(value: u32) -> comptime_type!(u32) { + intrinsic!(|_| value.constant().unwrap().as_u32()) +} diff --git a/crates/burn-cubecl/src/kernel/index/slice_assign.rs b/crates/burn-cubecl/src/kernel/index/slice_assign.rs new file mode 100644 index 0000000..49d0fa4 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/index/slice_assign.rs @@ -0,0 +1,259 @@ +use crate::{ + CubeRuntime, + kernel::utils::{address_type, shape_divmod}, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::{ + calculate_cube_count_elemwise, intrinsic, + prelude::*, + std::{FastDivmod, tensor::layout::linear::LinearView}, +}; + +#[cube(launch_unchecked, address_type = "dynamic")] +fn slice_assign_kernel( + input: &mut Tensor>, + value: LinearView<'_, Vector>, + slice_shape: Sequence>, + slice_offsets: Sequence, + #[define(E)] _dtype: StorageType, +) { + if !value.is_in_bounds(ABSOLUTE_POS) { + terminate!() + } + + let rank = comptime!(slice_shape.len()); + + let line_size = input.vector_size(); + let mut offset_remainder = ABSOLUTE_POS * line_size; + let mut offset_input = 0; + + #[allow(clippy::explicit_counter_loop)] + #[unroll] + for i in 0..rank { + let dim = rank - i - 1; + let (rem, offset_local) = slice_shape[dim].div_mod(offset_remainder); + + let range_start = slice_offsets[dim]; + let offset_local_input = offset_local + range_start; + + offset_input += offset_local_input * input.stride(dim); + offset_remainder = rem; + } + + // Value tensor is accessed linearly since it's a LinearView + input[offset_input / line_size] = value.read(ABSOLUTE_POS); +} + +/// Kernel for slice assign with steps +#[cube(launch_unchecked, address_type = "dynamic")] +fn slice_assign_with_steps_kernel( + input: &mut Tensor, + value: LinearView<'_, E>, + value_shape: Sequence>, + starts: Sequence, + ends: Sequence, + steps: Sequence, + #[define(E)] _dtype: StorageType, +) { + if !value.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + let rank = comptime![value_shape.len()]; + let mut value_offset = ABSOLUTE_POS; + let mut input_offset = 0; + + // Calculate the input offset based on value position and slice info + #[unroll] + for i in 0..rank { + // Iterate in reverse to use divmod + let dim = rank - i - 1; + let start = starts[dim]; + let end = ends[dim]; + let step = steps[dim]; + + let (rem, value_idx) = value_shape[dim].div_mod(value_offset); + value_offset = rem; + + let input_idx = if step > 0 { + // Forward stepping + start + value_idx * (step as usize) + } else if step < 0 { + // Backward stepping - start from end-1 + // For negative steps, we iterate backwards through the selected indices + let abs_step = (-step) as usize; + let end_minus_1 = end - 1; + end_minus_1 - value_idx * abs_step + } else { + // step == 0, shouldn't happen + value_idx + }; + + input_offset += input_idx * input.stride(dim); + } + + input[input_offset] = value.read(ABSOLUTE_POS); +} + +pub(crate) fn slice_assign( + tensor: CubeTensor, + indices: &[burn_backend::Slice], + value: CubeTensor, +) -> CubeTensor { + // Check if any slice has non-unit step + let has_non_unit_step = indices.iter().any(|s| s.step != 1 && s.step != 0); + + if has_non_unit_step { + // Use slice_assign_with_steps + return slice_assign_with_steps(tensor, indices, value); + } + + let client = tensor.client.clone(); + let tensor = match tensor.can_mut() && tensor.is_nonoverlapping() { + true => tensor, + false => tensor.copy(), + }; + let ndims = tensor.meta.num_dims(); + + let vector_size = + if tensor.meta.strides()[ndims - 1] == 1 && value.meta.strides()[ndims - 1] == 1 { + let last = indices + .get(ndims - 1) + .cloned() + .unwrap_or(burn_backend::Slice { + start: 0, + end: Some(tensor.meta.shape()[ndims - 1] as isize), + step: 1, + }); + let end = last.end.unwrap_or(tensor.meta.shape()[ndims - 1] as isize); + let shape = (end - last.start) as usize; + let offset = last.start as usize; + client + .io_optimized_vector_sizes(tensor.dtype.size()) + .filter(|&it| { + shape.is_multiple_of(it) + && strides_compatible(tensor.meta.strides(), it) + && strides_compatible(value.meta.strides(), it) + && offset.is_multiple_of(it) + }) + .max() + .unwrap_or(1) + } else { + 1 + }; + + let mut shape = SequenceArg::>::new(); + let mut offsets = SequenceArg::::new(); + + for i in 0..ndims { + let slice = indices.get(i).cloned().unwrap_or(burn_backend::Slice { + start: 0, + end: Some(tensor.meta.shape()[i] as isize), + step: 1, + }); + let start = slice.start as usize; + let end = slice.end.unwrap_or(tensor.meta.shape()[i] as isize); + let length = (end - slice.start) as usize; + + shape.push(length); + offsets.push(start); + } + + let working_units = value.meta.num_elements() / vector_size; + let cube_dim = CubeDim::new(&tensor.client, working_units); + let cube_count = calculate_cube_count_elemwise(&tensor.client, working_units, cube_dim); + + unsafe { + slice_assign_kernel::launch_unchecked( + &tensor.client, + cube_count, + cube_dim, + address_type!(tensor, value), + vector_size, + tensor.clone().into_tensor_arg(), + value.into_linear_view(), + shape, + offsets, + dtype_to_storage_type(tensor.dtype), + ) + }; + + tensor +} + +/// Slice assign with steps support +/// +/// This function handles slice assignment with arbitrary step values, including negative steps. +/// It follows NumPy/PyTorch semantics where values[i] is assigned to selected_indices[i]. +/// +/// For example, with s![0..6;-1] which selects indices [5,4,3,2,1,0]: +/// - values[0] goes to index 5 +/// - values[1] goes to index 4 +/// - etc. +pub(crate) fn slice_assign_with_steps( + tensor: CubeTensor, + slices: &[burn_backend::Slice], + value: CubeTensor, +) -> CubeTensor { + let tensor = match tensor.can_mut() && tensor.is_nonoverlapping() { + true => tensor, + false => tensor.copy(), + }; + + // Prepare sequences for kernel + let mut starts = SequenceArg::::new(); + let mut ends = SequenceArg::::new(); + let mut steps = SequenceArg::::new(); + + for (dim, slice) in slices.iter().enumerate() { + let range = slice.to_range(tensor.meta.shape()[dim]); + starts.push(range.start); + ends.push(range.end); + steps.push(slice.step as i32); + } + + // Pad with default values if needed to match tensor dimensions + for dim in slices.len()..tensor.meta.num_dims() { + starts.push(0); + ends.push(tensor.meta.shape[dim]); + steps.push(1); + } + + // Launch kernel + let working_units = value.meta.num_elements(); + let cube_dim = CubeDim::new(&tensor.client, working_units); + let cube_count = calculate_cube_count_elemwise(&tensor.client, working_units, cube_dim); + + let shape = shape_divmod(&value); + unsafe { + slice_assign_with_steps_kernel::launch_unchecked( + &tensor.client, + cube_count, + cube_dim, + address_type!(tensor, value), + tensor.clone().into_tensor_arg(), + value.into_linear_view(), + shape, + starts, + ends, + steps, + dtype_to_storage_type(tensor.dtype), + ); + } + + tensor +} + +fn strides_compatible(strides: &[usize], vec: usize) -> bool { + strides + .iter() + .all(|stride| *stride % vec == 0 || *stride == 1) +} + +/// Helper function for unwrap +#[allow(unused)] +#[cube] +fn unwrap(value: u32) -> comptime_type!(u32) { + intrinsic!(|_| value.constant().unwrap().as_u32()) +} diff --git a/crates/burn-cubecl/src/kernel/interpolate/base.rs b/crates/burn-cubecl/src/kernel/interpolate/base.rs new file mode 100644 index 0000000..97b63c6 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/interpolate/base.rs @@ -0,0 +1,168 @@ +use crate::{ + CubeRuntime, + kernel::{interpolate::interpolate_autotune, into_contiguous}, + ops::{numeric::empty_device_dtype, permute_nchw_to_nhwc, permute_nhwc_to_nchw}, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{Shape, TensorMetadata, ops::InterpolateMode, ops::InterpolateOptions}; +#[cfg(not(feature = "autotune"))] +use cubek::interpolate::definition::TileSize; +use cubek::interpolate::{ + definition::{ + InterpolateError, InterpolateMode as CubekInterpolateMode, + InterpolateOptions as CubekInterpolateOptions, NearestMode as CubekNearestMode, + }, + interpolate as cubek_interpolate, interpolate_backward as cubek_interpolate_backward, + launch::InterpolateStrategy as CubekInterpolateStrategy, + routines::{ + BlueprintStrategy, GlobalMemoryRoutine, GlobalMemoryStrategy, SharedMemoryRoutine, + SharedMemoryStrategy, + }, +}; + +#[derive(Debug)] +/// Strategy used to select which interpolate implementation to run. +pub enum InterpolateStrategy { + /// Default interpolate strategy. + GlobalMemory(GlobalMemoryStrategy), + + /// Use shared memory for caching tiles of the input and output. + SharedMemory(SharedMemoryStrategy), + + /// Automatically benchmark and select the best strategy at runtime. + #[cfg(feature = "autotune")] + Autotune, +} + +impl Default for InterpolateStrategy { + fn default() -> Self { + // if autotune is enabled, default to autotune + #[cfg(feature = "autotune")] + return InterpolateStrategy::Autotune; + + // if autotune is disabled, default to global memory with a 16x16 tile size + #[cfg(not(feature = "autotune"))] + InterpolateStrategy::GlobalMemory(GlobalMemoryStrategy { + tile_size: TileSize::new(16, 16), + }) + } +} + +/// Interpolate operation +/// +/// Supports nearest, bilinear, bicubic and lanczos3 modes +pub fn interpolate( + input: CubeTensor, + output_size: [usize; 2], + options: InterpolateOptions, + strategy: InterpolateStrategy, +) -> Result, InterpolateError> { + match strategy { + InterpolateStrategy::GlobalMemory(strategy) => execute_interpolate( + input, + output_size, + options, + CubekInterpolateStrategy::GlobalMemoryStrategy( + BlueprintStrategy::::Inferred(strategy), + ), + ), + InterpolateStrategy::SharedMemory(strategy) => execute_interpolate( + input, + output_size, + options, + CubekInterpolateStrategy::SharedMemoryStrategy( + BlueprintStrategy::::Inferred(strategy), + ), + ), + #[cfg(feature = "autotune")] + InterpolateStrategy::Autotune => Ok(interpolate_autotune(input, output_size, options)), + } +} + +/// Execute the given interpolate strategy without autotuning. This is used by the autotune implementation to run each candidate strategy. +pub fn execute_interpolate( + input: CubeTensor, + output_size: [usize; 2], + options: InterpolateOptions, + strategy: CubekInterpolateStrategy, +) -> Result, InterpolateError> { + let [batch_size, channels, _, _] = input.meta.shape().dims(); + let [out_height, out_width] = output_size; + + let input = into_contiguous(permute_nchw_to_nhwc(input)); + + let shape_out = Shape::new([batch_size, out_height, out_width, channels]); + let output = empty_device_dtype( + input.client.clone(), + input.device.clone(), + shape_out, + input.dtype, + ); + + cubek_interpolate( + &input.client.clone(), + input.clone().binding(), + output.clone().binding(), + map_options(options.clone()), + strategy, + dtype_to_storage_type(input.dtype), + )?; + + Ok(permute_nhwc_to_nchw(output)) +} + +/// Backward interpolate operation +/// +/// Note: only nearest mode is supported +pub fn interpolate_backward( + input: CubeTensor, + out_grad: CubeTensor, + _output_size: [usize; 2], + options: InterpolateOptions, +) -> CubeTensor { + let input = permute_nchw_to_nhwc(input); + let out_grad = permute_nchw_to_nhwc(out_grad); + + let output_shape = input.shape(); + let output = empty_device_dtype( + input.client.clone(), + input.device.clone(), + output_shape, + input.dtype, + ); + + cubek_interpolate_backward( + &input.client.clone(), + input.clone().binding(), + out_grad.binding(), + output.clone().binding(), + map_options(options.clone()), + dtype_to_storage_type(input.dtype), + ) + .unwrap_or_else(|e| { + panic!( + "interpolate_backward kernel failed (device={0:?}, dtype={1:?}, options={2:?}): {3}", + input.device, input.dtype, options, e + ) + }); + + permute_nhwc_to_nchw(output) +} + +pub(crate) fn map_mode(mode: InterpolateMode) -> CubekInterpolateMode { + match mode { + InterpolateMode::Nearest => CubekInterpolateMode::Nearest(CubekNearestMode::Floor), + InterpolateMode::NearestExact => CubekInterpolateMode::Nearest(CubekNearestMode::Exact), + InterpolateMode::Bilinear => CubekInterpolateMode::Bilinear, + InterpolateMode::Bicubic => CubekInterpolateMode::Bicubic, + InterpolateMode::Lanczos3 => CubekInterpolateMode::Lanczos3, + } +} + +pub(crate) fn map_options(options: InterpolateOptions) -> CubekInterpolateOptions { + CubekInterpolateOptions { + mode: map_mode(options.mode), + align_corners: options.align_corners, + } +} diff --git a/crates/burn-cubecl/src/kernel/interpolate/mod.rs b/crates/burn-cubecl/src/kernel/interpolate/mod.rs new file mode 100644 index 0000000..8ff38a9 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/interpolate/mod.rs @@ -0,0 +1,5 @@ +mod base; +mod tune; + +pub use base::*; +pub use tune::*; diff --git a/crates/burn-cubecl/src/kernel/interpolate/tune.rs b/crates/burn-cubecl/src/kernel/interpolate/tune.rs new file mode 100644 index 0000000..8162a56 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/interpolate/tune.rs @@ -0,0 +1,136 @@ +use crate::{ + CubeRuntime, CubeTuneId, kernel::interpolate::execute_interpolate, + kernel::interpolate::map_mode, tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_elem_type; +use burn_backend::ops::InterpolateOptions; +use cubecl::tune::{LocalTuner, Tunable, TunableSet, TuneGroup, local_tuner}; +use cubek::interpolate::{ + definition::TileSize, + launch::{InterpolateAutotuneKey, InterpolateStrategy}, + routines::{ + BlueprintStrategy, GlobalMemoryRoutine, GlobalMemoryStrategy, SharedMemoryRoutine, + SharedMemoryStrategy, + }, +}; + +/// Interpolate operation with autotuning. This benchmarks multiple strategies and selects the best one at runtime. +pub fn interpolate_autotune( + input: CubeTensor, + output_size: [usize; 2], + options: InterpolateOptions, +) -> CubeTensor { + let client = input.client.clone(); + + static TUNER: LocalTuner = local_tuner!(); + + let tunables = TUNER.init(|| { + const PRIORITY: i8 = 0; + + let global_memory = + TuneGroup::::new("global_memory", |_key| PRIORITY); + let shared_memory = + TuneGroup::::new("shared_memory", |_key| PRIORITY); + + let mut set = TunableSet::new(create_key::, input_gen::); + + let tile_sizes: [TileSize; 16] = [ + // Square shapes + TileSize::new(8, 8), + TileSize::new(16, 16), + TileSize::new(32, 32), + // Rectangular shapes + TileSize::new(8, 16), + TileSize::new(16, 8), + TileSize::new(16, 32), + TileSize::new(32, 16), + // Flat horizontal shapes + TileSize::new(1, 64), + TileSize::new(1, 128), + TileSize::new(1, 256), + TileSize::new(2, 128), + TileSize::new(1, 512), + TileSize::new(1, 1024), + // Flat vertical shapes + TileSize::new(64, 1), + TileSize::new(128, 1), + TileSize::new(256, 1), + ]; + + for tile_size in tile_sizes { + let name = format!( + "global_memory_tile_size_{}_{}", + tile_size.height(), + tile_size.width() + ); + set = set.with( + Tunable::new(&name, move |(input, output_size, options)| { + execute_interpolate::( + input, + output_size, + options, + InterpolateStrategy::GlobalMemoryStrategy(BlueprintStrategy::< + GlobalMemoryRoutine, + >::Inferred( + GlobalMemoryStrategy { tile_size }, + )), + ) + }) + .group(&global_memory, |_key| PRIORITY), + ); + + let name = format!( + "shared_memory_tile_size_{}_{}", + tile_size.height(), + tile_size.width() + ); + set = set.with( + Tunable::new(&name, move |(input, output_size, options)| { + execute_interpolate::( + input, + output_size, + options, + InterpolateStrategy::SharedMemoryStrategy(BlueprintStrategy::< + SharedMemoryRoutine, + >::Inferred( + SharedMemoryStrategy { tile_size }, + )), + ) + }) + .group(&shared_memory, |_key| PRIORITY), + ); + } + + set + }); + + TUNER.execute( + &CubeTuneId::new(&client, &input.device), + &client, + tunables, + (input, output_size, options), + ) +} + +fn create_key( + (input, output_size, options): &(CubeTensor, [usize; 2], InterpolateOptions), +) -> InterpolateAutotuneKey { + let elem_input = dtype_to_elem_type(input.dtype); + let elem_output = dtype_to_elem_type(input.dtype); + let mode = map_mode(options.mode.clone()); + + InterpolateAutotuneKey::generate( + elem_input, + elem_output, + mode, + input.meta.shape(), + output_size, + ) +} + +fn input_gen( + _key: &InterpolateAutotuneKey, + (input, output_size, options): &(CubeTensor, [usize; 2], InterpolateOptions), +) -> (CubeTensor, [usize; 2], InterpolateOptions) { + (input.clone(), *output_size, options.clone()) +} diff --git a/crates/burn-cubecl/src/kernel/mask/base.rs b/crates/burn-cubecl/src/kernel/mask/base.rs new file mode 100644 index 0000000..9c7057e --- /dev/null +++ b/crates/burn-cubecl/src/kernel/mask/base.rs @@ -0,0 +1,39 @@ +use burn_backend::DType; +use cubecl::prelude::InputScalar; + +use super::{MaskFillStrategy, mask_where::MaskWhereStrategy}; +use crate::{CubeRuntime, tensor::CubeTensor}; + +/// Execute the mask fill kernel. +pub(crate) fn mask_fill_auto( + tensor: CubeTensor, + mask: CubeTensor, + value: InputScalar, + dtype_bool: DType, +) -> CubeTensor { + let strategy = if tensor.can_mut() && tensor.is_nonoverlapping() { + MaskFillStrategy::Inplace + } else { + MaskFillStrategy::Readonly + }; + + super::mask_fill(tensor, mask, value, strategy, dtype_bool) +} + +/// Execute the mask where kernel. +pub(crate) fn mask_where_auto( + tensor: CubeTensor, + mask: CubeTensor, + value: CubeTensor, + dtype_bool: DType, +) -> CubeTensor { + let strategy = if tensor.can_mut_broadcast(&value) { + MaskWhereStrategy::InplaceLhs + } else if value.can_mut_broadcast(&tensor) { + MaskWhereStrategy::InplaceRhs + } else { + MaskWhereStrategy::Readonly + }; + + super::mask_where(tensor, mask, value, strategy, dtype_bool) +} diff --git a/crates/burn-cubecl/src/kernel/mask/mask_fill.rs b/crates/burn-cubecl/src/kernel/mask/mask_fill.rs new file mode 100644 index 0000000..901c8b0 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/mask/mask_fill.rs @@ -0,0 +1,99 @@ +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{DType, TensorMetadata}; +use cubecl::{ + calculate_cube_count_elemwise, + prelude::*, + std::tensor::layout::linear::{LinearView, LinearViewMut}, +}; + +use crate::{ + CubeRuntime, + kernel::utils::address_type, + ops::{max_vector_size_many, numeric::empty_device_dtype}, + tensor::CubeTensor, +}; + +#[cube(launch_unchecked, address_type = "dynamic")] +fn mask_fill_kernel( + input: LinearView<'_, Vector>, + mask: LinearView<'_, Vector>, + mut output: LinearViewMut<'_, Vector>, + value: InputScalar, + #[define(T, B)] _dtypes: [StorageType; 2], +) { + if !output.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + let mask = Vector::cast_from(mask.read(ABSOLUTE_POS)); + let input = input.read(ABSOLUTE_POS); + let value = Vector::new(value.get::()); + + output.write(ABSOLUTE_POS, select_many(mask, value, input)); +} + +#[derive(Clone, Copy, Debug)] +/// Define how to run the mask fill kernel. +/// +/// # Notes +/// +/// All assertions should be done before choosing the strategy. +pub enum MaskFillStrategy { + /// Don't mutate any input. + Readonly, + /// Reuse the input tensor inplace. + Inplace, +} + +/// Execute the mask fill kernel with the given strategy. +pub fn mask_fill( + input: CubeTensor, + mask: CubeTensor, + value: InputScalar, + strategy: MaskFillStrategy, + dtype_bool: DType, +) -> CubeTensor { + let ndims = input.meta.num_dims(); + let output = match strategy { + MaskFillStrategy::Readonly => empty_device_dtype( + input.client.clone(), + input.device.clone(), + input.shape(), + input.dtype, + ), + MaskFillStrategy::Inplace => input.clone(), + }; + + let vector_size = max_vector_size_many(&[&input, &mask], ndims - 1); + let working_units = input.meta.num_elements() / vector_size as usize; + let cube_dim = CubeDim::new(&input.client, working_units); + let cube_count = calculate_cube_count_elemwise(&input.client, working_units, cube_dim); + + let out_arg = match strategy { + MaskFillStrategy::Readonly => output.clone().into_linear_view(), + MaskFillStrategy::Inplace => output.as_linear_view_alias(0), + }; + + let at = address_type!(input, mask, output); + let mask = mask.into_linear_view_like(&input); + + unsafe { + mask_fill_kernel::launch_unchecked( + &output.client, + cube_count, + cube_dim, + at, + vector_size, + input.into_linear_view(), + mask, + out_arg, + value, + [ + dtype_to_storage_type(output.dtype), + dtype_to_storage_type(dtype_bool), + ], + ); + } + + output +} diff --git a/crates/burn-cubecl/src/kernel/mask/mask_where.rs b/crates/burn-cubecl/src/kernel/mask/mask_where.rs new file mode 100644 index 0000000..a9053c1 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/mask/mask_where.rs @@ -0,0 +1,104 @@ +use burn_backend::DType; +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::{ + calculate_cube_count_elemwise, + prelude::*, + std::tensor::layout::linear::{LinearView, LinearViewMut}, +}; + +use crate::{ + CubeRuntime, + kernel::utils::{address_type, broadcast_shape}, + ops::{max_vector_size_many, numeric::empty_device_dtype}, + tensor::CubeTensor, +}; + +#[cube(launch, address_type = "dynamic")] +fn mask_where_kernel( + input: LinearView<'_, Vector>, + value: LinearView<'_, Vector>, + mask: LinearView<'_, Vector>, + mut output: LinearViewMut<'_, Vector>, + #[define(T, B)] _dtypes: [StorageType; 2], +) { + let pos = ABSOLUTE_POS; + if !output.is_in_bounds(pos) { + terminate!(); + } + + output.write( + pos, + select_many( + Vector::cast_from(mask.read(pos)), + value.read(pos), + input.read(pos), + ), + ); +} + +#[derive(Clone, Copy, Debug)] +/// Define how to run the mask where kernel. +/// +/// # Notes +/// +/// All assertions should be done before choosing the strategy. +pub enum MaskWhereStrategy { + /// Don't mutate any input. + Readonly, + /// Reuse the lhs tensor inplace. + InplaceLhs, + /// Reuse the rhs tensor inplace. + InplaceRhs, +} + +/// Execute the mask where kernel with the given strategy. +pub fn mask_where( + input: CubeTensor, + mask: CubeTensor, + value: CubeTensor, + strategy: MaskWhereStrategy, + dtype_bool: DType, +) -> CubeTensor { + let vector_size = max_vector_size_many(&[&input, &mask, &value], input.meta.num_dims() - 1); + + let working_units = input.meta.num_elements() / vector_size as usize; + let cube_dim = CubeDim::new(&input.client, working_units); + let cube_count = calculate_cube_count_elemwise(&input.client, working_units, cube_dim); + + let out_shape = broadcast_shape(&[&input, &mask, &value]); + + let output = match strategy { + MaskWhereStrategy::Readonly => empty_device_dtype( + input.client.clone(), + input.device.clone(), + out_shape, + input.dtype, + ), + MaskWhereStrategy::InplaceLhs => input.clone(), + MaskWhereStrategy::InplaceRhs => value.clone(), + }; + + let out = match strategy { + MaskWhereStrategy::Readonly => output.clone().into_linear_view(), + MaskWhereStrategy::InplaceLhs => output.as_linear_view_alias(0), + MaskWhereStrategy::InplaceRhs => output.as_linear_view_alias(1), + }; + + mask_where_kernel::launch( + &output.client, + cube_count, + cube_dim, + address_type!(input, value, mask, output), + vector_size, + input.into_linear_view_like(&output), + value.into_linear_view_like(&output), + mask.into_linear_view_like(&output), + out, + [ + dtype_to_storage_type(output.dtype), + dtype_to_storage_type(dtype_bool), + ], + ); + + output +} diff --git a/crates/burn-cubecl/src/kernel/mask/mod.rs b/crates/burn-cubecl/src/kernel/mask/mod.rs new file mode 100644 index 0000000..0044101 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/mask/mod.rs @@ -0,0 +1,8 @@ +mod base; +mod mask_fill; +mod mask_where; + +pub(crate) use base::*; + +pub use mask_fill::*; +pub use mask_where::*; diff --git a/crates/burn-cubecl/src/kernel/matmul/base.rs b/crates/burn-cubecl/src/kernel/matmul/base.rs new file mode 100644 index 0000000..01e65e1 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/matmul/base.rs @@ -0,0 +1,189 @@ +use super::init_matmul_output; +use crate::{CubeRuntime, kernel::quantization::dequantize, tensor::CubeTensor}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{DType, TensorMetadata}; +use burn_std::{MatmulTransformAnalysis, MatmulTransformPolicy, QuantLevel}; +use cubek::{ + matmul::{ + definition::{MatmulElems, MatmulGlobalElems, MatmulSetupError}, + strategy::Strategy, + }, + std::InputBinding, +}; + +#[cfg(feature = "autotune")] +use super::matmul_autotune; + +/// The strategy to be used when launching a matmul kernel. +pub enum MatmulStrategy { + #[cfg(feature = "autotune")] + /// Using autotune to choose the best kernel based on runtime information. + Autotune, + /// Cube implementation of matmul. + Cube, +} + +impl Default for MatmulStrategy { + fn default() -> Self { + // if autotune is enabled, default to autotune + #[cfg(feature = "autotune")] + return MatmulStrategy::Autotune; + + #[cfg(not(feature = "autotune"))] + MatmulStrategy::Cube + } +} + +/// Launch a matmul kernel using the given strategy. +pub fn matmul( + mut lhs: CubeTensor, + rhs: CubeTensor, + out: Option>, + strategy: MatmulStrategy, + out_dtype: DType, +) -> Result, MatmulSetupError> { + let out = out.unwrap_or_else(|| init_matmul_output(&lhs, &rhs, out_dtype)); + + // A broadcast-rhs batched matmul that would tile poorly is folded into a + // single matmul: `[.., b, m, k] @ [.., 1, k, n]` runs as `[.., 1, b*m, k]` + // instead of `b` matmuls that each re-read the whole rhs. Pure metadata — + // the launch operands share the handles, the returned tensor keeps the + // broadcast shape. + let mut out_launch = out.clone(); + if lhs.qparams.is_none() { + let analysis = MatmulTransformAnalysis::from_metadata(&lhs.meta, &rhs.meta, &out.meta); + let action = MatmulTransformPolicy::default().action(&analysis); + action.apply(&mut lhs.meta); + action.apply(&mut out_launch.meta); + } + + match strategy { + MatmulStrategy::Cube => { + launch_matmul(&Default::default(), lhs, rhs, out_launch)?; + Ok(out) + } + #[cfg(feature = "autotune")] + MatmulStrategy::Autotune => { + matmul_autotune(lhs, rhs, Some(out_launch), out_dtype); + Ok(out) + } + } +} + +pub(crate) fn launch_matmul_naive( + strategy: &Strategy, + mut lhs: CubeTensor, + mut rhs: CubeTensor, + out: CubeTensor, +) -> Result<(), MatmulSetupError> { + // Naive has very specific layout requirements for block scaled tensors, so we need to manually + // dequantize if it fails to launch normally. This is because naive is assumed to always work. + if lhs.qparams.is_some() || rhs.qparams.is_some() { + match launch_matmul(strategy, lhs.clone(), rhs.clone(), out.clone()) { + Err(_) => { + if lhs.qparams.is_some() { + lhs = dequantize(lhs, out.dtype); + } + if rhs.qparams.is_some() { + rhs = dequantize(rhs, out.dtype); + } + launch_matmul(strategy, lhs, rhs, out) + } + Ok(_) => Ok(()), + } + } else { + launch_matmul(strategy, lhs, rhs, out) + } +} + +pub(crate) fn launch_matmul( + strategy: &Strategy, + lhs: CubeTensor, + mut rhs: CubeTensor, + out: CubeTensor, +) -> Result<(), MatmulSetupError> { + let client = &out.client; + + let lhs_quant_handles = lhs.quantized_handles(); + let out_dtype: DType = out.dtype; + + let (lhs_dtype, lhs_handle) = match lhs_quant_handles { + None => { + let lhs_dtype = lhs.dtype; + ( + lhs_dtype, + InputBinding::new(lhs.binding(), dtype_to_storage_type(lhs_dtype)), + ) + } + Some((data, scale)) => { + let scheme = lhs.scheme(); + let data_dtype = data.dtype; + let scale_dtype = scale.dtype; + ( + out_dtype, + InputBinding::quantized( + data.binding(), + scale.binding(), + lhs.meta.shape().clone(), + scheme, + dtype_to_storage_type(data_dtype), + dtype_to_storage_type(scale_dtype), + ), + ) + } + }; + + let rhs_quant_handles = rhs.quantized_handles(); + + let (rhs_dtype, rhs_handle) = match rhs_quant_handles { + None => ( + lhs_dtype, + InputBinding::new(rhs.binding(), dtype_to_storage_type(lhs_dtype)), + ), + Some((data, scale)) => { + // Extremely hacky fix to ensure naive can run in every case + if matches!(strategy, Strategy::Naive) + && matches!(rhs.scheme().level, QuantLevel::Block(_)) + { + rhs = dequantize(rhs.clone(), lhs_dtype); + let rhs_dtype = rhs.dtype; + ( + lhs_dtype, + InputBinding::new(rhs.binding(), dtype_to_storage_type(rhs_dtype)), + ) + } else { + let scheme = rhs.scheme(); + let data_dtype = data.dtype; + let scale_dtype = scale.dtype; + ( + out_dtype, + InputBinding::quantized( + data.binding(), + scale.binding(), + rhs.meta.shape().clone(), + scheme, + dtype_to_storage_type(data_dtype), + dtype_to_storage_type(scale_dtype), + ), + ) + } + } + }; + + let mut dtypes = MatmulElems::from_globals(&MatmulGlobalElems { + lhs: dtype_to_storage_type(lhs_dtype), + rhs: dtype_to_storage_type(rhs_dtype), + out: dtype_to_storage_type(out_dtype), + }); + + cubek::matmul::launch::launch_ref( + strategy, + client, + lhs_handle, + rhs_handle, + out.clone().binding(), + &mut dtypes, + )?; + + Ok(()) +} diff --git a/crates/burn-cubecl/src/kernel/matmul/mod.rs b/crates/burn-cubecl/src/kernel/matmul/mod.rs new file mode 100644 index 0000000..86df664 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/matmul/mod.rs @@ -0,0 +1,10 @@ +mod base; +mod tune; + +/// Contains utilities for matmul operation +pub mod utils; + +pub use base::*; +#[cfg(feature = "autotune")] +pub use tune::*; +pub use utils::*; diff --git a/crates/burn-cubecl/src/kernel/matmul/tune/base.rs b/crates/burn-cubecl/src/kernel/matmul/tune/base.rs new file mode 100644 index 0000000..1ae9e42 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/matmul/tune/base.rs @@ -0,0 +1,520 @@ +use crate::{ + CubeRuntime, CubeTuneId, + kernel::matmul::{launch_matmul, launch_matmul_naive, utils::init_matmul_output}, + tensor::CubeTensor, +}; +use burn_backend::DType; +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::{ + std::tensor::MatrixBatchLayout, + tune::{LocalTuner, Tunable, TunableSet, TuneGroup, local_tuner}, +}; +use cubek::matmul::{ + components::tile::TileMatmulKind, + definition::MatmulKind, + routines::{ + BlueprintStrategy, TileSizeSelection, + batch::{ + double_buffering::DoubleBufferingArgs, double_unit::DoubleUnitSelectionArgs, + ordered_double_buffering::OrderedSelectionArgs, simple::SimpleArgs, + simple_unit::SimpleUnitSelectionArgs, + }, + cpu_gemm::CpuGemmStrategy, + gemm::GemmStrategy, + }, + strategy::{MatmulAutotuneKey, MatmulGlobalScale, Strategy, should_tune_double_buffering}, +}; + +fn matmul_input_gen( + _key: &MatmulAutotuneKey, + (lhs, rhs, out): &(CubeTensor, CubeTensor, CubeTensor), +) -> (CubeTensor, CubeTensor, CubeTensor) { + (lhs.clone(), rhs.clone(), out.copy()) +} + +/// Executes autotune on matmul operations +pub fn matmul_autotune( + lhs: CubeTensor, + rhs: CubeTensor, + out: Option>, + out_dtype: DType, +) -> CubeTensor { + let output = out.unwrap_or_else(|| init_matmul_output(&lhs, &rhs, out_dtype)); + + // Short-circuit: zero-sized matmul produces a zero-sized output. + if lhs.meta.shape().iter().any(|&d| d == 0) || rhs.meta.shape().iter().any(|&d| d == 0) { + return output; + } + + let client = lhs.client.clone(); + let num_cpu_cores = client.properties().hardware.num_cpu_cores; + + static TUNER: LocalTuner = local_tuner!(); + + let tunables = TUNER.init(move || { + const PRIORITY_MAX: i8 = 3; + const PRIORITY_HIGH: i8 = 2; + const PRIORITY_MEDIUM: i8 = 1; + const PRIORITY_MIN: i8 = 0; + const PRIORITY_NEVER: i8 = -1; + + let accelerated = TuneGroup::::new("accelerated", |key| { + if matches!(key.analysis.kind, MatmulKind::General) { + match key.analysis.scale_global { + MatmulGlobalScale::Large => PRIORITY_MAX, + _ => PRIORITY_HIGH, + } + + // In some case when a relayout can be fused (no call to into_contiguous) it's better + // to use accelerated matmul. + // + // TODO: Actually implement good gemv with fused relayout. + } else if matches!(key.analysis.kind, MatmulKind::MatVec | MatmulKind::VecMat) { + PRIORITY_MAX + } else { + PRIORITY_MEDIUM + } + }); + + let unit = TuneGroup::::new("unit", |key| { + if !matches!(key.analysis.kind, MatmulKind::General) + || matches!(key.analysis.scale_global, MatmulGlobalScale::Small) + { + PRIORITY_HIGH + } else { + PRIORITY_MEDIUM + } + }); + + let tma = TuneGroup::::new("tma", |key| { + // Zero-sized matmuls cannot use TMA kernels. + if key.definition.m == 0 || key.definition.n == 0 || key.definition.k == 0 { + return PRIORITY_NEVER; + } + + // For large matmul, we set the max priority to TMA kernels, higher than any other + // matmuls, since they are the best kernels no matter what. + // + // But only when all axis are large. + let max_axis = usize::max(key.definition.m, key.definition.n); + let max_axis = usize::max(key.definition.k, max_axis); + + let min_axis = usize::min(key.definition.m, key.definition.n); + let min_axis = usize::min(key.definition.k, min_axis); + + let skewed_factor = max_axis / min_axis; + + let priority_max = if matches!(key.analysis.kind, MatmulKind::General) + && matches!(key.analysis.scale_global, MatmulGlobalScale::Large) + && skewed_factor < 4 + { + PRIORITY_MAX + } else { + PRIORITY_HIGH + }; + + if key.definition.lhs_stride_factor >= 4 && key.definition.rhs_stride_factor >= 4 { + priority_max + } else { + PRIORITY_NEVER + } + }); + + let gemv = TuneGroup::::new("gemv", move |key| { + if num_cpu_cores.is_some() { + return PRIORITY_MAX; + } + + if matches!(key.analysis.kind, MatmulKind::MatVec) { + // LHS is the matrix + match key.definition.matrix_layout_lhs { + MatrixBatchLayout::Contiguous => PRIORITY_MAX, + MatrixBatchLayout::MildlyPermuted { transposed, .. } => { + // We don't yet have algo which are good for col major matvec. + if transposed { + PRIORITY_HIGH + } else { + PRIORITY_MAX + } + } + // Every algo will need to relayout, in this case, we should take the optimal + // kernel with a gemv. + MatrixBatchLayout::HighlyPermuted => PRIORITY_MAX, + } + } else if matches!(key.analysis.kind, MatmulKind::VecMat) { + // RHS is the matrix + match key.definition.matrix_layout_rhs { + // We don't have good algos for row major vecmat. + MatrixBatchLayout::Contiguous => PRIORITY_HIGH, + MatrixBatchLayout::MildlyPermuted { transposed, .. } => { + // Best algo is col major vec mat. + if transposed { + PRIORITY_MAX + } else { + PRIORITY_HIGH + } + } + // TODO: Actually do the correct relayout here. + // + // Every algo will need to relayout, in this case, we should take the optimal + // kernel with a gemv. + MatrixBatchLayout::HighlyPermuted => PRIORITY_HIGH, + } + } else { + PRIORITY_NEVER + } + }); + + // CPU-only group + let cpu = + TuneGroup::::new("cpu", move |_key| match num_cpu_cores.is_some() { + true => PRIORITY_MAX, + false => PRIORITY_NEVER, + }); + + fn double_buffering_priority(key: &MatmulAutotuneKey, max: i8, min: i8) -> i8 { + if should_tune_double_buffering(false, key) { + max + } else { + min + } + } + + let mut set = TunableSet::new(create_key::, matmul_input_gen::); + + // First entry should always work, since it is considered the fallback. + set = set.with( + Tunable::new("matmul_naive", |(lhs, rhs, out)| { + launch_matmul_naive::(&Strategy::Naive, lhs, rhs, out) + .map_err(|err| std::format!("{err:?}")) + }) + .group(&unit, |key| { + if matches!(key.analysis.kind, MatmulKind::InnerProduct) { + PRIORITY_MAX + } else if matches!(key.analysis.scale_global, MatmulGlobalScale::Small) { + PRIORITY_HIGH + } else { + PRIORITY_MIN + } + }), + ); + + // Matrix Vector multiplication kernels. + for (strategy, double_buf) in [ + ( + Strategy::DoubleVecMat(BlueprintStrategy::Inferred(().into())), + true, + ), + ( + Strategy::SimpleVecMat(BlueprintStrategy::Inferred(().into())), + false, + ), + ( + Strategy::Gemm(BlueprintStrategy::Inferred(Default::default())), + false, + ), + ( + Strategy::GemvUnitPerpendicular(BlueprintStrategy::Inferred(Default::default())), + false, + ), + ] { + set = set.with( + Tunable::new(&strategy.to_string(), move |(lhs, rhs, out)| { + launch_matmul::(&strategy, lhs, rhs, out) + .map_err(|err| std::format!("{err:?}")) + }) + .group(&gemv, move |key| match double_buf { + false => PRIORITY_MAX, + true => double_buffering_priority(key, PRIORITY_MAX, PRIORITY_HIGH), + }), + ); + } + + // Unit matmuls + for tile_size in [ + TileSizeSelection::MaxTileSize, + TileSizeSelection::MinTileSize, + ] { + for (strategy, double_buf) in [ + ( + Strategy::SimpleUnit(BlueprintStrategy::Inferred(SimpleUnitSelectionArgs { + tile_size, + })), + false, + ), + ( + Strategy::DoubleUnit(BlueprintStrategy::Inferred(DoubleUnitSelectionArgs { + tile_size, + })), + true, + ), + ] { + set = set.with( + Tunable::new(&strategy.to_string(), move |(lhs, rhs, out)| { + launch_matmul::(&strategy, lhs, rhs, out) + .map_err(|err| format!("{err:?}")) + }) + .group(&unit, move |key| match double_buf { + false => PRIORITY_MAX, + true => double_buffering_priority(key, PRIORITY_MAX, PRIORITY_HIGH), + }), + ) + } + } + + // Gemm no stage + // In unit because not accelerated + let gemm_no_stage_strategy = Strategy::Gemm(BlueprintStrategy::Inferred(GemmStrategy { + target_num_planes: None, + })); + set = set.with( + Tunable::new( + &gemm_no_stage_strategy.to_string(), + move |(lhs, rhs, out)| { + launch_matmul::(&gemm_no_stage_strategy, lhs, rhs, out) + .map_err(|err| format!("{err:?}")) + }, + ) + .group(&unit, move |_key| PRIORITY_MAX), + ); + + // CPU GEMM (CPU-only via the `cpu` group; the size limit is specific to this strategy). + let cpu_gemm_strategy = + Strategy::CpuGemm(BlueprintStrategy::Inferred(CpuGemmStrategy::default())); + set = set.with( + Tunable::new(&cpu_gemm_strategy.to_string(), move |(lhs, rhs, out)| { + launch_matmul::(&cpu_gemm_strategy, lhs, rhs, out) + .map_err(|err| format!("{err:?}")) + }) + .group(&cpu, move |_key| PRIORITY_MAX), + ); + + // Accelerated matmuls + for (strategy, double_buf, group_extra, tile_group) in [ + ( + Strategy::SimpleCyclicCmma(BlueprintStrategy::Inferred(SimpleArgs { + multi_rows: false, + tile_matmul: TileMatmulKind::Cmma, + })), + false, + None, + &accelerated, + ), + ( + Strategy::SimpleCyclicMma(BlueprintStrategy::Inferred(SimpleArgs { + multi_rows: false, + tile_matmul: TileMatmulKind::Mma, + })), + false, + None, + &accelerated, + ), + ( + Strategy::SimpleCyclicCmma(BlueprintStrategy::Inferred(SimpleArgs { + multi_rows: true, + tile_matmul: TileMatmulKind::Cmma, + })), + false, + None, + &accelerated, + ), + ( + Strategy::SimpleCyclicMma(BlueprintStrategy::Inferred(SimpleArgs { + multi_rows: true, + tile_matmul: TileMatmulKind::Mma, + })), + false, + None, + &accelerated, + ), + ( + Strategy::OrderedDoubleCmma(BlueprintStrategy::Inferred(OrderedSelectionArgs { + partition_k: Some(2), + row_count: Some(4), + rows_per_plane: Some(2), + tile_matmul: TileMatmulKind::Cmma, + })), + true, + None, + &accelerated, + ), + ( + Strategy::OrderedDoubleMma(BlueprintStrategy::Inferred(OrderedSelectionArgs { + partition_k: Some(2), + row_count: Some(4), + rows_per_plane: Some(2), + tile_matmul: TileMatmulKind::Mma, + })), + true, + None, + &accelerated, + ), + ( + Strategy::OrderedDoubleCmma(BlueprintStrategy::Inferred(OrderedSelectionArgs { + partition_k: Some(2), + row_count: Some(8), + rows_per_plane: Some(2), + tile_matmul: TileMatmulKind::Cmma, + })), + true, + None, + &accelerated, + ), + ( + Strategy::OrderedDoubleMma(BlueprintStrategy::Inferred(OrderedSelectionArgs { + partition_k: Some(2), + row_count: Some(8), + rows_per_plane: Some(2), + tile_matmul: TileMatmulKind::Mma, + })), + true, + None, + &accelerated, + ), + ( + Strategy::DoubleCyclicCmma(BlueprintStrategy::Inferred(DoubleBufferingArgs { + specialized: false, + tile_matmul: TileMatmulKind::Cmma, + })), + true, + None, + &accelerated, + ), + ( + Strategy::DoubleCyclicMma(BlueprintStrategy::Inferred(DoubleBufferingArgs { + specialized: false, + tile_matmul: TileMatmulKind::Mma, + })), + true, + None, + &accelerated, + ), + ( + Strategy::DoubleCyclicCmma(BlueprintStrategy::Inferred(DoubleBufferingArgs { + specialized: true, + tile_matmul: TileMatmulKind::Cmma, + })), + true, + None, + &accelerated, + ), + ( + Strategy::DoubleCyclicMma(BlueprintStrategy::Inferred(DoubleBufferingArgs { + specialized: true, + tile_matmul: TileMatmulKind::Mma, + })), + true, + None, + &accelerated, + ), + ( + Strategy::SpecializedCyclicCmma(BlueprintStrategy::Inferred(().into())), + true, + None, + &accelerated, + ), + ( + Strategy::SpecializedCyclicMma(BlueprintStrategy::Inferred(().into())), + true, + None, + &accelerated, + ), + ( + Strategy::SimpleTmaCmma(BlueprintStrategy::Inferred(SimpleArgs { + multi_rows: false, + tile_matmul: TileMatmulKind::Cmma, + })), + false, + Some(&tma), + &accelerated, + ), + ( + Strategy::SimpleTmaMma(BlueprintStrategy::Inferred(SimpleArgs { + multi_rows: false, + tile_matmul: TileMatmulKind::Mma, + })), + false, + Some(&tma), + &accelerated, + ), + ( + Strategy::SimpleTmaCmma(BlueprintStrategy::Inferred(SimpleArgs { + multi_rows: true, + tile_matmul: TileMatmulKind::Cmma, + })), + false, + Some(&tma), + &accelerated, + ), + ( + Strategy::SimpleTmaMma(BlueprintStrategy::Inferred(SimpleArgs { + multi_rows: true, + tile_matmul: TileMatmulKind::Mma, + })), + false, + Some(&tma), + &accelerated, + ), + ( + Strategy::SpecializedTmaCmma(BlueprintStrategy::Inferred(().into())), + true, + Some(&tma), + &accelerated, + ), + ( + Strategy::SpecializedTmaMma(BlueprintStrategy::Inferred(().into())), + true, + Some(&tma), + &accelerated, + ), + ] { + let priority_within_group = |key: &MatmulAutotuneKey, double_buf: bool| match double_buf + { + false => PRIORITY_MAX, + true => double_buffering_priority(key, PRIORITY_MAX, PRIORITY_HIGH), + }; + let mut tunable = Tunable::new(&strategy.to_string(), move |(lhs, rhs, out)| { + launch_matmul::(&strategy, lhs, rhs, out).map_err(|err| format!("{err:?}")) + }); + + // tile group + tunable = tunable.group(tile_group, move |key| { + priority_within_group(key, double_buf) + }); + + // extra group + if let Some(group) = group_extra { + tunable = tunable.group(group, move |key| priority_within_group(key, double_buf)); + } + set = set.with(tunable); + } + + set + }); + + TUNER.execute( + &CubeTuneId::new(&lhs.client, &lhs.device), + &client, + tunables, + (lhs, rhs, output.clone()), + ); + + output +} + +fn create_key( + (lhs, rhs, out): &(CubeTensor, CubeTensor, CubeTensor), +) -> MatmulAutotuneKey { + MatmulAutotuneKey::generate( + &lhs.client, + lhs.meta.shape(), + rhs.meta.shape(), + lhs.meta.strides(), + rhs.meta.strides(), + dtype_to_storage_type(lhs.dtype), + dtype_to_storage_type(rhs.dtype), + dtype_to_storage_type(out.dtype), + lhs.try_scheme(), + rhs.try_scheme(), + ) +} diff --git a/crates/burn-cubecl/src/kernel/matmul/tune/mod.rs b/crates/burn-cubecl/src/kernel/matmul/tune/mod.rs new file mode 100644 index 0000000..87ee928 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/matmul/tune/mod.rs @@ -0,0 +1,5 @@ +#[cfg(feature = "autotune")] +mod base; + +#[cfg(feature = "autotune")] +pub use base::matmul_autotune; diff --git a/crates/burn-cubecl/src/kernel/matmul/utils.rs b/crates/burn-cubecl/src/kernel/matmul/utils.rs new file mode 100644 index 0000000..fbbf69f --- /dev/null +++ b/crates/burn-cubecl/src/kernel/matmul/utils.rs @@ -0,0 +1,16 @@ +use crate::{CubeRuntime, ops::numeric::empty_device_dtype, tensor::CubeTensor}; +use burn_backend::{DType, calculate_matmul_output}; + +/// Creates an empty output tensor with matmul output shape +pub fn init_matmul_output( + lhs: &CubeTensor, + rhs: &CubeTensor, + dtype: DType, +) -> CubeTensor { + empty_device_dtype( + lhs.client.clone(), + lhs.device.clone(), + calculate_matmul_output(lhs.meta.shape(), rhs.meta.shape()).unwrap(), + dtype, + ) +} diff --git a/crates/burn-cubecl/src/kernel/mod.rs b/crates/burn-cubecl/src/kernel/mod.rs new file mode 100644 index 0000000..d59d8a9 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/mod.rs @@ -0,0 +1,55 @@ +mod binary; +mod binary_float; +mod binary_int; +mod cast; +mod clamp; +mod comparison; +mod contiguous; +mod cross; +mod index; +mod mask; +mod unary_float; +mod unary_int; +mod unary_numeric; + +pub(crate) use binary::*; +pub(crate) use binary_float::*; +pub(crate) use binary_int::*; +pub use cast::*; +pub use contiguous::*; +pub(crate) use cross::*; +pub use mask::*; +pub(crate) use unary_float::*; +pub(crate) use unary_int::*; +pub(crate) use unary_numeric::*; + +pub use crate::cubecl::prelude::KernelMetadata; + +/// Attention kernels +pub mod attention; +/// Convolution kernels +pub mod conv; +/// CTC loss kernel +pub mod ctc; +/// FFT algorithms +pub mod fft; +/// Grid sampling kernels +pub mod grid_sample; +/// Interpolation kernels +pub mod interpolate; +/// Matmul kernels +pub mod matmul; +/// Pooling kernels +pub mod pool; +/// Pseudo-random number generator kernels +pub mod prng; +/// Quantization operations +pub mod quantization; +/// Reduction algorithms +pub mod reduce; + +pub(crate) use clamp::*; +pub(crate) use comparison::*; +pub use index::*; + +pub(crate) mod utils; diff --git a/crates/burn-cubecl/src/kernel/pool/base.rs b/crates/burn-cubecl/src/kernel/pool/base.rs new file mode 100644 index 0000000..382c83e --- /dev/null +++ b/crates/burn-cubecl/src/kernel/pool/base.rs @@ -0,0 +1,336 @@ +use crate::{ + CubeRuntime, + kernel::into_contiguous_aligned, + ops::{numeric::empty_device_dtype, permute_nchw_to_nhwc, permute_nhwc_to_nchw}, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{DType, Shape, ops::conv::calculate_pool_output_size}; +use cubek::pool::{ + definition::{AdaptiveAvgPoolOptions, AvgPoolOptions, MaxPoolOptions, PoolError, PoolMode}, + pool2d, pool2d_backward, pool2d_with_indices, pool2d_with_indices_backward, +}; + +pub(crate) fn max_pool2d( + x: CubeTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, +) -> CubeTensor { + let [batch_size, channels, height, width] = x.meta.shape().dims(); + + let size_0 = calculate_pool_output_size( + kernel_size[0], + stride[0], + padding[0], + dilation[0], + height, + ceil_mode, + ); + let size_1 = calculate_pool_output_size( + kernel_size[1], + stride[1], + padding[1], + dilation[1], + width, + ceil_mode, + ); + + let x = into_contiguous_aligned(permute_nchw_to_nhwc(x)); + + let shape_out = Shape::new([batch_size, size_0, size_1, channels]); + let output = empty_device_dtype(x.client.clone(), x.device.clone(), shape_out, x.dtype); + + let mode = PoolMode::from(MaxPoolOptions::new( + kernel_size, + stride, + padding, + dilation, + ceil_mode, + )); + + pool2d( + &output.client, + x.clone().binding(), + output.clone().binding(), + mode, + dtype_to_storage_type(output.dtype), + ) + .unwrap_or_else(|e| pool_panic("max_pool2d", &x, e)); + + permute_nhwc_to_nchw(output) +} + +pub(crate) fn max_pool2d_with_indices( + x: CubeTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + dtype_indices: DType, +) -> (CubeTensor, CubeTensor) { + let [batch_size, channels, size_0, size_1] = x.meta.shape().dims(); + + let size_0 = calculate_pool_output_size( + kernel_size[0], + stride[0], + padding[0], + dilation[0], + size_0, + ceil_mode, + ); + let size_1 = calculate_pool_output_size( + kernel_size[1], + stride[1], + padding[1], + dilation[1], + size_1, + ceil_mode, + ); + + let x = into_contiguous_aligned(permute_nchw_to_nhwc(x)); + + let shape_out = Shape::new([batch_size, size_0, size_1, channels]); + let output = empty_device_dtype( + x.client.clone(), + x.device.clone(), + shape_out.clone(), + x.dtype, + ); + let indices = empty_device_dtype(x.client.clone(), x.device.clone(), shape_out, dtype_indices); + + let mode = PoolMode::from(MaxPoolOptions::new( + kernel_size, + stride, + padding, + dilation, + ceil_mode, + )); + + pool2d_with_indices( + &output.client, + x.clone().binding(), + output.clone().binding(), + indices.clone().binding(), + mode, + dtype_to_storage_type(output.dtype), + ) + .unwrap_or_else(|e| pool_panic("max_pool2d_with_indices", &x, e)); + + let output = permute_nhwc_to_nchw(output); + let indices = permute_nhwc_to_nchw(indices); + (output, indices) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn max_pool2d_with_indices_backward( + x: CubeTensor, + grad: CubeTensor, + indices: CubeTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, +) -> CubeTensor { + let [batches, channels, height, width] = x.meta.shape().dims(); + let input = into_contiguous_aligned(permute_nchw_to_nhwc(x)); + let grad = into_contiguous_aligned(permute_nchw_to_nhwc(grad)); + let indices = into_contiguous_aligned(permute_nchw_to_nhwc(indices)); + + let out_shape = Shape::new([batches, height, width, channels]); + let output = empty_device_dtype( + input.client.clone(), + input.device.clone(), + out_shape, + input.dtype, + ); + + let mode = PoolMode::from(MaxPoolOptions::new( + kernel_size, + stride, + padding, + dilation, + ceil_mode, + )); + + pool2d_with_indices_backward( + &output.client, + input.clone().binding(), + grad.clone().binding(), + indices.clone().binding(), + output.clone().binding(), + mode, + dtype_to_storage_type(output.dtype), + dtype_to_storage_type(indices.dtype), + ) + .unwrap_or_else(|e| pool_panic("max_pool2d_with_indices_backward", &input, e)); + + permute_nhwc_to_nchw(output) +} + +pub(crate) fn avg_pool2d( + x: CubeTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, +) -> CubeTensor { + let [batch_size, channels, in_h, in_w] = x.meta.shape().dims(); + let dilation = 1; + + let size_0 = calculate_pool_output_size( + kernel_size[0], + stride[0], + padding[0], + dilation, + in_h, + ceil_mode, + ); + let size_1 = calculate_pool_output_size( + kernel_size[1], + stride[1], + padding[1], + dilation, + in_w, + ceil_mode, + ); + + let x = into_contiguous_aligned(permute_nchw_to_nhwc(x)); + + let shape_out = Shape::new([batch_size, size_0, size_1, channels]); + let output = empty_device_dtype(x.client.clone(), x.device.clone(), shape_out, x.dtype); + + let mode = PoolMode::from(AvgPoolOptions::new( + kernel_size, + stride, + padding, + ceil_mode, + count_include_pad, + )); + + pool2d( + &output.client, + x.clone().binding(), + output.clone().binding(), + mode, + dtype_to_storage_type(output.dtype), + ) + .unwrap_or_else(|e| pool_panic("avg_pool2d", &x, e)); + + permute_nhwc_to_nchw(output) +} + +pub(crate) fn avg_pool2d_backward( + x: CubeTensor, + grad: CubeTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, +) -> CubeTensor { + let [batches, channels, height, width] = x.meta.shape().dims(); + let input = into_contiguous_aligned(permute_nchw_to_nhwc(x)); + let grad = into_contiguous_aligned(permute_nchw_to_nhwc(grad)); + + let out_shape = Shape::new([batches, height, width, channels]); + let output = empty_device_dtype( + input.client.clone(), + input.device.clone(), + out_shape, + input.dtype, + ); + + let mode = PoolMode::from(AvgPoolOptions::new( + kernel_size, + stride, + padding, + ceil_mode, + count_include_pad, + )); + + pool2d_backward( + &output.client, + input.clone().binding(), + grad.clone().binding(), + output.clone().binding(), + mode, + dtype_to_storage_type(output.dtype), + ) + .unwrap_or_else(|e| pool_panic("avg_pool2d_backward", &input, e)); + + permute_nhwc_to_nchw(output) +} + +pub(crate) fn adaptive_avg_pool2d( + input: CubeTensor, + output_size: [usize; 2], +) -> CubeTensor { + let [batch_size, channels, _, _] = input.meta.shape().dims(); + let input = into_contiguous_aligned(permute_nchw_to_nhwc(input)); + + let output_shape = Shape::new([batch_size, output_size[0], output_size[1], channels]); + let output = empty_device_dtype( + input.client.clone(), + input.device.clone(), + output_shape, + input.dtype, + ); + + let mode = PoolMode::from(AdaptiveAvgPoolOptions::new(output_size)); + + pool2d( + &output.client, + input.clone().binding(), + output.clone().binding(), + mode, + dtype_to_storage_type(output.dtype), + ) + .unwrap_or_else(|e| pool_panic("adaptive_avg_pool2d", &input, e)); + + permute_nhwc_to_nchw(output) +} + +pub(crate) fn adaptive_avg_pool2d_backward( + x: CubeTensor, + out_grad: CubeTensor, +) -> CubeTensor { + let [batches, channels, height, width] = x.meta.shape().dims(); + let [_, _, out_h, out_w] = out_grad.meta.shape().dims(); + let input = into_contiguous_aligned(permute_nchw_to_nhwc(x)); + let out_grad = into_contiguous_aligned(permute_nchw_to_nhwc(out_grad)); + + let out_shape = Shape::new([batches, height, width, channels]); + let output = empty_device_dtype( + input.client.clone(), + input.device.clone(), + out_shape, + input.dtype, + ); + + let mode = PoolMode::from(AdaptiveAvgPoolOptions::new([out_h, out_w])); + + pool2d_backward( + &output.client, + input.clone().binding(), + out_grad.clone().binding(), + output.clone().binding(), + mode, + dtype_to_storage_type(output.dtype), + ) + .unwrap_or_else(|e| pool_panic("adaptive_avg_pool2d_backward", &input, e)); + + permute_nhwc_to_nchw(output) +} + +fn pool_panic(label: &str, input: &CubeTensor, error: PoolError) -> ! { + panic!( + "{0} kernel failed (device={1:?}, dtype={2:?}): {3}", + label, input.device, input.dtype, error + ) +} diff --git a/crates/burn-cubecl/src/kernel/pool/mod.rs b/crates/burn-cubecl/src/kernel/pool/mod.rs new file mode 100644 index 0000000..41e113f --- /dev/null +++ b/crates/burn-cubecl/src/kernel/pool/mod.rs @@ -0,0 +1,3 @@ +mod base; + +pub(crate) use base::*; diff --git a/crates/burn-cubecl/src/kernel/prng/bernoulli.rs b/crates/burn-cubecl/src/kernel/prng/bernoulli.rs new file mode 100644 index 0000000..395cf15 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/prng/bernoulli.rs @@ -0,0 +1,24 @@ +use crate::{CubeRuntime, ops::numeric::empty_device_dtype, tensor::CubeTensor}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{DType, Shape}; + +/// Pseudo-random generator with bernoulli distribution +pub fn random_bernoulli( + shape: Shape, + device: &R::Device, + probability: f32, + dtype: DType, +) -> CubeTensor { + let client = R::client(device); + let output = empty_device_dtype(client.clone(), device.clone(), shape, dtype); + + cubek::random::random_bernoulli( + &client, + probability, + output.clone().binding(), + dtype_to_storage_type(dtype), + ) + .expect("Kernel to never fail"); + + output +} diff --git a/crates/burn-cubecl/src/kernel/prng/mod.rs b/crates/burn-cubecl/src/kernel/prng/mod.rs new file mode 100644 index 0000000..10f2cfe --- /dev/null +++ b/crates/burn-cubecl/src/kernel/prng/mod.rs @@ -0,0 +1,7 @@ +mod bernoulli; +mod normal; +mod uniform; + +pub use bernoulli::*; +pub use normal::*; +pub use uniform::*; diff --git a/crates/burn-cubecl/src/kernel/prng/normal.rs b/crates/burn-cubecl/src/kernel/prng/normal.rs new file mode 100644 index 0000000..462580f --- /dev/null +++ b/crates/burn-cubecl/src/kernel/prng/normal.rs @@ -0,0 +1,26 @@ +use crate::{CubeRuntime, ops::numeric::empty_device_dtype, tensor::CubeTensor}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{DType, Shape}; + +/// Pseudo-random generator with uniform distribution +pub fn random_normal( + shape: Shape, + device: &R::Device, + mean: f32, + std: f32, + dtype: DType, +) -> CubeTensor { + let client = R::client(device); + let output = empty_device_dtype(client.clone(), device.clone(), shape, dtype); + + cubek::random::random_normal( + &client, + mean, + std, + output.clone().binding(), + dtype_to_storage_type(dtype), + ) + .expect("Kernel to never fail"); + + output +} diff --git a/crates/burn-cubecl/src/kernel/prng/uniform.rs b/crates/burn-cubecl/src/kernel/prng/uniform.rs new file mode 100644 index 0000000..e29a51f --- /dev/null +++ b/crates/burn-cubecl/src/kernel/prng/uniform.rs @@ -0,0 +1,43 @@ +use crate::{CubeRuntime, ops::numeric::empty_device_dtype, tensor::CubeTensor}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{DType, Shape, TensorMetadata}; + +/// Pseudo-random generator with uniform distribution +pub fn random_uniform( + shape: Shape, + device: &R::Device, + lower_bound: f32, + upper_bound: f32, + dtype: DType, +) -> CubeTensor { + let client = R::client(device); + let output = empty_device_dtype(client.clone(), device.clone(), shape, dtype); + + cubek::random::random_uniform( + &client, + lower_bound, + upper_bound, + output.clone().binding(), + dtype_to_storage_type(dtype), + ) + .expect("Kernel to never fail"); + + output +} + +/// Pseudo-random generator for uniform distribution, based on +/// another tensor. +pub fn random_like_uniform( + tensor: &CubeTensor, + lower_bound: f32, + upper_bound: f32, + dtype: DType, +) -> CubeTensor { + random_uniform( + tensor.shape(), + &tensor.device, + lower_bound, + upper_bound, + dtype, + ) +} diff --git a/crates/burn-cubecl/src/kernel/quantization/dequantize.rs b/crates/burn-cubecl/src/kernel/quantization/dequantize.rs new file mode 100644 index 0000000..667896d --- /dev/null +++ b/crates/burn-cubecl/src/kernel/quantization/dequantize.rs @@ -0,0 +1,35 @@ +use crate::tensor::CubeTensor; +use crate::{CubeRuntime, ops::numeric::empty_device_dtype}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{DType, TensorMetadata}; + +/// Convert the tensor back to a higher precision data type. +pub fn dequantize(tensor: CubeTensor, dtype: DType) -> CubeTensor +where + R: CubeRuntime, +{ + let scheme = match tensor.dtype { + DType::QFloat(scheme) => scheme, + _ => return tensor, + }; + + let output = empty_device_dtype( + tensor.client.clone(), + tensor.device.clone(), + tensor.shape(), + dtype, + ); + let (values, params) = tensor.quantized_handles().unwrap(); + + cubek::quantization::dequantize::launch_ref( + &output.client, + values.binding(), + output.clone().binding(), + params.binding(), + &scheme, + dtype_to_storage_type(dtype), + ) + .expect("Kernel to never fail"); + + output +} diff --git a/crates/burn-cubecl/src/kernel/quantization/mod.rs b/crates/burn-cubecl/src/kernel/quantization/mod.rs new file mode 100644 index 0000000..a0244df --- /dev/null +++ b/crates/burn-cubecl/src/kernel/quantization/mod.rs @@ -0,0 +1,5 @@ +mod dequantize; +mod quantize; + +pub use dequantize::*; +pub use quantize::*; diff --git a/crates/burn-cubecl/src/kernel/quantization/quantize.rs b/crates/burn-cubecl/src/kernel/quantization/quantize.rs new file mode 100644 index 0000000..25ca88b --- /dev/null +++ b/crates/burn-cubecl/src/kernel/quantization/quantize.rs @@ -0,0 +1,31 @@ +use crate::CubeRuntime; +use crate::{ops::empty_qtensor_optimized, tensor::CubeTensor}; +use burn_backend::cubecl::dtype_to_elem_type; +use burn_backend::{TensorMetadata, quantization::QuantScheme}; + +/// Convert the tensor to a lower precision data type based on the quantization scheme and parameters. +pub fn quantize( + tensor: CubeTensor, + scheme: &QuantScheme, + scale: CubeTensor, +) -> CubeTensor +where + R: CubeRuntime, +{ + let output = empty_qtensor_optimized(tensor.shape(), *scheme, &tensor.device); + let (out_values, out_params) = output.clone().quantized_handles().unwrap(); + let dtype = tensor.dtype; + + cubek::quantization::quantize::launch_ref( + &output.client, + tensor.binding(), + out_values.binding(), + scale.binding(), + out_params.binding(), + scheme, + dtype_to_elem_type(dtype), + ) + .expect("Kernel to never fail"); + + output +} diff --git a/crates/burn-cubecl/src/kernel/reduce/base.rs b/crates/burn-cubecl/src/kernel/reduce/base.rs new file mode 100644 index 0000000..96aa46b --- /dev/null +++ b/crates/burn-cubecl/src/kernel/reduce/base.rs @@ -0,0 +1,294 @@ +#[cfg(feature = "autotune")] +use super::{autotune_reduce, autotune_sum}; +use crate::{ + CubeRuntime, + ops::numeric::{empty_device_contiguous_dtype, zeros_client}, + tensor::CubeTensor, +}; +use burn_backend::cubecl::{dtype_to_elem_type, elem_type_to_dtype}; +use burn_backend::{DType, TensorMetadata}; +use burn_std::{BoolDType, Metadata}; +use cubecl::{AutotuneKey, client::ComputeClient, features::AtomicUsage, ir::Type}; +use cubek::reduce::{ + ReduceDtypes, ReduceError, ReduceStrategy, + components::instructions::ReduceOperationConfig, + launch::{RoutineStrategy, VectorizationStrategy}, + routines::{BlueprintStrategy, unit::UnitStrategy}, + shared_sum, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Hash, Eq, PartialEq, Debug, Clone, Serialize, Deserialize, AutotuneKey)] +/// Autotune key representative of sum versions +pub struct SumAutotuneKey { + /// The type of the tensor + dtype: burn_backend::DType, + /// The anchored length of the tensor + #[autotune(anchor)] + length: usize, +} + +/// Check if the client supports atomic add for the given element type. +fn supports_atomic_add(client: &ComputeClient, dtype: DType) -> bool { + client + .properties() + .atomic_type_usage(Type::atomic(dtype_to_elem_type(dtype))) + .contains(AtomicUsage::Add) +} + +/// [Sum](sum) with fallback when `client` doesn't support atomic add for the type `E`. +pub fn sum_fallback( + tensor: CubeTensor, + mut strategy: SumStrategy, +) -> Result, ReduceError> { + // Early check before creating output and fallback + if matches!(strategy, SumStrategy::OneShot(_)) + && !supports_atomic_add(&tensor.client, tensor.dtype) + { + strategy = SumStrategy::Chained(Default::default()); + } + sum(tensor, strategy) +} + +/// Specialize reduce function to compute the sum of all elements of the `input` tensor and return +/// the value into a single-element tensor of shape `1 x 1 x 1 x ...` with the same rank as `input`. +/// +/// This is expected to be faster for larger tensors than calling [reduce] with the `Sum` instruction. +/// +/// Return an error if the `client` doesn't support atomic add for the type `E`. +pub fn sum( + tensor: CubeTensor, + strategy: SumStrategy, +) -> Result, ReduceError> { + let client = tensor.client.clone(); + let device = tensor.device.clone(); + + match strategy { + SumStrategy::OneShot(cube_count) => { + let output = zeros_client(client.clone(), device, [1].into(), tensor.dtype); + let dtype = tensor.dtype; + + shared_sum::( + &client, + tensor.binding(), + output.clone().binding(), + cube_count, + dtype_to_elem_type(dtype), + )?; + + Ok(output) + } + SumStrategy::Chained(strategy) => { + reduce::(tensor, None, strategy, ReduceOperationConfig::Sum) + } + #[cfg(feature = "autotune")] + SumStrategy::Autotune => Ok(autotune_sum::(&client, tensor)), + } +} + +/// Select a strategy to perform a sum. +pub enum SumStrategy { + /// Run a single kernel with many cubes working in parallel to sum all elements. + /// The provided value is the number of elements summed per unit (up-to-rounding ) + OneShot(u32), + /// Use multiple kernels + Chained(KernelReduceStrategy), + /// Use autotune to find the best cube count given the hardware and the input. + #[cfg(feature = "autotune")] + Autotune, +} + +impl Default for SumStrategy { + fn default() -> Self { + #[cfg(feature = "autotune")] + return Self::Autotune; + + #[cfg(not(feature = "autotune"))] + return Self::OneShot(4); + } +} + +/// Reduce all elements of the `input` tensor using the instruction `Rd` and the given [Strategy](ReduceStrategy). +/// +/// Return an error if `strategy` is `Specific(strategy)` and the specified strategy is not supported by the `client`. +/// +/// If there is no error, the output is a tensor with decreasing strides +/// where the shape of reduced dim is set to 1 but all shape are similar to the input. +pub fn reduce( + mut tensor: CubeTensor, + output_dtype: Option, + strategy: KernelReduceStrategy, + config: ReduceOperationConfig, +) -> Result, cubek::reduce::ReduceError> { + // In practice, it looks like starting by the axis with the smallest shape + // and going in increasing order lead to the fastest calculation. + let sorted_axis = argsort(tensor.meta.shape()); + for axis in sorted_axis { + tensor = reduce_dim::(tensor, output_dtype, axis, strategy.clone(), config)?; + } + // reshape to scalar tensor + *tensor.meta = Metadata::new([1], [1]); + Ok(tensor) +} + +/// Reduce with a logical instruction ([`Any`](ReduceOperationConfig::Any) / +/// [`All`](ReduceOperationConfig::All)) and return the result as a boolean tensor. +/// +/// `Any` / `All` require the output dtype (like `Arg*` index outputs): the +/// kernel writes the `0/1` flags directly into the numeric backing of the +/// boolean storage (cubek has no bool elem), so the only step left here is the +/// kernel-free relabel to `Bool`. +/// +/// `dim == None` reduces the whole tensor to a scalar; `Some(dim)` reduces a +/// single axis, keeping it with length 1. +pub fn reduce_logical( + tensor: CubeTensor, + dim: Option, + config: ReduceOperationConfig, + out_dtype: BoolDType, +) -> CubeTensor { + debug_assert!( + matches!( + config, + ReduceOperationConfig::Any | ReduceOperationConfig::All + ), + "reduce_logical only supports Any / All, got {config:?}" + ); + let out_bool = DType::Bool(out_dtype); + let backing = elem_type_to_dtype(dtype_to_elem_type(out_bool)); + + let mut out = match dim { + Some(d) => reduce_dim::(tensor, Some(backing), d, Default::default(), config), + None => reduce::(tensor, Some(backing), Default::default(), config), + } + .expect("Any/All reduce on a valid axis cannot fail"); + + out.dtype = out_bool; // same storage, relabel as Bool (no kernel) + out +} + +fn argsort(shape: &[usize]) -> Vec { + let mut indices = (0..shape.len()).collect::>(); + indices.sort_by_key(|&i| &shape[i]); + indices +} + +/// Reduce the given `axis` of the `input` tensor using the instruction `Rd` and the given [Strategy](ReduceStrategy). +/// +/// Return an error if `strategy` is `Specific(strategy)` and the specified strategy is not supported by the `client`. +/// Also returns an error if the `axis` is larger than the `input` rank or if the shape of `output` is invalid. +/// +/// If there is no error, the output is a tensor with decreasing strides +/// where the shape of reduced dim is set to 1 but all shape are similar to the input. +pub fn reduce_dim( + input: CubeTensor, + output_dtype: Option, + dim: usize, + strategy: KernelReduceStrategy, + config: ReduceOperationConfig, +) -> Result, cubek::reduce::ReduceError> { + debug_assert!( + !matches!( + config, + ReduceOperationConfig::ArgMax + | ReduceOperationConfig::ArgMin + | ReduceOperationConfig::ArgTopK(_) + | ReduceOperationConfig::Any + | ReduceOperationConfig::All + ) || output_dtype.is_some(), + "The `output_dtype` has to be `Some` when the `config` is `ArgMax`, `ArgMin`, `ArgTopK`, `Any` or `All`. + " + ); + + let accumulator_len = match config { + ReduceOperationConfig::ArgTopK(k) => k, + ReduceOperationConfig::TopK(k) => k, + _ => 1, + }; + let dtypes = config.precision( + dtype_to_elem_type(input.dtype), + output_dtype.map(dtype_to_elem_type), + ); + let client = input.client.clone(); + let output = init_reduce_output::(&input, dim, &dtypes, accumulator_len).ok_or( + cubek::reduce::ReduceError::InvalidAxis { + axis: dim, + rank: input.meta.num_dims(), + }, + )?; + + let result = match strategy { + KernelReduceStrategy::Unspecified => cubek::reduce::reduce::( + &client, + input.binding(), + output.clone().binding(), + dim, + ReduceStrategy { + routine: RoutineStrategy::Unit(BlueprintStrategy::Inferred(UnitStrategy)), + vectorization: VectorizationStrategy { + parallel_output_vectorization: false, + }, + }, + config, + dtypes, + ), + KernelReduceStrategy::Specific(strategy) => cubek::reduce::reduce::( + &client, + input.binding(), + output.clone().binding(), + dim, + strategy, + config, + dtypes, + ), + #[cfg(feature = "autotune")] + KernelReduceStrategy::Autotune => { + autotune_reduce::(&client, input, output.clone(), dim, config, dtypes); + Ok(()) + } + }; + result.map(|_| output) +} + +/// Creates an empty output tensor with the proper shape and decreasing strides to reduce the given `axis` of `input` +/// or return `None` if `axis` is out-of-bound. +pub fn init_reduce_output( + input: &CubeTensor, + dim: usize, + dtypes: &ReduceDtypes, + accumulator_len: usize, +) -> Option> { + (dim < input.meta.num_dims()).then(|| { + let mut shape_out = input.shape(); + shape_out[dim] = accumulator_len; + empty_device_contiguous_dtype( + input.client.clone(), + input.device.clone(), + shape_out, + elem_type_to_dtype(dtypes.output.elem_type()), + ) + }) +} + +/// Select a strategy to perform a reduction. +#[derive(Clone, Debug)] +pub enum KernelReduceStrategy { + /// Use a best-effort strategy based on the hardware capacity. + /// This differs from Autotune as it doesn't try and compare many strategies to select the best. + Unspecified, + /// Fix the exact strategy for the reduction. + Specific(cubek::reduce::launch::ReduceStrategy), + /// Use autotune to find the best strategy given the hardware and the inputs. + #[cfg(feature = "autotune")] + Autotune, +} + +impl Default for KernelReduceStrategy { + fn default() -> Self { + #[cfg(feature = "autotune")] + return Self::Autotune; + + #[cfg(not(feature = "autotune"))] + return Self::Unspecified; + } +} diff --git a/crates/burn-cubecl/src/kernel/reduce/mod.rs b/crates/burn-cubecl/src/kernel/reduce/mod.rs new file mode 100644 index 0000000..13a6c23 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/reduce/mod.rs @@ -0,0 +1,7 @@ +mod base; +#[cfg(feature = "autotune")] +mod tune; + +pub use base::*; +#[cfg(feature = "autotune")] +pub use tune::*; diff --git a/crates/burn-cubecl/src/kernel/reduce/tune.rs b/crates/burn-cubecl/src/kernel/reduce/tune.rs new file mode 100644 index 0000000..8a60187 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/reduce/tune.rs @@ -0,0 +1,292 @@ +#![allow(missing_docs)] + +use super::SumAutotuneKey; +use crate::{CubeAutotuneKey, CubeRuntime, CubeTuneId, tensor::CubeTensor}; +use burn_backend::cubecl::dtype_to_elem_type; +use cubecl::{ + client::ComputeClient, + tune::{LocalTuner, Tunable, TunableSet, TuneGroup, local_tuner}, +}; +use cubek::reduce::{ + ReduceDtypes, ReduceStrategy, + components::instructions::ReduceOperationConfig, + launch::{RoutineStrategy, VectorizationStrategy, tune_key::ReduceAutotuneKey}, + routines::{BlueprintStrategy, cube::CubeStrategy, plane::PlaneStrategy, unit::UnitStrategy}, +}; + +/// Executes autotune on reduce operations. +pub fn autotune_reduce( + client: &ComputeClient, + input: CubeTensor, + output: CubeTensor, + axis: usize, + config: ReduceOperationConfig, + dtypes: ReduceDtypes, +) { + use reduce_ops::*; + + static TUNER: LocalTuner = local_tuner!("reduce-dim"); + + let tunables = TUNER.init(|| { + const PRIORITY_MAX: i8 = 2; + const PRIORITY_MIN: i8 = 1; + const PRIORITY_SKIP: i8 = -1; + + let mut set = TunableSet::new(create_key::, reduce_input_gen::); + + let default_group = + TuneGroup::::new("default_reduce", |_key| PRIORITY_MAX); + let vectorized_parallel_group = + TuneGroup::::new("vectorized_parallel_reduce", |key| { + if key.axis_is_contiguous { + PRIORITY_MAX + } else { + // We disable the tunable with the setting [vector_size.parallel_output_vectorization] + // when the reduce isn't parallel, since it would duplicate tunables. + PRIORITY_SKIP + } + }); + + enum ReduceProps { + GreatWithLowReduceCount, + GreatWithHighReduceCount, + Balanced, + } + + for (vectorization, vector_size_ident) in [ + ( + VectorizationStrategy { + parallel_output_vectorization: true, + }, + "_vectorized_parallel_reduce", + ), + ( + VectorizationStrategy { + parallel_output_vectorization: false, + }, + "", + ), + ] { + for (name, routine, props) in [ + ( + "unit", + RoutineStrategy::Unit(BlueprintStrategy::Inferred(UnitStrategy)), + ReduceProps::GreatWithHighReduceCount, + ), + ( + "plane", + RoutineStrategy::Plane(BlueprintStrategy::Inferred(PlaneStrategy { + independent: true, + })), + ReduceProps::Balanced, + ), + ( + "cube", + RoutineStrategy::Cube(BlueprintStrategy::Inferred(CubeStrategy { + use_planes: true, + })), + ReduceProps::GreatWithLowReduceCount, + ), + ] { + let name = format!("{name}{vector_size_ident}"); + let mut tunable = Tunable::new( + &name, + move |(input, output, axis, config, dtypes): ( + CubeTensor, + CubeTensor, + usize, + ReduceOperationConfig, + ReduceDtypes, + )| { + let strategy = ReduceStrategy { + routine: routine.clone(), + vectorization, + }; + cubek::reduce::reduce::( + &output.client, + input.binding(), + output.clone().binding(), + axis, + strategy, + config, + dtypes, + ) + .map_err(|e| format!("{e}")) + }, + ); + if vectorization.parallel_output_vectorization { + tunable = tunable.group(&vectorized_parallel_group, |_| PRIORITY_MAX); + } + + tunable = tunable.group(&default_group, move |key| match props { + ReduceProps::GreatWithLowReduceCount => { + if key.vector_count < 128 { + PRIORITY_MAX + } else { + // When you have a high level of vector to reduce, it is normally + // better to use another routine. + PRIORITY_MIN + } + } + ReduceProps::GreatWithHighReduceCount => { + if key.vector_count > 64 { + PRIORITY_MAX + } else { + // Bellow 64 it is normally better to use another routine + PRIORITY_MIN + } + } + ReduceProps::Balanced => PRIORITY_MAX, + }); + set = set.with(tunable); + } + } + + set + }); + + TUNER.execute( + &CubeTuneId::new(&input.client, &input.device), + client, + tunables, + (input, output, axis, config, dtypes), + ); +} + +pub(crate) fn create_key( + (input, output, axis, _config, dtypes): &( + CubeTensor, + CubeTensor, + usize, + ReduceOperationConfig, + ReduceDtypes, + ), +) -> ReduceAutotuneKey { + let elem_input = dtype_to_elem_type(input.dtype); + let elem_output = dtype_to_elem_type(output.dtype); + let elem_acc = dtypes.accumulation.elem_type(); + + ReduceAutotuneKey::generate( + elem_input, + elem_output, + elem_acc, + input.meta.shape(), + input.meta.strides()[*axis] == 1, + *axis, + ) +} + +mod reduce_ops { + #![allow(missing_docs)] + + use cubek::reduce::ReduceDtypes; + + use super::*; + + pub(crate) fn reduce_input_gen( + _key: &ReduceAutotuneKey, + (input, output, dim, config, dtypes): &( + CubeTensor, + CubeTensor, + usize, + ReduceOperationConfig, + ReduceDtypes, + ), + ) -> ( + CubeTensor, + CubeTensor, + usize, + ReduceOperationConfig, + ReduceDtypes, + ) { + (input.clone(), output.copy(), *dim, *config, *dtypes) + } +} + +/// Executes autotune on reduce operations. +#[cfg(feature = "autotune")] +pub fn autotune_sum( + client: &ComputeClient, + input: CubeTensor, +) -> CubeTensor { + use sum_ops::*; + + static TUNER: LocalTuner = local_tuner!("autotune-sum"); + + let tunables = TUNER.init(|| { + TunableSet::new(create_key_sum::, sum_input_gen::) + .with(Tunable::new("sum_chained", sum_chained::)) + .with(Tunable::new("sum_one_shot", sum_one_shot::)) + .with(Tunable::new("sum_one_shot", sum_one_shot::)) + .with(Tunable::new("sum_one_shot", sum_one_shot::)) + .with(Tunable::new("sum_one_shot", sum_one_shot::)) + .with(Tunable::new("sum_one_shot", sum_one_shot::)) + .with(Tunable::new("sum_one_shot", sum_one_shot::)) + .with(Tunable::new("sum_one_shot", sum_one_shot::)) + }); + + TUNER.execute( + &CubeTuneId::new(&input.client, &input.device), + client, + tunables, + input, + ) +} + +pub(crate) fn create_key_sum(input: &CubeTensor) -> CubeAutotuneKey { + CubeAutotuneKey::Sum(SumAutotuneKey::generate(input)) +} + +impl SumAutotuneKey { + #[allow(unused)] + pub(crate) fn generate(input: &CubeTensor) -> Self { + let dtype = input.dtype; + let length = input.meta.num_elements(); + Self::new(dtype, length) + } +} +mod sum_ops { + #![allow(missing_docs)] + use crate::ops::numeric::zeros_client; + + use super::*; + + pub(crate) fn sum_input_gen( + _key: &CubeAutotuneKey, + input: &CubeTensor, + ) -> CubeTensor { + input.clone() + } + + pub(crate) fn sum_one_shot( + input: CubeTensor, + ) -> Result, String> { + let client = input.client.clone(); + let device = input.device.clone(); + let output = zeros_client(client.clone(), device, [1].into(), input.dtype); + let dtype = input.dtype; + + cubek::reduce::shared_sum::( + &output.client, + input.binding(), + output.clone().binding(), + C, + dtype_to_elem_type(dtype), + ) + .map_err(|e| e.to_string()) + .map(|_| output) + } + + #[cfg(feature = "autotune")] + pub(crate) fn sum_chained( + input: CubeTensor, + ) -> Result, String> { + crate::kernel::reduce::reduce::( + input, + None, + crate::kernel::reduce::KernelReduceStrategy::Autotune, + cubek::reduce::components::instructions::ReduceOperationConfig::Sum, + ) + .map_err(|e| e.to_string()) + } +} diff --git a/crates/burn-cubecl/src/kernel/unary_float.rs b/crates/burn-cubecl/src/kernel/unary_float.rs new file mode 100644 index 0000000..7b9d8e5 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/unary_float.rs @@ -0,0 +1,202 @@ +use crate::{ + CubeRuntime, + kernel::utils::address_type, + ops::{max_vector_size, numeric::empty_device_dtype}, + tensor::CubeTensor, +}; +use burn_backend::TensorMetadata; +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::{ + calculate_cube_count_elemwise, + prelude::*, + std::tensor::layout::linear::{LinearView, LinearViewMut}, +}; + +pub(crate) trait FloatUnaryOpFamily: 'static + Send + Sync { + type Options: LaunchArg; + type Unary: FloatUnaryOp; +} + +#[cube] +pub(crate) trait FloatUnaryOp: 'static + Send + Sync { + type Options: LaunchArg; + + fn execute(input: Vector, options: &Self::Options) -> Vector; +} + +#[cube(launch_unchecked, address_type = "dynamic")] +pub(crate) fn unary_float( + input: LinearView<'_, Vector>, + mut output: LinearViewMut<'_, Vector>, + options: &O::Options, + #[define(F)] _dtype: StorageType, +) { + if !output.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + output.write( + ABSOLUTE_POS, + O::Unary::::execute(input.read(ABSOLUTE_POS), options), + ); +} + +pub(crate) fn launch_unary_float(tensor: CubeTensor, args: Args) -> CubeTensor +where + // Magic fix for lifetime, the closure is supposed to capture everything required to create the + // argument. + for<'a> Args: FnOnce(&'a ()) -> RuntimeArg, + R: CubeRuntime, + O: FloatUnaryOpFamily, +{ + let vector_size = max_vector_size(&tensor); + + let client = tensor.client.clone(); + let num_elems = tensor.meta.num_elements(); + + let working_units = num_elems / vector_size as usize; + let cube_dim = CubeDim::new(&tensor.client, working_units); + let cube_count = calculate_cube_count_elemwise(&tensor.client, working_units, cube_dim); + let dtype = tensor.dtype; + + unsafe { + if tensor.can_mut() && tensor.is_nonoverlapping() { + unary_float::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(tensor), + vector_size, + tensor.clone().into_linear_view(), + tensor.as_linear_view_alias(0), + args(&()), + dtype_to_storage_type(dtype), + ); + + tensor + } else { + let output = empty_device_dtype( + tensor.client.clone(), + tensor.device.clone(), + tensor.shape(), + tensor.dtype, + ); + + unary_float::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(tensor, output), + vector_size, + tensor.into_linear_view(), + output.clone().into_linear_view(), + args(&()), + dtype_to_storage_type(dtype), + ); + + output + } + } +} + +/// Use comptime enum to implement all unary operations that don't have any input argument in the +/// kernel definition. +pub(crate) mod unary_basic { + use cubecl::num_traits::{One, Zero}; + + use super::*; + + pub(crate) fn launch(tensor: CubeTensor, args: Args) -> CubeTensor + where + R: CubeRuntime, + for<'a> Args: FnOnce(&'a ()) -> BasicFloatUnaryKind, + { + launch_unary_float::(tensor, |input| { + BasicFloatUnaryOptionsLaunch::new(args(input)) + }) + } + + #[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, serde::Serialize, serde::Deserialize)] + pub enum BasicFloatUnaryKind { + Exp, + Log, + Log1p, + Sqrt, + Abs, + Sign, + ArcCos, + ArcCosh, + ArcSin, + ArcSinh, + ArcTan, + ArcTanh, + Cos, + Cosh, + Sin, + Sinh, + Tan, + Tanh, + Round, + Floor, + Ceil, + Trunc, + Erf, + Recip, + } + + #[derive(CubeLaunch, CubeType)] + struct BasicFloatUnaryOptions { + #[cube(comptime)] + kind: BasicFloatUnaryKind, + } + struct BasicFloatUnary; + + #[cube] + impl FloatUnaryOp for BasicFloatUnary { + type Options = BasicFloatUnaryOptions; + + fn execute(input: Vector, options: &Self::Options) -> Vector { + match comptime![options.kind] { + BasicFloatUnaryKind::Exp => Vector::exp(input), + BasicFloatUnaryKind::Log => Vector::ln(input), + BasicFloatUnaryKind::Log1p => Vector::log1p(input), + BasicFloatUnaryKind::Sqrt => Vector::sqrt(input), + BasicFloatUnaryKind::Abs => Vector::abs(input), + BasicFloatUnaryKind::Sign => { + let zero = Vector::zero(); + let one = Vector::one(); + let minus_one = Vector::new(F::new(-1.0_f32)); + + let is_positive = input.greater_than(&zero); + let is_negative = input.less_than(&zero); + let sign = select_many(is_negative, minus_one, zero); + + select_many(is_positive, one, sign) + } + BasicFloatUnaryKind::Cos => Vector::cos(input), + BasicFloatUnaryKind::Sin => Vector::sin(input), + BasicFloatUnaryKind::Tan => Vector::tan(input), + BasicFloatUnaryKind::Cosh => Vector::cosh(input), + BasicFloatUnaryKind::Sinh => Vector::sinh(input), + BasicFloatUnaryKind::Tanh => Vector::tanh(input), + BasicFloatUnaryKind::Round => Vector::round(input), + BasicFloatUnaryKind::Floor => Vector::floor(input), + BasicFloatUnaryKind::Ceil => Vector::ceil(input), + BasicFloatUnaryKind::Trunc => Vector::trunc(input), + BasicFloatUnaryKind::Erf => Vector::erf(input), + BasicFloatUnaryKind::Recip => Vector::recip(input), + BasicFloatUnaryKind::ArcCos => Vector::acos(input), + BasicFloatUnaryKind::ArcCosh => Vector::acosh(input), + BasicFloatUnaryKind::ArcSin => Vector::asin(input), + BasicFloatUnaryKind::ArcSinh => Vector::asinh(input), + BasicFloatUnaryKind::ArcTan => Vector::atan(input), + BasicFloatUnaryKind::ArcTanh => Vector::atanh(input), + } + } + } + + impl FloatUnaryOpFamily for BasicFloatUnary { + type Options = BasicFloatUnaryOptions; + type Unary = Self; + } +} diff --git a/crates/burn-cubecl/src/kernel/unary_int.rs b/crates/burn-cubecl/src/kernel/unary_int.rs new file mode 100644 index 0000000..d018389 --- /dev/null +++ b/crates/burn-cubecl/src/kernel/unary_int.rs @@ -0,0 +1,154 @@ +use crate::{ + CubeRuntime, + kernel::utils::address_type, + ops::{max_vector_size, numeric::empty_device_dtype}, + tensor::CubeTensor, +}; +use burn_backend::TensorMetadata; +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::{ + calculate_cube_count_elemwise, + prelude::*, + std::tensor::layout::linear::{LinearView, LinearViewMut}, +}; + +pub(crate) trait IntUnaryOpFamily: 'static + Send + Sync { + type Options: LaunchArg; + type Unary: IntUnaryOp; +} + +#[cube] +pub(crate) trait IntUnaryOp: 'static + Send + Sync { + type Options: LaunchArg; + + fn execute(input: Vector, options: &Self::Options) -> Vector; +} + +#[cube(launch_unchecked, address_type = "dynamic")] +pub(crate) fn unary_int( + input: LinearView<'_, Vector>, + mut output: LinearViewMut<'_, Vector>, + options: &O::Options, + #[define(I)] _dtype: StorageType, +) { + if !output.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + output.write( + ABSOLUTE_POS, + O::Unary::::execute(input.read(ABSOLUTE_POS), options), + ); +} + +pub(crate) fn launch_unary_int(tensor: CubeTensor, args: Args) -> CubeTensor +where + for<'a> Args: FnOnce(&'a ()) -> RuntimeArg, + R: CubeRuntime, + O: IntUnaryOpFamily, +{ + let vector_size = max_vector_size(&tensor); + let client = tensor.client.clone(); + let num_elems = tensor.meta.num_elements(); + + let working_units = num_elems / vector_size as usize; + let cube_dim = CubeDim::new(&tensor.client, working_units); + let cube_count = calculate_cube_count_elemwise(&tensor.client, working_units, cube_dim); + let dtype = tensor.dtype; + + unsafe { + if tensor.can_mut() && tensor.is_nonoverlapping() { + unary_int::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(tensor), + vector_size, + tensor.clone().into_linear_view(), + tensor.as_linear_view_alias(0), + args(&()), + dtype_to_storage_type(dtype), + ); + + tensor + } else { + let output = empty_device_dtype( + tensor.client.clone(), + tensor.device.clone(), + tensor.shape(), + tensor.dtype, + ); + + unary_int::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(tensor, output), + vector_size, + tensor.into_linear_view(), + output.clone().into_linear_view(), + args(&()), + dtype_to_storage_type(dtype), + ); + + output + } + } +} + +pub(crate) mod unary_basic_int { + + use cubecl::num_traits::{One, Zero}; + + use super::*; + + pub(crate) fn launch(tensor: CubeTensor, args: Args) -> CubeTensor + where + R: CubeRuntime, + for<'a> Args: FnOnce(&'a ()) -> BasicIntUnaryKind, + { + launch_unary_int::(tensor, |input| { + BasicIntUnaryOptionsLaunch::new(args(input)) + }) + } + + #[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, serde::Serialize, serde::Deserialize)] + pub enum BasicIntUnaryKind { + BitwiseNot, + Sign, + } + + #[derive(CubeLaunch, CubeType)] + struct BasicIntUnaryOptions { + #[cube(comptime)] + kind: BasicIntUnaryKind, + } + struct BasicIntUnary; + + #[cube] + impl IntUnaryOp for BasicIntUnary { + type Options = BasicIntUnaryOptions; + + fn execute(input: Vector, options: &Self::Options) -> Vector { + match comptime![options.kind] { + BasicIntUnaryKind::BitwiseNot => !input, + BasicIntUnaryKind::Sign => { + let zero = Vector::zero(); + let one = Vector::one(); + let minus_one = Vector::new(I::new(-1)); + + let is_positive = input.greater_than(&zero); + let is_negative = input.less_than(&zero); + let sign = select_many(is_negative, minus_one, zero); + + select_many(is_positive, one, sign) + } + } + } + } + + impl IntUnaryOpFamily for BasicIntUnary { + type Options = BasicIntUnaryOptions; + type Unary = Self; + } +} diff --git a/crates/burn-cubecl/src/kernel/unary_numeric.rs b/crates/burn-cubecl/src/kernel/unary_numeric.rs new file mode 100644 index 0000000..61486fd --- /dev/null +++ b/crates/burn-cubecl/src/kernel/unary_numeric.rs @@ -0,0 +1,99 @@ +use crate::{ + CubeRuntime, + kernel::utils::address_type, + ops::{max_vector_size, numeric::empty_device_dtype}, + tensor::CubeTensor, +}; +use burn_backend::TensorMetadata; +use burn_backend::cubecl::dtype_to_storage_type; +use cubecl::{ + calculate_cube_count_elemwise, + prelude::*, + std::tensor::layout::linear::{LinearView, LinearViewMut}, +}; + +pub(crate) trait NumericUnaryOpFamily: 'static + Send + Sync { + type Options: LaunchArg; + type Unary: NumericUnaryOp; +} + +#[cube] +pub(crate) trait NumericUnaryOp: 'static + Send + Sync { + type Options: LaunchArg; + + fn execute(input: Vector, options: &Self::Options) -> Vector; +} + +#[cube(launch_unchecked, address_type = "dynamic")] +pub(crate) fn unary_numeric( + input: LinearView<'_, Vector>, + mut output: LinearViewMut<'_, Vector>, + options: &O::Options, + #[define(T)] _dtype: StorageType, +) { + if !output.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + output.write( + ABSOLUTE_POS, + O::Unary::::execute(input.read(ABSOLUTE_POS), options), + ); +} + +pub(crate) fn launch_unary_numeric(tensor: CubeTensor, args: Args) -> CubeTensor +where + // Magic fix for lifetime, the closure is supposed to capture everything required to create the + // argument. + for<'a> Args: FnOnce(&'a ()) -> RuntimeArg, + R: CubeRuntime, + O: NumericUnaryOpFamily, +{ + let vector_size = max_vector_size(&tensor); + let client = tensor.client.clone(); + let num_elems = tensor.meta.num_elements(); + + let working_units = num_elems / vector_size as usize; + let cube_dim = CubeDim::new(&tensor.client, working_units); + let cube_count = calculate_cube_count_elemwise(&tensor.client, working_units, cube_dim); + let dtype = tensor.dtype; + + unsafe { + if tensor.can_mut() && tensor.is_nonoverlapping() { + unary_numeric::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(tensor), + vector_size, + tensor.clone().into_linear_view(), + tensor.as_linear_view_alias(0), + args(&()), + dtype_to_storage_type(dtype), + ); + + tensor + } else { + let output = empty_device_dtype( + tensor.client.clone(), + tensor.device.clone(), + tensor.shape(), + tensor.dtype, + ); + + unary_numeric::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(tensor, output), + vector_size, + tensor.into_linear_view(), + output.clone().into_linear_view(), + args(&()), + dtype_to_storage_type(dtype), + ); + + output + } + } +} diff --git a/crates/burn-cubecl/src/kernel/utils.rs b/crates/burn-cubecl/src/kernel/utils.rs new file mode 100644 index 0000000..1df2ced --- /dev/null +++ b/crates/burn-cubecl/src/kernel/utils.rs @@ -0,0 +1,136 @@ +use burn_backend::Shape; +use cubecl::prelude::SequenceArg; +use cubecl::{ + prelude::*, + std::{FastDivmod, FastDivmodInt}, +}; + +use crate::{CubeRuntime, tensor::CubeTensor}; + +pub fn shape_divmod(tensor: &CubeTensor) -> SequenceArg> { + let mut arg = SequenceArg::new(); + for dim in tensor.meta.shape().iter() { + arg.push(*dim); + } + arg +} + +pub fn shape_divmod_range( + tensor: &CubeTensor, + range: core::ops::Range, +) -> SequenceArg> { + let mut arg = SequenceArg::new(); + let shape = &tensor.meta.shape; + for i in range { + arg.push(shape[i]); + } + arg +} + +pub fn split_dim( + mut tensor: CubeTensor, + dim: usize, + shape: &[usize], +) -> CubeTensor { + let mut stride = tensor.meta.strides()[dim]; + tensor.meta.remove(dim); + + for size in shape.iter().rev() { + tensor.meta.insert(dim, *size, stride); + stride *= size; + } + + tensor +} + +pub fn broadcast_shape(tensors: &[&CubeTensor]) -> Shape { + let rank = tensors[0].meta.num_dims(); + debug_assert!( + tensors.iter().all(|it| it.meta.num_dims() == rank), + "Broadcast tensors must have the same rank" + ); + + let dims = (0..rank).map(|dim| { + let max = tensors.iter().map(|it| it.meta.shape()[dim]).max(); + let max = max.unwrap_or(1); + debug_assert!( + tensors + .iter() + .all(|it| it.meta.shape()[dim] == max || it.meta.shape()[dim] == 1), + "Broadcast dims must be size 1" + ); + max + }); + + Shape::from(dims) +} + +pub fn broadcast_strides( + reference: &CubeTensor, + tensor: &CubeTensor, +) -> SequenceArg { + if reference.meta.shape() != tensor.meta.shape() { + tensor + .meta + .strides() + .iter() + .zip( + tensor + .meta + .shape() + .iter() + .zip(reference.meta.shape().iter()), + ) + .map(|(stride, (shape, ref_shape))| if *shape == *ref_shape { *stride } else { 0 }) + .collect() + } else { + tensor.meta.strides().iter().copied().collect() + } +} + +#[cube] +pub(crate) fn decompose_linear( + pos: I, + shape: &Sequence>, +) -> (I, Sequence) { + let rank = comptime![shape.len()]; + let mut offs = pos; + let mut out = Sequence::new(); + + #[unroll] + for i in 0..rank { + let dim = comptime![rank - i - 1]; + let (rem, offs_local) = shape.index(dim).div_mod(offs); + out.push(offs_local); + offs = rem; + } + + (offs, out.reversed()) +} + +pub(crate) trait RequiredAddrType { + fn required_address_type(&self) -> AddressType; +} + +impl RequiredAddrType for CubeTensor { + fn required_address_type(&self) -> AddressType { + self.required_address_type() + } +} +impl RequiredAddrType for Option> { + fn required_address_type(&self) -> AddressType { + self.as_ref() + .map(|it| it.required_address_type()) + .unwrap_or_default() + } +} + +macro_rules! address_type { + ($($tensor: tt),*) => { + [$($crate::kernel::utils::RequiredAddrType::required_address_type(&$tensor)),*] + .into_iter() + .max() + .unwrap_or_default() + }; +} +pub(crate) use address_type; diff --git a/crates/burn-cubecl/src/lib.rs b/crates/burn-cubecl/src/lib.rs new file mode 100644 index 0000000..108b02d --- /dev/null +++ b/crates/burn-cubecl/src/lib.rs @@ -0,0 +1,50 @@ +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] + +//! Burn JIT Backend + +#[macro_use] +extern crate derive_new; +extern crate alloc; + +/// Utilities for implementing JIT kernels +pub mod ops; + +/// Kernel module +pub mod kernel; +/// Tensor module. +pub mod tensor; + +/// Elements for JIT backend +pub mod element; + +use cubecl::{CubeTask, Runtime}; +pub use element::{BoolElement, CubeElement, FloatElement, IntElement}; + +mod backend; + +pub use backend::*; + +// Re-export cubecl. +pub use cubecl; + +mod tune_key; +pub use tune_key::CubeAutotuneKey; + +#[cfg(any(feature = "fusion", test))] +/// Module for interacting with fusion +pub mod fusion; + +#[cfg(feature = "template")] +/// Module for compiling custom non-jit kernels +pub mod template; + +/// Just-in-Time runtime extending the [cube runtime](Runtime). +pub trait CubeRuntime: Runtime { + /// The device that should also implement [burn_backend::backend::DeviceOps]. + type CubeDevice: burn_backend::DeviceOps; + /// The cube server with the [CubeAutotuneKey]. + type CubeServer: cubecl::server::ComputeServer>>; +} + +pub use cubecl::CubeTuneId; diff --git a/crates/burn-cubecl/src/ops/activation.rs b/crates/burn-cubecl/src/ops/activation.rs new file mode 100644 index 0000000..ee0917a --- /dev/null +++ b/crates/burn-cubecl/src/ops/activation.rs @@ -0,0 +1,4 @@ +use crate::{CubeBackend, CubeRuntime}; +use burn_backend::ops::ActivationOps; + +impl ActivationOps for CubeBackend {} diff --git a/crates/burn-cubecl/src/ops/base.rs b/crates/burn-cubecl/src/ops/base.rs new file mode 100644 index 0000000..713b49d --- /dev/null +++ b/crates/burn-cubecl/src/ops/base.rs @@ -0,0 +1,540 @@ +use crate::{CubeRuntime, kernel, ops::numeric::empty_device_dtype, tensor::CubeTensor}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{ + DType, ExecutionError, Shape, TensorData, + quantization::{QuantLevel, QuantStore, params_shape}, +}; +use burn_backend::{TensorMetadata, ops::unfold::calculate_unfold_shape}; +use burn_std::{ + Metadata, QuantValue, ReshapeAnalysis, reshape_analysis, strides, + tensor::{ReshapeAction, contiguous_strides, reshape_action}, +}; +use cubecl::{ir::VectorSize, server::CopyDescriptor}; +use cubecl::{quant::scheme::BlockSize, tensor_vector_size_parallel}; + +pub(crate) fn from_data(data: TensorData, device: &R::Device) -> CubeTensor { + // `TensorData` may contain lazily materialized device-backed bytes produced + // by `into_data()`. These unnecessary round-trips should be avoided, but + // materializing before re-uploading avoids recursive runtime submission. + if data.bytes.property() == burn_std::AllocationProperty::Device { + let _ = data.bytes.read(burn_std::Reader::new()); + } + + let client = R::client(device); + let alloc = client.create_tensor(data.bytes, data.shape.clone(), data.dtype.size()); + let shape: Shape = (&data.shape).into(); + CubeTensor::new( + client, + alloc.memory, + Metadata::new(shape, alloc.strides), + device.clone(), + data.dtype, + ) +} + +pub(crate) async fn into_data( + tensor: CubeTensor, +) -> Result { + let tensor = kernel::into_contiguous_aligned(tensor); + + let elem_size = tensor.elem_size(); + let shape = tensor.meta.shape().clone(); + let strides = tensor.meta.strides().clone(); + let binding = CopyDescriptor::new(tensor.handle.binding(), shape, strides, elem_size); + // Under an async runtime (e.g. a remote server), a lazy read defers the device→host + // copy to first access, where a blocking read would run on — and starve — an executor + // worker. Read eagerly there so the copy happens inside this awaited future. On a + // sync/threaded runtime the lazy read keeps the streaming optimization. + let read = match burn_std::runtime_kind() { + burn_std::RuntimeKind::Async => tensor.client.read_one_tensor_async(binding).await, + _ => tensor.client.read_lazy_async(binding).await, + }; + let bytes = read.map_err(|err| ExecutionError::WithContext { + reason: format!("{err}"), + })?; + + Ok(TensorData::from_bytes( + bytes, + tensor.meta.shape.clone(), + tensor.dtype, + )) +} + +/// Read data from a `CubeTensor` synchronously +#[allow(unused, reason = "useful for debugging kernels")] +pub fn into_data_sync(tensor: CubeTensor) -> TensorData { + burn_std::future::block_on(into_data(tensor)).unwrap() +} + +#[cfg_attr( + feature = "tracing", + tracing::instrument(level = "trace", skip(tensor, device)) +)] +pub(crate) fn to_device( + tensor: CubeTensor, + device: &R::Device, +) -> CubeTensor { + if &tensor.device == device { + return tensor; + } + + let mut tensor = kernel::into_contiguous_aligned(tensor); + let client = R::client(device); + tensor.to_client(client, device.clone()) +} + +pub(crate) fn empty( + shape: Shape, + device: &R::Device, + dtype: DType, +) -> CubeTensor { + let client = R::client(device); + let alloc = client.empty_tensor(shape.clone(), dtype.size()); + + CubeTensor::new( + client, + alloc.memory, + Metadata::new(shape, alloc.strides), + device.clone(), + dtype, + ) +} + +pub(crate) fn swap_dims( + mut tensor: CubeTensor, + dim1: usize, + dim2: usize, +) -> CubeTensor { + tensor.meta.swap(dim1, dim2); + + if let DType::QFloat(scheme) = tensor.dtype + && let QuantLevel::Block(block_size) = scheme.level + { + let rank = tensor.rank(); + let qparams = tensor.qparams.as_mut().unwrap(); + let mut block_size = block_size.to_dim_vec(rank); + block_size.swap(dim1, dim2); + + // Truncate unit dims from the start + let block_size = BlockSize::new_trim(block_size); + if block_size.len() > BlockSize::MAX_DIMS { + panic!("Swapped block size would exceed max dims"); + } + + qparams.scales.metadata.swap(dim1, dim2); + + tensor.dtype = DType::QFloat(scheme.with_level(QuantLevel::Block(block_size))) + } + + if let DType::QFloat(scheme) = &mut tensor.dtype + && let QuantStore::PackedU32(packed_dim) | QuantStore::PackedNative(packed_dim) = + &mut scheme.store + { + let rank = tensor.meta.num_dims(); + + if *packed_dim == rank - dim1 - 1 { + *packed_dim = rank - dim2 - 1; + } else if *packed_dim == rank - dim2 - 1 { + *packed_dim = rank - dim1 - 1; + } + } + + tensor +} + +/// Permute a tensor's dimensions +pub fn permute(mut tensor: CubeTensor, axes: &[usize]) -> CubeTensor { + tensor.meta.permute(axes).unwrap(); + + if let DType::QFloat(scheme) = tensor.dtype + && let QuantLevel::Block(block_size) = scheme.level + { + let rank = tensor.rank(); + let qparams = tensor.qparams.as_mut().unwrap(); + + let mut block_size = block_size.to_dim_vec(rank); + block_size = axes.iter().map(|i| block_size[*i]).collect(); + + // Truncate unit dims from the start + let block_size = block_size + .into_iter() + .skip_while(|it| *it == 1) + .collect::>(); + if block_size.len() > BlockSize::MAX_DIMS { + panic!("Swapped block size would exceed max dims"); + } + + qparams.scales.metadata.permute(axes).unwrap(); + + tensor.dtype = DType::QFloat(scheme.with_level(QuantLevel::block(&block_size))) + } + + if let DType::QFloat(scheme) = &mut tensor.dtype + && let QuantStore::PackedU32(packed_dim) = &mut scheme.store + { + let rank = tensor.meta.num_dims(); + let new_pos = axes + .iter() + .position(|axis| *axis == rank - *packed_dim - 1) + .unwrap_or(0); + *packed_dim = rank - new_pos - 1; + } + + tensor +} + +/// Permute a tensor's dimensions from NCHW to NHWC, or the N-dimensional equivalent +pub fn permute_nchw_to_nhwc(tensor: CubeTensor) -> CubeTensor { + let rank = tensor.meta.num_dims(); + let c_dim = 1; + + let mut dims = vec![0]; + dims.extend(2..rank); + dims.push(c_dim); + + permute(tensor, &dims) +} + +/// Permute a shape's dimensions from NCHW to NHWC, or the N-dimensional equivalent +pub fn permute_nchw_to_nhwc_shape(shape: Shape) -> Shape { + let rank = shape.num_dims(); + let c_dim = 1; + + let mut dims = vec![0]; + dims.extend(2..rank); + dims.push(c_dim); + + shape.permuted(&dims).expect("Shape permute should succeed") +} + +/// Permute a tensor's dimensions from NHWC to NCHW, or the N-dimensional equivalent +pub fn permute_nhwc_to_nchw(tensor: CubeTensor) -> CubeTensor { + let rank = tensor.meta.num_dims(); + let c_dim = rank - 1; + + let mut dims = vec![0]; + dims.push(c_dim); + dims.extend(1..c_dim); + + permute(tensor, &dims) +} + +/// Permute a shape's dimensions from NHWC to NCHW, or the N-dimensional equivalent +pub fn permute_nhwc_to_nchw_shape(shape: Shape) -> Shape { + let rank = shape.num_dims(); + let c_dim = rank - 1; + + let mut dims = vec![0]; + dims.push(c_dim); + dims.extend(1..c_dim); + + shape.permuted(&dims).expect("Shape permute should succeed") +} + +pub(crate) fn expand(tensor: CubeTensor, target_shape: Shape) -> CubeTensor { + let ndims_in = tensor.meta.shape().num_dims(); + let ndims_out = target_shape.num_dims(); + + // Initialize new strides with zeros + let mut new_strides = strides![0usize; ndims_out]; + + // Calculate the difference in dimensions + let dim_diff = ndims_out.saturating_sub(ndims_in); + + // Compare dimensions from the end, setting strides for matching dimensions or broadcasted ones + let mut tensor_dim_iter = tensor.meta.shape().iter().rev(); + for i in (0..ndims_out).rev() { + if i >= dim_diff { + if let Some(&tensor_dim) = tensor_dim_iter.next() { + if tensor_dim == target_shape[i] || tensor_dim == 1 { + // Copy stride for non-broadcast dimensions or set to 0 for broadcast ones + new_strides[i] = if tensor_dim == target_shape[i] { + tensor.meta.strides()[i - dim_diff] + } else { + 0 + }; + } else { + // Error handling: Dimension mismatch for broadcasting + panic!( + "Dimension mismatch: cannot broadcast dimension {tensor_dim} of tensor to target shape" + ); + } + } else { + // If the input tensor has fewer dimensions, treat missing dimensions as 1 + // and set stride to 0 (broadcasting) + new_strides[i] = 0; + } + } else { + // For extra dimensions in the target shape, set stride to 0 (broadcasting) + new_strides[i] = 0; + } + } + + // Extra check to ensure block scales must be properly handled once they're added + if tensor.qparams.is_some() { + match tensor.scheme().level { + QuantLevel::Tensor => {} + QuantLevel::Block(_) => todo!(), + } + } + + CubeTensor { + client: tensor.client.clone(), + device: tensor.device.clone(), + meta: Box::new(Metadata::new(target_shape, new_strides)), + handle: tensor.handle.clone(), + dtype: tensor.dtype, + qparams: tensor.qparams.clone(), + } +} + +/// Reshape a jit tensor to a new shape +pub fn reshape(mut tensor: CubeTensor, shape: Shape) -> CubeTensor { + let analysis = reshape_action(tensor.meta.shape(), tensor.meta.strides(), &shape); + + match analysis { + ReshapeAction::UpdateStrides { strides } => { + *tensor.meta = Metadata::new(shape, strides); + return tensor; + } + ReshapeAction::NoChange => return tensor, + ReshapeAction::Recompute => (), + } + + let out = empty_device_dtype( + tensor.client.clone(), + tensor.device.clone(), + shape, + tensor.dtype, + ); + + cubecl::std::tensor::copy_into( + &out.client, + tensor.binding(), + out.clone().binding(), + dtype_to_storage_type(out.dtype), + ); + + out +} + +/// Reshape a jit tensor to a new shape +pub fn q_reshape(mut tensor: CubeTensor, shape: Shape) -> CubeTensor { + let scheme = tensor.scheme(); + let curr_shape = tensor.meta.shape(); + + let shape_values = match scheme.store { + QuantStore::Native => shape.clone(), + QuantStore::PackedNative(packed_dim) | QuantStore::PackedU32(packed_dim) => { + let rank = shape.num_dims(); + let mut shape = shape.clone(); + let packed_d = rank - packed_dim - 1; + let num_quants = scheme.num_quants(); + + if !shape[packed_d].is_multiple_of(num_quants) { + unimplemented!( + "Cannot reshape packed tensor: inner dimension {} is not aligned with packing factor {num_quants}", + shape[packed_d] + ); + } + + shape[packed_d] = shape[packed_d].div_ceil(num_quants); + shape + } + }; + + let (values, scales) = tensor.quantized_handles().unwrap(); + let analysis_values = reshape_analysis( + values.meta.shape(), + Some(values.meta.strides()), + &shape_values, + ); + let action_values = + analysis_values.action(values.meta.shape(), values.meta.strides(), &shape_values); + + let n_new_dims = shape.num_dims().saturating_sub(curr_shape.num_dims()); + let is_unsqueeze = n_new_dims > 0 && shape[n_new_dims..] == **curr_shape; + + if !is_unsqueeze + && matches!( + scheme.value, + QuantValue::Q4S | QuantValue::Q4F | QuantValue::Q2S | QuantValue::Q2F + ) + { + // FIXME + todo!("Reshape with sub-byte values is not supported") + } + + // Check valid reshapes + if let ReshapeAction::UpdateStrides { .. } = &action_values { + match analysis_values { + ReshapeAnalysis::IsContiguous => { + if let QuantLevel::Block(block_size) = scheme.level + && block_size.len() > 1 + && !is_unsqueeze + { + // General reshape (e.g. [32, 4] -> [16, 8]): only valid if + // reshaped dimension is aligned with the block boundaries. + unimplemented!("Reshape of ND block-quantized tensor is not yet supported."); + } + } + ReshapeAnalysis::Broadcasted => {} // only preprends unit dims + ReshapeAnalysis::Split => { + if let QuantLevel::Block(block_size) = scheme.level + && block_size.len() > 1 + { + // Split reshape (e.g. [32, 4] -> [32, 2, 2]): only valid if + // reshaped dimension is aligned with the block boundaries. + unimplemented!( + "Split reshape of ND block-quantized tensor is not yet supported." + ); + } + } + other => unreachable!("Reshape analysis {other:?} should not update strides."), + } + } + + let shape_last = *shape.last().unwrap(); + + let shape_scales = match scheme.level { + QuantLevel::Tensor => scales.meta.shape().clone(), // always [1], invariant under reshape + QuantLevel::Block(block_size) + if block_size.len() == 1 && shape_last < (block_size[0] as usize) => + { + // If the new last dimension is smaller than the block size, + // it means a single block now spans across multiple rows. + if scales.meta.shape().num_elements() > 1 { + unimplemented!("Reshape would split a block across multiple rows."); + } + // Exception: allow if there is exactly 1 block total (essentially per-tensor quantization) + scales.meta.shape().clone() + } + QuantLevel::Block(_) => { + // ND blocks: derive scales shape from the new tensor shape + params_shape(&shape, scheme.level) + } + }; + + let action_scales = reshape_action(scales.meta.shape(), scales.meta.strides(), &shape_scales); + + match (action_values, action_scales) { + ( + ReshapeAction::UpdateStrides { strides }, + ReshapeAction::UpdateStrides { + strides: scales_strides, + }, + ) => { + let qparams = tensor.qparams.as_mut().unwrap(); + + *tensor.meta = Metadata::new(shape, strides); + qparams.scales.metadata = Metadata::new(shape_scales, scales_strides); + } + (ReshapeAction::UpdateStrides { strides }, ReshapeAction::NoChange) => { + *tensor.meta = Metadata::new(shape, strides); + } + ( + ReshapeAction::NoChange, + ReshapeAction::UpdateStrides { + strides: scales_strides, + }, + ) => { + let qparams = tensor.qparams.as_mut().unwrap(); + + qparams.scales.metadata = Metadata::new(shape_scales, scales_strides); + } + // Any action to recompute + (ReshapeAction::Recompute, _) | (_, ReshapeAction::Recompute) => { + if let QuantLevel::Block(_) = scheme.level + && shape_scales.num_elements() > 1 + { + // Original block boundaries no longer align with the layout, would have to be recomputed + unimplemented!( + "Cannot reshape a block-quantized tensor when the reshape requires recomputing the buffer." + ); + } + + tensor = kernel::into_contiguous(tensor); + *tensor.meta = Metadata::new(shape, contiguous_strides(&shape_values)); + + let qparams = tensor.qparams.as_mut().unwrap(); + + let strides = contiguous_strides(&shape_scales); + qparams.scales.metadata = Metadata::new(shape_scales, strides); + } + (ReshapeAction::NoChange, ReshapeAction::NoChange) => {} + } + + tensor +} + +pub(crate) fn max_vector_size(tensor: &CubeTensor) -> VectorSize { + tensor_vector_size_parallel( + tensor.client.io_optimized_vector_sizes(tensor.dtype.size()), + tensor.meta.shape(), + tensor.meta.strides(), + tensor.meta.num_dims() - 1, + ) +} + +pub(crate) fn max_vector_size_many( + tensors: &[&CubeTensor], + axis: usize, +) -> VectorSize { + let vec = tensors + .iter() + .map(|tensor| { + tensor_vector_size_parallel( + tensor.client.io_optimized_vector_sizes(tensor.dtype.size()), + tensor.meta.shape(), + tensor.meta.strides(), + axis, + ) + }) + .min(); + + vec.unwrap_or(0) +} + +/// Unfold windows along a dimension. +/// +/// Returns a view of the tensor with all complete windows of size `size` in dimension `dim`; +/// where windows are advanced by `step` at each index. +/// +/// The number of windows is `max(0, (shape[dim] - size).ceil_div(step))`. +/// +/// The new view will have the unfolded dimension replaced by two dimensions; +/// one in the position of the original dimension, with size equal to the number of windows, +/// and one appended to the right-most position, with size equal to `size`. +/// +/// # Arguments +/// +/// * `tensor` - The input tensor to unfold; of shape ``[pre=..., dim shape, post=...]`` +/// * `dim` - the dimension to unfold. +/// * `size` - the size of each unfolded window. +/// * `step` - the step between each window. +/// +/// # Returns +/// +/// A tensor view with the shape ``[pre=..., windows, post=..., size]``. +pub fn unfold( + tensor: CubeTensor, + dim: usize, + size: usize, + step: usize, +) -> CubeTensor { + let shape = calculate_unfold_shape(tensor.shape(), dim, size, step); + + let d_stride = tensor.meta.strides()[dim]; + let mut strides = tensor.meta.strides.clone(); + strides[dim] = step * d_stride; + strides.push(d_stride); + + CubeTensor { + meta: Box::new(Metadata::new(shape, strides)), + client: tensor.client.clone(), + handle: tensor.handle.clone(), + device: tensor.device.clone(), + dtype: tensor.dtype, + qparams: tensor.qparams.clone(), + } +} diff --git a/crates/burn-cubecl/src/ops/bool_tensor.rs b/crates/burn-cubecl/src/ops/bool_tensor.rs new file mode 100644 index 0000000..df540a2 --- /dev/null +++ b/crates/burn-cubecl/src/ops/bool_tensor.rs @@ -0,0 +1,241 @@ +use crate::{ + CubeBackend, CubeRuntime, + element::BoolElement, + kernel::{self, AndOp, OrOp}, + tensor::CubeTensor, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{ + ExecutionError, Slice, + ops::BoolTensorOps, + tensor::{BoolTensor, Device, FloatTensor, IntTensor}, +}; +use burn_backend::{Scalar, Shape, TensorData}; +use burn_std::{BoolDType, BoolStore, DType, FloatDType, IntDType}; +use cubecl::prelude::InputScalar; +use cubek::reduce::components::instructions::ReduceOperationConfig; +use std::ops::Range; + +use super::{expand, numeric, permute, unfold}; + +/// The boolean storage of a cubecl bool tensor. Cubecl backends never use +/// native bool (see `CubeBackend::supports_dtype`), so it is always `U8`/`U32`. +fn bool_store(tensor: &CubeTensor) -> BoolDType { + match tensor.dtype { + DType::Bool(store) => store, + other => unreachable!("cubecl bool tensors are always Bool(_): {other:?}"), + } +} + +impl BoolTensorOps for CubeBackend { + fn bool_empty(shape: Shape, device: &Device, dtype: BoolDType) -> BoolTensor { + super::empty(shape, device, dtype.into()) + } + + fn bool_zeros(shape: Shape, device: &Device, dtype: BoolDType) -> BoolTensor { + numeric::zeros(device.clone(), shape, dtype.into()) + } + + fn bool_ones(shape: Shape, device: &Device, dtype: BoolDType) -> BoolTensor { + numeric::ones(device.clone(), shape, dtype.into()) + } + + async fn bool_into_data(tensor: BoolTensor) -> Result { + super::into_data(tensor).await + } + + fn bool_from_data(data: TensorData, device: &Device) -> BoolTensor { + if !matches!( + data.dtype, + DType::Bool(BoolStore::U8) | DType::Bool(BoolStore::U32) + ) { + unimplemented!("Unsupported dtype for `bool_from_data` {:?}", data.dtype); + } + super::from_data(data, device) + } + + fn bool_into_int(tensor: BoolTensor, out_dtype: IntDType) -> IntTensor { + kernel::bool_cast(tensor, out_dtype.into()) + } + + fn bool_to_device(tensor: BoolTensor, device: &Device) -> BoolTensor { + super::to_device(tensor, device) + } + + fn bool_reshape(tensor: BoolTensor, shape: Shape) -> BoolTensor { + super::reshape(tensor, shape) + } + + fn bool_slice(tensor: BoolTensor, slices: &[Slice]) -> BoolTensor { + // Check if all steps are 1 + let all_steps_one = slices.iter().all(|info| info.step == 1); + + if all_steps_one { + // Use optimized slice for step=1 + let simple_ranges: Vec> = slices + .iter() + .enumerate() + .map(|(i, slice)| slice.to_range(tensor.meta.shape()[i])) + .collect(); + + kernel::slice(tensor, &simple_ranges) + } else { + // Use slice with steps kernel + kernel::slice_with_steps(tensor, slices) + } + } + + fn bool_slice_assign( + tensor: BoolTensor, + ranges: &[Slice], + value: BoolTensor, + ) -> BoolTensor { + kernel::slice_assign(tensor, ranges, value) + } + + fn bool_equal(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + let dtype = lhs.dtype; + kernel::equal(lhs, rhs, dtype) + } + + fn bool_not(tensor: BoolTensor) -> BoolTensor { + let dtype = tensor.dtype; + let storage = dtype_to_storage_type(dtype); + let scalar = match dtype { + DType::Bool(BoolStore::U32) => InputScalar::new(u32::false_val(), storage), + DType::Bool(BoolStore::U8) => InputScalar::new(u8::false_val(), storage), + other => unimplemented!("Unsupported dtype for `bool_from_data` {other:?}"), + }; + kernel::equal_elem(tensor, scalar, dtype) + } + + fn bool_and(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + kernel::launch_binop::(lhs, rhs) + } + + fn bool_or(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + kernel::launch_binop::(lhs, rhs) + } + + fn bool_any(tensor: BoolTensor) -> BoolTensor { + let store = bool_store(&tensor); + kernel::reduce::reduce_logical(tensor, None, ReduceOperationConfig::Any, store) + } + + fn bool_any_dim(tensor: BoolTensor, dim: usize) -> BoolTensor { + let store = bool_store(&tensor); + kernel::reduce::reduce_logical(tensor, Some(dim), ReduceOperationConfig::Any, store) + } + + fn bool_all(tensor: BoolTensor) -> BoolTensor { + let store = bool_store(&tensor); + kernel::reduce::reduce_logical(tensor, None, ReduceOperationConfig::All, store) + } + + fn bool_all_dim(tensor: BoolTensor, dim: usize) -> BoolTensor { + let store = bool_store(&tensor); + kernel::reduce::reduce_logical(tensor, Some(dim), ReduceOperationConfig::All, store) + } + + fn bool_into_float(tensor: BoolTensor, out_dtype: FloatDType) -> FloatTensor { + kernel::bool_cast(tensor, out_dtype.into()) + } + + fn bool_swap_dims(mut tensor: BoolTensor, dim1: usize, dim2: usize) -> BoolTensor { + tensor.meta.swap(dim1, dim2); + + tensor + } + + fn bool_repeat_dim(tensor: BoolTensor, dim: usize, times: usize) -> BoolTensor { + kernel::repeat_dim(tensor, dim, times) + } + + fn bool_permute(tensor: BoolTensor, axes: &[usize]) -> BoolTensor { + permute(tensor, axes) + } + + fn bool_expand(tensor: BoolTensor, shape: Shape) -> BoolTensor { + expand(tensor, shape) + } + + fn bool_select( + tensor: BoolTensor, + dim: usize, + indices: IntTensor, + ) -> BoolTensor { + kernel::select(tensor, dim, indices) + } + + fn bool_select_or( + tensor: BoolTensor, + dim: usize, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor { + kernel::select_assign(tensor, dim, indices, value, true) + } + + fn bool_flip(tensor: BoolTensor, axes: &[usize]) -> BoolTensor { + let dtype = tensor.dtype; + kernel::flip(tensor, axes, dtype) + } + + fn bool_unfold( + tensor: FloatTensor, + dim: usize, + size: usize, + step: usize, + ) -> FloatTensor { + unfold(tensor, dim, size, step) + } + + fn bool_mask_where( + tensor: BoolTensor, + mask: BoolTensor, + value: BoolTensor, + ) -> BoolTensor { + let dtype = tensor.dtype; + kernel::mask_where_auto(tensor, mask, value, dtype) + } + + fn bool_mask_fill( + tensor: BoolTensor, + mask: BoolTensor, + value: Scalar, + ) -> BoolTensor { + let dtype = tensor.dtype; + kernel::mask_fill_auto( + tensor, + mask, + InputScalar::new(value, dtype_to_storage_type(dtype)), + dtype, + ) + } + + fn bool_gather( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + ) -> BoolTensor { + kernel::gather(dim, tensor, indices) + } + + fn bool_scatter_or( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor { + kernel::scatter(dim, tensor, indices, value, true) + } + + fn bool_equal_elem(lhs: BoolTensor, rhs: Scalar) -> BoolTensor { + let dtype = lhs.dtype; + kernel::equal_elem( + lhs, + InputScalar::new(rhs, dtype_to_storage_type(dtype)), + dtype, + ) + } +} diff --git a/crates/burn-cubecl/src/ops/distributed.rs b/crates/burn-cubecl/src/ops/distributed.rs new file mode 100644 index 0000000..9cd521a --- /dev/null +++ b/crates/burn-cubecl/src/ops/distributed.rs @@ -0,0 +1,56 @@ +use burn_backend::distributed::DistributedOps; + +use crate::{CubeBackend, CubeRuntime}; + +#[cfg(feature = "std")] +use crate::ops::numeric::{self, zeros_client}; +#[cfg(feature = "std")] +use burn_backend::{ + DeviceId, TensorMetadata, + cubecl::dtype_to_elem_type, + distributed::{CollectiveTensor, ReduceOperation}, + tensor::{Device, FloatTensor}, +}; + +impl DistributedOps for CubeBackend { + #[cfg(feature = "std")] + fn all_reduce( + tensor: FloatTensor, + op: ReduceOperation, + device_ids: Vec, + ) -> CollectiveTensor { + let device = &tensor.device.clone(); + let out_tensor = if tensor.handle.can_mut() && tensor.is_contiguous() { + tensor + } else { + let zeros_tensor = zeros_client::( + tensor.client.clone(), + device.clone(), + tensor.shape(), + tensor.dtype(), + ); + numeric::add(zeros_tensor, tensor) + }; + + let op = match op { + ReduceOperation::Sum => cubecl::server::ReduceOperation::Sum, + ReduceOperation::Mean => cubecl::server::ReduceOperation::Mean, + }; + + let mut client = R::client(device); + client.all_reduce( + out_tensor.handle.clone(), + out_tensor.handle.clone(), + dtype_to_elem_type(out_tensor.dtype), + device_ids.clone(), + op, + ); + CollectiveTensor::new(out_tensor) + } + + #[cfg(feature = "std")] + fn sync_collective(device: &Device) { + let client = R::client(device); + client.sync_collective(); + } +} diff --git a/crates/burn-cubecl/src/ops/int_tensor.rs b/crates/burn-cubecl/src/ops/int_tensor.rs new file mode 100644 index 0000000..a1ecfd1 --- /dev/null +++ b/crates/burn-cubecl/src/ops/int_tensor.rs @@ -0,0 +1,663 @@ +use self::unary_basic_int::BasicIntUnaryKind; +use burn_backend::cubecl::dtype_to_storage_type; + +use super::{expand, numeric, permute, unfold}; +use crate::kernel::prng::{random_bernoulli, random_normal, random_uniform}; +use crate::kernel::{ + BitwiseShlOp, BitwiseShrOp, NumericUnaryOp, NumericUnaryOpFamily, launch_binop_int, + launch_scalar_binop_int, launch_unary_numeric, reduce, unary_basic_int, +}; +use crate::{ + CubeBackend, CubeRuntime, + kernel::{ + self, + matmul::{MatmulStrategy, matmul}, + }, +}; +use burn_backend::tensor::{BoolTensor, Device, FloatTensor, IntTensor}; +use burn_backend::{DType, IntDType, Slice, ops::IntTensorOps}; +use burn_backend::{Distribution, ElementConversion, Shape, TensorData, get_device_settings}; +use burn_backend::{ExecutionError, Scalar}; +use burn_std::{BoolDType, FloatDType}; +use cubecl::frontend::Numeric; +use cubecl::prelude::*; +use cubek::reduce::components::instructions::ReduceOperationConfig; +use std::ops::Range; + +impl IntTensorOps for CubeBackend { + fn int_empty(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + let dtype = dtype.into(); + super::empty(shape, device, dtype) + } + + async fn int_into_data(tensor: IntTensor) -> Result { + super::into_data(tensor).await + } + + fn int_from_data(data: TensorData, device: &Device) -> IntTensor { + match data.dtype { + DType::I64 + | DType::I32 + | DType::I16 + | DType::I8 + | DType::U64 + | DType::U32 + | DType::U16 + | DType::U8 => super::from_data(data, device), + _ => unimplemented!("Unsupported dtype for `int_from_data`"), + } + } + + fn int_to_device(tensor: IntTensor, device: &Device) -> IntTensor { + super::to_device(tensor, device) + } + + fn int_reshape(tensor: IntTensor, shape: Shape) -> IntTensor { + super::reshape(tensor, shape) + } + + fn int_slice(tensor: IntTensor, slices: &[Slice]) -> IntTensor { + // Check if all steps are 1 + let all_steps_one = slices.iter().all(|info| info.step == 1); + + if all_steps_one { + // Use optimized slice for step=1 + let simple_ranges: Vec> = slices + .iter() + .enumerate() + .map(|(i, slice)| slice.to_range(tensor.meta.shape()[i])) + .collect(); + + kernel::slice(tensor, &simple_ranges) + } else { + // Use slice with steps kernel + kernel::slice_with_steps(tensor, slices) + } + } + + fn int_slice_assign( + tensor: IntTensor, + ranges: &[Slice], + value: IntTensor, + ) -> IntTensor { + kernel::slice_assign(tensor, ranges, value) + } + + fn int_matmul(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + let dtype = lhs.dtype; + matmul(lhs, rhs, None, MatmulStrategy::default(), dtype).unwrap() + } + + fn int_mask_where( + tensor: IntTensor, + mask: BoolTensor, + value: IntTensor, + ) -> IntTensor { + let bool_dtype = mask.dtype; + kernel::mask_where_auto(tensor, mask, value, bool_dtype) + } + + fn int_mask_fill( + tensor: IntTensor, + mask: BoolTensor, + value: Scalar, + ) -> IntTensor { + let dtype = tensor.dtype; + let bool_dtype = mask.dtype; + kernel::mask_fill_auto( + tensor, + mask, + InputScalar::new(value, dtype_to_storage_type(dtype)), + bool_dtype, + ) + } + + fn int_gather( + dim: usize, + tensor: IntTensor, + indices: IntTensor, + ) -> IntTensor { + kernel::gather(dim, tensor, indices) + } + + fn int_scatter_add( + dim: usize, + tensor: IntTensor, + indices: IntTensor, + value: IntTensor, + ) -> IntTensor { + kernel::scatter(dim, tensor, indices, value, false) + } + + fn int_scatter_nd( + data: IntTensor, + indices: IntTensor, + values: IntTensor, + reduction: burn_backend::tensor::IndexingUpdateOp, + ) -> IntTensor { + kernel::scatter_nd(data, indices, values, reduction) + } + + fn int_gather_nd(data: IntTensor, indices: IntTensor) -> IntTensor { + kernel::gather_nd(data, indices) + } + + fn int_select( + tensor: IntTensor, + dim: usize, + indices: IntTensor, + ) -> IntTensor { + kernel::select(tensor, dim, indices) + } + + fn int_select_add( + tensor: IntTensor, + dim: usize, + indices: IntTensor, + value: IntTensor, + ) -> IntTensor { + kernel::select_assign(tensor, dim, indices, value, false) + } + + fn int_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + kernel::equal(lhs, rhs, out_dtype.into()) + } + + fn int_equal_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor { + let dtype = lhs.dtype; + kernel::equal_elem( + lhs, + InputScalar::new(rhs, dtype_to_storage_type(dtype)), + out_dtype.into(), + ) + } + + fn int_greater( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + kernel::greater(lhs, rhs, out_dtype.into()) + } + + fn int_greater_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let dtype = lhs.dtype; + kernel::greater_elem( + lhs, + InputScalar::new(rhs, dtype_to_storage_type(dtype)), + out_dtype.into(), + ) + } + + fn int_greater_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + kernel::greater_equal(lhs, rhs, out_dtype.into()) + } + + fn int_greater_equal_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let dtype = lhs.dtype; + kernel::greater_equal_elem( + lhs, + InputScalar::new(rhs, dtype_to_storage_type(dtype)), + out_dtype.into(), + ) + } + + fn int_lower( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + kernel::lower(lhs, rhs, out_dtype.into()) + } + + fn int_lower_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor { + let dtype = lhs.dtype; + kernel::lower_elem( + lhs, + InputScalar::new(rhs, dtype_to_storage_type(dtype)), + out_dtype.into(), + ) + } + + fn int_lower_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + kernel::lower_equal(lhs, rhs, out_dtype.into()) + } + + fn int_lower_equal_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let dtype = lhs.dtype; + kernel::lower_equal_elem( + lhs, + InputScalar::new(rhs, dtype_to_storage_type(dtype)), + out_dtype.into(), + ) + } + + fn int_add(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + numeric::add(lhs, rhs) + } + + fn int_add_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let dtype = lhs.dtype; + numeric::add_scalar(lhs, InputScalar::new(rhs, dtype_to_storage_type(dtype))) + } + + fn int_sub(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + numeric::sub(lhs, rhs) + } + + fn int_sub_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let dtype = lhs.dtype; + numeric::sub_scalar(lhs, InputScalar::new(rhs, dtype_to_storage_type(dtype))) + } + + fn int_mul(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + numeric::mul(lhs, rhs) + } + + fn int_mul_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let dtype = lhs.dtype; + numeric::mul_scalar(lhs, InputScalar::new(rhs, dtype_to_storage_type(dtype))) + } + + fn int_div(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + numeric::div(lhs, rhs) + } + + fn int_div_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let dtype = lhs.dtype; + numeric::div_scalar(lhs, InputScalar::new(rhs, dtype_to_storage_type(dtype))) + } + + fn int_remainder(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + numeric::remainder(lhs, rhs) + } + + fn int_remainder_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let dtype = lhs.dtype; + numeric::remainder_scalar(lhs, InputScalar::new(rhs, dtype_to_storage_type(dtype))) + } + + fn int_zeros(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + let dtype = dtype.into(); + numeric::zeros(device.clone(), shape, dtype) + } + + fn int_ones(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + let dtype = dtype.into(); + numeric::ones(device.clone(), shape, dtype) + } + + fn int_full( + shape: Shape, + fill_value: Scalar, + device: &Device, + dtype: IntDType, + ) -> IntTensor { + let dtype: DType = dtype.into(); + let client = R::client(device); + numeric::full_device_dtype( + client, + shape, + device.clone(), + InputScalar::new(fill_value, dtype_to_storage_type(dtype)), + dtype, + ) + } + + fn int_sum(tensor: IntTensor) -> IntTensor { + reduce::sum_fallback(tensor, Default::default()).unwrap() + } + + fn int_sum_dim(tensor: IntTensor, dim: usize) -> IntTensor { + reduce::reduce_dim( + tensor, + None, + dim, + Default::default(), + ReduceOperationConfig::Sum, + ) + .unwrap() + } + + fn int_any(tensor: IntTensor, out_dtype: BoolDType) -> BoolTensor { + reduce::reduce_logical(tensor, None, ReduceOperationConfig::Any, out_dtype) + } + + fn int_any_dim(tensor: IntTensor, dim: usize, out_dtype: BoolDType) -> BoolTensor { + reduce::reduce_logical(tensor, Some(dim), ReduceOperationConfig::Any, out_dtype) + } + + fn int_all(tensor: IntTensor, out_dtype: BoolDType) -> BoolTensor { + reduce::reduce_logical(tensor, None, ReduceOperationConfig::All, out_dtype) + } + + fn int_all_dim(tensor: IntTensor, dim: usize, out_dtype: BoolDType) -> BoolTensor { + reduce::reduce_logical(tensor, Some(dim), ReduceOperationConfig::All, out_dtype) + } + + fn int_prod(tensor: IntTensor) -> IntTensor { + reduce::reduce( + tensor, + None, + Default::default(), + ReduceOperationConfig::Prod, + ) + .unwrap() + } + + fn int_prod_dim(tensor: IntTensor, dim: usize) -> IntTensor { + reduce::reduce_dim( + tensor, + None, + dim, + Default::default(), + ReduceOperationConfig::Prod, + ) + .unwrap() + } + + fn int_max(tensor: IntTensor) -> IntTensor { + reduce::reduce(tensor, None, Default::default(), ReduceOperationConfig::Max).unwrap() + } + + fn int_max_dim(tensor: IntTensor, dim: usize) -> IntTensor { + reduce::reduce_dim( + tensor, + None, + dim, + Default::default(), + ReduceOperationConfig::Max, + ) + .unwrap() + } + + fn int_topk(tensor: IntTensor, dim: usize, k: usize) -> IntTensor { + reduce::reduce_dim( + tensor, + None, + dim, + Default::default(), + ReduceOperationConfig::TopK(k), + ) + .unwrap() + } + + fn int_max_abs(tensor: IntTensor) -> IntTensor { + reduce::reduce( + tensor, + None, + Default::default(), + ReduceOperationConfig::MaxAbs, + ) + .unwrap() + } + + fn int_max_abs_dim(tensor: IntTensor, dim: usize) -> IntTensor { + reduce::reduce_dim( + tensor, + None, + dim, + Default::default(), + ReduceOperationConfig::MaxAbs, + ) + .unwrap() + } + + fn int_min(tensor: IntTensor) -> IntTensor { + reduce::reduce(tensor, None, Default::default(), ReduceOperationConfig::Min).unwrap() + } + + fn int_min_dim(tensor: IntTensor, dim: usize) -> IntTensor { + reduce::reduce_dim( + tensor, + None, + dim, + Default::default(), + ReduceOperationConfig::Min, + ) + .unwrap() + } + + fn int_mean_dim(tensor: IntTensor, dim: usize) -> IntTensor { + reduce::reduce_dim( + tensor, + None, + dim, + Default::default(), + ReduceOperationConfig::Mean, + ) + .unwrap() + } + + fn int_cumsum(tensor: IntTensor, dim: usize) -> IntTensor { + numeric::cumsum(tensor, dim) + } + + fn int_cumprod(tensor: IntTensor, dim: usize) -> IntTensor { + numeric::cumprod(tensor, dim) + } + + fn int_cummin(tensor: IntTensor, dim: usize) -> IntTensor { + numeric::cummin(tensor, dim) + } + + fn int_cummax(tensor: IntTensor, dim: usize) -> IntTensor { + numeric::cummax(tensor, dim) + } + + fn int_argmax(tensor: IntTensor, dim: usize) -> IntTensor { + let dtype = tensor.dtype; + reduce::reduce_dim( + tensor, + Some(dtype), + dim, + Default::default(), + ReduceOperationConfig::ArgMax, + ) + .unwrap() + } + + fn int_argtopk(tensor: IntTensor, dim: usize, k: usize) -> IntTensor { + let dtype = tensor.dtype; + reduce::reduce_dim( + tensor, + Some(dtype), + dim, + Default::default(), + ReduceOperationConfig::ArgTopK(k), + ) + .unwrap() + } + + fn int_argmin(tensor: IntTensor, dim: usize) -> IntTensor { + let dtype = tensor.dtype; + reduce::reduce_dim( + tensor, + Some(dtype), + dim, + Default::default(), + ReduceOperationConfig::ArgMin, + ) + .unwrap() + } + + fn int_clamp(tensor: IntTensor, min: Scalar, max: Scalar) -> IntTensor { + let dtype = tensor.dtype; + kernel::clamp( + tensor, + InputScalar::new(min, dtype_to_storage_type(dtype)), + InputScalar::new(max, dtype_to_storage_type(dtype)), + ) + } + + fn int_abs(tensor: IntTensor) -> IntTensor { + struct Abs; + + #[cube] + impl NumericUnaryOp for Abs { + type Options = (); + + fn execute(input: Vector, _options: &Self::Options) -> Vector { + Vector::abs(input) + } + } + + impl NumericUnaryOpFamily for Abs { + type Options = (); + type Unary = Self; + } + + launch_unary_numeric::(tensor, |_| ()) + } + + fn int_sign(tensor: IntTensor) -> IntTensor { + unary_basic_int::launch::(tensor, |_| BasicIntUnaryKind::Sign) + } + + fn int_into_float(tensor: IntTensor, out_dtype: FloatDType) -> FloatTensor { + kernel::cast(tensor, out_dtype.into()) + } + + fn int_swap_dims(mut tensor: IntTensor, dim1: usize, dim2: usize) -> IntTensor { + tensor.meta.swap(dim1, dim2); + + tensor + } + + fn int_repeat_dim(tensor: IntTensor, dim: usize, times: usize) -> IntTensor { + kernel::repeat_dim(tensor, dim, times) + } + + fn int_random( + shape: Shape, + distribution: Distribution, + device: &Device, + dtype: IntDType, + ) -> IntTensor { + let dtype = dtype.into(); + match distribution { + Distribution::Default => random_uniform(shape, device, 0., 255., dtype), + Distribution::Uniform(low, high) => { + random_uniform(shape, device, low.elem(), high.elem(), dtype) + } + Distribution::Bernoulli(prob) => random_bernoulli(shape, device, prob as f32, dtype), + Distribution::Normal(mean, std) => { + random_normal(shape, device, mean.elem(), std.elem(), dtype) + } + } + } + + fn int_permute(tensor: IntTensor, axes: &[usize]) -> IntTensor { + permute(tensor, axes) + } + + fn int_expand(tensor: IntTensor, shape: Shape) -> IntTensor { + expand(tensor, shape) + } + + fn int_flip(tensor: IntTensor, axes: &[usize]) -> IntTensor { + let bool_dtype = get_device_settings::(&tensor.device).bool_dtype; + kernel::flip(tensor, axes, bool_dtype.into()) + } + + fn bitwise_and(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + numeric::bitwise_and(lhs, rhs) + } + + fn bitwise_and_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let dtype = lhs.dtype; + numeric::bitwise_and_scalar(lhs, InputScalar::new(rhs, dtype_to_storage_type(dtype))) + } + + fn bitwise_or(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + numeric::bitwise_or(lhs, rhs) + } + + fn bitwise_or_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let dtype = lhs.dtype; + numeric::bitwise_or_scalar(lhs, InputScalar::new(rhs, dtype_to_storage_type(dtype))) + } + + fn bitwise_xor(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + numeric::bitwise_xor(lhs, rhs) + } + + fn bitwise_xor_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let dtype = lhs.dtype; + numeric::bitwise_xor_scalar(lhs, InputScalar::new(rhs, dtype_to_storage_type(dtype))) + } + + fn bitwise_not(tensor: IntTensor) -> IntTensor { + unary_basic_int::launch::(tensor, |_| BasicIntUnaryKind::BitwiseNot) + } + + fn bitwise_left_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + launch_binop_int::(lhs, rhs) + } + + fn bitwise_left_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let dtype = lhs.dtype; + launch_scalar_binop_int::( + lhs, + InputScalar::new(rhs, dtype_to_storage_type(dtype)), + ) + } + + fn bitwise_right_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + launch_binop_int::(lhs, rhs) + } + + fn bitwise_right_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let dtype = lhs.dtype; + launch_scalar_binop_int::( + lhs, + InputScalar::new(rhs, dtype_to_storage_type(dtype)), + ) + } + + fn int_cast(tensor: IntTensor, dtype: IntDType) -> IntTensor { + kernel::cast(tensor, dtype.into()) + } + + fn int_unfold( + tensor: FloatTensor, + dim: usize, + size: usize, + step: usize, + ) -> FloatTensor { + unfold(tensor, dim, size, step) + } + + // TODO + // fn int_powi(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + // todo!() + // } + + // fn int_powi_scalar_impl(lhs: IntTensor, rhs: Scalar) -> IntTensor { + // todo!() + // } +} diff --git a/crates/burn-cubecl/src/ops/mod.rs b/crates/burn-cubecl/src/ops/mod.rs new file mode 100644 index 0000000..bf6f213 --- /dev/null +++ b/crates/burn-cubecl/src/ops/mod.rs @@ -0,0 +1,16 @@ +mod activation; +mod bool_tensor; +mod int_tensor; +mod module; +mod qtensor; +mod tensor; +mod transaction; + +pub(crate) mod distributed; + +pub(crate) mod base; +pub use base::*; +pub use qtensor::*; + +/// Numeric utility functions for jit backends +pub mod numeric; diff --git a/crates/burn-cubecl/src/ops/module.rs b/crates/burn-cubecl/src/ops/module.rs new file mode 100644 index 0000000..4d333c9 --- /dev/null +++ b/crates/burn-cubecl/src/ops/module.rs @@ -0,0 +1,399 @@ +use crate::{ + CubeBackend, CubeRuntime, + kernel::{self, conv::ConvTranspose2dStrategy}, +}; +use burn_backend::tensor::{BoolTensor, FloatTensor, IntTensor}; +use burn_backend::{ + TensorMetadata, + ops::{ + AttentionModuleOptions, ConvOptions, ConvTransposeOptions, DeformConv2dBackward, + DeformConvOptions, InterpolateOptions, MaxPool2dBackward, MaxPool2dWithIndices, ModuleOps, + }, +}; +use burn_std::IntDType; + +impl ModuleOps for CubeBackend +where + R: CubeRuntime, +{ + fn conv1d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<1>, + ) -> FloatTensor { + kernel::conv::conv_forward::(x, weight, bias, options, Default::default()).unwrap() + } + + fn conv1d_x_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<1>, + ) -> FloatTensor { + kernel::conv::conv_data_backward( + output_grad, + weight, + x.shape(), + options, + Default::default(), + ) + .unwrap() + } + + fn conv1d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<1>, + ) -> FloatTensor { + kernel::conv::conv_weight_backward::( + x, + output_grad, + weight.shape(), + options, + Default::default(), + ) + .unwrap() + } + + fn conv2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<2>, + ) -> FloatTensor { + kernel::conv::conv_forward::(x, weight, bias, options, Default::default()).unwrap() + } + + fn conv2d_x_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<2>, + ) -> FloatTensor { + kernel::conv::conv_data_backward( + output_grad, + weight, + x.shape(), + options, + Default::default(), + ) + .unwrap() + } + + fn conv2d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<2>, + ) -> FloatTensor { + kernel::conv::conv_weight_backward::( + x, + output_grad, + weight.shape(), + options, + Default::default(), + ) + .unwrap() + } + + fn deform_conv2d( + x: FloatTensor, + offset: FloatTensor, + weight: FloatTensor, + mask: Option>, + bias: Option>, + options: DeformConvOptions<2>, + ) -> FloatTensor { + kernel::conv::deform_conv2d(x, offset, weight, mask, bias, options).unwrap() + } + + fn deform_conv2d_backward( + x: FloatTensor, + offset: FloatTensor, + weight: FloatTensor, + mask: Option>, + bias: Option>, + output_grad: FloatTensor, + options: DeformConvOptions<2>, + ) -> DeformConv2dBackward { + let (x, o, w, m, b) = kernel::conv::deform_conv2d_backward( + x, + offset, + weight, + mask, + bias, + output_grad, + options, + ) + .unwrap(); + DeformConv2dBackward::new(x, o, w, m, b) + } + + fn conv3d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<3>, + ) -> FloatTensor { + kernel::conv::conv_forward::(x, weight, bias, options, Default::default()).unwrap() + } + + fn conv3d_x_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<3>, + ) -> FloatTensor { + kernel::conv::conv_data_backward( + output_grad, + weight, + x.shape(), + options, + Default::default(), + ) + .unwrap() + } + + fn conv3d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<3>, + ) -> FloatTensor { + kernel::conv::conv_weight_backward::( + x, + output_grad, + weight.shape(), + options, + Default::default(), + ) + .unwrap() + } + + fn conv_transpose2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<2>, + ) -> FloatTensor { + kernel::conv::conv_transpose2d(x, weight, bias, options, ConvTranspose2dStrategy::default()) + .unwrap() + } + + fn conv_transpose3d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<3>, + ) -> FloatTensor { + kernel::conv::conv_transpose3d(x, weight, bias, options).expect("Kernel to never fail") + } + + fn avg_pool2d( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + kernel::pool::avg_pool2d( + x, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ) + } + + fn avg_pool2d_backward( + x: FloatTensor, + grad: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + kernel::pool::avg_pool2d_backward( + x, + grad, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ) + } + + fn max_pool2d( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + ) -> FloatTensor { + kernel::pool::max_pool2d(x, kernel_size, stride, padding, dilation, ceil_mode) + } + + fn max_pool2d_with_indices( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + indices_dtype: IntDType, + ) -> MaxPool2dWithIndices { + let (output, indices) = kernel::pool::max_pool2d_with_indices( + x, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + indices_dtype.into(), + ); + + MaxPool2dWithIndices::new(output, indices) + } + + fn max_pool2d_with_indices_backward( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + output_grad: FloatTensor, + indices: IntTensor, + ) -> MaxPool2dBackward { + MaxPool2dBackward::new(kernel::pool::max_pool2d_with_indices_backward( + x, + output_grad, + indices, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + )) + } + + fn adaptive_avg_pool2d(x: FloatTensor, output_size: [usize; 2]) -> FloatTensor { + kernel::pool::adaptive_avg_pool2d(x, output_size) + } + + fn adaptive_avg_pool2d_backward( + x: FloatTensor, + grad: FloatTensor, + ) -> FloatTensor { + kernel::pool::adaptive_avg_pool2d_backward(x, grad) + } + + fn interpolate( + x: FloatTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> FloatTensor { + kernel::interpolate::interpolate(x, output_size, options, Default::default()).unwrap() + } + + fn interpolate_backward( + x: FloatTensor, + grad: FloatTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> FloatTensor { + kernel::interpolate::interpolate_backward(x, grad, output_size, options) + } + + fn attention( + query: FloatTensor, + key: FloatTensor, + value: FloatTensor, + mask: Option>, + attn_bias: Option>, + options: AttentionModuleOptions, + ) -> FloatTensor { + // Fall back to naive attention for features the flash kernel doesn't support. + if attn_bias.is_some() || options.softcap.is_some() || options.scale.is_some() { + return burn_backend::ops::attention::attention_fallback::( + query, key, value, mask, attn_bias, options, + ); + } + + kernel::attention::attention( + query, + key, + value, + mask, + attn_bias, + options, + Default::default(), + ) + .expect("Kernel to never fail") + } + + fn has_ctc_loss_backward() -> bool { + true + } + + fn ctc_loss( + log_probs: FloatTensor, + targets: IntTensor, + input_lengths: IntTensor, + target_lengths: IntTensor, + blank: usize, + ) -> FloatTensor { + kernel::ctc::ctc_loss(log_probs, targets, input_lengths, target_lengths, blank) + } + + fn ctc_loss_backward( + log_probs: FloatTensor, + targets: IntTensor, + input_lengths: IntTensor, + target_lengths: IntTensor, + grad_loss: FloatTensor, + blank: usize, + ) -> FloatTensor { + let (log_alpha_full, log_beta_full, nll) = kernel::ctc::ctc_alpha_beta( + log_probs.clone(), + targets.clone(), + input_lengths.clone(), + target_lengths, + blank, + ); + burn_backend::ops::ctc::ctc_grad_from_alpha_beta_default::( + log_probs, + targets, + input_lengths, + grad_loss, + log_alpha_full, + log_beta_full, + nll, + blank, + ) + } + + fn rfft( + signal: FloatTensor, + dim: usize, + n: Option, + ) -> (FloatTensor, FloatTensor) { + kernel::fft::rfft(signal, dim, n) + } + + fn irfft( + spectrum_re: FloatTensor, + spectrum_im: FloatTensor, + dim: usize, + n: Option, + ) -> FloatTensor { + kernel::fft::irfft(spectrum_re, spectrum_im, dim, n) + } +} diff --git a/crates/burn-cubecl/src/ops/numeric.rs b/crates/burn-cubecl/src/ops/numeric.rs new file mode 100644 index 0000000..0052133 --- /dev/null +++ b/crates/burn-cubecl/src/ops/numeric.rs @@ -0,0 +1,472 @@ +use crate::{ + CubeRuntime, + kernel::utils::{address_type, shape_divmod}, +}; +use crate::{element::CubeElement, tensor::CubeTensor}; +use crate::{ + kernel::{ + AddOp, BitwiseAndOp, BitwiseOrOp, BitwiseXorOp, DivOp, MulOp, PowOp, RemainderOp, SubOp, + launch_binop, launch_binop_int, launch_scalar_binop, launch_scalar_binop_int, + }, + ops::max_vector_size, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::{DType, Shape, TensorMetadata}; +use burn_std::Metadata; +use cubecl::{ + calculate_cube_count_elemwise, prelude::*, std::tensor::layout::linear::LinearViewMut, +}; +use cubecl::{client::ComputeClient, server::MemoryLayout}; +use cubecl::{server::MemoryLayoutDescriptor, std::FastDivmod}; + +/// Creates a tensor filled with `value` +pub fn full( + shape: Shape, + device: &R::Device, + value: E, +) -> CubeTensor { + let client = R::client(device); + + full_client::(client, shape, device.clone(), value) +} + +/// Creates a tensor filled with `value` +pub fn full_client( + client: ComputeClient, + shape: Shape, + device: R::Device, + value: E, +) -> CubeTensor { + let dtype = E::dtype(); + full_device_dtype( + client, + shape, + device, + InputScalar::new(value, dtype_to_storage_type(dtype)), + dtype, + ) +} + +/// Creates a tensor filled with `value` +pub fn full_device_dtype( + client: ComputeClient, + shape: Shape, + device: R::Device, + value: InputScalar, + dtype: DType, +) -> CubeTensor { + let empty = empty_device_dtype(client, device, shape, dtype); + + #[cube(launch_unchecked, address_type = "dynamic")] + pub fn full_kernel( + mut tensor: LinearViewMut<'_, Vector>, + value: InputScalar, + #[define(C)] _dtype: StorageType, + ) { + if !tensor.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + tensor.write(ABSOLUTE_POS, Vector::new(value.get::())); + } + + let num_elems = empty.meta.num_elements(); + let vector_size = max_vector_size(&empty); + + let working_units = num_elems / vector_size as usize; + let cube_dim = CubeDim::new(&empty.client, working_units); + let cube_count = calculate_cube_count_elemwise(&empty.client, working_units, cube_dim); + + unsafe { + full_kernel::launch_unchecked( + &empty.client, + cube_count, + cube_dim, + address_type!(empty), + vector_size, + empty.clone().into_linear_view(), + value, + dtype_to_storage_type(empty.dtype), + ); + } + + empty +} + +/// Creates a tensor filled with zeros +pub fn zeros(device: R::Device, shape: Shape, dtype: DType) -> CubeTensor { + let client = R::client(&device); + full_device_dtype( + client, + shape, + device, + InputScalar::new(0u32, dtype_to_storage_type(dtype)), + dtype, + ) +} + +/// Creates a tensor filled with ones +pub fn ones(device: R::Device, shape: Shape, dtype: DType) -> CubeTensor { + let client = R::client(&device); + full_device_dtype( + client, + shape, + device, + InputScalar::new(1u32, dtype_to_storage_type(dtype)), + dtype, + ) +} + +/// Creates a tensor filled with zeros +pub fn zeros_client( + client: ComputeClient, + device: R::Device, + shape: Shape, + dtype: DType, +) -> CubeTensor { + full_device_dtype( + client, + shape, + device, + InputScalar::new(0u32, dtype_to_storage_type(dtype)), + dtype, + ) +} + +/// Creates a tensor filled with ones +pub fn ones_client( + client: ComputeClient, + device: R::Device, + shape: Shape, + dtype: DType, +) -> CubeTensor { + full_device_dtype( + client, + shape, + device, + InputScalar::new(1u32, dtype_to_storage_type(dtype)), + dtype, + ) +} + +/// Create a tensor with uninitialized memory +pub fn empty_device( + client: ComputeClient, + device: R::Device, + shape: Shape, +) -> CubeTensor { + let MemoryLayout { memory, strides } = client.empty_tensor(shape.clone(), size_of::()); + + CubeTensor::new( + client, + memory, + Metadata::new(shape, strides), + device, + E::dtype(), + ) +} + +/// Create a tensor with uninitialized memory +pub fn empty_device_dtype( + client: ComputeClient, + device: R::Device, + shape: Shape, + dtype: DType, +) -> CubeTensor { + let MemoryLayout { memory, strides } = client.empty_tensor(shape.clone(), dtype.size()); + + CubeTensor::new(client, memory, Metadata::new(shape, strides), device, dtype) +} + +/// Create a contiguous tensor with uninitialized memory +pub fn empty_device_contiguous_dtype( + client: ComputeClient, + device: R::Device, + shape: Shape, + dtype: DType, +) -> CubeTensor { + let descriptor = MemoryLayoutDescriptor::contiguous(shape.clone(), dtype.size()); + let MemoryLayout { memory, strides } = client.empty_tensors(vec![descriptor]).remove(0); + + CubeTensor::new(client, memory, Metadata::new(shape, strides), device, dtype) +} + +/// Add two tensors +pub fn add(lhs: CubeTensor, rhs: CubeTensor) -> CubeTensor { + launch_binop::(lhs, rhs) +} + +/// Add a tensor and a scalar +pub fn add_scalar(lhs: CubeTensor, rhs: InputScalar) -> CubeTensor { + launch_scalar_binop::(lhs, rhs) +} + +/// Subtract two tensors +pub fn sub(lhs: CubeTensor, rhs: CubeTensor) -> CubeTensor { + launch_binop::(lhs, rhs) +} + +/// Subtract a tensor and a scalar +pub fn sub_scalar(lhs: CubeTensor, rhs: InputScalar) -> CubeTensor { + launch_scalar_binop::(lhs, rhs) +} + +/// Multiply two tensors +pub fn mul(lhs: CubeTensor, rhs: CubeTensor) -> CubeTensor { + launch_binop::(lhs, rhs) +} + +/// Multiply a tensor and a scalar +pub fn mul_scalar(lhs: CubeTensor, rhs: InputScalar) -> CubeTensor { + launch_scalar_binop::(lhs, rhs) +} + +/// Divide two tensors +pub fn div(lhs: CubeTensor, rhs: CubeTensor) -> CubeTensor { + launch_binop::(lhs, rhs) +} + +/// Divide a tensor by a scalar +pub fn div_scalar(lhs: CubeTensor, rhs: InputScalar) -> CubeTensor { + launch_scalar_binop::(lhs, rhs) +} + +/// Calculate remainder of two tensors +pub fn remainder(lhs: CubeTensor, rhs: CubeTensor) -> CubeTensor { + launch_binop::(lhs, rhs) +} + +/// Calculate the remainder of a tensor with a scalar +pub fn remainder_scalar(lhs: CubeTensor, rhs: InputScalar) -> CubeTensor { + launch_scalar_binop::(lhs, rhs) +} + +/// Calculate the power of two tensors +pub fn pow(lhs: CubeTensor, rhs: CubeTensor) -> CubeTensor { + launch_binop::(lhs, rhs) +} + +/// Bitwise and two tensors +pub fn bitwise_and(lhs: CubeTensor, rhs: CubeTensor) -> CubeTensor { + launch_binop_int::(lhs, rhs) +} + +/// Bitwise and with a scalar +pub fn bitwise_and_scalar(lhs: CubeTensor, rhs: InputScalar) -> CubeTensor { + launch_scalar_binop_int::(lhs, rhs) +} + +/// Bitwise or two tensors +pub fn bitwise_or(lhs: CubeTensor, rhs: CubeTensor) -> CubeTensor { + launch_binop_int::(lhs, rhs) +} + +/// Bitwise or with a scalar +pub fn bitwise_or_scalar(lhs: CubeTensor, rhs: InputScalar) -> CubeTensor { + launch_scalar_binop_int::(lhs, rhs) +} + +/// Bitwise xor two tensors +pub fn bitwise_xor(lhs: CubeTensor, rhs: CubeTensor) -> CubeTensor { + launch_binop_int::(lhs, rhs) +} + +/// Bitwise xor with a scalar +pub fn bitwise_xor_scalar(lhs: CubeTensor, rhs: InputScalar) -> CubeTensor { + launch_scalar_binop_int::(lhs, rhs) +} + +/// Operation family trait for cumulative operations +pub(crate) trait CumulativeOpFamily: Send + Sync + 'static { + type CumulativeOp: CumulativeOp; +} + +/// Trait for cumulative operations +#[cube] +pub(crate) trait CumulativeOp: 'static + Send + Sync { + /// Execute a cumulative operation + fn execute(lhs: C, rhs: C) -> C; + + /// Get the initial value for the accumulator + fn init_value(first_element: C) -> C; +} + +// Operation types +struct SumOp; +struct ProdOp; +struct MaxOp; +struct MinOp; + +// Implement CumulativeOpFamily for each operation +impl CumulativeOpFamily for SumOp { + type CumulativeOp = Self; +} + +impl CumulativeOpFamily for ProdOp { + type CumulativeOp = Self; +} + +impl CumulativeOpFamily for MaxOp { + type CumulativeOp = Self; +} + +impl CumulativeOpFamily for MinOp { + type CumulativeOp = Self; +} + +// Implement CumulativeOp for each operation type +#[cube] +impl CumulativeOp for SumOp { + fn execute(lhs: N, rhs: N) -> N { + lhs + rhs + } + + fn init_value(_first_element: N) -> N { + N::zero() + } +} + +#[cube] +impl CumulativeOp for ProdOp { + fn execute(lhs: N, rhs: N) -> N { + lhs * rhs + } + + fn init_value(_first_element: N) -> N { + N::from_int(1) + } +} + +#[cube] +impl CumulativeOp for MaxOp { + fn execute(lhs: N, rhs: N) -> N { + max(lhs, rhs) + } + + fn init_value(first_element: N) -> N { + first_element + } +} + +#[cube] +impl CumulativeOp for MinOp { + fn execute(lhs: N, rhs: N) -> N { + min(lhs, rhs) + } + + fn init_value(first_element: N) -> N { + first_element + } +} + +/// Generic cumulative operation kernel +/// +/// # Limitations +/// +/// This is a **naive sequential implementation** along the cumulative dimension: +/// - Each output element sequentially reads all previous elements along the dimension +/// - Computational complexity: O(n^2) memory reads where n is the size of the cumulative dimension +/// - **Performance:** Suitable for small tensors or small dimensions. For large tensors, +/// performance will degrade significantly compared to an optimized parallel scan algorithm. +/// +/// # TODO +/// +/// Implement an efficient GPU-optimized parallel scan algorithm. +#[cube(launch_unchecked, address_type = "dynamic")] +fn cumulative_kernel( + input: &Tensor, + mut output: LinearViewMut<'_, C>, + shape: Sequence>, + #[comptime] dim: usize, + #[define(C)] _dtype: StorageType, +) { + if !output.is_in_bounds(ABSOLUTE_POS) { + terminate!(); + } + + let rank = comptime![shape.len()]; + let dim_stride = input.stride(dim); + + let mut remainder = ABSOLUTE_POS; + let mut offset = 0; + let mut dim_idx = 0; + + #[unroll] + for i in 0..shape.len() { + let i = comptime![rank - i - 1]; + let (rem, local_idx) = shape.index(i).div_mod(remainder); + remainder = rem; + if i == dim { + dim_idx = local_idx; + } else { + offset += local_idx * input.stride(i); + } + } + + // Read first element + let first_read_idx = offset + dim_idx * dim_stride; + let first_elem = input[first_read_idx]; + + // Initialize accumulator + let mut result = O::CumulativeOp::::init_value(first_elem); + + // Accumulate values + for i in 0..=dim_idx { + let read_idx = offset + i * dim_stride; + result = O::CumulativeOp::::execute(result, input[read_idx]); + } + output.write(ABSOLUTE_POS, result); +} + +/// Compute the cumulative sum along a dimension +pub fn cumsum(input: CubeTensor, dim: usize) -> CubeTensor { + cumulative_op::(input, dim) +} + +/// Compute the cumulative product along a dimension +pub fn cumprod(input: CubeTensor, dim: usize) -> CubeTensor { + cumulative_op::(input, dim) +} + +/// Compute the cumulative minimum along a dimension +pub fn cummin(input: CubeTensor, dim: usize) -> CubeTensor { + cumulative_op::(input, dim) +} + +/// Compute the cumulative maximum along a dimension +pub fn cummax(input: CubeTensor, dim: usize) -> CubeTensor { + cumulative_op::(input, dim) +} + +/// Generic cumulative operation function +fn cumulative_op( + input: CubeTensor, + dim: usize, +) -> CubeTensor { + let client = input.client.clone(); + let device = input.device.clone(); + + let output = empty_device_dtype(client.clone(), device, input.shape(), input.dtype); + + let num_elems = output.meta.num_elements(); + let working_units = num_elems; + let cube_dim = CubeDim::new(&client, working_units); + let cube_count = calculate_cube_count_elemwise(&client, working_units, cube_dim); + let shape = shape_divmod(&input); + + unsafe { + cumulative_kernel::launch_unchecked::( + &client, + cube_count, + cube_dim, + address_type!(input, output), + input.into_tensor_arg(), + output.clone().into_linear_view(), + shape, + dim, + dtype_to_storage_type(output.dtype), + ); + } + + output +} diff --git a/crates/burn-cubecl/src/ops/qtensor.rs b/crates/burn-cubecl/src/ops/qtensor.rs new file mode 100644 index 0000000..d0efb50 --- /dev/null +++ b/crates/burn-cubecl/src/ops/qtensor.rs @@ -0,0 +1,283 @@ +use burn_backend::{ + Bytes, DType, ExecutionError, Shape, SplitPolicy, TensorData, TensorMetadata, TensorPrimitive, + get_device_settings, + ops::QTensorOps, + quantization::{ + QParamTensor, QuantLevel, QuantMode, QuantParam, QuantPropagation, QuantScheme, QuantValue, + QuantizationParametersPrimitive, params_shape, + }, + tensor::{Device, FloatTensor, QuantizedTensor}, +}; +use burn_std::{FloatDType, Metadata}; +use cubecl::server::{MemoryLayout, MemoryLayoutDescriptor, MemoryLayoutStrategy}; +use cubecl::{e2m1x2, quant::scheme::QuantStore}; + +use crate::{ + CubeBackend, CubeRuntime, + kernel::{self, matmul::MatmulStrategy}, + tensor::{CubeTensor, QParams}, +}; + +use super::{into_data, permute, swap_dims}; + +/// Create a quantized tensor with packed values (u32). +fn new_qtensor_optimized( + data: Bytes, + shape: impl Into, + scheme: QuantScheme, + device: &R::Device, +) -> CubeTensor { + new_qtensor(data, shape, scheme, device, MemoryLayoutStrategy::Optimized) +} + +/// Create a quantized tensor with packed values (u32). +fn new_qtensor( + data: Bytes, + shape: impl Into, + scheme: QuantScheme, + device: &R::Device, + kind: MemoryLayoutStrategy, +) -> CubeTensor { + new_quantized(shape, scheme, device, Some(data), kind) +} + +/// Create an empty quantized tensor. +pub fn empty_qtensor_optimized( + shape: impl Into, + scheme: QuantScheme, + device: &R::Device, +) -> CubeTensor { + empty_qtensor(shape, scheme, device, MemoryLayoutStrategy::Optimized) +} + +/// Create an empty quantized tensor. +pub fn empty_qtensor( + shape: impl Into, + scheme: QuantScheme, + device: &R::Device, + kind: MemoryLayoutStrategy, +) -> CubeTensor { + new_quantized(shape, scheme, device, None, kind) +} + +fn new_quantized( + shape: impl Into, + scheme: QuantScheme, + device: &R::Device, + data: Option, + alloc_kind: MemoryLayoutStrategy, +) -> CubeTensor { + let client = R::client(device); + let shape: Shape = shape.into(); + let mut shape_value: Shape = shape.clone(); + + let rank = shape.rank(); + let shape_last = shape[rank - 1]; + let num_quants = scheme.num_quants(); + + let data_size = match scheme.store { + QuantStore::PackedU32(_) => { + if !shape_last.is_multiple_of(num_quants) { + panic!("Can't store in u32") + } + shape_value[rank - 1] = shape_last.div_ceil(num_quants); + size_of::() + } + QuantStore::Native => match scheme.value { + QuantValue::Q8F | QuantValue::Q8S | QuantValue::E4M3 | QuantValue::E5M2 => { + size_of::() + } + QuantValue::Q4F + | QuantValue::Q4S + | QuantValue::Q2F + | QuantValue::Q2S + | QuantValue::E2M1 => { + panic!("Can't store native sub-byte values") + } + }, + QuantStore::PackedNative(_) => match scheme.value { + QuantValue::E2M1 => size_of::(), + other => panic!("{other:?} doesn't support native packing"), + }, + }; + + let scales_dtype = match scheme.param { + QuantParam::F32 => DType::F32, + QuantParam::F16 => DType::F16, + QuantParam::BF16 => DType::BF16, + // Represented by U8 and reinterpreted in the kernel + QuantParam::UE8M0 | QuantParam::UE4M3 => DType::U8, + }; + + let scales_shape = params_shape(&shape, scheme.level); + let data_desc = MemoryLayoutDescriptor::new(alloc_kind, shape_value.clone(), data_size); + let scales_desc = + MemoryLayoutDescriptor::new(alloc_kind, scales_shape.clone(), scales_dtype.size()); + + let mut tensors = match data { + Some(data) => { + let num_bytes = shape_value.num_elements() * data_size; + + match data.split(num_bytes, SplitPolicy::Shared) { + Ok((bytes_data, bytes_scales)) => client + .create_tensors(vec![(data_desc, bytes_data), (scales_desc, bytes_scales)]), + Err((data, _)) => client.create_tensors_from_slices(vec![ + (data_desc, &data[..num_bytes]), + (scales_desc, &data[num_bytes..]), + ]), + } + } + None => client.empty_tensors(vec![data_desc, scales_desc]), + }; + let MemoryLayout { + memory: scales_handle, + strides: scales_strides, + } = tensors.remove(1); + let MemoryLayout { memory, strides } = tensors.remove(0); + + let scales = QParamTensor { + offset_start: scales_handle.offset_start.unwrap_or(0) as usize, + offset_end: scales_handle.offset_end.unwrap_or(0) as usize, + metadata: Metadata::new(scales_shape, scales_strides), + dtype: scales_dtype, + }; + let qparams = QParams { scales }; + + CubeTensor::new_quantized( + client, + memory, + shape, + device.clone(), + strides, + DType::QFloat(scheme), + qparams, + ) +} + +impl QTensorOps for CubeBackend { + fn q_from_data(data: TensorData, device: &Device) -> QuantizedTensor { + match data.dtype { + DType::QFloat(scheme) => match scheme { + QuantScheme { + level: QuantLevel::Tensor | QuantLevel::Block(_), + mode: QuantMode::Symmetric, + value: + QuantValue::Q8F + | QuantValue::Q8S + | QuantValue::Q4F + | QuantValue::Q4S + | QuantValue::Q2F + | QuantValue::Q2S + | QuantValue::E4M3 + | QuantValue::E5M2 + | QuantValue::E2M1, + .. + } => { + // TensorData quantized representation is the same, with multiple quantized values + // packed into u32 and quantization parameters appended to the bytes + new_qtensor_optimized(data.bytes, data.shape.clone(), scheme, device) + } + }, + _ => panic!( + "Invalid dtype (expected DType::QFloat, got {:?})", + data.dtype + ), + } + } + + // TODO: quantize_dynamic (we can compute min-max on the fly and scale, especially when not per-tensor) + + fn quantize( + tensor: FloatTensor, + scheme: &QuantScheme, + qparams: QuantizationParametersPrimitive, + ) -> QuantizedTensor { + kernel::quantization::quantize(tensor, scheme, qparams.scales) + } + + fn dequantize(tensor: QuantizedTensor, dtype: FloatDType) -> FloatTensor { + kernel::quantization::dequantize(tensor, dtype.into()) + } + + fn q_to_device(tensor: QuantizedTensor, device: &Device) -> QuantizedTensor { + super::to_device(tensor, device) + } + + fn q_reshape(tensor: QuantizedTensor, shape: Shape) -> QuantizedTensor { + super::q_reshape(tensor, shape) + } + + async fn q_into_data(tensor: QuantizedTensor) -> Result { + if tensor.qparams.is_none() { + return into_data(tensor).await; + } + + let (shape, dtype) = (tensor.shape(), tensor.dtype); + let (values, params) = tensor.quantized_handles().unwrap(); + + let mut data_values = into_data(values).await?; + let data_params = into_data(params).await?; + + data_values.bytes.extend_from_byte_slice(&data_params.bytes); + + Ok(TensorData { + bytes: data_values.bytes, + shape, + dtype, + }) + } + + fn q_swap_dims( + tensor: QuantizedTensor, + dim1: usize, + dim2: usize, + ) -> QuantizedTensor { + swap_dims(tensor, dim1, dim2) + } + + fn q_permute(tensor: QuantizedTensor, axes: &[usize]) -> QuantizedTensor { + permute(tensor, axes) + } + + fn q_flip(_tensor: QuantizedTensor, _axes: &[usize]) -> QuantizedTensor { + unimplemented!() + } + + fn q_matmul(lhs: TensorPrimitive, rhs: TensorPrimitive) -> TensorPrimitive { + let (settings, scheme) = match (&lhs, &rhs) { + (TensorPrimitive::QFloat(lhs), _) => { + (get_device_settings::(&lhs.device), lhs.scheme()) + } + (_, TensorPrimitive::QFloat(rhs)) => { + (get_device_settings::(&rhs.device), rhs.scheme()) + } + _ => unreachable!(), + }; + + // Inherit precision for mixed inputs, default to `FloatElem` for fully quantized. + let out_dtype = match (&lhs, &rhs) { + (TensorPrimitive::Float(lhs), _) => lhs.dtype, + (_, TensorPrimitive::Float(rhs)) => rhs.dtype, + _ => settings.float_dtype.into(), + }; + + let (_lhs_dtype, lhs) = match lhs { + TensorPrimitive::Float(lhs) => (lhs.dtype, lhs), + TensorPrimitive::QFloat(lhs) => (out_dtype, lhs), + }; + let (_rhs_dtype, rhs) = match rhs { + TensorPrimitive::Float(rhs) => (rhs.dtype, rhs), + TensorPrimitive::QFloat(rhs) => (out_dtype, rhs), + }; + + let out = + kernel::matmul::matmul(lhs, rhs, None, MatmulStrategy::default(), out_dtype).unwrap(); + + match settings.quantization.propagation { + QuantPropagation::Propagate => { + TensorPrimitive::QFloat(Self::quantize_dynamic(out, &scheme)) + } + QuantPropagation::Inhibit => TensorPrimitive::Float(out), + } + } +} diff --git a/crates/burn-cubecl/src/ops/tensor.rs b/crates/burn-cubecl/src/ops/tensor.rs new file mode 100644 index 0000000..cf20881 --- /dev/null +++ b/crates/burn-cubecl/src/ops/tensor.rs @@ -0,0 +1,757 @@ +use super::{expand, numeric, permute, unfold}; +use crate::CubeBackend; +use crate::CubeRuntime; +use crate::kernel::matmul::{MatmulStrategy, matmul}; +use crate::kernel::prng::{random_bernoulli, random_normal, random_uniform}; +use crate::kernel::unary_basic::BasicFloatUnaryKind; +use crate::kernel::{ + self, FloatUnaryOp, FloatUnaryOpFamily, launch_unary_float, reduce, unary_basic, +}; +use burn_backend::cubecl::dtype_to_storage_type; +use burn_backend::ops::GridSampleOptions; +use burn_backend::tensor::{BoolTensor, Device, FloatTensor, IntTensor}; +use burn_backend::{DType, ElementConversion, FloatDType, Slice}; +use burn_backend::{Distribution, Shape, TensorData, ops::FloatTensorOps}; +use burn_backend::{ExecutionError, Scalar, get_device_settings}; +use burn_std::{BoolDType, IntDType}; +use cubecl::prelude::*; +use cubek::reduce::components::instructions::ReduceOperationConfig; +use std::ops::Range; + +impl FloatTensorOps for CubeBackend +where + R: CubeRuntime, +{ + #[cfg_attr(feature = "tracing", tracing::instrument( + level="trace", + skip(data), + fields(?data.shape, ?data.dtype) + ))] + fn float_from_data(data: TensorData, device: &Device) -> FloatTensor { + match data.dtype { + DType::F64 | DType::F32 | DType::F16 | DType::BF16 => super::from_data(data, device), + _ => unimplemented!("Unsupported dtype for `float_from_data`"), + } + } + + fn float_random( + shape: Shape, + distribution: Distribution, + device: &Device, + dtype: FloatDType, + ) -> FloatTensor { + let dtype = dtype.into(); + match distribution { + Distribution::Default => random_uniform(shape, device, 0., 1., dtype), + Distribution::Uniform(low, high) => { + random_uniform(shape, device, low.elem(), high.elem(), dtype) + } + Distribution::Bernoulli(prob) => random_bernoulli(shape, device, prob as f32, dtype), + Distribution::Normal(mean, std) => { + random_normal(shape, device, mean.elem(), std.elem(), dtype) + } + } + } + + #[cfg_attr(feature = "tracing", tracing::instrument( + level="trace", + skip(tensor), + fields(from = ?tensor.device, meta = ?tensor.meta, dtype = ?tensor.dtype) + ))] + async fn float_into_data(tensor: FloatTensor) -> Result { + super::into_data(tensor).await + } + + #[cfg_attr(feature = "tracing", tracing::instrument( + level="trace", + skip(tensor), + fields(from = ?tensor.device, meta = ?tensor.meta, dtype = ?tensor.dtype) + ))] + fn float_to_device(tensor: FloatTensor, device: &Device) -> FloatTensor { + super::to_device(tensor, device) + } + + fn float_empty(shape: Shape, device: &Device, dtype: FloatDType) -> FloatTensor { + let dtype = dtype.into(); + super::empty(shape, device, dtype) + } + + fn float_add(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + numeric::add(lhs, rhs) + } + + fn float_add_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let dtype = lhs.dtype; + numeric::add_scalar(lhs, InputScalar::new(rhs, dtype_to_storage_type(dtype))) + } + + fn float_zeros(shape: Shape, device: &Device, dtype: FloatDType) -> FloatTensor { + let dtype = dtype.into(); + numeric::zeros(device.clone(), shape, dtype) + } + + fn float_full( + shape: Shape, + fill_value: Scalar, + device: &R::Device, + dtype: FloatDType, + ) -> FloatTensor { + let dtype: DType = dtype.into(); + let client = R::client(device); + numeric::full_device_dtype( + client, + shape, + device.clone(), + InputScalar::new(fill_value, dtype_to_storage_type(dtype)), + dtype, + ) + } + + fn float_ones(shape: Shape, device: &Device, dtype: FloatDType) -> FloatTensor { + let dtype = dtype.into(); + numeric::ones(device.clone(), shape, dtype) + } + + fn float_sub(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + numeric::sub(lhs, rhs) + } + + fn float_sub_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let dtype = lhs.dtype; + numeric::sub_scalar(lhs, InputScalar::new(rhs, dtype_to_storage_type(dtype))) + } + + fn float_mul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + numeric::mul(lhs, rhs) + } + + fn float_mul_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let dtype = lhs.dtype; + numeric::mul_scalar(lhs, InputScalar::new(rhs, dtype_to_storage_type(dtype))) + } + + fn float_div(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + numeric::div(lhs, rhs) + } + + fn float_div_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let dtype = lhs.dtype; + numeric::div_scalar(lhs, InputScalar::new(rhs, dtype_to_storage_type(dtype))) + } + + fn float_remainder(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + numeric::remainder(lhs, rhs) + } + + fn float_remainder_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let dtype = lhs.dtype; + numeric::remainder_scalar(lhs, InputScalar::new(rhs, dtype_to_storage_type(dtype))) + } + + fn float_matmul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + let dtype = lhs.dtype; + matmul(lhs, rhs, None, MatmulStrategy::default(), dtype).unwrap() + } + + fn float_cross( + lhs: FloatTensor, + rhs: FloatTensor, + dim: usize, + ) -> FloatTensor { + kernel::cross(lhs, rhs, dim) + } + + fn float_swap_dims(tensor: FloatTensor, dim1: usize, dim2: usize) -> FloatTensor { + super::swap_dims(tensor, dim1, dim2) + } + + fn float_reshape(tensor: FloatTensor, shape: Shape) -> FloatTensor { + super::reshape(tensor, shape) + } + + fn float_gather( + dim: usize, + tensor: FloatTensor, + indices: IntTensor, + ) -> FloatTensor { + kernel::gather(dim, tensor, indices) + } + + fn float_scatter_add( + dim: usize, + tensor: FloatTensor, + indices: IntTensor, + value: FloatTensor, + ) -> FloatTensor { + kernel::scatter(dim, tensor, indices, value, false) + } + + fn float_scatter_nd( + data: FloatTensor, + indices: IntTensor, + values: FloatTensor, + reduction: burn_backend::tensor::IndexingUpdateOp, + ) -> FloatTensor { + kernel::scatter_nd(data, indices, values, reduction) + } + + fn float_gather_nd(data: FloatTensor, indices: IntTensor) -> FloatTensor { + kernel::gather_nd(data, indices) + } + + fn float_select( + tensor: FloatTensor, + dim: usize, + indices: IntTensor, + ) -> FloatTensor { + kernel::select(tensor, dim, indices) + } + + fn float_select_add( + tensor: FloatTensor, + dim: usize, + indices: IntTensor, + value: FloatTensor, + ) -> FloatTensor { + kernel::select_assign(tensor, dim, indices, value, false) + } + + fn float_slice(tensor: FloatTensor, slices: &[Slice]) -> FloatTensor { + // Check if all steps are 1 + let all_steps_one = slices.iter().all(|info| info.step == 1); + + if all_steps_one { + // Use optimized slice for step=1 + let simple_ranges: Vec> = slices + .iter() + .enumerate() + .map(|(i, slice)| slice.to_range(tensor.meta.shape()[i])) + .collect(); + + kernel::slice(tensor, &simple_ranges) + } else { + // Use slice with steps kernel + kernel::slice_with_steps(tensor, slices) + } + } + + fn float_slice_assign( + tensor: FloatTensor, + ranges: &[Slice], + value: FloatTensor, + ) -> FloatTensor { + kernel::slice_assign(tensor, ranges, value) + } + + fn float_mask_where( + tensor: FloatTensor, + mask: BoolTensor, + value: FloatTensor, + ) -> FloatTensor { + let bool_dtype = mask.dtype; + kernel::mask_where_auto(tensor, mask, value, bool_dtype) + } + + fn float_mask_fill( + tensor: FloatTensor, + mask: BoolTensor, + value: Scalar, + ) -> FloatTensor { + let dtype = tensor.dtype; + let bool_dtype = mask.dtype; + kernel::mask_fill_auto( + tensor, + mask, + InputScalar::new(value, dtype_to_storage_type(dtype)), + bool_dtype, + ) + } + + fn float_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + kernel::equal(lhs, rhs, out_dtype.into()) + } + + fn float_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let dtype = lhs.dtype; + kernel::equal_elem( + lhs, + InputScalar::new(rhs, dtype_to_storage_type(dtype)), + out_dtype.into(), + ) + } + + fn float_greater( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + kernel::greater(lhs, rhs, out_dtype.into()) + } + + fn float_greater_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let dtype = lhs.dtype; + kernel::greater_elem( + lhs, + InputScalar::new(rhs, dtype_to_storage_type(dtype)), + out_dtype.into(), + ) + } + + fn float_greater_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + kernel::greater_equal(lhs, rhs, out_dtype.into()) + } + + fn float_greater_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let dtype = lhs.dtype; + kernel::greater_equal_elem( + lhs, + InputScalar::new(rhs, dtype_to_storage_type(dtype)), + out_dtype.into(), + ) + } + + fn float_lower( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + kernel::lower(lhs, rhs, out_dtype.into()) + } + + fn float_lower_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let dtype = lhs.dtype; + kernel::lower_elem( + lhs, + InputScalar::new(rhs, dtype_to_storage_type(dtype)), + out_dtype.into(), + ) + } + + fn float_lower_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + kernel::lower_equal(lhs, rhs, out_dtype.into()) + } + + fn float_lower_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let dtype = lhs.dtype; + kernel::lower_equal_elem( + lhs, + InputScalar::new(rhs, dtype_to_storage_type(dtype)), + out_dtype.into(), + ) + } + + fn float_sum(tensor: FloatTensor) -> FloatTensor { + reduce::sum_fallback(tensor, Default::default()).unwrap() + } + + fn float_max(tensor: FloatTensor) -> FloatTensor { + reduce::reduce(tensor, None, Default::default(), ReduceOperationConfig::Max).unwrap() + } + + fn float_max_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + reduce::reduce_dim( + tensor, + None, + dim, + Default::default(), + ReduceOperationConfig::Max, + ) + .unwrap() + } + + fn float_min(tensor: FloatTensor) -> FloatTensor { + reduce::reduce(tensor, None, Default::default(), ReduceOperationConfig::Min).unwrap() + } + + fn float_min_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + reduce::reduce_dim( + tensor, + None, + dim, + Default::default(), + ReduceOperationConfig::Min, + ) + .unwrap() + } + + fn float_max_abs(tensor: FloatTensor) -> FloatTensor { + reduce::reduce( + tensor, + None, + Default::default(), + ReduceOperationConfig::MaxAbs, + ) + .unwrap() + } + + fn float_max_abs_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + reduce::reduce_dim( + tensor, + None, + dim, + Default::default(), + ReduceOperationConfig::MaxAbs, + ) + .unwrap() + } + + fn float_any(tensor: FloatTensor, out_dtype: BoolDType) -> BoolTensor { + reduce::reduce_logical(tensor, None, ReduceOperationConfig::Any, out_dtype) + } + + fn float_any_dim( + tensor: FloatTensor, + dim: usize, + out_dtype: BoolDType, + ) -> BoolTensor { + reduce::reduce_logical(tensor, Some(dim), ReduceOperationConfig::Any, out_dtype) + } + + fn float_all(tensor: FloatTensor, out_dtype: BoolDType) -> BoolTensor { + reduce::reduce_logical(tensor, None, ReduceOperationConfig::All, out_dtype) + } + + fn float_all_dim( + tensor: FloatTensor, + dim: usize, + out_dtype: BoolDType, + ) -> BoolTensor { + reduce::reduce_logical(tensor, Some(dim), ReduceOperationConfig::All, out_dtype) + } + + fn float_sum_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + reduce::reduce_dim( + tensor, + None, + dim, + Default::default(), + ReduceOperationConfig::Sum, + ) + .unwrap() + } + + fn float_mean_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + reduce::reduce_dim( + tensor, + None, + dim, + Default::default(), + ReduceOperationConfig::Mean, + ) + .unwrap() + } + + fn float_mean(tensor: FloatTensor) -> FloatTensor { + reduce::reduce( + tensor, + None, + Default::default(), + ReduceOperationConfig::Mean, + ) + .unwrap() + } + + fn float_cumsum(tensor: FloatTensor, dim: usize) -> FloatTensor { + numeric::cumsum(tensor, dim) + } + + fn float_cumprod(tensor: FloatTensor, dim: usize) -> FloatTensor { + numeric::cumprod(tensor, dim) + } + + fn float_cummin(tensor: FloatTensor, dim: usize) -> FloatTensor { + numeric::cummin(tensor, dim) + } + + fn float_cummax(tensor: FloatTensor, dim: usize) -> FloatTensor { + numeric::cummax(tensor, dim) + } + + fn float_prod(tensor: FloatTensor) -> FloatTensor { + reduce::reduce( + tensor, + None, + Default::default(), + ReduceOperationConfig::Prod, + ) + .unwrap() + } + + fn float_prod_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + reduce::reduce_dim( + tensor, + None, + dim, + Default::default(), + ReduceOperationConfig::Prod, + ) + .unwrap() + } + + fn float_exp(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Exp) + } + + fn float_log(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Log) + } + + fn float_log1p(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Log1p) + } + + fn float_powf_scalar_impl(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + struct Powf; + + #[cube] + impl FloatUnaryOp for Powf { + type Options = InputScalar; + + fn execute(input: Vector, options: &Self::Options) -> Vector { + Vector::powf(input, Vector::new(options.get::())) + } + } + + impl FloatUnaryOpFamily for Powf { + type Options = InputScalar; + type Unary = Self; + } + + let dtype = lhs.dtype; + launch_unary_float::(lhs, |_| { + InputScalar::new(rhs, dtype_to_storage_type(dtype)) + }) + } + + fn float_sqrt(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Sqrt) + } + + fn float_abs(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Abs) + } + + fn float_sign(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Sign) + } + + fn float_cos(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Cos) + } + + fn float_sin(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Sin) + } + + fn float_tan(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Tan) + } + + fn float_cosh(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Cosh) + } + + fn float_sinh(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Sinh) + } + + fn float_tanh(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Tanh) + } + + fn float_acos(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::ArcCos) + } + + fn float_acosh(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::ArcCosh) + } + + fn float_asin(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::ArcSin) + } + + fn float_asinh(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::ArcSinh) + } + + fn float_atan(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::ArcTan) + } + + fn float_atanh(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::ArcTanh) + } + + fn float_atan2(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + crate::kernel::atan2::(lhs, rhs) + } + + fn float_round(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Round) + } + + fn float_floor(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Floor) + } + + fn float_ceil(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Ceil) + } + + fn float_trunc(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Trunc) + } + + fn float_erf(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Erf) + } + + fn float_argmax(tensor: FloatTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + reduce::reduce_dim( + tensor, + Some(out_dtype.into()), + dim, + Default::default(), + ReduceOperationConfig::ArgMax, + ) + .unwrap() + } + + fn float_argtopk( + tensor: FloatTensor, + dim: usize, + k: usize, + out_dtype: IntDType, + ) -> IntTensor { + reduce::reduce_dim( + tensor, + Some(out_dtype.into()), + dim, + Default::default(), + ReduceOperationConfig::ArgTopK(k), + ) + .unwrap() + } + + fn float_topk(tensor: FloatTensor, dim: usize, k: usize) -> FloatTensor { + reduce::reduce_dim( + tensor, + None, + dim, + Default::default(), + ReduceOperationConfig::TopK(k), + ) + .unwrap() + } + + fn float_argmin(tensor: FloatTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + reduce::reduce_dim( + tensor, + Some(out_dtype.into()), + dim, + Default::default(), + ReduceOperationConfig::ArgMin, + ) + .unwrap() + } + + fn float_into_int(tensor: FloatTensor, out_dtype: IntDType) -> IntTensor { + kernel::cast(tensor, out_dtype.into()) + } + + fn float_clamp(tensor: FloatTensor, min: Scalar, max: Scalar) -> FloatTensor { + let dtype = tensor.dtype; + kernel::clamp( + tensor, + InputScalar::new(min, dtype_to_storage_type(dtype)), + InputScalar::new(max, dtype_to_storage_type(dtype)), + ) + } + + fn float_recip(tensor: FloatTensor) -> FloatTensor { + unary_basic::launch::(tensor, |_| BasicFloatUnaryKind::Recip) + } + + fn float_repeat_dim(tensor: FloatTensor, dim: usize, times: usize) -> FloatTensor { + kernel::repeat_dim(tensor, dim, times) + } + + fn float_powf(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + numeric::pow(lhs, rhs) + } + + fn float_permute(tensor: FloatTensor, axes: &[usize]) -> FloatTensor { + permute(tensor, axes) + } + + fn float_expand(tensor: FloatTensor, shape: Shape) -> FloatTensor { + expand(tensor, shape) + } + + fn float_flip(tensor: FloatTensor, axes: &[usize]) -> FloatTensor { + let bool_dtype = get_device_settings::(&tensor.device).bool_dtype; + kernel::flip(tensor, axes, bool_dtype.into()) + } + + fn float_cast(tensor: FloatTensor, dtype: FloatDType) -> FloatTensor { + kernel::cast(tensor, dtype.into()) + } + + fn float_unfold( + tensor: FloatTensor, + dim: usize, + size: usize, + step: usize, + ) -> FloatTensor { + unfold(tensor, dim, size, step) + } + + fn float_is_nan(tensor: FloatTensor, out_dtype: BoolDType) -> BoolTensor { + kernel::is_nan(tensor, out_dtype.into()) + } + + fn float_is_inf(tensor: FloatTensor, out_dtype: BoolDType) -> BoolTensor { + kernel::is_inf(tensor, out_dtype.into()) + } + + fn float_grid_sample_2d( + tensor: FloatTensor, + grid: FloatTensor, + options: GridSampleOptions, + ) -> FloatTensor { + kernel::grid_sample::grid_sample(tensor, grid, options) + } +} diff --git a/crates/burn-cubecl/src/ops/transaction.rs b/crates/burn-cubecl/src/ops/transaction.rs new file mode 100644 index 0000000..ddc44dd --- /dev/null +++ b/crates/burn-cubecl/src/ops/transaction.rs @@ -0,0 +1,137 @@ +use burn_backend::{ + DType, TensorData, + backend::ExecutionError, + ops::{TransactionOps, TransactionPrimitive, TransactionPrimitiveData}, +}; +use burn_std::{Shape, Strides}; +use cubecl::server::{CopyDescriptor, Handle}; + +use crate::{CubeBackend, CubeRuntime}; + +impl TransactionOps for CubeBackend { + async fn tr_execute( + transaction: TransactionPrimitive, + ) -> Result { + let mut client = None; + + enum Kind { + Float, + Int, + Bool, + } + + #[derive(new)] + struct BindingData { + index: usize, + kind: Kind, + handle: Option, + shape: Shape, + strides: Strides, + dtype: DType, + } + + let mut num_bindings = 0; + + let mut kinds = Vec::new(); + + for t in transaction.read_floats.into_iter() { + if client.is_none() { + client = Some(t.client.clone()); + } + + let t = crate::kernel::into_contiguous_aligned(t); + let binding = BindingData::new( + num_bindings, + Kind::Float, + Some(t.handle.clone()), + t.meta.shape.clone(), + t.meta.strides.clone(), + t.dtype, + ); + + kinds.push(binding); + num_bindings += 1; + } + for t in transaction.read_ints.into_iter() { + if client.is_none() { + client = Some(t.client.clone()); + } + + let t = crate::kernel::into_contiguous_aligned(t); + let binding = BindingData::new( + num_bindings, + Kind::Int, + Some(t.handle.clone()), + t.meta.shape.clone(), + t.meta.strides.clone(), + t.dtype, + ); + + kinds.push(binding); + num_bindings += 1; + } + for t in transaction.read_bools.into_iter() { + if client.is_none() { + client = Some(t.client.clone()); + } + + let t = crate::kernel::into_contiguous_aligned(t); + let binding = BindingData::new( + num_bindings, + Kind::Bool, + Some(t.handle.clone()), + t.meta.shape.clone(), + t.meta.strides.clone(), + t.dtype, + ); + + kinds.push(binding); + num_bindings += 1; + } + + let client = client.unwrap(); + + let bindings = kinds + .iter_mut() + .map(|b| { + CopyDescriptor::new( + b.handle.take().unwrap().binding(), + b.shape.clone(), + b.strides.clone(), + b.dtype.size(), + ) + }) + .collect(); + + let mut data: Vec> = client + .read_tensor_async(bindings) + .await + .map_err(|err| ExecutionError::WithContext { + reason: format!("{err:?}"), + })? + .into_iter() + .map(Some) + .collect::>>(); + + let mut result = TransactionPrimitiveData::default(); + + for binding in kinds { + let bytes = data.get_mut(binding.index).unwrap().take().unwrap(); + let t_data = TensorData::from_bytes(bytes, binding.shape, binding.dtype); + + match binding.kind { + Kind::Float => { + result.read_floats.push(t_data); + } + Kind::Int => { + result.read_ints.push(t_data); + } + Kind::Bool => { + result.read_bools.push(t_data); + } + } + } + + Ok(result) + } +} diff --git a/crates/burn-cubecl/src/template/base.rs b/crates/burn-cubecl/src/template/base.rs new file mode 100644 index 0000000..0ec7c08 --- /dev/null +++ b/crates/burn-cubecl/src/template/base.rs @@ -0,0 +1,103 @@ +use super::SourceTemplate; +use crate::{CubeRuntime, tensor::CubeTensor}; +use cubecl::{CompilationError, Compiler, CubeTask, prelude::*}; + +/// Kernel source to create a [source](SourceTemplate) +pub trait KernelSource: Send + 'static + Sync { + /// Convert to [source](SourceTemplate) + fn source(&self) -> SourceTemplate; + /// Identifier for the kernel, used for caching kernel compilation. + fn id(&self) -> KernelId; +} + +#[derive(new)] +/// Wraps a [kernel source](KernelSource) into a [cube task](CubeTask). +pub struct SourceKernel { + kernel_source: K, + cube_dim: CubeDim, +} + +impl CubeTask for SourceKernel { + fn compile( + &self, + _compiler: &mut C, + _options: &C::CompilationOptions, + _mode: ExecutionMode, + _address_type: StorageType, + ) -> Result, CompilationError> { + let source_template = self.kernel_source.source(); + let source = source_template.complete(); + + Ok(CompiledKernel { + entrypoint_name: "main".to_string(), + debug_name: Some(core::any::type_name::()), + source, + cube_dim: self.cube_dim, + debug_info: None, + repr: None, + }) + } +} + +impl KernelMetadata for SourceKernel { + fn id(&self) -> KernelId { + self.kernel_source.id() + } + + fn address_type(&self) -> StorageType { + u32::as_type_native_unchecked().storage_type() + } +} + +/// Generates kernel source code by replacing some information using templating. +#[macro_export] +macro_rules! kernel_source { + ( + $struct:ident, + $file:expr + ) => { + /// Generated kernel from a source file. + #[derive(new)] + pub struct $struct; + + impl $struct { + fn source(&self) -> $crate::template::SourceTemplate { + $crate::template::SourceTemplate::new(include_str!($file)) + } + } + }; +} + +/// Create a vector containing the dimension, strides and shape of tensors. +/// +/// # Example +/// +/// With two tensors (lhs, rhs) +/// +/// | Indexes | Value | +/// |:------------------------:|:-----------:| +/// | 0..1 | D | +/// | 1..D + 1 | lhs strides | +/// | (D + 1)..(2 * D + 1) | rhs strides | +/// | (2 * D + 1)..(3 * D + 1) | lhs shape | +/// | (3 * D + 1)..(4 * D + 1) | rhs shape | +pub fn build_info(tensors: &[&CubeTensor]) -> Vec { + let ndims = tensors[0].meta.num_dims(); + let mut info: Vec = vec![0; tensors.len() * 2 * ndims + 1]; + info[0] = ndims as u32; + + let mut current = 1; + for tensor in tensors.iter() { + for d in 0..ndims { + info[current] = tensor.meta.strides()[d] as u32; + current += 1; + } + } + for tensor in tensors.iter() { + for d in 0..ndims { + info[current] = tensor.meta.shape()[d] as u32; + current += 1; + } + } + info +} diff --git a/crates/burn-cubecl/src/template/mod.rs b/crates/burn-cubecl/src/template/mod.rs new file mode 100644 index 0000000..8c34090 --- /dev/null +++ b/crates/burn-cubecl/src/template/mod.rs @@ -0,0 +1,5 @@ +mod base; +pub use base::*; + +mod source; +pub use source::*; diff --git a/crates/burn-cubecl/src/template/source.rs b/crates/burn-cubecl/src/template/source.rs new file mode 100644 index 0000000..b13c2f6 --- /dev/null +++ b/crates/burn-cubecl/src/template/source.rs @@ -0,0 +1,69 @@ +use std::collections::HashMap; + +/// Kernel source code abstraction allowing for templating. +/// +/// The templates can have text placeholders in the form {{ label }}. +/// They will be updated with their proper value when `generate` is called. +#[derive(Debug)] +pub struct SourceTemplate { + items: HashMap, + templates: Vec, +} + +impl SourceTemplate { + /// Create a new source template. + pub fn new(template: S) -> Self + where + S: Into, + { + Self { + items: HashMap::new(), + templates: vec![template.into()], + } + } + + /// Register the value for a placeholder item. + /// + /// # Notes + /// + /// The value can't have placeholders, since it would require recursive templating with + /// possibly circular dependencies. If you want to add a value that has some + /// placeholders, consider adding a new template to the source using + /// [add_template](SourceTemplate::add_template). The added template can be a function, and you can + /// register the function call instead. + pub fn register(mut self, name: Name, value: Value) -> Self + where + Name: Into, + Value: Into, + { + self.items.insert(name.into(), value.into()); + self + } + + /// Add a new template. + pub fn add_template(mut self, template: S) -> Self + where + S: Into, + { + self.templates.push(template.into()); + self + } + + /// Complete the template and returns the source code. + pub fn complete(mut self) -> String { + let mut source = self.templates.remove(0); + + for s in self.templates.into_iter() { + source.push_str(&s); + } + + let template = text_placeholder::Template::new(&source); + let mut context = HashMap::new(); + + for (key, value) in self.items.iter() { + context.insert(key.as_str(), value.as_str()); + } + + template.fill_with_hashmap(&context) + } +} diff --git a/crates/burn-cubecl/src/tensor/base.rs b/crates/burn-cubecl/src/tensor/base.rs new file mode 100644 index 0000000..7f835ac --- /dev/null +++ b/crates/burn-cubecl/src/tensor/base.rs @@ -0,0 +1,405 @@ +use crate::CubeRuntime; +use crate::kernel::{NumericUnaryOp, NumericUnaryOpFamily, launch_unary_numeric}; +use burn_backend::cubecl::{dtype_to_elem_type, dtype_to_storage_type}; +use burn_backend::quantization::QuantScheme; +use burn_backend::{DType, Shape, TensorMetadata}; +use burn_std::{Metadata, strides, tensor::is_contiguous}; +use cubecl::server::Handle; +use cubecl::std::tensor::TensorHandle; +use cubecl::{client::ComputeClient, std::tensor::layout::linear::LinearViewLaunch}; +use cubecl::{frontend::Numeric, std::tensor::layout::linear::LinearViewLayoutLaunch}; +use cubecl::{ + prelude::{TensorBinding, *}, + std::tensor::layout::linear::LinearViewLayout, +}; +use std::marker::PhantomData; + +use super::QParams; + +/// The basic tensor primitive struct. +pub struct CubeTensor { + /// Compute client for the [runtime](CubeRuntime). + pub client: ComputeClient, + /// The buffer where the data are stored. + pub handle: Handle, + /// The metadata of the tensor. + pub meta: Box, + /// The device of the tensor. + pub device: R::Device, + /// The datatype of the tensor. + pub dtype: DType, + /// Runtime quantization parameters, if applicable + pub qparams: Option, +} + +impl From> for TensorHandle { + fn from(val: CubeTensor) -> Self { + TensorHandle::new( + val.handle.clone(), + val.meta.shape().clone(), + val.meta.strides().clone(), + dtype_to_storage_type(val.dtype), + ) + } +} + +impl cubecl::tune::AutotuneOutput for CubeTensor { + #[cfg(feature = "autotune-checks")] + fn check_equivalence(&self, other: Self) { + use crate::ops::into_data_sync; + use burn_backend::Tolerance; + + let expected = into_data_sync::(self.clone()); + let actual = into_data_sync::(other); + expected.assert_approx_eq::(&actual, Tolerance::permissive()); + } +} + +// TODO: Needed to cleanup leaves tensor. +// +// Maybe not needed when fusion is activated, since we have a detector there. +// We could rely on basic GC strategy when not using fusion. +// +// impl Drop for CubeTensor { +// fn drop(&mut self) { +// todo!() +// } +// } + +impl core::fmt::Debug for CubeTensor +where + R: CubeRuntime, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_fmt(format_args!( + "CubeTensor {{ shape: {:?}, device: {:?}, strides: {:?}, elem: {}, runtime: {}}}", + self.meta.shape(), + self.device, + self.meta.strides(), + self.dtype.name(), + R::name(&self.client), + )) + } +} + +impl Clone for CubeTensor +where + R: CubeRuntime, +{ + fn clone(&self) -> Self { + Self { + client: self.client.clone(), + handle: self.handle.clone(), + meta: self.meta.clone(), + device: self.device.clone(), + dtype: self.dtype, + qparams: self.qparams.clone(), + } + } +} + +impl TensorMetadata for CubeTensor { + type Device = R::CubeDevice; + fn dtype(&self) -> DType { + self.dtype + } + + fn shape(&self) -> Shape { + self.meta.shape().clone() + } + + fn rank(&self) -> usize { + self.meta.rank() + } + + fn device(&self) -> Self::Device { + self.device.clone() + } + + fn can_mut(&self) -> bool { + self.handle.can_mut() + } +} + +impl CubeTensor +where + R: CubeRuntime, +{ + /// Create a new standard tensor + pub fn new( + client: ComputeClient, + handle: Handle, + metadata: Metadata, + device: R::Device, + dtype: DType, + ) -> Self { + CubeTensor { + client, + handle, + meta: Box::new(metadata), + device, + dtype, + qparams: None, + } + } + + /// Create a new tensor with a contiguous memory layout. + pub fn new_contiguous( + client: ComputeClient, + device: R::Device, + shape: Shape, + handle: Handle, + dtype: DType, + ) -> Self { + let ndims = shape.num_dims(); + let mut strides = strides![0; ndims]; + let mut current = 1; + + shape.iter().enumerate().rev().for_each(|(index, val)| { + strides[index] = current; + current *= val; + }); + + Self { + client, + handle, + meta: Box::new(Metadata::new(shape, strides)), + device, + dtype, + qparams: None, + } + } + + /// Change the context of the current tensor and return the newly transferred tensor. + pub fn to_client(&mut self, client: ComputeClient, device: R::Device) -> Self { + let desc = self.handle.clone().copy_descriptor( + self.meta.shape().clone(), + self.meta.strides().clone(), + self.elem_size(), + ); + let handle = self + .client + .to_client_tensor(desc, &client, dtype_to_elem_type(self.dtype)); + + Self { + client, + handle, + meta: Box::new(Metadata::new(self.shape(), self.meta.strides().clone())), + device, + dtype: self.dtype, + qparams: self.qparams.clone(), + } + } + + /// Return the reference to a tensor handle. + pub fn binding(self) -> TensorBinding { + TensorBinding { + handle: self.handle.binding(), + strides: self.meta.strides, + shape: self.meta.shape, + runtime: PhantomData, + } + } + + /// Returns the element size of this tensor + pub fn elem_size(&self) -> usize { + self.dtype.size() + } + + /// Return the reference to a tensor argument. + pub fn into_tensor_arg(self) -> TensorArg { + self.binding().into_tensor_arg() + } + + /// Return the reference to a buffer argument. + pub fn into_buffer_arg(self) -> BufferArg { + self.into_tensor_arg().into_buffer_arg() + } + + /// Returns a reference to the aliased tensor argument. + pub fn as_tensor_alias(&self, input_pos: usize) -> TensorArg { + TensorArg::Alias { + input_pos, + strides: self.meta.strides().clone(), + shape: self.meta.shape().clone(), + } + } + + /// Return a linear view of this tensor. + pub fn into_linear_view(self) -> LinearViewLaunch { + let layout = LinearViewLayoutLaunch::new(); + let buffer = self.into_tensor_arg(); + LinearViewLaunch::new_tensor::(buffer, layout) + } + + /// Return an aliased linear view of this tensor + pub fn as_linear_view_alias(&self, input_pos: usize) -> LinearViewLaunch { + let layout = LinearViewLayoutLaunch::new(); + let buffer = self.as_tensor_alias(input_pos); + LinearViewLaunch::new_tensor::(buffer, layout) + } + + /// Return a linear view broadcast to the reference tensor's shape + pub fn into_linear_view_like(self, reference: &Self) -> LinearViewLaunch { + let layout = LinearViewLayoutLaunch::from_reference_shape(reference.shape()); + let buffer = self.into_tensor_arg(); + LinearViewLaunch::new_tensor::(buffer, layout) + } + + /// Returns the address type required to index this tensor + pub fn required_address_type(&self) -> AddressType { + match self.try_scheme() { + Some(scheme) => { + let len = self.handle.size() as usize * 8 / scheme.size_bits_value(); + AddressType::from_len(len) + } + None => AddressType::from_len(self.handle.size() as usize / self.dtype.size()), + } + } + + /// Return the `QuantScheme` if present + pub fn try_scheme(&self) -> Option<&QuantScheme> { + match &self.dtype { + DType::QFloat(scheme) => Some(scheme), + _ => None, + } + } + + pub(crate) fn can_mut_broadcast(&self, rhs: &Self) -> bool { + if !self.handle.can_mut() || !self.is_nonoverlapping() { + return false; + } + let ndims = self.meta.num_dims(); + + for i in 0..ndims { + let shape_lhs = self.meta.shape()[i]; + let shape_rhs = rhs.meta.shape()[i]; + + // Output tensor will be different from the mutable tensor. + if shape_lhs < shape_rhs { + return false; + } + } + + true + } + + /// Copy the current tensor. + pub fn copy(&self) -> Self { + struct Copy; + + #[cube] + impl NumericUnaryOp for Copy { + type Options = (); + + fn execute(input: Vector, _options: &Self::Options) -> Vector { + input + } + } + + impl NumericUnaryOpFamily for Copy { + type Options = (); + type Unary = Self; + } + + let tensor = self.clone(); + launch_unary_numeric::(tensor, |_| ()) + } + + /// Check if the tensor is safe to mutate. + pub fn can_mut(&self) -> bool { + self.handle.can_mut() + } + + /// Assert that both tensors are on the same device. + pub fn assert_is_on_same_device(&self, other: &Self) { + if self.device != other.device { + panic!( + "Both tensors should be on the same device {:?} != {:?}", + self.device, other.device + ); + } + } + + /// Check if the current tensor is contiguous. + /// + /// A tensor is contiguous if the elements are stored in memory + /// if the strides in non-increasing order and the + /// strides at position k is equal to the product of the shapes + /// at all positions greater than k. However, all axes with a shape of 1 are ignored. + pub fn is_contiguous(&self) -> bool { + is_contiguous(self.meta.shape(), self.meta.strides()) + } + + /// Check if the current tensor has a contiguous backing buffer (no overlap and no empty memory + /// regions within the shape). + pub fn is_contiguous_buffer(&self) -> bool { + self.meta.shape().num_elements() * self.dtype.size() == self.handle.size() as usize + } + + /// Checks if the tensor is non-overlapping (can be safely written to). + pub fn is_nonoverlapping(&self) -> bool { + let shape = self.meta.shape(); + let strides = self.meta.strides(); + + if strides.contains(&0) { + return false; + } + let rank = self.rank(); + if rank > 1 { + let mut dims = shape.iter().zip(strides.iter()).collect::>(); + dims.sort_by_key(|(_, stride)| **stride); + + let mut max_offset = 0; + for (shape, stride) in dims.into_iter() { + if *stride <= max_offset && *shape != 1 { + return false; + } + + max_offset += (*shape - 1) * *stride; + } + } + true + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn is_contiguous_non_increasing() { + assert!(is_contiguous(&[3, 1], &[1, 1])); + } + + #[test] + fn is_contiguous_basic() { + assert!(is_contiguous(&[32, 32], &[32, 1])); + } + + #[test] + fn is_contiguous_permuted() { + assert!(!is_contiguous(&[32, 32], &[1, 32])); + } + + #[test] + fn is_contiguous_slice() { + assert!(!is_contiguous(&[32, 1, 64], &[32, 64, 1])); + } + + #[test] + fn is_contiguous_4d_positive() { + assert!(is_contiguous(&[8, 256, 32, 32], &[262144, 1024, 32, 1])); + } + + #[test] + fn is_contiguous_4d_negative() { + assert!(!is_contiguous(&[256, 8, 32, 32], &[1024, 262144, 32, 1])); + } + + /// Based on a bug encountered in interpolate_1d + #[test] + fn is_contiguous_4d_unit_shape() { + assert!(!is_contiguous(&[1, 1, 1, 9], &[72, 1, 72, 8])); + } +} diff --git a/crates/burn-cubecl/src/tensor/mod.rs b/crates/burn-cubecl/src/tensor/mod.rs new file mode 100644 index 0000000..013604d --- /dev/null +++ b/crates/burn-cubecl/src/tensor/mod.rs @@ -0,0 +1,5 @@ +mod base; +mod quantization; + +pub use base::*; +pub use quantization::*; diff --git a/crates/burn-cubecl/src/tensor/quantization.rs b/crates/burn-cubecl/src/tensor/quantization.rs new file mode 100644 index 0000000..c77b3ba --- /dev/null +++ b/crates/burn-cubecl/src/tensor/quantization.rs @@ -0,0 +1,122 @@ +use burn_backend::{DType, Shape, TensorMetadata as _, quantization::QParamTensor}; +use burn_std::{Metadata, Strides}; +use cubecl::quant::scheme::{QuantStore, QuantValue}; +use cubecl::{client::ComputeClient, server::Handle}; + +use crate::CubeRuntime; + +use super::CubeTensor; + +/// Runtime parameters for quantization. Can be used to construct a scales handle from the base +/// tensor handle. +pub type QParams = burn_backend::quantization::QParams; + +impl CubeTensor { + /// Create a new quantized tensor + pub fn new_quantized( + client: ComputeClient, + handle: Handle, + shape: Shape, + device: R::Device, + strides: Strides, + dtype: DType, + qparams: QParams, + ) -> Self { + CubeTensor { + client, + handle, + meta: Box::new(Metadata::new(shape, strides)), + device, + dtype, + qparams: Some(qparams), + } + } + + /// Returns the two tensors: (values, params) for a quantized tensor. + /// For the values, native types that aren't supported as a normal `DType` will be returned + /// as an unsigned integer tensor representing the bits. Should be reconstructed using `from_bits` + /// in kernels. + pub fn quantized_handles(&self) -> Option<(CubeTensor, CubeTensor)> { + let params = self.scales()?; + let scheme = match self.dtype { + DType::QFloat(sc) => sc, + _ => return None, + }; + let values = match scheme.store { + QuantStore::Native => match scheme.value { + QuantValue::Q8F | QuantValue::Q8S => CubeTensor { + client: self.client.clone(), + handle: self.handle.clone(), + meta: self.meta.clone(), + device: self.device.clone(), + dtype: DType::I8, + qparams: None, + }, + QuantValue::E4M3 | QuantValue::E5M2 => CubeTensor { + client: self.client.clone(), + handle: self.handle.clone(), + meta: self.meta.clone(), + device: self.device.clone(), + dtype: DType::U8, + qparams: None, + }, + QuantValue::Q4F + | QuantValue::Q4S + | QuantValue::Q2F + | QuantValue::Q2S + | QuantValue::E2M1 => { + panic!("Can't store native sub-byte values") + } + }, + QuantStore::PackedU32(packed_dim) => { + let packed_dim = self.rank() - packed_dim - 1; + let mut shape = self.shape(); + shape[packed_dim] = shape[packed_dim].div_ceil(scheme.num_quants()); + + CubeTensor { + client: self.client.clone(), + handle: self.handle.clone(), + meta: Box::new(Metadata::new(shape, self.meta.strides.clone())), + device: self.device.clone(), + dtype: DType::U32, + qparams: None, + } + } + QuantStore::PackedNative(packed_dim) => match scheme.value { + QuantValue::E2M1 => { + let packed_dim = self.rank() - packed_dim - 1; + let mut shape = self.shape(); + shape[packed_dim] = shape[packed_dim].div_ceil(scheme.num_quants()); + + CubeTensor { + client: self.client.clone(), + handle: self.handle.clone(), + meta: Box::new(Metadata::new(shape, self.meta.strides.clone())), + device: self.device.clone(), + dtype: DType::U8, + qparams: None, + } + } + other => panic!("{other:?} doesn't support native packing"), + }, + }; + + Some((values, params)) + } + + /// Construct a separate tensor for the quantization scales, if present + pub fn scales(&self) -> Option> { + let qparams = self.qparams.as_ref()?; + let mut handle = self.handle.clone(); + handle.offset_start = Some(qparams.scales.offset_start as u64); + handle.offset_end = Some(qparams.scales.offset_end as u64); + + Some(CubeTensor::new( + self.client.clone(), + handle, + qparams.scales.metadata.clone(), + self.device.clone(), + qparams.scales.dtype, + )) + } +} diff --git a/crates/burn-cubecl/src/tune_key.rs b/crates/burn-cubecl/src/tune_key.rs new file mode 100644 index 0000000..d6892c7 --- /dev/null +++ b/crates/burn-cubecl/src/tune_key.rs @@ -0,0 +1,30 @@ +use crate::kernel::{ + conv::{ConvAutotuneKey, ConvTranspose2dAutotuneKey}, + reduce::SumAutotuneKey, +}; +use cubecl::tune::AutotuneKey; +use serde::{Deserialize, Serialize}; +use std::fmt::Display; + +#[derive(Hash, Eq, PartialEq, Debug, Clone, Serialize, Deserialize)] +/// Key for all autotune-enabled operations +pub enum CubeAutotuneKey { + /// Key for sum operations + Sum(SumAutotuneKey), + /// Key for convolution operations + Conv(ConvAutotuneKey), + /// Key for transpose convolution operations + ConvTranspose(ConvTranspose2dAutotuneKey), +} + +impl Display for CubeAutotuneKey { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + CubeAutotuneKey::Sum(reduce_key) => std::fmt::Debug::fmt(&reduce_key, f), + CubeAutotuneKey::Conv(conv_key) => std::fmt::Debug::fmt(&conv_key, f), + CubeAutotuneKey::ConvTranspose(conv_key) => std::fmt::Debug::fmt(&conv_key, f), + } + } +} + +impl AutotuneKey for CubeAutotuneKey {} diff --git a/crates/burn-cuda/Cargo.toml b/crates/burn-cuda/Cargo.toml new file mode 100644 index 0000000..5591536 --- /dev/null +++ b/crates/burn-cuda/Cargo.toml @@ -0,0 +1,42 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science"] +description = "CUDA backend for the Burn framework" +documentation = "https://docs.rs/burn-cuda" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "gpu", "cuda"] +license.workspace = true +name = "burn-cuda" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-cuda" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["std", "fusion", "autotune", "burn-cubecl/default", "cubecl/default"] +doc = ["burn-cubecl/doc"] +std = ["burn-cubecl/std", "cubecl/std"] +tracing = [ + "burn-backend/tracing", + "burn-cubecl/tracing", + "burn-fusion?/tracing", + "cubecl/tracing", +] + +fusion = ["burn-fusion", "burn-cubecl/fusion"] +autotune = ["burn-cubecl/autotune"] +autotune-checks = ["burn-cubecl/autotune-checks"] + +[dependencies] +burn-fusion = { workspace = true, optional = true, features = ["default"] } +burn-cubecl = { workspace = true } +burn-backend = { workspace = true, features = [ + "cubecl-cuda", +] } +cubecl = { workspace = true, features = ["cuda"] } + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-cuda/README.md b/crates/burn-cuda/README.md new file mode 100644 index 0000000..bf23978 --- /dev/null +++ b/crates/burn-cuda/README.md @@ -0,0 +1,30 @@ +# Burn CUDA Backend + +[Burn](https://github.com/tracel-ai/burn) CUDA backend + +[![Current Crates.io Version](https://img.shields.io/crates/v/burn-cuda.svg)](https://crates.io/crates/burn-cuda) +[![license](https://shields.io/badge/license-MIT%2FApache--2.0-blue)](https://github.com/tracel-ai/burn-cuda/blob/master/README.md) + +This crate provides a CUDA backend for [Burn](https://github.com/tracel-ai/burn) using the +[cubecl](https://github.com/tracel-ai/cubecl.git) and [cudarc](https://github.com/coreylowman/cudarc.git) +crates. + +## Usage Example + +```rust +#[cfg(feature = "cuda")] +mod cuda { + use burn_autodiff::Autodiff; + use burn_cuda::{Cuda, CudaDevice}; + use mnist::training; + + pub fn run() { + let device = CudaDevice::default(); + training::run::>>(device); + } +} +``` + +## Dependencies + +Requires CUDA 12.x to be installed and on the `PATH`. \ No newline at end of file diff --git a/crates/burn-cuda/src/lib.rs b/crates/burn-cuda/src/lib.rs new file mode 100644 index 0000000..b2ae199 --- /dev/null +++ b/crates/burn-cuda/src/lib.rs @@ -0,0 +1,44 @@ +#![cfg_attr(docsrs, feature(doc_cfg))] + +extern crate alloc; + +use burn_cubecl::CubeBackend; +pub use cubecl::cuda::CudaDevice; +use cubecl::cuda::CudaRuntime; + +#[cfg(not(feature = "fusion"))] +pub type Cuda = CubeBackend; + +#[cfg(feature = "fusion")] +pub type Cuda = burn_fusion::Fusion>; + +#[cfg(all(test, not(target_os = "macos")))] +mod tests { + use super::*; + use burn_backend::{Backend, BoolStore, DType, DeviceOps}; + + #[test] + fn should_support_dtypes() { + type B = Cuda; + let device = CudaDevice::default(); + let scheme = device.defaults().quantization.scheme; + + assert!(B::supports_dtype(&device, DType::F32)); + assert!(B::supports_dtype(&device, DType::Flex32)); + assert!(B::supports_dtype(&device, DType::F16)); + assert!(B::supports_dtype(&device, DType::BF16)); + assert!(B::supports_dtype(&device, DType::I64)); + assert!(B::supports_dtype(&device, DType::I32)); + assert!(B::supports_dtype(&device, DType::I16)); + assert!(B::supports_dtype(&device, DType::I8)); + assert!(B::supports_dtype(&device, DType::U64)); + assert!(B::supports_dtype(&device, DType::U32)); + assert!(B::supports_dtype(&device, DType::U16)); + assert!(B::supports_dtype(&device, DType::U8)); + assert!(B::supports_dtype(&device, DType::Bool(BoolStore::Native))); + assert!(B::supports_dtype(&device, DType::QFloat(scheme))); + + // Currently not registered in supported types + assert!(!B::supports_dtype(&device, DType::F64)); + } +} diff --git a/crates/burn-dataset/Cargo.toml b/crates/burn-dataset/Cargo.toml new file mode 100644 index 0000000..2cbd2f6 --- /dev/null +++ b/crates/burn-dataset/Cargo.toml @@ -0,0 +1,85 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science"] +description = "Library with simple dataset APIs for creating ML data pipelines" +documentation = "https://docs.rs/burn-dataset" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "data"] +license.workspace = true +name = "burn-dataset" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-dataset" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["sqlite-bundled"] +doc = ["default"] +tracing = [ + "burn-std/tracing", +] + +audio = ["hound"] +builtin-sources = ["vision", "dep:tar", "nlp"] +fake = ["dep:fake"] +network = ["dep:burn-std"] +sqlite = ["__sqlite-shared", "dep:rusqlite"] +sqlite-bundled = ["__sqlite-shared", "rusqlite/bundled"] +vision = ["dep:flate2", "dep:globwalk", "dep:image", "network"] +nlp = ["dep:zip", "dep:encoding_rs"] +# internal +__sqlite-shared = [ + "dep:r2d2", + "dep:r2d2_sqlite", + "dep:serde_rusqlite", + "dep:image", + "dep:gix-tempfile", +] +dataframe = ["dep:polars", "dep:planus"] + +[dependencies] +burn-std = { workspace = true, optional = true, features = [ + "default", + "network", +] } +csv = { workspace = true } +derive-new = { workspace = true } +dirs = { workspace = true } +fake = { workspace = true, optional = true } +flate2 = { workspace = true, optional = true } +gix-tempfile = { workspace = true, optional = true } +globwalk = { workspace = true, optional = true } +hound = { workspace = true, optional = true } +image = { workspace = true, optional = true } +planus = { workspace = true, optional = true } +encoding_rs = { workspace = true, optional = true } +polars = { workspace = true, optional = true } +r2d2 = { workspace = true, optional = true } +r2d2_sqlite = { workspace = true, optional = true } +rand = { workspace = true, features = ["std", "sys_rng"] } +zip = { workspace = true, optional = true } +rmp-serde = { workspace = true } +rusqlite = { workspace = true, optional = true } +sanitize-filename = { workspace = true } +serde = { workspace = true, features = ["std", "derive"] } +serde_json = { workspace = true, features = ["std"] } +serde_rusqlite = { workspace = true, optional = true } +strum = { workspace = true } +tar = { workspace = true, optional = true } +tempfile = { workspace = true } +thiserror = { workspace = true } + + +[dev-dependencies] +fake = { workspace = true } +rayon = { workspace = true } +rstest = { workspace = true } + +[package.metadata.cargo-udeps.ignore] +normal = ["strum", "strum_macros"] + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-dataset/LICENSE-APACHE b/crates/burn-dataset/LICENSE-APACHE new file mode 120000 index 0000000..1cd601d --- /dev/null +++ b/crates/burn-dataset/LICENSE-APACHE @@ -0,0 +1 @@ +../../LICENSE-APACHE \ No newline at end of file diff --git a/crates/burn-dataset/LICENSE-MIT b/crates/burn-dataset/LICENSE-MIT new file mode 120000 index 0000000..b2cfbdc --- /dev/null +++ b/crates/burn-dataset/LICENSE-MIT @@ -0,0 +1 @@ +../../LICENSE-MIT \ No newline at end of file diff --git a/crates/burn-dataset/README.md b/crates/burn-dataset/README.md new file mode 100644 index 0000000..742aa23 --- /dev/null +++ b/crates/burn-dataset/README.md @@ -0,0 +1,17 @@ +# Burn Dataset + +> [Burn](https://github.com/tracel-ai/burn) dataset library + +[![Current Crates.io Version](https://img.shields.io/crates/v/burn-dataset.svg)](https://crates.io/crates/burn-dataset) +[![license](https://shields.io/badge/license-MIT%2FApache--2.0-blue)](https://github.com/tracel-ai/burn-dataset/blob/master/README.md) + +The Burn Dataset library is designed to streamline your machine learning (ML) data pipeline creation +process. It offers a variety of dataset implementations, transformation functions, and data sources. + +## Feature Flags + +- `audio` - enables audio dataset (SpeechCommandsDataset). Run the following example to try it out: + + ```shell + cargo run --example speech_commands --features audio + ``` diff --git a/crates/burn-dataset/examples/hf_dataset.rs b/crates/burn-dataset/examples/hf_dataset.rs new file mode 100644 index 0000000..07e6bcf --- /dev/null +++ b/crates/burn-dataset/examples/hf_dataset.rs @@ -0,0 +1,22 @@ +use burn_dataset::HuggingfaceDatasetLoader; +use burn_dataset::SqliteDataset; +use serde::Deserialize; + +#[derive(Deserialize, Debug, Clone)] +struct MnistItemRaw { + pub _image_bytes: Vec, + pub _label: usize, +} +fn main() { + // There are some datasets, such as https://huggingface.co/datasets/ylecun/mnist/tree/main that contains a script, + // In this cases you must enable trusting remote code execution if you want to use it. + let _train_ds: SqliteDataset = HuggingfaceDatasetLoader::new("ylecun/mnist") + .with_trust_remote_code(true) + .dataset("train") + .unwrap(); + + // However not all dataset requires it https://huggingface.co/datasets/Anthropic/hh-rlhf/tree/main + let _train_ds: SqliteDataset = HuggingfaceDatasetLoader::new("Anthropic/hh-rlhf") + .dataset("train") + .unwrap(); +} diff --git a/crates/burn-dataset/examples/speech_commands.rs b/crates/burn-dataset/examples/speech_commands.rs new file mode 100644 index 0000000..7efec10 --- /dev/null +++ b/crates/burn-dataset/examples/speech_commands.rs @@ -0,0 +1,23 @@ +#[cfg(feature = "audio")] +use burn_dataset::{Dataset, audio::SpeechCommandsDataset}; + +#[cfg(feature = "audio")] +fn speech_command() { + let index: usize = 4835; + let test = SpeechCommandsDataset::test(); + let item = test.get(index).unwrap(); + + println!("Item: {:?}", item); + println!("Item Length: {:?}", item.audio_samples.len()); + println!("Label: {}", item.label); + + assert_eq!(test.len(), 4890); + assert_eq!(item.label.to_string(), "Yes"); + assert_eq!(item.sample_rate, 16000); + assert_eq!(item.audio_samples.len(), 16000); +} + +fn main() { + #[cfg(feature = "audio")] + speech_command() +} diff --git a/crates/burn-dataset/src/audio/mod.rs b/crates/burn-dataset/src/audio/mod.rs new file mode 100644 index 0000000..5d357e9 --- /dev/null +++ b/crates/burn-dataset/src/audio/mod.rs @@ -0,0 +1,3 @@ +mod speech_commands; + +pub use speech_commands::*; diff --git a/crates/burn-dataset/src/audio/speech_commands.rs b/crates/burn-dataset/src/audio/speech_commands.rs new file mode 100644 index 0000000..d291a73 --- /dev/null +++ b/crates/burn-dataset/src/audio/speech_commands.rs @@ -0,0 +1,208 @@ +use crate::{ + Dataset, HuggingfaceDatasetLoader, SqliteDataset, + transform::{Mapper, MapperDataset}, +}; + +use hound::WavReader; +use serde::{Deserialize, Serialize}; +use strum::{Display, EnumCount, FromRepr}; + +type MappedDataset = MapperDataset, ConvertSamples, SpeechItemRaw>; + +/// Enum representing speech command classes in the Speech Commands dataset. +/// Class names are based on the Speech Commands dataset from Huggingface. +/// See [speech_commands](https://huggingface.co/datasets/google/speech_commands) +/// for more information. +#[allow(missing_docs)] +#[derive(Debug, Display, Clone, Copy, FromRepr, Serialize, Deserialize, EnumCount)] +pub enum SpeechCommandClass { + // Target command words + Yes = 0, + No = 1, + Up = 2, + Down = 3, + Left = 4, + Right = 5, + On = 6, + Off = 7, + Stop = 8, + Go = 9, + Zero = 10, + One = 11, + Two = 12, + Three = 13, + Four = 14, + Five = 15, + Six = 16, + Seven = 17, + Eight = 18, + Nine = 19, + + // Non-target words that can be grouped into "Other" + Bed = 20, + Bird = 21, + Cat = 22, + Dog = 23, + Happy = 24, + House = 25, + Marvin = 26, + Sheila = 27, + Tree = 28, + Wow = 29, + + // Commands from v2 dataset, that can be grouped into "Other" + Backward = 30, + Forward = 31, + Follow = 32, + Learn = 33, + Visual = 34, + + // Background noise + Silence = 35, + + // Other miscellaneous words + Other = 36, +} + +/// Struct containing raw speech data returned from a database. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct SpeechItemRaw { + /// Audio file bytes. + pub audio_bytes: Vec, + + /// Label index. + pub label: usize, + + /// Indicates if the label is unknown. + pub is_unknown: bool, +} + +/// Speech item with audio samples and label. +/// +/// The audio samples are floats in the range [-1.0, 1.0]. +/// The sample rate is in Hz. +/// The label is the class index (see [SpeechCommandClass]). +/// To convert to usize simply use `as usize`. To convert label to string use `.to_string()`. +/// +/// The original label is also stored in the `label_original` field for debugging and remapping if needed. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct SpeechItem { + /// Audio samples in the range [-1.0, 1.0]. + pub audio_samples: Vec, + + /// The sample rate of the audio. + pub sample_rate: usize, + + /// The label of the audio. + pub label: SpeechCommandClass, +} + +/// Speech Commands dataset from Huggingface v0.02. +/// See [Speech Commands dataset](https://huggingface.co/datasets/google/speech_commands). +/// +/// The data is downloaded from Huggingface and stored in a SQLite database (3.0 GB). +/// The dataset contains 99,720 audio samples of 2,607 people saying 35 different words. +/// +/// NOTE: The most samples are under 1 second long but there are some with pure background noise that +/// need splitting into shorter segmants. +/// +/// The labels are 20 target words, silence and other words. +/// +/// The dataset is split into 3 parts: +/// - train: 84,848 audio files +/// - test: 4,890 audio files +/// - validation: 9,982 audio files +pub struct SpeechCommandsDataset { + dataset: MappedDataset, +} + +impl SpeechCommandsDataset { + /// Create a new dataset with the given split. + pub fn new(split: &str) -> Self { + let dataset: SqliteDataset = + HuggingfaceDatasetLoader::new("google/speech_commands") + .with_subset("v0.02") + .dataset(split) + .unwrap(); + let dataset = MapperDataset::new(dataset, ConvertSamples); + Self { dataset } + } + + /// Create a new dataset with the train split. + pub fn train() -> Self { + Self::new("train") + } + + /// Create a new dataset with the test split. + pub fn test() -> Self { + Self::new("test") + } + + /// Create a new dataset with the validation split. + pub fn validation() -> Self { + Self::new("validation") + } + + /// Returns the number of classes in the dataset + pub fn num_classes() -> usize { + SpeechCommandClass::COUNT + } +} + +impl Dataset for SpeechCommandsDataset { + fn get(&self, index: usize) -> Option { + self.dataset.get(index) + } + + fn len(&self) -> usize { + self.dataset.len() + } +} + +/// Mapper converting audio bytes into audio samples and the label to enum class. +struct ConvertSamples; + +impl ConvertSamples { + /// Convert label to enum class. + fn to_speechcommandclass(label: usize) -> SpeechCommandClass { + SpeechCommandClass::from_repr(label).unwrap() + } + + /// Convert audio bytes into samples of floats [-1.0, 1.0]. + fn to_audiosamples(bytes: &Vec) -> (Vec, usize) { + let reader = WavReader::new(bytes.as_slice()).unwrap(); + let spec = reader.spec(); + + // Maximum value of the audio samples (using bit shift to raise 2 to the power of bits per sample). + let max_value = (1 << (spec.bits_per_sample - 1)) as f32; + + // The sample rate of the audio. + let sample_rate = spec.sample_rate as usize; + + // Convert the audio samples to floats [-1.0, 1.0]. + let audio_samples: Vec = reader + .into_samples::() + .filter_map(Result::ok) + .map(|sample| sample as f32 / max_value) + .collect(); + + (audio_samples, sample_rate) + } +} + +impl Mapper for ConvertSamples { + /// Convert audio bytes into samples of floats [-1.0, 1.0] + /// and the label to enum class with the target word, other and silence classes. + fn map(&self, item: &SpeechItemRaw) -> SpeechItem { + let (audio_samples, sample_rate) = Self::to_audiosamples(&item.audio_bytes); + + // Convert the label to enum class, with the target words, other and silence classes. + let label = Self::to_speechcommandclass(item.label); + + SpeechItem { + audio_samples, + sample_rate, + label, + } + } +} diff --git a/crates/burn-dataset/src/dataset/base.rs b/crates/burn-dataset/src/dataset/base.rs new file mode 100644 index 0000000..eb53980 --- /dev/null +++ b/crates/burn-dataset/src/dataset/base.rs @@ -0,0 +1,71 @@ +use std::sync::Arc; + +use crate::DatasetIterator; + +/// The dataset trait defines a basic collection of items with a predefined size. +pub trait Dataset: Send + Sync { + /// Gets the item at the given index. + fn get(&self, index: usize) -> Option; + + /// Gets the number of items in the dataset. + fn len(&self) -> usize; + + /// Checks if the dataset is empty. + fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Returns an iterator over the dataset. + fn iter(&self) -> DatasetIterator<'_, I> + where + Self: Sized, + { + DatasetIterator::new(self) + } +} + +impl Dataset for Arc +where + D: Dataset, +{ + fn get(&self, index: usize) -> Option { + self.as_ref().get(index) + } + + fn len(&self) -> usize { + self.as_ref().len() + } +} + +impl Dataset for Arc> { + fn get(&self, index: usize) -> Option { + self.as_ref().get(index) + } + + fn len(&self) -> usize { + self.as_ref().len() + } +} + +impl Dataset for Box +where + D: Dataset, +{ + fn get(&self, index: usize) -> Option { + self.as_ref().get(index) + } + + fn len(&self) -> usize { + self.as_ref().len() + } +} + +impl Dataset for Box> { + fn get(&self, index: usize) -> Option { + self.as_ref().get(index) + } + + fn len(&self) -> usize { + self.as_ref().len() + } +} diff --git a/crates/burn-dataset/src/dataset/dataframe.rs b/crates/burn-dataset/src/dataset/dataframe.rs new file mode 100644 index 0000000..c1617db --- /dev/null +++ b/crates/burn-dataset/src/dataset/dataframe.rs @@ -0,0 +1,464 @@ +use std::marker::PhantomData; + +use crate::Dataset; + +use polars::frame::row::Row; +use polars::prelude::*; +use serde::de::DeserializeSeed; +use serde::{ + Deserialize, + de::{self, DeserializeOwned, Deserializer, SeqAccess, Visitor}, + forward_to_deserialize_any, +}; + +/// Error type for DataframeDataset +#[derive(thiserror::Error, Debug)] +pub enum DataframeDatasetError { + /// Error occurred during deserialization or other operations + #[error("{0}")] + Other(String), +} + +impl de::Error for DataframeDatasetError { + fn custom(msg: T) -> Self { + DataframeDatasetError::Other(msg.to_string()) + } +} + +/// Dataset implementation for Polars DataFrame +/// +/// This struct provides a way to access data from a Polars DataFrame +/// as if it were a Dataset of type I. +pub struct DataframeDataset { + df: DataFrame, + len: usize, + column_name_mapping: Vec, + phantom: PhantomData, +} + +impl DataframeDataset +where + I: Clone + Send + Sync + DeserializeOwned, +{ + /// Create a new DataframeDataset from a Polars DataFrame + /// + /// # Arguments + /// + /// * `df` - A Polars DataFrame + /// + /// # Returns + /// + /// A Result containing the new DataframeDataset or a DataframeDatasetError + pub fn new(df: DataFrame) -> Result { + let len = df.height(); + let field_names = extract_field_names::(); + + let column_name_mapping = field_names + .iter() + .map(|name| { + df.schema() + .try_get_full(name) + .map(|(index, _, _)| index) + .map_err(|err| DataframeDatasetError::Other(err.to_string())) + }) + .collect::, _>>()?; + + Ok(DataframeDataset { + df, + len, + column_name_mapping, + phantom: PhantomData, + }) + } +} + +impl Dataset for DataframeDataset +where + I: Clone + Send + Sync + DeserializeOwned, +{ + /// Get an item from the dataset at the specified index + /// + /// # Arguments + /// + /// * `index` - The index of the item to retrieve + /// + /// # Returns + /// + /// An Option containing the item if it exists, or None if it doesn't + fn get(&self, index: usize) -> Option { + let row = self.df.get_row(index).ok()?; + + let mut deserializer = RowDeserializer::new(&row, &self.column_name_mapping); + I::deserialize(&mut deserializer).ok() + } + + /// Get the length of the dataset + fn len(&self) -> usize { + self.len + } + + /// Check if the dataset is empty + fn is_empty(&self) -> bool { + self.len == 0 + } +} + +/// A deserializer for Polars DataFrame rows +struct RowDeserializer<'a> { + row: &'a Row<'a>, + column_name_mapping: &'a Vec, + index: usize, +} + +impl<'a> RowDeserializer<'a> { + /// Create a new RowDeserializer + /// + /// # Arguments + /// + /// * `row` - A reference to a Polars DataFrame row + /// * `column_name_mapping` - A reference to a vector mapping field names to column indices + fn new(row: &'a Row, column_name_mapping: &'a Vec) -> RowDeserializer<'a> { + RowDeserializer { + row, + column_name_mapping, + index: 0, + } + } +} + +impl<'de, 'a> Deserializer<'de> for &'a mut RowDeserializer<'a> { + type Error = DataframeDatasetError; + + fn deserialize_any(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + let i = self.column_name_mapping[self.index]; + + let value = &self.row.0[i]; + match value { + AnyValue::Null => visitor.visit_none(), + AnyValue::Boolean(b) => visitor.visit_bool(*b), + AnyValue::Int8(i) => visitor.visit_i8(*i), + AnyValue::Int16(i) => visitor.visit_i16(*i), + AnyValue::Int32(i) => visitor.visit_i32(*i), + AnyValue::Int64(i) => visitor.visit_i64(*i), + AnyValue::UInt8(i) => visitor.visit_u8(*i), + AnyValue::UInt16(i) => visitor.visit_u16(*i), + AnyValue::UInt32(i) => visitor.visit_u32(*i), + AnyValue::UInt64(i) => visitor.visit_u64(*i), + AnyValue::Float32(f) => visitor.visit_f32(*f), + AnyValue::Float64(f) => visitor.visit_f64(*f), + AnyValue::Date(i) => visitor.visit_i32(*i), + AnyValue::String(s) => visitor.visit_string(s.to_string()), + AnyValue::Binary(b) => { + visitor.visit_seq(de::value::SeqDeserializer::new(b.iter().copied())) + } + AnyValue::Time(t) => visitor.visit_i64(*t), + ty => Err(DataframeDatasetError::Other( + format!("Unsupported type: {ty:?}").to_string(), + )), + } + } + + fn deserialize_struct( + self, + _name: &'static str, + _fields: &'static [&'static str], + visitor: V, + ) -> Result + where + V: Visitor<'de>, + { + visitor.visit_seq(self) + } + + forward_to_deserialize_any! { + bool i8 i16 i32 i64 u8 u16 u32 u64 f32 f64 char str string + bytes byte_buf option unit unit_struct newtype_struct seq tuple + tuple_struct map enum identifier ignored_any + } +} + +impl<'de, 'a> SeqAccess<'de> for RowDeserializer<'a> { + type Error = DataframeDatasetError; + + fn next_element_seed(&mut self, seed: T) -> Result, DataframeDatasetError> + where + T: DeserializeSeed<'de>, + { + if self.index >= self.row.0.len() { + return Ok(None); + } + let mut deserializer = RowDeserializer { + row: self.row, + column_name_mapping: self.column_name_mapping, + index: self.index, + }; + self.index += 1; + seed.deserialize(&mut deserializer).map(Some) + } +} + +struct FieldExtractor { + fields: Vec<&'static str>, +} + +impl<'de> Deserializer<'de> for &mut FieldExtractor { + type Error = de::value::Error; + + fn deserialize_any(self, _visitor: V) -> core::result::Result + where + V: Visitor<'de>, + { + Err(de::Error::custom("Field extractor")) + } + + fn deserialize_struct( + self, + _name: &'static str, + fields: &'static [&'static str], + _visitor: V, + ) -> core::result::Result + where + V: Visitor<'de>, + { + self.fields.extend_from_slice(fields); + Err(de::Error::custom("Field extractor")) + } + + forward_to_deserialize_any! { + bool i8 i16 i32 i64 u8 u16 u32 u64 f32 f64 char str string bytes + byte_buf option unit unit_struct newtype_struct seq tuple + tuple_struct map enum identifier ignored_any + } +} + +/// Extract field names from a type T that implements Deserialize +/// +/// # Returns +/// +/// A vector of field names as static string slices +fn extract_field_names<'de, T>() -> Vec<&'static str> +where + T: Deserialize<'de>, +{ + let mut extractor = FieldExtractor { fields: Vec::new() }; + let _ = T::deserialize(&mut extractor); + extractor.fields +} + +#[cfg(test)] +mod tests { + use polars::prelude::*; + use serde::Deserialize; + + use super::*; + #[derive(Clone, Debug, Deserialize, PartialEq)] + struct TestData { + int32: i32, + bool: bool, + float64: f64, + string: String, + int16: i16, + uint32: u32, + uint64: u64, + float32: f32, + int64: i64, + int8: i8, + binary: Vec, + } + + fn create_test_dataframe() -> DataFrame { + let s0 = Column::new("int32".into(), &[1i32, 2i32, 3i32]); + let s1 = Column::new("bool".into(), &[true, false, true]); + let s2 = Column::new("float64".into(), &[1.1f64, 2.2f64, 3.3f64]); + let s3 = Column::new("string".into(), &["Boo", "Boo2", "Boo3"]); + let s6 = Column::new("int16".into(), &[1i16, 2i16, 3i16]); + let s8 = Column::new("uint32".into(), &[1u32, 2u32, 3u32]); + let s9 = Column::new("uint64".into(), &[1u64, 2u64, 3u64]); + let s10 = Column::new("float32".into(), &[1.1f32, 2.2f32, 3.3f32]); + let s11 = Column::new("int64".into(), &[1i64, 2i64, 3i64]); + let s12 = Column::new("int8".into(), &[1i8, 2i8, 3i8]); + + let binary_data: Vec<&[u8]> = vec![&[1, 2, 3], &[4, 5, 6], &[7, 8, 9]]; + + let s13 = Column::new("binary".into(), binary_data); + DataFrame::new_infer_height(vec![s0, s1, s2, s3, s6, s8, s9, s10, s11, s12, s13]).unwrap() + } + + #[test] + fn test_dataframe_dataset_creation() { + let df = create_test_dataframe(); + let dataset = DataframeDataset::::new(df); + assert!(dataset.is_ok()); + } + + #[test] + fn test_dataframe_dataset_length() { + let df = create_test_dataframe(); + let dataset = DataframeDataset::::new(df).unwrap(); + assert_eq!(dataset.len(), 3); + assert!(!dataset.is_empty()); + } + + #[test] + fn test_dataframe_dataset_get() { + let df = create_test_dataframe(); + let dataset = DataframeDataset::::new(df).unwrap(); + + let expected_items = vec![ + TestData { + int32: 1, + bool: true, + float64: 1.1, + string: "Boo".to_string(), + int16: 1, + uint32: 1, + uint64: 1, + float32: 1.1, + int64: 1, + int8: 1, + binary: vec![1, 2, 3], + }, + TestData { + int32: 2, + bool: false, + float64: 2.2, + string: "Boo2".to_string(), + int16: 2, + uint32: 2, + uint64: 2, + float32: 2.2, + int64: 2, + int8: 2, + binary: vec![4, 5, 6], + }, + TestData { + int32: 3, + bool: true, + float64: 3.3, + string: "Boo3".to_string(), + int16: 3, + uint32: 3, + uint64: 3, + float32: 3.3, + int64: 3, + int8: 3, + binary: vec![7, 8, 9], + }, + ]; + + for (index, expected_item) in expected_items.iter().enumerate() { + let item = dataset.get(index).unwrap(); + assert_eq!(&item, expected_item); + } + } + + #[test] + fn test_dataframe_dataset_out_of_bounds() { + let df = create_test_dataframe(); + let dataset = DataframeDataset::::new(df).unwrap(); + assert!(dataset.get(3).is_none()); + } + + #[test] + fn test_dataframe_dataset() { + let df = create_test_dataframe(); + let dataset: DataframeDataset = DataframeDataset::new(df).unwrap(); + + assert_eq!(dataset.len(), 3); + assert!(!dataset.is_empty()); + + let item = dataset.get(1).unwrap(); + assert_eq!( + item, + TestData { + int32: 2, + bool: false, + float64: 2.2, + string: "Boo2".to_string(), + int16: 2, + uint32: 2, + uint64: 2, + float32: 2.2, + int64: 2, + int8: 2, + binary: vec![4, 5, 6], + } + ); + + let item = dataset.get(2).unwrap(); + + assert_eq!( + item, + TestData { + int32: 3, + bool: true, + float64: 3.3, + string: "Boo3".to_string(), + int16: 3, + uint32: 3, + uint64: 3, + float32: 3.3, + int64: 3, + int8: 3, + binary: vec![7, 8, 9], + } + ); + } + + #[test] + fn test_non_existing_struct_fields() { + #[derive(Clone, Debug, Deserialize, PartialEq)] + struct PartialTestData { + int32: i32, + bool: bool, + non_existent: String, + } + + let df = create_test_dataframe(); + let dataset = DataframeDataset::::new(df); + + assert!(dataset.is_err()); + if let Err(e) = dataset { + assert!(matches!(e, DataframeDatasetError::Other(_))); + } + } + + #[test] + fn test_partial_table() { + #[derive(Clone, Debug, Deserialize, PartialEq)] + struct PartialTestData { + int32: i32, + bool: bool, + string: String, + } + + let df = create_test_dataframe(); + let dataset = DataframeDataset::::new(df).unwrap(); + + assert_eq!(dataset.len(), 3); + assert!(!dataset.is_empty()); + + let item = dataset.get(1).unwrap(); + assert_eq!( + item, + PartialTestData { + int32: 2, + bool: false, + string: "Boo2".to_string(), + } + ); + + let item = dataset.get(2).unwrap(); + assert_eq!( + item, + PartialTestData { + int32: 3, + bool: true, + string: "Boo3".to_string(), + } + ); + } +} diff --git a/crates/burn-dataset/src/dataset/fake.rs b/crates/burn-dataset/src/dataset/fake.rs new file mode 100644 index 0000000..c27f8cf --- /dev/null +++ b/crates/burn-dataset/src/dataset/fake.rs @@ -0,0 +1,38 @@ +use crate::{Dataset, DatasetIterator, InMemDataset}; +use fake::{Dummy, Fake, Faker}; + +/// Dataset filled with fake items generated from the [fake](fake) crate. +pub struct FakeDataset { + dataset: InMemDataset, +} + +impl> FakeDataset { + /// Create a new fake dataset with the given size. + pub fn new(size: usize) -> Self { + let mut items = Vec::with_capacity(size); + for _ in 0..size { + items.push(Faker.fake()); + } + let dataset = InMemDataset::new(items); + + Self { dataset } + } +} + +impl Dataset for FakeDataset { + fn iter(&self) -> DatasetIterator<'_, I> { + DatasetIterator::new(self) + } + + fn get(&self, index: usize) -> Option { + self.dataset.get(index) + } + + fn len(&self) -> usize { + self.dataset.len() + } + + fn is_empty(&self) -> bool { + self.dataset.is_empty() + } +} diff --git a/crates/burn-dataset/src/dataset/in_memory.rs b/crates/burn-dataset/src/dataset/in_memory.rs new file mode 100644 index 0000000..13854cd --- /dev/null +++ b/crates/burn-dataset/src/dataset/in_memory.rs @@ -0,0 +1,192 @@ +use std::{ + fs::File, + io::{BufRead, BufReader}, + path::Path, +}; + +use serde::de::DeserializeOwned; + +use crate::Dataset; + +/// Dataset where all items are stored in ram. +pub struct InMemDataset { + items: Vec, +} + +impl InMemDataset { + /// Creates a new in memory dataset from the given items. + pub fn new(items: Vec) -> Self { + InMemDataset { items } + } +} + +impl Dataset for InMemDataset +where + I: Clone + Send + Sync, +{ + fn get(&self, index: usize) -> Option { + self.items.get(index).cloned() + } + fn len(&self) -> usize { + self.items.len() + } +} + +impl InMemDataset +where + I: Clone + DeserializeOwned, +{ + /// Create from a dataset. All items are loaded in memory. + pub fn from_dataset(dataset: &impl Dataset) -> Self { + let items: Vec = dataset.iter().collect(); + Self::new(items) + } + + /// Create from a json rows file (one json per line). + /// + /// [Supported field types](https://docs.rs/serde_json/latest/serde_json/value/enum.Value.html) + pub fn from_json_rows>(path: P) -> Result { + let file = File::open(path)?; + let reader = BufReader::new(file); + let mut items = Vec::new(); + + for line in reader.lines() { + let item = serde_json::from_str(line.unwrap().as_str()).unwrap(); + items.push(item); + } + + let dataset = Self::new(items); + + Ok(dataset) + } + + /// Create from a csv file. + /// + /// The provided `csv::ReaderBuilder` can be configured to fit your csv format. + /// + /// The supported field types are: String, integer, float, and bool. + /// + /// See: + /// - [Reading with Serde](https://docs.rs/csv/latest/csv/tutorial/index.html#reading-with-serde) + /// - [Delimiters, quotes and variable length records](https://docs.rs/csv/latest/csv/tutorial/index.html#delimiters-quotes-and-variable-length-records) + pub fn from_csv>( + path: P, + builder: &csv::ReaderBuilder, + ) -> Result { + let mut rdr = builder.from_path(path)?; + + let mut items = Vec::new(); + + for result in rdr.deserialize() { + let item: I = result?; + items.push(item); + } + + let dataset = Self::new(items); + + Ok(dataset) + } +} + +#[cfg(test)] +mod tests { + + use super::*; + use crate::{SqliteDataset, test_data}; + + use rstest::{fixture, rstest}; + use serde::{Deserialize, Serialize}; + + const DB_FILE: &str = "tests/data/sqlite-dataset.db"; + const JSON_FILE: &str = "tests/data/dataset.json"; + const CSV_FILE: &str = "tests/data/dataset.csv"; + const CSV_FMT_FILE: &str = "tests/data/dataset-fmt.csv"; + + type SqlDs = SqliteDataset; + + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] + pub struct Sample { + column_str: String, + column_bytes: Vec, + column_int: i64, + column_bool: bool, + column_float: f64, + } + + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] + pub struct SampleCsv { + column_str: String, + column_int: i64, + column_bool: bool, + column_float: f64, + } + + #[fixture] + fn train_dataset() -> SqlDs { + SqliteDataset::from_db_file(DB_FILE, "train").unwrap() + } + + #[rstest] + pub fn from_dataset(train_dataset: SqlDs) { + let dataset = InMemDataset::from_dataset(&train_dataset); + + let non_existing_record_index: usize = 10; + let record_index: usize = 0; + + assert_eq!(train_dataset.get(non_existing_record_index), None); + assert_eq!(dataset.get(record_index).unwrap().column_str, "HI1"); + } + + #[test] + pub fn from_json_rows() { + let dataset = InMemDataset::::from_json_rows(JSON_FILE).unwrap(); + + let non_existing_record_index: usize = 10; + let record_index: usize = 1; + + assert_eq!(dataset.get(non_existing_record_index), None); + assert_eq!(dataset.get(record_index).unwrap().column_str, "HI2"); + assert!(!dataset.get(record_index).unwrap().column_bool); + } + + #[test] + pub fn from_csv_rows() { + let rdr = csv::ReaderBuilder::new(); + let dataset = InMemDataset::::from_csv(CSV_FILE, &rdr).unwrap(); + + let non_existing_record_index: usize = 10; + let record_index: usize = 1; + + assert_eq!(dataset.get(non_existing_record_index), None); + assert_eq!(dataset.get(record_index).unwrap().column_str, "HI2"); + assert_eq!(dataset.get(record_index).unwrap().column_int, 1); + assert!(!dataset.get(record_index).unwrap().column_bool); + assert_eq!(dataset.get(record_index).unwrap().column_float, 1.0); + } + + #[test] + pub fn from_csv_rows_fmt() { + let mut rdr = csv::ReaderBuilder::new(); + let rdr = rdr.delimiter(b' ').has_headers(false); + let dataset = InMemDataset::::from_csv(CSV_FMT_FILE, rdr).unwrap(); + + let non_existing_record_index: usize = 10; + let record_index: usize = 1; + + assert_eq!(dataset.get(non_existing_record_index), None); + assert_eq!(dataset.get(record_index).unwrap().column_str, "HI2"); + assert_eq!(dataset.get(record_index).unwrap().column_int, 1); + assert!(!dataset.get(record_index).unwrap().column_bool); + assert_eq!(dataset.get(record_index).unwrap().column_float, 1.0); + } + + #[test] + pub fn given_in_memory_dataset_when_iterate_should_iterate_though_all_items() { + let items_original = test_data::string_items(); + let dataset = InMemDataset::new(items_original.clone()); + + let items: Vec = dataset.iter().collect(); + + assert_eq!(items_original, items); + } +} diff --git a/crates/burn-dataset/src/dataset/iterator.rs b/crates/burn-dataset/src/dataset/iterator.rs new file mode 100644 index 0000000..cbea476 --- /dev/null +++ b/crates/burn-dataset/src/dataset/iterator.rs @@ -0,0 +1,31 @@ +use crate::dataset::Dataset; +use std::iter::Iterator; + +/// Dataset iterator. +pub struct DatasetIterator<'a, I> { + current: usize, + dataset: &'a dyn Dataset, +} + +impl<'a, I> DatasetIterator<'a, I> { + /// Creates a new dataset iterator. + pub fn new(dataset: &'a D) -> Self + where + D: Dataset, + { + DatasetIterator { + current: 0, + dataset, + } + } +} + +impl Iterator for DatasetIterator<'_, I> { + type Item = I; + + fn next(&mut self) -> Option { + let item = self.dataset.get(self.current); + self.current += 1; + item + } +} diff --git a/crates/burn-dataset/src/dataset/mod.rs b/crates/burn-dataset/src/dataset/mod.rs new file mode 100644 index 0000000..9d9061e --- /dev/null +++ b/crates/burn-dataset/src/dataset/mod.rs @@ -0,0 +1,25 @@ +mod base; +mod in_memory; +mod iterator; + +pub use base::*; +pub use in_memory::*; +pub use iterator::*; + +#[cfg(any(test, feature = "fake"))] +mod fake; + +#[cfg(any(test, feature = "fake"))] +pub use self::fake::*; + +#[cfg(feature = "dataframe")] +mod dataframe; + +#[cfg(feature = "dataframe")] +pub use dataframe::*; + +#[cfg(any(feature = "sqlite", feature = "sqlite-bundled"))] +pub use sqlite::*; + +#[cfg(any(feature = "sqlite", feature = "sqlite-bundled"))] +mod sqlite; diff --git a/crates/burn-dataset/src/dataset/sqlite.rs b/crates/burn-dataset/src/dataset/sqlite.rs new file mode 100644 index 0000000..a3cca83 --- /dev/null +++ b/crates/burn-dataset/src/dataset/sqlite.rs @@ -0,0 +1,851 @@ +use std::{ + collections::HashSet, + fs, io, + marker::PhantomData, + path::{Path, PathBuf}, + sync::{Arc, RwLock}, +}; + +use crate::Dataset; + +use gix_tempfile::{ + AutoRemove, ContainingDirectory, Handle, + handle::{Writable, persist}, +}; +use r2d2::{Pool, PooledConnection}; +use r2d2_sqlite::{ + SqliteConnectionManager, + rusqlite::{OpenFlags, OptionalExtension}, +}; +use sanitize_filename::sanitize; +use serde::{Serialize, de::DeserializeOwned}; +use serde_rusqlite::{columns_from_statement, from_row_with_columns}; + +/// Result type for the sqlite dataset. +pub type Result = core::result::Result; + +/// Sqlite dataset error. +#[derive(thiserror::Error, Debug)] +pub enum SqliteDatasetError { + /// IO related error. + #[error("IO error: {0}")] + Io(#[from] io::Error), + + /// Sql related error. + #[error("Sql error: {0}")] + Sql(#[from] serde_rusqlite::rusqlite::Error), + + /// Serde related error. + #[error("Serde error: {0}")] + Serde(#[from] rmp_serde::encode::Error), + + /// The database file already exists error. + #[error("Overwrite flag is set to false and the database file already exists: {0}")] + FileExists(PathBuf), + + /// Error when creating the connection pool. + #[error("Failed to create connection pool: {0}")] + ConnectionPool(#[from] r2d2::Error), + + /// Error when persisting the temporary database file. + #[error("Could not persist the temporary database file: {0}")] + PersistDbFile(#[from] persist::Error), + + /// Any other error. + #[error("{0}")] + Other(&'static str), +} + +impl From<&'static str> for SqliteDatasetError { + fn from(s: &'static str) -> Self { + SqliteDatasetError::Other(s) + } +} + +/// This struct represents a dataset where all items are stored in an SQLite database. +/// Each instance of this struct corresponds to a specific table within the SQLite database, +/// and allows for interaction with the data stored in the table in a structured and typed manner. +/// +/// The SQLite database must contain a table with the same name as the `split` field. This table should +/// have a primary key column named `row_id`, which is used to index the rows in the table. The `row_id` +/// should start at 1, while the corresponding dataset `index` should start at 0, i.e., `row_id` = `index` + 1. +/// +/// Table columns can be represented in two ways: +/// +/// 1. The table can have a column for each field in the `I` struct. In this case, the column names in the table +/// should match the field names of the `I` struct. The field names can be a subset of column names and +/// can be in any order. +/// +/// For the supported field types, refer to: +/// - [Serialization field types](https://docs.rs/serde_rusqlite/latest/serde_rusqlite) +/// - [SQLite data types](https://www.sqlite.org/datatype3.html) +/// +/// 2. The fields in the `I` struct can be serialized into a single column `item` in the table. In this case, the table +/// should have a single column named `item` of type `BLOB`. This is useful when the `I` struct contains complex fields +/// that cannot be mapped to a SQLite type, such as nested structs, vectors, etc. The serialization is done using +/// [MessagePack](https://msgpack.org/). +/// +/// Note: The code automatically figures out which of the above two cases is applicable, and uses the appropriate +/// method to read the data from the table. +#[derive(Debug)] +pub struct SqliteDataset { + db_file: PathBuf, + split: String, + conn_pool: Pool, + columns: Vec, + len: usize, + select_statement: String, + row_serialized: bool, + phantom: PhantomData, +} + +impl SqliteDataset { + /// Initializes a `SqliteDataset` from a SQLite database file and a split name. + pub fn from_db_file>(db_file: P, split: &str) -> Result { + // Create a connection pool + let conn_pool = create_conn_pool(&db_file, false)?; + + // Determine how the table is stored + let row_serialized = Self::check_if_row_serialized(&conn_pool, split)?; + + // Create a select statement and save it + let select_statement = if row_serialized { + format!("select item from {split} where row_id = ?") + } else { + format!("select * from {split} where row_id = ?") + }; + + // Save the column names and the number of rows + let (columns, len) = fetch_columns_and_len(&conn_pool, &select_statement, split)?; + + Ok(SqliteDataset { + db_file: db_file.as_ref().to_path_buf(), + split: split.to_string(), + conn_pool, + columns, + len, + select_statement, + row_serialized, + phantom: PhantomData, + }) + } + + /// Returns true if table has two columns: row_id (integer) and item (blob). + /// + /// This is used to determine if the table is row serialized or not. + fn check_if_row_serialized( + conn_pool: &Pool, + split: &str, + ) -> Result { + // This struct is used to store the column name and type + struct Column { + name: String, + ty: String, + } + + const COLUMN_NAME: usize = 1; + const COLUMN_TYPE: usize = 2; + + let sql_statement = format!("PRAGMA table_info({split})"); + + let conn = conn_pool.get()?; + + let mut stmt = conn.prepare(sql_statement.as_str())?; + let column_iter = stmt.query_map([], |row| { + Ok(Column { + name: row + .get::(COLUMN_NAME) + .unwrap() + .to_lowercase(), + ty: row + .get::(COLUMN_TYPE) + .unwrap() + .to_lowercase(), + }) + })?; + + let mut columns: Vec = vec![]; + + for column in column_iter { + columns.push(column?); + } + + if columns.len() != 2 { + Ok(false) + } else { + // Check if the column names and types match the expected values + Ok(columns[0].name == "row_id" + && columns[0].ty == "integer" + && columns[1].name == "item" + && columns[1].ty == "blob") + } + } + + /// Get the database file name. + pub fn db_file(&self) -> PathBuf { + self.db_file.clone() + } + + /// Get the split name. + pub fn split(&self) -> &str { + self.split.as_str() + } +} + +impl Dataset for SqliteDataset +where + I: Clone + Send + Sync + DeserializeOwned, +{ + /// Get an item from the dataset. + fn get(&self, index: usize) -> Option { + // Row ids start with 1 (one) and index starts with 0 (zero) + let row_id = index + 1; + + // Get a connection from the pool + let connection = self.conn_pool.get().unwrap(); + let mut statement = connection.prepare(self.select_statement.as_str()).unwrap(); + + if self.row_serialized { + // Fetch with a single column `item` and deserialize it with MessagePack + statement + .query_row([row_id], |row| { + // Deserialize item (blob) with MessagePack (rmp-serde) + Ok( + rmp_serde::from_slice::(row.get_ref(0).unwrap().as_blob().unwrap()) + .unwrap(), + ) + }) + .optional() //Converts Error (not found) to None + .unwrap() + } else { + // Fetch a row with multiple columns and deserialize it serde_rusqlite + statement + .query_row([row_id], |row| { + // Deserialize the row with serde_rusqlite + Ok(from_row_with_columns::(row, &self.columns).unwrap()) + }) + .optional() //Converts Error (not found) to None + .unwrap() + } + } + + /// Return the number of rows in the dataset. + fn len(&self) -> usize { + self.len + } +} + +/// Fetch the column names and the number of rows from the database. +fn fetch_columns_and_len( + conn_pool: &Pool, + select_statement: &str, + split: &str, +) -> Result<(Vec, usize)> { + // Save the column names + let connection = conn_pool.get()?; + let statement = connection.prepare(select_statement)?; + let columns = columns_from_statement(&statement); + + // Count the number of rows and save it as len + // + // NOTE: Using coalesce(max(row_id), 0) instead of count(*) because count(*) is super slow for large tables. + // The coalesce(max(row_id), 0) returns 0 if the table is empty, otherwise it returns the max row_id, + // which corresponds to the number of rows in the table. + // The main assumption, which always holds true, is that the row_id is always increasing and there are no gaps. + // This is true for all the datasets that we are using, otherwise row_id will not correspond to the index. + let mut statement = + connection.prepare(format!("select coalesce(max(row_id), 0) from {split}").as_str())?; + + let len = statement.query_row([], |row| { + let len: usize = row.get(0)?; + Ok(len) + })?; + Ok((columns, len)) +} + +/// Helper function to create a connection pool +fn create_conn_pool>( + db_file: P, + write: bool, +) -> Result> { + let sqlite_flags = if write { + OpenFlags::SQLITE_OPEN_READ_WRITE | OpenFlags::SQLITE_OPEN_CREATE + } else { + OpenFlags::SQLITE_OPEN_READ_ONLY + }; + + let manager = SqliteConnectionManager::file(db_file).with_flags(sqlite_flags); + Pool::new(manager).map_err(SqliteDatasetError::ConnectionPool) +} + +/// The `SqliteDatasetStorage` struct represents a SQLite database for storing datasets. +/// It consists of an optional name, a database file path, and a base directory for storage. +#[derive(Clone, Debug)] +pub struct SqliteDatasetStorage { + name: Option, + db_file: Option, + base_dir: Option, +} + +impl SqliteDatasetStorage { + /// Creates a new instance of `SqliteDatasetStorage` using a dataset name. + /// + /// # Arguments + /// + /// * `name` - A string slice that holds the name of the dataset. + pub fn from_name(name: &str) -> Self { + SqliteDatasetStorage { + name: Some(name.to_string()), + db_file: None, + base_dir: None, + } + } + + /// Creates a new instance of `SqliteDatasetStorage` using a database file path. + /// + /// # Arguments + /// + /// * `db_file` - A reference to the Path that represents the database file path. + pub fn from_file>(db_file: P) -> Self { + SqliteDatasetStorage { + name: None, + db_file: Some(db_file.as_ref().to_path_buf()), + base_dir: None, + } + } + + /// Sets the base directory for storing the dataset. + /// + /// # Arguments + /// + /// * `base_dir` - A string slice that represents the base directory. + pub fn with_base_dir>(mut self, base_dir: P) -> Self { + self.base_dir = Some(base_dir.as_ref().to_path_buf()); + self + } + + /// Checks if the database file exists in the given path. + /// + /// # Returns + /// + /// * A boolean value indicating whether the file exists or not. + pub fn exists(&self) -> bool { + self.db_file().exists() + } + + /// Fetches the database file path. + /// + /// # Returns + /// + /// * A `PathBuf` instance representing the file path. + pub fn db_file(&self) -> PathBuf { + match &self.db_file { + Some(db_file) => db_file.clone(), + None => { + let name = sanitize(self.name.as_ref().expect("Name is not set")); + Self::base_dir(self.base_dir.to_owned()).join(format!("{name}.db")) + } + } + } + + /// Determines the base directory for storing the dataset. + /// + /// # Arguments + /// + /// * `base_dir` - An `Option` that may contain a `PathBuf` instance representing the base directory. + /// + /// # Returns + /// + /// * A `PathBuf` instance representing the base directory. + pub fn base_dir(base_dir: Option) -> PathBuf { + match base_dir { + Some(base_dir) => base_dir, + None => dirs::cache_dir() + .expect("Could not get cache directory") + .join("burn-dataset"), + } + } + + /// Provides a writer instance for the SQLite dataset. + /// + /// # Arguments + /// + /// * `overwrite` - A boolean indicating if the existing database file should be overwritten. + /// + /// # Returns + /// + /// * A `Result` which is `Ok` if the writer could be created, `Err` otherwise. + pub fn writer(&self, overwrite: bool) -> Result> + where + I: Clone + Send + Sync + Serialize + DeserializeOwned, + { + SqliteDatasetWriter::new(self.db_file(), overwrite) + } + + /// Provides a reader instance for the SQLite dataset. + /// + /// # Arguments + /// + /// * `split` - A string slice that defines the data split for reading (e.g., "train", "test"). + /// + /// # Returns + /// + /// * A `Result` which is `Ok` if the reader could be created, `Err` otherwise. + pub fn reader(&self, split: &str) -> Result> + where + I: Clone + Send + Sync + Serialize + DeserializeOwned, + { + if !self.exists() { + panic!("The database file does not exist"); + } + + SqliteDataset::from_db_file(self.db_file(), split) + } +} + +/// This `SqliteDatasetWriter` struct is a SQLite database writer dedicated to storing datasets. +/// It retains the current writer's state and its database connection. +/// +/// Being thread-safe, this writer can be concurrently used across multiple threads. +/// +/// Typical applications include: +/// +/// - Generation of a new dataset +/// - Storage of preprocessed data or metadata +/// - Enlargement of a dataset's item count post preprocessing +#[derive(Debug)] +pub struct SqliteDatasetWriter { + db_file: PathBuf, + db_file_tmp: Option>, + splits: Arc>>, + overwrite: bool, + conn_pool: Option>, + is_completed: Arc>, + phantom: PhantomData, +} + +impl SqliteDatasetWriter +where + I: Clone + Send + Sync + Serialize + DeserializeOwned, +{ + /// Creates a new instance of `SqliteDatasetWriter`. + /// + /// # Arguments + /// + /// * `db_file` - A reference to the Path that represents the database file path. + /// * `overwrite` - A boolean indicating if the existing database file should be overwritten. + /// + /// # Returns + /// + /// * A `Result` which is `Ok` if the writer could be created, `Err` otherwise. + pub fn new>(db_file: P, overwrite: bool) -> Result { + let writer = Self { + db_file: db_file.as_ref().to_path_buf(), + db_file_tmp: None, + splits: Arc::new(RwLock::new(HashSet::new())), + overwrite, + conn_pool: None, + is_completed: Arc::new(RwLock::new(false)), + phantom: PhantomData, + }; + + writer.init() + } + + /// Initializes the dataset writer by creating the database file, tables, and connection pool. + /// + /// # Returns + /// + /// * A `Result` which is `Ok` if the writer could be initialized, `Err` otherwise. + fn init(mut self) -> Result { + // Remove the db file if it already exists + if self.db_file.exists() { + if self.overwrite { + fs::remove_file(&self.db_file)?; + } else { + return Err(SqliteDatasetError::FileExists(self.db_file)); + } + } + + // Create the database file directory if it does not exist + let db_file_dir = self + .db_file + .parent() + .ok_or("Unable to get parent directory")?; + + if !db_file_dir.exists() { + fs::create_dir_all(db_file_dir)?; + } + + // Create a temp database file name as {base_dir}/{name}.db.tmp + let mut db_file_tmp = self.db_file.clone(); + db_file_tmp.set_extension("db.tmp"); + if db_file_tmp.exists() { + fs::remove_file(&db_file_tmp)?; + } + + // Create the temp database file and wrap it with a gix_tempfile::Handle + // This will ensure that the temp file is deleted when the writer is dropped + // or when process exits with SIGINT or SIGTERM (tempfile crate does not do this) + gix_tempfile::signal::setup(Default::default()); + self.db_file_tmp = Some(gix_tempfile::writable_at( + &db_file_tmp, + ContainingDirectory::Exists, + AutoRemove::Tempfile, + )?); + + let conn_pool = create_conn_pool(db_file_tmp, true)?; + self.conn_pool = Some(conn_pool); + + Ok(self) + } + + /// Serializes and writes an item to the database. The item is written to the table for the + /// specified split. If the table does not exist, it is created. If the table exists, the item + /// is appended to the table. The serialization is done using the [MessagePack](https://msgpack.org/) + /// + /// # Arguments + /// + /// * `split` - A string slice that defines the data split for writing (e.g., "train", "test"). + /// * `item` - A reference to the item to be written to the database. + /// + /// # Returns + /// + /// * A `Result` containing the index of the inserted row if successful, an error otherwise. + pub fn write(&self, split: &str, item: &I) -> Result { + // Acquire the read lock (wont't block other reads) + let is_completed = self.is_completed.read().unwrap(); + + // If the writer is completed, return an error + if *is_completed { + return Err(SqliteDatasetError::Other( + "Cannot save to a completed dataset writer", + )); + } + + // create the table for the split if it does not exist + if !self.splits.read().unwrap().contains(split) { + self.create_table(split)?; + } + + // Get a connection from the pool + let conn_pool = self.conn_pool.as_ref().unwrap(); + let conn = conn_pool.get()?; + + // Serialize the item using MessagePack + let serialized_item = rmp_serde::to_vec(item)?; + + // Turn off the synchronous and journal mode for speed up + // We are sacrificing durability for speed but it's okay because + // we always recreate the dataset if it is not completed. + pragma_update_with_error_handling(&conn, "synchronous", "OFF")?; + pragma_update_with_error_handling(&conn, "journal_mode", "OFF")?; + + // Insert the serialized item into the database + let insert_statement = format!("insert into {split} (item) values (?)"); + conn.execute(insert_statement.as_str(), [serialized_item])?; + + // Get the primary key of the last inserted row and convert to index (row_id-1) + let index = (conn.last_insert_rowid() - 1) as usize; + + Ok(index) + } + + /// Marks the dataset as completed and persists the temporary database file. + pub fn set_completed(&mut self) -> Result<()> { + let mut is_completed = self.is_completed.write().unwrap(); + + // Force close the connection pool + // This is required on Windows platform where the connection pool prevents + // from persisting the db by renaming the temp file. + if let Some(pool) = self.conn_pool.take() { + std::mem::drop(pool); + } + + // Rename the database file from tmp to db + let _file_result = self + .db_file_tmp + .take() // take ownership of the temporary file and set to None + .unwrap() // unwrap the temporary file + .persist(&self.db_file)? + .ok_or("Unable to persist the database file")?; + + *is_completed = true; + Ok(()) + } + + /// Creates table for the data split. + /// + /// Note: call is idempotent and thread-safe. + /// + /// # Arguments + /// + /// * `split` - A string slice that defines the data split for the table (e.g., "train", "test"). + /// + /// # Returns + /// + /// * A `Result` which is `Ok` if the table could be created, `Err` otherwise. + /// + /// TODO (@antimora): add support creating a table with columns corresponding to the item fields + fn create_table(&self, split: &str) -> Result<()> { + // Check if the split already exists + if self.splits.read().unwrap().contains(split) { + return Ok(()); + } + + let conn_pool = self.conn_pool.as_ref().unwrap(); + let connection = conn_pool.get()?; + let create_table_statement = format!( + "create table if not exists {split} (row_id integer primary key autoincrement not \ + null, item blob not null)" + ); + + connection.execute(create_table_statement.as_str(), [])?; + + // Add the split to the splits + self.splits.write().unwrap().insert(split.to_string()); + + Ok(()) + } +} + +/// Runs a pragma update and ignores the `ExecuteReturnedResults` error. +/// +/// Sometimes ExecuteReturnedResults is returned when running a pragma update. This is not an error +/// and can be ignored. This function runs the pragma update and ignores the error if it is +/// `ExecuteReturnedResults`. +fn pragma_update_with_error_handling( + conn: &PooledConnection, + setting: &str, + value: &str, +) -> Result<()> { + let result = conn.pragma_update(None, setting, value); + if let Err(error) = result + && error != rusqlite::Error::ExecuteReturnedResults + { + return Err(SqliteDatasetError::Sql(error)); + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use rayon::prelude::*; + use rstest::{fixture, rstest}; + use serde::{Deserialize, Serialize}; + use tempfile::{NamedTempFile, TempDir, tempdir}; + + use super::*; + + type SqlDs = SqliteDataset; + + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] + pub struct Sample { + column_str: String, + column_bytes: Vec, + column_int: i64, + column_bool: bool, + column_float: f64, + } + + #[fixture] + fn train_dataset() -> SqlDs { + SqliteDataset::::from_db_file("tests/data/sqlite-dataset.db", "train").unwrap() + } + + #[rstest] + pub fn len(train_dataset: SqlDs) { + assert_eq!(train_dataset.len(), 2); + } + + #[rstest] + pub fn get_some(train_dataset: SqlDs) { + let item = train_dataset.get(0).unwrap(); + assert_eq!(item.column_str, "HI1"); + assert_eq!(item.column_bytes, vec![55, 231, 159]); + assert_eq!(item.column_int, 1); + assert!(item.column_bool); + assert_eq!(item.column_float, 1.0); + } + + #[rstest] + pub fn get_none(train_dataset: SqlDs) { + assert_eq!(train_dataset.get(10), None); + } + + #[rstest] + pub fn multi_thread(train_dataset: SqlDs) { + let indices: Vec = vec![0, 1, 1, 3, 4, 5, 6, 0, 8, 1]; + let results: Vec> = + indices.par_iter().map(|&i| train_dataset.get(i)).collect(); + + let mut match_count = 0; + for (_index, result) in indices.iter().zip(results.iter()) { + if let Some(_val) = result { + match_count += 1 + } + } + + assert_eq!(match_count, 5); + } + + #[test] + fn sqlite_dataset_storage() { + // Test with non-existing file + let storage = SqliteDatasetStorage::from_file("non-existing.db"); + assert!(!storage.exists()); + + // Test with non-existing name + let storage = SqliteDatasetStorage::from_name("non-existing.db"); + assert!(!storage.exists()); + + // Test with existing file + let storage = SqliteDatasetStorage::from_file("tests/data/sqlite-dataset.db"); + assert!(storage.exists()); + let result = storage.reader::("train"); + assert!(result.is_ok()); + let train = result.unwrap(); + assert_eq!(train.len(), 2); + + // Test get writer + let temp_file = NamedTempFile::new().unwrap(); + let storage = SqliteDatasetStorage::from_file(temp_file.path()); + assert!(storage.exists()); + let result = storage.writer::(true); + assert!(result.is_ok()); + } + + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] + pub struct Complex { + column_str: String, + column_bytes: Vec, + column_int: i64, + column_bool: bool, + column_float: f64, + column_complex: Vec>>, + } + + /// Create a temporary directory. + #[fixture] + fn tmp_dir() -> TempDir { + // Create a TempDir. This object will be automatically + // deleted when it goes out of scope. + tempdir().unwrap() + } + type Writer = SqliteDatasetWriter; + + /// Create a SqliteDatasetWriter with a temporary directory. + /// Make sure to return the temporary directory so that it is not deleted. + #[fixture] + fn writer_fixture(tmp_dir: TempDir) -> (Writer, TempDir) { + let temp_dir_str = tmp_dir.path(); + let storage = SqliteDatasetStorage::from_name("preprocessed").with_base_dir(temp_dir_str); + let overwrite = true; + let result = storage.writer::(overwrite); + assert!(result.is_ok()); + let writer = result.unwrap(); + (writer, tmp_dir) + } + + #[test] + fn test_new() { + // Test that the constructor works with overwrite = true + let test_path = NamedTempFile::new().unwrap(); + let _writer = SqliteDatasetWriter::::new(&test_path, true).unwrap(); + assert!(!test_path.path().exists()); + + // Test that the constructor works with overwrite = false + let test_path = NamedTempFile::new().unwrap(); + let result = SqliteDatasetWriter::::new(&test_path, false); + assert!(result.is_err()); + + // Test that the constructor works with no existing file + let temp = NamedTempFile::new().unwrap(); + let test_path = temp.path().to_path_buf(); + assert!(temp.close().is_ok()); + assert!(!test_path.exists()); + let _writer = SqliteDatasetWriter::::new(&test_path, true).unwrap(); + assert!(!test_path.exists()); + } + + #[rstest] + pub fn sqlite_writer_write(writer_fixture: (Writer, TempDir)) { + // Get the dataset_saver from the fixture and tmp_dir (will be deleted after scope) + let (writer, _tmp_dir) = writer_fixture; + + assert!(writer.overwrite); + assert!(!writer.db_file.exists()); + + let new_item = Complex { + column_str: "HI1".to_string(), + column_bytes: vec![1_u8, 2, 3], + column_int: 0, + column_bool: true, + column_float: 1.0, + column_complex: vec![vec![vec![[1, 23_u8, 3]]]], + }; + + let index = writer.write("train", &new_item).unwrap(); + assert_eq!(index, 0); + + let mut writer = writer; + + writer.set_completed().expect("Failed to set completed"); + + assert!(writer.db_file.exists()); + assert!(writer.db_file_tmp.is_none()); + + let result = writer.write("train", &new_item); + + // Should fail because the writer is completed + assert!(result.is_err()); + + let dataset = SqliteDataset::::from_db_file(writer.db_file, "train").unwrap(); + + let fetched_item = dataset.get(0).unwrap(); + assert_eq!(fetched_item, new_item); + assert_eq!(dataset.len(), 1); + } + + #[rstest] + pub fn sqlite_writer_write_multi_thread(writer_fixture: (Writer, TempDir)) { + // Get the dataset_saver from the fixture and tmp_dir (will be deleted after scope) + let (writer, _tmp_dir) = writer_fixture; + + let writer = Arc::new(writer); + let record_count = 20; + + let splits = ["train", "test"]; + + (0..record_count).into_par_iter().for_each(|index: i64| { + let thread_id: std::thread::ThreadId = std::thread::current().id(); + let sample = Complex { + column_str: format!("test_{thread_id:?}_{index}"), + column_bytes: vec![index as u8, 2, 3], + column_int: index, + column_bool: true, + column_float: 1.0, + column_complex: vec![vec![vec![[1, index as u8, 3]]]], + }; + + // half for train and half for test + let split = splits[index as usize % 2]; + + let _index = writer.write(split, &sample).unwrap(); + }); + + let mut writer = Arc::try_unwrap(writer).unwrap(); + + writer + .set_completed() + .expect("Should set completed successfully"); + + let train = + SqliteDataset::::from_db_file(writer.db_file.clone(), "train").unwrap(); + let test = SqliteDataset::::from_db_file(writer.db_file, "test").unwrap(); + + assert_eq!(train.len(), record_count as usize / 2); + assert_eq!(test.len(), record_count as usize / 2); + } +} diff --git a/crates/burn-dataset/src/lib.rs b/crates/burn-dataset/src/lib.rs new file mode 100644 index 0000000..822a19d --- /dev/null +++ b/crates/burn-dataset/src/lib.rs @@ -0,0 +1,52 @@ +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] + +//! # Burn Dataset +//! +//! Burn Dataset is a library for creating and loading datasets. + +#[macro_use] +extern crate derive_new; + +extern crate alloc; +extern crate dirs; + +/// Sources for datasets. +pub mod source; + +pub mod transform; + +/// Audio datasets. +#[cfg(feature = "audio")] +pub mod audio; + +/// Vision datasets. +#[cfg(feature = "vision")] +pub mod vision; + +/// Natural language processing datasets. +#[cfg(feature = "nlp")] +pub mod nlp; + +/// Network dataset utilities. +#[cfg(feature = "network")] +pub mod network { + pub use burn_std::network::*; +} + +mod dataset; +pub use dataset::*; +#[cfg(any(feature = "sqlite", feature = "sqlite-bundled"))] +pub use source::huggingface::downloader::*; + +#[cfg(test)] +mod test_data { + pub fn string_items() -> Vec { + vec![ + "1 Item".to_string(), + "2 Items".to_string(), + "3 Items".to_string(), + "4 Items".to_string(), + ] + } +} diff --git a/crates/burn-dataset/src/nlp/ag_news.rs b/crates/burn-dataset/src/nlp/ag_news.rs new file mode 100644 index 0000000..dc151f1 --- /dev/null +++ b/crates/burn-dataset/src/nlp/ag_news.rs @@ -0,0 +1,211 @@ +//! AG NEWS Dataset Module +//! +//! This module provides functionality for loading the AG NEWS text classification dataset. +//! AG NEWS is a collection of news articles categorized into different topics. +//! The dataset is split into training (120,000 articles) and test (7,600 articles) sets. +//! +//! ## Dataset Details +//! - **Classes**: 4 categories (World, Sports, Business, Sci/Tech) +//! - **AG NEWS mirror**: [fastai](https://github.com/fastai/fastai/blob/master/fastai/data/external.py#L83) +//! - **License**: [Apache License](https://github.com/fastai/fastai/blob/master/LICENSE) +//! +//! ## Usage Example +//! ```rust +//! use burn_dataset::nlp::AgNewsDataset; +//! +//! // Create an AG NEWS dataset accessor +//! let dataset = AgNewsDataset::new(); +//! +//! // Access training and test sets +//! let train_dataset = dataset.train(); +//! let test_dataset = dataset.test(); +//! ``` + +use std::{path::PathBuf, sync::Mutex}; + +use flate2::read::GzDecoder; +use serde::{Deserialize, Serialize}; +use tar::Archive; + +use crate::InMemDataset; +use crate::network::downloader; + +/// AG NEWS mirror from [fastai](https://github.com/fastai/fastai/blob/master/fastai/data/external.py#L83). +/// Licensed under the [Apache License](https://github.com/fastai/fastai/blob/master/LICENSE). +const AG_NEWS_URL: &str = "https://s3.amazonaws.com/fast-ai-nlp/ag_news_csv.tgz"; + +/// Represents an item in the AG NEWS dataset. +/// +/// Each item contains a label, title, and content of a news article. +#[derive(Deserialize, Serialize, Debug, Clone)] +pub struct AgNewsItem { + /// The category label of the news article. + pub label: String, + /// The title of the news article. + pub title: String, + /// The content/body of the news article. + pub content: String, +} + +/// AG NEWS dataset accessor. +/// +/// This struct provides convenient access to the AG NEWS text classification dataset. +/// It automatically downloads (if not already downloaded), extracts, and loads the datasets. +/// +/// The dataset is split into training (120,000 articles) and test (7,600 articles) sets. +pub struct AgNewsDataset { + agnews_dir: PathBuf, +} + +/// AG NEWS dataset download lock. +/// +/// This lock ensures that only one thread downloads the AG NEWS dataset at a time. +static DOWNLOAD_LOCK: Mutex<()> = Mutex::new(()); + +impl AgNewsDataset { + /// Creates a new AG NEWS dataset accessor. + /// + /// This will download and extract the dataset if it's not already present. + pub fn new() -> Self { + Self { + agnews_dir: Self::download(), + } + } + + /// Downloads and extracts the AG NEWS dataset. + /// + /// # Returns + /// Path to the directory containing the extracted dataset. + fn download() -> PathBuf { + // Acquire the lock. This will block if another thread already holds the lock. + let _lock = DOWNLOAD_LOCK.lock().unwrap(); + + // Dataset files are stored in the burn-dataset cache directory + let cache_dir = dirs::cache_dir() + .expect("Could not get cache directory") + .join("burn-dataset"); + + // AG NEWS dataset directory + let agnews_dir = cache_dir.join("ag_news_csv"); + + // AG NEWS dataset url + let url = AG_NEWS_URL; + + // AG NEWS dataset archive filename + let filename = "ag_news_csv.tgz"; + + // Check for already downloaded content + if !agnews_dir.exists() { + // Download gzip file + let bytes = downloader::download_file_as_bytes(url, filename); + + // Decode gzip file content and unpack archive + let gz_buffer = GzDecoder::new(&bytes[..]); + let mut archive = Archive::new(gz_buffer); + archive.unpack(cache_dir).unwrap(); + } + + agnews_dir + } + + /// Parses a CSV file into an in-memory dataset. + /// + /// # Arguments + /// * `file_path` - Path to the CSV file to parse. + /// + /// # Returns + /// An `InMemDataset` containing the parsed data. + fn parse_csv(file_path: &str) -> InMemDataset { + let mut rdr = csv::ReaderBuilder::new(); + let rdr = rdr.has_headers(false); + + InMemDataset::from_csv(file_path, &rdr).expect("Failed to parse CSV file") + } + + /// Gets the training dataset. + /// + /// # Returns + /// An `InMemDataset` instance containing 120,000 training articles. + pub fn train(&self) -> InMemDataset { + let file_path = self.agnews_dir.join("train.csv"); + Self::parse_csv(file_path.to_str().unwrap()) + } + + /// Gets the test dataset. + /// + /// # Returns + /// An `InMemDataset` instance containing 7,600 test articles. + pub fn test(&self) -> InMemDataset { + let file_path = self.agnews_dir.join("test.csv"); + Self::parse_csv(file_path.to_str().unwrap()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Dataset; + + // AG NEWS dataset train and test dataset lengths + const TRAIN_DATASET_LEN: usize = 120000; + const TEST_DATASET_LEN: usize = 7600; + + #[test] + fn test_agnews_download() { + let agnews_dir = AgNewsDataset::download(); + assert!(agnews_dir.exists()); + } + + #[test] + fn test_agnews_len() { + let agnews = AgNewsDataset::new(); + let train_dataset = agnews.train(); + let test_dataset = agnews.test(); + assert_eq!(train_dataset.len(), TRAIN_DATASET_LEN); + assert_eq!(test_dataset.len(), TEST_DATASET_LEN); + } + + #[test] + fn test_agnews_first_and_last_item() { + let agnews = AgNewsDataset::new(); + + // Test the first and the last item in training dataset + let train_dataset = agnews.train(); + let first_item = train_dataset.get(0).unwrap(); + let last_item = train_dataset.get(train_dataset.len() - 1).unwrap(); + assert!(compare_item(&first_item, &("3".to_string(), "Wall St. Bears Claw Back Into the Black (Reuters)".to_string(), "Reuters - Short-sellers, Wall Street's dwindling\\band of ultra-cynics, are seeing green again.".to_string()))); + assert!(compare_item( + &last_item, + &( + "2".to_string(), + "Nets get Carter from Raptors".to_string(), + "INDIANAPOLIS -- All-Star Vince Carter was traded by the Toronto Raptors to the New Jersey Nets for Alonzo Mourning, Eric Williams, Aaron Williams, and a pair of first-round draft picks yesterday.".to_string() + ) + )); + + // Test the first and the last item in test dataset + let test_dataset = agnews.test(); + let first_item = test_dataset.get(0).unwrap(); + let last_item = test_dataset.get(test_dataset.len() - 1).unwrap(); + assert!(compare_item( + &first_item, + &( + "3".to_string(), + "Fears for T N pension after talks".to_string(), + "Unions representing workers at Turner Newall say they are 'disappointed' after talks with stricken parent firm Federal Mogul.".to_string() + ) + )); + assert!(compare_item( + &last_item, + &( + "3".to_string(), + "EBay gets into rentals".to_string(), + "EBay plans to buy the apartment and home rental service Rent.com for \\$415 million, adding to its already exhaustive breadth of offerings.".to_string() + ) + )); + } + + fn compare_item(item: &AgNewsItem, target: &(String, String, String)) -> bool { + item.label == target.0 && item.title == target.1 && item.content == target.2 + } +} diff --git a/crates/burn-dataset/src/nlp/mod.rs b/crates/burn-dataset/src/nlp/mod.rs new file mode 100644 index 0000000..52574b3 --- /dev/null +++ b/crates/burn-dataset/src/nlp/mod.rs @@ -0,0 +1,7 @@ +#[cfg(feature = "builtin-sources")] +mod ag_news; +mod text_folder; + +#[cfg(feature = "builtin-sources")] +pub use ag_news::*; +pub use text_folder::*; diff --git a/crates/burn-dataset/src/nlp/text_folder.rs b/crates/burn-dataset/src/nlp/text_folder.rs new file mode 100644 index 0000000..9df8850 --- /dev/null +++ b/crates/burn-dataset/src/nlp/text_folder.rs @@ -0,0 +1,421 @@ +use crate::transform::{Mapper, MapperDataset}; +use crate::{Dataset, InMemDataset}; + +use encoding_rs::{GB18030, GBK, UTF_8, UTF_16BE, UTF_16LE}; +use globwalk::{self, DirEntry}; +use std::collections::{HashMap, HashSet}; +use std::fs; +use std::io::Read; +use std::path::{Path, PathBuf}; +use thiserror::Error; + +const SUPPORTED_FILES: [&str; 1] = ["txt"]; + +/// Text data type. +#[derive(Debug, Clone, PartialEq)] +pub struct TextData { + /// The text content. + pub text: String, + + /// Original text source. + pub text_path: String, +} + +/// Text dataset item. +#[derive(Debug, Clone, PartialEq)] +pub struct TextDatasetItem { + /// Text content. + pub text: TextData, + + /// Label for the text. + pub label: usize, +} + +/// Raw text dataset item. +#[derive(Debug, Clone)] +struct TextDatasetItemRaw { + /// Text path. + text_path: PathBuf, + + /// Text label. + label: String, +} + +impl TextDatasetItemRaw { + fn new>(text_path: P, label: String) -> TextDatasetItemRaw { + TextDatasetItemRaw { + text_path: text_path.as_ref().to_path_buf(), + label, + } + } +} + +struct PathToTextDatasetItem { + classes: HashMap, +} + +/// Parse the text content from file with auto-detection of encoding. +fn parse_text_content(text_path: &PathBuf) -> String { + // Read raw bytes from disk + let mut file = fs::File::open(text_path).unwrap(); + let mut bytes = Vec::new(); + file.read_to_end(&mut bytes).unwrap(); + + // Try to detect encoding and decode text + // First try UTF-8 with BOM + if bytes.starts_with(&[0xEF, 0xBB, 0xBF]) && bytes.len() >= 3 { + let (result, _, had_errors) = UTF_8.decode(&bytes[3..]); + if !had_errors { + return result.into_owned(); + } + } + + // Try UTF-8 without BOM + let (result, _, had_errors) = UTF_8.decode(&bytes); + if !had_errors { + return result.into_owned(); + } + + // Try UTF-16LE with BOM + if bytes.starts_with(&[0xFF, 0xFE]) && bytes.len() >= 2 { + let (result, had_errors) = UTF_16LE.decode_with_bom_removal(&bytes[2..]); + if !had_errors { + return result.into_owned(); + } + } + + // Try UTF-16BE with BOM + if bytes.starts_with(&[0xFE, 0xFF]) && bytes.len() >= 2 { + let (result, had_errors) = UTF_16BE.decode_with_bom_removal(&bytes[2..]); + if !had_errors { + return result.into_owned(); + } + } + + // Try GB18030 encoding + let (result, _, had_errors) = GB18030.decode(&bytes); + if !had_errors { + return result.into_owned(); + } + + // Try GBK encoding + let (result, _, had_errors) = GBK.decode(&bytes); + if !had_errors { + return result.into_owned(); + } + + // Default fallback - use from_utf8_lossy for any remaining cases + String::from_utf8_lossy(&bytes).to_string() +} + +impl Mapper for PathToTextDatasetItem { + /// Convert a raw text dataset item (path-like) to text content with a target label. + fn map(&self, item: &TextDatasetItemRaw) -> TextDatasetItem { + let label = *self.classes.get(&item.label).unwrap(); + + // Load text from disk + let text_content = parse_text_content(&item.text_path); + + let text_data = TextData { + text: text_content, + text_path: item.text_path.display().to_string(), + }; + + TextDatasetItem { + text: text_data, + label, + } + } +} + +/// Error type for [TextFolderDataset](TextFolderDataset). +#[derive(Error, Debug)] +pub enum TextLoaderError { + /// Unknown error. + #[error("unknown: `{0}`")] + Unknown(String), + + /// I/O operation error. + #[error("I/O error: `{0}`")] + IOError(String), + + /// Invalid file error. + #[error("Invalid file extension: `{0}`")] + InvalidFileExtensionError(String), + + /// Encoding error. + #[error("Encoding error: `{0}`")] + EncodingError(String), +} + +type TextDatasetMapper = + MapperDataset, PathToTextDatasetItem, TextDatasetItemRaw>; + +/// A generic dataset to load texts from disk. +pub struct TextFolderDataset { + dataset: TextDatasetMapper, +} + +impl Dataset for TextFolderDataset { + fn get(&self, index: usize) -> Option { + self.dataset.get(index) + } + + fn len(&self) -> usize { + self.dataset.len() + } +} + +impl TextFolderDataset { + /// Create a text classification dataset from the root folder. + /// + /// # Arguments + /// + /// * `root` - Dataset root folder. + /// + /// # Returns + /// A new dataset instance. + pub fn new_classification>(root: P) -> Result { + // New dataset containing any of the supported file types + TextFolderDataset::new_classification_with(root, &SUPPORTED_FILES) + } + + /// Create a text classification dataset from the root folder. + /// The included texts are filtered based on the provided extensions. + /// + /// # Arguments + /// + /// * `root` - Dataset root folder. + /// * `extensions` - List of allowed extensions. + /// + /// # Returns + /// A new dataset instance. + pub fn new_classification_with(root: P, extensions: &[S]) -> Result + where + P: AsRef, + S: AsRef, + { + // Glob all texts with extensions + let walker = globwalk::GlobWalkerBuilder::from_patterns( + root.as_ref(), + &[format!( + "*.{{{}}}", // "*.{ext1,ext2,ext3} + extensions + .iter() + .map(Self::check_extension) + .collect::, _>>()? + .join(",") + )], + ) + .follow_links(true) + .sort_by(|p1: &DirEntry, p2: &DirEntry| p1.path().cmp(p2.path())) // order by path + .build() + .map_err(|err| TextLoaderError::Unknown(format!("{err:?}")))? + .filter_map(Result::ok); + + // Get all dataset items + let mut items = Vec::new(); + let mut classes = HashSet::new(); + for text in walker { + let text_path = text.path(); + + // Label name is represented by the parent folder name + let label = text_path + .parent() + .ok_or_else(|| { + TextLoaderError::IOError("Could not resolve text parent folder".to_string()) + })? + .file_name() + .ok_or_else(|| { + TextLoaderError::IOError( + "Could not resolve text parent folder name".to_string(), + ) + })? + .to_string_lossy() + .into_owned(); + + classes.insert(label.clone()); + + items.push(TextDatasetItemRaw::new(text_path, label)) + } + + // Sort class names + let mut classes = classes.into_iter().collect::>(); + classes.sort(); + + Self::with_items(items, &classes) + } + + /// Create a text classification dataset with the specified items. + /// + /// # Arguments + /// + /// * `items` - List of dataset items, each item represented by a tuple `(text path, label)`. + /// * `classes` - Dataset class names. + /// + /// # Returns + /// A new dataset instance. + pub fn new_classification_with_items, S: AsRef>( + items: Vec<(P, String)>, + classes: &[S], + ) -> Result { + // Parse items and check valid text extension types + let items = items + .into_iter() + .map(|(path, label)| { + // Map text path and label + let path = path.as_ref(); + let label = label; + + Self::check_extension(&path.extension().unwrap().to_str().unwrap())?; + + Ok(TextDatasetItemRaw::new(path, label)) + }) + .collect::, _>>()?; + + Self::with_items(items, classes) + } + + /// Create a text dataset with the specified items. + /// + /// # Arguments + /// + /// * `items` - Raw dataset items. + /// * `classes` - Dataset class names. + /// + /// # Returns + /// A new dataset instance. + fn with_items>( + items: Vec, + classes: &[S], + ) -> Result { + // NOTE: right now we don't need to validate the supported text files since + // the method is private. We assume it's already validated. + let dataset = InMemDataset::new(items); + + // Class names to index map + let classes = classes.iter().map(|c| c.as_ref()).collect::>(); + let classes_map: HashMap<_, _> = classes + .into_iter() + .enumerate() + .map(|(idx, cls)| (cls.to_string(), idx)) + .collect(); + + let mapper = PathToTextDatasetItem { + classes: classes_map, + }; + let dataset = MapperDataset::new(dataset, mapper); + + Ok(Self { dataset }) + } + + /// Check if extension is supported. + fn check_extension>(extension: &S) -> Result { + let extension = extension.as_ref(); + if !SUPPORTED_FILES.contains(&extension) { + Err(TextLoaderError::InvalidFileExtensionError( + extension.to_string(), + )) + } else { + Ok(extension.to_string()) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::path::Path; + + const TEXT_ROOT: &str = "tests/data/text_folder"; + + #[test] + fn test_text_folder_dataset() { + let dataset = TextFolderDataset::new_classification(TEXT_ROOT).unwrap(); + + // Dataset should have 4 elements (2 positive + 2 negative) + assert_eq!(dataset.len(), 4); + assert_eq!(dataset.get(4), None); + + // Check that we have items from both classes + let mut found_positive = false; + let mut found_negative = false; + + for i in 0..dataset.len() { + let item = dataset.get(i).unwrap(); + if item.label == 0 { + found_negative = true; + // Check that the text content is loaded correctly + assert!(!item.text.text.is_empty()); + assert!(item.text.text_path.contains("negative")); + } else if item.label == 1 { + found_positive = true; + // Check that the text content is loaded correctly + assert!(!item.text.text.is_empty()); + assert!(item.text.text_path.contains("positive")); + } + } + + // Verify we found items from both classes + assert!(found_positive); + assert!(found_negative); + } + + #[test] + fn test_text_folder_dataset_with_invalid_extension() { + // Try to create a dataset with an unsupported extension + let result = TextFolderDataset::new_classification_with(TEXT_ROOT, &["invalid"]); + assert!(result.is_err()); + } + + #[test] + fn test_text_folder_dataset_with_items() { + // Create the dataset + let root = Path::new(TEXT_ROOT); + let items = vec![ + ( + root.join("positive").join("sample1.txt"), + "positive".to_string(), + ), + ( + root.join("negative").join("sample2.txt"), + "negative".to_string(), + ), + ]; + let classes = vec!["positive", "negative"]; + let dataset = TextFolderDataset::new_classification_with_items(items, &classes).unwrap(); + + // Dataset should have 2 elements + assert_eq!(dataset.len(), 2); + assert_eq!(dataset.get(2), None); + + // Get items + let item0 = dataset.get(0).unwrap(); + let item1 = dataset.get(1).unwrap(); + + // Check item0 + assert!(compare_item( + &item0, + &( + "This is a positive text sample for testing the text folder dataset functionality." + .to_string(), + 0 + ) + )); + + // Check item1 + assert_eq!(item1.label, 1); + assert!(item1.text.text_path.contains("negative")); + assert!(compare_item( + &item1, + &( + "另一个负面文本样本，用以确保数据集能够处理同一类别中的多个文件。".to_string(), + 1 + ) + )); + } + + fn compare_item(item: &TextDatasetItem, target: &(String, usize)) -> bool { + item.text.text == target.0 && item.label == target.1 + } +} diff --git a/crates/burn-dataset/src/source/huggingface/downloader.rs b/crates/burn-dataset/src/source/huggingface/downloader.rs new file mode 100644 index 0000000..a97bd46 --- /dev/null +++ b/crates/burn-dataset/src/source/huggingface/downloader.rs @@ -0,0 +1,367 @@ +use std::fs::{self, create_dir_all}; +use std::path::{Path, PathBuf}; +use std::process::Command; + +use crate::{SqliteDataset, SqliteDatasetError, SqliteDatasetStorage}; + +use sanitize_filename::sanitize; +use serde::de::DeserializeOwned; +use thiserror::Error; + +const PYTHON_SOURCE: &str = include_str!("importer.py"); +#[cfg(not(target_os = "windows"))] +const VENV_BIN_PYTHON: &str = "bin/python3"; +#[cfg(target_os = "windows")] +const VENV_BIN_PYTHON: &str = "Scripts\\python"; + +/// Error type for [HuggingfaceDatasetLoader](HuggingfaceDatasetLoader). +#[derive(Error, Debug)] +pub enum ImporterError { + /// Unknown error. + #[error("unknown: `{0}`")] + Unknown(String), + + /// Fail to download python dependencies. + #[error("fail to download python dependencies: `{0}`")] + FailToDownloadPythonDependencies(String), + + /// Fail to create sqlite dataset. + #[error("sqlite dataset: `{0}`")] + SqliteDataset(#[from] SqliteDatasetError), + + /// python3 is not installed. + #[error("python3 is not installed")] + PythonNotInstalled, + + /// venv environment is not initialized. + #[error("venv environment is not initialized")] + VenvNotInitialized, +} + +/// Load a dataset from [huggingface datasets](https://huggingface.co/datasets). +/// +/// The dataset with all splits is stored in a single sqlite database (see [SqliteDataset](SqliteDataset)). +/// +/// # Example +/// ```no_run +/// use burn_dataset::HuggingfaceDatasetLoader; +/// use burn_dataset::SqliteDataset; +/// use serde::{Deserialize, Serialize}; +/// +/// #[derive(Deserialize, Debug, Clone)] +/// struct MnistItemRaw { +/// pub image_bytes: Vec, +/// pub label: usize, +/// } +/// +/// let train_ds:SqliteDataset = HuggingfaceDatasetLoader::new("mnist") +/// .dataset("train") +/// .unwrap(); +/// ``` +/// +/// # Note +/// This loader relies on the [`datasets` library by HuggingFace](https://huggingface.co/docs/datasets/index) +/// to download datasets. This is a Python library, so you must have an existing Python installation. +pub struct HuggingfaceDatasetLoader { + name: String, + subset: Option, + base_dir: Option, + huggingface_token: Option, + huggingface_cache_dir: Option, + huggingface_data_dir: Option, + trust_remote_code: bool, + use_python_venv: bool, +} + +impl HuggingfaceDatasetLoader { + /// Create a huggingface dataset loader. + pub fn new(name: &str) -> Self { + Self { + name: name.to_string(), + subset: None, + base_dir: None, + huggingface_token: None, + huggingface_cache_dir: None, + huggingface_data_dir: None, + trust_remote_code: false, + use_python_venv: true, + } + } + + /// Create a huggingface dataset loader for a subset of the dataset. + /// + /// The subset name must be one of the subsets listed in the dataset page. + /// + /// If no subset names are listed, then do not use this method. + pub fn with_subset(mut self, subset: &str) -> Self { + self.subset = Some(subset.to_string()); + self + } + + /// Specify a base directory to store the dataset. + /// + /// If not specified, the dataset will be stored in the system cache directory under `burn-dataset`. + pub fn with_base_dir(mut self, base_dir: &str) -> Self { + self.base_dir = Some(base_dir.into()); + self + } + + /// Specify a huggingface token to download datasets behind authentication. + /// + /// You can get a token from [tokens settings](https://huggingface.co/settings/tokens) + pub fn with_huggingface_token(mut self, huggingface_token: &str) -> Self { + self.huggingface_token = Some(huggingface_token.to_string()); + self + } + + /// Specify a huggingface cache directory to store the downloaded datasets. + /// + /// If not specified, the dataset will be stored in the system cache directory under `huggingface/datasets`. + pub fn with_huggingface_cache_dir(mut self, huggingface_cache_dir: &str) -> Self { + self.huggingface_cache_dir = Some(huggingface_cache_dir.to_string()); + self + } + + /// Specify a relative path to a subset of a dataset. This is used in some datasets for the + /// manual steps of dataset download process. + /// + /// Unless you've encountered a ManualDownloadError + /// when loading your dataset you probably don't have to worry about this setting. + pub fn with_huggingface_data_dir(mut self, huggingface_data_dir: &str) -> Self { + self.huggingface_data_dir = Some(huggingface_data_dir.to_string()); + self + } + + /// Specify whether or not to trust remote code. + /// + /// If not specified, trust remote code is set to true. + pub fn with_trust_remote_code(mut self, trust_remote_code: bool) -> Self { + self.trust_remote_code = trust_remote_code; + self + } + + /// Specify whether or not to use the burn-dataset Python + /// virtualenv for running the importer script. If false, local + /// `python3`'s environment is used. + /// + /// If not specified, the virtualenv is used. + pub fn with_use_python_venv(mut self, use_python_venv: bool) -> Self { + self.use_python_venv = use_python_venv; + self + } + + /// Load the dataset. + pub fn dataset( + self, + split: &str, + ) -> Result, ImporterError> { + let db_file = self.db_file()?; + let dataset = SqliteDataset::from_db_file(db_file, split)?; + Ok(dataset) + } + + /// Get the path to the sqlite database file. + /// + /// If the database file does not exist, it will be downloaded and imported. + pub fn db_file(self) -> Result { + // determine (and create if needed) the base directory + let base_dir = SqliteDatasetStorage::base_dir(self.base_dir); + + if !base_dir.exists() { + create_dir_all(&base_dir).expect("Failed to create base directory"); + } + + //sanitize the name and subset + let name = sanitize(self.name.as_str()); + + // create the db file path + let db_file_name = if let Some(subset) = self.subset.clone() { + format!("{name}-{}.db", sanitize(subset.as_str())) + } else { + format!("{name}.db") + }; + + let db_file = base_dir.join(db_file_name); + + // import the dataset if needed + if !Path::new(&db_file).exists() { + import( + self.name, + self.subset, + db_file.clone(), + base_dir, + self.huggingface_token, + self.huggingface_cache_dir, + self.huggingface_data_dir, + self.trust_remote_code, + self.use_python_venv, + )?; + } + + Ok(db_file) + } +} + +/// Import a dataset from huggingface. The transformed dataset is stored as sqlite database. +#[allow(clippy::too_many_arguments)] +fn import( + name: String, + subset: Option, + base_file: PathBuf, + base_dir: PathBuf, + huggingface_token: Option, + huggingface_cache_dir: Option, + huggingface_data_dir: Option, + trust_remote_code: bool, + use_python_venv: bool, +) -> Result<(), ImporterError> { + let python_path = if use_python_venv { + install_python_deps(&base_dir)? + } else { + get_python_name()?.into() + }; + + let mut command = Command::new(python_path); + + command.arg(importer_script_path(&base_dir)); + + command.arg("--name"); + command.arg(name); + + command.arg("--file"); + command.arg(base_file); + + if let Some(subset) = subset { + command.arg("--subset"); + command.arg(subset); + } + + if let Some(huggingface_token) = huggingface_token { + command.arg("--token"); + command.arg(huggingface_token); + } + + if let Some(huggingface_cache_dir) = huggingface_cache_dir { + command.arg("--cache_dir"); + command.arg(huggingface_cache_dir); + } + if let Some(huggingface_data_dir) = huggingface_data_dir { + command.arg("--data_dir"); + command.arg(huggingface_data_dir); + } + if trust_remote_code { + command.arg("--trust_remote_code"); + command.arg("True"); + } + let mut handle = command.spawn().unwrap(); + + let exit_status = handle + .wait() + .map_err(|err| ImporterError::Unknown(format!("{err:?}")))?; + + if !exit_status.success() { + return Err(ImporterError::Unknown(format!("{exit_status}"))); + } + + Ok(()) +} + +/// check python --version output is `Python 3.x.x` +fn check_python_version_is_3(python: &str) -> bool { + let output = Command::new(python).arg("--version").output(); + match output { + Ok(output) => { + if output.status.success() { + let version_string = String::from_utf8_lossy(&output.stdout); + if let Some(index) = version_string.find(' ') { + let version = &version_string[index + 1..]; + version.starts_with("3.") + } else { + false + } + } else { + false + } + } + Err(_error) => false, + } +} + +/// get python3 name `python` `python3` or `py` +fn get_python_name() -> Result<&'static str, ImporterError> { + let python_name_list = ["python3", "python", "py"]; + for python_name in python_name_list.iter() { + if check_python_version_is_3(python_name) { + return Ok(python_name); + } + } + Err(ImporterError::PythonNotInstalled) +} + +fn importer_script_path(base_dir: &Path) -> PathBuf { + let path_file = base_dir.join("importer.py"); + + fs::write(&path_file, PYTHON_SOURCE).expect("Write python dataset downloader"); + path_file +} + +fn install_python_deps(base_dir: &Path) -> Result { + let venv_dir = base_dir.join("venv"); + let venv_python_path = venv_dir.join(VENV_BIN_PYTHON); + // If the venv environment is already initialized, skip the initialization. + if !check_python_version_is_3(venv_python_path.to_str().unwrap()) { + let python_name = get_python_name()?; + let mut command = Command::new(python_name); + command.args([ + "-m", + "venv", + venv_dir + .as_os_str() + .to_str() + .expect("Path utf8 conversion should not fail"), + ]); + + // Spawn the venv creation process and wait for it to complete. + let mut handle = command.spawn().unwrap(); + + handle.wait().map_err(|err| { + ImporterError::FailToDownloadPythonDependencies(format!(" error: {err}")) + })?; + // Check if the venv environment can be used successfully." + if !check_python_version_is_3(venv_python_path.to_str().unwrap()) { + return Err(ImporterError::VenvNotInitialized); + } + } + + let mut ensurepip_cmd = Command::new(&venv_python_path); + ensurepip_cmd.args(["-m", "ensurepip", "--upgrade"]); + let status = ensurepip_cmd.status().map_err(|err| { + ImporterError::FailToDownloadPythonDependencies(format!("failed to run ensurepip: {err}")) + })?; + if !status.success() { + return Err(ImporterError::FailToDownloadPythonDependencies( + "ensurepip failed to initialize pip".to_string(), + )); + } + + let mut command = Command::new(&venv_python_path); + command.args([ + "-m", + "pip", + "--quiet", + "install", + "pyarrow", + "sqlalchemy", + "Pillow", + "soundfile", + "datasets", + ]); + + // Spawn the pip install process and wait for it to complete. + let mut handle = command.spawn().unwrap(); + handle + .wait() + .map_err(|err| ImporterError::FailToDownloadPythonDependencies(format!(" error: {err}")))?; + + Ok(venv_python_path) +} diff --git a/crates/burn-dataset/src/source/huggingface/importer.py b/crates/burn-dataset/src/source/huggingface/importer.py new file mode 100644 index 0000000..c3ce923 --- /dev/null +++ b/crates/burn-dataset/src/source/huggingface/importer.py @@ -0,0 +1,207 @@ +import argparse + +import pyarrow as pa +from datasets import Audio, Image, load_dataset +from sqlalchemy import Column, Integer, Table, create_engine, event, inspect +from sqlalchemy.types import LargeBinary + + +def download_and_export( + name: str, + subset: str, + db_file: str, + token: str, + cache_dir: str, + data_dir: str | None, + trust_remote_code: bool, +): + """ + Download a dataset from using HuggingFace dataset and export it to a sqlite database. + """ + + # TODO For media columns (Image and Audio) sometimes when decode=False, + # bytes can be none {'bytes': None, 'path': 'healthy_train.265.jpg'} + # We should handle this case, but unfortunately we did not come across this case yet to test it. + + print("*" * 80) + print("Starting huggingface dataset download and export") + print(f"Dataset Name: {name}") + print(f"Subset Name: {subset}") + print(f"Sqlite database file: {db_file}") + print(f"Trust remote code: {trust_remote_code}") + if cache_dir is None: + print(f"Custom cache dir: {cache_dir}") + print("*" * 80) + + # Load the dataset + dataset_all = load_dataset( + name, + subset, + cache_dir=cache_dir, + data_dir=data_dir, + use_auth_token=token, + trust_remote_code=trust_remote_code, + ) + + print(f"Dataset: {dataset_all}") + + # Create the database connection descriptor (sqlite) + engine = create_engine(f"sqlite:///{db_file}") + + # Set some sqlite pragmas to speed up the database + event.listen(engine, "connect", set_sqlite_pragma) + + # Add an row_id column to each table as primary key (datasets does not have API for this) + event.listen(Table, "before_create", add_pk_column) + + # Export each split in the dataset + for key in dataset_all.keys(): + dataset = dataset_all[key] + + # Disable decoding for audio and image fields + dataset = disable_decoding(dataset) + + # Flatten the dataset + dataset = dataset.flatten() + + # Rename columns to remove dots from the names + dataset = rename_columns(dataset) + + print(f"Saving dataset: {name} - {key}") + print(f"Dataset features: {dataset.features}") + + # Save the dataset to a sqlite database + dataset.to_sql( + key, # table name + engine, + # don't save the index, use row_id instead (index is not unique) + index=False, + dtype=blob_columns(dataset), # save binary columns as blob + ) + + # Print the schema of the database so we can reference the columns in the rust code + print_table_info(engine) + + +def disable_decoding(dataset): + """ + Disable decoding for audio and image fields. The fields will be saved as raw file bytes. + """ + for k, v in dataset.features.items(): + if isinstance(v, Audio): + dataset = dataset.cast_column(k, Audio(decode=False)) + elif isinstance(v, Image): + dataset = dataset.cast_column(k, Image(decode=False)) + + return dataset + + +def rename_columns(dataset): + """ + Rename columns to remove dots from the names. Dots appear in the column names because of the flattening. + Dots are not allowed in column names in rust and sql (unless quoted). So we replace them with underscores. + This way there is an easy name mapping between the rust and sql columns. + """ + + for name in dataset.features.keys(): + if "." in name: + dataset = dataset.rename_column(name, name.replace(".", "_")) + + return dataset + + +def blob_columns(dataset): + """ + Make sure all binary columns are blob columns in the database because + `to_sql` exports binary values as TEXT instead of BLOB. + """ + type_mapping = {} + for name, value in dataset.features.items(): + if value.pa_type is not None and pa.types.is_binary(value.pa_type): + type_mapping[name] = LargeBinary + return type_mapping + + +def set_sqlite_pragma(dbapi_connection, connection_record): + """ + Set some sqlite pragmas to speed up the database + """ + cursor = dbapi_connection.cursor() + cursor.execute("PRAGMA synchronous = OFF") + cursor.execute("PRAGMA journal_mode = OFF") + cursor.close() + + +def add_pk_column(target, connection, **kw): + """ + Add an id column to each table. + """ + target.append_column(Column("row_id", Integer, primary_key=True)) + + +def print_table_info(engine): + """ + Print the schema of the database so we can reference the columns in the rust code + """ + print(f"Printing table schema for sqlite3 db ({engine})") + inspector = inspect(engine) + for table_name in inspector.get_table_names(): + print(f"Table: {table_name}") + for column in inspector.get_columns(table_name): + print(f"Column: {column['name']} - {column['type']}") + print("") + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Huggingface datasets downloader to use with burn-dataset" + ) + parser.add_argument( + "--name", type=str, help="Name of the dataset to download", required=True + ) + parser.add_argument( + "--file", type=str, help="Base file name where the data is saved", required=True + ) + parser.add_argument( + "--subset", type=str, help="Subset name", required=False, default=None + ) + parser.add_argument( + "--token", + type=str, + help="HuggingFace authentication token", + required=False, + default=None, + ) + parser.add_argument( + "--cache_dir", type=str, help="Cache directory", required=False, default=None + ) + parser.add_argument( + "--data_dir", type=str, help="Relative path to a specific subset of your dataset", required=False, default=None + ) + parser.add_argument( + "--trust_remote_code", + type=bool, + help="Trust remote code", + required=False, + default=None, + ) + + return parser.parse_args() + + +def run(): + args = parse_args() + + download_and_export( + args.name, + args.subset, + args.file, + args.token, + args.data_dir, + args.cache_dir, + args.trust_remote_code, + ) + + +if __name__ == "__main__": + run() diff --git a/crates/burn-dataset/src/source/huggingface/mod.rs b/crates/burn-dataset/src/source/huggingface/mod.rs new file mode 100644 index 0000000..a6ccec9 --- /dev/null +++ b/crates/burn-dataset/src/source/huggingface/mod.rs @@ -0,0 +1,3 @@ +pub(crate) mod downloader; + +pub use downloader::*; diff --git a/crates/burn-dataset/src/source/mod.rs b/crates/burn-dataset/src/source/mod.rs new file mode 100644 index 0000000..b9ed762 --- /dev/null +++ b/crates/burn-dataset/src/source/mod.rs @@ -0,0 +1,3 @@ +/// Huggingface source +#[cfg(any(feature = "sqlite", feature = "sqlite-bundled"))] +pub mod huggingface; diff --git a/crates/burn-dataset/src/transform/composed.rs b/crates/burn-dataset/src/transform/composed.rs new file mode 100644 index 0000000..dfe0a9e --- /dev/null +++ b/crates/burn-dataset/src/transform/composed.rs @@ -0,0 +1,56 @@ +use crate::Dataset; + +/// Compose multiple datasets together to create a bigger one. +#[derive(new)] +pub struct ComposedDataset { + datasets: Vec, +} + +impl Dataset for ComposedDataset +where + D: Dataset, + I: Clone, +{ + fn get(&self, index: usize) -> Option { + let mut current_index = 0; + for dataset in self.datasets.iter() { + if index < dataset.len() + current_index { + return dataset.get(index - current_index); + } + current_index += dataset.len(); + } + None + } + fn len(&self) -> usize { + let mut total = 0; + for dataset in self.datasets.iter() { + total += dataset.len(); + } + total + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::FakeDataset; + + #[test] + fn test_composed_dataset() { + let dataset1 = FakeDataset::::new(10); + let dataset2 = FakeDataset::::new(5); + + let items1 = dataset1.iter().collect::>(); + let items2 = dataset2.iter().collect::>(); + + let composed = ComposedDataset::new(vec![dataset1, dataset2]); + + assert_eq!(composed.len(), 15); + + let expected_items: Vec = items1.iter().chain(items2.iter()).cloned().collect(); + + let items = composed.iter().collect::>(); + + assert_eq!(items, expected_items); + } +} diff --git a/crates/burn-dataset/src/transform/mapper.rs b/crates/burn-dataset/src/transform/mapper.rs new file mode 100644 index 0000000..f73eec8 --- /dev/null +++ b/crates/burn-dataset/src/transform/mapper.rs @@ -0,0 +1,60 @@ +use crate::Dataset; +use std::marker::PhantomData; + +/// Basic mapper trait to be used with the [mapper dataset](MapperDataset). +pub trait Mapper: Send + Sync { + /// Maps an item of type I to an item of type O. + fn map(&self, item: &I) -> O; +} + +/// Dataset mapping each element in an inner dataset to another element type lazily. +#[derive(new)] +pub struct MapperDataset { + dataset: D, + mapper: M, + input: PhantomData, +} + +impl Dataset for MapperDataset +where + D: Dataset, + M: Mapper + Send + Sync, + I: Send + Sync, + O: Send + Sync, +{ + fn get(&self, index: usize) -> Option { + let item = self.dataset.get(index); + item.map(|item| self.mapper.map(&item)) + } + + fn len(&self) -> usize { + self.dataset.len() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{InMemDataset, test_data}; + + #[test] + pub fn given_mapper_dataset_when_iterate_should_iterate_though_all_map_items() { + struct StringToFirstChar; + + impl Mapper for StringToFirstChar { + fn map(&self, item: &String) -> String { + let mut item = item.clone(); + item.truncate(1); + item + } + } + + let items_original = test_data::string_items(); + let dataset = InMemDataset::new(items_original); + let dataset = MapperDataset::new(dataset, StringToFirstChar); + + let items: Vec = dataset.iter().collect(); + + assert_eq!(vec!["1", "2", "3", "4"], items); + } +} diff --git a/crates/burn-dataset/src/transform/mod.rs b/crates/burn-dataset/src/transform/mod.rs new file mode 100644 index 0000000..a35f57d --- /dev/null +++ b/crates/burn-dataset/src/transform/mod.rs @@ -0,0 +1,30 @@ +//! # Dataset Transformations +//! +//! This module provides a collection of [`crate::Dataset`] composition wrappers; +//! providing composition, subset selection, sampling, random shuffling, and windowing. +//! +//! * [`ComposedDataset`] - composes a list of datasets. +//! * [`PartialDataset`] - selects a contiguous index range subset of a dataset. +//! * [`ShuffledDataset`] - a randomly shuffled / mutably shuffle-able dataset; +//! a thin wrapper around [`SelectionDataset`]. +//! * [`SamplerDataset`] - samples a dataset; support for with/without replacement, +//! and under/oversampling. +//! * [`SelectionDataset`] - selects a subset of a dataset via indices; support for shuffling. +//! * [`WindowsDataset`] - creates a sliding window over a dataset. +mod composed; +mod mapper; +mod options; +mod partial; +mod sampler; +mod selection; +mod shuffle; +mod window; + +pub use composed::*; +pub use mapper::*; +pub use options::*; +pub use partial::*; +pub use sampler::*; +pub use selection::*; +pub use shuffle::*; +pub use window::*; diff --git a/crates/burn-dataset/src/transform/options.rs b/crates/burn-dataset/src/transform/options.rs new file mode 100644 index 0000000..eee5dd3 --- /dev/null +++ b/crates/burn-dataset/src/transform/options.rs @@ -0,0 +1,199 @@ +use rand::SeedableRng; +use rand::prelude::StdRng; +use rand::rngs::SysRng; + +/// Defines a source for a `StdRng`. +/// +/// # Examples +/// +/// ```rust,no_run +/// use rand::rngs::StdRng; +/// use rand::SeedableRng; +/// use burn_dataset::transform::RngSource; +/// +/// // Default via `StdRng::from_os_rng()` (`RngSource::Default`) +/// let system: RngSource = RngSource::default(); +/// +/// // From a fixed seed (`RngSource::Seed`) +/// let seeded: RngSource = 42.into(); +/// +/// // From an existing rng (`RngSource::Rng`) +/// let rng = StdRng::seed_from_u64(123); +/// let with_rng: RngSource = rng.into(); +/// +/// // Forks the parent RNG to derive an independent, deterministic child RNG. +/// // The original `rng` is modified, and the resulting `RngSource` contains +/// // a new RNG starting from a unique state. +/// let mut rng = StdRng::seed_from_u64(123); +/// let forked: RngSource = (&mut rng).into(); +/// ``` +#[derive(Debug, Default, PartialEq, Eq)] +#[allow(clippy::large_enum_variant)] +pub enum RngSource { + /// Build a new rng from the system. + #[default] + Default, + + /// The rng is passed as a seed. + Seed(u64), + + /// The rng is passed as an option. + Rng(StdRng), +} + +impl From for StdRng { + fn from(source: RngSource) -> Self { + match source { + RngSource::Default => StdRng::try_from_rng(&mut SysRng).unwrap(), + RngSource::Rng(rng) => rng, + RngSource::Seed(seed) => StdRng::seed_from_u64(seed), + } + } +} + +impl From for RngSource { + fn from(seed: u64) -> Self { + Self::Seed(seed) + } +} + +impl From for RngSource { + fn from(rng: StdRng) -> Self { + Self::Rng(rng) + } +} + +/// Derive an independent RNG from a mutable parent RNG. +/// +/// This advances the parent RNG and creates a new RNG seeded from its output. +/// The derived RNG is *not* a clone of the parent's state, but an independent +/// stream (equivalent to `SeedableRng::fork`). +impl From<&mut StdRng> for RngSource { + fn from(rng: &mut StdRng) -> Self { + Self::Rng(rng.fork()) + } +} + +/// Helper option to describe the size of a wrapper, relative to a wrapped object. +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub enum SizeConfig { + /// Use the size of the source dataset. + #[default] + Default, + + /// Use the size as a ratio of the source dataset size. + /// + /// Must be >= 0. + Ratio(f64), + + /// Use a fixed size. + Fixed(usize), +} + +impl SizeConfig { + /// Construct a source which will have the same size as the source dataset. + pub fn source() -> Self { + Self::Default + } + + /// Resolve the effective size. + /// + /// ## Arguments + /// + /// - `source_size`: the size of the source dataset. + /// + /// ## Returns + /// + /// The resolved size of the wrapper dataset. + pub fn resolve(self, source_size: usize) -> usize { + match self { + SizeConfig::Default => source_size, + SizeConfig::Ratio(ratio) => { + assert!(ratio >= 0.0, "Ratio must be non-negative: {ratio}"); + ((source_size as f64) * ratio) as usize + } + SizeConfig::Fixed(size) => size, + } + } +} + +impl From for SizeConfig { + fn from(size: usize) -> Self { + Self::Fixed(size) + } +} + +impl From for SizeConfig { + fn from(ratio: f64) -> Self { + Self::Ratio(ratio) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rand::SeedableRng; + + #[test] + fn test_rng_source_default() { + let rng_source: RngSource = Default::default(); + assert_eq!(&rng_source, &RngSource::Default); + assert_eq!(&rng_source, &RngSource::default()); + + // Exercise the from_os_rng() call; but we don't know its seed; + let _rng: StdRng = rng_source.into(); + } + + #[test] + fn test_rng_source_seed() { + let rng_source = RngSource::from(42); + assert_eq!(&rng_source, &RngSource::Seed(42)); + + let rng: StdRng = rng_source.into(); + let expected = StdRng::seed_from_u64(42); + + assert_eq!(rng, expected); + } + + #[test] + fn test_rng_source_rng() { + // From StdRng (owned). + { + let original = StdRng::seed_from_u64(42); + + let rng_source = RngSource::from(original); + let rng: StdRng = rng_source.into(); + // No longer clone, but from <> into should not have advanced the state + let original = StdRng::seed_from_u64(42); + assert_eq!(rng, original); + } + + // From &mut StdRng (forks parent) + { + let mut original = StdRng::seed_from_u64(42); + let mut rng = StdRng::seed_from_u64(42); + let rng_forked = rng.fork(); + + let rng_source = RngSource::from(&mut original); + + // Ensure the original was advanced + assert_eq!(original, rng); + + // Ensure the sourced RNG matches the fork + let rng: StdRng = rng_source.into(); + assert_eq!(rng, rng_forked); + } + } + + #[test] + fn test_size_config() { + assert_eq!(SizeConfig::default(), SizeConfig::Default); + + assert_eq!(SizeConfig::from(42), SizeConfig::Fixed(42)); + + assert_eq!(SizeConfig::from(1.5), SizeConfig::Ratio(1.5)); + + assert_eq!(SizeConfig::source(), SizeConfig::Default); + assert_eq!(SizeConfig::source().resolve(50), 50); + } +} diff --git a/crates/burn-dataset/src/transform/partial.rs b/crates/burn-dataset/src/transform/partial.rs new file mode 100644 index 0000000..520aa8d --- /dev/null +++ b/crates/burn-dataset/src/transform/partial.rs @@ -0,0 +1,206 @@ +use crate::Dataset; +use std::{marker::PhantomData, sync::Arc}; + +/// Only use a fraction of an existing dataset lazily. +#[derive(new, Clone)] +pub struct PartialDataset { + dataset: D, + start_index: usize, + end_index: usize, + input: PhantomData, +} + +impl PartialDataset +where + D: Dataset, +{ + /// Splits a dataset into multiple partial datasets. + pub fn split(dataset: D, num: usize) -> Vec, I>> { + let dataset = Arc::new(dataset); // cheap cloning. + + let mut current = 0; + let mut datasets = Vec::with_capacity(num); + + let batch_size = dataset.len() / num; + + for i in 0..num { + let start = current; + let mut end = current + batch_size; + + if i == (num - 1) { + end = dataset.len(); + } + + let dataset = PartialDataset::new(dataset.clone(), start, end); + + current += batch_size; + datasets.push(dataset); + } + + datasets + } + + /// Splits a dataset by distributing complete chunks/batches across multiple partial datasets. + pub fn split_chunks( + dataset: D, + num: usize, + batch_size: usize, + ) -> Vec, I>> { + let dataset = Arc::new(dataset); // cheap cloning. + let total_items = dataset.len(); + + // Total number of complete batches + let total_batches = total_items.div_ceil(batch_size); + let batches_per_split = total_batches / num; + let extra_batches = total_batches % num; + + let mut datasets = Vec::with_capacity(num); + let mut current_batch = 0; + + for i in 0..num { + // Extra batches distributed across first splits + let split_batches = if i < extra_batches { + batches_per_split + 1 + } else { + batches_per_split + }; + + let start_batch = current_batch; + let end_batch = start_batch + split_batches; + + let start_index = start_batch * batch_size; + let end_index = core::cmp::min(end_batch * batch_size, total_items); + + if start_index < total_items { + datasets.push(PartialDataset::new(dataset.clone(), start_index, end_index)); + } + + current_batch = end_batch; + } + + datasets + } +} + +impl Dataset for PartialDataset +where + D: Dataset, + I: Clone + Send + Sync, +{ + fn get(&self, index: usize) -> Option { + let index = index + self.start_index; + if index < self.start_index || index >= self.end_index { + return None; + } + self.dataset.get(index) + } + + fn len(&self) -> usize { + usize::min(self.end_index - self.start_index, self.dataset.len()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::FakeDataset; + use std::collections::HashSet; + + #[test] + fn test_start_from_beginning() { + let dataset_original = FakeDataset::::new(27); + let mut items_original_1 = HashSet::new(); + let mut items_original_2 = HashSet::new(); + let mut items_partial = HashSet::new(); + dataset_original.iter().enumerate().for_each(|(i, item)| { + match i >= 10 { + true => items_original_2.insert(item), + false => items_original_1.insert(item), + }; + }); + + let dataset_partial = PartialDataset::new(dataset_original, 0, 10); + + for item in dataset_partial.iter() { + items_partial.insert(item); + } + + assert_eq!(dataset_partial.len(), 10); + assert_eq!(items_original_1, items_partial); + for item in items_original_2 { + assert!(!items_partial.contains(&item)); + } + } + + #[test] + fn test_start_inside() { + let dataset_original = FakeDataset::::new(27); + let mut items_original_1 = HashSet::new(); + let mut items_original_2 = HashSet::new(); + let mut items_partial = HashSet::new(); + + dataset_original.iter().enumerate().for_each(|(i, item)| { + match !(10..20).contains(&i) { + true => items_original_2.insert(item), + false => items_original_1.insert(item), + }; + }); + + let dataset_partial = PartialDataset::new(dataset_original, 10, 20); + for item in dataset_partial.iter() { + items_partial.insert(item); + } + + assert_eq!(dataset_partial.len(), 10); + assert_eq!(items_original_1, items_partial); + for item in items_original_2 { + assert!(!items_partial.contains(&item)); + } + } + + #[test] + fn test_split_contains_all_items_without_duplicates() { + let dataset_original = FakeDataset::::new(27); + let mut items_original = Vec::new(); + let mut items_partial = Vec::new(); + for item in dataset_original.iter() { + items_original.push(item); + } + + let dataset_partials = PartialDataset::split(dataset_original, 4); + let expected_len = [6, 6, 6, 9]; + + for (i, dataset) in dataset_partials.iter().enumerate() { + assert_eq!(dataset.len(), expected_len[i]); + for item in dataset.iter() { + items_partial.push(item); + } + } + + assert_eq!(items_original, items_partial); + } + + #[test] + fn test_split_chunks_contains_all_items_without_duplicates() { + let dataset_original = FakeDataset::::new(27); + let mut items_original = Vec::new(); + let mut items_partial = Vec::new(); + for item in dataset_original.iter() { + items_original.push(item); + } + + let dataset_partials = PartialDataset::split_chunks(dataset_original, 4, 5); + // [(2 * 5), (2 * 5), 5, 2] -> 5 complete chunks + 1 incomplete with 2 remaining items + // OTOH, `split(dataset, 4)` would yield [6, 6, 6, 9] -> 4 incomplete chunks + 4 incomplete with [1, 1, 1, 4] + let expected_len = [10, 10, 5, 2]; + + for (i, dataset) in dataset_partials.iter().enumerate() { + assert_eq!(dataset.len(), expected_len[i]); + for item in dataset.iter() { + items_partial.push(item); + } + } + + assert_eq!(items_original, items_partial); + } +} diff --git a/crates/burn-dataset/src/transform/sampler.rs b/crates/burn-dataset/src/transform/sampler.rs new file mode 100644 index 0000000..a4f0268 --- /dev/null +++ b/crates/burn-dataset/src/transform/sampler.rs @@ -0,0 +1,438 @@ +use crate::Dataset; +use crate::transform::{RngSource, SizeConfig}; +use rand::prelude::SliceRandom; +use rand::{RngExt, distr::Uniform, rngs::StdRng, seq::IteratorRandom}; +use std::{marker::PhantomData, ops::DerefMut, sync::Mutex}; + +/// Options to configure a [SamplerDataset]. +#[derive(Debug, PartialEq)] +pub struct SamplerDatasetOptions { + /// The sampling mode. + pub replace_samples: bool, + + /// The size source of the wrapper relative to the dataset. + pub size_config: SizeConfig, + + /// The source of the random number generator. + pub rng_source: RngSource, +} + +impl Default for SamplerDatasetOptions { + fn default() -> Self { + Self { + replace_samples: true, + size_config: SizeConfig::Default, + rng_source: RngSource::Default, + } + } +} + +impl From> for SamplerDatasetOptions +where + T: Into, +{ + fn from(option: Option) -> Self { + match option { + Some(option) => option.into(), + None => Self::default(), + } + } +} + +impl From for SamplerDatasetOptions { + fn from(size: usize) -> Self { + Self::default().with_replacement().with_fixed_size(size) + } +} + +impl SamplerDatasetOptions { + /// Set the replacement mode. + pub fn with_replace_samples(self, replace_samples: bool) -> Self { + Self { + replace_samples, + ..self + } + } + + /// Set the replacement mode to WithReplacement. + pub fn with_replacement(self) -> Self { + self.with_replace_samples(true) + } + + /// Set the replacement mode to WithoutReplacement. + pub fn without_replacement(self) -> Self { + self.with_replace_samples(false) + } + + /// Set the size source. + pub fn with_size(self, source: S) -> Self + where + S: Into, + { + Self { + size_config: source.into(), + ..self + } + } + + /// Set the size to the size of the source. + pub fn with_source_size(self) -> Self { + self.with_size(SizeConfig::Default) + } + + /// Set the size to a fixed size. + pub fn with_fixed_size(self, size: usize) -> Self { + self.with_size(size) + } + + /// Set the size to be a multiple of the ration and the source size. + pub fn with_size_ratio(self, size_ratio: f64) -> Self { + self.with_size(size_ratio) + } + + /// Set the `RngSource`. + pub fn with_rng(self, rng: R) -> Self + where + R: Into, + { + Self { + rng_source: rng.into(), + ..self + } + } + + /// Use the system rng. + pub fn with_system_rng(self) -> Self { + self.with_rng(RngSource::Default) + } + + /// Use a rng, built from a seed. + pub fn with_seed(self, seed: u64) -> Self { + self.with_rng(seed) + } +} + +/// Sample items from a dataset. +/// +/// This is a convenient way of modeling a dataset as a probability distribution of a fixed size. +/// You have multiple options to instantiate the dataset sampler. +/// +/// * With replacement (Default): This is the most efficient way of using the sampler because no state is +/// required to keep indices that have been selected. +/// +/// * Without replacement: This has a similar effect to using a +/// [shuffled dataset](crate::transform::ShuffledDataset), but with more flexibility since you can +/// set the dataset to an arbitrary size. Once every item has been used, a new cycle is +/// created with a new random suffle. +pub struct SamplerDataset { + dataset: D, + size: usize, + state: Mutex, + input: PhantomData, +} +enum SamplerState { + WithReplacement(StdRng), + WithoutReplacement(StdRng, Vec), +} + +impl SamplerDataset +where + D: Dataset, + I: Send + Sync, +{ + /// Creates a new sampler dataset with replacement. + /// + /// When the sample size is less than or equal to the source dataset size, + /// data will be sampled without replacement from the source dataset in + /// a uniformly shuffled order. + /// + /// When the sample size is greater than the source dataset size, + /// the entire source dataset will be sampled once for every multiple + /// of the size ratios; with the remaining samples taken without replacement + /// uniformly from the source. All samples will be returned uniformly shuffled. + /// + /// ## Arguments + /// + /// * `dataset`: the dataset to wrap. + /// * `options`: the options to configure the sampler dataset. + /// + /// ## Examples + /// ```rust,ignore + /// use burn_dataset::transform::{ + /// SamplerDataset, + /// SamplerDatasetOptions, + /// }; + /// + /// // Examples below assuming `dataset.len()` = `10`. + /// + /// // sample size: 5 + /// // WithReplacement + /// // rng: StdRng::from_os_rng() + /// SamplerDataset::new(dataset, 5); + /// + /// // sample size: 10 (source) + /// // WithReplacement + /// // rng: StdRng::from_os_rng() + /// SamplerDataset::new(dataset, SamplerDatasetOptions::default()); + /// + /// // sample size: 15 + /// // WithoutReplacement + /// // rng: StdRng::seed_from_u64(42) + /// SamplerDataset::new( + /// dataset, + /// SamplerDatasetOptions::default() + /// .with_size(1.5) + /// .without_replacement() + /// .with_rng(42), + /// ); + /// ``` + pub fn new(dataset: D, options: O) -> Self + where + O: Into, + { + let options = options.into(); + let size = options.size_config.resolve(dataset.len()); + let rng = options.rng_source.into(); + Self { + dataset, + size, + state: Mutex::new(match options.replace_samples { + true => SamplerState::WithReplacement(rng), + false => SamplerState::WithoutReplacement(rng, Vec::with_capacity(size)), + }), + input: PhantomData, + } + } + + /// Creates a new sampler dataset with replacement. + /// + /// # Arguments + /// + /// - `dataset`: the dataset to wrap. + /// - `size`: the effective size of the sampled dataset. + pub fn with_replacement(dataset: D, size: usize) -> Self { + Self::new( + dataset, + SamplerDatasetOptions::default() + .with_replacement() + .with_fixed_size(size), + ) + } + + /// Creates a new sampler dataset without replacement. + /// + /// When the sample size is less than or equal to the source dataset size, + /// data will be sampled without replacement from the source dataset in + /// a uniformly shuffled order. + /// + /// When the sample size is greater than the source dataset size, + /// the entire source dataset will be sampled once for every multiple + /// of the size ratios; with the remaining samples taken without replacement + /// uniformly from the source. All samples will be returned uniformly shuffled. + /// + /// # Arguments + /// - `dataset`: the dataset to wrap. + /// - `size`: the effective size of the sampled dataset. + pub fn without_replacement(dataset: D, size: usize) -> Self { + Self::new( + dataset, + SamplerDatasetOptions::default() + .without_replacement() + .with_fixed_size(size), + ) + } + + /// Determines if the sampler is using the "with replacement" strategy. + /// + /// # Returns + /// - `true`: If the sampler is configured to sample with replacement. + /// - `false`: If the sampler is configured to sample without replacement. + pub fn is_with_replacement(&self) -> bool { + match self.state.lock().unwrap().deref_mut() { + SamplerState::WithReplacement(_) => true, + SamplerState::WithoutReplacement(_, _) => false, + } + } + + fn index(&self) -> usize { + match self.state.lock().unwrap().deref_mut() { + SamplerState::WithReplacement(rng) => { + rng.sample(Uniform::new(0, self.dataset.len()).unwrap()) + } + SamplerState::WithoutReplacement(rng, indices) => { + if indices.is_empty() { + // Refill the state. + let idx_range = 0..self.dataset.len(); + for _ in 0..(self.size / self.dataset.len()) { + // No need to `.choose_multiple` here because we're using + // the entire source range; and `.choose_multiple` will + // not return a random sample anyway. + indices.extend(idx_range.clone()) + } + + // From `choose_multiple` documentation: + // > Although the elements are selected randomly, the order of elements in + // > the buffer is neither stable nor fully random. If random ordering is + // > desired, shuffle the result. + indices.extend(idx_range.sample(rng, self.size - indices.len())); + + // The real shuffling is done here. + indices.shuffle(rng); + } + + indices.pop().expect("Indices are refilled when empty.") + } + } + } +} + +impl Dataset for SamplerDataset +where + D: Dataset, + I: Send + Sync, +{ + fn get(&self, index: usize) -> Option { + if index >= self.size { + return None; + } + + self.dataset.get(self.index()) + } + + fn len(&self) -> usize { + self.size + } +} + +#[cfg(test)] +mod tests { + #![allow(clippy::bool_assert_comparison)] + + use super::*; + use crate::FakeDataset; + use rand::SeedableRng; + use std::collections::HashMap; + + #[test] + fn test_samplerdataset_options() { + let options = SamplerDatasetOptions::default(); + assert_eq!(options.replace_samples, true); + assert_eq!(options.size_config, SizeConfig::Default); + assert_eq!(options.rng_source, RngSource::Default); + + // ReplacementMode + let options = options.with_replace_samples(false); + assert_eq!(options.replace_samples, false); + let options = options.with_replacement(); + assert_eq!(options.replace_samples, true); + let options = options.without_replacement(); + assert_eq!(options.replace_samples, false); + + // SourceSize + let options = options.with_size(SizeConfig::Default); + assert_eq!(options.size_config, SizeConfig::Default); + let options = options.with_source_size(); + assert_eq!(options.size_config, SizeConfig::Default); + let options = options.with_fixed_size(10); + assert_eq!(options.size_config, SizeConfig::Fixed(10)); + let options = options.with_size_ratio(1.5); + assert_eq!(options.size_config, SizeConfig::Ratio(1.5)); + + // RngSource + let options = options.with_system_rng(); + assert_eq!(options.rng_source, RngSource::Default); + let options = options.with_seed(42); + assert_eq!(options.rng_source, RngSource::Seed(42)); + let rng = StdRng::seed_from_u64(9); + let options = options.with_rng(rng); + assert!(matches!(options.rng_source, RngSource::Rng(_))); + } + + #[test] + fn sampler_dataset_constructors_test() { + let ds = SamplerDataset::new(FakeDataset::::new(10), 15); + assert_eq!(ds.len(), 15); + assert_eq!(ds.dataset.len(), 10); + assert!(ds.is_with_replacement()); + + let ds = SamplerDataset::with_replacement(FakeDataset::::new(10), 15); + assert_eq!(ds.len(), 15); + assert_eq!(ds.dataset.len(), 10); + assert!(ds.is_with_replacement()); + + let ds = SamplerDataset::without_replacement(FakeDataset::::new(10), 15); + assert_eq!(ds.len(), 15); + assert_eq!(ds.dataset.len(), 10); + assert!(!ds.is_with_replacement()); + } + + #[test] + fn sampler_dataset_with_replacement_iter() { + let factor = 3; + let len_original = 10; + let dataset_sampler = SamplerDataset::with_replacement( + FakeDataset::::new(len_original), + len_original * factor, + ); + let mut total = 0; + + for _item in dataset_sampler.iter() { + total += 1; + } + + assert_eq!(total, factor * len_original); + } + + #[test] + fn sampler_dataset_without_replacement_bucket_test() { + let factor = 3; + let len_original = 10; + + let dataset_sampler = SamplerDataset::new( + FakeDataset::::new(len_original), + SamplerDatasetOptions::default() + .without_replacement() + .with_size_ratio(factor as f64), + ); + + let mut buckets = HashMap::new(); + + for item in dataset_sampler.iter() { + let count = match buckets.get(&item) { + Some(count) => count + 1, + None => 1, + }; + + buckets.insert(item, count); + } + + let mut total = 0; + for count in buckets.into_values() { + assert_eq!(count, factor); + total += count; + } + assert_eq!(total, factor * len_original); + } + + #[test] + fn sampler_dataset_without_replacement_uniform_order_test() { + // This is a reversion test on the indices.shuffle(rng) call in SamplerDataset::index(). + let size = 1000; + let dataset_sampler = + SamplerDataset::without_replacement(FakeDataset::::new(size), size); + + let indices: Vec<_> = (0..size).map(|_| dataset_sampler.index()).collect(); + let mean_delta = indices + .windows(2) + .map(|pair| pair[1].abs_diff(pair[0])) + .sum::() as f64 + / (size - 1) as f64; + + let expected = (size + 2) as f64 / 3.0; + + assert!( + (mean_delta - expected).abs() <= 0.25 * expected, + "Sampled indices are not uniformly distributed: mean_delta: {mean_delta}, expected: {expected}" + ); + } +} diff --git a/crates/burn-dataset/src/transform/selection.rs b/crates/burn-dataset/src/transform/selection.rs new file mode 100644 index 0000000..283be57 --- /dev/null +++ b/crates/burn-dataset/src/transform/selection.rs @@ -0,0 +1,374 @@ +use crate::Dataset; +use crate::transform::RngSource; +use rand::prelude::SliceRandom; +use rand::rngs::StdRng; +use std::marker::PhantomData; +use std::sync::Arc; + +/// Generates a vector of indices from 0 to size - 1. +/// +/// # Arguments +/// +/// * `size` - The size of the dataset. +/// +/// # Returns +/// +/// A vector containing indices from 0 to size - 1. +#[inline(always)] +pub fn iota(size: usize) -> Vec { + (0..size).collect() +} + +/// Generates a shuffled vector of indices up to a size. +/// +/// # Arguments +/// +/// * `size` - The size of the dataset to shuffle. +/// +/// # Returns +/// +/// A vector of shuffled indices. +#[inline(always)] +pub fn shuffled_indices(size: usize, rng: &mut StdRng) -> Vec { + let mut indices = iota(size); + indices.shuffle(rng); + indices +} + +/// A dataset that selects a subset of indices from an existing dataset. +/// +/// Indices may appear multiple times, but they must be within the bounds of the original dataset. +#[derive(Clone)] +pub struct SelectionDataset +where + D: Dataset, + I: Clone + Send + Sync, +{ + /// The wrapped dataset from which to select indices. + pub wrapped: Arc, + + /// The indices to select from the wrapped dataset. + pub indices: Vec, + + input: PhantomData, +} + +impl SelectionDataset +where + D: Dataset, + I: Clone + Send + Sync, +{ + /// Creates a new selection dataset with the given dataset and indices. + /// + /// Checks that all indices are within the bounds of the dataset. + /// + /// # Arguments + /// + /// * `dataset` - The original dataset to select from. + /// * `indices` - A slice of indices to select from the dataset. + /// These indices must be within the bounds of the dataset. + /// + /// # Panics + /// + /// Panics if any index is out of bounds for the dataset. + pub fn from_indices_checked(dataset: S, indices: Vec) -> Self + where + S: Into>, + { + let dataset = dataset.into(); + + let size = dataset.len(); + if let Some(idx) = indices.iter().find(|&i| *i >= size) { + panic!("Index out of bounds for wrapped dataset size: {idx} >= {size}"); + } + + Self::from_indices_unchecked(dataset, indices) + } + + /// Creates a new selection dataset with the given dataset and indices without checking bounds. + /// + /// # Arguments + /// + /// * `dataset` - The original dataset to select from. + /// * `indices` - A vector of indices to select from the dataset. + /// + /// # Safety + /// + /// This function does not check if the indices are within the bounds of the dataset. + pub fn from_indices_unchecked(dataset: S, indices: Vec) -> Self + where + S: Into>, + { + Self { + wrapped: dataset.into(), + indices, + input: PhantomData, + } + } + + /// Creates a new selection dataset that selects all indices from the dataset. + /// + /// This allocates a 1-to-1 mapping of indices to the dataset size, + /// essentially functioning as a no-op selection. This is only useful + /// when the dataset will later be shuffled or transformed in place. + /// + /// # Arguments + /// + /// * `dataset` - The original dataset to select from. + /// + /// # Returns + /// + /// A new `SelectionDataset` that selects all indices from the dataset. + pub fn new_select_all(dataset: S) -> Self + where + S: Into>, + { + let dataset = dataset.into(); + let size = dataset.len(); + Self::from_indices_unchecked(dataset, iota(size)) + } + + /// Creates a new selection dataset with shuffled indices. + /// + /// Selects every index of the dataset and shuffles them + /// with randomness from the provided random number generator. + /// + /// # Arguments + /// + /// * `dataset` - The original dataset to select from. + /// * `rng` - A mutable reference to a random number generator. + /// + /// # Returns + /// + /// A new `SelectionDataset` with shuffled indices. + pub fn new_shuffled(dataset: S, rng_source: R) -> Self + where + S: Into>, + R: Into, + { + let mut this = Self::new_select_all(dataset); + this.shuffle(rng_source); + this + } + + /// Shuffles the indices of the dataset using a mutable random number generator. + /// + /// This method modifies the dataset in place, shuffling the indices. + /// + /// # Arguments + /// + /// * `rng` - A mutable reference to a random number generator. + pub fn shuffle(&mut self, rng_source: R) + where + R: Into, + { + let mut rng: StdRng = rng_source.into().into(); + self.indices.shuffle(&mut rng) + } + + /// Creates a new dataset that is a slice of the current selection dataset. + /// + /// Slices the *selection indices* from ``[start..end]``. + /// + /// Independent of future shuffles on the parent, but shares the same wrapped dataset. + /// + /// + /// # Arguments + /// + /// * `start` - The start of the range. + /// * `end` - The end of the range (exclusive). + // TODO: SliceArg in burn-tensor should be lifted to burn-std; this should use SliceArg. + pub fn slice(&self, start: usize, end: usize) -> Self { + Self::from_indices_unchecked(self.wrapped.clone(), self.indices[start..end].to_vec()) + } + + /// Split into `num` datasets by slicing the selection indices evenly. + /// + /// Split is done via `slice`, so the datasets share the same wrapped dataset. + /// + /// Independent of future shuffles on the parent, but shares the same wrapped dataset. + /// + /// # Arguments + /// + /// * `num` - The number of datasets to split into. + /// + /// # Returns + /// + /// A vector of `SelectionDataset` instances, each containing a subset of the indices. + pub fn split(&self, num: usize) -> Vec { + let n = self.indices.len(); + + let mut current = 0; + let mut datasets = Vec::with_capacity(num); + + let batch_size = n / num; + for i in 0..num { + let start = current; + let mut end = current + batch_size; + + if i == (num - 1) { + end = n; + } + + let dataset = self.slice(start, end); + + current += batch_size; + datasets.push(dataset); + } + + datasets + } +} + +impl Dataset for SelectionDataset +where + D: Dataset, + I: Clone + Send + Sync, +{ + fn get(&self, index: usize) -> Option { + let index = self.indices.get(index)?; + self.wrapped.get(*index) + } + + fn len(&self) -> usize { + self.indices.len() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::FakeDataset; + use rand::SeedableRng; + + #[test] + fn test_iota() { + let size = 10; + let indices = iota(size); + assert_eq!(indices.len(), size); + assert_eq!(indices, vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); + } + + #[test] + fn test_shuffled_indices_same_seed_is_deterministic() { + let size = 10; + + let mut rng1 = StdRng::seed_from_u64(10); + // `StdRng` is no longer `Clone`, so its internal state cannot be duplicated. + // To test determinism, we must explicitly create a second RNG from the same seed. + let mut rng2 = StdRng::seed_from_u64(10); + + let mut expected = iota(size); + expected.shuffle(&mut rng1); + + let indices = shuffled_indices(size, &mut rng2); + + assert_eq!(indices, expected); + } + + #[test] + fn test_shuffled_indices_forked_rngs_differ() { + let size = 10; + + let mut rng1 = StdRng::seed_from_u64(10); + let mut rng2 = rng1.fork(); + + let mut a = iota(size); + let mut b = iota(size); + + a.shuffle(&mut rng1); + b.shuffle(&mut rng2); + + assert_ne!(a, b); + } + + #[should_panic(expected = "Index out of bounds for wrapped dataset size: 300 >= 27")] + #[test] + fn test_from_indices_checked_panics() { + let source_dataset = FakeDataset::::new(27); + let indices: Vec = vec![15, 1, 12, 300]; + SelectionDataset::from_indices_checked(source_dataset, indices); + } + + #[test] + fn test_checked_selection_dataset() { + let source_dataset = FakeDataset::::new(27); + + let indices: Vec = vec![15, 1, 12, 12]; + let expected: Vec = indices + .iter() + .map(|i| source_dataset.get(*i).unwrap()) + .collect(); + + let selection = SelectionDataset::from_indices_checked(source_dataset, indices.clone()); + + assert_eq!(&selection.indices, &indices); + + let items = selection.iter().collect::>(); + + assert_eq!(items, expected); + } + + #[test] + fn test_shuffled_dataset() { + let dataset = FakeDataset::::new(27); + let source_items = dataset.iter().collect::>(); + + let selection = SelectionDataset::new_shuffled(dataset, 42); + + let indices = shuffled_indices(source_items.len(), &mut StdRng::seed_from_u64(42)); + + assert_eq!(&selection.indices, &indices); + assert_eq!(selection.len(), source_items.len()); + + let expected_items: Vec<_> = indices + .iter() + .map(|&i| source_items[i].to_string()) + .collect(); + assert_eq!(&selection.iter().collect::>(), &expected_items); + } + + #[test] + fn test_slice() { + let dataset = FakeDataset::::new(27); + let source_items = dataset.iter().collect::>(); + + let selection = SelectionDataset::new_select_all(dataset); + + let start = 5; + let end = 15; + let sliced_selection = selection.slice(start, end); + + assert_eq!(sliced_selection.len(), end - start); + + #[allow(clippy::needless_range_loop)] + for i in start..end { + assert_eq!( + sliced_selection.get(i - start), + Some(source_items[i].to_string()) + ); + } + } + + #[test] + fn test_split() { + let dataset = FakeDataset::::new(28); + let source_items = dataset.iter().collect::>(); + + let selection = SelectionDataset::new_select_all(dataset); + + let split_contents: Vec> = selection + .split(3) + .iter() + .map(|d| d.iter().collect::>()) + .collect(); + assert_eq!( + split_contents, + vec![ + source_items[0..9].to_vec(), + source_items[9..18].to_vec(), + source_items[18..28].to_vec(), + ] + ); + } +} diff --git a/crates/burn-dataset/src/transform/shuffle.rs b/crates/burn-dataset/src/transform/shuffle.rs new file mode 100644 index 0000000..8e93e49 --- /dev/null +++ b/crates/burn-dataset/src/transform/shuffle.rs @@ -0,0 +1,109 @@ +use crate::Dataset; +use crate::transform::{RngSource, SelectionDataset}; + +/// A Shuffled a dataset. +/// +/// This is a thin wrapper around a [SelectionDataset] which selects and shuffles +/// the full indices of the original dataset. +/// +/// Consider using [SelectionDataset] if you are only interested in +/// shuffling mechanisms. +/// +/// Consider using [sampler dataset](crate::transform::SamplerDataset) if you +/// want a probability distribution which is computed lazily. +pub struct ShuffledDataset +where + D: Dataset, + I: Clone + Send + Sync, +{ + wrapped: SelectionDataset, +} + +impl ShuffledDataset +where + D: Dataset, + I: Clone + Send + Sync, +{ + /// Creates a new selection dataset with shuffled indices. + /// + /// This is a thin wrapper around `SelectionDataset::new_shuffled`. + /// + /// # Arguments + /// + /// * `dataset` - The original dataset to select from. + /// * `rng_source` - The source of the random number generator. + /// + /// # Returns + /// + /// A new `ShuffledDataset`. + pub fn new(dataset: D, rng_source: R) -> Self + where + R: Into, + { + Self { + wrapped: SelectionDataset::new_shuffled(dataset, rng_source), + } + } + + /// Creates a new selection dataset with shuffled indices using a fixed seed. + /// + /// This is a thin wrapper around `SelectionDataset::new_shuffled_with_seed`. + /// + /// # Arguments + /// + /// * `dataset` - The original dataset to select from. + /// * `seed` - A fixed seed for the random number generator. + /// + /// # Returns + /// + /// A new `ShuffledDataset`. + #[deprecated(since = "0.19.0", note = "Use `new(dataset, seed)` instead`")] + pub fn with_seed(dataset: D, seed: u64) -> Self { + Self::new(dataset, seed) + } +} + +impl Dataset for ShuffledDataset +where + D: Dataset, + I: Clone + Send + Sync, +{ + fn get(&self, index: usize) -> Option { + self.wrapped.get(index) + } + + fn len(&self) -> usize { + self.wrapped.len() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::FakeDataset; + use crate::transform::selection::shuffled_indices; + use rand::SeedableRng; + use rand::prelude::StdRng; + + #[test] + fn test_shuffled_dataset() { + let dataset = FakeDataset::::new(27); + let source_items = dataset.iter().collect::>(); + + let seed = 42; + + #[allow(deprecated)] + let shuffled = ShuffledDataset::with_seed(dataset, seed); + + let mut rng = StdRng::seed_from_u64(seed); + let indices = shuffled_indices(source_items.len(), &mut rng); + + assert_eq!(shuffled.len(), source_items.len()); + + let expected_items: Vec<_> = indices + .iter() + .map(|&i| source_items[i].to_string()) + .collect(); + assert_eq!(&shuffled.iter().collect::>(), &expected_items); + } +} diff --git a/crates/burn-dataset/src/transform/window.rs b/crates/burn-dataset/src/transform/window.rs new file mode 100644 index 0000000..e6bb8b9 --- /dev/null +++ b/crates/burn-dataset/src/transform/window.rs @@ -0,0 +1,290 @@ +use std::{cmp::max, marker::PhantomData, num::NonZeroUsize}; + +use crate::Dataset; + +/// Functionality to create a window. +pub trait Window { + /// Creates a window of a collection. + /// + /// # Returns + /// + /// A `Vec` representing the window. + fn window(&self, current: usize, size: NonZeroUsize) -> Option>; +} + +impl + ?Sized> Window for T { + fn window(&self, current: usize, size: NonZeroUsize) -> Option> { + (current..current + size.get()) + .map(|x| self.get(x)) + .collect() + } +} + +/// Functionality to create a `WindowsIterator`. +pub trait Windows { + /// Creates and returns an iterator over all the windows of length `size`. + fn windows(&self, size: usize) -> WindowsIterator<'_, I>; +} + +impl> Windows for T { + /// Is empty if the `Dataset` is shorter than `size`. + /// + /// # Panics + /// + /// Panics if `size` is 0. + /// + /// # Examples + /// + /// ``` + /// use crate::burn_dataset::{ + /// transform::{Windows, WindowsDataset}, + /// Dataset, InMemDataset, + /// }; + /// + /// let items = [1, 2, 3, 4].to_vec(); + /// let dataset = InMemDataset::new(items.clone()); + /// + /// for window in dataset.windows(2) { + /// // do sth with window + /// } + /// ``` + fn windows(&self, size: usize) -> WindowsIterator<'_, I> { + let size = NonZeroUsize::new(size).expect("window size must be non-zero"); + WindowsIterator::new(self, size) + } +} + +/// Overlapping windows iterator. +pub struct WindowsIterator<'a, I> { + /// The size of the windows. + pub size: NonZeroUsize, + current: usize, + dataset: &'a dyn Dataset, +} + +impl<'a, I> WindowsIterator<'a, I> { + /// Creates a new `WindowsIterator` instance. The windows overlap. + /// Is empty if the input `Dataset` is shorter than `size`. + /// + /// # Parameters + /// + /// - `dataset`: The dataset over which windows will be created. + /// - `size`: The size of the windows. + pub fn new(dataset: &'a dyn Dataset, size: NonZeroUsize) -> Self { + WindowsIterator { + current: 0, + dataset, + size, + } + } +} + +impl Iterator for WindowsIterator<'_, I> { + type Item = Vec; + + fn next(&mut self) -> Option> { + self.current += 1; + self.dataset.window(self.current - 1, self.size) + } +} + +impl Clone for WindowsIterator<'_, I> { + fn clone(&self) -> Self { + WindowsIterator { + size: self.size, + dataset: self.dataset, + current: self.current, + } + } +} + +/// Dataset designed to work with overlapping windows of data. +pub struct WindowsDataset { + /// The size of the windows. + pub size: NonZeroUsize, + dataset: D, + input: PhantomData, +} + +impl WindowsDataset +where + D: Dataset, +{ + /// Creates a new `WindowsDataset` instance. The windows overlap. + /// Is empty if the input `Dataset` is shorter than `size`. + /// + /// # Parameters + /// + /// - `dataset`: The dataset over which windows will be created. + /// - `size`: The size of the windows. + pub fn new(dataset: D, size: usize) -> Self + where + D:, + { + let size = NonZeroUsize::new(size).expect("window size must be non-zero"); + WindowsDataset:: { + size, + dataset, + input: PhantomData, + } + } +} + +impl Dataset> for WindowsDataset +where + D: Dataset, + I: Send + Sync, +{ + /// Retrieves a window of items from the dataset. + /// + /// # Parameters + /// + /// - `index`: The index of the window. + /// + /// # Returns + /// + /// A vector representing the window. + fn get(&self, index: usize) -> Option> { + self.dataset.window(index, self.size) + } + + /// Retrieves the number of windows in the dataset. + /// + /// # Returns + /// + /// A size representing the number of windows. + fn len(&self) -> usize { + let len = self.dataset.len() as isize - self.size.get() as isize + 1; + max(len, 0) as usize + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use crate::{ + Dataset, InMemDataset, + transform::{Windows, WindowsDataset}, + }; + + #[rstest] + pub fn windows_should_be_equal_to_vec_windows() { + let items = [1, 2, 3, 4, 5].to_vec(); + let dataset = InMemDataset::new(items.clone()); + let expected = items + .windows(3) + .map(|x| x.to_vec()) + .collect::>>(); + + let result = dataset.windows(3).collect::>>(); + + assert_eq!(result, expected); + } + + #[rstest] + pub fn windows_dataset_should_be_equal_to_vec_windows() { + let items = [1, 2, 3, 4, 5].to_vec(); + let dataset = InMemDataset::new(items.clone()); + let expected = items + .windows(3) + .map(|x| x.to_vec()) + .collect::>>(); + + let result = WindowsDataset::new(dataset, 3) + .iter() + .collect::>>(); + + assert_eq!(result, expected); + } + + #[rstest] + pub fn cloned_iterator_should_be_equal() { + let items = [1, 2, 3, 4, 5].to_vec(); + let dataset = InMemDataset::new(items.clone()); + let original = dataset.windows(4); + + let cloned = original.clone(); + + assert!(std::ptr::eq(cloned.dataset, original.dataset)); + assert_eq!(cloned.size, original.size); + assert_eq!(cloned.current, original.current); + } + + #[rstest] + pub fn cloned_iterator_should_be_unaffected() { + let items = [1, 2, 3, 4, 5].to_vec(); + let dataset = InMemDataset::new(items.clone()); + let mut original = dataset.windows(4); + + let cloned = original.clone(); + original.current = 2; + + assert_ne!(cloned.current, original.current); + } + + #[rstest] + #[should_panic(expected = "window size must be non-zero")] + pub fn windows_should_panic() { + let items = [1, 2].to_vec(); + let dataset = InMemDataset::new(items.clone()); + + dataset.windows(0); + } + + #[rstest] + #[should_panic(expected = "window size must be non-zero")] + pub fn new_window_dataset_should_panic() { + let items = [1, 2].to_vec(); + let dataset = InMemDataset::new(items.clone()); + + WindowsDataset::new(dataset, 0); + } + + #[rstest] + pub fn window_dataset_len_should_be_equal() { + let dataset = InMemDataset::new([1, 2, 3, 4].to_vec()); + + let result = WindowsDataset::new(dataset, 2).len(); + + assert_eq!(result, 3); + } + + #[rstest] + pub fn window_iterator_should_be_empty() { + let dataset = InMemDataset::new([1, 2].to_vec()); + let mut peekable = dataset.windows(4).peekable(); + + let result = peekable.peek(); + + assert_eq!(result, None); + } + + #[rstest] + pub fn window_dataset_len_should_be_zero() { + let dataset = InMemDataset::new([1, 2].to_vec()); + + let result = WindowsDataset::new(dataset, 4).len(); + + assert_eq!(result, 0); + } + + #[rstest] + pub fn window_dataset_get_should_be_equal() { + let dataset = InMemDataset::new([1, 2, 3, 4].to_vec()); + let expected = Some([1, 2, 3].to_vec()); + + let result = WindowsDataset::new(dataset, 3).get(0); + + assert_eq!(result, expected); + } + + #[rstest] + pub fn window_dataset_get_should_be_none() { + let dataset = InMemDataset::new([1, 2].to_vec()); + + let result = WindowsDataset::new(dataset, 4).get(0); + + assert_eq!(result, None); + } +} diff --git a/crates/burn-dataset/src/vision/cifar.rs b/crates/burn-dataset/src/vision/cifar.rs new file mode 100644 index 0000000..1598db2 --- /dev/null +++ b/crates/burn-dataset/src/vision/cifar.rs @@ -0,0 +1,241 @@ +//! CIFAR Dataset Module +//! +//! This module provides functionality for loading the CIFAR-10 and CIFAR-100 image classification datasets. +//! CIFAR (Canadian Institute For Advanced Research) datasets are widely used benchmarks in computer vision, +//! consisting of 32×32 pixel color images split into training (50,000 images) and test (10,000 images) sets. +//! +//! ## Dataset Variants +//! - **CIFAR-10**: Contains 10 distinct classes (e.g., airplane, automobile, bird, cat) +//! - CIFAR-10 mirror from [fastai](https://github.com/fastai/fastai/blob/master/fastai/data/external.py#L44). +//! - Licensed under the [Apache License](https://github.com/fastai/fastai/blob/master/LICENSE). +//! - **CIFAR-100**: Contains 100 fine-grained classes (e.g., beaver, dolphin, oak tree) +//! - CIFAR-100 mirror from [fastai](https://github.com/fastai/fastai/blob/master/fastai/data/external.py#L75). +//! - Licensed under the [Apache License](https://github.com/fastai/fastai/blob/master/LICENSE). +//! +//! ## Usage Example +//! ```rust +//! use burn_dataset::vision::CifarDataset; +//! use burn_dataset::vision::CifarType; +//! +//! // Create a CIFAR-10 dataset accessor +//! let dataset = CifarDataset::new(CifarType::Cifar10); +//! +//! // Access training and test sets +//! let train_dataset = dataset.train(); +//! let test_dataset = dataset.test(); +//! ``` +//! ```rust +//! use burn_dataset::vision::CifarDataset; +//! use burn_dataset::vision::CifarType; +//! +//! // Create a CIFAR-100 dataset accessor +//! let dataset = CifarDataset::new(CifarType::Cifar100); +//! +//! // Access training and test sets +//! let train_dataset = dataset.train(); +//! let test_dataset = dataset.test(); +//! ``` + +use std::{path::PathBuf, sync::Mutex}; + +use flate2::read::GzDecoder; +use tar::Archive; + +use crate::network::downloader; +use crate::vision::ImageFolderDataset; + +/// CIFAR-10 mirror from [fastai](https://github.com/fastai/fastai/blob/master/fastai/data/external.py#L44). +/// Licensed under the [Apache License](https://github.com/fastai/fastai/blob/master/LICENSE). +const CIFAR10_URL: &str = "https://s3.amazonaws.com/fast-ai-sample/cifar10.tgz"; + +/// CIFAR-100 mirror from [fastai](https://github.com/fastai/fastai/blob/master/fastai/data/external.py#L75). +/// Licensed under the [Apache License](https://github.com/fastai/fastai/blob/master/LICENSE). +const CIFAR100_URL: &str = "https://s3.amazonaws.com/fast-ai-imageclas/cifar100.tgz"; + +/// Enum representing the types of CIFAR datasets available. +/// +/// CIFAR (Canadian Institute For Advanced Research) datasets are widely used benchmarks for image classification. +/// This enum provides support for the two main CIFAR datasets. +#[derive(Debug, Clone, Copy)] +#[allow(dead_code)] +pub enum CifarType { + /// CIFAR-10 dataset containing 10 classes with 60,000 images in total. + Cifar10, + /// CIFAR-100 dataset containing 100 classes with 60,000 images in total. + Cifar100, +} + +/// CIFAR dataset accessor. +/// +/// This struct provides convenient access to the CIFAR-10 and CIFAR-100 image classification datasets. +/// It automatically downloads (if not already downloaded), extracts, and loads the datasets. +/// +/// All images in CIFAR datasets are 32×32 pixel color images, with 50,000 images in the training set +/// and 10,000 images in the test set. +/// +/// ## Differences between datasets +/// - **CIFAR-10**: Contains 10 mutually exclusive classes such as airplane, automobile, bird, cat, etc. +/// - **CIFAR-100**: Contains 100 fine-grained classes such as beaver, dolphin, etc. +pub struct CifarDataset { + cifar_dir: PathBuf, +} + +impl CifarDataset { + /// Creates a new CIFAR dataset accessor. + /// + /// # Arguments + /// * `cifar_type` - Specifies whether to use CIFAR-10 or CIFAR-100 dataset + pub fn new(cifar_type: CifarType) -> Self { + Self { + cifar_dir: download(&cifar_type), + } + } + + /// Gets the training dataset. + /// + /// # Returns + /// An `ImageFolderDataset` instance containing 50,000 training images + pub fn train(&self) -> ImageFolderDataset { + ImageFolderDataset::new_classification(self.cifar_dir.join("train")).unwrap() + } + + /// Gets the test dataset. + /// + /// # Returns + /// An `ImageFolderDataset` instance containing 10,000 test images + pub fn test(&self) -> ImageFolderDataset { + ImageFolderDataset::new_classification(self.cifar_dir.join("test")).unwrap() + } +} + +/// CIFAR dataset download lock. +/// +/// This lock ensures that only one thread downloads the CIFAR dataset at a time. +static DOWNLOAD_LOCK: Mutex<()> = Mutex::new(()); + +fn download(cifar_type: &CifarType) -> PathBuf { + // Acquire the lock. This will block if another thread already holds the lock. + let _lock = DOWNLOAD_LOCK.lock().unwrap(); + + // Dataset files are stored in the burn-dataset cache directory + let cache_dir = dirs::cache_dir() + .expect("Could not get cache directory") + .join("burn-dataset"); + + // Cifar store directory + let cifar_dir = match cifar_type { + CifarType::Cifar10 => cache_dir.join("cifar10"), + CifarType::Cifar100 => cache_dir.join("cifar100"), + }; + + // Cifar dataset url + let url = match cifar_type { + CifarType::Cifar10 => CIFAR10_URL, + CifarType::Cifar100 => CIFAR100_URL, + }; + + // Cifar dataset archive filename + let filename = match cifar_type { + CifarType::Cifar10 => "cifar10.tgz", + CifarType::Cifar100 => "cifar100.tgz", + }; + + // Check for already downloaded content + if !cifar_dir.exists() { + // Download gzip file + let bytes = downloader::download_file_as_bytes(url, filename); + + // Decode gzip file content and unpack archive + let gz_buffer = GzDecoder::new(&bytes[..]); + let mut archive = Archive::new(gz_buffer); + archive.unpack(cache_dir).unwrap(); + } + + cifar_dir +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{Dataset, vision::Annotation}; + + /// CIFAR dataset length + const TRAINDATASET_LEN: usize = 50000; + const TESTDATASET_LEN: usize = 10000; + + /// CIFAR-10 label range + const CIFAR10_LABEL_MIN: usize = 0; + const CIFAR10_LABEL_MAX: usize = 9; + + /// CIFAR-100 label range + const CIFAR100_LABEL_MIN: usize = 0; + const CIFAR100_LABEL_MAX: usize = 99; + + #[test] + fn test_cifar10_download() { + let cifar_dir = download(&CifarType::Cifar10); + assert!(cifar_dir.exists()); + } + + #[test] + fn test_cifar100_download() { + let cifar_dir = download(&CifarType::Cifar100); + assert!(cifar_dir.exists()); + } + + #[test] + fn test_cifar10_len() { + let dataset = CifarDataset::new(CifarType::Cifar10); + let train_dataset = dataset.train(); + let test_dataset = dataset.test(); + assert_eq!(train_dataset.len(), TRAINDATASET_LEN); + assert_eq!(test_dataset.len(), TESTDATASET_LEN); + } + + #[test] + fn test_cifar100_len() { + let dataset = CifarDataset::new(CifarType::Cifar100); + let train_dataset = dataset.train(); + let test_dataset = dataset.test(); + assert_eq!(train_dataset.len(), TRAINDATASET_LEN); + assert_eq!(test_dataset.len(), TESTDATASET_LEN); + } + + #[test] + fn test_cifar10_label_range() { + let dataset = CifarDataset::new(CifarType::Cifar10); + let test_dataset = dataset.test(); + let (min, max) = get_label_range(&test_dataset); + assert_eq!(min, CIFAR10_LABEL_MIN); + assert_eq!(max, CIFAR10_LABEL_MAX); + } + + #[test] + fn test_cifar100_label_range() { + let dataset = CifarDataset::new(CifarType::Cifar100); + let test_dataset = dataset.test(); + let (min, max) = get_label_range(&test_dataset); + assert_eq!(min, CIFAR100_LABEL_MIN); + assert_eq!(max, CIFAR100_LABEL_MAX); + } + + fn get_label_range(dataset: &ImageFolderDataset) -> (usize, usize) { + let labels: Vec<_> = dataset.iter().map(|item| item.annotation).collect(); + let mut min = 128; + let mut max = 0; + for label in labels { + let index = match label { + Annotation::Label(index) => index, + _ => 0, + }; + if index < min { + min = index; + } + if index > max { + max = index; + } + } + + (min, max) + } +} diff --git a/crates/burn-dataset/src/vision/image_folder.rs b/crates/burn-dataset/src/vision/image_folder.rs new file mode 100644 index 0000000..63d0ad5 --- /dev/null +++ b/crates/burn-dataset/src/vision/image_folder.rs @@ -0,0 +1,1124 @@ +use crate::transform::{Mapper, MapperDataset}; +use crate::{Dataset, InMemDataset}; + +use globwalk::{self, DirEntry}; +use image::{self, ColorType}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::{HashMap, HashSet}; +use std::fs; +use std::path::{Path, PathBuf}; +use thiserror::Error; + +const SUPPORTED_FILES: [&str; 4] = ["bmp", "jpg", "jpeg", "png"]; +const BBOX_MIN_NUM_VALUES: usize = 4; + +/// Image data type. +#[derive(Debug, Copy, Clone, PartialEq)] +pub enum PixelDepth { + /// 8-bit unsigned. + U8(u8), + /// 16-bit unsigned. + U16(u16), + /// 32-bit floating point. + F32(f32), +} + +impl TryFrom for u8 { + type Error = &'static str; + + fn try_from(value: PixelDepth) -> Result { + if let PixelDepth::U8(v) = value { + Ok(v) + } else { + Err("Value is not u8") + } + } +} + +impl TryFrom for u16 { + type Error = &'static str; + + fn try_from(value: PixelDepth) -> Result { + if let PixelDepth::U16(v) = value { + Ok(v) + } else { + Err("Value is not u16") + } + } +} + +impl TryFrom for f32 { + type Error = &'static str; + + fn try_from(value: PixelDepth) -> Result { + if let PixelDepth::F32(v) = value { + Ok(v) + } else { + Err("Value is not f32") + } + } +} + +/// Annotation type for different tasks. +#[derive(Debug, Clone, PartialEq)] +pub enum Annotation { + /// Image-level label. + Label(usize), + /// Multiple image-level labels. + MultiLabel(Vec), + /// Object bounding boxes. + BoundingBoxes(Vec), + /// Segmentation mask. + SegmentationMask(SegmentationMask), +} + +/// Segmentation mask annotation. +/// For semantic segmentation, a mask has a single channel (C = 1). +/// For instance segmentation, there may be multiple masks per image (C >= 1). +#[derive(Debug, Clone, PartialEq)] +pub struct SegmentationMask { + /// Segmentation mask. + pub mask: Vec, +} + +/// Object detection bounding box annotation. +#[derive(Deserialize, Serialize, Debug, Clone, PartialEq)] +pub struct BoundingBox { + /// Coordinates in [x_min, y_min, width, height] format. + pub coords: [f32; 4], + + /// Box class label. + pub label: usize, +} + +/// Image dataset item. +#[derive(Debug, Clone, PartialEq)] +pub struct ImageDatasetItem { + /// Image as a vector with a valid image type. + pub image: Vec, + + /// Original source image width. + pub image_width: usize, + + /// Original source image height. + pub image_height: usize, + + /// Annotation for the image. + pub annotation: Annotation, + + /// Original image source. + pub image_path: String, +} + +/// Raw annotation types. +#[derive(Deserialize, Serialize, Debug, Clone)] +enum AnnotationRaw { + Label(String), + MultiLabel(Vec), + BoundingBoxes(Vec), + SegmentationMask(PathBuf), +} + +#[derive(Deserialize, Serialize, Debug, Clone)] +struct ImageDatasetItemRaw { + /// Image path. + image_path: PathBuf, + + /// Image annotation. + annotation: AnnotationRaw, +} + +impl ImageDatasetItemRaw { + fn new>(image_path: P, annotation: AnnotationRaw) -> ImageDatasetItemRaw { + ImageDatasetItemRaw { + image_path: image_path.as_ref().to_path_buf(), + annotation, + } + } +} + +struct PathToImageDatasetItem { + classes: HashMap, +} + +fn segmentation_mask_to_vec_usize(mask_path: &PathBuf) -> Vec { + // Load image from disk + let image = image::open(mask_path).unwrap(); + + // Image as Vec + // if rgb8 or rgb16, keep only the first channel assuming all channels are the same + + match image.color() { + ColorType::L8 => image.into_luma8().iter().map(|&x| x as usize).collect(), + ColorType::L16 => image.into_luma16().iter().map(|&x| x as usize).collect(), + ColorType::Rgb8 => image + .into_rgb8() + .iter() + .step_by(3) + .map(|&x| x as usize) + .collect(), + ColorType::Rgb16 => image + .into_rgb16() + .iter() + .step_by(3) + .map(|&x| x as usize) + .collect(), + _ => panic!("Unrecognized image color type"), + } +} + +/// Parse the image annotation to the corresponding type. +fn parse_image_annotation( + annotation: &AnnotationRaw, + classes: &HashMap, +) -> Annotation { + // TODO: add support for other annotations + // - [ ] Object bounding boxes + // - [x] Segmentation mask + // For now, only image classification labels and segmentation are supported. + + // Map class string to label id + match annotation { + AnnotationRaw::Label(name) => Annotation::Label(*classes.get(name).unwrap()), + AnnotationRaw::MultiLabel(names) => Annotation::MultiLabel( + names + .iter() + .map(|name| *classes.get(name).unwrap()) + .collect(), + ), + AnnotationRaw::SegmentationMask(mask_path) => { + Annotation::SegmentationMask(SegmentationMask { + mask: segmentation_mask_to_vec_usize(mask_path), + }) + } + AnnotationRaw::BoundingBoxes(v) => Annotation::BoundingBoxes(v.clone()), + } +} + +/// Retrieve all available classes from the COCO JSON +fn parse_coco_classes( + json: &serde_json::Value, +) -> Result, ImageLoaderError> { + let mut classes = HashMap::new(); + + if let Some(json_classes) = json["categories"].as_array() { + for class in json_classes { + let id = class["id"] + .as_u64() + .ok_or_else(|| ImageLoaderError::ParsingError("Invalid class ID".to_string())) + .and_then(|v| { + usize::try_from(v).map_err(|_| { + ImageLoaderError::ParsingError("Class ID out of usize range".to_string()) + }) + })?; + + let name = class["name"] + .as_str() + .filter(|&s| !s.is_empty()) + .ok_or_else(|| ImageLoaderError::ParsingError("Invalid class name".to_string()))? + .to_string(); + + classes.insert(name, id); + } + } + + if classes.is_empty() { + return Err(ImageLoaderError::ParsingError( + "No classes found in annotations".to_string(), + )); + } + + Ok(classes) +} + +/// Retrieve annotations from COCO JSON +fn parse_coco_bbox_annotations( + json: &serde_json::Value, +) -> Result, ImageLoaderError> { + let mut annotations = HashMap::new(); + + if let Some(json_annotations) = json["annotations"].as_array() { + for annotation in json_annotations { + let image_id = annotation["image_id"].as_u64().ok_or_else(|| { + ImageLoaderError::ParsingError("Invalid image ID in annotation".into()) + })?; + + let class_id = annotation["category_id"] + .as_u64() + .ok_or_else(|| { + ImageLoaderError::ParsingError("Invalid class ID in annotations".to_string()) + }) + .and_then(|v| { + usize::try_from(v).map_err(|_| { + ImageLoaderError::ParsingError( + "Class ID in annotations out of usize range".to_string(), + ) + }) + })?; + + let bbox_coords = annotation["bbox"] + .as_array() + .ok_or_else(|| ImageLoaderError::ParsingError("missing bbox array".to_string()))? + .iter() + .map(|v| { + v.as_f64() + .ok_or_else(|| { + ImageLoaderError::ParsingError("invalid bbox value".to_string()) + }) + .map(|val| val as f32) + }) + .collect::, _>>()?; + + if bbox_coords.len() < BBOX_MIN_NUM_VALUES { + return Err(ImageLoaderError::ParsingError(format!( + "not enough bounding box coordinates in annotation for image {image_id}", + ))); + } + + let bbox = BoundingBox { + coords: [ + bbox_coords[0], + bbox_coords[1], + bbox_coords[2], + bbox_coords[3], + ], + label: class_id, + }; + + annotations + .entry(image_id) + .and_modify(|entry| { + if let AnnotationRaw::BoundingBoxes(bboxes) = entry { + bboxes.push(bbox.clone()); + } + }) + .or_insert_with(|| AnnotationRaw::BoundingBoxes(vec![bbox])); + } + } + + if annotations.is_empty() { + return Err(ImageLoaderError::ParsingError( + "no annotations found".to_string(), + )); + } + + Ok(annotations) +} + +/// Retrieve all available images from the COCO JSON +fn parse_coco_images>( + images_path: &P, + mut annotations: HashMap, + json: &serde_json::Value, +) -> Result, ImageLoaderError> { + let mut images = Vec::new(); + if let Some(json_images) = json["images"].as_array() { + for image in json_images { + let image_id = image["id"].as_u64().ok_or_else(|| { + ImageLoaderError::ParsingError("Invalid image ID in image list".to_string()) + })?; + + let file_name = image["file_name"] + .as_str() + .ok_or_else(|| ImageLoaderError::ParsingError("Invalid image ID".to_string()))? + .to_string(); + + let mut image_path = images_path.as_ref().to_path_buf(); + image_path.push(file_name); + + if !image_path.exists() { + return Err(ImageLoaderError::IOError(format!( + "Image {} not found", + image_path.display() + ))); + } + + let annotation = annotations + .remove(&image_id) + .unwrap_or_else(|| AnnotationRaw::BoundingBoxes(Vec::new())); + + images.push(ImageDatasetItemRaw { + annotation, + image_path, + }); + } + } + + if images.is_empty() { + return Err(ImageLoaderError::ParsingError( + "No images found in annotations".to_string(), + )); + } + + Ok(images) +} + +impl Mapper for PathToImageDatasetItem { + /// Convert a raw image dataset item (path-like) to a 3D image array with a target label. + fn map(&self, item: &ImageDatasetItemRaw) -> ImageDatasetItem { + let annotation = parse_image_annotation(&item.annotation, &self.classes); + + // Load image from disk + let image = image::open(&item.image_path).unwrap(); + + // Save image dimensions for manipulation + let img_width = image.width() as usize; + let img_height = image.height() as usize; + + // Image as Vec + let img_vec = match image.color() { + ColorType::L8 => image + .into_luma8() + .iter() + .map(|&x| PixelDepth::U8(x)) + .collect(), + ColorType::La8 => image + .into_luma_alpha8() + .iter() + .map(|&x| PixelDepth::U8(x)) + .collect(), + ColorType::L16 => image + .into_luma16() + .iter() + .map(|&x| PixelDepth::U16(x)) + .collect(), + ColorType::La16 => image + .into_luma_alpha16() + .iter() + .map(|&x| PixelDepth::U16(x)) + .collect(), + ColorType::Rgb8 => image + .into_rgb8() + .iter() + .map(|&x| PixelDepth::U8(x)) + .collect(), + ColorType::Rgba8 => image + .into_rgba8() + .iter() + .map(|&x| PixelDepth::U8(x)) + .collect(), + ColorType::Rgb16 => image + .into_rgb16() + .iter() + .map(|&x| PixelDepth::U16(x)) + .collect(), + ColorType::Rgba16 => image + .into_rgba16() + .iter() + .map(|&x| PixelDepth::U16(x)) + .collect(), + ColorType::Rgb32F => image + .into_rgb32f() + .iter() + .map(|&x| PixelDepth::F32(x)) + .collect(), + ColorType::Rgba32F => image + .into_rgba32f() + .iter() + .map(|&x| PixelDepth::F32(x)) + .collect(), + _ => panic!("Unrecognized image color type"), + }; + + ImageDatasetItem { + image: img_vec, + image_width: img_width, + image_height: img_height, + annotation, + image_path: item.image_path.display().to_string(), + } + } +} + +/// Error type for [ImageFolderDataset](ImageFolderDataset). +#[derive(Error, Debug)] +pub enum ImageLoaderError { + /// Unknown error. + #[error("unknown: `{0}`")] + Unknown(String), + + /// I/O operation error. + #[error("I/O error: `{0}`")] + IOError(String), + + /// Invalid file error. + #[error("Invalid file extension: `{0}`")] + InvalidFileExtensionError(String), + + /// Parsing error. + #[error("Parsing error: `{0}`")] + ParsingError(String), +} + +type ImageDatasetMapper = + MapperDataset, PathToImageDatasetItem, ImageDatasetItemRaw>; + +/// A generic dataset to load images from disk. +pub struct ImageFolderDataset { + dataset: ImageDatasetMapper, +} + +impl Dataset for ImageFolderDataset { + fn get(&self, index: usize) -> Option { + self.dataset.get(index) + } + + fn len(&self) -> usize { + self.dataset.len() + } +} + +impl ImageFolderDataset { + /// Create an image classification dataset from the root folder. + /// + /// # Arguments + /// + /// * `root` - Dataset root folder. + /// + /// # Returns + /// A new dataset instance. + pub fn new_classification>(root: P) -> Result { + // New dataset containing any of the supported file types + ImageFolderDataset::new_classification_with(root, &SUPPORTED_FILES) + } + + /// Create an image classification dataset from the root folder. + /// The included images are filtered based on the provided extensions. + /// + /// # Arguments + /// + /// * `root` - Dataset root folder. + /// * `extensions` - List of allowed extensions. + /// + /// # Returns + /// A new dataset instance. + pub fn new_classification_with( + root: P, + extensions: &[S], + ) -> Result + where + P: AsRef, + S: AsRef, + { + // Glob all images with extensions + let walker = globwalk::GlobWalkerBuilder::from_patterns( + root.as_ref(), + &[format!( + "*.{{{}}}", // "*.{ext1,ext2,ext3} + extensions + .iter() + .map(Self::check_extension) + .collect::, _>>()? + .join(",") + )], + ) + .follow_links(true) + .sort_by(|p1: &DirEntry, p2: &DirEntry| p1.path().cmp(p2.path())) // order by path + .build() + .map_err(|err| ImageLoaderError::Unknown(format!("{err:?}")))? + .filter_map(Result::ok); + + // Get all dataset items + let mut items = Vec::new(); + let mut classes = HashSet::new(); + for img in walker { + let image_path = img.path(); + + // Label name is represented by the parent folder name + let label = image_path + .parent() + .ok_or_else(|| { + ImageLoaderError::IOError("Could not resolve image parent folder".to_string()) + })? + .file_name() + .ok_or_else(|| { + ImageLoaderError::IOError( + "Could not resolve image parent folder name".to_string(), + ) + })? + .to_string_lossy() + .into_owned(); + + classes.insert(label.clone()); + + items.push(ImageDatasetItemRaw::new( + image_path, + AnnotationRaw::Label(label), + )) + } + + // Sort class names + let mut classes = classes.into_iter().collect::>(); + classes.sort(); + + Self::with_items(items, &classes) + } + + /// Create an image classification dataset with the specified items. + /// + /// # Arguments + /// + /// * `items` - List of dataset items, each item represented by a tuple `(image path, label)`. + /// * `classes` - Dataset class names. + /// + /// # Returns + /// A new dataset instance. + pub fn new_classification_with_items, S: AsRef>( + items: Vec<(P, String)>, + classes: &[S], + ) -> Result { + // Parse items and check valid image extension types + let items = items + .into_iter() + .map(|(path, label)| { + // Map image path and label + let path = path.as_ref(); + let label = AnnotationRaw::Label(label); + + Self::check_extension(&path.extension().unwrap().to_str().unwrap())?; + + Ok(ImageDatasetItemRaw::new(path, label)) + }) + .collect::, _>>()?; + + Self::with_items(items, classes) + } + + /// Create a multi-label image classification dataset with the specified items. + /// + /// # Arguments + /// + /// * `items` - List of dataset items, each item represented by a tuple `(image path, labels)`. + /// * `classes` - Dataset class names. + /// + /// # Returns + /// A new dataset instance. + pub fn new_multilabel_classification_with_items, S: AsRef>( + items: Vec<(P, Vec)>, + classes: &[S], + ) -> Result { + // Parse items and check valid image extension types + let items = items + .into_iter() + .map(|(path, labels)| { + // Map image path and multi-label + let path = path.as_ref(); + let labels = AnnotationRaw::MultiLabel(labels); + + Self::check_extension(&path.extension().unwrap().to_str().unwrap())?; + + Ok(ImageDatasetItemRaw::new(path, labels)) + }) + .collect::, _>>()?; + + Self::with_items(items, classes) + } + + /// Create an image segmentation dataset with the specified items. + /// + /// # Arguments + /// + /// * `items` - List of dataset items, each item represented by a tuple `(image path, annotation path)`. + /// * `classes` - Dataset class names. + /// + /// # Returns + /// A new dataset instance. + pub fn new_segmentation_with_items, S: AsRef>( + items: Vec<(P, P)>, + classes: &[S], + ) -> Result { + // Parse items and check valid image extension types + let items = items + .into_iter() + .map(|(image_path, mask_path)| { + // Map image path and segmentation mask path + let image_path = image_path.as_ref(); + let annotation = AnnotationRaw::SegmentationMask(mask_path.as_ref().to_path_buf()); + + Self::check_extension(&image_path.extension().unwrap().to_str().unwrap())?; + + Ok(ImageDatasetItemRaw::new(image_path, annotation)) + }) + .collect::, _>>()?; + + Self::with_items(items, classes) + } + + /// Create a COCO detection dataset based on the annotations JSON and image directory. + /// + /// # Arguments + /// + /// * `annotations_json` - Path to the JSON file containing annotations in COCO format (for + /// example instances_train2017.json). + /// + /// * `images_path` - Path containing the images matching the annotations JSON. + /// + /// # Returns + /// A new dataset instance. + pub fn new_coco_detection, I: AsRef>( + annotations_json: A, + images_path: I, + ) -> Result { + let file = fs::File::open(annotations_json) + .map_err(|e| ImageLoaderError::IOError(format!("Failed to open annotations: {e}")))?; + let json: Value = serde_json::from_reader(file).map_err(|e| { + ImageLoaderError::ParsingError(format!("Failed to parse annotations: {e}")) + })?; + + let classes = parse_coco_classes(&json)?; + let annotations = parse_coco_bbox_annotations(&json)?; + let items = parse_coco_images(&images_path, annotations, &json)?; + let dataset = InMemDataset::new(items); + let mapper = PathToImageDatasetItem { classes }; + let dataset = MapperDataset::new(dataset, mapper); + + Ok(Self { dataset }) + } + + /// Create an image dataset with the specified items. + /// + /// # Arguments + /// + /// * `items` - Raw dataset items. + /// * `classes` - Dataset class names. + /// + /// # Returns + /// A new dataset instance. + fn with_items>( + items: Vec, + classes: &[S], + ) -> Result { + // NOTE: right now we don't need to validate the supported image files since + // the method is private. We assume it's already validated. + let dataset = InMemDataset::new(items); + + // Class names to index map + let classes = classes.iter().map(|c| c.as_ref()).collect::>(); + let classes_map: HashMap<_, _> = classes + .into_iter() + .enumerate() + .map(|(idx, cls)| (cls.to_string(), idx)) + .collect(); + + let mapper = PathToImageDatasetItem { + classes: classes_map, + }; + let dataset = MapperDataset::new(dataset, mapper); + + Ok(Self { dataset }) + } + + /// Check if extension is supported. + fn check_extension>(extension: &S) -> Result { + let extension = extension.as_ref(); + if !SUPPORTED_FILES.contains(&extension) { + Err(ImageLoaderError::InvalidFileExtensionError( + extension.to_string(), + )) + } else { + Ok(extension.to_string()) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + const DATASET_ROOT: &str = "tests/data/image_folder"; + const SEGMASK_ROOT: &str = "tests/data/segmask_folder"; + const COCO_JSON: &str = "tests/data/dataset_coco.json"; + const COCO_IMAGES: &str = "tests/data/image_folder_coco"; + + #[test] + pub fn image_folder_dataset() { + let dataset = ImageFolderDataset::new_classification(DATASET_ROOT).unwrap(); + + // Dataset has 3 elements + assert_eq!(dataset.len(), 3); + assert_eq!(dataset.get(3), None); + + // Dataset elements should be: orange (0), red (1), red (1) + assert_eq!(dataset.get(0).unwrap().annotation, Annotation::Label(0)); + assert_eq!(dataset.get(1).unwrap().annotation, Annotation::Label(1)); + assert_eq!(dataset.get(2).unwrap().annotation, Annotation::Label(1)); + } + + #[test] + pub fn image_folder_dataset_filtered() { + let dataset = ImageFolderDataset::new_classification_with(DATASET_ROOT, &["jpg"]).unwrap(); + + // Filtered dataset has 2 elements + assert_eq!(dataset.len(), 2); + assert_eq!(dataset.get(2), None); + + // Dataset elements should be: orange (0), red (1) + assert_eq!(dataset.get(0).unwrap().annotation, Annotation::Label(0)); + assert_eq!(dataset.get(1).unwrap().annotation, Annotation::Label(1)); + } + + #[test] + pub fn image_folder_dataset_with_items_sizes() { + let root = Path::new(DATASET_ROOT); + let items = vec![ + (root.join("orange").join("dot.jpg"), "orange".to_string()), + (root.join("red").join("dot.jpg"), "red".to_string()), + (root.join("red").join("dot.png"), "red".to_string()), + ]; + let dataset = + ImageFolderDataset::new_classification_with_items(items, &["orange", "red"]).unwrap(); + + // Dataset has 3 elements + assert_eq!(dataset.len(), 3); + assert_eq!(dataset.get(3), None); + + // Test item sizes + + assert_eq!( + ( + dataset.get(0).unwrap().image_width, + dataset.get(0).unwrap().image_height + ), + (1, 1) + ); + assert_eq!( + ( + dataset.get(1).unwrap().image_width, + dataset.get(1).unwrap().image_height + ), + (1, 1) + ); + assert_eq!( + ( + dataset.get(2).unwrap().image_width, + dataset.get(2).unwrap().image_height + ), + (1, 1) + ); + } + + #[test] + pub fn image_folder_dataset_with_items() { + let root = Path::new(DATASET_ROOT); + let items = vec![ + (root.join("orange").join("dot.jpg"), "orange".to_string()), + (root.join("red").join("dot.jpg"), "red".to_string()), + (root.join("red").join("dot.png"), "red".to_string()), + ]; + let dataset = + ImageFolderDataset::new_classification_with_items(items, &["orange", "red"]).unwrap(); + + // Dataset has 3 elements + assert_eq!(dataset.len(), 3); + assert_eq!(dataset.get(3), None); + + // Dataset elements should be: orange (0), red (1), red (1) + assert_eq!(dataset.get(0).unwrap().annotation, Annotation::Label(0)); + assert_eq!(dataset.get(1).unwrap().annotation, Annotation::Label(1)); + assert_eq!(dataset.get(2).unwrap().annotation, Annotation::Label(1)); + } + + #[test] + pub fn image_folder_dataset_multilabel() { + let root = Path::new(DATASET_ROOT); + let items = vec![ + ( + root.join("orange").join("dot.jpg"), + vec!["dot".to_string(), "orange".to_string()], + ), + ( + root.join("red").join("dot.jpg"), + vec!["dot".to_string(), "red".to_string()], + ), + ( + root.join("red").join("dot.png"), + vec!["dot".to_string(), "red".to_string()], + ), + ]; + let dataset = ImageFolderDataset::new_multilabel_classification_with_items( + items, + &["dot", "orange", "red"], + ) + .unwrap(); + + // Dataset has 3 elements + assert_eq!(dataset.len(), 3); + assert_eq!(dataset.get(3), None); + + // Dataset elements should be: [dot, orange] (0, 1), [dot, red] (0, 2), [dot, red] (0, 2) + assert_eq!( + dataset.get(0).unwrap().annotation, + Annotation::MultiLabel(vec![0, 1]) + ); + assert_eq!( + dataset.get(1).unwrap().annotation, + Annotation::MultiLabel(vec![0, 2]) + ); + assert_eq!( + dataset.get(2).unwrap().annotation, + Annotation::MultiLabel(vec![0, 2]) + ); + } + + #[test] + #[should_panic] + pub fn image_folder_dataset_invalid_extension() { + // Some invalid file extension + let _ = ImageFolderDataset::new_classification_with(DATASET_ROOT, &["ico"]).unwrap(); + } + + #[test] + pub fn pixel_depth_try_into_u8() { + let val = u8::MAX; + let pix: u8 = PixelDepth::U8(val).try_into().unwrap(); + assert_eq!(pix, val); + } + + #[test] + #[should_panic] + pub fn pixel_depth_try_into_u8_invalid() { + let _: u8 = PixelDepth::U16(u8::MAX as u16 + 1).try_into().unwrap(); + } + + #[test] + pub fn pixel_depth_try_into_u16() { + let val = u16::MAX; + let pix: u16 = PixelDepth::U16(val).try_into().unwrap(); + assert_eq!(pix, val); + } + + #[test] + #[should_panic] + pub fn pixel_depth_try_into_u16_invalid() { + let _: u16 = PixelDepth::F32(u16::MAX as f32).try_into().unwrap(); + } + + #[test] + pub fn pixel_depth_try_into_f32() { + let val = f32::MAX; + let pix: f32 = PixelDepth::F32(val).try_into().unwrap(); + assert_eq!(pix, val); + } + + #[test] + #[should_panic] + pub fn pixel_depth_try_into_f32_invalid() { + let _: f32 = PixelDepth::U16(u16::MAX).try_into().unwrap(); + } + + #[test] + pub fn parse_image_annotation_label_string() { + let classes = HashMap::from([("0".to_string(), 0_usize), ("1".to_string(), 1_usize)]); + let anno = AnnotationRaw::Label("0".to_string()); + assert_eq!( + parse_image_annotation(&anno, &classes), + Annotation::Label(0) + ); + } + + #[test] + pub fn parse_image_annotation_multilabel_string() { + let classes = HashMap::from([ + ("0".to_string(), 0_usize), + ("1".to_string(), 1_usize), + ("2".to_string(), 2_usize), + ]); + let anno = AnnotationRaw::MultiLabel(vec!["0".to_string(), "2".to_string()]); + assert_eq!( + parse_image_annotation(&anno, &classes), + Annotation::MultiLabel(vec![0, 2]) + ); + } + + #[test] + pub fn segmask_image_path_to_vec_usize() { + let root = Path::new(SEGMASK_ROOT); + + // checkerboard mask + const TEST_CHECKERBOARD_MASK_PATTERN: [u8; 64] = [ + 1, 2, 1, 2, 1, 2, 1, 2, 2, 1, 2, 1, 2, 1, 2, 1, 1, 2, 1, 2, 1, 2, 1, 2, 2, 1, 2, 1, 2, + 1, 2, 1, 1, 2, 1, 2, 1, 2, 1, 2, 2, 1, 2, 1, 2, 1, 2, 1, 1, 2, 1, 2, 1, 2, 1, 2, 2, 1, + 2, 1, 2, 1, 2, 1, + ]; + assert_eq!( + TEST_CHECKERBOARD_MASK_PATTERN + .iter() + .map(|&x| x as usize) + .collect::>(), + segmentation_mask_to_vec_usize(&root.join("annotations").join("mask_checkerboard.png")), + ); + + // random 2 colors mask + const TEST_RANDOM2COLORS_MASK_PATTERN: [u8; 64] = [ + 1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1, 2, 2, 2, 1, 2, 1, 2, 2, 2, 2, 2, 2, 2, + 2, 1, 1, 2, 2, 2, 1, 2, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 1, 2, 2, 1, 2, 1, 2, 1, 2, 2, 1, + 1, 1, 1, 1, 1, 1, + ]; + assert_eq!( + TEST_RANDOM2COLORS_MASK_PATTERN + .iter() + .map(|&x| x as usize) + .collect::>(), + segmentation_mask_to_vec_usize( + &root.join("annotations").join("mask_random_2colors.png") + ), + ); + // random 3 colors mask + const TEST_RANDOM3COLORS_MASK_PATTERN: [u8; 64] = [ + 3, 1, 3, 3, 1, 1, 3, 2, 3, 3, 3, 3, 1, 3, 2, 1, 2, 2, 2, 2, 1, 1, 2, 2, 1, 1, 1, 3, 3, + 3, 2, 3, 2, 2, 3, 2, 3, 3, 1, 3, 1, 3, 3, 1, 1, 3, 2, 1, 2, 2, 2, 1, 2, 1, 2, 3, 3, 1, + 3, 3, 2, 1, 2, 2, + ]; + assert_eq!( + TEST_RANDOM3COLORS_MASK_PATTERN + .iter() + .map(|&x| x as usize) + .collect::>(), + segmentation_mask_to_vec_usize( + &root.join("annotations").join("mask_random_3colors.png") + ), + ); + } + + #[test] + pub fn segmask_folder_dataset() { + let root = Path::new(SEGMASK_ROOT); + + let items = vec![ + ( + root.join("images").join("image_checkerboard.png"), + root.join("annotations").join("mask_checkerboard.png"), + ), + ( + root.join("images").join("image_random_2colors.png"), + root.join("annotations").join("mask_random_2colors.png"), + ), + ( + root.join("images").join("image_random_3colors.png"), + root.join("annotations").join("mask_random_3colors.png"), + ), + ]; + let dataset = ImageFolderDataset::new_segmentation_with_items( + items, + &[ + "foo", // 0 + "bar", // 1 + "baz", // 2 + "qux", // 3 + ], + ) + .unwrap(); + + // Dataset has 3 elements; each (image, annotation) is a single item + assert_eq!(dataset.len(), 3); + assert_eq!(dataset.get(3), None); + + // checkerboard mask + const TEST_CHECKERBOARD_MASK_PATTERN: [u8; 64] = [ + 1, 2, 1, 2, 1, 2, 1, 2, 2, 1, 2, 1, 2, 1, 2, 1, 1, 2, 1, 2, 1, 2, 1, 2, 2, 1, 2, 1, 2, + 1, 2, 1, 1, 2, 1, 2, 1, 2, 1, 2, 2, 1, 2, 1, 2, 1, 2, 1, 1, 2, 1, 2, 1, 2, 1, 2, 2, 1, + 2, 1, 2, 1, 2, 1, + ]; + assert_eq!( + dataset.get(0).unwrap().annotation, + Annotation::SegmentationMask(SegmentationMask { + mask: TEST_CHECKERBOARD_MASK_PATTERN + .iter() + .map(|&x| x as usize) + .collect() + }) + ); + // random 2 colors mask + const TEST_RANDOM2COLORS_MASK_PATTERN: [u8; 64] = [ + 1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1, 2, 2, 2, 1, 2, 1, 2, 2, 2, 2, 2, 2, 2, + 2, 1, 1, 2, 2, 2, 1, 2, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 1, 2, 2, 1, 2, 1, 2, 1, 2, 2, 1, + 1, 1, 1, 1, 1, 1, + ]; + assert_eq!( + dataset.get(1).unwrap().annotation, + Annotation::SegmentationMask(SegmentationMask { + mask: TEST_RANDOM2COLORS_MASK_PATTERN + .iter() + .map(|&x| x as usize) + .collect() + }) + ); + // random 3 colors mask + const TEST_RANDOM3COLORS_MASK_PATTERN: [u8; 64] = [ + 3, 1, 3, 3, 1, 1, 3, 2, 3, 3, 3, 3, 1, 3, 2, 1, 2, 2, 2, 2, 1, 1, 2, 2, 1, 1, 1, 3, 3, + 3, 2, 3, 2, 2, 3, 2, 3, 3, 1, 3, 1, 3, 3, 1, 1, 3, 2, 1, 2, 2, 2, 1, 2, 1, 2, 3, 3, 1, + 3, 3, 2, 1, 2, 2, + ]; + assert_eq!( + dataset.get(2).unwrap().annotation, + Annotation::SegmentationMask(SegmentationMask { + mask: TEST_RANDOM3COLORS_MASK_PATTERN + .iter() + .map(|&x| x as usize) + .collect() + }) + ); + } + + #[test] + pub fn coco_detection_dataset() { + let dataset = ImageFolderDataset::new_coco_detection(COCO_JSON, COCO_IMAGES).unwrap(); + assert_eq!(dataset.len(), 3); // we have only three images defined + assert_eq!(dataset.get(3), None); + + const TWO_DOTS_AND_TRIANGLE_B1: BoundingBox = BoundingBox { + coords: [3.125_172, 18.090_784, 10.960_11, 10.740_027], + label: 0, + }; + + const TWO_DOTS_AND_TRIANGLE_B2: BoundingBox = BoundingBox { + coords: [3.257_221_5, 3.037_139, 10.563_961, 10.828_06], + label: 0, + }; + + const TWO_DOTS_AND_TRIANGLE_B3: BoundingBox = BoundingBox { + coords: [15.097_662, 3.389_271, 12.632_737, 11.180_193], + label: 1, + }; + + const DOTS_TRIANGLE_B1: BoundingBox = BoundingBox { + coords: [3.125_172, 17.914_719, 10.828_06, 11.004_127], + label: 0, + }; + + const DOTS_TRIANGLE_B2: BoundingBox = BoundingBox { + coords: [15.273_727, 3.301_238, 12.192_573, 11.708_39], + label: 1, + }; + + const ONE_DOT_B1: BoundingBox = BoundingBox { + coords: [10.079_78, 9.595_598, 10.960_11, 11.356_258], + label: 0, + }; + + for item in dataset.iter() { + let file_name = Path::new(&item.image_path).file_name().unwrap(); + match item.annotation { + // check if the number of bounding boxes is correct + Annotation::BoundingBoxes(v) => { + if file_name == "two_dots_and_triangle.jpg" { + assert_eq!(v.len(), 3); + assert!(v.contains(&TWO_DOTS_AND_TRIANGLE_B1)); + assert!(v.contains(&TWO_DOTS_AND_TRIANGLE_B2)); + assert!(v.contains(&TWO_DOTS_AND_TRIANGLE_B3)); + } else if file_name == "dot_triangle.jpg" { + assert_eq!(v.len(), 2); + assert!(v.contains(&DOTS_TRIANGLE_B1)); + assert!(v.contains(&DOTS_TRIANGLE_B2)); + } else if file_name == "one_dot.jpg" { + assert_eq!(v.len(), 1); + assert!(v.contains(&ONE_DOT_B1)); + } else { + panic!("{}", format!("unexpected image name: {}", item.image_path)); + } + } + _ => panic!("unexpected annotation"), + } + } + } +} diff --git a/crates/burn-dataset/src/vision/mnist.rs b/crates/burn-dataset/src/vision/mnist.rs new file mode 100644 index 0000000..909c779 --- /dev/null +++ b/crates/burn-dataset/src/vision/mnist.rs @@ -0,0 +1,221 @@ +use std::fs::{File, create_dir_all}; +use std::io::{Read, Seek, SeekFrom}; +use std::path::{Path, PathBuf}; + +use flate2::read::GzDecoder; +use serde::{Deserialize, Serialize}; + +use crate::{ + Dataset, InMemDataset, + transform::{Mapper, MapperDataset}, +}; + +use crate::network::downloader::download_file_as_bytes; + +// CVDF mirror of http://yann.lecun.com/exdb/mnist/ +const URL: &str = "https://storage.googleapis.com/cvdf-datasets/mnist/"; +const TRAIN_IMAGES: &str = "train-images-idx3-ubyte"; +const TRAIN_LABELS: &str = "train-labels-idx1-ubyte"; +const TEST_IMAGES: &str = "t10k-images-idx3-ubyte"; +const TEST_LABELS: &str = "t10k-labels-idx1-ubyte"; + +const WIDTH: usize = 28; +const HEIGHT: usize = 28; + +/// MNIST item. +#[derive(Deserialize, Serialize, Debug, Clone)] +pub struct MnistItem { + /// Image as a 2D array of floats. + pub image: [[f32; WIDTH]; HEIGHT], + + /// Label of the image. + pub label: u8, +} + +#[derive(Deserialize, Debug, Clone)] +struct MnistItemRaw { + pub image_bytes: Vec, + pub label: u8, +} + +struct BytesToImage; + +impl Mapper for BytesToImage { + /// Convert a raw MNIST item (image bytes) to a MNIST item (2D array image). + fn map(&self, item: &MnistItemRaw) -> MnistItem { + // Ensure the image dimensions are correct. + debug_assert_eq!(item.image_bytes.len(), WIDTH * HEIGHT); + + // Convert the image to a 2D array of floats. + let mut image_array = [[0f32; WIDTH]; HEIGHT]; + for (i, pixel) in item.image_bytes.iter().enumerate() { + let x = i % WIDTH; + let y = i / HEIGHT; + image_array[y][x] = *pixel as f32; + } + + MnistItem { + image: image_array, + label: item.label, + } + } +} + +type MappedDataset = MapperDataset, BytesToImage, MnistItemRaw>; + +/// The MNIST dataset consists of 70,000 28x28 black-and-white images in 10 classes (one for each digits), with 7,000 +/// images per class. There are 60,000 training images and 10,000 test images. +/// +/// The data is downloaded from the web from the [CVDF mirror](https://github.com/cvdfoundation/mnist). +pub struct MnistDataset { + dataset: MappedDataset, +} + +impl Dataset for MnistDataset { + fn get(&self, index: usize) -> Option { + self.dataset.get(index) + } + + fn len(&self) -> usize { + self.dataset.len() + } +} + +impl MnistDataset { + /// Creates a new train dataset. + pub fn train() -> Self { + Self::new("train") + } + + /// Creates a new test dataset. + pub fn test() -> Self { + Self::new("test") + } + + fn new(split: &str) -> Self { + // Download dataset + let root = MnistDataset::download(split); + + // MNIST is tiny so we can load it in-memory + // Train images (u8): 28 * 28 * 60000 = 47.04Mb + // Test images (u8): 28 * 28 * 10000 = 7.84Mb + let images = MnistDataset::read_images(&root, split); + let labels = MnistDataset::read_labels(&root, split); + + // Collect as vector of MnistItemRaw + let items: Vec<_> = images + .into_iter() + .zip(labels) + .map(|(image_bytes, label)| MnistItemRaw { image_bytes, label }) + .collect(); + + let dataset = InMemDataset::new(items); + let dataset = MapperDataset::new(dataset, BytesToImage); + + Self { dataset } + } + + /// Download the MNIST dataset files from the web. + /// Panics if the download cannot be completed or the content of the file cannot be written to disk. + fn download(split: &str) -> PathBuf { + // Dataset files are stored in the burn-dataset cache directory + let cache_dir = dirs::cache_dir() + .expect("Could not get cache directory") + .join("burn-dataset"); + let split_dir = cache_dir.join("mnist").join(split); + + if !split_dir.exists() { + create_dir_all(&split_dir).expect("Failed to create base directory"); + } + + // Download split files + match split { + "train" => { + MnistDataset::download_file(TRAIN_IMAGES, &split_dir); + MnistDataset::download_file(TRAIN_LABELS, &split_dir); + } + "test" => { + MnistDataset::download_file(TEST_IMAGES, &split_dir); + MnistDataset::download_file(TEST_LABELS, &split_dir); + } + _ => panic!("Invalid split specified {split}"), + }; + + split_dir + } + + /// Download a file from the MNIST dataset URL to the destination directory. + /// File download progress is reported with the help of a [progress bar](indicatif). + fn download_file>(name: &str, dest_dir: &P) -> PathBuf { + // Output file name + let file_name = dest_dir.as_ref().join(name); + + if !file_name.exists() { + // Download gzip file + let bytes = download_file_as_bytes(&format!("{URL}{name}.gz"), name); + + // Create file to write the downloaded content to + let mut output_file = File::create(&file_name).unwrap(); + + // Decode gzip file content and write to disk + let mut gz_buffer = GzDecoder::new(&bytes[..]); + std::io::copy(&mut gz_buffer, &mut output_file).unwrap(); + } + + file_name + } + + /// Read images at the provided path for the specified split. + /// Each image is a vector of bytes. + fn read_images>(root: &P, split: &str) -> Vec> { + let file_name = if split == "train" { + TRAIN_IMAGES + } else { + TEST_IMAGES + }; + let file_name = root.as_ref().join(file_name); + + // Read number of images from 16-byte header metadata + let mut f = File::open(file_name).unwrap(); + let mut buf = [0u8; 4]; + let _ = f.seek(SeekFrom::Start(4)).unwrap(); + f.read_exact(&mut buf) + .expect("Should be able to read image file header"); + let size = u32::from_be_bytes(buf); + + let mut buf_images: Vec = vec![0u8; WIDTH * HEIGHT * (size as usize)]; + let _ = f.seek(SeekFrom::Start(16)).unwrap(); + f.read_exact(&mut buf_images) + .expect("Should be able to read image file header"); + + buf_images + .chunks(WIDTH * HEIGHT) + .map(|chunk| chunk.to_vec()) + .collect() + } + + /// Read labels at the provided path for the specified split. + fn read_labels>(root: &P, split: &str) -> Vec { + let file_name = if split == "train" { + TRAIN_LABELS + } else { + TEST_LABELS + }; + let file_name = root.as_ref().join(file_name); + + // Read number of labels from 8-byte header metadata + let mut f = File::open(file_name).unwrap(); + let mut buf = [0u8; 4]; + let _ = f.seek(SeekFrom::Start(4)).unwrap(); + f.read_exact(&mut buf) + .expect("Should be able to read label file header"); + let size = u32::from_be_bytes(buf); + + let mut buf_labels: Vec = vec![0u8; size as usize]; + let _ = f.seek(SeekFrom::Start(8)).unwrap(); + f.read_exact(&mut buf_labels) + .expect("Should be able to read labels from file"); + + buf_labels + } +} diff --git a/crates/burn-dataset/src/vision/mod.rs b/crates/burn-dataset/src/vision/mod.rs new file mode 100644 index 0000000..948eb0b --- /dev/null +++ b/crates/burn-dataset/src/vision/mod.rs @@ -0,0 +1,9 @@ +#[cfg(feature = "builtin-sources")] +mod cifar; +mod image_folder; +mod mnist; + +#[cfg(feature = "builtin-sources")] +pub use cifar::*; +pub use image_folder::*; +pub use mnist::*; diff --git a/crates/burn-dataset/tests/data/dataset-fmt.csv b/crates/burn-dataset/tests/data/dataset-fmt.csv new file mode 100644 index 0000000..74a155b --- /dev/null +++ b/crates/burn-dataset/tests/data/dataset-fmt.csv @@ -0,0 +1,2 @@ +HI1 1 true 1.0 +HI2 1 false 1.0 diff --git a/crates/burn-dataset/tests/data/dataset.csv b/crates/burn-dataset/tests/data/dataset.csv new file mode 100644 index 0000000..2ec5fe8 --- /dev/null +++ b/crates/burn-dataset/tests/data/dataset.csv @@ -0,0 +1,3 @@ +column_str,column_int,column_bool,column_float +HI1,1,true,1.0 +HI2,1,false,1.0 diff --git a/crates/burn-dataset/tests/data/dataset.json b/crates/burn-dataset/tests/data/dataset.json new file mode 100644 index 0000000..a04eaa2 --- /dev/null +++ b/crates/burn-dataset/tests/data/dataset.json @@ -0,0 +1,2 @@ +{"column_str":"HI1","column_bytes":[1,2,3,3],"column_int":1,"column_bool":true,"column_float":1.0} +{"column_str":"HI2","column_bytes":[1,2,3,3],"column_int":1,"column_bool":false,"column_float":1.0} \ No newline at end of file diff --git a/crates/burn-dataset/tests/data/dataset_coco.json b/crates/burn-dataset/tests/data/dataset_coco.json new file mode 100644 index 0000000..6a75bf9 --- /dev/null +++ b/crates/burn-dataset/tests/data/dataset_coco.json @@ -0,0 +1,132 @@ +{ + "images": [ + { + "width": 32, + "height": 32, + "id": 0, + "file_name": "two_dots_and_triangle.jpg" + }, + { + "width": 32, + "height": 32, + "id": 1, + "file_name": "dot_triangle.jpg" + }, + { + "width": 32, + "height": 32, + "id": 2, + "file_name": "one_dot.jpg" + } + ], + "categories": [ + { + "id": 0, + "name": "dot" + }, + { + "id": 1, + "name": "triangle" + } + ], + "annotations": [ + { + "id": 0, + "image_id": 0, + "category_id": 0, + "segmentation": [], + "bbox": [ + 3.1251719394773056, + 18.0907840440165, + 10.96011004126548, + 10.740027510316379 + ], + "ignore": 0, + "iscrowd": 0, + "area": 117.71188335928603 + }, + { + "id": 1, + "image_id": 0, + "category_id": 0, + "segmentation": [], + "bbox": [ + 3.2572214580467658, + 3.0371389270976605, + 10.563961485557085, + 10.828060522696012 + ], + "ignore": 0, + "iscrowd": 0, + "area": 114.38721432504178 + }, + { + "id": 2, + "image_id": 0, + "category_id": 1, + "segmentation": [], + "bbox": [ + 15.097661623108666, + 3.3892709766162312, + 12.632737276478679, + 11.18019257221458 + ], + "ignore": 0, + "iscrowd": 0, + "area": 141.23643546522516 + }, + { + "id": 3, + "image_id": 1, + "category_id": 0, + "segmentation": [], + "bbox": [ + 3.125171939477304, + 17.914718019257222, + 10.82806052269601, + 11.004126547455297 + ], + "ignore": 0, + "iscrowd": 0, + "area": 119.15334825525184 + }, + { + "id": 4, + "image_id": 1, + "category_id": 1, + "segmentation": [], + "bbox": [ + 15.27372764786794, + 3.301237964236589, + 12.192572214580478, + 11.708390646492433 + ], + "ignore": 0, + "iscrowd": 0, + "area": 142.7553984738776 + }, + { + "id": 5, + "image_id": 2, + "category_id": 0, + "segmentation": [], + "bbox": [ + 10.07977991746905, + 9.59559834938102, + 10.960110041265464, + 11.356258596973863 + ], + "ignore": 0, + "iscrowd": 0, + "area": 124.46584387990049 + } + ], + "info": { + "year": 2024, + "version": "1.0", + "description": "", + "contributor": "", + "url": "", + "date_created": "2024-12-11 22:16:31.823494" + } +} diff --git a/crates/burn-dataset/tests/data/image_folder/orange/dot.jpg b/crates/burn-dataset/tests/data/image_folder/orange/dot.jpg new file mode 100755 index 0000000..3629f7f Binary files /dev/null and b/crates/burn-dataset/tests/data/image_folder/orange/dot.jpg differ diff --git a/crates/burn-dataset/tests/data/image_folder/red/dot.jpg b/crates/burn-dataset/tests/data/image_folder/red/dot.jpg new file mode 100755 index 0000000..be19ad3 Binary files /dev/null and b/crates/burn-dataset/tests/data/image_folder/red/dot.jpg differ diff --git a/crates/burn-dataset/tests/data/image_folder/red/dot.png b/crates/burn-dataset/tests/data/image_folder/red/dot.png new file mode 100755 index 0000000..93d9dd1 Binary files /dev/null and b/crates/burn-dataset/tests/data/image_folder/red/dot.png differ diff --git a/crates/burn-dataset/tests/data/image_folder_coco/dot_triangle.jpg b/crates/burn-dataset/tests/data/image_folder_coco/dot_triangle.jpg new file mode 100644 index 0000000..572ab2e Binary files /dev/null and b/crates/burn-dataset/tests/data/image_folder_coco/dot_triangle.jpg differ diff --git a/crates/burn-dataset/tests/data/image_folder_coco/one_dot.jpg b/crates/burn-dataset/tests/data/image_folder_coco/one_dot.jpg new file mode 100644 index 0000000..b719f53 Binary files /dev/null and b/crates/burn-dataset/tests/data/image_folder_coco/one_dot.jpg differ diff --git a/crates/burn-dataset/tests/data/image_folder_coco/two_dots_and_triangle.jpg b/crates/burn-dataset/tests/data/image_folder_coco/two_dots_and_triangle.jpg new file mode 100644 index 0000000..fe52aa6 Binary files /dev/null and b/crates/burn-dataset/tests/data/image_folder_coco/two_dots_and_triangle.jpg differ diff --git a/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_checkerboard.png b/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_checkerboard.png new file mode 100644 index 0000000..3c87252 Binary files /dev/null and b/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_checkerboard.png differ diff --git a/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_checkerboard.txt b/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_checkerboard.txt new file mode 100644 index 0000000..2e01635 --- /dev/null +++ b/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_checkerboard.txt @@ -0,0 +1,8 @@ +1 2 1 2 1 2 1 2 +2 1 2 1 2 1 2 1 +1 2 1 2 1 2 1 2 +2 1 2 1 2 1 2 1 +1 2 1 2 1 2 1 2 +2 1 2 1 2 1 2 1 +1 2 1 2 1 2 1 2 +2 1 2 1 2 1 2 1 diff --git a/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_random_2colors.png b/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_random_2colors.png new file mode 100644 index 0000000..f0a129a Binary files /dev/null and b/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_random_2colors.png differ diff --git a/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_random_2colors.txt b/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_random_2colors.txt new file mode 100644 index 0000000..4fa2b7c --- /dev/null +++ b/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_random_2colors.txt @@ -0,0 +1,8 @@ +1 2 1 1 1 2 1 1 +1 2 1 1 1 1 2 1 +2 2 2 1 2 1 2 2 +2 2 2 2 2 2 1 1 +2 2 2 1 2 1 1 1 +1 1 2 2 2 2 2 1 +2 2 1 2 1 2 1 2 +2 1 1 1 1 1 1 1 diff --git a/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_random_3colors.png b/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_random_3colors.png new file mode 100644 index 0000000..38984e7 Binary files /dev/null and b/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_random_3colors.png differ diff --git a/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_random_3colors.txt b/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_random_3colors.txt new file mode 100644 index 0000000..08f222b --- /dev/null +++ b/crates/burn-dataset/tests/data/segmask_folder/annotations/mask_random_3colors.txt @@ -0,0 +1,8 @@ +3 1 3 3 1 1 3 2 +3 3 3 3 1 3 2 1 +2 2 2 2 1 1 2 2 +1 1 1 3 3 3 2 3 +2 2 3 2 3 3 1 3 +1 3 3 1 1 3 2 1 +2 2 2 1 2 1 2 3 +3 1 3 3 2 1 2 2 diff --git a/crates/burn-dataset/tests/data/segmask_folder/images/image_checkerboard.png b/crates/burn-dataset/tests/data/segmask_folder/images/image_checkerboard.png new file mode 100644 index 0000000..2087501 Binary files /dev/null and b/crates/burn-dataset/tests/data/segmask_folder/images/image_checkerboard.png differ diff --git a/crates/burn-dataset/tests/data/segmask_folder/images/image_random_2colors.png b/crates/burn-dataset/tests/data/segmask_folder/images/image_random_2colors.png new file mode 100644 index 0000000..f433fd9 Binary files /dev/null and b/crates/burn-dataset/tests/data/segmask_folder/images/image_random_2colors.png differ diff --git a/crates/burn-dataset/tests/data/segmask_folder/images/image_random_3colors.png b/crates/burn-dataset/tests/data/segmask_folder/images/image_random_3colors.png new file mode 100644 index 0000000..880bac4 Binary files /dev/null and b/crates/burn-dataset/tests/data/segmask_folder/images/image_random_3colors.png differ diff --git a/crates/burn-dataset/tests/data/sqlite-dataset.db b/crates/burn-dataset/tests/data/sqlite-dataset.db new file mode 100644 index 0000000..a53bf20 Binary files /dev/null and b/crates/burn-dataset/tests/data/sqlite-dataset.db differ diff --git a/crates/burn-dataset/tests/data/text_folder/negative/sample1.txt b/crates/burn-dataset/tests/data/text_folder/negative/sample1.txt new file mode 100644 index 0000000..7e3d1f2 --- /dev/null +++ b/crates/burn-dataset/tests/data/text_folder/negative/sample1.txt @@ -0,0 +1 @@ +This is a negative text sample for testing the text folder dataset functionality. \ No newline at end of file diff --git a/crates/burn-dataset/tests/data/text_folder/negative/sample2.txt b/crates/burn-dataset/tests/data/text_folder/negative/sample2.txt new file mode 100644 index 0000000..a4abb75 --- /dev/null +++ b/crates/burn-dataset/tests/data/text_folder/negative/sample2.txt @@ -0,0 +1 @@ +另一个负面文本样本，用以确保数据集能够处理同一类别中的多个文件。 \ No newline at end of file diff --git a/crates/burn-dataset/tests/data/text_folder/positive/sample1.txt b/crates/burn-dataset/tests/data/text_folder/positive/sample1.txt new file mode 100644 index 0000000..6e18f0b --- /dev/null +++ b/crates/burn-dataset/tests/data/text_folder/positive/sample1.txt @@ -0,0 +1 @@ +This is a positive text sample for testing the text folder dataset functionality. \ No newline at end of file diff --git a/crates/burn-dataset/tests/data/text_folder/positive/sample2.txt b/crates/burn-dataset/tests/data/text_folder/positive/sample2.txt new file mode 100644 index 0000000..d025616 --- /dev/null +++ b/crates/burn-dataset/tests/data/text_folder/positive/sample2.txt @@ -0,0 +1 @@ +另一个正面文本样本，以确保数据集能够处理同一类别中的多个文件。 \ No newline at end of file diff --git a/crates/burn-derive/Cargo.toml b/crates/burn-derive/Cargo.toml new file mode 100644 index 0000000..f20d207 --- /dev/null +++ b/crates/burn-derive/Cargo.toml @@ -0,0 +1,23 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science"] +description = "Derive crate for the Burn framework" +edition.workspace = true +keywords = [] +license.workspace = true +name = "burn-derive" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-derive" +version.workspace = true + +[lints] +workspace = true + +[lib] +proc-macro = true + +[dependencies] +proc-macro2 = { workspace = true } +quote = { workspace = true } +syn = { workspace = true, features = ["visit"] } +derive-new = { workspace = true } diff --git a/crates/burn-derive/LICENSE-APACHE b/crates/burn-derive/LICENSE-APACHE new file mode 120000 index 0000000..1cd601d --- /dev/null +++ b/crates/burn-derive/LICENSE-APACHE @@ -0,0 +1 @@ +../../LICENSE-APACHE \ No newline at end of file diff --git a/crates/burn-derive/LICENSE-MIT b/crates/burn-derive/LICENSE-MIT new file mode 120000 index 0000000..b2cfbdc --- /dev/null +++ b/crates/burn-derive/LICENSE-MIT @@ -0,0 +1 @@ +../../LICENSE-MIT \ No newline at end of file diff --git a/crates/burn-derive/README.md b/crates/burn-derive/README.md new file mode 100644 index 0000000..bc34709 --- /dev/null +++ b/crates/burn-derive/README.md @@ -0,0 +1,6 @@ +# Burn Derive + +This crate should only be used with [burn](https://github.com/tracel-ai/burn). + +[![Current Crates.io Version](https://img.shields.io/crates/v/burn-derive.svg)](https://crates.io/crates/burn-derive) +[![license](https://shields.io/badge/license-MIT%2FApache--2.0-blue)](https://github.com/tracel-ai/burn-derive/blob/master/README.md) diff --git a/crates/burn-derive/src/config/analyzer.rs b/crates/burn-derive/src/config/analyzer.rs new file mode 100644 index 0000000..e5e6285 --- /dev/null +++ b/crates/burn-derive/src/config/analyzer.rs @@ -0,0 +1,87 @@ +use super::ConfigEnumAnalyzer; +use crate::config::ConfigStructAnalyzer; +use crate::shared::{attribute::AttributeItem, field::FieldTypeAnalyzer}; +use proc_macro2::TokenStream; +use quote::quote; +use syn::{Field, Ident}; + +pub struct ConfigAnalyzerFactory {} + +pub trait ConfigAnalyzer { + fn gen_new_fn(&self) -> TokenStream { + quote! {} + } + fn gen_builder_fns(&self) -> TokenStream { + quote! {} + } + fn gen_serde_impl(&self) -> TokenStream; + fn gen_clone_impl(&self) -> TokenStream; + fn gen_display_impl(&self) -> TokenStream; + fn gen_config_impl(&self) -> TokenStream; +} + +impl ConfigAnalyzerFactory { + pub fn new() -> Self { + Self {} + } + + pub fn create_analyzer(&self, item: &syn::DeriveInput) -> Box { + let name = item.ident.clone(); + let config_type = parse_asm(item); + + match config_type { + ConfigType::Struct(data) => Box::new(self.create_struct_analyzer(name, data)), + ConfigType::Enum(data) => Box::new(self.create_enum_analyzer(name, data)), + } + } + + fn create_struct_analyzer(&self, name: Ident, fields: Vec) -> ConfigStructAnalyzer { + let fields = fields.into_iter().map(FieldTypeAnalyzer::new); + + let mut fields_required = Vec::new(); + let mut fields_option = Vec::new(); + let mut fields_default = Vec::new(); + + for field in fields { + let attributes: Vec = field + .attributes() + .filter(|attr| attr.has_name("config")) + .map(|attr| attr.item()) + .collect(); + + if !attributes.is_empty() { + let item = attributes.first().unwrap().clone(); + fields_default.push((field.clone(), item)); + continue; + } + + if field.is_of_type(&["Option"]) { + fields_option.push(field.clone()); + continue; + } + + fields_required.push(field.clone()); + } + + ConfigStructAnalyzer::new(name, fields_required, fields_option, fields_default) + } + + fn create_enum_analyzer(&self, name: Ident, data: syn::DataEnum) -> ConfigEnumAnalyzer { + ConfigEnumAnalyzer::new(name, data) + } +} + +enum ConfigType { + Struct(Vec), + Enum(syn::DataEnum), +} + +fn parse_asm(ast: &syn::DeriveInput) -> ConfigType { + match &ast.data { + syn::Data::Struct(struct_data) => { + ConfigType::Struct(struct_data.fields.clone().into_iter().collect()) + } + syn::Data::Enum(enum_data) => ConfigType::Enum(enum_data.clone()), + syn::Data::Union(_) => panic!("Only struct and enum can be derived"), + } +} diff --git a/crates/burn-derive/src/config/analyzer_enum.rs b/crates/burn-derive/src/config/analyzer_enum.rs new file mode 100644 index 0000000..84ddd9b --- /dev/null +++ b/crates/burn-derive/src/config/analyzer_enum.rs @@ -0,0 +1,141 @@ +use crate::shared::enum_variant::map_enum_variant; + +use super::ConfigAnalyzer; +use proc_macro2::{Ident, TokenStream}; +use quote::quote; + +pub struct ConfigEnumAnalyzer { + name: Ident, + data: syn::DataEnum, +} + +impl ConfigEnumAnalyzer { + pub fn new(name: Ident, data: syn::DataEnum) -> Self { + Self { name, data } + } + + fn serde_enum_ident(&self) -> Ident { + Ident::new(&format!("{}Serde", self.name), self.name.span()) + } + + fn gen_serde_enum(&self) -> TokenStream { + let enum_name = self.serde_enum_ident(); + let data = &self.data.variants; + + quote! { + #[derive(burn::serde::Serialize, burn::serde::Deserialize)] + #[serde(crate = "burn::serde")] + enum #enum_name { + #data + } + + } + } + + fn gen_serialize_fn(&self) -> TokenStream { + let enum_name = self.serde_enum_ident(); + let variants = self.data.variants.iter().map(|variant| { + let variant_name = &variant.ident; + let (inputs, outputs) = map_enum_variant(variant, |ident| quote! { #ident.clone() }); + + quote! { Self::#variant_name #inputs => #enum_name::#variant_name #outputs } + }); + + let name = &self.name; + + quote! { + impl burn::serde::Serialize for #name { + fn serialize(&self, serializer: S) -> Result + where + S: burn::serde::Serializer { + let serde_state = match self { + #(#variants),* + }; + serde_state.serialize(serializer) + } + } + + } + } + + fn gen_deserialize_fn(&self) -> TokenStream { + let enum_name = self.serde_enum_ident(); + let variants = self.data.variants.iter().map(|variant| { + let variant_name = &variant.ident; + let (inputs, outputs) = map_enum_variant(variant, |ident| quote! { #ident.clone() }); + + quote! { #enum_name::#variant_name #inputs => Self::#variant_name #outputs } + }); + let name = &self.name; + + quote! { + impl<'de> burn::serde::Deserialize<'de> for #name { + fn deserialize(deserializer: D) -> Result + where + D: burn::serde::Deserializer<'de> { + let serde_state = #enum_name::deserialize(deserializer)?; + Ok(match serde_state { + #(#variants),* + }) + } + } + + } + } +} + +impl ConfigAnalyzer for ConfigEnumAnalyzer { + fn gen_serde_impl(&self) -> TokenStream { + let struct_gen = self.gen_serde_enum(); + let serialize_gen = self.gen_serialize_fn(); + let deserialize_gen = self.gen_deserialize_fn(); + + quote! { + #struct_gen + #serialize_gen + #deserialize_gen + } + } + + fn gen_clone_impl(&self) -> TokenStream { + let variants = self.data.variants.iter().map(|variant| { + let variant_name = &variant.ident; + let (inputs, outputs) = map_enum_variant(variant, |ident| quote! { #ident.clone() }); + + quote! { Self::#variant_name #inputs => Self::#variant_name #outputs } + }); + let name = &self.name; + + quote! { + impl Clone for #name { + fn clone(&self) -> Self { + match self { + #(#variants),* + } + } + } + + } + } + + fn gen_display_impl(&self) -> TokenStream { + let name = &self.name; + + quote! { + impl core::fmt::Display for #name { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_str(&burn::config::config_to_json(self)) + } + } + } + } + + fn gen_config_impl(&self) -> TokenStream { + let name = &self.name; + + quote! { + impl burn::config::Config for #name { + } + } + } +} diff --git a/crates/burn-derive/src/config/analyzer_struct.rs b/crates/burn-derive/src/config/analyzer_struct.rs new file mode 100644 index 0000000..d347f71 --- /dev/null +++ b/crates/burn-derive/src/config/analyzer_struct.rs @@ -0,0 +1,380 @@ +use super::ConfigAnalyzer; +use crate::shared::{attribute::AttributeItem, field::FieldTypeAnalyzer}; +use proc_macro2::{Ident, TokenStream}; +use quote::quote; + +pub struct ConfigStructAnalyzer { + name: Ident, + fields_required: Vec, + fields_option: Vec, + fields_default: Vec<(FieldTypeAnalyzer, AttributeItem)>, +} + +impl ConfigStructAnalyzer { + pub fn new( + name: Ident, + fields_required: Vec, + fields_option: Vec, + fields_default: Vec<(FieldTypeAnalyzer, AttributeItem)>, + ) -> Self { + Self { + name, + fields_required, + fields_option, + fields_default, + } + } + + fn wrap_impl_block(&self, tokens: TokenStream) -> TokenStream { + let name = &self.name; + + quote! { + impl #name { + #tokens + } + } + } + + fn names(&self) -> Vec { + let mut names = Vec::new(); + + for field in self.fields_required.iter() { + names.push(field.clone()); + } + + for field in self.fields_option.iter() { + names.push(field.clone()); + } + + for (field, _) in self.fields_default.iter() { + names.push(field.clone()); + } + + names + } + + fn name_types(&self, names: &[FieldTypeAnalyzer]) -> Vec { + let mut name_types = Vec::new(); + + for field in names.iter() { + let name = field.ident(); + let ty = &field.field.ty; + + name_types.push(quote! { + #name: #ty + }); + } + + name_types + } + + fn serde_struct_ident(&self) -> Ident { + Ident::new(&format!("{}Serde", self.name), self.name.span()) + } + + fn gen_serialize_fn( + &self, + struct_name: &Ident, + struct_gen: &TokenStream, + names: &[FieldTypeAnalyzer], + ) -> TokenStream { + let name = &self.name; + let names = names.iter().map(|name| { + let name = name.ident(); + quote! { #name: self.#name.clone() } + }); + + quote! { + impl burn::serde::Serialize for #name { + + fn serialize(&self, serializer: S) -> Result + where + S: burn::serde::Serializer { + #[derive(burn::serde::Serialize)] + #[serde(crate = "burn::serde")] + #struct_gen + + let serde_state = #struct_name { + #(#names),* + }; + serde_state.serialize(serializer) + } + } + + } + } + + fn gen_deserialize_fn( + &self, + struct_name: &Ident, + struct_gen: &TokenStream, + names: &[FieldTypeAnalyzer], + ) -> TokenStream { + let name = &self.name; + let names = names.iter().map(|name| { + let name = name.ident(); + quote! { #name: serde_state.#name } + }); + + quote! { + impl<'de> burn::serde::Deserialize<'de> for #name { + fn deserialize(deserializer: D) -> Result + where + D: burn::serde::Deserializer<'de> { + #[derive(burn::serde::Deserialize)] + #[serde(crate = "burn::serde")] + #struct_gen + + let serde_state = #struct_name::deserialize(deserializer)?; + Ok(#name { + #(#names),* + }) + } + } + + } + } + + fn gen_serde_struct(&self, names: &[TokenStream]) -> TokenStream { + let struct_name = self.serde_struct_ident(); + + quote! { + struct #struct_name { + #(#names),* + } + + } + } +} + +impl ConfigAnalyzer for ConfigStructAnalyzer { + fn gen_new_fn(&self) -> TokenStream { + let mut body = quote! {}; + let mut args = Vec::new(); + + let mut fn_docs = quote! {}; + let mut has_field_docs = false; + let mut has_required_docs = false; + let mut has_option_docs = false; + let mut has_default_docs = false; + let mut docs_header = |fn_docs: &mut TokenStream, + required_docs: bool, + option_docs: bool, + default_docs: bool| { + if !has_field_docs { + has_field_docs = true; + fn_docs.extend(quote! { + #[doc = "# Arguments"] + }); + } + if !has_required_docs && required_docs { + fn_docs.extend(quote! { + #[doc = "###### Required Arguments"] + }); + has_required_docs = true; + } + if !has_option_docs && option_docs { + fn_docs.extend(quote! { + #[doc = "###### Optional Arguments"] + }); + has_option_docs = true; + } + if !has_default_docs && default_docs { + fn_docs.extend(quote! { + #[doc = "###### Default Arguments"] + }); + has_default_docs = true; + } + }; + + for field in self.fields_required.iter() { + let name = field.ident(); + let ty = &field.field.ty; + let docs = field.docs(); + + body.extend(quote! { + #name: #name, + }); + args.push(quote! { + #name: #ty + }); + docs_header(&mut fn_docs, true, false, false); + let doc_str = format!("###### `{}`\n\n", quote!(#name)); + fn_docs.extend(quote! { + #[doc = #doc_str] + #(#docs)* + }); + } + + for field in self.fields_option.iter() { + let name = field.ident(); + let docs = field.docs(); + + body.extend(quote! { + #name: None, + }); + docs_header(&mut fn_docs, false, true, false); + let default_doc = "- Defaults to `None`"; + let doc_str = format!("###### `{}`\n", quote!(#name)); + fn_docs.extend(quote! { + #[doc = #doc_str] + #(#docs)* + #[doc = #default_doc] + }); + } + + for (field, attribute) in self.fields_default.iter() { + let name = field.ident(); + let value = &attribute.value; + let docs = field.docs(); + + match value { + syn::Lit::Str(value) => { + let stream: proc_macro2::TokenStream = value.value().parse().unwrap(); + + body.extend(quote! { + #name: #stream, + }); + } + _ => { + body.extend(quote! { + #name: #value, + }); + } + }; + docs_header(&mut fn_docs, false, false, true); + let default_doc = format!("- Defaults to `{}`", quote!(#value)); + let doc_str = format!("###### `{}`\n", quote!(#name)); + fn_docs.extend(quote! { + #[doc = #doc_str] + #(#docs)* + #[doc = #default_doc] + }); + } + + let body = quote! { + #[doc = "Create a new instance of the config."] + #fn_docs + #[allow(clippy::too_many_arguments)] + pub fn new( + #(#args),* + ) -> Self { + Self { #body } + } + }; + self.wrap_impl_block(body) + } + + fn gen_builder_fns(&self) -> TokenStream { + let mut body = quote! {}; + + for (field, attribute) in self.fields_default.iter() { + let name = field.ident(); + let ty = &field.field.ty; + let value = &attribute.value; + let docs = field.docs(); + let default_doc = format!("- Defaults to `{}`", quote!(#value)); + let doc_str = format!( + "Sets the value for the field [`{}`](Self::{0}).\n\n", + quote!(#name) + ); + let fn_docs = quote! { + #[doc = #doc_str] + #(#docs)* + #[doc = #default_doc] + }; + let fn_name = Ident::new(&format!("with_{name}"), name.span()); + + body.extend(quote! { + #fn_docs + pub fn #fn_name(mut self, #name: #ty) -> Self { + self.#name = #name; + self + } + }); + } + + for field in self.fields_option.iter() { + let name = field.ident(); + let ty = &field.field.ty; + let docs = field.docs(); + let default_doc = "- Defaults to `None`"; + let doc_str = format!( + "Sets the value for the field [`{}`](Self::{0}).\n\n", + quote!(#name) + ); + let fn_docs = quote! { + #[doc = #doc_str] + #(#docs)* + #[doc = #default_doc] + }; + let fn_name = Ident::new(&format!("with_{name}"), name.span()); + + body.extend(quote! { + #fn_docs + pub fn #fn_name(mut self, #name: #ty) -> Self { + self.#name = #name; + self + } + }); + } + + self.wrap_impl_block(body) + } + + fn gen_serde_impl(&self) -> TokenStream { + let names = self.names(); + + let struct_name = self.serde_struct_ident(); + let name_types = self.name_types(&names); + let struct_gen = self.gen_serde_struct(&name_types); + + let serialize_gen = self.gen_serialize_fn(&struct_name, &struct_gen, &names); + let deserialize_gen = self.gen_deserialize_fn(&struct_name, &struct_gen, &names); + + quote! { + #serialize_gen + #deserialize_gen + } + } + + fn gen_clone_impl(&self) -> TokenStream { + let name = &self.name; + let names = self.names().into_iter().map(|name| { + let name = name.ident(); + quote! { #name: self.#name.clone() } + }); + + quote! { + impl Clone for #name { + fn clone(&self) -> Self { + Self { + #(#names),* + } + } + } + + } + } + + fn gen_display_impl(&self) -> TokenStream { + let name = &self.name; + + quote! { + impl core::fmt::Display for #name { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_str(&burn::config::config_to_json(self)) + } + } + } + } + + fn gen_config_impl(&self) -> TokenStream { + let name = &self.name; + + quote! { + impl burn::config::Config for #name { + } + } + } +} diff --git a/crates/burn-derive/src/config/base.rs b/crates/burn-derive/src/config/base.rs new file mode 100644 index 0000000..cca3f04 --- /dev/null +++ b/crates/burn-derive/src/config/base.rs @@ -0,0 +1,24 @@ +use super::ConfigAnalyzerFactory; +use quote::quote; + +pub(crate) fn derive_impl(item: &syn::DeriveInput) -> proc_macro::TokenStream { + let factory = ConfigAnalyzerFactory::new(); + let analyzer = factory.create_analyzer(item); + + let constructor = analyzer.gen_new_fn(); + let builders = analyzer.gen_builder_fns(); + let serde = analyzer.gen_serde_impl(); + let clone = analyzer.gen_clone_impl(); + let display = analyzer.gen_display_impl(); + let config_impl = analyzer.gen_config_impl(); + + quote! { + #config_impl + #constructor + #builders + #serde + #clone + #display + } + .into() +} diff --git a/crates/burn-derive/src/config/mod.rs b/crates/burn-derive/src/config/mod.rs new file mode 100644 index 0000000..36a5c0e --- /dev/null +++ b/crates/burn-derive/src/config/mod.rs @@ -0,0 +1,9 @@ +mod analyzer; +mod analyzer_enum; +mod analyzer_struct; +mod base; + +pub(crate) use analyzer::*; +pub(crate) use analyzer_enum::*; +pub(crate) use analyzer_struct::*; +pub(crate) use base::*; diff --git a/crates/burn-derive/src/lib.rs b/crates/burn-derive/src/lib.rs new file mode 100644 index 0000000..8581732 --- /dev/null +++ b/crates/burn-derive/src/lib.rs @@ -0,0 +1,88 @@ +#![warn(missing_docs)] + +//! The derive crate of Burn. + +#[macro_use] +extern crate derive_new; + +use proc_macro::TokenStream; + +pub(crate) mod config; +pub(crate) mod module; +pub(crate) mod record_state; +pub(crate) mod shared; + +/// Derive macro for the `Module` trait. +/// +/// # Sub-modules +/// +/// By default, the macro automatically detects sub-modules and parameters as module types. +/// +/// Any field not recognized as a module type is assumed to be a non-module +/// and is skipped by the module system (not persistent, not visited). +/// +/// ## Generics +/// +/// Generic type parameters (e.g., `field: M`) are assumed to be sub-modules by default. +/// If a generic field represents some other runtime state or configuration, you can use +/// the `#[module(skip)]` attribute to provide a hint. +/// +/// # Field Attributes +/// +/// ## `#[module(skip)]` +/// +/// Explicitly marks a field to be ignored by the module derive. +/// +/// Skipped fields are not parameters, not modules, and are not persistent. +/// This is equivalent to the deprecated `Ignored` wrapper. +/// +/// ### Requirements +/// +/// The field must implement: `Debug + Clone + Send`. +/// +/// # Example +/// +/// ```ignore +/// #[derive(Module, Debug)] +/// pub struct MyModule { +/// /// A normal parameter. +/// weights: Param>, +/// /// A field configured at runtime. +/// dropout_prob: f64, +/// /// A field that is recomputed at runtime. +/// cached_mask: Option>, +/// /// A field that contains some debug state. +/// debug_state: String, +/// /// Treated as a module (default for generics). +/// inner: M, +/// /// Hint required: this generic is NOT a module. +/// #[module(skip)] +/// other: N, +/// } +/// ``` +#[proc_macro_derive(Module, attributes(module))] +pub fn module_derive(input: TokenStream) -> TokenStream { + let input = syn::parse(input).unwrap(); + module::derive_impl(&input) +} + +/// Derive macro for the config. +#[proc_macro_derive(Config, attributes(config))] +pub fn config_derive(input: TokenStream) -> TokenStream { + let item = syn::parse(input).unwrap(); + config::derive_impl(&item) +} + +/// Derive macro for a recordable state (optimizer or learning-rate scheduler), decomposing it +/// into named tensors and scalars for the burnpack format. +/// +/// Supported field shapes: `Tensor`, `Option>`, `Vec>`, scalars +/// (`usize`/`isize`/`u8`..`u64`/`i8`..`i64`/`f32`/`f64`/`bool`), `Option`, a nested +/// `RecordState`, and `Option`. Scalar fields must use a concrete primitive type, not a +/// type alias (e.g. `f64`, not `LearningRate`): classification is syntactic, so an alias is treated +/// as a nested state. +#[proc_macro_derive(RecordState)] +pub fn record_state_derive(input: TokenStream) -> TokenStream { + let input = syn::parse(input).unwrap(); + record_state::derive_impl(&input) +} diff --git a/crates/burn-derive/src/module/base.rs b/crates/burn-derive/src/module/base.rs new file mode 100644 index 0000000..e76daa8 --- /dev/null +++ b/crates/burn-derive/src/module/base.rs @@ -0,0 +1,35 @@ +use super::{ + codegen::{generate_module_standard, generate_module_stateless}, + codegen_enum::EnumModuleCodegen, + codegen_struct::StructModuleCodegen, +}; +use proc_macro::TokenStream; + +pub(crate) fn derive_impl(ast: &syn::DeriveInput) -> TokenStream { + match &ast.data { + syn::Data::Struct(_) => match StructModuleCodegen::from_ast(ast) { + Ok(struct_codegen) => { + if struct_codegen.has_module_fields() { + generate_module_standard(ast, struct_codegen) + } else { + generate_module_stateless(ast, struct_codegen) + } + } + Err(err) => err.to_compile_error(), + }, + syn::Data::Enum(_data) => match EnumModuleCodegen::from_ast(ast) { + Ok(enum_codegen) => { + if enum_codegen.has_module_fields() { + generate_module_standard(ast, enum_codegen) + } else { + generate_module_stateless(ast, enum_codegen) + } + } + Err(err) => err.to_compile_error(), + }, + syn::Data::Union(_) => { + syn::Error::new_spanned(ast, "Union modules aren't supported").to_compile_error() + } + } + .into() +} diff --git a/crates/burn-derive/src/module/codegen.rs b/crates/burn-derive/src/module/codegen.rs new file mode 100644 index 0000000..d0a9a30 --- /dev/null +++ b/crates/burn-derive/src/module/codegen.rs @@ -0,0 +1,217 @@ +use super::display; +use crate::{ + module::generics::{GenericKind, ModuleGenerics}, + shared::generics::GenericsHelper, +}; +use proc_macro2::TokenStream; +use quote::quote; +use syn::{Attribute, Generics, parse_quote}; + +/// Basic trait to be implemented for Module generation. +pub(crate) trait ModuleCodegen { + fn gen_num_params(&self) -> TokenStream; + fn gen_visit(&self) -> TokenStream; + fn gen_collect_devices(&self) -> TokenStream; + fn gen_to_device(&self) -> TokenStream; + fn gen_fork(&self) -> TokenStream; + fn gen_map(&self) -> TokenStream; + fn gen_valid(&self) -> TokenStream; + fn gen_from_inner(&self) -> TokenStream; + fn gen_clone(&self) -> TokenStream; + + fn gen_display(&self) -> TokenStream; + + fn module_generics(&self) -> &ModuleGenerics; +} + +pub(crate) fn generate_module_standard( + ast: &syn::DeriveInput, + codegen: Codegen, +) -> TokenStream { + let name = &ast.ident; + + let generics = GenericsParser::from_ast(&ast.generics, codegen.module_generics()); + + let display_fn = display::display_fn(ast); + let attributes_fn = codegen.gen_display(); + let num_params_fn = codegen.gen_num_params(); + let visit = codegen.gen_visit(); + let map_mut = codegen.gen_map(); + let collect_devices = codegen.gen_collect_devices(); + let to_device = codegen.gen_to_device(); + let fork = codegen.gen_fork(); + let valid_fn = codegen.gen_valid(); + let from_inner_fn = codegen.gen_from_inner(); + let clone_fn = codegen.gen_clone(); + + let (generics_module, generics_ty_module, generics_where_module) = + generics.module.split_for_impl(); + let (generics_module_autodiff, generics_ty_module_autodiff, generics_where_module_autodiff) = + generics.module_autodiff.split_for_impl(); + + let mut codegen = quote! { + + impl #generics_module burn::module::Module for #name #generics_ty_module #generics_where_module { + #num_params_fn + + #visit + #map_mut + + #collect_devices + #to_device + #fork + + } + + impl #generics_module_autodiff burn::module::AutodiffModule for #name #generics_ty_module_autodiff #generics_where_module_autodiff + { + #valid_fn + + #from_inner_fn + } + + impl #generics_module core::fmt::Display for #name #generics_ty_module #generics_where_module { + #display_fn + } + + impl #generics_module burn::module::ModuleDisplayDefault for #name #generics_ty_module #generics_where_module { + #attributes_fn + + fn num_params(&self) -> usize { + burn::module::Module::num_params(self) + } + } + + impl #generics_module Clone for #name #generics_ty_module #generics_where_module { + #clone_fn + } + }; + + if !has_custom_display(&ast.attrs) { + codegen.extend(quote! { + impl #generics_module burn::module::ModuleDisplay for #name #generics_ty_module #generics_where_module { + + } + }); + } + + codegen +} + +// When there is inner param or module, the type is considered stateless. +pub(crate) fn generate_module_stateless( + ast: &syn::DeriveInput, + codegen: Codegen, // Pass the codegen here +) -> TokenStream { + let name = &ast.ident; + let (generics, generics_ty, generics_where) = ast.generics.split_for_impl(); + + let display_fn = display::display_fn(ast); + let attributes_fn = codegen.gen_display(); // Use codegen for attributes too + let clone_fn = codegen.gen_clone(); // The automatic clone logic + + let mut codegen = quote! { + impl #generics burn::module::Module for #name #generics_ty #generics_where { + burn::empty!(module); + } + + impl #generics burn::module::AutodiffModule for #name #generics_ty #generics_where { + burn::empty!(ad_module, #name #generics_ty); + } + + impl #generics core::fmt::Display for #name #generics_ty #generics_where { + #display_fn + } + + impl #generics burn::module::ModuleDisplayDefault for #name #generics_ty #generics_where { + #attributes_fn + } + + impl #generics Clone for #name #generics_ty #generics_where { + #clone_fn + } + }; + + if !has_custom_display(&ast.attrs) { + codegen.extend(quote! { + impl #generics burn::module::ModuleDisplay for #name #generics_ty #generics_where { + + } + }); + } + + codegen +} + +struct GenericsParser { + module: Generics, + module_autodiff: Generics, +} + +impl GenericsParser { + fn from_ast(generics: &Generics, module_generics: &ModuleGenerics) -> Self { + let mut module = GenericsHelper::new(generics.clone()); + let mut module_autodiff = GenericsHelper::new(generics.clone()); + + module.types().into_iter().for_each(|ident| { + // By default, require module bound + let mut requires_module_bound = true; + let mut generic_kind = None; + if !module_generics.is_empty() { + generic_kind = module_generics.get_generic_kind(&ident); + let has_module_bound = matches!(generic_kind, Some(GenericKind::Module)); + let is_unbounded = matches!(generic_kind, Some(GenericKind::Plain)); + + requires_module_bound = has_module_bound || is_unbounded; + } + + if requires_module_bound { + module.add_predicate(parse_quote! { + #ident: burn::module::Module + }); + + module.add_predicate(parse_quote! { + #ident: burn::module::ModuleDisplay + }); + + module_autodiff.add_predicate(parse_quote! { + #ident: burn::module::AutodiffModule + }); + + module_autodiff.add_predicate(parse_quote! { + #ident: burn::module::ModuleDisplay + }); + } else { + // Add required bounds to impl + if let Some(GenericKind::Skip) = generic_kind { + module.add_predicate(parse_quote! { + #ident: Clone + core::fmt::Debug + Send + }); + module_autodiff.add_predicate(parse_quote! { + #ident: Clone + core::fmt::Debug + Send + }); + } + } + }); + + Self { + module: module.generics, + module_autodiff: module_autodiff.generics, + } + } +} + +fn has_custom_display(attrs: &[Attribute]) -> bool { + attrs.iter().any(|attr| { + attr.path().is_ident("module") + && attr + .parse_nested_meta(|meta| { + if meta.path.is_ident("custom_display") { + Ok(()) + } else { + Err(meta.error("unsupported attribute")) + } + }) + .is_ok() + }) +} diff --git a/crates/burn-derive/src/module/codegen_enum.rs b/crates/burn-derive/src/module/codegen_enum.rs new file mode 100644 index 0000000..bcd8dcd --- /dev/null +++ b/crates/burn-derive/src/module/codegen_enum.rs @@ -0,0 +1,303 @@ +use super::codegen::ModuleCodegen; +use crate::module::{ + codegen_struct::{ModuleFieldType, parse_module_field_type}, + generics::{ModuleGenerics, parse_module_generics}, +}; +use proc_macro2::{Ident, Span, TokenStream}; +use quote::quote; + +pub(crate) struct EnumModuleCodegen { + pub name: Ident, + pub variants: Vec, + pub generics: ModuleGenerics, +} + +impl ModuleCodegen for EnumModuleCodegen { + fn gen_num_params(&self) -> TokenStream { + let match_body = self.gen_variants_match_fn(|_| { + quote! { + burn::module::Module::num_params(module) + } + }); + + quote! { + fn num_params(&self) -> usize { + #match_body + } + } + } + + fn gen_visit(&self) -> TokenStream { + let enum_name = self.name.to_string(); + let container_type = format!("Enum:{}", enum_name); + let match_body = self.gen_variants_match_fn(|variant_name| { + let variant_str = variant_name.to_string(); + quote! { + { + visitor.enter_module(#variant_str, #container_type); + burn::module::Module::visit(module, visitor); + visitor.exit_module(#variant_str, #container_type); + } + } + }); + + quote! { + fn visit(&self, visitor: &mut Visitor) { + #match_body + } + } + } + + fn gen_collect_devices(&self) -> TokenStream { + let match_body = self.gen_variants_match_fn(|_| { + quote! { + burn::module::Module::collect_devices(module, devices) + } + }); + + quote! { + fn collect_devices( + &self, + devices: burn::module::Devices + ) -> burn::module::Devices { + #match_body + } + } + } + + fn gen_to_device(&self) -> TokenStream { + let match_body = self.gen_variants_match_fn(|variant| { + quote! { + Self::#variant(burn::module::Module::to_device(module, device)) + } + }); + + quote! { + fn to_device(self, device: &burn::tensor::Device) -> Self { + #match_body + } + } + } + + fn gen_fork(&self) -> TokenStream { + let match_body = self.gen_variants_match_fn(|variant| { + quote! { + Self::#variant(burn::module::Module::fork(module, device)) + } + }); + + quote! { + fn fork(self, device: &burn::tensor::Device) -> Self { + #match_body + } + } + } + + fn gen_map(&self) -> TokenStream { + let enum_name = self.name.to_string(); + let container_type = format!("Enum:{}", enum_name); + let match_body = self.gen_variants_match_fn(|variant| { + let variant_str = variant.to_string(); + quote! { + { + mapper.enter_module(#variant_str, #container_type); + let result = burn::module::Module::map(module, mapper); + mapper.exit_module(#variant_str, #container_type); + Self::#variant(result) + } + } + }); + + quote! { + fn map(self, mapper: &mut Mapper) -> Self { + #match_body + } + } + } + + fn gen_valid(&self) -> TokenStream { + let match_body = self.gen_variants_match_fn(|variant| { + quote! { + Self::#variant(burn::module::AutodiffModule::valid(module)) + } + }); + + quote! { + fn valid(&self) -> Self { + #match_body + } + } + } + + fn gen_from_inner(&self) -> TokenStream { + let match_body = self.gen_variants_match_fn_param("module", "Self::", |variant| { + quote! { + Self::#variant(burn::module::AutodiffModule::from_inner(module)) + } + }); + + quote! { + fn from_inner(module: Self) -> Self { + #match_body + } + } + } + + fn gen_clone(&self) -> TokenStream { + let match_body = self.gen_variants_match_fn(|variant| { + quote! { + Self::#variant(module.clone()) + } + }); + + quote! { + fn clone(&self) -> Self { + #match_body + } + } + } + + fn module_generics(&self) -> &ModuleGenerics { + &self.generics + } + + fn gen_display(&self) -> TokenStream { + // Only tuple enum variants with exactly one field are supported + let variant_prints = self.variants.iter().map(|variant| { + let variant_name = &variant.ident; + let field_names = + (0..1).map(|i| syn::Ident::new(&format!("_{i}"), proc_macro2::Span::call_site())); + + let field_prints = field_names.clone().map(|field_name| { + quote! { .add(stringify!(#field_name), #field_name) } + }); + quote! { + Self::#variant_name(#(#field_names),*) => { + content.set_top_level_type(&stringify!(#variant_name)) + #(#field_prints)* + .optional() + } + } + }); + quote! { + fn content(&self, mut content: burn::module::Content) -> Option { + match self { + #(#variant_prints)* + } + } + } + } +} + +impl EnumModuleCodegen { + pub fn from_ast(ast: &syn::DeriveInput) -> syn::Result { + let mut generics = parse_module_generics(&ast.generics); + Ok(Self { + name: ast.ident.clone(), + variants: parse_variants(ast, &mut generics)?, + generics, + }) + } + + /// Generate the enum variants' match arms with the provided function + fn gen_variants_match_fn(&self, func: F) -> TokenStream + where + F: Fn(Ident) -> TokenStream, + { + self.gen_variants_match_fn_param("self", "Self::", func) + } + + /// Generate a match expression over the given argument (e.g., `self`) + /// and using the provided prefix for variants (e.g., `Self::`) + fn gen_variants_match_fn_param(&self, arg: &str, prefix: &str, func: F) -> TokenStream + where + F: Fn(Ident) -> TokenStream, + { + let match_arms = self.variants.iter().map(|variant| { + let name = &variant.ident; + let full_variant = syn::parse_str::(&format!("{prefix}{name}")).unwrap(); + let arm_pattern = quote! { #full_variant(module) }; + let arm_code = func(name.clone()); + quote! { #arm_pattern => #arm_code, } + }); + + let arg = Ident::new(arg, Span::call_site()); + + quote! { + match #arg { + #(#match_arms)* + } + } + } + + /// Returns true if any field in any variant is considered a module. + pub fn has_module_fields(&self) -> bool { + self.variants + .iter() + .any(|variant| variant.field_type.is_module) + } +} + +/// Module enum variant +pub(crate) struct EnumVariant { + pub ident: syn::Ident, + pub field_type: ModuleFieldType, +} + +pub(crate) fn parse_variants( + ast: &syn::DeriveInput, + generics: &mut ModuleGenerics, +) -> syn::Result> { + let enum_data = match &ast.data { + syn::Data::Enum(data) => data, + _ => return Err(syn::Error::new_spanned(ast, "Only enums are supported")), + }; + + let mut variants = Vec::new(); + + for variant in enum_data.variants.iter() { + for attr in &variant.attrs { + if attr.path().is_ident("module") { + Err(syn::Error::new_spanned( + variant, + "Module attributes are not supported for enum variants.", + ))?; + } + } + + match &variant.fields { + syn::Fields::Unnamed(fields) if fields.unnamed.len() == 1 => { + let field = &fields.unnamed[0]; + + // USE THE SAME PARSER AS STRUCTS + // This gives us the is_module, is_param, etc. logic + let field_type = parse_module_field_type(field, generics)?; + + variants.push(EnumVariant { + ident: variant.ident.clone(), + field_type, + }); + } + syn::Fields::Unnamed(_) => { + return Err(syn::Error::new_spanned( + variant, + "Module derive only supports tuple enum variants with exactly one field.", + )); + } + syn::Fields::Named(_) => { + return Err(syn::Error::new_spanned( + variant, + "Module derive does not support struct enum variants.", + )); + } + syn::Fields::Unit => { + return Err(syn::Error::new_spanned( + variant, + "Module derive does not support unit enum variants.", + )); + } + } + } + + Ok(variants) +} diff --git a/crates/burn-derive/src/module/codegen_struct.rs b/crates/burn-derive/src/module/codegen_struct.rs new file mode 100644 index 0000000..fae8adf --- /dev/null +++ b/crates/burn-derive/src/module/codegen_struct.rs @@ -0,0 +1,498 @@ +use std::collections::HashSet; + +use crate::module::generics::{ + GenericKind, ModuleGenerics, parse_module_generics, parse_ty_generics, +}; + +use super::codegen::ModuleCodegen; +use proc_macro2::{Ident, TokenStream}; +use quote::{ToTokens, quote}; +use syn::Field; + +pub(crate) struct StructModuleCodegen { + pub name: Ident, + pub fields: Vec, + pub generics: ModuleGenerics, +} + +impl ModuleCodegen for StructModuleCodegen { + fn gen_num_params(&self) -> TokenStream { + let body = self.gen_fields_fn(|name, field_type| { + if field_type.is_parameter_module() || field_type.maybe_generic_module() { + quote! { + num_params += burn::module::Module::num_params(&self.#name); + } + } else { + quote! {} // other fields have 0 params + } + }); + + quote! { + fn num_params(&self) -> usize { + let mut num_params = 0; + #body + num_params + } + } + } + + fn gen_visit(&self) -> TokenStream { + let struct_name = self.name.to_string(); + let container_type = format!("Struct:{}", struct_name); + let body = self.gen_fields_fn(|name, field_type| { + if field_type.is_parameter_module() || field_type.maybe_generic_module() { + let name_str = name.to_string(); + quote! { + visitor.enter_module(#name_str, #container_type); + burn::module::Module::visit(&self.#name, visitor); + visitor.exit_module(#name_str, #container_type); + } + } else { + quote! {} + } + }); + + quote! { + fn visit(&self, visitor: &mut Visitor) { + #body + } + } + } + + fn gen_collect_devices(&self) -> TokenStream { + let body = self.gen_fields_fn(|name, field_type| { + if field_type.is_module || field_type.maybe_generic_module() { + quote! { + let devices = burn::module::Module::collect_devices(&self.#name, devices); + } + } else { + quote! {} + } + }); + + quote! { + fn collect_devices( + &self, + devices: burn::module::Devices + ) -> burn::module::Devices { + #body + devices + } + } + } + + fn gen_to_device(&self) -> TokenStream { + let (names, body) = self.gen_fields_fn_names(|name, field_type| { + if field_type.is_module || field_type.maybe_generic_module() { + quote! { + let #name = burn::module::Module::to_device(self.#name, device); + } + } else { + quote! { let #name = self.#name; } + } + }); + + quote! { + fn to_device(self, device: &burn::tensor::Device) -> Self { + #body + Self { #(#names),* } + } + } + } + + fn gen_fork(&self) -> TokenStream { + let (names, body) = self.gen_fields_fn_names(|name, field_type| { + if field_type.is_module || field_type.maybe_generic_module() { + quote! { + let #name = burn::module::Module::fork(self.#name, device); + } + } else { + quote! { let #name = self.#name; } + } + }); + + quote! { + fn fork(self, device: &burn::tensor::Device) -> Self { + #body + Self { #(#names),* } + } + } + } + + fn gen_map(&self) -> TokenStream { + let struct_name = self.name.to_string(); + let container_type = format!("Struct:{}", struct_name); + let (names, body) = self.gen_fields_fn_names(|name, field_type| { + if field_type.is_parameter_module() || field_type.maybe_generic_module() { + let name_str = name.to_string(); + quote! { + mapper.enter_module(#name_str, #container_type); + let #name = burn::module::Module::map(self.#name, mapper); + mapper.exit_module(#name_str, #container_type); + } + } else { + quote! { let #name = self.#name; } + } + }); + + quote! { + fn map(self, mapper: &mut Mapper) -> Self { + #body + Self { #(#names),* } + } + } + } + + fn gen_valid(&self) -> TokenStream { + let (names, body) = self.gen_fields_fn_names(|name, field_type| { + if field_type.is_module || field_type.maybe_generic_module() { + quote! { + let #name = burn::module::AutodiffModule::valid(&self.#name); + } + } else { + quote! { let #name = self.#name.clone(); } + } + }); + + quote! { + fn valid(&self) -> Self { + #body + Self { #(#names),* } + } + } + } + + fn gen_from_inner(&self) -> TokenStream { + let (names, body) = self.gen_fields_fn_names(|name, field_type| { + if field_type.is_module || field_type.maybe_generic_module() { + quote! { + let #name = burn::module::AutodiffModule::from_inner(#name); + } + } else { + quote! { let #name = #name; } + } + }); + + let destructure = quote! { + let Self { #(#names),* } = module; + }; + + quote! { + fn from_inner(module: Self) -> Self { + #destructure + #body + Self { #(#names),* } + } + } + } + + fn gen_clone(&self) -> TokenStream { + let (names, body) = self.gen_fields_fn_names(|name, _field_type| { + quote! { + let #name = self.#name.clone(); + } + }); + + quote! { + fn clone(&self) -> Self { + #body + Self { #(#names),* } + } + } + } + + fn module_generics(&self) -> &ModuleGenerics { + &self.generics + } + + fn gen_display(&self) -> TokenStream { + // let struct_name = self.name.to_string(); + let struct_ident = &self.name; + let field_prints = self.fields.iter().map(|field| { + let field_name = field.ident(); + if field.field_type.is_module || field.field_type.maybe_generic_module() { + // Standard module type, use underlying `ModuleDisplay` impl + quote! { .add(stringify!(#field_name), &self.#field_name) } + } else { + // Not a module, use the debug implementation + quote! { + .add_debug_attribute(stringify!(#field_name), &self.#field_name) + } + } + }); + quote! { + fn content(&self, mut content: burn::module::Content) -> Option { + content + .set_top_level_type(&stringify!(#struct_ident)) + #(#field_prints)* + .optional() + } + } + } +} + +impl StructModuleCodegen { + pub fn from_ast(ast: &syn::DeriveInput) -> syn::Result { + let mut generics = parse_module_generics(&ast.generics); + Ok(Self { + name: ast.ident.clone(), + fields: parse_module_fields(ast, &mut generics)?, + generics, + }) + } + + fn gen_fields_fn_names(&self, func: F) -> (Vec, TokenStream) + where + F: Fn(Ident, &ModuleFieldType) -> TokenStream, + { + let mut body = quote! {}; + let mut names = Vec::new(); + + for field in self.fields.iter() { + let name = field.ident(); + + names.push(name.clone()); + body.extend(func(name, &field.field_type)); + } + + (names, body) + } + + fn gen_fields_fn(&self, func: F) -> TokenStream + where + F: Fn(Ident, &ModuleFieldType) -> TokenStream, + { + let mut body = quote! {}; + + for field in self.fields.iter() { + body.extend(func(field.ident(), &field.field_type)); + } + + body + } + + pub fn has_module_fields(&self) -> bool { + // let has_module_fields = self.fields.iter().any(|f| f.field_type.is_module); + // eprintln!("[HAS MODULE FIELDS] {}: {has_module_fields:?}", self.name); + self.fields.iter().any(|f| f.field_type.is_module) + } +} + +#[derive(new, Debug)] +pub struct ModuleField { + pub field: Field, + pub field_type: ModuleFieldType, +} + +impl ModuleField { + pub fn ident(&self) -> Ident { + self.field.ident.clone().unwrap() + } +} + +#[derive(Debug)] +pub enum ModuleFieldAttribute { + Skip, +} + +#[derive(Default, Debug)] +pub struct ModuleFieldType { + pub is_module: bool, + pub attr: Option, + pub generic_idents: HashSet, +} + +impl ModuleFieldType { + /// Returns true if the field is a module with parameters + /// (i.e., a real module that is neither skipped nor constant). + pub fn is_parameter_module(&self) -> bool { + self.is_module && self.attr.is_none() + } + + /// Returns true for generic fields that are assumed to be modules. + pub fn maybe_generic_module(&self) -> bool { + // We assumed it might be a module generic if the field is not marked + // by any attributes (skip or constant) + !self.generic_idents.is_empty() && self.attr.is_none() + } +} + +pub(crate) fn parse_module_fields( + ast: &syn::DeriveInput, + generics: &mut ModuleGenerics, +) -> syn::Result> { + let mut fields = Vec::new(); + let mut module_generics = HashSet::new(); + let mut skip_generics = HashSet::new(); + + match &ast.data { + syn::Data::Struct(struct_data) => { + for field in struct_data.fields.iter() { + let field_type = parse_module_field_type(field, generics)?; + if field_type.is_module { + module_generics.extend(field_type.generic_idents.iter().cloned()); + } else if matches!(field_type.attr, Some(ModuleFieldAttribute::Skip)) { + skip_generics.extend(field_type.generic_idents.iter().cloned()); + } + fields.push(ModuleField::new(field.clone(), field_type)); + } + } + syn::Data::Enum(_) => panic!("Only struct can be derived"), + syn::Data::Union(_) => panic!("Only struct can be derived"), + }; + + for ident in module_generics.intersection(&skip_generics) { + if let Some(param) = ast.generics.params.iter().find_map(|p| match p { + syn::GenericParam::Type(tp) if tp.ident == *ident => Some(tp), + _ => None, + }) { + return Err(syn::Error::new( + param.ident.span(), + format!( + "Generic type `{ident}` should not be used on both a module field and a skipped field. \ + Consider removing `#[module(skip)]` or using a different type for one of the fields.", + ), + )); + } + } + + Ok(fields) +} + +pub(crate) fn parse_module_field_type( + field: &Field, + generics: &mut ModuleGenerics, +) -> syn::Result { + let mut field_type = ModuleFieldType::default(); + + // Check for generics + let mut has_module_bound = false; + let field_generics = parse_ty_generics(&field.ty, generics) + .into_iter() + .inspect(|ident| { + has_module_bound |= generics.is_bounded_module(ident); + }) + .collect::>(); + + // Infer if a field is a module + let is_primitive = is_primitive_type(&field.ty); + let is_param = is_param_type(&field.ty); + + for attr in &field.attrs { + if attr.path().is_ident("module") { + attr.parse_nested_meta(|meta| { + if meta.path.is_ident("skip") { + // Mark field attribute and generic + field_type.attr = Some(ModuleFieldAttribute::Skip); + for ty in &field_generics { + generics.update(ty, GenericKind::Skip); + } + Ok(()) + } else { + let path = meta.path.to_token_stream().to_string(); + Err(meta.error(format!("Unsupported module attribute: {}", path))) + }?; + + if is_param && field_type.attr.is_some() { + Err(meta.error("Fields of type 'Param' should not be marked as 'skip'. Use a 'Tensor' instead.")) + } else if has_module_bound && field_type.attr.is_some() { + Err(meta.error("Fields with module generics should not be marked as 'skip'.")) + } + else { + Ok(()) + } + })?; + } + } + + field_type.is_module = + !(is_primitive || matches!(field_type.attr, Some(ModuleFieldAttribute::Skip))); + field_type.generic_idents = field_generics; + + Ok(field_type) +} + +fn type_matches_ident(ty: &syn::Type, idents: &[&str]) -> bool { + if let syn::Type::Path(type_path) = ty { + // Look at the last segment of the path (e.g., 'Param' in 'burn::module::Param') + if let Some(segment) = type_path.path.segments.last() { + return idents.contains(&segment.ident.to_string().as_str()); + } + } + false +} + +fn is_primitive_ident(ident: &str) -> bool { + matches!( + ident, + "bool" + | "u8" + | "u16" + | "u32" + | "u64" + | "usize" + | "i8" + | "i16" + | "i32" + | "i64" + | "isize" + | "f32" + | "f64" + | "String" + ) +} + +fn is_primitive_type(ty: &syn::Type) -> bool { + match ty { + // e.g. usize, String, or Option + syn::Type::Path(type_path) => { + let segment = match type_path.path.segments.last() { + Some(seg) => seg, + None => return false, + }; + + let ident = segment.ident.to_string(); + + // Direct primitive + if is_primitive_ident(&ident) { + return true; + } + + // Generic types like Option, Vec, etc. + match &segment.arguments { + syn::PathArguments::AngleBracketed(args) => args.args.iter().all(|arg| { + if let syn::GenericArgument::Type(inner_ty) = arg { + is_primitive_type(inner_ty) + } else { + false + } + }), + _ => false, + } + } + + // e.g. (T, U, ...) + syn::Type::Tuple(tuple) => tuple.elems.iter().all(is_primitive_type), + + // e.g. [T; N] + syn::Type::Array(array) => is_primitive_type(&array.elem), + + // e.g. &T or &mut T + syn::Type::Reference(reference) => is_primitive_type(&reference.elem), + + // e.g. *const T / *mut T + syn::Type::Ptr(ptr) => is_primitive_type(&ptr.elem), + + // e.g. (T) + syn::Type::Paren(paren) => is_primitive_type(&paren.elem), + + // slices `[T]` + syn::Type::Slice(slice) => is_primitive_type(&slice.elem), + + _ => false, + } +} + +fn is_param_type(ty: &syn::Type) -> bool { + type_matches_ident(ty, &["Param"]) +} diff --git a/crates/burn-derive/src/module/display.rs b/crates/burn-derive/src/module/display.rs new file mode 100644 index 0000000..9b32dce --- /dev/null +++ b/crates/burn-derive/src/module/display.rs @@ -0,0 +1,10 @@ +use quote::quote; + +pub fn display_fn(_ast: &syn::DeriveInput) -> proc_macro2::TokenStream { + quote! { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let formatted = burn::module::ModuleDisplay::format(self, Default::default()); + write!(f, "{}", formatted) + } + } +} diff --git a/crates/burn-derive/src/module/generics.rs b/crates/burn-derive/src/module/generics.rs new file mode 100644 index 0000000..5281e51 --- /dev/null +++ b/crates/burn-derive/src/module/generics.rs @@ -0,0 +1,129 @@ +use std::collections::{HashMap, HashSet}; + +use proc_macro2::Ident; +use syn::{GenericParam, Generics, Type, TypeParamBound, WherePredicate, visit::Visit}; + +#[derive(Debug)] +pub enum GenericKind { + /// A generic with `Module` bound. + Module, + /// A generic used in a field marked by `#[module(skip)]`. + Skip, + /// A plain generic that does not fit any of the above conditions. + Plain, +} + +#[derive(Debug)] +pub struct ModuleGenerics { + kinds: HashMap, +} + +impl ModuleGenerics { + pub fn is_empty(&self) -> bool { + self.kinds.is_empty() + } + + pub fn get_generic_kind(&self, ident: &Ident) -> Option<&GenericKind> { + self.kinds.get(ident) + } + + pub fn is_bounded_module(&self, ident: &Ident) -> bool { + self.kinds + .get(ident) + .map(|kind| matches!(kind, GenericKind::Module)) + .unwrap_or(false) + } + + pub fn update(&mut self, ident: &Ident, kind: GenericKind) { + self.kinds.insert(ident.clone(), kind); + } + + pub fn contains(&self, ident: &Ident) -> bool { + self.kinds.contains_key(ident) + } +} + +pub fn parse_module_generics(generics: &Generics) -> ModuleGenerics { + let mut kinds = HashMap::new(); + + // Check inline bounds e.g. `M: Module` + for param in &generics.params { + if let GenericParam::Type(type_param) = param { + let ident = &type_param.ident; + if has_module_bound(&type_param.bounds) { + kinds.insert(ident.clone(), GenericKind::Module); + } else { + kinds.insert(ident.clone(), GenericKind::Plain); + } + } + } + + // Check `where` clauses + if let Some(where_clause) = &generics.where_clause { + for predicate in &where_clause.predicates { + if let WherePredicate::Type(pt) = predicate { + // We only care if the bounded type is a simple identifier (like 'M') + if let Type::Path(p) = &pt.bounded_ty + && let Some(ident) = p.path.get_ident() + { + if has_module_bound(&pt.bounds) { + kinds.insert(ident.clone(), GenericKind::Module); + } else { + kinds.insert(ident.clone(), GenericKind::Plain); + } + } + } + } + } + + ModuleGenerics { kinds } +} + +/// Helper to check if a list of bounds contains "Module". +fn has_module_bound( + bounds: &syn::punctuated::Punctuated, +) -> bool { + has_bound(bounds, "Module") +} + +/// Helper to check if a list of bounds contains the specified bound. +fn has_bound( + bounds: &syn::punctuated::Punctuated, + ident: &str, +) -> bool { + bounds.iter().any(|bound| { + if let TypeParamBound::Trait(trait_bound) = bound + && let Some(segment) = trait_bound.path.segments.last() + { + return segment.ident == ident; + } + false + }) +} + +pub fn parse_ty_generics(ty: &Type, declared: &ModuleGenerics) -> HashSet { + struct Collector<'a> { + generics: HashSet, + declared: &'a ModuleGenerics, + } + + impl<'ast, 'a> Visit<'ast> for Collector<'a> { + fn visit_type_path(&mut self, type_path: &'ast syn::TypePath) { + if type_path.qself.is_none() + && let Some(ident) = type_path.path.get_ident() + && (self.declared.contains(ident)) + { + self.generics.insert(ident.clone()); + } + + syn::visit::visit_type_path(self, type_path); + } + } + + let mut collector = Collector { + generics: HashSet::new(), + declared, + }; + collector.visit_type(ty); + collector.generics +} diff --git a/crates/burn-derive/src/module/mod.rs b/crates/burn-derive/src/module/mod.rs new file mode 100644 index 0000000..66b68d3 --- /dev/null +++ b/crates/burn-derive/src/module/mod.rs @@ -0,0 +1,9 @@ +pub(crate) mod codegen; +pub(crate) mod codegen_enum; +pub(crate) mod codegen_struct; +pub(crate) mod display; +pub(crate) mod generics; + +mod base; + +pub(crate) use base::*; diff --git a/crates/burn-derive/src/record_state/mod.rs b/crates/burn-derive/src/record_state/mod.rs new file mode 100644 index 0000000..8bb164b --- /dev/null +++ b/crates/burn-derive/src/record_state/mod.rs @@ -0,0 +1,257 @@ +use quote::quote; +use syn::{Data, DeriveInput, Fields, GenericArgument, PathArguments, Type}; + +/// How a state field is serialized. +enum FieldKind { + /// A `Tensor` leaf. + Tensor, + /// An `Option>` leaf (carrying the inner `Tensor` type). + OptionTensor(Type), + /// A `Vec>` field (carrying the inner `Tensor` type). + VecTensor(Type), + /// A scalar leaf serialized via `burn_pack::Scalar`'s `From`/`TryFrom` conversions. + Scalar, + /// An `Option` leaf (carrying the inner scalar type). + OptionScalar(Type), + /// A nested [`RecordState`] field. + Nested, + /// An `Option` field (carrying the inner nested type). + OptionNested(Type), +} + +/// Returns the identifier of the last path segment of a type, if any (`Tensor` -> `Tensor`). +fn head_ident(ty: &Type) -> Option { + match ty { + Type::Path(path) => path.path.segments.last().map(|s| s.ident.to_string()), + _ => None, + } +} + +/// Returns the single generic type argument of a type (the `T` in `Option` / `Vec`). +fn single_generic_arg(ty: &Type) -> Option { + let Type::Path(path) = ty else { return None }; + let segment = path.path.segments.last()?; + let PathArguments::AngleBracketed(args) = &segment.arguments else { + return None; + }; + args.args.iter().find_map(|arg| match arg { + GenericArgument::Type(inner) => Some(inner.clone()), + _ => None, + }) +} + +fn is_scalar(ident: &str) -> bool { + // Only the widths `burn_pack::Scalar` has `From`/`TryFrom` conversions for. `i128`/`u128` are + // intentionally excluded — `Scalar` tops out at 64-bit, so a 128-bit field is treated as a + // nested `RecordState` and fails with a clear trait-bound error rather than a silent truncation. + matches!( + ident, + "usize" + | "isize" + | "u8" + | "u16" + | "u32" + | "u64" + | "i8" + | "i16" + | "i32" + | "i64" + | "f32" + | "f64" + | "bool" + ) +} + +fn classify(ty: &Type) -> Result { + Ok(match head_ident(ty).as_deref() { + Some("Tensor") => FieldKind::Tensor, + Some("Vec") => match single_generic_arg(ty) { + Some(inner) if head_ident(&inner).as_deref() == Some("Tensor") => { + FieldKind::VecTensor(inner) + } + _ => { + return Err(syn::Error::new_spanned( + ty, + "RecordState only supports `Vec>` for sequence fields.", + )); + } + }, + Some("Option") => match single_generic_arg(ty) { + Some(inner) => match head_ident(&inner).as_deref() { + Some("Tensor") => FieldKind::OptionTensor(inner), + Some(ident) if is_scalar(ident) => FieldKind::OptionScalar(inner), + _ => FieldKind::OptionNested(inner), + }, + None => FieldKind::Nested, + }, + Some(ident) if is_scalar(ident) => FieldKind::Scalar, + _ => FieldKind::Nested, + }) +} + +pub(crate) fn derive_impl(ast: &DeriveInput) -> proc_macro::TokenStream { + match try_derive_impl(ast) { + Ok(tokens) => tokens, + Err(err) => err.to_compile_error().into(), + } +} + +fn try_derive_impl(ast: &DeriveInput) -> Result { + let name = &ast.ident; + let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl(); + + let fields = match &ast.data { + Data::Struct(data) => match &data.fields { + Fields::Named(named) => &named.named, + other => { + return Err(syn::Error::new_spanned( + other, + "RecordState can only be derived for structs with named fields.", + )); + } + }, + _ => { + return Err(syn::Error::new_spanned( + name, + "RecordState can only be derived for structs.", + )); + } + }; + + let mut flatten = Vec::new(); + let mut unflatten = Vec::new(); + let mut ctor = Vec::new(); + + for field in fields { + let ident = field.ident.as_ref().unwrap(); + let leaf = ident.to_string(); + let ty = &field.ty; + ctor.push(quote! { #ident }); + + match classify(ty)? { + FieldKind::Tensor => { + flatten.push(quote! { + out.push_tensor(prefix, #leaf, self.#ident.clone().into_data()); + }); + unflatten.push(quote! { + let #ident = <#ty>::from_data(src.take_tensor(prefix, #leaf)?, device); + }); + } + FieldKind::OptionTensor(inner) => { + flatten.push(quote! { + if let Some(value) = &self.#ident { + out.push_tensor(prefix, #leaf, value.clone().into_data()); + } + }); + unflatten.push(quote! { + let #ident = src + .take_tensor(prefix, #leaf) + .map(|data| <#inner>::from_data(data, device)); + }); + } + FieldKind::VecTensor(inner) => { + flatten.push(quote! { + out.push_scalar( + prefix, + concat!(#leaf, ".len"), + burn_pack::Scalar::from(self.#ident.len()), + ); + for (__index, __value) in self.#ident.iter().enumerate() { + let __leaf = crate::join_index(#leaf, __index); + out.push_tensor(prefix, &__leaf, __value.clone().into_data()); + } + }); + unflatten.push(quote! { + let #ident = { + let __len = >::try_from( + src.take_scalar(prefix, concat!(#leaf, ".len"))?, + ) + .ok()?; + let mut __items = ::alloc::vec::Vec::with_capacity(__len); + for __index in 0..__len { + let __leaf = crate::join_index(#leaf, __index); + __items.push(<#inner>::from_data( + src.take_tensor(prefix, &__leaf)?, + device, + )); + } + __items + }; + }); + } + FieldKind::Scalar => { + flatten.push(quote! { + out.push_scalar(prefix, #leaf, burn_pack::Scalar::from(self.#ident)); + }); + unflatten.push(quote! { + let #ident = <#ty as ::core::convert::TryFrom>::try_from( + src.take_scalar(prefix, #leaf)?, + ) + .ok()?; + }); + } + FieldKind::OptionScalar(inner) => { + flatten.push(quote! { + if let Some(value) = self.#ident { + out.push_scalar(prefix, #leaf, burn_pack::Scalar::from(value)); + } + }); + unflatten.push(quote! { + let #ident: Option<#inner> = src.take_scalar(prefix, #leaf).and_then(|__s| { + <#inner as ::core::convert::TryFrom>::try_from(__s).ok() + }); + }); + } + FieldKind::Nested => { + flatten.push(quote! { + { + let __path = crate::join_path(prefix, #leaf); + crate::RecordState::state_flatten(&self.#ident, &__path, out); + } + }); + unflatten.push(quote! { + let #ident = { + let __path = crate::join_path(prefix, #leaf); + <#ty as crate::RecordState>::state_unflatten(&__path, src, device)? + }; + }); + } + FieldKind::OptionNested(inner) => { + flatten.push(quote! { + if let Some(value) = &self.#ident { + let __path = crate::join_path(prefix, #leaf); + crate::RecordState::state_flatten(value, &__path, out); + } + }); + unflatten.push(quote! { + let #ident = { + let __path = crate::join_path(prefix, #leaf); + if src.has_under(&__path) { + <#inner as crate::RecordState>::state_unflatten(&__path, src, device) + } else { + None + } + }; + }); + } + } + } + + Ok(quote! { + impl #impl_generics crate::RecordState for #name #ty_generics #where_clause { + fn state_flatten(&self, prefix: &str, out: &mut crate::StateSink) { + #(#flatten)* + } + + fn state_unflatten( + prefix: &str, + src: &mut crate::StateSource, + device: &burn::tensor::Device, + ) -> Option { + #(#unflatten)* + Some(Self { #(#ctor),* }) + } + } + } + .into()) +} diff --git a/crates/burn-derive/src/shared/attribute.rs b/crates/burn-derive/src/shared/attribute.rs new file mode 100644 index 0000000..bc1f037 --- /dev/null +++ b/crates/burn-derive/src/shared/attribute.rs @@ -0,0 +1,49 @@ +use syn::{Attribute, Meta}; + +pub struct AttributeAnalyzer { + attr: Attribute, +} + +#[derive(Clone)] +pub struct AttributeItem { + pub value: syn::Lit, +} + +impl AttributeAnalyzer { + pub fn new(attr: Attribute) -> Self { + Self { attr } + } + + pub fn item(&self) -> AttributeItem { + let value = match &self.attr.meta { + Meta::List(val) => val.parse_args::().unwrap(), + Meta::NameValue(meta) => meta.clone(), + Meta::Path(_) => panic!("Path meta unsupported"), + }; + + let lit = match value.value { + syn::Expr::Lit(lit) => lit.lit, + _ => panic!("Only literal is supported"), + }; + + AttributeItem { value: lit } + } + + pub fn has_name(&self, name: &str) -> bool { + Self::path_syn_name(self.attr.path()) == name + } + + fn path_syn_name(path: &syn::Path) -> String { + let length = path.segments.len(); + let mut name = String::new(); + for (i, segment) in path.segments.iter().enumerate() { + if i == length - 1 { + name += segment.ident.to_string().as_str(); + } else { + let tmp = segment.ident.to_string() + "::"; + name += tmp.as_str(); + } + } + name + } +} diff --git a/crates/burn-derive/src/shared/enum_variant.rs b/crates/burn-derive/src/shared/enum_variant.rs new file mode 100644 index 0000000..ceed9fe --- /dev/null +++ b/crates/burn-derive/src/shared/enum_variant.rs @@ -0,0 +1,51 @@ +use proc_macro2::{Ident, Span, TokenStream}; +use quote::quote; +use syn::{FieldsNamed, Variant}; + +/// Process a variant of an enum where the output is the result of the given mapper. +pub(crate) fn map_enum_variant( + variant: &Variant, + mapper: Mapper, +) -> (TokenStream, TokenStream) +where + Mapper: Fn(&Ident) -> TokenStream, +{ + let gen_fields_unnamed = |num: usize| { + let mut inputs = Vec::new(); + let mut outputs = Vec::new(); + + for i in 0..num { + let arg_name = Ident::new(&format!("arg_{i}"), Span::call_site()); + let input = quote! { #arg_name }; + let output = mapper(&arg_name); + + inputs.push(input); + outputs.push(output); + } + + (quote! (( #(#inputs),* )), quote! (( #(#outputs),* ))) + }; + let gen_fields_named = |fields: &FieldsNamed| { + let mut inputs = Vec::new(); + let mut outputs = Vec::new(); + + fields.named.iter().for_each(|field| { + let ident = field.ident.as_ref().expect("Named field to have a name."); + let input = quote! { #ident }; + let output = mapper(ident); + + inputs.push(input); + outputs.push(quote! { + #ident: #output + }); + }); + + (quote! {{ #(#inputs),* }}, quote! {{ #(#outputs),* }}) + }; + + match &variant.fields { + syn::Fields::Named(fields) => gen_fields_named(fields), + syn::Fields::Unnamed(_) => gen_fields_unnamed(variant.fields.len()), + syn::Fields::Unit => (quote! {}, quote! {}), + } +} diff --git a/crates/burn-derive/src/shared/field.rs b/crates/burn-derive/src/shared/field.rs new file mode 100644 index 0000000..7dd5534 --- /dev/null +++ b/crates/burn-derive/src/shared/field.rs @@ -0,0 +1,84 @@ +use super::attribute::AttributeAnalyzer; +use proc_macro2::Ident; +use syn::{Field, Type, TypePath}; + +#[derive(Clone)] +pub struct FieldTypeAnalyzer { + pub field: Field, +} + +impl FieldTypeAnalyzer { + pub fn new(field: Field) -> Self { + FieldTypeAnalyzer { field } + } + + pub fn ident(&self) -> Ident { + self.field.ident.clone().unwrap() + } + + pub fn is_of_type(&self, paths: &[&str]) -> bool { + match &self.field.ty { + syn::Type::Path(path) => { + let name = Self::path_name(path); + paths.contains(&name.as_str()) + } + _ => false, + } + } + + #[allow(dead_code)] + pub fn first_generic_field(&self) -> TypePath { + let err = || panic!("Field {} as no generic", self.field.ident.clone().unwrap()); + match &self.field.ty { + syn::Type::Path(path) => Self::path_generic_argument(path), + _ => err(), + } + } + pub fn path_generic_argument(path: &TypePath) -> TypePath { + let segment = path.path.segments.last().unwrap(); + let err = || panic!("Path segment {} has no generic", segment.ident.clone(),); + match &segment.arguments { + syn::PathArguments::None => err(), + syn::PathArguments::AngleBracketed(param) => { + let first_param = param.args.first().unwrap(); + + if let syn::GenericArgument::Type(Type::Path(path)) = first_param { + path.clone() + } else { + err() + } + } + syn::PathArguments::Parenthesized(_) => err(), + } + } + + fn path_name(path: &TypePath) -> String { + let length = path.path.segments.len(); + let mut name = String::new(); + for (i, segment) in path.path.segments.iter().enumerate() { + if i == length - 1 { + name += segment.ident.to_string().as_str(); + } else { + let tmp = segment.ident.to_string() + "::"; + name += tmp.as_str(); + } + } + name + } + + /// Returns the docs of the field. + pub fn docs(&self) -> impl Iterator { + self.field + .attrs + .iter() + .filter(|attr| attr.path().is_ident("doc")) + } + + pub fn attributes(&self) -> impl Iterator { + self.field + .attrs + .clone() + .into_iter() + .map(AttributeAnalyzer::new) + } +} diff --git a/crates/burn-derive/src/shared/generics.rs b/crates/burn-derive/src/shared/generics.rs new file mode 100644 index 0000000..edbb019 --- /dev/null +++ b/crates/burn-derive/src/shared/generics.rs @@ -0,0 +1,30 @@ +use proc_macro2::Ident; +use syn::{Generics, WhereClause, WherePredicate, parse_quote}; + +#[derive(new)] +pub struct GenericsHelper { + pub(crate) generics: Generics, +} + +impl GenericsHelper { + pub fn add_predicate(&mut self, predicate: WherePredicate) { + let where_clause: WhereClause = match &self.generics.where_clause { + Some(val) => parse_quote! { + #val + #predicate, + }, + None => parse_quote! { + where + #predicate, + }, + }; + self.generics.where_clause = Some(where_clause); + } + + pub fn types(&self) -> Vec { + self.generics + .type_params() + .map(|tp| tp.ident.clone()) + .collect() + } +} diff --git a/crates/burn-derive/src/shared/mod.rs b/crates/burn-derive/src/shared/mod.rs new file mode 100644 index 0000000..009851f --- /dev/null +++ b/crates/burn-derive/src/shared/mod.rs @@ -0,0 +1,4 @@ +pub(crate) mod attribute; +pub(crate) mod enum_variant; +pub(crate) mod field; +pub(crate) mod generics; diff --git a/crates/burn-dispatch/Cargo.toml b/crates/burn-dispatch/Cargo.toml new file mode 100644 index 0000000..1ea02a4 --- /dev/null +++ b/crates/burn-dispatch/Cargo.toml @@ -0,0 +1,118 @@ +[package] +authors = [ + "laggui ", + "nathanielsimard ", +] +categories = ["science"] +description = "Backend dispatch for the Burn framework" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "data"] +license.workspace = true +name = "burn-dispatch" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-dispatch" +documentation = "https://docs.rs/burn-dispatch" +version.workspace = true + +[lints] +workspace = true + +[features] +default = [ + "std", + "burn-autodiff?/default", + "burn-cpu?/default", + "burn-cuda?/default", + "burn-flex/default", + "burn-ndarray?/default", + "burn-rocm?/default", + "burn-tch?/default", + "burn-wgpu?/default", +] +doc = ["default"] +std = [ + "burn-backend/std", + "burn-autodiff?/std", + "burn-cpu?/std", + "burn-cuda?/std", + "burn-flex/std", + "burn-ndarray?/std", + "burn-rocm?/std", + "burn-tch?/std", + "burn-wgpu?/std", +] +tracing = [ + "burn-autodiff?/tracing", + "burn-cpu?/tracing", + "burn-cuda?/tracing", + "burn-flex/tracing", + "burn-ndarray?/tracing", + "burn-rocm?/tracing", + "burn-tch?/tracing", + "burn-wgpu?/tracing", +] + +# Backends +cuda = ["burn-cuda"] +flex = [] # default +rocm = ["burn-rocm"] +ndarray = ["burn-ndarray"] +tch = ["burn-tch"] +vulkan = ["wgpu", "burn-wgpu/vulkan"] +webgpu = ["wgpu", "burn-wgpu/webgpu"] +metal = ["wgpu", "burn-wgpu/metal"] +wgpu = ["burn-wgpu"] +cpu = ["burn-cpu"] +autodiff = ["burn-autodiff"] +# Remote compute client over Iroh. +remote = ["std", "burn-remote", "burn-remote/iroh"] +# Host a remote compute server over Iroh. +remote-server = ["remote", "burn-remote/server"] +# Add the legacy WebSocket transport. +remote-websocket = ["remote", "burn-remote/websocket"] + +# Backend features +autotune = [ + "burn-wgpu?/autotune", + "burn-cuda?/autotune", + "burn-rocm?/autotune", + "burn-cpu?/autotune", +] +autotune-checks = [ + "burn-wgpu?/autotune-checks", + "burn-cuda?/autotune-checks", + "burn-rocm?/autotune-checks", + "burn-cpu?/autotune-checks", +] +fusion = [ + "burn-wgpu?/fusion", + "burn-cuda?/fusion", + "burn-rocm?/fusion", + "burn-cpu?/fusion", + "burn-remote?/fusion", +] + +[dependencies] +burn-backend = { workspace = true } + +# Backends +burn-autodiff = { workspace = true, optional = true } +burn-cpu = { workspace = true, optional = true } +burn-cuda = { workspace = true, optional = true } +burn-flex = { workspace = true } +burn-ndarray = { workspace = true, optional = true } +burn-tch = { workspace = true, optional = true } +burn-rocm = { workspace = true, optional = true } +burn-wgpu = { workspace = true, optional = true } + +burn-remote = { workspace = true, optional = true, features = ["client"] } + +# Op macros with `.as_$inner_kind()` +paste = { workspace = true } + +[target.'cfg(not(target_has_atomic = "ptr"))'.dependencies] +burn-flex = { workspace = true, features = ["critical-section"]} + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-dispatch/README.md b/crates/burn-dispatch/README.md new file mode 100644 index 0000000..7b81db5 --- /dev/null +++ b/crates/burn-dispatch/README.md @@ -0,0 +1,3 @@ +# Burn Backend Dispatch + +A multi-backend dispatch that forwards the tensor operations to the appropriate backend. diff --git a/crates/burn-dispatch/build.rs b/crates/burn-dispatch/build.rs new file mode 100644 index 0000000..8eca0d7 --- /dev/null +++ b/crates/burn-dispatch/build.rs @@ -0,0 +1,22 @@ +fn main() { + println!("cargo::rustc-check-cfg=cfg(default_backend)"); + + // If you try to build with `--no-default-features`, we enable a cpu backend by default + let cuda = cfg!(feature = "cuda"); + let flex = cfg!(feature = "flex"); + let rocm = cfg!(feature = "rocm"); + let ndarray = cfg!(feature = "ndarray"); + let tch = cfg!(feature = "tch"); + let cpu = cfg!(feature = "cpu"); + let metal = cfg!(feature = "metal"); + let vulkan = cfg!(feature = "vulkan"); + let webgpu = cfg!(feature = "webgpu"); + let wgpu = cfg!(feature = "wgpu"); + + let no_backend_enabled = + !(cuda || flex || rocm || ndarray || tch || cpu || metal || vulkan || webgpu || wgpu); + + if no_backend_enabled { + println!("cargo:rustc-cfg=default_backend"); + } +} diff --git a/crates/burn-dispatch/src/backend.rs b/crates/burn-dispatch/src/backend.rs new file mode 100644 index 0000000..4c5f314 --- /dev/null +++ b/crates/burn-dispatch/src/backend.rs @@ -0,0 +1,940 @@ +use alloc::format; +use alloc::string::String; +use alloc::vec::Vec; + +#[cfg(any( + feature = "cpu", + feature = "ndarray", + feature = "flex", + default_backend +))] +use alloc::vec; + +#[cfg(feature = "autodiff")] +use burn_backend::distributed::{DistributedParamId, DistributedParams}; +use burn_backend::{AutodiffBackend, Backend, BackendGraph, BackendTypes, DType, ExecutionError}; + +/// A captured graph from one of the dispatched backends (see +/// [`BackendTypes::GraphPrimitive`]). +/// +/// Like [`DispatchTensorKind`], one variant per enabled backend: the graph is +/// captured by, and can only replay on, the backend it was recorded on. +#[derive(Debug, Clone)] +pub enum DispatchGraph { + /// A graph captured on the [CPU backend](Cpu). + #[cfg(feature = "cpu")] + Cpu(BackendGraph), + + /// A graph captured on the [CUDA backend](Cuda). + #[cfg(feature = "cuda")] + Cuda(BackendGraph), + + /// A graph captured on the [Metal backend](Metal). + #[cfg(feature = "metal")] + Metal(BackendGraph), + + /// A graph captured on the [ROCm backend](Rocm). + #[cfg(feature = "rocm")] + Rocm(BackendGraph), + + /// A graph captured on the [Vulkan backend](Vulkan). + #[cfg(feature = "vulkan")] + Vulkan(BackendGraph), + + /// A graph captured on the [Wgpu backend](Wgpu). + #[cfg(feature = "wgpu")] + Wgpu(BackendGraph), + + /// A graph captured on the [WebGPU backend](WebGpu). + #[cfg(feature = "webgpu")] + WebGpu(BackendGraph), + + /// A graph captured on the [Flex backend](Flex). + #[cfg(any(feature = "flex", default_backend))] + Flex(BackendGraph), + + /// A graph captured on the [NdArray backend](NdArray). + #[cfg(feature = "ndarray")] + NdArray(BackendGraph), + + /// A graph captured on the [LibTorch backend](LibTorch). + #[cfg(feature = "tch")] + LibTorch(BackendGraph), + + /// A graph captured on the [Remote backend](Remote). + #[cfg(feature = "remote")] + Remote(BackendGraph), +} + +/// The error returned when a graph operation cannot be dispatched. +fn graph_dispatch_err(reason: alloc::string::String) -> ExecutionError { + ExecutionError::WithContext { reason } +} + +/// Match arm generator for [`Backend::graph_stop_capture`] on [`Dispatch`]: +/// each backend's captured graph is wrapped in its [`DispatchGraph`] variant. +macro_rules! graph_stop_capture_arms { + ($device:expr; $([$Backend:ident, $cfg:meta]),*) => { + match $device { + $( + #[cfg($cfg)] + $crate::DispatchDevice::$Backend(device) => { + <$crate::backends::$Backend as Backend>::graph_stop_capture(device) + .map(DispatchGraph::$Backend) + } + )* + #[allow(unreachable_patterns)] + other => Err(graph_dispatch_err(format!( + "Graph capture is not supported for device {other:?}" + ))), + } + }; +} + +/// Match arm generator for [`Backend::graph_replay`] on [`Dispatch`]: the graph +/// variant must match the device's backend, since a graph only replays on the +/// backend that captured it. +macro_rules! graph_replay_arms { + ($device:expr, $graph:expr; $([$Backend:ident, $cfg:meta]),*) => { + match ($device, $graph) { + $( + #[cfg($cfg)] + ($crate::DispatchDevice::$Backend(device), DispatchGraph::$Backend(graph)) => { + // Safety: forwarded verbatim from `Dispatch::graph_replay`'s + // own contract. + unsafe { + <$crate::backends::$Backend as Backend>::graph_replay(device, graph) + } + } + )* + #[allow(unreachable_patterns)] + (device, _) => Err(graph_dispatch_err(format!( + "The graph was not captured on the backend of device {device:?}" + ))), + } + }; +} + +#[cfg(feature = "autodiff")] +use alloc::boxed::Box; +#[cfg(feature = "autodiff")] +use burn_autodiff::grads::Gradients; + +#[allow(unused)] +use crate::DispatchDeviceId; +#[allow(unused)] +use crate::DispatchTensorKind; +use crate::backends::*; +use crate::{DispatchDevice, DispatchTensor}; + +/// The main execution backend in Burn. +/// +/// [`Dispatch`] acts as a global backend that can manage multiple underlying +/// backends (e.g., `Cpu`, `Cuda`, `Wgpu`, `Metal`, etc.). +/// It is responsible for: +/// - Dispatching tensor operations to the appropriate backend. +/// - Managing cross-backend tensor transfers. +/// +/// Essentially, [`Dispatch`] is the single entry point for executing tensor operations +/// in a backend-agnostic way. It allows Burn to provide a unified, global backend +/// for users while still leveraging multiple specialized backends under the hood. +/// +/// # Example +/// +/// ```ignore +/// use burn::Dispatch; +/// use burn::DispatchDevice; +/// +/// // Select the device to execute operations on +/// let device = DispatchDevice::Cuda(Default::default()); +/// +/// // Create a tensor using the global backend +/// let t = Tensor::::zeros([128, 128], &device); +/// ``` +#[derive(Debug, Default, Clone)] +pub struct Dispatch; + +impl BackendTypes for Dispatch { + type Device = DispatchDevice; + + type FloatTensorPrimitive = DispatchTensor; + type IntTensorPrimitive = DispatchTensor; + type BoolTensorPrimitive = DispatchTensor; + type QuantizedTensorPrimitive = DispatchTensor; + + type GraphPrimitive = DispatchGraph; +} + +impl Backend for Dispatch { + fn name(device: &Self::Device) -> String { + let inner = dispatch_device!(device, |device| B::name(device)); + format!("dispatch<{inner}>") + } + + fn seed(device: &Self::Device, seed: u64) { + dispatch_device!(device, |device| B::seed(device, seed)) + } + + fn sync(device: &Self::Device) -> Result<(), ExecutionError> { + dispatch_device!(device, |device| B::sync(device)) + } + + fn graph_prepare(device: &Self::Device) -> Result<(), ExecutionError> { + dispatch_device!(device, |device| B::graph_prepare(device)) + } + + fn graph_start_capture(device: &Self::Device) -> Result<(), ExecutionError> { + dispatch_device!(device, |device| B::graph_start_capture(device)) + } + + fn graph_stop_capture(device: &Self::Device) -> Result { + backend_list!(graph_stop_capture_arms, device) + } + + unsafe fn graph_replay( + device: &Self::Device, + graph: &DispatchGraph, + ) -> Result<(), ExecutionError> { + backend_list!(graph_replay_arms, device, graph) + } + + fn dtype_usage(device: &Self::Device, dtype: DType) -> burn_backend::DTypeUsageSet { + dispatch_device!(device, |device| B::dtype_usage(device, dtype)) + } + + fn ad_enabled(device: &Self::Device) -> bool { + match device { + #[cfg(feature = "autodiff")] + DispatchDevice::Autodiff(_) => true, + _ => false, + } + } + + fn device_count(type_id: u16) -> usize { + let (dispatch_id, backend_type_id) = DispatchDevice::decode_type_id(type_id); + match dispatch_id { + #[cfg(feature = "cpu")] + DispatchDeviceId::Cpu => Cpu::device_count(backend_type_id), + #[cfg(feature = "cuda")] + DispatchDeviceId::Cuda => Cuda::device_count(backend_type_id), + #[cfg(feature = "metal")] + DispatchDeviceId::Metal => Metal::device_count(backend_type_id), + #[cfg(feature = "rocm")] + DispatchDeviceId::Rocm => Rocm::device_count(backend_type_id), + #[cfg(feature = "vulkan")] + DispatchDeviceId::Vulkan => Vulkan::device_count(backend_type_id), + #[cfg(feature = "wgpu")] + DispatchDeviceId::Wgpu => Wgpu::device_count(backend_type_id), + #[cfg(feature = "webgpu")] + DispatchDeviceId::WebGpu => WebGpu::device_count(backend_type_id), + #[cfg(any(feature = "flex", default_backend))] + DispatchDeviceId::Flex => Flex::device_count(backend_type_id), + #[cfg(feature = "ndarray")] + DispatchDeviceId::NdArray => NdArray::device_count(backend_type_id), + #[cfg(feature = "tch")] + DispatchDeviceId::LibTorch => LibTorch::device_count(backend_type_id), + #[cfg(feature = "remote")] + DispatchDeviceId::Remote => Remote::device_count(backend_type_id), + _ => unreachable!("No backend feature enabled."), + } + } + + fn memory_persistent_allocations< + Output: Send, + Input: Send, + Func: Fn(Input) -> Output + Send, + >( + device: &Self::Device, + input: Input, + func: Func, + ) -> Output { + dispatch_device!(device, |device| B::memory_persistent_allocations( + device, input, func + )) + } + + fn memory_cleanup(device: &Self::Device) { + dispatch_device!(device, |device| B::memory_cleanup(device)) + } + + fn staging<'a, Iter>(data: Iter, device: &Self::Device) + where + Iter: Iterator, + { + dispatch_device!(device, |device| B::staging(data, device)) + } + + fn supports_dtype(device: &Self::Device, dtype: DType) -> bool { + dispatch_device!(device, |device| B::supports_dtype(device, dtype)) + } + + fn flush(device: &Self::Device) { + dispatch_device!(device, |device| B::flush(device)) + } +} + +#[cfg(feature = "autodiff")] +impl AutodiffBackend for Dispatch { + type InnerBackend = Dispatch; + + type Gradients = Gradients; + + fn backward(tensor: DispatchTensor) -> Self::Gradients { + let DispatchTensor { kind, .. } = tensor; + + match kind { + DispatchTensorKind::Autodiff(tensor) => match *tensor { + #[cfg(feature = "cpu")] + DispatchTensorKind::Cpu(tensor) => tensor.autodiff().backward(), + #[cfg(feature = "cuda")] + DispatchTensorKind::Cuda(tensor) => tensor.autodiff().backward(), + #[cfg(feature = "metal")] + DispatchTensorKind::Metal(tensor) => tensor.autodiff().backward(), + #[cfg(feature = "rocm")] + DispatchTensorKind::Rocm(tensor) => tensor.autodiff().backward(), + #[cfg(feature = "vulkan")] + DispatchTensorKind::Vulkan(tensor) => tensor.autodiff().backward(), + #[cfg(feature = "wgpu")] + DispatchTensorKind::Wgpu(tensor) => tensor.autodiff().backward(), + #[cfg(feature = "webgpu")] + DispatchTensorKind::WebGpu(tensor) => tensor.autodiff().backward(), + #[cfg(any(feature = "flex", default_backend))] + DispatchTensorKind::Flex(tensor) => tensor.autodiff().backward(), + #[cfg(feature = "ndarray")] + DispatchTensorKind::NdArray(tensor) => tensor.autodiff().backward(), + #[cfg(feature = "tch")] + DispatchTensorKind::LibTorch(tensor) => tensor.autodiff().backward(), + #[cfg(feature = "remote")] + DispatchTensorKind::Remote(tensor) => tensor.autodiff().backward(), + DispatchTensorKind::Autodiff(_) => { + panic!("Autodiff should not wrap an autodiff tensor.") + } + }, + _ => panic!("Requires autodiff tensor."), + } + } + + fn grad(tensor: &DispatchTensor, grads: &Self::Gradients) -> Option { + let DispatchTensor { + kind, + checkpointing, + } = tensor; + let grad: Option = match &kind { + DispatchTensorKind::Autodiff(inner_kind) => match &**inner_kind { + #[cfg(feature = "cpu")] + DispatchTensorKind::Cpu(tensor) => tensor + .as_autodiff() + .grad(grads) + .map(|t| DispatchTensorKind::Cpu(crate::BackendTensor::Float(t))), + #[cfg(feature = "cuda")] + DispatchTensorKind::Cuda(tensor) => tensor + .as_autodiff() + .grad(grads) + .map(|t| DispatchTensorKind::Cuda(crate::BackendTensor::Float(t))), + #[cfg(feature = "metal")] + DispatchTensorKind::Metal(tensor) => tensor + .as_autodiff() + .grad(grads) + .map(|t| DispatchTensorKind::Metal(crate::BackendTensor::Float(t))), + #[cfg(feature = "rocm")] + DispatchTensorKind::Rocm(tensor) => tensor + .as_autodiff() + .grad(grads) + .map(|t| DispatchTensorKind::Rocm(crate::BackendTensor::Float(t))), + #[cfg(feature = "vulkan")] + DispatchTensorKind::Vulkan(tensor) => tensor + .as_autodiff() + .grad(grads) + .map(|t| DispatchTensorKind::Vulkan(crate::BackendTensor::Float(t))), + #[cfg(feature = "wgpu")] + DispatchTensorKind::Wgpu(tensor) => tensor + .as_autodiff() + .grad(grads) + .map(|t| DispatchTensorKind::Wgpu(crate::BackendTensor::Float(t))), + #[cfg(feature = "webgpu")] + DispatchTensorKind::WebGpu(tensor) => tensor + .as_autodiff() + .grad(grads) + .map(|t| DispatchTensorKind::WebGpu(crate::BackendTensor::Float(t))), + #[cfg(any(feature = "flex", default_backend))] + DispatchTensorKind::Flex(tensor) => tensor + .as_autodiff() + .grad(grads) + .map(|t| DispatchTensorKind::Flex(crate::BackendTensor::Float(t))), + #[cfg(feature = "ndarray")] + DispatchTensorKind::NdArray(tensor) => tensor + .as_autodiff() + .grad(grads) + .map(|t| DispatchTensorKind::NdArray(crate::BackendTensor::Float(t))), + #[cfg(feature = "tch")] + DispatchTensorKind::LibTorch(tensor) => tensor + .as_autodiff() + .grad(grads) + .map(|t| DispatchTensorKind::LibTorch(crate::BackendTensor::Float(t))), + #[cfg(feature = "remote")] + DispatchTensorKind::Remote(tensor) => tensor + .as_autodiff() + .grad(grads) + .map(|t| DispatchTensorKind::Remote(crate::BackendTensor::Float(t))), + DispatchTensorKind::Autodiff(_) => { + panic!("Autodiff should not wrap an autodiff tensor.") + } + }, + _ => panic!("Requires autodiff tensor."), + }; + grad.map(|kind| DispatchTensor { + kind, + checkpointing: *checkpointing, + }) + } + + fn grad_remove(tensor: &DispatchTensor, grads: &mut Self::Gradients) -> Option { + let DispatchTensor { + kind, + checkpointing, + } = tensor; + let grad: Option = match &kind { + DispatchTensorKind::Autodiff(inner_kind) => match &**inner_kind { + #[cfg(feature = "cpu")] + DispatchTensorKind::Cpu(tensor) => tensor + .as_autodiff() + .grad_remove(grads) + .map(|t| DispatchTensorKind::Cpu(crate::BackendTensor::Float(t))), + #[cfg(feature = "cuda")] + DispatchTensorKind::Cuda(tensor) => tensor + .as_autodiff() + .grad_remove(grads) + .map(|t| DispatchTensorKind::Cuda(crate::BackendTensor::Float(t))), + #[cfg(feature = "metal")] + DispatchTensorKind::Metal(tensor) => tensor + .as_autodiff() + .grad_remove(grads) + .map(|t| DispatchTensorKind::Metal(crate::BackendTensor::Float(t))), + #[cfg(feature = "rocm")] + DispatchTensorKind::Rocm(tensor) => tensor + .as_autodiff() + .grad_remove(grads) + .map(|t| DispatchTensorKind::Rocm(crate::BackendTensor::Float(t))), + #[cfg(feature = "vulkan")] + DispatchTensorKind::Vulkan(tensor) => tensor + .as_autodiff() + .grad_remove(grads) + .map(|t| DispatchTensorKind::Vulkan(crate::BackendTensor::Float(t))), + #[cfg(feature = "wgpu")] + DispatchTensorKind::Wgpu(tensor) => tensor + .as_autodiff() + .grad_remove(grads) + .map(|t| DispatchTensorKind::Wgpu(crate::BackendTensor::Float(t))), + #[cfg(feature = "webgpu")] + DispatchTensorKind::WebGpu(tensor) => tensor + .as_autodiff() + .grad_remove(grads) + .map(|t| DispatchTensorKind::WebGpu(crate::BackendTensor::Float(t))), + #[cfg(any(feature = "flex", default_backend))] + DispatchTensorKind::Flex(tensor) => tensor + .as_autodiff() + .grad_remove(grads) + .map(|t| DispatchTensorKind::Flex(crate::BackendTensor::Float(t))), + #[cfg(feature = "ndarray")] + DispatchTensorKind::NdArray(tensor) => tensor + .as_autodiff() + .grad_remove(grads) + .map(|t| DispatchTensorKind::NdArray(crate::BackendTensor::Float(t))), + #[cfg(feature = "tch")] + DispatchTensorKind::LibTorch(tensor) => tensor + .as_autodiff() + .grad_remove(grads) + .map(|t| DispatchTensorKind::LibTorch(crate::BackendTensor::Float(t))), + #[cfg(feature = "remote")] + DispatchTensorKind::Remote(tensor) => tensor + .as_autodiff() + .grad_remove(grads) + .map(|t| DispatchTensorKind::Remote(crate::BackendTensor::Float(t))), + DispatchTensorKind::Autodiff(_) => { + panic!("Autodiff should not wrap an autodiff tensor.") + } + }, + _ => panic!("Requires autodiff tensor."), + }; + grad.map(|kind| DispatchTensor { + kind, + checkpointing: *checkpointing, + }) + } + + fn grad_replace(tensor: &DispatchTensor, grads: &mut Self::Gradients, grad: DispatchTensor) { + let DispatchTensor { + kind, + checkpointing, + } = tensor; + let DispatchTensor { + kind: grad, + checkpointing: grad_ckp, + } = grad; + debug_assert_eq!(checkpointing, &grad_ckp); + + match &kind { + DispatchTensorKind::Autodiff(inner_kind) => match (&**inner_kind, grad) { + #[cfg(feature = "cpu")] + (DispatchTensorKind::Cpu(tensor), DispatchTensorKind::Cpu(grad)) => { + tensor.as_autodiff().grad_replace(grads, grad.float()) + } + #[cfg(feature = "cuda")] + (DispatchTensorKind::Cuda(tensor), DispatchTensorKind::Cuda(grad)) => { + tensor.as_autodiff().grad_replace(grads, grad.float()) + } + #[cfg(feature = "metal")] + (DispatchTensorKind::Metal(tensor), DispatchTensorKind::Metal(grad)) => { + tensor.as_autodiff().grad_replace(grads, grad.float()) + } + #[cfg(feature = "rocm")] + (DispatchTensorKind::Rocm(tensor), DispatchTensorKind::Rocm(grad)) => { + tensor.as_autodiff().grad_replace(grads, grad.float()) + } + #[cfg(feature = "vulkan")] + (DispatchTensorKind::Vulkan(tensor), DispatchTensorKind::Vulkan(grad)) => { + tensor.as_autodiff().grad_replace(grads, grad.float()) + } + #[cfg(feature = "wgpu")] + (DispatchTensorKind::Wgpu(tensor), DispatchTensorKind::Wgpu(grad)) => { + tensor.as_autodiff().grad_replace(grads, grad.float()) + } + #[cfg(feature = "webgpu")] + (DispatchTensorKind::WebGpu(tensor), DispatchTensorKind::WebGpu(grad)) => { + tensor.as_autodiff().grad_replace(grads, grad.float()) + } + #[cfg(any(feature = "flex", default_backend))] + (DispatchTensorKind::Flex(tensor), DispatchTensorKind::Flex(grad)) => { + tensor.as_autodiff().grad_replace(grads, grad.float()) + } + #[cfg(feature = "ndarray")] + (DispatchTensorKind::NdArray(tensor), DispatchTensorKind::NdArray(grad)) => { + tensor.as_autodiff().grad_replace(grads, grad.float()) + } + #[cfg(feature = "remote")] + (DispatchTensorKind::Remote(tensor), DispatchTensorKind::Remote(grad)) => { + tensor.as_autodiff().grad_replace(grads, grad.float()) + } + (DispatchTensorKind::Autodiff(_), _) => { + panic!("Autodiff should not wrap an autodiff tensor.") + } + // TODO: distributed message? + (t, g) => panic!( + "The provided tensors are not on the same backend. Got backends {t:?} and {g:?}." + ), + }, + _ => panic!("Requires autodiff tensor."), + } + } + + fn inner(tensor: DispatchTensor) -> DispatchTensor { + let DispatchTensor { + kind, + checkpointing: _, + } = tensor; + + let kind = match kind { + DispatchTensorKind::Autodiff(inner_kind) => match *inner_kind { + #[cfg(feature = "cpu")] + DispatchTensorKind::Cpu(tensor) => DispatchTensorKind::Cpu( + crate::BackendTensor::Float(tensor.autodiff().primitive), + ), + #[cfg(feature = "cuda")] + DispatchTensorKind::Cuda(tensor) => DispatchTensorKind::Cuda( + crate::BackendTensor::Float(tensor.autodiff().primitive), + ), + #[cfg(feature = "metal")] + DispatchTensorKind::Metal(tensor) => DispatchTensorKind::Metal( + crate::BackendTensor::Float(tensor.autodiff().primitive), + ), + #[cfg(feature = "rocm")] + DispatchTensorKind::Rocm(tensor) => DispatchTensorKind::Rocm( + crate::BackendTensor::Float(tensor.autodiff().primitive), + ), + #[cfg(feature = "vulkan")] + DispatchTensorKind::Vulkan(tensor) => DispatchTensorKind::Vulkan( + crate::BackendTensor::Float(tensor.autodiff().primitive), + ), + #[cfg(feature = "wgpu")] + DispatchTensorKind::Wgpu(tensor) => DispatchTensorKind::Wgpu( + crate::BackendTensor::Float(tensor.autodiff().primitive), + ), + #[cfg(feature = "webgpu")] + DispatchTensorKind::WebGpu(tensor) => DispatchTensorKind::WebGpu( + crate::BackendTensor::Float(tensor.autodiff().primitive), + ), + #[cfg(any(feature = "flex", default_backend))] + DispatchTensorKind::Flex(tensor) => DispatchTensorKind::Flex( + crate::BackendTensor::Float(tensor.autodiff().primitive), + ), + #[cfg(feature = "ndarray")] + DispatchTensorKind::NdArray(tensor) => DispatchTensorKind::NdArray( + crate::BackendTensor::Float(tensor.autodiff().primitive), + ), + #[cfg(feature = "tch")] + DispatchTensorKind::LibTorch(tensor) => DispatchTensorKind::LibTorch( + crate::BackendTensor::Float(tensor.autodiff().primitive), + ), + #[cfg(feature = "remote")] + DispatchTensorKind::Remote(tensor) => DispatchTensorKind::Remote( + crate::BackendTensor::Float(tensor.autodiff().primitive), + ), + DispatchTensorKind::Autodiff(_) => { + panic!("Autodiff should not wrap an autodiff tensor.") + } + }, + _ => panic!("Requires autodiff tensor."), + }; + DispatchTensor { + kind, + checkpointing: None, + } + } + + fn int_inner(tensor: DispatchTensor) -> DispatchTensor { + tensor + } + + fn bool_inner(tensor: DispatchTensor) -> DispatchTensor { + tensor + } + + fn q_inner(tensor: DispatchTensor) -> DispatchTensor { + tensor + } + + fn from_inner(tensor: DispatchTensor) -> DispatchTensor { + let DispatchTensor { + kind, + checkpointing, + } = tensor; + + let kind = match kind { + #[cfg(feature = "cpu")] + DispatchTensorKind::Cpu(tensor) => { + DispatchTensorKind::Autodiff(Box::new(DispatchTensorKind::Cpu( + crate::BackendTensor::Autodiff(Autodiff::::from_inner(tensor.float())), + ))) + } + #[cfg(feature = "cuda")] + DispatchTensorKind::Cuda(tensor) => { + DispatchTensorKind::Autodiff(Box::new(DispatchTensorKind::Cuda( + crate::BackendTensor::Autodiff(Autodiff::::from_inner(tensor.float())), + ))) + } + #[cfg(feature = "metal")] + DispatchTensorKind::Metal(tensor) => { + DispatchTensorKind::Autodiff(Box::new(DispatchTensorKind::Metal( + crate::BackendTensor::Autodiff(Autodiff::::from_inner(tensor.float())), + ))) + } + #[cfg(feature = "rocm")] + DispatchTensorKind::Rocm(tensor) => { + DispatchTensorKind::Autodiff(Box::new(DispatchTensorKind::Rocm( + crate::BackendTensor::Autodiff(Autodiff::::from_inner(tensor.float())), + ))) + } + #[cfg(feature = "vulkan")] + DispatchTensorKind::Vulkan(tensor) => { + DispatchTensorKind::Autodiff(Box::new(DispatchTensorKind::Vulkan( + crate::BackendTensor::Autodiff(Autodiff::::from_inner(tensor.float())), + ))) + } + #[cfg(feature = "wgpu")] + DispatchTensorKind::Wgpu(tensor) => { + DispatchTensorKind::Autodiff(Box::new(DispatchTensorKind::Wgpu( + crate::BackendTensor::Autodiff(Autodiff::::from_inner(tensor.float())), + ))) + } + #[cfg(feature = "webgpu")] + DispatchTensorKind::WebGpu(tensor) => { + DispatchTensorKind::Autodiff(Box::new(DispatchTensorKind::WebGpu( + crate::BackendTensor::Autodiff(Autodiff::::from_inner(tensor.float())), + ))) + } + #[cfg(any(feature = "flex", default_backend))] + DispatchTensorKind::Flex(tensor) => { + DispatchTensorKind::Autodiff(Box::new(DispatchTensorKind::Flex( + crate::BackendTensor::Autodiff(Autodiff::::from_inner(tensor.float())), + ))) + } + #[cfg(feature = "ndarray")] + DispatchTensorKind::NdArray(tensor) => { + DispatchTensorKind::Autodiff(Box::new(DispatchTensorKind::NdArray( + crate::BackendTensor::Autodiff(Autodiff::::from_inner(tensor.float())), + ))) + } + #[cfg(feature = "tch")] + DispatchTensorKind::LibTorch(tensor) => DispatchTensorKind::Autodiff(Box::new( + DispatchTensorKind::LibTorch(crate::BackendTensor::Autodiff( + Autodiff::::from_inner(tensor.float()), + )), + )), + #[cfg(feature = "remote")] + DispatchTensorKind::Remote(tensor) => { + DispatchTensorKind::Autodiff(Box::new(DispatchTensorKind::Remote( + crate::BackendTensor::Autodiff(Autodiff::::from_inner(tensor.float())), + ))) + } + DispatchTensorKind::Autodiff(_) => { + panic!("Autodiff should not wrap an autodiff tensor.") + } + }; + + // TODO: should use C::STRATEGY + let checkpointing = if let Some(strategy) = checkpointing { + Some(strategy) + } else { + Some(crate::CheckpointingStrategy::None) + }; + DispatchTensor { + kind, + checkpointing, + } + } + + fn int_from_inner(tensor: DispatchTensor) -> DispatchTensor { + tensor + } + + fn bool_from_inner(tensor: DispatchTensor) -> DispatchTensor { + tensor + } + + fn q_from_inner(tensor: DispatchTensor) -> DispatchTensor { + tensor + } + + // Only the collective-capable backends (Cuda/Remote) carry distributed params; in builds + // without them the match arms cfg out, leaving the bindings unused and the tail unreachable. + #[allow(unused_variables, unreachable_code)] + fn set_distributed_params( + tensor: DispatchTensor, + param_id: DistributedParamId, + ) -> DispatchTensor { + let DispatchTensor { + kind, + checkpointing, + } = tensor; + + let kind = match kind { + DispatchTensorKind::Autodiff(inner_kind) => match *inner_kind { + #[cfg(feature = "cuda")] + DispatchTensorKind::Cuda(tensor) => { + DispatchTensorKind::Autodiff(Box::new(DispatchTensorKind::Cuda( + crate::BackendTensor::Autodiff(Autodiff::::set_distributed_params( + tensor.as_autodiff().clone(), + param_id, + )), + ))) + } + #[cfg(feature = "remote")] + DispatchTensorKind::Remote(tensor) => { + DispatchTensorKind::Autodiff(Box::new(DispatchTensorKind::Remote( + crate::BackendTensor::Autodiff(Autodiff::::set_distributed_params( + tensor.as_autodiff().clone(), + param_id, + )), + ))) + } + DispatchTensorKind::Autodiff(_) => { + panic!("Autodiff should not wrap an autodiff tensor.") + } + other => { + panic!("Distributed operations are not supported for tensor kind {other:?}") + } + }, + _ => panic!("Requires autodiff tensor."), + }; + + let checkpointing = if let Some(strategy) = checkpointing { + Some(strategy) + } else { + Some(crate::CheckpointingStrategy::None) + }; + DispatchTensor { + kind, + checkpointing, + } + } + + #[allow(unused_variables)] + fn distributed_params(tensor: &DispatchTensor) -> Option { + let DispatchTensor { + kind, + checkpointing: _, + } = tensor; + + match &kind { + DispatchTensorKind::Autodiff(inner_kind) => match &**inner_kind { + #[cfg(feature = "cuda")] + DispatchTensorKind::Cuda(tensor) => { + tensor.as_autodiff().node.distributed_params.clone() + } + #[cfg(feature = "remote")] + DispatchTensorKind::Remote(tensor) => { + tensor.as_autodiff().node.distributed_params.clone() + } + + DispatchTensorKind::Autodiff(_) => { + panic!("Autodiff should not wrap an autodiff tensor.") + } + // Backends without distributed support never carry distributed params. + _ => None, + }, + _ => panic!("Requires autodiff tensor."), + } + } + + #[allow(unused_variables)] + fn is_distributed(tensor: &DispatchTensor) -> bool { + let DispatchTensor { + kind, + checkpointing: _, + } = tensor; + + match &kind { + DispatchTensorKind::Autodiff(inner_kind) => match &**inner_kind { + #[cfg(feature = "cuda")] + DispatchTensorKind::Cuda(tensor) => { + tensor.as_autodiff().node.distributed_params.is_some() + } + #[cfg(feature = "remote")] + DispatchTensorKind::Remote(tensor) => { + tensor.as_autodiff().node.distributed_params.is_some() + } + + DispatchTensorKind::Autodiff(_) => { + panic!("Autodiff should not wrap an autodiff tensor.") + } + // Backends without distributed support are never distributed. + _ => false, + }, + _ => panic!("Requires autodiff tensor."), + } + } +} + +// NOTE: placeholder for autodiff module requirements +#[cfg(not(feature = "autodiff"))] +impl AutodiffBackend for Dispatch { + type InnerBackend = Dispatch; + + type Gradients = bool; + + fn backward(_tensor: DispatchTensor) -> Self::Gradients { + unimplemented!("Requires `autodiff` feature") + } + + fn grad(_tensor: &DispatchTensor, _grads: &Self::Gradients) -> Option { + unimplemented!("Requires `autodiff` feature") + } + + fn grad_remove( + _tensor: &DispatchTensor, + _grads: &mut Self::Gradients, + ) -> Option { + unimplemented!("Requires `autodiff` feature") + } + + fn grad_replace(_tensor: &DispatchTensor, _grads: &mut Self::Gradients, _grad: DispatchTensor) { + unimplemented!("Requires `autodiff` feature") + } + + fn inner(_tensor: DispatchTensor) -> DispatchTensor { + unimplemented!("Requires `autodiff` feature") + } + + fn int_inner(_tensor: DispatchTensor) -> DispatchTensor { + unimplemented!("Requires `autodiff` feature") + } + + fn bool_inner(_tensor: DispatchTensor) -> DispatchTensor { + unimplemented!("Requires `autodiff` feature") + } + + fn q_inner(_tensor: DispatchTensor) -> DispatchTensor { + unimplemented!("Requires `autodiff` feature") + } + + fn from_inner(_tensor: DispatchTensor) -> DispatchTensor { + unimplemented!("Requires `autodiff` feature") + } + + fn int_from_inner(_tensor: DispatchTensor) -> DispatchTensor { + unimplemented!("Requires `autodiff` feature") + } + + fn bool_from_inner(_tensor: DispatchTensor) -> DispatchTensor { + unimplemented!("Requires `autodiff` feature") + } + + fn q_from_inner(_tensor: DispatchTensor) -> DispatchTensor { + unimplemented!("Requires `autodiff` feature") + } +} + +impl Dispatch { + /// List all available devices of the specified [type id](DispatchDeviceId). + pub fn enumerate(type_id: DispatchDeviceId) -> Vec { + // TODO: right now this assumes `type_id = 0`, but WgpuDevice and LibTorchDevice have other types. + match type_id { + #[cfg(feature = "cpu")] + DispatchDeviceId::Cpu => vec![CpuDevice.into()], + #[cfg(feature = "cuda")] + DispatchDeviceId::Cuda => (0..Cuda::device_count(0)) + .map(|i| CudaDevice::new(i).into()) + .collect(), + #[cfg(feature = "metal")] + DispatchDeviceId::Metal => (0..Metal::device_count(0)) + .map(|i| DispatchDevice::Metal(WgpuDevice::DiscreteGpu(i))) + .collect(), + #[cfg(feature = "rocm")] + DispatchDeviceId::Rocm => (0..Rocm::device_count(0)) + .map(|i| RocmDevice::new(i).into()) + .collect(), + #[cfg(feature = "vulkan")] + DispatchDeviceId::Vulkan => (0..Vulkan::device_count(0)) + .map(|i| DispatchDevice::Vulkan(WgpuDevice::DiscreteGpu(i))) + .collect(), + #[cfg(feature = "wgpu")] + DispatchDeviceId::Wgpu => (0..Wgpu::device_count(0)) + .map(|i| DispatchDevice::Wgpu(WgpuDevice::DiscreteGpu(i))) + .collect(), + #[cfg(feature = "webgpu")] + DispatchDeviceId::WebGpu => (0..WebGpu::device_count(0)) + .map(|i| DispatchDevice::WebGpu(WgpuDevice::DiscreteGpu(i))) + .collect(), + #[cfg(any(feature = "flex", default_backend))] + DispatchDeviceId::Flex => vec![FlexDevice.into()], + #[cfg(feature = "ndarray")] + DispatchDeviceId::NdArray => vec![NdArrayDevice::Cpu.into()], + #[cfg(feature = "tch")] + DispatchDeviceId::LibTorch => (0..LibTorch::device_count(0)) + .map(|i| LibTorchDevice::Cuda(i).into()) + .collect(), + #[cfg(feature = "remote")] + // Remote devices are keyed by a network address, which the type-id-only + // `enumerate` can't carry. Use [`Dispatch::enumerate_remote_websocket`] to list the devices + // behind a given address. + DispatchDeviceId::Remote => Vec::new(), + _ => unreachable!("No backend feature enabled."), + } + } + + /// List every device hosted by the remote server at `address`. + /// + /// Unlike [`enumerate`](Self::enumerate), remote devices are identified by a network + /// address rather than enumerable local hardware, so they need a dedicated entry point. + /// Connecting to the server (required to learn its device count) happens here; see + /// [`RemoteDevice::enumerate_websocket`]. + /// + /// Websocket-only: Iroh peers are addressed by endpoint identity, not a URL string. + #[cfg(feature = "remote-websocket")] + pub fn enumerate_remote_websocket(address: &str) -> Vec { + RemoteDevice::enumerate_websocket(address) + .into_iter() + .map(DispatchDevice::Remote) + .collect() + } +} diff --git a/crates/burn-dispatch/src/device.rs b/crates/burn-dispatch/src/device.rs new file mode 100644 index 0000000..ca0722d --- /dev/null +++ b/crates/burn-dispatch/src/device.rs @@ -0,0 +1,654 @@ +use burn_backend::{DeviceId, DeviceOps, DeviceSettings}; + +use crate::backends::*; + +#[cfg(feature = "autodiff")] +use alloc::boxed::Box; + +/// Represents a device for the [`Dispatch`](crate::Dispatch). +/// +/// Each variant corresponds to a backend that the [`Dispatch`](crate::Dispatch) can dispatch operations to. +/// +/// # Example +/// +/// ```ignore +/// use burn::DispatchDevice; +/// +/// #[cfg(feature = "cpu")] +/// let cpu_device = DispatchDevice::Cpu(Default::default()); +/// +/// #[cfg(feature = "cuda")] +/// let cuda_device = DispatchDevice::Cuda(Default::default()); +/// ``` +#[derive(Clone, Eq)] +pub enum DispatchDevice { + /// The [CPU backend](Cpu) device. + #[cfg(feature = "cpu")] + Cpu(CpuDevice), + + /// The [CUDA backend](Cuda) device. + #[cfg(feature = "cuda")] + Cuda(CudaDevice), + + /// The [Metal backend](Metal) device (via WGPU runtime). + #[cfg(feature = "metal")] + Metal(WgpuDevice), + + /// The [ROCm backend](Rocm) device. + #[cfg(feature = "rocm")] + Rocm(RocmDevice), + + /// The [Vulkan backend](Vulkan) device. + #[cfg(feature = "vulkan")] + Vulkan(WgpuDevice), + + /// The [Wgpu backend](Wgpu) device (via WGPU runtime with auto-selected compiler). + #[cfg(feature = "wgpu")] + Wgpu(WgpuDevice), + + /// The [WebGPU backend](WebGpu) device (via WGPU runtime). + #[cfg(feature = "webgpu")] + WebGpu(WgpuDevice), + + /// The [Flex backend](Flex) device (CPU-only). + #[cfg(any(feature = "flex", default_backend))] + Flex(FlexDevice), + + /// The [NdArray backend](NdArray) device (CPU-only). + #[cfg(feature = "ndarray")] + NdArray(NdArrayDevice), + + /// The [LibTorch backend](LibTorch) device. + #[cfg(feature = "tch")] + LibTorch(LibTorchDevice), + + /// The [remote backend](Remote) device, identified by a network address. + #[cfg(feature = "remote")] + Remote(RemoteDevice), + + /// The [autodiff enabled backend](Autodiff) device. + #[cfg(feature = "autodiff")] + Autodiff(AutodiffDevice), +} + +#[cfg(feature = "autodiff")] +// This tuple struct mainly restricts users from creating Autodiff(Autodiff) devices. +/// A wrapper that enables automatic differentiation for a [`DispatchDevice`]. +/// +/// Use [`DispatchDevice::autodiff`] to construct this type. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AutodiffDevice { + pub(crate) inner: Box, + pub(crate) checkpointing: CheckpointingStrategy, +} + +#[cfg(feature = "autodiff")] +impl AutodiffDevice { + pub(crate) fn new(device: DispatchDevice, checkpointing: CheckpointingStrategy) -> Self { + Self { + inner: Box::new(device), + checkpointing, + } + } + + /// Returns the underlying device, removing the autodiff capability. + pub fn inner(self) -> DispatchDevice { + *self.inner + } +} + +#[cfg(feature = "autodiff")] +// Useful for match in dispatch macros +impl core::ops::Deref for AutodiffDevice { + type Target = DispatchDevice; + + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +#[allow(missing_docs)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +/// Checkpointing strategy for autodiff. +#[repr(u8)] +pub enum CheckpointingStrategy { + Balanced, + #[default] + None, +} + +#[cfg(feature = "autodiff")] +pub(crate) fn validate_checkpointing( + lhs: Option, + rhs: Option, +) -> Option { + match (lhs, rhs) { + (Some(lhs), Some(rhs)) => { + assert_eq!( + lhs, rhs, + "Autodiff strategy mismatch: {lhs:?} vs {rhs:?}. Tensors in the same operation must share a strategy." + ); + Some(lhs) + } + (None, None) => None, + // When tensors are created on non-autodiff device there is no checkpointing, but + // tensor created with autodiff which moved out (`tensor.inner()`) will still carry the state. + // In such cases, we can "promote" the checkpointing. + (None, rhs) => rhs, + (lhs, None) => lhs, + } +} + +impl core::fmt::Debug for DispatchDevice { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + #[cfg(feature = "cpu")] + Self::Cpu(device) => f.debug_tuple("Cpu").field(device).finish(), + #[cfg(feature = "cuda")] + Self::Cuda(device) => f.debug_tuple("Cuda").field(device).finish(), + #[cfg(feature = "metal")] + Self::Metal(device) => f.debug_tuple("Metal").field(device).finish(), + #[cfg(feature = "rocm")] + Self::Rocm(device) => f.debug_tuple("Rocm").field(device).finish(), + #[cfg(feature = "vulkan")] + Self::Vulkan(device) => f.debug_tuple("Vulkan").field(device).finish(), + #[cfg(feature = "wgpu")] + Self::Wgpu(device) => f.debug_tuple("Wgpu").field(device).finish(), + #[cfg(feature = "webgpu")] + Self::WebGpu(device) => f.debug_tuple("WebGpu").field(device).finish(), + #[cfg(any(feature = "flex", default_backend))] + Self::Flex(device) => f.debug_tuple("Flex").field(device).finish(), + #[cfg(feature = "ndarray")] + Self::NdArray(device) => f.debug_tuple("NdArray").field(device).finish(), + #[cfg(feature = "tch")] + Self::LibTorch(device) => f.debug_tuple("LibTorch").field(device).finish(), + #[cfg(feature = "remote")] + Self::Remote(device) => f.debug_tuple("Remote").field(device).finish(), + #[cfg(feature = "autodiff")] + // Format without `AutodiffDevice` wrapper + Self::Autodiff(device) => f.debug_tuple("Autodiff").field(&device.inner).finish(), + } + } +} + +impl Default for DispatchDevice { + #[allow(unreachable_code)] + fn default() -> Self { + // TODO: which priority? + // Single override e.g. `BURN_DEVICE=vulkan` forces Vulkan or panics if not available. + // Priority list e.g. `BURN_DEVICE_PRIORITY=cuda,vulkan,cpu` sets the order. + // Both could be tied into `burn.toml` config + // For now we just use `BURN_DEVICE` on CI to force a single device + + #[cfg(feature = "std")] + { + if let Ok(device_str) = std::env::var("BURN_DEVICE") { + match device_str.to_lowercase().as_str() { + "cuda" => { + #[cfg(feature = "cuda")] + return Self::Cuda(CudaDevice::default()); + panic!( + "BURN_DEVICE=cuda requested, but the 'cuda' feature is not enabled." + ); + } + "metal" => { + #[cfg(feature = "metal")] + return Self::Metal(burn_wgpu::WgpuDevice::default()); + panic!( + "BURN_DEVICE=metal requested, but the 'metal' feature is not enabled." + ); + } + "rocm" => { + #[cfg(feature = "rocm")] + return Self::Rocm(RocmDevice::default()); + panic!( + "BURN_DEVICE=rocm requested, but the 'rocm' feature is not enabled." + ); + } + "vulkan" => { + #[cfg(feature = "vulkan")] + return Self::Vulkan(burn_wgpu::WgpuDevice::default()); + panic!( + "BURN_DEVICE=vulkan requested, but the 'vulkan' feature is not enabled." + ); + } + "webgpu" => { + #[cfg(feature = "webgpu")] + return Self::WebGpu(burn_wgpu::WgpuDevice::default()); + panic!( + "BURN_DEVICE=webgpu requested, but the 'webgpu' feature is not enabled." + ); + } + "wgpu" => { + #[cfg(feature = "wgpu")] + return Self::Wgpu(burn_wgpu::WgpuDevice::default()); + panic!( + "BURN_DEVICE=wgpu requested, but the 'wgpu' feature is not enabled." + ); + } + "cpu" => { + #[cfg(feature = "cpu")] + return Self::Cpu(CpuDevice); + panic!("BURN_DEVICE=cpu requested, but the 'cpu' feature is not enabled."); + } + "tch" => { + #[cfg(feature = "tch")] + return Self::LibTorch(LibTorchDevice::default()); + panic!("BURN_DEVICE=tch requested, but the 'tch' feature is not enabled."); + } + "remote" => { + #[cfg(feature = "remote")] + return Self::Remote(RemoteDevice::default()); + panic!( + "BURN_DEVICE=remote requested, but the 'remote' feature is not enabled." + ); + } + "flex" => { + #[cfg(any(feature = "flex", default_backend))] + return Self::Flex(FlexDevice); + panic!( + "BURN_DEVICE=flex requested, but the 'flex' feature is not enabled." + ); + } + "ndarray" => { + #[cfg(feature = "ndarray")] + return Self::NdArray(NdArrayDevice::default()); + panic!( + "BURN_DEVICE=ndarray requested, but the 'ndarray' feature is not enabled." + ); + } + _ => panic!("Unknown BURN_DEVICE override: '{}'.", device_str), + } + } + } + + #[cfg(feature = "cuda")] + return Self::Cuda(CudaDevice::default()); + + #[cfg(feature = "metal")] + return Self::Metal(burn_wgpu::WgpuDevice::default()); + + #[cfg(feature = "rocm")] + return Self::Rocm(RocmDevice::default()); + + #[cfg(feature = "vulkan")] + return Self::Vulkan(burn_wgpu::WgpuDevice::default()); + + #[cfg(feature = "webgpu")] + return Self::WebGpu(burn_wgpu::WgpuDevice::default()); + + #[cfg(feature = "wgpu")] + return Self::Wgpu(burn_wgpu::WgpuDevice::default()); + + #[cfg(feature = "cpu")] + return Self::Cpu(CpuDevice); + + #[cfg(feature = "tch")] + return Self::LibTorch(LibTorchDevice::default()); + + // Prefer Flex over NdArray when both are enabled: Flex is the long-term + // CPU backend replacement and should win the default tie. + #[cfg(any(feature = "flex", default_backend))] + return Self::Flex(FlexDevice); + + #[cfg(feature = "remote")] + return Self::Remote(RemoteDevice::default()); + + #[cfg(feature = "ndarray")] + return Self::NdArray(NdArrayDevice::default()); + } +} + +impl PartialEq for DispatchDevice { + /// Compares devices based on hardware identity. + /// + /// Returns `true` if both devices represent the same compute resource. + /// Note that this comparison ignores autodiff and checkpointing settings. + fn eq(&self, other: &Self) -> bool { + match (self, other) { + // If both are Autodiff, compare the inner devices + #[cfg(feature = "autodiff")] + (DispatchDevice::Autodiff(a), DispatchDevice::Autodiff(b)) => a == b, + // If one is Autodiff, compare it to the raw device + #[cfg(feature = "autodiff")] + (DispatchDevice::Autodiff(a), b) => a.inner.as_ref() == b, + #[cfg(feature = "autodiff")] + (a, DispatchDevice::Autodiff(b)) => a == b.inner.as_ref(), + #[cfg(feature = "cpu")] + (Self::Cpu(a), Self::Cpu(b)) => a == b, + #[cfg(feature = "cuda")] + (Self::Cuda(a), Self::Cuda(b)) => a == b, + #[cfg(feature = "metal")] + (Self::Metal(a), Self::Metal(b)) => a == b, + #[cfg(feature = "rocm")] + (Self::Rocm(a), Self::Rocm(b)) => a == b, + #[cfg(feature = "vulkan")] + (Self::Vulkan(a), Self::Vulkan(b)) => a == b, + #[cfg(feature = "wgpu")] + (Self::Wgpu(a), Self::Wgpu(b)) => a == b, + #[cfg(feature = "webgpu")] + (Self::WebGpu(a), Self::WebGpu(b)) => a == b, + #[cfg(any(feature = "flex", default_backend))] + (Self::Flex(a), Self::Flex(b)) => a == b, + #[cfg(feature = "ndarray")] + (Self::NdArray(a), Self::NdArray(b)) => a == b, + #[cfg(feature = "tch")] + (Self::LibTorch(a), Self::LibTorch(b)) => a == b, + #[cfg(feature = "remote")] + (Self::Remote(a), Self::Remote(b)) => a == b, + #[allow(unreachable_patterns)] + (_, _) => false, + } + } +} + +const INTERNAL_ID_MASK: u16 = 0x00FF; +const BACKEND_SHIFT: u32 = 8; + +impl DispatchDevice { + #[cfg(feature = "autodiff")] + /// Creates a new [`DispatchDevice`] with [automatic differentiation](Autodiff) enabled. + pub fn autodiff(device: impl Into) -> DispatchDevice { + Self::autodiff_checkpointed(device, CheckpointingStrategy::None) + } + #[cfg(feature = "autodiff")] + /// Creates a new [`DispatchDevice`] with [automatic differentiation](Autodiff) enabled. + pub fn autodiff_checkpointed( + device: impl Into, + checkpointing: CheckpointingStrategy, + ) -> DispatchDevice { + let device = device.into(); + DispatchDevice::Autodiff(AutodiffDevice::new(device, checkpointing)) + } + + /// Returns the inner device, without autodiff (when enabled). + pub fn inner(self) -> Self { + #[cfg(feature = "autodiff")] + if let DispatchDevice::Autodiff(device) = self { + return *device.inner; + } + + self + } + + /// Returns a unique number per variant to encode into type_id. + fn backend_id(&self) -> DispatchDeviceId { + match self { + #[cfg(feature = "cpu")] + Self::Cpu(_) => DispatchDeviceId::Cpu, + #[cfg(feature = "cuda")] + Self::Cuda(_) => DispatchDeviceId::Cuda, + #[cfg(feature = "metal")] + Self::Metal(_) => DispatchDeviceId::Metal, + #[cfg(feature = "rocm")] + Self::Rocm(_) => DispatchDeviceId::Rocm, + #[cfg(feature = "vulkan")] + Self::Vulkan(_) => DispatchDeviceId::Vulkan, + #[cfg(feature = "wgpu")] + Self::Wgpu(_) => DispatchDeviceId::Wgpu, + #[cfg(feature = "webgpu")] + Self::WebGpu(_) => DispatchDeviceId::WebGpu, + #[cfg(any(feature = "flex", default_backend))] + Self::Flex(_) => DispatchDeviceId::Flex, + #[cfg(feature = "ndarray")] + Self::NdArray(_) => DispatchDeviceId::NdArray, + #[cfg(feature = "tch")] + Self::LibTorch(_) => DispatchDeviceId::LibTorch, + #[cfg(feature = "remote")] + Self::Remote(_) => DispatchDeviceId::Remote, + #[cfg(feature = "autodiff")] + Self::Autodiff(device) => device.inner.backend_id(), + } + } + + /// Encode variant ID and backend type ID into a unique `type_id`. + fn encode_type_id(&self, backend_type_id: u16) -> u16 { + // Use the lower 8 bits for the backend's internal type ID + let internal_type_id = backend_type_id & INTERNAL_ID_MASK; + // Use the upper 8 bits for the DispatchDevice/DispatchDeviceId + let backend = u16::from(self.backend_id()) << BACKEND_SHIFT; + backend | internal_type_id + } + + /// Decode an encoded `type_id` into variant ID and backend type ID. + pub(crate) fn decode_type_id(type_id: u16) -> (DispatchDeviceId, u16) { + let backend_raw = type_id >> BACKEND_SHIFT; + let internal_type_id = type_id & INTERNAL_ID_MASK; + + let backend = DispatchDeviceId::try_from(backend_raw).expect("Unknown DispatchDevice ID"); + + (backend, internal_type_id) + } +} + +#[allow(missing_docs)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u16)] +pub enum DispatchDeviceId { + Cpu = 0, + Cuda = 1, + Wgpu = 2, + Rocm = 3, + Flex = 4, + LibTorch = 5, + NdArray = 6, + Metal = 7, + Vulkan = 8, + WebGpu = 9, + Remote = 10, +} + +impl From for u16 { + fn from(variant: DispatchDeviceId) -> Self { + variant as u16 + } +} + +impl TryFrom for DispatchDeviceId { + type Error = (); + + fn try_from(value: u16) -> Result { + match value { + #[cfg(feature = "cpu")] + 0 => Ok(Self::Cpu), + #[cfg(feature = "cuda")] + 1 => Ok(Self::Cuda), + #[cfg(feature = "wgpu")] + 2 => Ok(Self::Wgpu), + #[cfg(feature = "rocm")] + 3 => Ok(Self::Rocm), + #[cfg(any(feature = "flex", default_backend))] + 4 => Ok(Self::Flex), + #[cfg(feature = "tch")] + 5 => Ok(Self::LibTorch), + #[cfg(feature = "ndarray")] + 6 => Ok(Self::NdArray), + #[cfg(feature = "metal")] + 7 => Ok(Self::Metal), + #[cfg(feature = "vulkan")] + 8 => Ok(Self::Vulkan), + #[cfg(feature = "webgpu")] + 9 => Ok(Self::WebGpu), + #[cfg(feature = "remote")] + 10 => Ok(Self::Remote), + _ => Err(()), + } + } +} + +impl DeviceOps for DispatchDevice { + fn defaults(&self) -> DeviceSettings { + match self { + #[cfg(feature = "cpu")] + Self::Cpu(device) => device.defaults(), + #[cfg(feature = "cuda")] + Self::Cuda(device) => device.defaults(), + #[cfg(feature = "metal")] + Self::Metal(device) => device.defaults(), + #[cfg(feature = "rocm")] + Self::Rocm(device) => device.defaults(), + #[cfg(feature = "vulkan")] + Self::Vulkan(device) => device.defaults(), + #[cfg(feature = "wgpu")] + Self::Wgpu(device) => device.defaults(), + #[cfg(feature = "webgpu")] + Self::WebGpu(device) => device.defaults(), + #[cfg(any(feature = "flex", default_backend))] + Self::Flex(device) => device.defaults(), + #[cfg(feature = "ndarray")] + Self::NdArray(device) => device.defaults(), + #[cfg(feature = "tch")] + Self::LibTorch(device) => device.defaults(), + #[cfg(feature = "remote")] + Self::Remote(device) => device.defaults(), + #[cfg(feature = "autodiff")] + Self::Autodiff(device) => device.inner.defaults(), + } + } +} + +impl burn_backend::Device for DispatchDevice { + fn from_id(mut device_id: DeviceId) -> Self { + let (dispatch_id, backend_type_id) = Self::decode_type_id(device_id.type_id); + device_id.type_id = backend_type_id; + + match dispatch_id { + #[cfg(feature = "cpu")] + DispatchDeviceId::Cpu => Self::Cpu(CpuDevice::from_id(device_id)), + #[cfg(feature = "cuda")] + DispatchDeviceId::Cuda => Self::Cuda(CudaDevice::from_id(device_id)), + #[cfg(feature = "metal")] + DispatchDeviceId::Metal => Self::Metal(WgpuDevice::from_id(device_id)), + #[cfg(feature = "rocm")] + DispatchDeviceId::Rocm => Self::Rocm(RocmDevice::from_id(device_id)), + #[cfg(feature = "vulkan")] + DispatchDeviceId::Vulkan => Self::Vulkan(WgpuDevice::from_id(device_id)), + #[cfg(feature = "wgpu")] + DispatchDeviceId::Wgpu => Self::Wgpu(WgpuDevice::from_id(device_id)), + #[cfg(feature = "webgpu")] + DispatchDeviceId::WebGpu => Self::WebGpu(WgpuDevice::from_id(device_id)), + #[cfg(any(feature = "flex", default_backend))] + DispatchDeviceId::Flex => Self::Flex(FlexDevice::from_id(device_id)), + #[cfg(feature = "ndarray")] + DispatchDeviceId::NdArray => Self::NdArray(NdArrayDevice::from_id(device_id)), + #[cfg(feature = "tch")] + DispatchDeviceId::LibTorch => Self::LibTorch(LibTorchDevice::from_id(device_id)), + #[cfg(feature = "remote")] + DispatchDeviceId::Remote => Self::Remote(RemoteDevice::from_id(device_id)), + _ => unreachable!("No backend feature enabled."), + } + } + + fn to_id(&self) -> DeviceId { + let mut device_id = match self { + #[cfg(feature = "cpu")] + Self::Cpu(device) => device.to_id(), + #[cfg(feature = "cuda")] + Self::Cuda(device) => device.to_id(), + #[cfg(feature = "metal")] + Self::Metal(device) => device.to_id(), + #[cfg(feature = "rocm")] + Self::Rocm(device) => device.to_id(), + #[cfg(feature = "vulkan")] + Self::Vulkan(device) => device.to_id(), + #[cfg(feature = "wgpu")] + Self::Wgpu(device) => device.to_id(), + #[cfg(feature = "webgpu")] + Self::WebGpu(device) => device.to_id(), + #[cfg(any(feature = "flex", default_backend))] + Self::Flex(device) => device.to_id(), + #[cfg(feature = "ndarray")] + Self::NdArray(device) => device.to_id(), + #[cfg(feature = "tch")] + Self::LibTorch(device) => device.to_id(), + #[cfg(feature = "remote")] + Self::Remote(device) => device.to_id(), + #[cfg(feature = "autodiff")] + Self::Autodiff(device) => device.inner.to_id(), + }; + device_id.type_id = self.encode_type_id(device_id.type_id); + device_id + } +} + +#[cfg(feature = "cpu")] +impl From for DispatchDevice { + fn from(device: CpuDevice) -> Self { + DispatchDevice::Cpu(device) + } +} + +#[cfg(feature = "cuda")] +impl From for DispatchDevice { + fn from(device: CudaDevice) -> Self { + DispatchDevice::Cuda(device) + } +} + +#[cfg(feature = "rocm")] +impl From for DispatchDevice { + fn from(device: RocmDevice) -> Self { + DispatchDevice::Rocm(device) + } +} + +// A bare `WgpuDevice` maps to the auto-compiler [`DispatchDevice::Wgpu`] variant. To target a +// specific wgpu specialization (Metal, Vulkan, WebGpu) construct the variant explicitly. +#[cfg(all( + feature = "wgpu", + not(any(feature = "metal", feature = "vulkan", feature = "webgpu")) +))] +impl From for DispatchDevice { + fn from(device: WgpuDevice) -> Self { + DispatchDevice::Wgpu(device) + } +} + +#[cfg(all(feature = "metal", not(any(feature = "vulkan", feature = "webgpu"))))] +impl From for DispatchDevice { + fn from(device: WgpuDevice) -> Self { + DispatchDevice::Metal(device) + } +} + +#[cfg(all(feature = "vulkan", not(any(feature = "metal", feature = "webgpu"))))] +impl From for DispatchDevice { + fn from(device: WgpuDevice) -> Self { + DispatchDevice::Vulkan(device) + } +} + +#[cfg(all(feature = "webgpu", not(any(feature = "metal", feature = "vulkan"))))] +impl From for DispatchDevice { + fn from(device: WgpuDevice) -> Self { + DispatchDevice::WebGpu(device) + } +} + +#[cfg(any(feature = "flex", default_backend))] +impl From for DispatchDevice { + fn from(device: FlexDevice) -> Self { + DispatchDevice::Flex(device) + } +} + +#[cfg(feature = "ndarray")] +impl From for DispatchDevice { + fn from(device: NdArrayDevice) -> Self { + DispatchDevice::NdArray(device) + } +} + +#[cfg(feature = "tch")] +impl From for DispatchDevice { + fn from(device: LibTorchDevice) -> Self { + DispatchDevice::LibTorch(device) + } +} + +#[cfg(feature = "remote")] +impl From for DispatchDevice { + fn from(device: RemoteDevice) -> Self { + DispatchDevice::Remote(device) + } +} diff --git a/crates/burn-dispatch/src/lib.rs b/crates/burn-dispatch/src/lib.rs new file mode 100644 index 0000000..5c50a58 --- /dev/null +++ b/crates/burn-dispatch/src/lib.rs @@ -0,0 +1,126 @@ +#![cfg_attr(not(feature = "std"), no_std)] +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] +#![recursion_limit = "138"] + +//! Burn multi-backend dispatch. +//! +//! # Available Backends +//! +//! The dispatch backend supports the following variants, each enabled via cargo features: +//! +//! | Backend | Feature | Description | +//! |------------|------------|-------------| +//! | `Cpu` | `cpu` | Rust CPU backend (MLIR + LLVM) | +//! | `Cuda` | `cuda` | NVIDIA CUDA backend | +//! | `Metal` | `metal` | Apple Metal backend via `wgpu` (MSL) | +//! | `Rocm` | `rocm` | AMD ROCm backend | +//! | `Vulkan` | `vulkan` | Vulkan backend via `wgpu` (SPIR-V) | +//! | `Wgpu` | `webgpu` | WebGPU backend via `wgpu` (WGSL) | +//! | `Flex` | `flex` | Pure Rust CPU backend using `burn-flex` | +//! | `NdArray` | `ndarray` | Pure Rust CPU backend using `ndarray` (legacy - prefer `flex`) | +//! | `LibTorch` | `tch` | Libtorch backend via `tch` | +//! | `Autodiff` | `autodiff` | Autodiff-enabled backend (used in combination with any of the backends above) | +//! +//! **Note:** All backends, including the WGPU-based ones (`wgpu`, `metal`, `vulkan`, `webgpu`), +//! can be combined freely. Each enabled wgpu backend appears as its own +//! [`DispatchDevice`] variant. + +#[macro_use] +mod macros; + +/// Dispatch backend module. +pub mod backend; +/// Dispatch device module. +pub mod device; +mod ops; +/// Dispatch tensor module. +pub mod tensor; + +/// Entry points for hosting a remote-execution server. +#[cfg(feature = "remote-server")] +pub mod remote_server; + +pub use backend::*; +pub use device::*; +pub use tensor::*; + +extern crate alloc; + +/// Backends and devices used. +pub mod backends { + #[cfg(feature = "autodiff")] + pub use burn_autodiff as autodiff; + #[cfg(feature = "autodiff")] + pub use burn_autodiff::Autodiff; // re-export for extensions + + #[cfg(feature = "cpu")] + pub use burn_cpu as cpu; + #[cfg(feature = "cpu")] + pub use burn_cpu::Cpu; + #[cfg(feature = "cuda")] + pub use burn_cuda as cuda; + #[cfg(feature = "cuda")] + pub use burn_cuda::Cuda; + #[cfg(feature = "rocm")] + pub use burn_rocm as rocm; + #[cfg(feature = "rocm")] + pub use burn_rocm::Rocm; + #[cfg(feature = "wgpu")] + pub use burn_wgpu as wgpu; + #[cfg(feature = "metal")] + pub use burn_wgpu::Metal; + #[cfg(feature = "vulkan")] + pub use burn_wgpu::Vulkan; + #[cfg(feature = "webgpu")] + pub use burn_wgpu::WebGpu; + #[cfg(feature = "wgpu")] + pub use burn_wgpu::Wgpu; + + #[cfg(any(feature = "flex", default_backend))] + pub use burn_flex as flex; + #[cfg(any(feature = "flex", default_backend))] + pub use burn_flex::Flex; + #[cfg(feature = "ndarray")] + pub use burn_ndarray as ndarray; + #[cfg(feature = "ndarray")] + pub use burn_ndarray::NdArray; + #[cfg(feature = "tch")] + pub use burn_tch as libtorch; + #[cfg(feature = "tch")] + pub use burn_tch::LibTorch; + + #[cfg(feature = "remote")] + pub use burn_remote as remote; + #[cfg(feature = "remote")] + pub use burn_remote::RemoteBackend as Remote; + + pub use super::devices::*; +} + +// Re-export devices + +/// Backend devices. +pub mod devices { + #[cfg(feature = "cpu")] + pub use burn_cpu::CpuDevice; + #[cfg(feature = "cuda")] + pub use burn_cuda::CudaDevice; + #[cfg(feature = "rocm")] + pub use burn_rocm::RocmDevice; + #[cfg(feature = "wgpu")] + pub use burn_wgpu::WgpuDevice; + + #[cfg(any(feature = "flex", default_backend))] + pub use burn_flex::FlexDevice; + #[cfg(feature = "ndarray")] + pub use burn_ndarray::NdArrayDevice; + #[cfg(feature = "tch")] + pub use burn_tch::LibTorchDevice; + + #[cfg(feature = "remote")] + pub use burn_remote::RemoteDevice; + + #[cfg(feature = "remote")] + pub use burn_remote::BURN_REMOTE_ALPN; +} diff --git a/crates/burn-dispatch/src/macros.rs b/crates/burn-dispatch/src/macros.rs new file mode 100644 index 0000000..9bbece3 --- /dev/null +++ b/crates/burn-dispatch/src/macros.rs @@ -0,0 +1,1466 @@ +/// Supplies a list of all supported backends and their corresponding feature flags +/// to a callback macro. This centralizes the backend registry. +macro_rules! backend_list { + ($callback:ident, $($extra:tt)*) => { + $callback! { + $($extra)*; + [Cpu, feature = "cpu"], + [Cuda, feature = "cuda"], + [Metal, feature = "metal"], + [Rocm, feature = "rocm"], + [Vulkan, feature = "vulkan"], + [Wgpu, feature = "wgpu"], + [WebGpu, feature = "webgpu"], + [Flex, any(feature = "flex", default_backend)], + [NdArray, feature = "ndarray"], + [LibTorch, feature = "tch"], + [Remote, feature = "remote"] + } + }; +} + +/// Supplies a list of all supported distributed backends and their corresponding feature flags +/// to a callback macro. This centralizes the backend registry. +macro_rules! distributed_backend_list { + ($callback:ident, $($extra:tt)*) => { + $callback! { + $($extra)*; + // The cubecl CudaServer implements the distributed communication directly; the remote + // backend forwards the collective operations to its server (which runs a real + // distributed backend). + [Cuda, feature = "cuda"], + [Remote, feature = "remote"] + } + }; +} + +/// Supplies a matrix of cross-backend combinations. Used for operations where the source and destination backends may differ. +macro_rules! backend_matrix { + ($callback:ident, $($extra:tt)*) => { + $callback! { + $($extra)*; + [Cpu, feature = "cpu"] => [[Cuda, feature = "cuda"], [Metal, feature = "metal"], [Rocm, feature = "rocm"], [Vulkan, feature = "vulkan"], [Wgpu, feature = "wgpu"], [WebGpu, feature = "webgpu"], [Flex, any(feature = "flex", default_backend)], [NdArray, feature = "ndarray"], [LibTorch, feature = "tch"], [Remote, feature = "remote"]]; + [Cuda, feature = "cuda"] => [[Cpu, feature = "cpu"], [Metal, feature = "metal"], [Rocm, feature = "rocm"], [Vulkan, feature = "vulkan"], [Wgpu, feature = "wgpu"], [WebGpu, feature = "webgpu"], [Flex, any(feature = "flex", default_backend)], [NdArray, feature = "ndarray"], [LibTorch, feature = "tch"], [Remote, feature = "remote"]]; + [Metal, feature = "metal"] => [[Cpu, feature = "cpu"], [Cuda, feature = "cuda"], [Rocm, feature = "rocm"], [Vulkan, feature = "vulkan"], [Wgpu, feature = "wgpu"], [WebGpu, feature = "webgpu"], [Flex, any(feature = "flex", default_backend)], [NdArray, feature = "ndarray"], [LibTorch, feature = "tch"], [Remote, feature = "remote"]]; + [Rocm, feature = "rocm"] => [[Cpu, feature = "cpu"], [Cuda, feature = "cuda"], [Metal, feature = "metal"], [Vulkan, feature = "vulkan"], [Wgpu, feature = "wgpu"], [WebGpu, feature = "webgpu"], [Flex, any(feature = "flex", default_backend)], [NdArray, feature = "ndarray"], [LibTorch, feature = "tch"], [Remote, feature = "remote"]]; + [Vulkan, feature = "vulkan"] => [[Cpu, feature = "cpu"], [Cuda, feature = "cuda"], [Metal, feature = "metal"], [Rocm, feature = "rocm"], [Wgpu, feature = "wgpu"], [WebGpu, feature = "webgpu"], [Flex, any(feature = "flex", default_backend)], [NdArray, feature = "ndarray"], [LibTorch, feature = "tch"], [Remote, feature = "remote"]]; + [Wgpu, feature = "wgpu"] => [[Cpu, feature = "cpu"], [Cuda, feature = "cuda"], [Metal, feature = "metal"], [Rocm, feature = "rocm"], [Vulkan, feature = "vulkan"], [WebGpu, feature = "webgpu"], [Flex, any(feature = "flex", default_backend)], [NdArray, feature = "ndarray"], [LibTorch, feature = "tch"], [Remote, feature = "remote"]]; + [WebGpu, feature = "webgpu"] => [[Cpu, feature = "cpu"], [Cuda, feature = "cuda"], [Metal, feature = "metal"], [Rocm, feature = "rocm"], [Vulkan, feature = "vulkan"], [Wgpu, feature = "wgpu"], [Flex, any(feature = "flex", default_backend)], [NdArray, feature = "ndarray"], [LibTorch, feature = "tch"], [Remote, feature = "remote"]]; + [Flex, any(feature = "flex", default_backend)] => [[Cpu, feature = "cpu"], [Cuda, feature = "cuda"], [Metal, feature = "metal"], [Rocm, feature = "rocm"], [Vulkan, feature = "vulkan"], [Wgpu, feature = "wgpu"], [WebGpu, feature = "webgpu"], [NdArray, feature = "ndarray"], [LibTorch, feature = "tch"], [Remote, feature = "remote"]]; + [NdArray, feature = "ndarray"] => [[Cpu, feature = "cpu"], [Cuda, feature = "cuda"], [Metal, feature = "metal"], [Rocm, feature = "rocm"], [Vulkan, feature = "vulkan"], [Wgpu, feature = "wgpu"], [WebGpu, feature = "webgpu"], [Flex, any(feature = "flex", default_backend)], [LibTorch, feature = "tch"], [Remote, feature = "remote"]]; + [LibTorch, feature = "tch"] => [[Cpu, feature = "cpu"], [Cuda, feature = "cuda"], [Metal, feature = "metal"], [Rocm, feature = "rocm"], [Vulkan, feature = "vulkan"], [Wgpu, feature = "wgpu"], [WebGpu, feature = "webgpu"], [Flex, any(feature = "flex", default_backend)], [NdArray, feature = "ndarray"], [Remote, feature = "remote"]]; + [Remote, feature = "remote"] => [[Cpu, feature = "cpu"], [Cuda, feature = "cuda"], [Metal, feature = "metal"], [Rocm, feature = "rocm"], [Vulkan, feature = "vulkan"], [Wgpu, feature = "wgpu"], [WebGpu, feature = "webgpu"], [Flex, any(feature = "flex", default_backend)], [NdArray, feature = "ndarray"], [LibTorch, feature = "tch"]] + } + }; +} + +#[cfg(feature = "autodiff")] +/// Helper to map the runtime strategy to the compile-time Autodiff generic. +macro_rules! with_autodiff_backend { + ($Backend:ident, $checkpointing:expr, |$B:ident| $body:expr) => { + match $checkpointing { + Some($crate::CheckpointingStrategy::Balanced) => { + type $B = $crate::backends::Autodiff< + $crate::backends::$Backend, + burn_autodiff::checkpoint::strategy::BalancedCheckpointing, + >; + $body + } + Some($crate::CheckpointingStrategy::None) => { + type $B = $crate::backends::Autodiff< + $crate::backends::$Backend, + burn_autodiff::checkpoint::strategy::NoCheckpointing, + >; + $body + } + None => unreachable!("Should only be called with autodiff."), + } + }; +} + +/// Match arm generator for `dispatch_device`. +/// Maps each backend variant to a block where the specific backend type is bound to `B`. +macro_rules! dispatch_device_arms { + ( + $device:expr, + |$inner:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => { + match $device { + // Autodiff arm first + #[cfg(feature = "autodiff")] + $crate::DispatchDevice::Autodiff(inner) => { + // Recursively dispatch on inner + dispatch_device_arms!( + @autodiff + &**inner, + |$inner| $body; + $([$Backend, $cfg]),* + ) + }, + $( + #[cfg($cfg)] + $crate::DispatchDevice::$Backend($inner) => { + type B = $crate::backends::$Backend; + $body + } + )* + #[allow(unreachable_patterns)] + other => panic!("Distributed operations are not supported for device {other:?}"), + } + }; + ( + @autodiff + $device:expr, + |$inner:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => { + match $device { + $( + #[cfg($cfg)] + $crate::DispatchDevice::$Backend($inner) => { + type B = $crate::backends::Autodiff<$crate::backends::$Backend>; + $body + } + )* + $crate::DispatchDevice::Autodiff(_) => unreachable!("Autodiff should not wrap an autodiff device."), + #[allow(unreachable_patterns)] + other => panic!("Distributed operations are not supported for device {other:?}"), + } + }; +} + +/// Dispatches an operation body based on the provided device. +macro_rules! dispatch_device { + ($device:expr, |$inner:ident| $body:expr) => { + dispatch_device!(@internal backend_list, $device, |$inner| $body) + }; + (@distributed $device:expr, |$inner:ident| $body:expr) => { + dispatch_device!(@internal distributed_backend_list, $device, |$inner| $body) + }; + (@internal $list_macro:ident, $device:expr, |$inner:ident| $body:expr) => { + $list_macro!(dispatch_device_arms, $device, |$inner| $body) + }; +} + +/// Match arm generator for `to_device`. +/// Handles the logic for same-backend transfers (fast path) and cross-backend +/// transfers by generating a grid of all device combinations provided via `backend_matrix`. +macro_rules! to_device_arms { + ( + $kind:ident, $inner_fn:ident, $tensor:expr, $device:expr, $to_device:ident, |$inner:ident, $device_ident:ident| $body:expr; + $( [$B1:ident, $src_cfg:meta] => [ $( [$B2:ident, $dst_cfg:meta] ),+ ] );* + ) => { + match ($tensor.kind, $device) { + // --- Same backend to_device --- + $( + #[cfg($src_cfg)] + ($crate::DispatchTensorKind::$B1(t), $crate::DispatchDevice::$B1(d)) => { + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$B1($crate::BackendTensor::$kind( + $crate::backends::$B1::$to_device(t.$inner_fn(), d) + )), + checkpointing: $tensor.checkpointing, + } + } + )* + + // --- Cross backend arms --- + // This loop generates the grid of combinations + $( + $( + #[cfg(all($src_cfg, $dst_cfg))] + ($crate::DispatchTensorKind::$B1(t), $crate::DispatchDevice::$B2($device_ident)) => { + type B1 = $crate::backends::$B1; + type B2 = $crate::backends::$B2; + let $inner = t.$inner_fn(); + + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$B2( + $crate::BackendTensor::$kind($body) + ), + checkpointing: $tensor.checkpointing, + } + } + )+ + )* + // --- To autodiff --- + // This can happen when moving a bool or int tensor to the device of a float autodiff tensor. + // We move it to the inner device and preserve the checkpointing strategy. + + // --- Same backend to_device --- + $( + #[cfg(all($src_cfg, feature = "autodiff"))] + ($crate::DispatchTensorKind::$B1(t), $crate::DispatchDevice::Autodiff(device_ad)) + if matches!(&*device_ad.inner, $crate::DispatchDevice::$B1(_)) => { + let $crate::DispatchDevice::$B1(d) = &*device_ad.inner else { unreachable!() }; + + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$B1($crate::BackendTensor::$kind( + <$crate::backends::$B1>::$to_device(t.$inner_fn(), d) + )), + checkpointing: Some(device_ad.checkpointing), + } + } + )* + + // --- Same backend to_device --- + $( + $( + #[cfg(all($src_cfg, $dst_cfg, feature = "autodiff"))] + ($crate::DispatchTensorKind::$B1(tensor), $crate::DispatchDevice::Autodiff(device_ad)) + if matches!(&*device_ad.inner, $crate::DispatchDevice::$B2(_)) => { + let $crate::DispatchDevice::$B2($device_ident) = &*device_ad.inner else { unreachable!() }; + type B1 = $crate::backends::$B1; + type B2 = $crate::backends::$B2; + let $inner = tensor.$inner_fn(); + + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$B2( + $crate::BackendTensor::$kind($body) + ), + checkpointing: Some(device_ad.checkpointing), + } + + }, + )+ + )* + #[cfg(feature = "autodiff")] + (_, $crate::DispatchDevice::Autodiff(_)) => unreachable!("Autodiff should not wrap an autodiff device."), + #[cfg(feature = "autodiff")] + ($crate::DispatchTensorKind::Autodiff(..), _) => panic!("Operation not marked for autodiff.") + } + }; +} + +/// Handles tensor movement between devices, supporting both same-backend transfers +/// and cross-backend dispatches. +macro_rules! to_device { + ($kind:ident, $inner_fn:ident, $tensor:expr, $device:expr, $to_device:ident, |$inner:ident, $device_ident:ident| $body:expr) => { + backend_matrix!( + to_device_arms, + $kind, + $inner_fn, + $tensor, + $device, + $to_device, + |$inner, $device_ident| $body + ) + }; +} + +/// Match arm generator for `float_to_device`. +/// +/// Similar to `to_device_arms`, but float tensors are checked for autodiff support. +macro_rules! float_to_device_arms { + ( + $tensor:expr, $device:expr, $to_device:ident, |$inner:ident, $device_ident:ident| $body:expr; + $( [$B1:ident, $src_cfg:meta] => [ $( [$B2:ident, $dst_cfg:meta] ),+ ] );* + ) => { + match ($tensor.kind, $device) { + #[cfg(feature = "autodiff")] + ($crate::DispatchTensorKind::Autodiff(kind), $crate::DispatchDevice::Autodiff(device)) => { + let ckp = $tensor.checkpointing; + float_to_device_arms!( + @autodiff + *kind, &**device, ckp, $to_device; + $([$B1, $src_cfg]);* + ) + + } + // --- Same backend to_device --- + $( + #[cfg($src_cfg)] + ($crate::DispatchTensorKind::$B1(kind), $crate::DispatchDevice::$B1(d)) => { + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$B1($crate::BackendTensor::Float( + $crate::backends::$B1::$to_device(kind.float(), d) + )), + checkpointing: $tensor.checkpointing, + } + } + )* + + // --- Cross backend arms --- + // This loop generates the grid of combinations + $( + $( + #[cfg(all($src_cfg, $dst_cfg))] + ($crate::DispatchTensorKind::$B1(kind), $crate::DispatchDevice::$B2($device_ident)) => { + type B1 = $crate::backends::$B1; + type B2 = $crate::backends::$B2; + let $inner = kind.float(); + + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$B2($crate::BackendTensor::Float($body)), + checkpointing: $tensor.checkpointing, + } + } + )+ + )* + #[cfg(feature = "autodiff")] + ($crate::DispatchTensorKind::Autodiff(..), _) | (_, $crate::DispatchDevice::Autodiff(_)) => panic!("Cannot move between autodiff and non-autodiff instances.") + } + }; + + // Autodiff(DispatchTensor) + ( + @autodiff + $tensor:expr, $device:expr, $ckp:expr, $to_device:ident; + $( [$B1:ident, $src_cfg:meta] );* + ) => {{ + match ($tensor, $device) { + // --- Same backend to_device --- + $( + #[cfg($src_cfg)] + ($crate::DispatchTensorKind::$B1(tensor), $crate::DispatchDevice::$B1(d)) => { + let kind = $crate::DispatchTensorKind::Autodiff(alloc::boxed::Box::new($crate::DispatchTensorKind::$B1($crate::BackendTensor::Autodiff( + with_autodiff_backend!($B1, $ckp, |B| { + B::$to_device(tensor.autodiff(), d) + }) + )))); + $crate::DispatchTensor {kind, checkpointing: $ckp} + } + )* + // TODO: should be possible + (_, _) => unimplemented!("Autodiff tensor cannot be moved between backends.") + } + }}; +} + +/// Handles float tensor movement between devices (that might support autodiff). +macro_rules! float_to_device { + ($kind:ident, $inner_fn:ident, $tensor:expr, $device:expr, $to_device:ident, |$inner:ident, $device_ident:ident| $body:expr) => { + backend_matrix!( + float_to_device_arms, + $tensor, + $device, + $to_device, + |$inner, $device_ident| $body + ) + }; +} + +/// Dispatches a tensor creation operation (e.g., zeros, ones) to the correct backend +/// based on the provided device. +macro_rules! creation_op { + ($kind:ident, $device:expr, |$inner:ident| $body:expr) => { + backend_list!(creation_op_arms, $kind, $device, |$inner| $body) + }; +} + +/// Match arm generator for `creation_float`. +/// +/// Similar to `creation_op_arms`, but float tensors are checked for autodiff support. +macro_rules! creation_op_arms { + ( + $kind:ident, + $device:expr, + |$inner:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => {{ + match $device { + // Autodiff arm first + #[cfg(feature = "autodiff")] + $crate::DispatchDevice::Autodiff(inner) => { + // Recursively dispatch on inner + creation_op_arms!( + @autodiff + $kind, + &**inner, + Some(inner.checkpointing), + |$inner| $body; + $([$Backend, $cfg]),* + ) + }, + $( + #[cfg($cfg)] + $crate::DispatchDevice::$Backend($inner) => { + type B = $crate::backends::$Backend; + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$Backend( + $crate::BackendTensor::$kind($body) + ), + checkpointing: None, + } + } + )* + } + }}; + + ( + @autodiff + $kind:ident, + $device:expr, + $ckp:expr, + |$inner:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => {{ + match $device { + $( + #[cfg($cfg)] + $crate::DispatchDevice::$Backend($inner) => { + with_autodiff_backend!($Backend, $ckp, |B| { + wrap_float!(@wrap_autodiff $kind, $Backend, $ckp, { $body }) + }) + } + )* + $crate::DispatchDevice::Autodiff(_) => unreachable!("Autodiff should not wrap an autodiff device.") + } + }}; +} + +/// Wrap the result in the backend tensor kind, handling float -> autodiff. +#[cfg(feature = "autodiff")] +macro_rules! wrap_float { + ( + @wrap_autodiff Float, + $Backend:ident, + $ckp:expr, + $expr:expr + ) => { + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::Autodiff(alloc::boxed::Box::new( + $crate::DispatchTensorKind::$Backend($crate::BackendTensor::Autodiff($expr)), + )), + checkpointing: $ckp, + } + }; + + ( + @wrap_autodiff $other:ident, + $Backend:ident, + $ckp:expr, + $expr:expr + ) => { + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$Backend($crate::BackendTensor::$other($expr)), + checkpointing: $ckp, + } + }; +} + +/// Match arm generator for `unary_op`. +/// Unwraps the inner tensor primitive (e.g., `inner.float()`) and provides the backend type `B` +/// for the operation. +/// +/// When the return kind is provided, the result is wrapped in the corresponding `DispatchTensor` variant. +macro_rules! unary_op_arms { + ( + Float, // handle autodiff wrapped tensors for float return kind (e.g. int_into_float) + $inner_kind:ident, + $tensor:expr, + |$inner:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => {{ + let checkpointing = $tensor.checkpointing; + + match $tensor.kind { + $( + #[cfg($cfg)] + $crate::DispatchTensorKind::$Backend($inner) => { + type B = $crate::backends::$Backend; + let $inner = $inner.$inner_kind(); + + #[cfg(feature = "autodiff")] + if checkpointing.is_some() { + with_autodiff_backend!($Backend, checkpointing, |B| { + wrap_float!(@wrap_autodiff Float, $Backend, checkpointing, { $body }) + }) + } else { + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$Backend($crate::BackendTensor::Float($body)), + checkpointing, + } + } + + #[cfg(not(feature = "autodiff"))] + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$Backend($crate::BackendTensor::Float($body)), + checkpointing, + } + } + )* + #[cfg(feature = "autodiff")] + $crate::DispatchTensorKind::Autodiff(..) => panic!("Operation not marked for autodiff.") + } + }}; + + ( + $kind:ident, + $inner_kind:ident, + $tensor:expr, + |$inner:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => {{ + let checkpointing = $tensor.checkpointing; + + match $tensor.kind { + $( + #[cfg($cfg)] + $crate::DispatchTensorKind::$Backend($inner) => { + type B = $crate::backends::$Backend; + let $inner = $inner.$inner_kind(); + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$Backend($crate::BackendTensor::$kind($body)), + checkpointing, + } + } + )* + #[cfg(feature = "autodiff")] + $crate::DispatchTensorKind::Autodiff(..) => panic!("Operation not marked for autodiff.") + } + }}; + + // Operations that do not return a tensor kind + ( + $inner_kind:ident, + $tensor:expr, + |$inner:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => {{ + match $tensor.kind { + $( + #[cfg($cfg)] + $crate::DispatchTensorKind::$Backend($inner) => { + type B = $crate::backends::$Backend; + let $inner = $inner.$inner_kind(); + $body + } + )* + #[cfg(feature = "autodiff")] + $crate::DispatchTensorKind::Autodiff(..) => panic!("Operation not marked for autodiff.") + } + }}; +} + +/// Backend dispatch for unary operations. +/// +/// When the return `=> Kind` is not provided, the operation output is not wrapped in a dispatch tensor (e.g., `into_data(..)`) +macro_rules! unary_op { + ($tensor:expr, $inner_kind:ident, |$inner:ident| $body:expr => $kind:ident) => { + backend_list!(unary_op_arms, $kind, $inner_kind, $tensor, |$inner| { + $body + }) + }; + ($tensor:expr, $inner_kind:ident, |$inner:ident| $body:expr) => { + backend_list!(unary_op_arms, $inner_kind, $tensor, |$inner| { $body }) + }; +} + +/// Match arm generator for `unary_float`. +/// +/// Similar to `unary_op_arms`, but float tensors are checked for autodiff support. +macro_rules! unary_float_arms { + ( + $mode:ident, // `owned` or `ref` + $kind:ident, + $inner_kind:ident, + $tensor:expr, + |$inner:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => {{ + let checkpointing = $tensor.checkpointing; + + match $tensor.kind { + #[cfg(feature = "autodiff")] + $crate::DispatchTensorKind::Autodiff(inner) => { + unary_float_arms!( + @autodiff $mode, + checkpointing, + $kind, + { if_mode!($mode, &**inner, *inner) }, + |$inner| $body; + $([$Backend, $cfg]),* + ) + }, + $( + #[cfg($cfg)] + $crate::DispatchTensorKind::$Backend($inner) => { + type B = $crate::backends::$Backend; + let $inner = unary_float_arms!(@unwrap $mode, $inner, $inner_kind); + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$Backend( + $crate::BackendTensor::$kind($body) + ), + checkpointing, + } + } + )* + #[allow(unreachable_patterns)] + other => panic!("Distributed operations are not supported for tensor kind {other:?}"), + } + }}; + + // --- Autodiff recursive arm --- + ( + @autodiff $mode:ident, + $ckp:expr, + $kind:ident, + $tensor:expr, + |$inner:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => {{ + match $tensor { + $( + #[cfg($cfg)] + $crate::DispatchTensorKind::$Backend($inner) => { + with_autodiff_backend!($Backend, $ckp, |B| { + let $inner = unary_float_arms!(@unwrap_ad $mode, $inner); + wrap_float!(@wrap_autodiff $kind, $Backend, $ckp, { $body }) + }) + } + )* + $crate::DispatchTensorKind::Autodiff(..) => unreachable!("Autodiff should not wrap an autodiff tensor."), + #[allow(unreachable_patterns)] + other => panic!("Distributed operations are not supported for tensor kind {other:?}"), + } + }}; + + // --- Non-wrapping arms (operations not returning a tensor) --- + ( + $mode:ident, + $inner_kind:ident, + $tensor:expr, + |$inner:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => {{ + #[cfg(feature = "autodiff")] + let checkpointing = &$tensor.checkpointing; + + match { if_mode!($mode, &$tensor.kind, $tensor.kind) } { + #[cfg(feature = "autodiff")] + $crate::DispatchTensorKind::Autodiff(inner) => { + unary_float_arms!( + @autodiff $mode, + checkpointing, + { if_mode!($mode, &**inner, *inner) }, + |$inner| $body; + $([$Backend, $cfg]),* + ) + }, + $( + #[cfg($cfg)] + $crate::DispatchTensorKind::$Backend($inner) => { + type B = $crate::backends::$Backend; + let $inner = unary_float_arms!(@unwrap $mode, $inner, $inner_kind); + $body + } + )* + #[allow(unreachable_patterns)] + other => panic!("Distributed operations are not supported for tensor kind {other:?}"), + } + }}; + ( + @autodiff $mode:ident, + $ckp:expr, + $tensor:expr, + |$inner:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => {{ + match $tensor { + $( + #[cfg($cfg)] + $crate::DispatchTensorKind::$Backend($inner) => { + with_autodiff_backend!($Backend, $ckp, |B| { + let $inner = unary_float_arms!(@unwrap_ad $mode, $inner); + $body + }) + } + )* + $crate::DispatchTensorKind::Autodiff(..) => unreachable!("Autodiff should not wrap an autodiff tensor."), + #[allow(unreachable_patterns)] + other => panic!("Distributed operations are not supported for tensor kind {other:?}"), + } + }}; + + // --- Helpers to unwarp the tensor based on owned/ref --- + (@unwrap owned, $inner:ident, $inner_kind:ident) => { $inner.$inner_kind() }; + (@unwrap ref, $inner:ident, $inner_kind:ident) => { + paste::paste! { $inner.[< as_ $inner_kind >]() } + }; + + (@unwrap_ad owned, $inner:ident) => { $inner.autodiff() }; + (@unwrap_ad ref, $inner:ident) => { $inner.as_autodiff() }; + +} + +/// Utility to pick a token based on mode +macro_rules! if_mode { + (ref, $if_ref:expr, $if_owned:expr) => { + $if_ref + }; + (owned, $if_ref:expr, $if_owned:expr) => { + $if_owned + }; +} + +/// Backend dispatch for float unary operations (that might support autodiff). +/// +/// When the return `=> Kind` is not provided, the operation output is not wrapped in a dispatch tensor (e.g., `into_data(..)`) +macro_rules! unary_float { + // --- Entrypoints that default to standard backend_list --- + ($tensor:expr, $inner_kind:ident, |$inner:ident| $body:expr => $kind:ident) => { + unary_float!(@internal backend_list, $tensor, $inner_kind, |$inner| $body => $kind) + }; + ($tensor:expr, $inner_kind:ident, |$inner:ident| $body:expr) => { + unary_float!(@internal backend_list, $tensor, $inner_kind, |$inner| $body) + }; + (ref $tensor:expr, $inner_kind:ident, |$inner:ident| $body:expr) => { + unary_float!(@internal backend_list, ref $tensor, $inner_kind, |$inner| $body) + }; + + // --- Distributed entrypoint --- + (@distributed $tensor:expr, $inner_kind:ident, |$inner:ident| $body:expr => $kind:ident) => { + unary_float!(@internal distributed_backend_list, $tensor, $inner_kind, |$inner| $body => $kind) + }; + + // --- Internal implementation workers that accept the list macro dynamically --- + (@internal $list_macro:ident, $tensor:expr, $inner_kind:ident, |$inner:ident| $body:expr => $kind:ident) => { + $list_macro!( + unary_float_arms, + owned, + $kind, + $inner_kind, + $tensor, + |$inner| { $body } + ) + }; + (@internal $list_macro:ident, $tensor:expr, $inner_kind:ident, |$inner:ident| $body:expr) => { + $list_macro!(unary_float_arms, owned, $inner_kind, $tensor, |$inner| { + $body + }) + }; + (@internal $list_macro:ident, ref $tensor:expr, $inner_kind:ident, |$inner:ident| $body:expr) => { + $list_macro!(unary_float_arms, ref, $inner_kind, $tensor, |$inner| { + $body + }) + }; +} + +/// Match arm generator for `binary_op`. +/// Matches two tensors to ensure they share the same backend before unwrapping them for the operation. +macro_rules! binary_op_arms { + ( + $kind:ident, + ($lhs:expr, $lhs_kind:ident), + ($rhs:expr, $rhs_kind:ident), + |$lhs_inner:ident, $rhs_inner:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => {{ + #[cfg(feature = "autodiff")] + let checkpointing = $crate::validate_checkpointing($lhs.checkpointing, $rhs.checkpointing); + #[cfg(not(feature = "autodiff"))] + let checkpointing = $lhs.checkpointing; + + match ($lhs.kind, $rhs.kind) { + $( + #[cfg($cfg)] + ($crate::DispatchTensorKind::$Backend($lhs_inner), $crate::DispatchTensorKind::$Backend($rhs_inner)) => { + type B = $crate::backends::$Backend; + let $lhs_inner = $lhs_inner.$lhs_kind(); + let $rhs_inner = $rhs_inner.$rhs_kind(); + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$Backend($crate::BackendTensor::$kind($body)), + + checkpointing, + } + } + )* + #[allow(unreachable_patterns)] + (lhs, rhs) => { + panic!( + "The provided tensors are not on the same backend. Got backends {:?} and {:?}.", lhs, rhs + ); + } + } + }}; +} + +/// Backend dispatch for binary operations. +/// Automatically verifies that both tensors reside on the same backend. +macro_rules! binary_op { + (($lhs:expr, $lhs_kind:ident), ($rhs:expr, $rhs_kind:ident), |$lhs_inner:ident, $rhs_inner:ident| $body:expr => $kind:ident) => { + backend_list!( + binary_op_arms, + $kind, + ($lhs, $lhs_kind), + ($rhs, $rhs_kind), + |$lhs_inner, $rhs_inner| { $body } + ) + }; +} + +/// Match arm generator for `binary_float`. +/// Matches two tensors to ensure they share the same backend before unwrapping them for the operation. +macro_rules! binary_float_arms { + // (float, float) binary op + ( + $kind:ident, + ($lhs:expr, float), + ($rhs:expr, float), + |$lhs_inner:ident, $rhs_inner:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => {{ + #[cfg(feature = "autodiff")] + let checkpointing = $crate::validate_checkpointing($lhs.checkpointing, $rhs.checkpointing); + #[cfg(not(feature = "autodiff"))] + let checkpointing = $lhs.checkpointing; + + match ($lhs.kind, $rhs.kind) { + // Autodiff arms first + #[cfg(feature = "autodiff")] + ($crate::DispatchTensorKind::Autodiff(lhs_inner), $crate::DispatchTensorKind::Autodiff(rhs_inner)) => { + // Recursively dispatch on inner + binary_float_arms!( + @autodiff + $kind, + (*lhs_inner, autodiff, checkpointing), + (*rhs_inner, autodiff, checkpointing), + |$lhs_inner, $rhs_inner| $body; + $([$Backend, $cfg]),* + ) + }, + $( + #[cfg($cfg)] + ($crate::DispatchTensorKind::$Backend($lhs_inner), $crate::DispatchTensorKind::$Backend($rhs_inner)) => { + type B = $crate::backends::$Backend; + let $lhs_inner = $lhs_inner.float(); + let $rhs_inner = $rhs_inner.float(); + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$Backend($crate::BackendTensor::$kind($body)), + checkpointing, + } + } + )* + #[allow(unreachable_patterns)] + (lhs, rhs) => { + panic!( + "The provided tensors are not on the same backend. Got backends {:?} and {:?}.", lhs, rhs + ); + } + } + }}; + // (float, any) binary op + ( + $kind:ident, + ($lhs:expr, float), + ($rhs:expr, $rhs_kind:ident), + |$lhs_inner:ident, $rhs_inner:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => {{ + #[cfg(feature = "autodiff")] + let checkpointing = $crate::validate_checkpointing($lhs.checkpointing, $rhs.checkpointing); + #[cfg(not(feature = "autodiff"))] + let checkpointing = $lhs.checkpointing; + + match ($lhs.kind, $rhs.kind) { + $( + // Autodiff arms first + #[cfg(all(feature = "autodiff", $cfg))] + ($crate::DispatchTensorKind::Autodiff(lhs_inner), $crate::DispatchTensorKind::$Backend($rhs_inner)) => { + // Match on inner + match *lhs_inner { + $crate::DispatchTensorKind::$Backend($lhs_inner) => { + with_autodiff_backend!($Backend, checkpointing, |B| { + let $lhs_inner = $lhs_inner.autodiff(); + let $rhs_inner = $rhs_inner.$rhs_kind(); + wrap_float!( + @wrap_autodiff + $kind, + $Backend, + checkpointing, + { $body } + ) + }) + } + $crate::DispatchTensorKind::Autodiff(..) => unreachable!("Autodiff should not wrap an autodiff tensor."), + #[allow(unreachable_patterns)] + _ => panic!("The provided tensors are not on the same backend.") + } + }, + + #[cfg($cfg)] + ($crate::DispatchTensorKind::$Backend($lhs_inner), $crate::DispatchTensorKind::$Backend($rhs_inner)) => { + type B = $crate::backends::$Backend; + let $lhs_inner = $lhs_inner.float(); + let $rhs_inner = $rhs_inner.$rhs_kind(); + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$Backend($crate::BackendTensor::$kind($body)), + checkpointing, + } + } + )* + #[allow(unreachable_patterns)] + (lhs, rhs) => { + panic!( + "The provided tensors are not on the same backend. Got backends {:?} and {:?}.", lhs, rhs + ); + } + } + }}; + ( + $kind:ident, + ($lhs:expr, $lhs_kind:ident), + ($rhs:expr, $rhs_kind:ident), + |$lhs_inner:ident, $rhs_inner:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => {{ + match ($lhs, $rhs) { + $( + #[cfg($cfg)] + ($crate::DispatchTensorKind::$Backend($lhs_inner), $crate::DispatchTensorKind::$Backend($rhs_inner)) => { + type B = $crate::backends::$Backend; + let $lhs_inner = $lhs_inner.$lhs_kind(); + let $rhs_inner = $rhs_inner.$rhs_kind(); + $crate::DispatchTensorKind::$Backend($crate::BackendTensor::$kind($body)) + } + )* + (lhs, rhs) => { + panic!( + "The provided tensors are not on the same backend. Got backends {:?} and {:?}.", lhs, rhs + ); + } + } + }}; + // Autodiff (lhs, rhs) tensors + ( + @autodiff + $kind:ident, + ($lhs:expr, $lhs_kind:ident, $ckp_lhs:expr), + ($rhs:expr, $rhs_kind:ident, $ckp_rhs:expr), + |$lhs_inner:ident, $rhs_inner:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => {{ + match ($lhs, $rhs) { + $( + #[cfg($cfg)] + ($crate::DispatchTensorKind::$Backend($lhs_inner), $crate::DispatchTensorKind::$Backend($rhs_inner)) => { + with_autodiff_backend!($Backend, $ckp_lhs, |B| { + let $lhs_inner = $lhs_inner.$lhs_kind(); + let $rhs_inner = $rhs_inner.$rhs_kind(); + wrap_float!( + @wrap_autodiff + $kind, + $Backend, + $ckp_lhs, + { $body } + ) + }) + } + )* + #[cfg(feature = "autodiff")] + ($crate::DispatchTensorKind::Autodiff(..), _) | (_, $crate::DispatchTensorKind::Autodiff(..)) => unreachable!("Autodiff should not wrap an autodiff tensor."), + #[allow(unreachable_patterns)] + (lhs, rhs) => { + panic!( + "The provided tensors are not on the same backend. Got backends {:?} and {:?}.", lhs, rhs + ); + } + } + }}; + +} + +/// Backend dispatch for binary operations. +/// Automatically verifies that both tensors reside on the same backend. +macro_rules! binary_float { + (($lhs:expr, $lhs_kind:ident), ($rhs:expr, $rhs_kind:ident), |$lhs_inner:ident, $rhs_inner:ident| $body:expr => $kind:ident) => { + backend_list!( + binary_float_arms, + $kind, + ($lhs, $lhs_kind), + ($rhs, $rhs_kind), + |$lhs_inner, $rhs_inner| { $body } + ) + }; +} + +/// The core logic for a single backend in a `multi_op`. +/// Handles the manual unwrapping of required/optional inputs and the +/// re-wrapping of multiple required/optional output tensors. +macro_rules! multi_op_arm { + ( + $Backend:ident, + $ckp:ident, + [ $( ($x:ident, $x_kind:ident) ),+ ], + [ $( ($opt_in:ident, $opt_kind:ident) ),* ], + [ $( ($out:ident, $out_kind:ident) ),+ ], + [ $( $opt_out:ident ),* ], + $body:expr + ) => {{ + type B = $crate::backends::$Backend; + + // Required inputs + $( + let $x = match $x.kind { + $crate::DispatchTensorKind::$Backend(inner) => inner.$x_kind(), + #[allow(unreachable_patterns)] + _ => panic!("Input tensor {} is on the wrong device", stringify!($x)), + }; + )+ + + // Optional inputs + $( + let $opt_in = $opt_in.map(|o| match o.kind { + $crate::DispatchTensorKind::$Backend(inner) => inner.$opt_kind(), + #[allow(unreachable_patterns)] + _ => panic!("Optional tensor {} is on the wrong device", stringify!($opt_in)), + }); + )* + + let ($($out),+, $($opt_out),*) = $body; + + // Outputs and optional outputs + ( + $( + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$Backend($crate::BackendTensor::$out_kind($out)), + checkpointing: $ckp, + } + ),+, + $( + $opt_out.map(|t| + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$Backend($crate::BackendTensor::Float(t)), + checkpointing: $ckp, + } + ) + ),* + ) + }}; +} + +#[cfg(feature = "autodiff")] +macro_rules! wrap_input_autodiff { + ($Backend:ident, $inner:expr, int) => { + $inner.int() + }; + ($Backend:ident, $inner:expr, bool) => { + $inner.bool() + }; + // Float tensors: wrap with autodiff + ($Backend:ident, $inner:expr, float) => { + $inner.autodiff() + }; +} + +#[cfg(feature = "autodiff")] +// DispatchTensorKind::Autodiff(DispatchTensorKind::$Backend(BackendTensor::Autodiff())) +macro_rules! multi_op_arm_autodiff { + ( + $Backend:ident, + $ckp:ident, + [ $( ($x:ident, $x_kind:ident) ),+ ], + [ $( ($opt_in:ident, $opt_kind:ident) ),* ], + [ $( ($out:ident, $out_kind:ident) ),+ ], + [ $( $opt_out:ident ),* ], + $body:expr + ) => {{ + // type B = Autodiff<$Backend>; + with_autodiff_backend!($Backend, $ckp, |B| { + // Required inputs + $( + let $x = match $x.kind { + $crate::DispatchTensorKind::Autodiff(inner) => { + match *inner { + $crate::DispatchTensorKind::$Backend(inner) => wrap_input_autodiff!($Backend, inner, $x_kind), + _ => panic!("Input tensor {} is on the wrong device", stringify!($x)), + } + }, + // Unreachable, except when input is int + $crate::DispatchTensorKind::$Backend(inner) => wrap_input_autodiff!($Backend, inner, $x_kind), + #[allow(unreachable_patterns)] + _ => panic!("Input tensor {} is on the wrong device", stringify!($x)), + }; + )+ + + // Optional inputs (always assumed to be float / autodiff) + $( + let $opt_in = $opt_in.map(|o| match o.kind { + $crate::DispatchTensorKind::Autodiff(inner) => { + match *inner { + $crate::DispatchTensorKind::$Backend(inner) => wrap_input_autodiff!($Backend, inner, $opt_kind), + _ => panic!("Input tensor {} is on the wrong device", stringify!($opt_in)), + } + }, + _ => panic!("Optional tensor {} is on the wrong device", stringify!($opt_in)), + }); + )* + + let ($($out),+, $($opt_out),*) = $body; + + // Outputs and optional outputs + ( + $( wrap_float!(@wrap_autodiff $out_kind, $Backend, $ckp, $out) ),+, + $( $opt_out.map(|t| wrap_float!(@wrap_autodiff Float, $Backend, $ckp, t)) ),* + ) + }) + }}; +} + +/// Helper to extract the first identifier from an input list. +/// Used to determine the device/backend for dispatching multi-tensor operations. +macro_rules! first_input { + ([ ($x:ident, $kind:ident) $(, $rest:tt)* ]) => { + $x + }; +} + +/// Match arm generator for `multi_op`. +/// Determines the backend based on the first input and delegates to `multi_op_arm` +/// to handle the repetition-heavy unwrapping and wrapping logic. +macro_rules! multi_op_arms_autodiff { + ( + $inputs:tt, + $opt_inputs:tt, + $outputs:tt, + $opt_outputs:tt, + $body:expr; + $( [$Backend:ident, $cfg:meta] ),* + ) => {{ + let first_input = &first_input!($inputs); + let checkpointing = first_input.checkpointing; + match &first_input.kind { + // Autodiff first + #[cfg(feature = "autodiff")] + $crate::DispatchTensorKind::Autodiff(inner) => { + match **inner { + $( + #[cfg($cfg)] + $crate::DispatchTensorKind::$Backend(_) => { + multi_op_arm_autodiff!( + $Backend, + checkpointing, + $inputs, + $opt_inputs, + $outputs, + $opt_outputs, + $body + ) + } + )* + $crate::DispatchTensorKind::Autodiff(..) => unreachable!("Autodiff should not wrap an autodiff tensor.") + } + }, + $( + #[cfg($cfg)] + $crate::DispatchTensorKind::$Backend(_) => { + multi_op_arm!( + $Backend, + checkpointing, + $inputs, + $opt_inputs, + $outputs, + $opt_outputs, + $body + ) + } + )* + } + }}; +} + +/// Match arm generator for `multi_op`. +/// +/// Similar to `multi_op_arms`, but skips autodiff checks. +macro_rules! multi_op_arms { + ( + $inputs:tt, + $opt_inputs:tt, + $outputs:tt, + $opt_outputs:tt, + $body:expr; + $( [$Backend:ident, $cfg:meta] ),* + ) => {{ + let first_input = &first_input!($inputs); + let checkpointing = first_input.checkpointing; + + match first_input.kind { + $( + #[cfg($cfg)] + $crate::DispatchTensorKind::$Backend(_) => { + multi_op_arm!( + $Backend, + checkpointing, + $inputs, + $opt_inputs, + $outputs, + $opt_outputs, + $body + ) + } + )* + #[cfg(feature = "autodiff")] + $crate::DispatchTensorKind::Autodiff(..) => panic!("Operation not marked for autodiff.") + } + }}; +} + +/// High-level macro for complex module operations (e.g., conv2d) and multi-tensor operations. +/// Handles variable numbers of required/optional inputs and wraps multiple outputs. +/// +/// Usage: +/// ```ignore +/// multi_op!( +/// inputs[(x, float), (weight, float)], +/// opt_inputs[(bias, float)], +/// => Float, +/// B::conv2d(x, weight, bias, options) +/// ) +/// ``` +macro_rules! multi_op { + // --- Single output shorthands --- + // Automatically wraps body in tuple and extracts .0 + ( + inputs[$( ($x:ident, $kind:ident) ),+], + => Float, + $body:expr + ) => { + multi_op!( + inputs[$( ($x, $kind) ),+], + opt_inputs[], + outputs[(out, Float)], + opt_outputs[], + { ($body,) } + ) + .0 + }; + ( + inputs[$( ($x:ident, $kind:ident) ),+], + opt_inputs[ $(($opt_in:ident, $opt_kind:ident)),* ], + => $out_kind:ident, + $body:expr + ) => { + multi_op!( + inputs[$( ($x, $kind) ),+], + opt_inputs[ $(($opt_in, $opt_kind)),* ], + outputs[(out, $out_kind)], + opt_outputs[], + { ($body,) } + ) + .0 + }; + // Int/Bool op specialization (not marked for autodiff) + ( + inputs[$( ($x:ident, $kind:ident) ),+], + => $out_kind:ident, + $body:expr + ) => { + backend_list!( + multi_op_arms, + [ $(($x, $kind)),+ ], + [], + [ (out, $out_kind) ], + [], + { ($body,) } + ).0 + }; + + // --- Required + optional for both inputs and outputs --- + ( + inputs[ $(($x:ident, $kind:ident)),+ ], + opt_inputs[ $(($opt_in:ident, $opt_kind:ident)),* ], + outputs[ $( ($out:ident, $out_kind:ident) ),+ ], + opt_outputs[ $($opt_out:ident),* ], + $body:expr + ) => { + backend_list!( + multi_op_arms_autodiff, + [ $(($x, $kind)),+ ], + [ $(($opt_in, $opt_kind)),* ], + [ $(($out, $out_kind)),+ ], + [ $($opt_out),* ], + $body + ) + }; + + ( + inputs[ $(($x:ident, $kind:ident)),+ ], + opt_inputs[ $(($opt_in:ident, $opt_kind:ident)),* ], + outputs[ $($out:ident),+ ], + $body:expr + ) => { + multi_op!( + inputs[ $(($x, $kind)),+ ], + opt_inputs[ $(($opt_in, $opt_kind)),* ], + outputs[ $(($out, Float)),+ ], + opt_outputs[], + $body + ) + }; + + ( + inputs[ $(($x:ident, $kind:ident)),+ ], + outputs[ $( ($out:ident, $out_kind:ident) ),+ ], + $body:expr + ) => { + multi_op!( + inputs[ $(($x, $kind)),+ ], + opt_inputs[], + outputs[ $(($out, $out_kind)),+ ], + opt_outputs[], + $body + ) + }; +} + +/// Unwraps a `Vec` for a known backend. +macro_rules! unwrap_vec { + ($Backend:ident, $vec:expr, $kind:ident) => { + $vec.into_iter() + .map(|t| match t.kind { + $crate::DispatchTensorKind::$Backend(inner) => inner.$kind(), + #[allow(unreachable_patterns)] + _ => panic!( + "Tensor is on the wrong backend (expected {}).", + stringify!($Backend) + ), + }) + .collect::>() + }; + + // Autodiff-wrapped backend + (@autodiff $Backend:ident, $vec:expr, $kind:ident) => { + $vec.into_iter() + .map(|t| match t.kind { + $crate::DispatchTensorKind::Autodiff(inner) => match *inner { + $crate::DispatchTensorKind::$Backend(inner) => inner.$kind(), + _ => panic!( + "Autodiff float tensor is on the wrong backend (expected {}).", + stringify!($Backend) + ), + }, + _ => panic!( + "Expected autodiff-wrapped float tensor for backend {}.", + stringify!($Backend) + ), + }) + .collect::>() + }; +} + +/// Match arm generator for `vec_op`. +macro_rules! vec_op_arms { + (Float, $inner_kind:ident, $tensors:expr, |$inner:ident| $body:expr; $([$Backend:ident, $cfg:meta]),*) => {{ + let first = &$tensors[0]; + let checkpointing = first.checkpointing; + + match &first.kind { + // Autodiff arm first + #[cfg(feature = "autodiff")] + $crate::DispatchTensorKind::Autodiff(inner) => { + // Recursively dispatch on inner + match **inner { + $( + #[cfg($cfg)] + $crate::DispatchTensorKind::$Backend(_) => { + with_autodiff_backend!($Backend, checkpointing, |B| { + let $inner = unwrap_vec!(@autodiff $Backend, $tensors, autodiff); + wrap_float!( @wrap_autodiff Float, $Backend, checkpointing, { $body } ) + }) + } + )* + $crate::DispatchTensorKind::Autodiff(..) => unreachable!("Autodiff should not wrap an autodiff tensor.") + } + }, + + $( + #[cfg($cfg)] + $crate::DispatchTensorKind::$Backend(_) => { + type B = $crate::backends::$Backend; + + let $inner = unwrap_vec!($Backend, $tensors, $inner_kind); + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$Backend($crate::BackendTensor::Float($body)), + checkpointing, + } + } + )* + } + }}; + ($kind:ident, $inner_kind:ident, $tensors:expr, |$inner:ident| $body:expr; $([$Backend:ident, $cfg:meta]),*) => {{ + let first = &$tensors[0]; + let checkpointing = first.checkpointing; + match first.kind { + $( + #[cfg($cfg)] + $crate::DispatchTensorKind::$Backend(_) => { + type B = $crate::backends::$Backend; + + let $inner = unwrap_vec!($Backend, $tensors, $inner_kind); + $crate::DispatchTensor { + kind: $crate::DispatchTensorKind::$Backend($crate::BackendTensor::$kind($body)), + checkpointing, + } + } + )* + #[cfg(feature = "autodiff")] + $crate::DispatchTensorKind::Autodiff(..) => panic!("Operation not marked for autodiff.") + } + }}; +} + +/// Backend dispatch for operations on multiple inputs (vec). +/// Automatically verifies that tensors reside on the first backend. +macro_rules! vec_op { + ($tensors:expr, $inner_kind:ident, |$inner:ident| $body:expr => $kind:ident) => { + backend_list!(vec_op_arms, $kind, $inner_kind, $tensors, |$inner| { + $body + }) + }; +} + +/// Match arm generator for `transaction_op`. +macro_rules! transaction_op_arms { + ($tx:ident, $first:expr; $([$Backend:ident, $cfg:meta]),*) => {{ + match &$first.kind { + // Autodiff arm first + #[cfg(feature = "autodiff")] + $crate::DispatchTensorKind::Autodiff(inner) => { + // Recursively dispatch on inner + match **inner { + $( + #[cfg($cfg)] + $crate::DispatchTensorKind::$Backend(_) => { + type B = $crate::backends::$Backend; + + // Unwrap vec + let floats = unwrap_vec!(@autodiff $Backend, $tx.read_floats, autodiff_inner); + let ints = unwrap_vec!($Backend, $tx.read_ints, int); + let bools = unwrap_vec!($Backend, $tx.read_bools, bool); + // Not supported + let qfloats = $tx.read_qfloats.into_iter().map(|_t| todo!("Quantization not supported yet")).collect(); + + B::tr_execute(TransactionPrimitive::new(floats, qfloats, ints, bools)).await + } + )* + $crate::DispatchTensorKind::Autodiff(..) => unreachable!("Autodiff should not wrap an autodiff tensor.") + } + }, + + $( + #[cfg($cfg)] + $crate::DispatchTensorKind::$Backend(_) => { + type B = $crate::backends::$Backend; + + // Unwrap vec + let floats = unwrap_vec!($Backend, $tx.read_floats, float); + let ints = unwrap_vec!($Backend, $tx.read_ints, int); + let bools = unwrap_vec!($Backend, $tx.read_bools, bool); + // Not supported + let qfloats = $tx.read_qfloats.into_iter().map(|_t| todo!("Quantization not supported yet")).collect(); + + B::tr_execute(TransactionPrimitive::new(floats, qfloats, ints, bools)).await + } + )* + } + }}; +} + +/// Helper to dispatch a transaction based on the first available tensor. +macro_rules! transaction_op { + ($tx:ident, $first:expr) => { + backend_list!(transaction_op_arms, $tx, $first) + }; +} diff --git a/crates/burn-dispatch/src/ops/activation.rs b/crates/burn-dispatch/src/ops/activation.rs new file mode 100644 index 0000000..2346697 --- /dev/null +++ b/crates/burn-dispatch/src/ops/activation.rs @@ -0,0 +1,61 @@ +use burn_backend::{Scalar, ops::ActivationOps, tensor::FloatTensor}; + +use crate::Dispatch; + +impl ActivationOps for Dispatch { + fn leaky_relu(tensor: FloatTensor, negative_slope: Scalar) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::leaky_relu(tensor, negative_slope) => Float) + } + + fn relu(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::relu(tensor) => Float) + } + + fn relu_backward(output: FloatTensor, grad: FloatTensor) -> FloatTensor { + binary_float!((output, float), (grad, float), |output, grad| B::relu_backward(output, grad) => Float) + } + + fn gelu(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::gelu(tensor) => Float) + } + + fn prelu(tensor: FloatTensor, alpha: FloatTensor) -> FloatTensor { + binary_float!((tensor, float), (alpha, float), |tensor, alpha| B::prelu(tensor, alpha) => Float) + } + + fn gelu_backward(x: FloatTensor, grad: FloatTensor) -> FloatTensor { + binary_float!((x, float), (grad, float), |x, grad| B::gelu_backward(x, grad) => Float) + } + + fn sigmoid(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::sigmoid(tensor) => Float) + } + + fn sigmoid_backward(output: FloatTensor, grad: FloatTensor) -> FloatTensor { + binary_float!((output, float), (grad, float), |output, grad| B::sigmoid_backward(output, grad) => Float) + } + + fn hard_sigmoid(tensor: FloatTensor, alpha: Scalar, beta: Scalar) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::hard_sigmoid(tensor, alpha, beta) => Float) + } + + fn softmax(tensor: FloatTensor, dim: usize) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::softmax(tensor, dim) => Float) + } + + fn log_softmax(tensor: FloatTensor, dim: usize) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::log_softmax(tensor, dim) => Float) + } + + fn softmin(tensor: FloatTensor, dim: usize) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::softmin(tensor, dim) => Float) + } + + fn log_sigmoid(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::log_sigmoid(tensor) => Float) + } + + fn log_sigmoid_backward(x: FloatTensor, grad: FloatTensor) -> FloatTensor { + binary_float!((x, float), (grad, float), |x, grad| B::log_sigmoid_backward(x, grad) => Float) + } +} diff --git a/crates/burn-dispatch/src/ops/bool_tensor.rs b/crates/burn-dispatch/src/ops/bool_tensor.rs new file mode 100644 index 0000000..ef8f0b4 --- /dev/null +++ b/crates/burn-dispatch/src/ops/bool_tensor.rs @@ -0,0 +1,217 @@ +use alloc::vec::Vec; +use burn_backend::{ + BoolDType, ExecutionError, FloatDType, IntDType, Scalar, Shape, Slice, TensorData, + ops::BoolTensorOps, + tensor::{BoolTensor, FloatTensor, IntTensor}, +}; + +use crate::{Dispatch, DispatchDevice}; + +impl BoolTensorOps for Dispatch { + fn bool_empty(shape: Shape, device: &DispatchDevice, dtype: BoolDType) -> BoolTensor { + creation_op!(Bool, device, |device| B::bool_empty(shape, device, dtype)) + } + + fn bool_zeros(shape: Shape, device: &DispatchDevice, dtype: BoolDType) -> BoolTensor { + creation_op!(Bool, device, |device| B::bool_zeros(shape, device, dtype)) + } + + fn bool_ones(shape: Shape, device: &DispatchDevice, dtype: BoolDType) -> BoolTensor { + creation_op!(Bool, device, |device| B::bool_ones(shape, device, dtype)) + } + + async fn bool_into_data(tensor: BoolTensor) -> Result { + unary_op!(tensor, bool, |tensor| B::bool_into_data(tensor).await) + } + + fn bool_from_data(data: TensorData, device: &DispatchDevice) -> BoolTensor { + creation_op!(Bool, device, |device| B::bool_from_data(data, device)) + } + + fn bool_into_int(tensor: BoolTensor, out_dtype: IntDType) -> IntTensor { + unary_op!(tensor, bool, |tensor| B::bool_into_int(tensor, out_dtype) => Int) + } + + fn bool_into_float(tensor: BoolTensor, out_dtype: FloatDType) -> FloatTensor { + unary_op!(tensor, bool, |tensor| B::bool_into_float(tensor, out_dtype) => Float) + } + + fn bool_to_device(tensor: BoolTensor, device: &DispatchDevice) -> BoolTensor { + to_device!( + Bool, + bool, + tensor, + device, + bool_to_device, + |inner, device| { + let data = + burn_backend::read_sync(B1::bool_into_data(inner)).expect("Should read data"); + B2::bool_from_data(data, device) + } + ) + } + + fn bool_reshape(tensor: BoolTensor, shape: Shape) -> BoolTensor { + unary_op!(tensor, bool, |tensor| B::bool_reshape(tensor, shape) => Bool) + } + + fn bool_slice(tensor: BoolTensor, slices: &[Slice]) -> BoolTensor { + unary_op!(tensor, bool, |tensor| B::bool_slice(tensor, slices) => Bool) + } + + fn bool_slice_assign( + tensor: BoolTensor, + slices: &[Slice], + value: BoolTensor, + ) -> BoolTensor { + binary_op!((tensor, bool), (value, bool), |tensor, value| B::bool_slice_assign(tensor, slices, value) => Bool) + } + + fn bool_mask_where( + tensor: BoolTensor, + mask: BoolTensor, + value: BoolTensor, + ) -> BoolTensor { + multi_op!( + inputs[(tensor, bool), (mask, bool), (value, bool)], => Bool, + B::bool_mask_where(tensor, mask, value) + ) + } + + fn bool_mask_fill( + tensor: BoolTensor, + mask: BoolTensor, + value: Scalar, + ) -> BoolTensor { + binary_op!((tensor, bool), (mask, bool), |tensor, mask| B::bool_mask_fill(tensor, mask, value) => Bool) + } + + fn bool_gather( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + ) -> BoolTensor { + binary_op!((tensor, bool), (indices, int), |tensor, indices| B::bool_gather(dim, tensor, indices) => Bool) + } + + fn bool_scatter_or( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor { + multi_op!( + inputs[(tensor, bool), (indices, int), (value, bool)], => Bool, + B::bool_scatter_or(dim, tensor, indices, value) + ) + } + + fn bool_equal(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + binary_op!((lhs, bool), (rhs, bool), |lhs, rhs| B::bool_equal(lhs, rhs) => Bool) + } + + fn bool_equal_elem(lhs: BoolTensor, rhs: Scalar) -> BoolTensor { + unary_op!(lhs, bool, |lhs| B::bool_equal_elem(lhs, rhs) => Bool) + } + + fn bool_not(tensor: BoolTensor) -> BoolTensor { + unary_op!(tensor, bool, |tensor| B::bool_not(tensor) => Bool) + } + + fn bool_and(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + binary_op!((lhs, bool), (rhs, bool), |lhs, rhs| B::bool_and(lhs, rhs) => Bool) + } + + fn bool_or(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + binary_op!((lhs, bool), (rhs, bool), |lhs, rhs| B::bool_or(lhs, rhs) => Bool) + } + + fn bool_swap_dims(tensor: BoolTensor, dim1: usize, dim2: usize) -> BoolTensor { + unary_op!(tensor, bool, |tensor| B::bool_swap_dims(tensor, dim1, dim2) => Bool) + } + + fn bool_permute(tensor: BoolTensor, axes: &[usize]) -> BoolTensor { + unary_op!(tensor, bool, |tensor| B::bool_permute(tensor, axes) => Bool) + } + + fn bool_flip(tensor: BoolTensor, axes: &[usize]) -> BoolTensor { + unary_op!(tensor, bool, |tensor| B::bool_flip(tensor, axes) => Bool) + } + + fn bool_expand(tensor: BoolTensor, shape: Shape) -> BoolTensor { + unary_op!(tensor, bool, |tensor| B::bool_expand(tensor, shape) => Bool) + } + + fn bool_unfold( + tensor: BoolTensor, + dim: usize, + size: usize, + step: usize, + ) -> BoolTensor { + unary_op!(tensor, bool, |tensor| B::bool_unfold(tensor, dim, size, step) => Bool) + } + + fn bool_select( + tensor: BoolTensor, + dim: usize, + indices: IntTensor, + ) -> BoolTensor { + binary_op!((tensor, bool), (indices, int), |tensor, indices| B::bool_select(tensor, dim, indices) => Bool) + } + + fn bool_select_or( + tensor: BoolTensor, + dim: usize, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor { + multi_op!( + inputs[(tensor, bool), (indices, int), (value, bool)], => Bool, + B::bool_select_or(tensor, dim, indices, value) + ) + } + + fn bool_repeat_dim(tensor: BoolTensor, dim: usize, times: usize) -> BoolTensor { + unary_op!(tensor, bool, |tensor| B::bool_repeat_dim(tensor, dim, times) => Bool) + } + + fn bool_cat(tensors: Vec>, dim: usize) -> BoolTensor { + vec_op!(tensors, bool, |tensors| B::bool_cat(tensors, dim) => Bool) + } + + fn bool_not_equal(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + binary_op!((lhs, bool), (rhs, bool), |lhs, rhs| B::bool_not_equal(lhs, rhs) => Bool) + } + + fn bool_not_equal_elem(lhs: BoolTensor, rhs: Scalar) -> BoolTensor { + unary_op!(lhs, bool, |lhs| B::bool_not_equal_elem(lhs, rhs) => Bool) + } + + fn bool_xor(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + binary_op!((lhs, bool), (rhs, bool), |lhs, rhs| B::bool_xor(lhs, rhs) => Bool) + } + + fn bool_transpose(tensor: BoolTensor) -> BoolTensor { + unary_op!(tensor, bool, |tensor| B::bool_transpose(tensor) => Bool) + } + + fn bool_any(tensor: BoolTensor) -> BoolTensor { + unary_op!(tensor, bool, |tensor| B::bool_any(tensor) => Bool) + } + + fn bool_any_dim(tensor: BoolTensor, dim: usize) -> BoolTensor { + unary_op!(tensor, bool, |tensor| B::bool_any_dim(tensor, dim) => Bool) + } + + fn bool_all(tensor: BoolTensor) -> BoolTensor { + unary_op!(tensor, bool, |tensor| B::bool_all(tensor) => Bool) + } + + fn bool_all_dim(tensor: BoolTensor, dim: usize) -> BoolTensor { + unary_op!(tensor, bool, |tensor| B::bool_all_dim(tensor, dim) => Bool) + } + + async fn bool_argwhere(tensor: BoolTensor, out_dtype: IntDType) -> IntTensor { + unary_op!(tensor, bool, |tensor| B::bool_argwhere(tensor, out_dtype).await => Int) + } +} diff --git a/crates/burn-dispatch/src/ops/distributed.rs b/crates/burn-dispatch/src/ops/distributed.rs new file mode 100644 index 0000000..7ce9c6d --- /dev/null +++ b/crates/burn-dispatch/src/ops/distributed.rs @@ -0,0 +1,171 @@ +use alloc::vec::Vec; + +use burn_backend::{ + DeviceId, + distributed::{ + CollectiveTensor, DistributedConfig, DistributedOps, DistributedParams, ReduceOperation, + TensorRef, + }, + tensor::FloatTensor, +}; + +use crate::{Dispatch, DispatchDevice}; + +macro_rules! dispatch_distributed_devices_arms { + ( + $device:expr, + $devices:expr, + |$inner_devices:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => { + match $device { + // Autodiff arm first + #[cfg(feature = "autodiff")] + $crate::DispatchDevice::Autodiff(inner) => { + let inner_devices = $devices + .iter() + .map(|d| { + match &d { + #[cfg(feature = "autodiff")] + $crate::DispatchDevice::Autodiff(d_inner) => *d_inner.inner.clone(), + _ => unreachable!("All devices are expected to be of the same variant."), + } + }) + .collect::>(); + let inner_devices = inner_devices.as_slice(); + // Recursively dispatch on inner + dispatch_distributed_devices_arms!( + @autodiff + &**inner, + &*inner_devices, + |$inner_devices| $body; + $([$Backend, $cfg]),* + ) + }, + $( + #[cfg($cfg)] + $crate::DispatchDevice::$Backend(_) => { + type B = $crate::backends::$Backend; + let $inner_devices = $devices + .iter() + .map(|d| { + let DispatchDevice::$Backend(dev) = d else { + unreachable!("All devices are expected to be of the same variant.") + }; + dev.clone() + }) + .collect::>(); + $body + } + )* + other => panic!("Distributed operations are not supported for device {other:?}"), + } + }; + ( + @autodiff + $device:expr, + $devices:expr, + |$inner_devices:ident| $body:expr; + $([$Backend:ident, $cfg:meta]),* + ) => { + match $device { + $( + #[cfg($cfg)] + $crate::DispatchDevice::$Backend(_) => { + type B = $crate::backends::Autodiff<$crate::backends::$Backend>; + let $inner_devices = $devices + .iter() + .map(|d| { + let DispatchDevice::$Backend(dev) = d else { + unreachable!("All devices are expected to be of the same variant.") + }; + dev.clone() + }) + .collect::>(); + $body + } + )* + $crate::DispatchDevice::Autodiff(_) => panic!("Autodiff should not wrap an autodiff device."), + other => panic!("Distributed operations are not supported for device {other:?}"), + } + }; +} + +/// Dispatches an operation body based on the provided devices. +macro_rules! dispatch_distributed_devices { + ($device:expr, $devices:expr, |$inner_devices:ident| $body:expr) => { + distributed_backend_list!( + dispatch_distributed_devices_arms, + $device, + $devices, + |$inner_devices| $body + ) + }; +} + +// In builds without a collective-capable backend (Cuda/Remote), the distributed dispatch arms +// are all cfg'd out, leaving only a diverging fallback — so the captured arguments and trailing +// expressions are intentionally unused/unreachable. +#[allow(unused_variables, unreachable_code)] +impl DistributedOps for Dispatch { + fn start_communication_server(devices: &[DispatchDevice], config: DistributedConfig) { + if !devices.is_empty() { + let first = &devices[0]; + dispatch_distributed_devices!(first, devices, |inner_devices| { + B::start_communication_server(&inner_devices, config) + }); + } + } + + fn close_communication_server(device: &DispatchDevice) { + dispatch_device!(@distributed device, |device| { + B::close_communication_server(device) + }) + } + + fn register_sync_parameters( + device: &DispatchDevice, + sharded_param_ids: Vec, + ) { + dispatch_device!(@distributed device, |device| B::register_sync_parameters( + device, + sharded_param_ids, + )) + } + + fn submit_sync_collective(device: &DispatchDevice) { + dispatch_device!(@distributed device, |device| B::submit_sync_collective(device)) + } + + fn submit_gradient_sync(_tensor: TensorRef, _distributed_params: DistributedParams) { + unimplemented!() + } + + fn all_reduce( + tensor: FloatTensor, + op: ReduceOperation, + device_ids: Vec, + ) -> CollectiveTensor { + // Safety: we call `assume_resolved` only to wrap it in a new `CollectiveTensor`. + // Explicit type: the distributed dispatch only emits arms for collective-capable + // backends (Cuda, Remote), so a build with none of them leaves only the diverging + // fallback and the match would otherwise infer `!`. + let tensor: FloatTensor = unary_float!(@distributed tensor, float, |tensor| { + let collective_tensor = B::all_reduce(tensor, op, device_ids); + unsafe { collective_tensor.assume_resolved() } + } => Float); + CollectiveTensor::new(tensor) + } + + fn sync_collective(device: &DispatchDevice) { + dispatch_device!(@distributed device, |device| B::sync_collective(device)) + } + + unsafe fn comm_device(_tensor: &TensorRef) -> DispatchDevice { + unimplemented!() + } + + unsafe fn float_from_ref(_tensor: &TensorRef) -> FloatTensor { + unimplemented!() + } +} diff --git a/crates/burn-dispatch/src/ops/int_tensor.rs b/crates/burn-dispatch/src/ops/int_tensor.rs new file mode 100644 index 0000000..ddf65c2 --- /dev/null +++ b/crates/burn-dispatch/src/ops/int_tensor.rs @@ -0,0 +1,560 @@ +use alloc::vec::Vec; +use burn_backend::{ + BoolDType, ExecutionError, FloatDType, IntDType, Scalar, Shape, Slice, TensorData, + ops::IntTensorOps, + tensor::{BoolTensor, FloatTensor, IntTensor}, +}; + +use crate::{Dispatch, DispatchDevice}; + +impl IntTensorOps for Dispatch { + fn int_empty(shape: Shape, device: &DispatchDevice, dtype: IntDType) -> IntTensor { + creation_op!(Int, device, |device| B::int_empty(shape, device, dtype)) + } + + async fn int_into_data(tensor: IntTensor) -> Result { + unary_op!(tensor, int, |tensor| B::int_into_data(tensor).await) + } + + fn int_from_data(data: TensorData, device: &DispatchDevice) -> IntTensor { + creation_op!(Int, device, |device| B::int_from_data(data, device)) + } + + fn int_to_device(tensor: IntTensor, device: &DispatchDevice) -> IntTensor { + to_device!(Int, int, tensor, device, int_to_device, |inner, device| { + let data = burn_backend::read_sync(B1::int_into_data(inner)).expect("Should read data"); + B2::int_from_data(data, device) + }) + } + + fn int_reshape(tensor: IntTensor, shape: Shape) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_reshape(tensor, shape) => Int) + } + + fn int_slice(tensor: IntTensor, slices: &[Slice]) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_slice(tensor, slices) => Int) + } + + fn int_slice_assign( + tensor: IntTensor, + slices: &[Slice], + value: IntTensor, + ) -> IntTensor { + binary_op!((tensor, int), (value, int), |tensor, value| B::int_slice_assign(tensor, slices, value) => Int) + } + + fn int_into_float(tensor: IntTensor, out_dtype: FloatDType) -> FloatTensor { + unary_op!(tensor, int, |tensor| B::int_into_float(tensor, out_dtype) => Float) + } + + fn int_mask_where( + tensor: IntTensor, + mask: BoolTensor, + value: IntTensor, + ) -> IntTensor { + multi_op!( + inputs[(tensor, int), (mask, bool), (value, int)], => Int, + B::int_mask_where(tensor, mask, value) + ) + } + + fn int_mask_fill( + tensor: IntTensor, + mask: BoolTensor, + value: Scalar, + ) -> IntTensor { + binary_op!((tensor, int), (mask, bool), |tensor, mask| B::int_mask_fill(tensor, mask, value) => Int) + } + + fn int_gather( + dim: usize, + tensor: IntTensor, + indices: IntTensor, + ) -> IntTensor { + binary_op!((tensor, int), (indices, int), |tensor, indices| B::int_gather(dim, tensor, indices) => Int) + } + + fn int_scatter_add( + dim: usize, + tensor: IntTensor, + indices: IntTensor, + value: IntTensor, + ) -> IntTensor { + multi_op!( + inputs[(tensor, int), (indices, int), (value, int)], => Int, + B::int_scatter_add(dim, tensor, indices, value) + ) + } + + fn int_scatter_nd( + data: IntTensor, + indices: IntTensor, + values: IntTensor, + reduction: burn_backend::tensor::IndexingUpdateOp, + ) -> IntTensor { + multi_op!( + inputs[(data, int), (indices, int), (values, int)], => Int, + B::int_scatter_nd(data, indices, values, reduction) + ) + } + + fn int_gather_nd(data: IntTensor, indices: IntTensor) -> IntTensor { + binary_op!((data, int), (indices, int), |data, indices| B::int_gather_nd(data, indices) => Int) + } + + fn int_select( + tensor: IntTensor, + dim: usize, + indices: IntTensor, + ) -> IntTensor { + binary_op!((tensor, int), (indices, int), |tensor, indices| B::int_select(tensor, dim, indices) => Int) + } + + fn int_select_add( + tensor: IntTensor, + dim: usize, + indices: IntTensor, + value: IntTensor, + ) -> IntTensor { + multi_op!( + inputs[(tensor, int), (indices, int), (value, int)], => Int, + B::int_select_add(tensor, dim, indices, value) + ) + } + + fn int_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::int_equal(lhs, rhs, out_dtype) => Bool) + } + + fn int_equal_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor { + unary_op!(lhs, int, |lhs| B::int_equal_elem(lhs, rhs, out_dtype) => Bool) + } + + fn int_greater( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::int_greater(lhs, rhs, out_dtype) => Bool) + } + + fn int_greater_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + unary_op!(lhs, int, |lhs| B::int_greater_elem(lhs, rhs, out_dtype) => Bool) + } + + fn int_greater_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::int_greater_equal(lhs, rhs, out_dtype) => Bool) + } + + fn int_greater_equal_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + unary_op!(lhs, int, |lhs| B::int_greater_equal_elem(lhs, rhs, out_dtype) => Bool) + } + + fn int_lower( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::int_lower(lhs, rhs, out_dtype) => Bool) + } + + fn int_lower_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor { + unary_op!(lhs, int, |lhs| B::int_lower_elem(lhs, rhs, out_dtype) => Bool) + } + + fn int_lower_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::int_lower_equal(lhs, rhs, out_dtype) => Bool) + } + + fn int_lower_equal_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + unary_op!(lhs, int, |lhs| B::int_lower_equal_elem(lhs, rhs, out_dtype) => Bool) + } + + fn int_add(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::int_add(lhs, rhs) => Int) + } + + fn int_add_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + unary_op!(lhs, int, |lhs| B::int_add_scalar(lhs, rhs) => Int) + } + + fn int_sub(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::int_sub(lhs, rhs) => Int) + } + + fn int_sub_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + unary_op!(lhs, int, |lhs| B::int_sub_scalar(lhs, rhs) => Int) + } + + fn int_mul(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::int_mul(lhs, rhs) => Int) + } + + fn int_mul_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + unary_op!(lhs, int, |lhs| B::int_mul_scalar(lhs, rhs) => Int) + } + + fn int_div(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::int_div(lhs, rhs) => Int) + } + + fn int_div_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + unary_op!(lhs, int, |lhs| B::int_div_scalar(lhs, rhs) => Int) + } + + fn int_remainder(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::int_remainder(lhs, rhs) => Int) + } + + fn int_remainder_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + unary_op!(lhs, int, |lhs| B::int_remainder_scalar(lhs, rhs) => Int) + } + + fn int_matmul(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::int_matmul(lhs, rhs) => Int) + } + + fn int_sum(tensor: IntTensor) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_sum(tensor) => Int) + } + + fn int_sum_dim(tensor: IntTensor, dim: usize) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_sum_dim(tensor, dim) => Int) + } + + fn int_prod(tensor: IntTensor) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_prod(tensor) => Int) + } + + fn int_prod_dim(tensor: IntTensor, dim: usize) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_prod_dim(tensor, dim) => Int) + } + + fn int_mean_dim(tensor: IntTensor, dim: usize) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_mean_dim(tensor, dim) => Int) + } + + fn int_cumsum(tensor: IntTensor, dim: usize) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_cumsum(tensor, dim) => Int) + } + + fn int_cumprod(tensor: IntTensor, dim: usize) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_cumprod(tensor, dim) => Int) + } + + fn int_cummin(tensor: IntTensor, dim: usize) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_cummin(tensor, dim) => Int) + } + + fn int_cummax(tensor: IntTensor, dim: usize) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_cummax(tensor, dim) => Int) + } + + fn int_argmax(tensor: IntTensor, dim: usize) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_argmax(tensor, dim) => Int) + } + + fn int_argtopk(tensor: IntTensor, dim: usize, k: usize) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_argtopk(tensor, dim, k) => Int) + } + + fn int_topk(tensor: IntTensor, dim: usize, k: usize) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_topk(tensor, dim, k) => Int) + } + + fn int_argmin(tensor: IntTensor, dim: usize) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_argmin(tensor, dim) => Int) + } + + fn int_abs(tensor: IntTensor) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_abs(tensor) => Int) + } + + fn int_swap_dims(tensor: IntTensor, dim1: usize, dim2: usize) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_swap_dims(tensor, dim1, dim2) => Int) + } + + fn int_permute(tensor: IntTensor, axes: &[usize]) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_permute(tensor, axes) => Int) + } + + fn int_flip(tensor: IntTensor, axes: &[usize]) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_flip(tensor, axes) => Int) + } + + fn int_random( + shape: Shape, + distribution: burn_backend::Distribution, + device: &DispatchDevice, + dtype: IntDType, + ) -> IntTensor { + creation_op!(Int, device, |device| { + B::int_random(shape, distribution, device, dtype) + }) + } + + fn int_expand(tensor: IntTensor, shape: Shape) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_expand(tensor, shape) => Int) + } + + fn bitwise_and(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::bitwise_and(lhs, rhs) => Int) + } + + fn bitwise_and_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + unary_op!(lhs, int, |lhs| B::bitwise_and_scalar(lhs, rhs) => Int) + } + + fn bitwise_or(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::bitwise_or(lhs, rhs) => Int) + } + + fn bitwise_or_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + unary_op!(lhs, int, |lhs| B::bitwise_or_scalar(lhs, rhs) => Int) + } + + fn bitwise_xor(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::bitwise_xor(lhs, rhs) => Int) + } + + fn bitwise_xor_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + unary_op!(lhs, int, |lhs| B::bitwise_xor_scalar(lhs, rhs) => Int) + } + + fn bitwise_not(tensor: IntTensor) -> IntTensor { + unary_op!(tensor, int, |tensor| B::bitwise_not(tensor) => Int) + } + + fn bitwise_left_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::bitwise_left_shift(lhs, rhs) => Int) + } + + fn bitwise_left_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + unary_op!(lhs, int, |lhs| B::bitwise_left_shift_scalar(lhs, rhs) => Int) + } + + fn bitwise_right_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::bitwise_right_shift(lhs, rhs) => Int) + } + + fn bitwise_right_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + unary_op!(lhs, int, |lhs| B::bitwise_right_shift_scalar(lhs, rhs) => Int) + } + + fn int_cast(tensor: IntTensor, dtype: IntDType) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_cast(tensor, dtype) => Int) + } + + fn int_unfold( + tensor: IntTensor, + dim: usize, + size: usize, + step: usize, + ) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_unfold(tensor, dim, size, step) => Int) + } + + fn int_repeat_dim(tensor: IntTensor, dim: usize, times: usize) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_repeat_dim(tensor, dim, times) => Int) + } + + fn int_cat(tensors: Vec>, dim: usize) -> IntTensor { + vec_op!(tensors, int, |tensors| B::int_cat(tensors, dim) => Int) + } + + fn int_not_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::int_not_equal(lhs, rhs, out_dtype) => Bool) + } + + fn int_not_equal_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + unary_op!(lhs, int, |lhs| B::int_not_equal_elem(lhs, rhs, out_dtype) => Bool) + } + + fn int_powi(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_op!((lhs, int), (rhs, int), |lhs, rhs| B::int_powi(lhs, rhs) => Int) + } + + fn int_powi_scalar_impl(lhs: IntTensor, rhs: Scalar) -> IntTensor { + unary_op!(lhs, int, |lhs| B::int_powi_scalar_impl(lhs, rhs) => Int) + } + + fn int_clamp_min(tensor: IntTensor, min: Scalar) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_clamp_min(tensor, min) => Int) + } + + fn int_clamp_max(tensor: IntTensor, max: Scalar) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_clamp_max(tensor, max) => Int) + } + + fn int_clamp(tensor: IntTensor, min: Scalar, max: Scalar) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_clamp(tensor, min, max) => Int) + } + + fn int_neg(tensor: IntTensor) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_neg(tensor) => Int) + } + + fn int_zeros(shape: Shape, device: &DispatchDevice, dtype: IntDType) -> IntTensor { + creation_op!(Int, device, |device| B::int_zeros(shape, device, dtype)) + } + + fn int_ones(shape: Shape, device: &DispatchDevice, dtype: IntDType) -> IntTensor { + creation_op!(Int, device, |device| B::int_ones(shape, device, dtype)) + } + + fn int_full( + shape: Shape, + fill_value: Scalar, + device: &DispatchDevice, + dtype: IntDType, + ) -> IntTensor { + creation_op!(Int, device, |device| B::int_full( + shape, fill_value, device, dtype + )) + } + + fn int_mean(tensor: IntTensor) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_mean(tensor) => Int) + } + + fn int_max(tensor: IntTensor) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_max(tensor) => Int) + } + + fn int_max_dim(tensor: IntTensor, dim: usize) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_max_dim(tensor, dim) => Int) + } + + fn int_max_dim_with_indices( + tensor: IntTensor, + dim: usize, + ) -> (IntTensor, IntTensor) { + multi_op!( + inputs[(tensor, int)], + outputs[(out, Int), (indices, Int)], + B::int_max_dim_with_indices(tensor, dim) + ) + } + + fn int_max_abs(tensor: IntTensor) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_max_abs(tensor) => Int) + } + + fn int_max_abs_dim(tensor: IntTensor, dim: usize) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_max_abs_dim(tensor, dim) => Int) + } + + fn int_min(tensor: IntTensor) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_min(tensor) => Int) + } + + fn int_min_dim(tensor: IntTensor, dim: usize) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_min_dim(tensor, dim) => Int) + } + + fn int_min_dim_with_indices( + tensor: IntTensor, + dim: usize, + ) -> (IntTensor, IntTensor) { + multi_op!( + inputs[(tensor, int)], + outputs[(out, Int), (indices, Int)], + B::int_min_dim_with_indices(tensor, dim) + ) + } + + fn int_transpose(tensor: IntTensor) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_transpose(tensor) => Int) + } + + fn int_arange_step( + range: core::ops::Range, + step: usize, + device: &DispatchDevice, + dtype: IntDType, + ) -> IntTensor { + creation_op!(Int, device, |device| B::int_arange_step( + range, step, device, dtype + )) + } + + fn int_arange( + range: core::ops::Range, + device: &DispatchDevice, + dtype: IntDType, + ) -> IntTensor { + creation_op!(Int, device, |device| B::int_arange(range, device, dtype)) + } + + fn int_any(tensor: IntTensor, out_dtype: BoolDType) -> BoolTensor { + unary_op!(tensor, int, |tensor| B::int_any(tensor, out_dtype) => Bool) + } + + fn int_any_dim(tensor: IntTensor, dim: usize, out_dtype: BoolDType) -> BoolTensor { + unary_op!(tensor, int, |tensor| B::int_any_dim(tensor, dim, out_dtype) => Bool) + } + + fn int_all(tensor: IntTensor, out_dtype: BoolDType) -> BoolTensor { + unary_op!(tensor, int, |tensor| B::int_all(tensor, out_dtype) => Bool) + } + + fn int_all_dim(tensor: IntTensor, dim: usize, out_dtype: BoolDType) -> BoolTensor { + unary_op!(tensor, int, |tensor| B::int_all_dim(tensor, dim, out_dtype) => Bool) + } + + fn int_sign(tensor: IntTensor) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_sign(tensor) => Int) + } + + fn int_sort(tensor: IntTensor, dim: usize, descending: bool) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_sort(tensor, dim, descending) => Int) + } + + fn int_sort_with_indices( + tensor: IntTensor, + dim: usize, + descending: bool, + ) -> (IntTensor, IntTensor) { + multi_op!( + inputs[(tensor, int)], + outputs[(out, Int), (indices, Int)], + B::int_sort_with_indices(tensor, dim, descending) + ) + } + + fn int_argsort(tensor: IntTensor, dim: usize, descending: bool) -> IntTensor { + unary_op!(tensor, int, |tensor| B::int_argsort(tensor, dim, descending) => Int) + } +} diff --git a/crates/burn-dispatch/src/ops/mod.rs b/crates/burn-dispatch/src/ops/mod.rs new file mode 100644 index 0000000..5ce9d98 --- /dev/null +++ b/crates/burn-dispatch/src/ops/mod.rs @@ -0,0 +1,8 @@ +mod activation; +mod bool_tensor; +mod distributed; +mod int_tensor; +mod module; +mod qtensor; +mod tensor; +mod transaction; diff --git a/crates/burn-dispatch/src/ops/module.rs b/crates/burn-dispatch/src/ops/module.rs new file mode 100644 index 0000000..2d2044c --- /dev/null +++ b/crates/burn-dispatch/src/ops/module.rs @@ -0,0 +1,722 @@ +use burn_backend::{ + IntDType, + ops::{ + DeformConv2dBackward, MaxPool1dBackward, MaxPool1dWithIndices, MaxPool2dBackward, + MaxPool2dWithIndices, ModuleOps, + }, + tensor::{FloatTensor, IntTensor}, +}; + +use crate::Dispatch; + +impl ModuleOps for Dispatch { + fn conv2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: burn_backend::ops::ConvOptions<2>, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (weight, float)], + opt_inputs[(bias, float)], + => Float, + B::conv2d(x, weight, bias, options) + ) + } + + fn deform_conv2d( + x: FloatTensor, + offset: FloatTensor, + weight: FloatTensor, + mask: Option>, + bias: Option>, + options: burn_backend::ops::DeformConvOptions<2>, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (offset, float), (weight, float)], + opt_inputs[(mask, float), (bias, float)], + => Float, + B::deform_conv2d(x, offset, weight, mask, bias, options) + ) + } + + fn deform_conv2d_backward( + x: FloatTensor, + offset: FloatTensor, + weight: FloatTensor, + mask: Option>, + bias: Option>, + output_grad: FloatTensor, + options: burn_backend::ops::DeformConvOptions<2>, + ) -> DeformConv2dBackward { + let (x_grad, offset_grad, weight_grad, mask_grad, bias_grad) = multi_op!( + inputs[(x, float), (offset, float), (weight, float), (output_grad, float)], + opt_inputs[(mask, float), (bias, float)], + outputs[(x_grad, Float), (offset_grad, Float), (weight_grad, Float)], + opt_outputs[mask_grad, bias_grad], + { + let res = B::deform_conv2d_backward(x, offset, weight, mask, bias, output_grad, options); + (res.x_grad, res.offset_grad, res.weight_grad, res.mask_grad, res.bias_grad) + } + ); + DeformConv2dBackward::new(x_grad, offset_grad, weight_grad, mask_grad, bias_grad) + } + + fn conv3d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: burn_backend::ops::ConvOptions<3>, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (weight, float)], + opt_inputs[(bias, float)], + => Float, + B::conv3d(x, weight, bias, options) + ) + } + + fn conv_transpose2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: burn_backend::ops::ConvTransposeOptions<2>, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (weight, float)], + opt_inputs[(bias, float)], + => Float, + B::conv_transpose2d(x, weight, bias, options) + ) + } + + fn conv_transpose3d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: burn_backend::ops::ConvTransposeOptions<3>, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (weight, float)], + opt_inputs[(bias, float)], + => Float, + B::conv_transpose3d(x, weight, bias, options) + ) + } + + fn avg_pool2d( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + multi_op!(inputs[(x, float)], + => Float, + B::avg_pool2d(x, kernel_size, stride, padding, count_include_pad, ceil_mode) + ) + } + + fn avg_pool2d_backward( + x: FloatTensor, + grad: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (grad, float)], + => Float, + B::avg_pool2d_backward(x, grad, kernel_size, stride, padding, count_include_pad, ceil_mode) + ) + } + + fn adaptive_avg_pool2d(x: FloatTensor, output_size: [usize; 2]) -> FloatTensor { + multi_op!( + inputs[(x, float)], + => Float, + B::adaptive_avg_pool2d(x, output_size) + ) + } + + fn adaptive_avg_pool2d_backward( + x: FloatTensor, + grad: FloatTensor, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (grad, float)], + => Float, + B::adaptive_avg_pool2d_backward(x, grad) + ) + } + + fn max_pool2d( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + ) -> FloatTensor { + multi_op!( + inputs[(x, float)], + => Float, + B::max_pool2d(x, kernel_size, stride, padding, dilation, ceil_mode) + ) + } + + fn max_pool2d_with_indices( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + indices_dtype: IntDType, + ) -> MaxPool2dWithIndices { + let (out, indices) = multi_op!( + inputs[(x, float)], + outputs[(out, Float), (indices, Int)], + { + let res = B::max_pool2d_with_indices(x, kernel_size, stride, padding, dilation, ceil_mode, indices_dtype); + (res.output, res.indices) + } + ); + MaxPool2dWithIndices::new(out, indices) + } + + fn max_pool2d_with_indices_backward( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + output_grad: FloatTensor, + indices: IntTensor, + ) -> MaxPool2dBackward { + let x_grad = multi_op!( + inputs[(x, float), (output_grad, float), (indices, int)], + => Float, + { + let res = B::max_pool2d_with_indices_backward(x, kernel_size, stride, padding, dilation, ceil_mode, output_grad, indices); + res.x_grad + } + ); + MaxPool2dBackward::new(x_grad) + } + + fn interpolate( + x: FloatTensor, + output_size: [usize; 2], + options: burn_backend::ops::InterpolateOptions, + ) -> FloatTensor { + multi_op!( + inputs[(x, float)], + => Float, + B::interpolate(x, output_size, options) + ) + } + + fn interpolate_backward( + x: FloatTensor, + grad: FloatTensor, + output_size: [usize; 2], + options: burn_backend::ops::InterpolateOptions, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (grad, float)], + => Float, + B::interpolate_backward(x, grad, output_size, options) + ) + } + + fn embedding(weights: FloatTensor, indices: IntTensor) -> FloatTensor { + multi_op!( + inputs[(weights, float), (indices, int)], + => Float, + B::embedding(weights, indices) + ) + } + + fn embedding_backward( + weights: FloatTensor, + output_grad: FloatTensor, + indices: IntTensor, + ) -> FloatTensor { + multi_op!( + inputs[(weights, float), (output_grad, float), (indices, int)], + => Float, + B::embedding_backward(weights, output_grad, indices) + ) + } + + fn conv1d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: burn_backend::ops::ConvOptions<1>, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (weight, float)], + opt_inputs[(bias, float)], + => Float, + B::conv1d(x, weight, bias, options) + ) + } + + fn conv1d_x_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: burn_backend::ops::ConvOptions<1>, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (weight, float), (output_grad, float)], + => Float, + B::conv1d_x_backward(x, weight, output_grad, options) + ) + } + + fn conv1d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: burn_backend::ops::ConvOptions<1>, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (weight, float), (output_grad, float)], + => Float, + B::conv1d_weight_backward(x, weight, output_grad, options) + ) + } + + fn conv1d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (bias, float), (output_grad, float)], + => Float, + B::conv1d_bias_backward(x, bias, output_grad) + ) + } + + fn conv2d_x_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: burn_backend::ops::ConvOptions<2>, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (weight, float), (output_grad, float)], + => Float, + B::conv2d_x_backward(x, weight, output_grad, options) + ) + } + + fn conv2d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: burn_backend::ops::ConvOptions<2>, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (weight, float), (output_grad, float)], + => Float, + B::conv2d_weight_backward(x, weight, output_grad, options) + ) + } + + fn conv2d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (bias, float), (output_grad, float)], + => Float, + B::conv2d_bias_backward(x, bias, output_grad) + ) + } + + fn conv3d_x_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: burn_backend::ops::ConvOptions<3>, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (weight, float), (output_grad, float)], + => Float, + B::conv3d_x_backward(x, weight, output_grad, options) + ) + } + + fn conv3d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: burn_backend::ops::ConvOptions<3>, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (weight, float), (output_grad, float)], + => Float, + B::conv3d_weight_backward(x, weight, output_grad, options) + ) + } + + fn conv3d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (bias, float), (output_grad, float)], + => Float, + B::conv3d_bias_backward(x, bias, output_grad) + ) + } + + fn conv_transpose1d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: burn_backend::ops::ConvTransposeOptions<1>, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (weight, float)], + opt_inputs[(bias, float)], + => Float, + B::conv_transpose1d(x, weight, bias, options) + ) + } + + fn conv_transpose1d_x_backward( + weight: FloatTensor, + output_grad: FloatTensor, + options: burn_backend::ops::ConvTransposeOptions<1>, + ) -> FloatTensor { + multi_op!( + inputs[(weight, float), (output_grad, float)], + => Float, + B::conv_transpose1d_x_backward(weight, output_grad, options) + ) + } + + fn conv_transpose1d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: burn_backend::ops::ConvTransposeOptions<1>, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (weight, float), (output_grad, float)], + => Float, + B::conv_transpose1d_weight_backward(x, weight, output_grad, options) + ) + } + + fn conv_transpose1d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (bias, float), (output_grad, float)], + => Float, + B::conv_transpose1d_bias_backward(x, bias, output_grad) + ) + } + + fn conv_transpose2d_x_backward( + weight: FloatTensor, + output_grad: FloatTensor, + options: burn_backend::ops::ConvTransposeOptions<2>, + ) -> FloatTensor { + multi_op!( + inputs[(weight, float), (output_grad, float)], + => Float, + B::conv_transpose2d_x_backward(weight, output_grad, options) + ) + } + + fn conv_transpose2d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: burn_backend::ops::ConvTransposeOptions<2>, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (weight, float), (output_grad, float)], + => Float, + B::conv_transpose2d_weight_backward(x, weight, output_grad, options) + ) + } + + fn conv_transpose2d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (bias, float), (output_grad, float)], + => Float, + B::conv_transpose2d_bias_backward(x, bias, output_grad) + ) + } + + fn conv_transpose3d_x_backward( + weight: FloatTensor, + output_grad: FloatTensor, + options: burn_backend::ops::ConvTransposeOptions<3>, + ) -> FloatTensor { + multi_op!( + inputs[(weight, float), (output_grad, float)], + => Float, + B::conv_transpose3d_x_backward(weight, output_grad, options) + ) + } + + fn conv_transpose3d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: burn_backend::ops::ConvTransposeOptions<3>, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (weight, float), (output_grad, float)], + => Float, + B::conv_transpose3d_weight_backward(x, weight, output_grad, options) + ) + } + + fn conv_transpose3d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (bias, float), (output_grad, float)], + => Float, + B::conv_transpose3d_bias_backward(x, bias, output_grad) + ) + } + + fn unfold4d( + x: FloatTensor, + kernel_size: [usize; 2], + options: burn_backend::ops::UnfoldOptions, + ) -> FloatTensor { + multi_op!(inputs[(x, float)], => Float, B::unfold4d(x, kernel_size, options)) + } + + fn avg_pool1d( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + multi_op!(inputs[(x, float)], => Float, + B::avg_pool1d(x, kernel_size, stride, padding, count_include_pad, ceil_mode) + ) + } + + fn avg_pool1d_backward( + x: FloatTensor, + grad: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (grad, float)], + => Float, + B::avg_pool1d_backward(x, grad, kernel_size, stride, padding, count_include_pad, ceil_mode) + ) + } + + fn adaptive_avg_pool1d(x: FloatTensor, output_size: usize) -> FloatTensor { + multi_op!(inputs[(x, float)], => Float, B::adaptive_avg_pool1d(x, output_size)) + } + + fn adaptive_avg_pool1d_backward( + x: FloatTensor, + grad: FloatTensor, + ) -> FloatTensor { + multi_op!( + inputs[(x, float), (grad, float)], + => Float, + B::adaptive_avg_pool1d_backward(x, grad) + ) + } + + fn max_pool1d( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + ) -> FloatTensor { + multi_op!(inputs[(x, float)], => Float, + B::max_pool1d(x, kernel_size, stride, padding, dilation, ceil_mode)) + } + + fn max_pool1d_with_indices( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + indices_dtype: IntDType, + ) -> MaxPool1dWithIndices { + let (out, indices) = multi_op!( + inputs[(x, float)], + outputs[(out, Float), (indices, Int)], + { + let res = B::max_pool1d_with_indices(x, kernel_size, stride, padding, dilation, ceil_mode, indices_dtype); + (res.output, res.indices) + } + ); + MaxPool1dWithIndices::new(out, indices) + } + + fn max_pool1d_with_indices_backward( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + output_grad: FloatTensor, + indices: IntTensor, + ) -> MaxPool1dBackward { + let x_grad = multi_op!( + inputs[(x, float), (output_grad, float), (indices, int)], + => Float, + { + let res = B::max_pool1d_with_indices_backward(x, kernel_size, stride, padding, dilation, ceil_mode, output_grad, indices); + res.x_grad + } + ); + MaxPool1dBackward::new(x_grad) + } + + fn attention( + query: FloatTensor, + key: FloatTensor, + value: FloatTensor, + mask: Option>, + attn_bias: Option>, + options: burn_backend::ops::AttentionModuleOptions, + ) -> FloatTensor { + multi_op!( + inputs[(query, float), (key, float), (value, float)], + opt_inputs[(mask, bool), (attn_bias, float)], + => Float, + B::attention(query, key, value, mask, attn_bias, options) + ) + } + + fn layer_norm( + tensor: FloatTensor, + gamma: FloatTensor, + beta: Option>, + epsilon: f64, + ) -> FloatTensor { + multi_op!( + inputs[(tensor, float), (gamma, float)], + opt_inputs[(beta, float)], + => Float, + B::layer_norm(tensor, gamma, beta, epsilon) + ) + } + + fn rfft( + signal: FloatTensor, + dim: usize, + n: Option, + ) -> (FloatTensor, FloatTensor) { + let (real, imag) = multi_op!( + inputs[(signal, float)], + outputs[(real, Float), (imag, Float)], + { + let res = B::rfft(signal, dim, n); + (res.0, res.1) + } + ); + + (real, imag) + } + + fn irfft( + spectrum_re: FloatTensor, + spectrum_im: FloatTensor, + dim: usize, + n: Option, + ) -> FloatTensor { + multi_op!( + inputs[(spectrum_re, float), (spectrum_im, float)], + => Float, + { + B::irfft(spectrum_re, spectrum_im, dim, n) + } + ) + } + + fn has_ctc_loss_backward() -> bool { + // Dispatch routes per-tensor at runtime, but autodiff queries this flag + // statically. Returning `false` makes autodiff differentiate through + // the default decomposed forward, which is safe for every inner + // backend regardless of whether it has its own ctc_loss_backward. + false + } + + fn ctc_loss( + log_probs: FloatTensor, + targets: IntTensor, + input_lengths: IntTensor, + target_lengths: IntTensor, + blank: usize, + ) -> FloatTensor { + multi_op!( + inputs[(log_probs, float), (targets, int), (input_lengths, int), (target_lengths, int)], + => Float, + B::ctc_loss(log_probs, targets, input_lengths, target_lengths, blank) + ) + } + + fn ctc_loss_backward( + log_probs: FloatTensor, + targets: IntTensor, + input_lengths: IntTensor, + target_lengths: IntTensor, + grad_loss: FloatTensor, + blank: usize, + ) -> FloatTensor { + multi_op!( + inputs[(log_probs, float), (targets, int), (input_lengths, int), (target_lengths, int), (grad_loss, float)], + => Float, + B::ctc_loss_backward(log_probs, targets, input_lengths, target_lengths, grad_loss, blank) + ) + } + + // TODO: linear ops + // fn linear( + // x: FloatTensor, + // weight: FloatTensor, + // bias: Option>, + // ) -> FloatTensor { + + // } +} diff --git a/crates/burn-dispatch/src/ops/qtensor.rs b/crates/burn-dispatch/src/ops/qtensor.rs new file mode 100644 index 0000000..b59dd88 --- /dev/null +++ b/crates/burn-dispatch/src/ops/qtensor.rs @@ -0,0 +1,210 @@ +use burn_backend::{ + DeviceOps, ExecutionError, FloatDType, Shape, Slice, TensorData, TensorMetadata, + TensorPrimitive, + ops::QTensorOps, + quantization::{QuantPropagation, QuantScheme, QuantizationParametersPrimitive}, + tensor::{FloatTensor, IntTensor, QuantizedTensor}, +}; + +use crate::{Dispatch, DispatchDevice}; + +impl QTensorOps for Dispatch { + fn q_from_data(data: TensorData, device: &DispatchDevice) -> QuantizedTensor { + creation_op!(Quantized, device, |device| B::q_from_data(data, device)) + } + + fn quantize( + tensor: FloatTensor, + scheme: &QuantScheme, + qparams: QuantizationParametersPrimitive, + ) -> QuantizedTensor { + binary_op!( + (tensor, float), + (qparams.scales, float), + |tensor, scales| { + B::quantize(tensor, scheme, QuantizationParametersPrimitive { scales }) + } => Quantized + ) + } + + fn dequantize(tensor: QuantizedTensor, dtype: FloatDType) -> FloatTensor { + unary_op!(tensor, quantized, |tensor| B::dequantize(tensor, dtype) => Float) + } + + fn q_to_device( + tensor: QuantizedTensor, + device: &DispatchDevice, + ) -> QuantizedTensor { + to_device!( + Quantized, + quantized, + tensor, + device, + q_to_device, + |inner, device| { + let data = + burn_backend::read_sync(B1::q_into_data(inner)).expect("Should read data"); + B2::q_from_data(data, device) + } + ) + } + + fn q_reshape(tensor: QuantizedTensor, shape: Shape) -> QuantizedTensor { + unary_op!(tensor, quantized, |tensor| B::q_reshape(tensor, shape) => Quantized) + } + + async fn q_into_data(tensor: QuantizedTensor) -> Result { + unary_op!(tensor, quantized, |tensor| B::q_into_data(tensor).await) + } + + fn q_expand(tensor: QuantizedTensor, shape: Shape) -> QuantizedTensor { + unary_op!(tensor, quantized, |tensor| B::q_expand(tensor, shape) => Quantized) + } + + fn q_swap_dims( + tensor: QuantizedTensor, + dim1: usize, + dim2: usize, + ) -> QuantizedTensor { + unary_op!(tensor, quantized, |tensor| B::q_swap_dims(tensor, dim1, dim2) => Quantized) + } + + fn q_permute(tensor: QuantizedTensor, axes: &[usize]) -> QuantizedTensor { + unary_op!(tensor, quantized, |tensor| B::q_permute(tensor, axes) => Quantized) + } + + fn q_flip(tensor: QuantizedTensor, axes: &[usize]) -> QuantizedTensor { + unary_op!(tensor, quantized, |tensor| B::q_flip(tensor, axes) => Quantized) + } + + fn q_select( + tensor: QuantizedTensor, + dim: usize, + indices: IntTensor, + ) -> QuantizedTensor { + binary_op!( + (tensor, quantized), + (indices, int), + |tensor, indices| B::q_select(tensor, dim, indices) => Quantized + ) + } + + fn q_slice(tensor: QuantizedTensor, slices: &[Slice]) -> QuantizedTensor { + unary_op!(tensor, quantized, |tensor| B::q_slice(tensor, slices) => Quantized) + } + + fn q_matmul(lhs: TensorPrimitive, rhs: TensorPrimitive) -> TensorPrimitive { + // TODO: this would be much cleaner if we consolidated tensor primitive types + match (lhs, rhs) { + (TensorPrimitive::QFloat(lhs), TensorPrimitive::QFloat(rhs)) => { + let propagation = lhs.device().defaults().quantization.propagation; + if matches!(propagation, QuantPropagation::Propagate) { + let out = binary_op!( + (lhs, quantized), + (rhs, quantized), + |lhs, rhs| { + if let TensorPrimitive::QFloat(out) = B::q_matmul( + TensorPrimitive::QFloat(lhs), + TensorPrimitive::QFloat(rhs), + ) { + out + } else { + unreachable!() + } + } => Quantized + ); + TensorPrimitive::QFloat(out) + } else { + let out = binary_op!( + (lhs, quantized), + (rhs, quantized), + |lhs, rhs| { + if let TensorPrimitive::Float(out) = B::q_matmul( + TensorPrimitive::QFloat(lhs), + TensorPrimitive::QFloat(rhs), + ) { + out + } else { + unreachable!() + } + } => Float + ); + TensorPrimitive::Float(out) + } + } + (TensorPrimitive::Float(lhs), TensorPrimitive::QFloat(rhs)) => { + let propagation = rhs.device().defaults().quantization.propagation; + if matches!(propagation, QuantPropagation::Propagate) { + let out = binary_op!( + (lhs, float), + (rhs, quantized), + |lhs, rhs| { + if let TensorPrimitive::QFloat(out) = B::q_matmul( + TensorPrimitive::Float(lhs), + TensorPrimitive::QFloat(rhs), + ) { + out + } else { + unreachable!() + } + } => Quantized + ); + TensorPrimitive::QFloat(out) + } else { + let out = binary_op!( + (lhs, float), + (rhs, quantized), + |lhs, rhs| { + if let TensorPrimitive::Float(out) = B::q_matmul( + TensorPrimitive::Float(lhs), + TensorPrimitive::QFloat(rhs), + ) { + out + } else { + unreachable!() + } + } => Float + ); + TensorPrimitive::Float(out) + } + } + (TensorPrimitive::QFloat(lhs), TensorPrimitive::Float(rhs)) => { + let propagation = lhs.device().defaults().quantization.propagation; + if matches!(propagation, QuantPropagation::Propagate) { + let out = binary_op!( + (lhs, quantized), + (rhs, float), + |lhs, rhs| { + if let TensorPrimitive::QFloat(out) = B::q_matmul( + TensorPrimitive::QFloat(lhs), + TensorPrimitive::Float(rhs), + ) { + out + } else { + unreachable!() + } + } => Quantized + ); + TensorPrimitive::QFloat(out) + } else { + let out = binary_op!( + (lhs, quantized), + (rhs, float), + |lhs, rhs| { + if let TensorPrimitive::Float(out) = B::q_matmul( + TensorPrimitive::QFloat(lhs), + TensorPrimitive::Float(rhs), + ) { + out + } else { + unreachable!() + } + } => Float + ); + TensorPrimitive::Float(out) + } + } + _ => unreachable!(), + } + } +} diff --git a/crates/burn-dispatch/src/ops/tensor.rs b/crates/burn-dispatch/src/ops/tensor.rs new file mode 100644 index 0000000..a6d58d3 --- /dev/null +++ b/crates/burn-dispatch/src/ops/tensor.rs @@ -0,0 +1,696 @@ +use alloc::vec::Vec; +use burn_backend::{ + BoolDType, ExecutionError, FloatDType, IntDType, Scalar, Shape, Slice, TensorData, + ops::FloatTensorOps, + tensor::{BoolTensor, FloatTensor, IntTensor}, +}; + +use crate::{Dispatch, DispatchDevice}; + +impl FloatTensorOps for Dispatch { + fn float_from_data( + data: burn_backend::TensorData, + device: &DispatchDevice, + ) -> FloatTensor { + creation_op!(Float, device, |device| B::float_from_data(data, device)) + } + + fn float_random( + shape: Shape, + distribution: burn_backend::Distribution, + device: &DispatchDevice, + dtype: FloatDType, + ) -> FloatTensor { + creation_op!(Float, device, |device| { + B::float_random(shape, distribution, device, dtype) + }) + } + + async fn float_into_data(tensor: FloatTensor) -> Result { + unary_float!(tensor, float, |tensor| B::float_into_data(tensor).await) + } + + fn float_to_device(tensor: FloatTensor, device: &DispatchDevice) -> FloatTensor { + // Relocating a non-tracked float tensor onto an autodiff device is a plain data move: + // place it on the underlying hardware device and leave the tensor non-tracked. The + // int/bool `to_device` paths already handle this case; only the float path used to + // panic. This is what lets gradient tensors — which are never autodiff-tracked — be + // moved onto the autodiff `device_main` during multi-device training. + #[cfg(feature = "autodiff")] + if let DispatchDevice::Autodiff(device_ad) = device + && !matches!(&tensor.kind, crate::DispatchTensorKind::Autodiff(_)) + { + return Self::float_to_device(tensor, &device_ad.inner); + } + + float_to_device!( + Float, + float, + tensor, + device, + float_to_device, + |inner, device| { + let data = + burn_backend::read_sync(B1::float_into_data(inner)).expect("Should read data"); + B2::float_from_data(data, device) + } + ) + } + + fn float_into_int(tensor: FloatTensor, dtype: burn_backend::IntDType) -> IntTensor { + unary_float!(tensor, float, |tensor| B::float_into_int(tensor, dtype) => Int) + } + + fn float_empty(shape: Shape, device: &DispatchDevice, dtype: FloatDType) -> FloatTensor { + creation_op!(Float, device, |device| B::float_empty(shape, device, dtype)) + } + + fn float_add(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float!((lhs, float), (rhs, float), |lhs, rhs| B::float_add(lhs, rhs) => Float) + } + + fn float_add_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + unary_float!(lhs, float, |lhs| B::float_add_scalar(lhs, rhs) => Float) + } + + fn float_sub(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float!((lhs, float), (rhs, float), |lhs, rhs| B::float_sub(lhs, rhs) => Float) + } + + fn float_sub_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + unary_float!(lhs, float, |lhs| B::float_sub_scalar(lhs, rhs) => Float) + } + + fn float_mul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float!((lhs, float), (rhs, float), |lhs, rhs| B::float_mul(lhs, rhs) => Float) + } + + fn float_mul_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + unary_float!(lhs, float, |lhs| B::float_mul_scalar(lhs, rhs) => Float) + } + + fn float_div(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float!((lhs, float), (rhs, float), |lhs, rhs| B::float_div(lhs, rhs) => Float) + } + + fn float_div_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + unary_float!(lhs, float, |lhs| B::float_div_scalar(lhs, rhs) => Float) + } + + fn float_remainder(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float!((lhs, float), (rhs, float), |lhs, rhs| B::float_remainder(lhs, rhs) => Float) + } + + fn float_remainder_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + unary_float!(lhs, float, |lhs| B::float_remainder_scalar(lhs, rhs) => Float) + } + + fn float_matmul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float!((lhs, float), (rhs, float), |lhs, rhs| B::float_matmul(lhs, rhs) => Float) + } + + fn float_cross( + lhs: FloatTensor, + rhs: FloatTensor, + dim: usize, + ) -> FloatTensor { + binary_float!((lhs, float), (rhs, float), |lhs, rhs| B::float_cross(lhs, rhs, dim) => Float) + } + + fn float_recip(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_recip(tensor) => Float) + } + + fn float_swap_dims(tensor: FloatTensor, dim1: usize, dim2: usize) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_swap_dims(tensor, dim1, dim2) => Float) + } + + fn float_permute(tensor: FloatTensor, axes: &[usize]) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_permute(tensor, axes) => Float) + } + + fn float_flip(tensor: FloatTensor, axes: &[usize]) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_flip(tensor, axes) => Float) + } + + fn float_reshape(tensor: FloatTensor, shape: Shape) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_reshape(tensor, shape) => Float) + } + + fn float_gather( + dim: usize, + tensor: FloatTensor, + indices: IntTensor, + ) -> FloatTensor { + binary_float!((tensor, float), (indices, int), |tensor, indices| B::float_gather(dim, tensor, indices) => Float) + } + + fn float_scatter_add( + dim: usize, + tensor: FloatTensor, + indices: IntTensor, + value: FloatTensor, + ) -> FloatTensor { + multi_op!( + inputs[(tensor, float), (indices, int), (value, float)], => Float, + B::float_scatter_add(dim, tensor, indices, value) + ) + } + + fn float_scatter_nd( + data: FloatTensor, + indices: IntTensor, + values: FloatTensor, + reduction: burn_backend::tensor::IndexingUpdateOp, + ) -> FloatTensor { + multi_op!( + inputs[(data, float), (indices, int), (values, float)], => Float, + B::float_scatter_nd(data, indices, values, reduction) + ) + } + + fn float_gather_nd(data: FloatTensor, indices: IntTensor) -> FloatTensor { + binary_float!((data, float), (indices, int), |data, indices| B::float_gather_nd(data, indices) => Float) + } + + fn float_select( + tensor: FloatTensor, + dim: usize, + indices: IntTensor, + ) -> FloatTensor { + binary_float!((tensor, float), (indices, int), |tensor, indices| B::float_select(tensor, dim, indices) => Float) + } + + fn float_select_add( + tensor: FloatTensor, + dim: usize, + indices: IntTensor, + value: FloatTensor, + ) -> FloatTensor { + multi_op!( + inputs[(tensor, float), (indices, int), (value, float)], => Float, + B::float_select_add(tensor, dim, indices, value) + ) + } + + fn float_slice(tensor: FloatTensor, slices: &[Slice]) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_slice(tensor, slices) => Float) + } + + fn float_slice_assign( + tensor: FloatTensor, + slices: &[Slice], + value: FloatTensor, + ) -> FloatTensor { + binary_float!((tensor, float), (value, float), |tensor, value| B::float_slice_assign(tensor, slices, value) => Float) + } + + fn float_mask_where( + tensor: FloatTensor, + mask: BoolTensor, + value: FloatTensor, + ) -> FloatTensor { + multi_op!( + inputs[(tensor, float), (mask, bool), (value, float)], => Float, + B::float_mask_where(tensor, mask, value) + ) + } + + fn float_mask_fill( + tensor: FloatTensor, + mask: BoolTensor, + value: Scalar, + ) -> FloatTensor { + binary_float!((tensor, float), (mask, bool), |tensor, mask| B::float_mask_fill(tensor, mask, value) => Float) + } + + fn float_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_float!((lhs, float), (rhs, float), |lhs, rhs| B::float_equal(lhs, rhs, out_dtype) => Bool) + } + + fn float_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + unary_float!(lhs, float, |lhs| B::float_equal_elem(lhs, rhs, out_dtype) => Bool) + } + + fn float_greater( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_float!((lhs, float), (rhs, float), |lhs, rhs| B::float_greater(lhs, rhs, out_dtype) => Bool) + } + + fn float_greater_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + unary_float!(lhs, float, |lhs| B::float_greater_elem(lhs, rhs, out_dtype) => Bool) + } + + fn float_greater_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_float!((lhs, float), (rhs, float), |lhs, rhs| B::float_greater_equal(lhs, rhs, out_dtype) => Bool) + } + + fn float_greater_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + unary_float!(lhs, float, |lhs| B::float_greater_equal_elem(lhs, rhs, out_dtype) => Bool) + } + + fn float_lower( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_float!((lhs, float), (rhs, float), |lhs, rhs| B::float_lower(lhs, rhs, out_dtype) => Bool) + } + + fn float_lower_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + unary_float!(lhs, float, |lhs| B::float_lower_elem(lhs, rhs, out_dtype) => Bool) + } + + fn float_lower_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_float!((lhs, float), (rhs, float), |lhs, rhs| B::float_lower_equal(lhs, rhs, out_dtype) => Bool) + } + + fn float_lower_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + unary_float!(lhs, float, |lhs| B::float_lower_equal_elem(lhs, rhs, out_dtype) => Bool) + } + + fn float_sum(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_sum(tensor) => Float) + } + + fn float_sum_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_sum_dim(tensor, dim) => Float) + } + + fn float_mean_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_mean_dim(tensor, dim) => Float) + } + + fn float_cumsum(tensor: FloatTensor, dim: usize) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_cumsum(tensor, dim) => Float) + } + + fn float_cumprod(tensor: FloatTensor, dim: usize) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_cumprod(tensor, dim) => Float) + } + + fn float_cummin(tensor: FloatTensor, dim: usize) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_cummin(tensor, dim) => Float) + } + + fn float_cummax(tensor: FloatTensor, dim: usize) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_cummax(tensor, dim) => Float) + } + + fn float_cast(tensor: FloatTensor, dtype: FloatDType) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_cast(tensor, dtype) => Float) + } + + fn float_exp(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_exp(tensor) => Float) + } + + fn float_log(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_log(tensor) => Float) + } + + fn float_log1p(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_log1p(tensor) => Float) + } + + fn float_powf(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float!((lhs, float), (rhs, float), |lhs, rhs| B::float_powf(lhs, rhs) => Float) + } + + fn float_powf_scalar_impl(tensor: FloatTensor, value: Scalar) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_powf_scalar_impl(tensor, value) => Float) + } + + fn float_sqrt(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_sqrt(tensor) => Float) + } + + fn float_abs(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_abs(tensor) => Float) + } + + fn float_cos(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_cos(tensor) => Float) + } + + fn float_sin(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_sin(tensor) => Float) + } + + fn float_tan(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_tan(tensor) => Float) + } + + fn float_cosh(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_cosh(tensor) => Float) + } + + fn float_sinh(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_sinh(tensor) => Float) + } + + fn float_tanh(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_tanh(tensor) => Float) + } + + fn float_acos(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_acos(tensor) => Float) + } + + fn float_acosh(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_acosh(tensor) => Float) + } + + fn float_asin(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_asin(tensor) => Float) + } + + fn float_asinh(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_asinh(tensor) => Float) + } + + fn float_atan(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_atan(tensor) => Float) + } + + fn float_atanh(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_atanh(tensor) => Float) + } + + fn float_atan2(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float!((lhs, float), (rhs, float), |lhs, rhs| B::float_atan2(lhs, rhs) => Float) + } + + fn float_round(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_round(tensor) => Float) + } + + fn float_floor(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_floor(tensor) => Float) + } + + fn float_ceil(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_ceil(tensor) => Float) + } + + fn float_trunc(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_trunc(tensor) => Float) + } + + fn float_erf(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_erf(tensor) => Float) + } + + fn float_argmax(tensor: FloatTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + unary_float!(tensor, float, |tensor| B::float_argmax(tensor, dim, out_dtype) => Int) + } + + fn float_argtopk( + tensor: FloatTensor, + dim: usize, + k: usize, + out_dtype: IntDType, + ) -> IntTensor { + unary_float!(tensor, float, |tensor| B::float_argtopk(tensor, dim, k, out_dtype) => Int) + } + + fn float_topk(tensor: FloatTensor, dim: usize, k: usize) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_topk(tensor, dim, k) => Float) + } + + fn float_argmin(tensor: FloatTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + unary_float!(tensor, float, |tensor| B::float_argmin(tensor, dim, out_dtype) => Int) + } + + fn float_expand(tensor: FloatTensor, shape: Shape) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_expand(tensor, shape) => Float) + } + + fn float_unfold( + tensor: FloatTensor, + dim: usize, + size: usize, + step: usize, + ) -> FloatTensor { + unary_float!(tensor, float, |tensor| { + B::float_unfold(tensor, dim, size, step) + } => Float) + } + + fn float_detach(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_detach(tensor) => Float) + } + + fn float_set_require_grad(tensor: FloatTensor, require_grad: bool) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_set_require_grad(tensor, require_grad) => Float) + } + + fn float_is_require_grad(tensor: &FloatTensor) -> bool { + unary_float!(ref tensor, float, |tensor| B::float_is_require_grad(tensor)) + } + + // Default implementation + fn float_zeros(shape: Shape, device: &DispatchDevice, dtype: FloatDType) -> FloatTensor { + creation_op!(Float, device, |device| B::float_zeros(shape, device, dtype)) + } + + fn float_ones(shape: Shape, device: &DispatchDevice, dtype: FloatDType) -> FloatTensor { + creation_op!(Float, device, |device| B::float_ones(shape, device, dtype)) + } + + fn float_full( + shape: Shape, + fill_value: Scalar, + device: &DispatchDevice, + dtype: FloatDType, + ) -> FloatTensor { + creation_op!(Float, device, |device| B::float_full( + shape, fill_value, device, dtype + )) + } + + fn float_repeat_dim(tensor: FloatTensor, dim: usize, times: usize) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_repeat_dim(tensor, dim, times) => Float) + } + + fn float_clamp_min(tensor: FloatTensor, min: Scalar) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_clamp_min(tensor, min) => Float) + } + + fn float_clamp_max(tensor: FloatTensor, max: Scalar) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_clamp_max(tensor, max) => Float) + } + + fn float_clamp(tensor: FloatTensor, min: Scalar, max: Scalar) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_clamp(tensor, min, max) => Float) + } + + fn float_neg(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_neg(tensor) => Float) + } + + fn float_transpose(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_transpose(tensor) => Float) + } + + fn float_not_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_float!((lhs, float), (rhs, float), |lhs, rhs| B::float_not_equal(lhs, rhs, out_dtype) => Bool) + } + + fn float_not_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + unary_float!(lhs, float, |lhs| B::float_not_equal_elem(lhs, rhs, out_dtype) => Bool) + } + + fn float_prod(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_prod(tensor) => Float) + } + + fn float_prod_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_prod_dim(tensor, dim) => Float) + } + + fn float_mean(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_mean(tensor) => Float) + } + + fn float_powi(lhs: FloatTensor, rhs: IntTensor) -> FloatTensor { + binary_float!((lhs, float), (rhs, int), |lhs, rhs| B::float_powi(lhs, rhs) => Float) + } + + fn float_powi_scalar_impl(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + unary_float!(lhs, float, |lhs| B::float_powi_scalar_impl(lhs, rhs) => Float) + } + + fn float_powf_scalar(tensor: FloatTensor, value: Scalar) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_powf_scalar(tensor, value) => Float) + } + + fn float_cat(tensors: Vec>, dim: usize) -> FloatTensor { + vec_op!(tensors, float, |tensors| B::float_cat(tensors, dim) => Float) + } + + fn float_max(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_max(tensor) => Float) + } + + fn float_max_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_max_dim(tensor, dim) => Float) + } + + fn float_max_dim_with_indices( + tensor: FloatTensor, + dim: usize, + indices_dtype: IntDType, + ) -> (FloatTensor, IntTensor) { + multi_op!( + inputs[(tensor, float)], + outputs[(out, Float), (indices, Int)], + B::float_max_dim_with_indices(tensor, dim, indices_dtype) + ) + } + + fn float_min(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_min(tensor) => Float) + } + + fn float_min_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_min_dim(tensor, dim) => Float) + } + + fn float_min_dim_with_indices( + tensor: FloatTensor, + dim: usize, + indices_dtype: IntDType, + ) -> (FloatTensor, IntTensor) { + multi_op!( + inputs[(tensor, float)], + outputs[(out, Float), (indices, Int)], + B::float_min_dim_with_indices(tensor, dim, indices_dtype) + ) + } + + fn float_max_abs(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_max_abs(tensor) => Float) + } + + fn float_max_abs_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_max_abs_dim(tensor, dim) => Float) + } + + fn float_any(tensor: FloatTensor, out_dtype: BoolDType) -> BoolTensor { + unary_float!(tensor, float, |tensor| B::float_any(tensor, out_dtype) => Bool) + } + + fn float_any_dim( + tensor: FloatTensor, + dim: usize, + out_dtype: BoolDType, + ) -> BoolTensor { + unary_float!(tensor, float, |tensor| B::float_any_dim(tensor, dim, out_dtype) => Bool) + } + + fn float_all(tensor: FloatTensor, out_dtype: BoolDType) -> BoolTensor { + unary_float!(tensor, float, |tensor| B::float_all(tensor, out_dtype) => Bool) + } + + fn float_all_dim( + tensor: FloatTensor, + dim: usize, + out_dtype: BoolDType, + ) -> BoolTensor { + unary_float!(tensor, float, |tensor| B::float_all_dim(tensor, dim, out_dtype) => Bool) + } + + fn float_sign(tensor: FloatTensor) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_sign(tensor) => Float) + } + + fn float_sort(tensor: FloatTensor, dim: usize, descending: bool) -> FloatTensor { + unary_float!(tensor, float, |tensor| B::float_sort(tensor, dim, descending) => Float) + } + + fn float_sort_with_indices( + tensor: FloatTensor, + dim: usize, + descending: bool, + indices_dtype: IntDType, + ) -> (FloatTensor, IntTensor) { + multi_op!( + inputs[(tensor, float)], + outputs[(out, Float), (indices, Int)], + B::float_sort_with_indices(tensor, dim, descending, indices_dtype) + ) + } + + fn float_argsort( + tensor: FloatTensor, + dim: usize, + descending: bool, + out_dtype: IntDType, + ) -> IntTensor { + unary_float!(tensor, float, |tensor| B::float_argsort(tensor, dim, descending, out_dtype) => Int) + } + + fn float_grid_sample_2d( + tensor: FloatTensor, + grid: FloatTensor, + options: burn_backend::ops::GridSampleOptions, + ) -> FloatTensor { + binary_float!((tensor, float), (grid, float), |tensor, grid| B::float_grid_sample_2d(tensor, grid, options) => Float) + } + + fn float_is_nan(tensor: FloatTensor, out_dtype: BoolDType) -> BoolTensor { + unary_float!(tensor, float, |tensor| B::float_is_nan(tensor, out_dtype) => Bool) + } + + fn float_is_inf(tensor: FloatTensor, out_dtype: BoolDType) -> BoolTensor { + unary_float!(tensor, float, |tensor| B::float_is_inf(tensor, out_dtype) => Bool) + } + + fn float_hypot(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float!((lhs, float), (rhs, float), |lhs, rhs| B::float_hypot(lhs, rhs) => Float) + } +} diff --git a/crates/burn-dispatch/src/ops/transaction.rs b/crates/burn-dispatch/src/ops/transaction.rs new file mode 100644 index 0000000..8da3a18 --- /dev/null +++ b/crates/burn-dispatch/src/ops/transaction.rs @@ -0,0 +1,26 @@ +use alloc::vec::Vec; +use burn_backend::{ + ExecutionError, + ops::{TransactionOps, TransactionPrimitive, TransactionPrimitiveData}, +}; + +use crate::Dispatch; + +impl TransactionOps for Dispatch { + async fn tr_execute( + transaction: TransactionPrimitive, + ) -> Result { + let first_tensor = transaction + .read_floats + .first() + .or(transaction.read_ints.first()) + .or(transaction.read_bools.first()); + + match first_tensor { + Some(tensor) => { + transaction_op!(transaction, tensor) + } + None => Ok(TransactionPrimitiveData::default()), + } + } +} diff --git a/crates/burn-dispatch/src/remote_server.rs b/crates/burn-dispatch/src/remote_server.rs new file mode 100644 index 0000000..b6e51bc --- /dev/null +++ b/crates/burn-dispatch/src/remote_server.rs @@ -0,0 +1,172 @@ +//! Concrete-backend dispatch for `burn-remote`'s server entry points. +//! +//! Lives in `burn-dispatch` because matching on [`DispatchDevice`] requires the +//! local `wgpu_metal`/`wgpu_vulkan`/`wgpu_webgpu` cfgs set by this crate's +//! `build.rs`, plus visibility of every in-tree `BackendIr` type. The user +//! surface (`Channel` enum, opaque `Device` argument) lives in `burn-tensor`. + +use std::sync::Arc; + +use burn_remote::Endpoint; +use burn_remote::server::{CustomOpRegistry, IrohRemoteProtocol, PeerAuthorizer, RemoteProtocol}; +use burn_remote::telemetry::TelemetryProbe; + +use crate::backends::*; +use crate::{Dispatch, DispatchDevice, DispatchDeviceId}; + +/// Transport used to serve remote clients. Re-exported from `burn-remote` so the whole stack +/// shares one definition. +pub use burn_remote::server::Channel; + +/// Collect every [`Device`] the host exposes for the backend that owns `$variant`, by +/// enumerating the backend (see [`Dispatch::enumerate`]) and unwrapping the matching variant. +/// `$id` is the [`DispatchDeviceId`] to enumerate; the result is `Vec>`, indexed by +/// hardware device index, which is exactly the index a client selects with `Device::remote`. +/// +/// Enumeration is only trustworthy when it finds **more than one** device: several backends (the +/// wgpu family — Vulkan/WebGpu/Wgpu — and the CPU-only ones like Flex) can't enumerate hardware +/// and report either nothing or a single placeholder index that isn't the device you'd actually +/// run on. In that case fall back to hosting a single backend-specific default device +/// (`Device::::default()`, e.g. `WgpuDevice::DefaultDevice`) rather than a possibly-empty or +/// bogus enumerated list. This generalizes what used to be a hardcoded Vulkan special-case. +macro_rules! host_devices { + ($id:expr, $variant:ident) => {{ + let devices = Dispatch::enumerate($id) + .into_iter() + .filter_map(|device| match device { + DispatchDevice::$variant(device) => Some(device), + _ => None, + }) + .collect::>(); + if devices.len() > 1 { + devices + } else { + vec![Default::default()] + } + }}; +} + +/// Run `$body` with the concrete backend that owns `$device`'s variant bound to the type alias +/// `$b` and that backend's host device list bound to `$devices`. +/// +/// This is the single source of truth for the `DispatchDevice` → concrete-`BackendIr` mapping. +/// The sync and async server entry points differ only in whether `$body` awaits the call, so they +/// share this one match instead of duplicating eleven `#[cfg]`-gated arms each. Backends without a +/// `BackendIr` impl (`LibTorch`, `Remote`) panic; `Autodiff` is already stripped by `.inner()`. +macro_rules! with_backend { + ($device:expr, |$b:ident, $devices:ident| $body:expr) => { + match $device.inner() { + #[cfg(feature = "cpu")] + DispatchDevice::Cpu(_) => { + type $b = Cpu; + let $devices = host_devices!(DispatchDeviceId::Cpu, Cpu); + $body + } + #[cfg(feature = "cuda")] + DispatchDevice::Cuda(_) => { + type $b = Cuda; + let $devices = host_devices!(DispatchDeviceId::Cuda, Cuda); + $body + } + #[cfg(feature = "metal")] + DispatchDevice::Metal(_) => { + type $b = Metal; + let $devices = host_devices!(DispatchDeviceId::Metal, Metal); + $body + } + #[cfg(feature = "rocm")] + DispatchDevice::Rocm(_) => { + type $b = Rocm; + let $devices = host_devices!(DispatchDeviceId::Rocm, Rocm); + $body + } + #[cfg(feature = "vulkan")] + DispatchDevice::Vulkan(_) => { + type $b = Vulkan; + let $devices = host_devices!(DispatchDeviceId::Vulkan, Vulkan); + $body + } + #[cfg(feature = "wgpu")] + DispatchDevice::Wgpu(_) => { + type $b = Wgpu; + let $devices = host_devices!(DispatchDeviceId::Wgpu, Wgpu); + $body + } + #[cfg(feature = "webgpu")] + DispatchDevice::WebGpu(_) => { + type $b = WebGpu; + let $devices = host_devices!(DispatchDeviceId::WebGpu, WebGpu); + $body + } + #[cfg(any(feature = "flex", default_backend))] + DispatchDevice::Flex(_) => { + type $b = Flex; + let $devices = host_devices!(DispatchDeviceId::Flex, Flex); + $body + } + #[cfg(feature = "ndarray")] + DispatchDevice::NdArray(_) => { + type $b = NdArray; + let $devices = host_devices!(DispatchDeviceId::NdArray, NdArray); + $body + } + #[cfg(feature = "tch")] + DispatchDevice::LibTorch(_) => { + panic!("LibTorch is not supported as a remote-server backend (no BackendIr impl)") + } + #[cfg(feature = "remote")] + DispatchDevice::Remote(_) => { + panic!("Cannot host a remote server on a remote device") + } + #[cfg(feature = "autodiff")] + DispatchDevice::Autodiff(_) => { + unreachable!("Autodiff stripped by .inner() above") + } + } + }; +} + +/// Start a remote-execution server for `device`'s backend, blocking the current thread. +/// +/// `device` selects the backend; `channel` selects the transport. The server hosts that backend's +/// devices, indexed by hardware device index. Use [`start_async`] for the async counterpart. +#[cfg(not(target_family = "wasm"))] +pub fn start(device: DispatchDevice, channel: Channel) { + with_backend!(device, |B, devices| { + burn_remote::server::RemoteServerBuilder::::new(devices) + .channel(channel) + .start() + }) +} + +/// Async counterpart of [`start`]; runs on the caller's tokio runtime instead of blocking. +#[cfg(not(target_family = "wasm"))] +pub async fn start_async(device: DispatchDevice, channel: Channel) { + with_backend!(device, |B, devices| { + burn_remote::server::RemoteServerBuilder::::new(devices) + .channel(channel) + .start_async() + .await + }) +} + +/// Build a backend-erased protocol handler for `device`'s backend. +/// +/// Resolves the dispatch device to its concrete backend type and returns a [`RemoteProtocol`] +/// the caller registers on its own Iroh router under [`BURN_REMOTE_ALPN`]. +pub fn remote_protocol( + device: DispatchDevice, + endpoint: &Endpoint, + probe: TelemetryProbe, + authorizer: Arc, +) -> RemoteProtocol { + with_backend!(device, |B, devices| { + RemoteProtocol::new(IrohRemoteProtocol::::new( + endpoint.clone(), + devices, + authorizer, + probe, + CustomOpRegistry::default(), + )) + }) +} diff --git a/crates/burn-dispatch/src/tensor.rs b/crates/burn-dispatch/src/tensor.rs new file mode 100644 index 0000000..ed02727 --- /dev/null +++ b/crates/burn-dispatch/src/tensor.rs @@ -0,0 +1,570 @@ +use crate::{DispatchDevice, backends::*}; + +#[cfg(feature = "autodiff")] +use burn_autodiff::checkpoint::strategy::{ + BalancedCheckpointing, CheckpointStrategy, NoCheckpointing, +}; +use burn_backend::{Backend, BackendTypes, DType, Shape, TensorMetadata}; + +use crate::CheckpointingStrategy; +#[cfg(feature = "autodiff")] +use alloc::boxed::Box; +#[cfg(feature = "autodiff")] +use burn_backend::tensor::FloatTensor; + +use alloc::{format, string::String}; + +// TODO: if we reduce the different associated types for float/int/bool/quantized tensor primitives down to a single +// `B::TensorPrimitive` we can simplify this. + +/// Tensor which points to a backend tensor primitive kind. +#[derive(Clone, Debug)] +pub enum BackendTensor { + /// Float tensor handle. + Float(B::FloatTensorPrimitive), + /// Int tensor handle. + Int(B::IntTensorPrimitive), + /// Bool tensor handle. + Bool(B::BoolTensorPrimitive), + /// Quantized tensor handle. + Quantized(B::QuantizedTensorPrimitive), + #[cfg(feature = "autodiff")] + /// Autodiff float tensor handle. + Autodiff(FloatTensor>), +} + +impl BackendTensor { + /// Returns the inner float tensor primitive. + pub fn float(self) -> B::FloatTensorPrimitive { + match self { + BackendTensor::Float(tensor) => tensor, + BackendTensor::Int(_) => panic!("Should be float, got int"), + BackendTensor::Bool(_) => panic!("Should be float, got bool"), + BackendTensor::Quantized(_) => panic!("Should be float, got quantized"), + #[cfg(feature = "autodiff")] + BackendTensor::Autodiff(_) => panic!("Should be float, got autodiff"), + } + } + /// Returns the inner float tensor primitive. + pub fn as_float(&self) -> &B::FloatTensorPrimitive { + match self { + BackendTensor::Float(tensor) => tensor, + BackendTensor::Int(_) => panic!("Should be float, got int"), + BackendTensor::Bool(_) => panic!("Should be float, got bool"), + BackendTensor::Quantized(_) => panic!("Should be float, got quantized"), + #[cfg(feature = "autodiff")] + BackendTensor::Autodiff(_) => panic!("Should be float, got autodiff"), + } + } + + /// Returns the inner int tensor primitive. + pub fn int(self) -> B::IntTensorPrimitive { + match self { + BackendTensor::Int(tensor) => tensor, + BackendTensor::Float(_) => panic!("Should be int, got float"), + BackendTensor::Bool(_) => panic!("Should be int, got bool"), + BackendTensor::Quantized(_) => panic!("Should be int, got quantized"), + #[cfg(feature = "autodiff")] + BackendTensor::Autodiff(_) => panic!("Should be int, got autodiff"), + } + } + + /// Returns the inner bool tensor primitive. + pub fn bool(self) -> B::BoolTensorPrimitive { + match self { + BackendTensor::Bool(tensor) => tensor, + BackendTensor::Float(_) => panic!("Should be bool, got float"), + BackendTensor::Int(_) => panic!("Should be bool, got int"), + BackendTensor::Quantized(_) => panic!("Should be bool, got quantized"), + #[cfg(feature = "autodiff")] + BackendTensor::Autodiff(_) => panic!("Should be bool, got autodiff"), + } + } + + /// Returns the inner quantized tensor primitive. + pub fn quantized(self) -> B::QuantizedTensorPrimitive { + match self { + BackendTensor::Quantized(tensor) => tensor, + _ => unreachable!(), + } + } + + #[cfg(feature = "autodiff")] + /// Returns the inner autodiff tensor primitive. + pub fn autodiff(self) -> FloatTensor> { + match self { + BackendTensor::Autodiff(tensor) => tensor, + // NOTE: this is the panicking code reached in tensor.rs:74:18: + _ => unreachable!(), + } + } + + #[cfg(feature = "autodiff")] + /// Returns the inner autodiff tensor primitive. + pub fn as_autodiff(&self) -> &FloatTensor> { + match self { + BackendTensor::Autodiff(tensor) => tensor, + _ => unreachable!(), + } + } + + #[cfg(feature = "autodiff")] + /// Returns the inner autodiff tensor primitive. + pub fn autodiff_inner(self) -> B::FloatTensorPrimitive { + match self { + BackendTensor::Autodiff(tensor) => tensor.primitive, + _ => unreachable!(), + } + } + + /// Returns the tensor primitive kind name. + pub fn name(&self) -> &'static str { + match self { + BackendTensor::Float(_) => "Float", + BackendTensor::Int(_) => "Int", + BackendTensor::Bool(_) => "Bool", + BackendTensor::Quantized(_) => "Quantized", + #[cfg(feature = "autodiff")] + BackendTensor::Autodiff(_) => "Autodiff", + } + } +} + +impl TensorMetadata for BackendTensor { + type Device = B::Device; + fn device(&self) -> Self::Device { + match self { + BackendTensor::Float(tensor) => tensor.device(), + BackendTensor::Int(tensor) => tensor.device(), + BackendTensor::Bool(tensor) => tensor.device(), + BackendTensor::Quantized(tensor) => tensor.device(), + #[cfg(feature = "autodiff")] + BackendTensor::Autodiff(tensor) => tensor.device(), + } + } + fn dtype(&self) -> DType { + match self { + BackendTensor::Float(tensor) => tensor.dtype(), + BackendTensor::Int(tensor) => tensor.dtype(), + BackendTensor::Bool(tensor) => tensor.dtype(), + BackendTensor::Quantized(tensor) => tensor.dtype(), + #[cfg(feature = "autodiff")] + BackendTensor::Autodiff(tensor) => tensor.dtype(), + } + } + + fn shape(&self) -> Shape { + match self { + BackendTensor::Float(tensor) => tensor.shape(), + BackendTensor::Int(tensor) => tensor.shape(), + BackendTensor::Bool(tensor) => tensor.shape(), + BackendTensor::Quantized(tensor) => tensor.shape(), + #[cfg(feature = "autodiff")] + BackendTensor::Autodiff(tensor) => tensor.shape(), + } + } + + fn can_mut(&self) -> bool { + match self { + BackendTensor::Float(tensor) => tensor.can_mut(), + BackendTensor::Int(tensor) => tensor.can_mut(), + BackendTensor::Bool(tensor) => tensor.can_mut(), + BackendTensor::Quantized(tensor) => tensor.can_mut(), + #[cfg(feature = "autodiff")] + BackendTensor::Autodiff(tensor) => tensor.can_mut(), + } + } +} + +/// A tensor that can dispatch operations to any enabled backend at runtime. +/// +/// When the `autodiff` feature is enabled, tensors may carry a checkpointing +/// strategy used to control gradient computation. This is derived from the +/// device used to create the tensor. +#[derive(Clone, Debug)] +pub struct DispatchTensor { + /// Tensor kind primitive. + pub kind: DispatchTensorKind, + // Technically more of a device property, but device is not a dispatch tensor field. + // Right now this is the easiest way to preserve the checkpointing strategy because primitives are not consolidated. + // Once float/int/bool primitives are consolidated into a single associative type, we could hold that + // property for all autodiff tensors. + /// Holds the autodiff checkpointing strategy. + /// - `None`: tensor is not tracked by autodiff + /// - `Some(strategy)`: tensor is tracked by autodiff, and uses the checkpointing `strategy` + pub checkpointing: Option, +} + +/// Internal representation of a [`DispatchTensor`]. +/// +/// This enum contains the concrete backend tensor for each enabled backend. +/// It is not intended to be used directly; instead, it is manipulated by +/// the dispatch system to route operations to the correct backend. +/// +/// Each variant corresponds to a specific backend implementation. +#[derive(Clone, Debug)] +pub enum DispatchTensorKind { + /// The [CPU backend](Cpu) tensor. + #[cfg(feature = "cpu")] + Cpu(BackendTensor), + + /// The [CUDA backend](Cuda) tensor. + #[cfg(feature = "cuda")] + Cuda(BackendTensor), + + /// The [Metal backend](Metal) tensor. + #[cfg(feature = "metal")] + Metal(BackendTensor), + + /// The [ROCm backend](Rocm) tensor. + #[cfg(feature = "rocm")] + Rocm(BackendTensor), + + /// The [Vulkan backend](Vulkan) tensor. + #[cfg(feature = "vulkan")] + Vulkan(BackendTensor), + + /// The [Wgpu backend](Wgpu) tensor. + #[cfg(feature = "wgpu")] + Wgpu(BackendTensor), + + /// The [WebGPU backend](Wgpu) tensor. + #[cfg(feature = "webgpu")] + WebGpu(BackendTensor), + + /// The [Flex backend](Flex) tensor. + #[cfg(any(feature = "flex", default_backend))] + Flex(BackendTensor), + + /// The [NdArray backend](NdArray) tensor. + #[cfg(feature = "ndarray")] + NdArray(BackendTensor), + + /// The [LibTorch backend](LibTorch) tensor. + #[cfg(feature = "tch")] + LibTorch(BackendTensor), + + /// The [Remote backend](Remote) tensor (lives on a remote server). + #[cfg(feature = "remote")] + Remote(BackendTensor), + + /// The [autodiff enabled backend](Autodiff) tensor. + #[cfg(feature = "autodiff")] + Autodiff(Box), +} + +impl TensorMetadata for DispatchTensorKind { + type Device = DispatchDevice; + + fn dtype(&self) -> DType { + match self { + #[cfg(feature = "cpu")] + Self::Cpu(tensor) => tensor.dtype(), + #[cfg(feature = "cuda")] + Self::Cuda(tensor) => tensor.dtype(), + #[cfg(feature = "metal")] + Self::Metal(tensor) => tensor.dtype(), + #[cfg(feature = "rocm")] + Self::Rocm(tensor) => tensor.dtype(), + #[cfg(feature = "vulkan")] + Self::Vulkan(tensor) => tensor.dtype(), + #[cfg(feature = "wgpu")] + Self::Wgpu(tensor) => tensor.dtype(), + #[cfg(feature = "webgpu")] + Self::WebGpu(tensor) => tensor.dtype(), + #[cfg(any(feature = "flex", default_backend))] + Self::Flex(tensor) => tensor.dtype(), + #[cfg(feature = "ndarray")] + Self::NdArray(tensor) => tensor.dtype(), + #[cfg(feature = "tch")] + Self::LibTorch(tensor) => tensor.dtype(), + #[cfg(feature = "remote")] + Self::Remote(tensor) => tensor.dtype(), + #[cfg(feature = "autodiff")] + Self::Autodiff(tensor) => tensor.dtype(), + } + } + + fn shape(&self) -> Shape { + match self { + #[cfg(feature = "cpu")] + Self::Cpu(tensor) => tensor.shape(), + #[cfg(feature = "cuda")] + Self::Cuda(tensor) => tensor.shape(), + #[cfg(feature = "metal")] + Self::Metal(tensor) => tensor.shape(), + #[cfg(feature = "rocm")] + Self::Rocm(tensor) => tensor.shape(), + #[cfg(feature = "vulkan")] + Self::Vulkan(tensor) => tensor.shape(), + #[cfg(feature = "wgpu")] + Self::Wgpu(tensor) => tensor.shape(), + #[cfg(feature = "webgpu")] + Self::WebGpu(tensor) => tensor.shape(), + #[cfg(any(feature = "flex", default_backend))] + Self::Flex(tensor) => tensor.shape(), + #[cfg(feature = "ndarray")] + Self::NdArray(tensor) => tensor.shape(), + #[cfg(feature = "tch")] + Self::LibTorch(tensor) => tensor.shape(), + #[cfg(feature = "remote")] + Self::Remote(tensor) => tensor.shape(), + #[cfg(feature = "autodiff")] + Self::Autodiff(tensor) => tensor.shape(), + } + } + + fn device(&self) -> DispatchDevice { + match self { + #[cfg(feature = "cpu")] + DispatchTensorKind::Cpu(tensor) => DispatchDevice::Cpu(tensor.device()), + #[cfg(feature = "cuda")] + DispatchTensorKind::Cuda(tensor) => DispatchDevice::Cuda(tensor.device()), + #[cfg(feature = "metal")] + DispatchTensorKind::Metal(tensor) => DispatchDevice::Metal(tensor.device()), + #[cfg(feature = "rocm")] + DispatchTensorKind::Rocm(tensor) => DispatchDevice::Rocm(tensor.device()), + #[cfg(feature = "vulkan")] + DispatchTensorKind::Vulkan(tensor) => DispatchDevice::Vulkan(tensor.device()), + #[cfg(feature = "wgpu")] + DispatchTensorKind::Wgpu(tensor) => DispatchDevice::Wgpu(tensor.device()), + #[cfg(feature = "webgpu")] + DispatchTensorKind::WebGpu(tensor) => DispatchDevice::WebGpu(tensor.device()), + #[cfg(any(feature = "flex", default_backend))] + DispatchTensorKind::Flex(tensor) => DispatchDevice::Flex(tensor.device()), + #[cfg(feature = "ndarray")] + DispatchTensorKind::NdArray(tensor) => DispatchDevice::NdArray(tensor.device()), + #[cfg(feature = "tch")] + DispatchTensorKind::LibTorch(tensor) => DispatchDevice::LibTorch(tensor.device()), + #[cfg(feature = "remote")] + DispatchTensorKind::Remote(tensor) => DispatchDevice::Remote(tensor.device()), + #[cfg(feature = "autodiff")] + DispatchTensorKind::Autodiff(tensor) => DispatchDevice::autodiff(tensor.device()), + } + } + + fn can_mut(&self) -> bool { + match self { + #[cfg(feature = "cpu")] + Self::Cpu(tensor) => tensor.can_mut(), + #[cfg(feature = "cuda")] + Self::Cuda(tensor) => tensor.can_mut(), + #[cfg(feature = "metal")] + Self::Metal(tensor) => tensor.can_mut(), + #[cfg(feature = "rocm")] + Self::Rocm(tensor) => tensor.can_mut(), + #[cfg(feature = "vulkan")] + Self::Vulkan(tensor) => tensor.can_mut(), + #[cfg(feature = "wgpu")] + Self::Wgpu(tensor) => tensor.can_mut(), + #[cfg(feature = "webgpu")] + Self::WebGpu(tensor) => tensor.can_mut(), + #[cfg(any(feature = "flex", default_backend))] + Self::Flex(tensor) => tensor.can_mut(), + #[cfg(feature = "ndarray")] + Self::NdArray(tensor) => tensor.can_mut(), + #[cfg(feature = "tch")] + Self::LibTorch(tensor) => tensor.can_mut(), + #[cfg(feature = "remote")] + Self::Remote(tensor) => tensor.can_mut(), + #[cfg(feature = "autodiff")] + Self::Autodiff(tensor) => tensor.can_mut(), + } + } +} + +impl TensorMetadata for DispatchTensor { + fn dtype(&self) -> DType { + self.kind.dtype() + } + + fn shape(&self) -> Shape { + self.kind.shape() + } + + fn can_mut(&self) -> bool { + self.kind.can_mut() + } + + type Device = DispatchDevice; + + fn device(&self) -> Self::Device { + #[allow(unused_mut)] + let mut device = self.kind.device(); + + // TODO: should int and bool kinds return an autodiff device? + // It would be much easier once there is a single underlying primitive type, which + // we can wrap with Autodiff in all cases. + + #[cfg(feature = "autodiff")] + if let DispatchDevice::Autodiff(device) = &mut device + && let Some(checkpointing) = &self.checkpointing + { + device.checkpointing = *checkpointing; + } + + device + } +} + +impl DispatchTensorKind { + /// Returns the backend tensor kind name. + pub(crate) fn name(&self) -> &'static str { + match self { + #[cfg(feature = "cpu")] + DispatchTensorKind::Cpu(_) => "Cpu", + #[cfg(feature = "cuda")] + DispatchTensorKind::Cuda(_) => "Cuda", + #[cfg(feature = "metal")] + DispatchTensorKind::Metal(_) => "Metal", + #[cfg(feature = "rocm")] + DispatchTensorKind::Rocm(_) => "Rocm", + #[cfg(feature = "vulkan")] + DispatchTensorKind::Vulkan(_) => "Vulkan", + #[cfg(feature = "wgpu")] + DispatchTensorKind::Wgpu(_) => "Wgpu", + #[cfg(feature = "webgpu")] + DispatchTensorKind::WebGpu(_) => "WebGpu", + #[cfg(any(feature = "flex", default_backend))] + DispatchTensorKind::Flex(_) => "Flex", + #[cfg(feature = "ndarray")] + DispatchTensorKind::NdArray(_) => "NdArray", + #[cfg(feature = "tch")] + DispatchTensorKind::LibTorch(_) => "LibTorch", + #[cfg(feature = "remote")] + DispatchTensorKind::Remote(_) => "Remote", + #[cfg(feature = "autodiff")] + DispatchTensorKind::Autodiff(_) => "Autodiff", + } + } +} + +#[cfg(feature = "autodiff")] +trait IntoCheckpointingStrategy { + const STRATEGY: CheckpointingStrategy; +} + +#[cfg(feature = "autodiff")] +impl IntoCheckpointingStrategy for NoCheckpointing { + const STRATEGY: CheckpointingStrategy = CheckpointingStrategy::None; +} + +#[cfg(feature = "autodiff")] +impl IntoCheckpointingStrategy for BalancedCheckpointing { + const STRATEGY: CheckpointingStrategy = CheckpointingStrategy::Balanced; +} + +/// Trait to execute runtime routing conversions between the dynamic dispatch layer and specific backends. +pub trait DispatchKindConversion { + /// Attempts to extract a backend-specific [`BackendTensor`] wrapper from a generic, dynamically-routed [`DispatchTensor`]. + /// + /// # Errors + /// + /// Returns an error if the dynamic routing state does not match the requested backend `B`. + fn try_into_backend(tensor: DispatchTensor) -> Result, String>; + + /// Encapsulates a backend-specific tensor variant back into a globally routing [`DispatchTensor`]. + fn from_backend(tensor: BackendTensor) -> DispatchTensor; +} + +macro_rules! impl_dispatch_conversion { + ($backend:ident, $cfg:meta) => { + #[cfg($cfg)] + impl DispatchKindConversion<$backend> for DispatchTensor { + fn try_into_backend(tensor: DispatchTensor) -> Result, String> { + // The catch-all is unreachable in single-backend builds (the enum then has one + // variant), but required when several backend features are enabled. + #[allow(unreachable_patterns)] + match tensor.kind { + DispatchTensorKind::$backend(t) => Ok(t), + other => Err(format!( + "Expected {} tensor, got variant: {}", + stringify!($backend), + other.name() + )), + } + } + + fn from_backend(tensor: BackendTensor<$backend>) -> DispatchTensor { + DispatchTensor { + kind: DispatchTensorKind::$backend(tensor), + checkpointing: None, + } + } + } + + #[cfg(all($cfg, feature = "autodiff"))] + impl + DispatchKindConversion> for DispatchTensor + { + fn try_into_backend( + tensor: DispatchTensor, + ) -> Result>, String> { + match tensor.kind { + DispatchTensorKind::Autodiff(t) => match *t { + DispatchTensorKind::$backend(t) => match t { + // Encode as `BackendTensor::Float` for `Autodiff` + BackendTensor::Autodiff(t) => Ok(BackendTensor::Float(t)), + other => Err(format!( + "Expected Autodiff {} float tensor, got Autodiff variant: {}", + stringify!($backend), + other.name() + )), + }, + other => Err(format!( + "Expected Autodiff {} tensor, got Autodiff variant: {}", + stringify!($backend), + other.name() + )), + }, + other => Err(format!( + "Expected Autodiff tensor, got backend: {}", + other.name() + )), + } + } + + fn from_backend(tensor: BackendTensor>) -> DispatchTensor { + // Unwrap the Autodiff backend representation back into the inner hardware representation + let kind = match tensor { + // Inverse: Wrap the `Float` variant back into the backend's `Autodiff` primitive variant + BackendTensor::Float(t) => { + let ad_tensor = BackendTensor::Autodiff(t); + // Wrap in the concrete backend's dispatch container + let inner_dispatch = DispatchTensorKind::$backend(ad_tensor); + // Re-apply the outer Autodiff dispatch wrapper + DispatchTensorKind::Autodiff(Box::new(inner_dispatch)) + } + + // Pass-throughs for non-differentiable types + BackendTensor::Int(t) => DispatchTensorKind::$backend(BackendTensor::Int(t)), + BackendTensor::Bool(t) => DispatchTensorKind::$backend(BackendTensor::Bool(t)), + BackendTensor::Quantized(t) => { + DispatchTensorKind::$backend(BackendTensor::Quantized(t)) + } + + BackendTensor::Autodiff(_) => { + panic!("Unexpected Autodiff variant provided to `from_backend`",) + } + }; + + DispatchTensor { + kind, + checkpointing: Some(C::STRATEGY), + } + } + } + }; +} + +impl_dispatch_conversion!(Flex, any(feature = "flex", default_backend)); +impl_dispatch_conversion!(Cpu, feature = "cpu"); +impl_dispatch_conversion!(Cuda, feature = "cuda"); +impl_dispatch_conversion!(Rocm, feature = "rocm"); +impl_dispatch_conversion!(Remote, feature = "remote"); +impl_dispatch_conversion!(Metal, feature = "metal"); +impl_dispatch_conversion!(Vulkan, feature = "vulkan"); +impl_dispatch_conversion!(Wgpu, feature = "wgpu"); +impl_dispatch_conversion!(WebGpu, feature = "webgpu"); +impl_dispatch_conversion!(NdArray, feature = "ndarray"); +impl_dispatch_conversion!(LibTorch, feature = "tch"); diff --git a/crates/burn-flex/ACKNOWLEDGMENTS.md b/crates/burn-flex/ACKNOWLEDGMENTS.md new file mode 100644 index 0000000..1d0915d --- /dev/null +++ b/crates/burn-flex/ACKNOWLEDGMENTS.md @@ -0,0 +1,30 @@ +# Acknowledgments + +burn-flex draws on ideas and techniques from several open-source projects. + +## ndarray + +[ndarray](https://github.com/rust-ndarray/ndarray) - N-dimensional array library for Rust. + +- **8-fold unrolled reduction loop**: Our `sum_f32` implementation in `simd/kernels.rs` uses + ndarray's `unrolled_fold` pattern, where eight independent accumulators allow LLVM to emit optimal + SIMD code without explicit intrinsic dispatch. + +## Candle + +[Candle](https://github.com/huggingface/candle) - Minimalist ML framework by Hugging Face. + +- **Tiled im2col for convolutions**: Our convolution implementation uses a tiled im2col approach + (TILE_SIZE=512) inspired by Candle, processing output in fixed-size tiles for better L2 cache + utilization and enabling tile-level parallelism. + +## gemm / macerator + +[gemm](https://github.com/sarah-ek/gemm) and [macerator](https://github.com/wingertge/macerator), +part of the [faer](https://github.com/sarah-ek/faer-rs) ecosystem. + +- **gemm**: Powers all matrix multiplication and convolution GEMM calls (matmul, conv im2col, + deformable conv). Provides strided memory access so we can multiply transposed tensors without + copying, and native f16 support. +- **macerator**: Provides portable SIMD dispatch for scatter-add reductions across NEON, AVX2, and + WASM SIMD128 with a scalar fallback. diff --git a/crates/burn-flex/ARCHITECTURE.md b/crates/burn-flex/ARCHITECTURE.md new file mode 100644 index 0000000..85434d3 --- /dev/null +++ b/crates/burn-flex/ARCHITECTURE.md @@ -0,0 +1,788 @@ +# burn-flex Architecture + +A pure-Rust CPU backend for [Burn](https://github.com/tracel-ai/burn). + +## Goals + +From README: + +- Fast, memory-efficient CPU backend +- Multi-threading, SIMD, optimized matrix multiplication +- Runs on std, no_std, and WebAssembly +- Supports f16/bf16 +- Zero-copy data loading +- Thread-safe by design (Arc-based COW) + +## Robustness + +burn-flex is tested for edge-case robustness to ensure safe behavior on embedded devices and in +production. This includes: + +- **Integer overflow safety**: `wrapping_abs`, `wrapping_neg`, `wrapping_shl/shr` for signed + integers at type boundaries (e.g. `i64::MIN`), matching PyTorch two's complement semantics +- **Rounding correctness**: Uses `num_traits::Float::round` with a ties-to-even correction, + correct for the full float range (values beyond integer precision have no fractional bits) +- **Input validation**: Hard assertions for invalid pooling parameters (zero kernel/stride) and + zero-sized reduce dimensions, preventing undefined behavior on malformed inputs +- **Negative index detection**: Debug assertions on gather/scatter index conversions +- **Index dtype correctness**: Index-producing ops (argmax, argmin, argsort, argwhere, + sort_with_indices) must respect `out_dtype`/`indices_dtype` parameters. Internally use + `isize` + `INDEX_DTYPE` for platform portability, then cast to the requested dtype via + `int_cast` if needed. Never hardcode `i64` for index outputs as it breaks on 32-bit targets. + +## Target Platform + +**Primary: Apple Silicon M3 (ARM64 + NEON)** + +- 128-bit SIMD registers (4x f32, 8x f16) +- Unified memory architecture +- Native f16 support in hardware + +**Secondary: x86_64 with AVX2/AVX-512** (via conditional compilation) + +--- + +## Design Principles + +1. **Leverage Burn** - Use `burn-backend` types and `burn-std` utilities wherever possible +2. **Portability first** - No platform-specific dependencies; std, no_std, WASM +3. **Zero C dependencies** - Pure Rust only (gemm crate for matrix multiplication) +4. **Simple and direct** - Eager execution, no lazy graphs, no fusion (use `burn-fusion` if needed) +5. **Memory reuse** - Minimize allocations through in-place ops and buffer reuse + +--- + +## Feature Flags + +```toml +default = ["std", "simd", "rayon"] +``` + +| Feature | Default | Description | +| ----------- | ------- | --------------------------------------------------------------------------- | +| `std` | Yes | Standard library support | +| `simd` | Yes | Portable SIMD via macerator (enables `macerator`, `aligned-vec`) | +| `rayon` | Yes | Parallel execution for large tensors (forwards `gemm/rayon`) | +| `x86-v4` | No | AVX-512 kernels in gemm for x86_64 (Sapphire Rapids, Zen 4/5, etc.) | +| `apple-amx` | No | Apple Silicon AMX matrix coprocessor in gemm (experimental upstream) | + +The `simd` feature also forwards `gemm/wasm-simd128-enable`, a no-op outside WASM. + +`gemm` is an always-on required dependency (not behind a feature flag). + +### Performance impact on Apple M3 Max (median speedup vs serial baseline) + +Measured via `cargo bench -p burn-flex --bench {matmul,attention,conv_ops}` with features +`std,simd` (serial), `std,simd,rayon` (default), and `std,simd,rayon,apple-amx`. + +| Workload | rayon vs serial | +apple-amx vs rayon | combined | +| --------------------------------------- | --------------- | ------------------- | -------- | +| matmul 1024×1024 f32 | 7.0x | 1.7x | **12.2x** | +| matmul 512×512 f32 | 3.8x | 1.5x | 5.8x | +| attention self b1·h32·s256·d128 | 1.0x | 2.0x | 2.0x | +| attention self b1·h12·s512·d64 | 1.0x | 1.6x | 1.6x | +| conv2d first_layer 4×3×224×224 k7×7 s2 | 9.8x | 1.2x | **11.6x** | +| conv2d large 16×128×64×64 k3×3 | 7.7x | 1.5x | 11.1x | +| conv2d k7×7 | 6.5x | 1.4x | 9.2x | + +Notes: +- Attention ops currently see no rayon uplift; the per-head matmul pipeline does not + propagate `Parallelism::Rayon` to gemm. AMX still delivers a standalone speedup. +- Small shapes (e.g. `batch8_64x64` matmul, `depthwise_k3_8x32x512` conv1d) can regress + under rayon due to thread-spawn overhead; a size-based gating in the matmul/conv + paths would recover those without losing the large-shape wins. +- AMX regresses on transposed operands (`both/rhs_transposed_256x256` matmul drop to + ~0.55x vs rayon). Avoid `apple-amx` for workloads dominated by transposed GEMM. + +--- + +## Memory Strategy + +Minimize allocations wherever possible: + +### In-Place Operations + +When tensor is contiguous at offset 0, mutate in place: + +```rust +fn neg_inplace(mut tensor: FlexTensor) -> FlexTensor { + if let Some((0, end)) = tensor.layout().contiguous_offsets() { + let slice: &mut [f32] = tensor.storage_mut(); + for x in slice[..end].iter_mut() { + *x = -*x; + } + tensor + } else { + // Allocate new buffer for non-contiguous + neg_copy(&tensor) + } +} +``` + +### Output Buffer Reuse + +For binary ops, reuse lhs buffer when contiguous at offset 0: + +```rust +fn add(mut lhs: FlexTensor, rhs: &FlexTensor) -> FlexTensor { + if let Some((0, l_end)) = lhs.layout().contiguous_offsets() { + if let Some((r_start, r_end)) = rhs.layout().contiguous_offsets() { + let lhs_storage: &mut [f32] = lhs.storage_mut(); + let rhs_storage: &[f32] = rhs.storage(); + for (l, &r) in lhs_storage[..l_end].iter_mut().zip(&rhs_storage[r_start..r_end]) { + *l = *l + r; + } + return lhs; + } + } + add_alloc(&lhs, rhs) +} +``` + +### When to Allocate + +Only allocate when necessary: + +- Shape changes (broadcast, concat, reshape of non-contiguous) +- Non-contiguous input that must become contiguous +- Views/slices with non-zero offset + +### Arc-based Copy-on-Write + +Tensor storage is wrapped in `Arc` for O(1) cloning and thread-safe COW: + +```rust +pub struct FlexTensor { + data: Arc, // O(1) clone via refcount increment + layout: Layout, + dtype: DType, +} + +impl FlexTensor { + /// Check if this tensor uniquely owns its data + pub fn is_unique(&self) -> bool { + Arc::strong_count(&self.data) == 1 + } + + /// Get mutable access, cloning data if shared (COW) + pub fn make_data_mut(&mut self) -> &mut Bytes { + Arc::make_mut(&mut self.data) + } +} +``` + +Benefits: + +- **O(1) cloning**: `Arc::clone` is just a refcount increment +- **Thread-safe sharing**: `Arc` is `Send + Sync` +- **COW semantics**: `Arc::make_mut` clones only when shared +- **Smart in-place ops**: `is_unique()` enables mutation without allocation + +This enables the optimization pattern used throughout: + +```rust +fn add_inplace(mut lhs: FlexTensor, rhs: &FlexTensor) -> FlexTensor { + if lhs.is_unique() && lhs.is_contiguous_at_offset_zero() { + // Mutate in place - no allocation needed + let storage = lhs.make_data_mut(); + // ... perform addition ... + lhs + } else { + // Allocate new buffer + add_alloc(&lhs, rhs) + } +} +``` + +Performance impact (vs previous non-Arc implementation): + +- Binary ops: **2.6-4.2x faster** than NdArray (was 1.4-1.8x) +- Scalar ops: **2.6x faster** (was 1.8x) +- Memory: 3x less allocation for binary ops (4.2 MB vs 12.6 MB for 1M elements) + +--- + +## Burn Infrastructure We Use + +From `burn-backend`: + +- `Shape` - tensor dimensions +- `TensorData` - serialized tensor format +- `DType` - runtime dtype enum +- `Element` trait - compile-time element types +- `Backend` trait - the interface we implement +- `*TensorOps` traits - operation interfaces + +From `burn-std`: + +- `Bytes` - aligned byte storage with COW semantics (our tensor backing store) +- `is_contiguous()` - stride validation +- Platform abstractions for no_std + +--- + +## Core Types + +### Layout + +Metadata for interpreting storage as an N-dimensional tensor: + +```rust +use burn_backend::Shape; + +pub struct Layout { + shape: Shape, + strides: Vec, // Signed strides for zero-copy flip + start_offset: usize, +} +``` + +**Signed Strides** + +Strides are `isize` (signed) to enable zero-copy flip operations. A negative stride means we iterate +backward through that dimension: + +```rust +// Original tensor [1, 2, 3, 4] with shape [4], stride [1], offset 0 +// Flipped tensor uses: +// - offset: 3 (point to last element) +// - stride: -1 (move backward) +// Iteration: indices 3, 2, 1, 0 -> values 4, 3, 2, 1 +``` + +Many operations are zero-copy (metadata changes only): + +- `transpose()` - swap strides +- `narrow()` - adjust offset +- `reshape()` - recompute strides if contiguous +- `broadcast()` - set stride to 0 +- `flip()` - negate stride, adjust offset +- `permute()` - reorder strides + +**Zero-Copy Flip** + +With signed strides, `flip(tensor, axes)` is O(1): + +```rust +pub fn flip(&self, axes: &[usize]) -> Self { + let mut new_strides = self.strides.clone(); + let mut offset_adjustment: isize = 0; + + for &axis in axes { + let dim_size = self.shape.dims[axis]; + if dim_size > 1 { + // Move start to the last element in this dimension + offset_adjustment += (dim_size as isize - 1) * self.strides[axis]; + // Negate stride to iterate backward + new_strides[axis] = -new_strides[axis]; + } + } + + let new_start = (self.start_offset as isize + offset_adjustment) as usize; + Self { shape: self.shape.clone(), strides: new_strides, start_offset: new_start } +} +``` + +This avoids the O(n) element-by-element copy that would be required with unsigned strides. + +### Tensor + +Uses `Arc` for O(1) cloning with COW semantics: + +```rust +use std::sync::Arc; +use burn_std::Bytes; +use burn_backend::DType; + +pub struct FlexTensor { + data: Arc, // O(1) clone, COW via Arc::make_mut + layout: Layout, + dtype: DType, +} + +impl FlexTensor { + /// Zero-copy typed view of full storage (for use with StridedIter) + pub fn storage(&self) -> &[E] { + bytemuck::cast_slice(&self.data) + } + + /// Mutable typed view for in-place operations + pub fn storage_mut(&mut self) -> &mut [E] { + bytemuck::cast_slice_mut(&mut self.data) + } +} +``` + +Operations dispatch on `dtype` and cast once at the boundary: + +```rust +fn add(a: &FlexTensor, b: &FlexTensor) -> FlexTensor { + match a.dtype { + DType::F32 => add_impl(a.as_slice::(), b.as_slice::()), + DType::F16 => add_impl(a.as_slice::(), b.as_slice::()), + // ... + } +} +``` + +--- + +## Backend Implementation + +```rust +use burn_backend::{Backend, DType}; + +#[derive(Clone, Copy, Debug, Default)] +pub struct Flex; + +impl Backend for Flex { + type Device = FlexDevice; + type FloatTensorPrimitive = FlexTensor; + type IntTensorPrimitive = FlexTensor; + type BoolTensorPrimitive = FlexTensor; + type QuantizedTensorPrimitive = FlexQTensor; + + fn name() -> String { "flex".into() } + + fn float_supported_dtypes() -> Vec { + vec![DType::F64, DType::F32, DType::F16, DType::BF16] + } + + fn int_supported_dtypes() -> Vec { + vec![DType::I64, DType::I32, DType::I16, DType::I8, + DType::U64, DType::U32, DType::U16, DType::U8] + } +} +``` + +--- + +## FusionBackend + +burn-flex does not implement `FusionBackend`. Without JIT compilation, fusion adds tracking overhead +with no performance benefit. Deferred operations would still execute one-by-one with intermediate +allocations. For CPU with fusion, use `burn-cpu` (which has cubecl's MLIR-based JIT runtime). + +--- + +## Execution Strategy + +### Contiguous Fast Path + +Most tensors are contiguous. Detect and use direct slice operations: + +```rust +fn unary_op(storage: &[T], layout: &Layout, f: F) -> Vec +where + T: Copy, + F: Fn(T) -> T, +{ + if let Some((start, end)) = layout.contiguous_offsets() { + storage[start..end].iter().map(|&x| f(x)).collect() + } else { + StridedIter::new(layout).map(|i| f(storage[i])).collect() + } +} +``` + +### SIMD Kernels + +Portable SIMD via macerator, with automatic dispatch per architecture (NEON, AVX2, SSE, WASM +SIMD128) and a scalar fallback module for unsupported platforms: + +```rust +use macerator::{Simd, with_simd, vload_unaligned, vstore_unaligned}; + +#[with_simd] +fn my_kernel(src: &[f32], dst: &mut [f32]) { + let lanes = f32::lanes::(); + // load/store vectors, use operator overloading for arithmetic +} + +// Dispatch: detects CPU features at runtime +my_kernel(src, dst); +``` + +The `simd/` module is organized as: + +- `portable.rs`: macerator-based binary, comparison, and boolean ops (auto-dispatches to + NEON/AVX2/SSE/SIMD128/scalar) +- `kernels.rs`: macerator-based reduction kernels (sum, scatter-add) +- `scalar.rs`: fallback for builds without the `simd` feature (bool ops only) +- `aligned.rs`: SIMD-aligned memory allocation + +### Parallel Execution + +Via rayon for large tensors: + +```rust +use rayon::prelude::*; + +fn parallel_unary(src: &[T], f: F) -> Vec +where + T: Copy + Send + Sync, + F: Fn(T) -> T + Send + Sync, +{ + src.par_iter().map(|&x| f(x)).collect() +} +``` + +### Linear Algebra + +gemm crate for matrix multiplication with rayon parallelism: + +```rust +use gemm::{gemm, Parallelism}; + +pub fn matmul_f32(lhs: &[f32], rhs: &[f32], out: &mut [f32], m: usize, n: usize, k: usize) { + let parallelism = if m * n * k >= 192 * 192 * 192 { + Parallelism::Rayon(0) // Use all available threads + } else { + Parallelism::None + }; + + unsafe { + gemm( + m, n, k, + out.as_mut_ptr(), n as isize, 1, + 1.0, // alpha + lhs.as_ptr(), k as isize, 1, + rhs.as_ptr(), n as isize, 1, + 0.0, // beta + parallelism, + ); + } +} +``` + +Performance: 1.3-3.4x faster than NdArray (which uses matrixmultiply crate). + +### Convolutions (im2col + gemm) + +All convolutions use the im2col transformation followed by matrix multiplication. This approach: + +- Converts convolution to a well-optimized GEMM operation +- Leverages the same gemm crate used for matmul +- Supports arbitrary strides, padding, dilation, and groups + +**Unified 3D Implementation** + +Rather than three separate implementations, conv1d and conv2d delegate to conv3d: + +``` +conv1d([B, C, W], kernel=[K_out, C_in, W_k]) + → expand dims → conv3d([B, C, 1, 1, W], kernel=[K_out, C_in, 1, 1, W_k]) + → squeeze → [B, K_out, W_out] + +conv2d([B, C, H, W], kernel=[K_out, C_in, H_k, W_k]) + → expand dims → conv3d([B, C, 1, H, W], kernel=[K_out, C_in, 1, H_k, W_k]) + → squeeze → [B, K_out, H_out, W_out] +``` + +Size-1 dimensions have negligible overhead since the gemm operation dominates runtime. + +**im2col Transformation** + +Rearranges input patches into columns for matrix multiplication: + +``` +Input: [B, C_in, D, H, W] +Kernel: [C_out, C_in/groups, K_d, K_h, K_w] + +im2col produces: [spatial_out, C_in/groups * K_d * K_h * K_w] + where spatial_out = D_out * H_out * W_out + +GEMM: W[C_out/groups, col_len] × col[col_len, spatial_out] + → output[C_out/groups, spatial_out] +``` + +**Dtype Support** + +| Dtype | Implementation | +| ----- | ------------------------------------- | +| f32 | Native gemm | +| f64 | Native gemm | +| f16 | Native gemm (since gemm v0.15) | +| bf16 | Convert to f32, compute, convert back | + +bf16 requires conversion because gemm doesn't have native bf16 support. + +**Current Optimizations** + +- **Rayon parallelism**: Batches and groups are parallelized via rayon +- **Tiled im2col**: Column buffer is tiled for better cache locality + +**Remaining Optimization Opportunities** + +1. **Direct convolution**: For small kernels (3x3), direct convolution without im2col can be faster + due to less memory movement + +### Pooling (Unified 3D) + +All pooling operations use the same unified 3D pattern as convolutions: + +``` +pool1d([B, C, W]) + → expand dims → pool3d([B, C, 1, 1, W]) + → squeeze → [B, C, W_out] + +pool2d([B, C, H, W]) + → expand dims → pool3d([B, C, 1, H, W]) + → squeeze → [B, C, H_out, W_out] +``` + +**Supported Operations** + +| Operation | Forward | Backward | +| ----------------- | ------- | ----------------- | +| max_pool | Yes | Yes (via indices) | +| avg_pool | Yes | Yes | +| adaptive_avg_pool | Yes | Yes | + +**Dtype Support** + +| Dtype | Implementation | +| ----- | ------------------------------------- | +| f32 | Native | +| f64 | Native | +| f16 | Native | +| bf16 | Convert to f32, compute, convert back | + +**Parallelization** + +Pooling uses rayon to parallelize over (batch, channel) pairs: + +```rust +(0..batch_size).into_par_iter().for_each(|b| { + (0..channels).into_par_iter().for_each(|c| { + // Process spatial dimensions for this (b, c) slice + }); +}); +``` + +Each (b, c) slice is independent with good cache locality. + +**Max Pool Indices** + +Max pool stores flat indices into input spatial dimensions (as i64): + +- Used by backward pass to route gradients to correct input positions +- Matches Burn's IntElem type for compatibility + +### Conv Transpose (Unified 3D) + +Transposed convolutions (deconvolutions) for upsampling. Uses the same unified 3D pattern: + +``` +conv_transpose1d([B, C_in, W]) + → expand dims → conv_transpose3d([B, C_in, 1, 1, W]) + → squeeze → [B, C_out, W_out] + +conv_transpose2d([B, C_in, H, W]) + → expand dims → conv_transpose3d([B, C_in, 1, H, W]) + → squeeze → [B, C_out, H_out, W_out] +``` + +**Algorithm** + +Unlike regular convolution (which gathers input into output), transposed convolution scatters: + +```rust +for each input position (id, ih, iw): + for each kernel position (kd, kh, kw): + od = id * stride_d + kd * dilation_d - padding_d + oh = ih * stride_h + kh * dilation_h - padding_h + ow = iw * stride_w + kw * dilation_w - padding_w + if (od, oh, ow) in bounds: + output[od, oh, ow] += input[id, ih, iw] * weight[kd, kh, kw] +``` + +**Weight Shape** + +Conv transpose weight shape is opposite of regular conv: + +- Regular conv: `[out_channels, in_channels_per_group, kd, kh, kw]` +- Transpose conv: `[in_channels, out_channels_per_group, kd, kh, kw]` + +**Output Size Formula** + +``` +output_size = (input - 1) * stride + dilation * (kernel - 1) + 1 + padding_out - 2 * padding +``` + +**Parallelization** + +Uses rayon over (batch, output_channel) pairs. For f32, uses atomic adds for thread-safe +accumulation: + +```rust +(0..batch_size * out_channels).into_par_iter().for_each(|k| { + // Scatter input values to output using atomic f32 adds +}); +``` + +**Dtype Support** + +| Dtype | Implementation | +| ----- | -------------------------------------- | +| f32 | Native with atomic adds | +| f64 | Native (sequential per output channel) | +| f16 | Native (sequential) | +| bf16 | Convert to f32, compute, convert back | + +### Attention (Scaled Dot-Product) + +Computes `softmax(Q @ K^T * scale + bias) @ V` with fused scale, softcap, masking (bool + causal), +and additive bias. Auto-selects between two strategies: + +**Naive attention** (seq_q * seq_kv <= 256K): Materializes the full [seq_q, seq_kv] score matrix. Per (batch, +head), issues two gemm calls: one for `Q @ K^T` and one for `softmax(scores) @ V`. The softmax loop +applies scale/softcap/mask/bias and normalizes in two passes (find-max, then exp-and-sum). NaN-safe: +fully-masked rows produce zero output, not NaN. + +**Flash attention** (seq_q * seq_kv > 256K): Tiles over the KV dimension in chunks of TILE_KV (64 on +native, 32 on WASM). Each tile does a small score gemm, online softmax update (running max/sum with +correction factor to rescale previous tiles), and a value accumulation gemm. Memory is +`O(seq_q * TILE_KV)` per head instead of `O(seq_q * seq_kv)`. + +**Why two strategies**: Benchmarks show naive is 5-10% faster for typical transformer shapes +(seq <= 512) because two large gemm calls amortize kernel dispatch overhead better than many small +tiled ones. Flash wins when the score matrix exceeds L2 cache. The threshold is `NAIVE_SCORE_BUDGET` +(256K elements = 1 MB for f32). + +Both paths share: gemm via `gemm::gemm`, dtype dispatch with f16/bf16 upcast to f32, scratch buffer +reuse across (batch, head) pairs. + +### Unfold (Zero-Copy Strided View) + +Unfold extracts sliding windows from a tensor along a dimension. Unlike most backends that copy +data, Flex implements unfold as a **zero-copy strided view**. + +**Output Shape** + +Given input with shape `[pre..., dim_size, post...]`, unfold along dimension `dim` produces: + +- Output shape: `[pre..., windows, post..., window_size]` +- Windows count: `(dim_size - window_size + step) / step` + +**Algorithm** + +Instead of copying window data, Flex manipulates strides: + +```rust +// Build output strides: +// - Dimension `dim` (now windows): input_stride[dim] * step +// - New window_size dimension (appended): input_stride[dim] +// - All other dimensions: same as input + +output_strides[dim] = input_strides[dim] * step; // Windows stride +output_strides.push(input_strides[dim]); // Within-window stride +``` + +This makes unfold O(1) regardless of tensor size, simply returning a view with new shape/strides. + +**Example** + +``` +Input: [1, 2, 3, 4, 5] shape [5], stride [1] +Unfold dim=0, size=3, step=1 + +Output shape: [3, 3] (3 windows of size 3) +Output strides: [1, 1] (window stride = 1*1, within-window stride = 1) + +Logical view: + Window 0: [1, 2, 3] (offsets 0, 1, 2) + Window 1: [2, 3, 4] (offsets 1, 2, 3) + Window 2: [3, 4, 5] (offsets 2, 3, 4) +``` + +**Performance** + +| Metric | Flex | NdArray | +| --------------- | ---------------------------- | ------------------------------ | +| Time complexity | O(1) | O(output_elements) | +| Memory | 56-136 bytes (metadata only) | Megabytes (copies all windows) | +| Speedup | **1,300-156,000x faster** | - | + +**Non-Contiguous Output** + +The returned tensor is non-contiguous (overlapping windows share storage). Operations that require +contiguous data call `to_contiguous()` internally. Many operations (reduce, matmul, conv) work +directly on strided tensors via `StridedIter`. + +### FFT (Real Forward and Inverse) + +**Location**: `ops/fft.rs` + +Forward (rfft) and inverse (irfft) real FFT via Cooley-Tukey with mixed radix-4/radix-2 DIT. + +**Key optimizations:** + +- **Complex packing**: For rfft, pack N real values as N/2 complex, run a half-size complex FFT, + then unpack using Hermitian symmetry. For irfft, reverse the process: repack spectrum, half-size + inverse FFT, de-interleave. This halves the work compared to a full N-point FFT. +- **Compile-time twiddle tables**: `const fn` Taylor-series sin/cos generates static twiddle factor + tables for N=2 through 65536. Zero runtime allocation for common sizes. Stored as split f32 + arrays for direct SIMD loads. +- **Unrolled small kernels**: Hardcoded butterfly networks for N=2, 4, 8 with compile-time twiddle + values (W_4=-i, W_8=sqrt2/2). Eliminates loop overhead for the small inner FFTs produced by + complex packing. +- **Mixed radix-4/radix-2**: Pairs of radix-2 stages are fused into radix-4 passes, halving the + number of data passes for better cache behavior. Odd-stage-count FFTs do one radix-2 pass first. +- **SIMD butterflies**: `#[macerator::with_simd]` vectorizes radix-4 butterfly passes across + consecutive elements within each stage. +- **Inverse via conjugation**: irfft computes IFFT as `(1/N)*conj(FFT(conj(X)))`, reusing the + forward FFT (with its SIMD path) rather than maintaining a separate inverse kernel. +- **Rayon parallelism**: Batched transforms (multiple independent fibers along the FFT dimension) + are distributed across threads. + +**Dtype support**: f32 (native with SIMD radix-4), f64 (rfft computes in f64 with widened f32 +twiddles; irfft truncates to f32 for computation), f16/bf16 (via f32 upcast/downcast). + +--- + +## Optimization Decisions + +### Implemented + +| Optimization | Benefit | Notes | +| ----------------------------- | ----------------------------------- | -------------------------------------------- | +| **Arc-based COW** | O(1) clone, 2.6-4.2x faster ops | `is_unique()` enables true in-place mutation | +| **Portable SIMD (macerator)** | ~1.5-1.7x for contiguous ops | Auto-dispatches to NEON/AVX2/SSE/SIMD128 | +| **Rayon parallelism** | Scales with cores for large tensors | Threshold: 4M elements (memory-bound ops) | +| **Row-based 2D iteration** | 5.9x faster for transposed tensors | Replaces per-element StridedIter | +| **In-place mutation** | Eliminates allocation | When tensor is unique and contiguous | + +### Considered but Skipped + +| Optimization | Why Skipped | +| -------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| **Cache blocking / loop tiling** | Requires architecture-specific tile sizes. M3 has 128KB L1, but optimal tile size varies by operation, data type, and cache hierarchy. Adds complexity without portable benefit. | +| **Software prefetching** | ARM64 `_prefetch` intrinsic is unstable (requires nightly Rust). Apple Silicon has excellent hardware prefetchers that detect strided access patterns automatically. Benefit likely marginal. | +| **Kernel fusion** | Outside burn-flex scope. Fusion is handled at the Burn framework level via `burn-fusion`. This backend focuses on single-operation efficiency. | +| **Hand-tuned intrinsics** | Portable SIMD via macerator covers NEON/AVX2/SSE/SIMD128 with a single implementation. Hand-tuned per-arch intrinsics add maintenance burden with marginal benefit for memory-bound ops. | + +### Why Element-wise Ops are Memory-Bound + +Element-wise operations (add, mul, etc.) perform ~1 FLOP per 4-8 bytes loaded. Modern CPUs can +execute 100+ FLOPs in the time it takes to load one cache line from RAM. This means: + +1. **SIMD helps marginally** - Reduces instruction count but doesn't change memory bandwidth +2. **Avoiding allocation matters more** - In-place mutation eliminates write-allocate traffic +3. **Simple loops auto-vectorize** - Compiler generates good SIMD code for predictable patterns +4. **Hardware prefetchers are effective** - M3 detects sequential and strided patterns automatically + +--- + +## Zero-Copy Loading + +`Bytes` from burn-std supports zero-copy scenarios (mmap, external buffers). `FlexTensor` wraps this +in `Arc` for cheap cloning while preserving zero-copy capabilities. + +## Thread Safety + +`Arc` provides thread-safe sharing with automatic COW: + +- `Arc` is `Send + Sync` for safe cross-thread sharing +- `Arc::make_mut` triggers copy only when data is shared +- `Arc::strong_count` enables `is_unique()` checks for in-place optimization diff --git a/crates/burn-flex/BENCHMARKS.md b/crates/burn-flex/BENCHMARKS.md new file mode 100644 index 0000000..3dec2b7 --- /dev/null +++ b/crates/burn-flex/BENCHMARKS.md @@ -0,0 +1,976 @@ +# Benchmarks: Flex vs NdArray + +All benchmarks run on Apple M3 Max, comparing burn-flex against burn-ndarray. Default features +enabled (`std`, `simd`, `rayon`); `gemm` is a required dependency. + +**Date**: 2026-04-06 + +## How to Read + +- **Median** time reported (lower is better) +- **Speedup** = NdArray median / Flex median +- **Mem** = peak allocation (`max alloc` from divan) +- Bold speedup means Flex wins; plain means tie or NdArray wins + +--- + +## Binary Operations (f32) + +| Operation | Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| ----------- | ---- | ------- | ------- | -------- | -------- | ----------- | +| add | 4K | 389 ns | 436 ns | **1.1x** | 16.4 KB | 16.4 KB | +| add | 64K | 7.36 us | 7.45 us | ~1x | 262 KB | 262 KB | +| add | 1M | 83.9 us | 115 us | **1.4x** | 4.19 MB | 4.19 MB | +| mul | 4K | 382 ns | 430 ns | **1.1x** | 16.4 KB | 16.4 KB | +| mul | 64K | 7.40 us | 7.40 us | ~1x | 262 KB | 262 KB | +| mul | 1M | 115 us | 115 us | ~1x | 4.19 MB | 4.19 MB | +| div | 1M | 115 us | 115 us | ~1x | 4.19 MB | 4.19 MB | +| add_scalar | 1M | 78.7 us | 87.8 us | **1.1x** | 4.19 MB | 4.19 MB | +| mul_scalar | 1M | 75.8 us | 87.5 us | **1.2x** | 4.19 MB | 4.19 MB | +| powf | 64K | 197 us | 199 us | ~1x | 262 KB | 262 KB | +| powf | 1M | 3.17 ms | 3.21 ms | ~1x | 4.19 MB | 4.19 MB | +| powf_scalar | 1M | 3.23 ms | 3.18 ms | ~1x | 4.19 MB | 4.19 MB | +| atan2 | 64K | 143 us | 142 us | ~1x | 262 KB | 262 KB | +| atan2 | 1M | 2.33 ms | 2.32 ms | ~1x | 4.19 MB | 4.19 MB | + +### Transposed + +| Operation | Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| --------- | --------- | ------- | ------- | ------- | -------- | ----------- | +| add | 256x256 | 48.5 us | 46.0 us | 0.95x | 262 KB | 262 KB | +| add | 1024x1024 | 1.00 ms | 990 us | ~1x | 4.19 MB | 4.19 MB | + +--- + +## Binary Operations (i64) + +| Operation | Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| -------------- | ---- | ------- | ------- | -------- | -------- | ----------- | +| int_add | 4K | 361 ns | 655 ns | **1.8x** | 16.5 KB | 32.8 KB | +| int_add | 64K | 7.40 us | 14.6 us | **2.0x** | 262 KB | 524 KB | +| int_add | 1M | 115 us | 230 us | **2.0x** | 4.19 MB | 8.39 MB | +| int_mul | 4K | 366 ns | 1.95 us | **5.3x** | 16.4 KB | 32.8 KB | +| int_mul | 64K | 7.40 us | 26.7 us | **3.6x** | 262 KB | 524 KB | +| int_mul | 1M | 115 us | 230 us | **2.0x** | 4.19 MB | 8.39 MB | +| int_div | 1M | 604 us | 698 us | **1.2x** | 4.19 MB | 8.39 MB | +| int_add_scalar | 1M | 75.8 us | 174 us | **2.3x** | 4.19 MB | 8.39 MB | +| int_mul_scalar | 1M | 75.7 us | 258 us | **3.4x** | 4.19 MB | 8.39 MB | + +### Int Power + +| Operation | Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| --------- | -------- | ------- | ------- | -------- | -------- | ----------- | +| int_powi | 256x256 | 95.6 us | 83.1 us | 0.87x | 262 KB | 524 KB | +| int_powi | 1024x256 | 336 us | 382 us | **1.1x** | 1.05 MB | 2.10 MB | + +### Transposed (i64) + +| Operation | Size | Flex | NdArray | Speedup | +| --------- | --------- | ------- | ------- | -------- | +| int_add | 256x256 | 55.6 us | 50.5 us | 0.91x | +| int_add | 1024x1024 | 996 us | 1.10 ms | **1.1x** | + +--- + +## Int Cast + +| Operation | Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| ---------- | --------- | ------- | ------- | ----------- | -------- | ----------- | +| i64 to i8 | 256x256 | 3.19 us | 20.2 us | **6.3x** | 65.6 KB | 65.6 KB | +| i64 to i32 | 64x64 | 16.8 ns | 1.37 us | **82x** | 16.0 B | 16.4 KB | +| i64 to i32 | 256x256 | 13.6 ns | 20.0 us | **~1475x** | 16.0 B | 262 KB | +| i64 to i32 | 1024x1024 | 13.6 ns | 352 us | **~25963x** | 16.0 B | 4.19 MB | + +--- + +## Int Random + +| Operation | Size | Flex | NdArray | Speedup | +| --------- | ---------- | ------- | ------- | -------- | +| uniform | 64x64 | 20.9 us | 31.6 us | **1.5x** | +| uniform | 256x256 | 334 us | 510 us | **1.5x** | +| uniform | 1024x1024 | 5.36 ms | 8.16 ms | **1.5x** | +| uniform | 16x128x128 | 1.35 ms | 2.04 ms | **1.5x** | + +--- + +## Matrix Multiplication + +### Square (f32) + +| Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| --------- | ------- | ------- | -------- | -------- | ----------- | +| 64x64 | 6.06 us | 18.9 us | **3.1x** | 33.6 KB | 49.3 KB | +| 128x128 | 43.8 us | 41.8 us | ~1x | 328 KB | 197 KB | +| 256x256 | 166 us | 138 us | 0.83x | 524 KB | 786 KB | +| 512x512 | 579 us | 840 us | **1.4x** | 2.10 MB | 3.15 MB | +| 1024x1024 | 2.69 ms | 5.83 ms | **2.2x** | 8.39 MB | 12.6 MB | + +### Rectangular (f32) + +| Shape | Flex | NdArray | Speedup | +| ------------------- | ------ | ------- | ------- | +| 512x64 x 64x512 | 167 us | 144 us | 0.87x | +| 256x512 x 512x256 | 265 us | 266 us | ~1x | +| 128x1024 x 1024x128 | 190 us | 199 us | ~1x | + +### Transposed (256x256) + +| Config | Flex | NdArray | Speedup | +| --------------- | ------ | ------- | -------- | +| LHS transposed | 140 us | 174 us | **1.2x** | +| RHS transposed | 158 us | 173 us | **1.1x** | +| Both transposed | 165 us | 210 us | **1.3x** | + +### Batched (f32) + +| Shape | Flex | NdArray | Speedup | +| ------------------ | ------- | ------- | -------- | +| 8x 64x64 | 57.6 us | 76.7 us | **1.3x** | +| 32x 64x64 | 67.0 us | 111 us | **1.7x** | +| 16x 128x128 | 267 us | 540 us | **2.0x** | +| 12x 512x64 (heads) | 777 us | 1.71 ms | **2.2x** | + +### Broadcast (f32) + +| Shape | Flex | NdArray | Speedup | +| ------------------------- | ------- | ------- | -------- | +| [1,64,64] x [8,64,64] | 47.9 us | 79.8 us | **1.7x** | +| [8,64,64] x [1,64,64] | 54.0 us | 78.5 us | **1.5x** | +| [2,1,32,32] x [1,4,32,32] | 7.00 us | 40.0 us | **5.7x** | +| [4,1,64,64] x [1,4,64,64] | 50.0 us | 66.4 us | **1.3x** | + +### Integer (i32) + +| Size | Flex | NdArray | Speedup | +| ------- | ------- | ------- | -------- | +| 64x64 | 30.2 us | 110 us | **3.7x** | +| 128x128 | 196 us | 971 us | **4.9x** | +| 256x256 | 1.90 ms | 10.1 ms | **5.3x** | +| 512x512 | 18.3 ms | 119 ms | **6.5x** | + +--- + +## Slice Operations + +### Basic Slicing + +| Operation | Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| --------- | --------- | ------ | ------- | --------- | -------- | ----------- | +| slice 1D | 1K | 112 ns | 235 ns | **2.1x** | 56.0 B | 2.15 KB | +| slice 1D | 1M | 105 ns | 26.4 us | **~252x** | 8.75 B | 2.10 MB | +| slice 2D | 256x256 | 120 ns | 3.61 us | **30x** | 19.0 B | 65.7 KB | +| slice 2D | 1024x1024 | 115 ns | 31.9 us | **~278x** | 17.5 B | 1.05 MB | +| slice 3D | 64x64x64 | 148 ns | 16.0 us | **~108x** | 28.5 B | 131 KB | + +### Narrow + +| Operation | Size | Flex | NdArray | Speedup | +| ----------- | --------- | ------ | ------- | --------- | +| narrow dim0 | 256x256 | 141 ns | 1.70 us | **12x** | +| narrow dim0 | 1024x1024 | 142 ns | 26.0 us | **~183x** | +| narrow dim1 | 256x256 | 128 ns | 6.84 us | **53x** | + +### Slice Assignment + +| Operation | Size | Flex | NdArray | Speedup | +| --------- | --------- | ------- | ------- | -------- | +| assign 1D | 1K | 304 ns | 395 ns | **1.3x** | +| assign 2D | 256x256 | 5.39 us | 5.69 us | **1.1x** | +| assign 2D | 1024x1024 | 74.9 us | 74.8 us | ~1x | + +### Transposed Slicing + +| Size | Flex | NdArray | Speedup | +| --------- | ------- | ------- | ---------- | +| 256x256 | 98.2 ns | 7.98 us | **81x** | +| 1024x1024 | 98.2 ns | 232 us | **~2363x** | + +### Slice with Step + +| Operation | Size | Flex | NdArray | Speedup | +| --------- | --------- | ------- | ------- | ---------- | +| step2 1D | 1K | 87.1 ns | 360 ns | **4.1x** | +| step2 1D | 1M | 76.7 ns | 142 us | **~1849x** | +| step2 2D | 1024x1024 | 103 ns | 86.8 us | **~845x** | +| step4 2D | 256x256 | 101 ns | 2.65 us | **26x** | + +--- + +## Concatenation + +### Cat (dim 0, contiguous memcpy fast path) + +| Tensors | Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| ------- | -------- | ------- | ------- | -------- | -------- | ----------- | +| 4x | 256x256 | 16.0 us | 33.1 us | **2.1x** | 1.05 MB | 2.10 MB | +| 4x | 1024x256 | 57.9 us | 132 us | **2.3x** | 4.20 MB | 8.39 MB | +| 16x | 64x64 | 4.11 us | 11.5 us | **2.8x** | 265 KB | 528 KB | +| 4x | 16K (1D) | 3.47 us | 10.1 us | **2.9x** | 263 KB | 525 KB | + +### Cat (dim 1, general path) + +| Tensors | Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| ------- | ------- | ------- | ------- | -------- | -------- | ----------- | +| 4x | 256x64 | 6.92 us | 59.2 us | **8.6x** | 263 KB | 525 KB | +| 4x | 1024x64 | 25.5 us | 366 us | **14x** | 1.05 MB | 2.10 MB | + +Dim-1 cat is much faster because NdArray's default uses N `slice_assign` calls while Flex copies +contiguous chunks directly. + +--- + +## Reduce Operations + +### Full Tensor Sum + +| Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| ---- | ------- | ------- | -------- | -------- | ----------- | +| 1K | 118 ns | 156 ns | **1.3x** | 76.2 B | 44.0 B | +| 64K | 3.23 us | 6.20 us | **1.9x** | 80.0 B | 44.0 B | +| 1M | 43.3 us | 97.0 us | **2.2x** | 84.0 B | 44.0 B | + +### Full Tensor Max + +| Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| ---- | ------- | ------- | -------- | -------- | ----------- | +| 1K | 207 ns | 662 ns | **3.2x** | 76.2 B | 44.0 B | +| 64K | 8.78 us | 32.5 us | **3.7x** | 84.0 B | 44.0 B | +| 1M | 139 us | 558 us | **4.0x** | 84.0 B | 44.0 B | + +### Full Tensor Min + +| Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| ---- | ------- | ------- | -------- | -------- | ----------- | +| 1K | 278 ns | 579 ns | **2.1x** | 84.0 B | 44.0 B | +| 64K | 9.15 us | 34.8 us | **3.8x** | 84.0 B | 44.0 B | +| 1M | 142 us | 540 us | **3.8x** | 84.0 B | 44.0 B | + +### Int Max + +| Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| --------- | ------- | ------- | -------- | -------- | ----------- | +| 256x256 | 2.84 us | 9.12 us | **3.2x** | 84.0 B | 48.0 B | +| 1024x1024 | 42.2 us | 145 us | **3.4x** | 92.0 B | 48.0 B | + +### Sum Along Dimension + +| Shape | Dim | Flex | NdArray | Speedup | +| --------- | --- | ------- | ------- | -------- | +| 256x256 | 0 | 5.06 us | 11.4 us | **2.3x** | +| 256x256 | 1 | 2.77 us | 4.61 us | **1.7x** | +| 1024x1024 | 0 | 80.0 us | 100 us | **1.3x** | +| 1024x1024 | 1 | 42.2 us | 82.0 us | **1.9x** | + +### 3D Sum (Batched) + +| Shape | Dim | Flex | NdArray | Speedup | +| ---------- | --- | ------- | ------- | -------- | +| 32x256x256 | 1 | 156 us | 212 us | **1.4x** | +| 32x256x256 | 2 | 86.2 us | 134 us | **1.6x** | + +### Sum Transposed + +| Size | Flex | NdArray | Speedup | +| --------- | ------- | ------- | -------- | +| 256x256 | 3.27 us | 6.20 us | **1.9x** | +| 1024x1024 | 41.4 us | 96.7 us | **2.3x** | + +### Sum Dim on Transposed + +| Size | Dim | Flex | NdArray | Speedup | +| --------- | --- | ------- | ------- | -------- | +| 256x256 | 0 | 2.79 us | 4.40 us | **1.6x** | +| 1024x1024 | 0 | 41.5 us | 81.9 us | **2.0x** | + +### Mean Along Dimension + +| Shape | Dim | Flex | NdArray | Speedup | +| --------- | --- | ------- | ------- | -------- | +| 256x256 | 1 | 2.90 us | 4.53 us | **1.6x** | +| 1024x1024 | 1 | 42.5 us | 82.4 us | **1.9x** | + +### Argmax + +| Shape | Dim | Flex | NdArray | Speedup | +| --------- | --- | ------- | ------- | -------- | +| 1K | - | 752 ns | 4.09 us | **5.4x** | +| 256x256 | 1 | 66.7 us | 242 us | **3.6x** | +| 1024x1024 | 1 | 120 us | 3.98 ms | **33x** | + +--- + +## Cumulative Operations + +### Cumsum + +| Shape | Dim | Flex | NdArray | Speedup | +| --------- | --- | ------- | ------- | -------- | +| 1K | 0 | 838 ns | 65.7 us | **78x** | +| 64K | 0 | 45.2 us | 4.25 ms | **94x** | +| 1M | 0 | 719 us | 68.1 ms | **95x** | +| 256x256 | 0 | 11.3 us | 34.2 us | **3.0x** | +| 256x256 | 1 | 42.6 us | 215 us | **5.0x** | +| 1024x1024 | 1 | 709 us | 5.51 ms | **7.8x** | + +### Cumprod + +| Shape | Dim | Flex | NdArray | Speedup | +| ------- | --- | ------- | ------- | -------- | +| 1K | 0 | 1.27 us | 66.0 us | **52x** | +| 256x256 | 1 | 66.3 us | 216 us | **3.3x** | + +### Cummin + +| Shape | Dim | Flex | NdArray | Speedup | +| --------- | --- | ------- | ------- | -------- | +| 1K | 0 | 1.79 us | 66.3 us | **37x** | +| 256x256 | 1 | 102 us | 204 us | **2.0x** | +| 1024x1024 | 1 | 1.71 ms | 5.53 ms | **3.2x** | + +### Cummax + +| Shape | Dim | Flex | NdArray | Speedup | +| --------- | --- | ------- | ------- | -------- | +| 1K | 0 | 1.82 us | 65.9 us | **36x** | +| 256x256 | 1 | 102 us | 123 us | **1.2x** | +| 1024x1024 | 1 | 1.70 ms | 3.60 ms | **2.1x** | + +### 3D Cumsum (Batched) + +| Shape | Dim | Flex | NdArray | Speedup | +| -------- | --- | ------- | ------- | -------- | +| 32x64x64 | 1 | 24.1 us | 83.7 us | **3.5x** | +| 32x64x64 | 2 | 68.0 us | 237 us | **3.5x** | + +--- + +## Gather/Scatter Operations + +### Gather + +| Shape | Dim | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| --------- | --- | ------- | ------- | -------- | -------- | ----------- | +| 256x256 | 0 | 32.9 us | 140 us | **4.3x** | 393 KB | 786 KB | +| 256x256 | 1 | 33.8 us | 87.1 us | **2.6x** | 393 KB | 786 KB | +| 1024x1024 | 1 | 273 us | 1.31 ms | **4.8x** | 6.29 MB | 12.6 MB | + +### Scatter Add + +| Shape | Dim | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| --------- | --- | ------- | ------- | -------- | -------- | ----------- | +| 256x256 | 1 | 35.7 us | 189 us | **5.3x** | 524 KB | 918 KB | +| 1024x1024 | 1 | 563 us | 2.83 ms | **5.0x** | 8.39 MB | 14.7 MB | + +### Select + +| Shape | Dim | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| --------- | --- | ------- | ------- | -------- | -------- | ----------- | +| 256x256 | 0 | 2.02 us | 12.9 us | **6.4x** | 132 KB | 143 KB | +| 256x256 | 1 | 26.3 us | 31.1 us | **1.2x** | 132 KB | 143 KB | +| 1024x1024 | 0 | 26.8 us | 88.1 us | **3.3x** | 2.10 MB | 2.15 MB | + +### Bool Select + +| Shape | Indices | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| -------- | ------- | ------- | ------- | ------- | -------- | ----------- | +| 256x256 | 128 | 935 ns | 12.0 us | **13x** | 33.9 KB | 45.0 KB | +| 1024x256 | 512 | 2.89 us | 47.4 us | **16x** | 135 KB | 180 KB | + +### Select Add + +| Shape | Dim | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| --------- | --- | ------- | ------- | -------- | -------- | ----------- | +| 256x256 | 0 | 7.35 us | 13.5 us | **1.8x** | 263 KB | 263 KB | +| 1024x1024 | 0 | 103 us | 126 us | **1.2x** | 4.20 MB | 4.20 MB | + +--- + +## Unary Operations + +### Basic Math + +| Operation | Size | Flex | NdArray | Speedup | +| --------- | ---- | ------- | ------- | -------- | +| exp | 4K | 5.07 us | 5.20 us | ~1x | +| exp | 64K | 80.5 us | 85.0 us | **1.1x** | +| exp | 1M | 1.31 ms | 1.35 ms | ~1x | +| log | 4K | 6.74 us | 6.87 us | ~1x | +| log | 64K | 106 us | 111 us | ~1x | +| log | 1M | 1.72 ms | 1.77 ms | ~1x | +| sqrt | 4K | 612 ns | 860 ns | **1.4x** | +| sqrt | 64K | 9.03 us | 12.8 us | **1.4x** | +| sqrt | 1M | 142 us | 195 us | **1.4x** | +| abs | 1M | 75.8 us | 75.8 us | ~1x | +| recip | 1M | 75.5 us | 75.6 us | ~1x | + +### Trigonometric + +| Operation | Size | Flex | NdArray | Speedup | +| --------- | ---- | ------- | ------- | -------- | +| sin | 4K | 5.65 us | 8.04 us | **1.4x** | +| sin | 64K | 89.2 us | 130 us | **1.5x** | +| sin | 1M | 1.45 ms | 2.10 ms | **1.4x** | +| cos | 4K | 6.57 us | 8.45 us | **1.3x** | +| cos | 1M | 1.68 ms | 2.21 ms | **1.3x** | +| tanh | 4K | 7.07 us | 13.7 us | **1.9x** | +| tanh | 64K | 112 us | 222 us | **2.0x** | +| tanh | 1M | 1.80 ms | 3.57 ms | **2.0x** | + +### Transposed (Non-contiguous) + +| Operation | Size | Flex | NdArray | Speedup | +| --------- | --------- | ------- | ------- | -------- | +| exp | 256x256 | 80.1 us | 84.8 us | **1.1x** | +| exp | 1024x1024 | 1.31 ms | 1.35 ms | ~1x | + +--- + +## Comparison & Boolean Operations + +### Tensor-Tensor Comparisons + +| Operation | Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| --------- | ---- | ------- | ------- | ------- | -------- | ----------- | +| greater | 4K | 431 ns | 398 ns | ~1x | 4.17 KB | 4.14 KB | +| greater | 64K | 6.48 us | 5.73 us | ~1x | 65.6 KB | 65.6 KB | +| greater | 1M | 93 us | 88 us | ~1x | 1.05 MB | 1.05 MB | +| equal | 4K | 433 ns | 403 ns | ~1x | 4.17 KB | 4.14 KB | +| equal | 1M | 86 us | 89 us | ~1x | 1.05 MB | 1.05 MB | +| lower | 1M | 92 us | 87 us | ~1x | 1.05 MB | 1.05 MB | + +### Scalar Comparisons + +| Operation | Size | Flex | NdArray | Speedup | +| ------------ | ---- | ----- | ------- | --------- | +| greater_elem | 1M | 56 us | 76 us | **1.36x** | + +### Transposed Comparisons + +| Operation | Size | Flex | NdArray | Speedup | +| --------- | --------- | ------- | ------- | ------- | +| greater | 256x256 | 53.6 us | 43.7 us | 0.82x | +| greater | 1024x1024 | 985 us | 990 us | ~1x | + +### Broadcast Comparisons + +| Operation | Shape | Flex | NdArray | Speedup | +| --------- | --------- | ------- | ------- | -------- | +| greater | 256x256 | 7.98 us | 25.6 us | **3.2x** | +| greater | 1024x1024 | 120 us | 317 us | **2.6x** | + +### Expand (Broadcasting) + +| Operation | Flex | NdArray | Speedup | +| ------------------- | ------ | ------- | ---------- | +| 1x1 to 1000x1000 | 126 ns | 291 us | **~2307x** | +| 1024x1 to 1024x1024 | 110 ns | 310 us | **~2803x** | +| 1x1024 to 1024x1024 | 126 ns | 78.6 us | **~623x** | + +### Boolean Operations + +| Operation | Size | Flex | NdArray | Speedup | +| --------- | ---- | ------- | ------- | ------- | +| bool_not | 1M | 24.1 us | 19.0 us | 0.79x | +| bool_and | 1M | 34.6 us | 28.8 us | 0.83x | + +--- + +## Convolutions + +### Kernel Size Comparison (4x64x56x56, 64 to 128 channels) + +| Kernel | Flex | NdArray | Speedup | +| ------ | ------- | ------- | -------- | +| 1x1 | 577 us | 803 us | **1.4x** | +| 3x3 | 3.65 ms | 9.46 ms | **2.6x** | +| 5x5 | 8.05 ms | 24.7 ms | **3.1x** | +| 7x7 | 15.7 ms | 49.8 ms | **3.2x** | + +### ResNet Layers (batch=1, 3x3) + +| Layer | Input | Channels | Flex | NdArray | Speedup | +| ------ | ----------- | -------------- | ------- | ------- | -------- | +| conv1 | 1x3x224x224 | 3 to 64 (k7s2) | 954 us | 1.26 ms | **1.3x** | +| layer1 | 1x64x56x56 | 64 to 64 | 986 us | 1.82 ms | **1.8x** | +| layer2 | 1x128x28x28 | 128 to 128 | 1.08 ms | 1.60 ms | **1.5x** | +| layer3 | 1x256x14x14 | 256 to 256 | 1.65 ms | 3.08 ms | **1.9x** | +| layer4 | 1x512x7x7 | 512 to 512 | 2.71 ms | 10.3 ms | **3.8x** | + +### Small (batch=1, 3x3) + +| Input | Channels | Flex | NdArray | Speedup | +| ---------- | -------- | ------- | ------- | -------- | +| 1x3x32x32 | 3 to 16 | 71.3 us | 79.4 us | **1.1x** | +| 1x16x32x32 | 16 to 32 | 219 us | 252 us | **1.2x** | +| 1x32x16x16 | 32 to 64 | 164 us | 338 us | **2.1x** | + +### Large Batched (batch=16, 3x3) + +| Input | Channels | Flex | NdArray | Speedup | +| ------------- | ---------- | ------- | ------- | -------- | +| 16x64x128x128 | 64 to 128 | 79.8 ms | 180 ms | **2.3x** | +| 16x128x64x64 | 128 to 256 | 59.8 ms | 218 ms | **3.7x** | + +### Medium Batched (batch=8, 3x3) + +| Input | Channels | Flex | NdArray | Speedup | +| ---------- | --------- | ------- | ------- | -------- | +| 8x3x64x64 | 3 to 64 | 925 us | 491 us | 0.53x | +| 8x32x64x64 | 32 to 64 | 4.67 ms | 6.44 ms | **1.4x** | +| 8x64x32x32 | 64 to 128 | 3.07 ms | 9.27 ms | **3.0x** | + +### Conv1d + +| Input | Kernel | Flex | NdArray | Speedup | +| ---------- | ------ | ------- | ------- | -------- | +| 1x16x256 | 3 | 31.4 us | 163 us | **5.2x** | +| 8x32x512 | 5 | 536 us | 2.32 ms | **4.3x** | +| 16x64x1024 | 7 | 5.17 ms | 50.7 ms | **9.8x** | + +--- + +## Pooling + +### Max Pool 2D + +| Input | Kernel | Flex | NdArray | Speedup | +| ------------ | ------ | ------- | ------- | -------- | +| 1x64x56x56 | 3x3 s2 | 135 us | 165 us | **1.2x** | +| 8x64x56x56 | 3x3 s2 | 683 us | 914 us | **1.3x** | +| 16x128x28x28 | 2x2 s2 | 406 us | 640 us | **1.6x** | +| 1x512x14x14 | 2x2 s2 | 90.5 us | 106 us | **1.2x** | + +### Max Pool 2D (ResNet) + +| Input | Kernel | Flex | NdArray | Speedup | +| ------------- | ------ | ------- | ------- | -------- | +| 1x64x112x112 | 3x3 s2 | 446 us | 520 us | **1.2x** | +| 8x64x112x112 | 3x3 s2 | 2.63 ms | 3.00 ms | **1.1x** | +| 16x64x112x112 | 3x3 s2 | 5.03 ms | 5.91 ms | **1.2x** | + +### Avg Pool 2D + +| Input | Kernel | Flex | NdArray | Speedup | +| ------------ | ------ | ------ | ------- | -------- | +| 1x64x56x56 | 3x3 s2 | 155 us | 149 us | ~1x | +| 8x64x56x56 | 3x3 s2 | 782 us | 889 us | **1.1x** | +| 16x128x28x28 | 2x2 s2 | 484 us | 480 us | ~1x | + +### Adaptive Avg Pool 2D + +| Input | Output | Flex | NdArray | Speedup | +| ----------- | ------ | ------- | ------- | -------- | +| 1x256x56x56 | 7x7 | 151 us | 142 us | 0.94x | +| 1x512x7x7 | 1x1 | 63.7 us | 68.2 us | **1.1x** | +| 8x512x7x7 | 1x1 | 112 us | 110 us | ~1x | +| 16x2048x7x7 | 1x1 | 286 us | 289 us | ~1x | + +### Max Pool 1D + +| Input | Kernel | Flex | NdArray | Speedup | +| ----------- | ------ | ------- | ------- | -------- | +| 1x64x256 | 3 s2 | 57.8 us | 79.4 us | **1.4x** | +| 8x128x512 | 3 s2 | 316 us | 828 us | **2.6x** | +| 16x256x1024 | 3 s2 | 1.73 ms | 5.41 ms | **3.1x** | + +### Kernel Size Comparison (4x64x56x56) + +| Kernel | Flex | NdArray | Speedup | +| ------ | ------- | ------- | -------- | +| 2x2 | 221 us | 317 us | **1.4x** | +| 3x3 | 375 us | 515 us | **1.4x** | +| 5x5 | 1.03 ms | 799 us | 0.78x | + +--- + +## Transposed Convolutions + +### Conv Transpose 2D + +| Input | Output | Flex | NdArray | Speedup | +| -------------- | ------ | ------- | ------- | ------- | +| 1x64x7x7 | 14x14 | 138 us | 1.67 ms | **12x** | +| 1x128x14x14 | 28x28 | 533 us | 12.9 ms | **24x** | +| 1x256x28x28 | 56x56 | 2.49 ms | 209 ms | **84x** | +| 1x512x7x7 k3s1 | 7x7 | 1.01 ms | 52.6 ms | **52x** | +| 8x64x14x14 | 28x28 | 3.41 ms | 49.7 ms | **15x** | + +### DCGAN Generator + +| Layer | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| -------------- | ------- | ------- | -------- | -------- | ----------- | +| 1x1 to 4x4 | 156 us | 1.43 ms | **9.2x** | 33.1 KB | 16.4 KB | +| 4x4 to 8x8 | 234 us | 3.83 ms | **16x** | 164 KB | 32.8 KB | +| 8x8 to 16x16 | 305 us | 4.24 ms | **14x** | 852 KB | 65.6 KB | +| 16x16 to 32x32 | 26.4 us | 1.47 ms | **56x** | 193 KB | 12.3 KB | + +### Conv Transpose 1D + +| Input | Flex | NdArray | Speedup | +| --------- | ------- | ------- | ------- | +| 1x64x32 | 17.5 us | 383 us | **22x** | +| 8x128x64 | 433 us | 9.02 ms | **21x** | +| 1x256x128 | 337 us | 8.95 ms | **27x** | + +### Conv Transpose 3D + +| Input | Output | Flex | NdArray | Speedup | +| ---------- | -------- | ------- | ------- | ------- | +| 1x32x4x4x4 | 8x8x8 | 245 us | 2.58 ms | **11x** | +| 1x64x8x8x8 | 16x16x16 | 1.41 ms | 47.4 ms | **34x** | + +--- + +## Interpolation + +### Nearest + +| Input | Output | Flex | NdArray | Speedup | +| ----------- | ------- | ------- | ------- | -------- | +| 1x3x64x64 | 128x128 | 22.7 us | 138 us | **6.1x** | +| 1x3x32x32 | 128x128 | 22.8 us | 143 us | **6.3x** | +| 1x3x256x256 | 128x128 | 22.5 us | 143 us | **6.3x** | +| 8x3x64x64 | 128x128 | 58.5 us | 307 us | **5.2x** | +| 1x64x32x32 | 64x64 | 57.9 us | 254 us | **4.4x** | + +### Bilinear + +| Input | Output | Flex | NdArray | Speedup | +| ----------- | ------- | ------- | ------- | -------- | +| 1x3x64x64 | 128x128 | 81.2 us | 161 us | **2.0x** | +| 1x3x32x32 | 128x128 | 86.5 us | 147 us | **1.7x** | +| 1x3x256x256 | 128x128 | 83.9 us | 157 us | **1.9x** | +| 8x3x64x64 | 128x128 | 181 us | 382 us | **2.1x** | +| 1x64x32x32 | 64x64 | 108 us | 309 us | **2.8x** | + +### Bicubic + +| Input | Output | Flex | NdArray | Speedup | +| ----------- | ------- | ------ | ------- | -------- | +| 1x3x64x64 | 128x128 | 159 us | 239 us | **1.5x** | +| 1x3x32x32 | 128x128 | 160 us | 232 us | **1.4x** | +| 1x3x256x256 | 128x128 | 159 us | 238 us | **1.5x** | +| 8x3x64x64 | 128x128 | 894 us | 994 us | **1.1x** | +| 1x64x32x32 | 64x64 | 616 us | 707 us | **1.1x** | + +--- + +## Grid Sample 2D + +| Input | Grid | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| ---------- | ----- | ------- | ------- | -------- | -------- | ----------- | +| 1x3x32x32 | 32x32 | 16.7 us | 87.6 us | **5.2x** | 12.5 KB | 12.3 KB | +| 1x3x64x64 | 64x64 | 66.9 us | 127 us | **1.9x** | 49.4 KB | 49.2 KB | +| 4x3x32x32 | 32x32 | 60.7 us | 152 us | **2.5x** | 49.4 KB | 49.2 KB | +| 1x16x64x64 | 64x64 | 291 us | 223 us | 0.77x | 262 KB | 262 KB | + +--- + +## Cross Product & Unfold + +### Cross Product + +| Shape | Flex | NdArray | Speedup | +| ------- | ------- | ------- | -------- | +| 1Kx3 | 31.1 us | 43.5 us | **1.4x** | +| 64Kx3 | 1.88 ms | 2.78 ms | **1.5x** | +| 256Kx3 | 7.58 ms | 11.1 ms | **1.5x** | +| 64x3x64 | 141 us | 290 us | **2.1x** | + +### Unfold (1D) + +| Input | Window | Step | Flex | NdArray | Speedup | +| ----- | ------ | ---- | ------- | ------- | ------------ | +| 1K | 8 | 1 | 62.2 ns | 120 us | **~1927x** | +| 64K | 8 | 1 | 64.1 ns | 7.55 ms | **~117686x** | +| 64K | 64 | 1 | 62.2 ns | 8.04 ms | **~129295x** | +| 64K | 64 | 32 | 62.2 ns | 254 us | **~4094x** | + +### Unfold (2D/3D) + +| Shape | Dim | Window | Step | Flex | NdArray | Speedup | +| -------- | --- | ------ | ---- | ------- | ------- | ----------- | +| 256x256 | 1 | 8 | 1 | 68.7 ns | 871 us | **~12685x** | +| 256x256 | 1 | 32 | 16 | 62.5 ns | 57.7 us | **~924x** | +| 1024x256 | 1 | 8 | 1 | 62.8 ns | 3.28 ms | **~52182x** | +| 32x64x64 | 2 | 8 | 4 | 77.1 ns | 424 us | **~5497x** | + +--- + +## Deformable Convolutions + +### Small/Tiny Inputs + +| Input | Config | Flex | NdArray | Speedup | +| --------- | ----------- | ------- | ------- | -------- | +| 1x3x8x8 | 3 to 8, k3 | 8.72 us | 92.6 us | **11x** | +| 1x3x8x8 | no mask | 7.98 us | 78.9 us | **9.9x** | +| 1x3x16x16 | 3 to 16, k3 | 36.4 us | 122 us | **3.4x** | +| 1x3x16x16 | stride 2 | 10.2 us | 80.8 us | **7.9x** | +| 2x8x16x16 | 8 to 16, k3 | 116 us | 246 us | **2.1x** | + +### Medium Inputs + +| Input | Config | Flex | NdArray | Speedup | +| ---------- | ------------ | ------ | ------- | ------- | +| 1x16x32x32 | 16 to 32, k3 | 826 us | 581 us | 0.70x | +| 1x16x32x32 | wg=4 | 840 us | 528 us | 0.63x | +| 1x16x32x32 | og=4 | 942 us | 606 us | 0.64x | + +--- + +## Attention (Scaled Dot-Product) + +Flex auto-selects between two gemm-backed strategies: + +- **Naive** (score matrix <= 256K elements): Materializes full [seq_q, seq_kv] score matrix. Two + large gemm calls per (batch, head) amortize dispatch overhead better than many small tiled calls. +- **Flash** (score matrix > 256K elements): Tiles over KV dimension with online softmax. + `O(seq_q * TILE_KV)` memory per head instead of `O(seq_q * seq_kv)`. + +Both fuse scale + softcap + masking + bias + softmax into a single pass, reducing intermediate +allocations from ~12 (NdArray fallback) to 3. + +### Self-Attention + +| Config | Flex | NdArray | Speedup | +| --------------- | ------- | ------- | -------- | +| h8, s64, d64 | 180 us | 534 us | **3.0x** | +| h12, s128, d64 | 989 us | 1.61 ms | **1.6x** | +| h12, s256, d64 | 3.79 ms | 6.03 ms | **1.6x** | +| h12, s512, d64 | 14.8 ms | 22.7 ms | **1.5x** | +| h32, s256, d128 | 15.1 ms | 17.5 ms | **1.2x** | +| b4, h12, s128 | 3.96 ms | 5.47 ms | **1.4x** | + +### Causal Attention + +| Config | Flex | NdArray | Speedup | +| -------------- | ------- | ------- | -------- | +| h12, s128, d64 | 1.01 ms | 1.72 ms | **1.7x** | +| h12, s256, d64 | 3.82 ms | 6.47 ms | **1.7x** | +| h12, s512, d64 | 14.8 ms | 23.8 ms | **1.6x** | + +### With Additive Bias (ALiBi-style) + +| Config | Flex | NdArray | Speedup | +| -------------- | ------- | ------- | -------- | +| h12, s128, d64 | 1.04 ms | 1.60 ms | **1.5x** | +| h12, s256, d64 | 4.00 ms | 6.15 ms | **1.5x** | + +### Cross-Attention (seq_q != seq_k) + +| Config | Flex | NdArray | Speedup | +| ----------------- | ------- | ------- | -------- | +| sq128, sk512, d64 | 3.82 ms | 6.14 ms | **1.6x** | +| sq32, sk1024, d64 | 2.05 ms | 3.72 ms | **1.8x** | + +--- + +## Quantized Tensor Operations + +All quantized ops (except layout ops) go through a dequantize-op-quantize cycle. Flex stores scales +separately and applies `scale * x_q` directly; NdArray reparses `QuantizedBytes` on every dequantize +call, which dominates the cost. + +### Quantize (float to i8) + +| Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| ---- | ------- | ------- | -------- | -------- | ----------- | +| 4K | 6.93 us | 10.4 us | **1.5x** | 20.6 KB | 24.7 KB | +| 64K | 109 us | 145 us | **1.3x** | 328 KB | 393 KB | +| 1M | 1.75 ms | 2.31 ms | **1.3x** | 5.24 MB | 6.29 MB | + +### Dequantize (i8 to float) + +| Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| ---- | ------- | ------- | --------- | -------- | ----------- | +| 4K | 399 ns | 48.8 us | **~122x** | 16.5 KB | 24.6 KB | +| 64K | 3.73 us | 801 us | **~215x** | 262 KB | 393 KB | +| 1M | 54.6 us | 13.0 ms | **~238x** | 4.19 MB | 6.29 MB | + +### q_add (dequant + add + requant) + +| Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| ---- | ------- | ------- | --------- | -------- | ----------- | +| 4K | 1.15 us | 101 us | **88x** | 20.6 KB | 41.0 KB | +| 64K | 14.3 us | 1.61 ms | **~113x** | 524 KB | 655 KB | +| 1M | 208 us | 25.9 ms | **~125x** | 8.39 MB | 10.5 MB | + +### q_matmul (dequant + matmul + requant) + +| Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| ------- | ------- | ------- | ------- | -------- | ----------- | +| 64x64 | 7.10 us | 137 us | **19x** | 66.5 KB | 49.3 KB | +| 256x256 | 147 us | 1.93 ms | **13x** | 1.05 MB | 788 KB | +| 512x512 | 641 us | 8.07 ms | **13x** | 4.19 MB | 3.15 MB | + +### q_sum (dequant + sum) + +| Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| ---- | ------- | ------- | --------- | -------- | ----------- | +| 4K | 605 ns | 50.7 us | **84x** | 2.13 KB | 24.6 KB | +| 64K | 7.71 us | 834 us | **~108x** | 262 KB | 393 KB | +| 1M | 194 us | 13.3 ms | **68x** | 4.19 MB | 6.29 MB | + +### q_permute (zero-copy layout op) + +| Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| --------- | ------- | ------- | ------- | -------- | ----------- | +| 256x256 | 75.5 ns | 66.1 ns | 0.88x | 20.5 B | 4.00 B | +| 1024x1024 | 77.5 ns | 66.1 ns | 0.85x | 20.5 B | 4.00 B | + +### q_argmax (operates on i8 directly) + +| Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| --------- | ------- | ------- | -------- | -------- | ----------- | +| 256x256 | 67.9 us | 106 us | **1.6x** | 3.25 KB | 4.16 KB | +| 1024x1024 | 137 us | 1.71 ms | **12x** | 12.5 KB | 16.4 KB | + +### q_argmin (operates on i8 directly) + +| Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| --------- | ------- | ------- | -------- | -------- | ----------- | +| 256x256 | 69.4 us | 106 us | **1.5x** | 3.25 KB | 4.16 KB | +| 1024x1024 | 139 us | 1.72 ms | **12x** | 12.5 KB | 16.4 KB | + +### q_gather (operates on i8 directly for tensor-level quant) + +| Size | Flex | NdArray | Speedup | Flex Mem | NdArray Mem | +| --------- | ------- | ------- | -------- | -------- | ----------- | +| 256x256 | 65.7 us | 155 us | **2.4x** | 590 KB | 721 KB | +| 1024x1024 | 393 us | 2.40 ms | **6.1x** | 9.44 MB | 11.5 MB | + +--- + +## Default Ops (sort, repeat, creation, embedding, predicates) + +These ops override burn's default trait implementations with direct storage operations. + +### Sort (f32, 1D) + +| Size | Flex | NdArray | Speedup | +| ---- | ------- | ------- | -------- | +| 4K | 53.6 us | 123 us | **2.3x** | +| 64K | 593 us | 1.59 ms | **2.7x** | +| 1M | 8.47 ms | 23.9 ms | **2.8x** | + +### Sort (f32, 2D along last dim) + +| Size | Flex | NdArray | Speedup | +| --------- | ------- | ------- | -------- | +| 64x64 | 5.72 us | 70.3 us | **12x** | +| 256x256 | 200 us | 1.26 ms | **6.3x** | +| 1024x1024 | 1.17 ms | 34.5 ms | **29x** | + +### Argsort (f32, 1D) + +| Size | Flex | NdArray | Speedup | +| ---- | ------- | ------- | -------- | +| 4K | 70.2 us | 123 us | **1.7x** | +| 1M | 12.8 ms | 26.2 ms | **2.1x** | + +### Repeat Dim (f32, 256x256) + +| Config | Flex | NdArray | Speedup | +| ----------------- | ------- | ------- | -------- | +| dim0 4x | 12.7 us | 134 us | **11x** | +| dim1 4x | 11.5 us | 138 us | **12x** | +| dim0 8x (512x512) | 98.4 us | 880 us | **8.9x** | + +### Tensor Creation (f32, 1M elements) + +| Operation | Flex | NdArray | Speedup | +| --------- | ------- | ------- | ------- | +| zeros | 17.7 us | 579 us | **33x** | +| ones | 35.7 us | 579 us | **16x** | +| full | 35.9 us | 579 us | **16x** | + +### Arange (i64) + +| Size | Flex | NdArray | Speedup | +| ---- | ------- | ------- | ------- | +| 4K | 1.21 us | 1.21 us | ~1x | +| 1M | 299 us | 280 us | 0.94x | + +### Embedding (f32) + +| Config | Flex | NdArray | Speedup | +| ----------------------- | ------- | ------- | -------- | +| 30k vocab, d=512, 8x128 | 26.1 us | 150 us | **5.8x** | +| 50k vocab, d=768, 4x256 | 38.9 us | 188 us | **4.8x** | + +### Predicates (f32, 1M elements) + +| Operation | Flex | NdArray | Speedup | +| --------- | ------- | ------- | -------- | +| is_nan | 46.4 us | 73.6 us | **1.6x** | +| is_inf | 52.6 us | 147 us | **2.8x** | + +--- + +## FFT (Real FFT) + +Compared against `realfft` v3 (backed by `rustfft` v6), the gold-standard pure-Rust FFT library. +NdArray does not implement rfft. `realfft` requires `std`; Flex works in `no_std`. + +### 1D rfft + +| Size | Flex (median) | realfft (median) | Ratio | +| ------- | ------------- | ---------------- | ----- | +| n=256 | 841 ns | 252 ns | 3.3x | +| n=1024 | 2.64 us | 991 ns | 2.7x | +| n=4096 | 10.6 us | 4.37 us | 2.4x | +| n=16384 | 45.4 us | 20.7 us | 2.2x | +| n=65536 | 231 us | 91.2 us | 2.5x | + +### Batched 2D rfft (along last dim) + +| Size | Flex (median) | realfft (median) | Ratio | +| --------- | ------------- | ---------------- | ----- | +| 16 x 1024 | 59.9 us | 15.9 us | 3.8x | +| 64 x 1024 | 114 us | 64.0 us | 1.8x | +| 256 x 256 | 191 us | 64.9 us | 2.9x | + +### 1D irfft (inverse) + +| Size | Flex (median) | realfft (median) | Ratio | +| ------- | ------------- | ---------------- | ----- | +| n=256 | 1.32 us | 207 ns | 6.4x | +| n=1024 | 4.01 us | 908 ns | 4.4x | +| n=4096 | 15.6 us | 4.24 us | 3.7x | +| n=16384 | 63.2 us | 21.1 us | 3.0x | +| n=65536 | 278 us | 93.0 us | 3.0x | + +Flex implementation: Cooley-Tukey with mixed radix-4/radix-2, complex packing (forward) / inverse +packing (inverse), compile-time twiddle tables via const fn, SIMD vectorization via macerator, +unrolled small kernels (N=2,4,8), and rayon parallelism across fibers. The remaining gap to rustfft +is due to their hand-tuned per-arch SIMD rewrites, split-radix algorithms, and strength-reduced +modular arithmetic. + +--- + +## Running Benchmarks + +```bash +cargo bench --bench attention +cargo bench --bench binary_ops +cargo bench --bench matmul +cargo bench --bench int_ops +cargo bench --bench slice_ops +cargo bench --bench reduce_ops +cargo bench --bench cumulative_ops +cargo bench --bench gather_scatter_ops +cargo bench --bench unary_ops +cargo bench --bench comparison_ops +cargo bench --bench conv_ops +cargo bench --bench pool_ops +cargo bench --bench conv_transpose_ops +cargo bench --bench interpolate_ops +cargo bench --bench cross_unfold_ops +cargo bench --bench deform_conv_ops +cargo bench --bench quantization_ops +cargo bench --bench cat_max_min_ops +cargo bench --bench default_ops +cargo bench --bench fft_ops +``` diff --git a/crates/burn-flex/COMPARISON.md b/crates/burn-flex/COMPARISON.md new file mode 100644 index 0000000..9163fe6 --- /dev/null +++ b/crates/burn-flex/COMPARISON.md @@ -0,0 +1,624 @@ +# burn-flex vs burn-ndarray: Comprehensive Comparison + +This document compares burn-flex (proposed replacement) against burn-ndarray (current CPU backend) +to demonstrate full coverage and the architectural differences between the two. + +## Executive Summary + +burn-flex is a from-scratch CPU backend built to replace burn-ndarray. The +[ndarray](https://crates.io/crates/ndarray) crate has been slow to evolve: it lacks f16/bf16 +support, is limited to 6 dimensions, uses unsigned-only strides (preventing zero-copy flip), and +simulates quantization rather than executing natively. burn-flex addresses all of these while +passing the full `burn-backend-tests` suite, all ONNX model checks, and real model inference +(ALBERT, MiniLM). + +Performance improvements fall into two categories: + +- **Compute gains** (1.1-9.7x): Better algorithms and libraries (gemm over matrixmultiply, Arc COW + for buffer reuse, SIMD reductions). +- **Structural improvements** (up to 166,000x): Operations that burn-ndarray eagerly materializes + (unfold, expand, slice, dequantize) are represented as zero-copy views or direct lookups in + burn-flex, avoiding the work entirely. + +burn-flex uses significantly less memory, supports f16/bf16 natively, runs on no_std/WASM/embedded, +and has no dimension limit. + +--- + +## 1. Architecture + +### Tensor Representation + +| Aspect | burn-flex | burn-ndarray | +| ----------- | ----------------------------------------------------- | ------------------------------------------------------------------------- | +| Storage | `Arc` (type-erased bytes) | `enum NdArrayTensor { F64(NdArrayStorage), F32(...), ... }` | +| Dtype | Runtime `DType` field on `FlexTensor` | Compile-time via enum variant | +| Dispatch | `match dtype` at op entry, cast once | `execute_with_dtype!` macro expands match for every op | +| Clone cost | O(1) Arc refcount increment | O(1) ArcArray refcount increment | +| COW | `Arc::make_mut` / `is_unique()` | `ArcArray::is_unique()` + `NdArrayStorage::Borrowed` always returns false | +| Metadata | `Layout { shape, strides: Vec, start_offset }` | ndarray's internal strides (`usize` only) | +| Stride sign | **Signed** (`isize`) for zero-copy flip | **Unsigned** (`usize`), flip requires data copy | + +**FlexTensor** (44 bytes without shape vec): + +```rust +struct FlexTensor { + data: Arc, // 8 bytes (pointer) + layout: Layout, // shape + strides + offset + dtype: DType, // 1 byte enum +} +``` + +**NdArrayTensor** (enum with 11 typed variants): + +```rust +enum NdArrayTensor { + F64(NdArrayStorage), + F32(NdArrayStorage), + // ... 9 more variants +} +``` + +**Key insight**: Flex uses one struct for all dtypes with runtime dispatch. NdArray uses a typed +enum with macro-based dispatch. Flex's approach is simpler (no macros, no generics plumbing) and +enables operations to handle all dtypes uniformly. + +### Backend Type + +| Aspect | burn-flex | burn-ndarray | +| ------------- | ---------------------------- | ------------------------------------------------------- | +| Type | `struct Flex;` (unit struct) | `struct NdArray` (3 generic params) | +| Float element | Runtime (f32/f64/f16/bf16) | Compile-time `E: FloatNdArrayElement` (f32 or f64 only) | +| Int element | Runtime (i8-i64, u8-u64) | Compile-time `I: IntNdArrayElement` | +| Quant element | Runtime | Compile-time `Q: QuantElement` | + +Flex eliminates generic parameters entirely. Users write `Flex` instead of `NdArray`. +Dtype selection happens at runtime via `DType`. + +--- + +## 2. Feature Coverage + +### Float Dtypes + +| Dtype | burn-flex | burn-ndarray | +| -------- | ----------------------------------------------------------- | --------------------- | +| f32 | Full support (native) | Full support (native) | +| f64 | Full support (native) | Full support (native) | +| **f16** | **Full support (native)** | **Not supported** | +| **bf16** | **Full support (via f32 conversion for compute-heavy ops)** | **Not supported** | +| Flex32 | Not applicable | Maps to f32 | + +burn-flex's f16 support is native for all operations. For matmul and convolution, the gemm crate has +native f16 kernels (since v0.15). bf16 converts to f32 for compute-heavy ops (matmul, conv) because +gemm lacks native bf16 support. + +### Integer Dtypes + +| Dtype | burn-flex | burn-ndarray | +| ----- | ------------ | ------------ | +| i64 | Full support | Full support | +| i32 | Full support | Full support | +| i16 | Full support | Full support | +| i8 | Full support | Full support | +| u64 | Full support | Full support | +| u32 | Full support | Full support | +| u16 | Full support | Full support | +| u8 | Full support | Full support | + +Both backends support the same integer dtypes. + +### Bool + +| Feature | burn-flex | burn-ndarray | +| ---------- | ------------------------- | --------------------------------------- | +| Storage | `u8` (1 byte per element) | `bool` (1 byte per element via ndarray) | +| Operations | All BoolTensorOps | All BoolTensorOps | + +### Quantization + +| Feature | burn-flex | burn-ndarray | +| -------------- | --------------------------------------------------------------- | --------------------------------------- | +| Quantize | Per-tensor and per-block symmetric | Per-tensor and per-block symmetric | +| Dequantize | `scale * x_q` (direct multiply, **135-232x faster**) | Reparses `QuantizedBytes` on every call | +| Scale storage | `Vec` stored separately | `QParams` in `NdArrayQTensor` | +| Q layout ops | **Zero-copy** (permute, flip, expand, slice, select) | Copies entire tensor | +| Q ordering ops | **Skip dequantization** (argmax, argmin, gather on i8 directly) | Dequantize to f32, then operate | +| QuantStore | Native | Native | +| QuantValue | Q8F, Q8S | Q8F, Q8S (+ Q4/Q2 for export_tests) | + +The fundamental difference is scale storage. Flex stores scales separately so dequantization is a +simple `scale * x_q` multiply. NdArray stores everything in `QuantizedBytes` which must be parsed on +every access, making it the bottleneck for all quantized operations. + +--- + +## 3. Operation Coverage + +### Tensor Operations (FloatTensorOps) + +All operations listed below are implemented by both backends unless marked otherwise. + +| Operation | burn-flex | burn-ndarray | Notes | +| ----------------------------------------------------------------- | --------- | ------------ | ----------------------------------------------------- | +| from_data | Yes | Yes | | +| into_data | Yes | Yes | | +| random | Yes | Yes | | +| empty/zeros/ones | Yes | Yes | | +| full | Yes | Yes | | +| add / sub / mul / div | Yes | Yes | | +| add_scalar / sub_scalar / mul_scalar / div_scalar | Yes | Yes | | +| remainder | Yes | Yes | | +| remainder_scalar | Yes | Yes | | +| matmul | Yes | Yes | Flex uses gemm, NdArray uses matrixmultiply | +| neg | Yes | Yes | | +| recip | Yes | Yes | | +| swap_dims / permute | Yes | Yes | Both zero-copy | +| reshape | Yes | Yes | Both zero-copy when contiguous | +| gather / scatter_add | Yes | Yes | | +| select / select_add | Yes | Yes | | +| slice / slice_assign | Yes | Yes | Flex: zero-copy view; NdArray: may copy | +| mask_fill / mask_where | Yes | Yes | | +| equal / not_equal / greater / lower / greater_equal / lower_equal | Yes | Yes | | +| equal_elem / not_equal_elem / greater_elem / lower_elem | Yes | Yes | | +| sum / sum_dim / mean / mean_dim / prod / prod_dim | Yes | Yes | | +| max / max_dim / max_dim_with_indices | Yes | Yes | | +| min / min_dim / min_dim_with_indices | Yes | Yes | | +| argmax / argmin | Yes | Yes | | +| any / any_dim / all / all_dim | Yes | Yes | | +| exp / log / log1p | Yes | Yes | | +| powf / powf_scalar / powi / powi_scalar | Yes | Yes | | +| sqrt / abs / sign | Yes | Yes | | +| cos / sin / tanh | Yes | Yes | | +| erf | Yes | Yes | | +| cat | Yes | Yes | | +| into_int / into_bool | Yes | Yes | | +| clamp / clamp_min / clamp_max | Yes | Yes | | +| expand | Yes | Yes | Flex: zero-copy; NdArray: copies | +| flip | Yes | Yes | Flex: zero-copy (signed strides); NdArray: copies | +| repeat_dim | Yes | Yes | | +| sort / sort_with_indices / argsort | Yes | Yes | | +| cumsum / cumprod / cummin / cummax | Yes | Yes | | +| narrow | Yes | Yes | Flex: zero-copy; NdArray: may copy | +| chunk | Yes | Yes | | +| cross | Yes | Yes | | +| unfold | Yes | Yes | Flex: zero-copy (strided view); NdArray: materializes | +| round / floor / ceil | Yes | Yes | | +| cast | Yes | Yes | | +| grid_sample_2d | Yes | Yes | | +| bool_select | Yes | Yes | | +| int_powi | Yes | Yes | | + +### Module Operations (ModuleOps) + +| Operation | burn-flex | burn-ndarray | Notes | +| -------------------------------- | --------- | ------------ | -------------------------------------------------------------------------------------- | +| conv1d | Yes | Yes | Flex: delegates to conv3d | +| conv2d | Yes | Yes | Flex: delegates to conv3d | +| conv3d | Yes | Yes | Flex: unified implementation | +| conv_transpose1d | Yes | Yes | Flex: delegates to conv_transpose3d | +| conv_transpose2d | Yes | Yes | Flex: delegates to conv_transpose3d | +| conv_transpose3d | Yes | Yes | Flex: unified implementation | +| deform_conv2d | Yes | Yes | | +| deform_conv2d_backward | Yes | Yes | | +| avg_pool2d | Yes | Yes | Flex: delegates to pool3d | +| avg_pool2d_backward | Yes | Yes | | +| max_pool2d | Yes | Yes | Flex: delegates to pool3d | +| max_pool2d_with_indices | Yes | Yes | | +| max_pool2d_with_indices_backward | Yes | Yes | | +| adaptive_avg_pool2d | Yes | Yes | | +| adaptive_avg_pool2d_backward | Yes | Yes | | +| interpolate | Yes | Yes | Nearest, bilinear, bicubic | +| attention (SDPA) | Yes | Yes | Flex: auto-selects naive or flash by score matrix size; NdArray: matmul + softmax | +| rfft | Yes | No | Flex: Cooley-Tukey with complex packing, radix-4, SIMD, compile-time twiddles. no_std. | +| irfft | Yes | No | Flex: Inverse packing trick, SIMD via conjugate-forward-conjugate. no_std. | + +### Int and Bool Operations + +Both backends implement all IntTensorOps and BoolTensorOps. The operations mirror float ops where +applicable (arithmetic, comparison, reduction, gather/scatter, slice, etc.) plus type-specific +operations (int_random uniform, bool_not, bool_and, bool_or, bool_xor). + +### Quantized Operations (QTensorOps) + +Both backends implement all QTensorOps. The ops follow a dequantize-op-requantize pattern for most +operations. Flex optimizes by: + +- Storing scales separately for O(1) dequantization access +- Zero-copy layout ops on quantized tensors (permute, flip, expand, slice, select) +- Skipping dequantization for ordering ops (argmax, argmin, gather with tensor-level quant) + +### Activation Operations (ActivationOps) + +Both backends implement all ActivationOps via the default trait implementations (relu, gelu, etc.). + +### Transaction Operations + +Both backends implement TransactionOps for batched tensor operations. + +--- + +## 4. Dimension Limits + +| Aspect | burn-flex | burn-ndarray | +| -------------- | -------------------------------- | ---------------------------------------------- | +| Max dimensions | **Unlimited** (arbitrary rank) | **6** (hardcoded in reshape macro) | +| Enforcement | Dynamic `Vec` for strides | Static `Dim<[usize; N]>` requires match on 1-6 | + +burn-ndarray's dimension limit comes from its `reshape!` macro which matches on dimensions 1-6: + +```rust +match $D { + 1 => reshape!(ty $ty, n 1, ...), + // ... + 6 => reshape!(ty $ty, n 6, ...), + _ => panic!("NdArray supports arrays up to 6 dimensions"), +} +``` + +burn-flex uses `IxDyn`-equivalent dynamic shapes with no upper bound. + +--- + +## 5. Zero-Copy Operations + +| Operation | burn-flex | burn-ndarray | +| -------------- | --------------------------------- | ------------------------------ | +| transpose | Zero-copy (swap strides) | Zero-copy (ndarray view) | +| permute | Zero-copy (reorder strides) | Zero-copy (ndarray view) | +| reshape | Zero-copy if contiguous | Zero-copy if standard layout | +| slice / narrow | Zero-copy (offset + strides) | May allocate depending on path | +| **flip** | **Zero-copy (negate stride)** | **Copies data** | +| **unfold** | **Zero-copy (O(1) strided view)** | **O(n) full materialization** | +| **expand** | **Zero-copy (set stride to 0)** | **Copies data** | + +Flex's signed strides (`isize`) enable zero-copy flip, which is impossible with ndarray's unsigned +strides. The unfold operation is especially dramatic: Flex returns a strided view in ~50ns +regardless of size, while NdArray copies all window data (milliseconds for large tensors). + +--- + +## 6. Memory Strategy + +### In-Place Mutation + +| Strategy | burn-flex | burn-ndarray | +| ------------------ | ----------------------------------------------------- | ---------------------------------- | +| Unique check | `Arc::strong_count(&data) == 1` | `ArcArray::is_unique()` | +| In-place threshold | Contiguous at offset 0 AND unique | Unique (via SIMD ops, not all ops) | +| Binary op reuse | Reuses lhs buffer when contiguous | Allocates new for most ops | +| Allocation savings | 3x less for binary ops (4.2 MB vs 12.6 MB for 1M f32) | Standard ndarray allocation | + +### Zero-Copy Loading + +Both backends support zero-copy loading from external sources (burnpack files, mmap'd data): + +| Feature | burn-flex | burn-ndarray | +| ----------- | ----------------------------------------- | ------------------------------------------------ | +| Mechanism | `Arc` wraps borrowed data directly | `NdArrayStorage::Borrowed` holds `Bytes` + shape | +| COW trigger | `Arc::make_mut` clones on shared mutation | `into_owned()` copies borrowed to ArcArray | +| View access | `storage::()` via bytemuck cast | `view()` via unsafe ArrayView from raw pointer | + +--- + +## 7. SIMD + +| Aspect | burn-flex | burn-ndarray | +| ------------ | ----------------------------------------------------------- | ---------------------------------------------- | +| Library | macerator (required with `simd` feature) | macerator (optional with `simd` feature) | +| Dispatch | `Arch::new().dispatch(kernel)` | Same macerator dispatch | +| ISAs | NEON, AVX2, AVX512, SSE, SIMD128, scalar fallback | NEON, AVX2, SSE, SIMD128, scalar fallback | +| Coverage | Binary ops, comparisons, boolean ops, reductions, unary ops | Binary ops, comparisons, unary ops, conv, pool | +| Without SIMD | Scalar fallback module (`simd/scalar.rs`) | Falls back to ndarray operations | + +Both use macerator for portable SIMD. NdArray additionally has SIMD-optimized conv and pool kernels. +Flex relies on the gemm crate's built-in SIMD for matmul/conv performance. + +--- + +## 8. Matrix Multiplication + +| Aspect | burn-flex | burn-ndarray | +| ----------- | -------------------------------------- | ---------------------------------------------------- | +| Library | `gemm` crate (v0.18) | `matrixmultiply` crate (via ndarray) | +| f32 | Native gemm kernel | matrixmultiply | +| f64 | Native gemm kernel | matrixmultiply | +| f16 | **Native gemm kernel (since v0.15)** | **Not supported** | +| bf16 | Convert to f32, gemm, convert back | **Not supported** | +| i32 matmul | Manual nested loop | Manual nested loop | +| Parallelism | Rayon via gemm (threshold: 192^3) | Rayon via iter_range_par macro | +| Batched | Parallel over batches + per-batch gemm | Parallel over batches + ndarray general_mat_mul | +| Broadcast | Handles batch broadcast natively | Handles batch broadcast via stride mapping | +| BLAS option | No (pure Rust only) | Yes (Accelerate, OpenBLAS, Netlib via feature flags) | + +burn-ndarray offers optional BLAS acceleration (Accelerate on macOS, OpenBLAS, Netlib) through +feature flags. burn-flex uses only the gemm crate, which is pure Rust but highly optimized with its +own SIMD kernels. The gemm crate consistently outperforms matrixmultiply by 1.3-3.4x on Apple M3 +Max. + +--- + +## 9. Convolutions + +| Aspect | burn-flex | burn-ndarray | +| ------------ | ----------------------------- | -------------------------------------------------- | +| Algorithm | im2col + gemm (unified 3D) | Direct computation (per-dimension implementations) | +| conv1d | Delegates to conv3d | Separate implementation | +| conv2d | Delegates to conv3d | Separate implementation | +| conv3d | Single unified implementation | Separate implementation | +| f16 support | **Native gemm** | **Not supported** | +| bf16 support | **Via f32 conversion** | **Not supported** | +| Parallelism | Rayon over batches and groups | iter_range_par over batches | +| SIMD conv | Via gemm SIMD kernels | macerator-based SIMD conv kernel | + +Flex's unified 3D approach means one implementation covers all dimensionalities. The tradeoff is +that 1D/2D convolutions expand dimensions (negligible overhead since gemm dominates). + +NdArray has dedicated SIMD conv/pool kernels via macerator, which can be faster for specific +patterns. Flex relies on the gemm crate's SIMD for all compute-heavy paths. + +--- + +## 10. Parallelism + +| Aspect | burn-flex | burn-ndarray | +| ---------------- | ------------------------------------- | ----------------------------------------- | +| Library | rayon (optional) | rayon (optional, called "multi-threads") | +| Feature flag | `rayon` | `multi-threads` | +| Threshold | 4M elements for memory-bound ops | Via `run_par!` / `iter_range_par!` macros | +| Scope | Large tensors, batch dims, pool, conv | Matmul batches, ops via macros | +| gemm parallelism | Rayon via `Parallelism::Rayon(0)` | matrixmultiply threading | +| Without feature | Single-threaded (all ops work) | Single-threaded (all ops work) | + +--- + +## 11. Platform Support + +| Target | burn-flex | burn-ndarray | +| ------------------------------- | ---------------------------------- | ------------------------- | +| x86_64 (std) | Yes | Yes | +| aarch64 (std) | Yes (primary target) | Yes | +| wasm32-unknown-unknown | **Yes (verified)** | Yes (claimed, categories) | +| thumbv6m-none-eabi (Cortex-M0+) | **Yes (verified, no atomic ptrs)** | Not verified | +| thumbv7m-none-eabi (Cortex-M3) | **Yes (verified)** | Not verified | +| no_std | **Yes (tested, MNIST inference)** | Yes (supported) | + +burn-flex has been explicitly tested on embedded targets with Burn's `burn-no-std-tests` integration +suite (MNIST model inference). + +--- + +## 12. Dependencies + +### burn-flex + +| Dependency | Purpose | Required | +| ------------ | ----------------------------------- | ------------------ | +| burn-backend | Backend traits, types | Always | +| burn-ir | BackendIr trait | Always | +| burn-std | Bytes, Shape, platform abstractions | Always | +| half | f16/bf16 types | Always | +| bytemuck | Zero-copy type casting | Always | +| num-traits | Numeric traits (libm for no_std) | Always | +| gemm | Matrix multiplication | Always | +| macerator | Portable SIMD | Optional (`simd`) | +| aligned-vec | SIMD-aligned allocation | Optional (`simd`) | +| rayon | Parallelism | Optional (`rayon`) | + +**Total: 7 required + 3 optional** + +### burn-ndarray + +| Dependency | Purpose | Required | +| -------------------- | -------------------------------- | -------------------------- | +| burn-backend | Backend traits, types | Always | +| burn-std | Platform abstractions | Always | +| burn-autodiff | Autodiff support | Optional (`std`) | +| burn-ir | IR types | Always | +| ndarray | N-dimensional array library | Always | +| matrixmultiply | Matrix multiplication | Always | +| atomic_float | Atomic f32/f64 | Always | +| const-random | Compile-time random | Always | +| libm | Math functions for no_std | Always | +| num-traits | Numeric traits | Always | +| paste | Macro utilities | Always | +| rand | Random number generation | Always | +| macerator | Portable SIMD | Optional (`simd`) | +| bytemuck | Type casting | Optional (`simd`) | +| itertools | Iterator utilities | Optional (`simd`) | +| seq-macro | Sequence macros | Optional (`simd`) | +| rayon | Parallelism | Optional (`multi-threads`) | +| blas-src | BLAS bindings | Optional (`blas-*`) | +| openblas-src | OpenBLAS | Optional (`blas-openblas`) | +| portable-atomic | Atomic for no-atomic-ptr targets | Conditional | +| portable-atomic-util | Atomic utilities | Conditional | + +**Total: 12 required + 9 optional + 2 conditional** + +burn-flex has significantly fewer dependencies, with no dependency on ndarray itself, no macro +utility crates, and no BLAS bindings. + +--- + +## 13. Codebase Size + +| Metric | burn-flex | burn-ndarray | +| -------------- | ------------- | ------------ | +| Source files | 38 | 37 | +| Total lines | ~23,500 | ~11,400 | +| ops/ directory | ~19,700 lines | ~8,200 lines | +| SIMD module | ~1,200 lines | ~2,100 lines | + +burn-flex has roughly 2x the code. This is because: + +1. Flex implements all ops from scratch (ndarray delegates to the ndarray crate's built-in ops) +2. Flex has dedicated optimized implementations (pool, conv, reduce, cumulative, gather/scatter) +3. Flex has more comprehensive dtype handling (f16/bf16 paths for every op) +4. Flex has explicit contiguous/non-contiguous fast paths throughout + +--- + +## 14. Testing + +| Aspect | burn-flex | burn-ndarray | +| --------------------- | ---------------------------------------------------------- | ----------------------- | +| burn-backend-tests | **All pass** (6 feature flag combos) | All pass | +| burn-no-std-tests | **Pass** (MNIST inference) | Not explicitly verified | +| ONNX model checks | **All pass** | All pass | +| Real model inference | **ALBERT, MiniLM** | Not documented | +| Feature combos tested | no-default, simd, std, std+simd, std+rayon, std+simd+rayon | Default | +| Edge-case robustness | Integer overflow, rounding, zero-size, invalid params | Standard | +| Embedded builds | thumbv6m, thumbv7m, wasm32 | wasm32 | + +--- + +## 15. Performance Summary + +All benchmarks on Apple M3 Max, default features enabled. + +### Compute Performance + +Genuine algorithmic and library improvements: + +| Category | Flex vs NdArray | Why | +| ---------------- | -------------------- | ----------------------------------------- | +| Binary ops (f32) | **2.4-3.9x faster** | Arc COW avoids allocation; 3x less memory | +| Binary ops (i64) | **1.5-6.4x faster** | Same COW benefits | +| Matmul (square) | **1.1-3.4x faster** | gemm > matrixmultiply | +| Matmul (batched) | **1.8-3.2x faster** | Better batch parallelism | +| Attention | **1.2-2.4x faster** | Flash attention, 2-8.5x lower peak memory | +| Conv2d | **1.2-4.0x faster** | im2col+gemm vs direct | +| Conv1d | **4.3-9.6x faster** | Unified 3D avoids overhead | +| Pooling | **1.2-3.1x faster** | Unified 3D, better parallelism | +| Interpolation | **1.2-3.6x faster** | Direct computation vs intermediates | +| Reductions | **1.6-5.1x faster** | Zero-alloc SIMD single-pass | +| Cumulative | **3.1-97x faster** | Blocked scan, scalar accumulator | +| Gather/scatter | **1.6-9.8x faster** | Direct indexing | +| Unary | **1.1-2.7x faster** | In-place mutation when possible | +| Comparisons | **2.1-3.9x faster** | SIMD + compact u8 output | +| Int cast | **5.0-7.6x faster** | Direct byte reinterpretation | +| Quantize | **1.6x faster** | Fused 2-pass implementation | +| Concatenation | **3.6-16.3x faster** | Direct memcpy vs slice_assign | + +### Structural Improvements + +These reflect changes in how operations are _represented and executed_, not pure compute speedups. +burn-ndarray eagerly materializes data where burn-flex uses zero-copy views or separated storage. + +| Category | Improvement | What changed | +| ------------- | ------------------ | ------------------------------------------------------------ | +| Dequantize | **135-232x** | Direct `scale * x_q` vs reparsing `QuantizedBytes` each call | +| Quantized ops | **2.9-117x** | Dominated by fast dequantize path | +| Slice/narrow | **2.1-2,100x** | Zero-copy strided view vs potential data copy | +| Unfold | **1,200-166,000x** | O(1) strided view vs O(n) full materialization | +| Expand | **550-2,600x** | Zero-copy broadcast (stride=0) vs data copy | + +> **Note on quantization**: burn-ndarray simulates quantization by dequantizing to f32 for most +> operations. The quantized speedups reflect the difference between simulated and native execution, +> not equivalent algorithms running at different speeds. + +### Where NdArray Wins + +| Category | NdArray advantage | Reason | +| -------------------------- | ----------------- | ------------------------------------------- | +| bool_not/bool_and | ~20% faster | ndarray's vectorized mapv is well-optimized | +| int_powf_scalar | ~10% faster | ndarray's vectorized internals | +| Transposed i64 add (large) | ~7% faster | ndarray handles non-contiguous well | +| Deform conv (medium) | ~30% faster | NdArray has optimized deform conv path | +| Max pool 5x5 | ~17% faster | Specific kernel size advantage | + +These are specific edge cases where NdArray's ndarray-based internals have an advantage. + +--- + +## 16. Why Replace burn-ndarray? + +The [ndarray](https://crates.io/crates/ndarray) crate has been slow to accept contributions and +evolve. Burn's CPU backend inherits these constraints: + +- **No f16/bf16**: Models using half-precision weights must convert to f32. An f16 PR has been open + for a long time with no clear timeline. +- **6-dimension limit**: Hard-coded in reshape macros, cannot be fixed without upstream changes. +- **Unsigned strides**: `usize`-only strides make zero-copy flip impossible. +- **Simulated quantization**: No native quantized storage; dequantize/requantize on every op. +- **COW limitations**: `NdArrayStorage::Borrowed` always returns false for `is_unique()`, preventing + in-place mutation of externally loaded data. + +burn-flex was built to address these gaps without waiting on upstream. It is not intended to compete +with CubeCL CPU, which targets high-performance computation through operator fusion and just-in-time +compilation via LLVM. The goal is to provide a lightweight, portable replacement for burn-ndarray +that works today on platforms CubeCL CPU cannot target (no_std, WASM, embedded). + +## 17. What burn-flex Adds + +1. **f16/bf16 support**: Native arithmetic on half-precision types. Enables running models that use + f16 weights without conversion. + +2. **No dimension limit**: Arbitrary tensor rank (ndarray is limited to 6). + +3. **Zero-copy flip/unfold/expand**: Signed strides enable O(1) flip. Unfold returns a strided view + instead of materializing all windows. + +4. **Unified 3D conv/pool**: Single implementation covers 1D/2D/3D, reducing code paths and + potential for inconsistencies. + +5. **Native quantization**: Stores scales separately for direct `scale * x_q` dequantization instead + of reparsing packed bytes on every access. Zero-copy layout ops on quantized tensors. + +6. **Fewer dependencies**: 7 required deps vs 12. No ndarray, no matrixmultiply, no paste, no + const-random, no BLAS bindings. + +7. **Simpler type system**: `Flex` vs `NdArray`. No generic parameters, no element trait + hierarchy (`FloatNdArrayElement`, `IntNdArrayElement`, `NdArrayElement`, `ExpElement`). + +8. **Real FFT**: Forward (rfft) and inverse (irfft) real FFT with complex packing, SIMD + butterflies, and compile-time twiddle tables. Works in `no_std` (rustfft/realfft require std). + NdArray implements neither. + +--- + +## 18. What burn-ndarray Has That burn-flex Does Not + +1. **BLAS acceleration**: Feature flags for Accelerate (macOS), OpenBLAS, and Netlib BLAS. These can + outperform gemm for very large matmuls on specific hardware. burn-flex relies solely on the gemm + crate. + +2. **SIMD conv/pool kernels**: burn-ndarray has dedicated macerator-based SIMD kernels for + convolution and pooling. burn-flex delegates to gemm's SIMD. + +3. **export_tests feature**: burn-ndarray serves as a reference implementation for some burn-cubecl + kernels via `export_tests`. + +--- + +## 19. Migration Path + +For Burn users switching from burn-ndarray to burn-flex: + +| Change | Details | +| -------------- | -------------------------------------------------- | +| Type parameter | `NdArray` becomes `Flex` | +| Device | `NdArrayDevice::Cpu` becomes `FlexDevice` | +| Feature flags | `multi-threads` becomes `rayon` | +| BLAS features | No equivalent (gemm handles matmul) | +| Autodiff | Use `burn_autodiff::Autodiff` (same pattern) | +| f16/bf16 | Works out of the box (new capability) | +| Quantization | Same API, faster execution | +| Tests | Same burn-backend-tests suite passes | + +--- + +## 20. Conclusion + +burn-flex is a from-scratch replacement for burn-ndarray, motivated by ndarray's lack of f16/bf16 +support, 6-dimension limit, simulated quantization, and slow pace of upstream development. It +implements all required Backend traits (FloatTensorOps, IntTensorOps, BoolTensorOps, QTensorOps, +ModuleOps, ActivationOps, TransactionOps) and passes the same test suite. + +Performance gains come in two forms: compute improvements (1.1-9.7x) from better libraries and +algorithms, and structural improvements (up to 166,000x) from representing operations as zero-copy +views instead of eagerly materializing data. Memory usage is significantly reduced through Arc-based +COW and in-place mutation. + +The only capabilities lost are optional BLAS acceleration (replaced by the gemm crate, which is +faster in most benchmarks) and the `export_tests` reference implementation feature. diff --git a/crates/burn-flex/Cargo.toml b/crates/burn-flex/Cargo.toml new file mode 100644 index 0000000..117dbfa --- /dev/null +++ b/crates/burn-flex/Cargo.toml @@ -0,0 +1,81 @@ +[package] +authors = ["Dilshod Tadjibaev (@antimora)"] +categories = ["science", "mathematics", "no-std"] +description = "A fast, portable CPU backend for the Burn framework" +documentation = "https://docs.rs/burn-flex" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "tensor", "simd"] +license.workspace = true +name = "burn-flex" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-flex" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["std", "simd", "rayon"] +doc = ["default"] +std = [ + "gemm/std", + "burn-backend/std", + "burn-std/std", + "burn-ir/std", + "half/std", + "num-traits/std", +] +simd = ["dep:macerator", "dep:aligned-vec", "gemm/wasm-simd128-enable"] +rayon = ["dep:rayon", "gemm/rayon"] +# Opt-in gemm codegen paths. Not enabled by default. +# AVX-512 kernels for x86_64 (Sapphire Rapids, Zen 4/5, etc.). +x86-v4 = ["gemm/x86-v4"] +# Apple AMX matrix coprocessor (M-series). Experimental upstream; requires std. +apple-amx = ["std", "gemm/experimental-apple-amx"] +tracing = ["burn-std/tracing", "burn-backend/tracing", "burn-ir/tracing"] +critical-section = [ + "dep:once_cell", + "once_cell/critical-section", + "dep:portable-atomic", + "portable-atomic/critical-section", + "dep:portable-atomic-util", +] + +[dependencies] + +# ** Please make sure all dependencies support no_std when std is disabled ** + +burn-backend = { workspace = true } +burn-ir = { workspace = true } +burn-std = { workspace = true } + +bytemuck = { workspace = true } +half = { workspace = true } +# libm for `erf` (f32/f64). num-traits does not expose erf, and Rust's +# f64/f32 have no builtin erf method even with `std`. libm is already +# in the dep graph transitively via num-traits, so listing it here +# doesn't add a new dependency. +libm = { workspace = true } +num-traits = { workspace = true } + +# Matmul +gemm = { workspace = true, features = ["f16"] } + +# SIMD +aligned-vec = { workspace = true, optional = true } +macerator = { workspace = true, optional = true } + +# Parallel +rayon = { workspace = true, optional = true } + +# no_std targets without atomic CAS (enabled via `critical-section` feature) +once_cell = { workspace = true, optional = true } +portable-atomic = { workspace = true, optional = true } +portable-atomic-util = { workspace = true, optional = true } + +[dev-dependencies] +realfft = { workspace = true } + +[package.metadata.docs.rs] +features = ["default"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-flex/README.md b/crates/burn-flex/README.md new file mode 100644 index 0000000..4c03d9a --- /dev/null +++ b/crates/burn-flex/README.md @@ -0,0 +1,187 @@ +# burn-flex + +A fast, memory-efficient CPU backend for Burn with multi-threading, SIMD, and optimized matrix +multiplication. Runs on std, no_std, and WebAssembly. Supports f16/bf16, zero-copy data loading, and +is thread-safe by design. + +> **[Detailed comparison with burn-ndarray](./COMPARISON.md)**: Full architecture, feature coverage, +> operation-by-operation analysis, and migration path. + +### Features + +- **Zero-Copy Operations**: Many operations return strided views without copying data: + - `transpose`, `permute`, `flip`, `narrow`, `slice` + - `unfold` (sliding windows as strided view instead of materialization) + - `expand` (broadcast via zero strides) +- **Arc-based Copy-on-Write**: O(1) tensor cloning with automatic COW semantics. In-place mutation + when uniquely owned. +- **Convolutions**: Unified 3D implementation with im2col + gemm. Conv1d/2d delegate to conv3d. + Direct conv path for small spatial 1D convs (decomposes into kw gemm calls on NCHW data, skipping + NHWC conversion and im2col). 1x1 pointwise fast path. Supports groups, dilation, padding. +- **Attention**: Auto-selecting between two gemm-backed strategies based on sequence length. Short + sequences (score matrix seq_q \* seq_kv <= 256K elements) use naive attention with two large gemm + calls for lower overhead. Larger shapes use tiled flash attention with online softmax for + O(TILE_KV) memory per row. Both support causal masking, additive bias (ALiBi), softcap, custom + scale, and cross-attention. +- **Pooling**: Max pool, avg pool, adaptive avg pool. All via unified 3D with backward pass support. +- **Conv Transpose**: Scatter-based transposed convolutions for upsampling. +- **Portable SIMD**: Uses [macerator](https://crates.io/crates/macerator) for automatic dispatch: + - aarch64: NEON + - x86_64: AVX2, AVX512, SSE + - wasm32: SIMD128 + - Embedded/other: Scalar fallback +- **Matrix Multiplication**: Optimized via [gemm](https://crates.io/crates/gemm) with native f16 + support. Strided batched matmul passes layout strides directly to gemm, so transposed views (e.g. + `q.matmul(k.swap_dims(-2,-1))`) run at contiguous speed with no copy. +- **FFT**: Forward (rfft) and inverse (irfft) real FFT via Cooley-Tukey with complex packing, + compile-time twiddle tables, SIMD butterflies, and unrolled small kernels. Works in `no_std`. +- **Parallel Execution**: Optional rayon for large tensors +- **Quantization**: Full quantize/dequantize support with per-tensor and per-block symmetric + schemes. All ~40 quantized ops (arithmetic, trig, reductions, sorting, etc.) work out of the box. + Layout ops on quantized tensors (permute, flip, expand, slice, select) are zero-copy. Stores + scales separately for direct `scale * x_q` dequantization instead of reparsing packed bytes. +- **Dtype Support**: f32, f64, f16 (native), bf16 (via f32 conversion), i8-i64, u8-u64 +- **Built on Burn**: Leverages Burn's native infrastructure (`Bytes`, `Shape`, `TensorData`, + `Element` trait) from burn-backend and burn-std + +### Feature Flags + +Default: `std`, `simd`, `rayon`. + +| Flag | Default | Description | +| ------------------- | ------- | ------------------------------------------------------------------- | +| `std` | Yes | Standard library support | +| `simd` | Yes | Portable SIMD via macerator; also enables `gemm/wasm-simd128-enable`| +| `rayon` | Yes | Parallel execution for large tensors (forwards `gemm/rayon`) | +| `x86-v4` | No | AVX-512 kernels in gemm for x86_64 (Sapphire Rapids, Zen 4/5) | +| `apple-amx` | No | Apple Silicon AMX matrix coprocessor in gemm (experimental) | +| `tracing` | No | Propagate `tracing` instrumentation | +| `critical-section` | No | Support for no_std targets without atomic CAS | + +Enable opt-in paths by passing them to Cargo, either directly on `burn-flex` or through the +top-level `burn` crate: + +```toml +# Direct +burn-flex = { version = "0.21", features = ["apple-amx"] } + +# Via burn +burn = { version = "0.21", features = ["flex", "apple-amx"] } +``` + +See [ARCHITECTURE.md#feature-flags](./ARCHITECTURE.md#feature-flags) for per-case benchmark impact. + +### Why replace burn-ndarray? + +burn-ndarray depends on the [ndarray](https://crates.io/crates/ndarray) crate, which has been slow +to accept contributions and evolve. + +burn-flex was built as a from-scratch replacement that addresses the gaps while maintaining full +compatibility with Burn's backend test suite. + +### Performance vs burn-ndarray (Apple M3 Max) + +#### Compute Performance + +burn-ndarray now uses macerator SIMD for f32 elementwise ops +([tracel-ai/burn#2851](https://github.com/tracel-ai/burn/pull/2851)), so contiguous f32 binary/unary +ops are at parity. Flex advantages come from gemm, integer ops (i32 vs i64), structural zero-copy, +and fused kernels. + +| Category | Speedup | Highlights | +| ----------------- | ------------ | --------------------------------------- | +| Binary ops (f32) | **~1x** | Both use macerator SIMD for f32 | +| Binary ops (i32) | **1.8-5.3x** | Flex uses i32, NdArray uses i64 | +| Matmul (square) | **1.4-3.1x** | gemm at small/large; tied at mid-sizes | +| Matmul (batched) | **1.3-2.2x** | Multi-head attention shapes | +| Matmul (int) | **3.7-6.5x** | gemm vs matrixmultiply for integers | +| Conv2d (3x3) | **1.1-3.7x** | Larger kernels and batches benefit most | +| Conv1d | **4.3-9.8x** | | +| Conv transpose | **9.2-84x** | Direct scatter vs im2col | +| Attention | **1.2-3.0x** | Fused softmax, 2-8x lower peak memory | +| Pooling | **1.1-3.1x** | | +| Interpolation | **1.1-6.3x** | Nearest 4-6x, bilinear 1.7-2.8x | +| Reductions | **1.3-5.4x** | Near-zero allocation for scalar results | +| Cumulative ops | **2.1-95x** | 1D cumsum: 95x faster | +| Gather/scatter | **1.2-6.4x** | | +| Unary (tanh, sin) | **1.3-2.0x** | tanh 2x, sin/cos 1.3-1.5x | +| Sort | **2.3-29x** | 2D sort up to 29x | +| Repeat dim | **8.9-12x** | Single alloc + memcpy vs N slice_assign | +| Tensor creation | **16-33x** | zeros/ones/full | +| Embedding | **4.8-5.8x** | | +| Quantize | **1.3-1.5x** | Fused 2-pass implementation | +| FFT (rfft/irfft) | **yes** | Native implementation, works in no_std | + +#### Structural Improvements + +These reflect better _operation representation_, not faster computation. burn-ndarray eagerly +materializes data for these operations; burn-flex avoids the work entirely through zero-copy views +and separated storage layouts. + +| Category | Improvement | What changed | +| ------------- | ---------------- | ------------------------------------------------------------ | +| Dequantize | **122-238x** | Direct `scale * x_q` vs reparsing `QuantizedBytes` each call | +| Quantized ops | **6.1-125x** | Dominated by fast dequantize path above | +| Slice/narrow | **2.1-2,400x** | Zero-copy strided view vs data copy | +| Unfold | **920-130,000x** | O(1) strided view vs O(n) full materialization | +| Expand | **620-2,800x** | Zero-copy broadcast (stride=0) vs data copy | +| Int cast | **6.3-26,000x** | Zero-copy reinterpret vs element-wise conversion | + +> **Note on quantization**: burn-ndarray simulates quantization by dequantizing to f32 for most +> operations. The quantized speedups reflect the difference between simulated and native execution, +> not equivalent algorithms running at different speeds. + +See [BENCHMARKS.md](./BENCHMARKS.md) for the full breakdown. + +### Performance vs candle-core (Apple M3 Max, pure-Rust, no BLAS) + +Per-op comparison against [candle-core](https://github.com/huggingface/candle), pure-Rust on both +sides. Across 11 bench files covering every flex op that intersects with candle's CPU API, flex is +as fast or faster on every operation category. + +| Category | Representative ratio | Notes | +| ----------------------------------- | -------------------- | -------------------------------- | +| Batched matmul | **8-11x** | Strided gemm, no copy | +| Conv1d (wav2vec2) | **1.4-7.6x** | Direct conv path | +| Conv2d (ResNet) | **1.3-4.0x** | 1x1 pointwise 4x | +| Conv transpose | **1.5-1.9x** | | +| Max/min reductions | **3.8-5.1x** | SIMD + zero-alloc | +| Pooling (k=3 s=2) | **1.8-2.5x** | | +| Layer norm (fused) | **1.6-3.4x** | Two-pass Welford kernel | +| Softmax (fused) | **1.4-1.7x** | Three-pass row kernel | +| Cat, gather, select | **1.3-2.5x** | | +| Nearest2d interpolation | **1.3-1.4x** | | +| Elementwise, matmul, gelu, view ops | tied | Both at memory bandwidth ceiling | + +### Status + +- All `burn-backend-tests` pass across all feature flag combinations: + - `no-default-features` (no_std, no SIMD, no rayon) + - `no-default-features + simd` (no_std with SIMD) + - `std` + - `std + simd` + - `std + rayon` + - `std + simd + rayon` (default) +- Burn's `burn-no-std-tests` integration suite passes (MNIST model inference in `#![no_std]`) +- Builds for embedded and WebAssembly targets: + - `thumbv6m-none-eabi` (ARM Cortex-M0+, no atomic pointers) + - `thumbv7m-none-eabi` (ARM Cortex-M3) + - `wasm32-unknown-unknown` +- Passes [Miri](https://github.com/rust-lang/miri) (undefined behavior detector) on all burn-flex + code (quantization tests skipped due to upstream UB in burn-std). Validates memory safety of + unsafe pointer arithmetic, bytemuck casts, and Send/Sync implementations. +- Tested for edge-case robustness: integer overflow at type boundaries, large-float rounding, + invalid pooling parameters, zero-sized dimensions. Safe for embedded devices. +- All ONNX model checks in `burn-onnx` pass +- Real model inference verified: + - [ALBERT](https://huggingface.co/albert/albert-base-v2) (masked language model, all v2 variants) + - [MiniLM](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2) (sentence embeddings, L6 + and L12) + +### Documentation + +- [COMPARISON.md](./COMPARISON.md) - Comprehensive comparison with burn-ndarray +- [ARCHITECTURE.md](./ARCHITECTURE.md) - Design decisions, memory strategy, and implementation + patterns +- [BENCHMARKS.md](./BENCHMARKS.md) - Full benchmark results (Flex vs NdArray) +- [ACKNOWLEDGMENTS.md](./ACKNOWLEDGMENTS.md) - Projects that influenced burn-flex diff --git a/crates/burn-flex/src/backend.rs b/crates/burn-flex/src/backend.rs new file mode 100644 index 0000000..16ebc9f --- /dev/null +++ b/crates/burn-flex/src/backend.rs @@ -0,0 +1,311 @@ +use alloc::string::String; +use burn_std::{BoolStore, DeviceSettings, QuantConfig, QuantScheme, QuantStore}; + +use burn_backend::{Backend, BackendTypes, DType, DTypeUsage, DTypeUsageSet, DeviceId, DeviceOps}; +use burn_ir::{BackendIr, HandleKind, TensorHandle}; +use burn_std::device::Device; +use burn_std::rand::{SeedableRng, StdRng}; +use burn_std::stub::Mutex; + +use crate::qtensor::FlexQTensor; +use crate::tensor::FlexTensor; + +/// Type alias for the RNG used by Flex. +pub type FlexRng = StdRng; + +/// Global seed storage for reproducible random number generation. +/// Uses Mutex for thread-safe RNG state management. +pub(crate) static SEED: Mutex> = Mutex::new(None); + +/// Fallback RNG when `SEED` is empty (consumed or never set). +/// +/// The seeding flow is: `Backend::seed()` stores a `FlexRng` in `SEED`. Random +/// ops (`float_random`, `int_random`) call `SEED.lock().take()`, consuming it for +/// that op and falling back to this function for subsequent calls. This function +/// delegates to burn_std's own entropy source. +pub(crate) fn get_seeded_rng() -> FlexRng { + burn_std::rand::get_seeded_rng() +} + +/// CPU device for the Flex backend. +/// +/// Unit struct since there's only one CPU device. +#[derive(Clone, Copy, Default, PartialEq, Eq, Hash)] +pub struct FlexDevice; + +impl Device for FlexDevice { + fn to_id(&self) -> DeviceId { + DeviceId::new(0, 0) + } + + fn from_id(_id: DeviceId) -> Self { + Self + } +} + +impl DeviceOps for FlexDevice { + fn defaults(&self) -> DeviceSettings { + DeviceSettings::new( + DType::F32, + DType::I32, + DType::Bool(BoolStore::Native), + QuantConfig::new( + QuantScheme::default().with_store(QuantStore::Native), + Default::default(), + ), + ) + } +} + +impl core::fmt::Display for FlexDevice { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + write!(f, "Cpu") + } +} + +impl core::fmt::Debug for FlexDevice { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + core::fmt::Display::fmt(self, f) + } +} + +/// The Flex backend, a fast, portable CPU backend for Burn. +/// +/// The `E` and `I` type parameters exist purely to match the shape of other Burn +/// backends (e.g. `NdArray`) so `Flex` slots into `burn-dispatch`'s +/// generic dispatch macros. The body of `Flex` uses runtime `DType` dispatch, so +/// both parameters are phantom and unused at runtime. +/// +/// # Limitations of the phantom generics +/// +/// The `Backend` impl is provided only for the default instantiation +/// `Flex`. Writing `Flex` (with no arguments) resolves to the default +/// and works exactly as before. Writing `Flex` or any other non-default +/// combination is a valid Rust type but will not satisfy trait bounds requiring +/// `Backend`, producing errors like: +/// +/// ```text +/// the trait bound `Flex: Backend` is not satisfied +/// ``` +/// +/// This is a deliberate compromise for the initial migration: making `Flex` +/// generic over element types at the trait-impl level is a follow-up that would +/// require rewriting all `impl FooOps for Flex` blocks plus internal +/// `Flex::method()` calls (tracked in +/// [#4762](https://github.com/tracel-ai/burn/issues/4762)). Until then, treat +/// the generic parameters as opaque shape placeholders; real element-type +/// selection happens at runtime via `DType`. +/// +/// The bound is locked in by a compile-fail doctest so that if someone later +/// makes the `Backend` impl generic over `E`/`I`, this documentation gets +/// flagged as out of date: +/// +/// ```compile_fail +/// use burn_backend::Backend; +/// use burn_flex::Flex; +/// fn requires_backend() {} +/// requires_backend::>(); +/// ``` +#[derive(Clone, Copy, Debug, Default)] +pub struct Flex {} + +impl BackendTypes for Flex { + type Device = FlexDevice; + + type FloatTensorPrimitive = FlexTensor; + type IntTensorPrimitive = FlexTensor; + type BoolTensorPrimitive = FlexTensor; + type QuantizedTensorPrimitive = FlexQTensor; + + type GraphPrimitive = burn_backend::GraphUnsupported; +} + +impl Backend for Flex { + fn name(_device: &Self::Device) -> String { + "flex".into() + } + + fn seed(_device: &Self::Device, seed: u64) { + let rng = FlexRng::seed_from_u64(seed); + let mut seed_lock = SEED.lock().unwrap(); + *seed_lock = Some(rng); + } + + fn device_count(_type_id: u16) -> usize { + 1 + } + + fn dtype_usage(_device: &Self::Device, dtype: DType) -> DTypeUsageSet { + match dtype { + // Full support for standard types + DType::F64 | DType::F32 | DType::F16 | DType::BF16 => { + DTypeUsage::Storage | DTypeUsage::Arithmetic + } + DType::I64 | DType::I32 | DType::I16 | DType::I8 => { + DTypeUsage::Storage | DTypeUsage::Arithmetic + } + DType::U64 | DType::U32 | DType::U16 | DType::U8 => { + DTypeUsage::Storage | DTypeUsage::Arithmetic + } + // Bool storage: flex stores bools as 1 byte per element, so Native and + // U8 are both supported (they share the same layout, only the tag + // differs). Bool(U32) would require 4-byte-per-element storage + // throughout the backend and is not yet implemented. + DType::Bool(burn_std::BoolStore::Native | burn_std::BoolStore::U8) => { + DTypeUsage::Storage | DTypeUsage::Arithmetic + } + DType::Bool(burn_std::BoolStore::U32) => DTypeUsageSet::empty(), + // Quantized types: storage only for now + DType::QFloat(_) => DTypeUsage::Storage.into(), + _ => DTypeUsageSet::empty(), + } + } + + fn flush(_device: &Self::Device) {} +} + +impl BackendIr for Flex { + type Handle = HandleKind; + + fn float_tensor(handle: TensorHandle) -> FlexTensor { + match handle.handle { + HandleKind::Float(t) => t, + _ => panic!("Expected float handle, got {}", handle.handle.name()), + } + } + + fn int_tensor(handle: TensorHandle) -> FlexTensor { + match handle.handle { + HandleKind::Int(t) => t, + _ => panic!("Expected int handle, got {}", handle.handle.name()), + } + } + + fn bool_tensor(handle: TensorHandle) -> FlexTensor { + match handle.handle { + HandleKind::Bool(t) => t, + _ => panic!("Expected bool handle, got {}", handle.handle.name()), + } + } + + fn quantized_tensor(handle: TensorHandle) -> FlexQTensor { + match handle.handle { + HandleKind::Quantized(t) => t, + _ => panic!("Expected quantized handle, got {}", handle.handle.name()), + } + } + + fn float_tensor_handle(tensor: FlexTensor) -> Self::Handle { + HandleKind::Float(tensor) + } + + fn int_tensor_handle(tensor: FlexTensor) -> Self::Handle { + HandleKind::Int(tensor) + } + + fn bool_tensor_handle(tensor: FlexTensor) -> Self::Handle { + HandleKind::Bool(tensor) + } + + fn quantized_tensor_handle(tensor: FlexQTensor) -> Self::Handle { + HandleKind::Quantized(tensor) + } +} + +// Ops traits are implemented in the ops module + +#[cfg(test)] +mod tests { + use burn_backend::{Backend, DType}; + use burn_std::BoolStore; + + use super::*; + + #[test] + fn supports_bool_native() { + let device = FlexDevice; + assert!(Flex::supports_dtype( + &device, + DType::Bool(BoolStore::Native) + )); + } + + #[test] + fn supports_bool_u8() { + let device = FlexDevice; + assert!(Flex::supports_dtype(&device, DType::Bool(BoolStore::U8))); + } + + #[test] + fn does_not_support_bool_u32() { + let device = FlexDevice; + assert!( + !Flex::supports_dtype(&device, DType::Bool(BoolStore::U32)), + "Bool(U32) should not be supported: flex stores bools as 1 byte per element" + ); + } + + #[test] + fn bool_empty_preserves_native_dtype() { + use burn_backend::ops::BoolTensorOps; + let shape = burn_std::Shape::from(alloc::vec![3]); + let t = Flex::bool_empty(shape, &FlexDevice, burn_std::BoolDType::Native); + assert_eq!(t.dtype(), DType::Bool(BoolStore::Native)); + } + + #[test] + fn bool_empty_preserves_u8_dtype() { + use burn_backend::ops::BoolTensorOps; + let shape = burn_std::Shape::from(alloc::vec![3]); + let t = Flex::bool_empty(shape, &FlexDevice, burn_std::BoolDType::U8); + assert_eq!(t.dtype(), DType::Bool(BoolStore::U8)); + } + + #[test] + fn device_prints_as_cpu() { + use alloc::format; + assert_eq!(format!("{:?}", FlexDevice), "Cpu"); + assert_eq!(format!("{}", FlexDevice), "Cpu"); + } + + #[test] + fn comparison_preserves_out_dtype_native() { + let lhs = FlexTensor::from_data(burn_backend::TensorData::from([1.0f32, 2.0, 3.0])); + let rhs = FlexTensor::from_data(burn_backend::TensorData::from([2.0f32, 2.0, 1.0])); + let result = crate::ops::comparison::greater(lhs, rhs, burn_std::BoolDType::Native); + assert_eq!(result.dtype(), DType::Bool(BoolStore::Native)); + } + + #[test] + fn comparison_preserves_out_dtype_u8() { + let lhs = FlexTensor::from_data(burn_backend::TensorData::from([1.0f32, 2.0, 3.0])); + let rhs = FlexTensor::from_data(burn_backend::TensorData::from([2.0f32, 2.0, 1.0])); + let result = crate::ops::comparison::greater(lhs, rhs, burn_std::BoolDType::U8); + assert_eq!(result.dtype(), DType::Bool(BoolStore::U8)); + } + + #[test] + #[should_panic(expected = "Bool(U32)")] + fn comparison_u32_panics() { + let lhs = FlexTensor::from_data(burn_backend::TensorData::from([1.0f32, 2.0])); + let rhs = FlexTensor::from_data(burn_backend::TensorData::from([2.0f32, 1.0])); + let _ = crate::ops::comparison::greater(lhs, rhs, burn_std::BoolDType::U32); + } + + #[test] + fn bool_not_preserves_u8_dtype() { + use burn_backend::ops::BoolTensorOps; + // Construct a Bool(U8) tensor directly to verify bool_not preserves + // the dtype tag across the op. from_data would produce Bool(Native), + // so we use make_bool_tensor to get the U8 tag. + let t_u8 = crate::ops::comparison::make_bool_tensor( + alloc::vec![1, 0, 1], + burn_std::Shape::from(alloc::vec![3]), + burn_std::BoolDType::U8, + ); + let result = Flex::bool_not(t_u8); + assert_eq!(result.dtype(), DType::Bool(BoolStore::U8)); + let data: &[u8] = result.bytes(); + assert_eq!(&data[..3], &[0, 1, 0]); + } +} diff --git a/crates/burn-flex/src/layout.rs b/crates/burn-flex/src/layout.rs new file mode 100644 index 0000000..cb3a0aa --- /dev/null +++ b/crates/burn-flex/src/layout.rs @@ -0,0 +1,660 @@ +use alloc::vec; +use alloc::vec::Vec; +use burn_std::{Shape, Slice}; + +/// Layout describes how to interpret a linear buffer as an N-dimensional tensor. +/// +/// Stores shape, strides (in elements, can be negative for flipped dimensions), +/// and a start offset for views/slices. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Layout { + shape: Shape, + /// Strides in elements. Negative strides enable zero-copy flip. + strides: Vec, + start_offset: usize, +} + +/// Compute row-major contiguous strides for a shape (as `usize`). +pub(crate) fn contiguous_strides_usize(shape: &Shape) -> Vec { + let ndims = shape.num_dims(); + let mut strides = vec![1usize; ndims]; + for i in (0..ndims.saturating_sub(1)).rev() { + strides[i] = strides[i + 1] * shape[i + 1]; + } + strides +} + +/// Compute the flat offset for the `slice_idx`-th 1D fiber along `dim`. +/// +/// Enumerates all index combinations for dimensions other than `dim`, +/// mapping the flat `slice_idx` (0..product of non-dim sizes) to the +/// corresponding starting offset in a contiguous buffer. +pub(crate) fn slice_base_offset( + slice_idx: usize, + shape: &Shape, + strides: &[usize], + dim: usize, +) -> usize { + let ndims = shape.num_dims(); + let mut offset = 0; + let mut remaining = slice_idx; + for d in (0..ndims).rev() { + if d == dim { + continue; + } + let s = shape[d]; + offset += (remaining % s) * strides[d]; + remaining /= s; + } + offset +} + +impl Layout { + /// Create a new contiguous layout (row-major/C-order). + pub fn contiguous(shape: Shape) -> Self { + let strides: Vec = contiguous_strides_usize(&shape) + .into_iter() + .map(|s| s as isize) + .collect(); + + Self { + shape, + strides, + start_offset: 0, + } + } + + /// Create a layout with explicit strides. + pub fn new(shape: Shape, strides: Vec, start_offset: usize) -> Self { + debug_assert_eq!(shape.num_dims(), strides.len()); + Self { + shape, + strides, + start_offset, + } + } + + /// The shape of the tensor. + pub fn shape(&self) -> &Shape { + &self.shape + } + + /// The strides in elements (can be negative for flipped dimensions). + pub fn strides(&self) -> &[isize] { + &self.strides + } + + /// The start offset for views/slices. + pub fn start_offset(&self) -> usize { + self.start_offset + } + + /// Number of dimensions. + pub fn num_dims(&self) -> usize { + self.shape.num_dims() + } + + /// Total number of elements. + pub fn num_elements(&self) -> usize { + self.shape.num_elements() + } + + /// Check if this layout is contiguous (row-major, positive strides). + pub fn is_contiguous(&self) -> bool { + if self.shape.num_dims() == 0 { + return true; + } + + let mut expected_stride = 1isize; + for i in (0..self.shape.num_dims()).rev() { + if self.strides[i] != expected_stride { + return false; + } + expected_stride *= self.shape[i] as isize; + } + true + } + + /// If contiguous, return (start, end) offsets for direct slice access. + pub fn contiguous_offsets(&self) -> Option<(usize, usize)> { + if self.is_contiguous() { + Some((self.start_offset, self.start_offset + self.num_elements())) + } else { + None + } + } + + /// Transpose: swap two dimensions (zero-copy, metadata only). + pub fn transpose(&self, dim1: usize, dim2: usize) -> Self { + let mut dims = self.shape.to_vec(); + let mut strides = self.strides.clone(); + dims.swap(dim1, dim2); + strides.swap(dim1, dim2); + Self { + shape: Shape::from(dims), + strides, + start_offset: self.start_offset, + } + } + + /// Permute: reorder dimensions according to axes (zero-copy, metadata only). + /// + /// `axes` must be a permutation of 0..ndim. + pub fn permute(&self, axes: &[usize]) -> Self { + debug_assert_eq!( + axes.len(), + self.num_dims(), + "permute: axes length must match number of dimensions" + ); + + let new_dims: Vec = axes.iter().map(|&i| self.shape[i]).collect(); + let new_strides: Vec = axes.iter().map(|&i| self.strides[i]).collect(); + + Self { + shape: Shape::from(new_dims), + strides: new_strides, + start_offset: self.start_offset, + } + } + + /// Flip: reverse elements along specified axes (zero-copy, metadata only). + /// + /// For each flipped axis, negates the stride and adjusts start_offset + /// to point to the last element along that dimension. + pub fn flip(&self, axes: &[usize]) -> Self { + let mut new_strides = self.strides.clone(); + let mut offset_adjustment: isize = 0; + + for &axis in axes { + debug_assert!( + axis < self.num_dims(), + "flip: axis {} out of bounds for {} dimensions", + axis, + self.num_dims() + ); + + let dim_size = self.shape[axis]; + if dim_size > 1 { + // Move start to last element along this axis + offset_adjustment += (dim_size as isize - 1) * self.strides[axis]; + // Negate stride to iterate backwards + new_strides[axis] = -new_strides[axis]; + } + } + + let new_start_isize = self.start_offset as isize + offset_adjustment; + debug_assert!(new_start_isize >= 0, "flip: negative offset"); + let new_start = new_start_isize as usize; + + Self { + shape: self.shape.clone(), + strides: new_strides, + start_offset: new_start, + } + } + + /// Narrow/slice along a dimension (zero-copy, metadata only). + pub fn narrow(&self, dim: usize, start: usize, len: usize) -> Self { + debug_assert!( + start + len <= self.shape[dim], + "narrow: start ({}) + len ({}) exceeds dimension size ({})", + start, + len, + self.shape[dim] + ); + let mut dims = self.shape.to_vec(); + dims[dim] = len; + + let new_offset_isize = self.start_offset as isize + self.strides[dim] * start as isize; + debug_assert!(new_offset_isize >= 0, "narrow: negative offset"); + let new_offset = new_offset_isize as usize; + + Self { + shape: Shape::from(dims), + strides: self.strides.clone(), + start_offset: new_offset, + } + } + + /// Apply slices to create a new layout. + /// + /// Returns `(new_layout, needs_copy)`: + /// - `needs_copy = false`: Can use zero-copy view with new layout + /// - `needs_copy = true`: Has negative steps requiring data reordering + pub fn slice(&self, slices: &[Slice]) -> (Self, bool) { + let ndims = self.num_dims(); + let mut new_dims = self.shape.to_vec(); + let mut new_strides = self.strides.clone(); + let mut new_offset = self.start_offset as isize; + let mut needs_copy = false; + + for (dim, slice) in slices.iter().enumerate() { + if dim >= ndims { + break; + } + + let dim_size = self.shape[dim] as isize; + let stride = self.strides[dim]; + + // Normalize start index (handle negative) + let start = if slice.start < 0 { + (dim_size + slice.start).max(0) as usize + } else { + (slice.start as usize).min(dim_size as usize) + }; + + // Normalize end index (handle negative and None) + // Note: Range [start, end) determines WHICH elements to select, + // step determines iteration ORDER + let end = match slice.end { + Some(e) if e < 0 => (dim_size + e).max(0) as usize, + Some(e) => (e as usize).min(dim_size as usize), + None => dim_size as usize, // Always full range when end is None + }; + + let step = slice.step; + let abs_step = step.unsigned_abs(); + + if step > 0 { + // Positive step: forward iteration + let len = if end > start { + (end - start).div_ceil(abs_step) + } else { + 0 + }; + new_dims[dim] = len; + new_strides[dim] = stride * step; + new_offset += stride * start as isize; + } else { + // Negative step: select range then iterate in reverse + // Requires copy to reorder elements + needs_copy = true; + let len = if end > start { + (end - start).div_ceil(abs_step) + } else { + 0 + }; + new_dims[dim] = len; + new_strides[dim] = stride; // Will be handled during copy + } + } + + debug_assert!(new_offset >= 0, "slice: negative offset"); + + ( + Self { + shape: Shape::from(new_dims), + strides: new_strides, + start_offset: new_offset as usize, + }, + needs_copy, + ) + } + + /// Reshape to a new shape. Only works if contiguous with zero offset. + /// + /// Returns None if not contiguous or has non-zero offset (would require data copy). + pub fn reshape(&self, new_shape: Shape) -> Option { + if !self.is_contiguous() || self.start_offset != 0 { + return None; + } + debug_assert_eq!( + self.num_elements(), + new_shape.num_elements(), + "reshape must preserve total elements" + ); + Some(Self::contiguous(new_shape)) + } + + /// Compute linear index from multi-dimensional indices. + pub fn index(&self, indices: &[usize]) -> usize { + debug_assert_eq!(indices.len(), self.num_dims()); + let mut offset = self.start_offset as isize; + for (i, &idx) in indices.iter().enumerate() { + offset += idx as isize * self.strides[i]; + } + debug_assert!(offset >= 0, "index: negative offset"); + offset as usize + } + + /// Get stride of the innermost (last) dimension. + /// Returns 1 for contiguous tensors, larger values for transposed. + /// Returns absolute value (ignores flip). + pub fn inner_stride(&self) -> usize { + self.strides.last().map(|s| s.unsigned_abs()).unwrap_or(1) + } + + /// Check if innermost dimension is contiguous (|stride| == 1). + /// This enables efficient vectorized inner loops. + pub fn has_contiguous_inner(&self) -> bool { + self.inner_stride() == 1 + } + + /// For 2D layouts, get (outer_size, inner_size, outer_stride, inner_stride). + /// Returns None if not 2D. + pub fn as_2d_strides(&self) -> Option<(usize, usize, isize, isize)> { + if self.num_dims() != 2 { + return None; + } + Some(( + self.shape[0], + self.shape[1], + self.strides[0], + self.strides[1], + )) + } + + /// Check if all strides are non-negative. + pub fn has_positive_strides(&self) -> bool { + self.strides.iter().all(|&s| s >= 0) + } + + /// Compute strided blocks for efficient iteration. + /// + /// Returns (block_len, num_blocks, block_stride) where: + /// - block_len: number of contiguous elements in each block + /// - num_blocks: total number of blocks + /// - block_stride: stride between consecutive blocks (0 if single block) + /// + /// For contiguous tensors: single block covering all elements. + /// For transposed/strided: multiple blocks of contiguous data. + pub fn strided_blocks(&self) -> StridedBlocks<'_> { + let n = self.num_elements(); + if n == 0 { + return StridedBlocks::Single { start: 0, len: 0 }; + } + + // Fast path: fully contiguous + if self.is_contiguous() { + return StridedBlocks::Single { + start: self.start_offset, + len: n, + }; + } + + // Find contiguous inner dimensions (only positive strides) + // Start from innermost and work outward while strides match contiguous pattern + let ndims = self.num_dims(); + let mut block_len = 1usize; + let mut expected_stride = 1isize; + + for i in (0..ndims).rev() { + if self.strides[i] == expected_stride { + block_len *= self.shape[i]; + expected_stride *= self.shape[i] as isize; + } else { + break; + } + } + + if block_len == n { + // All dimensions contiguous (just offset) + return StridedBlocks::Single { + start: self.start_offset, + len: n, + }; + } + + let num_blocks = n / block_len; + StridedBlocks::Multiple { + layout: self, + block_len, + num_blocks, + } + } +} + +/// Result of strided block analysis. +#[derive(Debug, Clone)] +pub enum StridedBlocks<'a> { + /// Single contiguous block - direct slice access. + Single { start: usize, len: usize }, + /// Multiple blocks requiring iteration. + Multiple { + layout: &'a Layout, + block_len: usize, + num_blocks: usize, + }, +} + +impl<'a> StridedBlocks<'a> { + /// Get the block length (elements per block). + pub fn block_len(&self) -> usize { + match self { + Self::Single { len, .. } => *len, + Self::Multiple { block_len, .. } => *block_len, + } + } + + /// Iterator over block start indices. + pub fn block_starts(&self) -> BlockStartIter<'_> { + match self { + Self::Single { start, .. } => BlockStartIter::Single { + start: *start, + done: false, + }, + Self::Multiple { + layout, + block_len, + num_blocks, + } => { + // Calculate dimensions for outer iteration (non-contiguous part) + let ndims = layout.num_dims(); + let mut outer_dims = 0; + let mut expected_stride = 1isize; + + for i in (0..ndims).rev() { + if layout.strides[i] == expected_stride { + expected_stride *= layout.shape[i] as isize; + } else { + outer_dims = i + 1; + break; + } + } + + BlockStartIter::Multiple { + layout, + multi_index: vec![0; outer_dims], + remaining: *num_blocks, + block_len: *block_len, + } + } + } + } +} + +/// Iterator over block start indices. +pub enum BlockStartIter<'a> { + Single { + start: usize, + done: bool, + }, + Multiple { + layout: &'a Layout, + multi_index: Vec, + remaining: usize, + block_len: usize, + }, +} + +impl Iterator for BlockStartIter<'_> { + type Item = usize; + + fn next(&mut self) -> Option { + match self { + Self::Single { start, done } => { + if *done { + None + } else { + *done = true; + Some(*start) + } + } + Self::Multiple { + layout, + multi_index, + remaining, + block_len: _, + } => { + if *remaining == 0 { + return None; + } + + // Compute current block start + let outer_dims = multi_index.len(); + let mut offset = layout.start_offset as isize; + for (i, &idx) in multi_index.iter().enumerate() { + offset += idx as isize * layout.strides[i]; + } + + *remaining -= 1; + + // Advance multi-index for next iteration + let shape = &layout.shape; + for d in (0..outer_dims).rev() { + multi_index[d] += 1; + if multi_index[d] < shape[d] { + break; + } + multi_index[d] = 0; + } + + Some(offset as usize) + } + } + } + + fn size_hint(&self) -> (usize, Option) { + let len = match self { + Self::Single { done, .. } => { + if *done { + 0 + } else { + 1 + } + } + Self::Multiple { remaining, .. } => *remaining, + }; + (len, Some(len)) + } +} + +impl ExactSizeIterator for BlockStartIter<'_> {} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_contiguous_layout() { + let layout = Layout::contiguous(Shape::from(vec![2, 3, 4])); + assert_eq!(layout.strides(), &[12, 4, 1]); + assert!(layout.is_contiguous()); + } + + #[test] + fn test_transpose() { + let layout = Layout::contiguous(Shape::from(vec![2, 3])); + let transposed = layout.transpose(0, 1); + assert_eq!(transposed.shape().to_vec(), vec![3, 2]); + assert_eq!(transposed.strides(), &[1, 3]); + assert!(!transposed.is_contiguous()); + } + + #[test] + fn test_narrow() { + let layout = Layout::contiguous(Shape::from(vec![4, 4])); + let narrowed = layout.narrow(0, 1, 2); + assert_eq!(narrowed.shape().to_vec(), vec![2, 4]); + assert_eq!(narrowed.start_offset(), 4); + } + + #[test] + fn test_contiguous_offsets() { + let layout = Layout::contiguous(Shape::from(vec![2, 3])); + assert_eq!(layout.contiguous_offsets(), Some((0, 6))); + } + + #[test] + fn test_index() { + let layout = Layout::contiguous(Shape::from(vec![2, 3])); + assert_eq!(layout.index(&[0, 0]), 0); + assert_eq!(layout.index(&[0, 2]), 2); + assert_eq!(layout.index(&[1, 0]), 3); + assert_eq!(layout.index(&[1, 2]), 5); + } + + #[test] + fn test_flip_1d() { + // Original: [0, 1, 2, 3] with strides [1] + // Flipped: strides [-1], start_offset = 3 + let layout = Layout::contiguous(Shape::from(vec![4])); + let flipped = layout.flip(&[0]); + + assert_eq!(flipped.shape().to_vec(), vec![4]); + assert_eq!(flipped.strides(), &[-1]); + assert_eq!(flipped.start_offset(), 3); + + // Verify indices: logical [0] -> physical [3], logical [1] -> physical [2], etc. + assert_eq!(flipped.index(&[0]), 3); + assert_eq!(flipped.index(&[1]), 2); + assert_eq!(flipped.index(&[2]), 1); + assert_eq!(flipped.index(&[3]), 0); + } + + #[test] + fn test_flip_2d_axis0() { + // [[0, 1, 2], [3, 4, 5]] with strides [3, 1] + // Flip axis 0: strides [-3, 1], start_offset = 3 + let layout = Layout::contiguous(Shape::from(vec![2, 3])); + let flipped = layout.flip(&[0]); + + assert_eq!(flipped.strides(), &[-3, 1]); + assert_eq!(flipped.start_offset(), 3); + + // Row 0 of flipped = Row 1 of original + assert_eq!(flipped.index(&[0, 0]), 3); + assert_eq!(flipped.index(&[0, 1]), 4); + assert_eq!(flipped.index(&[0, 2]), 5); + // Row 1 of flipped = Row 0 of original + assert_eq!(flipped.index(&[1, 0]), 0); + assert_eq!(flipped.index(&[1, 1]), 1); + assert_eq!(flipped.index(&[1, 2]), 2); + } + + #[test] + fn test_flip_2d_axis1() { + // [[0, 1, 2], [3, 4, 5]] with strides [3, 1] + // Flip axis 1: strides [3, -1], start_offset = 2 + let layout = Layout::contiguous(Shape::from(vec![2, 3])); + let flipped = layout.flip(&[1]); + + assert_eq!(flipped.strides(), &[3, -1]); + assert_eq!(flipped.start_offset(), 2); + + // Col 0 of flipped = Col 2 of original + assert_eq!(flipped.index(&[0, 0]), 2); + assert_eq!(flipped.index(&[0, 1]), 1); + assert_eq!(flipped.index(&[0, 2]), 0); + assert_eq!(flipped.index(&[1, 0]), 5); + assert_eq!(flipped.index(&[1, 1]), 4); + assert_eq!(flipped.index(&[1, 2]), 3); + } + + #[test] + fn test_flip_both_axes() { + // [[0, 1, 2], [3, 4, 5]] -> [[5, 4, 3], [2, 1, 0]] + let layout = Layout::contiguous(Shape::from(vec![2, 3])); + let flipped = layout.flip(&[0, 1]); + + assert_eq!(flipped.strides(), &[-3, -1]); + assert_eq!(flipped.start_offset(), 5); // 3 + 2 = 5 + + assert_eq!(flipped.index(&[0, 0]), 5); + assert_eq!(flipped.index(&[0, 1]), 4); + assert_eq!(flipped.index(&[0, 2]), 3); + assert_eq!(flipped.index(&[1, 0]), 2); + assert_eq!(flipped.index(&[1, 1]), 1); + assert_eq!(flipped.index(&[1, 2]), 0); + } +} diff --git a/crates/burn-flex/src/lib.rs b/crates/burn-flex/src/lib.rs new file mode 100644 index 0000000..a909bd8 --- /dev/null +++ b/crates/burn-flex/src/lib.rs @@ -0,0 +1,47 @@ +#![cfg_attr(not(feature = "std"), no_std)] + +//! # burn-flex +//! +//! A fast, portable CPU backend for [Burn](https://github.com/tracel-ai/burn). +//! +//! ## Features +//! +//! - Pure Rust (no C dependencies) +//! - f16/bf16 support +//! - SIMD acceleration via macerator (NEON, AVX2/AVX-512/SSE, SIMD128, scalar fallback) +//! - Zero-copy tensor views +//! - Thread-safe by design +//! +//! ## Usage +//! +//! ```ignore +//! use burn_flex::Flex; +//! use burn::tensor::Tensor; +//! +//! let tensor: Tensor = Tensor::from_data([[1.0, 2.0], [3.0, 4.0]], &Default::default()); +//! ``` + +extern crate alloc; + +#[cfg(all(not(target_has_atomic = "ptr"), not(feature = "critical-section")))] +compile_error!( + "This target lacks atomic CAS support. Enable the `critical-section` feature: \ + burn-flex = { ..., features = [\"critical-section\"] }" +); + +mod backend; +mod layout; +mod qtensor; +mod strided_index; +mod tensor; + +#[doc(hidden)] +pub mod ops; + +#[doc(hidden)] +pub mod simd; + +pub use backend::{Flex, FlexDevice}; +pub use layout::Layout; +pub use qtensor::FlexQTensor; +pub use tensor::FlexTensor; diff --git a/crates/burn-flex/src/ops/activation.rs b/crates/burn-flex/src/ops/activation.rs new file mode 100644 index 0000000..2dd4e6f --- /dev/null +++ b/crates/burn-flex/src/ops/activation.rs @@ -0,0 +1,1464 @@ +//! Activation function operations for the Flex backend. +//! +//! Each activation is implemented as a single-pass unary operation, +//! replacing the default multi-op compositions from Burn's trait defaults. + +use alloc::vec; +use alloc::vec::Vec; +use burn_backend::Scalar; +use burn_backend::ops::{ActivationOps, FloatTensorOps}; +use burn_backend::tensor::FloatTensor; +use burn_backend::{DType, TensorMetadata}; +use burn_std::{Bytes, bf16, f16}; +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float; +use num_traits::ToPrimitive; + +use crate::ops::binary::binary_op; +use crate::ops::unary::unary_op; +use crate::{Flex, FlexTensor, Layout}; + +impl ActivationOps for Flex { + fn relu(tensor: FloatTensor) -> FloatTensor { + unary_op(tensor, |x: f32| x.max(0.0), |x: f64| x.max(0.0)) + } + + fn relu_backward(output: FloatTensor, grad: FloatTensor) -> FloatTensor { + // grad * (output > 0): zero the gradient where output was zero + binary_op( + output, + grad, + |out: f32, g| if out > 0.0 { g } else { 0.0 }, + |out: f64, g| if out > 0.0 { g } else { 0.0 }, + None, + ) + } + + fn leaky_relu(tensor: FloatTensor, negative_slope: Scalar) -> FloatTensor { + let ns32 = negative_slope.to_f32().unwrap(); + let ns64 = negative_slope.to_f64().unwrap(); + unary_op( + tensor, + move |x: f32| if x >= 0.0 { x } else { ns32 * x }, + move |x: f64| if x >= 0.0 { x } else { ns64 * x }, + ) + } + + fn prelu(tensor: FloatTensor, alpha: FloatTensor) -> FloatTensor { + // x if x >= 0, alpha * x otherwise + binary_op( + tensor, + alpha, + |x: f32, a| if x >= 0.0 { x } else { a * x }, + |x: f64, a| if x >= 0.0 { x } else { a * x }, + None, + ) + } + + fn gelu(tensor: FloatTensor) -> FloatTensor { + // 0.5 * x * (1 + erf(x / sqrt(2))) + use crate::ops::unary::{erf_f32, erf_f64}; + let sqrt2_f32: f32 = core::f32::consts::SQRT_2; + let sqrt2_f64: f64 = core::f64::consts::SQRT_2; + unary_op( + tensor, + move |x: f32| 0.5 * x * (1.0 + erf_f32(x / sqrt2_f32)), + move |x: f64| 0.5 * x * (1.0 + erf_f64(x / sqrt2_f64)), + ) + } + + fn gelu_backward(x: FloatTensor, grad: FloatTensor) -> FloatTensor { + // d/dx[gelu(x)] = 0.5 * (1 + erf(x/sqrt(2))) + x * (1/sqrt(2*pi)) * exp(-x^2/2) + use crate::ops::unary::{erf_f32, erf_f64}; + let sqrt2_f32: f32 = core::f32::consts::SQRT_2; + let sqrt2_f64: f64 = core::f64::consts::SQRT_2; + let inv_sqrt_2pi_f32: f32 = 1.0 / (2.0 * core::f32::consts::PI).sqrt(); + let inv_sqrt_2pi_f64: f64 = 1.0 / (2.0 * core::f64::consts::PI).sqrt(); + binary_op( + x, + grad, + move |x: f32, g| { + let cdf = 0.5 * (1.0 + erf_f32(x / sqrt2_f32)); + let pdf = inv_sqrt_2pi_f32 * (-0.5 * x * x).exp(); + g * (cdf + x * pdf) + }, + move |x: f64, g| { + let cdf = 0.5 * (1.0 + erf_f64(x / sqrt2_f64)); + let pdf = inv_sqrt_2pi_f64 * (-0.5 * x * x).exp(); + g * (cdf + x * pdf) + }, + None, + ) + } + + fn sigmoid(tensor: FloatTensor) -> FloatTensor { + unary_op(tensor, sigmoid_f32, sigmoid_f64) + } + + fn sigmoid_backward(output: FloatTensor, grad: FloatTensor) -> FloatTensor { + // grad * output * (1 - output) + binary_op( + output, + grad, + |s: f32, g| g * s * (1.0 - s), + |s: f64, g| g * s * (1.0 - s), + None, + ) + } + + fn hard_sigmoid(tensor: FloatTensor, alpha: Scalar, beta: Scalar) -> FloatTensor { + let alpha32 = alpha.to_f32().unwrap(); + let beta32 = beta.to_f32().unwrap(); + let alpha64 = alpha.to_f64().unwrap(); + let beta64 = beta.to_f64().unwrap(); + unary_op( + tensor, + move |x: f32| (alpha32 * x + beta32).clamp(0.0, 1.0), + move |x: f64| (alpha64 * x + beta64).clamp(0.0, 1.0), + ) + } + + fn log_sigmoid(tensor: FloatTensor) -> FloatTensor { + // Numerically stable: -softplus(-x) = -log(1 + exp(-x)) + // For x >= 0: -log(1 + exp(-x)) (standard form, exp(-x) is small) + // For x < 0: x - log(1 + exp(x)) (avoids exp of large positive) + unary_op( + tensor, + |x: f32| { + if x >= 0.0 { + -((-x).exp().ln_1p()) + } else { + x - x.exp().ln_1p() + } + }, + |x: f64| { + if x >= 0.0 { + -((-x).exp().ln_1p()) + } else { + x - x.exp().ln_1p() + } + }, + ) + } + + fn log_sigmoid_backward(x: FloatTensor, grad: FloatTensor) -> FloatTensor { + // d/dx[log_sigmoid(x)] = sigmoid(-x) * (-1) * (-1) = 1 - sigmoid(x) = sigmoid(-x) + // So: grad * sigmoid(-x) + binary_op( + x, + grad, + |x: f32, g| g * sigmoid_f32(-x), + |x: f64, g| g * sigmoid_f64(-x), + None, + ) + } + + fn softmax(tensor: FloatTensor, dim: usize) -> FloatTensor { + softmax(tensor, dim) + } +} + +#[inline] +fn sigmoid_f32(x: f32) -> f32 { + if x >= 0.0 { + 1.0 / (1.0 + (-x).exp()) + } else { + let e = x.exp(); + e / (1.0 + e) + } +} + +#[inline] +fn sigmoid_f64(x: f64) -> f64 { + if x >= 0.0 { + 1.0 / (1.0 + (-x).exp()) + } else { + let e = x.exp(); + e / (1.0 + e) + } +} + +// ============================================================================ +// Fused softmax +// ============================================================================ +// +// `ActivationOps` does not currently expose a `softmax` hook, so +// `burn_tensor::activation::softmax` falls back to a 5-op decomposition +// (`max_dim`/`sub`/`exp`/`sum_dim`/`div`). This module provides a fused +// alternative users can opt into directly. + +/// Fused softmax along `dim`. +/// +/// Three-pass row-wise algorithm (max, exp+sum, normalize) keeping each row +/// cache-hot. Rows are processed in parallel via rayon. For axes other than +/// the last, the tensor is permuted to put `dim` last, the fused kernel runs, +/// and the result is permuted back (both permutes are metadata-only; the +/// fused kernel's internal `to_contiguous` materializes the permuted layout +/// once). +/// +/// # Panics +/// +/// * If `dim` is out of range for `input`. +/// * If `input`'s dtype is not one of `f32`/`f64`/`f16`/`bf16`. +pub fn softmax(tensor: FloatTensor, dim: usize) -> FloatTensor { + let rank = tensor.shape().num_dims(); + assert!( + dim < rank, + "softmax dim {} out of range for rank {}", + dim, + rank + ); + + if dim != rank - 1 { + let swapped = Flex::float_swap_dims(tensor, dim, rank - 1); + let normed = softmax_last(swapped); + return Flex::float_swap_dims(normed, dim, rank - 1); + } + + softmax_last(tensor) +} + +fn softmax_last(tensor: FloatTensor) -> FloatTensor { + let tensor = tensor.to_contiguous(); + match tensor.dtype() { + DType::F32 => softmax_last_f32(tensor), + DType::F64 => softmax_last_f64(tensor), + DType::F16 => softmax_last_f16(tensor), + DType::BF16 => softmax_last_bf16(tensor), + dtype => panic!("softmax: unsupported dtype {:?}", dtype), + } +} + +fn softmax_last_f32(tensor: FlexTensor) -> FlexTensor { + let shape = tensor.layout().shape().clone(); + let last = *shape.last().expect("softmax: empty shape"); + if last == 0 { + return tensor; + } + let input: &[f32] = tensor.storage(); + let n = input.len(); + + // Zero-initialize the output. The previous implementation used + // `Vec::with_capacity` + `spare_capacity_mut` + a raw-pointer cast to + // `&mut [f32]` to skip the memset, but forming a `&mut [f32]` over + // uninitialized memory violates Rust's validity invariant (references + // must point to initialized values of the correct type) even if every + // element is written before it is read. The sound zero-memset + // alternative would require threading `&mut [MaybeUninit]` through + // the row kernel, which does not compose with macerator's `#[with_simd]` + // signature. The memset is a streaming write on a bandwidth-bound + // kernel, so the overhead is small (~10% at the largest bench shape) + // and the fused path remains well ahead of decomposed and candle. + let mut output: Vec = vec![0.0; n]; + let out_slice = output.as_mut_slice(); + + // Row-parallel via rayon: one macerator dispatch per chunk of rows, + // amortized over all rows in the chunk. + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + const ROWS_PER_TASK: usize = 64; + let chunk_elems = ROWS_PER_TASK * last; + out_slice + .par_chunks_mut(chunk_elems) + .zip(input.par_chunks(chunk_elems)) + .for_each(|(o, i)| softmax_rows_f32(i, o, last)); + } + #[cfg(not(feature = "rayon"))] + { + softmax_rows_f32(input, out_slice, last); + } + + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(shape), + DType::F32, + ) +} + +/// Row sweep for f32 softmax. With the `simd` feature, delegates to the +/// `#[macerator::with_simd]` SIMD kernel (one dispatch per chunk of rows, +/// amortized over all rows in the chunk). Without `simd`, uses a scalar +/// row kernel. +#[inline] +fn softmax_rows_f32(input: &[f32], output: &mut [f32], row_len: usize) { + // Release-mode invariant checks. These run once per chunk of rows + // (dozens of times per call, not per-element), so the overhead is + // unmeasurable against the kernel work. A debug-only check would + // silently pass a short final chunk to the row kernel on release + // builds if a future refactor broke the row alignment at the call + // site, yielding wrong softmax output with no panic. + assert_eq!(input.len(), output.len()); + assert_eq!(input.len() % row_len, 0); + #[cfg(feature = "simd")] + softmax_rows_f32_simd(input, output, row_len); + #[cfg(not(feature = "simd"))] + { + for (in_row, out_row) in input.chunks(row_len).zip(output.chunks_mut(row_len)) { + softmax_row_f32_scalar(in_row, out_row); + } + } +} + +#[cfg(feature = "simd")] +#[macerator::with_simd] +fn softmax_rows_f32_simd(input: &[f32], output: &mut [f32], row_len: usize) { + debug_assert_eq!(input.len(), output.len()); + debug_assert_eq!(input.len() % row_len, 0); + for (in_row, out_row) in input.chunks(row_len).zip(output.chunks_mut(row_len)) { + softmax_row_f32_simd::(in_row, out_row); + } +} + +/// Scalar fallback row kernel for f32 softmax when the `simd` feature is +/// disabled. Uses the same 3-pass algorithm as the SIMD path; LLVM +/// autovectorizes the max-reduce and normalize loops on most targets. +#[cfg(not(feature = "simd"))] +#[inline] +fn softmax_row_f32_scalar(input: &[f32], output: &mut [f32]) { + let mut max_val = f32::NEG_INFINITY; + for &x in input { + if x > max_val { + max_val = x; + } + } + let mut sum = 0.0f32; + for (i, &x) in input.iter().enumerate() { + let e = (x - max_val).exp(); + output[i] = e; + sum += e; + } + let inv = 1.0f32 / sum; + for x in output.iter_mut() { + *x *= inv; + } +} + +/// Inner row kernel for a single softmax row. `#[inline(always)]` so it +/// inlines into `softmax_rows_f32_simd`'s loop body for each monomorphized S, +/// avoiding a per-row call boundary. +#[cfg(feature = "simd")] +#[inline(always)] +fn softmax_row_f32_simd(input: &[f32], output: &mut [f32]) { + use macerator::{Scalar, vload_unaligned, vstore_unaligned}; + let lanes = ::lanes::(); + let len = input.len(); + let simd_len = len / lanes * lanes; + + // Pass 1: row max for numerical stability. + // SIMD max-reduction across the row, scalar tail. + let (mut max_val, tail_start) = if simd_len >= lanes { + let mut max_vec = unsafe { vload_unaligned::(input.as_ptr()) }; + let mut j = lanes; + while j < simd_len { + let v = unsafe { vload_unaligned::(input.as_ptr().add(j)) }; + max_vec = max_vec.max(v); + j += lanes; + } + (max_vec.reduce_max(), simd_len) + } else { + (f32::NEG_INFINITY, 0) + }; + for &x in &input[tail_start..] { + if x > max_val { + max_val = x; + } + } + + // Pass 2: compute exp(x - max), store in output, accumulate sum. + // Scalar exp (no SIMD exp in macerator). This pass is the one that + // actually does memory reads + writes on the whole row, so scalar + // here still lands us at memory bandwidth. + let mut sum = 0.0f32; + for idx in 0..len { + let e = (input[idx] - max_val).exp(); + output[idx] = e; + sum += e; + } + + // Pass 3: normalize. + // SIMD splat + multiply, scalar tail. + let inv = 1.0f32 / sum; + let inv_vec = inv.splat::(); + let mut i = 0; + while i < simd_len { + unsafe { + let v = vload_unaligned::(output.as_ptr().add(i)); + vstore_unaligned::(output.as_mut_ptr().add(i), v * inv_vec); + } + i += lanes; + } + for x in &mut output[i..] { + *x *= inv; + } +} + +// f64, f16, bf16 softmax share the same row-parallel dispatcher shell and +// differ only in their row kernel (native f64 vs via-f32 for half +// precision). Generated via macros to keep the three variants in lockstep. +// Only f32 has a dedicated SIMD fast path above. + +macro_rules! softmax_last_dtype { + ($fn_name:ident, $T:ty, $zero:expr, $dtype:expr, $row_fn:ident) => { + fn $fn_name(tensor: FlexTensor) -> FlexTensor { + let shape = tensor.layout().shape().clone(); + let last = *shape.last().expect("softmax: empty shape"); + if last == 0 { + return tensor; + } + let input: &[$T] = tensor.storage(); + let mut output: Vec<$T> = vec![$zero; input.len()]; + + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + output + .par_chunks_mut(last) + .zip(input.par_chunks(last)) + .for_each(|(o, i)| $row_fn(i, o)); + } + #[cfg(not(feature = "rayon"))] + { + for (i, o) in input.chunks(last).zip(output.chunks_mut(last)) { + $row_fn(i, o); + } + } + + FlexTensor::new(Bytes::from_elems(output), Layout::contiguous(shape), $dtype) + } + }; +} + +/// Half-precision softmax row kernel. Accumulates in f32 for numerical +/// stability and converts back to the target type at each write. This +/// double-rounds across passes 2 and 3; acceptable for half precision. An +/// f32 scratch buffer would remove the double rounding at the cost of a +/// per-row allocation. +macro_rules! softmax_row_half { + ($fn_name:ident, $T:ty) => { + #[inline] + fn $fn_name(input: &[$T], output: &mut [$T]) { + let mut max_val = f32::NEG_INFINITY; + for &x in input { + let xf = x.to_f32(); + if xf > max_val { + max_val = xf; + } + } + let mut sum = 0.0f32; + for (i, &x) in input.iter().enumerate() { + let e = (x.to_f32() - max_val).exp(); + output[i] = <$T>::from_f32(e); + sum += e; + } + let inv = 1.0f32 / sum; + for x in output.iter_mut() { + *x = <$T>::from_f32(x.to_f32() * inv); + } + } + }; +} + +#[inline] +fn softmax_row_f64(input: &[f64], output: &mut [f64]) { + let mut max_val = f64::NEG_INFINITY; + for &x in input { + if x > max_val { + max_val = x; + } + } + let mut sum = 0.0f64; + for (i, &x) in input.iter().enumerate() { + let e = (x - max_val).exp(); + output[i] = e; + sum += e; + } + let inv = 1.0f64 / sum; + for x in output.iter_mut() { + *x *= inv; + } +} + +softmax_row_half!(softmax_row_f16, f16); +softmax_row_half!(softmax_row_bf16, bf16); + +softmax_last_dtype!(softmax_last_f64, f64, 0.0f64, DType::F64, softmax_row_f64); +softmax_last_dtype!( + softmax_last_f16, + f16, + f16::from_f32(0.0), + DType::F16, + softmax_row_f16 +); +softmax_last_dtype!( + softmax_last_bf16, + bf16, + bf16::from_f32(0.0), + DType::BF16, + softmax_row_bf16 +); + +// ============================================================================ +// Fused layer_norm +// ============================================================================ +// +// `burn::nn::LayerNorm::forward` decomposes into ~6 primitive tensor ops +// with intermediate allocations, and there is no backend trait hook for +// layer_norm. This module provides a fused alternative users can opt into +// directly. Two-pass row kernel (sum+sumsq sweep, then normalize+affine +// sweep), both vectorized via macerator. + +/// Fused layer normalization along the last axis. +/// +/// Applies `y = ((x - mean) / sqrt(var + eps)) * gamma + beta`, where +/// `mean` and `var` are computed per row along the last axis of `input`. +/// `gamma` and `beta` are 1-D tensors of length `input.shape()[-1]`; +/// `beta` is optional (set to `None` for a bias-free layer norm). +/// +/// Two-pass row kernel (mean/variance via a single sum+sum-of-squares +/// sweep, then one normalize+affine sweep). Both passes are SIMD via +/// macerator; each row stays cache-hot across both passes. +/// +/// Supports `f32` (SIMD-vectorized), `f64` (scalar + LLVM autovec), and +/// `f16`/`bf16` (via an f32 cast-fuse-cast shell; the f32 row kernel +/// already accumulates in f32, so this matches the precision a +/// half-precision-native kernel would produce). +/// +/// # Panics +/// +/// * If `input`'s dtype is not one of `f32`/`f64`/`f16`/`bf16`. +/// * If `input` has rank 0. +/// * If `gamma` (or `beta`, when present) is not a 1-D tensor of length +/// equal to the last dim of `input`. +pub fn layer_norm( + input: FloatTensor, + gamma: FloatTensor, + beta: Option>, + epsilon: f64, +) -> FloatTensor { + let rank = input.shape().num_dims(); + assert!(rank >= 1, "layer_norm: input must have at least one dim"); + // Keep gamma/beta dtypes aligned with the input. The half-precision path + // (see `layer_norm_via_f32`) ultimately accesses storage using the input's + // element type, and a mismatch would panic there; reject it up front with + // a clearer layer_norm-specific error message. + assert_eq!( + gamma.dtype(), + input.dtype(), + "layer_norm: gamma dtype {:?} does not match input dtype {:?}", + gamma.dtype(), + input.dtype(), + ); + if let Some(ref b) = beta { + assert_eq!( + b.dtype(), + input.dtype(), + "layer_norm: beta dtype {:?} does not match input dtype {:?}", + b.dtype(), + input.dtype(), + ); + } + let input = input.to_contiguous(); + let gamma = gamma.to_contiguous(); + let beta = beta.map(|b| b.to_contiguous()); + + let d_model = *input + .layout() + .shape() + .last() + .expect("layer_norm: empty shape"); + // Validate rank + length explicitly rather than just last-dim == d_model. + // A gamma shaped like `[2, d_model]` would pass a last-dim check but + // has 2*d_model elements, which would index the wrong data in the row + // kernel (caught by an inner assert, but with a confusing message). + let gamma_shape = gamma.layout().shape(); + assert!( + gamma_shape.len() == 1 && gamma_shape[0] == d_model, + "layer_norm: gamma must be a 1-D tensor of length equal to last dim of input \ + (got shape {:?}, expected [{}])", + gamma_shape, + d_model, + ); + if let Some(ref b) = beta { + let beta_shape = b.layout().shape(); + assert!( + beta_shape.len() == 1 && beta_shape[0] == d_model, + "layer_norm: beta must be a 1-D tensor of length equal to last dim of input \ + (got shape {:?}, expected [{}])", + beta_shape, + d_model, + ); + } + + match input.dtype() { + DType::F32 => layer_norm_f32(input, gamma, beta, epsilon as f32), + DType::F64 => layer_norm_f64(input, gamma, beta, epsilon), + DType::F16 => { + layer_norm_via_f32::(input, gamma, beta, epsilon, f16::to_f32, f16::from_f32) + } + DType::BF16 => { + layer_norm_via_f32::(input, gamma, beta, epsilon, bf16::to_f32, bf16::from_f32) + } + dtype => panic!("burn_flex::layer_norm: unsupported dtype {:?}", dtype), + } +} + +fn layer_norm_via_f32( + input: FlexTensor, + gamma: FlexTensor, + beta: Option, + epsilon: f64, + to_f32: fn(E) -> f32, + from_f32: fn(f32) -> E, +) -> FlexTensor { + let input_f32 = crate::ops::module::cast_to_f32::(input, to_f32); + let gamma_f32 = crate::ops::module::cast_to_f32::(gamma, to_f32); + let beta_f32 = beta.map(|b| crate::ops::module::cast_to_f32::(b, to_f32)); + let out = layer_norm_f32(input_f32, gamma_f32, beta_f32, epsilon as f32); + crate::ops::module::cast_from_f32::(out, from_f32) +} + +/// Fused f64 layer_norm. The Welford mean/variance pass is serial (the +/// mean update on iteration `k` depends on iteration `k-1`); the +/// normalize+affine pass autovectorizes on targets with f64 SIMD. A +/// macerator f64 path can be added if profiling shows it matters. +fn layer_norm_f64( + input: FlexTensor, + gamma: FlexTensor, + beta: Option, + epsilon: f64, +) -> FlexTensor { + let shape = input.layout().shape().clone(); + let d_model = *shape.last().expect("layer_norm: empty shape"); + if d_model == 0 { + return input; + } + let input_data: &[f64] = input.storage(); + let gamma_data: &[f64] = gamma.storage(); + let beta_data: Option<&[f64]> = beta.as_ref().map(|b| b.storage()); + let mut output: Vec = vec![0.0; input_data.len()]; + + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + const ROWS_PER_TASK: usize = 64; + let chunk_elems = ROWS_PER_TASK * d_model; + match beta_data { + Some(beta_slice) => { + output + .par_chunks_mut(chunk_elems) + .zip(input_data.par_chunks(chunk_elems)) + .for_each(|(o, i)| { + layer_norm_rows_f64_with_beta( + i, o, gamma_data, beta_slice, d_model, epsilon, + ); + }); + } + None => { + output + .par_chunks_mut(chunk_elems) + .zip(input_data.par_chunks(chunk_elems)) + .for_each(|(o, i)| { + layer_norm_rows_f64_no_beta(i, o, gamma_data, d_model, epsilon); + }); + } + } + } + #[cfg(not(feature = "rayon"))] + { + match beta_data { + Some(beta_slice) => layer_norm_rows_f64_with_beta( + input_data, + output.as_mut_slice(), + gamma_data, + beta_slice, + d_model, + epsilon, + ), + None => layer_norm_rows_f64_no_beta( + input_data, + output.as_mut_slice(), + gamma_data, + d_model, + epsilon, + ), + } + } + + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(shape), + DType::F64, + ) +} + +#[inline] +fn layer_norm_rows_f64_with_beta( + input: &[f64], + output: &mut [f64], + gamma: &[f64], + beta: &[f64], + d_model: usize, + epsilon: f64, +) { + for (in_row, out_row) in input.chunks(d_model).zip(output.chunks_mut(d_model)) { + let (mean, inv_std) = welford_f64(in_row, epsilon); + for (i, &x) in in_row.iter().enumerate() { + out_row[i] = (x - mean) * (inv_std * gamma[i]) + beta[i]; + } + } +} + +#[inline] +fn layer_norm_rows_f64_no_beta( + input: &[f64], + output: &mut [f64], + gamma: &[f64], + d_model: usize, + epsilon: f64, +) { + for (in_row, out_row) in input.chunks(d_model).zip(output.chunks_mut(d_model)) { + let (mean, inv_std) = welford_f64(in_row, epsilon); + for (i, &x) in in_row.iter().enumerate() { + out_row[i] = (x - mean) * (inv_std * gamma[i]); + } + } +} + +#[inline] +fn welford_f64(row: &[f64], epsilon: f64) -> (f64, f64) { + let mut mean = 0.0f64; + let mut m2 = 0.0f64; + for (k, &x) in row.iter().enumerate() { + let n_k = (k + 1) as f64; + let delta = x - mean; + mean += delta / n_k; + m2 += delta * (x - mean); + } + let var = m2 / row.len() as f64; + (mean, 1.0f64 / (var + epsilon).sqrt()) +} + +fn layer_norm_f32( + input: FlexTensor, + gamma: FlexTensor, + beta: Option, + epsilon: f32, +) -> FlexTensor { + let shape = input.layout().shape().clone(); + let d_model = *shape.last().expect("layer_norm: empty shape"); + if d_model == 0 { + return input; + } + + let input_data: &[f32] = input.storage(); + let gamma_data: &[f32] = gamma.storage(); + let beta_data: Option<&[f32]> = beta.as_ref().map(|b| b.storage()); + + let n = input_data.len(); + // See softmax_last_f32 for the rationale on zero-init instead of + // `spare_capacity_mut` + `&mut [f32]` cast: the latter creates a + // reference to uninitialized f32 values, which is UB under Rust's + // aliasing model even with no intervening read. + let mut output: Vec = vec![0.0; n]; + let out_slice = output.as_mut_slice(); + + // `#[macerator::with_simd]` can't auto-lifetime through + // `Option<&[T]>`, so we dispatch two separate monomorphized + // versions, one with beta and one without. Both call into the + // same shared row kernel. + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + const ROWS_PER_TASK: usize = 64; + let chunk_elems = ROWS_PER_TASK * d_model; + match beta_data { + Some(beta_slice) => { + out_slice + .par_chunks_mut(chunk_elems) + .zip(input_data.par_chunks(chunk_elems)) + .for_each(|(o, i)| { + layer_norm_rows_f32_with_beta( + i, o, gamma_data, beta_slice, d_model, epsilon, + ); + }); + } + None => { + out_slice + .par_chunks_mut(chunk_elems) + .zip(input_data.par_chunks(chunk_elems)) + .for_each(|(o, i)| { + layer_norm_rows_f32_no_beta(i, o, gamma_data, d_model, epsilon); + }); + } + } + } + #[cfg(not(feature = "rayon"))] + { + match beta_data { + Some(beta_slice) => layer_norm_rows_f32_with_beta( + input_data, out_slice, gamma_data, beta_slice, d_model, epsilon, + ), + None => { + layer_norm_rows_f32_no_beta(input_data, out_slice, gamma_data, d_model, epsilon) + } + } + } + + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(shape), + DType::F32, + ) +} + +/// Row sweep for f32 layer_norm with bias. Delegates to the SIMD kernel +/// when the `simd` feature is enabled; otherwise uses a scalar row loop. +#[inline] +fn layer_norm_rows_f32_with_beta( + input: &[f32], + output: &mut [f32], + gamma: &[f32], + beta: &[f32], + d_model: usize, + epsilon: f32, +) { + // Release-mode invariant checks; see softmax_rows_f32 for rationale. + assert_eq!(input.len(), output.len()); + assert_eq!(input.len() % d_model, 0); + assert_eq!(gamma.len(), d_model); + assert_eq!(beta.len(), d_model); + #[cfg(feature = "simd")] + layer_norm_rows_f32_with_beta_simd(input, output, gamma, beta, d_model, epsilon); + #[cfg(not(feature = "simd"))] + { + for (in_row, out_row) in input.chunks(d_model).zip(output.chunks_mut(d_model)) { + layer_norm_row_f32_scalar(in_row, out_row, gamma, Some(beta), epsilon); + } + } +} + +/// Row sweep for f32 layer_norm without bias. +#[inline] +fn layer_norm_rows_f32_no_beta( + input: &[f32], + output: &mut [f32], + gamma: &[f32], + d_model: usize, + epsilon: f32, +) { + // Release-mode invariant checks; see softmax_rows_f32 for rationale. + assert_eq!(input.len(), output.len()); + assert_eq!(input.len() % d_model, 0); + assert_eq!(gamma.len(), d_model); + #[cfg(feature = "simd")] + layer_norm_rows_f32_no_beta_simd(input, output, gamma, d_model, epsilon); + #[cfg(not(feature = "simd"))] + { + for (in_row, out_row) in input.chunks(d_model).zip(output.chunks_mut(d_model)) { + layer_norm_row_f32_scalar(in_row, out_row, gamma, None, epsilon); + } + } +} + +/// Scalar fallback row kernel for layer_norm when the `simd` feature is +/// disabled. Two-pass algorithm matching the SIMD version (sum+sumsq, +/// then normalize+affine). +#[cfg(not(feature = "simd"))] +#[inline] +fn layer_norm_row_f32_scalar( + input: &[f32], + output: &mut [f32], + gamma: &[f32], + beta: Option<&[f32]>, + epsilon: f32, +) { + // Welford's online algorithm for mean and variance, rather than the + // `sumsq / n - mean * mean` identity the SIMD path uses. The identity + // is vulnerable to catastrophic cancellation when the two terms are + // close in magnitude (large mean relative to variance). Welford's + // single-pass formulation avoids that by tracking the running mean + // and accumulating squared deviations from it. The scalar path is + // the contract used when `simd` is disabled, so we prefer numerical + // stability over bit-for-bit match with the SIMD tree reduction. + let len = input.len(); + let mut mean = 0.0f32; + let mut m2 = 0.0f32; + for (k, &x) in input.iter().enumerate() { + let n_k = (k + 1) as f32; + let delta = x - mean; + mean += delta / n_k; + let delta2 = x - mean; + m2 += delta * delta2; + } + let var = m2 / len as f32; + let inv_std = 1.0f32 / (var + epsilon).sqrt(); + for (i, &x) in input.iter().enumerate() { + let scale = inv_std * gamma[i]; + let normed = (x - mean) * scale; + output[i] = match beta { + Some(b) => normed + b[i], + None => normed, + }; + } +} + +/// SIMD-dispatched row sweep for f32 layer_norm with bias (beta). One +/// macerator dispatch per chunk of rows, amortized over the whole chunk. +#[cfg(feature = "simd")] +#[macerator::with_simd] +fn layer_norm_rows_f32_with_beta_simd( + input: &[f32], + output: &mut [f32], + gamma: &[f32], + beta: &[f32], + d_model: usize, + epsilon: f32, +) { + debug_assert_eq!(input.len(), output.len()); + debug_assert_eq!(input.len() % d_model, 0); + debug_assert_eq!(gamma.len(), d_model); + debug_assert_eq!(beta.len(), d_model); + for (in_row, out_row) in input.chunks(d_model).zip(output.chunks_mut(d_model)) { + layer_norm_row_f32_simd::(in_row, out_row, gamma, Some(beta), epsilon); + } +} + +/// SIMD-dispatched row sweep for f32 layer_norm without bias. +#[cfg(feature = "simd")] +#[macerator::with_simd] +fn layer_norm_rows_f32_no_beta_simd( + input: &[f32], + output: &mut [f32], + gamma: &[f32], + d_model: usize, + epsilon: f32, +) { + debug_assert_eq!(input.len(), output.len()); + debug_assert_eq!(input.len() % d_model, 0); + debug_assert_eq!(gamma.len(), d_model); + for (in_row, out_row) in input.chunks(d_model).zip(output.chunks_mut(d_model)) { + layer_norm_row_f32_simd::(in_row, out_row, gamma, None, epsilon); + } +} + +/// Single-row layer_norm kernel. Two vectorized passes. +#[cfg(feature = "simd")] +#[inline(always)] +fn layer_norm_row_f32_simd( + input: &[f32], + output: &mut [f32], + gamma: &[f32], + beta: Option<&[f32]>, + epsilon: f32, +) { + use macerator::{Scalar, vload_unaligned, vstore_unaligned}; + let lanes = ::lanes::(); + let len = input.len(); + let simd_len = len / lanes * lanes; + + // Pass 1: compute sum and sum-of-squares in one sweep, then derive + // mean and variance. Two independent SIMD accumulators (sum, sumsq) + // expose ILP to the two FMA ports. + let (sum, sumsq) = if simd_len >= lanes { + let mut acc_sum = 0.0f32.splat::(); + let mut acc_sumsq = 0.0f32.splat::(); + let mut i = 0; + while i < simd_len { + unsafe { + let v = vload_unaligned::(input.as_ptr().add(i)); + acc_sum += v; + // acc_sumsq += v * v; Vector::mul_add(self, a, b) = self*a + b, + // so v.mul_add(v, acc_sumsq) = v*v + acc_sumsq. + acc_sumsq = v.mul_add(v, acc_sumsq); + } + i += lanes; + } + let mut s = acc_sum.reduce_add(); + let mut sq = acc_sumsq.reduce_add(); + for &x in &input[simd_len..] { + s += x; + sq += x * x; + } + (s, sq) + } else { + let mut s = 0.0f32; + let mut sq = 0.0f32; + for &x in input { + s += x; + sq += x * x; + } + (s, sq) + }; + + let n = len as f32; + let mean = sum / n; + // Biased variance: E[x^2] - E[x]^2. Matches burn::nn::LayerNorm which + // uses var_mean_bias (the biased estimator) rather than Bessel's + // correction. + let var = (sumsq / n) - mean * mean; + let inv_std = 1.0f32 / (var + epsilon).sqrt(); + + // Pass 2: normalize and affine transform. + // out[i] = (x[i] - mean) * inv_std * gamma[i] + beta[i] + // mean_vec and inv_std_vec are hoisted outside the loop (one splat + // each per row). gamma and beta are read once per element; both + // fit in L1 and are shared across all rows within a rayon chunk. + let mean_vec = mean.splat::(); + let inv_std_vec = inv_std.splat::(); + let mut i = 0; + while i < simd_len { + unsafe { + let x = vload_unaligned::(input.as_ptr().add(i)); + let g = vload_unaligned::(gamma.as_ptr().add(i)); + // scale = inv_std * g + let scale = inv_std_vec * g; + // centered = x - mean + let centered = x - mean_vec; + // out = centered * scale (+ beta if present) + let normed = centered * scale; + let out = if let Some(b) = beta { + let b_vec = vload_unaligned::(b.as_ptr().add(i)); + normed + b_vec + } else { + normed + }; + vstore_unaligned::(output.as_mut_ptr().add(i), out); + } + i += lanes; + } + // Scalar tail + while i < len { + let centered = input[i] - mean; + let normed = centered * inv_std * gamma[i]; + output[i] = match beta { + Some(b) => normed + b[i], + None => normed, + }; + i += 1; + } +} + +// Tests kept here exercise flex-specific behavior: SIMD boundaries, rayon +// chunk boundaries, non-contiguous input handling, the flex-internal +// layer_norm op (no public API yet), and dtype-specific fused softmax +// paths (f16/bf16/f64). Plain activation/softmax smoke tests have been +// migrated to burn-backend-tests so they cover every backend. When adding +// new tests, keep them here only if they probe flex internals; otherwise +// add them to crates/burn-backend-tests/tests/tensor/float/activation/. +#[cfg(test)] +mod tests { + use alloc::vec; + use burn_backend::{DType, TensorData, TensorMetadata, Tolerance}; + use burn_std::{bf16, f16}; + use num_traits::Float; + + use crate::FlexTensor; + + // ============================================================================ + // Reference implementations (per-row, last-axis). + // + // These mirror the contract the fused kernel commits to: stable softmax via + // (x - max), layer_norm via (x - mean) * inv(sqrt(var + eps)) with optional + // affine. Written in plain Rust over f32/f64 slices so the tests avoid any + // tensor-library dependency. + // ============================================================================ + + fn softmax_row(row_in: &[T], row_out: &mut [T]) { + let max = row_in + .iter() + .copied() + .fold(T::neg_infinity(), |a, b| if a > b { a } else { b }); + let mut sum = T::zero(); + for (i, &x) in row_in.iter().enumerate() { + let e = (x - max).exp(); + row_out[i] = e; + sum = sum + e; + } + for v in row_out.iter_mut() { + *v = *v / sum; + } + } + + fn softmax_last_ref(data: &[T], row_len: usize) -> Vec { + let mut out = vec![T::zero(); data.len()]; + for (i, o) in data.chunks(row_len).zip(out.chunks_mut(row_len)) { + softmax_row(i, o); + } + out + } + + fn layer_norm_row( + row_in: &[T], + gamma: &[T], + beta: Option<&[T]>, + eps: T, + row_out: &mut [T], + ) { + let n = T::from(row_in.len()).unwrap(); + let mean = row_in.iter().copied().fold(T::zero(), |a, b| a + b) / n; + let var = row_in + .iter() + .map(|&x| (x - mean) * (x - mean)) + .fold(T::zero(), |a, b| a + b) + / n; + let inv_std = T::one() / (var + eps).sqrt(); + for (i, &x) in row_in.iter().enumerate() { + let normed = (x - mean) * inv_std; + let scaled = normed * gamma[i]; + row_out[i] = match beta { + Some(b) => scaled + b[i], + None => scaled, + }; + } + } + + fn layer_norm_last_ref( + data: &[T], + gamma: &[T], + beta: Option<&[T]>, + eps: T, + row_len: usize, + ) -> Vec { + let mut out = vec![T::zero(); data.len()]; + for (i, o) in data.chunks(row_len).zip(out.chunks_mut(row_len)) { + layer_norm_row(i, gamma, beta, eps, o); + } + out + } + + // ============================================================================ + // Helpers: FlexTensor constructors for typed inputs. + // ============================================================================ + + fn flex_f32(data: Vec, shape: &[usize]) -> FlexTensor { + FlexTensor::from_data(TensorData::new(data, shape.to_vec())) + } + + fn flex_f64(data: Vec, shape: &[usize]) -> FlexTensor { + FlexTensor::from_data(TensorData::new(data, shape.to_vec())) + } + + fn flex_half(data: Vec, shape: &[usize]) -> FlexTensor { + FlexTensor::from_data(TensorData::new(data, shape.to_vec())) + } + + // ============================================================================ + // layer_norm tests + // ============================================================================ + + #[test] + fn test_layer_norm_2d_with_beta() { + let t = flex_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], &[2, 4]); + let gamma = flex_f32(vec![1.0; 4], &[4]); + let beta = flex_f32(vec![0.0; 4], &[4]); + let out = crate::ops::activation::layer_norm(t, gamma, Some(beta), 1e-5); + + let expected: Vec = vec![ + -1.3416408, -0.4472136, 0.4472136, 1.3416408, -1.3416408, -0.4472136, 0.4472136, + 1.3416408, + ]; + out.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![2, 4]), + Tolerance::absolute(1e-4), + ); + } + + #[test] + fn test_layer_norm_with_affine() { + let t = flex_f32(vec![1.0, 2.0, 3.0, 4.0], &[1, 4]); + let gamma = flex_f32(vec![2.0, 0.5, 1.0, 3.0], &[4]); + let beta = flex_f32(vec![1.0, -1.0, 0.0, 2.0], &[4]); + let out = crate::ops::activation::layer_norm(t, gamma, Some(beta), 1e-5); + + // normalized = [-1.3416, -0.4472, 0.4472, 1.3416] + // affine: [-1.6833, -1.2236, 0.4472, 6.0249] + out.into_data().assert_approx_eq::( + &TensorData::new(vec![-1.6833, -1.2236, 0.4472, 6.0249], vec![1, 4]), + Tolerance::absolute(1e-3), + ); + } + + #[test] + fn test_layer_norm_no_beta() { + let t = flex_f32(vec![1.0, 2.0, 3.0, 4.0], &[1, 4]); + let gamma = flex_f32(vec![1.0; 4], &[4]); + let out = crate::ops::activation::layer_norm(t, gamma, None, 1e-5); + + out.into_data().assert_approx_eq::( + &TensorData::new( + vec![-1.3416408, -0.4472136, 0.4472136, 1.3416408], + vec![1, 4], + ), + Tolerance::absolute(1e-4), + ); + } + + // ============================================================================ + // softmax SIMD / rayon boundary tests + // ============================================================================ + + #[test] + fn test_softmax_simd_body_row() { + // Row length 32 ensures the SIMD body runs on every supported target: + // NEON (lanes=4), AVX2 (lanes=8), AVX-512 (lanes=16), SIMD128 (lanes=4). + let data: Vec = (0..32).map(|i| i as f32 * 0.1).collect(); + let expected = softmax_last_ref(&data, 32); + let fused = crate::ops::activation::softmax(flex_f32(data, &[1, 32]), 1); + fused.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![1, 32]), + Tolerance::absolute(1e-5), + ); + } + + #[test] + fn test_softmax_multi_chunk_rayon() { + // 100 rows > ROWS_PER_TASK (64) triggers the rayon par_chunks path. + let data: Vec = (0..100 * 16).map(|i| ((i % 17) as f32) * 0.05).collect(); + let expected = softmax_last_ref(&data, 16); + let fused = crate::ops::activation::softmax(flex_f32(data, &[100, 16]), 1); + fused.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![100, 16]), + Tolerance::absolute(1e-5), + ); + } + + #[test] + fn test_softmax_f64() { + // Exercises softmax_last_dtype! + softmax_row_native f64 path. + let data: Vec = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; + let expected = softmax_last_ref(&data, 4); + let fused = crate::ops::activation::softmax(flex_f64(data, &[2, 4]), 1); + fused.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![2, 4]), + Tolerance::absolute(1e-10), + ); + } + + #[test] + fn test_softmax_f16() { + // Exercises softmax_last_dtype! + softmax_row_half f16 path. + let source: Vec = vec![1.0, 2.0, 3.0, 4.0, 0.5, 0.5, 0.5, 0.5]; + let data: Vec = source.iter().map(|&x| f16::from_f32(x)).collect(); + let expected = softmax_last_ref(&data, 4); + let fused = crate::ops::activation::softmax(flex_half(data, &[2, 4]), 1); + fused.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![2, 4]), + Tolerance::absolute(1e-2), + ); + } + + #[test] + fn test_softmax_bf16() { + // Exercises softmax_last_dtype! + softmax_row_half bf16 path. + let source: Vec = vec![1.0, 2.0, 3.0, 4.0, 0.5, 0.5, 0.5, 0.5]; + let data: Vec = source.iter().map(|&x| bf16::from_f32(x)).collect(); + let expected = softmax_last_ref(&data, 4); + let fused = crate::ops::activation::softmax(flex_half(data, &[2, 4]), 1); + fused.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![2, 4]), + Tolerance::absolute(5e-2), + ); + } + + #[test] + fn test_layer_norm_multi_chunk_rayon() { + // 128 rows > ROWS_PER_TASK (64) triggers the rayon path. + let data: Vec = (0..128 * 16).map(|i| ((i % 19) as f32) * 0.03).collect(); + let gamma_data: Vec = vec![1.0; 16]; + let beta_data: Vec = vec![0.0; 16]; + let expected = layer_norm_last_ref(&data, &gamma_data, Some(&beta_data), 1e-5f32, 16); + let fused = crate::ops::activation::layer_norm( + flex_f32(data, &[128, 16]), + flex_f32(gamma_data, &[16]), + Some(flex_f32(beta_data, &[16])), + 1e-5, + ); + fused.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![128, 16]), + Tolerance::absolute(1e-4), + ); + } + + #[test] + fn test_softmax_empty_last_dim_returns_input() { + // shape [2, 0]: empty last dim should round-trip unchanged instead + // of producing NaN via 0/0. + let t = flex_f32(Vec::::new(), &[2, 0]); + let result = crate::ops::activation::softmax(t, 1); + assert_eq!(result.shape().as_slice(), &[2, 0]); + } + + #[test] + fn test_layer_norm_empty_last_dim_returns_input() { + let t = flex_f32(Vec::::new(), &[3, 0]); + let gamma = flex_f32(Vec::::new(), &[0]); + let beta = flex_f32(Vec::::new(), &[0]); + let result = crate::ops::activation::layer_norm(t, gamma, Some(beta), 1e-5); + assert_eq!(result.shape().as_slice(), &[3, 0]); + } + + #[test] + #[should_panic(expected = "gamma must be a 1-D tensor")] + fn test_layer_norm_gamma_length_mismatch_panics() { + let t = flex_f32(vec![1.0, 2.0, 3.0, 4.0], &[1, 4]); + let gamma = flex_f32(vec![1.0, 1.0, 1.0], &[3]); + let _ = crate::ops::activation::layer_norm(t, gamma, None, 1e-5); + } + + #[test] + #[should_panic(expected = "beta must be a 1-D tensor")] + fn test_layer_norm_beta_length_mismatch_panics() { + let t = flex_f32(vec![1.0, 2.0, 3.0, 4.0], &[1, 4]); + let gamma = flex_f32(vec![1.0, 1.0, 1.0, 1.0], &[4]); + let beta = flex_f32(vec![0.0, 0.0, 0.0], &[3]); + let _ = crate::ops::activation::layer_norm(t, gamma, Some(beta), 1e-5); + } + + #[test] + #[should_panic(expected = "gamma must be a 1-D tensor")] + fn test_layer_norm_gamma_rank_mismatch_panics() { + // gamma [2, 4] has matching last-dim but rank 2, so the old last-dim + // check alone would have accepted it and then indexed wrong storage. + let t = flex_f32(vec![1.0, 2.0, 3.0, 4.0], &[1, 4]); + let gamma = flex_f32(vec![1.0; 8], &[2, 4]); + let _ = crate::ops::activation::layer_norm(t, gamma, None, 1e-5); + } + + // Row length 17 leaves exactly one scalar-tail element after every common + // SIMD width (NEON/SSE f32x4: body=16, tail=1; AVX2 f32x8: body=16, tail=1; + // AVX-512 f32x16: body=16, tail=1). Row lengths that divide evenly by the + // SIMD width skip the tail branch entirely, so a bug in the scalar tail + // kernel would sail past CI without a test like this. + #[test] + fn test_softmax_simd_body_plus_scalar_tail() { + let data: Vec = (0..34).map(|i| (i as f32 * 0.137) - 2.3).collect(); + let expected = softmax_last_ref(&data, 17); + let fused = crate::ops::activation::softmax(flex_f32(data, &[2, 17]), 1); + fused.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![2, 17]), + Tolerance::absolute(1e-5), + ); + } + + #[test] + fn test_layer_norm_simd_body_plus_scalar_tail() { + let data: Vec = (0..34).map(|i| (i as f32 * 0.137) - 2.3).collect(); + let gamma_data: Vec = (0..17).map(|i| 1.0 + i as f32 * 0.05).collect(); + let beta_data: Vec = (0..17).map(|i| i as f32 * 0.01).collect(); + let expected = layer_norm_last_ref(&data, &gamma_data, Some(&beta_data), 1e-5f32, 17); + let fused = crate::ops::activation::layer_norm( + flex_f32(data, &[2, 17]), + flex_f32(gamma_data, &[17]), + Some(flex_f32(beta_data, &[17])), + 1e-5, + ); + fused.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![2, 17]), + Tolerance::absolute(1e-5), + ); + } + + #[test] + fn test_layer_norm_f64_with_beta_multi_chunk() { + // 80 rows > ROWS_PER_TASK (64) exercises the rayon multi-chunk f64 path. + let d_model = 16; + let n_rows = 80; + let data: Vec = (0..n_rows * d_model) + .map(|i| ((i % 13) as f64) * 0.07 - 0.3) + .collect(); + let gamma_data: Vec = vec![0.9; d_model]; + let beta_data: Vec = vec![0.05; d_model]; + let eps = 1e-5f64; + let expected = layer_norm_last_ref(&data, &gamma_data, Some(&beta_data), eps, d_model); + let fused = crate::ops::activation::layer_norm( + flex_f64(data, &[n_rows, d_model]), + flex_f64(gamma_data, &[d_model]), + Some(flex_f64(beta_data, &[d_model])), + eps, + ); + fused.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![n_rows, d_model]), + Tolerance::absolute(1e-10), + ); + } + + #[test] + fn test_layer_norm_f64_no_beta() { + let data: Vec = vec![1.0, 2.0, 3.0, 4.0, -1.0, 0.5, 1.5, -0.5]; + let gamma_data: Vec = vec![1.0; 4]; + let eps = 1e-5f64; + let expected = layer_norm_last_ref(&data, &gamma_data, None, eps, 4); + let fused = crate::ops::activation::layer_norm( + flex_f64(data, &[2, 4]), + flex_f64(gamma_data, &[4]), + None, + eps, + ); + fused.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![2, 4]), + Tolerance::absolute(1e-10), + ); + } + + // Shared body for f16/bf16 layer_norm tests. The fused half-precision + // kernel casts to f32 internally, so the reference is computed in f32 + // and compared back against the half output with an f32 tolerance. + fn check_layer_norm_half_precision(from_f32: fn(f32) -> E, dtype: DType) + where + E: burn_backend::Element + Float, + { + let rows_f32: [f32; 12] = [ + 1.0, 2.0, 3.0, 4.0, -1.0, 0.0, 1.0, 2.0, 0.5, -0.5, 1.5, -1.5, + ]; + let gamma_f32: [f32; 4] = [1.0, 0.5, 1.5, 1.0]; + let beta_f32: [f32; 4] = [0.1, -0.1, 0.0, 0.2]; + let eps = 1e-5f32; + + let expected_f32 = layer_norm_last_ref(&rows_f32, &gamma_f32, Some(&beta_f32), eps, 4); + + let data: Vec = rows_f32.iter().map(|&x| from_f32(x)).collect(); + let gamma_data: Vec = gamma_f32.iter().map(|&x| from_f32(x)).collect(); + let beta_data: Vec = beta_f32.iter().map(|&x| from_f32(x)).collect(); + assert_eq!(E::dtype(), dtype); + + let fused = crate::ops::activation::layer_norm( + flex_half(data, &[3, 4]), + flex_half(gamma_data, &[4]), + Some(flex_half(beta_data, &[4])), + eps as f64, + ); + fused.into_data().assert_approx_eq::( + &TensorData::new(expected_f32, vec![3, 4]), + Tolerance::absolute(3e-2), + ); + } + + #[test] + fn test_layer_norm_f16_via_f32_cast() { + check_layer_norm_half_precision::(f16::from_f32, DType::F16); + } + + #[test] + fn test_layer_norm_bf16_via_f32_cast() { + check_layer_norm_half_precision::(bf16::from_f32, DType::BF16); + } + + #[test] + #[should_panic(expected = "softmax dim")] + fn test_softmax_dim_out_of_range_panics() { + let t = flex_f32(vec![1.0, 2.0, 3.0, 4.0], &[2, 2]); + let _ = crate::ops::activation::softmax(t, 2); + } + + #[test] + #[should_panic(expected = "gamma dtype")] + fn test_layer_norm_gamma_dtype_mismatch_panics() { + // Input f32, gamma f64: layer_norm rejects the mismatch up front + // rather than panicking later inside the storage-typed access. + let t = flex_f32(vec![1.0, 2.0, 3.0, 4.0], &[1, 4]); + let gamma = flex_f64(vec![1.0; 4], &[4]); + let _ = crate::ops::activation::layer_norm(t, gamma, None, 1e-5); + } +} diff --git a/crates/burn-flex/src/ops/attention.rs b/crates/burn-flex/src/ops/attention.rs new file mode 100644 index 0000000..48a9c00 --- /dev/null +++ b/crates/burn-flex/src/ops/attention.rs @@ -0,0 +1,1488 @@ +//! Attention (scaled dot-product) for CPU. +//! +//! Computes: softmax(Q @ K^T * scale + bias) @ V +//! +//! Two strategies, auto-selected by `attention()` based on sequence length: +//! +//! - **Naive** (seq_kv <= 8*TILE_KV): materializes full score matrix, two +//! large gemm calls per (batch, head). Faster for short sequences. +//! - **Flash** (seq_kv > 8*TILE_KV): tiles over KV with online softmax, +//! O(TILE_KV) scratch per row. Better cache behavior for long sequences. + +use alloc::vec; +use alloc::vec::Vec; +use burn_backend::DType; +use burn_backend::ops::AttentionModuleOptions; +use burn_std::Bytes; +use bytemuck::Pod; +use num_traits::Float; + +use crate::{FlexTensor, Layout}; + +/// KV tile size for flash attention. +/// +/// Chosen so the score row (TILE_KV * 4 bytes for f32) and the V tile +/// [TILE_KV, val_dim] fit comfortably in L1. With val_dim=128 the V tile +/// is 32KB; this fits well on Apple Silicon (64-128KB L1) and modern x86 +/// (48KB+ L1d since Golden Cove / Zen 4). On older x86 with 32KB L1d the +/// tile saturates L1 but still benefits from L2 residency. +/// WASM targets use a smaller tile to stay within tighter cache budgets. +#[cfg(target_family = "wasm")] +const TILE_KV: usize = 32; +#[cfg(not(target_family = "wasm"))] +const TILE_KV: usize = 64; + +/// Max score matrix size (in elements) for the naive path. +/// +/// Naive attention materializes a [seq_q, seq_kv] score matrix per head. +/// When this exceeds the budget, flash attention is used instead. +/// 256K elements = 1 MB for f32, fits comfortably in L2. +const NAIVE_SCORE_BUDGET: usize = 256 * 1024; + +/// Auto-selecting attention: picks the fastest strategy based on sequence length. +/// +/// Uses naive attention when the score matrix (seq_q * seq_kv) fits within +/// `NAIVE_SCORE_BUDGET`. Falls back to flash attention for larger shapes. +pub fn attention( + query: FlexTensor, + key: FlexTensor, + value: FlexTensor, + mask: Option, + attn_bias: Option, + options: AttentionModuleOptions, +) -> FlexTensor { + debug_assert!( + query.layout().shape().num_dims() == 4, + "attention: query must be 4D, got {}D", + query.layout().shape().num_dims() + ); + debug_assert!( + key.layout().shape().num_dims() == 4, + "attention: key must be 4D, got {}D", + key.layout().shape().num_dims() + ); + let seq_q = query.layout().shape()[2]; + let seq_kv = key.layout().shape()[2]; + if seq_q * seq_kv <= NAIVE_SCORE_BUDGET { + return attention_naive(query, key, value, mask, attn_bias, options); + } + attention_flash(query, key, value, mask, attn_bias, options) +} + +/// Dispatch attention by dtype, casting f16/bf16 to f32 for computation. +macro_rules! dispatch_attention_dtype { + ($query:expr, $key:expr, $value:expr, $mask:expr, $attn_bias:expr, $options:expr, $impl_fn:ident) => {{ + let query = $query; + let key = $key; + let value = $value; + let mask = $mask; + let attn_bias = $attn_bias; + let options = $options; + let dtype = query.dtype(); + debug_assert_eq!(key.dtype(), dtype, "attention: key dtype mismatch"); + debug_assert_eq!(value.dtype(), dtype, "attention: value dtype mismatch"); + if let Some(ref b) = attn_bias { + debug_assert_eq!(b.dtype(), dtype, "attention: attn_bias dtype mismatch"); + } + match dtype { + DType::F32 => $impl_fn::(query, key, value, mask, attn_bias, options), + DType::F64 => $impl_fn::(query, key, value, mask, attn_bias, options), + DType::F16 => { + use burn_std::f16; + let r = $impl_fn::( + cast_to_f32(query, f16::to_f32), + cast_to_f32(key, f16::to_f32), + cast_to_f32(value, f16::to_f32), + mask, + attn_bias.map(|b| cast_to_f32(b, f16::to_f32)), + options, + ); + cast_from_f32(r, f16::from_f32) + } + DType::BF16 => { + use burn_std::bf16; + let r = $impl_fn::( + cast_to_f32(query, bf16::to_f32), + cast_to_f32(key, bf16::to_f32), + cast_to_f32(value, bf16::to_f32), + mask, + attn_bias.map(|b| cast_to_f32(b, bf16::to_f32)), + options, + ); + cast_from_f32(r, bf16::from_f32) + } + dtype => panic!("attention: unsupported dtype {:?}", dtype), + } + }}; +} + +/// Contiguous mask/bias tensor plus the per-batch and per-head element offsets the +/// inner loop should use to locate the `[seq_q, seq_kv]` tile for each `(batch, head)` +/// pair. When a leading dim (batch or heads) is `1` in the source, its step is `0`, so +/// the inner loop re-reads the same tile for every pair without allocating an expanded +/// copy. The tile length itself is always `seq_q * seq_kv` and is computed at the call +/// site, so it is not stored here. +struct BroadcastMaskBias { + tensor: FlexTensor, + batch_step: usize, + head_step: usize, +} + +/// Prepare an attention mask or bias for the inner loop, accepting ONNX Attention-23 +/// broadcast shapes. +/// +/// Stride-0 along the leading `[batch, heads]` dims is handled without materializing, +/// so the common ONNX patterns (`[1, 1, seq_q, seq_kv]`, `[batch, 1, seq_q, seq_kv]`, +/// `[1, heads, seq_q, seq_kv]`) stay zero-copy. That matters especially for the flash +/// path, where materializing an expanded mask/bias would allocate a full +/// `[batch, heads, seq_q, seq_kv]` buffer and negate flash attention's memory +/// efficiency. If the trailing `[seq_q, seq_kv]` dims are themselves broadcast (rare in +/// practice), we fall back to `expand` + `to_contiguous` so the tile stays contiguous +/// in memory for the inner loop's slice-based access. +fn broadcast_attn_mask_bias( + tensor: FlexTensor, + target: [usize; 4], + name: &'static str, +) -> BroadcastMaskBias { + let ndim = tensor.layout().shape().num_dims(); + assert!(ndim == 4, "attention: {name} must be 4D, got {ndim}D"); + let shape = tensor.layout().shape(); + let src = [shape[0], shape[1], shape[2], shape[3]]; + for i in 0..4 { + assert!( + src[i] == target[i] || src[i] == 1, + "attention: {name} dim {i} must be {} or 1, got {}", + target[i], + src[i] + ); + } + + let tile_len = target[2] * target[3]; + + // Broadcast on seq_q or seq_kv: the source's trailing tile has fewer elements + // than `tile_len`, so per-pair slice access would under-read. Materialize via + // expand + to_contiguous in that case. + if src[2] != target[2] || src[3] != target[3] { + let expanded = crate::ops::expand::expand(tensor, burn_std::Shape::new(target)); + return BroadcastMaskBias { + tensor: expanded.to_contiguous(), + batch_step: target[1] * tile_len, + head_step: tile_len, + }; + } + + // Trailing dims match the target. Keep the source at its own shape (size + // `src[0] * src[1] * tile_len`) and zero-out the step for any leading dim of 1. + BroadcastMaskBias { + tensor: tensor.to_contiguous(), + batch_step: if src[0] == 1 { 0 } else { src[1] * tile_len }, + head_step: if src[1] == 1 { 0 } else { tile_len }, + } +} + +/// Flash attention: tiled computation with online softmax. Use directly to bypass auto-selection. +pub fn attention_flash( + query: FlexTensor, + key: FlexTensor, + value: FlexTensor, + mask: Option, + attn_bias: Option, + options: AttentionModuleOptions, +) -> FlexTensor { + dispatch_attention_dtype!(query, key, value, mask, attn_bias, options, attention_impl) +} + +fn cast_to_f32( + tensor: FlexTensor, + to_f32: fn(E) -> f32, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let data: &[E] = tensor.storage(); + let f32_data: Vec = data.iter().map(|&v| to_f32(v)).collect(); + FlexTensor::new( + Bytes::from_elems(f32_data), + Layout::contiguous(shape), + DType::F32, + ) +} + +fn cast_from_f32( + tensor: FlexTensor, + from_f32: fn(f32) -> E, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let data: &[f32] = tensor.storage(); + let half_data: Vec = data.iter().map(|&v| from_f32(v)).collect(); + FlexTensor::new( + Bytes::from_elems(half_data), + Layout::contiguous(shape), + E::dtype(), + ) +} + +/// Flash attention: tiled computation that avoids materializing the full scores matrix. +/// +/// Input shapes (all 4D): +/// query: \[batch, heads, seq_q, head_dim\] +/// key: \[batch, heads, seq_kv, head_dim\] +/// value: \[batch, heads, seq_kv, val_dim\] +/// mask: \[batch, heads, seq_q, seq_kv\] (optional, u8 where nonzero = masked out) +/// attn_bias: \[batch, heads, seq_q, seq_kv\] (optional) +/// +/// Output: \[batch, heads, seq_q, val_dim\] +/// +/// Algorithm per (batch, head): +/// For each KV tile of size TILE_KV: +/// 1. Score matmul: scores\[seq_q, tile_kv\] = Q @ K_tile^T (gemm) +/// 2. Per query row: apply scale, softcap, mask, bias +/// 3. Per query row: online softmax update (running max/sum, rescale accumulator) +/// 4. Value matmul: output += P @ V_tile (gemm) +/// Final: output\[qi\] /= row_sum\[qi\] for each query row +fn attention_impl( + query: FlexTensor, + key: FlexTensor, + value: FlexTensor, + mask: Option, + attn_bias: Option, + options: AttentionModuleOptions, +) -> FlexTensor +where + T: FlashGemm + burn_backend::Element, +{ + if let Some(softcap) = options.softcap { + assert!(softcap > 0.0, "softcap must be positive, got {softcap}"); + } + + let query = query.to_contiguous(); + let key = key.to_contiguous(); + let value = value.to_contiguous(); + + let q_shape = query.layout().shape(); + let k_shape = key.layout().shape(); + let v_shape = value.layout().shape(); + assert!(q_shape.num_dims() == 4, "attention: query must be 4D"); + assert!(k_shape.num_dims() == 4, "attention: key must be 4D"); + assert!(v_shape.num_dims() == 4, "attention: value must be 4D"); + + let batch = q_shape[0]; + let heads = q_shape[1]; + let seq_q = q_shape[2]; + let head_dim = q_shape[3]; + assert!(head_dim > 0, "attention: head_dim must be non-zero"); + + let seq_kv = k_shape[2]; + let val_dim = v_shape[3]; + + assert_eq!(k_shape[0], batch, "attention: key batch mismatch"); + assert_eq!(k_shape[1], heads, "attention: key heads mismatch"); + assert_eq!(k_shape[3], head_dim, "attention: key head_dim mismatch"); + assert_eq!(v_shape[0], batch, "attention: value batch mismatch"); + assert_eq!(v_shape[1], heads, "attention: value heads mismatch"); + assert_eq!(v_shape[2], seq_kv, "attention: value seq_kv mismatch"); + + let target = [batch, heads, seq_q, seq_kv]; + let mask_bcast = mask.map(|m| broadcast_attn_mask_bias(m, target, "mask")); + let bias_bcast = attn_bias.map(|b| broadcast_attn_mask_bias(b, target, "bias")); + + let scale = T::from( + options + .scale + .unwrap_or_else(|| 1.0 / (head_dim as f64).sqrt()), + ) + .unwrap(); + let softcap: Option = options.softcap.map(|s| T::from(s).unwrap()); + let causal_offset = if options.is_causal { + Some(seq_kv as isize - seq_q as isize) + } else { + None + }; + + let q_data: &[T] = query.storage(); + let k_data: &[T] = key.storage(); + let v_data: &[T] = value.storage(); + let mask_data: Option<&[u8]> = mask_bcast.as_ref().map(|b| b.tensor.bytes()); + let bias_data: Option<&[T]> = bias_bcast.as_ref().map(|b| b.tensor.storage()); + let (mask_batch_step, mask_head_step) = mask_bcast + .as_ref() + .map(|b| (b.batch_step, b.head_step)) + .unwrap_or((0, 0)); + let (bias_batch_step, bias_head_step) = bias_bcast + .as_ref() + .map(|b| (b.batch_step, b.head_step)) + .unwrap_or((0, 0)); + + let mut output = vec![T::zero(); batch * heads * seq_q * val_dim]; + + // 4D strides for contiguous layout + let q_head_stride = seq_q * head_dim; + let q_batch_stride = heads * q_head_stride; + let k_head_stride = seq_kv * head_dim; + let k_batch_stride = heads * k_head_stride; + let v_head_stride = seq_kv * val_dim; + let v_batch_stride = heads * v_head_stride; + let o_head_stride = seq_q * val_dim; + let o_batch_stride = heads * o_head_stride; + let mask_tile_len = seq_q * seq_kv; + + let params = AttentionParams { + scale, + softcap, + causal_offset, + seq_q, + seq_kv, + head_dim, + val_dim, + }; + + // Allocate scratch buffers once and reuse across all (batch, head) pairs + let mut scratch = ScratchBuffers { + row_max: vec![T::neg_infinity(); seq_q], + row_sum: vec![T::zero(); seq_q], + scores: vec![T::zero(); seq_q * TILE_KV], + }; + + for b in 0..batch { + for h in 0..heads { + let q_off = b * q_batch_stride + h * q_head_stride; + let k_off = b * k_batch_stride + h * k_head_stride; + let v_off = b * v_batch_stride + h * v_head_stride; + let o_off = b * o_batch_stride + h * o_head_stride; + let mask_off = b * mask_batch_step + h * mask_head_step; + let bias_off = b * bias_batch_step + h * bias_head_step; + + flash_attention_head( + &q_data[q_off..q_off + q_head_stride], + &k_data[k_off..k_off + k_head_stride], + &v_data[v_off..v_off + v_head_stride], + &mut output[o_off..o_off + o_head_stride], + mask_data.map(|m| &m[mask_off..mask_off + mask_tile_len]), + bias_data.map(|b| &b[bias_off..bias_off + mask_tile_len]), + ¶ms, + &mut scratch, + ); + } + } + + let shape = burn_std::Shape::from(vec![batch, heads, seq_q, val_dim]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(shape), + T::dtype(), + ) +} + +/// Gemm dispatch for flash attention block matmuls. +/// +/// Wraps `gemm::gemm` for f32 and f64 so `flash_attention_head` stays generic. +/// Convention: dst = alpha * dst + beta * (lhs @ rhs) +trait FlashGemm: Float + Pod + Copy + core::ops::AddAssign { + /// Block matrix multiply used for score and value matmuls. + /// + /// # Safety + /// All pointers must be valid for the given dimensions and strides. + unsafe fn block_gemm(args: BlockGemmArgs); +} + +/// Arguments for a block matrix multiply: dst = alpha * dst + beta * (lhs @ rhs). +struct BlockGemmArgs { + m: usize, + n: usize, + k: usize, + dst: *mut T, + dst_cs: isize, + dst_rs: isize, + read_dst: bool, + lhs: *const T, + lhs_cs: isize, + lhs_rs: isize, + rhs: *const T, + rhs_cs: isize, + rhs_rs: isize, + alpha: T, + beta: T, +} + +macro_rules! impl_flash_gemm { + ($ty:ty) => { + impl FlashGemm for $ty { + unsafe fn block_gemm(a: BlockGemmArgs) { + unsafe { + gemm::gemm( + a.m, + a.n, + a.k, + a.dst, + a.dst_cs, + a.dst_rs, + a.read_dst, + a.lhs, + a.lhs_cs, + a.lhs_rs, + a.rhs, + a.rhs_cs, + a.rhs_rs, + a.alpha, + a.beta, + false, + false, + false, + gemm::Parallelism::None, + ); + } + } + } + }; +} + +impl_flash_gemm!(f32); +impl_flash_gemm!(f64); + +/// Scratch buffers reused across (batch, head) pairs to avoid per-head allocation. +struct ScratchBuffers { + row_max: Vec, + row_sum: Vec, + scores: Vec, +} + +/// Parameters for a single (batch, head) flash attention computation. +struct AttentionParams { + scale: T, + softcap: Option, + causal_offset: Option, + seq_q: usize, + seq_kv: usize, + head_dim: usize, + val_dim: usize, +} + +#[allow(clippy::too_many_arguments)] +/// Process a single (batch, head) pair with flash attention. +/// +/// Uses gemm for the two block matmuls per tile: +/// 1. Score matmul: scores\[seq_q, tile_kv\] = Q @ K_tile^T +/// 2. Value matmul: output += P @ V_tile (where P = exp(scores - max)) +/// +/// The online softmax (scale, mask, bias, exp, correction) is applied +/// row-by-row between the two gemm calls. +fn flash_attention_head( + q: &[T], + k: &[T], + v: &[T], + output: &mut [T], + mask: Option<&[u8]>, + bias: Option<&[T]>, + p: &AttentionParams, + scratch: &mut ScratchBuffers, +) { + debug_assert_eq!(q.len(), p.seq_q * p.head_dim); + debug_assert_eq!(k.len(), p.seq_kv * p.head_dim); + debug_assert_eq!(v.len(), p.seq_kv * p.val_dim); + debug_assert_eq!(output.len(), p.seq_q * p.val_dim); + + let neg_inf = T::neg_infinity(); + let AttentionParams { + scale, + softcap, + causal_offset, + seq_q, + seq_kv, + head_dim, + val_dim, + } = *p; + + // Reset scratch buffers for this (batch, head) pair + let row_max = &mut scratch.row_max; + row_max.fill(neg_inf); + let row_sum = &mut scratch.row_sum; + row_sum.fill(T::zero()); + let scores = &mut scratch.scores; + + let num_kv_tiles = seq_kv.div_ceil(TILE_KV); + + for tile_idx in 0..num_kv_tiles { + let kv_start = tile_idx * TILE_KV; + let kv_end = (kv_start + TILE_KV).min(seq_kv); + let tile_kv = kv_end - kv_start; + + // Step 1: Score matmul via gemm + // scores[seq_q, tile_kv] = Q[seq_q, head_dim] @ K_tile[tile_kv, head_dim]^T + // + // Q is row-major [seq_q, head_dim]: rs=head_dim, cs=1 + // K_tile^T: treat K[tile_kv, head_dim] row-major with swapped strides + // K row-major: rs=head_dim, cs=1 -> K^T: rs=1, cs=head_dim + // scores row-major [seq_q, tile_kv]: rs=tile_kv, cs=1 + unsafe { + T::block_gemm(BlockGemmArgs { + m: seq_q, + n: tile_kv, + k: head_dim, + dst: scores.as_mut_ptr(), + dst_cs: 1, + dst_rs: tile_kv as isize, + read_dst: false, + lhs: q.as_ptr(), + lhs_cs: 1, + lhs_rs: head_dim as isize, + rhs: k.as_ptr().add(kv_start * head_dim), + rhs_cs: head_dim as isize, + rhs_rs: 1, + alpha: T::zero(), + beta: T::one(), + }); + } + + // Step 2: Apply scale/softcap/mask/bias, online softmax, rescale output + for qi in 0..seq_q { + let score_row = &mut scores[qi * tile_kv..(qi + 1) * tile_kv]; + + let mut tile_max = neg_inf; + + for (ki, score) in score_row.iter_mut().enumerate() { + let kv_idx = kv_start + ki; + let mut val = *score * scale; + + if let Some(cap) = softcap { + val = cap * (val / cap).tanh(); + } + + if let Some(m) = mask + && m[qi * seq_kv + kv_idx] != 0 + { + val = neg_inf; + } + + if let Some(offset) = causal_offset + && (kv_idx as isize) > (qi as isize) + offset + { + val = neg_inf; + } + + if let Some(b) = bias { + val += b[qi * seq_kv + kv_idx]; + } + + *score = val; + if val > tile_max { + tile_max = val; + } + } + + if tile_max == neg_inf { + // All masked: zero out scores so gemm contributes nothing + for score in score_row.iter_mut() { + *score = T::zero(); + } + continue; + } + + let new_max = if row_max[qi] > tile_max { + row_max[qi] + } else { + tile_max + }; + + let mut tile_sum = T::zero(); + for score in score_row.iter_mut() { + let e = (*score - new_max).exp(); + *score = e; + tile_sum += e; + } + + let correction = if row_max[qi] == neg_inf { + T::zero() + } else { + (row_max[qi] - new_max).exp() + }; + + // Rescale existing output accumulator + let out_row = &mut output[qi * val_dim..(qi + 1) * val_dim]; + for o in out_row.iter_mut() { + *o = *o * correction; + } + + row_sum[qi] = row_sum[qi] * correction + tile_sum; + row_max[qi] = new_max; + } + + // Step 3: Value matmul via gemm + // output[seq_q, val_dim] += P[seq_q, tile_kv] @ V_tile[tile_kv, val_dim] + // + // P (scores buffer) row-major: rs=tile_kv, cs=1 + // V_tile row-major [tile_kv, val_dim]: rs=val_dim, cs=1 + // output row-major [seq_q, val_dim]: rs=val_dim, cs=1 + // + // dst = 1.0 * dst + 1.0 * P @ V (accumulate) + unsafe { + T::block_gemm(BlockGemmArgs { + m: seq_q, + n: val_dim, + k: tile_kv, + dst: output.as_mut_ptr(), + dst_cs: 1, + dst_rs: val_dim as isize, + read_dst: true, + lhs: scores.as_ptr(), + lhs_cs: 1, + lhs_rs: tile_kv as isize, + rhs: v.as_ptr().add(kv_start * val_dim), + rhs_cs: 1, + rhs_rs: val_dim as isize, + alpha: T::one(), + beta: T::one(), + }); + } + } + + // Final normalization: output /= row_sum + for qi in 0..seq_q { + let sum = row_sum[qi]; + if sum > T::zero() { + let inv_sum = T::one() / sum; + let out_row = &mut output[qi * val_dim..(qi + 1) * val_dim]; + for o in out_row.iter_mut() { + *o = *o * inv_sum; + } + } + // sum == 0 means all positions masked; output stays zero + } +} + +/// Naive attention: mathematically equivalent to flash but without KV tiling. +/// +/// Materializes the full [seq_q, seq_kv] score matrix, applies scale/softcap/mask/causal/bias, +/// row-wise softmax, then a single output matmul. Two large gemm calls per (batch, head). +/// +/// Faster than flash for short sequences. Also useful as a benchmarking baseline. +pub fn attention_naive( + query: FlexTensor, + key: FlexTensor, + value: FlexTensor, + mask: Option, + attn_bias: Option, + options: AttentionModuleOptions, +) -> FlexTensor { + dispatch_attention_dtype!( + query, + key, + value, + mask, + attn_bias, + options, + attention_naive_impl + ) +} + +fn attention_naive_impl( + query: FlexTensor, + key: FlexTensor, + value: FlexTensor, + mask: Option, + attn_bias: Option, + options: AttentionModuleOptions, +) -> FlexTensor +where + T: FlashGemm + burn_backend::Element, +{ + if let Some(softcap) = options.softcap { + assert!(softcap > 0.0, "softcap must be positive, got {softcap}"); + } + + let query = query.to_contiguous(); + let key = key.to_contiguous(); + let value = value.to_contiguous(); + + let q_shape = query.layout().shape(); + let k_shape = key.layout().shape(); + let v_shape = value.layout().shape(); + assert!(q_shape.num_dims() == 4, "attention_naive: query must be 4D"); + assert!(k_shape.num_dims() == 4, "attention_naive: key must be 4D"); + assert!(v_shape.num_dims() == 4, "attention_naive: value must be 4D"); + + let batch = q_shape[0]; + let heads = q_shape[1]; + let seq_q = q_shape[2]; + let head_dim = q_shape[3]; + assert!(head_dim > 0, "attention_naive: head_dim must be non-zero"); + + let seq_kv = k_shape[2]; + let val_dim = v_shape[3]; + + assert_eq!(k_shape[0], batch, "attention_naive: key batch mismatch"); + assert_eq!(k_shape[1], heads, "attention_naive: key heads mismatch"); + assert_eq!( + k_shape[3], head_dim, + "attention_naive: key head_dim mismatch" + ); + assert_eq!(v_shape[0], batch, "attention_naive: value batch mismatch"); + assert_eq!(v_shape[1], heads, "attention_naive: value heads mismatch"); + assert_eq!(v_shape[2], seq_kv, "attention_naive: value seq_kv mismatch"); + + let target = [batch, heads, seq_q, seq_kv]; + let mask_bcast = mask.map(|m| broadcast_attn_mask_bias(m, target, "mask")); + let bias_bcast = attn_bias.map(|b| broadcast_attn_mask_bias(b, target, "bias")); + + let scale = T::from( + options + .scale + .unwrap_or_else(|| 1.0 / (head_dim as f64).sqrt()), + ) + .unwrap(); + let softcap: Option = options.softcap.map(|s| T::from(s).unwrap()); + let causal_offset = if options.is_causal { + Some(seq_kv as isize - seq_q as isize) + } else { + None + }; + + let q_data: &[T] = query.storage(); + let k_data: &[T] = key.storage(); + let v_data: &[T] = value.storage(); + let mask_data: Option<&[u8]> = mask_bcast.as_ref().map(|b| b.tensor.bytes()); + let bias_data: Option<&[T]> = bias_bcast.as_ref().map(|b| b.tensor.storage()); + let (mask_batch_step, mask_head_step) = mask_bcast + .as_ref() + .map(|b| (b.batch_step, b.head_step)) + .unwrap_or((0, 0)); + let (bias_batch_step, bias_head_step) = bias_bcast + .as_ref() + .map(|b| (b.batch_step, b.head_step)) + .unwrap_or((0, 0)); + + let mut output = vec![T::zero(); batch * heads * seq_q * val_dim]; + let mut scores = vec![T::zero(); seq_q * seq_kv]; + + let q_head_stride = seq_q * head_dim; + let q_batch_stride = heads * q_head_stride; + let k_head_stride = seq_kv * head_dim; + let k_batch_stride = heads * k_head_stride; + let v_head_stride = seq_kv * val_dim; + let v_batch_stride = heads * v_head_stride; + let o_head_stride = seq_q * val_dim; + let o_batch_stride = heads * o_head_stride; + let mask_tile_len = seq_q * seq_kv; + + let params = AttentionParams { + scale, + softcap, + causal_offset, + seq_q, + seq_kv, + head_dim, + val_dim, + }; + + for b in 0..batch { + for h in 0..heads { + let q_off = b * q_batch_stride + h * q_head_stride; + let k_off = b * k_batch_stride + h * k_head_stride; + let v_off = b * v_batch_stride + h * v_head_stride; + let o_off = b * o_batch_stride + h * o_head_stride; + let mask_off = b * mask_batch_step + h * mask_head_step; + let bias_off = b * bias_batch_step + h * bias_head_step; + + naive_attention_head( + &q_data[q_off..q_off + q_head_stride], + &k_data[k_off..k_off + k_head_stride], + &v_data[v_off..v_off + v_head_stride], + &mut output[o_off..o_off + o_head_stride], + &mut scores, + ¶ms, + ( + mask_data.map(|m| &m[mask_off..mask_off + mask_tile_len]), + bias_data.map(|b| &b[bias_off..bias_off + mask_tile_len]), + ), + ); + } + } + + let shape = burn_std::Shape::from(vec![batch, heads, seq_q, val_dim]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(shape), + T::dtype(), + ) +} + +/// Process a single (batch, head) pair with naive (non-tiled) attention. +/// +/// 1. scores[seq_q, seq_kv] = Q @ K^T (one gemm) +/// 2. Apply scale/softcap/mask/causal/bias + softmax (per-row) +/// 3. output[seq_q, val_dim] = scores @ V (one gemm) +fn naive_attention_head( + q: &[T], + k: &[T], + v: &[T], + output: &mut [T], + scores: &mut [T], + p: &AttentionParams, + mask_bias: (Option<&[u8]>, Option<&[T]>), +) { + let (mask, bias) = mask_bias; + let neg_inf = T::neg_infinity(); + let AttentionParams { + scale, + softcap, + causal_offset, + seq_q, + seq_kv, + head_dim, + val_dim, + } = *p; + + // scores = Q @ K^T + unsafe { + T::block_gemm(BlockGemmArgs { + m: seq_q, + n: seq_kv, + k: head_dim, + dst: scores.as_mut_ptr(), + dst_cs: 1, + dst_rs: seq_kv as isize, + read_dst: false, + lhs: q.as_ptr(), + lhs_cs: 1, + lhs_rs: head_dim as isize, + rhs: k.as_ptr(), + rhs_cs: head_dim as isize, + rhs_rs: 1, + alpha: T::zero(), + beta: T::one(), + }); + } + + for qi in 0..seq_q { + let row = &mut scores[qi * seq_kv..(qi + 1) * seq_kv]; + + let mut row_max = neg_inf; + for (ki, s) in row.iter_mut().enumerate() { + let mut val = *s * scale; + + if let Some(cap) = softcap { + val = cap * (val / cap).tanh(); + } + + if let Some(m) = mask + && m[qi * seq_kv + ki] != 0 + { + val = neg_inf; + } + + if let Some(offset) = causal_offset + && (ki as isize) > (qi as isize) + offset + { + val = neg_inf; + } + + if let Some(b) = bias { + val += b[qi * seq_kv + ki]; + } + + *s = val; + if val > row_max { + row_max = val; + } + } + + if row_max == neg_inf { + row.fill(T::zero()); + continue; + } + + let mut sum = T::zero(); + for s in row.iter_mut() { + let e = (*s - row_max).exp(); + *s = e; + sum += e; + } + + let inv_sum = T::one() / sum; + for s in row.iter_mut() { + *s = *s * inv_sum; + } + } + + // output = softmax(scores) @ V + unsafe { + T::block_gemm(BlockGemmArgs { + m: seq_q, + n: val_dim, + k: seq_kv, + dst: output.as_mut_ptr(), + dst_cs: 1, + dst_rs: val_dim as isize, + read_dst: false, + lhs: scores.as_ptr(), + lhs_cs: 1, + lhs_rs: seq_kv as isize, + rhs: v.as_ptr(), + rhs_cs: 1, + rhs_rs: val_dim as isize, + alpha: T::zero(), + beta: T::one(), + }); + } +} + +// Tests kept here exercise flex-specific internals: direct calls into +// `attention_flash` / `attention_naive` (the public `attention()` dispatcher +// routes small shapes to naive so flash-path coverage requires a direct +// call), the `broadcast_attn_mask_bias` helper's rank/dim validation, +// flex-internal tiling and online-softmax correction paths, and +// dtype-specific kernels (f16/f64). Generic attention semantics (causal, +// custom scale, softcap, cross-attention, bool mask, additive bias, +// multi-batch/multi-head, single-element) live in +// crates/burn-backend-tests/tests/tensor/float/module/attention.rs, which +// exercises every backend. When adding new tests, keep them here only if +// they probe flex internals; otherwise add them to that suite. +#[cfg(test)] +mod tests { + use alloc::{vec, vec::Vec}; + use burn_backend::DType; + use burn_backend::ops::AttentionModuleOptions; + use burn_std::{BoolStore, Bytes, Shape, f16}; + + use crate::{FlexTensor, Layout}; + + fn flex_f32(data: Vec, shape: &[usize]) -> FlexTensor { + FlexTensor::new( + Bytes::from_elems(data), + Layout::contiguous(Shape::from(shape.to_vec())), + DType::F32, + ) + } + + fn flex_f64(data: Vec, shape: &[usize]) -> FlexTensor { + FlexTensor::new( + Bytes::from_elems(data), + Layout::contiguous(Shape::from(shape.to_vec())), + DType::F64, + ) + } + + fn flex_f16(data: Vec, shape: &[usize]) -> FlexTensor { + FlexTensor::new( + Bytes::from_elems(data), + Layout::contiguous(Shape::from(shape.to_vec())), + DType::F16, + ) + } + + fn flex_bool(data: Vec, shape: &[usize]) -> FlexTensor { + FlexTensor::new( + Bytes::from_elems(data), + Layout::contiguous(Shape::from(shape.to_vec())), + DType::Bool(BoolStore::Native), + ) + } + + /// Assert two attention-output slices are elementwise close. Used by the + /// flash broadcast tests below. + fn assert_attention_outputs_close(bcast: &[f32], full: &[f32], label: &str) { + assert_eq!(bcast.len(), full.len(), "{label}: length mismatch"); + for (i, (&a, &b)) in bcast.iter().zip(full).enumerate() { + assert!((a - b).abs() < 1e-5, "{label} mismatch at {i}: {a} vs {b}"); + } + } + + #[test] + fn test_flash_bias_broadcast_across_batch_and_heads() { + // Exercises the flash path for a `[1, 1, seq_q, seq_kv]` broadcast + // bias. The dispatcher in `attention()` routes to naive for + // `seq_q * seq_kv <= NAIVE_SCORE_BUDGET` (and the backend-tests + // attention suite uses small shapes), so the flash entry needs a + // direct call to stay covered. General broadcast semantics for the + // main `attention()` path live in + // crates/burn-backend-tests/tests/tensor/float/module/attention.rs. + let batch = 2; + let heads = 2; + let seq_q = 3; + let seq_kv = 5; + let head_dim = 4; + + let mk = |shape: &[usize], g: &dyn Fn(usize) -> f32| -> FlexTensor { + let len: usize = shape.iter().product(); + flex_f32((0..len).map(g).collect(), shape) + }; + + let q = mk(&[batch, heads, seq_q, head_dim], &|i| { + (i as f32 * 0.1).sin() + }); + let k = mk(&[batch, heads, seq_kv, head_dim], &|i| { + (i as f32 * 0.1 + 1.0).sin() + }); + let v = mk(&[batch, heads, seq_kv, head_dim], &|i| { + (i as f32 * 0.1 + 2.0).sin() + }); + + let bias_tile: Vec = (0..seq_q * seq_kv) + .map(|i| (i as f32 * 0.4).sin()) + .collect(); + let bias_bcast = flex_f32(bias_tile.clone(), &[1, 1, seq_q, seq_kv]); + let bias_full_vec: Vec = bias_tile + .iter() + .cloned() + .cycle() + .take(batch * heads * seq_q * seq_kv) + .collect(); + let bias_full = flex_f32(bias_full_vec, &[batch, heads, seq_q, seq_kv]); + + let out_bcast = super::attention_flash( + q.clone(), + k.clone(), + v.clone(), + None, + Some(bias_bcast), + Default::default(), + ); + let out_full = super::attention_flash(q, k, v, None, Some(bias_full), Default::default()); + + let bcast: &[f32] = out_bcast.storage(); + let full: &[f32] = out_full.storage(); + assert_attention_outputs_close(bcast, full, "flash bias[1,1,sq,skv]"); + } + + #[test] + fn test_flash_bool_mask_broadcast_across_batch_and_heads() { + // Flash-path counterpart to the bias broadcast test, but for a + // `[1, 1, seq_q, seq_kv]` bool mask. The mask (u8) slicing path with + // stride-0 batch/head steps is a different branch from the bias + // (f32) path, so both need direct flash coverage. + let batch = 2; + let heads = 2; + let seq_q = 3; + let seq_kv = 5; + let head_dim = 4; + + let mk = |shape: &[usize], g: &dyn Fn(usize) -> f32| -> FlexTensor { + let len: usize = shape.iter().product(); + flex_f32((0..len).map(g).collect(), shape) + }; + + let q = mk(&[batch, heads, seq_q, head_dim], &|i| { + (i as f32 * 0.1).sin() + }); + let k = mk(&[batch, heads, seq_kv, head_dim], &|i| { + (i as f32 * 0.1 + 1.0).sin() + }); + let v = mk(&[batch, heads, seq_kv, head_dim], &|i| { + (i as f32 * 0.1 + 2.0).sin() + }); + + // Mask out columns 3 and 4 (1 == masked out). + let mask_tile: Vec = (0..seq_q * seq_kv) + .map(|i| if (i % seq_kv) >= 3 { 1u8 } else { 0u8 }) + .collect(); + let mask_bcast = flex_bool(mask_tile.clone(), &[1, 1, seq_q, seq_kv]); + let mask_full_vec: Vec = mask_tile + .iter() + .copied() + .cycle() + .take(batch * heads * seq_q * seq_kv) + .collect(); + let mask_full = flex_bool(mask_full_vec, &[batch, heads, seq_q, seq_kv]); + + let out_bcast = super::attention_flash( + q.clone(), + k.clone(), + v.clone(), + Some(mask_bcast), + None, + Default::default(), + ); + let out_full = super::attention_flash(q, k, v, Some(mask_full), None, Default::default()); + + let bcast: &[f32] = out_bcast.storage(); + let full: &[f32] = out_full.storage(); + assert_attention_outputs_close(bcast, full, "flash bool mask[1,1,sq,skv]"); + } + + #[test] + #[should_panic(expected = "must be 4D")] + fn test_mask_wrong_rank_panics() { + // Contract: the broadcast helper rejects non-4D mask/bias up-front + // with a clearer message than `expand`'s rank prepending would produce. + let mask = flex_bool(vec![0u8; 6], &[2, 3]); + super::broadcast_attn_mask_bias(mask, [1, 1, 2, 3], "mask"); + } + + #[test] + #[should_panic(expected = "bias dim 1 must be 3 or 1, got 2")] + fn test_bias_incompatible_dim_panics() { + // Contract: a dim that is neither equal to target nor `1` is a hard + // error. Here heads=2 in the bias but target heads=3. + let bias = flex_f32(vec![0.0f32; 2 * 2 * 4 * 5], &[2, 2, 4, 5]); + super::broadcast_attn_mask_bias(bias, [2, 3, 4, 5], "bias"); + } + + #[test] + fn test_multi_tile_seq_kv() { + // seq_kv > TILE_KV (64) exercises tiling with online softmax + // correction. 128 keys = exactly 2 tiles. + let seq_q = 2; + let seq_kv = 128; + let head_dim = 4; + let val_dim = 2; + + // Q: first query selects dim 0, second selects dim 1. + let mut q_data = vec![0.0f32; seq_q * head_dim]; + q_data[0] = 1.0; + q_data[head_dim + 1] = 1.0; + + // K: all keys identical so all scores are equal. + let k_data = vec![0.1f32; seq_kv * head_dim]; + + // V: linearly increasing values. + let mut v_data = vec![0.0f32; seq_kv * val_dim]; + for i in 0..seq_kv { + v_data[i * val_dim] = i as f32; + v_data[i * val_dim + 1] = (seq_kv - 1 - i) as f32; + } + + let q = flex_f32(q_data, &[1, 1, seq_q, head_dim]); + let k = flex_f32(k_data, &[1, 1, seq_kv, head_dim]); + let v = flex_f32(v_data, &[1, 1, seq_kv, val_dim]); + + let result = super::attention(q, k, v, None, None, Default::default()); + let data: &[f32] = result.storage(); + + // All scores equal -> output = mean of 0..127 = 63.5 for both cols. + assert_eq!(data.len(), seq_q * val_dim); + assert!( + (data[0] - 63.5).abs() < 0.1, + "expected ~63.5, got {}", + data[0] + ); + assert!( + (data[1] - 63.5).abs() < 0.1, + "expected ~63.5, got {}", + data[1] + ); + } + + #[test] + fn test_multi_tile_causal() { + // seq_kv=128 with causal mask: each query sees only up to + // causal_offset + its own index worth of keys. + let seq_q = 4; + let seq_kv = 128; + let head_dim = 2; + let val_dim = 1; + + // All queries/keys [1, 0] so all visible scores equal. + let mut q_data = vec![0.0f32; seq_q * head_dim]; + for i in 0..seq_q { + q_data[i * head_dim] = 1.0; + } + let mut k_data = vec![0.0f32; seq_kv * head_dim]; + for i in 0..seq_kv { + k_data[i * head_dim] = 1.0; + } + let v_data: Vec = (0..seq_kv).map(|i| i as f32).collect(); + + let q = flex_f32(q_data, &[1, 1, seq_q, head_dim]); + let k = flex_f32(k_data, &[1, 1, seq_kv, head_dim]); + let v = flex_f32(v_data, &[1, 1, seq_kv, val_dim]); + + let opts = AttentionModuleOptions { + is_causal: true, + ..Default::default() + }; + let result = super::attention(q, k, v, None, None, opts); + let data: &[f32] = result.storage(); + + // causal_offset = seq_kv - seq_q = 124 + // q[0] sees 0..=124 -> mean 62.0; q[1] 62.5; q[2] 63.0; q[3] 63.5. + assert_eq!(data.len(), seq_q); + assert!((data[0] - 62.0).abs() < 0.1, "q0: got {}", data[0]); + assert!((data[1] - 62.5).abs() < 0.1, "q1: got {}", data[1]); + assert!((data[2] - 63.0).abs() < 0.1, "q2: got {}", data[2]); + assert!((data[3] - 63.5).abs() < 0.1, "q3: got {}", data[3]); + } + + #[test] + fn test_tile_boundary_mask() { + // Mask falls exactly on a tile boundary: first 64 keys masked, + // next 64 visible. + let seq_q = 1; + let seq_kv = 128; + let head_dim = 2; + let val_dim = 1; + + let q_data = vec![1.0f32, 0.0]; + let k_data = [1.0f32, 0.0].repeat(seq_kv); + let v_data: Vec = (0..seq_kv).map(|i| i as f32).collect(); + // true == masked out. + let mask_data: Vec = (0..seq_kv).map(|i| (i < 64) as u8).collect(); + + let q = flex_f32(q_data, &[1, 1, seq_q, head_dim]); + let k = flex_f32(k_data, &[1, 1, seq_kv, head_dim]); + let v = flex_f32(v_data, &[1, 1, seq_kv, val_dim]); + let mask = flex_bool(mask_data, &[1, 1, seq_q, seq_kv]); + + let result = super::attention(q, k, v, Some(mask), None, Default::default()); + let data: &[f32] = result.storage(); + + // Visible = 64..128 -> mean of 64..=127 = 95.5. + assert!( + (data[0] - 95.5).abs() < 0.1, + "expected ~95.5, got {}", + data[0] + ); + } + + #[test] + fn test_non_uniform_scores_across_tiles() { + // Forces the online softmax correction path: tile 1 has much + // larger scores than tile 0, so the correction factor + // exp(old_max - new_max) < 1. + let seq_q = 1; + let seq_kv = 128; + let head_dim = 1; + let val_dim = 1; + + let q_data = vec![1.0f32]; + let mut k_data = vec![0.0f32; seq_kv]; + for k in k_data.iter_mut().take(64) { + *k = 0.1; + } + for k in k_data.iter_mut().skip(64) { + *k = 5.0; + } + let mut v_data = vec![0.0f32; seq_kv]; + for v in v_data.iter_mut().skip(64) { + *v = 1.0; + } + + let q = flex_f32(q_data, &[1, 1, seq_q, head_dim]); + let k = flex_f32(k_data, &[1, 1, seq_kv, head_dim]); + let v = flex_f32(v_data, &[1, 1, seq_kv, val_dim]); + + let result = super::attention(q, k, v, None, None, Default::default()); + let data: &[f32] = result.storage(); + + // Large-score keys (tile 2) dominate softmax -> output ~ 1.0. + assert!(data[0] > 0.99, "expected ~1.0, got {}", data[0]); + } + + #[test] + fn test_partial_last_tile() { + // seq_kv=100 is not a multiple of TILE_KV=64, so the last tile has + // 36 elements -> exercises partial-tile gemm and score buffer + // indexing. + let seq_q = 2; + let seq_kv = 100; + let head_dim = 2; + let val_dim = 1; + + let q_data = [0.1f32, 0.1].repeat(seq_q); + let k_data = [0.1f32, 0.1].repeat(seq_kv); + let v_data: Vec = (0..seq_kv).map(|i| i as f32).collect(); + + let q = flex_f32(q_data, &[1, 1, seq_q, head_dim]); + let k = flex_f32(k_data, &[1, 1, seq_kv, head_dim]); + let v = flex_f32(v_data, &[1, 1, seq_kv, val_dim]); + + let result = super::attention(q, k, v, None, None, Default::default()); + let data: &[f32] = result.storage(); + + // Uniform attention -> mean of 0..99 = 49.5. + assert_eq!(data.len(), seq_q); + assert!( + (data[0] - 49.5).abs() < 0.1, + "expected ~49.5, got {}", + data[0] + ); + assert!( + (data[1] - 49.5).abs() < 0.1, + "expected ~49.5, got {}", + data[1] + ); + } + + /// Verify naive attention produces the same results as flash attention + /// across various configurations. + #[test] + fn test_naive_matches_flash() { + /// Deterministic tensor for cross-validation tests. + /// + /// Uses `i * 997 % N` as a cheap hash: 997 is prime and coprime to + /// any power-of-two length, so the sequence visits all residues + /// before repeating. For f32 this gives values in [-0.5, 0.5]. For + /// bool masks it gives ~30% density with an irregular pattern that + /// varies across rows (unlike a regular `i % 3` stride). + fn make_tensor(shape: &[usize], dtype: DType) -> FlexTensor { + let len: usize = shape.iter().product(); + match dtype { + DType::F32 => { + let data: Vec = + (0..len).map(|i| ((i % 997) as f32 / 997.0) - 0.5).collect(); + flex_f32(data, shape) + } + DType::Bool(_) => { + let data: Vec = (0..len) + .map(|i| (i.wrapping_mul(997) % 100 < 30) as u8) + .collect(); + flex_bool(data, shape) + } + _ => unreachable!(), + } + } + + #[allow(clippy::too_many_arguments)] + fn run_both( + batch: usize, + heads: usize, + seq_q: usize, + seq_kv: usize, + head_dim: usize, + val_dim: usize, + with_mask: bool, + with_bias: bool, + options: AttentionModuleOptions, + label: &str, + ) { + let f32_dt = DType::F32; + let bool_dt = DType::Bool(BoolStore::Native); + let q = make_tensor(&[batch, heads, seq_q, head_dim], f32_dt); + let k = make_tensor(&[batch, heads, seq_kv, head_dim], f32_dt); + let v = make_tensor(&[batch, heads, seq_kv, val_dim], f32_dt); + let score_shape = [batch, heads, seq_q, seq_kv]; + let mask = with_mask.then(|| make_tensor(&score_shape, bool_dt)); + let bias = with_bias.then(|| make_tensor(&score_shape, f32_dt)); + + let flash = super::attention_flash( + q.clone(), + k.clone(), + v.clone(), + mask.clone(), + bias.clone(), + options, + ); + let naive = super::attention_naive(q, k, v, mask, bias, options); + + let flash_data: &[f32] = flash.storage(); + let naive_data: &[f32] = naive.storage(); + assert_eq!( + flash_data.len(), + naive_data.len(), + "{label}: length mismatch" + ); + + for (i, (&f, &n)) in flash_data.iter().zip(naive_data.iter()).enumerate() { + let diff = (f - n).abs(); + let tol = 1e-4 * f.abs().max(n.abs()).max(1.0); + assert!( + diff < tol, + "{label}: position {i}: flash={f} vs naive={n}, diff={diff}" + ); + } + } + + let default = AttentionModuleOptions::default(); + let causal = AttentionModuleOptions { + is_causal: true, + ..Default::default() + }; + let all_opts = AttentionModuleOptions { + scale: Some(0.05), + softcap: Some(30.0), + is_causal: true, + }; + + run_both(1, 1, 4, 4, 8, 8, false, false, default, "basic_4x4"); + run_both( + 2, + 4, + 8, + 8, + 16, + 16, + false, + false, + default, + "multi_head_batch", + ); + run_both(1, 2, 4, 32, 16, 16, false, false, default, "cross_attn"); + run_both(1, 1, 4, 128, 16, 16, false, false, default, "multi_tile"); + run_both(1, 2, 16, 16, 32, 32, false, false, causal, "causal"); + run_both(2, 2, 16, 16, 32, 32, false, false, all_opts, "all_options"); + run_both(1, 1, 32, 256, 64, 64, false, false, causal, "large_causal"); + run_both(1, 2, 8, 8, 16, 16, true, false, default, "with_mask"); + run_both(1, 2, 8, 8, 16, 16, false, true, default, "with_bias"); + run_both( + 2, + 2, + 16, + 128, + 32, + 32, + true, + true, + causal, + "mask_bias_causal", + ); + run_both(1, 1, 4, 100, 16, 16, false, false, default, "partial_tile"); + run_both( + 1, + 2, + 8, + 100, + 16, + 16, + false, + false, + causal, + "partial_tile_causal", + ); + } + + #[test] + fn test_f64_flash_attention() { + let q = flex_f64(vec![1.0f64, 0.0, 0.0, 1.0], &[1, 1, 2, 2]); + let k = q.clone(); + let v = flex_f64(vec![10.0f64, 20.0], &[1, 1, 2, 1]); + + let result = super::attention(q, k, v, None, None, Default::default()); + let data: &[f64] = result.storage(); + + assert!((data[0] - 13.30).abs() < 0.1, "got {}", data[0]); + assert!((data[1] - 16.70).abs() < 0.1, "got {}", data[1]); + } + + /// Verify f16 cast-to-f32 round-trip produces correct results for both + /// paths. + #[test] + fn test_f16_attention() { + let q_vec: Vec = [1.0f32, 0.0, 0.0, 1.0] + .iter() + .map(|&v| f16::from_f32(v)) + .collect(); + let v_vec: Vec = [10.0f32, 20.0].iter().map(|&v| f16::from_f32(v)).collect(); + let q = flex_f16(q_vec, &[1, 1, 2, 2]); + let k = q.clone(); + let v = flex_f16(v_vec, &[1, 1, 2, 1]); + + let flash = super::attention_flash( + q.clone(), + k.clone(), + v.clone(), + None, + None, + Default::default(), + ); + let naive = super::attention_naive(q, k, v, None, None, Default::default()); + + let flash_data: &[f16] = flash.storage(); + let naive_data: &[f16] = naive.storage(); + + // softmax([1/sqrt(2), 0]) = [0.670, 0.330] -> row0: ~13.3, row1: ~16.7 + for (label, data) in [("flash", flash_data), ("naive", naive_data)] { + let r0 = data[0].to_f32(); + let r1 = data[1].to_f32(); + assert!((r0 - 13.30).abs() < 0.2, "{label} row0: got {r0}"); + assert!((r1 - 16.70).abs() < 0.2, "{label} row1: got {r1}"); + } + } +} diff --git a/crates/burn-flex/src/ops/binary.rs b/crates/burn-flex/src/ops/binary.rs new file mode 100644 index 0000000..aaca5b5 --- /dev/null +++ b/crates/burn-flex/src/ops/binary.rs @@ -0,0 +1,1124 @@ +//! Binary tensor operations (add, sub, mul, div). + +use alloc::vec::Vec; +use burn_backend::{DType, Element}; +use burn_std::{Bytes, Shape, bf16, f16}; + +use crate::FlexTensor; +use crate::layout::Layout; +use crate::strided_index::StridedIter; + +#[cfg(feature = "simd")] +use crate::simd; + +/// Operation type for SIMD dispatch. +#[derive(Clone, Copy)] +pub enum BinaryOp { + Add, + Sub, + Mul, + Div, +} + +/// Apply a binary operation element-wise to two tensors. +/// +/// Requires tensors to have the same shape. Uses SIMD acceleration for f32 +/// when available and both tensors are contiguous. +/// +/// Pass `simd_hint` to enable direct SIMD dispatch for standard ops (add/sub/mul/div). +/// Pass `None` for custom operations that have no SIMD fast path. +pub fn binary_op( + lhs: FlexTensor, + rhs: FlexTensor, + f32_op: F32Op, + f64_op: F64Op, + simd_hint: Option, +) -> FlexTensor +where + F32Op: Fn(f32, f32) -> f32 + Copy, + F64Op: Fn(f64, f64) -> f64 + Copy, +{ + debug_assert_eq!(lhs.dtype(), rhs.dtype(), "binary_op: dtype mismatch"); + + // Broadcast tensors to the same shape if needed + let (lhs, rhs) = crate::ops::expand::broadcast_binary(lhs, rhs); + + let dtype = lhs.dtype(); + + match dtype { + DType::F32 => binary_op_f32(lhs, &rhs, f32_op, simd_hint), + DType::F64 => binary_op_typed(lhs, &rhs, f64_op), + DType::F16 => binary_op_typed(lhs, &rhs, |a: f16, b: f16| { + f16::from_f32(f32_op(a.to_f32(), b.to_f32())) + }), + DType::BF16 => binary_op_typed(lhs, &rhs, |a: bf16, b: bf16| { + bf16::from_f32(f32_op(a.to_f32(), b.to_f32())) + }), + _ => panic!("binary_op: unsupported dtype {:?}", dtype), + } +} + +/// Specialized binary operation for f32 with SIMD fast path. +#[cfg(feature = "simd")] +fn binary_op_f32( + mut lhs: FlexTensor, + rhs: &FlexTensor, + op: Op, + simd_hint: Option, +) -> FlexTensor +where + Op: Fn(f32, f32) -> f32, +{ + // Permuted lhs + broadcast rhs (e.g. `x.permute(...) - mean`): + // the generic path would walk the permuted lhs with a scalar + // `StridedIter`, an order of magnitude slower than the SIMD fast + // path. Pay one memcpy to materialize lhs contiguous so the fast + // paths below can take over. + // + // Gate on `simd_hint.is_some()` so custom ops like `atan2` or + // `powf` (which have no SIMD fast path and go straight to + // `binary_op_typed`) don't pay for a memcpy they can't benefit + // from. Their strided fallback handles non-contig lhs directly. + if simd_hint.is_some() && !lhs.layout().is_contiguous() && rhs.layout().strides().contains(&0) { + lhs = lhs.to_contiguous(); + } + + // In-place SIMD fast path: lhs unique contiguous at offset 0, rhs + // contiguous (no broadcast). + if let Some(simd_op) = simd_hint + && lhs.is_unique() + && let (Some((0, l_end)), Some((r_start, r_end))) = ( + lhs.layout().contiguous_offsets(), + rhs.layout().contiguous_offsets(), + ) + { + let r_slice: &[f32] = &rhs.storage()[r_start..r_end]; + let lhs_storage: &mut [f32] = lhs.storage_mut(); + let l_slice = &mut lhs_storage[..l_end]; + + match simd_op { + BinaryOp::Add => simd::add_inplace_f32(l_slice, r_slice), + BinaryOp::Sub => simd::sub_inplace_f32(l_slice, r_slice), + BinaryOp::Mul => simd::mul_inplace_f32(l_slice, r_slice), + BinaryOp::Div => simd::div_inplace_f32(l_slice, r_slice), + } + return lhs; + } + + // Broadcast SIMD fast path: rhs is broadcast via stride-0 dims in + // one of two hot shapes that dominate layer_norm decomposition -- + // shared-row (`gamma.unsqueeze() * x`) or per-row scalar + // (`x - x.mean_dim(-1)`). + if let Some(simd_op) = simd_hint + && let Some(pattern) = detect_broadcast_pattern(lhs.layout(), rhs) + { + return apply_broadcast_pattern_f32(lhs, rhs, simd_op, pattern); + } + + binary_op_typed(lhs, rhs, op) +} + +/// Categorization of the two broadcast patterns we can accelerate. +/// +/// Consumers assume `lhs` and `rhs` already share the same logical +/// shape and that `lhs` is row-contiguous at offset 0. +#[cfg(feature = "simd")] +#[derive(Debug, Clone, Copy)] +enum BroadcastView { + /// rhs's inner `row_len` elements form a contiguous row that is + /// shared across `outer_count` outer positions. Starts at + /// `rhs_row_offset` in rhs's storage. + SharedRow { + outer_count: usize, + row_len: usize, + rhs_row_offset: usize, + }, + /// rhs's inner `row_len` elements are all the same scalar, + /// stepping through `outer_count` scalars along the outer dims + /// starting at `rhs_scalar_base` in rhs's storage. + PerRowScalar { + outer_count: usize, + row_len: usize, + rhs_scalar_base: usize, + }, +} + +/// Detect whether rhs can be handled as one of the accelerated +/// broadcast patterns, returning `None` if the stride pattern doesn't +/// fit either bucket or the resulting offsets would leave rhs's +/// storage. +#[cfg(feature = "simd")] +fn detect_broadcast_pattern(lhs: &Layout, rhs: &FlexTensor) -> Option { + let rhs_layout = rhs.layout(); + let rhs_storage_elems = rhs.storage::().len(); + // Require lhs to be row-contiguous at offset 0. The broadcast kernel + // below uses linear offsets into lhs's storage; relaxing this would + // complicate the indexing without helping the hot layer_norm path. + let (l_start, _) = lhs.contiguous_offsets()?; + if l_start != 0 { + return None; + } + let ndims = lhs.num_dims(); + if ndims == 0 || rhs_layout.num_dims() != ndims { + return None; + } + let lhs_shape = lhs.shape(); + let rhs_strides = rhs_layout.strides(); + + let last_stride = rhs_strides[ndims - 1]; + + // Case A: shared row. Innermost rhs stride is 1, and every outer + // dim either has stride 0 (a broadcast dim) or size 1 (stride + // doesn't matter since the dim never advances past index 0). + if last_stride == 1 { + let outer_ok = (0..ndims - 1).all(|d| rhs_strides[d] == 0 || lhs_shape[d] == 1); + if outer_ok { + let outer_count: usize = (0..ndims - 1).map(|d| lhs_shape[d]).product(); + let row_len = lhs_shape[ndims - 1]; + if outer_count == 0 || row_len == 0 { + return None; + } + let rhs_row_offset = rhs_layout.start_offset(); + // Bounds: kernel reads `rhs_storage[off..off+row_len]`. + if rhs_row_offset.checked_add(row_len)? > rhs_storage_elems { + return None; + } + return Some(BroadcastView::SharedRow { + outer_count, + row_len, + rhs_row_offset, + }); + } + } + + // Case B: per-row scalar. Innermost dims all have stride 0 in + // rhs and outer dims walk rhs contiguously in row-major order. + if last_stride == 0 { + // Count the trailing stride-0 dims to find the inner scalar span. + let mut inner_dims = 0usize; + let mut row_len: usize = 1; + for d in (0..ndims).rev() { + if rhs_strides[d] == 0 { + inner_dims += 1; + row_len *= lhs_shape[d]; + } else { + break; + } + } + if inner_dims == 0 { + return None; + } + // The outer dims must walk rhs's storage contiguously in + // row-major order. + let outer_ndims = ndims - inner_dims; + let mut expected: isize = 1; + for d in (0..outer_ndims).rev() { + if rhs_strides[d] != expected { + return None; + } + expected *= lhs_shape[d] as isize; + } + let outer_count: usize = (0..outer_ndims).map(|d| lhs_shape[d]).product(); + if outer_count == 0 || row_len == 0 { + return None; + } + let rhs_scalar_base = rhs_layout.start_offset(); + // Bounds: kernel reads `rhs_storage[base..base+outer_count]`. + if rhs_scalar_base.checked_add(outer_count)? > rhs_storage_elems { + return None; + } + return Some(BroadcastView::PerRowScalar { + outer_count, + row_len, + rhs_scalar_base, + }); + } + + None +} + +/// Execute a detected broadcast pattern for f32. Writes in-place into +/// lhs when unique; otherwise allocates a fresh contiguous output. +#[cfg(feature = "simd")] +fn apply_broadcast_pattern_f32( + mut lhs: FlexTensor, + rhs: &FlexTensor, + simd_op: BinaryOp, + pattern: BroadcastView, +) -> FlexTensor { + let numel = lhs.layout().num_elements(); + let rhs_storage = rhs.storage::(); + + if lhs.is_unique() { + let dst = &mut lhs.storage_mut::()[..numel]; + run_broadcast_pattern_f32(dst, rhs_storage, simd_op, pattern); + lhs + } else { + // Copy lhs once, then apply the broadcast in place. The + // memcpy is cheaper than the StridedIter fallback it replaces. + let mut out: Vec = lhs.storage::()[..numel].to_vec(); + run_broadcast_pattern_f32(&mut out, rhs_storage, simd_op, pattern); + make_tensor(out, lhs.layout().shape().clone(), lhs.dtype()) + } +} + +/// Shared kernel: run the chosen broadcast pattern against a mutable +/// destination buffer (which already holds lhs's values) and rhs's +/// storage slice. +#[cfg(feature = "simd")] +fn run_broadcast_pattern_f32( + dst: &mut [f32], + rhs_storage: &[f32], + simd_op: BinaryOp, + pattern: BroadcastView, +) { + match pattern { + BroadcastView::SharedRow { + outer_count, + row_len, + rhs_row_offset, + } => { + let rhs_row: &[f32] = &rhs_storage[rhs_row_offset..rhs_row_offset + row_len]; + let total = outer_count * row_len; + // One SIMD dispatch covers the whole outer walk. The kernel + // keeps `rhs_row` in registers across rows for small row + // lengths, and pays the macerator feature-detection cost + // exactly once. + let dst_full = &mut dst[..total]; + match simd_op { + BinaryOp::Add => simd::add_shared_row_inplace_f32(dst_full, rhs_row), + BinaryOp::Sub => simd::sub_shared_row_inplace_f32(dst_full, rhs_row), + BinaryOp::Mul => simd::mul_shared_row_inplace_f32(dst_full, rhs_row), + BinaryOp::Div => simd::div_shared_row_inplace_f32(dst_full, rhs_row), + } + } + BroadcastView::PerRowScalar { + outer_count, + row_len, + rhs_scalar_base, + } => { + let scalars = &rhs_storage[rhs_scalar_base..rhs_scalar_base + outer_count]; + // One monomorphized helper per op. The closure is statically + // known at each call site so LLVM still autovectorizes the + // inner scalar loop, and the outer op dispatch happens once. + match simd_op { + BinaryOp::Add => per_row_scalar_apply(dst, scalars, row_len, |a, b| a + b), + BinaryOp::Sub => per_row_scalar_apply(dst, scalars, row_len, |a, b| a - b), + BinaryOp::Mul => per_row_scalar_apply(dst, scalars, row_len, |a, b| a * b), + BinaryOp::Div => per_row_scalar_apply(dst, scalars, row_len, |a, b| a / b), + } + } + } +} + +/// Apply `dst[r * row_len + j] = op(dst[r * row_len + j], scalars[r])` +/// for `r in 0..scalars.len(), j in 0..row_len`. Generic over `Op` so +/// each call site gets a monomorphized, autovectorizable inner loop. +#[cfg(feature = "simd")] +#[inline] +fn per_row_scalar_apply(dst: &mut [f32], scalars: &[f32], row_len: usize, op: Op) +where + Op: Fn(f32, f32) -> f32, +{ + for (i, &scalar) in scalars.iter().enumerate() { + let start = i * row_len; + for x in dst[start..start + row_len].iter_mut() { + *x = op(*x, scalar); + } + } +} + +/// Fallback when SIMD is disabled. +#[cfg(not(feature = "simd"))] +fn binary_op_f32( + lhs: FlexTensor, + rhs: &FlexTensor, + op: Op, + _simd_hint: Option, +) -> FlexTensor +where + Op: Fn(f32, f32) -> f32, +{ + binary_op_typed(lhs, rhs, op) +} + +/// Binary operation with in-place optimization for Pod types. +pub(crate) fn binary_op_typed(mut lhs: FlexTensor, rhs: &FlexTensor, op: Op) -> FlexTensor +where + E: Element + bytemuck::Pod, + Op: Fn(E, E) -> E, +{ + let rhs_storage: &[E] = rhs.storage(); + + // In-place fast path: lhs unique, contiguous at offset 0, rhs contiguous + if lhs.is_unique() + && let (Some((0, l_end)), Some((r_start, r_end))) = ( + lhs.layout().contiguous_offsets(), + rhs.layout().contiguous_offsets(), + ) + { + let lhs_storage: &mut [E] = lhs.storage_mut(); + let r_slice = &rhs_storage[r_start..r_end]; + for (l, &r) in lhs_storage[..l_end].iter_mut().zip(r_slice) { + *l = op(*l, r); + } + return lhs; + } + + // Allocating path + let shape = lhs.layout().shape().clone(); + let dtype = lhs.dtype(); + let lhs_storage: &[E] = lhs.storage(); + + let result: Vec = match ( + lhs.layout().contiguous_offsets(), + rhs.layout().contiguous_offsets(), + ) { + // Both contiguous (but lhs not at offset 0) + (Some((l_start, l_end)), Some((r_start, r_end))) => { + let l_slice = &lhs_storage[l_start..l_end]; + let r_slice = &rhs_storage[r_start..r_end]; + l_slice + .iter() + .zip(r_slice) + .map(|(&a, &b)| op(a, b)) + .collect() + } + // Fast path for 2D non-contiguous (common for transpose) + _ if lhs.layout().num_dims() == 2 => { + apply_2d_strided(lhs_storage, rhs_storage, lhs.layout(), rhs.layout(), op) + } + // General fallback + _ => { + let lhs_iter = StridedIter::new(lhs.layout()); + let rhs_iter = StridedIter::new(rhs.layout()); + lhs_iter + .zip(rhs_iter) + .map(|(li, ri)| op(lhs_storage[li], rhs_storage[ri])) + .collect() + } + }; + + make_tensor(result, shape, dtype) +} + +/// Fast 2D strided binary operation using row-based iteration. +#[inline] +pub(crate) fn apply_2d_strided( + lhs: &[E], + rhs: &[E], + lhs_layout: &Layout, + rhs_layout: &Layout, + op: Op, +) -> Vec +where + E: Copy, + Op: Fn(E, E) -> R, +{ + let (rows, cols, l_row_stride, l_col_stride) = lhs_layout.as_2d_strides().unwrap(); + let (_, _, r_row_stride, r_col_stride) = rhs_layout.as_2d_strides().unwrap(); + let l_offset = lhs_layout.start_offset() as isize; + let r_offset = rhs_layout.start_offset() as isize; + + let mut result = Vec::with_capacity(rows * cols); + + for row in 0..rows { + let l_row_start = l_offset + row as isize * l_row_stride; + let r_row_start = r_offset + row as isize * r_row_stride; + for col in 0..cols { + let l_idx = (l_row_start + col as isize * l_col_stride) as usize; + let r_idx = (r_row_start + col as isize * r_col_stride) as usize; + result.push(op(lhs[l_idx], rhs[r_idx])); + } + } + + result +} + +/// Apply a scalar operation to each element of a tensor. +/// +/// Attempts in-place mutation when tensor is contiguous at offset 0. +pub fn scalar_op( + tensor: FlexTensor, + scalar: f64, + f32_op: F32Op, + f64_op: F64Op, +) -> FlexTensor +where + F32Op: Fn(f32, f32) -> f32 + Copy, + F64Op: Fn(f64, f64) -> f64 + Copy, +{ + let dtype = tensor.dtype(); + + match dtype { + DType::F32 => scalar_op_typed(tensor, scalar as f32, f32_op), + DType::F64 => scalar_op_typed(tensor, scalar, f64_op), + DType::F16 => { + let scalar_f16 = f16::from_f32(scalar as f32); + let s = scalar_f16.to_f32(); + scalar_op_typed(tensor, scalar_f16, |a: f16, _| { + f16::from_f32(f32_op(a.to_f32(), s)) + }) + } + DType::BF16 => { + let scalar_bf16 = bf16::from_f32(scalar as f32); + let s = scalar_bf16.to_f32(); + scalar_op_typed(tensor, scalar_bf16, |a: bf16, _| { + bf16::from_f32(f32_op(a.to_f32(), s)) + }) + } + _ => panic!("scalar_op: unsupported dtype {:?}", dtype), + } +} + +pub(crate) fn scalar_op_typed(mut tensor: FlexTensor, scalar: E, op: Op) -> FlexTensor +where + E: Element + bytemuck::Pod, + Op: Fn(E, E) -> E, +{ + // In-place fast path: unique, contiguous at offset 0 + if tensor.is_unique() + && let Some((0, end)) = tensor.layout().contiguous_offsets() + { + let storage: &mut [E] = tensor.storage_mut(); + for x in storage[..end].iter_mut() { + *x = op(*x, scalar); + } + return tensor; + } + + // Allocating path + let shape = tensor.layout().shape().clone(); + let dtype = tensor.dtype(); + let storage: &[E] = tensor.storage(); + + let result: Vec = match tensor.layout().contiguous_offsets() { + Some((start, end)) => storage[start..end].iter().map(|&x| op(x, scalar)).collect(), + None => StridedIter::new(tensor.layout()) + .map(|i| op(storage[i], scalar)) + .collect(), + }; + + make_tensor(result, shape, dtype) +} + +/// Helper to construct a tensor from result data. +fn make_tensor( + data: Vec, + shape: Shape, + dtype: DType, +) -> FlexTensor { + let bytes = Bytes::from_elems(data); + let layout = Layout::contiguous(shape); + FlexTensor::new(bytes, layout, dtype) +} + +/// Apply a binary operation element-wise to two integer tensors. +/// +/// Supports all integer dtypes: I64, I32, I16, I8, U64, U32, U16, U8. +pub fn int_binary_op(lhs: FlexTensor, rhs: FlexTensor, op: Op) -> FlexTensor +where + Op: Fn(i64, i64) -> i64 + Copy, +{ + debug_assert_eq!(lhs.dtype(), rhs.dtype(), "int_binary_op: dtype mismatch"); + + // Broadcast tensors to the same shape if needed + let (lhs, rhs) = crate::ops::expand::broadcast_binary(lhs, rhs); + + let dtype = lhs.dtype(); + + match dtype { + DType::I64 => binary_op_typed(lhs, &rhs, op), + DType::I32 => binary_op_typed(lhs, &rhs, |a: i32, b: i32| op(a as i64, b as i64) as i32), + DType::I16 => binary_op_typed(lhs, &rhs, |a: i16, b: i16| op(a as i64, b as i64) as i16), + DType::I8 => binary_op_typed(lhs, &rhs, |a: i8, b: i8| op(a as i64, b as i64) as i8), + // u64 values > i64::MAX wrap to negative i64. This is correct for + // add/sub/mul/bitwise (two's complement). Div/rem are handled at the call site. + DType::U64 => binary_op_typed(lhs, &rhs, |a: u64, b: u64| op(a as i64, b as i64) as u64), + DType::U32 => binary_op_typed(lhs, &rhs, |a: u32, b: u32| op(a as i64, b as i64) as u32), + DType::U16 => binary_op_typed(lhs, &rhs, |a: u16, b: u16| op(a as i64, b as i64) as u16), + DType::U8 => binary_op_typed(lhs, &rhs, |a: u8, b: u8| op(a as i64, b as i64) as u8), + _ => panic!("int_binary_op: unsupported dtype {:?}", dtype), + } +} + +/// Apply a scalar operation to each element of an integer tensor. +/// Note: scalar is truncated to target dtype (matches PyTorch). +pub fn int_scalar_op(tensor: FlexTensor, scalar: i64, op: Op) -> FlexTensor +where + Op: Fn(i64, i64) -> i64 + Copy, +{ + let dtype = tensor.dtype(); + + match dtype { + DType::I64 => scalar_op_typed(tensor, scalar, op), + DType::I32 => scalar_op_typed(tensor, scalar as i32, |a: i32, b: i32| { + op(a as i64, b as i64) as i32 + }), + DType::I16 => scalar_op_typed(tensor, scalar as i16, |a: i16, b: i16| { + op(a as i64, b as i64) as i16 + }), + DType::I8 => scalar_op_typed(tensor, scalar as i8, |a: i8, b: i8| { + op(a as i64, b as i64) as i8 + }), + DType::U64 => scalar_op_typed(tensor, scalar as u64, |a: u64, b: u64| { + op(a as i64, b as i64) as u64 + }), + DType::U32 => scalar_op_typed(tensor, scalar as u32, |a: u32, b: u32| { + op(a as i64, b as i64) as u32 + }), + DType::U16 => scalar_op_typed(tensor, scalar as u16, |a: u16, b: u16| { + op(a as i64, b as i64) as u16 + }), + DType::U8 => scalar_op_typed(tensor, scalar as u8, |a: u8, b: u8| { + op(a as i64, b as i64) as u8 + }), + _ => panic!("int_scalar_op: unsupported dtype {:?}", dtype), + } +} + +// Tests kept here exercise flex-specific behavior of `binary_op` / +// `scalar_op`: non-contiguous (transposed/narrowed/permuted) strides, +// flex f16/bf16 half-precision storage paths, and broadcast patterns +// that probe the flex layout system. Plain contiguous add/sub/mul/div +// and scalar-op smoke tests have been dropped in favor of the +// equivalent coverage in burn-backend-tests, which exercises every +// backend. When adding new tests, keep them here only if they probe +// flex-internal dispatch; otherwise add them to +// crates/burn-backend-tests/tests/tensor/float/ops/. +#[cfg(test)] +#[allow(clippy::needless_range_loop)] +mod tests { + use super::*; + use alloc::vec; + use burn_backend::{TensorData, Tolerance}; + + // =================== + // F16 tests + // =================== + + #[test] + fn test_binary_add_f16() { + let a_vals: Vec = vec![1.0, 2.0, 3.0, 4.0] + .into_iter() + .map(f16::from_f32) + .collect(); + let b_vals: Vec = vec![5.0, 6.0, 7.0, 8.0] + .into_iter() + .map(f16::from_f32) + .collect(); + + let a = FlexTensor::from_data(TensorData::new(a_vals, vec![2, 2])); + let b = FlexTensor::from_data(TensorData::new(b_vals, vec![2, 2])); + + let result = binary_op(a, b, |x, y| x + y, |x, y| x + y, None); + let expected: Vec = vec![6.0, 8.0, 10.0, 12.0] + .into_iter() + .map(f16::from_f32) + .collect(); + result.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![2, 2]), + Tolerance::absolute(f16::from_f32(0.01)), + ); + } + + #[test] + fn test_binary_mul_f16() { + let a_vals: Vec = vec![2.0, 3.0, 4.0, 5.0] + .into_iter() + .map(f16::from_f32) + .collect(); + let b_vals: Vec = vec![1.0, 2.0, 3.0, 4.0] + .into_iter() + .map(f16::from_f32) + .collect(); + + let a = FlexTensor::from_data(TensorData::new(a_vals, vec![2, 2])); + let b = FlexTensor::from_data(TensorData::new(b_vals, vec![2, 2])); + + let result = binary_op(a, b, |x, y| x * y, |x, y| x * y, None); + let expected: Vec = vec![2.0, 6.0, 12.0, 20.0] + .into_iter() + .map(f16::from_f32) + .collect(); + result.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![2, 2]), + Tolerance::absolute(f16::from_f32(0.01)), + ); + } + + #[test] + fn test_binary_f16_transposed() { + let a_vals: Vec = vec![1.0, 2.0, 3.0, 4.0] + .into_iter() + .map(f16::from_f32) + .collect(); + let b_vals: Vec = vec![10.0, 20.0, 30.0, 40.0] + .into_iter() + .map(f16::from_f32) + .collect(); + + let a = FlexTensor::from_data(TensorData::new(a_vals, vec![2, 2])).transpose(0, 1); + let b = FlexTensor::from_data(TensorData::new(b_vals, vec![2, 2])).transpose(0, 1); + + // a_t = [[1,3], [2,4]], b_t = [[10,30], [20,40]] + // result = [[11,33], [22,44]] + let result = binary_op(a, b, |x, y| x + y, |x, y| x + y, None); + let expected: Vec = vec![11.0, 33.0, 22.0, 44.0] + .into_iter() + .map(f16::from_f32) + .collect(); + result.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![2, 2]), + Tolerance::absolute(f16::from_f32(0.1)), + ); + } + + #[test] + fn test_scalar_f16() { + let a_vals: Vec = vec![1.0, 2.0, 3.0].into_iter().map(f16::from_f32).collect(); + let a = FlexTensor::from_data(TensorData::new(a_vals, vec![3])); + + let result = scalar_op(a, 10.0, |x, y| x + y, |x, y| x + y); + let expected: Vec = vec![11.0, 12.0, 13.0] + .into_iter() + .map(f16::from_f32) + .collect(); + result.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![3]), + Tolerance::absolute(f16::from_f32(0.01)), + ); + } + + // =================== + // BF16 tests + // =================== + + #[test] + fn test_binary_add_bf16() { + let a_vals: Vec = vec![1.0, 2.0, 3.0, 4.0] + .into_iter() + .map(bf16::from_f32) + .collect(); + let b_vals: Vec = vec![5.0, 6.0, 7.0, 8.0] + .into_iter() + .map(bf16::from_f32) + .collect(); + + let a = FlexTensor::from_data(TensorData::new(a_vals, vec![2, 2])); + let b = FlexTensor::from_data(TensorData::new(b_vals, vec![2, 2])); + + let result = binary_op(a, b, |x, y| x + y, |x, y| x + y, None); + let expected: Vec = vec![6.0, 8.0, 10.0, 12.0] + .into_iter() + .map(bf16::from_f32) + .collect(); + result.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![2, 2]), + Tolerance::absolute(bf16::from_f32(0.1)), + ); + } + + #[test] + fn test_binary_mul_bf16() { + let a_vals: Vec = vec![2.0, 3.0, 4.0, 5.0] + .into_iter() + .map(bf16::from_f32) + .collect(); + let b_vals: Vec = vec![1.0, 2.0, 3.0, 4.0] + .into_iter() + .map(bf16::from_f32) + .collect(); + + let a = FlexTensor::from_data(TensorData::new(a_vals, vec![2, 2])); + let b = FlexTensor::from_data(TensorData::new(b_vals, vec![2, 2])); + + let result = binary_op(a, b, |x, y| x * y, |x, y| x * y, None); + let expected: Vec = vec![2.0, 6.0, 12.0, 20.0] + .into_iter() + .map(bf16::from_f32) + .collect(); + result.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![2, 2]), + Tolerance::absolute(bf16::from_f32(0.1)), + ); + } + + #[test] + fn test_binary_bf16_transposed() { + let a_vals: Vec = vec![1.0, 2.0, 3.0, 4.0] + .into_iter() + .map(bf16::from_f32) + .collect(); + let b_vals: Vec = vec![10.0, 20.0, 30.0, 40.0] + .into_iter() + .map(bf16::from_f32) + .collect(); + + let a = FlexTensor::from_data(TensorData::new(a_vals, vec![2, 2])).transpose(0, 1); + let b = FlexTensor::from_data(TensorData::new(b_vals, vec![2, 2])).transpose(0, 1); + + // a_t = [[1,3], [2,4]], b_t = [[10,30], [20,40]] + // result = [[11,33], [22,44]] + let result = binary_op(a, b, |x, y| x + y, |x, y| x + y, None); + let expected: Vec = vec![11.0, 33.0, 22.0, 44.0] + .into_iter() + .map(bf16::from_f32) + .collect(); + result.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![2, 2]), + Tolerance::absolute(bf16::from_f32(0.5)), + ); + } + + #[test] + fn test_scalar_bf16() { + let a_vals: Vec = vec![1.0, 2.0, 3.0] + .into_iter() + .map(bf16::from_f32) + .collect(); + let a = FlexTensor::from_data(TensorData::new(a_vals, vec![3])); + + let result = scalar_op(a, 10.0, |x, y| x + y, |x, y| x + y); + let expected: Vec = vec![11.0, 12.0, 13.0] + .into_iter() + .map(bf16::from_f32) + .collect(); + result.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![3]), + Tolerance::absolute(bf16::from_f32(0.1)), + ); + } + + #[test] + fn test_scalar_f16_non_representable() { + // 0.1 is not exactly representable in f16; verify the scalar is rounded + // to f16 precision before the op (matching dtype semantics). + let a_vals: Vec = vec![1.0, 2.0, 3.0].into_iter().map(f16::from_f32).collect(); + let a = FlexTensor::from_data(TensorData::new(a_vals, vec![3])); + + let s_f16 = f16::from_f32(0.1); + let result = scalar_op(a, 0.1, |x, y| x * y, |x, y| x * y); + let expected: Vec = vec![1.0, 2.0, 3.0] + .into_iter() + .map(|v| f16::from_f32(f16::from_f32(v).to_f32() * s_f16.to_f32())) + .collect(); + result.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![3]), + Tolerance::absolute(f16::from_f32(0.001)), + ); + } + + #[test] + fn test_scalar_bf16_non_representable() { + // 1.1 is not exactly representable in bf16; verify dtype rounding. + let a_vals: Vec = vec![1.0, 2.0, 3.0] + .into_iter() + .map(bf16::from_f32) + .collect(); + let a = FlexTensor::from_data(TensorData::new(a_vals, vec![3])); + + let s_bf16 = bf16::from_f32(1.1); + let result = scalar_op(a, 1.1, |x, y| x * y, |x, y| x * y); + let expected: Vec = vec![1.0, 2.0, 3.0] + .into_iter() + .map(|v| bf16::from_f32(bf16::from_f32(v).to_f32() * s_bf16.to_f32())) + .collect(); + result.into_data().assert_approx_eq::( + &TensorData::new(expected, vec![3]), + Tolerance::absolute(bf16::from_f32(0.01)), + ); + } + + // ============================================================================ + // Broadcast binary-op fast paths + // ============================================================================ + + /// Shared-row broadcast: 1-D gamma reshaped + expanded, with the + /// size-1 outer dim exemption in play. + #[test] + fn test_binary_shared_row_broadcast_f32() { + let a = FlexTensor::from_data(TensorData::new( + vec![ + 1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, + ], + vec![1, 3, 4], + )); + // 1D gamma broadcast over rows. + let gamma = + FlexTensor::from_data(TensorData::new(vec![10.0f32, 20.0, 30.0, 40.0], vec![4])); + let gamma_unsqueezed = gamma.reshape(Shape::from(vec![1, 1, 4])); + let result = binary_op( + a, + gamma_unsqueezed, + |a, b| a * b, + |a, b| a * b, + Some(BinaryOp::Mul), + ); + let data = result.into_data(); + let expected = vec![ + 10.0f32, 40.0, 90.0, 160.0, 50.0, 120.0, 210.0, 320.0, 90.0, 200.0, 330.0, 480.0, + ]; + assert_eq!(data.as_slice::().unwrap(), expected.as_slice()); + } + + /// Per-row scalar broadcast: `[1, 3, 4] - mean_dim(-1)` shape, + /// rhs expands to strides `[3, 1, 0]`. + #[test] + fn test_binary_per_row_scalar_broadcast_f32() { + let a = FlexTensor::from_data(TensorData::new( + vec![ + 1.0f32, 2.0, 3.0, 4.0, 10.0, 20.0, 30.0, 40.0, 100.0, 200.0, 300.0, 400.0, + ], + vec![1, 3, 4], + )); + // Scalars shaped like `mean_dim(-1)` output. + let mean = FlexTensor::from_data(TensorData::new(vec![2.5f32, 25.0, 250.0], vec![1, 3, 1])); + let result = binary_op(a, mean, |a, b| a - b, |a, b| a - b, Some(BinaryOp::Sub)); + let data = result.into_data(); + let expected = vec![ + -1.5f32, -0.5, 0.5, 1.5, -15.0, -5.0, 5.0, 15.0, -150.0, -50.0, 50.0, 150.0, + ]; + assert_eq!(data.as_slice::().unwrap(), expected.as_slice()); + } + + /// Non-contig lhs + shared-row broadcast: must materialize lhs + /// contiguous then dispatch to the shared-row kernel. + #[test] + fn test_binary_permuted_lhs_broadcast_rhs() { + let a = FlexTensor::from_data(TensorData::new( + (0..24).map(|i| i as f32).collect::>(), + vec![2, 3, 4], + )); + let a_permuted = a.transpose(1, 2); // shape [2, 4, 3] + let gamma = FlexTensor::from_data(TensorData::new(vec![1.0f32, 10.0, 100.0], vec![3])); + let gamma_expanded = gamma.reshape(Shape::from(vec![1, 1, 3])); + + let result = binary_op( + a_permuted, + gamma_expanded, + |a, b| a * b, + |a, b| a * b, + Some(BinaryOp::Mul), + ); + + // Compare against a naive reference computed on the permuted + // values. + let reference: Vec = { + // original[b, r, c] = b*12 + r*4 + c + // permuted[b, c, r] = original[b, r, c] + // result[b, c, r] = permuted[b, c, r] * gamma_row[r] + let gamma_vals = [1.0f32, 10.0, 100.0]; + let mut out = Vec::with_capacity(24); + for b in 0..2 { + for c in 0..4 { + for r in 0..3 { + let orig_val = (b * 12 + r * 4 + c) as f32; + out.push(orig_val * gamma_vals[r]); + } + } + } + out + }; + + let data = result.into_data(); + assert_eq!(data.as_slice::().unwrap(), reference.as_slice()); + } + + /// Non-contig lhs + per-row-scalar broadcast-sub. + #[test] + fn test_binary_permuted_lhs_per_row_scalar_sub() { + let a = FlexTensor::from_data(TensorData::new( + (1..=24).map(|i| i as f32).collect::>(), + vec![2, 3, 4], + )); + let a_permuted = a.transpose(1, 2); // shape [2, 4, 3] + + // Per-row scalar in the permuted layout: shape [2, 4, 1]. + let mean = FlexTensor::from_data(TensorData::new( + (0..8).map(|i| i as f32).collect::>(), + vec![2, 4, 1], + )); + + let result = binary_op( + a_permuted, + mean, + |a, b| a - b, + |a, b| a - b, + Some(BinaryOp::Sub), + ); + + // Reference computation. + let reference: Vec = { + let mut out = Vec::with_capacity(24); + for b in 0..2 { + for c in 0..4 { + let mean_val = (b * 4 + c) as f32; + for r in 0..3 { + let orig_val = (b * 12 + r * 4 + c + 1) as f32; + out.push(orig_val - mean_val); + } + } + } + out + }; + + let data = result.into_data(); + assert_eq!(data.as_slice::().unwrap(), reference.as_slice()); + } + + /// Exercise every `(op, pattern)` combination of the broadcast fast path: + /// Add/Sub/Mul/Div crossed with SharedRow and PerRowScalar. The existing + /// targeted tests only cover a subset, so a sign error in + /// `div_shared_row_inplace_f32` or `add_per_row_scalar` would ship green. + #[test] + fn test_binary_broadcast_all_ops_and_patterns_f32() { + fn build_shared() -> (FlexTensor, FlexTensor) { + let a = FlexTensor::from_data(TensorData::new( + vec![4.0f32, 8.0, 12.0, 20.0, 30.0, 60.0], + vec![2, 3], + )); + let b = FlexTensor::from_data(TensorData::new(vec![2.0f32, 4.0, 3.0], vec![3])) + .reshape(Shape::from(vec![1, 3])); + (a, b) + } + fn build_perrow() -> (FlexTensor, FlexTensor) { + let a = FlexTensor::from_data(TensorData::new( + vec![4.0f32, 8.0, 12.0, 20.0, 30.0, 60.0], + vec![2, 3], + )); + let b = FlexTensor::from_data(TensorData::new(vec![2.0f32, 5.0], vec![2, 1])); + (a, b) + } + + let run = |name: &str, + build: fn() -> (FlexTensor, FlexTensor), + simd_op: BinaryOp, + op_fn: fn(f32, f32) -> f32, + expected: &[f32]| { + let (a, b) = build(); + let result = binary_op( + a, + b, + op_fn, + |x: f64, y: f64| op_fn(x as f32, y as f32) as f64, + Some(simd_op), + ); + let data = result.into_data(); + assert_eq!( + data.as_slice::().unwrap(), + expected, + "case {name} produced wrong values" + ); + }; + + // SharedRow expected: lhs[i][j] OP rhs[j] + run( + "shared_add", + build_shared, + BinaryOp::Add, + |a, b| a + b, + &[6.0, 12.0, 15.0, 22.0, 34.0, 63.0], + ); + run( + "shared_sub", + build_shared, + BinaryOp::Sub, + |a, b| a - b, + &[2.0, 4.0, 9.0, 18.0, 26.0, 57.0], + ); + run( + "shared_mul", + build_shared, + BinaryOp::Mul, + |a, b| a * b, + &[8.0, 32.0, 36.0, 40.0, 120.0, 180.0], + ); + run( + "shared_div", + build_shared, + BinaryOp::Div, + |a, b| a / b, + &[2.0, 2.0, 4.0, 10.0, 7.5, 20.0], + ); + // PerRowScalar expected: lhs[i][j] OP rhs[i] + run( + "perrow_add", + build_perrow, + BinaryOp::Add, + |a, b| a + b, + &[6.0, 10.0, 14.0, 25.0, 35.0, 65.0], + ); + run( + "perrow_sub", + build_perrow, + BinaryOp::Sub, + |a, b| a - b, + &[2.0, 6.0, 10.0, 15.0, 25.0, 55.0], + ); + run( + "perrow_mul", + build_perrow, + BinaryOp::Mul, + |a, b| a * b, + &[8.0, 16.0, 24.0, 100.0, 150.0, 300.0], + ); + run( + "perrow_div", + build_perrow, + BinaryOp::Div, + |a, b| a / b, + &[2.0, 4.0, 6.0, 4.0, 6.0, 12.0], + ); + } + + /// Non-unique lhs: `apply_broadcast_pattern_f32` takes the allocating + /// branch instead of writing in place. Clone the lhs so its Arc refcount + /// is > 1, then run a broadcast op and verify the result matches the + /// unique path. Without this test, a regression in the allocating branch + /// would only fire on shared-lhs call sites which are rare in bench code. + #[test] + fn test_binary_broadcast_non_unique_lhs_f32() { + let a = FlexTensor::from_data(TensorData::new( + vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], + vec![2, 3], + )); + let _keep_alive = a.clone(); // bump Arc refcount so lhs is shared + let b = FlexTensor::from_data(TensorData::new(vec![10.0f32, 20.0, 30.0], vec![3])) + .reshape(Shape::from(vec![1, 3])); + let result = binary_op(a, b, |a, b| a + b, |a, b| a + b, Some(BinaryOp::Add)); + let data = result.into_data(); + assert_eq!( + data.as_slice::().unwrap(), + &[11.0f32, 22.0, 33.0, 14.0, 25.0, 36.0] + ); + } + + /// Fully-broadcast scalar: rhs strides all 0, PerRowScalar with + /// empty outer walk, applies one scalar across the whole dst. + #[test] + fn test_binary_fully_broadcast_scalar_f32() { + let a = FlexTensor::from_data(TensorData::new( + (0..12).map(|i| i as f32).collect::>(), + vec![2, 2, 3], + )); + // 1-element tensor expanded to lhs's full shape. All strides + // become 0. + let scalar_tensor = FlexTensor::from_data(TensorData::new(vec![100.0f32], [1])); + let scalar_expanded = crate::ops::expand::expand(scalar_tensor, Shape::from(vec![2, 2, 3])); + // Sanity check: every stride is 0. + assert!(scalar_expanded.layout().strides().iter().all(|&s| s == 0)); + + let result = binary_op( + a, + scalar_expanded, + |a, b| a + b, + |a, b| a + b, + Some(BinaryOp::Add), + ); + + let expected: Vec = (0..12).map(|i| i as f32 + 100.0).collect(); + let data = result.into_data(); + assert_eq!(data.as_slice::().unwrap(), expected.as_slice()); + } +} diff --git a/crates/burn-flex/src/ops/bool.rs b/crates/burn-flex/src/ops/bool.rs new file mode 100644 index 0000000..753f5ec --- /dev/null +++ b/crates/burn-flex/src/ops/bool.rs @@ -0,0 +1,559 @@ +//! Bool tensor operations for the Flex backend. + +use alloc::vec; +use alloc::vec::Vec; +use burn_backend::{ + DType, ExecutionError, TensorData, + ops::{BoolTensorOps, IntTensorOps}, + tensor::{BoolTensor, Device, FloatTensor, IntTensor}, +}; +use burn_std::{Bytes, FloatDType, IntDType, Shape, Slice, bf16, f16}; + +use crate::{Flex, FlexTensor, Layout}; + +impl BoolTensorOps for Flex { + fn bool_from_data(data: TensorData, _device: &Device) -> BoolTensor { + FlexTensor::from_data(data) + } + + async fn bool_into_data(tensor: BoolTensor) -> Result { + Ok(tensor.into_data()) + } + + fn bool_to_device(tensor: BoolTensor, _device: &Device) -> BoolTensor { + tensor + } + + fn bool_cat(tensors: Vec>, dim: usize) -> BoolTensor { + crate::ops::cat::cat(tensors, dim) + } + + fn bool_reshape(tensor: BoolTensor, shape: Shape) -> BoolTensor { + tensor.reshape(shape) + } + + fn bool_slice(tensor: BoolTensor, slices: &[Slice]) -> BoolTensor { + crate::ops::slice::slice(tensor, slices) + } + + fn bool_empty( + shape: Shape, + _device: &Device, + dtype: burn_std::BoolDType, + ) -> BoolTensor { + FlexTensor::empty(shape, DType::from(dtype)) + } + + fn bool_slice_assign( + tensor: BoolTensor, + slices: &[Slice], + value: BoolTensor, + ) -> BoolTensor { + crate::ops::slice::slice_assign(tensor, slices, value) + } + + fn bool_into_int(tensor: BoolTensor, out_dtype: burn_std::IntDType) -> IntTensor { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let out_dt = DType::from(out_dtype); + let bools = tensor.bytes(); + + macro_rules! convert { + ($int_ty:ty) => {{ + let data: Vec<$int_ty> = + bools.iter().map(|&x| if x != 0 { 1 } else { 0 }).collect(); + FlexTensor::new(Bytes::from_elems(data), Layout::contiguous(shape), out_dt) + }}; + } + + match out_dtype { + IntDType::I64 => convert!(i64), + IntDType::I32 => convert!(i32), + IntDType::I16 => convert!(i16), + IntDType::I8 => convert!(i8), + IntDType::U64 => convert!(u64), + IntDType::U32 => convert!(u32), + IntDType::U16 => convert!(u16), + IntDType::U8 => convert!(u8), + } + } + + fn bool_into_float( + tensor: BoolTensor, + out_dtype: burn_std::FloatDType, + ) -> FloatTensor { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let out_dt = DType::from(out_dtype); + let bools = tensor.bytes(); + + match out_dtype { + FloatDType::F64 => { + let data: Vec = bools + .iter() + .map(|&x| if x != 0 { 1.0 } else { 0.0 }) + .collect(); + FlexTensor::new(Bytes::from_elems(data), Layout::contiguous(shape), out_dt) + } + FloatDType::F32 | FloatDType::Flex32 => { + let data: Vec = bools + .iter() + .map(|&x| if x != 0 { 1.0 } else { 0.0 }) + .collect(); + FlexTensor::new(Bytes::from_elems(data), Layout::contiguous(shape), out_dt) + } + FloatDType::F16 => { + let one = f16::from_f32(1.0); + let zero = f16::from_f32(0.0); + let data: Vec = bools + .iter() + .map(|&x| if x != 0 { one } else { zero }) + .collect(); + FlexTensor::new(Bytes::from_elems(data), Layout::contiguous(shape), out_dt) + } + FloatDType::BF16 => { + let one = bf16::from_f32(1.0); + let zero = bf16::from_f32(0.0); + let data: Vec = bools + .iter() + .map(|&x| if x != 0 { one } else { zero }) + .collect(); + FlexTensor::new(Bytes::from_elems(data), Layout::contiguous(shape), out_dt) + } + } + } + + fn bool_swap_dims(tensor: BoolTensor, dim1: usize, dim2: usize) -> BoolTensor { + tensor.transpose(dim1, dim2) + } + + fn bool_permute(tensor: BoolTensor, axes: &[usize]) -> BoolTensor { + tensor.permute(axes) + } + + fn bool_flip(tensor: BoolTensor, axes: &[usize]) -> BoolTensor { + crate::ops::flip::flip(tensor, axes) + } + + fn bool_equal(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + use crate::strided_index::StridedIter; + + // Broadcast to a common shape before comparing. The contiguous fast + // path below uses `zip`, which silently truncates to the shorter + // operand; and the output shape is taken from lhs, so mismatched + // operands would otherwise produce a result vec shorter than the + // output layout claims. + let (lhs, rhs) = crate::ops::expand::broadcast_binary(lhs, rhs); + + let out_dtype = burn_std::BoolDType::from(lhs.dtype()); + let shape = lhs.layout().shape().clone(); + let lhs_storage: &[u8] = lhs.bytes(); + let rhs_storage: &[u8] = rhs.bytes(); + + let result: Vec = match ( + lhs.layout().contiguous_offsets(), + rhs.layout().contiguous_offsets(), + ) { + (Some((l_start, l_end)), Some((r_start, r_end))) => { + let l_slice = &lhs_storage[l_start..l_end]; + let r_slice = &rhs_storage[r_start..r_end]; + l_slice + .iter() + .zip(r_slice) + .map(|(&a, &b)| (a == b) as u8) + .collect() + } + _ => { + let lhs_iter = StridedIter::new(lhs.layout()); + let rhs_iter = StridedIter::new(rhs.layout()); + lhs_iter + .zip(rhs_iter) + .map(|(li, ri)| (lhs_storage[li] == rhs_storage[ri]) as u8) + .collect() + } + }; + + crate::ops::comparison::make_bool_tensor(result, shape, out_dtype) + } + + fn bool_not(mut tensor: BoolTensor) -> BoolTensor { + use crate::strided_index::StridedIter; + + debug_assert!( + matches!( + tensor.dtype(), + DType::Bool(burn_std::BoolStore::Native | burn_std::BoolStore::U8) + ), + "bool_not: only Bool(Native) and Bool(U8) are supported, got {:?}", + tensor.dtype() + ); + + // Fast path: in-place for unique, contiguous tensors at offset 0. This + // preserves the input tensor's dtype tag implicitly (the in-place SIMD + // ops flip bytes without touching the dtype tag). + if tensor.is_unique() + && tensor.layout().is_contiguous() + && tensor.layout().start_offset() == 0 + { + let storage = tensor.storage_mut::(); + crate::simd::bool_not_inplace_u8(storage); + return tensor; + } + + // Allocating path for shared, non-contiguous, or offset tensors: + // preserve the input's bool dtype for the new tensor. + let out_dtype = burn_std::BoolDType::from(tensor.dtype()); + let shape = tensor.layout().shape().clone(); + let storage: &[u8] = tensor.bytes(); + + let result: Vec = match tensor.layout().contiguous_offsets() { + Some((start, end)) => { + let slice = &storage[start..end]; + let mut out = vec![0u8; slice.len()]; + crate::simd::bool_not_u8(slice, &mut out); + out + } + None => StridedIter::new(tensor.layout()) + .map(|idx| (storage[idx] == 0) as u8) + .collect(), + }; + + crate::ops::comparison::make_bool_tensor(result, shape, out_dtype) + } + + fn bool_and(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + bool_binary_op_simd(lhs, rhs, BoolBinaryOp::And) + } + + fn bool_or(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + bool_binary_op_simd(lhs, rhs, BoolBinaryOp::Or) + } + + fn bool_xor(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + bool_binary_op_simd(lhs, rhs, BoolBinaryOp::Xor) + } + + fn bool_expand(tensor: BoolTensor, shape: Shape) -> BoolTensor { + crate::ops::expand::expand(tensor, shape) + } + + // Missing methods + fn bool_zeros( + shape: Shape, + device: &Device, + dtype: burn_std::BoolDType, + ) -> BoolTensor { + Self::bool_empty(shape, device, dtype) + } + + fn bool_ones( + shape: Shape, + _device: &Device, + dtype: burn_std::BoolDType, + ) -> BoolTensor { + let num_elements = shape.num_elements(); + let data = vec![1u8; num_elements]; + crate::ops::comparison::make_bool_tensor(data, shape, dtype) + } + + fn bool_mask_where( + tensor: BoolTensor, + mask: BoolTensor, + value: BoolTensor, + ) -> BoolTensor { + crate::ops::mask::mask_where_bool(tensor, mask, value) + } + + fn bool_mask_fill( + tensor: BoolTensor, + mask: BoolTensor, + value: burn_backend::Scalar, + ) -> BoolTensor { + let value: bool = value.elem(); + crate::ops::mask::mask_fill_bool(tensor, mask, value) + } + + fn bool_gather( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + ) -> BoolTensor { + crate::ops::gather_scatter::gather_bool(tensor, dim, indices) + } + + fn bool_scatter_or( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor { + crate::ops::gather_scatter::scatter_or(tensor, dim, indices, value) + } + + fn bool_equal_elem(lhs: BoolTensor, rhs: burn_backend::Scalar) -> BoolTensor { + use crate::strided_index::StridedIter; + + let out_dtype = burn_std::BoolDType::from(lhs.dtype()); + let shape = lhs.layout().shape().clone(); + let storage: &[u8] = lhs.bytes(); + let rhs_bool: bool = rhs.elem(); + let rhs_val = rhs_bool as u8; + + let result: Vec = match lhs.layout().contiguous_offsets() { + Some((start, end)) => storage[start..end] + .iter() + .map(|&v| (v == rhs_val) as u8) + .collect(), + None => StridedIter::new(lhs.layout()) + .map(|idx| (storage[idx] == rhs_val) as u8) + .collect(), + }; + + crate::ops::comparison::make_bool_tensor(result, shape, out_dtype) + } + + fn bool_unfold( + tensor: BoolTensor, + dim: usize, + size: usize, + step: usize, + ) -> BoolTensor { + crate::ops::unfold::unfold_bool(tensor, dim, size, step) + } + + fn bool_not_equal(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + let out_dtype = burn_std::BoolDType::from(lhs.dtype()); + crate::ops::comparison::bool_not_equal(lhs, rhs, out_dtype) + } + + fn bool_not_equal_elem(lhs: BoolTensor, rhs: burn_backend::Scalar) -> BoolTensor { + let out_dtype = burn_std::BoolDType::from(lhs.dtype()); + let rhs: bool = rhs.elem(); + crate::ops::comparison::bool_not_equal_elem(lhs, rhs, out_dtype) + } + + fn bool_any(tensor: BoolTensor) -> BoolTensor { + let out_dtype = burn_std::BoolDType::from(tensor.dtype()); + crate::ops::comparison::any_bool(tensor, out_dtype) + } + + fn bool_any_dim(tensor: BoolTensor, dim: usize) -> BoolTensor { + let out_dtype = burn_std::BoolDType::from(tensor.dtype()); + crate::ops::comparison::any_bool_dim(tensor, dim, out_dtype) + } + + fn bool_all(tensor: BoolTensor) -> BoolTensor { + let out_dtype = burn_std::BoolDType::from(tensor.dtype()); + crate::ops::comparison::all_bool(tensor, out_dtype) + } + + fn bool_all_dim(tensor: BoolTensor, dim: usize) -> BoolTensor { + let out_dtype = burn_std::BoolDType::from(tensor.dtype()); + crate::ops::comparison::all_bool_dim(tensor, dim, out_dtype) + } + + fn bool_select( + tensor: BoolTensor, + dim: usize, + indices: IntTensor, + ) -> BoolTensor { + crate::ops::gather_scatter::select::(tensor, dim, indices) + } + + fn bool_select_or( + tensor: BoolTensor, + dim: usize, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor { + let mut result = crate::ops::gather_scatter::select_add::(tensor, dim, indices, value); + // Clamp to 0/1: select_add sums u8 values, but bool OR saturates at 1 + let storage: &mut [u8] = result.storage_mut(); + for v in storage.iter_mut() { + if *v > 1 { + *v = 1; + } + } + result + } + + fn bool_transpose(tensor: BoolTensor) -> BoolTensor { + let ndims = tensor.layout().num_dims(); + if ndims < 2 { + return tensor; + } + tensor.transpose(ndims - 2, ndims - 1) + } + + fn bool_repeat_dim(tensor: BoolTensor, dim: usize, times: usize) -> BoolTensor { + crate::ops::repeat_dim::repeat_dim(tensor, dim, times) + } + + async fn bool_argwhere(tensor: BoolTensor, out_dtype: IntDType) -> IntTensor { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let ndims = shape.num_dims(); + let data: &[u8] = tensor.storage(); + let n = shape.num_elements(); + + let count = data[..n].iter().filter(|&&v| v != 0).count(); + let mut coords: Vec = Vec::with_capacity(count * ndims); + let strides = crate::layout::contiguous_strides_usize(&shape); + + for (flat_idx, &val) in data[..n].iter().enumerate() { + if val != 0 { + let mut remaining = flat_idx; + for &s in &strides { + coords.push((remaining / s) as isize); + remaining %= s; + } + } + } + + let out_shape = Shape::from(vec![count, ndims]); + let result = FlexTensor::new( + Bytes::from_elems(coords), + Layout::contiguous(out_shape), + crate::ops::INDEX_DTYPE, + ); + if result.dtype() != DType::from(out_dtype) { + Flex::int_cast(result, out_dtype) + } else { + result + } + } +} + +/// Boolean binary operation type. +#[derive(Clone, Copy)] +enum BoolBinaryOp { + And, + Or, + Xor, +} + +fn bool_binary_op_simd(lhs: FlexTensor, rhs: FlexTensor, op: BoolBinaryOp) -> FlexTensor { + use crate::strided_index::StridedIter; + + debug_assert_eq!(lhs.dtype(), rhs.dtype(), "bool_binary_op: dtype mismatch"); + + // Broadcast to a common shape before dispatching. The scalar/SIMD helpers + // below assume equal-length operands; without this, mismatched shapes + // either silently keep the lhs shape or OOB-panic inside the helpers. + let (mut lhs, mut rhs) = crate::ops::expand::broadcast_binary(lhs, rhs); + + // Preserve the input bool dtype (taken from lhs; rhs is assumed to match + // in dtype, checked above). + let out_dtype = burn_std::BoolDType::from(lhs.dtype()); + let shape = lhs.layout().shape().clone(); + let l_offsets = lhs.layout().contiguous_offsets(); + let r_offsets = rhs.layout().contiguous_offsets(); + + // Fast path 1: lhs is unique and contiguous at offset 0 -> in-place on lhs + if lhs.is_unique() + && let (Some((0, l_end)), Some((r_start, r_end))) = (l_offsets, r_offsets) + { + let rhs_storage: &[u8] = rhs.bytes(); + let r_slice = &rhs_storage[r_start..r_end]; + let lhs_storage: &mut [u8] = lhs.storage_mut(); + let l_slice = &mut lhs_storage[..l_end]; + + match op { + BoolBinaryOp::And => crate::simd::bool_and_inplace_u8(l_slice, r_slice), + BoolBinaryOp::Or => crate::simd::bool_or_inplace_u8(l_slice, r_slice), + BoolBinaryOp::Xor => crate::simd::bool_xor_inplace_u8(l_slice, r_slice), + } + return lhs; + } + + // Fast path 2: rhs is unique and contiguous at offset 0 -> in-place on rhs + // (And/Or/Xor are commutative, so we can swap operands) + if rhs.is_unique() + && let (Some((l_start, l_end)), Some((0, r_end))) = (l_offsets, r_offsets) + { + let lhs_storage: &[u8] = lhs.bytes(); + let l_slice = &lhs_storage[l_start..l_end]; + let rhs_storage: &mut [u8] = rhs.storage_mut(); + let r_slice = &mut rhs_storage[..r_end]; + + match op { + BoolBinaryOp::And => crate::simd::bool_and_inplace_u8(r_slice, l_slice), + BoolBinaryOp::Or => crate::simd::bool_or_inplace_u8(r_slice, l_slice), + BoolBinaryOp::Xor => crate::simd::bool_xor_inplace_u8(r_slice, l_slice), + } + return rhs; + } + + // Allocating path: neither tensor is suitable for in-place + let lhs_storage: &[u8] = lhs.bytes(); + let rhs_storage: &[u8] = rhs.bytes(); + + let result: Vec = match (l_offsets, r_offsets) { + (Some((l_start, l_end)), Some((r_start, r_end))) => { + let l_slice = &lhs_storage[l_start..l_end]; + let r_slice = &rhs_storage[r_start..r_end]; + let mut out = vec![0u8; l_slice.len()]; + match op { + BoolBinaryOp::And => crate::simd::bool_and_u8(l_slice, r_slice, &mut out), + BoolBinaryOp::Or => crate::simd::bool_or_u8(l_slice, r_slice, &mut out), + BoolBinaryOp::Xor => crate::simd::bool_xor_u8(l_slice, r_slice, &mut out), + } + out + } + _ => { + let lhs_iter = StridedIter::new(lhs.layout()); + let rhs_iter = StridedIter::new(rhs.layout()); + match op { + BoolBinaryOp::And => lhs_iter + .zip(rhs_iter) + .map(|(li, ri)| lhs_storage[li] & rhs_storage[ri]) + .collect(), + BoolBinaryOp::Or => lhs_iter + .zip(rhs_iter) + .map(|(li, ri)| lhs_storage[li] | rhs_storage[ri]) + .collect(), + BoolBinaryOp::Xor => lhs_iter + .zip(rhs_iter) + .map(|(li, ri)| lhs_storage[li] ^ rhs_storage[ri]) + .collect(), + } + } + }; + + crate::ops::comparison::make_bool_tensor(result, shape, out_dtype) +} + +// Tests kept here exercise flex-specific dtype storage selection via +// explicit IntDType/FloatDType. Plain bool ops, bool-to-int/float +// casts, and negative-stride (flipped) bool coverage have been migrated +// to crates/burn-backend-tests/tests/tensor/bool/ops/{logical,cast}.rs +// so they run against every backend. When adding new tests, keep them +// here only if they probe flex dtype dispatch; otherwise add them +// there. +#[cfg(test)] +mod tests { + use alloc::vec; + use burn_backend::TensorData; + use burn_backend::ops::BoolTensorOps; + use burn_std::{FloatDType, IntDType}; + + use crate::{Flex, FlexTensor}; + + #[test] + fn test_bool_into_int_u8() { + let t = FlexTensor::from_data(TensorData::from([true, false, true])); + let result = Flex::bool_into_int(t, IntDType::U8); + assert_eq!(result.dtype(), burn_backend::DType::U8); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![1u8, 0, 1]); + } + + #[test] + fn test_bool_into_float_f64() { + let t = FlexTensor::from_data(TensorData::from([true, false, true])); + let result = Flex::bool_into_float(t, FloatDType::F64); + assert_eq!(result.dtype(), burn_backend::DType::F64); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![1.0f64, 0.0, 1.0]); + } +} diff --git a/crates/burn-flex/src/ops/cat.rs b/crates/burn-flex/src/ops/cat.rs new file mode 100644 index 0000000..38177fc --- /dev/null +++ b/crates/burn-flex/src/ops/cat.rs @@ -0,0 +1,119 @@ +//! Concatenation operations for FlexTensor. + +use alloc::vec::Vec; +use burn_backend::{DType, Element}; +use burn_std::{Bytes, Shape, bf16, f16}; + +use crate::{FlexTensor, Layout}; + +/// Concatenate tensors along the given dimension. +pub fn cat(tensors: Vec, dim: usize) -> FlexTensor { + assert!(!tensors.is_empty(), "cat: cannot concatenate empty list"); + if tensors.len() == 1 { + return tensors.into_iter().next().unwrap(); + } + + let dtype = tensors[0].dtype(); + match dtype { + DType::F32 => cat_impl::(tensors, dim), + DType::F64 => cat_impl::(tensors, dim), + DType::F16 => cat_impl::(tensors, dim), + DType::BF16 => cat_impl::(tensors, dim), + DType::I64 => cat_impl::(tensors, dim), + DType::I32 => cat_impl::(tensors, dim), + DType::I16 => cat_impl::(tensors, dim), + DType::I8 => cat_impl::(tensors, dim), + DType::U64 => cat_impl::(tensors, dim), + DType::U32 => cat_impl::(tensors, dim), + DType::U16 => cat_impl::(tensors, dim), + DType::U8 | DType::Bool(_) => cat_impl::(tensors, dim), + _ => panic!("cat: unsupported dtype {:?}", dtype), + } +} + +fn cat_impl(tensors: Vec, dim: usize) -> FlexTensor { + let dtype = tensors[0].dtype(); + let first_shape = tensors[0].layout().shape(); + let ndims = first_shape.num_dims(); + + assert!( + dim < ndims, + "cat: dim {} out of bounds for {} dimensions", + dim, + ndims + ); + + // Compute output shape: sum along cat dim, others must match + let mut out_dims = first_shape.to_vec(); + out_dims[dim] = 0; + for t in &tensors { + assert_eq!( + t.dtype(), + dtype, + "cat: dtype mismatch: expected {:?}, got {:?}", + dtype, + t.dtype() + ); + let s = t.layout().shape(); + assert_eq!(s.num_dims(), ndims, "cat: dimension count mismatch"); + for (d, out_d) in out_dims.iter_mut().enumerate() { + if d == dim { + *out_d += s[d]; + } else { + assert_eq!( + s[d], first_shape[d], + "cat: shape mismatch at dim {d}: expected {}, got {}", + first_shape[d], s[d] + ); + } + } + } + + let out_shape = Shape::from(out_dims.clone()); + let total_elements: usize = out_shape.num_elements(); + + if total_elements == 0 { + let bytes = Bytes::from_elems::(Vec::new()); + return FlexTensor::new(bytes, Layout::contiguous(out_shape), dtype); + } + + let mut output: Vec = Vec::with_capacity(total_elements); + + // Fast path: dim 0, all contiguous + if dim == 0 + && tensors + .iter() + .all(|t| t.layout().contiguous_offsets().is_some()) + { + for t in &tensors { + let (start, end) = t.layout().contiguous_offsets().unwrap(); + let data: &[E] = t.storage(); + output.extend_from_slice(&data[start..end]); + } + } else { + // General path: iterate over outer/inner chunks + // outer_size = product of dims before `dim` + // inner_size = product of dims after `dim` + let outer_size: usize = out_dims[..dim].iter().product(); + let inner_size: usize = out_dims[dim + 1..].iter().product(); + + // Make all tensors contiguous for simple indexing + let contiguous: Vec = tensors.into_iter().map(|t| t.to_contiguous()).collect(); + + for outer in 0..outer_size { + for t in &contiguous { + let data: &[E] = t.storage(); + let t_dim_size = t.layout().shape()[dim]; + let t_start = t.layout().start_offset(); + // Each tensor's chunk for this outer index: + // offset = t_start + outer * t_dim_size * inner_size + let chunk_start = t_start + outer * t_dim_size * inner_size; + let chunk_len = t_dim_size * inner_size; + output.extend_from_slice(&data[chunk_start..chunk_start + chunk_len]); + } + } + } + + let bytes = Bytes::from_elems(output); + FlexTensor::new(bytes, Layout::contiguous(out_shape), dtype) +} diff --git a/crates/burn-flex/src/ops/comparison.rs b/crates/burn-flex/src/ops/comparison.rs new file mode 100644 index 0000000..4658480 --- /dev/null +++ b/crates/burn-flex/src/ops/comparison.rs @@ -0,0 +1,1041 @@ +//! Comparison operations returning boolean tensors. + +use alloc::boxed::Box; +#[cfg(feature = "simd")] +use alloc::vec; +use alloc::vec::Vec; +use burn_backend::{DType, Element}; +use burn_std::{BoolDType, BoolStore, Bytes, Shape, bf16, f16}; +use bytemuck::Pod; + +use crate::strided_index::StridedIter; +use crate::{FlexTensor, Layout}; + +use crate::simd; + +/// Comparison operation type for SIMD dispatch. +pub use simd::CmpOp as CompareOp; + +/// Compare two tensors element-wise, returning a boolean tensor with the +/// requested output dtype. +pub fn compare( + lhs: FlexTensor, + rhs: FlexTensor, + out_dtype: BoolDType, + f32_cmp: F32Cmp, + f64_cmp: F64Cmp, + simd_hint: Option, +) -> FlexTensor +where + F32Cmp: Fn(f32, f32) -> bool + Copy, + F64Cmp: Fn(f64, f64) -> bool + Copy, +{ + debug_assert_eq!(lhs.dtype(), rhs.dtype(), "compare: dtype mismatch"); + + // Broadcast to same shape if needed + let (lhs, rhs) = crate::ops::expand::broadcast_binary(lhs, rhs); + + let dtype = lhs.dtype(); + + match dtype { + DType::F32 => compare_f32(lhs, &rhs, out_dtype, f32_cmp, simd_hint), + DType::F64 => compare_typed(lhs, &rhs, out_dtype, f64_cmp), + DType::F16 => compare_typed(lhs, &rhs, out_dtype, |a: f16, b: f16| { + f32_cmp(a.to_f32(), b.to_f32()) + }), + DType::BF16 => compare_typed(lhs, &rhs, out_dtype, |a: bf16, b: bf16| { + f32_cmp(a.to_f32(), b.to_f32()) + }), + _ => panic!("compare: unsupported dtype {:?}", dtype), + } +} + +/// Specialized comparison for f32 with SIMD fast path. +#[cfg(feature = "simd")] +fn compare_f32( + lhs: FlexTensor, + rhs: &FlexTensor, + out_dtype: BoolDType, + cmp: Cmp, + simd_hint: Option, +) -> FlexTensor +where + Cmp: Fn(f32, f32) -> bool, +{ + // SIMD fast path: both tensors contiguous + if let (Some((l_start, l_end)), Some((r_start, r_end))) = ( + lhs.layout().contiguous_offsets(), + rhs.layout().contiguous_offsets(), + ) && let Some(simd_op) = simd_hint + { + let shape = lhs.layout().shape().clone(); + let lhs_storage: &[f32] = lhs.storage(); + let rhs_storage: &[f32] = rhs.storage(); + + let l_slice = &lhs_storage[l_start..l_end]; + let r_slice = &rhs_storage[r_start..r_end]; + + let mut result = vec![0u8; l_slice.len()]; + simd::cmp_f32(l_slice, r_slice, &mut result, simd_op); + + return make_bool_tensor(result, shape, out_dtype); + } + + // Optimized broadcast path for outer-product style broadcasting + // Pattern: [N, 1] vs [1, M] -> [N, M] where one has stride 0 in inner dim + if lhs.layout().num_dims() == 2 + && let Some(simd_op) = simd_hint + && let Some((result, shape)) = try_broadcast_cmp_f32(&lhs, rhs, simd_op) + { + return make_bool_tensor(result, shape, out_dtype); + } + + // Fallback to generic path + compare_typed(lhs, rhs, out_dtype, cmp) +} + +/// Try optimized outer-product style broadcast comparison. +/// Returns Some((result, shape)) if the pattern matches. +#[cfg(feature = "simd")] +fn try_broadcast_cmp_f32( + lhs: &FlexTensor, + rhs: &FlexTensor, + op: simd::CmpOp, +) -> Option<(Vec, Shape)> { + let lhs_strides = lhs.layout().strides(); + let rhs_strides = rhs.layout().strides(); + let shape = lhs.layout().shape().clone(); + let [rows, cols] = shape[..] else { + return None; + }; + + // Pattern 1: lhs has stride 0 in dim 1 (column broadcast), rhs contiguous + // lhs[i,j] = lhs_data[i*stride], rhs[i,j] = rhs_data[i*cols + j] + if lhs_strides[1] == 0 && rhs_strides == [cols as isize, 1] { + let lhs_storage: &[f32] = lhs.storage(); + let rhs_storage: &[f32] = rhs.storage(); + let l_offset = lhs.layout().start_offset() as isize; + let l_stride = lhs_strides[0]; + let r_offset = rhs.layout().start_offset(); + + let mut result = vec![0u8; rows * cols]; + for row in 0..rows { + let a_val = lhs_storage[(l_offset + row as isize * l_stride) as usize]; + let r_row_start = r_offset + row * cols; + let r_slice = &rhs_storage[r_row_start..r_row_start + cols]; + let out_start = row * cols; + simd::cmp_scalar_f32( + r_slice, + a_val, + &mut result[out_start..out_start + cols], + swap_cmp_op(op), + ); + } + return Some((result, shape)); + } + + // Pattern 2: rhs has stride 0 in dim 0 (row broadcast), lhs contiguous + // lhs[i,j] = lhs_data[i*cols + j], rhs[i,j] = rhs_data[j*stride] + if rhs_strides[0] == 0 && lhs_strides == [cols as isize, 1] { + let lhs_storage: &[f32] = lhs.storage(); + let rhs_storage: &[f32] = rhs.storage(); + let l_offset = lhs.layout().start_offset(); + let r_offset = rhs.layout().start_offset() as isize; + let r_stride = rhs_strides[1]; + + // Build the broadcast rhs values once + let rhs_row: Vec = (0..cols) + .map(|j| rhs_storage[(r_offset + j as isize * r_stride) as usize]) + .collect(); + + let mut result = vec![0u8; rows * cols]; + for row in 0..rows { + let l_row_start = l_offset + row * cols; + let l_slice = &lhs_storage[l_row_start..l_row_start + cols]; + let out_start = row * cols; + // Compare row with broadcast values + for (j, (&lv, &rv)) in l_slice.iter().zip(rhs_row.iter()).enumerate() { + result[out_start + j] = match op { + simd::CmpOp::Gt => (lv > rv) as u8, + simd::CmpOp::Ge => (lv >= rv) as u8, + simd::CmpOp::Lt => (lv < rv) as u8, + simd::CmpOp::Le => (lv <= rv) as u8, + simd::CmpOp::Eq => (lv == rv) as u8, + simd::CmpOp::Ne => (lv != rv) as u8, + }; + } + } + return Some((result, shape)); + } + + // Pattern 3: Outer product - lhs stride 0 in dim 1, rhs stride 0 in dim 0 + // This is the [N,1] vs [1,M] case + if lhs_strides[1] == 0 && rhs_strides[0] == 0 { + let lhs_storage: &[f32] = lhs.storage(); + let rhs_storage: &[f32] = rhs.storage(); + let l_offset = lhs.layout().start_offset() as isize; + let l_stride = lhs_strides[0]; + let r_offset = rhs.layout().start_offset() as isize; + let r_stride = rhs_strides[1]; + + // Build the broadcast rhs row once + let rhs_row: Vec = (0..cols) + .map(|j| rhs_storage[(r_offset + j as isize * r_stride) as usize]) + .collect(); + + let mut result = vec![0u8; rows * cols]; + for row in 0..rows { + let a_val = lhs_storage[(l_offset + row as isize * l_stride) as usize]; + let out_start = row * cols; + simd::cmp_scalar_f32( + &rhs_row, + a_val, + &mut result[out_start..out_start + cols], + swap_cmp_op(op), + ); + } + return Some((result, shape)); + } + + None +} + +/// Swap comparison operation for reversed operand order. +#[cfg(feature = "simd")] +fn swap_cmp_op(op: simd::CmpOp) -> simd::CmpOp { + match op { + simd::CmpOp::Gt => simd::CmpOp::Lt, // a > b becomes b < a + simd::CmpOp::Ge => simd::CmpOp::Le, + simd::CmpOp::Lt => simd::CmpOp::Gt, + simd::CmpOp::Le => simd::CmpOp::Ge, + simd::CmpOp::Eq => simd::CmpOp::Eq, // symmetric + simd::CmpOp::Ne => simd::CmpOp::Ne, + } +} + +/// Fallback when SIMD is disabled. +#[cfg(not(feature = "simd"))] +fn compare_f32( + lhs: FlexTensor, + rhs: &FlexTensor, + out_dtype: BoolDType, + cmp: Cmp, + _simd_hint: Option, +) -> FlexTensor +where + Cmp: Fn(f32, f32) -> bool, +{ + compare_typed(lhs, rhs, out_dtype, cmp) +} + +/// Compare tensor with scalar, returning a boolean tensor with the requested +/// output dtype. +pub fn compare_elem( + lhs: FlexTensor, + rhs: f64, + out_dtype: BoolDType, + f32_cmp: F32Cmp, + f64_cmp: F64Cmp, + simd_hint: Option, +) -> FlexTensor +where + F32Cmp: Fn(f32, f32) -> bool + Copy, + F64Cmp: Fn(f64, f64) -> bool + Copy, +{ + let dtype = lhs.dtype(); + + match dtype { + DType::F32 => compare_elem_f32(lhs, rhs as f32, out_dtype, f32_cmp, simd_hint), + DType::F64 => compare_elem_typed(lhs, rhs, out_dtype, f64_cmp), + DType::F16 => { + let scalar = f16::from_f64(rhs); + compare_elem_typed(lhs, scalar, out_dtype, |a: f16, b: f16| { + f32_cmp(a.to_f32(), b.to_f32()) + }) + } + DType::BF16 => { + let scalar = bf16::from_f64(rhs); + compare_elem_typed(lhs, scalar, out_dtype, |a: bf16, b: bf16| { + f32_cmp(a.to_f32(), b.to_f32()) + }) + } + _ => panic!("compare_elem: unsupported dtype {:?}", dtype), + } +} + +/// Specialized scalar comparison for f32 with SIMD fast path. +#[cfg(feature = "simd")] +fn compare_elem_f32( + lhs: FlexTensor, + rhs: f32, + out_dtype: BoolDType, + cmp: Cmp, + simd_hint: Option, +) -> FlexTensor +where + Cmp: Fn(f32, f32) -> bool, +{ + // SIMD fast path: tensor is contiguous + if let Some((start, end)) = lhs.layout().contiguous_offsets() + && let Some(simd_op) = simd_hint + { + let shape = lhs.layout().shape().clone(); + let lhs_storage: &[f32] = lhs.storage(); + let l_slice = &lhs_storage[start..end]; + + let mut result = vec![0u8; l_slice.len()]; + simd::cmp_scalar_f32(l_slice, rhs, &mut result, simd_op); + + return make_bool_tensor(result, shape, out_dtype); + } + + // Fallback to generic path + compare_elem_typed(lhs, rhs, out_dtype, cmp) +} + +/// Fallback when SIMD is disabled. +#[cfg(not(feature = "simd"))] +fn compare_elem_f32( + lhs: FlexTensor, + rhs: f32, + out_dtype: BoolDType, + cmp: Cmp, + _simd_hint: Option, +) -> FlexTensor +where + Cmp: Fn(f32, f32) -> bool, +{ + compare_elem_typed(lhs, rhs, out_dtype, cmp) +} + +fn compare_typed( + lhs: FlexTensor, + rhs: &FlexTensor, + out_dtype: BoolDType, + cmp: Cmp, +) -> FlexTensor +where + E: Element + Pod, + Cmp: Fn(E, E) -> bool, +{ + let shape = lhs.layout().shape().clone(); + let lhs_storage: &[E] = lhs.storage(); + let rhs_storage: &[E] = rhs.storage(); + + let result: Vec = match ( + lhs.layout().contiguous_offsets(), + rhs.layout().contiguous_offsets(), + ) { + (Some((l_start, l_end)), Some((r_start, r_end))) => { + let l_slice = &lhs_storage[l_start..l_end]; + let r_slice = &rhs_storage[r_start..r_end]; + l_slice + .iter() + .zip(r_slice) + .map(|(&a, &b)| cmp(a, b) as u8) + .collect() + } + // Fast path for 2D non-contiguous (common for transpose) + _ if lhs.layout().num_dims() == 2 => crate::ops::binary::apply_2d_strided( + lhs_storage, + rhs_storage, + lhs.layout(), + rhs.layout(), + |a, b| cmp(a, b) as u8, + ), + _ => { + let lhs_iter = StridedIter::new(lhs.layout()); + let rhs_iter = StridedIter::new(rhs.layout()); + lhs_iter + .zip(rhs_iter) + .map(|(li, ri)| cmp(lhs_storage[li], rhs_storage[ri]) as u8) + .collect() + } + }; + + make_bool_tensor(result, shape, out_dtype) +} + +fn compare_elem_typed(lhs: FlexTensor, rhs: E, out_dtype: BoolDType, cmp: Cmp) -> FlexTensor +where + E: Element + Pod + Copy, + Cmp: Fn(E, E) -> bool, +{ + let shape = lhs.layout().shape().clone(); + let lhs_storage: &[E] = lhs.storage(); + + let result: Vec = match lhs.layout().contiguous_offsets() { + Some((start, end)) => lhs_storage[start..end] + .iter() + .map(|&a| cmp(a, rhs) as u8) + .collect(), + None => StridedIter::new(lhs.layout()) + .map(|idx| cmp(lhs_storage[idx], rhs) as u8) + .collect(), + }; + + make_bool_tensor(result, shape, out_dtype) +} + +/// Build a bool `FlexTensor` from a `Vec` of 0/1 bytes, tagged with the +/// requested output dtype. +/// +/// burn-flex stores bools as 1 byte per element, so only Native and U8 are +/// supported. `Bool(U32)` would require 4-byte-per-element storage throughout +/// the backend; `dtype_usage` declares it unsupported and this function panics +/// if it's requested. +pub(crate) fn make_bool_tensor(data: Vec, shape: Shape, out_dtype: BoolDType) -> FlexTensor { + let store = match out_dtype { + BoolDType::Native => BoolStore::Native, + BoolDType::U8 => BoolStore::U8, + BoolDType::U32 => panic!( + "burn-flex does not support Bool(U32) storage (only Native and U8). \ + Use a backend that declares Bool(U32) support, or work with Bool(Native)/Bool(U8)." + ), + }; + let bytes = Bytes::from_elems(data); + FlexTensor::new(bytes, Layout::contiguous(shape), DType::Bool(store)) +} + +// Specific comparison functions + +pub fn greater(lhs: FlexTensor, rhs: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + compare( + lhs, + rhs, + out_dtype, + |a, b| a > b, + |a, b| a > b, + Some(CompareOp::Gt), + ) +} + +pub fn greater_elem(lhs: FlexTensor, rhs: f64, out_dtype: BoolDType) -> FlexTensor { + compare_elem( + lhs, + rhs, + out_dtype, + |a, b| a > b, + |a, b| a > b, + Some(CompareOp::Gt), + ) +} + +pub fn greater_equal(lhs: FlexTensor, rhs: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + compare( + lhs, + rhs, + out_dtype, + |a, b| a >= b, + |a, b| a >= b, + Some(CompareOp::Ge), + ) +} + +pub fn greater_equal_elem(lhs: FlexTensor, rhs: f64, out_dtype: BoolDType) -> FlexTensor { + compare_elem( + lhs, + rhs, + out_dtype, + |a, b| a >= b, + |a, b| a >= b, + Some(CompareOp::Ge), + ) +} + +pub fn lower(lhs: FlexTensor, rhs: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + compare( + lhs, + rhs, + out_dtype, + |a, b| a < b, + |a, b| a < b, + Some(CompareOp::Lt), + ) +} + +pub fn lower_elem(lhs: FlexTensor, rhs: f64, out_dtype: BoolDType) -> FlexTensor { + compare_elem( + lhs, + rhs, + out_dtype, + |a, b| a < b, + |a, b| a < b, + Some(CompareOp::Lt), + ) +} + +pub fn lower_equal(lhs: FlexTensor, rhs: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + compare( + lhs, + rhs, + out_dtype, + |a, b| a <= b, + |a, b| a <= b, + Some(CompareOp::Le), + ) +} + +pub fn lower_equal_elem(lhs: FlexTensor, rhs: f64, out_dtype: BoolDType) -> FlexTensor { + compare_elem( + lhs, + rhs, + out_dtype, + |a, b| a <= b, + |a, b| a <= b, + Some(CompareOp::Le), + ) +} + +pub fn equal(lhs: FlexTensor, rhs: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + compare( + lhs, + rhs, + out_dtype, + |a, b| a == b, + |a, b| a == b, + Some(CompareOp::Eq), + ) +} + +pub fn equal_elem(lhs: FlexTensor, rhs: f64, out_dtype: BoolDType) -> FlexTensor { + compare_elem( + lhs, + rhs, + out_dtype, + |a, b| a == b, + |a, b| a == b, + Some(CompareOp::Eq), + ) +} + +pub fn not_equal(lhs: FlexTensor, rhs: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + compare( + lhs, + rhs, + out_dtype, + |a, b| a != b, + |a, b| a != b, + Some(CompareOp::Ne), + ) +} + +pub fn not_equal_elem(lhs: FlexTensor, rhs: f64, out_dtype: BoolDType) -> FlexTensor { + compare_elem( + lhs, + rhs, + out_dtype, + |a, b| a != b, + |a, b| a != b, + Some(CompareOp::Ne), + ) +} + +// Integer comparison functions + +fn compare_int( + lhs: FlexTensor, + rhs: FlexTensor, + out_dtype: BoolDType, + i64_cmp: I64Cmp, + u64_cmp: U64Cmp, +) -> FlexTensor +where + I64Cmp: Fn(i64, i64) -> bool, + U64Cmp: Fn(u64, u64) -> bool, +{ + let (lhs, rhs) = crate::ops::expand::broadcast_binary(lhs, rhs); + + match lhs.dtype() { + DType::I64 => compare_typed(lhs, &rhs, out_dtype, i64_cmp), + DType::U64 => compare_typed(lhs, &rhs, out_dtype, u64_cmp), + DType::I32 => compare_typed(lhs, &rhs, out_dtype, |a: i32, b: i32| { + i64_cmp(a as i64, b as i64) + }), + DType::I16 => compare_typed(lhs, &rhs, out_dtype, |a: i16, b: i16| { + i64_cmp(a as i64, b as i64) + }), + DType::I8 => compare_typed(lhs, &rhs, out_dtype, |a: i8, b: i8| { + i64_cmp(a as i64, b as i64) + }), + DType::U32 => compare_typed(lhs, &rhs, out_dtype, |a: u32, b: u32| { + i64_cmp(a as i64, b as i64) + }), + DType::U16 => compare_typed(lhs, &rhs, out_dtype, |a: u16, b: u16| { + i64_cmp(a as i64, b as i64) + }), + DType::U8 => compare_typed(lhs, &rhs, out_dtype, |a: u8, b: u8| { + i64_cmp(a as i64, b as i64) + }), + other => panic!("compare_int: unsupported dtype {:?}", other), + } +} + +fn compare_int_elem( + lhs: FlexTensor, + i64_rhs: i64, + u64_rhs: u64, + out_dtype: BoolDType, + i64_cmp: I64Cmp, + u64_cmp: U64Cmp, +) -> FlexTensor +where + I64Cmp: Fn(i64, i64) -> bool, + U64Cmp: Fn(u64, u64) -> bool, +{ + match lhs.dtype() { + DType::I64 => compare_elem_typed(lhs, i64_rhs, out_dtype, i64_cmp), + DType::U64 => compare_elem_typed(lhs, u64_rhs, out_dtype, u64_cmp), + DType::I32 => compare_elem_typed(lhs, i64_rhs as i32, out_dtype, |a: i32, b: i32| { + i64_cmp(a as i64, b as i64) + }), + DType::I16 => compare_elem_typed(lhs, i64_rhs as i16, out_dtype, |a: i16, b: i16| { + i64_cmp(a as i64, b as i64) + }), + DType::I8 => compare_elem_typed(lhs, i64_rhs as i8, out_dtype, |a: i8, b: i8| { + i64_cmp(a as i64, b as i64) + }), + DType::U32 => compare_elem_typed(lhs, i64_rhs as u32, out_dtype, |a: u32, b: u32| { + i64_cmp(a as i64, b as i64) + }), + DType::U16 => compare_elem_typed(lhs, i64_rhs as u16, out_dtype, |a: u16, b: u16| { + i64_cmp(a as i64, b as i64) + }), + DType::U8 => compare_elem_typed(lhs, i64_rhs as u8, out_dtype, |a: u8, b: u8| { + i64_cmp(a as i64, b as i64) + }), + other => panic!("compare_int_elem: unsupported dtype {:?}", other), + } +} + +pub fn int_greater(lhs: FlexTensor, rhs: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + compare_int(lhs, rhs, out_dtype, |a, b| a > b, |a, b| a > b) +} + +pub fn int_greater_elem( + lhs: FlexTensor, + i64_rhs: i64, + u64_rhs: u64, + out_dtype: BoolDType, +) -> FlexTensor { + compare_int_elem(lhs, i64_rhs, u64_rhs, out_dtype, |a, b| a > b, |a, b| a > b) +} + +pub fn int_greater_equal(lhs: FlexTensor, rhs: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + compare_int(lhs, rhs, out_dtype, |a, b| a >= b, |a, b| a >= b) +} + +pub fn int_greater_equal_elem( + lhs: FlexTensor, + i64_rhs: i64, + u64_rhs: u64, + out_dtype: BoolDType, +) -> FlexTensor { + compare_int_elem( + lhs, + i64_rhs, + u64_rhs, + out_dtype, + |a, b| a >= b, + |a, b| a >= b, + ) +} + +pub fn int_lower(lhs: FlexTensor, rhs: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + compare_int(lhs, rhs, out_dtype, |a, b| a < b, |a, b| a < b) +} + +pub fn int_lower_elem( + lhs: FlexTensor, + i64_rhs: i64, + u64_rhs: u64, + out_dtype: BoolDType, +) -> FlexTensor { + compare_int_elem(lhs, i64_rhs, u64_rhs, out_dtype, |a, b| a < b, |a, b| a < b) +} + +pub fn int_lower_equal(lhs: FlexTensor, rhs: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + compare_int(lhs, rhs, out_dtype, |a, b| a <= b, |a, b| a <= b) +} + +pub fn int_lower_equal_elem( + lhs: FlexTensor, + i64_rhs: i64, + u64_rhs: u64, + out_dtype: BoolDType, +) -> FlexTensor { + compare_int_elem( + lhs, + i64_rhs, + u64_rhs, + out_dtype, + |a, b| a <= b, + |a, b| a <= b, + ) +} + +pub fn int_equal(lhs: FlexTensor, rhs: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + compare_int(lhs, rhs, out_dtype, |a, b| a == b, |a, b| a == b) +} + +pub fn int_equal_elem( + lhs: FlexTensor, + i64_rhs: i64, + u64_rhs: u64, + out_dtype: BoolDType, +) -> FlexTensor { + compare_int_elem( + lhs, + i64_rhs, + u64_rhs, + out_dtype, + |a, b| a == b, + |a, b| a == b, + ) +} + +pub fn int_not_equal(lhs: FlexTensor, rhs: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + compare_int(lhs, rhs, out_dtype, |a, b| a != b, |a, b| a != b) +} + +pub fn int_not_equal_elem( + lhs: FlexTensor, + i64_rhs: i64, + u64_rhs: u64, + out_dtype: BoolDType, +) -> FlexTensor { + compare_int_elem( + lhs, + i64_rhs, + u64_rhs, + out_dtype, + |a, b| a != b, + |a, b| a != b, + ) +} + +pub fn bool_not_equal(lhs: FlexTensor, rhs: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + let (lhs, rhs) = crate::ops::expand::broadcast_binary(lhs, rhs); + let shape = lhs.layout().shape().clone(); + let lhs_data: &[u8] = lhs.bytes(); + let rhs_data: &[u8] = rhs.bytes(); + let result: Vec = match ( + lhs.layout().contiguous_offsets(), + rhs.layout().contiguous_offsets(), + ) { + (Some((ls, le)), Some((rs, re))) => lhs_data[ls..le] + .iter() + .zip(&rhs_data[rs..re]) + .map(|(&a, &b)| if a != b { 1 } else { 0 }) + .collect(), + _ => { + let lhs = lhs.to_contiguous(); + let rhs = rhs.to_contiguous(); + lhs.bytes() + .iter() + .zip(rhs.bytes()) + .map(|(&a, &b)| if a != b { 1 } else { 0 }) + .collect() + } + }; + make_bool_tensor(result, shape, out_dtype) +} + +pub fn bool_not_equal_elem(lhs: FlexTensor, rhs: bool, out_dtype: BoolDType) -> FlexTensor { + let rhs_val: u8 = if rhs { 1 } else { 0 }; + let shape = lhs.layout().shape().clone(); + let lhs = lhs.to_contiguous(); + let data: &[u8] = lhs.bytes(); + let result: Vec = data + .iter() + .map(|&a| if a != rhs_val { 1 } else { 0 }) + .collect(); + make_bool_tensor(result, shape, out_dtype) +} + +// ============================================================================ +// any / all operations +// ============================================================================ + +/// Check if any element is non-zero (float tensors). +pub fn any_float(tensor: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + let has_any = match tensor.dtype() { + DType::F32 => iter_elements::(&tensor).any(|x| x != 0.0), + DType::F64 => iter_elements::(&tensor).any(|x| x != 0.0), + DType::F16 => iter_elements::(&tensor).any(|x: f16| x.to_f32() != 0.0), + DType::BF16 => iter_elements::(&tensor).any(|x: bf16| x.to_f32() != 0.0), + _ => panic!("any_float: unsupported dtype {:?}", tensor.dtype()), + }; + bool_scalar(has_any, out_dtype) +} + +/// Check if any element along a dimension is non-zero (float tensors). +pub fn any_float_dim(tensor: FlexTensor, dim: usize, out_dtype: BoolDType) -> FlexTensor { + reduce_bool_dim(&tensor, dim, false, |a, b| a || b, out_dtype) +} + +/// Check if all elements are non-zero (float tensors). +pub fn all_float(tensor: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + let all = match tensor.dtype() { + DType::F32 => iter_elements::(&tensor).all(|x| x != 0.0), + DType::F64 => iter_elements::(&tensor).all(|x| x != 0.0), + DType::F16 => iter_elements::(&tensor).all(|x: f16| x.to_f32() != 0.0), + DType::BF16 => iter_elements::(&tensor).all(|x: bf16| x.to_f32() != 0.0), + _ => panic!("all_float: unsupported dtype {:?}", tensor.dtype()), + }; + bool_scalar(all, out_dtype) +} + +/// Check if all elements along a dimension are non-zero (float tensors). +pub fn all_float_dim(tensor: FlexTensor, dim: usize, out_dtype: BoolDType) -> FlexTensor { + reduce_bool_dim(&tensor, dim, true, |a, b| a && b, out_dtype) +} + +/// Check if any element is non-zero (int tensors). +pub fn any_int(tensor: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + let has_any = match tensor.dtype() { + DType::I64 => iter_elements::(&tensor).any(|x| x != 0), + DType::I32 => iter_elements::(&tensor).any(|x| x != 0), + DType::I16 => iter_elements::(&tensor).any(|x| x != 0), + DType::I8 => iter_elements::(&tensor).any(|x| x != 0), + DType::U64 => iter_elements::(&tensor).any(|x| x != 0), + DType::U32 => iter_elements::(&tensor).any(|x| x != 0), + DType::U16 => iter_elements::(&tensor).any(|x| x != 0), + DType::U8 => iter_elements::(&tensor).any(|x| x != 0), + _ => panic!("any_int: unsupported dtype {:?}", tensor.dtype()), + }; + bool_scalar(has_any, out_dtype) +} + +/// Check if any element along a dimension is non-zero (int tensors). +pub fn any_int_dim(tensor: FlexTensor, dim: usize, out_dtype: BoolDType) -> FlexTensor { + reduce_bool_dim_int(&tensor, dim, false, |a, b| a || b, out_dtype) +} + +/// Check if all elements are non-zero (int tensors). +pub fn all_int(tensor: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + let all = match tensor.dtype() { + DType::I64 => iter_elements::(&tensor).all(|x| x != 0), + DType::I32 => iter_elements::(&tensor).all(|x| x != 0), + DType::I16 => iter_elements::(&tensor).all(|x| x != 0), + DType::I8 => iter_elements::(&tensor).all(|x| x != 0), + DType::U64 => iter_elements::(&tensor).all(|x| x != 0), + DType::U32 => iter_elements::(&tensor).all(|x| x != 0), + DType::U16 => iter_elements::(&tensor).all(|x| x != 0), + DType::U8 => iter_elements::(&tensor).all(|x| x != 0), + _ => panic!("all_int: unsupported dtype {:?}", tensor.dtype()), + }; + bool_scalar(all, out_dtype) +} + +/// Check if all elements along a dimension are non-zero (int tensors). +pub fn all_int_dim(tensor: FlexTensor, dim: usize, out_dtype: BoolDType) -> FlexTensor { + reduce_bool_dim_int(&tensor, dim, true, |a, b| a && b, out_dtype) +} + +/// Check if any bool element is true. +pub fn any_bool(tensor: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let data: &[u8] = tensor.bytes(); + bool_scalar(data.iter().any(|&x| x != 0), out_dtype) +} + +/// Check if any bool element along a dimension is true. +pub fn any_bool_dim(tensor: FlexTensor, dim: usize, out_dtype: BoolDType) -> FlexTensor { + reduce_bool_dim_raw(&tensor, dim, false, |a, b| a || b, out_dtype) +} + +/// Check if all bool elements are true. +pub fn all_bool(tensor: FlexTensor, out_dtype: BoolDType) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let data: &[u8] = tensor.bytes(); + bool_scalar(data.iter().all(|&x| x != 0), out_dtype) +} + +/// Check if all bool elements along a dimension are true. +pub fn all_bool_dim(tensor: FlexTensor, dim: usize, out_dtype: BoolDType) -> FlexTensor { + reduce_bool_dim_raw(&tensor, dim, true, |a, b| a && b, out_dtype) +} + +// ============================================================================ +// Helpers for any/all +// ============================================================================ + +fn bool_scalar(val: bool, out_dtype: BoolDType) -> FlexTensor { + let byte: u8 = if val { 1 } else { 0 }; + make_bool_tensor(alloc::vec![byte], Shape::from(alloc::vec![1]), out_dtype) +} + +fn iter_elements<'a, E: Element + Pod + 'a>( + tensor: &'a FlexTensor, +) -> Box + 'a> { + let data: &[E] = tensor.storage(); + match tensor.layout().contiguous_offsets() { + Some((start, end)) => Box::new(data[start..end].iter().copied()), + None => Box::new(StridedIter::new(tensor.layout()).map(move |idx| data[idx])), + } +} + +/// Reduce along a dimension producing a bool tensor. +/// +/// The `is_nonzero` closure reads the data slice at a given index and returns +/// whether the element is nonzero. +fn reduce_bool_dim_with( + tensor: &FlexTensor, + dim: usize, + init: bool, + combine: fn(bool, bool) -> bool, + out_dtype: BoolDType, + is_nonzero: impl Fn(usize) -> bool, +) -> FlexTensor { + debug_assert!(tensor.is_contiguous() && tensor.layout().start_offset() == 0); + let shape = tensor.layout().shape(); + let ndims = shape.num_dims(); + assert!(dim < ndims); + + let dim_size = shape[dim]; + let mut out_shape: Vec = shape.to_vec(); + out_shape[dim] = 1; + let outer_size: usize = shape[..dim].iter().product(); + let inner_size: usize = shape[dim + 1..].iter().product(); + + let out_size = outer_size.max(1) * inner_size.max(1); + let mut result: Vec = Vec::with_capacity(out_size); + + for outer in 0..outer_size.max(1) { + for inner in 0..inner_size.max(1) { + let mut acc = init; + for d in 0..dim_size { + let idx = outer * dim_size * inner_size + d * inner_size + inner; + acc = combine(acc, is_nonzero(idx)); + } + result.push(if acc { 1 } else { 0 }); + } + } + + make_bool_tensor(result, Shape::from(out_shape), out_dtype) +} + +/// Reduce along a dimension producing a bool tensor (for float any/all_dim). +fn reduce_bool_dim( + tensor: &FlexTensor, + dim: usize, + init: bool, + combine: fn(bool, bool) -> bool, + out_dtype: BoolDType, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + match tensor.dtype() { + DType::F32 => { + let data: &[f32] = tensor.storage(); + reduce_bool_dim_with(&tensor, dim, init, combine, out_dtype, |idx| { + data[idx] != 0.0 + }) + } + DType::F64 => { + let data: &[f64] = tensor.storage(); + reduce_bool_dim_with(&tensor, dim, init, combine, out_dtype, |idx| { + data[idx] != 0.0 + }) + } + DType::F16 => { + let data: &[f16] = tensor.storage(); + reduce_bool_dim_with(&tensor, dim, init, combine, out_dtype, |idx| { + data[idx].to_f32() != 0.0 + }) + } + DType::BF16 => { + let data: &[bf16] = tensor.storage(); + reduce_bool_dim_with(&tensor, dim, init, combine, out_dtype, |idx| { + data[idx].to_f32() != 0.0 + }) + } + _ => panic!("reduce_bool_dim: unsupported dtype {:?}", tensor.dtype()), + } +} + +/// Reduce along a dimension producing a bool tensor (for int any/all_dim). +fn reduce_bool_dim_int( + tensor: &FlexTensor, + dim: usize, + init: bool, + combine: fn(bool, bool) -> bool, + out_dtype: BoolDType, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + macro_rules! dispatch { + ($ty:ty) => {{ + let data: &[$ty] = tensor.storage(); + reduce_bool_dim_with(&tensor, dim, init, combine, out_dtype, |idx| data[idx] != 0) + }}; + } + match tensor.dtype() { + DType::I64 => dispatch!(i64), + DType::I32 => dispatch!(i32), + DType::I16 => dispatch!(i16), + DType::I8 => dispatch!(i8), + DType::U64 => dispatch!(u64), + DType::U32 => dispatch!(u32), + DType::U16 => dispatch!(u16), + DType::U8 => dispatch!(u8), + other => panic!("reduce_bool_dim_int: unsupported dtype {:?}", other), + } +} + +/// Reduce along a dimension producing a bool tensor (for bool any/all_dim). +fn reduce_bool_dim_raw( + tensor: &FlexTensor, + dim: usize, + init: bool, + combine: fn(bool, bool) -> bool, + out_dtype: BoolDType, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let data: &[u8] = tensor.bytes(); + reduce_bool_dim_with(&tensor, dim, init, combine, out_dtype, |idx| data[idx] != 0) +} + +// Tests kept here probe flex-internal `reduce_bool_dim_with` dispatch on +// non-contiguous inputs (stale-pointer-read regression, see prior incident +// in `any_float_dim`). Plain comparison ops and stride variants (flipped +// / transposed / narrowed) have been migrated to burn-backend-tests at +// tensor/{float,int}/ops/comparison.rs so every backend is exercised. +#[cfg(test)] +mod tests { + use super::*; + use burn_backend::TensorData; + + fn tensor_2d(data: Vec, rows: usize, cols: usize) -> FlexTensor { + FlexTensor::from_data(TensorData::new(data, vec![rows, cols])) + } + + #[test] + fn test_any_float_dim_transposed() { + // [[0, 1], [0, 0]] transposed -> [[0, 0], [1, 0]] + // any_dim along dim 1: [false, true] -> [0, 1] + // Before the C2 fix, to_contiguous was called inside reduce_bool_dim_with + // AFTER taking a reference to the original non-contiguous storage, + // causing stale pointer reads. + let tensor = tensor_2d(vec![0.0, 1.0, 0.0, 0.0], 2, 2); + let transposed = tensor.transpose(0, 1); + assert!(!transposed.is_contiguous()); + + let result = any_float_dim(transposed, 1, BoolDType::Native); + let data: &[u8] = result.bytes(); + assert_eq!(data, &[0, 1]); // row 0: all zeros; row 1: has a 1 + } + + #[test] + fn test_any_float_dim_narrowed() { + // [0, 5, 0, 3, 0, 0] narrowed to [[5, 0], [3, 0]] (shape [2,2]) + // any_dim along dim 1: [true, true] -> [1, 1] + let tensor = FlexTensor::from_data(TensorData::new( + vec![0.0f32, 5.0, 0.0, 3.0, 0.0, 0.0], + [3, 2], + )); + let narrowed = tensor.narrow(0, 0, 2); // first 2 rows: [[0, 5], [0, 3]] + let result = any_float_dim(narrowed, 1, BoolDType::Native); + let data: &[u8] = result.bytes(); + assert_eq!(data, &[1, 1]); + } +} diff --git a/crates/burn-flex/src/ops/conv.rs b/crates/burn-flex/src/ops/conv.rs new file mode 100644 index 0000000..69fae22 --- /dev/null +++ b/crates/burn-flex/src/ops/conv.rs @@ -0,0 +1,2969 @@ +//! Forward convolution operations using tiled im2col + gemm approach. +//! +//! All convolutions (1D, 2D, 3D) use a unified 3D implementation: +//! - conv1d: adds two size-1 dimensions, calls conv3d, squeezes output +//! - conv2d: adds one size-1 dimension, calls conv3d, squeezes output +//! - conv3d: native implementation +//! +//! Optimizations: +//! - Tiled im2col: Process output in tiles for better cache usage and parallelism +//! - NHWC layout: Convert to channels-last for cache-friendly access +//! - Nested parallelism: Batch and tile dimensions run in parallel via rayon +//! - 1x1 fast path: Skip im2col for pointwise convolutions +//! - Depthwise fast path: For canonical depthwise (groups == c_in == c_out, +//! channels_per_group == 1), skip NHWC conversion, im2col, and gemm entirely. +//! Uses a direct per-(b, c) accumulate with analytic bounds so the inner +//! spatial loop has no padding checks and autovectorizes. +//! - Small-channel fast path: For groups=1 convs with very few input channels +//! (e.g. 3-channel Sobel-style edge filters), reuse the depthwise kernel by +//! accumulating over input channels. Skips the same NHWC/im2col/gemm overhead +//! as the depthwise path. Wins when `channels_in` is small enough that +//! gemm's per-dispatch setup cost dominates over the tiny inner compute. +//! - Direct conv path: For small-spatial 1D-like convolutions, decompose into +//! per-kernel-position gemm calls on NCHW data, skipping NHWC conversion and im2col +//! +//! Supported dtypes: f32, f64, f16 (native gemm), bf16 (via f32 conversion) + +use alloc::vec; +use alloc::vec::Vec; +use burn_backend::DType; +use burn_backend::ops::ConvOptions; +use burn_backend::ops::conv::calculate_conv_output_size; +use burn_std::{Bytes, Shape, f16}; + +use crate::{FlexTensor, Layout}; + +use super::conv_common::{add_bias, squeeze_3d_to_1d, squeeze_3d_to_2d}; + +// ============================================================================ +// Macros for forward conv +// ============================================================================ + +/// Generates a conv3d_1x1 function that uses the optimized gemm fast path. +macro_rules! conv3d_1x1_typed { + ($fn_name:ident, $T:ty, $dtype:expr, $zero:expr, $one:expr, $add_fn:expr) => { + fn $fn_name( + x: FlexTensor, + weight: FlexTensor, + bias: Option, + options: &ConvOptions<3>, + ) -> FlexTensor { + conv3d_1x1_impl::<$T>(x, weight, bias, options, $dtype, $zero, $one, $add_fn) + } + }; +} + +/// Generates a conv3d typed function with 1x1, depthwise, small-channel, and +/// direct fast-path checks. +macro_rules! conv3d_typed { + ($fn_name:ident, $T:ty, $dtype:expr, $zero:expr, $gemm_fn:ident, $add_fn:expr, $fn_1x1:ident, $fn_depthwise:ident, $fn_small_channel:ident $(, $fn_direct:ident)?) => { + pub fn $fn_name( + x: FlexTensor, + weight: FlexTensor, + bias: Option, + options: &ConvOptions<3>, + ) -> FlexTensor { + let w_shape = weight.layout().shape(); + if is_1x1_conv(w_shape[2], w_shape[3], w_shape[4], options) { + return $fn_1x1(x, weight, bias, options); + } + let x_shape = x.layout().shape(); + if should_use_depthwise_conv(x_shape, w_shape, options) { + return $fn_depthwise(x, weight, bias, options); + } + if should_use_small_channel_conv(x_shape, w_shape, options) { + return $fn_small_channel(x, weight, bias, options); + } + $( + if should_use_direct_conv(x_shape, w_shape, options) { + return $fn_direct(x, weight, bias, options); + } + )? + conv3d_impl::<$T>(x, weight, bias, options, $dtype, $zero, $gemm_fn, $add_fn) + } + }; +} + +// ============================================================================ +// Conv1d - delegates to conv3d +// ============================================================================ + +conv_nd_via_3d!( + conv1d_f32, + conv3d_f32, + expand_1d_to_3d, + squeeze_3d_to_1d, + 1, + ConvOptions +); +conv_nd_via_3d!( + conv1d_f64, + conv3d_f64, + expand_1d_to_3d, + squeeze_3d_to_1d, + 1, + ConvOptions +); +conv_nd_via_3d!( + conv1d_f16, + conv3d_f16, + expand_1d_to_3d, + squeeze_3d_to_1d, + 1, + ConvOptions +); +bf16_via_f32!(conv1d_bf16, conv1d_f32, 1, ConvOptions); + +fn expand_1d_to_3d( + x: &FlexTensor, + weight: &FlexTensor, + options: &ConvOptions<1>, +) -> (FlexTensor, FlexTensor, ConvOptions<3>) { + let x_shape = x.layout().shape(); + let x_3d = x.reshape(Shape::from(vec![x_shape[0], x_shape[1], 1, 1, x_shape[2]])); + + let w_shape = weight.layout().shape(); + let weight_3d = weight.reshape(Shape::from(vec![w_shape[0], w_shape[1], 1, 1, w_shape[2]])); + + let options_3d = ConvOptions::new( + [1, 1, options.stride[0]], + [0, 0, options.padding[0]], + [1, 1, options.dilation[0]], + options.groups, + ); + + (x_3d, weight_3d, options_3d) +} + +// ============================================================================ +// Conv2d - delegates to conv3d +// ============================================================================ + +conv_nd_via_3d!( + conv2d_f32, + conv3d_f32, + expand_2d_to_3d, + squeeze_3d_to_2d, + 2, + ConvOptions +); +conv_nd_via_3d!( + conv2d_f64, + conv3d_f64, + expand_2d_to_3d, + squeeze_3d_to_2d, + 2, + ConvOptions +); +conv_nd_via_3d!( + conv2d_f16, + conv3d_f16, + expand_2d_to_3d, + squeeze_3d_to_2d, + 2, + ConvOptions +); +bf16_via_f32!(conv2d_bf16, conv2d_f32, 2, ConvOptions); + +fn expand_2d_to_3d( + x: &FlexTensor, + weight: &FlexTensor, + options: &ConvOptions<2>, +) -> (FlexTensor, FlexTensor, ConvOptions<3>) { + let x_shape = x.layout().shape(); + let x_3d = x.reshape(Shape::from(vec![ + x_shape[0], x_shape[1], 1, x_shape[2], x_shape[3], + ])); + + let w_shape = weight.layout().shape(); + let weight_3d = weight.reshape(Shape::from(vec![ + w_shape[0], w_shape[1], 1, w_shape[2], w_shape[3], + ])); + + let options_3d = ConvOptions::new( + [1, options.stride[0], options.stride[1]], + [0, options.padding[0], options.padding[1]], + [1, options.dilation[0], options.dilation[1]], + options.groups, + ); + + (x_3d, weight_3d, options_3d) +} + +// ============================================================================ +// Conv3d - native implementations +// ============================================================================ + +conv3d_typed!( + conv3d_f32, + f32, + DType::F32, + 0.0f32, + gemm_f32, + |a, b| a + b, + conv3d_1x1_f32, + conv3d_depthwise_f32, + conv3d_small_channel_f32, + conv3d_direct_f32 +); +conv3d_typed!( + conv3d_f64, + f64, + DType::F64, + 0.0f64, + gemm_f64, + |a, b| a + b, + conv3d_1x1_f64, + conv3d_depthwise_f64, + conv3d_small_channel_f64, + conv3d_direct_f64 +); +conv3d_typed!( + conv3d_f16, + f16, + DType::F16, + f16::from_f32(0.0), + gemm_f16, + |a: f16, b: f16| f16::from_f32(a.to_f32() + b.to_f32()), + conv3d_1x1_f16, + conv3d_depthwise_f16, + conv3d_small_channel_f16 +); +bf16_via_f32!(conv3d_bf16, conv3d_f32, 3, ConvOptions); + +/// Generic 3D convolution implementation using tiled im2col. +/// +/// Tiled approach processes output in TILE_SIZE chunks: +/// - Reduces memory usage (smaller im2col buffer per tile) +/// - Enables tile-level parallelism +/// - Improves cache utilization +#[allow(clippy::too_many_arguments, clippy::type_complexity)] +fn conv3d_impl( + x: FlexTensor, + weight: FlexTensor, + bias: Option, + options: &ConvOptions<3>, + dtype: DType, + zero: T, + gemm_fn: fn(&[T], &[T], usize, usize, usize) -> Vec, + add_fn: fn(T, T) -> T, +) -> FlexTensor { + let x = x.to_contiguous(); + let weight = weight.to_contiguous(); + + let x_shape = x.layout().shape(); + let w_shape = weight.layout().shape(); + + let batch_size = x_shape[0]; + let channels_in = x_shape[1]; + let in_d = x_shape[2]; + let in_h = x_shape[3]; + let in_w = x_shape[4]; + + let channels_out = w_shape[0]; + let channels_per_group = w_shape[1]; + let kernel_d = w_shape[2]; + let kernel_h = w_shape[3]; + let kernel_w = w_shape[4]; + + let [stride_d, stride_h, stride_w] = options.stride; + let [pad_d, pad_h, pad_w] = options.padding; + let groups = options.groups; + let out_channels_per_group = channels_out / groups; + + let out_d = calculate_conv_output_size(kernel_d, stride_d, pad_d, options.dilation[0], in_d); + let out_h = calculate_conv_output_size(kernel_h, stride_h, pad_h, options.dilation[1], in_h); + let out_w = calculate_conv_output_size(kernel_w, stride_w, pad_w, options.dilation[2], in_w); + + // Validate sizes won't overflow index calculations + let _total = [batch_size, channels_out, out_d, out_h, out_w] + .iter() + .try_fold(1usize, |acc, &x| acc.checked_mul(x)) + .expect("conv: output tensor dimensions would overflow index calculations"); + let _col_total = [channels_per_group, kernel_d, kernel_h, kernel_w] + .iter() + .try_fold(1usize, |acc, &x| acc.checked_mul(x)) + .expect("conv: kernel dimensions would overflow index calculations"); + + let x_data: &[T] = x.storage(); + let w_data: &[T] = weight.storage(); + + let col_len = channels_per_group * kernel_d * kernel_h * kernel_w; + let spatial_out = out_d * out_h * out_w; + + let [dilation_d, dilation_h, dilation_w] = options.dilation; + + // Tile size for processing output pixels. Larger = better GEMM utilization, + // smaller = more parallelism and better cache usage. 512 is a good balance. + const TILE_SIZE: usize = 512; + let num_tiles = spatial_out.div_ceil(TILE_SIZE); + + // Flatten kernel [c_out, c_in, kd, kh, kw] -> [c_out, kd, kh, kw, c_in] + // for GEMM. Expressed as a 2D transpose of (c_in, k_spatial) per c_out + // so the compiler can unroll the k_spatial inner loop; the older 5- + // nested formulation had a 10-term index expression LLVM wouldn't + // autovectorize. + let k_spatial = kernel_d * kernel_h * kernel_w; + let mut w_flat = vec![zero; channels_out * col_len]; + for c_out in 0..channels_out { + let src_base = c_out * channels_per_group * k_spatial; + let dst_base = c_out * col_len; + for c_in in 0..channels_per_group { + let src_row = src_base + c_in * k_spatial; + for k in 0..k_spatial { + w_flat[dst_base + k * channels_per_group + c_in] = w_data[src_row + k]; + } + } + } + + // Convert input to NHWC layout for cache-friendly access in im2col. + // Loop order (c innermost, spatial outer) is intentional: on aarch64 + // M3 Max, strided loads + contiguous stores beat the inverse by + // 10-35% per layer. Load prefetchers outrun the store write buffer. + let nhwc_stride = ( + in_d * in_h * in_w * channels_in, + in_h * in_w * channels_in, + in_w * channels_in, + channels_in, + 1, + ); + let mut x_nhwc = vec![zero; batch_size * in_d * in_h * in_w * channels_in]; + for b in 0..batch_size { + for d in 0..in_d { + for h in 0..in_h { + for w in 0..in_w { + for c in 0..channels_in { + let src_idx = b * channels_in * in_d * in_h * in_w + + c * in_d * in_h * in_w + + d * in_h * in_w + + h * in_w + + w; + let dst_idx = b * nhwc_stride.0 + + d * nhwc_stride.1 + + h * nhwc_stride.2 + + w * nhwc_stride.3 + + c; + x_nhwc[dst_idx] = x_data[src_idx]; + } + } + } + } + } + + // Use tiled parallel execution + let output = { + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + + let mut dst = vec![zero; batch_size * channels_out * spatial_out]; + let dst_ptr = crate::ops::SendMutPtr::new(dst.as_mut_ptr()); + + // Process batches and tiles in parallel (nested parallelism) + (0..batch_size).into_par_iter().for_each(|b| { + (0..num_tiles).into_par_iter().for_each(|tile_idx| { + let tile_start = tile_idx * TILE_SIZE; + let tile_end = (tile_start + TILE_SIZE).min(spatial_out); + let tile_size = tile_end - tile_start; + + // Process each group separately + for g in 0..groups { + let in_c_start = g * channels_per_group; + let out_c_start = g * out_channels_per_group; + + // Build im2col for this tile and group + let mut col_tile = vec![zero; col_len * tile_size]; + + im2col_3d_tile( + &mut col_tile, + &x_nhwc, + tile_start, + tile_end, + out_h, + out_w, + kernel_d, + kernel_h, + kernel_w, + stride_d, + stride_h, + stride_w, + dilation_d, + dilation_h, + dilation_w, + pad_d, + pad_h, + pad_w, + in_d, + in_h, + in_w, + channels_per_group, + col_len, + b, + in_c_start, + nhwc_stride, + ); + + // Get weight slice for this group + let w_start = out_c_start * col_len; + let w_end = w_start + out_channels_per_group * col_len; + let w_group = &w_flat[w_start..w_end]; + + // GEMM: w_group[out_c_per_group, col_len] @ col_tile[tile_size, col_len]^T + let result = gemm_fn( + w_group, + &col_tile, + out_channels_per_group, + col_len, + tile_size, + ); + + // Write results to output for this group's output channels + for (local_idx, global_idx) in (tile_start..tile_end).enumerate() { + for c_out in 0..out_channels_per_group { + let dst_idx = b * channels_out * spatial_out + + (out_c_start + c_out) * spatial_out + + global_idx; + let res_idx = c_out * tile_size + local_idx; + unsafe { + debug_assert!( + dst_idx < batch_size * channels_out * spatial_out + ); + dst_ptr.write(dst_idx, result[res_idx]); + } + } + } + } + }); + }); + dst + } + #[cfg(not(feature = "rayon"))] + { + // Sequential path with tiling + let mut output = vec![zero; batch_size * channels_out * spatial_out]; + + for b in 0..batch_size { + for tile_idx in 0..num_tiles { + let tile_start = tile_idx * TILE_SIZE; + let tile_end = (tile_start + TILE_SIZE).min(spatial_out); + let tile_size = tile_end - tile_start; + + // Process each group separately + for g in 0..groups { + let in_c_start = g * channels_per_group; + let out_c_start = g * out_channels_per_group; + + let mut col_tile = vec![zero; col_len * tile_size]; + + im2col_3d_tile( + &mut col_tile, + &x_nhwc, + tile_start, + tile_end, + out_h, + out_w, + kernel_d, + kernel_h, + kernel_w, + stride_d, + stride_h, + stride_w, + dilation_d, + dilation_h, + dilation_w, + pad_d, + pad_h, + pad_w, + in_d, + in_h, + in_w, + channels_per_group, + col_len, + b, + in_c_start, + nhwc_stride, + ); + + // Get weight slice for this group + let w_start = out_c_start * col_len; + let w_end = w_start + out_channels_per_group * col_len; + let w_group = &w_flat[w_start..w_end]; + + // GEMM for this group + let result = gemm_fn( + w_group, + &col_tile, + out_channels_per_group, + col_len, + tile_size, + ); + + // Write results to output for this group's output channels + for (local_idx, global_idx) in (tile_start..tile_end).enumerate() { + for c_out in 0..out_channels_per_group { + let dst_idx = b * channels_out * spatial_out + + (out_c_start + c_out) * spatial_out + + global_idx; + let res_idx = c_out * tile_size + local_idx; + output[dst_idx] = result[res_idx]; + } + } + } + } + } + output + } + }; + + let mut output = output; + if let Some(bias) = bias { + let bias = bias.to_contiguous(); + let bias_data: &[T] = bias.storage(); + add_bias( + &mut output, + bias_data, + batch_size, + channels_out, + spatial_out, + add_fn, + ); + } + + let out_shape = Shape::from(vec![batch_size, channels_out, out_d, out_h, out_w]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + dtype, + ) +} + +/// Build im2col tile for a range of output positions. +/// +/// Fills `col_tile` with shape [tile_size, col_len] where each row is a flattened +/// patch from the NHWC input for one output position. +#[allow(clippy::too_many_arguments)] +fn im2col_3d_tile( + col_tile: &mut [T], + x_nhwc: &[T], + tile_start: usize, + tile_end: usize, + out_h: usize, + out_w: usize, + kernel_d: usize, + kernel_h: usize, + kernel_w: usize, + stride_d: usize, + stride_h: usize, + stride_w: usize, + dilation_d: usize, + dilation_h: usize, + dilation_w: usize, + pad_d: usize, + pad_h: usize, + pad_w: usize, + in_d: usize, + in_h: usize, + in_w: usize, + channels_per_group: usize, + col_len: usize, + b: usize, + in_c_start: usize, + nhwc_stride: (usize, usize, usize, usize, usize), +) { + for (local_idx, global_idx) in (tile_start..tile_end).enumerate() { + let out_d_idx = global_idx / (out_h * out_w); + let rem = global_idx % (out_h * out_w); + let out_h_idx = rem / out_w; + let out_w_idx = rem % out_w; + + let mut col_offset = 0; + for kd in 0..kernel_d { + let id = (out_d_idx * stride_d + kd * dilation_d) as isize - pad_d as isize; + for kh in 0..kernel_h { + let ih = (out_h_idx * stride_h + kh * dilation_h) as isize - pad_h as isize; + for kw in 0..kernel_w { + let iw = (out_w_idx * stride_w + kw * dilation_w) as isize - pad_w as isize; + + if id >= 0 + && id < in_d as isize + && ih >= 0 + && ih < in_h as isize + && iw >= 0 + && iw < in_w as isize + { + let id = id as usize; + let ih = ih as usize; + let iw = iw as usize; + let inp_base = b * nhwc_stride.0 + + id * nhwc_stride.1 + + ih * nhwc_stride.2 + + iw * nhwc_stride.3 + + in_c_start; + for c in 0..channels_per_group { + col_tile[local_idx * col_len + col_offset] = x_nhwc[inp_base + c]; + col_offset += 1; + } + } else { + // Padding: skip channels_per_group positions (already zero) + col_offset += channels_per_group; + } + } + } + } + } +} + +/// Check if this is a 1x1 convolution that can use the fast path. +fn is_1x1_conv( + kernel_d: usize, + kernel_h: usize, + kernel_w: usize, + options: &ConvOptions<3>, +) -> bool { + kernel_d == 1 + && kernel_h == 1 + && kernel_w == 1 + && options.stride == [1, 1, 1] + && options.padding == [0, 0, 0] +} + +/// Optimized 1x1 convolution: skip im2col, call gemm directly on NCHW data. +/// +/// For 1x1 conv with stride=1 and padding=0, the input for each (batch, group) is +/// already [channels_per_group, spatial] row-major in NCHW layout. We pass it directly +/// to gemm as the RHS with appropriate strides, avoiding the transpose allocation +/// and the intermediate result buffer. +#[allow(clippy::too_many_arguments, clippy::type_complexity)] +fn conv3d_1x1_impl( + x: FlexTensor, + weight: FlexTensor, + bias: Option, + options: &ConvOptions<3>, + dtype: DType, + zero: T, + one: T, + add_fn: fn(T, T) -> T, +) -> FlexTensor { + let x = x.to_contiguous(); + let weight = weight.to_contiguous(); + + let x_shape = x.layout().shape(); + let w_shape = weight.layout().shape(); + + let batch_size = x_shape[0]; + let channels_in = x_shape[1]; + let spatial = x_shape[2] * x_shape[3] * x_shape[4]; + + let channels_out = w_shape[0]; + let channels_per_group = w_shape[1]; + let groups = options.groups; + let out_channels_per_group = channels_out / groups; + + let m = out_channels_per_group; + let k = channels_per_group; + let n = spatial; + + // Validate output size won't overflow + let total_output = [batch_size, channels_out, spatial] + .iter() + .try_fold(1usize, |acc, &x| acc.checked_mul(x)) + .expect("conv 1x1: output tensor dimensions would overflow index calculations"); + + let x_data: &[T] = x.storage(); + let w_data: &[T] = weight.storage(); + + // Gemm: C[M,N] = W[M,K] * X[K,N] + // W is [out_channels_per_group, channels_per_group] row-major + // X is [channels_per_group, spatial] row-major (directly from NCHW layout) + // C is [out_channels_per_group, spatial] row-major + #[cfg(feature = "rayon")] + let parallelism = if m.saturating_mul(n).saturating_mul(k) >= 192 * 192 * 192 { + gemm::Parallelism::Rayon(0) + } else { + gemm::Parallelism::None + }; + #[cfg(not(feature = "rayon"))] + let parallelism = gemm::Parallelism::None; + + let mut output = vec![zero; total_output]; + + { + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + let dst_ptr = crate::ops::SendMutPtr::new(output.as_mut_ptr()); + + (0..batch_size).into_par_iter().for_each(|b| { + for g in 0..groups { + let x_offset = b * channels_in * spatial + g * k * spatial; + let w_offset = g * out_channels_per_group * k; + let out_offset = + b * channels_out * spatial + g * out_channels_per_group * spatial; + + unsafe { + gemm::gemm( + m, + n, + k, + dst_ptr.ptr_add(out_offset), + 1, // dst_cs + n as isize, // dst_rs + false, // read_dst + w_data.as_ptr().add(w_offset), + 1, // lhs_cs: W row-major + k as isize, // lhs_rs + x_data.as_ptr().add(x_offset), + 1, // rhs_cs: X row-major + n as isize, // rhs_rs + zero, + one, + false, + false, + false, + parallelism, + ); + } + } + }); + } + #[cfg(not(feature = "rayon"))] + { + for b in 0..batch_size { + for g in 0..groups { + let x_offset = b * channels_in * spatial + g * k * spatial; + let w_offset = g * out_channels_per_group * k; + let out_offset = + b * channels_out * spatial + g * out_channels_per_group * spatial; + + unsafe { + gemm::gemm( + m, + n, + k, + output.as_mut_ptr().add(out_offset), + 1, + n as isize, + false, + w_data.as_ptr().add(w_offset), + 1, + k as isize, + x_data.as_ptr().add(x_offset), + 1, + n as isize, + zero, + one, + false, + false, + false, + parallelism, + ); + } + } + } + } + } + + if let Some(bias) = bias { + let bias = bias.to_contiguous(); + let bias_data: &[T] = bias.storage(); + add_bias( + &mut output, + bias_data, + batch_size, + channels_out, + spatial, + add_fn, + ); + } + + let out_shape = Shape::from(vec![ + batch_size, + channels_out, + x_shape[2], + x_shape[3], + x_shape[4], + ]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + dtype, + ) +} + +conv3d_1x1_typed!(conv3d_1x1_f32, f32, DType::F32, 0.0f32, 1.0f32, |a, b| a + + b); +conv3d_1x1_typed!(conv3d_1x1_f64, f64, DType::F64, 0.0f64, 1.0f64, |a, b| a + + b); +conv3d_1x1_typed!( + conv3d_1x1_f16, + f16, + DType::F16, + f16::from_f32(0.0), + f16::from_f32(1.0), + |a: f16, b: f16| f16::from_f32(a.to_f32() + b.to_f32()) +); + +// ============================================================================ +// Depthwise conv fast path (no NHWC, no im2col, no gemm) +// ============================================================================ +// +// For canonical depthwise convolutions where every input channel maps to +// exactly one output channel through its own filter (groups == c_in == c_out, +// channels_per_group == 1), the generic im2col + per-group gemm path pays +// enormous overhead: +// +// - Per (tile, group) heap allocation of a col_tile buffer. +// - Per (tile, group) gemm dispatch with M=1, K=kernel_spatial, N=tile_size. +// The gemm kernel's setup cost dominates the tiny inner compute. +// - Full NHWC conversion of the input even though no channel mixing occurs. +// +// The depthwise path drops all three. It walks NCHW data directly and, for +// each (batch, channel), accumulates the convolution by looping over kernel +// positions on the outside and spatial positions on the inside. The valid +// output range for each kernel position is computed analytically so the +// inner spatial loop has no padding checks and LLVM autovectorizes it in +// the common stride=1, dilation=1 case. + +/// Decide whether to use the depthwise fast path. +/// +/// Triggers on canonical depthwise 1D or 2D convolutions (restricted to +/// `kd == 1 && in_d == 1` after the 3D expansion used by conv1d/conv2d). +fn should_use_depthwise_conv( + x_shape: &[usize], + w_shape: &[usize], + options: &ConvOptions<3>, +) -> bool { + let channels_in = x_shape[1]; + let channels_out = w_shape[0]; + let channels_per_group = w_shape[1]; + let groups = options.groups; + + // Canonical depthwise: one input channel per group, one output channel per group. + if channels_per_group != 1 || groups != channels_in || channels_out != channels_in { + return false; + } + + // Only the 1D/2D-in-3D shapes (kd=1, in_d=1) produced by the conv1d/conv2d + // expansion. Pure 3D depthwise would need additional loop nesting. + if w_shape[2] != 1 || x_shape[2] != 1 { + return false; + } + if options.stride[0] != 1 || options.padding[0] != 0 || options.dilation[0] != 1 { + return false; + } + + true +} + +/// Compute the half-open range `[out_start, out_end)` of output positions `o` +/// for which the corresponding input index `o * stride + k * dilation - pad` +/// lies inside `[0, in_size)`. +#[inline] +fn valid_out_range( + k: usize, + dilation: usize, + pad: usize, + stride: usize, + in_size: usize, + out_size: usize, +) -> (usize, usize) { + debug_assert!(stride >= 1, "stride must be >= 1"); + let offset = k * dilation; + + // Lower: smallest o with o*stride + offset >= pad. + let out_start = if offset >= pad { + 0 + } else { + (pad - offset).div_ceil(stride) + }; + + // Upper (exclusive): smallest o with o*stride + offset - pad >= in_size, + // i.e. smallest o with o*stride >= in_size + pad - offset. + let threshold = in_size + pad; + let out_end = if offset >= threshold { + 0 + } else { + (threshold - offset).div_ceil(stride) + }; + + let out_end = out_end.min(out_size); + let out_start = out_start.min(out_end); + (out_start, out_end) +} + +/// Output-plane element count above which `conv_plane_accumulate` switches +/// from the `kh, kw, oh, ow` "kh-outer" loop order to `oh, kh, kw, ow` +/// "oh-outer". +/// +/// Below the threshold the whole plane fits comfortably in L1 (~32 KB on +/// most modern CPUs), so the hardware prefetcher tracks the regular +/// `oh`-stride access and the kh-outer order amortizes per-kernel-position +/// setup over long inner runs. Above the threshold the plane no longer +/// fits, and the kh-outer order refetches the output from L2/DRAM once +/// per `(kh, kw)` pair: a `kh * kw`-fold amplification of output memory +/// traffic. Flipping to oh-outer pins one output row in L1 across every +/// kernel position and traverses the plane exactly once. +/// +/// 8192 f32 elements = 32 KB, i.e. L1 data-cache size on M3/M4 Max and +/// most modern x86 server parts. The gap in the conv benchmarks is clean: +/// every depthwise shape that regressed under pure oh-outer was <= 3136 +/// elements, while the Sobel / preproc / mask shapes that win under +/// oh-outer are all > 65000 elements. +const CONV_PLANE_OH_OUTER_THRESHOLD: usize = 8192; + +/// Accumulate one 2D conv plane: `out_plane += conv2d(in_plane, w_plane)` using +/// the precomputed analytic `oh_ranges`/`ow_ranges` to skip padding checks in +/// the inner loop. +/// +/// **Precondition**: `out_plane` must already hold the running accumulator +/// (zero on the first call, or whatever partial sum has been built up across +/// previous `ci` iterations). The function reads each output element before +/// writing, so an uninitialized buffer produces silent garbage. +/// +/// Dispatches to one of two loop orders based on `out_plane.len()`; see +/// `CONV_PLANE_OH_OUTER_THRESHOLD` for the tradeoff. Shared by +/// `conv3d_depthwise_impl` and `conv3d_small_channel_impl`. The dispatcher +/// is `#[inline]` (not `inline(always)`) so the runtime length branch lives +/// once at each call site instead of inlining both variants; the variants +/// themselves stay `inline(always)` so LLVM sees the concrete inner loop +/// pattern and emits SIMD fmuladd. Using `num_traits::Float` bounds (rather +/// than fn-pointer arithmetic) is load-bearing for vectorization. +#[inline] +#[allow(clippy::too_many_arguments)] +fn conv_plane_accumulate( + out_plane: &mut [T], + in_plane: &[T], + w_plane: &[T], + kernel_h: usize, + kernel_w: usize, + in_w: usize, + out_w: usize, + stride_h: usize, + stride_w: usize, + pad_h: usize, + pad_w: usize, + dilation_h: usize, + dilation_w: usize, + oh_ranges: &[(usize, usize)], + ow_ranges: &[(usize, usize)], +) { + if out_plane.len() > CONV_PLANE_OH_OUTER_THRESHOLD { + conv_plane_accumulate_oh_outer( + out_plane, in_plane, w_plane, kernel_h, kernel_w, in_w, out_w, stride_h, stride_w, + pad_h, pad_w, dilation_h, dilation_w, oh_ranges, ow_ranges, + ); + } else { + conv_plane_accumulate_kh_outer( + out_plane, in_plane, w_plane, kernel_h, kernel_w, in_w, out_w, stride_h, stride_w, + pad_h, pad_w, dilation_h, dilation_w, oh_ranges, ow_ranges, + ); + } +} + +/// `oh`-outermost variant. Pins one output row in L1 across every kernel +/// position that accumulates into it, so each row is touched exactly once +/// regardless of how large the full output plane is. Selected when the +/// plane exceeds L1 (see `CONV_PLANE_OH_OUTER_THRESHOLD`). +#[inline(always)] +#[allow(clippy::too_many_arguments)] +fn conv_plane_accumulate_oh_outer( + out_plane: &mut [T], + in_plane: &[T], + w_plane: &[T], + kernel_h: usize, + kernel_w: usize, + in_w: usize, + out_w: usize, + stride_h: usize, + stride_w: usize, + pad_h: usize, + pad_w: usize, + dilation_h: usize, + dilation_w: usize, + oh_ranges: &[(usize, usize)], + ow_ranges: &[(usize, usize)], +) { + // An empty plane or zero-width output is a trivial no-op. This also + // guards the divide below against `out_w == 0`, which is not reachable + // from the in-tree callers (burn's `calculate_conv_output_size` is + // always >= 1 for valid inputs) but would otherwise panic if some + // future caller handed us a degenerate slice. + if out_plane.is_empty() || out_w == 0 { + return; + } + + // out_plane is a per-(batch, channel) slice of shape [out_h, out_w] + // stored contiguously; both call sites produce it by disjoint splitting + // of a `vec![zero; batch * channels * out_h * out_w]` allocation, so + // the length is always an exact multiple of `out_w`. The `debug_assert` + // turns that documentary invariant into an enforceable one. + debug_assert_eq!( + out_plane.len() % out_w, + 0, + "out_plane length must be a whole number of rows" + ); + let out_h = out_plane.len() / out_w; + + for oh in 0..out_h { + let out_row = &mut out_plane[oh * out_w..(oh + 1) * out_w]; + + for kh in 0..kernel_h { + let (oh_start, oh_end) = oh_ranges[kh]; + // Skip kernel rows that fall outside the padded image at this oh. + if oh < oh_start || oh >= oh_end { + continue; + } + let ih = oh * stride_h + kh * dilation_h - pad_h; + let in_row = &in_plane[ih * in_w..(ih + 1) * in_w]; + + for kw in 0..kernel_w { + let (ow_start, ow_end) = ow_ranges[kw]; + if ow_start >= ow_end { + continue; + } + let w_val = w_plane[kh * kernel_w + kw]; + // All terms are non-negative because ow_start was chosen so + // that the corresponding `iw` is in bounds. + let iw_start = ow_start * stride_w + kw * dilation_w - pad_w; + + if stride_w == 1 { + let run_len = ow_end - ow_start; + let in_slice = &in_row[iw_start..iw_start + run_len]; + let out_slice = &mut out_row[ow_start..ow_end]; + for (o, &xv) in out_slice.iter_mut().zip(in_slice.iter()) { + *o = *o + w_val * xv; + } + } else { + let mut iw = iw_start; + for o in &mut out_row[ow_start..ow_end] { + *o = *o + w_val * in_row[iw]; + iw += stride_w; + } + } + } + } + } +} + +/// `kh, kw`-outermost variant. Amortizes per-kernel-position setup over +/// long regular `oh` runs that the hardware prefetcher can track. +/// Selected when the whole output plane already fits in L1 (see +/// `CONV_PLANE_OH_OUTER_THRESHOLD`): the oh-outer trick buys nothing +/// there and the extra per-iteration bookkeeping slightly hurts. +#[inline(always)] +#[allow(clippy::too_many_arguments)] +fn conv_plane_accumulate_kh_outer( + out_plane: &mut [T], + in_plane: &[T], + w_plane: &[T], + kernel_h: usize, + kernel_w: usize, + in_w: usize, + out_w: usize, + stride_h: usize, + stride_w: usize, + pad_h: usize, + pad_w: usize, + dilation_h: usize, + dilation_w: usize, + oh_ranges: &[(usize, usize)], + ow_ranges: &[(usize, usize)], +) { + for kh in 0..kernel_h { + let (oh_start, oh_end) = oh_ranges[kh]; + if oh_start >= oh_end { + continue; + } + for kw in 0..kernel_w { + let (ow_start, ow_end) = ow_ranges[kw]; + if ow_start >= ow_end { + continue; + } + let w_val = w_plane[kh * kernel_w + kw]; + let iw_start = ow_start * stride_w + kw * dilation_w - pad_w; + let run_len = ow_end - ow_start; + for oh in oh_start..oh_end { + let ih = oh * stride_h + kh * dilation_h - pad_h; + let in_row = &in_plane[ih * in_w..(ih + 1) * in_w]; + let out_row = &mut out_plane[oh * out_w..(oh + 1) * out_w]; + if stride_w == 1 { + let in_slice = &in_row[iw_start..iw_start + run_len]; + let out_slice = &mut out_row[ow_start..ow_end]; + for (o, &xv) in out_slice.iter_mut().zip(in_slice.iter()) { + *o = *o + w_val * xv; + } + } else { + let mut iw = iw_start; + for o in &mut out_row[ow_start..ow_end] { + *o = *o + w_val * in_row[iw]; + iw += stride_w; + } + } + } + } + } +} + +macro_rules! conv3d_depthwise_typed { + ($fn_name:ident, $T:ty, $dtype:expr) => { + fn $fn_name( + x: FlexTensor, + weight: FlexTensor, + bias: Option, + options: &ConvOptions<3>, + ) -> FlexTensor { + conv3d_depthwise_impl::<$T>(x, weight, bias, options, $dtype) + } + }; +} + +conv3d_depthwise_typed!(conv3d_depthwise_f32, f32, DType::F32); +conv3d_depthwise_typed!(conv3d_depthwise_f64, f64, DType::F64); +conv3d_depthwise_typed!(conv3d_depthwise_f16, f16, DType::F16); + +/// Depthwise conv3d: one filter per channel, no channel mixing. +/// +/// Preconditions (checked by `should_use_depthwise_conv`): +/// - `channels_per_group == 1` +/// - `groups == channels_in == channels_out` +/// - `kernel_d == 1`, `in_d == 1`, trivial d-axis options. +/// +/// Uses `num_traits::Float` bounds so the inner multiply-accumulate compiles +/// down to direct `fmul`/`fadd` instructions that LLVM can autovectorize. +/// Function-pointer arithmetic (`fn(T, T) -> T`) would prevent vectorization +/// because each call is an indirect branch that blocks the loop pattern match. +fn conv3d_depthwise_impl( + x: FlexTensor, + weight: FlexTensor, + bias: Option, + options: &ConvOptions<3>, + dtype: DType, +) -> FlexTensor +where + T: num_traits::Float + bytemuck::Pod + Clone + Copy + burn_backend::Element + Send + Sync, +{ + let zero = ::zero(); + let x = x.to_contiguous(); + let weight = weight.to_contiguous(); + + let x_shape = x.layout().shape(); + let w_shape = weight.layout().shape(); + + let batch_size = x_shape[0]; + let channels = x_shape[1]; + let in_h = x_shape[3]; + let in_w = x_shape[4]; + + let kernel_h = w_shape[3]; + let kernel_w = w_shape[4]; + + let [_, stride_h, stride_w] = options.stride; + let [_, pad_h, pad_w] = options.padding; + let [_, dilation_h, dilation_w] = options.dilation; + + let out_h = calculate_conv_output_size(kernel_h, stride_h, pad_h, dilation_h, in_h); + let out_w = calculate_conv_output_size(kernel_w, stride_w, pad_w, dilation_w, in_w); + + let total = [batch_size, channels, out_h, out_w] + .iter() + .try_fold(1usize, |acc, &x| acc.checked_mul(x)) + .expect("conv depthwise: output dimensions would overflow"); + + let x_data: &[T] = x.storage(); + let w_data: &[T] = weight.storage(); + + let in_spatial = in_h * in_w; + let out_spatial = out_h * out_w; + let k_spatial = kernel_h * kernel_w; + + // Valid oh/ow ranges are identical for every (b, c), so precompute once. + let oh_ranges: Vec<(usize, usize)> = (0..kernel_h) + .map(|kh| valid_out_range(kh, dilation_h, pad_h, stride_h, in_h, out_h)) + .collect(); + let ow_ranges: Vec<(usize, usize)> = (0..kernel_w) + .map(|kw| valid_out_range(kw, dilation_w, pad_w, stride_w, in_w, out_w)) + .collect(); + + let mut output = vec![zero; total]; + + // Per-plane work for one `(batch, channel)` pair. Output slice is passed + // in so the rayon and sequential dispatch paths can source it differently + // (from a raw pointer or a direct borrow) without duplicating this body. + let plane_work = |bc: usize, out_plane: &mut [T]| { + let c = bc % channels; + let in_base = bc * in_spatial; + let w_base = c * k_spatial; + conv_plane_accumulate( + out_plane, + &x_data[in_base..in_base + in_spatial], + &w_data[w_base..w_base + k_spatial], + kernel_h, + kernel_w, + in_w, + out_w, + stride_h, + stride_w, + pad_h, + pad_w, + dilation_h, + dilation_w, + &oh_ranges, + &ow_ranges, + ); + }; + + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + + let dst_ptr = crate::ops::SendMutPtr::new(output.as_mut_ptr()); + (0..batch_size * channels).into_par_iter().for_each(|bc| { + // SAFETY: disjoint `[bc * out_spatial, (bc+1) * out_spatial)` + // ranges tile the output buffer. + let out_plane: &mut [T] = unsafe { + core::slice::from_raw_parts_mut(dst_ptr.ptr_add(bc * out_spatial), out_spatial) + }; + plane_work(bc, out_plane); + }); + } + #[cfg(not(feature = "rayon"))] + { + for bc in 0..batch_size * channels { + let out_base = bc * out_spatial; + plane_work(bc, &mut output[out_base..out_base + out_spatial]); + } + } + + if let Some(bias) = bias { + let bias = bias.to_contiguous(); + let bias_data: &[T] = bias.storage(); + assert_eq!( + bias_data.len(), + channels, + "conv depthwise: bias length ({}) must equal channels ({channels})", + bias_data.len() + ); + add_bias( + &mut output, + bias_data, + batch_size, + channels, + out_spatial, + |a, b| a + b, + ); + } + + let out_shape = Shape::from(vec![batch_size, channels, 1, out_h, out_w]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + dtype, + ) +} + +// ============================================================================ +// Small-channel conv fast path (no NHWC, no im2col, no gemm) +// ============================================================================ +// +// For `groups=1` convs with very few input channels (e.g. 3-channel Sobel +// filters, single-channel mask networks, early-stage image preprocessors), +// the generic `conv3d_impl` pays the same kind of overhead as the depthwise +// case: gemm is dispatched with `M = channels_out, K = channels_in * k_spatial, +// N = tile_size`, and at small `K` the gemm kernel's pack + select + dispatch +// cost dominates the actual FMAs. On top of that, the full NHWC conversion +// of the input is pure memory traffic that buys nothing when the channel +// dimension is already tiny. +// +// The small-channel path walks NCHW data directly. For each `(batch, out_ch)` +// pair it accumulates contributions from every input channel by calling the +// same `conv_plane_accumulate` helper used by the depthwise path. Compared +// to the depthwise case this adds an outer `for ci in 0..channels_in` loop +// around the accumulate; everything else (analytic output ranges, parallel +// fan-out, bias add) is identical. + +/// Threshold on `channels_in` for the small-channel fast path. +/// +/// Picked empirically: at `channels_in <= 4`, gemm dispatch overhead exceeds +/// the inner compute by a clear margin. Tuning higher risks regressing shapes +/// where gemm's data reuse across output channels wins. +const SMALL_CHANNEL_IN_THRESHOLD: usize = 4; + +/// Threshold on `channels_out` for the small-channel fast path. +/// +/// The small-channel path loops `for co: for ci:` and re-reads each input +/// channel once per output channel. At large `channels_out`, that redundant +/// memory traffic dominates over gemm's register-tiled data reuse. Picked +/// empirically: classic ImageNet first-layer (`3 -> 64`) runs ~5x slower on +/// the direct path than on gemm, so we cap at 16 to stay safely on the right +/// side of that cliff. Sobel-style filters (`3 -> 3..8`) are well under this +/// and keep their 1.1-1.4x win over burn-ndarray. +const SMALL_CHANNEL_OUT_THRESHOLD: usize = 16; + +/// Decide whether to use the small-channel fast path. +/// +/// Triggers on `groups=1` 2D-in-3D convs (or 1D via the conv1d expansion) +/// with both small `channels_in` and small `channels_out`. Non-depthwise +/// grouped convs and the many-channel case both go through the generic +/// `conv3d_impl`. +fn should_use_small_channel_conv( + x_shape: &[usize], + w_shape: &[usize], + options: &ConvOptions<3>, +) -> bool { + if options.groups != 1 { + return false; + } + + // Only the 1D/2D-in-3D shapes (kd=1, in_d=1). Pure 3D convs do not benefit + // and would require adding a d-axis loop. + if w_shape[2] != 1 || x_shape[2] != 1 { + return false; + } + if options.stride[0] != 1 || options.padding[0] != 0 || options.dilation[0] != 1 { + return false; + } + + let channels_in = x_shape[1]; + let channels_out = w_shape[0]; + channels_in > 0 + && channels_in <= SMALL_CHANNEL_IN_THRESHOLD + && channels_out > 0 + && channels_out <= SMALL_CHANNEL_OUT_THRESHOLD +} + +macro_rules! conv3d_small_channel_typed { + ($fn_name:ident, $T:ty, $dtype:expr) => { + fn $fn_name( + x: FlexTensor, + weight: FlexTensor, + bias: Option, + options: &ConvOptions<3>, + ) -> FlexTensor { + conv3d_small_channel_impl::<$T>(x, weight, bias, options, $dtype) + } + }; +} + +conv3d_small_channel_typed!(conv3d_small_channel_f32, f32, DType::F32); +conv3d_small_channel_typed!(conv3d_small_channel_f64, f64, DType::F64); +conv3d_small_channel_typed!(conv3d_small_channel_f16, f16, DType::F16); + +/// Small-channel conv3d: `groups=1` with small `channels_in`. +/// +/// Preconditions (checked by `should_use_small_channel_conv`): +/// - `options.groups == 1` +/// - `channels_in <= SMALL_CHANNEL_IN_THRESHOLD` +/// - `kernel_d == 1`, `in_d == 1`, trivial d-axis options. +fn conv3d_small_channel_impl( + x: FlexTensor, + weight: FlexTensor, + bias: Option, + options: &ConvOptions<3>, + dtype: DType, +) -> FlexTensor +where + T: num_traits::Float + bytemuck::Pod + Clone + Copy + burn_backend::Element + Send + Sync, +{ + let zero = ::zero(); + let x = x.to_contiguous(); + let weight = weight.to_contiguous(); + + let x_shape = x.layout().shape(); + let w_shape = weight.layout().shape(); + + let batch_size = x_shape[0]; + let channels_in = x_shape[1]; + let in_h = x_shape[3]; + let in_w = x_shape[4]; + + let channels_out = w_shape[0]; + let kernel_h = w_shape[3]; + let kernel_w = w_shape[4]; + + let [_, stride_h, stride_w] = options.stride; + let [_, pad_h, pad_w] = options.padding; + let [_, dilation_h, dilation_w] = options.dilation; + + let out_h = calculate_conv_output_size(kernel_h, stride_h, pad_h, dilation_h, in_h); + let out_w = calculate_conv_output_size(kernel_w, stride_w, pad_w, dilation_w, in_w); + + let total = [batch_size, channels_out, out_h, out_w] + .iter() + .try_fold(1usize, |acc, &x| acc.checked_mul(x)) + .expect("conv small-channel: output dimensions would overflow"); + + let x_data: &[T] = x.storage(); + let w_data: &[T] = weight.storage(); + + let in_spatial = in_h * in_w; + let out_spatial = out_h * out_w; + let k_spatial = kernel_h * kernel_w; + let w_co_stride = channels_in * k_spatial; + let x_batch_stride = channels_in * in_spatial; + + let oh_ranges: Vec<(usize, usize)> = (0..kernel_h) + .map(|kh| valid_out_range(kh, dilation_h, pad_h, stride_h, in_h, out_h)) + .collect(); + let ow_ranges: Vec<(usize, usize)> = (0..kernel_w) + .map(|kw| valid_out_range(kw, dilation_w, pad_w, stride_w, in_w, out_w)) + .collect(); + + let mut output = vec![zero; total]; + + // Per-plane work for one `(batch, out_channel)` pair. Accumulates + // `sum over ci of conv2d(in[b, ci], w[co, ci])` into `out_plane` by + // calling the shared helper once per input channel. + let plane_work = |b_co: usize, out_plane: &mut [T]| { + let b = b_co / channels_out; + let co = b_co % channels_out; + for ci in 0..channels_in { + let in_base = b * x_batch_stride + ci * in_spatial; + let w_base = co * w_co_stride + ci * k_spatial; + conv_plane_accumulate( + out_plane, + &x_data[in_base..in_base + in_spatial], + &w_data[w_base..w_base + k_spatial], + kernel_h, + kernel_w, + in_w, + out_w, + stride_h, + stride_w, + pad_h, + pad_w, + dilation_h, + dilation_w, + &oh_ranges, + &ow_ranges, + ); + } + }; + + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + + let dst_ptr = crate::ops::SendMutPtr::new(output.as_mut_ptr()); + (0..batch_size * channels_out) + .into_par_iter() + .for_each(|b_co| { + // SAFETY: disjoint `[b_co * out_spatial, (b_co+1) * out_spatial)` + // ranges tile the output buffer. + let out_plane: &mut [T] = unsafe { + core::slice::from_raw_parts_mut( + dst_ptr.ptr_add(b_co * out_spatial), + out_spatial, + ) + }; + plane_work(b_co, out_plane); + }); + } + #[cfg(not(feature = "rayon"))] + { + for b_co in 0..batch_size * channels_out { + let out_base = b_co * out_spatial; + plane_work(b_co, &mut output[out_base..out_base + out_spatial]); + } + } + + if let Some(bias) = bias { + let bias = bias.to_contiguous(); + let bias_data: &[T] = bias.storage(); + assert_eq!( + bias_data.len(), + channels_out, + "conv small-channel: bias length ({}) must equal channels_out ({channels_out})", + bias_data.len() + ); + add_bias( + &mut output, + bias_data, + batch_size, + channels_out, + out_spatial, + |a, b| a + b, + ); + } + + let out_shape = Shape::from(vec![batch_size, channels_out, 1, out_h, out_w]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + dtype, + ) +} + +// ============================================================================ +// Direct conv path for small spatial sizes (no im2col) +// ============================================================================ + +/// Decide whether to use the direct conv path instead of tiled im2col + GEMM. +/// +/// The direct path decomposes the conv into kw gemm calls with strided pointers +/// into NCHW data, eliminating both the NHWC conversion and im2col buffer. +/// This wins when spatial_out is small enough that im2col allocation and fill +/// overhead is significant relative to compute. +fn should_use_direct_conv(x_shape: &[usize], w_shape: &[usize], options: &ConvOptions<3>) -> bool { + // Only for groups=1, no padding, dilation=1 (the wav2vec2 case). + if options.groups != 1 || options.padding != [0, 0, 0] || options.dilation != [1, 1, 1] { + return false; + } + + let kernel_d = w_shape[2]; + let kernel_h = w_shape[3]; + let kernel_w = w_shape[4]; + + // Only for "1D-like" convolutions expanded to 3D + if kernel_d != 1 || kernel_h != 1 { + return false; + } + if x_shape[2] != 1 || x_shape[3] != 1 { + return false; + } + + let channels_in = x_shape[1]; + let in_w = x_shape[4]; + let out_w = calculate_conv_output_size(kernel_w, options.stride[2], 0, 1, in_w); + + // The direct path eliminates im2col buffer allocation (kw * c_in * tile_size floats) + // and the copy into it. This matters when spatial_out is small relative to c_in * kw. + channels_in >= 32 && kernel_w <= 8 && out_w <= 800 +} + +macro_rules! conv3d_direct_typed { + ($fn_name:ident, $T:ty, $dtype:expr, $zero:expr, $one:expr, $add_fn:expr) => { + fn $fn_name( + x: FlexTensor, + weight: FlexTensor, + bias: Option, + options: &ConvOptions<3>, + ) -> FlexTensor { + conv3d_direct_impl::<$T>(x, weight, bias, options, $dtype, $zero, $one, $add_fn) + } + }; +} + +conv3d_direct_typed!( + conv3d_direct_f32, + f32, + DType::F32, + 0.0f32, + 1.0f32, + |a, b| a + b +); +conv3d_direct_typed!( + conv3d_direct_f64, + f64, + DType::F64, + 0.0f64, + 1.0f64, + |a, b| a + b +); + +/// Direct conv3d: decompose into kw gemm calls on NCHW data directly. +/// +/// The conv operation is: out[co, o] = sum_k sum_c w[co, c, k] * x[c, o*stride+k] +/// +/// This decomposes into kw matrix multiplies, one per kernel position: +/// out += W_k @ X_k where W_k = w[:, :, k] and X_k = x[:, o*stride+k] +/// +/// Each gemm reads directly from NCHW input with strided pointers, eliminating +/// both the NHWC conversion and im2col buffer entirely. +/// +/// Constraints: groups=1, padding=0, dilation=1, 1D-like (d=1, h=1). +#[allow(clippy::too_many_arguments, clippy::type_complexity)] +fn conv3d_direct_impl( + x: FlexTensor, + weight: FlexTensor, + bias: Option, + options: &ConvOptions<3>, + dtype: DType, + zero: T, + one: T, + add_fn: fn(T, T) -> T, +) -> FlexTensor { + let x = x.to_contiguous(); + let weight = weight.to_contiguous(); + + let x_shape = x.layout().shape(); + let w_shape = weight.layout().shape(); + + let batch_size = x_shape[0]; + let channels_in = x_shape[1]; + let in_w = x_shape[4]; + + let channels_out = w_shape[0]; + let kernel_w = w_shape[4]; + let stride_w = options.stride[2]; + + let out_w = calculate_conv_output_size(kernel_w, stride_w, 0, 1, in_w); + + let x_data: &[T] = x.storage(); + let w_data: &[T] = weight.storage(); + + // Weight layout: [c_out, c_in, kw] contiguous (kd=kh=1). + // For kernel position k: W_k[co, c] = w_data[co * c_in * kw + c * kw + k] + let lhs_rs = (channels_in * kernel_w) as isize; + let lhs_cs = kernel_w as isize; + + // For kernel position k: X_k[c_in, o] = x_data[c_in * in_w + o * stride + k] + let rhs_rs = in_w as isize; + let rhs_cs = stride_w as isize; + + let m = channels_out; + let gemm_k = channels_in; + let n = out_w; + + #[cfg(feature = "rayon")] + let parallelism = if m.saturating_mul(n).saturating_mul(gemm_k) >= 192 * 192 * 192 { + gemm::Parallelism::Rayon(0) + } else { + gemm::Parallelism::None + }; + #[cfg(not(feature = "rayon"))] + let parallelism = gemm::Parallelism::None; + + let total_output = batch_size + .checked_mul(channels_out) + .and_then(|x| x.checked_mul(out_w)) + .expect("conv direct: output dimensions overflow"); + let mut output = vec![zero; total_output]; + + let batch_x_len = channels_in * in_w; + + { + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + let dst_ptr = crate::ops::SendMutPtr::new(output.as_mut_ptr()); + + (0..batch_size).into_par_iter().for_each(|b| { + let out_offset = b * channels_out * out_w; + + for k in 0..kernel_w { + unsafe { + let x_base = x_data.as_ptr().add(b * batch_x_len); + gemm::gemm( + m, + n, + gemm_k, + dst_ptr.ptr_add(out_offset), + 1, + n as isize, + k > 0, + w_data.as_ptr().add(k), + lhs_cs, + lhs_rs, + x_base.add(k), + rhs_cs, + rhs_rs, + one, + one, + false, + false, + false, + parallelism, + ); + } + } + }); + } + #[cfg(not(feature = "rayon"))] + { + for b in 0..batch_size { + let out_offset = b * channels_out * out_w; + + for k in 0..kernel_w { + unsafe { + let x_base = x_data.as_ptr().add(b * batch_x_len); + gemm::gemm( + m, + n, + gemm_k, + output.as_mut_ptr().add(out_offset), + 1, + n as isize, + k > 0, + w_data.as_ptr().add(k), + lhs_cs, + lhs_rs, + x_base.add(k), + rhs_cs, + rhs_rs, + one, + one, + false, + false, + false, + parallelism, + ); + } + } + } + } + } + + if let Some(bias) = bias { + let bias = bias.to_contiguous(); + let bias_data: &[T] = bias.storage(); + add_bias( + &mut output, + bias_data, + batch_size, + channels_out, + out_w, + add_fn, + ); + } + + let out_shape = Shape::from(vec![batch_size, channels_out, 1, 1, out_w]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + dtype, + ) +} + +// ============================================================================ +// gemm implementations +// ============================================================================ +// +// One thin wrapper per element type around `gemm::gemm` with the strides and +// parallelism heuristic fixed to what conv3d_impl needs. The wrappers share +// every line of logic aside from the numeric type, zero, and one literals, +// so they're generated via a macro. + +macro_rules! gemm_typed { + ($fn_name:ident, $T:ty, $zero:expr, $one:expr) => { + fn $fn_name(a: &[$T], b: &[$T], m: usize, k: usize, n: usize) -> Vec<$T> { + let mut c = vec![$zero; m * n]; + #[cfg(feature = "rayon")] + let parallelism = if m * n * k >= 192 * 192 * 192 { + gemm::Parallelism::Rayon(0) + } else { + gemm::Parallelism::None + }; + #[cfg(not(feature = "rayon"))] + let parallelism = gemm::Parallelism::None; + unsafe { + gemm::gemm( + m, + n, + k, + c.as_mut_ptr(), + 1, + n as isize, + false, + a.as_ptr(), + 1, + k as isize, + b.as_ptr(), + k as isize, + 1, + $zero, + $one, + false, + false, + false, + parallelism, + ); + } + c + } + }; +} + +gemm_typed!(gemm_f32, f32, 0.0f32, 1.0f32); +gemm_typed!(gemm_f64, f64, 0.0f64, 1.0f64); +gemm_typed!(gemm_f16, f16, f16::from_f32(0.0), f16::from_f32(1.0)); + +// ============================================================================ +// Tests +// ============================================================================ + +// Tests kept here exercise flex-specific dtype storage paths (f64/f16/bf16) +// and flex-internal fast-path dispatch (1x1, depthwise, small-channel, +// direct, oh-outer) plus validation/panic checks on flex helpers. Plain +// conv shape/stride/padding/groups correctness is covered at the public +// API level by crates/burn-backend-tests/tests/tensor/float/module/ +// conv{1,2,3}d.rs, so those tests were removed from here. When adding +// new tests, keep them here only if they probe flex dtype dispatch or a +// flex-internal fast path; otherwise add them there. +#[cfg(test)] +mod tests { + use super::*; + use burn_backend::TensorData; + use burn_std::bf16; + + #[test] + fn test_conv1d_direct_path() { + // c_in=64 >= 32, kernel=3, stride=2, out_w=49 <= 800 -> hits direct path + let c_in = 64; + let c_out = 32; + let in_w = 100; + let kw = 3; + let stride = 2; + let out_w = (in_w - kw) / stride + 1; // 49 + + // Deterministic input and weights + let x_data: Vec = (0..c_in * in_w) + .map(|i| ((i % 100) as f32 / 100.0) - 0.5) + .collect(); + let w_data: Vec = (0..c_out * c_in * kw) + .map(|i| ((i % 50) as f32 / 50.0) - 0.5) + .collect(); + + let x = FlexTensor::from_data(TensorData::new(x_data.clone(), vec![1, c_in, in_w])); + let weight = FlexTensor::from_data(TensorData::new(w_data.clone(), vec![c_out, c_in, kw])); + let options = ConvOptions::new([stride], [0], [1], 1); + let result = conv1d_f32(x, weight, None, &options); + let out: Vec = result.into_data().to_vec().unwrap(); + + assert_eq!(out.len(), c_out * out_w); + + // Verify against naive reference implementation + for co in 0..c_out { + for o in 0..out_w { + let mut expected = 0.0f32; + for ci in 0..c_in { + for k in 0..kw { + expected += w_data[co * c_in * kw + ci * kw + k] + * x_data[ci * in_w + o * stride + k]; + } + } + let actual = out[co * out_w + o]; + assert!( + (actual - expected).abs() < 1e-3, + "mismatch at co={co}, o={o}: expected {expected}, got {actual}" + ); + } + } + } + + #[test] + fn test_conv1d_direct_path_kw2() { + // Test with kw=2 (L5/L6 shapes use kw=2) + let c_in = 64; + let c_out = 32; + let in_w = 50; + let kw = 2; + let stride = 2; + let out_w = (in_w - kw) / stride + 1; // 24 + + let x_data: Vec = (0..c_in * in_w) + .map(|i| ((i % 100) as f32 / 100.0) - 0.5) + .collect(); + let w_data: Vec = (0..c_out * c_in * kw) + .map(|i| ((i % 50) as f32 / 50.0) - 0.5) + .collect(); + + let x = FlexTensor::from_data(TensorData::new(x_data.clone(), vec![1, c_in, in_w])); + let weight = FlexTensor::from_data(TensorData::new(w_data.clone(), vec![c_out, c_in, kw])); + let options = ConvOptions::new([stride], [0], [1], 1); + let result = conv1d_f32(x, weight, None, &options); + let out: Vec = result.into_data().to_vec().unwrap(); + + for co in 0..c_out { + for o in 0..out_w { + let mut expected = 0.0f32; + for ci in 0..c_in { + for k in 0..kw { + expected += w_data[co * c_in * kw + ci * kw + k] + * x_data[ci * in_w + o * stride + k]; + } + } + let actual = out[co * out_w + o]; + assert!( + (actual - expected).abs() < 1e-3, + "mismatch at co={co}, o={o}: expected {expected}, got {actual}" + ); + } + } + } + + #[test] + fn test_conv1d_direct_path_f64() { + let c_in = 64; + let c_out = 16; + let in_w = 50; + let kw = 3; + let stride = 2; + let out_w = (in_w - kw) / stride + 1; + + let x_data: Vec = (0..c_in * in_w) + .map(|i| ((i % 100) as f64 / 100.0) - 0.5) + .collect(); + let w_data: Vec = (0..c_out * c_in * kw) + .map(|i| ((i % 50) as f64 / 50.0) - 0.5) + .collect(); + + let x = FlexTensor::from_data(TensorData::new(x_data.clone(), vec![1, c_in, in_w])); + let weight = FlexTensor::from_data(TensorData::new(w_data.clone(), vec![c_out, c_in, kw])); + let options = ConvOptions::new([stride], [0], [1], 1); + let result = conv1d_f64(x, weight, None, &options); + let out: Vec = result.into_data().to_vec().unwrap(); + + for co in 0..c_out { + for o in 0..out_w { + let mut expected = 0.0f64; + for ci in 0..c_in { + for k in 0..kw { + expected += w_data[co * c_in * kw + ci * kw + k] + * x_data[ci * in_w + o * stride + k]; + } + } + let actual = out[co * out_w + o]; + assert!( + (actual - expected).abs() < 1e-10, + "f64 mismatch at co={co}, o={o}: expected {expected}, got {actual}" + ); + } + } + } + + #[test] + fn test_conv1d_direct_path_with_bias() { + let c_in = 64; + let c_out = 32; + let in_w = 50; + let kw = 2; + let stride = 2; + let out_w = (in_w - kw) / stride + 1; + + let x_data: Vec = (0..c_in * in_w) + .map(|i| ((i % 100) as f32 / 100.0) - 0.5) + .collect(); + let w_data: Vec = (0..c_out * c_in * kw) + .map(|i| ((i % 50) as f32 / 50.0) - 0.5) + .collect(); + let bias_data: Vec = (0..c_out).map(|i| i as f32 * 0.1).collect(); + + let x = FlexTensor::from_data(TensorData::new(x_data.clone(), vec![1, c_in, in_w])); + let weight = FlexTensor::from_data(TensorData::new(w_data.clone(), vec![c_out, c_in, kw])); + let bias = FlexTensor::from_data(TensorData::new(bias_data.clone(), vec![c_out])); + let options = ConvOptions::new([stride], [0], [1], 1); + let result = conv1d_f32(x, weight, Some(bias), &options); + let out: Vec = result.into_data().to_vec().unwrap(); + + for co in 0..c_out { + for o in 0..out_w { + let mut expected = bias_data[co]; + for ci in 0..c_in { + for k in 0..kw { + expected += w_data[co * c_in * kw + ci * kw + k] + * x_data[ci * in_w + o * stride + k]; + } + } + let actual = out[co * out_w + o]; + assert!( + (actual - expected).abs() < 1e-3, + "bias mismatch at co={co}, o={o}: expected {expected}, got {actual}" + ); + } + } + } + + #[test] + fn test_conv2d_f64() { + let x_data: Vec = (1..=16).map(|x| x as f64).collect(); + let x = FlexTensor::from_data(TensorData::new(x_data, vec![1, 1, 4, 4])); + let w_data = vec![1.0f64; 4]; + let weight = FlexTensor::from_data(TensorData::new(w_data, vec![1, 1, 2, 2])); + let options = ConvOptions::new([1, 1], [0, 0], [1, 1], 1); + let result = conv2d_f64(x, weight, None, &options); + let out: Vec = result.into_data().to_vec().unwrap(); + assert_eq!( + out, + vec![14.0, 18.0, 22.0, 30.0, 34.0, 38.0, 46.0, 50.0, 54.0] + ); + } + + #[test] + fn test_conv2d_f16() { + let x_data: Vec = (1..=16).map(|x| f16::from_f32(x as f32)).collect(); + let x = FlexTensor::from_data(TensorData::new(x_data, vec![1, 1, 4, 4])); + let w_data: Vec = vec![f16::from_f32(1.0); 4]; + let weight = FlexTensor::from_data(TensorData::new(w_data, vec![1, 1, 2, 2])); + let options = ConvOptions::new([1, 1], [0, 0], [1, 1], 1); + let result = conv2d_f16(x, weight, None, &options); + let out: Vec = result.into_data().to_vec().unwrap(); + let expected = vec![14.0, 18.0, 22.0, 30.0, 34.0, 38.0, 46.0, 50.0, 54.0]; + for (a, e) in out.iter().zip(expected.iter()) { + assert!((a.to_f32() - e).abs() < 0.5); + } + } + + #[test] + fn test_conv2d_bf16() { + let x_data: Vec = (1..=16).map(|x| bf16::from_f32(x as f32)).collect(); + let x = FlexTensor::from_data(TensorData::new(x_data, vec![1, 1, 4, 4])); + let w_data: Vec = vec![bf16::from_f32(1.0); 4]; + let weight = FlexTensor::from_data(TensorData::new(w_data, vec![1, 1, 2, 2])); + let options = ConvOptions::new([1, 1], [0, 0], [1, 1], 1); + let result = conv2d_bf16(x, weight, None, &options); + let out: Vec = result.into_data().to_vec().unwrap(); + let expected = vec![14.0, 18.0, 22.0, 30.0, 34.0, 38.0, 46.0, 50.0, 54.0]; + for (a, e) in out.iter().zip(expected.iter()) { + assert!((a.to_f32() - e).abs() < 0.5); + } + } + + // ======================================================================== + // Depthwise conv fast-path tests + // ======================================================================== + // + // These tests exercise `conv3d_depthwise_impl` and cross-check against a + // naive reference to catch any indexing, bounds, or accumulation mistakes. + + /// Naive NCHW depthwise conv2d reference implementation. + /// Returns output with shape `[batch, channels, out_h, out_w]`. + #[allow(clippy::too_many_arguments)] + fn naive_depthwise_conv2d_f32( + x: &[f32], + w: &[f32], + bias: Option<&[f32]>, + batch: usize, + channels: usize, + in_h: usize, + in_w: usize, + kernel_h: usize, + kernel_w: usize, + stride_h: usize, + stride_w: usize, + pad_h: usize, + pad_w: usize, + dilation_h: usize, + dilation_w: usize, + ) -> (Vec, usize, usize) { + let out_h = (in_h + 2 * pad_h - dilation_h * (kernel_h - 1) - 1) / stride_h + 1; + let out_w = (in_w + 2 * pad_w - dilation_w * (kernel_w - 1) - 1) / stride_w + 1; + let mut out = vec![0.0f32; batch * channels * out_h * out_w]; + for b in 0..batch { + for c in 0..channels { + for oh in 0..out_h { + for ow in 0..out_w { + let mut acc = 0.0f32; + for kh in 0..kernel_h { + let ih = oh as isize * stride_h as isize + + kh as isize * dilation_h as isize + - pad_h as isize; + if ih < 0 || ih >= in_h as isize { + continue; + } + for kw in 0..kernel_w { + let iw = ow as isize * stride_w as isize + + kw as isize * dilation_w as isize + - pad_w as isize; + if iw < 0 || iw >= in_w as isize { + continue; + } + let x_idx = + ((b * channels + c) * in_h + ih as usize) * in_w + iw as usize; + let w_idx = (c * kernel_h + kh) * kernel_w + kw; + acc += x[x_idx] * w[w_idx]; + } + } + if let Some(bias) = bias { + acc += bias[c]; + } + let o_idx = ((b * channels + c) * out_h + oh) * out_w + ow; + out[o_idx] = acc; + } + } + } + } + (out, out_h, out_w) + } + + /// Build deterministic pseudo-random input for a depthwise test. + fn seeded_vec_f32(n: usize, seed: u32) -> Vec { + (0..n) + .map(|i| { + let v = ((i as u32).wrapping_mul(2654435761).wrapping_add(seed)) & 0xffff; + (v as f32 / 32768.0) - 1.0 + }) + .collect() + } + + /// Shared helper: run conv2d via the public dispatch (which must pick the + /// depthwise path for these shapes) and compare to the naive reference. + #[allow(clippy::too_many_arguments)] + fn check_depthwise_conv2d_f32( + batch: usize, + channels: usize, + in_h: usize, + in_w: usize, + kernel_h: usize, + kernel_w: usize, + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + with_bias: bool, + ) { + let x_vec = seeded_vec_f32(batch * channels * in_h * in_w, 1); + let w_vec = seeded_vec_f32(channels * kernel_h * kernel_w, 2); + let bias_vec = if with_bias { + Some(seeded_vec_f32(channels, 3)) + } else { + None + }; + + let (expected, out_h, out_w) = naive_depthwise_conv2d_f32( + &x_vec, + &w_vec, + bias_vec.as_deref(), + batch, + channels, + in_h, + in_w, + kernel_h, + kernel_w, + stride[0], + stride[1], + padding[0], + padding[1], + dilation[0], + dilation[1], + ); + + let x = FlexTensor::from_data(TensorData::new(x_vec, vec![batch, channels, in_h, in_w])); + let weight = FlexTensor::from_data(TensorData::new( + w_vec, + vec![channels, 1, kernel_h, kernel_w], + )); + let bias = bias_vec.map(|v| FlexTensor::from_data(TensorData::new(v, vec![channels]))); + let options = ConvOptions::new(stride, padding, dilation, channels); + let result = conv2d_f32(x, weight, bias, &options); + + assert_eq!( + result.layout().shape().to_vec(), + vec![batch, channels, out_h, out_w], + "output shape mismatch" + ); + + let out: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(out.len(), expected.len()); + for (i, (a, e)) in out.iter().zip(expected.iter()).enumerate() { + assert!( + (a - e).abs() < 1e-4, + "mismatch at {i}: got {a}, expected {e}" + ); + } + } + + #[test] + fn test_conv2d_depthwise_3x3_no_pad() { + // Canonical depthwise: groups == channels_in == channels_out = 8, 3x3. + check_depthwise_conv2d_f32(2, 8, 16, 16, 3, 3, [1, 1], [0, 0], [1, 1], false); + } + + #[test] + fn test_conv2d_depthwise_3x3_pad1() { + // Same input and output size via padding=1. + check_depthwise_conv2d_f32(2, 8, 16, 16, 3, 3, [1, 1], [1, 1], [1, 1], false); + } + + #[test] + fn test_conv2d_depthwise_3x3_stride2_pad1() { + // Halve spatial via stride 2, with padding. + check_depthwise_conv2d_f32(1, 16, 32, 32, 3, 3, [2, 2], [1, 1], [1, 1], false); + } + + #[test] + fn test_conv2d_depthwise_5x5_pad2() { + check_depthwise_conv2d_f32(2, 4, 10, 10, 5, 5, [1, 1], [2, 2], [1, 1], false); + } + + #[test] + fn test_conv2d_depthwise_7x7_pad3() { + // Large kernel like ConvNeXt's 7x7 depthwise. + check_depthwise_conv2d_f32(2, 24, 14, 14, 7, 7, [1, 1], [3, 3], [1, 1], false); + } + + #[test] + fn test_conv2d_depthwise_dilated() { + // Dilation 2 means the effective receptive field is 5x5 but with gaps. + check_depthwise_conv2d_f32(1, 8, 12, 12, 3, 3, [1, 1], [2, 2], [2, 2], false); + } + + #[test] + fn test_conv2d_depthwise_with_bias() { + check_depthwise_conv2d_f32(2, 8, 8, 8, 3, 3, [1, 1], [1, 1], [1, 1], true); + } + + #[test] + fn test_conv2d_depthwise_single_channel() { + // groups = channels_in = 1 is a degenerate depthwise which should + // still go through this path (it also looks identical to an + // ungrouped 1-channel conv, but validates the range math). + check_depthwise_conv2d_f32(1, 1, 5, 5, 3, 3, [1, 1], [1, 1], [1, 1], false); + } + + #[test] + fn test_conv2d_depthwise_asymmetric_kernel() { + // Non-square kernel. + check_depthwise_conv2d_f32(1, 4, 8, 12, 3, 5, [1, 1], [1, 2], [1, 1], false); + } + + #[test] + fn test_conv2d_depthwise_f64() { + // Smoke test f64 dispatch through depthwise path. + let x_data: Vec = (0..2 * 4 * 5 * 5).map(|i| (i as f64) * 0.1).collect(); + let w_data: Vec = (0..4 * 3 * 3).map(|i| (i as f64) * 0.01).collect(); + let x = FlexTensor::from_data(TensorData::new(x_data.clone(), vec![2, 4, 5, 5])); + let weight = FlexTensor::from_data(TensorData::new(w_data.clone(), vec![4, 1, 3, 3])); + let options = ConvOptions::new([1, 1], [1, 1], [1, 1], 4); + let result = conv2d_f64(x, weight, None, &options); + let out: Vec = result.into_data().to_vec().unwrap(); + + // Verify against a naive f64 reference for one element (center of channel 2). + let b = 0usize; + let c = 2usize; + let oh = 2usize; + let ow = 2usize; + let mut expected = 0.0f64; + for kh in 0..3 { + for kw in 0..3 { + let ih = oh as isize + kh as isize - 1; + let iw = ow as isize + kw as isize - 1; + if (0..5).contains(&ih) && (0..5).contains(&iw) { + let x_idx = ((b * 4 + c) * 5 + ih as usize) * 5 + iw as usize; + let w_idx = (c * 3 + kh) * 3 + kw; + expected += x_data[x_idx] * w_data[w_idx]; + } + } + } + let out_idx = ((b * 4 + c) * 5 + oh) * 5 + ow; + assert!((out[out_idx] - expected).abs() < 1e-10); + } + + #[test] + fn test_conv2d_depthwise_f16() { + // Validate f16 depthwise dispatch against a naive f32 reference. + // Shape check alone would miss indexing/accumulation bugs that still + // happen to produce the right shape. + use burn_std::f16; + let x_data_f32: Vec = (0..16).map(|i| i as f32 * 0.1).collect(); + let w_data_f32: Vec = (0..16).map(|i| i as f32 * 0.01).collect(); + let x_data: Vec = x_data_f32.iter().copied().map(f16::from_f32).collect(); + let w_data: Vec = w_data_f32.iter().copied().map(f16::from_f32).collect(); + let x = FlexTensor::from_data(TensorData::new(x_data, vec![1, 4, 2, 2])); + let weight = FlexTensor::from_data(TensorData::new(w_data, vec![4, 1, 2, 2])); + let options = ConvOptions::new([1, 1], [0, 0], [1, 1], 4); + let result = conv2d_f16(x, weight, None, &options); + assert_eq!(result.layout().shape().to_vec(), vec![1, 4, 1, 1]); + let out: Vec = result.into_data().to_vec().unwrap(); + + // Depthwise: out[c] = sum over (kh, kw) of x[c, kh, kw] * w[c, 0, kh, kw]. + // The input per-channel is 4 elements (2x2) and the kernel is 2x2, so + // there's exactly one output pixel per channel. + for c in 0..4 { + let mut expected = 0.0f32; + for k in 0..4 { + // c * 4 indexes into channel c's 2x2 plane; both x and w share + // the same [c, ...] offset because the depthwise weight layout + // is [c_out, 1, kh, kw]. + expected += x_data_f32[c * 4 + k] * w_data_f32[c * 4 + k]; + } + let actual = out[c].to_f32(); + assert!( + (actual - expected).abs() < 1e-2, + "f16 depthwise mismatch at c={c}: expected {expected}, got {actual}" + ); + } + } + + #[test] + fn test_conv1d_depthwise() { + // conv1d -> conv3d expansion (kd=kh=1), depthwise on the last axis. + let channels = 4; + let in_w = 16; + let kw = 3; + let x_data = seeded_vec_f32(channels * in_w, 10); + let w_data = seeded_vec_f32(channels * kw, 20); + let x = FlexTensor::from_data(TensorData::new(x_data.clone(), vec![1, channels, in_w])); + let weight = FlexTensor::from_data(TensorData::new(w_data.clone(), vec![channels, 1, kw])); + let options = ConvOptions::new([1], [1], [1], channels); + let result = conv1d_f32(x, weight, None, &options); + let out_w = in_w; + assert_eq!(result.layout().shape().to_vec(), vec![1, channels, out_w]); + let out: Vec = result.into_data().to_vec().unwrap(); + + // Naive reference. + for c in 0..channels { + for o in 0..out_w { + let mut expected = 0.0f32; + for k in 0..kw { + let i = o as isize + k as isize - 1; + if i >= 0 && i < in_w as isize { + expected += x_data[c * in_w + i as usize] * w_data[c * kw + k]; + } + } + let actual = out[c * out_w + o]; + assert!( + (actual - expected).abs() < 1e-5, + "conv1d depthwise mismatch at c={c}, o={o}: expected {expected}, got {actual}" + ); + } + } + } + + #[test] + fn test_conv1d_depthwise_stride_batch_bias() { + // Depthwise conv1d with batch > 1, stride > 1, padding, and bias. + // conv1d flows through the same conv3d_depthwise_impl as conv2d + // because conv1d expands to a 3D shape with kd=kh=1. This test + // validates the full dispatch path with a non-trivial stride. + let batch = 3; + let channels = 6; + let in_w = 32; + let kw = 5; + let stride = 2; + let pad = 2; + let out_w = (in_w + 2 * pad - kw) / stride + 1; + + let x_data = seeded_vec_f32(batch * channels * in_w, 30); + let w_data = seeded_vec_f32(channels * kw, 40); + let bias_data = seeded_vec_f32(channels, 50); + + let x = FlexTensor::from_data(TensorData::new(x_data.clone(), vec![batch, channels, in_w])); + let weight = FlexTensor::from_data(TensorData::new(w_data.clone(), vec![channels, 1, kw])); + let bias = FlexTensor::from_data(TensorData::new(bias_data.clone(), vec![channels])); + let options = ConvOptions::new([stride], [pad], [1], channels); + let result = conv1d_f32(x, weight, Some(bias), &options); + assert_eq!( + result.layout().shape().to_vec(), + vec![batch, channels, out_w] + ); + let out: Vec = result.into_data().to_vec().unwrap(); + + // Naive reference. + for b in 0..batch { + for c in 0..channels { + for o in 0..out_w { + let mut expected = bias_data[c]; + for k in 0..kw { + let i = (o as isize * stride as isize) + k as isize - pad as isize; + if i >= 0 && i < in_w as isize { + let x_idx = (b * channels + c) * in_w + i as usize; + let w_idx = c * kw + k; + expected += x_data[x_idx] * w_data[w_idx]; + } + } + let actual = out[(b * channels + c) * out_w + o]; + assert!( + (actual - expected).abs() < 1e-4, + "conv1d depthwise mismatch at b={b}, c={c}, o={o}: expected {expected}, got {actual}" + ); + } + } + } + } + + // ======================================================================== + // Small-channel conv fast-path tests + // ======================================================================== + // + // These tests exercise `conv3d_small_channel_impl` for groups=1 convs + // with `channels_in <= SMALL_CHANNEL_IN_THRESHOLD`. They cross-check the + // output against a naive NCHW reference that sums over both input + // channels and kernel positions. + + #[allow(clippy::too_many_arguments)] + fn naive_conv2d_f32( + x: &[f32], + w: &[f32], + bias: Option<&[f32]>, + batch: usize, + channels_in: usize, + channels_out: usize, + in_h: usize, + in_w: usize, + kernel_h: usize, + kernel_w: usize, + stride_h: usize, + stride_w: usize, + pad_h: usize, + pad_w: usize, + dilation_h: usize, + dilation_w: usize, + ) -> (Vec, usize, usize) { + let out_h = (in_h + 2 * pad_h - dilation_h * (kernel_h - 1) - 1) / stride_h + 1; + let out_w = (in_w + 2 * pad_w - dilation_w * (kernel_w - 1) - 1) / stride_w + 1; + let mut out = vec![0.0f32; batch * channels_out * out_h * out_w]; + for b in 0..batch { + for co in 0..channels_out { + for oh in 0..out_h { + for ow in 0..out_w { + let mut acc = 0.0f32; + for ci in 0..channels_in { + for kh in 0..kernel_h { + let ih = oh as isize * stride_h as isize + + kh as isize * dilation_h as isize + - pad_h as isize; + if ih < 0 || ih >= in_h as isize { + continue; + } + for kw in 0..kernel_w { + let iw = ow as isize * stride_w as isize + + kw as isize * dilation_w as isize + - pad_w as isize; + if iw < 0 || iw >= in_w as isize { + continue; + } + let x_idx = ((b * channels_in + ci) * in_h + ih as usize) + * in_w + + iw as usize; + let w_idx = + ((co * channels_in + ci) * kernel_h + kh) * kernel_w + kw; + acc += x[x_idx] * w[w_idx]; + } + } + } + if let Some(bias) = bias { + acc += bias[co]; + } + let o_idx = ((b * channels_out + co) * out_h + oh) * out_w + ow; + out[o_idx] = acc; + } + } + } + } + (out, out_h, out_w) + } + + #[allow(clippy::too_many_arguments)] + fn check_small_channel_conv2d_f32( + batch: usize, + channels_in: usize, + channels_out: usize, + in_h: usize, + in_w: usize, + kernel_h: usize, + kernel_w: usize, + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + with_bias: bool, + ) { + let x_vec = seeded_vec_f32(batch * channels_in * in_h * in_w, 100); + let w_vec = seeded_vec_f32(channels_out * channels_in * kernel_h * kernel_w, 200); + let bias_vec = if with_bias { + Some(seeded_vec_f32(channels_out, 300)) + } else { + None + }; + + let (expected, out_h, out_w) = naive_conv2d_f32( + &x_vec, + &w_vec, + bias_vec.as_deref(), + batch, + channels_in, + channels_out, + in_h, + in_w, + kernel_h, + kernel_w, + stride[0], + stride[1], + padding[0], + padding[1], + dilation[0], + dilation[1], + ); + + let x = FlexTensor::from_data(TensorData::new(x_vec, vec![batch, channels_in, in_h, in_w])); + let weight = FlexTensor::from_data(TensorData::new( + w_vec, + vec![channels_out, channels_in, kernel_h, kernel_w], + )); + let bias = bias_vec.map(|v| FlexTensor::from_data(TensorData::new(v, vec![channels_out]))); + let options = ConvOptions::new(stride, padding, dilation, 1); + let result = conv2d_f32(x, weight, bias, &options); + + assert_eq!( + result.layout().shape().to_vec(), + vec![batch, channels_out, out_h, out_w], + "output shape mismatch" + ); + + let out: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(out.len(), expected.len()); + for (i, (a, e)) in out.iter().zip(expected.iter()).enumerate() { + assert!( + (a - e).abs() < 1e-3, + "mismatch at {i}: got {a}, expected {e}" + ); + } + } + + #[test] + fn test_conv2d_small_channel_3in_8out_k3x3_pad1() { + // Sobel-like: 3 input channels, several output channels, 3x3 kernel. + check_small_channel_conv2d_f32(2, 3, 8, 16, 16, 3, 3, [1, 1], [1, 1], [1, 1], false); + } + + #[test] + fn test_conv2d_small_channel_3in_3out_k3x3_no_pad() { + // Channels_in == channels_out == 3, same count but groups=1 so each + // output channel combines all input channels (not depthwise). + check_small_channel_conv2d_f32(1, 3, 3, 10, 10, 3, 3, [1, 1], [0, 0], [1, 1], false); + } + + #[test] + fn test_conv2d_small_channel_3in_16out_k5x5_pad2() { + check_small_channel_conv2d_f32(2, 3, 16, 12, 12, 5, 5, [1, 1], [2, 2], [1, 1], false); + } + + #[test] + fn test_conv2d_small_channel_4in_8out_k3x3_stride2() { + // Threshold channels_in == 4. + check_small_channel_conv2d_f32(1, 4, 8, 16, 16, 3, 3, [2, 2], [1, 1], [1, 1], false); + } + + #[test] + fn test_conv2d_small_channel_2in_4out_dilated() { + check_small_channel_conv2d_f32(1, 2, 4, 12, 12, 3, 3, [1, 1], [2, 2], [2, 2], false); + } + + #[test] + fn test_conv2d_small_channel_with_bias() { + check_small_channel_conv2d_f32(2, 3, 8, 8, 8, 3, 3, [1, 1], [1, 1], [1, 1], true); + } + + #[test] + fn test_conv2d_small_channel_asymmetric_kernel() { + // 1x3 and 3x1 kernels are common for separable edge filters. + check_small_channel_conv2d_f32(1, 3, 6, 16, 16, 1, 3, [1, 1], [0, 1], [1, 1], false); + check_small_channel_conv2d_f32(1, 3, 6, 16, 16, 3, 1, [1, 1], [1, 0], [1, 1], false); + } + + #[test] + fn test_conv2d_small_channel_f64() { + // Smoke test f64 dispatch. + let x_data: Vec = (0..2 * 3 * 5 * 5).map(|i| (i as f64) * 0.1).collect(); + let w_data: Vec = (0..4 * 3 * 3 * 3).map(|i| (i as f64) * 0.01).collect(); + let x = FlexTensor::from_data(TensorData::new(x_data.clone(), vec![2, 3, 5, 5])); + let weight = FlexTensor::from_data(TensorData::new(w_data.clone(), vec![4, 3, 3, 3])); + let options = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + let result = conv2d_f64(x, weight, None, &options); + let out: Vec = result.into_data().to_vec().unwrap(); + + // Verify against a naive reference for one element (center of channel 2). + let b = 0usize; + let co = 2usize; + let oh = 2usize; + let ow = 2usize; + let mut expected = 0.0f64; + for ci in 0..3 { + for kh in 0..3 { + for kw in 0..3 { + let ih = oh as isize + kh as isize - 1; + let iw = ow as isize + kw as isize - 1; + if (0..5).contains(&ih) && (0..5).contains(&iw) { + let x_idx = ((b * 3 + ci) * 5 + ih as usize) * 5 + iw as usize; + let w_idx = ((co * 3 + ci) * 3 + kh) * 3 + kw; + expected += x_data[x_idx] * w_data[w_idx]; + } + } + } + } + let out_idx = ((b * 4 + co) * 5 + oh) * 5 + ow; + assert!( + (out[out_idx] - expected).abs() < 1e-10, + "got {}, expected {expected}", + out[out_idx] + ); + } + + #[test] + fn test_conv2d_small_channel_f16() { + // Validate f16 dispatch through the small-channel path by running + // the same shape through the f32 path and comparing element-wise. + // This catches indexing/accumulation bugs in addition to regressions + // in the `num_traits::Float` monomorphization for f16. + use burn_std::f16; + let x_data_f32: Vec = (0..3 * 4 * 4).map(|i| i as f32 * 0.1).collect(); + let w_data_f32: Vec = (0..4 * 3 * 3 * 3).map(|i| i as f32 * 0.01).collect(); + + let x_data_f16: Vec = x_data_f32.iter().copied().map(f16::from_f32).collect(); + let w_data_f16: Vec = w_data_f32.iter().copied().map(f16::from_f32).collect(); + + let x_f16 = FlexTensor::from_data(TensorData::new(x_data_f16, vec![1, 3, 4, 4])); + let weight_f16 = FlexTensor::from_data(TensorData::new(w_data_f16, vec![4, 3, 3, 3])); + let x_f32 = FlexTensor::from_data(TensorData::new(x_data_f32, vec![1, 3, 4, 4])); + let weight_f32 = FlexTensor::from_data(TensorData::new(w_data_f32, vec![4, 3, 3, 3])); + + let options = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + let result_f16 = conv2d_f16(x_f16, weight_f16, None, &options); + let result_f32 = conv2d_f32(x_f32, weight_f32, None, &options); + assert_eq!(result_f16.layout().shape().to_vec(), vec![1, 4, 4, 4]); + assert_eq!(result_f32.layout().shape().to_vec(), vec![1, 4, 4, 4]); + + let out_f16: Vec = result_f16.into_data().to_vec().unwrap(); + let out_f32: Vec = result_f32.into_data().to_vec().unwrap(); + assert_eq!(out_f16.len(), out_f32.len()); + + // f16 has ~11 bits of mantissa (~0.1% relative precision). With + // accumulation across `c_in * k_spatial = 27` FMAs the relative + // error can grow to a few times that. Use a relative tolerance + // scaled by the expected magnitude with a small absolute floor for + // values near zero. + let rel_tol = 3e-3f32; + let abs_tol = 1e-2f32; + for (i, (actual, expected)) in out_f16.iter().zip(out_f32.iter()).enumerate() { + let actual_f32 = actual.to_f32(); + let bound = (expected.abs() * rel_tol).max(abs_tol); + assert!( + !actual_f32.is_nan() && (actual_f32 - expected).abs() <= bound, + "f16 small-channel mismatch at {i}: got {actual_f32}, expected {expected}, bound {bound}" + ); + } + } + + #[test] + fn test_conv2d_small_channel_bias_length_mismatch_panics() { + // The small-channel impl must panic loudly when bias length does not + // match channels_out. A silent truncation here would make models with + // misconfigured bias silently produce wrong output. + let x = FlexTensor::from_data(TensorData::new(vec![0.0f32; 48], vec![1, 3, 4, 4])); + let weight = FlexTensor::from_data(TensorData::new(vec![0.0f32; 108], vec![4, 3, 3, 3])); + // Bias has only 2 elements but there are 4 output channels. + let bias = FlexTensor::from_data(TensorData::new(vec![0.0f32, 0.0], vec![2])); + let options = ConvOptions::new([1, 1], [1, 1], [1, 1], 1); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + conv2d_f32(x, weight, Some(bias), &options) + })); + assert!(result.is_err(), "expected panic on bias length mismatch"); + } + + #[test] + fn test_conv2d_depthwise_bias_length_mismatch_panics() { + // Mirror of the small-channel bias check for the depthwise path. + // Depthwise dispatch requires groups == channels_in == channels_out, so + // we construct a 3->3 depthwise conv and deliberately underspecify the + // bias length to exercise the assert_eq in conv3d_depthwise_impl. + let x = FlexTensor::from_data(TensorData::new(vec![0.0f32; 48], vec![1, 3, 4, 4])); + let weight = FlexTensor::from_data(TensorData::new(vec![0.0f32; 27], vec![3, 1, 3, 3])); + // Bias has only 2 elements but there are 3 depthwise channels. + let bias = FlexTensor::from_data(TensorData::new(vec![0.0f32, 0.0], vec![2])); + let options = ConvOptions::new([1, 1], [1, 1], [1, 1], 3); // groups=3 => depthwise + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + conv2d_f32(x, weight, Some(bias), &options) + })); + assert!(result.is_err(), "expected panic on bias length mismatch"); + } + + #[test] + fn test_conv_plane_accumulate_accumulates_into_prefilled() { + // Pin the documented precondition that `conv_plane_accumulate` + // accumulates into `out_plane` rather than overwriting it. Both + // in-tree callers pre-zero before the first call, then call it + // repeatedly across input channels with the result acting as the + // running accumulator. A regression that overwrites instead of + // accumulating would silently drop every input-channel contribution + // except the last. + // + // Test shape: single 2x2 input -> single 2x2 output, 1x1 kernel with + // weight 1.0. Pre-fill out_plane with [10, 20, 30, 40] so the + // expected post-call value is x + prefill. + let x = vec![1.0f32, 2.0, 3.0, 4.0]; + let w = vec![1.0f32]; // 1x1 kernel, weight = 1 + let mut out_plane = vec![10.0f32, 20.0, 30.0, 40.0]; + + // 1x1 kernel, no padding: the only kernel position contributes to + // every output row and column (indices [0, out_h) and [0, out_w)). + let oh_ranges = [(0usize, 2usize)]; // per kh: (out_start, out_end) + let ow_ranges = [(0usize, 2usize)]; // per kw: (out_start, out_end) + + super::conv_plane_accumulate::( + &mut out_plane, + &x, + &w, + /* kernel_h */ 1, + /* kernel_w */ 1, + /* in_w */ 2, + /* out_w */ 2, + /* stride_h */ 1, + /* stride_w */ 1, + /* pad_h */ 0, + /* pad_w */ 0, + /* dilation_h */ 1, + /* dilation_w */ 1, + &oh_ranges, + &ow_ranges, + ); + + // Expected: prefill + (x * 1.0) element-wise. + assert_eq!(out_plane, vec![11.0f32, 22.0, 33.0, 44.0]); + } + + #[test] + fn test_conv1d_small_channel_3in() { + // conv1d -> conv3d expansion with groups=1 and 3 input channels. + let batch = 2; + let channels_in = 3; + let channels_out = 5; + let in_w = 24; + let kw = 5; + let stride = 1; + let pad = 2; + let out_w = in_w; // same padding + + let x_data = seeded_vec_f32(batch * channels_in * in_w, 400); + let w_data = seeded_vec_f32(channels_out * channels_in * kw, 500); + let x = FlexTensor::from_data(TensorData::new( + x_data.clone(), + vec![batch, channels_in, in_w], + )); + let weight = FlexTensor::from_data(TensorData::new( + w_data.clone(), + vec![channels_out, channels_in, kw], + )); + let options = ConvOptions::new([stride], [pad], [1], 1); + let result = conv1d_f32(x, weight, None, &options); + assert_eq!( + result.layout().shape().to_vec(), + vec![batch, channels_out, out_w] + ); + let out: Vec = result.into_data().to_vec().unwrap(); + + for b in 0..batch { + for co in 0..channels_out { + for o in 0..out_w { + let mut expected = 0.0f32; + for ci in 0..channels_in { + for k in 0..kw { + let i = o as isize + k as isize - pad as isize; + if i >= 0 && i < in_w as isize { + let x_idx = (b * channels_in + ci) * in_w + i as usize; + let w_idx = (co * channels_in + ci) * kw + k; + expected += x_data[x_idx] * w_data[w_idx]; + } + } + } + let actual = out[(b * channels_out + co) * out_w + o]; + assert!( + (actual - expected).abs() < 1e-4, + "mismatch at b={b}, co={co}, o={o}: expected {expected}, got {actual}" + ); + } + } + } + } + + #[test] + fn test_conv2d_small_channel_single_input() { + // 1 input channel is allowed (<= threshold) but this also matches + // many existing conv tests that passed through conv3d_impl before + // this path existed. Cross-check against the naive reference. + check_small_channel_conv2d_f32(1, 1, 4, 8, 8, 3, 3, [1, 1], [1, 1], [1, 1], false); + } + + #[test] + fn test_conv2d_small_channel_threshold_exact() { + // Verify the thresholds are honored: + // - channels_in in [1, 4] AND channels_out in [1, 16] triggers the + // small-channel path + // - channels_in == 5 does not + // - channels_out == 17 does not + // - groups != 1 does not + // + // conv2d options are expanded to 3D as `[0, pad_h, pad_w]` and + // `[1, stride_h, stride_w]` etc. - the d-axis is always trivial. + + // c_in == 4, c_out == 16: exactly at the thresholds, should trigger. + assert!(should_use_small_channel_conv( + &[1, 4, 1, 8, 8], + &[16, 4, 1, 3, 3], + &ConvOptions::new([1, 1, 1], [0, 1, 1], [1, 1, 1], 1), + )); + // c_in == 5: excluded. + assert!(!should_use_small_channel_conv( + &[1, 5, 1, 8, 8], + &[2, 5, 1, 3, 3], + &ConvOptions::new([1, 1, 1], [0, 1, 1], [1, 1, 1], 1), + )); + // c_in == 3, c_out == 17: excluded (large output channel count). + assert!(!should_use_small_channel_conv( + &[1, 3, 1, 8, 8], + &[17, 3, 1, 3, 3], + &ConvOptions::new([1, 1, 1], [0, 1, 1], [1, 1, 1], 1), + )); + // c_in == 3, c_out == 64 (ImageNet first layer): excluded. + assert!(!should_use_small_channel_conv( + &[1, 3, 1, 224, 224], + &[64, 3, 1, 7, 7], + &ConvOptions::new([1, 2, 2], [0, 3, 3], [1, 1, 1], 1), + )); + // Non-groups=1 is excluded. + assert!(!should_use_small_channel_conv( + &[1, 4, 1, 8, 8], + &[4, 1, 1, 3, 3], + &ConvOptions::new([1, 1, 1], [0, 1, 1], [1, 1, 1], 4), + )); + + // Cross-check c_in == 5 against naive reference (this goes through + // the generic conv3d_impl path, not the small-channel path). + check_small_channel_conv2d_f32(1, 5, 4, 8, 8, 3, 3, [1, 1], [1, 1], [1, 1], false); + } + + // The tests below exercise `conv_plane_accumulate_oh_outer`, the variant + // selected when the output plane exceeds `CONV_PLANE_OH_OUTER_THRESHOLD` + // (8192 elements). Every test above this point uses shapes where the plane + // is <= 256 elements and stays on the kh-outer variant; without these, + // the oh-outer code path ships uncovered. + // + // 96x96 = 9216 > 8192, so one element past the threshold. 97x97 = 9409 + // leaves a little more slack in case someone nudges the threshold. + + #[test] + fn test_conv2d_depthwise_oh_outer_k3x3() { + // Depthwise path, oh-outer dispatch, 3x3 kernel, stride 1. + // Plane = 96 * 96 = 9216 elements. + check_depthwise_conv2d_f32(1, 2, 96, 96, 3, 3, [1, 1], [1, 1], [1, 1], false); + } + + #[test] + fn test_conv2d_depthwise_oh_outer_k3x3_stride2() { + // Depthwise, oh-outer, stride 2: exercises the `stride_w != 1` + // inner branch inside the oh-outer variant. Plane = 96*96 = 9216. + check_depthwise_conv2d_f32(1, 2, 192, 192, 3, 3, [2, 2], [1, 1], [1, 1], false); + } + + #[test] + fn test_conv2d_depthwise_oh_outer_k5x1() { + // Depthwise, oh-outer, 5x1 asymmetric kernel: the exact shape + // pattern that motivated the loop reorder (Sobel-style separable + // filter on a large plane). Plane = 97*97 = 9409. + check_depthwise_conv2d_f32(1, 2, 97, 97, 5, 1, [1, 1], [2, 0], [1, 1], false); + } + + #[test] + fn test_conv2d_depthwise_oh_outer_k1x5() { + // Depthwise, oh-outer, 1x5 asymmetric kernel. Plane = 97*97 = 9409. + check_depthwise_conv2d_f32(1, 2, 97, 97, 1, 5, [1, 1], [0, 2], [1, 1], false); + } + + #[test] + fn test_conv2d_small_channel_oh_outer_k3x3() { + // Small-channel path, oh-outer dispatch, stride 1. + // Plane = 96 * 96 = 9216. + check_small_channel_conv2d_f32(1, 3, 4, 96, 96, 3, 3, [1, 1], [1, 1], [1, 1], false); + } + + #[test] + fn test_conv2d_small_channel_oh_outer_k3x3_stride2() { + // Small-channel, oh-outer, stride 2: exercises the stride != 1 + // inner branch through small-channel dispatch. Plane = 96*96. + check_small_channel_conv2d_f32(1, 3, 4, 192, 192, 3, 3, [2, 2], [1, 1], [1, 1], false); + } + + #[test] + fn test_conv2d_small_channel_oh_outer_k5x1_sobel() { + // Small-channel, oh-outer, 5x1 asymmetric kernel: the shape from + // the user-reported Sobel regression on RGB, on a plane large + // enough to cross the threshold. Plane = 97*97 = 9409. + check_small_channel_conv2d_f32(1, 3, 3, 97, 97, 5, 1, [1, 1], [2, 0], [1, 1], false); + } + + #[test] + fn test_conv2d_small_channel_oh_outer_k1x5_sobel() { + // Small-channel, oh-outer, 1x5 asymmetric kernel. Plane = 97*97. + check_small_channel_conv2d_f32(1, 3, 3, 97, 97, 1, 5, [1, 1], [0, 2], [1, 1], false); + } + + #[test] + fn test_conv2d_small_channel_oh_outer_with_bias_and_dilation() { + // Small-channel, oh-outer, with bias and dilation > 1. + // Plane = 96*96 = 9216. + check_small_channel_conv2d_f32(1, 3, 8, 100, 100, 3, 3, [1, 1], [2, 2], [2, 2], true); + } + + #[test] + fn test_conv2d_depthwise_predicate_triggers() { + // Directly assert `should_use_depthwise_conv` returns true for + // representative canonical depthwise shapes, and false for shapes + // that look depthwise-ish but are not. Without this, a future change + // that tightens the predicate could silently revert depthwise shapes + // to the generic path while the correctness tests still pass. + // + // conv2d options expand to 3D as `[1, stride_h, stride_w]` etc., with + // the d-axis always trivial. conv1d expands with `kd == kh == 1` and + // `in_d == in_h == 1`. + + // Canonical depthwise 3x3, groups == c_in == c_out == 32. + assert!(should_use_depthwise_conv( + &[4, 32, 1, 56, 56], + &[32, 1, 1, 3, 3], + &ConvOptions::new([1, 1, 1], [0, 1, 1], [1, 1, 1], 32), + )); + // Canonical depthwise 7x7, the ConvNeXt shape Thomas reported. + assert!(should_use_depthwise_conv( + &[4, 48, 1, 56, 56], + &[48, 1, 1, 7, 7], + &ConvOptions::new([1, 1, 1], [0, 3, 3], [1, 1, 1], 48), + )); + // Conv1d depthwise (kh == 1 via the conv1d -> conv3d expansion). + assert!(should_use_depthwise_conv( + &[8, 64, 1, 1, 1024], + &[64, 1, 1, 1, 3], + &ConvOptions::new([1, 1, 1], [0, 0, 1], [1, 1, 1], 64), + )); + // Strided + dilated depthwise. + assert!(should_use_depthwise_conv( + &[1, 16, 1, 32, 32], + &[16, 1, 1, 3, 3], + &ConvOptions::new([1, 2, 2], [0, 1, 1], [1, 2, 2], 16), + )); + + // Not depthwise: groups == 1. + assert!(!should_use_depthwise_conv( + &[1, 8, 1, 16, 16], + &[16, 8, 1, 3, 3], + &ConvOptions::new([1, 1, 1], [0, 1, 1], [1, 1, 1], 1), + )); + // Not depthwise: channels_per_group > 1 (grouped but not depthwise). + assert!(!should_use_depthwise_conv( + &[1, 8, 1, 16, 16], + &[8, 4, 1, 3, 3], + &ConvOptions::new([1, 1, 1], [0, 1, 1], [1, 1, 1], 2), + )); + // Not depthwise: groups == c_in but c_out != c_in (depth multiplier + // > 1; the canonical depthwise path is restricted to multiplier 1). + assert!(!should_use_depthwise_conv( + &[1, 8, 1, 16, 16], + &[16, 1, 1, 3, 3], + &ConvOptions::new([1, 1, 1], [0, 1, 1], [1, 1, 1], 8), + )); + // Not depthwise: pure 3D with kd > 1 (the `d`-axis restriction). + assert!(!should_use_depthwise_conv( + &[1, 8, 4, 16, 16], + &[8, 1, 3, 3, 3], + &ConvOptions::new([1, 1, 1], [1, 1, 1], [1, 1, 1], 8), + )); + } + + #[test] + fn test_valid_out_range_basics() { + // Sanity checks for the analytic range helper. + // 3x3, no pad, stride 1: every k position produces the full range. + let (s, e) = valid_out_range(0, 1, 0, 1, 5, 3); + assert_eq!((s, e), (0, 3)); + let (s, e) = valid_out_range(2, 1, 0, 1, 5, 3); + assert_eq!((s, e), (0, 3)); + // 3x3, pad 1, stride 1: kernel position 0 needs o*1 + 0 >= 1 -> o >= 1. + let (s, e) = valid_out_range(0, 1, 1, 1, 5, 5); + assert_eq!((s, e), (1, 5)); + // kernel position 2 needs o*1 + 2 - 1 < 5 -> o < 4. + let (s, e) = valid_out_range(2, 1, 1, 1, 5, 5); + assert_eq!((s, e), (0, 4)); + // stride 2: o*2 + 0 in [0, in). With in=5, out=3: o in [0, 3). + let (s, e) = valid_out_range(0, 1, 0, 2, 5, 3); + assert_eq!((s, e), (0, 3)); + // dilation 2, pad 2, stride 1: kernel position 0 iw = o*1 + 0 - 2 -> o >= 2. + let (s, e) = valid_out_range(0, 2, 2, 1, 5, 5); + assert_eq!((s, e), (2, 5)); + } +} diff --git a/crates/burn-flex/src/ops/conv_common.rs b/crates/burn-flex/src/ops/conv_common.rs new file mode 100644 index 0000000..d8f7dc1 --- /dev/null +++ b/crates/burn-flex/src/ops/conv_common.rs @@ -0,0 +1,110 @@ +//! Shared macros and helpers for forward and transposed convolution. + +use alloc::vec::Vec; +use burn_backend::DType; +use burn_std::{Bytes, Shape, bf16}; + +use crate::{FlexTensor, Layout}; + +// ============================================================================ +// Macros for dtype wrappers +// ============================================================================ + +/// Generates a convNd function that delegates to a conv3d function via expand/squeeze. +macro_rules! conv_nd_via_3d { + ($fn_name:ident, $conv3d_fn:ident, $expand_fn:ident, $squeeze_fn:ident, $dim:literal, $Options:ident) => { + pub fn $fn_name( + x: FlexTensor, + weight: FlexTensor, + bias: Option, + options: &$Options<$dim>, + ) -> FlexTensor { + let (x_3d, weight_3d, options_3d) = $expand_fn(&x, &weight, options); + let result_3d = $conv3d_fn(x_3d, weight_3d, bias, &options_3d); + $squeeze_fn(result_3d) + } + }; +} + +/// Generates a bf16 function that converts to f32, calls the f32 variant, converts back. +macro_rules! bf16_via_f32 { + ($bf16_fn:ident, $f32_fn:ident, $dim:literal, $Options:ident) => { + pub fn $bf16_fn( + x: FlexTensor, + weight: FlexTensor, + bias: Option, + options: &$Options<$dim>, + ) -> FlexTensor { + let x_f32 = $crate::ops::conv_common::convert_bf16_to_f32(&x); + let weight_f32 = $crate::ops::conv_common::convert_bf16_to_f32(&weight); + let bias_f32 = bias.map(|b| $crate::ops::conv_common::convert_bf16_to_f32(&b)); + let result_f32 = $f32_fn(x_f32, weight_f32, bias_f32, options); + $crate::ops::conv_common::convert_f32_to_bf16(&result_f32) + } + }; +} + +// ============================================================================ +// Squeeze helpers (used by both conv and conv_transpose 1d/2d paths) +// ============================================================================ + +pub(crate) fn squeeze_3d_to_1d(tensor: FlexTensor) -> FlexTensor { + let shape = tensor.layout().shape(); + tensor.reshape(Shape::from(alloc::vec![shape[0], shape[1], shape[4]])) +} + +pub(crate) fn squeeze_3d_to_2d(tensor: FlexTensor) -> FlexTensor { + let shape = tensor.layout().shape(); + tensor.reshape(Shape::from(alloc::vec![ + shape[0], shape[1], shape[3], shape[4] + ])) +} + +// ============================================================================ +// bf16 conversion helpers +// ============================================================================ + +pub(crate) fn convert_bf16_to_f32(tensor: &FlexTensor) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let data: &[bf16] = tensor.storage(); + let f32_data: Vec = data.iter().map(|x| x.to_f32()).collect(); + FlexTensor::new( + Bytes::from_elems(f32_data), + Layout::contiguous(tensor.layout().shape().clone()), + DType::F32, + ) +} + +pub(crate) fn convert_f32_to_bf16(tensor: &FlexTensor) -> FlexTensor { + let data: &[f32] = tensor.storage(); + let bf16_data: Vec = data.iter().map(|x| bf16::from_f32(*x)).collect(); + FlexTensor::new( + Bytes::from_elems(bf16_data), + Layout::contiguous(tensor.layout().shape().clone()), + DType::BF16, + ) +} + +// ============================================================================ +// Bias addition +// ============================================================================ + +#[allow(clippy::needless_range_loop)] +pub(crate) fn add_bias( + output: &mut [T], + bias: &[T], + batch: usize, + channels: usize, + spatial: usize, + add_fn: fn(T, T) -> T, +) { + for b in 0..batch { + for c in 0..channels { + let offset = b * channels * spatial + c * spatial; + let bias_val = bias[c]; + for i in 0..spatial { + output[offset + i] = add_fn(output[offset + i], bias_val); + } + } + } +} diff --git a/crates/burn-flex/src/ops/conv_transpose.rs b/crates/burn-flex/src/ops/conv_transpose.rs new file mode 100644 index 0000000..48f2f3d --- /dev/null +++ b/crates/burn-flex/src/ops/conv_transpose.rs @@ -0,0 +1,474 @@ +//! Transposed convolution operations using GEMM + col2im approach. +//! +//! All transposed convolutions (1D, 2D, 3D) use a unified 3D implementation: +//! - conv_transpose1d: adds two size-1 dimensions, calls conv_transpose3d, squeezes output +//! - conv_transpose2d: adds one size-1 dimension, calls conv_transpose3d, squeezes output +//! - conv_transpose3d: native implementation +//! +//! For each (batch, group), computes: columns = W_g^T @ X_g, then scatters +//! the columns matrix into the output via col2im. The GEMM handles the heavy +//! channel-reduction multiply-adds; col2im is a lightweight spatial scatter. +//! +//! Supported dtypes: f32, f64, f16 (native gemm), bf16 (via f32 conversion) + +use alloc::vec; +use burn_backend::DType; +use burn_backend::ops::ConvTransposeOptions; +use burn_backend::ops::conv::calculate_conv_transpose_output_size; +use burn_std::{Bytes, Shape, f16}; + +use crate::{FlexTensor, Layout}; + +use super::conv_common::{add_bias, squeeze_3d_to_1d, squeeze_3d_to_2d}; + +/// Generates a conv_transpose3d typed function. +macro_rules! conv_transpose3d_typed { + ($fn_name:ident, $T:ty, $dtype:expr, $zero:expr, $gemm_fn:ident, $add_fn:expr) => { + pub fn $fn_name( + x: FlexTensor, + weight: FlexTensor, + bias: Option, + options: &ConvTransposeOptions<3>, + ) -> FlexTensor { + conv_transpose3d_impl::<$T>(x, weight, bias, options, $dtype, $zero, $gemm_fn, $add_fn) + } + }; +} + +// ============================================================================ +// Conv Transpose 1d - delegates to conv_transpose3d +// ============================================================================ + +conv_nd_via_3d!( + conv_transpose1d_f32, + conv_transpose3d_f32, + expand_transpose_1d_to_3d, + squeeze_3d_to_1d, + 1, + ConvTransposeOptions +); +conv_nd_via_3d!( + conv_transpose1d_f64, + conv_transpose3d_f64, + expand_transpose_1d_to_3d, + squeeze_3d_to_1d, + 1, + ConvTransposeOptions +); +conv_nd_via_3d!( + conv_transpose1d_f16, + conv_transpose3d_f16, + expand_transpose_1d_to_3d, + squeeze_3d_to_1d, + 1, + ConvTransposeOptions +); +bf16_via_f32!( + conv_transpose1d_bf16, + conv_transpose1d_f32, + 1, + ConvTransposeOptions +); + +fn expand_transpose_1d_to_3d( + x: &FlexTensor, + weight: &FlexTensor, + options: &ConvTransposeOptions<1>, +) -> (FlexTensor, FlexTensor, ConvTransposeOptions<3>) { + // x: [N, C_in, L] -> [N, C_in, 1, 1, L] + let x_shape = x.layout().shape(); + let x_3d = x.reshape(Shape::from(vec![x_shape[0], x_shape[1], 1, 1, x_shape[2]])); + + // weight: [C_in, C_out, K] -> [C_in, C_out, 1, 1, K] + let w_shape = weight.layout().shape(); + let weight_3d = weight.reshape(Shape::from(vec![w_shape[0], w_shape[1], 1, 1, w_shape[2]])); + + let options_3d = ConvTransposeOptions::new( + [1, 1, options.stride[0]], + [0, 0, options.padding[0]], + [0, 0, options.padding_out[0]], + [1, 1, options.dilation[0]], + options.groups, + ); + + (x_3d, weight_3d, options_3d) +} + +// ============================================================================ +// Conv Transpose 2d - delegates to conv_transpose3d +// ============================================================================ + +conv_nd_via_3d!( + conv_transpose2d_f32, + conv_transpose3d_f32, + expand_transpose_2d_to_3d, + squeeze_3d_to_2d, + 2, + ConvTransposeOptions +); +conv_nd_via_3d!( + conv_transpose2d_f64, + conv_transpose3d_f64, + expand_transpose_2d_to_3d, + squeeze_3d_to_2d, + 2, + ConvTransposeOptions +); +conv_nd_via_3d!( + conv_transpose2d_f16, + conv_transpose3d_f16, + expand_transpose_2d_to_3d, + squeeze_3d_to_2d, + 2, + ConvTransposeOptions +); +bf16_via_f32!( + conv_transpose2d_bf16, + conv_transpose2d_f32, + 2, + ConvTransposeOptions +); + +fn expand_transpose_2d_to_3d( + x: &FlexTensor, + weight: &FlexTensor, + options: &ConvTransposeOptions<2>, +) -> (FlexTensor, FlexTensor, ConvTransposeOptions<3>) { + // x: [N, C_in, H, W] -> [N, C_in, 1, H, W] + let x_shape = x.layout().shape(); + let x_3d = x.reshape(Shape::from(vec![ + x_shape[0], x_shape[1], 1, x_shape[2], x_shape[3], + ])); + + // weight: [C_in, C_out, Kh, Kw] -> [C_in, C_out, 1, Kh, Kw] + let w_shape = weight.layout().shape(); + let weight_3d = weight.reshape(Shape::from(vec![ + w_shape[0], w_shape[1], 1, w_shape[2], w_shape[3], + ])); + + let options_3d = ConvTransposeOptions::new( + [1, options.stride[0], options.stride[1]], + [0, options.padding[0], options.padding[1]], + [0, options.padding_out[0], options.padding_out[1]], + [1, options.dilation[0], options.dilation[1]], + options.groups, + ); + + (x_3d, weight_3d, options_3d) +} + +// ============================================================================ +// Conv Transpose 3d - core implementation +// ============================================================================ + +conv_transpose3d_typed!( + conv_transpose3d_f32, + f32, + DType::F32, + 0.0f32, + conv_transpose_gemm_f32, + |a, b| a + b +); +conv_transpose3d_typed!( + conv_transpose3d_f64, + f64, + DType::F64, + 0.0f64, + conv_transpose_gemm_f64, + |a, b| a + b +); +conv_transpose3d_typed!( + conv_transpose3d_f16, + f16, + DType::F16, + f16::from_f32(0.0), + conv_transpose_gemm_f16, + |a: f16, b: f16| f16::from_f32(a.to_f32() + b.to_f32()) +); +bf16_via_f32!( + conv_transpose3d_bf16, + conv_transpose3d_f32, + 3, + ConvTransposeOptions +); + +/// GEMM for conv_transpose: writes `c = a^T @ b` where a is [k,m] and b is [k,n]. +type ConvTransposeGemmFn = fn(&mut [T], &[T], &[T], usize, usize, usize); + +/// 3D transposed convolution via GEMM + col2im. +#[allow(clippy::too_many_arguments)] +fn conv_transpose3d_impl( + x: FlexTensor, + weight: FlexTensor, + bias: Option, + options: &ConvTransposeOptions<3>, + dtype: DType, + zero: T, + gemm_fn: ConvTransposeGemmFn, + add_fn: fn(T, T) -> T, +) -> FlexTensor { + let x = x.to_contiguous(); + let weight = weight.to_contiguous(); + + let x_shape = x.layout().shape(); + let w_shape = weight.layout().shape(); + + let batch_size = x_shape[0]; + let in_channels = x_shape[1]; + let in_d = x_shape[2]; + let in_h = x_shape[3]; + let in_w = x_shape[4]; + + // Weight shape for transpose: [in_channels, out_channels_per_group, kd, kh, kw] + let out_channels_per_group = w_shape[1]; + let kernel_d = w_shape[2]; + let kernel_h = w_shape[3]; + let kernel_w = w_shape[4]; + + let [stride_d, stride_h, stride_w] = options.stride; + let [pad_d, pad_h, pad_w] = options.padding; + let [pad_out_d, pad_out_h, pad_out_w] = options.padding_out; + let [dilation_d, dilation_h, dilation_w] = options.dilation; + let groups = options.groups; + + let out_channels = out_channels_per_group * groups; + let in_channels_per_group = in_channels / groups; + + let out_d = calculate_conv_transpose_output_size( + kernel_d, stride_d, pad_d, pad_out_d, dilation_d, in_d, + ); + let out_h = calculate_conv_transpose_output_size( + kernel_h, stride_h, pad_h, pad_out_h, dilation_h, in_h, + ); + let out_w = calculate_conv_transpose_output_size( + kernel_w, stride_w, pad_w, pad_out_w, dilation_w, in_w, + ); + + let x_data: &[T] = x.storage(); + let w_data: &[T] = weight.storage(); + + let k_spatial = [kernel_d, kernel_h, kernel_w] + .iter() + .try_fold(1usize, |acc, &x| acc.checked_mul(x)) + .expect("conv_transpose: kernel dimensions would overflow"); + let in_spatial = [in_d, in_h, in_w] + .iter() + .try_fold(1usize, |acc, &x| acc.checked_mul(x)) + .expect("conv_transpose: input spatial dimensions would overflow"); + let out_spatial = out_d * out_h * out_w; + let col_ch = out_channels_per_group + .checked_mul(k_spatial) + .expect("conv_transpose: columns dimensions would overflow"); + let columns_len = col_ch + .checked_mul(in_spatial) + .expect("conv_transpose: columns buffer size would overflow"); + + let output_size = [batch_size, out_channels, out_d, out_h, out_w] + .iter() + .try_fold(1usize, |acc, &x| acc.checked_mul(x)) + .expect("conv_transpose: output dimensions would overflow"); + let mut output = vec![zero; output_size]; + + // Reuse columns buffer across (batch, group) iterations; GEMM overwrites it fully. + let mut columns = vec![zero; columns_len]; + + for b in 0..batch_size { + for g in 0..groups { + let ic_start = g * in_channels_per_group; + let oc_start = g * out_channels_per_group; + + let x_offset = b * in_channels * in_spatial + ic_start * in_spatial; + let w_offset = ic_start * out_channels_per_group * k_spatial; + + let x_group = &x_data[x_offset..x_offset + in_channels_per_group * in_spatial]; + let w_group = &w_data[w_offset..w_offset + in_channels_per_group * col_ch]; + + gemm_fn( + &mut columns, + w_group, + x_group, + col_ch, + in_channels_per_group, + in_spatial, + ); + + // col2im: scatter columns into output for this (batch, group) + let out_base = b * out_channels * out_spatial; + for oc in 0..out_channels_per_group { + let out_ch_base = out_base + (oc_start + oc) * out_spatial; + let oc_col_base = oc * k_spatial; + + for kd in 0..kernel_d { + for kh in 0..kernel_h { + for kw in 0..kernel_w { + let k_idx = kd * kernel_h * kernel_w + kh * kernel_w + kw; + let col_base = (oc_col_base + k_idx) * in_spatial; + + for id in 0..in_d { + let od_raw = id * stride_d + kd * dilation_d; + if od_raw < pad_d { + continue; + } + let od = od_raw - pad_d; + if od >= out_d { + continue; + } + + for ih in 0..in_h { + let oh_raw = ih * stride_h + kh * dilation_h; + if oh_raw < pad_h { + continue; + } + let oh = oh_raw - pad_h; + if oh >= out_h { + continue; + } + + for iw in 0..in_w { + let ow_raw = iw * stride_w + kw * dilation_w; + if ow_raw < pad_w { + continue; + } + let ow = ow_raw - pad_w; + if ow >= out_w { + continue; + } + + let s = id * in_h * in_w + ih * in_w + iw; + let val = columns[col_base + s]; + let out_idx = + out_ch_base + od * out_h * out_w + oh * out_w + ow; + output[out_idx] = add_fn(output[out_idx], val); + } + } + } + } + } + } + } + } + } + + // Add bias if present + if let Some(bias) = bias { + let bias = bias.to_contiguous(); + let bias_data: &[T] = bias.storage(); + add_bias( + &mut output, + bias_data, + batch_size, + out_channels, + out_spatial, + add_fn, + ); + } + + let out_shape = Shape::from(vec![batch_size, out_channels, out_d, out_h, out_w]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + dtype, + ) +} + +// ============================================================================ +// gemm for conv_transpose: C[m,n] = A[k,m]^T @ B[k,n] +// ============================================================================ + +macro_rules! conv_transpose_gemm_typed { + ($fn_name:ident, $T:ty, $zero:expr, $one:expr) => { + fn $fn_name(c: &mut [$T], a: &[$T], b: &[$T], m: usize, k: usize, n: usize) { + debug_assert_eq!(c.len(), m * n); + debug_assert_eq!(a.len(), k * m); + debug_assert_eq!(b.len(), k * n); + #[cfg(feature = "rayon")] + let parallelism = if m * n * k >= 192 * 192 * 192 { + gemm::Parallelism::Rayon(0) + } else { + gemm::Parallelism::None + }; + #[cfg(not(feature = "rayon"))] + let parallelism = gemm::Parallelism::None; + unsafe { + gemm::gemm( + m, + n, + k, + c.as_mut_ptr(), + 1, // dst_cs + n as isize, // dst_rs + false, + a.as_ptr(), + m as isize, // lhs_cs: A^T column stride = row length of A + 1, // lhs_rs: A^T row stride = 1 + b.as_ptr(), + 1, // rhs_cs + n as isize, // rhs_rs + $zero, + $one, + false, + false, + false, + parallelism, + ); + } + } + }; +} + +conv_transpose_gemm_typed!(conv_transpose_gemm_f32, f32, 0.0f32, 1.0f32); +conv_transpose_gemm_typed!(conv_transpose_gemm_f64, f64, 0.0f64, 1.0f64); +conv_transpose_gemm_typed!( + conv_transpose_gemm_f16, + f16, + f16::from_f32(0.0), + f16::from_f32(1.0) +); + +// ============================================================================ +// Tests +// ============================================================================ + +// Tests kept here exercise flex-specific dtype storage paths (f64/f16) +// through the conv_transpose dispatch. Plain conv_transpose shape/stride/ +// padding/groups/dilation/multichannel/batch correctness is covered at +// the public API level by crates/burn-backend-tests/tests/tensor/float/ +// module/conv_transpose{1,2,3}d.rs, so those tests were removed from here. +// When adding new tests, keep them here only if they probe flex dtype +// dispatch or a flex-internal fast path; otherwise add them there. +#[cfg(test)] +mod tests { + use super::*; + use burn_backend::TensorData; + + #[test] + fn test_conv_transpose2d_f64() { + let x = FlexTensor::from_data(TensorData::new( + vec![1.0f64, 2.0, 3.0, 4.0], + vec![1, 1, 2, 2], + )); + let w = FlexTensor::from_data(TensorData::new(vec![1.0f64; 4], vec![1, 1, 2, 2])); + let opts = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [1, 1], 1); + let result = conv_transpose2d_f64(x, w, None, &opts); + let out: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(out, vec![1.0, 3.0, 2.0, 4.0, 10.0, 6.0, 3.0, 7.0, 4.0]); + } + + #[test] + fn test_conv_transpose2d_f16() { + let x_data: Vec = [1.0f32, 2.0, 3.0, 4.0] + .iter() + .map(|&v| f16::from_f32(v)) + .collect(); + let w_data: Vec = [1.0f32; 4].iter().map(|&v| f16::from_f32(v)).collect(); + let x = FlexTensor::from_data(TensorData::new(x_data, vec![1, 1, 2, 2])); + let w = FlexTensor::from_data(TensorData::new(w_data, vec![1, 1, 2, 2])); + let opts = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [1, 1], 1); + let result = conv_transpose2d_f16(x, w, None, &opts); + let out: Vec = result.into_data().to_vec().unwrap(); + let expected = [1.0f32, 3.0, 2.0, 4.0, 10.0, 6.0, 3.0, 7.0, 4.0]; + for (a, e) in out.iter().zip(expected.iter()) { + assert!((a.to_f32() - e).abs() < 0.1); + } + } +} diff --git a/crates/burn-flex/src/ops/cumulative.rs b/crates/burn-flex/src/ops/cumulative.rs new file mode 100644 index 0000000..ebd3780 --- /dev/null +++ b/crates/burn-flex/src/ops/cumulative.rs @@ -0,0 +1,264 @@ +//! Cumulative operations along a dimension. + +use alloc::vec; +use burn_backend::Element; +use burn_std::Bytes; +use bytemuck::Pod; +use num_traits::{Bounded, Num}; + +use crate::{FlexTensor, Layout}; + +/// Cumulative sum along a dimension. +/// +/// For each position along `dim`, output contains the sum of all elements +/// from index 0 up to and including that position. +pub fn cumsum( + tensor: FlexTensor, + dim: usize, +) -> FlexTensor { + cumulative_op(tensor, dim, E::zero(), |acc, val| acc + val) +} + +/// Cumulative product along a dimension. +pub fn cumprod( + tensor: FlexTensor, + dim: usize, +) -> FlexTensor { + cumulative_op(tensor, dim, E::one(), |acc, val| acc * val) +} + +/// Generic cumulative operation along a dimension. +/// +/// Uses blocked iteration for cache-friendly access: processes contiguous +/// inner blocks together rather than striding across memory one element at +/// a time. +fn cumulative_op( + tensor: FlexTensor, + dim: usize, + init: E, + op: F, +) -> FlexTensor +where + F: Fn(E, E) -> E, +{ + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let ndims = shape.num_dims(); + + assert!( + dim < ndims, + "dim {} out of bounds for {} dimensions", + dim, + ndims + ); + + let data: &[E] = tensor.storage(); + let total_size = shape.num_elements(); + let mut result = vec![E::default(); total_size]; + + let dim_size = shape[dim]; + // Contiguous block size after the cumulative dimension + let inner_size: usize = shape[dim + 1..].iter().product(); + // Number of outer blocks (dimensions before the cumulative dimension) + let outer_size: usize = shape[..dim].iter().product(); + let block_size = dim_size * inner_size; + + if inner_size == 1 { + // Scalar accumulator path: accumulator stays in a register. + for outer in 0..outer_size { + let base = outer * dim_size; + let mut acc = init; + for i in 0..dim_size { + acc = op(acc, data[base + i]); + result[base + i] = acc; + } + } + } else { + // Blocked path: process contiguous inner blocks together for + // cache-friendly access when the cumulative dim is not last. + let mut acc = vec![init; inner_size]; + for outer in 0..outer_size { + let base = outer * block_size; + acc.fill(init); + for i in 0..dim_size { + let offset = base + i * inner_size; + for j in 0..inner_size { + acc[j] = op(acc[j], data[offset + j]); + result[offset + j] = acc[j]; + } + } + } + } + + let bytes = Bytes::from_elems(result); + FlexTensor::new(bytes, Layout::contiguous(shape), E::dtype()) +} + +/// Cumulative operation for half-precision types, accumulating in f32. +fn cumulative_op_half( + tensor: FlexTensor, + dim: usize, + init: f32, + op: F, + to_f32: fn(E) -> f32, + from_f32: fn(f32) -> E, +) -> FlexTensor +where + F: Fn(f32, f32) -> f32, +{ + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let ndims = shape.num_dims(); + + assert!( + dim < ndims, + "dim {} out of bounds for {} dimensions", + dim, + ndims + ); + + let data: &[E] = tensor.storage(); + let total_size = shape.num_elements(); + let mut result = vec![E::default(); total_size]; + + let dim_size = shape[dim]; + let inner_size: usize = shape[dim + 1..].iter().product(); + let outer_size: usize = shape[..dim].iter().product(); + let block_size = dim_size * inner_size; + + // Accumulator buffer for f32 intermediate values + let mut acc = vec![init; inner_size]; + + for outer in 0..outer_size { + let base = outer * block_size; + + // Reset accumulators + acc.fill(init); + + for i in 0..dim_size { + let offset = base + i * inner_size; + for j in 0..inner_size { + acc[j] = op(acc[j], to_f32(data[offset + j])); + result[offset + j] = from_f32(acc[j]); + } + } + } + + let bytes = Bytes::from_elems(result); + FlexTensor::new(bytes, Layout::contiguous(shape), E::dtype()) +} + +// Type-specific wrappers + +pub fn cumsum_f32(tensor: FlexTensor, dim: usize) -> FlexTensor { + cumsum::(tensor, dim) +} + +pub fn cumsum_f64(tensor: FlexTensor, dim: usize) -> FlexTensor { + cumsum::(tensor, dim) +} + +pub fn cumprod_f32(tensor: FlexTensor, dim: usize) -> FlexTensor { + cumprod::(tensor, dim) +} + +pub fn cumprod_f64(tensor: FlexTensor, dim: usize) -> FlexTensor { + cumprod::(tensor, dim) +} + +/// Cumulative min along a dimension. +pub fn cummin( + tensor: FlexTensor, + dim: usize, +) -> FlexTensor { + cumulative_op(tensor, dim, E::max_value(), Ord::min) +} + +pub fn cummin_f32(tensor: FlexTensor, dim: usize) -> FlexTensor { + cumulative_op(tensor, dim, f32::INFINITY, |acc, val| { + if val.is_nan() || val < acc { val } else { acc } + }) +} + +pub fn cummin_f64(tensor: FlexTensor, dim: usize) -> FlexTensor { + cumulative_op(tensor, dim, f64::INFINITY, |acc, val| { + if val.is_nan() || val < acc { val } else { acc } + }) +} + +/// Cumulative max along a dimension. +pub fn cummax( + tensor: FlexTensor, + dim: usize, +) -> FlexTensor { + cumulative_op(tensor, dim, E::min_value(), Ord::max) +} + +pub fn cummax_f32(tensor: FlexTensor, dim: usize) -> FlexTensor { + cumulative_op(tensor, dim, f32::NEG_INFINITY, |acc, val| { + if val.is_nan() || val > acc { val } else { acc } + }) +} + +pub fn cummax_f64(tensor: FlexTensor, dim: usize) -> FlexTensor { + cumulative_op(tensor, dim, f64::NEG_INFINITY, |acc, val| { + if val.is_nan() || val > acc { val } else { acc } + }) +} + +pub fn cumsum_half( + tensor: FlexTensor, + dim: usize, + to_f32: fn(E) -> f32, + from_f32: fn(f32) -> E, +) -> FlexTensor { + cumulative_op_half(tensor, dim, 0.0, |acc, val| acc + val, to_f32, from_f32) +} + +pub fn cumprod_half( + tensor: FlexTensor, + dim: usize, + to_f32: fn(E) -> f32, + from_f32: fn(f32) -> E, +) -> FlexTensor { + cumulative_op_half(tensor, dim, 1.0, |acc, val| acc * val, to_f32, from_f32) +} + +pub fn cummin_half( + tensor: FlexTensor, + dim: usize, + to_f32: fn(E) -> f32, + from_f32: fn(f32) -> E, +) -> FlexTensor { + cumulative_op_half( + tensor, + dim, + f32::INFINITY, + |acc, val| if val.is_nan() || val < acc { val } else { acc }, + to_f32, + from_f32, + ) +} + +pub fn cummax_half( + tensor: FlexTensor, + dim: usize, + to_f32: fn(E) -> f32, + from_f32: fn(f32) -> E, +) -> FlexTensor { + cumulative_op_half( + tensor, + dim, + f32::NEG_INFINITY, + |acc, val| if val.is_nan() || val > acc { val } else { acc }, + to_f32, + from_f32, + ) +} + +// Cumsum / cumprod / cummin / cummax coverage (basic, stride variants, +// NaN propagation, int dtype) lives in +// crates/burn-backend-tests/tests/tensor/{float,int}/ops/cumulative.rs so +// every backend is exercised. When adding new tests, keep them here only +// if they probe flex-specific kernel internals; otherwise add them +// there. diff --git a/crates/burn-flex/src/ops/deform_conv.rs b/crates/burn-flex/src/ops/deform_conv.rs new file mode 100644 index 0000000..31e07de --- /dev/null +++ b/crates/burn-flex/src/ops/deform_conv.rs @@ -0,0 +1,1021 @@ +//! Deformable convolution implementation using im2col + GEMM. +//! +//! Deformable convolution applies learned offsets to the sampling grid, +//! allowing the network to adaptively adjust its receptive field. +//! +//! Optimization approach: +//! - Build deformable im2col matrix with bilinear-interpolated samples +//! - Use optimized GEMM for the actual convolution +//! - Rayon parallelism over batch dimension + +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float; + +use alloc::vec; +use alloc::vec::Vec; +use burn_backend::DType; +use burn_std::{Bytes, Shape}; + +use crate::{FlexTensor, Layout}; + +/// Build deformable im2col matrix for one batch sample and weight group. +/// +/// Fills `col` with shape [col_len, spatial_out] where each column holds +/// bilinear-interpolated and optionally masked samples for one output position. +#[allow(clippy::too_many_arguments)] +fn deform_im2col_f32( + col: &mut [f32], + x_data: &[f32], + offset_data: &[f32], + mask_data: Option<&[f32]>, + b: usize, + ic_start: usize, + channels_per_weight_group: usize, + channels_per_offset_group: usize, + channels_in: usize, + offset_groups: usize, + offset_channels: usize, + kernel_h: usize, + kernel_w: usize, + out_h: usize, + out_w: usize, + in_h: usize, + in_w: usize, + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + spatial_out: usize, +) { + for oh in 0..out_h { + for ow in 0..out_w { + let spatial_idx = oh * out_w + ow; + + for kh in 0..kernel_h { + for kw in 0..kernel_w { + let base_h = (oh * stride[0] + kh * dilation[0]) as f32 - padding[0] as f32; + let base_w = (ow * stride[1] + kw * dilation[1]) as f32 - padding[1] as f32; + + for ic in 0..channels_per_weight_group { + let global_ic = ic_start + ic; + let offset_group = global_ic / channels_per_offset_group; + + let kernel_idx = kh * kernel_w + kw; + let offset_idx_h = offset_group * kernel_h * kernel_w * 2 + kernel_idx * 2; + let offset_idx_w = offset_idx_h + 1; + + let offset_h_flat = b * offset_channels * spatial_out + + offset_idx_h * spatial_out + + spatial_idx; + let offset_w_flat = b * offset_channels * spatial_out + + offset_idx_w * spatial_out + + spatial_idx; + + let offset_h = offset_data[offset_h_flat]; + let offset_w = offset_data[offset_w_flat]; + + let sample_h = base_h + offset_h; + let sample_w = base_w + offset_w; + + let mut val = bilinear_interpolate( + x_data, + b, + global_ic, + in_h, + in_w, + channels_in, + sample_h, + sample_w, + ); + + if let Some(md) = mask_data { + let mask_idx_base = offset_group * kernel_h * kernel_w + kernel_idx; + let mask_idx = b * (offset_groups * kernel_h * kernel_w) * spatial_out + + mask_idx_base * spatial_out + + spatial_idx; + val *= md[mask_idx]; + } + + let col_row = kh * kernel_w * channels_per_weight_group + + kw * channels_per_weight_group + + ic; + col[col_row * spatial_out + spatial_idx] = val; + } + } + } + } + } +} + +/// Deformable 2D convolution using im2col + GEMM. +/// +/// # Arguments +/// * `x` - Input tensor [batch, channels_in, height, width] +/// * `offset` - Offset tensor \[batch, offset_groups * kernel_h * kernel_w * 2, out_h, out_w\] +/// * `weight` - Weight tensor \[channels_out, channels_in/weight_groups, kernel_h, kernel_w\] +/// * `mask` - Optional mask tensor \[batch, offset_groups * kernel_h * kernel_w, out_h, out_w\] +/// * `bias` - Optional bias tensor \[channels_out\] +/// * `stride` - Stride \[stride_h, stride_w\] +/// * `padding` - Padding \[pad_h, pad_w\] +/// * `dilation` - Dilation \[dil_h, dil_w\] +/// * `weight_groups` - Number of weight groups +/// * `offset_groups` - Number of offset groups +#[allow(clippy::too_many_arguments)] +pub fn deform_conv2d_f32( + x: FlexTensor, + offset: FlexTensor, + weight: FlexTensor, + mask: Option, + bias: Option, + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + weight_groups: usize, + offset_groups: usize, +) -> FlexTensor { + let x = x.to_contiguous(); + let offset = offset.to_contiguous(); + let weight = weight.to_contiguous(); + let mask = mask.map(|m| m.to_contiguous()); + let bias = bias.map(|b| b.to_contiguous()); + + let x_shape = x.layout().shape(); + let weight_shape = weight.layout().shape(); + let offset_shape = offset.layout().shape(); + + let batch = x_shape[0]; + let channels_in = x_shape[1]; + let in_h = x_shape[2]; + let in_w = x_shape[3]; + + let channels_out = weight_shape[0]; + let channels_per_weight_group = weight_shape[1]; // channels_in / weight_groups + let kernel_h = weight_shape[2]; + let kernel_w = weight_shape[3]; + + let out_h = offset_shape[2]; + let out_w = offset_shape[3]; + + let x_data: &[f32] = x.storage(); + let offset_data: &[f32] = offset.storage(); + let weight_data: &[f32] = weight.storage(); + let mask_data: Option<&[f32]> = mask.as_ref().map(|m| m.storage()); + let bias_data: Option<&[f32]> = bias.as_ref().map(|b| b.storage()); + + let channels_per_offset_group = channels_in / offset_groups; + let out_channels_per_weight_group = channels_out / weight_groups; + let spatial_out = out_h * out_w; + let col_len = channels_per_weight_group * kernel_h * kernel_w; + let offset_channels = offset_shape[1]; + + // Flatten weights to [channels_out, col_len] for GEMM + // Layout: [oc, ic, kh, kw] -> [oc, kh * kw * ic] + let mut w_flat = vec![0.0f32; channels_out * col_len]; + for oc in 0..channels_out { + for kh in 0..kernel_h { + for kw in 0..kernel_w { + for ic in 0..channels_per_weight_group { + let w_idx = oc * channels_per_weight_group * kernel_h * kernel_w + + ic * kernel_h * kernel_w + + kh * kernel_w + + kw; + let flat_idx = oc * col_len + + kh * kernel_w * channels_per_weight_group + + kw * channels_per_weight_group + + ic; + w_flat[flat_idx] = weight_data[w_idx]; + } + } + } + } + + #[cfg(feature = "rayon")] + let output = { + use rayon::prelude::*; + + let results: Vec> = (0..batch) + .into_par_iter() + .map(|b| { + let mut batch_output = vec![0.0f32; channels_out * spatial_out]; + + // Process each weight group + for g in 0..weight_groups { + let ic_start = g * channels_per_weight_group; + let oc_start = g * out_channels_per_weight_group; + + // Build deformable im2col for this batch and group + let mut col = vec![0.0f32; col_len * spatial_out]; + + deform_im2col_f32( + &mut col, + x_data, + offset_data, + mask_data, + b, + ic_start, + channels_per_weight_group, + channels_per_offset_group, + channels_in, + offset_groups, + offset_channels, + kernel_h, + kernel_w, + out_h, + out_w, + in_h, + in_w, + stride, + padding, + dilation, + spatial_out, + ); + + // GEMM: w_group[out_c_per_wg, col_len] @ col[col_len, spatial_out] + // Result: [out_c_per_wg, spatial_out] + let w_start = oc_start * col_len; + let w_end = w_start + out_channels_per_weight_group * col_len; + let w_group = &w_flat[w_start..w_end]; + + let result = gemm_f32( + w_group, + &col, + out_channels_per_weight_group, + col_len, + spatial_out, + ); + + // Copy to output + for oc_local in 0..out_channels_per_weight_group { + let oc = oc_start + oc_local; + for s in 0..spatial_out { + batch_output[oc * spatial_out + s] = result[oc_local * spatial_out + s]; + } + } + } + + // Add bias + if let Some(bd) = bias_data { + for oc in 0..channels_out { + for s in 0..spatial_out { + batch_output[oc * spatial_out + s] += bd[oc]; + } + } + } + + batch_output + }) + .collect(); + + // Flatten results + let mut output = vec![0.0f32; batch * channels_out * spatial_out]; + for (b, batch_out) in results.into_iter().enumerate() { + let start = b * channels_out * spatial_out; + output[start..start + channels_out * spatial_out].copy_from_slice(&batch_out); + } + output + }; + + #[cfg(not(feature = "rayon"))] + let output = { + let mut output = vec![0.0f32; batch * channels_out * spatial_out]; + + for b in 0..batch { + // Process each weight group + for g in 0..weight_groups { + let ic_start = g * channels_per_weight_group; + let oc_start = g * out_channels_per_weight_group; + + // Build deformable im2col for this batch and group + let mut col = vec![0.0f32; col_len * spatial_out]; + + deform_im2col_f32( + &mut col, + x_data, + offset_data, + mask_data, + b, + ic_start, + channels_per_weight_group, + channels_per_offset_group, + channels_in, + offset_groups, + offset_channels, + kernel_h, + kernel_w, + out_h, + out_w, + in_h, + in_w, + stride, + padding, + dilation, + spatial_out, + ); + + // GEMM + let w_start = oc_start * col_len; + let w_end = w_start + out_channels_per_weight_group * col_len; + let w_group = &w_flat[w_start..w_end]; + + let result = gemm_f32( + w_group, + &col, + out_channels_per_weight_group, + col_len, + spatial_out, + ); + + // Copy to output + for oc_local in 0..out_channels_per_weight_group { + let oc = oc_start + oc_local; + for s in 0..spatial_out { + let out_idx = b * channels_out * spatial_out + oc * spatial_out + s; + output[out_idx] = result[oc_local * spatial_out + s]; + } + } + } + + // Add bias + if let Some(bd) = bias_data { + #[allow(clippy::needless_range_loop)] + for oc in 0..channels_out { + for s in 0..spatial_out { + let idx = b * channels_out * spatial_out + oc * spatial_out + s; + output[idx] += bd[oc]; + } + } + } + } + + output + }; + + let out_shape = Shape::from(vec![batch, channels_out, out_h, out_w]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + DType::F32, + ) +} + +/// GEMM: C = A @ B where A is [m, k], B is [k, n], result C is [m, n] +#[inline] +fn gemm_f32(a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec { + let mut c = vec![0.0f32; m * n]; + unsafe { + gemm::gemm( + m, + n, + k, + c.as_mut_ptr(), + 1, // dst_cs: column stride = 1 (row-major) + n as isize, // dst_rs: row stride = n (row-major) + false, + a.as_ptr(), + 1, // lhs_cs: column stride = 1 (row-major) + k as isize, // lhs_rs: row stride = k (row-major) + b.as_ptr(), + 1, // rhs_cs: column stride = 1 (row-major) + n as isize, // rhs_rs: row stride = n (row-major) + 0.0, // alpha: dst = alpha*dst + beta*lhs*rhs + 1.0, // beta + false, + false, + false, + gemm::Parallelism::None, + ); + } + c +} + +/// Bilinear interpolation for sampling at fractional coordinates. +#[inline] +#[allow(clippy::too_many_arguments)] +fn bilinear_interpolate( + data: &[f32], + batch: usize, + channel: usize, + height: usize, + width: usize, + channels: usize, + h: f32, + w: f32, +) -> f32 { + // Out of bounds check + if h <= -1.0 || h >= height as f32 || w <= -1.0 || w >= width as f32 { + return 0.0; + } + + let h_low = h.floor(); + let w_low = w.floor(); + let h_high = (h_low + 1.0) as usize; + let w_high = (w_low + 1.0) as usize; + + let base = batch * channels * height * width + channel * height * width; + + let v1 = if h_low >= 0.0 && w_low >= 0.0 { + data[base + (h_low as usize) * width + (w_low as usize)] + } else { + 0.0 + }; + let v2 = if h_low >= 0.0 && w_high < width { + data[base + (h_low as usize) * width + w_high] + } else { + 0.0 + }; + let v3 = if h_high < height && w_low >= 0.0 { + data[base + h_high * width + (w_low as usize)] + } else { + 0.0 + }; + let v4 = if h_high < height && w_high < width { + data[base + h_high * width + w_high] + } else { + 0.0 + }; + + let lh = h - h_low; + let lw = w - w_low; + let hh = 1.0 - lh; + let hw = 1.0 - lw; + + hh * hw * v1 + hh * lw * v2 + lh * hw * v3 + lh * lw * v4 +} + +/// Backward pass for deformable 2D convolution (f32). +/// +/// Returns (x_grad, offset_grad, weight_grad, mask_grad, bias_grad). +#[allow(clippy::too_many_arguments)] +pub fn deform_conv2d_backward_f32( + x: FlexTensor, + offset: FlexTensor, + weight: FlexTensor, + mask: Option, + bias: Option, + output_grad: FlexTensor, + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + weight_groups: usize, + offset_groups: usize, +) -> ( + FlexTensor, + FlexTensor, + FlexTensor, + Option, + Option, +) { + let x = x.to_contiguous(); + let offset = offset.to_contiguous(); + let weight = weight.to_contiguous(); + let mask = mask.map(|m| m.to_contiguous()); + let output_grad = output_grad.to_contiguous(); + + let x_shape = x.layout().shape(); + let weight_shape = weight.layout().shape(); + let offset_shape = offset.layout().shape(); + let out_grad_shape = output_grad.layout().shape(); + + let batch = x_shape[0]; + let channels_in = x_shape[1]; + let in_h = x_shape[2]; + let in_w = x_shape[3]; + + let channels_out = weight_shape[0]; + let kernel_h = weight_shape[2]; + let kernel_w = weight_shape[3]; + + let out_h = out_grad_shape[2]; + let out_w = out_grad_shape[3]; + + let x_data: &[f32] = x.storage(); + let offset_data: &[f32] = offset.storage(); + let weight_data: &[f32] = weight.storage(); + let mask_data: Option<&[f32]> = mask.as_ref().map(|m| m.storage()); + let out_grad_data: &[f32] = output_grad.storage(); + + let channels_per_offset_group = channels_in / offset_groups; + let channels_per_weight_group = channels_in / weight_groups; + let out_channels_per_weight_group = channels_out / weight_groups; + + // Initialize gradients + let mut x_grad = vec![0.0f32; batch * channels_in * in_h * in_w]; + let mut offset_grad = vec![0.0f32; batch * offset_shape[1] * out_h * out_w]; + let mut weight_grad = vec![0.0f32; weight_shape.num_elements()]; + let mut mask_grad = mask + .as_ref() + .map(|m| vec![0.0f32; m.layout().shape().num_elements()]); + let mut bias_grad = bias.as_ref().map(|_| vec![0.0f32; channels_out]); + + // Compute bias gradient (sum over batch and spatial dimensions) + if let Some(ref mut bg) = bias_grad { + for b in 0..batch { + for (oc, bg_oc) in bg.iter_mut().enumerate() { + for oh in 0..out_h { + for ow in 0..out_w { + let idx = + b * channels_out * out_h * out_w + oc * out_h * out_w + oh * out_w + ow; + *bg_oc += out_grad_data[idx]; + } + } + } + } + } + + // Main backward loop + for b in 0..batch { + for oc in 0..channels_out { + let weight_group = oc / out_channels_per_weight_group; + + for oh in 0..out_h { + for ow in 0..out_w { + let out_idx = + b * channels_out * out_h * out_w + oc * out_h * out_w + oh * out_w + ow; + let grad_out = out_grad_data[out_idx]; + + let ic_start = weight_group * channels_per_weight_group; + let ic_end = ic_start + channels_per_weight_group; + + for ic in ic_start..ic_end { + let offset_group = ic / channels_per_offset_group; + + for kh in 0..kernel_h { + for kw in 0..kernel_w { + let base_h = (oh * stride[0]) as f32 + (kh * dilation[0]) as f32 + - padding[0] as f32; + let base_w = (ow * stride[1]) as f32 + (kw * dilation[1]) as f32 + - padding[1] as f32; + + let kernel_idx = kh * kernel_w + kw; + let offset_idx_h = + offset_group * kernel_h * kernel_w * 2 + kernel_idx * 2; + let offset_idx_w = offset_idx_h + 1; + + let offset_h_flat = b * offset_shape[1] * out_h * out_w + + offset_idx_h * out_h * out_w + + oh * out_w + + ow; + let offset_w_flat = b * offset_shape[1] * out_h * out_w + + offset_idx_w * out_h * out_w + + oh * out_w + + ow; + + let off_h = offset_data[offset_h_flat]; + let off_w = offset_data[offset_w_flat]; + + let sample_h = base_h + off_h; + let sample_w = base_w + off_w; + + // Get mask value + let (mask_val, mask_flat_idx) = if let Some(md) = mask_data { + let mask_idx_base = + offset_group * kernel_h * kernel_w + kernel_idx; + let mask_idx = + b * (offset_groups * kernel_h * kernel_w) * out_h * out_w + + mask_idx_base * out_h * out_w + + oh * out_w + + ow; + (md[mask_idx], Some(mask_idx)) + } else { + (1.0, None) + }; + + // Weight index + let weight_ic = ic - ic_start; + let weight_idx = oc + * (channels_per_weight_group * kernel_h * kernel_w) + + weight_ic * kernel_h * kernel_w + + kh * kernel_w + + kw; + + let w = weight_data[weight_idx]; + + // Get interpolated value for weight gradient + let interp_val = bilinear_interpolate( + x_data, + b, + ic, + in_h, + in_w, + channels_in, + sample_h, + sample_w, + ); + + // Weight gradient + weight_grad[weight_idx] += grad_out * mask_val * interp_val; + + // Mask gradient + if let (Some(mg), Some(midx)) = (&mut mask_grad, mask_flat_idx) { + mg[midx] += grad_out * w * interp_val; + } + + // Input and offset gradients via bilinear interpolation backward + let grad_val = grad_out * mask_val * w; + bilinear_interpolate_backward_f32( + &mut x_grad, + &mut offset_grad, + x_data, + b, + ic, + in_h, + in_w, + channels_in, + sample_h, + sample_w, + grad_val, + offset_h_flat, + offset_w_flat, + ); + } + } + } + } + } + } + } + + // Build output tensors + let x_grad_tensor = FlexTensor::new( + Bytes::from_elems(x_grad), + Layout::contiguous(x_shape.clone()), + DType::F32, + ); + let offset_grad_tensor = FlexTensor::new( + Bytes::from_elems(offset_grad), + Layout::contiguous(offset_shape.clone()), + DType::F32, + ); + let weight_grad_tensor = FlexTensor::new( + Bytes::from_elems(weight_grad), + Layout::contiguous(weight_shape.clone()), + DType::F32, + ); + let mask_grad_tensor = mask_grad.map(|mg| { + FlexTensor::new( + Bytes::from_elems(mg), + Layout::contiguous(mask.as_ref().unwrap().layout().shape().clone()), + DType::F32, + ) + }); + let bias_grad_tensor = bias_grad.map(|bg| { + FlexTensor::new( + Bytes::from_elems(bg), + Layout::contiguous(Shape::from(vec![channels_out])), + DType::F32, + ) + }); + + ( + x_grad_tensor, + offset_grad_tensor, + weight_grad_tensor, + mask_grad_tensor, + bias_grad_tensor, + ) +} + +/// Backward pass for bilinear interpolation - computes gradients for input and offsets. +#[allow(clippy::too_many_arguments)] +#[inline] +fn bilinear_interpolate_backward_f32( + x_grad: &mut [f32], + offset_grad: &mut [f32], + x_data: &[f32], + batch: usize, + channel: usize, + height: usize, + width: usize, + channels: usize, + h: f32, + w: f32, + grad_val: f32, + offset_h_flat: usize, + offset_w_flat: usize, +) { + // Out of bounds check + if h <= -1.0 || h >= height as f32 || w <= -1.0 || w >= width as f32 { + return; + } + + let h_low = h.floor(); + let w_low = w.floor(); + let h_high = h_low + 1.0; + let w_high = w_low + 1.0; + + let lh = h - h_low; + let lw = w - w_low; + let hh = 1.0 - lh; + let hw = 1.0 - lw; + + let base = batch * channels * height * width + channel * height * width; + + // Get input values for offset gradient computation + let v1 = if h_low >= 0.0 && w_low >= 0.0 { + x_data[base + (h_low as usize) * width + (w_low as usize)] + } else { + 0.0 + }; + let v2 = if h_low >= 0.0 && (w_high as usize) < width { + x_data[base + (h_low as usize) * width + (w_high as usize)] + } else { + 0.0 + }; + let v3 = if (h_high as usize) < height && w_low >= 0.0 { + x_data[base + (h_high as usize) * width + (w_low as usize)] + } else { + 0.0 + }; + let v4 = if (h_high as usize) < height && (w_high as usize) < width { + x_data[base + (h_high as usize) * width + (w_high as usize)] + } else { + 0.0 + }; + + // Input gradient (distribute grad_val to the 4 corners) + if h_low >= 0.0 && w_low >= 0.0 { + let idx = base + (h_low as usize) * width + (w_low as usize); + x_grad[idx] += hh * hw * grad_val; + } + if h_low >= 0.0 && (w_high as usize) < width { + let idx = base + (h_low as usize) * width + (w_high as usize); + x_grad[idx] += hh * lw * grad_val; + } + if (h_high as usize) < height && w_low >= 0.0 { + let idx = base + (h_high as usize) * width + (w_low as usize); + x_grad[idx] += lh * hw * grad_val; + } + if (h_high as usize) < height && (w_high as usize) < width { + let idx = base + (h_high as usize) * width + (w_high as usize); + x_grad[idx] += lh * lw * grad_val; + } + + // Offset gradient (derivative of bilinear interpolation w.r.t. coordinates) + let grad_h = hw * (v3 - v1) + lw * (v4 - v2); + let grad_w = hh * (v2 - v1) + lh * (v4 - v3); + + offset_grad[offset_h_flat] += grad_val * grad_h; + offset_grad[offset_w_flat] += grad_val * grad_w; +} + +/// f64 version of deform_conv2d using im2col + GEMM. +#[allow(clippy::too_many_arguments)] +pub fn deform_conv2d_f64( + x: FlexTensor, + offset: FlexTensor, + weight: FlexTensor, + mask: Option, + bias: Option, + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + weight_groups: usize, + offset_groups: usize, +) -> FlexTensor { + let x = x.to_contiguous(); + let offset = offset.to_contiguous(); + let weight = weight.to_contiguous(); + let mask = mask.map(|m| m.to_contiguous()); + let bias = bias.map(|b| b.to_contiguous()); + + let x_shape = x.layout().shape(); + let weight_shape = weight.layout().shape(); + let offset_shape = offset.layout().shape(); + + let batch = x_shape[0]; + let channels_in = x_shape[1]; + let in_h = x_shape[2]; + let in_w = x_shape[3]; + + let channels_out = weight_shape[0]; + let channels_per_weight_group = weight_shape[1]; + let kernel_h = weight_shape[2]; + let kernel_w = weight_shape[3]; + + let out_h = offset_shape[2]; + let out_w = offset_shape[3]; + + let x_data: &[f64] = x.storage(); + let offset_data: &[f64] = offset.storage(); + let weight_data: &[f64] = weight.storage(); + let mask_data: Option<&[f64]> = mask.as_ref().map(|m| m.storage()); + let bias_data: Option<&[f64]> = bias.as_ref().map(|b| b.storage()); + + let channels_per_offset_group = channels_in / offset_groups; + let out_channels_per_weight_group = channels_out / weight_groups; + let spatial_out = out_h * out_w; + let col_len = channels_per_weight_group * kernel_h * kernel_w; + + // Flatten weights + let mut w_flat = vec![0.0f64; channels_out * col_len]; + for oc in 0..channels_out { + for kh in 0..kernel_h { + for kw in 0..kernel_w { + for ic in 0..channels_per_weight_group { + let w_idx = oc * channels_per_weight_group * kernel_h * kernel_w + + ic * kernel_h * kernel_w + + kh * kernel_w + + kw; + let flat_idx = oc * col_len + + kh * kernel_w * channels_per_weight_group + + kw * channels_per_weight_group + + ic; + w_flat[flat_idx] = weight_data[w_idx]; + } + } + } + } + + let mut output = vec![0.0f64; batch * channels_out * spatial_out]; + + for b in 0..batch { + for g in 0..weight_groups { + let ic_start = g * channels_per_weight_group; + let oc_start = g * out_channels_per_weight_group; + + let mut col = vec![0.0f64; col_len * spatial_out]; + + for oh in 0..out_h { + for ow in 0..out_w { + let spatial_idx = oh * out_w + ow; + + for kh in 0..kernel_h { + for kw in 0..kernel_w { + let base_h = + (oh * stride[0] + kh * dilation[0]) as f64 - padding[0] as f64; + let base_w = + (ow * stride[1] + kw * dilation[1]) as f64 - padding[1] as f64; + + for ic in 0..channels_per_weight_group { + let global_ic = ic_start + ic; + let offset_group = global_ic / channels_per_offset_group; + + let kernel_idx = kh * kernel_w + kw; + let offset_idx_h = + offset_group * kernel_h * kernel_w * 2 + kernel_idx * 2; + let offset_idx_w = offset_idx_h + 1; + + let offset_h_flat = b * offset_shape[1] * spatial_out + + offset_idx_h * spatial_out + + spatial_idx; + let offset_w_flat = b * offset_shape[1] * spatial_out + + offset_idx_w * spatial_out + + spatial_idx; + + let offset_h = offset_data[offset_h_flat]; + let offset_w = offset_data[offset_w_flat]; + + let sample_h = base_h + offset_h; + let sample_w = base_w + offset_w; + + let mut val = bilinear_interpolate_f64( + x_data, + b, + global_ic, + in_h, + in_w, + channels_in, + sample_h, + sample_w, + ); + + if let Some(md) = mask_data { + let mask_idx_base = + offset_group * kernel_h * kernel_w + kernel_idx; + let mask_idx = + b * (offset_groups * kernel_h * kernel_w) * spatial_out + + mask_idx_base * spatial_out + + spatial_idx; + val *= md[mask_idx]; + } + + let col_row = kh * kernel_w * channels_per_weight_group + + kw * channels_per_weight_group + + ic; + col[col_row * spatial_out + spatial_idx] = val; + } + } + } + } + } + + // GEMM + let w_start = oc_start * col_len; + let w_end = w_start + out_channels_per_weight_group * col_len; + let w_group = &w_flat[w_start..w_end]; + + let result = gemm_f64( + w_group, + &col, + out_channels_per_weight_group, + col_len, + spatial_out, + ); + + for oc_local in 0..out_channels_per_weight_group { + let oc = oc_start + oc_local; + for s in 0..spatial_out { + let out_idx = b * channels_out * spatial_out + oc * spatial_out + s; + output[out_idx] = result[oc_local * spatial_out + s]; + } + } + } + + if let Some(bd) = bias_data { + for (oc, &bias_val) in bd.iter().enumerate() { + for s in 0..spatial_out { + let idx = b * channels_out * spatial_out + oc * spatial_out + s; + output[idx] += bias_val; + } + } + } + } + + let out_shape = Shape::from(vec![batch, channels_out, out_h, out_w]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + DType::F64, + ) +} + +#[inline] +fn gemm_f64(a: &[f64], b: &[f64], m: usize, k: usize, n: usize) -> Vec { + let mut c = vec![0.0f64; m * n]; + unsafe { + gemm::gemm( + m, + n, + k, + c.as_mut_ptr(), + 1, + n as isize, + false, + a.as_ptr(), + 1, + k as isize, + b.as_ptr(), + 1, + n as isize, + 0.0, + 1.0, + false, + false, + false, + gemm::Parallelism::None, + ); + } + c +} + +#[allow(clippy::too_many_arguments)] +#[inline] +fn bilinear_interpolate_f64( + data: &[f64], + batch: usize, + channel: usize, + height: usize, + width: usize, + channels: usize, + h: f64, + w: f64, +) -> f64 { + if h <= -1.0 || h >= height as f64 || w <= -1.0 || w >= width as f64 { + return 0.0; + } + + let h_low = h.floor(); + let w_low = w.floor(); + let h_high = (h_low + 1.0) as usize; + let w_high = (w_low + 1.0) as usize; + + let base = batch * channels * height * width + channel * height * width; + + let v1 = if h_low >= 0.0 && w_low >= 0.0 { + data[base + (h_low as usize) * width + (w_low as usize)] + } else { + 0.0 + }; + let v2 = if h_low >= 0.0 && w_high < width { + data[base + (h_low as usize) * width + w_high] + } else { + 0.0 + }; + let v3 = if h_high < height && w_low >= 0.0 { + data[base + h_high * width + (w_low as usize)] + } else { + 0.0 + }; + let v4 = if h_high < height && w_high < width { + data[base + h_high * width + w_high] + } else { + 0.0 + }; + + let lh = h - h_low; + let lw = w - w_low; + let hh = 1.0 - lh; + let hw = 1.0 - lw; + + hh * hw * v1 + hh * lw * v2 + lh * hw * v3 + lh * lw * v4 +} diff --git a/crates/burn-flex/src/ops/expand.rs b/crates/burn-flex/src/ops/expand.rs new file mode 100644 index 0000000..5caa3d8 --- /dev/null +++ b/crates/burn-flex/src/ops/expand.rs @@ -0,0 +1,312 @@ +//! Expand operation for broadcasting tensors to larger shapes. + +use alloc::vec; +use alloc::vec::Vec; +use burn_std::Shape; + +use crate::{FlexTensor, Layout}; + +/// Compute the broadcast shape of two tensors. +/// +/// Returns the shape that both tensors can be expanded to for element-wise operations. +pub fn broadcast_shape(lhs: &Shape, rhs: &Shape) -> Shape { + let max_dims = lhs.num_dims().max(rhs.num_dims()); + let mut result = vec![0; max_dims]; + + for (i, out) in result.iter_mut().enumerate() { + let lhs_idx = i as isize + lhs.num_dims() as isize - max_dims as isize; + let rhs_idx = i as isize + rhs.num_dims() as isize - max_dims as isize; + + let lhs_dim = if lhs_idx >= 0 { + lhs[lhs_idx as usize] + } else { + 1 + }; + let rhs_dim = if rhs_idx >= 0 { + rhs[rhs_idx as usize] + } else { + 1 + }; + + if lhs_dim == rhs_dim { + *out = lhs_dim; + } else if lhs_dim == 1 { + *out = rhs_dim; + } else if rhs_dim == 1 { + *out = lhs_dim; + } else { + panic!( + "broadcast_shape: incompatible dimensions {} and {} at position {}", + lhs_dim, rhs_dim, i + ); + } + } + + Shape::from(result) +} + +/// Broadcast two tensors to the same shape for binary operations. +pub fn broadcast_binary(lhs: FlexTensor, rhs: FlexTensor) -> (FlexTensor, FlexTensor) { + let lhs_shape = lhs.layout().shape().clone(); + let rhs_shape = rhs.layout().shape().clone(); + + if lhs_shape == rhs_shape { + return (lhs, rhs); + } + + let target = broadcast_shape(&lhs_shape, &rhs_shape); + + let lhs_expanded = if lhs_shape == target { + lhs + } else { + expand(lhs, target.clone()) + }; + let rhs_expanded = if rhs_shape == target { + rhs + } else { + expand(rhs, target) + }; + + (lhs_expanded, rhs_expanded) +} + +/// Expand a tensor to a larger shape by broadcasting. +/// +/// Dimensions of size 1 can be expanded to any size. The result is a view +/// that doesn't copy data - it uses stride 0 for expanded dimensions. +pub fn expand(tensor: FlexTensor, target_shape: Shape) -> FlexTensor { + // Capture values we need before consuming tensor + let src_dims = tensor.layout().shape().to_vec(); + let src_strides = tensor.layout().strides().to_vec(); + let start_offset = tensor.layout().start_offset(); + let dtype = tensor.dtype(); + + let src_ndims = src_dims.len(); + let target_ndims = target_shape.num_dims(); + + // Broadcasting only prepends new dims; it never drops existing ones. + // A target with fewer dims than the source would silently discard the + // trailing source strides and produce an invalid layout. + assert!( + target_ndims >= src_ndims, + "expand: target rank ({}) must be >= source rank ({}); \ + broadcasting cannot drop dimensions", + target_ndims, + src_ndims + ); + + // Prepend 1s to source shape if needed (for broadcasting like [3] -> [2, 3]) + let dim_diff = target_ndims - src_ndims; + + let mut new_strides = Vec::with_capacity(target_ndims); + + for i in 0..target_ndims { + let target_dim = target_shape[i]; + + if i < dim_diff { + // New dimension prepended - must be broadcastable from size 1 + new_strides.push(0); + } else { + let src_idx = i - dim_diff; + let src_dim = src_dims[src_idx]; + let src_stride = src_strides[src_idx]; + + if src_dim == target_dim { + // Same size - keep stride + new_strides.push(src_stride); + } else if src_dim == 1 { + // Broadcast dimension - stride becomes 0 + new_strides.push(0); + } else { + panic!( + "expand: cannot expand dimension {} from {} to {}", + i, src_dim, target_dim + ); + } + } + } + + let new_layout = Layout::new(target_shape, new_strides, start_offset); + FlexTensor::from_arc(tensor.data_arc(), new_layout, dtype) +} + +// Tests kept here probe flex-internal expand behavior: stride metadata +// (stride 0 on broadcast dims, preservation of negative strides on +// flipped inputs, preserved start-offset on narrowed inputs) and the +// flex-only `broadcast_binary` helper. Public-API expand coverage for +// transpose/flip/narrow variants lives in +// crates/burn-backend-tests/tests/tensor/{float,int,bool}/ops/expand.rs +// so it runs on every backend. +#[cfg(test)] +mod tests { + use super::*; + use burn_backend::TensorData; + + #[test] + fn test_expand_1d_to_2d() { + // [3] -> [2, 3] + let tensor = FlexTensor::from_data(TensorData::new(vec![1.0f32, 2.0, 3.0], [3])); + let expanded = expand(tensor, Shape::new([2, 3])); + + assert_eq!(expanded.layout().shape().to_vec(), vec![2, 3]); + assert_eq!(expanded.layout().strides(), &[0, 1]); + } + + #[test] + fn test_expand_broadcast_dim() { + // [3, 1] -> [3, 4] + let tensor = FlexTensor::from_data(TensorData::new(vec![1.0f32, 2.0, 3.0], [3, 1])); + let expanded = expand(tensor, Shape::new([3, 4])); + + assert_eq!(expanded.layout().shape().to_vec(), vec![3, 4]); + // Original strides for [3, 1] would be [1, 1] + // After expand to [3, 4], stride for dim 1 becomes 0 + assert_eq!(expanded.layout().strides()[1], 0); + } + + #[test] + fn test_expand_same_shape() { + // [2, 3] -> [2, 3] (no change) + let tensor = FlexTensor::from_data(TensorData::new( + vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], + [2, 3], + )); + let original_strides = tensor.layout().strides().to_vec(); + let expanded = expand(tensor, Shape::new([2, 3])); + + assert_eq!(expanded.layout().shape().to_vec(), vec![2, 3]); + assert_eq!(expanded.layout().strides(), &original_strides); + } + + // === Non-contiguous tensor tests === + + #[test] + fn test_expand_transposed() { + // [[1, 2], [3, 4]] transposed -> [[1, 3], [2, 4]] with strides [1, 2] + // Expand [2, 2] -> [3, 2, 2] by prepending dimension + let tensor = FlexTensor::from_data(TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0], [2, 2])); + let transposed = tensor.transpose(0, 1); + assert!(!transposed.is_contiguous()); + assert_eq!(transposed.layout().strides(), &[1, 2]); + + let expanded = expand(transposed, Shape::new([3, 2, 2])); + assert_eq!(expanded.layout().shape().to_vec(), vec![3, 2, 2]); + // New dim with stride 0, original strides preserved + assert_eq!(expanded.layout().strides(), &[0, 1, 2]); + + // Verify content: should see same transposed values repeated 3 times + let data: Vec = expanded.into_data().to_vec().unwrap(); + // [[1, 3], [2, 4]] repeated 3 times + assert_eq!( + data, + vec![1.0, 3.0, 2.0, 4.0, 1.0, 3.0, 2.0, 4.0, 1.0, 3.0, 2.0, 4.0] + ); + } + + #[test] + fn test_expand_flipped_1d() { + // [1, 2, 3] flipped -> [3, 2, 1] with negative stride + // Expand [3] -> [2, 3] + let tensor = FlexTensor::from_data(TensorData::new(vec![1.0f32, 2.0, 3.0], [3])); + let flipped = crate::ops::flip::flip(tensor, &[0]); + assert!(flipped.layout().strides()[0] < 0); + + let expanded = expand(flipped, Shape::new([2, 3])); + assert_eq!(expanded.layout().shape().to_vec(), vec![2, 3]); + // Stride 0 for new broadcast dim, negative stride preserved + assert_eq!(expanded.layout().strides()[0], 0); + assert!(expanded.layout().strides()[1] < 0); + + // Verify content: [3, 2, 1] repeated twice + let data: Vec = expanded.into_data().to_vec().unwrap(); + assert_eq!(data, vec![3.0, 2.0, 1.0, 3.0, 2.0, 1.0]); + } + + #[test] + fn test_expand_flipped_2d_preserves_negative_stride() { + // [[1, 2], [3, 4]] with axis 0 flipped -> [[3, 4], [1, 2]] + // Shape [2, 2] with strides [-2, 1] (negative on axis 0) + // Expand [2, 2] -> [3, 2, 2] + let tensor = FlexTensor::from_data(TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0], [2, 2])); + let flipped = crate::ops::flip::flip(tensor, &[0]); + assert!(flipped.layout().strides()[0] < 0); + assert_eq!(flipped.layout().strides()[1], 1); + + let expanded = expand(flipped, Shape::new([3, 2, 2])); + assert_eq!(expanded.layout().shape().to_vec(), vec![3, 2, 2]); + // Stride 0 for broadcast, negative stride preserved for axis 1, positive for axis 2 + assert_eq!(expanded.layout().strides()[0], 0); + assert!(expanded.layout().strides()[1] < 0); + assert_eq!(expanded.layout().strides()[2], 1); + + // Verify content + let data: Vec = expanded.into_data().to_vec().unwrap(); + // [[3, 4], [1, 2]] repeated 3 times + assert_eq!( + data, + vec![3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0] + ); + } + + #[test] + fn test_expand_narrowed_preserves_offset() { + // [0, 1, 2, 3, 4] narrowed to [1, 2, 3] with offset 1 + // Expand [3] -> [2, 3] + let tensor = FlexTensor::from_data(TensorData::new(vec![0.0f32, 1.0, 2.0, 3.0, 4.0], [5])); + let narrowed = tensor.narrow(0, 1, 3); + assert_eq!(narrowed.layout().start_offset(), 1); + + let expanded = expand(narrowed, Shape::new([2, 3])); + assert_eq!(expanded.layout().shape().to_vec(), vec![2, 3]); + // Start offset preserved + assert_eq!(expanded.layout().start_offset(), 1); + + // Verify content: [1, 2, 3] repeated twice + let data: Vec = expanded.into_data().to_vec().unwrap(); + assert_eq!(data, vec![1.0, 2.0, 3.0, 1.0, 2.0, 3.0]); + } + + #[test] + #[should_panic(expected = "broadcasting cannot drop dimensions")] + fn test_expand_to_fewer_dims_panics() { + // Shape [2, 3, 4] expanded to target [2, 3] would silently drop + // the trailing dim and produce an invalid layout; must panic. + let tensor = FlexTensor::from_data(TensorData::new(alloc::vec![0.0f32; 24], [2, 3, 4])); + let _ = expand(tensor, Shape::new([2, 3])); + } + + #[test] + #[should_panic(expected = "broadcasting cannot drop dimensions")] + fn test_expand_1d_to_scalar_panics() { + // 1D -> 0D exercises the saturating_sub edge we removed: dim_diff + // used to come out as 0 and the loop would iterate zero times, + // silently producing a 0-dim layout over 1D data. + let tensor = FlexTensor::from_data(TensorData::new(alloc::vec![1.0f32, 2.0, 3.0], [3])); + let _ = expand(tensor, Shape::from(alloc::vec::Vec::::new())); + } + + #[test] + fn test_broadcast_binary_with_flipped() { + // One tensor flipped, broadcast to same shape + // lhs: [1, 2, 3, 4] flipped -> [4, 3, 2, 1], shape [4] + // rhs: [[1], [2]], shape [2, 1] -> broadcast to [2, 4] + // After broadcast: lhs [2, 4], rhs [2, 4] + let lhs = FlexTensor::from_data(TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0], [4])); + let lhs = crate::ops::flip::flip(lhs, &[0]); + assert!(lhs.layout().strides()[0] < 0); + + let rhs = FlexTensor::from_data(TensorData::new(vec![10.0f32, 20.0], [2, 1])); + + let (lhs_bc, rhs_bc) = broadcast_binary(lhs, rhs); + assert_eq!(lhs_bc.layout().shape().to_vec(), vec![2, 4]); + assert_eq!(rhs_bc.layout().shape().to_vec(), vec![2, 4]); + + // lhs should have stride 0 in dim 0 (broadcast), negative stride in dim 1 (from flip) + assert_eq!(lhs_bc.layout().strides()[0], 0); + assert!(lhs_bc.layout().strides()[1] < 0); + + // Verify lhs content: [4, 3, 2, 1] repeated twice + let lhs_data: Vec = lhs_bc.into_data().to_vec().unwrap(); + assert_eq!(lhs_data, vec![4.0, 3.0, 2.0, 1.0, 4.0, 3.0, 2.0, 1.0]); + } +} diff --git a/crates/burn-flex/src/ops/fft.rs b/crates/burn-flex/src/ops/fft.rs new file mode 100644 index 0000000..5f61da7 --- /dev/null +++ b/crates/burn-flex/src/ops/fft.rs @@ -0,0 +1,2168 @@ +//! Real FFT (rfft) via Cooley-Tukey with aggressive optimization. +//! +//! Key optimizations: +//! - Real FFT via complex packing: pack N real values as N/2 complex, +//! do a half-size complex FFT, then unpack using Hermitian symmetry (~2x) +//! - Compile-time twiddle tables via const fn Taylor-series sin/cos +//! - Unrolled small complex FFT kernels for N=2, 4, 8 +//! - Mixed radix-4/radix-2 butterfly stages (halves passes over data) +//! - SIMD-vectorized butterflies via macerator +//! - Rayon parallelism across independent fibers + +use alloc::vec; +use alloc::vec::Vec; +use burn_std::{Bytes, Shape}; + +use crate::layout::{contiguous_strides_usize, slice_base_offset}; +use crate::{FlexTensor, Layout}; + +// ============================================================================ +// Const-evaluable sin/cos via Taylor series (13 terms, ~13 digit accuracy) +// ============================================================================ + +const PI: f64 = core::f64::consts::PI; + +const fn const_sin(x: f64) -> f64 { + let mut x = x; + x = x - ((x / (2.0 * PI)) as i64 as f64) * 2.0 * PI; + if x > PI { + x -= 2.0 * PI; + } else if x < -PI { + x += 2.0 * PI; + } + let x2 = x * x; + let mut term = x; + let mut sum = x; + let mut i = 1u32; + while i <= 12 { + term *= -x2 / ((2 * i) as f64 * (2 * i + 1) as f64); + sum += term; + i += 1; + } + sum +} + +const fn const_cos(x: f64) -> f64 { + const_sin(x + PI / 2.0) +} + +// ============================================================================ +// Compile-time twiddle table +// ============================================================================ + +struct TwiddleTable { + re: [f32; M], + im: [f32; M], + offsets: [usize; 18], + num_stages: usize, +} + +const fn make_twiddle_table() -> TwiddleTable { + let mut re = [0.0f32; M]; + let mut im = [0.0f32; M]; + let mut offsets = [0usize; 18]; + let num_stages = N.trailing_zeros() as usize; + let mut pos = 0usize; + let mut len = 2usize; + let mut stage = 0usize; + while stage < num_stages { + offsets[stage] = pos; + let half = len / 2; + let angle_step = -2.0 * PI / len as f64; + let mut k = 0usize; + while k < half { + let angle = angle_step * k as f64; + re[pos] = const_cos(angle) as f32; + im[pos] = const_sin(angle) as f32; + pos += 1; + k += 1; + } + len <<= 1; + stage += 1; + } + offsets[num_stages] = pos; + TwiddleTable { + re, + im, + offsets, + num_stages, + } +} + +macro_rules! def_twiddle { + ($name:ident, $n:expr) => { + static $name: TwiddleTable<{ $n - 1 }> = make_twiddle_table::<$n, { $n - 1 }>(); + }; +} + +def_twiddle!(TW_2, 2); +def_twiddle!(TW_4, 4); +def_twiddle!(TW_8, 8); +def_twiddle!(TW_16, 16); +def_twiddle!(TW_32, 32); +def_twiddle!(TW_64, 64); +def_twiddle!(TW_128, 128); +def_twiddle!(TW_256, 256); +def_twiddle!(TW_512, 512); +def_twiddle!(TW_1024, 1024); +def_twiddle!(TW_2048, 2048); +def_twiddle!(TW_4096, 4096); +def_twiddle!(TW_8192, 8192); +def_twiddle!(TW_16384, 16384); +def_twiddle!(TW_32768, 32768); +def_twiddle!(TW_65536, 65536); + +enum TwiddleRef { + Static { + re: &'static [f32], + im: &'static [f32], + offsets: &'static [usize], + }, + Owned { + re: Vec, + im: Vec, + offsets: Vec, + }, +} + +impl TwiddleRef { + fn re(&self) -> &[f32] { + match self { + Self::Static { re, .. } => re, + Self::Owned { re, .. } => re, + } + } + fn im(&self) -> &[f32] { + match self { + Self::Static { im, .. } => im, + Self::Owned { im, .. } => im, + } + } + fn offsets(&self) -> &[usize] { + match self { + Self::Static { offsets, .. } => offsets, + Self::Owned { offsets, .. } => offsets, + } + } +} + +fn get_twiddles(n: usize) -> TwiddleRef { + macro_rules! match_static { + ($($size:expr => $table:ident),+ $(,)?) => { + match n { + 0 | 1 => TwiddleRef::Static { re: &[], im: &[], offsets: &[0] }, + $($size => TwiddleRef::Static { + re: &$table.re, im: &$table.im, + offsets: &$table.offsets[..$table.num_stages + 1], + },)+ + _ => { + let (re, im, offsets) = precompute_twiddles_runtime(n); + TwiddleRef::Owned { re, im, offsets } + } + } + }; + } + match_static!( + 2 => TW_2, 4 => TW_4, 8 => TW_8, 16 => TW_16, + 32 => TW_32, 64 => TW_64, 128 => TW_128, 256 => TW_256, + 512 => TW_512, 1024 => TW_1024, 2048 => TW_2048, 4096 => TW_4096, + 8192 => TW_8192, 16384 => TW_16384, 32768 => TW_32768, 65536 => TW_65536, + ) +} + +fn precompute_twiddles_runtime(n: usize) -> (Vec, Vec, Vec) { + let num_stages = n.trailing_zeros() as usize; + let total = n - 1; + let mut re = Vec::with_capacity(total); + let mut im = Vec::with_capacity(total); + let mut offsets = Vec::with_capacity(num_stages + 1); + let mut len = 2; + for _ in 0..num_stages { + offsets.push(re.len()); + let half = len / 2; + let angle_step = -2.0 * core::f64::consts::PI / len as f64; + for k in 0..half { + let angle = angle_step * k as f64; + re.push(const_cos(angle) as f32); + im.push(const_sin(angle) as f32); + } + len <<= 1; + } + offsets.push(re.len()); + (re, im, offsets) +} + +// ============================================================================ +// Bit-reversal permutation +// ============================================================================ + +#[inline] +fn bit_reverse_permute(re: &mut [f32], im: &mut [f32], n: usize) { + let mut j = 0usize; + for i in 1..n { + let mut bit = n >> 1; + while j & bit != 0 { + j ^= bit; + bit >>= 1; + } + j ^= bit; + if i < j { + re.swap(i, j); + im.swap(i, j); + } + } +} + +// ============================================================================ +// Unrolled small complex FFT kernels +// ============================================================================ + +/// Complex FFT of size 2: single butterfly, no twiddles. +#[inline(always)] +fn complex_fft_2(re: &mut [f32], im: &mut [f32]) { + let (r0, r1) = (re[0], re[1]); + let (i0, i1) = (im[0], im[1]); + re[0] = r0 + r1; + re[1] = r0 - r1; + im[0] = i0 + i1; + im[1] = i0 - i1; +} + +/// Complex FFT of size 4: 2 stages, fully unrolled. +/// Twiddle for stage 1, k=1 is W_4^1 = -i. +#[inline(always)] +fn complex_fft_4(re: &mut [f32], im: &mut [f32]) { + // Bit-reversal: swap indices 1 and 2 + re.swap(1, 2); + im.swap(1, 2); + + // Stage 0: two size-2 butterflies + let (r0, r1) = (re[0] + re[1], re[0] - re[1]); + let (i0, i1) = (im[0] + im[1], im[0] - im[1]); + let (r2, r3) = (re[2] + re[3], re[2] - re[3]); + let (i2, i3) = (im[2] + im[3], im[2] - im[3]); + + // Stage 1: size-4 butterfly + // k=0: W=1, butterfly (0,2) + re[0] = r0 + r2; + im[0] = i0 + i2; + re[2] = r0 - r2; + im[2] = i0 - i2; + // k=1: W=-i, butterfly (1,3): -i*(r3+i*i3) = (i3, -r3) + re[1] = r1 + i3; + im[1] = i1 - r3; + re[3] = r1 - i3; + im[3] = i1 + r3; +} + +/// Complex FFT of size 8: 3 stages, fully unrolled. +#[inline(always)] +fn complex_fft_8(re: &mut [f32], im: &mut [f32]) { + // Bit-reversal for n=8: [0,4,2,6,1,5,3,7] + re.swap(1, 4); + im.swap(1, 4); + re.swap(3, 6); + im.swap(3, 6); + + // Stage 0: four size-2 butterflies + macro_rules! butterfly2 { + ($a:expr, $b:expr) => { + let (ra, rb) = (re[$a] + re[$b], re[$a] - re[$b]); + let (ia, ib) = (im[$a] + im[$b], im[$a] - im[$b]); + re[$a] = ra; + re[$b] = rb; + im[$a] = ia; + im[$b] = ib; + }; + } + butterfly2!(0, 1); + butterfly2!(2, 3); + butterfly2!(4, 5); + butterfly2!(6, 7); + + // Stage 1: two size-4 butterflies + // Group [0,1,2,3]: k=0 W=1, k=1 W=-i + { + let (r0, r2) = (re[0] + re[2], re[0] - re[2]); + let (i0, i2) = (im[0] + im[2], im[0] - im[2]); + re[0] = r0; + im[0] = i0; + re[2] = r2; + im[2] = i2; + // k=1: W=-i → (im[3], -re[3]) + let (t_re, t_im) = (im[3], -re[3]); + let (r1a, r1b) = (re[1] + t_re, re[1] - t_re); + let (i1a, i1b) = (im[1] + t_im, im[1] - t_im); + re[1] = r1a; + re[3] = r1b; + im[1] = i1a; + im[3] = i1b; + } + // Group [4,5,6,7]: same pattern + { + let (r4, r6) = (re[4] + re[6], re[4] - re[6]); + let (i4, i6) = (im[4] + im[6], im[4] - im[6]); + re[4] = r4; + im[4] = i4; + re[6] = r6; + im[6] = i6; + let (t_re, t_im) = (im[7], -re[7]); + let (r5a, r5b) = (re[5] + t_re, re[5] - t_re); + let (i5a, i5b) = (im[5] + t_im, im[5] - t_im); + re[5] = r5a; + re[7] = r5b; + im[5] = i5a; + im[7] = i5b; + } + + // Stage 2: one size-8 butterfly + // k=0: W=1 + { + let (a, b) = (re[0] + re[4], re[0] - re[4]); + let (c, d) = (im[0] + im[4], im[0] - im[4]); + re[0] = a; + re[4] = b; + im[0] = c; + im[4] = d; + } + // k=1: W_8^1 = (sqrt2/2, -sqrt2/2) + { + const W: f32 = core::f32::consts::FRAC_1_SQRT_2; // 0.7071... + let t_re = W * re[5] - (-W) * im[5]; // W*re + W*im + let t_im = W * im[5] + (-W) * re[5]; // W*im - W*re + re[5] = re[1] - t_re; + im[5] = im[1] - t_im; + re[1] += t_re; + im[1] += t_im; + } + // k=2: W_8^2 = -i + { + let (t_re, t_im) = (im[6], -re[6]); + re[6] = re[2] - t_re; + im[6] = im[2] - t_im; + re[2] += t_re; + im[2] += t_im; + } + // k=3: W_8^3 = (-sqrt2/2, -sqrt2/2) + { + const W: f32 = core::f32::consts::FRAC_1_SQRT_2; + let t_re = -W * re[7] - (-W) * im[7]; // -W*re + W*im + let t_im = -W * im[7] + (-W) * re[7]; // -W*im - W*re + re[7] = re[3] - t_re; + im[7] = im[3] - t_im; + re[3] += t_re; + im[3] += t_im; + } +} + +// ============================================================================ +// General complex FFT: mixed radix-4/radix-2 with SIMD +// ============================================================================ + +/// Complex FFT of size n (power of 2) using precomputed twiddles. +#[inline] +fn complex_fft(re: &mut [f32], im: &mut [f32], n: usize, tw: &TwiddleRef) { + match n { + 0 | 1 => return, + 2 => { + complex_fft_2(re, im); + return; + } + 4 => { + complex_fft_4(re, im); + return; + } + 8 => { + complex_fft_8(re, im); + return; + } + _ => {} + } + + bit_reverse_permute(re, im, n); + + let tw_re = tw.re(); + let tw_im = tw.im(); + let offsets = tw.offsets(); + let num_stages = offsets.len() - 1; + + // For odd number of stages, do one radix-2 pass first so the + // remaining stages can be processed in radix-4 pairs. + // Stage 0 twiddle is always W_2^0 = 1, so just add/sub. + let start_stage = if num_stages % 2 == 1 { + let mut start = 0; + while start < n { + let (a, b) = (re[start] + re[start + 1], re[start] - re[start + 1]); + let (c, d) = (im[start] + im[start + 1], im[start] - im[start + 1]); + re[start] = a; + re[start + 1] = b; + im[start] = c; + im[start + 1] = d; + start += 2; + } + 1 + } else { + 0 + }; + + #[cfg(feature = "simd")] + { + simd_fft::radix4_simd(re, im, n, tw_re, tw_im, offsets, start_stage, num_stages); + } + #[cfg(not(feature = "simd"))] + { + radix4_scalar(re, im, n, tw_re, tw_im, offsets, start_stage, num_stages); + } +} + +/// Correct DIT radix-4: fuse two radix-2 stages into one pass. +/// +/// For each pair of stages (s, s+1) with quarter = 2^s: +/// 1. Apply W_{2q}^k to x[p1] and x[p3] (inner stage twiddle, same for both) +/// 2. Inner butterflies: a=p0+tw1, b=p0-tw1, c=p2+tw3, d=p2-tw3 +/// 3. Apply W_{4q}^k to c and d; d also gets -i rotation +/// 4. Outer butterflies: p0=a+tc, p2=a-tc, p1=b+(-i*td), p3=b-(-i*td) +#[cfg(not(feature = "simd"))] +#[allow(clippy::too_many_arguments)] +fn radix4_scalar( + re: &mut [f32], + im: &mut [f32], + n: usize, + tw_re: &[f32], + tw_im: &[f32], + offsets: &[usize], + start_stage: usize, + num_stages: usize, +) { + let mut stage = start_stage; + while stage + 1 < num_stages { + let quarter = 1 << stage; + let group_size = quarter << 2; + let tw_off_inner = offsets[stage]; // W_{2q}^k + let tw_off_outer = offsets[stage + 1]; // W_{4q}^k + + let mut group_start = 0; + while group_start < n { + for k in 0..quarter { + let p0 = group_start + k; + let p1 = p0 + quarter; + let p2 = p1 + quarter; + let p3 = p2 + quarter; + + // Inner twiddle: W_{2q}^k applied to p1 and p3 + let wi_re = tw_re[tw_off_inner + k]; + let wi_im = tw_im[tw_off_inner + k]; + let tw1_re = wi_re * re[p1] - wi_im * im[p1]; + let tw1_im = wi_re * im[p1] + wi_im * re[p1]; + let tw3_re = wi_re * re[p3] - wi_im * im[p3]; + let tw3_im = wi_re * im[p3] + wi_im * re[p3]; + + // Inner butterflies + let a_re = re[p0] + tw1_re; + let a_im = im[p0] + tw1_im; + let b_re = re[p0] - tw1_re; + let b_im = im[p0] - tw1_im; + let c_re = re[p2] + tw3_re; + let c_im = im[p2] + tw3_im; + let d_re = re[p2] - tw3_re; + let d_im = im[p2] - tw3_im; + + // Outer twiddle: W_{4q}^k applied to c and d + let wo_re = tw_re[tw_off_outer + k]; + let wo_im = tw_im[tw_off_outer + k]; + let tc_re = wo_re * c_re - wo_im * c_im; + let tc_im = wo_re * c_im + wo_im * c_re; + let td_re = wo_re * d_re - wo_im * d_im; + let td_im = wo_re * d_im + wo_im * d_re; + + // Outer butterflies (-i*(td_re+i*td_im) = (td_im, -td_re)) + re[p0] = a_re + tc_re; + im[p0] = a_im + tc_im; + re[p2] = a_re - tc_re; + im[p2] = a_im - tc_im; + re[p1] = b_re + td_im; + im[p1] = b_im - td_re; + re[p3] = b_re - td_im; + im[p3] = b_im + td_re; + } + group_start += group_size; + } + stage += 2; + } +} + +#[cfg(feature = "simd")] +mod simd_fft { + use macerator::{Simd, vload_unaligned, vstore_unaligned}; + + /// Scalar radix-4 butterfly for the SIMD tail path. + #[allow(clippy::too_many_arguments)] + #[inline(always)] + fn scalar_radix4( + re: &mut [f32], + im: &mut [f32], + p0: usize, + quarter: usize, + tw_re: &[f32], + tw_im: &[f32], + tw_off_inner: usize, + tw_off_outer: usize, + k: usize, + ) { + let p1 = p0 + quarter; + let p2 = p1 + quarter; + let p3 = p2 + quarter; + + let wi_r = tw_re[tw_off_inner + k]; + let wi_i = tw_im[tw_off_inner + k]; + let tw1_re = wi_r * re[p1] - wi_i * im[p1]; + let tw1_im = wi_r * im[p1] + wi_i * re[p1]; + let tw3_re = wi_r * re[p3] - wi_i * im[p3]; + let tw3_im = wi_r * im[p3] + wi_i * re[p3]; + + let a_re = re[p0] + tw1_re; + let a_im = im[p0] + tw1_im; + let b_re = re[p0] - tw1_re; + let b_im = im[p0] - tw1_im; + let c_re = re[p2] + tw3_re; + let c_im = im[p2] + tw3_im; + let d_re = re[p2] - tw3_re; + let d_im = im[p2] - tw3_im; + + let wo_r = tw_re[tw_off_outer + k]; + let wo_i = tw_im[tw_off_outer + k]; + let tc_re = wo_r * c_re - wo_i * c_im; + let tc_im = wo_r * c_im + wo_i * c_re; + let td_re = wo_r * d_re - wo_i * d_im; + let td_im = wo_r * d_im + wo_i * d_re; + + re[p0] = a_re + tc_re; + im[p0] = a_im + tc_im; + re[p2] = a_re - tc_re; + im[p2] = a_im - tc_im; + re[p1] = b_re + td_im; + im[p1] = b_im - td_re; + re[p3] = b_re - td_im; + im[p3] = b_im + td_re; + } + + /// SIMD radix-4 butterfly passes (pairs of radix-2 stages). + #[macerator::with_simd] + #[allow(clippy::too_many_arguments)] + pub fn radix4_simd( + re: &mut [f32], + im: &mut [f32], + n: usize, + tw_re: &[f32], + tw_im: &[f32], + offsets: &[usize], + start_stage: usize, + num_stages: usize, + ) { + let lanes = S::lanes32(); + let mut stage = start_stage; + + while stage + 1 < num_stages { + let quarter = 1 << stage; + let group_size = quarter << 2; + let tw_off_inner = offsets[stage]; + let tw_off_outer = offsets[stage + 1]; + + if quarter >= lanes { + let mut group_start = 0; + while group_start < n { + let mut k = 0; + while k + lanes <= quarter { + unsafe { + // Inner twiddle: W_{2q}^k + let wi_r = + vload_unaligned::(tw_re.as_ptr().add(tw_off_inner + k)); + let wi_i = + vload_unaligned::(tw_im.as_ptr().add(tw_off_inner + k)); + + let p0 = group_start + k; + let p1 = p0 + quarter; + let p2 = p1 + quarter; + let p3 = p2 + quarter; + + let r0 = vload_unaligned::(re.as_ptr().add(p0)); + let i0 = vload_unaligned::(im.as_ptr().add(p0)); + let r1 = vload_unaligned::(re.as_ptr().add(p1)); + let i1 = vload_unaligned::(im.as_ptr().add(p1)); + let r2 = vload_unaligned::(re.as_ptr().add(p2)); + let i2 = vload_unaligned::(im.as_ptr().add(p2)); + let r3 = vload_unaligned::(re.as_ptr().add(p3)); + let i3 = vload_unaligned::(im.as_ptr().add(p3)); + + // Apply inner twiddle to p1 and p3 + let tw1_re = wi_r * r1 - wi_i * i1; + let tw1_im = wi_r * i1 + wi_i * r1; + let tw3_re = wi_r * r3 - wi_i * i3; + let tw3_im = wi_r * i3 + wi_i * r3; + + // Inner butterflies + let a_re = r0 + tw1_re; + let a_im = i0 + tw1_im; + let b_re = r0 - tw1_re; + let b_im = i0 - tw1_im; + let c_re = r2 + tw3_re; + let c_im = i2 + tw3_im; + let d_re = r2 - tw3_re; + let d_im = i2 - tw3_im; + + // Outer twiddle: W_{4q}^k + let wo_r = + vload_unaligned::(tw_re.as_ptr().add(tw_off_outer + k)); + let wo_i = + vload_unaligned::(tw_im.as_ptr().add(tw_off_outer + k)); + let tc_re = wo_r * c_re - wo_i * c_im; + let tc_im = wo_r * c_im + wo_i * c_re; + let td_re = wo_r * d_re - wo_i * d_im; + let td_im = wo_r * d_im + wo_i * d_re; + + // Outer butterflies: -i*(td_re+i*td_im) = (td_im, -td_re) + vstore_unaligned::(re.as_mut_ptr().add(p0), a_re + tc_re); + vstore_unaligned::(im.as_mut_ptr().add(p0), a_im + tc_im); + vstore_unaligned::(re.as_mut_ptr().add(p2), a_re - tc_re); + vstore_unaligned::(im.as_mut_ptr().add(p2), a_im - tc_im); + vstore_unaligned::(re.as_mut_ptr().add(p1), b_re + td_im); + vstore_unaligned::(im.as_mut_ptr().add(p1), b_im - td_re); + vstore_unaligned::(re.as_mut_ptr().add(p3), b_re - td_im); + vstore_unaligned::(im.as_mut_ptr().add(p3), b_im + td_re); + } + k += lanes; + } + while k < quarter { + scalar_radix4( + re, + im, + group_start + k, + quarter, + tw_re, + tw_im, + tw_off_inner, + tw_off_outer, + k, + ); + k += 1; + } + group_start += group_size; + } + } else { + let mut group_start = 0; + while group_start < n { + for k in 0..quarter { + scalar_radix4( + re, + im, + group_start + k, + quarter, + tw_re, + tw_im, + tw_off_inner, + tw_off_outer, + k, + ); + } + group_start += group_size; + } + } + stage += 2; + } + } +} + +// ============================================================================ +// Real FFT unpacking +// ============================================================================ + +/// Unpack N/2-point complex FFT result into N/2+1 real FFT bins. +/// +/// Given Z = FFT(pack(x)), recovers X = FFT(x) using: +/// Xe[k] = (Z[k] + conj(Z[N/2-k])) / 2 +/// Xo[k] = -i * (Z[k] - conj(Z[N/2-k])) / 2 +/// X[k] = Xe[k] + W_N^k * Xo[k] +fn unpack_rfft( + z_re: &[f32], + z_im: &[f32], + half: usize, + unpack_tw_re: &[f32], + unpack_tw_im: &[f32], + out_re: &mut [f32], + out_im: &mut [f32], +) { + // k=0: X[0] = Z_re[0] + Z_im[0] (real) + out_re[0] = z_re[0] + z_im[0]; + out_im[0] = 0.0; + + // k=N/2: X[N/2] = Z_re[0] - Z_im[0] (real) + out_re[half] = z_re[0] - z_im[0]; + out_im[half] = 0.0; + + // k=1..half-1 + for k in 1..half { + let j = half - k; + let (zk_re, zk_im) = (z_re[k], z_im[k]); + let (zj_re, zj_im) = (z_re[j], z_im[j]); + + // Xe = (Z[k] + conj(Z[j])) / 2 + let xe_re = (zk_re + zj_re) * 0.5; + let xe_im = (zk_im - zj_im) * 0.5; + + // Xo = -i * (Z[k] - conj(Z[j])) / 2 + // diff = Z[k] - conj(Z[j]) = (zk_re - zj_re, zk_im + zj_im) + // -i * diff = (diff_im, -diff_re) + let xo_re = (zk_im + zj_im) * 0.5; + let xo_im = (zj_re - zk_re) * 0.5; + + // X[k] = Xe + W * Xo + let wr = unpack_tw_re[k]; + let wi = unpack_tw_im[k]; + + out_re[k] = xe_re + wr * xo_re - wi * xo_im; + out_im[k] = xe_im + wr * xo_im + wi * xo_re; + } +} + +// ============================================================================ +// Tensor helpers +// ============================================================================ + +fn make_tensors_typed( + re: Vec, + im: Vec, + shape: Shape, +) -> (FlexTensor, FlexTensor) { + let dtype = E::dtype(); + let re_t = FlexTensor::new( + Bytes::from_elems(re), + Layout::contiguous(shape.clone()), + dtype, + ); + let im_t = FlexTensor::new(Bytes::from_elems(im), Layout::contiguous(shape), dtype); + (re_t, im_t) +} + +// ============================================================================ +// Top-level rfft: real FFT via complex packing +// ============================================================================ + +/// Process a single fiber: pack real signal as complex, FFT, unpack. +#[allow(clippy::too_many_arguments)] +#[inline] +fn rfft_fiber( + signal: &[f32], + in_stride: usize, + n: usize, + sig_len: usize, + out_re: &mut [f32], + out_im: &mut [f32], + tw_half: &TwiddleRef, + unpack_tw_re: &[f32], + unpack_tw_im: &[f32], + z_re: &mut [f32], + z_im: &mut [f32], +) { + let half = n / 2; + + if n == 1 { + out_re[0] = if sig_len >= 1 { signal[0] } else { 0.0 }; + out_im[0] = 0.0; + return; + } + + if sig_len >= n { + if in_stride == 1 { + for k in 0..half { + z_re[k] = signal[2 * k]; + z_im[k] = signal[2 * k + 1]; + } + } else { + for k in 0..half { + z_re[k] = signal[(2 * k) * in_stride]; + z_im[k] = signal[(2 * k + 1) * in_stride]; + } + } + } else { + for k in 0..half { + let even = 2 * k; + let odd = 2 * k + 1; + z_re[k] = if even < sig_len { + signal[even * in_stride] + } else { + 0.0 + }; + z_im[k] = if odd < sig_len { + signal[odd * in_stride] + } else { + 0.0 + }; + } + } + + complex_fft(z_re, z_im, half, tw_half); + unpack_rfft(z_re, z_im, half, unpack_tw_re, unpack_tw_im, out_re, out_im); +} + +pub fn rfft_f32(tensor: FlexTensor, dim: usize, n: Option) -> (FlexTensor, FlexTensor) { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + assert!( + dim < shape.num_dims(), + "rfft: dim {dim} out of bounds for {}-D tensor", + shape.num_dims() + ); + + let requested_n = n.unwrap_or_else(|| { + let sig_len = shape[dim]; + assert!( + sig_len > 0 && sig_len.is_power_of_two(), + "rfft: dimension size must be a power of 2, got {sig_len}" + ); + sig_len + }); + let fft_size = requested_n.next_power_of_two(); + let sig_len = shape[dim].min(requested_n); + + let n = fft_size; + let out_len = n / 2 + 1; + + let mut out_dims: Vec = shape.as_slice().to_vec(); + out_dims[dim] = out_len; + let out_shape = Shape::from(out_dims); + let total_out = out_shape.num_elements(); + let num_fibers = shape.num_elements() / shape[dim]; + + let data: &[f32] = tensor.storage(); + let in_strides = contiguous_strides_usize(&shape); + let out_strides = contiguous_strides_usize(&out_shape); + + // N=1: each element is its own DFT, no twiddles needed + if n == 1 { + let mut re_out = vec![0.0f32; total_out]; + let im_out = vec![0.0f32; total_out]; + if sig_len >= 1 { + for fiber_idx in 0..num_fibers { + let base = slice_base_offset(fiber_idx, &shape, &in_strides, dim); + let out_base = slice_base_offset(fiber_idx, &out_shape, &out_strides, dim); + re_out[out_base] = data[base]; + } + } + return make_tensors_typed(re_out, im_out, out_shape); + } + + let half = n / 2; + let tw_half = get_twiddles(half); + + let tw_full = get_twiddles(n); + let full_offsets = tw_full.offsets(); + let last_stage_off = if full_offsets.len() >= 2 { + full_offsets[full_offsets.len() - 2] + } else { + 0 + }; + let unpack_tw_re = &tw_full.re()[last_stage_off..]; + let unpack_tw_im = &tw_full.im()[last_stage_off..]; + + let mut re_out = vec![0.0f32; total_out]; + let mut im_out = vec![0.0f32; total_out]; + + let in_stride = in_strides[dim]; + let out_stride = out_strides[dim]; + + #[cfg(feature = "rayon")] + if num_fibers >= 4 && n >= 64 { + use rayon::prelude::*; + + let fiber_results: Vec<(usize, Vec, Vec)> = (0..num_fibers) + .into_par_iter() + .map(|fiber_idx| { + let base_offset = slice_base_offset(fiber_idx, &shape, &in_strides, dim); + let mut z_re = vec![0.0f32; half.max(1)]; + let mut z_im = vec![0.0f32; half.max(1)]; + let mut fiber_re = vec![0.0f32; out_len]; + let mut fiber_im = vec![0.0f32; out_len]; + + rfft_fiber( + &data[base_offset..], + in_stride, + n, + sig_len, + &mut fiber_re, + &mut fiber_im, + &tw_half, + unpack_tw_re, + unpack_tw_im, + &mut z_re, + &mut z_im, + ); + (fiber_idx, fiber_re, fiber_im) + }) + .collect(); + + for (fiber_idx, fiber_re, fiber_im) in fiber_results { + let out_base = slice_base_offset(fiber_idx, &out_shape, &out_strides, dim); + for k in 0..out_len { + re_out[out_base + k * out_stride] = fiber_re[k]; + im_out[out_base + k * out_stride] = fiber_im[k]; + } + } + + let (re, im) = make_tensors_typed(re_out, im_out, out_shape); + return (re, im); + } + + let mut z_re_buf = vec![0.0f32; half.max(1)]; + let mut z_im_buf = vec![0.0f32; half.max(1)]; + let mut fiber_re = vec![0.0f32; out_len]; + let mut fiber_im = vec![0.0f32; out_len]; + + for fiber_idx in 0..num_fibers { + let base_offset = slice_base_offset(fiber_idx, &shape, &in_strides, dim); + let out_base = slice_base_offset(fiber_idx, &out_shape, &out_strides, dim); + + rfft_fiber( + &data[base_offset..], + in_stride, + n, + sig_len, + &mut fiber_re, + &mut fiber_im, + &tw_half, + unpack_tw_re, + unpack_tw_im, + &mut z_re_buf, + &mut z_im_buf, + ); + + for k in 0..out_len { + re_out[out_base + k * out_stride] = fiber_re[k]; + im_out[out_base + k * out_stride] = fiber_im[k]; + } + } + + let (re, im) = make_tensors_typed(re_out, im_out, out_shape); + (re, im) +} + +#[allow(clippy::too_many_arguments)] +fn rfft_fiber_f64( + signal: &[f64], + in_stride: usize, + n: usize, + sig_len: usize, + half: usize, + out_re: &mut [f64], + out_im: &mut [f64], + tw_re: &[f32], + tw_im: &[f32], + tw_offsets: &[usize], + unpack_re: &[f32], + unpack_im: &[f32], + z_re: &mut [f64], + z_im: &mut [f64], +) { + if n == 1 { + out_re[0] = if sig_len >= 1 { signal[0] } else { 0.0 }; + out_im[0] = 0.0; + return; + } + + if sig_len >= n { + for k in 0..half { + z_re[k] = signal[(2 * k) * in_stride]; + z_im[k] = signal[(2 * k + 1) * in_stride]; + } + } else { + for k in 0..half { + let even = 2 * k; + let odd = 2 * k + 1; + z_re[k] = if even < sig_len { + signal[even * in_stride] + } else { + 0.0 + }; + z_im[k] = if odd < sig_len { + signal[odd * in_stride] + } else { + 0.0 + }; + } + } + + fft_f64_inplace(z_re, z_im, half, tw_re, tw_im, tw_offsets); + + out_re[0] = z_re[0] + z_im[0]; + out_im[0] = 0.0; + out_re[half] = z_re[0] - z_im[0]; + out_im[half] = 0.0; + + for k in 1..half { + let j = half - k; + let (zk_re, zk_im) = (z_re[k], z_im[k]); + let (zj_re, zj_im) = (z_re[j], z_im[j]); + + let xe_re = (zk_re + zj_re) * 0.5; + let xe_im = (zk_im - zj_im) * 0.5; + let xo_re = (zk_im + zj_im) * 0.5; + let xo_im = (zj_re - zk_re) * 0.5; + + let wr = unpack_re[k] as f64; + let wi = unpack_im[k] as f64; + + out_re[k] = xe_re + wr * xo_re - wi * xo_im; + out_im[k] = xe_im + wr * xo_im + wi * xo_re; + } +} + +pub fn rfft_f64(tensor: FlexTensor, dim: usize, n: Option) -> (FlexTensor, FlexTensor) { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + assert!( + dim < shape.num_dims(), + "rfft: dim {dim} out of bounds for {}-D tensor", + shape.num_dims() + ); + + let requested_n = n.unwrap_or_else(|| { + let sig_len = shape[dim]; + assert!( + sig_len > 0 && sig_len.is_power_of_two(), + "rfft: dimension size must be a power of 2, got {sig_len}" + ); + sig_len + }); + let fft_size = requested_n.next_power_of_two(); + let sig_len = shape[dim].min(requested_n); + + let n = fft_size; + let out_len = n / 2 + 1; + + let mut out_dims: Vec = shape.as_slice().to_vec(); + out_dims[dim] = out_len; + let out_shape = Shape::from(out_dims); + let total_out = out_shape.num_elements(); + let num_fibers = shape.num_elements() / shape[dim]; + + let data: &[f64] = tensor.storage(); + let in_strides = contiguous_strides_usize(&shape); + let out_strides = contiguous_strides_usize(&out_shape); + let half = n / 2; + + let tw_half = get_twiddles(half); + let tw_full = get_twiddles(n); + let full_offsets = tw_full.offsets(); + let last_stage_off = if full_offsets.len() >= 2 { + full_offsets[full_offsets.len() - 2] + } else { + 0 + }; + let unpack_re = &tw_full.re()[last_stage_off..]; + let unpack_im = &tw_full.im()[last_stage_off..]; + + let mut re_out = vec![0.0f64; total_out]; + let mut im_out = vec![0.0f64; total_out]; + let in_stride = in_strides[dim]; + let out_stride = out_strides[dim]; + + let tw_half_re = tw_half.re(); + let tw_half_im = tw_half.im(); + let tw_half_offsets = tw_half.offsets(); + + #[cfg(feature = "rayon")] + if num_fibers >= 4 && n >= 64 { + use rayon::prelude::*; + + let fiber_results: Vec<(usize, Vec, Vec)> = (0..num_fibers) + .into_par_iter() + .map(|fiber_idx| { + let base_offset = slice_base_offset(fiber_idx, &shape, &in_strides, dim); + let mut z_re = vec![0.0f64; half.max(1)]; + let mut z_im = vec![0.0f64; half.max(1)]; + let mut fiber_re = vec![0.0f64; out_len]; + let mut fiber_im = vec![0.0f64; out_len]; + + rfft_fiber_f64( + &data[base_offset..], + in_stride, + n, + sig_len, + half, + &mut fiber_re, + &mut fiber_im, + tw_half_re, + tw_half_im, + tw_half_offsets, + unpack_re, + unpack_im, + &mut z_re, + &mut z_im, + ); + (fiber_idx, fiber_re, fiber_im) + }) + .collect(); + + for (fiber_idx, fiber_re, fiber_im) in fiber_results { + let out_base = slice_base_offset(fiber_idx, &out_shape, &out_strides, dim); + for k in 0..out_len { + re_out[out_base + k * out_stride] = fiber_re[k]; + im_out[out_base + k * out_stride] = fiber_im[k]; + } + } + + let (re, im) = make_tensors_typed(re_out, im_out, out_shape); + return (re, im); + } + + let mut z_re = vec![0.0f64; half.max(1)]; + let mut z_im = vec![0.0f64; half.max(1)]; + let mut fiber_re = vec![0.0f64; out_len]; + let mut fiber_im = vec![0.0f64; out_len]; + + for fiber_idx in 0..num_fibers { + let base_offset = slice_base_offset(fiber_idx, &shape, &in_strides, dim); + let out_base = slice_base_offset(fiber_idx, &out_shape, &out_strides, dim); + + rfft_fiber_f64( + &data[base_offset..], + in_stride, + n, + sig_len, + half, + &mut fiber_re, + &mut fiber_im, + tw_half_re, + tw_half_im, + tw_half_offsets, + unpack_re, + unpack_im, + &mut z_re, + &mut z_im, + ); + + for k in 0..out_len { + re_out[out_base + k * out_stride] = fiber_re[k]; + im_out[out_base + k * out_stride] = fiber_im[k]; + } + } + + let (re, im) = make_tensors_typed(re_out, im_out, out_shape); + (re, im) +} + +/// f64 complex FFT using f32 twiddle table (widened in inner loop). +/// Twiddle precision is limited to ~7 digits (f32), so output accuracy +/// is below full f64 precision for large N. +fn fft_f64_inplace( + re: &mut [f64], + im: &mut [f64], + n: usize, + tw_re: &[f32], + tw_im: &[f32], + offsets: &[usize], +) { + if n <= 1 { + return; + } + + // Bit-reversal + let mut j = 0usize; + for i in 1..n { + let mut bit = n >> 1; + while j & bit != 0 { + j ^= bit; + bit >>= 1; + } + j ^= bit; + if i < j { + re.swap(i, j); + im.swap(i, j); + } + } + + // Scalar radix-2 passes + let num_stages = offsets.len() - 1; + let mut len = 2; + for &tw_off in &offsets[..num_stages] { + let half = len / 2; + let mut start = 0; + while start < n { + for k in 0..half { + let wr = tw_re[tw_off + k] as f64; + let wi = tw_im[tw_off + k] as f64; + let even = start + k; + let odd = even + half; + let t_re = wr * re[odd] - wi * im[odd]; + let t_im = wr * im[odd] + wi * re[odd]; + re[odd] = re[even] - t_re; + im[odd] = im[even] - t_im; + re[even] += t_re; + im[even] += t_im; + } + start += len; + } + len <<= 1; + } +} + +pub fn rfft_f16(tensor: FlexTensor, dim: usize, n: Option) -> (FlexTensor, FlexTensor) { + use burn_std::f16; + let tensor = super::module::cast_to_f32(tensor, f16::to_f32); + let (re, im) = rfft_f32(tensor, dim, n); + ( + super::module::cast_from_f32(re, f16::from_f32), + super::module::cast_from_f32(im, f16::from_f32), + ) +} + +pub fn rfft_bf16(tensor: FlexTensor, dim: usize, n: Option) -> (FlexTensor, FlexTensor) { + use burn_std::bf16; + let tensor = super::module::cast_to_f32(tensor, bf16::to_f32); + let (re, im) = rfft_f32(tensor, dim, n); + ( + super::module::cast_from_f32(re, bf16::from_f32), + super::module::cast_from_f32(im, bf16::from_f32), + ) +} + +// ============================================================================ +// Inverse real FFT (irfft) +// ============================================================================ + +/// Inverse complex FFT: IFFT(X) = (1/N) * conj(FFT(conj(X))). +/// +/// Conjugates input, runs the forward FFT (with SIMD), conjugates +/// output, and scales by 1/N. +#[inline] +fn inverse_complex_fft(re: &mut [f32], im: &mut [f32], n: usize, tw: &TwiddleRef) { + if n <= 1 { + return; + } + + // Conjugate input + for v in im.iter_mut() { + *v = -*v; + } + + // Forward FFT (mixed radix-4/radix-2, SIMD when available) + complex_fft(re, im, n, tw); + + // Conjugate output and scale by 1/N + let scale = 1.0 / n as f32; + for v in re.iter_mut() { + *v *= scale; + } + for v in im.iter_mut() { + *v = -*v * scale; + } +} + +/// Repack N/2+1 spectrum bins into N/2 complex values Z[k], +/// reversing rfft's unpack step. +/// +/// Z[k] = Xe[k] + i*conj(W)*D[k] where: +/// Xe = (X[k] + conj(X[half-k])) / 2 +/// D = (X[k] - conj(X[half-k])) / 2 +/// W = W_N^k (same twiddle used in rfft unpack) +fn repack_irfft( + x_re: &[f32], + x_im: &[f32], + half: usize, + tw_re: &[f32], + tw_im: &[f32], + z_re: &mut [f32], + z_im: &mut [f32], +) { + // k=0: Z[0] = (X[0] + X[half])/2 + i*(X[0] - X[half])/2 + z_re[0] = (x_re[0] + x_re[half]) * 0.5; + z_im[0] = (x_re[0] - x_re[half]) * 0.5; + + for k in 1..half { + let j = half - k; + let (xk_re, xk_im) = (x_re[k], x_im[k]); + let (xj_re, xj_im) = (x_re[j], x_im[j]); + + // Xe = (X[k] + conj(X[j])) / 2 + let a_re = (xk_re + xj_re) * 0.5; + let a_im = (xk_im - xj_im) * 0.5; + + // D = (X[k] - conj(X[j])) / 2 + let d_re = (xk_re - xj_re) * 0.5; + let d_im = (xk_im + xj_im) * 0.5; + + // i*conj(W)*D where W = (wr, wi), conj(W) = (wr, -wi) + // conj(W)*D = (wr*d_re + wi*d_im, wr*d_im - wi*d_re) + // i*(...) = (-(wr*d_im - wi*d_re), wr*d_re + wi*d_im) + let wr = tw_re[k]; + let wi = tw_im[k]; + + z_re[k] = a_re - wr * d_im + wi * d_re; + z_im[k] = a_im + wr * d_re + wi * d_im; + } +} + +/// Process a single irfft fiber via inverse packing trick. +/// +/// Repacks N/2+1 spectrum bins into N/2 complex values, runs N/2-point +/// inverse complex FFT, then de-interleaves to N real output values. +#[allow(clippy::too_many_arguments)] +#[inline] +fn irfft_fiber( + re_in: &[f32], + im_in: &[f32], + in_stride: usize, + half: usize, + spec_bins: usize, + signal_out: &mut [f32], + out_stride: usize, + tw_half: &TwiddleRef, + unpack_tw_re: &[f32], + unpack_tw_im: &[f32], + z_re: &mut [f32], + z_im: &mut [f32], + spec_re: &mut [f32], + spec_im: &mut [f32], +) { + if spec_bins > half { + for k in 0..=half { + spec_re[k] = re_in[k * in_stride]; + spec_im[k] = im_in[k * in_stride]; + } + } else { + for k in 0..=half { + spec_re[k] = if k < spec_bins { + re_in[k * in_stride] + } else { + 0.0 + }; + spec_im[k] = if k < spec_bins { + im_in[k * in_stride] + } else { + 0.0 + }; + } + } + + repack_irfft( + spec_re, + spec_im, + half, + unpack_tw_re, + unpack_tw_im, + z_re, + z_im, + ); + + // Inverse complex FFT of size N/2 + inverse_complex_fft(z_re, z_im, half, tw_half); + + // De-interleave: signal[2k] = z_re[k], signal[2k+1] = z_im[k] + if out_stride == 1 { + for k in 0..half { + signal_out[2 * k] = z_re[k]; + signal_out[2 * k + 1] = z_im[k]; + } + } else { + for k in 0..half { + signal_out[(2 * k) * out_stride] = z_re[k]; + signal_out[(2 * k + 1) * out_stride] = z_im[k]; + } + } +} + +pub fn irfft_f32( + spectrum_re: FlexTensor, + spectrum_im: FlexTensor, + dim: usize, + n: Option, +) -> FlexTensor { + let spectrum_re = spectrum_re.to_contiguous(); + let spectrum_im = spectrum_im.to_contiguous(); + let shape = spectrum_re.layout().shape().clone(); + assert!( + *spectrum_im.layout().shape() == shape, + "irfft: spectrum_re and spectrum_im shapes must match" + ); + assert!( + dim < shape.num_dims(), + "irfft: dim {dim} out of bounds for {}-D tensor", + shape.num_dims() + ); + let half_plus_1 = shape[dim]; + assert!( + half_plus_1 >= 1, + "irfft: spectrum dimension cannot be empty" + ); + + let spec_bins = half_plus_1; + + let requested_n = n.unwrap_or_else(|| { + let sig_len = (half_plus_1 - 1) * 2; + assert!( + sig_len.is_power_of_two(), + "irfft: reconstructed signal length must be a power of 2, got {sig_len}" + ); + sig_len + }); + let fft_size = requested_n.next_power_of_two(); + + // N=1: single DC bin, output is just the real value along `dim`. + // If caller's spectrum has more bins, take the DC (bin 0) only. + if fft_size <= 1 { + let out = if spectrum_re.layout().shape()[dim] != 1 { + spectrum_re.narrow(dim, 0, 1) + } else { + spectrum_re + }; + return out; + } + + let half = fft_size / 2; + let n = fft_size; + + let mut out_dims: Vec = shape.as_slice().to_vec(); + out_dims[dim] = n; + let out_shape = Shape::from(out_dims); + let total_out = out_shape.num_elements(); + let num_fibers = shape.num_elements() / half_plus_1; + + let re_data: &[f32] = spectrum_re.storage(); + let im_data: &[f32] = spectrum_im.storage(); + let in_strides = contiguous_strides_usize(&shape); + let out_strides = contiguous_strides_usize(&out_shape); + + // Twiddles for N/2-point inverse complex FFT + let tw_half = get_twiddles(half); + + // Unpack twiddles: last stage of size-N table (same as rfft) + let tw_full = get_twiddles(n); + let full_offsets = tw_full.offsets(); + let last_stage_off = if full_offsets.len() >= 2 { + full_offsets[full_offsets.len() - 2] + } else { + 0 + }; + let unpack_tw_re = &tw_full.re()[last_stage_off..]; + let unpack_tw_im = &tw_full.im()[last_stage_off..]; + + let mut signal_out = vec![0.0f32; total_out]; + let in_stride = in_strides[dim]; + let out_stride = out_strides[dim]; + + #[cfg(feature = "rayon")] + if num_fibers >= 4 && n >= 64 { + use rayon::prelude::*; + + let fiber_results: Vec<(usize, Vec)> = (0..num_fibers) + .into_par_iter() + .map(|fiber_idx| { + let re_base = slice_base_offset(fiber_idx, &shape, &in_strides, dim); + let mut z_re = vec![0.0f32; half.max(1)]; + let mut z_im = vec![0.0f32; half.max(1)]; + let mut spec_re = vec![0.0f32; half + 1]; + let mut spec_im = vec![0.0f32; half + 1]; + let mut fiber_out = vec![0.0f32; n]; + + irfft_fiber( + &re_data[re_base..], + &im_data[re_base..], + in_stride, + half, + spec_bins, + &mut fiber_out, + 1, + &tw_half, + unpack_tw_re, + unpack_tw_im, + &mut z_re, + &mut z_im, + &mut spec_re, + &mut spec_im, + ); + (fiber_idx, fiber_out) + }) + .collect(); + + for (fiber_idx, fiber_out) in fiber_results { + let out_base = slice_base_offset(fiber_idx, &out_shape, &out_strides, dim); + for k in 0..n { + signal_out[out_base + k * out_stride] = fiber_out[k]; + } + } + + let result = FlexTensor::new( + Bytes::from_elems(signal_out), + Layout::contiguous(out_shape), + burn_backend::DType::F32, + ); + return if fft_size > requested_n { + result.narrow(dim, 0, requested_n) + } else { + result + }; + } + + let mut z_re = vec![0.0f32; half.max(1)]; + let mut z_im = vec![0.0f32; half.max(1)]; + let mut spec_re = vec![0.0f32; half + 1]; + let mut spec_im = vec![0.0f32; half + 1]; + let mut fiber_out = vec![0.0f32; n]; + + for fiber_idx in 0..num_fibers { + let re_base = slice_base_offset(fiber_idx, &shape, &in_strides, dim); + let out_base = slice_base_offset(fiber_idx, &out_shape, &out_strides, dim); + + irfft_fiber( + &re_data[re_base..], + &im_data[re_base..], + in_stride, + half, + spec_bins, + &mut fiber_out, + 1, + &tw_half, + unpack_tw_re, + unpack_tw_im, + &mut z_re, + &mut z_im, + &mut spec_re, + &mut spec_im, + ); + + for k in 0..n { + signal_out[out_base + k * out_stride] = fiber_out[k]; + } + } + + let result = FlexTensor::new( + Bytes::from_elems(signal_out), + Layout::contiguous(out_shape), + burn_backend::DType::F32, + ); + if fft_size > requested_n { + result.narrow(dim, 0, requested_n) + } else { + result + } +} + +pub fn irfft_f64( + spectrum_re: FlexTensor, + spectrum_im: FlexTensor, + dim: usize, + n: Option, +) -> FlexTensor { + use burn_backend::DType; + match spectrum_re.dtype() { + DType::F64 => { + let re_f32 = super::module::cast_to_f32::(spectrum_re, |v| v as f32); + let im_f32 = super::module::cast_to_f32::(spectrum_im, |v| v as f32); + let result = irfft_f32(re_f32, im_f32, dim, n); + super::module::cast_from_f32::(result, |v| v as f64) + } + _ => irfft_f32(spectrum_re, spectrum_im, dim, n), + } +} + +pub fn irfft_f16( + spectrum_re: FlexTensor, + spectrum_im: FlexTensor, + dim: usize, + n: Option, +) -> FlexTensor { + use burn_std::f16; + let re = super::module::cast_to_f32(spectrum_re, f16::to_f32); + let im = super::module::cast_to_f32(spectrum_im, f16::to_f32); + let result = irfft_f32(re, im, dim, n); + super::module::cast_from_f32(result, f16::from_f32) +} + +pub fn irfft_bf16( + spectrum_re: FlexTensor, + spectrum_im: FlexTensor, + dim: usize, + n: Option, +) -> FlexTensor { + use burn_std::bf16; + let re = super::module::cast_to_f32(spectrum_re, bf16::to_f32); + let im = super::module::cast_to_f32(spectrum_im, bf16::to_f32); + let result = irfft_f32(re, im, dim, n); + super::module::cast_from_f32(result, bf16::from_f32) +} + +// Tests kept here exercise flex-specific internals: the FFT kernels +// (`rfft_f32`/`_f64`/`_f16`, `irfft_*`, `complex_fft`, `inverse_complex_fft`) +// across sizes that span the radix-4 and complex packing paths (N=1, 2, 4, 8, +// 256, 1024, 4096), f16/f64 dtype handling, twiddle accuracy, Parseval's +// theorem on synthetic inputs, and a reference cross-check against realfft. +#[cfg(test)] +mod tests { + use super::*; + use burn_backend::{DType, TensorData, Tolerance}; + + fn make_f32(data: Vec, shape: Vec) -> FlexTensor { + FlexTensor::from_data(TensorData::new(data, shape)) + } + + fn make_f64(data: Vec, shape: Vec) -> FlexTensor { + FlexTensor::from_data(TensorData::new(data, shape)) + } + + fn assert_approx(tensor: FlexTensor, expected: &[f32], tol: f32) { + let shape = tensor.layout().shape().as_slice().to_vec(); + tensor.into_data().assert_approx_eq::( + &TensorData::new(expected.to_vec(), shape), + Tolerance::absolute(tol), + ); + } + + fn assert_approx_f64(tensor: FlexTensor, expected: &[f64], tol: f64) { + tensor + .into_data() + .assert_approx_eq::(&TensorData::from(expected), Tolerance::absolute(tol)); + } + + // ---- N=1 ---- + + #[test] + fn rfft_n1() { + let signal = make_f32(vec![5.0], vec![1]); + let (re, im) = rfft_f32(signal, 0, None); + assert_approx(re, &[5.0], 1e-6); + assert_approx(im, &[0.0], 1e-6); + } + + // ---- N=2 ---- + + #[test] + fn rfft_n2() { + let signal = make_f32(vec![1.0, -1.0], vec![2]); + let (re, im) = rfft_f32(signal, 0, None); + assert_approx(re, &[0.0, 2.0], 1e-6); + assert_approx(im, &[0.0, 0.0], 1e-6); + } + + // ---- N=4: known DFT of [1,0,0,0] = [1,1,1] (all real) ---- + + #[test] + fn rfft_n4_impulse() { + let signal = make_f32(vec![1.0, 0.0, 0.0, 0.0], vec![4]); + let (re, im) = rfft_f32(signal, 0, None); + assert_approx(re, &[1.0, 1.0, 1.0], 1e-6); + assert_approx(im, &[0.0, 0.0, 0.0], 1e-6); + } + + // ---- N=4: constant signal [1,1,1,1] -> DC only ---- + + #[test] + fn rfft_n4_constant() { + let signal = make_f32(vec![1.0, 1.0, 1.0, 1.0], vec![4]); + let (re, im) = rfft_f32(signal, 0, None); + assert_approx(re, &[4.0, 0.0, 0.0], 1e-6); + assert_approx(im, &[0.0, 0.0, 0.0], 1e-6); + } + + // ---- N=4: zeros ---- + + #[test] + fn rfft_n4_zeros() { + let signal = make_f32(vec![0.0; 4], vec![4]); + let (re, im) = rfft_f32(signal, 0, None); + assert_approx(re, &[0.0, 0.0, 0.0], 1e-6); + assert_approx(im, &[0.0, 0.0, 0.0], 1e-6); + } + + // ---- N=8 ---- + + #[test] + fn rfft_n8_impulse() { + let mut signal = vec![0.0f32; 8]; + signal[0] = 1.0; + let (re, im) = rfft_f32(make_f32(signal, vec![8]), 0, None); + // DFT of impulse is all 1s + assert_approx(re, &[1.0, 1.0, 1.0, 1.0, 1.0], 1e-6); + assert_approx(im, &[0.0, 0.0, 0.0, 0.0, 0.0], 1e-6); + } + + #[test] + fn rfft_n8_cosine() { + // cos(2*pi*k/8) for k=0..7 -> energy at bin 1 + let signal: Vec = (0..8) + .map(|k| (2.0 * std::f32::consts::PI * k as f32 / 8.0).cos()) + .collect(); + let (re, im) = rfft_f32(make_f32(signal, vec![8]), 0, None); + // Bin 1 should have amplitude 4 (real), rest ~0 + assert_approx(re, &[0.0, 4.0, 0.0, 0.0, 0.0], 1e-4); + assert_approx(im, &[0.0, 0.0, 0.0, 0.0, 0.0], 1e-4); + } + + // ---- Larger size: N=256 ---- + + #[test] + fn rfft_n256_impulse() { + let mut signal = vec![0.0f32; 256]; + signal[0] = 1.0; + let (re, im) = rfft_f32(make_f32(signal, vec![256]), 0, None); + let re_data = re.into_data(); + let im_data = im.into_data(); + let re_vals = re_data.as_slice::().unwrap(); + let im_vals = im_data.as_slice::().unwrap(); + assert_eq!(re_vals.len(), 129); + for &v in re_vals { + assert!((v - 1.0).abs() < 1e-5, "re bin should be 1.0, got {v}"); + } + for &v in im_vals { + assert!(v.abs() < 1e-5, "im bin should be 0.0, got {v}"); + } + } + + // ---- Multi-dimensional: FFT along dim 1 ---- + + #[test] + fn rfft_2d_dim1() { + // 2 rows, each of length 4: impulse and constant + let data = vec![ + 1.0, 0.0, 0.0, 0.0, // row 0: impulse + 1.0, 1.0, 1.0, 1.0, // row 1: constant + ]; + let signal = make_f32(data, vec![2, 4]); + let (re, im) = rfft_f32(signal, 1, None); + // Shape should be [2, 3] + let re_data = re.into_data(); + let im_data = im.into_data(); + let re_vals = re_data.as_slice::().unwrap(); + let im_vals = im_data.as_slice::().unwrap(); + assert_eq!(re_vals.len(), 6); // 2 * 3 + // Row 0 (impulse): [1, 1, 1] + assert!((re_vals[0] - 1.0).abs() < 1e-5); + assert!((re_vals[1] - 1.0).abs() < 1e-5); + assert!((re_vals[2] - 1.0).abs() < 1e-5); + // Row 1 (constant): [4, 0, 0] + assert!((re_vals[3] - 4.0).abs() < 1e-5); + assert!((re_vals[4]).abs() < 1e-5); + assert!((re_vals[5]).abs() < 1e-5); + // All imaginary should be ~0 + for &v in im_vals { + assert!(v.abs() < 1e-5); + } + } + + // ---- FFT along dim 0 ---- + + #[test] + fn rfft_2d_dim0() { + // 4 rows, 2 cols: impulse in each column + let data = vec![1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]; + let signal = make_f32(data, vec![4, 2]); + let (re, _im) = rfft_f32(signal, 0, None); + // Shape should be [3, 2] + let re_data = re.into_data(); + let re_vals = re_data.as_slice::().unwrap(); + assert_eq!(re_vals.len(), 6); + // Each column is an impulse -> all bins = 1 + for &v in re_vals { + assert!((v - 1.0).abs() < 1e-5, "expected 1.0, got {v}"); + } + } + + // ---- f64 dtype ---- + + #[test] + fn rfft_f64_n4_impulse() { + let signal = make_f64(vec![1.0, 0.0, 0.0, 0.0], vec![4]); + let (re, im) = rfft_f64(signal, 0, None); + assert_approx_f64(re, &[1.0, 1.0, 1.0], 1e-10); + assert_approx_f64(im, &[0.0, 0.0, 0.0], 1e-10); + } + + #[test] + fn rfft_f64_n8_cosine() { + let signal: Vec = (0..8) + .map(|k| (2.0 * std::f64::consts::PI * k as f64 / 8.0).cos()) + .collect(); + let (re, im) = rfft_f64(make_f64(signal, vec![8]), 0, None); + assert_approx_f64(re, &[0.0, 4.0, 0.0, 0.0, 0.0], 1e-6); + assert_approx_f64(im, &[0.0, 0.0, 0.0, 0.0, 0.0], 1e-6); + } + + // ---- f16 dtype ---- + + #[test] + fn rfft_f16_n4_impulse() { + use burn_std::f16; + let f16_data = vec![ + f16::from_f32(1.0), + f16::from_f32(0.0), + f16::from_f32(0.0), + f16::from_f32(0.0), + ]; + let signal = FlexTensor::new( + Bytes::from_elems(f16_data), + Layout::contiguous(Shape::from(vec![4])), + DType::F16, + ); + let (re, _im) = rfft_f16(signal, 0, None); + // Verify via round-trip to f32 + let re_f32 = super::super::module::cast_to_f32(re, f16::to_f32); + let re_data = re_f32.into_data(); + let re_vals = re_data.as_slice::().unwrap(); + assert_eq!(re_vals.len(), 3); + for &v in re_vals { + assert!((v - 1.0).abs() < 0.01, "expected ~1.0, got {v}"); + } + } + + // ---- Const twiddle accuracy ---- + + #[test] + fn const_sin_cos_accuracy() { + let test_angles = [0.0, 0.1, 0.5, 1.0, 2.0, 3.0, -1.0, -3.0, 6.0]; + for &angle in &test_angles { + let cs = const_sin(angle); + let cc = const_cos(angle); + let rs = angle.sin(); + let rc = angle.cos(); + assert!( + (cs - rs).abs() < 1e-12, + "const_sin({angle}) = {cs}, expected {rs}" + ); + assert!( + (cc - rc).abs() < 1e-12, + "const_cos({angle}) = {cc}, expected {rc}" + ); + } + } + + // ---- N=1024 round-trip with known property: Parseval's theorem ---- + // Sum of |x|^2 = (1/N) * Sum of |X|^2 + + #[test] + fn rfft_n1024_parseval() { + let n = 1024; + let signal: Vec = (0..n).map(|i| (i as f32 * 0.37).sin()).collect(); + let time_energy: f64 = signal.iter().map(|&x| (x as f64) * (x as f64)).sum(); + + let (re, im) = rfft_f32(make_f32(signal, vec![n]), 0, None); + let re_data = re.into_data(); + let im_data = im.into_data(); + let re_vals = re_data.as_slice::().unwrap(); + let im_vals = im_data.as_slice::().unwrap(); + + // Frequency energy: DC and Nyquist count once, others count double + let out_len = n / 2 + 1; + let mut freq_energy = 0.0f64; + for k in 0..out_len { + let mag2 = (re_vals[k] as f64).powi(2) + (im_vals[k] as f64).powi(2); + if k == 0 || k == n / 2 { + freq_energy += mag2; + } else { + freq_energy += 2.0 * mag2; + } + } + freq_energy /= n as f64; + + let rel_err = (freq_energy - time_energy).abs() / time_energy; + assert!( + rel_err < 1e-4, + "Parseval's theorem violated: time={time_energy}, freq={freq_energy}, rel_err={rel_err}" + ); + } + + // ---- irfft tests ---- + + #[test] + fn irfft_roundtrip_n4() { + let signal = make_f32(vec![1.0, 2.0, 3.0, 4.0], vec![4]); + let (re, im) = rfft_f32(signal.clone(), 0, None); + let reconstructed = irfft_f32(re, im, 0, None); + assert_approx(reconstructed, &[1.0, 2.0, 3.0, 4.0], 1e-5); + } + + #[test] + fn irfft_roundtrip_n8() { + let data: Vec = (0..8).map(|i| (i as f32 * 0.3).sin()).collect(); + let signal = make_f32(data.clone(), vec![8]); + let (re, im) = rfft_f32(signal, 0, None); + let reconstructed = irfft_f32(re, im, 0, None); + assert_approx(reconstructed, &data, 1e-5); + } + + #[test] + fn rfft_vs_realfft() { + // Verify our rfft matches realfft (rustfft-backed) for non-trivial input + // at sizes that exercise radix-4 (n>=16) and the complex packing trick. + use realfft::RealFftPlanner; + + let mut planner = RealFftPlanner::::new(); + + for &n in &[4, 8, 16, 32, 64, 256, 1024, 4096] { + let data: Vec = (0..n).map(|i| (i as f32 * 0.37).sin() + 0.5).collect(); + + // Our rfft + let signal = make_f32(data.clone(), vec![n]); + let (re_out, im_out) = rfft_f32(signal, 0, None); + let re_data = re_out.into_data(); + let im_data = im_out.into_data(); + let our_re = re_data.as_slice::().unwrap(); + let our_im = im_data.as_slice::().unwrap(); + + // Reference: realfft + let r2c = planner.plan_fft_forward(n); + let mut input = data.clone(); + let mut spectrum = r2c.make_output_vec(); + r2c.process(&mut input, &mut spectrum).unwrap(); + + let out_len = n / 2 + 1; + assert_eq!(our_re.len(), out_len); + assert_eq!(spectrum.len(), out_len); + + let max_re_err = our_re + .iter() + .zip(spectrum.iter()) + .map(|(&a, b)| (a - b.re).abs()) + .fold(0.0f32, f32::max); + let max_im_err = our_im + .iter() + .zip(spectrum.iter()) + .map(|(&a, b)| (a - b.im).abs()) + .fold(0.0f32, f32::max); + assert!( + max_re_err < 1e-3 && max_im_err < 1e-3, + "rfft vs realfft mismatch at n={n}: max_re_err={max_re_err}, max_im_err={max_im_err}" + ); + } + } + + #[test] + fn irfft_vs_realfft() { + // Verify our irfft matches realfft's inverse for non-trivial spectra. + use realfft::RealFftPlanner; + + let mut planner = RealFftPlanner::::new(); + + for &n in &[4, 8, 16, 32, 64, 256, 1024, 4096] { + // Generate a spectrum via realfft forward + let r2c = planner.plan_fft_forward(n); + let c2r = planner.plan_fft_inverse(n); + let data: Vec = (0..n).map(|i| (i as f32 * 0.37).sin() + 0.5).collect(); + let mut input = data.clone(); + let mut spectrum = r2c.make_output_vec(); + r2c.process(&mut input, &mut spectrum).unwrap(); + + // Our irfft + let out_len = n / 2 + 1; + let spec_re: Vec = spectrum.iter().map(|c| c.re).collect(); + let spec_im: Vec = spectrum.iter().map(|c| c.im).collect(); + let re_tensor = make_f32(spec_re, vec![out_len]); + let im_tensor = make_f32(spec_im, vec![out_len]); + let our_result = irfft_f32(re_tensor, im_tensor, 0, None); + let our_data = our_result.into_data(); + let our_vals = our_data.as_slice::().unwrap(); + + // Reference: realfft inverse (note: realfft doesn't normalize, so scale) + let mut spec_copy = spectrum.clone(); + let mut ref_output = c2r.make_output_vec(); + c2r.process(&mut spec_copy, &mut ref_output).unwrap(); + let scale = 1.0 / n as f32; + let ref_scaled: Vec = ref_output.iter().map(|&v| v * scale).collect(); + + let max_err = our_vals + .iter() + .zip(ref_scaled.iter()) + .map(|(&a, &b)| (a - b).abs()) + .fold(0.0f32, f32::max); + assert!( + max_err < 1e-3, + "irfft vs realfft mismatch at n={n}: max_err={max_err}" + ); + } + } + + #[test] + fn forward_complex_fft_impulse() { + // DFT of impulse [1,0,0,...,0] should be all 1s + for &n in &[4, 8, 16, 32, 64] { + let tw = get_twiddles(n); + let mut re = vec![0.0f32; n]; + let mut im = vec![0.0f32; n]; + re[0] = 1.0; + complex_fft(&mut re, &mut im, n, &tw); + let max_re_err = re.iter().map(|&v| (v - 1.0).abs()).fold(0.0f32, f32::max); + let max_im_err = im.iter().map(|&v| v.abs()).fold(0.0f32, f32::max); + assert!( + max_re_err < 1e-5 && max_im_err < 1e-5, + "forward FFT impulse n={n}: max_re_err={max_re_err}, max_im_err={max_im_err}" + ); + } + } + + #[test] + fn inverse_complex_fft_roundtrip() { + for &n in &[4, 8, 16, 32, 64, 256] { + let tw = get_twiddles(n); + let mut re: Vec = (0..n).map(|i| (i as f32 * 0.3).sin()).collect(); + let mut im = vec![0.0f32; n]; + let orig_re = re.clone(); + + complex_fft(&mut re, &mut im, n, &tw); + inverse_complex_fft(&mut re, &mut im, n, &tw); + + let max_err = re + .iter() + .zip(orig_re.iter()) + .map(|(&got, &expected)| (got - expected).abs()) + .fold(0.0f32, f32::max); + assert!( + max_err < 1e-4, + "inverse_complex_fft roundtrip n={n}: max error {max_err}" + ); + } + } + + #[test] + fn irfft_roundtrip_n256() { + let data: Vec = (0..256).map(|i| (i as f32 * 0.1).cos()).collect(); + let signal = make_f32(data.clone(), vec![256]); + let (re, im) = rfft_f32(signal, 0, None); + let reconstructed = irfft_f32(re, im, 0, None); + let result = reconstructed.into_data(); + let vals = result.as_slice::().unwrap(); + let max_err = vals + .iter() + .zip(data.iter()) + .map(|(&got, &expected)| (got - expected).abs()) + .fold(0.0f32, f32::max); + assert!( + max_err < 5e-3, + "irfft_roundtrip_n256: max error {max_err} exceeds tolerance" + ); + } + + #[test] + fn irfft_roundtrip_2d_dim1() { + let data = vec![ + 1.0, 2.0, 3.0, 4.0, // row 0 + 5.0, 6.0, 7.0, 8.0, // row 1 + ]; + let signal = make_f32(data.clone(), vec![2, 4]); + let (re, im) = rfft_f32(signal, 1, None); + let reconstructed = irfft_f32(re, im, 1, None); + assert_approx(reconstructed, &data, 1e-5); + } + + #[test] + fn irfft_roundtrip_2d_dim0() { + let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; + let signal = make_f32(data.clone(), vec![4, 2]); + let (re, im) = rfft_f32(signal, 0, None); + let reconstructed = irfft_f32(re, im, 0, None); + assert_approx(reconstructed, &data, 1e-5); + } + + #[test] + fn irfft_known_spectrum() { + // DC=4, all others zero -> constant signal [1,1,1,1] + let re = make_f32(vec![4.0, 0.0, 0.0], vec![3]); + let im = make_f32(vec![0.0, 0.0, 0.0], vec![3]); + let signal = irfft_f32(re, im, 0, None); + assert_approx(signal, &[1.0, 1.0, 1.0, 1.0], 1e-5); + } + + #[test] + fn irfft_f64_roundtrip() { + // irfft_f64 truncates to f32 internally, so tolerance is f32-level + let data: Vec = (0..8).map(|i| (i as f64 * 0.3).sin()).collect(); + let signal = make_f64(data.clone(), vec![8]); + let (re, im) = rfft_f64(signal, 0, None); + let reconstructed = irfft_f64(re, im, 0, None); + assert_approx_f64(reconstructed, &data, 1e-5); + } + + // Coverage for the n=Some(pow2) path on flex. + + #[test] + fn rfft_n_larger_than_signal_zero_pads() { + // n=8 zero-pads the signal; output has 8/2+1 = 5 bins. + let signal = make_f32(vec![1.0, 0.0, 0.0, 0.0], vec![4]); + let (re, im) = rfft_f32(signal, 0, Some(8)); + assert_eq!(re.layout().shape()[0], 5); + assert_eq!(im.layout().shape()[0], 5); + let re_vals = re.into_data().as_slice::().unwrap().to_vec(); + for (k, v) in re_vals.iter().enumerate() { + assert!( + (v - 1.0).abs() < 1e-5, + "impulse DFT re[{k}] should be 1.0, got {v}" + ); + } + } + + #[test] + fn rfft_n_smaller_than_signal_truncates_first() { + // Signal length 8, n=4 -> truncate to 4, compute 4-point DFT of [1,0,0,0]. + let signal = make_f32(vec![1.0, 0.0, 0.0, 0.0, 99.0, 99.0, 99.0, 99.0], vec![8]); + let (re, _im) = rfft_f32(signal, 0, Some(4)); + assert_eq!(re.layout().shape()[0], 3); + let re_vals = re.into_data().as_slice::().unwrap().to_vec(); + for v in &re_vals { + assert!((v - 1.0).abs() < 1e-5, "expected 1.0, got {v}"); + } + } + + #[test] + fn rfft_irfft_roundtrip_with_pow2_n() { + let data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; + let signal = make_f32(data.clone(), vec![8]); + let (re, im) = rfft_f32(signal, 0, Some(8)); + let reconstructed = irfft_f32(re, im, 0, Some(8)); + assert_approx(reconstructed, &data, 1e-4); + } + + #[test] + fn rfft_f64_with_pow2_n_and_truncation() { + // Signal length 8, n=4. Output has 4/2+1 = 3 bins. + let data: Vec = (0..8).map(|i| (i as f64 * 0.3).sin()).collect(); + let signal = make_f64(data, vec![8]); + let (re, im) = rfft_f64(signal, 0, Some(4)); + assert_eq!(re.layout().shape()[0], 3); + assert_eq!(im.layout().shape()[0], 3); + } + + #[test] + fn rfft_vs_realfft_with_pow2_n_and_padding() { + use realfft::RealFftPlanner; + + let mut planner = RealFftPlanner::::new(); + // Signal shorter than n (pow2); backend zero-pads before the FFT. + for &(sig_len, n) in &[(3usize, 4usize), (5, 8), (6, 8), (9, 16)] { + let data: Vec = (0..sig_len) + .map(|i| (i as f32 * 0.41).cos() - 0.2) + .collect(); + let mut padded = data.clone(); + padded.resize(n, 0.0); + + let r2c = planner.plan_fft_forward(n); + let mut input = padded; + let mut ref_spec = r2c.make_output_vec(); + r2c.process(&mut input, &mut ref_spec).unwrap(); + + let signal = make_f32(data, vec![sig_len]); + let (re, im) = rfft_f32(signal, 0, Some(n)); + let re_v = re.into_data().as_slice::().unwrap().to_vec(); + let im_v = im.into_data().as_slice::().unwrap().to_vec(); + + assert_eq!(re_v.len(), n / 2 + 1); + for (k, refc) in ref_spec.iter().enumerate() { + let err_re = (re_v[k] - refc.re).abs(); + let err_im = (im_v[k] - refc.im).abs(); + assert!( + err_re < 1e-3 && err_im < 1e-3, + "sig_len={sig_len} n={n} bin={k}: got ({}, {}), ref ({}, {})", + re_v[k], + im_v[k], + refc.re, + refc.im + ); + } + } + } +} diff --git a/crates/burn-flex/src/ops/flip.rs b/crates/burn-flex/src/ops/flip.rs new file mode 100644 index 0000000..a423b9a --- /dev/null +++ b/crates/burn-flex/src/ops/flip.rs @@ -0,0 +1,43 @@ +//! Flip operation for reversing tensor elements along axes. +//! +//! With signed strides, flip is a zero-copy operation that simply negates +//! the stride and adjusts the start offset for each flipped axis. + +use crate::FlexTensor; + +/// Flip tensor elements along specified axes. +/// +/// This is a zero-copy operation using negative strides. +pub fn flip(tensor: FlexTensor, axes: &[usize]) -> FlexTensor { + if axes.is_empty() { + return tensor; + } + + let new_layout = tensor.layout().flip(axes); + tensor.with_layout(new_layout) +} + +// Tests kept here exercise flex-specific behavior: flip is a zero-copy +// stride-only operation in the flex backend, and the test below verifies +// the underlying buffer pointer is shared across the flip. Correctness +// tests for flip along various axes live in +// crates/burn-backend-tests/tests/tensor/{float,int,bool}/ops/flip.rs and +// run against every backend. +#[cfg(test)] +mod tests { + use super::*; + use burn_backend::TensorData; + + #[test] + fn test_flip_is_zero_copy() { + // Verify flip doesn't copy data by checking it shares the same underlying storage + let tensor = FlexTensor::from_data(TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0], [4])); + let tensor_ptr = tensor.bytes().as_ptr(); + let flipped = flip(tensor, &[0]); + let flipped_ptr = flipped.bytes().as_ptr(); + assert_eq!( + tensor_ptr, flipped_ptr, + "flip should share underlying storage" + ); + } +} diff --git a/crates/burn-flex/src/ops/float.rs b/crates/burn-flex/src/ops/float.rs new file mode 100644 index 0000000..36e6339 --- /dev/null +++ b/crates/burn-flex/src/ops/float.rs @@ -0,0 +1,1220 @@ +//! Float tensor operations for the Flex backend. + +use alloc::vec; +use alloc::vec::Vec; +use burn_backend::{ + DType, Distribution, ExecutionError, FloatDType, Scalar, TensorData, TensorMetadata, + ops::{FloatTensorOps, GridSampleOptions, IntTensorOps}, + tensor::{BoolTensor, Device, FloatTensor, IntTensor}, +}; +use burn_std::{Bytes, IntDType, Shape, Slice, bf16, f16}; +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float; + +use crate::Layout; +use num_traits::ToPrimitive; + +use crate::ops::binary::{BinaryOp, binary_op, scalar_op}; +use crate::ops::matmul; +use crate::ops::unary; +use crate::{Flex, FlexTensor}; + +impl FloatTensorOps for Flex { + fn float_from_data(data: TensorData, _device: &Device) -> FloatTensor { + FlexTensor::from_data(data) + } + + fn float_random( + shape: Shape, + distribution: Distribution, + _device: &Device, + dtype: FloatDType, + ) -> FloatTensor { + let mut seed = crate::backend::SEED.lock().unwrap(); + let mut rng = seed.take().unwrap_or_else(crate::backend::get_seeded_rng); + let data = match dtype { + FloatDType::F64 => TensorData::random::(shape, distribution, &mut rng), + FloatDType::F32 | FloatDType::Flex32 => { + TensorData::random::(shape, distribution, &mut rng) + } + FloatDType::F16 => TensorData::random::(shape, distribution, &mut rng), + FloatDType::BF16 => TensorData::random::(shape, distribution, &mut rng), + }; + *seed = Some(rng); + FlexTensor::from_data(data) + } + + async fn float_into_data(tensor: FloatTensor) -> Result { + Ok(tensor.into_data()) + } + + fn float_to_device(tensor: FloatTensor, _device: &Device) -> FloatTensor { + // CPU backend: no-op, tensors are always on CPU + tensor + } + + fn float_detach(tensor: FloatTensor) -> FloatTensor { + tensor + } + + fn float_into_int(tensor: FloatTensor, out_dtype: burn_std::IntDType) -> IntTensor { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let src = tensor.dtype(); + let out_dt = DType::from(out_dtype); + + // Read source floats as f64 (lossless for f32/f16/bf16). + macro_rules! read_floats { + (|$x:ident| $conv:expr) => { + match src { + DType::F32 => tensor + .storage::() + .iter() + .map(|v| { + let $x = *v as f64; + $conv + }) + .collect(), + DType::F64 => tensor + .storage::() + .iter() + .map(|v| { + let $x = *v; + $conv + }) + .collect(), + DType::F16 => tensor + .storage::() + .iter() + .map(|v| { + let $x = f32::from(*v) as f64; + $conv + }) + .collect(), + DType::BF16 => tensor + .storage::() + .iter() + .map(|v| { + let $x = f32::from(*v) as f64; + $conv + }) + .collect(), + _ => panic!("float_into_int: unsupported source dtype {:?}", src), + } + }; + } + + macro_rules! convert { + ($int_ty:ty) => {{ + let data: Vec<$int_ty> = read_floats!(|x| x as $int_ty); + FlexTensor::new(Bytes::from_elems(data), Layout::contiguous(shape), out_dt) + }}; + } + + match out_dtype { + IntDType::I64 => convert!(i64), + IntDType::I32 => convert!(i32), + IntDType::I16 => convert!(i16), + IntDType::I8 => convert!(i8), + IntDType::U64 => convert!(u64), + IntDType::U32 => convert!(u32), + IntDType::U16 => convert!(u16), + IntDType::U8 => convert!(u8), + } + } + + fn float_empty(shape: Shape, _device: &Device, dtype: FloatDType) -> FloatTensor { + FlexTensor::empty(shape, dtype.into()) + } + + fn float_add(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_op(lhs, rhs, |a, b| a + b, |a, b| a + b, Some(BinaryOp::Add)) + } + + fn float_add_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let rhs_val = rhs.to_f64().unwrap(); + scalar_op(lhs, rhs_val, |a, b| a + b, |a, b| a + b) + } + + fn float_sub(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_op(lhs, rhs, |a, b| a - b, |a, b| a - b, Some(BinaryOp::Sub)) + } + + fn float_sub_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let rhs_val = rhs.to_f64().unwrap(); + scalar_op(lhs, rhs_val, |a, b| a - b, |a, b| a - b) + } + + fn float_mul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_op(lhs, rhs, |a, b| a * b, |a, b| a * b, Some(BinaryOp::Mul)) + } + + fn float_mul_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let rhs_val = rhs.to_f64().unwrap(); + scalar_op(lhs, rhs_val, |a, b| a * b, |a, b| a * b) + } + + fn float_div(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_op(lhs, rhs, |a, b| a / b, |a, b| a / b, Some(BinaryOp::Div)) + } + + fn float_div_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let rhs_val = rhs.to_f64().unwrap(); + scalar_op(lhs, rhs_val, |a, b| a / b, |a, b| a / b) + } + + fn float_remainder(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + // Python/PyTorch-style remainder: result has same sign as divisor + binary_op( + lhs, + rhs, + |a, b| ((a % b) + b) % b, + |a, b| ((a % b) + b) % b, + None, + ) + } + + fn float_remainder_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let rhs_val = rhs.to_f64().unwrap(); + // Python/PyTorch-style remainder: result has same sign as divisor + scalar_op( + lhs, + rhs_val, + |a, b| ((a % b) + b) % b, + |a, b| ((a % b) + b) % b, + ) + } + + fn float_matmul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + matmul::matmul(lhs, rhs) + } + + fn float_cross( + lhs: FloatTensor, + rhs: FloatTensor, + dim: usize, + ) -> FloatTensor { + let shape = lhs.layout().shape(); + let ndims = shape.num_dims(); + assert_eq!( + shape[dim], 3, + "cross product requires dimension {} to have size 3, got {}", + dim, shape[dim] + ); + + // Helper to create slices that select index `idx` along `dim` + let make_slices = |idx: usize| -> alloc::vec::Vec { + (0..ndims) + .map(|d| { + if d == dim { + Slice::new(idx as isize, Some((idx + 1) as isize), 1) + } else { + Slice::new(0, None, 1) + } + }) + .collect() + }; + + // Extract components along the dimension + // a = [a0, a1, a2], b = [b0, b1, b2] + let a0 = Self::float_slice(lhs.clone(), &make_slices(0)); + let a1 = Self::float_slice(lhs.clone(), &make_slices(1)); + let a2 = Self::float_slice(lhs, &make_slices(2)); + + let b0 = Self::float_slice(rhs.clone(), &make_slices(0)); + let b1 = Self::float_slice(rhs.clone(), &make_slices(1)); + let b2 = Self::float_slice(rhs, &make_slices(2)); + + // Cross product: c = a × b + // c0 = a1*b2 - a2*b1 + // c1 = a2*b0 - a0*b2 + // c2 = a0*b1 - a1*b0 + let c0 = Self::float_sub( + Self::float_mul(a1.clone(), b2.clone()), + Self::float_mul(a2.clone(), b1.clone()), + ); + let c1 = Self::float_sub( + Self::float_mul(a2, b0.clone()), + Self::float_mul(a0.clone(), b2), + ); + let c2 = Self::float_sub(Self::float_mul(a0, b1), Self::float_mul(a1, b0)); + + // Concatenate along the dimension + Self::float_cat(vec![c0, c1, c2], dim) + } + + fn float_recip(tensor: FloatTensor) -> FloatTensor { + unary::recip(tensor) + } + + fn float_swap_dims(tensor: FloatTensor, dim1: usize, dim2: usize) -> FloatTensor { + tensor.transpose(dim1, dim2) + } + + fn float_permute(tensor: FloatTensor, axes: &[usize]) -> FloatTensor { + tensor.permute(axes) + } + + fn float_flip(tensor: FloatTensor, axes: &[usize]) -> FloatTensor { + crate::ops::flip::flip(tensor, axes) + } + + fn float_cat(tensors: Vec>, dim: usize) -> FloatTensor { + crate::ops::cat::cat(tensors, dim) + } + + fn float_reshape(tensor: FloatTensor, shape: Shape) -> FloatTensor { + tensor.reshape(shape) + } + + fn float_gather( + dim: usize, + tensor: FloatTensor, + indices: IntTensor, + ) -> FloatTensor { + match tensor.dtype() { + DType::F32 => crate::ops::gather_scatter::gather::(tensor, dim, indices), + DType::F64 => crate::ops::gather_scatter::gather::(tensor, dim, indices), + DType::F16 => crate::ops::gather_scatter::gather::(tensor, dim, indices), + DType::BF16 => crate::ops::gather_scatter::gather::(tensor, dim, indices), + _ => panic!("float_gather: unsupported dtype {:?}", tensor.dtype()), + } + } + + fn float_scatter_add( + dim: usize, + tensor: FloatTensor, + indices: IntTensor, + value: FloatTensor, + ) -> FloatTensor { + match tensor.dtype() { + DType::F32 => { + crate::ops::gather_scatter::scatter_add::(tensor, dim, indices, value) + } + DType::F64 => { + crate::ops::gather_scatter::scatter_add::(tensor, dim, indices, value) + } + DType::F16 => { + crate::ops::gather_scatter::scatter_add::(tensor, dim, indices, value) + } + DType::BF16 => { + crate::ops::gather_scatter::scatter_add::(tensor, dim, indices, value) + } + _ => panic!("float_scatter_add: unsupported dtype {:?}", tensor.dtype()), + } + } + + fn float_scatter_nd( + data: FloatTensor, + indices: IntTensor, + values: FloatTensor, + reduction: burn_backend::tensor::IndexingUpdateOp, + ) -> FloatTensor { + match data.dtype() { + DType::F32 => { + crate::ops::gather_scatter::scatter_nd::(data, indices, values, reduction) + } + DType::F64 => { + crate::ops::gather_scatter::scatter_nd::(data, indices, values, reduction) + } + DType::F16 => { + crate::ops::gather_scatter::scatter_nd::(data, indices, values, reduction) + } + DType::BF16 => { + crate::ops::gather_scatter::scatter_nd::(data, indices, values, reduction) + } + _ => panic!("float_scatter_nd: unsupported dtype {:?}", data.dtype()), + } + } + + fn float_gather_nd(data: FloatTensor, indices: IntTensor) -> FloatTensor { + match data.dtype() { + DType::F32 => crate::ops::gather_scatter::gather_nd::(data, indices), + DType::F64 => crate::ops::gather_scatter::gather_nd::(data, indices), + DType::F16 => crate::ops::gather_scatter::gather_nd::(data, indices), + DType::BF16 => crate::ops::gather_scatter::gather_nd::(data, indices), + _ => panic!("float_gather_nd: unsupported dtype {:?}", data.dtype()), + } + } + + fn float_select( + tensor: FloatTensor, + dim: usize, + indices: IntTensor, + ) -> FloatTensor { + match tensor.dtype() { + DType::F32 => crate::ops::gather_scatter::select::(tensor, dim, indices), + DType::F64 => crate::ops::gather_scatter::select::(tensor, dim, indices), + DType::F16 => crate::ops::gather_scatter::select::(tensor, dim, indices), + DType::BF16 => crate::ops::gather_scatter::select::(tensor, dim, indices), + _ => panic!("float_select: unsupported dtype {:?}", tensor.dtype()), + } + } + + fn float_select_add( + tensor: FloatTensor, + dim: usize, + indices: IntTensor, + value: FloatTensor, + ) -> FloatTensor { + match tensor.dtype() { + DType::F32 => { + crate::ops::gather_scatter::select_add::(tensor, dim, indices, value) + } + DType::F64 => { + crate::ops::gather_scatter::select_add::(tensor, dim, indices, value) + } + DType::F16 => { + crate::ops::gather_scatter::select_add::(tensor, dim, indices, value) + } + DType::BF16 => { + crate::ops::gather_scatter::select_add::(tensor, dim, indices, value) + } + _ => panic!("float_select_add: unsupported dtype {:?}", tensor.dtype()), + } + } + + fn float_slice(tensor: FloatTensor, slices: &[Slice]) -> FloatTensor { + crate::ops::slice::slice(tensor, slices) + } + + fn float_slice_assign( + tensor: FloatTensor, + slices: &[Slice], + value: FloatTensor, + ) -> FloatTensor { + crate::ops::slice::slice_assign(tensor, slices, value) + } + + fn float_mask_where( + tensor: FloatTensor, + mask: BoolTensor, + value: FloatTensor, + ) -> FloatTensor { + match tensor.dtype() { + DType::F32 => crate::ops::mask::mask_where_f32(tensor, mask, value), + DType::F64 => crate::ops::mask::mask_where_f64(tensor, mask, value), + DType::F16 => crate::ops::mask::mask_where_f16(tensor, mask, value), + DType::BF16 => crate::ops::mask::mask_where_bf16(tensor, mask, value), + dtype => panic!("float_mask_where: unsupported dtype {:?}", dtype), + } + } + + fn float_mask_fill( + tensor: FloatTensor, + mask: BoolTensor, + value: Scalar, + ) -> FloatTensor { + match tensor.dtype() { + DType::F32 => crate::ops::mask::mask_fill_f32(tensor, mask, value.to_f32().unwrap()), + DType::F64 => crate::ops::mask::mask_fill_f64(tensor, mask, value.to_f64().unwrap()), + DType::F16 => crate::ops::mask::mask_fill_f16( + tensor, + mask, + f16::from_f64(value.to_f64().unwrap()), + ), + DType::BF16 => crate::ops::mask::mask_fill_bf16( + tensor, + mask, + bf16::from_f64(value.to_f64().unwrap()), + ), + dtype => panic!("float_mask_fill: unsupported dtype {:?}", dtype), + } + } + + fn float_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::equal(lhs, rhs, out_dtype) + } + + fn float_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::equal_elem(lhs, rhs.to_f64().unwrap(), out_dtype) + } + + fn float_greater( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::greater(lhs, rhs, out_dtype) + } + + fn float_greater_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::greater_elem(lhs, rhs.to_f64().unwrap(), out_dtype) + } + + fn float_greater_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::greater_equal(lhs, rhs, out_dtype) + } + + fn float_greater_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::greater_equal_elem(lhs, rhs.to_f64().unwrap(), out_dtype) + } + + fn float_lower( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::lower(lhs, rhs, out_dtype) + } + + fn float_lower_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::lower_elem(lhs, rhs.to_f64().unwrap(), out_dtype) + } + + fn float_lower_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::lower_equal(lhs, rhs, out_dtype) + } + + fn float_lower_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::lower_equal_elem(lhs, rhs.to_f64().unwrap(), out_dtype) + } + + fn float_not_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::not_equal(lhs, rhs, out_dtype) + } + + fn float_not_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::not_equal_elem(lhs, rhs.to_f64().unwrap(), out_dtype) + } + + fn float_neg(tensor: FloatTensor) -> FloatTensor { + unary::unary_op(tensor, |x: f32| -x, |x: f64| -x) + } + + fn float_clamp(tensor: FloatTensor, min: Scalar, max: Scalar) -> FloatTensor { + let min32 = min.to_f32().unwrap(); + let max32 = max.to_f32().unwrap(); + let min64 = min.to_f64().unwrap(); + let max64 = max.to_f64().unwrap(); + unary::unary_op( + tensor, + move |x: f32| x.clamp(min32, max32), + move |x: f64| x.clamp(min64, max64), + ) + } + + fn float_clamp_min(tensor: FloatTensor, min: Scalar) -> FloatTensor { + let min32 = min.to_f32().unwrap(); + let min64 = min.to_f64().unwrap(); + unary::unary_op( + tensor, + move |x: f32| x.max(min32), + move |x: f64| x.max(min64), + ) + } + + fn float_clamp_max(tensor: FloatTensor, max: Scalar) -> FloatTensor { + let max32 = max.to_f32().unwrap(); + let max64 = max.to_f64().unwrap(); + unary::unary_op( + tensor, + move |x: f32| x.min(max32), + move |x: f64| x.min(max64), + ) + } + + fn float_sign(tensor: FloatTensor) -> FloatTensor { + unary::unary_op( + tensor, + |x: f32| { + if x.is_nan() { + x + } else if x > 0.0 { + 1.0 + } else if x < 0.0 { + -1.0 + } else { + 0.0 + } + }, + |x: f64| { + if x.is_nan() { + x + } else if x > 0.0 { + 1.0 + } else if x < 0.0 { + -1.0 + } else { + 0.0 + } + }, + ) + } + + fn float_mean(tensor: FloatTensor) -> FloatTensor { + crate::ops::reduce::mean(tensor) + } + + fn float_max(tensor: FloatTensor) -> FloatTensor { + crate::ops::reduce::max(tensor) + } + + fn float_max_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + crate::ops::reduce::max_dim(tensor, dim) + } + + fn float_min(tensor: FloatTensor) -> FloatTensor { + crate::ops::reduce::min(tensor) + } + + fn float_min_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + crate::ops::reduce::min_dim(tensor, dim) + } + + fn float_max_dim_with_indices( + tensor: FloatTensor, + dim: usize, + indices_dtype: burn_std::IntDType, + ) -> (FloatTensor, IntTensor) { + let (values, indices) = crate::ops::reduce::max_dim_with_indices(tensor, dim); + if indices.dtype() != DType::from(indices_dtype) { + (values, Flex::int_cast(indices, indices_dtype)) + } else { + (values, indices) + } + } + + fn float_min_dim_with_indices( + tensor: FloatTensor, + dim: usize, + indices_dtype: burn_std::IntDType, + ) -> (FloatTensor, IntTensor) { + let (values, indices) = crate::ops::reduce::min_dim_with_indices(tensor, dim); + if indices.dtype() != DType::from(indices_dtype) { + (values, Flex::int_cast(indices, indices_dtype)) + } else { + (values, indices) + } + } + + fn float_any(tensor: FloatTensor, out_dtype: burn_std::BoolDType) -> BoolTensor { + crate::ops::comparison::any_float(tensor, out_dtype) + } + + fn float_any_dim( + tensor: FloatTensor, + dim: usize, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::any_float_dim(tensor, dim, out_dtype) + } + + fn float_all(tensor: FloatTensor, out_dtype: burn_std::BoolDType) -> BoolTensor { + crate::ops::comparison::all_float(tensor, out_dtype) + } + + fn float_all_dim( + tensor: FloatTensor, + dim: usize, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::all_float_dim(tensor, dim, out_dtype) + } + + fn float_sum(tensor: FloatTensor) -> FloatTensor { + crate::ops::reduce::sum(tensor) + } + + fn float_sum_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + crate::ops::reduce::sum_dim(tensor, dim) + } + + fn float_mean_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + crate::ops::reduce::mean_dim(tensor, dim) + } + + fn float_prod(tensor: FloatTensor) -> FloatTensor { + crate::ops::reduce::prod(tensor) + } + + fn float_prod_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + crate::ops::reduce::prod_dim(tensor, dim) + } + + fn float_cumsum(tensor: FloatTensor, dim: usize) -> FloatTensor { + match tensor.dtype() { + DType::F32 => crate::ops::cumulative::cumsum_f32(tensor, dim), + DType::F64 => crate::ops::cumulative::cumsum_f64(tensor, dim), + DType::F16 => { + crate::ops::cumulative::cumsum_half(tensor, dim, f16::to_f32, f16::from_f32) + } + DType::BF16 => { + crate::ops::cumulative::cumsum_half(tensor, dim, bf16::to_f32, bf16::from_f32) + } + _ => panic!("float_cumsum: unsupported dtype {:?}", tensor.dtype()), + } + } + + fn float_cumprod(tensor: FloatTensor, dim: usize) -> FloatTensor { + match tensor.dtype() { + DType::F32 => crate::ops::cumulative::cumprod_f32(tensor, dim), + DType::F64 => crate::ops::cumulative::cumprod_f64(tensor, dim), + DType::F16 => { + crate::ops::cumulative::cumprod_half(tensor, dim, f16::to_f32, f16::from_f32) + } + DType::BF16 => { + crate::ops::cumulative::cumprod_half(tensor, dim, bf16::to_f32, bf16::from_f32) + } + _ => panic!("float_cumprod: unsupported dtype {:?}", tensor.dtype()), + } + } + + fn float_cummin(tensor: FloatTensor, dim: usize) -> FloatTensor { + match tensor.dtype() { + DType::F32 => crate::ops::cumulative::cummin_f32(tensor, dim), + DType::F64 => crate::ops::cumulative::cummin_f64(tensor, dim), + DType::F16 => { + crate::ops::cumulative::cummin_half(tensor, dim, f16::to_f32, f16::from_f32) + } + DType::BF16 => { + crate::ops::cumulative::cummin_half(tensor, dim, bf16::to_f32, bf16::from_f32) + } + _ => panic!("float_cummin: unsupported dtype {:?}", tensor.dtype()), + } + } + + fn float_cummax(tensor: FloatTensor, dim: usize) -> FloatTensor { + match tensor.dtype() { + DType::F32 => crate::ops::cumulative::cummax_f32(tensor, dim), + DType::F64 => crate::ops::cumulative::cummax_f64(tensor, dim), + DType::F16 => { + crate::ops::cumulative::cummax_half(tensor, dim, f16::to_f32, f16::from_f32) + } + DType::BF16 => { + crate::ops::cumulative::cummax_half(tensor, dim, bf16::to_f32, bf16::from_f32) + } + _ => panic!("float_cummax: unsupported dtype {:?}", tensor.dtype()), + } + } + + fn float_cast(tensor: FloatTensor, dtype: FloatDType) -> FloatTensor { + use crate::Layout; + use burn_std::{Bytes, bf16, f16}; + + let src_dtype = tensor.dtype(); + let target_dtype = DType::from(dtype); + + // No-op if already the same dtype + if src_dtype == target_dtype { + return tensor; + } + + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + + // Convert to f64 intermediate, then to target + let f64_values: Vec = match src_dtype { + DType::F32 => { + let src: &[f32] = tensor.storage(); + src.iter().map(|&v| v as f64).collect() + } + DType::F64 => { + let src: &[f64] = tensor.storage(); + src.to_vec() + } + DType::F16 => { + let src: &[f16] = tensor.storage(); + src.iter().map(|&v| v.to_f32() as f64).collect() + } + DType::BF16 => { + let src: &[bf16] = tensor.storage(); + src.iter().map(|&v| v.to_f32() as f64).collect() + } + _ => panic!("float_cast: unsupported source dtype {:?}", src_dtype), + }; + + // Convert from f64 to target dtype + match target_dtype { + DType::F32 => { + let result: Vec = f64_values.iter().map(|&v| v as f32).collect(); + let bytes = Bytes::from_elems(result); + FlexTensor::new(bytes, Layout::contiguous(shape), DType::F32) + } + DType::F64 => { + let bytes = Bytes::from_elems(f64_values); + FlexTensor::new(bytes, Layout::contiguous(shape), DType::F64) + } + DType::F16 => { + let result: Vec = f64_values.iter().map(|&v| f16::from_f64(v)).collect(); + let bytes = Bytes::from_elems(result); + FlexTensor::new(bytes, Layout::contiguous(shape), DType::F16) + } + DType::BF16 => { + let result: Vec = f64_values.iter().map(|&v| bf16::from_f64(v)).collect(); + let bytes = Bytes::from_elems(result); + FlexTensor::new(bytes, Layout::contiguous(shape), DType::BF16) + } + _ => panic!("float_cast: unsupported target dtype {:?}", target_dtype), + } + } + + fn float_exp(tensor: FloatTensor) -> FloatTensor { + unary::exp(tensor) + } + + fn float_log(tensor: FloatTensor) -> FloatTensor { + unary::log(tensor) + } + + fn float_log1p(tensor: FloatTensor) -> FloatTensor { + unary::log1p(tensor) + } + + fn float_powf(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_op(lhs, rhs, |a: f32, b| a.powf(b), |a: f64, b| a.powf(b), None) + } + + fn float_powf_scalar_impl(tensor: FloatTensor, value: Scalar) -> FloatTensor { + let exp = value.to_f64().unwrap(); + scalar_op(tensor, exp, |a: f32, b| a.powf(b), |a: f64, b| a.powf(b)) + } + + fn float_sqrt(tensor: FloatTensor) -> FloatTensor { + unary::sqrt(tensor) + } + + fn float_abs(tensor: FloatTensor) -> FloatTensor { + unary::abs(tensor) + } + + fn float_cos(tensor: FloatTensor) -> FloatTensor { + unary::cos(tensor) + } + + fn float_sin(tensor: FloatTensor) -> FloatTensor { + unary::sin(tensor) + } + + fn float_tan(tensor: FloatTensor) -> FloatTensor { + unary::tan(tensor) + } + + fn float_cosh(tensor: FloatTensor) -> FloatTensor { + unary::cosh(tensor) + } + + fn float_sinh(tensor: FloatTensor) -> FloatTensor { + unary::sinh(tensor) + } + + fn float_tanh(tensor: FloatTensor) -> FloatTensor { + unary::tanh(tensor) + } + + fn float_acos(tensor: FloatTensor) -> FloatTensor { + unary::acos(tensor) + } + + fn float_acosh(tensor: FloatTensor) -> FloatTensor { + unary::acosh(tensor) + } + + fn float_asin(tensor: FloatTensor) -> FloatTensor { + unary::asin(tensor) + } + + fn float_asinh(tensor: FloatTensor) -> FloatTensor { + unary::asinh(tensor) + } + + fn float_atan(tensor: FloatTensor) -> FloatTensor { + unary::atan(tensor) + } + + fn float_atanh(tensor: FloatTensor) -> FloatTensor { + unary::atanh(tensor) + } + + fn float_atan2(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_op( + lhs, + rhs, + |a: f32, b| a.atan2(b), + |a: f64, b| a.atan2(b), + None, + ) + } + + fn float_round(tensor: FloatTensor) -> FloatTensor { + unary::round(tensor) + } + + fn float_floor(tensor: FloatTensor) -> FloatTensor { + unary::floor(tensor) + } + + fn float_ceil(tensor: FloatTensor) -> FloatTensor { + unary::ceil(tensor) + } + + fn float_trunc(tensor: FloatTensor) -> FloatTensor { + unary::trunc(tensor) + } + + fn float_erf(tensor: FloatTensor) -> FloatTensor { + unary::erf(tensor) + } + + fn float_argmax( + tensor: FloatTensor, + dim: usize, + out_dtype: burn_std::IntDType, + ) -> IntTensor { + let result = crate::ops::reduce::argmax(tensor, dim); + if result.dtype() != DType::from(out_dtype) { + Flex::int_cast(result, out_dtype) + } else { + result + } + } + + fn float_argtopk( + _tensor: FloatTensor, + _dim: usize, + _k: usize, + _out_dtype: burn_std::IntDType, + ) -> IntTensor { + unimplemented!("float_argtopk not implemented for flex") + } + + fn float_argmin( + tensor: FloatTensor, + dim: usize, + out_dtype: burn_std::IntDType, + ) -> IntTensor { + let result = crate::ops::reduce::argmin(tensor, dim); + if result.dtype() != DType::from(out_dtype) { + Flex::int_cast(result, out_dtype) + } else { + result + } + } + + fn float_expand(tensor: FloatTensor, shape: Shape) -> FloatTensor { + crate::ops::expand::expand(tensor, shape) + } + + fn float_unfold( + tensor: FloatTensor, + dim: usize, + size: usize, + step: usize, + ) -> FloatTensor { + // unfold is now type-agnostic (zero-copy strided view) + crate::ops::unfold::unfold(tensor, dim, size, step) + } + + fn float_grid_sample_2d( + tensor: FloatTensor, + grid: FloatTensor, + options: GridSampleOptions, + ) -> FloatTensor { + crate::ops::grid_sample::grid_sample_2d(tensor, grid, options) + } + + fn float_zeros(shape: Shape, _device: &Device, dtype: FloatDType) -> FloatTensor { + FlexTensor::zeros(shape, dtype.into()) + } + + fn float_ones(shape: Shape, _device: &Device, dtype: FloatDType) -> FloatTensor { + let dt: burn_backend::DType = dtype.into(); + match dt { + DType::F32 => FlexTensor::filled_typed(shape, dt, 1.0f32), + DType::F64 => FlexTensor::filled_typed(shape, dt, 1.0f64), + DType::F16 => FlexTensor::filled_typed(shape, dt, f16::ONE), + DType::BF16 => FlexTensor::filled_typed(shape, dt, bf16::ONE), + _ => unreachable!(), + } + } + + fn float_full( + shape: Shape, + fill_value: Scalar, + _device: &Device, + dtype: FloatDType, + ) -> FloatTensor { + let dt: burn_backend::DType = dtype.into(); + match dt { + DType::F32 => FlexTensor::filled_typed(shape, dt, fill_value.to_f32().unwrap()), + DType::F64 => FlexTensor::filled_typed(shape, dt, fill_value.to_f64().unwrap()), + DType::F16 => { + FlexTensor::filled_typed(shape, dt, f16::from_f32(fill_value.to_f32().unwrap())) + } + DType::BF16 => { + FlexTensor::filled_typed(shape, dt, bf16::from_f32(fill_value.to_f32().unwrap())) + } + _ => unreachable!(), + } + } + + fn float_transpose(tensor: FloatTensor) -> FloatTensor { + let ndims = tensor.layout().num_dims(); + if ndims < 2 { + return tensor; + } + tensor.transpose(ndims - 2, ndims - 1) + } + + fn float_repeat_dim(tensor: FloatTensor, dim: usize, times: usize) -> FloatTensor { + crate::ops::repeat_dim::repeat_dim(tensor, dim, times) + } + + fn float_sort(tensor: FloatTensor, dim: usize, descending: bool) -> FloatTensor { + crate::ops::sort::sort(tensor, dim, descending) + } + + fn float_sort_with_indices( + tensor: FloatTensor, + dim: usize, + descending: bool, + indices_dtype: burn_std::IntDType, + ) -> (FloatTensor, IntTensor) { + let (values, indices) = crate::ops::sort::sort_with_indices(tensor, dim, descending); + let indices = if indices.dtype() != DType::from(indices_dtype) { + Flex::int_cast(indices, indices_dtype) + } else { + indices + }; + (values, indices) + } + + fn float_argsort( + tensor: FloatTensor, + dim: usize, + descending: bool, + out_dtype: burn_std::IntDType, + ) -> IntTensor { + let indices = crate::ops::sort::argsort(tensor, dim, descending); + if indices.dtype() != DType::from(out_dtype) { + Flex::int_cast(indices, out_dtype) + } else { + indices + } + } + + fn float_powi(lhs: FloatTensor, rhs: IntTensor) -> FloatTensor { + let dtype = lhs.dtype(); + Self::float_powf(lhs, Flex::int_into_float(rhs, dtype.into())) + } + + fn float_powi_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + match rhs.to_i64().unwrap() { + 0 => Self::float_ones(lhs.shape(), &Default::default(), lhs.dtype().into()), + 1 => lhs, + 2 => Self::float_mul(lhs.clone(), lhs), + -1 => Self::float_recip(lhs), + -2 => Self::float_recip(Self::float_mul(lhs.clone(), lhs)), + _ => Self::float_powf_scalar_impl(lhs, rhs), + } + } + + fn float_powf_scalar(tensor: FloatTensor, value: Scalar) -> FloatTensor { + if let Some(exp) = value.try_as_integer() { + Self::float_powi_scalar(tensor, exp) + } else { + Self::float_powf_scalar_impl(tensor, value) + } + } + + fn float_max_abs(tensor: FloatTensor) -> FloatTensor { + let abs = unary::abs(tensor); + crate::ops::reduce::max(abs) + } + + fn float_max_abs_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + let abs = unary::abs(tensor); + crate::ops::reduce::max_dim(abs, dim) + } + + fn float_is_nan(tensor: FloatTensor, out_dtype: burn_std::BoolDType) -> BoolTensor { + unary::float_predicate(tensor, out_dtype, |x: f32| x.is_nan(), |x: f64| x.is_nan()) + } + + fn float_is_inf(tensor: FloatTensor, out_dtype: burn_std::BoolDType) -> BoolTensor { + unary::float_predicate( + tensor, + out_dtype, + |x: f32| x.is_infinite(), + |x: f64| x.is_infinite(), + ) + } + + fn float_hypot(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_op( + lhs, + rhs, + |a: f32, b| a.hypot(b), + |a: f64, b| a.hypot(b), + None, + ) + } +} + +// Tests kept here exercise flex-specific behavior: direct `Flex::` +// backend-op calls with explicit IntDType/FloatDType to pin dtype storage +// selection (U8/I32/I64, F16/F64). Plain arithmetic, math, cast, cross, +// unfold, and random smoke tests have been dropped in favor of the +// equivalent coverage in burn-backend-tests, which exercises every backend. +// When adding new tests, keep them here only if they probe flex dtype +// storage or flex internals; otherwise add them to +// crates/burn-backend-tests/tests/tensor/float/ops/. +#[cfg(test)] +mod tests { + use burn_backend::TensorData; + + use crate::Flex; + + #[test] + fn test_float_into_int_i32() { + use burn_backend::ops::FloatTensorOps; + use burn_std::IntDType; + + let t = crate::FlexTensor::from_data(TensorData::from([1.5f32, -2.7, 0.0, 255.9])); + let result = Flex::float_into_int(t, IntDType::I32); + assert_eq!(result.dtype(), burn_backend::DType::I32); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![1, -2, 0, 255]); + } + + #[test] + fn test_float_into_int_u8() { + use burn_backend::ops::FloatTensorOps; + use burn_std::IntDType; + + let t = crate::FlexTensor::from_data(TensorData::from([0.0f32, 1.9, 127.5, 255.0])); + let result = Flex::float_into_int(t, IntDType::U8); + assert_eq!(result.dtype(), burn_backend::DType::U8); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![0, 1, 127, 255]); + } + + #[test] + fn test_float_argmax_i32_out_dtype() { + use burn_backend::ops::FloatTensorOps; + use burn_std::IntDType; + + let t = crate::FlexTensor::from_data(TensorData::from([[1.0f32, 3.0, 2.0]])); + let result = Flex::float_argmax(t, 1, IntDType::I32); + assert_eq!(result.dtype(), burn_backend::DType::I32); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![1]); + } + + #[test] + fn test_float_argmin_i32_out_dtype() { + use burn_backend::ops::FloatTensorOps; + use burn_std::IntDType; + + let t = crate::FlexTensor::from_data(TensorData::from([[3.0f32, 1.0, 2.0]])); + let result = Flex::float_argmin(t, 1, IntDType::I32); + assert_eq!(result.dtype(), burn_backend::DType::I32); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![1]); + } + + #[test] + fn test_float_argmax_i64_out_dtype() { + use burn_backend::ops::FloatTensorOps; + use burn_std::IntDType; + + let t = crate::FlexTensor::from_data(TensorData::from([[1.0f32, 3.0, 2.0]])); + let result = Flex::float_argmax(t, 1, IntDType::I64); + assert_eq!(result.dtype(), burn_backend::DType::I64); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![1]); + } + + #[test] + fn test_float_max_dim_with_indices_i32() { + use burn_backend::ops::FloatTensorOps; + use burn_std::IntDType; + + let t = crate::FlexTensor::from_data(TensorData::from([[1.0f32, 5.0], [3.0, 2.0]])); + let (values, indices) = Flex::float_max_dim_with_indices(t, 1, IntDType::I32); + assert_eq!(indices.dtype(), burn_backend::DType::I32); + let idx: Vec = indices.into_data().to_vec().unwrap(); + assert_eq!(idx, vec![1, 0]); + let vals: Vec = values.into_data().to_vec().unwrap(); + assert_eq!(vals, vec![5.0, 3.0]); + } + + #[test] + fn test_float_min_dim_with_indices_i32() { + use burn_backend::ops::FloatTensorOps; + use burn_std::IntDType; + + let t = crate::FlexTensor::from_data(TensorData::from([[1.0f32, 5.0], [3.0, 2.0]])); + let (values, indices) = Flex::float_min_dim_with_indices(t, 1, IntDType::I32); + assert_eq!(indices.dtype(), burn_backend::DType::I32); + let idx: Vec = indices.into_data().to_vec().unwrap(); + assert_eq!(idx, vec![0, 1]); + let vals: Vec = values.into_data().to_vec().unwrap(); + assert_eq!(vals, vec![1.0, 2.0]); + } + + #[test] + fn test_float_random_f64() { + use burn_backend::{DType, FloatDType, ops::FloatTensorOps}; + + let shape = burn_std::Shape::from(vec![100]); + let dist = burn_backend::Distribution::Uniform(0.0, 1.0); + let device = crate::FlexDevice; + let t = Flex::float_random(shape, dist, &device, FloatDType::F64); + assert_eq!(t.dtype(), DType::F64); + let data: Vec = t.into_data().to_vec().unwrap(); + assert!(data.iter().all(|&v| (0.0..=1.0).contains(&v))); + } + + #[test] + fn test_float_random_f16() { + use burn_backend::{DType, FloatDType, ops::FloatTensorOps}; + + let shape = burn_std::Shape::from(vec![100]); + let dist = burn_backend::Distribution::Uniform(0.0, 1.0); + let device = crate::FlexDevice; + let t = Flex::float_random(shape, dist, &device, FloatDType::F16); + assert_eq!(t.dtype(), DType::F16); + } +} diff --git a/crates/burn-flex/src/ops/gather_scatter.rs b/crates/burn-flex/src/ops/gather_scatter.rs new file mode 100644 index 0000000..8fef585 --- /dev/null +++ b/crates/burn-flex/src/ops/gather_scatter.rs @@ -0,0 +1,1251 @@ +//! Gather and scatter operations for indexed tensor access. + +use alloc::borrow::Cow; +use alloc::vec; +use alloc::vec::Vec; +use burn_backend::{DType, Element}; +use burn_std::{Bytes, Shape}; +use bytemuck::Pod; + +#[cfg(feature = "rayon")] +use rayon::prelude::*; + +use crate::{FlexTensor, Layout}; + +/// Read indices from a tensor as `isize`, the native offset type used by the +/// gather/scatter/select kernels in this module. +/// +/// This is the internal index layer for burn-flex: every indexed op +/// ([`gather`], [`scatter_add`], [`select`], [`select_add`], and the +/// [`scatter_min`]/[`scatter_max`] variants) routes its index tensor through +/// this helper before touching the element buffer. Normalising to `isize` +/// lets the kernels use a single inner-loop signature regardless of how the +/// caller's index tensor was dtyped. +/// +/// # Accepted widths +/// +/// Any of the integer DTypes `I8`, `I16`, `I32`, `I64`, `U8`, `U16`, `U32`, +/// `U64` is accepted. This is intentional: burn-flex's default `IntElem` is +/// I32 rather than the I64 convention used by other backends, and users can +/// also pin index tensors to any width they want via +/// `Tensor::from_data(.., (&device, DType::Ix))`. Whichever width lands here +/// is converted to `isize` on the fly. +/// +/// # Zero-copy vs. owned +/// +/// The return type is `Cow<'_, [isize]>` because only one width is zero-copy: +/// the one matching the host pointer width. On 64-bit targets, I64 indices +/// can be borrowed directly via `bytemuck::cast_slice` (both are 8 bytes). +/// Every other width requires an owned `Vec` with an element-wise +/// cast. U64 indices additionally go through a `try_from` to surface values +/// that would wrap when cast to `isize`. +/// +/// # History +/// +/// Earlier versions of `int_gather`, `int_scatter_add`, `int_select`, and +/// `int_select_add` carried a `debug_assert_eq!(indices.dtype(), DType::I64, +/// ..)` that contradicted this helper's contract. The asserts were dropped +/// in tracel-ai/burn#4776 once it was confirmed that `read_indices` had +/// always handled every supported width correctly at runtime. If you're +/// tempted to re-add a dtype check here, don't - the float siblings +/// ([`gather_f32`], [`select_f32`], ...) already share this helper without a +/// check, and asymmetry between the int and float paths was what surfaced +/// the bug. +fn read_indices(tensor: &FlexTensor) -> Cow<'_, [isize]> { + match tensor.dtype() { + #[cfg(target_pointer_width = "64")] + DType::I64 => { + const { assert!(size_of::() == size_of::()) }; + let data = tensor.storage::(); + Cow::Borrowed(bytemuck::cast_slice(data)) + } + #[cfg(target_pointer_width = "32")] + DType::I64 => Cow::Owned( + tensor + .storage::() + .iter() + .map(|&v| { + isize::try_from(v).unwrap_or_else(|_| { + panic!("read_indices: i64 index {v} out of isize range") + }) + }) + .collect(), + ), + #[cfg(target_pointer_width = "64")] + DType::I32 => Cow::Owned( + tensor + .storage::() + .iter() + .map(|&v| v as isize) + .collect(), + ), + #[cfg(target_pointer_width = "32")] + DType::I32 => { + const { assert!(size_of::() == size_of::()) }; + let data = tensor.storage::(); + Cow::Borrowed(bytemuck::cast_slice(data)) + } + DType::I16 => Cow::Owned( + tensor + .storage::() + .iter() + .map(|&v| v as isize) + .collect(), + ), + DType::I8 => Cow::Owned(tensor.storage::().iter().map(|&v| v as isize).collect()), + DType::U64 => Cow::Owned( + tensor + .storage::() + .iter() + .map(|&v| { + isize::try_from(v).unwrap_or_else(|_| { + panic!("read_indices: u64 index {v} out of isize range") + }) + }) + .collect(), + ), + #[cfg(target_pointer_width = "64")] + DType::U32 => Cow::Owned( + tensor + .storage::() + .iter() + .map(|&v| v as isize) + .collect(), + ), + #[cfg(target_pointer_width = "32")] + DType::U32 => Cow::Owned( + tensor + .storage::() + .iter() + .map(|&v| { + isize::try_from(v).unwrap_or_else(|_| { + panic!("read_indices: u32 index {v} out of isize range") + }) + }) + .collect(), + ), + DType::U16 => Cow::Owned( + tensor + .storage::() + .iter() + .map(|&v| v as isize) + .collect(), + ), + DType::U8 => Cow::Owned(tensor.storage::().iter().map(|&v| v as isize).collect()), + other => panic!("read_indices: unsupported index dtype {:?}", other), + } +} + +#[cold] +#[inline(never)] +fn index_oob(raw: isize, dim_size: usize) -> ! { + panic!("index {raw} out of bounds for dimension of size {dim_size}"); +} + +/// Validate an index is non-negative and within bounds, panicking with a clear message otherwise. +#[inline(always)] +fn checked_index(raw: isize, dim_size: usize) -> usize { + if raw < 0 || raw as usize >= dim_size { + index_oob(raw, dim_size); + } + raw as usize +} + +/// Gather values from tensor along a dimension using indices. +/// +/// For a 2D tensor with dim=1: +/// output[i, j] = tensor[i, indices[i, j]] +/// +/// The output has the same shape as indices. +pub fn gather( + tensor: FlexTensor, + dim: usize, + indices: FlexTensor, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let indices = indices.to_contiguous(); + + let tensor_shape = tensor.layout().shape(); + let indices_shape = indices.layout().shape(); + let ndims = tensor_shape.num_dims(); + + assert!( + dim < ndims, + "dim {} out of bounds for {} dimensions", + dim, + ndims + ); + + // Validate shapes: all dims except `dim` must match between tensor and indices + for i in 0..ndims { + if i != dim { + assert_eq!( + tensor_shape[i], indices_shape[i], + "gather: shape mismatch at dim {}: tensor {} vs indices {}", + i, tensor_shape[i], indices_shape[i] + ); + } + } + + let tensor_data: &[E] = tensor.storage(); + let indices_data = read_indices(&indices); + + // Calculate strides for tensor (row-major) + let tensor_strides: Vec = compute_strides(tensor_shape); + let indices_strides: Vec = compute_strides(indices_shape); + + let output_size = indices_shape.num_elements(); + + // Use specialized 2D implementation for common case + if ndims == 2 { + let result = gather_2d::( + tensor_data, + &indices_data, + tensor_shape[0], + tensor_shape[1], + indices_shape[0], + indices_shape[1], + dim, + ); + let bytes = Bytes::from_elems(result); + return FlexTensor::new(bytes, Layout::contiguous(indices_shape.clone()), E::dtype()); + } + + // General N-D case with pre-allocated coordinates + let dim_stride = tensor_strides[dim]; + + let gather_dim_size = tensor_shape[dim]; + + #[cfg(feature = "rayon")] + let result: Vec = (0..output_size) + .into_par_iter() + .map(|out_idx| { + let index_val = checked_index(indices_data[out_idx], gather_dim_size); + let src_idx = compute_gather_index( + out_idx, + index_val, + dim, + dim_stride, + &indices_strides, + &tensor_strides, + ndims, + ); + tensor_data[src_idx] + }) + .collect(); + + #[cfg(not(feature = "rayon"))] + let result: Vec = (0..output_size) + .map(|out_idx| { + let index_val = checked_index(indices_data[out_idx], gather_dim_size); + let src_idx = compute_gather_index( + out_idx, + index_val, + dim, + dim_stride, + &indices_strides, + &tensor_strides, + ndims, + ); + tensor_data[src_idx] + }) + .collect(); + + let bytes = Bytes::from_elems(result); + FlexTensor::new(bytes, Layout::contiguous(indices_shape.clone()), E::dtype()) +} + +/// Optimized 2D gather implementation. +#[inline] +fn gather_2d( + tensor_data: &[E], + indices_data: &[isize], + tensor_rows: usize, + tensor_cols: usize, + indices_rows: usize, + indices_cols: usize, + dim: usize, +) -> Vec { + let output_size = indices_rows * indices_cols; + let dim_size = if dim == 0 { tensor_rows } else { tensor_cols }; + + let mut result = vec![E::default(); output_size]; + + #[cfg(feature = "rayon")] + const PARALLEL_THRESHOLD: usize = 256 * 1024; + + #[cfg(feature = "rayon")] + if output_size >= PARALLEL_THRESHOLD { + if dim == 0 { + result + .par_chunks_mut(indices_cols) + .enumerate() + .for_each(|(i, row)| { + for j in 0..indices_cols { + let src_row = checked_index(indices_data[i * indices_cols + j], dim_size); + row[j] = tensor_data[src_row * tensor_cols + j]; + } + }); + } else { + result + .par_chunks_mut(indices_cols) + .enumerate() + .for_each(|(i, row)| { + for j in 0..indices_cols { + let src_col = checked_index(indices_data[i * indices_cols + j], dim_size); + row[j] = tensor_data[i * tensor_cols + src_col]; + } + }); + } + } else if dim == 0 { + for i in 0..indices_rows { + for j in 0..indices_cols { + let src_row = checked_index(indices_data[i * indices_cols + j], dim_size); + result[i * indices_cols + j] = tensor_data[src_row * tensor_cols + j]; + } + } + } else { + for i in 0..indices_rows { + for j in 0..indices_cols { + let src_col = checked_index(indices_data[i * indices_cols + j], dim_size); + result[i * indices_cols + j] = tensor_data[i * tensor_cols + src_col]; + } + } + } + + #[cfg(not(feature = "rayon"))] + { + if dim == 0 { + for i in 0..indices_rows { + for j in 0..indices_cols { + let src_row = checked_index(indices_data[i * indices_cols + j], dim_size); + result[i * indices_cols + j] = tensor_data[src_row * tensor_cols + j]; + } + } + } else { + for i in 0..indices_rows { + for j in 0..indices_cols { + let src_col = checked_index(indices_data[i * indices_cols + j], dim_size); + result[i * indices_cols + j] = tensor_data[i * tensor_cols + src_col]; + } + } + } + } + + result +} + +/// Compute source index for gather operation (N-D case). +#[inline] +fn compute_gather_index( + out_idx: usize, + index_val: usize, + dim: usize, + dim_stride: usize, + indices_strides: &[usize], + tensor_strides: &[usize], + ndims: usize, +) -> usize { + let mut src_idx = index_val * dim_stride; + let mut remaining = out_idx; + + for d in 0..ndims { + if d != dim { + let coord = remaining / indices_strides[d]; + remaining %= indices_strides[d]; + src_idx += coord * tensor_strides[d]; + } else { + remaining %= indices_strides[d]; + } + } + src_idx +} + +/// Scatter add: adds values to tensor at positions specified by indices. +pub fn scatter_add( + tensor: FlexTensor, + dim: usize, + indices: FlexTensor, + value: FlexTensor, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let indices = indices.to_contiguous(); + let value = value.to_contiguous(); + + let tensor_shape = tensor.layout().shape().clone(); + let indices_shape = indices.layout().shape(); + let value_shape = value.layout().shape(); + let ndims = tensor_shape.num_dims(); + + assert!( + dim < ndims, + "dim {} out of bounds for {} dimensions", + dim, + ndims + ); + assert_eq!( + indices_shape, + value_shape, + "scatter_add: indices shape {:?} must match value shape {:?}", + indices_shape.to_vec(), + value_shape.to_vec() + ); + + for i in 0..ndims { + if i != dim { + assert_eq!( + tensor_shape[i], indices_shape[i], + "scatter_add: shape mismatch at dim {}: tensor {} vs indices {}", + i, tensor_shape[i], indices_shape[i] + ); + } + } + + let tensor_data: &[E] = tensor.storage(); + let indices_data = read_indices(&indices); + let value_data: &[E] = value.storage(); + + let mut result: Vec = tensor_data.to_vec(); + + let tensor_strides: Vec = compute_strides(&tensor_shape); + let indices_strides: Vec = compute_strides(indices_shape); + + let num_elements = indices_shape.num_elements(); + + // Use specialized 2D implementation + if ndims == 2 { + scatter_add_2d( + &mut result, + &indices_data, + value_data, + tensor_shape[0], + tensor_shape[1], + indices_shape[0], + indices_shape[1], + dim, + ); + } else { + // General N-D case (sequential due to potential index conflicts) + let dim_stride = tensor_strides[dim]; + let scatter_dim_size = tensor_shape[dim]; + for idx in 0..num_elements { + let index_val = checked_index(indices_data[idx], scatter_dim_size); + let dst_idx = compute_gather_index( + idx, + index_val, + dim, + dim_stride, + &indices_strides, + &tensor_strides, + ndims, + ); + result[dst_idx] += value_data[idx]; + } + } + + let bytes = Bytes::from_elems(result); + FlexTensor::new(bytes, Layout::contiguous(tensor_shape), E::dtype()) +} + +/// Optimized 2D scatter_add implementation. +#[inline] +#[allow(clippy::too_many_arguments)] +fn scatter_add_2d( + result: &mut [E], + indices_data: &[isize], + value_data: &[E], + tensor_rows: usize, + tensor_cols: usize, + indices_rows: usize, + indices_cols: usize, + dim: usize, +) { + let dim_size = if dim == 0 { tensor_rows } else { tensor_cols }; + if dim == 0 { + for i in 0..indices_rows { + for j in 0..indices_cols { + let idx = i * indices_cols + j; + let dst_row = checked_index(indices_data[idx], dim_size); + result[dst_row * tensor_cols + j] += value_data[idx]; + } + } + } else { + for i in 0..indices_rows { + for j in 0..indices_cols { + let idx = i * indices_cols + j; + let dst_col = checked_index(indices_data[idx], dim_size); + result[i * tensor_cols + dst_col] += value_data[idx]; + } + } + } +} + +/// Select slices from tensor along a dimension using 1D indices. +/// +/// Unlike gather, indices is 1D and selects entire slices. +/// For a 2D tensor with dim=0 and indices=[2, 0]: +/// output[0, :] = tensor[2, :] +/// output[1, :] = tensor[0, :] +pub fn select( + tensor: FlexTensor, + dim: usize, + indices: FlexTensor, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let indices = indices.to_contiguous(); + + let tensor_shape = tensor.layout().shape(); + let ndims = tensor_shape.num_dims(); + + assert!( + dim < ndims, + "dim {} out of bounds for {} dimensions", + dim, + ndims + ); + assert_eq!( + indices.layout().num_dims(), + 1, + "select: indices must be 1D, got {} dims", + indices.layout().num_dims() + ); + + let tensor_data: &[E] = tensor.storage(); + let indices_data = read_indices(&indices); + let num_indices = indices_data.len(); + + // Build output shape: replace dim with num_indices + let mut output_dims = tensor_shape.to_vec(); + output_dims[dim] = num_indices; + let output_shape = Shape::from(output_dims); + + // Use optimized 2D implementation with bulk copies + if ndims == 2 { + let result = select_2d::( + tensor_data, + &indices_data, + tensor_shape[0], + tensor_shape[1], + num_indices, + dim, + ); + let bytes = Bytes::from_elems(result); + return FlexTensor::new(bytes, Layout::contiguous(output_shape), E::dtype()); + } + + // General N-D case + let tensor_strides: Vec = compute_strides(tensor_shape); + let output_strides: Vec = compute_strides(&output_shape); + let output_size = output_shape.num_elements(); + + // Calculate slice size (elements after dim) + let slice_size: usize = tensor_strides[dim]; + + let select_dim_size = tensor_shape[dim]; + + // If dim is the last dimension or we can use bulk copies + if dim == ndims - 1 || slice_size == 1 { + // Element-wise with parallelism + #[cfg(feature = "rayon")] + let result: Vec = (0..output_size) + .into_par_iter() + .map(|out_idx| { + let mut remaining = out_idx; + let mut src_idx = 0; + for d in 0..ndims { + let coord = remaining / output_strides[d]; + remaining %= output_strides[d]; + if d == dim { + let index_val = checked_index(indices_data[coord], select_dim_size); + src_idx += index_val * tensor_strides[d]; + } else { + src_idx += coord * tensor_strides[d]; + } + } + tensor_data[src_idx] + }) + .collect(); + + #[cfg(not(feature = "rayon"))] + #[allow(clippy::needless_range_loop)] + let result: Vec = { + let mut result = vec![E::default(); output_size]; + for out_idx in 0..output_size { + let mut remaining = out_idx; + let mut src_idx = 0; + for d in 0..ndims { + let coord = remaining / output_strides[d]; + remaining %= output_strides[d]; + if d == dim { + let index_val = checked_index(indices_data[coord], select_dim_size); + src_idx += index_val * tensor_strides[d]; + } else { + src_idx += coord * tensor_strides[d]; + } + } + result[out_idx] = tensor_data[src_idx]; + } + result + }; + + let bytes = Bytes::from_elems(result); + return FlexTensor::new(bytes, Layout::contiguous(output_shape), E::dtype()); + } + + // Use bulk copies for contiguous slices + let mut result = vec![E::default(); output_size]; + + // For each position in dimensions before `dim` + let outer_count = if dim == 0 { + 1 + } else { + tensor_shape[..dim].iter().product() + }; + + for outer in 0..outer_count { + let outer_offset_tensor = outer * tensor_strides[if dim == 0 { 0 } else { dim - 1 }]; + let outer_offset_output = outer * output_strides[if dim == 0 { 0 } else { dim - 1 }]; + + for (i, &idx) in indices_data.iter().enumerate() { + let index_val = checked_index(idx, select_dim_size); + let src_start = outer_offset_tensor + index_val * tensor_strides[dim]; + let dst_start = outer_offset_output + i * output_strides[dim]; + result[dst_start..dst_start + slice_size] + .copy_from_slice(&tensor_data[src_start..src_start + slice_size]); + } + } + + let bytes = Bytes::from_elems(result); + FlexTensor::new(bytes, Layout::contiguous(output_shape), E::dtype()) +} + +/// Optimized 2D select with bulk row copies when dim=0. +#[inline] +fn select_2d( + tensor_data: &[E], + indices_data: &[isize], + tensor_rows: usize, + tensor_cols: usize, + num_indices: usize, + dim: usize, +) -> Vec { + let dim_size = if dim == 0 { tensor_rows } else { tensor_cols }; + let (output_rows, output_cols) = if dim == 0 { + (num_indices, tensor_cols) + } else { + (tensor_rows, num_indices) + }; + let output_size = output_rows * output_cols; + + // Minimum bytes of output before we consider rayon. Below this, a + // single-threaded loop is faster because there is not enough work to + // amortize the work-stealing dispatch overhead. + #[cfg(feature = "rayon")] + const PARALLEL_THRESHOLD_BYTES: usize = 4 * 1024 * 1024; + + // Minimum elements per rayon task. Without batching, par_chunks_mut + // creates one task per row (e.g. 512 single-row tasks of 4 KB each) + // whose dispatch overhead dominates the actual copy. + #[cfg(feature = "rayon")] + const MIN_ELEMS_PER_TASK: usize = 64 * 1024; + + if dim == 0 { + // SAFETY: the output has exactly num_indices * tensor_cols elements. + // Both the parallel and serial paths below write every element exactly + // once via non-overlapping row copies, so no element is left uninitialized. + let mut result = Vec::with_capacity(output_size); + #[allow(clippy::uninit_vec)] + unsafe { + result.set_len(output_size) + }; + + #[cfg(feature = "rayon")] + if output_size * size_of::() >= PARALLEL_THRESHOLD_BYTES { + // Batch multiple rows per rayon task so each task copies at + // least MIN_ELEMS_PER_TASK elements. + let rows_per_chunk = (MIN_ELEMS_PER_TASK / tensor_cols).max(1); + let elems_per_chunk = rows_per_chunk * tensor_cols; + result.par_chunks_mut(elems_per_chunk).enumerate().for_each( + |(chunk_idx, dst_chunk)| { + let start_row = chunk_idx * rows_per_chunk; + let chunk_rows = dst_chunk.len() / tensor_cols; + for i in 0..chunk_rows { + let src_row_idx = checked_index(indices_data[start_row + i], dim_size); + let src_start = src_row_idx * tensor_cols; + let dst_start = i * tensor_cols; + dst_chunk[dst_start..dst_start + tensor_cols] + .copy_from_slice(&tensor_data[src_start..src_start + tensor_cols]); + } + }, + ); + } else { + for (i, &idx) in indices_data.iter().enumerate() { + let src_row_idx = checked_index(idx, dim_size); + let src_start = src_row_idx * tensor_cols; + let dst_start = i * tensor_cols; + result[dst_start..dst_start + tensor_cols] + .copy_from_slice(&tensor_data[src_start..src_start + tensor_cols]); + } + } + + #[cfg(not(feature = "rayon"))] + { + for (i, &idx) in indices_data.iter().enumerate() { + let src_row_idx = checked_index(idx, dim_size); + let src_start = src_row_idx * tensor_cols; + let dst_start = i * tensor_cols; + result[dst_start..dst_start + tensor_cols] + .copy_from_slice(&tensor_data[src_start..src_start + tensor_cols]); + } + } + + result + } else { + // dim == 1: gather individual elements per row (not contiguous). + // Zero-init is fine here since the inner loop is per-element anyway. + let mut result = vec![E::default(); output_size]; + + #[cfg(feature = "rayon")] + if output_size * size_of::() >= PARALLEL_THRESHOLD_BYTES { + let rows_per_chunk = (MIN_ELEMS_PER_TASK / output_cols).max(1); + let elems_per_chunk = rows_per_chunk * output_cols; + result.par_chunks_mut(elems_per_chunk).enumerate().for_each( + |(chunk_idx, dst_chunk)| { + let start_row = chunk_idx * rows_per_chunk; + let chunk_rows = dst_chunk.len() / output_cols; + for r in 0..chunk_rows { + let row = start_row + r; + let dst_base = r * output_cols; + for (j, &idx) in indices_data.iter().enumerate() { + let src_col = checked_index(idx, dim_size); + dst_chunk[dst_base + j] = tensor_data[row * tensor_cols + src_col]; + } + } + }, + ); + } else { + for row in 0..output_rows { + for (j, &idx) in indices_data.iter().enumerate() { + let src_col = checked_index(idx, dim_size); + result[row * output_cols + j] = tensor_data[row * tensor_cols + src_col]; + } + } + } + + #[cfg(not(feature = "rayon"))] + { + for row in 0..output_rows { + for (j, &idx) in indices_data.iter().enumerate() { + let src_col = checked_index(idx, dim_size); + result[row * output_cols + j] = tensor_data[row * tensor_cols + src_col]; + } + } + } + + result + } +} + +/// Select add: adds values back to tensor at positions specified by 1D indices. +pub fn select_add( + tensor: FlexTensor, + dim: usize, + indices: FlexTensor, + value: FlexTensor, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let indices = indices.to_contiguous(); + let value = value.to_contiguous(); + + let tensor_shape = tensor.layout().shape().clone(); + let value_shape = value.layout().shape(); + let ndims = tensor_shape.num_dims(); + + assert!( + dim < ndims, + "dim {} out of bounds for {} dimensions", + dim, + ndims + ); + assert_eq!( + indices.layout().num_dims(), + 1, + "select_add: indices must be 1D" + ); + + let tensor_data: &[E] = tensor.storage(); + let indices_data = read_indices(&indices); + let value_data: &[E] = value.storage(); + let num_indices = indices_data.len(); + + // Validate value shape + for d in 0..ndims { + if d == dim { + assert_eq!( + value_shape[d], num_indices, + "select_add: value dim {} should be {} (num indices), got {}", + d, num_indices, value_shape[d] + ); + } else { + assert_eq!( + value_shape[d], tensor_shape[d], + "select_add: value dim {} should match tensor dim {}, got {}", + d, tensor_shape[d], value_shape[d] + ); + } + } + + let mut result: Vec = tensor_data.to_vec(); + + // Use optimized 2D implementation + if ndims == 2 { + select_add_2d( + &mut result, + &indices_data, + value_data, + tensor_shape[0], + tensor_shape[1], + num_indices, + dim, + ); + let bytes = Bytes::from_elems(result); + return FlexTensor::new(bytes, Layout::contiguous(tensor_shape), E::dtype()); + } + + // General N-D case + let tensor_strides: Vec = compute_strides(&tensor_shape); + let value_strides: Vec = compute_strides(value_shape); + let select_add_dim_size = tensor_shape[dim]; + + for (val_idx, &val) in value_data.iter().enumerate() { + let mut remaining = val_idx; + let mut dst_idx = 0; + for d in 0..ndims { + let coord = remaining / value_strides[d]; + remaining %= value_strides[d]; + if d == dim { + let index_val = checked_index(indices_data[coord], select_add_dim_size); + dst_idx += index_val * tensor_strides[d]; + } else { + dst_idx += coord * tensor_strides[d]; + } + } + result[dst_idx] += val; + } + + let bytes = Bytes::from_elems(result); + FlexTensor::new(bytes, Layout::contiguous(tensor_shape), E::dtype()) +} + +/// Optimized 2D select_add. +#[inline] +fn select_add_2d( + result: &mut [E], + indices_data: &[isize], + value_data: &[E], + tensor_rows: usize, + tensor_cols: usize, + num_indices: usize, + dim: usize, +) { + let dim_size = if dim == 0 { tensor_rows } else { tensor_cols }; + if dim == 0 { + for (i, &idx) in indices_data.iter().enumerate() { + let dst_row = checked_index(idx, dim_size); + let dst_start = dst_row * tensor_cols; + let src_start = i * tensor_cols; + for j in 0..tensor_cols { + result[dst_start + j] += value_data[src_start + j]; + } + } + } else { + for row in 0..tensor_rows { + for (j, &idx) in indices_data.iter().enumerate() { + let dst_col = checked_index(idx, dim_size); + result[row * tensor_cols + dst_col] += value_data[row * num_indices + j]; + } + } + } +} + +/// Compute row-major strides for a shape. +#[inline] +fn compute_strides(dims: &[usize]) -> Vec { + let ndims = dims.len(); + let mut strides = vec![1usize; ndims]; + for i in (0..ndims.saturating_sub(1)).rev() { + strides[i] = strides[i + 1] * dims[i + 1]; + } + strides +} + +/// Multi-dimensional scatter: update `data` at locations specified by N-dimensional index tuples. +pub fn scatter_nd< + E: Element + Pod + Default + Copy + core::ops::AddAssign + core::ops::Mul + PartialOrd, +>( + data: FlexTensor, + indices: FlexTensor, + values: FlexTensor, + reduction: burn_backend::tensor::IndexingUpdateOp, +) -> FlexTensor { + use burn_backend::tensor::IndexingUpdateOp; + + let data = data.to_contiguous(); + let indices = indices.to_contiguous(); + let values = values.to_contiguous(); + + let data_shape: Vec = data.layout().shape().to_vec(); + let idx_shape: Vec = indices.layout().shape().to_vec(); + let m = idx_shape.len(); + let k = idx_shape[m - 1]; + + let num_indices: usize = idx_shape[..m - 1].iter().product(); + let slice_size: usize = data_shape[k..].iter().product(); + + let data_data: &[E] = data.storage(); + let idx_data = read_indices(&indices); + let val_data: &[E] = values.storage(); + + let mut result: Vec = data_data.to_vec(); + + let strides = compute_strides(&data_shape); + + for n in 0..num_indices { + let mut base_offset = 0usize; + for j in 0..k { + let idx_val = idx_data[n * k + j] as usize; + base_offset += idx_val * strides[j]; + } + + let val_offset = n * slice_size; + match reduction { + IndexingUpdateOp::Assign => { + result[base_offset..(base_offset + slice_size)] + .copy_from_slice(&val_data[val_offset..(val_offset + slice_size)]); + } + IndexingUpdateOp::Add => { + for s in 0..slice_size { + result[base_offset + s] += val_data[val_offset + s]; + } + } + IndexingUpdateOp::Mul => { + for s in 0..slice_size { + result[base_offset + s] = result[base_offset + s] * val_data[val_offset + s]; + } + } + IndexingUpdateOp::Min => { + for s in 0..slice_size { + let b = val_data[val_offset + s]; + if b < result[base_offset + s] { + result[base_offset + s] = b; + } + } + } + IndexingUpdateOp::Max => { + for s in 0..slice_size { + let b = val_data[val_offset + s]; + if b > result[base_offset + s] { + result[base_offset + s] = b; + } + } + } + } + } + + let shape = Shape::from(data_shape); + let bytes = Bytes::from_elems(result); + FlexTensor::new(bytes, Layout::contiguous(shape), E::dtype()) +} + +/// Multi-dimensional gather: collect slices from `data` at locations specified by N-dimensional +/// index tuples. +pub fn gather_nd( + data: FlexTensor, + indices: FlexTensor, +) -> FlexTensor { + let data = data.to_contiguous(); + let indices = indices.to_contiguous(); + + let data_shape: Vec = data.layout().shape().to_vec(); + let idx_shape: Vec = indices.layout().shape().to_vec(); + let m = idx_shape.len(); + let k = idx_shape[m - 1]; + + let num_indices: usize = idx_shape[..m - 1].iter().product(); + let slice_size: usize = data_shape[k..].iter().product(); + + let data_data: &[E] = data.storage(); + let idx_data = read_indices(&indices); + + let mut out_shape_vec: Vec = idx_shape[..m - 1].to_vec(); + out_shape_vec.extend_from_slice(&data_shape[k..]); + + let strides = compute_strides(&data_shape); + + let total = num_indices * slice_size; + let mut result = vec![E::default(); total]; + + for n in 0..num_indices { + let mut base_offset = 0usize; + for j in 0..k { + let idx_val = idx_data[n * k + j] as usize; + base_offset += idx_val * strides[j]; + } + let out_offset = n * slice_size; + result[out_offset..(out_offset + slice_size)] + .copy_from_slice(&data_data[base_offset..(base_offset + slice_size)]); + } + + let shape = Shape::from(out_shape_vec); + let bytes = Bytes::from_elems(result); + FlexTensor::new(bytes, Layout::contiguous(shape), E::dtype()) +} + +// Type-specific wrappers + +pub fn gather_f32(tensor: FlexTensor, dim: usize, indices: FlexTensor) -> FlexTensor { + gather::(tensor, dim, indices) +} + +pub fn gather_f64(tensor: FlexTensor, dim: usize, indices: FlexTensor) -> FlexTensor { + gather::(tensor, dim, indices) +} + +pub fn gather_i64(tensor: FlexTensor, dim: usize, indices: FlexTensor) -> FlexTensor { + gather::(tensor, dim, indices) +} + +pub fn scatter_add_f32( + tensor: FlexTensor, + dim: usize, + indices: FlexTensor, + value: FlexTensor, +) -> FlexTensor { + scatter_add::(tensor, dim, indices, value) +} + +pub fn scatter_add_f64( + tensor: FlexTensor, + dim: usize, + indices: FlexTensor, + value: FlexTensor, +) -> FlexTensor { + scatter_add::(tensor, dim, indices, value) +} + +pub fn scatter_add_i64( + tensor: FlexTensor, + dim: usize, + indices: FlexTensor, + value: FlexTensor, +) -> FlexTensor { + scatter_add::(tensor, dim, indices, value) +} + +pub fn select_f32(tensor: FlexTensor, dim: usize, indices: FlexTensor) -> FlexTensor { + select::(tensor, dim, indices) +} + +pub fn select_f64(tensor: FlexTensor, dim: usize, indices: FlexTensor) -> FlexTensor { + select::(tensor, dim, indices) +} + +pub fn select_i64(tensor: FlexTensor, dim: usize, indices: FlexTensor) -> FlexTensor { + select::(tensor, dim, indices) +} + +pub fn select_add_f32( + tensor: FlexTensor, + dim: usize, + indices: FlexTensor, + value: FlexTensor, +) -> FlexTensor { + select_add::(tensor, dim, indices, value) +} + +pub fn select_add_f64( + tensor: FlexTensor, + dim: usize, + indices: FlexTensor, + value: FlexTensor, +) -> FlexTensor { + select_add::(tensor, dim, indices, value) +} + +pub fn select_add_i64( + tensor: FlexTensor, + dim: usize, + indices: FlexTensor, + value: FlexTensor, +) -> FlexTensor { + select_add::(tensor, dim, indices, value) +} + +// Bool-specific operations + +pub fn gather_bool(tensor: FlexTensor, dim: usize, indices: FlexTensor) -> FlexTensor { + gather::(tensor, dim, indices) +} + +/// Scatter OR for bool tensors: ORs values into tensor at indexed positions. +pub fn scatter_or( + tensor: FlexTensor, + dim: usize, + indices: FlexTensor, + value: FlexTensor, +) -> FlexTensor { + // Preserve the input tensor's bool dtype for the output. + let out_dtype = burn_std::BoolDType::from(tensor.dtype()); + let tensor = tensor.to_contiguous(); + let indices = indices.to_contiguous(); + let value = value.to_contiguous(); + + let tensor_shape = tensor.layout().shape().clone(); + let indices_shape = indices.layout().shape(); + let value_shape = value.layout().shape(); + let ndims = tensor_shape.num_dims(); + + assert!( + dim < ndims, + "dim {} out of bounds for {} dimensions", + dim, + ndims + ); + assert_eq!( + indices_shape, + value_shape, + "scatter_or: indices shape {:?} must match value shape {:?}", + indices_shape.to_vec(), + value_shape.to_vec() + ); + + for i in 0..ndims { + if i != dim { + assert_eq!( + tensor_shape[i], indices_shape[i], + "scatter_or: shape mismatch at dim {}: tensor {} vs indices {}", + i, tensor_shape[i], indices_shape[i] + ); + } + } + + let tensor_data: &[u8] = tensor.storage(); + let indices_data = read_indices(&indices); + let value_data: &[u8] = value.storage(); + + let mut result: Vec = tensor_data.to_vec(); + + let tensor_strides = compute_strides(&tensor_shape); + let indices_strides = compute_strides(indices_shape); + + let num_elements = indices_shape.num_elements(); + + let scatter_or_dim_size = tensor_shape[dim]; + + // Use 2D specialized path + if ndims == 2 { + let tensor_cols = tensor_shape[1]; + let indices_rows = indices_shape[0]; + let indices_cols = indices_shape[1]; + + if dim == 0 { + for i in 0..indices_rows { + for j in 0..indices_cols { + let idx = i * indices_cols + j; + let dst_row = checked_index(indices_data[idx], scatter_or_dim_size); + result[dst_row * tensor_cols + j] |= value_data[idx]; + } + } + } else { + for i in 0..indices_rows { + for j in 0..indices_cols { + let idx = i * indices_cols + j; + let dst_col = checked_index(indices_data[idx], scatter_or_dim_size); + result[i * tensor_cols + dst_col] |= value_data[idx]; + } + } + } + } else { + let dim_stride = tensor_strides[dim]; + for idx in 0..num_elements { + let index_val = checked_index(indices_data[idx], scatter_or_dim_size); + let dst_idx = compute_gather_index( + idx, + index_val, + dim, + dim_stride, + &indices_strides, + &tensor_strides, + ndims, + ); + result[dst_idx] |= value_data[idx]; + } + } + + crate::ops::comparison::make_bool_tensor(result, tensor_shape, out_dtype) +} + +// Tests kept here probe flex-specific behavior: non-I64 index dtype +// acceptance through the internal `read_indices` path and the uninit- +// buffer + rayon-chunked `select` kernel. Plain gather/scatter/select +// tests (including an empty-indices edge case) live in +// crates/burn-backend-tests/tests/tensor/float/ops/{gather_scatter,select}.rs +// so every backend is exercised. When adding new tests, keep them here +// only if they probe flex internals; otherwise add them there. +#[cfg(test)] +mod tests { + use super::*; + use burn_backend::TensorData; + + #[test] + fn test_gather_with_i32_indices() { + let tensor = FlexTensor::from_data(TensorData::new(vec![10.0f32, 20.0, 30.0, 40.0], [4])); + let indices = FlexTensor::from_data(TensorData::new(vec![3i32, 0, 2], [3])); + + let result = gather::(tensor, 0, indices); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![40.0, 10.0, 30.0]); + } + + #[test] + fn test_select_with_i32_indices() { + let tensor = FlexTensor::from_data(TensorData::new( + vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], + [3, 2], + )); + let indices = FlexTensor::from_data(TensorData::new(vec![2i32, 0], [2])); + + let result = select::(tensor, 0, indices); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![5.0, 6.0, 1.0, 2.0]); + } + + /// Exercises the uninit buffer path (dim=0) and, when rayon is enabled, + /// the chunked parallel path (output > 4 MB). + #[test] + fn test_select_2d_dim0_large() { + let rows = 2048; + let cols = 1024; + let data: Vec = (0..rows * cols).map(|i| i as f32).collect(); + let tensor = FlexTensor::from_data(TensorData::new(data.clone(), [rows, cols])); + + // Select every other row in reverse order. + let idx: Vec = (0..rows as i64).rev().step_by(2).collect(); + let num_idx = idx.len(); + let indices = FlexTensor::from_data(TensorData::new(idx.clone(), [num_idx])); + + let result = select::(tensor, 0, indices); + assert_eq!(result.layout().shape().to_vec(), vec![num_idx, cols]); + let out: Vec = result.into_data().to_vec().unwrap(); + + for (i, &row_idx) in idx.iter().enumerate() { + let expected_start = row_idx as usize * cols; + let actual = &out[i * cols..(i + 1) * cols]; + let expected = &data[expected_start..expected_start + cols]; + assert_eq!( + actual, expected, + "mismatch at output row {i} (src row {row_idx})" + ); + } + } +} diff --git a/crates/burn-flex/src/ops/grid_sample.rs b/crates/burn-flex/src/ops/grid_sample.rs new file mode 100644 index 0000000..3f233af --- /dev/null +++ b/crates/burn-flex/src/ops/grid_sample.rs @@ -0,0 +1,244 @@ +//! Grid sampling operations for FlexTensor. +//! +//! Supported dtypes: f32, f64, f16, bf16. All dtypes share a single f64 +//! compute path. This is required for f16/bf16 correctness (so coordinate +//! math, bilinear weights, and accumulated samples keep full precision) and +//! incidentally gives f32 a small accuracy bump at the cost of extra casts. + +use alloc::vec; +use alloc::vec::Vec; +use burn_backend::element::cast::ToElement; +use burn_backend::ops::{GridSampleOptions, GridSamplePaddingMode, InterpolateMode}; +use burn_backend::{DType, Element}; +use burn_std::{Bytes, Shape, bf16, f16}; + +use num_traits::{Float, NumCast}; + +use crate::{FlexTensor, Layout}; + +/// Grid sample 2D (bilinear and nearest-neighbor interpolation). +/// +/// Input tensor shape: [N, C, H_in, W_in] +/// Grid shape: [N, H_out, W_out, 2] (x, y normalized to [-1, 1]) +/// Output shape: [N, C, H_out, W_out] +pub fn grid_sample_2d( + tensor: FlexTensor, + grid: FlexTensor, + options: GridSampleOptions, +) -> FlexTensor { + match options.mode { + InterpolateMode::Bilinear | InterpolateMode::Nearest => {} + other => panic!("grid_sample_2d: {:?} mode is not supported", other), + } + + let tensor = tensor.to_contiguous(); + let grid = grid.to_contiguous(); + + match tensor.dtype() { + DType::F32 => grid_sample_2d_impl::(tensor, grid, options), + DType::F64 => grid_sample_2d_impl::(tensor, grid, options), + DType::F16 => grid_sample_2d_impl::(tensor, grid, options), + DType::BF16 => grid_sample_2d_impl::(tensor, grid, options), + _ => panic!("grid_sample_2d: unsupported dtype {:?}", tensor.dtype()), + } +} + +fn grid_sample_2d_impl( + tensor: FlexTensor, + grid: FlexTensor, + options: GridSampleOptions, +) -> FlexTensor +where + T: Float + Element + bytemuck::Pod, +{ + let t_shape = tensor.layout().shape(); + let g_shape = grid.layout().shape(); + + assert_eq!(t_shape.num_dims(), 4, "grid_sample_2d: input must be 4D"); + assert_eq!(g_shape.num_dims(), 4, "grid_sample_2d: grid must be 4D"); + assert_eq!(g_shape[3], 2, "grid_sample_2d: grid last dim must be 2"); + assert_eq!( + t_shape[0], g_shape[0], + "grid_sample_2d: batch size mismatch" + ); + + let batch_size = t_shape[0]; + let channels = t_shape[1]; + let h_in = t_shape[2]; + let w_in = t_shape[3]; + let h_out = g_shape[1]; + let w_out = g_shape[2]; + + let tensor_data: &[T] = tensor.storage(); + let grid_data: &[T] = grid.storage(); + + let out_shape = Shape::from(vec![batch_size, channels, h_out, w_out]); + let out_len = batch_size * channels * h_out * w_out; + let mut output: Vec = vec![T::zero(); out_len]; + + let align = options.align_corners; + let pad_mode = options.padding_mode; + + let t_stride_n = channels * h_in * w_in; + let t_stride_c = h_in * w_in; + let t_stride_h = w_in; + + let g_stride_n = h_out * w_out * 2; + let g_stride_h = w_out * 2; + + let o_stride_n = channels * h_out * w_out; + let o_stride_c = h_out * w_out; + let o_stride_h = w_out; + + // Low-precision types (f16/bf16) are widened to f64 for all arithmetic so + // that coordinate math, weights, and accumulated samples keep full precision. + // + // The from_f64 unwrap is unreachable for any well-formed input: bilinear + // is a convex combination of finite samples so the result stays bounded by + // the sample envelope. The `half` crate's `NumCast` impl for f16/bf16 + // forwards through `to_f32` and maps non-finite inputs to `Some(inf/nan)`; + // `num_traits`'s f32/f64 impls do the same. The message-bearing panic is a + // diagnostic hatch for future dtypes where this invariant does not hold. + let to_f64 = |x: T| -> f64 { ToElement::to_f64(&x) }; + let from_f64 = |x: f64| -> T { + ::from(x).unwrap_or_else(|| { + panic!( + "grid_sample_2d: NumCast::from({x:?}) to {:?} returned None", + T::dtype() + ) + }) + }; + + for b in 0..batch_size { + for y in 0..h_out { + for x in 0..w_out { + let g_idx = b * g_stride_n + y * g_stride_h + x * 2; + let sample_x = to_f64(grid_data[g_idx]); + let sample_y = to_f64(grid_data[g_idx + 1]); + + let (px, py) = if align { + let px = (sample_x + 1.0) * ((w_in - 1) as f64) / 2.0; + let py = (sample_y + 1.0) * ((h_in - 1) as f64) / 2.0; + (px, py) + } else { + let px = (sample_x + 1.0) * (w_in as f64) / 2.0 - 0.5; + let py = (sample_y + 1.0) * (h_in as f64) / 2.0 - 0.5; + (px, py) + }; + + let (px, py) = apply_padding(px, py, w_in, h_in, pad_mode, align); + + let read = |t_base: usize, xi: i64, yi: i64| -> f64 { + match pad_mode { + GridSamplePaddingMode::Zeros => { + if xi >= 0 && xi < w_in as i64 && yi >= 0 && yi < h_in as i64 { + to_f64(tensor_data[t_base + yi as usize * t_stride_h + xi as usize]) + } else { + 0.0 + } + } + GridSamplePaddingMode::Border | GridSamplePaddingMode::Reflection => { + let xi = xi.clamp(0, (w_in - 1) as i64) as usize; + let yi = yi.clamp(0, (h_in - 1) as i64) as usize; + to_f64(tensor_data[t_base + yi * t_stride_h + xi]) + } + } + }; + + for c in 0..channels { + let t_base = b * t_stride_n + c * t_stride_c; + let o_idx = b * o_stride_n + c * o_stride_c + y * o_stride_h + x; + + let val = if matches!(options.mode, InterpolateMode::Nearest) { + let xi = px.round() as i64; + let yi = py.round() as i64; + read(t_base, xi, yi) + } else { + // Bilinear + let x0 = px.floor() as i64; + let y0 = py.floor() as i64; + let x1 = x0 + 1; + let y1 = y0 + 1; + + let x_frac = px - px.floor(); + let y_frac = py - py.floor(); + + let w00 = (1.0 - x_frac) * (1.0 - y_frac); + let w01 = (1.0 - x_frac) * y_frac; + let w10 = x_frac * (1.0 - y_frac); + let w11 = x_frac * y_frac; + + read(t_base, x0, y0) * w00 + + read(t_base, x0, y1) * w01 + + read(t_base, x1, y0) * w10 + + read(t_base, x1, y1) * w11 + }; + + output[o_idx] = from_f64(val); + } + } + } + } + + let bytes = Bytes::from_elems(output); + FlexTensor::new(bytes, Layout::contiguous(out_shape), T::dtype()) +} + +fn apply_padding( + px: f64, + py: f64, + w: usize, + h: usize, + mode: GridSamplePaddingMode, + align_corners: bool, +) -> (f64, f64) { + if !px.is_finite() || !py.is_finite() { + return match mode { + GridSamplePaddingMode::Border => { + let cx = ((w - 1) as f64 / 2.0).clamp(0.0, (w - 1) as f64); + let cy = ((h - 1) as f64 / 2.0).clamp(0.0, (h - 1) as f64); + (cx, cy) + } + _ => (px, py), + }; + } + + match mode { + GridSamplePaddingMode::Zeros => (px, py), + GridSamplePaddingMode::Border => { + let px = px.clamp(0.0, (w - 1) as f64); + let py = py.clamp(0.0, (h - 1) as f64); + (px, py) + } + GridSamplePaddingMode::Reflection => { + let px = reflect_coordinate(px, w, align_corners); + let py = reflect_coordinate(py, h, align_corners); + (px, py) + } + } +} + +fn reflect_coordinate(coord: f64, size: usize, align_corners: bool) -> f64 { + let size_f = size as f64; + let (min_val, max_val) = if align_corners { + (0.0, size_f - 1.0) + } else { + (-0.5, size_f - 0.5) + }; + + let span = max_val - min_val; + if span <= 0.0 { + return min_val; + } + + let period = 2.0 * span; + let x = (coord - min_val).abs(); + let x_mod = x - (x / period).floor() * period; + span - (x_mod - span).abs() + min_val +} + +// Correctness of grid_sample_2d (nearest/bilinear, zeros/border/reflection +// padding, align_corners on/off) is covered by the cross-backend suite in +// crates/burn-backend-tests/tests/tensor/float/ops/grid_sample.rs, which +// reaches `grid_sample_2d` via the `FloatTensorOps` trait for flex. No +// flex-specific tests remain here. diff --git a/crates/burn-flex/src/ops/int.rs b/crates/burn-flex/src/ops/int.rs new file mode 100644 index 0000000..9804190 --- /dev/null +++ b/crates/burn-flex/src/ops/int.rs @@ -0,0 +1,1336 @@ +//! Int tensor operations for the Flex backend. + +use alloc::vec::Vec; +use burn_backend::{ + DType, Distribution, ExecutionError, FloatDType, Scalar, TensorData, TensorMetadata, + ops::IntTensorOps, + tensor::{BoolTensor, Device, FloatTensor, IntTensor}, +}; +use burn_std::{Bytes, IntDType, Shape, Slice, bf16, f16}; +use num_traits::ToPrimitive; + +use crate::Layout; +use crate::ops::binary::{binary_op_typed, int_binary_op, int_scalar_op, scalar_op_typed}; +use crate::{Flex, FlexTensor, ops::matmul}; + +/// Convert a Scalar to (i64, u64) pair for the given dtype. +/// Only the matching type's conversion is validated; the other gets a dummy 0. +fn scalar_to_int_pair(dtype: DType, rhs: &Scalar) -> (i64, u64) { + if dtype == DType::U64 { + (0, rhs.to_u64().unwrap()) + } else { + (rhs.to_i64().unwrap(), 0) + } +} + +impl IntTensorOps for Flex { + fn int_from_data(data: TensorData, _device: &Device) -> IntTensor { + FlexTensor::from_data(data) + } + + async fn int_into_data(tensor: IntTensor) -> Result { + Ok(tensor.into_data()) + } + + fn int_to_device(tensor: IntTensor, _device: &Device) -> IntTensor { + tensor + } + + fn int_cat(tensors: Vec>, dim: usize) -> IntTensor { + crate::ops::cat::cat(tensors, dim) + } + + fn int_reshape(tensor: IntTensor, shape: Shape) -> IntTensor { + tensor.reshape(shape) + } + + fn int_slice(tensor: IntTensor, slices: &[Slice]) -> IntTensor { + crate::ops::slice::slice(tensor, slices) + } + + fn int_empty(shape: Shape, _device: &Device, dtype: IntDType) -> IntTensor { + FlexTensor::empty(shape, dtype.into()) + } + + fn int_mask_where( + tensor: IntTensor, + mask: BoolTensor, + value: IntTensor, + ) -> IntTensor { + debug_assert_eq!( + tensor.dtype(), + value.dtype(), + "int_mask_where: dtype mismatch" + ); + match tensor.dtype() { + DType::I64 => crate::ops::mask::mask_where::(tensor, mask, value), + DType::I32 => crate::ops::mask::mask_where::(tensor, mask, value), + DType::I16 => crate::ops::mask::mask_where::(tensor, mask, value), + DType::I8 => crate::ops::mask::mask_where::(tensor, mask, value), + DType::U64 => crate::ops::mask::mask_where::(tensor, mask, value), + DType::U32 => crate::ops::mask::mask_where::(tensor, mask, value), + DType::U16 => crate::ops::mask::mask_where::(tensor, mask, value), + DType::U8 => crate::ops::mask::mask_where::(tensor, mask, value), + dt => panic!("int_mask_where: unsupported dtype {:?}", dt), + } + } + + fn int_mask_fill( + tensor: IntTensor, + mask: BoolTensor, + value: Scalar, + ) -> IntTensor { + match tensor.dtype() { + DType::I64 => crate::ops::mask::mask_fill(tensor, mask, value.to_i64().unwrap()), + DType::I32 => crate::ops::mask::mask_fill(tensor, mask, value.to_i64().unwrap() as i32), + DType::I16 => crate::ops::mask::mask_fill(tensor, mask, value.to_i64().unwrap() as i16), + DType::I8 => crate::ops::mask::mask_fill(tensor, mask, value.to_i64().unwrap() as i8), + DType::U64 => crate::ops::mask::mask_fill(tensor, mask, value.to_u64().unwrap()), + DType::U32 => crate::ops::mask::mask_fill(tensor, mask, value.to_u64().unwrap() as u32), + DType::U16 => crate::ops::mask::mask_fill(tensor, mask, value.to_u64().unwrap() as u16), + DType::U8 => crate::ops::mask::mask_fill(tensor, mask, value.to_u64().unwrap() as u8), + dt => panic!("int_mask_fill: unsupported dtype {:?}", dt), + } + } + + fn int_slice_assign( + tensor: IntTensor, + slices: &[Slice], + value: IntTensor, + ) -> IntTensor { + crate::ops::slice::slice_assign(tensor, slices, value) + } + + /// Gather ints along `dim` at the given indices. + /// + /// The `tensor` dispatches on its own int dtype (I8/I16/I32/I64 signed or + /// U8/U16/U32/U64 unsigned). The `indices` tensor may be any of those + /// widths too - it's normalised to `isize` by the shared `read_indices` + /// helper in `ops::gather_scatter` before the kernel runs, so callers are + /// not required to pre-convert to I64. + fn int_gather( + dim: usize, + tensor: IntTensor, + indices: IntTensor, + ) -> IntTensor { + match tensor.dtype() { + DType::I64 => crate::ops::gather_scatter::gather::(tensor, dim, indices), + DType::I32 => crate::ops::gather_scatter::gather::(tensor, dim, indices), + DType::I16 => crate::ops::gather_scatter::gather::(tensor, dim, indices), + DType::I8 => crate::ops::gather_scatter::gather::(tensor, dim, indices), + DType::U64 => crate::ops::gather_scatter::gather::(tensor, dim, indices), + DType::U32 => crate::ops::gather_scatter::gather::(tensor, dim, indices), + DType::U16 => crate::ops::gather_scatter::gather::(tensor, dim, indices), + DType::U8 => crate::ops::gather_scatter::gather::(tensor, dim, indices), + dt => panic!("int_gather: unsupported dtype {:?}", dt), + } + } + + /// Scatter-add int values at the given indices along `dim`. + /// + /// `tensor` and `value` must share the same int dtype; `indices` may be + /// any supported int width. See [`int_gather`](Self::int_gather) for the + /// full index-width policy. + fn int_scatter_add( + dim: usize, + tensor: IntTensor, + indices: IntTensor, + value: IntTensor, + ) -> IntTensor { + debug_assert_eq!( + tensor.dtype(), + value.dtype(), + "int_scatter_add: dtype mismatch" + ); + match tensor.dtype() { + DType::I64 => { + crate::ops::gather_scatter::scatter_add::(tensor, dim, indices, value) + } + DType::I32 => { + crate::ops::gather_scatter::scatter_add::(tensor, dim, indices, value) + } + DType::I16 => { + crate::ops::gather_scatter::scatter_add::(tensor, dim, indices, value) + } + DType::I8 => crate::ops::gather_scatter::scatter_add::(tensor, dim, indices, value), + DType::U64 => { + crate::ops::gather_scatter::scatter_add::(tensor, dim, indices, value) + } + DType::U32 => { + crate::ops::gather_scatter::scatter_add::(tensor, dim, indices, value) + } + DType::U16 => { + crate::ops::gather_scatter::scatter_add::(tensor, dim, indices, value) + } + DType::U8 => crate::ops::gather_scatter::scatter_add::(tensor, dim, indices, value), + dt => panic!("int_scatter_add: unsupported dtype {:?}", dt), + } + } + + fn int_scatter_nd( + data: IntTensor, + indices: IntTensor, + values: IntTensor, + reduction: burn_backend::tensor::IndexingUpdateOp, + ) -> IntTensor { + match data.dtype() { + DType::I64 => { + crate::ops::gather_scatter::scatter_nd::(data, indices, values, reduction) + } + DType::I32 => { + crate::ops::gather_scatter::scatter_nd::(data, indices, values, reduction) + } + DType::I16 => { + crate::ops::gather_scatter::scatter_nd::(data, indices, values, reduction) + } + DType::I8 => { + crate::ops::gather_scatter::scatter_nd::(data, indices, values, reduction) + } + DType::U64 => { + crate::ops::gather_scatter::scatter_nd::(data, indices, values, reduction) + } + DType::U32 => { + crate::ops::gather_scatter::scatter_nd::(data, indices, values, reduction) + } + DType::U16 => { + crate::ops::gather_scatter::scatter_nd::(data, indices, values, reduction) + } + DType::U8 => { + crate::ops::gather_scatter::scatter_nd::(data, indices, values, reduction) + } + dt => panic!("int_scatter_nd: unsupported dtype {:?}", dt), + } + } + + fn int_gather_nd(data: IntTensor, indices: IntTensor) -> IntTensor { + match data.dtype() { + DType::I64 => crate::ops::gather_scatter::gather_nd::(data, indices), + DType::I32 => crate::ops::gather_scatter::gather_nd::(data, indices), + DType::I16 => crate::ops::gather_scatter::gather_nd::(data, indices), + DType::I8 => crate::ops::gather_scatter::gather_nd::(data, indices), + DType::U64 => crate::ops::gather_scatter::gather_nd::(data, indices), + DType::U32 => crate::ops::gather_scatter::gather_nd::(data, indices), + DType::U16 => crate::ops::gather_scatter::gather_nd::(data, indices), + DType::U8 => crate::ops::gather_scatter::gather_nd::(data, indices), + dt => panic!("int_gather_nd: unsupported dtype {:?}", dt), + } + } + + /// Select ints along `dim` by a 1D index tensor. + /// + /// The `indices` tensor may be any supported int width. See + /// [`int_gather`](Self::int_gather) for the full index-width policy. + fn int_select( + tensor: IntTensor, + dim: usize, + indices: IntTensor, + ) -> IntTensor { + match tensor.dtype() { + DType::I64 => crate::ops::gather_scatter::select::(tensor, dim, indices), + DType::I32 => crate::ops::gather_scatter::select::(tensor, dim, indices), + DType::I16 => crate::ops::gather_scatter::select::(tensor, dim, indices), + DType::I8 => crate::ops::gather_scatter::select::(tensor, dim, indices), + DType::U64 => crate::ops::gather_scatter::select::(tensor, dim, indices), + DType::U32 => crate::ops::gather_scatter::select::(tensor, dim, indices), + DType::U16 => crate::ops::gather_scatter::select::(tensor, dim, indices), + DType::U8 => crate::ops::gather_scatter::select::(tensor, dim, indices), + dt => panic!("int_select: unsupported dtype {:?}", dt), + } + } + + /// Select-add int values at a 1D index tensor along `dim`. + /// + /// `tensor` and `value` must share the same int dtype; `indices` may be + /// any supported int width. See [`int_gather`](Self::int_gather) for the + /// full index-width policy. + fn int_select_add( + tensor: IntTensor, + dim: usize, + indices: IntTensor, + value: IntTensor, + ) -> IntTensor { + debug_assert_eq!( + tensor.dtype(), + value.dtype(), + "int_select_add: dtype mismatch" + ); + match tensor.dtype() { + DType::I64 => { + crate::ops::gather_scatter::select_add::(tensor, dim, indices, value) + } + DType::I32 => { + crate::ops::gather_scatter::select_add::(tensor, dim, indices, value) + } + DType::I16 => { + crate::ops::gather_scatter::select_add::(tensor, dim, indices, value) + } + DType::I8 => crate::ops::gather_scatter::select_add::(tensor, dim, indices, value), + DType::U64 => { + crate::ops::gather_scatter::select_add::(tensor, dim, indices, value) + } + DType::U32 => { + crate::ops::gather_scatter::select_add::(tensor, dim, indices, value) + } + DType::U16 => { + crate::ops::gather_scatter::select_add::(tensor, dim, indices, value) + } + DType::U8 => crate::ops::gather_scatter::select_add::(tensor, dim, indices, value), + dt => panic!("int_select_add: unsupported dtype {:?}", dt), + } + } + + fn int_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::int_equal(lhs, rhs, out_dtype) + } + + fn int_equal_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + let (i, u) = scalar_to_int_pair(lhs.dtype(), &rhs); + crate::ops::comparison::int_equal_elem(lhs, i, u, out_dtype) + } + + fn int_greater( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::int_greater(lhs, rhs, out_dtype) + } + + fn int_greater_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + let (i, u) = scalar_to_int_pair(lhs.dtype(), &rhs); + crate::ops::comparison::int_greater_elem(lhs, i, u, out_dtype) + } + + fn int_greater_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::int_greater_equal(lhs, rhs, out_dtype) + } + + fn int_greater_equal_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + let (i, u) = scalar_to_int_pair(lhs.dtype(), &rhs); + crate::ops::comparison::int_greater_equal_elem(lhs, i, u, out_dtype) + } + + fn int_lower( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::int_lower(lhs, rhs, out_dtype) + } + + fn int_lower_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + let (i, u) = scalar_to_int_pair(lhs.dtype(), &rhs); + crate::ops::comparison::int_lower_elem(lhs, i, u, out_dtype) + } + + fn int_lower_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::int_lower_equal(lhs, rhs, out_dtype) + } + + fn int_lower_equal_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + let (i, u) = scalar_to_int_pair(lhs.dtype(), &rhs); + crate::ops::comparison::int_lower_equal_elem(lhs, i, u, out_dtype) + } + + fn int_add(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + int_binary_op(lhs, rhs, |a, b| a + b) + } + + fn int_add_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + if lhs.dtype() == DType::U64 { + return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| { + a.wrapping_add(b) + }); + } + int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a + b) + } + + fn int_sub(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + int_binary_op(lhs, rhs, |a, b| a - b) + } + + fn int_sub_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + if lhs.dtype() == DType::U64 { + return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| { + a.wrapping_sub(b) + }); + } + int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a - b) + } + + fn int_mul(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + int_binary_op(lhs, rhs, |a, b| a * b) + } + + fn int_mul_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + if lhs.dtype() == DType::U64 { + return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| { + a.wrapping_mul(b) + }); + } + int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a * b) + } + + fn int_div(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + // U64 values > i64::MAX produce wrong results through i64 cast + if lhs.dtype() == DType::U64 { + let (lhs, rhs) = crate::ops::expand::broadcast_binary(lhs, rhs); + return binary_op_typed(lhs, &rhs, |a: u64, b: u64| a / b); + } + int_binary_op(lhs, rhs, |a, b| a / b) + } + + fn int_div_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + if lhs.dtype() == DType::U64 { + return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| a / b); + } + int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a / b) + } + + fn int_remainder(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + // U64 values > i64::MAX produce wrong results through i64 cast + if lhs.dtype() == DType::U64 { + let (lhs, rhs) = crate::ops::expand::broadcast_binary(lhs, rhs); + return binary_op_typed(lhs, &rhs, |a: u64, b: u64| a % b); + } + // Python/PyTorch-style remainder: result has same sign as divisor + int_binary_op(lhs, rhs, |a, b| ((a % b) + b) % b) + } + + fn int_remainder_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + if lhs.dtype() == DType::U64 { + return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| a % b); + } + // Python/PyTorch-style remainder: result has same sign as divisor + int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| ((a % b) + b) % b) + } + + // Precision limits: i64/u64 > 2^24 for f32/f16/bf16, > 2^53 for f64. + fn int_into_float( + tensor: IntTensor, + out_dtype: burn_std::FloatDType, + ) -> FloatTensor { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let src = tensor.dtype(); + let out_dt = DType::from(out_dtype); + + // Read source ints, applying conversion per-element. + // Each arm binds `$x` to the native int value; `$conv` must work for all int types. + macro_rules! read_ints { + (|$x:ident| $conv:expr) => { + match src { + DType::I64 => tensor.storage::().iter().map(|&$x| $conv).collect(), + DType::I32 => tensor.storage::().iter().map(|&$x| $conv).collect(), + DType::I16 => tensor.storage::().iter().map(|&$x| $conv).collect(), + DType::I8 => tensor.storage::().iter().map(|&$x| $conv).collect(), + DType::U64 => tensor.storage::().iter().map(|&$x| $conv).collect(), + DType::U32 => tensor.storage::().iter().map(|&$x| $conv).collect(), + DType::U16 => tensor.storage::().iter().map(|&$x| $conv).collect(), + DType::U8 => tensor.storage::().iter().map(|&$x| $conv).collect(), + _ => panic!("int_into_float: unsupported source dtype {:?}", src), + } + }; + } + + match out_dtype { + FloatDType::F64 => { + let data: Vec = read_ints!(|x| x as f64); + FlexTensor::new(Bytes::from_elems(data), Layout::contiguous(shape), out_dt) + } + FloatDType::F32 | FloatDType::Flex32 => { + let data: Vec = read_ints!(|x| x as f32); + FlexTensor::new(Bytes::from_elems(data), Layout::contiguous(shape), out_dt) + } + FloatDType::F16 => { + let data: Vec = read_ints!(|x| f16::from_f32(x as f32)); + FlexTensor::new(Bytes::from_elems(data), Layout::contiguous(shape), out_dt) + } + FloatDType::BF16 => { + let data: Vec = read_ints!(|x| bf16::from_f32(x as f32)); + FlexTensor::new(Bytes::from_elems(data), Layout::contiguous(shape), out_dt) + } + } + } + + fn int_swap_dims(tensor: IntTensor, dim1: usize, dim2: usize) -> IntTensor { + tensor.transpose(dim1, dim2) + } + + fn int_permute(tensor: IntTensor, axes: &[usize]) -> IntTensor { + tensor.permute(axes) + } + + fn int_flip(tensor: IntTensor, axes: &[usize]) -> IntTensor { + crate::ops::flip::flip(tensor, axes) + } + + fn int_random( + shape: Shape, + distribution: Distribution, + _device: &Device, + dtype: IntDType, + ) -> IntTensor { + let mut seed = crate::backend::SEED.lock().unwrap(); + let mut rng = seed.take().unwrap_or_else(crate::backend::get_seeded_rng); + let data = match dtype { + IntDType::I64 => TensorData::random::(shape, distribution, &mut rng), + IntDType::I32 => TensorData::random::(shape, distribution, &mut rng), + IntDType::I16 => TensorData::random::(shape, distribution, &mut rng), + IntDType::I8 => TensorData::random::(shape, distribution, &mut rng), + IntDType::U64 => TensorData::random::(shape, distribution, &mut rng), + IntDType::U32 => TensorData::random::(shape, distribution, &mut rng), + IntDType::U16 => TensorData::random::(shape, distribution, &mut rng), + IntDType::U8 => TensorData::random::(shape, distribution, &mut rng), + }; + *seed = Some(rng); + FlexTensor::from_data(data) + } + + fn int_expand(tensor: IntTensor, shape: Shape) -> IntTensor { + crate::ops::expand::expand(tensor, shape) + } + + fn int_matmul(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + matmul::int_matmul(lhs, rhs) + } + + fn int_sum(tensor: IntTensor) -> IntTensor { + crate::ops::reduce::sum(tensor) + } + + fn int_sum_dim(tensor: IntTensor, dim: usize) -> IntTensor { + crate::ops::reduce::sum_dim(tensor, dim) + } + + fn int_prod(tensor: IntTensor) -> IntTensor { + crate::ops::reduce::prod(tensor) + } + + fn int_prod_dim(tensor: IntTensor, dim: usize) -> IntTensor { + crate::ops::reduce::prod_dim(tensor, dim) + } + + fn int_mean_dim(tensor: IntTensor, dim: usize) -> IntTensor { + crate::ops::reduce::mean_dim(tensor, dim) + } + + fn int_cumsum(tensor: IntTensor, dim: usize) -> IntTensor { + match tensor.dtype() { + DType::I64 => crate::ops::cumulative::cumsum::(tensor, dim), + DType::I32 => crate::ops::cumulative::cumsum::(tensor, dim), + DType::I16 => crate::ops::cumulative::cumsum::(tensor, dim), + DType::I8 => crate::ops::cumulative::cumsum::(tensor, dim), + DType::U64 => crate::ops::cumulative::cumsum::(tensor, dim), + DType::U32 => crate::ops::cumulative::cumsum::(tensor, dim), + DType::U16 => crate::ops::cumulative::cumsum::(tensor, dim), + DType::U8 => crate::ops::cumulative::cumsum::(tensor, dim), + dt => panic!("int_cumsum: unsupported dtype {:?}", dt), + } + } + + fn int_cumprod(tensor: IntTensor, dim: usize) -> IntTensor { + match tensor.dtype() { + DType::I64 => crate::ops::cumulative::cumprod::(tensor, dim), + DType::I32 => crate::ops::cumulative::cumprod::(tensor, dim), + DType::I16 => crate::ops::cumulative::cumprod::(tensor, dim), + DType::I8 => crate::ops::cumulative::cumprod::(tensor, dim), + DType::U64 => crate::ops::cumulative::cumprod::(tensor, dim), + DType::U32 => crate::ops::cumulative::cumprod::(tensor, dim), + DType::U16 => crate::ops::cumulative::cumprod::(tensor, dim), + DType::U8 => crate::ops::cumulative::cumprod::(tensor, dim), + dt => panic!("int_cumprod: unsupported dtype {:?}", dt), + } + } + + fn int_cummin(tensor: IntTensor, dim: usize) -> IntTensor { + match tensor.dtype() { + DType::I64 => crate::ops::cumulative::cummin::(tensor, dim), + DType::I32 => crate::ops::cumulative::cummin::(tensor, dim), + DType::I16 => crate::ops::cumulative::cummin::(tensor, dim), + DType::I8 => crate::ops::cumulative::cummin::(tensor, dim), + DType::U64 => crate::ops::cumulative::cummin::(tensor, dim), + DType::U32 => crate::ops::cumulative::cummin::(tensor, dim), + DType::U16 => crate::ops::cumulative::cummin::(tensor, dim), + DType::U8 => crate::ops::cumulative::cummin::(tensor, dim), + dt => panic!("int_cummin: unsupported dtype {:?}", dt), + } + } + + fn int_cummax(tensor: IntTensor, dim: usize) -> IntTensor { + match tensor.dtype() { + DType::I64 => crate::ops::cumulative::cummax::(tensor, dim), + DType::I32 => crate::ops::cumulative::cummax::(tensor, dim), + DType::I16 => crate::ops::cumulative::cummax::(tensor, dim), + DType::I8 => crate::ops::cumulative::cummax::(tensor, dim), + DType::U64 => crate::ops::cumulative::cummax::(tensor, dim), + DType::U32 => crate::ops::cumulative::cummax::(tensor, dim), + DType::U16 => crate::ops::cumulative::cummax::(tensor, dim), + DType::U8 => crate::ops::cumulative::cummax::(tensor, dim), + dt => panic!("int_cummax: unsupported dtype {:?}", dt), + } + } + + fn int_argmax(tensor: IntTensor, dim: usize) -> IntTensor { + crate::ops::reduce::argmax(tensor, dim) + } + + fn int_argtopk(_tensor: IntTensor, _dim: usize, _k: usize) -> IntTensor { + panic!("argtopk not implemented for flex") + } + + fn int_argmin(tensor: IntTensor, dim: usize) -> IntTensor { + crate::ops::reduce::argmin(tensor, dim) + } + + fn int_abs(tensor: IntTensor) -> IntTensor { + crate::ops::unary::int_abs(tensor) + } + + fn bitwise_and(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + int_binary_op(lhs, rhs, |a, b| a & b) + } + + fn bitwise_and_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + if lhs.dtype() == DType::U64 { + return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| a & b); + } + int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a & b) + } + + fn bitwise_or(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + int_binary_op(lhs, rhs, |a, b| a | b) + } + + fn bitwise_or_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + if lhs.dtype() == DType::U64 { + return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| a | b); + } + int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a | b) + } + + fn bitwise_xor(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + int_binary_op(lhs, rhs, |a, b| a ^ b) + } + + fn bitwise_xor_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + if lhs.dtype() == DType::U64 { + return scalar_op_typed(lhs, rhs.to_u64().unwrap(), |a: u64, b: u64| a ^ b); + } + int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a ^ b) + } + + fn bitwise_not(tensor: IntTensor) -> IntTensor { + // Use scalar op with dummy value, only applying NOT to lhs + int_scalar_op(tensor, 0, |a, _| !a) + } + + // Shift amounts masked to type width via wrapping_shl/wrapping_shr. + fn bitwise_left_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + int_binary_op(lhs, rhs, |a, b| a.wrapping_shl(b as u32)) + } + + fn bitwise_left_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a.wrapping_shl(b as u32)) + } + + fn bitwise_right_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + int_binary_op(lhs, rhs, |a, b| a.wrapping_shr(b as u32)) + } + + fn bitwise_right_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + int_scalar_op(lhs, rhs.to_i64().unwrap(), |a, b| a.wrapping_shr(b as u32)) + } + + fn int_cast(tensor: IntTensor, dtype: IntDType) -> IntTensor { + let target_dtype: DType = dtype.into(); + + // If already the target dtype, return as-is + if tensor.dtype() == target_dtype { + return tensor; + } + + // Make contiguous for easier iteration + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + + // Helper macro to convert between types + macro_rules! cast_impl { + ($src_type:ty, $dst_type:ty, $dst_dtype:expr) => {{ + let src: &[$src_type] = tensor.storage(); + let dst: Vec<$dst_type> = src.iter().map(|&x| x as $dst_type).collect(); + FlexTensor::new( + Bytes::from_elems(dst), + Layout::contiguous(shape), + $dst_dtype, + ) + }}; + } + + // Match source dtype to target dtype + match (tensor.dtype(), target_dtype) { + // From I64 + (DType::I64, DType::I32) => cast_impl!(i64, i32, DType::I32), + (DType::I64, DType::I16) => cast_impl!(i64, i16, DType::I16), + (DType::I64, DType::I8) => cast_impl!(i64, i8, DType::I8), + (DType::I64, DType::U64) => cast_impl!(i64, u64, DType::U64), + (DType::I64, DType::U32) => cast_impl!(i64, u32, DType::U32), + (DType::I64, DType::U16) => cast_impl!(i64, u16, DType::U16), + (DType::I64, DType::U8) => cast_impl!(i64, u8, DType::U8), + + // From I32 + (DType::I32, DType::I64) => cast_impl!(i32, i64, DType::I64), + (DType::I32, DType::I16) => cast_impl!(i32, i16, DType::I16), + (DType::I32, DType::I8) => cast_impl!(i32, i8, DType::I8), + (DType::I32, DType::U64) => cast_impl!(i32, u64, DType::U64), + (DType::I32, DType::U32) => cast_impl!(i32, u32, DType::U32), + (DType::I32, DType::U16) => cast_impl!(i32, u16, DType::U16), + (DType::I32, DType::U8) => cast_impl!(i32, u8, DType::U8), + + // From I16 + (DType::I16, DType::I64) => cast_impl!(i16, i64, DType::I64), + (DType::I16, DType::I32) => cast_impl!(i16, i32, DType::I32), + (DType::I16, DType::I8) => cast_impl!(i16, i8, DType::I8), + (DType::I16, DType::U64) => cast_impl!(i16, u64, DType::U64), + (DType::I16, DType::U32) => cast_impl!(i16, u32, DType::U32), + (DType::I16, DType::U16) => cast_impl!(i16, u16, DType::U16), + (DType::I16, DType::U8) => cast_impl!(i16, u8, DType::U8), + + // From I8 + (DType::I8, DType::I64) => cast_impl!(i8, i64, DType::I64), + (DType::I8, DType::I32) => cast_impl!(i8, i32, DType::I32), + (DType::I8, DType::I16) => cast_impl!(i8, i16, DType::I16), + (DType::I8, DType::U64) => cast_impl!(i8, u64, DType::U64), + (DType::I8, DType::U32) => cast_impl!(i8, u32, DType::U32), + (DType::I8, DType::U16) => cast_impl!(i8, u16, DType::U16), + (DType::I8, DType::U8) => cast_impl!(i8, u8, DType::U8), + + // From U64 + (DType::U64, DType::I64) => cast_impl!(u64, i64, DType::I64), + (DType::U64, DType::I32) => cast_impl!(u64, i32, DType::I32), + (DType::U64, DType::I16) => cast_impl!(u64, i16, DType::I16), + (DType::U64, DType::I8) => cast_impl!(u64, i8, DType::I8), + (DType::U64, DType::U32) => cast_impl!(u64, u32, DType::U32), + (DType::U64, DType::U16) => cast_impl!(u64, u16, DType::U16), + (DType::U64, DType::U8) => cast_impl!(u64, u8, DType::U8), + + // From U32 + (DType::U32, DType::I64) => cast_impl!(u32, i64, DType::I64), + (DType::U32, DType::I32) => cast_impl!(u32, i32, DType::I32), + (DType::U32, DType::I16) => cast_impl!(u32, i16, DType::I16), + (DType::U32, DType::I8) => cast_impl!(u32, i8, DType::I8), + (DType::U32, DType::U64) => cast_impl!(u32, u64, DType::U64), + (DType::U32, DType::U16) => cast_impl!(u32, u16, DType::U16), + (DType::U32, DType::U8) => cast_impl!(u32, u8, DType::U8), + + // From U16 + (DType::U16, DType::I64) => cast_impl!(u16, i64, DType::I64), + (DType::U16, DType::I32) => cast_impl!(u16, i32, DType::I32), + (DType::U16, DType::I16) => cast_impl!(u16, i16, DType::I16), + (DType::U16, DType::I8) => cast_impl!(u16, i8, DType::I8), + (DType::U16, DType::U64) => cast_impl!(u16, u64, DType::U64), + (DType::U16, DType::U32) => cast_impl!(u16, u32, DType::U32), + (DType::U16, DType::U8) => cast_impl!(u16, u8, DType::U8), + + // From U8 + (DType::U8, DType::I64) => cast_impl!(u8, i64, DType::I64), + (DType::U8, DType::I32) => cast_impl!(u8, i32, DType::I32), + (DType::U8, DType::I16) => cast_impl!(u8, i16, DType::I16), + (DType::U8, DType::I8) => cast_impl!(u8, i8, DType::I8), + (DType::U8, DType::U64) => cast_impl!(u8, u64, DType::U64), + (DType::U8, DType::U32) => cast_impl!(u8, u32, DType::U32), + (DType::U8, DType::U16) => cast_impl!(u8, u16, DType::U16), + + _ => panic!( + "int_cast: unsupported conversion from {:?} to {:?}", + tensor.dtype(), + target_dtype + ), + } + } + + fn int_unfold( + tensor: IntTensor, + dim: usize, + size: usize, + step: usize, + ) -> IntTensor { + crate::ops::unfold::unfold_int(tensor, dim, size, step) + } + + fn int_neg(tensor: IntTensor) -> IntTensor { + int_scalar_op(tensor, 0i64, |a, _| a.wrapping_neg()) + } + + fn int_clamp(tensor: IntTensor, min: Scalar, max: Scalar) -> IntTensor { + if tensor.dtype() == DType::U64 { + let min_val = min.to_u64().unwrap(); + let max_val = max.to_u64().unwrap(); + return scalar_op_typed(tensor, 0u64, move |x: u64, _| x.clamp(min_val, max_val)); + } + let min_val = min.to_i64().unwrap(); + let max_val = max.to_i64().unwrap(); + int_scalar_op(tensor, 0i64, move |x, _| x.clamp(min_val, max_val)) + } + + fn int_clamp_min(tensor: IntTensor, min: Scalar) -> IntTensor { + if tensor.dtype() == DType::U64 { + let min_val = min.to_u64().unwrap(); + return scalar_op_typed(tensor, 0u64, move |x: u64, _| x.max(min_val)); + } + let min_val = min.to_i64().unwrap(); + int_scalar_op(tensor, 0i64, move |x, _| x.max(min_val)) + } + + fn int_clamp_max(tensor: IntTensor, max: Scalar) -> IntTensor { + if tensor.dtype() == DType::U64 { + let max_val = max.to_u64().unwrap(); + return scalar_op_typed(tensor, 0u64, move |x: u64, _| x.min(max_val)); + } + let max_val = max.to_i64().unwrap(); + int_scalar_op(tensor, 0i64, move |x, _| x.min(max_val)) + } + + fn int_sign(tensor: IntTensor) -> IntTensor { + if tensor.dtype() == DType::U64 { + return scalar_op_typed(tensor, 0u64, |x: u64, _| if x > 0 { 1 } else { 0 }); + } + int_scalar_op(tensor, 0i64, |x, _| { + if x > 0 { + 1 + } else if x < 0 { + -1 + } else { + 0 + } + }) + } + + fn int_mean(tensor: IntTensor) -> IntTensor { + let n = tensor.layout().num_elements(); + assert!(n > 0, "int_mean: cannot take mean of empty tensor"); + let dtype = tensor.dtype(); + let sum_result = crate::ops::reduce::sum(tensor); + // Compute in i64 to avoid truncation of n for small int types + macro_rules! compute_mean { + ($ty:ty) => {{ + let data: &[$ty] = sum_result.storage(); + let mean_val = (data[0] as i64 / n as i64) as $ty; + FlexTensor::new( + Bytes::from_elems(alloc::vec![mean_val]), + Layout::contiguous(Shape::from(alloc::vec![1])), + dtype, + ) + }}; + } + match dtype { + DType::I64 => compute_mean!(i64), + DType::I32 => compute_mean!(i32), + DType::I16 => compute_mean!(i16), + DType::I8 => compute_mean!(i8), + other => panic!("int_mean: unsupported dtype {:?}", other), + } + } + + fn int_max(tensor: IntTensor) -> IntTensor { + crate::ops::reduce::max(tensor) + } + + fn int_max_dim(tensor: IntTensor, dim: usize) -> IntTensor { + crate::ops::reduce::max_dim(tensor, dim) + } + + fn int_min(tensor: IntTensor) -> IntTensor { + crate::ops::reduce::min(tensor) + } + + fn int_min_dim(tensor: IntTensor, dim: usize) -> IntTensor { + crate::ops::reduce::min_dim(tensor, dim) + } + + fn int_max_dim_with_indices( + tensor: IntTensor, + dim: usize, + ) -> (IntTensor, IntTensor) { + crate::ops::reduce::max_dim_with_indices(tensor, dim) + } + + fn int_min_dim_with_indices( + tensor: IntTensor, + dim: usize, + ) -> (IntTensor, IntTensor) { + crate::ops::reduce::min_dim_with_indices(tensor, dim) + } + + fn int_any(tensor: IntTensor, out_dtype: burn_std::BoolDType) -> BoolTensor { + crate::ops::comparison::any_int(tensor, out_dtype) + } + + fn int_any_dim( + tensor: IntTensor, + dim: usize, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::any_int_dim(tensor, dim, out_dtype) + } + + fn int_all(tensor: IntTensor, out_dtype: burn_std::BoolDType) -> BoolTensor { + crate::ops::comparison::all_int(tensor, out_dtype) + } + + fn int_all_dim( + tensor: IntTensor, + dim: usize, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::all_int_dim(tensor, dim, out_dtype) + } + + fn int_powi(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + int_binary_op(lhs, rhs, |a, b| a.wrapping_pow(b as u32)) + } + + fn int_zeros(shape: Shape, _device: &Device, dtype: IntDType) -> IntTensor { + FlexTensor::zeros(shape, dtype.into()) + } + + fn int_ones(shape: Shape, _device: &Device, dtype: IntDType) -> IntTensor { + let dt: DType = dtype.into(); + match dt { + DType::I64 => FlexTensor::filled_typed(shape, dt, 1i64), + DType::I32 => FlexTensor::filled_typed(shape, dt, 1i32), + DType::I16 => FlexTensor::filled_typed(shape, dt, 1i16), + DType::I8 => FlexTensor::filled_typed(shape, dt, 1i8), + DType::U64 => FlexTensor::filled_typed(shape, dt, 1u64), + DType::U32 => FlexTensor::filled_typed(shape, dt, 1u32), + DType::U16 => FlexTensor::filled_typed(shape, dt, 1u16), + DType::U8 => FlexTensor::filled_typed(shape, dt, 1u8), + _ => unreachable!(), + } + } + + fn int_full( + shape: Shape, + fill_value: burn_backend::Scalar, + _device: &Device, + dtype: IntDType, + ) -> IntTensor { + let dt: DType = dtype.into(); + let v = fill_value.to_i64().unwrap(); + match dt { + DType::I64 => FlexTensor::filled_typed(shape, dt, v), + DType::I32 => FlexTensor::filled_typed(shape, dt, v as i32), + DType::I16 => FlexTensor::filled_typed(shape, dt, v as i16), + DType::I8 => FlexTensor::filled_typed(shape, dt, v as i8), + DType::U64 => FlexTensor::filled_typed(shape, dt, v as u64), + DType::U32 => FlexTensor::filled_typed(shape, dt, v as u32), + DType::U16 => FlexTensor::filled_typed(shape, dt, v as u16), + DType::U8 => FlexTensor::filled_typed(shape, dt, v as u8), + _ => unreachable!(), + } + } + + fn int_transpose(tensor: IntTensor) -> IntTensor { + let ndims = tensor.layout().num_dims(); + if ndims < 2 { + return tensor; + } + tensor.transpose(ndims - 2, ndims - 1) + } + + fn int_repeat_dim(tensor: IntTensor, dim: usize, times: usize) -> IntTensor { + crate::ops::repeat_dim::repeat_dim(tensor, dim, times) + } + + fn int_not_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + crate::ops::comparison::int_not_equal(lhs, rhs, out_dtype) + } + + fn int_not_equal_elem( + lhs: IntTensor, + rhs: burn_backend::Scalar, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + let (i, u) = scalar_to_int_pair(lhs.dtype(), &rhs); + crate::ops::comparison::int_not_equal_elem(lhs, i, u, out_dtype) + } + + fn int_sort(tensor: IntTensor, dim: usize, descending: bool) -> IntTensor { + crate::ops::sort::sort(tensor, dim, descending) + } + + fn int_sort_with_indices( + tensor: IntTensor, + dim: usize, + descending: bool, + ) -> (IntTensor, IntTensor) { + crate::ops::sort::sort_with_indices(tensor, dim, descending) + } + + fn int_argsort(tensor: IntTensor, dim: usize, descending: bool) -> IntTensor { + crate::ops::sort::argsort(tensor, dim, descending) + } + + fn int_powi_scalar(lhs: IntTensor, rhs: burn_backend::Scalar) -> IntTensor { + use num_traits::ToPrimitive; + match rhs.to_i64().unwrap() { + 0 => Self::int_ones(lhs.shape(), &Default::default(), lhs.dtype().into()), + 1 => lhs, + 2 => Self::int_mul(lhs.clone(), lhs), + _ => Self::int_powi_scalar_impl(lhs, rhs), + } + } + + fn int_powi_scalar_impl(lhs: IntTensor, rhs: burn_backend::Scalar) -> IntTensor { + use num_traits::ToPrimitive; + let exp = rhs.to_i64().unwrap() as u32; + if lhs.dtype() == DType::U64 { + return scalar_op_typed(lhs, exp as u64, move |x: u64, _| x.wrapping_pow(exp)); + } + int_scalar_op(lhs, exp as i64, move |x, _| x.wrapping_pow(exp)) + } + + fn int_max_abs(tensor: IntTensor) -> IntTensor { + let abs = Self::int_abs(tensor); + crate::ops::reduce::max(abs) + } + + fn int_max_abs_dim(tensor: IntTensor, dim: usize) -> IntTensor { + let abs = Self::int_abs(tensor); + crate::ops::reduce::max_dim(abs, dim) + } + + fn int_arange( + range: core::ops::Range, + _device: &Device, + dtype: IntDType, + ) -> IntTensor { + Self::int_arange_step(range, 1, &Default::default(), dtype) + } + + fn int_arange_step( + range: core::ops::Range, + step: usize, + _device: &Device, + dtype: IntDType, + ) -> IntTensor { + let dt: DType = dtype.into(); + + macro_rules! arange_typed { + ($ty:ty) => {{ + let data: Vec<$ty> = range.step_by(step).map(|v| v as $ty).collect(); + let shape = Shape::from(alloc::vec![data.len()]); + FlexTensor::new(Bytes::from_elems(data), Layout::contiguous(shape), dt) + }}; + } + + match dt { + DType::I64 => arange_typed!(i64), + DType::I32 => arange_typed!(i32), + DType::I16 => arange_typed!(i16), + DType::I8 => arange_typed!(i8), + DType::U64 => arange_typed!(u64), + DType::U32 => arange_typed!(u32), + DType::U16 => arange_typed!(u16), + DType::U8 => arange_typed!(u8), + _ => unreachable!(), + } + } +} + +// Tests kept here exercise flex-specific behavior: dtype storage +// selection for every int width (I16/I32/U8/U16/U32/I64/U64), and edge +// cases of the dtype-specific kernels (u64 wrap, i64::MIN abs/neg, bit +// shift at width). Plain int arithmetic, scalar ops, bool->int cast +// smokes, and negative-stride (flipped/transposed) variants have been +// migrated to burn-backend-tests so they run against every backend. +// When adding new tests, keep them here only if they probe flex dtype +// storage; otherwise add them to +// crates/burn-backend-tests/tests/tensor/int/ops/. +#[cfg(test)] +mod tests { + use alloc::vec; + use burn_backend::TensorData; + use burn_backend::ops::IntTensorOps; + + use crate::Flex; + use crate::FlexTensor; + + #[test] + fn test_u64_div_large_values() { + let a = FlexTensor::from_data(TensorData::new(vec![u64::MAX], [1])); + let b = FlexTensor::from_data(TensorData::new(vec![2u64], [1])); + let result = Flex::int_div(a, b); + let values: Vec = bytemuck::cast_slice(&result.into_data().bytes).to_vec(); + assert_eq!(values[0], u64::MAX / 2); + } + + #[test] + fn test_u64_remainder_large_values() { + let a = FlexTensor::from_data(TensorData::new(vec![u64::MAX], [1])); + let b = FlexTensor::from_data(TensorData::new(vec![2u64], [1])); + let result = Flex::int_remainder(a, b); + let values: Vec = bytemuck::cast_slice(&result.into_data().bytes).to_vec(); + assert_eq!(values[0], u64::MAX % 2); + } + + #[test] + fn test_int_abs_min_value() { + // i64::MIN.abs() panics in debug; wrapping_abs returns MIN (matches PyTorch) + let a = FlexTensor::from_data(TensorData::new(vec![i64::MIN], [1])); + let result = Flex::int_abs(a); + let values: Vec = bytemuck::cast_slice(&result.into_data().bytes).to_vec(); + assert_eq!(values[0], i64::MIN.wrapping_abs()); + } + + #[test] + fn test_int_neg_min_value() { + // i64::MIN negation panics in debug; wrapping_neg returns MIN (matches PyTorch) + let a = FlexTensor::from_data(TensorData::new(vec![i64::MIN], [1])); + let result = Flex::int_neg(a); + let values: Vec = bytemuck::cast_slice(&result.into_data().bytes).to_vec(); + assert_eq!(values[0], i64::MIN.wrapping_neg()); + } + + #[test] + fn test_int_shift_large_amount() { + // Shift by >= bit width panics without wrapping; should not crash + let a = FlexTensor::from_data(TensorData::new(vec![1i64], [1])); + let b = FlexTensor::from_data(TensorData::new(vec![64i64], [1])); + let _left = Flex::bitwise_left_shift(a.clone(), b.clone()); + let _right = Flex::bitwise_right_shift(a, b); + } + + #[test] + fn test_int_into_float_f64() { + use burn_backend::ops::IntTensorOps; + use burn_std::FloatDType; + + let t = FlexTensor::from_data(TensorData::new(vec![1i64, 2, -3], [3])); + let result = Flex::int_into_float(t, FloatDType::F64); + assert_eq!(result.dtype(), burn_backend::DType::F64); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![1.0f64, 2.0, -3.0]); + } + + #[test] + fn test_u64_add_scalar_large() { + let t = FlexTensor::from_data(TensorData::new(vec![1u64, 2, 3], [3])); + let big: u64 = (i64::MAX as u64) + 100; + let result = Flex::int_add_scalar(t, burn_backend::Scalar::from(big)); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![big + 1, big + 2, big + 3]); + } + + #[test] + fn test_u64_greater_elem_large() { + let big: u64 = (i64::MAX as u64) + 100; + let t = FlexTensor::from_data(TensorData::new(vec![big, big + 1, big - 1], [3])); + let result = Flex::int_greater_elem( + t, + burn_backend::Scalar::from(big), + burn_std::BoolStore::Native, + ); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![false, true, false]); + } + + #[test] + fn test_int_mask_fill_i32() { + let t = FlexTensor::from_data(TensorData::new(vec![1i32, 2, 3, 4], [4])); + let mask = FlexTensor::from_data(TensorData::new(vec![true, false, true, false], [4])); + let result = Flex::int_mask_fill(t, mask, burn_backend::Scalar::from(0i64)); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![0, 2, 0, 4]); + } + + #[test] + fn test_int_mask_fill_i16() { + let t = FlexTensor::from_data(TensorData::new(vec![10i16, 20, 30, 40], [4])); + let mask = FlexTensor::from_data(TensorData::new(vec![false, true, false, true], [4])); + let result = Flex::int_mask_fill(t, mask, burn_backend::Scalar::from(-1i64)); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![10, -1, 30, -1]); + } + + #[test] + fn test_int_mask_fill_u8() { + let t = FlexTensor::from_data(TensorData::new(vec![1u8, 2, 3, 4], [4])); + let mask = FlexTensor::from_data(TensorData::new(vec![true, true, false, false], [4])); + let result = Flex::int_mask_fill(t, mask, burn_backend::Scalar::from(255i64)); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![255, 255, 3, 4]); + } + + #[test] + fn test_int_mask_fill_u32() { + let t = FlexTensor::from_data(TensorData::new(vec![100u32, 200, 300], [3])); + let mask = FlexTensor::from_data(TensorData::new(vec![true, false, true], [3])); + let result = Flex::int_mask_fill(t, mask, burn_backend::Scalar::from(0i64)); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![0, 200, 0]); + } + + #[test] + fn test_int_mask_where_i32() { + let t = FlexTensor::from_data(TensorData::new(vec![1i32, 2, 3, 4], [4])); + let mask = FlexTensor::from_data(TensorData::new(vec![true, false, true, false], [4])); + let v = FlexTensor::from_data(TensorData::new(vec![10i32, 20, 30, 40], [4])); + let result = Flex::int_mask_where(t, mask, v); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![10, 2, 30, 4]); + } + + #[test] + fn test_int_mask_where_u8() { + let t = FlexTensor::from_data(TensorData::new(vec![1u8, 2, 3, 4], [4])); + let mask = FlexTensor::from_data(TensorData::new(vec![false, true, false, true], [4])); + let v = FlexTensor::from_data(TensorData::new(vec![10u8, 20, 30, 40], [4])); + let result = Flex::int_mask_where(t, mask, v); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![1, 20, 3, 40]); + } + + #[test] + fn test_int_gather_i32() { + let t = FlexTensor::from_data(TensorData::new(vec![10i32, 20, 30, 40, 50, 60], [2, 3])); + let indices = FlexTensor::from_data(TensorData::new(vec![2i64, 0, 1, 2], [2, 2])); + let result = Flex::int_gather(1, t, indices); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![30, 10, 50, 60]); + } + + #[test] + fn test_int_select_u16() { + let t = FlexTensor::from_data(TensorData::new(vec![10u16, 20, 30, 40, 50, 60], [2, 3])); + let indices = FlexTensor::from_data(TensorData::new(vec![0i64, 1], [2])); + let result = Flex::int_select(t, 1, indices); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![10, 20, 40, 50]); + } + + #[test] + fn test_int_cumsum_i32() { + let t = FlexTensor::from_data(TensorData::new(vec![1i32, 2, 3, 4], [4])); + let result = Flex::int_cumsum(t, 0); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![1, 3, 6, 10]); + } + + #[test] + fn test_int_cumprod_u8() { + let t = FlexTensor::from_data(TensorData::new(vec![1u8, 2, 3, 4], [4])); + let result = Flex::int_cumprod(t, 0); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![1, 2, 6, 24]); + } + + #[test] + fn test_int_cummin_i32() { + let t = FlexTensor::from_data(TensorData::new(vec![3i32, 1, 4, 1, 5], [5])); + let result = Flex::int_cummin(t, 0); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![3, 1, 1, 1, 1]); + } + + #[test] + fn test_int_cummax_u16() { + let t = FlexTensor::from_data(TensorData::new(vec![3u16, 1, 4, 1, 5], [5])); + let result = Flex::int_cummax(t, 0); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![3, 3, 4, 4, 5]); + } + + #[test] + fn test_int_scatter_add_i32() { + let t = FlexTensor::from_data(TensorData::new(vec![0i32, 0, 0], [1, 3])); + let indices = FlexTensor::from_data(TensorData::new(vec![0i64, 2, 1], [1, 3])); + let values = FlexTensor::from_data(TensorData::new(vec![10i32, 20, 30], [1, 3])); + let result = Flex::int_scatter_add(1, t, indices, values); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![10, 30, 20]); + } + + #[test] + fn test_int_select_add_u8() { + let t = FlexTensor::from_data(TensorData::new(vec![1u8, 2, 3], [3])); + let indices = FlexTensor::from_data(TensorData::new(vec![0i64, 2], [2])); + let values = FlexTensor::from_data(TensorData::new(vec![10u8, 20], [2])); + let result = Flex::int_select_add(t, 0, indices, values); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![11, 2, 23]); + } + + #[test] + fn test_int_random_i32() { + use burn_backend::{DType, Distribution, ops::IntTensorOps}; + use burn_std::{IntDType, Shape}; + + let shape = Shape::from(vec![100]); + let dist = Distribution::Uniform(0.0, 10.0); + let device = crate::FlexDevice; + let t = Flex::int_random(shape, dist, &device, IntDType::I32); + assert_eq!(t.dtype(), DType::I32); + let data: Vec = t.into_data().to_vec().unwrap(); + assert!(data.iter().all(|&v| (0..=10).contains(&v))); + } + + #[test] + fn test_int_random_u8() { + use burn_backend::{DType, Distribution, ops::IntTensorOps}; + use burn_std::{IntDType, Shape}; + + let shape = Shape::from(vec![50]); + let dist = Distribution::Uniform(0.0, 100.0); + let device = crate::FlexDevice; + let t = Flex::int_random(shape, dist, &device, IntDType::U8); + assert_eq!(t.dtype(), DType::U8); + } + + #[test] + fn test_int_mean_i32() { + use burn_backend::{DType, ops::IntTensorOps}; + + let t = FlexTensor::from_data(TensorData::new(vec![10i32, 20, 30], [3])); + let result = Flex::int_mean(t); + assert_eq!(result.dtype(), DType::I32); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![20]); // (10 + 20 + 30) / 3 = 20 + } +} diff --git a/crates/burn-flex/src/ops/interpolate.rs b/crates/burn-flex/src/ops/interpolate.rs new file mode 100644 index 0000000..d42a54f --- /dev/null +++ b/crates/burn-flex/src/ops/interpolate.rs @@ -0,0 +1,1187 @@ +//! Interpolation operations for image resizing. +//! +//! Supported modes: +//! - Nearest: Floor-based coordinate mapping (fastest) +//! - Bilinear: 4-point weighted average (good quality/speed tradeoff) +//! - Bicubic: 16-point cubic convolution (highest quality) +//! +//! Optimizations: +//! - Rayon parallelism over (batch, channel) pairs +//! - Precomputed coordinate mappings where beneficial +//! +//! Supported dtypes: f32, f64, f16 (native), bf16 (via f32 conversion) + +use alloc::vec; +use alloc::vec::Vec; +use burn_backend::DType; +use burn_std::{Bytes, Shape, bf16, f16}; +use num_traits::Float; + +use crate::{FlexTensor, Layout}; + +// ============================================================================ +// Macros for dtype wrappers +// ============================================================================ + +/// Generates an interpolation forward typed dispatcher. +macro_rules! interpolate_typed { + ($fn_name:ident, $impl_fn:ident, $T:ty) => { + pub fn $fn_name(x: FlexTensor, output_size: [usize; 2], align_corners: bool) -> FlexTensor { + $impl_fn::<$T>(x, output_size, align_corners) + } + }; +} + +/// Generates an interpolation bf16 forward wrapper via f32 conversion. +macro_rules! interpolate_bf16 { + ($bf16_fn:ident, $f32_fn:ident) => { + pub fn $bf16_fn(x: FlexTensor, output_size: [usize; 2], align_corners: bool) -> FlexTensor { + let x_f32 = convert_bf16_to_f32(&x); + let result_f32 = $f32_fn(x_f32, output_size, align_corners); + convert_f32_to_bf16(&result_f32) + } + }; +} + +/// Generates an interpolation backward typed dispatcher. +macro_rules! interpolate_backward_typed { + ($fn_name:ident, $impl_fn:ident, $T:ty) => { + pub fn $fn_name( + x: FlexTensor, + grad: FlexTensor, + output_size: [usize; 2], + align_corners: bool, + ) -> FlexTensor { + $impl_fn::<$T>(x, grad, output_size, align_corners) + } + }; +} + +/// Generates an interpolation bf16 backward wrapper via f32 conversion. +macro_rules! interpolate_backward_bf16 { + ($bf16_fn:ident, $f32_fn:ident) => { + pub fn $bf16_fn( + x: FlexTensor, + grad: FlexTensor, + output_size: [usize; 2], + align_corners: bool, + ) -> FlexTensor { + let x_f32 = convert_bf16_to_f32(&x); + let grad_f32 = convert_bf16_to_f32(&grad); + let result_f32 = $f32_fn(x_f32, grad_f32, output_size, align_corners); + convert_f32_to_bf16(&result_f32) + } + }; +} + +// ============================================================================ +// Public API - dtype dispatch +// ============================================================================ + +interpolate_typed!(interpolate_nearest_f32, interpolate_nearest_impl, f32); +interpolate_typed!(interpolate_nearest_f64, interpolate_nearest_impl, f64); +interpolate_typed!(interpolate_nearest_f16, interpolate_nearest_impl, f16); +interpolate_bf16!(interpolate_nearest_bf16, interpolate_nearest_f32); + +interpolate_typed!(interpolate_bilinear_f32, interpolate_bilinear_impl, f32); +interpolate_typed!(interpolate_bilinear_f64, interpolate_bilinear_impl, f64); +interpolate_typed!(interpolate_bilinear_f16, interpolate_bilinear_impl, f16); +interpolate_bf16!(interpolate_bilinear_bf16, interpolate_bilinear_f32); + +interpolate_typed!(interpolate_bicubic_f32, interpolate_bicubic_impl, f32); +interpolate_typed!(interpolate_bicubic_f64, interpolate_bicubic_impl, f64); +interpolate_typed!(interpolate_bicubic_f16, interpolate_bicubic_impl, f16); +interpolate_bf16!(interpolate_bicubic_bf16, interpolate_bicubic_f32); + +interpolate_typed!(interpolate_lanczos3_f32, interpolate_lanczos3_impl, f32); +interpolate_typed!(interpolate_lanczos3_f64, interpolate_lanczos3_impl, f64); +interpolate_typed!(interpolate_lanczos3_f16, interpolate_lanczos3_impl, f16); +interpolate_bf16!(interpolate_lanczos3_bf16, interpolate_lanczos3_f32); + +// ============================================================================ +// Backward pass - dtype dispatch +// ============================================================================ + +interpolate_backward_typed!( + interpolate_nearest_backward_f32, + interpolate_nearest_backward_impl, + f32 +); +interpolate_backward_typed!( + interpolate_nearest_backward_f64, + interpolate_nearest_backward_impl, + f64 +); +interpolate_backward_typed!( + interpolate_nearest_backward_f16, + interpolate_nearest_backward_impl, + f16 +); +interpolate_backward_bf16!( + interpolate_nearest_backward_bf16, + interpolate_nearest_backward_f32 +); + +interpolate_backward_typed!( + interpolate_bilinear_backward_f32, + interpolate_bilinear_backward_impl, + f32 +); +interpolate_backward_typed!( + interpolate_bilinear_backward_f64, + interpolate_bilinear_backward_impl, + f64 +); +interpolate_backward_typed!( + interpolate_bilinear_backward_f16, + interpolate_bilinear_backward_impl, + f16 +); +interpolate_backward_bf16!( + interpolate_bilinear_backward_bf16, + interpolate_bilinear_backward_f32 +); + +interpolate_backward_typed!( + interpolate_bicubic_backward_f32, + interpolate_bicubic_backward_impl, + f32 +); +interpolate_backward_typed!( + interpolate_bicubic_backward_f64, + interpolate_bicubic_backward_impl, + f64 +); +interpolate_backward_typed!( + interpolate_bicubic_backward_f16, + interpolate_bicubic_backward_impl, + f16 +); +interpolate_backward_bf16!( + interpolate_bicubic_backward_bf16, + interpolate_bicubic_backward_f32 +); + +// ============================================================================ +// Generic implementations with rayon parallelism +// ============================================================================ + +/// Compute coordinate mapping parameters. +/// +/// align_corners=true: ratio = (in_size - 1) / (out_size - 1), coord = out * ratio +/// align_corners=false: ratio = in_size / out_size, coord = (out + 0.5) * ratio - 0.5 +fn coord_ratio(in_size: usize, out_size: usize, align_corners: bool) -> f64 { + if align_corners { + (in_size as f64 - 1.0) / (out_size.max(1) - 1).max(1) as f64 + } else { + in_size as f64 / out_size as f64 + } +} + +/// Map an output coordinate to input coordinate. +#[inline] +fn map_coord(out_coord: usize, ratio: f64, align_corners: bool) -> f64 { + if align_corners { + out_coord as f64 * ratio + } else { + (out_coord as f64 + 0.5) * ratio - 0.5 + } +} + +/// Precompute nearest-neighbor source index lookup table. +/// Returns a Vec where `map[out_coord]` is the corresponding input coordinate. +fn nearest_index_map(in_size: usize, out_size: usize) -> Vec { + let ratio = in_size as f64 / out_size as f64; + let max = in_size - 1; + (0..out_size) + .map(|o| ((o as f64 * ratio).floor() as usize).min(max)) + .collect() +} + +/// Nearest neighbor interpolation. +/// Maps output coordinates to input using floor(ratio * out_coord) via precomputed lookup tables. +fn interpolate_nearest_impl( + x: FlexTensor, + output_size: [usize; 2], + _align_corners: bool, +) -> FlexTensor +where + T: Float + burn_backend::Element + bytemuck::Pod + Send + Sync, +{ + let x = x.to_contiguous(); + let input = x.storage::(); + let shape = x.layout().shape(); + + let batch = shape[0]; + let channels = shape[1]; + let in_height = shape[2]; + let in_width = shape[3]; + assert!( + in_height > 0 && in_width > 0, + "interpolate: input spatial dimensions must be > 0" + ); + let [out_height, out_width] = output_size; + + let y_map = nearest_index_map(in_height, out_height); + let x_map = nearest_index_map(in_width, out_width); + + let out_numel = batch * channels * out_height * out_width; + let in_hw = in_height * in_width; + let out_hw = out_height * out_width; + + // Per-plane nearest gather using precomputed index maps. + #[inline] + fn gather_plane( + input: &[T], + in_base: usize, + output: &mut [T], + in_width: usize, + out_width: usize, + y_map: &[usize], + x_map: &[usize], + ) { + for (oh, &ih) in y_map.iter().enumerate() { + let in_row = in_base + ih * in_width; + let out_row_start = oh * out_width; + for (ow, &iw) in x_map.iter().enumerate() { + output[out_row_start + ow] = input[in_row + iw]; + } + } + } + + let output = { + let mut output: Vec = Vec::with_capacity(out_numel); + #[allow(clippy::uninit_vec)] + unsafe { + output.set_len(out_numel); + } + + let bc = batch * channels; + + // Each (batch, channel) plane is independent, so parallelize across planes. + #[cfg(feature = "rayon")] + if out_numel >= super::PARALLEL_THRESHOLD { + use rayon::prelude::*; + + output + .par_chunks_mut(out_hw) + .enumerate() + .for_each(|(bc_idx, out_plane)| { + let in_base = bc_idx * in_hw; + gather_plane( + input, in_base, out_plane, in_width, out_width, &y_map, &x_map, + ); + }); + } else { + for bc_idx in 0..bc { + let in_base = bc_idx * in_hw; + let out_start = bc_idx * out_hw; + gather_plane( + input, + in_base, + &mut output[out_start..out_start + out_hw], + in_width, + out_width, + &y_map, + &x_map, + ); + } + } + #[cfg(not(feature = "rayon"))] + for bc_idx in 0..bc { + let in_base = bc_idx * in_hw; + let out_start = bc_idx * out_hw; + gather_plane( + input, + in_base, + &mut output[out_start..out_start + out_hw], + in_width, + out_width, + &y_map, + &x_map, + ); + } + output + }; + + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(Shape::from(vec![batch, channels, out_height, out_width])), + x.dtype(), + ) +} + +/// Bilinear interpolation. +/// Uses 4-point weighted average based on distance to neighbors. +fn interpolate_bilinear_impl( + x: FlexTensor, + output_size: [usize; 2], + align_corners: bool, +) -> FlexTensor +where + T: Float + burn_backend::Element + bytemuck::Pod + Send + Sync, +{ + let x = x.to_contiguous(); + let input = x.storage::(); + let shape = x.layout().shape(); + + let batch = shape[0]; + let channels = shape[1]; + let in_height = shape[2]; + let in_width = shape[3]; + assert!( + in_height > 0 && in_width > 0, + "interpolate: input spatial dimensions must be > 0" + ); + let [out_height, out_width] = output_size; + + let y_ratio = coord_ratio(in_height, out_height, align_corners); + let x_ratio = coord_ratio(in_width, out_width, align_corners); + + let out_numel = batch * channels * out_height * out_width; + let in_hw = in_height * in_width; + let out_hw = out_height * out_width; + + let output = { + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + + let mut output = vec![T::zero(); out_numel]; + let out_ptr = crate::ops::SendMutPtr::new(output.as_mut_ptr()); + + (0..batch).into_par_iter().for_each(|b| { + (0..channels).into_par_iter().for_each(|c| { + let in_base = b * channels * in_hw + c * in_hw; + let out_base = b * channels * out_hw + c * out_hw; + + for oh in 0..out_height { + let y_in = map_coord(oh, y_ratio, align_corners); + let y_low = (y_in.floor().max(0.0)) as usize; + let y_high = (y_low + 1).min(in_height - 1); + let y_weight = T::from((y_in - y_low as f64).max(0.0)).unwrap(); + + for ow in 0..out_width { + let x_in = map_coord(ow, x_ratio, align_corners); + let x_low = (x_in.floor().max(0.0)) as usize; + let x_high = (x_low + 1).min(in_width - 1); + let x_weight = T::from((x_in - x_low as f64).max(0.0)).unwrap(); + + let p_a = input[in_base + y_low * in_width + x_low]; + let p_b = input[in_base + y_low * in_width + x_high]; + let p_c = input[in_base + y_high * in_width + x_low]; + let p_d = input[in_base + y_high * in_width + x_high]; + + let one = T::one(); + let result = p_a * (one - x_weight) * (one - y_weight) + + p_b * x_weight * (one - y_weight) + + p_c * (one - x_weight) * y_weight + + p_d * x_weight * y_weight; + + let out_idx = out_base + oh * out_width + ow; + unsafe { + out_ptr.write(out_idx, result); + } + } + } + }); + }); + output + } + #[cfg(not(feature = "rayon"))] + { + let mut output = vec![T::zero(); out_numel]; + + for b in 0..batch { + for c in 0..channels { + let in_base = b * channels * in_hw + c * in_hw; + let out_base = b * channels * out_hw + c * out_hw; + + for oh in 0..out_height { + let y_in = map_coord(oh, y_ratio, align_corners); + let y_low = (y_in.floor().max(0.0)) as usize; + let y_high = (y_low + 1).min(in_height - 1); + let y_weight = T::from((y_in - y_low as f64).max(0.0)).unwrap(); + + for ow in 0..out_width { + let x_in = map_coord(ow, x_ratio, align_corners); + let x_low = (x_in.floor().max(0.0)) as usize; + let x_high = (x_low + 1).min(in_width - 1); + let x_weight = T::from((x_in - x_low as f64).max(0.0)).unwrap(); + + let p_a = input[in_base + y_low * in_width + x_low]; + let p_b = input[in_base + y_low * in_width + x_high]; + let p_c = input[in_base + y_high * in_width + x_low]; + let p_d = input[in_base + y_high * in_width + x_high]; + + let one = T::one(); + let result = p_a * (one - x_weight) * (one - y_weight) + + p_b * x_weight * (one - y_weight) + + p_c * (one - x_weight) * y_weight + + p_d * x_weight * y_weight; + + output[out_base + oh * out_width + ow] = result; + } + } + } + } + output + } + }; + + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(Shape::from(vec![batch, channels, out_height, out_width])), + x.dtype(), + ) +} + +/// Bicubic interpolation using cubic convolution. +fn interpolate_bicubic_impl( + x: FlexTensor, + output_size: [usize; 2], + align_corners: bool, +) -> FlexTensor +where + T: Float + burn_backend::Element + bytemuck::Pod + Send + Sync, +{ + let x = x.to_contiguous(); + let input = x.storage::(); + let shape = x.layout().shape(); + + let batch = shape[0]; + let channels = shape[1]; + let in_height = shape[2]; + let in_width = shape[3]; + assert!( + in_height > 0 && in_width > 0, + "interpolate: input spatial dimensions must be > 0" + ); + let [out_height, out_width] = output_size; + + let y_ratio = coord_ratio(in_height, out_height, align_corners); + let x_ratio = coord_ratio(in_width, out_width, align_corners); + + let out_numel = batch * channels * out_height * out_width; + let in_hw = in_height * in_width; + let out_hw = out_height * out_width; + let a = -0.75_f64; + + let output = { + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + + let mut output = vec![T::zero(); out_numel]; + let out_ptr = crate::ops::SendMutPtr::new(output.as_mut_ptr()); + let num_bc_pairs = batch * channels; + + // Adaptive parallelization: if few (batch, channel) pairs, parallelize rows too + if num_bc_pairs < 8 { + // Fine-grained: parallelize over (batch, channel, row) for better CPU utilization + let total_rows = batch * channels * out_height; + (0..total_rows).into_par_iter().for_each(|id| { + let b = id / (channels * out_height); + let remainder = id % (channels * out_height); + let c = remainder / out_height; + let oh = remainder % out_height; + + let in_base = b * channels * in_hw + c * in_hw; + let out_base = b * channels * out_hw + c * out_hw; + + let y_in = map_coord(oh, y_ratio, align_corners); + let y0 = y_in.floor() as isize; + + for ow in 0..out_width { + let x_in = map_coord(ow, x_ratio, align_corners); + let x0 = x_in.floor() as isize; + + let mut sum = 0.0_f64; + + for dy in -1..=2_isize { + let y = y0 + dy; + let y_clamped = y.clamp(0, in_height as isize - 1) as usize; + let ty = (y_in - y0 as f64) - dy as f64; + let wy = cubic_weight(ty, a); + + for dx in -1..=2_isize { + let x = x0 + dx; + let x_clamped = x.clamp(0, in_width as isize - 1) as usize; + let tx = (x_in - x0 as f64) - dx as f64; + let wx = cubic_weight(tx, a); + + let val = input[in_base + y_clamped * in_width + x_clamped]; + let val_f64 = + ::to_f64(&val).unwrap_or(0.0); + sum += val_f64 * wx * wy; + } + } + + let out_idx = out_base + oh * out_width + ow; + unsafe { + out_ptr.write(out_idx, T::from(sum).unwrap()); + } + } + }); + } else { + // Coarse-grained: parallelize over (batch, channel) for cache efficiency + (0..batch).into_par_iter().for_each(|b| { + (0..channels).into_par_iter().for_each(|c| { + let in_base = b * channels * in_hw + c * in_hw; + let out_base = b * channels * out_hw + c * out_hw; + + for oh in 0..out_height { + let y_in = map_coord(oh, y_ratio, align_corners); + let y0 = y_in.floor() as isize; + + for ow in 0..out_width { + let x_in = map_coord(ow, x_ratio, align_corners); + let x0 = x_in.floor() as isize; + + let mut sum = 0.0_f64; + + for dy in -1..=2_isize { + let y = y0 + dy; + let y_clamped = y.clamp(0, in_height as isize - 1) as usize; + let ty = (y_in - y0 as f64) - dy as f64; + let wy = cubic_weight(ty, a); + + for dx in -1..=2_isize { + let x = x0 + dx; + let x_clamped = x.clamp(0, in_width as isize - 1) as usize; + let tx = (x_in - x0 as f64) - dx as f64; + let wx = cubic_weight(tx, a); + + let val = input[in_base + y_clamped * in_width + x_clamped]; + let val_f64 = ::to_f64(&val) + .unwrap_or(0.0); + sum += val_f64 * wx * wy; + } + } + + let out_idx = out_base + oh * out_width + ow; + unsafe { + out_ptr.write(out_idx, T::from(sum).unwrap()); + } + } + } + }); + }); + } + output + } + #[cfg(not(feature = "rayon"))] + { + let mut output = vec![T::zero(); out_numel]; + + for b in 0..batch { + for c in 0..channels { + let in_base = b * channels * in_hw + c * in_hw; + let out_base = b * channels * out_hw + c * out_hw; + + for oh in 0..out_height { + let y_in = map_coord(oh, y_ratio, align_corners); + let y0 = y_in.floor() as isize; + + for ow in 0..out_width { + let x_in = map_coord(ow, x_ratio, align_corners); + let x0 = x_in.floor() as isize; + + let mut sum = 0.0_f64; + + for dy in -1..=2_isize { + let y = y0 + dy; + let y_clamped = y.clamp(0, in_height as isize - 1) as usize; + let ty = (y_in - y0 as f64) - dy as f64; + let wy = cubic_weight(ty, a); + + for dx in -1..=2_isize { + let x = x0 + dx; + let x_clamped = x.clamp(0, in_width as isize - 1) as usize; + let tx = (x_in - x0 as f64) - dx as f64; + let wx = cubic_weight(tx, a); + + let val = input[in_base + y_clamped * in_width + x_clamped]; + let val_f64 = + ::to_f64(&val).unwrap_or(0.0); + sum += val_f64 * wx * wy; + } + } + + output[out_base + oh * out_width + ow] = T::from(sum).unwrap(); + } + } + } + } + output + } + }; + + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(Shape::from(vec![batch, channels, out_height, out_width])), + x.dtype(), + ) +} + +/// Cubic convolution weight function. +/// Uses the Keys cubic interpolation kernel. +#[inline] +fn cubic_weight(t: f64, a: f64) -> f64 { + let t_abs = t.abs(); + if t_abs < 1.0 { + ((a + 2.0) * t_abs - (a + 3.0)) * t_abs * t_abs + 1.0 + } else if t_abs < 2.0 { + ((a * t_abs - 5.0 * a) * t_abs + 8.0 * a) * t_abs - 4.0 * a + } else { + 0.0 + } +} + +/// Lanczos3 interpolation weight function. +/// Uses a sinc-windowed sinc kernel with a=3. +#[inline] +fn lanczos3_weight(x: f64) -> f64 { + if x == 0.0 { + return 1.0; + } + let abs_x = x.abs(); + if abs_x >= 3.0 { + return 0.0; + } + let pi = core::f64::consts::PI; + let pi_x = pi * x; + let pi_x_over_3 = pi_x / 3.0; + (pi_x.sin() * pi_x_over_3.sin()) / (pi_x * pi_x_over_3) +} + +/// Lanczos3 interpolation (6x6 sinc-windowed kernel). +/// +/// Uses skip-and-renormalize boundary handling: out-of-bounds samples are +/// excluded and weights are renormalized. This avoids edge ringing artifacts +/// from replicated boundary pixels (unlike bicubic which clamps to edge). +/// Matches the ndarray reference implementation. +fn interpolate_lanczos3_impl( + x: FlexTensor, + output_size: [usize; 2], + align_corners: bool, +) -> FlexTensor +where + T: Float + burn_backend::Element + bytemuck::Pod + Send + Sync, +{ + let x = x.to_contiguous(); + let input = x.storage::(); + let shape = x.layout().shape(); + + let batch = shape[0]; + let channels = shape[1]; + let in_height = shape[2]; + let in_width = shape[3]; + assert!( + in_height > 0 && in_width > 0, + "interpolate: input spatial dimensions must be > 0" + ); + let [out_height, out_width] = output_size; + + let y_ratio = coord_ratio(in_height, out_height, align_corners); + let x_ratio = coord_ratio(in_width, out_width, align_corners); + + let out_numel = batch * channels * out_height * out_width; + let in_hw = in_height * in_width; + let out_hw = out_height * out_width; + let max_h = in_height as isize - 1; + let max_w = in_width as isize - 1; + + /// Compute one Lanczos3 output pixel given pre-computed y coordinates. + #[inline] + #[allow(clippy::too_many_arguments)] + fn lanczos3_sample( + input: &[T], + in_base: usize, + in_width: usize, + y_in: f64, + y0: f64, + x_in: f64, + max_h: isize, + max_w: isize, + ) -> T { + let x0 = x_in.floor(); + let mut result = 0.0_f64; + let mut weight_sum = 0.0_f64; + + for ky in -2..=3_isize { + let yi = y0 as isize + ky; + if yi < 0 || yi > max_h { + continue; + } + let wy = lanczos3_weight(y_in - (y0 + ky as f64)); + for kx in -2..=3_isize { + let xi = x0 as isize + kx; + if xi < 0 || xi > max_w { + continue; + } + let wx = lanczos3_weight(x_in - (x0 + kx as f64)); + let w = wy * wx; + let val = input[in_base + yi as usize * in_width + xi as usize]; + let val_f64 = ::to_f64(&val).unwrap_or(0.0); + result += val_f64 * w; + weight_sum += w; + } + } + + if weight_sum != 0.0 { + result /= weight_sum; + } + T::from(result).unwrap() + } + + let output = { + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + + let mut output = vec![T::zero(); out_numel]; + let out_ptr = crate::ops::SendMutPtr::new(output.as_mut_ptr()); + + let total_rows = batch * channels * out_height; + (0..total_rows).into_par_iter().for_each(|id| { + let b = id / (channels * out_height); + let remainder = id % (channels * out_height); + let c = remainder / out_height; + let oh = remainder % out_height; + + let in_base = b * channels * in_hw + c * in_hw; + let out_base = b * channels * out_hw + c * out_hw; + + let y_in = map_coord(oh, y_ratio, align_corners); + let y0 = y_in.floor(); + + for ow in 0..out_width { + let x_in = map_coord(ow, x_ratio, align_corners); + let out_idx = out_base + oh * out_width + ow; + unsafe { + out_ptr.write( + out_idx, + lanczos3_sample(input, in_base, in_width, y_in, y0, x_in, max_h, max_w), + ); + } + } + }); + output + } + #[cfg(not(feature = "rayon"))] + { + let mut output = vec![T::zero(); out_numel]; + + for b in 0..batch { + for c in 0..channels { + let in_base = b * channels * in_hw + c * in_hw; + let out_base = b * channels * out_hw + c * out_hw; + + for oh in 0..out_height { + let y_in = map_coord(oh, y_ratio, align_corners); + let y0 = y_in.floor(); + + for ow in 0..out_width { + let x_in = map_coord(ow, x_ratio, align_corners); + output[out_base + oh * out_width + ow] = lanczos3_sample( + input, in_base, in_width, y_in, y0, x_in, max_h, max_w, + ); + } + } + } + } + output + } + }; + + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(Shape::from(vec![batch, channels, out_height, out_width])), + x.dtype(), + ) +} + +// ============================================================================ +// Backward implementations +// ============================================================================ + +/// Nearest neighbor backward: accumulates gradients at source positions. +fn interpolate_nearest_backward_impl( + x: FlexTensor, + grad: FlexTensor, + output_size: [usize; 2], + _align_corners: bool, +) -> FlexTensor +where + T: Float + burn_backend::Element + bytemuck::Pod + Send + Sync, +{ + let grad = grad.to_contiguous(); + let grad_data = grad.storage::(); + let shape = x.layout().shape(); + + let batch = shape[0]; + let channels = shape[1]; + let in_height = shape[2]; + let in_width = shape[3]; + assert!( + in_height > 0 && in_width > 0, + "interpolate: input spatial dimensions must be > 0" + ); + let [out_height, out_width] = output_size; + + let y_map = nearest_index_map(in_height, out_height); + let x_map = nearest_index_map(in_width, out_width); + + let in_numel = batch * channels * in_height * in_width; + let in_hw = in_height * in_width; + let out_hw = out_height * out_width; + + // Scatter-add gradients from one output plane into one input gradient plane. + #[inline] + fn scatter_plane( + grad_data: &[T], + grad_base: usize, + input_grad: &mut [T], + in_width: usize, + out_width: usize, + y_map: &[usize], + x_map: &[usize], + ) { + for (oh, &ih) in y_map.iter().enumerate() { + let grad_row = grad_base + oh * out_width; + for (ow, &iw) in x_map.iter().enumerate() { + input_grad[ih * in_width + iw] = + input_grad[ih * in_width + iw] + grad_data[grad_row + ow]; + } + } + } + + let mut input_grad = vec![T::zero(); in_numel]; + let bc = batch * channels; + + // Each (batch, channel) plane is independent, so parallelize across planes. + // Gate on output size since the work is proportional to iterating output pixels. + #[cfg(feature = "rayon")] + if bc * out_hw >= super::PARALLEL_THRESHOLD { + use rayon::prelude::*; + + input_grad + .par_chunks_mut(in_hw) + .enumerate() + .for_each(|(bc_idx, grad_plane)| { + let grad_base = bc_idx * out_hw; + scatter_plane( + grad_data, grad_base, grad_plane, in_width, out_width, &y_map, &x_map, + ); + }); + } else { + for bc_idx in 0..bc { + let grad_base = bc_idx * out_hw; + let in_start = bc_idx * in_hw; + scatter_plane( + grad_data, + grad_base, + &mut input_grad[in_start..in_start + in_hw], + in_width, + out_width, + &y_map, + &x_map, + ); + } + } + #[cfg(not(feature = "rayon"))] + for bc_idx in 0..bc { + let grad_base = bc_idx * out_hw; + let in_start = bc_idx * in_hw; + scatter_plane( + grad_data, + grad_base, + &mut input_grad[in_start..in_start + in_hw], + in_width, + out_width, + &y_map, + &x_map, + ); + } + + FlexTensor::new( + Bytes::from_elems(input_grad), + Layout::contiguous(Shape::from(vec![batch, channels, in_height, in_width])), + x.dtype(), + ) +} + +/// Bilinear backward: distributes gradients to 4 source positions weighted by bilinear coefficients. +fn interpolate_bilinear_backward_impl( + x: FlexTensor, + grad: FlexTensor, + output_size: [usize; 2], + align_corners: bool, +) -> FlexTensor +where + T: Float + burn_backend::Element + bytemuck::Pod, +{ + let grad = grad.to_contiguous(); + let grad_data = grad.storage::(); + let shape = x.layout().shape(); + + let batch = shape[0]; + let channels = shape[1]; + let in_height = shape[2]; + let in_width = shape[3]; + assert!( + in_height > 0 && in_width > 0, + "interpolate: input spatial dimensions must be > 0" + ); + let [out_height, out_width] = output_size; + + let y_ratio = coord_ratio(in_height, out_height, align_corners); + let x_ratio = coord_ratio(in_width, out_width, align_corners); + + let in_numel = batch * channels * in_height * in_width; + let mut input_grad = vec![T::zero(); in_numel]; + + let in_hw = in_height * in_width; + let out_hw = out_height * out_width; + + for b in 0..batch { + for c in 0..channels { + let in_base = b * channels * in_hw + c * in_hw; + let out_base = b * channels * out_hw + c * out_hw; + + for oh in 0..out_height { + let y_in = map_coord(oh, y_ratio, align_corners); + let y_low = (y_in.floor().max(0.0)) as usize; + let y_high = (y_low + 1).min(in_height - 1); + let y_weight = T::from((y_in - y_low as f64).max(0.0)).unwrap(); + + for ow in 0..out_width { + let x_in = map_coord(ow, x_ratio, align_corners); + let x_low = (x_in.floor().max(0.0)) as usize; + let x_high = (x_low + 1).min(in_width - 1); + let x_weight = T::from((x_in - x_low as f64).max(0.0)).unwrap(); + + let grad_val = grad_data[out_base + oh * out_width + ow]; + let one = T::one(); + + input_grad[in_base + y_low * in_width + x_low] = input_grad + [in_base + y_low * in_width + x_low] + + grad_val * (one - x_weight) * (one - y_weight); + input_grad[in_base + y_low * in_width + x_high] = input_grad + [in_base + y_low * in_width + x_high] + + grad_val * x_weight * (one - y_weight); + input_grad[in_base + y_high * in_width + x_low] = input_grad + [in_base + y_high * in_width + x_low] + + grad_val * (one - x_weight) * y_weight; + input_grad[in_base + y_high * in_width + x_high] = input_grad + [in_base + y_high * in_width + x_high] + + grad_val * x_weight * y_weight; + } + } + } + } + + FlexTensor::new( + Bytes::from_elems(input_grad), + Layout::contiguous(Shape::from(vec![batch, channels, in_height, in_width])), + x.dtype(), + ) +} + +/// Bicubic backward: distributes gradients to 16 source positions weighted by cubic coefficients. +fn interpolate_bicubic_backward_impl( + x: FlexTensor, + grad: FlexTensor, + output_size: [usize; 2], + align_corners: bool, +) -> FlexTensor +where + T: Float + burn_backend::Element + bytemuck::Pod, +{ + let grad = grad.to_contiguous(); + let grad_data = grad.storage::(); + let shape = x.layout().shape(); + + let batch = shape[0]; + let channels = shape[1]; + let in_height = shape[2]; + let in_width = shape[3]; + assert!( + in_height > 0 && in_width > 0, + "interpolate: input spatial dimensions must be > 0" + ); + let [out_height, out_width] = output_size; + + let y_ratio = coord_ratio(in_height, out_height, align_corners); + let x_ratio = coord_ratio(in_width, out_width, align_corners); + + let in_numel = batch * channels * in_height * in_width; + let mut input_grad = vec![T::zero(); in_numel]; + + let in_hw = in_height * in_width; + let out_hw = out_height * out_width; + let a = -0.75_f64; + + for b in 0..batch { + for c in 0..channels { + let in_base = b * channels * in_hw + c * in_hw; + let out_base = b * channels * out_hw + c * out_hw; + + for oh in 0..out_height { + let y_in = map_coord(oh, y_ratio, align_corners); + let y0 = y_in.floor() as isize; + + for ow in 0..out_width { + let x_in = map_coord(ow, x_ratio, align_corners); + let x0 = x_in.floor() as isize; + + let grad_val = ::to_f64( + &grad_data[out_base + oh * out_width + ow], + ) + .unwrap_or(0.0); + + for dy in -1..=2_isize { + let y = y0 + dy; + let y_idx = y.clamp(0, in_height as isize - 1) as usize; + let ty = (y_in - y0 as f64) - dy as f64; + let wy = cubic_weight(ty, a); + + for dx in -1..=2_isize { + let x = x0 + dx; + let x_idx = x.clamp(0, in_width as isize - 1) as usize; + let tx = (x_in - x0 as f64) - dx as f64; + let wx = cubic_weight(tx, a); + + let weight = wx * wy * grad_val; + input_grad[in_base + y_idx * in_width + x_idx] = input_grad + [in_base + y_idx * in_width + x_idx] + + T::from(weight).unwrap(); + } + } + } + } + } + } + + FlexTensor::new( + Bytes::from_elems(input_grad), + Layout::contiguous(Shape::from(vec![batch, channels, in_height, in_width])), + x.dtype(), + ) +} + +// ============================================================================ +// Dtype conversion helpers +// ============================================================================ + +fn convert_bf16_to_f32(x: &FlexTensor) -> FlexTensor { + let x = x.clone().to_contiguous(); + let input = x.storage::(); + let output: Vec = input.iter().map(|v| v.to_f32()).collect(); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(x.layout().shape().clone()), + DType::F32, + ) +} + +fn convert_f32_to_bf16(x: &FlexTensor) -> FlexTensor { + let x = x.clone().to_contiguous(); + let input = x.storage::(); + let output: Vec = input.iter().map(|v| bf16::from_f32(*v)).collect(); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(x.layout().shape().clone()), + DType::BF16, + ) +} + +// Tests kept here exercise flex-specific behavior: the f32-typed +// internal helpers (`interpolate_nearest_f32`, `interpolate_bilinear_f32`, +// `interpolate_bicubic_f32`, `interpolate_nearest_backward_f32`) that +// back the dtype-dispatching public ops. End-to-end interpolate +// correctness across backends lives in the shared module tests under +// crates/burn-backend-tests/tests/tensor/float/module/{bicubic,bilinear, +// lanczos3,nearest}_interpolate.rs. +#[cfg(test)] +mod tests { + use super::*; + + fn make_input_f32(batch: usize, channels: usize, height: usize, width: usize) -> FlexTensor { + let numel = batch * channels * height * width; + let data: Vec = (0..numel).map(|i| i as f32).collect(); + FlexTensor::new( + Bytes::from_elems(data), + Layout::contiguous(Shape::from(vec![batch, channels, height, width])), + DType::F32, + ) + } + + #[test] + fn test_nearest_upsample_2x() { + let data = vec![1.0f32, 2.0, 3.0, 4.0]; + let x = FlexTensor::new( + Bytes::from_elems(data), + Layout::contiguous(Shape::from(vec![1, 1, 2, 2])), + DType::F32, + ); + + let result = interpolate_nearest_f32(x, [4, 4], true); + let output = result.storage::(); + + assert_eq!(output.len(), 16); + assert_eq!(output[0], 1.0); + assert_eq!(output[1], 1.0); + assert_eq!(output[2], 2.0); + assert_eq!(output[3], 2.0); + } + + #[test] + fn test_bilinear_upsample_2x() { + let data = vec![0.0f32, 1.0, 1.0, 0.0]; + let x = FlexTensor::new( + Bytes::from_elems(data), + Layout::contiguous(Shape::from(vec![1, 1, 2, 2])), + DType::F32, + ); + + let result = interpolate_bilinear_f32(x, [4, 4], true); + let output = result.storage::(); + + assert!((output[0] - 0.0).abs() < 1e-5); + assert!((output[3] - 1.0).abs() < 1e-5); + assert!((output[12] - 1.0).abs() < 1e-5); + assert!((output[15] - 0.0).abs() < 1e-5); + } + + #[test] + fn test_bicubic_basic() { + let x = make_input_f32(1, 1, 4, 4); + let result = interpolate_bicubic_f32(x, [8, 8], true); + assert_eq!(result.layout().shape().to_vec(), vec![1, 1, 8, 8]); + } + + #[test] + fn test_downsample() { + let x = make_input_f32(1, 1, 4, 4); + let result = interpolate_nearest_f32(x, [2, 2], true); + assert_eq!(result.layout().shape().to_vec(), vec![1, 1, 2, 2]); + } + + #[test] + fn test_nearest_backward() { + let x = make_input_f32(1, 1, 2, 2); + let grad = FlexTensor::new( + Bytes::from_elems(vec![1.0f32; 16]), + Layout::contiguous(Shape::from(vec![1, 1, 4, 4])), + DType::F32, + ); + + let result = interpolate_nearest_backward_f32(x, grad, [4, 4], true); + let output = result.storage::(); + + assert_eq!(output.len(), 4); + assert!((output[0] - 4.0).abs() < 1e-5); + } +} diff --git a/crates/burn-flex/src/ops/mask.rs b/crates/burn-flex/src/ops/mask.rs new file mode 100644 index 0000000..6e7cfad --- /dev/null +++ b/crates/burn-flex/src/ops/mask.rs @@ -0,0 +1,318 @@ +//! Mask operations for conditional element replacement. + +use alloc::vec::Vec; +use burn_backend::Element; +use burn_std::{Bytes, bf16, f16}; + +use crate::{FlexTensor, Layout}; + +/// Allocate a Vec of given length without zeroing. +/// The caller must write every element before reading. +#[cfg(feature = "simd")] +#[inline] +fn uninit_vec(len: usize) -> Vec { + let mut v = Vec::with_capacity(len); + #[allow(clippy::uninit_vec)] + unsafe { + v.set_len(len); + } + v +} + +/// Fill tensor elements with a value where mask is true. +/// +/// mask_fill(tensor, mask, value) -> tensor with elements replaced where mask is true +pub fn mask_fill(tensor: FlexTensor, mask: FlexTensor, value: T) -> FlexTensor +where + T: Element + bytemuck::Pod + Copy, +{ + let dtype = tensor.dtype(); + + // Broadcast mask to tensor shape if needed + let (tensor, mask) = crate::ops::expand::broadcast_binary(tensor, mask); + + let tensor = tensor.to_contiguous(); + let mask = mask.to_contiguous(); + + let shape = tensor.layout().shape().clone(); + let tensor_data: &[T] = tensor.storage(); + let mask_data: &[u8] = mask.bytes(); + + let result: Vec = tensor_data + .iter() + .zip(mask_data.iter()) + .map(|(&elem, &m)| if m != 0 { value } else { elem }) + .collect(); + + FlexTensor::new(Bytes::from_elems(result), Layout::contiguous(shape), dtype) +} + +/// Mask fill for f32 (SIMD-accelerated). +pub fn mask_fill_f32(tensor: FlexTensor, mask: FlexTensor, value: f32) -> FlexTensor { + #[cfg(feature = "simd")] + { + let dtype = tensor.dtype(); + let (tensor, mask) = crate::ops::expand::broadcast_binary(tensor, mask); + let tensor = tensor.to_contiguous(); + let mask = mask.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let len = tensor.storage::().len(); + let mut out = uninit_vec::(len); + crate::simd::mask_fill_f32(tensor.storage(), mask.bytes(), value, &mut out); + FlexTensor::new(Bytes::from_elems(out), Layout::contiguous(shape), dtype) + } + #[cfg(not(feature = "simd"))] + { + mask_fill(tensor, mask, value) + } +} + +/// Mask fill for f64 (SIMD-accelerated). +pub fn mask_fill_f64(tensor: FlexTensor, mask: FlexTensor, value: f64) -> FlexTensor { + #[cfg(feature = "simd")] + { + let dtype = tensor.dtype(); + let (tensor, mask) = crate::ops::expand::broadcast_binary(tensor, mask); + let tensor = tensor.to_contiguous(); + let mask = mask.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let len = tensor.storage::().len(); + let mut out = uninit_vec::(len); + crate::simd::mask_fill_f64(tensor.storage(), mask.bytes(), value, &mut out); + FlexTensor::new(Bytes::from_elems(out), Layout::contiguous(shape), dtype) + } + #[cfg(not(feature = "simd"))] + { + mask_fill(tensor, mask, value) + } +} + +/// Mask fill for f16. +pub fn mask_fill_f16(tensor: FlexTensor, mask: FlexTensor, value: f16) -> FlexTensor { + mask_fill(tensor, mask, value) +} + +/// Mask fill for bf16. +pub fn mask_fill_bf16(tensor: FlexTensor, mask: FlexTensor, value: bf16) -> FlexTensor { + mask_fill(tensor, mask, value) +} + +/// Mask fill for i64 (SIMD-accelerated). +pub fn mask_fill_i64(tensor: FlexTensor, mask: FlexTensor, value: i64) -> FlexTensor { + #[cfg(feature = "simd")] + { + let dtype = tensor.dtype(); + let (tensor, mask) = crate::ops::expand::broadcast_binary(tensor, mask); + let tensor = tensor.to_contiguous(); + let mask = mask.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let len = tensor.storage::().len(); + let mut out = uninit_vec::(len); + crate::simd::mask_fill_i64(tensor.storage(), mask.bytes(), value, &mut out); + FlexTensor::new(Bytes::from_elems(out), Layout::contiguous(shape), dtype) + } + #[cfg(not(feature = "simd"))] + { + mask_fill(tensor, mask, value) + } +} + +/// Mask fill for u64. +pub fn mask_fill_u64(tensor: FlexTensor, mask: FlexTensor, value: u64) -> FlexTensor { + mask_fill(tensor, mask, value) +} + +/// Mask fill for bool tensors (SIMD-accelerated). +pub fn mask_fill_bool(tensor: FlexTensor, mask: FlexTensor, value: bool) -> FlexTensor { + // Preserve the input tensor's bool dtype for the output. + let out_dtype = burn_std::BoolDType::from(tensor.dtype()); + #[cfg(feature = "simd")] + { + let (tensor, mask) = crate::ops::expand::broadcast_binary(tensor, mask); + let tensor = tensor.to_contiguous(); + let mask = mask.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let len = tensor.bytes().len(); + let mut out = uninit_vec::(len); + crate::simd::mask_fill_u8(tensor.bytes(), mask.bytes(), value as u8, &mut out); + crate::ops::comparison::make_bool_tensor(out, shape, out_dtype) + } + #[cfg(not(feature = "simd"))] + { + let (tensor, mask) = crate::ops::expand::broadcast_binary(tensor, mask); + let tensor = tensor.to_contiguous(); + let mask = mask.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let tensor_data: &[u8] = tensor.bytes(); + let mask_data: &[u8] = mask.bytes(); + let value_u8 = value as u8; + let result: Vec = tensor_data + .iter() + .zip(mask_data.iter()) + .map(|(&elem, &m)| if m != 0 { value_u8 } else { elem }) + .collect(); + crate::ops::comparison::make_bool_tensor(result, shape, out_dtype) + } +} + +/// Replace elements from value tensor where mask is true. +/// +/// mask_where(tensor, mask, value) -> tensor with elements from value where mask is true +pub fn mask_where(tensor: FlexTensor, mask: FlexTensor, value: FlexTensor) -> FlexTensor +where + T: Element + bytemuck::Pod + Copy, +{ + let dtype = tensor.dtype(); + let (tensor, mask, value) = broadcast_three(tensor, mask, value); + + let shape = tensor.layout().shape().clone(); + let tensor_data: &[T] = tensor.storage(); + let mask_data: &[u8] = mask.bytes(); + let value_data: &[T] = value.storage(); + + let result: Vec = tensor_data + .iter() + .zip(mask_data.iter()) + .zip(value_data.iter()) + .map(|((&t, &m), &v)| if m != 0 { v } else { t }) + .collect(); + + FlexTensor::new(Bytes::from_elems(result), Layout::contiguous(shape), dtype) +} + +/// Helper to broadcast three tensors to the same shape. +fn broadcast_three( + tensor: FlexTensor, + mask: FlexTensor, + value: FlexTensor, +) -> (FlexTensor, FlexTensor, FlexTensor) { + let target_shape = + crate::ops::expand::broadcast_shape(tensor.layout().shape(), mask.layout().shape()); + let target_shape = crate::ops::expand::broadcast_shape(&target_shape, value.layout().shape()); + + let tensor = if tensor.layout().shape() == &target_shape { + tensor + } else { + crate::ops::expand::expand(tensor, target_shape.clone()) + }; + let mask = if mask.layout().shape() == &target_shape { + mask + } else { + crate::ops::expand::expand(mask, target_shape.clone()) + }; + let value = if value.layout().shape() == &target_shape { + value + } else { + crate::ops::expand::expand(value, target_shape) + }; + + ( + tensor.to_contiguous(), + mask.to_contiguous(), + value.to_contiguous(), + ) +} + +/// Mask where for f32 (SIMD-accelerated). +pub fn mask_where_f32(tensor: FlexTensor, mask: FlexTensor, value: FlexTensor) -> FlexTensor { + #[cfg(feature = "simd")] + { + let dtype = tensor.dtype(); + let (tensor, mask, value) = broadcast_three(tensor, mask, value); + let shape = tensor.layout().shape().clone(); + let len = tensor.storage::().len(); + let mut out = uninit_vec::(len); + crate::simd::mask_where_f32(tensor.storage(), mask.bytes(), value.storage(), &mut out); + FlexTensor::new(Bytes::from_elems(out), Layout::contiguous(shape), dtype) + } + #[cfg(not(feature = "simd"))] + { + mask_where::(tensor, mask, value) + } +} + +/// Mask where for f64 (SIMD-accelerated). +pub fn mask_where_f64(tensor: FlexTensor, mask: FlexTensor, value: FlexTensor) -> FlexTensor { + #[cfg(feature = "simd")] + { + let dtype = tensor.dtype(); + let (tensor, mask, value) = broadcast_three(tensor, mask, value); + let shape = tensor.layout().shape().clone(); + let len = tensor.storage::().len(); + let mut out = uninit_vec::(len); + crate::simd::mask_where_f64(tensor.storage(), mask.bytes(), value.storage(), &mut out); + FlexTensor::new(Bytes::from_elems(out), Layout::contiguous(shape), dtype) + } + #[cfg(not(feature = "simd"))] + { + mask_where::(tensor, mask, value) + } +} + +/// Mask where for f16. +pub fn mask_where_f16(tensor: FlexTensor, mask: FlexTensor, value: FlexTensor) -> FlexTensor { + mask_where::(tensor, mask, value) +} + +/// Mask where for bf16. +pub fn mask_where_bf16(tensor: FlexTensor, mask: FlexTensor, value: FlexTensor) -> FlexTensor { + mask_where::(tensor, mask, value) +} + +/// Mask where for i64 (SIMD-accelerated). +pub fn mask_where_i64(tensor: FlexTensor, mask: FlexTensor, value: FlexTensor) -> FlexTensor { + #[cfg(feature = "simd")] + { + let dtype = tensor.dtype(); + let (tensor, mask, value) = broadcast_three(tensor, mask, value); + let shape = tensor.layout().shape().clone(); + let len = tensor.storage::().len(); + let mut out = uninit_vec::(len); + crate::simd::mask_where_i64(tensor.storage(), mask.bytes(), value.storage(), &mut out); + FlexTensor::new(Bytes::from_elems(out), Layout::contiguous(shape), dtype) + } + #[cfg(not(feature = "simd"))] + { + mask_where::(tensor, mask, value) + } +} + +/// Mask where for bool tensors (SIMD-accelerated). +pub fn mask_where_bool(tensor: FlexTensor, mask: FlexTensor, value: FlexTensor) -> FlexTensor { + // Preserve the input tensor's bool dtype for the output. + let out_dtype = burn_std::BoolDType::from(tensor.dtype()); + #[cfg(feature = "simd")] + { + let (tensor, mask, value) = broadcast_three(tensor, mask, value); + let shape = tensor.layout().shape().clone(); + let len = tensor.bytes().len(); + let mut out = uninit_vec::(len); + crate::simd::mask_where_u8(tensor.bytes(), mask.bytes(), value.bytes(), &mut out); + crate::ops::comparison::make_bool_tensor(out, shape, out_dtype) + } + #[cfg(not(feature = "simd"))] + { + let (tensor, mask, value) = broadcast_three(tensor, mask, value); + let shape = tensor.layout().shape().clone(); + let tensor_data: &[u8] = tensor.bytes(); + let mask_data: &[u8] = mask.bytes(); + let value_data: &[u8] = value.bytes(); + let result: Vec = tensor_data + .iter() + .zip(mask_data.iter()) + .zip(value_data.iter()) + .map(|((&t, &m), &v)| if m != 0 { v } else { t }) + .collect(); + crate::ops::comparison::make_bool_tensor(result, shape, out_dtype) + } +} + +// All mask_fill / mask_where tests, including negative-stride (flipped / +// transposed / narrowed) variants, have been migrated to +// burn-backend-tests (float/ops/mask.rs, int/ops/mask.rs) so they cover +// every backend. The flex-internal dispatchers `mask_fill_f32` / +// `mask_where_f32` / `mask_fill_i64` exercised there indirectly via +// `Flex::float_mask_fill`, `Flex::int_mask_where`, etc. When adding new +// tests, keep them here only if they probe flex-specific behavior that +// cannot be expressed through the public backend API; otherwise add +// them to burn-backend-tests. diff --git a/crates/burn-flex/src/ops/matmul.rs b/crates/burn-flex/src/ops/matmul.rs new file mode 100644 index 0000000..db70258 --- /dev/null +++ b/crates/burn-flex/src/ops/matmul.rs @@ -0,0 +1,976 @@ +//! Matrix multiplication via gemm crate. +//! +//! Optimizations: +//! - Strided gemm for f32/f64/f16 avoids copying non-contiguous tensors +//! - Enables parallelism for large matrices (with rayon feature) +//! - Batched matmul parallelized across batch dimension + +use alloc::vec; +use alloc::vec::Vec; +use burn_backend::{DType, Element}; +use burn_std::{Bytes, Shape, bf16, f16}; + +use crate::{FlexTensor, Layout}; + +/// Types that can be used with gemm-based matmul. +/// Only implement for types that `gemm::gemm` dispatches on via TypeId (f32, f64, f16). +trait GemmScalar: Element + bytemuck::Pod { + fn zero() -> Self; + fn one() -> Self; +} + +impl GemmScalar for f32 { + fn zero() -> Self { + 0.0 + } + fn one() -> Self { + 1.0 + } +} + +impl GemmScalar for f64 { + fn zero() -> Self { + 0.0 + } + fn one() -> Self { + 1.0 + } +} + +impl GemmScalar for f16 { + fn zero() -> Self { + f16::from_f32(0.0) + } + fn one() -> Self { + f16::from_f32(1.0) + } +} + +/// Checked multiplication for matrix sizes, panics on overflow. +#[inline] +fn checked_size(a: usize, b: usize) -> usize { + a.checked_mul(b) + .unwrap_or_else(|| panic!("matmul: matrix size overflow: {a} * {b}")) +} + +/// Threshold for enabling parallelism (M*N*K operations). +/// 192^3 = ~7M ops - balance between 128x128 (no parallel) and 256x256 (parallel) +const PARALLEL_THRESHOLD: usize = 192 * 192 * 192; + +/// Threshold for batch-level parallelism (total ops across all batches). +/// Use batch parallelism when individual matrices are small but total work is large. +#[cfg(feature = "rayon")] +const BATCH_PARALLEL_THRESHOLD: usize = 128 * 128 * 128; // ~2M ops total + +/// Get parallelism setting based on matrix size. +fn get_parallelism(m: usize, n: usize, k: usize) -> gemm::Parallelism { + let ops = m.saturating_mul(n).saturating_mul(k); + if ops >= PARALLEL_THRESHOLD { + #[cfg(feature = "rayon")] + { + gemm::Parallelism::Rayon(0) // 0 = use all available threads + } + #[cfg(not(feature = "rayon"))] + { + gemm::Parallelism::None + } + } else { + gemm::Parallelism::None + } +} + +/// Dispatch matrix multiplication based on dtype. +pub fn matmul(lhs: FlexTensor, rhs: FlexTensor) -> FlexTensor { + assert_eq!(lhs.dtype(), rhs.dtype(), "matmul: dtype mismatch"); + + let lhs_shape = lhs.layout().shape(); + let rhs_shape = rhs.layout().shape(); + let lhs_rank = lhs_shape.num_dims(); + let rhs_rank = rhs_shape.num_dims(); + + assert!(lhs_rank >= 2, "matmul requires at least 2D tensors"); + assert!(rhs_rank >= 2, "matmul requires at least 2D tensors"); + + // Check inner dimensions match: lhs[..., M, K] x rhs[..., K, N] + let k_lhs = lhs_shape[lhs_rank - 1]; + let k_rhs = rhs_shape[rhs_rank - 2]; + assert_eq!(k_lhs, k_rhs, "matmul: inner dimensions must match"); + + match lhs.dtype() { + DType::F32 => matmul_gemm::(lhs, rhs), + DType::F64 => matmul_gemm::(lhs, rhs), + DType::F16 => matmul_gemm::(lhs, rhs), + DType::BF16 => matmul_bf16(lhs, rhs), + _ => panic!("matmul: unsupported dtype {:?}", lhs.dtype()), + } +} + +/// Extract 2D matrix strides from a tensor layout. +/// Returns (row_stride, col_stride) for the last two dimensions. +fn get_2d_strides(layout: &Layout) -> (isize, isize) { + let strides = layout.strides(); + let ndim = strides.len(); + let row_stride = strides[ndim - 2]; + let col_stride = strides[ndim - 1]; + (row_stride, col_stride) +} + +/// Compute broadcast batch dimensions for batched matmul. +/// Returns (broadcast_shape, lhs_strides, rhs_strides) where strides map +/// output batch index to input batch offset (in matrices). +fn broadcast_batch_dims( + lhs_batch: &[usize], + rhs_batch: &[usize], +) -> (Vec, Vec, Vec) { + // Pad shorter batch dims with 1s on the left + let max_len = lhs_batch.len().max(rhs_batch.len()); + let lhs_padded: Vec = (0..max_len) + .map(|i| { + if i < max_len - lhs_batch.len() { + 1 + } else { + lhs_batch[i - (max_len - lhs_batch.len())] + } + }) + .collect(); + let rhs_padded: Vec = (0..max_len) + .map(|i| { + if i < max_len - rhs_batch.len() { + 1 + } else { + rhs_batch[i - (max_len - rhs_batch.len())] + } + }) + .collect(); + + // Compute broadcast shape and strides + let mut broadcast_shape = Vec::with_capacity(max_len); + let mut lhs_strides = Vec::with_capacity(max_len); + let mut rhs_strides = Vec::with_capacity(max_len); + + // Compute strides from right to left + let mut lhs_stride = 1usize; + let mut rhs_stride = 1usize; + for i in (0..max_len).rev() { + let ld = lhs_padded[i]; + let rd = rhs_padded[i]; + debug_assert!( + ld == rd || ld == 1 || rd == 1, + "matmul: batch dimensions not broadcastable: {:?} vs {:?}", + lhs_batch, + rhs_batch + ); + broadcast_shape.push(ld.max(rd)); + // Stride is 0 if dimension is 1 (broadcast), otherwise actual stride + lhs_strides.push(if ld == 1 { 0 } else { lhs_stride }); + rhs_strides.push(if rd == 1 { 0 } else { rhs_stride }); + lhs_stride *= ld; + rhs_stride *= rd; + } + + // Reverse to get correct order + broadcast_shape.reverse(); + lhs_strides.reverse(); + rhs_strides.reverse(); + + (broadcast_shape, lhs_strides, rhs_strides) +} + +/// Convert a flat batch index to input batch offset using broadcast strides. +#[inline] +fn batch_index_to_offset(b: usize, broadcast_shape: &[usize], strides: &[usize]) -> usize { + let mut offset = 0; + let mut remaining = b; + for i in (0..broadcast_shape.len()).rev() { + let idx = remaining % broadcast_shape[i]; + offset += idx * strides[i]; + remaining /= broadcast_shape[i]; + } + offset +} + +/// Compute element-level batch strides for a tensor in a broadcast context. +/// Uses the actual layout strides so non-contiguous (transposed/sliced) tensors +/// work without a copy. Dimensions that are broadcast (size 1) get stride 0. +#[allow(clippy::needless_range_loop)] +fn broadcast_batch_elem_strides( + batch_shape: &[usize], + layout_strides: &[isize], + broadcast_len: usize, +) -> Vec { + let batch_ndim = batch_shape.len(); + debug_assert!(broadcast_len >= batch_ndim); + let mut result = vec![0isize; broadcast_len]; + + for i in 0..broadcast_len { + let batch_idx = i as isize - (broadcast_len as isize - batch_ndim as isize); + if batch_idx >= 0 { + let bi = batch_idx as usize; + if batch_shape[bi] > 1 { + result[i] = layout_strides[bi]; + } + } + } + + result +} + +/// Convert a flat batch index to an element offset using element-level strides. +#[inline] +fn batch_elem_offset(b: usize, broadcast_shape: &[usize], elem_strides: &[isize]) -> isize { + let mut offset: isize = 0; + let mut remaining = b; + for i in (0..broadcast_shape.len()).rev() { + let idx = remaining % broadcast_shape[i]; + offset += idx as isize * elem_strides[i]; + remaining /= broadcast_shape[i]; + } + offset +} + +// ============================================================================ +// Generic gemm-based matmul (f32, f64, f16) +// ============================================================================ + +fn matmul_gemm(lhs: FlexTensor, rhs: FlexTensor) -> FlexTensor { + let lhs_rank = lhs.layout().shape().num_dims(); + let rhs_rank = rhs.layout().shape().num_dims(); + + if lhs_rank == 2 && rhs_rank == 2 { + matmul_2d_strided::(lhs, rhs) + } else { + matmul_batched_gemm::(lhs, rhs) + } +} + +/// 2D matmul with strided support: [M, K] x [K, N] -> [M, N] +fn matmul_2d_strided(lhs: FlexTensor, rhs: FlexTensor) -> FlexTensor { + let lhs_shape = lhs.layout().shape(); + let rhs_shape = rhs.layout().shape(); + + let m = lhs_shape[0]; + let k = lhs_shape[1]; + let n = rhs_shape[1]; + + let (lhs_row_stride, lhs_col_stride) = get_2d_strides(lhs.layout()); + let (rhs_row_stride, rhs_col_stride) = get_2d_strides(rhs.layout()); + + let lhs_data: &[T] = lhs.storage(); + let rhs_data: &[T] = rhs.storage(); + let lhs_ptr = unsafe { lhs_data.as_ptr().add(lhs.layout().start_offset()) }; + let rhs_ptr = unsafe { rhs_data.as_ptr().add(rhs.layout().start_offset()) }; + + let out_shape = Shape::from(vec![m, n]); + let mut output = FlexTensor::empty(out_shape, T::dtype()); + let out_data: &mut [T] = output.storage_mut(); + + let parallelism = get_parallelism(m, n, k); + + unsafe { + gemm_call( + m, + n, + k, + out_data.as_mut_ptr(), + 1, + n as isize, + lhs_ptr, + lhs_col_stride, + lhs_row_stride, + rhs_ptr, + rhs_col_stride, + rhs_row_stride, + parallelism, + ); + } + + output +} + +/// Strided gemm call for one matrix. Wraps `gemm::gemm` with GemmScalar zero/one. +#[inline] +#[allow(clippy::too_many_arguments)] +unsafe fn gemm_call( + m: usize, + n: usize, + k: usize, + out: *mut T, + out_cs: isize, + out_rs: isize, + lhs: *const T, + lhs_cs: isize, + lhs_rs: isize, + rhs: *const T, + rhs_cs: isize, + rhs_rs: isize, + parallelism: gemm::Parallelism, +) { + unsafe { + gemm::gemm( + m, + n, + k, + out, + out_cs, + out_rs, + false, + lhs, + lhs_cs, + lhs_rs, + rhs, + rhs_cs, + rhs_rs, + T::zero(), + T::one(), + false, + false, + false, + parallelism, + ); + } +} + +/// Batched matmul: [B..., M, K] x [B..., K, N] -> [B..., M, N] +/// Supports broadcasting on batch dimensions and strided (non-contiguous) inputs. +fn matmul_batched_gemm(lhs: FlexTensor, rhs: FlexTensor) -> FlexTensor { + let lhs_shape = lhs.layout().shape(); + let rhs_shape = rhs.layout().shape(); + let lhs_rank = lhs_shape.num_dims(); + let rhs_rank = rhs_shape.num_dims(); + + let m = lhs_shape[lhs_rank - 2]; + let k = lhs_shape[lhs_rank - 1]; + let n = rhs_shape[rhs_rank - 1]; + + let lhs_batch: Vec = lhs_shape[..lhs_rank - 2].to_vec(); + let rhs_batch: Vec = rhs_shape[..rhs_rank - 2].to_vec(); + + let (broadcast_shape, _, _) = broadcast_batch_dims(&lhs_batch, &rhs_batch); + let batch_size: usize = broadcast_shape.iter().product(); + let broadcast_len = broadcast_shape.len(); + + let lhs_batch_strides = + broadcast_batch_elem_strides(&lhs_batch, lhs.layout().strides(), broadcast_len); + let rhs_batch_strides = + broadcast_batch_elem_strides(&rhs_batch, rhs.layout().strides(), broadcast_len); + + let (lhs_row_stride, lhs_col_stride) = get_2d_strides(lhs.layout()); + let (rhs_row_stride, rhs_col_stride) = get_2d_strides(rhs.layout()); + + let out_matrix_size = checked_size(m, n); + + let mut out_dims = broadcast_shape.clone(); + out_dims.push(m); + out_dims.push(n); + let out_shape = Shape::from(out_dims); + + let mut output = FlexTensor::empty(out_shape, T::dtype()); + + let lhs_data: &[T] = lhs.storage(); + let rhs_data: &[T] = rhs.storage(); + let lhs_start = lhs.layout().start_offset() as isize; + let rhs_start = rhs.layout().start_offset() as isize; + let out_data: &mut [T] = output.storage_mut(); + + let per_matrix_ops = m.saturating_mul(n).saturating_mul(k); + + // Closure: run gemm for one batch slice at the given pointers + let run_one = |out_ptr: *mut T, b: usize, parallelism: gemm::Parallelism| { + let lhs_off = lhs_start + batch_elem_offset(b, &broadcast_shape, &lhs_batch_strides); + let rhs_off = rhs_start + batch_elem_offset(b, &broadcast_shape, &rhs_batch_strides); + unsafe { + gemm_call::( + m, + n, + k, + out_ptr, + 1, + n as isize, + lhs_data.as_ptr().offset(lhs_off), + lhs_col_stride, + lhs_row_stride, + rhs_data.as_ptr().offset(rhs_off), + rhs_col_stride, + rhs_row_stride, + parallelism, + ); + } + }; + + // Strategy: + // 1. Large matrices: let gemm parallelize internally + // 2. Small matrices, large batch: parallelize batch loop + // 3. Small total work: single-threaded + #[cfg(feature = "rayon")] + { + let total_ops = batch_size.saturating_mul(per_matrix_ops); + let prefer_batch_parallel = batch_size >= 4 && total_ops >= BATCH_PARALLEL_THRESHOLD; + + if per_matrix_ops >= PARALLEL_THRESHOLD && !prefer_batch_parallel { + let parallelism = gemm::Parallelism::Rayon(0); + for b in 0..batch_size { + run_one(out_data[b * out_matrix_size..].as_mut_ptr(), b, parallelism); + } + } else if total_ops >= BATCH_PARALLEL_THRESHOLD && batch_size > 1 { + use rayon::prelude::*; + + out_data + .par_chunks_mut(out_matrix_size) + .enumerate() + .for_each(|(b, out_chunk)| { + run_one(out_chunk.as_mut_ptr(), b, gemm::Parallelism::None); + }); + } else { + for b in 0..batch_size { + run_one( + out_data[b * out_matrix_size..].as_mut_ptr(), + b, + gemm::Parallelism::None, + ); + } + } + } + + #[cfg(not(feature = "rayon"))] + { + let _ = per_matrix_ops; + for b in 0..batch_size { + run_one( + out_data[b * out_matrix_size..].as_mut_ptr(), + b, + gemm::Parallelism::None, + ); + } + } + + output +} + +// ============================================================================ +// bf16 matmul (via f32 conversion) +// ============================================================================ + +fn matmul_bf16(lhs: FlexTensor, rhs: FlexTensor) -> FlexTensor { + let lhs = lhs.to_contiguous(); + let rhs = rhs.to_contiguous(); + + let lhs_shape = lhs.layout().shape(); + let rhs_shape = rhs.layout().shape(); + + // Convert bf16 -> f32 + let lhs_f32: Vec = lhs.storage::().iter().map(|x| x.to_f32()).collect(); + let rhs_f32: Vec = rhs.storage::().iter().map(|x| x.to_f32()).collect(); + + // Create f32 tensors + let lhs_f32_tensor = FlexTensor::new( + Bytes::from_elems(lhs_f32), + Layout::contiguous(lhs_shape.clone()), + DType::F32, + ); + let rhs_f32_tensor = FlexTensor::new( + Bytes::from_elems(rhs_f32), + Layout::contiguous(rhs_shape.clone()), + DType::F32, + ); + + // Compute matmul in f32 + let result_f32 = matmul_gemm::(lhs_f32_tensor, rhs_f32_tensor); + + // Convert f32 -> bf16 + let result_bf16: Vec = result_f32 + .storage::() + .iter() + .map(|x| bf16::from_f32(*x)) + .collect(); + + FlexTensor::new( + Bytes::from_elems(result_bf16), + result_f32.layout().clone(), + DType::BF16, + ) +} + +// ============================================================================ +// Integer matmul (naive, with optional SIMD for i32) +// ============================================================================ + +/// Integer matrix multiplication dispatch. +pub fn int_matmul(lhs: FlexTensor, rhs: FlexTensor) -> FlexTensor { + assert_eq!(lhs.dtype(), rhs.dtype(), "int_matmul: dtype mismatch"); + + let lhs_shape = lhs.layout().shape(); + let rhs_shape = rhs.layout().shape(); + let lhs_rank = lhs_shape.num_dims(); + let rhs_rank = rhs_shape.num_dims(); + + assert!(lhs_rank >= 2, "int_matmul requires at least 2D tensors"); + assert!(rhs_rank >= 2, "int_matmul requires at least 2D tensors"); + + let k_lhs = lhs_shape[lhs_rank - 1]; + let k_rhs = rhs_shape[rhs_rank - 2]; + assert_eq!(k_lhs, k_rhs, "int_matmul: inner dimensions must match"); + + match lhs.dtype() { + DType::I32 => matmul_i32(lhs, rhs), + DType::I64 => matmul_i64(lhs, rhs), + _ => panic!("int_matmul: unsupported dtype {:?}", lhs.dtype()), + } +} + +/// i32 matmul using naive triple loop with SIMD dot product. +fn matmul_i32(lhs: FlexTensor, rhs: FlexTensor) -> FlexTensor { + let lhs = lhs.to_contiguous(); + let rhs = rhs.to_contiguous(); + + let lhs_shape = lhs.layout().shape(); + let rhs_shape = rhs.layout().shape(); + let lhs_rank = lhs_shape.num_dims(); + let rhs_rank = rhs_shape.num_dims(); + + if lhs_rank == 2 && rhs_rank == 2 { + matmul_2d_i32(&lhs, &rhs) + } else { + matmul_batched_i32(lhs, rhs) + } +} + +/// 2D i32 matmul: [M, K] x [K, N] -> [M, N] +/// Transposes rhs to enable contiguous access for dot product. +fn matmul_2d_i32(lhs: &FlexTensor, rhs: &FlexTensor) -> FlexTensor { + let lhs_shape = lhs.layout().shape(); + let rhs_shape = rhs.layout().shape(); + + let m = lhs_shape[0]; + let k = lhs_shape[1]; + let n = rhs_shape[1]; + + let lhs_data: &[i32] = lhs.storage(); + let rhs_data: &[i32] = rhs.storage(); + + // Transpose rhs [K, N] -> [N, K] for contiguous column access + let mut rhs_t = vec![0i32; k * n]; + for i in 0..k { + for j in 0..n { + rhs_t[j * k + i] = rhs_data[i * n + j]; + } + } + + let mut output = vec![0i32; m * n]; + + // Now both lhs rows and rhs columns (transposed rows) are contiguous + for i in 0..m { + let lhs_row = &lhs_data[i * k..(i + 1) * k]; + for j in 0..n { + let rhs_col = &rhs_t[j * k..(j + 1) * k]; + output[i * n + j] = dot_i32(lhs_row, rhs_col); + } + } + + let out_shape = Shape::from(vec![m, n]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + DType::I32, + ) +} + +/// Dot product for i32 slices. Uses macerator SIMD when the `simd` feature is enabled. +#[inline] +fn dot_i32(a: &[i32], b: &[i32]) -> i32 { + debug_assert_eq!(a.len(), b.len()); + + #[cfg(feature = "simd")] + { + dot_i32_simd(a, b) + } + + #[cfg(not(feature = "simd"))] + { + dot_i32_scalar(a, b) + } +} + +#[cfg(not(feature = "simd"))] +#[inline] +fn dot_i32_scalar(a: &[i32], b: &[i32]) -> i32 { + let mut sum = 0i32; + for i in 0..a.len() { + sum = sum.wrapping_add(a[i].wrapping_mul(b[i])); + } + sum +} + +#[cfg(feature = "simd")] +#[macerator::with_simd] +fn dot_i32_simd(a: &[i32], b: &[i32]) -> i32 { + use macerator::{Scalar, VMulAdd, vload_unaligned}; + + let lanes = i32::lanes::(); + let len = a.len(); + let simd_len = len / lanes * lanes; + let mut acc = 0i32.splat::(); + + let mut i = 0; + while i < simd_len { + let va = unsafe { vload_unaligned(a.as_ptr().add(i)) }; + let vb = unsafe { vload_unaligned(b.as_ptr().add(i)) }; + acc = i32::vmul_add(va, vb, acc); + i += lanes; + } + + let mut sum = acc.reduce_add(); + while i < len { + sum = sum.wrapping_add(a[i].wrapping_mul(b[i])); + i += 1; + } + sum +} + +/// Batched i32 matmul: [B..., M, K] x [B..., K, N] -> [B..., M, N] +/// +/// Uses naive triple-loop with SIMD dot product and batch-level parallelism. +fn matmul_batched_i32(lhs: FlexTensor, rhs: FlexTensor) -> FlexTensor { + let lhs_shape = lhs.layout().shape(); + let rhs_shape = rhs.layout().shape(); + let lhs_rank = lhs_shape.num_dims(); + let rhs_rank = rhs_shape.num_dims(); + + let m = lhs_shape[lhs_rank - 2]; + let k = lhs_shape[lhs_rank - 1]; + let n = rhs_shape[rhs_rank - 1]; + + let lhs_batch: Vec = lhs_shape[..lhs_rank - 2].to_vec(); + let rhs_batch: Vec = rhs_shape[..rhs_rank - 2].to_vec(); + + let (broadcast_shape, lhs_strides, rhs_strides) = broadcast_batch_dims(&lhs_batch, &rhs_batch); + + let batch_size: usize = broadcast_shape.iter().product(); + let rhs_batch_size: usize = rhs_batch.iter().product(); + let lhs_matrix_size = checked_size(m, k); + let rhs_matrix_size = checked_size(k, n); + let out_matrix_size = checked_size(m, n); + + let mut out_dims = broadcast_shape.clone(); + out_dims.push(m); + out_dims.push(n); + let out_shape = Shape::from(out_dims); + + let lhs_data: &[i32] = lhs.storage(); + let rhs_data: &[i32] = rhs.storage(); + + // Transpose rhs per actual rhs batch: [B_rhs, K, N] -> [B_rhs, N, K] + let mut rhs_transposed = vec![0i32; rhs_batch_size * n * k]; + for b in 0..rhs_batch_size { + let src_offset = b * rhs_matrix_size; + let dst_offset = b * n * k; + for i in 0..k { + for j in 0..n { + rhs_transposed[dst_offset + j * k + i] = rhs_data[src_offset + i * n + j]; + } + } + } + + let mut output = vec![0i32; batch_size * out_matrix_size]; + + let run_one = |b: usize, out_slice: &mut [i32]| { + let lhs_batch_idx = batch_index_to_offset(b, &broadcast_shape, &lhs_strides); + let rhs_batch_idx = batch_index_to_offset(b, &broadcast_shape, &rhs_strides); + let lhs_offset = lhs_batch_idx * lhs_matrix_size; + let rhs_t_offset = rhs_batch_idx * n * k; + + let lhs_slice = &lhs_data[lhs_offset..lhs_offset + lhs_matrix_size]; + let rhs_t_slice = &rhs_transposed[rhs_t_offset..rhs_t_offset + n * k]; + + for i in 0..m { + let lhs_row = &lhs_slice[i * k..(i + 1) * k]; + for j in 0..n { + let rhs_col = &rhs_t_slice[j * k..(j + 1) * k]; + out_slice[i * n + j] = dot_i32(lhs_row, rhs_col); + } + } + }; + + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + output + .par_chunks_mut(out_matrix_size) + .enumerate() + .for_each(|(b, out_slice)| run_one(b, out_slice)); + } + + #[cfg(not(feature = "rayon"))] + { + for b in 0..batch_size { + let offset = b * out_matrix_size; + run_one(b, &mut output[offset..offset + out_matrix_size]); + } + } + + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + DType::I32, + ) +} + +/// i64 matmul using naive triple loop. +fn matmul_i64(lhs: FlexTensor, rhs: FlexTensor) -> FlexTensor { + let lhs = lhs.to_contiguous(); + let rhs = rhs.to_contiguous(); + + let lhs_shape = lhs.layout().shape(); + let rhs_shape = rhs.layout().shape(); + let lhs_rank = lhs_shape.num_dims(); + let rhs_rank = rhs_shape.num_dims(); + + if lhs_rank == 2 && rhs_rank == 2 { + matmul_2d_i64(&lhs, &rhs) + } else { + matmul_batched_i64(lhs, rhs) + } +} + +/// 2D i64 matmul: [M, K] x [K, N] -> [M, N] +fn matmul_2d_i64(lhs: &FlexTensor, rhs: &FlexTensor) -> FlexTensor { + let lhs_shape = lhs.layout().shape(); + let rhs_shape = rhs.layout().shape(); + + let m = lhs_shape[0]; + let k = lhs_shape[1]; + let n = rhs_shape[1]; + + let lhs_data: &[i64] = lhs.storage(); + let rhs_data: &[i64] = rhs.storage(); + + let mut output = vec![0i64; m * n]; + + for i in 0..m { + for j in 0..n { + let mut sum = 0i64; + for l in 0..k { + sum = sum.wrapping_add(lhs_data[i * k + l].wrapping_mul(rhs_data[l * n + j])); + } + output[i * n + j] = sum; + } + } + + let out_shape = Shape::from(vec![m, n]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + DType::I64, + ) +} + +/// Batched i64 matmul with broadcast support +fn matmul_batched_i64(lhs: FlexTensor, rhs: FlexTensor) -> FlexTensor { + let lhs_shape = lhs.layout().shape(); + let rhs_shape = rhs.layout().shape(); + let lhs_rank = lhs_shape.num_dims(); + let rhs_rank = rhs_shape.num_dims(); + + let m = lhs_shape[lhs_rank - 2]; + let k = lhs_shape[lhs_rank - 1]; + let n = rhs_shape[rhs_rank - 1]; + + let lhs_batch: Vec = lhs_shape[..lhs_rank - 2].to_vec(); + let rhs_batch: Vec = rhs_shape[..rhs_rank - 2].to_vec(); + + // Compute broadcast batch dimensions + let (broadcast_shape, lhs_strides, rhs_strides) = broadcast_batch_dims(&lhs_batch, &rhs_batch); + + let batch_size: usize = broadcast_shape.iter().product(); + let lhs_matrix_size = checked_size(m, k); + let rhs_matrix_size = checked_size(k, n); + let out_matrix_size = checked_size(m, n); + + let mut out_dims = broadcast_shape.clone(); + out_dims.push(m); + out_dims.push(n); + let out_shape = Shape::from(out_dims); + + let lhs_data: &[i64] = lhs.storage(); + let rhs_data: &[i64] = rhs.storage(); + + let mut output = vec![0i64; batch_size * out_matrix_size]; + + for b in 0..batch_size { + let lhs_batch_idx = batch_index_to_offset(b, &broadcast_shape, &lhs_strides); + let rhs_batch_idx = batch_index_to_offset(b, &broadcast_shape, &rhs_strides); + let lhs_offset = lhs_batch_idx * lhs_matrix_size; + let rhs_offset = rhs_batch_idx * rhs_matrix_size; + let out_offset = b * out_matrix_size; + + for i in 0..m { + for j in 0..n { + let mut sum = 0i64; + for l in 0..k { + let lhs_idx = lhs_offset + i * k + l; + let rhs_idx = rhs_offset + l * n + j; + sum = sum.wrapping_add(lhs_data[lhs_idx].wrapping_mul(rhs_data[rhs_idx])); + } + output[out_offset + i * n + j] = sum; + } + } + } + + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + DType::I64, + ) +} + +// ============================================================================ +// Tests +// ============================================================================ + +// Tests kept here exercise flex-specific behavior of the matmul kernel: +// dtype-specific storage paths (F64, F16, BF16) that the generic +// FloatElem-parameterized backend-tests cannot reach. Plain contiguous +// F32/I32/I64 matmul, stride-through-matmul variants (transposed / +// swap_dims / broadcast-transposed), and other generic coverage have +// been migrated to burn-backend-tests so they run against every backend. +// When adding new tests, keep them here only if they probe a flex +// dtype-storage path; otherwise add them to +// crates/burn-backend-tests/tests/tensor/float/ops/matmul.rs. +#[cfg(test)] +mod tests { + use alloc::vec; + use burn_backend::TensorData; + use burn_backend::ops::FloatTensorOps; + use burn_std::{bf16, f16}; + + use crate::{Flex, FlexTensor}; + + #[test] + fn test_matmul_f64() { + let lhs = FlexTensor::from_data(TensorData::new(vec![1.0f64, 2.0, 3.0, 4.0], [2, 2])); + let rhs = FlexTensor::from_data(TensorData::new(vec![5.0f64, 6.0, 7.0, 8.0], [2, 2])); + + let result = Flex::float_matmul(lhs, rhs); + let values: Vec = result.into_data().to_vec().unwrap(); + + assert_eq!(values, vec![19.0, 22.0, 43.0, 50.0]); + } + + #[test] + fn test_matmul_f16() { + let lhs_vals: Vec = [1.0f32, 2.0, 3.0, 4.0] + .iter() + .copied() + .map(f16::from_f32) + .collect(); + let rhs_vals: Vec = [5.0f32, 6.0, 7.0, 8.0] + .iter() + .copied() + .map(f16::from_f32) + .collect(); + + let lhs = FlexTensor::from_data(TensorData::new(lhs_vals, [2, 2])); + let rhs = FlexTensor::from_data(TensorData::new(rhs_vals, [2, 2])); + + let result = Flex::float_matmul(lhs, rhs); + let values: Vec = result.into_data().to_vec().unwrap(); + + let expected = [19.0f32, 22.0, 43.0, 50.0]; + for (a, e) in values.iter().zip(expected.iter()) { + assert!((a.to_f32() - e).abs() < 0.1, "f16 matmul mismatch"); + } + } + + #[test] + fn test_matmul_bf16() { + let lhs_vals: Vec = [1.0f32, 2.0, 3.0, 4.0] + .iter() + .copied() + .map(bf16::from_f32) + .collect(); + let rhs_vals: Vec = [5.0f32, 6.0, 7.0, 8.0] + .iter() + .copied() + .map(bf16::from_f32) + .collect(); + + let lhs = FlexTensor::from_data(TensorData::new(lhs_vals, [2, 2])); + let rhs = FlexTensor::from_data(TensorData::new(rhs_vals, [2, 2])); + + let result = Flex::float_matmul(lhs, rhs); + let values: Vec = result.into_data().to_vec().unwrap(); + + let expected = [19.0f32, 22.0, 43.0, 50.0]; + for (a, e) in values.iter().zip(expected.iter()) { + assert!((a.to_f32() - e).abs() < 0.5, "bf16 matmul mismatch"); + } + } + + #[test] + fn test_matmul_batched_transposed_f64() { + // Non-contiguous (swap_dims) batched matmul on the F64 dtype path. + let q_data = TensorData::new(vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], [2, 2, 2]); + let k_data = TensorData::new(vec![1.0f64, 0.0, 0.0, 1.0, 2.0, 0.0, 0.0, 2.0], [2, 2, 2]); + + let q = FlexTensor::from_data(q_data.clone()); + let k = FlexTensor::from_data(k_data.clone()); + let k_t = k.transpose(1, 2); + let result = Flex::float_matmul(q, k_t); + + let q2 = FlexTensor::from_data(q_data); + let k2 = FlexTensor::from_data(k_data) + .transpose(1, 2) + .to_contiguous(); + let expected = Flex::float_matmul(q2, k2); + + let values: Vec = result.into_data().to_vec().unwrap(); + let expected: Vec = expected.into_data().to_vec().unwrap(); + assert_eq!(values, expected); + } + + #[test] + fn test_matmul_batched_transposed_f16() { + // Non-contiguous (swap_dims) batched matmul on the F16 dtype path. + let f = f16::from_f32; + let q_data = TensorData::new( + vec![ + f(1.0), + f(2.0), + f(3.0), + f(4.0), + f(5.0), + f(6.0), + f(7.0), + f(8.0), + ], + [2, 2, 2], + ); + let k_data = TensorData::new( + vec![ + f(1.0), + f(0.0), + f(0.0), + f(1.0), + f(2.0), + f(0.0), + f(0.0), + f(2.0), + ], + [2, 2, 2], + ); + + let q = FlexTensor::from_data(q_data.clone()); + let k = FlexTensor::from_data(k_data.clone()); + let k_t = k.transpose(1, 2); + let result = Flex::float_matmul(q, k_t); + + let q2 = FlexTensor::from_data(q_data); + let k2 = FlexTensor::from_data(k_data) + .transpose(1, 2) + .to_contiguous(); + let expected = Flex::float_matmul(q2, k2); + + let values: Vec = result.into_data().to_vec().unwrap(); + let expected: Vec = expected.into_data().to_vec().unwrap(); + assert_eq!(values, expected); + } +} diff --git a/crates/burn-flex/src/ops/mod.rs b/crates/burn-flex/src/ops/mod.rs new file mode 100644 index 0000000..35f754a --- /dev/null +++ b/crates/burn-flex/src/ops/mod.rs @@ -0,0 +1,122 @@ +//! Backend operations implementations. + +use alloc::borrow::Cow; +use burn_backend::DType; +use burn_std::{bf16, f16}; + +use crate::FlexTensor; + +/// The `DType` that matches `isize` on the current platform. +#[cfg(target_pointer_width = "64")] +pub(crate) const INDEX_DTYPE: DType = DType::I64; +/// The `DType` that matches `isize` on the current platform. +#[cfg(target_pointer_width = "32")] +pub(crate) const INDEX_DTYPE: DType = DType::I32; + +/// Minimum total element count for rayon fan-out on per-row or per-chunk loops. +/// Below this, task-dispatch overhead dominates the per-unit work. +#[cfg(feature = "rayon")] +pub(crate) const PARALLEL_THRESHOLD: usize = 256 * 1024; + +/// Wrapper for raw mutable pointers that can be sent across rayon threads. +/// +/// # Safety +/// +/// The caller must ensure: +/// - The pointer remains valid for the lifetime of all uses +/// - No two threads write to the same offset +/// - No references to the underlying data exist during writes +#[cfg(feature = "rayon")] +pub(crate) struct SendMutPtr(*mut T); + +#[cfg(feature = "rayon")] +unsafe impl Send for SendMutPtr {} +#[cfg(feature = "rayon")] +unsafe impl Sync for SendMutPtr {} + +#[cfg(feature = "rayon")] +impl SendMutPtr { + pub(crate) fn new(ptr: *mut T) -> Self { + Self(ptr) + } + + /// Write `val` at the given element offset. + /// + /// # Safety + /// Offset must be in bounds and no other thread may write to the same offset. + pub(crate) unsafe fn write(&self, offset: usize, val: T) { + unsafe { self.0.add(offset).write(val) } + } + + /// Returns the raw pointer offset by `offset` elements. + /// + /// # Safety + /// Offset must be in bounds. + pub(crate) unsafe fn ptr_add(&self, offset: usize) -> *mut T { + unsafe { self.0.add(offset) } + } +} + +/// Read a float tensor's storage as f32 values, regardless of source dtype. +/// Returns a borrowed slice for F32 (zero-copy) and an owned Vec for other +/// float dtypes (F64, F16, BF16). +/// +/// Returns elements in underlying buffer order. Callers needing logical +/// iteration order must call `to_contiguous()` first. +/// +/// # Panics +/// Panics if the tensor's dtype is not one of F32, F64, F16, or BF16. +pub(crate) fn float_storage_as_f32(tensor: &FlexTensor) -> Cow<'_, [f32]> { + match tensor.dtype() { + DType::F32 => Cow::Borrowed(tensor.storage::()), + DType::F64 => Cow::Owned(tensor.storage::().iter().map(|&x| x as f32).collect()), + DType::F16 => Cow::Owned( + tensor + .storage::() + .iter() + .map(|x| f32::from(*x)) + .collect(), + ), + DType::BF16 => Cow::Owned( + tensor + .storage::() + .iter() + .map(|x| f32::from(*x)) + .collect(), + ), + other => panic!("float_storage_as_f32: unsupported dtype {:?}", other), + } +} + +pub mod activation; +pub mod attention; +pub mod binary; +mod bool; +pub mod cat; +pub mod comparison; +#[macro_use] +mod conv_common; +pub mod conv; +pub mod conv_transpose; +pub mod cumulative; +pub mod deform_conv; +pub mod expand; +pub mod fft; +pub mod flip; +mod float; +pub mod gather_scatter; +pub mod grid_sample; +mod int; +pub mod interpolate; +pub mod mask; +pub mod matmul; +mod module; +pub mod pool; +mod qtensor; +pub mod reduce; +pub mod repeat_dim; +pub mod slice; +pub mod sort; +mod transaction; +pub mod unary; +pub mod unfold; diff --git a/crates/burn-flex/src/ops/module.rs b/crates/burn-flex/src/ops/module.rs new file mode 100644 index 0000000..4d80409 --- /dev/null +++ b/crates/burn-flex/src/ops/module.rs @@ -0,0 +1,776 @@ +//! Module operations for the Flex backend. +//! +//! These operations power neural network modules like convolutions and pooling. + +use crate::ops::{conv, conv_transpose, deform_conv, interpolate, pool}; +use crate::{Flex, FlexTensor, Layout}; +use burn_backend::{ + DType, Element, TensorMetadata, + ops::{ + AttentionModuleOptions, ConvOptions, ConvTransposeOptions, DeformConv2dBackward, + DeformConvOptions, FloatTensorOps, IntTensorOps, InterpolateMode, InterpolateOptions, + MaxPool2dBackward, MaxPool2dWithIndices, ModuleOps, + }, + tensor::{BoolTensor, FloatTensor, IntTensor}, +}; +use burn_std::{Bytes, IntDType, Shape}; +use bytemuck::Pod; + +/// Cast a tensor from half-precision type E to f32. +pub(crate) fn cast_to_f32( + tensor: FlexTensor, + to_f32: fn(E) -> f32, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let data: &[E] = tensor.storage(); + let f32_data: alloc::vec::Vec = data.iter().map(|&v| to_f32(v)).collect(); + let bytes = Bytes::from_elems(f32_data); + FlexTensor::new(bytes, Layout::contiguous(shape), DType::F32) +} + +/// Cast a tensor from f32 back to half-precision type E. +pub(crate) fn cast_from_f32( + tensor: FlexTensor, + from_f32: fn(f32) -> E, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let data: &[f32] = tensor.storage(); + let half_data: alloc::vec::Vec = data.iter().map(|&v| from_f32(v)).collect(); + let bytes = Bytes::from_elems(half_data); + FlexTensor::new(bytes, Layout::contiguous(shape), E::dtype()) +} + +impl ModuleOps for Flex { + fn conv1d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<1>, + ) -> FloatTensor { + match x.dtype() { + DType::F32 => conv::conv1d_f32(x, weight, bias, &options), + DType::F64 => conv::conv1d_f64(x, weight, bias, &options), + DType::F16 => conv::conv1d_f16(x, weight, bias, &options), + DType::BF16 => conv::conv1d_bf16(x, weight, bias, &options), + dtype => panic!("conv1d: unsupported dtype {:?}", dtype), + } + } + + fn conv2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<2>, + ) -> FloatTensor { + match x.dtype() { + DType::F32 => conv::conv2d_f32(x, weight, bias, &options), + DType::F64 => conv::conv2d_f64(x, weight, bias, &options), + DType::F16 => conv::conv2d_f16(x, weight, bias, &options), + DType::BF16 => conv::conv2d_bf16(x, weight, bias, &options), + dtype => panic!("conv2d: unsupported dtype {:?}", dtype), + } + } + + fn deform_conv2d( + x: FloatTensor, + offset: FloatTensor, + weight: FloatTensor, + mask: Option>, + bias: Option>, + options: DeformConvOptions<2>, + ) -> FloatTensor { + match x.dtype() { + DType::F32 => deform_conv::deform_conv2d_f32( + x, + offset, + weight, + mask, + bias, + options.stride, + options.padding, + options.dilation, + options.weight_groups, + options.offset_groups, + ), + DType::F64 => deform_conv::deform_conv2d_f64( + x, + offset, + weight, + mask, + bias, + options.stride, + options.padding, + options.dilation, + options.weight_groups, + options.offset_groups, + ), + DType::F16 => { + use burn_std::f16; + let result = deform_conv::deform_conv2d_f32( + cast_to_f32(x, f16::to_f32), + cast_to_f32(offset, f16::to_f32), + cast_to_f32(weight, f16::to_f32), + mask.map(|m| cast_to_f32(m, f16::to_f32)), + bias.map(|b| cast_to_f32(b, f16::to_f32)), + options.stride, + options.padding, + options.dilation, + options.weight_groups, + options.offset_groups, + ); + cast_from_f32(result, f16::from_f32) + } + DType::BF16 => { + use burn_std::bf16; + let result = deform_conv::deform_conv2d_f32( + cast_to_f32(x, bf16::to_f32), + cast_to_f32(offset, bf16::to_f32), + cast_to_f32(weight, bf16::to_f32), + mask.map(|m| cast_to_f32(m, bf16::to_f32)), + bias.map(|b| cast_to_f32(b, bf16::to_f32)), + options.stride, + options.padding, + options.dilation, + options.weight_groups, + options.offset_groups, + ); + cast_from_f32(result, bf16::from_f32) + } + dtype => panic!("deform_conv2d: unsupported dtype {:?}", dtype), + } + } + + fn deform_conv2d_backward( + x: FloatTensor, + offset: FloatTensor, + weight: FloatTensor, + mask: Option>, + bias: Option>, + output_grad: FloatTensor, + options: DeformConvOptions<2>, + ) -> DeformConv2dBackward { + let (x_grad, offset_grad, weight_grad, mask_grad, bias_grad) = match x.dtype() { + DType::F32 => deform_conv::deform_conv2d_backward_f32( + x, + offset, + weight, + mask, + bias, + output_grad, + options.stride, + options.padding, + options.dilation, + options.weight_groups, + options.offset_groups, + ), + DType::F16 => { + use burn_std::f16; + let (xg, og, wg, mg, bg) = deform_conv::deform_conv2d_backward_f32( + cast_to_f32(x, f16::to_f32), + cast_to_f32(offset, f16::to_f32), + cast_to_f32(weight, f16::to_f32), + mask.map(|m| cast_to_f32(m, f16::to_f32)), + bias.map(|b| cast_to_f32(b, f16::to_f32)), + cast_to_f32(output_grad, f16::to_f32), + options.stride, + options.padding, + options.dilation, + options.weight_groups, + options.offset_groups, + ); + ( + cast_from_f32(xg, f16::from_f32), + cast_from_f32(og, f16::from_f32), + cast_from_f32(wg, f16::from_f32), + mg.map(|m| cast_from_f32(m, f16::from_f32)), + bg.map(|b| cast_from_f32(b, f16::from_f32)), + ) + } + DType::BF16 => { + use burn_std::bf16; + let (xg, og, wg, mg, bg) = deform_conv::deform_conv2d_backward_f32( + cast_to_f32(x, bf16::to_f32), + cast_to_f32(offset, bf16::to_f32), + cast_to_f32(weight, bf16::to_f32), + mask.map(|m| cast_to_f32(m, bf16::to_f32)), + bias.map(|b| cast_to_f32(b, bf16::to_f32)), + cast_to_f32(output_grad, bf16::to_f32), + options.stride, + options.padding, + options.dilation, + options.weight_groups, + options.offset_groups, + ); + ( + cast_from_f32(xg, bf16::from_f32), + cast_from_f32(og, bf16::from_f32), + cast_from_f32(wg, bf16::from_f32), + mg.map(|m| cast_from_f32(m, bf16::from_f32)), + bg.map(|b| cast_from_f32(b, bf16::from_f32)), + ) + } + // f64 backward computed via f32: precision loss for large/small values. + // A native f64 implementation would require duplicating ~400 lines of + // deform_conv2d_backward. f64 deform_conv is rare in practice. + DType::F64 => { + let to = |v: f64| v as f32; + let from = |v: f32| v as f64; + let (xg, og, wg, mg, bg) = deform_conv::deform_conv2d_backward_f32( + cast_to_f32(x, to), + cast_to_f32(offset, to), + cast_to_f32(weight, to), + mask.map(|m| cast_to_f32(m, to)), + bias.map(|b| cast_to_f32(b, to)), + cast_to_f32(output_grad, to), + options.stride, + options.padding, + options.dilation, + options.weight_groups, + options.offset_groups, + ); + ( + cast_from_f32(xg, from), + cast_from_f32(og, from), + cast_from_f32(wg, from), + mg.map(|m| cast_from_f32(m, from)), + bg.map(|b| cast_from_f32(b, from)), + ) + } + dtype => panic!("deform_conv2d_backward: unsupported dtype {:?}", dtype), + }; + DeformConv2dBackward::new(x_grad, offset_grad, weight_grad, mask_grad, bias_grad) + } + + fn conv3d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<3>, + ) -> FloatTensor { + match x.dtype() { + DType::F32 => conv::conv3d_f32(x, weight, bias, &options), + DType::F64 => conv::conv3d_f64(x, weight, bias, &options), + DType::F16 => conv::conv3d_f16(x, weight, bias, &options), + DType::BF16 => conv::conv3d_bf16(x, weight, bias, &options), + dtype => panic!("conv3d: unsupported dtype {:?}", dtype), + } + } + + fn conv_transpose1d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<1>, + ) -> FloatTensor { + match x.dtype() { + DType::F32 => conv_transpose::conv_transpose1d_f32(x, weight, bias, &options), + DType::F64 => conv_transpose::conv_transpose1d_f64(x, weight, bias, &options), + DType::F16 => conv_transpose::conv_transpose1d_f16(x, weight, bias, &options), + DType::BF16 => conv_transpose::conv_transpose1d_bf16(x, weight, bias, &options), + dtype => panic!("conv_transpose1d: unsupported dtype {:?}", dtype), + } + } + + fn conv_transpose2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<2>, + ) -> FloatTensor { + match x.dtype() { + DType::F32 => conv_transpose::conv_transpose2d_f32(x, weight, bias, &options), + DType::F64 => conv_transpose::conv_transpose2d_f64(x, weight, bias, &options), + DType::F16 => conv_transpose::conv_transpose2d_f16(x, weight, bias, &options), + DType::BF16 => conv_transpose::conv_transpose2d_bf16(x, weight, bias, &options), + dtype => panic!("conv_transpose2d: unsupported dtype {:?}", dtype), + } + } + + fn conv_transpose3d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<3>, + ) -> FloatTensor { + match x.dtype() { + DType::F32 => conv_transpose::conv_transpose3d_f32(x, weight, bias, &options), + DType::F64 => conv_transpose::conv_transpose3d_f64(x, weight, bias, &options), + DType::F16 => conv_transpose::conv_transpose3d_f16(x, weight, bias, &options), + DType::BF16 => conv_transpose::conv_transpose3d_bf16(x, weight, bias, &options), + dtype => panic!("conv_transpose3d: unsupported dtype {:?}", dtype), + } + } + + fn avg_pool2d( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + match x.dtype() { + DType::F32 => pool::avg_pool2d_f32( + x, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ), + DType::F64 => pool::avg_pool2d_f64( + x, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ), + DType::F16 => pool::avg_pool2d_f16( + x, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ), + DType::BF16 => pool::avg_pool2d_bf16( + x, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ), + dtype => panic!("avg_pool2d: unsupported dtype {:?}", dtype), + } + } + + fn avg_pool2d_backward( + x: FloatTensor, + grad: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + _divisor_override: bool, + ) -> FloatTensor { + match x.dtype() { + DType::F32 => pool::avg_pool2d_backward_f32( + x, + grad, + kernel_size, + stride, + padding, + count_include_pad, + ), + DType::F64 => pool::avg_pool2d_backward_f64( + x, + grad, + kernel_size, + stride, + padding, + count_include_pad, + ), + DType::F16 => pool::avg_pool2d_backward_f16( + x, + grad, + kernel_size, + stride, + padding, + count_include_pad, + ), + DType::BF16 => pool::avg_pool2d_backward_bf16( + x, + grad, + kernel_size, + stride, + padding, + count_include_pad, + ), + dtype => panic!("avg_pool2d_backward: unsupported dtype {:?}", dtype), + } + } + + fn adaptive_avg_pool2d(x: FloatTensor, output_size: [usize; 2]) -> FloatTensor { + match x.dtype() { + DType::F32 => pool::adaptive_avg_pool2d_f32(x, output_size), + DType::F64 => pool::adaptive_avg_pool2d_f64(x, output_size), + DType::F16 => pool::adaptive_avg_pool2d_f16(x, output_size), + DType::BF16 => pool::adaptive_avg_pool2d_bf16(x, output_size), + dtype => panic!("adaptive_avg_pool2d: unsupported dtype {:?}", dtype), + } + } + + fn adaptive_avg_pool2d_backward( + x: FloatTensor, + grad: FloatTensor, + ) -> FloatTensor { + match x.dtype() { + DType::F32 => pool::adaptive_avg_pool2d_backward_f32(x, grad), + DType::F64 => pool::adaptive_avg_pool2d_backward_f64(x, grad), + DType::F16 => pool::adaptive_avg_pool2d_backward_f16(x, grad), + DType::BF16 => pool::adaptive_avg_pool2d_backward_bf16(x, grad), + dtype => panic!( + "adaptive_avg_pool2d_backward: unsupported dtype {:?}", + dtype + ), + } + } + + fn max_pool2d( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + ) -> FloatTensor { + match x.dtype() { + DType::F32 => { + pool::max_pool2d_f32(x, kernel_size, stride, padding, dilation, ceil_mode) + } + DType::F64 => { + pool::max_pool2d_f64(x, kernel_size, stride, padding, dilation, ceil_mode) + } + DType::F16 => { + pool::max_pool2d_f16(x, kernel_size, stride, padding, dilation, ceil_mode) + } + DType::BF16 => { + pool::max_pool2d_bf16(x, kernel_size, stride, padding, dilation, ceil_mode) + } + dtype => panic!("max_pool2d: unsupported dtype {:?}", dtype), + } + } + + fn max_pool2d_with_indices( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + indices_dtype: IntDType, + ) -> MaxPool2dWithIndices { + let (output, mut indices) = match x.dtype() { + DType::F32 => pool::max_pool2d_with_indices_f32( + x, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + ), + DType::F64 => pool::max_pool2d_with_indices_f64( + x, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + ), + DType::F16 => pool::max_pool2d_with_indices_f16( + x, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + ), + DType::BF16 => pool::max_pool2d_with_indices_bf16( + x, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + ), + dtype => panic!("max_pool2d_with_indices: unsupported dtype {:?}", dtype), + }; + if indices.dtype() != DType::from(indices_dtype) { + indices = Flex::int_cast(indices, indices_dtype); + } + MaxPool2dWithIndices::new(output, indices) + } + + fn max_pool2d_with_indices_backward( + x: FloatTensor, + _kernel_size: [usize; 2], + _stride: [usize; 2], + _padding: [usize; 2], + _dilation: [usize; 2], + _ceil_mode: bool, + output_grad: FloatTensor, + indices: IntTensor, + ) -> MaxPool2dBackward { + let x_grad = match x.dtype() { + DType::F32 => pool::max_pool2d_backward_f32(x, output_grad, indices), + DType::F64 => pool::max_pool2d_backward_f64(x, output_grad, indices), + DType::F16 => pool::max_pool2d_backward_f16(x, output_grad, indices), + DType::BF16 => pool::max_pool2d_backward_bf16(x, output_grad, indices), + dtype => panic!( + "max_pool2d_with_indices_backward: unsupported dtype {:?}", + dtype + ), + }; + MaxPool2dBackward::new(x_grad) + } + + fn interpolate( + x: FloatTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> FloatTensor { + match (options.mode, x.dtype()) { + (InterpolateMode::Nearest, DType::F32) => { + interpolate::interpolate_nearest_f32(x, output_size, options.align_corners) + } + (InterpolateMode::Nearest, DType::F64) => { + interpolate::interpolate_nearest_f64(x, output_size, options.align_corners) + } + (InterpolateMode::Nearest, DType::F16) => { + interpolate::interpolate_nearest_f16(x, output_size, options.align_corners) + } + (InterpolateMode::Nearest, DType::BF16) => { + interpolate::interpolate_nearest_bf16(x, output_size, options.align_corners) + } + (InterpolateMode::Bilinear, DType::F32) => { + interpolate::interpolate_bilinear_f32(x, output_size, options.align_corners) + } + (InterpolateMode::Bilinear, DType::F64) => { + interpolate::interpolate_bilinear_f64(x, output_size, options.align_corners) + } + (InterpolateMode::Bilinear, DType::F16) => { + interpolate::interpolate_bilinear_f16(x, output_size, options.align_corners) + } + (InterpolateMode::Bilinear, DType::BF16) => { + interpolate::interpolate_bilinear_bf16(x, output_size, options.align_corners) + } + (InterpolateMode::Bicubic, DType::F32) => { + interpolate::interpolate_bicubic_f32(x, output_size, options.align_corners) + } + (InterpolateMode::Bicubic, DType::F64) => { + interpolate::interpolate_bicubic_f64(x, output_size, options.align_corners) + } + (InterpolateMode::Bicubic, DType::F16) => { + interpolate::interpolate_bicubic_f16(x, output_size, options.align_corners) + } + (InterpolateMode::Bicubic, DType::BF16) => { + interpolate::interpolate_bicubic_bf16(x, output_size, options.align_corners) + } + (InterpolateMode::Lanczos3, DType::F32) => { + interpolate::interpolate_lanczos3_f32(x, output_size, options.align_corners) + } + (InterpolateMode::Lanczos3, DType::F64) => { + interpolate::interpolate_lanczos3_f64(x, output_size, options.align_corners) + } + (InterpolateMode::Lanczos3, DType::F16) => { + interpolate::interpolate_lanczos3_f16(x, output_size, options.align_corners) + } + (InterpolateMode::Lanczos3, DType::BF16) => { + interpolate::interpolate_lanczos3_bf16(x, output_size, options.align_corners) + } + (mode, dtype) => panic!( + "interpolate: unsupported mode {:?} / dtype {:?}", + mode, dtype + ), + } + } + + fn interpolate_backward( + x: FloatTensor, + grad: FloatTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> FloatTensor { + match (options.mode, x.dtype()) { + (InterpolateMode::Nearest, DType::F32) => { + interpolate::interpolate_nearest_backward_f32( + x, + grad, + output_size, + options.align_corners, + ) + } + (InterpolateMode::Nearest, DType::F64) => { + interpolate::interpolate_nearest_backward_f64( + x, + grad, + output_size, + options.align_corners, + ) + } + (InterpolateMode::Nearest, DType::F16) => { + interpolate::interpolate_nearest_backward_f16( + x, + grad, + output_size, + options.align_corners, + ) + } + (InterpolateMode::Nearest, DType::BF16) => { + interpolate::interpolate_nearest_backward_bf16( + x, + grad, + output_size, + options.align_corners, + ) + } + (InterpolateMode::Bilinear, DType::F32) => { + interpolate::interpolate_bilinear_backward_f32( + x, + grad, + output_size, + options.align_corners, + ) + } + (InterpolateMode::Bilinear, DType::F64) => { + interpolate::interpolate_bilinear_backward_f64( + x, + grad, + output_size, + options.align_corners, + ) + } + (InterpolateMode::Bilinear, DType::F16) => { + interpolate::interpolate_bilinear_backward_f16( + x, + grad, + output_size, + options.align_corners, + ) + } + (InterpolateMode::Bilinear, DType::BF16) => { + interpolate::interpolate_bilinear_backward_bf16( + x, + grad, + output_size, + options.align_corners, + ) + } + (InterpolateMode::Bicubic, DType::F32) => { + interpolate::interpolate_bicubic_backward_f32( + x, + grad, + output_size, + options.align_corners, + ) + } + (InterpolateMode::Bicubic, DType::F64) => { + interpolate::interpolate_bicubic_backward_f64( + x, + grad, + output_size, + options.align_corners, + ) + } + (InterpolateMode::Bicubic, DType::F16) => { + interpolate::interpolate_bicubic_backward_f16( + x, + grad, + output_size, + options.align_corners, + ) + } + (InterpolateMode::Bicubic, DType::BF16) => { + interpolate::interpolate_bicubic_backward_bf16( + x, + grad, + output_size, + options.align_corners, + ) + } + (mode, dtype) => { + panic!( + "interpolate_backward: unsupported mode {:?} / dtype {:?}", + mode, dtype + ) + } + } + } + + fn attention( + query: FloatTensor, + key: FloatTensor, + value: FloatTensor, + mask: Option>, + attn_bias: Option>, + options: AttentionModuleOptions, + ) -> FloatTensor { + crate::ops::attention::attention(query, key, value, mask, attn_bias, options) + } + + fn rfft( + signal: FloatTensor, + dim: usize, + n: Option, + ) -> (FloatTensor, FloatTensor) { + match signal.dtype() { + DType::F32 => crate::ops::fft::rfft_f32(signal, dim, n), + DType::F64 => crate::ops::fft::rfft_f64(signal, dim, n), + DType::F16 => crate::ops::fft::rfft_f16(signal, dim, n), + DType::BF16 => crate::ops::fft::rfft_bf16(signal, dim, n), + dtype => panic!("rfft: unsupported dtype {:?}", dtype), + } + } + + fn irfft( + spectrum_re: FloatTensor, + spectrum_im: FloatTensor, + dim: usize, + n: Option, + ) -> FloatTensor { + match spectrum_re.dtype() { + DType::F32 => crate::ops::fft::irfft_f32(spectrum_re, spectrum_im, dim, n), + DType::F64 => crate::ops::fft::irfft_f64(spectrum_re, spectrum_im, dim, n), + DType::F16 => crate::ops::fft::irfft_f16(spectrum_re, spectrum_im, dim, n), + DType::BF16 => crate::ops::fft::irfft_bf16(spectrum_re, spectrum_im, dim, n), + dtype => panic!("irfft: unsupported dtype {:?}", dtype), + } + } + + fn embedding(weights: FloatTensor, indices: IntTensor) -> FloatTensor { + let [batch_size, seq_length] = indices.shape().dims(); + let [_, d_model] = weights.shape().dims(); + + let indices = Flex::int_reshape(indices, Shape::from(alloc::vec![batch_size * seq_length])); + let output = Flex::float_select(weights, 0, indices); + Flex::float_reshape( + output, + Shape::from(alloc::vec![batch_size, seq_length, d_model]), + ) + } + + fn layer_norm( + tensor: FloatTensor, + gamma: FloatTensor, + beta: Option>, + epsilon: f64, + ) -> FloatTensor { + crate::ops::activation::layer_norm(tensor, gamma, beta, epsilon) + } + + fn embedding_backward( + weights: FloatTensor, + output_grad: FloatTensor, + indices: IntTensor, + ) -> FloatTensor { + let [batch_size, seq_length] = indices.shape().dims(); + let [n_embeddings, d_model] = weights.shape().dims(); + let dtype = output_grad.dtype(); + + let indices = Flex::int_reshape(indices, Shape::from(alloc::vec![batch_size * seq_length])); + let output_grad = Flex::float_reshape( + output_grad, + Shape::from(alloc::vec![batch_size * seq_length, d_model]), + ); + let grad = Flex::float_zeros( + Shape::from(alloc::vec![n_embeddings, d_model]), + &Default::default(), + dtype.into(), + ); + Flex::float_select_add(grad, 0, indices, output_grad) + } +} diff --git a/crates/burn-flex/src/ops/pool.rs b/crates/burn-flex/src/ops/pool.rs new file mode 100644 index 0000000..f0c0538 --- /dev/null +++ b/crates/burn-flex/src/ops/pool.rs @@ -0,0 +1,1891 @@ +//! Pooling operations using a unified 3D implementation. +//! +//! All pooling (1D, 2D, 3D) uses a unified 3D implementation: +//! - pool1d: adds two size-1 dimensions, calls pool3d, squeezes output +//! - pool2d: adds one size-1 dimension, calls pool3d, squeezes output +//! - pool3d: native implementation +//! +//! Supported dtypes: f32, f64, f16 (native), bf16 (via f32 conversion) + +use alloc::vec; +use alloc::vec::Vec; +use burn_backend::{DType, Element}; +use burn_std::{Bytes, Shape, bf16, f16}; + +use crate::{FlexTensor, Layout}; + +// ============================================================================ +// Macros for dtype wrappers +// ============================================================================ + +/// Generates max_pool3d_with_indices typed dispatchers. +macro_rules! max_pool3d_with_indices_typed { + ($fn_name:ident, $T:ty, $dtype:expr, $neg_inf:expr) => { + pub fn $fn_name( + x: FlexTensor, + kernel_size: [usize; 3], + stride: [usize; 3], + padding: [usize; 3], + dilation: [usize; 3], + ceil_mode: bool, + ) -> (FlexTensor, FlexTensor) { + max_pool3d_with_indices_impl::<$T>( + x, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + $dtype, + $neg_inf, + ) + } + }; +} + +/// Generates avg_pool3d typed dispatchers. +macro_rules! avg_pool3d_typed { + ($fn_name:ident, $T:ty, $dtype:expr, $zero:expr, $add_fn:expr, $div_fn:expr) => { + pub fn $fn_name( + x: FlexTensor, + kernel_size: [usize; 3], + stride: [usize; 3], + padding: [usize; 3], + count_include_pad: bool, + ceil_mode: bool, + ) -> FlexTensor { + avg_pool3d_impl::<$T>( + x, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + $dtype, + $zero, + $add_fn, + $div_fn, + ) + } + }; +} + +/// Generates adaptive_avg_pool3d typed dispatchers. +macro_rules! adaptive_avg_pool3d_typed { + ($fn_name:ident, $T:ty, $dtype:expr, $zero:expr, $add_fn:expr, $div_fn:expr) => { + pub fn $fn_name(x: FlexTensor, output_size: [usize; 3]) -> FlexTensor { + adaptive_avg_pool3d_impl::<$T>(x, output_size, $dtype, $zero, $add_fn, $div_fn) + } + }; +} + +/// Generates max_pool3d_backward typed dispatchers. +macro_rules! max_pool3d_backward_typed { + ($fn_name:ident, $T:ty, $dtype:expr, $zero:expr, $add_fn:expr) => { + pub fn $fn_name(x: FlexTensor, grad: FlexTensor, indices: FlexTensor) -> FlexTensor { + match indices.dtype() { + DType::I64 => { + max_pool3d_backward_impl::<$T, i64>(x, grad, indices, $dtype, $zero, $add_fn) + } + DType::I32 => { + max_pool3d_backward_impl::<$T, i32>(x, grad, indices, $dtype, $zero, $add_fn) + } + DType::I16 => { + max_pool3d_backward_impl::<$T, i16>(x, grad, indices, $dtype, $zero, $add_fn) + } + DType::I8 => { + max_pool3d_backward_impl::<$T, i8>(x, grad, indices, $dtype, $zero, $add_fn) + } + other => panic!("max_pool3d_backward: unsupported index dtype {other:?}",), + } + } + }; +} + +/// Generates avg_pool3d_backward typed dispatchers. +macro_rules! avg_pool3d_backward_typed { + ($fn_name:ident, $T:ty, $dtype:expr, $zero:expr, $add_fn:expr, $div_fn:expr) => { + pub fn $fn_name( + x: FlexTensor, + grad: FlexTensor, + kernel_size: [usize; 3], + stride: [usize; 3], + padding: [usize; 3], + count_include_pad: bool, + ) -> FlexTensor { + avg_pool3d_backward_impl::<$T>( + x, + grad, + kernel_size, + stride, + padding, + count_include_pad, + $dtype, + $zero, + $add_fn, + $div_fn, + ) + } + }; +} + +/// Generates adaptive_avg_pool3d_backward typed dispatchers. +macro_rules! adaptive_avg_pool3d_backward_typed { + ($fn_name:ident, $T:ty, $dtype:expr, $zero:expr, $add_fn:expr, $div_fn:expr) => { + pub fn $fn_name(x: FlexTensor, grad: FlexTensor) -> FlexTensor { + adaptive_avg_pool3d_backward_impl::<$T>(x, grad, $dtype, $zero, $add_fn, $div_fn) + } + }; +} + +// ============================================================================ +// Output size calculation +// ============================================================================ + +/// Calculate pooling output size for a single dimension. +fn pool_output_size( + input: usize, + kernel: usize, + padding: usize, + stride: usize, + dilation: usize, + ceil_mode: bool, +) -> usize { + assert!(kernel > 0, "pool: kernel size must be > 0"); + assert!(stride > 0, "pool: stride must be > 0"); + let effective_kernel = dilation * (kernel - 1) + 1; + let padded = input + 2 * padding; + if padded < effective_kernel { + return if ceil_mode { 1 } else { 0 }; + } + let numerator = padded - effective_kernel; + if ceil_mode { + numerator.div_ceil(stride) + 1 + } else { + numerator / stride + 1 + } +} + +// ============================================================================ +// Max Pool 3D - core implementation +// ============================================================================ + +max_pool3d_with_indices_typed!( + max_pool3d_with_indices_f32, + f32, + DType::F32, + f32::NEG_INFINITY +); +max_pool3d_with_indices_typed!( + max_pool3d_with_indices_f64, + f64, + DType::F64, + f64::NEG_INFINITY +); +max_pool3d_with_indices_typed!( + max_pool3d_with_indices_f16, + f16, + DType::F16, + f16::NEG_INFINITY +); + +pub fn max_pool3d_with_indices_bf16( + x: FlexTensor, + kernel_size: [usize; 3], + stride: [usize; 3], + padding: [usize; 3], + dilation: [usize; 3], + ceil_mode: bool, +) -> (FlexTensor, FlexTensor) { + let x_f32 = convert_bf16_to_f32(&x); + let (output_f32, indices) = + max_pool3d_with_indices_f32(x_f32, kernel_size, stride, padding, dilation, ceil_mode); + (convert_f32_to_bf16(&output_f32), indices) +} + +/// Generic 3D max pooling with indices implementation. +#[allow(clippy::too_many_arguments)] +fn max_pool3d_with_indices_impl( + x: FlexTensor, + kernel_size: [usize; 3], + stride: [usize; 3], + padding: [usize; 3], + dilation: [usize; 3], + ceil_mode: bool, + dtype: DType, + neg_inf: T, +) -> (FlexTensor, FlexTensor) +where + T: bytemuck::Pod + Copy + PartialOrd + Send + Sync + Element, +{ + let x = x.to_contiguous(); + let x_shape = x.layout().shape(); + + let batch_size = x_shape[0]; + let channels = x_shape[1]; + let in_d = x_shape[2]; + let in_h = x_shape[3]; + let in_w = x_shape[4]; + + let [kernel_d, kernel_h, kernel_w] = kernel_size; + let [stride_d, stride_h, stride_w] = stride; + let [pad_d, pad_h, pad_w] = padding; + let [dilation_d, dilation_h, dilation_w] = dilation; + + let out_d = pool_output_size(in_d, kernel_d, pad_d, stride_d, dilation_d, ceil_mode); + let out_h = pool_output_size(in_h, kernel_h, pad_h, stride_h, dilation_h, ceil_mode); + let out_w = pool_output_size(in_w, kernel_w, pad_w, stride_w, dilation_w, ceil_mode); + + let spatial_out = out_d * out_h * out_w; + let x_data: &[T] = x.storage(); + + let (output, indices) = { + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + + let mut output = vec![neg_inf; batch_size * channels * spatial_out]; + let mut indices = vec![-1i64; batch_size * channels * spatial_out]; + let out_ptr = crate::ops::SendMutPtr::new(output.as_mut_ptr()); + let idx_ptr = crate::ops::SendMutPtr::new(indices.as_mut_ptr()); + + // Flatten batch*channels into a single par_iter so rayon chooses + // the right granularity instead of creating one task per channel. + let bc_total = batch_size * channels; + (0..bc_total).into_par_iter().for_each(|bc| { + let b = bc / channels; + let c = bc % channels; + let x_offset = b * channels * in_d * in_h * in_w + c * in_d * in_h * in_w; + let out_offset = bc * spatial_out; + + for od in 0..out_d { + for oh in 0..out_h { + for ow in 0..out_w { + let out_idx = out_offset + od * out_h * out_w + oh * out_w + ow; + let mut max_val = neg_inf; + let mut max_idx: i64 = -1; + + for kd in 0..kernel_d { + let id = + (od * stride_d + kd * dilation_d) as isize - pad_d as isize; + if id < 0 || id >= in_d as isize { + continue; + } + let id = id as usize; + + for kh in 0..kernel_h { + let ih = + (oh * stride_h + kh * dilation_h) as isize - pad_h as isize; + if ih < 0 || ih >= in_h as isize { + continue; + } + let ih = ih as usize; + + for kw in 0..kernel_w { + let iw = (ow * stride_w + kw * dilation_w) as isize + - pad_w as isize; + if iw < 0 || iw >= in_w as isize { + continue; + } + let iw = iw as usize; + + let x_idx = x_offset + id * in_h * in_w + ih * in_w + iw; + let val = x_data[x_idx]; + + if max_idx < 0 || val > max_val { + max_val = val; + max_idx = (id * in_h * in_w + ih * in_w + iw) as i64; + } + } + } + } + + unsafe { + out_ptr.write(out_idx, max_val); + idx_ptr.write(out_idx, max_idx); + } + } + } + } + }); + (output, indices) + } + #[cfg(not(feature = "rayon"))] + { + let mut output = vec![neg_inf; batch_size * channels * spatial_out]; + let mut indices = vec![-1i64; batch_size * channels * spatial_out]; + + for b in 0..batch_size { + for c in 0..channels { + let x_offset = b * channels * in_d * in_h * in_w + c * in_d * in_h * in_w; + let out_offset = b * channels * spatial_out + c * spatial_out; + + for od in 0..out_d { + for oh in 0..out_h { + for ow in 0..out_w { + let out_idx = out_offset + od * out_h * out_w + oh * out_w + ow; + let mut max_val = neg_inf; + let mut max_idx: i64 = -1; + + for kd in 0..kernel_d { + let id = + (od * stride_d + kd * dilation_d) as isize - pad_d as isize; + if id < 0 || id >= in_d as isize { + continue; + } + let id = id as usize; + + for kh in 0..kernel_h { + let ih = (oh * stride_h + kh * dilation_h) as isize + - pad_h as isize; + if ih < 0 || ih >= in_h as isize { + continue; + } + let ih = ih as usize; + + for kw in 0..kernel_w { + let iw = (ow * stride_w + kw * dilation_w) as isize + - pad_w as isize; + if iw < 0 || iw >= in_w as isize { + continue; + } + let iw = iw as usize; + + let x_idx = + x_offset + id * in_h * in_w + ih * in_w + iw; + let val = x_data[x_idx]; + + if max_idx < 0 || val > max_val { + max_val = val; + max_idx = + (id * in_h * in_w + ih * in_w + iw) as i64; + } + } + } + } + + output[out_idx] = max_val; + indices[out_idx] = max_idx; + } + } + } + } + } + (output, indices) + } + }; + + let out_shape = Shape::from(vec![batch_size, channels, out_d, out_h, out_w]); + let output_tensor = FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape.clone()), + dtype, + ); + let indices_tensor = FlexTensor::new( + Bytes::from_elems(indices), + Layout::contiguous(out_shape), + DType::I64, + ); + + (output_tensor, indices_tensor) +} + +/// 3D max pooling (without returning indices) for f32. +pub fn max_pool3d_f32( + x: FlexTensor, + kernel_size: [usize; 3], + stride: [usize; 3], + padding: [usize; 3], + dilation: [usize; 3], + ceil_mode: bool, +) -> FlexTensor { + max_pool3d_with_indices_f32(x, kernel_size, stride, padding, dilation, ceil_mode).0 +} + +/// 3D max pooling (without returning indices) for f64. +pub fn max_pool3d_f64( + x: FlexTensor, + kernel_size: [usize; 3], + stride: [usize; 3], + padding: [usize; 3], + dilation: [usize; 3], + ceil_mode: bool, +) -> FlexTensor { + max_pool3d_with_indices_f64(x, kernel_size, stride, padding, dilation, ceil_mode).0 +} + +/// 3D max pooling (without returning indices) for f16. +pub fn max_pool3d_f16( + x: FlexTensor, + kernel_size: [usize; 3], + stride: [usize; 3], + padding: [usize; 3], + dilation: [usize; 3], + ceil_mode: bool, +) -> FlexTensor { + max_pool3d_with_indices_f16(x, kernel_size, stride, padding, dilation, ceil_mode).0 +} + +/// 3D max pooling (without returning indices) for bf16. +pub fn max_pool3d_bf16( + x: FlexTensor, + kernel_size: [usize; 3], + stride: [usize; 3], + padding: [usize; 3], + dilation: [usize; 3], + ceil_mode: bool, +) -> FlexTensor { + max_pool3d_with_indices_bf16(x, kernel_size, stride, padding, dilation, ceil_mode).0 +} + +// ============================================================================ +// Max Pool 2D - delegates to 3D +// ============================================================================ + +/// 2D max pooling with indices for f32. +pub fn max_pool2d_with_indices_f32( + x: FlexTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, +) -> (FlexTensor, FlexTensor) { + let x_3d = expand_2d_to_3d(&x); + let (output, indices) = max_pool3d_with_indices_f32( + x_3d, + [1, kernel_size[0], kernel_size[1]], + [1, stride[0], stride[1]], + [0, padding[0], padding[1]], + [1, dilation[0], dilation[1]], + ceil_mode, + ); + (squeeze_3d_to_2d(output), squeeze_3d_to_2d(indices)) +} + +/// 2D max pooling with indices for f64. +pub fn max_pool2d_with_indices_f64( + x: FlexTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, +) -> (FlexTensor, FlexTensor) { + let x_3d = expand_2d_to_3d(&x); + let (output, indices) = max_pool3d_with_indices_f64( + x_3d, + [1, kernel_size[0], kernel_size[1]], + [1, stride[0], stride[1]], + [0, padding[0], padding[1]], + [1, dilation[0], dilation[1]], + ceil_mode, + ); + (squeeze_3d_to_2d(output), squeeze_3d_to_2d(indices)) +} + +/// 2D max pooling with indices for f16. +pub fn max_pool2d_with_indices_f16( + x: FlexTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, +) -> (FlexTensor, FlexTensor) { + let x_3d = expand_2d_to_3d(&x); + let (output, indices) = max_pool3d_with_indices_f16( + x_3d, + [1, kernel_size[0], kernel_size[1]], + [1, stride[0], stride[1]], + [0, padding[0], padding[1]], + [1, dilation[0], dilation[1]], + ceil_mode, + ); + (squeeze_3d_to_2d(output), squeeze_3d_to_2d(indices)) +} + +/// 2D max pooling with indices for bf16. +pub fn max_pool2d_with_indices_bf16( + x: FlexTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, +) -> (FlexTensor, FlexTensor) { + let x_3d = expand_2d_to_3d(&x); + let (output, indices) = max_pool3d_with_indices_bf16( + x_3d, + [1, kernel_size[0], kernel_size[1]], + [1, stride[0], stride[1]], + [0, padding[0], padding[1]], + [1, dilation[0], dilation[1]], + ceil_mode, + ); + (squeeze_3d_to_2d(output), squeeze_3d_to_2d(indices)) +} + +/// 2D max pooling (without indices) for f32. +pub fn max_pool2d_f32( + x: FlexTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, +) -> FlexTensor { + max_pool2d_with_indices_f32(x, kernel_size, stride, padding, dilation, ceil_mode).0 +} + +/// 2D max pooling (without indices) for f64. +pub fn max_pool2d_f64( + x: FlexTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, +) -> FlexTensor { + max_pool2d_with_indices_f64(x, kernel_size, stride, padding, dilation, ceil_mode).0 +} + +/// 2D max pooling (without indices) for f16. +pub fn max_pool2d_f16( + x: FlexTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, +) -> FlexTensor { + max_pool2d_with_indices_f16(x, kernel_size, stride, padding, dilation, ceil_mode).0 +} + +/// 2D max pooling (without indices) for bf16. +pub fn max_pool2d_bf16( + x: FlexTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, +) -> FlexTensor { + max_pool2d_with_indices_bf16(x, kernel_size, stride, padding, dilation, ceil_mode).0 +} + +// ============================================================================ +// Max Pool 1D - delegates to 3D +// ============================================================================ + +/// 1D max pooling with indices for f32. +pub fn max_pool1d_with_indices_f32( + x: FlexTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, +) -> (FlexTensor, FlexTensor) { + let x_3d = expand_1d_to_3d(&x); + let (output, indices) = max_pool3d_with_indices_f32( + x_3d, + [1, 1, kernel_size], + [1, 1, stride], + [0, 0, padding], + [1, 1, dilation], + ceil_mode, + ); + (squeeze_3d_to_1d(output), squeeze_3d_to_1d(indices)) +} + +/// 1D max pooling (without indices) for f32. +pub fn max_pool1d_f32( + x: FlexTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, +) -> FlexTensor { + max_pool1d_with_indices_f32(x, kernel_size, stride, padding, dilation, ceil_mode).0 +} + +// ============================================================================ +// Avg Pool 3D - core implementation +// ============================================================================ + +avg_pool3d_typed!( + avg_pool3d_f32, + f32, + DType::F32, + 0.0f32, + |a, b| a + b, + |sum, count| sum / count as f32 +); +avg_pool3d_typed!( + avg_pool3d_f64, + f64, + DType::F64, + 0.0f64, + |a, b| a + b, + |sum, count| sum / count as f64 +); +avg_pool3d_typed!( + avg_pool3d_f16, + f16, + DType::F16, + f16::from_f32(0.0), + |a: f16, b: f16| f16::from_f32(a.to_f32() + b.to_f32()), + |sum: f16, count| f16::from_f32(sum.to_f32() / count as f32) +); + +pub fn avg_pool3d_bf16( + x: FlexTensor, + kernel_size: [usize; 3], + stride: [usize; 3], + padding: [usize; 3], + count_include_pad: bool, + ceil_mode: bool, +) -> FlexTensor { + let x_f32 = convert_bf16_to_f32(&x); + let result_f32 = avg_pool3d_f32( + x_f32, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ); + convert_f32_to_bf16(&result_f32) +} + +/// Generic 3D average pooling implementation. +#[allow(clippy::too_many_arguments)] +fn avg_pool3d_impl( + x: FlexTensor, + kernel_size: [usize; 3], + stride: [usize; 3], + padding: [usize; 3], + count_include_pad: bool, + ceil_mode: bool, + dtype: DType, + zero: T, + add_fn: fn(T, T) -> T, + div_fn: fn(T, usize) -> T, +) -> FlexTensor +where + T: bytemuck::Pod + Copy + Send + Sync + Element, +{ + let x = x.to_contiguous(); + let x_shape = x.layout().shape(); + + let batch_size = x_shape[0]; + let channels = x_shape[1]; + let in_d = x_shape[2]; + let in_h = x_shape[3]; + let in_w = x_shape[4]; + + let [kernel_d, kernel_h, kernel_w] = kernel_size; + let [stride_d, stride_h, stride_w] = stride; + let [pad_d, pad_h, pad_w] = padding; + + // Avg pool doesn't use dilation in typical implementations + let dilation_d = 1; + let dilation_h = 1; + let dilation_w = 1; + + let out_d = pool_output_size(in_d, kernel_d, pad_d, stride_d, dilation_d, ceil_mode); + let out_h = pool_output_size(in_h, kernel_h, pad_h, stride_h, dilation_h, ceil_mode); + let out_w = pool_output_size(in_w, kernel_w, pad_w, stride_w, dilation_w, ceil_mode); + + let spatial_out = out_d * out_h * out_w; + let x_data: &[T] = x.storage(); + let _kernel_volume = kernel_d * kernel_h * kernel_w; + + let output = { + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + + let mut output = vec![zero; batch_size * channels * spatial_out]; + let out_ptr = crate::ops::SendMutPtr::new(output.as_mut_ptr()); + + let bc_total = batch_size * channels; + (0..bc_total).into_par_iter().for_each(|bc| { + let b = bc / channels; + let c = bc % channels; + let x_offset = b * channels * in_d * in_h * in_w + c * in_d * in_h * in_w; + let out_offset = bc * spatial_out; + + for od in 0..out_d { + for oh in 0..out_h { + for ow in 0..out_w { + let out_idx = out_offset + od * out_h * out_w + oh * out_w + ow; + let mut sum = zero; + let mut count = 0usize; + let mut pad_count = 0usize; + + for kd in 0..kernel_d { + let id = (od * stride_d + kd) as isize - pad_d as isize; + let id_in_bounds = + id >= -(pad_d as isize) && id < (in_d + pad_d) as isize; + if !id_in_bounds { + continue; + } + let id_valid = id >= 0 && id < in_d as isize; + + for kh in 0..kernel_h { + let ih = (oh * stride_h + kh) as isize - pad_h as isize; + let ih_in_bounds = + ih >= -(pad_h as isize) && ih < (in_h + pad_h) as isize; + if !ih_in_bounds { + continue; + } + let ih_valid = ih >= 0 && ih < in_h as isize; + + for kw in 0..kernel_w { + let iw = (ow * stride_w + kw) as isize - pad_w as isize; + let iw_in_bounds = + iw >= -(pad_w as isize) && iw < (in_w + pad_w) as isize; + if !iw_in_bounds { + continue; + } + + pad_count += 1; + + let iw_valid = iw >= 0 && iw < in_w as isize; + if !id_valid || !ih_valid || !iw_valid { + continue; + } + + let id = id as usize; + let ih = ih as usize; + let iw = iw as usize; + let x_idx = x_offset + id * in_h * in_w + ih * in_w + iw; + sum = add_fn(sum, x_data[x_idx]); + count += 1; + } + } + } + + let divisor = if count_include_pad { + pad_count.max(1) + } else { + count.max(1) + }; + + unsafe { + out_ptr.write(out_idx, div_fn(sum, divisor)); + } + } + } + } + }); + output + } + #[cfg(not(feature = "rayon"))] + { + let mut output = vec![zero; batch_size * channels * spatial_out]; + + for b in 0..batch_size { + for c in 0..channels { + let x_offset = b * channels * in_d * in_h * in_w + c * in_d * in_h * in_w; + let out_offset = b * channels * spatial_out + c * spatial_out; + + for od in 0..out_d { + for oh in 0..out_h { + for ow in 0..out_w { + let out_idx = out_offset + od * out_h * out_w + oh * out_w + ow; + let mut sum = zero; + let mut count = 0usize; + + // Track count for count_include_pad (positions within padded bounds) + let mut pad_count = 0usize; + + for kd in 0..kernel_d { + let id = (od * stride_d + kd) as isize - pad_d as isize; + // Check if within padded bounds (not ceil_mode extension) + let id_in_bounds = + id >= -(pad_d as isize) && id < (in_d + pad_d) as isize; + if !id_in_bounds { + continue; // ceil_mode extension - skip entirely + } + let id_valid = id >= 0 && id < in_d as isize; + + for kh in 0..kernel_h { + let ih = (oh * stride_h + kh) as isize - pad_h as isize; + let ih_in_bounds = + ih >= -(pad_h as isize) && ih < (in_h + pad_h) as isize; + if !ih_in_bounds { + continue; + } + let ih_valid = ih >= 0 && ih < in_h as isize; + + for kw in 0..kernel_w { + let iw = (ow * stride_w + kw) as isize - pad_w as isize; + let iw_in_bounds = iw >= -(pad_w as isize) + && iw < (in_w + pad_w) as isize; + if !iw_in_bounds { + continue; + } + + // Position is within padded bounds + pad_count += 1; + + let iw_valid = iw >= 0 && iw < in_w as isize; + if !id_valid || !ih_valid || !iw_valid { + continue; // In padding zone - count but don't add + } + + let id = id as usize; + let ih = ih as usize; + let iw = iw as usize; + let x_idx = + x_offset + id * in_h * in_w + ih * in_w + iw; + sum = add_fn(sum, x_data[x_idx]); + count += 1; + } + } + } + + let divisor = if count_include_pad { + pad_count.max(1) // Positions within padded bounds + } else { + count.max(1) // Only actual valid positions + }; + + output[out_idx] = div_fn(sum, divisor); + } + } + } + } + } + output + } + }; + + let out_shape = Shape::from(vec![batch_size, channels, out_d, out_h, out_w]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + dtype, + ) +} + +// ============================================================================ +// Avg Pool 2D - delegates to 3D +// ============================================================================ + +/// 2D average pooling for f32. +pub fn avg_pool2d_f32( + x: FlexTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, +) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let result = avg_pool3d_f32( + x_3d, + [1, kernel_size[0], kernel_size[1]], + [1, stride[0], stride[1]], + [0, padding[0], padding[1]], + count_include_pad, + ceil_mode, + ); + squeeze_3d_to_2d(result) +} + +/// 2D average pooling for f64. +pub fn avg_pool2d_f64( + x: FlexTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, +) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let result = avg_pool3d_f64( + x_3d, + [1, kernel_size[0], kernel_size[1]], + [1, stride[0], stride[1]], + [0, padding[0], padding[1]], + count_include_pad, + ceil_mode, + ); + squeeze_3d_to_2d(result) +} + +/// 2D average pooling for f16. +pub fn avg_pool2d_f16( + x: FlexTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, +) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let result = avg_pool3d_f16( + x_3d, + [1, kernel_size[0], kernel_size[1]], + [1, stride[0], stride[1]], + [0, padding[0], padding[1]], + count_include_pad, + ceil_mode, + ); + squeeze_3d_to_2d(result) +} + +/// 2D average pooling for bf16. +pub fn avg_pool2d_bf16( + x: FlexTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, +) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let result = avg_pool3d_bf16( + x_3d, + [1, kernel_size[0], kernel_size[1]], + [1, stride[0], stride[1]], + [0, padding[0], padding[1]], + count_include_pad, + ceil_mode, + ); + squeeze_3d_to_2d(result) +} + +// ============================================================================ +// Avg Pool 1D - delegates to 3D +// ============================================================================ + +/// 1D average pooling for f32. +pub fn avg_pool1d_f32( + x: FlexTensor, + kernel_size: usize, + stride: usize, + padding: usize, + count_include_pad: bool, + ceil_mode: bool, +) -> FlexTensor { + let x_3d = expand_1d_to_3d(&x); + let result = avg_pool3d_f32( + x_3d, + [1, 1, kernel_size], + [1, 1, stride], + [0, 0, padding], + count_include_pad, + ceil_mode, + ); + squeeze_3d_to_1d(result) +} + +// ============================================================================ +// Adaptive Avg Pool 3D - core implementation +// ============================================================================ + +adaptive_avg_pool3d_typed!( + adaptive_avg_pool3d_f32, + f32, + DType::F32, + 0.0f32, + |a, b| a + b, + |sum, count| sum / count as f32 +); +adaptive_avg_pool3d_typed!( + adaptive_avg_pool3d_f64, + f64, + DType::F64, + 0.0f64, + |a, b| a + b, + |sum, count| sum / count as f64 +); +adaptive_avg_pool3d_typed!( + adaptive_avg_pool3d_f16, + f16, + DType::F16, + f16::from_f32(0.0), + |a: f16, b: f16| f16::from_f32(a.to_f32() + b.to_f32()), + |sum: f16, count| f16::from_f32(sum.to_f32() / count as f32) +); + +pub fn adaptive_avg_pool3d_bf16(x: FlexTensor, output_size: [usize; 3]) -> FlexTensor { + let x_f32 = convert_bf16_to_f32(&x); + let result_f32 = adaptive_avg_pool3d_f32(x_f32, output_size); + convert_f32_to_bf16(&result_f32) +} + +/// Generic 3D adaptive average pooling implementation. +fn adaptive_avg_pool3d_impl( + x: FlexTensor, + output_size: [usize; 3], + dtype: DType, + zero: T, + add_fn: fn(T, T) -> T, + div_fn: fn(T, usize) -> T, +) -> FlexTensor +where + T: bytemuck::Pod + Copy + Send + Sync + Element, +{ + let x = x.to_contiguous(); + let x_shape = x.layout().shape(); + + let batch_size = x_shape[0]; + let channels = x_shape[1]; + let in_d = x_shape[2]; + let in_h = x_shape[3]; + let in_w = x_shape[4]; + + let [out_d, out_h, out_w] = output_size; + let spatial_out = out_d * out_h * out_w; + let x_data: &[T] = x.storage(); + + let output = { + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + + let mut output = vec![zero; batch_size * channels * spatial_out]; + let out_ptr = crate::ops::SendMutPtr::new(output.as_mut_ptr()); + + (0..batch_size).into_par_iter().for_each(|b| { + (0..channels).into_par_iter().for_each(|c| { + let x_offset = b * channels * in_d * in_h * in_w + c * in_d * in_h * in_w; + let out_offset = b * channels * spatial_out + c * spatial_out; + + for od in 0..out_d { + // Compute input range for this output position + // start = floor(out * in / out_size), end = ceil((out+1) * in / out_size) + let d_start = (od * in_d) / out_d; + let d_end = ((od + 1) * in_d).div_ceil(out_d); + + for oh in 0..out_h { + let h_start = (oh * in_h) / out_h; + let h_end = ((oh + 1) * in_h).div_ceil(out_h); + + for ow in 0..out_w { + let w_start = (ow * in_w) / out_w; + let w_end = ((ow + 1) * in_w).div_ceil(out_w); + + let out_idx = out_offset + od * out_h * out_w + oh * out_w + ow; + let mut sum = zero; + let mut count = 0usize; + + for id in d_start..d_end { + for ih in h_start..h_end { + for iw in w_start..w_end { + let x_idx = + x_offset + id * in_h * in_w + ih * in_w + iw; + sum = add_fn(sum, x_data[x_idx]); + count += 1; + } + } + } + + unsafe { + out_ptr.write(out_idx, div_fn(sum, count.max(1))); + } + } + } + } + }); + }); + output + } + #[cfg(not(feature = "rayon"))] + { + let mut output = vec![zero; batch_size * channels * spatial_out]; + + for b in 0..batch_size { + for c in 0..channels { + let x_offset = b * channels * in_d * in_h * in_w + c * in_d * in_h * in_w; + let out_offset = b * channels * spatial_out + c * spatial_out; + + for od in 0..out_d { + let d_start = (od * in_d) / out_d; + let d_end = ((od + 1) * in_d).div_ceil(out_d); + + for oh in 0..out_h { + let h_start = (oh * in_h) / out_h; + let h_end = ((oh + 1) * in_h).div_ceil(out_h); + + for ow in 0..out_w { + let w_start = (ow * in_w) / out_w; + let w_end = ((ow + 1) * in_w).div_ceil(out_w); + + let out_idx = out_offset + od * out_h * out_w + oh * out_w + ow; + let mut sum = zero; + let mut count = 0usize; + + for id in d_start..d_end { + for ih in h_start..h_end { + for iw in w_start..w_end { + let x_idx = + x_offset + id * in_h * in_w + ih * in_w + iw; + sum = add_fn(sum, x_data[x_idx]); + count += 1; + } + } + } + + output[out_idx] = div_fn(sum, count.max(1)); + } + } + } + } + } + output + } + }; + + let out_shape = Shape::from(vec![batch_size, channels, out_d, out_h, out_w]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + dtype, + ) +} + +// ============================================================================ +// Adaptive Avg Pool 2D - delegates to 3D +// ============================================================================ + +/// 2D adaptive average pooling for f32. +pub fn adaptive_avg_pool2d_f32(x: FlexTensor, output_size: [usize; 2]) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let result = adaptive_avg_pool3d_f32(x_3d, [1, output_size[0], output_size[1]]); + squeeze_3d_to_2d(result) +} + +/// 2D adaptive average pooling for f64. +pub fn adaptive_avg_pool2d_f64(x: FlexTensor, output_size: [usize; 2]) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let result = adaptive_avg_pool3d_f64(x_3d, [1, output_size[0], output_size[1]]); + squeeze_3d_to_2d(result) +} + +/// 2D adaptive average pooling for f16. +pub fn adaptive_avg_pool2d_f16(x: FlexTensor, output_size: [usize; 2]) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let result = adaptive_avg_pool3d_f16(x_3d, [1, output_size[0], output_size[1]]); + squeeze_3d_to_2d(result) +} + +/// 2D adaptive average pooling for bf16. +pub fn adaptive_avg_pool2d_bf16(x: FlexTensor, output_size: [usize; 2]) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let result = adaptive_avg_pool3d_bf16(x_3d, [1, output_size[0], output_size[1]]); + squeeze_3d_to_2d(result) +} + +// ============================================================================ +// Adaptive Avg Pool 1D - delegates to 3D +// ============================================================================ + +/// 1D adaptive average pooling for f32. +pub fn adaptive_avg_pool1d_f32(x: FlexTensor, output_size: usize) -> FlexTensor { + let x_3d = expand_1d_to_3d(&x); + let result = adaptive_avg_pool3d_f32(x_3d, [1, 1, output_size]); + squeeze_3d_to_1d(result) +} + +// ============================================================================ +// Backward passes +// ============================================================================ + +/// Max pool 2D backward using stored indices. +pub fn max_pool2d_backward_f32(x: FlexTensor, grad: FlexTensor, indices: FlexTensor) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let grad_3d = expand_2d_to_3d(&grad); + let indices_3d = expand_2d_to_3d(&indices); + let result = max_pool3d_backward_f32(x_3d, grad_3d, indices_3d); + squeeze_3d_to_2d(result) +} + +/// Max pool 2D backward for f64. +pub fn max_pool2d_backward_f64(x: FlexTensor, grad: FlexTensor, indices: FlexTensor) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let grad_3d = expand_2d_to_3d(&grad); + let indices_3d = expand_2d_to_3d(&indices); + let result = max_pool3d_backward_f64(x_3d, grad_3d, indices_3d); + squeeze_3d_to_2d(result) +} + +/// Max pool 2D backward for f16. +pub fn max_pool2d_backward_f16(x: FlexTensor, grad: FlexTensor, indices: FlexTensor) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let grad_3d = expand_2d_to_3d(&grad); + let indices_3d = expand_2d_to_3d(&indices); + let result = max_pool3d_backward_f16(x_3d, grad_3d, indices_3d); + squeeze_3d_to_2d(result) +} + +/// Max pool 2D backward for bf16. +pub fn max_pool2d_backward_bf16( + x: FlexTensor, + grad: FlexTensor, + indices: FlexTensor, +) -> FlexTensor { + let x_f32 = convert_bf16_to_f32(&x); + let grad_f32 = convert_bf16_to_f32(&grad); + let result_f32 = max_pool2d_backward_f32(x_f32, grad_f32, indices); + convert_f32_to_bf16(&result_f32) +} + +max_pool3d_backward_typed!(max_pool3d_backward_f32, f32, DType::F32, 0.0f32, |a, b| a + + b); +max_pool3d_backward_typed!(max_pool3d_backward_f64, f64, DType::F64, 0.0f64, |a, b| a + + b); +max_pool3d_backward_typed!( + max_pool3d_backward_f16, + f16, + DType::F16, + f16::from_f32(0.0), + |a: f16, b: f16| f16::from_f32(a.to_f32() + b.to_f32()) +); + +/// Generic max pool 3D backward implementation. +fn max_pool3d_backward_impl( + x: FlexTensor, + grad: FlexTensor, + indices: FlexTensor, + dtype: DType, + zero: T, + add_fn: fn(T, T) -> T, +) -> FlexTensor +where + T: bytemuck::Pod + Copy + Send + Sync + Element, + I: bytemuck::Pod + Copy + Send + Sync + Element, +{ + let x_shape = x.layout().shape(); + let grad = grad.to_contiguous(); + let indices = indices.to_contiguous(); + + let batch_size = x_shape[0]; + let channels = x_shape[1]; + let in_d = x_shape[2]; + let in_h = x_shape[3]; + let in_w = x_shape[4]; + let spatial_in = in_d * in_h * in_w; + + let grad_shape = grad.layout().shape(); + let out_d = grad_shape[2]; + let out_h = grad_shape[3]; + let out_w = grad_shape[4]; + let spatial_out = out_d * out_h * out_w; + + let grad_data: &[T] = grad.storage(); + let indices_data: &[I] = indices.storage(); + + // Accumulate gradients back to input positions + let mut output = vec![zero; batch_size * channels * spatial_in]; + + for b in 0..batch_size { + for c in 0..channels { + let grad_offset = b * channels * spatial_out + c * spatial_out; + let out_offset = b * channels * spatial_in + c * spatial_in; + + for i in 0..spatial_out { + let idx = indices_data[grad_offset + i].elem::(); + if idx >= 0 { + let input_idx = out_offset + idx as usize; + output[input_idx] = add_fn(output[input_idx], grad_data[grad_offset + i]); + } + } + } + } + + let out_shape = Shape::from(vec![batch_size, channels, in_d, in_h, in_w]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + dtype, + ) +} + +/// Avg pool 2D backward. +pub fn avg_pool2d_backward_f32( + x: FlexTensor, + grad: FlexTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, +) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let grad_3d = expand_2d_to_3d(&grad); + let result = avg_pool3d_backward_f32( + x_3d, + grad_3d, + [1, kernel_size[0], kernel_size[1]], + [1, stride[0], stride[1]], + [0, padding[0], padding[1]], + count_include_pad, + ); + squeeze_3d_to_2d(result) +} + +/// Avg pool 2D backward for f64. +pub fn avg_pool2d_backward_f64( + x: FlexTensor, + grad: FlexTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, +) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let grad_3d = expand_2d_to_3d(&grad); + let result = avg_pool3d_backward_f64( + x_3d, + grad_3d, + [1, kernel_size[0], kernel_size[1]], + [1, stride[0], stride[1]], + [0, padding[0], padding[1]], + count_include_pad, + ); + squeeze_3d_to_2d(result) +} + +/// Avg pool 2D backward for f16. +pub fn avg_pool2d_backward_f16( + x: FlexTensor, + grad: FlexTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, +) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let grad_3d = expand_2d_to_3d(&grad); + let result = avg_pool3d_backward_f16( + x_3d, + grad_3d, + [1, kernel_size[0], kernel_size[1]], + [1, stride[0], stride[1]], + [0, padding[0], padding[1]], + count_include_pad, + ); + squeeze_3d_to_2d(result) +} + +/// Avg pool 2D backward for bf16. +pub fn avg_pool2d_backward_bf16( + x: FlexTensor, + grad: FlexTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, +) -> FlexTensor { + let x_f32 = convert_bf16_to_f32(&x); + let grad_f32 = convert_bf16_to_f32(&grad); + let result_f32 = avg_pool2d_backward_f32( + x_f32, + grad_f32, + kernel_size, + stride, + padding, + count_include_pad, + ); + convert_f32_to_bf16(&result_f32) +} + +avg_pool3d_backward_typed!( + avg_pool3d_backward_f32, + f32, + DType::F32, + 0.0f32, + |a, b| a + b, + |val, count| val / count as f32 +); +avg_pool3d_backward_typed!( + avg_pool3d_backward_f64, + f64, + DType::F64, + 0.0f64, + |a, b| a + b, + |val, count| val / count as f64 +); +avg_pool3d_backward_typed!( + avg_pool3d_backward_f16, + f16, + DType::F16, + f16::from_f32(0.0), + |a: f16, b: f16| f16::from_f32(a.to_f32() + b.to_f32()), + |val: f16, count| f16::from_f32(val.to_f32() / count as f32) +); + +/// Generic avg pool 3D backward implementation. +#[allow(clippy::too_many_arguments)] +fn avg_pool3d_backward_impl( + x: FlexTensor, + grad: FlexTensor, + kernel_size: [usize; 3], + stride: [usize; 3], + padding: [usize; 3], + count_include_pad: bool, + dtype: DType, + zero: T, + add_fn: fn(T, T) -> T, + div_fn: fn(T, usize) -> T, +) -> FlexTensor +where + T: bytemuck::Pod + Copy + Send + Sync + Element, +{ + let x_shape = x.layout().shape(); + let grad = grad.to_contiguous(); + + let batch_size = x_shape[0]; + let channels = x_shape[1]; + let in_d = x_shape[2]; + let in_h = x_shape[3]; + let in_w = x_shape[4]; + let spatial_in = in_d * in_h * in_w; + + let [kernel_d, kernel_h, kernel_w] = kernel_size; + let [stride_d, stride_h, stride_w] = stride; + let [pad_d, pad_h, pad_w] = padding; + let kernel_volume = kernel_d * kernel_h * kernel_w; + + let grad_shape = grad.layout().shape(); + let out_d = grad_shape[2]; + let out_h = grad_shape[3]; + let out_w = grad_shape[4]; + let spatial_out = out_d * out_h * out_w; + + let grad_data: &[T] = grad.storage(); + + // Distribute gradient equally across window + let mut output = vec![zero; batch_size * channels * spatial_in]; + + for b in 0..batch_size { + for c in 0..channels { + let grad_offset = b * channels * spatial_out + c * spatial_out; + let out_offset = b * channels * spatial_in + c * spatial_in; + + for od in 0..out_d { + for oh in 0..out_h { + for ow in 0..out_w { + let grad_idx = grad_offset + od * out_h * out_w + oh * out_w + ow; + let grad_val = grad_data[grad_idx]; + + // Count valid positions for this output + let mut count = 0usize; + for kd in 0..kernel_d { + let id = (od * stride_d + kd) as isize - pad_d as isize; + if id >= 0 && id < in_d as isize { + for kh in 0..kernel_h { + let ih = (oh * stride_h + kh) as isize - pad_h as isize; + if ih >= 0 && ih < in_h as isize { + for kw in 0..kernel_w { + let iw = (ow * stride_w + kw) as isize - pad_w as isize; + if iw >= 0 && iw < in_w as isize { + count += 1; + } + } + } + } + } + } + + let divisor = if count_include_pad { + kernel_volume + } else { + count.max(1) + }; + + // Distribute gradient + let distributed = div_fn(grad_val, divisor); + for kd in 0..kernel_d { + let id = (od * stride_d + kd) as isize - pad_d as isize; + if id < 0 || id >= in_d as isize { + continue; + } + let id = id as usize; + + for kh in 0..kernel_h { + let ih = (oh * stride_h + kh) as isize - pad_h as isize; + if ih < 0 || ih >= in_h as isize { + continue; + } + let ih = ih as usize; + + for kw in 0..kernel_w { + let iw = (ow * stride_w + kw) as isize - pad_w as isize; + if iw < 0 || iw >= in_w as isize { + continue; + } + let iw = iw as usize; + + let input_idx = out_offset + id * in_h * in_w + ih * in_w + iw; + output[input_idx] = add_fn(output[input_idx], distributed); + } + } + } + } + } + } + } + } + + let out_shape = Shape::from(vec![batch_size, channels, in_d, in_h, in_w]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + dtype, + ) +} + +/// Adaptive avg pool 2D backward. +pub fn adaptive_avg_pool2d_backward_f32(x: FlexTensor, grad: FlexTensor) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let grad_3d = expand_2d_to_3d(&grad); + let result = adaptive_avg_pool3d_backward_f32(x_3d, grad_3d); + squeeze_3d_to_2d(result) +} + +/// Adaptive avg pool 2D backward for f64. +pub fn adaptive_avg_pool2d_backward_f64(x: FlexTensor, grad: FlexTensor) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let grad_3d = expand_2d_to_3d(&grad); + let result = adaptive_avg_pool3d_backward_f64(x_3d, grad_3d); + squeeze_3d_to_2d(result) +} + +/// Adaptive avg pool 2D backward for f16. +pub fn adaptive_avg_pool2d_backward_f16(x: FlexTensor, grad: FlexTensor) -> FlexTensor { + let x_3d = expand_2d_to_3d(&x); + let grad_3d = expand_2d_to_3d(&grad); + let result = adaptive_avg_pool3d_backward_f16(x_3d, grad_3d); + squeeze_3d_to_2d(result) +} + +/// Adaptive avg pool 2D backward for bf16. +pub fn adaptive_avg_pool2d_backward_bf16(x: FlexTensor, grad: FlexTensor) -> FlexTensor { + let x_f32 = convert_bf16_to_f32(&x); + let grad_f32 = convert_bf16_to_f32(&grad); + let result_f32 = adaptive_avg_pool2d_backward_f32(x_f32, grad_f32); + convert_f32_to_bf16(&result_f32) +} + +adaptive_avg_pool3d_backward_typed!( + adaptive_avg_pool3d_backward_f32, + f32, + DType::F32, + 0.0f32, + |a, b| a + b, + |val, count| val / count as f32 +); +adaptive_avg_pool3d_backward_typed!( + adaptive_avg_pool3d_backward_f64, + f64, + DType::F64, + 0.0f64, + |a, b| a + b, + |val, count| val / count as f64 +); +adaptive_avg_pool3d_backward_typed!( + adaptive_avg_pool3d_backward_f16, + f16, + DType::F16, + f16::from_f32(0.0), + |a: f16, b: f16| f16::from_f32(a.to_f32() + b.to_f32()), + |val: f16, count| f16::from_f32(val.to_f32() / count as f32) +); + +/// Generic adaptive avg pool 3D backward implementation. +fn adaptive_avg_pool3d_backward_impl( + x: FlexTensor, + grad: FlexTensor, + dtype: DType, + zero: T, + add_fn: fn(T, T) -> T, + div_fn: fn(T, usize) -> T, +) -> FlexTensor +where + T: bytemuck::Pod + Copy + Send + Sync + Element, +{ + let x_shape = x.layout().shape(); + let grad = grad.to_contiguous(); + + let batch_size = x_shape[0]; + let channels = x_shape[1]; + let in_d = x_shape[2]; + let in_h = x_shape[3]; + let in_w = x_shape[4]; + let spatial_in = in_d * in_h * in_w; + + let grad_shape = grad.layout().shape(); + let out_d = grad_shape[2]; + let out_h = grad_shape[3]; + let out_w = grad_shape[4]; + let spatial_out = out_d * out_h * out_w; + + let grad_data: &[T] = grad.storage(); + + let mut output = vec![zero; batch_size * channels * spatial_in]; + + for b in 0..batch_size { + for c in 0..channels { + let grad_offset = b * channels * spatial_out + c * spatial_out; + let out_offset = b * channels * spatial_in + c * spatial_in; + + for od in 0..out_d { + let d_start = (od * in_d) / out_d; + let d_end = ((od + 1) * in_d).div_ceil(out_d); + + for oh in 0..out_h { + let h_start = (oh * in_h) / out_h; + let h_end = ((oh + 1) * in_h).div_ceil(out_h); + + for ow in 0..out_w { + let w_start = (ow * in_w) / out_w; + let w_end = ((ow + 1) * in_w).div_ceil(out_w); + + let grad_idx = grad_offset + od * out_h * out_w + oh * out_w + ow; + let grad_val = grad_data[grad_idx]; + + let count = (d_end - d_start) * (h_end - h_start) * (w_end - w_start); + let distributed = div_fn(grad_val, count.max(1)); + + for id in d_start..d_end { + for ih in h_start..h_end { + for iw in w_start..w_end { + let input_idx = out_offset + id * in_h * in_w + ih * in_w + iw; + output[input_idx] = add_fn(output[input_idx], distributed); + } + } + } + } + } + } + } + } + + let out_shape = Shape::from(vec![batch_size, channels, in_d, in_h, in_w]); + FlexTensor::new( + Bytes::from_elems(output), + Layout::contiguous(out_shape), + dtype, + ) +} + +// ============================================================================ +// Dimension expansion/squeeze helpers +// ============================================================================ + +/// Expand 2D tensor [N, C, H, W] to 3D [N, C, 1, H, W]. +fn expand_2d_to_3d(x: &FlexTensor) -> FlexTensor { + let shape = x.layout().shape(); + x.reshape(Shape::from(vec![shape[0], shape[1], 1, shape[2], shape[3]])) +} + +/// Squeeze 3D tensor [N, C, 1, H, W] to 2D [N, C, H, W]. +fn squeeze_3d_to_2d(x: FlexTensor) -> FlexTensor { + let shape = x.layout().shape(); + x.reshape(Shape::from(vec![shape[0], shape[1], shape[3], shape[4]])) +} + +/// Expand 1D tensor [N, C, L] to 3D [N, C, 1, 1, L]. +fn expand_1d_to_3d(x: &FlexTensor) -> FlexTensor { + let shape = x.layout().shape(); + x.reshape(Shape::from(vec![shape[0], shape[1], 1, 1, shape[2]])) +} + +/// Squeeze 3D tensor [N, C, 1, 1, L] to 1D [N, C, L]. +fn squeeze_3d_to_1d(x: FlexTensor) -> FlexTensor { + let shape = x.layout().shape(); + x.reshape(Shape::from(vec![shape[0], shape[1], shape[4]])) +} + +// ============================================================================ +// bf16 conversion helpers +// ============================================================================ + +fn convert_bf16_to_f32(tensor: &FlexTensor) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let data: &[bf16] = tensor.storage(); + let f32_data: Vec = data.iter().map(|x| x.to_f32()).collect(); + FlexTensor::new( + Bytes::from_elems(f32_data), + Layout::contiguous(tensor.layout().shape().clone()), + DType::F32, + ) +} + +fn convert_f32_to_bf16(tensor: &FlexTensor) -> FlexTensor { + let data: &[f32] = tensor.storage(); + let bf16_data: Vec = data.iter().map(|x| bf16::from_f32(*x)).collect(); + FlexTensor::new( + Bytes::from_elems(bf16_data), + Layout::contiguous(tensor.layout().shape().clone()), + DType::BF16, + ) +} + +// ============================================================================ +// Tests +// ============================================================================ + +// Tests kept here probe flex internals: the `pool_output_size` helper +// (including zero-kernel/stride panics), dtype storage paths for max_pool2d +// (f16/bf16/f64), the backward kernels (max/avg/adaptive), and flex's +// count_include_pad avg_pool2d semantics. Plain forward-pass pool tests +// (max/avg/adaptive 2d, pool1d/3d delegation) live in +// crates/burn-backend-tests/tests/tensor/float/module/{maxpool,avgpool, +// adaptive_avgpool}*.rs and run on every backend. +#[cfg(test)] +mod tests { + use super::*; + use burn_backend::TensorData; + + #[test] + fn test_pool_output_size() { + // Basic: input=4, kernel=2, padding=0, stride=2, dilation=1 + // output = (4 + 0 - 2) / 2 + 1 = 2 + assert_eq!(pool_output_size(4, 2, 0, 2, 1, false), 2); + + // With padding: input=4, kernel=2, padding=1, stride=2 + // output = (4 + 2 - 2) / 2 + 1 = 3 + assert_eq!(pool_output_size(4, 2, 1, 2, 1, false), 3); + + // With ceil mode: input=5, kernel=2, padding=0, stride=2 + // floor: (5 - 2) / 2 + 1 = 2 + // ceil: ceil(3/2) + 1 = 3 + assert_eq!(pool_output_size(5, 2, 0, 2, 1, false), 2); + assert_eq!(pool_output_size(5, 2, 0, 2, 1, true), 3); + + // With dilation: input=7, kernel=2, dilation=2 + // effective_kernel = 2*(2-1)+1 = 3 + // output = (7 - 3) / 1 + 1 = 5 + assert_eq!(pool_output_size(7, 2, 0, 1, 2, false), 5); + } + + #[test] + fn test_avg_pool2d_count_include_pad() { + // 3x3 input with padding 1, kernel 2x2, stride 2 + let x_data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]; + let x = FlexTensor::from_data(TensorData::new(x_data, vec![1, 1, 3, 3])); + + // count_include_pad = true: divide by kernel size (4) + let result_include = avg_pool2d_f32(x.clone(), [2, 2], [2, 2], [1, 1], true, false); + let out_include: Vec = result_include.into_data().to_vec().unwrap(); + + // count_include_pad = false: divide by actual count + let result_exclude = avg_pool2d_f32(x, [2, 2], [2, 2], [1, 1], false, false); + let out_exclude: Vec = result_exclude.into_data().to_vec().unwrap(); + + // Corner position with padding: only 1 valid element + // With count_include_pad: 1.0 / 4 = 0.25 + // Without: 1.0 / 1 = 1.0 + assert!((out_include[0] - 0.25).abs() < 1e-5); + assert!((out_exclude[0] - 1.0).abs() < 1e-5); + } + + #[test] + fn test_max_pool2d_f64() { + let x_data: Vec = (1..=16).map(|x| x as f64).collect(); + let x = FlexTensor::from_data(TensorData::new(x_data, vec![1, 1, 4, 4])); + + let result = max_pool2d_f64(x, [2, 2], [2, 2], [0, 0], [1, 1], false); + let out: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(out, vec![6.0, 8.0, 14.0, 16.0]); + } + + #[test] + fn test_max_pool2d_f16() { + let x_data: Vec = (1..=16).map(|x| f16::from_f32(x as f32)).collect(); + let x = FlexTensor::from_data(TensorData::new(x_data, vec![1, 1, 4, 4])); + + let result = max_pool2d_f16(x, [2, 2], [2, 2], [0, 0], [1, 1], false); + let out: Vec = result.into_data().to_vec().unwrap(); + + assert!((out[0].to_f32() - 6.0).abs() < 0.1); + assert!((out[1].to_f32() - 8.0).abs() < 0.1); + assert!((out[2].to_f32() - 14.0).abs() < 0.1); + assert!((out[3].to_f32() - 16.0).abs() < 0.1); + } + + #[test] + fn test_max_pool2d_bf16() { + let x_data: Vec = (1..=16).map(|x| bf16::from_f32(x as f32)).collect(); + let x = FlexTensor::from_data(TensorData::new(x_data, vec![1, 1, 4, 4])); + + let result = max_pool2d_bf16(x, [2, 2], [2, 2], [0, 0], [1, 1], false); + let out: Vec = result.into_data().to_vec().unwrap(); + + assert!((out[0].to_f32() - 6.0).abs() < 0.5); + assert!((out[1].to_f32() - 8.0).abs() < 0.5); + } + + #[test] + fn test_max_pool_backward() { + // Forward pass + let x_data: Vec = (1..=16).map(|x| x as f32).collect(); + let x = FlexTensor::from_data(TensorData::new(x_data.clone(), vec![1, 1, 4, 4])); + let (_output, indices) = + max_pool2d_with_indices_f32(x.clone(), [2, 2], [2, 2], [0, 0], [1, 1], false); + + // Backward pass: gradient of 1.0 for each output + let grad = FlexTensor::from_data(TensorData::new(vec![1.0f32; 4], vec![1, 1, 2, 2])); + + let x_grad = max_pool2d_backward_f32(x, grad, indices); + let grad_data: Vec = x_grad.into_data().to_vec().unwrap(); + + // Gradient should be 1.0 at max positions, 0.0 elsewhere + // Max positions were: 5, 7, 13, 15 (0-indexed) + assert_eq!(grad_data[5], 1.0); + assert_eq!(grad_data[7], 1.0); + assert_eq!(grad_data[13], 1.0); + assert_eq!(grad_data[15], 1.0); + assert_eq!(grad_data[0], 0.0); + } + + #[test] + fn test_avg_pool_backward() { + let x_data: Vec = (1..=16).map(|x| x as f32).collect(); + let x = FlexTensor::from_data(TensorData::new(x_data, vec![1, 1, 4, 4])); + + // Backward with gradient of 4.0 for each output (will distribute as 1.0 each) + let grad = FlexTensor::from_data(TensorData::new(vec![4.0f32; 4], vec![1, 1, 2, 2])); + + let x_grad = avg_pool2d_backward_f32(x, grad, [2, 2], [2, 2], [0, 0], false); + let grad_data: Vec = x_grad.into_data().to_vec().unwrap(); + + // Each position in input should receive grad/4 = 1.0 + assert!(grad_data.iter().all(|&v| (v - 1.0).abs() < 1e-5)); + } + + #[test] + fn test_adaptive_avg_pool_backward() { + let x_data: Vec = (1..=16).map(|x| x as f32).collect(); + let x = FlexTensor::from_data(TensorData::new(x_data, vec![1, 1, 4, 4])); + + // Backward with gradient of 4.0 for each output element + let grad = FlexTensor::from_data(TensorData::new(vec![4.0f32; 4], vec![1, 1, 2, 2])); + let x_grad = adaptive_avg_pool2d_backward_f32(x, grad); + let grad_data: Vec = x_grad.into_data().to_vec().unwrap(); + + // Each input position receives gradient from its output region + assert!(grad_data.iter().all(|&v| (v - 1.0).abs() < 1e-5)); + } + + #[test] + #[should_panic(expected = "kernel size must be > 0")] + fn test_pool_output_size_zero_kernel_panics() { + pool_output_size(4, 0, 0, 1, 1, false); + } + + #[test] + #[should_panic(expected = "stride must be > 0")] + fn test_pool_output_size_zero_stride_panics() { + pool_output_size(4, 2, 0, 0, 1, false); + } +} diff --git a/crates/burn-flex/src/ops/qtensor.rs b/crates/burn-flex/src/ops/qtensor.rs new file mode 100644 index 0000000..bdb41e5 --- /dev/null +++ b/crates/burn-flex/src/ops/qtensor.rs @@ -0,0 +1,763 @@ +//! Quantized tensor operations for the Flex backend. + +use alloc::vec::Vec; +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float; + +use burn_backend::{ + DType, ExecutionError, FloatDType, TensorData, TensorMetadata, + ops::{IntTensorOps, QTensorOps}, + quantization::{ + QuantLevel, QuantScheme, QuantStore, QuantizationParametersPrimitive, QuantizedBytes, + }, + tensor::{Device, FloatTensor, IntTensor, QuantizedTensor}, +}; +use burn_std::{Bytes, Shape, Slice, bf16, f16}; + +use super::float_storage_as_f32; +use crate::{Flex, FlexQTensor, FlexTensor, Layout}; + +impl QTensorOps for Flex { + fn q_from_data(data: TensorData, _device: &Device) -> QuantizedTensor { + let scheme = match data.dtype { + DType::QFloat(scheme) => scheme, + _ => panic!("Expected quantized dtype, got {:?}", data.dtype), + }; + + let shape = data.shape.clone(); + let num_elements = data.num_elements(); + + let q_bytes = QuantizedBytes { + bytes: data.into_bytes(), + scheme, + num_elements, + }; + + let (values, qparams) = q_bytes.into_vec_i8(); + let tensor_data = TensorData::new(values, shape); + let tensor = FlexTensor::from_data(tensor_data); + + // Use native storage since we've unpacked to i8 + let scheme = scheme.with_store(QuantStore::Native); + + FlexQTensor::new(tensor, scheme, qparams.scales) + } + + fn quantize_dynamic(tensor: FloatTensor, scheme: &QuantScheme) -> QuantizedTensor { + let shape = tensor.shape(); + let tensor = tensor.to_contiguous(); + let float_data = float_storage_as_f32(&tensor); + let (a, b) = scheme.value.range(); + let range = b - a; + + let (quantized, scales) = match scheme.level { + QuantLevel::Tensor => { + // Pass 1: find alpha = max(|min|, |max|) + let mut alpha: f32 = 0.0; + for &x in &*float_data { + let abs = x.abs(); + if abs > alpha { + alpha = abs; + } + } + let scale = validated_scale(2.0 * alpha / range); + let inv_scale = 1.0 / scale; + + // Pass 2: quantize + let quantized = float_data + .iter() + .map(|&x| (x * inv_scale).round().clamp(a, b) as i8) + .collect::>(); + + (quantized, alloc::vec![scale]) + } + QuantLevel::Block(block_size) => { + let block_elems = block_size.num_elements(); + debug_assert!( + float_data.len().is_multiple_of(block_elems), + "tensor length {} not divisible by block size {}", + float_data.len(), + block_elems + ); + let num_blocks = float_data.len() / block_elems; + let mut scales = Vec::with_capacity(num_blocks); + let mut quantized = Vec::with_capacity(float_data.len()); + + for block in float_data.chunks(block_elems) { + // Find alpha for this block + let mut alpha: f32 = 0.0; + for &x in block { + let abs = x.abs(); + if abs > alpha { + alpha = abs; + } + } + let scale = validated_scale(2.0 * alpha / range); + let inv_scale = 1.0 / scale; + scales.push(scale); + + // Quantize this block + for &x in block { + quantized.push((x * inv_scale).round().clamp(a, b) as i8); + } + } + + (quantized, scales) + } + }; + + let bytes = Bytes::from_elems(quantized); + let layout = Layout::contiguous(shape); + let qt = FlexTensor::new(bytes, layout, DType::I8); + + FlexQTensor::new(qt, scheme.with_store(QuantStore::Native), scales) + } + + fn quantize( + tensor: FloatTensor, + scheme: &QuantScheme, + qparams: QuantizationParametersPrimitive, + ) -> QuantizedTensor { + let shape = tensor.shape(); + let tensor = tensor.to_contiguous(); + let float_data = float_storage_as_f32(&tensor); + + // Extract and validate scales from the qparams tensor. The scales tensor + // shares its dtype with the float element type, which can be any of + // f32/f64/f16/bf16, so we normalise via float_storage_as_f32 instead of + // assuming f32 storage. + let scales_tensor = qparams.scales.to_contiguous(); + let scales_data = float_storage_as_f32(&scales_tensor); + let scales: Vec = scales_data.iter().copied().map(validated_scale).collect(); + + let (a, b) = scheme.value.range(); + + let quantized = match scheme.level { + QuantLevel::Tensor => { + let inv_scale = 1.0 / scales[0]; + float_data + .iter() + .map(|&x| (x * inv_scale).round().clamp(a, b) as i8) + .collect::>() + } + QuantLevel::Block(block_size) => { + let block_elems = block_size.num_elements(); + debug_assert!( + float_data.len().is_multiple_of(block_elems), + "tensor length {} not divisible by block size {}", + float_data.len(), + block_elems + ); + let mut quantized = Vec::with_capacity(float_data.len()); + for (block, &scale) in float_data.chunks(block_elems).zip(scales.iter()) { + let inv_scale = 1.0 / scale; + for &x in block { + quantized.push((x * inv_scale).round().clamp(a, b) as i8); + } + } + quantized + } + }; + + let bytes = Bytes::from_elems(quantized); + let layout = Layout::contiguous(shape); + let qt = FlexTensor::new(bytes, layout, DType::I8); + + FlexQTensor::new(qt, scheme.with_store(QuantStore::Native), scales) + } + + fn dequantize(tensor: QuantizedTensor, dtype: FloatDType) -> FloatTensor { + let shape = tensor.tensor.shape(); + let qt = tensor.tensor.to_contiguous(); + let q_data: &[i8] = qt.storage(); + + let dequantized = match tensor.scheme.level { + QuantLevel::Tensor => { + let scale = tensor.scales[0]; + q_data + .iter() + .map(|&x_q| scale * x_q as f32) + .collect::>() + } + QuantLevel::Block(block_size) => { + let block_elems = block_size.num_elements(); + q_data + .chunks(block_elems) + .zip(tensor.scales.iter()) + .flat_map(|(block, &scale)| block.iter().map(move |&x_q| scale * x_q as f32)) + .collect::>() + } + }; + + let layout = Layout::contiguous(shape); + match dtype { + FloatDType::F32 | FloatDType::Flex32 => { + FlexTensor::new(Bytes::from_elems(dequantized), layout, DType::F32) + } + FloatDType::F64 => { + let data: Vec = dequantized.iter().map(|&v| v as f64).collect(); + FlexTensor::new(Bytes::from_elems(data), layout, DType::F64) + } + FloatDType::F16 => { + let data: Vec = dequantized.iter().map(|&v| f16::from_f32(v)).collect(); + FlexTensor::new(Bytes::from_elems(data), layout, DType::F16) + } + FloatDType::BF16 => { + let data: Vec = dequantized.iter().map(|&v| bf16::from_f32(v)).collect(); + FlexTensor::new(Bytes::from_elems(data), layout, DType::BF16) + } + } + } + + fn q_to_device(tensor: QuantizedTensor, _device: &Device) -> QuantizedTensor { + tensor + } + + fn q_reshape(tensor: QuantizedTensor, shape: Shape) -> QuantizedTensor { + block_safe_layout_op(tensor, |t| t.reshape(shape)) + } + + async fn q_into_data(tensor: QuantizedTensor) -> Result { + let shape = tensor.tensor.shape(); + let scheme = tensor.scheme; + let qt = tensor.tensor.to_contiguous(); + let values: Vec = qt.storage::().to_vec(); + + Ok(TensorData::quantized( + values, + shape.to_vec(), + scheme, + &tensor.scales, + )) + } + + fn q_swap_dims( + tensor: QuantizedTensor, + dim1: usize, + dim2: usize, + ) -> QuantizedTensor { + block_safe_layout_op(tensor, |t| t.transpose(dim1, dim2)) + } + + fn q_permute(tensor: QuantizedTensor, axes: &[usize]) -> QuantizedTensor { + block_safe_layout_op(tensor, |t| t.permute(axes)) + } + + fn q_flip(tensor: QuantizedTensor, axes: &[usize]) -> QuantizedTensor { + block_safe_layout_op(tensor, |t| crate::ops::flip::flip(t, axes)) + } + + fn q_expand(tensor: QuantizedTensor, shape: Shape) -> QuantizedTensor { + block_safe_layout_op(tensor, |t| crate::ops::expand::expand(t, shape)) + } + + fn q_select( + tensor: QuantizedTensor, + dim: usize, + indices: IntTensor, + ) -> QuantizedTensor { + match tensor.scheme.level { + QuantLevel::Tensor => FlexQTensor::new( + crate::ops::gather_scatter::select::(tensor.tensor, dim, indices), + tensor.scheme, + tensor.scales, + ), + QuantLevel::Block(_) => { + let scheme = tensor.scheme; + let float_tensor = Flex::dequantize(tensor, FloatDType::F32); + let result = crate::ops::gather_scatter::select::(float_tensor, dim, indices); + Flex::quantize_dynamic(result, &scheme) + } + } + } + + fn q_slice(tensor: QuantizedTensor, slices: &[Slice]) -> QuantizedTensor { + block_safe_layout_op(tensor, |t| crate::ops::slice::slice(t, slices)) + } + + fn q_argmax( + tensor: QuantizedTensor, + dim: usize, + out_dtype: burn_std::IntDType, + ) -> IntTensor { + let result = crate::ops::reduce::argmax(tensor.tensor, dim); + if result.dtype() != DType::from(out_dtype) { + Flex::int_cast(result, out_dtype) + } else { + result + } + } + + fn q_argmin( + tensor: QuantizedTensor, + dim: usize, + out_dtype: burn_std::IntDType, + ) -> IntTensor { + let result = crate::ops::reduce::argmin(tensor.tensor, dim); + if result.dtype() != DType::from(out_dtype) { + Flex::int_cast(result, out_dtype) + } else { + result + } + } + + fn q_gather( + dim: usize, + tensor: QuantizedTensor, + indices: IntTensor, + ) -> QuantizedTensor { + match tensor.scheme.level { + QuantLevel::Tensor => FlexQTensor::new( + crate::ops::gather_scatter::gather::(tensor.tensor, dim, indices), + tensor.scheme, + tensor.scales, + ), + QuantLevel::Block(_) => { + let scheme = tensor.scheme; + let float_tensor = Flex::dequantize(tensor, FloatDType::F32); + let result = crate::ops::gather_scatter::gather::(float_tensor, dim, indices); + Flex::quantize_dynamic(result, &scheme) + } + } + } +} + +/// Apply a layout operation to a quantized tensor. +/// For block-quantized tensors, dequantizes and requantizes to preserve +/// correct scale-to-block mapping. +fn block_safe_layout_op( + qtensor: FlexQTensor, + op: impl FnOnce(FlexTensor) -> FlexTensor, +) -> FlexQTensor { + match qtensor.scheme.level { + QuantLevel::Tensor => FlexQTensor::new(op(qtensor.tensor), qtensor.scheme, qtensor.scales), + QuantLevel::Block(_) => { + let scheme = qtensor.scheme; + let float_tensor = Flex::dequantize(qtensor, FloatDType::F32); + let result = op(float_tensor); + Flex::quantize_dynamic(result, &scheme) + } + } +} + +/// Ensure scale is finite and nonzero to avoid division by zero or NaN propagation. +fn validated_scale(scale: f32) -> f32 { + if scale.is_normal() { + scale + } else { + f32::MIN_POSITIVE + } +} + +// Tests kept here exercise flex-specific behavior: quantization scheme +// roundtrips, per-block / dynamic quantization, block-quantized layout +// ops (transpose / select / flip dequantize), and f16/f64 dequantize +// dtype paths. Plain layout-preservation / select / slice / argmax / +// argmin / gather tests are covered generically in +// crates/burn-backend-tests/tests/tensor/float/quantization/ops/extended/ +// so they run on every backend. +#[cfg(test)] +mod tests { + use super::*; + use burn_backend::{TensorMetadata, quantization::QuantValue}; + + #[test] + fn test_quantize_dequantize_roundtrip() { + // Create a float tensor + let values = vec![0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0]; + let tensor = FlexTensor::from_data(TensorData::new(values.clone(), [2, 3])); + + let scheme = QuantScheme::default() + .with_value(QuantValue::Q8S) + .with_store(QuantStore::Native); + + // Compute scale: symmetric, so scale = 2 * max(|min|, |max|) / (b - a) + // max_abs = 5.0, range = 127 - (-127) = 254 + // scale = 2 * 5.0 / 254 = 0.03937008 + let scale: f32 = 2.0 * 5.0 / 254.0; + let scales_tensor = FlexTensor::from_data(TensorData::new(vec![scale], [1])); + + let qparams = QuantizationParametersPrimitive { + scales: scales_tensor, + }; + + // Quantize + let qtensor = Flex::quantize(tensor, &scheme, qparams); + assert_eq!(qtensor.tensor.shape().to_vec(), vec![2, 3]); + assert_eq!(qtensor.tensor.dtype(), DType::I8); + + // Check quantized values + let q_vals: &[i8] = qtensor.tensor.storage(); + // 0 / 0.03937 = 0, 1 / 0.03937 = 25.4 -> 25, etc. + assert_eq!(q_vals[0], 0); + assert_eq!(q_vals[1], 25); + assert_eq!(q_vals[5], 127); + + // Dequantize + let result = Flex::dequantize(qtensor, FloatDType::F32); + assert_eq!(result.shape().to_vec(), vec![2, 3]); + assert_eq!(result.dtype(), DType::F32); + + let result_vals: &[f32] = result.storage(); + // Values should be approximately equal (quantization introduces small errors) + for (orig, deq) in values.iter().zip(result_vals.iter()) { + assert!((orig - deq).abs() < 0.05, "orig={orig}, dequantized={deq}"); + } + } + + #[test] + fn test_quantize_dequantize_negative_values() { + let values = vec![-3.0f32, -1.5, 0.0, 1.5, 3.0]; + let tensor = FlexTensor::from_data(TensorData::new(values.clone(), [5])); + + let scheme = QuantScheme::default() + .with_value(QuantValue::Q8S) + .with_store(QuantStore::Native); + + let scale: f32 = 2.0 * 3.0 / 254.0; + let scales_tensor = FlexTensor::from_data(TensorData::new(vec![scale], [1])); + + let qparams = QuantizationParametersPrimitive { + scales: scales_tensor, + }; + + let qtensor = Flex::quantize(tensor, &scheme, qparams); + let result = Flex::dequantize(qtensor, FloatDType::F32); + let result_vals: &[f32] = result.storage(); + + for (orig, deq) in values.iter().zip(result_vals.iter()) { + assert!((orig - deq).abs() < 0.05, "orig={orig}, dequantized={deq}"); + } + } + + #[test] + fn test_q_from_data_into_data_roundtrip() { + // Create quantized TensorData the standard way + let values = vec![0i8, 25, 51, 76, 102, 127]; + let scale = 0.03937008f32; + let scheme = QuantScheme::default() + .with_value(QuantValue::Q8S) + .with_store(QuantStore::Native); + + let data = TensorData::quantized(values.clone(), [2, 3], scheme, &[scale]); + + // Load into FlexQTensor + let qtensor = Flex::q_from_data(data, &Default::default()); + assert_eq!(qtensor.tensor.shape().to_vec(), vec![2, 3]); + assert_eq!(qtensor.scales, vec![scale]); + + // Dequantize and check values + let float_tensor = Flex::dequantize(qtensor, FloatDType::F32); + let result: &[f32] = float_tensor.storage(); + assert!((result[0]).abs() < 0.01); // 0 * scale ~ 0 + assert!((result[5] - 5.0).abs() < 0.05); // 127 * scale ~ 5.0 + } + + #[test] + fn test_quantize_zero_tensor() { + let values = vec![0.0f32; 4]; + let tensor = FlexTensor::from_data(TensorData::new(values, [4])); + + let scheme = QuantScheme::default() + .with_value(QuantValue::Q8S) + .with_store(QuantStore::Native); + + // Scale of 0 should be handled gracefully + let scales_tensor = FlexTensor::from_data(TensorData::new(vec![0.0f32], [1])); + let qparams = QuantizationParametersPrimitive { + scales: scales_tensor, + }; + + let qtensor = Flex::quantize(tensor, &scheme, qparams); + let q_vals: &[i8] = qtensor.tensor.storage(); + assert_eq!(q_vals, &[0, 0, 0, 0]); + } + + #[test] + fn test_quantize_dynamic_roundtrip() { + let values = vec![-3.0f32, -1.5, 0.0, 1.5, 3.0, 4.5]; + let tensor = FlexTensor::from_data(TensorData::new(values.clone(), [2, 3])); + + let scheme = QuantScheme::default() + .with_value(QuantValue::Q8S) + .with_store(QuantStore::Native); + + let qtensor = Flex::quantize_dynamic(tensor, &scheme); + assert_eq!(qtensor.tensor.shape().to_vec(), vec![2, 3]); + assert_eq!(qtensor.scales.len(), 1); + + // Scale should be 2 * 4.5 / 254 + let expected_scale: f32 = 2.0 * 4.5 / 254.0; + assert!( + (qtensor.scales[0] - expected_scale).abs() < 1e-6, + "scale={}, expected={}", + qtensor.scales[0], + expected_scale + ); + + let result = Flex::dequantize(qtensor, FloatDType::F32); + let result_vals: &[f32] = result.storage(); + for (orig, deq) in values.iter().zip(result_vals.iter()) { + assert!((orig - deq).abs() < 0.1, "orig={orig}, dequantized={deq}"); + } + } + + #[test] + fn test_per_block_quantize_dequantize() { + use burn_std::quantization::BlockSize; + + let values = vec![0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]; + let tensor = FlexTensor::from_data(TensorData::new(values.clone(), [8])); + + let block_size = BlockSize::new([4]); + let scheme = QuantScheme::default() + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::Block(block_size)) + .with_store(QuantStore::Native); + + // Block 1: [0, 1, 2, 3] -> max_abs=3, scale = 6/254 + // Block 2: [4, 5, 6, 7] -> max_abs=7, scale = 14/254 + let scale_1: f32 = 2.0 * 3.0 / 254.0; + let scale_2: f32 = 2.0 * 7.0 / 254.0; + let scales_tensor = FlexTensor::from_data(TensorData::new(vec![scale_1, scale_2], [2])); + + let qparams = QuantizationParametersPrimitive { + scales: scales_tensor, + }; + + let qtensor = Flex::quantize(tensor, &scheme, qparams); + assert_eq!(qtensor.scales.len(), 2); + + let result = Flex::dequantize(qtensor, FloatDType::F32); + let result_vals: &[f32] = result.storage(); + + for (orig, deq) in values.iter().zip(result_vals.iter()) { + assert!((orig - deq).abs() < 0.1, "orig={orig}, dequantized={deq}"); + } + } + + #[test] + fn test_quantize_dynamic_block() { + use burn_std::quantization::BlockSize; + + let values = vec![-2.0f32, -1.0, 0.0, 1.0, 4.0, 5.0, 6.0, 7.0]; + let tensor = FlexTensor::from_data(TensorData::new(values.clone(), [8])); + + let block_size = BlockSize::new([4]); + let scheme = QuantScheme::default() + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::Block(block_size)) + .with_store(QuantStore::Native); + + let qtensor = Flex::quantize_dynamic(tensor, &scheme); + assert_eq!(qtensor.scales.len(), 2); + + // Block 1: [-2, -1, 0, 1] -> alpha=2, scale = 4/254 + // Block 2: [4, 5, 6, 7] -> alpha=7, scale = 14/254 + let expected_scale_1: f32 = 2.0 * 2.0 / 254.0; + let expected_scale_2: f32 = 2.0 * 7.0 / 254.0; + assert!((qtensor.scales[0] - expected_scale_1).abs() < 1e-6); + assert!((qtensor.scales[1] - expected_scale_2).abs() < 1e-6); + + let result = Flex::dequantize(qtensor, FloatDType::F32); + let result_vals: &[f32] = result.storage(); + for (orig, deq) in values.iter().zip(result_vals.iter()) { + assert!((orig - deq).abs() < 0.1, "orig={orig}, dequantized={deq}"); + } + } + + #[test] + fn test_quantize_dynamic_q8f() { + // Q8F uses asymmetric range [-128, 127] + let values = vec![-5.0f32, -2.5, 0.0, 2.5, 5.0, 7.5]; + let tensor = FlexTensor::from_data(TensorData::new(values.clone(), [6])); + + let scheme = QuantScheme::default() + .with_value(QuantValue::Q8F) + .with_store(QuantStore::Native); + + let qtensor = Flex::quantize_dynamic(tensor, &scheme); + + // Q8F range: [-128, 127], so range = 255 + // alpha = 7.5, scale = 2 * 7.5 / 255 + let expected_scale: f32 = 2.0 * 7.5 / 255.0; + assert!( + (qtensor.scales[0] - expected_scale).abs() < 1e-6, + "scale={}, expected={}", + qtensor.scales[0], + expected_scale + ); + + let result = Flex::dequantize(qtensor, FloatDType::F32); + let result_vals: &[f32] = result.storage(); + for (orig, deq) in values.iter().zip(result_vals.iter()) { + assert!((orig - deq).abs() < 0.1, "orig={orig}, dequantized={deq}"); + } + } + + #[test] + fn test_block_quantized_transpose_dequantize() { + use burn_std::quantization::BlockSize; + + // 2x4 tensor, 2 blocks of 4 + let values = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; + let tensor = FlexTensor::from_data(TensorData::new(values, [2, 4])); + + let block_size = BlockSize::new([4]); + let scheme = QuantScheme::default() + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::Block(block_size)) + .with_store(QuantStore::Native); + + let qtensor = Flex::quantize_dynamic(tensor, &scheme); + + // Transpose to [4, 2], then dequantize + let transposed = Flex::q_swap_dims(qtensor, 0, 1); + assert_eq!(transposed.tensor.shape().to_vec(), vec![4, 2]); + + let result = Flex::dequantize(transposed, FloatDType::F32); + let result_vals: &[f32] = result.storage(); + + // Original [[1,2,3,4],[5,6,7,8]] transposed to [[1,5],[2,6],[3,7],[4,8]] + let expected = [1.0f32, 5.0, 2.0, 6.0, 3.0, 7.0, 4.0, 8.0]; + for (exp, deq) in expected.iter().zip(result_vals.iter()) { + assert!( + (exp - deq).abs() < 0.15, + "expected={exp}, dequantized={deq}" + ); + } + } + + #[test] + fn test_block_quantized_select() { + use burn_std::quantization::BlockSize; + + // 2x4 tensor, 2 blocks of 4 + let values = vec![1.0f32, 2.0, 3.0, 4.0, 10.0, 20.0, 30.0, 40.0]; + let tensor = FlexTensor::from_data(TensorData::new(values, [2, 4])); + + let block_size = BlockSize::new([4]); + let scheme = QuantScheme::default() + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::Block(block_size)) + .with_store(QuantStore::Native); + + let qtensor = Flex::quantize_dynamic(tensor, &scheme); + + // Select row 1 -> [10, 20, 30, 40] + let indices = FlexTensor::from_data(TensorData::new(vec![1i64], [1])); + let selected = Flex::q_select(qtensor, 0, indices); + assert_eq!(selected.tensor.shape().to_vec(), vec![1, 4]); + + let result = Flex::dequantize(selected, FloatDType::F32); + let result_vals: &[f32] = result.storage(); + let expected = [10.0f32, 20.0, 30.0, 40.0]; + for (exp, deq) in expected.iter().zip(result_vals.iter()) { + assert!((exp - deq).abs() < 0.5, "expected={exp}, dequantized={deq}"); + } + } + + #[test] + fn test_block_quantized_flip_dequantize() { + use burn_std::quantization::BlockSize; + + let values = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; + let tensor = FlexTensor::from_data(TensorData::new(values, [2, 4])); + + let block_size = BlockSize::new([4]); + let scheme = QuantScheme::default() + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::Block(block_size)) + .with_store(QuantStore::Native); + + let qtensor = Flex::quantize_dynamic(tensor, &scheme); + + // Flip along axis 0: [[5,6,7,8],[1,2,3,4]] + let flipped = Flex::q_flip(qtensor, &[0]); + assert_eq!(flipped.tensor.shape().to_vec(), vec![2, 4]); + + let result = Flex::dequantize(flipped, FloatDType::F32); + let result_vals: &[f32] = result.storage(); + let expected = [5.0f32, 6.0, 7.0, 8.0, 1.0, 2.0, 3.0, 4.0]; + for (exp, deq) in expected.iter().zip(result_vals.iter()) { + assert!( + (exp - deq).abs() < 0.15, + "expected={exp}, dequantized={deq}" + ); + } + } + + #[test] + fn test_quantize_dynamic_f64_tensor() { + use burn_backend::quantization::QuantValue; + + let values = vec![0.0f64, 1.0, 2.0, 3.0, 4.0, 5.0]; + let tensor = FlexTensor::new( + Bytes::from_elems(values), + Layout::contiguous([6].into()), + DType::F64, + ); + assert_eq!(tensor.dtype(), DType::F64); + + let scheme = QuantScheme::default() + .with_value(QuantValue::Q8S) + .with_store(QuantStore::Native); + + let qtensor = Flex::quantize_dynamic(tensor, &scheme); + assert_eq!(qtensor.tensor.dtype(), DType::I8); + + // Dequantize and verify round-trip accuracy + let result = Flex::dequantize(qtensor, FloatDType::F32); + let result_vals: &[f32] = result.storage(); + let expected = [0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0]; + for (exp, deq) in expected.iter().zip(result_vals.iter()) { + assert!( + (exp - deq).abs() < 0.15, + "expected={exp}, dequantized={deq}" + ); + } + } + + #[test] + fn test_dequantize_f64() { + let values = vec![0.0f32, 1.0, 2.0, 3.0]; + let tensor = FlexTensor::from_data(TensorData::new(values.clone(), [4])); + + let scheme = QuantScheme::default() + .with_value(QuantValue::Q8S) + .with_store(QuantStore::Native); + + let qtensor = Flex::quantize_dynamic(tensor, &scheme); + let result = Flex::dequantize(qtensor, FloatDType::F64); + assert_eq!(result.dtype(), DType::F64); + let result_vals: &[f64] = result.storage(); + for (orig, deq) in values.iter().zip(result_vals.iter()) { + assert!( + (*orig as f64 - deq).abs() < 0.05, + "orig={orig}, dequantized={deq}" + ); + } + } + + #[test] + fn test_dequantize_f16() { + let values = vec![0.0f32, 1.0, 2.0, 3.0]; + let tensor = FlexTensor::from_data(TensorData::new(values.clone(), [4])); + + let scheme = QuantScheme::default() + .with_value(QuantValue::Q8S) + .with_store(QuantStore::Native); + + let qtensor = Flex::quantize_dynamic(tensor, &scheme); + let result = Flex::dequantize(qtensor, FloatDType::F16); + assert_eq!(result.dtype(), DType::F16); + let result_vals: &[f16] = result.storage(); + for (orig, deq) in values.iter().zip(result_vals.iter()) { + assert!( + (*orig - f32::from(*deq)).abs() < 0.05, + "orig={orig}, dequantized={deq}" + ); + } + } +} diff --git a/crates/burn-flex/src/ops/reduce.rs b/crates/burn-flex/src/ops/reduce.rs new file mode 100644 index 0000000..93b682f --- /dev/null +++ b/crates/burn-flex/src/ops/reduce.rs @@ -0,0 +1,2559 @@ +//! Reduction operations for FlexTensor. +//! +//! Optimized with: +//! - Strided iteration (no copy for non-contiguous tensors) +//! - Portable SIMD via macerator (NEON, AVX2, SIMD128, scalar fallback) +//! - Rayon parallelism for large tensors + +use alloc::vec; +use alloc::vec::Vec; +use burn_backend::{DType, Element}; +use burn_std::{Bytes, Shape, bf16, f16}; + +use crate::strided_index::StridedIter; +use crate::{FlexTensor, Layout}; + +use super::{INDEX_DTYPE, float_storage_as_f32}; + +/// Assert that a dimension size fits in `isize`, which is required for index-producing +/// operations (argmax, argmin, *_with_indices) that store dimension indices as `isize`. +#[inline(always)] +fn assert_dim_fits_isize(dim_size: usize, dim: usize) { + assert!( + dim_size <= isize::MAX as usize, + "dimension {dim} has size {dim_size} which exceeds isize::MAX" + ); +} + +#[cfg(feature = "simd")] +use crate::simd::kernels; + +#[cfg(feature = "simd")] +use crate::simd::aligned; + +#[cfg(feature = "rayon")] +use rayon::prelude::*; + +/// Truncate an i64 to a smaller Pod type, keeping the low-order bytes. +/// Endian-safe: works correctly on both little-endian and big-endian targets. +fn truncate_i64_to_pod(value: i64) -> E { + let bytes = value.to_ne_bytes(); + let size = core::mem::size_of::(); + debug_assert!(size <= core::mem::size_of::()); + let offset = if cfg!(target_endian = "big") { + core::mem::size_of::() - size + } else { + 0 + }; + bytemuck::pod_read_unaligned(&bytes[offset..offset + size]) +} + +// ============================================================================ +// Sum (all elements) +// ============================================================================ + +/// Sum all elements in a tensor, returning a scalar tensor. +pub fn sum(tensor: FlexTensor) -> FlexTensor { + match tensor.dtype() { + DType::F32 => sum_f32(&tensor), + DType::F64 => sum_impl::(&tensor), + DType::F16 => reduce_scalar_half(&tensor, |a, b| a + b, 0.0, f16::to_f32, f16::from_f32), + DType::BF16 => reduce_scalar_half(&tensor, |a, b| a + b, 0.0, bf16::to_f32, bf16::from_f32), + DType::I8 => sum_impl_widening::(&tensor), + DType::I16 => sum_impl_widening::(&tensor), + DType::I32 => sum_impl_widening::(&tensor), + DType::I64 => sum_impl::(&tensor), + DType::U8 => sum_impl_widening::(&tensor), + DType::U16 => sum_impl_widening::(&tensor), + DType::U32 => sum_impl_widening::(&tensor), + DType::U64 => sum_impl::(&tensor), + _ => panic!("sum: unsupported dtype {:?}", tensor.dtype()), + } +} + +/// Optimized f32 sum with SIMD and parallelism. +fn sum_f32(tensor: &FlexTensor) -> FlexTensor { + let result = match tensor.layout().contiguous_offsets() { + Some((start, end)) => { + let data: &[f32] = tensor.storage(); + let slice = &data[start..end]; + sum_f32_contiguous(slice) + } + None => { + // Non-contiguous: check if we can sum the buffer directly. + // For transposed tensors that use all elements (no slicing), + // the sum is the same regardless of element order. + let data: &[f32] = tensor.storage(); + let elem_count = tensor.layout().num_elements(); + + if data.len() == elem_count { + // Tensor uses entire buffer - sum directly (order doesn't matter for sum) + sum_f32_contiguous(data) + } else { + // Sliced or partial view - must use strided iteration + StridedIter::new(tensor.layout()).map(|idx| data[idx]).sum() + } + } + }; + + let bytes = Bytes::from_elems(vec![result]); + FlexTensor::new(bytes, Layout::contiguous(Shape::from(vec![1])), DType::F32) +} + +/// SIMD + parallel sum for contiguous f32 slice. +/// +/// The parallel threshold is higher than the general `PARALLEL_THRESHOLD` because +/// sum is memory-bound and L2-resident data (< ~16 MiB / 4M f32 elements on +/// Apple M-series) doesn't benefit from rayon's task dispatch overhead. +#[inline] +fn sum_f32_contiguous(data: &[f32]) -> f32 { + #[cfg(feature = "rayon")] + if data.len() >= 4 * 1024 * 1024 { + return sum_f32_parallel(data); + } + + #[cfg(feature = "simd")] + { + kernels::sum_f32(data) + } + + #[cfg(not(feature = "simd"))] + { + data.iter().copied().sum() + } +} + +/// Parallel sum using rayon with SIMD per chunk. +#[cfg(feature = "rayon")] +#[inline] +fn sum_f32_parallel(data: &[f32]) -> f32 { + const CHUNK_SIZE: usize = 64 * 1024; // 64K elements per chunk + + data.par_chunks(CHUNK_SIZE) + .map(|chunk| { + #[cfg(feature = "simd")] + { + kernels::sum_f32(chunk) + } + #[cfg(not(feature = "simd"))] + { + chunk.iter().copied().sum::() + } + }) + .sum() +} + +fn sum_impl( + tensor: &FlexTensor, +) -> FlexTensor { + let result: E = match tensor.layout().contiguous_offsets() { + Some((start, end)) => { + let data: &[E] = tensor.storage(); + data[start..end].iter().copied().sum() + } + None => { + let data: &[E] = tensor.storage(); + StridedIter::new(tensor.layout()).map(|idx| data[idx]).sum() + } + }; + + let bytes = Bytes::from_elems(vec![result]); + FlexTensor::new( + bytes, + Layout::contiguous(Shape::from(vec![1])), + tensor.dtype(), + ) +} + +/// Widening scalar reduction for small integer types: accumulate in i64 to avoid overflow. +macro_rules! widening_scalar_reduce { + ($name:ident, $fold:expr, $init:expr) => { + fn $name(tensor: &FlexTensor) -> FlexTensor + where + E: Element + bytemuck::Pod + Default, + i64: From, + { + let total: i64 = match tensor.layout().contiguous_offsets() { + Some((start, end)) => { + let data: &[E] = tensor.storage(); + data[start..end] + .iter() + .fold($init, |acc, x| ($fold)(acc, i64::from(*x))) + } + None => { + let data: &[E] = tensor.storage(); + StridedIter::new(tensor.layout()) + .fold($init, |acc, idx| ($fold)(acc, i64::from(data[idx]))) + } + }; + // Truncate back to target type (wrapping, matches PyTorch) + let result: E = truncate_i64_to_pod(total); + let bytes = Bytes::from_elems(vec![result]); + FlexTensor::new( + bytes, + Layout::contiguous(Shape::from(vec![1])), + tensor.dtype(), + ) + } + }; +} + +widening_scalar_reduce!( + sum_impl_widening, + |acc: i64, x: i64| acc.wrapping_add(x), + 0i64 +); +widening_scalar_reduce!( + prod_impl_widening, + |acc: i64, x: i64| acc.wrapping_mul(x), + 1i64 +); + +/// Scalar reduction for half-precision types, accumulating in f32. +fn reduce_scalar_half( + tensor: &FlexTensor, + fold: fn(f32, f32) -> f32, + init: f32, + to_f32: fn(E) -> f32, + from_f32: fn(f32) -> E, +) -> FlexTensor +where + E: Element + bytemuck::Pod, +{ + let result: f32 = match tensor.layout().contiguous_offsets() { + Some((start, end)) => { + let data: &[E] = tensor.storage(); + data[start..end] + .iter() + .fold(init, |acc, x| fold(acc, to_f32(*x))) + } + None => { + let data: &[E] = tensor.storage(); + StridedIter::new(tensor.layout()).fold(init, |acc, idx| fold(acc, to_f32(data[idx]))) + } + }; + + let bytes = Bytes::from_elems(vec![from_f32(result)]); + FlexTensor::new(bytes, Layout::contiguous(Shape::from(vec![1])), E::dtype()) +} + +// ============================================================================ +// Sum along dimension +// ============================================================================ + +/// Sum along a dimension, keeping the dimension with size 1. +pub fn sum_dim(tensor: FlexTensor, dim: usize) -> FlexTensor { + match tensor.dtype() { + DType::F32 => reduce_dim_f32(&tensor, dim, ReduceOp::Sum), + DType::F64 => reduce_dim_impl::(&tensor, dim, 0.0, |acc, x| acc + x), + DType::F16 => reduce_dim_half( + &tensor, + dim, + 0.0, + |acc, x| acc + x, + f16::to_f32, + f16::from_f32, + ), + DType::BF16 => reduce_dim_half( + &tensor, + dim, + 0.0, + |acc, x| acc + x, + bf16::to_f32, + bf16::from_f32, + ), + DType::I8 => reduce_dim_widening::(&tensor, dim, 0, |acc, x| acc.wrapping_add(x)), + DType::I16 => reduce_dim_widening::(&tensor, dim, 0, |acc, x| acc.wrapping_add(x)), + DType::I32 => reduce_dim_widening::(&tensor, dim, 0, |acc, x| acc.wrapping_add(x)), + DType::I64 => reduce_dim_impl::(&tensor, dim, 0, |acc, x| acc + x), + DType::U8 => reduce_dim_widening::(&tensor, dim, 0, |acc, x| acc.wrapping_add(x)), + DType::U16 => reduce_dim_widening::(&tensor, dim, 0, |acc, x| acc.wrapping_add(x)), + DType::U32 => reduce_dim_widening::(&tensor, dim, 0, |acc, x| acc.wrapping_add(x)), + DType::U64 => reduce_dim_impl::(&tensor, dim, 0, |acc, x| acc + x), + _ => panic!("sum_dim: unsupported dtype {:?}", tensor.dtype()), + } +} + +/// Mean along a dimension, keeping the dimension with size 1. +pub fn mean_dim(tensor: FlexTensor, dim: usize) -> FlexTensor { + let dim_size = tensor.layout().shape()[dim]; + assert!( + dim_size > 0, + "mean_dim: cannot take mean of empty dimension" + ); + let dtype = tensor.dtype(); + + // Half-precision types fuse sum+divide in f32 to avoid overflow when the + // intermediate sum exceeds f16::MAX, so they don't go through sum_dim. + match dtype { + DType::F16 => return mean_dim_half::(&tensor, dim), + DType::BF16 => return mean_dim_half::(&tensor, dim), + _ => {} + } + + let sum_result = sum_dim(tensor, dim); + + // Divide by dimension size + match dtype { + DType::F32 => scalar_div::(sum_result, dim_size as f32), + DType::F64 => scalar_div::(sum_result, dim_size as f64), + DType::I8 => { + let divisor = dim_size as i32; + let mut tensor = sum_result; + let data: &mut [i8] = tensor.storage_mut(); + for x in data.iter_mut() { + *x = ((*x as i32) / divisor) as i8; + } + tensor + } + DType::I16 => { + let divisor = dim_size as i32; + let mut tensor = sum_result; + let data: &mut [i16] = tensor.storage_mut(); + for x in data.iter_mut() { + *x = ((*x as i32) / divisor) as i16; + } + tensor + } + DType::I32 => scalar_div::(sum_result, dim_size as i32), + DType::I64 => scalar_div::(sum_result, dim_size as i64), + DType::U8 => { + let divisor = dim_size as u32; + let mut tensor = sum_result; + let data: &mut [u8] = tensor.storage_mut(); + for x in data.iter_mut() { + *x = ((*x as u32) / divisor) as u8; + } + tensor + } + DType::U16 => { + let divisor = dim_size as u32; + let mut tensor = sum_result; + let data: &mut [u16] = tensor.storage_mut(); + for x in data.iter_mut() { + *x = ((*x as u32) / divisor) as u16; + } + tensor + } + DType::U32 => scalar_div::(sum_result, dim_size as u32), + DType::U64 => scalar_div::(sum_result, dim_size as u64), + _ => panic!("mean_dim: unsupported dtype {:?}", dtype), + } +} + +/// Product of all elements in a tensor, returning a scalar tensor. +pub fn prod(tensor: FlexTensor) -> FlexTensor { + match tensor.dtype() { + DType::F32 => prod_impl::(&tensor), + DType::F64 => prod_impl::(&tensor), + DType::F16 => reduce_scalar_half(&tensor, |a, b| a * b, 1.0, f16::to_f32, f16::from_f32), + DType::BF16 => reduce_scalar_half(&tensor, |a, b| a * b, 1.0, bf16::to_f32, bf16::from_f32), + DType::I8 => prod_impl_widening::(&tensor), + DType::I16 => prod_impl_widening::(&tensor), + DType::I32 => prod_impl_widening::(&tensor), + DType::I64 => prod_impl::(&tensor), + DType::U8 => prod_impl_widening::(&tensor), + DType::U16 => prod_impl_widening::(&tensor), + DType::U32 => prod_impl_widening::(&tensor), + DType::U64 => prod_impl::(&tensor), + _ => panic!("prod: unsupported dtype {:?}", tensor.dtype()), + } +} + +fn prod_impl( + tensor: &FlexTensor, +) -> FlexTensor { + let result: E = match tensor.layout().contiguous_offsets() { + Some((start, end)) => { + let data: &[E] = tensor.storage(); + data[start..end].iter().copied().product() + } + None => { + let data: &[E] = tensor.storage(); + StridedIter::new(tensor.layout()) + .map(|idx| data[idx]) + .product() + } + }; + + let bytes = Bytes::from_elems(vec![result]); + FlexTensor::new( + bytes, + Layout::contiguous(Shape::from(vec![1])), + tensor.dtype(), + ) +} + +/// Product along a dimension, keeping the dimension with size 1. +pub fn prod_dim(tensor: FlexTensor, dim: usize) -> FlexTensor { + match tensor.dtype() { + DType::F32 => reduce_dim_f32(&tensor, dim, ReduceOp::Prod), + DType::F64 => reduce_dim_impl::(&tensor, dim, 1.0, |acc, x| acc * x), + DType::F16 => reduce_dim_half( + &tensor, + dim, + 1.0, + |acc, x| acc * x, + f16::to_f32, + f16::from_f32, + ), + DType::BF16 => reduce_dim_half( + &tensor, + dim, + 1.0, + |acc, x| acc * x, + bf16::to_f32, + bf16::from_f32, + ), + DType::I8 => reduce_dim_widening::(&tensor, dim, 1, |acc, x| acc.wrapping_mul(x)), + DType::I16 => reduce_dim_widening::(&tensor, dim, 1, |acc, x| acc.wrapping_mul(x)), + DType::I32 => reduce_dim_widening::(&tensor, dim, 1, |acc, x| acc.wrapping_mul(x)), + DType::I64 => reduce_dim_impl::(&tensor, dim, 1, |acc, x| acc * x), + DType::U8 => reduce_dim_widening::(&tensor, dim, 1, |acc, x| acc.wrapping_mul(x)), + DType::U16 => reduce_dim_widening::(&tensor, dim, 1, |acc, x| acc.wrapping_mul(x)), + DType::U32 => reduce_dim_widening::(&tensor, dim, 1, |acc, x| acc.wrapping_mul(x)), + DType::U64 => reduce_dim_impl::(&tensor, dim, 1, |acc, x| acc * x), + _ => panic!("prod_dim: unsupported dtype {:?}", tensor.dtype()), + } +} + +// ============================================================================ +// Max / Min (all elements) +// ============================================================================ + +/// Max of all elements, returning a scalar tensor of shape \[1\]. +pub fn max(tensor: FlexTensor) -> FlexTensor { + match tensor.dtype() { + DType::F32 => max_f32_reduce(&tensor), + DType::F64 => max_impl::(&tensor), + DType::F16 => reduce_scalar_half( + &tensor, + f32::max, + f32::NEG_INFINITY, + f16::to_f32, + f16::from_f32, + ), + DType::BF16 => reduce_scalar_half( + &tensor, + f32::max, + f32::NEG_INFINITY, + bf16::to_f32, + bf16::from_f32, + ), + DType::I8 => max_impl::(&tensor), + DType::I16 => max_impl::(&tensor), + DType::I32 => max_impl::(&tensor), + DType::I64 => max_impl::(&tensor), + DType::U8 => max_impl::(&tensor), + DType::U16 => max_impl::(&tensor), + DType::U32 => max_impl::(&tensor), + DType::U64 => max_impl::(&tensor), + _ => panic!("max: unsupported dtype {:?}", tensor.dtype()), + } +} + +/// Min of all elements, returning a scalar tensor of shape \[1\]. +pub fn min(tensor: FlexTensor) -> FlexTensor { + match tensor.dtype() { + DType::F32 => min_f32_reduce(&tensor), + DType::F64 => min_impl::(&tensor), + DType::F16 => { + reduce_scalar_half(&tensor, f32::min, f32::INFINITY, f16::to_f32, f16::from_f32) + } + DType::BF16 => reduce_scalar_half( + &tensor, + f32::min, + f32::INFINITY, + bf16::to_f32, + bf16::from_f32, + ), + DType::I8 => min_impl::(&tensor), + DType::I16 => min_impl::(&tensor), + DType::I32 => min_impl::(&tensor), + DType::I64 => min_impl::(&tensor), + DType::U8 => min_impl::(&tensor), + DType::U16 => min_impl::(&tensor), + DType::U32 => min_impl::(&tensor), + DType::U64 => min_impl::(&tensor), + _ => panic!("min: unsupported dtype {:?}", tensor.dtype()), + } +} + +fn max_f32_reduce(tensor: &FlexTensor) -> FlexTensor { + let result = match tensor.layout().contiguous_offsets() { + Some((start, end)) => { + let data: &[f32] = tensor.storage(); + max_f32_contiguous(&data[start..end]) + } + None => { + let data: &[f32] = tensor.storage(); + let elem_count = tensor.layout().num_elements(); + if data.len() == elem_count { + // Non-contiguous but uses all elements (e.g., transposed) + max_f32_contiguous(data) + } else { + StridedIter::new(tensor.layout()) + .map(|idx| data[idx]) + .reduce(|a, b| if a >= b { a } else { b }) + .expect("max: tensor must not be empty") + } + } + }; + + let bytes = Bytes::from_elems(vec![result]); + FlexTensor::new(bytes, Layout::contiguous(Shape::from(vec![1])), DType::F32) +} + +#[inline] +fn max_f32_contiguous(data: &[f32]) -> f32 { + #[cfg(feature = "simd")] + { + kernels::max_f32(data) + } + + #[cfg(not(feature = "simd"))] + { + data.iter() + .copied() + .reduce(|a, b| if a >= b { a } else { b }) + .expect("max: tensor must not be empty") + } +} + +fn min_f32_reduce(tensor: &FlexTensor) -> FlexTensor { + let result = match tensor.layout().contiguous_offsets() { + Some((start, end)) => { + let data: &[f32] = tensor.storage(); + min_f32_contiguous(&data[start..end]) + } + None => { + let data: &[f32] = tensor.storage(); + let elem_count = tensor.layout().num_elements(); + if data.len() == elem_count { + min_f32_contiguous(data) + } else { + StridedIter::new(tensor.layout()) + .map(|idx| data[idx]) + .reduce(|a, b| if a <= b { a } else { b }) + .expect("min: tensor must not be empty") + } + } + }; + + let bytes = Bytes::from_elems(vec![result]); + FlexTensor::new(bytes, Layout::contiguous(Shape::from(vec![1])), DType::F32) +} + +#[inline] +fn min_f32_contiguous(data: &[f32]) -> f32 { + #[cfg(feature = "simd")] + { + kernels::min_f32(data) + } + + #[cfg(not(feature = "simd"))] + { + data.iter() + .copied() + .reduce(|a, b| if a <= b { a } else { b }) + .expect("min: tensor must not be empty") + } +} + +fn max_impl(tensor: &FlexTensor) -> FlexTensor { + let result: E = match tensor.layout().contiguous_offsets() { + Some((start, end)) => { + let data: &[E] = tensor.storage(); + data[start..end] + .iter() + .copied() + .reduce(|a, b| if a >= b { a } else { b }) + .expect("max: tensor must not be empty") + } + None => { + let data: &[E] = tensor.storage(); + StridedIter::new(tensor.layout()) + .map(|idx| data[idx]) + .reduce(|a, b| if a >= b { a } else { b }) + .expect("max: tensor must not be empty") + } + }; + + let bytes = Bytes::from_elems(vec![result]); + FlexTensor::new( + bytes, + Layout::contiguous(Shape::from(vec![1])), + tensor.dtype(), + ) +} + +fn min_impl(tensor: &FlexTensor) -> FlexTensor { + let result: E = match tensor.layout().contiguous_offsets() { + Some((start, end)) => { + let data: &[E] = tensor.storage(); + data[start..end] + .iter() + .copied() + .reduce(|a, b| if a <= b { a } else { b }) + .expect("min: tensor must not be empty") + } + None => { + let data: &[E] = tensor.storage(); + StridedIter::new(tensor.layout()) + .map(|idx| data[idx]) + .reduce(|a, b| if a <= b { a } else { b }) + .expect("min: tensor must not be empty") + } + }; + + let bytes = Bytes::from_elems(vec![result]); + FlexTensor::new( + bytes, + Layout::contiguous(Shape::from(vec![1])), + tensor.dtype(), + ) +} + +// ============================================================================ +// Argmax / Argmin +// ============================================================================ + +/// Argmax along a dimension, returning indices as isize (INDEX_DTYPE). +pub fn argmax(tensor: FlexTensor, dim: usize) -> FlexTensor { + assert_dim_fits_isize(tensor.layout().shape()[dim], dim); + // f32 last-dim fast path: 2-pass SIMD for large rows, 1-pass scalar for small rows + if tensor.dtype() == DType::F32 && dim == tensor.layout().shape().num_dims() - 1 { + #[cfg(feature = "simd")] + if tensor.layout().shape()[dim] >= EXTREMUM_SIMD_ROW_THRESHOLD { + return extremum_indices_f32_last_simd(&tensor, dim, kernels::max_f32); + } + return extremum_indices_f32_last_scalar(&tensor, dim, |a, b| a > b); + } + match tensor.dtype() { + DType::F32 => { + extremum_dim_with_indices::(&tensor, dim, |a, b| { + !b.is_nan() && (a.is_nan() || a > b) + }) + .1 + } + DType::F64 => { + extremum_dim_with_indices::(&tensor, dim, |a, b| { + !b.is_nan() && (a.is_nan() || a > b) + }) + .1 + } + DType::F16 => { + extremum_dim_with_indices_half::( + &tensor, + dim, + |a, b| !b.is_nan() && (a.is_nan() || a > b), + f16::to_f32, + f16::from_f32, + ) + .1 + } + DType::BF16 => { + extremum_dim_with_indices_half::( + &tensor, + dim, + |a, b| !b.is_nan() && (a.is_nan() || a > b), + bf16::to_f32, + bf16::from_f32, + ) + .1 + } + DType::I8 => extremum_dim_with_indices::(&tensor, dim, |a, b| a > b).1, + DType::I16 => extremum_dim_with_indices::(&tensor, dim, |a, b| a > b).1, + DType::I32 => extremum_dim_with_indices::(&tensor, dim, |a, b| a > b).1, + DType::I64 => extremum_dim_with_indices::(&tensor, dim, |a, b| a > b).1, + _ => panic!("argmax: unsupported dtype {:?}", tensor.dtype()), + } +} + +/// Argmin along a dimension, returning indices as isize (INDEX_DTYPE). +pub fn argmin(tensor: FlexTensor, dim: usize) -> FlexTensor { + assert_dim_fits_isize(tensor.layout().shape()[dim], dim); + // f32 last-dim fast path: 2-pass SIMD for large rows, 1-pass scalar for small rows + if tensor.dtype() == DType::F32 && dim == tensor.layout().shape().num_dims() - 1 { + #[cfg(feature = "simd")] + if tensor.layout().shape()[dim] >= EXTREMUM_SIMD_ROW_THRESHOLD { + return extremum_indices_f32_last_simd(&tensor, dim, kernels::min_f32); + } + return extremum_indices_f32_last_scalar(&tensor, dim, |a, b| a < b); + } + match tensor.dtype() { + DType::F32 => { + extremum_dim_with_indices::(&tensor, dim, |a, b| { + !b.is_nan() && (a.is_nan() || a < b) + }) + .1 + } + DType::F64 => { + extremum_dim_with_indices::(&tensor, dim, |a, b| { + !b.is_nan() && (a.is_nan() || a < b) + }) + .1 + } + DType::F16 => { + extremum_dim_with_indices_half::( + &tensor, + dim, + |a, b| !b.is_nan() && (a.is_nan() || a < b), + f16::to_f32, + f16::from_f32, + ) + .1 + } + DType::BF16 => { + extremum_dim_with_indices_half::( + &tensor, + dim, + |a, b| !b.is_nan() && (a.is_nan() || a < b), + bf16::to_f32, + bf16::from_f32, + ) + .1 + } + DType::I8 => extremum_dim_with_indices::(&tensor, dim, |a, b| a < b).1, + DType::I16 => extremum_dim_with_indices::(&tensor, dim, |a, b| a < b).1, + DType::I32 => extremum_dim_with_indices::(&tensor, dim, |a, b| a < b).1, + DType::I64 => extremum_dim_with_indices::(&tensor, dim, |a, b| a < b).1, + _ => panic!("argmin: unsupported dtype {:?}", tensor.dtype()), + } +} + +// ============================================================================ +// Dimension reduction helpers +// ============================================================================ + +#[derive(Clone, Copy)] +enum ReduceOp { + Sum, + Prod, +} + +/// Optimized f32 dimension reduction with SIMD. +fn reduce_dim_f32(tensor: &FlexTensor, dim: usize, op: ReduceOp) -> FlexTensor { + let ndims = tensor.layout().shape().num_dims(); + + assert!( + dim < ndims, + "dim {} out of bounds for {} dimensions", + dim, + ndims + ); + + // Copy to contiguous only when the flattened stride assumption breaks: + // non-contiguous tensor with 2+ outer dims or 2+ inner dims. + let outer_dims = dim; + let inner_dims = ndims - dim - 1; + let needs_copy = !tensor.is_contiguous() && (outer_dims > 1 || inner_dims > 1); + let tensor = if needs_copy { + tensor.to_contiguous() + } else { + tensor.clone() + }; + let shape = tensor.layout().shape(); + let strides = tensor.layout().strides(); + + let dim_size = shape[dim]; + let mut out_shape: Vec = shape.to_vec(); + out_shape[dim] = 1; + let out_size: usize = out_shape.iter().product(); + + // Empty output: any zero-sized non-reduced dim means the result has no + // elements. Return early so the SIMD kernels and fallback loops never + // see `outer_size == 0` or `inner_size == 0`. + if out_size == 0 { + return FlexTensor::new( + Bytes::from_elems(Vec::::new()), + Layout::contiguous(Shape::from(out_shape)), + DType::F32, + ); + } + + let outer_size: usize = shape[..dim].iter().product(); + let inner_size: usize = shape[dim + 1..].iter().product(); + + let data: &[f32] = tensor.storage(); + let start_offset = tensor.layout().start_offset(); + let dim_stride = strides[dim]; + + let (init, reduce_fn): (f32, fn(f32, f32) -> f32) = match op { + ReduceOp::Sum => (0.0, |a, b| a + b), + ReduceOp::Prod => (1.0, |a, b| a * b), + }; + + // Check for negative strides (from flip operations) - fall back to general case + let has_negative_strides = strides.iter().any(|&s| s < 0); + + // Check if inner dimension is contiguous (stride = 1) and no negative strides + let inner_contiguous = !has_negative_strides && (dim + 1 >= ndims || strides[ndims - 1] == 1); + + let result: Vec = if inner_contiguous && dim == ndims - 1 && dim_stride == 1 { + // Reducing last dimension with contiguous data: use SIMD. + // `reduce_last_dim_f32` reads each row as `&data[start..start + dim_size]`, + // which only matches the logical row when the reduce dim itself has + // stride 1. Transposed views (e.g. shape [3,2] strides [1,3]) would + // otherwise read contiguous storage and return wrong sums. + reduce_last_dim_f32(data, start_offset, outer_size, dim_size, strides, dim, op) + } else if dim == 0 && inner_contiguous && matches!(op, ReduceOp::Sum) { + // First-dim reduction with contiguous inner: use cache-friendly accumulation + reduce_first_dim_f32(data, start_offset, dim_size, inner_size, dim_stride) + } else if dim > 0 && dim < ndims - 1 && inner_contiguous && matches!(op, ReduceOp::Sum) { + // Middle-dim reduction (e.g., [B, M, K] reducing dim=1): cache-friendly accumulation + let outer_stride = strides[dim - 1]; + reduce_middle_dim_f32( + data, + start_offset, + outer_size, + dim_size, + inner_size, + outer_stride, + dim_stride, + ) + } else if dim_stride == 1 && matches!(op, ReduceOp::Sum) && outer_size == 1 { + // Reduction dimension is contiguous, no outer batch (e.g., transposed 2D reducing dim=0) + // Storage is [inner_size rows of dim_size elements each] - use sum_rows_f32 + #[cfg(feature = "simd")] + { + let mut result = vec![0.0f32; inner_size]; + kernels::sum_rows_f32( + &data[start_offset..], + &mut result, + inner_size, // number of rows (output positions) + dim_size, // elements per row (to sum) + ); + result + } + #[cfg(not(feature = "simd"))] + { + let inner_stride: isize = if dim + 1 < ndims { strides[dim + 1] } else { 1 }; + let mut result = Vec::with_capacity(out_size); + for inner in 0..inner_size { + let base = (start_offset as isize + inner as isize * inner_stride) as usize; + let slice = &data[base..base + dim_size]; + result.push(slice.iter().copied().sum()); + } + result + } + } else if dim_stride == 1 && matches!(op, ReduceOp::Sum) { + // Reduction dimension is contiguous but with outer batches + let outer_stride: isize = if dim > 0 { strides[dim - 1] } else { 0 }; + let inner_stride: isize = if dim + 1 < ndims { strides[dim + 1] } else { 1 }; + + let mut result = Vec::with_capacity(out_size); + for outer in 0..outer_size { + for inner in 0..inner_size { + let base = (start_offset as isize + + outer as isize * outer_stride + + inner as isize * inner_stride) as usize; + let slice = &data[base..base + dim_size]; + #[cfg(feature = "simd")] + let acc = kernels::sum_f32(slice); + #[cfg(not(feature = "simd"))] + let acc = slice.iter().copied().sum(); + result.push(acc); + } + } + result + } else if tensor.is_contiguous() { + // Contiguous: use flat index arithmetic (safe for any ndims). + // outer_size and inner_size are guaranteed positive by the out_size == 0 + // early return above. + let mut result = Vec::with_capacity(out_size); + for outer in 0..outer_size { + for inner in 0..inner_size { + let mut acc = init; + for d in 0..dim_size { + let idx = start_offset + outer * dim_size * inner_size + d * inner_size + inner; + acc = reduce_fn(acc, data[idx]); + } + result.push(acc); + } + } + result + } else { + // Non-contiguous with at most 1 outer + 1 inner dim (e.g., flipped 2D) + let outer_stride: isize = if dim > 0 { strides[dim - 1] } else { 0 }; + let inner_stride: isize = if dim + 1 < ndims { strides[dim + 1] } else { 1 }; + + let mut result = Vec::with_capacity(out_size); + for outer in 0..outer_size { + for inner in 0..inner_size { + let base = start_offset as isize + + outer as isize * outer_stride + + inner as isize * inner_stride; + let mut acc = init; + for d in 0..dim_size { + let idx = (base + d as isize * dim_stride) as usize; + acc = reduce_fn(acc, data[idx]); + } + result.push(acc); + } + } + result + }; + + let bytes = Bytes::from_elems(result); + FlexTensor::new( + bytes, + Layout::contiguous(Shape::from(out_shape)), + DType::F32, + ) +} + +/// Reduce middle dimension (e.g., [B, M, K] reducing dim=1) with cache-friendly iteration. +/// For each batch, iterate over rows (dim to reduce) sequentially and accumulate into columns. +#[inline] +fn reduce_middle_dim_f32( + data: &[f32], + start_offset: usize, + outer_size: usize, // batch size + dim_size: usize, // rows to sum + inner_size: usize, // columns (output per batch) + outer_stride: isize, + dim_stride: isize, +) -> Vec { + let out_size = outer_size * inner_size; + + #[cfg(feature = "simd")] + { + // Use aligned allocation for optimal SIMD scatter-add + let mut result = aligned::alloc_aligned_zeroed::(out_size); + kernels::scatter_add_batched( + &data[start_offset..], + &mut result, + outer_size, + dim_size, + inner_size, + outer_stride as usize, + dim_stride as usize, + ); + aligned::to_vec(result) + } + + #[cfg(not(feature = "simd"))] + { + let mut result = vec![0.0f32; out_size]; + let start = start_offset as isize; + for batch in 0..outer_size { + let batch_start = (start + batch as isize * outer_stride) as usize; + let out_batch_start = batch * inner_size; + + for row in 0..dim_size { + let row_start = (batch_start as isize + row as isize * dim_stride) as usize; + for c in 0..inner_size { + result[out_batch_start + c] += data[row_start + c]; + } + } + } + result + } +} + +/// Reduce first dimension with cache-friendly row iteration. +/// Instead of iterating per-output (col) and gathering from rows (cache-unfriendly), +/// iterate over rows (sequential access) and scatter-accumulate into outputs. +#[inline] +fn reduce_first_dim_f32( + data: &[f32], + start_offset: usize, + dim_size: usize, // number of rows to sum + inner_size: usize, // number of columns (output positions) + dim_stride: isize, // stride between rows +) -> Vec { + #[cfg(feature = "simd")] + { + // Use aligned allocation for optimal SIMD scatter-add + let mut result = aligned::alloc_aligned_zeroed::(inner_size); + kernels::scatter_add_f32( + &data[start_offset..], + &mut result, + dim_size, + inner_size, + dim_stride as usize, + ); + aligned::to_vec(result) + } + + #[cfg(not(feature = "simd"))] + { + let mut result = vec![0.0f32; inner_size]; + let start = start_offset as isize; + for row in 0..dim_size { + let row_start = (start + row as isize * dim_stride) as usize; + for c in 0..inner_size { + result[c] += data[row_start + c]; + } + } + result + } +} + +/// Reduce last dimension with SIMD. +/// +/// For contiguous Sum: batches all rows in a single kernel call using +/// 4-accumulator SIMD to hide add latency. +#[inline] +fn reduce_last_dim_f32( + data: &[f32], + start_offset: usize, + outer_size: usize, + dim_size: usize, + strides: &[isize], + dim: usize, + op: ReduceOp, +) -> Vec { + let outer_stride: isize = if dim > 0 { + strides[dim - 1] + } else { + dim_size as isize + }; + + // `outer_size > 0` is guaranteed by the out_size == 0 early return in + // `reduce_dim_f32`. + let rows = outer_size; + + // Contiguous Sum: batch all rows in one kernel call to avoid per-row overhead. + #[cfg(feature = "simd")] + if matches!(op, ReduceOp::Sum) && outer_stride == dim_size as isize { + let mut result = vec![0.0f32; rows]; + kernels::sum_rows_f32(&data[start_offset..], &mut result, rows, dim_size); + return result; + } + + // Fallback: non-contiguous strides or Prod. + let mut result = Vec::with_capacity(rows); + for outer in 0..rows { + let row_start = (start_offset as isize + outer as isize * outer_stride) as usize; + let row = &data[row_start..row_start + dim_size]; + + let val = match op { + ReduceOp::Sum => { + #[cfg(feature = "simd")] + { + kernels::sum_f32(row) + } + #[cfg(not(feature = "simd"))] + { + row.iter().copied().sum() + } + } + ReduceOp::Prod => row.iter().copied().product(), + }; + result.push(val); + } + result +} + +/// Generic dimension reduction implementation. +fn reduce_dim_impl(tensor: &FlexTensor, dim: usize, init: E, reduce_fn: F) -> FlexTensor +where + E: Element + bytemuck::Pod + Copy, + F: Fn(E, E) -> E, +{ + let ndims = tensor.layout().shape().num_dims(); + assert!( + dim < ndims, + "dim {} out of bounds for {} dimensions", + dim, + ndims + ); + + // Copy to contiguous only when the flattened stride assumption breaks: + // non-contiguous tensor with 2+ outer dims or 2+ inner dims. + let outer_dims = dim; + let inner_dims = ndims - dim - 1; + let needs_copy = !tensor.is_contiguous() && (outer_dims > 1 || inner_dims > 1); + let tensor = if needs_copy { + tensor.to_contiguous() + } else { + tensor.clone() + }; + let shape = tensor.layout().shape(); + let strides = tensor.layout().strides(); + + let dim_size = shape[dim]; + let mut out_shape: Vec = shape.to_vec(); + out_shape[dim] = 1; + let out_size: usize = out_shape.iter().product(); + + // Empty output: any zero-sized non-reduced dim means the result has no + // elements. Returning early keeps the loops below from producing phantom + // outputs when `outer_size == 0` or `inner_size == 0`. + if out_size == 0 { + return FlexTensor::new( + Bytes::from_elems(Vec::::new()), + Layout::contiguous(Shape::from(out_shape)), + tensor.dtype(), + ); + } + + let outer_size: usize = shape[..dim].iter().product(); + let inner_size: usize = shape[dim + 1..].iter().product(); + + let data: &[E] = tensor.storage(); + let start_offset = tensor.layout().start_offset(); + + let mut result: Vec = Vec::with_capacity(out_size); + + if tensor.is_contiguous() { + // Contiguous: use flat index arithmetic (safe for any ndims). + // outer_size and inner_size are positive by the early return above. + for outer in 0..outer_size { + for inner in 0..inner_size { + let mut acc = init; + for d in 0..dim_size { + let idx = start_offset + outer * dim_size * inner_size + d * inner_size + inner; + acc = reduce_fn(acc, data[idx]); + } + result.push(acc); + } + } + } else { + // Non-contiguous with at most 1 outer + 1 inner dim + let dim_stride = strides[dim]; + let outer_stride: isize = if dim > 0 { strides[dim - 1] } else { 0 }; + let inner_stride: isize = if dim + 1 < ndims { strides[dim + 1] } else { 1 }; + + for outer in 0..outer_size { + for inner in 0..inner_size { + let base = start_offset as isize + + outer as isize * outer_stride + + inner as isize * inner_stride; + let mut acc = init; + for d in 0..dim_size { + let idx = (base + d as isize * dim_stride) as usize; + acc = reduce_fn(acc, data[idx]); + } + result.push(acc); + } + } + } + + let bytes = Bytes::from_elems(result); + FlexTensor::new( + bytes, + Layout::contiguous(Shape::from(out_shape)), + tensor.dtype(), + ) +} + +/// Widening dimension reduction for small integer types: accumulate in i64 to avoid overflow. +fn reduce_dim_widening(tensor: &FlexTensor, dim: usize, init: i64, reduce_fn: F) -> FlexTensor +where + E: Element + bytemuck::Pod, + i64: From, + F: Fn(i64, i64) -> i64, +{ + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape(); + let ndims = shape.num_dims(); + + assert!( + dim < ndims, + "dim {} out of bounds for {} dimensions", + dim, + ndims + ); + + let dim_size = shape[dim]; + let mut out_shape: Vec = shape.to_vec(); + out_shape[dim] = 1; + let out_size: usize = out_shape.iter().product(); + + // Empty output: skip the loop entirely for zero-sized non-reduced dims. + if out_size == 0 { + return FlexTensor::new( + Bytes::from_elems(Vec::::new()), + Layout::contiguous(Shape::from(out_shape)), + tensor.dtype(), + ); + } + + let outer_size: usize = shape[..dim].iter().product(); + let inner_size: usize = shape[dim + 1..].iter().product(); + + let data: &[E] = tensor.storage(); + let start_offset = tensor.layout().start_offset(); + + let mut result: Vec = Vec::with_capacity(out_size); + + for outer in 0..outer_size { + for inner in 0..inner_size { + let mut acc = init; + for d in 0..dim_size { + let idx = start_offset + outer * dim_size * inner_size + d * inner_size + inner; + acc = reduce_fn(acc, i64::from(data[idx])); + } + // Truncate back to target type (wrapping, matches PyTorch) + let val: E = truncate_i64_to_pod(acc); + result.push(val); + } + } + + let bytes = Bytes::from_elems(result); + FlexTensor::new( + bytes, + Layout::contiguous(Shape::from(out_shape)), + tensor.dtype(), + ) +} + +/// Half-precision dimension reduction with f32 accumulation. +/// +/// Works for both f16 and bf16 via the `to_f32`/`from_f32` closures. +fn reduce_dim_half( + tensor: &FlexTensor, + dim: usize, + init: f32, + reduce_fn: F, + to_f32: fn(E) -> f32, + from_f32: fn(f32) -> E, +) -> FlexTensor +where + E: Element + bytemuck::Pod, + F: Fn(f32, f32) -> f32, +{ + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape(); + let ndims = shape.num_dims(); + + assert!( + dim < ndims, + "dim {} out of bounds for {} dimensions", + dim, + ndims + ); + + let dim_size = shape[dim]; + let mut out_shape: Vec = shape.to_vec(); + out_shape[dim] = 1; + let out_size: usize = out_shape.iter().product(); + + // Empty output: skip the loop entirely for zero-sized non-reduced dims. + if out_size == 0 { + return FlexTensor::new( + Bytes::from_elems(Vec::::new()), + Layout::contiguous(Shape::from(out_shape)), + E::dtype(), + ); + } + + let outer_size: usize = shape[..dim].iter().product(); + let inner_size: usize = shape[dim + 1..].iter().product(); + + let data: &[E] = tensor.storage(); + let start_offset = tensor.layout().start_offset(); + + let mut result: Vec = Vec::with_capacity(out_size); + + for outer in 0..outer_size { + for inner in 0..inner_size { + let mut acc = init; + for d in 0..dim_size { + let idx = start_offset + outer * dim_size * inner_size + d * inner_size + inner; + acc = reduce_fn(acc, to_f32(data[idx])); + } + result.push(from_f32(acc)); + } + } + + let bytes = Bytes::from_elems(result); + FlexTensor::new( + bytes, + Layout::contiguous(Shape::from(out_shape)), + E::dtype(), + ) +} + +/// Sum along `dim` for an already-contiguous row-major f32 slice, producing +/// one output per (outer, inner) position. +/// +/// The caller owns the contiguity guarantee: `data.len()` must equal +/// `outer_size * dim_size * inner_size` in logical row-major order. Dispatches +/// to the same SIMD kernels `reduce_dim_f32` uses, but without the stride +/// bookkeeping. +fn sum_dim_contiguous_f32( + data: &[f32], + outer_size: usize, + dim_size: usize, + inner_size: usize, +) -> Vec { + // Empty output: if any non-reduced dim has size 0, the result is empty. + // Returning early avoids indexing past an empty `data` slice in the SIMD + // kernels and keeps the output length in sync with the caller's out_shape. + if outer_size == 0 || inner_size == 0 { + return Vec::new(); + } + + // Last-dim: each output is the sum of a contiguous run of dim_size elements. + if inner_size == 1 { + let rows = outer_size; + #[cfg(feature = "simd")] + { + let mut result = vec![0.0f32; rows]; + kernels::sum_rows_f32(data, &mut result, rows, dim_size); + return result; + } + #[cfg(not(feature = "simd"))] + { + return (0..rows) + .map(|i| data[i * dim_size..(i + 1) * dim_size].iter().sum()) + .collect(); + } + } + + // First-dim (or equivalent: any collapsed-outer case): scatter-add dim_size + // rows of inner_size cols into a single inner_size accumulator. + if outer_size == 1 { + #[cfg(feature = "simd")] + { + let mut result = aligned::alloc_aligned_zeroed::(inner_size); + kernels::scatter_add_f32(data, &mut result, dim_size, inner_size, inner_size); + return aligned::to_vec(result); + } + #[cfg(not(feature = "simd"))] + { + let mut result = vec![0.0f32; inner_size]; + for row in 0..dim_size { + let row_start = row * inner_size; + for c in 0..inner_size { + result[c] += data[row_start + c]; + } + } + return result; + } + } + + // Middle-dim: batched scatter-add. For each outer batch, sum dim_size rows + // of inner_size cols into a per-batch accumulator. + let out_size = outer_size * inner_size; + #[cfg(feature = "simd")] + { + let mut result = aligned::alloc_aligned_zeroed::(out_size); + kernels::scatter_add_batched( + data, + &mut result, + outer_size, + dim_size, + inner_size, + dim_size * inner_size, + inner_size, + ); + aligned::to_vec(result) + } + #[cfg(not(feature = "simd"))] + { + let mut result = vec![0.0f32; out_size]; + for outer in 0..outer_size { + let out_base = outer * inner_size; + for d in 0..dim_size { + let in_base = outer * dim_size * inner_size + d * inner_size; + for c in 0..inner_size { + result[out_base + c] += data[in_base + c]; + } + } + } + result + } +} + +/// Mean along a dimension for half-precision types, fusing sum and divide in f32. +/// +/// A naive `sum_dim` + `scalar_div` implementation can overflow to +inf when the +/// intermediate sum exceeds `f16::MAX` (65504), even if the final mean fits. This +/// function keeps the entire reduction and division in f32 and only narrows to +/// f16/bf16 on store via `E::from_elem`. +fn mean_dim_half(tensor: &FlexTensor, dim: usize) -> FlexTensor +where + E: Element + bytemuck::Pod, +{ + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape(); + let ndims = shape.num_dims(); + + assert!( + dim < ndims, + "dim {} out of bounds for {} dimensions", + dim, + ndims + ); + + let dim_size = shape[dim]; + assert!( + dim_size > 0, + "mean_dim: cannot take mean of empty dimension" + ); + let mut out_shape: Vec = shape.to_vec(); + out_shape[dim] = 1; + + let outer_size: usize = shape[..dim].iter().product(); + let inner_size: usize = shape[dim + 1..].iter().product(); + + let data = float_storage_as_f32(&tensor); + let divisor = dim_size as f32; + + let sums = sum_dim_contiguous_f32(&data, outer_size, dim_size, inner_size); + let result: Vec = sums + .into_iter() + .map(|s| E::from_elem(s / divisor)) + .collect(); + + let bytes = Bytes::from_elems(result); + FlexTensor::new( + bytes, + Layout::contiguous(Shape::from(out_shape)), + E::dtype(), + ) +} + +/// Scalar mean for half-precision types, fusing sum and divide in f32. +/// Avoids f16 overflow when the total sum exceeds `f16::MAX`. Empty input +/// produces NaN to match the f32/f64 path in `mean()`. +fn mean_scalar_half(tensor: &FlexTensor) -> FlexTensor +where + E: Element + bytemuck::Pod, +{ + let tensor = tensor.to_contiguous(); + let n = tensor.layout().num_elements(); + let data = float_storage_as_f32(&tensor); + // Route through `sum_f32_contiguous` to pick up SIMD + rayon for the + // f32 reduction. The half-precision narrowing happens after the divide. + let acc = sum_f32_contiguous(&data); + + let mean = acc / (n as f32); + let bytes = Bytes::from_elems(vec![E::from_elem(mean)]); + FlexTensor::new(bytes, Layout::contiguous(Shape::from(vec![1])), E::dtype()) +} + +// ============================================================================ +// Mean (all elements) +// ============================================================================ + +/// Mean of all elements, returning a scalar tensor. +pub fn mean(tensor: FlexTensor) -> FlexTensor { + let dtype = tensor.dtype(); + + // Half-precision types fuse sum+divide in f32 to avoid overflow when the + // total sum exceeds f16::MAX. + match dtype { + DType::F16 => return mean_scalar_half::(&tensor), + DType::BF16 => return mean_scalar_half::(&tensor), + _ => {} + } + + let n = tensor.layout().num_elements(); + let sum_result = sum(tensor); + match dtype { + DType::F32 => scalar_div::(sum_result, n as f32), + DType::F64 => scalar_div::(sum_result, n as f64), + _ => panic!("mean: unsupported dtype {:?}", dtype), + } +} + +// ============================================================================ +// Max/Min along dimension (value + optional indices in a single pass) +// ============================================================================ + +/// Max along a dimension, returning only values. +pub fn max_dim(tensor: FlexTensor, dim: usize) -> FlexTensor { + assert!( + tensor.layout().shape()[dim] > 0, + "max_dim: dimension {dim} has size 0" + ); + if tensor.dtype() == DType::F32 && dim == tensor.layout().shape().num_dims() - 1 { + #[cfg(feature = "simd")] + if tensor.layout().shape()[dim] >= EXTREMUM_SIMD_ROW_THRESHOLD { + return extremum_dim_f32_last_simd(&tensor, dim, kernels::max_f32); + } + return extremum_f32_last_scalar(&tensor, dim, |a, b| a > b); + } + match tensor.dtype() { + DType::F32 => { + extremum_dim::(&tensor, dim, |a, b| !b.is_nan() && (a.is_nan() || a > b)) + } + DType::F64 => { + extremum_dim::(&tensor, dim, |a, b| !b.is_nan() && (a.is_nan() || a > b)) + } + DType::F16 => extremum_dim_half::( + &tensor, + dim, + |a, b| !b.is_nan() && (a.is_nan() || a > b), + f16::to_f32, + f16::from_f32, + ), + DType::BF16 => extremum_dim_half::( + &tensor, + dim, + |a, b| !b.is_nan() && (a.is_nan() || a > b), + bf16::to_f32, + bf16::from_f32, + ), + DType::I64 => extremum_dim::(&tensor, dim, |a, b| a > b), + DType::I32 => extremum_dim::(&tensor, dim, |a, b| a > b), + DType::I16 => extremum_dim::(&tensor, dim, |a, b| a > b), + DType::I8 => extremum_dim::(&tensor, dim, |a, b| a > b), + DType::U64 => extremum_dim::(&tensor, dim, |a, b| a > b), + DType::U32 => extremum_dim::(&tensor, dim, |a, b| a > b), + DType::U16 => extremum_dim::(&tensor, dim, |a, b| a > b), + DType::U8 => extremum_dim::(&tensor, dim, |a, b| a > b), + _ => panic!("max_dim: unsupported dtype {:?}", tensor.dtype()), + } +} + +/// Min along a dimension, returning only values. +pub fn min_dim(tensor: FlexTensor, dim: usize) -> FlexTensor { + assert!( + tensor.layout().shape()[dim] > 0, + "min_dim: dimension {dim} has size 0" + ); + if tensor.dtype() == DType::F32 && dim == tensor.layout().shape().num_dims() - 1 { + #[cfg(feature = "simd")] + if tensor.layout().shape()[dim] >= EXTREMUM_SIMD_ROW_THRESHOLD { + return extremum_dim_f32_last_simd(&tensor, dim, kernels::min_f32); + } + return extremum_f32_last_scalar(&tensor, dim, |a, b| a < b); + } + match tensor.dtype() { + DType::F32 => { + extremum_dim::(&tensor, dim, |a, b| !b.is_nan() && (a.is_nan() || a < b)) + } + DType::F64 => { + extremum_dim::(&tensor, dim, |a, b| !b.is_nan() && (a.is_nan() || a < b)) + } + DType::F16 => extremum_dim_half::( + &tensor, + dim, + |a, b| !b.is_nan() && (a.is_nan() || a < b), + f16::to_f32, + f16::from_f32, + ), + DType::BF16 => extremum_dim_half::( + &tensor, + dim, + |a, b| !b.is_nan() && (a.is_nan() || a < b), + bf16::to_f32, + bf16::from_f32, + ), + DType::I64 => extremum_dim::(&tensor, dim, |a, b| a < b), + DType::I32 => extremum_dim::(&tensor, dim, |a, b| a < b), + DType::I16 => extremum_dim::(&tensor, dim, |a, b| a < b), + DType::I8 => extremum_dim::(&tensor, dim, |a, b| a < b), + DType::U64 => extremum_dim::(&tensor, dim, |a, b| a < b), + DType::U32 => extremum_dim::(&tensor, dim, |a, b| a < b), + DType::U16 => extremum_dim::(&tensor, dim, |a, b| a < b), + DType::U8 => extremum_dim::(&tensor, dim, |a, b| a < b), + _ => panic!("min_dim: unsupported dtype {:?}", tensor.dtype()), + } +} + +/// Max along a dimension with indices, returning (values, indices) in a single pass. +pub fn max_dim_with_indices(tensor: FlexTensor, dim: usize) -> (FlexTensor, FlexTensor) { + let dim_len = tensor.layout().shape()[dim]; + assert!( + dim_len > 0, + "max_dim_with_indices: dimension {dim} has size 0" + ); + assert_dim_fits_isize(dim_len, dim); + if tensor.dtype() == DType::F32 && dim == tensor.layout().shape().num_dims() - 1 { + #[cfg(feature = "simd")] + if tensor.layout().shape()[dim] >= EXTREMUM_SIMD_ROW_THRESHOLD { + return extremum_dim_with_indices_f32_last_simd(&tensor, dim, kernels::max_f32); + } + return extremum_with_indices_f32_last_scalar(&tensor, dim, |a, b| a > b); + } + match tensor.dtype() { + DType::F32 => extremum_dim_with_indices::(&tensor, dim, |a, b| { + !b.is_nan() && (a.is_nan() || a > b) + }), + DType::F64 => extremum_dim_with_indices::(&tensor, dim, |a, b| { + !b.is_nan() && (a.is_nan() || a > b) + }), + DType::F16 => extremum_dim_with_indices_half::( + &tensor, + dim, + |a, b| !b.is_nan() && (a.is_nan() || a > b), + f16::to_f32, + f16::from_f32, + ), + DType::BF16 => extremum_dim_with_indices_half::( + &tensor, + dim, + |a, b| !b.is_nan() && (a.is_nan() || a > b), + bf16::to_f32, + bf16::from_f32, + ), + DType::I64 => extremum_dim_with_indices::(&tensor, dim, |a, b| a > b), + DType::I32 => extremum_dim_with_indices::(&tensor, dim, |a, b| a > b), + DType::I16 => extremum_dim_with_indices::(&tensor, dim, |a, b| a > b), + DType::I8 => extremum_dim_with_indices::(&tensor, dim, |a, b| a > b), + DType::U64 => extremum_dim_with_indices::(&tensor, dim, |a, b| a > b), + DType::U32 => extremum_dim_with_indices::(&tensor, dim, |a, b| a > b), + DType::U16 => extremum_dim_with_indices::(&tensor, dim, |a, b| a > b), + DType::U8 => extremum_dim_with_indices::(&tensor, dim, |a, b| a > b), + _ => panic!( + "max_dim_with_indices: unsupported dtype {:?}", + tensor.dtype() + ), + } +} + +/// Min along a dimension with indices, returning (values, indices) in a single pass. +pub fn min_dim_with_indices(tensor: FlexTensor, dim: usize) -> (FlexTensor, FlexTensor) { + let dim_len = tensor.layout().shape()[dim]; + assert!( + dim_len > 0, + "min_dim_with_indices: dimension {dim} has size 0" + ); + assert_dim_fits_isize(dim_len, dim); + if tensor.dtype() == DType::F32 && dim == tensor.layout().shape().num_dims() - 1 { + #[cfg(feature = "simd")] + if tensor.layout().shape()[dim] >= EXTREMUM_SIMD_ROW_THRESHOLD { + return extremum_dim_with_indices_f32_last_simd(&tensor, dim, kernels::min_f32); + } + return extremum_with_indices_f32_last_scalar(&tensor, dim, |a, b| a < b); + } + match tensor.dtype() { + DType::F32 => extremum_dim_with_indices::(&tensor, dim, |a, b| { + !b.is_nan() && (a.is_nan() || a < b) + }), + DType::F64 => extremum_dim_with_indices::(&tensor, dim, |a, b| { + !b.is_nan() && (a.is_nan() || a < b) + }), + DType::F16 => extremum_dim_with_indices_half::( + &tensor, + dim, + |a, b| !b.is_nan() && (a.is_nan() || a < b), + f16::to_f32, + f16::from_f32, + ), + DType::BF16 => extremum_dim_with_indices_half::( + &tensor, + dim, + |a, b| !b.is_nan() && (a.is_nan() || a < b), + bf16::to_f32, + bf16::from_f32, + ), + DType::I64 => extremum_dim_with_indices::(&tensor, dim, |a, b| a < b), + DType::I32 => extremum_dim_with_indices::(&tensor, dim, |a, b| a < b), + DType::I16 => extremum_dim_with_indices::(&tensor, dim, |a, b| a < b), + DType::I8 => extremum_dim_with_indices::(&tensor, dim, |a, b| a < b), + DType::U64 => extremum_dim_with_indices::(&tensor, dim, |a, b| a < b), + DType::U32 => extremum_dim_with_indices::(&tensor, dim, |a, b| a < b), + DType::U16 => extremum_dim_with_indices::(&tensor, dim, |a, b| a < b), + DType::U8 => extremum_dim_with_indices::(&tensor, dim, |a, b| a < b), + _ => panic!( + "min_dim_with_indices: unsupported dtype {:?}", + tensor.dtype() + ), + } +} + +// ============================================================================ +// Extremum helpers (SIMD fast paths + generic scalar, parallelized with rayon) +// ============================================================================ + +// Lower threshold than the global PARALLEL_THRESHOLD (256K) because the per-element +// work (a single comparison + conditional store) is cheap enough that rayon overhead +// is amortized at smaller sizes. 32K elements * ~1.5ns/elem = ~48µs of serial work, +// enough to justify thread-pool dispatch. +#[cfg(feature = "rayon")] +const EXTREMUM_PARALLEL_THRESHOLD: usize = 32 * 1024; + +/// Minimum row length for the 2-pass SIMD extremum path (reduce + scan). +/// Below this, single-pass scalar is faster since it reads each element once. +#[cfg(feature = "simd")] +const EXTREMUM_SIMD_ROW_THRESHOLD: usize = 512; + +/// Scalar single-pass f32 last-dim extremum (values only). +/// Avoids the generic closure path by operating directly on contiguous f32 rows. +fn extremum_f32_last_scalar(tensor: &FlexTensor, dim: usize, is_better: F) -> FlexTensor +where + F: Fn(f32, f32) -> bool + Send + Sync, +{ + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape(); + let dim_size = shape[dim]; + let outer_size: usize = shape[..dim].iter().product(); + let data: &[f32] = tensor.storage(); + let start = tensor.layout().start_offset(); + + let mut out_shape: Vec = shape.to_vec(); + out_shape[dim] = 1; + + let reduce_row = |outer: usize| -> f32 { + let row_start = start + outer * dim_size; + let row = &data[row_start..row_start + dim_size]; + // Single left-to-right scan: first NaN wins, otherwise track the + // extremum. Starting `best = row[0]` means the i=0 `is_better` + // branch is a no-op for strict comparisons. + let mut best = row[0]; + for &v in row { + if v.is_nan() { + return f32::NAN; + } + if is_better(v, best) { + best = v; + } + } + best + }; + + #[cfg(feature = "rayon")] + let values: Vec = if outer_size * dim_size >= EXTREMUM_PARALLEL_THRESHOLD { + (0..outer_size).into_par_iter().map(&reduce_row).collect() + } else { + (0..outer_size).map(reduce_row).collect() + }; + + #[cfg(not(feature = "rayon"))] + let values: Vec = (0..outer_size).map(reduce_row).collect(); + + FlexTensor::new( + Bytes::from_elems(values), + Layout::contiguous(Shape::from(out_shape)), + DType::F32, + ) +} + +/// Scalar single-pass f32 last-dim argmax/argmin (indices only). +/// Uses direct pointer access and simple comparisons instead of the generic closure path. +fn extremum_indices_f32_last_scalar(tensor: &FlexTensor, dim: usize, is_better: F) -> FlexTensor +where + F: Fn(f32, f32) -> bool + Send + Sync, +{ + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape(); + let dim_size = shape[dim]; + let outer_size: usize = shape[..dim].iter().product(); + let data: &[f32] = tensor.storage(); + let start = tensor.layout().start_offset(); + + let mut out_shape: Vec = shape.to_vec(); + out_shape[dim] = 1; + + let find_row = |outer: usize| -> isize { + let row_start = start + outer * dim_size; + let row = &data[row_start..row_start + dim_size]; + // Single left-to-right scan: first NaN wins, otherwise track the + // extremum. Starting `best = row[0]` means the i=0 `is_better` + // branch is a no-op (strict comparisons are false on equal + // operands), so we don't need a separate row[0] check. + let mut best = row[0]; + let mut best_idx: isize = 0; + for (i, &v) in row.iter().enumerate() { + if v.is_nan() { + return i as isize; + } + if is_better(v, best) { + best = v; + best_idx = i as isize; + } + } + best_idx + }; + + #[cfg(feature = "rayon")] + let indices: Vec = if outer_size * dim_size >= EXTREMUM_PARALLEL_THRESHOLD { + (0..outer_size).into_par_iter().map(find_row).collect() + } else { + (0..outer_size).map(find_row).collect() + }; + + #[cfg(not(feature = "rayon"))] + let indices: Vec = (0..outer_size).map(find_row).collect(); + + FlexTensor::new( + Bytes::from_elems(indices), + Layout::contiguous(Shape::from(out_shape)), + INDEX_DTYPE, + ) +} + +/// Scalar single-pass f32 last-dim extremum with indices (values + indices). +fn extremum_with_indices_f32_last_scalar( + tensor: &FlexTensor, + dim: usize, + is_better: F, +) -> (FlexTensor, FlexTensor) +where + F: Fn(f32, f32) -> bool + Send + Sync, +{ + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape(); + let dim_size = shape[dim]; + let outer_size: usize = shape[..dim].iter().product(); + let data: &[f32] = tensor.storage(); + let start = tensor.layout().start_offset(); + + let mut out_shape: Vec = shape.to_vec(); + out_shape[dim] = 1; + + let find_row = |outer: usize| -> (f32, isize) { + let row_start = start + outer * dim_size; + let row = &data[row_start..row_start + dim_size]; + let mut best = row[0]; + let mut best_idx: isize = 0; + for (i, &v) in row.iter().enumerate() { + if v.is_nan() { + return (f32::NAN, i as isize); + } + if is_better(v, best) { + best = v; + best_idx = i as isize; + } + } + (best, best_idx) + }; + + #[cfg(feature = "rayon")] + let (values, indices): (Vec, Vec) = + if outer_size * dim_size >= EXTREMUM_PARALLEL_THRESHOLD { + (0..outer_size).into_par_iter().map(find_row).unzip() + } else { + (0..outer_size).map(find_row).unzip() + }; + + #[cfg(not(feature = "rayon"))] + let (values, indices): (Vec, Vec) = (0..outer_size).map(find_row).unzip(); + + let val_tensor = FlexTensor::new( + Bytes::from_elems(values), + Layout::contiguous(Shape::from(out_shape.clone())), + DType::F32, + ); + let idx_tensor = FlexTensor::new( + Bytes::from_elems(indices), + Layout::contiguous(Shape::from(out_shape)), + INDEX_DTYPE, + ); + (val_tensor, idx_tensor) +} + +/// Uses macerator SIMD reduction per contiguous row, with NaN propagation. +#[cfg(feature = "simd")] +fn extremum_dim_f32_last_simd( + tensor: &FlexTensor, + dim: usize, + simd_reduce: fn(&[f32]) -> f32, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape(); + let dim_size = shape[dim]; + let outer_size: usize = shape[..dim].iter().product(); + let data: &[f32] = tensor.storage(); + let start = tensor.layout().start_offset(); + + let mut out_shape: Vec = shape.to_vec(); + out_shape[dim] = 1; + + let reduce_row = |outer: usize| -> f32 { + let row_start = start + outer * dim_size; + let row = &data[row_start..row_start + dim_size]; + let ext = simd_reduce(row); + // SIMD max/min may silently drop NaN (architecture-dependent). + // If the result is already NaN, we're done. Otherwise, scan to + // check for any NaN the SIMD op missed. + if ext.is_nan() { + return f32::NAN; + } + for &v in row { + if v.is_nan() { + return f32::NAN; + } + } + ext + }; + + #[cfg(feature = "rayon")] + let values: Vec = if outer_size * dim_size >= EXTREMUM_PARALLEL_THRESHOLD { + (0..outer_size).into_par_iter().map(reduce_row).collect() + } else { + (0..outer_size).map(reduce_row).collect() + }; + + #[cfg(not(feature = "rayon"))] + let values: Vec = (0..outer_size).map(reduce_row).collect(); + + FlexTensor::new( + Bytes::from_elems(values), + Layout::contiguous(Shape::from(out_shape)), + DType::F32, + ) +} + +/// SIMD fast path for f32 last-dim extremum with indices. +/// Per row: SIMD reduction finds the extremum value, then a linear scan +/// locates the first NaN or first matching index. +#[cfg(feature = "simd")] +fn extremum_dim_with_indices_f32_last_simd( + tensor: &FlexTensor, + dim: usize, + simd_reduce: fn(&[f32]) -> f32, +) -> (FlexTensor, FlexTensor) { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape(); + let dim_size = shape[dim]; + let outer_size: usize = shape[..dim].iter().product(); + let data: &[f32] = tensor.storage(); + let start = tensor.layout().start_offset(); + + let mut out_shape: Vec = shape.to_vec(); + out_shape[dim] = 1; + + let find_row = |outer: usize| -> (f32, isize) { + let row_start = start + outer * dim_size; + let row = &data[row_start..row_start + dim_size]; + let ext = simd_reduce(row); + // Single scan: return first NaN (with NaN value) or first match of ext. + for (i, &v) in row.iter().enumerate() { + if v.is_nan() { + return (f32::NAN, i as isize); + } + if v == ext { + return (ext, i as isize); + } + } + (ext, 0) + }; + + #[cfg(feature = "rayon")] + let (values, indices): (Vec, Vec) = + if outer_size * dim_size >= EXTREMUM_PARALLEL_THRESHOLD { + (0..outer_size).into_par_iter().map(find_row).unzip() + } else { + (0..outer_size).map(find_row).unzip() + }; + + #[cfg(not(feature = "rayon"))] + let (values, indices): (Vec, Vec) = (0..outer_size).map(find_row).unzip(); + + let val_tensor = FlexTensor::new( + Bytes::from_elems(values), + Layout::contiguous(Shape::from(out_shape.clone())), + DType::F32, + ); + let idx_tensor = FlexTensor::new( + Bytes::from_elems(indices), + Layout::contiguous(Shape::from(out_shape)), + INDEX_DTYPE, + ); + (val_tensor, idx_tensor) +} + +/// SIMD fast path for f32 last-dim argmax/argmin (indices only, no values allocation). +#[cfg(feature = "simd")] +fn extremum_indices_f32_last_simd( + tensor: &FlexTensor, + dim: usize, + simd_reduce: fn(&[f32]) -> f32, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape(); + let dim_size = shape[dim]; + let outer_size: usize = shape[..dim].iter().product(); + let data: &[f32] = tensor.storage(); + let start = tensor.layout().start_offset(); + + let mut out_shape: Vec = shape.to_vec(); + out_shape[dim] = 1; + + let find_row = |outer: usize| -> isize { + let row_start = start + outer * dim_size; + let row = &data[row_start..row_start + dim_size]; + let ext = simd_reduce(row); + for (i, &v) in row.iter().enumerate() { + if v.is_nan() || v == ext { + return i as isize; + } + } + 0 + }; + + #[cfg(feature = "rayon")] + let indices: Vec = if outer_size * dim_size >= EXTREMUM_PARALLEL_THRESHOLD { + (0..outer_size).into_par_iter().map(find_row).collect() + } else { + (0..outer_size).map(find_row).collect() + }; + + #[cfg(not(feature = "rayon"))] + let indices: Vec = (0..outer_size).map(find_row).collect(); + + FlexTensor::new( + Bytes::from_elems(indices), + Layout::contiguous(Shape::from(out_shape)), + INDEX_DTYPE, + ) +} + +/// Find extremum value along a dimension. `is_better(new, current) -> bool`. +fn extremum_dim(tensor: &FlexTensor, dim: usize, is_better: F) -> FlexTensor +where + E: Element + bytemuck::Pod + Send + Sync, + F: Fn(E, E) -> bool + Send + Sync, +{ + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape(); + let ndims = shape.num_dims(); + assert!(dim < ndims); + + let dim_size = shape[dim]; + let mut out_shape: Vec = shape.to_vec(); + out_shape[dim] = 1; + let outer_size: usize = shape[..dim].iter().product::(); + let inner_size: usize = shape[dim + 1..].iter().product::(); + let out_size = outer_size * inner_size; + let data: &[E] = tensor.storage(); + let start_offset = tensor.layout().start_offset(); + + let find = |flat_idx: usize| -> E { + let outer = flat_idx / inner_size; + let inner = flat_idx % inner_size; + let base = start_offset + outer * dim_size * inner_size + inner; + let mut best = data[base]; + for d in 1..dim_size { + let val = data[base + d * inner_size]; + if is_better(val, best) { + best = val; + } + } + best + }; + + #[cfg(feature = "rayon")] + let values: Vec = if out_size * dim_size >= EXTREMUM_PARALLEL_THRESHOLD { + (0..out_size).into_par_iter().map(&find).collect() + } else { + (0..out_size).map(find).collect() + }; + + #[cfg(not(feature = "rayon"))] + let values: Vec = (0..out_size).map(find).collect(); + + FlexTensor::new( + Bytes::from_elems(values), + Layout::contiguous(Shape::from(out_shape)), + E::dtype(), + ) +} + +/// Find extremum value and its index along a dimension. `is_better(new, current) -> bool`. +fn extremum_dim_with_indices( + tensor: &FlexTensor, + dim: usize, + is_better: F, +) -> (FlexTensor, FlexTensor) +where + E: Element + bytemuck::Pod + Send + Sync, + F: Fn(E, E) -> bool + Send + Sync, +{ + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape(); + let ndims = shape.num_dims(); + assert!(dim < ndims); + + let dim_size = shape[dim]; + let mut out_shape: Vec = shape.to_vec(); + out_shape[dim] = 1; + let outer_size: usize = shape[..dim].iter().product::(); + let inner_size: usize = shape[dim + 1..].iter().product::(); + let out_size = outer_size * inner_size; + let data: &[E] = tensor.storage(); + let start_offset = tensor.layout().start_offset(); + + let find = |flat_idx: usize| -> (E, isize) { + let outer = flat_idx / inner_size; + let inner = flat_idx % inner_size; + let base = start_offset + outer * dim_size * inner_size + inner; + let mut best = data[base]; + let mut best_idx: isize = 0; + for d in 1..dim_size { + let val = data[base + d * inner_size]; + if is_better(val, best) { + best = val; + best_idx = d as isize; + } + } + (best, best_idx) + }; + + #[cfg(feature = "rayon")] + let (values, indices): (Vec, Vec) = + if out_size * dim_size >= EXTREMUM_PARALLEL_THRESHOLD { + (0..out_size).into_par_iter().map(&find).unzip() + } else { + (0..out_size).map(find).unzip() + }; + + #[cfg(not(feature = "rayon"))] + let (values, indices): (Vec, Vec) = (0..out_size).map(find).unzip(); + + let val_tensor = FlexTensor::new( + Bytes::from_elems(values), + Layout::contiguous(Shape::from(out_shape.clone())), + E::dtype(), + ); + let idx_tensor = FlexTensor::new( + Bytes::from_elems(indices), + Layout::contiguous(Shape::from(out_shape)), + INDEX_DTYPE, + ); + (val_tensor, idx_tensor) +} + +/// Find extremum value along a dimension for half-precision types (compared via f32). +fn extremum_dim_half( + tensor: &FlexTensor, + dim: usize, + is_better: F, + to_f32: fn(E) -> f32, + from_f32: fn(f32) -> E, +) -> FlexTensor +where + E: Element + bytemuck::Pod + Send + Sync, + F: Fn(f32, f32) -> bool + Send + Sync, +{ + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape(); + let ndims = shape.num_dims(); + assert!(dim < ndims); + + let dim_size = shape[dim]; + let mut out_shape: Vec = shape.to_vec(); + out_shape[dim] = 1; + let outer_size: usize = shape[..dim].iter().product::(); + let inner_size: usize = shape[dim + 1..].iter().product::(); + let out_size = outer_size * inner_size; + let data: &[E] = tensor.storage(); + let start_offset = tensor.layout().start_offset(); + + let find = |flat_idx: usize| -> E { + let outer = flat_idx / inner_size; + let inner = flat_idx % inner_size; + let base = start_offset + outer * dim_size * inner_size + inner; + let mut best = to_f32(data[base]); + for d in 1..dim_size { + let val = to_f32(data[base + d * inner_size]); + if is_better(val, best) { + best = val; + } + } + from_f32(best) + }; + + #[cfg(feature = "rayon")] + let values: Vec = if out_size * dim_size >= EXTREMUM_PARALLEL_THRESHOLD { + (0..out_size).into_par_iter().map(&find).collect() + } else { + (0..out_size).map(find).collect() + }; + + #[cfg(not(feature = "rayon"))] + let values: Vec = (0..out_size).map(find).collect(); + + FlexTensor::new( + Bytes::from_elems(values), + Layout::contiguous(Shape::from(out_shape)), + E::dtype(), + ) +} + +/// Find extremum value and index along a dimension for half-precision types (compared via f32). +fn extremum_dim_with_indices_half( + tensor: &FlexTensor, + dim: usize, + is_better: F, + to_f32: fn(E) -> f32, + from_f32: fn(f32) -> E, +) -> (FlexTensor, FlexTensor) +where + E: Element + bytemuck::Pod + Send + Sync, + F: Fn(f32, f32) -> bool + Send + Sync, +{ + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape(); + let ndims = shape.num_dims(); + assert!(dim < ndims); + + let dim_size = shape[dim]; + let mut out_shape: Vec = shape.to_vec(); + out_shape[dim] = 1; + let outer_size: usize = shape[..dim].iter().product::(); + let inner_size: usize = shape[dim + 1..].iter().product::(); + let out_size = outer_size * inner_size; + let data: &[E] = tensor.storage(); + let start_offset = tensor.layout().start_offset(); + + let find = |flat_idx: usize| -> (E, isize) { + let outer = flat_idx / inner_size; + let inner = flat_idx % inner_size; + let base = start_offset + outer * dim_size * inner_size + inner; + let mut best = to_f32(data[base]); + let mut best_idx: isize = 0; + for d in 1..dim_size { + let val = to_f32(data[base + d * inner_size]); + if is_better(val, best) { + best = val; + best_idx = d as isize; + } + } + (from_f32(best), best_idx) + }; + + #[cfg(feature = "rayon")] + let (values, indices): (Vec, Vec) = + if out_size * dim_size >= EXTREMUM_PARALLEL_THRESHOLD { + (0..out_size).into_par_iter().map(&find).unzip() + } else { + (0..out_size).map(find).unzip() + }; + + #[cfg(not(feature = "rayon"))] + let (values, indices): (Vec, Vec) = (0..out_size).map(find).unzip(); + + let val_tensor = FlexTensor::new( + Bytes::from_elems(values), + Layout::contiguous(Shape::from(out_shape.clone())), + E::dtype(), + ); + let idx_tensor = FlexTensor::new( + Bytes::from_elems(indices), + Layout::contiguous(Shape::from(out_shape)), + INDEX_DTYPE, + ); + (val_tensor, idx_tensor) +} + +// ============================================================================ +// Scalar division helpers +// ============================================================================ + +fn scalar_div + Copy>( + mut tensor: FlexTensor, + divisor: E, +) -> FlexTensor { + let data: &mut [E] = tensor.storage_mut(); + for x in data.iter_mut() { + *x = *x / divisor; + } + tensor +} + +// ============================================================================ +// Tests +// ============================================================================ + +// Tests kept here exercise flex-specific behavior: dtype storage selection +// for every numeric width (i8/i16/i32/i64/u8/u16/u32/u64, bf16/f16/f32/f64) +// and edge cases of the dtype-specific reduction kernels (zero-sized +// non-reduced dims across contiguous/half/widening paths, dim sizes that +// exceed the element's max, f16 mean overflow fusion, zero-size panics on +// max_dim/min_dim). Plain 1d/2d smokes, NaN propagation, negative-stride +// (flipped/transposed/narrowed) variants, 4D middle-dim reductions, and +// permuted-argmax regressions have been migrated to burn-backend-tests so +// they run against every backend. When adding new tests, keep them here +// only if they probe flex dtype storage or panic on flex internals; +// otherwise add them to crates/burn-backend-tests/tests/tensor/{float,int}/ops/. +#[cfg(test)] +mod tests { + use alloc::vec; + use burn_backend::TensorData; + use burn_backend::ops::{FloatTensorOps, IntTensorOps}; + use burn_std::{bf16, f16}; + + use crate::{Flex, FlexTensor}; + + #[test] + fn test_mean_f16_overflow_intermediate_sum() { + // Scalar `mean()` for f16 must fuse sum+divide on the f32 accumulator. + // Sum of 0..1024 is 523776, well above f16::MAX (65504), so a naive + // sum-then-divide that materialises the intermediate in f16 would clip + // to inf. The final mean (511.5) fits f16 comfortably. + let data: Vec = (0..1024).map(|i| f16::from_f32(i as f32)).collect(); + let tensor = FlexTensor::from_data(TensorData::new(data, [1024])); + + let result = Flex::float_mean(tensor); + let result_data = result.into_data(); + let values: &[f16] = bytemuck::cast_slice(&result_data.bytes); + + assert_eq!(values.len(), 1); + let mean = values[0].to_f32(); + assert!(mean.is_finite(), "mean overflowed to {mean}"); + assert!((mean - 511.5).abs() < 0.5, "expected ~511.5, got {mean}"); + } + + #[test] + fn test_mean_dim_f16_zero_outer_dim() { + // Regression test for mean_dim_half / sum_dim_contiguous_f32 clamping + // `outer_size.max(1)` in the last-dim branch: shape [0, 4] reducing + // dim=1 has outer_size=0, and the old clamp produced rows=1 which + // ran sum_rows_f32 past the end of an empty buffer. Result should be + // an empty tensor with shape [0, 1]. + let data: Vec = Vec::new(); + let tensor = FlexTensor::from_data(TensorData::new(data, [0, 4])); + let result = Flex::float_mean_dim(tensor, 1); + + assert_eq!(result.layout().shape().to_vec(), vec![0, 1]); + let result_data = result.into_data(); + let values: &[f16] = bytemuck::cast_slice(&result_data.bytes); + assert!(values.is_empty()); + } + + #[test] + fn test_sum_dim_f32_zero_outer_dim() { + // Regression: the f32 last-dim SIMD path (`reduce_last_dim_f32`) used + // `rows = outer_size.max(1)` which would index past the end of an + // empty data buffer for shape [0, K]. Guarded by the `out_size == 0` + // early return in `reduce_dim_f32`. + let data: Vec = Vec::new(); + let tensor = FlexTensor::from_data(TensorData::new(data, [0, 4])); + let result = Flex::float_sum_dim(tensor, 1); + + assert_eq!(result.layout().shape().to_vec(), vec![0, 1]); + assert!(result.into_data().bytes.is_empty()); + } + + #[test] + fn test_sum_dim_f32_zero_inner_dim() { + // Mirror of the outer-zero case: shape [3, 0] reducing dim=0 has + // inner_size=0. + let data: Vec = Vec::new(); + let tensor = FlexTensor::from_data(TensorData::new(data, [3, 0])); + let result = Flex::float_sum_dim(tensor, 0); + + assert_eq!(result.layout().shape().to_vec(), vec![1, 0]); + assert!(result.into_data().bytes.is_empty()); + } + + #[test] + fn test_sum_dim_f64_zero_outer_dim() { + // Covers `reduce_dim_impl` (generic contiguous/non-contiguous path) + // for the zero-sized non-reduced dim case. + let data: Vec = Vec::new(); + let tensor = FlexTensor::from_data(TensorData::new(data, [0, 4])); + let result = Flex::float_sum_dim(tensor, 1); + + assert_eq!(result.layout().shape().to_vec(), vec![0, 1]); + assert!(result.into_data().bytes.is_empty()); + } + + #[test] + fn test_sum_dim_i8_zero_outer_dim() { + // Covers `reduce_dim_widening` (i8/i16/u8/u16 path that accumulates + // in i64) for the zero-sized non-reduced dim case. + let data: Vec = Vec::new(); + let tensor = FlexTensor::from_data(TensorData::new(data, [0, 4])); + let result = Flex::int_sum_dim(tensor, 1); + + assert_eq!(result.layout().shape().to_vec(), vec![0, 1]); + assert!(result.into_data().bytes.is_empty()); + } + + #[test] + fn test_sum_dim_bf16_zero_outer_dim() { + // Covers `reduce_dim_half` (bf16/f16 sum_dim/prod_dim path) for the + // zero-sized non-reduced dim case. + let data: Vec = Vec::new(); + let tensor = FlexTensor::from_data(TensorData::new(data, [0, 4])); + let result = Flex::float_sum_dim(tensor, 1); + + assert_eq!(result.layout().shape().to_vec(), vec![0, 1]); + assert!(result.into_data().bytes.is_empty()); + } + + #[test] + fn test_mean_dim_i8_large_dimension() { + // dim_size=200 exceeds i8::MAX (127). Before the fix, 200 as i8 = -56, + // causing wrong results (or 256 as i8 = 0 causing div-by-zero). + let mut data: Vec = vec![0i8; 200]; + data[0] = 100; + let tensor = FlexTensor::from_data(TensorData::new(data, [1, 200])); + let result = Flex::int_mean_dim(tensor, 1); + + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + // integer division: 100 / 200 = 0 + assert_eq!(values, vec![0]); + } + + #[test] + fn test_mean_dim_i16_large_dimension() { + // dim_size=40000 exceeds i16::MAX (32767). + let mut data: Vec = vec![0i16; 40000]; + data[0] = 32000; + let tensor = FlexTensor::from_data(TensorData::new(data, [1, 40000])); + let result = Flex::int_mean_dim(tensor, 1); + + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + assert_eq!(values, vec![0]); + } + + #[test] + fn test_sum_i32() { + let data: Vec = vec![1, 2, 3, 4, 5]; + let tensor = FlexTensor::from_data(TensorData::new(data, [5])); + let result = Flex::int_sum(tensor); + + assert_eq!(result.layout().shape().to_vec(), vec![1]); + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + assert_eq!(values, vec![15]); + } + + #[test] + fn test_sum_dim_i32() { + let data: Vec = vec![1, 2, 3, 4, 5, 6]; + let tensor = FlexTensor::from_data(TensorData::new(data, [2, 3])); + let result = Flex::int_sum_dim(tensor, 1); + + assert_eq!(result.layout().shape().to_vec(), vec![2, 1]); + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + assert_eq!(values, vec![6, 15]); + } + + #[test] + fn test_argmax_i32() { + let data: Vec = vec![1, 5, 3, 2, 4]; + let tensor = FlexTensor::from_data(TensorData::new(data, [5])); + let result = Flex::int_argmax(tensor, 0); + + assert_eq!(result.layout().shape().to_vec(), vec![1]); + let result_data = result.into_data(); + #[cfg(target_pointer_width = "64")] + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + #[cfg(target_pointer_width = "32")] + let values: Vec = bytemuck::cast_slice::(&result_data.bytes) + .iter() + .map(|&v| v as i64) + .collect(); + assert_eq!(values, vec![1]); + } + + #[test] + #[should_panic(expected = "dimension 0 has size 0")] + fn test_max_dim_zero_size_panics() { + let tensor = FlexTensor::from_data(TensorData::new(Vec::::new(), [0, 3])); + Flex::float_max_dim(tensor, 0); + } + + #[test] + #[should_panic(expected = "dimension 1 has size 0")] + fn test_min_dim_zero_size_panics() { + let tensor = FlexTensor::from_data(TensorData::new(Vec::::new(), [3, 0])); + Flex::float_min_dim(tensor, 1); + } + + // === Unsigned integer dtype tests === + + #[test] + fn test_sum_u32() { + let tensor = FlexTensor::from_data(TensorData::new(vec![10u32, 20, 30], [3])); + let result = Flex::int_sum(tensor); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![60]); + } + + #[test] + fn test_sum_u64() { + let tensor = FlexTensor::from_data(TensorData::new(vec![100u64, 200, 300], [3])); + let result = Flex::int_sum(tensor); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![600]); + } + + #[test] + fn test_sum_dim_u8() { + let tensor = FlexTensor::from_data(TensorData::new(vec![1u8, 2, 3, 4], [2, 2])); + let result = Flex::int_sum_dim(tensor, 1); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![3, 7]); + } + + #[test] + fn test_prod_u16() { + let tensor = FlexTensor::from_data(TensorData::new(vec![2u16, 3, 5], [3])); + let result = Flex::int_prod(tensor); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![30]); + } + + #[test] + fn test_max_u32() { + let tensor = FlexTensor::from_data(TensorData::new(vec![5u32, 100, 42], [3])); + let result = Flex::int_max(tensor); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![100]); + } + + #[test] + fn test_min_u8() { + let tensor = FlexTensor::from_data(TensorData::new(vec![5u8, 1, 42], [3])); + let result = Flex::int_min(tensor); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![1]); + } + + #[test] + fn test_max_dim_u64() { + let tensor = FlexTensor::from_data(TensorData::new(vec![10u64, 20, 30, 5], [2, 2])); + let result = Flex::int_max_dim(tensor, 1); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![20, 30]); + } + + #[test] + fn test_min_dim_u16() { + let tensor = FlexTensor::from_data(TensorData::new(vec![10u16, 2, 30, 5], [2, 2])); + let result = Flex::int_min_dim(tensor, 1); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![2, 5]); + } + + #[test] + fn test_mean_dim_u8() { + let tensor = FlexTensor::from_data(TensorData::new(vec![10u8, 20, 30, 40], [2, 2])); + let result = Flex::int_mean_dim(tensor, 1); + let data: Vec = result.into_data().to_vec().unwrap(); + assert_eq!(data, vec![15, 35]); + } + + #[test] + fn test_max_dim_with_indices_u32() { + let tensor = FlexTensor::from_data(TensorData::new(vec![5u32, 10, 3, 8], [2, 2])); + let (values, indices) = Flex::int_max_dim_with_indices(tensor, 1); + let vals: Vec = values.into_data().to_vec().unwrap(); + let idxs: Vec = bytemuck::cast_slice(&indices.into_data().bytes).to_vec(); + assert_eq!(vals, vec![10, 8]); + assert_eq!(idxs, vec![1, 1]); + } + + // Cross-path consistency: `argmax`/`argmin` route short rows to the + // scalar kernel and rows of length >= EXTREMUM_SIMD_ROW_THRESHOLD to + // the SIMD kernel. Both kernels must agree on "first NaN wins". + // This is a flex-internal dispatch concern; the behavioral NaN + // propagation contract itself is exercised in burn-backend-tests + // under the `flex` feature gate (see issue #4814). + #[test] + fn test_argmax_scalar_and_simd_paths_agree_on_leading_nan() { + let short = + FlexTensor::from_data(TensorData::new(vec![f32::NAN, f32::NAN, f32::NAN], [1, 3])); + let short_idxs: Vec = + bytemuck::cast_slice(&super::argmax(short, 1).into_data().bytes).to_vec(); + + let mut long_data = alloc::vec![1.0f32; 600]; + long_data[0] = f32::NAN; + long_data[1] = f32::NAN; + long_data[300] = 5.0; + let long = FlexTensor::from_data(TensorData::new(long_data, [1, 600])); + let long_idxs: Vec = + bytemuck::cast_slice(&super::argmax(long, 1).into_data().bytes).to_vec(); + + assert_eq!(short_idxs, vec![0], "scalar path"); + assert_eq!(long_idxs, vec![0], "SIMD path"); + } +} diff --git a/crates/burn-flex/src/ops/repeat_dim.rs b/crates/burn-flex/src/ops/repeat_dim.rs new file mode 100644 index 0000000..706bf98 --- /dev/null +++ b/crates/burn-flex/src/ops/repeat_dim.rs @@ -0,0 +1,54 @@ +//! Repeat a tensor along a dimension. + +use alloc::vec::Vec; +use burn_std::{Bytes, Shape}; + +use crate::{FlexTensor, Layout}; + +/// Repeat `tensor` along `dim` by `times`. +pub fn repeat_dim(tensor: FlexTensor, dim: usize, times: usize) -> FlexTensor { + if times == 1 { + return tensor; + } + + let ndims = tensor.layout().num_dims(); + assert!( + dim < ndims, + "repeat_dim: dim {} out of bounds for tensor with {} dimensions", + dim, + ndims + ); + + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let dtype = tensor.dtype(); + let elem_size = crate::tensor::dtype_size(dtype); + + let mut new_dims: Vec = shape.iter().cloned().collect(); + new_dims[dim] *= times; + let new_shape = Shape::from(new_dims); + + let src: &[u8] = tensor.bytes(); + let n = new_shape.num_elements() * elem_size; + let mut dst: Vec = Vec::with_capacity(n); + + let inner: usize = shape.iter().skip(dim + 1).product(); + let dim_size = shape[dim]; + let chunk_bytes = dim_size * inner * elem_size; + let outer: usize = shape.iter().take(dim).product(); + + for o in 0..outer { + let start = o * chunk_bytes; + let end = start + chunk_bytes; + for _t in 0..times { + dst.extend_from_slice(&src[start..end]); + } + } + + debug_assert_eq!(dst.len(), n); + FlexTensor::new( + Bytes::from_bytes_vec(dst), + Layout::contiguous(new_shape), + dtype, + ) +} diff --git a/crates/burn-flex/src/ops/slice.rs b/crates/burn-flex/src/ops/slice.rs new file mode 100644 index 0000000..4844b89 --- /dev/null +++ b/crates/burn-flex/src/ops/slice.rs @@ -0,0 +1,690 @@ +//! Slice operations for FlexTensor. + +use alloc::vec; +use alloc::vec::Vec; +use burn_backend::{DType, Element}; +use burn_std::{Bytes, Shape, Slice, bf16, f16}; + +use crate::{FlexTensor, Layout}; + +/// Slice a tensor according to the given slice parameters. +/// +/// For positive steps, this is zero-copy (metadata only). +/// For negative steps, data is copied to handle the reversal. +pub fn slice(tensor: FlexTensor, slices: &[Slice]) -> FlexTensor { + let (new_layout, needs_copy) = tensor.layout().slice(slices); + + if !needs_copy { + // Zero-copy: share data with new layout + FlexTensor::from_arc(tensor.data_arc(), new_layout, tensor.dtype()) + } else { + // Needs copy due to negative steps + slice_with_copy(&tensor, slices) + } +} + +/// Slice with data copy (handles negative steps). +fn slice_with_copy(tensor: &FlexTensor, slices: &[Slice]) -> FlexTensor { + match tensor.dtype() { + DType::F32 => slice_copy_impl::(tensor, slices), + DType::F64 => slice_copy_impl::(tensor, slices), + DType::F16 => slice_copy_impl::(tensor, slices), + DType::BF16 => slice_copy_impl::(tensor, slices), + DType::I32 => slice_copy_impl::(tensor, slices), + DType::I64 => slice_copy_impl::(tensor, slices), + DType::I16 => slice_copy_impl::(tensor, slices), + DType::I8 => slice_copy_impl::(tensor, slices), + DType::U32 => slice_copy_impl::(tensor, slices), + DType::U64 => slice_copy_impl::(tensor, slices), + DType::U16 => slice_copy_impl::(tensor, slices), + DType::U8 => slice_copy_impl::(tensor, slices), + DType::Bool(_) => slice_copy_impl::(tensor, slices), + _ => panic!("slice: unsupported dtype {:?}", tensor.dtype()), + } +} + +/// Generic slice implementation with copy. +fn slice_copy_impl( + tensor: &FlexTensor, + slices: &[Slice], +) -> FlexTensor { + let src = tensor.storage::(); + let src_layout = tensor.layout(); + let ndims = src_layout.num_dims(); + + // Calculate output shape and collect normalized slice info + let mut out_shape = Vec::with_capacity(ndims); + let mut slice_info: Vec<(usize, usize, isize)> = Vec::with_capacity(ndims); // (start, len, step) + + for dim in 0..ndims { + let dim_size = src_layout.shape()[dim] as isize; + + let slice = if dim < slices.len() { + &slices[dim] + } else { + // Default: full range + &Slice::new(0, None, 1) + }; + + let (start, len, step) = compute_slice_info(slice, dim_size); + out_shape.push(len); + slice_info.push((start, len, step)); + } + + let out_layout = Layout::contiguous(Shape::from(out_shape.clone())); + let num_elements = out_layout.num_elements(); + + if num_elements == 0 { + let bytes = Bytes::from_elems::(Vec::new()); + return FlexTensor::new(bytes, out_layout, tensor.dtype()); + } + + // Allocate output + let mut out_data: Vec = Vec::with_capacity(num_elements); + + // Use recursive iteration for arbitrary dimensions + let mut indices = vec![0usize; ndims]; + copy_slice_recursive(src, src_layout, &slice_info, &mut out_data, &mut indices, 0); + + let bytes = Bytes::from_elems(out_data); + FlexTensor::new(bytes, out_layout, tensor.dtype()) +} + +/// Recursively copy sliced elements. +fn copy_slice_recursive( + src: &[E], + src_layout: &Layout, + slice_info: &[(usize, usize, isize)], + out: &mut Vec, + indices: &mut [usize], + dim: usize, +) { + let ndims = src_layout.num_dims(); + + if dim == ndims { + // Base case: copy single element + let src_idx = compute_src_index(src_layout, slice_info, indices); + out.push(src[src_idx]); + return; + } + + let (_, len, _) = slice_info[dim]; + + for i in 0..len { + indices[dim] = i; + copy_slice_recursive(src, src_layout, slice_info, out, indices, dim + 1); + } +} + +/// Compute source index from output indices and slice info. +fn compute_src_index( + layout: &Layout, + slice_info: &[(usize, usize, isize)], + out_indices: &[usize], +) -> usize { + let mut idx = layout.start_offset() as isize; + for (dim, &out_i) in out_indices.iter().enumerate() { + let (start, _, step) = slice_info[dim]; + let src_i = if step > 0 { + start + out_i * step as usize + } else { + // Negative step: start from high index, go down + let result = start as isize - (out_i as isize) * (-step); + debug_assert!(result >= 0, "slice: negative source index at dim {dim}"); + result as usize + }; + idx += src_i as isize * layout.strides()[dim]; + } + debug_assert!(idx >= 0, "slice: negative final index"); + idx as usize +} + +/// Normalize a potentially negative index to a positive one. +fn normalize_index(idx: isize, dim_size: isize) -> usize { + if idx < 0 { + (dim_size + idx).max(0) as usize + } else { + idx as usize + } +} + +/// Assign values to a slice of a tensor. +pub fn slice_assign(tensor: FlexTensor, slices: &[Slice], value: FlexTensor) -> FlexTensor { + match tensor.dtype() { + DType::F32 => slice_assign_impl::(tensor, slices, value), + DType::F64 => slice_assign_impl::(tensor, slices, value), + DType::F16 => slice_assign_impl::(tensor, slices, value), + DType::BF16 => slice_assign_impl::(tensor, slices, value), + DType::I32 => slice_assign_impl::(tensor, slices, value), + DType::I64 => slice_assign_impl::(tensor, slices, value), + DType::I16 => slice_assign_impl::(tensor, slices, value), + DType::I8 => slice_assign_impl::(tensor, slices, value), + DType::U32 => slice_assign_impl::(tensor, slices, value), + DType::U64 => slice_assign_impl::(tensor, slices, value), + DType::U16 => slice_assign_impl::(tensor, slices, value), + DType::U8 => slice_assign_impl::(tensor, slices, value), + DType::Bool(_) => slice_assign_impl::(tensor, slices, value), + _ => panic!("slice_assign: unsupported dtype {:?}", tensor.dtype()), + } +} + +/// Generic slice assign implementation. +fn slice_assign_impl( + tensor: FlexTensor, + slices: &[Slice], + value: FlexTensor, +) -> FlexTensor { + // Broadcast-scalar fast path: if `value` is a fully-broadcast + // scalar (all strides zero), read the scalar once instead of + // materializing the expansion via `to_contiguous`. The + // `num_elements > 0` gate also guards the storage read against + // zero-sized sources where `iter().all(...)` would be vacuously + // true. + if value.layout().num_elements() > 0 && value.layout().strides().iter().all(|&s| s == 0) { + let scalar = value.storage::()[value.layout().start_offset()]; + return slice_write_impl::(tensor, slices, WriteSource::Scalar(scalar)); + } + + let value = value.to_contiguous(); + let val_src: &[E] = value.storage::(); + slice_write_impl::(tensor, slices, WriteSource::Slice(val_src)) +} + +/// Source to splat into a sliced region of a destination tensor. The +/// two variants drive the same dispatch tree; `Scalar` hits the +/// broadcast-scalar fast path (no value buffer), `Slice` hits the +/// memcpy-style assign path (advances through `val_src`). +#[derive(Copy, Clone)] +enum WriteSource<'a, E: Copy> { + Scalar(E), + Slice(&'a [E]), +} + +impl<'a, E: Copy> WriteSource<'a, E> { + /// Write a contiguous span of `dst` starting at `dst_offset` for + /// `len` elements. For `Slice`, `src_offset` is the current + /// position in the value buffer. + #[inline] + fn write_span(self, dst: &mut [E], dst_offset: usize, len: usize, src_offset: usize) { + match self { + WriteSource::Scalar(s) => dst[dst_offset..dst_offset + len].fill(s), + WriteSource::Slice(src) => dst[dst_offset..dst_offset + len] + .copy_from_slice(&src[src_offset..src_offset + len]), + } + } + + /// Write a single element. `src_idx` is only read in the `Slice` + /// variant. + #[inline] + fn write_element(self, dst: &mut [E], dst_idx: usize, src_idx: usize) { + match self { + WriteSource::Scalar(s) => dst[dst_idx] = s, + WriteSource::Slice(src) => dst[dst_idx] = src[src_idx], + } + } +} + +/// Unified slice writer used by both `slice_assign_impl` and the +/// scalar-broadcast fast path. Walks the destination's sliced region +/// (1D / 2D inner-contig / ND inner-contig / strided fallback) and +/// pulls values from the given [`WriteSource`]. +fn slice_write_impl( + tensor: FlexTensor, + slices: &[Slice], + source: WriteSource<'_, E>, +) -> FlexTensor { + let mut tensor = tensor.to_contiguous(); + let dst_layout = tensor.layout().clone(); + let ndims = dst_layout.num_dims(); + + let slice_info: Vec<(usize, usize, isize)> = (0..ndims) + .map(|dim| { + let dim_size = dst_layout.shape()[dim] as isize; + let slice = if dim < slices.len() { + &slices[dim] + } else { + &Slice::new(0, None, 1) + }; + compute_slice_info(slice, dim_size) + }) + .collect(); + + let dst = tensor.storage_mut::(); + + let inner_contiguous = slice_info + .last() + .map(|(_, _, step)| *step == 1) + .unwrap_or(false); + + if ndims == 0 { + // Rank 0: single scalar destination. Only reachable from the + // scalar fast path; `slice_assign` on a rank-0 tensor with a + // rank-0 source also ends up here. + if !dst.is_empty() { + source.write_element(dst, 0, 0); + } + } else if ndims == 1 { + let (start, len, step) = slice_info[0]; + if step == 1 { + source.write_span(dst, start, len, 0); + } else { + for i in 0..len { + let dst_i = if step > 0 { + start + i * step as usize + } else { + (start as isize - (i as isize) * (-step)) as usize + }; + source.write_element(dst, dst_i, i); + } + } + } else if ndims == 2 && inner_contiguous { + let (row_start, row_len, row_step) = slice_info[0]; + let (col_start, col_len, _) = slice_info[1]; + let dst_cols = dst_layout.shape()[1]; + + let mut val_offset = 0; + for r in 0..row_len { + let row_idx = if row_step > 0 { + row_start + r * row_step as usize + } else { + (row_start as isize - (r as isize) * (-row_step)) as usize + }; + let dst_row_start = row_idx * dst_cols + col_start; + source.write_span(dst, dst_row_start, col_len, val_offset); + val_offset += col_len; + } + } else if inner_contiguous { + let inner_len = slice_info[ndims - 1].1; + let outer_dims = ndims - 1; + let dst_strides = dst_layout.strides(); + + let outer_count: usize = slice_info.iter().take(outer_dims).map(|i| i.1).product(); + + let mut outer_indices = vec![0usize; outer_dims]; + let mut val_offset = 0; + + for _ in 0..outer_count { + let mut dst_offset = dst_layout.start_offset() as isize; + for (dim, &idx) in outer_indices.iter().enumerate() { + let (start, _, step) = slice_info[dim]; + let src_i = if step > 0 { + start + idx * step as usize + } else { + (start as isize - (idx as isize) * (-step)) as usize + }; + dst_offset += src_i as isize * dst_strides[dim]; + } + dst_offset += slice_info[ndims - 1].0 as isize * dst_strides[ndims - 1]; + let dst_offset = dst_offset as usize; + + source.write_span(dst, dst_offset, inner_len, val_offset); + val_offset += inner_len; + + // Odometer increment over outer dims. + for dim in (0..outer_dims).rev() { + outer_indices[dim] += 1; + if outer_indices[dim] < slice_info[dim].1 { + break; + } + outer_indices[dim] = 0; + } + } + } else { + let total_elements: usize = slice_info.iter().map(|(_, len, _)| len).product(); + let dst_strides = dst_layout.strides(); + let mut indices = vec![0usize; ndims]; + + for i in 0..total_elements { + let mut dst_offset = dst_layout.start_offset() as isize; + for (dim, &idx) in indices.iter().enumerate() { + let (start, _, step) = slice_info[dim]; + let src_i = if step > 0 { + start + idx * step as usize + } else { + (start as isize - (idx as isize) * (-step)) as usize + }; + dst_offset += src_i as isize * dst_strides[dim]; + } + + source.write_element(dst, dst_offset as usize, i); + + for dim in (0..ndims).rev() { + indices[dim] += 1; + if indices[dim] < slice_info[dim].1 { + break; + } + indices[dim] = 0; + } + } + } + + tensor +} + +/// Compute slice info (start, len, step) for a dimension. +/// For negative step: start is the LAST index in the range (end-1), iterating down. +fn compute_slice_info(slice: &Slice, dim_size: isize) -> (usize, usize, isize) { + let step = slice.step; + let abs_step = step.unsigned_abs(); + + // Normalize start and end to [0, dim_size] + let range_start = normalize_index(slice.start, dim_size); + let range_end = match slice.end { + Some(e) => normalize_index(e, dim_size).min(dim_size as usize), + None => dim_size as usize, + }; + + let len = if range_end > range_start { + (range_end - range_start).div_ceil(abs_step) + } else { + 0 + }; + + if step > 0 { + // Forward: start at low index, go up + (range_start, len, step) + } else { + // Reverse: start at end-1 (highest index in range), go down + // For s![2..8;-2]: start from index 7, go to 5, then 3 + let reverse_start = if range_end > range_start { + range_end - 1 + } else { + range_start + }; + (reverse_start, len, step) + } +} + +// Tests kept here exercise flex-specific behavior: the internal +// `slice` / `slice_assign` helpers, the broadcast-scalar fast paths for +// `slice_fill` (1D contiguous, 2D inner-contig, 3D inner-contig, ND +// strided fallback, stepped-row 2D inner-contig), and non-f32 dtype +// coverage. General slice correctness across backends is covered by +// crates/burn-backend-tests/tests/tensor/float/ops/{slice,slice_assign}.rs. +#[cfg(test)] +mod tests { + use super::*; + use burn_backend::TensorData; + + #[test] + fn test_slice_basic() { + // Create a 2x3 tensor: [[0, 1, 2], [3, 4, 5]] + let data: Vec = vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0]; + let tensor = FlexTensor::from_data(TensorData::new(data, [2, 3])); + + // Slice [0:1, 1:3] -> [[1, 2]] + let slices = vec![Slice::new(0, Some(1), 1), Slice::new(1, Some(3), 1)]; + let result = slice(tensor, &slices); + + assert_eq!(result.layout().shape().to_vec(), vec![1, 2]); + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + assert_eq!(values, vec![1.0, 2.0]); + } + + #[test] + fn test_slice_with_step() { + // Create a 1D tensor: [0, 1, 2, 3, 4, 5] + let data: Vec = vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0]; + let tensor = FlexTensor::from_data(TensorData::new(data, [6])); + + // Slice [0:6:2] -> [0, 2, 4] + let slices = vec![Slice::new(0, Some(6), 2)]; + let result = slice(tensor, &slices); + + assert_eq!(result.layout().shape().to_vec(), vec![3]); + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + assert_eq!(values, vec![0.0, 2.0, 4.0]); + } + + #[test] + fn test_slice_negative_index() { + // Create a 1D tensor: [0, 1, 2, 3, 4] + let data: Vec = vec![0.0, 1.0, 2.0, 3.0, 4.0]; + let tensor = FlexTensor::from_data(TensorData::new(data, [5])); + + // Slice [-3:] -> [2, 3, 4] + let slices = vec![Slice::new(-3, None, 1)]; + let result = slice(tensor, &slices); + + assert_eq!(result.layout().shape().to_vec(), vec![3]); + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + assert_eq!(values, vec![2.0, 3.0, 4.0]); + } + + #[test] + fn test_slice_negative_step() { + // Create a 1D tensor: [0, 1, 2, 3, 4] + let data: Vec = vec![0.0, 1.0, 2.0, 3.0, 4.0]; + let tensor = FlexTensor::from_data(TensorData::new(data, [5])); + + // Slice [0..;-1] -> [4, 3, 2, 1, 0] (reverse full range) + // In Burn's semantics: range selects elements, step determines order + let slices = vec![Slice::new(0, None, -1)]; + let result = slice(tensor, &slices); + + assert_eq!(result.layout().shape().to_vec(), vec![5]); + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + assert_eq!(values, vec![4.0, 3.0, 2.0, 1.0, 0.0]); + } + + #[test] + fn test_slice_assign_1d() { + // Create a 1D tensor: [0, 1, 2, 3, 4] + let data: Vec = vec![0.0, 1.0, 2.0, 3.0, 4.0]; + let tensor = FlexTensor::from_data(TensorData::new(data, [5])); + + // Assign [10, 11, 12] to positions [1:4] + let value_data: Vec = vec![10.0, 11.0, 12.0]; + let value = FlexTensor::from_data(TensorData::new(value_data, [3])); + let slices = vec![Slice::new(1, Some(4), 1)]; + let result = slice_assign(tensor, &slices, value); + + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + assert_eq!(values, vec![0.0, 10.0, 11.0, 12.0, 4.0]); + } + + #[test] + fn test_slice_assign_2d() { + // Create a 3x3 tensor + let data: Vec = vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; + let tensor = FlexTensor::from_data(TensorData::new(data, [3, 3])); + + // Assign [[10, 11], [12, 13]] to [1:3, 1:3] + let value_data: Vec = vec![10.0, 11.0, 12.0, 13.0]; + let value = FlexTensor::from_data(TensorData::new(value_data, [2, 2])); + let slices = vec![Slice::new(1, Some(3), 1), Slice::new(1, Some(3), 1)]; + let result = slice_assign(tensor, &slices, value); + + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + assert_eq!( + values, + vec![0.0, 1.0, 2.0, 3.0, 10.0, 11.0, 6.0, 12.0, 13.0,] + ); + } + + #[test] + fn test_slice_assign_2d_full_row() { + // Create a 3x4 tensor + let data: Vec = (0..12).map(|i| i as f32).collect(); + let tensor = FlexTensor::from_data(TensorData::new(data, [3, 4])); + + // Assign [100, 101, 102, 103] to row 1 + let value_data: Vec = vec![100.0, 101.0, 102.0, 103.0]; + let value = FlexTensor::from_data(TensorData::new(value_data, [1, 4])); + let slices = vec![Slice::new(1, Some(2), 1), Slice::new(0, None, 1)]; + let result = slice_assign(tensor, &slices, value); + + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + assert_eq!( + values, + vec![ + 0.0, 1.0, 2.0, 3.0, 100.0, 101.0, 102.0, 103.0, 8.0, 9.0, 10.0, 11.0, + ] + ); + } + + // Broadcast-scalar fast path tests: mimic what + // `Tensor::slice_fill` produces (a 1-element source expanded to + // the slice shape with all strides zero). + + fn broadcast_scalar_f32(value: f32, target_shape: &[usize]) -> FlexTensor { + let scalar_tensor = FlexTensor::from_data(TensorData::new(vec![value], [1])); + crate::ops::expand::expand(scalar_tensor, Shape::from(target_shape.to_vec())) + } + + #[test] + fn test_slice_assign_broadcast_scalar_1d_contiguous() { + let data: Vec = vec![0.0, 1.0, 2.0, 3.0, 4.0]; + let tensor = FlexTensor::from_data(TensorData::new(data, [5])); + let value = broadcast_scalar_f32(7.0, &[3]); + let slices = vec![Slice::new(1, Some(4), 1)]; + let result = slice_assign(tensor, &slices, value); + + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + assert_eq!(values, vec![0.0, 7.0, 7.0, 7.0, 4.0]); + } + + #[test] + fn test_slice_assign_broadcast_scalar_2d_inner_contiguous() { + let data: Vec = (0..16).map(|i| i as f32).collect(); + let tensor = FlexTensor::from_data(TensorData::new(data, [4, 4])); + let value = broadcast_scalar_f32(-1.0, &[2, 2]); + let slices = vec![Slice::new(1, Some(3), 1), Slice::new(1, Some(3), 1)]; + let result = slice_assign(tensor, &slices, value); + + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + assert_eq!( + values, + vec![ + 0.0, 1.0, 2.0, 3.0, 4.0, -1.0, -1.0, 7.0, 8.0, -1.0, -1.0, 11.0, 12.0, 13.0, 14.0, + 15.0, + ] + ); + } + + #[test] + fn test_slice_assign_broadcast_scalar_3d_inner_contiguous() { + // 3D case that matched the user-reported regression shape. + let data: Vec = (0..24).map(|i| i as f32).collect(); + let tensor = FlexTensor::from_data(TensorData::new(data, [2, 3, 4])); + let value = broadcast_scalar_f32(9.0, &[1, 2, 2]); + let slices = vec![ + Slice::new(0, Some(1), 1), + Slice::new(0, Some(2), 1), + Slice::new(1, Some(3), 1), + ]; + let result = slice_assign(tensor, &slices, value); + + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + // Region [0..1, 0..2, 1..3] fills positions (0, 0, 1), (0, 0, 2), + // (0, 1, 1), (0, 1, 2) with 9.0. Linear indices: 1, 2, 5, 6. + let mut expected: Vec = (0..24).map(|i| i as f32).collect(); + for &i in &[1usize, 2, 5, 6] { + expected[i] = 9.0; + } + assert_eq!(values, expected); + } + + #[test] + fn test_slice_assign_broadcast_scalar_strided_fallback() { + // Stepped slice: hits the strided fallback (not inner-contiguous). + let data: Vec = (0..10).map(|i| i as f32).collect(); + let tensor = FlexTensor::from_data(TensorData::new(data, [10])); + // Source needs to match slice_info length: s![0..10;2] selects + // 5 positions (0, 2, 4, 6, 8), so the expand target is [5]. + let value = broadcast_scalar_f32(0.0, &[5]); + let slices = vec![Slice::new(0, Some(10), 2)]; + let result = slice_assign(tensor, &slices, value); + + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + assert_eq!( + values, + vec![0.0, 1.0, 0.0, 3.0, 0.0, 5.0, 0.0, 7.0, 0.0, 9.0] + ); + } + + /// Broadcast-scalar fast path on a non-f32 dtype. + #[test] + fn test_slice_assign_broadcast_scalar_i64() { + fn broadcast_scalar_i64(value: i64, target_shape: &[usize]) -> FlexTensor { + let scalar_tensor = FlexTensor::from_data(TensorData::new(vec![value], [1])); + crate::ops::expand::expand(scalar_tensor, Shape::from(target_shape.to_vec())) + } + + let data: Vec = (0..12).collect(); + let tensor = FlexTensor::from_data(TensorData::new(data, [3, 4])); + let value = broadcast_scalar_i64(-7, &[2, 2]); + let slices = vec![Slice::new(0, Some(2), 1), Slice::new(1, Some(3), 1)]; + let result = slice_assign(tensor, &slices, value); + + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + // Positions (0, 1), (0, 2), (1, 1), (1, 2): linear indices 1, + // 2, 5, 6 get replaced by -7. + assert_eq!(values, vec![0, -7, -7, 3, 4, -7, -7, 7, 8, 9, 10, 11]); + } + + /// ND strided fallback path: 3D slice with a stepped inner dim. + #[test] + fn test_slice_assign_broadcast_scalar_nd_strided_fallback() { + let data: Vec = (0..24).map(|i| i as f32).collect(); + let tensor = FlexTensor::from_data(TensorData::new(data, [2, 3, 4])); + // Step 2 on the innermost dim means slice_info's last step != 1, + // so inner_contiguous is false and we take the ND fallback. + let value = broadcast_scalar_f32(9.0, &[2, 3, 2]); + let slices = vec![ + Slice::new(0, Some(2), 1), + Slice::new(0, Some(3), 1), + Slice::new(0, Some(4), 2), + ]; + let result = slice_assign(tensor, &slices, value); + + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + // Step-2 on dim 2 picks indices 0 and 2, so within each + // `[b, r, :]` row the 0 and 2 positions get 9.0. + let mut expected: Vec = (0..24).map(|i| i as f32).collect(); + for b in 0..2 { + for r in 0..3 { + for c in [0, 2] { + expected[b * 12 + r * 4 + c] = 9.0; + } + } + } + assert_eq!(values, expected); + } + + /// 2D inner-contig branch with `row_step > 1`. + #[test] + fn test_slice_assign_broadcast_scalar_2d_stepped_rows() { + let data: Vec = (0..25).map(|i| i as f32).collect(); + let tensor = FlexTensor::from_data(TensorData::new(data, [5, 5])); + // Step-2 rows pick out rows 0, 2, 4 (3 rows). Slice the inner + // dim as a contiguous range so the 2D-inner-contig branch with + // row_step != 1 fires. + let value = broadcast_scalar_f32(-1.0, &[3, 3]); + let slices = vec![Slice::new(0, Some(5), 2), Slice::new(1, Some(4), 1)]; + let result = slice_assign(tensor, &slices, value); + + let result_data = result.into_data(); + let values: Vec = bytemuck::cast_slice(&result_data.bytes).to_vec(); + let mut expected: Vec = (0..25).map(|i| i as f32).collect(); + for r in [0, 2, 4] { + for c in 1..4 { + expected[r * 5 + c] = -1.0; + } + } + assert_eq!(values, expected); + } +} diff --git a/crates/burn-flex/src/ops/sort.rs b/crates/burn-flex/src/ops/sort.rs new file mode 100644 index 0000000..0604e55 --- /dev/null +++ b/crates/burn-flex/src/ops/sort.rs @@ -0,0 +1,715 @@ +//! Sort and argsort operations for FlexTensor. +//! +//! Operates directly on storage without TensorData round-trips. + +use alloc::vec; +use alloc::vec::Vec; +use burn_backend::{DType, Element}; +use burn_std::{Bytes, Shape, bf16, f16}; +use bytemuck::Pod; + +#[cfg(feature = "rayon")] +use rayon::prelude::*; + +use crate::{FlexTensor, Layout}; + +use super::INDEX_DTYPE; +#[cfg(feature = "rayon")] +use super::PARALLEL_THRESHOLD; + +/// Validate sort dimension and check for empty tensors. +/// Returns `true` if the tensor is empty (caller should return early). +fn validate_sort_args(shape: &Shape, dim: usize) -> bool { + assert!( + dim < shape.num_dims(), + "sort: dim {} out of bounds for tensor with {} dimensions", + dim, + shape.num_dims() + ); + let dim_size = shape[dim]; + assert!( + dim_size <= isize::MAX as usize, + "sort: dimension {} has size {} which exceeds isize::MAX", + dim, + dim_size + ); + shape.num_elements() == 0 +} + +/// Sort elements along a dimension, returning the sorted tensor. +pub fn sort(tensor: FlexTensor, dim: usize, descending: bool) -> FlexTensor { + match tensor.dtype() { + DType::F32 => sort_typed::(tensor, dim, descending, f32::total_cmp), + DType::F64 => sort_typed::(tensor, dim, descending, f64::total_cmp), + DType::F16 => sort_half(tensor, dim, descending, f16::to_f32, f16::from_f32), + DType::BF16 => sort_half(tensor, dim, descending, bf16::to_f32, bf16::from_f32), + DType::I64 => sort_typed::(tensor, dim, descending, Ord::cmp), + DType::I32 => sort_typed::(tensor, dim, descending, Ord::cmp), + DType::I16 => sort_typed::(tensor, dim, descending, Ord::cmp), + DType::I8 => sort_typed::(tensor, dim, descending, Ord::cmp), + DType::U64 => sort_typed::(tensor, dim, descending, Ord::cmp), + DType::U32 => sort_typed::(tensor, dim, descending, Ord::cmp), + DType::U16 => sort_typed::(tensor, dim, descending, Ord::cmp), + DType::U8 => sort_typed::(tensor, dim, descending, Ord::cmp), + dt => panic!("sort: unsupported dtype {:?}", dt), + } +} + +/// Sort elements along a dimension, returning (sorted_values, indices). +pub fn sort_with_indices( + tensor: FlexTensor, + dim: usize, + descending: bool, +) -> (FlexTensor, FlexTensor) { + match tensor.dtype() { + DType::F32 => sort_with_indices_typed::(tensor, dim, descending, f32::total_cmp), + DType::F64 => sort_with_indices_typed::(tensor, dim, descending, f64::total_cmp), + DType::F16 => sort_with_indices_half(tensor, dim, descending, f16::to_f32, f16::from_f32), + DType::BF16 => { + sort_with_indices_half(tensor, dim, descending, bf16::to_f32, bf16::from_f32) + } + DType::I64 => sort_with_indices_typed::(tensor, dim, descending, Ord::cmp), + DType::I32 => sort_with_indices_typed::(tensor, dim, descending, Ord::cmp), + DType::I16 => sort_with_indices_typed::(tensor, dim, descending, Ord::cmp), + DType::I8 => sort_with_indices_typed::(tensor, dim, descending, Ord::cmp), + DType::U64 => sort_with_indices_typed::(tensor, dim, descending, Ord::cmp), + DType::U32 => sort_with_indices_typed::(tensor, dim, descending, Ord::cmp), + DType::U16 => sort_with_indices_typed::(tensor, dim, descending, Ord::cmp), + DType::U8 => sort_with_indices_typed::(tensor, dim, descending, Ord::cmp), + dt => panic!("sort_with_indices: unsupported dtype {:?}", dt), + } +} + +/// Argsort along a dimension, returning indices that would sort the tensor. +pub fn argsort(tensor: FlexTensor, dim: usize, descending: bool) -> FlexTensor { + match tensor.dtype() { + DType::F32 => argsort_typed::(tensor, dim, descending, f32::total_cmp), + DType::F64 => argsort_typed::(tensor, dim, descending, f64::total_cmp), + DType::F16 => argsort_half(tensor, dim, descending, f16::to_f32), + DType::BF16 => argsort_half(tensor, dim, descending, bf16::to_f32), + DType::I64 => argsort_typed::(tensor, dim, descending, Ord::cmp), + DType::I32 => argsort_typed::(tensor, dim, descending, Ord::cmp), + DType::I16 => argsort_typed::(tensor, dim, descending, Ord::cmp), + DType::I8 => argsort_typed::(tensor, dim, descending, Ord::cmp), + DType::U64 => argsort_typed::(tensor, dim, descending, Ord::cmp), + DType::U32 => argsort_typed::(tensor, dim, descending, Ord::cmp), + DType::U16 => argsort_typed::(tensor, dim, descending, Ord::cmp), + DType::U8 => argsort_typed::(tensor, dim, descending, Ord::cmp), + dt => panic!("argsort: unsupported dtype {:?}", dt), + } +} + +// --------------------------------------------------------------------------- +// Typed sort (operates directly on storage) +// --------------------------------------------------------------------------- + +fn sort_typed( + tensor: FlexTensor, + dim: usize, + descending: bool, + cmp: fn(&E, &E) -> core::cmp::Ordering, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let dtype = tensor.dtype(); + if validate_sort_args(&shape, dim) { + return tensor; + } + + let mut data: Vec = tensor.storage::().to_vec(); + + if shape.num_dims() == 1 { + if descending { + data.sort_unstable_by(|a, b| cmp(b, a)); + } else { + data.sort_unstable_by(cmp); + } + } else { + sort_along_dim(&mut data, &shape, dim, descending, cmp); + } + + FlexTensor::new(Bytes::from_elems(data), Layout::contiguous(shape), dtype) +} + +fn sort_with_indices_typed( + tensor: FlexTensor, + dim: usize, + descending: bool, + cmp: fn(&E, &E) -> core::cmp::Ordering, +) -> (FlexTensor, FlexTensor) { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let dtype = tensor.dtype(); + let n = shape.num_elements(); + if validate_sort_args(&shape, dim) { + let idx = make_index_tensor(Vec::new(), shape.clone()); + return (tensor, idx); + } + + let src: &[E] = tensor.storage(); + let mut values: Vec = src.to_vec(); + let mut indices: Vec = vec![0; n]; + + if shape.num_dims() == 1 { + sort_1d_with_indices(&mut values, &mut indices, descending, cmp); + } else { + sort_along_dim_with_indices(&mut values, &mut indices, &shape, dim, descending, cmp); + } + + let idx_tensor = make_index_tensor(indices, shape.clone()); + let val_tensor = FlexTensor::new(Bytes::from_elems(values), Layout::contiguous(shape), dtype); + (val_tensor, idx_tensor) +} + +/// Argsort without materializing sorted values. +fn argsort_typed( + tensor: FlexTensor, + dim: usize, + descending: bool, + cmp: fn(&E, &E) -> core::cmp::Ordering, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let n = shape.num_elements(); + if validate_sort_args(&shape, dim) { + return make_index_tensor(Vec::new(), shape); + } + + let src: &[E] = tensor.storage(); + let mut indices: Vec = vec![0; n]; + + if shape.num_dims() == 1 { + let mut idx_vec: Vec = (0..n).collect(); + if descending { + idx_vec.sort_unstable_by(|&a, &b| cmp(&src[b], &src[a])); + } else { + idx_vec.sort_unstable_by(|&a, &b| cmp(&src[a], &src[b])); + } + for (out_i, &orig_i) in idx_vec.iter().enumerate() { + indices[out_i] = orig_i as isize; + } + } else { + argsort_along_dim(src, &mut indices, &shape, dim, descending, cmp); + } + + make_index_tensor(indices, shape) +} + +/// 1D sort with index tracking, shared by typed and half-precision paths. +fn sort_1d_with_indices( + values: &mut [E], + indices: &mut [isize], + descending: bool, + cmp: fn(&E, &E) -> core::cmp::Ordering, +) { + let n = values.len(); + let mut idx_vec: Vec = (0..n).collect(); + if descending { + idx_vec.sort_unstable_by(|&a, &b| cmp(&values[b], &values[a])); + } else { + idx_vec.sort_unstable_by(|&a, &b| cmp(&values[a], &values[b])); + } + // Apply permutation in one pass using the sorted index order + let old_values = values.to_vec(); + for (out_i, &orig_i) in idx_vec.iter().enumerate() { + values[out_i] = old_values[orig_i]; + indices[out_i] = orig_i as isize; + } +} + +/// Sort along a given dimension for N-D tensors. +fn sort_along_dim( + data: &mut [E], + shape: &Shape, + dim: usize, + descending: bool, + cmp: fn(&E, &E) -> core::cmp::Ordering, +) { + let strides = contiguous_strides(shape); + let dim_size = shape[dim]; + let dim_stride = strides[dim]; + let num_slices = data.len() / dim_size; + + // Fast path: last dimension (stride==1). Rows are contiguous at + // offsets `slice_idx * dim_size`, so `chunks_exact_mut(dim_size)` + // walks them directly. Parallelized with rayon above the threshold. + if dim_stride == 1 { + // chunks_exact silently drops any remainder, so lock the + // "length is a multiple of dim_size" invariant that holds for + // contiguous tensors by construction. + debug_assert_eq!(data.len() % dim_size, 0); + let sort_row = |row: &mut [E]| { + if descending { + row.sort_unstable_by(|a, b| cmp(b, a)); + } else { + row.sort_unstable_by(cmp); + } + }; + + #[cfg(feature = "rayon")] + if data.len() >= PARALLEL_THRESHOLD { + data.par_chunks_exact_mut(dim_size).for_each(sort_row); + return; + } + + data.chunks_exact_mut(dim_size).for_each(sort_row); + return; + } + + let mut slice_buf: Vec = vec![data[0]; dim_size]; + + for slice_idx in 0..num_slices { + let base = slice_base_offset(slice_idx, shape, &strides, dim); + + for i in 0..dim_size { + slice_buf[i] = data[base + i * dim_stride]; + } + + if descending { + slice_buf.sort_unstable_by(|a, b| cmp(b, a)); + } else { + slice_buf.sort_unstable_by(cmp); + } + + for i in 0..dim_size { + data[base + i * dim_stride] = slice_buf[i]; + } + } +} + +/// Sort along a dimension, tracking original indices. +fn sort_along_dim_with_indices( + data: &mut [E], + indices: &mut [isize], + shape: &Shape, + dim: usize, + descending: bool, + cmp: fn(&E, &E) -> core::cmp::Ordering, +) { + let strides = contiguous_strides(shape); + let dim_size = shape[dim]; + let dim_stride = strides[dim]; + let num_slices = data.len() / dim_size; + + // Fast path: last dimension (stride==1). Values and indices rows + // are both contiguous at `slice_idx * dim_size`, so we can zip + // matching chunks and avoid the per-row stride arithmetic. + if dim_stride == 1 { + // zip silently truncates to the shorter iterator and chunks_exact + // silently drops remainders, so lock both invariants. + debug_assert_eq!(data.len(), indices.len()); + debug_assert_eq!(data.len() % dim_size, 0); + // Buffer is reused across rows (one per thread under rayon via + // `for_each_init`) to avoid a heap alloc per row. + let sort_row = |pairs: &mut Vec<(usize, E)>, (row, idx_row): (&mut [E], &mut [isize])| { + pairs.clear(); + pairs.extend((0..dim_size).map(|i| (i, row[i]))); + if descending { + pairs.sort_unstable_by(|a, b| cmp(&b.1, &a.1)); + } else { + pairs.sort_unstable_by(|a, b| cmp(&a.1, &b.1)); + } + for (i, &(orig_idx, val)) in pairs.iter().enumerate() { + row[i] = val; + idx_row[i] = orig_idx as isize; + } + }; + + #[cfg(feature = "rayon")] + if data.len() >= PARALLEL_THRESHOLD { + data.par_chunks_exact_mut(dim_size) + .zip(indices.par_chunks_exact_mut(dim_size)) + .for_each_init(|| Vec::with_capacity(dim_size), sort_row); + return; + } + + let mut pairs: Vec<(usize, E)> = Vec::with_capacity(dim_size); + data.chunks_exact_mut(dim_size) + .zip(indices.chunks_exact_mut(dim_size)) + .for_each(|row_and_idx| sort_row(&mut pairs, row_and_idx)); + return; + } + + let mut pairs: Vec<(usize, E)> = Vec::with_capacity(dim_size); + + for slice_idx in 0..num_slices { + let base = slice_base_offset(slice_idx, shape, &strides, dim); + + pairs.clear(); + for i in 0..dim_size { + pairs.push((i, data[base + i * dim_stride])); + } + + if descending { + pairs.sort_unstable_by(|a, b| cmp(&b.1, &a.1)); + } else { + pairs.sort_unstable_by(|a, b| cmp(&a.1, &b.1)); + } + + for (i, &(orig_idx, val)) in pairs.iter().enumerate() { + let offset = base + i * dim_stride; + data[offset] = val; + indices[offset] = orig_idx as isize; + } + } +} + +/// Argsort along a dimension without writing sorted values. +fn argsort_along_dim( + data: &[E], + indices: &mut [isize], + shape: &Shape, + dim: usize, + descending: bool, + cmp: fn(&E, &E) -> core::cmp::Ordering, +) { + let strides = contiguous_strides(shape); + let dim_size = shape[dim]; + let dim_stride = strides[dim]; + let num_slices = data.len() / dim_size; + + // Fast path: last dimension (stride==1). Both input rows and + // output index rows are contiguous at `slice_idx * dim_size`. + if dim_stride == 1 { + // zip silently truncates to the shorter iterator and chunks_exact + // silently drops remainders, so lock both invariants. + debug_assert_eq!(data.len(), indices.len()); + debug_assert_eq!(data.len() % dim_size, 0); + // Buffer is reused across rows (one per thread under rayon via + // `for_each_init`) to avoid a heap alloc per row. + let sort_row = |idx_buf: &mut Vec, (row, idx_row): (&[E], &mut [isize])| { + idx_buf.clear(); + idx_buf.extend(0..dim_size); + if descending { + idx_buf.sort_unstable_by(|&a, &b| cmp(&row[b], &row[a])); + } else { + idx_buf.sort_unstable_by(|&a, &b| cmp(&row[a], &row[b])); + } + for (i, &orig_idx) in idx_buf.iter().enumerate() { + idx_row[i] = orig_idx as isize; + } + }; + + #[cfg(feature = "rayon")] + if data.len() >= PARALLEL_THRESHOLD { + data.par_chunks_exact(dim_size) + .zip(indices.par_chunks_exact_mut(dim_size)) + .for_each_init(|| Vec::with_capacity(dim_size), sort_row); + return; + } + + let mut idx_buf: Vec = Vec::with_capacity(dim_size); + data.chunks_exact(dim_size) + .zip(indices.chunks_exact_mut(dim_size)) + .for_each(|row_and_idx| sort_row(&mut idx_buf, row_and_idx)); + return; + } + + let mut idx_buf: Vec = (0..dim_size).collect(); + + for slice_idx in 0..num_slices { + let base = slice_base_offset(slice_idx, shape, &strides, dim); + + idx_buf.clear(); + idx_buf.extend(0..dim_size); + + if descending { + idx_buf.sort_unstable_by(|&a, &b| { + cmp(&data[base + b * dim_stride], &data[base + a * dim_stride]) + }); + } else { + idx_buf.sort_unstable_by(|&a, &b| { + cmp(&data[base + a * dim_stride], &data[base + b * dim_stride]) + }); + } + + for (i, &orig_idx) in idx_buf.iter().enumerate() { + indices[base + i * dim_stride] = orig_idx as isize; + } + } +} + +// --------------------------------------------------------------------------- +// Half-precision sort (convert to f32, sort, convert back) +// --------------------------------------------------------------------------- + +fn sort_half( + tensor: FlexTensor, + dim: usize, + descending: bool, + to_f32: fn(H) -> f32, + from_f32: fn(f32) -> H, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let dtype = tensor.dtype(); + if validate_sort_args(&shape, dim) { + return tensor; + } + let src: &[H] = tensor.storage(); + let mut f32_data: Vec = src.iter().map(|&v| to_f32(v)).collect(); + + if shape.num_dims() == 1 { + if descending { + f32_data.sort_unstable_by(|a, b| f32::total_cmp(b, a)); + } else { + f32_data.sort_unstable_by(f32::total_cmp); + } + } else { + sort_along_dim(&mut f32_data, &shape, dim, descending, f32::total_cmp); + } + + let result: Vec = f32_data.iter().map(|&v| from_f32(v)).collect(); + FlexTensor::new(Bytes::from_elems(result), Layout::contiguous(shape), dtype) +} + +fn sort_with_indices_half( + tensor: FlexTensor, + dim: usize, + descending: bool, + to_f32: fn(H) -> f32, + from_f32: fn(f32) -> H, +) -> (FlexTensor, FlexTensor) { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let dtype = tensor.dtype(); + let n = shape.num_elements(); + if validate_sort_args(&shape, dim) { + let idx = make_index_tensor(Vec::new(), shape.clone()); + return (tensor, idx); + } + let src: &[H] = tensor.storage(); + let mut f32_data: Vec = src.iter().map(|&v| to_f32(v)).collect(); + let mut indices: Vec = vec![0; n]; + + if shape.num_dims() == 1 { + sort_1d_with_indices(&mut f32_data, &mut indices, descending, f32::total_cmp); + } else { + sort_along_dim_with_indices( + &mut f32_data, + &mut indices, + &shape, + dim, + descending, + f32::total_cmp, + ); + } + + let result: Vec = f32_data.iter().map(|&v| from_f32(v)).collect(); + let val_tensor = FlexTensor::new( + Bytes::from_elems(result), + Layout::contiguous(shape.clone()), + dtype, + ); + let idx_tensor = make_index_tensor(indices, shape); + (val_tensor, idx_tensor) +} + +fn argsort_half( + tensor: FlexTensor, + dim: usize, + descending: bool, + to_f32: fn(H) -> f32, +) -> FlexTensor { + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let n = shape.num_elements(); + if validate_sort_args(&shape, dim) { + return make_index_tensor(Vec::new(), shape); + } + let src: &[H] = tensor.storage(); + let f32_data: Vec = src.iter().map(|&v| to_f32(v)).collect(); + let mut indices: Vec = vec![0; n]; + + if shape.num_dims() == 1 { + let mut idx_vec: Vec = (0..n).collect(); + if descending { + idx_vec.sort_unstable_by(|&a, &b| f32::total_cmp(&f32_data[b], &f32_data[a])); + } else { + idx_vec.sort_unstable_by(|&a, &b| f32::total_cmp(&f32_data[a], &f32_data[b])); + } + for (out_i, &orig_i) in idx_vec.iter().enumerate() { + indices[out_i] = orig_i as isize; + } + } else { + argsort_along_dim( + &f32_data, + &mut indices, + &shape, + dim, + descending, + f32::total_cmp, + ); + } + + make_index_tensor(indices, shape) +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +fn contiguous_strides(shape: &Shape) -> Vec { + crate::layout::contiguous_strides_usize(shape) +} + +fn slice_base_offset(slice_idx: usize, shape: &Shape, strides: &[usize], dim: usize) -> usize { + crate::layout::slice_base_offset(slice_idx, shape, strides, dim) +} + +fn make_index_tensor(indices: Vec, shape: Shape) -> FlexTensor { + let bytes = Bytes::from_elems(indices); + FlexTensor::new(bytes, Layout::contiguous(shape), INDEX_DTYPE) +} + +// Tests kept here exercise flex-specific behavior: the internal +// `sort_along_dim` / `sort_along_dim_with_indices` / `argsort_along_dim` +// helpers at sizes that straddle `PARALLEL_THRESHOLD`, so both the serial +// and rayon-parallel branches are covered. End-to-end sort/argsort +// correctness across backends lives in +// crates/burn-backend-tests/tests/tensor/float/ops/sort_argsort.rs. +#[cfg(test)] +mod tests { + use super::*; + + // Exercise both below and above the parallel threshold so serial and + // rayon fast paths in sort_along_dim agree on row-wise sort results. + fn check_sort_last_dim(rows: usize, cols: usize) { + let n = rows * cols; + // Deterministic non-monotonic input with repeats (mirrors bench fill % 1000). + let src: Vec = (0..n) + .map(|i| ((i * 1664525 + 1013904223) % 1000) as f32) + .collect(); + + let mut data = src.clone(); + let shape = Shape::new([rows, cols]); + sort_along_dim(&mut data, &shape, 1, false, f32::total_cmp); + + for r in 0..rows { + let row = &data[r * cols..(r + 1) * cols]; + for w in row.windows(2) { + assert!(w[0] <= w[1], "row {r} not sorted: {:?}", row); + } + let mut expected: Vec = src[r * cols..(r + 1) * cols].to_vec(); + expected.sort_unstable_by(f32::total_cmp); + assert_eq!(row, expected.as_slice()); + } + } + + #[test] + fn sort_along_last_dim_small_serial() { + // 64*64 = 4K elements, well under PARALLEL_THRESHOLD. + check_sort_last_dim(64, 64); + } + + #[cfg(feature = "rayon")] + #[test] + fn sort_along_last_dim_large_parallel() { + // Just above PARALLEL_THRESHOLD so the rayon path is exercised + // without paying for a much larger input under debug test builds. + let cols = 1024; + let rows = (PARALLEL_THRESHOLD / cols) + 1; + check_sort_last_dim(rows, cols); + } + + #[test] + fn sort_along_last_dim_descending() { + let mut data: Vec = (0..4096).map(|i| (i % 17) as f32).collect(); + let shape = Shape::new([128, 32]); + sort_along_dim(&mut data, &shape, 1, true, f32::total_cmp); + for r in 0..128 { + let row = &data[r * 32..(r + 1) * 32]; + for w in row.windows(2) { + assert!(w[0] >= w[1]); + } + } + } + + fn check_sort_with_indices_last_dim(rows: usize, cols: usize, descending: bool) { + let src: Vec = (0..rows * cols).map(|i| (i as f32 * 0.37).sin()).collect(); + let mut values = src.clone(); + let mut indices = vec![0isize; rows * cols]; + let shape = Shape::new([rows, cols]); + sort_along_dim_with_indices( + &mut values, + &mut indices, + &shape, + 1, + descending, + f32::total_cmp, + ); + for r in 0..rows { + let vs = &values[r * cols..(r + 1) * cols]; + let idx_row = &indices[r * cols..(r + 1) * cols]; + let orig = &src[r * cols..(r + 1) * cols]; + let want_order = if descending { + core::cmp::Ordering::Less + } else { + core::cmp::Ordering::Greater + }; + for w in vs.windows(2) { + assert_ne!(f32::total_cmp(&w[0], &w[1]), want_order); + } + // Indices must reconstruct the sorted values from the original row + // and must be a valid permutation of 0..cols (each index appears once). + let mut seen = vec![false; cols]; + for (i, &orig_idx) in idx_row.iter().enumerate() { + let j = orig_idx as usize; + assert_eq!(vs[i], orig[j]); + assert!(!seen[j], "row {r}: index {j} repeated"); + seen[j] = true; + } + } + } + + #[test] + fn sort_with_indices_last_dim_ascending() { + // 512*512 = 256K, exactly PARALLEL_THRESHOLD (hits parallel branch). + check_sort_with_indices_last_dim(512, 512, false); + } + + #[test] + fn sort_with_indices_last_dim_descending() { + check_sort_with_indices_last_dim(512, 512, true); + } + + fn check_argsort_last_dim(rows: usize, cols: usize, descending: bool) { + let src: Vec = (0..rows * cols) + .map(|i| ((i * 7919) % 997) as f32) + .collect(); + let mut indices = vec![0isize; rows * cols]; + let shape = Shape::new([rows, cols]); + argsort_along_dim(&src, &mut indices, &shape, 1, descending, f32::total_cmp); + for r in 0..rows { + let idx_row = &indices[r * cols..(r + 1) * cols]; + let orig = &src[r * cols..(r + 1) * cols]; + let sorted: Vec = idx_row.iter().map(|&i| orig[i as usize]).collect(); + for w in sorted.windows(2) { + if descending { + assert!(w[0] >= w[1]); + } else { + assert!(w[0] <= w[1]); + } + } + // idx_row must be a permutation of 0..cols; with heavy duplicates + // in src this catches parallel bugs that emit the same index twice. + let mut seen = vec![false; cols]; + for &i in idx_row { + let j = i as usize; + assert!(!seen[j], "row {r}: index {j} repeated"); + seen[j] = true; + } + } + } + + #[test] + fn argsort_last_dim_ascending() { + // 200*1500 = 300K, above PARALLEL_THRESHOLD. + check_argsort_last_dim(200, 1500, false); + } + + #[test] + fn argsort_last_dim_descending() { + check_argsort_last_dim(200, 1500, true); + } +} diff --git a/crates/burn-flex/src/ops/transaction.rs b/crates/burn-flex/src/ops/transaction.rs new file mode 100644 index 0000000..ec2227c --- /dev/null +++ b/crates/burn-flex/src/ops/transaction.rs @@ -0,0 +1,11 @@ +//! Transaction operations for the Flex backend. + +use crate::Flex; +use burn_backend::distributed::DistributedOps; +use burn_backend::ops::TransactionOps; + +// TransactionOps has default implementations. +impl TransactionOps for Flex {} + +// DistributedOps has default implementations; Flex does not support collective operations. +impl DistributedOps for Flex {} diff --git a/crates/burn-flex/src/ops/unary.rs b/crates/burn-flex/src/ops/unary.rs new file mode 100644 index 0000000..09f7b50 --- /dev/null +++ b/crates/burn-flex/src/ops/unary.rs @@ -0,0 +1,463 @@ +//! Unary tensor operations (exp, log, sqrt, sin, cos, etc.). + +use alloc::vec::Vec; +use burn_backend::DType; +use burn_std::{BoolDType, Bytes, bf16, f16}; +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float; + +use crate::layout::StridedBlocks; +use crate::{FlexTensor, Layout}; + +/// Apply a float predicate element-wise, producing a boolean tensor. +/// +/// Delegates to make_bool_tensor for output +/// construction, so it shares the same `BoolDType` support (Native/U8 only, +/// panics on U32). +pub fn float_predicate( + tensor: FlexTensor, + out_dtype: BoolDType, + f32_pred: F32P, + f64_pred: F64P, +) -> FlexTensor +where + F32P: Fn(f32) -> bool + Copy, + F64P: Fn(f64) -> bool + Copy, +{ + let tensor = tensor.to_contiguous(); + let shape = tensor.layout().shape().clone(); + let n = shape.num_elements(); + + let result: Vec = match tensor.dtype() { + DType::F32 => { + let s: &[f32] = tensor.storage(); + s[..n].iter().map(|&x| f32_pred(x) as u8).collect() + } + DType::F64 => { + let s: &[f64] = tensor.storage(); + s[..n].iter().map(|&x| f64_pred(x) as u8).collect() + } + DType::F16 => { + let s: &[f16] = tensor.storage(); + s[..n].iter().map(|&x| f32_pred(x.to_f32()) as u8).collect() + } + DType::BF16 => { + let s: &[bf16] = tensor.storage(); + s[..n].iter().map(|&x| f32_pred(x.to_f32()) as u8).collect() + } + dt => panic!("float_predicate: expected float dtype, got {:?}", dt), + }; + + crate::ops::comparison::make_bool_tensor(result, shape, out_dtype) +} + +/// Apply a unary operation element-wise to a tensor. +pub fn unary_op(tensor: FlexTensor, f32_op: F32Op, f64_op: F64Op) -> FlexTensor +where + F32Op: Fn(f32) -> f32 + Copy, + F64Op: Fn(f64) -> f64 + Copy, +{ + let dtype = tensor.dtype(); + + match dtype { + DType::F32 => unary_op_typed(tensor, f32_op), + DType::F64 => unary_op_typed(tensor, f64_op), + DType::F16 => unary_op_typed(tensor, |x: f16| f16::from_f32(f32_op(x.to_f32()))), + DType::BF16 => unary_op_typed(tensor, |x: bf16| bf16::from_f32(f32_op(x.to_f32()))), + _ => panic!("unary_op: unsupported dtype {:?}", dtype), + } +} + +/// Generic unary operation for any element type. +fn unary_op_typed(mut tensor: FlexTensor, op: Op) -> FlexTensor +where + E: burn_backend::Element + bytemuck::Pod, + Op: Fn(E) -> E, +{ + let n = tensor.layout().num_elements(); + + // In-place fast path: unique, contiguous tensor at offset 0 + if tensor.is_unique() && tensor.layout().is_contiguous() && tensor.layout().start_offset() == 0 + { + let storage: &mut [E] = tensor.storage_mut(); + for x in storage[..n].iter_mut() { + *x = op(*x); + } + return tensor; + } + + // Allocating path for non-contiguous or offset tensors + let layout = tensor.layout().clone(); + let src: &[E] = tensor.storage(); + + // Check for negative strides (from flip operations) + let has_negative_strides = layout.strides().iter().any(|&s| s < 0); + + // Fast path: storage exactly matches tensor view (covers transposed tensors) + // Iterate in storage order (contiguous) and preserve original layout. + // Only valid when all strides are positive. + if !has_negative_strides && layout.start_offset() == 0 && src.len() == n { + let result: Vec = src.iter().map(|&x| op(x)).collect(); + let bytes = Bytes::from_elems(result); + return FlexTensor::new(bytes, layout, E::dtype()); + } + + // Fallback for negative strides: use StridedIter for correct element order + if has_negative_strides { + let result: Vec = crate::strided_index::StridedIter::new(&layout) + .map(|idx| op(src[idx])) + .collect(); + let bytes = Bytes::from_elems(result); + return FlexTensor::new( + bytes, + Layout::contiguous(layout.shape().clone()), + E::dtype(), + ); + } + + // General path for views/slices with offset or extra storage + let blocks = layout.strided_blocks(); + let result = match &blocks { + // Single contiguous block (with offset) + StridedBlocks::Single { start, len } => { + src[*start..*start + *len].iter().map(|&x| op(x)).collect() + } + // Strided: iterate over contiguous blocks + StridedBlocks::Multiple { + block_len, + num_blocks, + .. + } => { + let block_len = *block_len; + let num_blocks = *num_blocks; + let mut result = Vec::with_capacity(n); + + if block_len == 1 { + for block_start in blocks.block_starts() { + result.push(op(src[block_start])); + } + } else { + for block_start in blocks.block_starts() { + for i in 0..block_len { + result.push(op(src[block_start + i])); + } + } + } + debug_assert_eq!(result.len(), num_blocks * block_len); + result + } + }; + + let bytes = Bytes::from_elems(result); + FlexTensor::new( + bytes, + Layout::contiguous(layout.shape().clone()), + E::dtype(), + ) +} + +// ============================================================================ +// Specific unary operations +// ============================================================================ + +/// Exponential: e^x +pub fn exp(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.exp(), |x| x.exp()) +} + +/// Natural logarithm: ln(x) +pub fn log(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.ln(), |x| x.ln()) +} + +/// Natural logarithm of (1 + x): ln(1 + x) +pub fn log1p(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.ln_1p(), |x| x.ln_1p()) +} + +/// Square root +pub fn sqrt(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.sqrt(), |x| x.sqrt()) +} + +/// Absolute value (float) +pub fn abs(tensor: FlexTensor) -> FlexTensor { + #[cfg(feature = "simd")] + if tensor.dtype() == DType::F32 + && tensor.is_unique() + && tensor.layout().is_contiguous() + && tensor.layout().start_offset() == 0 + { + let n = tensor.layout().num_elements(); + let mut tensor = tensor; + let storage: &mut [f32] = tensor.storage_mut(); + crate::simd::abs_inplace_f32(&mut storage[..n]); + return tensor; + } + unary_op(tensor, |x| x.abs(), |x| x.abs()) +} + +/// Absolute value (integer) +pub fn int_abs(tensor: FlexTensor) -> FlexTensor { + let dtype = tensor.dtype(); + match dtype { + DType::I64 => unary_op_typed::(tensor, |x| x.wrapping_abs()), + DType::I32 => unary_op_typed::(tensor, |x| x.wrapping_abs()), + DType::I16 => unary_op_typed::(tensor, |x| x.wrapping_abs()), + DType::I8 => unary_op_typed::(tensor, |x| x.wrapping_abs()), + // Unsigned integers: abs is identity + DType::U64 | DType::U32 | DType::U16 | DType::U8 => tensor, + _ => panic!("int_abs: unsupported dtype {:?}", dtype), + } +} + +/// Reciprocal: 1/x +pub fn recip(tensor: FlexTensor) -> FlexTensor { + #[cfg(feature = "simd")] + if tensor.dtype() == DType::F32 + && tensor.is_unique() + && tensor.layout().is_contiguous() + && tensor.layout().start_offset() == 0 + { + let n = tensor.layout().num_elements(); + let mut tensor = tensor; + let storage: &mut [f32] = tensor.storage_mut(); + crate::simd::recip_inplace_f32(&mut storage[..n]); + return tensor; + } + unary_op(tensor, |x| 1.0 / x, |x| 1.0 / x) +} + +/// Cosine +pub fn cos(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.cos(), |x| x.cos()) +} + +/// Sine +pub fn sin(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.sin(), |x| x.sin()) +} + +/// Tangent +pub fn tan(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.tan(), |x| x.tan()) +} + +/// Hyperbolic cosine +pub fn cosh(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.cosh(), |x| x.cosh()) +} + +/// Hyperbolic sine +pub fn sinh(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.sinh(), |x| x.sinh()) +} + +/// Hyperbolic tangent +pub fn tanh(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.tanh(), |x| x.tanh()) +} + +/// Inverse cosine (arccos) +pub fn acos(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.acos(), |x| x.acos()) +} + +/// Inverse hyperbolic cosine +pub fn acosh(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.acosh(), |x| x.acosh()) +} + +/// Inverse sine (arcsin) +pub fn asin(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.asin(), |x| x.asin()) +} + +/// Inverse hyperbolic sine +pub fn asinh(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.asinh(), |x| x.asinh()) +} + +/// Inverse tangent (arctan) +pub fn atan(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.atan(), |x| x.atan()) +} + +/// Inverse hyperbolic tangent +pub fn atanh(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.atanh(), |x| x.atanh()) +} + +/// Round to nearest integer (ties to even / banker's rounding) +pub fn round(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, round_ties_even_f32, round_ties_even_f64) +} + +fn round_ties_even_f32(x: f32) -> f32 { + round_ties_even(x) +} + +fn round_ties_even_f64(x: f64) -> f64 { + round_ties_even(x) +} + +/// Round to nearest integer, ties to even (banker's rounding). +/// +/// `num_traits::Float::round` rounds ties away from zero. This corrects +/// the halfway case to round to the nearest even integer instead. +/// +/// Safety of the `to_i64` path: for f32, values with magnitude >= 2^23 +/// have no fractional bits, so `(x - r).abs()` is always 0.0 (never 0.5). +/// For f64, the threshold is 2^52. The halfway check therefore only +/// triggers for values well within i64 range. +fn round_ties_even(x: F) -> F { + let r = x.round(); + if (x - r).abs() == F::from(0.5).unwrap() { + // Ties: round to even by checking if r is odd + match r.to_i64() { + Some(ri) if ri % 2 != 0 => r - x.signum(), + _ => r, + } + } else { + r + } +} + +/// Floor (round down) +pub fn floor(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.floor(), |x| x.floor()) +} + +/// Ceiling (round up) +pub fn ceil(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.ceil(), |x| x.ceil()) +} + +/// Truncate (round towards zero) +pub fn trunc(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, |x| x.trunc(), |x| x.trunc()) +} + +/// Error function +pub fn erf(tensor: FlexTensor) -> FlexTensor { + unary_op(tensor, erf_f32, erf_f64) +} + +// ============================================================================ +// Error function implementation +// ============================================================================ +// +// Delegates to libm for full f32 / f64 precision. The previous Abramowitz +// and Stegun 7.1.26 formula was capped at ~1.5e-7 max absolute error, +// adequate for f32 but well short of f64's ~2.2e-16 precision. libm's erf +// is derived from fdlibm and lands within a few ulps across the full +// range in both precisions. + +/// Error function for f32. +pub fn erf_f32(x: f32) -> f32 { + libm::erff(x) +} + +/// Error function for f64. +pub fn erf_f64(x: f64) -> f64 { + libm::erf(x) +} + +// Tests kept here probe flex-internal unary helpers that don't flow through +// the generic FloatTensorOps API. Specifically, the manual round_ties_even +// implementation has a subtle to_i64 conversion path that must not trip on +// values outside i64 range. Plain elementwise unary smokes (exp, log, sqrt, +// abs, sin/cos, tanh, round/floor/ceil, erf) and their stride-through-op +// variants (transposed/flipped/narrowed/sliced) have been migrated to +// burn-backend-tests so they run against every backend. When adding new +// tests, keep them here only if they probe flex-internal helpers +// (erf_f32/f64, round_ties_even); otherwise add them to +// crates/burn-backend-tests/tests/tensor/float/ops/. +#[cfg(test)] +#[allow(clippy::excessive_precision)] +mod tests { + use alloc::vec; + use burn_backend::TensorData; + + use crate::FlexTensor; + + #[test] + fn test_round_ties_even_large_float() { + // Values outside i32 range used to break the manual round_ties_even impl. + let data = vec![2e18_f32, -2e18_f32, f32::MAX, f32::MIN]; + let tensor = FlexTensor::from_data(TensorData::new(data.clone(), [data.len()])); + let result = super::round(tensor); + let out: Vec = result.into_data().to_vec().unwrap(); + for (a, b) in out.iter().zip(data.iter()) { + assert_eq!(a.to_bits(), b.to_bits()); + } + } + + // The previous A&S 7.1.26 erf implementation had ~1.5e-7 max absolute + // error regardless of arithmetic precision. That was acceptable for + // f32 but orders of magnitude worse than f64's ~2.2e-16 precision. + // These tests guard near-ulp accuracy across the domain, including + // libm's piecewise-rational switch points near |x|=0.84375 and the + // saturation region around |x|>=6 where erf(x) rounds to ±1. + // Reference values from Wolfram Alpha / DLMF. + #[test] + fn test_erf_f64_is_full_precision() { + // Reference values computed with mpmath at 25-digit precision. + let cases = [ + (0.0f64, 0.0f64), + (1e-10, 1.128379167095512574e-10), // small-x linear regime + (0.5, 0.5204998778130465377), + (0.84, 0.7651427114549945347), // just below libm's 0.84375 boundary + (0.85, 0.7706680576083525324), // just above + (1.0, 0.8427007929497148693), + (1.5, 0.9661051464753107271), + (2.0, 0.9953222650189527342), + (3.0, 0.9999779095030014145), + (6.0, 0.9999999999999999785), // near saturation + (30.0, 1.0), // fully saturated + (-0.5, -0.5204998778130465377), + (-1.0, -0.8427007929497148693), + (-3.0, -0.9999779095030014145), + ]; + for (x, expected) in cases { + let got = super::erf_f64(x); + let err = (got - expected).abs(); + assert!( + err < 1e-14, + "erf_f64({}) = {} expected {} (err {:e})", + x, + got, + expected, + err + ); + } + } + + // f32 erf must stay accurate to within f32 precision (~6e-8). Covers + // the same piecewise boundaries as the f64 test. + #[test] + fn test_erf_f32_full_f32_precision() { + let cases = [ + (0.0f32, 0.0f32), + (0.5, 0.520_499_9), + (0.84, 0.765_142_7), + (0.85, 0.770_668_06), + (1.0, 0.842_700_8), + (2.0, 0.995_322_24), + (6.0, 1.0), + (-0.5, -0.520_499_9), + (-1.0, -0.842_700_8), + ]; + for (x, expected) in cases { + let got = super::erf_f32(x); + assert!( + (got - expected).abs() < 1e-6, + "erf_f32({}) = {} expected {}", + x, + got, + expected + ); + } + } +} diff --git a/crates/burn-flex/src/ops/unfold.rs b/crates/burn-flex/src/ops/unfold.rs new file mode 100644 index 0000000..c90b0c0 --- /dev/null +++ b/crates/burn-flex/src/ops/unfold.rs @@ -0,0 +1,118 @@ +//! Unfold operation for sliding window extraction. +//! +//! Implemented as a zero-copy strided view. The output tensor shares storage +//! with the input, using strides to represent overlapping windows. + +use alloc::vec::Vec; + +use burn_std::Shape; + +use crate::{FlexTensor, Layout}; + +/// Calculate the number of windows that can be extracted from a dimension. +#[inline] +fn calculate_windows(dim_size: usize, window_size: usize, step: usize) -> usize { + assert!(step > 0, "step must be positive"); + if dim_size + step < window_size { + 0 + } else { + (dim_size + step - window_size) / step + } +} + +/// Unfold: extract sliding windows from a tensor along a dimension. +/// +/// Given a tensor with shape `[pre..., dim_size, post...]`, extracts windows of +/// `size` elements along dimension `dim`, stepping by `step`. +/// +/// Returns a tensor with shape `[pre..., windows, post..., size]` where: +/// - `windows = (dim_size - size + step) / step` +/// - The `size` dimension is appended at the end +/// +/// This is a zero-copy operation that returns a strided view of the input. +/// The same storage elements may appear in multiple windows (overlapping). +pub fn unfold(tensor: FlexTensor, dim: usize, size: usize, step: usize) -> FlexTensor { + let input_layout = tensor.layout(); + let shape = input_layout.shape(); + let input_strides = input_layout.strides(); + let start_offset = input_layout.start_offset(); + let ndims = shape.num_dims(); + let dtype = tensor.dtype(); + + assert!( + dim < ndims, + "dim {} out of bounds for {} dimensions", + dim, + ndims + ); + assert!(size > 0, "window size must be positive"); + assert!(step > 0, "step must be positive"); + assert!( + shape[dim] >= size, + "dimension {} has size {} which is smaller than window size {}", + dim, + shape[dim], + size + ); + + let dim_size = shape[dim]; + let windows = calculate_windows(dim_size, size, step); + + // Build output shape: [pre..., windows, post..., size] + let mut output_dims: Vec = Vec::with_capacity(ndims + 1); + for (d, &s) in shape.iter().enumerate() { + if d == dim { + output_dims.push(windows); + } else { + output_dims.push(s); + } + } + output_dims.push(size); // Append size at the end + + // Build output strides: + // - Dimensions before `dim`: same stride as input + // - Dimension `dim` (now windows): input_stride[dim] * step + // - Dimensions after `dim`: same stride as input + // - New size dimension (appended): input_stride[dim] + let mut output_strides: Vec = Vec::with_capacity(ndims + 1); + for (d, &s) in input_strides.iter().enumerate() { + if d == dim { + // Windows dimension: stride = original_stride * step + output_strides.push(s * step as isize); + } else { + output_strides.push(s); + } + } + // Append stride for the size dimension (position within window) + output_strides.push(input_strides[dim]); + + let output_shape = Shape::from(output_dims); + let output_layout = Layout::new(output_shape, output_strides, start_offset); + + // Zero-copy: reuse the same storage with new layout + FlexTensor::from_arc(tensor.data_arc(), output_layout, dtype) +} + +// Type-specific wrappers (all delegate to the generic unfold which is now type-agnostic) + +pub fn unfold_f32(tensor: FlexTensor, dim: usize, size: usize, step: usize) -> FlexTensor { + unfold(tensor, dim, size, step) +} + +pub fn unfold_f64(tensor: FlexTensor, dim: usize, size: usize, step: usize) -> FlexTensor { + unfold(tensor, dim, size, step) +} + +pub fn unfold_bool(tensor: FlexTensor, dim: usize, size: usize, step: usize) -> FlexTensor { + unfold(tensor, dim, size, step) +} + +pub fn unfold_int(tensor: FlexTensor, dim: usize, size: usize, step: usize) -> FlexTensor { + unfold(tensor, dim, size, step) +} + +// Correctness of unfold across dtypes and shapes is covered by the +// cross-backend suite in +// crates/burn-backend-tests/tests/tensor/{float,int,bool}/ops/unfold.rs, +// which exercises the flex backend through the public `unfold` op. No +// flex-specific tests remain here. diff --git a/crates/burn-flex/src/qtensor.rs b/crates/burn-flex/src/qtensor.rs new file mode 100644 index 0000000..aa95db5 --- /dev/null +++ b/crates/burn-flex/src/qtensor.rs @@ -0,0 +1,77 @@ +use alloc::vec::Vec; + +use burn_backend::{DType, TensorMetadata}; +use burn_std::{QuantScheme, Shape}; + +use crate::{FlexDevice, tensor::FlexTensor}; + +/// Quantized tensor for the Flex backend. +/// +/// Stores quantized i8 values in the tensor and keeps scales separately +/// for efficient dequantization without reparsing bytes. +#[derive(Clone, Debug)] +pub struct FlexQTensor { + /// The underlying quantized data (stored as i8). + pub(crate) tensor: FlexTensor, + /// Quantization scheme. + pub(crate) scheme: QuantScheme, + /// Per-tensor or per-block scale factors. + pub(crate) scales: Vec, +} + +impl FlexQTensor { + /// Create a new quantized tensor. + /// + /// The tensor must store i8 data and scales must be non-empty. + pub fn new(tensor: FlexTensor, scheme: QuantScheme, scales: Vec) -> Self { + assert_eq!( + tensor.dtype(), + DType::I8, + "quantized tensor must store i8 data, got {:?}", + tensor.dtype() + ); + assert!( + !scales.is_empty(), + "quantized tensor must have at least one scale factor" + ); + Self { + tensor, + scheme, + scales, + } + } + + /// Get the underlying tensor. + pub fn tensor(&self) -> &FlexTensor { + &self.tensor + } + + /// Get the quantization scales. + pub fn scales(&self) -> &[f32] { + &self.scales + } +} + +impl TensorMetadata for FlexQTensor { + type Device = FlexDevice; + + fn dtype(&self) -> DType { + DType::QFloat(self.scheme) + } + + fn shape(&self) -> Shape { + self.tensor.shape() + } + + fn rank(&self) -> usize { + self.tensor.rank() + } + + fn device(&self) -> Self::Device { + FlexDevice + } + + fn can_mut(&self) -> bool { + self.tensor.is_unique() + } +} diff --git a/crates/burn-flex/src/simd/aligned.rs b/crates/burn-flex/src/simd/aligned.rs new file mode 100644 index 0000000..fd6b3c0 --- /dev/null +++ b/crates/burn-flex/src/simd/aligned.rs @@ -0,0 +1,70 @@ +//! SIMD-aligned memory allocation utilities. +//! +//! Provides aligned vectors for optimal SIMD performance. Uses 64-byte alignment +//! to support all SIMD instruction sets (NEON: 16B, AVX2: 32B, AVX-512: 64B). + +use aligned_vec::{AVec, ConstAlign}; +use alloc::vec::Vec; + +/// Alignment for SIMD operations (64 bytes = AVX-512 width). +/// Also works for NEON (16B) and AVX2 (32B) since 64 is a multiple of both. +pub const SIMD_ALIGN: usize = 64; + +/// Type alias for 64-byte aligned vector. +pub type AlignedVec = AVec>; + +/// Allocate a SIMD-aligned vector with the given capacity. +#[inline] +pub fn alloc_aligned(capacity: usize) -> AlignedVec { + AVec::with_capacity(SIMD_ALIGN, capacity) +} + +/// Allocate a SIMD-aligned vector filled with zeros. +#[inline] +pub fn alloc_aligned_zeroed(len: usize) -> AlignedVec { + let mut vec = AVec::with_capacity(SIMD_ALIGN, len); + vec.resize(len, T::zeroed()); + vec +} + +/// Convert an aligned vector to a regular Vec. +/// +/// This may copy the data if the alignment requirements differ. +#[inline] +pub fn to_vec(aligned: AlignedVec) -> Vec { + aligned.to_vec() +} + +/// Convert a slice to an aligned vector (copies data). +#[inline] +pub fn from_slice(slice: &[T]) -> AlignedVec { + let mut vec = AVec::with_capacity(SIMD_ALIGN, slice.len()); + vec.extend_from_slice(slice); + vec +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_alignment() { + let vec: AlignedVec = alloc_aligned(1000); + assert_eq!(vec.as_ptr() as usize % SIMD_ALIGN, 0); + } + + #[test] + fn test_zeroed() { + let vec: AlignedVec = alloc_aligned_zeroed(100); + assert_eq!(vec.len(), 100); + assert!(vec.iter().all(|&x| x == 0.0)); + } + + #[test] + fn test_from_slice() { + let data = [1.0f32, 2.0, 3.0, 4.0]; + let aligned = from_slice(&data); + assert_eq!(aligned.as_ptr() as usize % SIMD_ALIGN, 0); + assert_eq!(&aligned[..], &data[..]); + } +} diff --git a/crates/burn-flex/src/simd/kernels.rs b/crates/burn-flex/src/simd/kernels.rs new file mode 100644 index 0000000..43429d6 --- /dev/null +++ b/crates/burn-flex/src/simd/kernels.rs @@ -0,0 +1,317 @@ +//! Portable SIMD kernels using macerator. +//! +//! These kernels work across all architectures supported by macerator: +//! - aarch64 (NEON) +//! - x86_64 (AVX2, AVX512, SSE) +//! - wasm32 (SIMD128) +//! - Scalar fallback for embedded/other platforms + +use core::iter::Sum; +use core::ops::AddAssign; + +use macerator::{ + ReduceAdd, ReduceMax, ReduceMin, Simd, VAdd, VOrd, vload_unaligned, vstore_unaligned, +}; + +// ============================================================================ +// Sum reduction +// ============================================================================ + +/// Sum all elements in a f32 slice using SIMD with 4 accumulators. +#[inline] +pub fn sum_f32(data: &[f32]) -> f32 { + macerator_sum(data) +} + +/// 8-accumulator SIMD sum. Independent accumulator chains let the CPU +/// pipeline floating-point adds and hide L2 cache latency. +#[macerator::with_simd] +fn macerator_sum(mut xs: &[F]) -> F { + let lanes = F::lanes::(); + let stride = lanes * 8; + let zero = F::default().splat::(); + let (mut s0, mut s1, mut s2, mut s3) = (zero, zero, zero, zero); + let (mut s4, mut s5, mut s6, mut s7) = (zero, zero, zero, zero); + + while xs.len() >= stride { + unsafe { + let p = xs.as_ptr(); + s0 += vload_unaligned(p); + s1 += vload_unaligned(p.add(lanes)); + s2 += vload_unaligned(p.add(lanes * 2)); + s3 += vload_unaligned(p.add(lanes * 3)); + s4 += vload_unaligned(p.add(lanes * 4)); + s5 += vload_unaligned(p.add(lanes * 5)); + s6 += vload_unaligned(p.add(lanes * 6)); + s7 += vload_unaligned(p.add(lanes * 7)); + } + xs = &xs[stride..]; + } + + // Combine 8 accumulators into one, then drain remaining full vectors + let mut sum = ((s0 + s1) + (s2 + s3)) + ((s4 + s5) + (s6 + s7)); + while xs.len() >= lanes { + sum += unsafe { vload_unaligned(xs.as_ptr()) }; + xs = &xs[lanes..]; + } + + sum.reduce_add() + xs.iter().copied().sum() +} + +// ============================================================================ +// Scatter-add for dimension reductions +// ============================================================================ + +/// Scatter-add: for each row, add to corresponding output positions. +/// Used for cache-friendly first-dim and middle-dim reductions. +/// +/// # Arguments +/// * `src` - Source data pointer +/// * `dst` - Destination accumulator (must be pre-zeroed) +/// * `num_rows` - Number of rows to sum +/// * `row_len` - Length of each row (columns) +/// * `src_row_stride` - Stride between source rows +#[macerator::with_simd] +pub fn scatter_add_f32( + src: &[F], + dst: &mut [F], + num_rows: usize, + row_len: usize, + src_row_stride: usize, +) { + let lanes = F::lanes::(); + + for row in 0..num_rows { + let row_start = row * src_row_stride; + let row_data = &src[row_start..row_start + row_len]; + + let simd_len = row_len / lanes * lanes; + + // SIMD accumulate + let mut i = 0; + while i < simd_len { + unsafe { + let s = vload_unaligned(row_data.as_ptr().add(i)); + let d = vload_unaligned(dst.as_ptr().add(i)); + vstore_unaligned::(dst.as_mut_ptr().add(i), d + s); + } + i += lanes; + } + + // Scalar tail + for j in simd_len..row_len { + dst[j] += row_data[j]; + } + } +} + +/// Batched scatter-add for middle-dim reductions. +/// For tensors like [B, M, K] reducing dim=1. +#[macerator::with_simd] +pub fn scatter_add_batched( + src: &[F], + dst: &mut [F], + num_batches: usize, + num_rows: usize, + row_len: usize, + batch_stride: usize, + row_stride: usize, +) { + let lanes = F::lanes::(); + + for batch in 0..num_batches { + let batch_src_start = batch * batch_stride; + let batch_dst_start = batch * row_len; + let batch_dst = &mut dst[batch_dst_start..batch_dst_start + row_len]; + + for row in 0..num_rows { + let row_start = batch_src_start + row * row_stride; + let row_data = &src[row_start..row_start + row_len]; + + let simd_len = row_len / lanes * lanes; + + let mut i = 0; + while i < simd_len { + unsafe { + let s = vload_unaligned(row_data.as_ptr().add(i)); + let d = vload_unaligned(batch_dst.as_ptr().add(i)); + vstore_unaligned::(batch_dst.as_mut_ptr().add(i), d + s); + } + i += lanes; + } + + for j in simd_len..row_len { + batch_dst[j] += row_data[j]; + } + } + } +} + +// ============================================================================ +// Row-wise sum (last-dim reduction) +// ============================================================================ + +/// Sum each row, storing results in output slice. +/// Used for last-dim reductions. +#[inline] +pub fn sum_rows_f32(src: &[f32], dst: &mut [f32], num_rows: usize, row_len: usize) { + debug_assert_eq!(dst.len(), num_rows, "dst length must equal num_rows"); + debug_assert!( + src.len() >= num_rows * row_len, + "src too short: need {} elements, got {}", + num_rows * row_len, + src.len() + ); + for (row, dst_val) in dst.iter_mut().enumerate() { + let row_start = row * row_len; + let row_data = &src[row_start..row_start + row_len]; + *dst_val = macerator_sum(row_data); + } +} + +// ============================================================================ +// Max/Min reduction +// ============================================================================ + +/// Find the maximum element in a f32 slice using SIMD. +#[inline] +pub fn max_f32(data: &[f32]) -> f32 { + macerator_max(data, f32::NEG_INFINITY) +} + +/// Find the minimum element in a f32 slice using SIMD. +#[inline] +pub fn min_f32(data: &[f32]) -> f32 { + macerator_min(data, f32::INFINITY) +} + +#[macerator::with_simd] +fn macerator_max(mut xs: &[F], init: F) -> F { + let lanes = F::lanes::(); + let mut acc = init.splat::(); + + while xs.len() >= lanes { + let v = unsafe { vload_unaligned(xs.as_ptr()) }; + acc = acc.max(v); + xs = &xs[lanes..]; + } + + let mut result = acc.reduce_max(); + for &x in xs { + if x > result { + result = x; + } + } + result +} + +#[macerator::with_simd] +fn macerator_min(mut xs: &[F], init: F) -> F { + let lanes = F::lanes::(); + let mut acc = init.splat::(); + + while xs.len() >= lanes { + let v = unsafe { vload_unaligned(xs.as_ptr()) }; + acc = acc.min(v); + xs = &xs[lanes..]; + } + + let mut result = acc.reduce_min(); + for &x in xs { + if x < result { + result = x; + } + } + result +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_sum_f32() { + let data: Vec = (0..1000).map(|i| i as f32).collect(); + let expected: f32 = data.iter().sum(); + let result = sum_f32(&data); + assert!((result - expected).abs() < 0.01); + } + + #[test] + fn test_sum_f32_empty() { + let data: Vec = vec![]; + assert_eq!(sum_f32(&data), 0.0); + } + + #[test] + fn test_sum_f32_small() { + let data = vec![1.0, 2.0, 3.0]; + assert_eq!(sum_f32(&data), 6.0); + } + + #[test] + fn test_scatter_add_f32() { + // Simulate reducing [3, 4] along dim=0 -> [1, 4] + let src = vec![ + 1.0, 2.0, 3.0, 4.0, // row 0 + 5.0, 6.0, 7.0, 8.0, // row 1 + 9.0, 10.0, 11.0, 12.0, // row 2 + ]; + let mut dst = vec![0.0; 4]; + + scatter_add_f32(&src, &mut dst, 3, 4, 4); + + assert_eq!(dst, vec![15.0, 18.0, 21.0, 24.0]); + } + + #[test] + fn test_sum_rows_f32() { + // Simulate reducing [3, 4] along dim=1 -> [3, 1] + let src = vec![ + 1.0, 2.0, 3.0, 4.0, // row 0 -> 10 + 5.0, 6.0, 7.0, 8.0, // row 1 -> 26 + 9.0, 10.0, 11.0, 12.0, // row 2 -> 42 + ]; + let mut dst = vec![0.0; 3]; + + sum_rows_f32(&src, &mut dst, 3, 4); + + assert_eq!(dst, vec![10.0, 26.0, 42.0]); + } + + #[test] + fn test_max_f32() { + let data: Vec = (0..1000).map(|i| i as f32).collect(); + assert_eq!(max_f32(&data), 999.0); + } + + #[test] + fn test_max_f32_small() { + let data = vec![3.0, 1.0, 4.0, 1.0, 5.0]; + assert_eq!(max_f32(&data), 5.0); + } + + #[test] + fn test_max_f32_negative() { + let data = vec![-3.0, -1.0, -4.0, -1.0, -5.0]; + assert_eq!(max_f32(&data), -1.0); + } + + #[test] + fn test_min_f32() { + let data: Vec = (0..1000).map(|i| i as f32).collect(); + assert_eq!(min_f32(&data), 0.0); + } + + #[test] + fn test_min_f32_small() { + let data = vec![3.0, 1.0, 4.0, 1.0, 5.0]; + assert_eq!(min_f32(&data), 1.0); + } + + #[test] + fn test_min_f32_negative() { + let data = vec![-3.0, -1.0, -4.0, -1.0, -5.0]; + assert_eq!(min_f32(&data), -5.0); + } +} diff --git a/crates/burn-flex/src/simd/mod.rs b/crates/burn-flex/src/simd/mod.rs new file mode 100644 index 0000000..9b5440c --- /dev/null +++ b/crates/burn-flex/src/simd/mod.rs @@ -0,0 +1,42 @@ +//! SIMD-optimized kernels for tensor operations. +//! +//! Provides portable SIMD implementations via `macerator` with automatic +//! dispatch to the best available instruction set: +//! - aarch64: NEON +//! - x86_64: AVX2, AVX512, SSE +//! - wasm32: SIMD128 +//! - Other: Scalar fallback +//! +//! Enable with the `simd` feature flag (enabled by default). + +// Portable SIMD kernels using macerator (reductions, scatter-add) +#[cfg(feature = "simd")] +pub mod kernels; + +// SIMD-aligned memory allocation +#[cfg(feature = "simd")] +pub mod aligned; + +// When simd feature enabled: use portable macerator for binary/comparison/bool ops +#[cfg(feature = "simd")] +mod portable; + +#[cfg(feature = "simd")] +pub use portable::{ + CmpOp, abs_inplace_f32, add_inplace_f32, add_shared_row_inplace_f32, bool_and_inplace_u8, + bool_and_u8, bool_not_inplace_u8, bool_not_u8, bool_or_inplace_u8, bool_or_u8, + bool_xor_inplace_u8, bool_xor_u8, cmp_f32, cmp_scalar_f32, div_inplace_f32, + div_shared_row_inplace_f32, mask_fill_f32, mask_fill_f64, mask_fill_i64, mask_fill_u8, + mask_where_f32, mask_where_f64, mask_where_i64, mask_where_u8, mul_inplace_f32, + mul_shared_row_inplace_f32, recip_inplace_f32, sub_inplace_f32, sub_shared_row_inplace_f32, +}; + +// When simd feature disabled: use scalar fallback (bool ops + CmpOp only) +#[cfg(not(feature = "simd"))] +mod scalar; + +#[cfg(not(feature = "simd"))] +pub use scalar::{ + CmpOp, bool_and_inplace_u8, bool_and_u8, bool_not_inplace_u8, bool_not_u8, bool_or_inplace_u8, + bool_or_u8, bool_xor_inplace_u8, bool_xor_u8, +}; diff --git a/crates/burn-flex/src/simd/portable.rs b/crates/burn-flex/src/simd/portable.rs new file mode 100644 index 0000000..4596777 --- /dev/null +++ b/crates/burn-flex/src/simd/portable.rs @@ -0,0 +1,1387 @@ +//! Portable SIMD kernels using macerator. +//! +//! Replaces platform-specific implementations (neon.rs) with a single +//! portable implementation that auto-dispatches to NEON/AVX2/SSE/SIMD128/scalar. + +use macerator::{Scalar, Simd, VBitAnd, VBitOr, VBitXor, vload_unaligned, vstore_unaligned}; + +#[cfg(feature = "rayon")] +use rayon::prelude::*; + +/// Threshold for parallel execution (elements). +/// For memory-bound operations, parallelism helps when data exceeds L3 cache. +#[cfg(feature = "rayon")] +const PARALLEL_THRESHOLD: usize = 4 * 1024 * 1024; + +#[cfg(feature = "rayon")] +const CHUNK_SIZE: usize = 4096; + +// ============================================================================ +// f32 in-place binary ops +// ============================================================================ + +macro_rules! define_inplace_f32_op { + ($pub_fn:ident, $seq_fn:ident, $par_fn:ident, $op:tt) => { + #[inline] + pub fn $pub_fn(a: &mut [f32], b: &[f32]) { + debug_assert_eq!(a.len(), b.len()); + + #[cfg(feature = "rayon")] + if a.len() >= PARALLEL_THRESHOLD { + $par_fn(a, b); + return; + } + + $seq_fn(a, b); + } + + #[macerator::with_simd] + #[allow(clippy::assign_op_pattern)] + fn $seq_fn(a: &mut [f32], b: &[f32]) { + let lanes = S::lanes32(); + let len = a.len(); + let simd_len = len / lanes * lanes; + + let mut i = 0; + while i < simd_len { + unsafe { + let va = vload_unaligned(a.as_ptr().add(i)); + let vb = vload_unaligned(b.as_ptr().add(i)); + vstore_unaligned::(a.as_mut_ptr().add(i), va $op vb); + } + i += lanes; + } + + for j in simd_len..len { + a[j] = a[j] $op b[j]; + } + } + + #[cfg(feature = "rayon")] + fn $par_fn(a: &mut [f32], b: &[f32]) { + a.par_chunks_mut(CHUNK_SIZE) + .zip(b.par_chunks(CHUNK_SIZE)) + .for_each(|(a_chunk, b_chunk)| { + $seq_fn(a_chunk, b_chunk); + }); + } + }; +} + +define_inplace_f32_op!(add_inplace_f32, add_inplace_f32_seq, add_inplace_f32_par, +); +define_inplace_f32_op!(sub_inplace_f32, sub_inplace_f32_seq, sub_inplace_f32_par, -); +define_inplace_f32_op!(mul_inplace_f32, mul_inplace_f32_seq, mul_inplace_f32_par, *); +define_inplace_f32_op!(div_inplace_f32, div_inplace_f32_seq, div_inplace_f32_par, /); + +// ============================================================================ +// f32 shared-row broadcast binary ops +// ============================================================================ +// +// Pattern: `dst[r * row_len + j] = dst[r * row_len + j] op row[j]` for +// all `r in 0..num_rows, j in 0..row_len`. That is, a single contiguous +// `row` is broadcast and combined with each of `num_rows` consecutive +// rows in `dst`. This is what `gamma.unsqueeze() * x` and friends +// produce after broadcast expansion. +// +// We call the SIMD dispatch once for the entire walk so that +// `#[macerator::with_simd]`'s feature detection pays once, not once per +// row. Calling the normal `add_inplace_f32` in a `0..num_rows` loop +// gives the right answer but costs ~55k dispatches on the typical +// layer_norm shape; see issue #64 item 2 benchmarks. + +macro_rules! define_shared_row_f32_op { + ($pub_fn:ident, $seq_fn:ident, $par_fn:ident, $op:tt) => { + #[inline] + pub fn $pub_fn(dst: &mut [f32], row: &[f32]) { + let row_len = row.len(); + if row_len == 0 || dst.is_empty() { + return; + } + debug_assert_eq!( + dst.len() % row_len, + 0, + "shared-row broadcast: dst length must be a multiple of row length" + ); + + #[cfg(feature = "rayon")] + if dst.len() >= PARALLEL_THRESHOLD { + $par_fn(dst, row); + return; + } + + $seq_fn(dst, row); + } + + #[macerator::with_simd] + #[allow(clippy::assign_op_pattern)] + fn $seq_fn(dst: &mut [f32], row: &[f32]) { + let lanes = S::lanes32(); + let row_len = row.len(); + let simd_len = row_len / lanes * lanes; + let num_rows = dst.len() / row_len; + + for r in 0..num_rows { + let base = r * row_len; + let mut i = 0; + while i < simd_len { + unsafe { + let va = vload_unaligned::(dst.as_ptr().add(base + i)); + let vb = vload_unaligned::(row.as_ptr().add(i)); + vstore_unaligned::( + dst.as_mut_ptr().add(base + i), + va $op vb, + ); + } + i += lanes; + } + for j in simd_len..row_len { + dst[base + j] = dst[base + j] $op row[j]; + } + } + } + + #[cfg(feature = "rayon")] + fn $par_fn(dst: &mut [f32], row: &[f32]) { + let row_len = row.len(); + // Chunk on row boundaries so each worker sees whole rows. + let rows_per_chunk = (CHUNK_SIZE / row_len).max(1); + let chunk_elems = rows_per_chunk * row_len; + dst.par_chunks_mut(chunk_elems).for_each(|chunk| { + $seq_fn(chunk, row); + }); + } + }; +} + +define_shared_row_f32_op!( + add_shared_row_inplace_f32, + add_shared_row_inplace_f32_seq, + add_shared_row_inplace_f32_par, + + +); +define_shared_row_f32_op!( + sub_shared_row_inplace_f32, + sub_shared_row_inplace_f32_seq, + sub_shared_row_inplace_f32_par, + - +); +define_shared_row_f32_op!( + mul_shared_row_inplace_f32, + mul_shared_row_inplace_f32_seq, + mul_shared_row_inplace_f32_par, + * +); +define_shared_row_f32_op!( + div_shared_row_inplace_f32, + div_shared_row_inplace_f32_seq, + div_shared_row_inplace_f32_par, + / +); + +// ============================================================================ +// f32 in-place unary ops (SIMD-accelerated) +// ============================================================================ + +#[inline] +pub fn abs_inplace_f32(a: &mut [f32]) { + #[cfg(feature = "rayon")] + if a.len() >= PARALLEL_THRESHOLD { + abs_inplace_f32_par(a); + return; + } + + abs_inplace_f32_seq(a); +} + +#[macerator::with_simd] +fn abs_inplace_f32_seq(a: &mut [f32]) { + let lanes = S::lanes32(); + let len = a.len(); + let simd_len = len / lanes * lanes; + + let mut i = 0; + while i < simd_len { + unsafe { + let v = vload_unaligned::(a.as_ptr().add(i)); + vstore_unaligned::(a.as_mut_ptr().add(i), v.abs()); + } + i += lanes; + } + + for v in &mut a[simd_len..len] { + *v = v.abs(); + } +} + +#[cfg(feature = "rayon")] +fn abs_inplace_f32_par(a: &mut [f32]) { + a.par_chunks_mut(CHUNK_SIZE).for_each(|chunk| { + abs_inplace_f32_seq(chunk); + }); +} + +#[inline] +pub fn recip_inplace_f32(a: &mut [f32]) { + #[cfg(feature = "rayon")] + if a.len() >= PARALLEL_THRESHOLD { + recip_inplace_f32_par(a); + return; + } + + recip_inplace_f32_seq(a); +} + +#[macerator::with_simd] +fn recip_inplace_f32_seq(a: &mut [f32]) { + let lanes = S::lanes32(); + let len = a.len(); + let simd_len = len / lanes * lanes; + + // Use exact SIMD division (not VRecip which is approximate on NEON/SSE) + let ones = 1.0f32.splat::(); + + let mut i = 0; + while i < simd_len { + unsafe { + let v = vload_unaligned::(a.as_ptr().add(i)); + vstore_unaligned::(a.as_mut_ptr().add(i), ones / v); + } + i += lanes; + } + + for v in &mut a[simd_len..len] { + *v = 1.0 / *v; + } +} + +#[cfg(feature = "rayon")] +fn recip_inplace_f32_par(a: &mut [f32]) { + a.par_chunks_mut(CHUNK_SIZE).for_each(|chunk| { + recip_inplace_f32_seq(chunk); + }); +} + +// ============================================================================ +// f32 comparison ops +// ============================================================================ + +#[derive(Clone, Copy)] +pub enum CmpOp { + Gt, + Ge, + Lt, + Le, + Eq, + Ne, +} + +#[inline] +pub fn cmp_f32(a: &[f32], b: &[f32], out: &mut [u8], op: CmpOp) { + debug_assert_eq!(a.len(), b.len()); + debug_assert_eq!(a.len(), out.len()); + + #[cfg(feature = "rayon")] + if a.len() >= PARALLEL_THRESHOLD { + cmp_f32_par(a, b, out, op); + return; + } + + cmp_f32_seq(a, b, out, op); +} + +/// Comparison kernel using simple loops that LLVM autovectorizes. +/// +/// Autovectorization outperforms explicit SIMD here because comparisons +/// produce u8 output from f32 input (4:1 size ratio). LLVM batches 16+ +/// comparisons and packs results into a single wide vector store, while +/// explicit SIMD (store_as_bool) writes only `lanes` bytes per iteration. +#[inline] +fn cmp_f32_seq(a: &[f32], b: &[f32], out: &mut [u8], op: CmpOp) { + match op { + CmpOp::Gt => { + for ((a, b), o) in a.iter().zip(b).zip(out.iter_mut()) { + *o = (*a > *b) as u8; + } + } + CmpOp::Ge => { + for ((a, b), o) in a.iter().zip(b).zip(out.iter_mut()) { + *o = (*a >= *b) as u8; + } + } + CmpOp::Lt => { + for ((a, b), o) in a.iter().zip(b).zip(out.iter_mut()) { + *o = (*a < *b) as u8; + } + } + CmpOp::Le => { + for ((a, b), o) in a.iter().zip(b).zip(out.iter_mut()) { + *o = (*a <= *b) as u8; + } + } + CmpOp::Eq => { + for ((a, b), o) in a.iter().zip(b).zip(out.iter_mut()) { + *o = (*a == *b) as u8; + } + } + CmpOp::Ne => { + for ((a, b), o) in a.iter().zip(b).zip(out.iter_mut()) { + *o = (*a != *b) as u8; + } + } + } +} + +#[cfg(feature = "rayon")] +fn cmp_f32_par(a: &[f32], b: &[f32], out: &mut [u8], op: CmpOp) { + out.par_chunks_mut(CHUNK_SIZE) + .enumerate() + .for_each(|(chunk_idx, out_chunk)| { + let start = chunk_idx * CHUNK_SIZE; + let end = (start + CHUNK_SIZE).min(a.len()); + cmp_f32_seq(&a[start..end], &b[start..end], out_chunk, op); + }); +} + +#[inline] +pub fn cmp_scalar_f32(a: &[f32], scalar: f32, out: &mut [u8], op: CmpOp) { + debug_assert_eq!(a.len(), out.len()); + + #[cfg(feature = "rayon")] + if a.len() >= PARALLEL_THRESHOLD { + cmp_scalar_f32_par(a, scalar, out, op); + return; + } + + cmp_scalar_f32_seq(a, scalar, out, op); +} + +/// Scalar comparison kernel using simple loops that LLVM autovectorizes. +/// See `cmp_f32_seq` for rationale. +#[inline] +fn cmp_scalar_f32_seq(a: &[f32], scalar: f32, out: &mut [u8], op: CmpOp) { + match op { + CmpOp::Gt => { + for (a, o) in a.iter().zip(out.iter_mut()) { + *o = (*a > scalar) as u8; + } + } + CmpOp::Ge => { + for (a, o) in a.iter().zip(out.iter_mut()) { + *o = (*a >= scalar) as u8; + } + } + CmpOp::Lt => { + for (a, o) in a.iter().zip(out.iter_mut()) { + *o = (*a < scalar) as u8; + } + } + CmpOp::Le => { + for (a, o) in a.iter().zip(out.iter_mut()) { + *o = (*a <= scalar) as u8; + } + } + CmpOp::Eq => { + for (a, o) in a.iter().zip(out.iter_mut()) { + *o = (*a == scalar) as u8; + } + } + CmpOp::Ne => { + for (a, o) in a.iter().zip(out.iter_mut()) { + *o = (*a != scalar) as u8; + } + } + } +} + +#[cfg(feature = "rayon")] +fn cmp_scalar_f32_par(a: &[f32], scalar: f32, out: &mut [u8], op: CmpOp) { + out.par_chunks_mut(CHUNK_SIZE) + .enumerate() + .for_each(|(chunk_idx, out_chunk)| { + let start = chunk_idx * CHUNK_SIZE; + let end = (start + CHUNK_SIZE).min(a.len()); + cmp_scalar_f32_seq(&a[start..end], scalar, out_chunk, op); + }); +} + +// ============================================================================ +// u8 boolean ops +// ============================================================================ + +macro_rules! define_bool_binary_u8_op { + ($pub_fn:ident, $seq_fn:ident, $par_fn:ident, + $inplace_pub:ident, $inplace_seq:ident, $inplace_par:ident, + $trait:ident, $method:ident, $op:tt) => { + #[inline] + pub fn $pub_fn(a: &[u8], b: &[u8], out: &mut [u8]) { + debug_assert_eq!(a.len(), b.len()); + debug_assert_eq!(a.len(), out.len()); + + #[cfg(feature = "rayon")] + if a.len() >= PARALLEL_THRESHOLD { + $par_fn(a, b, out); + return; + } + + $seq_fn(a, b, out); + } + + #[macerator::with_simd] + fn $seq_fn(a: &[u8], b: &[u8], out: &mut [u8]) { + let lanes = S::lanes8(); + let len = a.len(); + let simd_len = len / lanes * lanes; + + let mut i = 0; + while i < simd_len { + unsafe { + let va = vload_unaligned::(a.as_ptr().add(i)); + let vb = vload_unaligned::(b.as_ptr().add(i)); + vstore_unaligned::( + out.as_mut_ptr().add(i), + ::$method::(va, vb), + ); + } + i += lanes; + } + + for j in simd_len..len { + out[j] = a[j] $op b[j]; + } + } + + #[cfg(feature = "rayon")] + fn $par_fn(a: &[u8], b: &[u8], out: &mut [u8]) { + out.par_chunks_mut(CHUNK_SIZE) + .enumerate() + .for_each(|(chunk_idx, out_chunk)| { + let start = chunk_idx * CHUNK_SIZE; + let end = (start + CHUNK_SIZE).min(a.len()); + $seq_fn(&a[start..end], &b[start..end], out_chunk); + }); + } + + #[inline] + pub fn $inplace_pub(a: &mut [u8], b: &[u8]) { + debug_assert_eq!(a.len(), b.len()); + + #[cfg(feature = "rayon")] + if a.len() >= PARALLEL_THRESHOLD { + $inplace_par(a, b); + return; + } + + $inplace_seq(a, b); + } + + #[allow(clippy::assign_op_pattern)] + #[macerator::with_simd] + fn $inplace_seq(a: &mut [u8], b: &[u8]) { + let lanes = S::lanes8(); + let len = a.len(); + let simd_len = len / lanes * lanes; + + let mut i = 0; + while i < simd_len { + unsafe { + let va = vload_unaligned::(a.as_ptr().add(i)); + let vb = vload_unaligned::(b.as_ptr().add(i)); + vstore_unaligned::( + a.as_mut_ptr().add(i), + ::$method::(va, vb), + ); + } + i += lanes; + } + + for j in simd_len..len { + a[j] = a[j] $op b[j]; + } + } + + #[cfg(feature = "rayon")] + fn $inplace_par(a: &mut [u8], b: &[u8]) { + a.par_chunks_mut(CHUNK_SIZE) + .zip(b.par_chunks(CHUNK_SIZE)) + .for_each(|(a_chunk, b_chunk)| { + $inplace_seq(a_chunk, b_chunk); + }); + } + }; +} + +define_bool_binary_u8_op!( + bool_and_u8, bool_and_u8_seq, bool_and_u8_par, + bool_and_inplace_u8, bool_and_inplace_u8_seq, bool_and_inplace_u8_par, + VBitAnd, vbitand, &); +define_bool_binary_u8_op!( + bool_or_u8, bool_or_u8_seq, bool_or_u8_par, + bool_or_inplace_u8, bool_or_inplace_u8_seq, bool_or_inplace_u8_par, + VBitOr, vbitor, |); +define_bool_binary_u8_op!( + bool_xor_u8, bool_xor_u8_seq, bool_xor_u8_par, + bool_xor_inplace_u8, bool_xor_inplace_u8_seq, bool_xor_inplace_u8_par, + VBitXor, vbitxor, ^); + +// Boolean NOT is special (unary), implemented separately. + +#[inline] +pub fn bool_not_u8(a: &[u8], out: &mut [u8]) { + debug_assert_eq!(a.len(), out.len()); + + #[cfg(feature = "rayon")] + if a.len() >= PARALLEL_THRESHOLD { + bool_not_u8_par(a, out); + return; + } + + bool_not_u8_seq(a, out); +} + +#[macerator::with_simd] +fn bool_not_u8_seq(a: &[u8], out: &mut [u8]) { + let lanes = S::lanes8(); + let len = a.len(); + let simd_len = len / lanes * lanes; + + let zeros = 0u8.splat::(); + + let mut i = 0; + while i < simd_len { + unsafe { + let va = vload_unaligned::(a.as_ptr().add(i)); + let mask = va.eq(zeros); + mask.store_as_bool(out.as_mut_ptr().add(i) as *mut bool); + } + i += lanes; + } + + for j in simd_len..len { + out[j] = (a[j] == 0) as u8; + } +} + +#[cfg(feature = "rayon")] +fn bool_not_u8_par(a: &[u8], out: &mut [u8]) { + out.par_chunks_mut(CHUNK_SIZE) + .enumerate() + .for_each(|(chunk_idx, out_chunk)| { + let start = chunk_idx * CHUNK_SIZE; + let end = (start + CHUNK_SIZE).min(a.len()); + bool_not_u8_seq(&a[start..end], out_chunk); + }); +} + +#[inline] +pub fn bool_not_inplace_u8(a: &mut [u8]) { + #[cfg(feature = "rayon")] + if a.len() >= PARALLEL_THRESHOLD { + bool_not_inplace_u8_par(a); + return; + } + + bool_not_inplace_u8_seq(a); +} + +#[macerator::with_simd] +fn bool_not_inplace_u8_seq(a: &mut [u8]) { + let lanes = S::lanes8(); + let len = a.len(); + let simd_len = len / lanes * lanes; + + let zeros = 0u8.splat::(); + + let mut i = 0; + while i < simd_len { + unsafe { + let va = vload_unaligned::(a.as_ptr().add(i)); + let mask = va.eq(zeros); + mask.store_as_bool(a.as_mut_ptr().add(i) as *mut bool); + } + i += lanes; + } + + for v in &mut a[simd_len..len] { + *v = (*v == 0) as u8; + } +} + +#[cfg(feature = "rayon")] +fn bool_not_inplace_u8_par(a: &mut [u8]) { + a.par_chunks_mut(CHUNK_SIZE).for_each(|chunk| { + bool_not_inplace_u8_seq(chunk); + }); +} + +// ============================================================================ +// Mask select (mask_where / mask_fill) via bitwise blend +// ============================================================================ + +/// Conditional select: `out[i] = if mask[i] != 0 { value[i] } else { tensor[i] }` +/// +/// Operates on f32 data reinterpreted as u32 for bitwise blend. +/// Uses SIMD bitwise ops: `(mask & value) | (!mask & tensor)`. +#[inline] +pub fn mask_where_f32(tensor: &[f32], mask: &[u8], value: &[f32], out: &mut [f32]) { + debug_assert_eq!(tensor.len(), mask.len()); + debug_assert_eq!(tensor.len(), value.len()); + debug_assert_eq!(tensor.len(), out.len()); + + let t = bytemuck::cast_slice::(tensor); + let v = bytemuck::cast_slice::(value); + let o = bytemuck::cast_slice_mut::(out); + + #[cfg(feature = "rayon")] + if tensor.len() >= PARALLEL_THRESHOLD { + mask_where_u32_par(t, mask, v, o); + return; + } + + mask_blend_u32(t, mask, v, o); +} + +/// Conditional select for f64. +#[inline] +pub fn mask_where_f64(tensor: &[f64], mask: &[u8], value: &[f64], out: &mut [f64]) { + debug_assert_eq!(tensor.len(), mask.len()); + debug_assert_eq!(tensor.len(), value.len()); + debug_assert_eq!(tensor.len(), out.len()); + + let t = bytemuck::cast_slice::(tensor); + let v = bytemuck::cast_slice::(value); + let o = bytemuck::cast_slice_mut::(out); + + #[cfg(feature = "rayon")] + if tensor.len() >= PARALLEL_THRESHOLD { + mask_where_u64_par(t, mask, v, o); + return; + } + + mask_blend_u64(t, mask, v, o); +} + +/// Conditional select for i64. +#[inline] +pub fn mask_where_i64(tensor: &[i64], mask: &[u8], value: &[i64], out: &mut [i64]) { + debug_assert_eq!(tensor.len(), mask.len()); + debug_assert_eq!(tensor.len(), value.len()); + debug_assert_eq!(tensor.len(), out.len()); + + let t = bytemuck::cast_slice::(tensor); + let v = bytemuck::cast_slice::(value); + let o = bytemuck::cast_slice_mut::(out); + + #[cfg(feature = "rayon")] + if tensor.len() >= PARALLEL_THRESHOLD { + mask_where_u64_par(t, mask, v, o); + return; + } + + mask_blend_u64(t, mask, v, o); +} + +/// Conditional select for u8 (bool tensors). +#[inline] +pub fn mask_where_u8(tensor: &[u8], mask: &[u8], value: &[u8], out: &mut [u8]) { + debug_assert_eq!(tensor.len(), mask.len()); + debug_assert_eq!(tensor.len(), value.len()); + debug_assert_eq!(tensor.len(), out.len()); + + #[cfg(feature = "rayon")] + if tensor.len() >= PARALLEL_THRESHOLD { + mask_where_u8_par(tensor, mask, value, out); + return; + } + + mask_where_u8_seq(tensor, mask, value, out); +} + +/// Conditional fill: `out[i] = if mask[i] != 0 { fill_value } else { tensor[i] }` +#[inline] +pub fn mask_fill_f32(tensor: &[f32], mask: &[u8], fill_value: f32, out: &mut [f32]) { + debug_assert_eq!(tensor.len(), mask.len()); + debug_assert_eq!(tensor.len(), out.len()); + + let t = bytemuck::cast_slice::(tensor); + let o = bytemuck::cast_slice_mut::(out); + let fill_bits = fill_value.to_bits(); + + #[cfg(feature = "rayon")] + if tensor.len() >= PARALLEL_THRESHOLD { + mask_fill_u32_par(t, mask, fill_bits, o); + return; + } + + mask_blend_fill_u32(t, mask, fill_bits, o); +} + +/// Conditional fill for f64. +#[inline] +pub fn mask_fill_f64(tensor: &[f64], mask: &[u8], fill_value: f64, out: &mut [f64]) { + debug_assert_eq!(tensor.len(), mask.len()); + debug_assert_eq!(tensor.len(), out.len()); + + let t = bytemuck::cast_slice::(tensor); + let o = bytemuck::cast_slice_mut::(out); + let fill_bits = fill_value.to_bits(); + + #[cfg(feature = "rayon")] + if tensor.len() >= PARALLEL_THRESHOLD { + mask_fill_u64_par(t, mask, fill_bits, o); + return; + } + + mask_blend_fill_u64(t, mask, fill_bits, o); +} + +/// Conditional fill for i64. +#[inline] +pub fn mask_fill_i64(tensor: &[i64], mask: &[u8], fill_value: i64, out: &mut [i64]) { + debug_assert_eq!(tensor.len(), mask.len()); + debug_assert_eq!(tensor.len(), out.len()); + + let t = bytemuck::cast_slice::(tensor); + let o = bytemuck::cast_slice_mut::(out); + let fill_bits = fill_value as u64; + + #[cfg(feature = "rayon")] + if tensor.len() >= PARALLEL_THRESHOLD { + mask_fill_u64_par(t, mask, fill_bits, o); + return; + } + + mask_blend_fill_u64(t, mask, fill_bits, o); +} + +/// Conditional fill for u8 (bool tensors). +#[inline] +pub fn mask_fill_u8(tensor: &[u8], mask: &[u8], fill_value: u8, out: &mut [u8]) { + debug_assert_eq!(tensor.len(), mask.len()); + debug_assert_eq!(tensor.len(), out.len()); + + #[cfg(feature = "rayon")] + if tensor.len() >= PARALLEL_THRESHOLD { + mask_fill_u8_par(tensor, mask, fill_value, out); + return; + } + + mask_fill_u8_seq(tensor, mask, fill_value, out); +} + +// -- Branchless bitwise blend kernels -- +// +// These use a tight branchless loop that LLVM autovectorizes with native +// u8->u32/u64 widening instructions (NEON ushll, AVX2 vpmovzxbd). +// +// Mask values must be exactly 0 or 1 (Burn bool tensor invariant). +// wrapping_sub(0, 0)=0x00, wrapping_sub(0, 1)=0xFF..FF. Other values +// produce partial masks and corrupt output. + +#[inline] +fn mask_blend_u32(tensor: &[u32], mask: &[u8], value: &[u32], out: &mut [u32]) { + for i in 0..tensor.len() { + let m = 0u32.wrapping_sub(mask[i] as u32); + out[i] = (value[i] & m) | (tensor[i] & !m); + } +} + +#[inline] +fn mask_blend_fill_u32(tensor: &[u32], mask: &[u8], fill_bits: u32, out: &mut [u32]) { + for i in 0..tensor.len() { + let m = 0u32.wrapping_sub(mask[i] as u32); + out[i] = (fill_bits & m) | (tensor[i] & !m); + } +} + +#[inline] +fn mask_blend_u64(tensor: &[u64], mask: &[u8], value: &[u64], out: &mut [u64]) { + for i in 0..tensor.len() { + let m = 0u64.wrapping_sub(mask[i] as u64); + out[i] = (value[i] & m) | (tensor[i] & !m); + } +} + +#[inline] +fn mask_blend_fill_u64(tensor: &[u64], mask: &[u8], fill_bits: u64, out: &mut [u64]) { + for i in 0..tensor.len() { + let m = 0u64.wrapping_sub(mask[i] as u64); + out[i] = (fill_bits & m) | (tensor[i] & !m); + } +} + +#[cfg(feature = "rayon")] +fn mask_where_u32_par(tensor: &[u32], mask: &[u8], value: &[u32], out: &mut [u32]) { + out.par_chunks_mut(CHUNK_SIZE) + .enumerate() + .for_each(|(chunk_idx, out_chunk)| { + let start = chunk_idx * CHUNK_SIZE; + let end = (start + CHUNK_SIZE).min(tensor.len()); + mask_blend_u32( + &tensor[start..end], + &mask[start..end], + &value[start..end], + out_chunk, + ); + }); +} + +#[cfg(feature = "rayon")] +fn mask_fill_u32_par(tensor: &[u32], mask: &[u8], fill_bits: u32, out: &mut [u32]) { + out.par_chunks_mut(CHUNK_SIZE) + .enumerate() + .for_each(|(chunk_idx, out_chunk)| { + let start = chunk_idx * CHUNK_SIZE; + let end = (start + CHUNK_SIZE).min(tensor.len()); + mask_blend_fill_u32(&tensor[start..end], &mask[start..end], fill_bits, out_chunk); + }); +} + +#[cfg(feature = "rayon")] +fn mask_where_u64_par(tensor: &[u64], mask: &[u8], value: &[u64], out: &mut [u64]) { + out.par_chunks_mut(CHUNK_SIZE) + .enumerate() + .for_each(|(chunk_idx, out_chunk)| { + let start = chunk_idx * CHUNK_SIZE; + let end = (start + CHUNK_SIZE).min(tensor.len()); + mask_blend_u64( + &tensor[start..end], + &mask[start..end], + &value[start..end], + out_chunk, + ); + }); +} + +#[cfg(feature = "rayon")] +fn mask_fill_u64_par(tensor: &[u64], mask: &[u8], fill_bits: u64, out: &mut [u64]) { + out.par_chunks_mut(CHUNK_SIZE) + .enumerate() + .for_each(|(chunk_idx, out_chunk)| { + let start = chunk_idx * CHUNK_SIZE; + let end = (start + CHUNK_SIZE).min(tensor.len()); + mask_blend_fill_u64(&tensor[start..end], &mask[start..end], fill_bits, out_chunk); + }); +} + +// -- u8 SIMD kernels (for bool tensors) -- + +#[macerator::with_simd] +fn mask_where_u8_seq(tensor: &[u8], mask: &[u8], value: &[u8], out: &mut [u8]) { + let lanes = S::lanes8(); + let len = tensor.len(); + let simd_len = len / lanes * lanes; + + // SIMD subtract wraps: 0-0=0x00, 0-1=0xFF + let zeros = 0u8.splat::(); + + let mut i = 0; + while i < simd_len { + unsafe { + let vm_raw = vload_unaligned::(mask.as_ptr().add(i)); + let vm = zeros - vm_raw; // 0->0x00, 1->0xFF + let vt = vload_unaligned::(tensor.as_ptr().add(i)); + let vv = vload_unaligned::(value.as_ptr().add(i)); + let selected = (vm & vv) | (!vm & vt); + vstore_unaligned::(out.as_mut_ptr().add(i), selected); + } + i += lanes; + } + + for j in simd_len..len { + let m = 0u8.wrapping_sub(mask[j]); + out[j] = (m & value[j]) | (!m & tensor[j]); + } +} + +#[cfg(feature = "rayon")] +fn mask_where_u8_par(tensor: &[u8], mask: &[u8], value: &[u8], out: &mut [u8]) { + out.par_chunks_mut(CHUNK_SIZE) + .enumerate() + .for_each(|(chunk_idx, out_chunk)| { + let start = chunk_idx * CHUNK_SIZE; + let end = (start + CHUNK_SIZE).min(tensor.len()); + mask_where_u8_seq( + &tensor[start..end], + &mask[start..end], + &value[start..end], + out_chunk, + ); + }); +} + +#[macerator::with_simd] +fn mask_fill_u8_seq(tensor: &[u8], mask: &[u8], fill_value: u8, out: &mut [u8]) { + let lanes = S::lanes8(); + let len = tensor.len(); + let simd_len = len / lanes * lanes; + let vfill = fill_value.splat::(); + let zeros = 0u8.splat::(); + + let mut i = 0; + while i < simd_len { + unsafe { + let vm_raw = vload_unaligned::(mask.as_ptr().add(i)); + let vm = zeros - vm_raw; // 0->0x00, 1->0xFF + let vt = vload_unaligned::(tensor.as_ptr().add(i)); + let selected = (vm & vfill) | (!vm & vt); + vstore_unaligned::(out.as_mut_ptr().add(i), selected); + } + i += lanes; + } + + for j in simd_len..len { + let m = 0u8.wrapping_sub(mask[j]); + out[j] = (m & fill_value) | (!m & tensor[j]); + } +} + +#[cfg(feature = "rayon")] +fn mask_fill_u8_par(tensor: &[u8], mask: &[u8], fill_value: u8, out: &mut [u8]) { + out.par_chunks_mut(CHUNK_SIZE) + .enumerate() + .for_each(|(chunk_idx, out_chunk)| { + let start = chunk_idx * CHUNK_SIZE; + let end = (start + CHUNK_SIZE).min(tensor.len()); + mask_fill_u8_seq( + &tensor[start..end], + &mask[start..end], + fill_value, + out_chunk, + ); + }); +} + +// ============================================================================ +// Tests +// ============================================================================ + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_add_inplace_f32() { + let mut a = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]; + let b = [10.0f32, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0]; + add_inplace_f32(&mut a, &b); + assert_eq!(a, [11.0, 22.0, 33.0, 44.0, 55.0, 66.0, 77.0]); + } + + #[test] + fn test_sub_inplace_f32() { + let mut a = [10.0f32, 20.0, 30.0, 40.0, 50.0]; + let b = [1.0f32, 2.0, 3.0, 4.0, 5.0]; + sub_inplace_f32(&mut a, &b); + assert_eq!(a, [9.0, 18.0, 27.0, 36.0, 45.0]); + } + + #[test] + fn test_mul_inplace_f32() { + let mut a = [1.0f32, 2.0, 3.0, 4.0, 5.0]; + let b = [2.0f32, 2.0, 2.0, 2.0, 2.0]; + mul_inplace_f32(&mut a, &b); + assert_eq!(a, [2.0, 4.0, 6.0, 8.0, 10.0]); + } + + #[test] + fn test_div_inplace_f32() { + let mut a = [10.0f32, 20.0, 30.0, 40.0]; + let b = [2.0f32, 4.0, 5.0, 8.0]; + div_inplace_f32(&mut a, &b); + assert_eq!(a, [5.0, 5.0, 6.0, 5.0]); + } + + #[test] + fn test_cmp_gt_f32() { + let a = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]; + let b = [2.0f32, 2.0, 2.0, 4.0, 4.0, 4.0, 4.0]; + let mut out = [0u8; 7]; + cmp_f32(&a, &b, &mut out, CmpOp::Gt); + assert_eq!(out, [0, 0, 1, 0, 1, 1, 1]); + } + + #[test] + fn test_cmp_ge_f32() { + let a = [1.0f32, 2.0, 3.0, 4.0]; + let b = [2.0f32, 2.0, 2.0, 5.0]; + let mut out = [0u8; 4]; + cmp_f32(&a, &b, &mut out, CmpOp::Ge); + assert_eq!(out, [0, 1, 1, 0]); + } + + #[test] + fn test_cmp_eq_f32() { + let a = [1.0f32, 2.0, 3.0, 4.0, 5.0]; + let b = [1.0f32, 3.0, 3.0, 5.0, 5.0]; + let mut out = [0u8; 5]; + cmp_f32(&a, &b, &mut out, CmpOp::Eq); + assert_eq!(out, [1, 0, 1, 0, 1]); + } + + #[test] + fn test_cmp_ne_f32() { + let a = [1.0f32, 2.0, 3.0, 4.0, 5.0]; + let b = [1.0f32, 3.0, 3.0, 5.0, 5.0]; + let mut out = [0u8; 5]; + cmp_f32(&a, &b, &mut out, CmpOp::Ne); + assert_eq!(out, [0, 1, 0, 1, 0]); + } + + #[test] + fn test_cmp_scalar_gt_f32() { + let a = [1.0f32, 2.0, 3.0, 4.0, 5.0]; + let mut out = [0u8; 5]; + cmp_scalar_f32(&a, 3.0, &mut out, CmpOp::Gt); + assert_eq!(out, [0, 0, 0, 1, 1]); + } + + #[test] + fn test_bool_not_u8() { + let a = [1u8, 0, 1, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 1, 0, 1, 0]; + let mut out = [0u8; 18]; + bool_not_u8(&a, &mut out); + let expected = [0u8, 1, 0, 1, 0, 1, 1, 0, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1]; + assert_eq!(out, expected); + } + + #[test] + fn test_bool_not_inplace_u8() { + let mut a = [1u8, 0, 1, 0]; + bool_not_inplace_u8(&mut a); + assert_eq!(a, [0, 1, 0, 1]); + } + + #[test] + fn test_bool_and_u8() { + let a = [1u8, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 0]; + let b = [1u8, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1]; + let mut out = [0u8; 18]; + bool_and_u8(&a, &b, &mut out); + let expected = [1u8, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0]; + assert_eq!(out, expected); + } + + #[test] + fn test_bool_or_u8() { + let a = [1u8, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 0]; + let b = [1u8, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 1]; + let mut out = [0u8; 18]; + bool_or_u8(&a, &b, &mut out); + let expected = [1u8, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1]; + assert_eq!(out, expected); + } + + #[test] + fn test_bool_xor_u8() { + let a = [1u8, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 0]; + let b = [1u8, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1]; + let mut out = [0u8; 18]; + bool_xor_u8(&a, &b, &mut out); + let expected = [0u8, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1]; + assert_eq!(out, expected); + } + + #[test] + fn test_bool_and_inplace_u8() { + let mut a = [1u8, 1, 0, 0]; + let b = [1u8, 0, 1, 0]; + bool_and_inplace_u8(&mut a, &b); + assert_eq!(a, [1, 0, 0, 0]); + } + + #[test] + fn test_bool_or_inplace_u8() { + let mut a = [1u8, 1, 0, 0]; + let b = [1u8, 0, 1, 0]; + bool_or_inplace_u8(&mut a, &b); + assert_eq!(a, [1, 1, 1, 0]); + } + + #[test] + fn test_bool_xor_inplace_u8() { + let mut a = [1u8, 1, 0, 0]; + let b = [1u8, 0, 1, 0]; + bool_xor_inplace_u8(&mut a, &b); + assert_eq!(a, [0, 1, 1, 0]); + } + + #[test] + fn test_abs_inplace_f32() { + let mut a = [-3.0f32, -1.0, 0.0, 1.0, 3.0, -5.0, 7.0]; + abs_inplace_f32(&mut a); + assert_eq!(a, [3.0, 1.0, 0.0, 1.0, 3.0, 5.0, 7.0]); + } + + #[test] + fn test_recip_inplace_f32() { + let mut a = [1.0f32, 2.0, 4.0, 0.5, 10.0]; + recip_inplace_f32(&mut a); + assert_eq!(a, [1.0, 0.5, 0.25, 2.0, 0.1]); + } + + // ================================================================ + // mask_where / mask_fill tests + // ================================================================ + + #[test] + fn test_mask_where_f32_basic() { + let tensor = [1.0f32, 2.0, 3.0, 4.0]; + let mask = [1u8, 0, 1, 0]; + let value = [10.0f32, 20.0, 30.0, 40.0]; + let mut out = [0.0f32; 4]; + mask_where_f32(&tensor, &mask, &value, &mut out); + assert_eq!(out, [10.0, 2.0, 30.0, 4.0]); + } + + #[test] + fn test_mask_where_f32_all_true() { + let tensor = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]; + let mask = [1u8; 9]; + let value = [10.0f32, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0]; + let mut out = [0.0f32; 9]; + mask_where_f32(&tensor, &mask, &value, &mut out); + assert_eq!(out, value); + } + + #[test] + fn test_mask_where_f32_all_false() { + let tensor = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]; + let mask = [0u8; 9]; + let value = [10.0f32, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0]; + let mut out = [0.0f32; 9]; + mask_where_f32(&tensor, &mask, &value, &mut out); + assert_eq!(out, tensor); + } + + #[test] + fn test_mask_where_f64_basic() { + let tensor = [1.0f64, 2.0, 3.0]; + let mask = [0u8, 1, 0]; + let value = [10.0f64, 20.0, 30.0]; + let mut out = [0.0f64; 3]; + mask_where_f64(&tensor, &mask, &value, &mut out); + assert_eq!(out, [1.0, 20.0, 3.0]); + } + + #[test] + fn test_mask_where_i64_basic() { + let tensor = [10i64, 20, 30, 40, 50]; + let mask = [1u8, 0, 1, 0, 1]; + let value = [-1i64, -2, -3, -4, -5]; + let mut out = [0i64; 5]; + mask_where_i64(&tensor, &mask, &value, &mut out); + assert_eq!(out, [-1, 20, -3, 40, -5]); + } + + #[test] + fn test_mask_where_u8_basic() { + let tensor = [0u8, 1, 0, 1]; + let mask = [1u8, 1, 0, 0]; + let value = [1u8, 0, 1, 0]; + let mut out = [0u8; 4]; + mask_where_u8(&tensor, &mask, &value, &mut out); + assert_eq!(out, [1, 0, 0, 1]); + } + + #[test] + fn test_mask_fill_f32_basic() { + let tensor = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]; + let mask = [1u8, 0, 1, 0, 1, 0, 1]; + let mut out = [0.0f32; 7]; + mask_fill_f32(&tensor, &mask, -1.0, &mut out); + assert_eq!(out, [-1.0, 2.0, -1.0, 4.0, -1.0, 6.0, -1.0]); + } + + #[test] + fn test_mask_fill_f64_basic() { + let tensor = [1.0f64, 2.0, 3.0]; + let mask = [0u8, 1, 0]; + let mut out = [0.0f64; 3]; + mask_fill_f64(&tensor, &mask, 99.0, &mut out); + assert_eq!(out, [1.0, 99.0, 3.0]); + } + + #[test] + fn test_mask_fill_i64_basic() { + let tensor = [10i64, 20, 30, 40]; + let mask = [1u8, 0, 0, 1]; + let mut out = [0i64; 4]; + mask_fill_i64(&tensor, &mask, -1, &mut out); + assert_eq!(out, [-1, 20, 30, -1]); + } + + #[test] + fn test_mask_fill_u8_basic() { + let tensor = [0u8, 1, 0, 1, 0]; + let mask = [1u8, 1, 0, 0, 1]; + let mut out = [0u8; 5]; + mask_fill_u8(&tensor, &mask, 1, &mut out); + assert_eq!(out, [1, 1, 0, 1, 1]); + } + + #[test] + fn test_mask_where_f32_nan() { + let tensor = [f32::NAN, 2.0, 3.0, f32::NAN]; + let mask = [1u8, 0, 1, 0]; + let value = [10.0f32, 20.0, 30.0, 40.0]; + let mut out = [0.0f32; 4]; + mask_where_f32(&tensor, &mask, &value, &mut out); + // mask=1 picks value, mask=0 picks tensor (including NaN) + assert_eq!(out[0], 10.0); + assert_eq!(out[1], 2.0); + assert_eq!(out[2], 30.0); + assert!(out[3].is_nan()); + } + + // Lane-boundary tests: sizes that exercise SIMD + scalar tail on all ISAs. + // 17 elements for f32 = 4 NEON iters + 1 tail, or 2 AVX2 iters + 1 tail. + + #[test] + fn test_mask_where_f32_lane_boundary() { + let n = 17; + let tensor: Vec = (0..n).map(|i| i as f32).collect(); + let value: Vec = (0..n).map(|i| (i as f32) * 10.0).collect(); + let mask: Vec = (0..n).map(|i| (i % 2) as u8).collect(); + let mut out = vec![0.0f32; n]; + mask_where_f32(&tensor, &mask, &value, &mut out); + for i in 0..n { + let expected = if i % 2 != 0 { value[i] } else { tensor[i] }; + assert_eq!(out[i], expected, "mismatch at index {i}"); + } + } + + #[test] + fn test_mask_fill_f32_lane_boundary() { + let n = 17; + let tensor: Vec = (0..n).map(|i| i as f32).collect(); + let mask: Vec = (0..n).map(|i| (i % 3 == 0) as u8).collect(); + let mut out = vec![0.0f32; n]; + mask_fill_f32(&tensor, &mask, -1.0, &mut out); + for i in 0..n { + let expected = if i % 3 == 0 { -1.0 } else { tensor[i] }; + assert_eq!(out[i], expected, "mismatch at index {i}"); + } + } + + #[test] + fn test_mask_where_u8_lane_boundary() { + // 33 elements for u8: exercises 2 NEON iters + 1 tail, or 1 AVX2 iter + 1 tail + let n = 33; + let tensor: Vec = (0..n).map(|i| (i % 2) as u8).collect(); + let value: Vec = (0..n).map(|i| ((i + 1) % 2) as u8).collect(); + let mask: Vec = (0..n).map(|i| (i % 3 == 0) as u8).collect(); + let mut out = vec![0u8; n]; + mask_where_u8(&tensor, &mask, &value, &mut out); + for i in 0..n { + let expected = if i % 3 == 0 { value[i] } else { tensor[i] }; + assert_eq!(out[i], expected, "mismatch at index {i}"); + } + } + + #[test] + fn test_mask_where_f64_lane_boundary() { + // 9 elements for f64: 4 NEON iters + 1 tail (2 lanes), or 2 AVX2 iters + 1 tail (4 lanes) + let n = 9; + let tensor: Vec = (0..n).map(|i| i as f64).collect(); + let value: Vec = (0..n).map(|i| -(i as f64)).collect(); + let mask: Vec = (0..n).map(|i| (i % 2) as u8).collect(); + let mut out = vec![0.0f64; n]; + mask_where_f64(&tensor, &mask, &value, &mut out); + for i in 0..n { + let expected = if i % 2 != 0 { value[i] } else { tensor[i] }; + assert_eq!(out[i], expected, "mismatch at index {i}"); + } + } + + #[test] + fn test_mask_where_empty() { + let mut out = vec![0.0f32; 0]; + mask_where_f32(&[], &[], &[], &mut out); + assert!(out.is_empty()); + } + + #[test] + fn test_mask_fill_empty() { + let mut out = vec![0.0f32; 0]; + mask_fill_f32(&[], &[], 1.0, &mut out); + assert!(out.is_empty()); + } + + // bool_not_u8 writes into a `*mut bool` via macerator's store_as_bool. + // Rust's bool is only valid as 0x00 or 0x01 (any other byte is UB when + // read back as bool). A previous audit flagged that SIMD mask stores + // might emit 0xFF. Verify every output byte is normalized to 0/1. + #[test] + fn bool_not_u8_output_is_normalized_0_or_1() { + // Spans SIMD body + scalar tail on any realistic lane width: + // 17 elements -> NEON 16-byte SIMD + 1 tail; AVX2 32 spills to tail. + // 127 elements -> SIMD body + 15 tail for 16-byte lanes. + for &len in &[1usize, 8, 15, 16, 17, 31, 32, 63, 127, 256] { + let a: Vec = (0..len).map(|i| (i % 2) as u8).collect(); + let mut out = vec![0xAAu8; len]; + super::bool_not_u8(&a, &mut out); + for (i, &b) in out.iter().enumerate() { + assert!( + b == 0 || b == 1, + "len={}: out[{}] = 0x{:02x}, expected 0x00 or 0x01", + len, + i, + b + ); + let expected = if a[i] == 0 { 1 } else { 0 }; + assert_eq!( + b, expected, + "len={}: out[{}] = {}, expected {}", + len, i, b, expected + ); + } + } + } + + #[test] + fn bool_not_inplace_u8_output_is_normalized_0_or_1() { + for &len in &[1usize, 8, 15, 16, 17, 31, 32, 63, 127, 256] { + let mut a: Vec = (0..len).map(|i| (i % 2) as u8).collect(); + let original = a.clone(); + super::bool_not_inplace_u8(&mut a); + for (i, &b) in a.iter().enumerate() { + assert!( + b == 0 || b == 1, + "len={}: a[{}] = 0x{:02x}, expected 0x00 or 0x01", + len, + i, + b + ); + let expected = if original[i] == 0 { 1 } else { 0 }; + assert_eq!( + b, expected, + "len={}: a[{}] = {}, expected {}", + len, i, b, expected + ); + } + } + } + + // Edge cases: empty input, homogeneous all-zero, homogeneous all-one. + // Homogeneous inputs exercise the SIMD mask-to-byte conversion for + // all-true and all-false cases, which alternating inputs do not. + #[test] + fn bool_not_u8_edge_cases() { + // Empty input. + let mut out: Vec = Vec::new(); + super::bool_not_u8(&[], &mut out); + assert!(out.is_empty()); + + // All zeros -> all ones. + let a = alloc::vec![0u8; 32]; + let mut out = alloc::vec![0xAAu8; 32]; + super::bool_not_u8(&a, &mut out); + assert!(out.iter().all(|&b| b == 1)); + + // All ones -> all zeros. + let a = alloc::vec![1u8; 32]; + let mut out = alloc::vec![0xAAu8; 32]; + super::bool_not_u8(&a, &mut out); + assert!(out.iter().all(|&b| b == 0)); + } +} diff --git a/crates/burn-flex/src/simd/scalar.rs b/crates/burn-flex/src/simd/scalar.rs new file mode 100644 index 0000000..6299ef6 --- /dev/null +++ b/crates/burn-flex/src/simd/scalar.rs @@ -0,0 +1,74 @@ +/// Comparison operation type. +#[derive(Clone, Copy)] +pub enum CmpOp { + Gt, + Ge, + Lt, + Le, + Eq, + Ne, +} + +/// Scalar boolean NOT: out[i] = !a[i] (0 becomes 1, non-zero becomes 0) +#[inline] +pub fn bool_not_u8(a: &[u8], out: &mut [u8]) { + for i in 0..a.len() { + out[i] = (a[i] == 0) as u8; + } +} + +/// Scalar boolean NOT in-place: a[i] = !a[i] +#[inline] +pub fn bool_not_inplace_u8(a: &mut [u8]) { + for x in a.iter_mut() { + *x = (*x == 0) as u8; + } +} + +/// Scalar boolean AND: out[i] = a[i] & b[i] +#[inline] +pub fn bool_and_u8(a: &[u8], b: &[u8], out: &mut [u8]) { + for i in 0..a.len() { + out[i] = a[i] & b[i]; + } +} + +/// Scalar boolean OR: out[i] = a[i] | b[i] +#[inline] +pub fn bool_or_u8(a: &[u8], b: &[u8], out: &mut [u8]) { + for i in 0..a.len() { + out[i] = a[i] | b[i]; + } +} + +/// Scalar boolean XOR: out[i] = a[i] ^ b[i] +#[inline] +pub fn bool_xor_u8(a: &[u8], b: &[u8], out: &mut [u8]) { + for i in 0..a.len() { + out[i] = a[i] ^ b[i]; + } +} + +/// Scalar boolean AND in-place: a[i] &= b[i] +#[inline] +pub fn bool_and_inplace_u8(a: &mut [u8], b: &[u8]) { + for i in 0..a.len() { + a[i] &= b[i]; + } +} + +/// Scalar boolean OR in-place: a[i] |= b[i] +#[inline] +pub fn bool_or_inplace_u8(a: &mut [u8], b: &[u8]) { + for i in 0..a.len() { + a[i] |= b[i]; + } +} + +/// Scalar boolean XOR in-place: a[i] ^= b[i] +#[inline] +pub fn bool_xor_inplace_u8(a: &mut [u8], b: &[u8]) { + for i in 0..a.len() { + a[i] ^= b[i]; + } +} diff --git a/crates/burn-flex/src/strided_index.rs b/crates/burn-flex/src/strided_index.rs new file mode 100644 index 0000000..49e3125 --- /dev/null +++ b/crates/burn-flex/src/strided_index.rs @@ -0,0 +1,99 @@ +use crate::layout::Layout; +use alloc::vec; +use alloc::vec::Vec; + +/// Iterator that yields linear indices for strided tensor access. +/// +/// Handles non-contiguous tensors (including those with negative strides from flip) +/// by computing the correct storage index for each logical element position. +pub struct StridedIter<'a> { + /// Current storage index (signed to handle negative strides) + storage_index: isize, + multi_index: Vec, + layout: &'a Layout, + remaining: usize, +} + +impl<'a> StridedIter<'a> { + /// Create a new strided iterator for the given layout. + pub fn new(layout: &'a Layout) -> Self { + let ndims = layout.num_dims(); + Self { + storage_index: layout.start_offset() as isize, + multi_index: vec![0; ndims], + layout, + remaining: layout.num_elements(), + } + } +} + +impl Iterator for StridedIter<'_> { + type Item = usize; + + fn next(&mut self) -> Option { + if self.remaining == 0 { + return None; + } + + debug_assert!( + self.storage_index >= 0, + "StridedIter: negative storage index" + ); + let idx = self.storage_index as usize; + self.remaining -= 1; + + // Advance multi-index (last dimension first, like odometer) + let shape = self.layout.shape(); + let strides = self.layout.strides(); + + for d in (0..shape.num_dims()).rev() { + self.multi_index[d] += 1; + if self.multi_index[d] < shape[d] { + self.storage_index += strides[d]; + break; + } + // Wrap around this dimension + self.multi_index[d] = 0; + self.storage_index -= (shape[d] as isize - 1) * strides[d]; + } + + Some(idx) + } + + fn size_hint(&self) -> (usize, Option) { + (self.remaining, Some(self.remaining)) + } +} + +impl ExactSizeIterator for StridedIter<'_> {} + +#[cfg(test)] +mod tests { + use super::*; + use burn_std::Shape; + + #[test] + fn test_contiguous_iteration() { + let layout = Layout::contiguous(Shape::from(vec![2, 3])); + let indices: Vec<_> = StridedIter::new(&layout).collect(); + assert_eq!(indices, vec![0, 1, 2, 3, 4, 5]); + } + + #[test] + fn test_transposed_iteration() { + // Original: [[0, 1, 2], [3, 4, 5]] (shape [2, 3], strides [3, 1]) + // Transposed: [[0, 3], [1, 4], [2, 5]] (shape [3, 2], strides [1, 3]) + let layout = Layout::contiguous(Shape::from(vec![2, 3])).transpose(0, 1); + let indices: Vec<_> = StridedIter::new(&layout).collect(); + assert_eq!(indices, vec![0, 3, 1, 4, 2, 5]); + } + + #[test] + fn test_narrowed_iteration() { + // Original: [[0, 1, 2, 3], [4, 5, 6, 7]] (shape [2, 4]) + // Narrowed dim=1, start=1, len=2: [[1, 2], [5, 6]] + let layout = Layout::contiguous(Shape::from(vec![2, 4])).narrow(1, 1, 2); + let indices: Vec<_> = StridedIter::new(&layout).collect(); + assert_eq!(indices, vec![1, 2, 5, 6]); + } +} diff --git a/crates/burn-flex/src/tensor.rs b/crates/burn-flex/src/tensor.rs new file mode 100644 index 0000000..82ab187 --- /dev/null +++ b/crates/burn-flex/src/tensor.rs @@ -0,0 +1,1016 @@ +#[cfg(target_has_atomic = "ptr")] +use alloc::sync::Arc; +use alloc::vec::Vec; +use core::fmt; +#[cfg(not(target_has_atomic = "ptr"))] +use portable_atomic_util::Arc; + +use burn_backend::{DType, Element, TensorData, TensorMetadata}; +use burn_std::{Bytes, Shape, bf16, f16}; + +use crate::{FlexDevice, layout::Layout}; + +/// CPU tensor primitive for the Flex backend. +/// +/// Uses type-erased byte storage with runtime dtype and Arc-based sharing. +/// Clone is O(1) (refcount increment). Copy-on-write for mutations. +#[derive(Clone)] +pub struct FlexTensor { + /// Shared byte storage. Clone increments refcount. + data: Arc, + /// Layout describing shape, strides, and offset. + layout: Layout, + /// Runtime data type. + dtype: DType, +} + +impl fmt::Debug for FlexTensor { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("FlexTensor") + .field("shape", self.layout.shape()) + .field("dtype", &self.dtype) + .field("contiguous", &self.layout.is_contiguous()) + .field("unique", &self.is_unique()) + .finish() + } +} + +impl FlexTensor { + /// Create a new tensor from bytes, layout, and dtype. + pub fn new(data: Bytes, layout: Layout, dtype: DType) -> Self { + Self { + data: Arc::new(data), + layout, + dtype, + } + } + + /// Create a tensor from TensorData. + pub fn from_data(data: TensorData) -> Self { + let shape = data.shape.clone(); + let layout = Layout::contiguous(shape); + let dtype = data.dtype; + Self { + data: Arc::new(data.bytes), + layout, + dtype, + } + } + + /// Convert tensor to TensorData. + /// + /// If non-contiguous or shared, this will copy data. + pub fn into_data(self) -> TensorData { + if self.layout.is_contiguous() && self.layout.start_offset() == 0 { + let expected_bytes = self.layout.num_elements() * dtype_size(self.dtype); + assert!( + expected_bytes <= self.data.len(), + "into_data: buffer ({} bytes) too small for {} elements of {:?}", + self.data.len(), + self.layout.num_elements(), + self.dtype + ); + if self.data.len() == expected_bytes { + // Buffer exactly matches logical size; try zero-copy unwrap + match Arc::try_unwrap(self.data) { + Ok(bytes) => TensorData { + bytes, + shape: self.layout.shape().clone(), + dtype: self.dtype, + }, + Err(arc) => { + let bytes = Bytes::from_bytes_vec((*arc)[..expected_bytes].to_vec()); + TensorData { + bytes, + shape: self.layout.shape().clone(), + dtype: self.dtype, + } + } + } + } else { + // Contiguous at offset 0 but buffer is oversized (e.g., narrowed view). + // Truncate to exact logical size. + let bytes = Bytes::from_bytes_vec(self.data[..expected_bytes].to_vec()); + TensorData { + bytes, + shape: self.layout.shape().clone(), + dtype: self.dtype, + } + } + } else { + // Non-contiguous or non-zero offset: copy to contiguous layout + self.to_contiguous().into_data() + } + } + + /// Check if this tensor has exclusive ownership of its data. + /// + /// When true, in-place mutations are safe without copying. + #[inline] + pub fn is_unique(&self) -> bool { + Arc::strong_count(&self.data) == 1 + } + + /// Get the layout. + pub fn layout(&self) -> &Layout { + &self.layout + } + + /// Create a new tensor with a different layout but sharing the same data. + /// + /// This is a zero-copy operation used for operations like flip, transpose, etc. + pub fn with_layout(self, layout: Layout) -> Self { + Self { + data: self.data, + layout, + dtype: self.dtype, + } + } + + /// Get the dtype. + pub fn dtype(&self) -> DType { + self.dtype + } + + /// Check if tensor is contiguous. + pub fn is_contiguous(&self) -> bool { + self.layout.is_contiguous() + } + + /// Get the raw bytes (read-only). + pub fn bytes(&self) -> &[u8] { + &self.data + } + + /// Get a clone of the Arc for sharing data with a new layout. + /// + /// Use this for zero-copy view operations (reshape, transpose, slice). + pub fn data_arc(&self) -> Arc { + Arc::clone(&self.data) + } + + /// Create a tensor from shared data, layout, and dtype. + /// + /// Use this for zero-copy view operations. + pub fn from_arc(data: Arc, layout: Layout, dtype: DType) -> Self { + Self { + data, + layout, + dtype, + } + } + + /// Zero-copy typed view of the full storage buffer. + /// + /// Use with `StridedIter` for non-contiguous access, or with + /// `layout().contiguous_offsets()` for the contiguous fast path. + /// + /// # Panics + /// Panics if `E::dtype()` doesn't match the tensor's dtype. + /// Note: Bool tensors are stored as u8, so both Bool(Native) and Bool(U8) + /// dtypes accept u8 access. + pub fn storage(&self) -> &[E] { + assert!( + E::dtype() == self.dtype + || (matches!( + self.dtype, + DType::Bool(burn_std::BoolStore::Native | burn_std::BoolStore::U8) + ) && E::dtype() == DType::U8), + "storage: dtype mismatch (expected {:?}, got {:?})", + self.dtype, + E::dtype() + ); + bytemuck::cast_slice(&self.data) + } + + /// Mutable typed view with copy-on-write semantics. + /// + /// If the tensor is shared (refcount > 1), this will copy the data first. + /// For in-place operations, prefer `try_storage_mut()` which returns None + /// if shared, allowing you to choose an alternative strategy. + /// + /// # Panics + /// Panics if `E::dtype()` doesn't match the tensor's dtype. + /// Note: Bool tensors are stored as u8, so both Bool(Native) and Bool(U8) + /// dtypes accept u8 access. + pub fn storage_mut(&mut self) -> &mut [E] { + assert!( + E::dtype() == self.dtype + || (matches!( + self.dtype, + DType::Bool(burn_std::BoolStore::Native | burn_std::BoolStore::U8) + ) && E::dtype() == DType::U8), + "storage_mut: dtype mismatch (expected {:?}, got {:?})", + self.dtype, + E::dtype() + ); + // COW: clone data if shared + let bytes = Arc::make_mut(&mut self.data); + bytemuck::cast_slice_mut(bytes) + } + + /// Try to get mutable storage without copying. + /// + /// Returns `Some` if tensor is uniquely owned, `None` if shared. + /// Use this when you want to avoid the implicit copy in `storage_mut()`. + /// Note: Bool tensors are stored as u8, so both Bool(Native) and Bool(U8) + /// dtypes accept u8 access. + pub fn try_storage_mut(&mut self) -> Option<&mut [E]> { + assert!( + E::dtype() == self.dtype + || (matches!( + self.dtype, + DType::Bool(burn_std::BoolStore::Native | burn_std::BoolStore::U8) + ) && E::dtype() == DType::U8), + "try_storage_mut: dtype mismatch (expected {:?}, got {:?})", + self.dtype, + E::dtype() + ); + if self.is_unique() { + // Safe: we're the only owner + let bytes = Arc::get_mut(&mut self.data)?; + Some(bytemuck::cast_slice_mut(bytes)) + } else { + None + } + } + + /// Get typed slice view (zero-cost if contiguous and offset is 0). + /// + /// Returns None if dtype doesn't match E or tensor is non-contiguous. + pub fn as_slice(&self) -> Option<&[E]> { + if E::dtype() != self.dtype { + return None; + } + let storage: &[E] = self.storage(); + self.layout + .contiguous_offsets() + .map(|(start, end)| &storage[start..end]) + } + + /// Create an empty tensor with given shape and dtype. + pub fn empty(shape: Shape, dtype: DType) -> Self { + let num_elements = shape.num_elements(); + let elem_size = dtype_size(dtype); + let bytes = Bytes::from_bytes_vec(alloc::vec![0u8; num_elements * elem_size]); + let layout = Layout::contiguous(shape); + Self { + data: Arc::new(bytes), + layout, + dtype, + } + } + + /// Create a tensor filled with zeros. + pub fn zeros(shape: Shape, dtype: DType) -> Self { + Self::empty(shape, dtype) + } + + /// Create a tensor filled with `n` copies of a typed value. + pub fn filled_typed( + shape: Shape, + dtype: DType, + value: E, + ) -> Self { + assert_eq!( + dtype_size(dtype), + core::mem::size_of::(), + "filled_typed: dtype size mismatch" + ); + let n = shape.num_elements(); + let data = alloc::vec![value; n]; + let bytes = Bytes::from_elems(data); + Self { + data: Arc::new(bytes), + layout: Layout::contiguous(shape), + dtype, + } + } + + /// Copy to contiguous layout if needed. + pub fn to_contiguous(&self) -> Self { + // Fast path requires the logical tensor to cover the whole buffer. + // A contiguous prefix view (e.g. [8, 5] sliced to [5, 5]) has + // canonical strides and offset 0 but an oversized buffer, and would + // otherwise mislead callers that read `storage()` / `bytes()` by + // length (e.g. the SIMD `mask_fill_*` kernels). + if self.is_contiguous() + && self.layout.start_offset() == 0 + && self.data.len() == self.layout.num_elements() * dtype_size(self.dtype) + { + return self.clone(); + } + + // Copy data to new contiguous buffer + match self.dtype { + DType::F64 => self.copy_contiguous::(), + DType::F32 => self.copy_contiguous::(), + DType::F16 => self.copy_contiguous::(), + DType::BF16 => self.copy_contiguous::(), + DType::I64 => self.copy_contiguous::(), + DType::I32 => self.copy_contiguous::(), + DType::I16 => self.copy_contiguous::(), + DType::I8 => self.copy_contiguous::(), + DType::U64 => self.copy_contiguous::(), + DType::U32 => self.copy_contiguous::(), + DType::U16 => self.copy_contiguous::(), + DType::U8 => self.copy_contiguous::(), + DType::Bool(burn_std::BoolStore::Native | burn_std::BoolStore::U8) => { + self.copy_contiguous::() + } + DType::Bool(burn_std::BoolStore::U32) => { + panic!("burn-flex: Bool(U32) storage is not yet supported") + } + _ => panic!("Unsupported dtype for contiguous copy: {:?}", self.dtype), + } + } + + fn copy_contiguous(&self) -> Self { + let src: &[E] = bytemuck::cast_slice(&self.data); + let n = self.layout.num_elements(); + let mut dst = Vec::with_capacity(n); + + // Squeeze size-1 dims and merge adjacent stride-contiguous + // runs so e.g. a permuted `[N, H, W, C]` ConvNeXt layer-norm + // input becomes a plain 2D `[H*W, C]` transpose that the + // tiled copy below handles at near-memcpy speed. Without the + // collapse, the 4D ND fallback scalar-walks the tensor. + let collapsed = collapse_for_copy(self.layout.shape(), self.layout.strides()); + let (shape, strides) = collapsed.as_slices(); + let offset = self.layout.start_offset() as isize; + let all_positive = strides.iter().all(|&s| s >= 0); + + if shape.len() <= 1 && all_positive { + // 0-D scalar or 1-D run with a uniform stride. Empty + // collapsed shape means rank 0 (numel 1); otherwise + // numel is the single dim's size (which may be 0 for + // zero-sized 1D tensors, so don't clamp via `.max(1)`). + let collapsed_numel = if shape.is_empty() { 1 } else { shape[0] }; + debug_assert_eq!(n, collapsed_numel); + // SAFETY: capacity is n; we fill every position below. + unsafe { dst.set_len(n) }; + if shape.is_empty() { + if n > 0 { + dst[0] = src[offset as usize]; + } + } else { + let len = shape[0]; + let stride = strides[0]; + if stride == 1 { + dst[..len].copy_from_slice(&src[offset as usize..offset as usize + len]); + } else { + for (i, slot) in dst.iter_mut().take(len).enumerate() { + let idx = (offset + i as isize * stride) as usize; + *slot = src[idx]; + } + } + } + } else if shape.len() == 2 && all_positive { + // 2D positive-stride (transpose-like): tile both dims so + // reads stay in cache. The loop-nesting chooser inside + // `copy_2d_tiled` picks whichever ordering puts the + // smaller source stride on the innermost loop. + debug_assert_eq!(shape[0] * shape[1], n, "2D strides must cover all elements"); + // SAFETY: capacity is n; `copy_2d_tiled` writes every + // `(row, col)` position exactly once. + unsafe { dst.set_len(n) }; + copy_2d_tiled( + &mut dst, src, offset, shape[0], shape[1], strides[0], strides[1], + ); + } else if all_positive && strides.last().copied() == Some(1) { + // Since the inner stride is 1, we can bulk-copy + // the contiguous inner run for each outer index. + unsafe { dst.set_len(n) }; + copy_inner_contiguous_run(&mut dst, src, offset, shape, strides); + } else { + // General fallback: covers negative strides (flipped + // tensors) and ND layouts that can't collapse to ≤2D. + for idx in crate::strided_index::StridedIter::new(&self.layout) { + dst.push(src[idx]); + } + } + + let bytes = Bytes::from_elems(dst); + let layout = Layout::contiguous(self.layout.shape().clone()); + Self { + data: Arc::new(bytes), + layout, + dtype: self.dtype, + } + } + + /// Reshape tensor. Zero-copy if contiguous. + pub fn reshape(&self, new_shape: Shape) -> Self { + assert_eq!( + self.layout.num_elements(), + new_shape.num_elements(), + "reshape must preserve total elements" + ); + + if let Some(new_layout) = self.layout.reshape(new_shape.clone()) { + Self { + data: Arc::clone(&self.data), + layout: new_layout, + dtype: self.dtype, + } + } else { + // Non-contiguous: copy first + self.to_contiguous().reshape(new_shape) + } + } + + /// Transpose two dimensions. Zero-copy (metadata only). + pub fn transpose(&self, dim1: usize, dim2: usize) -> Self { + Self { + data: Arc::clone(&self.data), + layout: self.layout.transpose(dim1, dim2), + dtype: self.dtype, + } + } + + /// Narrow/slice along a dimension. Zero-copy (metadata only). + pub fn narrow(&self, dim: usize, start: usize, len: usize) -> Self { + Self { + data: Arc::clone(&self.data), + layout: self.layout.narrow(dim, start, len), + dtype: self.dtype, + } + } + + /// Permute dimensions according to axes. Zero-copy (metadata only). + pub fn permute(&self, axes: &[usize]) -> Self { + Self { + data: Arc::clone(&self.data), + layout: self.layout.permute(axes), + dtype: self.dtype, + } + } +} + +impl TensorMetadata for FlexTensor { + type Device = FlexDevice; + + fn dtype(&self) -> DType { + self.dtype + } + + fn shape(&self) -> Shape { + self.layout.shape().clone() + } + + fn rank(&self) -> usize { + self.layout.num_dims() + } + + fn device(&self) -> Self::Device { + FlexDevice + } + + fn can_mut(&self) -> bool { + self.is_unique() + } +} + +/// Max rank we're willing to handle without falling back to the +/// strided iterator. Burn tensors are capped at 8 dims in practice. +const COLLAPSE_MAX_RANK: usize = 8; + +/// Collapsed layout result of [`collapse_for_copy`], stored in stack +/// arrays so `to_contiguous()` doesn't have to hit the allocator on +/// its hot path. +#[derive(Debug, Clone, Copy)] +struct CollapsedLayout { + ndim: usize, + shape: [usize; COLLAPSE_MAX_RANK], + strides: [isize; COLLAPSE_MAX_RANK], +} + +impl CollapsedLayout { + #[inline] + fn as_slices(&self) -> (&[usize], &[isize]) { + (&self.shape[..self.ndim], &self.strides[..self.ndim]) + } +} + +/// Collapse a shape/stride pair into the minimum-rank equivalent +/// layout for a contiguous copy: +/// +/// 1. Squeeze size-1 dims (their stride never gets stepped past 0). +/// 2. Merge adjacent dims `(i, i+1)` when +/// `stride[i] == stride[i+1] * shape[i+1]`, which means the two +/// dims form a single logical run through memory. +/// +/// Canonical example: a 4D ConvNeXt input `[1, 244, 224, 48]` with +/// strides `[2_623_488, 224, 1, 54656]` (from +/// `[N, C, H, W].permute([0, 2, 3, 1])`) collapses to 2D +/// `[54656, 48]` with strides `[1, 54656]`. +/// +/// If the input rank exceeds [`COLLAPSE_MAX_RANK`] the result is +/// left at rank > 2 so the caller falls through to its generic +/// strided path. If the input is rank > `COLLAPSE_MAX_RANK`, we +/// return the original (un-collapsed) layout truncated, which the +/// caller will reject via its `shape.len() == 2` gate. +/// +/// PRECONDITION: the caller must gate on all-positive strides before +/// using the collapsed layout. The merge rule assumes positive +/// strides and will produce iteration-order-incorrect results for +/// flipped tensors. +fn collapse_for_copy(shape: &[usize], strides: &[isize]) -> CollapsedLayout { + let mut out = CollapsedLayout { + ndim: 0, + shape: [0; COLLAPSE_MAX_RANK], + strides: [0; COLLAPSE_MAX_RANK], + }; + + // Bail out to the caller's fallback if the rank is too large to + // fit our stack buffer. In practice this never triggers (burn + // tensors are ≤8 dims), but leaving the `ndim` high signals the + // caller to take the generic strided path. + if shape.len() > COLLAPSE_MAX_RANK { + out.ndim = shape.len().min(COLLAPSE_MAX_RANK); + return out; + } + + // Single forward sweep: squeeze size-1 dims and merge whenever + // the current dim's `stride * size` equals the previous output + // dim's stride (i.e. the two form a contiguous run). + // + // Use `checked_mul` so a pathological layout whose stride math + // would overflow `isize` simply fails to merge rather than + // wrapping into an incorrect merge decision. Real tensors can't + // hit this (total numel is bounded by `isize::MAX`), but + // hand-built layouts passed through the test paths could. + for (&s, &st) in shape.iter().zip(strides.iter()) { + if s == 1 { + continue; + } + let merge = out.ndim > 0 + && (s as isize) + .checked_mul(st) + .is_some_and(|run| out.strides[out.ndim - 1] == run); + if merge { + out.shape[out.ndim - 1] *= s; + out.strides[out.ndim - 1] = st; + } else { + out.shape[out.ndim] = s; + out.strides[out.ndim] = st; + out.ndim += 1; + } + } + + out +} + +/// Copy a strided source into a contiguous destination by treating the +/// layout as a set of strided outer dims wrapping a contiguous inner run. +/// +/// Precondition: strides are all positive and the innermost stride is 1 +#[inline] +fn copy_inner_contiguous_run( + dst: &mut [E], + src: &[E], + offset: isize, + shape: &[usize], + strides: &[isize], +) { + debug_assert!(!shape.is_empty()); + debug_assert_eq!(*strides.last().unwrap(), 1, "innermost stride must be 1"); + + // We want the maximal trailing run that is contiguous in src. + // Walk inward-out, a dim dim_idx extends the contiguous run if AND ONLY IF ("iff") + // strides[dim_idx] == product of inner shapes seen so far. + let mut expected_stride: isize = 1; + let mut split = shape.len(); // first OUTER index (dims [split..] are inner) + while split > 0 { + let dim_idx = split - 1; + if strides[dim_idx] == expected_stride { + expected_stride = expected_stride + .checked_mul(shape[dim_idx] as isize) + .expect("contiguous run length overflows isize"); + split -= 1; + } else { + break; + } + } + + let outer_shape = &shape[..split]; + let outer_strides = &strides[..split]; + let inner_run_length: usize = shape[split..].iter().product(); + + // Whole tensor is one contiguous run! We can bulk-copy the whole thing. + if outer_shape.is_empty() { + let src_index = offset as usize; + dst[..inner_run_length].copy_from_slice(&src[src_index..src_index + inner_run_length]); + return; + } + + if inner_run_length == 0 || outer_shape.contains(&0) { + return; + } + + // Walk the outer index space in row-major order via an odometer: + // each iteration ticks the rightmost counter, and so on overflow we + // reset it and carry +1 into the dim to its left. + // Row-major is what we want because dst is contiguous (consecutive outer-index + // tuples land in consecutive inner_run_length slots, so we just bump + // dst_position each step instead of recomputing it) + // We use this instead of N nested for loops because the rank isn't + // known at compile time. + let mut counter = [0usize; COLLAPSE_MAX_RANK]; + let outer_dim_count = outer_shape.len(); + let mut dst_position = 0usize; + + loop { + let mut src_start = offset; + for dim_idx in 0..outer_dim_count { + src_start += counter[dim_idx] as isize * outer_strides[dim_idx]; + } + let src_index = src_start as usize; + dst[dst_position..dst_position + inner_run_length] + .copy_from_slice(&src[src_index..src_index + inner_run_length]); + dst_position += inner_run_length; + + let mut dim_idx = outer_dim_count; + loop { + if dim_idx == 0 { + return; + } + dim_idx -= 1; + counter[dim_idx] += 1; + if counter[dim_idx] < outer_shape[dim_idx] { + break; + } + counter[dim_idx] = 0; + } + } +} + +/// Tiled 2D copy from a strided source into a contiguous destination. +/// The loop nesting is chosen so the innermost read walks whichever +/// source stride is smaller, which keeps the hot loop in cache even +/// for transpose-like layouts. +#[inline] +fn copy_2d_tiled( + dst: &mut [E], + src: &[E], + offset: isize, + rows: usize, + cols: usize, + row_stride: isize, + col_stride: isize, +) { + const TILE: usize = 16; + + if row_stride <= col_stride { + // row-inside-col: the inner loop walks `row_stride` (smaller). + for col_tile in (0..cols).step_by(TILE) { + let col_end = (col_tile + TILE).min(cols); + for row_tile in (0..rows).step_by(TILE) { + let row_end = (row_tile + TILE).min(rows); + for col in col_tile..col_end { + let col_base = offset + col as isize * col_stride; + for row in row_tile..row_end { + let idx = (col_base + row as isize * row_stride) as usize; + // SAFETY: caller set `dst.len() == rows * cols` + // and each `(row, col)` is visited once. + unsafe { + *dst.get_unchecked_mut(row * cols + col) = src[idx]; + } + } + } + } + } + } else { + // col-inside-row: the inner loop walks `col_stride` (smaller). + for row_tile in (0..rows).step_by(TILE) { + let row_end = (row_tile + TILE).min(rows); + for col_tile in (0..cols).step_by(TILE) { + let col_end = (col_tile + TILE).min(cols); + for row in row_tile..row_end { + let row_base = + offset + row as isize * row_stride + col_tile as isize * col_stride; + let dst_base = row * cols + col_tile; + for c in 0..(col_end - col_tile) { + let idx = (row_base + c as isize * col_stride) as usize; + // SAFETY: same as above. + unsafe { + *dst.get_unchecked_mut(dst_base + c) = src[idx]; + } + } + } + } + } + } +} + +/// Get the size in bytes for a dtype element. +/// +/// Matches `burn_std::DType::size()` semantics: Bool(Native) and Bool(U8) are +/// 1 byte, Bool(U32) is 4 bytes. This makes buffer-size validation correct +/// regardless of which BoolStore variant the dtype carries. +/// +/// # Panics +/// +/// Panics if the dtype has a zero-byte element size. `burn_std::DType::size()` +/// returns 0 for sub-byte quantized dtypes (Q4F, Q4S, Q2F, Q2S, and most +/// `QuantStore::PackedNative` variants). burn-flex does not yet support these +/// packed quantization formats; passing them here would silently produce +/// empty allocations in `FlexTensor::empty`, truncated buffers in `into_data`, +/// and zero-byte memcpys in `repeat_dim`. The panic turns all three into a +/// loud, actionable failure at the dispatch boundary. +pub(crate) fn dtype_size(dtype: DType) -> usize { + // Delegate to burn-std's canonical size to stay in sync. + let size = dtype.size(); + assert!( + size > 0, + "burn-flex: dtype {:?} has zero-byte element size (sub-byte packed \ + quantization is not yet supported)", + dtype + ); + size +} + +#[cfg(test)] +mod tests { + use super::*; + use alloc::vec; + + #[test] + fn test_from_data_roundtrip() { + let data = TensorData::from([1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]); + let tensor = FlexTensor::from_data(data.clone()); + let result = tensor.into_data(); + assert_eq!(data.shape, result.shape); + assert_eq!(data.dtype, result.dtype); + } + + #[test] + fn test_collapse_for_copy_squeezes_size1_and_merges_contig() { + // Permuted ConvNeXt input: [1, 48, 244, 224].permute([0,2,3,1]). + let shape = vec![1, 244, 224, 48]; + let strides = vec![2_623_488_isize, 224, 1, 54656]; + let collapsed = collapse_for_copy(&shape, &strides); + let (s, st) = collapsed.as_slices(); + assert_eq!(s, &[54656, 48]); + assert_eq!(st, &[1, 54656]); + } + + #[test] + fn test_collapse_for_copy_already_contiguous_3d() { + let collapsed = collapse_for_copy(&[2, 3, 4], &[12, 4, 1]); + let (s, st) = collapsed.as_slices(); + assert_eq!(s, &[24]); + assert_eq!(st, &[1]); + } + + #[test] + fn test_collapse_for_copy_transpose_2d() { + let collapsed = collapse_for_copy(&[5, 3], &[1, 5]); + let (s, st) = collapsed.as_slices(); + assert_eq!(s, &[5, 3]); + assert_eq!(st, &[1, 5]); + } + + #[test] + fn test_collapse_for_copy_all_size1() { + let collapsed = collapse_for_copy(&[1, 1, 1], &[0, 0, 0]); + let (s, st) = collapsed.as_slices(); + assert!(s.is_empty()); + assert!(st.is_empty()); + } + + /// Regression: an empty 1D view produced by `narrow` at a + /// non-zero offset forces `copy_contiguous` to run (it can't + /// early-return via the contiguous-at-offset-0 shortcut). The + /// old `debug_assert_eq!(n, shape.product().max(1))` tripped + /// for this shape because `.max(1)` produced 1 while the true + /// numel is 0. + #[test] + fn test_to_contiguous_zero_sized_narrowed() { + let t = FlexTensor::from_data(TensorData::new( + (0..6).map(|i| i as f32).collect::>(), + vec![6], + )); + // narrow(dim, start=3, len=0): shape [0], start_offset 3. + let empty_view = t.narrow(0, 3, 0); + assert_eq!(empty_view.shape().to_vec(), vec![0]); + assert_ne!(empty_view.layout().start_offset(), 0); + + let contig = empty_view.to_contiguous(); + assert_eq!(contig.shape().to_vec(), vec![0]); + assert_eq!(contig.layout().start_offset(), 0); + assert_eq!(contig.into_data().bytes.len(), 0); + } + + /// Regression for #4855: a prefix view (e.g. `narrow(dim, 0, n)`) has + /// canonical contiguous strides and start_offset 0, but its underlying + /// buffer is still the larger original. `to_contiguous` must materialize + /// a right-sized copy so callers keying off `storage().len()` (like the + /// SIMD `mask_fill_*` kernels reached from `triu`/`tril` in LU on tall + /// matrices) don't walk past the logical shape. + #[test] + fn test_to_contiguous_prefix_view_shrinks_buffer() { + let data: Vec = (0..40).map(|i| i as f32).collect(); + let t = FlexTensor::from_data(TensorData::new(data, vec![8, 5])); + + let prefix = t.narrow(0, 0, 5); + assert_eq!(prefix.shape().to_vec(), vec![5, 5]); + assert_eq!(prefix.layout().strides(), &[5, 1]); + assert_eq!(prefix.layout().start_offset(), 0); + assert!(prefix.is_contiguous()); + assert_eq!(prefix.storage::().len(), 40); + + let contig = prefix.to_contiguous(); + assert_eq!(contig.storage::().len(), 25); + assert_eq!(contig.layout().num_elements(), 25); + assert_eq!( + contig.storage::(), + &(0..5) + .flat_map(|r| (0..5).map(move |c| (r * 5 + c) as f32)) + .collect::>()[..] + ); + } + + /// 4D permuted layout round-trips through the collapse + tiled + /// copy path. Mirrors the ConvNeXt channels-last permute. + #[test] + fn test_to_contiguous_4d_permuted_matches_naive() { + let dims = [1, 48, 4, 5]; + let n: usize = dims.iter().product(); + let data: Vec = (0..n).map(|i| i as f32).collect(); + let t = FlexTensor::from_data(TensorData::new(data.clone(), dims.to_vec())); + let permuted = t.permute(&[0, 2, 3, 1]); + assert!(!permuted.is_contiguous()); + + let contig = permuted.to_contiguous(); + assert!(contig.is_contiguous()); + assert_eq!(contig.shape().to_vec(), vec![1, 4, 5, 48]); + + // Expected via manual strided walk of the source. + let mut expected = Vec::with_capacity(n); + for h in 0..4 { + for w in 0..5 { + for c in 0..48 { + let idx = c * 20 + h * 5 + w; + expected.push(data[idx]); + } + } + } + + let result_data = contig.into_data(); + let values = result_data.as_slice::().unwrap(); + assert_eq!(values, expected.as_slice()); + } + + /// Regression: the ShuffleNet channel shuffle (reshape -> permute([0,2,1,3,4]) -> reshape) + /// collapses to a 3D layout with stride-1 innermost. + /// Before this fix that went into the StridedIter scalar walk in the final else in copy_contiguous + /// Now, we use a new inner-contiguous branch that bulk-copies the trailing run + #[test] + fn test_to_contiguous_3d_inner_stride1_matches_naive() { + let dims = [1, 2, 4, 3, 3]; // [N, G, C, H, W] + let n: usize = dims.iter().product(); + let data: Vec = (0..n).map(|i| i as f32).collect(); + let t = FlexTensor::from_data(TensorData::new(data.clone(), dims.to_vec())); + let permuted = t.permute(&[0, 2, 1, 3, 4]); // swap G and C from the original dims in the permute + assert!(!permuted.is_contiguous()); + + let contiguous_data = permuted.to_contiguous(); + assert!(contiguous_data.is_contiguous()); + assert_eq!(contiguous_data.shape().to_vec(), vec![1, 4, 2, 3, 3]); + + // Now we manually do a strided walk to be able to check that against a naive solution. + // output[0, c, g, h, w] = input[0, g, c, h, w]. + let mut expected = Vec::with_capacity(n); + for c in 0..4 { + for g in 0..2 { + for h in 0..3 { + for w in 0..3 { + let idx = g * 36 + c * 9 + h * 3 + w; + expected.push(data[idx]); + } + } + } + } + + // Verify the naive solution matches (the manual strided walk). + let result_data = contiguous_data.into_data(); + let values = result_data.as_slice::().unwrap(); + assert_eq!(values, expected.as_slice()); + } + + /// Exercise the `row_stride > col_stride` branch of the 2D tiled + /// copy (the ConvNeXt case hits the other branch). + #[test] + fn test_to_contiguous_2d_row_stride_gt_col_stride() { + // `slice(s![..;2, ..])` on a [6, 3] contiguous tensor gives a + // [3, 3] view with strides [6, 1] that doesn't collapse, so + // the 2D branch runs with row_stride > col_stride. + let data: Vec = (0..18).map(|i| i as f32).collect(); + let t = FlexTensor::from_data(TensorData::new(data, vec![6, 3])); + let stepped = crate::ops::slice::slice( + t, + &[ + burn_std::Slice::new(0, Some(6), 2), + burn_std::Slice::new(0, None, 1), + ], + ); + // Verify the layout matches what the branch requires. + assert_eq!(stepped.layout().shape().to_vec(), vec![3, 3]); + assert_eq!(stepped.layout().strides(), &[6, 1]); + assert!(!stepped.layout().is_contiguous()); + + let contig = stepped.to_contiguous(); + assert!(contig.is_contiguous()); + assert_eq!(contig.shape().to_vec(), vec![3, 3]); + + let result_data = contig.into_data(); + let values = result_data.as_slice::().unwrap(); + // Expected: rows 0, 2, 4 of the original 6x3 tensor. + let expected = vec![ + 0.0f32, 1.0, 2.0, // row 0 + 6.0, 7.0, 8.0, // row 2 + 12.0, 13.0, 14.0, // row 4 + ]; + assert_eq!(values, expected.as_slice()); + } + + #[test] + fn test_reshape() { + let data = TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]); + let tensor = FlexTensor::from_data(data); + let reshaped = tensor.reshape(Shape::from(vec![3, 2])); + assert_eq!(reshaped.shape().to_vec(), vec![3, 2]); + } + + #[test] + fn test_transpose() { + let data = TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]); + let tensor = FlexTensor::from_data(data); + let transposed = tensor.transpose(0, 1); + assert_eq!(transposed.shape().to_vec(), vec![3, 2]); + assert!(!transposed.is_contiguous()); + } + + #[test] + fn test_clone_is_cheap() { + let data = TensorData::from([1.0f32, 2.0, 3.0, 4.0]); + let tensor = FlexTensor::from_data(data); + + // Before clone, tensor is unique + assert!(tensor.is_unique()); + + // Clone shares data + let cloned = tensor.clone(); + assert!(!tensor.is_unique()); + assert!(!cloned.is_unique()); + + // Both point to same data + assert!(core::ptr::eq( + tensor.bytes().as_ptr(), + cloned.bytes().as_ptr(), + )); + } + + #[test] + fn test_cow_on_mutation() { + let data = TensorData::from([1.0f32, 2.0, 3.0, 4.0]); + let tensor = FlexTensor::from_data(data); + let mut cloned = tensor.clone(); + + // Both share data + assert!(!tensor.is_unique()); + assert!(!cloned.is_unique()); + + // Mutate cloned - triggers COW + let storage: &mut [f32] = cloned.storage_mut(); + storage[0] = 99.0; + + // Now cloned has its own copy, tensor is unique again + assert!(tensor.is_unique()); + assert!(cloned.is_unique()); + + // Data is different + assert_ne!(tensor.bytes().as_ptr(), cloned.bytes().as_ptr()); + assert_eq!(tensor.storage::()[0], 1.0); + assert_eq!(cloned.storage::()[0], 99.0); + } + + #[test] + fn test_into_data_narrowed_at_offset_zero() { + // [1, 2, 3, 4, 5, 6] shape [2, 3] + let data = TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]); + let tensor = FlexTensor::from_data(data); + // narrow to first row: shape [1, 3], offset 0, contiguous + let narrowed = tensor.narrow(0, 0, 1); + assert!(narrowed.is_contiguous()); + assert_eq!(narrowed.layout().start_offset(), 0); + + let result = narrowed.into_data(); + assert_eq!(result.shape.to_vec(), vec![1, 3]); + // Must have exactly 3 f32s = 12 bytes, not 24 + assert_eq!(result.bytes.len(), 3 * core::mem::size_of::()); + let values: Vec = result.to_vec().unwrap(); + assert_eq!(values, vec![1.0, 2.0, 3.0]); + } +} diff --git a/crates/burn-fusion/Cargo.toml b/crates/burn-fusion/Cargo.toml new file mode 100644 index 0000000..50bc328 --- /dev/null +++ b/crates/burn-fusion/Cargo.toml @@ -0,0 +1,48 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science"] +description = "Kernel fusion backend decorator for the Burn framework" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "data"] +license.workspace = true +name = "burn-fusion" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-fusion" +documentation = "https://docs.rs/burn-fusion" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["std", "tracing"] +std = ["serde/std", "tracing?/std"] +doc = ["default"] +memory-checks = ["std"] + +tracing = [ + "dep:tracing", + "burn-backend/tracing", + "burn-ir/tracing", +] + +# Exposes the fusion spy API for tests that need to assert on fusion decisions. +# Not intended for production use. +test-util = ["std"] + +[dependencies] +burn-backend = { workspace = true, features = ["default"] } +burn-std = { workspace = true, features = ["default"] } +burn-ir = { workspace = true, features = ["default"] } +tracing = { workspace = true, optional = true, features = ["attributes"] } +smallvec = { workspace = true } + +hashbrown = { workspace = true } +derive-new = { workspace = true } +spin = { workspace = true } +log = { workspace = true } +serde = { workspace = true } + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-fusion/LICENSE-APACHE b/crates/burn-fusion/LICENSE-APACHE new file mode 120000 index 0000000..1cd601d --- /dev/null +++ b/crates/burn-fusion/LICENSE-APACHE @@ -0,0 +1 @@ +../../LICENSE-APACHE \ No newline at end of file diff --git a/crates/burn-fusion/LICENSE-MIT b/crates/burn-fusion/LICENSE-MIT new file mode 120000 index 0000000..b2cfbdc --- /dev/null +++ b/crates/burn-fusion/LICENSE-MIT @@ -0,0 +1 @@ +../../LICENSE-MIT \ No newline at end of file diff --git a/crates/burn-fusion/README.md b/crates/burn-fusion/README.md new file mode 100644 index 0000000..8c508f5 --- /dev/null +++ b/crates/burn-fusion/README.md @@ -0,0 +1,3 @@ +# Burn Fusion + +A kernel fusion backend decorator for Burn. diff --git a/crates/burn-fusion/src/backend.rs b/crates/burn-fusion/src/backend.rs new file mode 100644 index 0000000..4909883 --- /dev/null +++ b/crates/burn-fusion/src/backend.rs @@ -0,0 +1,319 @@ +use crate::{ + FusionTensor, + client::GlobalFusionClient, + stream::{Context, OrderedExecution}, +}; +use burn_backend::{ + Backend, BackendGraph, BackendTypes, DType, DeviceOps, ExecutionError, + tensor::{BoolTensor, Device, FloatTensor, IntTensor, QuantizedTensor}, +}; +use burn_ir::{BackendIr, HandleContainer, OperationIr, TensorHandle, TensorIr}; +use serde::{Serialize, de::DeserializeOwned}; +use std::marker::PhantomData; + +/// Get the client for the given device. +pub fn get_client(device: &Device) -> Client { + GlobalFusionClient::load(device) +} + +/// Enable dynamic operation fusion on a backend that implements [fusion backend](crate::FusionBackend). +#[derive(Clone, Debug, Default)] +pub struct Fusion { + _backend: PhantomData, +} + +impl BackendTypes for Fusion { + type Device = B::Device; + + type FloatTensorPrimitive = FusionTensor; + type IntTensorPrimitive = FusionTensor; + type BoolTensorPrimitive = FusionTensor; + type QuantizedTensorPrimitive = FusionTensor; + + // Fusion adds nothing to a captured graph: the fused kernels are recorded + // on the inner backend's stream, so the graph handle is the inner one. + type GraphPrimitive = B::GraphPrimitive; +} + +impl Backend for Fusion { + fn name(device: &Self::Device) -> String { + format!("fusion<{}>", B::name(device)) + } + + fn seed(device: &B::Device, seed: u64) { + let client = GlobalFusionClient::::load(device); + let device = device.clone(); + client.sync(move || B::seed(&device, seed)); + } + + fn sync(device: &Self::Device) -> Result<(), ExecutionError> { + let client = GlobalFusionClient::::load(device); + let device = device.clone(); + client.sync(move || B::sync(&device)) + } + + fn ad_enabled(_device: &Self::Device) -> bool { + false + } + + fn memory_persistent_allocations< + Output: Send, + Input: Send, + Func: Fn(Input) -> Output + Send, + >( + device: &Self::Device, + input: Input, + func: Func, + ) -> Output { + B::memory_persistent_allocations(device, input, func) + } + + fn memory_cleanup(device: &Self::Device) { + B::memory_cleanup(device) + } + + fn staging<'a, Iter>(data: Iter, device: &Self::Device) + where + Iter: Iterator, + { + B::staging(data, device); + } + + fn supports_dtype(device: &Self::Device, dtype: DType) -> bool { + B::supports_dtype(device, dtype) + } + + fn dtype_usage(device: &Self::Device, dtype: DType) -> burn_backend::DTypeUsageSet { + B::dtype_usage(device, dtype) + } + + fn device_count(type_id: u16) -> usize { + B::device_count(type_id) + } + + fn flush(device: &Self::Device) { + let client = GlobalFusionClient::::load(device); + let device = device.clone(); + client.sync(move || B::flush(&device)) + } + + fn graph_prepare(device: &Self::Device) -> Result<(), ExecutionError> { + let client = GlobalFusionClient::::load(device); + let device = device.clone(); + // `client.sync` drains the lazy queue and runs the closure without a + // device sync — the inner backend then arms persistent allocation. + client.sync(move || B::graph_prepare(&device)) + } + + fn graph_start_capture(device: &Self::Device) -> Result<(), ExecutionError> { + let client = GlobalFusionClient::::load(device); + let device = device.clone(); + // Drain any pending fused operations so nothing leaks into the capture + // window before it opens, then begin capture on the same stream. + client.sync(move || B::graph_start_capture(&device)) + } + + fn graph_stop_capture(device: &Self::Device) -> Result, ExecutionError> { + let client = GlobalFusionClient::::load(device); + let device = device.clone(); + // Drain the capture run's fused kernels onto the captured stream (this + // is where they actually launch), then close the capture. + client.sync(move || B::graph_stop_capture(&device)) + } + + unsafe fn graph_replay( + device: &Self::Device, + graph: &BackendGraph, + ) -> Result<(), ExecutionError> { + let client = GlobalFusionClient::::load(device); + // Drain any queued fused operations onto the stream first (no device + // sync), then launch the captured graph after them. + client.sync(|| ()); + // Safety: forwarded verbatim from this method's own contract. + unsafe { B::graph_replay(device, graph) } + } +} + +/// The status of a [fuser](OperationFuser). +#[derive(Clone, Debug, Copy, PartialEq, Eq)] +pub enum FuserStatus { + /// No more operations can be fused. + Closed, + /// More operations can be fused. + Open, +} + +/// The properties of a [fuser](OperationFuser). +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct FuserProperties { + /// The score of the optimization, higher is better. + pub score: u64, + /// If the operation is ready to be executed. + pub ready: bool, +} + +/// The fusion operation abstraction allows implementations to fuse many +/// [tensor operations](OperationIr) into one, improving the performance of the backend. +/// +/// +/// # Notes +/// +/// The implementations are free to execute the registered operations the way they want to improve +/// the speed and efficiency of the computational graph. It doesn't mean that all registered +/// operations should be fused, but that another way of executing them is more efficient. +/// +/// Also, it is important to return (FuserStatus::Closed) when no more registered operation can +/// improve the performance. +pub trait OperationFuser: Send { + /// Register a new [tensor operation](OperationIr). + fn fuse(&mut self, operation: &OperationIr); + /// Finish the optimization and create a fusion operation. + fn finish(&mut self) -> O; + /// Reset the state. + fn reset(&mut self); + /// Return the builder [status](FuserStatus). + fn status(&self) -> FuserStatus; + /// Return the builder [properties](FuserProperties). + fn properties(&self) -> FuserProperties; + /// The number of operation fused. + fn len(&self) -> usize; + /// If no operations are fused. + fn is_empty(&self) -> bool { + self.len() == 0 + } + /// Clone the optimization builder. + fn clone_dyn(&self) -> Box>; +} + +/// The number of operations contained in the data structure. +pub trait NumOperations: core::fmt::Debug { + /// The number of registered operations. + fn len(&self) -> usize; + /// If the current optimization is empty. + fn is_empty(&self) -> bool { + self.len() == 0 + } + /// The name of the optimization. + fn name(&self) -> &'static str; +} + +/// The optimization created from a [fuser](OperationFuser). +pub trait Optimization: Send + NumOperations { + /// Execute the optimization. + fn execute(&mut self, context: &mut Context, execution: &OrderedExecution); + + /// Returns the state that can be serialized. + fn to_state(&self) -> R::OptimizationState; + /// Create the optimization from the state. + fn from_state(device: &R::FusionDevice, state: R::OptimizationState) -> Self; +} + +/// Type alias for `::FusionDevice`. +pub type FusionDevice = ::FusionDevice; +/// Type alias for `::FusionHandle`. +pub type FusionHandle = ::FusionHandle; +/// Client alias. +pub type Client = GlobalFusionClient; + +/// Trait that defines a runtime that will benefits from fused operations. +pub trait FusionRuntime: Send + Sync + Sized + core::fmt::Debug + 'static { + /// The state that can be serialized for an optimization. + type OptimizationState: Serialize + DeserializeOwned; + /// Optimization type for the backend. + type Optimization: Optimization; + /// Handle used to store tensor dynamically. + type FusionHandle: Clone + Send; + /// Device used by the runtime. + type FusionDevice: DeviceOps; + + /// The list of fusers that will be used to optimize the computational graph. + fn fusers(device: Self::FusionDevice) -> Vec>>; + + /// Create a cross-stream alias of `handle` to register under a fresh tensor id. + /// + /// Called by [`MultiStream::tag_shared_view`](crate::stream::MultiStream::tag_shared_view) when + /// a tensor is shared from one stream to another. The new handle must be an *independent* + /// container entry over the *same* backing buffer, so that consuming one alias (a `ReadWrite` + /// last-use that frees its handle) never frees the buffer out from under the other stream. + /// + /// The default just clones the handle, which is exactly right for local backends whose handle + /// is an `Arc`-style refcount over a device buffer — the clone is a new map entry sharing the + /// allocation. Backends whose handle is a *remote* resource (the router/remote backend, where + /// the handle is a thin id pointing at a server-side tensor) must override this: a bare clone + /// keeps the same server id, so all aliases collapse to one server handle and the first consume + /// frees it for everyone. Such backends allocate a fresh server id aliasing the same buffer. + fn alias_handle(handle: &Self::FusionHandle) -> Self::FusionHandle { + handle.clone() + } + + /// Reclaim a `ReadWrite` (last-use) handle as the fusion engine drains a block. + /// + /// The default removes the container entry; the handle's own `Drop` releases the backend + /// resource. This is correct for local backends, whose handle is an `Arc`-style buffer refcount. + /// + /// A runtime whose handle is a *remote* resource must override this. The router/remote backend + /// frees a tensor by registering an `OperationIr::Drop`, but the drained block **already** freed + /// it server-side — every consumed op (including any `Drop`) is replayed on the server, popping + /// its `ReadWrite` inputs. So letting the handle's `Drop` run here registers a *second*, + /// redundant `Drop` for the same id (the "unfused drop" traffic). The override removes the + /// container entry while suppressing that re-registration. + fn free_handle(handles: &mut HandleContainer, tensor: &TensorIr) { + handles.free(tensor); + } +} + +/// Trait that allows an existing [backend](Backend) to specify graph optimizations using +/// [operation fuser](crate::OperationFuser). +/// +/// Collective operations on the remote/router interpreter path are driven through the +/// [`DistributedOps`](burn_backend::distributed::DistributedOps) supertrait of [`Backend`], which +/// every fusion backend satisfies unconditionally. +pub trait FusionBackend: + BackendIr, Device = FusionDevice> +{ + /// The runtime used for this backend. + type FusionRuntime: FusionRuntime; + + /// Cast a float tensor and returns the resulting handle. + fn cast_float(tensor: FloatTensor, dtype: DType) -> Self::Handle; + + /// Pointer to the full precision fusion backend. + type FullPrecisionBackend: FusionBackend; +} + +// Fusion implements `BackendIr` to enable router backend usage. +impl BackendIr for Fusion { + type Handle = FusionTensor; + + fn float_tensor(handle: TensorHandle) -> FloatTensor { + handle.handle + } + + fn int_tensor(handle: TensorHandle) -> IntTensor { + handle.handle + } + + fn bool_tensor(handle: TensorHandle) -> BoolTensor { + handle.handle + } + + fn quantized_tensor(handle: TensorHandle) -> QuantizedTensor { + handle.handle + } + + fn float_tensor_handle(tensor: FloatTensor) -> Self::Handle { + tensor + } + + fn int_tensor_handle(tensor: IntTensor) -> Self::Handle { + tensor + } + + fn bool_tensor_handle(tensor: BoolTensor) -> Self::Handle { + tensor + } + + fn quantized_tensor_handle(tensor: QuantizedTensor) -> Self::Handle { + tensor + } +} diff --git a/crates/burn-fusion/src/client.rs b/crates/burn-fusion/src/client.rs new file mode 100644 index 0000000..cd4998b --- /dev/null +++ b/crates/burn-fusion/src/client.rs @@ -0,0 +1,434 @@ +use crate::{ + FusionBackend, FusionDevice, FusionHandle, FusionRuntime, FusionServer, FusionTensor, + FusionUtilities, UnfusedOp, + stream::{StreamId, execution::Operation}, +}; +use burn_backend::{Device, DeviceHandle, DeviceId, DeviceService, DeviceServiceStage}; +use burn_std::CommunicationId; + +use burn_backend::{TensorData, backend::ExecutionError}; +use burn_ir::{OperationIr, TensorId, TensorIr}; +use burn_std::stub::RwLock; +use hashbrown::HashSet; +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Use a mutex to communicate with the fusion server. +pub struct GlobalFusionClient { + server: DeviceHandle>, + device: FusionDevice, +} + +impl DeviceService for FusionServer { + fn init(device_id: DeviceId) -> Self { + let device = FusionDevice::::from_id(device_id); + let utilities = FusionUtilities { + initialized_comms: RwLock::new(HashSet::default()), + }; + FusionServer::new(device, utilities) + } + + fn utilities(&self) -> burn_backend::ServerUtilitiesHandle { + self.utilities.clone() + } + fn stage() -> DeviceServiceStage { + DeviceServiceStage::Upstream + } +} + +impl Clone for GlobalFusionClient +where + R: FusionRuntime, +{ + fn clone(&self) -> Self { + Self { + server: self.server.clone(), + device: self.device.clone(), + } + } +} +impl GlobalFusionClient +where + R: FusionRuntime + 'static, +{ + /// Loads the client from the given device. + pub fn load(device: &FusionDevice) -> Self { + Self { + device: device.clone(), + server: DeviceHandle::new(device.to_id()), + } + } +} + +static COUNTER: AtomicU64 = AtomicU64::new(0); + +impl GlobalFusionClient +where + R: FusionRuntime + 'static, +{ + /// Create a new client for the given [device](FusionRuntime::FusionDevice). + pub fn new(device: FusionDevice) -> Self { + Self { + device: device.clone(), + server: DeviceHandle::new(device.to_id()), + } + } + + pub(crate) fn tag_shared_view(&self, src_stream: StreamId, src: TensorId, dst: TensorId) { + let _ = self + .server + .submit_blocking(move |server| { + server.tag_shared_view(src_stream, src, dst); + }) + .ok(); + } + + /// Register a new [tensor operation intermediate representation](OperationIr) on `stream`. + /// + /// Returns the new (uninitialized) output tensor(s) generated by the registered operation. + /// + /// Cross-stream tensor sharing is resolved out-of-band by `tag_shared_view` before + /// this call (see [`MultiStream`](crate::stream::MultiStream) docs for the full + /// strategy), so every input tensor in `repr` is assumed to live on `stream`. Pass + /// `StreamId::current()` for normal ops and the tensor's own stream for `Drop`. + pub fn register( + &self, + stream: StreamId, + repr: OperationIr, + operation: O, + ) -> Vec> + where + O: Operation + 'static, + { + // Create output tensors returned by this operation + let outputs = repr + .outputs() + .map(|output| { + FusionTensor::new( + output.id, + output.shape.clone(), + output.dtype, + self.clone(), + stream, + ) + }) + .collect(); + + // By doing this comparison, we reduce the number of bytes transferred in the device handle + // queue. + if size_of::() < size_of::>() { + // Here the [`O`] type is smaller than the [`UnfusedOp`] type, so it's better to + // transfer it directly. + self.server.submit(move |server| { + let operation = UnfusedOp::new(operation, stream); + server.register(stream, repr, operation); + }); + } else { + // Here the [`O`] type is larger than the [`UnfusedOp`] type, so it's better to + // first create the [`UnfusedOp`] before transferring it to the server. + let operation = UnfusedOp::new(operation, stream); + + self.server.submit(move |server| { + server.register(stream, repr, operation); + }); + } + + outputs + } + + /// Register all lazy computation. + pub fn sync(&self, sync_fn: impl FnOnce() -> Re + Send + 'static) -> Re { + let id = StreamId::current(); + self.server + .submit_blocking(move |server| { + server.drain_stream(id); + sync_fn() + }) + .unwrap() + } + + /// Flush the operations queue. + pub fn flush_queue(&self) { + self.server.flush_queue(); + } + + /// Create a new (uninitialized) empty tensor handle and returns its corresponding [tensor id](TensorId). + pub fn create_empty_handle(&self) -> TensorId { + let value = COUNTER.fetch_add(1, Ordering::Relaxed); + TensorId::new(value) + } + + /// Get the current device used by all operations handled by this client. + pub fn device(&self) -> &FusionDevice { + &self.device + } + + /// Create a tensor with the given handle and returns its corresponding [tensor id](TensorId). + pub fn register_tensor_handle(&self, handle: FusionHandle) -> TensorId { + let id = self.create_empty_handle(); + + self.server + .submit(move |server| server.handles.register_handle(id, handle)); + + id + } + + /// Read the values contained by a float tensor. + pub fn read_tensor_float( + self, + tensor: TensorIr, + stream: StreamId, + ) -> impl Future> + Send + where + B: FusionBackend, + { + self.server + .submit_blocking(move |server| server.float_data::(tensor, stream)) + .unwrap() + } + + /// Read the values contained by an int tensor. + pub fn read_tensor_int( + self, + tensor: TensorIr, + stream: StreamId, + ) -> impl Future> + Send + where + B: FusionBackend, + { + self.server + .submit_blocking(move |server| server.int_data::(tensor, stream)) + .unwrap() + } + + /// Read the values contained by a bool tensor. + pub fn read_tensor_bool( + self, + tensor: TensorIr, + stream: StreamId, + ) -> impl Future> + Send + where + B: FusionBackend, + { + self.server + .submit_blocking(move |server| server.bool_data::(tensor, stream)) + .unwrap() + } + + /// Read the values contained by a quantized tensor. + pub fn read_tensor_quantized( + self, + tensor: TensorIr, + stream: StreamId, + ) -> impl Future> + Send + where + B: FusionBackend, + { + self.server + .submit_blocking(move |server| server.quantized_data::(tensor, stream)) + .unwrap() + } + + /// Change the client of the given float tensor. + pub fn change_client_float( + tensor: TensorIr, + client_src: Self, + client_dst: Self, + stream: StreamId, + ) -> FusionTensor + where + B: FusionBackend, + { + let dtype = tensor.dtype; + let client_dst_cloned = client_dst.clone(); + let shape = tensor.shape.clone(); + let id = client_src.create_empty_handle(); + + let float_tensor = client_src + .server + .submit_blocking(move |server_src| server_src.read_float::(tensor, stream)) + .unwrap(); + let float_tensor = B::float_to_device(float_tensor, client_dst.device()); + client_dst.server.submit(move |server_dst| { + server_dst + .handles + .register_float_tensor::(&id, float_tensor.clone()); + }); + + FusionTensor::new(id, shape, dtype, client_dst_cloned, StreamId::current()) + } + + /// Change the client of the given int tensor. + pub fn change_client_int( + tensor: TensorIr, + client_src: Self, + client_dst: Self, + stream: StreamId, + ) -> FusionTensor + where + B: FusionBackend, + { + let dtype = tensor.dtype; + let client_dst_cloned = client_dst.clone(); + let shape = tensor.shape.clone(); + let id = client_src.create_empty_handle(); + + let int_tensor = client_src + .server + .submit_blocking(move |server_src| server_src.read_int::(tensor, stream)) + .unwrap(); + let int_tensor = B::int_to_device(int_tensor, client_dst.device()); + client_dst.server.submit(move |server_dst| { + server_dst + .handles + .register_int_tensor::(&id, int_tensor.clone()); + }); + + FusionTensor::new(id, shape, dtype, client_dst_cloned, StreamId::current()) + } + + /// Change the client of the given bool tensor. + pub fn change_client_bool( + tensor: TensorIr, + client_src: Self, + client_dst: Self, + stream: StreamId, + ) -> FusionTensor + where + B: FusionBackend, + { + let dtype = tensor.dtype; + let client_dst_cloned = client_dst.clone(); + let shape = tensor.shape.clone(); + let id = client_src.create_empty_handle(); + + let bool_tensor = client_src + .server + .submit_blocking(move |server_src| server_src.read_bool::(tensor, stream)) + .unwrap(); + let bool_tensor = B::bool_to_device(bool_tensor, client_dst.device()); + client_dst.server.submit(move |server_dst| { + server_dst + .handles + .register_bool_tensor::(&id, bool_tensor.clone()); + }); + + FusionTensor::new(id, shape, dtype, client_dst_cloned, StreamId::current()) + } + + /// Change the client of the given quantized tensor. + pub fn change_client_quantized( + tensor: TensorIr, + client_src: Self, + client_dst: Self, + stream: StreamId, + ) -> FusionTensor + where + B: FusionBackend, + { + let dtype = tensor.dtype; + let client_dst_cloned = client_dst.clone(); + let shape = tensor.shape.clone(); + let id = client_src.create_empty_handle(); + + let q_tensor = client_src + .server + .submit_blocking(move |server_src| server_src.read_quantized::(tensor, stream)) + .unwrap(); + let q_tensor = B::q_to_device(q_tensor, client_dst.device()); + client_dst.server.submit(move |server_dst| { + server_dst + .handles + .register_quantized_tensor::(&id, q_tensor.clone()); + }); + + FusionTensor::new(id, shape, dtype, client_dst_cloned, StreamId::current()) + } + + /// Resolve the given float tensor to a primitive tensor. + pub fn resolve_tensor_float(&self, tensor: FusionTensor) -> B::FloatTensorPrimitive + where + B: FusionBackend, + { + // Must evaluate `into_ir()` on the frontend thread. If evaluated inside + // the closure (server thread), a cross-stream tensor will trigger `shared_view` -> + // `submit_blocking`, causing a fatal channel re-entrancy panic (`BorrowMutError`). + let stream = tensor.stream; + let tensor = tensor.into_ir(); + self.server + .submit_blocking(move |server| { + server.drain_stream(stream); + server.resolve_server_float::(&tensor) + }) + .unwrap() + } + + /// Resolve the given int tensor to a primitive tensor. + pub fn resolve_tensor_int(&self, tensor: FusionTensor) -> B::IntTensorPrimitive + where + B: FusionBackend, + { + // Must evaluate `into_ir()` on the frontend thread. If evaluated inside + // the closure (server thread), a cross-stream tensor will trigger `shared_view` -> + // `submit_blocking`, causing a fatal channel re-entrancy panic (`BorrowMutError`). + let stream = tensor.stream; + let tensor = tensor.into_ir(); + self.server + .submit_blocking(move |server| { + server.drain_stream(stream); + server.resolve_server_int::(&tensor) + }) + .unwrap() + } + + /// Resolve the given bool tensor to a primitive tensor. + pub fn resolve_tensor_bool(&self, tensor: FusionTensor) -> B::BoolTensorPrimitive + where + B: FusionBackend, + { + // Must evaluate `into_ir()` on the frontend thread. If evaluated inside + // the closure (server thread), a cross-stream tensor will trigger `shared_view` -> + // `submit_blocking`, causing a fatal channel re-entrancy panic (`BorrowMutError`). + let stream = tensor.stream; + let tensor = tensor.into_ir(); + self.server + .submit_blocking(move |server| { + server.drain_stream(stream); + server.resolve_server_bool::(&tensor) + }) + .unwrap() + } + + /// Synchronize the collective operations. + pub fn sync_collective(&self, device: &B::Device) + where + B: FusionBackend, + { + // Ensure that all operations are resolved before calling sync_collective. + self.sync(|| ()); + B::sync_collective(device) + } + + /// Ensure that communication between the given devices is initialized. + /// Initializing communication is generally blocking, so we make sure to flush those operations. + pub fn ensure_collective_init(&self, device_ids: Vec) + where + B: FusionBackend, + { + let utilities = self + .server + .utilities() + .downcast::() + .expect("Can downcast to `FusionUtilities`"); + let id = CommunicationId::from(device_ids); + + if !utilities.initialized_comms.read().unwrap().contains(&id) { + // Communication initialization is blocking for the server, so we need to flush right away to make sure other devices + // aren't waiting indefinitely on this initialization call. + // This is already handled by cubecl, but since fusion adds another layer of streams and asynchronous submits, + // we also needed to add some logic here to flush the fusion server. + self.flush_queue(); + let mut initialized_comms = utilities.initialized_comms.write().unwrap(); + initialized_comms.insert(id); + } + } +} diff --git a/crates/burn-fusion/src/inspect.rs b/crates/burn-fusion/src/inspect.rs new file mode 100644 index 0000000..732a2d9 --- /dev/null +++ b/crates/burn-fusion/src/inspect.rs @@ -0,0 +1,521 @@ +//! Test-only introspection into fusion runtime behavior. +//! +//! When a [`FusionInspector`] is installed for a given [`StreamId`], every execution +//! plan that runs on that stream is captured as a [`FusionReport`] describing which +//! [`OperationIr`]s ended up fused into a single kernel and which ran unfused. +//! Handle snapshots from the backing [`HandleContainer`](burn_ir::HandleContainer) +//! are also recorded at quiescent points, enabling leak detection. +//! +//! [`FusionBlock`] and [`BlockKind`] are the same types the logger renders into its +//! Full-level execution table — see [`crate::stream::execution::trace`]. +//! +//! # Usage +//! +//! ```ignore +//! # use burn_fusion::inspect::FusionInspector; +//! # use burn_tensor::StreamId; +//! let stream = StreamId::current(); +//! let inspector = FusionInspector::install(stream); +//! // ... run tensor ops with a fusion-wrapped backend, then sync ... +//! let reports = inspector.drain(); +//! let block = reports[0].assert_single_fused_block(); +//! assert_eq!(block.fuser_name(), Some("ElementWise")); +//! ``` +//! +//! # Threading +//! +//! The fusion server runs on a background thread per device. Reports are captured on +//! that thread and deposited into a process-global registry keyed by [`StreamId`]. +//! Multiple inspectors for *different* streams may coexist freely. Each only sees +//! reports from its own stream, so concurrent tests on different streams are naturally +//! isolated without any serialization. +//! +//! Installing two inspectors for the *same* stream will panic. Use `test_stream()` +//! (which returns a fresh unique [`StreamId`] per call) to avoid accidental sharing. +//! +//! ## Handle leak detection +//! +//! [`HandleContainer`](burn_ir::HandleContainer) is shared per-device across the +//! whole process, so handle snapshots are not stream-scoped. Every live handle on +//! the device appears in the snapshot regardless of which stream created it. +//! Tests that assert on handle counts via [`FusionInspector::new_handles_since_baseline`] +//! must therefore be serialized against each other with `#[serial]` +//! to prevent handles from concurrent tests from appearing as false leaks. + +use burn_ir::{OperationIr, TensorId}; +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, Mutex, OnceLock}; + +use crate::stream::StreamId; +use crate::stream::execution::trace::format_table; +pub use crate::stream::execution::trace::{BlockKind, FusionBlock}; + +/// The list of fusion decisions captured for one execution plan. +#[derive(Debug, Clone)] +pub struct FusionReport { + /// One entry per leaf of the executed strategy, in execution order. + pub blocks: Vec, +} + +impl FusionReport { + /// Render the report as the same human-readable table that fusion logging emits at + /// [`FusionLogLevel::Full`](burn_std::config::fusion::FusionLogLevel). Useful for + /// rich panic messages from inspector-based tests. + pub fn format_table(&self) -> String { + format_table(&self.blocks) + } + + /// Iterate over the fused blocks only. + pub fn fused_blocks(&self) -> impl Iterator { + self.blocks + .iter() + .filter(|b| matches!(b.kind, BlockKind::Fused { .. })) + } + + /// Iterate over the unfused blocks only. + pub fn unfused_blocks(&self) -> impl Iterator { + self.blocks + .iter() + .filter(|b| matches!(b.kind, BlockKind::Unfused)) + } + + /// The total number of operations across all blocks. + pub fn total_operations(&self) -> usize { + self.blocks.iter().map(|b| b.operations.len()).sum() + } + + /// Assert that exactly one block exists and that it is fused. Returns it. + /// + /// # Panics + /// Panics if the report does not contain exactly one fused block. + pub fn assert_single_fused_block(&self) -> &FusionBlock { + assert_eq!( + self.blocks.len(), + 1, + "expected exactly one block, got {}: {:#?}", + self.blocks.len(), + self.blocks, + ); + let block = &self.blocks[0]; + assert!( + matches!(block.kind, BlockKind::Fused { .. }), + "expected block to be fused, got {:?}", + block.kind, + ); + block + } +} + +/// Test-facing helpers on [`FusionBlock`]. Defined here so they aren't part of the +/// always-compiled logging data model. +impl FusionBlock { + /// The fuser name if this block was fused, `None` otherwise. + pub fn fuser_name(&self) -> Option<&'static str> { + match self.kind { + BlockKind::Fused { name, .. } => Some(name), + BlockKind::Unfused => None, + } + } + + /// Returns `true` if the block's operations match the given matchers one-for-one, + /// in order. + pub fn ops_match(&self, matchers: &[OpMatcher]) -> bool { + if self.operations.len() != matchers.len() { + return false; + } + self.operations + .iter() + .zip(matchers.iter()) + .all(|(op, m)| m(op)) + } +} + +/// A boxed predicate over an [`OperationIr`]. The matchers returned from +/// [`matchers`] have this type. +pub type OpMatcher = Box bool + Send + Sync>; + +/// A guard installed via [`FusionInspector::install`]. Holds the shared buffers and +/// clears the global sink on drop. +#[must_use = "the inspector is uninstalled as soon as this guard is dropped"] +pub struct FusionInspector { + buffer: Arc>>, + /// The full set of [`TensorId`]s alive in the + /// [`HandleContainer`](burn_ir::HandleContainer) at the most recent quiescent + /// point. Used to assert stream-isolated leak detection: a TensorId that appears + /// here but wasn't in the baseline (see [`FusionInspector::set_baseline`]) was + /// born during this test and still hasn't been freed. + live_handles: Arc>>>, + /// Snapshot of `live_handles` captured by [`FusionInspector::set_baseline`]. Any + /// handle here is assumed to belong to unrelated work (other running tests, + /// pre-existing state) and excluded from the leak set. + baseline_handles: Arc>>>, + /// Track which stream this inspector belongs to. + stream_id: StreamId, +} + +struct Sink { + buffer: Arc>>, + live_handles: Arc>>>, +} + +/// Maps each observed [`StreamId`] to its active [`Sink`]. +fn registry() -> &'static Mutex> { + static REGISTRY: OnceLock>> = OnceLock::new(); + REGISTRY.get_or_init(|| Mutex::new(HashMap::new())) +} + +impl FusionInspector { + /// Install an inspector for `stream_id`. + /// + /// # Panics + /// Panics if an inspector for the same stream is already active. Two inspectors + /// for *different* streams may coexist freely. + pub fn install(stream_id: StreamId) -> Self { + let buffer = Arc::new(Mutex::new(Vec::new())); + let live_handles = Arc::new(Mutex::new(None)); + let baseline_handles = Arc::new(Mutex::new(None)); + + { + let mut reg = registry().lock().unwrap_or_else(|p| p.into_inner()); + assert!( + !reg.contains_key(&stream_id), + "FusionInspector: a sink for stream {stream_id:?} is already installed. Drop the previous inspector before installing a new one.", + ); + reg.insert( + stream_id, + Sink { + buffer: buffer.clone(), + live_handles: live_handles.clone(), + }, + ); + } + + Self { + stream_id, + buffer, + live_handles, + baseline_handles, + } + } + + /// Take all captured reports, clearing the buffer. + pub fn drain(&self) -> Vec { + let mut buf = self.buffer.lock().unwrap_or_else(|p| p.into_inner()); + core::mem::take(&mut *buf) + } + + /// Peek at currently captured reports without clearing. + pub fn reports(&self) -> Vec { + self.buffer + .lock() + .unwrap_or_else(|p| p.into_inner()) + .clone() + } + + /// The most recent [`HandleContainer`](burn_ir::HandleContainer) size observed by + /// the inspector, or `None` if no snapshot has been taken yet. + /// + /// Snapshots are emitted by the fusion runtime at the same quiescent points as + /// `memory-checks` (after every registered op and after every drain), so calling + /// this immediately after a `Backend::sync` reflects the post-drain state. + pub fn last_handle_count(&self) -> Option { + self.live_handles + .lock() + .unwrap_or_else(|p| p.into_inner()) + .as_ref() + .map(|s| s.len()) + } + + /// Snapshot of the [`TensorId`]s currently in the + /// [`HandleContainer`](burn_ir::HandleContainer), or `None` if nothing has been + /// observed yet. Mirrors [`Self::last_handle_count`] but returns the full ID set + /// so callers can compute stream-isolated deltas. + pub fn last_live_handles(&self) -> Option> { + self.live_handles + .lock() + .unwrap_or_else(|p| p.into_inner()) + .clone() + } + + /// Capture the most recent live-handle snapshot as a baseline. Handles present + /// here are excluded from [`Self::new_handles_since_baseline`]. + /// + /// Call this *after* `Backend::sync` so the fusion server has had a chance to + /// emit a snapshot reflecting the pre-test quiescent state. Tests that use this + /// method should be marked `#[serial]` to prevent handles from concurrent tests + /// from polluting the baseline window. + pub fn set_baseline(&self) { + let current = self + .live_handles + .lock() + .unwrap_or_else(|p| p.into_inner()) + .clone() + .unwrap_or_default(); + *self + .baseline_handles + .lock() + .unwrap_or_else(|p| p.into_inner()) = Some(current); + } + + /// The set of [`TensorId`]s present in the most recent snapshot that were *not* + /// in the baseline captured by [`Self::set_baseline`] — i.e., handles born + /// during this test's window that are still alive. + /// + /// If no baseline was set, every live handle is considered "new". An empty set + /// after the test syncs means no handles leaked. + pub fn new_handles_since_baseline(&self) -> HashSet { + let live = self + .live_handles + .lock() + .unwrap_or_else(|p| p.into_inner()) + .clone() + .unwrap_or_default(); + let baseline = self + .baseline_handles + .lock() + .unwrap_or_else(|p| p.into_inner()) + .clone() + .unwrap_or_default(); + live.difference(&baseline).copied().collect() + } +} + +impl Drop for FusionInspector { + fn drop(&mut self) { + if let Ok(mut reg) = registry().lock() { + reg.remove(&self.stream_id); + } + } +} + +/// Whether an inspector is currently installed. Called from the fusion server thread. +pub(crate) fn is_installed() -> bool { + registry().lock().map(|r| !r.is_empty()).unwrap_or(false) +} + +/// Push a report into the sink registered for `stream_id`, if any. +/// Called from the fusion server thread. +pub(crate) fn emit(stream_id: StreamId, blocks: &[FusionBlock]) { + let reg = match registry().lock() { + Ok(r) => r, + Err(_) => return, + }; + let Some(sink) = reg.get(&stream_id) else { + return; + }; + let report = FusionReport { + blocks: blocks.to_vec(), + }; + if let Ok(mut buf) = sink.buffer.lock() { + buf.push(report); + } +} + +/// Record the live [`TensorId`]s for `stream_id` at a quiescent point. +pub(crate) fn emit_handle_snapshot(stream_id: StreamId, ids: impl IntoIterator) { + let reg = match registry().lock() { + Ok(r) => r, + Err(_) => return, + }; + let Some(sink) = reg.get(&stream_id) else { + return; + }; + let set: HashSet = ids.into_iter().collect(); + if let Ok(mut slot) = sink.live_handles.lock() { + *slot = Some(set); + } +} + +/// Small library of matchers for common operations, to keep tests readable. +pub mod matchers { + use super::OpMatcher; + use burn_backend::DType; + use burn_ir::{ + BaseOperationIr, FloatOperationIr, IntOperationIr, NumericOperationIr, OperationIr, + }; + + /// Matches a float add (`a + b`) on the given dtype. + pub fn is_add_float(dtype: DType) -> OpMatcher { + Box::new(move |op| { + matches!( + op, + OperationIr::NumericFloat(d, NumericOperationIr::Add(_)) if *d == dtype + ) + }) + } + + /// Matches a float multiply (`a * b`) on the given dtype. + pub fn is_mul_float(dtype: DType) -> OpMatcher { + Box::new(move |op| { + matches!( + op, + OperationIr::NumericFloat(d, NumericOperationIr::Mul(_)) if *d == dtype + ) + }) + } + + /// Matches a float multiply by scalar on the given dtype. + pub fn is_mul_scalar_float(dtype: DType) -> OpMatcher { + Box::new(move |op| { + matches!( + op, + OperationIr::NumericFloat(d, NumericOperationIr::MulScalar(_)) if *d == dtype + ) + }) + } + + /// Matches an int bitwise and. + pub fn is_bitwise_and_int() -> OpMatcher { + Box::new(|op| matches!(op, OperationIr::Int(IntOperationIr::BitwiseAnd(_)))) + } + + /// Matches an int bitwise and by scalar. + pub fn is_bitwise_and_scalar_int() -> OpMatcher { + Box::new(|op| matches!(op, OperationIr::Int(IntOperationIr::BitwiseAndScalar(_)))) + } + + /// Matches an int bitwise or. + pub fn is_bitwise_or_int() -> OpMatcher { + Box::new(|op| matches!(op, OperationIr::Int(IntOperationIr::BitwiseOr(_)))) + } + + /// Matches an int bitwise or by scalar. + pub fn is_bitwise_or_scalar_int() -> OpMatcher { + Box::new(|op| matches!(op, OperationIr::Int(IntOperationIr::BitwiseOrScalar(_)))) + } + + /// Matches an int bitwise xor. + pub fn is_bitwise_xor_int() -> OpMatcher { + Box::new(|op| matches!(op, OperationIr::Int(IntOperationIr::BitwiseXor(_)))) + } + + /// Matches an int bitwise xor by scalar. + pub fn is_bitwise_xor_scalar_int() -> OpMatcher { + Box::new(|op| matches!(op, OperationIr::Int(IntOperationIr::BitwiseXorScalar(_)))) + } + + /// Matches an int bitwise not. + pub fn is_bitwise_not_int() -> OpMatcher { + Box::new(|op| matches!(op, OperationIr::Int(IntOperationIr::BitwiseNot(_)))) + } + + /// Matches an int bitwise left shift. + pub fn is_bitwise_left_shift_int() -> OpMatcher { + Box::new(|op| matches!(op, OperationIr::Int(IntOperationIr::BitwiseLeftShift(_)))) + } + + /// Matches an int bitwise left shift by scalar. + pub fn is_bitwise_left_shift_scalar_int() -> OpMatcher { + Box::new(|op| { + matches!( + op, + OperationIr::Int(IntOperationIr::BitwiseLeftShiftScalar(_)) + ) + }) + } + + /// Matches an int bitwise right shift. + pub fn is_bitwise_right_shift_int() -> OpMatcher { + Box::new(|op| matches!(op, OperationIr::Int(IntOperationIr::BitwiseRightShift(_)))) + } + + /// Matches an int bitwise right shift by scalar. + pub fn is_bitwise_right_shift_scalar_int() -> OpMatcher { + Box::new(|op| { + matches!( + op, + OperationIr::Int(IntOperationIr::BitwiseRightShiftScalar(_)) + ) + }) + } + + /// Matches an int-to-float cast. + pub fn is_int_into_float() -> OpMatcher { + Box::new(|op| matches!(op, OperationIr::Int(IntOperationIr::IntoFloat(_)))) + } + + /// Matches a float subtraction (`a - b`) on the given dtype. + pub fn is_sub_float(dtype: DType) -> OpMatcher { + Box::new(move |op| { + matches!( + op, + OperationIr::NumericFloat(d, NumericOperationIr::Sub(_)) if *d == dtype + ) + }) + } + + /// Matches a `sum_dim` reduction on float tensors of the given dtype. + pub fn is_sum_dim_float(dtype: DType) -> OpMatcher { + Box::new(move |op| { + matches!( + op, + OperationIr::NumericFloat(d, NumericOperationIr::SumDim(_)) if *d == dtype + ) + }) + } + + /// Matches `Exp` on the given float dtype. + pub fn is_exp(dtype: DType) -> OpMatcher { + Box::new(move |op| { + matches!( + op, + OperationIr::Float(d, FloatOperationIr::Exp(_)) if *d == dtype + ) + }) + } + + /// Matches `Log` on the given float dtype. + pub fn is_log(dtype: DType) -> OpMatcher { + Box::new(move |op| { + matches!( + op, + OperationIr::Float(d, FloatOperationIr::Log(_)) if *d == dtype + ) + }) + } + + /// Matches a `reshape` on any element type (`BaseFloat`/`BaseInt`/`BaseBool`). + pub fn is_reshape() -> OpMatcher { + use burn_ir::BaseOperationIr; + Box::new(|op| { + matches!( + op, + OperationIr::BaseFloat(BaseOperationIr::Reshape(_)) + | OperationIr::BaseInt(BaseOperationIr::Reshape(_)) + | OperationIr::BaseBool(BaseOperationIr::Reshape(_)) + ) + }) + } + + /// Matches a `swap_dims` on any element type (`BaseFloat`/`BaseInt`/`BaseBool`). + pub fn is_swap_dims() -> OpMatcher { + use burn_ir::BaseOperationIr; + Box::new(|op| { + matches!( + op, + OperationIr::BaseFloat(BaseOperationIr::SwapDims(_)) + | OperationIr::BaseInt(BaseOperationIr::SwapDims(_)) + | OperationIr::BaseBool(BaseOperationIr::SwapDims(_)) + ) + }) + } + + /// Matches a `Cat` (concatenation) on any tensor kind. + pub fn is_cat() -> OpMatcher { + Box::new(|op| { + matches!( + op, + OperationIr::BaseFloat(BaseOperationIr::Cat(_)) + | OperationIr::BaseInt(BaseOperationIr::Cat(_)) + | OperationIr::BaseBool(BaseOperationIr::Cat(_)) + ) + }) + } + + /// Matches the memory-bookkeeping `Drop` op emitted when a tensor is deallocated. + /// These aren't really operations — useful to filter out when counting ops in + /// fusion-shape tests. + pub fn is_drop() -> OpMatcher { + Box::new(|op| matches!(op, OperationIr::Drop(_))) + } +} diff --git a/crates/burn-fusion/src/lib.rs b/crates/burn-fusion/src/lib.rs new file mode 100644 index 0000000..404c9cb --- /dev/null +++ b/crates/burn-fusion/src/lib.rs @@ -0,0 +1,36 @@ +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] + +//! # Burn Fusion +//! +//! This library is a part of the Burn project. It is a standalone crate that +//! can be used to perform automatic operation fusion on backends that support it. + +#[macro_use] +extern crate derive_new; + +/// Client module exposing types to communicate with the fusion server. +pub mod client; +/// Stream module exposing all tensor operations that can be optimized. +pub mod stream; + +/// Search module for stream optimizations. +pub(crate) mod search; + +mod backend; +mod op; +mod ops; +mod server; +mod tensor; + +/// Test-only introspection into fusion runtime behavior — see +/// [`inspect::FusionInspector`]. +#[cfg(feature = "test-util")] +pub mod inspect; + +pub use op::UnfusedOp; +pub(crate) use server::*; + +pub use backend::*; +pub use ops::NoOp; +pub use tensor::*; diff --git a/crates/burn-fusion/src/op.rs b/crates/burn-fusion/src/op.rs new file mode 100644 index 0000000..828cfa7 --- /dev/null +++ b/crates/burn-fusion/src/op.rs @@ -0,0 +1,116 @@ +use burn_backend::StreamId; +use burn_ir::HandleContainer; +use burn_std::arena::ReservedMemory; +use std::{cell::RefCell, sync::Arc}; + +use crate::{FusionRuntime, stream::Operation}; + +const MAX_ITEM_COUNT: usize = 256; +const MAX_ITEM_SIZE: usize = 512; + +type Data = burn_std::arena::Bytes; +type Arena = burn_std::arena::Arena; + +std::thread_local! { + static ARENA: RefCell = const {RefCell::new(Arena::new())}; +} + +/// An [operation](Operation) that isn't fused. +/// +/// This can be executed with [Self::execute]. +pub struct UnfusedOp { + kind: UnfusedOpKind, + stream_id: StreamId, +} + +impl UnfusedOp { + /// Creates a new unfused [operation](Operation) that will execute on the given [StreamId]. + pub fn new + 'static>(op: O, stream_id: StreamId) -> Self { + let arena_item = match Arena::accept::() { + true => ARENA.with_borrow_mut(|arena| arena.reserve()), + false => None, + }; + + let reserved = match arena_item { + Some(val) => val, + None => { + return UnfusedOp { + kind: UnfusedOpKind::Alloc(Arc::new(op)), + stream_id, + }; + } + }; + let reserved = reserved.init(op); + let ptr_execute = shim_execute::; + + UnfusedOp { + kind: UnfusedOpKind::Arena(UnfusedOpInArena { + reserved, + ptr_execute, + }), + stream_id, + } + } + + /// Executes the [operation](Operation) and modifies the given handles. + pub fn execute(&self, handles: &mut HandleContainer) { + self.stream_id.executes(|| match &self.kind { + UnfusedOpKind::Arena(o) => o.execute(handles), + UnfusedOpKind::Alloc(o) => o.execute(handles), + }) + } +} + +#[derive(Debug)] +struct UnfusedOpInArena { + /// The data pointer. + reserved: ReservedMemory, + /// The execute function pointer. + ptr_execute: fn(*const Data, handles: &mut HandleContainer), +} + +impl Clone for UnfusedOpInArena { + fn clone(&self) -> Self { + Self { + reserved: self.reserved.clone(), + ptr_execute: self.ptr_execute, + } + } +} + +impl UnfusedOpInArena { + fn execute(&self, handles: &mut HandleContainer) { + (self.ptr_execute)(self.reserved.as_ref(), handles); + } +} + +fn shim_execute>( + ptr_data: *const Data, + handles: &mut HandleContainer, +) { + let operation: &O = unsafe { (ptr_data as *const O).as_ref().unwrap() }; + operation.execute(handles); +} + +impl Clone for UnfusedOp { + fn clone(&self) -> Self { + Self { + kind: self.kind.clone(), + stream_id: self.stream_id, + } + } +} + +enum UnfusedOpKind { + Arena(UnfusedOpInArena), + Alloc(Arc>), +} + +impl Clone for UnfusedOpKind { + fn clone(&self) -> Self { + match self { + Self::Arena(o) => Self::Arena(o.clone()), + Self::Alloc(o) => Self::Alloc(o.clone()), + } + } +} diff --git a/crates/burn-fusion/src/ops/activation.rs b/crates/burn-fusion/src/ops/activation.rs new file mode 100644 index 0000000..ec95a9a --- /dev/null +++ b/crates/burn-fusion/src/ops/activation.rs @@ -0,0 +1,4 @@ +use crate::{Fusion, FusionBackend}; +use burn_backend::ops::ActivationOps; + +impl ActivationOps for Fusion {} diff --git a/crates/burn-fusion/src/ops/base.rs b/crates/burn-fusion/src/ops/base.rs new file mode 100644 index 0000000..90ccc78 --- /dev/null +++ b/crates/burn-fusion/src/ops/base.rs @@ -0,0 +1,15 @@ +use crate::{FusionBackend, stream::Operation}; +use burn_ir::HandleContainer; +use std::marker::PhantomData; + +/// A no-operation placeholder for the fusion backend. +/// +/// `NoOp` is an implementation of [`Operation`] that doesn't execute anything. +#[derive(new, Clone, Debug)] +pub struct NoOp { + _b: PhantomData, +} + +impl Operation for NoOp { + fn execute(&self, _handles: &mut HandleContainer) {} +} diff --git a/crates/burn-fusion/src/ops/binary.rs b/crates/burn-fusion/src/ops/binary.rs new file mode 100644 index 0000000..16594cc --- /dev/null +++ b/crates/burn-fusion/src/ops/binary.rs @@ -0,0 +1,117 @@ +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! binary_float_ops { + ( + $name:ident, + $ops:expr + ) => { + #[derive(Debug)] + struct $name { + desc: BinaryOpIr, + _b: PhantomData, + } + + impl $name { + fn new(desc: BinaryOpIr) -> Self { + Self { + desc, + _b: PhantomData, + } + } + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = handles.get_float_tensor::(&self.desc.lhs); + let rhs = handles.get_float_tensor::(&self.desc.rhs); + let output = $ops(lhs, rhs); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + }; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! binary_float_cmp_ops { + ( + $name:ident, + $ops:expr + ) => { + #[derive(new, Debug)] + struct $name { + desc: BinaryOpIr, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = handles.get_float_tensor::(&self.desc.lhs); + let rhs = handles.get_float_tensor::(&self.desc.rhs); + let output = $ops(lhs, rhs, self.desc.out.dtype.into()); + + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + }; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! binary_int_cmp_ops { + ( + $name:ident, + $ops:expr + ) => { + #[derive(Debug)] + struct $name { + desc: BinaryOpIr, + _b: PhantomData, + } + + impl $name { + fn new(desc: BinaryOpIr) -> Self { + Self { + desc, + _b: PhantomData, + } + } + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = handles.get_int_tensor::(&self.desc.lhs); + let rhs = handles.get_int_tensor::(&self.desc.rhs); + let output = $ops(lhs, rhs, self.desc.out.dtype.into()); + + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + }; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! binary_int_ops { + ( + $name:ident, + $ops:expr + ) => { + #[derive(new, Debug)] + struct $name { + desc: BinaryOpIr, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = handles.get_int_tensor::(&self.desc.lhs); + let rhs = handles.get_int_tensor::(&self.desc.rhs); + let output = $ops(lhs, rhs); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + }; +} diff --git a/crates/burn-fusion/src/ops/bool_tensor.rs b/crates/burn-fusion/src/ops/bool_tensor.rs new file mode 100644 index 0000000..80f25af --- /dev/null +++ b/crates/burn-fusion/src/ops/bool_tensor.rs @@ -0,0 +1,945 @@ +use crate::{ + Fusion, FusionBackend, + client::GlobalFusionClient, + get_client, + stream::{StreamId, execution::Operation}, +}; +use burn_backend::{ + BoolDType, ExecutionError, FloatDType, IntDType, Scalar, Shape, Slice, TensorData, + ops::BoolTensorOps, + tensor::{BoolTensor, Device, FloatTensor, IndexingUpdateOp, IntTensor}, +}; +use burn_ir::{ + BaseOperationIr, BinaryOpIr, BoolOperationIr, CastOpIr, CatOpIr, CreationOpIr, FlipOpIr, + GatherOpIr, HandleContainer, InitOperationIr, MaskFillOpIr, MaskWhereOpIr, OperationIr, + OperationOutput, PermuteOpIr, RepeatDimOpIr, ScalarOpIr, ScatterOpIr, SelectAssignOpIr, + SelectOpIr, ShapeOpIr, SliceAssignOpIr, SliceOpIr, SwapDimsOpIr, TensorIr, UnaryOpIr, + UnfoldOpIr, +}; +use std::marker::PhantomData; + +use super::NoOp; + +impl BoolTensorOps for Fusion { + fn bool_empty(shape: Shape, device: &Device, dtype: BoolDType) -> BoolTensor { + #[derive(new, Debug)] + struct EmptyOps { + desc: TensorIr, + device: Device, + } + + impl Operation for EmptyOps { + fn execute(&self, handles: &mut HandleContainer) { + let output = B::bool_empty( + self.desc.shape.clone(), + &self.device, + self.desc.dtype.into(), + ); + handles.register_bool_tensor::(&self.desc.id, output); + } + } + + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register( + StreamId::current(), + OperationIr::BaseBool(BaseOperationIr::Empty(desc.clone())), + EmptyOps::::new(desc.out, device.clone()), + ) + .output() + } + + fn bool_zeros(shape: Shape, device: &Device, dtype: BoolDType) -> BoolTensor { + #[derive(new, Debug)] + struct ZerosOps { + desc: TensorIr, + device: Device, + } + + impl Operation for ZerosOps { + fn execute(&self, handles: &mut HandleContainer) { + let output = B::bool_zeros( + self.desc.shape.clone(), + &self.device, + self.desc.dtype.into(), + ); + handles.register_bool_tensor::(&self.desc.id, output); + } + } + + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register( + StreamId::current(), + OperationIr::BaseBool(BaseOperationIr::Zeros(desc.clone())), + ZerosOps::::new(desc.out, device.clone()), + ) + .output() + } + + fn bool_ones(shape: Shape, device: &Device, dtype: BoolDType) -> BoolTensor { + #[derive(new, Debug)] + struct OnesOps { + desc: TensorIr, + device: Device, + } + + impl Operation for OnesOps { + fn execute(&self, handles: &mut HandleContainer) { + let output = B::bool_ones( + self.desc.shape.clone(), + &self.device, + self.desc.dtype.into(), + ); + handles.register_bool_tensor::(&self.desc.id, output); + } + } + + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register( + StreamId::current(), + OperationIr::BaseBool(BaseOperationIr::Ones(desc.clone())), + OnesOps::::new(desc.out, device.clone()), + ) + .output() + } + + async fn bool_into_data(tensor: BoolTensor) -> Result { + tensor.bool_into_data::().await + } + + fn bool_from_data(data: burn_backend::TensorData, device: &Device) -> BoolTensor { + let client = get_client::(device); + let dtype = data.dtype; + let tensor = B::bool_from_data(data, device); + let shape = burn_backend::TensorMetadata::shape(&tensor); + + let handle = B::bool_tensor_handle(tensor); + let desc = InitOperationIr::create(shape, dtype, || client.register_tensor_handle(handle)); + + client + .register( + StreamId::current(), + OperationIr::Init(desc), + NoOp::::new(), + ) + .output() + } + + fn bool_into_int(tensor: BoolTensor, out_dtype: IntDType) -> IntTensor { + #[derive(new, Debug)] + struct IntoIntOps { + desc: CastOpIr, + _b: PhantomData, + } + + impl Operation for IntoIntOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_bool_tensor::(&self.desc.input); + let output = B::bool_into_int(input, self.desc.out.dtype.into()); + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = CastOpIr::create(tensor.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Bool(BoolOperationIr::IntoInt(desc.clone())), + IntoIntOps::::new(desc), + ) + .output() + } + + fn bool_into_float(tensor: BoolTensor, out_dtype: FloatDType) -> FloatTensor { + #[derive(new, Debug)] + struct IntoFloatOps { + desc: CastOpIr, + _b: PhantomData, + } + + impl Operation for IntoFloatOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_bool_tensor::(&self.desc.input); + let output = B::bool_into_float(input, self.desc.out.dtype.into()); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = CastOpIr::create(tensor.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Bool(BoolOperationIr::IntoFloat(desc.clone())), + IntoFloatOps::::new(desc), + ) + .output() + } + + fn bool_to_device(tensor: BoolTensor, device_dst: &Device) -> BoolTensor { + let device_src: &B::Device = tensor.client.device(); + + if device_src == device_dst { + return tensor; + } + + let id = tensor.stream; + let client_dst = get_client::(device_dst); + let client_src = tensor.client.clone(); + + GlobalFusionClient::change_client_bool::(tensor.into_ir(), client_src, client_dst, id) + } + + fn bool_reshape(tensor: BoolTensor, shape: Shape) -> BoolTensor { + if tensor.shape == shape { + return tensor; + } + + #[derive(new, Debug)] + struct ReshapeDimsOps { + desc: ShapeOpIr, + _b: PhantomData, + } + + impl Operation for ReshapeDimsOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_bool_tensor::(&self.desc.input); + let output = B::bool_reshape(input, self.desc.out.shape.clone()); + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ShapeOpIr::reshape(tensor.into_ir(), shape, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::Reshape(desc.clone())), + ReshapeDimsOps::::new(desc), + ) + .output() + } + + fn bool_slice(tensor: BoolTensor, slices: &[Slice]) -> BoolTensor { + #[derive(new, Debug)] + struct SliceOps { + desc: SliceOpIr, + _b: PhantomData, + } + + impl Operation for SliceOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_bool_tensor::(&self.desc.tensor); + + let output = B::bool_slice(tensor, self.desc.ranges.as_slice()); + + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = SliceOpIr::create(tensor.into_ir(), slices.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::Slice(desc.clone())), + SliceOps::::new(desc), + ) + .output() + } + + fn bool_slice_assign( + tensor: BoolTensor, + slices: &[Slice], + value: BoolTensor, + ) -> BoolTensor { + #[derive(new, Debug)] + struct SliceAssignOps { + desc: SliceAssignOpIr, + _b: PhantomData, + } + + impl Operation for SliceAssignOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_bool_tensor::(&self.desc.tensor); + let value = handles.get_bool_tensor::(&self.desc.value); + + let output = B::bool_slice_assign(tensor, self.desc.ranges.as_slice(), value); + + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = + SliceAssignOpIr::create(tensor.into_ir(), slices.into(), value.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::SliceAssign(desc.clone())), + SliceAssignOps::::new(desc), + ) + .output() + } + + fn bool_cat(tensors: Vec>, dim: usize) -> BoolTensor { + #[derive(new, Debug)] + struct CatOps { + desc: CatOpIr, + _b: PhantomData, + } + + impl Operation for CatOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensors = self + .desc + .tensors + .iter() + .map(|tensor| handles.get_bool_tensor::(tensor)) + .collect(); + + let output = B::bool_cat(tensors, self.desc.dim); + + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensors.first().unwrap().client.clone(); + let tensors = tensors.into_iter().map(|t| t.into_ir()).collect(); + let desc = CatOpIr::create(tensors, dim, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::Cat(desc.clone())), + CatOps::::new(desc), + ) + .output() + } + + fn bool_equal(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + #[derive(new, Debug)] + struct EqualOps { + desc: BinaryOpIr, + _b: PhantomData, + } + + impl Operation for EqualOps { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = handles.get_bool_tensor::(&self.desc.lhs); + let rhs = handles.get_bool_tensor::(&self.desc.rhs); + let output = B::bool_equal(lhs, rhs); + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::Equal(desc.clone())), + EqualOps::::new(desc), + ) + .output() + } + + fn bool_not(tensor: BoolTensor) -> BoolTensor { + #[derive(new, Debug)] + struct NotOps { + desc: UnaryOpIr, + _b: PhantomData, + } + + impl Operation for NotOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_bool_tensor::(&self.desc.input); + let output = B::bool_not(input); + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Bool(BoolOperationIr::Not(desc.clone())), + NotOps::::new(desc), + ) + .output() + } + + fn bool_and(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + #[derive(new, Debug)] + struct AndOps { + desc: BinaryOpIr, + _b: PhantomData, + } + + impl Operation for AndOps { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = handles.get_bool_tensor::(&self.desc.lhs); + let rhs = handles.get_bool_tensor::(&self.desc.rhs); + let output = B::bool_and(lhs, rhs); + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Bool(BoolOperationIr::And(desc.clone())), + AndOps::::new(desc), + ) + .output() + } + + fn bool_or(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + #[derive(new, Debug)] + struct OrOps { + desc: BinaryOpIr, + _b: PhantomData, + } + + impl Operation for OrOps { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = handles.get_bool_tensor::(&self.desc.lhs); + let rhs = handles.get_bool_tensor::(&self.desc.rhs); + let output = B::bool_or(lhs, rhs); + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + client + .register( + streams, + OperationIr::Bool(BoolOperationIr::Or(desc.clone())), + OrOps::::new(desc), + ) + .output() + } + + fn bool_swap_dims(tensor: BoolTensor, dim1: usize, dim2: usize) -> BoolTensor { + #[derive(new, Debug)] + struct SwapDimsOps { + desc: SwapDimsOpIr, + _b: PhantomData, + } + + impl Operation for SwapDimsOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_bool_tensor::(&self.desc.input); + let output = B::bool_swap_dims(input, self.desc.dim1, self.desc.dim2); + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = SwapDimsOpIr::create(tensor.into_ir(), dim1, dim2, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::SwapDims(desc.clone())), + SwapDimsOps::::new(desc), + ) + .output() + } + + fn bool_permute(tensor: BoolTensor, axes: &[usize]) -> BoolTensor { + #[derive(new, Debug)] + struct PermuteDimsOps { + desc: PermuteOpIr, + _b: PhantomData, + } + + impl Operation for PermuteDimsOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_bool_tensor::(&self.desc.input); + let output = B::bool_permute(input, self.desc.axes.as_slice()); + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = PermuteOpIr::create(tensor.into_ir(), axes.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::Permute(desc.clone())), + PermuteDimsOps::::new(desc), + ) + .output() + } + + fn bool_expand(tensor: BoolTensor, shape: Shape) -> BoolTensor { + #[derive(new, Debug)] + struct ExpandOps { + desc: ShapeOpIr, + _b: PhantomData, + } + + impl Operation for ExpandOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_bool_tensor::(&self.desc.input); + let output = B::bool_expand(input, self.desc.out.shape.clone()); + + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ShapeOpIr::expand(tensor.into_ir(), shape, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::Expand(desc.clone())), + ExpandOps::::new(desc), + ) + .output() + } + + fn bool_flip(tensor: BoolTensor, axes: &[usize]) -> BoolTensor { + #[derive(new, Debug)] + struct FlipOps { + desc: FlipOpIr, + _b: PhantomData, + } + + impl Operation for FlipOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_bool_tensor::(&self.desc.input); + let output = B::bool_flip(input, self.desc.axes.as_slice()); + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = FlipOpIr::create(tensor.into_ir(), axes.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::Flip(desc.clone())), + FlipOps::::new(desc), + ) + .output() + } + + fn bool_repeat_dim(tensor: BoolTensor, dim: usize, times: usize) -> BoolTensor { + #[derive(new, Debug)] + struct RepeatDimOps { + desc: RepeatDimOpIr, + _b: PhantomData, + } + + impl Operation for RepeatDimOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_bool_tensor::(&self.desc.tensor); + + let output = B::bool_repeat_dim(tensor, self.desc.dim, self.desc.times); + + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = RepeatDimOpIr::create(tensor.into_ir(), dim, times, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::RepeatDim(desc.clone())), + RepeatDimOps::::new(desc), + ) + .output() + } + + fn bool_unfold( + tensor: BoolTensor, + dim: usize, + size: usize, + step: usize, + ) -> BoolTensor { + #[derive(new, Debug)] + struct UnfoldOps { + desc: UnfoldOpIr, + _b: PhantomData, + } + + impl Operation for UnfoldOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_bool_tensor::(&self.desc.input); + let output = B::bool_unfold(input, self.desc.dim, self.desc.size, self.desc.step); + + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnfoldOpIr::create(tensor.into_ir(), dim, size, step, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::Unfold(desc.clone())), + UnfoldOps::::new(desc), + ) + .output() + } + + fn bool_mask_where( + tensor: BoolTensor, + mask: BoolTensor, + value: BoolTensor, + ) -> BoolTensor { + #[derive(new, Debug)] + struct MaskWhereOps { + desc: MaskWhereOpIr, + _b: PhantomData, + } + + impl Operation for MaskWhereOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_bool_tensor::(&self.desc.tensor); + let value = handles.get_bool_tensor::(&self.desc.value); + let mask = handles.get_bool_tensor::(&self.desc.mask); + + let output = B::bool_mask_where(tensor, mask, value); + + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = MaskWhereOpIr::create(tensor.into_ir(), mask.into_ir(), value.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::MaskWhere(desc.clone())), + MaskWhereOps::::new(desc), + ) + .output() + } + + fn bool_mask_fill( + tensor: BoolTensor, + mask: BoolTensor, + value: Scalar, + ) -> BoolTensor { + #[derive(new, Debug)] + struct MaskFillOps { + desc: MaskFillOpIr, + _b: PhantomData, + } + + impl Operation for MaskFillOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_bool_tensor::(&self.desc.tensor); + let mask = handles.get_bool_tensor::(&self.desc.mask); + + let output = B::bool_mask_fill(tensor, mask, self.desc.value.into()); + + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let value = value.into(); + let desc = MaskFillOpIr::create(tensor.into_ir(), mask.into_ir(), value, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::MaskFill(desc.clone())), + MaskFillOps::::new(desc), + ) + .output() + } + + fn bool_gather( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + ) -> BoolTensor { + #[derive(new, Debug)] + struct GatherOps { + desc: GatherOpIr, + _b: PhantomData, + } + + impl Operation for GatherOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_bool_tensor::(&self.desc.tensor); + let indices = handles.get_int_tensor::(&self.desc.indices); + + let output = B::bool_gather(self.desc.dim, tensor, indices); + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = GatherOpIr::create(tensor.into_ir(), dim, indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::Gather(desc.clone())), + GatherOps::::new(desc), + ) + .output() + } + + fn bool_scatter_or( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor { + #[derive(new, Debug)] + struct ScatterOps { + desc: ScatterOpIr, + _b: PhantomData, + } + + impl Operation for ScatterOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_bool_tensor::(&self.desc.tensor); + let indices = handles.get_int_tensor::(&self.desc.indices); + let value = handles.get_bool_tensor::(&self.desc.value); + + let output = B::bool_scatter_or(self.desc.dim, tensor, indices, value); + + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ScatterOpIr::create( + tensor.into_ir(), + dim, + indices.into_ir(), + value.into_ir(), + IndexingUpdateOp::Add, + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::Scatter(desc.clone())), + ScatterOps::::new(desc), + ) + .output() + } + + fn bool_equal_elem(lhs: BoolTensor, rhs: Scalar) -> BoolTensor { + #[derive(new, Debug)] + struct EqualElemOps { + desc: ScalarOpIr, + _b: PhantomData, + } + impl Operation for EqualElemOps { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = handles.get_bool_tensor::(&self.desc.lhs); + let output = B::bool_equal_elem(lhs, self.desc.rhs.into()); + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let dtype = lhs.dtype; + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, dtype, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::EqualElem(desc.clone())), + EqualElemOps::::new(desc), + ) + .output() + } + + fn bool_select( + tensor: BoolTensor, + dim: usize, + indices: IntTensor, + ) -> BoolTensor { + #[derive(new, Debug)] + struct SelectOps { + desc: SelectOpIr, + _b: PhantomData, + } + + impl Operation for SelectOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_bool_tensor::(&self.desc.tensor); + let indices = handles.get_int_tensor::(&self.desc.indices); + + let output = B::bool_select(tensor, self.desc.dim, indices); + + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = SelectOpIr::create(tensor.into_ir(), dim, indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::Select(desc.clone())), + SelectOps::::new(desc), + ) + .output() + } + + fn bool_select_or( + tensor: BoolTensor, + dim: usize, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor { + #[derive(new, Debug)] + struct SelectAssignOps { + desc: SelectAssignOpIr, + _b: PhantomData, + } + + impl Operation for SelectAssignOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_bool_tensor::(&self.desc.tensor); + let indices = handles.get_int_tensor::(&self.desc.indices); + let value = handles.get_bool_tensor::(&self.desc.value); + + let output = B::bool_select_or(tensor, self.desc.dim, indices, value); + + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = SelectAssignOpIr::create( + tensor.into_ir(), + dim, + indices.into_ir(), + value.into_ir(), + IndexingUpdateOp::Add, + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::BaseBool(BaseOperationIr::SelectAssign(desc.clone())), + SelectAssignOps::::new(desc), + ) + .output() + } +} diff --git a/crates/burn-fusion/src/ops/distributed.rs b/crates/burn-fusion/src/ops/distributed.rs new file mode 100644 index 0000000..344f05f --- /dev/null +++ b/crates/burn-fusion/src/ops/distributed.rs @@ -0,0 +1,67 @@ +use std::marker::PhantomData; + +use burn_backend::{ + DeviceId, + distributed::{CollectiveTensor, DistributedOps, ReduceOperation}, + tensor::{Device, FloatTensor}, +}; +use burn_ir::{AllReduceOpIr, DistributedOperationIr, HandleContainer, OperationIr}; + +use crate::{ + Fusion, FusionBackend, get_client, + stream::{Operation, StreamId}, +}; +use burn_ir::OperationOutput; + +impl DistributedOps for Fusion { + fn all_reduce( + tensor: FloatTensor, + op: ReduceOperation, + device_ids: Vec, + ) -> CollectiveTensor { + #[derive(new, Debug)] + struct AllReduceOps { + desc: AllReduceOpIr, + op: ReduceOperation, + device_ids: Vec, + _b: PhantomData, + } + + impl Operation for AllReduceOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_float_tensor::(&self.desc.tensor); + let output = B::all_reduce(tensor, self.op, self.device_ids.clone()); + handles.register_float_tensor::(&self.desc.out.id, unsafe { + output.assume_resolved() + }); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = AllReduceOpIr::create( + tensor.into_ir(), + op, + device_ids.iter().map(|id| (*id).into()).collect(), + || client.create_empty_handle(), + ); + + let output = client + .register( + streams, + OperationIr::Distributed(DistributedOperationIr::AllReduce(desc.clone())), + AllReduceOps::::new(desc, op, device_ids.clone()), + ) + .output(); + + client.ensure_collective_init::(device_ids); + + CollectiveTensor::new(output) + } + + fn sync_collective(device: &Device) { + let client = get_client::(device); + client.sync_collective::(device); + } +} diff --git a/crates/burn-fusion/src/ops/int_tensor.rs b/crates/burn-fusion/src/ops/int_tensor.rs new file mode 100644 index 0000000..2c74e22 --- /dev/null +++ b/crates/burn-fusion/src/ops/int_tensor.rs @@ -0,0 +1,2185 @@ +use super::NoOp; +use crate::{ + Fusion, FusionBackend, binary_int_cmp_ops, binary_int_ops, + client::GlobalFusionClient, + get_client, reduce_int_ops, scalar_int_cmp_ops, scalar_int_ops, + stream::{StreamId, execution::Operation}, + unary_int_ops, +}; +use burn_backend::{ + BoolDType, Distribution, ExecutionError, FloatDType, IntDType, Scalar, Shape, Slice, + TensorData, + ops::IntTensorOps, + tensor::{BoolTensor, Device, FloatTensor, IndexingUpdateOp, IntTensor}, +}; +use burn_ir::*; +use std::marker::PhantomData; + +impl IntTensorOps for Fusion { + fn int_empty(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + #[derive(new, Debug)] + struct EmptyOps { + desc: TensorIr, + device: Device, + } + + impl Operation for EmptyOps { + fn execute(&self, handles: &mut HandleContainer) { + let output = B::int_empty( + self.desc.shape.clone(), + &self.device, + self.desc.dtype.into(), + ); + handles.register_int_tensor::(&self.desc.id, output); + } + } + + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register( + StreamId::current(), + OperationIr::BaseInt(BaseOperationIr::Empty(desc.clone())), + EmptyOps::::new(desc.out, device.clone()), + ) + .output() + } + + async fn int_into_data(tensor: IntTensor) -> Result { + tensor.int_into_data::().await + } + + fn int_from_data(data: TensorData, device: &Device) -> IntTensor { + let client = get_client::(device); + let dtype = data.dtype; + let tensor = B::int_from_data(data, device); + let shape = burn_backend::TensorMetadata::shape(&tensor); + + let handle = B::int_tensor_handle(tensor); + let desc = InitOperationIr::create(shape, dtype, || client.register_tensor_handle(handle)); + + client + .register( + StreamId::current(), + OperationIr::Init(desc), + NoOp::::new(), + ) + .output() + } + + fn int_to_device(tensor: IntTensor, device_dst: &Device) -> IntTensor { + let device_src: &B::Device = tensor.client.device(); + + if device_src == device_dst { + return tensor; + } + + let id = tensor.stream; + let client_dst = get_client::(device_dst); + let client_src = tensor.client.clone(); + + GlobalFusionClient::change_client_int::(tensor.into_ir(), client_src, client_dst, id) + } + + fn int_reshape(tensor: IntTensor, shape: Shape) -> IntTensor { + if tensor.shape == shape { + return tensor; + } + + #[derive(new, Debug)] + struct ReshapeDimsOps { + desc: ShapeOpIr, + _b: PhantomData, + } + + impl Operation for ReshapeDimsOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_int_tensor::(&self.desc.input); + let output = B::int_reshape(input, self.desc.out.shape.clone()); + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ShapeOpIr::reshape(tensor.into_ir(), shape, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::Reshape(desc.clone())), + ReshapeDimsOps::::new(desc), + ) + .output() + } + + fn int_slice(tensor: IntTensor, slices: &[Slice]) -> IntTensor { + #[derive(new, Debug)] + struct SliceOps { + desc: SliceOpIr, + _b: PhantomData, + } + + impl Operation for SliceOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_int_tensor::(&self.desc.tensor); + + let output = B::int_slice(tensor, self.desc.ranges.as_slice()); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = SliceOpIr::create(tensor.into_ir(), slices.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::Slice(desc.clone())), + SliceOps::::new(desc), + ) + .output() + } + + fn int_slice_assign( + tensor: IntTensor, + slices: &[burn_backend::Slice], + value: IntTensor, + ) -> IntTensor { + #[derive(new, Debug)] + struct SliceAssignOps { + desc: SliceAssignOpIr, + _b: PhantomData, + } + + impl Operation for SliceAssignOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_int_tensor::(&self.desc.tensor); + let value = handles.get_int_tensor::(&self.desc.value); + + let output = B::int_slice_assign(tensor, self.desc.ranges.as_slice(), value); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = + SliceAssignOpIr::create(tensor.into_ir(), slices.into(), value.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::SliceAssign(desc.clone())), + SliceAssignOps::::new(desc), + ) + .output() + } + + fn int_matmul(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_int_ops!(MatmulOps, B::int_matmul); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = MatmulOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Matmul(desc.clone())), + MatmulOps::::new(desc.into()), + ) + .output() + } + + fn int_mask_where( + tensor: IntTensor, + mask: BoolTensor, + value: IntTensor, + ) -> IntTensor { + #[derive(new, Debug)] + struct MaskWhereOps { + desc: MaskWhereOpIr, + _b: PhantomData, + } + + impl Operation for MaskWhereOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_int_tensor::(&self.desc.tensor); + let value = handles.get_int_tensor::(&self.desc.value); + let mask = handles.get_bool_tensor::(&self.desc.mask); + + let output = B::int_mask_where(tensor, mask, value); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = MaskWhereOpIr::create(tensor.into_ir(), mask.into_ir(), value.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::MaskWhere(desc.clone())), + MaskWhereOps::::new(desc), + ) + .output() + } + + fn int_mask_fill( + tensor: IntTensor, + mask: BoolTensor, + value: Scalar, + ) -> IntTensor { + #[derive(new, Debug)] + struct MaskFillOps { + desc: MaskFillOpIr, + _b: PhantomData, + } + + impl Operation for MaskFillOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_int_tensor::(&self.desc.tensor); + let mask = handles.get_bool_tensor::(&self.desc.mask); + + let output = B::int_mask_fill(tensor, mask, self.desc.value.into()); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let value = value.into(); + let desc = MaskFillOpIr::create(tensor.into_ir(), mask.into_ir(), value, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::MaskFill(desc.clone())), + MaskFillOps::::new(desc), + ) + .output() + } + + fn int_gather( + dim: usize, + tensor: IntTensor, + indices: IntTensor, + ) -> IntTensor { + #[derive(new, Debug)] + struct GatherOps { + desc: GatherOpIr, + _b: PhantomData, + } + + impl Operation for GatherOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_int_tensor::(&self.desc.tensor); + let indices = handles.get_int_tensor::(&self.desc.indices); + + let output = B::int_gather(self.desc.dim, tensor, indices); + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = GatherOpIr::create(tensor.into_ir(), dim, indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::Gather(desc.clone())), + GatherOps::::new(desc), + ) + .output() + } + + fn int_scatter_add( + dim: usize, + tensor: IntTensor, + indices: IntTensor, + value: IntTensor, + ) -> IntTensor { + #[derive(new, Debug)] + struct ScatterOps { + desc: ScatterOpIr, + _b: PhantomData, + } + + impl Operation for ScatterOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_int_tensor::(&self.desc.tensor); + let indices = handles.get_int_tensor::(&self.desc.indices); + let value = handles.get_int_tensor::(&self.desc.value); + + let output = B::int_scatter_add(self.desc.dim, tensor, indices, value); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ScatterOpIr::create( + tensor.into_ir(), + dim, + indices.into_ir(), + value.into_ir(), + IndexingUpdateOp::Add, + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::Scatter(desc.clone())), + ScatterOps::::new(desc), + ) + .output() + } + + fn int_scatter_nd( + data: IntTensor, + indices: IntTensor, + values: IntTensor, + reduction: IndexingUpdateOp, + ) -> IntTensor { + #[derive(new, Debug)] + struct ScatterNdOps { + desc: ScatterNdOpIr, + _b: PhantomData, + } + + impl Operation for ScatterNdOps { + fn execute(&self, handles: &mut HandleContainer) { + let data = handles.get_int_tensor::(&self.desc.data); + let indices = handles.get_int_tensor::(&self.desc.indices); + let values = handles.get_int_tensor::(&self.desc.values); + + let output = B::int_scatter_nd(data, indices, values, self.desc.reduction); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = data.client.clone(); + let desc = ScatterNdOpIr::create( + data.into_ir(), + indices.into_ir(), + values.into_ir(), + reduction, + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::ScatterNd(desc.clone())), + ScatterNdOps::::new(desc), + ) + .output() + } + + fn int_gather_nd(data: IntTensor, indices: IntTensor) -> IntTensor { + #[derive(new, Debug)] + struct GatherNdOps { + desc: GatherNdOpIr, + _b: PhantomData, + } + + impl Operation for GatherNdOps { + fn execute(&self, handles: &mut HandleContainer) { + let data = handles.get_int_tensor::(&self.desc.data); + let indices = handles.get_int_tensor::(&self.desc.indices); + + let output = B::int_gather_nd(data, indices); + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = data.client.clone(); + let desc = GatherNdOpIr::create(data.into_ir(), indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::GatherNd(desc.clone())), + GatherNdOps::::new(desc), + ) + .output() + } + + fn int_select( + tensor: IntTensor, + dim: usize, + indices: IntTensor, + ) -> IntTensor { + #[derive(new, Debug)] + struct SelectOps { + desc: SelectOpIr, + _b: PhantomData, + } + + impl Operation for SelectOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_int_tensor::(&self.desc.tensor); + let indices = handles.get_int_tensor::(&self.desc.indices); + + let output = B::int_select(tensor, self.desc.dim, indices); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = SelectOpIr::create(tensor.into_ir(), dim, indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::Select(desc.clone())), + SelectOps::::new(desc), + ) + .output() + } + + fn int_select_add( + tensor: IntTensor, + dim: usize, + indices: IntTensor, + value: IntTensor, + ) -> IntTensor { + #[derive(new, Debug)] + struct SelectAssignOps { + desc: SelectAssignOpIr, + _b: PhantomData, + } + + impl Operation for SelectAssignOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_int_tensor::(&self.desc.tensor); + let indices = handles.get_int_tensor::(&self.desc.indices); + let value = handles.get_int_tensor::(&self.desc.value); + + let output = B::int_select_add(tensor, self.desc.dim, indices, value); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = SelectAssignOpIr::create( + tensor.into_ir(), + dim, + indices.into_ir(), + value.into_ir(), + IndexingUpdateOp::Add, + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::SelectAssign(desc.clone())), + SelectAssignOps::::new(desc), + ) + .output() + } + + fn int_cat(tensors: Vec>, dim: usize) -> IntTensor { + #[derive(new, Debug)] + struct CatOps { + desc: CatOpIr, + _b: PhantomData, + } + + impl Operation for CatOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensors = self + .desc + .tensors + .iter() + .map(|tensor| handles.get_int_tensor::(tensor)) + .collect(); + + let output = B::int_cat(tensors, self.desc.dim); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensors.first().unwrap().client.clone(); + let tensors = tensors.into_iter().map(|t| t.into_ir()).collect(); + let desc = CatOpIr::create(tensors, dim, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::Cat(desc.clone())), + CatOps::::new(desc), + ) + .output() + } + + fn int_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_int_cmp_ops!(EqualOps, B::int_equal); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::Equal(desc.clone())), + EqualOps::::new(desc), + ) + .output() + } + + fn int_equal_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor { + scalar_int_cmp_ops!(EqualElemOps, B::int_equal_elem); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::EqualElem(desc.clone())), + EqualElemOps::::new(desc), + ) + .output() + } + + fn int_greater( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_int_cmp_ops!(GreaterOps, B::int_greater); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericInt(desc.lhs.dtype, NumericOperationIr::Greater(desc.clone())), + GreaterOps::::new(desc), + ) + .output() + } + + fn int_greater_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + scalar_int_cmp_ops!(GreaterElemOps, B::int_greater_elem); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericInt( + desc.lhs.dtype, + NumericOperationIr::GreaterElem(desc.clone()), + ), + GreaterElemOps::::new(desc), + ) + .output() + } + + fn int_greater_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_int_cmp_ops!(GreaterEqualOps, B::int_greater_equal); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericInt( + desc.lhs.dtype, + NumericOperationIr::GreaterEqual(desc.clone()), + ), + GreaterEqualOps::::new(desc), + ) + .output() + } + + fn int_greater_equal_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + scalar_int_cmp_ops!(GreaterEqualElemOps, B::int_greater_equal_elem); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericInt( + desc.lhs.dtype, + NumericOperationIr::GreaterEqualElem(desc.clone()), + ), + GreaterEqualElemOps::::new(desc), + ) + .output() + } + + fn int_lower( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_int_cmp_ops!(LowerOps, B::int_lower); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericInt(desc.lhs.dtype, NumericOperationIr::Lower(desc.clone())), + LowerOps::::new(desc), + ) + .output() + } + + fn int_lower_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor { + scalar_int_cmp_ops!(LowerElemOps, B::int_lower_elem); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericInt( + desc.lhs.dtype, + NumericOperationIr::LowerElem(desc.clone()), + ), + LowerElemOps::::new(desc), + ) + .output() + } + + fn int_lower_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_int_cmp_ops!(LowerEqualOps, B::int_lower_equal); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericInt( + desc.lhs.dtype, + NumericOperationIr::LowerEqual(desc.clone()), + ), + LowerEqualOps::::new(desc), + ) + .output() + } + + fn int_lower_equal_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + scalar_int_cmp_ops!(LowerEqualElemOps, B::int_lower_equal_elem); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericInt( + desc.lhs.dtype, + NumericOperationIr::LowerEqualElem(desc.clone()), + ), + LowerEqualElemOps::::new(desc), + ) + .output() + } + + fn int_add(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_int_ops!(AddOps, B::int_add); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::Add(desc.clone())), + AddOps::::new(desc), + ) + .output() + } + + fn int_add_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + scalar_int_ops!(AddOps, B::int_add_scalar); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::AddScalar(desc.clone()), + ), + AddOps::::new(desc), + ) + .output() + } + + fn int_sub(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_int_ops!(SubOps, B::int_sub); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::Sub(desc.clone())), + SubOps::::new(desc), + ) + .output() + } + + fn int_sub_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + scalar_int_ops!(SubOps, B::int_sub_scalar); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::SubScalar(desc.clone()), + ), + SubOps::::new(desc), + ) + .output() + } + + fn int_mul(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_int_ops!(MulOps, B::int_mul); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::Mul(desc.clone())), + MulOps::::new(desc), + ) + .output() + } + + fn int_mul_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + scalar_int_ops!(MulOps, B::int_mul_scalar); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::MulScalar(desc.clone()), + ), + MulOps::::new(desc), + ) + .output() + } + + fn int_div(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_int_ops!(DivOps, B::int_div); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::Div(desc.clone())), + DivOps::::new(desc), + ) + .output() + } + + fn int_div_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + scalar_int_ops!(DivOps, B::int_div_scalar); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::DivScalar(desc.clone()), + ), + DivOps::::new(desc), + ) + .output() + } + + fn int_remainder(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_int_ops!(ModOps, B::int_remainder); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::Rem(desc.clone())), + ModOps::::new(desc), + ) + .output() + } + + fn int_remainder_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + scalar_int_ops!(ModOps, B::int_remainder_scalar); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::RemScalar(desc.clone()), + ), + ModOps::::new(desc), + ) + .output() + } + + fn int_zeros(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + #[derive(new, Debug)] + struct ZerosOps { + desc: TensorIr, + device: Device, + } + + impl Operation for ZerosOps { + fn execute(&self, handles: &mut HandleContainer) { + let shape = self.desc.shape.clone(); + let output = B::int_zeros(shape, &self.device, self.desc.dtype.into()); + handles.register_int_tensor::(&self.desc.id, output); + } + } + + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register( + StreamId::current(), + OperationIr::BaseInt(BaseOperationIr::Zeros(desc.clone())), + ZerosOps::::new(desc.out, device.clone()), + ) + .output() + } + + fn int_ones(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + #[derive(new, Debug)] + struct OnesOps { + desc: TensorIr, + device: Device, + } + + impl Operation for OnesOps { + fn execute(&self, handles: &mut HandleContainer) { + let shape = self.desc.shape.clone(); + let output = B::int_ones(shape, &self.device, self.desc.dtype.into()); + handles.register_int_tensor::(&self.desc.id, output); + } + } + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register( + StreamId::current(), + OperationIr::BaseInt(BaseOperationIr::Ones(desc.clone())), + OnesOps::::new(desc.out, device.clone()), + ) + .output() + } + + fn int_full( + shape: Shape, + fill_value: Scalar, + device: &Device, + dtype: IntDType, + ) -> IntTensor { + #[derive(new, Debug)] + struct FullOps { + out: TensorIr, + elem: ScalarIr, + device: Device, + } + + impl Operation for FullOps { + fn execute(&self, handles: &mut HandleContainer) { + let shape = self.out.shape.clone(); + let output = + B::int_full(shape, self.elem.into(), &self.device, self.out.dtype.into()); + handles.register_int_tensor::(&self.out.id, output); + } + } + + let client = get_client::(device); + let dtype = dtype.into(); + let value = fill_value.into(); + let desc = FullOpIr::create(shape, dtype, value, || client.create_empty_handle()); + + client + .register( + StreamId::current(), + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::Full(desc.clone())), + FullOps::::new(desc.out, desc.value, device.clone()), + ) + .output() + } + + fn int_sum(tensor: IntTensor) -> IntTensor { + unary_int_ops!(SumOps, B::int_sum, reduce); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::Sum(desc.clone())), + SumOps::::new(desc.into()), + ) + .output() + } + + fn int_sum_dim(tensor: IntTensor, axis: usize) -> IntTensor { + reduce_int_ops!(SumDimOps, |tensor, axis, _| B::int_sum_dim(tensor, axis)); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = + ReduceDimOpIr::create(tensor.into_ir(), axis, 1, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::SumDim(desc.clone())), + SumDimOps::::new(desc), + ) + .output() + } + + fn int_prod(tensor: IntTensor) -> IntTensor { + unary_int_ops!(ProdOps, B::int_prod, reduce); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::Prod(desc.clone())), + ProdOps::::new(desc.into()), + ) + .output() + } + + fn int_prod_dim(tensor: IntTensor, dim: usize) -> IntTensor { + reduce_int_ops!(ProdDimOps, |tensor, axis, _| B::int_prod_dim(tensor, axis)); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::ProdDim(desc.clone())), + ProdDimOps::::new(desc), + ) + .output() + } + + fn int_mean(tensor: IntTensor) -> IntTensor { + unary_int_ops!(MeanOps, B::int_mean, reduce); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::Mean(desc.clone())), + MeanOps::::new(desc.into()), + ) + .output() + } + + fn int_mean_dim(tensor: IntTensor, dim: usize) -> IntTensor { + reduce_int_ops!(MeanDimOps, |tensor, axis, _| B::int_mean_dim(tensor, axis)); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::MeanDim(desc.clone())), + MeanDimOps::::new(desc), + ) + .output() + } + + fn int_cumsum(tensor: IntTensor, dim: usize) -> IntTensor { + #[derive(new, Debug)] + struct CumsumOps { + desc: DimOpIr, + _b: PhantomData, + } + + impl Operation for CumsumOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_int_tensor::(&self.desc.input); + let output = B::int_cumsum(input, self.desc.axis); + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::CumSum(desc.clone())), + CumsumOps::::new(desc), + ) + .output() + } + + fn int_arange_step( + range: std::ops::Range, + step: usize, + device: &Device, + dtype: IntDType, + ) -> IntTensor { + // Build `arange` from fusable ops (`int_full` of ones → `cumsum` → affine) instead of the + // host-materialized `from_data` default. This keeps it *inside* the fusion graph: for the + // remote backend it means no per-call `Init` op shipping the index data over the network — + // the indices are computed on the device and ride the cached graph like any other op. + // (`int_arange` delegates here, so this covers both.) + // + // Match the default impl's contract: `range.step_by(0)` panics, so a zero step is a misuse + // we surface loudly rather than silently returning an empty tensor. + assert!(step != 0, "arange step must be non-zero"); + if range.end <= range.start { + return Self::int_full(Shape::new([0]), 0i64.into(), device, dtype); + } + + let count = ((range.end - range.start) as usize).div_ceil(step); + // `ones.cumsum` is `[1, 2, ..., count]`; map the 1-based index `k` to `start + (k - 1) * step`. + let ones = Self::int_full(Shape::new([count]), 1i64.into(), device, dtype); + let counting = Self::int_cumsum(ones, 0); + let scaled = Self::int_mul_scalar(counting, (step as i64).into()); + Self::int_add_scalar(scaled, (range.start - step as i64).into()) + } + + fn int_cumprod(tensor: IntTensor, dim: usize) -> IntTensor { + #[derive(new, Debug)] + struct CumprodOps { + desc: DimOpIr, + _b: PhantomData, + } + + impl Operation for CumprodOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_int_tensor::(&self.desc.input); + let output = B::int_cumprod(input, self.desc.axis); + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::CumProd(desc.clone())), + CumprodOps::::new(desc), + ) + .output() + } + + fn int_cummin(tensor: IntTensor, dim: usize) -> IntTensor { + #[derive(new, Debug)] + struct CumminOps { + desc: DimOpIr, + _b: PhantomData, + } + + impl Operation for CumminOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_int_tensor::(&self.desc.input); + let output = B::int_cummin(input, self.desc.axis); + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::CumMin(desc.clone())), + CumminOps::::new(desc), + ) + .output() + } + + fn int_cummax(tensor: IntTensor, dim: usize) -> IntTensor { + #[derive(new, Debug)] + struct CummaxOps { + desc: DimOpIr, + _b: PhantomData, + } + + impl Operation for CummaxOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_int_tensor::(&self.desc.input); + let output = B::int_cummax(input, self.desc.axis); + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::CumMax(desc.clone())), + CummaxOps::::new(desc), + ) + .output() + } + + fn int_argmax(tensor: IntTensor, dim: usize) -> IntTensor { + reduce_int_ops!(ArgMaxOps, |tensor, axis, _| B::int_argmax(tensor, axis)); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::ArgMax(desc.clone())), + ArgMaxOps::::new(desc), + ) + .output() + } + + fn int_argtopk(tensor: IntTensor, dim: usize, k: usize) -> IntTensor { + reduce_int_ops!(ArgTopKOps, B::int_argtopk); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, k, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::ArgTopK(desc.clone())), + ArgTopKOps::::new(desc), + ) + .output() + } + + fn int_topk(tensor: IntTensor, dim: usize, k: usize) -> IntTensor { + reduce_int_ops!(TopKOps, B::int_topk); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, k, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::TopK(desc.clone())), + TopKOps::::new(desc), + ) + .output() + } + + fn int_argmin(tensor: IntTensor, dim: usize) -> IntTensor { + reduce_int_ops!(ArgMinOps, |tensor, axis, _| B::int_argmin(tensor, axis)); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::ArgMin(desc.clone())), + ArgMinOps::::new(desc), + ) + .output() + } + + fn int_clamp(tensor: IntTensor, min: Scalar, max: Scalar) -> IntTensor { + #[derive(new, Debug)] + struct ClampOps { + desc: ClampOpIr, + _b: PhantomData, + } + + impl Operation for ClampOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_int_tensor::(&self.desc.tensor); + let output = B::int_clamp(input, self.desc.min.into(), self.desc.max.into()); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let min = min.into(); + let max = max.into(); + let desc = ClampOpIr::create(tensor.into_ir(), min, max, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::Clamp(desc.clone())), + ClampOps::::new(desc), + ) + .output() + } + + fn int_abs(tensor: IntTensor) -> IntTensor { + unary_int_ops!(AbsOps, B::int_abs); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::Abs(desc.clone())), + AbsOps::::new(desc), + ) + .output() + } + + fn int_into_float(tensor: IntTensor, out_dtype: FloatDType) -> FloatTensor { + #[derive(new, Debug)] + struct IntoFloatOps { + desc: CastOpIr, + _b: PhantomData, + } + + impl Operation for IntoFloatOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_int_tensor::(&self.desc.input); + let output = B::int_into_float(input, self.desc.out.dtype.into()); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = CastOpIr::create(tensor.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Int(IntOperationIr::IntoFloat(desc.clone())), + IntoFloatOps::::new(desc), + ) + .output() + } + + fn int_swap_dims(tensor: IntTensor, dim1: usize, dim2: usize) -> IntTensor { + #[derive(new, Debug)] + struct SwapDimsOps { + desc: SwapDimsOpIr, + _b: PhantomData, + } + + impl Operation for SwapDimsOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_int_tensor::(&self.desc.input); + let output = B::int_swap_dims(input, self.desc.dim1, self.desc.dim2); + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = SwapDimsOpIr::create(tensor.into_ir(), dim1, dim2, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::SwapDims(desc.clone())), + SwapDimsOps::::new(desc), + ) + .output() + } + + fn int_max(tensor: IntTensor) -> IntTensor { + unary_int_ops!(MaxOps, B::int_max, reduce); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::Max(desc.clone())), + MaxOps::::new(desc.into()), + ) + .output() + } + + fn int_max_dim(tensor: IntTensor, dim: usize) -> IntTensor { + reduce_int_ops!(MaxDimOps, |tensor, axis, _| B::int_max_dim(tensor, axis)); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::MaxDim(desc.clone())), + MaxDimOps::::new(desc), + ) + .output() + } + + fn int_max_dim_with_indices( + tensor: IntTensor, + dim: usize, + ) -> (IntTensor, IntTensor) { + #[derive(new, Debug)] + struct MaxDimWithIndicesOps { + desc: ReduceDimWithIndicesOpIr, + _b: PhantomData, + } + + impl Operation for MaxDimWithIndicesOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_int_tensor::(&self.desc.tensor); + let (output, indices) = B::int_max_dim_with_indices(tensor, self.desc.dim); + + handles.register_int_tensor::(&self.desc.out.id, output); + handles.register_int_tensor::(&self.desc.out_indices.id, indices); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let dtype = tensor.dtype; + let desc = ReduceDimWithIndicesOpIr::create(tensor.into_ir(), dim, dtype, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericInt(dtype, NumericOperationIr::MaxDimWithIndices(desc.clone())), + MaxDimWithIndicesOps::::new(desc), + ) + .outputs() + .into() + } + + fn int_min(tensor: IntTensor) -> IntTensor { + unary_int_ops!(MinOps, B::int_min, reduce); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::Min(desc.clone())), + MinOps::::new(desc.into()), + ) + .output() + } + + fn int_max_abs(tensor: IntTensor) -> IntTensor { + unary_int_ops!(MaxAbsOps, B::int_max_abs, reduce); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::MaxAbs(desc.clone())), + MaxAbsOps::::new(desc.into()), + ) + .output() + } + + fn int_max_abs_dim(tensor: IntTensor, dim: usize) -> IntTensor { + reduce_int_ops!(MaxAbsDimOps, |tensor, axis, _| B::int_max_abs_dim( + tensor, axis + )); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::MaxAbsDim(desc.clone()), + ), + MaxAbsDimOps::::new(desc), + ) + .output() + } + + fn int_min_dim(tensor: IntTensor, dim: usize) -> IntTensor { + reduce_int_ops!(MinDimOps, |tensor, axis, _| B::int_min_dim(tensor, axis)); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::MinDim(desc.clone())), + MinDimOps::::new(desc), + ) + .output() + } + + fn int_min_dim_with_indices( + tensor: IntTensor, + dim: usize, + ) -> (IntTensor, IntTensor) { + #[derive(new, Debug)] + struct MinDimWithIndicesOps { + desc: ReduceDimWithIndicesOpIr, + _b: PhantomData, + } + + impl Operation for MinDimWithIndicesOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_int_tensor::(&self.desc.tensor); + let (output, indices) = B::int_min_dim_with_indices(tensor, self.desc.dim); + + handles.register_int_tensor::(&self.desc.out.id, output); + handles.register_int_tensor::(&self.desc.out_indices.id, indices); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let dtype = tensor.dtype; + let desc = ReduceDimWithIndicesOpIr::create(tensor.into_ir(), dim, dtype, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericInt(dtype, NumericOperationIr::MinDimWithIndices(desc.clone())), + MinDimWithIndicesOps::::new(desc), + ) + .outputs() + .into() + } + + fn int_random( + shape: Shape, + distribution: Distribution, + device: &Device, + dtype: IntDType, + ) -> IntTensor { + #[derive(new, Debug)] + struct IntRandomOps { + desc: RandomOpIr, + device: Device, + } + + impl Operation for IntRandomOps { + fn execute(&self, handles: &mut HandleContainer) { + let shape = self.desc.out.shape.clone(); + let output = B::int_random( + shape, + self.desc.distribution, + &self.device, + self.desc.out.dtype.into(), + ); + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let dtype = dtype.into(); + let client = get_client::(device); + let desc = RandomOpIr::create(shape, dtype, distribution, || client.create_empty_handle()); + + client + .register( + StreamId::current(), + OperationIr::NumericInt(dtype, NumericOperationIr::IntRandom(desc.clone())), + IntRandomOps::::new(desc, device.clone()), + ) + .output() + } + + fn int_permute(tensor: IntTensor, axes: &[usize]) -> IntTensor { + #[derive(new, Debug)] + struct PermuteDimsOps { + desc: PermuteOpIr, + _b: PhantomData, + } + + impl Operation for PermuteDimsOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_int_tensor::(&self.desc.input); + let output = B::int_permute(input, self.desc.axes.as_slice()); + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = PermuteOpIr::create(tensor.into_ir(), axes.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::Permute(desc.clone())), + PermuteDimsOps::::new(desc), + ) + .output() + } + + fn int_expand(tensor: IntTensor, shape: Shape) -> IntTensor { + #[derive(new, Debug)] + struct ExpandOps { + desc: ShapeOpIr, + _b: PhantomData, + } + + impl Operation for ExpandOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_int_tensor::(&self.desc.input); + let output = B::int_expand(input, self.desc.out.shape.clone()); + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ShapeOpIr::expand(tensor.into_ir(), shape, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::Expand(desc.clone())), + ExpandOps::::new(desc), + ) + .output() + } + + fn int_flip(tensor: IntTensor, axes: &[usize]) -> IntTensor { + #[derive(new, Debug)] + struct FlipDimsOps { + desc: FlipOpIr, + _b: PhantomData, + } + + impl Operation for FlipDimsOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_int_tensor::(&self.desc.input); + let axes = &self.desc.axes; + let output = B::int_flip(input, axes); + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = FlipOpIr::create(tensor.into_ir(), axes.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::Flip(desc.clone())), + FlipDimsOps::::new(desc), + ) + .output() + } + + fn int_repeat_dim(tensor: IntTensor, dim: usize, times: usize) -> IntTensor { + #[derive(new, Debug)] + struct RepeatDimOps { + desc: RepeatDimOpIr, + _b: PhantomData, + } + + impl Operation for RepeatDimOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_int_tensor::(&self.desc.tensor); + + let output = B::int_repeat_dim(tensor, self.desc.dim, self.desc.times); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = RepeatDimOpIr::create(tensor.into_ir(), dim, times, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::RepeatDim(desc.clone())), + RepeatDimOps::::new(desc), + ) + .output() + } + + fn bitwise_and(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_int_ops!(BitwiseAndOps, B::bitwise_and); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Int(IntOperationIr::BitwiseAnd(desc.clone())), + BitwiseAndOps::::new(desc), + ) + .output() + } + + fn bitwise_and_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + scalar_int_ops!(BitwiseAndOps, B::bitwise_and_scalar); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Int(IntOperationIr::BitwiseAndScalar(desc.clone())), + BitwiseAndOps::::new(desc), + ) + .output() + } + + fn bitwise_or(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_int_ops!(BitwiseOrOps, B::bitwise_or); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Int(IntOperationIr::BitwiseOr(desc.clone())), + BitwiseOrOps::::new(desc), + ) + .output() + } + + fn bitwise_or_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + scalar_int_ops!(BitwiseOrOps, B::bitwise_or_scalar); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Int(IntOperationIr::BitwiseOrScalar(desc.clone())), + BitwiseOrOps::::new(desc), + ) + .output() + } + + fn bitwise_xor(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_int_ops!(BitwiseXorOps, B::bitwise_xor); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Int(IntOperationIr::BitwiseXor(desc.clone())), + BitwiseXorOps::::new(desc), + ) + .output() + } + + fn bitwise_xor_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + scalar_int_ops!(BitwiseXorOps, B::bitwise_xor_scalar); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Int(IntOperationIr::BitwiseXorScalar(desc.clone())), + BitwiseXorOps::::new(desc), + ) + .output() + } + + fn bitwise_not(tensor: IntTensor) -> IntTensor { + unary_int_ops!(BitwiseNotOps, B::bitwise_not); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Int(IntOperationIr::BitwiseNot(desc.clone())), + BitwiseNotOps::::new(desc), + ) + .output() + } + + fn bitwise_left_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_int_ops!(BitwiseLeftShiftOps, B::bitwise_left_shift); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Int(IntOperationIr::BitwiseLeftShift(desc.clone())), + BitwiseLeftShiftOps::::new(desc), + ) + .output() + } + + fn bitwise_left_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + scalar_int_ops!(BitwiseLeftShiftOps, B::bitwise_left_shift_scalar); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Int(IntOperationIr::BitwiseLeftShiftScalar(desc.clone())), + BitwiseLeftShiftOps::::new(desc), + ) + .output() + } + + fn bitwise_right_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_int_ops!(BitwiseRightShiftOps, B::bitwise_right_shift); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Int(IntOperationIr::BitwiseRightShift(desc.clone())), + BitwiseRightShiftOps::::new(desc), + ) + .output() + } + + fn bitwise_right_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + scalar_int_ops!(BitwiseRightShiftOps, B::bitwise_right_shift_scalar); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Int(IntOperationIr::BitwiseRightShiftScalar(desc.clone())), + BitwiseRightShiftOps::::new(desc), + ) + .output() + } + + fn int_cast(tensor: IntTensor, dtype: burn_backend::IntDType) -> IntTensor { + #[derive(new, Debug)] + struct CastOps { + desc: CastOpIr, + dtype: burn_backend::IntDType, + _b: PhantomData, + } + + impl Operation for CastOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_int_tensor::(&self.desc.input); + let output: B::IntTensorPrimitive = B::int_cast(input, self.dtype); + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = CastOpIr::create(tensor.into_ir(), dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::Cast(desc.clone())), + CastOps::::new(desc, dtype), + ) + .output() + } + + fn int_unfold( + tensor: IntTensor, + dim: usize, + size: usize, + step: usize, + ) -> IntTensor { + #[derive(new, Debug)] + struct UnfoldOps { + desc: UnfoldOpIr, + _b: PhantomData, + } + + impl Operation for UnfoldOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_int_tensor::(&self.desc.input); + let output = B::int_unfold(input, self.desc.dim, self.desc.size, self.desc.step); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnfoldOpIr::create(tensor.into_ir(), dim, size, step, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::Unfold(desc.clone())), + UnfoldOps::::new(desc), + ) + .output() + } + + fn int_powi(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + binary_int_ops!(PowOps, B::int_powi); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericInt(desc.out.dtype, NumericOperationIr::Powi(desc.clone())), + PowOps::::new(desc), + ) + .output() + } + + fn int_powi_scalar_impl(lhs: IntTensor, rhs: Scalar) -> IntTensor { + scalar_int_ops!(PowiOps, B::int_powi_scalar); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::PowiScalar(desc.clone()), + ), + PowiOps::::new(desc), + ) + .output() + } +} diff --git a/crates/burn-fusion/src/ops/mod.rs b/crates/burn-fusion/src/ops/mod.rs new file mode 100644 index 0000000..327815d --- /dev/null +++ b/crates/burn-fusion/src/ops/mod.rs @@ -0,0 +1,13 @@ +mod activation; +mod binary; +mod bool_tensor; +mod distributed; +mod int_tensor; +mod module; +mod qtensor; +mod tensor; +mod transaction; +mod unary; + +mod base; +pub use base::NoOp; diff --git a/crates/burn-fusion/src/ops/module.rs b/crates/burn-fusion/src/ops/module.rs new file mode 100644 index 0000000..feef822 --- /dev/null +++ b/crates/burn-fusion/src/ops/module.rs @@ -0,0 +1,1709 @@ +use crate::{ + Fusion, FusionBackend, + stream::{StreamId, execution::Operation}, +}; +use burn_backend::{ + ops::{ + ConvOptions, ConvTransposeOptions, DeformConv2dBackward, DeformConvOptions, + InterpolateOptions, MaxPool1dBackward, MaxPool1dWithIndices, MaxPool2dBackward, + MaxPool2dWithIndices, ModuleOps, + }, + tensor::{FloatTensor, IntTensor}, +}; +use burn_ir::*; +use burn_std::IntDType; +use std::marker::PhantomData; + +macro_rules! make_ops { + ($name:ident, $desc:ty, $fn:expr) => { + #[derive(new, Debug)] + struct $name { + desc: $desc, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + #[allow(clippy::redundant_closure_call)] + $fn(&self.desc, handles) + } + } + }; +} + +impl ModuleOps> for Fusion { + // linear and its backward ops fall back to the default ModuleOps impl, + // which decomposes into matmul + add / matmul + sum. This preserves + // downstream fusion in burn-cubecl-fusion, which matches on those + // primitive IR nodes. + fn conv1d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<1>, + ) -> FloatTensor { + make_ops!(Conv1dOps, Conv1dOpIr, |desc: &Conv1dOpIr, + handles: &mut HandleContainer< + B::Handle, + >| { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let bias = desc + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + let output = B::conv1d(x, weight, bias, desc.options.clone().into()); + handles.register_float_tensor::(&desc.out.id, output); + }); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = Conv1dOpIr::create( + x.into_ir(), + weight.into_ir(), + bias.map(|bias| bias.into_ir()), + options.into(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::Conv1d(desc.clone())), + Conv1dOps::::new(desc), + ) + .output() + } + + fn conv1d_x_backward( + x: FloatTensor>, + weight: FloatTensor>, + output_grad: FloatTensor>, + options: ConvOptions<1>, + ) -> FloatTensor> { + make_ops!( + Conv1dXBackwardOps, + Conv1dXBackwardOpIr, + |desc: &Conv1dXBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + let output = + B::conv1d_x_backward(x, weight, output_grad, desc.options.clone().into()); + handles.register_float_tensor::(&desc.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = Conv1dXBackwardOpIr::create( + x.into_ir(), + weight.into_ir(), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::Conv1dXBackward(desc.clone())), + Conv1dXBackwardOps::::new(desc), + ) + .output() + } + + fn conv1d_weight_backward( + x: FloatTensor>, + weight: FloatTensor>, + output_grad: FloatTensor>, + options: ConvOptions<1>, + ) -> FloatTensor> { + make_ops!( + Conv1dWeightBackwardOps, + Conv1dWeightBackwardOpIr, + |desc: &Conv1dWeightBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + let output = + B::conv1d_weight_backward(x, weight, output_grad, desc.options.clone().into()); + handles.register_float_tensor::(&desc.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = Conv1dWeightBackwardOpIr::create( + x.into_ir(), + weight.into_ir(), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::Conv1dWeightBackward(desc.clone())), + Conv1dWeightBackwardOps::::new(desc), + ) + .output() + } + + fn conv1d_bias_backward( + x: FloatTensor>, + bias: FloatTensor>, + output_grad: FloatTensor>, + ) -> FloatTensor> { + make_ops!( + Conv1dBiasBackwardOps, + Conv1dBiasBackwardOpIr, + |desc: &Conv1dBiasBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&desc.x); + let bias = handles.get_float_tensor::(&desc.bias); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + let output = B::conv1d_bias_backward(x, bias, output_grad); + handles.register_float_tensor::(&desc.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = Conv1dBiasBackwardOpIr::create( + x.into_ir(), + bias.into_ir(), + output_grad.into_ir(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::Conv1dBiasBackward(desc.clone())), + Conv1dBiasBackwardOps::::new(desc), + ) + .output() + } + + fn conv2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<2>, + ) -> FloatTensor { + make_ops!(Conv2dOps, Conv2dOpIr, |args: &Conv2dOpIr, + handles: &mut HandleContainer< + B::Handle, + >| { + let x = handles.get_float_tensor::(&args.x); + let weight = handles.get_float_tensor::(&args.weight); + let bias = args + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + + let output = B::conv2d(x, weight, bias, args.options.clone().into()); + + handles.register_float_tensor::(&args.out.id, output); + }); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = Conv2dOpIr::create( + x.into_ir(), + weight.into_ir(), + bias.map(|bias| bias.into_ir()), + options.into(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::Conv2d(desc.clone())), + Conv2dOps::::new(desc), + ) + .output() + } + + fn conv2d_x_backward( + x: FloatTensor>, + weight: FloatTensor>, + output_grad: FloatTensor>, + options: ConvOptions<2>, + ) -> FloatTensor> { + make_ops!( + Conv2dXBackwardOps, + Conv2dXBackwardOpIr, + |desc: &Conv2dXBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + let output = + B::conv2d_x_backward(x, weight, output_grad, desc.options.clone().into()); + handles.register_float_tensor::(&desc.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = Conv2dXBackwardOpIr::create( + x.into_ir(), + weight.into_ir(), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::Conv2dXBackward(desc.clone())), + Conv2dXBackwardOps::::new(desc), + ) + .output() + } + + fn conv2d_weight_backward( + x: FloatTensor>, + weight: FloatTensor>, + output_grad: FloatTensor>, + options: ConvOptions<2>, + ) -> FloatTensor> { + make_ops!( + Conv2dWeightBackwardOps, + Conv2dWeightBackwardOpIr, + |desc: &Conv2dWeightBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + let output = + B::conv2d_weight_backward(x, weight, output_grad, desc.options.clone().into()); + handles.register_float_tensor::(&desc.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = Conv2dWeightBackwardOpIr::create( + x.into_ir(), + weight.into_ir(), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::Conv2dWeightBackward(desc.clone())), + Conv2dWeightBackwardOps::::new(desc), + ) + .output() + } + + fn conv2d_bias_backward( + x: FloatTensor>, + bias: FloatTensor>, + output_grad: FloatTensor>, + ) -> FloatTensor> { + make_ops!( + Conv2dBiasBackwardOps, + Conv2dBiasBackwardOpIr, + |desc: &Conv2dBiasBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&desc.x); + let bias = handles.get_float_tensor::(&desc.bias); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + let output = B::conv2d_bias_backward(x, bias, output_grad); + handles.register_float_tensor::(&desc.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = Conv2dBiasBackwardOpIr::create( + x.into_ir(), + bias.into_ir(), + output_grad.into_ir(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::Conv2dBiasBackward(desc.clone())), + Conv2dBiasBackwardOps::::new(desc), + ) + .output() + } + + fn deform_conv2d( + x: FloatTensor, + offset: FloatTensor, + weight: FloatTensor, + mask: Option>, + bias: Option>, + options: DeformConvOptions<2>, + ) -> FloatTensor { + make_ops!( + DeformConv2dOps, + DeformConv2dOpIr, + |args: &DeformConv2dOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let offset = handles.get_float_tensor::(&args.offset); + let weight = handles.get_float_tensor::(&args.weight); + let mask = args + .mask + .as_ref() + .map(|mask| handles.get_float_tensor::(mask)); + let bias = args + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + + let output = + B::deform_conv2d(x, offset, weight, mask, bias, args.options.clone().into()); + + handles.register_float_tensor::(&args.out.id, output); + } + ); + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = DeformConv2dOpIr::create( + x.into_ir(), + offset.into_ir(), + weight.into_ir(), + mask.map(|mask| mask.into_ir()), + bias.map(|bias| bias.into_ir()), + options.into(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::DeformableConv2d(Box::new(desc.clone()))), + DeformConv2dOps::::new(desc), + ) + .output() + } + + fn deform_conv2d_backward( + x: FloatTensor, + offset: FloatTensor, + weight: FloatTensor, + mask: Option>, + bias: Option>, + output_grad: FloatTensor, + options: DeformConvOptions<2>, + ) -> DeformConv2dBackward { + make_ops!( + DeformConv2dBackwardOps, + DeformConv2dBackwardOpIr, + |args: &DeformConv2dBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let offset = handles.get_float_tensor::(&args.offset); + let weight = handles.get_float_tensor::(&args.weight); + let mask = args + .mask + .as_ref() + .map(|mask| handles.get_float_tensor::(mask)); + let bias = args + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + let output_grad = handles.get_float_tensor::(&args.out_grad); + + let output = B::deform_conv2d_backward( + x, + offset, + weight, + mask, + bias, + output_grad, + args.options.clone().into(), + ); + + handles.register_float_tensor::(&args.input_grad.id, output.x_grad); + handles.register_float_tensor::(&args.offset_grad.id, output.offset_grad); + handles.register_float_tensor::(&args.weight_grad.id, output.weight_grad); + if let Some((mask_grad, field)) = output.mask_grad.zip(args.mask_grad.as_ref()) { + handles.register_float_tensor::(&field.id, mask_grad); + } + if let Some((bias_grad, field)) = output.bias_grad.zip(args.bias_grad.as_ref()) { + handles.register_float_tensor::(&field.id, bias_grad); + } + } + ); + + let has_bias = bias.is_some(); + let has_mask = mask.is_some(); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = DeformConv2dBackwardOpIr::create( + x.into_ir(), + offset.into_ir(), + weight.into_ir(), + mask.map(|mask| mask.into_ir()), + bias.map(|bias| bias.into_ir()), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + + let mut outputs = client + .register( + streams, + OperationIr::Module(ModuleOperationIr::DeformableConv2dBackward(Box::new( + desc.clone(), + ))), + DeformConv2dBackwardOps::::new(desc), + ) + .into_iter(); + + // When the number of outputs is variable, the order is important + let input_grad = outputs.next().unwrap(); + let offset_grad = outputs.next().unwrap(); + let weight_grad = outputs.next().unwrap(); + let mask_grad = has_mask.then(|| outputs.next().unwrap()); + let bias_grad = has_bias.then(|| outputs.next().unwrap()); + + DeformConv2dBackward::new(input_grad, offset_grad, weight_grad, mask_grad, bias_grad) + } + + fn conv3d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<3>, + ) -> FloatTensor { + make_ops!(Conv3dOps, Conv3dOpIr, |args: &Conv3dOpIr, + handles: &mut HandleContainer< + B::Handle, + >| { + let x = handles.get_float_tensor::(&args.x); + let weight = handles.get_float_tensor::(&args.weight); + let bias = args + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + + let output = B::conv3d(x, weight, bias, args.options.clone().into()); + + handles.register_float_tensor::(&args.out.id, output); + }); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = Conv3dOpIr::create( + x.into_ir(), + weight.into_ir(), + bias.map(|bias| bias.into_ir()), + options.into(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::Conv3d(desc.clone())), + Conv3dOps::::new(desc), + ) + .output() + } + + fn conv3d_x_backward( + x: FloatTensor>, + weight: FloatTensor>, + output_grad: FloatTensor>, + options: ConvOptions<3>, + ) -> FloatTensor> { + make_ops!( + Conv3dXBackwardOps, + Conv3dXBackwardOpIr, + |desc: &Conv3dXBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + let output = + B::conv3d_x_backward(x, weight, output_grad, desc.options.clone().into()); + handles.register_float_tensor::(&desc.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = Conv3dXBackwardOpIr::create( + x.into_ir(), + weight.into_ir(), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::Conv3dXBackward(desc.clone())), + Conv3dXBackwardOps::::new(desc), + ) + .output() + } + + fn conv3d_weight_backward( + x: FloatTensor>, + weight: FloatTensor>, + output_grad: FloatTensor>, + options: ConvOptions<3>, + ) -> FloatTensor> { + make_ops!( + Conv3dWeightBackwardOps, + Conv3dWeightBackwardOpIr, + |desc: &Conv3dWeightBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + let output = + B::conv3d_weight_backward(x, weight, output_grad, desc.options.clone().into()); + handles.register_float_tensor::(&desc.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = Conv3dWeightBackwardOpIr::create( + x.into_ir(), + weight.into_ir(), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::Conv3dWeightBackward(desc.clone())), + Conv3dWeightBackwardOps::::new(desc), + ) + .output() + } + + fn conv3d_bias_backward( + x: FloatTensor>, + bias: FloatTensor>, + output_grad: FloatTensor>, + ) -> FloatTensor> { + make_ops!( + Conv3dBiasBackwardOps, + Conv3dBiasBackwardOpIr, + |desc: &Conv3dBiasBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&desc.x); + let bias = handles.get_float_tensor::(&desc.bias); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + let output = B::conv3d_bias_backward(x, bias, output_grad); + handles.register_float_tensor::(&desc.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = Conv3dBiasBackwardOpIr::create( + x.into_ir(), + bias.into_ir(), + output_grad.into_ir(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::Conv3dBiasBackward(desc.clone())), + Conv3dBiasBackwardOps::::new(desc), + ) + .output() + } + + fn conv_transpose1d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<1>, + ) -> FloatTensor { + make_ops!( + ConvTranspose1dOps, + ConvTranspose1dOpIr, + |args: &ConvTranspose1dOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let weight = handles.get_float_tensor::(&args.weight); + let bias = args + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + + let output = B::conv_transpose1d(x, weight, bias, args.options.clone().into()); + + handles.register_float_tensor::(&args.out.id, output); + } + ); + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = ConvTranspose1dOpIr::create( + x.into_ir(), + weight.into_ir(), + bias.map(|bias| bias.into_ir()), + options.into(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::ConvTranspose1d(desc.clone())), + ConvTranspose1dOps::::new(desc), + ) + .output() + } + + fn conv_transpose2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<2>, + ) -> FloatTensor { + make_ops!( + ConvTranspose2dOps, + ConvTranspose2dOpIr, + |args: &ConvTranspose2dOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let weight = handles.get_float_tensor::(&args.weight); + let bias = args + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + + let output = B::conv_transpose2d(x, weight, bias, args.options.clone().into()); + + handles.register_float_tensor::(&args.out.id, output); + } + ); + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = ConvTranspose2dOpIr::create( + x.into_ir(), + weight.into_ir(), + bias.map(|bias| bias.into_ir()), + options.into(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::ConvTranspose2d(desc.clone())), + ConvTranspose2dOps::::new(desc), + ) + .output() + } + + fn conv_transpose3d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<3>, + ) -> FloatTensor { + make_ops!( + ConvTranspose3dOps, + ConvTranspose3dOpIr, + |args: &ConvTranspose3dOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let weight = handles.get_float_tensor::(&args.weight); + let bias = args + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + + let output = B::conv_transpose3d(x, weight, bias, args.options.clone().into()); + + handles.register_float_tensor::(&args.out.id, output); + } + ); + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = ConvTranspose3dOpIr::create( + x.into_ir(), + weight.into_ir(), + bias.map(|bias| bias.into_ir()), + options.into(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::ConvTranspose3d(desc.clone())), + ConvTranspose3dOps::::new(desc), + ) + .output() + } + + fn avg_pool1d( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + make_ops!( + AvgPool1dOps, + AvgPool1dOpIr, + |args: &AvgPool1dOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let output = B::avg_pool1d( + x, + args.kernel_size, + args.stride, + args.padding, + args.count_include_pad, + args.ceil_mode, + ); + + handles.register_float_tensor::(&args.out.id, output); + } + ); + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = AvgPool1dOpIr::create( + x.into_ir(), + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::AvgPool1d(desc.clone())), + AvgPool1dOps::::new(desc), + ) + .output() + } + + fn avg_pool2d( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + make_ops!( + AvgPool2dOps, + AvgPool2dOpIr, + |args: &AvgPool2dOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let output = B::avg_pool2d( + x, + args.kernel_size, + args.stride, + args.padding, + args.count_include_pad, + args.ceil_mode, + ); + + handles.register_float_tensor::(&args.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = AvgPool2dOpIr::create( + x.into_ir(), + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::AvgPool2d(desc.clone())), + AvgPool2dOps::::new(desc), + ) + .output() + } + + fn avg_pool1d_backward( + x: FloatTensor, + grad: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + make_ops!( + AvgPool1dBackwardOps, + AvgPool1dBackwardOpIr, + |args: &AvgPool1dBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let grad = handles.get_float_tensor::(&args.grad); + let output = B::avg_pool1d_backward( + x, + grad, + args.kernel_size, + args.stride, + args.padding, + args.count_include_pad, + args.ceil_mode, + ); + + handles.register_float_tensor::(&args.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = AvgPool1dBackwardOpIr::create( + x.into_ir(), + grad.into_ir(), + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::AvgPool1dBackward(desc.clone())), + AvgPool1dBackwardOps::::new(desc), + ) + .output() + } + + fn avg_pool2d_backward( + x: FloatTensor, + grad: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + make_ops!( + AvgPool2dBackwardOps, + AvgPool2dBackwardOpIr, + |args: &AvgPool2dBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let grad = handles.get_float_tensor::(&args.grad); + let output = B::avg_pool2d_backward( + x, + grad, + args.kernel_size, + args.stride, + args.padding, + args.count_include_pad, + args.ceil_mode, + ); + + handles.register_float_tensor::(&args.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = AvgPool2dBackwardOpIr::create( + x.into_ir(), + grad.into_ir(), + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::AvgPool2dBackward(desc.clone())), + AvgPool2dBackwardOps::::new(desc), + ) + .output() + } + + fn max_pool1d( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + ) -> FloatTensor { + make_ops!( + MaxPool1dOps, + MaxPool1dOpIr, + |args: &MaxPool1dOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let output = B::max_pool1d( + x, + args.kernel_size, + args.stride, + args.padding, + args.dilation, + args.ceil_mode, + ); + + handles.register_float_tensor::(&args.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = MaxPool1dOpIr::create( + x.into_ir(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::MaxPool1d(desc.clone())), + MaxPool1dOps::::new(desc), + ) + .output() + } + + fn max_pool2d( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + ) -> FloatTensor { + make_ops!( + MaxPool2dOps, + MaxPool2dOpIr, + |args: &MaxPool2dOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let output = B::max_pool2d( + x, + args.kernel_size, + args.stride, + args.padding, + args.dilation, + args.ceil_mode, + ); + + handles.register_float_tensor::(&args.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = MaxPool2dOpIr::create( + x.into_ir(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::MaxPool2d(desc.clone())), + MaxPool2dOps::::new(desc), + ) + .output() + } + + fn max_pool1d_with_indices( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + indices_dtype: IntDType, + ) -> MaxPool1dWithIndices { + make_ops!( + MaxPool1dWithIndicesOps, + MaxPool1dWithIndicesOpIr, + |args: &MaxPool1dWithIndicesOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let output = B::max_pool1d_with_indices( + x, + args.kernel_size, + args.stride, + args.padding, + args.dilation, + args.ceil_mode, + args.out_indices.dtype.into(), + ); + + handles.register_float_tensor::(&args.out.id, output.output); + handles.register_int_tensor::(&args.out_indices.id, output.indices); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = MaxPool1dWithIndicesOpIr::create( + x.into_ir(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + indices_dtype.into(), + || client.create_empty_handle(), + ); + + let [out, out_indices] = client + .register( + streams, + OperationIr::Module(ModuleOperationIr::MaxPool1dWithIndices(desc.clone())), + MaxPool1dWithIndicesOps::::new(desc), + ) + .outputs(); + + MaxPool1dWithIndices::new(out, out_indices) + } + + fn max_pool2d_with_indices( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + indices_dtype: IntDType, + ) -> MaxPool2dWithIndices { + make_ops!( + MaxPool2dWithIndicesOps, + MaxPool2dWithIndicesOpIr, + |args: &MaxPool2dWithIndicesOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let output = B::max_pool2d_with_indices( + x, + args.kernel_size, + args.stride, + args.padding, + args.dilation, + args.ceil_mode, + args.out_indices.dtype.into(), + ); + + handles.register_float_tensor::(&args.out.id, output.output); + handles.register_int_tensor::(&args.out_indices.id, output.indices); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = MaxPool2dWithIndicesOpIr::create( + x.into_ir(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + indices_dtype.into(), + || client.create_empty_handle(), + ); + + let [out, out_indices] = client + .register( + streams, + OperationIr::Module(ModuleOperationIr::MaxPool2dWithIndices(desc.clone())), + MaxPool2dWithIndicesOps::::new(desc), + ) + .outputs(); + + MaxPool2dWithIndices::new(out, out_indices) + } + + fn max_pool1d_with_indices_backward( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + output_grad: FloatTensor, + indices: IntTensor, + ) -> MaxPool1dBackward { + make_ops!( + MaxPool1dWithIndicesBackwardOps, + MaxPool1dWithIndicesBackwardOpIr, + |args: &MaxPool1dWithIndicesBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let grad = handles.get_float_tensor::(&args.grad); + let indices = handles.get_int_tensor::(&args.indices); + let output = B::max_pool1d_with_indices_backward( + x, + args.kernel_size, + args.stride, + args.padding, + args.dilation, + args.ceil_mode, + grad, + indices, + ); + + handles.register_float_tensor::(&args.out.id, output.x_grad); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = MaxPool1dWithIndicesBackwardOpIr::create( + x.into_ir(), + output_grad.into_ir(), + indices.into_ir(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + || client.create_empty_handle(), + ); + + let out = client + .register( + streams, + OperationIr::Module(ModuleOperationIr::MaxPool1dWithIndicesBackward( + desc.clone(), + )), + MaxPool1dWithIndicesBackwardOps::::new(desc), + ) + .output(); + + MaxPool1dBackward::new(out) + } + + fn max_pool2d_with_indices_backward( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + output_grad: FloatTensor, + indices: IntTensor, + ) -> MaxPool2dBackward { + make_ops!( + MaxPool2dWithIndicesBackwardOps, + MaxPool2dWithIndicesBackwardOpIr, + |args: &MaxPool2dWithIndicesBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let grad = handles.get_float_tensor::(&args.grad); + let indices = handles.get_int_tensor::(&args.indices); + let output = B::max_pool2d_with_indices_backward( + x, + args.kernel_size, + args.stride, + args.padding, + args.dilation, + args.ceil_mode, + grad, + indices, + ); + + handles.register_float_tensor::(&args.out.id, output.x_grad); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = MaxPool2dWithIndicesBackwardOpIr::create( + x.into_ir(), + output_grad.into_ir(), + indices.into_ir(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + || client.create_empty_handle(), + ); + + let out = client + .register( + streams, + OperationIr::Module(ModuleOperationIr::MaxPool2dWithIndicesBackward( + desc.clone(), + )), + MaxPool2dWithIndicesBackwardOps::::new(desc), + ) + .output(); + + MaxPool2dBackward::new(out) + } + + fn adaptive_avg_pool1d(x: FloatTensor, output_size: usize) -> FloatTensor { + make_ops!( + AdaptiveAvgPool1dOps, + AdaptiveAvgPool1dOpIr, + |args: &AdaptiveAvgPool1dOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let output = B::adaptive_avg_pool1d(x, args.output_size); + + handles.register_float_tensor::(&args.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = AdaptiveAvgPool1dOpIr::create(x.into_ir(), output_size, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::AdaptiveAvgPool1d(desc.clone())), + AdaptiveAvgPool1dOps::::new(desc), + ) + .output() + } + + fn adaptive_avg_pool2d(x: FloatTensor, output_size: [usize; 2]) -> FloatTensor { + make_ops!( + AdaptiveAvgPool2dOps, + AdaptiveAvgPool2dOpIr, + |args: &AdaptiveAvgPool2dOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let output = B::adaptive_avg_pool2d(x, args.output_size); + + handles.register_float_tensor::(&args.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = AdaptiveAvgPool2dOpIr::create(x.into_ir(), output_size, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::AdaptiveAvgPool2d(desc.clone())), + AdaptiveAvgPool2dOps::::new(desc), + ) + .output() + } + + fn adaptive_avg_pool1d_backward( + x: FloatTensor, + grad: FloatTensor, + ) -> FloatTensor { + make_ops!( + AdaptiveAvgPool1dBackwardOps, + AdaptiveAvgPool1dBackwardOpIr, + |args: &AdaptiveAvgPool1dBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let grad = handles.get_float_tensor::(&args.grad); + let output = B::adaptive_avg_pool1d_backward(x, grad); + + handles.register_float_tensor::(&args.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = AdaptiveAvgPool1dBackwardOpIr::create(x.into_ir(), grad.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::AdaptiveAvgPool1dBackward(desc.clone())), + AdaptiveAvgPool1dBackwardOps::::new(desc), + ) + .output() + } + + fn adaptive_avg_pool2d_backward( + x: FloatTensor, + grad: FloatTensor, + ) -> FloatTensor { + make_ops!( + AdaptiveAvgPool2dBackwardOps, + AdaptiveAvgPool2dBackwardOpIr, + |args: &AdaptiveAvgPool2dBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let grad = handles.get_float_tensor::(&args.grad); + let output = B::adaptive_avg_pool2d_backward(x, grad); + + handles.register_float_tensor::(&args.out.id, output); + } + ); + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = AdaptiveAvgPool2dBackwardOpIr::create(x.into_ir(), grad.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::AdaptiveAvgPool2dBackward(desc.clone())), + AdaptiveAvgPool2dBackwardOps::::new(desc), + ) + .output() + } + + fn interpolate( + x: FloatTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> FloatTensor { + make_ops!( + InterpolateOps, + InterpolateOpIr, + |args: &InterpolateOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let output = B::interpolate(x, args.output_size, args.options.clone().into()); + handles.register_float_tensor::(&args.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = InterpolateOpIr::create(x.into_ir(), output_size, options.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::Interpolate(desc.clone())), + InterpolateOps::::new(desc), + ) + .output() + } + + fn interpolate_backward( + x: FloatTensor, + grad: FloatTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> FloatTensor { + make_ops!( + InterpolateBackwardOps, + InterpolateBackwardOpIr, + |args: &InterpolateBackwardOpIr, handles: &mut HandleContainer| { + let x = handles.get_float_tensor::(&args.x); + let grad = handles.get_float_tensor::(&args.grad); + let output = + B::interpolate_backward(x, grad, args.output_size, args.options.clone().into()); + + handles.register_float_tensor::(&args.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = x.client.clone(); + let desc = InterpolateBackwardOpIr::create( + x.into_ir(), + grad.into_ir(), + output_size, + options.into(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::InterpolateBackward(desc.clone())), + InterpolateBackwardOps::::new(desc), + ) + .output() + } + + fn attention( + query: FloatTensor>, + key: FloatTensor>, + value: FloatTensor>, + mask: Option>>, + attn_bias: Option>>, + options: burn_backend::ops::AttentionModuleOptions, + ) -> FloatTensor> { + make_ops!( + AttentionOps, + AttentionOpIr, + |args: &AttentionOpIr, handles: &mut HandleContainer| { + let query = handles.get_float_tensor::(&args.query); + let key = handles.get_float_tensor::(&args.key); + let value = handles.get_float_tensor::(&args.value); + let mask = args.mask.as_ref().map(|m| handles.get_bool_tensor::(m)); + let attn_bias = args + .attn_bias + .as_ref() + .map(|ab| handles.get_float_tensor::(ab)); + + let output = B::attention( + query, + key, + value, + mask, + attn_bias, + args.options.clone().into(), + ); + + handles.register_float_tensor::(&args.out.id, output); + } + ); + + let streams = StreamId::current(); + + let client = query.client.clone(); + let desc = AttentionOpIr::create( + query.into_ir(), + key.into_ir(), + value.into_ir(), + mask.map(|m| m.into_ir()), + attn_bias.map(|ab| ab.into_ir()), + options.into(), + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::Attention(desc.clone())), + AttentionOps::::new(desc), + ) + .output() + } + + fn rfft( + signal: FloatTensor>, + dim: usize, + n: Option, + ) -> (FloatTensor>, FloatTensor>) { + make_ops!(RfftOps, RfftOpIr, |desc: &RfftOpIr, + handles: &mut HandleContainer< + B::Handle, + >| { + let signal = handles.get_float_tensor::(&desc.signal); + let (re, im) = B::rfft(signal, desc.dim, desc.n); + + handles.register_float_tensor::(&desc.out_re.id, re); + handles.register_float_tensor::(&desc.out_im.id, im); + }); + + let streams = StreamId::current(); + let client = signal.client.clone(); + + let desc = RfftOpIr::create(signal.into_ir(), dim, n, || client.create_empty_handle()); + + let mut outputs = client + .register( + streams, + OperationIr::Module(ModuleOperationIr::Rfft(desc.clone())), + RfftOps::::new(desc), + ) + .into_iter(); + + (outputs.next().unwrap(), outputs.next().unwrap()) + } + + fn irfft( + spectrum_re: FloatTensor>, + spectrum_im: FloatTensor>, + dim: usize, + n: Option, + ) -> FloatTensor> { + make_ops!(IRfftOps, IRfftOpIr, |desc: &IRfftOpIr, + handles: &mut HandleContainer< + B::Handle, + >| { + let input_re = handles.get_float_tensor::(&desc.input_re); + let input_im = handles.get_float_tensor::(&desc.input_im); + + let signal = B::irfft(input_re, input_im, desc.dim, desc.n); + handles.register_float_tensor::(&desc.out_signal.id, signal); + }); + + let streams = StreamId::current(); + let client = spectrum_re.client.clone(); + + let desc = IRfftOpIr::create(spectrum_re.into_ir(), spectrum_im.into_ir(), dim, n, || { + client.create_empty_handle() + }); + + let mut outputs = client + .register( + streams, + OperationIr::Module(ModuleOperationIr::IRfft(desc.clone())), + IRfftOps::::new(desc), + ) + .into_iter(); + + outputs.next().unwrap() + } + + fn has_ctc_loss_backward() -> bool { + B::has_ctc_loss_backward() + } + + fn ctc_loss( + log_probs: FloatTensor>, + targets: IntTensor>, + input_lengths: IntTensor>, + target_lengths: IntTensor>, + blank: usize, + ) -> FloatTensor> { + // CTC is treated as its own non-fuseable IR node, the same way + // `attention` and `conv2d` are. The execute callback drains the input + // handles and dispatches to `B::ctc_loss` on the inner backend, which + // either runs a native kernel (cubecl, libtorch) or the decomposed + // default - either way it executes on raw inner-backend tensors, + // never re-entering the fusion stream. + make_ops!(CtcLossOps, CtcLossOpIr, |args: &CtcLossOpIr, + handles: &mut HandleContainer< + B::Handle, + >| { + let log_probs = handles.get_float_tensor::(&args.log_probs); + let targets = handles.get_int_tensor::(&args.targets); + let input_lengths = handles.get_int_tensor::(&args.input_lengths); + let target_lengths = handles.get_int_tensor::(&args.target_lengths); + let output = B::ctc_loss( + log_probs, + targets, + input_lengths, + target_lengths, + args.blank, + ); + handles.register_float_tensor::(&args.out.id, output); + }); + + let streams = StreamId::current(); + let client = log_probs.client.clone(); + let desc = CtcLossOpIr::create( + log_probs.into_ir(), + targets.into_ir(), + input_lengths.into_ir(), + target_lengths.into_ir(), + blank, + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::CtcLoss(desc.clone())), + CtcLossOps::::new(desc), + ) + .output() + } + + fn ctc_loss_backward( + log_probs: FloatTensor>, + targets: IntTensor>, + input_lengths: IntTensor>, + target_lengths: IntTensor>, + grad_loss: FloatTensor>, + blank: usize, + ) -> FloatTensor> { + // Mirrors `ctc_loss`: a typed IR node that dispatches to the inner + // backend's native backward kernel. + make_ops!( + CtcLossBackwardOps, + CtcLossBackwardOpIr, + |args: &CtcLossBackwardOpIr, handles: &mut HandleContainer| { + let log_probs = handles.get_float_tensor::(&args.log_probs); + let targets = handles.get_int_tensor::(&args.targets); + let input_lengths = handles.get_int_tensor::(&args.input_lengths); + let target_lengths = handles.get_int_tensor::(&args.target_lengths); + let grad_loss = handles.get_float_tensor::(&args.grad_loss); + let output = B::ctc_loss_backward( + log_probs, + targets, + input_lengths, + target_lengths, + grad_loss, + args.blank, + ); + handles.register_float_tensor::(&args.out.id, output); + } + ); + + let streams = StreamId::current(); + let client = log_probs.client.clone(); + let desc = CtcLossBackwardOpIr::create( + log_probs.into_ir(), + targets.into_ir(), + input_lengths.into_ir(), + target_lengths.into_ir(), + grad_loss.into_ir(), + blank, + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::Module(ModuleOperationIr::CtcLossBackward(desc.clone())), + CtcLossBackwardOps::::new(desc), + ) + .output() + } +} diff --git a/crates/burn-fusion/src/ops/qtensor.rs b/crates/burn-fusion/src/ops/qtensor.rs new file mode 100644 index 0000000..c0ffece --- /dev/null +++ b/crates/burn-fusion/src/ops/qtensor.rs @@ -0,0 +1,502 @@ +use std::marker::PhantomData; + +use burn_backend::{ + DType, ExecutionError, FloatDType, Shape, Slice, TensorData, TensorMetadata, TensorPrimitive, + get_device_settings, + ops::QTensorOps, + quantization::{QuantPropagation, QuantScheme, QuantizationParametersPrimitive}, + tensor::{Device, FloatTensor, IntTensor, QuantizedTensor}, +}; +use burn_ir::{ + BaseOperationIr, DequantizeOpIr, FlipOpIr, FloatOperationIr, GatherOpIr, HandleContainer, + InitOperationIr, MatmulOpIr, OperationIr, OperationOutput, PermuteOpIr, + QuantizationParametersIr, QuantizeOpIr, SelectOpIr, ShapeOpIr, SliceOpIr, SwapDimsOpIr, +}; + +use crate::{ + Fusion, FusionBackend, + client::GlobalFusionClient, + get_client, + stream::{StreamId, execution::Operation}, +}; + +use super::NoOp; + +impl QTensorOps for Fusion { + fn q_from_data(data: TensorData, device: &Device) -> QuantizedTensor { + let client = get_client::(device); + let dtype = data.dtype; + let tensor = B::q_from_data(data, device); + let shape = burn_backend::TensorMetadata::shape(&tensor); + + let handle = B::quantized_tensor_handle(tensor); + let desc = InitOperationIr::create(shape, dtype, || client.register_tensor_handle(handle)); + + client + .register( + StreamId::current(), + OperationIr::Init(desc), + NoOp::::new(), + ) + .output() + } + + fn quantize( + tensor: FloatTensor, + scheme: &QuantScheme, + qparams: QuantizationParametersPrimitive, + ) -> QuantizedTensor { + #[derive(new, Debug)] + struct QuantizeOp { + desc: QuantizeOpIr, + _b: PhantomData, + } + + impl Operation for QuantizeOp { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_float_tensor::(&self.desc.tensor); + let scales = handles.get_float_tensor::(&self.desc.qparams.scales); + + let qparams = QuantizationParametersPrimitive { scales }; + let output = B::quantize(tensor, &self.desc.scheme, qparams); + handles.register_quantized_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let qparams = QuantizationParametersIr { + scales: qparams.scales.into_ir(), + }; + let desc = QuantizeOpIr::create(tensor.into_ir(), qparams, *scheme, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Float(desc.tensor.dtype, FloatOperationIr::Quantize(desc.clone())), + QuantizeOp::::new(desc), + ) + .output() + } + + fn dequantize(tensor: QuantizedTensor, dtype: FloatDType) -> FloatTensor { + #[derive(new, Debug)] + struct DequantizeOp { + desc: DequantizeOpIr, + _b: PhantomData, + } + + impl Operation for DequantizeOp { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_quantized_tensor::(&self.desc.input); + + let output = B::dequantize(tensor, self.desc.out.dtype.into()); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let dtype = dtype.into(); + let desc = DequantizeOpIr::create(tensor.into_ir(), dtype, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(dtype, FloatOperationIr::Dequantize(desc.clone())), + DequantizeOp::::new(desc), + ) + .output() + } + + fn q_to_device( + tensor: QuantizedTensor, + device_dst: &Device, + ) -> QuantizedTensor { + let device_src: &B::Device = tensor.client.device(); + let device_dst: B::Device = device_dst.clone(); + + if device_src == &device_dst { + return tensor; + } + + let id = tensor.stream; + let client_dst = get_client::(&device_dst); + let client_src = tensor.client.clone(); + + GlobalFusionClient::change_client_quantized::( + tensor.into_ir(), + client_src, + client_dst, + id, + ) + } + + fn q_reshape(tensor: QuantizedTensor, shape: Shape) -> QuantizedTensor { + if tensor.shape == shape { + return tensor; + } + + #[derive(new, Debug)] + struct ReshapeDimsOps { + desc: ShapeOpIr, + _b: PhantomData, + } + + impl Operation for ReshapeDimsOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_quantized_tensor::(&self.desc.input); + let output = B::q_reshape(input, self.desc.out.shape.clone()); + handles.register_quantized_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ShapeOpIr::reshape(tensor.into_ir(), shape, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Reshape(desc.clone())), + ReshapeDimsOps::::new(desc), + ) + .output() + } + + async fn q_into_data(tensor: QuantizedTensor) -> Result { + tensor.q_into_data::().await + } + + fn q_swap_dims( + tensor: QuantizedTensor, + dim1: usize, + dim2: usize, + ) -> QuantizedTensor { + #[derive(new, Debug)] + struct SwapDimsOps { + desc: SwapDimsOpIr, + _b: PhantomData, + } + + impl Operation for SwapDimsOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_quantized_tensor::(&self.desc.input); + let output = B::q_swap_dims(input, self.desc.dim1, self.desc.dim2); + handles.register_quantized_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = SwapDimsOpIr::create(tensor.into_ir(), dim1, dim2, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::SwapDims(desc.clone())), + SwapDimsOps::::new(desc), + ) + .output() + } + + fn q_permute(tensor: QuantizedTensor, axes: &[usize]) -> QuantizedTensor { + #[derive(new, Debug)] + struct PermuteDimsOps { + desc: PermuteOpIr, + _b: PhantomData, + } + + impl Operation for PermuteDimsOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_quantized_tensor::(&self.desc.input); + let output = B::q_permute(input, self.desc.axes.as_slice()); + handles.register_quantized_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = PermuteOpIr::create(tensor.into_ir(), axes.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Permute(desc.clone())), + PermuteDimsOps::::new(desc), + ) + .output() + } + + fn q_flip(tensor: QuantizedTensor, axes: &[usize]) -> QuantizedTensor { + #[derive(new, Debug)] + struct FlipOps { + desc: FlipOpIr, + _b: PhantomData, + } + + impl Operation for FlipOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_quantized_tensor::(&self.desc.input); + let output = B::q_flip(input, &self.desc.axes); + handles.register_quantized_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = FlipOpIr::create(tensor.into_ir(), axes.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Flip(desc.clone())), + FlipOps::::new(desc), + ) + .output() + } + + fn q_gather( + dim: usize, + tensor: QuantizedTensor, + indices: IntTensor, + ) -> QuantizedTensor { + #[derive(new, Debug)] + struct GatherOps { + desc: GatherOpIr, + _b: PhantomData, + } + + impl Operation for GatherOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_quantized_tensor::(&self.desc.tensor); + let indices = handles.get_int_tensor::(&self.desc.indices); + + let output = B::q_gather(self.desc.dim, tensor, indices); + handles.register_quantized_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = GatherOpIr::create(tensor.into_ir(), dim, indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Gather(desc.clone())), + GatherOps::::new(desc), + ) + .output() + } + + fn q_select( + tensor: QuantizedTensor, + dim: usize, + indices: IntTensor, + ) -> QuantizedTensor { + #[derive(new, Debug)] + struct SelectOps { + desc: SelectOpIr, + _b: PhantomData, + } + + impl Operation for SelectOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_quantized_tensor::(&self.desc.tensor); + let indices = handles.get_int_tensor::(&self.desc.indices); + + let output = B::q_select(tensor, self.desc.dim, indices); + + handles.register_quantized_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = SelectOpIr::create(tensor.into_ir(), dim, indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Select(desc.clone())), + SelectOps::::new(desc), + ) + .output() + } + + fn q_slice(tensor: QuantizedTensor, slices: &[Slice]) -> QuantizedTensor { + #[derive(new, Debug)] + struct SliceOps { + desc: SliceOpIr, + _b: PhantomData, + } + + impl Operation for SliceOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_quantized_tensor::(&self.desc.tensor); + + let output = B::q_slice(tensor, self.desc.ranges.as_slice()); + + handles.register_quantized_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = SliceOpIr::create(tensor.into_ir(), slices.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Slice(desc.clone())), + SliceOps::::new(desc), + ) + .output() + } + + fn q_expand(tensor: QuantizedTensor, shape: Shape) -> QuantizedTensor { + #[derive(new, Debug)] + struct ExpandOps { + desc: ShapeOpIr, + _b: PhantomData, + } + + impl Operation for ExpandOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_quantized_tensor::(&self.desc.input); + let output = B::q_expand(input, self.desc.out.shape.clone()); + + handles.register_quantized_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ShapeOpIr::expand(tensor.into_ir(), shape, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Expand(desc.clone())), + ExpandOps::::new(desc), + ) + .output() + } + + fn q_matmul(lhs: TensorPrimitive, rhs: TensorPrimitive) -> TensorPrimitive { + #[derive(new, Debug)] + struct MatmulOps { + desc: MatmulOpIr, + lhs_quantized: bool, + rhs_quantized: bool, + _b: PhantomData, + } + + impl Operation for MatmulOps { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = match self.lhs_quantized { + true => { + TensorPrimitive::QFloat(handles.get_quantized_tensor::(&self.desc.lhs)) + } + false => TensorPrimitive::Float(handles.get_float_tensor::(&self.desc.lhs)), + }; + let rhs = match self.rhs_quantized { + true => { + TensorPrimitive::QFloat(handles.get_quantized_tensor::(&self.desc.rhs)) + } + false => TensorPrimitive::Float(handles.get_float_tensor::(&self.desc.rhs)), + }; + let output = B::q_matmul(lhs, rhs); + match output { + TensorPrimitive::Float(output) => { + handles.register_float_tensor::(&self.desc.out.id, output); + } + TensorPrimitive::QFloat(output) => { + handles.register_quantized_tensor::(&self.desc.out.id, output); + } + } + } + } + + let mut propagation = QuantPropagation::Inhibit; + let mut scheme = QuantScheme::default(); + let streams = StreamId::current(); + let mut lhs_quantized = false; + let mut rhs_quantized = false; + + let client = match &lhs { + TensorPrimitive::Float(lhs) => lhs.client.clone(), + TensorPrimitive::QFloat(lhs) => lhs.client.clone(), + }; + + let settings = get_device_settings::(client.device()); + + if let TensorPrimitive::QFloat(lhs) = &lhs { + propagation = settings.quantization.propagation; + scheme = lhs.scheme(); + lhs_quantized = true; + } + if let TensorPrimitive::QFloat(rhs) = &rhs { + propagation = settings.quantization.propagation; + scheme = rhs.scheme(); + rhs_quantized = true; + } + + let dtype = match propagation { + QuantPropagation::Propagate => DType::QFloat(scheme), + QuantPropagation::Inhibit => match (&lhs, &rhs) { + (TensorPrimitive::Float(t), _) | (_, TensorPrimitive::Float(t)) => t.dtype, + _ => settings.float_dtype.into(), + }, + }; + + let lhs = match lhs { + TensorPrimitive::Float(lhs) => lhs.into_ir(), + TensorPrimitive::QFloat(lhs) => lhs.into_ir(), + }; + let rhs = match rhs { + TensorPrimitive::Float(rhs) => rhs.into_ir(), + TensorPrimitive::QFloat(rhs) => rhs.into_ir(), + }; + + let desc = MatmulOpIr::create_mixed(lhs, rhs, dtype, || client.create_empty_handle()); + + let out = client + .register( + streams, + OperationIr::Float(dtype, FloatOperationIr::Matmul(desc.clone())), + MatmulOps::::new(desc, lhs_quantized, rhs_quantized), + ) + .output(); + + match propagation { + QuantPropagation::Propagate => TensorPrimitive::QFloat(out), + QuantPropagation::Inhibit => TensorPrimitive::Float(out), + } + } +} diff --git a/crates/burn-fusion/src/ops/tensor.rs b/crates/burn-fusion/src/ops/tensor.rs new file mode 100644 index 0000000..c47ac05 --- /dev/null +++ b/crates/burn-fusion/src/ops/tensor.rs @@ -0,0 +1,2592 @@ +use super::NoOp; +use crate::{ + Fusion, FusionBackend, binary_float_cmp_ops, binary_float_ops, + client::GlobalFusionClient, + get_client, reduce_float_ops, reduce_float2int_ops, scalar_float_cmp_ops, scalar_float_ops, + stream::{StreamId, execution::Operation}, + unary_float_ops, +}; +use burn_backend::{ + BoolDType, Distribution, ExecutionError, FloatDType, IntDType, Scalar, Shape, Slice, + TensorData, + ops::{FloatTensorOps, GridSampleOptions}, + tensor::{BoolTensor, Device, FloatTensor, IndexingUpdateOp, IntTensor}, +}; +use burn_ir::*; +use std::marker::PhantomData; + +impl FloatTensorOps for Fusion { + #[cfg_attr(feature = "tracing", tracing::instrument( + level="trace", + skip(data), + fields(?data.shape, ?data.dtype) + ))] + fn float_from_data(data: TensorData, device: &Device) -> FloatTensor { + let client = get_client::(device); + let dtype = data.dtype; + let tensor = B::float_from_data(data, device); + let shape = burn_backend::TensorMetadata::shape(&tensor); + + let handle = B::float_tensor_handle(tensor); + let desc = InitOperationIr::create(shape, dtype, || client.register_tensor_handle(handle)); + + client + .register( + StreamId::current(), + OperationIr::Init(desc), + NoOp::::new(), + ) + .output() + } + + fn float_random( + shape: Shape, + distribution: Distribution, + device: &Device, + dtype: FloatDType, + ) -> FloatTensor { + #[derive(new, Debug)] + struct RandomOps { + desc: RandomOpIr, + device: Device, + } + + impl Operation for RandomOps { + fn execute(&self, handles: &mut HandleContainer) { + let output: B::FloatTensorPrimitive = B::float_random( + self.desc.out.shape.clone(), + self.desc.distribution, + &self.device, + self.desc.out.dtype.into(), + ); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let dtype = dtype.into(); + let client = get_client::(device); + let desc = RandomOpIr::create(shape, dtype, distribution, || client.create_empty_handle()); + + client + .register( + StreamId::current(), + OperationIr::Float(dtype, FloatOperationIr::Random(desc.clone())), + RandomOps::::new(desc, device.clone()), + ) + .output() + } + + fn float_zeros(shape: Shape, device: &Device, dtype: FloatDType) -> FloatTensor { + #[derive(new, Debug)] + struct ZerosOps { + out: TensorIr, + device: Device, + } + + impl Operation for ZerosOps { + fn execute(&self, handles: &mut HandleContainer) { + let shape = self.out.shape.clone(); + let output = B::float_zeros(shape, &self.device, self.out.dtype.into()); + handles.register_float_tensor::(&self.out.id, output); + } + } + + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register( + StreamId::current(), + OperationIr::BaseFloat(BaseOperationIr::Zeros(desc.clone())), + ZerosOps::::new(desc.out, device.clone()), + ) + .output() + } + + fn float_ones(shape: Shape, device: &Device, dtype: FloatDType) -> FloatTensor { + #[derive(new, Debug)] + struct OnesOps { + out: TensorIr, + device: Device, + } + + impl Operation for OnesOps { + fn execute(&self, handles: &mut HandleContainer) { + let shape = self.out.shape.clone(); + let output = B::float_ones(shape, &self.device, self.out.dtype.into()); + handles.register_float_tensor::(&self.out.id, output); + } + } + + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register( + StreamId::current(), + OperationIr::BaseFloat(BaseOperationIr::Ones(desc.clone())), + OnesOps::::new(desc.out, device.clone()), + ) + .output() + } + + fn float_full( + shape: Shape, + fill_value: Scalar, + device: &Device, + dtype: FloatDType, + ) -> FloatTensor { + #[derive(new, Debug)] + struct FullOps { + out: TensorIr, + elem: ScalarIr, + device: Device, + } + + impl Operation for FullOps { + fn execute(&self, handles: &mut HandleContainer) { + let shape = self.out.shape.clone(); + let dtype = self.out.dtype.into(); + let output: B::FloatTensorPrimitive = + B::float_full(shape, self.elem.into(), &self.device, dtype); + handles.register_float_tensor::(&self.out.id, output); + } + } + + let dtype = dtype.into(); + let client = get_client::(device); + let value = fill_value.into(); + let desc = FullOpIr::create(shape, dtype, value, || client.create_empty_handle()); + + client + .register( + StreamId::current(), + OperationIr::NumericFloat(dtype, NumericOperationIr::Full(desc.clone())), + FullOps::::new(desc.out, desc.value, device.clone()), + ) + .output() + } + + #[cfg_attr(feature = "tracing", tracing::instrument( + level="trace", + skip(tensor), + fields( + from = ?tensor.client.device(), + shape = ?tensor.shape, + dtype = ?tensor.dtype + ) + ))] + async fn float_into_data(tensor: FloatTensor) -> Result { + tensor.into_data::().await + } + + #[cfg_attr(feature = "tracing", tracing::instrument( + level="trace", + skip(tensor), + fields( + from = ?tensor.client.device(), + shape = ?tensor.shape, + dtype = ?tensor.dtype, + ) + ))] + fn float_to_device(tensor: FloatTensor, device_dst: &Device) -> FloatTensor { + let device_src: &B::Device = tensor.client.device(); + + if device_src == device_dst { + return tensor; + } + + let id = tensor.stream; + let client_dst = get_client::(device_dst); + let client_src = tensor.client.clone(); + + GlobalFusionClient::change_client_float::(tensor.into_ir(), client_src, client_dst, id) + } + + fn float_into_int(tensor: FloatTensor, dtype: IntDType) -> IntTensor { + #[derive(new, Debug)] + struct IntoIntOps { + desc: CastOpIr, + _b: PhantomData, + } + + impl Operation for IntoIntOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = B::float_into_int(input, self.desc.out.dtype.into()); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = CastOpIr::create(tensor.into_ir(), dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Float(desc.input.dtype, FloatOperationIr::IntoInt(desc.clone())), + IntoIntOps::::new(desc), + ) + .output() + } + + fn float_empty(shape: Shape, device: &Device, dtype: FloatDType) -> FloatTensor { + #[derive(new, Debug)] + struct EmptyOps { + desc: TensorIr, + device: Device, + } + + impl Operation for EmptyOps { + fn execute(&self, handles: &mut HandleContainer) { + let output = B::float_empty( + self.desc.shape.clone(), + &self.device, + self.desc.dtype.into(), + ); + handles.register_float_tensor::(&self.desc.id, output); + } + } + + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register( + StreamId::current(), + OperationIr::BaseFloat(BaseOperationIr::Empty(desc.clone())), + EmptyOps::::new(desc.out, device.clone()), + ) + .output() + } + + fn float_add(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float_ops!(AddOps, B::float_add); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::Add(desc.clone())), + AddOps::::new(desc), + ) + .output() + } + + fn float_add_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + scalar_float_ops!(AddOps, B::float_add_scalar); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::AddScalar(desc.clone()), + ), + AddOps::::new(desc), + ) + .output() + } + + fn float_clamp(tensor: FloatTensor, min: Scalar, max: Scalar) -> FloatTensor { + #[derive(new, Debug)] + struct ClampOps { + desc: ClampOpIr, + _b: PhantomData, + } + + impl Operation for ClampOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.tensor); + let output = B::float_clamp(input, self.desc.min.into(), self.desc.max.into()); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let min = min.into(); + let max = max.into(); + let desc = ClampOpIr::create(tensor.into_ir(), min, max, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.tensor.dtype, + NumericOperationIr::Clamp(desc.clone()), + ), + ClampOps::::new(desc), + ) + .output() + } + + fn float_sub(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float_ops!(SubOps, B::float_sub); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::Sub(desc.clone())), + SubOps::::new(desc), + ) + .output() + } + + fn float_sub_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + scalar_float_ops!(SubOps, B::float_sub_scalar); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::SubScalar(desc.clone()), + ), + SubOps::::new(desc), + ) + .output() + } + + fn float_mul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float_ops!(MulOps, B::float_mul); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::Mul(desc.clone())), + MulOps::::new(desc), + ) + .output() + } + + fn float_mul_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + scalar_float_ops!(MulOps, B::float_mul_scalar); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::MulScalar(desc.clone()), + ), + MulOps::::new(desc), + ) + .output() + } + + fn float_div(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float_ops!(DivOps, B::float_div); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::Div(desc.clone())), + DivOps::::new(desc), + ) + .output() + } + + fn float_div_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + scalar_float_ops!(DivOps, B::float_div_scalar); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::DivScalar(desc.clone()), + ), + DivOps::::new(desc), + ) + .output() + } + + fn float_remainder(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float_ops!(ModOps, B::float_remainder); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::Rem(desc.clone())), + ModOps::::new(desc), + ) + .output() + } + + fn float_remainder_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + scalar_float_ops!(ModOps, B::float_remainder_scalar); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::RemScalar(desc.clone()), + ), + ModOps::::new(desc), + ) + .output() + } + + fn float_matmul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float_ops!(MatmulOps, B::float_matmul); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = MatmulOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Matmul(desc.clone())), + MatmulOps::::new(desc.into()), + ) + .output() + } + + fn float_cross( + lhs: FloatTensor, + rhs: FloatTensor, + dim: usize, + ) -> FloatTensor { + #[derive(new, Debug)] + struct CrossOps { + desc: CrossOpIr, + _b: PhantomData, + } + + impl Operation for CrossOps { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = handles.get_float_tensor::(&self.desc.lhs); + let rhs = handles.get_float_tensor::(&self.desc.rhs); + let output = B::float_cross(lhs, rhs, self.desc.dim); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = CrossOpIr::create(lhs.into_ir(), rhs.into_ir(), dim, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Cross(desc.clone())), + CrossOps::::new(desc), + ) + .output() + } + + fn float_swap_dims(tensor: FloatTensor, dim1: usize, dim2: usize) -> FloatTensor { + #[derive(new, Debug)] + struct SwapDimsOps { + desc: SwapDimsOpIr, + _b: PhantomData, + } + + impl Operation for SwapDimsOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = B::float_swap_dims(input, self.desc.dim1, self.desc.dim2); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = SwapDimsOpIr::create(tensor.into_ir(), dim1, dim2, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::SwapDims(desc.clone())), + SwapDimsOps::::new(desc), + ) + .output() + } + + fn float_reshape(tensor: FloatTensor, shape: Shape) -> FloatTensor { + if tensor.shape == shape { + return tensor; + } + + #[derive(new, Debug)] + struct ReshapeDimsOps { + desc: ShapeOpIr, + _b: PhantomData, + } + + impl Operation for ReshapeDimsOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = B::float_reshape(input, self.desc.out.shape.clone()); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ShapeOpIr::reshape(tensor.into_ir(), shape, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Reshape(desc.clone())), + ReshapeDimsOps::::new(desc), + ) + .output() + } + + fn float_gather( + dim: usize, + tensor: FloatTensor, + indices: IntTensor, + ) -> FloatTensor { + #[derive(new, Debug)] + struct GatherOps { + desc: GatherOpIr, + _b: PhantomData, + } + + impl Operation for GatherOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_float_tensor::(&self.desc.tensor); + let indices = handles.get_int_tensor::(&self.desc.indices); + + let output = B::float_gather(self.desc.dim, tensor, indices); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = GatherOpIr::create(tensor.into_ir(), dim, indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Gather(desc.clone())), + GatherOps::::new(desc), + ) + .output() + } + + fn float_scatter_add( + dim: usize, + tensor: FloatTensor, + indices: IntTensor, + value: FloatTensor, + ) -> FloatTensor { + #[derive(new, Debug)] + struct ScatterOps { + desc: ScatterOpIr, + _b: PhantomData, + } + + impl Operation for ScatterOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_float_tensor::(&self.desc.tensor); + let indices = handles.get_int_tensor::(&self.desc.indices); + let value = handles.get_float_tensor::(&self.desc.value); + + let output = B::float_scatter_add(self.desc.dim, tensor, indices, value); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ScatterOpIr::create( + tensor.into_ir(), + dim, + indices.into_ir(), + value.into_ir(), + IndexingUpdateOp::Add, + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Scatter(desc.clone())), + ScatterOps::::new(desc), + ) + .output() + } + + fn float_scatter_nd( + data: FloatTensor, + indices: IntTensor, + values: FloatTensor, + reduction: IndexingUpdateOp, + ) -> FloatTensor { + #[derive(new, Debug)] + struct ScatterNdOps { + desc: ScatterNdOpIr, + _b: PhantomData, + } + + impl Operation for ScatterNdOps { + fn execute(&self, handles: &mut HandleContainer) { + let data = handles.get_float_tensor::(&self.desc.data); + let indices = handles.get_int_tensor::(&self.desc.indices); + let values = handles.get_float_tensor::(&self.desc.values); + + let output = B::float_scatter_nd(data, indices, values, self.desc.reduction); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = data.client.clone(); + let desc = ScatterNdOpIr::create( + data.into_ir(), + indices.into_ir(), + values.into_ir(), + reduction, + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::ScatterNd(desc.clone())), + ScatterNdOps::::new(desc), + ) + .output() + } + + fn float_gather_nd(data: FloatTensor, indices: IntTensor) -> FloatTensor { + #[derive(new, Debug)] + struct GatherNdOps { + desc: GatherNdOpIr, + _b: PhantomData, + } + + impl Operation for GatherNdOps { + fn execute(&self, handles: &mut HandleContainer) { + let data = handles.get_float_tensor::(&self.desc.data); + let indices = handles.get_int_tensor::(&self.desc.indices); + + let output = B::float_gather_nd(data, indices); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = data.client.clone(); + let desc = GatherNdOpIr::create(data.into_ir(), indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::GatherNd(desc.clone())), + GatherNdOps::::new(desc), + ) + .output() + } + + fn float_select( + tensor: FloatTensor, + dim: usize, + indices: IntTensor, + ) -> FloatTensor { + #[derive(new, Debug)] + struct SelectOps { + desc: SelectOpIr, + _b: PhantomData, + } + + impl Operation for SelectOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_float_tensor::(&self.desc.tensor); + let indices = handles.get_int_tensor::(&self.desc.indices); + + let output = B::float_select(tensor, self.desc.dim, indices); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = SelectOpIr::create(tensor.into_ir(), dim, indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Select(desc.clone())), + SelectOps::::new(desc), + ) + .output() + } + + fn float_select_add( + tensor: FloatTensor, + dim: usize, + indices: IntTensor, + value: FloatTensor, + ) -> FloatTensor { + #[derive(new, Debug)] + struct SelectAssignOps { + desc: SelectAssignOpIr, + _b: PhantomData, + } + + impl Operation for SelectAssignOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_float_tensor::(&self.desc.tensor); + let indices = handles.get_int_tensor::(&self.desc.indices); + let value = handles.get_float_tensor::(&self.desc.value); + + let output = B::float_select_add(tensor, self.desc.dim, indices, value); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = SelectAssignOpIr::create( + tensor.into_ir(), + dim, + indices.into_ir(), + value.into_ir(), + IndexingUpdateOp::Add, + || client.create_empty_handle(), + ); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::SelectAssign(desc.clone())), + SelectAssignOps::::new(desc), + ) + .output() + } + + fn float_slice(tensor: FloatTensor, slices: &[Slice]) -> FloatTensor { + #[derive(new, Debug)] + struct SliceOps { + desc: SliceOpIr, + _b: PhantomData, + } + + impl Operation for SliceOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_float_tensor::(&self.desc.tensor); + + let output = B::float_slice(tensor, self.desc.ranges.as_slice()); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = SliceOpIr::create(tensor.into_ir(), slices.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Slice(desc.clone())), + SliceOps::::new(desc), + ) + .output() + } + + fn float_slice_assign( + tensor: FloatTensor, + slices: &[burn_backend::Slice], + value: FloatTensor, + ) -> FloatTensor { + #[derive(new, Debug)] + struct SliceAssignOps { + desc: SliceAssignOpIr, + _b: PhantomData, + } + + impl Operation for SliceAssignOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_float_tensor::(&self.desc.tensor); + let value = handles.get_float_tensor::(&self.desc.value); + + let output = B::float_slice_assign(tensor, self.desc.ranges.as_slice(), value); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = + SliceAssignOpIr::create(tensor.into_ir(), slices.into(), value.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::SliceAssign(desc.clone())), + SliceAssignOps::::new(desc), + ) + .output() + } + + fn float_mask_where( + tensor: FloatTensor, + mask: BoolTensor, + value: FloatTensor, + ) -> FloatTensor { + #[derive(new, Debug)] + struct MaskWhereOps { + desc: MaskWhereOpIr, + _b: PhantomData, + } + + impl Operation for MaskWhereOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_float_tensor::(&self.desc.tensor); + let value = handles.get_float_tensor::(&self.desc.value); + let mask = handles.get_bool_tensor::(&self.desc.mask); + + let output = B::float_mask_where(tensor, mask, value); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = MaskWhereOpIr::create(tensor.into_ir(), mask.into_ir(), value.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::MaskWhere(desc.clone())), + MaskWhereOps::::new(desc), + ) + .output() + } + + fn float_mask_fill( + tensor: FloatTensor, + mask: BoolTensor, + value: Scalar, + ) -> FloatTensor { + #[derive(new, Debug)] + struct MaskFillOps { + desc: MaskFillOpIr, + _b: PhantomData, + } + + impl Operation for MaskFillOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_float_tensor::(&self.desc.tensor); + let mask = handles.get_bool_tensor::(&self.desc.mask); + + let output = B::float_mask_fill(tensor, mask, self.desc.value.into()); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let value = value.into(); + let desc = MaskFillOpIr::create(tensor.into_ir(), mask.into_ir(), value, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::MaskFill(desc.clone())), + MaskFillOps::::new(desc), + ) + .output() + } + + fn float_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_float_cmp_ops!(EqualOps, B::float_equal); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Equal(desc.clone())), + EqualOps::::new(desc), + ) + .output() + } + + fn float_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + scalar_float_cmp_ops!(EqualElemOps, B::float_equal_elem); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::EqualElem(desc.clone())), + EqualElemOps::::new(desc), + ) + .output() + } + + fn float_greater( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_float_cmp_ops!(GreaterOps, B::float_greater); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.lhs.dtype, + NumericOperationIr::Greater(desc.clone()), + ), + GreaterOps::::new(desc), + ) + .output() + } + + fn float_greater_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + scalar_float_cmp_ops!(GreaterElemOps, B::float_greater_elem); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.lhs.dtype, + NumericOperationIr::GreaterElem(desc.clone()), + ), + GreaterElemOps::::new(desc), + ) + .output() + } + + fn float_greater_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_float_cmp_ops!(GreaterEqualOps, B::float_greater_equal); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.lhs.dtype, + NumericOperationIr::GreaterEqual(desc.clone()), + ), + GreaterEqualOps::::new(desc), + ) + .output() + } + + fn float_greater_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + scalar_float_cmp_ops!(GreaterEqualElemOps, B::float_greater_equal_elem); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.lhs.dtype, + NumericOperationIr::GreaterEqualElem(desc.clone()), + ), + GreaterEqualElemOps::::new(desc), + ) + .output() + } + + fn float_lower( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_float_cmp_ops!(LowerOps, B::float_lower); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat(desc.lhs.dtype, NumericOperationIr::Lower(desc.clone())), + LowerOps::::new(desc), + ) + .output() + } + + fn float_lower_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + scalar_float_cmp_ops!(LowerElemOps, B::float_lower_elem); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.lhs.dtype, + NumericOperationIr::LowerElem(desc.clone()), + ), + LowerElemOps::::new(desc), + ) + .output() + } + + fn float_lower_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + binary_float_cmp_ops!(LowerEqualOps, B::float_lower_equal); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.lhs.dtype, + NumericOperationIr::LowerEqual(desc.clone()), + ), + LowerEqualOps::::new(desc), + ) + .output() + } + + fn float_lower_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + scalar_float_cmp_ops!(LowerEqualElemOps, B::float_lower_equal_elem); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.lhs.dtype, + NumericOperationIr::LowerEqualElem(desc.clone()), + ), + LowerEqualElemOps::::new(desc), + ) + .output() + } + + fn float_sum(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(SumOps, B::float_sum, reduce); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::Sum(desc.clone())), + SumOps::::new(desc.into()), + ) + .output() + } + + fn float_sum_dim(tensor: FloatTensor, axis: usize) -> FloatTensor { + reduce_float_ops!(SumDimOps, |tensor, axis, _| B::float_sum_dim(tensor, axis)); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = + ReduceDimOpIr::create(tensor.into_ir(), axis, 1, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::SumDim(desc.clone())), + SumDimOps::::new(desc), + ) + .output() + } + + fn float_prod(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(ProdOps, B::float_prod, reduce); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::Prod(desc.clone())), + ProdOps::::new(desc.into()), + ) + .output() + } + + fn float_prod_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + reduce_float_ops!(ProdDimOps, |tensor, axis, _| B::float_prod_dim( + tensor, axis + )); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::ProdDim(desc.clone()), + ), + ProdDimOps::::new(desc), + ) + .output() + } + + fn float_mean(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(MeanOps, B::float_mean, reduce); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::Mean(desc.clone())), + MeanOps::::new(desc.into()), + ) + .output() + } + + fn float_mean_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + reduce_float_ops!(MeanDimOps, |tensor, axis, _| B::float_mean_dim( + tensor, axis + )); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::MeanDim(desc.clone()), + ), + MeanDimOps::::new(desc), + ) + .output() + } + + fn float_cumsum(tensor: FloatTensor, dim: usize) -> FloatTensor { + #[derive(new, Debug)] + struct CumsumOps { + desc: DimOpIr, + _b: PhantomData, + } + + impl Operation for CumsumOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = B::float_cumsum(input, self.desc.axis); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::CumSum(desc.clone())), + CumsumOps::::new(desc), + ) + .output() + } + + fn float_cumprod(tensor: FloatTensor, dim: usize) -> FloatTensor { + #[derive(new, Debug)] + struct CumprodOps { + desc: DimOpIr, + _b: PhantomData, + } + + impl Operation for CumprodOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = B::float_cumprod(input, self.desc.axis); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::CumProd(desc.clone()), + ), + CumprodOps::::new(desc), + ) + .output() + } + + fn float_cummin(tensor: FloatTensor, dim: usize) -> FloatTensor { + #[derive(new, Debug)] + struct CumminOps { + desc: DimOpIr, + _b: PhantomData, + } + + impl Operation for CumminOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = B::float_cummin(input, self.desc.axis); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::CumMin(desc.clone())), + CumminOps::::new(desc), + ) + .output() + } + + fn float_cummax(tensor: FloatTensor, dim: usize) -> FloatTensor { + #[derive(new, Debug)] + struct CummaxOps { + desc: DimOpIr, + _b: PhantomData, + } + + impl Operation for CummaxOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = B::float_cummax(input, self.desc.axis); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::CumMax(desc.clone())), + CummaxOps::::new(desc), + ) + .output() + } + + fn float_exp(lhs: FloatTensor) -> FloatTensor { + unary_float_ops!(ExpOps, B::float_exp); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = UnaryOpIr::create(lhs.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Exp(desc.clone())), + ExpOps::::new(desc), + ) + .output() + } + + fn float_log(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(LogOps, B::float_log); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Log(desc.clone())), + LogOps::::new(desc), + ) + .output() + } + + fn float_log1p(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(Log1pOps, B::float_log1p); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Log1p(desc.clone())), + Log1pOps::::new(desc), + ) + .output() + } + + fn float_powf_scalar_impl(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + scalar_float_ops!(PowfOps, B::float_powf_scalar_impl); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::PowfScalar(desc.clone())), + PowfOps::::new(desc), + ) + .output() + } + + fn float_sqrt(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(SqrtOps, B::float_sqrt); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Sqrt(desc.clone())), + SqrtOps::::new(desc), + ) + .output() + } + + fn float_abs(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(AbsOps, B::float_abs); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::Abs(desc.clone())), + AbsOps::::new(desc), + ) + .output() + } + + fn float_cos(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(CosOps, B::float_cos); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Cos(desc.clone())), + CosOps::::new(desc), + ) + .output() + } + + fn float_sin(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(SinOps, B::float_sin); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Sin(desc.clone())), + SinOps::::new(desc), + ) + .output() + } + + fn float_tan(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(TanOps, B::float_tan); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Tan(desc.clone())), + TanOps::::new(desc), + ) + .output() + } + + fn float_cosh(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(CoshOps, B::float_cosh); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Cosh(desc.clone())), + CoshOps::::new(desc), + ) + .output() + } + + fn float_sinh(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(SinhOps, B::float_sinh); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Sinh(desc.clone())), + SinhOps::::new(desc), + ) + .output() + } + + fn float_tanh(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(TanhOps, B::float_tanh); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Tanh(desc.clone())), + TanhOps::::new(desc), + ) + .output() + } + + fn float_acos(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(ArcCosOps, B::float_acos); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::ArcCos(desc.clone())), + ArcCosOps::::new(desc), + ) + .output() + } + + fn float_acosh(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(ArcCoshOps, B::float_acosh); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::ArcCosh(desc.clone())), + ArcCoshOps::::new(desc), + ) + .output() + } + + fn float_asin(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(ArcSinOps, B::float_asin); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::ArcSin(desc.clone())), + ArcSinOps::::new(desc), + ) + .output() + } + + fn float_asinh(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(ArcSinhOps, B::float_asinh); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::ArcSinh(desc.clone())), + ArcSinhOps::::new(desc), + ) + .output() + } + + fn float_atan(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(ArcTanOps, B::float_atan); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::ArcTan(desc.clone())), + ArcTanOps::::new(desc), + ) + .output() + } + + fn float_atanh(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(ArcTanhOps, B::float_atanh); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::ArcTanh(desc.clone())), + ArcTanhOps::::new(desc), + ) + .output() + } + + fn float_atan2(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float_ops!(ArcTan2Ops, B::float_atan2); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::ArcTan2(desc.clone())), + ArcTan2Ops::::new(desc), + ) + .output() + } + + fn float_recip(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(Recip, B::float_recip); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Recip(desc.clone())), + Recip::::new(desc), + ) + .output() + } + + fn float_erf(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(TanhOps, B::float_erf); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Erf(desc.clone())), + TanhOps::::new(desc), + ) + .output() + } + + fn float_cat(tensors: Vec>, dim: usize) -> FloatTensor { + #[derive(new, Debug)] + struct CatOps { + desc: CatOpIr, + _b: PhantomData, + } + + impl Operation for CatOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensors = self + .desc + .tensors + .iter() + .map(|tensor| handles.get_float_tensor::(tensor)) + .collect(); + + let output = B::float_cat(tensors, self.desc.dim); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensors.first().unwrap().client.clone(); + let tensors = tensors.into_iter().map(|t| t.into_ir()).collect(); + let desc = CatOpIr::create(tensors, dim, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Cat(desc.clone())), + CatOps::::new(desc), + ) + .output() + } + + fn float_argmax(tensor: FloatTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + reduce_float2int_ops!(ArgMaxOps, |input, axis, _, dtype| B::float_argmax( + input, axis, dtype + )); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create_arg(tensor.into_ir(), dim, 1, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.input.dtype, + NumericOperationIr::ArgMax(desc.clone()), + ), + ArgMaxOps::::new(desc), + ) + .output() + } + + fn float_argtopk( + tensor: FloatTensor, + dim: usize, + k: usize, + out_dtype: IntDType, + ) -> IntTensor { + reduce_float2int_ops!(ArgTopKOps, B::float_argtopk); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create_arg(tensor.into_ir(), dim, k, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.input.dtype, + NumericOperationIr::ArgTopK(desc.clone()), + ), + ArgTopKOps::::new(desc), + ) + .output() + } + + fn float_topk(tensor: FloatTensor, dim: usize, k: usize) -> FloatTensor { + reduce_float_ops!(TopKOps, B::float_topk); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, k, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::TopK(desc.clone())), + TopKOps::::new(desc), + ) + .output() + } + + fn float_repeat_dim(tensor: FloatTensor, dim: usize, times: usize) -> FloatTensor { + #[derive(new, Debug)] + struct RepeatDimOps { + desc: RepeatDimOpIr, + _b: PhantomData, + } + + impl Operation for RepeatDimOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_float_tensor::(&self.desc.tensor); + + let output = B::float_repeat_dim(tensor, self.desc.dim, self.desc.times); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = RepeatDimOpIr::create(tensor.into_ir(), dim, times, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::RepeatDim(desc.clone())), + RepeatDimOps::::new(desc), + ) + .output() + } + + fn float_argmin(tensor: FloatTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + reduce_float2int_ops!(ArgMinOps, |input, axis, _, dtype| B::float_argmin( + input, axis, dtype + )); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create_arg(tensor.into_ir(), dim, 1, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.input.dtype, + NumericOperationIr::ArgMin(desc.clone()), + ), + ArgMinOps::::new(desc), + ) + .output() + } + + fn float_max(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(MaxOps, B::float_max, reduce); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::Max(desc.clone())), + MaxOps::::new(desc.into()), + ) + .output() + } + + fn float_max_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + reduce_float_ops!(MaxDimOps, |tensor, axis, _| B::float_max_dim(tensor, axis)); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::MaxDim(desc.clone())), + MaxDimOps::::new(desc), + ) + .output() + } + + fn float_max_dim_with_indices( + tensor: FloatTensor, + dim: usize, + indices_dtype: IntDType, + ) -> (FloatTensor, IntTensor) { + #[derive(new, Debug)] + struct MaxDimWithIndicesOps { + desc: ReduceDimWithIndicesOpIr, + _b: PhantomData, + } + + impl Operation for MaxDimWithIndicesOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_float_tensor::(&self.desc.tensor); + let (output, indices) = B::float_max_dim_with_indices( + tensor, + self.desc.dim, + self.desc.out_indices.dtype.into(), + ); + + handles.register_float_tensor::(&self.desc.out.id, output); + handles.register_int_tensor::(&self.desc.out_indices.id, indices); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = + ReduceDimWithIndicesOpIr::create(tensor.into_ir(), dim, indices_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.tensor.dtype, + NumericOperationIr::MaxDimWithIndices(desc.clone()), + ), + MaxDimWithIndicesOps::::new(desc), + ) + .outputs() + .into() + } + + fn float_min(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(MinOps, B::float_min, reduce); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::Min(desc.clone())), + MinOps::::new(desc.into()), + ) + .output() + } + + fn float_min_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + reduce_float_ops!(MinDimOps, |tensor, axis, _| B::float_min_dim(tensor, axis)); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::MinDim(desc.clone())), + MinDimOps::::new(desc), + ) + .output() + } + + fn float_min_dim_with_indices( + tensor: FloatTensor, + dim: usize, + indices_dtype: IntDType, + ) -> (FloatTensor, IntTensor) { + #[derive(new, Debug)] + struct MinDimWithIndicesOps { + desc: ReduceDimWithIndicesOpIr, + _b: PhantomData, + } + + impl Operation for MinDimWithIndicesOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_float_tensor::(&self.desc.tensor); + let (output, indices) = B::float_min_dim_with_indices( + tensor, + self.desc.dim, + self.desc.out_indices.dtype.into(), + ); + + handles.register_float_tensor::(&self.desc.out.id, output); + handles.register_int_tensor::(&self.desc.out_indices.id, indices); + } + } + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = + ReduceDimWithIndicesOpIr::create(tensor.into_ir(), dim, indices_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.tensor.dtype, + NumericOperationIr::MinDimWithIndices(desc.clone()), + ), + MinDimWithIndicesOps::::new(desc), + ) + .outputs() + .into() + } + + fn float_max_abs(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(MaxAbsOps, B::float_max_abs, reduce); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat(desc.out.dtype, NumericOperationIr::MaxAbs(desc.clone())), + MaxAbsOps::::new(desc.into()), + ) + .output() + } + + fn float_max_abs_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + reduce_float_ops!(MaxAbsDimOps, |tensor, axis, _| B::float_max_abs_dim( + tensor, axis + )); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::MaxAbsDim(desc.clone()), + ), + MaxAbsDimOps::::new(desc), + ) + .output() + } + + // TODO: float_powi w/ burn-cubecl-fusion impl + fn float_powf(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float_ops!(PowOps, B::float_powf); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Powf(desc.clone())), + PowOps::::new(desc), + ) + .output() + } + + fn float_permute(tensor: FloatTensor, axes: &[usize]) -> FloatTensor { + #[derive(new, Debug)] + struct PermuteDimsOps { + desc: PermuteOpIr, + _b: PhantomData, + } + + impl Operation for PermuteDimsOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = B::float_permute(input, self.desc.axes.as_slice()); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = PermuteOpIr::create(tensor.into_ir(), axes.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::Permute(desc.clone())), + PermuteDimsOps::::new(desc), + ) + .output() + } + + fn float_expand(tensor: FloatTensor, shape: Shape) -> FloatTensor { + #[derive(new, Debug)] + struct ExpandOps { + desc: ShapeOpIr, + _b: PhantomData, + } + + impl Operation for ExpandOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = B::float_expand(input, self.desc.out.shape.clone()); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = ShapeOpIr::expand(tensor.into_ir(), shape, || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Expand(desc.clone())), + ExpandOps::::new(desc), + ) + .output() + } + + fn float_flip(tensor: FloatTensor, axes: &[usize]) -> FloatTensor { + #[derive(new, Debug)] + struct FlipOps { + desc: FlipOpIr, + _b: PhantomData, + } + + impl Operation for FlipOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = B::float_flip(input, &self.desc.axes); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = FlipOpIr::create(tensor.into_ir(), axes.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseInt(BaseOperationIr::Flip(desc.clone())), + FlipOps::::new(desc), + ) + .output() + } + + fn float_round(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(RoundOps, B::float_round); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Round(desc.clone())), + RoundOps::::new(desc), + ) + .output() + } + + fn float_floor(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(FloorOps, B::float_floor); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Floor(desc.clone())), + FloorOps::::new(desc), + ) + .output() + } + + fn float_ceil(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(CeilOps, B::float_ceil); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Ceil(desc.clone())), + CeilOps::::new(desc), + ) + .output() + } + + fn float_trunc(tensor: FloatTensor) -> FloatTensor { + unary_float_ops!(TruncOps, B::float_trunc); + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Trunc(desc.clone())), + TruncOps::::new(desc), + ) + .output() + } + + fn float_cast(tensor: FloatTensor, dtype: burn_backend::FloatDType) -> FloatTensor { + #[derive(new, Debug)] + struct CastOps { + desc: CastOpIr, + dtype: burn_backend::FloatDType, + _b: PhantomData, + } + + impl Operation for CastOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output: B::FloatTensorPrimitive = B::float_cast(input, self.dtype); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = CastOpIr::create(tensor.into_ir(), dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Cast(desc.clone())), + CastOps::::new(desc, dtype), + ) + .output() + } + + fn float_unfold( + tensor: FloatTensor, + dim: usize, + size: usize, + step: usize, + ) -> FloatTensor { + #[derive(new, Debug)] + struct UnfoldOps { + desc: UnfoldOpIr, + _b: PhantomData, + } + + impl Operation for UnfoldOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = B::float_unfold(input, self.desc.dim, self.desc.size, self.desc.step); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnfoldOpIr::create(tensor.into_ir(), dim, size, step, || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::BaseFloat(BaseOperationIr::Unfold(desc.clone())), + UnfoldOps::::new(desc), + ) + .output() + } + + fn float_is_nan(tensor: FloatTensor, out_dtype: BoolDType) -> BoolTensor { + #[derive(new, Debug)] + struct IsNanOps { + desc: UnaryOpIr, + _b: PhantomData, + } + impl Operation for IsNanOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = B::float_is_nan(input, self.desc.out.dtype.into()); + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create_comparison(tensor.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Float(desc.input.dtype, FloatOperationIr::IsNan(desc.clone())), + IsNanOps::::new(desc), + ) + .output() + } + + fn float_is_inf(tensor: FloatTensor, out_dtype: BoolDType) -> BoolTensor { + #[derive(new, Debug)] + struct IsInfOps { + desc: UnaryOpIr, + _b: PhantomData, + } + impl Operation for IsInfOps { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = B::float_is_inf(input, self.desc.out.dtype.into()); + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = UnaryOpIr::create_comparison(tensor.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Float(desc.input.dtype, FloatOperationIr::IsInf(desc.clone())), + IsInfOps::::new(desc), + ) + .output() + } + + fn float_grid_sample_2d( + tensor: FloatTensor, + grid: FloatTensor, + options: GridSampleOptions, + ) -> FloatTensor { + #[derive(new, Debug)] + struct GridSample2dOps { + desc: GridSample2dOpIr, + _b: PhantomData, + } + + impl Operation for GridSample2dOps { + fn execute(&self, handles: &mut HandleContainer) { + let tensor = handles.get_float_tensor::(&self.desc.tensor); + let grid = handles.get_float_tensor::(&self.desc.grid); + let output = + B::float_grid_sample_2d(tensor, grid, self.desc.options.clone().into()); + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + + let streams = StreamId::current(); + + let client = tensor.client.clone(); + let desc = + GridSample2dOpIr::create(tensor.into_ir(), grid.into_ir(), options.into(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::GridSample2d(desc.clone())), + GridSample2dOps::::new(desc), + ) + .output() + } + + fn float_hypot(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + binary_float_ops!(HypotOps, B::float_hypot); + + let streams = StreamId::current(); + + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register( + streams, + OperationIr::Float(desc.out.dtype, FloatOperationIr::Hypot(desc.clone())), + HypotOps::::new(desc), + ) + .output() + } +} diff --git a/crates/burn-fusion/src/ops/transaction.rs b/crates/burn-fusion/src/ops/transaction.rs new file mode 100644 index 0000000..186a1d8 --- /dev/null +++ b/crates/burn-fusion/src/ops/transaction.rs @@ -0,0 +1,36 @@ +use burn_backend::{ + backend::ExecutionError, + ops::{TransactionOps, TransactionPrimitive}, +}; + +use crate::{Fusion, FusionBackend}; + +impl TransactionOps> for Fusion { + async fn tr_execute( + transaction: TransactionPrimitive, + ) -> Result { + B::tr_execute(TransactionPrimitive::new( + transaction + .read_floats + .into_iter() + .map(|t| t.client.clone().resolve_tensor_float::(t)) + .collect(), + transaction + .read_qfloats + .into_iter() + .map(|_t| todo!("Quantization not supported yet")) + .collect(), + transaction + .read_ints + .into_iter() + .map(|t| t.client.clone().resolve_tensor_int::(t)) + .collect(), + transaction + .read_bools + .into_iter() + .map(|t| t.client.clone().resolve_tensor_bool::(t)) + .collect(), + )) + .await + } +} diff --git a/crates/burn-fusion/src/ops/unary.rs b/crates/burn-fusion/src/ops/unary.rs new file mode 100644 index 0000000..da3bd19 --- /dev/null +++ b/crates/burn-fusion/src/ops/unary.rs @@ -0,0 +1,324 @@ +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! scalar_float_ops { + ( + $name:ident, + $ops:expr + ) => { + #[derive(new, Debug)] + struct $name { + desc: ScalarOpIr, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = handles.get_float_tensor::(&self.desc.lhs); + let output = $ops(lhs, self.desc.rhs.into()); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + }; + ( + $name:ident, + $ops:expr, + noconvert + ) => { + #[derive(new, Debug)] + struct $name { + desc: ScalarOpIr, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = handles.get_float_tensor::(&self.desc.lhs); + let output = $ops(lhs, self.desc.rhs); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + }; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! reduce_float_ops { + ( + $name:ident, + $ops:expr + ) => { + #[derive(new, Debug)] + struct $name { + desc: ReduceDimOpIr, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = $ops(input, self.desc.axis, self.desc.accumulator_len); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + }; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! reduce_float2int_ops { + ( + $name:ident, + $ops:expr + ) => { + #[derive(new, Debug)] + struct $name { + desc: ReduceDimOpIr, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = $ops( + input, + self.desc.axis, + self.desc.accumulator_len, + self.desc.out.dtype.into(), + ); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + }; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! reduce_int_ops { + ( + $name:ident, + $ops:expr + ) => { + #[derive(new, Debug)] + struct $name { + desc: ReduceDimOpIr, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_int_tensor::(&self.desc.input); + let output = $ops(input, self.desc.axis, self.desc.accumulator_len); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + }; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! scalar_float2int_ops { + ( + $name:ident, + $ops:expr, + ) => { + #[derive(new, Debug)] + struct $name { + desc: ScalarOpIr, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = handles.get_float_tensor::(&self.desc.lhs); + let output = $ops(lhs, self.desc.rhs.clone()); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + }; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! unary_float_ops { + ( + $name:ident, + $ops:expr + ) => { + #[derive(new, Debug)] + struct $name { + desc: UnaryOpIr, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = $ops(input); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + }; + ( + $name:ident, + $ops:expr, + reduce + ) => { + #[derive(new, Debug)] + struct $name { + desc: UnaryOpIr, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_float_tensor::(&self.desc.input); + let output = $ops(input); + + handles.register_float_tensor::(&self.desc.out.id, output); + } + } + }; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! unary_int_ops { + ( + $name:ident, + $ops:expr + ) => { + #[derive(new, Debug)] + struct $name { + desc: UnaryOpIr, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_int_tensor::(&self.desc.input); + let output = $ops(input); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + }; + ( + $name:ident, + $ops:expr, + reduce + ) => { + #[derive(new, Debug)] + struct $name { + desc: UnaryOpIr, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let input = handles.get_int_tensor::(&self.desc.input); + let output = $ops(input); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + }; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! scalar_float_cmp_ops { + ( + $name:ident, + $ops:expr + ) => { + #[derive(new, Debug)] + struct $name { + desc: ScalarOpIr, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = handles.get_float_tensor::(&self.desc.lhs); + let output = $ops(lhs, self.desc.rhs.into(), self.desc.out.dtype.into()); + + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + }; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! scalar_int_cmp_ops { + ( + $name:ident, + $ops:expr + ) => { + #[derive(new, Debug)] + struct $name { + desc: ScalarOpIr, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = handles.get_int_tensor::(&self.desc.lhs); + let output = $ops(lhs, self.desc.rhs.into(), self.desc.out.dtype.into()); + + handles.register_bool_tensor::(&self.desc.out.id, output); + } + } + }; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! scalar_int_ops { + ( + $name:ident, + $ops:expr + ) => { + #[derive(new, Debug)] + struct $name { + desc: ScalarOpIr, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = handles.get_int_tensor::(&self.desc.lhs); + let output = $ops(lhs, self.desc.rhs.into()); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + }; + ( + $name:ident, + $ops:expr, + noconvert + ) => { + #[derive(new, Debug)] + struct $name { + desc: ScalarOpIr, + _b: PhantomData, + } + + impl Operation for $name { + fn execute(&self, handles: &mut HandleContainer) { + let lhs = handles.get_int_tensor::(&self.desc.lhs); + let output = $ops(lhs, self.desc.rhs); + + handles.register_int_tensor::(&self.desc.out.id, output); + } + } + }; +} diff --git a/crates/burn-fusion/src/search/block.rs b/crates/burn-fusion/src/search/block.rs new file mode 100644 index 0000000..a09cfbd --- /dev/null +++ b/crates/burn-fusion/src/search/block.rs @@ -0,0 +1,639 @@ +use crate::{ + FuserStatus, NumOperations, OperationFuser, + search::graph::{GraphNode, SubGraph, is_valid_execution_order}, + stream::store::ExecutionStrategy, +}; +use burn_ir::{OperationIr, TensorId, TensorIr, TensorStatus}; +use std::{collections::HashSet, sync::Arc}; + +/// A block represents a list of operations, not necessarily in the same order as the execution +/// stream. +/// +/// The start and end position of the relative execution stream are tracked in the block alongside +/// the ordering. +pub struct Block { + builders: Vec>>, + operations: Vec, + ids: HashSet, + /// Tensor ids produced (as an output) by an operation of this block. + produced: HashSet, + /// Tensor ids consumed by this block but produced elsewhere (external inputs). + read: HashSet, + /// Tensor ids this block reads with [ReadWrite](TensorStatus::ReadWrite) status — i.e. it is + /// the last use and frees/reuses the buffer in place. + freed: HashSet, + /// The original blocks this block subsumes. + /// + /// Seeded to a single node by the optimizer before a merge pass and unioned on + /// [merge](Self::merge), so the [Reachability](crate::search::graph::Reachability) guard can + /// reason about the original dependency graph through the recursive merging. + constituents: SubGraph, + ordering: Vec, + /// The start position in the relative execution stream. + pub start_pos: usize, + /// The end position in the relative execution stream. + pub end_pos: usize, +} + +/// The result of [registering](Block::register) an [operation](OperationIr). +pub enum RegistrationResult { + /// If the [operation](OperationIr) is correctly registered. + Accepted, + /// If the [operation](OperationIr) isn't part of the graph. + /// + /// In this case the operation isn't registered. + NotPartOfTheGraph, +} + +/// The optimization found for a [block](Block). +#[derive(Debug, new)] +pub struct BlockOptimization { + /// The [execution strategy](ExecutionStrategy) to be used to execute the [block](Block). + pub strategy: ExecutionStrategy, + /// The ordering of each operation in the relative execution stream. + pub ordering: Vec, +} + +impl Block { + /// The original blocks this block subsumes. + pub fn constituents(&self) -> &SubGraph { + &self.constituents + } + + /// Seed this block's [constituents](Self::constituents) to a single original block index. + pub fn seed_constituent(&mut self, index: usize) { + self.constituents = SubGraph::single(index); + } +} + +/// A single operation at its stream position, viewed as a [GraphNode]. +pub struct OperationNode<'a> { + /// The operation. + pub operation: &'a OperationIr, + /// The stream position of the operation. + pub position: usize, +} + +impl GraphNode for OperationNode<'_> { + type Resource = TensorId; + + fn produced(&self) -> impl Iterator { + self.operation.outputs().map(|tensor| tensor.id) + } + + fn read(&self) -> impl Iterator { + self.operation.inputs().map(|tensor| tensor.id) + } + + fn freed(&self) -> impl Iterator { + self.operation + .inputs() + .filter(|tensor| matches!(tensor.status, TensorStatus::ReadWrite)) + .map(|tensor| tensor.id) + } + + fn position(&self) -> usize { + self.position + } +} + +/// The dependency edges between blocks are derived from what each block reads, produces, and +/// frees — see [GraphNode]. +impl GraphNode for Block { + type Resource = TensorId; + + fn produced(&self) -> impl Iterator { + self.produced.iter().copied() + } + + fn read(&self) -> impl Iterator { + self.read.iter().copied() + } + + fn freed(&self) -> impl Iterator { + self.freed.iter().copied() + } + + fn produces(&self, resource: TensorId) -> bool { + self.produced.contains(&resource) + } + + fn reads(&self, resource: TensorId) -> bool { + self.read.contains(&resource) + } + + fn position(&self) -> usize { + self.start_pos + } +} + +impl Block { + /// Create a new block that will be optimized with the provided [optimization builders](OptimizationBuilder). + pub fn new(builders: &[Box>]) -> Self { + Self { + builders: builders.iter().map(|o| o.clone_dyn()).collect(), + operations: Vec::new(), + ids: HashSet::new(), + produced: HashSet::new(), + read: HashSet::new(), + freed: HashSet::new(), + constituents: SubGraph::empty(), + ordering: Vec::new(), + start_pos: usize::MAX, + end_pos: usize::MIN, + } + } + + /// Sort the [blocks](Block) based on the start position. + pub fn sort(blocks: &mut [Self]) { + blocks.sort_by_key(|a| a.start_pos); + } + + /// Optimize the block. + pub fn optimize(mut self) -> BlockOptimization { + match find_best_optimization_index(&mut self.builders) { + BestOptimization::Found { index, score } => { + let opt = self.builders[index].finish(); + let opt_len = opt.len(); + if opt_len < self.operations.len() { + self.ordering.drain(opt_len..); + } + + let strategy = ExecutionStrategy::Optimization { + ordering: Arc::new(self.ordering.clone()), + opt, + score, + }; + BlockOptimization::new(strategy, self.ordering) + } + BestOptimization::NotFound => { + let strategy = ExecutionStrategy::Operations { + ordering: Arc::new(self.ordering.clone()), + }; + BlockOptimization::new(strategy, self.ordering) + } + } + } + + /// Returns if the block contains any of the provided [tensors](TensorIr). + pub fn contains_tensors(&self, tensors: &[&TensorIr]) -> bool { + for node in tensors { + if self.ids.contains(&node.id) { + return true; + } + } + + false + } + + /// Merge the current block with the other one and returns if the operation is successful. + /// + /// # Warning + /// + /// This will modify the current block even if the other block isn't correctly merged. + pub fn merge(&mut self, other: &Block) -> bool { + // A block executes as one contiguous unit in registration order (fused kernels replay + // the fusion order). Appending the other block's operations can place a consumer before + // its producer — or a free before a read — when the blocks depend on each other. Reject + // such merges before mutating anything; the caller can retry in the other direction. + if !self.can_append(other) { + return false; + } + + let self_ready = self.has_ready_optimization(); + let other_ready = other.has_ready_optimization(); + + // Absorb the other block's provenance so the merge guard sees the combined set on any + // subsequent `can_contract` check. + self.constituents.union_with(&other.constituents); + + for (op, pos) in other.operations.iter().zip(&other.ordering) { + self.register(op, *pos, true); + } + + // If the merged block can still be improved, keep it — that's the + // usual lazy-optimization signal. + if self.still_optimizing() { + return true; + } + + // Otherwise, only accept the merge when it *creates* a ready fusion + // that didn't exist separately. If either side already had a ready + // fusion before the merge, merging would collapse them and hide one + // of the fusions — keep the blocks separate instead. + !self_ready && !other_ready && self.has_ready_optimization() + } + + fn has_ready_optimization(&self) -> bool { + self.builders.iter().any(|b| b.properties().ready) + } + + /// Whether appending the other block's operations after this block's — the registration + /// order a [merge](Self::merge) in this direction produces — respects every tensor lifetime + /// (no read before the producing operation, no read after the freeing operation). + /// + /// This is the validity check `merge` performs, exposed so callers can test a merge + /// direction *before* paying for a deep clone of the base block. + pub fn can_append<'a>(&'a self, other: &'a Block) -> bool { + is_valid_execution_order(self.operation_nodes().chain(other.operation_nodes())) + } + + /// The block's operations in registration order, viewed as [graph nodes](OperationNode). + fn operation_nodes(&self) -> impl Iterator> + Clone { + self.operations + .iter() + .zip(&self.ordering) + .map(|(operation, &position)| OperationNode { + operation, + position, + }) + } + + /// Register an [operation](OperationIr) in the current block. + /// + /// You need to provide the order of the operation as well as a force flag. + /// + /// When the force flag is true, the builder will always accept the operation, otherwise it + /// might refuse it if the operation [isn't part of the graph](RegistrationResult::NotPartOfTheGraph). + /// + /// Forcing is useful to fuse operations that are part of different graphs, but included + /// in the same optimization. + pub fn register( + &mut self, + operation: &OperationIr, + order: usize, + force: bool, + ) -> RegistrationResult { + if self.ids.is_empty() { + self.register_op(operation, order); + return RegistrationResult::Accepted; + } + let mut contains = false; + for node in operation.nodes() { + contains = self.ids.contains(&node.id); + + if contains { + break; + } + } + + if !contains && !force { + return RegistrationResult::NotPartOfTheGraph; + } + + self.register_op(operation, order); + RegistrationResult::Accepted + } + + /// If the block can still be optimized further. + pub fn still_optimizing(&self) -> bool { + let mut num_stopped = 0; + + for optimization in self.builders.iter() { + if let FuserStatus::Closed = optimization.status() { + num_stopped += 1 + } + } + + num_stopped < self.builders.len() + } + + fn register_op(&mut self, operation: &OperationIr, pos: usize) { + self.operations.push(operation.clone()); + self.ordering.push(pos); + + if pos < self.start_pos { + self.start_pos = pos; + } + if pos + 1 > self.end_pos { + self.end_pos = pos + 1; + } + + for builder in self.builders.iter_mut() { + builder.fuse(operation); + } + + for node in operation.nodes() { + self.ids.insert(node.id); + } + + // Maintain the produced / external-read sets incrementally. A consumed id counts as an + // external read only while nothing in the block has produced it; producing an id makes + // it internal (and clears any earlier external record). This is order-independent, so it + // stays correct when `merge` folds ops in a non-causal order. + for node in operation.inputs() { + if !self.produced.contains(&node.id) { + self.read.insert(node.id); + } + if let TensorStatus::ReadWrite = node.status { + self.freed.insert(node.id); + } + } + for node in operation.outputs() { + self.produced.insert(node.id); + self.read.remove(&node.id); + } + } +} + +impl BlockOptimization { + /// Maps the ordering of the current block optimization using the given mapping. + pub fn map_ordering(&mut self, mapping: &[usize]) { + for i in self.ordering.iter_mut() { + *i = mapping[*i]; + } + self.strategy.map_ordering(mapping); + } + + /// Extend the optimization to cover all `num_operations` of its segment: the + /// positions the search left uncovered (the drained tail of a block whose best + /// fusion stopped early) are appended as un-fused + /// [operations](ExecutionStrategy::Operations) in a + /// [composed](ExecutionStrategy::Composed) strategy. + /// + /// Leaving the tail out is right in lazy mode — it seeds the next exploration + /// round, where more incoming operations may fuse with it. At a sync the segment + /// is final: no later round can pick the tail up, and a plan covering only part + /// of the segment can never be matched again (plan lookup at a sync compares the + /// whole remaining segment). Covering the full segment makes the cached plan + /// reusable on every subsequent identical sync. + /// + /// The uncovered positions are always a trailing run of the segment (interior + /// holes are re-optimized by the stream optimizer before this is called — enforced + /// by a debug assertion), so appending them last preserves the hazard-respecting + /// execution order. + pub fn include_trailing(&mut self, num_operations: usize) { + if self.ordering.len() >= num_operations { + return; + } + + #[cfg(debug_assertions)] + { + let mut resolved = vec![false; num_operations]; + for &position in &self.ordering { + resolved[position] = true; + } + debug_assert!( + resolved[..self.ordering.len()].iter().all(|&r| r), + "uncovered positions must be a trailing run of the segment" + ); + } + let trailing: Vec = (self.ordering.len()..num_operations).collect(); + + let tail = ExecutionStrategy::Operations { + ordering: Arc::new(trailing.clone()), + }; + let strategy = core::mem::replace(&mut self.strategy, ExecutionStrategy::Composed(vec![])); + self.strategy = match strategy { + ExecutionStrategy::Composed(mut items) => { + items.push(Box::new(tail)); + ExecutionStrategy::Composed(items) + } + single => ExecutionStrategy::Composed(vec![Box::new(single), Box::new(tail)]), + }; + self.ordering.extend(trailing); + } +} + +impl ExecutionStrategy { + /// Maps the ordering of the current execution strategy using the given mapping. + pub fn map_ordering(&mut self, mapping: &[usize]) { + match self { + ExecutionStrategy::Optimization { ordering, .. } => { + let mut ordering_mapped = ordering.to_vec(); + + for o in ordering_mapped.iter_mut() { + *o = mapping[*o]; + } + *ordering = Arc::new(ordering_mapped); + } + ExecutionStrategy::Operations { ordering } => { + let mut ordering_mapped = ordering.to_vec(); + + for o in ordering_mapped.iter_mut() { + *o = mapping[*o]; + } + + *ordering = Arc::new(ordering_mapped); + } + ExecutionStrategy::Composed(items) => { + for item in items.iter_mut() { + item.map_ordering(mapping); + } + } + } + } +} + +enum BestOptimization { + NotFound, + Found { index: usize, score: u64 }, +} + +fn find_best_optimization_index( + optimizations: &mut [Box>], +) -> BestOptimization { + let mut best_index = BestOptimization::NotFound; + let mut best_score = 0; + + for (i, optimization) in optimizations.iter().enumerate() { + let properties = optimization.properties(); + + // A score of zero is worse than fusing. + if properties.ready && properties.score > best_score { + best_index = BestOptimization::Found { + index: i, + score: properties.score, + }; + best_score = properties.score; + } + } + + best_index +} + +impl PartialEq for Block { + fn eq(&self, other: &Self) -> bool { + // Since the ordering can be seen as operation ids, we can use it to compare + // blocks. + let mut sorted_a = self.ordering.clone(); + let mut sorted_b = other.ordering.clone(); + sorted_a.sort(); + sorted_b.sort(); + + sorted_a == sorted_b + } +} + +impl core::fmt::Debug for Block { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_fmt(format_args!( + "Block {{ pos: [{:?}, {:?}; {:?}] }}", + self.start_pos, + self.end_pos, + self.ordering.len(), + )) + } +} + +impl Clone for Block { + fn clone(&self) -> Self { + Self { + builders: self.builders.iter().map(|b| b.clone_dyn()).collect(), + operations: self.operations.clone(), + ids: self.ids.clone(), + produced: self.produced.clone(), + read: self.read.clone(), + freed: self.freed.clone(), + constituents: self.constituents.clone(), + ordering: self.ordering.clone(), + start_pos: self.start_pos, + end_pos: self.end_pos, + } + } +} + +/// Integration of [Block] with the [graph](crate::search::graph) algorithms: dependency edges +/// between blocks are derived from the tensors they read, produce, and free. +#[cfg(test)] +mod tests { + use super::*; + use crate::search::graph::Dag; + use crate::search::testing::{add, add_rw}; + use crate::stream::execution::tests::TestOptimization; + + /// A block owning the given `(operation, position)` pairs. No builders needed — the + /// dependency analysis only reads the data-flow sets and `start_pos`. + fn block(ops: &[(OperationIr, usize)]) -> Block { + let builders: Vec>> = Vec::new(); + let mut block = Block::new(&builders); + for (op, pos) in ops { + block.register(op, *pos, true); + } + block + } + + #[test] + fn dependency_beats_start_pos_in_topological_order() { + // Block 0 (start_pos 0) consumes tensor 50, which block 1 (start_pos 5) produces. + let blocks = vec![ + block(&[(add(50, 1, 2), 0)]), // Consumes 50 -> depends on block 1. + block(&[(add(9, 8, 50), 5)]), // Produces 50. + ]; + + assert_eq!(Dag::new(&blocks).topological_order(), Some(vec![1, 0])); + } + + #[test] + fn mutual_dependency_between_blocks_is_a_cycle() { + // Block 0 produces 100 (consumed by block 1) and consumes 200 (produced by block 1). + let blocks = vec![ + block(&[(add(1, 2, 100), 0), (add(200, 3, 101), 1)]), + block(&[(add(100, 4, 200), 2)]), + ]; + + assert!(!Dag::new(&blocks).is_acyclic()); + } + + #[test] + fn freeing_block_runs_after_reading_block() { + // Both blocks read tensor 100 (produced elsewhere): block 0 read-only, block 1 frees it + // (ReadWrite). The freer must run after the reader. + let blocks = vec![ + block(&[(add(100, 1, 2), 0)]), // Reads 100 read-only. + block(&[(add_rw(100, 3, 4), 1)]), // Frees 100. + ]; + + assert_eq!(Dag::new(&blocks).topological_order(), Some(vec![0, 1])); + } + + #[test] + fn merge_guard_blocks_contraction_around_intermediate_block() { + // A(0) -> C(1) -> B(2): C sits between A and B, so merging A and B would create a cycle. + let mut blocks = vec![ + block(&[(add(1, 2, 100), 0)]), // A: produces 100. + block(&[(add(100, 3, 200), 1)]), // C: consumes 100, produces 200. + block(&[(add(200, 4, 300), 2)]), // B: consumes 200. + ]; + for (i, block) in blocks.iter_mut().enumerate() { + block.seed_constituent(i); + } + let guard = Dag::new(&blocks).reachability(); + + // A and B cannot merge (C is between them)... + assert!(!guard.can_contract(blocks[0].constituents(), blocks[2].constituents())); + // ...but each direct edge is fine. + assert!(guard.can_contract(blocks[0].constituents(), blocks[1].constituents())); + assert!(guard.can_contract(blocks[1].constituents(), blocks[2].constituents())); + } + + #[test] + fn merge_guard_refuses_unseeded_blocks() { + let blocks = vec![block(&[(add(1, 2, 3), 0)]), block(&[(add(4, 5, 6), 1)])]; + let guard = Dag::new(&blocks).reachability(); + + assert!(!guard.can_contract(blocks[0].constituents(), blocks[1].constituents())); + } + + /// A block with a builder open to any prefix of `pattern`, owning the given ops. + /// + /// The builder keeps the block `still_optimizing` as long as the registered ops follow the + /// pattern, so merge acceptance is driven by the pattern and the tests below can isolate the + /// internal-order validation. + fn block_with_pattern( + pattern: &[OperationIr], + ops: &[(OperationIr, usize)], + ) -> Block { + use crate::stream::execution::tests::TestOptimizationBuilder; + + let builders: Vec>> = + vec![Box::new(TestOptimizationBuilder::new(0, pattern.to_vec()))]; + let mut block = Block::new(&builders); + for (op, pos) in ops { + block.register(op, *pos, true); + } + block + } + + #[test] + fn merge_rejects_consumer_fused_before_producer() { + let x = add(1, 2, 10); + let y = add(50, 3, 11); // Reads 50... + let z = add(4, 5, 50); // ...which this op produces. + + // The pattern accepts the forward merge order [x, y, z], so only the internal-order + // validation can reject it: y would be fused before the op producing its input. + let pattern = [x.clone(), y.clone(), z.clone(), x.clone()]; + let mut base = block_with_pattern(&pattern, &[(x, 0), (y, 1)]); + let producer = block_with_pattern(&pattern, &[(z, 2)]); + + assert!(!base.merge(&producer)); + } + + #[test] + fn merge_accepts_producer_fused_before_consumer() { + let x = add(1, 2, 10); + let y = add(50, 3, 11); + let z = add(4, 5, 50); + + // Same blocks as above, merged in the other direction: [z, x, y] is causal. + let pattern = [z.clone(), x.clone(), y.clone(), x.clone()]; + let consumers = block_with_pattern(&pattern, &[(x, 0), (y, 1)]); + let mut base = block_with_pattern(&pattern, &[(z, 2)]); + + assert!(base.merge(&consumers)); + } + + #[test] + fn merge_rejects_free_fused_before_read() { + let reader = add(100, 1, 2); // Reads 100. + let freer = add_rw(100, 3, 4); // Frees 100. + + // Fusing the free before the read would release the tensor under the reader. + let pattern = [freer.clone(), reader.clone(), reader.clone()]; + let mut base = block_with_pattern(&pattern, &[(freer, 1)]); + let other = block_with_pattern(&pattern, &[(reader, 0)]); + + assert!(!base.merge(&other)); + } +} diff --git a/crates/burn-fusion/src/search/graph/dag.rs b/crates/burn-fusion/src/search/graph/dag.rs new file mode 100644 index 0000000..feccfd5 --- /dev/null +++ b/crates/burn-fusion/src/search/graph/dag.rs @@ -0,0 +1,190 @@ +use super::{GraphNode, SubGraph}; +use std::cmp::Reverse; +use std::collections::BinaryHeap; + +/// A dependency graph over a fixed list of nodes, with edges derived from the nodes' data-flow +/// sets (see [GraphNode]). +/// +/// The graph is not guaranteed to be acyclic — a cycle means the nodes cannot be linearized and +/// is reported by [topological_order](Self::topological_order) returning `None`. +pub struct Dag { + /// `dependencies[i]` = nodes that must execute before node `i` (direct edges only). + dependencies: Vec, + /// `dependents[j]` = nodes that directly depend on node `j` (transpose of `dependencies`). + dependents: Vec, + /// Tie-break position of each node (see [GraphNode::position]). + positions: Vec, +} + +impl Dag { + /// Build the dependency graph of the given nodes. + pub fn new(nodes: &[N]) -> Self { + let n = nodes.len(); + let mut dependencies = vec![SubGraph::empty(); n]; + let mut dependents = vec![SubGraph::empty(); n]; + + for (i, node) in nodes.iter().enumerate() { + for (j, other) in nodes.iter().enumerate() { + if i != j && depends_on(node, other) { + dependencies[i].insert(j); + dependents[j].insert(i); + } + } + } + + Self { + dependencies, + dependents, + positions: nodes.iter().map(|n| n.position()).collect(), + } + } + + /// The number of nodes in the graph. + pub fn len(&self) -> usize { + self.dependencies.len() + } + + /// Whether node `i` directly depends on node `j` (`j` must execute before `i`). + #[cfg(test)] + pub fn depends_on(&self, i: usize, j: usize) -> bool { + self.dependencies[i].contains(j) + } + + /// A valid execution order of the nodes (dependencies first), or `None` if the graph has a + /// cycle and cannot be linearized. + /// + /// Among the nodes ready at each step, the one with the smallest `(position, index)` is + /// emitted first, so an edge-free graph reproduces the original program order. + pub fn topological_order(&self) -> Option> { + let n = self.len(); + let mut pending: Vec = self.dependencies.iter().map(SubGraph::len).collect(); + + // Min-heap of the nodes ready to execute, keyed by `(position, index)`. + let mut ready: BinaryHeap> = (0..n) + .filter(|&i| pending[i] == 0) + .map(|i| Reverse((self.positions[i], i))) + .collect(); + + let mut order = Vec::with_capacity(n); + while let Some(Reverse((_, i))) = ready.pop() { + order.push(i); + for dependent in self.dependents[i].iter() { + pending[dependent] -= 1; + if pending[dependent] == 0 { + ready.push(Reverse((self.positions[dependent], dependent))); + } + } + } + + // A node still pending at the end sits on a cycle. + (order.len() == n).then_some(order) + } + + /// Whether the graph can be linearized. + pub fn is_acyclic(&self) -> bool { + self.topological_order().is_some() + } + + /// Compute the transitive [Reachability] of the graph. + pub fn reachability(&self) -> Reachability { + let n = self.len(); + let mut ancestors = self.dependencies.clone(); + + // A node absorbs the (final) ancestor sets of its direct dependencies. In execution + // order every dependency is processed before its dependents, so a single pass is + // complete. Taking the set out and putting it back sidesteps the aliasing between + // `ancestors[i]` and `ancestors[j]` without any clone (`j != i`: no self-edges). + match self.topological_order() { + Some(order) => { + for i in order { + let mut acc = core::mem::take(&mut ancestors[i]); + for j in self.dependencies[i].iter() { + acc.union_with(&ancestors[j]); + } + ancestors[i] = acc; + } + } + // On a cycle no linear pass exists (defensive — the merge guards keep the block set + // acyclic): relax the direct edges to a fixpoint instead. Unions only grow, so a + // length change is an exact change signal. + None => loop { + let mut changed = false; + for i in 0..n { + let mut acc = core::mem::take(&mut ancestors[i]); + let before = acc.len(); + for j in self.dependencies[i].iter() { + acc.union_with(&ancestors[j]); + } + changed |= acc.len() != before; + ancestors[i] = acc; + } + if !changed { + break; + } + }, + } + + // Descendants are the transpose of ancestors. + let mut descendants = vec![SubGraph::empty(); n]; + for (i, ancestors) in ancestors.iter().enumerate() { + for j in ancestors.iter() { + descendants[j].insert(i); + } + } + + Reachability { + ancestors, + descendants, + } + } +} + +/// Transitive reachability over a [Dag]. +/// +/// Answers whether contracting a set of nodes into a single node keeps the graph acyclic, which +/// is the legality condition for merging fusion blocks that may depend on each other. +pub struct Reachability { + /// `ancestors[i]` = nodes that must execute before node `i` (transitive). + ancestors: Vec, + /// `descendants[i]` = nodes that must execute after node `i` (transitive). + descendants: Vec, +} + +impl Reachability { + /// Whether contracting the two subgraphs into a single node keeps the graph acyclic. + /// + /// Contraction is illegal iff some node outside `a ∪ b` is both an ancestor and a descendant + /// of the union — it would have to execute both before and after the contracted node. + /// + /// Empty subgraphs are refused conservatively: they identify nodes the reachability was not + /// built from. + pub fn can_contract(&self, a: &SubGraph, b: &SubGraph) -> bool { + if a.is_empty() || b.is_empty() { + return false; + } + + let mut union = a.clone(); + union.union_with(b); + + let mut ancestors = SubGraph::empty(); + let mut descendants = SubGraph::empty(); + for i in union.iter() { + ancestors.union_with(&self.ancestors[i]); + descendants.union_with(&self.descendants[i]); + } + + let mut between = ancestors; + between.intersect_with(&descendants); + between.subtract(&union); + between.is_empty() + } +} + +/// Whether `node` must execute after `other`: read-after-write (`node` reads a resource `other` +/// produces) or write-after-read (`node` frees a resource `other` reads). +/// +/// Works directly on the nodes' own data-flow sets through the [GraphNode] membership queries — +/// no intermediate collection. +fn depends_on(node: &N, other: &N) -> bool { + node.read().any(|r| other.produces(r)) || node.freed().any(|r| other.reads(r)) +} diff --git a/crates/burn-fusion/src/search/graph/lifetime.rs b/crates/burn-fusion/src/search/graph/lifetime.rs new file mode 100644 index 0000000..fa734e4 --- /dev/null +++ b/crates/burn-fusion/src/search/graph/lifetime.rs @@ -0,0 +1,46 @@ +use super::GraphNode; +use std::collections::HashSet; + +/// Whether executing the nodes in the given order keeps every read resource live. +/// +/// The nodes are taken already in execution order — pass a (cloneable, cheap) iterator over the +/// candidate order rather than materializing it; the function walks it twice. +/// +/// Simulates resource lifetimes: every resource a node reads must be live when the node +/// executes, a freed resource dies after the freeing node, and produced resources become live. +/// A resource read before any ordered node produced it is external (from an earlier segment) +/// and assumed live — unless some ordered node *does* produce it, in which case the order reads +/// it before its producer and is invalid. +pub fn is_valid_execution_order(order: impl Iterator + Clone) -> bool { + let mut produced = HashSet::new(); + for node in order.clone() { + produced.extend(node.produced()); + } + + let mut alive = HashSet::new(); + let mut dead = HashSet::new(); + for node in order { + for resource in node.read() { + if dead.contains(&resource) { + return false; // Read after the resource was freed: bad order. + } + if !alive.contains(&resource) { + if produced.contains(&resource) { + return false; // Produced in this segment but not live yet: bad order. + } + alive.insert(resource); // External resource from a prior segment. + } + } + for resource in node.freed() { + alive.remove(&resource); + dead.insert(resource); + } + for resource in node.produced() { + // Producing redefines the resource, even if an id were ever reused after a free. + dead.remove(&resource); + alive.insert(resource); + } + } + + true +} diff --git a/crates/burn-fusion/src/search/graph/mod.rs b/crates/burn-fusion/src/search/graph/mod.rs new file mode 100644 index 0000000..475d196 --- /dev/null +++ b/crates/burn-fusion/src/search/graph/mod.rs @@ -0,0 +1,18 @@ +//! A small, self-contained dependency-graph toolkit used by the fusion search. +//! +//! Everything here is generic over the [GraphNode] trait — nothing depends on the tensor IR — so +//! the algorithms can be unit-tested in isolation and reused at different granularities: fusion +//! blocks, execution-strategy chunks, or single operations. + +mod dag; +mod lifetime; +mod node; +mod subgraph; + +pub use dag::*; +pub use lifetime::*; +pub use node::*; +pub use subgraph::*; + +#[cfg(test)] +mod tests; diff --git a/crates/burn-fusion/src/search/graph/node.rs b/crates/burn-fusion/src/search/graph/node.rs new file mode 100644 index 0000000..edb620e --- /dev/null +++ b/crates/burn-fusion/src/search/graph/node.rs @@ -0,0 +1,69 @@ +use core::hash::Hash; + +/// Data-flow description of a node in a dependency graph. +/// +/// A node reads and produces resources (tensors, buffers, …) identified by +/// [Resource](Self::Resource). Dependencies between nodes are never declared explicitly: they are +/// derived from the resource sets by [Dag::new](super::Dag::new) using two hazard rules. +/// +/// - **Read-after-write**: a node reading a resource depends on the node producing it. +/// - **Write-after-read**: a node freeing a resource depends on every other node reading it — +/// executing the freeing node first would release the resource out from under the readers. +pub trait GraphNode { + /// Identifies a resource read, produced, or freed by a node. + type Resource: Copy + Eq + Hash; + + /// The resources produced by this node. + fn produced(&self) -> impl Iterator; + + /// The resources this node reads that are produced by other nodes. + fn read(&self) -> impl Iterator; + + /// The resources this node reads for the last time, releasing the underlying storage + /// (in-place reuse or deallocation). + fn freed(&self) -> impl Iterator; + + /// Whether this node produces the resource — the membership form of + /// [produced](Self::produced). Nodes that already hold their resources in a set should + /// override this with a direct lookup. + fn produces(&self, resource: Self::Resource) -> bool { + self.produced().any(|r| r == resource) + } + + /// Whether this node reads the resource — the membership form of [read](Self::read). + fn reads(&self, resource: Self::Resource) -> bool { + self.read().any(|r| r == resource) + } + + /// The position of the node in the original program order, used to break ties between + /// independent nodes when ordering them. + fn position(&self) -> usize; +} + +impl GraphNode for &N { + type Resource = N::Resource; + + fn produced(&self) -> impl Iterator { + (*self).produced() + } + + fn read(&self) -> impl Iterator { + (*self).read() + } + + fn freed(&self) -> impl Iterator { + (*self).freed() + } + + fn produces(&self, resource: Self::Resource) -> bool { + (*self).produces(resource) + } + + fn reads(&self, resource: Self::Resource) -> bool { + (*self).reads(resource) + } + + fn position(&self) -> usize { + (*self).position() + } +} diff --git a/crates/burn-fusion/src/search/graph/subgraph.rs b/crates/burn-fusion/src/search/graph/subgraph.rs new file mode 100644 index 0000000..365be84 --- /dev/null +++ b/crates/burn-fusion/src/search/graph/subgraph.rs @@ -0,0 +1,125 @@ +/// A set of nodes of a [Dag](super::Dag), identified by their node index. +/// +/// The first 64 node indices live in an inline word, so subgraphs over small graphs — the common +/// case, given the block cap — never touch the heap. Larger graphs spill into a word vector kept +/// canonical (no trailing zero words), which makes the derived equality structural. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct SubGraph { + /// Nodes `0..64`. + word0: u64, + /// Nodes `64..`: word `i` holds nodes `(i + 1) * 64..(i + 2) * 64`. + spill: Vec, +} + +const BITS: usize = u64::BITS as usize; + +impl SubGraph { + /// The empty subgraph. + pub fn empty() -> Self { + Self::default() + } + + /// The subgraph containing a single node. + pub fn single(node: usize) -> Self { + let mut set = Self::empty(); + set.insert(node); + set + } + + /// Whether the subgraph contains no node. + pub fn is_empty(&self) -> bool { + self.word0 == 0 && self.spill.is_empty() + } + + /// The number of nodes in the subgraph. + pub fn len(&self) -> usize { + let spill: u32 = self.spill.iter().map(|word| word.count_ones()).sum(); + (self.word0.count_ones() + spill) as usize + } + + /// Add a node to the subgraph. + pub fn insert(&mut self, node: usize) { + if node < BITS { + self.word0 |= 1 << node; + return; + } + let word = node / BITS - 1; + if word >= self.spill.len() { + self.spill.resize(word + 1, 0); + } + self.spill[word] |= 1 << (node % BITS); + } + + /// Whether the subgraph contains the node. + #[cfg(test)] + pub fn contains(&self, node: usize) -> bool { + if node < BITS { + return self.word0 & (1 << node) != 0; + } + self.spill + .get(node / BITS - 1) + .is_some_and(|word| word & (1 << (node % BITS)) != 0) + } + + /// Add every node of the other subgraph. + pub fn union_with(&mut self, other: &Self) { + self.word0 |= other.word0; + if other.spill.len() > self.spill.len() { + self.spill.resize(other.spill.len(), 0); + } + for (word, other) in self.spill.iter_mut().zip(&other.spill) { + *word |= other; + } + } + + /// Keep only the nodes also contained in the other subgraph. + pub fn intersect_with(&mut self, other: &Self) { + self.word0 &= other.word0; + self.spill.truncate(other.spill.len()); + for (word, other) in self.spill.iter_mut().zip(&other.spill) { + *word &= other; + } + self.trim(); + } + + /// Remove every node of the other subgraph. + pub fn subtract(&mut self, other: &Self) { + self.word0 &= !other.word0; + for (word, other) in self.spill.iter_mut().zip(&other.spill) { + *word &= !other; + } + self.trim(); + } + + /// Whether the two subgraphs share at least one node. + #[cfg(test)] + pub fn intersects(&self, other: &Self) -> bool { + self.word0 & other.word0 != 0 + || self.spill.iter().zip(&other.spill).any(|(a, b)| a & b != 0) + } + + /// The nodes of the subgraph, in ascending index order. + pub fn iter(&self) -> impl Iterator + '_ { + core::iter::once(self.word0) + .chain(self.spill.iter().copied()) + .enumerate() + .flat_map(|(i, word)| { + let mut bits = word; + core::iter::from_fn(move || { + if bits == 0 { + return None; + } + let bit = bits.trailing_zeros() as usize; + bits &= bits - 1; + Some(i * BITS + bit) + }) + }) + } + + /// Restore the canonical form (no trailing zero spill words). + fn trim(&mut self) { + while self.spill.last() == Some(&0) { + self.spill.pop(); + } + } +} diff --git a/crates/burn-fusion/src/search/graph/tests.rs b/crates/burn-fusion/src/search/graph/tests.rs new file mode 100644 index 0000000..cfdd079 --- /dev/null +++ b/crates/burn-fusion/src/search/graph/tests.rs @@ -0,0 +1,331 @@ +use super::*; + +/// A plain node for exercising the graph algorithms without any tensor machinery. +#[derive(Default)] +struct TestNode { + position: usize, + produced: Vec, + read: Vec, + freed: Vec, +} + +impl TestNode { + fn new(position: usize) -> Self { + Self { + position, + ..Default::default() + } + } + + fn produces(mut self, resources: impl IntoIterator) -> Self { + self.produced.extend(resources); + self + } + + fn reads(mut self, resources: impl IntoIterator) -> Self { + self.read.extend(resources); + self + } + + /// Reads the resources for the last time (both read and freed). + fn frees(mut self, resources: impl IntoIterator) -> Self { + for resource in resources { + self.read.push(resource); + self.freed.push(resource); + } + self + } +} + +impl GraphNode for TestNode { + type Resource = u32; + + fn produced(&self) -> impl Iterator { + self.produced.iter().copied() + } + + fn read(&self) -> impl Iterator { + self.read.iter().copied() + } + + fn freed(&self) -> impl Iterator { + self.freed.iter().copied() + } + + fn position(&self) -> usize { + self.position + } +} + +mod subgraph { + use super::*; + + #[test] + fn insert_contains_iter() { + let mut set = SubGraph::empty(); + assert!(set.is_empty()); + + set.insert(3); + set.insert(70); // Spills past the inline word. + assert!(set.contains(3)); + assert!(set.contains(70)); + assert!(!set.contains(4)); + assert_eq!(set.len(), 2); + assert_eq!(set.iter().collect::>(), vec![3, 70]); + } + + #[test] + fn set_operations_keep_canonical_form() { + let mut a = SubGraph::single(2); + a.insert(100); + + // Subtracting the high bit must trim the trailing word so equality stays structural. + a.subtract(&SubGraph::single(100)); + assert_eq!(a, SubGraph::single(2)); + + let mut b = SubGraph::single(2); + b.union_with(&SubGraph::single(65)); + assert!(b.intersects(&SubGraph::single(65))); + assert!(!b.intersects(&SubGraph::single(64))); + + b.intersect_with(&SubGraph::single(2)); + assert_eq!(b, SubGraph::single(2)); + } +} + +mod dag { + use super::*; + + #[test] + fn independent_nodes_ordered_by_position() { + let nodes = vec![ + TestNode::new(2).reads([200, 201]).produces([202]), + TestNode::new(0).reads([100, 101]).produces([102]), + ]; + let dag = Dag::new(&nodes); + + assert!(!dag.depends_on(0, 1)); + assert!(!dag.depends_on(1, 0)); + assert_eq!(dag.topological_order(), Some(vec![1, 0])); + } + + #[test] + fn read_after_write_beats_position() { + // Node 0 (position 0) reads resource 50, which node 1 (position 5) produces. + let nodes = vec![ + TestNode::new(0).reads([50]).produces([2]), + TestNode::new(5).reads([9]).produces([50]), + ]; + let dag = Dag::new(&nodes); + + assert!(dag.depends_on(0, 1)); + assert_eq!(dag.topological_order(), Some(vec![1, 0])); + } + + #[test] + fn mutual_dependency_is_a_cycle() { + // Node 0 produces 100 (read by node 1) and reads 200 (produced by node 1). + let nodes = vec![ + TestNode::new(0).reads([200]).produces([100]), + TestNode::new(1).reads([100]).produces([200]), + ]; + let dag = Dag::new(&nodes); + + assert!(!dag.is_acyclic()); + assert_eq!(dag.topological_order(), None); + } + + #[test] + fn write_after_read_orders_freer_after_reader() { + // Both nodes read resource 100 (produced elsewhere); node 1 frees it. The freer must + // run after the reader, even though nothing flows between them. + let nodes = vec![ + TestNode::new(0).reads([100]).produces([2]), + TestNode::new(1).frees([100]).produces([4]), + ]; + let dag = Dag::new(&nodes); + + assert!(dag.depends_on(1, 0)); + assert!(!dag.depends_on(0, 1)); + assert_eq!(dag.topological_order(), Some(vec![0, 1])); + } + + #[test] + fn diamond_topological_order() { + // 0 + // / \ + // 1 2 + // \ / + // 3 + let nodes = vec![ + TestNode::new(0).produces([10]), + TestNode::new(1).reads([10]).produces([11]), + TestNode::new(2).reads([10]).produces([12]), + TestNode::new(3).reads([11, 12]).produces([13]), + ]; + let dag = Dag::new(&nodes); + + assert_eq!(dag.topological_order(), Some(vec![0, 1, 2, 3])); + } +} + +mod reachability { + use super::*; + + /// A → C → B chain: contracting A and B would trap C between them. + fn chain() -> Vec { + vec![ + TestNode::new(0).produces([100]), // A + TestNode::new(1).reads([100]).produces([200]), // C + TestNode::new(2).reads([200]).produces([300]), // B + ] + } + + #[test] + fn direct_edge_can_contract() { + let reachability = Dag::new(&chain()).reachability(); + + assert!(reachability.can_contract(&SubGraph::single(0), &SubGraph::single(1))); + assert!(reachability.can_contract(&SubGraph::single(1), &SubGraph::single(2))); + } + + #[test] + fn intermediate_node_blocks_contraction() { + let reachability = Dag::new(&chain()).reachability(); + + assert!(!reachability.can_contract(&SubGraph::single(0), &SubGraph::single(2))); + } + + #[test] + fn contraction_including_the_intermediate_is_legal() { + let reachability = Dag::new(&chain()).reachability(); + + // Contracting {A, C} with {B} is fine: nothing remains between them. + let mut a_and_c = SubGraph::single(0); + a_and_c.union_with(&SubGraph::single(1)); + assert!(reachability.can_contract(&a_and_c, &SubGraph::single(2))); + } + + #[test] + fn empty_subgraphs_are_refused() { + let reachability = Dag::new(&chain()).reachability(); + + assert!(!reachability.can_contract(&SubGraph::empty(), &SubGraph::empty())); + assert!(!reachability.can_contract(&SubGraph::single(0), &SubGraph::empty())); + } + + #[test] + fn transitive_ancestors_block_contraction() { + // 0 → 1 → 2 → 3: node 0 and node 3 have no direct edge, but both 1 and 2 sit between. + let nodes = vec![ + TestNode::new(0).produces([10]), + TestNode::new(1).reads([10]).produces([11]), + TestNode::new(2).reads([11]).produces([12]), + TestNode::new(3).reads([12]).produces([13]), + ]; + let reachability = Dag::new(&nodes).reachability(); + + assert!(!reachability.can_contract(&SubGraph::single(0), &SubGraph::single(3))); + assert!(!reachability.can_contract(&SubGraph::single(0), &SubGraph::single(2))); + } + + #[test] + fn cyclic_graph_falls_back_to_fixpoint() { + // 0 ↔ 1 feed each other (a cycle); 2 reads from 1. + let nodes = vec![ + TestNode::new(0).reads([200]).produces([100]), + TestNode::new(1).reads([100]).produces([200]), + TestNode::new(2).reads([200]).produces([300]), + ]; + let dag = Dag::new(&nodes); + assert!(!dag.is_acyclic()); + + let reachability = dag.reachability(); + // Contracting the cycle itself is legal — it disappears inside the contracted node... + assert!(reachability.can_contract(&SubGraph::single(0), &SubGraph::single(1))); + // ...but node 1 still sits between node 0 and node 2. + assert!(!reachability.can_contract(&SubGraph::single(0), &SubGraph::single(2))); + } + + #[test] + fn supports_more_than_64_nodes() { + // A chain of 70 nodes: node i reads what node i-1 produces. + let nodes = (0..70) + .map(|i| { + let node = TestNode::new(i).produces([i as u32 + 1]); + match i { + 0 => node, + _ => node.reads([i as u32]), + } + }) + .collect::>(); + let dag = Dag::new(&nodes); + let reachability = dag.reachability(); + + assert_eq!(dag.topological_order(), Some((0..70).collect::>())); + assert!(reachability.can_contract(&SubGraph::single(68), &SubGraph::single(69))); + assert!(!reachability.can_contract(&SubGraph::single(0), &SubGraph::single(69))); + } +} + +mod lifetime { + use super::*; + + /// Run the lifetime simulation on `nodes` executed in the given order. + fn valid(nodes: &[TestNode], order: &[usize]) -> bool { + is_valid_execution_order(order.iter().map(|&i| &nodes[i])) + } + + #[test] + fn program_order_is_valid() { + let nodes = vec![ + TestNode::new(0).reads([1]).produces([10]), + TestNode::new(1).frees([10]).produces([11]), + ]; + + assert!(valid(&nodes, &[0, 1])); + } + + #[test] + fn read_before_produce_is_invalid() { + let nodes = vec![ + TestNode::new(0).reads([1]).produces([10]), + TestNode::new(1).reads([10]).produces([11]), + ]; + + assert!(!valid(&nodes, &[1, 0])); + } + + #[test] + fn read_after_free_is_invalid() { + let nodes = vec![ + TestNode::new(0).frees([1]).produces([10]), + TestNode::new(1).reads([1]).produces([11]), + ]; + + assert!(!valid(&nodes, &[0, 1])); + assert!(valid(&nodes, &[1, 0])); + } + + #[test] + fn external_reads_are_assumed_live() { + // Resource 1 is produced by no ordered node: it comes from an earlier segment. + let nodes = vec![ + TestNode::new(0).reads([1]).produces([10]), + TestNode::new(1).reads([1, 10]).produces([11]), + ]; + + assert!(valid(&nodes, &[0, 1])); + } + + #[test] + fn order_may_cover_a_subset_of_nodes() { + // Only node 1 is ordered; the resource node 0 produces counts as external. + let nodes = vec![ + TestNode::new(0).produces([10]), + TestNode::new(1).reads([10]).produces([11]), + ]; + + assert!(valid(&nodes, &[1])); + } +} diff --git a/crates/burn-fusion/src/search/merging.rs b/crates/burn-fusion/src/search/merging.rs new file mode 100644 index 0000000..ffceffc --- /dev/null +++ b/crates/burn-fusion/src/search/merging.rs @@ -0,0 +1,517 @@ +use super::Block; +use crate::{NumOperations, search::graph::Reachability}; + +#[derive(Debug, PartialEq)] +/// The result of [merging](merge_blocks) [blocks](Block). +// `Full`/`Partial` carry owned blocks by design; the size gap between variants is expected. +#[allow(clippy::large_enum_variant)] +pub enum MergeBlocksResult { + /// All [blocks](Block) merged into one. + Full(Block), + /// Some [blocks](Block) merged and some failed. + Partial { + merged: Vec>, + failed: Vec>, + }, + /// All [blocks](Block) failed to merge. + Fail, +} + +/// Merge multiple [block](Block) together. +/// +/// The resulting [blocks](Block) might be sorted if the flag is true, otherwise the order isn't +/// guarantee. This is mostly useful for testing. +/// +/// # Strategy +/// +/// The merging strategy is in two steps: +/// +/// 1. The first step is to recursively try to merge adjacent blocks. This has the advantage of +/// trying multiple blocks ordering, therefore trying multiple permutation of the blocks. +/// However, it has the downside of not trying to merge blocks that are further away in the list +/// of blocks. Since trying all combinations possible is exponential, therefore not possible, we +/// fallback on the second strategy. +/// 2. The second step is to reduce blocks by setting an accumulator block, then sequentially +/// trying to merge the remaining blocks. We try some permutations based on the result from +/// step1. +pub fn merge_blocks( + blocks: &[&Block], + sorted: bool, + guard: &Reachability, +) -> MergeBlocksResult { + if blocks.is_empty() { + return MergeBlocksResult::Fail; + } + + if blocks.len() == 1 { + return MergeBlocksResult::Full(blocks[0].clone()); + } + + if blocks.len() == 2 { + let block0 = blocks[0]; + let block1 = blocks[1]; + + return match merge_two(block0, block1, guard) { + Some(result) => MergeBlocksResult::Full(result), + None => MergeBlocksResult::Fail, + }; + } + + let mut step1 = merge_blocks_step1(blocks, guard); + + if step1.full.len() == 1 && step1.failed.is_empty() && step1.partial.is_empty() { + MergeBlocksResult::Full(step1.full.remove(0)) + } else if step1.partial.len() == 1 && step1.failed.is_empty() && step1.full.is_empty() { + MergeBlocksResult::Full(step1.partial.remove(0)) + } else { + let result = merge_blocks_step2(step1, guard); + + if !sorted { + return result; + } + + match result { + MergeBlocksResult::Full(block) => MergeBlocksResult::Full(block), + MergeBlocksResult::Partial { + mut merged, + mut failed, + } => { + Block::sort(&mut merged); + Block::sort(&mut failed); + + MergeBlocksResult::Partial { merged, failed } + } + MergeBlocksResult::Fail => MergeBlocksResult::Fail, + } + } +} + +struct MergeBlockStep1 { + full: Vec>, + partial: Vec>, + failed: Vec>, +} + +impl Default for MergeBlockStep1 { + fn default() -> Self { + Self { + full: Default::default(), + partial: Default::default(), + failed: Default::default(), + } + } +} + +fn merge_blocks_step1( + blocks: &[&Block], + guard: &Reachability, +) -> MergeBlockStep1 { + let step_size = blocks.len() / 2; + let num_steps = f32::ceil(blocks.len() as f32 / step_size as f32) as usize; + + let mut result = MergeBlockStep1::default(); + + for i in 0..num_steps { + let start = i * step_size; + let end = usize::min(start + step_size, blocks.len()); + + match merge_blocks(&blocks[start..end], false, guard) { + MergeBlocksResult::Full(block) => { + result.full.push(block); + } + MergeBlocksResult::Partial { + mut merged, + mut failed, + } => { + result.partial.append(&mut merged); + result.failed.append(&mut failed); + } + MergeBlocksResult::Fail => { + for b in &blocks[start..end] { + result.failed.push((*b).clone()); + } + } + } + } + + result +} + +fn merge_blocks_step2( + mut step1: MergeBlockStep1, + guard: &Reachability, +) -> MergeBlocksResult { + // First let's try to merge partial graphs. + if step1.partial.len() > 1 { + match merge_accumulator(&step1.partial[0], &step1.partial[1..], guard) { + MergeBlocksResult::Full(block) => { + step1.partial = vec![block]; + } + MergeBlocksResult::Partial { merged, mut failed } => { + step1.partial = merged; + step1.failed.append(&mut failed); + } + MergeBlocksResult::Fail => {} + } + } + + // Then let's try to merge partial graphs with failed merges. + if !step1.failed.is_empty() { + step1.partial.append(&mut step1.failed); + match merge_accumulator(&step1.partial[0], &step1.partial[1..], guard) { + MergeBlocksResult::Full(block) => { + step1.partial = vec![block]; + } + MergeBlocksResult::Partial { merged, mut failed } => { + step1.partial = merged; + step1.failed.append(&mut failed); + } + MergeBlocksResult::Fail => {} + } + } + + // Then let's try to merge full graphs. + if step1.full.len() > 1 { + match merge_accumulator(&step1.full[0], &step1.full[1..], guard) { + MergeBlocksResult::Full(block) => { + step1.full = vec![block]; + } + MergeBlocksResult::Partial { merged, mut failed } => { + step1.full = merged; + step1.failed.append(&mut failed); + } + MergeBlocksResult::Fail => {} + } + } + + // Then let's try to merge full graphs with failed graphs. + if !step1.full.is_empty() { + step1.full.append(&mut step1.failed); + match merge_accumulator(&step1.full[0], &step1.full[1..], guard) { + MergeBlocksResult::Full(block) => { + step1.full = vec![block]; + } + MergeBlocksResult::Partial { merged, mut failed } => { + step1.full = merged; + step1.failed.append(&mut failed); + } + MergeBlocksResult::Fail => {} + } + } + + // Then let's try to merge full graphs with partial graphs. + if !step1.full.is_empty() || !step1.partial.is_empty() { + step1.full.append(&mut step1.partial); + match merge_accumulator(&step1.full[0], &step1.full[1..], guard) { + MergeBlocksResult::Full(block) => { + step1.full = vec![block]; + } + MergeBlocksResult::Partial { merged, mut failed } => { + step1.full = merged; + step1.failed.append(&mut failed); + } + MergeBlocksResult::Fail => { + // We do nothing. + } + } + } + + if step1.full.is_empty() { + MergeBlocksResult::Fail + } else if step1.failed.is_empty() { + if step1.full.len() == 1 { + MergeBlocksResult::Full(step1.full.remove(0)) + } else { + MergeBlocksResult::Partial { + merged: step1.full, + failed: vec![], + } + } + } else { + MergeBlocksResult::Partial { + merged: step1.full, + failed: step1.failed, + } + } +} + +fn merge_accumulator( + base: &Block, + blocks: &[Block], + guard: &Reachability, +) -> MergeBlocksResult { + let mut base = base.clone(); + let mut merged_failed = Vec::>::new(); + let mut merged_success = false; + + for block in blocks { + // `merge_two` checks the cycle guard and tries both merge directions — the accumulator + // may depend on the block, in which case only folding the accumulator *into* the block + // yields a valid operation order. + match merge_two(&base, block, guard) { + Some(merged) => { + merged_success = true; + base = merged; + } + None => { + merged_failed.push(block.clone()); + } + } + } + + if merged_success { + if merged_failed.is_empty() { + MergeBlocksResult::Full(base) + } else { + MergeBlocksResult::Partial { + merged: vec![base], + failed: merged_failed, + } + } + } else { + MergeBlocksResult::Fail + } +} + +fn merge_two( + a: &Block, + b: &Block, + guard: &Reachability, +) -> Option> { + if !guard.can_contract(a.constituents(), b.constituents()) { + return None; + } + + // Test each direction's operation order before paying for a deep clone of the base block + // (operations, builders, and data-flow sets all clone). + if a.can_append(b) { + let mut base = a.clone(); + if base.merge(b) { + return Some(base); + } + } + + if b.can_append(a) { + let mut base = b.clone(); + if base.merge(a) { + return Some(base); + } + } + + None +} + +#[cfg(test)] +mod tests { + use super::*; + pub use crate::stream::execution::tests::{TestOptimization, TestOptimizationBuilder}; + use crate::{ + OperationFuser, + search::graph::Dag, + stream::tests::{operation_1, operation_2, operation_3}, + }; + + /// Seed constituents by index, build the [Reachability] guard from the blocks, and merge + /// (sorted). + fn merge(mut blocks: Vec>) -> MergeBlocksResult { + for (i, block) in blocks.iter_mut().enumerate() { + block.seed_constituent(i); + } + let refs = blocks.iter().collect::>(); + let guard = Dag::new(&refs).reachability(); + merge_blocks(&refs, true, &guard) + } + + #[test] + fn test_merge_blocks_no_block() { + let actual = merge(Vec::>::new()); + + assert_eq!(actual, MergeBlocksResult::Fail); + } + + #[test] + fn test_merge_blocks_single() { + let builders = builders(); + let block = Block::new(&builders); + let actual = merge(vec![block.clone()]); + + assert_eq!(actual, MergeBlocksResult::Full(block)); + } + + #[test] + fn test_merge_blocks_two_blocks() { + let builders = builders(); + let mut block1 = Block::new(&builders); + let mut block2 = Block::new(&builders); + block1.register(&operation_1(), 0, false); + block1.register(&operation_1(), 1, false); + block2.register(&operation_1(), 2, false); + block2.register(&operation_1(), 3, false); + + let actual = merge(vec![block1, block2]); + + let mut expected = Block::new(&builders); + expected.register(&operation_1(), 0, false); + expected.register(&operation_1(), 1, false); + expected.register(&operation_1(), 2, false); + expected.register(&operation_1(), 3, false); + + assert_eq!(actual, MergeBlocksResult::Full(expected)); + } + + #[test] + fn test_merge_blocks_three_blocks() { + let builders = builders(); + let mut block1 = Block::new(&builders); + let mut block2 = Block::new(&builders); + let mut block3 = Block::new(&builders); + block1.register(&operation_1(), 0, false); + block2.register(&operation_1(), 1, false); + block3.register(&operation_1(), 2, false); + + let actual = merge(vec![block1, block2, block3]); + + let mut expected = Block::new(&builders); + expected.register(&operation_1(), 0, false); + expected.register(&operation_1(), 1, false); + expected.register(&operation_1(), 2, false); + + assert_eq!(actual, MergeBlocksResult::Full(expected)); + } + + #[test] + fn test_merge_blocks_three_blocks_partial() { + let builders = builders(); + let mut block1 = Block::new(&builders); + let mut block2 = Block::new(&builders); + let mut block3 = Block::new(&builders); + block1.register(&operation_1(), 0, false); + block2.register(&operation_2(), 1, false); + block3.register(&operation_1(), 2, false); + + let actual = merge(vec![block1, block2, block3]); + + let mut expected1 = Block::new(&builders); + let mut expected2 = Block::new(&builders); + expected1.register(&operation_1(), 0, false); + expected1.register(&operation_1(), 2, false); + expected2.register(&operation_2(), 1, false); + + assert_eq!( + actual, + MergeBlocksResult::Partial { + merged: vec![expected1, expected2], + failed: vec![] + } + ); + } + + #[test] + fn test_merge_blocks_four_blocks_partial_with_failure() { + let builders = builders(); + let mut block1 = Block::new(&builders); + let mut block2 = Block::new(&builders); + let mut block3 = Block::new(&builders); + let mut block4 = Block::new(&builders); + block1.register(&operation_1(), 0, false); + block2.register(&operation_2(), 1, false); + block3.register(&operation_1(), 2, false); + block4.register(&operation_3(), 3, false); + + let actual = merge(vec![block1, block2, block3, block4]); + + let mut expected1 = Block::new(&builders); + let mut expected2 = Block::new(&builders); + let mut failed = Block::new(&builders); + expected1.register(&operation_1(), 0, false); + expected1.register(&operation_1(), 2, false); + expected2.register(&operation_2(), 1, false); + failed.register(&operation_3(), 3, false); + + assert_eq!( + actual, + MergeBlocksResult::Partial { + merged: vec![expected1], + failed: vec![expected2, failed] + } + ); + } + + #[test] + fn test_merge_blocks_five_blocks_partial_with_failure() { + let builders = builders(); + let mut block1 = Block::new(&builders); + let mut block2 = Block::new(&builders); + let mut block3 = Block::new(&builders); + let mut block4 = Block::new(&builders); + let mut block5 = Block::new(&builders); + block1.register(&operation_1(), 0, false); + block2.register(&operation_2(), 1, false); + block3.register(&operation_1(), 2, false); + block4.register(&operation_3(), 3, false); + block5.register(&operation_2(), 4, false); + + let actual = merge(vec![block1, block2, block3, block4, block5]); + + let mut expected1 = Block::new(&builders); + let mut expected2 = Block::new(&builders); + let mut failed = Block::new(&builders); + expected1.register(&operation_1(), 0, false); + expected1.register(&operation_1(), 2, false); + expected2.register(&operation_2(), 1, false); + expected2.register(&operation_2(), 4, false); + failed.register(&operation_3(), 3, false); + + assert_eq!( + actual, + MergeBlocksResult::Partial { + merged: vec![expected1, expected2], + failed: vec![failed] + } + ); + } + + /// The accumulator may depend on a block it tries to absorb: folding that block's ops at the + /// end would fuse a consumer before its producer. The merge must then happen in the other + /// direction (the accumulator folded into the block), preserving a full merge instead of + /// degrading to a partial one. + #[test] + fn test_merge_blocks_accumulator_reverses_direction() { + use crate::search::testing::add; + + let x = add(1, 2, 10); + let y = add(50, 3, 11); // Reads 50... + let z = add(4, 5, 50); // ...which this op produces. + let w = add(6, 7, 12); // Independent. + + // The builder accepts exactly the reversed merge order [z, x, y, w]. + let pattern = vec![z.clone(), x.clone(), y.clone(), w.clone(), x.clone()]; + let builders: Vec>> = + vec![Box::new(TestOptimizationBuilder::new(0, pattern))]; + + let mut block1 = Block::new(&builders); + block1.register(&x, 0, true); + block1.register(&y, 5, true); + let mut block2 = Block::new(&builders); + block2.register(&z, 3, true); + let mut block3 = Block::new(&builders); + block3.register(&w, 6, true); + + let actual = merge(vec![block1, block2, block3]); + + let mut expected = Block::new(&builders); + expected.register(&z, 3, true); + expected.register(&x, 0, true); + expected.register(&y, 5, true); + expected.register(&w, 6, true); + + assert_eq!(actual, MergeBlocksResult::Full(expected)); + } + + fn builders() -> Vec>> { + let builder_1 = TestOptimizationBuilder::new(0, vec![operation_1(); 10]); + let builder_2 = TestOptimizationBuilder::new(1, vec![operation_2(); 10]); + + vec![Box::new(builder_1), Box::new(builder_2)] + } +} diff --git a/crates/burn-fusion/src/search/mod.rs b/crates/burn-fusion/src/search/mod.rs new file mode 100644 index 0000000..3ff1ada --- /dev/null +++ b/crates/burn-fusion/src/search/mod.rs @@ -0,0 +1,10 @@ +mod block; +mod optimization; + +pub(crate) mod graph; +pub(super) mod merging; +#[cfg(test)] +pub(crate) mod testing; +pub(super) use block::*; + +pub use optimization::*; diff --git a/crates/burn-fusion/src/search/optimization/blocks.rs b/crates/burn-fusion/src/search/optimization/blocks.rs new file mode 100644 index 0000000..17809f9 --- /dev/null +++ b/crates/burn-fusion/src/search/optimization/blocks.rs @@ -0,0 +1,162 @@ +use burn_std::config::{fusion::FusionLogLevel, log_fusion}; + +use crate::{ + NumOperations, + search::{ + Block, BlockOptimization, + graph::Dag, + merging::{MergeBlocksResult, merge_blocks}, + }, + stream::store::ExecutionStrategy, +}; + +/// Try to optimize a list of [blocks](Block) into a [block optimization](BlockOptimization). +/// +/// # Notes +/// +/// The blocks form a dependency DAG: a block may depend on another when it consumes a tensor the +/// other produces. Blocks with no dependency between them can still be merged/reordered freely; +/// dependent blocks must be emitted in a topological order (dependencies first). +/// +/// The contract is that the length of operations executed must include all operations. If we don't +/// find an optimization that can be executed with that constraint, we return a +/// [BlocksOptimizerResult::WithHoles]. +pub struct BlocksOptimizer { + blocks: Vec>, + num_ops: usize, +} + +/// When we can't find a proper optimization for the provided list of [blocks](Block). +pub enum BlocksOptimizerResult { + /// When an optimization fill the hole stream. + Full(BlockOptimization), + /// The optimization found with the holes indices. + WithHoles { + strategies: Vec>>, + ordering: Vec, + holes: Vec, + }, +} + +impl BlocksOptimizer { + /// Create a new optimizer with the given blocks. + pub fn new(blocks: Vec>) -> Self { + let num_ops: usize = blocks.iter().map(|g| g.end_pos).max().unwrap(); + + Self { blocks, num_ops } + } + + /// Optimizes the blocks. + /// + /// Strategy: + /// 1. Try to merge blocks together — the merge guard rejects any contraction that would + /// create a dependency cycle, so a surviving merge is always order-safe. + /// 2. Emit the blocks in a topological order (dependencies first) and ask every block for its + /// best optimization (or the fallback [Operations](ExecutionStrategy::Operations) if no + /// builder matched), concatenating the strategies. + /// 3. An *interior* unresolved position — an op that no block's final + /// ordering covered, but which sits before some other resolved + /// position — is a hole that must be filled by a second optimization + /// pass. Trailing unresolved positions, on the other hand, are the + /// natural tail of a drained block and are left in the queue for the + /// processor to handle in the next round. + pub fn optimize(mut self) -> BlocksOptimizerResult { + self = self.merging_pass(); + + let num_ops = self.num_ops; + let blocks = core::mem::take(&mut self.blocks); + + // Emit blocks in a valid execution order: a dependency must run before its dependents. + // The set is acyclic (register keeps it so, and `merging_pass` only accepts cycle-free + // merges), so `topological_order` always succeeds; fall back to the current order if not. + let order = Dag::new(&blocks) + .topological_order() + .unwrap_or_else(|| (0..blocks.len()).collect::>()); + let mut slots: Vec>> = blocks.into_iter().map(Some).collect(); + let blocks: Vec> = order + .into_iter() + .map(|i| slots[i].take().expect("each block taken once")) + .collect(); + + let mut strategies: Vec>> = Vec::with_capacity(blocks.len()); + let mut ordering = Vec::new(); + let mut resolved = vec![false; num_ops]; + + for block in blocks { + let mut block_opt = block.optimize(); + for pos in block_opt.ordering.iter() { + resolved[*pos] = true; + } + ordering.append(&mut block_opt.ordering); + strategies.push(Box::new(block_opt.strategy)); + } + + // An unresolved position is a hole only if some position *after* it + // is resolved (it's interleaved, not trailing). A trailing run of + // unresolved positions is a drained tail — left for the processor. + let last_resolved_end = resolved + .iter() + .rposition(|&r| r) + .map(|i| i + 1) + .unwrap_or(0); + let holes: Vec = (0..last_resolved_end).filter(|i| !resolved[*i]).collect(); + + let num_strategies = strategies.len(); + log_fusion(FusionLogLevel::Basic, move || { + if num_strategies > 1 { + format!("selected composed strategy ({num_strategies} sub-strategies)") + } else { + "selected single strategy".to_string() + } + }); + + if holes.is_empty() { + let strategy = if strategies.len() > 1 { + ExecutionStrategy::Composed(strategies) + } else { + *strategies.remove(0) + }; + BlocksOptimizerResult::Full(BlockOptimization::new(strategy, ordering)) + } else { + BlocksOptimizerResult::WithHoles { + strategies, + ordering, + holes, + } + } + } + + /// Try to merge blocks together. + fn merging_pass(mut self) -> Self { + if self.blocks.len() == 1 { + return self; + } + + Block::sort(&mut self.blocks); + + // Seed constituents so the guard can reason about the original dependency DAG through the + // recursive merging, then reject any contraction that would create a cycle. + for (i, block) in self.blocks.iter_mut().enumerate() { + block.seed_constituent(i); + } + let blocks = self.blocks.iter().collect::>(); + let guard = Dag::new(&blocks).reachability(); + + match merge_blocks(&blocks, false, &guard) { + MergeBlocksResult::Full(block) => { + self.blocks = vec![block]; + } + MergeBlocksResult::Partial { + mut merged, + mut failed, + } => { + merged.append(&mut failed); + self.blocks = merged; + Block::sort(&mut self.blocks); + } + MergeBlocksResult::Fail => {} + } + + self + } +} diff --git a/crates/burn-fusion/src/search/optimization/mod.rs b/crates/burn-fusion/src/search/optimization/mod.rs new file mode 100644 index 0000000..9731cd0 --- /dev/null +++ b/crates/burn-fusion/src/search/optimization/mod.rs @@ -0,0 +1,7 @@ +mod blocks; +mod stream; + +#[cfg(test)] +mod tests; + +pub use stream::*; diff --git a/crates/burn-fusion/src/search/optimization/stream.rs b/crates/burn-fusion/src/search/optimization/stream.rs new file mode 100644 index 0000000..7e324b5 --- /dev/null +++ b/crates/burn-fusion/src/search/optimization/stream.rs @@ -0,0 +1,827 @@ +use super::blocks::BlocksOptimizer; +use crate::{ + NumOperations, OperationFuser, + search::{ + Block, BlockOptimization, OperationNode, RegistrationResult, + graph::{Dag, GraphNode, is_valid_execution_order}, + merging::{MergeBlocksResult, merge_blocks}, + optimization::blocks::BlocksOptimizerResult, + }, + stream::{execution::op_kind, store::ExecutionStrategy}, +}; +use burn_ir::{OperationIr, TensorId, TensorStatus}; +use burn_std::config::{config, fusion::FusionLogLevel, log_fusion}; +use std::collections::HashSet; +use std::sync::Arc; + +/// Optimize a stream of [operations](OperationIr) using a list of [builders](OptimizationBuilder). +pub struct StreamOptimizer { + builders: Vec>>, + blocks: Vec>, + length: usize, + stopped: bool, + max_blocks: Option, +} + +impl StreamOptimizer { + /// Create a new stream optimizer. + pub fn new(builders: Vec>>) -> Self { + // Too high and it may break the fusion cache always retriggering explorations. It also + // bounds the per-op cost of the dependency analysis: every registration that touches + // several blocks rebuilds the block DAG and its transitive reachability, which scale + // quadratically (and worse) with the number of blocks. + let max_blocks = Some(config().fusion().beam_search.max_blocks); + Self { + builders, + blocks: Vec::new(), + length: 0, + stopped: false, + max_blocks, + } + } + + /// Register a new [operation](OperationIr) in the optimizer. + /// + /// You can use the function [Self::still_optimizing] to know if the operations are actually + /// being registered. + pub fn register(&mut self, operation: &OperationIr) { + if self.stopped { + let length = self.length; + log_fusion(FusionLogLevel::Full, || { + format!( + "[stream] {} dropped (optimizer stopped at op {length})", + op_kind(operation) + ) + }); + return; + } + + if self.blocks.is_empty() { + self.on_new_block(operation); + self.length += 1; + return; + } + + match self.merge_blocks(operation, false) { + MergeBlockStep::Full | MergeBlockStep::NoNeed => {} + MergeBlockStep::Fail | MergeBlockStep::Partial => { + // The operation ties together blocks that couldn't be merged into one. Instead of + // giving up on the segment, give the operation its own block that *depends* on the + // blocks it reads from — the block set becomes a dependency DAG. + self.on_dependent_op(operation); + if self.stopped { + return; + } + self.length += 1; + return; + } + } + + if let Some(max_blocks) = self.max_blocks { + if self.register_max_block(operation, max_blocks) { + self.length += 1; + } else { + let length = self.length; + log_fusion(FusionLogLevel::Medium, || { + format!( + "[stream] stopped (max_blocks={max_blocks} reached) at op {length} ({})", + op_kind(operation) + ) + }); + self.stopped = true; + } + return; + } + + let added_count = self.register_inner(operation, false); + if added_count == 0 { + self.on_new_block(operation); + } else { + self.log_accepted(operation, added_count); + } + + self.length += 1; + } + + /// Optimize the current stream on the given [operations](OperationIr). + /// + /// # Notes + /// + /// The operations provided are the same as the ones used in the [register](Self::register) + /// method, this simply remove the need for the current type to also keep track of the list of + /// operations. + pub fn optimize(&self, operations: &[OperationIr]) -> BlockOptimization { + let result = BlocksOptimizer::new(self.blocks.clone()).optimize(); + + let out = match result { + BlocksOptimizerResult::Full(block_optimization) => block_optimization, + BlocksOptimizerResult::WithHoles { + mut strategies, + mut ordering, + mut holes, + } => { + loop { + let mut search = self.new_empty_search(); + + let mut operations_holes = Vec::with_capacity(holes.len()); + + for index in holes.iter() { + let op = &operations[*index]; + operations_holes.push(op.clone()); + search.register(op); + } + + let mut optimization_of_holes = search.optimize(&operations_holes); + + optimization_of_holes.map_ordering(&holes); + + // Append the re-optimized holes as their own chunk; `repair_order` (below) puts + // every chunk into a hazard-respecting execution order, so the placement here + // doesn't need to be correct — only complete. + let consumed = optimization_of_holes.ordering.len(); + strategies.push(Box::new(optimization_of_holes.strategy)); + ordering.append(&mut optimization_of_holes.ordering); + holes.drain(0..consumed); + + if holes.is_empty() { + break; + } + } + + BlockOptimization::new(ExecutionStrategy::Composed(strategies), ordering) + } + }; + + repair_order(out, operations) + } + + /// Reset the state of the optimizer. + pub fn reset(&mut self) { + self.builders.iter_mut().for_each(|b| b.reset()); + self.length = 0; + self.blocks.clear(); + self.stopped = false; + } + + /// Returns if some optimizations are still possible within the stream. + pub fn still_optimizing(&self) -> bool { + if self.stopped { + return false; + } + if self.blocks.is_empty() { + return true; + } + + let mut num_stopped = 0; + + for block in self.blocks.iter() { + if !block.still_optimizing() { + num_stopped += 1 + } + } + + num_stopped < self.blocks.len() + } + + fn register_max_block(&mut self, operation: &OperationIr, max_blocks: usize) -> bool { + if max_blocks == 1 { + // Register in the single block with a force. + self.register_inner(operation, true); + return true; + } + let added_count = self.register_inner(operation, false); + + if added_count > 0 { + self.log_accepted(operation, added_count); + return true; + } + + if added_count == 0 && self.blocks.len() < max_blocks { + self.on_new_block(operation); + return true; + } + + self.merge_blocks(operation, true); + + if self.blocks.len() >= max_blocks { + self.stopped = true; + return false; + } + + let added_count = self.register_inner(operation, false); + + if added_count == 0 { + self.on_new_block(operation); + } else { + self.log_accepted(operation, added_count); + } + + true + } + + fn log_accepted(&self, operation: &OperationIr, added_count: usize) { + let length = self.length; + let num_blocks = self.blocks.len(); + log_fusion(FusionLogLevel::Full, || { + format!( + "[stream] op {length} {} → accepted in {added_count}/{num_blocks} block(s)", + op_kind(operation) + ) + }); + } + + fn register_inner(&mut self, operation: &OperationIr, force: bool) -> usize { + let mut added_count = 0; + for block in self.blocks.iter_mut() { + match block.register(operation, self.length, force) { + RegistrationResult::Accepted => { + added_count += 1; + } + RegistrationResult::NotPartOfTheGraph => {} + } + } + added_count + } + + fn new_empty_search(&self) -> Self { + Self::new( + self.builders + .iter() + .map(|b| { + let mut b = b.clone_dyn(); + b.reset(); + b + }) + .collect(), + ) + } + + fn merge_blocks(&mut self, operation: &OperationIr, all: bool) -> MergeBlockStep { + let nodes = operation.nodes(); + let mut block_merges = Vec::new(); + + for (i, block) in self.blocks.iter().enumerate() { + if all || block.contains_tensors(&nodes) { + block_merges.push(i); + } + } + + if block_merges.len() <= 1 { + return MergeBlockStep::NoNeed; + } + + // Seed each block with its index, then build the cycle guard over ALL blocks — a block + // that "sits between" two merge candidates need not itself be a candidate, so the guard + // must see the full dependency graph, not just the subset being merged. + for (i, block) in self.blocks.iter_mut().enumerate() { + block.seed_constituent(i); + } + let all_blocks = self.blocks.iter().collect::>(); + let guard = Dag::new(&all_blocks).reachability(); + + let blocks_to_merge = self + .blocks + .iter() + .enumerate() + .filter_map(|(i, g)| match block_merges.contains(&i) { + true => Some(g), + false => None, + }) + .collect::>(); + + let merged = merge_blocks(&blocks_to_merge, false, &guard); + + let mut clear_blocks = || { + let mut indices = block_merges.to_vec(); + indices.sort(); + + for g in indices.into_iter().rev() { + self.blocks.remove(g); + } + }; + + match merged { + MergeBlocksResult::Full(block) => { + clear_blocks(); + self.blocks.push(block); + Block::sort(&mut self.blocks); + MergeBlockStep::Full + } + MergeBlocksResult::Partial { + mut merged, + mut failed, + } => { + clear_blocks(); + self.blocks.append(&mut merged); + self.blocks.append(&mut failed); + Block::sort(&mut self.blocks); + MergeBlockStep::Partial + } + MergeBlocksResult::Fail => MergeBlockStep::Fail, + } + } + + fn on_new_block(&mut self, operation: &OperationIr) { + let mut block = Block::new(&self.builders); + block.register(operation, self.length, true); + self.blocks.push(block); + + let length = self.length; + let num_blocks = self.blocks.len(); + log_fusion(FusionLogLevel::Full, || { + format!( + "[stream] op {length} {} → new block (total: {num_blocks})", + op_kind(operation) + ) + }); + } + + /// Give the operation its own block that depends on the blocks it reads from. + /// + /// The new block owns the operation's outputs and records its upstream inputs as external — + /// those external inputs form the dependency edges. A freshly created block is a pure sink, so + /// it cannot create a cycle, except through an in-place output that re-produces an upstream + /// tensor; that case is caught by the defensive acyclicity check. + fn on_dependent_op(&mut self, operation: &OperationIr) { + // Dependent blocks count toward the cap. With no room left, try to free a slot by + // force-merging mergeable blocks; if that fails, stop the segment (the op is deferred). + if let Some(max_blocks) = self.max_blocks + && self.blocks.len() >= max_blocks + { + self.merge_blocks(operation, true); + + if self.blocks.len() >= max_blocks { + let length = self.length; + log_fusion(FusionLogLevel::Medium, || { + format!( + "[stream] stopped (max_blocks={max_blocks} reached on dependency) at op {length} ({})", + op_kind(operation) + ) + }); + self.stopped = true; + return; + } + } + + let mut block = Block::new(&self.builders); + block.register(operation, self.length, true); + self.blocks.push(block); + + if !Dag::new(&self.blocks).is_acyclic() { + self.blocks.pop(); + let length = self.length; + log_fusion(FusionLogLevel::Medium, || { + format!( + "[stream] stopped (unresolvable dependency cycle) at op {length} ({})", + op_kind(operation) + ) + }); + self.stopped = true; + return; + } + + let length = self.length; + let num_blocks = self.blocks.len(); + log_fusion(FusionLogLevel::Full, || { + format!( + "[stream] op {length} {} → new dependent block (total: {num_blocks})", + op_kind(operation) + ) + }); + } +} + +enum MergeBlockStep { + Full, + Partial, + Fail, + NoNeed, +} + +/// Reorder the top-level strategy chunks so the execution order respects tensor handle lifetimes. +/// +/// Blocks and re-optimized holes are placed by separate, incremental heuristics that can't see the +/// whole picture (a hole may depend on another hole spliced later). This final pass treats each +/// top-level strategy as an atomic [Chunk] node and orders the chunks topologically (see +/// [GraphNode] for the hazard rules deriving the edges). Ties keep the earliest stream position +/// first, reproducing the historical order when there are no hazards. +/// +/// If the chunk graph has a cycle — the chunking can't be linearized as atomic units — unfused +/// chunks are [split](repair_order_split) to break the cycle before giving up on fusion. +fn repair_order(opt: BlockOptimization, operations: &[OperationIr]) -> BlockOptimization { + let (strategies, ordering) = match opt.strategy { + ExecutionStrategy::Composed(items) => (items, opt.ordering), + // A single strategy has nothing to reorder across, but a merge can still leave its ops + // internally out of stream order — validate, and unfuse in stream order if it's broken. + single => { + if ordering_is_valid(&opt.ordering, operations) { + return BlockOptimization::new(single, opt.ordering); + } + return unfused_stream_order(opt.ordering); + } + }; + + // Split the concatenated ordering back into one chunk per top-level strategy. + let mut chunks = Vec::with_capacity(strategies.len()); + let mut offset = 0; + for strategy in &strategies { + let len = strategy_len(strategy); + chunks.push(Chunk::new( + ordering[offset..offset + len].to_vec(), + operations, + )); + offset += len; + } + + let order = match Dag::new(&chunks).topological_order() { + Some(order) => order, + // The chunks can't be linearized as atomic units. Unfused chunks are free to split + // (their ops execute individually), which may break the cycle while every fused + // optimization survives intact. + None => return repair_order_split(strategies, chunks, operations), + }; + + assemble(strategies, &chunks, &order, operations) +} + +/// Retry [repair_order] with every [Operations](ExecutionStrategy::Operations) chunk split into +/// single-operation chunks, keeping fused optimizations atomic. Splitting only removes ordering +/// constraints, so this resolves any cycle that isn't between fused chunks themselves. +fn repair_order_split( + strategies: Vec>>, + chunks: Vec, + operations: &[OperationIr], +) -> BlockOptimization { + let mut split_strategies = Vec::new(); + let mut split_chunks = Vec::new(); + for (strategy, chunk) in strategies.into_iter().zip(chunks) { + match *strategy { + ExecutionStrategy::Operations { .. } => { + for position in chunk.positions { + split_strategies.push(Box::new(ExecutionStrategy::Operations { + ordering: Arc::new(vec![position]), + })); + split_chunks.push(Chunk::new(vec![position], operations)); + } + } + strategy => { + split_strategies.push(Box::new(strategy)); + split_chunks.push(chunk); + } + } + } + + match Dag::new(&split_chunks).topological_order() { + Some(order) => assemble(split_strategies, &split_chunks, &order, operations), + // Fused chunks are mutually dependent: the segment isn't executable as built — run + // everything unfused in stream order, which is always valid. + None => unfused_stream_order( + split_chunks + .into_iter() + .flat_map(|chunk| chunk.positions) + .collect(), + ), + } +} + +/// Emit the strategies in the given chunk order, validating the final operation-level ordering. +fn assemble( + strategies: Vec>>, + chunks: &[Chunk], + order: &[usize], + operations: &[OperationIr], +) -> BlockOptimization { + let mut slots: Vec>>> = + strategies.into_iter().map(Some).collect(); + let mut new_strategies = Vec::with_capacity(order.len()); + let mut new_ordering = Vec::new(); + for &k in order { + new_strategies.push(slots[k].take().expect("each chunk taken once")); + new_ordering.extend_from_slice(&chunks[k].positions); + } + + // A chunk that is internally miss-ordered can still violate handle lifetimes even after the + // chunk-level sort. If the result isn't executable, fall back to running every operation + // unfused in stream order. + if !ordering_is_valid(&new_ordering, operations) { + return unfused_stream_order(new_ordering); + } + + BlockOptimization::new(ExecutionStrategy::Composed(new_strategies), new_ordering) +} + +/// Run every operation unfused, in ascending stream order — always a valid execution order. +/// +/// This is the last-resort fallback: every fusion in the segment is dropped. Logged so a segment +/// that silently degrades shows up when investigating fusion regressions. +fn unfused_stream_order(mut positions: Vec) -> BlockOptimization { + let num_ops = positions.len(); + log_fusion(FusionLogLevel::Medium, || { + format!("[repair] falling back to unfused stream order ({num_ops} ops)") + }); + + positions.sort_unstable(); + let ordering = Arc::new(positions.clone()); + BlockOptimization::new(ExecutionStrategy::Operations { ordering }, positions) +} + +/// Whether the execution order respects tensor handle lifetimes: every operation's inputs must be +/// live when it runs (see [is_valid_execution_order]). +fn ordering_is_valid(ordering: &[usize], operations: &[OperationIr]) -> bool { + is_valid_execution_order(ordering.iter().map(|&position| OperationNode { + operation: &operations[position], + position, + })) +} + +/// Number of stream positions a strategy covers (its share of the concatenated ordering). +fn strategy_len(strategy: &ExecutionStrategy) -> usize { + match strategy { + ExecutionStrategy::Optimization { ordering, .. } => ordering.len(), + ExecutionStrategy::Operations { ordering } => ordering.len(), + ExecutionStrategy::Composed(items) => items.iter().map(|s| strategy_len(s)).sum(), + } +} + +/// One top-level strategy of a composed optimization, viewed as a single atomic [GraphNode] +/// covering the stream positions of its operations. +struct Chunk { + positions: Vec, + produced: HashSet, + read: HashSet, + freed: HashSet, +} + +impl Chunk { + fn new(positions: Vec, operations: &[OperationIr]) -> Self { + let mut produced = HashSet::new(); + for &position in &positions { + for tensor in operations[position].outputs() { + produced.insert(tensor.id); + } + } + + let mut read = HashSet::new(); + let mut freed = HashSet::new(); + for &position in &positions { + for tensor in operations[position].inputs() { + // A ReadWrite read frees the tensor even when this chunk also produced it — the + // write-after-read hazard against readers in *other* chunks still applies + // (mirrors [Block::register_op]). + if let TensorStatus::ReadWrite = tensor.status { + freed.insert(tensor.id); + } + if produced.contains(&tensor.id) { + continue; + } + read.insert(tensor.id); + } + } + + Self { + positions, + produced, + read, + freed, + } + } +} + +impl GraphNode for Chunk { + type Resource = TensorId; + + fn produced(&self) -> impl Iterator { + self.produced.iter().copied() + } + + fn read(&self) -> impl Iterator { + self.read.iter().copied() + } + + fn freed(&self) -> impl Iterator { + self.freed.iter().copied() + } + + fn produces(&self, resource: TensorId) -> bool { + self.produced.contains(&resource) + } + + fn reads(&self, resource: TensorId) -> bool { + self.read.contains(&resource) + } + + fn position(&self) -> usize { + self.positions.iter().copied().min().unwrap_or(0) + } +} + +/// Tests for the [repair_order] pass (driven directly with hand-built chunkings) and the +/// [on_dependent_op](StreamOptimizer::on_dependent_op) stop paths. +#[cfg(test)] +mod tests { + use super::*; + use crate::search::testing::{add, add_rw}; + use crate::stream::execution::tests::{TestOptimization, TestOptimizationBuilder}; + + fn fused_with( + builder_id: usize, + ordering: Vec, + score: u64, + ) -> ExecutionStrategy { + ExecutionStrategy::Optimization { + opt: TestOptimization::new(builder_id, ordering.len()), + ordering: Arc::new(ordering), + score, + } + } + + fn fused(ordering: Vec) -> ExecutionStrategy { + fused_with(0, ordering, 0) + } + + fn unfused(ordering: Vec) -> ExecutionStrategy { + ExecutionStrategy::Operations { + ordering: Arc::new(ordering), + } + } + + fn composed( + parts: Vec>, + ) -> ExecutionStrategy { + ExecutionStrategy::Composed(parts.into_iter().map(Box::new).collect()) + } + + /// A chunk producing AND freeing the same tensor still constrains readers in other chunks: + /// the free must be a write-after-read edge. With the edge, the chunk-level cycle is + /// detected and the unfused chunk is split around the fused one, salvaging the fusion; + /// without it, the invalid order is only caught by the final validation, which unfuses the + /// whole segment. + #[test] + fn splits_unfused_produce_free_chunk_instead_of_unfusing_everything() { + let ops = vec![ + add(1, 2, 10), // 0: produces 10 (unfused chunk) + add(10, 3, 100), // 1: fused, reads 10 + add(100, 4, 101), // 2: fused + add_rw(10, 5, 11), // 3: frees 10 (same unfused chunk as op 0) + ]; + let opt = BlockOptimization::new( + composed(vec![fused(vec![1, 2]), unfused(vec![0, 3])]), + vec![1, 2, 0, 3], + ); + + let repaired = repair_order(opt, &ops); + + // The fused pair survives; the produce/free chunk is split around it. + assert_eq!(repaired.ordering, vec![0, 1, 2, 3]); + assert_eq!( + repaired.strategy, + composed(vec![unfused(vec![0]), fused(vec![1, 2]), unfused(vec![3])]), + ); + } + + /// A chunk appended after a chunk it feeds (a hole depending on a later-spliced hole) is + /// moved before its consumer. + #[test] + fn reorders_chunks_by_dependency() { + let ops = vec![ + add(1, 2, 5), // 0: fused + add(5, 3, 10), // 1: fused, produces 10 + add(10, 4, 20), // 2: unfused, reads 10 + ]; + // The consumer chunk was placed first. + let opt = BlockOptimization::new( + composed(vec![unfused(vec![2]), fused(vec![0, 1])]), + vec![2, 0, 1], + ); + + let repaired = repair_order(opt, &ops); + + assert_eq!(repaired.ordering, vec![0, 1, 2]); + assert_eq!( + repaired.strategy, + composed(vec![fused(vec![0, 1]), unfused(vec![2])]), + ); + } + + /// Two fused chunks that feed each other cannot be linearized as atomic units and cannot be + /// split — the whole segment falls back to unfused stream order. + #[test] + fn falls_back_to_unfused_when_fused_chunks_cycle() { + let ops = vec![ + add(1, 2, 10), // 0: chunk A, produces 10 + add(10, 6, 20), // 1: chunk B, reads 10, produces 20 + add(20, 7, 21), // 2: chunk B + add(20, 5, 11), // 3: chunk A, reads 20 + ]; + let opt = BlockOptimization::new( + composed(vec![fused(vec![0, 3]), fused(vec![1, 2])]), + vec![0, 3, 1, 2], + ); + + let repaired = repair_order(opt, &ops); + + assert_eq!(repaired.ordering, vec![0, 1, 2, 3]); + assert_eq!(repaired.strategy, unfused(vec![0, 1, 2, 3])); + } + + /// A single (non-composed) strategy whose internal order breaks a tensor lifetime is + /// replaced by unfused stream order. + #[test] + fn single_strategy_with_invalid_order_unfuses() { + let ops = vec![ + add(1, 2, 10), // 0: produces 10 + add(10, 3, 11), // 1: reads 10 + ]; + // The consumer is fused before its producer. + let opt = BlockOptimization::new(fused(vec![1, 0]), vec![1, 0]); + + let repaired = repair_order(opt, &ops); + + assert_eq!(repaired.ordering, vec![0, 1]); + assert_eq!(repaired.strategy, unfused(vec![0, 1])); + } + + /// A single strategy with a valid order passes through untouched. + #[test] + fn single_strategy_with_valid_order_passes_through() { + let ops = vec![add(1, 2, 10), add(10, 3, 11)]; + let opt = BlockOptimization::new(fused(vec![0, 1]), vec![0, 1]); + + let repaired = repair_order(opt, &ops); + + assert_eq!(repaired.ordering, vec![0, 1]); + assert_eq!(repaired.strategy, fused(vec![0, 1])); + } + + /// Build a [StreamOptimizer] with one [TestOptimizationBuilder] per pattern. + fn optimizer(patterns: Vec>) -> StreamOptimizer { + let builders = patterns + .into_iter() + .enumerate() + .map(|(i, p)| { + Box::new(TestOptimizationBuilder::new(i, p)) + as Box> + }) + .collect(); + StreamOptimizer::new(builders) + } + + /// A dependent op arrives while the block cap is full of un-mergeable (ready) blocks: the + /// segment stops, the op is deferred, and the blocks registered so far still optimize. + #[test] + fn dependent_op_at_max_blocks_stops_when_no_merge_possible() { + let a1 = add(100, 101, 102); + let a2 = add(102, 103, 104); + let b1 = add(200, 201, 202); + let b2 = add(202, 203, 204); + let c = add(104, 204, 105); // Reads A's 104 and B's 204. + + let mut opt = optimizer(vec![ + vec![a1.clone(), a2.clone()], + vec![b1.clone(), b2.clone()], + ]); + opt.max_blocks = Some(2); + + for op in [&a1, &a2, &b1, &b2] { + opt.register(op); + } + assert!(!opt.stopped); + + // Both blocks hold their own ready fusion, so force-merging cannot free a slot. + opt.register(&c); + assert!(opt.stopped); + + let ops = vec![a1, a2, b1, b2]; + let res = opt.optimize(&ops); + assert_eq!(res.ordering, vec![0, 1, 2, 3]); + assert_eq!( + res.strategy, + composed(vec![ + fused_with(0, vec![0, 1], 2), + fused_with(1, vec![2, 3], 2), + ]), + ); + } + + /// The defensive cycle check in [on_dependent_op](StreamOptimizer::on_dependent_op): an op + /// that reads a tensor a block frees AND a tensor the same block produces would have to run + /// both before and after that block. Unreachable from a valid stream (it reads a freed + /// tensor), so the path is driven directly. + #[test] + fn dependent_op_creating_a_cycle_stops_segment() { + let a = add_rw(1, 2, 3); // Block A: frees 1, produces 3. + let b = add(10, 11, 12); // Block B: unrelated. + let n = add(1, 3, 20); // Reads 1 (freed by A) and 3 (produced by A). + + let mut opt = optimizer(vec![]); + opt.max_blocks = None; + opt.register(&a); + opt.register(&b); + assert_eq!(opt.blocks.len(), 2); + + opt.on_dependent_op(&n); + + assert!(opt.stopped); + assert_eq!(opt.blocks.len(), 2, "the cyclic block was rejected"); + } +} diff --git a/crates/burn-fusion/src/search/optimization/tests.rs b/crates/burn-fusion/src/search/optimization/tests.rs new file mode 100644 index 0000000..2b46cb5 --- /dev/null +++ b/crates/burn-fusion/src/search/optimization/tests.rs @@ -0,0 +1,387 @@ +//! Tests for the [StreamOptimizer] and the underlying [BlocksOptimizer]. +//! +//! These tests focus on whether the optimizer reorders operations across +//! independent blocks to find better fusion opportunities. Two operations end +//! up in distinct blocks when their tensor ids are disjoint — meaning they +//! have no data dependencies and can safely be reordered. + +use std::sync::Arc; + +use super::StreamOptimizer; +use crate::search::testing::add; +use crate::stream::execution::tests::{TestOptimization, TestOptimizationBuilder}; +use crate::{OperationFuser, search::BlockOptimization, stream::store::ExecutionStrategy}; +use burn_ir::OperationIr; + +/// Build a [StreamOptimizer] with one [TestOptimizationBuilder] per pattern. +fn optimizer(patterns: Vec>) -> StreamOptimizer { + let builders = patterns + .into_iter() + .enumerate() + .map(|(i, p)| { + Box::new(TestOptimizationBuilder::new(i, p)) + as Box> + }) + .collect(); + StreamOptimizer::new(builders) +} + +/// Register every op in the stream and return the resulting [BlockOptimization]. +fn run( + ops: &[OperationIr], + patterns: Vec>, +) -> BlockOptimization { + let mut opt = optimizer(patterns); + for op in ops { + opt.register(op); + } + opt.optimize(ops) +} + +// --- Strategy constructors (make expected values short and readable) ------- + +fn optimization( + builder_id: usize, + size: usize, + ordering: Vec, + score: u64, +) -> ExecutionStrategy { + ExecutionStrategy::Optimization { + opt: TestOptimization::new(builder_id, size), + ordering: Arc::new(ordering), + score, + } +} + +fn operations(ordering: Vec) -> ExecutionStrategy { + ExecutionStrategy::Operations { + ordering: Arc::new(ordering), + } +} + +fn composed( + parts: Vec>, +) -> ExecutionStrategy { + ExecutionStrategy::Composed(parts.into_iter().map(Box::new).collect()) +} + +// --------------------------------------------------------------------------- +// Baseline: a single block with a builder whose pattern matches the stream. +// --------------------------------------------------------------------------- + +#[test] +fn single_block_fuses_full_stream() { + let a1 = add(100, 101, 102); + let a2 = add(102, 103, 104); + let ops = vec![a1.clone(), a2.clone()]; + + let res = run(&ops, vec![vec![a1, a2]]); + + assert_eq!(res.ordering, vec![0, 1]); + assert_eq!(res.strategy, optimization(0, 2, vec![0, 1], 2)); +} + +// --------------------------------------------------------------------------- +// Two independent blocks, registered contiguously (all A then all B). +// Both should fuse — no reordering required. +// --------------------------------------------------------------------------- + +#[test] +fn contiguous_independent_blocks_fuse_both() { + let a1 = add(100, 101, 102); + let a2 = add(102, 103, 104); + let b1 = add(200, 201, 202); + let b2 = add(202, 203, 204); + let ops = vec![a1.clone(), a2.clone(), b1.clone(), b2.clone()]; + + let res = run( + &ops, + vec![vec![a1.clone(), a2.clone()], vec![b1.clone(), b2.clone()]], + ); + + assert_eq!(res.ordering, vec![0, 1, 2, 3]); + assert_eq!( + res.strategy, + composed(vec![ + optimization(0, 2, vec![0, 1], 2), + optimization(1, 2, vec![2, 3], 2), + ]), + ); +} + +// --------------------------------------------------------------------------- +// Two independent blocks *interleaved* in the stream. +// +// Block A has no shared tensors with block B, so the two blocks have no data +// dependency — reordering them is safe. Both blocks fuse, each block's ops +// grouped contiguously. +// --------------------------------------------------------------------------- + +#[test] +fn interleaved_independent_blocks_fuse_both() { + let a1 = add(100, 101, 102); + let a2 = add(102, 103, 104); + let b1 = add(200, 201, 202); + let b2 = add(202, 203, 204); + // Stream is A1, B1, A2, B2 — each pair belongs to its own independent block. + let ops = vec![a1.clone(), b1.clone(), a2.clone(), b2.clone()]; + + let res = run( + &ops, + vec![vec![a1.clone(), a2.clone()], vec![b1.clone(), b2.clone()]], + ); + + assert_eq!(res.ordering, vec![0, 2, 1, 3]); + assert_eq!( + res.strategy, + composed(vec![ + optimization(0, 2, vec![0, 2], 2), + optimization(1, 2, vec![1, 3], 2), + ]), + ); +} + +// --------------------------------------------------------------------------- +// Interleaved independent blocks with only ONE builder that matches the +// *combined*, reordered stream [A1, A2, B1, B2]. The merging pass combines +// the two blocks so the builder sees its full pattern and fuses everything +// into one optimization. +// --------------------------------------------------------------------------- + +#[test] +fn interleaved_blocks_with_merged_pattern_fuses_full() { + let a1 = add(100, 101, 102); + let a2 = add(102, 103, 104); + let b1 = add(200, 201, 202); + let b2 = add(202, 203, 204); + let ops = vec![a1.clone(), b1.clone(), a2.clone(), b2.clone()]; + + // Single builder whose pattern is the combined, reordered stream. + let res = run( + &ops, + vec![vec![a1.clone(), a2.clone(), b1.clone(), b2.clone()]], + ); + + assert_eq!(res.ordering, vec![0, 2, 1, 3]); + assert_eq!(res.strategy, optimization(0, 4, vec![0, 2, 1, 3], 4)); +} + +// --------------------------------------------------------------------------- +// Interleaved independent blocks, only one has a matching builder. Block A +// fuses; block B's ops survive as un-fused Operations in the composed +// strategy. +// --------------------------------------------------------------------------- + +#[test] +fn interleaved_blocks_single_builder_preserves_other_ops() { + let a1 = add(100, 101, 102); + let a2 = add(102, 103, 104); + let b1 = add(200, 201, 202); + let b2 = add(202, 203, 204); + let ops = vec![a1.clone(), b1.clone(), a2.clone(), b2.clone()]; + + // Only a builder for block A — B's ops survive as Operations. + let res = run(&ops, vec![vec![a1.clone(), a2.clone()]]); + + assert_eq!(res.ordering, vec![0, 2, 1, 3]); + assert_eq!( + res.strategy, + composed(vec![ + optimization(0, 2, vec![0, 2], 2), + operations(vec![1, 3]), + ]), + ); +} + +// --------------------------------------------------------------------------- +// Within a single block, operations are registered in stream order. A builder +// whose pattern would only match a *reordered* version of the stream cannot +// fuse. This documents that reordering does NOT happen inside a block (unlike +// across blocks). +// --------------------------------------------------------------------------- + +#[test] +fn within_block_order_is_not_reordered() { + let a1 = add(100, 101, 102); + let a2 = add(102, 103, 104); + let ops = vec![a1.clone(), a2.clone()]; + + // Builder wants [a2, a1] — only reachable if we reorder within the block. + let res = run(&ops, vec![vec![a2.clone(), a1.clone()]]); + + assert_eq!(res.ordering, vec![0, 1]); + assert_eq!(res.strategy, operations(vec![0, 1])); +} + +// --------------------------------------------------------------------------- +// Lazy-optimization semantics. +// +// `still_optimizing()` is a "keep feeding me" signal for the caller: it +// returns true if at least one builder might still match a longer pattern. +// `optimize()` is the commit point — it picks the best *ready* fusion, even +// if other builders are still open (this is what happens on a forced sync). +// --------------------------------------------------------------------------- + +/// A builder whose pattern is longer than the stream so far. The builder's +/// actual ops match the pattern prefix, so the builder stays Open and +/// `still_optimizing` keeps signaling "feed me more". A forced `optimize` +/// (simulating a sync) falls back to unfused operations because no builder +/// has reached `ready`. +#[test] +fn pattern_longer_than_stream_stays_optimizing() { + let a1 = add(100, 101, 102); + let a2 = add(102, 103, 104); + let a3 = add(104, 105, 106); + + let mut opt = optimizer(vec![vec![a1.clone(), a2.clone(), a3.clone()]]); + opt.register(&a1); + opt.register(&a2); + + assert!(opt.still_optimizing(), "builder still waiting for a3"); + + let ops = vec![a1, a2]; + let res = opt.optimize(&ops); + assert_eq!(res.ordering, vec![0, 1]); + assert_eq!(res.strategy, operations(vec![0, 1])); +} + +/// One builder is ready at two ops, another is still open waiting for a +/// third. `still_optimizing` stays true so the caller has the option to keep +/// feeding and pick up the longer fusion. A forced commit takes the ready +/// one. +#[test] +fn ready_and_open_coexist_keeps_optimizing() { + let a1 = add(100, 101, 102); + let a2 = add(102, 103, 104); + let a3 = add(104, 105, 106); + + let mut opt = optimizer(vec![ + vec![a1.clone(), a2.clone()], // ready after 2 ops + vec![a1.clone(), a2.clone(), a3.clone()], // still open after 2 ops + ]); + opt.register(&a1); + opt.register(&a2); + + assert!(opt.still_optimizing(), "builder 2 still waiting for a3"); + + let ops = vec![a1, a2]; + let res = opt.optimize(&ops); + assert_eq!(res.ordering, vec![0, 1]); + assert_eq!(res.strategy, optimization(0, 2, vec![0, 1], 2)); +} + +/// Multiple builders are ready at commit time — the highest-scoring fusion +/// wins. +#[test] +fn commit_picks_highest_score_ready_builder() { + let a1 = add(100, 101, 102); + let a2 = add(102, 103, 104); + let a3 = add(104, 105, 106); + + let mut opt = optimizer(vec![ + vec![a1.clone(), a2.clone()], // score 2 + vec![a1.clone(), a2.clone(), a3.clone()], // score 3 + ]); + opt.register(&a1); + opt.register(&a2); + opt.register(&a3); + + // Every builder is closed (everyone has seen at least as many ops as its + // pattern) — there's nothing more to gain by feeding. + assert!(!opt.still_optimizing()); + + let ops = vec![a1, a2, a3]; + let res = opt.optimize(&ops); + assert_eq!(res.ordering, vec![0, 1, 2]); + assert_eq!(res.strategy, optimization(1, 3, vec![0, 1, 2], 3)); +} + +/// Every builder is closed with `ready=false` (the stream diverged from +/// every pattern). `still_optimizing` flips to false — caller should +/// commit. Forced commit returns unfused operations. +#[test] +fn all_closed_without_ready_bails_out() { + let a1 = add(100, 101, 102); + let a2 = add(102, 103, 104); + // Builder wants [a1, b1]: stream diverges at the second op. + let b1 = add(200, 201, 202); + + let mut opt = optimizer(vec![vec![a1.clone(), b1.clone()]]); + opt.register(&a1); + opt.register(&a2); + + assert!(!opt.still_optimizing(), "mismatch closes the builder"); + + let ops = vec![a1, a2]; + let res = opt.optimize(&ops); + assert_eq!(res.ordering, vec![0, 1]); + assert_eq!(res.strategy, operations(vec![0, 1])); +} + +/// A single pattern that spans two independent blocks, but the last op of +/// the pattern hasn't arrived yet. The merge succeeds (builder stays Open in +/// the merged block → `merge` returns true), so `still_optimizing` stays +/// true. A forced commit falls back to unfused operations in the merged +/// block's registration order. +#[test] +fn pattern_straddles_incomplete_merged_blocks() { + let a1 = add(100, 101, 102); + let a2 = add(102, 103, 104); + let b1 = add(200, 201, 202); + let b2 = add(202, 203, 204); + + // Stream is [a1, b1, a2] — builder needs b2 too. + let mut opt = optimizer(vec![vec![a1.clone(), a2.clone(), b1.clone(), b2.clone()]]); + opt.register(&a1); + opt.register(&b1); + opt.register(&a2); + + assert!(opt.still_optimizing(), "builder waiting for b2"); + + let ops = vec![a1, b1, a2]; + let res = opt.optimize(&ops); + // After merging, the merged block's registration order is + // [a1 (pos 0), a2 (pos 2), b1 (pos 1)]. + assert_eq!(res.ordering, vec![0, 2, 1]); + assert_eq!(res.strategy, operations(vec![0, 2, 1])); +} + +// --------------------------------------------------------------------------- +// Dependent blocks. +// +// Block A and block B each hold their own *ready* fusion. Merging them would +// hide one of the two fusions, so `Block::merge` refuses. Op C then reads an +// output of A and an output of B, so it references two un-mergeable blocks. +// +// Previously this stopped the search entirely. Now C gets its own block that +// *depends on* A and B; the search continues and every op is covered. C's +// (unfused) block is emitted after both blocks it depends on. +// --------------------------------------------------------------------------- + +#[test] +fn dependent_block_created_instead_of_stopping() { + let a1 = add(100, 101, 102); + let a2 = add(102, 103, 104); // block A produces 104 + let b1 = add(200, 201, 202); + let b2 = add(202, 203, 204); // block B produces 204 + let c = add(104, 204, 105); // reads A's 104 and B's 204 + + let ops = vec![a1.clone(), a2.clone(), b1.clone(), b2.clone(), c.clone()]; + + let res = run( + &ops, + vec![vec![a1.clone(), a2.clone()], vec![b1.clone(), b2.clone()]], + ); + + // All five ops are covered (the search did not stop at C) and C runs after + // both A and B — a valid topological order. + assert_eq!(res.ordering, vec![0, 1, 2, 3, 4]); + assert_eq!( + res.strategy, + composed(vec![ + optimization(0, 2, vec![0, 1], 2), + optimization(1, 2, vec![2, 3], 2), + operations(vec![4]), + ]), + ); +} diff --git a/crates/burn-fusion/src/search/testing.rs b/crates/burn-fusion/src/search/testing.rs new file mode 100644 index 0000000..c964764 --- /dev/null +++ b/crates/burn-fusion/src/search/testing.rs @@ -0,0 +1,41 @@ +//! Shared IR constructors for the search tests. + +use burn_backend::{DType, Shape}; +use burn_ir::{BinaryOpIr, NumericOperationIr, OperationIr, TensorId, TensorIr, TensorStatus}; + +/// A `[32, 32]` float tensor with the given id, read with `ReadOnly` status. +pub(crate) fn tensor(id: u64) -> TensorIr { + TensorIr { + id: TensorId::new(id), + shape: Shape::new([32, 32]), + status: TensorStatus::ReadOnly, + dtype: DType::F32, + } +} + +/// An add operation reading `lhs` and `rhs` and producing `out`. +pub(crate) fn add(lhs: u64, rhs: u64, out: u64) -> OperationIr { + OperationIr::NumericFloat( + DType::F32, + NumericOperationIr::Add(BinaryOpIr { + lhs: tensor(lhs), + rhs: tensor(rhs), + out: tensor(out), + }), + ) +} + +/// Like [add] but reads `lhs` with [ReadWrite](TensorStatus::ReadWrite) status — it is the last +/// use and frees the tensor. +pub(crate) fn add_rw(lhs: u64, rhs: u64, out: u64) -> OperationIr { + let mut lhs = tensor(lhs); + lhs.status = TensorStatus::ReadWrite; + OperationIr::NumericFloat( + DType::F32, + NumericOperationIr::Add(BinaryOpIr { + lhs, + rhs: tensor(rhs), + out: tensor(out), + }), + ) +} diff --git a/crates/burn-fusion/src/server.rs b/crates/burn-fusion/src/server.rs new file mode 100644 index 0000000..790ea82 --- /dev/null +++ b/crates/burn-fusion/src/server.rs @@ -0,0 +1,166 @@ +use std::sync::Arc; + +use crate::{ + FusionBackend, FusionRuntime, UnfusedOp, + stream::{MultiStream, StreamId}, +}; +use burn_backend::{TensorData, backend::ExecutionError}; +use burn_ir::{HandleContainer, OperationIr, TensorId, TensorIr}; +use burn_std::{CommunicationId, stub::RwLock}; +use hashbrown::HashSet; + +pub(crate) struct FusionUtilities { + // Used in client using a downcast. + #[allow(dead_code)] + pub(crate) initialized_comms: RwLock>, +} + +pub struct FusionServer { + streams: MultiStream, + pub(crate) handles: HandleContainer, + pub(crate) utilities: Arc, +} + +impl FusionServer +where + R: FusionRuntime, +{ + pub fn new(device: R::FusionDevice, utilities: FusionUtilities) -> Self { + Self { + streams: MultiStream::new(device.clone()), + handles: HandleContainer::new(), + utilities: Arc::new(utilities), + } + } + + pub fn register(&mut self, stream: StreamId, repr: OperationIr, operation: UnfusedOp) { + self.streams + .register(stream, repr, operation, &mut self.handles) + } + + pub fn tag_shared_view(&mut self, src_stream: StreamId, src: TensorId, dst: TensorId) { + self.streams + .tag_shared_view(src_stream, src, dst, &mut self.handles) + } + + pub fn drain_stream(&mut self, id: StreamId) { + self.streams.drain(&mut self.handles, id) + } + + pub fn read_float(&mut self, tensor: TensorIr, id: StreamId) -> B::FloatTensorPrimitive + where + B: FusionBackend, + { + // Make sure all registered operations are executed. + // The underlying backend can still be async. + self.drain_stream(id); + let tensor_float = self.handles.get_float_tensor::(&tensor); + self.streams.mark_read(id, &tensor, &self.handles); + tensor_float + } + + pub fn read_int(&mut self, tensor: TensorIr, id: StreamId) -> B::IntTensorPrimitive + where + B: FusionBackend, + { + // Make sure all registered operations are executed. + // The underlying backend can still be async. + self.drain_stream(id); + let tensor_int = self.handles.get_int_tensor::(&tensor); + self.streams.mark_read(id, &tensor, &self.handles); + tensor_int + } + + pub fn read_bool(&mut self, tensor: TensorIr, id: StreamId) -> B::BoolTensorPrimitive + where + B: FusionBackend, + { + // Make sure all registered operations are executed. + // The underlying backend can still be async. + self.drain_stream(id); + let tensor_bool = self.handles.get_bool_tensor::(&tensor); + self.streams.mark_read(id, &tensor, &self.handles); + tensor_bool + } + + pub fn read_quantized( + &mut self, + tensor: TensorIr, + id: StreamId, + ) -> B::QuantizedTensorPrimitive + where + B: FusionBackend, + { + // Make sure all registered operations are executed. + // The underlying backend can still be async. + self.drain_stream(id); + let tensor_q = self.handles.get_quantized_tensor::(&tensor); + self.streams.mark_read(id, &tensor, &self.handles); + tensor_q + } + + pub fn float_data( + &mut self, + tensor: TensorIr, + id: StreamId, + ) -> impl Future> + Send + use + where + B: FusionBackend, + { + B::float_into_data(self.read_float::(tensor, id)) + } + + pub fn int_data( + &mut self, + tensor: TensorIr, + id: StreamId, + ) -> impl Future> + Send + use + where + B: FusionBackend, + { + B::int_into_data(self.read_int::(tensor, id)) + } + + pub fn bool_data( + &mut self, + tensor: TensorIr, + id: StreamId, + ) -> impl Future> + Send + use + where + B: FusionBackend, + { + B::bool_into_data(self.read_bool::(tensor, id)) + } + + pub fn quantized_data( + &mut self, + tensor: TensorIr, + id: StreamId, + ) -> impl Future> + Send + use + where + B: FusionBackend, + { + B::q_into_data(self.read_quantized::(tensor, id)) + } + + pub fn resolve_server_float(&mut self, tensor: &TensorIr) -> B::FloatTensorPrimitive + where + B: FusionBackend, + { + self.handles.get_float_tensor::(tensor) + } + + pub fn resolve_server_int(&mut self, tensor: &TensorIr) -> B::IntTensorPrimitive + where + B: FusionBackend, + { + self.handles.get_int_tensor::(tensor) + } + + pub fn resolve_server_bool(&mut self, tensor: &TensorIr) -> B::BoolTensorPrimitive + where + B: FusionBackend, + { + self.handles.get_bool_tensor::(tensor) + } +} diff --git a/crates/burn-fusion/src/stream/base.rs b/crates/burn-fusion/src/stream/base.rs new file mode 100644 index 0000000..39336fd --- /dev/null +++ b/crates/burn-fusion/src/stream/base.rs @@ -0,0 +1 @@ +pub use burn_backend::StreamId; diff --git a/crates/burn-fusion/src/stream/context.rs b/crates/burn-fusion/src/stream/context.rs new file mode 100644 index 0000000..a8c38e6 --- /dev/null +++ b/crates/burn-fusion/src/stream/context.rs @@ -0,0 +1,1792 @@ +use burn_backend::{Shape, Slice}; +use burn_ir::*; +use hashbrown::HashMap; + +/// The context contains the relative graph tensor mapping so that a relative tensor id can be +/// mapped to an existing tensor that can be fetched and updated with the +/// [handle container](HandleContainer). +/// +/// It also contains all scalar values, which can change even for the same graph. They are sorted +/// in the order in which they appear in the graph. +pub struct Context { + /// The tensor mapping where local tensor id points to the updated tensor representation. + pub tensors: HashMap, + /// Handle container to retrieve tensors based on their representation. + pub handles: HandleContainer, + /// Scalars found in the graph in the order they appeared. + pub scalars: HashMap, + /// Shape mapping from relative shape ids to global (real) shape ids. + pub shapes_relative2global: HashMap, + /// Concrete slice ranges found in the graph, indexed by the placeholder id assigned during + /// relativization (the value carried in a relativized range's `start` field). Ranges can change + /// for the same relative graph — like scalars — so they are kept out of the relative form and + /// rebound per invocation. + pub ranges: Vec, +} + +impl Context { + /// Fork the context into an independent owned copy. Used by autotuning to give each + /// benchmark run a sandbox — mutations stay local until the caller merges new handles + /// back. `HandleContainer::fork` clones the id→handle map; actual GPU buffers are + /// refcounted so only the map itself is duplicated. + pub fn fork(&self) -> Self { + Self { + tensors: self.tensors.clone(), + handles: self.handles.fork(), + scalars: self.scalars.clone(), + shapes_relative2global: self.shapes_relative2global.clone(), + ranges: self.ranges.clone(), + } + } +} + +#[derive(Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord, Debug)] +/// Scalar unique identifier. +pub struct ScalarId { + /// The value. + pub value: u64, +} + +pub(crate) struct OperationConverter { + tensors_relative2global: HashMap, + tensors_global2relative: HashMap, + shapes_global2relative: HashMap, + shapes_relative2global: HashMap, + scalars: HashMap, + ranges: Vec, +} + +impl Default for OperationConverter { + fn default() -> Self { + let mut val = Self { + tensors_relative2global: Default::default(), + tensors_global2relative: Default::default(), + shapes_global2relative: Default::default(), + shapes_relative2global: Default::default(), + scalars: Default::default(), + ranges: Default::default(), + }; + + // global 1 is always shape id 0. + val.shapes_global2relative.insert(1, 0); + val.shapes_relative2global.insert(0, 1); + + val + } +} + +/// RAII guard that temporarily lends three [`OperationConverter`] fields plus a +/// [`HandleContainer`] into a single owned [`Context`]. On drop, the guard moves every +/// field back into its original location. +pub(crate) struct ContextGuard<'a, H> { + context: Option>, + converter: &'a mut OperationConverter, + handles: &'a mut HandleContainer, +} + +impl<'a, H> ContextGuard<'a, H> { + /// Move the converter's per-block exposed state and the server's handle container into a + /// fresh [`Context`]. The originals are left holding `Default` placeholders; the guard + /// will swap them back on drop. + pub(crate) fn new( + converter: &'a mut OperationConverter, + handles: &'a mut HandleContainer, + ) -> Self { + let context = Context { + tensors: core::mem::take(&mut converter.tensors_relative2global), + scalars: core::mem::take(&mut converter.scalars), + shapes_relative2global: core::mem::take(&mut converter.shapes_relative2global), + ranges: core::mem::take(&mut converter.ranges), + handles: core::mem::take(handles), + }; + + Self { + context: Some(context), + converter, + handles, + } + } +} + +impl core::ops::Deref for ContextGuard<'_, H> { + type Target = Context; + + fn deref(&self) -> &Self::Target { + self.context.as_ref().expect("context guard is alive") + } +} + +impl core::ops::DerefMut for ContextGuard<'_, H> { + fn deref_mut(&mut self) -> &mut Self::Target { + self.context.as_mut().expect("context guard is alive") + } +} + +impl Drop for ContextGuard<'_, H> { + fn drop(&mut self) { + if let Some(ctx) = self.context.take() { + self.converter.tensors_relative2global = ctx.tensors; + self.converter.scalars = ctx.scalars; + self.converter.shapes_relative2global = ctx.shapes_relative2global; + self.converter.ranges = ctx.ranges; + *self.handles = ctx.handles; + } + } +} + +pub(crate) trait RelativeOps { + /// Convert (usually an [`OperationIr`]) to a relative form. + /// + /// The id and the shape of tensors will be computed relative to existing + /// operations in the queue. We do this because we want to fuse operations + /// that have similar shapes, but we do not care about the exact values. + /// + /// Similar we do not care about the exact ids of the tensor, but about their + /// relative ids (how close they are in the operation queue) + fn to_relative(&self, converter: &mut OperationConverter) -> Self; +} + +impl OperationConverter { + pub(crate) fn clear(&mut self) { + self.tensors_relative2global.clear(); + self.tensors_global2relative.clear(); + + self.shapes_global2relative.clear(); + self.shapes_relative2global.clear(); + + // global 1 is always shape id 0. + self.shapes_global2relative.insert(1, 0); + self.shapes_relative2global.insert(0, 1); + + self.scalars.clear(); + self.ranges.clear(); + } + + /// Relativize a slice range: stash the concrete bounds (they vary per invocation) and return a + /// placeholder whose `start` carries the binding id, mirroring how scalars become + /// `ScalarIr::UInt(placeholder)`. The relative form is thus invariant to the actual bounds, so + /// graphs that differ only in slice ranges still match — and graph replay restores the concrete + /// range from [`Context::ranges`] / `GraphBindings::ranges`. + fn relative_range(&mut self, range: &Slice) -> Slice { + let id = self.ranges.len(); + self.ranges.push(*range); + Slice { + start: id as isize, + end: None, + step: 1, + } + } +} + +impl RelativeOps for OperationIr { + fn to_relative(&self, converter: &mut OperationConverter) -> Self { + match self { + OperationIr::BaseFloat(ops) => OperationIr::BaseFloat(ops.to_relative(converter)), + OperationIr::BaseInt(ops) => OperationIr::BaseInt(ops.to_relative(converter)), + OperationIr::BaseBool(ops) => OperationIr::BaseBool(ops.to_relative(converter)), + OperationIr::NumericFloat(dtype, ops) => { + OperationIr::NumericFloat(*dtype, ops.to_relative(converter)) + } + OperationIr::NumericInt(dtype, ops) => { + OperationIr::NumericInt(*dtype, ops.to_relative(converter)) + } + OperationIr::Bool(ops) => OperationIr::Bool(ops.to_relative(converter)), + OperationIr::Int(ops) => OperationIr::Int(ops.to_relative(converter)), + OperationIr::Float(dtype, ops) => { + OperationIr::Float(*dtype, ops.to_relative(converter)) + } + OperationIr::Module(ops) => OperationIr::Module(ops.to_relative(converter)), + OperationIr::Custom(ops) => OperationIr::Custom(ops.to_relative(converter)), + OperationIr::Init(ops) => OperationIr::Init(ops.to_relative(converter)), + OperationIr::Drop(tensor) => OperationIr::Drop(tensor.to_relative(converter)), + OperationIr::Distributed(ops) => OperationIr::Distributed(ops.to_relative(converter)), + OperationIr::Activation(ops) => OperationIr::Activation(ops.to_relative(converter)), + } + } +} + +impl RelativeOps for ModuleOperationIr { + fn to_relative(&self, converter: &mut OperationConverter) -> Self { + match self { + ModuleOperationIr::Embedding(desc) => ModuleOperationIr::Embedding(EmbeddingOpIr { + weights: desc.weights.to_relative(converter), + indices: desc.indices.to_relative(converter), + out: desc.out.to_relative(converter), + }), + ModuleOperationIr::EmbeddingBackward(desc) => { + ModuleOperationIr::EmbeddingBackward(EmbeddingBackwardOpIr { + weights: desc.weights.to_relative(converter), + out_grad: desc.out_grad.to_relative(converter), + indices: desc.indices.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::Linear(desc) => ModuleOperationIr::Linear(LinearOpIr { + x: desc.x.to_relative(converter), + weight: desc.weight.to_relative(converter), + bias: desc.bias.as_ref().map(|t| t.to_relative(converter)), + out: desc.out.to_relative(converter), + }), + ModuleOperationIr::LinearXBackward(desc) => { + ModuleOperationIr::LinearXBackward(LinearXBackwardOpIr { + weight: desc.weight.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::LinearWeightBackward(desc) => { + ModuleOperationIr::LinearWeightBackward(LinearWeightBackwardOpIr { + x: desc.x.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::LinearBiasBackward(desc) => { + ModuleOperationIr::LinearBiasBackward(LinearBiasBackwardOpIr { + output_grad: desc.output_grad.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::Conv1d(desc) => ModuleOperationIr::Conv1d(Conv1dOpIr { + x: desc.x.to_relative(converter), + weight: desc.weight.to_relative(converter), + bias: desc.bias.as_ref().map(|t| t.to_relative(converter)), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }), + ModuleOperationIr::Conv1dXBackward(desc) => { + ModuleOperationIr::Conv1dXBackward(Conv1dXBackwardOpIr { + x: desc.x.to_relative(converter), + weight: desc.weight.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::Conv1dWeightBackward(desc) => { + ModuleOperationIr::Conv1dWeightBackward(Conv1dWeightBackwardOpIr { + x: desc.x.to_relative(converter), + weight: desc.weight.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::Conv1dBiasBackward(desc) => { + ModuleOperationIr::Conv1dBiasBackward(Conv1dBiasBackwardOpIr { + x: desc.x.to_relative(converter), + bias: desc.bias.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::Conv2d(desc) => ModuleOperationIr::Conv2d(Conv2dOpIr { + x: desc.x.to_relative(converter), + weight: desc.weight.to_relative(converter), + bias: desc.bias.as_ref().map(|t| t.to_relative(converter)), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }), + ModuleOperationIr::Conv2dXBackward(desc) => { + ModuleOperationIr::Conv2dXBackward(Conv2dXBackwardOpIr { + x: desc.x.to_relative(converter), + weight: desc.weight.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::Conv2dWeightBackward(desc) => { + ModuleOperationIr::Conv2dWeightBackward(Conv2dWeightBackwardOpIr { + x: desc.x.to_relative(converter), + weight: desc.weight.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::Conv2dBiasBackward(desc) => { + ModuleOperationIr::Conv2dBiasBackward(Conv2dBiasBackwardOpIr { + x: desc.x.to_relative(converter), + bias: desc.bias.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::Conv3d(desc) => ModuleOperationIr::Conv3d(Conv3dOpIr { + x: desc.x.to_relative(converter), + weight: desc.weight.to_relative(converter), + bias: desc.bias.as_ref().map(|t| t.to_relative(converter)), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }), + ModuleOperationIr::Conv3dXBackward(desc) => { + ModuleOperationIr::Conv3dXBackward(Conv3dXBackwardOpIr { + x: desc.x.to_relative(converter), + weight: desc.weight.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::Conv3dWeightBackward(desc) => { + ModuleOperationIr::Conv3dWeightBackward(Conv3dWeightBackwardOpIr { + x: desc.x.to_relative(converter), + weight: desc.weight.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::Conv3dBiasBackward(desc) => { + ModuleOperationIr::Conv3dBiasBackward(Conv3dBiasBackwardOpIr { + x: desc.x.to_relative(converter), + bias: desc.bias.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::DeformableConv2d(desc) => { + ModuleOperationIr::DeformableConv2d(Box::new(DeformConv2dOpIr { + x: desc.x.to_relative(converter), + offset: desc.offset.to_relative(converter), + weight: desc.weight.to_relative(converter), + mask: desc.mask.as_ref().map(|t| t.to_relative(converter)), + bias: desc.bias.as_ref().map(|t| t.to_relative(converter)), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + })) + } + ModuleOperationIr::DeformableConv2dBackward(desc) => { + ModuleOperationIr::DeformableConv2dBackward(Box::new(DeformConv2dBackwardOpIr { + x: desc.x.to_relative(converter), + offset: desc.offset.to_relative(converter), + weight: desc.weight.to_relative(converter), + mask: desc.mask.as_ref().map(|t| t.to_relative(converter)), + bias: desc.bias.as_ref().map(|t| t.to_relative(converter)), + out_grad: desc.out_grad.to_relative(converter), + options: desc.options.clone(), + input_grad: desc.input_grad.to_relative(converter), + offset_grad: desc.offset_grad.to_relative(converter), + weight_grad: desc.weight_grad.to_relative(converter), + mask_grad: desc.mask_grad.as_ref().map(|t| t.to_relative(converter)), + bias_grad: desc.bias_grad.as_ref().map(|t| t.to_relative(converter)), + })) + } + ModuleOperationIr::ConvTranspose1d(desc) => { + ModuleOperationIr::ConvTranspose1d(ConvTranspose1dOpIr { + x: desc.x.to_relative(converter), + weight: desc.weight.to_relative(converter), + bias: desc.bias.as_ref().map(|t| t.to_relative(converter)), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::ConvTranspose2d(desc) => { + ModuleOperationIr::ConvTranspose2d(ConvTranspose2dOpIr { + x: desc.x.to_relative(converter), + weight: desc.weight.to_relative(converter), + bias: desc.bias.as_ref().map(|t| t.to_relative(converter)), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::ConvTranspose3d(desc) => { + ModuleOperationIr::ConvTranspose3d(ConvTranspose3dOpIr { + x: desc.x.to_relative(converter), + weight: desc.weight.to_relative(converter), + bias: desc.bias.as_ref().map(|t| t.to_relative(converter)), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::AvgPool1d(desc) => ModuleOperationIr::AvgPool1d(AvgPool1dOpIr { + x: desc.x.to_relative(converter), + kernel_size: desc.kernel_size, + stride: desc.stride, + padding: desc.padding, + count_include_pad: desc.count_include_pad, + ceil_mode: desc.ceil_mode, + out: desc.out.to_relative(converter), + }), + ModuleOperationIr::AvgPool2d(desc) => ModuleOperationIr::AvgPool2d(AvgPool2dOpIr { + x: desc.x.to_relative(converter), + kernel_size: desc.kernel_size, + stride: desc.stride, + padding: desc.padding, + count_include_pad: desc.count_include_pad, + ceil_mode: desc.ceil_mode, + out: desc.out.to_relative(converter), + }), + ModuleOperationIr::AvgPool1dBackward(desc) => { + ModuleOperationIr::AvgPool1dBackward(AvgPool1dBackwardOpIr { + x: desc.x.to_relative(converter), + grad: desc.grad.to_relative(converter), + kernel_size: desc.kernel_size, + stride: desc.stride, + padding: desc.padding, + count_include_pad: desc.count_include_pad, + ceil_mode: desc.ceil_mode, + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::AvgPool2dBackward(desc) => { + ModuleOperationIr::AvgPool2dBackward(AvgPool2dBackwardOpIr { + x: desc.x.to_relative(converter), + grad: desc.grad.to_relative(converter), + kernel_size: desc.kernel_size, + stride: desc.stride, + padding: desc.padding, + count_include_pad: desc.count_include_pad, + ceil_mode: desc.ceil_mode, + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::AdaptiveAvgPool1d(desc) => { + ModuleOperationIr::AdaptiveAvgPool1d(AdaptiveAvgPool1dOpIr { + x: desc.x.to_relative(converter), + output_size: desc.output_size, + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::AdaptiveAvgPool2d(desc) => { + ModuleOperationIr::AdaptiveAvgPool2d(AdaptiveAvgPool2dOpIr { + x: desc.x.to_relative(converter), + output_size: desc.output_size, + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::AdaptiveAvgPool1dBackward(desc) => { + ModuleOperationIr::AdaptiveAvgPool1dBackward(AdaptiveAvgPool1dBackwardOpIr { + x: desc.x.to_relative(converter), + grad: desc.grad.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::AdaptiveAvgPool2dBackward(desc) => { + ModuleOperationIr::AdaptiveAvgPool2dBackward(AdaptiveAvgPool2dBackwardOpIr { + x: desc.x.to_relative(converter), + grad: desc.grad.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::MaxPool1d(desc) => ModuleOperationIr::MaxPool1d(MaxPool1dOpIr { + x: desc.x.to_relative(converter), + kernel_size: desc.kernel_size, + stride: desc.stride, + padding: desc.padding, + dilation: desc.dilation, + ceil_mode: desc.ceil_mode, + out: desc.out.to_relative(converter), + }), + ModuleOperationIr::MaxPool1dWithIndices(desc) => { + ModuleOperationIr::MaxPool1dWithIndices(MaxPool1dWithIndicesOpIr { + x: desc.x.to_relative(converter), + kernel_size: desc.kernel_size, + stride: desc.stride, + padding: desc.padding, + dilation: desc.dilation, + ceil_mode: desc.ceil_mode, + out: desc.out.to_relative(converter), + out_indices: desc.out_indices.to_relative(converter), + }) + } + ModuleOperationIr::MaxPool1dWithIndicesBackward(desc) => { + ModuleOperationIr::MaxPool1dWithIndicesBackward(MaxPool1dWithIndicesBackwardOpIr { + x: desc.x.to_relative(converter), + grad: desc.grad.to_relative(converter), + indices: desc.indices.to_relative(converter), + kernel_size: desc.kernel_size, + stride: desc.stride, + padding: desc.padding, + dilation: desc.dilation, + ceil_mode: desc.ceil_mode, + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::MaxPool2d(desc) => ModuleOperationIr::MaxPool2d(MaxPool2dOpIr { + x: desc.x.to_relative(converter), + kernel_size: desc.kernel_size, + stride: desc.stride, + padding: desc.padding, + dilation: desc.dilation, + ceil_mode: desc.ceil_mode, + out: desc.out.to_relative(converter), + }), + ModuleOperationIr::MaxPool2dWithIndices(desc) => { + ModuleOperationIr::MaxPool2dWithIndices(MaxPool2dWithIndicesOpIr { + x: desc.x.to_relative(converter), + kernel_size: desc.kernel_size, + stride: desc.stride, + padding: desc.padding, + dilation: desc.dilation, + ceil_mode: desc.ceil_mode, + out: desc.out.to_relative(converter), + out_indices: desc.out_indices.to_relative(converter), + }) + } + ModuleOperationIr::MaxPool2dWithIndicesBackward(desc) => { + ModuleOperationIr::MaxPool2dWithIndicesBackward(MaxPool2dWithIndicesBackwardOpIr { + x: desc.x.to_relative(converter), + grad: desc.grad.to_relative(converter), + indices: desc.indices.to_relative(converter), + kernel_size: desc.kernel_size, + stride: desc.stride, + padding: desc.padding, + dilation: desc.dilation, + ceil_mode: desc.ceil_mode, + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::Interpolate(desc) => { + ModuleOperationIr::Interpolate(InterpolateOpIr { + x: desc.x.to_relative(converter), + output_size: desc.output_size, + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::InterpolateBackward(desc) => { + ModuleOperationIr::InterpolateBackward(InterpolateBackwardOpIr { + x: desc.x.to_relative(converter), + grad: desc.grad.to_relative(converter), + output_size: desc.output_size, + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::Rfft(desc) => ModuleOperationIr::Rfft(RfftOpIr { + signal: desc.signal.to_relative(converter), + dim: desc.dim, + n: desc.n, + out_re: desc.out_re.to_relative(converter), + out_im: desc.out_im.to_relative(converter), + }), + ModuleOperationIr::IRfft(desc) => ModuleOperationIr::IRfft(IRfftOpIr { + input_re: desc.input_re.to_relative(converter), + input_im: desc.input_im.to_relative(converter), + dim: desc.dim, + n: desc.n, + out_signal: desc.out_signal.to_relative(converter), + }), + ModuleOperationIr::Attention(desc) => ModuleOperationIr::Attention(AttentionOpIr { + query: desc.query.to_relative(converter), + key: desc.key.to_relative(converter), + value: desc.value.to_relative(converter), + mask: desc.mask.as_ref().map(|m| m.to_relative(converter)), + attn_bias: desc.attn_bias.as_ref().map(|ab| ab.to_relative(converter)), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }), + ModuleOperationIr::CtcLoss(desc) => ModuleOperationIr::CtcLoss(CtcLossOpIr { + log_probs: desc.log_probs.to_relative(converter), + targets: desc.targets.to_relative(converter), + input_lengths: desc.input_lengths.to_relative(converter), + target_lengths: desc.target_lengths.to_relative(converter), + blank: desc.blank, + out: desc.out.to_relative(converter), + }), + ModuleOperationIr::CtcLossBackward(desc) => { + ModuleOperationIr::CtcLossBackward(CtcLossBackwardOpIr { + log_probs: desc.log_probs.to_relative(converter), + targets: desc.targets.to_relative(converter), + input_lengths: desc.input_lengths.to_relative(converter), + target_lengths: desc.target_lengths.to_relative(converter), + grad_loss: desc.grad_loss.to_relative(converter), + blank: desc.blank, + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::LayerNorm(desc) => ModuleOperationIr::LayerNorm(LayerNormOpIr { + input: desc.input.to_relative(converter), + gamma: desc.gamma.to_relative(converter), + beta: desc.beta.as_ref().map(|t| t.to_relative(converter)), + epsilon: desc.epsilon, + out: desc.out.to_relative(converter), + }), + ModuleOperationIr::Unfold4d(desc) => ModuleOperationIr::Unfold4d(Unfold4dOpIr { + x: desc.x.to_relative(converter), + kernel_size: desc.kernel_size, + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }), + ModuleOperationIr::ConvTranspose1dWeightBackward(desc) => { + ModuleOperationIr::ConvTranspose1dWeightBackward( + ConvTranspose1dWeightBackwardOpIr { + x: desc.x.to_relative(converter), + weight: desc.weight.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }, + ) + } + ModuleOperationIr::ConvTranspose1dBiasBackward(desc) => { + ModuleOperationIr::ConvTranspose1dBiasBackward(ConvTranspose1dBiasBackwardOpIr { + x: desc.x.to_relative(converter), + bias: desc.bias.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::ConvTranspose2dWeightBackward(desc) => { + ModuleOperationIr::ConvTranspose2dWeightBackward( + ConvTranspose2dWeightBackwardOpIr { + x: desc.x.to_relative(converter), + weight: desc.weight.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }, + ) + } + ModuleOperationIr::ConvTranspose2dBiasBackward(desc) => { + ModuleOperationIr::ConvTranspose2dBiasBackward(ConvTranspose2dBiasBackwardOpIr { + x: desc.x.to_relative(converter), + bias: desc.bias.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ModuleOperationIr::ConvTranspose3dWeightBackward(desc) => { + ModuleOperationIr::ConvTranspose3dWeightBackward( + ConvTranspose3dWeightBackwardOpIr { + x: desc.x.to_relative(converter), + weight: desc.weight.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }, + ) + } + ModuleOperationIr::ConvTranspose3dBiasBackward(desc) => { + ModuleOperationIr::ConvTranspose3dBiasBackward(ConvTranspose3dBiasBackwardOpIr { + x: desc.x.to_relative(converter), + bias: desc.bias.to_relative(converter), + output_grad: desc.output_grad.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + } + } +} + +impl RelativeOps for FloatOperationIr { + fn to_relative(&self, converter: &mut OperationConverter) -> Self { + match self { + FloatOperationIr::Exp(desc) => FloatOperationIr::Exp(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Log(desc) => FloatOperationIr::Log(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Log1p(desc) => FloatOperationIr::Log1p(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Erf(desc) => FloatOperationIr::Erf(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Powf(desc) => FloatOperationIr::Powf(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Hypot(desc) => FloatOperationIr::Hypot(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::PowfScalar(desc) => FloatOperationIr::PowfScalar(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Sqrt(desc) => FloatOperationIr::Sqrt(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Cos(desc) => FloatOperationIr::Cos(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Sin(desc) => FloatOperationIr::Sin(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Tanh(desc) => FloatOperationIr::Tanh(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Tan(desc) => FloatOperationIr::Tan(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Cosh(desc) => FloatOperationIr::Cosh(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Sinh(desc) => FloatOperationIr::Sinh(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::ArcCos(desc) => FloatOperationIr::ArcCos(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::ArcCosh(desc) => FloatOperationIr::ArcCosh(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::ArcSin(desc) => FloatOperationIr::ArcSin(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::ArcSinh(desc) => FloatOperationIr::ArcSinh(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::ArcTan(desc) => FloatOperationIr::ArcTan(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::ArcTanh(desc) => FloatOperationIr::ArcTanh(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::ArcTan2(desc) => FloatOperationIr::ArcTan2(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::IntoInt(desc) => FloatOperationIr::IntoInt(CastOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Matmul(desc) => FloatOperationIr::Matmul(MatmulOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Cross(desc) => FloatOperationIr::Cross(CrossOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + dim: desc.dim, + }), + FloatOperationIr::Random(desc) => FloatOperationIr::Random(RandomOpIr { + out: desc.out.to_relative(converter), + distribution: desc.distribution, + }), + FloatOperationIr::Recip(desc) => FloatOperationIr::Recip(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Quantize(desc) => FloatOperationIr::Quantize(QuantizeOpIr { + tensor: desc.tensor.to_relative(converter), + qparams: QuantizationParametersIr { + scales: desc.qparams.scales.to_relative(converter), + }, + scheme: desc.scheme, + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Dequantize(desc) => FloatOperationIr::Dequantize(DequantizeOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Round(desc) => FloatOperationIr::Round(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Floor(desc) => FloatOperationIr::Floor(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Ceil(desc) => FloatOperationIr::Ceil(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::Trunc(desc) => FloatOperationIr::Trunc(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::IsNan(desc) => FloatOperationIr::IsNan(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::IsInf(desc) => FloatOperationIr::IsInf(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + FloatOperationIr::GridSample2d(desc) => { + FloatOperationIr::GridSample2d(GridSample2dOpIr { + tensor: desc.tensor.to_relative(converter), + grid: desc.grid.to_relative(converter), + options: desc.options.clone(), + out: desc.out.to_relative(converter), + }) + } + } + } +} + +impl RelativeOps for ActivationOperationIr { + fn to_relative(&self, converter: &mut OperationConverter) -> Self { + match self { + ActivationOperationIr::Relu(desc) => ActivationOperationIr::Relu(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + ActivationOperationIr::ReluBackward(desc) => { + ActivationOperationIr::ReluBackward(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ActivationOperationIr::LeakyRelu(desc) => { + ActivationOperationIr::LeakyRelu(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ActivationOperationIr::PRelu(desc) => ActivationOperationIr::PRelu(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + ActivationOperationIr::Gelu(desc) => ActivationOperationIr::Gelu(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + ActivationOperationIr::GeluBackward(desc) => { + ActivationOperationIr::GeluBackward(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ActivationOperationIr::Sigmoid(desc) => ActivationOperationIr::Sigmoid(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + ActivationOperationIr::SigmoidBackward(desc) => { + ActivationOperationIr::SigmoidBackward(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ActivationOperationIr::HardSigmoid(desc) => { + ActivationOperationIr::HardSigmoid(HardSigmoidOpIr { + tensor: desc.tensor.to_relative(converter), + alpha: desc.alpha.to_relative(converter), + beta: desc.beta.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ActivationOperationIr::LogSigmoid(desc) => { + ActivationOperationIr::LogSigmoid(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ActivationOperationIr::LogSigmoidBackward(desc) => { + ActivationOperationIr::LogSigmoidBackward(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + ActivationOperationIr::Softmax(desc) => ActivationOperationIr::Softmax(DimOpIr { + input: desc.input.to_relative(converter), + axis: desc.axis, + out: desc.out.to_relative(converter), + }), + ActivationOperationIr::LogSoftmax(desc) => ActivationOperationIr::LogSoftmax(DimOpIr { + input: desc.input.to_relative(converter), + axis: desc.axis, + out: desc.out.to_relative(converter), + }), + ActivationOperationIr::Softmin(desc) => ActivationOperationIr::Softmin(DimOpIr { + input: desc.input.to_relative(converter), + axis: desc.axis, + out: desc.out.to_relative(converter), + }), + } + } +} + +impl RelativeOps for BoolOperationIr { + fn to_relative(&self, converter: &mut OperationConverter) -> Self { + match self { + BoolOperationIr::IntoFloat(desc) => BoolOperationIr::IntoFloat(CastOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BoolOperationIr::IntoInt(desc) => BoolOperationIr::IntoInt(CastOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BoolOperationIr::Not(desc) => BoolOperationIr::Not(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BoolOperationIr::And(desc) => BoolOperationIr::And(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BoolOperationIr::Or(desc) => BoolOperationIr::Or(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BoolOperationIr::Xor(desc) => BoolOperationIr::Xor(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + } + } +} + +impl RelativeOps for IntOperationIr { + fn to_relative(&self, converter: &mut OperationConverter) -> Self { + match self { + IntOperationIr::IntoFloat(desc) => IntOperationIr::IntoFloat(CastOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + IntOperationIr::Matmul(desc) => IntOperationIr::Matmul(MatmulOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + IntOperationIr::BitwiseAnd(desc) => IntOperationIr::BitwiseAnd(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + IntOperationIr::BitwiseAndScalar(desc) => { + IntOperationIr::BitwiseAndScalar(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + IntOperationIr::BitwiseOr(desc) => IntOperationIr::BitwiseOr(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + IntOperationIr::BitwiseOrScalar(desc) => IntOperationIr::BitwiseOrScalar(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + IntOperationIr::BitwiseXor(desc) => IntOperationIr::BitwiseXor(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + IntOperationIr::BitwiseXorScalar(desc) => { + IntOperationIr::BitwiseXorScalar(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + IntOperationIr::BitwiseNot(desc) => IntOperationIr::BitwiseNot(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + IntOperationIr::BitwiseLeftShift(desc) => { + IntOperationIr::BitwiseLeftShift(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + IntOperationIr::BitwiseLeftShiftScalar(desc) => { + IntOperationIr::BitwiseLeftShiftScalar(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + IntOperationIr::BitwiseRightShift(desc) => { + IntOperationIr::BitwiseRightShift(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + IntOperationIr::BitwiseRightShiftScalar(desc) => { + IntOperationIr::BitwiseRightShiftScalar(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + } + } +} + +impl RelativeOps for CustomOpIr { + fn to_relative(&self, converter: &mut OperationConverter) -> CustomOpIr { + let id = self.id.clone(); + + CustomOpIr { + id, + inputs: self + .inputs + .iter() + .map(|x| x.to_relative(converter)) + .collect(), + outputs: self + .outputs + .iter() + .map(|x| x.to_relative(converter)) + .collect(), + scalars: self + .scalars + .iter() + .map(|x| x.to_relative(converter)) + .collect(), + } + } +} + +impl RelativeOps for NumericOperationIr { + fn to_relative(&self, converter: &mut OperationConverter) -> Self { + match self { + NumericOperationIr::Add(desc) => NumericOperationIr::Add(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::AddScalar(desc) => NumericOperationIr::AddScalar(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::Sub(desc) => NumericOperationIr::Sub(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::SubScalar(desc) => NumericOperationIr::SubScalar(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::Div(desc) => NumericOperationIr::Div(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::DivScalar(desc) => NumericOperationIr::DivScalar(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::Rem(desc) => NumericOperationIr::Rem(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::RemScalar(desc) => NumericOperationIr::RemScalar(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::Mul(desc) => NumericOperationIr::Mul(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::MulScalar(desc) => NumericOperationIr::MulScalar(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::Abs(desc) => NumericOperationIr::Abs(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::Full(desc) => NumericOperationIr::Full(FullOpIr { + out: desc.out.to_relative(converter), + value: desc.value.to_relative(converter), + }), + NumericOperationIr::MeanDim(desc) => NumericOperationIr::MeanDim(ReduceDimOpIr { + input: desc.input.to_relative(converter), + axis: desc.axis, + accumulator_len: desc.accumulator_len, + out: desc.out.to_relative(converter), + }), + NumericOperationIr::Mean(desc) => NumericOperationIr::Mean(ReduceOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::Sum(desc) => NumericOperationIr::Sum(ReduceOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::SumDim(desc) => { + NumericOperationIr::SumDim(ReduceDimOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + axis: desc.axis, // Axis should stay the same. + accumulator_len: desc.accumulator_len, + }) + } + NumericOperationIr::Prod(desc) => NumericOperationIr::Prod(ReduceOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::ProdDim(desc) => NumericOperationIr::ProdDim(ReduceDimOpIr { + input: desc.input.to_relative(converter), + axis: desc.axis, + accumulator_len: desc.accumulator_len, + out: desc.out.to_relative(converter), + }), + NumericOperationIr::Greater(desc) => NumericOperationIr::Greater(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::GreaterElem(desc) => NumericOperationIr::GreaterElem(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::GreaterEqual(desc) => { + NumericOperationIr::GreaterEqual(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + NumericOperationIr::GreaterEqualElem(desc) => { + NumericOperationIr::GreaterEqualElem(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + NumericOperationIr::Lower(desc) => NumericOperationIr::Lower(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::LowerElem(desc) => NumericOperationIr::LowerElem(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::LowerEqual(desc) => NumericOperationIr::LowerEqual(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::LowerEqualElem(desc) => { + NumericOperationIr::LowerEqualElem(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }) + } + NumericOperationIr::ArgMax(desc) => NumericOperationIr::ArgMax(ReduceDimOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + axis: desc.axis, // Axis should stay the same. + accumulator_len: desc.accumulator_len, + }), + NumericOperationIr::ArgTopK(desc) => NumericOperationIr::ArgTopK(ReduceDimOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + axis: desc.axis, // Axis should stay the same. + accumulator_len: desc.accumulator_len, + }), + NumericOperationIr::TopK(desc) => NumericOperationIr::TopK(ReduceDimOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + axis: desc.axis, // Axis should stay the same. + accumulator_len: desc.accumulator_len, + }), + NumericOperationIr::ArgMin(desc) => NumericOperationIr::ArgMin(ReduceDimOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + axis: desc.axis, // Axis should stay the same. + accumulator_len: desc.accumulator_len, + }), + NumericOperationIr::Max(desc) => NumericOperationIr::Max(ReduceOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::MaxDimWithIndices(desc) => { + NumericOperationIr::MaxDimWithIndices(ReduceDimWithIndicesOpIr { + tensor: desc.tensor.to_relative(converter), + dim: desc.dim, + out: desc.out.to_relative(converter), + out_indices: desc.out_indices.to_relative(converter), + }) + } + NumericOperationIr::MinDimWithIndices(desc) => { + NumericOperationIr::MinDimWithIndices(ReduceDimWithIndicesOpIr { + tensor: desc.tensor.to_relative(converter), + dim: desc.dim, + out: desc.out.to_relative(converter), + out_indices: desc.out_indices.to_relative(converter), + }) + } + NumericOperationIr::Min(desc) => NumericOperationIr::Min(ReduceOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::MaxDim(desc) => NumericOperationIr::MaxDim(ReduceDimOpIr { + input: desc.input.to_relative(converter), + axis: desc.axis, + accumulator_len: desc.accumulator_len, + out: desc.out.to_relative(converter), + }), + NumericOperationIr::MinDim(desc) => NumericOperationIr::MinDim(ReduceDimOpIr { + input: desc.input.to_relative(converter), + axis: desc.axis, + accumulator_len: desc.accumulator_len, + out: desc.out.to_relative(converter), + }), + NumericOperationIr::MaxAbs(desc) => NumericOperationIr::MaxAbs(ReduceOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::MaxAbsDim(desc) => NumericOperationIr::MaxAbsDim(ReduceDimOpIr { + input: desc.input.to_relative(converter), + axis: desc.axis, + accumulator_len: desc.accumulator_len, + out: desc.out.to_relative(converter), + }), + NumericOperationIr::Clamp(desc) => NumericOperationIr::Clamp(ClampOpIr { + tensor: desc.tensor.to_relative(converter), + min: desc.min.to_relative(converter), + max: desc.max.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::IntRandom(desc) => NumericOperationIr::IntRandom(RandomOpIr { + out: desc.out.to_relative(converter), + distribution: desc.distribution, + }), + NumericOperationIr::Powi(desc) => NumericOperationIr::Powi(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::PowiScalar(desc) => NumericOperationIr::PowiScalar(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::CumSum(desc) => NumericOperationIr::CumSum(DimOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + axis: desc.axis, + }), + NumericOperationIr::CumProd(desc) => NumericOperationIr::CumProd(DimOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + axis: desc.axis, + }), + NumericOperationIr::CumMin(desc) => NumericOperationIr::CumMin(DimOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + axis: desc.axis, + }), + NumericOperationIr::CumMax(desc) => NumericOperationIr::CumMax(DimOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + axis: desc.axis, + }), + NumericOperationIr::Neg(desc) => NumericOperationIr::Neg(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::Sign(desc) => NumericOperationIr::Sign(UnaryOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::ClampMin(desc) => NumericOperationIr::ClampMin(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::ClampMax(desc) => NumericOperationIr::ClampMax(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + NumericOperationIr::Sort(desc) => NumericOperationIr::Sort(SortOpIr { + input: desc.input.to_relative(converter), + dim: desc.dim, + descending: desc.descending, + out: desc.out.to_relative(converter), + }), + NumericOperationIr::SortWithIndices(desc) => { + NumericOperationIr::SortWithIndices(SortWithIndicesOpIr { + input: desc.input.to_relative(converter), + dim: desc.dim, + descending: desc.descending, + out: desc.out.to_relative(converter), + out_indices: desc.out_indices.to_relative(converter), + }) + } + NumericOperationIr::ArgSort(desc) => NumericOperationIr::ArgSort(SortOpIr { + input: desc.input.to_relative(converter), + dim: desc.dim, + descending: desc.descending, + out: desc.out.to_relative(converter), + }), + } + } +} + +impl RelativeOps for BaseOperationIr { + fn to_relative(&self, converter: &mut OperationConverter) -> Self { + match self { + BaseOperationIr::Reshape(desc) => BaseOperationIr::Reshape(ShapeOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BaseOperationIr::SwapDims(desc) => BaseOperationIr::SwapDims(SwapDimsOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + dim1: desc.dim1, + dim2: desc.dim2, + }), + BaseOperationIr::Permute(desc) => BaseOperationIr::Permute(PermuteOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + axes: desc.axes.clone(), + }), + BaseOperationIr::Expand(desc) => BaseOperationIr::Expand(ShapeOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BaseOperationIr::Unfold(desc) => BaseOperationIr::Unfold(UnfoldOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + dim: desc.dim, + size: desc.size, + step: desc.step, + }), + BaseOperationIr::Flip(desc) => BaseOperationIr::Flip(FlipOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + axes: desc.axes.clone(), + }), + BaseOperationIr::Slice(desc) => BaseOperationIr::Slice(SliceOpIr { + tensor: desc.tensor.to_relative(converter), + ranges: desc + .ranges + .iter() + .map(|r| converter.relative_range(r)) + .collect(), + out: desc.out.to_relative(converter), + }), + BaseOperationIr::SliceAssign(desc) => BaseOperationIr::SliceAssign(SliceAssignOpIr { + tensor: desc.tensor.to_relative(converter), + ranges: desc + .ranges + .iter() + .map(|r| converter.relative_range(r)) + .collect(), + value: desc.value.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BaseOperationIr::Gather(desc) => BaseOperationIr::Gather(GatherOpIr { + tensor: desc.tensor.to_relative(converter), + dim: desc.dim, + indices: desc.indices.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BaseOperationIr::Scatter(desc) => BaseOperationIr::Scatter(ScatterOpIr { + tensor: desc.tensor.to_relative(converter), + dim: desc.dim, + indices: desc.indices.to_relative(converter), + value: desc.value.to_relative(converter), + update: desc.update, + out: desc.out.to_relative(converter), + }), + BaseOperationIr::ScatterNd(desc) => BaseOperationIr::ScatterNd(ScatterNdOpIr { + data: desc.data.to_relative(converter), + indices: desc.indices.to_relative(converter), + values: desc.values.to_relative(converter), + reduction: desc.reduction, + out: desc.out.to_relative(converter), + }), + BaseOperationIr::GatherNd(desc) => BaseOperationIr::GatherNd(GatherNdOpIr { + data: desc.data.to_relative(converter), + indices: desc.indices.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BaseOperationIr::Select(desc) => BaseOperationIr::Select(SelectOpIr { + tensor: desc.tensor.to_relative(converter), + dim: desc.dim, + indices: desc.indices.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BaseOperationIr::SelectAssign(desc) => { + BaseOperationIr::SelectAssign(SelectAssignOpIr { + tensor: desc.tensor.to_relative(converter), + dim: desc.dim, + indices: desc.indices.to_relative(converter), + value: desc.value.to_relative(converter), + update: desc.update, + out: desc.out.to_relative(converter), + }) + } + BaseOperationIr::MaskWhere(desc) => BaseOperationIr::MaskWhere(MaskWhereOpIr { + tensor: desc.tensor.to_relative(converter), + mask: desc.mask.to_relative(converter), + value: desc.value.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BaseOperationIr::MaskFill(desc) => BaseOperationIr::MaskFill(MaskFillOpIr { + tensor: desc.tensor.to_relative(converter), + mask: desc.mask.to_relative(converter), + value: desc.value.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BaseOperationIr::Equal(desc) => BaseOperationIr::Equal(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BaseOperationIr::EqualElem(desc) => BaseOperationIr::EqualElem(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BaseOperationIr::RepeatDim(desc) => BaseOperationIr::RepeatDim(RepeatDimOpIr { + tensor: desc.tensor.to_relative(converter), + dim: desc.dim, + times: desc.times, + out: desc.out.to_relative(converter), + }), + BaseOperationIr::Cat(desc) => BaseOperationIr::Cat(CatOpIr { + tensors: desc + .tensors + .iter() + .map(|tensor| tensor.to_relative(converter)) + .collect(), + dim: desc.dim, + out: desc.out.to_relative(converter), + }), + BaseOperationIr::Cast(desc) => BaseOperationIr::Cast(CastOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BaseOperationIr::Empty(desc) => BaseOperationIr::Empty(desc.to_relative(converter)), + BaseOperationIr::Ones(desc) => BaseOperationIr::Ones(desc.to_relative(converter)), + BaseOperationIr::Zeros(desc) => BaseOperationIr::Zeros(desc.to_relative(converter)), + BaseOperationIr::NotEqual(desc) => BaseOperationIr::NotEqual(BinaryOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BaseOperationIr::NotEqualElem(desc) => BaseOperationIr::NotEqualElem(ScalarOpIr { + lhs: desc.lhs.to_relative(converter), + rhs: desc.rhs.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BaseOperationIr::All(desc) => BaseOperationIr::All(ReduceOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BaseOperationIr::Any(desc) => BaseOperationIr::Any(ReduceOpIr { + input: desc.input.to_relative(converter), + out: desc.out.to_relative(converter), + }), + BaseOperationIr::AllDim(desc) => BaseOperationIr::AllDim(ReduceDimOpIr { + input: desc.input.to_relative(converter), + axis: desc.axis, + accumulator_len: desc.accumulator_len, + out: desc.out.to_relative(converter), + }), + BaseOperationIr::AnyDim(desc) => BaseOperationIr::AnyDim(ReduceDimOpIr { + input: desc.input.to_relative(converter), + axis: desc.axis, + accumulator_len: desc.accumulator_len, + out: desc.out.to_relative(converter), + }), + } + } +} + +impl RelativeOps for InitOperationIr { + fn to_relative(&self, converter: &mut OperationConverter) -> Self { + Self { + out: self.out.to_relative(converter), + } + } +} + +impl RelativeOps for CreationOpIr { + fn to_relative(&self, converter: &mut OperationConverter) -> Self { + Self { + out: self.out.to_relative(converter), + } + } +} + +impl RelativeOps for TensorIr { + fn to_relative(&self, converter: &mut OperationConverter) -> Self { + let relative_id = self.id.to_relative(converter); + + // We can create relative shapes by mapping each shape found to an ID, which is a `usize`. + let mut relative_shape = Vec::with_capacity(self.shape.rank()); + for dim in self.shape.iter() { + if let Some(dim_id) = converter.shapes_global2relative.get(dim) { + // We already saw that dim value before, so we retrieve its ID. + relative_shape.push(*dim_id); + } else { + // We never saw this dim value before, therefore we create a new ID. + let dim_id = converter.shapes_global2relative.len(); + relative_shape.push(dim_id); + + converter.shapes_global2relative.insert(*dim, dim_id); + converter.shapes_relative2global.insert(dim_id, *dim); + } + } + + // We create the relative tensor. + let relative_tensor = TensorIr { + id: relative_id, + shape: Shape::from(relative_shape), + status: self.status, + dtype: self.dtype, + }; + + // We update both mappings. + converter + .tensors_relative2global + .insert(relative_id, self.clone()); + converter + .tensors_global2relative + .insert(self.id, relative_tensor.clone()); + + relative_tensor + } +} + +impl RelativeOps for DistributedOperationIr { + fn to_relative(&self, converter: &mut OperationConverter) -> Self { + match self { + DistributedOperationIr::AllReduce(desc) => { + DistributedOperationIr::AllReduce(AllReduceOpIr { + tensor: desc.tensor.to_relative(converter), + out: desc.out.to_relative(converter), + op: desc.op, + device_ids: desc.device_ids.clone(), + }) + } + DistributedOperationIr::SyncCollective => DistributedOperationIr::SyncCollective, + } + } +} + +impl RelativeOps for TensorId { + fn to_relative(&self, converter: &mut OperationConverter) -> Self { + if let Some(value) = converter.tensors_global2relative.get(self) { + // If we already have the same tensor registered, we have to update its value, but not + // its id. + value.id + } else { + // We create a new relative id since we never seen this tensor in the graph before. + TensorId::new(converter.tensors_relative2global.len() as u64) + } + } +} + +impl RelativeOps for ScalarIr { + fn to_relative(&self, converter: &mut OperationConverter) -> Self { + // Every scalar variant (including `Bool`) is stashed verbatim and replaced by a positional + // `UInt` placeholder; replay rebinds it by index from `GraphBindings::scalars`, so the + // concrete type is irrelevant to relativization. + let id = ScalarId { + value: converter.scalars.len() as u64, + }; + + converter.scalars.insert(id, *self); + ScalarIr::UInt(id.value) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn_backend::DType; + use burn_ir::{TensorId, TensorIr, TensorStatus}; + + /// Helper to build a minimal [`Context`] with string handles for testing fork behavior. + fn make_test_context() -> Context { + let mut ctx = Context { + tensors: HashMap::new(), + handles: HandleContainer::new(), + scalars: HashMap::new(), + shapes_relative2global: HashMap::new(), + ranges: Vec::new(), + }; + + let id_input = TensorId::new(1); + ctx.tensors.insert( + id_input, + TensorIr { + id: id_input, + shape: Shape::new([4, 4]), + status: TensorStatus::ReadOnly, + dtype: DType::F32, + }, + ); + ctx.handles + .register_handle(id_input, "input_handle".to_string()); + + ctx + } + + #[test] + fn context_fork_output_handles_are_isolated() { + // Output handles registered in a forked context are NOT visible in the original — + // forks are independent sandboxes. `burn-cubecl-fusion`'s `TuneInput` layer is + // responsible for merging new handles back into the real context when appropriate. + let original = make_test_context(); + let output_id = TensorId::new(100); + + { + let mut fork = original.fork(); + fork.handles + .register_handle(output_id, "output_handle".to_string()); + + // The fork has the output handle. + assert!(fork.handles.get_handle_ref(&output_id).is_some()); + } + + // But the original does NOT — isolation is the point of `fork()`. + assert!(original.handles.get_handle_ref(&output_id).is_none()); + } + + #[test] + fn context_fork_preserves_input_handles() { + let original = make_test_context(); + let input_id = TensorId::new(1); + + let fork = original.fork(); + + // Fork should have a copy of the input handle. + assert_eq!( + fork.handles.get_handle_ref(&input_id), + Some(&"input_handle".to_string()) + ); + // Original is unchanged. + assert_eq!( + original.handles.get_handle_ref(&input_id), + Some(&"input_handle".to_string()) + ); + } + + #[test] + fn context_double_fork_fully_isolated() { + // Forking a fork creates a second level of isolation — a mutation in `fork2` is + // invisible to both `fork1` and the original. + let original = make_test_context(); + + let fork1 = original.fork(); + let mut fork2 = fork1.fork(); + + let deep_output_id = TensorId::new(200); + fork2 + .handles + .register_handle(deep_output_id, "deep_output".to_string()); + + // Neither the first fork nor the original see the deeply-nested output. + assert!(fork1.handles.get_handle_ref(&deep_output_id).is_none()); + assert!(original.handles.get_handle_ref(&deep_output_id).is_none()); + } +} + +#[cfg(test)] +mod tests_ir { + use super::*; + use burn_backend::DType; + use burn_ir::{TensorId, TensorIr, TensorStatus}; + + #[test] + fn tensor_description_to_relative() { + let tensor1 = TensorIr { + id: TensorId::new(500), + shape: Shape::new([512, 32, 2048]), + status: TensorStatus::ReadOnly, + dtype: DType::F32, + }; + let tensor2 = TensorIr { + id: TensorId::new(501), + shape: Shape::new([512, 128, 2048]), + status: TensorStatus::ReadOnly, + dtype: DType::F32, + }; + let mut converter = OperationConverter::default(); + let tensor1_local = tensor1.to_relative(&mut converter); + let tensor2_local = tensor2.to_relative(&mut converter); + + assert_eq!( + tensor1_local, + TensorIr { + id: TensorId::new(0), + shape: Shape::new([1, 2, 3]), + status: TensorStatus::ReadOnly, + dtype: DType::F32 + } + ); + assert_eq!( + tensor2_local, + TensorIr { + id: TensorId::new(1), + shape: Shape::new([1, 4, 3]), + status: TensorStatus::ReadOnly, + dtype: DType::F32 + } + ); + } + + #[test] + fn scalar_ir_to_relative() { + let scalar1 = ScalarIr::Float(1.0); + let scalar2 = ScalarIr::UInt(1); + let mut converter = OperationConverter::default(); + let scalar1_local = scalar1.to_relative(&mut converter); + let scalar2_local = scalar2.to_relative(&mut converter); + + assert_eq!(scalar1_local, ScalarIr::UInt(0)); + assert_eq!(scalar2_local, ScalarIr::UInt(1)); + } + + #[test] + fn scalar_ir_bool_to_relative() { + // A `Bool` scalar used to `todo!()`-panic here, which crashed any custom op carrying a bool + // argument the moment it was fused into a cached graph. It must relativize like every other + // scalar: replaced by a positional `UInt` placeholder, with the concrete value stashed so + // replay can rebind it from `GraphBindings::scalars`. + let mut converter = OperationConverter::default(); + let relative = ScalarIr::Bool(true).to_relative(&mut converter); + + assert_eq!(relative, ScalarIr::UInt(0)); + assert_eq!( + converter.scalars.get(&ScalarId { value: 0 }), + Some(&ScalarIr::Bool(true)) + ); + } +} diff --git a/crates/burn-fusion/src/stream/execution/base.rs b/crates/burn-fusion/src/stream/execution/base.rs new file mode 100644 index 0000000..d76231e --- /dev/null +++ b/crates/burn-fusion/src/stream/execution/base.rs @@ -0,0 +1,16 @@ +use burn_ir::HandleContainer; + +use crate::FusionRuntime; + +/// The mode in which the execution is done. +#[derive(Clone, Copy, Debug)] +pub(crate) enum ExecutionMode { + Lazy, + Sync, +} + +/// General trait to abstract how a single operation is executed. +pub trait Operation: Send + Sync + core::fmt::Debug { + /// Execute the operation. + fn execute(&self, handles: &mut HandleContainer); +} diff --git a/crates/burn-fusion/src/stream/execution/explorer.rs b/crates/burn-fusion/src/stream/execution/explorer.rs new file mode 100644 index 0000000..1ffe32f --- /dev/null +++ b/crates/burn-fusion/src/stream/execution/explorer.rs @@ -0,0 +1,163 @@ +use std::sync::Arc; + +use burn_ir::OperationIr; +use burn_std::config::{config, fusion::FusionLogLevel, log_fusion}; + +use super::{ExecutionMode, op_kind}; +use crate::{ + NumOperations, OperationFuser, + search::{BlockOptimization, StreamOptimizer}, + stream::store::ExecutionStrategy, +}; + +/// Explore and create new optimization. +pub struct Explorer { + optimizer: StreamOptimizer, + num_deferred: usize, + num_explored: usize, + is_still_optimizing: bool, + /// Number of optimizations actually built so far (one per cache miss that ran the optimizer). + num_explorations: usize, + /// Stop exploring once `num_explorations` reaches this; see [`BeamSearchConfig::max_explorations`]. + max_explorations: Option, +} + +/// The result of an exploration done by the [explorer](Explorer). +pub enum ExplorationAction { + /// Found a new optimization. + Completed(BlockOptimization), + /// We should continue exploring before arriving at a conclusion. + Continue, + /// Exploration is disabled (the cap was reached): execute the segment unfused, without caching. + Unfused(BlockOptimization), +} + +impl Explorer { + /// Create a new explorer. + pub(crate) fn new(optimizations: Vec>>) -> Self { + Self { + optimizer: StreamOptimizer::new(optimizations), + num_deferred: 0, + num_explored: 0, + is_still_optimizing: true, + num_explorations: 0, + max_explorations: config().fusion().beam_search.max_explorations, + } + } + + /// Indicate that a new operation is added. + pub(crate) fn on_new_operation(&mut self) { + self.num_deferred += 1; + } + + /// If the explorer is up to date. + pub(crate) fn is_up_to_date(&self) -> bool { + self.num_deferred == 0 + } + + /// Explore the provided operations. + pub(crate) fn explore( + &mut self, + operations: &[OperationIr], + mode: ExecutionMode, + ) -> ExplorationAction { + let total_ops = operations.len(); + let deferred = self.num_deferred; + let mode_dbg = match mode { + ExecutionMode::Lazy => "lazy", + ExecutionMode::Sync => "sync", + }; + + log_fusion(FusionLogLevel::Full, move || { + format!("[explorer] explore ({mode_dbg}): {deferred} deferred of {total_ops} queued") + }); + + // Exploration cap reached: skip the optimizer (both the incremental block-register in + // `update` and the search in `optimize`) and run the segment unfused. We mark the deferred + // ops as consumed so the processor sees the explorer as up-to-date and stops looping. + if let Some(max) = self.max_explorations + && self.num_explorations >= max + { + self.num_deferred = 0; + self.is_still_optimizing = false; + let ordering: Vec = (0..operations.len()).collect(); + let strategy = ExecutionStrategy::Operations { + ordering: Arc::new(ordering.clone()), + }; + return ExplorationAction::Unfused(BlockOptimization::new(strategy, ordering)); + } + + self.update(operations); + + // Can only continue exploration when not sync. + if let ExecutionMode::Lazy = mode + && self.is_still_optimizing + { + return ExplorationAction::Continue; + } + + let mut optimization = self.optimizer.optimize(operations); + self.num_explorations += 1; + + // At a sync the segment is final: fold the drained tail the search left for + // "the next round" into the plan as un-fused operations, so the cached plan + // covers the whole segment and can be matched by the policy on the next + // identical sync instead of re-exploring the same graph every time. + if let ExecutionMode::Sync = mode { + optimization.include_trailing(operations.len()); + } + + ExplorationAction::Completed(optimization) + } + + /// Number of optimizations built so far; a cached plan that is reused does not + /// count. Used by tests to assert that recurring graphs stop exploring. + #[cfg(test)] + pub(crate) fn num_explorations(&self) -> usize { + self.num_explorations + } + + /// Reset the state of the explorer to the provided list of operations. + pub(crate) fn reset(&mut self, operations: &[OperationIr]) { + self.optimizer.reset(); + self.num_explored = 0; + self.num_deferred = operations.len(); + self.is_still_optimizing = true; + } + + /// Register any operations that we had deferred + fn update(&mut self, operations: &[OperationIr]) { + let start_explored = self.num_explored; + for i in (0..self.num_deferred).rev() { + if !self.is_still_optimizing { + let remaining = i + 1; + let seen = self.num_explored - start_explored; + log_fusion(FusionLogLevel::Full, move || { + format!( + "[explorer] stopped optimizing after {seen} new op(s); {remaining} deferred op(s) left unprocessed" + ) + }); + break; + } + let index = operations.len() - 1 - i; + let relative = &operations[index]; + + self.optimizer.register(relative); + self.num_explored += 1; + + let before = self.is_still_optimizing; + self.is_still_optimizing = self.optimizer.still_optimizing(); + if before && !self.is_still_optimizing { + let explored = self.num_explored; + log_fusion(FusionLogLevel::Full, || { + format!( + "[explorer] still_optimizing → false after op {} (explored {explored} ops)", + op_kind(relative) + ) + }); + } + } + + self.num_deferred = 0; + } +} diff --git a/crates/burn-fusion/src/stream/execution/mod.rs b/crates/burn-fusion/src/stream/execution/mod.rs new file mode 100644 index 0000000..fbe22de --- /dev/null +++ b/crates/burn-fusion/src/stream/execution/mod.rs @@ -0,0 +1,20 @@ +pub(crate) mod validator; + +mod base; +mod explorer; +mod ordering; +mod policy; +mod processor; +pub(crate) mod trace; + +pub(crate) use trace::{log_execution_table, op_kind}; + +pub use base::*; +pub use ordering::*; + +pub(crate) use explorer::*; +pub(crate) use policy::*; +pub(crate) use processor::*; + +#[cfg(test)] +pub(crate) mod tests; diff --git a/crates/burn-fusion/src/stream/execution/ordering.rs b/crates/burn-fusion/src/stream/execution/ordering.rs new file mode 100644 index 0000000..9ce8538 --- /dev/null +++ b/crates/burn-fusion/src/stream/execution/ordering.rs @@ -0,0 +1,69 @@ +use std::sync::Arc; + +use burn_ir::HandleContainer; + +use crate::{FusionRuntime, NumOperations, Optimization, UnfusedOp, stream::Context}; + +/// Manage the execution of potentially multiple optimizations and operations out of order. +pub struct OrderedExecution { + operations: Vec>, + num_executed: usize, + ordering: Option>>, +} + +impl OrderedExecution { + /// Returns the operation that can be executed without impacting the state of the execution. + /// + /// This is useful to implement fallback for optimizations. + #[allow(clippy::borrowed_box)] + pub fn operation_within_optimization(&self, index: usize) -> UnfusedOp { + match &self.ordering { + Some(val) => { + let index = val[index]; + self.operations[index].clone() + } + None => panic!("No ordering provided"), + } + } + + pub(crate) fn new(operations: Vec>) -> Self { + Self { + operations, + num_executed: 0, + ordering: None, + } + } + + pub(crate) fn finish(mut self) -> (Vec>, usize) { + self.operations.drain(0..self.num_executed); + (self.operations, self.num_executed) + } + + pub(crate) fn execute_optimization( + &mut self, + optimization: &mut R::Optimization, + context: &mut Context, + ordering: Arc>, + ) { + if ordering.len() > self.operations.len() { + panic!("Ordering is bigger than operations"); + } + self.ordering = Some(ordering); + let num_drained = optimization.len(); + optimization.execute(context, self); + self.num_executed += num_drained; + } + + pub(crate) fn execute_operations( + &mut self, + handles: &mut HandleContainer, + ordering: &[usize], + ) { + self.num_executed += ordering.len(); + + for id in ordering { + let op = &self.operations[*id]; + op.execute(handles); + } + } +} diff --git a/crates/burn-fusion/src/stream/execution/policy.rs b/crates/burn-fusion/src/stream/execution/policy.rs new file mode 100644 index 0000000..14c21b8 --- /dev/null +++ b/crates/burn-fusion/src/stream/execution/policy.rs @@ -0,0 +1,572 @@ +use burn_ir::OperationIr; + +use super::ExecutionMode; +use super::validator::{ + ExecutionPlanOperationsStore, TriggerOperationsStore, TriggerProgress, TriggerValidator, + ValidatorState, +}; +use crate::stream::execution::validator::OperationsValidator; +use crate::stream::store::{ExecutionPlanId, ExecutionPlanStore, ExecutionTrigger, SearchQuery}; +use std::marker::PhantomData; + +/// The policy keeps track of all possible execution plans for the current operations. +/// +/// # Details +/// +/// We keep track of each new operation added and invalidate potential execution plans +/// when we see a different operation is added. +/// +/// Therefore, the overhead is very minimal, since the time-complexity of checking for existing +/// execution plans scales with the number of concurrent potential plans for the current operations, +/// which isn't supposed to be big at any time. +pub(crate) struct Policy { + /// List of potential execution plans that are compatible with current stream segment + candidates: Vec>, + /// List of candidate execution plans that have been found; we can still keep searching + /// to potentially find a better one. + availables: Vec, + /// The found execution plan that should be executed, along with the number of operations + /// in the plan. + found: Option<(ExecutionPlanId, usize)>, + /// The number of operations that have been analyzed + num_operations: usize, + _item_type: PhantomData, +} + +#[derive(new)] +struct AvailableItem { + id: ExecutionPlanId, + size: usize, + triggers: Vec, +} + +/// Action to be made depending on the stream. +#[derive(PartialEq, Eq, Debug)] +pub enum Action { + /// Continue exploring using the [builder](crate::OptimizationBuilder). + Explore, + /// The current policy indicates that an exploration may be possible in the future, so the + /// best action is to defer any execution. + /// + /// Sometimes, it can be a false positive and a new exploration should be built from scratch. + /// Therefore it's important to keep the previous operations to rebuild the state if it + /// happens. + Defer, + /// An exploration has been found, and the best action is to execute it! + Execute(ExecutionPlanId), +} + +impl Policy { + /// Create a new policy. + pub(crate) fn new() -> Self { + Self { + candidates: Vec::new(), + availables: Vec::new(), + found: None, + num_operations: 0, + _item_type: PhantomData, + } + } + + /// Returns the [action](Action) that should be taken given the state of the policy. + pub fn action( + &self, + store: &ExecutionPlanStore, + operations: &[OperationIr], + mode: ExecutionMode, + ) -> Action { + if self.num_operations < operations.len() { + panic!( + "Internal Error: Can't retrieve the policy action on a list of operations bigger than what is analyzed." + ); + } + + if let Some((id, _length)) = self.found { + return Action::Execute(id); + } + + match mode { + ExecutionMode::Lazy => self.action_lazy(operations), + ExecutionMode::Sync => self.action_sync(operations, store), + } + } + + /// Update the policy state. + pub fn update(&mut self, store: &ExecutionPlanStore, operation: &OperationIr) { + // reset the candidates to contain all execution plans starting with the operation. + if self.num_operations == 0 { + self.candidates = store + .find(SearchQuery::PlansStartingWith(operation)) + .into_iter() + .map(OperationsValidator::new) + .collect(); + } + + self.update_candidates(store, operation); + self.check_candidates(store); + + self.update_availables(store, operation); + self.check_availables(); + self.num_operations += 1; + } + + // Reset the state of the policy. + pub fn reset(&mut self) { + self.candidates.clear(); + self.availables.clear(); + + self.num_operations = 0; + self.found = None; + } + + /// Check which candidates can be removed, and which one can go from + /// 'candidate' to 'available' + fn check_candidates(&mut self, store: &ExecutionPlanStore) { + let mut candidates_to_remove = Vec::new(); + + for candidate in self.candidates.iter() { + match candidate.state { + ValidatorState::Found { size } => { + let item = store.get_unchecked(candidate.id); + let mut triggers = Vec::with_capacity(item.triggers.len()); + + for (index, trigger) in item.triggers.iter().enumerate() { + triggers.push(match trigger { + ExecutionTrigger::OnOperations(_) => TriggerValidator::OnOperations { + matching: OperationsValidator::new(index), + progress: TriggerProgress::NotInit, + }, + ExecutionTrigger::OnSync => TriggerValidator::OnSync, + ExecutionTrigger::Always => TriggerValidator::Always, + }); + } + + self.availables + .push(AvailableItem::new(candidate.id, size, triggers)); + candidates_to_remove.push(candidate.id); + } + ValidatorState::Invalidated => { + candidates_to_remove.push(candidate.id); + } + ValidatorState::Validating => {} + }; + } + + let mut updated_candidates = Vec::new(); + core::mem::swap(&mut updated_candidates, &mut self.candidates); + + self.candidates = updated_candidates + .into_iter() + .filter(|candidate| !candidates_to_remove.iter().any(|id| id == &candidate.id)) + .collect(); + } + + fn check_availables(&mut self) { + for available in self.availables.iter() { + for trigger in available.triggers.iter() { + match trigger { + TriggerValidator::OnOperations { + matching, + progress: _, + } => { + if let ValidatorState::Found { + size: _size_of_trigger, + } = matching.state + { + self.found = Some((available.id, available.size)); + return; + } + } + TriggerValidator::Always => { + self.found = Some((available.id, available.size)); + return; + } + TriggerValidator::OnSync => { + // Does nothing during an update. + } + } + } + } + } + + fn update_candidates(&mut self, store: &ExecutionPlanStore, operation: &OperationIr) { + let main_store = ExecutionPlanOperationsStore::new(store); + + self.candidates + .iter_mut() + .for_each(|candidate| candidate.update(operation, self.num_operations, &main_store)); + } + + fn update_availables(&mut self, store: &ExecutionPlanStore, operation: &OperationIr) { + self.availables.iter_mut().for_each(|available| { + let store_trigger = TriggerOperationsStore::new(available.id, store); + + available.triggers.iter_mut().for_each(|trigger| { + if let TriggerValidator::OnOperations { matching, progress } = trigger { + match progress { + TriggerProgress::NotInit => { + *progress = TriggerProgress::NumChecked(0); + } + TriggerProgress::NumChecked(num_check) => { + matching.update(operation, *num_check, &store_trigger); + *num_check += 1; + } + } + } + }); + }); + } + + fn action_lazy(&self, operations: &[OperationIr]) -> Action { + if !self.candidates.is_empty() { + return Action::Defer; + } + + for available in self.availables.iter() { + if available.size == operations.len() { + return Action::Defer; + } + + for trigger in available.triggers.iter() { + if let TriggerValidator::OnOperations { + matching, + progress: _, + } = trigger + && let ValidatorState::Validating = matching.state + { + return Action::Defer; + } + } + } + + Action::Explore + } + + fn action_sync(&self, operations: &[OperationIr], store: &ExecutionPlanStore) -> Action { + for available in self.availables.iter() { + if available.size == operations.len() { + return Action::Execute(available.id); + } + } + + for candidate in self.candidates.iter() { + let item = store.get_unchecked(candidate.id); + + if item.operations.len() == operations.len() { + return Action::Execute(candidate.id); + } + } + + Action::Explore + } +} + +#[cfg(test)] +mod tests { + use burn_backend::{DType, Shape}; + use burn_ir::{FloatOperationIr, TensorId, TensorIr, TensorStatus, UnaryOpIr}; + + use super::*; + use crate::{ + search::BlockOptimization, + stream::store::{ExecutionPlan, ExecutionStrategy, ExecutionTrigger}, + }; + use std::ops::Range; + + #[test] + fn given_no_optimization_should_explore() { + let store = ExecutionPlanStore::default(); + let mut policy = Policy::new(); + let stream = TestStream::new(3); + + stream.assert_updates( + &store, + &mut policy, + AssertUpdatesOptions::OperationsIndex(0..3), + Action::Explore, + ); + } + + #[test] + fn given_existing_optimizations_when_sync_should_execute_one_when_available() { + let mut store = ExecutionPlanStore::default(); + let mut policy = Policy::new(); + let stream = TestStream::new(3); + + let id_1 = store.add(ExecutionPlan { + operations: stream.operations[0..2].to_vec(), + triggers: Vec::new(), + optimization: BlockOptimization::new(ExecutionStrategy::operations(2), Vec::new()), + }); + let _id_2 = store.add(ExecutionPlan { + operations: stream.operations[0..3].to_vec(), + triggers: Vec::new(), + optimization: BlockOptimization::new(ExecutionStrategy::operations(3), Vec::new()), + }); + + stream.assert_updates( + &store, + &mut policy, + AssertUpdatesOptions::OperationsIndex(0..2), + Action::Defer, + ); + + let action = policy.action(&store, &stream.operations[0..2], ExecutionMode::Sync); + assert_eq!(action, Action::Execute(id_1)); + } + + #[test] + fn given_existing_plan_when_found_trigger_should_execute_plan() { + let mut store = ExecutionPlanStore::default(); + let mut policy = Policy::new(); + + let stream = TestStream::new(3); + let id = store.add(ExecutionPlan { + operations: stream.operations[0..2].to_vec(), + triggers: stream.operations[2..3] + .iter() + .map(|desc| ExecutionTrigger::OnOperations(vec![desc.clone()])) + .collect(), + optimization: BlockOptimization::new(ExecutionStrategy::operations(2), Vec::new()), + }); + + stream.assert_updates( + &store, + &mut policy, + AssertUpdatesOptions::OperationsIndex(0..2), + Action::Defer, + ); + stream.assert_updates( + &store, + &mut policy, + AssertUpdatesOptions::OperationsIndex(2..3), + Action::Execute(id), + ); + } + + #[test] + fn should_support_multiple_triggers() { + let mut store = ExecutionPlanStore::default(); + let mut policy_1 = Policy::new(); + let mut policy_2 = Policy::new(); + + let mut stream_1 = TestStream::new(2); + let mut stream_2 = TestStream::new(2); + + // Create different end operation for each stream. + let trigger_id_1 = 5; + let trigger_id_2 = 6; + stream_1.new_ops(trigger_id_1); + stream_2.new_ops(trigger_id_2); + + let id = store.add(ExecutionPlan { + operations: stream_1.operations[0..2].to_vec(), + triggers: vec![ + ExecutionTrigger::OnOperations(vec![stream_1.operations[2].clone()]), + ExecutionTrigger::OnOperations(vec![stream_2.operations[2].clone()]), + ], + optimization: BlockOptimization::new(ExecutionStrategy::operations(2), Vec::new()), + }); + + stream_1.assert_updates( + &store, + &mut policy_1, + AssertUpdatesOptions::OperationsIndex(0..2), + Action::Defer, + ); + stream_2.assert_updates( + &store, + &mut policy_2, + AssertUpdatesOptions::OperationsIndex(0..2), + Action::Defer, + ); + + stream_1.assert_updates( + &store, + &mut policy_1, + AssertUpdatesOptions::OperationsIndex(2..3), // First trigger. + Action::Execute(id), + ); + stream_2.assert_updates( + &store, + &mut policy_2, + AssertUpdatesOptions::OperationsIndex(2..3), // Second trigger. + Action::Execute(id), + ); + } + + #[test] + fn should_select_right_optimization() { + let mut store = ExecutionPlanStore::default(); + let mut policy_1 = Policy::new(); + let mut policy_2 = Policy::new(); + + let mut stream_1 = TestStream::new(2); + let mut stream_2 = TestStream::new(2); + + // Create different streams after op 2. + stream_1.new_ops(4); + stream_1.new_ops(5); + + stream_2.new_ops(5); + stream_2.new_ops(6); + + let optimization_stream_1 = store.add(ExecutionPlan { + operations: stream_1.operations[0..3].to_vec(), + triggers: stream_1.operations[3..4] + .iter() + .map(|desc| ExecutionTrigger::OnOperations(vec![desc.clone()])) + .collect(), + optimization: BlockOptimization::new(ExecutionStrategy::operations(3), Vec::new()), + }); + let optimization_stream_2 = store.add(ExecutionPlan { + operations: stream_2.operations[0..3].to_vec(), + triggers: stream_2.operations[3..4] + .iter() + .map(|desc| ExecutionTrigger::OnOperations(vec![desc.clone()])) + .collect(), + optimization: BlockOptimization::new(ExecutionStrategy::operations(3), Vec::new()), + }); + assert_ne!(optimization_stream_1, optimization_stream_2); + + stream_1.assert_updates( + &store, + &mut policy_1, + AssertUpdatesOptions::OperationsIndex(0..3), + Action::Defer, + ); + stream_2.assert_updates( + &store, + &mut policy_2, + AssertUpdatesOptions::OperationsIndex(0..3), + Action::Defer, + ); + + stream_1.assert_updates( + &store, + &mut policy_1, + AssertUpdatesOptions::OperationsIndex(3..4), + Action::Execute(optimization_stream_1), + ); + stream_2.assert_updates( + &store, + &mut policy_2, + AssertUpdatesOptions::OperationsIndex(3..4), + Action::Execute(optimization_stream_2), + ); + } + + #[test] + fn should_invalidate_wrong_optimizations() { + let mut store = ExecutionPlanStore::default(); + let stream_1 = TestStream::new(4); + let mut stream_2 = TestStream::new(2); + stream_2.new_ops(6); + stream_2.new_ops(7); + + store.add(ExecutionPlan { + operations: stream_1.operations[0..3].to_vec(), + triggers: stream_1.operations[3..4] + .iter() + .map(|desc| ExecutionTrigger::OnOperations(vec![desc.clone()])) + .collect(), + optimization: BlockOptimization::new(ExecutionStrategy::operations(3), Vec::new()), + }); + + let mut policy = Policy::new(); + // Same path as stream 1 + stream_2.assert_updates( + &store, + &mut policy, + AssertUpdatesOptions::OperationsIndex(0..2), + Action::Defer, + ); + + // But is different. + stream_2.assert_updates( + &store, + &mut policy, + AssertUpdatesOptions::OperationsIndex(2..4), + Action::Explore, + ); + } + + #[derive(Default, Debug)] + struct TestStream { + tensors: Vec, + operations: Vec, + } + + #[derive(Debug)] + enum AssertUpdatesOptions { + OperationsIndex(Range), + } + + impl TestStream { + /// Create a new test stream with `num_ops` operations registered. + pub fn new(num_ops: usize) -> Self { + let mut stream = Self::default(); + for id in 0..num_ops { + stream.new_ops(id as u64 + 1); + } + + stream + } + + /// The first follow should only be cache miss. + pub fn assert_updates( + &self, + optimizations: &ExecutionPlanStore<()>, + policy: &mut Policy<()>, + options: AssertUpdatesOptions, + action: Action, + ) { + match options { + AssertUpdatesOptions::OperationsIndex(range) => { + for i in range { + let stream = &self.operations[0..i]; + let next_ops = &self.operations[i]; + policy.update(optimizations, next_ops); + let result = policy.action(optimizations, stream, ExecutionMode::Lazy); + + assert_eq!(result, action); + } + } + } + } + + /// Add a simple operation to the stream. + pub fn new_ops(&mut self, out_id: u64) { + if self.tensors.is_empty() { + // Root node. + self.new_empty_node(0); + } + + // Out node. + self.new_empty_node(out_id); + + self.operations.push(OperationIr::Float( + DType::F32, + FloatOperationIr::Log(self.unary_description()), + )); + } + + fn new_empty_node(&mut self, id: u64) { + self.tensors.push(TensorIr { + id: TensorId::new(id), + shape: Shape::new([32, 32, 1]), + status: TensorStatus::NotInit, + dtype: DType::F32, + }); + } + + fn unary_description(&self) -> UnaryOpIr { + let size = self.tensors.len(); + + UnaryOpIr { + input: self.tensors[size - 2].clone(), + out: self.tensors[size - 1].clone(), + } + } + } +} diff --git a/crates/burn-fusion/src/stream/execution/processor.rs b/crates/burn-fusion/src/stream/execution/processor.rs new file mode 100644 index 0000000..ef36f3f --- /dev/null +++ b/crates/burn-fusion/src/stream/execution/processor.rs @@ -0,0 +1,224 @@ +use burn_ir::OperationIr; +use burn_std::config::{fusion::FusionLogLevel, log_fusion}; + +use super::{ExecutionMode, ExplorationAction, Explorer}; +use crate::search::BlockOptimization; +use crate::stream::execution::{Action, Policy}; +use crate::stream::store::{ExecutionPlan, ExecutionPlanId, ExecutionPlanStore, ExecutionTrigger}; +use crate::{NumOperations, OperationFuser}; + +/// Process a [stream segment](StreamSegment) following a [policy](Policy). +pub(crate) struct Processor { + policy: Policy, + explorer: Explorer, +} + +/// A part of a stream that can be executed partially using [execution plan](ExecutionPlan). +pub(crate) trait StreamSegment { + /// The operations in the segment. + fn operations(&self) -> &[OperationIr]; + /// Execute part of the segment using the given plan id. + fn execute(&mut self, id: ExecutionPlanId, store: &mut ExecutionPlanStore); + /// Execute the segment with a one-off optimization, without caching it in the store. + fn execute_unfused(&mut self, optimization: BlockOptimization); +} + +impl Processor { + /// Create a new stream processor. + pub fn new(optimizations: Vec>>) -> Self { + Self { + policy: Policy::new(), + explorer: Explorer::new(optimizations), + } + } + + /// Process the [stream segment](StreamSegment) with the provided [mode](ExecutionMode). + pub fn process( + &mut self, + mut segment: Segment, + store: &mut ExecutionPlanStore, + mode: ExecutionMode, + ) where + Segment: StreamSegment, + { + // We assume that we always register a new operation in lazy mode. + if let ExecutionMode::Lazy = mode { + self.on_new_operation(&segment, store); + } + + loop { + if segment.operations().is_empty() { + break; + } + + let action = self.policy.action(store, segment.operations(), mode); + + match action { + Action::Explore => { + self.explore(&mut segment, store, mode); + + if self.explorer.is_up_to_date() { + break; + } + } + Action::Defer => { + match mode { + ExecutionMode::Lazy => break, + ExecutionMode::Sync => panic!("Can't defer while sync"), + }; + } + Action::Execute(id) => { + let mode_dbg = match mode { + ExecutionMode::Lazy => "lazy", + ExecutionMode::Sync => "sync", + }; + let num_ops = segment.operations().len(); + log_fusion(FusionLogLevel::Full, move || { + format!( + "[plan] cache hit: execute plan #{id} ({mode_dbg}, segment has {num_ops} ops)" + ) + }); + + if let ExecutionMode::Sync = mode { + store.add_trigger(id, ExecutionTrigger::OnSync); + } + + segment.execute(id, store); + self.reset(store, segment.operations()); + } + }; + } + } + + fn on_new_operation(&mut self, segment: &Segment, store: &mut ExecutionPlanStore) + where + Segment: StreamSegment, + { + self.policy.update( + store, + segment + .operations() + .last() + .expect("At least one operation in the operation list."), + ); + self.explorer.on_new_operation(); + } + + fn explore>( + &mut self, + item: &mut Item, + store: &mut ExecutionPlanStore, + mode: ExecutionMode, + ) { + match self.explorer.explore(item.operations(), mode) { + ExplorationAction::Completed(optim) => { + let id = Self::on_exploration_completed( + &self.policy, + item.operations(), + store, + optim, + mode, + ); + item.execute(id, store); + self.reset(store, item.operations()); + } + ExplorationAction::Continue => { + if let ExecutionMode::Sync = mode { + panic!("Can't continue exploring when sync.") + } + } + ExplorationAction::Unfused(optim) => { + // Exploration is capped: run the ops unfused and don't touch the store, so a + // never-cacheable dynamic graph stops both re-exploring and growing the cache. + item.execute_unfused(optim); + self.reset(store, item.operations()); + } + } + } + + /// Number of optimizations built so far; a cached plan that is reused does not + /// count. Used by tests to assert that recurring graphs stop exploring. + #[cfg(test)] + pub(crate) fn num_explorations(&self) -> usize { + self.explorer.num_explorations() + } + + fn reset(&mut self, store: &mut ExecutionPlanStore, operations: &[OperationIr]) { + self.explorer.reset(operations); + self.policy.reset(); + + // Reset the policy state with the remaining operations + for operation in operations.iter() { + self.policy.update(store, operation); + } + } + + /// We found an optimization (i.e. a new execution plan). + /// Cache it in the store. + fn on_exploration_completed( + policy: &Policy, + operations: &[OperationIr], + store: &mut ExecutionPlanStore, + optimization: BlockOptimization, + mode: ExecutionMode, + ) -> ExecutionPlanId { + let num_optimized = optimization.ordering.len(); + let relative = &operations[0..num_optimized]; + + { + let total_ops = operations.len(); + let mode_dbg = match mode { + ExecutionMode::Lazy => "lazy", + ExecutionMode::Sync => "sync", + }; + log_fusion(FusionLogLevel::Full, move || { + format!( + "[plan] exploration completed: {mode_dbg}, {num_optimized}/{total_ops} ops optimized" + ) + }); + } + + match mode { + ExecutionMode::Lazy => { + let next_ops = &operations[num_optimized..operations.len()]; + + let trigger = if next_ops.is_empty() { + // Happens if the next ops is included in the fused operation, and there is no + // way the builder can still continue fusing. + ExecutionTrigger::Always + } else { + ExecutionTrigger::OnOperations(next_ops.to_vec()) + }; + + match policy.action(store, relative, ExecutionMode::Sync) { + Action::Execute(id) => { + store.add_trigger(id, trigger); + id + } + _ => { + let plan = ExecutionPlan { + operations: relative.to_vec(), + triggers: vec![trigger], + optimization, + }; + store.add(plan) + } + } + } + ExecutionMode::Sync => match policy.action(store, relative, ExecutionMode::Sync) { + Action::Execute(id) => { + store.add_trigger(id, ExecutionTrigger::OnSync); + id + } + _ => { + let plan = ExecutionPlan { + operations: relative.to_vec(), + triggers: vec![ExecutionTrigger::OnSync], + optimization, + }; + store.add(plan) + } + }, + } + } +} diff --git a/crates/burn-fusion/src/stream/execution/tests.rs b/crates/burn-fusion/src/stream/execution/tests.rs new file mode 100644 index 0000000..8740b50 --- /dev/null +++ b/crates/burn-fusion/src/stream/execution/tests.rs @@ -0,0 +1,765 @@ +//! A testing module that ensures the correctness of the explorer, policy, and processor. +//! +//! The primary focus is on validating the seamless interaction between these three components to +//! execute and optimize a stream of operations accurately. +//! +//! To test these components effectively, we create mock types for the stream, optimization, +//! optimization builder, and stream segment. These mock types aid in comprehensively +//! understanding the process of optimizing streams. +use std::sync::Arc; + +use burn_backend::{DType, Shape}; +use burn_ir::{ + BinaryOpIr, FloatOperationIr, NumericOperationIr, OperationIr, ScalarIr, ScalarOpIr, TensorId, + TensorIr, TensorStatus, UnaryOpIr, +}; + +use crate::{ + FuserProperties, FuserStatus, NumOperations, OperationFuser, + search::BlockOptimization, + stream::store::{ + ExecutionPlan, ExecutionPlanId, ExecutionPlanStore, ExecutionStrategy, ExecutionTrigger, + }, +}; + +use super::*; + +/// A fake stream of operations for testing purpose. +pub struct TestStream { + processor: Processor, + store: ExecutionPlanStore, + executed: Vec, + operations: Vec, +} + +/// A fake [optimization builder](OptimizationBuilder) for testing purpose. +/// +/// The optimizer tries to fuse only the `expected_operations` if they appear +/// in the operations queue +#[derive(Clone)] +pub struct TestOptimizationBuilder { + builder_id: usize, + expected_operations: Vec, + actual: Vec, +} + +/// A fake optimization for testing purpose. +#[derive(new, Debug, PartialEq)] +pub struct TestOptimization { + builder_id: usize, + size: usize, +} + +impl NumOperations for TestOptimization { + fn len(&self) -> usize { + self.size + } + fn name(&self) -> &'static str { + "TestOptimization" + } +} + +/// A fake [stream segment](StreamSegment) for testing purpose. +#[derive(new)] +pub struct TestSegment<'i> { + operations: &'i mut Vec, + executed: &'i mut Vec, +} + +impl ExecutionStrategy { + /// Create an ordered execution strategy with the given size. + pub fn operations(size: usize) -> Self { + Self::Operations { + ordering: Arc::new((0..size).collect()), + } + } +} + +impl ExecutionStrategy { + /// Only use it for testing, to easily create ordered strategies. + pub fn optimization(opt: TestOptimization) -> Self { + let ordering = Arc::new((0..opt.size).collect()); + Self::Optimization { + opt, + ordering, + score: 1, + } + } +} + +/// This is a substantial test case that examines a lengthy scenario with a diverse set of conditions. +/// +/// While it's usually preferable to split tests into multiple independent scenarios, in this case, it is +/// crucial to verify that the stream's state is correctly updated when various cases occur consecutively. +#[test] +fn should_support_complex_stream() { + // We have 2 different optimization builders in this test case. + let builder_id_1 = 0; + let builder_id_2 = 1; + + // We will have a total of 3 execution plans to execute. + let plan_id_1 = 0; + let plan_id_2 = 1; + let plan_id_3 = 2; + + let builder_1 = TestOptimizationBuilder::new(builder_id_1, vec![operation_1(), operation_2()]); + let builder_2 = TestOptimizationBuilder::new(builder_id_2, vec![operation_2(), operation_2()]); + let mut stream = TestStream::new(vec![Box::new(builder_1), Box::new(builder_2)]); + + // builder_1 is still waiting to see next op is operation_2 + // builder_2 is closed because it's not the right operation + stream.add(operation_1()); + stream.assert_number_of_operations(1); + stream.assert_number_of_executions(0); + + // No optimization found for the first two operations. + stream.add(operation_1()); + stream.assert_number_of_operations(0); + stream.assert_number_of_executions(1); + stream.assert_last_executed(plan_id_1); + stream.assert_plan( + plan_id_1, + ExecutionPlan { + operations: vec![operation_1(), operation_1()], + triggers: vec![ExecutionTrigger::Always], + optimization: BlockOptimization::new(ExecutionStrategy::operations(2), Vec::new()), + }, + ); + + // Nothing to execute. + stream.add(operation_1()); + stream.assert_number_of_operations(1); + stream.assert_number_of_executions(1); + + // Now we should trigger the first optimization builder. + stream.add(operation_2()); + stream.assert_number_of_operations(0); + stream.assert_number_of_executions(2); + stream.assert_last_executed(plan_id_2); + stream.assert_plan( + plan_id_2, + ExecutionPlan { + operations: vec![operation_1(), operation_2()], + triggers: vec![ExecutionTrigger::Always], + optimization: BlockOptimization::new( + ExecutionStrategy::optimization(TestOptimization::new(builder_id_1, 2)), + vec![0, 1], + ), + }, + ); + + // Nothing to execute. + stream.add(operation_2()); + stream.assert_number_of_operations(1); + stream.assert_number_of_executions(2); + + // Now we should trigger the second optimization builder. + stream.add(operation_2()); + stream.assert_number_of_operations(0); + stream.assert_number_of_executions(3); + stream.assert_last_executed(plan_id_3); + stream.assert_plan( + plan_id_3, + ExecutionPlan { + operations: vec![operation_2(), operation_2()], + triggers: vec![ExecutionTrigger::Always], + optimization: BlockOptimization { + strategy: ExecutionStrategy::optimization(TestOptimization::new(builder_id_2, 2)), + ordering: vec![0, 1], + }, + }, + ); + + // Nothing to execute. + stream.add(operation_1()); + stream.assert_number_of_operations(1); + stream.assert_number_of_executions(3); + + // Now we should trigger the first optimization builder (second plan). + stream.add(operation_2()); + stream.assert_number_of_operations(0); + stream.assert_number_of_executions(4); + stream.assert_last_executed(plan_id_2); + stream.assert_plan( + plan_id_2, + ExecutionPlan { + operations: vec![operation_1(), operation_2()], + triggers: vec![ExecutionTrigger::Always], + optimization: BlockOptimization { + strategy: ExecutionStrategy::optimization(TestOptimization::new(builder_id_1, 2)), + ordering: vec![0, 1], + }, + }, + ); + + // Nothing to execute. + stream.add(operation_2()); + stream.assert_number_of_operations(1); + stream.assert_number_of_executions(4); + + // Now we should trigger the first optimization builder (third plan). + stream.add(operation_2()); + stream.assert_number_of_operations(0); + stream.assert_number_of_executions(5); + stream.assert_last_executed(plan_id_3); +} + +/// In this scenario we will never use an optimization, but we check that we reuse the execution plan stored. +#[test] +fn should_reuse_basic_operations() { + let builder_id_1 = 0; + let plan_id_1 = 0; + let plan_id_2 = 1; + + let builder_1 = TestOptimizationBuilder::new(builder_id_1, vec![operation_1(), operation_2()]); + let mut stream = TestStream::new(vec![Box::new(builder_1)]); + + stream.add(operation_3()); + stream.assert_last_executed(plan_id_1); + stream.assert_number_of_operations(0); + stream.assert_plan( + plan_id_1, + ExecutionPlan { + operations: vec![operation_3()], + triggers: vec![ExecutionTrigger::Always], + optimization: BlockOptimization { + strategy: ExecutionStrategy::operations(1), + ordering: vec![0], + }, + }, + ); + + stream.add(operation_3()); + stream.assert_last_executed(plan_id_1); + stream.assert_number_of_operations(0); + stream.assert_plan( + plan_id_1, + ExecutionPlan { + operations: vec![operation_3()], + triggers: vec![ExecutionTrigger::Always], + optimization: BlockOptimization { + strategy: ExecutionStrategy::operations(1), + ordering: vec![0], + }, + }, + ); + + // Lazy try to build optimization 1. + stream.add(operation_1()); + // But not possible. + stream.add(operation_3()); + + // Creates a new plan with both operations. + stream.assert_plan( + plan_id_2, + ExecutionPlan { + operations: vec![operation_1(), operation_3()], + triggers: vec![ExecutionTrigger::Always], + optimization: BlockOptimization { + strategy: ExecutionStrategy::operations(2), + ordering: vec![0], + }, + }, + ); + stream.assert_number_of_operations(0); + stream.assert_last_executed(plan_id_2); +} + +// In this scenario we validate that we support multiple optimization builders with overlapping +// operations. +// +// This is a very long scenario that validates a lot of things. +#[test] +fn should_support_overlapping_optimizations() { + // We have 2 different optimization builders in this test case. + let builder_id_1 = 0; + let builder_id_2 = 0; + + // We will have a total of 5 execution plans to execute. + let plan_id_1 = 0; + let plan_id_2 = 1; + let plan_id_3 = 2; + let plan_id_4 = 3; + let plan_id_5 = 4; + + let builder_1 = TestOptimizationBuilder::new(builder_id_1, vec![operation_1(), operation_2()]); + let builder_2 = TestOptimizationBuilder::new( + builder_id_2, + vec![operation_1(), operation_2(), operation_1(), operation_1()], + ); + let mut stream = TestStream::new(vec![Box::new(builder_1), Box::new(builder_2)]); + + stream.add(operation_1()); + stream.assert_number_of_operations(1); + stream.assert_number_of_executions(0); + + stream.add(operation_2()); + stream.assert_number_of_operations(2); + stream.assert_number_of_executions(0); + + stream.add(operation_1()); + stream.assert_number_of_operations(3); + stream.assert_number_of_executions(0); + + stream.add(operation_2()); + stream.assert_number_of_operations(2); + stream.assert_number_of_executions(1); + stream.assert_last_executed(plan_id_1); + stream.assert_plan( + plan_id_1, + ExecutionPlan { + operations: vec![operation_1(), operation_2()], + triggers: vec![ExecutionTrigger::OnOperations(vec![ + operation_1(), + operation_2(), + ])], + optimization: BlockOptimization { + strategy: ExecutionStrategy::optimization(TestOptimization::new(builder_id_1, 2)), + ordering: vec![0, 1], + }, + }, + ); + + stream.add(operation_2()); + stream.assert_number_of_operations(0); + stream.assert_number_of_executions(3); + stream.assert_plan( + plan_id_1, + ExecutionPlan { + operations: vec![operation_1(), operation_2()], + triggers: vec![ + ExecutionTrigger::OnOperations(vec![operation_1(), operation_2()]), + ExecutionTrigger::OnOperations(vec![operation_2()]), + ], + optimization: BlockOptimization { + strategy: ExecutionStrategy::optimization(TestOptimization::new(builder_id_1, 2)), + ordering: vec![0, 1], + }, + }, + ); + stream.assert_plan( + plan_id_2, + ExecutionPlan { + operations: vec![operation_2()], + triggers: vec![ExecutionTrigger::Always], + optimization: BlockOptimization { + strategy: ExecutionStrategy::operations(1), + ordering: vec![0], + }, + }, + ); + + stream.add(operation_1()); + stream.assert_number_of_operations(1); + stream.assert_number_of_executions(3); + + stream.add(operation_2()); + stream.assert_number_of_operations(2); + stream.assert_number_of_executions(3); + + stream.add(operation_1()); + stream.assert_number_of_operations(3); + stream.assert_number_of_executions(3); + + stream.add(operation_1()); + stream.assert_number_of_operations(0); + stream.assert_number_of_executions(4); + + stream.assert_plan( + plan_id_3, + ExecutionPlan { + operations: vec![operation_1(), operation_2(), operation_1(), operation_1()], + triggers: vec![ExecutionTrigger::Always], + optimization: BlockOptimization { + strategy: ExecutionStrategy::optimization(TestOptimization::new(builder_id_1, 4)), + ordering: vec![0], + }, + }, + ); + + stream.add(operation_1()); + stream.assert_number_of_operations(1); + stream.assert_number_of_executions(4); + + stream.add(operation_2()); + stream.assert_number_of_operations(2); + stream.assert_number_of_executions(4); + + stream.add(operation_1()); + stream.assert_number_of_operations(3); + stream.assert_number_of_executions(4); + + stream.sync(); + stream.assert_number_of_operations(0); + // The sync executes a single plan covering the whole remaining segment: the + // fused pair composed with the un-fused trailing operation. + stream.assert_number_of_executions(5); + stream.assert_plan( + plan_id_1, + ExecutionPlan { + operations: vec![operation_1(), operation_2()], + triggers: vec![ + ExecutionTrigger::OnOperations(vec![operation_1(), operation_2()]), + ExecutionTrigger::OnOperations(vec![operation_2()]), + ], + optimization: BlockOptimization { + strategy: ExecutionStrategy::optimization(TestOptimization::new(builder_id_1, 2)), + ordering: vec![0, 1], + }, + }, + ); + stream.assert_plan( + plan_id_4, + ExecutionPlan { + operations: vec![operation_1(), operation_2(), operation_1()], + triggers: vec![ExecutionTrigger::OnSync], + optimization: BlockOptimization { + strategy: ExecutionStrategy::Composed(vec![ + Box::new(ExecutionStrategy::optimization(TestOptimization::new( + builder_id_1, + 2, + ))), + Box::new(ExecutionStrategy::Operations { + ordering: Arc::new(vec![2]), + }), + ]), + ordering: vec![0, 1, 2], + }, + }, + ); + + stream.add(operation_3()); + stream.assert_last_executed(plan_id_5); + stream.assert_plan( + plan_id_5, + ExecutionPlan { + operations: vec![operation_3()], + triggers: vec![ExecutionTrigger::Always], + optimization: BlockOptimization { + strategy: ExecutionStrategy::operations(1), + ordering: vec![0], + }, + }, + ); + + stream.add(operation_3()); + stream.assert_last_executed(plan_id_5); +} + +/// A plan discovered at a sync point must cover the whole segment and be reusable on +/// the next identical pass without re-exploring. +/// +/// This is the shape of an autoregressive inference step: the graph defers past the +/// fusable prefix because some fuser is still open, and the step ends on a sync. The +/// search leaves the still-open tail out of its optimization (in lazy mode the tail +/// seeds the next round), but at a sync there is no next round: the tail is folded +/// into the plan as un-fused operations in a composed strategy. The plan then matches +/// the whole segment, so every subsequent identical step executes it directly instead +/// of re-running the whole exploration. +#[test] +fn should_reuse_sync_plan_on_next_pass() { + let plan_id = 0; + + let pair_builder = TestOptimizationBuilder::new(0, vec![operation_1(), operation_2()]); + // A builder that wants a longer sequence: it stays open past the tail op, which is + // what defers the whole segment until the sync. + let hungry_builder = TestOptimizationBuilder::new( + 1, + vec![operation_1(), operation_2(), operation_3(), operation_1()], + ); + let mut stream = TestStream::new(vec![Box::new(pair_builder), Box::new(hungry_builder)]); + + // First pass: everything defers, the sync explores once and executes one plan + // covering the whole segment. + stream.add(operation_1()); + stream.add(operation_2()); + stream.add(operation_3()); + stream.assert_number_of_executions(0); + stream.sync(); + stream.assert_number_of_operations(0); + stream.assert_number_of_executions(1); + stream.assert_last_executed(plan_id); + stream.assert_number_of_explorations(1); + stream.assert_plan( + plan_id, + ExecutionPlan { + operations: vec![operation_1(), operation_2(), operation_3()], + triggers: vec![ExecutionTrigger::OnSync], + optimization: BlockOptimization { + strategy: ExecutionStrategy::Composed(vec![ + Box::new(ExecutionStrategy::optimization(TestOptimization::new(0, 2))), + Box::new(ExecutionStrategy::Operations { + ordering: Arc::new(vec![2]), + }), + ]), + ordering: vec![0, 1, 2], + }, + }, + ); + + // Second pass: the identical stream ends at the same sync; the cached plan + // matches the whole segment and executes without any re-exploration. + stream.add(operation_1()); + stream.add(operation_2()); + stream.add(operation_3()); + stream.assert_number_of_executions(1); + stream.sync(); + stream.assert_number_of_operations(0); + stream.assert_number_of_executions(2); + stream.assert_last_executed(plan_id); + stream.assert_number_of_explorations(1); +} + +impl TestStream { + /// Create a new stream with the given optimization builders. + fn new(optimizations: Vec>>) -> Self { + Self { + processor: Processor::::new(optimizations), + store: ExecutionPlanStore::::new(), + executed: Vec::new(), + operations: Vec::new(), + } + } + + /// Add an operation to the stream. + fn add(&mut self, operation: OperationIr) { + self.operations.push(operation); + self.processor.process( + TestSegment::new(&mut self.operations, &mut self.executed), + &mut self.store, + ExecutionMode::Lazy, + ); + } + + /// Sync the stream. + fn sync(&mut self) { + self.processor.process( + TestSegment::new(&mut self.operations, &mut self.executed), + &mut self.store, + ExecutionMode::Sync, + ); + } + + /// Assert that the plan has been executed as provided. + fn assert_plan(&self, id: ExecutionPlanId, expected: ExecutionPlan) { + let actual = self.store.get_unchecked(id); + assert_eq!(actual.operations, expected.operations, "Same operations"); + assert_eq!(actual.triggers, expected.triggers, "Same triggers"); + } + + /// Assert that the given plan id has been the last executed. + fn assert_last_executed(&self, id: ExecutionPlanId) { + match self.executed.last() { + Some(last_id) => assert_eq!(*last_id, id), + None => panic!("No plan has been executed"), + } + } + + /// Assert the number of executions since the start of the stream. + fn assert_number_of_executions(&self, number: usize) { + assert_eq!(self.executed.len(), number, "Number of execution match"); + } + + /// Assert the number of operations queued. + fn assert_number_of_operations(&self, number: usize) { + assert_eq!(self.operations.len(), number); + } + + /// Assert the number of optimizations built since the start of the stream; + /// reusing a cached plan does not count. + fn assert_number_of_explorations(&self, number: usize) { + assert_eq!( + self.processor.num_explorations(), + number, + "Number of explorations match" + ); + } +} + +impl TestOptimizationBuilder { + /// Create a new optimization builder that follows a pattern with a trigger. + pub fn new(builder_id: usize, operations: Vec) -> Self { + Self { + builder_id, + expected_operations: operations, + actual: Vec::new(), + } + } +} + +impl OperationFuser for TestOptimizationBuilder { + /// Register a new operation. + fn fuse(&mut self, operation: &OperationIr) { + self.actual.push(operation.clone()); + } + + /// Build the optimization. + fn finish(&mut self) -> TestOptimization { + TestOptimization::new(self.builder_id, self.len()) + } + + /// Reset the state. + fn reset(&mut self) { + self.actual.clear(); + } + + /// Return the optimization status. + fn status(&self) -> FuserStatus { + if self.actual.len() < self.expected_operations.len() { + let operations = &self.expected_operations[0..self.actual.len()]; + + return match self.actual == operations { + // Still optimizing. + true => FuserStatus::Open, + // Never gonna be possible on that stream. + false => FuserStatus::Closed, + }; + } + + FuserStatus::Closed + } + + /// Return the properties of this optimization. + fn properties(&self) -> FuserProperties { + if self.actual.len() < self.expected_operations.len() { + // Optimization not possible. + return FuserProperties { + score: 0, + ready: false, + }; + } + + let stream_is_ok = + self.actual[0..self.expected_operations.len()] == self.expected_operations; + + if !stream_is_ok { + // Optimization not possible. + return FuserProperties { + score: 0, + ready: false, + }; + } + + // Optimization possible. + FuserProperties { + score: self.expected_operations.len() as u64, + ready: true, + } + } + + // The number of operations that should be handle by the optimization. + fn len(&self) -> usize { + self.expected_operations.len() + } + fn clone_dyn(&self) -> Box> { + Box::new(self.clone()) + } +} + +impl StreamSegment for TestSegment<'_> { + // The operations in the process. + fn operations(&self) -> &[OperationIr] { + self.operations + } + + // Execute the process. + fn execute(&mut self, id: ExecutionPlanId, store: &mut ExecutionPlanStore) { + let execution_plan = store.get_unchecked(id); + + self.execute_strategy(&execution_plan.optimization.strategy); + + self.executed.push(id); + } + + fn execute_unfused(&mut self, optimization: BlockOptimization) { + self.execute_strategy(&optimization.strategy); + } +} + +impl TestSegment<'_> { + fn execute_strategy(&mut self, strategy: &ExecutionStrategy) { + match strategy { + ExecutionStrategy::Optimization { opt, .. } => { + self.operations.drain(0..opt.size); + } + ExecutionStrategy::Operations { ordering } => { + self.operations.drain(0..ordering.len()); + } + ExecutionStrategy::Composed(strategies) => { + for strategy in strategies { + self.execute_strategy(strategy); + } + } + } + } +} + +/// Just a simple operation. +pub fn operation_1() -> OperationIr { + OperationIr::NumericFloat( + DType::F32, + NumericOperationIr::Add(BinaryOpIr { + lhs: TensorIr { + id: TensorId::new(0), + shape: Shape::new([32, 32]), + status: TensorStatus::ReadOnly, + dtype: DType::F32, + }, + rhs: TensorIr { + id: TensorId::new(1), + shape: Shape::new([32, 32]), + status: TensorStatus::ReadOnly, + dtype: DType::F32, + }, + out: TensorIr { + id: TensorId::new(2), + shape: Shape::new([32, 32]), + status: TensorStatus::NotInit, + dtype: DType::F32, + }, + }), + ) +} + +/// Just a simple operation. +pub fn operation_2() -> OperationIr { + OperationIr::NumericFloat( + DType::F32, + NumericOperationIr::AddScalar(ScalarOpIr { + lhs: TensorIr { + id: TensorId::new(0), + shape: Shape::new([32, 32]), + status: TensorStatus::ReadOnly, + dtype: DType::F32, + }, + rhs: ScalarIr::Float(5.0), + out: TensorIr { + id: TensorId::new(2), + shape: Shape::new([32, 32]), + status: TensorStatus::NotInit, + dtype: DType::F32, + }, + }), + ) +} + +/// Just a simple operation. +pub fn operation_3() -> OperationIr { + OperationIr::Float( + DType::F32, + FloatOperationIr::Log(UnaryOpIr { + input: TensorIr { + id: TensorId::new(0), + shape: Shape::new([32, 32]), + status: TensorStatus::ReadOnly, + dtype: DType::F32, + }, + out: TensorIr { + id: TensorId::new(0), + shape: Shape::new([32, 32]), + status: TensorStatus::NotInit, + dtype: DType::F32, + }, + }), + ) +} diff --git a/crates/burn-fusion/src/stream/execution/trace.rs b/crates/burn-fusion/src/stream/execution/trace.rs new file mode 100644 index 0000000..c91fea0 --- /dev/null +++ b/crates/burn-fusion/src/stream/execution/trace.rs @@ -0,0 +1,383 @@ +//! Full-level fusion logging and the shared data model used by both the logger and +//! the test-only [`FusionInspector`](crate::inspect::FusionInspector). +//! +//! A single execution of a [`ExecutionStrategy`] is flattened into a sequence of +//! [`FusionBlock`]s — one per strategy leaf — each tagged with a [`BlockKind`]. The +//! logger renders these as a human-readable table; the inspector stores them in +//! `FusionReport`s so tests can assert on what fused together. + +use burn_ir::{OperationIr, TensorIr}; +use burn_std::config::{fusion::FusionLogLevel, log_fusion}; +use core::fmt::Write; + +use crate::NumOperations; +use crate::stream::{StreamId, store::ExecutionStrategy}; + +/// One contiguous run of operations — the output of a single [`ExecutionStrategy`] +/// leaf. A `Composed` strategy flattens into a sequence of these. +#[derive(Debug, Clone)] +pub struct FusionBlock { + /// Whether this block ran fused or unfused. + pub kind: BlockKind, + /// The operations contained in the block, in execution order. + pub operations: Vec, +} + +/// Whether a [`FusionBlock`] was fused into a single kernel or ran op-by-op. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum BlockKind { + /// The operations were fused into a single kernel produced by the named fuser. + Fused { + /// The name of the fuser / optimization (e.g. `"ElementWise"`). + name: &'static str, + /// The score the fuser reported for this kernel. + score: u64, + }, + /// The operations ran individually; no fusion happened for this block. + Unfused, +} + +/// Emit an execution table for the given `strategy` at [`FusionLogLevel::Full`]. +/// +/// `global` is the `OperationQueue::global` slice at the time of execution — it holds the +/// real tensor IDs and shapes that the indices inside the strategy's `ordering` fields +/// point into. The table is built lazily: no allocation happens when the log level is +/// below `Full` (except when the inspector is installed, which eagerly captures the +/// structured blocks for test assertions). +pub(crate) fn log_execution_table( + _id: StreamId, + strategy: &ExecutionStrategy, + global: &[OperationIr], +) { + // When the inspector is active, build the blocks eagerly so we can emit a + // structured report even when fusion logging is disabled. Otherwise stay lazy. + #[cfg(feature = "test-util")] + let blocks_for_inspector: Option> = if crate::inspect::is_installed() { + let mut blocks = Vec::new(); + collect_blocks(strategy, global, &mut blocks); + crate::inspect::emit(_id, &blocks); + Some(blocks) + } else { + None + }; + + log_fusion(FusionLogLevel::Full, || { + #[cfg(feature = "test-util")] + if let Some(blocks) = &blocks_for_inspector { + return format_table(blocks); + } + let mut blocks = Vec::new(); + collect_blocks(strategy, global, &mut blocks); + format_table(&blocks) + }); +} + +pub(crate) fn collect_blocks( + strategy: &ExecutionStrategy, + global: &[OperationIr], + blocks: &mut Vec, +) { + match strategy { + ExecutionStrategy::Optimization { + opt, + ordering, + score, + } => { + blocks.push(FusionBlock { + kind: BlockKind::Fused { + name: opt.name(), + score: *score, + }, + operations: ordering.iter().map(|&i| global[i].clone()).collect(), + }); + } + ExecutionStrategy::Operations { ordering } => { + blocks.push(FusionBlock { + kind: BlockKind::Unfused, + operations: ordering.iter().map(|&i| global[i].clone()).collect(), + }); + } + ExecutionStrategy::Composed(items) => { + for item in items { + collect_blocks(item, global, blocks); + } + } + } +} + +pub(crate) fn format_table(blocks: &[FusionBlock]) -> String { + let total: usize = blocks.iter().map(|b| b.operations.len()).sum(); + if total == 0 { + return String::from("fusion block: "); + } + + // One row per op across all blocks. The `block` column carries the block header + // on the first row of each block and is left blank for the rest, so each fused + // block is visually grouped but the name/score appear exactly once. + struct Row { + idx: String, + block: String, + op: String, + inputs: String, + outputs: String, + } + + let headers = ["idx", "block", "op", "inputs", "outputs"]; + let mut widths = headers.map(str::len); + + let mut rows: Vec = Vec::with_capacity(total); + let mut global_idx: usize = 0; + for block in blocks { + let header = block_header(&block.kind, block.operations.len()); + for (i, op) in block.operations.iter().enumerate() { + let row = Row { + idx: global_idx.to_string(), + block: if i == 0 { + header.clone() + } else { + String::new() + }, + op: op_kind(op), + inputs: format_tensors(op.inputs()), + outputs: format_tensors(op.outputs()), + }; + widths[0] = widths[0].max(row.idx.len()); + widths[1] = widths[1].max(row.block.len()); + widths[2] = widths[2].max(row.op.len()); + widths[3] = widths[3].max(row.inputs.len()); + widths[4] = widths[4].max(row.outputs.len()); + rows.push(row); + global_idx += 1; + } + } + + // Two-space separators between the 5 columns. + let table_width = widths.iter().sum::() + 2 * (widths.len() - 1); + + let mut out = String::new(); + writeln!(out, "{}", "═".repeat(table_width)).unwrap(); + writeln!( + out, + "fusion trace ({total} op{}, {} block{})", + if total == 1 { "" } else { "s" }, + blocks.len(), + if blocks.len() == 1 { "" } else { "s" }, + ) + .unwrap(); + + let write_row = |out: &mut String, cols: [&str; 5]| { + for (i, c) in cols.iter().enumerate() { + if i > 0 { + out.push_str(" "); + } + write!(out, "{: String { + let ops_suffix = if size == 1 { "op" } else { "ops" }; + match kind { + BlockKind::Fused { name, score } => { + format!("▸ fused {name} (score={score}, {size} {ops_suffix})") + } + BlockKind::Unfused => format!("▸ un-fused ({size} {ops_suffix})"), + } +} + +/// Take the top-level variant name from a Debug representation (`"Foo(..)"` → `"Foo"`). +fn debug_head(s: &str) -> &str { + let end = s.find(['(', '{', ' ']).unwrap_or(s.len()); + &s[..end] +} + +/// Produce a "Outer::Inner" kind string for every variant that has an inner enum, +/// so the table shows the concrete operation (e.g. `BaseFloat::Reshape`) rather than +/// just the category. +pub(crate) fn op_kind(op: &OperationIr) -> String { + fn inner(s: String) -> String { + debug_head(&s).to_string() + } + match op { + OperationIr::BaseFloat(x) => format!("BaseFloat::{}", inner(format!("{x:?}"))), + OperationIr::BaseInt(x) => format!("BaseInt::{}", inner(format!("{x:?}"))), + OperationIr::BaseBool(x) => format!("BaseBool::{}", inner(format!("{x:?}"))), + OperationIr::NumericFloat(_, x) => format!("NumericFloat::{}", inner(format!("{x:?}"))), + OperationIr::NumericInt(_, x) => format!("NumericInt::{}", inner(format!("{x:?}"))), + OperationIr::Bool(x) => format!("Bool::{}", inner(format!("{x:?}"))), + OperationIr::Int(x) => format!("Int::{}", inner(format!("{x:?}"))), + OperationIr::Float(_, x) => format!("Float::{}", inner(format!("{x:?}"))), + OperationIr::Module(x) => format!("Module::{}", inner(format!("{x:?}"))), + OperationIr::Init(x) => format!("Init::{}", inner(format!("{x:?}"))), + OperationIr::Custom(_) => "Custom".to_string(), + OperationIr::Drop(_) => "Drop".to_string(), + OperationIr::Distributed(x) => format!("Distributed::{}", inner(format!("{x:?}"))), + OperationIr::Activation(x) => format!("Activation::{}", inner(format!("{x:?}"))), + } +} + +fn format_tensors<'a, I: Iterator>(tensors: I) -> String { + let mut out = String::new(); + let mut first = true; + for t in tensors { + if !first { + out.push_str(", "); + } + first = false; + write!(out, "t{}:", t.id.value()).unwrap(); + write_dtype(&mut out, &t.dtype); + write_shape(&mut out, &t.shape); + } + out +} + +fn write_dtype(out: &mut String, dtype: &burn_backend::DType) { + // Debug gives "F32", "Bool(Native)", etc. Lowercase the initial identifier so the + // output reads like a Rust type annotation; keep anything after the first paren/brace + // (e.g. Bool variants or QFloat scheme details) verbatim. + let dbg = format!("{dtype:?}"); + let split = dbg.find(['(', '{', ' ']).unwrap_or(dbg.len()); + let (head, tail) = dbg.split_at(split); + for c in head.chars() { + out.push(c.to_ascii_lowercase()); + } + out.push_str(tail); +} + +fn write_shape(out: &mut String, shape: &burn_backend::Shape) { + out.push('['); + let mut first = true; + for d in shape.iter() { + if !first { + out.push(','); + } + first = false; + write!(out, "{d}").unwrap(); + } + out.push(']'); +} + +#[cfg(test)] +mod tests { + use super::*; + use burn_backend::{DType, Shape}; + use burn_ir::{BaseOperationIr, TensorId, TensorIr, TensorStatus, UnaryOpIr}; + + fn tensor(id: u64, dims: &[usize]) -> TensorIr { + TensorIr { + id: TensorId::new(id), + shape: Shape::new_raw(dims.iter().copied().collect()), + status: TensorStatus::ReadOnly, + dtype: DType::F32, + } + } + + #[test] + fn tensor_format_is_compact() { + let t = tensor(42, &[2, 3, 4]); + let out = format_tensors(core::iter::once(&t)); + assert_eq!(out, "t42:f32[2,3,4]"); + } + + #[test] + fn tensor_format_multiple_tensors_comma_separated() { + let a = tensor(1, &[8]); + let b = tensor(2, &[]); + let out = format_tensors([&a, &b].into_iter()); + assert_eq!(out, "t1:f32[8], t2:f32[]"); + } + + #[test] + fn op_kind_includes_inner_variant() { + let input = tensor(1, &[4, 2]); + let out = tensor(2, &[2, 4]); + let op = OperationIr::BaseFloat(BaseOperationIr::Permute(burn_ir::PermuteOpIr { + input, + out, + axes: vec![1, 0], + })); + assert_eq!(op_kind(&op), "BaseFloat::Permute"); + } + + #[test] + fn op_kind_handles_unary_inside_float() { + let input = tensor(1, &[4]); + let out = tensor(2, &[4]); + let op = OperationIr::Float( + DType::F32, + burn_ir::FloatOperationIr::Exp(UnaryOpIr { input, out }), + ); + assert_eq!(op_kind(&op), "Float::Exp"); + } + + #[test] + fn format_table_puts_name_and_score_in_block_header() { + let op = OperationIr::Float( + DType::F32, + burn_ir::FloatOperationIr::Exp(UnaryOpIr { + input: tensor(1, &[4]), + out: tensor(2, &[4]), + }), + ); + let blocks = vec![ + FusionBlock { + kind: BlockKind::Fused { + name: "FusedKernel", + score: 42, + }, + operations: vec![op.clone(), op.clone()], + }, + FusionBlock { + kind: BlockKind::Unfused, + operations: vec![op], + }, + ]; + + let table = format_table(&blocks); + + // Top separator precedes the title line. + let first_line = table.lines().next().unwrap(); + assert!( + first_line.chars().all(|c| c == '═'), + "expected top line of ═, got {first_line:?}" + ); + assert!(table.contains("\nfusion trace (3 ops, 2 blocks)\n")); + // Fused header names the optimization and its score exactly once. + assert!(table.contains("▸ fused FusedKernel (score=42, 2 ops)")); + // Un-fused header is tagged and sized. + assert!(table.contains("▸ un-fused (1 op)")); + // No per-row repetition of "FusedKernel" or "42": they appear exactly once. + assert_eq!(table.matches("FusedKernel").count(), 1); + assert_eq!(table.matches("score=42").count(), 1); + // Indices are global (0, 1, 2 across blocks). + assert!(table.contains("\n0 ")); + assert!(table.contains("\n1 ")); + assert!(table.contains("\n2 ")); + } +} diff --git a/crates/burn-fusion/src/stream/execution/validator.rs b/crates/burn-fusion/src/stream/execution/validator.rs new file mode 100644 index 0000000..2095a9d --- /dev/null +++ b/crates/burn-fusion/src/stream/execution/validator.rs @@ -0,0 +1,136 @@ +use burn_ir::OperationIr; + +use crate::stream::store::{ExecutionPlanId, ExecutionPlanStore, ExecutionTrigger}; + +/// Compare each operation in the list of operations provided by the [store](OperationsStore) +/// to verify if the newly added operations match the original list. +/// +/// It is used by the [policy](crate::stream::execution::Policy) to check each candidate as well +/// as to verify if a list of operations is optimal to execute based on their triggers. +#[derive(Debug)] +pub(crate) struct OperationsValidator { + /// The ID used to retrieve the operation list. + pub(crate) id: ID, + /// The current [state](MatchingState). + pub(crate) state: ValidatorState, +} + +/// The state of the validator. +#[derive(Debug)] +pub(crate) enum ValidatorState { + /// A matching operation list has been found. + Found { size: usize }, + /// No matching operation list has been found. + Invalidated, + /// Potentially going to find a matching operation list when more operations are added. + Validating, +} + +/// Provides a list of operations based on an Id. +pub(crate) trait OperationsStore { + /// The type used for the identifier. + type Id: Copy; + + /// retrieve the list of operations corresponding on the provided id. + fn get(&self, id: Self::Id) -> &[OperationIr]; +} + +impl OperationsValidator { + /// Create a new validator. + pub(crate) fn new(id: ID) -> Self { + Self { + id, + state: ValidatorState::Validating, + } + } + + /// Update the state of the validator based on the newly added operation. + pub(crate) fn update(&mut self, added: &OperationIr, added_position: usize, store: &S) + where + S: OperationsStore, + ID: PartialEq + Copy, + { + match &self.state { + ValidatorState::Found { size: _ } => return, + ValidatorState::Invalidated => return, + ValidatorState::Validating => {} + }; + + let item = store.get(self.id); + let operation_candidate = match item.get(added_position) { + Some(val) => val, + None => { + self.state = ValidatorState::Invalidated; + return; + } + }; + + if operation_candidate != added { + self.state = ValidatorState::Invalidated; + return; + } + + // Finished + if item.len() == added_position + 1 { + self.state = ValidatorState::Found { size: item.len() }; + } + } +} + +/// [Operations store](OperationsStore) used to retrieve the list of operations for a trigger. +#[derive(new)] +pub(crate) struct TriggerOperationsStore<'a, O> { + id: ExecutionPlanId, + store: &'a ExecutionPlanStore, +} + +/// Validates when operations match a trigger. +#[derive(Debug)] +pub(crate) enum TriggerValidator { + OnOperations { + matching: OperationsValidator, + progress: TriggerProgress, + }, + Always, + OnSync, +} + +/// The progress made into the trigger validation process. +#[derive(Debug)] +pub(crate) enum TriggerProgress { + /// When the validation hasn't started. + NotInit, + /// The number of operations that have been checked. + NumChecked(usize), +} + +/// An execution plan can have many triggers, so we use the position in the list to identify a +/// trigger. +pub(crate) type TriggerId = usize; + +impl OperationsStore for TriggerOperationsStore<'_, O> { + type Id = TriggerId; + + fn get(&self, id: Self::Id) -> &[OperationIr] { + match &self.store.get_unchecked(self.id).triggers[id] { + ExecutionTrigger::OnOperations(operations) => operations, + ExecutionTrigger::OnSync => &[], + ExecutionTrigger::Always => &[], + } + } +} + +/// [Operations store](OperationsStore) used to retrieve the list of operations for an +/// [execution plan](crate::stream::store::ExecutionPlan). +#[derive(new)] +pub(crate) struct ExecutionPlanOperationsStore<'a, O> { + store: &'a ExecutionPlanStore, +} + +impl OperationsStore for ExecutionPlanOperationsStore<'_, O> { + type Id = ExecutionPlanId; + + fn get(&self, id: Self::Id) -> &[OperationIr] { + &self.store.get_unchecked(id).operations + } +} diff --git a/crates/burn-fusion/src/stream/memory_checks.rs b/crates/burn-fusion/src/stream/memory_checks.rs new file mode 100644 index 0000000..069d7e8 --- /dev/null +++ b/crates/burn-fusion/src/stream/memory_checks.rs @@ -0,0 +1,247 @@ +use hashbrown::HashMap; +use std::{ + fmt::Display, + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + mpsc::SyncSender, + }, + thread::JoinHandle, + time::Duration, +}; + +use burn_ir::{HandleContainer, TensorId, TensorStatus}; +use burn_std::id::StreamId; + +use crate::FusionRuntime; + +use super::Stream; + +/// Memory checks struct to validate there is no memory leak with the fusion runtime. +#[derive(Clone)] +pub(crate) struct MemoryChecks { + sender: SyncSender, + num_queued: Arc, + // Keeps track of its thread. + _handle: Arc>, +} + +enum Message { + Register(StreamAnalyses), + Check(SyncSender), +} + +enum MemoryReport { + Success, + NotReady, + NotStarted, + Fail(String), +} + +#[derive(Default)] +struct StreamAnalyses { + streams: HashMap, + num_handles: usize, +} + +impl Display for StreamAnalyses { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("\n==== Fusion Memory Report ====\n")?; + f.write_fmt(format_args!(" - Handles: {}\n", self.num_handles))?; + f.write_fmt(format_args!(" - Streams: {}\n", self.streams.len()))?; + + for (id, analysis) in self.streams.iter() { + f.write_fmt(format_args!( + " - {} => operations: {} cursor: {}\n", + id, analysis.num_operations, analysis.cursor + ))?; + for (tid, status) in analysis.variables.iter() { + f.write_fmt(format_args!(" - {tid} => status: {status:?}\n"))?; + } + } + + f.write_str("==============================\n") + } +} + +#[derive(Default, Debug)] +struct Analysis { + variables: HashMap, + num_operations: usize, + cursor: u64, +} + +#[macro_export] +/// Export memory checks tests. +macro_rules! memory_checks { + () => { + #[cfg(test)] + mod memory_checks { + #[test] + fn test_memory_leaks() { + burn_fusion::stream::memory_checks::check_memory_leaks(); + } + } + }; +} + +static INSTANCE: spin::Mutex> = spin::Mutex::new(None); + +/// Performs memory checks and panics if a leak is discovered. +pub fn check_memory_leaks() { + let mut num_try_uninit = 0; + let max_try = 25; + + loop { + let report = fetch_memory_report(); + match report { + MemoryReport::Success => return, + MemoryReport::NotReady => { + num_try_uninit = 0; + std::thread::sleep(Duration::from_millis(100)) + } + MemoryReport::NotStarted => { + if num_try_uninit >= max_try { + // Nothing is running on the fusion runtime. + return; + } + num_try_uninit += 1; + std::thread::sleep(Duration::from_millis(100)) + } + MemoryReport::Fail(msg) => panic!("{msg}"), + } + } +} + +fn fetch_memory_report() -> MemoryReport { + let report = INSTANCE.lock(); + + let report = match report.as_ref() { + Some(client) => client, + None => return MemoryReport::NotStarted, + }; + + let (sender, rec) = std::sync::mpsc::sync_channel(1); + match report.sender.send(Message::Check(sender)) { + Ok(_) => {} + Err(err) => { + panic!("Channel closed can't send the check call: {err:?}") + } + }; + + match rec.recv() { + Ok(report) => report, + Err(err) => panic!("Received an error from fetching check results: {err}"), + } +} + +impl Default for MemoryChecks { + fn default() -> Self { + let mut instance = INSTANCE.lock(); + let result = match instance.as_mut() { + Some(client) => client.clone(), + None => { + let this = Self::spawn_new(); + *instance = Some(this.clone()); + this + } + }; + core::mem::drop(instance); + result + } +} + +impl MemoryChecks { + pub(crate) fn check( + &mut self, + streams: &HashMap>, + handles: &HandleContainer, + ) { + let mut analyses = StreamAnalyses { + num_handles: handles.num_handles(), + streams: Default::default(), + }; + + for (id, s) in streams.iter() { + let analysis = Analysis { + variables: s.queue.variables.clone(), + num_operations: s.queue.global.len(), + cursor: s.cursor, + }; + analyses.streams.insert(*id, analysis); + } + + self.num_queued.fetch_add(1, Ordering::Relaxed); + match self.sender.send(Message::Register(analyses)) { + Ok(..) => {} + Err(err) => { + panic!("Can't register memory checks analysis: {err:?}") + } + } + } + + fn spawn_new() -> Self { + let (sender, rec) = std::sync::mpsc::sync_channel(100); + let num_queued = Arc::new(AtomicU64::new(0)); + let num_queued_moved = num_queued.clone(); + + let handle = std::thread::spawn(move || { + let mut last_analyses = None; + + loop { + let payload = match rec.recv() { + Err(_err) => { + // A client has panic, safe to skip as it may be normal. + continue; + } + Ok(payload) => payload, + }; + match payload { + Message::Register(payload) => { + last_analyses = Some(payload); + num_queued_moved.fetch_sub(1, Ordering::Relaxed); + } + Message::Check(callback) => { + if num_queued_moved.load(Ordering::Relaxed) > 1 { + callback.send(MemoryReport::NotReady).unwrap(); + continue; + } + + // We assume that if nothing has been registered in the last second + // while being at a count of 1, it's the end. + std::thread::sleep(Duration::from_secs(5)); + + if num_queued_moved.load(Ordering::Relaxed) <= 1 { + match last_analyses.take() { + Some(val) => { + callback.send(Self::final_check(val)).unwrap(); + } + None => { + callback + .send(MemoryReport::Fail("No analyses".into())) + .unwrap(); + } + } + } else { + callback.send(MemoryReport::NotReady).unwrap(); + } + } + } + } + }); + + Self { + sender, + num_queued, + _handle: Arc::new(handle), + } + } + + fn final_check(analyses: StreamAnalyses) -> MemoryReport { + if !analyses.streams.is_empty() || analyses.num_handles > 0 { + return MemoryReport::Fail(format!("{analyses}")); + } + + MemoryReport::Success + } +} diff --git a/crates/burn-fusion/src/stream/mod.rs b/crates/burn-fusion/src/stream/mod.rs new file mode 100644 index 0000000..8562e56 --- /dev/null +++ b/crates/burn-fusion/src/stream/mod.rs @@ -0,0 +1,32 @@ +pub(crate) mod execution; +pub(crate) mod queue; +pub(crate) mod store; + +#[cfg(feature = "memory-checks")] +/// Memory checks module. +pub mod memory_checks; + +#[cfg(not(feature = "memory-checks"))] +#[macro_export] +/// Export memory checks tests. +macro_rules! memory_checks { + () => { + #[cfg(test)] + mod memory_checks { + #[ignore = "'memory-checks' disabled"] + #[test] + fn test_memory_leaks() { + // + } + } + }; +} + +mod base; +mod context; +mod multi; + +pub use base::*; +pub use context::*; +pub use execution::*; +pub use multi::*; diff --git a/crates/burn-fusion/src/stream/multi.rs b/crates/burn-fusion/src/stream/multi.rs new file mode 100644 index 0000000..e299648 --- /dev/null +++ b/crates/burn-fusion/src/stream/multi.rs @@ -0,0 +1,301 @@ +use super::{ + StreamId, + execution::{ExecutionMode, Processor, StreamSegment}, + queue::OperationQueue, + store::{ExecutionPlanId, ExecutionPlanStore}, +}; +use crate::{FusionRuntime, UnfusedOp, search::BlockOptimization}; +use burn_ir::{HandleContainer, OperationIr, TensorId}; +use hashbrown::{HashMap, HashSet}; + +/// Keep track of multiple concurrent lazy streams of operations. +/// +/// # Why this exists +/// +/// Each `Stream` holds a lazy queue of [`OperationIr`]s whose inputs are assumed +/// to live on that stream. That makes single-stream execution simple — every +/// `TensorId` in a queue is resolvable from the same handle map and the same +/// pending op chain. But a [`FusionTensor`](crate::FusionTensor) is `Send + Clone`, +/// so user code can move or clone a tensor from one thread (= one [`StreamId`]) to +/// another. The receiving thread will then submit ops whose inputs reference a +/// tensor whose home is a *different* stream's queue. This struct is what makes +/// that case behave correctly without giving up the stream-locality invariant. +/// +/// # Strategy: shared views +/// +/// We never let a foreign-stream tensor id appear in another stream's queue. +/// Instead, when [`FusionTensor::clone`](crate::FusionTensor::clone) or +/// [`FusionTensor::into_ir`](crate::FusionTensor::into_ir) detects that +/// `self.stream != StreamId::current()`, it allocates a fresh id (`dst`) and calls +/// `tag_shared_view` with `(src_stream, src, dst)`. That call does two +/// things, in order: +/// +/// 1. **Materialise `src`.** The id `src` might be the output of an op still +/// pending on `src_stream`. We need its backing handle to actually exist before +/// we can alias it. If [`HandleContainer::get_handle_ref`] returns `None` — +/// meaning no op has produced a handle for `src` yet — we drain `src_stream` +/// synchronously, forcing every pending op to run (and thus the handle to be +/// registered). We also record `src` in `shared_sources` so that any +/// *next* share of the same `src` can skip the drain: once registered, a +/// handle stays put (see invariants below). +/// +/// 2. **Alias the handle under `dst`.** [`HandleContainer::register_handle`] is +/// called with `dst` and a `clone()` of `src`'s backend handle. Cubecl handles +/// are `Arc`-style reference counters over a backing buffer, so `clone()` is +/// cheap and the buffer survives until the last alias drops. After this call, +/// `handles[src]` and `handles[dst]` are two distinct map entries that both +/// point at the same allocation. +/// +/// `shared_view` then returns a new `FusionTensor` carrying `(id = dst, stream = +/// current)`. Every subsequent op on that tensor enqueues on `current` like any +/// other local tensor — the rest of the fusion engine sees no special case. +/// +/// # Freeing +/// +/// Each `FusionTensor::drop` enqueues an `OperationIr::Drop(ir)` on **its own** +/// `stream` field (the home stream of that particular alias), not the calling +/// thread's stream. So: +/// +/// - The original tensor's drop targets `src_stream` and removes `handles[src]`. +/// - The alias tensor's drop targets the stream that minted it and removes +/// `handles[dst]`. +/// +/// Each removal decrements the backend handle's `Arc` refcount; the underlying +/// buffer is freed only after the last side drops. No cross-stream coordination +/// is needed. +/// +/// # Bounding `shared_sources` +/// +/// Naively the set would grow forever, since `tag_shared_view` only ever inserts. +/// Cleanup happens in `register`: as soon as we see an +/// `OperationIr::Drop(ir)` come through, we remove `ir.id` from +/// `shared_sources` immediately — without waiting for the queued `Drop` +/// to actually execute. This is safe because a `Drop` op is registered only +/// after the last live `FusionTensor` with that id has been dropped, so no +/// future `tag_shared_view` can possibly receive that id as a `src`. Removing +/// the entry therefore cannot trigger a redundant drain on any subsequent call. +/// +/// # The SSA-like invariant +/// +/// Skipping the drain on subsequent shares of the same `src` relies on a +/// property of the fusion IR: **every op output uses a fresh `TensorId` +/// allocated by [`crate::Client::create_empty_handle`], never the id of an +/// input.** Once `handles[src]` is set, no later op overwrites it; the data +/// behind `src` is effectively immutable from the IR's point of view. +/// +/// The cubecl-fusion engine *does* sometimes reuse the backing buffer of an +/// input for an output (in-place fusion), but that path is gated by +/// `handle.can_mut()`, which returns false the moment another reference exists. +/// Calling `handle.clone()` in step 2 above is precisely that extra reference, +/// so aliased sources are never eligible for in-place reuse — the engine +/// allocates a fresh output buffer instead. +/// +/// # The chained-share fast path +/// +/// When a share is itself re-shared (owner → peer → grandpeer), the second +/// `tag_shared_view` call has `src = peer's id`. That id was set up directly by +/// the previous call (via `register_handle`), not by an op enqueued on the peer +/// stream, so `handles.get_handle_ref(&src)` is already `Some` the moment we +/// look. The drain check therefore short-circuits — even though `peer's id` was +/// never added to `shared_sources` (only sources that *required* a +/// drain are tracked there), the handle-existence test alone is sufficient. +pub struct MultiStream { + /// Tensor ids that have been the source of a cross-stream share *and* + /// required a drain when first shared. Used by `tag_shared_view` to + /// skip the drain on subsequent shares of the same source. Bounded by + /// pruning in `register` when a `Drop` op for the id is enqueued — + /// see the struct-level docs for the full strategy. + shared_sources: HashSet, + streams: HashMap>, + optimizations: ExecutionPlanStore, + device: R::FusionDevice, + #[cfg(feature = "memory-checks")] + memory_checks: super::memory_checks::MemoryChecks, +} + +impl MultiStream { + pub(crate) fn new(device: R::FusionDevice) -> Self { + Self { + shared_sources: HashSet::new(), + streams: HashMap::new(), + optimizations: ExecutionPlanStore::new(), + device, + #[cfg(feature = "memory-checks")] + memory_checks: super::memory_checks::MemoryChecks::default(), + } + } + + /// Register a new tensor operation on the given `stream`. + pub(crate) fn register( + &mut self, + stream: StreamId, + repr: OperationIr, + operation: UnfusedOp, + handles: &mut HandleContainer, + ) { + // Bound `shared_sources` (see struct-level docs). When the last `FusionTensor` + // for an id is dropped, a `Drop` op is registered here. At that point no live + // `FusionTensor` holds this id, so no future `tag_shared_view` can use it as + // a source — it is safe to drop the entry immediately, without waiting for + // the queued `Drop` op to actually execute. + if let OperationIr::Drop(ir) = &repr { + self.shared_sources.remove(&ir.id); + } + + self.enqueue_operation(stream, repr, operation, handles); + + #[cfg(feature = "memory-checks")] + self.memory_checks.check(&self.streams, handles); + #[cfg(feature = "test-util")] + crate::inspect::emit_handle_snapshot(stream, handles.handle_ids().copied()); + } + + /// Set up a cross-stream alias `dst` for the foreign tensor `src` that lives on + /// `src_stream`. Called when [`FusionTensor::clone`](crate::FusionTensor::clone) + /// or [`FusionTensor::into_ir`](crate::FusionTensor::into_ir) detects that the + /// tensor's home stream is not the current stream. + /// + /// See the [`MultiStream`] struct-level docs for the full strategy. In short: + /// + /// - If `src`'s handle isn't materialised yet, drain `src_stream` so the + /// producing op runs and registers it. Skip the drain on subsequent shares + /// of the same source by remembering it in `shared_sources` *or* by + /// observing that the handle is already in the container (which covers the + /// chained-share case where `src` is itself a previously-aliased view). + /// - Then alias the backing handle under `dst`. `register_handle` clones the + /// cubecl handle (`Arc`-style), so both ids share refcount on the buffer + /// until each side's own `Drop` op runs. + pub fn tag_shared_view( + &mut self, + src_stream: StreamId, + src: TensorId, + dst: TensorId, + handles: &mut HandleContainer, + ) { + // Drain only when neither short-circuit applies: `shared_sources` records ids + // we already drained for, and a `Some` handle means `src` is materialised + // (e.g., it was itself set up by an earlier `tag_shared_view` call). We + // record `src` only when we actually drain — the handle-existence path is + // naturally idempotent on later calls. + if !self.shared_sources.contains(&src) && handles.get_handle_ref(&src).is_none() { + self.shared_sources.insert(src); + self.drain(handles, src_stream); + } + + if let Some(handle) = handles.get_handle_ref(&src) { + // Not a bare `clone()`: remote backends need a fresh server-side handle over the same + // buffer so consuming one alias doesn't free it for the other stream. Local backends' + // `alias_handle` default *is* `clone()`. See `FusionRuntime::alias_handle`. + let alias = R::alias_handle(handle); + handles.register_handle(dst, alias); + } + } + + /// Enqueue an operation on the queue for `stream` and run the lazy processor. + fn enqueue_operation( + &mut self, + stream: StreamId, + repr: OperationIr, + operation: UnfusedOp, + handles: &mut HandleContainer, + ) { + let s = self + .streams + .entry(stream) + .or_insert_with(|| Stream::new(self.device.clone())); + s.queue.add(repr, operation); + + let len_before = s.queue.global.len(); + s.processor.process( + Segment::new(&mut s.queue, handles, stream), + &mut self.optimizations, + ExecutionMode::Lazy, + ); + let len_after = s.queue.global.len(); + s.cursor += (len_before - len_after) as u64; + } + + /// Mark a tensor as read. + #[allow(unused_variables)] + pub fn mark_read( + &mut self, + id: StreamId, + ir: &burn_ir::TensorIr, + handles: &HandleContainer, + ) { + if !matches!(ir.status, burn_ir::TensorStatus::ReadWrite) { + return; + }; + + let stream = match self.streams.get_mut(&id) { + Some(val) => val, + None => return, + }; + + stream.queue.variables.remove(&ir.id); + + if stream.queue.variables.is_empty() { + self.streams.remove(&id); + } + + #[cfg(feature = "memory-checks")] + self.memory_checks.check(&self.streams, handles); + #[cfg(feature = "test-util")] + crate::inspect::emit_handle_snapshot(id, handles.handle_ids().copied()); + } + + /// Drain a stream. + pub fn drain(&mut self, handles: &mut HandleContainer, id: StreamId) { + id.executes(|| { + if let Some(stream) = self.streams.get_mut(&id) { + let num_executed = stream.queue.global.len(); + stream.processor.process( + Segment::new(&mut stream.queue, handles, id), + &mut self.optimizations, + ExecutionMode::Sync, + ); + stream.cursor += num_executed as u64; + } + }); + #[cfg(feature = "test-util")] + crate::inspect::emit_handle_snapshot(id, handles.handle_ids().copied()); + } +} + +pub(crate) struct Stream { + pub(crate) queue: OperationQueue, + processor: Processor, + pub(crate) cursor: u64, +} + +#[derive(new)] +struct Segment<'a, R: FusionRuntime> { + queue: &'a mut OperationQueue, + handles: &'a mut HandleContainer, + id: StreamId, +} + +impl StreamSegment for Segment<'_, R> { + fn operations(&self) -> &[OperationIr] { + &self.queue.relative + } + + fn execute(&mut self, id: ExecutionPlanId, store: &mut ExecutionPlanStore) { + self.queue.execute(id, self.handles, store, self.id) + } + + fn execute_unfused(&mut self, optimization: BlockOptimization) { + self.queue + .execute_unfused(optimization, self.handles, self.id) + } +} + +impl Stream { + fn new(device: R::FusionDevice) -> Self { + Self { + processor: Processor::new(R::fusers(device)), + queue: OperationQueue::new(), + cursor: 0, + } + } +} diff --git a/crates/burn-fusion/src/stream/queue/base.rs b/crates/burn-fusion/src/stream/queue/base.rs new file mode 100644 index 0000000..b6c573c --- /dev/null +++ b/crates/burn-fusion/src/stream/queue/base.rs @@ -0,0 +1,87 @@ +use crate::stream::{OperationConverter, RelativeOps}; +use crate::{FusionRuntime, UnfusedOp}; +use burn_ir::{OperationIr, TensorId, TensorStatus}; + +use hashbrown::HashMap; + +/// A growing list of [tensor operation descriptions](OperationIr). +/// +/// Every queue is associated with a single [`StreamId`](super::super::StreamId) — every tensor it +/// references is local to that stream (cross-stream sharing is handled out-of-band by +/// [`MultiStream::tag_shared_view`](super::super::MultiStream::tag_shared_view), which aliases +/// the source handle under a fresh local tensor id before the op is enqueued). +pub struct OperationQueue { + /// List of operation descriptions. These contain the exact tensor IDs + /// and shapes so that kernels can be run correctly. + /// + /// The length of this list is the same as the length of the `operations` list. + pub(crate) global: Vec, + /// List of operation descriptions. The tensor IDs and shapes are relative + /// because we don't need to know the exact values, but they are sufficient to + /// determine which operations can be fused. + pub(crate) relative: Vec, + pub(crate) converter: OperationConverter, + pub(crate) operations: Vec>, + pub(crate) variables: HashMap, +} + +impl Default for OperationQueue { + fn default() -> Self { + Self::new() + } +} + +impl OperationQueue { + /// Create a new empty queue. + pub fn new() -> Self { + Self { + global: Vec::new(), + relative: Vec::new(), + converter: OperationConverter::default(), + operations: Vec::new(), + variables: HashMap::new(), + } + } + + /// Add a new tensor operation to the queue. + /// + /// The new [operation intermediate representation](OperationIr) will be converted to a local + /// representation that can be reused when the same pattern emerge in different but similar + /// scenario, so that the same optimization can be used. + pub fn add(&mut self, global: OperationIr, operation: UnfusedOp) { + for node in global.nodes() { + self.variables.insert(node.id, node.status); + } + let relative = global.to_relative(&mut self.converter); + self.relative.push(relative); + self.global.push(global); + self.operations.push(operation); + } +} + +#[cfg(all(test, feature = "std"))] +mod tests { + use burn_backend::StreamId; + + #[test] + fn stream_id_from_different_threads() { + let current = StreamId::current(); + + let thread1 = std::thread::spawn(|| (StreamId::current(), StreamId::current())); + let thread2 = std::thread::spawn(StreamId::current); + + let (stream_1, stream_11) = thread1.join().unwrap(); + let stream_2 = thread2.join().unwrap(); + + assert_ne!(current, stream_1, "Should be different from thread 1"); + assert_ne!(current, stream_2, "Should be different from thread 2"); + assert_ne!( + stream_1, stream_2, + "Should be different from different threads" + ); + assert_eq!( + stream_1, stream_11, + "Should be the same, since same thread." + ); + } +} diff --git a/crates/burn-fusion/src/stream/queue/execution.rs b/crates/burn-fusion/src/stream/queue/execution.rs new file mode 100644 index 0000000..b7d46c9 --- /dev/null +++ b/crates/burn-fusion/src/stream/queue/execution.rs @@ -0,0 +1,125 @@ +use burn_ir::{HandleContainer, TensorStatus}; + +use crate::{ + FusionRuntime, UnfusedOp, + search::BlockOptimization, + stream::{ + Context, ContextGuard, OperationConverter, OrderedExecution, RelativeOps, StreamId, + execution::log_execution_table, + store::{ExecutionPlanId, ExecutionPlanStore, ExecutionStrategy}, + }, +}; + +use super::OperationQueue; + +impl OperationQueue { + /// Execute the queue partially following the execution strategy from the plan. + pub(crate) fn execute( + &mut self, + id: ExecutionPlanId, + handles: &mut HandleContainer, + store: &mut ExecutionPlanStore, + stream_id: StreamId, + ) { + let plan = store.get_mut_unchecked(id); + self.execute_block_optimization(&mut plan.optimization, handles, stream_id); + } + + /// Execute the queue with a one-off [`BlockOptimization`] that isn't stored in the cache. + /// + /// Used when fusion exploration is capped (see [`Explorer`](crate::stream::execution)): a + /// cache-missing segment runs unfused without paying the optimizer cost or growing the store. + pub(crate) fn execute_unfused( + &mut self, + mut optimization: BlockOptimization, + handles: &mut HandleContainer, + stream_id: StreamId, + ) { + self.execute_block_optimization(&mut optimization, handles, stream_id); + } + + fn execute_block_optimization( + &mut self, + step: &mut BlockOptimization, + handles: &mut HandleContainer, + stream_id: StreamId, + ) { + log_execution_table(stream_id, &step.strategy, &self.global); + + let mut operations = Vec::new(); + core::mem::swap(&mut operations, &mut self.operations); + + let (operations, num_drained) = + run_strategy(step, &mut self.converter, handles, operations); + + self.operations = operations; + self.drain_queue(num_drained, handles); + } + + /// Bookkeeping after executing `num_drained` operations from the queue. + fn drain_queue(&mut self, num_drained: usize, handles: &mut HandleContainer) { + self.global[0..num_drained] + .iter() + .flat_map(|desc| desc.nodes()) + .for_each(|tensor| { + if tensor.status == TensorStatus::ReadWrite { + self.variables.remove(&tensor.id); + }; + R::free_handle(handles, tensor) + }); + + self.global.drain(0..num_drained); + + self.reset_relative(); + } + + fn reset_relative(&mut self) { + self.relative.clear(); + self.converter.clear(); + + for node in self.global.iter() { + let relative = node.to_relative(&mut self.converter); + self.relative.push(relative); + } + } +} + +/// Drive one block's execution strategy. +/// +/// Wraps the converter's per-block fields and the handle container into a single owned +/// [`Context`] via [`ContextGuard`] for the duration of this call, then threads `&mut Context` +/// through the recursive strategy walk. Operations-only strategies just grab +/// `&mut ctx.handles`; optimization strategies hand `&mut ctx` to the fused op. +fn run_strategy( + optimization: &mut BlockOptimization, + converter: &mut OperationConverter, + handles: &mut HandleContainer, + operations: Vec>, +) -> (Vec>, usize) { + let mut execution = OrderedExecution::new(operations); + { + let mut guard = ContextGuard::new(converter, handles); + execute_strategy::(&mut optimization.strategy, &mut guard, &mut execution); + } + execution.finish() +} + +fn execute_strategy( + strategy: &mut ExecutionStrategy, + context: &mut Context, + execution: &mut OrderedExecution, +) { + match strategy { + ExecutionStrategy::Optimization { ordering, opt, .. } => { + execution.execute_optimization(opt, context, ordering.clone()); + } + ExecutionStrategy::Operations { ordering } => { + execution.execute_operations(&mut context.handles, ordering); + } + ExecutionStrategy::Composed(items) => { + for item in items.iter_mut() { + execute_strategy::(item, context, execution); + } + } + } +} diff --git a/crates/burn-fusion/src/stream/queue/mod.rs b/crates/burn-fusion/src/stream/queue/mod.rs new file mode 100644 index 0000000..450cba0 --- /dev/null +++ b/crates/burn-fusion/src/stream/queue/mod.rs @@ -0,0 +1,4 @@ +mod base; +mod execution; + +pub use base::*; diff --git a/crates/burn-fusion/src/stream/store/base.rs b/crates/burn-fusion/src/stream/store/base.rs new file mode 100644 index 0000000..7de0308 --- /dev/null +++ b/crates/burn-fusion/src/stream/store/base.rs @@ -0,0 +1,98 @@ +use std::sync::Arc; + +use crate::search::BlockOptimization; + +use super::{ExecutionPlanIndex, InsertQuery, SearchQuery}; +use burn_ir::OperationIr; +use serde::{Deserialize, Serialize}; + +/// The store that contains all explorations done on a device. +#[derive(Default)] +pub(crate) struct ExecutionPlanStore { + plans: Vec>, + index: ExecutionPlanIndex, +} + +/// How a list of operations should be executed. +#[derive(PartialEq, Debug, Clone)] +pub(crate) enum ExecutionStrategy { + /// An optimization was found, and therefore should be executed. + Optimization { + opt: O, + ordering: Arc>, + score: u64, + }, + /// No optimization was found, each operation should be executed individually. + Operations { ordering: Arc> }, + /// A composition of multiple execution strategies. + Composed(Vec>), +} + +/// The trigger that indicates when to stop exploring. +#[derive(Debug, PartialEq, Serialize, Deserialize)] +pub(crate) enum ExecutionTrigger { + OnOperations(Vec), + OnSync, + Always, +} + +/// The unique identifier for an exploration that was executed. +pub(crate) type ExecutionPlanId = usize; + +/// The outcome of an exploration that can be stored. +#[derive(Debug)] +pub(crate) struct ExecutionPlan { + /// The operations on which the exploration is related to. + pub(crate) operations: Vec, + /// The criteria that signal when this plan should be executed. Only one trigger is necessary. + pub(crate) triggers: Vec, + /// The optimization that should be used when executing this plan. + pub(crate) optimization: BlockOptimization, +} + +impl ExecutionPlanStore { + pub fn new() -> Self { + Self { + plans: Vec::new(), + index: ExecutionPlanIndex::default(), + } + } + + pub fn find(&self, query: SearchQuery<'_>) -> Vec { + self.index.find(query) + } + + pub fn add(&mut self, exploration: ExecutionPlan) -> ExecutionPlanId { + if exploration.operations.is_empty() { + panic!("Can't add an empty optimization."); + } + + let id = self.plans.len(); + + self.index.insert(InsertQuery::NewPlan { + operations: &exploration.operations, + id, + }); + + self.plans.push(exploration); + + id + } + + pub fn get_mut_unchecked(&mut self, id: ExecutionPlanId) -> &mut ExecutionPlan { + &mut self.plans[id] + } + + pub fn get_unchecked(&self, id: ExecutionPlanId) -> &ExecutionPlan { + &self.plans[id] + } + + /// Add a new end condition for an optimization. + pub fn add_trigger(&mut self, id: ExecutionPlanId, trigger: ExecutionTrigger) { + let criteria = &mut self.plans[id].triggers; + + if !criteria.contains(&trigger) { + criteria.push(trigger); + } + } +} diff --git a/crates/burn-fusion/src/stream/store/index.rs b/crates/burn-fusion/src/stream/store/index.rs new file mode 100644 index 0000000..014639c --- /dev/null +++ b/crates/burn-fusion/src/stream/store/index.rs @@ -0,0 +1,296 @@ +use crate::stream::store::ExecutionPlanId; +use burn_ir::OperationIr; +use serde::{Deserialize, Serialize}; +use std::{ + collections::{HashMap, hash_map::DefaultHasher}, + hash::{Hash, Hasher}, +}; + +/// Index used to search optimizations. +#[derive(Default, Serialize, Deserialize, Clone)] +pub struct ExecutionPlanIndex { + /// We can't use `HashMap>` since `OperationIr` + /// doesn't implement [`Eq`](core::cmp::Eq). + /// + /// `OperationIr` can't implement `Eq` since float types don't implement it. + /// + /// We rely instead on [`PartialEq`](core::cmp::PartialEq) to manually handle hash collisions. + /// This is OK because we use `relative` operations where any scalar values are set to zeros, + /// see [`RelativeStreamConverter`](crate::stream::RelativeStreamConverter). + /// + /// Map from the hash of the `OperationIr` to a list of `(OperationIr, index)` pairs, + /// where `index` is the index of all the execution plans that start with the `OperationIr` + /// in the `starters` list. + mapping: HashMap>, + starters: Vec>, +} + +pub enum SearchQuery<'a> { + PlansStartingWith(&'a OperationIr), +} + +pub enum InsertQuery<'a> { + NewPlan { + operations: &'a [OperationIr], + id: ExecutionPlanId, + }, +} + +impl ExecutionPlanIndex { + /// Search optimizations with the given [query](SearchQuery). + pub fn find(&self, query: SearchQuery<'_>) -> Vec { + match query { + SearchQuery::PlansStartingWith(ops) => self.find_starting_with(ops), + } + } + + /// Register a new optimization with the given [query](InsertQuery). + pub fn insert(&mut self, query: InsertQuery<'_>) { + match query { + InsertQuery::NewPlan { operations, id } => { + if let Some(operation) = operations.first() { + self.insert_new_operation(operation, id) + } + } + } + } + + /// Find execution plans starting with the `OperationIr` + fn find_starting_with(&self, operation: &OperationIr) -> Vec { + let key = self.operation_key(operation); + let values = match self.mapping.get(&key) { + Some(val) => val, + None => return Vec::new(), + }; + + if values.is_empty() { + return Vec::new(); + } + + let (_, index) = match values.iter().find(|value| &value.0 == operation) { + Some(val) => val, + None => return Vec::new(), + }; + + match self.starters.get(*index) { + Some(value) => value.clone(), + None => Vec::new(), + } + } + + /// Update the index for an execution plan starting with operation `ops` + fn insert_new_operation(&mut self, ops: &OperationIr, new_id: ExecutionPlanId) { + let key = self.operation_key(ops); + let values = match self.mapping.get_mut(&key) { + Some(val) => val, + None => { + // New starter ops. + let index = self.starters.len(); + self.starters.push(vec![new_id]); + self.mapping.insert(key, vec![(ops.clone(), index)]); + + return; + } + }; + let (_, index) = match values.iter_mut().find(|value| &value.0 == ops) { + Some(val) => val, + None => { + // New with hash collision. + let index = self.starters.len(); + self.starters.push(vec![new_id]); + values.push((ops.clone(), index)); + return; + } + }; + + // New optimization for an existing starter. + self.starters + .get_mut(*index) + .expect("Should exist") + .push(new_id); + } + + // Hash the value of the first operation in a list. + fn operation_key(&self, ops: &OperationIr) -> u64 { + let mut hasher = DefaultHasher::new(); + ops.hash(&mut hasher); + hasher.finish() + } +} + +#[cfg(test)] +mod tests { + use burn_backend::{DType, Shape}; + use burn_ir::{ + BinaryOpIr, NumericOperationIr, ScalarIr, ScalarOpIr, TensorId, TensorIr, TensorStatus, + }; + + use super::*; + + #[test] + fn should_find_optimization_id_based_on_tensor_ops() { + let mut index = ExecutionPlanIndex::default(); + let stream_1 = [ops_1()]; + let optimization_id_1 = 0; + + index.insert(InsertQuery::NewPlan { + operations: &stream_1, + id: optimization_id_1, + }); + + let found = index.find(SearchQuery::PlansStartingWith(&stream_1[0])); + + assert_eq!(found, vec![optimization_id_1]); + } + + #[test] + fn should_support_multiple_optimization_ids_with_same_starting_ops() { + let mut index = ExecutionPlanIndex::default(); + let stream_1 = [ops_1(), ops_2(), ops_1()]; + let stream_2 = [ops_1(), ops_1(), ops_2()]; + let optimization_id_1 = 0; + let optimization_id_2 = 1; + + index.insert(InsertQuery::NewPlan { + operations: &stream_1, + id: optimization_id_1, + }); + index.insert(InsertQuery::NewPlan { + operations: &stream_2, + id: optimization_id_2, + }); + + let found = index.find(SearchQuery::PlansStartingWith(&stream_1[0])); + + assert_eq!(found, vec![optimization_id_1, optimization_id_2]); + } + + #[test] + fn should_only_find_optimization_with_correct_starting_ops() { + let mut index = ExecutionPlanIndex::default(); + let stream_1 = [ops_1(), ops_1()]; + let stream_2 = [ops_2(), ops_1()]; + let optimization_id_1 = 0; + let optimization_id_2 = 1; + + index.insert(InsertQuery::NewPlan { + operations: &stream_1, + id: optimization_id_1, + }); + index.insert(InsertQuery::NewPlan { + operations: &stream_2, + id: optimization_id_2, + }); + + let found = index.find(SearchQuery::PlansStartingWith(&stream_1[0])); + + assert_eq!(found, vec![optimization_id_1]); + } + + #[test] + fn should_handle_hash_collisions() { + let mut index = ExecutionPlanIndex::default(); + let stream_1 = [ops_1(), ops_1()]; + let stream_2 = [ops_3(), ops_1()]; + let optimization_id_1 = 0; + let optimization_id_2 = 1; + + let stream_1_key = index.operation_key(&stream_1[0]); + let stream_2_key = index.operation_key(&stream_2[0]); + + assert_ne!( + stream_1_key, stream_2_key, + "Ops 1 and Ops 3 should not have the same hash" + ); // ops 1 and 3 have different variants, so the hash differs + assert_ne!( + stream_1[0], stream_2[0], + "Ops 1 and Ops 3 should be different." + ); + + index.insert(InsertQuery::NewPlan { + operations: &stream_1, + id: optimization_id_1, + }); + index.insert(InsertQuery::NewPlan { + operations: &stream_2, + id: optimization_id_2, + }); + + let found = index.find(SearchQuery::PlansStartingWith(&stream_1[0])); + + assert_eq!(found, vec![optimization_id_1]); + } + + fn ops_1() -> OperationIr { + OperationIr::NumericFloat( + DType::F32, + NumericOperationIr::Add(BinaryOpIr { + lhs: TensorIr { + id: TensorId::new(0), + shape: Shape::new([32, 32]), + status: TensorStatus::ReadOnly, + dtype: DType::F32, + }, + rhs: TensorIr { + id: TensorId::new(1), + shape: Shape::new([32, 32]), + status: TensorStatus::ReadOnly, + dtype: DType::F32, + }, + out: TensorIr { + id: TensorId::new(2), + shape: Shape::new([32, 32]), + status: TensorStatus::NotInit, + dtype: DType::F32, + }, + }), + ) + } + + fn ops_2() -> OperationIr { + OperationIr::NumericFloat( + DType::F32, + NumericOperationIr::AddScalar(ScalarOpIr { + lhs: TensorIr { + id: TensorId::new(0), + shape: Shape::new([32, 32]), + status: TensorStatus::ReadOnly, + dtype: DType::F32, + }, + rhs: ScalarIr::Float(5.0), + out: TensorIr { + id: TensorId::new(2), + shape: Shape::new([32, 32]), + status: TensorStatus::NotInit, + dtype: DType::F32, + }, + }), + ) + } + + fn ops_3() -> OperationIr { + OperationIr::NumericFloat( + DType::F32, + NumericOperationIr::Sub(BinaryOpIr { + lhs: TensorIr { + id: TensorId::new(0), + shape: Shape::new([32, 32]), + status: TensorStatus::ReadOnly, + dtype: DType::F32, + }, + rhs: TensorIr { + id: TensorId::new(1), + shape: Shape::new([32, 32]), + status: TensorStatus::ReadOnly, + dtype: DType::F32, + }, + out: TensorIr { + id: TensorId::new(2), + shape: Shape::new([32, 32]), + status: TensorStatus::NotInit, + dtype: DType::F32, + }, + }), + ) + } +} diff --git a/crates/burn-fusion/src/stream/store/mod.rs b/crates/burn-fusion/src/stream/store/mod.rs new file mode 100644 index 0000000..f14aff0 --- /dev/null +++ b/crates/burn-fusion/src/stream/store/mod.rs @@ -0,0 +1,5 @@ +mod base; +mod index; + +pub(crate) use base::*; +pub(super) use index::*; diff --git a/crates/burn-fusion/src/tensor.rs b/crates/burn-fusion/src/tensor.rs new file mode 100644 index 0000000..cc41e0d --- /dev/null +++ b/crates/burn-fusion/src/tensor.rs @@ -0,0 +1,266 @@ +use crate::{ + Client, FusionBackend, FusionRuntime, + stream::{Operation, StreamId}, +}; +use burn_backend::{DType, ExecutionError, Shape, TensorData, TensorMetadata}; +use burn_ir::{OperationIr, TensorId, TensorIr, TensorStatus}; +use std::sync::{ + Arc, + atomic::{AtomicU32, Ordering}, +}; + +/// Tensor primitive for the [fusion backend](crate::FusionBackend) for all kind. +pub struct FusionTensor { + /// Tensor id. + pub id: TensorId, + /// The shape of the tensor. + pub shape: Shape, + /// The fusion client. + pub client: Client, + /// The datatype of the tensor. + pub dtype: DType, + /// The current stream id this tensor is on. + pub stream: StreamId, + pub(crate) count: Arc, +} + +impl Clone for FusionTensor { + fn clone(&self) -> Self { + let current = StreamId::current(); + if self.stream != current { + return self.shared_view(current); + } + + self.count.fetch_add(1, Ordering::Acquire); + + Self { + id: self.id, + shape: self.shape.clone(), + client: self.client.clone(), + dtype: self.dtype, + stream: self.stream, + count: self.count.clone(), + } + } +} + +impl core::fmt::Debug for FusionTensor { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str( + format!( + "{{ id: {:?}, shape: {:?}, device: {:?} }}", + self.id, + self.shape, + self.client.device(), + ) + .as_str(), + ) + } +} + +impl TensorMetadata for FusionTensor { + type Device = R::FusionDevice; + fn dtype(&self) -> DType { + self.dtype + } + + fn shape(&self) -> Shape { + self.shape.clone() + } + + fn rank(&self) -> usize { + self.shape.num_dims() + } + + fn device(&self) -> Self::Device { + self.client.device().clone() + } + + fn can_mut(&self) -> bool { + // Same rule as `status` at drain time: a handle shared on its stream + // (count > 1) is read-only, a unique one is read-write and the fused + // kernel may write its buffer in place. + matches!( + self.status(self.count.load(Ordering::Acquire)), + TensorStatus::ReadWrite + ) + } +} + +impl FusionTensor { + pub(crate) fn new( + id: TensorId, + shape: Shape, + dtype: DType, + client: Client, + stream: StreamId, + ) -> Self { + Self { + id, + shape, + client, + dtype, + stream, + count: Arc::new(AtomicU32::new(1)), + } + } + + fn status(&self, count: u32) -> TensorStatus { + if count <= 1 { + TensorStatus::ReadWrite + } else { + TensorStatus::ReadOnly + } + } + + /// Intermediate representation to be used when using an uninitialized tensor as output. + pub fn to_ir_out(&self) -> TensorIr { + TensorIr { + status: TensorStatus::NotInit, + shape: self.shape.clone(), + id: self.id, + dtype: self.dtype, + } + } + + /// Intermediate representation to be used when using an initialized tensor used as input. + pub fn into_ir(mut self) -> TensorIr { + let current = StreamId::current(); + if self.stream != current { + self = self.shared_view(current); + } + + let count = self.count.load(Ordering::Acquire); + let status = self.status(count); + + let mut shape_out = Shape::from(Vec::::new()); + core::mem::swap(&mut self.shape, &mut shape_out); + + if let TensorStatus::ReadWrite = status { + // Avoids an unwanted drop on the same thread. + // + // Since `drop` is called after `into_ir`, we must not register a drop if the tensor + // was consumed with a `ReadWrite` status. + self.count.fetch_add(1, Ordering::Acquire); + } + + TensorIr { + status, + shape: shape_out, + id: self.id, + dtype: self.dtype, + } + } + + /// Create a fresh `FusionTensor` on `current` that aliases the same backing + /// handle as `self`. Used by [`Clone`] and [`Self::into_ir`] when the tensor is + /// crossing stream boundaries — the rest of the pipeline only ever sees ids + /// whose home stream is the calling stream. + /// + /// The cross-stream coordination (draining the source stream so the handle + /// exists, then aliasing it under a fresh id) is done by + /// [`MultiStream::tag_shared_view`](crate::stream::MultiStream::tag_shared_view). + /// See that type's docs for the full strategy — how shares are tagged, how the + /// buffer's lifetime is managed across the two sides, and why a single drain + /// per source is enough. + fn shared_view(&self, current: StreamId) -> Self { + let new_id = self.client.create_empty_handle(); + + self.client.tag_shared_view(self.stream, self.id, new_id); + + Self::new( + new_id, + self.shape.clone(), + self.dtype, + self.client.clone(), + current, + ) + } + + pub(crate) async fn into_data(self) -> Result + where + B: FusionBackend, + { + let id = self.stream; + let client = self.client.clone(); + let desc = self.into_ir(); + client.read_tensor_float::(desc, id).await + } + + pub(crate) async fn q_into_data(self) -> Result + where + B: FusionBackend, + { + if let DType::QFloat(_scheme) = self.dtype { + let id = self.stream; + let client = self.client.clone(); + let desc = self.into_ir(); + client.read_tensor_quantized::(desc, id).await + } else { + panic!("Expected quantized float dtype, got {:?}", self.dtype) + } + } + + pub(crate) async fn int_into_data(self) -> Result + where + B: FusionBackend, + { + let id = self.stream; + let client = self.client.clone(); + let desc = self.into_ir(); + client.read_tensor_int::(desc, id).await + } + + pub(crate) async fn bool_into_data(self) -> Result + where + B: FusionBackend, + { + let id = self.stream; + let client = self.client.clone(); + let desc = self.into_ir(); + client.read_tensor_bool::(desc, id).await + } +} + +#[derive(new, Debug)] +pub(crate) struct DropOp { + pub(crate) id: TensorId, +} + +impl Operation for DropOp { + fn execute(&self, handles: &mut burn_ir::HandleContainer) { + handles.remove_handle(self.id); + } +} + +impl Drop for FusionTensor { + fn drop(&mut self) { + let count = self.count.fetch_sub(1, Ordering::Acquire); + + // Workaround to prevent segfaults when an operation panics + if std::thread::panicking() { + return; + } + + match self.status(count) { + TensorStatus::ReadWrite => { + let mut shape = Shape::from(Vec::::new()); + core::mem::swap(&mut shape, &mut self.shape); + + let ir = TensorIr { + id: self.id, + shape, + status: TensorStatus::ReadWrite, + dtype: self.dtype, + }; + + // Drop is targeted at the tensor's home stream so it runs after any pending ops + // on this id, regardless of the thread we happen to be dropping from. + self.client + .register(self.stream, OperationIr::Drop(ir), DropOp { id: self.id }); + } + TensorStatus::ReadOnly => {} + TensorStatus::NotInit => {} + } + } +} diff --git a/crates/burn-ir/Cargo.toml b/crates/burn-ir/Cargo.toml new file mode 100644 index 0000000..0248a63 --- /dev/null +++ b/crates/burn-ir/Cargo.toml @@ -0,0 +1,33 @@ +[package] +authors = ["laggui ", "nathanielsimard "] +categories = ["science"] +description = "Intermediate representation for the Burn framework" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "tensor"] +license.workspace = true +name = "burn-ir" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-ir" +documentation = "https://docs.rs/burn-ir" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["std"] +std = ["burn-backend/std"] +doc = ["default"] +tracing = [ + "burn-backend/tracing", +] + +[dependencies] +serde = { workspace = true } +hashbrown = { workspace = true } # no_std compatible + +burn-backend = { workspace = true } + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-ir/README.md b/crates/burn-ir/README.md new file mode 100644 index 0000000..3e26709 --- /dev/null +++ b/crates/burn-ir/README.md @@ -0,0 +1,7 @@ +# Burn Intermediate Representation + +Defines an Intermediate Representation (IR) used to represent tensors and operations. + +The abstraction over computation allows execution across different targets (e.g., remote backend). +It also enables optimization and transformation of tensor computations before execution (e.g., +operator fusion). diff --git a/crates/burn-ir/src/backend.rs b/crates/burn-ir/src/backend.rs new file mode 100644 index 0000000..b655600 --- /dev/null +++ b/crates/burn-ir/src/backend.rs @@ -0,0 +1,63 @@ +use burn_backend::{ + Backend, Shape, + tensor::{BoolTensor, FloatTensor, IntTensor, QuantizedTensor}, +}; + +/// A tensor representation containing a reference to a tensor resource with a given shape. +#[derive(Clone)] +pub struct TensorHandle { + /// The type that can be used to point to a tensor of any kind. + pub handle: H, + /// The shape associated to the tensor. + pub shape: Shape, +} + +/// Backend extension trait that allows an existing [backend](Backend) to use the Burn tensor +/// intermediate representation for compilation purpose or other... +pub trait BackendIr: Backend { + /// The type that can be used to point to a tensor of any kind. + type Handle: Sync + Send + Clone; + + /// Convert a [handle](BackendIr::Handle) to a [float tensor](burn_backend::BackendTypes::FloatTensorPrimitive). + fn float_tensor(handle: TensorHandle) -> FloatTensor; + /// Convert a [handle](BackendIr::Handle) to an [int tensor](burn_backend::BackendTypes::IntTensorPrimitive). + fn int_tensor(handle: TensorHandle) -> IntTensor; + /// Convert a [handle](BackendIr::Handle) to a [bool tensor](burn_backend::BackendTypes::BoolTensorPrimitive). + fn bool_tensor(handle: TensorHandle) -> BoolTensor; + /// Convert a [handle](BackendIr::Handle) to a [quantized tensor](burn_backend::BackendTypes::QuantizedTensorPrimitive). + fn quantized_tensor(handle: TensorHandle) -> QuantizedTensor; + + /// Convert a [float tensor](burn_backend::BackendTypes::FloatTensorPrimitive) to a [handle](BackendIr::Handle). + fn float_tensor_handle(tensor: FloatTensor) -> Self::Handle; + /// Convert an [int tensor](burn_backend::BackendTypes::IntTensorPrimitive) to a [handle](BackendIr::Handle). + fn int_tensor_handle(tensor: IntTensor) -> Self::Handle; + /// Convert a [bool tensor](burn_backend::BackendTypes::BoolTensorPrimitive) to a [handle](BackendIr::Handle). + fn bool_tensor_handle(tensor: BoolTensor) -> Self::Handle; + /// Convert a [quantized tensor](burn_backend::BackendTypes::QuantizedTensorPrimitive) to a [handle](BackendIr::Handle). + fn quantized_tensor_handle(tensor: QuantizedTensor) -> Self::Handle; +} + +/// Handle which points to a backend tensor primitive kind. +#[derive(Clone, Debug)] +pub enum HandleKind { + /// Float tensor handle. + Float(B::FloatTensorPrimitive), + /// Int tensor handle. + Int(B::IntTensorPrimitive), + /// Bool tensor handle. + Bool(B::BoolTensorPrimitive), + /// Quantized tensor handle. + Quantized(B::QuantizedTensorPrimitive), +} + +impl HandleKind { + /// Returns the handle kind name. + pub fn name(&self) -> &str { + match self { + HandleKind::Float(_) => "float", + HandleKind::Int(_) => "int", + HandleKind::Bool(_) => "bool", + HandleKind::Quantized(_) => "quantized", + } + } +} diff --git a/crates/burn-ir/src/builder.rs b/crates/burn-ir/src/builder.rs new file mode 100644 index 0000000..6ef0f19 --- /dev/null +++ b/crates/burn-ir/src/builder.rs @@ -0,0 +1,1479 @@ +#![allow(missing_docs)] + +use alloc::vec::Vec; +use burn_backend::{ + DType, Distribution, Shape, Slice, SliceOps, calculate_matmul_output, + ops::{ + conv::{ + calculate_conv_output_shape, calculate_conv_transpose_output_shape, + calculate_pool_output_shape, + }, + unfold::calculate_unfold_shape, + }, + quantization::QuantScheme, + tensor::IndexingUpdateOp, +}; + +use crate::{ScalarIr, TensorId, TensorIr}; + +use super::operation::*; + +impl CreationOpIr { + pub fn create(shape: Shape, dtype: DType, new_id: impl FnOnce() -> TensorId) -> Self { + let out = TensorIr::uninit(new_id(), shape, dtype); + + CreationOpIr { out } + } +} + +impl InitOperationIr { + pub fn create(shape: Shape, dtype: DType, new_id: impl FnOnce() -> TensorId) -> Self { + let out = TensorIr::uninit(new_id(), shape, dtype); + + InitOperationIr { out } + } +} + +impl RandomOpIr { + pub fn create( + shape: Shape, + dtype: DType, + distribution: Distribution, + new_id: impl FnOnce() -> TensorId, + ) -> Self { + let out = TensorIr::uninit(new_id(), shape, dtype); + + RandomOpIr { out, distribution } + } +} + +impl FullOpIr { + pub fn create( + shape: Shape, + dtype: DType, + value: ScalarIr, + new_id: impl FnOnce() -> TensorId, + ) -> Self { + // TODO: check that ScalarIr dtype matches dtype? + let out = TensorIr::uninit(new_id(), shape, dtype); + + FullOpIr { out, value } + } +} + +impl CastOpIr { + pub fn create(input: TensorIr, dtype: DType, new_id: impl FnOnce() -> TensorId) -> Self { + let out = TensorIr::uninit(new_id(), input.shape.clone(), dtype); + CastOpIr { input, out } + } +} + +impl ShapeOpIr { + pub fn expand(input: TensorIr, shape: Shape, new_id: impl FnOnce() -> TensorId) -> Self { + let shape = input.shape.expand(shape).unwrap(); + Self::create(input, shape, new_id) + } + + pub fn reshape(input: TensorIr, shape: Shape, new_id: impl FnOnce() -> TensorId) -> Self { + let shape = input.shape.reshape(shape).unwrap(); + Self::create(input, shape, new_id) + } + + fn create(input: TensorIr, shape: Shape, new_id: impl FnOnce() -> TensorId) -> Self { + let out = TensorIr::uninit(new_id(), shape, input.dtype); + ShapeOpIr { input, out } + } +} + +// "Lower" specific operations into a binary or unary op representation. +// Useful when collecting inputs and outputs and don't care about the other semantics. +impl From for BinaryOpIr { + fn from(value: MatmulOpIr) -> Self { + Self { + lhs: value.lhs, + rhs: value.rhs, + out: value.out, + } + } +} + +impl From for UnaryOpIr { + fn from(value: ReduceOpIr) -> Self { + Self { + input: value.input, + out: value.out, + } + } +} + +#[derive(Debug)] +#[allow(missing_docs)] +pub enum IrError { + DTypeMismatch, +} + +fn dtype_compat(lhs: &DType, rhs: &DType) -> bool { + let lhs_qfloat = matches!(lhs, DType::QFloat(_)); + let rhs_qfloat = matches!(rhs, DType::QFloat(_)); + if lhs_qfloat && (rhs_qfloat || rhs.is_float()) + || lhs.is_float() && (rhs_qfloat || rhs.is_float()) + { + true + } else { + lhs == rhs + } +} + +fn output_check<'a, I>(inputs: I, compat: impl Fn(&DType, &DType) -> bool) -> Result +where + I: IntoIterator, +{ + let mut iter = inputs.into_iter(); + let first = iter.next().unwrap(); + for d in iter { + if !compat(first, d) { + return Err(IrError::DTypeMismatch); + } + } + Ok(*first) +} + +fn output_dtype<'a, I: IntoIterator>(inputs: I) -> Result { + output_check(inputs, |a, b| a == b) +} + +fn output_dtype_mixed<'a, I: IntoIterator>(inputs: I) -> Result { + output_check(inputs, dtype_compat) +} + +/// Macro to implement `create` constructors for operations with a single output. +/// +/// Supports shape and dtype validation. +macro_rules! impl_ir_create { + // Leaf: the default `create`, inferring the output dtype from the inputs. + (@create_fn $op:ident { $( $field:ident : $ty:ty ),* $(,)? } , $shape:expr, $dtype:expr) => { + #[doc = "Create a new operation IR from the given inputs."] + #[doc = "`new_id` should generate a unique `TensorId` for the uninitialized output tensor."] + #[allow(clippy::too_many_arguments)] + pub fn create($( $field : $ty ),*, new_id: impl FnOnce() -> crate::TensorId) -> $op { + let shape = $shape; + let dtype = $dtype; + let out = TensorIr::uninit(new_id(), shape, dtype); + $op { $( $field ),*, out } + } + }; + + // Leaf: a single additional constructor that takes an explicit output dtype. + (@extra_fn $op:ident { $( $field:ident : $ty:ty ),* $(,)? } , $shape:expr, $dtype:expr, $fn_name:ident ( $extra:ident : $extra_ty:ty )) => { + #[doc = "Create a new operation IR from the given inputs and the given output dtype."] + #[allow(clippy::too_many_arguments)] + pub fn $fn_name($( $field : $ty ),*, $extra: $extra_ty, new_id: impl FnOnce() -> crate::TensorId) -> $op { + let shape = $shape; + let _ = $dtype; // still validates dtype if needed + let out = TensorIr::uninit(new_id(), shape, $extra); + $op { $( $field ),*, out } + } + }; + + // Recursively emit each additional constructor in its own `impl` block. The field list is + // forwarded as an opaque token tree (`$fields`) so it is not iterated alongside the + // constructor list (the two repeat independently). + (@extras $op:ident $fields:tt, $shape:expr, $dtype:expr, $fn_name:ident ( $extra:ident : $extra_ty:ty ) $(, $rest_fn:ident ( $rest_extra:ident : $rest_extra_ty:ty ) )* $(,)?) => { + impl $op { + impl_ir_create!(@extra_fn $op $fields, $shape, $dtype, $fn_name ( $extra : $extra_ty )); + } + impl_ir_create!(@extras $op $fields, $shape, $dtype, $( $rest_fn ( $rest_extra : $rest_extra_ty ) ),*); + }; + (@extras $op:ident $fields:tt, $shape:expr, $dtype:expr, $(,)?) => {}; + + // Case: simple op, single `create` + ( + $op:ident $fields:tt, + shape = $shape:expr, + dtype = $dtype:expr $(,)? + ) => { + impl $op { + impl_ir_create!(@create_fn $op $fields, $shape, $dtype); + } + }; + + // Case: op with one or more additional constructors that accept an explicit output dtype + ( + $op:ident $fields:tt, + shape = $shape:expr, + dtype = $dtype:expr, + $fn_name:ident ( $extra:ident : $extra_ty:ty ) + $(, $rest_fn:ident ( $rest_extra:ident : $rest_extra_ty:ty ) )* $(,)? + ) => { + impl $op { + impl_ir_create!(@create_fn $op $fields, $shape, $dtype); + } + impl_ir_create!(@extras $op $fields, $shape, $dtype, $fn_name ( $extra : $extra_ty ) $(, $rest_fn ( $rest_extra : $rest_extra_ty ) )*); + }; +} + +impl_ir_create!( + UnaryOpIr { input: TensorIr }, + shape = input.shape.clone(), + dtype = input.dtype, + // Additional constructor for unary comparisons + create_comparison(bool_dtype: DType) +); + +impl_ir_create!( + BinaryOpIr { + lhs: TensorIr, + rhs: TensorIr + }, + shape = lhs.shape.broadcast(&rhs.shape).unwrap(), + dtype = output_dtype([&lhs.dtype, &rhs.dtype]).unwrap(), + // Additional constructor for binary comparisons + create_comparison(bool_dtype: DType) +); + +impl_ir_create!( + ScalarOpIr { + lhs: TensorIr, + rhs: ScalarIr + }, + shape = lhs.shape.clone(), + dtype = lhs.dtype, + // Additional constructor for scalar comparisons + create_comparison(bool_dtype: DType) +); + +impl_ir_create!( + MatmulOpIr { + lhs: TensorIr, + rhs: TensorIr + }, + shape = calculate_matmul_output(&lhs.shape, &rhs.shape).unwrap(), + dtype = output_dtype_mixed([&lhs.dtype, &rhs.dtype]).unwrap(), + // Additional constructor for mixed dtypes + create_mixed(out_dtype: DType) +); + +impl_ir_create!( + SwapDimsOpIr { + input: TensorIr, + dim1: usize, + dim2: usize + }, + shape = input.shape.clone().swapped(dim1, dim2).unwrap(), + dtype = input.dtype +); + +impl_ir_create!( + PermuteOpIr { input: TensorIr, axes: Vec }, + shape = input.shape.clone().permuted(&axes).unwrap(), + dtype = input.dtype +); + +impl_ir_create!( + RepeatDimOpIr { + tensor: TensorIr, + dim: usize, + times: usize + }, + shape = tensor.shape.clone().repeat(dim, times).unwrap(), + dtype = tensor.dtype +); + +impl_ir_create!( + FlipOpIr { input: TensorIr, axes: Vec }, + shape = input.shape.clone(), // TODO: check if axes are within the tensor dimensions + dtype = input.dtype +); + +impl_ir_create!( + CatOpIr { tensors: Vec, dim: usize }, + shape = Shape::cat(tensors.iter().map(|t| &t.shape), dim).unwrap(), + dtype = output_dtype(tensors.iter().map(|t| &t.dtype)).unwrap() +); + +impl_ir_create!( + AllReduceOpIr { + tensor: TensorIr, + op: burn_backend::distributed::ReduceOperation, + device_ids: Vec + }, + shape = tensor.shape.clone(), + dtype = tensor.dtype +); + +impl_ir_create!( + GatherOpIr { + tensor: TensorIr, + dim: usize, + indices: TensorIr + }, + shape = indices.shape.clone(), // TODO: check dims compat between tensor and indices + dtype = tensor.dtype +); + +impl_ir_create!( + ScatterOpIr { + tensor: TensorIr, + dim: usize, + indices: TensorIr, + value: TensorIr, + update: IndexingUpdateOp + }, + shape = tensor.shape.clone(), // TODO: check dims compat between tensor and indices + dtype = output_dtype([&tensor.dtype, &value.dtype]).unwrap() +); + +impl_ir_create!( + ScatterNdOpIr { + data: TensorIr, + indices: TensorIr, + values: TensorIr, + reduction: IndexingUpdateOp + }, + shape = data.shape.clone(), + dtype = output_dtype([&data.dtype, &values.dtype]).unwrap() +); + +impl GatherNdOpIr { + /// Create a new GatherNd IR operation. + pub fn create( + data: TensorIr, + indices: TensorIr, + new_id: impl FnOnce() -> crate::TensorId, + ) -> Self { + let m = indices.shape.num_dims(); + let k = indices.shape[m - 1]; + let mut dims = indices.shape.as_slice()[..m - 1].to_vec(); + dims.extend_from_slice(&data.shape.as_slice()[k..]); + let shape = Shape::from(dims); + let dtype = data.dtype; + let out = TensorIr::uninit(new_id(), shape, dtype); + GatherNdOpIr { data, indices, out } + } +} + +impl_ir_create!( + ReduceOpIr { input: TensorIr }, + shape = [1].into(), + dtype = input.dtype, + // Additional constructor for reduce-all/reduce-any (bool output) + create_bool(bool_dtype: DType) +); + +fn reduce_output_shape(mut output_shape: Shape, axis: usize, accumulator_len: usize) -> Shape { + assert!(output_shape.rank() > axis); + output_shape[axis] = accumulator_len; + output_shape +} + +impl_ir_create!( + ReduceDimOpIr { + input: TensorIr, + axis: usize, + accumulator_len: usize, + }, + shape = reduce_output_shape(input.shape.clone(), axis, accumulator_len), + dtype = input.dtype, + // Additional constructor for argument reduction + create_arg(ind_dtype: DType), + // Additional constructor for reduce-all/reduce-any along a dim (bool output) + create_bool(bool_dtype: DType) +); + +impl_ir_create!( + DimOpIr { + input: TensorIr, + axis: usize + }, + shape = input.shape.clone(), // TODO: check dims within rank + dtype = input.dtype +); + +impl_ir_create!( + SelectOpIr { + tensor: TensorIr, + dim: usize, + indices: TensorIr + }, + // TODO: shape.select? + shape = { + let mut s = tensor.shape.clone(); + s[dim] = indices.shape[0]; + s + }, + dtype = tensor.dtype +); + +impl_ir_create!( + SelectAssignOpIr { + tensor: TensorIr, + dim: usize, + indices: TensorIr, + value: TensorIr, + update: IndexingUpdateOp + }, + // TODO: check value and indices shape match for dim + shape = tensor.shape.clone(), + dtype = output_dtype([&tensor.dtype, &value.dtype]).unwrap() +); + +impl_ir_create!( + SliceOpIr { + tensor: TensorIr, + ranges: Vec, + }, + shape = tensor.shape.clone().slice(&ranges).unwrap(), + dtype = tensor.dtype +); + +impl_ir_create!( + SliceAssignOpIr { + tensor: TensorIr, + ranges: Vec, + value: TensorIr + }, + // TODO: check slice and value number of elements match + shape = tensor.shape.clone(), + dtype = output_dtype([&tensor.dtype, &value.dtype]).unwrap() +); + +impl_ir_create!( + MaskWhereOpIr { + tensor: TensorIr, + mask: TensorIr, + value: TensorIr + }, + shape = Shape::broadcast_many([&tensor.shape, &mask.shape, &value.shape]).unwrap(), + dtype = output_dtype([&tensor.dtype, &value.dtype]).unwrap() +); + +impl_ir_create!( + MaskFillOpIr { + tensor: TensorIr, + mask: TensorIr, + value: ScalarIr + }, + shape = tensor.shape.broadcast(&mask.shape).unwrap(), + dtype = tensor.dtype +); + +impl_ir_create!( + ClampOpIr { + tensor: TensorIr, + min: ScalarIr, + max: ScalarIr + }, + shape = tensor.shape.clone(), + dtype = tensor.dtype +); + +impl_ir_create!( + HardSigmoidOpIr { + tensor: TensorIr, + alpha: ScalarIr, + beta: ScalarIr + }, + shape = tensor.shape.clone(), + dtype = tensor.dtype +); + +impl_ir_create!( + SortOpIr { + input: TensorIr, + dim: usize, + descending: bool + }, + shape = input.shape.clone(), + dtype = input.dtype, + // Additional constructor for argsort (output is the indices tensor) + create_arg(ind_dtype: DType) +); + +impl SortWithIndicesOpIr { + /// Create a sort-with-indices IR. + pub fn create( + input: TensorIr, + dim: usize, + descending: bool, + indices_dtype: DType, + mut new_id: impl FnMut() -> TensorId, + ) -> Self { + let shape = input.shape.clone(); + let dtype = input.dtype; + let out = TensorIr::uninit(new_id(), shape.clone(), dtype); + let out_indices = TensorIr::uninit(new_id(), shape, indices_dtype); + SortWithIndicesOpIr { + input, + dim, + descending, + out, + out_indices, + } + } +} + +impl LayerNormOpIr { + /// Create a layer-norm IR. + pub fn create( + input: TensorIr, + gamma: TensorIr, + beta: Option, + epsilon: f64, + new_id: impl FnOnce() -> TensorId, + ) -> Self { + let dtype = output_dtype( + [ + Some(&input.dtype), + Some(&gamma.dtype), + beta.as_ref().map(|b| &b.dtype), + ] + .iter() + .filter_map(|&d| d), + ) + .unwrap(); + let out = TensorIr::uninit(new_id(), input.shape.clone(), dtype); + LayerNormOpIr { + input, + gamma, + beta, + epsilon: ScalarIr::Float(epsilon), + out, + } + } +} + +impl Unfold4dOpIr { + /// Create an unfold4d IR. + pub fn create( + x: TensorIr, + kernel_size: [usize; 2], + options: Unfold4dOptionsIr, + new_id: impl FnOnce() -> TensorId, + ) -> Self { + // Output shape mirrors the implementation in `unfold4d_using_unfold`: + // [N, C * kH * kW, num_blocks_h * num_blocks_w] + let dilation = options.dilation; + let padding = options.padding; + let stride = options.stride; + let weight_shape = Shape::from([ + x.shape[1] * kernel_size[0] * kernel_size[1], + x.shape[1], + kernel_size[0], + kernel_size[1], + ]); + let conv_options = burn_backend::ops::ConvOptions::new(stride, padding, dilation, 1); + let out_shape = burn_backend::ops::conv::calculate_conv_output_shape( + &x.shape, + &weight_shape, + &conv_options.stride, + &conv_options.padding, + &conv_options.dilation, + ) + .unwrap(); + let shape = Shape::from([ + x.shape[0], + x.shape[1] * kernel_size[0] * kernel_size[1], + out_shape[2] * out_shape[3], + ]); + let out = TensorIr::uninit(new_id(), shape, x.dtype); + Unfold4dOpIr { + x, + kernel_size, + options, + out, + } + } +} + +impl_ir_create!( + ConvTranspose1dWeightBackwardOpIr { + x: TensorIr, + weight: TensorIr, + output_grad: TensorIr, + options: ConvTranspose1dOptionsIr + }, + shape = weight.shape.clone(), + dtype = output_grad.dtype +); + +impl_ir_create!( + ConvTranspose1dBiasBackwardOpIr { + x: TensorIr, + bias: TensorIr, + output_grad: TensorIr, + }, + shape = bias.shape.clone(), + dtype = output_grad.dtype +); + +impl_ir_create!( + ConvTranspose2dWeightBackwardOpIr { + x: TensorIr, + weight: TensorIr, + output_grad: TensorIr, + options: ConvTranspose2dOptionsIr + }, + shape = weight.shape.clone(), + dtype = output_grad.dtype +); + +impl_ir_create!( + ConvTranspose2dBiasBackwardOpIr { + x: TensorIr, + bias: TensorIr, + output_grad: TensorIr, + }, + shape = bias.shape.clone(), + dtype = output_grad.dtype +); + +impl_ir_create!( + ConvTranspose3dWeightBackwardOpIr { + x: TensorIr, + weight: TensorIr, + output_grad: TensorIr, + options: ConvTranspose3dOptionsIr + }, + shape = weight.shape.clone(), + dtype = output_grad.dtype +); + +impl_ir_create!( + ConvTranspose3dBiasBackwardOpIr { + x: TensorIr, + bias: TensorIr, + output_grad: TensorIr, + }, + shape = bias.shape.clone(), + dtype = output_grad.dtype +); + +impl_ir_create!( + AvgPool1dOpIr { + x: TensorIr, + kernel_size: usize, + stride: usize, + padding: usize, + count_include_pad: bool, + ceil_mode: bool + }, + shape = calculate_pool_output_shape( + &x.shape, + &[kernel_size], + &[stride], + &[padding], + &[1], + ceil_mode + ) + .unwrap(), + dtype = x.dtype +); + +impl_ir_create!( + AvgPool1dBackwardOpIr { + x: TensorIr, + grad: TensorIr, + kernel_size: usize, + stride: usize, + padding: usize, + count_include_pad: bool, + ceil_mode: bool + }, + shape = x.shape.clone(), + dtype = x.dtype +); + +impl_ir_create!( + AvgPool2dOpIr { + x: TensorIr, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool + }, + shape = calculate_pool_output_shape( + &x.shape, + &kernel_size, + &stride, + &padding, + &[1, 1], + ceil_mode + ) + .unwrap(), + dtype = x.dtype +); + +impl_ir_create!( + AvgPool2dBackwardOpIr { + x: TensorIr, + grad: TensorIr, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool + }, + shape = x.shape.clone(), + dtype = x.dtype +); + +impl_ir_create!( + MaxPool1dOpIr { + x: TensorIr, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool + }, + shape = calculate_pool_output_shape( + &x.shape, + &[kernel_size], + &[stride], + &[padding], + &[dilation], + ceil_mode + ) + .unwrap(), + dtype = x.dtype +); + +impl_ir_create!( + MaxPool2dOpIr { + x: TensorIr, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool + }, + shape = calculate_pool_output_shape( + &x.shape, + &kernel_size, + &stride, + &padding, + &dilation, + ceil_mode + ) + .unwrap(), + dtype = x.dtype +); + +impl_ir_create!( + MaxPool1dWithIndicesBackwardOpIr { + x: TensorIr, + grad: TensorIr, + indices: TensorIr, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool + }, + shape = x.shape.clone(), + dtype = x.dtype +); + +impl_ir_create!( + MaxPool2dWithIndicesBackwardOpIr { + x: TensorIr, + grad: TensorIr, + indices: TensorIr, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool + }, + shape = x.shape.clone(), + dtype = x.dtype +); + +impl_ir_create!( + AdaptiveAvgPool1dOpIr { + x: TensorIr, + output_size: usize + }, + shape = Shape::new([x.shape[0], x.shape[1], output_size]), + dtype = x.dtype +); + +impl_ir_create!( + AdaptiveAvgPool2dOpIr { + x: TensorIr, + output_size: [usize; 2] + }, + shape = Shape::new([x.shape[0], x.shape[1], output_size[0], output_size[1]]), + dtype = x.dtype +); + +impl_ir_create!( + AdaptiveAvgPool1dBackwardOpIr { + x: TensorIr, + grad: TensorIr, + }, + shape = x.shape.clone(), + dtype = x.dtype +); + +impl_ir_create!( + AdaptiveAvgPool2dBackwardOpIr { + x: TensorIr, + grad: TensorIr, + }, + shape = x.shape.clone(), + dtype = x.dtype +); + +impl_ir_create!( + InterpolateOpIr { + x: TensorIr, + output_size: [usize; 2], + options: InterpolateOptionsIr + }, + shape = Shape::new([x.shape[0], x.shape[1], output_size[0], output_size[1]]), + dtype = x.dtype +); + +impl_ir_create!( + InterpolateBackwardOpIr { + x: TensorIr, + grad: TensorIr, + output_size: [usize; 2], + options: InterpolateOptionsIr + }, + shape = x.shape.clone(), + dtype = x.dtype +); + +impl_ir_create!( + GridSample2dOpIr { + tensor: TensorIr, + grid: TensorIr, + options: GridSampleOptionsIr + }, + // Input tensor: [N, C, H_in, W_in] + // Grid: [N, H_out, W_out, 2] + // Output: [N, C, H_out, W_out] + shape = Shape::new([ + tensor.shape[0], + tensor.shape[1], + grid.shape[1], + grid.shape[2] + ]), + dtype = tensor.dtype +); + +impl_ir_create!( + EmbeddingOpIr { + weights: TensorIr, + indices: TensorIr, + }, + shape = { + // weights: [n_embeddings, d_model] + // indices: [batch_size, seq_length] + // output: [batch_size, seq_length, d_model] + let d_model = weights.shape[1]; + Shape::from(alloc::vec![indices.shape[0], indices.shape[1], d_model]) + }, + dtype = weights.dtype +); + +impl_ir_create!( + EmbeddingBackwardOpIr { + weights: TensorIr, + out_grad: TensorIr, + indices: TensorIr, + }, + shape = weights.shape.clone(), + dtype = output_dtype([&weights.dtype, &out_grad.dtype]).unwrap() +); + +impl_ir_create!( + LinearOpIr { + x: TensorIr, + weight: TensorIr, + bias: Option + }, + shape = { + // output: [..., d_output] where x is [..., d_input] and weight is [d_input, d_output] + let n = x.shape.num_dims(); + let mut dims: Vec = (0..n).map(|i| x.shape[i]).collect(); + dims[n - 1] = weight.shape[1]; + Shape::from(dims) + }, + dtype = output_dtype( + [ + Some(&x.dtype), + Some(&weight.dtype), + bias.as_ref().map(|b| &b.dtype), + ] + .iter() + .filter_map(|&d| d), + ) + .unwrap() +); + +impl_ir_create!( + LinearXBackwardOpIr { + weight: TensorIr, + output_grad: TensorIr, + }, + shape = { + // dx = output_grad @ weight^T + // output_grad: [..., d_output], weight: [d_input, d_output] + // result: [..., d_input] + let n = output_grad.shape.num_dims(); + let mut dims: Vec = (0..n).map(|i| output_grad.shape[i]).collect(); + dims[n - 1] = weight.shape[0]; + Shape::from(dims) + }, + dtype = output_grad.dtype +); + +impl_ir_create!( + LinearWeightBackwardOpIr { + x: TensorIr, + output_grad: TensorIr, + }, + shape = { + // dW: [d_input, d_output] + let d_input = x.shape[x.shape.num_dims() - 1]; + let d_output = output_grad.shape[output_grad.shape.num_dims() - 1]; + Shape::from(alloc::vec![d_input, d_output]) + }, + dtype = output_grad.dtype +); + +impl_ir_create!( + LinearBiasBackwardOpIr { + output_grad: TensorIr, + }, + shape = { + // db: [d_output] + let d_output = output_grad.shape[output_grad.shape.num_dims() - 1]; + Shape::from(alloc::vec![d_output]) + }, + dtype = output_grad.dtype +); + +impl_ir_create!( + Conv1dOpIr { + x: TensorIr, + weight: TensorIr, + bias: Option, + options: Conv1dOptionsIr + }, + shape = calculate_conv_output_shape( + &x.shape, + &weight.shape, + &options.stride, + &options.padding, + &options.dilation, + ) + .unwrap(), + dtype = output_dtype( + [ + Some(&x.dtype), + Some(&weight.dtype), + bias.as_ref().map(|b| &b.dtype), + ] + .iter() + .filter_map(|&d| d), + ) + .unwrap() +); + +impl_ir_create!( + Conv1dXBackwardOpIr { + x: TensorIr, + weight: TensorIr, + output_grad: TensorIr, + options: Conv1dOptionsIr + }, + shape = x.shape.clone(), + dtype = output_grad.dtype +); + +impl_ir_create!( + Conv1dWeightBackwardOpIr { + x: TensorIr, + weight: TensorIr, + output_grad: TensorIr, + options: Conv1dOptionsIr + }, + shape = weight.shape.clone(), + dtype = output_grad.dtype +); + +impl_ir_create!( + Conv1dBiasBackwardOpIr { + x: TensorIr, + bias: TensorIr, + output_grad: TensorIr, + }, + shape = bias.shape.clone(), + dtype = output_grad.dtype +); + +impl_ir_create!( + Conv2dOpIr { + x: TensorIr, + weight: TensorIr, + bias: Option, + options: Conv2dOptionsIr + }, + shape = calculate_conv_output_shape( + &x.shape, + &weight.shape, + &options.stride, + &options.padding, + &options.dilation, + ) + .unwrap(), + dtype = output_dtype( + [ + Some(&x.dtype), + Some(&weight.dtype), + bias.as_ref().map(|b| &b.dtype), + ] + .iter() + .filter_map(|&d| d), + ) + .unwrap() +); + +impl_ir_create!( + Conv2dXBackwardOpIr { + x: TensorIr, + weight: TensorIr, + output_grad: TensorIr, + options: Conv2dOptionsIr + }, + shape = x.shape.clone(), + dtype = output_grad.dtype +); + +impl_ir_create!( + Conv2dWeightBackwardOpIr { + x: TensorIr, + weight: TensorIr, + output_grad: TensorIr, + options: Conv2dOptionsIr + }, + shape = weight.shape.clone(), + dtype = output_grad.dtype +); + +impl_ir_create!( + Conv2dBiasBackwardOpIr { + x: TensorIr, + bias: TensorIr, + output_grad: TensorIr, + }, + shape = bias.shape.clone(), + dtype = output_grad.dtype +); + +impl_ir_create!( + Conv3dOpIr { + x: TensorIr, + weight: TensorIr, + bias: Option, + options: Conv3dOptionsIr + }, + shape = calculate_conv_output_shape( + &x.shape, + &weight.shape, + &options.stride, + &options.padding, + &options.dilation, + ) + .unwrap(), + dtype = output_dtype( + [ + Some(&x.dtype), + Some(&weight.dtype), + bias.as_ref().map(|b| &b.dtype), + ] + .iter() + .filter_map(|&d| d), + ) + .unwrap() +); + +impl_ir_create!( + Conv3dXBackwardOpIr { + x: TensorIr, + weight: TensorIr, + output_grad: TensorIr, + options: Conv3dOptionsIr + }, + shape = x.shape.clone(), + dtype = output_grad.dtype +); + +impl_ir_create!( + Conv3dWeightBackwardOpIr { + x: TensorIr, + weight: TensorIr, + output_grad: TensorIr, + options: Conv3dOptionsIr + }, + shape = weight.shape.clone(), + dtype = output_grad.dtype +); + +impl_ir_create!( + Conv3dBiasBackwardOpIr { + x: TensorIr, + bias: TensorIr, + output_grad: TensorIr, + }, + shape = bias.shape.clone(), + dtype = output_grad.dtype +); + +impl_ir_create!( + DeformConv2dOpIr { + x: TensorIr, + offset: TensorIr, + weight: TensorIr, + mask: Option, + bias: Option, + options: DeformableConv2dOptionsIr + }, + shape = calculate_conv_output_shape( + &x.shape, + &weight.shape, + &options.stride, + &options.padding, + &options.dilation, + ) + .unwrap(), + dtype = output_dtype( + [ + Some(&x.dtype), + Some(&offset.dtype), + Some(&weight.dtype), + mask.as_ref().map(|m| &m.dtype), + bias.as_ref().map(|b| &b.dtype), + ] + .iter() + .filter_map(|&d| d), + ) + .unwrap() +); + +impl_ir_create!( + ConvTranspose1dOpIr { + x: TensorIr, + weight: TensorIr, + bias: Option, + options: ConvTranspose1dOptionsIr + }, + shape = calculate_conv_transpose_output_shape( + &x.shape, + &weight.shape, + &options.stride, + &options.padding, + &options.padding_out, + &options.dilation, + options.groups, + ) + .unwrap(), + dtype = output_dtype( + [ + Some(&x.dtype), + Some(&weight.dtype), + bias.as_ref().map(|b| &b.dtype), + ] + .iter() + .filter_map(|&d| d), + ) + .unwrap() +); + +impl_ir_create!( + ConvTranspose2dOpIr { + x: TensorIr, + weight: TensorIr, + bias: Option, + options: ConvTranspose2dOptionsIr + }, + shape = calculate_conv_transpose_output_shape( + &x.shape, + &weight.shape, + &options.stride, + &options.padding, + &options.padding_out, + &options.dilation, + options.groups, + ) + .unwrap(), + dtype = output_dtype( + [ + Some(&x.dtype), + Some(&weight.dtype), + bias.as_ref().map(|b| &b.dtype), + ] + .iter() + .filter_map(|&d| d), + ) + .unwrap() +); + +impl_ir_create!( + ConvTranspose3dOpIr { + x: TensorIr, + weight: TensorIr, + bias: Option, + options: ConvTranspose3dOptionsIr + }, + shape = calculate_conv_transpose_output_shape( + &x.shape, + &weight.shape, + &options.stride, + &options.padding, + &options.padding_out, + &options.dilation, + options.groups, + ) + .unwrap(), + dtype = output_dtype( + [ + Some(&x.dtype), + Some(&weight.dtype), + bias.as_ref().map(|b| &b.dtype), + ] + .iter() + .filter_map(|&d| d), + ) + .unwrap() +); + +impl_ir_create!( + UnfoldOpIr { + input: TensorIr, + dim: usize, + size: usize, + step: usize + }, + shape = calculate_unfold_shape(input.shape.clone(), dim, size, step), + dtype = input.dtype +); + +impl_ir_create!( + CrossOpIr { + lhs: TensorIr, + rhs: TensorIr, + dim: usize + }, + shape = lhs.shape.broadcast(&rhs.shape).unwrap(), + dtype = output_dtype([&lhs.dtype, &rhs.dtype]).unwrap() +); + +impl_ir_create!( + QuantizeOpIr { + tensor: TensorIr, + qparams: QuantizationParametersIr, + scheme: QuantScheme + }, + shape = tensor.shape.clone(), + dtype = DType::QFloat(scheme) +); + +impl_ir_create!( + AttentionOpIr { + query: TensorIr, + key: TensorIr, + value: TensorIr, + mask: Option, + attn_bias: Option, + options: AttentionOptionsIr, + }, + shape = Shape::new([query.shape[0], query.shape[1], query.shape[2], value.shape[3]]), + dtype = query.dtype +); + +impl_ir_create!( + CtcLossOpIr { + log_probs: TensorIr, + targets: TensorIr, + input_lengths: TensorIr, + target_lengths: TensorIr, + blank: usize, + }, + shape = Shape::new([log_probs.shape[1]]), + dtype = log_probs.dtype +); + +impl_ir_create!( + CtcLossBackwardOpIr { + log_probs: TensorIr, + targets: TensorIr, + input_lengths: TensorIr, + target_lengths: TensorIr, + grad_loss: TensorIr, + blank: usize, + }, + shape = log_probs.shape.clone(), + dtype = log_probs.dtype +); + +impl DequantizeOpIr { + pub fn create(input: TensorIr, dtype: DType, new_id: impl FnOnce() -> TensorId) -> Self { + let out = TensorIr::uninit(new_id(), input.shape.clone(), dtype); + + DequantizeOpIr { input, out } + } +} + +// Operations with multiple outputs + +impl ReduceDimWithIndicesOpIr { + pub fn create( + tensor: TensorIr, + dim: usize, + dtype_indices: DType, + mut new_id: impl FnMut() -> TensorId, + ) -> Self { + let mut shape = tensor.shape.clone(); + shape[dim] = 1; + let out = TensorIr::uninit(new_id(), shape.clone(), tensor.dtype); + let out_indices = TensorIr::uninit(new_id(), shape.clone(), dtype_indices); + + ReduceDimWithIndicesOpIr { + tensor, + dim, + out, + out_indices, + } + } +} + +impl DeformConv2dBackwardOpIr { + #[allow(clippy::too_many_arguments)] + pub fn create( + x: TensorIr, + offset: TensorIr, + weight: TensorIr, + mask: Option, + bias: Option, + out_grad: TensorIr, + options: DeformableConv2dOptionsIr, + mut new_id: impl FnMut() -> TensorId, + ) -> Self { + let dtype = output_dtype( + [ + Some(&x.dtype), + Some(&weight.dtype), + mask.as_ref().map(|m| &m.dtype), + bias.as_ref().map(|b| &b.dtype), + ] + .iter() + .filter_map(|&d| d), + ) + .unwrap(); + + let input_grad = TensorIr::uninit(new_id(), x.shape.clone(), dtype); + let offset_grad = TensorIr::uninit(new_id(), offset.shape.clone(), dtype); + let weight_grad = TensorIr::uninit(new_id(), weight.shape.clone(), dtype); + let mask_grad = mask + .as_ref() + .map(|t| TensorIr::uninit(new_id(), t.shape.clone(), dtype)); + let bias_grad = bias + .as_ref() + .map(|t| TensorIr::uninit(new_id(), t.shape.clone(), dtype)); + + DeformConv2dBackwardOpIr { + x, + offset, + weight, + mask, + bias, + out_grad, + options, + input_grad, + offset_grad, + weight_grad, + mask_grad, + bias_grad, + } + } +} + +impl MaxPool1dWithIndicesOpIr { + #[allow(clippy::too_many_arguments)] + pub fn create( + x: TensorIr, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + dtype_indices: DType, + mut new_id: impl FnMut() -> TensorId, + ) -> Self { + let shape = calculate_pool_output_shape( + &x.shape, + &[kernel_size], + &[stride], + &[padding], + &[dilation], + ceil_mode, + ) + .unwrap(); + let out = TensorIr::uninit(new_id(), shape.clone(), x.dtype); + let out_indices = TensorIr::uninit(new_id(), shape, dtype_indices); + + MaxPool1dWithIndicesOpIr { + x, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + out, + out_indices, + } + } +} + +impl MaxPool2dWithIndicesOpIr { + #[allow(clippy::too_many_arguments)] + pub fn create( + x: TensorIr, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + dtype_indices: DType, + mut new_id: impl FnMut() -> TensorId, + ) -> Self { + let shape = calculate_pool_output_shape( + &x.shape, + &kernel_size, + &stride, + &padding, + &dilation, + ceil_mode, + ) + .unwrap(); + let out = TensorIr::uninit(new_id(), shape.clone(), x.dtype); + let out_indices = TensorIr::uninit(new_id(), shape, dtype_indices); + + MaxPool2dWithIndicesOpIr { + x, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + out, + out_indices, + } + } +} diff --git a/crates/burn-ir/src/graph.rs b/crates/burn-ir/src/graph.rs new file mode 100644 index 0000000..c5c58b9 --- /dev/null +++ b/crates/burn-ir/src/graph.rs @@ -0,0 +1,39 @@ +use crate::{ScalarIr, TensorId}; +use alloc::vec::Vec; +use burn_backend::Slice; +use serde::{Deserialize, Serialize}; + +/// Identifier for a cached, reusable group of operations (a graph). +/// +/// A router backend that supports graph caching (e.g. the remote backend) registers a relative +/// op-graph once under this id, then replays it by id with only the changing [bindings](GraphBindings). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct GraphId(pub u64); + +/// Per-invocation bindings used to specialize a cached graph to concrete tensors. +/// +/// The cached graph is in *relative* form (positional tensor ids, relative shape-dim ids, scalar +/// placeholders). The bindings carry only the graph's *boundary* — its inputs and surviving +/// outputs — plus the dense shape-dim table and the scalar values. The replay reconstructs the +/// concrete shape of **every** tensor (including intermediates) from [`shapes`](Self::shapes), and +/// allocates fresh ids for intermediate tensors itself, so the payload stays small regardless of +/// how many ops the graph contains. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct GraphBindings { + /// Boundary tensors only: `(relative id, concrete id)` for each graph input and surviving + /// output. Intermediate tensors are *not* listed — the replaying backend allocates their ids. + pub tensors: Vec<(TensorId, TensorId)>, + /// Concrete dim value for each relative shape-dim id, indexed directly by the id. Relative dim + /// ids are dense (`0..N`), so this is a plain table rather than a map, and each distinct dim + /// value appears once however many tensors share it. Lets the replay rebuild every tensor's + /// concrete shape — inputs, outputs and intermediates alike. + pub shapes: Vec, + /// Concrete scalar values indexed by their placeholder id (the value carried in a relativized + /// `ScalarIr::UInt(placeholder)`). + pub scalars: Vec, + /// Concrete slice ranges indexed by their placeholder id (the value carried in a relativized + /// range's `start` field). The relative graph keeps every `Slice` range as a positional + /// placeholder — its actual bounds are discarded by relativization (they vary per invocation, + /// like scalars) — so the replay restores each from here. + pub ranges: Vec, +} diff --git a/crates/burn-ir/src/handle.rs b/crates/burn-ir/src/handle.rs new file mode 100644 index 0000000..8c12b6f --- /dev/null +++ b/crates/burn-ir/src/handle.rs @@ -0,0 +1,312 @@ +use hashbrown::HashMap; + +use crate::{BackendIr, TensorHandle, TensorId, TensorIr, TensorStatus}; + +/// Keep all [tensor handles](BackendIr::Handle) in one place and ensure that all resources +/// are used optimally. +pub struct HandleContainer { + handles: HashMap>, + counter: u64, +} + +// Hand-written perfect derive as we don't require `H: Default`. +impl Default for HandleContainer { + fn default() -> Self { + Self { + handles: HashMap::new(), + counter: 0, + } + } +} + +impl HandleContainer { + /// Fork the container, useful for autotune. + pub fn fork(&self) -> Self { + let mut handles = HashMap::with_capacity(self.handles.len()); + + for (id, handle) in self.handles.iter() { + handles.insert(*id, handle.clone()); + } + + Self { + handles, + counter: self.counter, + } + } +} + +impl core::fmt::Debug for HandleContainer { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("HandleContainer") + .field("handles", &self.handles.keys()) // only care about the IDs when debugging + .field("counter", &self.counter) + .finish() + } +} + +/// Backend [tensor handle](BackendIr::Handle) wrapper tracking their creation state +#[derive(Clone)] +pub enum Handle { + /// No [tensor handle](BackendIr::Handle) has been created yet + NotInit, + /// A [tensor handle](BackendIr::Handle) has been created + Existing(H), +} + +impl HandleContainer { + /// Create a new HandleContainer + pub fn new() -> Self { + Self { + handles: HashMap::new(), + counter: 0, + } + } + + /// Register a handle for the given [tensor id](TensorId). + pub fn register_handle(&mut self, id: TensorId, handle: H) { + self.handles.insert(id, Handle::Existing(handle)); + } + + /// Whether a handle exists. + pub fn has_handle(&self, id: &TensorId) -> bool { + self.handles.contains_key(id) + } + + /// Get the reference to a handle. + pub fn get_handle_ref(&self, id: &TensorId) -> Option<&H> { + self.handles + .get(id) + .filter(|h| !matches!(h, Handle::NotInit)) + .map(|h| match h { + Handle::Existing(handle) => handle, + Handle::NotInit => unreachable!(), + }) + } + + /// Get the handle for the given [tensor id](TensorId). The status is used to determine if the + /// tensor should be popped out of the current tensor map, necessary for inplace operations. + /// + /// # Warnings + /// + /// Make sure the status corresponds to the operation you want to execute the handle on, + /// otherwise you might remove a tensor handle that will be required in the future. + pub fn get_handle(&mut self, id: &TensorId, status: &TensorStatus) -> H { + let (id, handle) = self + .handles + .remove_entry(id) + .unwrap_or_else(|| panic!("Should have handle for tensor {id:?}")); + + match handle { + Handle::Existing(handle) => match status { + TensorStatus::ReadOnly => { + self.handles.insert(id, Handle::Existing(handle.clone())); + handle + } + TensorStatus::ReadWrite => handle, + TensorStatus::NotInit => panic!( + "Cannot get uninitialized tensor {id:?}. Tensor exist but with wrong status" + ), + }, + Handle::NotInit => panic!("Cannot get uninitialized handle {id:?}."), + } + } + + /// Get the tensor handle for the given [tensor intermediate representation](TensorIr). + pub fn get_tensor_handle(&mut self, tensor: &TensorIr) -> TensorHandle { + TensorHandle { + handle: self.get_handle(&tensor.id, &tensor.status), + shape: tensor.shape.clone(), + } + } + + /// Get the [float tensor](burn_backend::backend::BackendTypes::FloatTensorPrimitive) corresponding to the + /// given [tensor intermediate representation](TensorIr). + pub fn get_float_tensor(&mut self, tensor: &TensorIr) -> B::FloatTensorPrimitive + where + B: BackendIr, + { + B::float_tensor(self.get_tensor_handle(tensor)) + } + + /// Get the [int tensor](burn_backend::backend::BackendTypes::IntTensorPrimitive) corresponding to the + /// given [tensor intermediate representation](TensorIr). + pub fn get_int_tensor(&mut self, tensor: &TensorIr) -> B::IntTensorPrimitive + where + B: BackendIr, + { + B::int_tensor(self.get_tensor_handle(tensor)) + } + + /// Get the [bool tensor](burn_backend::backend::BackendTypes::BoolTensorPrimitive) corresponding to the + /// given [tensor intermediate representation](TensorIr). + pub fn get_bool_tensor(&mut self, tensor: &TensorIr) -> B::BoolTensorPrimitive + where + B: BackendIr, + { + B::bool_tensor(self.get_tensor_handle(tensor)) + } + + /// Get the [quantized tensor](burn_backend::backend::BackendTypes::QuantizedTensorPrimitive) corresponding to the + /// given [tensor intermediate representation](TensorIr). + pub fn get_quantized_tensor(&mut self, tensor: &TensorIr) -> B::QuantizedTensorPrimitive + where + B: BackendIr, + { + B::quantized_tensor(self.get_tensor_handle(tensor)) + } + + /// Register a new [float tensor](burn_backend::backend::BackendTypes::FloatTensorPrimitive) with the corresponding [tensor id](TensorId). + pub fn register_float_tensor(&mut self, id: &TensorId, tensor: B::FloatTensorPrimitive) + where + B: BackendIr, + { + let handle = B::float_tensor_handle(tensor); + self.handles.insert(*id, Handle::Existing(handle)); + } + + /// Register a new [quantized tensor](burn_backend::backend::BackendTypes::QuantizedTensorPrimitive) with the corresponding [tensor ids](TensorId). + pub fn register_quantized_tensor( + &mut self, + id: &TensorId, + tensor: B::QuantizedTensorPrimitive, + ) where + B: BackendIr, + { + let handle = B::quantized_tensor_handle(tensor); + self.handles.insert(*id, Handle::Existing(handle)); + } + + /// Register a new [int tensor](burn_backend::backend::BackendTypes::IntTensorPrimitive) with the corresponding [tensor id](TensorId). + pub fn register_int_tensor(&mut self, id: &TensorId, tensor: B::IntTensorPrimitive) + where + B: BackendIr, + { + let handle = B::int_tensor_handle(tensor); + self.handles.insert(*id, Handle::Existing(handle)); + } + + /// Register a new [bool tensor](burn_backend::backend::BackendTypes::BoolTensorPrimitive) with the corresponding [tensor id](TensorId). + pub fn register_bool_tensor(&mut self, id: &TensorId, tensor: B::BoolTensorPrimitive) + where + B: BackendIr, + { + let handle = B::bool_tensor_handle(tensor); + self.handles.insert(*id, Handle::Existing(handle)); + } + + /// Remove tensor handle from container. + pub fn remove_handle(&mut self, id: TensorId) -> Option> { + self.handles.remove(&id) + } + + /// Remove tensor handle from container if writable + pub fn free(&mut self, tensor: &TensorIr) { + match tensor.status { + TensorStatus::ReadOnly => (), + TensorStatus::NotInit => (), + TensorStatus::ReadWrite => { + self.handles.remove(&tensor.id); + } + }; + } + + /// Returns the number of handles. + pub fn num_handles(&self) -> usize { + self.handles.len() + } + + /// Returns the IDs of all currently registered handles. + /// + /// Useful for snapshotting which handles exist at a point in time (e.g., before + /// executing on a forked context) so that newly registered output handles can + /// be detected afterwards. + pub fn handle_ids(&self) -> impl Iterator { + self.handles.keys() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::TensorId; + + /// Helper to create a TensorId for tests. + fn tid(value: u64) -> TensorId { + TensorId::new(value) + } + + #[test] + fn fork_clones_existing_handles() { + let mut container = HandleContainer::::new(); + container.register_handle(tid(1), "input_a".to_string()); + container.register_handle(tid(2), "input_b".to_string()); + + let fork = container.fork(); + + assert_eq!(fork.num_handles(), 2); + assert!(fork.get_handle_ref(&tid(1)).is_some()); + assert!(fork.get_handle_ref(&tid(2)).is_some()); + } + + #[test] + fn fork_is_isolated_from_original() { + // This test documents the core of the autotune clone bug: + // output handles registered in a fork do NOT appear in the original. + let mut container = HandleContainer::::new(); + container.register_handle(tid(1), "input_a".to_string()); + + let mut fork = container.fork(); + + // Simulate an optimization registering output handles in the fork. + fork.register_handle(tid(100), "output_x".to_string()); + fork.register_handle(tid(101), "output_y".to_string()); + + // The fork has the output handles. + assert_eq!(fork.num_handles(), 3); + assert!(fork.get_handle_ref(&tid(100)).is_some()); + assert!(fork.get_handle_ref(&tid(101)).is_some()); + + // But the original does NOT — these output handles are lost. + assert_eq!(container.num_handles(), 1); + assert!(container.get_handle_ref(&tid(100)).is_none()); + assert!(container.get_handle_ref(&tid(101)).is_none()); + } + + #[test] + fn fork_mutations_do_not_affect_original() { + let mut container = HandleContainer::::new(); + container.register_handle(tid(1), "original_value".to_string()); + + let mut fork = container.fork(); + + // Overwrite a handle in the fork (e.g., inplace output reuse). + fork.register_handle(tid(1), "modified_in_fork".to_string()); + + // Original is unchanged. + assert_eq!( + container.get_handle_ref(&tid(1)), + Some(&"original_value".to_string()) + ); + assert_eq!( + fork.get_handle_ref(&tid(1)), + Some(&"modified_in_fork".to_string()) + ); + } + + #[test] + fn double_fork_is_fully_isolated() { + // Simulates what happens when UnsafeTuneContext::get() is called on a Fork: + // it forks again, creating a second level of isolation. + let mut container = HandleContainer::::new(); + container.register_handle(tid(1), "input".to_string()); + + let fork1 = container.fork(); + let mut fork2 = fork1.fork(); + + fork2.register_handle(tid(200), "deep_output".to_string()); + + assert!(fork1.get_handle_ref(&tid(200)).is_none()); + assert!(container.get_handle_ref(&tid(200)).is_none()); + assert!(fork2.get_handle_ref(&tid(200)).is_some()); + } +} diff --git a/crates/burn-ir/src/lib.rs b/crates/burn-ir/src/lib.rs new file mode 100644 index 0000000..709b381 --- /dev/null +++ b/crates/burn-ir/src/lib.rs @@ -0,0 +1,23 @@ +#![cfg_attr(not(feature = "std"), no_std)] +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] + +//! Burn intermediate representation. + +extern crate alloc; + +mod backend; +mod builder; +mod graph; +mod handle; +mod operation; +mod scalar; +mod tensor; + +pub use backend::*; +pub use builder::*; +pub use graph::*; +pub use handle::*; +pub use operation::*; +pub use scalar::*; +pub use tensor::*; diff --git a/crates/burn-ir/src/operation.rs b/crates/burn-ir/src/operation.rs new file mode 100644 index 0000000..0acf555 --- /dev/null +++ b/crates/burn-ir/src/operation.rs @@ -0,0 +1,5132 @@ +use burn_backend::ops::AttentionModuleOptions; +use burn_backend::tensor::IndexingUpdateOp; +use core::hash::Hash; +use serde::{Deserialize, Serialize}; + +use alloc::borrow::ToOwned; +use alloc::boxed::Box; +use alloc::{string::String, vec::Vec}; + +use burn_backend::{ + DType, Distribution, Slice, + ops::{ + ConvOptions, ConvTransposeOptions, DeformConvOptions, GridSampleOptions, + GridSamplePaddingMode, InterpolateMode, InterpolateOptions, + }, + quantization::QuantScheme, +}; + +use crate::{ScalarIr, TensorId, TensorIr, TensorStatus}; + +/// Visitor for mutating the components of an [`OperationIr`] in place. +pub trait IrVisitorMut { + /// Visit a [`TensorIr`] mutably. + fn visit_tensor_mut(&mut self, _tensor: &mut TensorIr) {} + /// Visit a [`ScalarIr`] mutably. + fn visit_scalar_mut(&mut self, _scalar: &mut ScalarIr) {} + /// Visit a slice range mutably. + fn visit_range_mut(&mut self, _range: &mut Slice) {} +} + +/// Custom operation in fusion stream, declaring its inputs, outputs and scalar arguments. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct CustomOpIr { + /// Unique identifier of the operation. + pub id: String, + /// Input tensors used in the custom operation. + pub inputs: Vec, + /// Output tensors used in the custom operation. + pub outputs: Vec, + /// Non-tensor scalar arguments, in declaration order. + /// + /// Carried through fusion's relativization like any other scalar (see the + /// `RelativeOps for CustomOpIr` impl), so a cached graph replays with fresh scalar values. The + /// remote backend relies on these to ship a custom op's scalar arguments to the server, where + /// the registered handler reads them back with [`ScalarIr::elem`]. + pub scalars: Vec, +} + +impl CustomOpIr { + /// Create a new custom operation intermediate representation (without scalar arguments). + pub fn new(id: &'static str, inputs: &[TensorIr], outputs: &[TensorIr]) -> Self { + Self { + id: id.to_owned(), + inputs: inputs.to_vec(), + outputs: outputs.to_vec(), + scalars: Vec::new(), + } + } + + /// Create a new custom operation intermediate representation with scalar arguments. + pub fn with_scalars( + id: &'static str, + inputs: &[TensorIr], + outputs: &[TensorIr], + scalars: Vec, + ) -> Self { + Self { + id: id.to_owned(), + inputs: inputs.to_vec(), + outputs: outputs.to_vec(), + scalars, + } + } + + /// Cast the intermediate representation, and get the in and output tensors. + pub fn as_fixed( + &self, + ) -> (&[TensorIr; N_IN], &[TensorIr; N_OUT]) { + ( + self.inputs.as_slice().try_into().expect( + "Wrong number of inputs expected (expected {D}, is {}), check your implementation", + ), + self.outputs.as_slice().try_into().expect( + "Wrong number of outputs expected (expected {D}, is {}), check your implementation", + ), + ) + } + + fn inputs(&self) -> Box + '_> { + Box::new(self.inputs.iter()) + } + + fn outputs(&self) -> Box + '_> { + Box::new(self.outputs.iter()) + } + + fn visit_mut(&mut self, v: &mut impl IrVisitorMut) { + for t in self.inputs.iter_mut() { + v.visit_tensor_mut(t); + } + for t in self.outputs.iter_mut() { + v.visit_tensor_mut(t); + } + for s in self.scalars.iter_mut() { + v.visit_scalar_mut(s); + } + } +} + +/// Describe all tensor operations possible. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(clippy::large_enum_variant)] +pub enum OperationIr { + /// Basic operation on a float tensor. + BaseFloat(BaseOperationIr), + /// Basic operation on an int tensor. + BaseInt(BaseOperationIr), + /// Basic operation on a bool tensor. + BaseBool(BaseOperationIr), + /// Numeric operation on a float tensor. + NumericFloat(DType, NumericOperationIr), + /// Numeric operation on an int tensor. + NumericInt(DType, NumericOperationIr), + /// Operation specific to a bool tensor. + Bool(BoolOperationIr), + /// Operation specific to an int tensor. + Int(IntOperationIr), + /// Operation specific to a float tensor. + Float(DType, FloatOperationIr), + /// Module operation. + Module(ModuleOperationIr), + /// Initialize operation. + Init(InitOperationIr), + /// A custom operation. + Custom(CustomOpIr), + /// A tensor is dropped. + Drop(TensorIr), + /// Operation specific to a distributed tensor. + Distributed(DistributedOperationIr), + /// Activation function operation (relu, gelu, sigmoid, softmax, …). + Activation(ActivationOperationIr), +} + +/// Operation intermediate representation specific to a float tensor. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub enum FloatOperationIr { + /// Operation corresponding to [exp](burn_backend::ops::FloatTensorOps::float_exp). + Exp(UnaryOpIr), + /// Operation corresponding to [log](burn_backend::ops::FloatTensorOps::float_log). + Log(UnaryOpIr), + /// Operation corresponding to [log1p](burn_backend::ops::FloatTensorOps::float_log1p). + Log1p(UnaryOpIr), + /// Operation corresponding to [erf](burn_backend::ops::FloatTensorOps::float_erf). + Erf(UnaryOpIr), + /// Operation corresponding to [powf_scalar](burn_backend::ops::FloatTensorOps::float_powf_scalar). + PowfScalar(ScalarOpIr), + /// Operation corresponding to [sqrt](burn_backend::ops::FloatTensorOps::float_sqrt). + Sqrt(UnaryOpIr), + /// Operation corresponding to [cos](burn_backend::ops::FloatTensorOps::float_cos). + Cos(UnaryOpIr), + /// Operation corresponding to [cosh](burn_backend::ops::FloatTensorOps::float_cosh). + Cosh(UnaryOpIr), + /// Operation corresponding to [sin](burn_backend::ops::FloatTensorOps::float_sin). + Sin(UnaryOpIr), + /// Operation corresponding to [sin](burn_backend::ops::FloatTensorOps::float_sinh). + Sinh(UnaryOpIr), + /// Operation corresponding to [tan](burn_backend::ops::FloatTensorOps::float_tan). + Tan(UnaryOpIr), + /// Operation corresponding to [tanh](burn_backend::ops::FloatTensorOps::float_tanh). + Tanh(UnaryOpIr), + /// Operation corresponding to [acos](burn_backend::ops::FloatTensorOps::float_acos). + ArcCos(UnaryOpIr), + /// Operation corresponding to [acosh](burn_backend::ops::FloatTensorOps::float_acosh). + ArcCosh(UnaryOpIr), + /// Operation corresponding to [asin](burn_backend::ops::FloatTensorOps::float_asin). + ArcSin(UnaryOpIr), + /// Operation corresponding to [asinh](burn_backend::ops::FloatTensorOps::float_asinh). + ArcSinh(UnaryOpIr), + /// Operation corresponding to [atan](burn_backend::ops::FloatTensorOps::float_atan). + ArcTan(UnaryOpIr), + /// Operation corresponding to [atanh](burn_backend::ops::FloatTensorOps::float_atanh). + ArcTanh(UnaryOpIr), + /// Operation corresponding to [atan2](burn_backend::ops::FloatTensorOps::float_atan2). + ArcTan2(BinaryOpIr), + /// Operation corresponding to [round](burn_backend::ops::FloatTensorOps::float_round). + Round(UnaryOpIr), + /// Operation corresponding to [floor](burn_backend::ops::FloatTensorOps::float_floor). + Floor(UnaryOpIr), + /// Operation corresponding to [ceil](burn_backend::ops::FloatTensorOps::float_ceil). + Ceil(UnaryOpIr), + /// Operation corresponding to [trunc](burn_backend::ops::FloatTensorOps::float_trunc). + Trunc(UnaryOpIr), + /// Operation corresponding to [into_int](burn_backend::ops::FloatTensorOps::float_into_int). + IntoInt(CastOpIr), + /// Operation corresponding to [matmul](burn_backend::ops::FloatTensorOps::float_matmul). + Matmul(MatmulOpIr), + /// Operation corresponding to [cross](burn_backend::ops::FloatTensorOps::float_cross). + Cross(CrossOpIr), + /// Operation corresponding to [random](burn_backend::ops::FloatTensorOps::float_random). + Random(RandomOpIr), + /// Operation corresponding to [recip](burn_backend::ops::FloatTensorOps::float_recip). + Recip(UnaryOpIr), + /// Operation corresponding to [is_nan](burn_backend::ops::FloatTensorOps::float_is_nan). + IsNan(UnaryOpIr), + /// Operation corresponding to [is_nan](burn_backend::ops::FloatTensorOps::float_is_inf). + IsInf(UnaryOpIr), + /// Operation corresponding to [quantize](burn_backend::ops::QTensorOps::quantize). + Quantize(QuantizeOpIr), + /// Operation corresponding to [dequantize](burn_backend::ops::QTensorOps::dequantize). + Dequantize(DequantizeOpIr), + /// Operation corresponding to [grid_sample_2d](burn_backend::ops::FloatTensorOps::float_grid_sample_2d). + GridSample2d(GridSample2dOpIr), + /// Operation corresponding to [powf](burn_backend::ops::FloatTensorOps::float_powi). + Powf(BinaryOpIr), + /// Operation corresponding to [hypot](burn_backend::ops::FloatTensorOps::float_hypot). + Hypot(BinaryOpIr), +} + +/// Operation intermediate representation specific to module. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub enum ModuleOperationIr { + /// Operation corresponding to [embedding](burn_backend::ops::ModuleOps::embedding). + Embedding(EmbeddingOpIr), + /// Operation corresponding to [embedding_backward](burn_backend::ops::ModuleOps::embedding_backward). + EmbeddingBackward(EmbeddingBackwardOpIr), + /// Operation corresponding to [linear](burn_backend::ops::ModuleOps::linear). + Linear(LinearOpIr), + /// Operation corresponding to [linear_x_backward](burn_backend::ops::ModuleOps::linear_x_backward). + LinearXBackward(LinearXBackwardOpIr), + /// Operation corresponding to [linear_weight_backward](burn_backend::ops::ModuleOps::linear_weight_backward). + LinearWeightBackward(LinearWeightBackwardOpIr), + /// Operation corresponding to [linear_bias_backward](burn_backend::ops::ModuleOps::linear_bias_backward). + LinearBiasBackward(LinearBiasBackwardOpIr), + /// Operation corresponding to [conv1d](burn_backend::ops::ModuleOps::conv1d). + Conv1d(Conv1dOpIr), + /// Operation corresponding to [conv1d_x_backward](burn_backend::ops::ModuleOps::conv1d_x_backward). + Conv1dXBackward(Conv1dXBackwardOpIr), + /// Operation corresponding to [conv1d_weight_backward](burn_backend::ops::ModuleOps::conv1d_weight_backward). + Conv1dWeightBackward(Conv1dWeightBackwardOpIr), + /// Operation corresponding to [conv1d_bias_backward](burn_backend::ops::ModuleOps::conv1d_bias_backward). + Conv1dBiasBackward(Conv1dBiasBackwardOpIr), + /// Operation corresponding to [conv2d](burn_backend::ops::ModuleOps::conv2d). + Conv2d(Conv2dOpIr), + /// Operation corresponding to [conv2d_x_backward](burn_backend::ops::ModuleOps::conv2d_x_backward). + Conv2dXBackward(Conv2dXBackwardOpIr), + /// Operation corresponding to [conv2d_weight_backward](burn_backend::ops::ModuleOps::conv2d_weight_backward). + Conv2dWeightBackward(Conv2dWeightBackwardOpIr), + /// Operation corresponding to [conv2d_bias_backward](burn_backend::ops::ModuleOps::conv2d_bias_backward). + Conv2dBiasBackward(Conv2dBiasBackwardOpIr), + /// Operation corresponding to [conv3d](burn_backend::ops::ModuleOps::conv3d). + Conv3d(Conv3dOpIr), + /// Operation corresponding to [conv3d_x_backward](burn_backend::ops::ModuleOps::conv3d_x_backward). + Conv3dXBackward(Conv3dXBackwardOpIr), + /// Operation corresponding to [conv3d_weight_backward](burn_backend::ops::ModuleOps::conv3d_weight_backward). + Conv3dWeightBackward(Conv3dWeightBackwardOpIr), + /// Operation corresponding to [conv3d_bias_backward](burn_backend::ops::ModuleOps::conv3d_bias_backward). + Conv3dBiasBackward(Conv3dBiasBackwardOpIr), + /// Operation corresponding to [deform_conv2d](burn_backend::ops::ModuleOps::deform_conv2d) + DeformableConv2d(Box), + /// Operation corresponding to [deform_conv2d_backward](burn_backend::ops::ModuleOps::deform_conv2d_backward) + DeformableConv2dBackward(Box), + /// Operation corresponding to [conv transpose 1d](burn_backend::ops::ModuleOps::conv_transpose1d). + ConvTranspose1d(ConvTranspose1dOpIr), + /// Operation corresponding to [conv transpose 2d](burn_backend::ops::ModuleOps::conv_transpose2d). + ConvTranspose2d(ConvTranspose2dOpIr), + /// Operation corresponding to [conv transpose 3d](burn_backend::ops::ModuleOps::conv_transpose3d). + ConvTranspose3d(ConvTranspose3dOpIr), + /// Operation corresponding to [avg pool 1d](burn_backend::ops::ModuleOps::avg_pool1d). + AvgPool1d(AvgPool1dOpIr), + /// Operation corresponding to [avg pool 2d](burn_backend::ops::ModuleOps::avg_pool2d). + AvgPool2d(AvgPool2dOpIr), + /// Operation corresponding to + /// [avg pool 1d backward](burn_backend::ops::ModuleOps::avg_pool1d_backward). + AvgPool1dBackward(AvgPool1dBackwardOpIr), + /// Operation corresponding to + /// [avg pool 2d backward](burn_backend::ops::ModuleOps::avg_pool2d_backward). + AvgPool2dBackward(AvgPool2dBackwardOpIr), + /// Operation corresponding to + /// [adaptive avg pool 1d](burn_backend::ops::ModuleOps::adaptive_avg_pool1d). + AdaptiveAvgPool1d(AdaptiveAvgPool1dOpIr), + /// Operation corresponding to + /// [adaptive avg pool 2d](burn_backend::ops::ModuleOps::adaptive_avg_pool2d). + AdaptiveAvgPool2d(AdaptiveAvgPool2dOpIr), + /// Operation corresponding to + /// [adaptive avg pool 1d backward](burn_backend::ops::ModuleOps::adaptive_avg_pool1d_backward). + AdaptiveAvgPool1dBackward(AdaptiveAvgPool1dBackwardOpIr), + /// Operation corresponding to + /// [adaptive avg pool 2d backward](burn_backend::ops::ModuleOps::adaptive_avg_pool2d_backward). + AdaptiveAvgPool2dBackward(AdaptiveAvgPool2dBackwardOpIr), + /// Operation corresponding to + /// [max pool 1d](burn_backend::ops::ModuleOps::max_pool1d). + MaxPool1d(MaxPool1dOpIr), + /// Operation corresponding to + /// [max pool 1d with indices](burn_backend::ops::ModuleOps::max_pool1d_with_indices). + MaxPool1dWithIndices(MaxPool1dWithIndicesOpIr), + /// Operation corresponding to + /// [max pool 1d with indices backward](burn_backend::ops::ModuleOps::max_pool1d_with_indices_backward). + MaxPool1dWithIndicesBackward(MaxPool1dWithIndicesBackwardOpIr), + /// Operation corresponding to + /// [max pool 2d](burn_backend::ops::ModuleOps::max_pool1d). + MaxPool2d(MaxPool2dOpIr), + /// Operation corresponding to + /// [max pool 2d with indices](burn_backend::ops::ModuleOps::max_pool2d_with_indices). + MaxPool2dWithIndices(MaxPool2dWithIndicesOpIr), + /// Operation corresponding to + /// [max pool 2d with indices backward](burn_backend::ops::ModuleOps::max_pool2d_with_indices_backward). + MaxPool2dWithIndicesBackward(MaxPool2dWithIndicesBackwardOpIr), + /// Operation corresponding to [interpolate](burn_backend::ops::ModuleOps::interpolate). + Interpolate(InterpolateOpIr), + /// Operation corresponding to [interpolate backward](burn_backend::ops::ModuleOps::interpolate_backward). + InterpolateBackward(InterpolateBackwardOpIr), + /// Operation corresponding to [rfft](burn_backend::ops::ModuleOps::rfft) + Rfft(RfftOpIr), + /// Operation corresponding to [irfft](burn_backend::ops::ModuleOps::irfft) + IRfft(IRfftOpIr), + /// Operation corresponding to [attention](burn_backend::ops::ModuleOps::attention). + Attention(AttentionOpIr), + /// Operation corresponding to [ctc_loss](burn_backend::ops::ModuleOps::ctc_loss). + CtcLoss(CtcLossOpIr), + /// Operation corresponding to + /// [ctc_loss_backward](burn_backend::ops::ModuleOps::ctc_loss_backward). + CtcLossBackward(CtcLossBackwardOpIr), + /// Operation corresponding to [layer_norm](burn_backend::ops::ModuleOps::layer_norm). + LayerNorm(LayerNormOpIr), + /// Operation corresponding to [unfold4d](burn_backend::ops::ModuleOps::unfold4d). + Unfold4d(Unfold4dOpIr), + /// Operation corresponding to + /// [conv_transpose1d_weight_backward](burn_backend::ops::ModuleOps::conv_transpose1d_weight_backward). + ConvTranspose1dWeightBackward(ConvTranspose1dWeightBackwardOpIr), + /// Operation corresponding to + /// [conv_transpose1d_bias_backward](burn_backend::ops::ModuleOps::conv_transpose1d_bias_backward). + ConvTranspose1dBiasBackward(ConvTranspose1dBiasBackwardOpIr), + /// Operation corresponding to + /// [conv_transpose2d_weight_backward](burn_backend::ops::ModuleOps::conv_transpose2d_weight_backward). + ConvTranspose2dWeightBackward(ConvTranspose2dWeightBackwardOpIr), + /// Operation corresponding to + /// [conv_transpose2d_bias_backward](burn_backend::ops::ModuleOps::conv_transpose2d_bias_backward). + ConvTranspose2dBiasBackward(ConvTranspose2dBiasBackwardOpIr), + /// Operation corresponding to + /// [conv_transpose3d_weight_backward](burn_backend::ops::ModuleOps::conv_transpose3d_weight_backward). + ConvTranspose3dWeightBackward(ConvTranspose3dWeightBackwardOpIr), + /// Operation corresponding to + /// [conv_transpose3d_bias_backward](burn_backend::ops::ModuleOps::conv_transpose3d_bias_backward). + ConvTranspose3dBiasBackward(ConvTranspose3dBiasBackwardOpIr), +} + +/// Basic operations that can be done on any tensor type. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub enum BaseOperationIr { + /// Operation corresponding to: + /// + /// Float => [reshape](burn_backend::ops::FloatTensorOps::float_reshape). + /// Int => [reshape](burn_backend::ops::IntTensorOps::int_reshape). + /// Bool => [reshape](burn_backend::ops::BoolTensorOps::bool_reshape). + Reshape(ShapeOpIr), + + /// Operation corresponding to: + /// + /// Float => [swap_dims](burn_backend::ops::FloatTensorOps::float_swap_dims). + /// Int => [swap_dims](burn_backend::ops::IntTensorOps::int_swap_dims). + /// Bool => [swap_dims](burn_backend::ops::BoolTensorOps::bool_swap_dims). + SwapDims(SwapDimsOpIr), + + /// Operation corresponding to: + /// + /// Float => [permute](burn_backend::ops::FloatTensorOps::float_permute). + /// Int => [permute](burn_backend::ops::IntTensorOps::int_permute). + /// Bool => [permute](burn_backend::ops::BoolTensorOps::bool_permute). + Permute(PermuteOpIr), + + /// Operation corresponding to: + /// Float => [flip](burn_backend::ops::FloatTensorOps::float_flip). + /// Int => [flip](burn_backend::ops::IntTensorOps::int_flip). + /// Bool => [flip](burn_backend::ops::BoolTensorOps::bool_flip). + Flip(FlipOpIr), + + /// Operation corresponding to: + /// + /// Float => [expand](burn_backend::ops::FloatTensorOps::float_expand). + /// Int => [expand](burn_backend::ops::IntTensorOps::int_expand). + /// Bool => [expand](burn_backend::ops::BoolTensorOps::bool_expand). + Expand(ShapeOpIr), + + /// Unfold windows along an axis. + /// + Unfold(UnfoldOpIr), + + /// Operation corresponding to: + /// + /// Float => [slice](burn_backend::ops::FloatTensorOps::float_slice). + /// Int => [slice](burn_backend::ops::IntTensorOps::int_slice). + /// Bool => [slice](burn_backend::ops::BoolTensorOps::bool_slice). + Slice(SliceOpIr), + /// Operation corresponding to: + /// + /// Float => [slice assign](burn_backend::ops::FloatTensorOps::float_slice_assign). + /// Int => [slice assign](burn_backend::ops::IntTensorOps::int_slice_assign). + /// Bool => [slice assign](burn_backend::ops::BoolTensorOps::bool_slice_assign). + SliceAssign(SliceAssignOpIr), + /// Operation corresponding to: + /// + /// Float => [select](burn_backend::ops::FloatTensorOps::float_select). + /// Int => [select](burn_backend::ops::IntTensorOps::int_select). + /// Bool => [select](burn_backend::ops::BoolTensorOps::bool_select). + Select(SelectOpIr), + /// Operation corresponding to: + /// + /// Float => [select assign](burn_backend::ops::FloatTensorOps::float_select_add). + /// Int => [select assign](burn_backend::ops::IntTensorOps::int_select_add). + /// Bool => [select assign](burn_backend::ops::BoolTensorOps::bool_select_or). + SelectAssign(SelectAssignOpIr), + /// Operation corresponding to: + /// + /// Float => [mask where](burn_backend::ops::FloatTensorOps::float_mask_where). + /// Int => [mask where](burn_backend::ops::IntTensorOps::int_mask_where). + /// Bool => [mask where](burn_backend::ops::BoolTensorOps::bool_mask_where). + MaskWhere(MaskWhereOpIr), + /// Operation corresponding to: + /// + /// Float => [mask fill](burn_backend::ops::FloatTensorOps::float_mask_fill). + /// Int => [mask fill](burn_backend::ops::IntTensorOps::int_mask_fill). + /// Bool => [mask fill](burn_backend::ops::BoolTensorOps::bool_mask_fill). + MaskFill(MaskFillOpIr), + /// Operation corresponding to: + /// + /// Float => [gather](burn_backend::ops::FloatTensorOps::float_gather). + /// Int => [gather](burn_backend::ops::IntTensorOps::int_gather). + /// Bool => [gather](burn_backend::ops::BoolTensorOps::bool_gather). + Gather(GatherOpIr), + /// Operation corresponding to: + /// + /// Float => [scatter](burn_backend::ops::FloatTensorOps::float_scatter_add). + /// Int => [scatter](burn_backend::ops::IntTensorOps::int_scatter_add). + /// Bool => [scatter](burn_backend::ops::BoolTensorOps::bool_scatter_or). + Scatter(ScatterOpIr), + /// Multi-dimensional scatter operation. + ScatterNd(ScatterNdOpIr), + /// Multi-dimensional gather operation. + GatherNd(GatherNdOpIr), + /// Operation corresponding to: + /// + /// Float => [equal](burn_backend::ops::FloatTensorOps::float_equal). + /// Int => [equal](burn_backend::ops::IntTensorOps::int_equal). + /// Bool => [equal](burn_backend::ops::BoolTensorOps::bool_equal). + Equal(BinaryOpIr), + /// Operation corresponding to: + /// + /// Float => [equal elem](burn_backend::ops::FloatTensorOps::float_equal_elem). + /// Int => [equal elem](burn_backend::ops::IntTensorOps::int_equal_elem). + /// Bool => [equal elem](burn_backend::ops::BoolTensorOps::bool_equal_elem). + EqualElem(ScalarOpIr), + /// Operation corresponding to: + /// + /// Float => [repeat dim](burn_backend::ops::FloatTensorOps::float_repeat_dim). + /// Int => [repeat dim](burn_backend::ops::IntTensorOps::int_repeat_dim). + /// Bool => [repeat dim](burn_backend::ops::BoolTensorOps::bool_repeat_dim). + RepeatDim(RepeatDimOpIr), + /// Operation corresponding to: + /// + /// Float => [cat](burn_backend::ops::FloatTensorOps::float_cat). + /// Int => [cat](burn_backend::ops::IntTensorOps::int_cat). + /// Bool => [cat](burn_backend::ops::BoolTensorOps::bool_cat). + Cat(CatOpIr), + /// Cast operation, no direct operation and should be supported by fusion backend. + Cast(CastOpIr), + /// Operation corresponding to: + /// + /// Float => [empty](burn_backend::ops::FloatTensorOps::float_empty). + /// Int => [empty](burn_backend::ops::IntTensorOps::int_empty). + /// Bool => [empty](burn_backend::ops::BoolTensorOps::bool_empty). + Empty(CreationOpIr), + /// Operation corresponding to: + /// + /// Float => [ones](burn_backend::ops::FloatTensorOps::float_ones). + /// Int => [ones](burn_backend::ops::IntTensorOps::int_ones). + /// Bool => [ones](burn_backend::ops::BoolTensorOps::bool_ones). + Ones(CreationOpIr), + /// Operation corresponding to: + /// + /// Float => [zeros](burn_backend::ops::FloatTensorOps::float_zeros). + /// Int => [zeros](burn_backend::ops::IntTensorOps::int_zeros). + /// Bool => [zeros](burn_backend::ops::BoolTensorOps::bool_zeros). + Zeros(CreationOpIr), + /// Operation corresponding to: + /// + /// Float => [not_equal](burn_backend::ops::FloatTensorOps::float_not_equal). + /// Int => [not_equal](burn_backend::ops::IntTensorOps::int_not_equal). + /// Bool => [not_equal](burn_backend::ops::BoolTensorOps::bool_not_equal). + NotEqual(BinaryOpIr), + /// Operation corresponding to: + /// + /// Float => [not_equal_elem](burn_backend::ops::FloatTensorOps::float_not_equal_elem). + /// Int => [not_equal_elem](burn_backend::ops::IntTensorOps::int_not_equal_elem). + /// Bool => [not_equal_elem](burn_backend::ops::BoolTensorOps::bool_not_equal_elem). + NotEqualElem(ScalarOpIr), + /// Reduce-`all` over the input tensor. + /// + /// Float/Int input is treated as a non-zero check; bool input is the value itself. + /// Output is a single-element bool tensor. + All(ReduceOpIr), + /// Reduce-`any` over the input tensor. + Any(ReduceOpIr), + /// Reduce-`all` along a dim. + AllDim(ReduceDimOpIr), + /// Reduce-`any` along a dim. + AnyDim(ReduceDimOpIr), +} + +/// Numeric operations on int and float tensors. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub enum NumericOperationIr { + /// Operation corresponding to: + /// + /// Float => [add](burn_backend::ops::FloatTensorOps::float_add). + /// Int => [add](burn_backend::ops::IntTensorOps::int_add). + Add(BinaryOpIr), + /// Operation corresponding to: + /// + /// Float => [add scalar](burn_backend::ops::FloatTensorOps::float_add_scalar). + /// Int => [add scalar](burn_backend::ops::IntTensorOps::int_add_scalar). + AddScalar(ScalarOpIr), + /// Operation corresponding to: + /// + /// Float => [sub](burn_backend::ops::FloatTensorOps::float_sub). + /// Int => [sub](burn_backend::ops::IntTensorOps::int_sub). + Sub(BinaryOpIr), + /// Operation corresponding to: + /// + /// Float => [sub scalar](burn_backend::ops::FloatTensorOps::float_sub_scalar). + /// Int => [sub scalar](burn_backend::ops::IntTensorOps::int_sub_scalar). + SubScalar(ScalarOpIr), + /// Operation corresponding to: + /// + /// Float => [div](burn_backend::ops::FloatTensorOps::float_div). + /// Int => [div](burn_backend::ops::IntTensorOps::int_div). + Div(BinaryOpIr), + /// Operation corresponding to: + /// + /// Float => [div scalar](burn_backend::ops::FloatTensorOps::float_div_scalar). + /// Int => [div scalar](burn_backend::ops::IntTensorOps::int_div_scalar). + DivScalar(ScalarOpIr), + /// Operation corresponding to: + /// + /// Float => [rem](burn_backend::ops::FloatTensorOps::float_remainder). + /// Int => [rem](burn_backend::ops::IntTensorOps::int_remainder). + Rem(BinaryOpIr), + /// Operation corresponding to: + /// + /// Float => [rem scalar](burn_backend::ops::FloatTensorOps::float_remainder_scalar). + /// Int => [rem scalar](burn_backend::ops::IntTensorOps::int_remainder_scalar). + RemScalar(ScalarOpIr), + /// Operation corresponding to: + /// + /// Float => [mul](burn_backend::ops::FloatTensorOps::float_mul). + /// Int => [mul](burn_backend::ops::IntTensorOps::int_mul). + Mul(BinaryOpIr), + /// Operation corresponding to: + /// + /// Float => [mul scalar](burn_backend::ops::FloatTensorOps::float_mul_scalar). + /// Int => [mul scalar](burn_backend::ops::IntTensorOps::int_mul_scalar). + MulScalar(ScalarOpIr), + /// Operation corresponding to: + /// + /// Float => [abs](burn_backend::ops::FloatTensorOps::float_abs). + /// Int => [abs](burn_backend::ops::IntTensorOps::int_abs). + Abs(UnaryOpIr), + /// Operation corresponding to: + /// + /// Float => [full](burn_backend::ops::FloatTensorOps::float_full). + /// Int => [full](burn_backend::ops::IntTensorOps::int_full). + Full(FullOpIr), + /// Operation corresponding to: + /// + /// Float => [mean dim](burn_backend::ops::FloatTensorOps::float_mean_dim). + /// Int => [mean dim](burn_backend::ops::IntTensorOps::int_mean_dim). + MeanDim(ReduceDimOpIr), + /// Operation corresponding to: + /// + /// Float => [mean](burn_backend::ops::FloatTensorOps::float_mean). + /// Int => [mean](burn_backend::ops::IntTensorOps::int_mean). + Mean(ReduceOpIr), + /// Operation corresponding to: + /// + /// Float => [sum](burn_backend::ops::FloatTensorOps::float_sum). + /// Int => [sum](burn_backend::ops::IntTensorOps::int_sum). + Sum(ReduceOpIr), + /// Operation corresponding to: + /// + /// Float => [sum dim](burn_backend::ops::FloatTensorOps::float_sum_dim). + /// Int => [sum dim](burn_backend::ops::IntTensorOps::int_sum_dim). + SumDim(ReduceDimOpIr), + /// Operation corresponding to: + /// + /// Float => [prod](burn_backend::ops::FloatTensorOps::float_prod). + /// Int => [prod](burn_backend::ops::IntTensorOps::int_prod). + Prod(ReduceOpIr), + /// Operation corresponding to: + /// + /// Float => [prod dim](burn_backend::ops::FloatTensorOps::float_prod_dim). + /// Int => [prod dim](burn_backend::ops::IntTensorOps::int_prod_dim). + ProdDim(ReduceDimOpIr), + /// Operation corresponding to: + /// + /// Float => [greater](burn_backend::ops::FloatTensorOps::float_greater). + /// Int => [greater](burn_backend::ops::IntTensorOps::int_greater). + Greater(BinaryOpIr), + /// Operation corresponding to: + /// + /// Float => [greater elem](burn_backend::ops::FloatTensorOps::float_greater_elem). + /// Int => [greater elem](burn_backend::ops::IntTensorOps::int_greater_elem). + GreaterElem(ScalarOpIr), + /// Operation corresponding to: + /// + /// Float => [greater equal](burn_backend::ops::FloatTensorOps::float_greater_elem). + /// Int => [greater elem](burn_backend::ops::IntTensorOps::int_greater_elem). + GreaterEqual(BinaryOpIr), + /// Operation corresponding to: + /// + /// Float => [greater equal elem](burn_backend::ops::FloatTensorOps::float_greater_equal_elem). + /// Int => [greater equal elem](burn_backend::ops::IntTensorOps::int_greater_equal_elem). + GreaterEqualElem(ScalarOpIr), + /// Operation corresponding to: + /// + /// Float => [lower](burn_backend::ops::FloatTensorOps::float_lower). + /// Int => [lower](burn_backend::ops::IntTensorOps::int_lower). + Lower(BinaryOpIr), + /// Operation corresponding to: + /// + /// Float => [lower elem](burn_backend::ops::FloatTensorOps::float_lower_elem). + /// Int => [lower elem](burn_backend::ops::IntTensorOps::int_lower_elem). + LowerElem(ScalarOpIr), + /// Operation corresponding to: + /// + /// Float => [lower equal](burn_backend::ops::FloatTensorOps::float_lower_equal). + /// Int => [lower equal](burn_backend::ops::IntTensorOps::int_lower_equal). + LowerEqual(BinaryOpIr), + /// Operation corresponding to: + /// + /// Float => [lower equal elem](burn_backend::ops::FloatTensorOps::float_lower_equal_elem). + /// Int => [lower equal elem](burn_backend::ops::IntTensorOps::int_lower_equal_elem). + LowerEqualElem(ScalarOpIr), + /// Operation corresponding to: + /// + /// Float => [argmax](burn_backend::ops::FloatTensorOps::float_argmax). + /// Int => [argmax](burn_backend::ops::IntTensorOps::int_argmax). + ArgMax(ReduceDimOpIr), + /// Operation corresponding to: + /// + /// Float => [argtopk](burn_backend::ops::FloatTensorOps::float_argtopk). + /// Int => [argtopk](burn_backend::ops::IntTensorOps::int_argtopk). + ArgTopK(ReduceDimOpIr), + /// Operation corresponding to: + /// + /// Float => [topk](burn_backend::ops::FloatTensorOps::float_topk). + /// Int => [topk](burn_backend::ops::IntTensorOps::int_topk). + TopK(ReduceDimOpIr), + /// Operation corresponding to: + /// + /// Float => [argmin](burn_backend::ops::FloatTensorOps::float_argmin). + /// Int => [argmin](burn_backend::ops::IntTensorOps::int_argmin). + ArgMin(ReduceDimOpIr), + /// Operation corresponding to: + /// + /// Float => [max](burn_backend::ops::FloatTensorOps::float_max). + /// Int => [max](burn_backend::ops::IntTensorOps::int_max). + Max(ReduceOpIr), + /// Operation corresponding to: + /// + /// Float => [max dim with indices](burn_backend::ops::FloatTensorOps::float_max_dim_with_indices). + /// Int => [max dim with indices](burn_backend::ops::IntTensorOps::int_max_dim_with_indices). + MaxDimWithIndices(ReduceDimWithIndicesOpIr), + /// Operation corresponding to: + /// + /// Float => [min dim with indices](burn_backend::ops::FloatTensorOps::float_min_dim_with_indices). + /// Int => [min dim with indices](burn_backend::ops::IntTensorOps::int_min_dim_with_indices). + MinDimWithIndices(ReduceDimWithIndicesOpIr), + /// Operation corresponding to: + /// + /// Float => [min](burn_backend::ops::FloatTensorOps::float_min). + /// Int => [min](burn_backend::ops::IntTensorOps::int_min). + Min(ReduceOpIr), + /// Operation corresponding to: + /// + /// Float => [max dim](burn_backend::ops::FloatTensorOps::float_max_dim). + /// Int => [max dim](burn_backend::ops::IntTensorOps::int_max_dim). + MaxDim(ReduceDimOpIr), + /// Operation corresponding to: + /// + /// Float => [min dim](burn_backend::ops::FloatTensorOps::float_min_dim). + /// Int => [min dim](burn_backend::ops::IntTensorOps::int_min_dim). + MinDim(ReduceDimOpIr), + /// Operation corresponding to: + /// + /// Float => [max_abs](burn_backend::ops::FloatTensorOps::float_max_abs). + /// Int => [max_abs](burn_backend::ops::IntTensorOps::int_max_abs). + MaxAbs(ReduceOpIr), + /// Operation corresponding to: + /// + /// Float => [max_abs dim](burn_backend::ops::FloatTensorOps::float_max_abs_dim). + /// Int => [max_abs dim](burn_backend::ops::IntTensorOps::int_max_abs_dim). + MaxAbsDim(ReduceDimOpIr), + /// Operation corresponding to: + /// + /// Float => [clamp](burn_backend::ops::FloatTensorOps::float_clamp). + /// Int => [clamp](burn_backend::ops::IntTensorOps::int_clamp). + Clamp(ClampOpIr), + /// Operation corresponding to: + /// + /// Int => [random](burn_backend::ops::IntTensorOps::int_random). + IntRandom(RandomOpIr), + /// Operation corresponding to: + /// + /// Float => [powf](burn_backend::ops::FloatTensorOps::float_powi). + /// Int => [powf](burn_backend::ops::IntTensorOps::int_powi). + Powi(BinaryOpIr), + /// Operation corresponding to: + /// + /// Float => [powi_scalar](burn_backend::ops::FloatTensorOps::float_powi_scalar). + /// Int => [powi_scalar](burn_backend::ops::IntTensorOps::int_powi_scalar). + PowiScalar(ScalarOpIr), + /// Operation corresponding to: + /// + /// Float => [cumsum](burn_backend::ops::FloatTensorOps::float_cumsum). + /// Int => [cumsum](burn_backend::ops::IntTensorOps::int_cumsum). + CumSum(DimOpIr), + /// Operation corresponding to: + /// + /// Float => [cumprod](burn_backend::ops::FloatTensorOps::float_cumprod). + /// Int => [cumprod](burn_backend::ops::IntTensorOps::int_cumprod). + CumProd(DimOpIr), + /// Operation corresponding to: + /// + /// Float => [cummin](burn_backend::ops::FloatTensorOps::float_cummin). + /// Int => [cummin](burn_backend::ops::IntTensorOps::int_cummin). + CumMin(DimOpIr), + /// Operation corresponding to: + /// + /// Float => [cummax](burn_backend::ops::FloatTensorOps::float_cummax). + /// Int => [cummax](burn_backend::ops::IntTensorOps::int_cummax). + CumMax(DimOpIr), + /// Operation corresponding to: + /// + /// Float => [neg](burn_backend::ops::FloatTensorOps::float_neg). + /// Int => [neg](burn_backend::ops::IntTensorOps::int_neg). + Neg(UnaryOpIr), + /// Operation corresponding to: + /// + /// Float => [sign](burn_backend::ops::FloatTensorOps::float_sign). + /// Int => [sign](burn_backend::ops::IntTensorOps::int_sign). + Sign(UnaryOpIr), + /// Operation corresponding to: + /// + /// Float => [clamp_min](burn_backend::ops::FloatTensorOps::float_clamp_min). + /// Int => [clamp_min](burn_backend::ops::IntTensorOps::int_clamp_min). + ClampMin(ScalarOpIr), + /// Operation corresponding to: + /// + /// Float => [clamp_max](burn_backend::ops::FloatTensorOps::float_clamp_max). + /// Int => [clamp_max](burn_backend::ops::IntTensorOps::int_clamp_max). + ClampMax(ScalarOpIr), + /// Sort along a dim. Output shares the input's shape and dtype. + Sort(SortOpIr), + /// Sort along a dim, also returning the source indices. + SortWithIndices(SortWithIndicesOpIr), + /// Sort along a dim and return only the indices (i.e., the argsort). + /// + /// Shares [`SortOpIr`] with [`Sort`](Self::Sort); only the output dtype differs. + ArgSort(SortOpIr), +} + +/// Operation intermediate representation specific to an int tensor. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub enum IntOperationIr { + /// Operation corresponding to [into float](burn_backend::ops::IntTensorOps::int_into_float). + IntoFloat(CastOpIr), + /// Operation corresponding to: + /// + /// Int => [bitwise and](burn_backend::ops::IntTensorOps::bitwise_and). + BitwiseAnd(BinaryOpIr), + /// Operation corresponding to: + /// + /// Int => [bitwise and scalar](burn_backend::ops::IntTensorOps::bitwise_and_scalar). + BitwiseAndScalar(ScalarOpIr), + /// Operation corresponding to: + /// + /// Int => [bitwise or](burn_backend::ops::IntTensorOps::bitwise_or). + BitwiseOr(BinaryOpIr), + /// Operation corresponding to: + /// + /// Int => [bitwise or scalar](burn_backend::ops::IntTensorOps::bitwise_or_scalar). + BitwiseOrScalar(ScalarOpIr), + /// Operation corresponding to: + /// + /// Int => [bitwise xor](burn_backend::ops::IntTensorOps::bitwise_xor). + BitwiseXor(BinaryOpIr), + /// Operation corresponding to: + /// + /// Int => [bitwise xor scalar](burn_backend::ops::IntTensorOps::bitwise_xor_scalar). + BitwiseXorScalar(ScalarOpIr), + /// Operation corresponding to: + /// + /// Int => [bitwise not](burn_backend::ops::IntTensorOps::bitwise_not). + BitwiseNot(UnaryOpIr), + /// Operation corresponding to: + /// + /// Int => [bitwise left shift](burn_backend::ops::IntTensorOps::bitwise_left_shift). + BitwiseLeftShift(BinaryOpIr), + /// Operation corresponding to: + /// + /// Int => [bitwise left shift scalar](burn_backend::ops::IntTensorOps::bitwise_left_shift_scalar). + BitwiseLeftShiftScalar(ScalarOpIr), + /// Operation corresponding to: + /// + /// Int => [bitwise right shift](burn_backend::ops::IntTensorOps::bitwise_right_shift). + BitwiseRightShift(BinaryOpIr), + /// Operation corresponding to: + /// + /// Int => [bitwise right shift scalar](burn_backend::ops::IntTensorOps::bitwise_right_shift_scalar). + BitwiseRightShiftScalar(ScalarOpIr), + /// Operation corresponding to [matmul](burn_backend::ops::IntTensorOps::int_matmul). + Matmul(MatmulOpIr), +} + +/// Operation intermediate representation specific to a bool tensor. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub enum BoolOperationIr { + /// Operation corresponding to [into float](burn_backend::ops::BoolTensorOps::bool_into_float). + IntoFloat(CastOpIr), + /// Operation corresponding to [into int](burn_backend::ops::BoolTensorOps::bool_into_int). + IntoInt(CastOpIr), + /// Operation corresponding to [not](burn_backend::ops::BoolTensorOps::bool_not). + Not(UnaryOpIr), + /// Operation corresponding to [and](burn_backend::ops::BoolTensorOps::bool_and). + And(BinaryOpIr), + /// Operation corresponding to [or](burn_backend::ops::BoolTensorOps::bool_or). + Or(BinaryOpIr), + /// Operation corresponding to [xor](burn_backend::ops::BoolTensorOps::bool_xor). + Xor(BinaryOpIr), +} + +/// Operations that can be done on distributed tensors. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(clippy::large_enum_variant)] +pub enum DistributedOperationIr { + /// Operation corresponding to: + /// [all_reduce](burn_backend::distributed::DistributedOps::all_reduce). + AllReduce(AllReduceOpIr), + /// Resolve the pending collective operations on the executing device. Corresponds to + /// [sync_collective](burn_backend::distributed::DistributedOps::sync_collective). + /// + /// Fire-and-forget and payload-free: it syncs whichever device the interpreter is bound to. + /// Modeled as an operation (not a side-channel call) so it travels the normal op stream + /// alongside [`AllReduce`](DistributedOperationIr::AllReduce) and is ordered against it. + SyncCollective, +} + +/// Swap dim operation intermediate representation. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct SwapDimsOpIr { + /// Input tensor intermediate representation. + pub input: TensorIr, + /// Output tensor intermediate representation. + pub out: TensorIr, + /// The first dim to swap. + pub dim1: usize, + /// The second dim to swap. + pub dim2: usize, +} + +/// Permute operation intermediate representation. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct PermuteOpIr { + /// Input tensor intermediate representation. + pub input: TensorIr, + /// Output tensor intermediate representation. + pub out: TensorIr, + /// The new order of the dimensions. + pub axes: Vec, +} + +/// Shape operation intermediate representation. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct ShapeOpIr { + /// Input tensor intermediate representation. + pub input: TensorIr, + /// Output tensor intermediate representation with the new shape. + pub out: TensorIr, +} + +/// Unfold operation intermediate representation. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct UnfoldOpIr { + /// Input tensor intermediate representation. + pub input: TensorIr, + /// Output tensor intermediate representation. + pub out: TensorIr, + + /// The selected dim. + pub dim: usize, + /// The window size. + pub size: usize, + /// The window step along dim. + pub step: usize, +} + +/// Flip operation intermediate representation. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct FlipOpIr { + /// Input tensor intermediate representation. + pub input: TensorIr, + /// Output tensor intermediate representation. + pub out: TensorIr, + /// The dimensions to flip. + pub axes: Vec, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct RandomOpIr { + pub out: TensorIr, + pub distribution: Distribution, +} + +/// Creation operation intermediate representation. +/// As opposed to [InitOperationIr], creation operations are lazy initialized. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct CreationOpIr { + /// Output tensor intermediate representation. + pub out: TensorIr, +} + +/// Full operation intermediate representation. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct FullOpIr { + /// Output tensor intermediate representation. + pub out: TensorIr, + /// Fill value. + pub value: ScalarIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +/// Declares a tensor has been initialized. +/// +/// It is necessary to register for proper orphan detection and avoid memory leak. +pub struct InitOperationIr { + /// The initialized tensor. + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct BinaryOpIr { + pub lhs: TensorIr, + pub rhs: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct MatmulOpIr { + pub lhs: TensorIr, + pub rhs: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct CrossOpIr { + pub lhs: TensorIr, + pub rhs: TensorIr, + pub out: TensorIr, + pub dim: usize, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct UnaryOpIr { + pub input: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct ScalarOpIr { + pub lhs: TensorIr, + // TODO: Make that an enum with `Value` and `Id` variants for relative/global + // conversion. + pub rhs: ScalarIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Hash)] +#[allow(missing_docs)] +pub struct ReduceOpIr { + pub input: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Hash)] +#[allow(missing_docs)] +pub struct ReduceDimOpIr { + pub input: TensorIr, + pub out: TensorIr, + pub axis: usize, + pub accumulator_len: usize, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct CastOpIr { + pub input: TensorIr, + pub out: TensorIr, +} + +/// IR for operations that operate along a dimension without reducing it. +/// Unlike `ReduceDimOpIr`, the output shape is the same as the input shape. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, Hash)] +#[allow(missing_docs)] +pub struct DimOpIr { + pub input: TensorIr, + pub out: TensorIr, + pub axis: usize, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct GatherOpIr { + pub tensor: TensorIr, + pub dim: usize, + pub indices: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct ScatterOpIr { + pub tensor: TensorIr, + pub dim: usize, + pub indices: TensorIr, + pub value: TensorIr, + pub update: IndexingUpdateOp, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct ScatterNdOpIr { + pub data: TensorIr, + pub indices: TensorIr, + pub values: TensorIr, + pub reduction: IndexingUpdateOp, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct GatherNdOpIr { + pub data: TensorIr, + pub indices: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct SelectOpIr { + pub tensor: TensorIr, + pub dim: usize, + pub indices: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct SelectAssignOpIr { + pub tensor: TensorIr, + pub dim: usize, + pub indices: TensorIr, + pub value: TensorIr, + pub update: IndexingUpdateOp, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct SliceOpIr { + pub tensor: TensorIr, + pub ranges: Vec, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct SliceAssignOpIr { + pub tensor: TensorIr, + pub ranges: Vec, + pub value: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct MaskWhereOpIr { + pub tensor: TensorIr, + pub mask: TensorIr, + pub value: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct MaskFillOpIr { + pub tensor: TensorIr, + pub mask: TensorIr, + pub value: ScalarIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct ClampOpIr { + pub tensor: TensorIr, + pub min: ScalarIr, + pub max: ScalarIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct RepeatDimOpIr { + pub tensor: TensorIr, + pub dim: usize, + pub times: usize, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct CatOpIr { + pub tensors: Vec, + pub dim: usize, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct AllReduceOpIr { + pub tensor: TensorIr, + pub out: TensorIr, + /// How to reduce the values across the participating devices. + pub op: burn_backend::distributed::ReduceOperation, + /// The devices participating in the collective operation. + pub device_ids: Vec, +} + +/// Serializable representation of a [device id](burn_backend::DeviceId). +/// +/// The intermediate representation is part of the wire protocol (e.g. the remote backend), so it +/// cannot store `burn_backend::DeviceId` directly since that type is not serializable. +#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)] +pub struct DeviceIdIr { + /// Identifies the type of the device. + pub type_id: u16, + /// Identifies the device number. + pub index_id: u16, +} + +impl From for DeviceIdIr { + fn from(value: burn_backend::DeviceId) -> Self { + Self { + type_id: value.type_id, + index_id: value.index_id, + } + } +} + +impl From for burn_backend::DeviceId { + fn from(value: DeviceIdIr) -> Self { + burn_backend::DeviceId::new(value.type_id, value.index_id) + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct ReduceDimWithIndicesOpIr { + pub tensor: TensorIr, + pub dim: usize, + pub out: TensorIr, + pub out_indices: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct EmbeddingOpIr { + pub weights: TensorIr, + pub indices: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct EmbeddingBackwardOpIr { + pub weights: TensorIr, + pub out_grad: TensorIr, + pub indices: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct LinearOpIr { + pub x: TensorIr, + pub weight: TensorIr, + pub bias: Option, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct LinearXBackwardOpIr { + pub weight: TensorIr, + pub output_grad: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct LinearWeightBackwardOpIr { + pub x: TensorIr, + pub output_grad: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct LinearBiasBackwardOpIr { + pub output_grad: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct Conv1dOpIr { + pub x: TensorIr, + pub weight: TensorIr, + pub bias: Option, + pub options: Conv1dOptionsIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct Conv1dXBackwardOpIr { + pub x: TensorIr, + pub weight: TensorIr, + pub output_grad: TensorIr, + pub options: Conv1dOptionsIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct Conv1dWeightBackwardOpIr { + pub x: TensorIr, + pub weight: TensorIr, + pub output_grad: TensorIr, + pub options: Conv1dOptionsIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct Conv1dBiasBackwardOpIr { + pub x: TensorIr, + pub bias: TensorIr, + pub output_grad: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct Conv2dOpIr { + pub x: TensorIr, + pub weight: TensorIr, + pub bias: Option, + pub options: Conv2dOptionsIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct Conv2dXBackwardOpIr { + pub x: TensorIr, + pub weight: TensorIr, + pub output_grad: TensorIr, + pub options: Conv2dOptionsIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct Conv2dWeightBackwardOpIr { + pub x: TensorIr, + pub weight: TensorIr, + pub output_grad: TensorIr, + pub options: Conv2dOptionsIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct Conv2dBiasBackwardOpIr { + pub x: TensorIr, + pub bias: TensorIr, + pub output_grad: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct DeformConv2dOpIr { + pub x: TensorIr, + pub offset: TensorIr, + pub weight: TensorIr, + pub mask: Option, + pub bias: Option, + pub options: DeformableConv2dOptionsIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct DeformConv2dBackwardOpIr { + pub x: TensorIr, + pub offset: TensorIr, + pub weight: TensorIr, + pub mask: Option, + pub bias: Option, + pub out_grad: TensorIr, + pub options: DeformableConv2dOptionsIr, + pub input_grad: TensorIr, + pub offset_grad: TensorIr, + pub weight_grad: TensorIr, + pub mask_grad: Option, + pub bias_grad: Option, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct Conv3dOpIr { + pub x: TensorIr, + pub weight: TensorIr, + pub bias: Option, + pub options: Conv3dOptionsIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct Conv3dXBackwardOpIr { + pub x: TensorIr, + pub weight: TensorIr, + pub output_grad: TensorIr, + pub options: Conv3dOptionsIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct Conv3dWeightBackwardOpIr { + pub x: TensorIr, + pub weight: TensorIr, + pub output_grad: TensorIr, + pub options: Conv3dOptionsIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct Conv3dBiasBackwardOpIr { + pub x: TensorIr, + pub bias: TensorIr, + pub output_grad: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct ConvTranspose1dOpIr { + pub x: TensorIr, + pub weight: TensorIr, + pub bias: Option, + pub options: ConvTranspose1dOptionsIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct ConvTranspose2dOpIr { + pub x: TensorIr, + pub weight: TensorIr, + pub bias: Option, + pub options: ConvTranspose2dOptionsIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct ConvTranspose3dOpIr { + pub x: TensorIr, + pub weight: TensorIr, + pub bias: Option, + pub options: ConvTranspose3dOptionsIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct Conv1dOptionsIr { + pub stride: [usize; 1], + pub padding: [usize; 1], + pub dilation: [usize; 1], + pub groups: usize, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct Conv2dOptionsIr { + pub stride: [usize; 2], + pub padding: [usize; 2], + pub dilation: [usize; 2], + pub groups: usize, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct DeformableConv2dOptionsIr { + pub stride: [usize; 2], + pub padding: [usize; 2], + pub dilation: [usize; 2], + pub weight_groups: usize, + pub offset_groups: usize, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct Conv3dOptionsIr { + pub stride: [usize; 3], + pub padding: [usize; 3], + pub dilation: [usize; 3], + pub groups: usize, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct ConvTranspose1dOptionsIr { + pub stride: [usize; 1], + pub padding: [usize; 1], + pub padding_out: [usize; 1], + pub dilation: [usize; 1], + pub groups: usize, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct ConvTranspose2dOptionsIr { + pub stride: [usize; 2], + pub padding: [usize; 2], + pub padding_out: [usize; 2], + pub dilation: [usize; 2], + pub groups: usize, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct ConvTranspose3dOptionsIr { + pub stride: [usize; 3], + pub padding: [usize; 3], + pub padding_out: [usize; 3], + pub dilation: [usize; 3], + pub groups: usize, +} + +/// Quantization parameters intermediate representation. +#[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize, Deserialize)] +pub struct QuantizationParametersIr { + /// The scaling factor. + pub scales: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct QuantizeOpIr { + pub tensor: TensorIr, + pub qparams: QuantizationParametersIr, + pub scheme: QuantScheme, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct DequantizeOpIr { + pub input: TensorIr, + pub out: TensorIr, +} + +impl From> for Conv1dOptionsIr { + fn from(value: ConvOptions<1>) -> Self { + Self { + stride: value.stride, + padding: value.padding, + dilation: value.dilation, + groups: value.groups, + } + } +} + +impl From> for Conv2dOptionsIr { + fn from(value: ConvOptions<2>) -> Self { + Self { + stride: value.stride, + padding: value.padding, + dilation: value.dilation, + groups: value.groups, + } + } +} + +impl From> for Conv3dOptionsIr { + fn from(value: ConvOptions<3>) -> Self { + Self { + stride: value.stride, + padding: value.padding, + dilation: value.dilation, + groups: value.groups, + } + } +} + +impl From> for DeformableConv2dOptionsIr { + fn from(value: DeformConvOptions<2>) -> Self { + Self { + stride: value.stride, + padding: value.padding, + dilation: value.dilation, + weight_groups: value.weight_groups, + offset_groups: value.offset_groups, + } + } +} + +impl From> for ConvTranspose1dOptionsIr { + fn from(value: ConvTransposeOptions<1>) -> Self { + Self { + stride: value.stride, + padding: value.padding, + padding_out: value.padding_out, + dilation: value.dilation, + groups: value.groups, + } + } +} + +impl From> for ConvTranspose2dOptionsIr { + fn from(value: ConvTransposeOptions<2>) -> Self { + Self { + stride: value.stride, + padding: value.padding, + padding_out: value.padding_out, + dilation: value.dilation, + groups: value.groups, + } + } +} + +impl From> for ConvTranspose3dOptionsIr { + fn from(value: ConvTransposeOptions<3>) -> Self { + Self { + stride: value.stride, + padding: value.padding, + padding_out: value.padding_out, + dilation: value.dilation, + groups: value.groups, + } + } +} + +impl From for ConvOptions<1> { + fn from(val: Conv1dOptionsIr) -> Self { + ConvOptions { + stride: val.stride, + padding: val.padding, + dilation: val.dilation, + groups: val.groups, + } + } +} + +impl From for ConvOptions<2> { + fn from(val: Conv2dOptionsIr) -> Self { + ConvOptions { + stride: val.stride, + padding: val.padding, + dilation: val.dilation, + groups: val.groups, + } + } +} + +impl From for ConvOptions<3> { + fn from(val: Conv3dOptionsIr) -> Self { + ConvOptions { + stride: val.stride, + padding: val.padding, + dilation: val.dilation, + groups: val.groups, + } + } +} + +impl From for DeformConvOptions<2> { + fn from(value: DeformableConv2dOptionsIr) -> Self { + DeformConvOptions { + stride: value.stride, + padding: value.padding, + dilation: value.dilation, + weight_groups: value.weight_groups, + offset_groups: value.offset_groups, + } + } +} + +impl From for ConvTransposeOptions<1> { + fn from(val: ConvTranspose1dOptionsIr) -> Self { + ConvTransposeOptions { + stride: val.stride, + padding: val.padding, + padding_out: val.padding_out, + dilation: val.dilation, + groups: val.groups, + } + } +} + +impl From for ConvTransposeOptions<2> { + fn from(val: ConvTranspose2dOptionsIr) -> Self { + ConvTransposeOptions { + stride: val.stride, + padding: val.padding, + padding_out: val.padding_out, + dilation: val.dilation, + groups: val.groups, + } + } +} + +impl From for ConvTransposeOptions<3> { + fn from(val: ConvTranspose3dOptionsIr) -> Self { + ConvTransposeOptions { + stride: val.stride, + padding: val.padding, + padding_out: val.padding_out, + dilation: val.dilation, + groups: val.groups, + } + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct AvgPool1dOpIr { + pub x: TensorIr, + pub kernel_size: usize, + pub stride: usize, + pub padding: usize, + pub count_include_pad: bool, + pub ceil_mode: bool, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct AvgPool2dOpIr { + pub x: TensorIr, + pub kernel_size: [usize; 2], + pub stride: [usize; 2], + pub padding: [usize; 2], + pub count_include_pad: bool, + pub ceil_mode: bool, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct AvgPool1dBackwardOpIr { + pub x: TensorIr, + pub grad: TensorIr, + pub kernel_size: usize, + pub stride: usize, + pub padding: usize, + pub count_include_pad: bool, + pub ceil_mode: bool, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct AvgPool2dBackwardOpIr { + pub x: TensorIr, + pub grad: TensorIr, + pub kernel_size: [usize; 2], + pub stride: [usize; 2], + pub padding: [usize; 2], + pub count_include_pad: bool, + pub ceil_mode: bool, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct AdaptiveAvgPool1dOpIr { + pub x: TensorIr, + pub output_size: usize, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct AdaptiveAvgPool2dOpIr { + pub x: TensorIr, + pub output_size: [usize; 2], + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct AdaptiveAvgPool1dBackwardOpIr { + pub x: TensorIr, + pub grad: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct AdaptiveAvgPool2dBackwardOpIr { + pub x: TensorIr, + pub grad: TensorIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct MaxPool1dOpIr { + pub x: TensorIr, + pub kernel_size: usize, + pub stride: usize, + pub padding: usize, + pub dilation: usize, + pub ceil_mode: bool, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct MaxPool1dWithIndicesOpIr { + pub x: TensorIr, + pub kernel_size: usize, + pub stride: usize, + pub padding: usize, + pub dilation: usize, + pub ceil_mode: bool, + pub out: TensorIr, + pub out_indices: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct MaxPool1dWithIndicesBackwardOpIr { + pub x: TensorIr, + pub grad: TensorIr, + pub indices: TensorIr, + pub kernel_size: usize, + pub stride: usize, + pub padding: usize, + pub dilation: usize, + pub ceil_mode: bool, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct MaxPool2dOpIr { + pub x: TensorIr, + pub kernel_size: [usize; 2], + pub stride: [usize; 2], + pub padding: [usize; 2], + pub dilation: [usize; 2], + pub ceil_mode: bool, + pub out: TensorIr, +} + +#[allow(missing_docs)] +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct MaxPool2dWithIndicesOpIr { + pub x: TensorIr, + pub kernel_size: [usize; 2], + pub stride: [usize; 2], + pub padding: [usize; 2], + pub dilation: [usize; 2], + pub ceil_mode: bool, + pub out: TensorIr, + pub out_indices: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct MaxPool2dWithIndicesBackwardOpIr { + pub x: TensorIr, + pub grad: TensorIr, + pub indices: TensorIr, + pub kernel_size: [usize; 2], + pub stride: [usize; 2], + pub padding: [usize; 2], + pub dilation: [usize; 2], + pub ceil_mode: bool, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub enum InterpolateModeIr { + Nearest, + NearestExact, + Bilinear, + Bicubic, + Lanczos3, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct InterpolateOptionsIr { + pub mode: InterpolateModeIr, + pub align_corners: bool, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct InterpolateOpIr { + pub x: TensorIr, + pub output_size: [usize; 2], + pub options: InterpolateOptionsIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct RfftOpIr { + pub signal: TensorIr, + pub dim: usize, + pub n: Option, + pub out_re: TensorIr, + pub out_im: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct IRfftOpIr { + pub input_re: TensorIr, + pub input_im: TensorIr, + pub dim: usize, + pub n: Option, + pub out_signal: TensorIr, +} + +#[allow(missing_docs)] +impl RfftOpIr { + pub fn create(signal: TensorIr, dim: usize, n: Option, mut new_id: F) -> Self + where + F: FnMut() -> crate::TensorId, + { + // `n` is required to be a power of two at the public API boundary, so + // the output has `n / 2 + 1` bins (matching scipy/torch for pow2 n). + let mut shape = signal.shape.clone(); + let fft_len = n.unwrap_or(shape[dim]); + shape[dim] = fft_len / 2 + 1; + let dtype = signal.dtype; + + Self { + signal, + dim, + n, + out_re: TensorIr::uninit(new_id(), shape.clone(), dtype), + out_im: TensorIr::uninit(new_id(), shape, dtype), + } + } +} + +#[allow(missing_docs)] +impl IRfftOpIr { + pub fn create( + input_re: TensorIr, + input_im: TensorIr, + dim: usize, + n: Option, + mut new_id: F, + ) -> Self + where + F: FnMut() -> crate::TensorId, + { + debug_assert!( + input_re.shape[dim] >= 1, + "IRfftOpIr: input spectrum dimension must be >= 1" + ); + debug_assert!( + !matches!(n, Some(0)), + "IRfftOpIr: n must be >= 1 when specified" + ); + let mut shape = input_re.shape.clone(); + shape[dim] = n.unwrap_or((shape[dim] - 1) * 2); + let dtype = input_re.dtype; + + Self { + input_re, + input_im, + dim, + n, + out_signal: TensorIr::uninit(new_id(), shape, dtype), + } + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct AttentionOptionsIr { + pub scale: Option, + pub softcap: Option, + pub is_causal: bool, +} + +impl From for AttentionModuleOptions { + fn from(ir: AttentionOptionsIr) -> Self { + AttentionModuleOptions { + scale: ir.scale.map(|s| s.elem()), + softcap: ir.softcap.map(|s| s.elem()), + is_causal: ir.is_causal, + } + } +} + +impl From for AttentionOptionsIr { + fn from(ir: AttentionModuleOptions) -> Self { + AttentionOptionsIr { + scale: ir.scale.map(ScalarIr::Float), + softcap: ir.softcap.map(ScalarIr::Float), + is_causal: ir.is_causal, + } + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct AttentionOpIr { + pub query: TensorIr, + pub key: TensorIr, + pub value: TensorIr, + pub mask: Option, + pub attn_bias: Option, + pub options: AttentionOptionsIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct CtcLossOpIr { + pub log_probs: TensorIr, + pub targets: TensorIr, + pub input_lengths: TensorIr, + pub target_lengths: TensorIr, + pub blank: usize, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct CtcLossBackwardOpIr { + pub log_probs: TensorIr, + pub targets: TensorIr, + pub input_lengths: TensorIr, + pub target_lengths: TensorIr, + pub grad_loss: TensorIr, + pub blank: usize, + pub out: TensorIr, +} + +impl From for InterpolateMode { + fn from(val: InterpolateModeIr) -> Self { + match val { + InterpolateModeIr::Nearest => Self::Nearest, + InterpolateModeIr::NearestExact => Self::NearestExact, + InterpolateModeIr::Bilinear => Self::Bilinear, + InterpolateModeIr::Bicubic => Self::Bicubic, + InterpolateModeIr::Lanczos3 => Self::Lanczos3, + } + } +} + +impl From for InterpolateOptions { + fn from(val: InterpolateOptionsIr) -> Self { + Self::new(val.mode.into()).with_align_corners(val.align_corners) + } +} + +impl From for InterpolateModeIr { + fn from(val: InterpolateMode) -> Self { + match val { + InterpolateMode::Nearest => Self::Nearest, + InterpolateMode::NearestExact => Self::Nearest, + InterpolateMode::Bilinear => Self::Bilinear, + InterpolateMode::Bicubic => Self::Bicubic, + InterpolateMode::Lanczos3 => Self::Lanczos3, + } + } +} + +impl From for InterpolateOptionsIr { + fn from(val: InterpolateOptions) -> Self { + Self { + mode: val.mode.into(), + align_corners: val.align_corners, + } + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct InterpolateBackwardOpIr { + pub x: TensorIr, + pub grad: TensorIr, + pub output_size: [usize; 2], + pub options: InterpolateOptionsIr, + pub out: TensorIr, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub enum GridSamplePaddingModeIr { + Zeros, + Border, + Reflection, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct GridSampleOptionsIr { + pub mode: InterpolateModeIr, + pub padding_mode: GridSamplePaddingModeIr, + pub align_corners: bool, +} + +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub struct GridSample2dOpIr { + pub tensor: TensorIr, + pub grid: TensorIr, + pub options: GridSampleOptionsIr, + pub out: TensorIr, +} + +impl From for GridSamplePaddingMode { + fn from(val: GridSamplePaddingModeIr) -> Self { + match val { + GridSamplePaddingModeIr::Zeros => Self::Zeros, + GridSamplePaddingModeIr::Border => Self::Border, + GridSamplePaddingModeIr::Reflection => Self::Reflection, + } + } +} + +impl From for GridSamplePaddingModeIr { + fn from(val: GridSamplePaddingMode) -> Self { + match val { + GridSamplePaddingMode::Zeros => Self::Zeros, + GridSamplePaddingMode::Border => Self::Border, + GridSamplePaddingMode::Reflection => Self::Reflection, + } + } +} + +impl From for GridSampleOptions { + fn from(val: GridSampleOptionsIr) -> Self { + Self { + mode: val.mode.into(), + padding_mode: val.padding_mode.into(), + align_corners: val.align_corners, + } + } +} + +impl From for GridSampleOptionsIr { + fn from(val: GridSampleOptions) -> Self { + Self { + mode: val.mode.into(), + padding_mode: val.padding_mode.into(), + align_corners: val.align_corners, + } + } +} + +impl OperationIr { + /// Get all input [tensors](TensorIr) involved with the current operation. + pub fn inputs(&self) -> impl Iterator { + match self { + OperationIr::BaseFloat(repr) => repr.inputs(), + OperationIr::BaseInt(repr) => repr.inputs(), + OperationIr::BaseBool(repr) => repr.inputs(), + OperationIr::NumericFloat(_dtype, repr) => repr.inputs(), + OperationIr::NumericInt(_dtype, repr) => repr.inputs(), + OperationIr::Bool(repr) => repr.inputs(), + OperationIr::Int(repr) => repr.inputs(), + OperationIr::Float(_dtype, repr) => repr.inputs(), + OperationIr::Module(repr) => repr.inputs(), + OperationIr::Init(repr) => repr.inputs(), + OperationIr::Custom(repr) => repr.inputs(), + OperationIr::Drop(repr) => Box::new([repr].into_iter()), + OperationIr::Distributed(repr) => repr.inputs(), + OperationIr::Activation(repr) => repr.inputs(), + } + } + + /// Get all output [tensors](TensorIr) involved with the current operation. + pub fn outputs(&self) -> impl Iterator { + match self { + OperationIr::BaseFloat(repr) => repr.outputs(), + OperationIr::BaseInt(repr) => repr.outputs(), + OperationIr::BaseBool(repr) => repr.outputs(), + OperationIr::NumericFloat(_dtype, repr) => repr.outputs(), + OperationIr::NumericInt(_dtype, repr) => repr.outputs(), + OperationIr::Bool(repr) => repr.outputs(), + OperationIr::Int(repr) => repr.outputs(), + OperationIr::Float(_dtype, repr) => repr.outputs(), + OperationIr::Module(repr) => repr.outputs(), + OperationIr::Init(repr) => repr.outputs(), + OperationIr::Custom(repr) => repr.outputs(), + OperationIr::Drop(_repr) => Box::new([].into_iter()), + OperationIr::Distributed(repr) => repr.outputs(), + OperationIr::Activation(repr) => repr.outputs(), + } + } + + /// Get all [tensor](TensorIr) involved with the current operation. + pub fn nodes(&self) -> Vec<&TensorIr> { + self.inputs().chain(self.outputs()).collect() + } + + /// Set the given nodes that are [read write](super::TensorStatus::ReadWrite) to + /// [read only](super::TensorStatus::ReadOnly) in the current operation. + /// + /// Returns the tensor that were updated with their original representation. + pub fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec { + match self { + OperationIr::BaseFloat(repr) => repr.mark_read_only(nodes), + OperationIr::BaseInt(repr) => repr.mark_read_only(nodes), + OperationIr::BaseBool(repr) => repr.mark_read_only(nodes), + OperationIr::NumericFloat(_dtype, repr) => repr.mark_read_only(nodes), + OperationIr::NumericInt(_dtype, repr) => repr.mark_read_only(nodes), + OperationIr::Bool(repr) => repr.mark_read_only(nodes), + OperationIr::Int(repr) => repr.mark_read_only(nodes), + OperationIr::Float(_dtype, repr) => repr.mark_read_only(nodes), + OperationIr::Module(repr) => repr.mark_read_only(nodes), + OperationIr::Init(_) => Vec::new(), + OperationIr::Drop(repr) => { + let mut output = Vec::new(); + repr.mark_read_only(nodes, &mut output); + output + } + OperationIr::Custom(repr) => { + let mut output = Vec::new(); + + for input in repr.inputs.iter_mut() { + input.mark_read_only(nodes, &mut output); + } + + output + } + OperationIr::Distributed(repr) => repr.mark_read_only(nodes), + OperationIr::Activation(repr) => repr.mark_read_only(nodes), + } + } + + /// Visit every mutable component of this operation in place. + /// + /// For each variant, visits its [tensors](TensorIr) (in the same order as + /// [`inputs`](Self::inputs) then [`outputs`](Self::outputs) combined), then its + /// [scalars](ScalarIr), then its slice ranges. Only `Slice`/`SliceAssign` (under the `Base*` + /// variants) carry ranges; every other category visits no ranges. + pub fn visit_mut(&mut self, v: &mut impl IrVisitorMut) { + match self { + OperationIr::BaseFloat(repr) => repr.visit_mut(v), + OperationIr::BaseInt(repr) => repr.visit_mut(v), + OperationIr::BaseBool(repr) => repr.visit_mut(v), + OperationIr::NumericFloat(_dtype, repr) => repr.visit_mut(v), + OperationIr::NumericInt(_dtype, repr) => repr.visit_mut(v), + OperationIr::Bool(repr) => repr.visit_mut(v), + OperationIr::Int(repr) => repr.visit_mut(v), + OperationIr::Float(_dtype, repr) => repr.visit_mut(v), + OperationIr::Module(repr) => repr.visit_mut(v), + OperationIr::Init(repr) => repr.visit_mut(v), + OperationIr::Custom(repr) => repr.visit_mut(v), + OperationIr::Drop(repr) => v.visit_tensor_mut(repr), + OperationIr::Distributed(repr) => repr.visit_mut(v), + OperationIr::Activation(repr) => repr.visit_mut(v), + } + } +} + +impl BaseOperationIr { + fn inputs(&self) -> Box + '_> { + match self { + BaseOperationIr::Reshape(repr) => Box::new([&repr.input].into_iter()), + BaseOperationIr::SwapDims(repr) => Box::new([&repr.input].into_iter()), + BaseOperationIr::Permute(repr) => Box::new([&repr.input].into_iter()), + BaseOperationIr::Expand(repr) => Box::new([&repr.input].into_iter()), + BaseOperationIr::Flip(repr) => Box::new([&repr.input].into_iter()), + BaseOperationIr::Slice(repr) => Box::new([&repr.tensor].into_iter()), + BaseOperationIr::SliceAssign(repr) => Box::new([&repr.tensor, &repr.value].into_iter()), + BaseOperationIr::Gather(repr) => Box::new([&repr.tensor, &repr.indices].into_iter()), + BaseOperationIr::Scatter(repr) => { + Box::new([&repr.tensor, &repr.indices, &repr.value].into_iter()) + } + BaseOperationIr::ScatterNd(repr) => { + Box::new([&repr.data, &repr.indices, &repr.values].into_iter()) + } + BaseOperationIr::GatherNd(repr) => Box::new([&repr.data, &repr.indices].into_iter()), + BaseOperationIr::Select(repr) => Box::new([&repr.tensor, &repr.indices].into_iter()), + BaseOperationIr::SelectAssign(repr) => { + Box::new([&repr.tensor, &repr.indices, &repr.value].into_iter()) + } + BaseOperationIr::MaskWhere(repr) => { + Box::new([&repr.tensor, &repr.mask, &repr.value].into_iter()) + } + BaseOperationIr::MaskFill(repr) => Box::new([&repr.tensor, &repr.mask].into_iter()), + BaseOperationIr::Equal(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + BaseOperationIr::EqualElem(repr) => Box::new([&repr.lhs].into_iter()), + BaseOperationIr::RepeatDim(repr) => Box::new([&repr.tensor].into_iter()), + BaseOperationIr::Cat(repr) => Box::new(repr.tensors.iter()), + BaseOperationIr::Cast(repr) => Box::new([&repr.input].into_iter()), + BaseOperationIr::Unfold(repr) => Box::new([&repr.input].into_iter()), + BaseOperationIr::Empty(_repr) => Box::new([].into_iter()), + BaseOperationIr::Ones(_repr) => Box::new([].into_iter()), + BaseOperationIr::Zeros(_repr) => Box::new([].into_iter()), + BaseOperationIr::NotEqual(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + BaseOperationIr::NotEqualElem(repr) => Box::new([&repr.lhs].into_iter()), + BaseOperationIr::All(repr) => Box::new([&repr.input].into_iter()), + BaseOperationIr::Any(repr) => Box::new([&repr.input].into_iter()), + BaseOperationIr::AllDim(repr) => Box::new([&repr.input].into_iter()), + BaseOperationIr::AnyDim(repr) => Box::new([&repr.input].into_iter()), + } + } + + fn outputs(&self) -> Box + '_> { + match self { + BaseOperationIr::Reshape(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::SwapDims(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::Permute(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::Expand(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::Flip(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::Slice(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::SliceAssign(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::Gather(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::Scatter(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::ScatterNd(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::GatherNd(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::Select(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::SelectAssign(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::MaskWhere(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::MaskFill(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::Equal(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::EqualElem(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::RepeatDim(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::Cat(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::Cast(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::Unfold(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::Empty(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::Ones(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::Zeros(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::NotEqual(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::NotEqualElem(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::All(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::Any(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::AllDim(repr) => Box::new([&repr.out].into_iter()), + BaseOperationIr::AnyDim(repr) => Box::new([&repr.out].into_iter()), + } + } + + fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec { + let mut output = Vec::new(); + + match self { + BaseOperationIr::Reshape(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + BaseOperationIr::SwapDims(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + BaseOperationIr::Permute(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + + BaseOperationIr::Expand(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + + BaseOperationIr::Flip(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + BaseOperationIr::Slice(repr) => { + repr.tensor.mark_read_only(nodes, &mut output); + } + BaseOperationIr::SliceAssign(repr) => { + repr.tensor.mark_read_only(nodes, &mut output); + repr.value.mark_read_only(nodes, &mut output); + } + BaseOperationIr::Gather(repr) => { + repr.tensor.mark_read_only(nodes, &mut output); + repr.indices.mark_read_only(nodes, &mut output); + } + BaseOperationIr::Scatter(repr) => { + repr.tensor.mark_read_only(nodes, &mut output); + repr.indices.mark_read_only(nodes, &mut output); + repr.value.mark_read_only(nodes, &mut output); + } + BaseOperationIr::ScatterNd(repr) => { + repr.data.mark_read_only(nodes, &mut output); + repr.indices.mark_read_only(nodes, &mut output); + repr.values.mark_read_only(nodes, &mut output); + } + BaseOperationIr::GatherNd(repr) => { + repr.data.mark_read_only(nodes, &mut output); + repr.indices.mark_read_only(nodes, &mut output); + } + BaseOperationIr::Select(repr) => { + repr.tensor.mark_read_only(nodes, &mut output); + repr.indices.mark_read_only(nodes, &mut output); + } + BaseOperationIr::SelectAssign(repr) => { + repr.tensor.mark_read_only(nodes, &mut output); + repr.indices.mark_read_only(nodes, &mut output); + repr.value.mark_read_only(nodes, &mut output); + } + BaseOperationIr::MaskWhere(repr) => { + repr.tensor.mark_read_only(nodes, &mut output); + repr.mask.mark_read_only(nodes, &mut output); + repr.value.mark_read_only(nodes, &mut output); + } + BaseOperationIr::MaskFill(repr) => { + repr.tensor.mark_read_only(nodes, &mut output); + repr.mask.mark_read_only(nodes, &mut output); + } + BaseOperationIr::Equal(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + BaseOperationIr::EqualElem(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + BaseOperationIr::RepeatDim(repr) => { + repr.tensor.mark_read_only(nodes, &mut output); + } + BaseOperationIr::Cat(repr) => { + for t in repr.tensors.iter_mut() { + t.mark_read_only(nodes, &mut output); + } + } + BaseOperationIr::Cast(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + BaseOperationIr::Unfold(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + BaseOperationIr::Empty(_) => {} + BaseOperationIr::Zeros(_) => {} + BaseOperationIr::Ones(_) => {} + BaseOperationIr::NotEqual(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + BaseOperationIr::NotEqualElem(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + BaseOperationIr::All(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + BaseOperationIr::Any(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + BaseOperationIr::AllDim(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + BaseOperationIr::AnyDim(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + }; + + output + } + + fn visit_mut(&mut self, v: &mut impl IrVisitorMut) { + match self { + BaseOperationIr::Reshape(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::SwapDims(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::Permute(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::Expand(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::Flip(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::Slice(repr) => { + v.visit_tensor_mut(&mut repr.tensor); + v.visit_tensor_mut(&mut repr.out); + repr.ranges.iter_mut().for_each(|r| v.visit_range_mut(r)); + } + BaseOperationIr::SliceAssign(repr) => { + v.visit_tensor_mut(&mut repr.tensor); + v.visit_tensor_mut(&mut repr.value); + v.visit_tensor_mut(&mut repr.out); + repr.ranges.iter_mut().for_each(|r| v.visit_range_mut(r)); + } + BaseOperationIr::Gather(repr) => { + v.visit_tensor_mut(&mut repr.tensor); + v.visit_tensor_mut(&mut repr.indices); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::Scatter(repr) => { + v.visit_tensor_mut(&mut repr.tensor); + v.visit_tensor_mut(&mut repr.indices); + v.visit_tensor_mut(&mut repr.value); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::ScatterNd(repr) => { + v.visit_tensor_mut(&mut repr.data); + v.visit_tensor_mut(&mut repr.indices); + v.visit_tensor_mut(&mut repr.values); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::GatherNd(repr) => { + v.visit_tensor_mut(&mut repr.data); + v.visit_tensor_mut(&mut repr.indices); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::Select(repr) => { + v.visit_tensor_mut(&mut repr.tensor); + v.visit_tensor_mut(&mut repr.indices); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::SelectAssign(repr) => { + v.visit_tensor_mut(&mut repr.tensor); + v.visit_tensor_mut(&mut repr.indices); + v.visit_tensor_mut(&mut repr.value); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::MaskWhere(repr) => { + v.visit_tensor_mut(&mut repr.tensor); + v.visit_tensor_mut(&mut repr.mask); + v.visit_tensor_mut(&mut repr.value); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::MaskFill(repr) => { + v.visit_tensor_mut(&mut repr.tensor); + v.visit_tensor_mut(&mut repr.mask); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.value); + } + BaseOperationIr::Equal(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::EqualElem(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + BaseOperationIr::RepeatDim(repr) => { + v.visit_tensor_mut(&mut repr.tensor); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::Cat(repr) => { + for t in repr.tensors.iter_mut() { + v.visit_tensor_mut(t); + } + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::Cast(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::Unfold(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::Empty(repr) => { + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::Ones(repr) => { + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::Zeros(repr) => { + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::NotEqual(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::NotEqualElem(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + BaseOperationIr::All(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::Any(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::AllDim(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + BaseOperationIr::AnyDim(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + } + } +} + +impl NumericOperationIr { + fn inputs(&self) -> Box + '_> { + match self { + NumericOperationIr::Add(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + NumericOperationIr::AddScalar(repr) => Box::new([&repr.lhs].into_iter()), + NumericOperationIr::Sub(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + NumericOperationIr::SubScalar(repr) => Box::new([&repr.lhs].into_iter()), + NumericOperationIr::Mul(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + NumericOperationIr::MulScalar(repr) => Box::new([&repr.lhs].into_iter()), + NumericOperationIr::Div(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + NumericOperationIr::DivScalar(repr) => Box::new([&repr.lhs].into_iter()), + NumericOperationIr::Rem(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + NumericOperationIr::RemScalar(repr) => Box::new([&repr.lhs].into_iter()), + NumericOperationIr::GreaterElem(repr) => Box::new([&repr.lhs].into_iter()), + NumericOperationIr::GreaterEqualElem(repr) => Box::new([&repr.lhs].into_iter()), + NumericOperationIr::LowerElem(repr) => Box::new([&repr.lhs].into_iter()), + NumericOperationIr::LowerEqualElem(repr) => Box::new([&repr.lhs].into_iter()), + NumericOperationIr::Greater(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + NumericOperationIr::GreaterEqual(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + NumericOperationIr::Lower(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + NumericOperationIr::LowerEqual(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + NumericOperationIr::ArgMax(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::ArgTopK(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::TopK(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::ArgMin(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::Clamp(repr) => Box::new([&repr.tensor].into_iter()), + NumericOperationIr::Abs(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::Full(_repr) => Box::new([].into_iter()), + NumericOperationIr::MeanDim(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::Mean(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::Sum(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::SumDim(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::Prod(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::ProdDim(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::Max(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::MaxDimWithIndices(repr) => Box::new([&repr.tensor].into_iter()), + NumericOperationIr::MinDimWithIndices(repr) => Box::new([&repr.tensor].into_iter()), + NumericOperationIr::Min(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::MaxDim(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::MinDim(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::MaxAbs(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::MaxAbsDim(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::IntRandom(_repr) => Box::new([].into_iter()), + NumericOperationIr::Powi(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + NumericOperationIr::PowiScalar(repr) => Box::new([&repr.lhs].into_iter()), + NumericOperationIr::CumMin(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::CumMax(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::CumProd(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::CumSum(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::Neg(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::Sign(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::ClampMin(repr) => Box::new([&repr.lhs].into_iter()), + NumericOperationIr::ClampMax(repr) => Box::new([&repr.lhs].into_iter()), + NumericOperationIr::Sort(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::SortWithIndices(repr) => Box::new([&repr.input].into_iter()), + NumericOperationIr::ArgSort(repr) => Box::new([&repr.input].into_iter()), + } + } + + fn outputs(&self) -> Box + '_> { + match self { + NumericOperationIr::Add(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::AddScalar(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Sub(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::SubScalar(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Mul(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::MulScalar(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Div(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::DivScalar(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Rem(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::RemScalar(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::GreaterElem(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::GreaterEqualElem(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::LowerElem(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::LowerEqualElem(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Greater(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::GreaterEqual(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Lower(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::LowerEqual(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::ArgMax(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::ArgTopK(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::TopK(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::ArgMin(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Clamp(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Abs(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Full(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::MeanDim(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Mean(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Sum(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::SumDim(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Prod(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::ProdDim(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Max(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::MaxDimWithIndices(repr) => { + Box::new([&repr.out, &repr.out_indices].into_iter()) + } + NumericOperationIr::MinDimWithIndices(repr) => { + Box::new([&repr.out, &repr.out_indices].into_iter()) + } + NumericOperationIr::Min(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::MaxDim(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::MinDim(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::MaxAbs(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::MaxAbsDim(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::IntRandom(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Powi(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::PowiScalar(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::CumMin(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::CumMax(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::CumProd(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::CumSum(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Neg(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Sign(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::ClampMin(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::ClampMax(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::Sort(repr) => Box::new([&repr.out].into_iter()), + NumericOperationIr::SortWithIndices(repr) => { + Box::new([&repr.out, &repr.out_indices].into_iter()) + } + NumericOperationIr::ArgSort(repr) => Box::new([&repr.out].into_iter()), + } + } + fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec { + let mut output = Vec::new(); + + match self { + NumericOperationIr::Add(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::AddScalar(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Sub(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::SubScalar(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Mul(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::MulScalar(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Div(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::DivScalar(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Rem(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::RemScalar(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::GreaterElem(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::GreaterEqualElem(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::LowerElem(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::LowerEqualElem(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Greater(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::GreaterEqual(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Lower(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::LowerEqual(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::ArgMax(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::ArgTopK(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::TopK(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::ArgMin(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Clamp(repr) => { + repr.tensor.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Abs(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Full(_) => {} + NumericOperationIr::MeanDim(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Mean(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Sum(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::SumDim(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Prod(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::ProdDim(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Max(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::MaxDimWithIndices(repr) => { + repr.tensor.mark_read_only(nodes, &mut output); + } + NumericOperationIr::MinDimWithIndices(repr) => { + repr.tensor.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Min(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::MaxDim(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::MinDim(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::MaxAbs(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::MaxAbsDim(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::IntRandom(_) => {} + NumericOperationIr::Powi(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::PowiScalar(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::CumSum(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::CumProd(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::CumMin(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::CumMax(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Neg(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Sign(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::ClampMin(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::ClampMax(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + NumericOperationIr::Sort(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::SortWithIndices(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + NumericOperationIr::ArgSort(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + }; + + output + } + + fn visit_mut(&mut self, v: &mut impl IrVisitorMut) { + match self { + NumericOperationIr::Add(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::AddScalar(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + NumericOperationIr::Sub(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::SubScalar(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + NumericOperationIr::Mul(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::MulScalar(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + NumericOperationIr::Div(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::DivScalar(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + NumericOperationIr::Rem(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::RemScalar(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + NumericOperationIr::GreaterElem(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + NumericOperationIr::GreaterEqualElem(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + NumericOperationIr::LowerElem(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + NumericOperationIr::LowerEqualElem(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + NumericOperationIr::Greater(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::GreaterEqual(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::Lower(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::LowerEqual(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::ArgMax(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::ArgTopK(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::TopK(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::ArgMin(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::Clamp(repr) => { + v.visit_tensor_mut(&mut repr.tensor); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.min); + v.visit_scalar_mut(&mut repr.max); + } + NumericOperationIr::Abs(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::Full(repr) => { + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.value); + } + NumericOperationIr::MeanDim(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::Mean(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::Sum(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::SumDim(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::Prod(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::ProdDim(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::Max(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::MaxDimWithIndices(repr) => { + v.visit_tensor_mut(&mut repr.tensor); + v.visit_tensor_mut(&mut repr.out); + v.visit_tensor_mut(&mut repr.out_indices); + } + NumericOperationIr::MinDimWithIndices(repr) => { + v.visit_tensor_mut(&mut repr.tensor); + v.visit_tensor_mut(&mut repr.out); + v.visit_tensor_mut(&mut repr.out_indices); + } + NumericOperationIr::Min(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::MaxDim(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::MinDim(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::MaxAbs(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::MaxAbsDim(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::IntRandom(repr) => { + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::Powi(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::PowiScalar(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + NumericOperationIr::CumMin(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::CumMax(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::CumProd(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::CumSum(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::Neg(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::Sign(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::ClampMin(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + NumericOperationIr::ClampMax(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + NumericOperationIr::Sort(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + NumericOperationIr::SortWithIndices(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + v.visit_tensor_mut(&mut repr.out_indices); + } + NumericOperationIr::ArgSort(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + } + } +} + +impl FloatOperationIr { + fn inputs(&self) -> Box + '_> { + match self { + FloatOperationIr::Matmul(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + FloatOperationIr::Cross(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + FloatOperationIr::Random(_repr) => Box::new([].into_iter()), + FloatOperationIr::Exp(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::Log(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::Log1p(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::Erf(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::Recip(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::PowfScalar(repr) => Box::new([&repr.lhs].into_iter()), + FloatOperationIr::Sqrt(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::Cos(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::Sin(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::Tanh(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::Round(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::Floor(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::Ceil(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::Trunc(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::IntoInt(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::Quantize(repr) => { + Box::new([&repr.tensor, &repr.qparams.scales].into_iter()) + } + FloatOperationIr::Dequantize(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::IsNan(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::IsInf(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::GridSample2d(repr) => { + Box::new([&repr.tensor, &repr.grid].into_iter()) + } + FloatOperationIr::Tan(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::Cosh(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::Sinh(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::ArcCos(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::ArcCosh(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::ArcSin(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::ArcSinh(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::ArcTan(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::ArcTanh(repr) => Box::new([&repr.input].into_iter()), + FloatOperationIr::ArcTan2(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + FloatOperationIr::Powf(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + FloatOperationIr::Hypot(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + } + } + fn outputs(&self) -> Box + '_> { + match self { + FloatOperationIr::Matmul(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Cross(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Random(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Exp(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Log(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Log1p(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Erf(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Recip(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::PowfScalar(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Sqrt(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Cos(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Sin(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Tanh(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Round(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Floor(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Ceil(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Trunc(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::IntoInt(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Quantize(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Dequantize(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::IsNan(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::IsInf(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::GridSample2d(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Tan(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Cosh(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Sinh(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::ArcCos(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::ArcCosh(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::ArcSin(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::ArcSinh(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::ArcTan(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::ArcTanh(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::ArcTan2(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Powf(repr) => Box::new([&repr.out].into_iter()), + FloatOperationIr::Hypot(repr) => Box::new([&repr.out].into_iter()), + } + } + + fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec { + let mut output = Vec::new(); + + match self { + FloatOperationIr::Matmul(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Cross(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Random(_) => {} + FloatOperationIr::Exp(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Log(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Log1p(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Erf(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Recip(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::PowfScalar(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Sqrt(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Cos(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Sin(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Tanh(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Round(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Floor(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Ceil(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Trunc(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Quantize(repr) => { + repr.tensor.mark_read_only(nodes, &mut output); + repr.qparams.scales.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Dequantize(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::IntoInt(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::IsNan(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::IsInf(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + FloatOperationIr::GridSample2d(repr) => { + repr.tensor.mark_read_only(nodes, &mut output); + repr.grid.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Tan(repr) => repr.input.mark_read_only(nodes, &mut output), + FloatOperationIr::Cosh(repr) => repr.input.mark_read_only(nodes, &mut output), + FloatOperationIr::Sinh(repr) => repr.input.mark_read_only(nodes, &mut output), + FloatOperationIr::ArcCos(repr) => repr.input.mark_read_only(nodes, &mut output), + FloatOperationIr::ArcCosh(repr) => repr.input.mark_read_only(nodes, &mut output), + FloatOperationIr::ArcSin(repr) => repr.input.mark_read_only(nodes, &mut output), + FloatOperationIr::ArcSinh(repr) => repr.input.mark_read_only(nodes, &mut output), + FloatOperationIr::ArcTan(repr) => repr.input.mark_read_only(nodes, &mut output), + FloatOperationIr::ArcTanh(repr) => repr.input.mark_read_only(nodes, &mut output), + FloatOperationIr::ArcTan2(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Powf(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + FloatOperationIr::Hypot(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + }; + + output + } + + fn visit_mut(&mut self, v: &mut impl IrVisitorMut) { + match self { + FloatOperationIr::Matmul(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Cross(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Random(repr) => { + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Exp(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Log(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Log1p(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Erf(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Recip(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::PowfScalar(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + FloatOperationIr::Sqrt(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Cos(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Sin(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Tanh(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Round(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Floor(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Ceil(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Trunc(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::IntoInt(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Quantize(repr) => { + v.visit_tensor_mut(&mut repr.tensor); + v.visit_tensor_mut(&mut repr.qparams.scales); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Dequantize(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::IsNan(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::IsInf(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::GridSample2d(repr) => { + v.visit_tensor_mut(&mut repr.tensor); + v.visit_tensor_mut(&mut repr.grid); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Tan(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Cosh(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Sinh(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::ArcCos(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::ArcCosh(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::ArcSin(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::ArcSinh(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::ArcTan(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::ArcTanh(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::ArcTan2(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Powf(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + FloatOperationIr::Hypot(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + } + } +} + +impl IntOperationIr { + fn inputs(&self) -> Box + '_> { + match self { + IntOperationIr::Matmul(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + IntOperationIr::IntoFloat(repr) => Box::new([&repr.input].into_iter()), + IntOperationIr::BitwiseAnd(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + IntOperationIr::BitwiseAndScalar(repr) => Box::new([&repr.lhs].into_iter()), + IntOperationIr::BitwiseOr(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + IntOperationIr::BitwiseOrScalar(repr) => Box::new([&repr.lhs].into_iter()), + IntOperationIr::BitwiseXor(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + IntOperationIr::BitwiseXorScalar(repr) => Box::new([&repr.lhs].into_iter()), + IntOperationIr::BitwiseNot(repr) => Box::new([&repr.input].into_iter()), + IntOperationIr::BitwiseLeftShift(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + IntOperationIr::BitwiseLeftShiftScalar(repr) => Box::new([&repr.lhs].into_iter()), + IntOperationIr::BitwiseRightShift(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + IntOperationIr::BitwiseRightShiftScalar(repr) => Box::new([&repr.lhs].into_iter()), + } + } + + fn outputs(&self) -> Box + '_> { + match self { + IntOperationIr::Matmul(repr) => Box::new([&repr.out].into_iter()), + IntOperationIr::IntoFloat(repr) => Box::new([&repr.out].into_iter()), + IntOperationIr::BitwiseAnd(repr) => Box::new([&repr.out].into_iter()), + IntOperationIr::BitwiseAndScalar(repr) => Box::new([&repr.out].into_iter()), + IntOperationIr::BitwiseOr(repr) => Box::new([&repr.out].into_iter()), + IntOperationIr::BitwiseOrScalar(repr) => Box::new([&repr.out].into_iter()), + IntOperationIr::BitwiseXor(repr) => Box::new([&repr.out].into_iter()), + IntOperationIr::BitwiseXorScalar(repr) => Box::new([&repr.out].into_iter()), + IntOperationIr::BitwiseNot(repr) => Box::new([&repr.out].into_iter()), + IntOperationIr::BitwiseLeftShift(repr) => Box::new([&repr.out].into_iter()), + IntOperationIr::BitwiseLeftShiftScalar(repr) => Box::new([&repr.out].into_iter()), + IntOperationIr::BitwiseRightShift(repr) => Box::new([&repr.out].into_iter()), + IntOperationIr::BitwiseRightShiftScalar(repr) => Box::new([&repr.out].into_iter()), + } + } + + fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec { + let mut output = Vec::new(); + + match self { + IntOperationIr::Matmul(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + IntOperationIr::IntoFloat(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + IntOperationIr::BitwiseAnd(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + IntOperationIr::BitwiseAndScalar(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + IntOperationIr::BitwiseOr(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + IntOperationIr::BitwiseOrScalar(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + IntOperationIr::BitwiseXor(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + IntOperationIr::BitwiseXorScalar(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + IntOperationIr::BitwiseNot(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + IntOperationIr::BitwiseLeftShift(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + IntOperationIr::BitwiseLeftShiftScalar(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + IntOperationIr::BitwiseRightShift(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + IntOperationIr::BitwiseRightShiftScalar(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + } + }; + + output + } + + fn visit_mut(&mut self, v: &mut impl IrVisitorMut) { + match self { + IntOperationIr::Matmul(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + IntOperationIr::IntoFloat(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + IntOperationIr::BitwiseAnd(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + IntOperationIr::BitwiseAndScalar(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + IntOperationIr::BitwiseOr(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + IntOperationIr::BitwiseOrScalar(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + IntOperationIr::BitwiseXor(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + IntOperationIr::BitwiseXorScalar(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + IntOperationIr::BitwiseNot(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + IntOperationIr::BitwiseLeftShift(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + IntOperationIr::BitwiseLeftShiftScalar(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + IntOperationIr::BitwiseRightShift(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + IntOperationIr::BitwiseRightShiftScalar(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + } + } +} + +impl BoolOperationIr { + fn inputs(&self) -> Box + '_> { + match self { + BoolOperationIr::IntoFloat(repr) => Box::new([&repr.input].into_iter()), + BoolOperationIr::IntoInt(repr) => Box::new([&repr.input].into_iter()), + BoolOperationIr::Not(repr) => Box::new([&repr.input].into_iter()), + BoolOperationIr::And(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + BoolOperationIr::Or(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + BoolOperationIr::Xor(repr) => Box::new([&repr.lhs, &repr.rhs].into_iter()), + } + } + fn outputs(&self) -> Box + '_> { + match self { + BoolOperationIr::IntoFloat(repr) => Box::new([&repr.out].into_iter()), + BoolOperationIr::IntoInt(repr) => Box::new([&repr.out].into_iter()), + BoolOperationIr::Not(repr) => Box::new([&repr.out].into_iter()), + BoolOperationIr::And(repr) => Box::new([&repr.out].into_iter()), + BoolOperationIr::Or(repr) => Box::new([&repr.out].into_iter()), + BoolOperationIr::Xor(repr) => Box::new([&repr.out].into_iter()), + } + } + fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec { + let mut output = Vec::new(); + + match self { + BoolOperationIr::IntoFloat(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + BoolOperationIr::IntoInt(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + BoolOperationIr::Not(repr) => { + repr.input.mark_read_only(nodes, &mut output); + } + BoolOperationIr::And(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + BoolOperationIr::Or(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + BoolOperationIr::Xor(repr) => { + repr.lhs.mark_read_only(nodes, &mut output); + repr.rhs.mark_read_only(nodes, &mut output); + } + }; + + output + } + + fn visit_mut(&mut self, v: &mut impl IrVisitorMut) { + match self { + BoolOperationIr::IntoFloat(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + BoolOperationIr::IntoInt(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + BoolOperationIr::Not(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + BoolOperationIr::And(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + BoolOperationIr::Or(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + BoolOperationIr::Xor(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + } + } +} + +impl ModuleOperationIr { + fn inputs(&self) -> Box + '_> { + match self { + ModuleOperationIr::Embedding(repr) => { + Box::new([&repr.weights, &repr.indices].into_iter()) + } + ModuleOperationIr::EmbeddingBackward(repr) => { + Box::new([&repr.weights, &repr.out_grad, &repr.indices].into_iter()) + } + ModuleOperationIr::Linear(repr) => { + if let Some(bias) = &repr.bias { + Box::new([&repr.x, &repr.weight, bias].into_iter()) + } else { + Box::new([&repr.x, &repr.weight].into_iter()) + } + } + ModuleOperationIr::LinearXBackward(repr) => { + Box::new([&repr.weight, &repr.output_grad].into_iter()) + } + ModuleOperationIr::LinearWeightBackward(repr) => { + Box::new([&repr.x, &repr.output_grad].into_iter()) + } + ModuleOperationIr::LinearBiasBackward(repr) => { + Box::new([&repr.output_grad].into_iter()) + } + ModuleOperationIr::Conv1d(repr) => { + if let Some(bias) = &repr.bias { + Box::new([&repr.x, &repr.weight, bias].into_iter()) + } else { + Box::new([&repr.x, &repr.weight].into_iter()) + } + } + ModuleOperationIr::Conv1dXBackward(repr) => { + Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter()) + } + ModuleOperationIr::Conv1dWeightBackward(repr) => { + Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter()) + } + ModuleOperationIr::Conv1dBiasBackward(repr) => { + Box::new([&repr.x, &repr.bias, &repr.output_grad].into_iter()) + } + ModuleOperationIr::Conv2d(repr) => { + if let Some(bias) = &repr.bias { + Box::new([&repr.x, &repr.weight, bias].into_iter()) + } else { + Box::new([&repr.x, &repr.weight].into_iter()) + } + } + ModuleOperationIr::Conv2dXBackward(repr) => { + Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter()) + } + ModuleOperationIr::Conv2dWeightBackward(repr) => { + Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter()) + } + ModuleOperationIr::Conv2dBiasBackward(repr) => { + Box::new([&repr.x, &repr.bias, &repr.output_grad].into_iter()) + } + ModuleOperationIr::Conv3d(repr) => { + if let Some(bias) = &repr.bias { + Box::new([&repr.x, &repr.weight, bias].into_iter()) + } else { + Box::new([&repr.x, &repr.weight].into_iter()) + } + } + ModuleOperationIr::Conv3dXBackward(repr) => { + Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter()) + } + ModuleOperationIr::Conv3dWeightBackward(repr) => { + Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter()) + } + ModuleOperationIr::Conv3dBiasBackward(repr) => { + Box::new([&repr.x, &repr.bias, &repr.output_grad].into_iter()) + } + ModuleOperationIr::DeformableConv2d(repr) => match (&repr.mask, &repr.bias) { + (Some(mask), Some(bias)) => { + Box::new([&repr.x, &repr.offset, &repr.weight, mask, bias].into_iter()) + } + (Some(mask), None) => { + Box::new([&repr.x, &repr.offset, &repr.weight, mask].into_iter()) + } + (None, Some(bias)) => { + Box::new([&repr.x, &repr.offset, &repr.weight, bias].into_iter()) + } + (None, None) => Box::new([&repr.x, &repr.offset, &repr.weight].into_iter()), + }, + ModuleOperationIr::DeformableConv2dBackward(repr) => match (&repr.mask, &repr.bias) { + (Some(mask), Some(bias)) => Box::new( + [ + &repr.x, + &repr.offset, + &repr.weight, + &repr.out_grad, + mask, + bias, + ] + .into_iter(), + ), + (Some(mask), None) => Box::new( + [&repr.x, &repr.offset, &repr.weight, &repr.out_grad, mask].into_iter(), + ), + (None, Some(bias)) => Box::new( + [&repr.x, &repr.offset, &repr.weight, &repr.out_grad, bias].into_iter(), + ), + (None, None) => { + Box::new([&repr.x, &repr.offset, &repr.weight, &repr.out_grad].into_iter()) + } + }, + ModuleOperationIr::ConvTranspose1d(repr) => { + if let Some(bias) = &repr.bias { + Box::new([&repr.x, &repr.weight, bias].into_iter()) + } else { + Box::new([&repr.x, &repr.weight].into_iter()) + } + } + ModuleOperationIr::ConvTranspose2d(repr) => { + if let Some(bias) = &repr.bias { + Box::new([&repr.x, &repr.weight, bias].into_iter()) + } else { + Box::new([&repr.x, &repr.weight].into_iter()) + } + } + ModuleOperationIr::ConvTranspose3d(repr) => { + if let Some(bias) = &repr.bias { + Box::new([&repr.x, &repr.weight, bias].into_iter()) + } else { + Box::new([&repr.x, &repr.weight].into_iter()) + } + } + ModuleOperationIr::AvgPool1d(repr) => Box::new([&repr.x].into_iter()), + ModuleOperationIr::AvgPool2d(repr) => Box::new([&repr.x].into_iter()), + ModuleOperationIr::AvgPool1dBackward(repr) => { + Box::new([&repr.x, &repr.grad].into_iter()) + } + ModuleOperationIr::AvgPool2dBackward(repr) => { + Box::new([&repr.x, &repr.grad].into_iter()) + } + ModuleOperationIr::AdaptiveAvgPool1d(repr) => Box::new([&repr.x].into_iter()), + ModuleOperationIr::AdaptiveAvgPool2d(repr) => Box::new([&repr.x].into_iter()), + ModuleOperationIr::AdaptiveAvgPool1dBackward(repr) => { + Box::new([&repr.x, &repr.grad].into_iter()) + } + ModuleOperationIr::AdaptiveAvgPool2dBackward(repr) => { + Box::new([&repr.x, &repr.grad].into_iter()) + } + ModuleOperationIr::MaxPool1d(repr) => Box::new([&repr.x].into_iter()), + ModuleOperationIr::MaxPool1dWithIndices(repr) => Box::new([&repr.x].into_iter()), + ModuleOperationIr::MaxPool1dWithIndicesBackward(repr) => { + Box::new([&repr.x, &repr.indices, &repr.grad].into_iter()) + } + ModuleOperationIr::MaxPool2d(repr) => Box::new([&repr.x].into_iter()), + ModuleOperationIr::MaxPool2dWithIndices(repr) => Box::new([&repr.x].into_iter()), + ModuleOperationIr::MaxPool2dWithIndicesBackward(repr) => { + Box::new([&repr.x, &repr.indices, &repr.grad].into_iter()) + } + ModuleOperationIr::Interpolate(repr) => Box::new([&repr.x].into_iter()), + ModuleOperationIr::InterpolateBackward(repr) => { + Box::new([&repr.x, &repr.grad].into_iter()) + } + ModuleOperationIr::Rfft(repr) => Box::new([&repr.signal].into_iter()), + ModuleOperationIr::IRfft(repr) => { + Box::new([&repr.input_re, &repr.input_im].into_iter()) + } + ModuleOperationIr::Attention(repr) => { + if let Some(mask) = &repr.mask { + if let Some(attn_bias) = &repr.attn_bias { + Box::new([&repr.query, &repr.key, &repr.value, mask, attn_bias].into_iter()) + } else { + Box::new([&repr.query, &repr.key, &repr.value, mask].into_iter()) + } + } else if let Some(attn_bias) = &repr.attn_bias { + Box::new([&repr.query, &repr.key, &repr.value, attn_bias].into_iter()) + } else { + Box::new([&repr.query, &repr.key, &repr.value].into_iter()) + } + } + ModuleOperationIr::CtcLoss(repr) => Box::new( + [ + &repr.log_probs, + &repr.targets, + &repr.input_lengths, + &repr.target_lengths, + ] + .into_iter(), + ), + ModuleOperationIr::CtcLossBackward(repr) => Box::new( + [ + &repr.log_probs, + &repr.targets, + &repr.input_lengths, + &repr.target_lengths, + &repr.grad_loss, + ] + .into_iter(), + ), + ModuleOperationIr::LayerNorm(repr) => match &repr.beta { + Some(beta) => Box::new([&repr.input, &repr.gamma, beta].into_iter()), + None => Box::new([&repr.input, &repr.gamma].into_iter()), + }, + ModuleOperationIr::Unfold4d(repr) => Box::new([&repr.x].into_iter()), + ModuleOperationIr::ConvTranspose1dWeightBackward(repr) => { + Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter()) + } + ModuleOperationIr::ConvTranspose1dBiasBackward(repr) => { + Box::new([&repr.x, &repr.bias, &repr.output_grad].into_iter()) + } + ModuleOperationIr::ConvTranspose2dWeightBackward(repr) => { + Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter()) + } + ModuleOperationIr::ConvTranspose2dBiasBackward(repr) => { + Box::new([&repr.x, &repr.bias, &repr.output_grad].into_iter()) + } + ModuleOperationIr::ConvTranspose3dWeightBackward(repr) => { + Box::new([&repr.x, &repr.weight, &repr.output_grad].into_iter()) + } + ModuleOperationIr::ConvTranspose3dBiasBackward(repr) => { + Box::new([&repr.x, &repr.bias, &repr.output_grad].into_iter()) + } + } + } + fn outputs(&self) -> Box + '_> { + match self { + ModuleOperationIr::Embedding(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::EmbeddingBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::Linear(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::LinearXBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::LinearWeightBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::LinearBiasBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::Conv1d(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::Conv1dXBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::Conv1dWeightBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::Conv1dBiasBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::Conv2d(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::Conv2dXBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::Conv2dWeightBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::Conv2dBiasBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::Conv3d(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::Conv3dXBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::Conv3dWeightBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::Conv3dBiasBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::DeformableConv2d(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::DeformableConv2dBackward(repr) => { + match (&repr.mask_grad, &repr.bias_grad) { + (Some(mask_grad), Some(bias_grad)) => Box::new( + [ + &repr.input_grad, + &repr.offset_grad, + &repr.weight_grad, + mask_grad, + bias_grad, + ] + .into_iter(), + ), + (Some(mask_grad), None) => Box::new( + [ + &repr.input_grad, + &repr.offset_grad, + &repr.weight_grad, + mask_grad, + ] + .into_iter(), + ), + (None, Some(bias_grad)) => Box::new( + [ + &repr.input_grad, + &repr.offset_grad, + &repr.weight_grad, + bias_grad, + ] + .into_iter(), + ), + (None, None) => Box::new( + [&repr.input_grad, &repr.offset_grad, &repr.weight_grad].into_iter(), + ), + } + } + ModuleOperationIr::ConvTranspose1d(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::ConvTranspose2d(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::ConvTranspose3d(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::AvgPool1d(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::AvgPool2d(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::AvgPool1dBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::AvgPool2dBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::AdaptiveAvgPool1d(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::AdaptiveAvgPool2d(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::AdaptiveAvgPool1dBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::AdaptiveAvgPool2dBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::MaxPool1d(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::MaxPool1dWithIndices(repr) => { + Box::new([&repr.out, &repr.out_indices].into_iter()) + } + ModuleOperationIr::MaxPool1dWithIndicesBackward(repr) => { + Box::new([&repr.out].into_iter()) + } + ModuleOperationIr::MaxPool2d(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::MaxPool2dWithIndices(repr) => { + Box::new([&repr.out, &repr.out_indices].into_iter()) + } + ModuleOperationIr::MaxPool2dWithIndicesBackward(repr) => { + Box::new([&repr.out].into_iter()) + } + ModuleOperationIr::Interpolate(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::InterpolateBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::Rfft(repr) => Box::new([&repr.out_re, &repr.out_im].into_iter()), + ModuleOperationIr::IRfft(repr) => Box::new([&repr.out_signal].into_iter()), + ModuleOperationIr::Attention(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::CtcLoss(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::CtcLossBackward(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::LayerNorm(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::Unfold4d(repr) => Box::new([&repr.out].into_iter()), + ModuleOperationIr::ConvTranspose1dWeightBackward(repr) => { + Box::new([&repr.out].into_iter()) + } + ModuleOperationIr::ConvTranspose1dBiasBackward(repr) => { + Box::new([&repr.out].into_iter()) + } + ModuleOperationIr::ConvTranspose2dWeightBackward(repr) => { + Box::new([&repr.out].into_iter()) + } + ModuleOperationIr::ConvTranspose2dBiasBackward(repr) => { + Box::new([&repr.out].into_iter()) + } + ModuleOperationIr::ConvTranspose3dWeightBackward(repr) => { + Box::new([&repr.out].into_iter()) + } + ModuleOperationIr::ConvTranspose3dBiasBackward(repr) => { + Box::new([&repr.out].into_iter()) + } + } + } + + fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec { + let mut output = Vec::new(); + + match self { + ModuleOperationIr::Embedding(repr) => { + repr.weights.mark_read_only(nodes, &mut output); + repr.indices.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::EmbeddingBackward(repr) => { + repr.weights.mark_read_only(nodes, &mut output); + repr.out_grad.mark_read_only(nodes, &mut output); + repr.indices.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::Linear(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + + if let Some(bias) = &mut repr.bias { + bias.mark_read_only(nodes, &mut output); + } + } + ModuleOperationIr::LinearXBackward(repr) => { + repr.weight.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::LinearWeightBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::LinearBiasBackward(repr) => { + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::Conv1d(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + + if let Some(bias) = &mut repr.bias { + bias.mark_read_only(nodes, &mut output); + } + } + ModuleOperationIr::Conv1dXBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::Conv1dWeightBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::Conv1dBiasBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.bias.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::Conv2d(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + + if let Some(bias) = &mut repr.bias { + bias.mark_read_only(nodes, &mut output); + } + } + ModuleOperationIr::Conv2dXBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::Conv2dWeightBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::Conv2dBiasBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.bias.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::Conv3d(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + + if let Some(bias) = &mut repr.bias { + bias.mark_read_only(nodes, &mut output); + } + } + ModuleOperationIr::Conv3dXBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::Conv3dWeightBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::Conv3dBiasBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.bias.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::DeformableConv2d(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + repr.offset.mark_read_only(nodes, &mut output); + + match (&mut repr.mask, &mut repr.bias) { + (Some(mask), Some(bias)) => { + mask.mark_read_only(nodes, &mut output); + bias.mark_read_only(nodes, &mut output); + } + (Some(mask), None) => { + mask.mark_read_only(nodes, &mut output); + } + (None, Some(bias)) => { + bias.mark_read_only(nodes, &mut output); + } + (None, None) => {} + }; + } + ModuleOperationIr::DeformableConv2dBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + repr.offset.mark_read_only(nodes, &mut output); + repr.out_grad.mark_read_only(nodes, &mut output); + + if let Some(mask) = repr.mask.as_mut() { + mask.mark_read_only(nodes, &mut output); + } + if let Some(bias) = repr.bias.as_mut() { + bias.mark_read_only(nodes, &mut output); + } + } + ModuleOperationIr::ConvTranspose1d(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + + if let Some(bias) = &mut repr.bias { + bias.mark_read_only(nodes, &mut output); + } + } + ModuleOperationIr::ConvTranspose2d(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + + if let Some(bias) = &mut repr.bias { + bias.mark_read_only(nodes, &mut output); + } + } + ModuleOperationIr::ConvTranspose3d(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + + if let Some(bias) = &mut repr.bias { + bias.mark_read_only(nodes, &mut output); + } + } + ModuleOperationIr::AvgPool1d(repr) => { + repr.x.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::AvgPool2d(repr) => { + repr.x.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::AvgPool1dBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::AvgPool2dBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::AdaptiveAvgPool1d(repr) => { + repr.x.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::AdaptiveAvgPool2d(repr) => { + repr.x.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::AdaptiveAvgPool1dBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::AdaptiveAvgPool2dBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::MaxPool1d(repr) => { + repr.x.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::MaxPool1dWithIndices(repr) => { + repr.x.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::MaxPool1dWithIndicesBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::MaxPool2d(repr) => { + repr.x.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::MaxPool2dWithIndices(repr) => { + repr.x.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::MaxPool2dWithIndicesBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::Interpolate(repr) => { + repr.x.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::InterpolateBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::Rfft(repr) => { + repr.signal.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::IRfft(repr) => { + repr.input_re.mark_read_only(nodes, &mut output); + repr.input_im.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::Attention(repr) => { + repr.query.mark_read_only(nodes, &mut output); + repr.key.mark_read_only(nodes, &mut output); + repr.value.mark_read_only(nodes, &mut output); + if let Some(mask) = &mut repr.mask { + mask.mark_read_only(nodes, &mut output); + } + if let Some(attn_bias) = &mut repr.attn_bias { + attn_bias.mark_read_only(nodes, &mut output); + } + } + ModuleOperationIr::CtcLoss(repr) => { + repr.log_probs.mark_read_only(nodes, &mut output); + repr.targets.mark_read_only(nodes, &mut output); + repr.input_lengths.mark_read_only(nodes, &mut output); + repr.target_lengths.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::CtcLossBackward(repr) => { + repr.log_probs.mark_read_only(nodes, &mut output); + repr.targets.mark_read_only(nodes, &mut output); + repr.input_lengths.mark_read_only(nodes, &mut output); + repr.target_lengths.mark_read_only(nodes, &mut output); + repr.grad_loss.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::LayerNorm(repr) => { + repr.input.mark_read_only(nodes, &mut output); + repr.gamma.mark_read_only(nodes, &mut output); + if let Some(beta) = &mut repr.beta { + beta.mark_read_only(nodes, &mut output); + } + } + ModuleOperationIr::Unfold4d(repr) => { + repr.x.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::ConvTranspose1dWeightBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::ConvTranspose1dBiasBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.bias.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::ConvTranspose2dWeightBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::ConvTranspose2dBiasBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.bias.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::ConvTranspose3dWeightBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.weight.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + ModuleOperationIr::ConvTranspose3dBiasBackward(repr) => { + repr.x.mark_read_only(nodes, &mut output); + repr.bias.mark_read_only(nodes, &mut output); + repr.output_grad.mark_read_only(nodes, &mut output); + } + }; + + output + } + + fn visit_mut(&mut self, v: &mut impl IrVisitorMut) { + match self { + ModuleOperationIr::Embedding(repr) => { + v.visit_tensor_mut(&mut repr.weights); + v.visit_tensor_mut(&mut repr.indices); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::EmbeddingBackward(repr) => { + v.visit_tensor_mut(&mut repr.weights); + v.visit_tensor_mut(&mut repr.out_grad); + v.visit_tensor_mut(&mut repr.indices); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::Linear(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.weight); + if let Some(bias) = &mut repr.bias { + v.visit_tensor_mut(bias); + } + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::LinearXBackward(repr) => { + v.visit_tensor_mut(&mut repr.weight); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::LinearWeightBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::LinearBiasBackward(repr) => { + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::Conv1d(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.weight); + if let Some(bias) = &mut repr.bias { + v.visit_tensor_mut(bias); + } + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::Conv1dXBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.weight); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::Conv1dWeightBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.weight); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::Conv1dBiasBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.bias); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::Conv2d(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.weight); + if let Some(bias) = &mut repr.bias { + v.visit_tensor_mut(bias); + } + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::Conv2dXBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.weight); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::Conv2dWeightBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.weight); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::Conv2dBiasBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.bias); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::Conv3d(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.weight); + if let Some(bias) = &mut repr.bias { + v.visit_tensor_mut(bias); + } + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::Conv3dXBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.weight); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::Conv3dWeightBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.weight); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::Conv3dBiasBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.bias); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::DeformableConv2d(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.offset); + v.visit_tensor_mut(&mut repr.weight); + if let Some(mask) = &mut repr.mask { + v.visit_tensor_mut(mask); + } + if let Some(bias) = &mut repr.bias { + v.visit_tensor_mut(bias); + } + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::DeformableConv2dBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.offset); + v.visit_tensor_mut(&mut repr.weight); + v.visit_tensor_mut(&mut repr.out_grad); + if let Some(mask) = &mut repr.mask { + v.visit_tensor_mut(mask); + } + if let Some(bias) = &mut repr.bias { + v.visit_tensor_mut(bias); + } + v.visit_tensor_mut(&mut repr.input_grad); + v.visit_tensor_mut(&mut repr.offset_grad); + v.visit_tensor_mut(&mut repr.weight_grad); + if let Some(mask_grad) = &mut repr.mask_grad { + v.visit_tensor_mut(mask_grad); + } + if let Some(bias_grad) = &mut repr.bias_grad { + v.visit_tensor_mut(bias_grad); + } + } + ModuleOperationIr::ConvTranspose1d(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.weight); + if let Some(bias) = &mut repr.bias { + v.visit_tensor_mut(bias); + } + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::ConvTranspose2d(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.weight); + if let Some(bias) = &mut repr.bias { + v.visit_tensor_mut(bias); + } + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::ConvTranspose3d(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.weight); + if let Some(bias) = &mut repr.bias { + v.visit_tensor_mut(bias); + } + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::AvgPool1d(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::AvgPool2d(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::AvgPool1dBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::AvgPool2dBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::AdaptiveAvgPool1d(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::AdaptiveAvgPool2d(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::AdaptiveAvgPool1dBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::AdaptiveAvgPool2dBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::MaxPool1d(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::MaxPool1dWithIndices(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.out); + v.visit_tensor_mut(&mut repr.out_indices); + } + ModuleOperationIr::MaxPool1dWithIndicesBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.indices); + v.visit_tensor_mut(&mut repr.grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::MaxPool2d(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::MaxPool2dWithIndices(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.out); + v.visit_tensor_mut(&mut repr.out_indices); + } + ModuleOperationIr::MaxPool2dWithIndicesBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.indices); + v.visit_tensor_mut(&mut repr.grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::Interpolate(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::InterpolateBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::Rfft(repr) => { + v.visit_tensor_mut(&mut repr.signal); + v.visit_tensor_mut(&mut repr.out_re); + v.visit_tensor_mut(&mut repr.out_im); + } + ModuleOperationIr::IRfft(repr) => { + v.visit_tensor_mut(&mut repr.input_re); + v.visit_tensor_mut(&mut repr.input_im); + v.visit_tensor_mut(&mut repr.out_signal); + } + ModuleOperationIr::Attention(repr) => { + v.visit_tensor_mut(&mut repr.query); + v.visit_tensor_mut(&mut repr.key); + v.visit_tensor_mut(&mut repr.value); + if let Some(mask) = &mut repr.mask { + v.visit_tensor_mut(mask); + } + if let Some(attn_bias) = &mut repr.attn_bias { + v.visit_tensor_mut(attn_bias); + } + v.visit_tensor_mut(&mut repr.out); + if let Some(scale) = &mut repr.options.scale { + v.visit_scalar_mut(scale); + } + if let Some(softcap) = &mut repr.options.softcap { + v.visit_scalar_mut(softcap); + } + } + ModuleOperationIr::CtcLoss(repr) => { + v.visit_tensor_mut(&mut repr.log_probs); + v.visit_tensor_mut(&mut repr.targets); + v.visit_tensor_mut(&mut repr.input_lengths); + v.visit_tensor_mut(&mut repr.target_lengths); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::CtcLossBackward(repr) => { + v.visit_tensor_mut(&mut repr.log_probs); + v.visit_tensor_mut(&mut repr.targets); + v.visit_tensor_mut(&mut repr.input_lengths); + v.visit_tensor_mut(&mut repr.target_lengths); + v.visit_tensor_mut(&mut repr.grad_loss); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::LayerNorm(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.gamma); + if let Some(beta) = &mut repr.beta { + v.visit_tensor_mut(beta); + } + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.epsilon); + } + ModuleOperationIr::Unfold4d(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::ConvTranspose1dWeightBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.weight); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::ConvTranspose1dBiasBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.bias); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::ConvTranspose2dWeightBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.weight); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::ConvTranspose2dBiasBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.bias); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::ConvTranspose3dWeightBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.weight); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + ModuleOperationIr::ConvTranspose3dBiasBackward(repr) => { + v.visit_tensor_mut(&mut repr.x); + v.visit_tensor_mut(&mut repr.bias); + v.visit_tensor_mut(&mut repr.output_grad); + v.visit_tensor_mut(&mut repr.out); + } + } + } +} + +impl DistributedOperationIr { + fn inputs(&self) -> Box + '_> { + match self { + DistributedOperationIr::AllReduce(repr) => Box::new([&repr.tensor].into_iter()), + DistributedOperationIr::SyncCollective => Box::new([].into_iter()), + } + } + + fn outputs(&self) -> Box + '_> { + match self { + DistributedOperationIr::AllReduce(repr) => Box::new([&repr.out].into_iter()), + DistributedOperationIr::SyncCollective => Box::new([].into_iter()), + } + } + + fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec { + let mut output = Vec::new(); + + match self { + DistributedOperationIr::AllReduce(repr) => { + repr.tensor.mark_read_only(nodes, &mut output); + } + DistributedOperationIr::SyncCollective => {} + } + + output + } + + fn visit_mut(&mut self, v: &mut impl IrVisitorMut) { + match self { + DistributedOperationIr::AllReduce(repr) => { + v.visit_tensor_mut(&mut repr.tensor); + v.visit_tensor_mut(&mut repr.out); + } + DistributedOperationIr::SyncCollective => {} + } + } +} + +impl InitOperationIr { + fn inputs(&self) -> Box + '_> { + Box::new([].into_iter()) + } + fn outputs(&self) -> Box + '_> { + Box::new([&self.out].into_iter()) + } + + fn visit_mut(&mut self, v: &mut impl IrVisitorMut) { + v.visit_tensor_mut(&mut self.out); + } +} + +impl TensorIr { + fn mark_read_only(&mut self, nodes: &[TensorId], output: &mut Vec) { + if self.status == TensorStatus::ReadWrite && nodes.contains(&self.id) { + output.push(self.clone()); + self.status = TensorStatus::ReadOnly; + } + } +} + +impl core::hash::Hash for RandomOpIr { + fn hash(&self, state: &mut H) { + self.out.hash(state); + + match self.distribution { + Distribution::Default => 1u8.hash(state), + Distribution::Bernoulli(_) => 2u8.hash(state), + Distribution::Uniform(_, _) => 3u8.hash(state), + Distribution::Normal(_, _) => 4u8.hash(state), + } + } +} + +/// Extension trait to extract outputs when registering an operation. +pub trait OperationOutput { + /// Extract a single output. + fn output(self) -> O; + + /// Extract a fixed number of outputs. + fn outputs(self) -> [O; N]; +} + +impl OperationOutput for Vec { + fn output(self) -> O { + let [tensor] = self.outputs(); + tensor + } + + fn outputs(self) -> [O; N] { + self.try_into().unwrap() + } +} + +/// Operation IR for sort along a dim. The output preserves the input shape/dtype. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct SortOpIr { + /// Input tensor. + pub input: TensorIr, + /// Dim along which to sort. + pub dim: usize, + /// Sort descending. + pub descending: bool, + /// Output tensor (same shape/dtype as input). + pub out: TensorIr, +} + +/// Operation IR for sort-with-indices: returns sorted values + source indices. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct SortWithIndicesOpIr { + /// Input tensor. + pub input: TensorIr, + /// Dim along which to sort. + pub dim: usize, + /// Sort descending. + pub descending: bool, + /// Output tensor with sorted values. + pub out: TensorIr, + /// Output tensor with the indices into the original input. + pub out_indices: TensorIr, +} + +/// Operation IR for layer normalization with optional bias. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct LayerNormOpIr { + /// Input tensor. + pub input: TensorIr, + /// Scale (gamma) parameter. + pub gamma: TensorIr, + /// Optional shift (beta) parameter. + pub beta: Option, + /// Numerical-stability epsilon. + pub epsilon: ScalarIr, + /// Output tensor. + pub out: TensorIr, +} + +/// Operation IR for unfold4d (a 4d sliding-window kernel-as-explicit-columns reshape). +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct Unfold4dOpIr { + /// Input tensor of shape `[N, C, H, W]`. + pub x: TensorIr, + /// Kernel size `[kH, kW]`. + pub kernel_size: [usize; 2], + /// Conv-like options (stride, padding, dilation). + pub options: Unfold4dOptionsIr, + /// Output tensor. + pub out: TensorIr, +} + +/// Options for [`Unfold4dOpIr`] — mirrors the backend's `UnfoldOptions`. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct Unfold4dOptionsIr { + /// Stride `[sH, sW]`. + pub stride: [usize; 2], + /// Padding `[pH, pW]`. + pub padding: [usize; 2], + /// Dilation `[dH, dW]`. + pub dilation: [usize; 2], +} + +/// Operation IR for the weight-gradient of `conv_transpose1d`. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct ConvTranspose1dWeightBackwardOpIr { + /// Forward input. + pub x: TensorIr, + /// Convolution weights. + pub weight: TensorIr, + /// Upstream gradient. + pub output_grad: TensorIr, + /// Convolution options. + pub options: ConvTranspose1dOptionsIr, + /// Output: gradient w.r.t. `weight`. + pub out: TensorIr, +} + +/// Operation IR for the bias-gradient of `conv_transpose1d`. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct ConvTranspose1dBiasBackwardOpIr { + /// Forward input. + pub x: TensorIr, + /// Bias tensor (sets the output shape). + pub bias: TensorIr, + /// Upstream gradient. + pub output_grad: TensorIr, + /// Output: gradient w.r.t. `bias`. + pub out: TensorIr, +} + +/// Operation IR for the weight-gradient of `conv_transpose2d`. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct ConvTranspose2dWeightBackwardOpIr { + /// Forward input. + pub x: TensorIr, + /// Convolution weights. + pub weight: TensorIr, + /// Upstream gradient. + pub output_grad: TensorIr, + /// Convolution options. + pub options: ConvTranspose2dOptionsIr, + /// Output: gradient w.r.t. `weight`. + pub out: TensorIr, +} + +/// Operation IR for the bias-gradient of `conv_transpose2d`. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct ConvTranspose2dBiasBackwardOpIr { + /// Forward input. + pub x: TensorIr, + /// Bias tensor. + pub bias: TensorIr, + /// Upstream gradient. + pub output_grad: TensorIr, + /// Output: gradient w.r.t. `bias`. + pub out: TensorIr, +} + +/// Operation IR for the weight-gradient of `conv_transpose3d`. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct ConvTranspose3dWeightBackwardOpIr { + /// Forward input. + pub x: TensorIr, + /// Convolution weights. + pub weight: TensorIr, + /// Upstream gradient. + pub output_grad: TensorIr, + /// Convolution options. + pub options: ConvTranspose3dOptionsIr, + /// Output: gradient w.r.t. `weight`. + pub out: TensorIr, +} + +/// Operation IR for the bias-gradient of `conv_transpose3d`. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct ConvTranspose3dBiasBackwardOpIr { + /// Forward input. + pub x: TensorIr, + /// Bias tensor. + pub bias: TensorIr, + /// Upstream gradient. + pub output_grad: TensorIr, + /// Output: gradient w.r.t. `bias`. + pub out: TensorIr, +} + +/// Operation IR for the hard-sigmoid activation function, which takes two scalars +/// (`alpha` and `beta`) in addition to the input tensor. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub struct HardSigmoidOpIr { + /// Input tensor. + pub tensor: TensorIr, + /// Alpha — multiplied with the tensor before adding `beta`. + pub alpha: ScalarIr, + /// Beta — added after multiplying by `alpha`. + pub beta: ScalarIr, + /// Output tensor. + pub out: TensorIr, +} + +/// Operation intermediate representation for activation functions. +/// +/// Activations are kept in their own enum (rather than smeared across `FloatOperationIr`) +/// so backends can match a single arm to dispatch the fused implementation, instead of +/// receiving the function decomposed into 5–20 primitive ops. +#[derive(Clone, Debug, Hash, PartialEq, Serialize, Deserialize)] +pub enum ActivationOperationIr { + /// [relu](burn_backend::ops::ActivationOps::relu). + Relu(UnaryOpIr), + /// [relu_backward](burn_backend::ops::ActivationOps::relu_backward). + /// `lhs` = output of the forward pass, `rhs` = upstream gradient. + ReluBackward(BinaryOpIr), + /// [leaky_relu](burn_backend::ops::ActivationOps::leaky_relu). + /// `lhs` = input, `rhs` = `negative_slope` scalar. + LeakyRelu(ScalarOpIr), + /// [prelu](burn_backend::ops::ActivationOps::prelu). + /// `lhs` = input, `rhs` = alpha tensor. + PRelu(BinaryOpIr), + /// [gelu](burn_backend::ops::ActivationOps::gelu). + Gelu(UnaryOpIr), + /// [gelu_backward](burn_backend::ops::ActivationOps::gelu_backward). + /// `lhs` = forward input, `rhs` = upstream gradient. + GeluBackward(BinaryOpIr), + /// [sigmoid](burn_backend::ops::ActivationOps::sigmoid). + Sigmoid(UnaryOpIr), + /// [sigmoid_backward](burn_backend::ops::ActivationOps::sigmoid_backward). + /// `lhs` = output of the forward pass, `rhs` = upstream gradient. + SigmoidBackward(BinaryOpIr), + /// [hard_sigmoid](burn_backend::ops::ActivationOps::hard_sigmoid). + HardSigmoid(HardSigmoidOpIr), + /// [log_sigmoid](burn_backend::ops::ActivationOps::log_sigmoid). + LogSigmoid(UnaryOpIr), + /// [log_sigmoid_backward](burn_backend::ops::ActivationOps::log_sigmoid_backward). + /// `lhs` = forward input, `rhs` = upstream gradient. + LogSigmoidBackward(BinaryOpIr), + /// [softmax](burn_backend::ops::ActivationOps::softmax). + Softmax(DimOpIr), + /// [log_softmax](burn_backend::ops::ActivationOps::log_softmax). + LogSoftmax(DimOpIr), + /// [softmin](burn_backend::ops::ActivationOps::softmin). + Softmin(DimOpIr), +} + +/// Generate [`ActivationOperationIr`]'s `inputs`/`outputs`/`mark_read_only` from a single +/// variant → input-field table, so the three methods can't drift out of sync. The `out` +/// field is implicit (every variant's op IR has one). Each entry lists the variant and the +/// names of its input tensor fields (`input`, `lhs`/`rhs`, `tensor`, …). +macro_rules! activation_ir_tensor_access { + ($( $variant:ident => [ $($field:ident),+ ] ),+ $(,)?) => { + impl ActivationOperationIr { + fn inputs(&self) -> Box + '_> { + match self { + $( Self::$variant(repr) => Box::new([$(&repr.$field),+].into_iter()), )+ + } + } + + fn outputs(&self) -> Box + '_> { + match self { + $( Self::$variant(repr) => Box::new([&repr.out].into_iter()), )+ + } + } + + fn mark_read_only(&mut self, nodes: &[TensorId]) -> Vec { + let mut output = Vec::new(); + match self { + $( Self::$variant(repr) => { + $( repr.$field.mark_read_only(nodes, &mut output); )+ + } )+ + } + output + } + } + }; +} + +impl ActivationOperationIr { + fn visit_mut(&mut self, v: &mut impl IrVisitorMut) { + match self { + ActivationOperationIr::Relu(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + ActivationOperationIr::ReluBackward(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + ActivationOperationIr::LeakyRelu(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.rhs); + } + ActivationOperationIr::PRelu(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + ActivationOperationIr::Gelu(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + ActivationOperationIr::GeluBackward(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + ActivationOperationIr::Sigmoid(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + ActivationOperationIr::SigmoidBackward(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + ActivationOperationIr::HardSigmoid(repr) => { + v.visit_tensor_mut(&mut repr.tensor); + v.visit_tensor_mut(&mut repr.out); + v.visit_scalar_mut(&mut repr.alpha); + v.visit_scalar_mut(&mut repr.beta); + } + ActivationOperationIr::LogSigmoid(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + ActivationOperationIr::LogSigmoidBackward(repr) => { + v.visit_tensor_mut(&mut repr.lhs); + v.visit_tensor_mut(&mut repr.rhs); + v.visit_tensor_mut(&mut repr.out); + } + ActivationOperationIr::Softmax(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + ActivationOperationIr::LogSoftmax(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + ActivationOperationIr::Softmin(repr) => { + v.visit_tensor_mut(&mut repr.input); + v.visit_tensor_mut(&mut repr.out); + } + } + } +} + +activation_ir_tensor_access! { + Relu => [input], + ReluBackward => [lhs, rhs], + LeakyRelu => [lhs], + PRelu => [lhs, rhs], + Gelu => [input], + GeluBackward => [lhs, rhs], + Sigmoid => [input], + SigmoidBackward => [lhs, rhs], + HardSigmoid => [tensor], + LogSigmoid => [input], + LogSigmoidBackward => [lhs, rhs], + Softmax => [input], + LogSoftmax => [input], + Softmin => [input], +} + +#[cfg(test)] +mod visit_mut_tests { + use super::*; + use burn_backend::{DType, Shape}; + + fn tensor(id: u64) -> TensorIr { + TensorIr::uninit(TensorId::new(id), Shape::from([2, 2]), DType::F32) + } + + /// Bumps every visited tensor id by 100 and collects (and rewrites) every visited scalar. + #[derive(Default)] + struct CollectVisitor { + scalars: Vec, + rewrite_scalar: Option, + } + + impl IrVisitorMut for CollectVisitor { + fn visit_tensor_mut(&mut self, tensor: &mut TensorIr) { + tensor.id = TensorId::new(tensor.id.value() + 100); + } + + fn visit_scalar_mut(&mut self, scalar: &mut ScalarIr) { + self.scalars.push(*scalar); + if let Some(value) = self.rewrite_scalar { + *scalar = value; + } + } + } + + #[test] + fn visit_mut_visits_all_tensors_and_scalars() { + // NumericFloat MulScalar (a ScalarOpIr with a scalar `rhs`). + let mut mul = OperationIr::NumericFloat( + DType::F32, + NumericOperationIr::MulScalar(ScalarOpIr { + lhs: tensor(1), + rhs: ScalarIr::Float(2.0), + out: tensor(2), + }), + ); + + // Bump every tensor id by 100 and rewrite the scalar to 9.0. + let mut visitor = CollectVisitor { + rewrite_scalar: Some(ScalarIr::Float(9.0)), + ..Default::default() + }; + mul.visit_mut(&mut visitor); + + let ids: Vec = mul + .inputs() + .chain(mul.outputs()) + .map(|t| t.id.value()) + .collect(); + assert_eq!(ids, vec![101, 102]); + // The scalar must have been visited (and was rewritable). + assert_eq!(visitor.scalars, vec![ScalarIr::Float(2.0)]); + + // Visiting again observes the rewritten scalar. + let mut after = CollectVisitor::default(); + mul.visit_mut(&mut after); + assert_eq!(after.scalars, vec![ScalarIr::Float(9.0)]); + + // BaseFloat Reshape (a ShapeOpIr, input + out, no scalars). + let mut reshape = OperationIr::BaseFloat(BaseOperationIr::Reshape(ShapeOpIr { + input: tensor(10), + out: tensor(11), + })); + + let mut visitor = CollectVisitor::default(); + reshape.visit_mut(&mut visitor); + let ids: Vec = reshape + .inputs() + .chain(reshape.outputs()) + .map(|t| t.id.value()) + .collect(); + assert_eq!(ids, vec![110, 111]); + assert_eq!(visitor.scalars.len(), 0); + } +} diff --git a/crates/burn-ir/src/scalar.rs b/crates/burn-ir/src/scalar.rs new file mode 100644 index 0000000..3434776 --- /dev/null +++ b/crates/burn-ir/src/scalar.rs @@ -0,0 +1,77 @@ +use burn_backend::{DType, Scalar}; +use burn_backend::{Element, ElementConversion}; +use core::hash::Hash; +use serde::{Deserialize, Serialize}; + +/// A scalar representation. +#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)] +#[allow(missing_docs)] +pub enum ScalarIr { + Float(f64), + Int(i64), + UInt(u64), + Bool(bool), +} + +impl Hash for ScalarIr { + fn hash(&self, state: &mut H) { + match self { + ScalarIr::Float(x) => x.to_bits().hash(state), + ScalarIr::Int(x) => x.hash(state), + ScalarIr::UInt(x) => x.hash(state), + ScalarIr::Bool(x) => x.hash(state), + } + } +} + +impl ScalarIr { + /// Creates a scalar with the specified data type. + pub fn new(value: E, dtype: &DType) -> Self { + if dtype.is_float() { + Self::Float(value.elem()) + } else if dtype.is_int() { + Self::Int(value.elem()) + } else if dtype.is_uint() { + Self::UInt(value.elem()) + } else if dtype.is_bool() { + Self::Bool(value.elem()) + } else { + unimplemented!("Scalar not supported for {dtype:?}") + } + } + + /// Converts and returns the converted element. + pub fn elem(self) -> E { + match self { + ScalarIr::Float(x) => x.elem(), + ScalarIr::Int(x) => x.elem(), + ScalarIr::UInt(x) => x.elem(), + ScalarIr::Bool(x) => x.elem(), + } + } +} + +// The enums are similar, but both types have different roles: +// - `Scalar`: runtime literal value +// - `ScalarIr`: serializable literal representation (used for IR) +impl From for ScalarIr { + fn from(value: Scalar) -> Self { + match value { + Scalar::Float(x) => Self::Float(x), + Scalar::Int(x) => Self::Int(x), + Scalar::UInt(x) => Self::UInt(x), + Scalar::Bool(x) => Self::Bool(x), + } + } +} + +impl From for Scalar { + fn from(value: ScalarIr) -> Self { + match value { + ScalarIr::Float(x) => Self::Float(x), + ScalarIr::Int(x) => Self::Int(x), + ScalarIr::UInt(x) => Self::UInt(x), + ScalarIr::Bool(x) => Self::Bool(x), + } + } +} diff --git a/crates/burn-ir/src/tensor.rs b/crates/burn-ir/src/tensor.rs new file mode 100644 index 0000000..bd5d89f --- /dev/null +++ b/crates/burn-ir/src/tensor.rs @@ -0,0 +1,72 @@ +use serde::{Deserialize, Serialize}; + +use burn_backend::{DType, Shape}; + +/// The tensor unique identifier. +#[derive(Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord, Debug, Serialize, Deserialize)] +pub struct TensorId { + value: u64, +} + +impl core::fmt::Display for TensorId { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_fmt(format_args!("TensorId({:?})", self.value)) + } +} + +/// The status of the current tensor. +#[derive(Hash, Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum TensorStatus { + /// The tensor can be read, but not written. + ReadOnly, + /// The tensor can be mutated inplace. + ReadWrite, + /// No handle exists for that tensor. + NotInit, +} + +/// A tensor definition represents a snapshot of a tensor when it was used. +/// +/// # Example +/// +/// A tensor that is used multiple times has its status updated for each operation. +/// +/// 1. Status::NotInit +/// 2. Status::ReadOnly +/// 3. Status::ReadOnly +/// 4. Status::ReadWrite +#[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize, Deserialize)] +pub struct TensorIr { + /// The [tensor id](TensorId). + pub id: TensorId, + /// The shape of the tensor. + pub shape: Shape, + /// The [status](TensorStatus) of the tensor when it was used. + pub status: TensorStatus, + /// The [type](DType) of the tensor. + pub dtype: DType, +} + +impl TensorId { + /// Create a new tensor id. + pub fn new(value: u64) -> Self { + Self { value } + } + + /// Return the underlying numeric value of this id. + pub fn value(&self) -> u64 { + self.value + } +} + +impl TensorIr { + /// Create a new tensor that is not already initialized. + pub fn uninit(id: TensorId, shape: Shape, dtype: DType) -> Self { + Self { + id, + status: TensorStatus::NotInit, + shape, + dtype, + } + } +} diff --git a/crates/burn-ndarray/Cargo.toml b/crates/burn-ndarray/Cargo.toml new file mode 100644 index 0000000..b227128 --- /dev/null +++ b/crates/burn-ndarray/Cargo.toml @@ -0,0 +1,97 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science", "no-std", "embedded", "wasm"] +description = "Ndarray backend for the Burn framework" +documentation = "https://docs.rs/burn-ndarray" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "data"] +license.workspace = true +name = "burn-ndarray" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-ndarray" +version.workspace = true + +[lints] +workspace = true + +[features] +blas-accelerate = [ + "blas-src/accelerate", # Accelerate framework (macOS only) + "ndarray/blas", +] +blas-netlib = ["blas-src/netlib", "ndarray/blas"] +blas-openblas = ["blas-src/openblas", "ndarray/blas", "openblas-src"] +blas-openblas-system = [ + "blas-src/openblas", + "ndarray/blas", + "openblas-src/system", +] +default = ["std", "simd", "multi-threads"] +doc = ["default"] +multi-threads = [ + "rayon", + "ndarray/rayon", + "matrixmultiply/threading", +] +simd = ["macerator", "bytemuck", "seq-macro", "itertools"] +std = [ + "burn-autodiff", + "burn-std/std", + "burn-backend/std", + "burn-ir/std", + "ndarray/std", + "matrixmultiply/std", + "rand/std", + "rand/std_rng", + "num-traits/std", + "macerator/std", +] +tracing = [ + "burn-autodiff?/tracing", + "burn-std/tracing", + "burn-backend/tracing", + "burn-ir/tracing", +] + +# Serves as a ref impl for some burn-cubecl kernels +export_tests = [] + +[dependencies] + +# ** Please make sure all dependencies support no_std when std is disabled ** + +burn-autodiff = { workspace = true, optional = true } +burn-std = { workspace = true } +burn-ir = { workspace = true } +burn-backend = { workspace = true } + +atomic_float = { workspace = true } +blas-src = { workspace = true, default-features = false, optional = true } # no-std compatible +const-random = { workspace = true } +libm = { workspace = true } +matrixmultiply = { workspace = true, default-features = false } +ndarray = { workspace = true } +num-traits = { workspace = true } +openblas-src = { workspace = true, optional = true } +paste = { workspace = true } +rand = { workspace = true, default-features = false } + +# SIMD +bytemuck = { workspace = true, optional = true } +itertools = { version = "0.14", optional = true } +macerator = { workspace = true, optional = true } +seq-macro = { version = "0.3", optional = true } + +# Parallel +rayon = { workspace = true, optional = true } + +[target.'cfg(not(target_has_atomic = "ptr"))'.dependencies] +portable-atomic = { workspace = true } +portable-atomic-util = { workspace = true } + +[dev-dependencies] +bytes = { workspace = true } + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-ndarray/LICENSE-APACHE b/crates/burn-ndarray/LICENSE-APACHE new file mode 120000 index 0000000..1cd601d --- /dev/null +++ b/crates/burn-ndarray/LICENSE-APACHE @@ -0,0 +1 @@ +../../LICENSE-APACHE \ No newline at end of file diff --git a/crates/burn-ndarray/LICENSE-MIT b/crates/burn-ndarray/LICENSE-MIT new file mode 120000 index 0000000..b2cfbdc --- /dev/null +++ b/crates/burn-ndarray/LICENSE-MIT @@ -0,0 +1 @@ +../../LICENSE-MIT \ No newline at end of file diff --git a/crates/burn-ndarray/README.md b/crates/burn-ndarray/README.md new file mode 100644 index 0000000..788a6c0 --- /dev/null +++ b/crates/burn-ndarray/README.md @@ -0,0 +1,36 @@ +# Burn NdArray + +> [Burn](https://github.com/tracel-ai/burn) ndarray backend + +[![Current Crates.io Version](https://img.shields.io/crates/v/burn-ndarray.svg)](https://crates.io/crates/burn-ndarray) +[![license](https://shields.io/badge/license-MIT%2FApache--2.0-blue)](https://github.com/tracel-ai/burn-ndarray/blob/master/README.md) + +> **New projects should use [`burn-flex`](https://crates.io/crates/burn-flex).** It is a +> from-scratch pure-Rust CPU backend that replaces `burn-ndarray` with faster gemm, zero-copy view +> operations, native quantization, and full support for `std`, `no_std`, and WebAssembly. See +> [`burn-flex/COMPARISON.md`](../burn-flex/COMPARISON.md) for a migration path and +> operation-by-operation benchmarks. + +## Feature Flags + +This crate can be used without the standard library (`#![no_std]`) with `alloc` by disabling the +default `std` feature. + +The following flags support various BLAS options: + +- `blas-accelerate` - Accelerate framework (macOS only) +- `blas-netlib` - Netlib +- `blas-openblas` - OpenBLAS static linked +- `blas-openblas-system` - OpenBLAS from the system + +Note: under the `no_std` mode, the seed is fixed if the seed is not +initialized by `Backend::seed` method. + +### Platform Support + +| Option | CPU | GPU | Linux | MacOS | Windows | Android | iOS | WASM | +| :--------- | :-: | :-: | :---: | :---: | :-----: | :-----: | :-: | :--: | +| Pure Rust | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | +| Accelerate | Yes | No | No | Yes | No | No | Yes | No | +| Netlib | Yes | No | Yes | Yes | Yes | No | No | No | +| Openblas | Yes | No | Yes | Yes | Yes | Yes | Yes | No | diff --git a/crates/burn-ndarray/build.rs b/crates/burn-ndarray/build.rs new file mode 100644 index 0000000..dcb4354 --- /dev/null +++ b/crates/burn-ndarray/build.rs @@ -0,0 +1,6 @@ +fn main() { + // https://github.com/rust-ndarray/ndarray/issues/1197 + if cfg!(feature = "blas-accelerate") { + println!("cargo:rustc-link-lib=framework=Accelerate"); + } +} diff --git a/crates/burn-ndarray/src/backend.rs b/crates/burn-ndarray/src/backend.rs new file mode 100644 index 0000000..07c7664 --- /dev/null +++ b/crates/burn-ndarray/src/backend.rs @@ -0,0 +1,211 @@ +use crate::rand::NdArrayRng; +use crate::{NdArrayQTensor, NdArrayTensor}; +use alloc::string::String; +use burn_backend::quantization::{QuantLevel, QuantMode, QuantScheme, QuantStore, QuantValue}; +use burn_backend::tensor::{BoolTensor, FloatTensor, IntTensor, QuantizedTensor}; +use burn_backend::{Backend, BackendTypes, DType, DeviceId, DeviceOps}; +use burn_ir::{BackendIr, HandleKind, TensorHandle}; +use burn_std::stub::Mutex; +use burn_std::{BoolStore, DeviceSettings, QuantConfig}; +use rand::SeedableRng; + +pub(crate) static SEED: Mutex> = Mutex::new(None); + +/// The device type for the ndarray backend. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +pub enum NdArrayDevice { + /// The CPU device. + #[default] + Cpu, +} + +impl DeviceOps for NdArrayDevice { + fn defaults(&self) -> DeviceSettings { + // E = f32, I = i64 + DeviceSettings::new( + DType::F32, + DType::I64, + DType::Bool(BoolStore::Native), + QuantConfig::new( + QuantScheme::default().with_store(QuantStore::Native), + Default::default(), + ), + ) + } +} + +impl burn_backend::Device for NdArrayDevice { + fn from_id(_device_id: DeviceId) -> Self { + Self::Cpu + } + + fn to_id(&self) -> DeviceId { + DeviceId { + type_id: 0, + index_id: 0, + } + } +} + +/// Tensor backend that uses the [ndarray](ndarray) crate for executing tensor operations. +/// +/// This backend is compatible with CPUs and can be compiled for almost any platform, including +/// `wasm`, `arm`, and `x86`. +#[derive(Clone, Copy, Default, Debug)] +pub struct NdArray; + +impl BackendTypes for NdArray { + type Device = NdArrayDevice; + + type FloatTensorPrimitive = NdArrayTensor; + type IntTensorPrimitive = NdArrayTensor; + type BoolTensorPrimitive = NdArrayTensor; + type QuantizedTensorPrimitive = NdArrayQTensor; + + type GraphPrimitive = burn_backend::GraphUnsupported; +} + +impl Backend for NdArray { + fn ad_enabled(_device: &Self::Device) -> bool { + false + } + + fn name(_device: &Self::Device) -> String { + String::from("ndarray") + } + + fn seed(_device: &Self::Device, seed: u64) { + let rng = NdArrayRng::seed_from_u64(seed); + let mut seed = SEED.lock().unwrap(); + *seed = Some(rng); + } + + fn dtype_usage(_device: &Self::Device, dtype: DType) -> burn_backend::DTypeUsageSet { + match dtype { + DType::F64 + | DType::F32 + | DType::Flex32 + | DType::I64 + | DType::I32 + | DType::I16 + | DType::I8 + | DType::U64 + | DType::U32 + | DType::U16 + | DType::U8 + | DType::Bool(BoolStore::Native) => burn_backend::DTypeUsage::general(), + DType::F16 | DType::BF16 | DType::Bool(_) => burn_backend::DTypeUsageSet::empty(), + DType::QFloat(scheme) => { + match scheme { + QuantScheme { + level: QuantLevel::Tensor | QuantLevel::Block(_), + mode: QuantMode::Symmetric, + #[cfg(not(feature = "export_tests"))] + value: QuantValue::Q8F | QuantValue::Q8S, + // For tests, "native" sub-byte quant serves as a reference for value equality. + // Values are stored as i8 regardless. + #[cfg(feature = "export_tests")] + value: + QuantValue::Q8F + | QuantValue::Q8S + | QuantValue::Q4F + | QuantValue::Q4S + | QuantValue::Q2F + | QuantValue::Q2S, + store: QuantStore::Native, + .. + } => burn_backend::DTypeUsage::general(), + _scheme => burn_backend::DTypeUsageSet::empty(), + } + } + } + } + + fn device_count(_: u16) -> usize { + 1 + } + + fn flush(_device: &Self::Device) {} +} + +impl BackendIr for NdArray { + type Handle = HandleKind; + + fn float_tensor(handle: TensorHandle) -> FloatTensor { + match handle.handle { + HandleKind::Float(handle) => handle, + _ => panic!("Expected float handle, got {}", handle.handle.name()), + } + } + + fn int_tensor(handle: TensorHandle) -> IntTensor { + match handle.handle { + HandleKind::Int(handle) => handle, + _ => panic!("Expected int handle, got {}", handle.handle.name()), + } + } + + fn bool_tensor(handle: TensorHandle) -> BoolTensor { + match handle.handle { + HandleKind::Bool(handle) => handle, + _ => panic!("Expected bool handle, got {}", handle.handle.name()), + } + } + + fn quantized_tensor(handle: TensorHandle) -> QuantizedTensor { + match handle.handle { + HandleKind::Quantized(handle) => handle, + _ => panic!("Expected quantized handle, got {}", handle.handle.name()), + } + } + + fn float_tensor_handle(tensor: FloatTensor) -> Self::Handle { + HandleKind::Float(tensor) + } + + fn int_tensor_handle(tensor: IntTensor) -> Self::Handle { + HandleKind::Int(tensor) + } + + fn bool_tensor_handle(tensor: BoolTensor) -> Self::Handle { + HandleKind::Bool(tensor) + } + + fn quantized_tensor_handle(tensor: QuantizedTensor) -> Self::Handle { + HandleKind::Quantized(tensor) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn should_support_dtypes() { + type B = NdArray; + let device = NdArrayDevice::Cpu; + let scheme = device.defaults().quantization.scheme; + + assert!(B::supports_dtype(&device, DType::F64)); + assert!(B::supports_dtype(&device, DType::F32)); + assert!(B::supports_dtype(&device, DType::Flex32)); + assert!(B::supports_dtype(&device, DType::I64)); + assert!(B::supports_dtype(&device, DType::I32)); + assert!(B::supports_dtype(&device, DType::I16)); + assert!(B::supports_dtype(&device, DType::I8)); + assert!(B::supports_dtype(&device, DType::U64)); + assert!(B::supports_dtype(&device, DType::U32)); + assert!(B::supports_dtype(&device, DType::U16)); + assert!(B::supports_dtype(&device, DType::U8)); + assert!(B::supports_dtype(&device, DType::Bool(BoolStore::Native))); + assert!(B::supports_dtype(&device, DType::QFloat(scheme))); + + assert!(!B::supports_dtype(&device, DType::F16)); + assert!(!B::supports_dtype(&device, DType::BF16)); + // QuantStore::U32 not supported + assert!(!B::supports_dtype( + &device, + DType::QFloat(QuantScheme::default()) + )); + } +} diff --git a/crates/burn-ndarray/src/element.rs b/crates/burn-ndarray/src/element.rs new file mode 100644 index 0000000..8485352 --- /dev/null +++ b/crates/burn-ndarray/src/element.rs @@ -0,0 +1,207 @@ +use burn_backend::Element; +use num_traits::Signed; + +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float; + +use num_traits::Pow; + +use libm::{log1p, log1pf}; + +/// A float element for ndarray backend. +pub trait FloatNdArrayElement: NdArrayElement + Signed + core::cmp::PartialOrd +where + Self: Sized, +{ +} + +/// An int element for ndarray backend. +pub trait IntNdArrayElement: NdArrayElement + core::cmp::PartialOrd {} + +/// A general element for ndarray backend. +pub trait NdArrayElement: + Element + + ndarray::LinalgScalar + + ndarray::ScalarOperand + + ExpElement + + AddAssignElement + + num_traits::FromPrimitive + + core::ops::AddAssign + + core::cmp::PartialEq + + core::ops::Rem +{ +} + +/// A element for ndarray backend that supports exp ops. +pub trait ExpElement { + /// Exponent + fn exp_elem(self) -> Self; + /// Log + fn log_elem(self) -> Self; + /// Log1p + fn log1p_elem(self) -> Self; + /// Powf + fn powf_elem(self, value: f32) -> Self; + /// Powi + fn powi_elem(self, value: i32) -> Self; + /// Sqrt + fn sqrt_elem(self) -> Self; + /// Abs + fn abs_elem(self) -> Self; +} + +/// The addition assignment operator implemented for ndarray elements. +pub trait AddAssignElement { + /// Performs the addition assignment operation. + /// + /// For `bool`, this corresponds to logical OR assignment. + fn add_assign(&mut self, rhs: Rhs); +} + +impl AddAssignElement for E { + fn add_assign(&mut self, rhs: Self) { + *self += rhs; + } +} + +impl AddAssignElement for bool { + fn add_assign(&mut self, rhs: Self) { + *self = *self || rhs; // logical OR for bool + } +} + +/// A quantized element for the ndarray backend. +pub trait QuantElement: NdArrayElement {} + +impl QuantElement for i8 {} + +impl FloatNdArrayElement for f64 {} +impl FloatNdArrayElement for f32 {} + +impl IntNdArrayElement for i64 {} +impl IntNdArrayElement for i32 {} +impl IntNdArrayElement for i16 {} +impl IntNdArrayElement for i8 {} + +impl IntNdArrayElement for u64 {} +impl IntNdArrayElement for u32 {} +impl IntNdArrayElement for u16 {} +impl IntNdArrayElement for u8 {} + +macro_rules! make_float { + ( + $ty:ty, + $log1p:expr + ) => { + impl NdArrayElement for $ty {} + + #[allow(clippy::cast_abs_to_unsigned)] + impl ExpElement for $ty { + #[inline(always)] + fn exp_elem(self) -> Self { + self.exp() + } + + #[inline(always)] + fn log_elem(self) -> Self { + self.ln() + } + + #[inline(always)] + fn log1p_elem(self) -> Self { + $log1p(self) + } + + #[inline(always)] + fn powf_elem(self, value: f32) -> Self { + self.pow(value) + } + + #[inline(always)] + fn powi_elem(self, value: i32) -> Self { + #[cfg(feature = "std")] + let val = self.powi(value); + + #[cfg(not(feature = "std"))] + let val = Self::powf_elem(self, value as f32); + + val + } + + #[inline(always)] + fn sqrt_elem(self) -> Self { + self.sqrt() + } + + #[inline(always)] + fn abs_elem(self) -> Self { + self.abs() + } + } + }; +} +macro_rules! make_int { + ( + $ty:ty, + $abs:expr + ) => { + impl NdArrayElement for $ty {} + + #[allow(clippy::cast_abs_to_unsigned)] + impl ExpElement for $ty { + #[inline(always)] + fn exp_elem(self) -> Self { + (self as f32).exp() as $ty + } + + #[inline(always)] + fn log_elem(self) -> Self { + (self as f32).ln() as $ty + } + + #[inline(always)] + fn log1p_elem(self) -> Self { + log1pf(self as f32) as $ty + } + + #[inline(always)] + fn powf_elem(self, value: f32) -> Self { + (self as f32).pow(value) as $ty + } + + #[inline(always)] + fn powi_elem(self, value: i32) -> Self { + #[cfg(feature = "std")] + let val = f32::powi(self as f32, value) as $ty; + + #[cfg(not(feature = "std"))] + let val = Self::powf_elem(self, value as f32); + + val + } + + #[inline(always)] + fn sqrt_elem(self) -> Self { + (self as f32).sqrt() as $ty + } + + #[inline(always)] + fn abs_elem(self) -> Self { + $abs(self) + } + } + }; +} + +make_float!(f64, log1p); +make_float!(f32, log1pf); + +make_int!(i64, i64::wrapping_abs); +make_int!(i32, i32::wrapping_abs); +make_int!(i16, i16::wrapping_abs); +make_int!(i8, i8::wrapping_abs); +make_int!(u64, |x| x); +make_int!(u32, |x| x); +make_int!(u16, |x| x); +make_int!(u8, |x| x); diff --git a/crates/burn-ndarray/src/lib.rs b/crates/burn-ndarray/src/lib.rs new file mode 100644 index 0000000..34a4625 --- /dev/null +++ b/crates/burn-ndarray/src/lib.rs @@ -0,0 +1,29 @@ +#![cfg_attr(not(feature = "std"), no_std)] +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] + +//! Burn ndarray backend. + +#[cfg(any( + feature = "blas-netlib", + feature = "blas-openblas", + feature = "blas-openblas-system", +))] +extern crate blas_src; + +mod backend; +mod element; +mod ops; +mod parallel; +mod rand; +mod sharing; +mod storage; +mod tensor; + +pub use backend::*; +pub use element::*; +pub(crate) use sharing::*; +pub(crate) use storage::*; +pub use tensor::*; + +extern crate alloc; diff --git a/crates/burn-ndarray/src/ops/activation.rs b/crates/burn-ndarray/src/ops/activation.rs new file mode 100644 index 0000000..d5c856e --- /dev/null +++ b/crates/burn-ndarray/src/ops/activation.rs @@ -0,0 +1,8 @@ +use crate::{NdArray, execute_with_numeric_dtype, ops::NdArrayMathOps}; +use burn_backend::{ElementConversion, TensorMetadata, ops::ActivationOps, tensor::FloatTensor}; + +impl ActivationOps for NdArray { + fn relu(tensor: FloatTensor) -> FloatTensor { + execute_with_numeric_dtype!(tensor, |array| NdArrayMathOps::clamp_min(array, 0.elem())) + } +} diff --git a/crates/burn-ndarray/src/ops/adaptive_avgpool.rs b/crates/burn-ndarray/src/ops/adaptive_avgpool.rs new file mode 100644 index 0000000..baaee09 --- /dev/null +++ b/crates/burn-ndarray/src/ops/adaptive_avgpool.rs @@ -0,0 +1,103 @@ +use crate::{ + SharedArray, element::FloatNdArrayElement, iter_range_par, run_par, sharing::UnsafeSharedRef, +}; +use burn_backend::ElementConversion; +use ndarray::Array4; + +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float; + +pub(crate) fn adaptive_avg_pool2d( + x: SharedArray, + output_size: [usize; 2], +) -> SharedArray { + let [batch_size, channels, input_height, input_width] = x.shape().try_into().unwrap(); + + let mut output = Array4::from_elem( + (batch_size, channels, output_size[0], output_size[1]), + 0.elem(), + ); + let unsafe_shared_out = UnsafeSharedRef::new(&mut output); + + run_par!(|| { + iter_range_par!(0, batch_size * channels).for_each(|k| unsafe { + let b = k / channels; + let c = k % channels; + + let output = unsafe_shared_out.get(); + for h in 0..output_size[0] { + for w in 0..output_size[1] { + let ih_start = start_index(h, output_size[0], input_height); + let ih_end = end_index(h, output_size[0], input_height); + let iw_start = start_index(w, output_size[1], input_width); + let iw_end = end_index(w, output_size[1], input_width); + + let mut sum_val: E = 0.elem(); + + for ih in ih_start..ih_end { + for iw in iw_start..iw_end { + sum_val += x[[b, c, ih, iw]]; + } + } + + let count: E = (((ih_end - ih_start) * (iw_end - iw_start)) as i32).elem(); + output[[b, c, h, w]] = sum_val / count.elem(); + } + } + }) + }); + + output.into_dyn().into_shared() +} + +pub(crate) fn adaptive_avg_pool2d_backward( + x: SharedArray, + grad: SharedArray, +) -> SharedArray { + let [_, _, input_height, input_width] = x.shape().try_into().unwrap(); + let [batch_size, channels, output_height, output_width] = grad.shape().try_into().unwrap(); + + let mut output_grad = + Array4::from_elem((batch_size, channels, input_height, input_width), 0.elem()); + let unsafe_shared_out = UnsafeSharedRef::new(&mut output_grad); + + run_par!(|| { + iter_range_par!(0, batch_size * channels).for_each(|k| unsafe { + let b = k / channels; + let c = k % channels; + + let output_grad = unsafe_shared_out.get(); + for oh in 0..output_height { + for ow in 0..output_width { + let ih_start = start_index(oh, output_height, input_height); + let ih_end = end_index(oh, output_height, input_height); + + let iw_start = start_index(ow, output_width, input_width); + let iw_end = end_index(ow, output_width, input_width); + + let count: E = (((ih_end - ih_start) * (iw_end - iw_start)) as i32).elem(); + + for ih in ih_start..ih_end { + for iw in iw_start..iw_end { + output_grad[[b, c, ih, iw]] += grad[[b, c, oh, ow]] / count.elem(); + } + } + } + } + }) + }); + + output_grad.into_dyn().into_shared() +} + +fn start_index(output_size_index: usize, output_size: usize, input_size: usize) -> usize { + ((output_size_index as f32 * input_size as f32) / output_size as f32).floor() as usize +} + +fn end_index(output_size_index: usize, output_size: usize, input_size: usize) -> usize { + let index = + (((output_size_index + 1) as f32 * input_size as f32) / output_size as f32).ceil() as usize; + + usize::min(index, input_size) +} diff --git a/crates/burn-ndarray/src/ops/avgpool.rs b/crates/burn-ndarray/src/ops/avgpool.rs new file mode 100644 index 0000000..4d015dd --- /dev/null +++ b/crates/burn-ndarray/src/ops/avgpool.rs @@ -0,0 +1,172 @@ +use crate::{ + SharedArray, element::FloatNdArrayElement, iter_range_par, run_par, sharing::UnsafeSharedRef, +}; + +use burn_backend::ElementConversion; +use burn_backend::ops::conv::calculate_pool_output_size; +use ndarray::Array4; + +pub(crate) fn avg_pool2d( + x: SharedArray, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, +) -> SharedArray { + let [kernel_height, kernel_width] = kernel_size; + let [padding_height, padding_width] = padding; + let [stride_height, stride_width] = stride; + let [batch_size, channels, x_height, x_width] = x.shape().try_into().unwrap(); + + let out_height = calculate_pool_output_size( + kernel_height, + stride_height, + padding_height, + 1, + x_height, + ceil_mode, + ); + let out_width = calculate_pool_output_size( + kernel_width, + stride_width, + padding_width, + 1, + x_width, + ceil_mode, + ); + + // Padded input bounds (for count_include_pad calculation) + let padded_height = x_height + 2 * padding_height; + let padded_width = x_width + 2 * padding_width; + + let mut output = Array4::from_elem((batch_size, channels, out_height, out_width), 0.elem()); + let unsafe_shared_out = UnsafeSharedRef::new(&mut output); + + run_par!(|| { + iter_range_par!(0, batch_size * channels).for_each(|k| unsafe { + let b = k / channels; + let c = k % channels; + + let output = unsafe_shared_out.get(); + + for oh in 0..out_height { + for ow in 0..out_width { + let mut sum_val: E = 0.elem(); + let mut valid_count = 0usize; + let mut padded_count = 0usize; + + for kh in 0..kernel_height { + let ih = oh * stride_height + kh; + + for kw in 0..kernel_width { + let iw = ow * stride_width + kw; + + // Check if within padded bounds (excludes ceil_mode extensions) + if ih < padded_height && iw < padded_width { + padded_count += 1; + + // Check if within valid (non-padding) input bounds + if ih >= padding_height + && ih < x_height + padding_height + && iw >= padding_width + && iw < x_width + padding_width + { + let ih_valid = ih - padding_height; + let iw_valid = iw - padding_width; + sum_val += x[[b, c, ih_valid, iw_valid]]; + valid_count += 1; + } + } + } + } + + // count_include_pad: count positions within padded bounds (not ceil_mode extensions) + // !count_include_pad: count only valid (non-padding) positions + let count: E = if count_include_pad { + (padded_count as i32).elem() + } else { + (valid_count as i32).elem() + }; + + output[[b, c, oh, ow]] = sum_val / count; + } + } + }) + }); + + output.into_dyn().into_shared() +} + +pub(crate) fn avg_pool2d_backward( + x: SharedArray, + grad: SharedArray, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + _ceil_mode: bool, +) -> SharedArray { + let [kernel_height, kernel_width] = kernel_size; + let [stride_height, stride_width] = stride; + let [padding_height, padding_width] = padding; + let [batch_size, channels, x_height, x_width] = x.shape().try_into().unwrap(); + let [_batch_size, _channels, out_height, out_width] = grad.shape().try_into().unwrap(); + + // Padded input bounds (for count_include_pad calculation) + let padded_height = x_height + 2 * padding_height; + let padded_width = x_width + 2 * padding_width; + + let mut output_grad = Array4::from_elem((batch_size, channels, x_height, x_width), 0.elem()); + let unsafe_shared_grad = UnsafeSharedRef::new(&mut output_grad); + + run_par!(|| { + iter_range_par!(0, batch_size * channels).for_each(|k| unsafe { + let b = k / channels; + let c = k % channels; + + let output_grad = unsafe_shared_grad.get(); + + for oh in 0..out_height { + for ow in 0..out_width { + let ih_start_kernel = oh * stride_height; + let iw_start_kernel = ow * stride_width; + + let ih_end_kernel = ih_start_kernel + kernel_height; + let iw_end_kernel = iw_start_kernel + kernel_width; + + // Clip to valid input bounds (for gradient distribution) + let ih_start = usize::max(ih_start_kernel, padding_height); + let iw_start = usize::max(iw_start_kernel, padding_width); + let ih_end = usize::min(ih_end_kernel, x_height + padding_height); + let iw_end = usize::min(iw_end_kernel, x_width + padding_width); + + // Calculate count based on count_include_pad + let count = if count_include_pad { + // Count positions within padded bounds (not ceil_mode extensions) + let ih_start_padded = ih_start_kernel; + let iw_start_padded = iw_start_kernel; + let ih_end_padded = usize::min(ih_end_kernel, padded_height); + let iw_end_padded = usize::min(iw_end_kernel, padded_width); + (ih_end_padded - ih_start_padded) * (iw_end_padded - iw_start_padded) + } else { + // Count only valid (non-padding) positions + (ih_end - ih_start) * (iw_end - iw_start) + }; + + for ih in ih_start..ih_end { + for iw in iw_start..iw_end { + let ih = ih - padding_height; + let iw = iw - padding_width; + + output_grad[[b, c, ih, iw]] += + grad[[b, c, oh, ow]] / (count as i32).elem(); + } + } + } + } + }) + }); + + output_grad.into_dyn().into_shared() +} diff --git a/crates/burn-ndarray/src/ops/base.rs b/crates/burn-ndarray/src/ops/base.rs new file mode 100644 index 0000000..d34fbcd --- /dev/null +++ b/crates/burn-ndarray/src/ops/base.rs @@ -0,0 +1,1665 @@ +use alloc::{vec, vec::Vec}; +use burn_backend::element::{Element, ElementConversion}; +#[cfg(feature = "simd")] +use burn_backend::{DType, quantization::QuantValue}; +use core::fmt::Debug; +use core::marker::PhantomData; +use ndarray::Dimension; +use ndarray::IntoDimension; +use ndarray::SliceInfo; +use ndarray::Zip; +use ndarray::s; +use ndarray::{Array2, ArrayD}; +use num_traits::Signed; +#[cfg(feature = "simd")] +use paste::paste; + +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float; + +#[cfg(feature = "simd")] +use crate::ops::simd::{ + binary::try_binary_simd, + binary_elemwise::{ + VecAdd, VecBitAnd, VecBitOr, VecBitXor, VecClamp, VecDiv, VecMax, VecMin, VecMul, VecSub, + try_binary_scalar_simd, + }, + cmp::{ + VecEquals, VecGreater, VecGreaterEq, VecLower, VecLowerEq, try_cmp_scalar_simd, + try_cmp_simd, + }, + unary::{RecipVec, VecAbs, VecBitNot, try_unary_simd}, +}; +use crate::reshape; +use crate::{ + IntNdArrayElement, ShapeOps, + ops::macros::{ + cummax_dim, cummin_dim, cumprod_dim, cumsum_dim, keepdim, mean_dim, prod_dim, sum_dim, + }, +}; +use crate::{SharedArray, element::NdArrayElement}; +use burn_backend::ops::unfold::calculate_unfold_shape; +use burn_backend::{Shape, Slice}; +use ndarray::ArrayView; +use ndarray::Axis; +use ndarray::Dim; +use ndarray::IxDyn; +use ndarray::SliceInfoElem; + +pub struct NdArrayOps { + e: PhantomData, +} + +pub(crate) struct NdArrayMathOps { + e: PhantomData, +} + +impl NdArrayOps +where + E: Copy + Debug + Element + crate::AddAssignElement, +{ + pub fn slice(tensor: ArrayView, slices: &[Slice]) -> SharedArray { + let slices = Self::to_slice_args_with_steps(slices, tensor.shape().num_dims()); + tensor.slice_move(slices.as_slice()).to_shared() + } + + pub fn slice_assign( + tensor: SharedArray, + slices: &[Slice], + value: SharedArray, + ) -> SharedArray { + let slices = Self::to_slice_args_with_steps(slices, tensor.shape().num_dims()); + let mut array = tensor.into_owned(); + array.slice_mut(slices.as_slice()).assign(&value); + array.into_shared() + } + + pub fn mask_where( + tensor: SharedArray, + mask: SharedArray, + source: SharedArray, + ) -> SharedArray { + let tensor = tensor.broadcast(mask.dim()).unwrap(); + let source = source.broadcast(mask.dim()).unwrap(); + Zip::from(&tensor) + .and(&mask) + .and(&source) + .map_collect(|&x, &mask_val, &y| if mask_val { y } else { x }) + .into_shared() + } + + pub fn mask_fill(tensor: SharedArray, mask: SharedArray, value: E) -> SharedArray { + // Use into_owned() instead of clone() - only copies if shared, avoids copy if unique + let mut output = tensor.into_owned(); + let broadcast_mask = mask.broadcast(output.dim()).unwrap(); + Zip::from(&mut output) + .and(&broadcast_mask) + .for_each(|out, &mask_val| { + if mask_val { + *out = value; + } + }); + output.into_shared() + } + + pub fn gather( + dim: usize, + mut tensor: SharedArray, + mut indices: SharedArray, + ) -> SharedArray { + let ndims = tensor.shape().num_dims(); + if dim != ndims - 1 { + tensor.swap_axes(ndims - 1, dim); + indices.swap_axes(ndims - 1, dim); + } + let (shape_tensor, shape_indices) = (tensor.shape(), indices.shape().into_shape()); + let (size_tensor, size_index) = (shape_tensor[ndims - 1], shape_indices[ndims - 1]); + let batch_size = Self::gather_batch_size(shape_tensor, &shape_indices); + + let indices = NdArrayOps::reshape(indices, Shape::new([batch_size, size_index])); + let tensor = NdArrayOps::reshape(tensor, Shape::new([batch_size, size_tensor])); + let mut output = Array2::from_elem((batch_size, size_index), 0.elem::()); + + for b in 0..batch_size { + let indices = indices.slice(s!(b, ..)); + for (i, index) in indices.iter().enumerate() { + output[[b, i]] = tensor[[b, index.elem::() as usize]]; + } + } + + let mut output = NdArrayOps::reshape(output.into_shared().into_dyn(), shape_indices); + + if dim != ndims - 1 { + output.swap_axes(ndims - 1, dim); + } + + output + } + + pub fn scatter( + dim: usize, + mut tensor: SharedArray, + mut indices: SharedArray, + mut value: SharedArray, + ) -> SharedArray { + let ndims = tensor.shape().num_dims(); + if dim != ndims - 1 { + tensor.swap_axes(ndims - 1, dim); + indices.swap_axes(ndims - 1, dim); + value.swap_axes(ndims - 1, dim); + } + + let (shape_tensor, shape_indices, shape_value) = + (tensor.shape().into_shape(), indices.shape(), value.shape()); + let (size_tensor, size_index, size_value) = ( + shape_tensor[ndims - 1], + shape_indices[ndims - 1], + shape_value[ndims - 1], + ); + let batch_size = Self::gather_batch_size(&shape_tensor, shape_indices); + + if shape_value != shape_indices { + panic!( + "Invalid dimension: the shape of the index tensor should be the same as the value \ + tensor: Index {:?} value {:?}", + shape_indices, shape_value + ); + } + + let indices = NdArrayOps::reshape(indices, Shape::new([batch_size, size_index])); + let value = NdArrayOps::reshape(value, Shape::new([batch_size, size_value])); + let mut tensor = NdArrayOps::reshape(tensor, Shape::new([batch_size, size_tensor])); + + for b in 0..batch_size { + let indices = indices.slice(s!(b, ..)); + + for (i, index) in indices.iter().enumerate() { + let index = index.elem::() as usize; + tensor[[b, index]].add_assign(value[[b, i]]); + } + } + + let mut output = NdArrayOps::reshape(tensor.into_shared().into_dyn(), shape_tensor); + if dim != ndims - 1 { + output.swap_axes(ndims - 1, dim); + } + output + } + + pub fn scatter_nd( + data: SharedArray, + indices: SharedArray, + values: SharedArray, + reduction: burn_backend::tensor::IndexingUpdateOp, + ) -> SharedArray + where + E: core::ops::Mul + PartialOrd, + { + use burn_backend::tensor::IndexingUpdateOp; + + let data_shape: Vec = data.shape().to_vec(); + let idx_shape: Vec = indices.shape().to_vec(); + let m = idx_shape.len(); + let k = idx_shape[m - 1]; + + // Number of index tuples = product of batch dims (first M-1 dims of indices) + let num_indices: usize = idx_shape[..m - 1].iter().product(); + // Size of each slice to scatter = product of data.shape[K..] + let slice_size: usize = data_shape[k..].iter().product(); + + let mut output = data.into_owned(); + let output_flat = output + .as_slice_mut() + .expect("ndarray scatter_nd requires contiguous data"); + + // Flatten indices to [num_indices, K] + let idx_flat = indices + .as_slice() + .expect("ndarray scatter_nd requires contiguous indices"); + + // Flatten values to [num_indices, slice_size] + let val_flat = values + .as_slice() + .expect("ndarray scatter_nd requires contiguous values"); + + let strides: Vec = { + let mut s = vec![0usize; k]; + if k > 0 { + s[k - 1] = slice_size; + for i in (0..k - 1).rev() { + s[i] = s[i + 1] * data_shape[i + 1]; + } + } + s + }; + + for n in 0..num_indices { + // Compute flat base offset from the K-dimensional index + let mut base_offset = 0usize; + for j in 0..k { + let idx_val = idx_flat[n * k + j].elem::() as usize; + base_offset += idx_val * strides[j]; + } + + // Apply reduction over the slice + let val_offset = n * slice_size; + match reduction { + IndexingUpdateOp::Assign => { + output_flat[base_offset..(base_offset + slice_size)] + .copy_from_slice(&val_flat[val_offset..(val_offset + slice_size)]); + } + IndexingUpdateOp::Add => { + for s in 0..slice_size { + output_flat[base_offset + s].add_assign(val_flat[val_offset + s]); + } + } + IndexingUpdateOp::Mul => { + for s in 0..slice_size { + output_flat[base_offset + s] = + output_flat[base_offset + s] * val_flat[val_offset + s]; + } + } + IndexingUpdateOp::Min => { + for s in 0..slice_size { + let b = val_flat[val_offset + s]; + if b < output_flat[base_offset + s] { + output_flat[base_offset + s] = b; + } + } + } + IndexingUpdateOp::Max => { + for s in 0..slice_size { + let b = val_flat[val_offset + s]; + if b > output_flat[base_offset + s] { + output_flat[base_offset + s] = b; + } + } + } + } + } + + output.into_shared() + } + + pub fn gather_nd( + data: SharedArray, + indices: SharedArray, + ) -> SharedArray { + let data_shape: Vec = data.shape().to_vec(); + let idx_shape: Vec = indices.shape().to_vec(); + let m = idx_shape.len(); + let k = idx_shape[m - 1]; + + // Number of index tuples + let num_indices: usize = idx_shape[..m - 1].iter().product(); + // Size of each output slice + let slice_size: usize = data_shape[k..].iter().product(); + + // Output shape: idx_shape[..m-1] ++ data_shape[k..] + let mut out_shape_vec: Vec = idx_shape[..m - 1].to_vec(); + out_shape_vec.extend_from_slice(&data_shape[k..]); + let out_total = num_indices * slice_size; + + let data_flat = data + .as_slice() + .expect("ndarray gather_nd requires contiguous data"); + + let idx_flat = indices + .as_slice() + .expect("ndarray gather_nd requires contiguous indices"); + + let strides: Vec = { + let mut s = vec![0usize; k]; + if k > 0 { + s[k - 1] = slice_size; + for i in (0..k - 1).rev() { + s[i] = s[i + 1] * data_shape[i + 1]; + } + } + s + }; + + let mut output_vec: Vec = vec![0.elem::(); out_total]; + + for n in 0..num_indices { + let mut base_offset = 0usize; + for j in 0..k { + let idx_val = idx_flat[n * k + j].elem::() as usize; + base_offset += idx_val * strides[j]; + } + + let out_offset = n * slice_size; + output_vec[out_offset..(out_offset + slice_size)] + .copy_from_slice(&data_flat[base_offset..(base_offset + slice_size)]); + } + + let out_shape = Shape::from(out_shape_vec); + let output = ArrayD::from_shape_vec(out_shape.as_slice(), output_vec) + .expect("gather_nd: shape mismatch"); + + output.into_shared() + } + + fn gather_batch_size(shape_tensor: &[usize], shape_indices: &[usize]) -> usize { + let ndims = shape_tensor.num_dims(); + let mut batch_size = 1; + + for i in 0..ndims - 1 { + if shape_tensor[i] != shape_indices[i] { + panic!( + "Unsupported dimension, only the last dimension can differ: Tensor {:?} Index \ + {:?}", + shape_tensor, shape_indices + ); + } + batch_size *= shape_indices[i]; + } + + batch_size + } + + pub fn reshape(tensor: SharedArray, shape: Shape) -> SharedArray { + reshape!( + ty E, + shape shape, + array tensor, + d shape.num_dims() + ) + } + + pub(crate) fn concatenate( + arrays: &[ndarray::ArrayView], + dim: usize, + ) -> SharedArray { + let array = ndarray::concatenate(Axis(dim), arrays) + .unwrap() + .into_shared(); + + // Transform column-major layout into row-major (standard) layout. (fix #1053) + // Get shape first (via reference), then pass ownership to avoid clone + let shape = array.shape().into_shape(); + Self::reshape(array, shape) + } + + pub fn cat(tensors: Vec>, dim: usize) -> SharedArray { + let arrays: Vec<_> = tensors.iter().map(|t| t.view()).collect(); + Self::concatenate(&arrays, dim) + } + + #[allow(clippy::wrong_self_convention)] + fn to_slice_args_with_steps( + burn_slices: &[burn_backend::Slice], + ndims: usize, + ) -> Vec { + let mut slices = vec![SliceInfoElem::NewAxis; ndims]; + + for i in 0..ndims { + slices[i] = if i < burn_slices.len() { + let slice = &burn_slices[i]; + + // Check for empty range (would result in no elements) + if let Some(end) = slice.end + && slice.start == end + { + SliceInfoElem::Slice { + start: 0, + end: Some(0), + step: 1, + } + } else { + // Pass slice parameters directly to ndarray + // ndarray handles both positive and negative steps correctly: + // - Positive step: iterates forward from start + // - Negative step: iterates backward from the last element in range + SliceInfoElem::Slice { + start: slice.start, + end: slice.end, + step: slice.step, + } + } + } else { + // Dimension not specified in slices - use full range + SliceInfoElem::Slice { + start: 0, + end: None, + step: 1, + } + } + } + + slices + } + + pub fn swap_dims(mut tensor: SharedArray, dim1: usize, dim2: usize) -> SharedArray { + tensor.swap_axes(dim1, dim2); + + tensor + } + + pub fn permute(tensor: SharedArray, axes: &[usize]) -> SharedArray { + tensor.permuted_axes(axes.into_dimension()) + } + + /// Broadcasts the tensor to the given shape + pub(crate) fn expand(tensor: SharedArray, shape: Shape) -> SharedArray { + tensor + .broadcast(shape.into_dimension()) + .expect("The shapes should be broadcastable") + // need to convert view to owned array because NdArrayTensor expects owned array + // and try_into_owned_nocopy() panics for broadcasted arrays (zero strides) + .into_owned() + .into_shared() + } + + pub fn flip(tensor: SharedArray, axes: &[usize]) -> SharedArray { + let slice_items: Vec<_> = (0..tensor.shape().num_dims()) + .map(|i| { + if axes.contains(&i) { + SliceInfoElem::Slice { + start: 0, + end: None, + step: -1, + } + } else { + SliceInfoElem::Slice { + start: 0, + end: None, + step: 1, + } + } + }) + .collect(); + let slice_info = + SliceInfo::, IxDyn, IxDyn>::try_from(slice_items).unwrap(); + tensor.slice(slice_info).into_owned().into_shared() + } + + /// Unfold windows along a dimension. + /// + /// # Warning + /// + /// This is a copy impl; `ndarray` doesn't expose the layout machinery + /// necessary to build the stride view. + /// + /// Returns a copy of the tensor with all complete windows of size `size` in dimension `dim`; + /// where windows are advanced by `step` at each index. + /// + /// The number of windows is `max(0, (shape[dim] - size).ceil_div(step))`. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor to unfold; of shape ``[pre=..., dim shape, post=...]`` + /// * `dim` - the dimension to unfold. + /// * `size` - the size of each unfolded window. + /// * `step` - the step between each window. + /// + /// # Returns + /// + /// A tensor view with shape ``[pre=..., windows, post=..., size]``. + #[allow(unused)] + pub(crate) fn unfold( + tensor: SharedArray, + dim: usize, + size: usize, + step: usize, + ) -> SharedArray { + let result_shape = calculate_unfold_shape(tensor.shape(), dim, size, step); + let windows = result_shape[dim]; + + let mut slices = vec![Slice::new(0, None, 1); tensor.shape().len()]; + let new_axis = slices.len(); + + let mut stack = Vec::with_capacity(windows); + for widx in 0..windows { + let start = widx * step; + let end = start + size; + slices[dim] = Slice::new(start as isize, Some(end as isize), 1); + + let mut window_slice = + tensor.slice(Self::to_slice_args_with_steps(&slices, slices.len()).as_slice()); + window_slice.insert_axis_inplace(Axis(new_axis)); + window_slice.swap_axes(dim, new_axis); + + stack.push(window_slice); + } + Self::concatenate(&stack, dim) + } +} + +#[cfg(feature = "simd")] +macro_rules! dispatch_binary_simd { + (noq, $elem: ty, $op: ty, $lhs: expr, $rhs: expr, $($ty: ty),*) => {{ + paste! { + let simd = match $elem::dtype() { + $(DType::[<$ty:upper>] => try_binary_simd::<$elem, $elem, $ty, $ty, $op>($lhs, $rhs),)* + _ => Err(($lhs, $rhs)), + }; + match simd { + Ok(out) => return out, + Err(args) => args, + } + } + }}; + ($elem: ty, $op: ty, $lhs: expr, $rhs: expr, $($ty: ty),*) => {{ + paste! { + let simd = match $elem::dtype() { + $(DType::[<$ty:upper>] => try_binary_simd::<$elem, $elem, $ty, $ty, $op>($lhs, $rhs),)* + DType::QFloat(strategy) => match strategy.value { + QuantValue::Q8F | QuantValue::Q8S => try_binary_simd::<$elem, $elem, i8, i8, $op>($lhs, $rhs), + _ => Err(($lhs, $rhs)), + }, + _ => Err(($lhs, $rhs)), + }; + match simd { + Ok(out) => return out, + Err(args) => args, + } + } + }}; +} + +#[cfg(not(feature = "simd"))] +macro_rules! dispatch_binary_simd { + (noq, $elem: ty, $op: ty, $lhs: expr, $rhs: expr, $($ty: ty),*) => {{ ($lhs, $rhs) }}; + ($elem: ty, $op: ty, $lhs: expr, $rhs: expr, $($ty: ty),*) => {{ ($lhs, $rhs) }}; +} + +#[cfg(feature = "simd")] +macro_rules! dispatch_binary_scalar_simd { + (noq, $elem: ty, $op: ty, $lhs: expr, $rhs: expr, $($ty: ty),*) => {{ + paste! { + let simd = match $elem::dtype() { + $(DType::[<$ty:upper>] => try_binary_scalar_simd::<$elem, $elem, $ty, $ty, $op>($lhs, $rhs),)* + _ => Err($lhs), + }; + match simd { + Ok(out) => return out, + Err(args) => args, + } + } + }}; + ($elem: ty, $op: ty, $lhs: expr, $rhs: expr, $($ty: ty),*) => {{ + paste! { + let simd = match $elem::dtype() { + $(DType::[<$ty:upper>] => try_binary_scalar_simd::<$elem, $elem, $ty, $ty, $op>($lhs, $rhs),)* + DType::QFloat(strategy) => match strategy.value { + QuantValue::Q8F | QuantValue::Q8S => try_binary_scalar_simd::<$elem, $elem, i8, i8, $op>($lhs, $rhs), + QuantValue::Q4F | QuantValue::Q4S | QuantValue::Q2F | QuantValue::Q2S | QuantValue::E4M3 | QuantValue::E5M2 | QuantValue::E2M1 => Err($lhs) + }, + _ => Err($lhs), + }; + match simd { + Ok(out) => return out, + Err(args) => args, + } + } + }}; +} + +#[cfg(not(feature = "simd"))] +macro_rules! dispatch_binary_scalar_simd { + (noq, $elem: ty, $op: ty, $lhs: expr, $rhs: expr, $($ty: ty),*) => {{ $lhs }}; + ($elem: ty, $op: ty, $lhs: expr, $rhs: expr, $($ty: ty),*) => {{ $lhs }}; +} + +#[cfg(feature = "simd")] +macro_rules! dispatch_cmp_simd { + ($elem: ty, $op: ty, $lhs: expr, $rhs: expr, $($ty: ty),*) => {{ + paste! { + let simd = match $elem::dtype() { + $(DType::[<$ty:upper>] => try_cmp_simd::<$elem, $ty, $op>($lhs, $rhs),)* + DType::QFloat(strategy) => match strategy.value { + QuantValue::Q8F | QuantValue::Q8S => try_cmp_simd::<$elem, i8, $op>($lhs, $rhs), + QuantValue::Q4F | QuantValue::Q4S | QuantValue::Q2F | QuantValue::Q2S | QuantValue::E4M3 | QuantValue::E5M2 | QuantValue::E2M1 => Err(($lhs, $rhs)) + }, + _ => Err(($lhs, $rhs)), + }; + match simd { + Ok(out) => return out, + Err(args) => args, + } + } + }}; +} + +#[cfg(not(feature = "simd"))] +macro_rules! dispatch_cmp_simd { + ($elem: ty, $op: ty, $lhs: expr, $rhs: expr, $($ty: ty),*) => {{ ($lhs, $rhs) }}; +} + +#[cfg(feature = "simd")] +macro_rules! dispatch_cmp_scalar_simd { + ($elem: ty, $op: ty, $lhs: expr, $rhs: expr, $($ty: ty),*) => {{ + paste! { + let simd = match $elem::dtype() { + $(DType::[<$ty:upper>] => try_cmp_scalar_simd::<$elem, $ty, $op>($lhs, $rhs),)* + DType::QFloat(strategy) => match strategy.value { + QuantValue::Q8F | QuantValue::Q8S => try_cmp_scalar_simd::<$elem, i8, $op>($lhs, $rhs), + QuantValue::Q4F | QuantValue::Q4S | QuantValue::Q2F | QuantValue::Q2S | QuantValue::E4M3 | QuantValue::E5M2 | QuantValue::E2M1 => Err($lhs) + }, + _ => Err($lhs), + }; + match simd { + Ok(out) => return out, + Err(args) => args, + } + } + }}; +} + +#[cfg(not(feature = "simd"))] +macro_rules! dispatch_cmp_scalar_simd { + ($elem: ty, $op: ty, $lhs: expr, $rhs: expr, $($ty: ty),*) => {{ $lhs }}; +} + +#[cfg(feature = "simd")] +macro_rules! dispatch_unary_simd { + ($elem: ty, $op: ty, $lhs: expr, $($ty: ty),*) => {{ + paste! { + let simd = match $elem::dtype() { + $(DType::[<$ty:upper>] => try_unary_simd::<$elem, $elem, $ty, $ty, $op>($lhs),)* + _ => Err($lhs), + }; + match simd { + Ok(out) => return out, + Err(args) => args, + } + } + }}; +} + +#[cfg(not(feature = "simd"))] +macro_rules! dispatch_unary_simd { + ($elem: ty, $op: ty, $lhs: expr, $($ty: ty),*) => {{ $lhs }}; +} + +// Helper function to broadcast two tensors to a common shape for binary operations +// Returns broadcasted views that can be safely zipped +fn broadcast_for_binary_ops<'a, E: Copy, S1, S2>( + lhs: &'a ndarray::ArrayBase, + rhs: &'a ndarray::ArrayBase, +) -> ( + ndarray::ArrayView<'a, E, ndarray::IxDyn>, + ndarray::ArrayView<'a, E, ndarray::IxDyn>, +) +where + S1: ndarray::Data, + S2: ndarray::Data, +{ + // Get shapes + let lhs_shape = lhs.shape(); + let rhs_shape = rhs.shape(); + + // Compute broadcast shape using ndarray's broadcast compatibility rules + let ndims = lhs_shape.len().max(rhs_shape.len()); + let mut broadcast_shape = vec![1; ndims]; + + for i in 0..ndims { + let lhs_dim = if i < lhs_shape.len() { + lhs_shape[lhs_shape.len() - 1 - i] + } else { + 1 + }; + let rhs_dim = if i < rhs_shape.len() { + rhs_shape[rhs_shape.len() - 1 - i] + } else { + 1 + }; + + if lhs_dim == rhs_dim { + broadcast_shape[ndims - 1 - i] = lhs_dim; + } else if lhs_dim == 1 { + broadcast_shape[ndims - 1 - i] = rhs_dim; + } else if rhs_dim == 1 { + broadcast_shape[ndims - 1 - i] = lhs_dim; + } else { + panic!( + "Incompatible shapes for broadcasting: {:?} and {:?}", + lhs_shape, rhs_shape + ); + } + } + + // Create IxDyn from broadcast shape + let broadcast_dim = ndarray::IxDyn(&broadcast_shape); + + // Broadcast both arrays + let lhs_broadcast = lhs + .broadcast(broadcast_dim.clone()) + .expect("Failed to broadcast lhs"); + let rhs_broadcast = rhs + .broadcast(broadcast_dim) + .expect("Failed to broadcast rhs"); + + (lhs_broadcast, rhs_broadcast) +} + +impl NdArrayMathOps +where + E: Copy + NdArrayElement, +{ + pub fn add(lhs: SharedArray, rhs: SharedArray) -> SharedArray { + let (lhs, rhs) = dispatch_binary_simd!( + E, VecAdd, lhs, rhs, u8, i8, u16, i16, u32, i32, f32, u64, i64, f64 + ); + + let array = &lhs + &rhs; + array.into_shared() + } + + pub fn add_scalar(lhs: SharedArray, rhs: E) -> SharedArray { + let lhs = dispatch_binary_scalar_simd!( + E, + VecAdd, + lhs, + rhs.elem(), + u8, + i8, + u16, + i16, + u32, + i32, + f32, + u64, + i64, + f64 + ); + + let array = lhs + rhs; + array.into_shared() + } + + pub fn sub(lhs: SharedArray, rhs: SharedArray) -> SharedArray { + let (lhs, rhs) = dispatch_binary_simd!( + E, VecSub, lhs, rhs, u8, i8, u16, i16, u32, i32, f32, u64, i64, f64 + ); + + let array = lhs - rhs; + array.into_shared() + } + + pub fn sub_scalar(lhs: SharedArray, rhs: E) -> SharedArray { + let lhs = dispatch_binary_scalar_simd!( + E, + VecSub, + lhs, + rhs.elem(), + u8, + i8, + u16, + i16, + u32, + i32, + f32, + u64, + i64, + f64 + ); + + let array = lhs - rhs; + array.into_shared() + } + + pub fn mul(lhs: SharedArray, rhs: SharedArray) -> SharedArray { + let (lhs, rhs) = + dispatch_binary_simd!(noq, E, VecMul, lhs, rhs, u16, i16, u32, i32, f32, f64); + + let array = lhs * rhs; + array.into_shared() + } + + pub fn mul_scalar(lhs: SharedArray, rhs: E) -> SharedArray { + let lhs = dispatch_binary_scalar_simd!( + noq, + E, + VecMul, + lhs, + rhs.elem(), + u16, + i16, + u32, + i32, + f32, + f64 + ); + + let array = lhs * rhs; + array.into_shared() + } + + pub fn div(lhs: SharedArray, rhs: SharedArray) -> SharedArray { + let (lhs, rhs) = dispatch_binary_simd!(noq, E, VecDiv, lhs, rhs, f32, f64); + + let array = lhs / rhs; + array.into_shared() + } + + pub fn div_scalar(lhs: SharedArray, rhs: E) -> SharedArray { + let lhs = dispatch_binary_scalar_simd!(noq, E, VecDiv, lhs, rhs.elem(), f32, f64); + + let array = lhs / rhs; + array.into_shared() + } + + pub fn remainder(lhs: SharedArray, rhs: SharedArray) -> SharedArray { + let (lhs, rhs) = broadcast_for_binary_ops(&lhs, &rhs); + + Zip::from(&lhs) + .and(&rhs) + .map_collect(|&a, &b| { + let a_f = a.to_f64(); + let b_f = b.to_f64(); + let r = a_f - b_f * (a_f / b_f).floor(); + r.elem() + }) + .into_shared() + } + + pub fn remainder_scalar(lhs: SharedArray, rhs: E) -> SharedArray + where + E: core::ops::Rem, + { + let array = lhs.mapv(|x| ((x % rhs) + rhs) % rhs); + array.into_shared() + } + + pub fn recip(tensor: SharedArray) -> SharedArray { + let tensor = dispatch_unary_simd!(E, RecipVec, tensor, f32); + + let array = tensor.map(|x| 1.elem::() / *x); + array.into_shared() + } + + /// Sum all elements - zero-copy for borrowed storage. + pub fn sum_view(view: ArrayView<'_, E, IxDyn>) -> SharedArray { + let sum = view.sum(); + ArrayD::from_elem(IxDyn(&[1]), sum).into_shared() + } + + /// Mean of all elements - zero-copy for borrowed storage. + pub fn mean_view(view: ArrayView<'_, E, IxDyn>) -> SharedArray { + let mean = view.mean().unwrap(); + ArrayD::from_elem(IxDyn(&[1]), mean).into_shared() + } + + /// Product of all elements - zero-copy for borrowed storage. + pub fn prod_view(view: ArrayView<'_, E, IxDyn>) -> SharedArray { + let prod = view.iter().fold(E::one(), |acc, &x| acc * x); + ArrayD::from_elem(IxDyn(&[1]), prod).into_shared() + } + + pub fn mean_dim(tensor: SharedArray, dim: usize) -> SharedArray { + let ndims = tensor.shape().num_dims(); + match ndims { + d if (1..=6).contains(&d) => keepdim!(dim, tensor, mean), + _ => panic!("Dim not supported {ndims}"), + } + } + + pub fn sum_dim(tensor: SharedArray, dim: usize) -> SharedArray { + let ndims = tensor.shape().num_dims(); + match ndims { + d if (1..=6).contains(&d) => keepdim!(dim, tensor, sum), + _ => panic!("Dim not supported {ndims}"), + } + } + + pub fn prod_dim(tensor: SharedArray, dim: usize) -> SharedArray { + let ndims = tensor.shape().num_dims(); + match ndims { + d if (1..=6).contains(&d) => keepdim!(dim, tensor, prod), + _ => panic!("Dim not supported {ndims}"), + } + } + + pub fn cumsum(tensor: SharedArray, dim: usize) -> SharedArray { + cumsum_dim(tensor, dim) + } + + pub fn cumprod(tensor: SharedArray, dim: usize) -> SharedArray { + cumprod_dim(tensor, dim) + } + + pub fn select( + tensor: SharedArray, + dim: usize, + indices: SharedArray, + ) -> SharedArray { + // Read the indices from a contiguous slice rather than through the + // array's own iterator: the slice walks a pointer, while the + // iterator advances a dynamic-rank index on every element and costs + // more than the selection itself. The standardization is a view, and + // hence free, for contiguous indices; otherwise, it pays that + // iteration once, which the direct consumption would pay anyway. + let indices = indices.as_standard_layout(); + let array = tensor.select( + Axis(dim), + &indices + .as_slice() + .unwrap() + .iter() + .map(|i| i.elem::() as usize) + .collect::>(), + ); + + array.into_shared() + } + + pub fn select_assign( + tensor: SharedArray, + dim: usize, + indices: SharedArray, + value: SharedArray, + ) -> SharedArray { + let mut output_array = tensor.into_owned(); + + for (index_value, index) in indices.into_iter().enumerate() { + let mut view = output_array.index_axis_mut(Axis(dim), index.elem::() as usize); + let value = value.index_axis(Axis(dim), index_value); + + view.zip_mut_with(&value, |a, b| *a += *b); + } + + output_array.into_shared() + } + + fn broadcast_dims(lhs: D::Pattern, rhs: D::Pattern) -> Option + where + D: Dimension + IntoDimension, + { + if lhs == rhs { + Some(lhs) + } else { + let lhs = lhs.into_dimension(); + let rhs = rhs.into_dimension(); + let lhs = lhs.slice(); + let rhs = rhs.slice(); + + let total_dims = core::cmp::max(lhs.len(), rhs.len()); + let mut target_shape = vec![0; total_dims]; + // Iterate backwards (from the trailing dimensions inward) + for i in 0..total_dims { + let lhs_dim = lhs + .len() + .checked_sub(1 + i) + .map(|idx| lhs[idx]) + .unwrap_or(1); + let rhs_dim = rhs + .len() + .checked_sub(1 + i) + .map(|idx| rhs[idx]) + .unwrap_or(1); + + if lhs_dim == rhs_dim { + target_shape[total_dims - 1 - i] = lhs_dim; + } else if lhs_dim == 1 { + target_shape[total_dims - 1 - i] = rhs_dim; + } else if rhs_dim == 1 { + target_shape[total_dims - 1 - i] = lhs_dim; + } else { + // Dimensions are completely incompatible (e.g., trying to match 3 and 5) + return None; + } + } + + let dyn_dim = IxDyn(&target_shape); + D::Dim::from_dimension(&dyn_dim).map(|d| d.into_pattern()) + } + } + + pub(crate) fn elementwise_op( + lhs: SharedArray, + rhs: SharedArray, + var_name: impl FnMut(&E, &OtherE) -> E, + ) -> SharedArray { + if let Some(target) = Self::broadcast_dims::(lhs.dim(), rhs.dim()) { + let lhs = lhs.broadcast(target.clone()).unwrap(); + let rhs = rhs.broadcast(target).unwrap(); + + Zip::from(lhs).and(rhs).map_collect(var_name).into_shared() + } else { + panic!( + "Incompatible shapes for broadcasting: {:?} and {:?}", + lhs.shape(), + rhs.shape() + ); + } + } + + pub(crate) fn elementwise_op_scalar( + lhs: SharedArray, + var_name: impl FnMut(E) -> E, + ) -> SharedArray { + lhs.mapv(var_name).into_shared() + } + + pub(crate) fn abs(tensor: SharedArray) -> SharedArray { + let tensor = dispatch_unary_simd!(E, VecAbs, tensor, i8, i16, i32, f32, f64); + + tensor.mapv_into(|a| a.abs_elem()).into_shared() + } + + pub(crate) fn equal(lhs: SharedArray, rhs: SharedArray) -> SharedArray { + let (lhs, rhs) = dispatch_cmp_simd!( + E, VecEquals, lhs, rhs, u8, i8, u16, i16, u32, f32, i32, u64, i64, f64 + ); + + // Use the helper to broadcast both arrays to a common shape + let (lhs_broadcast, rhs_broadcast) = broadcast_for_binary_ops(&lhs, &rhs); + // Now we can safely zip and compare + Zip::from(&lhs_broadcast) + .and(&rhs_broadcast) + .map_collect(|&lhs, &rhs| lhs == rhs) + .into_shared() + } + + pub(crate) fn equal_elem(lhs: SharedArray, rhs: E) -> SharedArray { + let lhs = dispatch_cmp_scalar_simd!( + E, + VecEquals, + lhs, + rhs.elem(), + u8, + i8, + u16, + i16, + u32, + f32, + i32, + u64, + i64, + f64 + ); + + lhs.mapv(|a| a == rhs).into_shared() + } + + pub(crate) fn sign_op(tensor: SharedArray) -> SharedArray + where + E: Signed, + { + let zero = 0.elem(); + let one = 1.elem::(); + + tensor + .mapv(|x| { + if x == zero { + zero + } else { + match x.is_positive() { + true => one, + false => -one, + } + } + }) + .into_shared() + } +} + +impl NdArrayMathOps +where + E: Copy + NdArrayElement + PartialOrd, +{ + /// Max of all elements - zero-copy for borrowed storage. + pub fn max_view(view: ArrayView<'_, E, IxDyn>) -> SharedArray { + let max = view + .iter() + .copied() + .reduce(|a, b| if a > b { a } else { b }) + .expect("Cannot compute max of empty tensor"); + ArrayD::from_elem(IxDyn(&[1]), max).into_shared() + } + + /// Min of all elements - zero-copy for borrowed storage. + pub fn min_view(view: ArrayView<'_, E, IxDyn>) -> SharedArray { + let min = view + .iter() + .copied() + .reduce(|a, b| if a < b { a } else { b }) + .expect("Cannot compute min of empty tensor"); + ArrayD::from_elem(IxDyn(&[1]), min).into_shared() + } + + /// Argmax along dimension - zero-copy for borrowed storage. + pub fn argmax_view( + view: ArrayView<'_, E, IxDyn>, + dim: usize, + ) -> SharedArray { + arg_view(view, dim, CmpType::Max) + } + + /// Argmin along dimension - zero-copy for borrowed storage. + pub fn argmin_view( + view: ArrayView<'_, E, IxDyn>, + dim: usize, + ) -> SharedArray { + arg_view(view, dim, CmpType::Min) + } + + pub fn cummin(tensor: SharedArray, dim: usize) -> SharedArray { + cummin_dim(tensor, dim) + } + + pub fn cummax(tensor: SharedArray, dim: usize) -> SharedArray { + cummax_dim(tensor, dim) + } + + pub fn argmax( + tensor: SharedArray, + dim: usize, + ) -> SharedArray { + arg(tensor, dim, CmpType::Max) + } + + pub fn argmin( + tensor: SharedArray, + dim: usize, + ) -> SharedArray { + arg(tensor, dim, CmpType::Min) + } + + pub fn clamp_min(tensor: SharedArray, min: E) -> SharedArray { + let mut tensor = dispatch_binary_scalar_simd!( + E, + VecMax, + tensor, + min.elem(), + u8, + i8, + u16, + i16, + u32, + i32, + f32, + u64, + i64, + f64 + ); + + tensor.mapv_inplace(|x| match x < min { + true => min, + false => x, + }); + + tensor + } + + pub fn clamp_max(tensor: SharedArray, max: E) -> SharedArray { + let mut tensor = dispatch_binary_scalar_simd!( + E, + VecMin, + tensor, + max.elem(), + u8, + i8, + u16, + i16, + u32, + i32, + f32, + u64, + i64, + f64 + ); + + tensor.mapv_inplace(|x| match x > max { + true => max, + false => x, + }); + + tensor + } + + pub fn clamp(tensor: SharedArray, min: E, max: E) -> SharedArray { + let mut tensor = dispatch_binary_scalar_simd!( + E, + VecClamp, + tensor, + (min.elem(), max.elem()), + u8, + i8, + u16, + i16, + u32, + i32, + f32, + u64, + i64, + f64 + ); + + tensor.mapv_inplace(|x| match x < min { + true => min, + false => match x > max { + true => max, + false => x, + }, + }); + + tensor + } + + pub(crate) fn greater(lhs: SharedArray, rhs: SharedArray) -> SharedArray { + let (lhs, rhs) = dispatch_cmp_simd!( + E, VecGreater, lhs, rhs, u8, i8, u16, i16, u32, f32, i32, u64, i64, f64 + ); + + // Use the helper to broadcast both arrays to a common shape + let (lhs_broadcast, rhs_broadcast) = broadcast_for_binary_ops(&lhs, &rhs); + // Now we can safely zip and compare + Zip::from(&lhs_broadcast) + .and(&rhs_broadcast) + .map_collect(|&lhs, &rhs| lhs > rhs) + .into_shared() + } + + pub(crate) fn greater_elem(lhs: SharedArray, rhs: E) -> SharedArray { + let lhs = dispatch_cmp_scalar_simd!( + E, + VecGreater, + lhs, + rhs.elem(), + u8, + i8, + u16, + i16, + u32, + f32, + i32, + u64, + i64, + f64 + ); + + lhs.mapv(|a| a > rhs).into_shared() + } + + pub(crate) fn greater_equal(lhs: SharedArray, rhs: SharedArray) -> SharedArray { + let (lhs, rhs) = dispatch_cmp_simd!( + E, + VecGreaterEq, + lhs, + rhs, + u8, + i8, + u16, + i16, + u32, + f32, + i32, + u64, + i64, + f64 + ); + + // Use the helper to broadcast both arrays to a common shape + let (lhs_broadcast, rhs_broadcast) = broadcast_for_binary_ops(&lhs, &rhs); + // Now we can safely zip and compare + Zip::from(&lhs_broadcast) + .and(&rhs_broadcast) + .map_collect(|&lhs, &rhs| lhs >= rhs) + .into_shared() + } + + pub(crate) fn greater_equal_elem(lhs: SharedArray, rhs: E) -> SharedArray { + let lhs = dispatch_cmp_scalar_simd!( + E, + VecGreaterEq, + lhs, + rhs.elem(), + u8, + i8, + u16, + i16, + u32, + f32, + i32, + u64, + i64, + f64 + ); + + lhs.mapv(|a| a >= rhs).into_shared() + } + + pub(crate) fn lower_equal(lhs: SharedArray, rhs: SharedArray) -> SharedArray { + let (lhs, rhs) = dispatch_cmp_simd!( + E, VecLowerEq, lhs, rhs, u8, i8, u16, i16, u32, f32, i32, u64, i64, f64 + ); + + // Use the helper to broadcast both arrays to a common shape + let (lhs_broadcast, rhs_broadcast) = broadcast_for_binary_ops(&lhs, &rhs); + // Now we can safely zip and compare + Zip::from(&lhs_broadcast) + .and(&rhs_broadcast) + .map_collect(|&lhs, &rhs| lhs <= rhs) + .into_shared() + } + + pub(crate) fn lower_equal_elem(lhs: SharedArray, rhs: E) -> SharedArray { + let lhs = dispatch_cmp_scalar_simd!( + E, + VecLowerEq, + lhs, + rhs.elem(), + u8, + i8, + u16, + i16, + u32, + f32, + i32, + u64, + i64, + f64 + ); + + lhs.mapv(|a| a <= rhs).into_shared() + } + + pub(crate) fn lower(lhs: SharedArray, rhs: SharedArray) -> SharedArray { + let (lhs, rhs) = dispatch_cmp_simd!( + E, VecLower, lhs, rhs, u8, i8, u16, i16, u32, f32, i32, u64, i64, f64 + ); + + // Use the helper to broadcast both arrays to a common shape + let (lhs_broadcast, rhs_broadcast) = broadcast_for_binary_ops(&lhs, &rhs); + + // Now we can safely zip and compare + Zip::from(&lhs_broadcast) + .and(&rhs_broadcast) + .map_collect(|&lhs, &rhs| lhs < rhs) + .into_shared() + } + + pub(crate) fn lower_elem(lhs: SharedArray, rhs: E) -> SharedArray { + let lhs = dispatch_cmp_scalar_simd!( + E, + VecLower, + lhs, + rhs.elem(), + u8, + i8, + u16, + i16, + u32, + f32, + i32, + u64, + i64, + f64 + ); + + lhs.mapv(|a| a < rhs).into_shared() + } +} + +pub struct NdArrayBitOps(PhantomData); + +impl NdArrayBitOps { + pub(crate) fn bitand(lhs: SharedArray, rhs: SharedArray) -> SharedArray { + let (lhs, rhs) = + dispatch_binary_simd!(I, VecBitAnd, lhs, rhs, i8, u8, i16, u16, i32, u32, i64, u64); + + NdArrayMathOps::elementwise_op(lhs, rhs, |a: &I, b: &I| { + (a.elem::() & (b.elem::())).elem() + }) + } + + pub(crate) fn bitand_scalar(lhs: SharedArray, rhs: I) -> SharedArray { + let lhs = dispatch_binary_scalar_simd!( + I, + VecBitAnd, + lhs, + rhs.elem(), + i8, + u8, + i16, + u16, + i32, + u32, + i64, + u64 + ); + + NdArrayMathOps::elementwise_op_scalar(lhs, |a: I| { + (a.elem::() & rhs.elem::()).elem() + }) + } + + pub(crate) fn bitor(lhs: SharedArray, rhs: SharedArray) -> SharedArray { + let (lhs, rhs) = + dispatch_binary_simd!(I, VecBitOr, lhs, rhs, i8, u8, i16, u16, i32, u32, i64, u64); + + NdArrayMathOps::elementwise_op(lhs, rhs, |a: &I, b: &I| { + (a.elem::() | (b.elem::())).elem() + }) + } + + pub(crate) fn bitor_scalar(lhs: SharedArray, rhs: I) -> SharedArray { + let lhs = dispatch_binary_scalar_simd!( + I, + VecBitOr, + lhs, + rhs.elem(), + i8, + u8, + i16, + u16, + i32, + u32, + i64, + u64 + ); + + NdArrayMathOps::elementwise_op_scalar(lhs, |a: I| { + (a.elem::() | rhs.elem::()).elem() + }) + } + + pub(crate) fn bitxor(lhs: SharedArray, rhs: SharedArray) -> SharedArray { + let (lhs, rhs) = + dispatch_binary_simd!(I, VecBitXor, lhs, rhs, i8, u8, i16, u16, i32, u32, i64, u64); + + NdArrayMathOps::elementwise_op(lhs, rhs, |a: &I, b: &I| { + (a.elem::() ^ (b.elem::())).elem() + }) + } + + pub(crate) fn bitxor_scalar(lhs: SharedArray, rhs: I) -> SharedArray { + let lhs = dispatch_binary_scalar_simd!( + I, + VecBitXor, + lhs, + rhs.elem(), + i8, + u8, + i16, + u16, + i32, + u32, + i64, + u64 + ); + + NdArrayMathOps::elementwise_op_scalar(lhs, |a: I| { + (a.elem::() ^ rhs.elem::()).elem() + }) + } + + pub(crate) fn bitnot(tensor: SharedArray) -> SharedArray { + let tensor = + dispatch_unary_simd!(I, VecBitNot, tensor, i8, u8, i16, u16, i32, u32, i64, u64); + + NdArrayMathOps::elementwise_op_scalar(tensor, |a: I| (!a.elem::()).elem()) + } +} + +pub struct NdArrayBoolOps; + +// Rust booleans are either `00000000` or `00000001`, so bitwise and/or is fine, but bitwise not would +// produce invalid values. +impl NdArrayBoolOps { + pub(crate) fn equal(lhs: SharedArray, rhs: SharedArray) -> SharedArray { + #[cfg(feature = "simd")] + let (lhs, rhs) = match try_cmp_simd::(lhs, rhs) { + Ok(out) => return out, + Err(args) => args, + }; + + // Use the helper to broadcast both arrays to a common shape + let (lhs_broadcast, rhs_broadcast) = broadcast_for_binary_ops(&lhs, &rhs); + // Now we can safely zip and compare + Zip::from(&lhs_broadcast) + .and(&rhs_broadcast) + .map_collect(|&lhs, &rhs| lhs == rhs) + .into_shared() + } + + pub(crate) fn equal_elem(lhs: SharedArray, rhs: bool) -> SharedArray { + #[cfg(feature = "simd")] + let lhs = match try_cmp_scalar_simd::(lhs, rhs.elem()) { + Ok(out) => return out, + Err(args) => args, + }; + + lhs.mapv(|a| a == rhs).into_shared() + } + + pub(crate) fn and(lhs: SharedArray, rhs: SharedArray) -> SharedArray { + #[cfg(feature = "simd")] + let (lhs, rhs) = match try_binary_simd::(lhs, rhs) { + Ok(out) => return out, + Err(args) => args, + }; + + // Use the helper to broadcast both arrays to a common shape + let (lhs_broadcast, rhs_broadcast) = broadcast_for_binary_ops(&lhs, &rhs); + // Now we can safely zip and compare + Zip::from(&lhs_broadcast) + .and(&rhs_broadcast) + .map_collect(|&lhs, &rhs| lhs && rhs) + .into_shared() + } + + pub(crate) fn or(lhs: SharedArray, rhs: SharedArray) -> SharedArray { + #[cfg(feature = "simd")] + let (lhs, rhs) = match try_binary_simd::(lhs, rhs) { + Ok(out) => return out, + Err(args) => args, + }; + + // Use the helper to broadcast both arrays to a common shape + let (lhs_broadcast, rhs_broadcast) = broadcast_for_binary_ops(&lhs, &rhs); + // Now we can safely zip and compare + Zip::from(&lhs_broadcast) + .and(&rhs_broadcast) + .map_collect(|&lhs, &rhs| lhs || rhs) + .into_shared() + } + + /// Any element is true - zero-copy for borrowed storage. + pub fn any_view(view: ArrayView<'_, bool, IxDyn>) -> bool { + view.iter().any(|&x| x) + } + + /// All elements are true - zero-copy for borrowed storage. + pub fn all_view(view: ArrayView<'_, bool, IxDyn>) -> bool { + view.iter().all(|&x| x) + } +} + +enum CmpType { + Min, + Max, +} + +fn arg( + tensor: SharedArray, + dim: usize, + cmp: CmpType, +) -> SharedArray { + arg_view(tensor.view(), dim, cmp) +} + +/// View-based argmax/argmin - zero-copy for borrowed storage. +fn arg_view( + view: ArrayView<'_, E, IxDyn>, + dim: usize, + cmp: CmpType, +) -> SharedArray { + let mut reshape = view.shape().to_vec(); + reshape[dim] = 1; + + let output = view.map_axis(Axis(dim), |arr| { + // Find the min/max value in the array, and return its index. + let (_e, idx) = arr.indexed_iter().fold((arr[0], 0usize), |acc, (idx, e)| { + let cmp = match cmp { + CmpType::Min => e < &acc.0, + CmpType::Max => e > &acc.0, + }; + + if cmp { (*e, idx) } else { acc } + }); + + (idx as i64).elem() + }); + + let output = output.to_shape(Dim(reshape.as_slice())).unwrap(); + + output.into_shared() +} + +#[cfg(test)] +mod tests { + use burn_backend::TensorData; + + use crate::NdArrayTensor; + + use super::*; + + #[test] + fn should_generate_row_major_layout_for_cat() { + let expected_shape: &[usize] = &[4, 6, 2]; + let expected_strides: &[isize] = &[12, 2, 1]; + let NdArrayTensor::I32(expected_storage) = NdArrayTensor::from_data(TensorData::from([ + [[1, 0], [2, 0], [3, 0], [4, 0], [5, 0], [6, 0]], + [[7, 0], [8, 0], [9, 0], [10, 0], [11, 0], [12, 0]], + [[13, 0], [14, 0], [15, 0], [16, 0], [17, 0], [18, 0]], + [[19, 0], [20, 0], [21, 0], [22, 0], [23, 0], [24, 0]], + ])) else { + panic!() + }; + let expected_array = expected_storage.into_shared(); + + let NdArrayTensor::I32(tensor_storage) = NdArrayTensor::from_data(TensorData::from([ + [1, 2, 3, 4, 5, 6], + [7, 8, 9, 10, 11, 12], + [13, 14, 15, 16, 17, 18], + [19, 20, 21, 22, 23, 24], + ])) else { + panic!() + }; + let tensor = tensor_storage.into_shared(); + + // unsqueeze dim on the outermost axis + let array = NdArrayOps::reshape(tensor, Shape::from([4, 6, 1])); + let NdArrayTensor::I32(zeros_storage) = + NdArrayTensor::from_data(TensorData::zeros::([4, 6, 1])) + else { + panic!() + }; + let zeros = zeros_storage.into_shared(); + // make `ndarray` concatenates array on the outermost axis + let array = NdArrayOps::cat([array, zeros].to_vec(), 2); + + assert!(array.is_standard_layout()); + assert_eq!(array.shape(), expected_shape); + assert_eq!(array.strides(), expected_strides); + assert_eq!( + array.into_iter().collect::>(), + expected_array.into_iter().collect::>(), + ); + } +} diff --git a/crates/burn-ndarray/src/ops/bool_tensor.rs b/crates/burn-ndarray/src/ops/bool_tensor.rs new file mode 100644 index 0000000..67913a4 --- /dev/null +++ b/crates/burn-ndarray/src/ops/bool_tensor.rs @@ -0,0 +1,219 @@ +// Language +use alloc::vec; +use alloc::vec::Vec; +use burn_backend::Scalar; +use burn_backend::{ElementConversion, TensorMetadata, tensor::FloatTensor}; +use burn_backend::{ + backend::ExecutionError, + ops::BoolTensorOps, + tensor::{BoolTensor, IntTensor}, +}; +use burn_std::{BoolDType, FloatDType, IntDType}; +use ndarray::IntoDimension; + +// Current crate +use crate::{NdArray, execute_with_int_dtype, tensor::NdArrayTensor}; +use crate::{ + NdArrayDevice, SharedArray, execute_with_float_out_dtype, execute_with_int_out_dtype, slice, +}; + +// Workspace crates +use burn_backend::{Shape, TensorData}; + +use super::{NdArrayBoolOps, NdArrayOps}; + +impl BoolTensorOps for NdArray { + fn bool_from_data(data: TensorData, _device: &NdArrayDevice) -> NdArrayTensor { + if !data.dtype.is_bool() { + unimplemented!("Unsupported dtype for `bool_from_data`") + } + NdArrayTensor::from_data(data) + } + + async fn bool_into_data(tensor: NdArrayTensor) -> Result { + Ok(tensor.into_data()) + } + + fn bool_to_device(tensor: NdArrayTensor, _device: &NdArrayDevice) -> NdArrayTensor { + tensor + } + + fn bool_reshape(tensor: NdArrayTensor, shape: Shape) -> NdArrayTensor { + NdArrayOps::reshape(tensor.bool(), shape).into() + } + + fn bool_slice(tensor: NdArrayTensor, slices: &[burn_backend::Slice]) -> NdArrayTensor { + slice!(tensor, slices) + } + + fn bool_into_int(tensor: NdArrayTensor, out_dtype: IntDType) -> NdArrayTensor { + // Use mapv directly instead of collecting to Vec and going through TensorData + execute_with_int_out_dtype!( + out_dtype, + I, + tensor.bool().mapv(|b| b.elem::()).into_shared().into() + ) + } + + fn bool_empty(shape: Shape, _device: &NdArrayDevice, dtype: BoolDType) -> NdArrayTensor { + Self::bool_zeros(shape, _device, dtype) + } + + fn bool_zeros(shape: Shape, _device: &NdArrayDevice, _dtype: BoolDType) -> NdArrayTensor { + let values = vec![false; shape.num_elements()]; + NdArrayTensor::from_data(TensorData::new(values, shape)) + } + + fn bool_ones(shape: Shape, _device: &NdArrayDevice, _dtype: BoolDType) -> NdArrayTensor { + let values = vec![true; shape.num_elements()]; + NdArrayTensor::from_data(TensorData::new(values, shape)) + } + + fn bool_slice_assign( + tensor: NdArrayTensor, + slices: &[burn_backend::Slice], + value: NdArrayTensor, + ) -> NdArrayTensor { + NdArrayOps::slice_assign(tensor.bool(), slices, value.bool()).into() + } + + fn bool_cat(tensors: Vec, dim: usize) -> NdArrayTensor { + NdArrayOps::cat(tensors.into_iter().map(|it| it.bool()).collect(), dim).into() + } + + fn bool_equal(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor { + NdArrayBoolOps::equal(lhs.bool(), rhs.bool()).into() + } + + fn bool_not(tensor: NdArrayTensor) -> NdArrayTensor { + tensor.bool().mapv(|a| !a).into_shared().into() + } + + fn bool_and(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor { + NdArrayBoolOps::and(lhs.bool(), rhs.bool()).into() + } + + fn bool_or(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor { + NdArrayBoolOps::or(lhs.bool(), rhs.bool()).into() + } + + fn bool_into_float(tensor: NdArrayTensor, out_dtype: FloatDType) -> FloatTensor { + execute_with_float_out_dtype!( + out_dtype, + E, + tensor.bool().mapv(|b| b.elem::()).into_shared().into() + ) + } + + fn bool_swap_dims(tensor: NdArrayTensor, dim1: usize, dim2: usize) -> NdArrayTensor { + NdArrayOps::swap_dims(tensor.bool(), dim1, dim2).into() + } + + fn bool_permute(tensor: NdArrayTensor, axes: &[usize]) -> NdArrayTensor { + tensor.bool().permuted_axes(axes.into_dimension()).into() + } + + fn bool_expand(tensor: NdArrayTensor, shape: Shape) -> NdArrayTensor { + NdArrayOps::expand(tensor.bool(), shape).into() + } + + fn bool_select(tensor: NdArrayTensor, dim: usize, indices: NdArrayTensor) -> NdArrayTensor { + execute_with_int_dtype!(indices, I, |indices: SharedArray| -> NdArrayTensor { + let tensor_bool = tensor.bool(); + let indices_vec: Vec = indices + .into_iter() + .map(|i| i.elem::() as usize) + .collect(); + + let selected = tensor_bool.select(ndarray::Axis(dim), &indices_vec); + selected.into_shared().into() + }) + } + + fn bool_select_or( + tensor: NdArrayTensor, + dim: usize, + indices: NdArrayTensor, + value: NdArrayTensor, + ) -> NdArrayTensor { + execute_with_int_dtype!(indices, I, |indices: SharedArray| -> NdArrayTensor { + let mut output_array = tensor.bool().into_owned(); + let value_bool = value.bool(); + + for (index_value, index) in indices.into_iter().enumerate() { + let index_usize = index.elem::() as usize; + let mut view = output_array.index_axis_mut(ndarray::Axis(dim), index_usize); + let value_slice = value_bool.index_axis(ndarray::Axis(dim), index_value); + // For boolean tensors, select_assign should use logical OR operation + view.zip_mut_with(&value_slice, |a, b| *a = *a || *b); + } + output_array.into_shared().into() + }) + } + + fn bool_flip(tensor: NdArrayTensor, axes: &[usize]) -> NdArrayTensor { + NdArrayOps::flip(tensor.bool(), axes).into() + } + + fn bool_unfold(tensor: NdArrayTensor, dim: usize, size: usize, step: usize) -> NdArrayTensor { + NdArrayOps::unfold(tensor.bool(), dim, size, step).into() + } + + fn bool_mask_where( + tensor: BoolTensor, + mask: BoolTensor, + value: BoolTensor, + ) -> BoolTensor { + NdArrayOps::mask_where(tensor.bool(), mask.bool(), value.bool()).into() + } + + fn bool_mask_fill( + tensor: BoolTensor, + mask: BoolTensor, + value: Scalar, + ) -> BoolTensor { + NdArrayOps::mask_fill(tensor.bool(), mask.bool(), value.elem()).into() + } + + fn bool_gather( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + ) -> BoolTensor { + execute_with_int_dtype!(indices, |indices| NdArrayOps::gather( + dim, + tensor.bool(), + indices + )) + } + + fn bool_scatter_or( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor { + execute_with_int_dtype!(indices, |indices| NdArrayOps::scatter( + dim, + tensor.bool(), + indices, + value.bool() + )) + } + + fn bool_equal_elem(lhs: BoolTensor, rhs: Scalar) -> BoolTensor { + NdArrayBoolOps::equal_elem(lhs.bool(), rhs.elem()).into() + } + + fn bool_any(tensor: BoolTensor) -> BoolTensor { + // Use view() for zero-copy on borrowed storage with short-circuit evaluation + let result = NdArrayBoolOps::any_view(tensor.bool().view()); + NdArrayTensor::from_data(TensorData::new(vec![result], Shape::new([1]))) + } + + fn bool_all(tensor: BoolTensor) -> BoolTensor { + // Use view() for zero-copy on borrowed storage with short-circuit evaluation + let result = NdArrayBoolOps::all_view(tensor.bool().view()); + NdArrayTensor::from_data(TensorData::new(vec![result], Shape::new([1]))) + } +} diff --git a/crates/burn-ndarray/src/ops/conv.rs b/crates/burn-ndarray/src/ops/conv.rs new file mode 100644 index 0000000..5fb2cad --- /dev/null +++ b/crates/burn-ndarray/src/ops/conv.rs @@ -0,0 +1,574 @@ +use burn_backend::{ + ElementConversion, + ops::{ + ConvOptions, ConvTransposeOptions, + conv::{calculate_conv_output_size, calculate_conv_transpose_output_size}, + }, +}; +use ndarray::{ + Array3, Array4, Array5, ArrayView2, ArrayView3, ArrayViewMut2, ArrayViewMut3, Axis, Dim, s, +}; + +use crate::{ + NdArrayElement, SharedArray, iter_par, iter_range_par, + ops::padding::{apply_padding_4d, apply_padding_5d}, + run_par, + sharing::UnsafeSharedRef, + tensor::NdArrayTensor, +}; + +#[inline(always)] +fn conv2d_mad_inner( + mut output: ArrayViewMut2, + x: ArrayView2, + k: E, + k_xy: (usize, usize), + out_xy: (usize, usize), + stride: (usize, usize), + dilation: (usize, usize), +) { + let (kh, kw) = k_xy; + let (out_width, out_height) = out_xy; + let (stride_width, stride_height) = stride; + let (dilation_width, dilation_height) = dilation; + + for oh in 0..out_height { + // Construct a sub-slice view of the input row. + // This is done upfront so that rustc does not have to emit bounds checks + // in the hot loop below. + let ir = x + .row(oh * stride_height + kh * dilation_height) + .to_slice() + .unwrap(); + + // Ditto. Construct a sub-slice view of the output row, and explicitly specify + // the bounds upfront as 0..out_width so that rustc can make the assumption + // that all accesses are in-bounds in the below loop. + let mut or = output.row_mut(oh); + let or = &mut or.as_slice_mut().unwrap()[0..out_width]; + + #[allow(clippy::needless_range_loop)] + for ow in 0..out_width { + let iw = ow * stride_width + kw * dilation_width; + or[ow] += ir[iw] * k; + } + } +} + +#[inline(always)] +fn conv3d_mad_inner( + mut output: ArrayViewMut3, + x: ArrayView3, + k: E, + k_xyz: (usize, usize, usize), + out_xyz: (usize, usize, usize), + stride: (usize, usize, usize), + dilation: (usize, usize, usize), +) { + let (kd, kh, kw) = k_xyz; + let (out_width, out_height, out_depth) = out_xyz; + let (stride_width, stride_height, stride_depth) = stride; + let (dilation_width, dilation_height, dilation_depth) = dilation; + + for od in 0..out_depth { + let id = od * stride_depth + kd * dilation_depth; + + for oh in 0..out_height { + let ih = oh * stride_height + kh * dilation_height; + + // Construct a sub-slice view of the input row. + // This is done upfront so that rustc does not have to emit bounds checks + // in the hot loop below. + let ir = x.slice(s![id, ih, ..]).to_slice().unwrap(); + + // Ditto. Construct a sub-slice view of the output row, and explicitly specify + // the bounds upfront as 0..out_width so that rustc can make the assumption + // that all accesses are in-bounds in the below loop. + let or = &mut output + .slice_mut(s![od, oh, 0..out_width]) + .into_slice() + .unwrap()[0..out_width]; + + #[allow(clippy::needless_range_loop)] + for ow in 0..out_width { + let iw = ow * stride_width + kw * dilation_width; + or[ow] += ir[iw] * k; + } + } + } +} + +pub(crate) fn conv2d( + x: SharedArray, + weight: SharedArray, + bias: Option>, + options: ConvOptions<2>, +) -> SharedArray +where + NdArrayTensor: From>, +{ + let [dilation_height, dilation_width] = options.dilation; + let [padding_height, padding_width] = options.padding; + let [stride_height, stride_width] = options.stride; + let [batch_size, _in_channels, in_height, in_width] = x.shape().try_into().unwrap(); + let [out_channels, in_channels, kernel_height, kernel_width] = + weight.shape().try_into().unwrap(); + let channels_per_group = out_channels / options.groups; + + let out_height = calculate_conv_output_size( + kernel_height, + stride_height, + padding_height, + dilation_height, + in_height, + ); + let out_width = calculate_conv_output_size( + kernel_width, + stride_width, + padding_width, + dilation_width, + in_width, + ); + + let x = apply_padding_4d::(x, options.padding, 0i32.elem()); + + // Convert inputs from dynamic indexes to static to improve perf. + let x = x.into_dimensionality::().unwrap(); + let weights = weight.into_dimensionality::().unwrap(); + + let mut output = Array3::zeros(Dim([batch_size * out_channels, out_height, out_width])); + + run_par!(|| { + iter_par!(output.axis_iter_mut(Axis(0))) + .enumerate() + .for_each( + #[inline(never)] + |(k, mut output)| { + let b = k / out_channels; + let oc = k % out_channels; + let g = oc / channels_per_group; + + for ic in (in_channels * g)..(in_channels * (g + 1)) { + let weight_ic = ic - (g * in_channels); + + let x = x.slice(s![b, ic, .., ..]); + let k = weights.slice(s![oc, weight_ic, .., ..]); + + for kh in 0..kernel_height { + for kw in 0..kernel_width { + let k = k[[kh, kw]]; + + // NOTE: This function call is duplicated twice so that the compiler can perform auto-vectorization + // in the case that the stride/dilation is 1. + #[allow(clippy::if_same_then_else)] + if (1, 1, 1, 1) + == ( + stride_width, + stride_height, + dilation_width, + dilation_height, + ) + { + conv2d_mad_inner( + output.view_mut(), + x.view(), + k, + (kh, kw), + (out_width, out_height), + (stride_width, stride_height), + (dilation_width, dilation_height), + ); + } else { + conv2d_mad_inner( + output.view_mut(), + x.view(), + k, + (kh, kw), + (out_width, out_height), + (stride_width, stride_height), + (dilation_width, dilation_height), + ); + } + } + } + } + + if let Some(bias) = &bias { + let bias = bias[oc]; + + for oh in 0..out_height { + // Get a mutable slice reference to the row we're looping over. + // We explicitly define the bounds to 0..out_width so that rustc can make + // the assumption that all accesses are in-bounds. + let mut or = output.row_mut(oh); + let or = &mut or.as_slice_mut().unwrap()[0..out_width]; + + #[allow(clippy::needless_range_loop)] + for ow in 0..out_width { + or[ow] += bias; + } + } + } + }, + ); + }); + + output + .to_shape([batch_size, out_channels, out_height, out_width]) + .unwrap() + .into_dyn() + .into_shared() +} + +pub(crate) fn conv_transpose2d( + x: SharedArray, + weight: SharedArray, + bias: Option>, + options: ConvTransposeOptions<2>, +) -> SharedArray { + let [dilation_height, dilation_width] = options.dilation; + let [padding_height, padding_width] = options.padding; + let [stride_height, stride_width] = options.stride; + let [out_padding_height, out_padding_width] = options.padding_out; + let [batch_size, _in_channels, in_height, in_width] = x.shape().try_into().unwrap(); + let [in_channels, out_channels, kernel_height, kernel_width] = + weight.shape().try_into().unwrap(); + + let out_height = calculate_conv_transpose_output_size( + kernel_height, + stride_height, + padding_height, + out_padding_height, + dilation_height, + in_height, + ); + let out_width = calculate_conv_transpose_output_size( + kernel_width, + stride_width, + padding_width, + out_padding_width, + dilation_width, + in_width, + ); + + let x = x; + let mut output = Array4::zeros(Dim([ + batch_size, + out_channels * options.groups, + out_height, + out_width, + ])); + + let unsafe_shared_out = UnsafeSharedRef::new(&mut output); + + run_par!(|| { + iter_range_par!(0, batch_size * out_channels * options.groups).for_each(|k| unsafe { + let b = k / (out_channels * options.groups); + let oc = k % out_channels; + let g = (k / out_channels) % options.groups; + + let output = unsafe_shared_out.get(); + + let oc_out = oc + (out_channels * g); + let ic_start = g * (in_channels / options.groups); + let ic_end = ic_start + in_channels / options.groups; + + for ic in ic_start..ic_end { + for ih in 0..in_height { + for iw in 0..in_width { + for kh in 0..kernel_height { + for kw in 0..kernel_width { + let oh = ih * stride_height + kh * dilation_height; + let ow = iw * stride_width + kw * dilation_width; + + if oh >= out_height + padding_height + || ow >= out_width + padding_width + || oh < padding_height + || ow < padding_width + { + continue; + } + + let oh = oh - padding_height; + let ow = ow - padding_width; + + output[[b, oc_out, oh, ow]] += + x[[b, ic, ih, iw]] * weight[[ic, oc, kh, kw]]; + } + } + } + } + } + + if let Some(bias) = &bias { + for oh in 0..out_height { + for ow in 0..out_width { + output[[b, oc_out, oh, ow]] += bias[oc_out]; + } + } + } + }); + }); + + output.into_dyn().into_shared() +} + +pub(crate) fn conv3d( + x: SharedArray, + weight: SharedArray, + bias: Option>, + options: ConvOptions<3>, +) -> SharedArray +where + NdArrayTensor: From>, +{ + let [dilation_depth, dilation_height, dilation_width] = options.dilation; + let [padding_depth, padding_height, padding_width] = options.padding; + let [stride_depth, stride_height, stride_width] = options.stride; + let [batch_size, _in_channels, in_depth, in_height, in_width] = x.shape().try_into().unwrap(); + let [ + out_channels, + in_channels, + kernel_depth, + kernel_height, + kernel_width, + ] = weight.shape().try_into().unwrap(); + let out_c_per_group = out_channels / options.groups; + + let out_depth = calculate_conv_output_size( + kernel_depth, + stride_depth, + padding_depth, + dilation_depth, + in_depth, + ); + let out_height = calculate_conv_output_size( + kernel_height, + stride_height, + padding_height, + dilation_height, + in_height, + ); + let out_width = calculate_conv_output_size( + kernel_width, + stride_width, + padding_width, + dilation_width, + in_width, + ); + + let x = apply_padding_5d::(x, options.padding, 0i32.elem()); + + // Convert inputs from dynamic indexes to static to improve perf. + let x = x.into_dimensionality::().unwrap(); + let weights = weight.into_dimensionality::().unwrap(); + + let mut output = Array4::zeros(Dim([ + batch_size * out_channels, + out_depth, + out_height, + out_width, + ])); + + run_par!(|| { + iter_par!(output.axis_iter_mut(Axis(0))) + .enumerate() + .for_each( + #[inline(never)] + |(k, mut output)| { + let b = k / out_channels; + let oc = k % out_channels; + let g = oc / out_c_per_group; + + for ic in (in_channels * g)..(in_channels * (g + 1)) { + let weight_ic = ic - (g * in_channels); + + let x = x.slice(s![b, ic, .., .., ..]); + let k = weights.slice(s![oc, weight_ic, .., .., ..]); + + for kd in 0..kernel_depth { + for kh in 0..kernel_height { + for kw in 0..kernel_width { + let k = k[[kd, kh, kw]]; + + // NOTE: This function call is duplicated twice so that the compiler can perform auto-vectorization + // in the case that the stride/dilation is 1. + #[allow(clippy::if_same_then_else)] + if (1, 1, 1, 1, 1, 1) + == ( + stride_width, + stride_height, + stride_depth, + dilation_width, + dilation_height, + dilation_depth, + ) + { + conv3d_mad_inner( + output.view_mut(), + x.view(), + k, + (kd, kh, kw), + (out_width, out_height, out_depth), + (stride_width, stride_height, stride_depth), + (dilation_width, dilation_height, dilation_depth), + ); + } else { + conv3d_mad_inner( + output.view_mut(), + x.view(), + k, + (kd, kh, kw), + (out_width, out_height, out_depth), + (stride_width, stride_height, stride_depth), + (dilation_width, dilation_height, dilation_depth), + ); + } + } + } + } + } + + if let Some(bias) = &bias { + let bias = bias[oc]; + + // Get a mutable iterator to the row we're looping over. + let orows = output.rows_mut(); + for mut or in orows { + // We explicitly define the bounds to 0..out_width so that rustc can make + // the assumption that all accesses are in-bounds. + let or = &mut or.as_slice_mut().unwrap()[0..out_width]; + + #[allow(clippy::needless_range_loop)] + for ow in 0..out_width { + or[ow] += bias; + } + } + } + }, + ); + }); + + output + .to_shape([batch_size, out_channels, out_depth, out_height, out_width]) + .unwrap() + .into_dyn() + .into_shared() +} + +pub(crate) fn conv_transpose3d( + x: SharedArray, + weight: SharedArray, + bias: Option>, + options: ConvTransposeOptions<3>, +) -> SharedArray { + let [dilation_depth, dilation_height, dilation_width] = options.dilation; + let [padding_depth, padding_height, padding_width] = options.padding; + let [stride_depth, stride_height, stride_width] = options.stride; + let [out_padding_depth, out_padding_height, out_padding_width] = options.padding_out; + let [batch_size, _in_channels, in_depth, in_height, in_width] = x.shape().try_into().unwrap(); + let [ + in_channels, + out_channels, + kernel_depth, + kernel_height, + kernel_width, + ] = weight.shape().try_into().unwrap(); + + let out_depth = calculate_conv_transpose_output_size( + kernel_depth, + stride_depth, + padding_depth, + out_padding_depth, + dilation_depth, + in_depth, + ); + let out_height = calculate_conv_transpose_output_size( + kernel_height, + stride_height, + padding_height, + out_padding_height, + dilation_height, + in_height, + ); + let out_width = calculate_conv_transpose_output_size( + kernel_width, + stride_width, + padding_width, + out_padding_width, + dilation_width, + in_width, + ); + + let x = x; + let mut output = Array5::zeros(Dim([ + batch_size, + out_channels * options.groups, + out_depth, + out_height, + out_width, + ])); + + let unsafe_shared_out = UnsafeSharedRef::new(&mut output); + + run_par!(|| { + iter_range_par!(0, batch_size * out_channels * options.groups).for_each(|k| unsafe { + let b = k / (out_channels * options.groups); + let oc = k % out_channels; + let g = (k / out_channels) % options.groups; + + let output = unsafe_shared_out.get(); + + let oc_out = oc + (out_channels * g); + let ic_start = g * (in_channels / options.groups); + let ic_end = ic_start + in_channels / options.groups; + + for ic in ic_start..ic_end { + for id in 0..in_depth { + for ih in 0..in_height { + for iw in 0..in_width { + for kd in 0..kernel_depth { + for kh in 0..kernel_height { + for kw in 0..kernel_width { + let od = id * stride_depth + kd * dilation_depth; + let oh = ih * stride_height + kh * dilation_height; + let ow = iw * stride_width + kw * dilation_width; + + if od >= out_depth + padding_depth + || oh >= out_height + padding_height + || ow >= out_width + padding_width + || od < padding_depth + || oh < padding_height + || ow < padding_width + { + continue; + } + + let od = od - padding_depth; + let oh = oh - padding_height; + let ow = ow - padding_width; + + output[[b, oc_out, od, oh, ow]] += + x[[b, ic, id, ih, iw]] * weight[[ic, oc, kd, kh, kw]]; + } + } + } + } + } + } + } + + if let Some(bias) = &bias { + for od in 0..out_depth { + for oh in 0..out_height { + for ow in 0..out_width { + output[[b, oc_out, od, oh, ow]] += bias[oc_out]; + } + } + } + } + }); + }); + + output.into_dyn().into_shared() +} diff --git a/crates/burn-ndarray/src/ops/deform_conv.rs b/crates/burn-ndarray/src/ops/deform_conv.rs new file mode 100644 index 0000000..390010b --- /dev/null +++ b/crates/burn-ndarray/src/ops/deform_conv.rs @@ -0,0 +1,662 @@ +use burn_backend::ops::{DeformConvOptions, conv::calculate_conv_output_size}; +use core::ops::AddAssign; +use ndarray::{ + Array2, Array4, ArrayView2, ArrayView3, ArrayView4, ArrayView6, ArrayViewMut2, Axis, Dim, Ix4, + Zip, s, +}; + +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float; + +use crate::{FloatNdArrayElement, NdArrayTensor, ShapeOps, SharedArray, iter_par, run_par}; + +use super::matmul::matmul; + +#[inline(always)] +#[allow(clippy::too_many_arguments)] +fn deform_im2col_kernel( + out_y: usize, + out_x: usize, + input: ArrayView2, + offset: ArrayView3, + mask: Option>, + mut columns: ArrayViewMut2, + args: DeformConvOptions<2>, + (kernel_h, kernel_w): (usize, usize), +) { + // position shape: [in_channels, batch_size, out_h, out_w] + // columns shape: [[in_channels, kernel_h, kernel_w], [batch_size, out_h, out_w]] + + let (height, width) = input.dim(); + + for kernel_y in 0..kernel_h { + for kernel_x in 0..kernel_w { + let mask_value = mask + .map(|it| it[[kernel_y, kernel_x]]) + .unwrap_or_else(|| F::from_elem(1.0)); + + let offset = offset.slice(s![kernel_y, kernel_x, ..]); + let y = F::from_elem(out_y * args.stride[0] + kernel_y * args.dilation[0]) + - F::from_elem(args.padding[0]) + + offset[0]; + let x = F::from_elem(out_x * args.stride[1] + kernel_x * args.dilation[1]) + - F::from_elem(args.padding[1]) + + offset[1]; + + let interpolated = bilinear_interpolate(input, height, width, y, x); + + columns[[kernel_y, kernel_x]] = mask_value * interpolated; + } + } +} + +fn bilinear_interpolate( + input: ArrayView2, + height: usize, + width: usize, + y: F, + x: F, +) -> F { + // To simplify code + let y = y.to_f32(); + let x = x.to_f32(); + + let mut result = F::from_elem(0.0); + if y > -1.0 && height as f32 > y && x > -1.0 && width as f32 > x { + let y_low = f32::floor(y); + let x_low = f32::floor(x); + let y_high = (y_low + 1.) as usize; + let x_high = (x_low + 1.) as usize; + + let zero = F::from_elem(0.0); + let v1: F = if y_low >= 0. && x_low >= 0. { + input[[y_low as usize, x_low as usize]] + } else { + zero + }; + let v2: F = if y_low >= 0. && x_high < width { + input[[y_low as usize, x_high]] + } else { + zero + }; + let v3: F = if y_high < height && x_low >= 0. { + input[[y_high, x_low as usize]] + } else { + zero + }; + let v4: F = if y_high < height && x_high < width { + input[[y_high, x_high]] + } else { + zero + }; + + let l_y = y - y_low; + let l_x = x - x_low; + let h_y = 1.0 - l_y; + let h_x = 1.0 - l_x; + + let w1 = F::from_elem(h_y * h_x); + let w2 = F::from_elem(h_y * l_x); + let w3 = F::from_elem(l_y * h_x); + let w4 = F::from_elem(l_y * l_x); + + result = w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4; + } + result +} + +pub(crate) fn deform_conv2d( + input: SharedArray, + offset: SharedArray, + weight: SharedArray, + mask: Option>, + bias: Option>, + args: DeformConvOptions<2>, +) -> SharedArray +where + NdArrayTensor: From>, +{ + let [batch_size, _, in_height, in_width] = input.shape().dims(); + let [out_channels, _, kernel_h, kernel_w] = weight.shape().dims(); + let groups = args.weight_groups; + + let weight = weight.as_standard_layout(); + + let out_h = calculate_conv_output_size( + kernel_h, + args.stride[0], + args.padding[0], + args.dilation[0], + in_height, + ); + let out_w = calculate_conv_output_size( + kernel_w, + args.stride[1], + args.padding[1], + args.dilation[1], + in_width, + ); + let out_dims = (out_h, out_w); + + let input = input.into_dimensionality::().unwrap(); + let offset = offset.into_dimensionality::().unwrap(); + let mask = mask.as_ref().map(|it| { + it.to_shape(( + batch_size, + args.offset_groups, + kernel_h, + kernel_w, + out_h, + out_w, + )) + .unwrap() + }); + + let columns = deform_im2col( + input.view(), + offset.view(), + mask.as_ref().map(|it| it.view()), + args, + out_dims, + (kernel_h, kernel_w), + ); + + let (col_size_0, col_size_1) = columns.dim(); + let col_size_0 = col_size_0 / groups; + let out_c_per_group = out_channels / groups; + + let weight = weight + .to_shape((groups, out_c_per_group, col_size_0)) + .unwrap(); + let columns = columns.to_shape((groups, col_size_0, col_size_1)).unwrap(); + let out = matmul( + weight.to_owned().into_dyn().into_shared(), + columns.to_owned().into_dyn().into_shared(), + ); + + let mut out = out + .into_shape_with_order((out_channels, batch_size, out_h, out_w)) + .unwrap(); + out.swap_axes(0, 1); + + if let Some(bias) = bias { + let bias = bias.to_shape((1, out_channels, 1, 1)).unwrap(); + out.add_assign(&bias); + } + + out.into_dyn().into_shared() +} + +pub(crate) fn deform_im2col( + input: ArrayView4, + offset: ArrayView4, + mask: Option>, + args: DeformConvOptions<2>, + out_dims: (usize, usize), + kernel_dims: (usize, usize), +) -> Array2 { + let (batch_size, in_channels, _, _) = input.dim(); + let (kernel_h, kernel_w) = kernel_dims; + let (out_h, out_w) = out_dims; + let channels_per_offset_group = in_channels / args.offset_groups; + + let mut columns = Array4::zeros(Dim([ + in_channels, + kernel_h, + kernel_w, + batch_size * out_h * out_w, + ])); + + let groups = args.offset_groups; + + run_par!(|| { + iter_par!(columns.axis_iter_mut(Axis(3))) + .enumerate() + .for_each(|(index, mut columns)| { + let out_x = index % out_w; + let out_y = (index / out_w) % out_h; + let batch = (index / (out_w * out_h)) % batch_size; + let offset = offset.slice(s![batch, .., out_y, out_x]); + let offset = offset.to_shape((groups, kernel_h, kernel_w, 2)).unwrap(); + let mask = mask + .as_ref() + .map(|it| it.slice(s![batch, .., .., .., out_y, out_x])); + columns + .axis_iter_mut(Axis(0)) + .enumerate() + .for_each(|(in_channel, mut columns)| { + let group_index = in_channel / channels_per_offset_group; + deform_im2col_kernel( + out_y, + out_x, + input.slice(s![batch, in_channel, .., ..]), + offset.slice(s![group_index, .., .., ..]), + mask.as_ref().map(|it| it.slice(s![group_index, .., ..])), + columns.view_mut(), + args.clone(), + kernel_dims, + ); + }); + }); + }); + + columns + // Columns is created here, so we know it's contiguous + .into_shape_with_order(( + in_channels * kernel_h * kernel_w, + batch_size * out_h * out_w, + )) + .unwrap() +} + +pub mod backward { + #[cfg(target_has_atomic = "32")] + use core::sync::atomic::Ordering; + + use atomic_float::AtomicF32; + use ndarray::{Array1, Array5, ArrayView4, ArrayView6, Ix4}; + + use super::*; + + pub(crate) type DeformConv2dBackward = ( + SharedArray, + SharedArray, + SharedArray, + Option>, + Option>, + ); + + /// Calculate the [deformable 2D convolution](crate::ops::ModuleOps::deform_conv2d) backward pass using convolutions. + pub(crate) fn deform_conv2d_backward( + input: SharedArray, + offset: SharedArray, + weight: SharedArray, + mask: Option>, + bias: Option>, + out_grad: SharedArray, + args: DeformConvOptions<2>, + ) -> DeformConv2dBackward { + let [batch_size, out_channels, out_h, out_w] = out_grad.shape().dims(); + let [_, _, kernel_h, kernel_w] = weight.shape().dims(); + let groups = args.weight_groups; + let out_c_per_group = out_channels / groups; + let col_shape_1 = batch_size * out_h * out_w; + let mut out_grad = out_grad.into_dimensionality::().unwrap(); + + let gradient_bias = bias.map(|_| { + let out_grad = out_grad + .clone() + .sum_axis(Axis(0)) + .sum_axis(Axis(1)) + .sum_axis(Axis(1)); + + out_grad.into_dyn().into_shared() + }); + + out_grad.swap_axes(0, 1); + let out_grad = out_grad + .to_shape((groups, out_c_per_group, col_shape_1)) + .unwrap(); + + let input = input.into_dimensionality::().unwrap(); + let offset = offset.into_dimensionality::().unwrap(); + let mask = mask.map(|it| { + it.into_shape_with_order(( + batch_size, + args.offset_groups, + kernel_h, + kernel_w, + out_h, + out_w, + )) + .unwrap() + }); + + let (input_gradient, offset_gradient, mask_gradient) = backward_gradient_inputs( + input.view(), + weight, + offset.view(), + mask.as_ref().map(|it| it.view()), + out_grad.view(), + &args, + (kernel_h, kernel_w), + ); + + let weight_grad = compute_weight_grad( + input.view(), + offset.view(), + mask.as_ref().map(|it| it.view()), + out_grad.view(), + args, + (kernel_h, kernel_w), + (out_h, out_w), + ); + + ( + input_gradient, + offset_gradient, + weight_grad, + mask_gradient, + gradient_bias, + ) + } + + fn compute_weight_grad( + input: ArrayView4, + offset: ArrayView4, + mask: Option>, + out_grad: ArrayView3, + options: DeformConvOptions<2>, + kernel_dims: (usize, usize), + out_dims: (usize, usize), + ) -> SharedArray { + let in_channels = input.dim().1; + let (groups, out_c_per_group, _) = out_grad.dim(); + let (kernel_h, kernel_w) = kernel_dims; + + let in_c_per_group = in_channels / groups; + + let columns = deform_im2col(input, offset, mask, options, out_dims, kernel_dims); + let (col_size_0, col_size_1) = columns.dim(); + let col_size_0 = col_size_0 / groups; + + let mut columns = columns.to_shape((groups, col_size_0, col_size_1)).unwrap(); + columns.swap_axes(1, 2); + + let grad_weight = matmul( + out_grad.to_owned().into_dyn().into_shared(), + columns.to_owned().into_dyn().into_shared(), + ); + + let grad_weight = grad_weight + .into_shape_with_order((out_c_per_group * groups, in_c_per_group, kernel_h, kernel_w)) + .unwrap(); + grad_weight.into_dyn().into_shared() + } + + type InputGradients = (SharedArray, SharedArray, Option>); + + fn backward_gradient_inputs( + image: ArrayView4, + weight: SharedArray, + offset: ArrayView4, + mask: Option>, + out_grad: ArrayView3, + args: &DeformConvOptions<2>, + kernel_dims: (usize, usize), + ) -> InputGradients { + let input_shape = image.dim(); + let in_channels = input_shape.1; + let [out_channels, in_c_per_group, kernel_h, kernel_w] = weight.shape().dims(); + let (batch_size, _, out_h, out_w) = offset.dim(); + + let groups = args.weight_groups; + let out_c_per_group = out_channels / groups; + + let col_shape_0 = in_c_per_group * kernel_h * kernel_w; + + let mut weight = weight + .to_shape((groups, out_c_per_group, col_shape_0)) + .unwrap(); + weight.swap_axes(1, 2); + let columns = matmul( + weight.to_owned().into_dyn().into_shared(), + out_grad.to_owned().into_dyn().into_shared(), + ); + + let columns = columns + .to_shape((in_channels, kernel_h, kernel_w, batch_size, out_h, out_w)) + .unwrap(); + + let (offset_gradient, mask_gradient) = compute_offset_and_mask_gradient( + columns.view(), + image.view(), + offset, + mask, + args, + kernel_dims, + ); + + let input_gradient = + compute_input_grad(columns.view(), offset, mask, args, kernel_dims, input_shape); + + (input_gradient, offset_gradient, mask_gradient) + } + + fn compute_offset_and_mask_gradient( + columns: ArrayView6, + image: ArrayView4, + offset: ArrayView4, + mask: Option>, + args: &DeformConvOptions<2>, + kernel_dims: (usize, usize), + ) -> (SharedArray, Option>) { + let (kernel_h, kernel_w) = kernel_dims; + let (_, in_channels, height, width) = image.dim(); + let (batch_size, offset_channels, out_h, out_w) = offset.dim(); + let offs_groups = args.offset_groups; + let channels_per_offset_group = in_channels / args.offset_groups; + + let mut grad_offset = Array5::zeros(( + offs_groups, + kernel_h, + kernel_w, + 2, + batch_size * out_h * out_w, + )); + let mut grad_mask = + Array4::zeros((offs_groups, kernel_h, kernel_w, batch_size * out_h * out_w)); + + grad_mask + .axis_iter_mut(Axis(3)) + .zip(grad_offset.axis_iter_mut(Axis(4))) + .enumerate() + .for_each(|(index, (mut grad_mask, mut grad_offset))| { + let out_x = index % out_w; + let out_y = (index / out_w) % out_h; + let batch = index / (out_w * out_h); + let offset = offset.slice(s![batch, .., out_y, out_x]); + let offset = offset + .to_shape((offs_groups, kernel_h, kernel_w, 2)) + .unwrap(); + let mask: Option> = mask + .as_ref() + .map(|mask| mask.slice(s![batch, .., .., .., out_y, out_x])); + let columns = columns.slice(s![.., .., .., batch, out_y, out_x]); + let image = image.slice(s![batch, .., .., ..]); + + for ((group, kernel_y, kernel_x), grad_mask) in grad_mask.indexed_iter_mut() { + let grad_mask: &mut F = grad_mask; + let mut grad_offset = grad_offset.slice_mut(s![group, kernel_y, kernel_x, ..]); + let offset = offset.slice(s![group, kernel_y, kernel_x, ..]); + let mask = mask.map(|it| it[[group, kernel_y, kernel_x]]); + let columns = columns.slice(s![.., kernel_y, kernel_x]); + let group_offset = group * channels_per_offset_group; + let image = image.slice(s![group_offset.., .., ..]); + let y = F::from_elem(out_y * args.stride[0] + kernel_y * args.dilation[0]) + - F::from_elem(args.padding[0]) + + offset[0]; + let x = F::from_elem(out_x * args.stride[1] + kernel_x * args.dilation[1]) + - F::from_elem(args.padding[1]) + + offset[1]; + for (i, grad_offset) in grad_offset.iter_mut().enumerate() { + let is_y_direction = i % 2 == 0; + let use_mask = mask.is_some(); + + for channel in 0..channels_per_offset_group { + let mask = mask.unwrap_or_else(|| F::one()); + let image = image.index_axis(Axis(0), channel); + let weight = + get_coordinate_weight(image, height, width, y, x, is_y_direction); + *grad_offset += mask * weight * columns[channel]; + if use_mask && is_y_direction { + *grad_mask += columns[channel] + * bilinear_interpolate(image, height, width, y, x); + } + } + } + } + }); + + let mask_gradient = mask.map(|_| { + let mut grad_mask = grad_mask + .into_shape_with_order((offset_channels / 2, batch_size, out_h, out_w)) + .unwrap(); + grad_mask.swap_axes(0, 1); + grad_mask.into_dyn().into_shared() + }); + let mut grad_offset = grad_offset + .into_shape_with_order((offset_channels, batch_size, out_h, out_w)) + .unwrap(); + grad_offset.swap_axes(0, 1); + let offset_gradient = grad_offset.into_dyn().into_shared(); + (offset_gradient, mask_gradient) + } + + fn get_coordinate_weight( + input: ArrayView2, + height: usize, + width: usize, + y: F, + x: F, + is_y_direction: bool, + ) -> F { + let y = y.to_f32(); + let x = x.to_f32(); + + let y_low = f32::floor(y); + let x_low = f32::floor(x); + let y_high = y_low + 1.; + let x_high = x_low + 1.; + + let valid_y_low = y_low >= 0. && y_low < height as f32; + let valid_y_high = y_high >= 0. && y_high < height as f32; + let valid_x_low = x_low >= 0. && x_low < width as f32; + let valid_x_high = x_high >= 0. && x_high < width as f32; + + let bottom_left = if valid_y_low && valid_x_low { + input[[y_low as usize, x_low as usize]] + } else { + F::zero() + }; + let bottom_right = if valid_y_low && valid_x_high { + input[[y_low as usize, x_high as usize]] + } else { + F::zero() + }; + let top_left = if valid_y_high && valid_x_low { + input[[y_high as usize, x_low as usize]] + } else { + F::zero() + }; + let top_right = if valid_y_high && valid_x_high { + input[[y_high as usize, x_high as usize]] + } else { + F::zero() + }; + + if is_y_direction { + let delta_x = F::from_elem(x - x_low); + delta_x * (top_right - bottom_right) + (F::one() - delta_x) * (top_left - bottom_left) + } else { + let delta_y = F::from_elem(y - y_low); + delta_y * (top_right - top_left) + (F::one() - delta_y) * (bottom_right - bottom_left) + } + } + + fn compute_input_grad( + columns: ArrayView6, + offset: ArrayView4, + mask: Option>, + args: &DeformConvOptions<2>, + kernel_dims: (usize, usize), + input_shape: (usize, usize, usize, usize), + ) -> SharedArray { + let (batch_size, in_channels, height, width) = input_shape; + let (kernel_h, kernel_w) = kernel_dims; + let offs_groups = args.offset_groups; + let channels_per_offset_group = in_channels / offs_groups; + + let grad_in = + Array4::from_shape_simple_fn((batch_size, in_channels, height, width), || { + AtomicF32::new(0.0) + }); + + let compute_for_each = |(in_channel, kernel_y, kernel_x, batch, out_y, out_x), col: &F| { + let group = in_channel / channels_per_offset_group; + let offset = offset.slice(s![batch, .., out_y, out_x]); + let offset = offset + .to_shape((offs_groups, kernel_h, kernel_w, 2)) + .unwrap(); + let offset = offset.slice(s![group, kernel_y, kernel_x, ..]); + let offset = [offset[0], offset[1]]; + let mask = mask + .as_ref() + .map(|it| it[[batch, group, kernel_y, kernel_x, out_y, out_x]].to_f32()); + let y = F::from_elem(out_y * args.stride[0] + kernel_y * args.dilation[0]) + - F::from_elem(args.padding[0]) + + offset[0]; + let x = F::from_elem(out_x * args.stride[1] + kernel_x * args.dilation[1]) + - F::from_elem(args.padding[1]) + + offset[1]; + let grad_in = grad_in.slice(s![batch, in_channel, .., ..]); + deform_col2img_kernel(y.to_f32(), x.to_f32(), mask, col.to_f32(), grad_in); + }; + + // `for_each` expects a 2-tuple argument with `.into_par_iter()`, but 2 separate arguments otherwise + #[cfg(feature = "multi-threads")] + run_par!(|| { + iter_par!(Zip::indexed(columns)) + .for_each(|(args0, args1)| compute_for_each(args0, args1)) + }); + + #[cfg(not(feature = "multi-threads"))] + run_par!(|| { iter_par!(Zip::indexed(columns)).for_each(&compute_for_each) }); + + let grad_in: Array1 = grad_in + .into_iter() + .map(|it| F::from_elem(it.into_inner())) + .collect(); + let grad_in = grad_in + .into_shape_with_order((batch_size, in_channels, height, width)) + .unwrap(); + grad_in.into_dyn().into_shared() + } + + fn deform_col2img_kernel( + y: f32, + x: f32, + mask: Option, + col: f32, + grad_input: ArrayView2, + ) { + let (height, width) = grad_input.dim(); + let mask_value = mask.unwrap_or(1.0); + + for dy in -1..=1 { + for dx in -1..=1 { + let yp = f32::floor(y) + dy as f32; + let xp = f32::floor(x) + dx as f32; + + if yp >= 0.0 + && yp < height as f32 + && xp >= 0.0 + && xp < width as f32 + && f32::abs(y - yp) < 1.0 + && f32::abs(x - xp) < 1.0 + { + let weight = (1.0 - f32::abs(y - yp)) * (1.0 - f32::abs(x - xp)); + + #[cfg_attr(not(target_has_atomic = "32"), allow(unused))] + let value = mask_value * weight * col; + + #[cfg(target_has_atomic = "32")] + grad_input[[yp as usize, xp as usize]].fetch_add(value, Ordering::AcqRel); + #[cfg(not(target_has_atomic = "32"))] + panic!("Can't use deformable convolution backwards pass without atomics"); + } + } + } + } +} diff --git a/crates/burn-ndarray/src/ops/grid_sample.rs b/crates/burn-ndarray/src/ops/grid_sample.rs new file mode 100644 index 0000000..256c2fd --- /dev/null +++ b/crates/burn-ndarray/src/ops/grid_sample.rs @@ -0,0 +1,214 @@ +use burn_backend::ElementConversion; +use burn_backend::ops::{GridSampleOptions, GridSamplePaddingMode, InterpolateMode}; +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float; + +use ndarray::Array4; + +use crate::SharedArray; +use crate::{FloatNdArrayElement, UnsafeSharedRef, iter_range_par, run_par}; + +/// Sample a tensor using grid-based sampling. +/// +/// # Arguments +/// +/// * `tensor` - The tensor being sampled from, must be contiguous with shape (N, C, H_in, W_in) +/// * `grid` - A tensor of locations, with shape (N, H_out, W_out, 2). Values are [-1, 1]. +/// A [x = -1, y = -1] means top-left, and [x = 1, y = 1] means bottom-right +/// * `options` - Grid sampling options (mode, padding_mode, align_corners) +/// +/// # Returns +/// +/// A tensor with shape (N, C, H_out, W_out) +pub(crate) fn grid_sample_2d( + tensor: SharedArray, + grid: SharedArray, + options: GridSampleOptions, +) -> SharedArray { + match options.mode { + InterpolateMode::Bilinear => (), + _ => todo!( + "grid_sample_2d with {:?} mode is not implemented", + options.mode + ), + } + + let tensor = tensor.into_dimensionality::().unwrap(); + let grid = grid.into_dimensionality::().unwrap(); + + let (batch_size, channels, height_in, width_in) = tensor.dim(); + let (b, height_out, width_out, d) = grid.dim(); + assert!(batch_size == b); + assert!(2 == d); + + let mut output = Array4::zeros((batch_size, channels, height_out, width_out)); + let unsafe_shared_out = UnsafeSharedRef::new(&mut output); + + let sample_count = batch_size * channels * height_out * width_out; + let strides = ( + channels * height_out * width_out, + height_out * width_out, + width_out, + ); + + let align = options.align_corners; + let pad_mode = options.padding_mode; + + run_par!(|| { + iter_range_par!(0, sample_count).for_each(|id| { + let (b, c, y, x) = ( + id / strides.0, + id % strides.0 / strides.1, + id % strides.1 / strides.2, + id % strides.2, + ); + + let sample_x = grid[(b, y, x, 0)].elem::(); + let sample_y = grid[(b, y, x, 1)].elem::(); + + // Convert normalized grid coordinates [-1, 1] to pixel coordinates + let (px, py) = if align { + // align_corners=true: x_pixel = (x_norm + 1) * (width - 1) / 2 + // Maps -1 to 0 and 1 to width - 1 + let px = (sample_x + 1.0) * ((width_in - 1) as f64) / 2.0; + let py = (sample_y + 1.0) * ((height_in - 1) as f64) / 2.0; + (px, py) + } else { + // align_corners=false: x_pixel = (x_norm + 1) * width / 2 - 0.5 + // Maps -1 to -0.5 and 1 to width - 0.5 + let px = (sample_x + 1.0) * (width_in as f64) / 2.0 - 0.5; + let py = (sample_y + 1.0) * (height_in as f64) / 2.0 - 0.5; + (px, py) + }; + + // Bilinear interpolation with the specified padding mode + let val = + bilinear_interpolate(&tensor, b, c, px, py, width_in, height_in, pad_mode, align); + + unsafe { + let output = unsafe_shared_out.get(); + output[(b, c, y, x)] = val.elem(); + } + }); + }); + + output.into_dyn().into_shared() +} + +/// Bilinear interpolation at a point with configurable padding mode. +#[allow(clippy::too_many_arguments)] +fn bilinear_interpolate( + source: &ndarray::ArrayBase>, + b: usize, + c: usize, + x: f64, + y: f64, + width: usize, + height: usize, + padding_mode: GridSamplePaddingMode, + align_corners: bool, +) -> f64 +where + E: FloatNdArrayElement, + S: ndarray::Data, +{ + // Handle inf/nan coordinates + if !x.is_finite() || !y.is_finite() { + return match padding_mode { + GridSamplePaddingMode::Zeros => 0.0, + GridSamplePaddingMode::Border => { + // Clamp to center of image for inf/nan + let cx = ((width - 1) as f64 / 2.0).clamp(0.0, (width - 1) as f64); + let cy = ((height - 1) as f64 / 2.0).clamp(0.0, (height - 1) as f64); + source[(b, c, cy as usize, cx as usize)].elem::() + } + GridSamplePaddingMode::Reflection => 0.0, // Simplified: treat as zeros for inf/nan + }; + } + + // Apply padding mode to get actual sampling coordinates + let (x, y) = match padding_mode { + GridSamplePaddingMode::Border => { + // Clamp coordinates to valid range [0, size-1] + let x = x.clamp(0.0, (width - 1) as f64); + let y = y.clamp(0.0, (height - 1) as f64); + (x, y) + } + GridSamplePaddingMode::Reflection => { + // Reflect coordinates at boundaries + let x = reflect_coordinate(x, width, align_corners); + let y = reflect_coordinate(y, height, align_corners); + (x, y) + } + GridSamplePaddingMode::Zeros => (x, y), // Keep as-is, handle out-of-bounds in read + }; + + // Get the four corner indices + let x0 = x.floor() as i64; + let y0 = y.floor() as i64; + let x1 = x0.saturating_add(1); + let y1 = y0.saturating_add(1); + + // Compute interpolation weights (fractional part) + let x_frac = x - x.floor(); + let y_frac = y - y.floor(); + + // Helper to read a value based on padding mode + let read_value = |xi: i64, yi: i64| -> f64 { + match padding_mode { + GridSamplePaddingMode::Zeros => { + // Return 0 for out-of-bounds + if xi >= 0 && xi < width as i64 && yi >= 0 && yi < height as i64 { + source[(b, c, yi as usize, xi as usize)].elem::() + } else { + 0.0 + } + } + GridSamplePaddingMode::Border | GridSamplePaddingMode::Reflection => { + // Coordinates should already be in valid range after clamping/reflection + let xi = xi.clamp(0, (width - 1) as i64) as usize; + let yi = yi.clamp(0, (height - 1) as i64) as usize; + source[(b, c, yi, xi)].elem::() + } + } + }; + + // Read the four corners + let v00 = read_value(x0, y0); + let v01 = read_value(x0, y1); + let v10 = read_value(x1, y0); + let v11 = read_value(x1, y1); + + // Bilinear interpolation weights + let w00 = (1.0 - x_frac) * (1.0 - y_frac); + let w01 = (1.0 - x_frac) * y_frac; + let w10 = x_frac * (1.0 - y_frac); + let w11 = x_frac * y_frac; + + v00 * w00 + v01 * w01 + v10 * w10 + v11 * w11 +} + +/// Reflect a coordinate at the boundaries using a triangle wave pattern. +/// +/// For align_corners=true: reflects within [0, size-1] +/// For align_corners=false: reflects within [-0.5, size-0.5] +fn reflect_coordinate(coord: f64, size: usize, align_corners: bool) -> f64 { + let size_f = size as f64; + let (min_val, max_val) = if align_corners { + (0.0, size_f - 1.0) + } else { + (-0.5, size_f - 0.5) + }; + + let span = max_val - min_val; + if span <= 0.0 { + return min_val; + } + + // Triangle wave formula: span - |((x mod 2*span) - span)| + let period = 2.0 * span; + let x = (coord - min_val).abs(); + let x_mod = x - (x / period).floor() * period; + span - (x_mod - span).abs() + min_val +} diff --git a/crates/burn-ndarray/src/ops/int_tensor.rs b/crates/burn-ndarray/src/ops/int_tensor.rs new file mode 100644 index 0000000..ddc676b --- /dev/null +++ b/crates/burn-ndarray/src/ops/int_tensor.rs @@ -0,0 +1,519 @@ +// Language +use crate::rand::get_seeded_rng; +use alloc::vec::Vec; +use burn_backend::backend::ExecutionError; +use burn_backend::ops::IntTensorOps; +use burn_backend::tensor::{FloatTensor, IntTensor}; +use burn_backend::{Distribution, IntDType, Scalar, TensorMetadata}; + +use burn_backend::ElementConversion; +use burn_std::{BoolDType, FloatDType}; + +// Current crate +use crate::SharedArray; +use crate::execute_with_int_dtype; +use crate::ops::matmul::matmul; +use crate::{ExpElement, NdArrayDevice, SEED, execute_with_int_out_dtype, slice}; +use crate::{NdArray, cast_to_dtype, execute_with_dtype, tensor::NdArrayTensor}; +use crate::{cat_with_dtype, execute_with_float_out_dtype}; + +// Workspace crates +use super::{NdArrayBitOps, NdArrayMathOps, NdArrayOps}; +use burn_backend::{DType, Shape, TensorData}; + +impl IntTensorOps for NdArray { + fn int_from_data(data: TensorData, _device: &NdArrayDevice) -> NdArrayTensor { + if data.dtype.is_int() || data.dtype.is_uint() { + NdArrayTensor::from_data(data) + } else { + unimplemented!("Unsupported dtype for `int_from_data`: {:?}", data.dtype) + } + } + + async fn int_into_data(tensor: NdArrayTensor) -> Result { + Ok(tensor.into_data()) + } + + fn int_to_device(tensor: NdArrayTensor, _device: &NdArrayDevice) -> NdArrayTensor { + tensor + } + + fn int_reshape(tensor: NdArrayTensor, shape: Shape) -> NdArrayTensor { + execute_with_int_dtype!(tensor, |array| NdArrayOps::reshape(array, shape)) + } + + fn int_slice(tensor: NdArrayTensor, slices: &[burn_backend::Slice]) -> NdArrayTensor { + slice!(tensor, slices) + } + + fn int_empty(shape: Shape, device: &NdArrayDevice, dtype: IntDType) -> NdArrayTensor { + Self::int_zeros(shape, device, dtype) + } + + fn int_matmul(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + execute_with_int_dtype!((lhs, rhs), matmul) + } + + fn int_mask_where( + tensor: NdArrayTensor, + mask: NdArrayTensor, + source: NdArrayTensor, + ) -> NdArrayTensor { + execute_with_int_dtype!((tensor, source), |tensor, source| { + NdArrayOps::mask_where(tensor, mask.bool(), source) + }) + } + + fn int_mask_fill(tensor: NdArrayTensor, mask: NdArrayTensor, value: Scalar) -> NdArrayTensor { + execute_with_int_dtype!(tensor, |array| NdArrayOps::mask_fill( + array, + mask.bool(), + value.elem() + )) + } + + fn int_slice_assign( + tensor: NdArrayTensor, + slices: &[burn_backend::Slice], + value: NdArrayTensor, + ) -> NdArrayTensor { + execute_with_int_dtype!((tensor, value), |tensor, value| NdArrayOps::slice_assign( + tensor, slices, value + )) + } + + fn int_cat(tensors: Vec, dim: usize) -> NdArrayTensor { + cat_with_dtype!(tensors, dim, [I64, I32, I16, I8, U64, U32, U16, U8]) + } + + fn int_equal(lhs: NdArrayTensor, rhs: NdArrayTensor, _out_dtype: BoolDType) -> NdArrayTensor { + execute_with_int_dtype!((lhs, rhs), NdArrayMathOps::equal) + } + + fn int_equal_elem(lhs: NdArrayTensor, rhs: Scalar, _out_dtype: BoolDType) -> NdArrayTensor { + execute_with_int_dtype!(lhs, |array| NdArrayMathOps::equal_elem(array, rhs.elem())) + } + + fn int_greater(lhs: NdArrayTensor, rhs: NdArrayTensor, _out_dtype: BoolDType) -> NdArrayTensor { + execute_with_int_dtype!((lhs, rhs), NdArrayMathOps::greater) + } + + fn int_greater_elem(lhs: NdArrayTensor, rhs: Scalar, _out_dtype: BoolDType) -> NdArrayTensor { + execute_with_int_dtype!(lhs, |array| NdArrayMathOps::greater_elem(array, rhs.elem())) + } + + fn int_greater_equal( + lhs: NdArrayTensor, + rhs: NdArrayTensor, + _out_dtype: BoolDType, + ) -> NdArrayTensor { + execute_with_int_dtype!((lhs, rhs), NdArrayMathOps::greater_equal) + } + + fn int_greater_equal_elem( + lhs: NdArrayTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> NdArrayTensor { + execute_with_int_dtype!(lhs, |array| NdArrayMathOps::greater_equal_elem( + array, + rhs.elem() + )) + } + + fn int_lower(lhs: NdArrayTensor, rhs: NdArrayTensor, _out_dtype: BoolDType) -> NdArrayTensor { + execute_with_int_dtype!((lhs, rhs), NdArrayMathOps::lower) + } + + fn int_lower_elem(lhs: NdArrayTensor, rhs: Scalar, _out_dtype: BoolDType) -> NdArrayTensor { + execute_with_int_dtype!(lhs, |array| NdArrayMathOps::lower_elem(array, rhs.elem())) + } + + fn int_lower_equal( + lhs: NdArrayTensor, + rhs: NdArrayTensor, + _out_dtype: BoolDType, + ) -> NdArrayTensor { + execute_with_int_dtype!((lhs, rhs), NdArrayMathOps::lower_equal) + } + + fn int_lower_equal_elem( + lhs: NdArrayTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> NdArrayTensor { + execute_with_int_dtype!(lhs, |array| NdArrayMathOps::lower_equal_elem( + array, + rhs.elem() + )) + } + + fn int_add(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor { + execute_with_int_dtype!((lhs, rhs), NdArrayMathOps::add) + } + + fn int_add_scalar(lhs: NdArrayTensor, rhs: Scalar) -> NdArrayTensor { + execute_with_int_dtype!(lhs, |array| NdArrayMathOps::add_scalar(array, rhs.elem())) + } + + fn int_sub(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor { + execute_with_int_dtype!((lhs, rhs), NdArrayMathOps::sub) + } + + fn int_sub_scalar(lhs: NdArrayTensor, rhs: Scalar) -> NdArrayTensor { + execute_with_int_dtype!(lhs, |array| NdArrayMathOps::sub_scalar(array, rhs.elem())) + } + + fn int_mul(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor { + execute_with_int_dtype!((lhs, rhs), NdArrayMathOps::mul) + } + + fn int_mul_scalar(lhs: NdArrayTensor, rhs: Scalar) -> NdArrayTensor { + execute_with_int_dtype!(lhs, |array| NdArrayMathOps::mul_scalar(array, rhs.elem())) + } + + fn int_div(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor { + execute_with_int_dtype!((lhs, rhs), NdArrayMathOps::div) + } + + fn int_div_scalar(lhs: NdArrayTensor, rhs: Scalar) -> NdArrayTensor { + execute_with_int_dtype!(lhs, |array| NdArrayMathOps::div_scalar(array, rhs.elem())) + } + + fn int_remainder(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor { + execute_with_int_dtype!((lhs, rhs), NdArrayMathOps::remainder) + } + + fn int_remainder_scalar(lhs: NdArrayTensor, rhs: Scalar) -> NdArrayTensor { + execute_with_int_dtype!(lhs, |array| NdArrayMathOps::remainder_scalar( + array, + rhs.elem() + )) + } + + fn int_sum(tensor: NdArrayTensor) -> NdArrayTensor { + // Use view() for zero-copy on borrowed storage + execute_with_int_dtype!(tensor, E, |array: SharedArray| NdArrayMathOps::sum_view( + array.view() + )) + } + + fn int_sum_dim(tensor: NdArrayTensor, dim: usize) -> NdArrayTensor { + execute_with_int_dtype!(tensor, |array| NdArrayMathOps::sum_dim(array, dim)) + } + + fn int_prod(tensor: NdArrayTensor) -> NdArrayTensor { + // Use view() for zero-copy on borrowed storage + execute_with_int_dtype!( + tensor, + E, + |array: SharedArray| NdArrayMathOps::prod_view(array.view()) + ) + } + + fn int_prod_dim(tensor: NdArrayTensor, dim: usize) -> NdArrayTensor { + execute_with_int_dtype!(tensor, |array| NdArrayMathOps::prod_dim(array, dim)) + } + + fn int_mean(tensor: NdArrayTensor) -> NdArrayTensor { + // Use view() for zero-copy on borrowed storage + execute_with_int_dtype!( + tensor, + E, + |array: SharedArray| NdArrayMathOps::mean_view(array.view()) + ) + } + + fn int_mean_dim(tensor: NdArrayTensor, dim: usize) -> NdArrayTensor { + execute_with_int_dtype!(tensor, |array| NdArrayMathOps::mean_dim(array, dim)) + } + + fn int_max(tensor: NdArrayTensor) -> NdArrayTensor { + // Use view() for zero-copy on borrowed storage + execute_with_int_dtype!(tensor, E, |array: SharedArray| NdArrayMathOps::max_view( + array.view() + )) + } + + fn int_min(tensor: NdArrayTensor) -> NdArrayTensor { + // Use view() for zero-copy on borrowed storage + execute_with_int_dtype!(tensor, E, |array: SharedArray| NdArrayMathOps::min_view( + array.view() + )) + } + + fn int_cumsum(tensor: NdArrayTensor, dim: usize) -> NdArrayTensor { + execute_with_int_dtype!(tensor, |array| NdArrayMathOps::cumsum(array, dim)) + } + + fn int_cumprod(tensor: NdArrayTensor, dim: usize) -> NdArrayTensor { + execute_with_int_dtype!(tensor, |array| NdArrayMathOps::cumprod(array, dim)) + } + + fn int_cummin(tensor: NdArrayTensor, dim: usize) -> NdArrayTensor { + execute_with_int_dtype!(tensor, |array| NdArrayMathOps::cummin(array, dim)) + } + + fn int_cummax(tensor: NdArrayTensor, dim: usize) -> NdArrayTensor { + execute_with_int_dtype!(tensor, |array| NdArrayMathOps::cummax(array, dim)) + } + + fn int_gather(dim: usize, tensor: NdArrayTensor, indices: NdArrayTensor) -> NdArrayTensor { + execute_with_int_dtype!(tensor, E, |array| -> NdArrayTensor { + execute_with_int_dtype!(indices, |idx_array| NdArrayOps::gather( + dim, array, idx_array + )) + }) + } + + fn int_scatter_add( + dim: usize, + tensor: NdArrayTensor, + indices: NdArrayTensor, + value: NdArrayTensor, + ) -> NdArrayTensor { + execute_with_int_dtype!((tensor, value), I, |tensor, value| -> NdArrayTensor { + execute_with_int_dtype!(indices, |idx_array| NdArrayOps::::scatter( + dim, tensor, idx_array, value + )) + }) + } + + fn int_scatter_nd( + data: NdArrayTensor, + indices: NdArrayTensor, + values: NdArrayTensor, + reduction: burn_backend::tensor::IndexingUpdateOp, + ) -> NdArrayTensor { + execute_with_int_dtype!((data, values), I, |data, values| -> NdArrayTensor { + execute_with_int_dtype!(indices, |idx_array| NdArrayOps::::scatter_nd( + data, idx_array, values, reduction + )) + }) + } + + fn int_gather_nd(data: NdArrayTensor, indices: NdArrayTensor) -> NdArrayTensor { + execute_with_int_dtype!(data, E, |array| -> NdArrayTensor { + execute_with_int_dtype!(indices, |idx_array| NdArrayOps::gather_nd(array, idx_array)) + }) + } + + fn int_select(tensor: NdArrayTensor, dim: usize, indices: NdArrayTensor) -> NdArrayTensor { + execute_with_int_dtype!(tensor, E, |array| -> NdArrayTensor { + execute_with_int_dtype!(indices, |idx_array| NdArrayMathOps::select( + array, dim, idx_array + )) + }) + } + + fn int_select_add( + tensor: NdArrayTensor, + dim: usize, + indices: NdArrayTensor, + value: NdArrayTensor, + ) -> NdArrayTensor { + execute_with_int_dtype!((tensor, value), I, |tensor, value| -> NdArrayTensor { + execute_with_int_dtype!(indices, |idx_array| NdArrayMathOps::::select_assign( + tensor, dim, idx_array, value + )) + }) + } + fn int_argmax(tensor: NdArrayTensor, dim: usize) -> NdArrayTensor { + // Use view() for zero-copy on borrowed storage + execute_with_int_dtype!(tensor, E, |array: SharedArray| { + NdArrayMathOps::argmax_view::(array.view(), dim) + }) + } + + fn int_argtopk(_tensor: NdArrayTensor, _dim: usize, _k: usize) -> NdArrayTensor { + unimplemented!("argtopk not implemented for ndarray"); + } + + fn int_argmin(tensor: NdArrayTensor, dim: usize) -> NdArrayTensor { + // Use view() for zero-copy on borrowed storage + execute_with_int_dtype!(tensor, E, |array: SharedArray| { + NdArrayMathOps::argmin_view::(array.view(), dim) + }) + } + + fn int_clamp_min(tensor: NdArrayTensor, min: Scalar) -> NdArrayTensor { + execute_with_int_dtype!(tensor, |array| NdArrayMathOps::clamp_min(array, min.elem())) + } + + fn int_clamp_max(tensor: NdArrayTensor, max: Scalar) -> NdArrayTensor { + execute_with_int_dtype!(tensor, |array| NdArrayMathOps::clamp_max(array, max.elem())) + } + + fn int_clamp(tensor: NdArrayTensor, min: Scalar, max: Scalar) -> NdArrayTensor { + execute_with_int_dtype!(tensor, |array| NdArrayMathOps::clamp( + array, + min.elem(), + max.elem() + )) + } + + fn int_abs(tensor: NdArrayTensor) -> NdArrayTensor { + match tensor.dtype() { + DType::I64 | DType::I32 | DType::I16 | DType::I8 => { + execute_with_dtype!(tensor, I, NdArrayMathOps::abs, [ + I64 => i64, I32 => i32, I16 => i16, I8 => i8 + ]) + } + // Already unsigned + DType::U64 | DType::U32 | DType::U16 | DType::U8 => tensor, + other => panic!("Unsupported dtype: {other:?}"), + } + } + + fn int_into_float(tensor: NdArrayTensor, out_dtype: FloatDType) -> FloatTensor { + execute_with_float_out_dtype!(out_dtype, F, { + execute_with_int_dtype!(tensor, IntElem, |array: SharedArray| { + array.mapv(|a: IntElem| a.elem::()).into_shared() + }) + }) + } + + fn int_swap_dims(tensor: NdArrayTensor, dim1: usize, dim2: usize) -> NdArrayTensor { + execute_with_int_dtype!(tensor, |array| NdArrayOps::swap_dims(array, dim1, dim2)) + } + + fn int_random( + shape: Shape, + distribution: Distribution, + device: &NdArrayDevice, + dtype: IntDType, + ) -> NdArrayTensor { + let mut seed = SEED.lock().unwrap(); + let mut rng = seed.take().unwrap_or_else(get_seeded_rng); + + let effective_distribution = if distribution == Distribution::Default { + Distribution::Uniform(0.0, 255.0) // Assuming UniformInt is the integer variant + } else { + distribution + }; + + let tensor = execute_with_int_out_dtype!( + dtype, + I, + Self::int_from_data( + TensorData::random::(shape, effective_distribution, &mut rng), + device, + ) + ); + *seed = Some(rng); + tensor + } + + fn int_powi(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor { + execute_with_int_dtype!((lhs, rhs), I, |lhs, rhs| NdArrayMathOps::elementwise_op( + lhs, + rhs, + |a: &I, b: &I| { (a.elem::().pow(b.elem::())).elem() } + )) + } + + fn int_permute(tensor: NdArrayTensor, axes: &[usize]) -> NdArrayTensor { + execute_with_int_dtype!(tensor, |array| NdArrayOps::permute(array, axes)) + } + + fn int_flip(tensor: NdArrayTensor, axes: &[usize]) -> NdArrayTensor { + execute_with_int_dtype!(tensor, |array| NdArrayOps::flip(array, axes)) + } + + fn int_sign(tensor: NdArrayTensor) -> NdArrayTensor { + match tensor.dtype() { + DType::I64 | DType::I32 | DType::I16 | DType::I8 => { + execute_with_dtype!(tensor, I, NdArrayMathOps::sign_op, [ + I64 => i64, I32 => i32, I16 => i16, I8 => i8 + ]) + } + DType::U64 | DType::U32 | DType::U16 | DType::U8 => { + Self::int_greater_elem(tensor, 0.into(), BoolDType::Native) + } + other => panic!("Unsupported dtype: {other:?}"), + } + } + + fn int_expand(tensor: NdArrayTensor, shape: Shape) -> NdArrayTensor { + execute_with_int_dtype!(tensor, |array| NdArrayOps::expand(array, shape)) + } + + fn bitwise_and(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor { + execute_with_int_dtype!((lhs, rhs), NdArrayBitOps::bitand) + } + + fn bitwise_and_scalar(lhs: NdArrayTensor, rhs: Scalar) -> NdArrayTensor { + execute_with_int_dtype!(lhs, |array| NdArrayBitOps::bitand_scalar(array, rhs.elem())) + } + + fn bitwise_or(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor { + execute_with_int_dtype!((lhs, rhs), NdArrayBitOps::bitor) + } + + fn bitwise_or_scalar(lhs: NdArrayTensor, rhs: Scalar) -> NdArrayTensor { + execute_with_int_dtype!(lhs, |array| NdArrayBitOps::bitor_scalar(array, rhs.elem())) + } + + fn bitwise_xor(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor { + execute_with_int_dtype!((lhs, rhs), NdArrayBitOps::bitxor) + } + + fn bitwise_xor_scalar(lhs: NdArrayTensor, rhs: Scalar) -> NdArrayTensor { + execute_with_int_dtype!(lhs, |array| NdArrayBitOps::bitxor_scalar(array, rhs.elem())) + } + + fn bitwise_not(tensor: NdArrayTensor) -> NdArrayTensor { + execute_with_int_dtype!(tensor, NdArrayBitOps::bitnot) + } + + fn bitwise_left_shift(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor { + execute_with_int_dtype!((lhs, rhs), I, |lhs, rhs| { + NdArrayMathOps::elementwise_op(lhs, rhs, |a: &I, b: &I| { + (a.elem::() << (b.elem::())).elem() + }) + }) + } + + fn bitwise_left_shift_scalar(lhs: NdArrayTensor, rhs: Scalar) -> NdArrayTensor { + execute_with_int_dtype!(lhs, I, |array| { + NdArrayMathOps::elementwise_op_scalar(array, |a: I| { + (a.elem::() << rhs.elem::()).elem() + }) + }) + } + + fn bitwise_right_shift(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor { + execute_with_int_dtype!((lhs, rhs), I, |lhs, rhs| { + NdArrayMathOps::elementwise_op(lhs, rhs, |a: &I, b: &I| { + (a.elem::() >> (b.elem::())).elem() + }) + }) + } + + fn bitwise_right_shift_scalar(lhs: NdArrayTensor, rhs: Scalar) -> NdArrayTensor { + execute_with_int_dtype!(lhs, I, |array| { + NdArrayMathOps::elementwise_op_scalar(array, |a: I| { + (a.elem::() >> rhs.elem::()).elem() + }) + }) + } + + fn int_cast(tensor: IntTensor, dtype: IntDType) -> IntTensor { + execute_with_int_dtype!(tensor, |array| cast_to_dtype(array, dtype.into())) + } + + fn int_unfold( + tensor: IntTensor, + dim: usize, + size: usize, + step: usize, + ) -> IntTensor { + execute_with_int_dtype!(tensor, |array| NdArrayOps::unfold(array, dim, size, step)) + } + + fn int_powi_scalar_impl(lhs: IntTensor, rhs: Scalar) -> IntTensor { + execute_with_int_dtype!(lhs, I, |array| { + NdArrayMathOps::elementwise_op_scalar(array, |a: I| a.powi_elem(rhs.elem())) + }) + } +} diff --git a/crates/burn-ndarray/src/ops/interpolate.rs b/crates/burn-ndarray/src/ops/interpolate.rs new file mode 100644 index 0000000..af9d50d --- /dev/null +++ b/crates/burn-ndarray/src/ops/interpolate.rs @@ -0,0 +1,397 @@ +use burn_backend::ElementConversion; +use ndarray::{Array4, ArrayBase, DataOwned}; +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float; + +use crate::{FloatNdArrayElement, ShapeOps, SharedArray, UnsafeSharedRef, iter_range_par, run_par}; + +pub(crate) fn nearest_interpolate( + x: SharedArray, + output_size: [usize; 2], +) -> SharedArray { + let x = x.into_dimensionality::().unwrap(); + + let (batch_size, channels, in_height, in_width) = x.dim(); + let [out_height, out_width] = output_size; + + let y_ratio = (in_height as f64) / (out_height as f64); + let x_ratio = (in_width as f64) / (out_width as f64); + + let out_element_num = batch_size * channels * out_height * out_width; + let strides = ( + channels * out_height * out_width, + out_height * out_width, + out_width, + ); + + let mut output = Array4::zeros((batch_size, channels, out_height, out_width)); + let unsafe_shared_out = UnsafeSharedRef::new(&mut output); + + run_par!(|| { + iter_range_par!(0, out_element_num).for_each(|id| { + let (b, c, h, w) = ( + id / strides.0, + id % strides.0 / strides.1, + id % strides.1 / strides.2, + id % strides.2, + ); + + let y_in = (y_ratio * h as f64).floor() as usize; + let x_in = (x_ratio * w as f64).floor() as usize; + + unsafe { + let output = unsafe_shared_out.get(); + output[(b, c, h, w)] = x[(b, c, y_in, x_in)]; + } + }); + }); + + output.into_dyn().into_shared() +} + +pub(crate) fn nearest_interpolate_backward( + x: SharedArray, + grad: SharedArray, + output_size: [usize; 2], +) -> SharedArray { + let [batch_size, channels, input_height, input_width] = x.shape().dims(); + let [output_height, output_width] = output_size; + + let mut output_grad = + Array4::from_elem((batch_size, channels, input_height, input_width), 0.elem()); + let unsafe_shared_out = UnsafeSharedRef::new(&mut output_grad); + + run_par!(|| { + iter_range_par!(0, batch_size * channels).for_each(|k| unsafe { + let b = k / channels; + let c = k % channels; + + let output_grad = unsafe_shared_out.get(); + + for oh in 0..output_height { + for ow in 0..output_width { + let ih = start_index(oh, output_height, input_height); + let iw = start_index(ow, output_width, input_width); + + output_grad[[b, c, ih, iw]] += grad[[b, c, oh, ow]] + } + } + }) + }); + + output_grad.into_dyn().into_shared() +} + +fn start_index(output_size_index: usize, output_size: usize, input_size: usize) -> usize { + ((output_size_index as f32 * input_size as f32) / output_size as f32).floor() as usize +} + +// clamp ceil(frac) to stay within bounds in case of floating-point imprecision +pub(crate) fn ceil_clamp(frac: f64, max: usize) -> f64 { + frac.ceil().min(max as f64) +} + +pub(crate) fn bilinear_interpolate( + x: SharedArray, + output_size: [usize; 2], + align_corners: bool, +) -> SharedArray { + let x = x.into_dimensionality::().unwrap(); + + let (batch_size, channels, in_height, in_width) = x.dim(); + let [out_height, out_width] = output_size; + + let out_element_num = batch_size * channels * out_height * out_width; + let strides = ( + channels * out_height * out_width, + out_height * out_width, + out_width, + ); + + let mut output = Array4::zeros((batch_size, channels, out_height, out_width)); + let unsafe_shared_out = UnsafeSharedRef::new(&mut output); + + run_par!(|| { + iter_range_par!(0, out_element_num).for_each(|id| { + let (b, c, h, w) = ( + id / strides.0, + id % strides.0 / strides.1, + id % strides.1 / strides.2, + id % strides.2, + ); + + let (y_frac, x_frac) = if align_corners { + let y_ratio = ((in_height - 1) as f64) / (core::cmp::max(out_height - 1, 1) as f64); + let x_ratio = ((in_width - 1) as f64) / (core::cmp::max(out_width - 1, 1) as f64); + (y_ratio * h as f64, x_ratio * w as f64) + } else { + let y_frac = (h as f64 + 0.5) * (in_height as f64 / out_height as f64) - 0.5; + let x_frac = (w as f64 + 0.5) * (in_width as f64 / out_width as f64) - 0.5; + ( + y_frac.clamp(0.0, (in_height - 1) as f64), + x_frac.clamp(0.0, (in_width - 1) as f64), + ) + }; + let val = + bilinear_interpolate_single(&x, b, c, x_frac, y_frac, in_width - 1, in_height - 1); + + unsafe { + let output = unsafe_shared_out.get(); + output[(b, c, h, w)] = val.elem(); + } + }); + }); + + output.into_dyn().into_shared() +} + +pub(crate) fn bicubic_interpolate( + x: SharedArray, + output_size: [usize; 2], + align_corners: bool, +) -> SharedArray { + fn cubic_interp1d(x0: f64, x1: f64, x2: f64, x3: f64, t: f64) -> f64 { + fn cubic_convolution1(x: f64, a: f64) -> f64 { + ((a + 2.0) * x - (a + 3.0)) * x * x + 1.0 + } + + fn cubic_convolution2(x: f64, a: f64) -> f64 { + ((a * x - 5.0 * a) * x + 8.0 * a) * x - 4.0 * a + } + + let coeffs = [ + cubic_convolution2(t + 1.0, -0.75), + cubic_convolution1(t, -0.75), + cubic_convolution1(1.0 - t, -0.75), + cubic_convolution2(2.0 - t, -0.75), + ]; + + x0 * coeffs[0] + x1 * coeffs[1] + x2 * coeffs[2] + x3 * coeffs[3] + } + + let x = x.into_dimensionality::().unwrap(); + + let (batch_size, channels, in_height, in_width) = x.dim(); + let [out_height, out_width] = output_size; + + let out_element_num = batch_size * channels * out_height * out_width; + let strides = ( + channels * out_height * out_width, + out_height * out_width, + out_width, + ); + + let mut output = Array4::zeros((batch_size, channels, out_height, out_width)); + let unsafe_shared_out = UnsafeSharedRef::new(&mut output); + + run_par!(|| { + iter_range_par!(0, out_element_num).for_each(|id| { + let (b, c, h, w) = ( + id / strides.0, + id % strides.0 / strides.1, + id % strides.1 / strides.2, + id % strides.2, + ); + + let (y_frac, x_frac) = if align_corners { + let y_ratio = ((in_height - 1) as f64) / (core::cmp::max(out_height - 1, 1) as f64); + let x_ratio = ((in_width - 1) as f64) / (core::cmp::max(out_width - 1, 1) as f64); + (y_ratio * h as f64, x_ratio * w as f64) + } else { + let y_frac = (h as f64 + 0.5) * (in_height as f64 / out_height as f64) - 0.5; + let x_frac = (w as f64 + 0.5) * (in_width as f64 / out_width as f64) - 0.5; + (y_frac, x_frac) + }; + let y0 = y_frac.floor(); + let yw = y_frac - y0; + let y_in = y0 as isize; + + let x0 = x_frac.floor(); + let xw = x_frac - x0; + let x_in = x0 as isize; + + let max_h = (in_height - 1) as isize; + let max_w = (in_width - 1) as isize; + + let ys_in = [ + (y_in - 1).clamp(0, max_h) as usize, + y_in.clamp(0, max_h) as usize, + (y_in + 1).clamp(0, max_h) as usize, + (y_in + 2).clamp(0, max_h) as usize, + ]; + + let xs_in = [ + (x_in - 1).clamp(0, max_w) as usize, + x_in.clamp(0, max_w) as usize, + (x_in + 1).clamp(0, max_w) as usize, + (x_in + 2).clamp(0, max_w) as usize, + ]; + + let coefficients = ys_in.map(|y| { + cubic_interp1d( + x[(b, c, y, xs_in[0])].elem(), + x[(b, c, y, xs_in[1])].elem(), + x[(b, c, y, xs_in[2])].elem(), + x[(b, c, y, xs_in[3])].elem(), + xw, + ) + }); + + let result = cubic_interp1d( + coefficients[0], + coefficients[1], + coefficients[2], + coefficients[3], + yw, + ) + .elem(); + + unsafe { + let output = unsafe_shared_out.get(); + output[(b, c, h, w)] = result; + } + }); + }); + + output.into_dyn().into_shared() +} + +pub(crate) fn lanczos3_interpolate( + x: SharedArray, + output_size: [usize; 2], + align_corners: bool, +) -> SharedArray { + fn lanczos3_weight(x: f64) -> f64 { + if x == 0.0 { + return 1.0; + } + let abs_x = x.abs(); + if abs_x >= 3.0 { + return 0.0; + } + let pi = core::f64::consts::PI; + let pi_x = pi * x; + let pi_x_over_3 = pi_x / 3.0; + (pi_x.sin() * pi_x_over_3.sin()) / (pi_x * pi_x_over_3) + } + + let x = x.into_dimensionality::().unwrap(); + + let (batch_size, channels, in_height, in_width) = x.dim(); + let [out_height, out_width] = output_size; + + let out_element_num = batch_size * channels * out_height * out_width; + let strides = ( + channels * out_height * out_width, + out_height * out_width, + out_width, + ); + + let mut output = Array4::zeros((batch_size, channels, out_height, out_width)); + let unsafe_shared_out = UnsafeSharedRef::new(&mut output); + + run_par!(|| { + iter_range_par!(0, out_element_num).for_each(|id| { + let (b, c, h, w) = ( + id / strides.0, + id % strides.0 / strides.1, + id % strides.1 / strides.2, + id % strides.2, + ); + + let (y_frac, x_frac) = if align_corners { + let y_ratio = ((in_height - 1) as f64) / (core::cmp::max(out_height - 1, 1) as f64); + let x_ratio = ((in_width - 1) as f64) / (core::cmp::max(out_width - 1, 1) as f64); + (y_ratio * h as f64, x_ratio * w as f64) + } else { + let y_frac = (h as f64 + 0.5) * (in_height as f64 / out_height as f64) - 0.5; + let x_frac = (w as f64 + 0.5) * (in_width as f64 / out_width as f64) - 0.5; + (y_frac, x_frac) + }; + + let y0 = y_frac.floor(); + let x0 = x_frac.floor(); + let max_h = (in_height - 1) as isize; + let max_w = (in_width - 1) as isize; + + // 6x6 separable Lanczos3 filter (skip out-of-bounds positions) + let mut result = 0.0; + let mut weight_sum = 0.0; + for ky in -2..=3 { + let yi = y0 as isize + ky; + if yi < 0 || yi > max_h { + continue; + } + let y_idx = yi as usize; + let wy = lanczos3_weight(y_frac - (y0 + ky as f64)); + for kx in -2..=3 { + let xi = x0 as isize + kx; + if xi < 0 || xi > max_w { + continue; + } + let x_idx = xi as usize; + let wx = lanczos3_weight(x_frac - (x0 + kx as f64)); + let w = wy * wx; + let pixel: f64 = x[(b, c, y_idx, x_idx)].elem(); + result += pixel * w; + weight_sum += w; + } + } + if weight_sum != 0.0 { + result /= weight_sum; + } + + unsafe { + let output = unsafe_shared_out.get(); + output[(b, c, h, w)] = result.elem(); + } + }); + }); + + output.into_dyn().into_shared() +} + +/// Sample an element of the source array with bilinear interpolation +/// +/// * `source` - The tensor to read from. Has shape (batch_size, channels, height, width) +/// * `b` - The batch to read from +/// * `c` - The channel to read from +/// * `x` - The x position to read in the array +/// * `y` - The y position to read in the array +/// * `x_max` - The max x position (inclusive) +/// * `y_max` - The max y position (inclusive) +/// +/// # Returns +/// +/// The interpolated value read from the array +pub(crate) fn bilinear_interpolate_single( + source: &ArrayBase>, + b: usize, + c: usize, + x: f64, + y: f64, + x_max: usize, + y_max: usize, +) -> f64 +where + E: FloatNdArrayElement, + S: DataOwned, +{ + let y0 = y.floor(); + let y1 = ceil_clamp(y, y_max); + let yw = y - y0; + + let x0 = x.floor(); + let x1 = ceil_clamp(x, x_max); + let xw = x - x0; + + let (x0, x1, y0, y1) = (x0 as usize, x1 as usize, y0 as usize, y1 as usize); + + let p_a = source[(b, c, y0, x0)].elem::() * (1.0 - xw) * (1.0 - yw); + let p_b = source[(b, c, y0, x1)].elem::() * xw * (1.0 - yw); + let p_c = source[(b, c, y1, x0)].elem::() * (1.0 - xw) * yw; + let p_d = source[(b, c, y1, x1)].elem::() * xw * yw; + + p_a + p_b + p_c + p_d +} diff --git a/crates/burn-ndarray/src/ops/macros.rs b/crates/burn-ndarray/src/ops/macros.rs new file mode 100644 index 0000000..b3ac4f9 --- /dev/null +++ b/crates/burn-ndarray/src/ops/macros.rs @@ -0,0 +1,107 @@ +macro_rules! keepdim { + ( + $dim:expr, + $self:expr, + mean + ) => {{ + // Get shape first (via reference), then pass ownership to avoid clone + let mut shape = $self.shape().into_shape(); + shape[$dim] = 1; + let tensor: SharedArray = mean_dim($self, $dim); + NdArrayOps::reshape(tensor, shape) + }}; + ( + $dim:expr, + $self:expr, + sum + ) => {{ + // Get shape first (via reference), then pass ownership to avoid clone + let mut shape = $self.shape().into_shape(); + shape[$dim] = 1; + let tensor: SharedArray = sum_dim($self, $dim); + NdArrayOps::reshape(tensor, shape) + }}; + ( + $dim:expr, + $self:expr, + prod + ) => {{ + // Get shape first (via reference), then pass ownership to avoid clone + let mut shape = $self.shape().into_shape(); + shape[$dim] = 1; + let tensor: SharedArray = prod_dim($self, $dim); + NdArrayOps::reshape(tensor, shape) + }}; +} + +use burn_backend::ElementConversion; +pub(crate) use keepdim; +use ndarray::{Axis, Zip}; + +use crate::{SharedArray, element::NdArrayElement}; + +pub(crate) fn mean_dim(tensor: SharedArray, dim: usize) -> SharedArray { + tensor.mean_axis(Axis(dim)).unwrap().into_shared() +} + +pub(crate) fn sum_dim(tensor: SharedArray, dim: usize) -> SharedArray { + tensor.sum_axis(Axis(dim)).into_shared() +} + +pub(crate) fn prod_dim(tensor: SharedArray, dim: usize) -> SharedArray { + tensor + .fold_axis(Axis(dim), 1.elem::(), |acc, &x| acc.mul(x.elem())) + .into_shared() +} + +/// Generic cumulative operation function with closure-based operation. +pub(crate) fn cumulative_with_op(tensor: SharedArray, dim: usize, op: F) -> SharedArray +where + E: NdArrayElement, + F: Fn(&mut E, &E), +{ + let axis = Axis(dim); + let shape = tensor.shape().to_vec(); + // Use into_owned() instead of to_owned() - only copies if shared, avoids copy if unique + let mut result = tensor.into_owned(); + let dim_size = shape[dim]; + + for i in 1..dim_size { + let prev = result.index_axis(axis, i - 1).to_owned(); + let mut current = result.index_axis_mut(axis, i); + Zip::from(&mut current).and(&prev).for_each(&op); + } + + result.into_shared() +} + +// Define all cumulative operation functions using the generic function +pub(crate) fn cumsum_dim(tensor: SharedArray, dim: usize) -> SharedArray { + cumulative_with_op(tensor, dim, |c, &p| *c = c.add(p.elem())) +} + +pub(crate) fn cumprod_dim(tensor: SharedArray, dim: usize) -> SharedArray { + cumulative_with_op(tensor, dim, |c, &p| *c = c.mul(p.elem())) +} + +pub(crate) fn cummin_dim>( + tensor: SharedArray, + dim: usize, +) -> SharedArray { + cumulative_with_op(tensor, dim, |c, &p| { + if p < *c { + *c = p; + } + }) +} + +pub(crate) fn cummax_dim>( + tensor: SharedArray, + dim: usize, +) -> SharedArray { + cumulative_with_op(tensor, dim, |c, &p| { + if p > *c { + *c = p; + } + }) +} diff --git a/crates/burn-ndarray/src/ops/matmul.rs b/crates/burn-ndarray/src/ops/matmul.rs new file mode 100644 index 0000000..6b17a43 --- /dev/null +++ b/crates/burn-ndarray/src/ops/matmul.rs @@ -0,0 +1,368 @@ +use crate::UnsafeSharedRef; +use crate::{NdArrayElement, ShapeOps, SharedArray, iter_range_par, ops::NdArrayOps, run_par}; + +use alloc::{vec, vec::Vec}; +use burn_backend::ElementConversion; +use burn_backend::Shape; +use ndarray::{IxDyn, s}; + +pub(crate) fn matmul( + lhs: SharedArray, + rhs: SharedArray, +) -> SharedArray { + let shape_lhs = lhs.shape(); + let shape_rhs = rhs.shape(); + let ndims = shape_lhs.num_dims(); + let m = shape_lhs[ndims - 2]; // # of left rows + let k = shape_rhs[ndims - 2]; // # of left cols and right rows + let n = shape_rhs[ndims - 1]; // # of right cols + + let (out_shape, strides_lhs, strides_rhs, strides_out) = output_shape(shape_lhs, shape_rhs); + let l_mat_size = m * k; // size of matrix component of left array + let r_mat_size = k * n; // size of matrix component of right array + let out_mat_size = m * n; // size of matrix component of output array + + let num_l_batches = shape_lhs.num_elements() / l_mat_size; + let num_r_batches = shape_rhs.num_elements() / r_mat_size; + let num_out_batches = out_shape.num_elements() / out_mat_size; + + let lhs_array = NdArrayOps::reshape(lhs, Shape::new([num_l_batches, m, k])); + let rhs_array = NdArrayOps::reshape(rhs, Shape::new([num_r_batches, k, n])); + + let alpha: E = 1.0.elem(); + let beta: E = 0.0.elem(); + + let out = run_par!(|| { + let mut out_array = ndarray::Array3::::zeros((num_out_batches, m, n)); + let unsafe_shared_out_array = UnsafeSharedRef::new(&mut out_array); + + iter_range_par!(0, num_out_batches).for_each(|out_batch| { + // Here, we: + // 1. Un-flatten the output batch into a component-based batch index. + // 2. Use the strides for left and right batch indices to convert it to a flattened + // batch for left and right. + let out_index = strides_out.unflatten(out_batch); + let l_batch = strides_lhs.flatten(&out_index); + let r_batch = strides_rhs.flatten(&out_index); + + let lhs_slice = lhs_array.slice(s!(l_batch, .., ..)); + let rhs_slice = rhs_array.slice(s!(r_batch, .., ..)); + + unsafe { + let mut out_slice = unsafe_shared_out_array + .get() + .slice_mut(s!(out_batch, .., ..)); + + ndarray::linalg::general_mat_mul( + alpha, + &lhs_slice, + &rhs_slice, + beta, + &mut out_slice, + ) + } + }); + + out_array.into_shared().into_dyn() + }); + + NdArrayOps::reshape(out, out_shape) +} + +#[derive(Debug, PartialEq)] +struct Strides { + strides: Vec, +} +impl Strides { + fn new(strides: Vec) -> Self { + Strides { strides } + } + + fn unflatten(&self, linear_index: usize) -> Vec { + let mut coord = Vec::with_capacity(self.strides.len()); + let mut rem = linear_index; + for stride in self.strides.iter() { + coord.push(rem / stride); + rem %= stride; + } + coord + } + + fn flatten(&self, index: &Vec) -> usize { + assert_eq!(self.strides.len(), index.len()); + self.strides + .iter() + .zip(index) + .map(|(stride, index)| stride * index) + .sum() + } +} + +/// Compute the (broadcasted) output shape of matrix multiplication, along with strides for +/// the non-matrix dimensions of all arrays. +/// +/// # Arguments +/// * `lsh`: Shape of the first (left-hand) matrix multiplication argument. +/// * `rsh`: Shape of the second (right-hand) matrix multiplication argument. +/// +/// # Panics +/// * If `D` is not at least 2. +/// * If the matrix multiplication dimensions (last 2) are incompatible. +/// * If any other dimension is not the same for both tensors, or equal to 1. (Any dimension where +/// one dim is equal to 1 is broadcast.) +fn output_shape(lsh: &[usize], rsh: &[usize]) -> (Shape, Strides, Strides, Strides) { + let ndims = lsh.num_dims(); + if ndims < 2 { + panic!( + "Matrix multiplication requires an array with at least 2 dimensions. Got Rank {}", + ndims + ); + } + + // Fetch matrix dimensions and check compatibility. + let l_rows = lsh[ndims - 2]; + let l_cols = lsh[ndims - 1]; + let r_rows = rsh[ndims - 2]; + let r_cols = rsh[ndims - 1]; + if l_cols != r_rows { + panic!( + "Dimensions are incompatible for matrix multiplication: LHS columns ({}) != ({})", + l_cols, r_rows + ); + } + // Set matrix dimensions of the output shape. + let mut osh = vec![0; ndims]; + osh[ndims - 2] = l_rows; + osh[ndims - 1] = r_cols; + + // Set other array dimensions, broadcasting as necessary. + // Compute the strides inline. + let mut cur_l_stride: usize = 1; + let mut cur_r_stride: usize = 1; + let mut cur_o_stride: usize = 1; + let mut l_strides = Vec::with_capacity(ndims - 2); + let mut r_strides = Vec::with_capacity(ndims - 2); + let mut o_strides = Vec::with_capacity(ndims - 2); + for i in (0..ndims - 2).rev() { + let l_dim = lsh[i]; + let r_dim = rsh[i]; + + // Compatible dimensions are: + // 1. Both dimensions are equal. + // 2. One of the dimensions is equal to 1. + let o_dim: usize; + if l_dim == r_dim { + o_dim = l_dim; // both dimensions are equal + l_strides.push(cur_l_stride); + r_strides.push(cur_r_stride); + } else if l_dim == 1 { + o_dim = r_dim; // broadcast the left + l_strides.push(0); + r_strides.push(cur_r_stride); + } else if r_dim == 1 { + o_dim = l_dim; // broadcast the right + l_strides.push(cur_l_stride); + r_strides.push(0); + } else { + panic!("Dimensions differ and cannot be broadcasted."); + } + osh[i] = o_dim; + o_strides.push(cur_o_stride); + cur_o_stride *= o_dim; + + cur_l_stride *= l_dim; + cur_r_stride *= r_dim; + } + l_strides.reverse(); + r_strides.reverse(); + o_strides.reverse(); + + ( + Shape::from(osh), + Strides::new(l_strides), + Strides::new(r_strides), + Strides::new(o_strides), + ) +} + +pub(crate) fn cross( + lhs: SharedArray, + rhs: SharedArray, + dim: usize, +) -> SharedArray { + let shape_lhs = lhs.shape(); + let shape_rhs = rhs.shape(); + let ndims = shape_lhs.num_dims(); + + // Broadcast the shapes except along dim + let mut broadcast_shape = vec![0; ndims]; + for i in 0..ndims { + if i == dim { + broadcast_shape[i] = shape_lhs[i]; // already checked to be 3 + } else { + let l = shape_lhs[i]; + let r = shape_rhs[i]; + if l == r { + broadcast_shape[i] = l; + } else if l == 1 { + broadcast_shape[i] = r; + } else if r == 1 { + broadcast_shape[i] = l; + } else { + panic!("Tensors are not broadcastable along dimension {}", i); + } + } + } + + // Broadcast lhs and rhs + let lhs_broadcast = if shape_lhs == broadcast_shape.as_slice() { + lhs + } else { + NdArrayOps::expand(lhs, Shape::from(broadcast_shape.clone())) + }; + let rhs_broadcast = if shape_rhs == broadcast_shape.as_slice() { + rhs + } else { + NdArrayOps::expand(rhs, Shape::from(broadcast_shape.clone())) + }; + + // Now, move dim to the last dimension + let mut perm = (0..ndims).collect::>(); + perm.remove(dim); + perm.push(dim); + + let lhs_permuted = NdArrayOps::permute(lhs_broadcast, &perm); + let rhs_permuted = NdArrayOps::permute(rhs_broadcast, &perm); + + // Reshape to (*, 3) + let total_elements = lhs_permuted.shape().num_elements(); + let batch_size = total_elements / 3; + let lhs_reshaped = NdArrayOps::reshape(lhs_permuted, Shape::new([batch_size, 3])); + let rhs_reshaped = NdArrayOps::reshape(rhs_permuted, Shape::new([batch_size, 3])); + + // Compute cross product + let mut result = ndarray::ArrayD::::zeros(IxDyn(&[batch_size, 3])); + for i in 0..batch_size { + let a1 = lhs_reshaped[IxDyn(&[i, 0])]; + let a2 = lhs_reshaped[IxDyn(&[i, 1])]; + let a3 = lhs_reshaped[IxDyn(&[i, 2])]; + let b1 = rhs_reshaped[IxDyn(&[i, 0])]; + let b2 = rhs_reshaped[IxDyn(&[i, 1])]; + let b3 = rhs_reshaped[IxDyn(&[i, 2])]; + result[IxDyn(&[i, 0])] = a2.mul(b3).sub(a3.mul(b2)); + result[IxDyn(&[i, 1])] = a3.mul(b1).sub(a1.mul(b3)); + result[IxDyn(&[i, 2])] = a1.mul(b2).sub(a2.mul(b1)); + } + + let result_shared = result.into_shared(); + + // Reshape back to the broadcast shape with dim at the end + let mut result_shape = broadcast_shape; + result_shape.remove(dim); + result_shape.push(3); + let result_reshaped = NdArrayOps::reshape(result_shared, Shape::from(result_shape)); + + // Permute back + let mut inv_perm = vec![0; ndims]; + for (i, &p) in perm.iter().enumerate() { + inv_perm[p] = i; + } + NdArrayOps::permute(result_reshaped, &inv_perm) +} + +#[cfg(test)] +mod tests { + use super::*; + + impl Strides { + fn empty() -> Self { + Strides { + strides: Vec::with_capacity(0), + } + } + } + + #[test] + fn test_output_shape() { + // plain matrix multiply + assert_eq!( + output_shape(&[5, 3], &[3, 7]), + ( + Shape::from([5, 7]), + Strides::empty(), + Strides::empty(), + Strides::empty() + ) + ); + // matrix multiply with one extra stack dimension + assert_eq!( + output_shape(&[4, 5, 3], &[4, 3, 7]), + ( + Shape::from([4, 5, 7]), + Strides::new(vec![1]), + Strides::new(vec![1]), + Strides::new(vec![1]) + ) + ); + // rank 3, broadcast left + assert_eq!( + output_shape(&[1, 5, 3], &[4, 3, 7]), + ( + Shape::from([4, 5, 7]), + Strides::new(vec![0]), + Strides::new(vec![1]), + Strides::new(vec![1]) + ) + ); + // rank 3, broadcast right + assert_eq!( + output_shape(&[4, 5, 3], &[1, 3, 7]), + ( + Shape::from([4, 5, 7]), + Strides::new(vec![1]), + Strides::new(vec![0]), + Strides::new(vec![1]) + ) + ); + // rank 4, multi broadcast + assert_eq!( + output_shape(&[1, 4, 5, 3], &[8, 1, 3, 7]), + ( + Shape::from([8, 4, 5, 7]), + Strides::new(vec![0, 1]), + Strides::new(vec![1, 0]), + Strides::new(vec![4, 1]) + ) + ); + // rank 5, multi-broadcast + assert_eq!( + output_shape(&[1, 3, 4, 5, 3], &[8, 3, 1, 3, 7]), + ( + Shape::from([8, 3, 4, 5, 7]), + Strides::new(vec![0, 4, 1]), + Strides::new(vec![3, 1, 0]), + Strides::new(vec![12, 4, 1]) + ) + ) + } + + #[test] + #[should_panic( + expected = "Matrix multiplication requires an array with at least 2 dimensions." + )] + fn test_output_shape_too_small() { + output_shape(&[4], &[4]); + } + + #[test] + #[should_panic(expected = "Dimensions are incompatible for matrix multiplication: LHS columns")] + fn test_output_shape_bad_matrix_dims() { + output_shape(&[5, 3], &[4, 7]); + } + + #[test] + #[should_panic(expected = "Dimensions differ and cannot be broadcasted.")] + fn test_output_shape_non_broadcast() { + output_shape(&[4, 5, 3], &[2, 3, 7]); + } +} diff --git a/crates/burn-ndarray/src/ops/maxpool.rs b/crates/burn-ndarray/src/ops/maxpool.rs new file mode 100644 index 0000000..2a162cf --- /dev/null +++ b/crates/burn-ndarray/src/ops/maxpool.rs @@ -0,0 +1,247 @@ +use crate::{ + ShapeOps, SharedArray, + element::{FloatNdArrayElement, IntNdArrayElement}, + iter_range_par, + ops::padding::apply_padding_4d, + run_par, + sharing::UnsafeSharedRef, +}; + +use burn_backend::ElementConversion; +use burn_backend::ops::conv::calculate_pool_output_size; +use ndarray::Array4; + +pub(crate) fn max_pool2d( + x: SharedArray, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, +) -> SharedArray { + let [kernel_height, kernel_width] = kernel_size; + let [padding_height, padding_width] = padding; + let [stride_height, stride_width] = stride; + let [dilation_height, dilation_width] = dilation; + let [batch_size, channels, x_height, x_width] = x.shape().dims(); + let inf = (-f32::INFINITY).elem::(); + + let out_height = calculate_pool_output_size( + kernel_height, + stride_height, + padding_height, + dilation_height, + x_height, + ceil_mode, + ); + let out_width = calculate_pool_output_size( + kernel_width, + stride_width, + padding_width, + dilation_width, + x_width, + ceil_mode, + ); + + // Calculate extra padding needed for ceil_mode + // The maximum input position accessed is: (out_size - 1) * stride + (kernel_size - 1) * dilation + // This must be < input_size + 2 * total_padding + let max_ih = + (out_height.saturating_sub(1)) * stride_height + (kernel_height - 1) * dilation_height; + let max_iw = (out_width.saturating_sub(1)) * stride_width + (kernel_width - 1) * dilation_width; + let padded_height = x_height + 2 * padding_height; + let padded_width = x_width + 2 * padding_width; + let extra_pad_h = max_ih.saturating_sub(padded_height.saturating_sub(1)); + let extra_pad_w = max_iw.saturating_sub(padded_width.saturating_sub(1)); + let total_padding = [padding_height + extra_pad_h, padding_width + extra_pad_w]; + + let x = apply_padding_4d::(x, total_padding, inf); + + // Offset to account for extra padding (extra_pad is added on both sides by apply_padding_4d) + let offset_h = extra_pad_h; + let offset_w = extra_pad_w; + + let mut output = Array4::from_elem((batch_size, channels, out_height, out_width), inf); + let unsafe_shared_out = UnsafeSharedRef::new(&mut output); + + run_par!(|| { + iter_range_par!(0, batch_size * channels).for_each(|k| unsafe { + let b = k / channels; + let c = k % channels; + + let output = unsafe_shared_out.get(); + + for oh in 0..out_height { + for ow in 0..out_width { + let mut max_val = inf; + + for kh in 0..kernel_height { + let ih = offset_h + oh * stride_height + kh * dilation_height; + + for kw in 0..kernel_width { + let iw = offset_w + ow * stride_width + kw * dilation_width; + + let val = x[[b, c, ih, iw]]; + + if val > max_val { + max_val = val; + } + } + } + + output[[b, c, oh, ow]] = max_val; + } + } + }) + }); + + output.into_dyn().into_shared() +} + +pub(crate) fn max_pool2d_with_indices( + x: SharedArray, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, +) -> (SharedArray, SharedArray) { + let [kernel_height, kernel_width] = kernel_size; + let [padding_height, padding_width] = padding; + let [stride_height, stride_width] = stride; + let [dilation_height, dilation_width] = dilation; + let [batch_size, channels, x_height, x_width] = x.shape().dims(); + let inf = (-f32::INFINITY).elem::(); + + let out_height = calculate_pool_output_size( + kernel_height, + stride_height, + padding_height, + dilation_height, + x_height, + ceil_mode, + ); + let out_width = calculate_pool_output_size( + kernel_width, + stride_width, + padding_width, + dilation_width, + x_width, + ceil_mode, + ); + + // Calculate extra padding needed for ceil_mode + let max_ih = + (out_height.saturating_sub(1)) * stride_height + (kernel_height - 1) * dilation_height; + let max_iw = (out_width.saturating_sub(1)) * stride_width + (kernel_width - 1) * dilation_width; + let padded_height = x_height + 2 * padding_height; + let padded_width = x_width + 2 * padding_width; + let extra_pad_h = max_ih.saturating_sub(padded_height.saturating_sub(1)); + let extra_pad_w = max_iw.saturating_sub(padded_width.saturating_sub(1)); + let total_padding = [padding_height + extra_pad_h, padding_width + extra_pad_w]; + + let x = apply_padding_4d::(x, total_padding, inf); + + // Offset to account for extra padding + let offset_h = extra_pad_h; + let offset_w = extra_pad_w; + + let mut output = Array4::from_elem((batch_size, channels, out_height, out_width), inf); + let mut indices = Array4::::zeros((batch_size, channels, out_height, out_width)); + + let unsafe_shared_out = UnsafeSharedRef::new(&mut output); + let unsafe_shared_indices = UnsafeSharedRef::new(&mut indices); + + run_par!(|| { + iter_range_par!(0, batch_size * channels).for_each(|k| unsafe { + let b = k / channels; + let c = k % channels; + + let output = unsafe_shared_out.get(); + let indices = unsafe_shared_indices.get(); + + for oh in 0..out_height { + for ow in 0..out_width { + let mut max_val = inf; + let mut index = 0; + + for kh in 0..kernel_height { + let ih = offset_h + oh * stride_height + kh * dilation_height; + + for kw in 0..kernel_width { + let iw = offset_w + ow * stride_width + kw * dilation_width; + let val = x[[b, c, ih, iw]]; + + if val > max_val { + max_val = val; + + // Calculate index in original (unpadded) input + let ih_orig = ih as i64 - (total_padding[0]) as i64; + let iw_orig = iw as i64 - (total_padding[1]) as i64; + + // Clamp to valid range for index calculation + let ih_clamped = ih_orig.max(0).min(x_height as i64 - 1); + let iw_clamped = iw_orig.max(0).min(x_width as i64 - 1); + + index = ih_clamped * x_width as i64 + iw_clamped; + } + } + } + + output[[b, c, oh, ow]] = max_val; + indices[[b, c, oh, ow]] = index.elem(); + } + } + }) + }); + + let output = output.into_dyn().into_shared(); + let indices = indices.into_dyn().into_shared(); + + (output, indices) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn max_pool2d_backward( + x: SharedArray, + _kernel_size: [usize; 2], + _stride: [usize; 2], + _padding: [usize; 2], + _dilation: [usize; 2], + _ceil_mode: bool, + output_grad: SharedArray, + indices: SharedArray, +) -> SharedArray { + let [_batch_size, _channels, height, width] = output_grad.shape().dims(); + let [batch_size, channels, height_x, width_x] = x.shape().dims(); + + let output_grad = output_grad; + let indices = indices; + + let mut output = Array4::zeros((batch_size, channels, height_x, width_x)); + + let unsafe_shared_out = UnsafeSharedRef::new(&mut output); + + run_par!(|| { + iter_range_par!(0, batch_size * channels).for_each(|k| unsafe { + let b = k / channels; + let c = k % channels; + + let output = unsafe_shared_out.get(); + + for h in 0..height { + for w in 0..width { + let index = indices[[b, c, h, w]].elem::(); + let grad = output_grad[[b, c, h, w]]; + + let index_h = index as usize / width_x; + let index_w = index as usize % width_x; + + output[[b, c, index_h, index_w]] += grad; + } + } + }); + }); + + output.into_dyn().into_shared() +} diff --git a/crates/burn-ndarray/src/ops/mod.rs b/crates/burn-ndarray/src/ops/mod.rs new file mode 100644 index 0000000..f4f215e --- /dev/null +++ b/crates/burn-ndarray/src/ops/mod.rs @@ -0,0 +1,24 @@ +mod activation; +mod base; +mod bool_tensor; +mod int_tensor; +mod module; +mod qtensor; +#[cfg(feature = "simd")] +mod simd; +mod tensor; +mod transaction; + +pub(crate) mod adaptive_avgpool; +pub(crate) mod avgpool; +pub(crate) mod conv; +pub(crate) mod deform_conv; +pub(crate) mod grid_sample; +pub(crate) mod interpolate; +pub(crate) mod macros; +pub(crate) mod matmul; +pub(crate) mod maxpool; +pub(crate) mod padding; +pub(crate) mod quantization; + +pub(crate) use base::*; diff --git a/crates/burn-ndarray/src/ops/module.rs b/crates/burn-ndarray/src/ops/module.rs new file mode 100644 index 0000000..59f32ca --- /dev/null +++ b/crates/burn-ndarray/src/ops/module.rs @@ -0,0 +1,396 @@ +use super::{ + adaptive_avgpool::{adaptive_avg_pool2d, adaptive_avg_pool2d_backward}, + avgpool::{avg_pool2d, avg_pool2d_backward}, + conv::{conv_transpose2d, conv_transpose3d, conv2d, conv3d}, + deform_conv::{backward::deform_conv2d_backward, deform_conv2d}, + interpolate::{ + bicubic_interpolate, bilinear_interpolate, lanczos3_interpolate, nearest_interpolate, + }, + maxpool::{max_pool2d, max_pool2d_backward, max_pool2d_with_indices}, +}; +use crate::ops::interpolate::nearest_interpolate_backward; +#[cfg(feature = "simd")] +use crate::ops::simd::{ + avgpool::try_avg_pool2d_simd, conv::try_conv2d_simd, maxpool::try_max_pool2d_simd, +}; +use crate::{ + NdArray, SharedArray, execute_with_int_dtype, execute_with_int_out_dtype, tensor::NdArrayTensor, +}; +use burn_backend::{ + TensorMetadata, + ops::{attention::attention_fallback, *}, + tensor::FloatTensor, +}; +use burn_std::IntDType; + +macro_rules! module_op { + // Module op with inputs (inp), optional (opt) and arguments (args). + // Converts NdArrayStorage to SharedArray for compatibility with existing operations. + (inp($($x:tt),+), opt($($opt:tt),*), $element:ident, $op:expr) => {{ + #[allow(unused_parens, unreachable_patterns)] + match ($($x),+) { + ($(NdArrayTensor::F32($x)),+) => { + type $element = f32; + $op( + $($x.into_shared()),+ + $(, $opt.map(|o| match o { NdArrayTensor::F32(val) => val.into_shared(), _ => panic!("Optional argument type mismatch") }))* + ) + } + ($(NdArrayTensor::F64($x)),+) => { + type $element = f64; + $op( + $($x.into_shared()),+ + $(, $opt.map(|o| match o { NdArrayTensor::F64(val) => val.into_shared(), _ => panic!("Optional argument type mismatch") }))* + ) + } + _ => panic!("Data type mismatch"), + } + }}; +} + +impl ModuleOps for NdArray { + fn conv2d( + x: NdArrayTensor, + weight: NdArrayTensor, + bias: Option, + options: ConvOptions<2>, + ) -> NdArrayTensor { + module_op!(inp(x, weight), opt(bias), E, |x, weight, bias| { + #[cfg(feature = "simd")] + let (x, weight, bias) = match try_conv2d_simd(x, weight, bias, options.clone()) { + Ok(out) => return out.into(), + Err(args) => args, + }; + conv2d::(x, weight, bias, options).into() + }) + } + + fn deform_conv2d( + x: FloatTensor, + offset: FloatTensor, + weight: FloatTensor, + mask: Option>, + bias: Option>, + options: DeformConvOptions<2>, + ) -> FloatTensor { + module_op!( + inp(x, offset, weight), + opt(mask, bias), + E, + |x, offset, weight, mask, bias| deform_conv2d::( + x, offset, weight, mask, bias, options + ) + .into() + ) + } + + fn deform_conv2d_backward( + x: FloatTensor, + offset: FloatTensor, + weight: FloatTensor, + mask: Option>, + bias: Option>, + output_grad: FloatTensor, + options: DeformConvOptions<2>, + ) -> DeformConv2dBackward { + module_op!( + inp(x, offset, weight, output_grad), + opt(mask, bias), + E, + |x, offset, weight, output_grad, mask, bias| { + let (x, offset, weight, mask, bias) = deform_conv2d_backward::( + x, + offset, + weight, + mask, + bias, + output_grad, + options, + ); + DeformConv2dBackward::new( + x.into(), + offset.into(), + weight.into(), + mask.map(|m| m.into()), + bias.map(|b| b.into()), + ) + } + ) + } + + fn conv_transpose2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<2>, + ) -> FloatTensor { + module_op!(inp(x, weight), opt(bias), E, |x, weight, bias| { + conv_transpose2d::(x, weight, bias, options).into() + }) + } + + fn avg_pool2d( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + module_op!(inp(x), opt(), E, |x| { + #[cfg(feature = "simd")] + let x = match if ceil_mode { + // SIMD path doesn't support ceil_mode yet, skip it + Err(x) + } else { + try_avg_pool2d_simd(x, kernel_size, stride, padding, count_include_pad) + } { + Ok(out) => return out.into(), + Err(x) => x, + }; + avg_pool2d::( + x, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ) + .into() + }) + } + + fn avg_pool2d_backward( + x: FloatTensor, + grad: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + module_op!(inp(x, grad), opt(), E, |x, grad| avg_pool2d_backward::( + x, + grad, + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode + ) + .into()) + } + + fn max_pool2d( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + ) -> FloatTensor { + module_op!(inp(x), opt(), E, |x| { + #[cfg(feature = "simd")] + let x = match if ceil_mode { + // SIMD path doesn't support ceil_mode yet, skip it + Err(x) + } else { + try_max_pool2d_simd(x, kernel_size, stride, padding, dilation) + } { + Ok(out) => return out.into(), + Err(x) => x, + }; + max_pool2d::(x, kernel_size, stride, padding, dilation, ceil_mode).into() + }) + } + + fn max_pool2d_with_indices( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + indices_dtype: IntDType, + ) -> MaxPool2dWithIndices { + execute_with_int_out_dtype!(indices_dtype, I, { + module_op!(inp(x), opt(), E, |x| { + let (output, indices) = max_pool2d_with_indices::( + x, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + ); + MaxPool2dWithIndices::new(output.into(), indices.into()) + }) + }) + } + + fn max_pool2d_with_indices_backward( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + output_grad: FloatTensor, + indices: NdArrayTensor, + ) -> MaxPool2dBackward { + execute_with_int_dtype!(indices, IntElem, |idx_s: SharedArray| { + module_op!(inp(x, output_grad), opt(), E, |x, output_grad| { + let output = max_pool2d_backward::( + x, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + output_grad, + idx_s, + ); + MaxPool2dBackward::new(output.into()) + }) + }) + } + + fn adaptive_avg_pool2d(x: FloatTensor, output_size: [usize; 2]) -> FloatTensor { + module_op!(inp(x), opt(), E, |x| adaptive_avg_pool2d::( + x, + output_size + ) + .into()) + } + + fn adaptive_avg_pool2d_backward( + x: FloatTensor, + grad: FloatTensor, + ) -> FloatTensor { + module_op!(inp(x, grad), opt(), E, |x, grad| { + adaptive_avg_pool2d_backward::(x, grad).into() + }) + } + + fn interpolate( + x: FloatTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> FloatTensor { + match options.mode { + InterpolateMode::Nearest => { + module_op!(inp(x), opt(), E, |x| nearest_interpolate::( + x, + output_size + ) + .into()) + } + InterpolateMode::NearestExact => { + panic!("nearest exact interpolation is not supported for ndarray backend") + } + InterpolateMode::Bilinear => { + let align_corners = options.align_corners; + module_op!(inp(x), opt(), E, |x| bilinear_interpolate::( + x, + output_size, + align_corners + ) + .into()) + } + InterpolateMode::Bicubic => { + let align_corners = options.align_corners; + module_op!(inp(x), opt(), E, |x| bicubic_interpolate::( + x, + output_size, + align_corners + ) + .into()) + } + InterpolateMode::Lanczos3 => { + let align_corners = options.align_corners; + module_op!(inp(x), opt(), E, |x| lanczos3_interpolate::( + x, + output_size, + align_corners + ) + .into()) + } + } + } + + fn interpolate_backward( + x: FloatTensor, + grad: FloatTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> FloatTensor { + match options.mode { + InterpolateMode::Nearest => module_op!(inp(x, grad), opt(), E, |x, grad| { + nearest_interpolate_backward::(x, grad, output_size).into() + }), + InterpolateMode::NearestExact => { + panic!("nearest exact interpolation backward is not supported for ndarray backend") + } + InterpolateMode::Bilinear => { + panic!("bilinear interpolation backward is not supported for ndarray backend") + } + InterpolateMode::Bicubic => { + panic!("bicubic interpolation backward is not supported for ndarray backend") + } + InterpolateMode::Lanczos3 => { + panic!("lanczos3 interpolation backward is not supported for ndarray backend") + } + } + } + + fn conv3d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<3>, + ) -> FloatTensor { + module_op!(inp(x, weight), opt(bias), E, |x, weight, bias| conv3d::( + x, weight, bias, options + ) + .into()) + } + + fn conv_transpose3d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<3>, + ) -> FloatTensor { + module_op!(inp(x, weight), opt(bias), E, |x, weight, bias| { + conv_transpose3d::(x, weight, bias, options).into() + }) + } + + fn attention( + query: FloatTensor, + key: FloatTensor, + value: FloatTensor, + mask: Option>, + attn_bias: Option>, + options: AttentionModuleOptions, + ) -> FloatTensor { + attention_fallback::(query, key, value, mask, attn_bias, options) + } + + fn rfft( + _signal: FloatTensor, + _dim: usize, + _n: Option, + ) -> (FloatTensor, FloatTensor) { + todo!("rfft is not supported for ndarray") + } + + fn irfft( + _spectrum_re: FloatTensor, + _spectrum_im: FloatTensor, + _dim: usize, + _n: Option, + ) -> FloatTensor { + todo!("irfft is not supported for ndarray") + } +} diff --git a/crates/burn-ndarray/src/ops/padding.rs b/crates/burn-ndarray/src/ops/padding.rs new file mode 100644 index 0000000..d9c6fd3 --- /dev/null +++ b/crates/burn-ndarray/src/ops/padding.rs @@ -0,0 +1,72 @@ +use crate::{NdArrayElement, SharedArray}; +use ndarray::{Array4, Array5}; + +use super::NdArrayOps; + +pub(crate) fn apply_padding_4d( + x: SharedArray, + padding: [usize; 2], + elem: E, +) -> SharedArray { + let [batch_size, input_channels, height, width] = x.shape().try_into().unwrap(); + let [padding_height, padding_width] = padding; + let padded_height = height + 2 * padding_height; + let padded_width = width + 2 * padding_width; + + let x_new = Array4::from_elem( + (batch_size, input_channels, padded_height, padded_width), + elem, + ); + let mut x_new = x_new.into_shared().into_dyn(); + + x_new = NdArrayOps::slice_assign( + x_new, + &[ + burn_backend::Slice::from(0..batch_size), + burn_backend::Slice::from(0..input_channels), + burn_backend::Slice::from(padding_height..height + padding_height), + burn_backend::Slice::from(padding_width..width + padding_width), + ], + x, + ); + + x_new +} + +pub(crate) fn apply_padding_5d( + x: SharedArray, + padding: [usize; 3], + elem: E, +) -> SharedArray { + let [batch_size, input_channels, depth, height, width] = x.shape().try_into().unwrap(); + let [padding_depth, padding_height, padding_width] = padding; + let padded_depth = depth + 2 * padding_depth; + let padded_height = height + 2 * padding_height; + let padded_width = width + 2 * padding_width; + + let x_new = Array5::from_elem( + ( + batch_size, + input_channels, + padded_depth, + padded_height, + padded_width, + ), + elem, + ); + let mut x_new = x_new.into_shared().into_dyn(); + + x_new = NdArrayOps::slice_assign( + x_new, + &[ + burn_backend::Slice::from(0..batch_size), + burn_backend::Slice::from(0..input_channels), + burn_backend::Slice::from(padding_depth..depth + padding_depth), + burn_backend::Slice::from(padding_height..height + padding_height), + burn_backend::Slice::from(padding_width..width + padding_width), + ], + x, + ); + + x_new +} diff --git a/crates/burn-ndarray/src/ops/qtensor.rs b/crates/burn-ndarray/src/ops/qtensor.rs new file mode 100644 index 0000000..543458b --- /dev/null +++ b/crates/burn-ndarray/src/ops/qtensor.rs @@ -0,0 +1,343 @@ +use alloc::{vec, vec::Vec}; + +use burn_backend::{ + DType, ExecutionError, Shape, TensorData, TensorMetadata, + ops::QTensorOps, + quantization::{ + QParams, QuantLevel, QuantMode, QuantScheme, QuantStore, QuantValue, + QuantizationParametersPrimitive, QuantizedBytes, + }, + tensor::{FloatTensor, IntTensor, QuantizedTensor}, +}; +use burn_std::{FloatDType, IntDType}; + +use crate::{ + NdArray, NdArrayDevice, NdArrayQTensor, NdArrayTensor, SharedArray, element::QuantElement, + execute_with_dtype, execute_with_int_dtype, execute_with_int_out_dtype, + execute_with_numeric_dtype, slice, +}; + +use super::quantization::{QuantizationStrategy, SymmetricQuantization}; +use super::{NdArrayMathOps, NdArrayOps}; + +impl QTensorOps for NdArray { + fn q_from_data(data: TensorData, _device: &NdArrayDevice) -> QuantizedTensor { + match data.dtype { + DType::QFloat(scheme) => { + let shape = data.shape.clone(); + let num_elements = data.num_elements(); + let q_bytes = QuantizedBytes { + bytes: data.into_bytes(), + scheme, + num_elements, + }; + + match scheme { + QuantScheme { + level: QuantLevel::Tensor | QuantLevel::Block(_), + mode: QuantMode::Symmetric, + value: QuantValue::Q8F | QuantValue::Q8S, + .. + } => { + // We can load QuantStore::U32 w/ QuantizedBytes impl + let (values, qparams) = q_bytes.into_vec_i8(); + let data = TensorData::new(values, shape); + // Overwrite storage + let scheme = scheme.with_store(QuantStore::Native); + + let qparams = qparams + .scales + .into_iter() + .map(|scales| QParams { scales }) + .collect(); + + NdArrayQTensor { + qtensor: NdArrayTensor::from_data(data), + scheme, + qparams, + } + } + QuantScheme { + value: + QuantValue::Q4F + | QuantValue::Q4S + | QuantValue::Q2F + | QuantValue::Q2S + | QuantValue::E2M1 + | QuantValue::E4M3 + | QuantValue::E5M2, + .. + } => unimplemented!("from_data not supported for scheme {scheme:?}"), + } + } + _ => panic!( + "Invalid dtype (expected DType::QFloat, got {:?})", + data.dtype + ), + } + } + + fn quantize( + tensor: FloatTensor, + scheme: &QuantScheme, + qparams: QuantizationParametersPrimitive, + ) -> QuantizedTensor { + let shape = tensor.shape(); + let data_f = tensor.into_data(); + let scales = qparams.scales.into_data().convert::(); + + // Implement with ndarray instead of QuantizationStrategy? + let (data, qparams) = match scheme { + QuantScheme { + level: QuantLevel::Tensor, + mode: QuantMode::Symmetric, + #[cfg(not(feature = "export_tests"))] + value: QuantValue::Q8F | QuantValue::Q8S, + // For tests, "native" sub-byte quant serves as a reference for value equality. + // Values are stored as i8 regardless. + #[cfg(feature = "export_tests")] + value: + QuantValue::Q8F + | QuantValue::Q8S + | QuantValue::Q4F + | QuantValue::Q4S + | QuantValue::Q2F + | QuantValue::Q2S, + store: QuantStore::Native, + .. + } => { + let scales = scales.iter().next().unwrap(); + let strategy = QuantizationStrategy::PerTensorSymmetric( + SymmetricQuantization::init(scales, scheme.value), + ); + let values = strategy.quantize(data_f.as_slice().unwrap()); + ( + TensorData::quantized(values, shape.clone(), *scheme, &[scales]), + vec![QParams { scales }], + ) + } + QuantScheme { + level: QuantLevel::Block(block_size), + mode: QuantMode::Symmetric, + #[cfg(not(feature = "export_tests"))] + value: QuantValue::Q8F | QuantValue::Q8S, + #[cfg(feature = "export_tests")] + value: + QuantValue::Q8F + | QuantValue::Q8S + | QuantValue::Q4F + | QuantValue::Q4S + | QuantValue::Q2F + | QuantValue::Q2S, + store: QuantStore::Native, + .. + } => { + let scales = scales.as_slice().unwrap(); + let (strategy, qparams) = scales + .iter() + .map(|&s| { + ( + SymmetricQuantization::init(s, scheme.value), + QParams { scales: s }, + ) + }) + .unzip(); + let strategy = QuantizationStrategy::PerBlockSymmetric(strategy, *block_size); + let values = strategy.quantize(data_f.as_slice().unwrap()); + ( + TensorData::quantized(values, shape.clone(), *scheme, scales), + qparams, + ) + } + scheme => unimplemented!("Quantization not supported for scheme {scheme:?}"), + }; + + let num_elements = data.num_elements(); + let q_bytes = QuantizedBytes { + bytes: data.into_bytes(), + scheme: *scheme, + num_elements, + }; + let (values, _) = q_bytes.into_vec_i8(); + let data = TensorData::new(values, shape); + + NdArrayQTensor { + qtensor: NdArrayTensor::from_data(data), + scheme: *scheme, + qparams, + } + } + + fn dequantize(tensor: QuantizedTensor, dtype: FloatDType) -> FloatTensor { + let strategy = tensor.strategy(); + let scheme = tensor.scheme; + let shape = tensor.shape(); + let data = match tensor.qtensor { + NdArrayTensor::I8(storage) => { + let data = storage.into_shared().into_iter().collect(); + dequantize(data, shape, scheme, &strategy, dtype.into()) + } + _ => unreachable!(), + }; + NdArrayTensor::from_data(data) + } + + fn q_to_device( + tensor: QuantizedTensor, + _device: &NdArrayDevice, + ) -> QuantizedTensor { + tensor + } + + fn q_reshape(tensor: QuantizedTensor, shape: Shape) -> QuantizedTensor { + NdArrayQTensor { + qtensor: execute_with_dtype!(tensor.qtensor, E, |array: SharedArray| { + NdArrayOps::reshape(array, shape) + }), + scheme: tensor.scheme, + qparams: tensor.qparams, + } + } + + async fn q_into_data(tensor: QuantizedTensor) -> Result { + let shape = tensor.qtensor.shape(); + let scales = tensor.qparams.iter().map(|q| q.scales).collect::>(); + Ok(execute_with_numeric_dtype!( + tensor.qtensor, + E, + |array: SharedArray| { + let values = array.into_iter().collect(); + TensorData::quantized(values, shape, tensor.scheme, &scales) + } + )) + } + + fn q_swap_dims( + tensor: QuantizedTensor, + dim1: usize, + dim2: usize, + ) -> QuantizedTensor { + NdArrayQTensor { + qtensor: execute_with_dtype!(tensor.qtensor, E, |array: SharedArray| { + NdArrayOps::swap_dims(array, dim1, dim2) + }), + scheme: tensor.scheme, + qparams: tensor.qparams, + } + } + + fn q_permute(tensor: QuantizedTensor, axes: &[usize]) -> QuantizedTensor { + NdArrayQTensor { + qtensor: execute_with_dtype!(tensor.qtensor, E, |array: SharedArray| { + NdArrayOps::permute(array, axes) + }), + scheme: tensor.scheme, + qparams: tensor.qparams, + } + } + + fn q_flip(tensor: QuantizedTensor, axes: &[usize]) -> QuantizedTensor { + NdArrayQTensor { + qtensor: execute_with_dtype!(tensor.qtensor, E, |array: SharedArray| { + NdArrayOps::flip(array, axes) + }), + scheme: tensor.scheme, + qparams: tensor.qparams, + } + } + + fn q_gather( + dim: usize, + tensor: QuantizedTensor, + indices: IntTensor, + ) -> QuantizedTensor { + let qtensor = execute_with_int_dtype!(indices, IntElem, |idx_array: SharedArray< + IntElem, + >| + -> NdArrayTensor { + execute_with_numeric_dtype!(tensor.qtensor, E, |array: SharedArray| { + NdArrayOps::gather(dim, array, idx_array) + }) + }); + NdArrayQTensor { + qtensor, + scheme: tensor.scheme, + qparams: tensor.qparams, + } + } + + fn q_select( + tensor: QuantizedTensor, + dim: usize, + indices: IntTensor, + ) -> QuantizedTensor { + let qtensor = execute_with_int_dtype!(indices, IntElem, |idx_array: SharedArray< + IntElem, + >| + -> NdArrayTensor { + execute_with_numeric_dtype!(tensor.qtensor, E, |array: SharedArray| { + NdArrayMathOps::select(array, dim, idx_array) + }) + }); + NdArrayQTensor { + qtensor, + scheme: tensor.scheme, + qparams: tensor.qparams, + } + } + + fn q_slice( + tensor: QuantizedTensor, + slices: &[burn_backend::Slice], + ) -> QuantizedTensor { + NdArrayQTensor { + qtensor: slice!(tensor.qtensor, slices), + scheme: tensor.scheme, + qparams: tensor.qparams, + } + } + + fn q_argmax(tensor: QuantizedTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + execute_with_int_out_dtype!(out_dtype, I, { + execute_with_numeric_dtype!(tensor.qtensor, E, |array: SharedArray| { + NdArrayMathOps::argmax::(array, dim) + }) + }) + } + + fn q_argmin(tensor: QuantizedTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + execute_with_int_out_dtype!(out_dtype, I, { + execute_with_numeric_dtype!(tensor.qtensor, E, |array: SharedArray| { + NdArrayMathOps::argmin::(array, dim) + }) + }) + } + + fn q_expand(tensor: QuantizedTensor, shape: Shape) -> QuantizedTensor { + NdArrayQTensor { + qtensor: execute_with_dtype!(tensor.qtensor, E, |array: SharedArray| { + NdArrayOps::expand(array, shape) + }), + scheme: tensor.scheme, + qparams: tensor.qparams, + } + } +} + +fn dequantize( + data: Vec, + shape: Shape, + scheme: QuantScheme, + strategy: &QuantizationStrategy, + dtype: DType, +) -> TensorData { + let qparams = match strategy { + QuantizationStrategy::PerTensorSymmetric(quant) => vec![quant.scale], + QuantizationStrategy::PerBlockSymmetric(quant, _block_size) => { + quant.iter().map(|q| q.scale).collect() + } + }; + let q_bytes = QuantizedBytes::new(data, scheme, &qparams); + let (values, _qparams) = q_bytes.into_vec_i8(); + TensorData::new(strategy.dequantize(&values), shape).convert_dtype(dtype) +} diff --git a/crates/burn-ndarray/src/ops/quantization.rs b/crates/burn-ndarray/src/ops/quantization.rs new file mode 100644 index 0000000..adaf1b1 --- /dev/null +++ b/crates/burn-ndarray/src/ops/quantization.rs @@ -0,0 +1,218 @@ +use alloc::vec::Vec; +use num_traits::{Float, PrimInt}; + +use burn_backend::quantization::{BlockSize, QuantValue}; + +// NOTE: this mainly serves as a simple reference implementation. +// The de/quantization ops should be refactored to use ndarray. + +/// Quantization strategy. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum QuantizationStrategy { + /// Per-tensor symmetric quantization. + PerTensorSymmetric(SymmetricQuantization), + /// Per-block symmetric quantization. + PerBlockSymmetric(Vec>, BlockSize), +} + +impl QuantizationStrategy { + /// Quantize the values to a lower precision data type. + pub fn quantize(&self, values: &[f32]) -> Vec { + match self { + QuantizationStrategy::PerTensorSymmetric(strategy) => strategy.quantize(values), + QuantizationStrategy::PerBlockSymmetric(strategy, block_size) => { + let block_elems = block_size.num_elements(); + let num_blocks = strategy.len(); + let numel = values.len(); + assert_eq!( + numel / block_elems, + num_blocks, + "Invalid per-block quantization with num blocks {num_blocks} and {numel} values" + ); + values + .chunks(block_elems) + .enumerate() + .flat_map(|(block_id, block)| strategy[block_id].quantize(block)) + .collect() + } + } + } + + /// Dequantize the values to a higher precision data type. + pub fn dequantize(&self, values: &[i8]) -> Vec { + match self { + QuantizationStrategy::PerTensorSymmetric(strategy) => strategy.dequantize(values), + QuantizationStrategy::PerBlockSymmetric(strategy, block_size) => { + let block_elems = block_size.num_elements(); + let num_blocks = strategy.len(); + let numel = values.len(); + assert_eq!( + numel / block_elems, + num_blocks, + "Invalid per-block quantization with block size {block_elems}, num blocks {num_blocks} and {numel} values" + ); + values + .chunks(block_elems) + .enumerate() + .flat_map(|(block_id, block)| strategy[block_id].dequantize(block)) + .collect() + } + } + } +} + +/// Quantization scheme to convert elements of a higher precision data type `E` to a lower precision +/// data type `Q` and vice-versa. +pub trait Quantization { + /// Returns the quantization range `[a, b]`. + fn range(&self) -> (E, E); + /// Convert the values to a lower precision data type. + fn quantize(&self, values: &[E]) -> Vec; + /// Convert a single value to a lower precision data type. + fn quantize_one(&self, value: E) -> Q; + /// Convert the values back to a higher precision data type. + fn dequantize(&self, values: &[Q]) -> Vec; + /// Convert a single value back to a higher precision data type. + fn dequantize_one(&self, value: Q) -> E; +} + +fn valid_scale(mut scale: E) -> E { + // If scale is 0 (most likely due to a tensor full of zeros), we arbitrarily adjust the + // scale to 0.1 to avoid division by zero. + if scale.eq(&E::zero()) { + scale = E::from(0.1).unwrap(); + } + scale +} + +/// Symmetric quantization scheme. +#[derive(Debug, Clone, Copy)] +pub struct SymmetricQuantization { + /// The scaling factor. + pub scale: E, + // The quantization value data type. + value: QuantValue, +} + +impl SymmetricQuantization { + /// Initialize a symmetric quantization scheme with the given parameters. + pub fn init(scale: E, value: QuantValue) -> Self { + Self { + scale: valid_scale(scale), + value, + } + } + + #[allow(dead_code)] + /// Create a new quantization scheme for an input range `[alpha, beta]`. + fn new(alpha: E, beta: E, value: QuantValue) -> Self { + let (a, b) = value.range(); + let a = E::from(a).unwrap(); + let b = E::from(b).unwrap(); + + // Compute scale to convert a floating point value in range `[-alpha, alpha]` to the quantized range + let alpha = alpha.abs().max(beta.abs()); + let scale = valid_scale((alpha + alpha) / (b - a)); + Self { scale, value } + } +} + +impl Quantization for SymmetricQuantization { + fn quantize(&self, values: &[E]) -> Vec { + values.iter().map(|x| self.quantize_one(*x)).collect() + } + + fn dequantize(&self, values: &[Q]) -> Vec { + values.iter().map(|x_q| self.dequantize_one(*x_q)).collect() + } + + fn quantize_one(&self, value: E) -> Q { + let (a, b) = self.range(); + + // x_q = clamp(round(x / scale), a, b) + Q::from(value.div(self.scale).round().clamp(a, b)).unwrap() + } + + fn dequantize_one(&self, value: Q) -> E { + // x = scale * x_q + self.scale * E::from(value).unwrap() + } + + fn range(&self) -> (E, E) { + let (a, b) = self.value.range(); + let a = E::from(a).unwrap(); + let b = E::from(b).unwrap(); + (a, b) + } +} + +impl PartialEq for SymmetricQuantization { + fn eq(&self, other: &Self) -> bool { + self.scale == other.scale + } +} + +impl Eq for SymmetricQuantization {} + +#[cfg(test)] +mod tests { + use burn_backend::TensorData; + + use super::*; + use alloc::vec; + + #[test] + fn test_int8_symmetric_quantization() { + let x: [f32; 4] = [-1.8, -1.0, 0.0, 0.5]; + let expected_q = vec![-127, -71, 0, 35]; + let expected_d = vec![-1.8, -1.0062993, 0.0, 0.496063]; + + let symmetric = SymmetricQuantization::::new(-1.8, 0.5, QuantValue::Q8S); + + let q: Vec = symmetric.quantize(&x); + assert_eq!(q, expected_q); + + let d = symmetric.dequantize(&expected_q); + + assert_eq!(d, expected_d); + } + + #[test] + fn test_int8_symmetric_quantization_per_block() { + let x: [f32; 8] = [-1.8, -1.0, 0.0, 0.5, -1.8, -1.0, 0.0, 0.5]; + let expected_q = vec![-127, -71, 0, 35, -127, -71, 0, 35]; + let expected_d = vec![ + -1.8, -1.0062993, 0.0, 0.496063, -1.8, -1.0062993, 0.0, 0.496063, + ]; + + let symmetric = SymmetricQuantization::::new(-1.8, 0.5, QuantValue::Q8S); + let strategy = QuantizationStrategy::PerBlockSymmetric( + vec![symmetric, symmetric], + BlockSize::new([4]), + ); + + let q: Vec = strategy.quantize(&x); + assert_eq!(q, expected_q); + + let d = symmetric.dequantize(&expected_q); + + assert_eq!(d, expected_d); + } + + #[test] + fn should_support_dequantize() { + let strategy = QuantizationStrategy::PerTensorSymmetric(SymmetricQuantization { + scale: 0.1, + value: QuantValue::Q8S, + }); + + let output = strategy.dequantize(&[-127i8, -77, -26, 25, 76, 127]); + + let output = TensorData::new(output, [2, 3]); + + output.assert_approx_eq::( + &TensorData::from([[-12.7, -7.7, -2.6], [2.5, 7.6, 12.7]]), + Default::default(), + ); + } +} diff --git a/crates/burn-ndarray/src/ops/simd/avgpool.rs b/crates/burn-ndarray/src/ops/simd/avgpool.rs new file mode 100644 index 0000000..41d5ba6 --- /dev/null +++ b/crates/burn-ndarray/src/ops/simd/avgpool.rs @@ -0,0 +1,443 @@ +use core::{marker::PhantomData, mem::transmute}; + +use crate::{SharedArray, iter_range_par, run_par, sharing::UnsafeSharedRef}; + +use burn_backend::DType; +use burn_backend::{Element, ElementConversion}; +use bytemuck::Zeroable; +use macerator::{Simd, VAdd, VDiv}; +use ndarray::{Array4, s}; +use nhwc::avg_pool_nhwc; + +use super::should_use_simd; + +#[macerator::with_simd] +fn is_accelerated(_x: PhantomData) -> bool { + ::is_accelerated::() && ::is_accelerated::() +} + +pub(crate) fn try_avg_pool2d_simd( + x: SharedArray, + ksize: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + with_pad: bool, +) -> Result, SharedArray> { + // Strides must be unit, dilation isn't supported, rows must be contiguous + if x.strides()[1] != 1 || !should_use_simd(x.shape()[1]) { + return Err(x); + } + + match E::dtype() { + DType::F64 if is_accelerated::(PhantomData) => Ok(cast(avg_pool_nhwc::( + cast(x), + ksize, + stride, + padding, + with_pad, + ))), + DType::F32 if is_accelerated::(PhantomData) => Ok(cast(avg_pool_nhwc::( + cast(x), + ksize, + stride, + padding, + with_pad, + ))), + _ => Err(x), + } +} + +fn cast(tensor: SharedArray) -> SharedArray { + unsafe { transmute::, SharedArray>(tensor) } +} + +mod nhwc { + use itertools::Itertools; + use macerator::{Simd, Vector, vload_unaligned, vstore_unaligned}; + use ndarray::{ArrayView3, ArrayViewMut3}; + use seq_macro::seq; + + use crate::ops::simd::lanes; + + use super::*; + + // Until you can use associated constants as array size, we need to hardcode this. + // The most common config (x86-v3) has 16 registers, so use half of them for accumulators. + const BLOCK_REGISTERS: usize = 8; + + pub(crate) fn avg_pool_nhwc( + x: SharedArray, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + with_pad: bool, + ) -> SharedArray { + let [kernel_height, kernel_width] = kernel_size; + let [pad_h, pad_w] = padding; + let [stride_height, stride_width] = stride; + let [batch_size, channels, x_height, x_width] = x.shape().try_into().unwrap(); + let lanes = lanes::(); + + let ch_block = lanes * BLOCK_REGISTERS; + + let out_height = ((x_height + 2 * pad_h - (kernel_height - 1) - 1) / stride_height) + 1; + let out_width = ((x_width + 2 * pad_w - (kernel_width - 1) - 1) / stride_width) + 1; + + let mut output = unsafe { + Array4::::uninit((batch_size, out_height, out_width, channels)).assume_init() + }; + let unsafe_shared_out = UnsafeSharedRef::new(&mut output); + let x = x.view(); + let x = x.permuted_axes(vec![0, 2, 3, 1]); + + // Floor division ensures `blocks * lanes * blocking factor` is always `<= out_channels`. + // An exclusive loop will always have `lanes * blocking factor` elements in bounds. + let blocks = channels / ch_block; + let blocks_end = blocks * ch_block; + // Floor division means simd_end is always divisible by `lanes` and `<= out_channels`. An + // exclusive loop will always have `lanes` elements in bounds. + let simd_end = channels / lanes * lanes; + let num_simd_unblocked = (simd_end - blocks_end) / lanes; + let remainder = channels - simd_end; + + run_par!(|| { + // SAFETY: Loop ranges are non-overlapping, so the unsafe shared reference is safe. + iter_range_par!(0, batch_size * blocks).for_each(|k| unsafe { + let block = k % blocks; + let b = k / blocks; + + let output = unsafe_shared_out.get(); + + let x = x.slice(s![b, .., .., ..]); + let out = output.slice_mut(s![b, .., .., ..]); + + loop_blocked(x, out, kernel_size, stride, padding, with_pad, block); + }); + // SAFETY: See `loop_unblocked` + iter_range_par!(0, batch_size * num_simd_unblocked).for_each(|k| unsafe { + let ch = (k % num_simd_unblocked) * lanes + blocks_end; + let b = k / num_simd_unblocked; + + let output = unsafe_shared_out.get(); + + let x = x.slice(s![b, .., .., ..]); + let out = output.slice_mut(s![b, .., .., ..]); + + loop_unblocked(x, out, kernel_size, stride, padding, with_pad, ch); + }); + // SAFETY: Loop ranges are non-overlapping, so the unsafe shared reference is safe. + iter_range_par!(0, batch_size * remainder).for_each(|k| unsafe { + let ch = (k % remainder) + simd_end; + let b = k / remainder; + + let output = unsafe_shared_out.get(); + + let x = x.slice(s![b, .., .., ..]); + let out = output.slice_mut(s![b, .., .., ..]); + + loop_scalar(x, out, kernel_size, stride, padding, with_pad, ch); + }); + }); + + output = output.permuted_axes([0, 3, 1, 2]); + + output.into_dyn().into_shared() + } + + /// Execute the blocked (unrolled) portion of the pool. + #[allow( + clippy::too_many_arguments, + clippy::erasing_op, + clippy::identity_op, + unused_mut + )] + #[macerator::with_simd] + fn loop_blocked<'a, S: Simd, E: Element + VAdd + VDiv>( + x: ArrayView3<'a, E>, + mut out: ArrayViewMut3<'a, E>, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + with_pad: bool, + block: usize, + ) where + 'a: 'a, + { + let [kernel_height, kernel_width] = kernel_size; + let [pad_h, pad_w] = padding; + let [stride_height, stride_width] = stride; + + let (x_height, x_width, _) = x.dim(); + let (out_height, out_width, _) = out.dim(); + let lanes = E::lanes::(); + + let ch_block = lanes * BLOCK_REGISTERS; + + // If pixels are more than `padding` from the edges, the in pixel cannot be out of bounds + for oh in pad_h..out_height.saturating_sub(pad_h) { + for ow in pad_w..out_width.saturating_sub(pad_w) { + seq!(N in 0..8 { + let mut sum~N: Vector = Zeroable::zeroed(); + }); + let ch = block * ch_block; + let ch_end = ch + ch_block; + let mut out = out.slice_mut(s![oh, ow, ch..ch_end]); + + for kh in 0..kernel_height { + let ih = oh * stride_height + kh - pad_h; + + for kw in 0..kernel_width { + let iw = ow * stride_width + kw - pad_w; + let x = x.slice(s![ih, iw, ch..ch_end]); + + seq!(N in 0..8 { + // SAFETY: + // Load a full vector from x[N * lanes]. This is bounds checked by the + // slice above. + sum~N += unsafe { vload_unaligned(&x[N * lanes]) }; + }); + } + } + + let count = kernel_height * kernel_width; + let count = (count as u64).elem::(); + let count_v = count.splat(); + seq!(N in 0..8 { + let s~N = sum~N / count_v; + // SAFETY: + // Store a full vector to out[N * lanes]. This is bounds checked by the + // slice above. + unsafe { vstore_unaligned(&mut out[N * lanes], s~N) }; + }); + } + } + + // Border pixels need bounds checks + if (pad_h, pad_w) != (0, 0) { + let v_borders = (0..pad_h) + .chain(out_height.saturating_sub(pad_h)..out_height) + .cartesian_product(0..out_width); + let h_borders = (0..out_height) + .cartesian_product((0..pad_w).chain(out_width.saturating_sub(pad_w)..out_width)); + + for (oh, ow) in v_borders.chain(h_borders) { + seq!(N in 0..8 { + let mut sum~N: Vector = Zeroable::zeroed(); + }); + let mut count: usize = 0; + let ch = block * ch_block; + let ch_end = ch + ch_block; + let mut out = out.slice_mut(s![oh, ow, ch..ch_end]); + + for kh in 0..kernel_height { + let ih = oh * stride_height + kh; + if ih < pad_h || ih >= x_height + pad_h { + continue; + } + let ih = ih - pad_h; + + for kw in 0..kernel_width { + let iw = ow * stride_width + kw; + if iw < pad_w || iw >= x_width + pad_w { + continue; + } + let iw = iw - pad_w; + count += 1; + + let x = x.slice(s![ih, iw, ch..ch_end]); + + seq!(N in 0..8 { + // SAFETY: + // Load a full vector from x[N * lanes]. This is bounds checked by the + // slice above. + sum~N += unsafe { vload_unaligned(&x[N * lanes]) }; + }); + } + } + + if with_pad { + count = kernel_height * kernel_width; + } + + let count = (count as u64).elem::(); + let count_v = count.splat(); + seq!(N in 0..8 { + let s~N = sum~N / count_v; + // SAFETY: + // Store a full vector to out[N * lanes]. This is bounds checked by the + // slice above. + unsafe { vstore_unaligned(&mut out[N * lanes], s~N) }; + }); + } + } + } + + /// Execute the unblocked (not unrolled) portion of the pool. + /// + /// SAFETY: Safe as long as `ch + simd_lanes <= out_channels`. + #[allow(clippy::too_many_arguments, unused_mut)] + #[macerator::with_simd] + unsafe fn loop_unblocked<'a, S: Simd, E: Element + VAdd + VDiv>( + x: ArrayView3<'a, E>, + mut out: ArrayViewMut3<'a, E>, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + with_pad: bool, + ch: usize, + ) where + 'a: 'a, + { + let [kernel_height, kernel_width] = kernel_size; + let [pad_h, pad_w] = padding; + let [stride_height, stride_width] = stride; + + let (x_height, x_width, _) = x.dim(); + let (out_height, out_width, _) = out.dim(); + + // If pixels are not within padding range, bounds checks are always true + for oh in pad_h..out_height - pad_h { + for ow in pad_w..out_width - pad_w { + let mut sum: Vector = Zeroable::zeroed(); + + for kh in 0..kernel_height { + let ih = oh * stride_height + kh - pad_h; + + for kw in 0..kernel_width { + let iw = ow * stride_width + kw - pad_w; + // Load a full vector from `x`. In bounds as long as `out_channels >= ch + lanes` + let s0 = unsafe { vload_unaligned(&x[[ih, iw, ch]]) }; + sum += s0; + } + } + + let count = kernel_height * kernel_width; + let count: E = (count as u64).elem(); + let count_v = count.splat(); + let s0 = sum / count_v; + // Store a full vector to `out`. In bounds as long as `out_channels >= ch + lanes`. + unsafe { vstore_unaligned(&mut out[[oh, ow, ch]], s0) }; + } + } + + // Border pixels need bounds checks + if (pad_h, pad_w) != (0, 0) { + let v_borders = (0..pad_h) + .chain(out_height.saturating_sub(pad_h)..out_height) + .cartesian_product(0..out_width); + let h_borders = (0..out_height) + .cartesian_product((0..pad_w).chain(out_width.saturating_sub(pad_w)..out_width)); + + for (oh, ow) in v_borders.chain(h_borders) { + let mut sum: Vector = Zeroable::zeroed(); + let mut count: usize = 0; + + for kh in 0..kernel_height { + let ih = oh * stride_height + kh; + if ih < pad_h || ih >= x_height + pad_h { + continue; + } + let ih = ih - pad_h; + + for kw in 0..kernel_width { + let iw = ow * stride_width + kw; + if iw < pad_w || iw >= x_width + pad_w { + continue; + } + let iw = iw - pad_w; + count += 1; + + // Load a full vector from `x`. In bounds as long as `out_channels >= ch + lanes` + sum += unsafe { vload_unaligned(&x[[ih, iw, ch]]) }; + } + } + + if with_pad { + count = kernel_height * kernel_width; + } + + let count = (count as u64).elem::(); + let count_v = count.splat(); + let s0 = sum / count_v; + // Store a full vector to `out`. In bounds as long as `out_channels >= ch + lanes`. + unsafe { vstore_unaligned(&mut out[[oh, ow, ch]], s0) }; + } + } + } + + /// Execute scalar portion of the pooling + #[allow(clippy::too_many_arguments)] + fn loop_scalar( + x: ArrayView3<'_, E>, + mut out: ArrayViewMut3<'_, E>, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + with_pad: bool, + ch: usize, + ) { + let [kernel_height, kernel_width] = kernel_size; + let [pad_h, pad_w] = padding; + let [stride_height, stride_width] = stride; + + let (x_height, x_width, _) = x.dim(); + let (out_height, out_width, _) = out.dim(); + + // If pixels are not within padding range, bounds checks are always true + for oh in pad_h..out_height.saturating_sub(pad_h) { + for ow in pad_w..out_width.saturating_sub(pad_w) { + let mut sum: E = Zeroable::zeroed(); + + for kh in 0..kernel_height { + let ih = oh * stride_height + kh - pad_h; + + for kw in 0..kernel_width { + let iw = ow * stride_width + kw - pad_w; + sum = sum + x[[ih, iw, ch]]; + } + } + + let count = (kernel_height * kernel_width) as u64; + out[[oh, ow, ch]] = sum / count.elem(); + } + } + + // Border pixels need bounds checks + if (pad_h, pad_w) != (0, 0) { + let v_borders = (0..pad_h) + .chain(out_height.saturating_sub(pad_h)..out_height) + .cartesian_product(0..out_width); + let h_borders = (0..out_height) + .cartesian_product((0..pad_w).chain(out_width.saturating_sub(pad_w)..out_width)); + + for (oh, ow) in v_borders.chain(h_borders) { + let mut sum: E = Zeroable::zeroed(); + let mut count: usize = 0; + + for kh in 0..kernel_height { + let ih = oh * stride_height + kh; + if ih < pad_h || ih >= x_height + pad_h { + continue; + } + let ih = ih - pad_h; + + for kw in 0..kernel_width { + let iw = ow * stride_width + kw; + if iw < pad_w || iw >= x_width + pad_w { + continue; + } + let iw = iw - pad_w; + count += 1; + sum = sum + x[[ih, iw, ch]]; + } + } + + if with_pad { + count = kernel_height * kernel_width; + } + + out[[oh, ow, ch]] = sum / (count as u64).elem(); + } + } + } +} diff --git a/crates/burn-ndarray/src/ops/simd/base.rs b/crates/burn-ndarray/src/ops/simd/base.rs new file mode 100644 index 0000000..005316f --- /dev/null +++ b/crates/burn-ndarray/src/ops/simd/base.rs @@ -0,0 +1,115 @@ +use core::{marker::PhantomData, mem::MaybeUninit}; + +use macerator::{Arch, Scalar, Simd}; +use ndarray::{ArcArray, ArrayD, IxDyn, ShapeBuilder}; + +/// Whether SIMD instructions are worth using +#[cfg(all( + any( + target_arch = "x86", + target_arch = "x86_64", + target_arch = "aarch64", + target_arch = "wasm32", + target_arch = "loongarch64" + ), + not(test) +))] +pub fn should_use_simd(len: usize) -> bool { + len >= 32 +} + +/// Whether SIMD instructions are worth using +#[cfg(all( + not(any( + target_arch = "x86", + target_arch = "x86_64", + target_arch = "aarch64", + target_arch = "wasm32", + target_arch = "loongarch64" + )), + not(test) +))] +pub fn should_use_simd(_len: usize) -> bool { + false +} + +#[cfg(test)] +pub fn should_use_simd(_len: usize) -> bool { + true +} + +pub(crate) fn lanes() -> usize { + #[allow(non_camel_case_types)] + struct lanes<__T0>(__T0); + + impl ::macerator::WithSimd for lanes> { + type Output = usize; + #[inline(always)] + fn with_simd<__S: ::macerator::Simd>(self) -> ::Output { + let Self(__ty) = self; + #[allow(unused_unsafe)] + unsafe { + lanes_simd::<__S, E>(__ty) + } + } + } + (Arch::new()).dispatch(lanes(PhantomData::)) +} + +fn lanes_simd(_ty: PhantomData) -> usize { + E::lanes::() +} + +pub(crate) fn uninit_array_like(reference: &ArcArray) -> ArrayD { + let shape = reference.raw_dim(); + let strides = reference.strides(); + let strides = strides.iter().map(|it| *it as usize).collect::>(); + let shape_strides = shape.strides(IxDyn(&strides)); + let size = reference.len(); + let mut out_data: Vec> = Vec::with_capacity(size); + unsafe { out_data.set_len(size) }; + unsafe { ArrayD::from_shape_vec_unchecked(shape_strides, out_data).assume_init() } +} + +pub trait MinMax { + fn min(self, other: Self) -> Self; + fn max(self, other: Self) -> Self; +} + +macro_rules! impl_minmax { + ($ty: ty) => { + impl MinMax for $ty { + fn min(self, other: Self) -> Self { + Ord::min(self, other) + } + fn max(self, other: Self) -> Self { + Ord::max(self, other) + } + } + }; + ($($ty: ty),*) => { + $(impl_minmax!($ty);)* + } +} + +impl_minmax!(u8, i8, u16, i16, u32, i32, u64, i64); + +impl MinMax for f32 { + fn min(self, other: Self) -> Self { + self.min(other) + } + + fn max(self, other: Self) -> Self { + self.max(other) + } +} + +impl MinMax for f64 { + fn min(self, other: Self) -> Self { + self.min(other) + } + + fn max(self, other: Self) -> Self { + self.max(other) + } +} diff --git a/crates/burn-ndarray/src/ops/simd/binary.rs b/crates/burn-ndarray/src/ops/simd/binary.rs new file mode 100644 index 0000000..dae3ed5 --- /dev/null +++ b/crates/burn-ndarray/src/ops/simd/binary.rs @@ -0,0 +1,299 @@ +use core::{marker::PhantomData, slice}; + +use burn_backend::Element; +use macerator::{ + Scalar, Simd, VAdd, VBitAnd, VBitOr, VBitXor, VDiv, VMul, VOrd, VSub, Vector, vload_unaligned, + vstore_unaligned, +}; +use ndarray::ArrayD; +use seq_macro::seq; + +use crate::{NdArrayElement, SharedArray, ops::simd::uninit_array_like}; + +use super::{ + MinMax, + binary_elemwise::{ + VecAdd, VecBitAnd, VecBitOr, VecBitXor, VecDiv, VecMax, VecMin, VecMul, VecSub, + }, + should_use_simd, +}; + +pub trait SimdBinop { + fn apply_vec(lhs: Vector, rhs: Vector) -> Vector; + fn apply(lhs: T, rhs: T) -> Out; + fn is_accelerated() -> bool; +} + +impl SimdBinop for VecAdd { + fn apply_vec(lhs: Vector, rhs: Vector) -> Vector { + lhs + rhs + } + + fn apply(lhs: T, rhs: T) -> T { + lhs + rhs + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +impl SimdBinop for VecDiv { + fn apply_vec(lhs: Vector, rhs: Vector) -> Vector { + lhs / rhs + } + + fn apply(lhs: T, rhs: T) -> T { + lhs / rhs + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +impl SimdBinop for VecMul { + fn apply_vec(lhs: Vector, rhs: Vector) -> Vector { + lhs * rhs + } + + fn apply(lhs: T, rhs: T) -> T { + lhs * rhs + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +impl SimdBinop for VecSub { + fn apply_vec(lhs: Vector, rhs: Vector) -> Vector { + lhs - rhs + } + + fn apply(lhs: T, rhs: T) -> T { + lhs - rhs + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +impl SimdBinop for VecMin { + fn apply_vec(lhs: Vector, rhs: Vector) -> Vector { + lhs.min(rhs) + } + + fn apply(lhs: T, rhs: T) -> T { + MinMax::min(lhs, rhs) + } + + fn is_accelerated() -> bool { + ::is_min_max_accelerated::() + } +} + +impl SimdBinop for VecMax { + fn apply_vec(lhs: Vector, rhs: Vector) -> Vector { + lhs.max(rhs) + } + + fn apply(lhs: T, rhs: T) -> T { + MinMax::max(lhs, rhs) + } + + fn is_accelerated() -> bool { + ::is_min_max_accelerated::() + } +} + +impl SimdBinop for VecBitAnd { + fn apply_vec(lhs: Vector, rhs: Vector) -> Vector { + lhs & rhs + } + + fn apply(lhs: T, rhs: T) -> T { + lhs.bitand(rhs) + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +impl SimdBinop for VecBitOr { + fn apply_vec(lhs: Vector, rhs: Vector) -> Vector { + lhs | rhs + } + + fn apply(lhs: T, rhs: T) -> T { + lhs.bitor(rhs) + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +impl SimdBinop for VecBitXor { + fn apply_vec(lhs: Vector, rhs: Vector) -> Vector { + lhs ^ rhs + } + + fn apply(lhs: T, rhs: T) -> T { + lhs.bitxor(rhs) + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +#[macerator::with_simd] +fn is_accelerated>( + _x: PhantomData<(T, Out, Op)>, +) -> bool { + Op::is_accelerated::() +} + +#[allow(clippy::result_large_err)] +pub fn try_binary_simd< + E: Element, + EOut: Element, + T: NdArrayElement + Scalar, + Out: NdArrayElement + Scalar, + Op: SimdBinop, +>( + lhs: SharedArray, + rhs: SharedArray, +) -> Result, (SharedArray, SharedArray)> { + let lhs_len = lhs.len(); + let rhs_len = rhs.len(); + if !should_use_simd(lhs_len.max(rhs_len)) + || !lhs.is_standard_layout() + || !rhs.is_standard_layout() + || lhs.shape() != rhs.shape() + || !is_accelerated::(PhantomData) + { + return Err((lhs, rhs)); + } + // Used to assert traits based on the dynamic `DType`. + let lhs = unsafe { core::mem::transmute::, SharedArray>(lhs) }; + let rhs = unsafe { core::mem::transmute::, SharedArray>(rhs) }; + let out = binary_simd_same::(lhs, rhs); + + // Used to assert traits based on the dynamic `DType`. + let out = unsafe { core::mem::transmute::, SharedArray>(out) }; + Ok(out) +} + +fn binary_simd_same< + T: NdArrayElement + Scalar, + Out: NdArrayElement + Scalar, + Op: SimdBinop, +>( + lhs: SharedArray, + rhs: SharedArray, +) -> SharedArray { + let out = if lhs.is_unique() { + let mut buf = lhs.into_owned(); + let lhs = buf.as_slice_mut().unwrap(); + let rhs = rhs.as_slice().unwrap(); + let out = + unsafe { core::mem::transmute::<&mut [T], &mut [Out]>(unsafe_alias_slice_mut(lhs)) }; + binary(lhs, rhs, out, PhantomData::); + unsafe { core::mem::transmute::, ArrayD>(buf) } + } else if rhs.is_unique() { + let mut buf = rhs.into_owned(); + let lhs = lhs.as_slice().unwrap(); + let rhs = buf.as_slice_mut().unwrap(); + let out = + unsafe { core::mem::transmute::<&mut [T], &mut [Out]>(unsafe_alias_slice_mut(rhs)) }; + binary(lhs, rhs, out, PhantomData::); + unsafe { core::mem::transmute::, ArrayD>(buf) } + } else { + let mut out = uninit_array_like(&lhs); + let lhs = lhs.as_slice().unwrap(); + let rhs = rhs.as_slice().unwrap(); + let out_slice = out.as_slice_mut().unwrap(); + binary(lhs, rhs, out_slice, PhantomData::); + out + }; + out.into_shared() +} + +#[allow(clippy::erasing_op, clippy::identity_op)] +#[macerator::with_simd] +fn binary< + 'a, + S: Simd, + T: NdArrayElement + Scalar, + Out: NdArrayElement + Scalar, + Op: SimdBinop, +>( + lhs: &'a [T], + rhs: &'a [T], + out: &'a mut [Out], + _op: PhantomData, +) where + 'a: 'a, +{ + let lanes = T::lanes::(); + let mut chunks_lhs = lhs.chunks_exact(8 * lanes); + let mut chunks_rhs = rhs.chunks_exact(8 * lanes); + let mut chunks_out = out.chunks_exact_mut(8 * lanes); + while let Some(((lhs, rhs), out)) = chunks_lhs + .next() + .zip(chunks_rhs.next()) + .zip(chunks_out.next()) + { + seq!(N in 0..8 { + // Load one full vector from `lhs`. + // SAFETY: Guaranteed to be in bounds because `len == 8 * lanes` + let lhs~N = unsafe { vload_unaligned::(&lhs[N * lanes]) }; + // Load one full vector from `rhs`. + // SAFETY: Guaranteed to be in bounds because `len == 8 * lanes` + let rhs~N = unsafe { vload_unaligned(&rhs[N * lanes]) }; + let s~N = Op::apply_vec(lhs~N, rhs~N); + // Store one full vector to `out`. + // SAFETY: Guaranteed to be in bounds because `len == 8 * lanes` + unsafe { vstore_unaligned(&mut out[N * lanes], s~N) }; + }); + } + let mut chunks_lhs = chunks_lhs.remainder().chunks_exact(lanes); + let mut chunks_rhs = chunks_rhs.remainder().chunks_exact(lanes); + let mut chunks_out = chunks_out.into_remainder().chunks_exact_mut(lanes); + while let Some(((lhs, rhs), out)) = chunks_lhs + .next() + .zip(chunks_rhs.next()) + .zip(chunks_out.next()) + { + // Load one full vector from `lhs`. + // SAFETY: Guaranteed to be in bounds because `len == lanes` + let lhs0 = unsafe { vload_unaligned::(lhs.as_ptr()) }; + // Load one full vector from `rhs`. + // SAFETY: Guaranteed to be in bounds because `len == lanes` + let rhs0 = unsafe { vload_unaligned(rhs.as_ptr()) }; + let s0 = Op::apply_vec(lhs0, rhs0); + // Store one full vector to `out`. + // SAFETY: Guaranteed to be in bounds because `len == lanes` + unsafe { vstore_unaligned(out.as_mut_ptr(), s0) }; + } + + for ((lhs, rhs), out) in chunks_lhs + .remainder() + .iter() + .zip(chunks_rhs.remainder()) + .zip(chunks_out.into_remainder()) + { + *out = Op::apply(*lhs, *rhs) + } +} + +/// Unsafely alias a slice to use as an inline argument +fn unsafe_alias_slice_mut<'a, T>(slice: &mut [T]) -> &'a mut [T] { + let ptr = slice.as_mut_ptr(); + let len = slice.len(); + unsafe { slice::from_raw_parts_mut(ptr, len) } +} diff --git a/crates/burn-ndarray/src/ops/simd/binary_elemwise.rs b/crates/burn-ndarray/src/ops/simd/binary_elemwise.rs new file mode 100644 index 0000000..7534da5 --- /dev/null +++ b/crates/burn-ndarray/src/ops/simd/binary_elemwise.rs @@ -0,0 +1,419 @@ +use core::marker::PhantomData; + +use bytemuck::cast; +use macerator::{ + Scalar, Simd, VAdd, VBitAnd, VBitOr, VBitXor, VDiv, VMul, VOrd, VSub, Vector, vload, + vload_unaligned, vstore, vstore_unaligned, +}; +use ndarray::ArrayD; +use seq_macro::seq; + +use crate::{NdArrayElement, SharedArray, ops::simd::uninit_array_like}; + +use super::{MinMax, should_use_simd}; + +pub trait ScalarSimdBinop { + type Rhs: Copy; + type RhsVec: Copy; + fn splat(rhs: Self::Rhs) -> Self::RhsVec; + fn apply_vec(lhs: Vector, rhs: Self::RhsVec) -> Vector; + fn apply(lhs: T, rhs: Self::Rhs) -> Out; + fn is_accelerated() -> bool; +} + +pub struct VecAdd; +pub struct VecDiv; +pub struct VecMul; +pub struct VecSub; +pub struct VecMin; +pub struct VecMax; +pub struct VecClamp; +pub struct VecBitAnd; +pub struct VecBitOr; +pub struct VecBitXor; + +impl ScalarSimdBinop for VecAdd { + type Rhs = T; + type RhsVec = Vector; + + fn splat(rhs: Self::Rhs) -> Self::RhsVec { + rhs.splat() + } + + fn apply_vec(lhs: Vector, rhs: Self::RhsVec) -> Vector { + lhs + rhs + } + + fn apply(lhs: T, rhs: T) -> T { + lhs + rhs + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +impl ScalarSimdBinop for VecDiv { + type Rhs = T; + type RhsVec = Vector; + + fn splat(rhs: Self::Rhs) -> Self::RhsVec { + rhs.splat() + } + + fn apply_vec(lhs: Vector, rhs: Self::RhsVec) -> Vector { + lhs / rhs + } + + fn apply(lhs: T, rhs: T) -> T { + lhs / rhs + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +impl ScalarSimdBinop for VecMul { + type Rhs = T; + type RhsVec = Vector; + + fn splat(rhs: Self::Rhs) -> Self::RhsVec { + rhs.splat() + } + + fn apply_vec(lhs: Vector, rhs: Self::RhsVec) -> Vector { + lhs * rhs + } + + fn apply(lhs: T, rhs: T) -> T { + lhs * rhs + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +impl ScalarSimdBinop for VecSub { + type Rhs = T; + type RhsVec = Vector; + + fn splat(rhs: Self::Rhs) -> Self::RhsVec { + rhs.splat() + } + + fn apply_vec(lhs: Vector, rhs: Self::RhsVec) -> Vector { + lhs - rhs + } + + fn apply(lhs: T, rhs: T) -> T { + lhs - rhs + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +impl ScalarSimdBinop for VecMin { + type Rhs = T; + type RhsVec = Vector; + + fn splat(rhs: Self::Rhs) -> Self::RhsVec { + rhs.splat() + } + + fn apply_vec(lhs: Vector, rhs: Self::RhsVec) -> Vector { + lhs.min(rhs) + } + + fn apply(lhs: T, rhs: T) -> T { + lhs.min(rhs) + } + + fn is_accelerated() -> bool { + ::is_min_max_accelerated::() + } +} + +impl ScalarSimdBinop for VecMax { + type Rhs = T; + type RhsVec = Vector; + + fn splat(rhs: Self::Rhs) -> Self::RhsVec { + rhs.splat() + } + + fn apply_vec(lhs: Vector, rhs: Self::RhsVec) -> Vector { + lhs.max(rhs) + } + + fn apply(lhs: T, rhs: T) -> T { + lhs.max(rhs) + } + + fn is_accelerated() -> bool { + ::is_min_max_accelerated::() + } +} + +impl ScalarSimdBinop for VecClamp { + type Rhs = (T, T); + type RhsVec = (Vector, Vector); + + fn splat((min, max): Self::Rhs) -> Self::RhsVec { + (min.splat(), max.splat()) + } + + fn apply_vec(lhs: Vector, (min, max): Self::RhsVec) -> Vector { + lhs.min(max).max(min) + } + + fn apply(lhs: T, (min, max): Self::Rhs) -> T { + lhs.min(max).max(min) + } + + fn is_accelerated() -> bool { + ::is_min_max_accelerated::() + } +} + +impl ScalarSimdBinop for VecBitAnd { + type Rhs = T; + type RhsVec = Vector; + + fn splat(rhs: Self::Rhs) -> Self::RhsVec { + rhs.splat() + } + + fn apply_vec(lhs: Vector, rhs: Self::RhsVec) -> Vector { + lhs & rhs + } + + fn apply(lhs: T, rhs: Self::Rhs) -> T { + lhs & rhs + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +impl ScalarSimdBinop for VecBitOr { + type Rhs = T; + type RhsVec = Vector; + + fn splat(rhs: Self::Rhs) -> Self::RhsVec { + rhs.splat() + } + + fn apply_vec(lhs: Vector, rhs: Self::RhsVec) -> Vector { + lhs | rhs + } + + fn apply(lhs: T, rhs: Self::Rhs) -> T { + lhs | rhs + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +impl ScalarSimdBinop for VecBitXor { + type Rhs = T; + type RhsVec = Vector; + + fn splat(rhs: Self::Rhs) -> Self::RhsVec { + rhs.splat() + } + + fn apply_vec(lhs: Vector, rhs: Self::RhsVec) -> Vector { + lhs ^ rhs + } + + fn apply(lhs: T, rhs: Self::Rhs) -> T { + lhs ^ rhs + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +#[macerator::with_simd] +fn is_accelerated>( + _x: PhantomData<(T, Out, Op)>, +) -> bool { + Op::is_accelerated::() +} + +pub fn try_binary_scalar_simd< + E: NdArrayElement, + EOut: NdArrayElement, + T: NdArrayElement + Scalar, + Out: NdArrayElement + Scalar, + Op: ScalarSimdBinop, +>( + input: SharedArray, + elem: Op::Rhs, +) -> Result, SharedArray> { + if !should_use_simd(input.len()) + || input.as_slice_memory_order().is_none() + || !is_accelerated::(PhantomData) + { + return Err(input); + } + // Used to assert traits based on the dynamic `DType`. + let input = unsafe { core::mem::transmute::, SharedArray>(input) }; + let out = if size_of::() == size_of::() + && align_of::() >= align_of::() + && input.is_unique() + { + unsafe { binary_scalar_simd_inplace::(input, elem) } + } else { + binary_scalar_simd_owned::(input, elem) + }; + // Used to assert traits based on the dynamic `DType`. + let out = unsafe { core::mem::transmute::, SharedArray>(out) }; + Ok(out) +} + +/// Execute operation in place on an owned tensor +/// SAFETY: +/// Must ensure `size_of:: == size_of::` and `align_of:: >= align_of::`. +unsafe fn binary_scalar_simd_inplace< + T: NdArrayElement + Scalar, + Out: NdArrayElement + Scalar, + Op: ScalarSimdBinop, +>( + input: SharedArray, + elem: Op::Rhs, +) -> SharedArray { + let mut buffer = input.into_owned(); + let slice = buffer.as_slice_memory_order_mut().unwrap(); + unsafe { binary_scalar_slice_inplace::(slice, elem, PhantomData) }; + // Buffer has the same elem size and is filled with the operation output, so this is safe + let out = unsafe { core::mem::transmute::, ArrayD>(buffer) }; + out.into_shared() +} + +/// Create a new copy of the tensor as the output +fn binary_scalar_simd_owned< + T: NdArrayElement + Scalar, + Out: NdArrayElement + Scalar, + Op: ScalarSimdBinop, +>( + input: SharedArray, + elem: Op::Rhs, +) -> SharedArray { + let mut out = uninit_array_like(&input); + let input = input.as_slice_memory_order().unwrap(); + let out_slice = out.as_slice_memory_order_mut().unwrap(); + binary_scalar_slice::(input, out_slice, elem, PhantomData); + out.into_shared() +} + +#[inline(always)] +#[allow(clippy::erasing_op, clippy::identity_op)] +#[macerator::with_simd] +fn binary_scalar_slice< + 'a, + S: Simd, + T: NdArrayElement + Scalar, + Out: NdArrayElement + Scalar, + Op: ScalarSimdBinop, +>( + input: &'a [T], + out: &'a mut [Out], + rhs: Op::Rhs, + _op: PhantomData, +) where + 'a: 'a, +{ + let lanes = T::lanes::(); + let mut chunks_input = input.chunks_exact(8 * lanes); + let mut chunks_out = out.chunks_exact_mut(8 * lanes); + let rhs_vec = Op::splat::(rhs); + while let Some((input, out)) = chunks_input.next().zip(chunks_out.next()) { + seq!(N in 0..8 { + // Load one full vector from `input`. + // SAFETY: Guaranteed to be in bounds because `len == 8 * lanes` + let s~N = unsafe { vload_unaligned(&input[N * lanes]) }; + let s~N = Op::apply_vec(s~N, rhs_vec); + // Store one full vector to `out`. + // SAFETY: Guaranteed to be in bounds because `len == 8 * lanes` + unsafe { vstore_unaligned(&mut out[N * lanes], s~N) }; + }); + } + let mut chunks_input = chunks_input.remainder().chunks_exact(lanes); + let mut chunks_out = chunks_out.into_remainder().chunks_exact_mut(lanes); + while let Some((input, out)) = chunks_input.next().zip(chunks_out.next()) { + // Load one full vector from `input`. + // SAFETY: Guaranteed to be in bounds because `len == lanes` + let s0 = unsafe { vload_unaligned(input.as_ptr()) }; + let s0 = Op::apply_vec(s0, rhs_vec); + // Store one full vector to `out`. + // SAFETY: Guaranteed to be in bounds because `len == lanes` + unsafe { vstore_unaligned(out.as_mut_ptr(), s0) }; + } + + for (input, out) in chunks_input + .remainder() + .iter() + .zip(chunks_out.into_remainder()) + { + *out = Op::apply(*input, rhs) + } +} + +/// Execute operation in line. +/// SAFETY: +/// Must ensure `size_of:: == size_of::` and `align_of:: >= align_of::`. +#[inline(always)] +#[macerator::with_simd] +unsafe fn binary_scalar_slice_inplace< + 'a, + S: Simd, + T: NdArrayElement + Scalar, + Out: NdArrayElement + Scalar, + Op: ScalarSimdBinop, +>( + buf: &'a mut [T], + rhs: Op::Rhs, + _op: PhantomData<(Out, Op)>, +) where + 'a: 'a, +{ + let (head, main, tail) = unsafe { buf.align_to_mut::>() }; + for elem in head.iter_mut().chain(tail) { + *elem = cast(Op::apply(*elem, rhs)); + } + let mut chunks = main.chunks_exact_mut(8); + let rhs = Op::splat::(rhs); + for elem in chunks.by_ref() { + seq!(N in 0..8 { + // Load a full vector from the aligned portion of the buffer. + // SAFETY: `align_to_mut` guarantees we're aligned to `T::Vector`'s size, and there is + // always a full vector in bounds. + let s~N = unsafe { vload(&elem[N] as *const _ as *const T) }; + let s~N = Op::apply_vec(s~N, rhs); + // Store a full vector at the same position as the input. Cast is safe because `Out` is + // size and align compatible + unsafe { vstore_unaligned(&mut elem[N] as *mut _ as *mut Out, s~N) }; + }); + } + + for elem in chunks.into_remainder() { + // Load a full vector from the aligned portion of the buffer. + // SAFETY: `align_to_mut` guarantees we're aligned to `T::Vector`'s size, and there is + // always a full vector in bounds. + let s0 = unsafe { vload(elem as *const _ as *const T) }; + + let s0 = Op::apply_vec(s0, rhs); + // Store a full vector at the same position as the input. Cast is safe because `Out` is + // size and align compatible + unsafe { vstore(elem as *mut _ as *mut Out, s0) }; + } +} diff --git a/crates/burn-ndarray/src/ops/simd/cmp.rs b/crates/burn-ndarray/src/ops/simd/cmp.rs new file mode 100644 index 0000000..c9f8c0e --- /dev/null +++ b/crates/burn-ndarray/src/ops/simd/cmp.rs @@ -0,0 +1,374 @@ +use core::{marker::PhantomData, slice}; + +use burn_backend::Element; +use macerator::{Mask, Scalar, Simd, VEq, VOrd, Vector, vload_unaligned}; +use ndarray::ArrayD; +use seq_macro::seq; + +use crate::{NdArrayElement, SharedArray, ops::simd::uninit_array_like}; + +use super::should_use_simd; + +pub trait SimdCmpOp { + fn apply_vec(lhs: Vector, rhs: Vector) -> Mask; + fn apply(lhs: T, rhs: T) -> bool; + fn is_accelerated() -> bool; +} + +pub struct VecEquals; + +impl SimdCmpOp for VecEquals { + fn apply_vec(lhs: Vector, rhs: Vector) -> Mask { + lhs.eq(rhs) + } + + fn apply(lhs: T, rhs: T) -> bool { + lhs == rhs + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +pub struct VecGreater; + +impl SimdCmpOp for VecGreater { + fn apply_vec(lhs: Vector, rhs: Vector) -> Mask { + lhs.gt(rhs) + } + + fn apply(lhs: T, rhs: T) -> bool { + lhs > rhs + } + + fn is_accelerated() -> bool { + ::is_cmp_accelerated::() + } +} + +pub struct VecGreaterEq; + +impl SimdCmpOp for VecGreaterEq { + fn apply_vec(lhs: Vector, rhs: Vector) -> Mask { + lhs.ge(rhs) + } + + fn apply(lhs: T, rhs: T) -> bool { + lhs >= rhs + } + + fn is_accelerated() -> bool { + ::is_cmp_accelerated::() + } +} + +pub struct VecLowerEq; + +impl SimdCmpOp for VecLowerEq { + fn apply_vec(lhs: Vector, rhs: Vector) -> Mask { + lhs.le(rhs) + } + + fn apply(lhs: T, rhs: T) -> bool { + lhs <= rhs + } + + fn is_accelerated() -> bool { + ::is_cmp_accelerated::() + } +} + +pub struct VecLower; + +impl SimdCmpOp for VecLower { + fn apply_vec(lhs: Vector, rhs: Vector) -> Mask { + lhs.lt(rhs) + } + + fn apply(lhs: T, rhs: T) -> bool { + lhs < rhs + } + + fn is_accelerated() -> bool { + ::is_cmp_accelerated::() + } +} + +#[macerator::with_simd] +fn is_accelerated>(_x: PhantomData<(T, Op)>) -> bool { + Op::is_accelerated::() +} + +#[allow(clippy::result_large_err)] +pub fn try_cmp_simd>( + lhs: SharedArray, + rhs: SharedArray, +) -> Result, (SharedArray, SharedArray)> { + let lhs_len = lhs.len(); + let rhs_len = rhs.len(); + if !should_use_simd(lhs_len.max(rhs_len)) + || !lhs.is_standard_layout() + || !rhs.is_standard_layout() + || lhs.shape() != rhs.shape() + || !is_accelerated::(PhantomData) + { + return Err((lhs, rhs)); + } + // Used to assert traits based on the dynamic `DType`. + let lhs = unsafe { core::mem::transmute::, SharedArray>(lhs) }; + let rhs = unsafe { core::mem::transmute::, SharedArray>(rhs) }; + let out = cmp_simd_same::(lhs, rhs); + + Ok(out) +} + +fn cmp_simd_same>( + lhs: SharedArray, + rhs: SharedArray, +) -> SharedArray { + let out = if lhs.is_unique() && size_of::() == size_of::() { + let mut buf = lhs.into_owned(); + let lhs = buf.as_slice_mut().unwrap(); + let rhs = rhs.as_slice().unwrap(); + let out = + unsafe { core::mem::transmute::<&mut [T], &mut [bool]>(unsafe_alias_slice_mut(lhs)) }; + cmp(lhs, rhs, out, PhantomData::); + unsafe { core::mem::transmute::, ArrayD>(buf) } + } else if rhs.is_unique() && size_of::() == size_of::() { + let mut buf = rhs.into_owned(); + let lhs = lhs.as_slice().unwrap(); + let rhs = buf.as_slice_mut().unwrap(); + let out = + unsafe { core::mem::transmute::<&mut [T], &mut [bool]>(unsafe_alias_slice_mut(rhs)) }; + cmp(lhs, rhs, out, PhantomData::); + unsafe { core::mem::transmute::, ArrayD>(buf) } + } else { + let mut out = uninit_array_like(&lhs); + let lhs = lhs.as_slice().unwrap(); + let rhs = rhs.as_slice().unwrap(); + let out_slice = out.as_slice_mut().unwrap(); + cmp(lhs, rhs, out_slice, PhantomData::); + out + }; + out.into_shared() +} + +#[allow(clippy::erasing_op, clippy::identity_op)] +#[macerator::with_simd] +fn cmp<'a, S: Simd, T: NdArrayElement + Scalar, Op: SimdCmpOp>( + lhs: &'a [T], + rhs: &'a [T], + out: &'a mut [bool], + _op: PhantomData, +) where + 'a: 'a, +{ + let lanes = T::lanes::(); + let mut chunks_lhs = lhs.chunks_exact(8 * lanes); + let mut chunks_rhs = rhs.chunks_exact(8 * lanes); + let mut chunks_out = out.chunks_exact_mut(8 * lanes); + while let Some(((lhs, rhs), out)) = chunks_lhs + .next() + .zip(chunks_rhs.next()) + .zip(chunks_out.next()) + { + seq!(N in 0..8 { + // Load one full vector from `lhs`. + // SAFETY: Guaranteed to be in bounds because `len == 8 * lanes` + let lhs~N = unsafe { vload_unaligned::(&lhs[N * lanes]) }; + // Load one full vector from `rhs`. + // SAFETY: Guaranteed to be in bounds because `len == 8 * lanes` + let rhs~N = unsafe { vload_unaligned(&rhs[N * lanes]) }; + let s~N = Op::apply_vec(lhs~N, rhs~N); + // Store one full vector to `out`. + // SAFETY: Guaranteed to be in bounds because `len == 8 * lanes` + unsafe { T::mask_store_as_bool(&mut out[N * lanes], s~N) }; + }); + } + let mut chunks_lhs = chunks_lhs.remainder().chunks_exact(lanes); + let mut chunks_rhs = chunks_rhs.remainder().chunks_exact(lanes); + let mut chunks_out = chunks_out.into_remainder().chunks_exact_mut(lanes); + while let Some(((lhs, rhs), out)) = chunks_lhs + .next() + .zip(chunks_rhs.next()) + .zip(chunks_out.next()) + { + // Load one full vector from `lhs`. + // SAFETY: Guaranteed to be in bounds because `len == lanes` + let lhs0 = unsafe { vload_unaligned::(lhs.as_ptr()) }; + // Load one full vector from `rhs`. + // SAFETY: Guaranteed to be in bounds because `len == lanes` + let rhs0 = unsafe { vload_unaligned(rhs.as_ptr()) }; + let s0 = Op::apply_vec(lhs0, rhs0); + // Store one full vector to `out`. + // SAFETY: Guaranteed to be in bounds because `len == lanes` + unsafe { T::mask_store_as_bool(out.as_mut_ptr(), s0) }; + } + + for ((lhs, rhs), out) in chunks_lhs + .remainder() + .iter() + .zip(chunks_rhs.remainder()) + .zip(chunks_out.into_remainder()) + { + *out = Op::apply(*lhs, *rhs) + } +} + +/// Unsafely alias a slice to use as an inline argument +fn unsafe_alias_slice_mut<'a, T>(slice: &mut [T]) -> &'a mut [T] { + let ptr = slice.as_mut_ptr(); + let len = slice.len(); + unsafe { slice::from_raw_parts_mut(ptr, len) } +} + +pub use elemwise::try_cmp_scalar_simd; + +mod elemwise { + use bytemuck::cast; + use macerator::vload; + + use super::*; + + pub fn try_cmp_scalar_simd>( + input: SharedArray, + elem: T, + ) -> Result, SharedArray> { + if !should_use_simd(input.len()) + || input.as_slice_memory_order().is_none() + || !is_accelerated::(PhantomData) + { + return Err(input); + } + // Used to assert traits based on the dynamic `DType`. + let input = unsafe { core::mem::transmute::, SharedArray>(input) }; + let out = if size_of::() == size_of::() + && align_of::() >= align_of::() + && input.is_unique() + { + unsafe { cmp_scalar_simd_inplace::(input, elem) } + } else { + cmp_scalar_simd_owned::(input, elem) + }; + Ok(out) + } + + /// Execute operation in place on an owned tensor + /// SAFETY: + /// Must ensure `size_of:: == size_of::` and `align_of:: >= align_of::`. + unsafe fn cmp_scalar_simd_inplace>( + input: SharedArray, + elem: T, + ) -> SharedArray { + let mut buffer = input.into_owned(); + let slice = buffer.as_slice_memory_order_mut().unwrap(); + unsafe { cmp_scalar_slice_inplace::(slice, elem, PhantomData) }; + // Buffer has the same elem size and is filled with the operation output, so this is safe + let out = unsafe { core::mem::transmute::, ArrayD>(buffer) }; + out.into_shared() + } + + /// Create a new copy of the tensor as the output + fn cmp_scalar_simd_owned>( + input: SharedArray, + elem: T, + ) -> SharedArray { + let mut out = uninit_array_like(&input); + let input = input.as_slice_memory_order().unwrap(); + let out_slice = out.as_slice_memory_order_mut().unwrap(); + cmp_scalar_slice::(input, out_slice, elem, PhantomData); + out.into_shared() + } + + #[inline(always)] + #[allow(clippy::erasing_op, clippy::identity_op)] + #[macerator::with_simd] + fn cmp_scalar_slice<'a, S: Simd, T: NdArrayElement + Scalar, Op: SimdCmpOp>( + input: &'a [T], + out: &'a mut [bool], + rhs: T, + _op: PhantomData, + ) where + 'a: 'a, + { + let lanes = T::lanes::(); + let mut chunks_input = input.chunks_exact(8 * lanes); + let mut chunks_out = out.chunks_exact_mut(8 * lanes); + let rhs_vec = rhs.splat::(); + while let Some((input, out)) = chunks_input.next().zip(chunks_out.next()) { + seq!(N in 0..8 { + // Load one full vector from `input`. + // SAFETY: Guaranteed to be in bounds because `len == 8 * lanes` + let s~N = unsafe { vload_unaligned(&input[N * lanes]) }; + let s~N = Op::apply_vec(s~N, rhs_vec); + // Store one full vector to `out`. + // SAFETY: Guaranteed to be in bounds because `len == 8 * lanes` + unsafe { T::mask_store_as_bool(&mut out[N * lanes], s~N) }; + }); + } + let mut chunks_input = chunks_input.remainder().chunks_exact(lanes); + let mut chunks_out = chunks_out.into_remainder().chunks_exact_mut(lanes); + while let Some((input, out)) = chunks_input.next().zip(chunks_out.next()) { + // Load one full vector from `input`. + // SAFETY: Guaranteed to be in bounds because `len == lanes` + let s0 = unsafe { vload_unaligned(input.as_ptr()) }; + let s0 = Op::apply_vec(s0, rhs_vec); + // Store one full vector to `out`. + // SAFETY: Guaranteed to be in bounds because `len == lanes` + unsafe { T::mask_store_as_bool(out.as_mut_ptr(), s0) }; + } + + for (input, out) in chunks_input + .remainder() + .iter() + .zip(chunks_out.into_remainder()) + { + *out = Op::apply(*input, rhs) + } + } + + /// Execute operation in line. + /// SAFETY: + /// Must ensure `size_of:: == size_of::` and `align_of:: >= align_of::`. + #[inline(always)] + #[macerator::with_simd] + unsafe fn cmp_scalar_slice_inplace<'a, S: Simd, T: NdArrayElement + Scalar, Op: SimdCmpOp>( + buf: &'a mut [T], + rhs: T, + _op: PhantomData, + ) where + 'a: 'a, + { + let (head, main, tail) = unsafe { buf.align_to_mut::>() }; + for elem in head.iter_mut().chain(tail) { + *elem = cast(Op::apply(*elem, rhs)); + } + let mut chunks = main.chunks_exact_mut(8); + let rhs = rhs.splat::(); + for elem in chunks.by_ref() { + seq!(N in 0..8 { + // Load a full vector from the aligned portion of the buffer. + // SAFETY: `align_to_mut` guarantees we're aligned to `T::Vector`'s size, and there is + // always a full vector in bounds. + let s~N = unsafe { vload(&elem[N] as *const _ as *const T) }; + let s~N = Op::apply_vec(s~N, rhs); + // Store a full vector at the same position as the input. Cast is safe because `Out` is + // size and align compatible + unsafe { T::mask_store_as_bool(&mut elem[N] as *mut _ as *mut bool, s~N) }; + }); + } + + for elem in chunks.into_remainder() { + // Load a full vector from the aligned portion of the buffer. + // SAFETY: `align_to_mut` guarantees we're aligned to `T::Vector`'s size, and there is + // always a full vector in bounds. + let s0 = unsafe { vload(elem as *const _ as *const T) }; + + let s0 = Op::apply_vec(s0, rhs); + // Store a full vector at the same position as the input. Cast is safe because `Out` is + // size and align compatible + unsafe { T::mask_store_as_bool(elem as *mut _ as *mut bool, s0) }; + } + } +} diff --git a/crates/burn-ndarray/src/ops/simd/conv.rs b/crates/burn-ndarray/src/ops/simd/conv.rs new file mode 100644 index 0000000..496a1ad --- /dev/null +++ b/crates/burn-ndarray/src/ops/simd/conv.rs @@ -0,0 +1,494 @@ +use core::{marker::PhantomData, mem::transmute}; + +use burn_backend::{ + DType, Element, + ops::{ConvOptions, conv::calculate_conv_output_size}, +}; +use bytemuck::Zeroable; +use macerator::{Simd, VMulAdd, Vector, vload_unaligned, vstore_unaligned}; +use ndarray::{ + ArcArray1, Array4, ArrayView3, ArrayView4, ArrayViewMut2, ArrayViewMut3, Dim, Ix1, Ix4, s, +}; +use seq_macro::seq; + +use crate::{FloatNdArrayElement, SharedArray, UnsafeSharedRef, iter_range_par, run_par}; + +type Args = (SharedArray, SharedArray, Option>); + +#[allow(clippy::result_large_err)] +pub fn try_conv2d_simd( + x: SharedArray, + weight: SharedArray, + bias: Option>, + options: ConvOptions<2>, +) -> Result, Args> { + match E::dtype() { + DType::F64 => conv2d::(x, weight, bias, options, PhantomData), + DType::F32 => conv2d::(x, weight, bias, options, PhantomData), + DType::I64 => conv2d::(x, weight, bias, options, PhantomData), + DType::I32 => conv2d::(x, weight, bias, options, PhantomData), + DType::I16 => conv2d::(x, weight, bias, options, PhantomData), + DType::U64 => conv2d::(x, weight, bias, options, PhantomData), + DType::U32 => conv2d::(x, weight, bias, options, PhantomData), + DType::U16 => conv2d::(x, weight, bias, options, PhantomData), + _ => Err((x, weight, bias)), + } +} + +fn cast(tensor: SharedArray) -> SharedArray { + unsafe { transmute::, SharedArray>(tensor) } +} + +/// Out-channel last SIMD accelerated direct convolution. Loop order and register blocking based on +/// E. Georganas, S. Avancha, K. Banerjee, D. Kalamkar, G. Henry, H. Pabst, A. Heinecke (2018). +/// Anatomy Of High-Performance Deep Learning Convolutions On SIMD Architectures. +/// SC '18, Article 6, pp. 1-12. arXiv:1808.05567. . +#[allow(clippy::result_large_err)] +fn conv2d( + x: SharedArray, + weight: SharedArray, + bias: Option>, + options: ConvOptions<2>, + _ty: PhantomData, +) -> Result, Args> { + let [out_channels, _, k_height, k_width] = weight.shape().try_into().unwrap(); + let channels_per_group = out_channels / options.groups; + + #[macerator::with_simd] + fn precheck(_ty: PhantomData) -> (usize, bool) { + (E::lanes::(), E::is_accelerated::()) + } + + let (lanes, accelerated) = precheck::(PhantomData); + + if !accelerated || !channels_per_group.is_multiple_of(lanes) { + return Err((x, weight, bias)); + } + + let x = cast::<_, E>(x); + let weight = cast::<_, E>(weight); + let bias = bias.map(|bias| cast::<_, E>(bias)); + + let [batch_size, _in_channels, in_height, in_width] = x.shape().try_into().unwrap(); + let [dilate_h, dilate_w] = options.dilation; + let [stride_h, stride_w] = options.stride; + let [pad_h, pad_w] = options.padding; + let padded = options.padding != [0, 0]; + let strided = options.stride != [1, 1] || options.dilation != [1, 1]; + let grouped = options.groups != 1; + + let out_height = calculate_conv_output_size(k_height, stride_h, pad_h, dilate_h, in_height); + let out_width = calculate_conv_output_size(k_width, stride_w, pad_w, dilate_w, in_width); + + let x = x.into_dimensionality::().unwrap(); + let weights = weight.into_dimensionality::().unwrap(); + let weights = weights.permuted_axes([1, 2, 3, 0]); + let weights = weights.as_standard_layout(); + let bias = bias.map(|bias| bias.into_dimensionality::().unwrap()); + // floor division means `(oc_blocks - 1) * lanes` can never be greater than `out_channels - lanes`. + let oc_blocks = out_channels / lanes; + + let mut out = unsafe { + Array4::::uninit(Dim([batch_size, out_height, out_width, out_channels])).assume_init() + }; + let unsafe_shared_out = UnsafeSharedRef::new(&mut out); + + run_par!(|| { + // SAFETY: Slices are guaranteed to be non-overlapping, so having an unsafe shared reference + // is safe. `oc_blocks * lanes` must be `<= out_channels` to satisfy safety of inner function. + iter_range_par!(0, batch_size * oc_blocks).for_each(|k| unsafe { + let b = k / oc_blocks; + let ob = k % oc_blocks; + let x = x.slice(s![b, .., .., ..]); + let out = unsafe_shared_out.get(); + let mut out = out.slice_mut(s![b, .., .., ..]); + let w = weights.view(); + + match (padded, strided, grouped) { + (true, true, true) => { + conv2d_launch::(x, w, &bias, &mut out, &options, ob) + } + (true, false, true) => { + conv2d_launch::(x, w, &bias, &mut out, &options, ob) + } + (false, true, true) => { + conv2d_launch::(x, w, &bias, &mut out, &options, ob) + } + (false, false, true) => { + conv2d_launch::(x, w, &bias, &mut out, &options, ob) + } + (true, true, false) => { + conv2d_launch::(x, w, &bias, &mut out, &options, ob) + } + (true, false, false) => { + conv2d_launch::(x, w, &bias, &mut out, &options, ob) + } + (false, true, false) => { + conv2d_launch::(x, w, &bias, &mut out, &options, ob) + } + (false, false, false) => { + conv2d_launch::(x, w, &bias, &mut out, &options, ob) + } + } + }); + }); + + let output = out.permuted_axes([0, 3, 1, 2]); + Ok(cast(output.into_dyn().into_shared())) +} + +/// Size of register blocks, we need to hardcode this because Rust and the `seq` macro don't support +/// using associated constants as constant parameters. 8 works for all semi-modern CPUs but might +/// not be perfectly optimized for AVX-512 capable CPUs (which probably should use 16). +/// This should always be conservative, since oversizing it will cause register spills and that's +/// **much** worse than the performance lost with lower values. +const REGISTER_BLOCK: usize = 8; +inner_with_register_blocking_size!(8); + +/// Run a loop of conv2d. +/// # SAFETY +/// See `conv2d_inner_nopad`, `conv2d_inner_nopad_nostride`, `conv2d_remainder`. +/// Required preconditions: `ob * simd_lanes` must be `<= out_channels - simd_lanes`, `weights` and +/// `out` must have unit stride for the out channels. +#[inline(always)] +#[macerator::with_simd] +unsafe fn conv2d_launch< + 'a, + S: Simd, + E: VMulAdd, + const PAD: bool, + const STRIDE: bool, + const GROUPS: bool, +>( + x: ArrayView3<'a, E>, + weights: ArrayView4<'a, E>, + bias: &'a Option>, + out: &'a mut ArrayViewMut3<'a, E>, + options: &'a ConvOptions<2>, + ob: usize, +) where + 'a: 'a, +{ + let (in_channels, k_height, k_width, out_channels) = weights.dim(); + let (out_height, out_width, _) = out.dim(); + let channels_per_group = out_channels / options.groups; + let lanes = E::lanes::(); + + let [mut pad_h, mut pad_w] = options.padding; + let [stride_h, stride_w] = options.stride; + let [dilate_h, dilate_w] = options.dilation; + + // Trick compiler into inlining 0 to padding + if !PAD { + pad_h = 0; + pad_w = 0; + } + + let oc_b = channels_per_group.min(lanes); + let ow_b = REGISTER_BLOCK; + + let ow_start = pad_w.min(out_width); + let ow_width = out_width.saturating_sub(2 * pad_w); + let oh_start = pad_h.min(out_height); + let oh_end = out_height.saturating_sub(pad_h); + + let ow_blocks = ow_width / ow_b; + let oc = ob * oc_b; + let group = oc / channels_per_group; + let mut ic_off = group * in_channels; + if !GROUPS { + ic_off = 0; + } + + unsafe { + let bias = if let Some(bias) = &bias { + vload_unaligned::(&bias[oc]) + } else { + Zeroable::zeroed() + }; + + for oh in oh_start..oh_end { + let mut out = out.slice_mut(s![oh, .., ..]); + for ow_block in 0..ow_blocks { + let ow = ow_block * ow_b + ow_start; + + #[allow(clippy::if_same_then_else)] + if STRIDE { + conv2d_inner_nopad( + &x, &weights, &mut out, bias, oh, ow, oc, ic_off, stride_h, stride_w, + dilate_h, dilate_w, k_height, k_width, pad_h, pad_w, + ); + } else { + conv2d_inner_nopad_nostride( + &x, &weights, &mut out, bias, oh, ow, oc, ic_off, k_height, k_width, pad_h, + pad_w, + ); + } + } + } + conv2d_remainder( + x, + weights, + out, + bias, + oc, + ic_off, + ow_blocks * ow_b, + stride_h, + stride_w, + dilate_h, + dilate_w, + pad_h, + pad_w, + k_height, + k_width, + ); + } +} + +/// Execute the non-unrolled and/or padded portion of the convolution. This has more checks and is +/// much slower, so we want to minimize the amount of pixels that need to be processed by this +/// +/// SAFETY: `oc` must be an index that's at most `out_channels - simd_lanes`, so the full vector +/// is in bounds. Weights and `out` must be channels last (with `stride == 1`). +#[allow(clippy::too_many_arguments)] +#[inline(always)] +unsafe fn conv2d_remainder( + x: ArrayView3, + weights: ArrayView4, + out: &mut ArrayViewMut3, + bias: Vector, + oc: usize, + ic_off: usize, + owb_end: usize, + stride_h: usize, + stride_w: usize, + dilate_h: usize, + dilate_w: usize, + pad_h: usize, + pad_w: usize, + k_height: usize, + k_width: usize, +) { + let in_channels = weights.shape()[0]; + let (_, in_height, in_width) = x.dim(); + let (out_height, out_width, _) = out.dim(); + let oh_start = pad_h.min(out_height); + let oh_end = out_height.saturating_sub(pad_h); + let ow_start = pad_w.min(out_width); + + let height1 = in_height + pad_h; + let width1 = in_width + pad_w; + + for oh in (0..oh_start).chain(oh_end..out_height) { + for ow in 0..out_width { + let mut acc = bias; + + for ic in 0..in_channels { + for kh in 0..k_height { + let ih = oh * stride_h + kh * dilate_h; + if (ih < pad_h) | (ih >= height1) { + continue; + } + let ih = ih - pad_h; + + for kw in 0..k_width { + let iw = ow * stride_w + kw * dilate_w; + if (iw < pad_w) | (iw >= width1) { + continue; + } + let iw = iw - pad_w; + + // Load a full vector from the weights. This is guaranteed to be in bounds + // as long as `oc <= out_channels - simd_lanes` and out channels are last. + // We need to ensure the weights are reshaped appropriately. + let f0 = unsafe { vload_unaligned(&weights[[ic, kh, kw, oc]]) }; + + // The loop bounds ensure `ic`, `ih` and `iw` are always in bounds, but the + // compiler can't prove this. We can't use `as_slice` with fixed bounds + // because we want to support arbitrary input layouts. So an unchecked load + // is used. + let i0 = unsafe { x.uget([ic_off + ic, ih, iw]) }.splat::(); + acc = i0.mul_add(f0, acc); + } + } + } + + // Store a full vector from the output. This is guaranteed to be in bounds + // as long as `oc <= out_channels - simd_lanes` and oc stride is 1. We create `out` with + // channels last, so this always holds. + unsafe { vstore_unaligned(&mut out[[oh, ow, oc]], acc) }; + } + } + for ow in (0..ow_start).chain(owb_end..out_width) { + for oh in 0..out_height { + let mut acc = bias; + + for ic in 0..in_channels { + for kh in 0..k_height { + let ih = oh * stride_h + kh * dilate_h; + if (ih < pad_h) | (ih >= height1) { + continue; + } + let ih = ih - pad_h; + + for kw in 0..k_width { + let iw = ow * stride_w + kw * dilate_w; + if (iw < pad_w) | (iw >= width1) { + continue; + } + let iw = iw - pad_w; + + // Load a full vector from the weights. This is guaranteed to be in bounds + // as long as `oc <= out_channels - simd_lanes` and out channels are last. + // We need to ensure the weights are reshaped appropriately. + let f0 = unsafe { vload_unaligned(&weights[[ic, kh, kw, oc]]) }; + + // The loop bounds ensure `ic`, `ih` and `iw` are always in bounds, but the + // compiler can't prove this. We can't use `as_slice` with fixed bounds + // because we want to support arbitrary input layouts. So an unchecked load + // is used. + let i0 = unsafe { x.uget([ic_off + ic, ih, iw]) }.splat::(); + acc = i0.mul_add(f0, acc); + } + } + } + + // Store a full vector from the output. This is guaranteed to be in bounds + // as long as `oc <= out_channels - simd_lanes` and oc stride is 1. We create `out` with + // channels last, so this always holds. + unsafe { vstore_unaligned(&mut out[[oh, ow, oc]], acc) }; + } + } +} + +macro_rules! inner_with_register_blocking_size { + ($rb: literal) => { + /// Execute the unrolled and unpadded portion of the convolution. Any pixel that is more than + /// `pad_h` away from the horizontal border, and `pad_w` away from the vertical border is + /// guaranteed to always be in bounds (because of the way out size is calculated). + /// + /// SAFETY: `oc` must be an index that's at most `out_channels - simd_lanes`, so the full vector + /// is in bounds. Weights and `out` must be channels last (with `stride == 1`). + #[allow(clippy::erasing_op, clippy::identity_op, clippy::too_many_arguments)] + #[inline(always)] + unsafe fn conv2d_inner_nopad( + x: &ArrayView3, + weights: &ArrayView4, + out: &mut ArrayViewMut2, + bias: Vector, + oh: usize, + ow: usize, + oc: usize, + ic_off: usize, + stride_h: usize, + stride_w: usize, + dilate_h: usize, + dilate_w: usize, + k_height: usize, + k_width: usize, + pad_h: usize, + pad_w: usize, + ) { + let in_channels = weights.shape()[0]; + + seq!(N in 0..$rb { + let mut acc~N = bias; + }); + + for ic in 0..in_channels { + for kh in 0..k_height { + let ih = oh * stride_h + kh * dilate_h - pad_h; + + for kw in 0..k_width { + // Load a full vector from the weights. This is guaranteed to be in bounds + // as long as `oc <= out_channels - simd_lanes` and out channels are last. + // We need to ensure the weights are reshaped appropriately. + let f0 = unsafe { vload_unaligned(&weights[[ic, kh, kw, oc]]) }; + let iw = ow * stride_w + kw * dilate_w - pad_w; + + seq!(N in 0..$rb { + // The loop bounds ensure `ic`, `ih` and `iw` are always in bounds, but the + // compiler can't prove this. We can't use `as_slice` with fixed bounds + // because we want to support arbitrary input layouts. So an unchecked load + // is used. + let i~N = unsafe { x.uget([ic + ic_off, ih, iw + N * stride_w]) }.splat::(); + }); + seq!(N in 0..$rb { + acc~N = i~N.mul_add(f0, acc~N); + }); + } + } + } + + seq!(N in 0..$rb { + // Store a full vector from the output. This is guaranteed to be in bounds + // as long as `oc <= out_channels - simd_lanes` and oc stride is 1. We create `out` with + // channels last, so this always holds. + unsafe { vstore_unaligned(&mut out[[ow + N, oc]], acc~N) }; + }); + } + + /// Execute the unrolled and unpadded portion of the convolution. Any pixel that is more than + /// `pad_h` away from the horizontal border, and `pad_w` away from the vertical border is + /// guaranteed to always be in bounds (because of the way out size is calculated). + /// + /// SAFETY: `oc` must be an index that's at most `out_channels - simd_lanes`, so the full vector + /// is in bounds. Weights and `out` must be channels last (with `stride == 1`). + #[allow(clippy::erasing_op, clippy::identity_op, clippy::too_many_arguments)] + #[inline(always)] + unsafe fn conv2d_inner_nopad_nostride( + x: &ArrayView3, + weights: &ArrayView4, + out: &mut ArrayViewMut2, + bias: Vector, + oh: usize, + ow: usize, + oc: usize, + ic_off: usize, + k_height: usize, + k_width: usize, + pad_h: usize, + pad_w: usize, + ) { + let in_channels = weights.shape()[0]; + + seq!(N in 0..$rb { + let mut acc~N = bias; + }); + + for ic in 0..in_channels { + for kh in 0..k_height { + let ih = oh + kh - pad_h; + + for kw in 0..k_width { + // Load a full vector from the weights. This is guaranteed to be in bounds + // as long as `oc <= out_channels - simd_lanes` and out channels are last. + // We need to ensure the weights are reshaped appropriately. + let f0 = unsafe { vload_unaligned(&weights[[ic, kh, kw, oc]]) }; + let iw = ow + kw - pad_w; + + seq!(N in 0..$rb { + // The loop bounds ensure `ic`, `ih` and `iw` are always in bounds, but the + // compiler can't prove this. We can't use `as_slice` with fixed bounds + // because we want to support arbitrary input layouts. So an unchecked load + // is used. + let i~N = unsafe { x.uget([ic + ic_off, ih, iw + N]) }.splat::(); + }); + seq!(N in 0..$rb { + acc~N = i~N.mul_add(f0, acc~N); + }); + } + } + } + + seq!(N in 0..$rb { + // Store a full vector from the output. This is guaranteed to be in bounds + // as long as `oc <= out_channels - simd_lanes` and oc stride is 1. We create `out` with + // channels last, so this always holds. + unsafe { vstore_unaligned(&mut out[[ow + N, oc]], acc~N) }; + }); + } + }; +} +pub(crate) use inner_with_register_blocking_size; diff --git a/crates/burn-ndarray/src/ops/simd/maxpool.rs b/crates/burn-ndarray/src/ops/simd/maxpool.rs new file mode 100644 index 0000000..804f5e9 --- /dev/null +++ b/crates/burn-ndarray/src/ops/simd/maxpool.rs @@ -0,0 +1,395 @@ +use core::{marker::PhantomData, mem::transmute}; + +use crate::{SharedArray, iter_range_par, run_par, sharing::UnsafeSharedRef}; + +use burn_backend::{BoolStore, DType, Element, quantization::QuantValue}; +use macerator::{Simd, VOrd}; +use ndarray::{Array4, s}; +use nhwc::max_pool2d_nhwc; + +use super::{MinMax, should_use_simd}; + +#[macerator::with_simd] +fn is_accelerated_impl(_x: PhantomData) -> bool { + ::is_min_max_accelerated::() +} + +fn is_accelerated() -> bool { + is_accelerated_impl::(PhantomData) +} + +macro_rules! launch_kernel { + ($ty: ty, $func: ident, $x: expr, $($arg: expr),*) => { + match <$ty as Element>::dtype() { + DType::F64 if is_accelerated::() => Ok(cast($func::(cast($x), $($arg),*))), + DType::F32 if is_accelerated::() => Ok(cast($func::(cast($x), $($arg),*))), + DType::I64 if is_accelerated::() => Ok(cast($func::(cast($x), $($arg),*))), + DType::I32 if is_accelerated::() => Ok(cast($func::(cast($x), $($arg),*))), + DType::I16 if is_accelerated::() => Ok(cast($func::(cast($x), $($arg),*))), + DType::I8 if is_accelerated::() => Ok(cast($func::(cast($x), $($arg),*))), + DType::U64 if is_accelerated::() => Ok(cast($func::(cast($x), $($arg),*))), + DType::U32 if is_accelerated::() => Ok(cast($func::(cast($x), $($arg),*))), + DType::U16 if is_accelerated::() => Ok(cast($func::(cast($x), $($arg),*))), + DType::U8 if is_accelerated::() => Ok(cast($func::(cast($x), $($arg),*))), + DType::Bool(BoolStore::Native) if is_accelerated::() => Ok(cast($func::(cast($x), $($arg),*))), + DType::QFloat(scheme) => match scheme.value { + QuantValue::Q8F | QuantValue::Q8S if is_accelerated::() => Ok(cast($func::(cast($x), $($arg),*))), + _ => Err($x) + }, + _ => Err($x), + } + }; +} + +pub(crate) fn try_max_pool2d_simd( + x: SharedArray, + ksize: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], +) -> Result, SharedArray> { + let [_, c, _, _] = x.shape().try_into().unwrap(); + if !should_use_simd(c) || x.strides()[1] != 1 { + return Err(x); + } + + launch_kernel!(E, max_pool2d_nhwc, x, ksize, stride, padding, dilation) +} + +fn cast(tensor: SharedArray) -> SharedArray { + unsafe { transmute::, SharedArray>(tensor) } +} + +mod nhwc { + use burn_backend::ElementOrdered; + use itertools::Itertools; + use macerator::{Simd, vload_unaligned, vstore_unaligned}; + use ndarray::{ArrayView3, ArrayViewMut3, Ix4}; + use seq_macro::seq; + + use crate::ops::simd::lanes; + + use super::*; + + // Until you can use associated constants as array size, we need to hardcode this. + // The most common config (x86-v3) has 16 registers, so use half of them for accumulators. + const BLOCK_REGISTERS: usize = 8; + + pub(crate) fn max_pool2d_nhwc( + x: SharedArray, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ) -> SharedArray { + let [kernel_height, kernel_width] = kernel_size; + let [pad_h, pad_w] = padding; + let [stride_height, stride_width] = stride; + let [dilation_height, dilation_width] = dilation; + let [batch_size, channels, x_height, x_width] = x.shape().try_into().unwrap(); + let lanes = lanes::(); + + let ch_block = lanes * BLOCK_REGISTERS; + + let out_height = ((x_height + 2 * pad_h - dilation_height * (kernel_height - 1) - 1) + / stride_height) + + 1; + let out_width = + ((x_width + 2 * pad_w - dilation_width * (kernel_width - 1) - 1) / stride_width) + 1; + + let mut output = unsafe { + Array4::::uninit((batch_size, out_height, out_width, channels)).assume_init() + }; + let unsafe_shared_out = UnsafeSharedRef::new(&mut output); + + let x = x.into_dimensionality::().unwrap(); + let x = x.view(); + let x = x.permuted_axes([0, 2, 3, 1]); + + // Floor division ensures `blocks * lanes * blocking factor` is always `<= out_channels`. + // An exclusive loop will always have `lanes * blocking factor` elements in bounds. + let blocks = channels / ch_block; + let blocks_end = blocks * ch_block; + // Floor division means simd_end is always divisible by `lanes` and `<= out_channels`. An + // exclusive loop will always have `lanes` elements in bounds. + let simd_end = channels / lanes * lanes; + let simd_unblocked = (simd_end - blocks_end) / lanes; + let remainder = channels - simd_end; + + run_par!(|| { + // SAFETY: Loop ranges are non-overlapping, so the unsafe shared reference is safe. + iter_range_par!(0, batch_size * blocks).for_each(|k| unsafe { + let block = k % blocks; + let b = k / blocks; + + let output = unsafe_shared_out.get(); + let x = x.slice(s![b, .., .., ..]); + let out = output.slice_mut(s![b, .., .., ..]); + loop_blocked(x, out, kernel_size, stride, padding, dilation, block); + }); + // SAFETY: See `loop_unblocked` + iter_range_par!(0, batch_size * simd_unblocked).for_each(|k| unsafe { + let ch = (k % simd_unblocked) * lanes + blocks_end; + let b = k / simd_unblocked; + + let output = unsafe_shared_out.get(); + let x = x.slice(s![b, .., .., ..]); + let out = output.slice_mut(s![b, .., .., ..]); + loop_unblocked(x, out, kernel_size, stride, padding, dilation, ch); + }); + // SAFETY: Loop ranges are non-overlapping, so the unsafe shared reference is safe. + iter_range_par!(0, batch_size * remainder).for_each(|k| unsafe { + let ch = (k % remainder) + simd_end; + let b = k / remainder; + + let output = unsafe_shared_out.get(); + let x = x.slice(s![b, .., .., ..]); + let out = output.slice_mut(s![b, .., .., ..]); + loop_scalar(x, out, kernel_size, stride, padding, dilation, ch); + }); + }); + + output = output.permuted_axes([0, 3, 1, 2]); + + output.into_dyn().into_shared() + } + + /// Execute the blocked (unrolled) portion of the pool. + #[allow( + clippy::too_many_arguments, + clippy::erasing_op, + clippy::identity_op, + unused_mut + )] + #[inline(always)] + #[macerator::with_simd] + fn loop_blocked<'a, S: Simd, E: ElementOrdered + VOrd + MinMax>( + x: ArrayView3<'a, E>, + mut out: ArrayViewMut3<'a, E>, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + block: usize, + ) where + 'a: 'a, + { + let [kernel_height, kernel_width] = kernel_size; + let [pad_h, pad_w] = padding; + let [stride_height, stride_width] = stride; + let [dilation_height, dilation_width] = dilation; + + let (x_height, x_width, _) = x.dim(); + let (out_height, out_width, _) = out.dim(); + let lanes = E::lanes::(); + let ch_block = lanes * BLOCK_REGISTERS; + + let min = E::MIN.splat::(); + // If outside padding area, kernels are guaranteed to be in bounds + for oh in pad_h..out_height.saturating_sub(pad_h) { + for ow in pad_w..out_width.saturating_sub(pad_w) { + seq!(N in 0..8 { + let mut acc~N = min; + }); + let ch = block * ch_block; + let ch_end = ch + ch_block; + let mut out = out.slice_mut(s![oh, ow, ch..ch_end]); + + for kh in 0..kernel_height { + let ih = oh * stride_height + kh * dilation_height - pad_h; + + for kw in 0..kernel_width { + let iw = ow * stride_width + kw * dilation_width - pad_w; + let x = x.slice(s![ih, iw, ch..ch_end]); + + seq!(N in 0..8 { + // SAFETY: + // Load a full vector from x[N * lanes]. This is bounds checked by the + // slice above. + acc~N = acc~N.max(unsafe { vload_unaligned(&x[N * lanes]) }); + }); + } + } + + seq!(N in 0..8 { + // SAFETY: + // Store a full vector to out[N * lanes]. This is bounds checked by the + // slice above. + unsafe { vstore_unaligned(&mut out[N * lanes], acc~N) }; + }); + } + } + + // Border pixels need bounds checks + if (pad_h, pad_w) != (0, 0) { + let v_borders = (0..pad_h) + .chain(out_height.saturating_sub(pad_h)..out_height) + .cartesian_product(0..out_width); + let h_borders = (0..out_height) + .cartesian_product((0..pad_w).chain(out_width.saturating_sub(pad_w)..out_width)); + + for (oh, ow) in v_borders.chain(h_borders) { + seq!(N in 0..8 { + let mut acc~N = min; + }); + let ch = block * ch_block; + let ch_end = ch + ch_block; + let mut out = out.slice_mut(s![oh, ow, ch..ch_end]); + + for kh in 0..kernel_height { + let ih = oh * stride_height + kh * dilation_height; + if ih < pad_h || ih >= x_height + pad_h { + continue; + } + let ih = ih - pad_h; + + for kw in 0..kernel_width { + let iw = ow * stride_width + kw * dilation_width; + if iw < pad_w || iw >= x_width + pad_w { + continue; + } + let iw = iw - pad_w; + + let x = x.slice(s![ih, iw, ch..ch_end]); + + seq!(N in 0..8 { + // SAFETY: + // Load a full vector from x[N * lanes]. This is bounds checked by the + // slice above. + acc~N = acc~N.max(unsafe { vload_unaligned(&x[N * lanes]) }); + }); + } + } + + seq!(N in 0..8 { + // SAFETY: + // Store a full vector to out[N * lanes]. This is bounds checked by the + // slice above. + unsafe { vstore_unaligned(&mut out[N * lanes], acc~N) }; + }); + } + } + } + + /// Execute the unblocked (not unrolled) portion of the pool. + /// + /// SAFETY: Safe as long as `ch + simd_lanes <= out_channels`. + #[allow(clippy::too_many_arguments, unused_mut)] + #[inline(always)] + #[macerator::with_simd] + unsafe fn loop_unblocked<'a, S: Simd, E: ElementOrdered + VOrd + MinMax>( + x: ArrayView3<'a, E>, + mut out: ArrayViewMut3<'a, E>, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ch: usize, + ) where + 'a: 'a, + { + let [kernel_height, kernel_width] = kernel_size; + let [pad_h, pad_w] = padding; + let [stride_height, stride_width] = stride; + let [dilation_height, dilation_width] = dilation; + + let (x_height, x_width, _) = x.dim(); + let (out_height, out_width, _) = out.dim(); + + for oh in pad_h..out_height.saturating_sub(pad_h) { + for ow in pad_w..out_width.saturating_sub(pad_w) { + let mut acc = E::MIN.splat::(); + let out = &mut out[[oh, ow, ch]]; + + for kh in 0..kernel_height { + let ih = oh * stride_height + kh * dilation_height - pad_h; + + for kw in 0..kernel_width { + let iw = ow * stride_width + kw * dilation_width - pad_w; + // Load a full vector from `x`. In bounds as long as `out_channels >= ch + lanes` + acc = acc.max(unsafe { vload_unaligned(&x[[ih, iw, ch]]) }); + } + } + // Store a full vector to `out`. In bounds as long as `out_channels >= ch + lanes`. + unsafe { vstore_unaligned(out, acc) }; + } + } + + // Border pixels need bounds checks + if (pad_h, pad_w) != (0, 0) { + let v_borders = (0..pad_h) + .chain(out_height.saturating_sub(pad_h)..out_height) + .cartesian_product(0..out_width); + let h_borders = (0..out_height) + .cartesian_product((0..pad_w).chain(out_width.saturating_sub(pad_w)..out_width)); + + for (oh, ow) in v_borders.chain(h_borders) { + let mut acc = E::MIN.splat::(); + let out = &mut out[[oh, ow, ch]]; + + for kh in 0..kernel_height { + let ih = oh * stride_height + kh * dilation_height; + if ih < pad_h || ih >= x_height + pad_h { + continue; + } + let ih = ih - pad_h; + + for kw in 0..kernel_width { + let iw = ow * stride_width + kw * dilation_width; + if iw < pad_w || iw >= x_width + pad_w { + continue; + } + let iw = iw - pad_w; + // Load a full vector from `x`. In bounds as long as `out_channels >= ch + lanes` + acc = acc.max(unsafe { vload_unaligned(&x[[ih, iw, ch]]) }); + } + } + // Store a full vector to `out`. In bounds as long as `out_channels >= ch + lanes`. + unsafe { vstore_unaligned(out, acc) }; + } + } + } + + fn loop_scalar( + x: ArrayView3<'_, E>, + mut out: ArrayViewMut3<'_, E>, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ch: usize, + ) { + let [kernel_height, kernel_width] = kernel_size; + let [pad_h, pad_w] = padding; + let [stride_height, stride_width] = stride; + let [dilation_height, dilation_width] = dilation; + + let (x_height, x_width, _) = x.dim(); + let (out_height, out_width, _) = out.dim(); + + for oh in 0..out_height { + for ow in 0..out_width { + let mut acc = E::MIN; + + for kh in 0..kernel_height { + let ih = oh * stride_height + kh * dilation_height; + if ih < pad_h || ih >= x_height + pad_h { + continue; + } + let ih = ih - pad_h; + + for kw in 0..kernel_width { + let iw = ow * stride_width + kw * dilation_width; + if iw < pad_w || iw >= x_width + pad_w { + continue; + } + let iw = iw - pad_w; + acc = acc.max(x[[ih, iw, ch]]); + } + } + + out[[oh, ow, ch]] = acc; + } + } + } +} diff --git a/crates/burn-ndarray/src/ops/simd/mod.rs b/crates/burn-ndarray/src/ops/simd/mod.rs new file mode 100644 index 0000000..2032f30 --- /dev/null +++ b/crates/burn-ndarray/src/ops/simd/mod.rs @@ -0,0 +1,10 @@ +pub(crate) mod avgpool; +mod base; +pub(crate) mod binary; +pub(crate) mod binary_elemwise; +pub(crate) mod cmp; +pub(crate) mod conv; +pub(crate) mod maxpool; +pub(crate) mod unary; + +pub use base::*; diff --git a/crates/burn-ndarray/src/ops/simd/unary.rs b/crates/burn-ndarray/src/ops/simd/unary.rs new file mode 100644 index 0000000..68d2626 --- /dev/null +++ b/crates/burn-ndarray/src/ops/simd/unary.rs @@ -0,0 +1,234 @@ +use core::marker::PhantomData; + +use bytemuck::cast; +use macerator::{ + Scalar, Simd, VAbs, VBitNot, VRecip, Vector, vload, vload_unaligned, vstore, vstore_unaligned, +}; +use ndarray::ArrayD; +use num_traits::Signed; +use seq_macro::seq; + +use crate::{NdArrayElement, SharedArray}; + +use super::should_use_simd; + +pub trait SimdUnop { + fn apply_vec(input: Vector) -> Vector; + fn apply(input: T) -> Out; + fn is_accelerated() -> bool; +} + +pub struct RecipVec; + +impl SimdUnop for RecipVec { + fn apply_vec(input: Vector) -> Vector { + input.recip() + } + + fn apply(input: f32) -> f32 { + input.recip() + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +pub struct VecAbs; + +impl SimdUnop for VecAbs { + fn apply_vec(input: Vector) -> Vector { + input.abs() + } + + fn apply(input: T) -> T { + input.abs() + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +pub struct VecBitNot; + +impl SimdUnop for VecBitNot { + fn apply_vec(input: Vector) -> Vector { + !input + } + + fn apply(input: T) -> T { + input.not() + } + + fn is_accelerated() -> bool { + ::is_accelerated::() + } +} + +#[macerator::with_simd] +fn is_accelerated>( + _x: PhantomData<(T, Out, Op)>, +) -> bool { + Op::is_accelerated::() +} + +pub fn try_unary_simd< + E: NdArrayElement, + EOut: NdArrayElement, + T: NdArrayElement + Scalar, + Out: NdArrayElement + Scalar, + Op: SimdUnop, +>( + input: SharedArray, +) -> Result, SharedArray> { + if !should_use_simd(input.len()) + || input.as_slice_memory_order().is_none() + || !is_accelerated::(PhantomData) + { + return Err(input); + } + // Used to assert traits based on the dynamic `DType`. + let input = unsafe { core::mem::transmute::, SharedArray>(input) }; + let out = if size_of::() == size_of::() + && align_of::() >= align_of::() + && input.is_unique() + { + unsafe { unary_scalar_simd_inplace::(input) } + } else { + unary_scalar_simd_owned::(input) + }; + // Used to assert traits based on the dynamic `DType`. + let out = unsafe { core::mem::transmute::, SharedArray>(out) }; + Ok(out) +} + +/// Execute operation in line. +/// SAFETY: +/// Must ensure `size_of:: == size_of::` and `align_of:: >= align_of::`. +unsafe fn unary_scalar_simd_inplace< + T: NdArrayElement + Scalar, + Out: NdArrayElement + Scalar, + Op: SimdUnop, +>( + input: SharedArray, +) -> SharedArray { + let mut buffer = input.into_owned(); + let slice = buffer.as_slice_memory_order_mut().unwrap(); + // This is only called when in and out have the same size, so it's safe + unsafe { unary_slice_inplace::(slice, PhantomData) }; + // Buffer has the same elem size and is filled with the operation output, so this is safe + let out = unsafe { core::mem::transmute::, ArrayD>(buffer) }; + out.into_shared() +} + +fn unary_scalar_simd_owned< + T: NdArrayElement + Scalar, + Out: NdArrayElement + Scalar, + Op: SimdUnop, +>( + input: SharedArray, +) -> SharedArray { + let mut out = unsafe { ArrayD::uninit(input.shape()).assume_init() }; + let input = input.as_slice_memory_order().unwrap(); + let out_slice = out.as_slice_memory_order_mut().unwrap(); + unary_slice::(input, out_slice, PhantomData); + out.into_shared() +} + +#[allow(clippy::erasing_op, clippy::identity_op)] +#[macerator::with_simd] +fn unary_slice< + 'a, + S: Simd, + T: NdArrayElement + Scalar, + Out: NdArrayElement + Scalar, + Op: SimdUnop, +>( + input: &'a [T], + out: &'a mut [Out], + _op: PhantomData, +) where + 'a: 'a, +{ + let lanes = T::lanes::(); + let mut chunks_input = input.chunks_exact(8 * lanes); + let mut chunks_out = out.chunks_exact_mut(8 * lanes); + while let Some((input, out)) = chunks_input.next().zip(chunks_out.next()) { + seq!(N in 0..8 { + // Load one full vector from `input`. + // SAFETY: Guaranteed to be in bounds because `len == 8 * lanes` + let s~N = unsafe { vload_unaligned(&input[N * lanes]) }; + let s~N = Op::apply_vec::(s~N); + // Store one full vector to `out`. + // SAFETY: Guaranteed to be in bounds because `len == 8 * lanes` + unsafe { vstore_unaligned(&mut out[N * lanes], s~N) }; + }); + } + let mut chunks_input = chunks_input.remainder().chunks_exact(lanes); + let mut chunks_out = chunks_out.into_remainder().chunks_exact_mut(lanes); + while let Some((input, out)) = chunks_input.next().zip(chunks_out.next()) { + // Load one full vector from `input`. + // SAFETY: Guaranteed to be in bounds because `len == lanes` + let s0 = unsafe { vload_unaligned(input.as_ptr()) }; + let s0 = Op::apply_vec::(s0); + // Store one full vector to `out`. + // SAFETY: Guaranteed to be in bounds because `len == lanes` + unsafe { vstore_unaligned(out.as_mut_ptr(), s0) }; + } + + for (input, out) in chunks_input + .remainder() + .iter() + .zip(chunks_out.into_remainder()) + { + *out = Op::apply(*input) + } +} + +/// Execute operation in line. +/// SAFETY: +/// Must ensure `size_of:: == size_of::` and `align_of:: >= align_of::`. +#[macerator::with_simd] +unsafe fn unary_slice_inplace< + 'a, + S: Simd, + T: NdArrayElement + Scalar, + Out: NdArrayElement + Scalar, + Op: SimdUnop, +>( + buf: &'a mut [T], + _op: PhantomData<(Out, Op)>, +) where + 'a: 'a, +{ + let (head, main, tail) = unsafe { buf.align_to_mut::>() }; + for elem in head.iter_mut().chain(tail) { + *elem = cast(Op::apply(*elem)); + } + let mut chunks = main.chunks_exact_mut(8); + for elem in chunks.by_ref() { + seq!(N in 0..8 { + // Load a full vector from the aligned portion of the buffer. + // SAFETY: `align_to_mut` guarantees we're aligned to `T::Vector`'s size, and there is + // always a full vector in bounds. + let s~N = unsafe { vload(&elem[N] as *const _ as *const T) }; + let s~N = Op::apply_vec::(s~N); + // Store a full vector at the same position as the input. Cast is safe because `Out` is + // size and align compatible + unsafe { vstore(&mut elem[N] as *mut _ as *mut Out, s~N) }; + }); + } + + for elem in chunks.into_remainder() { + // Load a full vector from the aligned portion of the buffer. + // SAFETY: `align_to_mut` guarantees we're aligned to `T::Vector`'s size, and there is + // always a full vector in bounds. + let s0 = unsafe { vload(elem as *const _ as *const T) }; + + let s0 = Op::apply_vec::(s0); + // Store a full vector at the same position as the input. Cast is safe because `Out` is + // size and align compatible + unsafe { vstore(elem as *mut _ as *mut Out, s0) }; + } +} diff --git a/crates/burn-ndarray/src/ops/tensor.rs b/crates/burn-ndarray/src/ops/tensor.rs new file mode 100644 index 0000000..ea00483 --- /dev/null +++ b/crates/burn-ndarray/src/ops/tensor.rs @@ -0,0 +1,769 @@ +// Language +use alloc::vec::Vec; +use burn_backend::backend::ExecutionError; +use burn_backend::ops::GridSampleOptions; +use burn_backend::tensor::FloatTensor; +use burn_backend::{TensorMetadata, element::cast::ToElement}; +use burn_std::{BoolDType, IntDType}; + +// Current crate +use super::{ + NdArrayMathOps, NdArrayOps, + matmul::{cross, matmul}, +}; +use crate::{ + NdArray, cast_to_dtype, cat_with_dtype, execute_with_int_dtype, tensor::NdArrayTensor, +}; +use crate::{NdArrayDevice, SEED, execute_with_float_out_dtype, execute_with_int_out_dtype, slice}; +use crate::{SharedArray, element::ExpElement}; +use crate::{execute_with_float_dtype, ops::grid_sample::grid_sample_2d}; + +// Workspace crates +use crate::rand::get_seeded_rng; +use burn_backend::{Distribution, FloatDType, Scalar}; +use burn_backend::{ElementConversion, Shape, TensorData, ops::FloatTensorOps}; + +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float; + +use libm::erf; + +#[cfg(feature = "std")] +#[allow(dead_code)] +fn round_ties_even_wrapper(x: f64) -> f64 { + x.round_ties_even() +} + +#[cfg(not(feature = "std"))] +#[allow(dead_code)] +fn round_ties_even_wrapper(x: f64) -> f64 { + if (x - x.floor()) == 0.5 { + (x * 0.5).round() * 2.0 + } else { + x.round() + } +} + +impl FloatTensorOps for NdArray { + fn float_from_data(data: TensorData, _device: &NdArrayDevice) -> FloatTensor { + NdArrayTensor::from_data(data) + } + + fn float_random( + shape: Shape, + distribution: Distribution, + device: &NdArrayDevice, + dtype: FloatDType, + ) -> FloatTensor { + let mut seed = SEED.lock().unwrap(); + let mut rng = seed.take().unwrap_or_else(get_seeded_rng); + let tensor = execute_with_float_out_dtype!( + dtype, + E, + Self::float_from_data( + TensorData::random::(shape, distribution, &mut rng), + device, + ) + ); + + *seed = Some(rng); + tensor + } + + async fn float_into_data(tensor: FloatTensor) -> Result { + Ok(tensor.into_data()) + } + + fn float_to_device(tensor: FloatTensor, _device: &NdArrayDevice) -> FloatTensor { + tensor + } + + fn float_empty(shape: Shape, device: &NdArrayDevice, dtype: FloatDType) -> FloatTensor { + Self::float_zeros(shape, device, dtype) + } + + fn float_add(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + execute_with_float_dtype!((lhs, rhs), NdArrayMathOps::add) + } + + fn float_add_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + execute_with_float_dtype!(lhs, FloatElem, |array: SharedArray| { + NdArrayMathOps::add_scalar(array, rhs.elem()) + }) + } + + fn float_sub(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + execute_with_float_dtype!((lhs, rhs), NdArrayMathOps::sub) + } + + fn float_sub_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + execute_with_float_dtype!(lhs, FloatElem, |array: SharedArray| { + NdArrayMathOps::sub_scalar(array, rhs.elem()) + }) + } + + fn float_mul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + execute_with_float_dtype!((lhs, rhs), NdArrayMathOps::mul) + } + + fn float_mul_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + execute_with_float_dtype!(lhs, FloatElem, |array: SharedArray| { + NdArrayMathOps::mul_scalar(array, rhs.elem()) + }) + } + + fn float_div(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + execute_with_float_dtype!((lhs, rhs), NdArrayMathOps::div) + } + + fn float_div_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + execute_with_float_dtype!(lhs, FloatElem, |array: SharedArray| { + NdArrayMathOps::div_scalar(array, rhs.elem()) + }) + } + + fn float_remainder(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + execute_with_float_dtype!((lhs, rhs), NdArrayMathOps::remainder) + } + + fn float_remainder_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + execute_with_float_dtype!(lhs, FloatElem, |array: SharedArray| { + NdArrayMathOps::remainder_scalar(array, rhs.elem()) + }) + } + + fn float_matmul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + execute_with_float_dtype!((lhs, rhs), matmul) + } + + fn float_cross( + lhs: FloatTensor, + rhs: FloatTensor, + dim: usize, + ) -> FloatTensor { + execute_with_float_dtype!((lhs, rhs), |lhs, rhs| cross(lhs, rhs, dim)) + } + + fn float_recip(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::recip(array) + }) + } + + fn float_swap_dims(tensor: FloatTensor, dim1: usize, dim2: usize) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayOps::swap_dims(array, dim1, dim2) + }) + } + + fn float_reshape(tensor: FloatTensor, shape: Shape) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayOps::reshape(array, shape) + }) + } + + fn float_gather( + dim: usize, + tensor: FloatTensor, + indices: NdArrayTensor, + ) -> FloatTensor { + execute_with_int_dtype!( + indices, + IntElem, + |idx_array: SharedArray| -> NdArrayTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayOps::gather(dim, array, idx_array) + }) + } + ) + } + + fn float_scatter_add( + dim: usize, + tensor: FloatTensor, + indices: NdArrayTensor, + value: FloatTensor, + ) -> FloatTensor { + execute_with_int_dtype!( + indices, + IntElem, + |idx_array: SharedArray| -> NdArrayTensor { + execute_with_float_dtype!((tensor, value), |tensor, value| NdArrayOps::scatter( + dim, tensor, idx_array, value + )) + } + ) + } + + fn float_scatter_nd( + data: FloatTensor, + indices: NdArrayTensor, + values: FloatTensor, + reduction: burn_backend::tensor::IndexingUpdateOp, + ) -> FloatTensor { + execute_with_int_dtype!( + indices, + IntElem, + |idx_array: SharedArray| -> NdArrayTensor { + execute_with_float_dtype!((data, values), |data, values| NdArrayOps::scatter_nd( + data, idx_array, values, reduction + )) + } + ) + } + + fn float_gather_nd(data: FloatTensor, indices: NdArrayTensor) -> FloatTensor { + execute_with_int_dtype!( + indices, + IntElem, + |idx_array: SharedArray| -> NdArrayTensor { + execute_with_float_dtype!(data, FloatElem, |array: SharedArray| { + NdArrayOps::gather_nd(array, idx_array) + }) + } + ) + } + + fn float_select( + tensor: FloatTensor, + dim: usize, + indices: NdArrayTensor, + ) -> FloatTensor { + execute_with_int_dtype!( + indices, + IntElem, + |idx_array: SharedArray| -> NdArrayTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::select(array, dim, idx_array) + }) + } + ) + } + + fn float_select_add( + tensor: FloatTensor, + dim: usize, + indices: NdArrayTensor, + value: FloatTensor, + ) -> FloatTensor { + execute_with_int_dtype!( + indices, + IntElem, + |idx_array: SharedArray| -> NdArrayTensor { + execute_with_float_dtype!((tensor, value), |tensor, value| { + NdArrayMathOps::select_assign(tensor, dim, idx_array, value) + }) + } + ) + } + + fn float_slice(tensor: FloatTensor, slices: &[burn_backend::Slice]) -> FloatTensor { + slice!(tensor, slices) + } + + fn float_slice_assign( + tensor: FloatTensor, + slices: &[burn_backend::Slice], + value: FloatTensor, + ) -> FloatTensor { + execute_with_float_dtype!((tensor, value), |tensor, value| { + NdArrayOps::slice_assign(tensor, slices, value) + }) + } + + fn float_mask_where( + tensor: FloatTensor, + mask: NdArrayTensor, + value: FloatTensor, + ) -> FloatTensor { + execute_with_float_dtype!((tensor, value), |tensor, value| { + NdArrayOps::mask_where(tensor, mask.bool(), value) + }) + } + + fn float_mask_fill( + tensor: FloatTensor, + mask: NdArrayTensor, + value: Scalar, + ) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayOps::mask_fill(array, mask.bool(), value.elem()) + }) + } + + fn float_equal( + lhs: FloatTensor, + rhs: FloatTensor, + _out_dtype: BoolDType, + ) -> NdArrayTensor { + execute_with_float_dtype!((lhs, rhs), |lhs, rhs| { NdArrayMathOps::equal(lhs, rhs) }) + } + + fn float_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> NdArrayTensor { + execute_with_float_dtype!(lhs, FloatElem, |array: SharedArray| { + NdArrayMathOps::equal_elem(array, rhs.elem()) + }) + } + + fn float_greater( + lhs: FloatTensor, + rhs: FloatTensor, + _out_dtype: BoolDType, + ) -> NdArrayTensor { + execute_with_float_dtype!((lhs, rhs), |lhs, rhs| { NdArrayMathOps::greater(lhs, rhs) }) + } + + fn float_greater_elem( + lhs: FloatTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> NdArrayTensor { + execute_with_float_dtype!(lhs, FloatElem, |array: SharedArray| { + NdArrayMathOps::greater_elem(array, rhs.elem()) + }) + } + + fn float_greater_equal( + lhs: FloatTensor, + rhs: FloatTensor, + _out_dtype: BoolDType, + ) -> NdArrayTensor { + execute_with_float_dtype!((lhs, rhs), |lhs, rhs| { + NdArrayMathOps::greater_equal(lhs, rhs) + }) + } + + fn float_greater_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> NdArrayTensor { + execute_with_float_dtype!(lhs, FloatElem, |array: SharedArray| { + NdArrayMathOps::greater_equal_elem(array, rhs.elem()) + }) + } + + fn float_lower( + lhs: FloatTensor, + rhs: FloatTensor, + _out_dtype: BoolDType, + ) -> NdArrayTensor { + execute_with_float_dtype!((lhs, rhs), |lhs, rhs| { NdArrayMathOps::lower(lhs, rhs) }) + } + + fn float_lower_elem( + lhs: FloatTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> NdArrayTensor { + execute_with_float_dtype!(lhs, FloatElem, |array: SharedArray| { + NdArrayMathOps::lower_elem(array, rhs.elem()) + }) + } + + fn float_lower_equal( + lhs: FloatTensor, + rhs: FloatTensor, + _out_dtype: BoolDType, + ) -> NdArrayTensor { + execute_with_float_dtype!((lhs, rhs), |lhs, rhs| { + NdArrayMathOps::lower_equal(lhs, rhs) + }) + } + + fn float_lower_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + _out_dtype: BoolDType, + ) -> NdArrayTensor { + execute_with_float_dtype!(lhs, FloatElem, |array: SharedArray| { + NdArrayMathOps::lower_equal_elem(array, rhs.elem()) + }) + } + + fn float_detach(tensor: FloatTensor) -> FloatTensor { + tensor + } + + fn float_mean(tensor: FloatTensor) -> FloatTensor { + // Use view() for zero-copy on borrowed storage + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::mean_view(array.view()) + }) + } + + fn float_sum(tensor: FloatTensor) -> FloatTensor { + // Use view() for zero-copy on borrowed storage + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::sum_view(array.view()) + }) + } + + fn float_mean_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::mean_dim(array, dim) + }) + } + + fn float_cumsum(tensor: FloatTensor, dim: usize) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::cumsum(array, dim) + }) + } + + fn float_cumprod(tensor: FloatTensor, dim: usize) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::cumprod(array, dim) + }) + } + + fn float_cummin(tensor: FloatTensor, dim: usize) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::cummin(array, dim) + }) + } + + fn float_cummax(tensor: FloatTensor, dim: usize) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::cummax(array, dim) + }) + } + + fn float_sum_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::sum_dim(array, dim) + }) + } + + fn float_argmax(tensor: FloatTensor, dim: usize, out_dtype: IntDType) -> NdArrayTensor { + // Use view() for zero-copy on borrowed storage + execute_with_int_out_dtype!(out_dtype, I, { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::argmax_view::(array.view(), dim) + }) + }) + } + + fn float_argtopk( + _tensor: FloatTensor, + _dim: usize, + _k: usize, + _out_dtype: IntDType, + ) -> NdArrayTensor { + unimplemented!("float_argtopk not implemented for ndarray") + } + + fn float_argmin(tensor: FloatTensor, dim: usize, out_dtype: IntDType) -> NdArrayTensor { + // Use view() for zero-copy on borrowed storage + execute_with_int_out_dtype!(out_dtype, I, { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::argmin_view::(array.view(), dim) + }) + }) + } + + fn float_exp(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array.mapv_into(|a: FloatElem| a.exp_elem()).into_shared() + }) + } + + fn float_log(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array.mapv_into(|a: FloatElem| a.log_elem()).into_shared() + }) + } + + fn float_prod(tensor: FloatTensor) -> FloatTensor { + // Use view() for zero-copy on borrowed storage + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::prod_view(array.view()) + }) + } + + fn float_prod_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::prod_dim(array, dim) + }) + } + + fn float_max(tensor: FloatTensor) -> FloatTensor { + // Use view() for zero-copy on borrowed storage + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::max_view(array.view()) + }) + } + + fn float_min(tensor: FloatTensor) -> FloatTensor { + // Use view() for zero-copy on borrowed storage + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::min_view(array.view()) + }) + } + + fn float_log1p(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array.mapv_into(|a: FloatElem| a.log1p_elem()).into_shared() + }) + } + + fn float_powf_scalar_impl(tensor: FloatTensor, value: Scalar) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| a.powf_elem(value.elem())) + .into_shared() + }) + } + + fn float_sqrt(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array.mapv_into(|a: FloatElem| a.sqrt_elem()).into_shared() + }) + } + + fn float_abs(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::abs(array) + }) + } + + fn float_cos(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| (a.to_f64()).cos().elem()) + .into_shared() + }) + } + + fn float_cosh(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| (a.to_f64()).cosh().elem()) + .into_shared() + }) + } + + fn float_sin(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| (a.to_f64()).sin().elem()) + .into_shared() + }) + } + + fn float_sinh(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| (a.to_f64()).sinh().elem()) + .into_shared() + }) + } + + fn float_tan(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| (a.to_f64()).tan().elem()) + .into_shared() + }) + } + + fn float_tanh(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| (a.to_f64()).tanh().elem()) + .into_shared() + }) + } + + fn float_acos(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| (a.to_f64()).acos().elem()) + .into_shared() + }) + } + + fn float_acosh(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| (a.to_f64()).acosh().elem()) + .into_shared() + }) + } + + fn float_asin(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| (a.to_f64()).asin().elem()) + .into_shared() + }) + } + + fn float_asinh(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| (a.to_f64()).asinh().elem()) + .into_shared() + }) + } + + fn float_atan(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| (a.to_f64()).atan().elem()) + .into_shared() + }) + } + + fn float_atanh(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| (a.to_f64()).atanh().elem()) + .into_shared() + }) + } + + fn float_atan2(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + execute_with_float_dtype!((lhs, rhs), FloatElem, |lhs, rhs| { + NdArrayMathOps::elementwise_op(lhs, rhs, |a: &FloatElem, b: &FloatElem| a.atan2(*b)) + }) + } + + fn float_round(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| round_ties_even_wrapper(a.to_f64()).elem()) + .into_shared() + }) + } + + fn float_floor(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| (a.to_f64()).floor().elem()) + .into_shared() + }) + } + + fn float_ceil(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| (a.to_f64()).ceil().elem()) + .into_shared() + }) + } + + fn float_trunc(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| (a.to_f64()).trunc().elem()) + .into_shared() + }) + } + + fn float_erf(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array + .mapv_into(|a: FloatElem| erf(a.to_f64()).elem()) + .into_shared() + }) + } + + fn float_cat(tensors: Vec>, dim: usize) -> FloatTensor { + cat_with_dtype!(tensors, dim, [F64, F32]) + } + + fn float_clamp_min(tensor: FloatTensor, min: Scalar) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::clamp_min(array, min.elem()) + }) + } + + fn float_clamp_max(tensor: FloatTensor, max: Scalar) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::clamp_max(array, max.elem()) + }) + } + + fn float_clamp(tensor: FloatTensor, min: Scalar, max: Scalar) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::clamp(array, min.elem(), max.elem()) + }) + } + + fn float_into_int(tensor: FloatTensor, out_dtype: IntDType) -> NdArrayTensor { + execute_with_int_out_dtype!(out_dtype, I, { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + array.mapv(|a: FloatElem| a.elem::()).into_shared() + }) + }) + } + + fn float_powf(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + execute_with_float_dtype!((lhs, rhs), FloatElem, |lhs, rhs| { + NdArrayMathOps::elementwise_op(lhs, rhs, |a: &FloatElem, b: &FloatElem| a.powf(*b)) + }) + } + + fn float_permute(tensor: FloatTensor, axes: &[usize]) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayOps::permute(array, axes) + }) + } + + fn float_flip(tensor: FloatTensor, axes: &[usize]) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayOps::flip(array, axes) + }) + } + + fn float_sign(tensor: FloatTensor) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayMathOps::sign_op(array) + }) + } + + fn float_expand(tensor: FloatTensor, shape: Shape) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayOps::expand(array, shape) + }) + } + + fn float_cast(tensor: FloatTensor, dtype: FloatDType) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + cast_to_dtype(array, dtype.into()) + }) + } + + fn float_grid_sample_2d( + tensor: FloatTensor, + grid: FloatTensor, + options: GridSampleOptions, + ) -> FloatTensor { + execute_with_float_dtype!((tensor, grid), |tensor, grid| grid_sample_2d( + tensor, grid, options + )) + } + + fn float_unfold( + tensor: FloatTensor, + dim: usize, + size: usize, + step: usize, + ) -> FloatTensor { + execute_with_float_dtype!(tensor, FloatElem, |array: SharedArray| { + NdArrayOps::unfold(array, dim, size, step) + }) + } + + fn float_hypot(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + execute_with_float_dtype!((lhs, rhs), FloatElem, |lhs, rhs| { + NdArrayMathOps::elementwise_op(lhs, rhs, |a: &FloatElem, b: &FloatElem| a.hypot(*b)) + }) + } +} diff --git a/crates/burn-ndarray/src/ops/transaction.rs b/crates/burn-ndarray/src/ops/transaction.rs new file mode 100644 index 0000000..c482e8f --- /dev/null +++ b/crates/burn-ndarray/src/ops/transaction.rs @@ -0,0 +1,8 @@ +use crate::NdArray; +use burn_backend::distributed::DistributedOps; +use burn_backend::ops::TransactionOps; + +impl TransactionOps for NdArray {} + +// DistributedOps has default implementations; NdArray does not support collective operations. +impl DistributedOps for NdArray {} diff --git a/crates/burn-ndarray/src/parallel.rs b/crates/burn-ndarray/src/parallel.rs new file mode 100644 index 0000000..a665761 --- /dev/null +++ b/crates/burn-ndarray/src/parallel.rs @@ -0,0 +1,76 @@ +/// Macro for running a function in parallel. +#[cfg(feature = "multi-threads")] +#[macro_export(local_inner_macros)] +macro_rules! run_par { + ( + $func:expr + ) => {{ + use rayon::prelude::*; + + #[allow(clippy::redundant_closure_call)] + rayon::scope(|_| $func()) + }}; +} + +/// Macro for running a function in parallel. +#[cfg(not(feature = "multi-threads"))] +#[macro_export(local_inner_macros)] +macro_rules! run_par { + ( + $func:expr + ) => {{ $func() }}; +} + +/// Macro for iterating in parallel. +#[cfg(not(feature = "multi-threads"))] +#[macro_export(local_inner_macros)] +macro_rules! iter_par { + ( + $iter:expr + ) => {{ $iter }}; +} + +/// Macro for iterating in parallel. +#[cfg(feature = "multi-threads")] +#[macro_export(local_inner_macros)] +macro_rules! iter_par { + ( + $iter:expr + ) => {{ $iter.into_par_iter() }}; +} + +/// Macro for iterating in parallel. +#[cfg(feature = "multi-threads")] +#[macro_export(local_inner_macros)] +macro_rules! iter_slice_par { + ( + $slice:expr + ) => {{ $slice.into_par_iter() }}; +} + +/// Macro for iterating in parallel. +#[cfg(not(feature = "multi-threads"))] +#[macro_export(local_inner_macros)] +macro_rules! iter_slice_par { + ( + $slice:expr + ) => {{ $slice.iter() }}; +} + +/// Macro for iterating over a range in parallel. +#[cfg(feature = "multi-threads")] +#[macro_export(local_inner_macros)] +macro_rules! iter_range_par { + ( + $start:expr, $end:expr + ) => {{ ($start..$end).into_par_iter() }}; +} + +/// Macro for iterating over a range in parallel. +#[cfg(not(feature = "multi-threads"))] +#[macro_export(local_inner_macros)] +macro_rules! iter_range_par { + ( + $start:expr, $end:expr + ) => {{ ($start..$end) }}; +} diff --git a/crates/burn-ndarray/src/rand.rs b/crates/burn-ndarray/src/rand.rs new file mode 100644 index 0000000..94b9bcd --- /dev/null +++ b/crates/burn-ndarray/src/rand.rs @@ -0,0 +1,36 @@ +//! Random number generation utilities for burn-ndarray + +#[cfg(not(feature = "std"))] +use rand::rngs::SmallRng; +#[cfg(feature = "std")] +use rand::rngs::StdRng; + +/// Type alias for the RNG used by burn-ndarray +#[cfg(feature = "std")] +pub type NdArrayRng = StdRng; +#[cfg(not(feature = "std"))] +pub type NdArrayRng = SmallRng; + +#[cfg(not(feature = "std"))] +use rand::SeedableRng; + +/// Get a seeded random number generator +/// +/// For std builds, uses OS entropy. +/// For no_std builds, uses a compile-time random seed. +#[cfg(feature = "std")] +pub fn get_seeded_rng() -> NdArrayRng { + // Use the standard implementation from burn-std + burn_std::rand::get_seeded_rng() +} + +/// Get a seeded random number generator +/// +/// For std builds, uses OS entropy. +/// For no_std builds, uses a compile-time random seed. +#[cfg(not(feature = "std"))] +pub fn get_seeded_rng() -> NdArrayRng { + // Use compile-time random seed for no_std + const SEED: u64 = const_random::const_random!(u64); + SmallRng::seed_from_u64(SEED) +} diff --git a/crates/burn-ndarray/src/sharing.rs b/crates/burn-ndarray/src/sharing.rs new file mode 100644 index 0000000..75d5142 --- /dev/null +++ b/crates/burn-ndarray/src/sharing.rs @@ -0,0 +1,19 @@ +use core::cell::UnsafeCell; + +/// Similar to `SyncUnsafeCell` see [Rust issues](https://github.com/rust-lang/rust/issues/95439). +pub(crate) struct UnsafeSharedRef<'a, T> { + cell: UnsafeCell<&'a mut T>, +} + +unsafe impl Sync for UnsafeSharedRef<'_, T> {} + +impl<'a, T> UnsafeSharedRef<'a, T> { + pub fn new(data: &'a mut T) -> Self { + Self { + cell: UnsafeCell::new(data), + } + } + pub unsafe fn get(&self) -> &'a mut T { + unsafe { core::ptr::read(self.cell.get()) } + } +} diff --git a/crates/burn-ndarray/src/storage.rs b/crates/burn-ndarray/src/storage.rs new file mode 100644 index 0000000..7eeca47 --- /dev/null +++ b/crates/burn-ndarray/src/storage.rs @@ -0,0 +1,506 @@ +//! Copy-on-write storage for zero-copy tensor loading. +//! +//! This module provides `NdArrayStorage`, which enables true zero-copy loading +//! from burnpack files. When data is borrowed from external memory (like mmap'd files +//! or static data), it remains zero-copy until a mutating operation is performed, +//! at which point it's copied (copy-on-write semantics). +//! +//! This integrates with ndarray's existing COW patterns - operations that check +//! `is_unique()` will see borrowed data as non-unique, triggering the allocation path. + +use burn_backend::Element; +use burn_std::{Bytes, Shape}; +use core::mem; +use ndarray::{ArcArray, ArrayView, IxDyn}; + +/// Storage that supports both owned data and borrowed (zero-copy) data. +/// +/// # Copy-on-Write Semantics +/// +/// - **Borrowed**: Data from external source (burnpack, mmap, static). +/// Reports `is_unique() == false` to trigger copy on mutation. +/// - **Owned**: Standard `ArcArray` with built-in COW via Arc refcount. +/// +/// # Example +/// +/// ```ignore +/// // Zero-copy load +/// let storage = NdArrayStorage::from_borrowed(bytes, shape); +/// storage.is_unique(); // false - will copy on mutation +/// +/// // Read operations use view() - zero-copy +/// let view = storage.view(); +/// +/// // Mutation converts to owned +/// let owned = storage.into_owned(); // Copies here +/// ``` +#[derive(Debug)] +pub enum NdArrayStorage { + /// Borrowed from external source (e.g., burnpack zero-copy load). + /// Keeps `Bytes` alive to ensure the referenced memory is valid. + Borrowed { + /// Source bytes - keeps external memory alive via reference counting + bytes: Bytes, + /// Shape of the tensor + shape: Shape, + }, + + /// Standard owned storage with ArcArray COW semantics. + Owned(ArcArray), +} + +impl Clone for NdArrayStorage { + fn clone(&self) -> Self { + match self { + // For borrowed data, clone the Bytes (cheap Arc clone) and shape + Self::Borrowed { bytes, shape } => Self::Borrowed { + bytes: bytes.clone(), + shape: shape.clone(), + }, + // For owned data, clone the ArcArray (cheap Arc clone) + Self::Owned(arr) => Self::Owned(arr.clone()), + } + } +} + +impl NdArrayStorage { + /// Create borrowed storage from external bytes. + /// + /// Returns the bytes and shape back on failure (misaligned or too small), + /// enabling zero-copy even for native allocations by avoiding defensive cloning. + /// + /// # Requirements + /// + /// The caller must ensure that: + /// - The `Bytes` contain valid data for the element type `E` + /// - The data is contiguous in row-major (C) order matching the provided shape + /// + /// These requirements are upheld when loading from `TensorData` (burnpack, etc.) + /// which always stores data contiguously in row-major order. + pub fn from_borrowed(bytes: Bytes, shape: impl Into) -> Result { + let shape = shape.into(); + // Validate alignment + let ptr = bytes.as_ptr(); + if !(ptr as usize).is_multiple_of(mem::align_of::()) { + return Err((bytes, shape)); + } + + // Validate size (using checked arithmetic to prevent overflow) + let num_elements = match shape + .iter() + .try_fold(1usize, |acc, &dim| acc.checked_mul(dim)) + { + Some(n) => n, + None => return Err((bytes, shape)), + }; + let expected_size = match num_elements.checked_mul(mem::size_of::()) { + Some(s) => s, + None => return Err((bytes, shape)), + }; + if bytes.len() < expected_size { + return Err((bytes, shape)); + } + + Ok(Self::Borrowed { bytes, shape }) + } + + /// Create owned storage from an ArcArray. + #[inline] + pub fn from_owned(array: ArcArray) -> Self { + Self::Owned(array) + } + + /// Returns whether this storage is uniquely owned and can be mutated in-place. + /// + /// - **Borrowed**: Always returns `false` to trigger copy-on-write. + /// - **Owned**: Delegates to `ArcArray::is_unique()`. + /// + /// This integrates with existing SIMD code patterns like: + /// ```ignore + /// if tensor.is_unique() { + /// // mutate in place + /// } else { + /// // allocate new + /// } + /// ``` + #[inline] + pub fn is_unique(&self) -> bool { + match self { + Self::Borrowed { .. } => false, // Force copy path + Self::Owned(arr) => arr.is_unique(), + } + } + + /// Get a read-only view of the data. + /// + /// This is zero-copy for both borrowed and owned variants. + #[inline] + pub fn view(&self) -> ArrayView<'_, E, IxDyn> { + match self { + Self::Borrowed { bytes, shape } => { + let ptr = bytes.as_ptr() as *const E; + let dim = IxDyn(shape); + // SAFETY: + // - `bytes` is kept alive for the lifetime of `self` + // - Alignment was validated in `from_borrowed` + // - Size was validated in `from_borrowed` + unsafe { ArrayView::from_shape_ptr(dim, ptr) } + } + Self::Owned(arr) => arr.view(), + } + } + + /// Convert to owned ArcArray. + /// + /// - **Borrowed**: Copies the data into a new ArcArray. + /// - **Owned + unique**: Returns the array without copying. + /// - **Owned + shared**: Clones the data. + pub fn into_owned(self) -> ArcArray { + match self { + Self::Borrowed { bytes, shape } => { + let ptr = bytes.as_ptr() as *const E; + let dim = IxDyn(&shape); + // SAFETY: Same as view() - bytes is valid for this scope + let view = unsafe { ArrayView::from_shape_ptr(dim, ptr) }; + view.to_owned().into_shared() + } + Self::Owned(arr) => arr, + } + } + + /// Convert to shared ArcArray, suitable for returning from operations. + /// + /// This is equivalent to `into_owned()` but named for clarity. + #[inline] + pub fn into_shared(self) -> ArcArray { + self.into_owned() + } + + /// Get the shape of the tensor. + pub fn shape(&self) -> &[usize] { + match self { + Self::Borrowed { shape, .. } => shape, + Self::Owned(arr) => arr.shape(), + } + } + + /// Get the number of dimensions. + #[inline] + pub fn ndim(&self) -> usize { + self.shape().len() + } + + /// Get the total number of elements. + #[inline] + pub fn len(&self) -> usize { + self.shape().iter().product() + } + + /// Check if the tensor is empty. + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Returns `true` if this is borrowed (zero-copy) storage. + #[inline] + pub fn is_borrowed(&self) -> bool { + matches!(self, Self::Borrowed { .. }) + } + + /// Returns `true` if this is owned storage. + #[inline] + pub fn is_owned(&self) -> bool { + matches!(self, Self::Owned(_)) + } + + /// Ensure owned and return mutable reference to the ArcArray. + /// + /// Converts borrowed to owned if necessary. + pub fn ensure_owned(&mut self) -> &mut ArcArray { + if let Self::Borrowed { bytes, shape } = self { + let ptr = bytes.as_ptr() as *const E; + let dim = IxDyn(shape); + // SAFETY: Same as view() + let view = unsafe { ArrayView::from_shape_ptr(dim, ptr) }; + *self = Self::Owned(view.to_owned().into_shared()); + } + match self { + Self::Owned(arr) => arr, + Self::Borrowed { .. } => unreachable!(), + } + } +} + +/// Convert from ArcArray to NdArrayStorage. +impl From> for NdArrayStorage { + fn from(array: ArcArray) -> Self { + Self::Owned(array) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use alloc::{vec, vec::Vec}; + use burn_std::Bytes; + + #[test] + fn test_borrowed_is_not_unique() { + let data: Vec = vec![1.0, 2.0, 3.0, 4.0]; + let bytes = Bytes::from_elems(data); + let storage = NdArrayStorage::::from_borrowed(bytes, [2, 2]).expect("should create"); + + assert!(!storage.is_unique()); + assert!(storage.is_borrowed()); + } + + #[test] + fn test_owned_unique_when_single_ref() { + let array = ndarray::ArrayD::from_elem(IxDyn(&[2, 2]), 1.0f32).into_shared(); + let storage = NdArrayStorage::from_owned(array); + + assert!(storage.is_unique()); + assert!(storage.is_owned()); + } + + #[test] + fn test_owned_not_unique_when_cloned() { + let array = ndarray::ArrayD::from_elem(IxDyn(&[2, 2]), 1.0f32).into_shared(); + let storage = NdArrayStorage::from_owned(array); + let _clone = storage.clone(); + + assert!(!storage.is_unique()); + } + + #[test] + fn test_view_zero_copy() { + let data: Vec = vec![1.0, 2.0, 3.0, 4.0]; + let bytes = Bytes::from_elems(data); + let storage = NdArrayStorage::::from_borrowed(bytes, [2, 2]).expect("should create"); + + let view = storage.view(); + assert_eq!(view[[0, 0]], 1.0); + assert_eq!(view[[1, 1]], 4.0); + } + + #[test] + fn test_into_owned_copies_borrowed() { + let data: Vec = vec![1.0, 2.0, 3.0, 4.0]; + let bytes = Bytes::from_elems(data); + let storage = NdArrayStorage::::from_borrowed(bytes, [2, 2]).expect("should create"); + + let owned = storage.into_owned(); + assert_eq!(owned[[0, 0]], 1.0); + assert_eq!(owned[[1, 1]], 4.0); + } + + #[test] + fn test_from_borrowed_validates_alignment() { + use burn_std::AllocationProperty; + + // Test 1: Properly aligned data should succeed + let aligned_data: Vec = vec![1.0, 2.0, 3.0, 4.0]; + let aligned_bytes = Bytes::from_elems(aligned_data); + + // Verify test setup - should be 4-byte aligned for f32 + assert_eq!( + (aligned_bytes.as_ptr() as usize) % core::mem::align_of::(), + 0, + "Test setup: f32 data should be properly aligned" + ); + + let result = NdArrayStorage::::from_borrowed(aligned_bytes, [2, 2]); + assert!( + result.is_ok(), + "from_borrowed should succeed for properly aligned data" + ); + + // Test 2: Misaligned data should fail + // Create a buffer large enough to find a misaligned offset + // (static data placement varies by platform, so we find an offset dynamically) + let buffer: &[u8] = &[0u8; 32]; + let shared = bytes::Bytes::from_static(buffer); + let base = shared.as_ptr() as usize; + let align = core::mem::align_of::(); + + // Find an offset in 1..align that produces misalignment (at least one must exist) + let misalign_offset = (1..align) + .find(|&off| !(base + off).is_multiple_of(align)) + .expect("Should find a misaligned offset"); + + let sliced = shared.slice(misalign_offset..(misalign_offset + 16)); + let misaligned_bytes = Bytes::from_shared(sliced, AllocationProperty::Other); + + // Verify test setup - should NOT be 4-byte aligned + assert_ne!( + (misaligned_bytes.as_ptr() as usize) % align, + 0, + "Test setup: sliced data should be misaligned for f32" + ); + + let result = NdArrayStorage::::from_borrowed(misaligned_bytes, [4]); + assert!( + result.is_err(), + "from_borrowed should return Err for misaligned data" + ); + } + + #[test] + fn test_insufficient_size_returns_err() { + // Create bytes that are too small for the requested shape + let data: Vec = vec![1.0, 2.0]; // 8 bytes + let bytes = Bytes::from_elems(data); + + // Try to create storage for 4 elements (needs 16 bytes) + let result = NdArrayStorage::::from_borrowed(bytes, [4]); + assert!( + result.is_err(), + "from_borrowed should return Err when bytes are too small" + ); + } + + // ========================================================================== + // Zero-copy hardening tests + // These tests verify the zero-copy guarantee is maintained. If any of these + // fail, it indicates a regression in zero-copy functionality. + // ========================================================================== + + #[test] + fn test_zero_copy_native_allocation() { + // CRITICAL: Verify that native allocations (Bytes::from_elems) are zero-copy + // on initial load. The view() must return a pointer to the SAME memory. + // + // Note: Native allocations copy on clone (this is expected), but the initial + // load is still zero-copy, avoiding an extra copy in the common case where + // the tensor is used without cloning. + let data: Vec = vec![1.0, 2.0, 3.0, 4.0]; + let bytes = Bytes::from_elems(data); + let original_ptr = bytes.as_ptr(); + + let storage = NdArrayStorage::::from_borrowed(bytes, [2, 2]).expect("should create"); + + // Initial load must be zero-copy + let view = storage.view(); + let view_ptr = view.as_ptr() as *const u8; + + assert_eq!( + original_ptr, view_ptr, + "ZERO-COPY REGRESSION: native allocation view() must return pointer to original bytes" + ); + + // Verify data integrity + assert_eq!(view[[0, 0]], 1.0); + assert_eq!(view[[0, 1]], 2.0); + assert_eq!(view[[1, 0]], 3.0); + assert_eq!(view[[1, 1]], 4.0); + } + + #[test] + fn test_zero_copy_shared_bytes_pointer_identity() { + // CRITICAL: Test with SharedBytesAllocationController for true zero-copy. + // This simulates the actual burnpack/mmap loading path. + use burn_std::AllocationProperty; + + // Create static-like data using bytes::Bytes + let data: &[u8] = &[ + 0, 0, 128, 63, // 1.0f32 in little-endian + 0, 0, 0, 64, // 2.0f32 + 0, 0, 64, 64, // 3.0f32 + 0, 0, 128, 64, // 4.0f32 + ]; + let shared = bytes::Bytes::from_static(data); + let original_ptr = shared.as_ptr(); + + // Create Bytes with SharedBytesAllocationController + let bytes = Bytes::from_shared(shared, AllocationProperty::Other); + + let storage = NdArrayStorage::::from_borrowed(bytes, [2, 2]).expect("should create"); + + // Verify pointer identity + let view_ptr = storage.view().as_ptr() as *const u8; + assert_eq!( + original_ptr, view_ptr, + "ZERO-COPY REGRESSION: SharedBytes view must point to original static data" + ); + + // Clone should also share the same memory + let cloned = storage.clone(); + let cloned_ptr = cloned.view().as_ptr() as *const u8; + assert_eq!( + original_ptr, cloned_ptr, + "ZERO-COPY REGRESSION: SharedBytes clone must share memory" + ); + } + + #[test] + fn test_clone_borrowed_stays_borrowed() { + // Verify that cloning borrowed storage produces another borrowed storage. + // Note: The underlying Bytes may or may not share memory depending on + // the allocation controller (native allocations copy, file-backed may share). + let data: Vec = vec![1.0, 2.0, 3.0, 4.0]; + let bytes = Bytes::from_elems(data); + + let storage = NdArrayStorage::::from_borrowed(bytes, [2, 2]).expect("should create"); + let cloned = storage.clone(); + + // Both should still be borrowed (the storage type is preserved) + assert!( + storage.is_borrowed(), + "ZERO-COPY REGRESSION: original should remain borrowed after clone" + ); + assert!( + cloned.is_borrowed(), + "ZERO-COPY REGRESSION: clone should be borrowed type" + ); + + // Both should report not unique (important for COW behavior) + assert!( + !storage.is_unique(), + "ZERO-COPY REGRESSION: original should not be unique after clone" + ); + assert!( + !cloned.is_unique(), + "ZERO-COPY REGRESSION: clone should not be unique" + ); + + // Data should be identical + assert_eq!(storage.view(), cloned.view(), "Clone should have same data"); + } + + #[test] + fn test_zero_copy_triggers_copy_on_mutation() { + // Verify that into_owned() on borrowed data creates a NEW allocation + // (this is the "copy" in copy-on-write) + let data: Vec = vec![1.0, 2.0, 3.0, 4.0]; + let bytes = Bytes::from_elems(data); + let original_ptr = bytes.as_ptr(); + + let storage = NdArrayStorage::::from_borrowed(bytes, [2, 2]).expect("should create"); + + assert!(storage.is_borrowed(), "should start as borrowed"); + + let owned = storage.into_owned(); + let owned_ptr = owned.as_ptr() as *const u8; + + assert_ne!( + original_ptr, owned_ptr, + "into_owned() on borrowed data MUST allocate new memory (copy-on-write)" + ); + } + + #[test] + fn test_borrowed_reports_not_unique() { + // CRITICAL: Borrowed storage must report is_unique() == false + // This is what triggers copy-on-write in mutation operations + let data: Vec = vec![1.0, 2.0, 3.0, 4.0]; + let bytes = Bytes::from_elems(data); + let storage = NdArrayStorage::::from_borrowed(bytes, [2, 2]).expect("should create"); + + assert!( + !storage.is_unique(), + "ZERO-COPY REGRESSION: borrowed storage MUST report is_unique() == false \ + to trigger copy-on-write. If this is true, mutations will corrupt shared data!" + ); + } +} diff --git a/crates/burn-ndarray/src/tensor.rs b/crates/burn-ndarray/src/tensor.rs new file mode 100644 index 0000000..baa3df7 --- /dev/null +++ b/crates/burn-ndarray/src/tensor.rs @@ -0,0 +1,971 @@ +use burn_backend::{ + AllocationProperty, DType, Element, Shape, TensorData, TensorMetadata, + quantization::{QParams, QuantLevel, QuantMode, QuantScheme, QuantValue}, +}; +use burn_std::BoolStore; + +use crate::ops::quantization::{QuantizationStrategy, SymmetricQuantization}; +use crate::{NdArrayDevice, NdArrayStorage}; +use alloc::vec::Vec; +use ndarray::{ArcArray, ArrayD, IxDyn}; + +/// Concrete storage type for ndarray (owned with COW semantics via Arc) +pub type SharedArray = ArcArray; + +/// Tensor primitive used by the [ndarray backend](crate::NdArray). +/// +/// Supports both owned and borrowed (zero-copy) data via `NdArrayStorage`. +/// When data is borrowed from external sources (like burnpack files), +/// it remains zero-copy until a mutating operation is performed. +#[derive(Debug, Clone)] +#[allow(missing_docs)] +pub enum NdArrayTensor { + F64(NdArrayStorage), + F32(NdArrayStorage), + I64(NdArrayStorage), + I32(NdArrayStorage), + I16(NdArrayStorage), + I8(NdArrayStorage), + U64(NdArrayStorage), + U32(NdArrayStorage), + U16(NdArrayStorage), + U8(NdArrayStorage), + Bool(NdArrayStorage), +} + +impl NdArrayTensor { + /// Extract bool array, converting to owned if necessary. + pub(crate) fn bool(self) -> SharedArray { + match self { + NdArrayTensor::Bool(storage) => storage.into_shared(), + _ => unimplemented!("Expected bool tensor, got {:?}", self.dtype()), + } + } + + /// Returns true if this tensor uses borrowed (zero-copy) storage. + #[inline] + pub fn is_borrowed(&self) -> bool { + macro_rules! check { + ($($variant:ident),*) => { + match self { + $(NdArrayTensor::$variant(s) => s.is_borrowed(),)* + } + }; + } + check!(F64, F32, I64, I32, I16, I8, U64, U32, U16, U8, Bool) + } +} + +pub(crate) fn cast_to_dtype(array: SharedArray, dtype: DType) -> NdArrayTensor +where + NdArrayTensor: From>, +{ + fn cast(array: SharedArray) -> SharedArray { + array.mapv(|a| a.elem()).into_shared() + } + + if E1::dtype() == dtype { + return array.into(); + } + + match dtype { + DType::F64 => cast::(array).into(), + DType::F32 => cast::(array).into(), + DType::Flex32 => cast::(array).into(), + DType::I64 => cast::(array).into(), + DType::I32 => cast::(array).into(), + DType::I16 => cast::(array).into(), + DType::I8 => cast::(array).into(), + DType::U64 => cast::(array).into(), + DType::U32 => cast::(array).into(), + DType::U16 => cast::(array).into(), + DType::U8 => cast::(array).into(), + DType::Bool(BoolStore::Native) => cast::(array).into(), + dtype => panic!("Unsupported dtype: {dtype:?}"), + } +} + +macro_rules! impl_from { + ($($ty: ty => $dtype: ident),*) => { + // From SharedArray (owned) -> NdArrayTensor + $(impl From> for NdArrayTensor { + fn from(value: SharedArray<$ty>) -> NdArrayTensor { + NdArrayTensor::$dtype(NdArrayStorage::from_owned(value)) + } + })* + + // From NdArrayStorage -> NdArrayTensor + $(impl From> for NdArrayTensor { + fn from(value: NdArrayStorage<$ty>) -> NdArrayTensor { + NdArrayTensor::$dtype(value) + } + })* + }; +} + +impl_from!( + f64 => F64, f32 => F32, + i64 => I64, i32 => I32, i16 => I16, i8 => I8, + u64 => U64, u32 => U32, u16 => U16, u8 => U8, + bool => Bool +); + +/// Macro to execute an operation on a given element type. +/// +/// Extracts the storage from NdArrayTensor, converts to SharedArray, and passes to operation. +/// +/// # Panics +/// Since there is no automatic type cast at this time, binary operations for different +/// floating point precision data types will panic with a data type mismatch. +#[macro_export] +macro_rules! execute_with_dtype { + (($lhs:expr, $rhs:expr),$element:ident, $op:expr, [$($dtype: ident => $ty: ty),*]) => {{ + let lhs_dtype = burn_backend::TensorMetadata::dtype(&$lhs); + let rhs_dtype = burn_backend::TensorMetadata::dtype(&$rhs); + match ($lhs, $rhs) { + $( + ($crate::NdArrayTensor::$dtype(lhs), $crate::NdArrayTensor::$dtype(rhs)) => { + #[allow(unused)] + type $element = $ty; + // Convert storage to SharedArray for compatibility with existing operations + $op(lhs.into_shared(), rhs.into_shared()).into() + } + )* + _ => panic!( + "Data type mismatch (lhs: {:?}, rhs: {:?})", + lhs_dtype, rhs_dtype + ), + } + }}; + // Binary op: type automatically inferred by the compiler + (($lhs:expr, $rhs:expr), $op:expr) => {{ + $crate::execute_with_dtype!(($lhs, $rhs), E, $op) + }}; + + // Binary op: generic type cannot be inferred for an operation + (($lhs:expr, $rhs:expr), $element:ident, $op:expr) => {{ + $crate::execute_with_dtype!(($lhs, $rhs), $element, $op, [ + F64 => f64, F32 => f32, + I64 => i64, I32 => i32, I16 => i16, I8 => i8, + U64 => u64, U32 => u32, U16 => u16, U8 => u8, + Bool => bool + ]) + }}; + + ($tensor:expr, $element:ident, $op:expr, [$($dtype: ident => $ty: ty),*]) => {{ + match $tensor { + $( + $crate::NdArrayTensor::$dtype(storage) => { + #[allow(unused)] + type $element = $ty; + // Convert to SharedArray for compatibility with most operations + $op(storage.into_shared()).into() + } + )* + #[allow(unreachable_patterns)] + other => unimplemented!("unsupported dtype: {:?}", other.dtype()) + } + }}; + // Unary op: type automatically inferred by the compiler + ($tensor:expr, $op:expr) => {{ + $crate::execute_with_dtype!($tensor, E, $op) + }}; + + // Unary op: generic type cannot be inferred for an operation + ($tensor:expr, $element:ident, $op:expr) => {{ + $crate::execute_with_dtype!($tensor, $element, $op, [ + F64 => f64, F32 => f32, + I64 => i64, I32 => i32, I16 => i16, I8 => i8, + U64 => u64, U32 => u32, U16 => u16, U8 => u8, + Bool => bool + ]) + }}; +} + +/// Macro to execute an operation a given element type. +/// Only handles float types. +/// +/// # Panics +/// Since there is no automatic type cast at this time, binary operations for different +/// floating point precision data types will panic with a data type mismatch. +#[macro_export] +macro_rules! execute_with_float_dtype { + // Binary op: type automatically inferred by the compiler + (($lhs:expr, $rhs:expr), $op:expr) => {{ + $crate::execute_with_float_dtype!(($lhs, $rhs), E, $op) + }}; + + // Binary op: generic type cannot be inferred for an operation + (($lhs:expr, $rhs:expr), $element:ident, $op:expr) => {{ + $crate::execute_with_dtype!(($lhs, $rhs), $element, $op, [ + F64 => f64, F32 => f32 + ]) + }}; + + // Unary op: type automatically inferred by the compiler + ($tensor:expr, $op:expr) => {{ + $crate::execute_with_float_dtype!($tensor, E, $op) + }}; + + // Unary op: generic type cannot be inferred for an operation + ($tensor:expr, $element:ident, $op:expr) => {{ + $crate::execute_with_dtype!($tensor, $element, $op, [ + F64 => f64, F32 => f32 + ]) + }}; +} + +/// Macro to execute an operation a given element type. +/// Only handles int types. +/// +/// # Panics +/// Since there is no automatic type cast at this time, binary operations for different +/// floating point precision data types will panic with a data type mismatch. +#[macro_export] +macro_rules! execute_with_int_dtype { + // Binary op: type automatically inferred by the compiler + (($lhs:expr, $rhs:expr), $op:expr) => {{ + $crate::execute_with_int_dtype!(($lhs, $rhs), E, $op) + }}; + + // Binary op: generic type cannot be inferred for an operation + (($lhs:expr, $rhs:expr), $element:ident, $op:expr) => {{ + $crate::execute_with_dtype!(($lhs, $rhs), $element, $op, [ + I64 => i64, I32 => i32, I16 => i16, I8 => i8, + U64 => u64, U32 => u32, U16 => u16, U8 => u8 + ]) + }}; + + // Unary op: type automatically inferred by the compiler + ($tensor:expr, $op:expr) => {{ + $crate::execute_with_int_dtype!($tensor, E, $op) + }}; + + // Unary op: generic type cannot be inferred for an operation + ($tensor:expr, $element:ident, $op:expr) => {{ + $crate::execute_with_dtype!($tensor, $element, $op, [ + I64 => i64, I32 => i32, I16 => i16, I8 => i8, + U64 => u64, U32 => u32, U16 => u16, U8 => u8 + ]) + }}; +} + +/// Macro to execute an operation a given element type. +/// Only handles numeric types +/// +/// # Panics +/// Since there is no automatic type cast at this time, binary operations for different +/// floating point precision data types will panic with a data type mismatch. +#[macro_export] +macro_rules! execute_with_numeric_dtype { + // Binary op: type automatically inferred by the compiler + (($lhs:expr, $rhs:expr), $op:expr) => {{ + $crate::execute_with_numeric_dtype!(($lhs, $rhs), E, $op) + }}; + + // Binary op: generic type cannot be inferred for an operation + (($lhs:expr, $rhs:expr), $element:ident, $op:expr) => {{ + $crate::execute_with_dtype!(($lhs, $rhs), $element, $op, [ + F64 => f64, F32 => f32, + I64 => i64, I32 => i32, I16 => i16, I8 => i8, + U64 => u64, U32 => u32, U16 => u16, U8 => u8 + ]) + }}; + + // Unary op: type automatically inferred by the compiler + ($tensor:expr, $op:expr) => {{ + $crate::execute_with_numeric_dtype!($tensor, E, $op) + }}; + + // Unary op: generic type cannot be inferred for an operation + ($tensor:expr, $element:ident, $op:expr) => {{ + $crate::execute_with_dtype!($tensor, $element, $op, [ + F64 => f64, F32 => f32, + I64 => i64, I32 => i32, I16 => i16, I8 => i8, + U64 => u64, U32 => u32, U16 => u16, U8 => u8 + ]) + }}; +} + +/// Macro to execute a cat operation on a given set of element types. +/// +/// Uses zero-copy views from storage for concatenation. +/// +/// # Panics +/// Since there is no automatic type cast at this time, binary operations for different +/// floating point precision data types will panic with a data type mismatch. +#[macro_export] +macro_rules! cat_with_dtype { + ($tensors: expr, $dim: expr, [$($dtype: ident),*]) => { + match &$tensors[0] { + $(NdArrayTensor::$dtype(_) => { + let tensors = $tensors + .iter() + .map(|t| { + if let NdArrayTensor::$dtype(storage) = t { + // Use storage.view() for zero-copy access + storage.view() + } else { + panic!("Concatenate data type mismatch (expected {:?}, got {:?})", $tensors[0].dtype(), t.dtype()) + } + }) + .collect::>(); + NdArrayOps::concatenate(&tensors, $dim).into() + })* + _ => panic!("Unsupported dtype: {:?}", $tensors[0].dtype()) + } + }; +} + +/// Macro to execute an operation that returns a given element type. +#[macro_export] +macro_rules! execute_with_float_out_dtype { + ($out_dtype:expr, $element:ident, $op:expr, [$($dtype: ident => $ty: ty),*]) => {{ + match $out_dtype { + $( + burn_std::FloatDType::$dtype => { + #[allow(unused)] + type $element = $ty; + $op + } + )* + #[allow(unreachable_patterns)] + other => unimplemented!("unsupported dtype: {other:?}") + } + }}; + // Unary op: type automatically inferred by the compiler + ($out_dtype:expr, $op:expr) => {{ + $crate::execute_with_float_out_dtype!($out_dtype, E, $op) + }}; + + // Unary op: generic type cannot be inferred for an operation + ($out_dtype:expr, $element:ident, $op:expr) => {{ + $crate::execute_with_float_out_dtype!($out_dtype, $element, $op, [ + F64 => f64, F32 => f32 + ]) + }}; +} + +/// Macro to execute an operation that returns a given element type. +#[macro_export] +macro_rules! execute_with_int_out_dtype { + ($out_dtype:expr, $element:ident, $op:expr, [$($dtype: ident => $ty: ty),*]) => {{ + match $out_dtype { + $( + burn_std::IntDType::$dtype => { + #[allow(unused)] + type $element = $ty; + $op + } + )* + #[allow(unreachable_patterns)] + other => unimplemented!("unsupported dtype: {other:?}") + } + }}; + // Unary op: type automatically inferred by the compiler + ($out_dtype:expr, $op:expr) => {{ + $crate::execute_with_int_out_dtype!($out_dtype, E, $op) + }}; + + // Unary op: generic type cannot be inferred for an operation + ($out_dtype:expr, $element:ident, $op:expr) => {{ + $crate::execute_with_int_out_dtype!($out_dtype, $element, $op, [ + I64 => i64, I32 => i32, I16 => i16, I8 => i8, + U64 => u64, U32 => u32, U16 => u16, U8 => u8 + ]) + }}; +} + +impl TensorMetadata for NdArrayTensor { + type Device = NdArrayDevice; + fn dtype(&self) -> DType { + match self { + NdArrayTensor::F64(_) => DType::F64, + NdArrayTensor::F32(_) => DType::F32, + NdArrayTensor::I64(_) => DType::I64, + NdArrayTensor::I32(_) => DType::I32, + NdArrayTensor::I16(_) => DType::I16, + NdArrayTensor::I8(_) => DType::I8, + NdArrayTensor::U64(_) => DType::U64, + NdArrayTensor::U32(_) => DType::U32, + NdArrayTensor::U16(_) => DType::U16, + NdArrayTensor::U8(_) => DType::U8, + NdArrayTensor::Bool(_) => DType::Bool(BoolStore::Native), + } + } + + fn shape(&self) -> Shape { + // Use storage's shape method (works for both borrowed and owned) + macro_rules! get_shape { + ($($variant:ident),*) => { + match self { + $(NdArrayTensor::$variant(storage) => Shape::from(storage.shape().to_vec()),)* + } + }; + } + get_shape!(F64, F32, I64, I32, I16, I8, U64, U32, U16, U8, Bool) + } + + fn rank(&self) -> usize { + self.shape().num_dims() + } + + fn device(&self) -> NdArrayDevice { + NdArrayDevice::Cpu + } + + fn can_mut(&self) -> bool { + // NdArray storage is copy-on-write (`ArcArray`) without a public + // uniqueness check at this level; in-place ops resolve sharing + // themselves, so conservatively report the buffer as shared. + false + } +} + +pub(crate) trait ShapeOps { + fn num_dims(self) -> usize; + fn num_elements(self) -> usize; + fn dims(self) -> [usize; N]; + fn into_shape(self) -> Shape; +} + +impl ShapeOps for &[usize] { + fn num_dims(self) -> usize { + self.len() + } + + fn num_elements(self) -> usize { + self.iter().product() + } + + fn dims(self) -> [usize; N] { + self.try_into().unwrap() + } + + fn into_shape(self) -> Shape { + Shape::from(self) + } +} + +mod utils { + use burn_std::tensor::is_contiguous; + + use super::*; + + impl NdArrayTensor { + pub(crate) fn into_data(self) -> TensorData { + let shape = self.shape(); + let contiguous = self.is_contiguous(); + + fn inner( + shape: Shape, + is_contiguous: bool, + array: ArcArray, + ) -> TensorData { + let vec = if is_contiguous { + match array.try_into_owned_nocopy() { + Ok(owned) => { + let (mut vec, offset) = owned.into_raw_vec_and_offset(); + if let Some(offset) = offset { + vec.drain(..offset); + } + if vec.len() > shape.num_elements() { + vec.drain(shape.num_elements()..vec.len()); + } + vec + } + Err(array) => array.into_iter().collect(), + } + } else { + array.into_iter().collect() + }; + + TensorData::new(vec, shape) + } + + // Convert storage to owned array before extracting data + execute_with_dtype!(self, |arr| inner(shape, contiguous, arr)) + } + + pub(crate) fn is_contiguous(&self) -> bool { + // For borrowed data, we assume it's contiguous (it came from TensorData which is contiguous) + // For owned data, we check the strides + macro_rules! check_contiguous { + ($($variant:ident),*) => { + match self { + $(NdArrayTensor::$variant(storage) => { + match storage { + NdArrayStorage::Borrowed { .. } => { + // Borrowed storage requires contiguous row-major data + // (see NdArrayStorage::from_borrowed documentation) + true + } + NdArrayStorage::Owned(array) => { + let shape = array.shape(); + let mut strides = Vec::with_capacity(array.strides().len()); + for &stride in array.strides() { + if stride <= 0 { + return false; + } + strides.push(stride as usize); + } + is_contiguous(shape, &strides) + } + } + })* + } + }; + } + check_contiguous!(F64, F32, I64, I32, I16, I8, U64, U32, U16, U8, Bool) + } + } +} + +/// Converts a slice of usize to a typed dimension. +#[macro_export(local_inner_macros)] +macro_rules! to_typed_dims { + ( + $n:expr, + $dims:expr, + justdim + ) => {{ + let mut dims = [0; $n]; + for i in 0..$n { + dims[i] = $dims[i]; + } + let dim: Dim<[usize; $n]> = Dim(dims); + dim + }}; +} + +/// Reshapes an array into a tensor. +#[macro_export(local_inner_macros)] +macro_rules! reshape { + ( + ty $ty:ty, + n $n:expr, + shape $shape:expr, + array $array:expr + ) => {{ + let dim = $crate::to_typed_dims!($n, $shape, justdim); + let array = match $array.is_standard_layout() { + // Move the array into the new shape rather than going through + // `to_shape`: the latter returns a borrowed view here, which + // `into_shared` then clones, copying the buffer on every reshape. + // Moving rewrites the dimensions in place, and the buffer stays + // shared for copy-on-write like in any other operation. + true => { + match $array.into_shape_with_order(dim) { + Ok(val) => val, + Err(err) => { + core::panic!("Shape should be compatible shape={dim:?}: {err:?}"); + } + } + }, + false => $array.to_shape(dim).unwrap().as_standard_layout().into_shared(), + }; + array.into_dyn() + }}; + ( + ty $ty:ty, + shape $shape:expr, + array $array:expr, + d $D:expr + ) => {{ + match $D { + 1 => reshape!(ty $ty, n 1, shape $shape, array $array), + 2 => reshape!(ty $ty, n 2, shape $shape, array $array), + 3 => reshape!(ty $ty, n 3, shape $shape, array $array), + 4 => reshape!(ty $ty, n 4, shape $shape, array $array), + 5 => reshape!(ty $ty, n 5, shape $shape, array $array), + 6 => reshape!(ty $ty, n 6, shape $shape, array $array), + _ => core::panic!("NdArray supports arrays up to 6 dimensions, received: {}", $D), + } + }}; +} + +/// Slice a tensor +#[macro_export] +macro_rules! slice { + ($tensor:expr, $slices:expr) => { + slice!($tensor, $slices, F64, F32, I64, I32, I16, I8, U64, U32, U16, U8, Bool) + }; + ($tensor:expr, $slices:expr, $($variant:ident),*) => { + match $tensor { + $(NdArrayTensor::$variant(s) => { NdArrayOps::slice(s.view(), $slices).into() })* + } + }; +} + +impl NdArrayTensor { + /// Create a new [ndarray tensor](NdArrayTensor) from [data](TensorData). + /// + /// This method attempts zero-copy loading when possible. If the data has properly + /// aligned bytes that can be borrowed, it creates a borrowed tensor. Otherwise, + /// it falls back to copying the data. + /// + /// Zero-copy loading works when: + /// - The data's bytes are properly aligned for the element type + /// - The bytes can be borrowed (e.g., from mmap'd file or static data) + pub fn from_data(data: TensorData) -> NdArrayTensor { + // Only use Borrowed storage for non-native allocations (e.g., burnpack mmap/file). + // For native Rust heap allocations (the common case), go directly to owned storage: + // `from_data_owned` reclaims the Vec zero-copy via `into_vec`, while + // Borrowed storage would trigger a full memcopy on every single operation. + if data.bytes.property() != AllocationProperty::Native { + match Self::try_from_data_borrowed(data) { + Ok(tensor) => return tensor, + Err(data) => return Self::from_data_owned(data), + } + } + Self::from_data_owned(data) + } + + /// Try to create a tensor with borrowed storage (zero-copy). + /// + /// Takes ownership of TensorData and returns it back on failure. + /// No cloning occurs - bytes are moved into storage or returned on failure. + /// + /// Returns `Err(data)` if borrowing is not possible (e.g., misaligned data). + fn try_from_data_borrowed(data: TensorData) -> Result { + let TensorData { + bytes, + shape, + dtype, + } = data; + + macro_rules! try_borrow { + ($ty:ty, $variant:ident, $bytes:expr, $shape:expr) => { + match NdArrayStorage::<$ty>::from_borrowed($bytes, $shape) { + Ok(storage) => return Ok(NdArrayTensor::$variant(storage)), + Err((bytes, shape)) => (bytes, shape), + } + }; + } + + // Try to create borrowed storage; get bytes back on failure + let (bytes, shape) = match dtype { + DType::F64 => try_borrow!(f64, F64, bytes, shape), + DType::F32 => try_borrow!(f32, F32, bytes, shape), + DType::I64 => try_borrow!(i64, I64, bytes, shape), + DType::I32 => try_borrow!(i32, I32, bytes, shape), + DType::I16 => try_borrow!(i16, I16, bytes, shape), + DType::I8 => try_borrow!(i8, I8, bytes, shape), + DType::U64 => try_borrow!(u64, U64, bytes, shape), + DType::U32 => try_borrow!(u32, U32, bytes, shape), + DType::U16 => try_borrow!(u16, U16, bytes, shape), + DType::U8 => try_borrow!(u8, U8, bytes, shape), + DType::Bool(BoolStore::Native) => try_borrow!(bool, Bool, bytes, shape), + _ => (bytes, shape), // QFloat not supported for zero-copy + }; + + Err(TensorData { + bytes, + shape, + dtype, + }) + } + + /// Create a tensor with owned storage. + /// + /// This may or may not copy data depending on whether the underlying bytes + /// can be reclaimed (via `try_into_vec`). If bytes are uniquely owned, + /// no copy occurs; otherwise data is copied to a new allocation. + fn from_data_owned(data: TensorData) -> NdArrayTensor { + let shape = data.shape.to_vec(); // TODO: into_vec + + macro_rules! execute { + ($data: expr, [$($dtype: pat => $ty: ty),*]) => { + match $data.dtype { + $( $dtype => { + match data.into_vec::<$ty>() { + Ok(vec) => unsafe { ArrayD::from_shape_vec_unchecked(shape, vec) }.into_shared(), + Err(err) => panic!("Data should have the same element type as the tensor {err:?}"), + }.into() + }, )* + other => unimplemented!("Unsupported dtype {other:?}"), + } + }; + } + + execute!(data, [ + DType::F64 => f64, DType::F32 => f32, + DType::I64 => i64, DType::I32 => i32, DType::I16 => i16, DType::I8 => i8, + DType::U64 => u64, DType::U32 => u32, DType::U16 => u16, DType::U8 => u8, + DType::Bool(BoolStore::Native) => bool + ]) + } +} + +/// A quantized tensor for the ndarray backend. +#[derive(Clone, Debug)] +pub struct NdArrayQTensor { + /// The quantized tensor. + pub qtensor: NdArrayTensor, + /// The quantization scheme. + pub scheme: QuantScheme, + /// The quantization parameters. + pub qparams: Vec>, +} + +impl NdArrayQTensor { + /// Returns the quantization strategy, including quantization parameters, for the given tensor. + pub fn strategy(&self) -> QuantizationStrategy { + match self.scheme { + QuantScheme { + level: QuantLevel::Tensor, + mode: QuantMode::Symmetric, + value: + QuantValue::Q8F + | QuantValue::Q8S + | QuantValue::E4M3 + | QuantValue::E5M2 + | QuantValue::Q4F + | QuantValue::Q4S + | QuantValue::E2M1 + | QuantValue::Q2F + | QuantValue::Q2S, + .. + } => QuantizationStrategy::PerTensorSymmetric(SymmetricQuantization::init( + self.qparams[0].scales, + self.scheme.value, + )), + QuantScheme { + level: QuantLevel::Block(block_size), + mode: QuantMode::Symmetric, + value: + QuantValue::Q8F + | QuantValue::Q8S + | QuantValue::E4M3 + | QuantValue::E5M2 + | QuantValue::Q4F + | QuantValue::Q4S + | QuantValue::E2M1 + | QuantValue::Q2F + | QuantValue::Q2S, + .. + } => QuantizationStrategy::PerBlockSymmetric( + self.qparams + .iter() + .map(|q| SymmetricQuantization::init(q.scales, self.scheme.value)) + .collect(), + block_size, + ), + } + } +} + +impl TensorMetadata for NdArrayQTensor { + type Device = NdArrayDevice; + fn dtype(&self) -> DType { + DType::QFloat(self.scheme) + } + + fn shape(&self) -> Shape { + self.qtensor.shape() + } + + fn rank(&self) -> usize { + self.shape().num_dims() + } + + fn device(&self) -> Self::Device { + NdArrayDevice::Cpu + } + + fn can_mut(&self) -> bool { + self.qtensor.can_mut() + } +} + +#[cfg(test)] +mod tests { + use crate::NdArray; + use alloc::vec; + + use super::*; + use burn_backend::{ + Distribution, + ops::{FloatTensorOps, QTensorOps}, + quantization::{QuantStore, QuantizationParametersPrimitive}, + }; + use burn_std::rand::get_seeded_rng; + + #[test] + fn should_support_into_and_from_data_1d() { + let data_expected = TensorData::random::( + Shape::new([3]), + Distribution::Default, + &mut get_seeded_rng(), + ); + let tensor = NdArrayTensor::from_data(data_expected.clone()); + + let data_actual = tensor.into_data(); + + assert_eq!(data_expected, data_actual); + } + + #[test] + fn should_support_into_and_from_data_2d() { + let data_expected = TensorData::random::( + Shape::new([2, 3]), + Distribution::Default, + &mut get_seeded_rng(), + ); + let tensor = NdArrayTensor::from_data(data_expected.clone()); + + let data_actual = tensor.into_data(); + + assert_eq!(data_expected, data_actual); + } + + #[test] + fn should_support_into_and_from_data_3d() { + let data_expected = TensorData::random::( + Shape::new([2, 3, 4]), + Distribution::Default, + &mut get_seeded_rng(), + ); + let tensor = NdArrayTensor::from_data(data_expected.clone()); + + let data_actual = tensor.into_data(); + + assert_eq!(data_expected, data_actual); + } + + #[test] + fn should_support_into_and_from_data_4d() { + let data_expected = TensorData::random::( + Shape::new([2, 3, 4, 2]), + Distribution::Default, + &mut get_seeded_rng(), + ); + let tensor = NdArrayTensor::from_data(data_expected.clone()); + + let data_actual = tensor.into_data(); + + assert_eq!(data_expected, data_actual); + } + + #[test] + fn should_support_qtensor_strategy() { + type B = NdArray; + let scale: f32 = 0.009_019_608; + let device = Default::default(); + + let tensor = B::float_from_data(TensorData::from([-1.8f32, -1.0, 0.0, 0.5]), &device); + let scheme = QuantScheme::default() + .with_value(QuantValue::Q8S) + .with_store(QuantStore::Native); + let qparams = QuantizationParametersPrimitive { + scales: B::float_from_data(TensorData::from([scale]), &device), + }; + let qtensor: NdArrayQTensor = B::quantize(tensor, &scheme, qparams); + + assert_eq!(qtensor.scheme(), scheme); + assert_eq!( + qtensor.strategy(), + QuantizationStrategy::PerTensorSymmetric(SymmetricQuantization::init( + scale, + QuantValue::Q8S + )) + ); + } + + // ========================================================================== + // Zero-copy integration tests + // These tests verify end-to-end zero-copy behavior through NdArrayTensor. + // ========================================================================== + + #[test] + fn zero_copy_creates_borrowed_storage_for_non_native() { + // Verify that from_data creates borrowed storage for non-native allocations + // (e.g. burnpack mmap/file data tagged with AllocationProperty::Other or File). + // Native heap allocations intentionally use Owned storage for performance. + use burn_backend::AllocationProperty; + use burn_std::Bytes; + + let data: Vec = vec![1.0, 2.0, 3.0, 4.0]; + let bytes = Bytes::from_elems(data); + // Tag as Other to simulate burnpack / mmap data (non-native backing storage) + let non_native_bytes = Bytes::from_shared( + bytes::Bytes::copy_from_slice(&bytes), + AllocationProperty::Other, + ); + let tensor_data = TensorData::from_bytes(non_native_bytes, Shape::new([2, 2]), DType::F32); + + let tensor = NdArrayTensor::from_data(tensor_data); + + match &tensor { + NdArrayTensor::F32(storage) => { + assert!( + storage.is_borrowed(), + "ZERO-COPY REGRESSION: from_data should create borrowed storage \ + for non-native (e.g. burnpack) TensorData" + ); + assert!( + !storage.is_unique(), + "ZERO-COPY REGRESSION: borrowed storage must report is_unique() == false" + ); + } + _ => panic!("Expected F32 tensor"), + } + } + + #[test] + fn native_alloc_creates_owned_storage() { + // Native heap allocations must use Owned storage to avoid the memcpy. + use burn_std::Bytes; + + let data: Vec = vec![1.0, 2.0, 3.0, 4.0]; + let bytes = Bytes::from_elems(data); // AllocationProperty::Native + let tensor_data = TensorData::from_bytes(bytes, Shape::new([2, 2]), DType::F32); + + let tensor = NdArrayTensor::from_data(tensor_data); + + match &tensor { + NdArrayTensor::F32(storage) => { + assert!( + !storage.is_borrowed(), + "PERF REGRESSION: from_data must NOT create borrowed storage \ + for native TensorData" + ); + } + _ => panic!("Expected F32 tensor"), + } + } + + #[test] + fn zero_copy_data_integrity() { + // Verify data is correctly accessible through borrowed storage + use burn_std::Bytes; + + let data: Vec = vec![1.0, 2.0, 3.0, 4.0]; + let bytes = Bytes::from_elems(data); + let tensor_data = TensorData::from_bytes(bytes, Shape::new([2, 2]), DType::F32); + + let tensor = NdArrayTensor::from_data(tensor_data); + + match &tensor { + NdArrayTensor::F32(storage) => { + let view = storage.view(); + assert_eq!(view[[0, 0]], 1.0); + assert_eq!(view[[0, 1]], 2.0); + assert_eq!(view[[1, 0]], 3.0); + assert_eq!(view[[1, 1]], 4.0); + } + _ => panic!("Expected F32 tensor"), + } + } + + #[test] + fn zero_copy_fallback_when_bytes_owned() { + // When TensorData owns bytes exclusively, it may use the copy path + // This is expected behavior - verify it still works correctly + let data = TensorData::from([1.0f32, 2.0, 3.0, 4.0]); + let tensor = NdArrayTensor::from_data(data.clone()); + let result = tensor.into_data(); + + assert_eq!(data, result, "Data should round-trip correctly"); + } +} diff --git a/crates/burn-nn/Cargo.toml b/crates/burn-nn/Cargo.toml new file mode 100644 index 0000000..dd4d09b --- /dev/null +++ b/crates/burn-nn/Cargo.toml @@ -0,0 +1,41 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science", "no-std", "embedded", "wasm"] +description = "Neural network building blocks for the Burn deep learning framework" +documentation = "https://docs.rs/burn-nn" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "tensor", "pytorch", "ndarray"] +license.workspace = true +name = "burn-nn" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-nn" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["std", "burn-core/default"] +doc = [ + "std", + # Doc features + "burn-core/doc", +] +std = ["burn-core/std", "num-traits/std"] + +tracing = ["burn-core/tracing"] + +[dependencies] + +# ** Please make sure all dependencies support no_std when std is disabled ** +burn-core = { workspace = true } + +num-traits = { workspace = true } + +[dev-dependencies] +burn-core = { workspace = true, features = ["autodiff"] } +rstest = { workspace = true } + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-nn/README.md b/crates/burn-nn/README.md new file mode 100644 index 0000000..daef9bc --- /dev/null +++ b/crates/burn-nn/README.md @@ -0,0 +1,3 @@ +# Burn Neural Networks + +Core building blocks for Burn neural networks. \ No newline at end of file diff --git a/crates/burn-nn/src/activation/activation_wrapper.rs b/crates/burn-nn/src/activation/activation_wrapper.rs new file mode 100644 index 0000000..9beb962 --- /dev/null +++ b/crates/burn-nn/src/activation/activation_wrapper.rs @@ -0,0 +1,621 @@ +use burn_core as burn; + +use crate::Identity; +use crate::activation::{ + Celu, CeluConfig, Elu, EluConfig, Gelu, HardShrink, HardShrinkConfig, HardSigmoid, + HardSigmoidConfig, HardSwish, LeakyRelu, LeakyReluConfig, PRelu, PReluConfig, Relu, Selu, + Shrink, ShrinkConfig, Sigmoid, SoftShrink, SoftShrinkConfig, Softplus, SoftplusConfig, + Softsign, SwiGlu, SwiGluConfig, Tanh, ThresholdedRelu, ThresholdedReluConfig, +}; +use burn::config::Config; +use burn::module::Module; +use burn::tensor::Device; +use burn::tensor::Tensor; + +/// [`Activation`] Configuration. +#[derive(Config, Debug)] +#[non_exhaustive] +pub enum ActivationConfig { + /// Identity activation layer. + Identity, + + /// [`Gelu`] activation layer. + Gelu, + + /// [`Gelu`] activation layer with tanh approximation. + GeluApproximate, + + /// [`PRelu`] activation layer. + PRelu(PReluConfig), + + /// [`Relu`] activation layer. + Relu, + + /// [`LeakyRelu`] activation layer. + LeakyRelu(LeakyReluConfig), + + /// [`SwiGlu`] activation layer. + SwiGlu(SwiGluConfig), + + /// [`Selu`] activation layer. + Selu, + + /// [`Sigmoid`] activation layer. + Sigmoid, + + /// [`Tanh`] activation layer. + Tanh, + + /// [`HardSigmoid`] activation layer. + HardSigmoid(HardSigmoidConfig), + + /// [`HardSwish`] activation layer. + HardSwish, + + /// [`Softplus`] activation layer. + Softplus(SoftplusConfig), + + /// [`Softsign`] activation layer. + Softsign, + + /// [`Elu`] activation layer. + Elu(EluConfig), + + /// [`Celu`] activation layer. + Celu(CeluConfig), + + /// [`ThresholdedRelu`] activation layer. + ThresholdedRelu(ThresholdedReluConfig), + + /// [`HardShrink`] activation layer. + HardShrink(HardShrinkConfig), + + /// [`SoftShrink`] activation layer. + SoftShrink(SoftShrinkConfig), + + /// [`Shrink`] activation layer. + Shrink(ShrinkConfig), +} + +impl From for ActivationConfig { + fn from(config: PReluConfig) -> Self { + Self::PRelu(config) + } +} + +impl From for ActivationConfig { + fn from(config: LeakyReluConfig) -> Self { + Self::LeakyRelu(config) + } +} + +impl From for ActivationConfig { + fn from(config: SwiGluConfig) -> Self { + Self::SwiGlu(config) + } +} + +impl From for ActivationConfig { + fn from(config: HardSigmoidConfig) -> Self { + Self::HardSigmoid(config) + } +} + +impl From for ActivationConfig { + fn from(config: SoftplusConfig) -> Self { + Self::Softplus(config) + } +} + +impl From for ActivationConfig { + fn from(config: EluConfig) -> Self { + Self::Elu(config) + } +} + +impl From for ActivationConfig { + fn from(config: CeluConfig) -> Self { + Self::Celu(config) + } +} + +impl From for ActivationConfig { + fn from(config: ThresholdedReluConfig) -> Self { + Self::ThresholdedRelu(config) + } +} + +impl From for ActivationConfig { + fn from(config: HardShrinkConfig) -> Self { + Self::HardShrink(config) + } +} + +impl From for ActivationConfig { + fn from(config: SoftShrinkConfig) -> Self { + Self::SoftShrink(config) + } +} + +impl From for ActivationConfig { + fn from(config: ShrinkConfig) -> Self { + Self::Shrink(config) + } +} + +impl ActivationConfig { + /// Initialize a wrapped activation layer. + pub fn init(&self, device: &Device) -> Activation { + match self { + ActivationConfig::Identity => Activation::Identity(Identity::new()), + ActivationConfig::Relu => Relu.into(), + ActivationConfig::LeakyRelu(conf) => conf.init().into(), + ActivationConfig::Gelu => Gelu::new().into(), + ActivationConfig::GeluApproximate => Gelu::new_approximate().into(), + ActivationConfig::PRelu(conf) => conf.init(device).into(), + ActivationConfig::SwiGlu(conf) => conf.init(device).into(), + ActivationConfig::HardSigmoid(conf) => conf.init().into(), + ActivationConfig::HardSwish => HardSwish.into(), + ActivationConfig::Softplus(conf) => conf.init().into(), + ActivationConfig::Selu => Selu.into(), + ActivationConfig::Sigmoid => Sigmoid.into(), + ActivationConfig::Tanh => Tanh.into(), + ActivationConfig::Softsign => Softsign.into(), + ActivationConfig::Elu(conf) => conf.init().into(), + ActivationConfig::Celu(conf) => conf.init().into(), + ActivationConfig::HardShrink(conf) => conf.init().into(), + ActivationConfig::SoftShrink(conf) => conf.init().into(), + ActivationConfig::Shrink(conf) => conf.init().into(), + ActivationConfig::ThresholdedRelu(conf) => conf.init().into(), + } + } +} + +/// Activation Layer Wrapper. +/// +/// Provides support for many in-built `burn::nn` activations. +#[derive(Module, Debug)] +#[non_exhaustive] +#[allow(clippy::large_enum_variant)] +pub enum Activation { + /// Identity activation layer. + Identity(Identity), + + /// [`Gelu`] activation layer. + Gelu(Gelu), + + /// [`PRelu`] activation layer. + PRelu(PRelu), + + /// [`Relu`] activation layer. + Relu(Relu), + + /// [`LeakyRelu`] activation layer. + LeakyRelu(LeakyRelu), + + /// [`SwiGlu`] activation layer. + SwiGlu(SwiGlu), + + /// [`Selu`] activation layer. + Selu(Selu), + + /// [`Sigmoid`] activation layer. + Sigmoid(Sigmoid), + + /// [`Tanh`] activation layer. + Tanh(Tanh), + + /// [`HardSigmoid`] activation layer. + HardSigmoid(HardSigmoid), + + /// [`HardSwish`] activation layer. + HardSwish(HardSwish), + + /// [`Softplus`] activation layer. + Softplus(Softplus), + + /// [`Softsign`] activation layer. + Softsign(Softsign), + + /// [`Elu`] activation layer. + Elu(Elu), + + /// [`Celu`] activation layer. + Celu(Celu), + + /// [`ThresholdedRelu`] activation layer. + ThresholdedRelu(ThresholdedRelu), + + /// [`HardShrink`] activation layer. + HardShrink(HardShrink), + + /// [`SoftShrink`] activation layer. + SoftShrink(SoftShrink), + + /// [`Shrink`] activation layer. + Shrink(Shrink), +} + +impl From for Activation { + fn from(layer: Identity) -> Self { + Self::Identity(layer) + } +} + +impl From for Activation { + fn from(layer: Gelu) -> Self { + Self::Gelu(layer) + } +} + +impl From for Activation { + fn from(layer: PRelu) -> Self { + Self::PRelu(layer) + } +} + +impl From for Activation { + fn from(layer: Relu) -> Self { + Self::Relu(layer) + } +} + +impl From for Activation { + fn from(layer: LeakyRelu) -> Self { + Self::LeakyRelu(layer) + } +} + +impl From for Activation { + fn from(layer: SwiGlu) -> Self { + Self::SwiGlu(layer) + } +} + +impl From for Activation { + fn from(layer: Selu) -> Self { + Self::Selu(layer) + } +} + +impl From for Activation { + fn from(layer: Sigmoid) -> Self { + Self::Sigmoid(layer) + } +} + +impl From for Activation { + fn from(layer: Tanh) -> Self { + Self::Tanh(layer) + } +} + +impl From for Activation { + fn from(layer: HardSigmoid) -> Self { + Self::HardSigmoid(layer) + } +} + +impl From for Activation { + fn from(layer: HardSwish) -> Self { + Self::HardSwish(layer) + } +} + +impl From for Activation { + fn from(layer: Softplus) -> Self { + Self::Softplus(layer) + } +} + +impl From for Activation { + fn from(layer: Softsign) -> Self { + Self::Softsign(layer) + } +} + +impl From for Activation { + fn from(layer: Elu) -> Self { + Self::Elu(layer) + } +} + +impl From for Activation { + fn from(layer: Celu) -> Self { + Self::Celu(layer) + } +} + +impl From for Activation { + fn from(layer: ThresholdedRelu) -> Self { + Self::ThresholdedRelu(layer) + } +} + +impl From for Activation { + fn from(layer: HardShrink) -> Self { + Self::HardShrink(layer) + } +} + +impl From for Activation { + fn from(layer: SoftShrink) -> Self { + Self::SoftShrink(layer) + } +} + +impl From for Activation { + fn from(layer: Shrink) -> Self { + Self::Shrink(layer) + } +} + +impl Activation { + /// Forward pass. + pub fn forward(&self, input: Tensor) -> Tensor { + match self { + Activation::Identity(layer) => layer.forward(input), + Activation::Relu(layer) => layer.forward(input), + Activation::LeakyRelu(layer) => layer.forward(input), + Activation::Gelu(layer) => layer.forward(input), + Activation::PRelu(layer) => layer.forward(input), + Activation::SwiGlu(layer) => layer.forward(input), + Activation::HardSigmoid(layer) => layer.forward(input), + Activation::HardSwish(layer) => layer.forward(input), + Activation::Softplus(layer) => layer.forward(input), + Activation::Selu(layer) => layer.forward(input), + Activation::Sigmoid(layer) => layer.forward(input), + Activation::Tanh(layer) => layer.forward(input), + Activation::Softsign(layer) => layer.forward(input), + Activation::Elu(layer) => layer.forward(input), + Activation::Celu(layer) => layer.forward(input), + Activation::ThresholdedRelu(layer) => layer.forward(input), + Activation::HardShrink(layer) => layer.forward(input), + Activation::SoftShrink(layer) => layer.forward(input), + Activation::Shrink(layer) => layer.forward(input), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::module::Module; + + fn make_input(device: &Device) -> Tensor<2> { + Tensor::from_data([[-1.0, -0.5, 0.0], [1.0, 0.5, 0.0]], device) + } + + fn expect_tensor(actual: Tensor, expected: Tensor) { + actual.to_data().assert_eq(&expected.to_data(), true); + } + + fn check_stateless_config_output( + config: ActivationConfig, + input: Tensor, + expected: Tensor, + device: &Device, + ) { + let act = config.init(device); + let output = act.forward(input); + expect_tensor(output, expected); + } + + #[test] + fn test_identity() { + let device = Default::default(); + let input = make_input(&device); + + let expected = input.clone(); + + check_stateless_config_output(ActivationConfig::Identity, input, expected, &device) + } + + #[test] + fn test_gelu() { + let device = Default::default(); + let input = make_input(&device); + + let expected = Gelu::new().forward(input.clone()); + + check_stateless_config_output(ActivationConfig::Gelu, input, expected, &device) + } + + #[test] + fn test_gelu_approximate() { + let device = Default::default(); + let input = make_input(&device); + + let expected = Gelu::new_approximate().forward(input.clone()); + + check_stateless_config_output(ActivationConfig::GeluApproximate, input, expected, &device) + } + + #[test] + fn test_prelu() { + let device = Default::default(); + let input = make_input(&device); + + let inner_config = PReluConfig::new(); + let expected = inner_config.init(&device).forward(input.clone()); + + check_stateless_config_output(inner_config.into(), input, expected, &device) + } + + #[test] + fn test_relu() { + let device = Default::default(); + let input = make_input(&device); + + let expected = Relu.forward(input.clone()); + + check_stateless_config_output(ActivationConfig::Relu, input, expected, &device) + } + + #[test] + fn test_leaky_relu() { + let device = Default::default(); + let input = make_input(&device); + + let inner_config = LeakyReluConfig::new(); + let expected = inner_config.init().forward(input.clone()); + + check_stateless_config_output(inner_config.into(), input, expected, &device) + } + + #[test] + fn test_swi_glu() { + let device = Default::default(); + let input = make_input(&device); + + let d_input = input.shape()[1]; + let d_output = 2 * d_input; + + let inner_config = SwiGluConfig::new(d_input, d_output); + let mut reference = inner_config.init(&device); + + let config: ActivationConfig = inner_config.into(); + let layer = config.init(&device); + + // Access tensors via forward pass to trigger lazy initialization, then clone weights. + let layer_output = layer.forward(input.clone()); + + match &layer { + Activation::SwiGlu(inner) => { + let state = inner.clone().into_record(); + reference = reference.load_record(state); + } + _ => unreachable!(), + }; + + expect_tensor(layer_output, reference.forward(input.clone())) + } + + #[test] + fn test_selu() { + let device = Default::default(); + let input = make_input(&device); + + let expected = Selu.forward(input.clone()); + + check_stateless_config_output(ActivationConfig::Selu, input, expected, &device) + } + + #[test] + fn test_sigmoid() { + let device = Default::default(); + let input = make_input(&device); + + let expected = Sigmoid.forward(input.clone()); + + check_stateless_config_output(ActivationConfig::Sigmoid, input, expected, &device) + } + + #[test] + fn test_tanh() { + let device = Default::default(); + let input = make_input(&device); + + let expected = Tanh.forward(input.clone()); + + check_stateless_config_output(ActivationConfig::Tanh, input, expected, &device) + } + + #[test] + fn test_hard_sigmoid() { + let device = Default::default(); + let input = make_input(&device); + + let inner_config = HardSigmoidConfig::new(); + let expected = inner_config.init().forward(input.clone()); + + check_stateless_config_output(inner_config.into(), input, expected, &device) + } + + #[test] + fn test_softsign() { + let device = Default::default(); + let input = make_input(&device); + + let expected = Softsign.forward(input.clone()); + + check_stateless_config_output(ActivationConfig::Softsign, input, expected, &device) + } + + #[test] + fn test_elu() { + let device = Default::default(); + let input = make_input(&device); + + let inner_config = EluConfig::new(); + let expected = inner_config.init().forward(input.clone()); + + check_stateless_config_output(inner_config.into(), input, expected, &device) + } + + #[test] + fn test_softplus() { + let device = Default::default(); + let input = make_input(&device); + + let inner_config = SoftplusConfig::new(); + let expected = inner_config.init().forward(input.clone()); + + check_stateless_config_output(inner_config.into(), input, expected, &device) + } + + #[test] + fn test_celu() { + let device = Default::default(); + let input = make_input(&device); + + let inner_config = CeluConfig::new(); + let expected = inner_config.init().forward(input.clone()); + + check_stateless_config_output(inner_config.into(), input, expected, &device) + } + + #[test] + fn test_thresholded_relu() { + let device = Default::default(); + let input = make_input(&device); + + let inner_config = ThresholdedReluConfig::new(); + let expected = inner_config.init().forward(input.clone()); + + check_stateless_config_output(inner_config.into(), input, expected, &device) + } + + #[test] + fn test_hard_shrink() { + let device = Default::default(); + let input = make_input(&device); + + let inner_config = HardShrinkConfig::new(); + let expected = inner_config.init().forward(input.clone()); + + check_stateless_config_output(inner_config.into(), input, expected, &device) + } + + #[test] + fn test_soft_shrink() { + let device = Default::default(); + let input = make_input(&device); + + let inner_config = SoftShrinkConfig::new(); + let expected = inner_config.init().forward(input.clone()); + + check_stateless_config_output(inner_config.into(), input, expected, &device) + } + + #[test] + fn test_shrink() { + let device = Default::default(); + let input = make_input(&device); + + let inner_config = ShrinkConfig::new(); + let expected = inner_config.init().forward(input.clone()); + + check_stateless_config_output(inner_config.into(), input, expected, &device) + } +} diff --git a/crates/burn-nn/src/activation/celu.rs b/crates/burn-nn/src/activation/celu.rs new file mode 100644 index 0000000..f0f99f7 --- /dev/null +++ b/crates/burn-nn/src/activation/celu.rs @@ -0,0 +1,101 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::tensor::activation::celu; + +/// CELU (Continuously Differentiable Exponential Linear Unit) layer. +/// +/// Applies the CELU function element-wise: +/// `celu(x) = max(0, x) + min(0, alpha * (exp(x / alpha) - 1))` +/// +/// Should be created with [CeluConfig](CeluConfig). +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Celu { + /// The alpha value for the CELU formulation. + pub alpha: f64, +} + +/// Configuration to create a [Celu](Celu) layer using the [init function](CeluConfig::init). +#[derive(Config, Debug)] +pub struct CeluConfig { + /// The alpha value for the CELU formulation. Default is 1.0 + #[config(default = "1.0")] + pub alpha: f64, +} + +impl CeluConfig { + /// Initialize a new [Celu](Celu) Layer + pub fn init(&self) -> Celu { + Celu { alpha: self.alpha } + } +} + +impl ModuleDisplay for Celu { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content.add("alpha", &self.alpha).optional() + } +} + +impl Celu { + /// Forward pass for the Celu layer. + /// + /// See [celu](burn::tensor::activation::celu) for more information. + /// + /// # Shapes + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + celu(input, self.alpha) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn test_celu_forward() { + let device = Default::default(); + let model = CeluConfig::new().init(); + let input = Tensor::<2>::from_data(TensorData::from([[0.5, -0.5, -1.0]]), &device); + let out = model.forward(input); + // celu(0.5, 1) = 0.5 + // celu(-0.5, 1) = 1 * (exp(-0.5) - 1) = -0.393469 + // celu(-1.0, 1) = 1 * (exp(-1) - 1) = -0.632121 + let expected = TensorData::from([[0.5, -0.393469, -0.632121]]); + out.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn test_celu_with_alpha() { + let device = Default::default(); + let model = CeluConfig::new().with_alpha(2.0).init(); + let input = Tensor::<2>::from_data(TensorData::from([[0.0, -2.0]]), &device); + let out = model.forward(input); + // celu(0, 2) = 0 + // celu(-2, 2) = 2 * (exp(-1) - 1) = -1.264241 + let expected = TensorData::from([[0.0, -1.264241]]); + out.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn display() { + let config = CeluConfig::new().init(); + assert_eq!(alloc::format!("{config}"), "Celu {alpha: 1}"); + } +} diff --git a/crates/burn-nn/src/activation/elu.rs b/crates/burn-nn/src/activation/elu.rs new file mode 100644 index 0000000..beb35ef --- /dev/null +++ b/crates/burn-nn/src/activation/elu.rs @@ -0,0 +1,82 @@ +use burn::config::Config; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; +use burn_core as burn; + +use burn::tensor::activation::elu; + +/// ELU (Exponential Linear Unit) layer. +/// +/// Should be created with [EluConfig](EluConfig). +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Elu { + /// The alpha value. + pub alpha: f64, +} +/// Configuration to create an [Elu](Elu) layer using the [init function](EluConfig::init). +#[derive(Config, Debug)] +pub struct EluConfig { + /// The alpha value. Default is 1.0 + #[config(default = "1.0")] + pub alpha: f64, +} +impl EluConfig { + /// Initialize a new [Elu](Elu) Layer + pub fn init(&self) -> Elu { + Elu { alpha: self.alpha } + } +} + +impl ModuleDisplay for Elu { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content.add("alpha", &self.alpha).optional() + } +} + +impl Elu { + /// Forward pass for the ELU layer. + /// + /// See [elu](burn::tensor::activation::elu) for more information. + /// + /// # Shapes + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + elu(input, self.alpha) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn test_elu_forward() { + let device = Default::default(); + let model = EluConfig::new().init(); + let input = Tensor::<2>::from_data(TensorData::from([[0.4410, -0.2507]]), &device); + let out = model.forward(input); + // elu(0.4410, 1.0) = 0.4410 + // elu(-0.2507, 1.0) = 1.0 * (exp(-0.2507) - 1) = -0.22186 + let expected = TensorData::from([[0.4410, -0.22186]]); + out.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn display() { + let config = EluConfig::new().init(); + assert_eq!(alloc::format!("{config}"), "Elu {alpha: 1}"); + } +} diff --git a/crates/burn-nn/src/activation/gelu.rs b/crates/burn-nn/src/activation/gelu.rs new file mode 100644 index 0000000..5a0f917 --- /dev/null +++ b/crates/burn-nn/src/activation/gelu.rs @@ -0,0 +1,78 @@ +use burn_core as burn; + +use burn::module::Module; +use burn::tensor::Tensor; + +/// Applies the Gaussian Error Linear Units function element-wise. +/// +/// See also [gelu](burn::tensor::activation::gelu) +/// +/// When `approximate` is true, uses the tanh approximation: +/// `0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))` +#[derive(Module, Debug, Default)] +pub struct Gelu { + /// Whether to use tanh approximation. + pub approximate: bool, +} + +impl Gelu { + /// Create the module with exact GELU. + pub fn new() -> Self { + Self::default() + } + + /// Create the module with tanh approximation. + pub fn new_approximate() -> Self { + Self { approximate: true } + } + + /// Applies the forward pass on the input tensor. + /// + /// # Shapes + /// + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + if self.approximate { + burn::tensor::activation::gelu_approximate(input) + } else { + burn::tensor::activation::gelu(input) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Tolerance; + + type FT = f32; + + #[test] + fn display() { + let layer = Gelu::new(); + + assert_eq!(alloc::format!("{layer}"), "Gelu {\n approximate: false\n}"); + } + + #[test] + fn forward_approximate() { + let device = Default::default(); + let input = Tensor::<2>::from_data([[-1.0, 0.0, 1.0], [0.5, -0.5, 2.0]], &device); + + let output = Gelu::new_approximate().forward(input); + + // PyTorch: torch.nn.functional.gelu(x, approximate="tanh") + let expected = Tensor::<2>::from_data( + [ + [-0.1588079929, 0.0000000000, 0.8411920071], + [0.3457140028, -0.1542859972, 1.9545977116], + ], + &device, + ); + + output + .into_data() + .assert_approx_eq::(&expected.into_data(), Tolerance::rel_abs(1e-5, 1e-5)); + } +} diff --git a/crates/burn-nn/src/activation/glu.rs b/crates/burn-nn/src/activation/glu.rs new file mode 100644 index 0000000..d596ff3 --- /dev/null +++ b/crates/burn-nn/src/activation/glu.rs @@ -0,0 +1,50 @@ +use burn_core as burn; + +use burn::module::Module; +use burn::tensor::Tensor; + +/// Applies the gated linear unit function. +/// +/// See also [glu](burn::tensor::activation::glu) +#[derive(Module, Debug, Default)] +pub struct GLU { + dim: usize, +} + +impl GLU { + /// Create the module. + /// + /// # Arguments + /// * `dim` - The dimension on which to split the input. + pub fn new(dim: usize) -> Self { + Self { dim } + } + + /// Applies the gated linear unit function. + /// + /// GLU(a,b)=a⊗σ(b) where `a` is the first half of the input matrices and `b` is the second half. + /// + /// **Note**: + /// * The size of the input tensor along `dim` must be divisible by 2. + /// + /// ### Arguments + /// * `tensor` - The input tensor. + /// + /// ### Returns + /// * A tensor with the same shape as the input, except the size along `dim` is halved. + pub fn forward(&self, input: Tensor) -> Tensor { + burn::tensor::activation::glu(input, self.dim) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn display() { + let layer = GLU::new(1); + + assert_eq!(alloc::format!("{layer}"), "GLU {\n dim: 1\n}"); + } +} diff --git a/crates/burn-nn/src/activation/hard_shrink.rs b/crates/burn-nn/src/activation/hard_shrink.rs new file mode 100644 index 0000000..fada3ea --- /dev/null +++ b/crates/burn-nn/src/activation/hard_shrink.rs @@ -0,0 +1,94 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::tensor::activation::hard_shrink; + +/// Hard Shrink layer. +/// +/// Applies the Hard Shrink function element-wise: +/// `hard_shrink(x) = x if |x| > lambda else 0` +/// +/// Should be created with [HardShrinkConfig](HardShrinkConfig). +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct HardShrink { + /// The lambda value for the Hard Shrink formulation. + pub lambda: f64, +} + +/// Configuration to create a [HardShrink](HardShrink) layer using the [init function](HardShrinkConfig::init). +#[derive(Config, Debug)] +pub struct HardShrinkConfig { + /// The lambda value for the Hard Shrink formulation. Default is 0.5 + #[config(default = "0.5")] + pub lambda: f64, +} + +impl HardShrinkConfig { + /// Initialize a new [HardShrink](HardShrink) Layer + pub fn init(&self) -> HardShrink { + HardShrink { + lambda: self.lambda, + } + } +} + +impl ModuleDisplay for HardShrink { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content.add("lambda", &self.lambda).optional() + } +} + +impl HardShrink { + /// Forward pass for the Hard Shrink layer. + /// + /// See [hard_shrink](burn::tensor::activation::hard_shrink) for more information. + /// + /// # Shapes + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + hard_shrink(input, self.lambda) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + + #[test] + fn test_hard_shrink_forward() { + let device = Default::default(); + let model = HardShrinkConfig::new().init(); + let input = Tensor::<2>::from_data([[0.5, -0.5, -1.0], [8.0, 0.3, 0.0]], &device); + let out = model.forward(input); + let expected = TensorData::from([[0.0_f32, 0.0, -1.0], [8.0, 0.0, 0.0]]); + assert_eq!(out.into_data(), expected); + } + + #[test] + fn test_hard_shrink_with_lambda() { + let device = Default::default(); + let model = HardShrinkConfig::new().with_lambda(0.2).init(); + let input = Tensor::<2>::from_data([[0.1, -0.1, -0.3], [0.5, 0.1, 0.0]], &device); + let out = model.forward(input); + let expected = TensorData::from([[0.0_f32, 0.0, -0.3], [0.5, 0.0, 0.0]]); + assert_eq!(out.into_data(), expected); + } + + #[test] + fn display() { + let config = HardShrinkConfig::new().init(); + assert_eq!(alloc::format!("{config}"), "HardShrink {lambda: 0.5}"); + } +} diff --git a/crates/burn-nn/src/activation/hard_sigmoid.rs b/crates/burn-nn/src/activation/hard_sigmoid.rs new file mode 100644 index 0000000..a28014a --- /dev/null +++ b/crates/burn-nn/src/activation/hard_sigmoid.rs @@ -0,0 +1,94 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::tensor::activation::hard_sigmoid; + +/// Hard Sigmoid layer. +/// +/// Should be created with [HardSigmoidConfig](HardSigmoidConfig). +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct HardSigmoid { + /// The alpha value. + pub alpha: f64, + /// The beta value. + pub beta: f64, +} +/// Configuration to create a [Hard Sigmoid](HardSigmoid) layer using the [init function](HardSigmoidConfig::init). +#[derive(Config, Debug)] +pub struct HardSigmoidConfig { + /// The alpha value. Default is 0.2 + #[config(default = "0.2")] + pub alpha: f64, + /// The beta value. Default is 0.5 + #[config(default = "0.5")] + pub beta: f64, +} +impl HardSigmoidConfig { + /// Initialize a new [Hard Sigmoid](HardSigmoid) Layer + pub fn init(&self) -> HardSigmoid { + HardSigmoid { + alpha: self.alpha, + beta: self.beta, + } + } +} + +impl ModuleDisplay for HardSigmoid { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("alpha", &self.alpha) + .add("beta", &self.beta) + .optional() + } +} + +impl HardSigmoid { + /// Forward pass for the Hard Sigmoid layer. + /// + /// See [hard_sigmoid](burn::tensor::activation::hard_sigmoid) for more information. + /// + /// # Shapes + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + hard_sigmoid(input, self.alpha, self.beta) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn test_hard_sigmoid_forward() { + let device = Default::default(); + let model = HardSigmoidConfig::new().init(); + let input = Tensor::<2>::from_data(TensorData::from([[0.4410, -0.2507]]), &device); + let out = model.forward(input); + let expected = TensorData::from([[0.5882, 0.44986]]); + out.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn display() { + let config = HardSigmoidConfig::new().init(); + assert_eq!( + alloc::format!("{config}"), + "HardSigmoid {alpha: 0.2, beta: 0.5}" + ); + } +} diff --git a/crates/burn-nn/src/activation/hard_swish.rs b/crates/burn-nn/src/activation/hard_swish.rs new file mode 100644 index 0000000..5db7e3e --- /dev/null +++ b/crates/burn-nn/src/activation/hard_swish.rs @@ -0,0 +1,53 @@ +use burn_core as burn; + +use burn::module::Module; +use burn::tensor::Tensor; +use burn::tensor::activation::hard_swish; + +/// Hard Swish layer. +#[derive(Module, Debug, Default)] +pub struct HardSwish; + +impl HardSwish { + /// Create the module. + pub fn new() -> Self { + Self + } + + /// Forward pass for the Hard Swish layer. + /// + /// See [hard_swish](burn::tensor::activation::hard_swish) for more information. + /// + /// # Shapes + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + hard_swish(input) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn test_hard_swish_forward() { + let device = Default::default(); + let model = HardSwish::new(); + + let input = Tensor::<2>::from_data(TensorData::from([[3.0f32, -3.0], [0.0, 1.0]]), &device); + let out = model.forward(input); + let expected = TensorData::from([[3.0f32, 0.0], [0.0, 0.6666667]]); + out.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn display() { + let layer = HardSwish::new(); + assert_eq!(alloc::format!("{layer}"), "HardSwish"); + } +} diff --git a/crates/burn-nn/src/activation/leaky_relu.rs b/crates/burn-nn/src/activation/leaky_relu.rs new file mode 100644 index 0000000..a424ced --- /dev/null +++ b/crates/burn-nn/src/activation/leaky_relu.rs @@ -0,0 +1,121 @@ +use burn::config::Config; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; +use burn_core as burn; + +use burn::tensor::activation::leaky_relu; + +/// Leaky ReLu layer. +/// +/// Should be created with [LeakyReluConfig](LeakyReluConfig). +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct LeakyRelu { + /// The negative slope. + pub negative_slope: f64, +} +/// Configuration to create a [Leaky Relu](LeakyRelu) layer using the [init function](LeakyReluConfig::init). +#[derive(Config, Debug)] +pub struct LeakyReluConfig { + /// The negative slope. Default is 0.01 + #[config(default = "0.01")] + pub negative_slope: f64, +} +impl LeakyReluConfig { + /// Initialize a new [Leaky Relu](LeakyRelu) Layer + pub fn init(&self) -> LeakyRelu { + LeakyRelu { + negative_slope: self.negative_slope, + } + } +} + +impl ModuleDisplay for LeakyRelu { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("negative_slope", &self.negative_slope) + .optional() + } +} + +impl LeakyRelu { + /// Forward pass for the Leaky ReLu layer. + /// + /// See [leaky_relu](burn::tensor::activation::leaky_relu) for more information. + /// + /// # Shapes + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + leaky_relu(input, self.negative_slope) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn test_leaky_relu_forward() { + let device = Default::default(); + let model: LeakyRelu = LeakyReluConfig::new().init(); + let input = Tensor::<2>::from_data(TensorData::from([[0.4410, -0.2507]]), &device); + let out = model.forward(input); + let expected = TensorData::from([[0.4410, -0.002507]]); + out.to_data().assert_eq(&expected, false); + } + #[test] + fn test_leaky_relu_forward_multi_dim() { + let input = [ + [ + [-1.0222, 1.5810, 0.3457, -1.3530], + [0.0231, 0.8681, 0.2473, -0.0377], + [0.3520, -1.1199, 1.2219, 0.2804], + ], + [ + [1.0002, 0.7259, 0.8779, 0.2084], + [1.5615, -0.1057, -0.4886, -1.5184], + [-0.5523, -0.2741, -0.0210, -1.1352], + ], + ]; + let expected = TensorData::from([ + [ + [-1.0222e-02, 1.5810e+00, 3.457e-01, -1.3530e-02], + [2.31e-02, 8.681e-01, 2.473e-01, -3.77e-04], + [3.52e-01, -1.1199e-02, 1.2219e+00, 2.804e-01], + ], + [ + [1.0002e+00, 7.259e-01, 8.779e-01, 2.084e-01], + [1.5615e+00, -1.057e-03, -4.886e-03, -1.5184e-02], + [-5.523e-03, -2.741e-03, -2.1e-04, -1.1352e-02], + ], + ]); + + let device = Default::default(); + let model = LeakyReluConfig::new().init(); + let input_data = Tensor::<3>::from_data(TensorData::from(input), &device); + let actual_output = model.forward(input_data); + actual_output + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()) + } + + #[test] + fn display() { + let config = LeakyReluConfig::new().init(); + assert_eq!( + alloc::format!("{config}"), + "LeakyRelu {negative_slope: 0.01}" + ); + } +} diff --git a/crates/burn-nn/src/activation/mod.rs b/crates/burn-nn/src/activation/mod.rs new file mode 100644 index 0000000..ced6eef --- /dev/null +++ b/crates/burn-nn/src/activation/mod.rs @@ -0,0 +1,68 @@ +//! # Activation Layers +//! +//! Users who desire a selectable activation function should +//! consider [`Activation`], which provides an abstraction over: +//! * [`Relu`] - the default, +//! * ['PRelu'] +//! * [`Gelu`] +//! * [`LeakyRelu`] +//! * [`SwiGlu`] +//! * [`Selu`] +//! * [`Sigmoid`] +//! * [`HardSigmoid`] +//! * [`HardSwish`] +//! * [`Softplus`] +//! * [`Softsign`] +//! * [`Tanh`] +//! * [`Elu`] +//! * [`Celu`] +//! * [`ThresholdedRelu`] +//! +//! The activation layer [`GLU`] has shape-changing behaviors +//! not compatible with the common API, and is not included +//! in the abstraction wrappers. + +mod activation_wrapper; + +// These are pub(crate) for dual-export in `nn` without re-exporting +// all of `nn.activation`, or manually listing each symbol. +pub(crate) mod celu; +pub(crate) mod elu; +pub(crate) mod gelu; +pub(crate) mod glu; +pub(crate) mod hard_shrink; +pub(crate) mod hard_sigmoid; +pub(crate) mod hard_swish; +pub(crate) mod leaky_relu; +pub(crate) mod prelu; +pub(crate) mod relu; +pub(crate) mod selu; +pub(crate) mod shrink; +pub(crate) mod sigmoid; +pub(crate) mod soft_shrink; +pub(crate) mod softplus; +pub(crate) mod softsign; +pub(crate) mod swiglu; +pub(crate) mod tanh; +pub(crate) mod thresholded_relu; + +pub use activation_wrapper::*; +pub use celu::*; +pub use elu::*; +pub use gelu::*; +pub use glu::*; +pub use hard_shrink::*; +pub use hard_sigmoid::*; +pub use hard_swish::*; +pub use leaky_relu::*; +pub use prelu::*; +pub use relu::*; +pub use selu::*; +pub use shrink::*; +pub use sigmoid::*; +pub use soft_shrink::*; +pub use softplus::*; +pub use softsign::*; +pub use swiglu::*; +pub use tanh::*; +pub use thresholded_relu::*; diff --git a/crates/burn-nn/src/activation/prelu.rs b/crates/burn-nn/src/activation/prelu.rs new file mode 100644 index 0000000..62e88a2 --- /dev/null +++ b/crates/burn-nn/src/activation/prelu.rs @@ -0,0 +1,87 @@ +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Initializer, Module, ModuleDisplay, Param}; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn_core as burn; + +/// Parametric Relu layer. +/// +/// Should be created using [PReluConfig] +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct PRelu { + /// the weights learnt for PReLu. can be of shape \[1\] or \[num_parameters\] in which case it must + /// be the same as number of channels in the input tensor + pub alpha: Param>, + + /// Alpha value for the PRelu layer + pub alpha_value: f64, +} + +impl ModuleDisplay for PRelu { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let [num_parameters] = self.alpha.shape().dims(); + + content + .add("num_parameters", &num_parameters) + .add("alpha_value", &self.alpha_value) + .optional() + } +} + +/// Configuration to create a [Parametric Relu](PRelu) layer using the [init function](PReluConfig::init). +#[derive(Config, Debug)] +pub struct PReluConfig { + /// The number of parameters. + #[config(default = "1")] + pub num_parameters: usize, + /// The learnable weight alpha. Default is 0.25 + #[config(default = "0.25")] + pub alpha: f64, +} + +impl PReluConfig { + /// Initialize a new [Parametric Relu](PRelu) Layer + pub fn init(&self, device: &Device) -> PRelu { + PRelu { + // alpha is a tensor of length num_parameters + alpha: Initializer::Constant { value: self.alpha }.init([self.num_parameters], device), + alpha_value: self.alpha, + } + } +} + +impl PRelu { + /// Applies the forward pass on the input tensor. + /// + /// # Shapes + /// + /// - input: `[..., any]` + /// - output: `[..., any]` + /// + /// See also [prelu](burn::tensor::activation::prelu) for more information. + pub fn forward(&self, input: Tensor) -> Tensor { + burn::tensor::activation::prelu(input, self.alpha.val()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn display() { + let layer = PReluConfig::new().init(&Default::default()); + + assert_eq!( + alloc::format!("{layer}"), + "PRelu {num_parameters: 1, alpha_value: 0.25, params: 1}" + ); + } +} diff --git a/crates/burn-nn/src/activation/relu.rs b/crates/burn-nn/src/activation/relu.rs new file mode 100644 index 0000000..16537de --- /dev/null +++ b/crates/burn-nn/src/activation/relu.rs @@ -0,0 +1,38 @@ +use burn_core as burn; + +use burn::module::Module; +use burn::tensor::Tensor; + +/// Applies the rectified linear unit function element-wise +/// See also [relu](burn::tensor::activation::relu) +/// +#[derive(Module, Debug, Default)] +pub struct Relu; + +impl Relu { + /// Create the module. + pub fn new() -> Self { + Self {} + } + /// Applies the forward pass on the input tensor. + /// + /// # Shapes + /// + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + burn::tensor::activation::relu(input) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn display() { + let layer = Relu::new(); + + assert_eq!(alloc::format!("{layer}"), "Relu"); + } +} diff --git a/crates/burn-nn/src/activation/selu.rs b/crates/burn-nn/src/activation/selu.rs new file mode 100644 index 0000000..55bc617 --- /dev/null +++ b/crates/burn-nn/src/activation/selu.rs @@ -0,0 +1,37 @@ +use burn_core as burn; + +use burn::module::Module; +use burn::tensor::Tensor; + +/// Applies the Scaled Exponential Linear Unit function element-wise. +/// See also [selu](burn::tensor::activation::selu) +#[derive(Module, Debug, Default)] +pub struct Selu; + +impl Selu { + /// Create the module. + pub fn new() -> Self { + Self {} + } + /// Applies the forward pass on the input tensor. + /// + /// # Shapes + /// + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + burn::tensor::activation::selu(input) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn display() { + let layer = Selu::new(); + + assert_eq!(alloc::format!("{layer}"), "Selu"); + } +} diff --git a/crates/burn-nn/src/activation/shrink.rs b/crates/burn-nn/src/activation/shrink.rs new file mode 100644 index 0000000..6f6fbd5 --- /dev/null +++ b/crates/burn-nn/src/activation/shrink.rs @@ -0,0 +1,110 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::tensor::activation::shrink; + +/// Shrink layer. +/// +/// Applies the Shrink function element-wise: +/// `shrink(x) = x - bias if x > lambda, x + bias if x < -lambda, 0 otherwise` +/// +/// Should be created with [ShrinkConfig](ShrinkConfig). +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Shrink { + /// The lambda value for the Shrink formulation. + pub lambda: f64, + /// The bias value for the Shrink formulation. + // Usually bias = lambda, but need this to handle onnx spec https://onnx.ai/onnx/operators/onnx__Shrink.html + pub bias: f64, +} + +/// Configuration to create a [Shrink](Shrink) layer using the [init function](ShrinkConfig::init). +#[derive(Config, Debug)] +pub struct ShrinkConfig { + /// The lambda value for the Shrink formulation. Default is 0.5 + #[config(default = "0.5")] + pub lambda: f64, + /// The bias value for the Shrink formulation. Default is 0.5. + #[config(default = "0.5")] + pub bias: f64, +} + +impl ShrinkConfig { + /// Initialize a new [Shrink](Shrink) Layer + pub fn init(&self) -> Shrink { + Shrink { + lambda: self.lambda, + bias: self.bias, + } + } +} + +impl ModuleDisplay for Shrink { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("lambda", &self.lambda) + .add("bias", &self.bias) + .optional() + } +} + +impl Shrink { + /// Forward pass for the Shrink layer. + /// + /// See [shrink](burn::tensor::activation::shrink) for more information. + /// + /// # Shapes + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + shrink(input, self.lambda, self.bias) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + + #[test] + fn test_shrink_forward() { + let device = Default::default(); + let model = ShrinkConfig::new().init(); + let input = Tensor::<2>::from_data([[0.5, -0.5, -1.0], [8.0, 0.3, 0.0]], &device); + let out = model.forward(input); + let expected = TensorData::from([[0.0_f32, 0.0, -0.5], [7.5, 0.0, 0.0]]); + assert_eq!(out.into_data(), expected); + } + + #[test] + fn test_shrink_with_lambda_and_bias() { + let device = Default::default(); + let model = ShrinkConfig::new() + .with_lambda(0.25) + .with_bias(0.125) + .init(); + let input = Tensor::<2>::from_data([[0.125, -0.125, -0.5], [0.75, 0.1, 0.0]], &device); + let out = model.forward(input); + let expected = TensorData::from([[0.0_f32, 0.0, -0.375], [0.625, 0.0, 0.0]]); + assert_eq!(out.into_data(), expected); + } + + #[test] + fn display() { + let config = ShrinkConfig::new().init(); + assert_eq!( + alloc::format!("{config}"), + "Shrink {lambda: 0.5, bias: 0.5}" + ); + } +} diff --git a/crates/burn-nn/src/activation/sigmoid.rs b/crates/burn-nn/src/activation/sigmoid.rs new file mode 100644 index 0000000..ce9d517 --- /dev/null +++ b/crates/burn-nn/src/activation/sigmoid.rs @@ -0,0 +1,37 @@ +use burn_core as burn; + +use burn::module::Module; +use burn::tensor::Tensor; + +/// Applies the sigmoid function element-wise +/// See also [sigmoid](burn::tensor::activation::sigmoid) +#[derive(Module, Debug, Default)] +pub struct Sigmoid; + +impl Sigmoid { + /// Create the module. + pub fn new() -> Self { + Self {} + } + /// Applies the forward pass on the input tensor. + /// + /// # Shapes + /// + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + burn::tensor::activation::sigmoid(input) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn display() { + let layer = Sigmoid::new(); + + assert_eq!(alloc::format!("{layer}"), "Sigmoid"); + } +} diff --git a/crates/burn-nn/src/activation/soft_shrink.rs b/crates/burn-nn/src/activation/soft_shrink.rs new file mode 100644 index 0000000..255329b --- /dev/null +++ b/crates/burn-nn/src/activation/soft_shrink.rs @@ -0,0 +1,94 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::tensor::activation::soft_shrink; + +/// Soft Shrink layer. +/// +/// Applies the Soft Shrink function element-wise: +/// `soft_shrink(x) = x - lambda if x > lambda, x + lambda if x < -lambda, 0 otherwise` +/// +/// Should be created with [SoftShrinkConfig](SoftShrinkConfig). +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct SoftShrink { + /// The lambda value for the Soft Shrink formulation. + pub lambda: f64, +} + +/// Configuration to create a [SoftShrink](SoftShrink) layer using the [init function](SoftShrinkConfig::init). +#[derive(Config, Debug)] +pub struct SoftShrinkConfig { + /// The lambda value for the Soft Shrink formulation. Default is 0.5 + #[config(default = "0.5")] + pub lambda: f64, +} + +impl SoftShrinkConfig { + /// Initialize a new [SoftShrink](SoftShrink) Layer + pub fn init(&self) -> SoftShrink { + SoftShrink { + lambda: self.lambda, + } + } +} + +impl ModuleDisplay for SoftShrink { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content.add("lambda", &self.lambda).optional() + } +} + +impl SoftShrink { + /// Forward pass for the Soft Shrink layer. + /// + /// See [soft_shrink](burn::tensor::activation::soft_shrink) for more information. + /// + /// # Shapes + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + soft_shrink(input, self.lambda) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + + #[test] + fn test_soft_shrink_forward() { + let device = Default::default(); + let model = SoftShrinkConfig::new().init(); + let input = Tensor::<2>::from_data([[0.5, -0.5, -1.0], [8.0, 0.3, 0.0]], &device); + let out = model.forward(input); + let expected = TensorData::from([[0.0_f32, 0.0, -0.5], [7.5, 0.0, 0.0]]); + assert_eq!(out.into_data(), expected); + } + + #[test] + fn test_soft_shrink_with_lambda() { + let device = Default::default(); + let model = SoftShrinkConfig::new().with_lambda(0.25).init(); + let input = Tensor::<2>::from_data([[0.125, -0.125, -0.5], [0.75, 0.1, 0.0]], &device); + let out = model.forward(input); + let expected = TensorData::from([[0.0_f32, 0.0, -0.25], [0.5, 0.0, 0.0]]); + assert_eq!(out.into_data(), expected); + } + + #[test] + fn display() { + let config = SoftShrinkConfig::new().init(); + assert_eq!(alloc::format!("{config}"), "SoftShrink {lambda: 0.5}"); + } +} diff --git a/crates/burn-nn/src/activation/softplus.rs b/crates/burn-nn/src/activation/softplus.rs new file mode 100644 index 0000000..5358f66 --- /dev/null +++ b/crates/burn-nn/src/activation/softplus.rs @@ -0,0 +1,102 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::tensor::activation::softplus; + +/// Softplus layer. +/// +/// Applies the softplus function element-wise: +/// `softplus(x) = (1/beta) * log(1 + exp(beta * x))` +/// +/// Should be created with [SoftplusConfig](SoftplusConfig). +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Softplus { + /// The beta value. + pub beta: f64, +} + +/// Configuration to create a [Softplus](Softplus) layer using the [init function](SoftplusConfig::init). +#[derive(Config, Debug)] +pub struct SoftplusConfig { + /// The beta value. Default is 1.0 + #[config(default = "1.0")] + pub beta: f64, +} + +impl SoftplusConfig { + /// Initialize a new [Softplus](Softplus) Layer + pub fn init(&self) -> Softplus { + Softplus { beta: self.beta } + } +} + +impl ModuleDisplay for Softplus { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content.add("beta", &self.beta).optional() + } +} + +impl Softplus { + /// Forward pass for the Softplus layer. + /// + /// See [softplus](burn::tensor::activation::softplus) for more information. + /// + /// # Shapes + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + softplus(input, self.beta) + } +} + +#[cfg(test)] +#[allow(clippy::approx_constant)] +mod tests { + use super::*; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn test_softplus_forward() { + let device = Default::default(); + let model = SoftplusConfig::new().init(); + let input = Tensor::<2>::from_data(TensorData::from([[0.0, 1.0, -1.0]]), &device); + let out = model.forward(input); + // softplus(0) = log(2) ≈ 0.6931 + // softplus(1) = log(1 + e) ≈ 1.3133 + // softplus(-1) = log(1 + e^-1) ≈ 0.3133 + let expected = TensorData::from([[0.6931, 1.3133, 0.3133]]); + out.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn test_softplus_with_beta() { + let device = Default::default(); + let model = SoftplusConfig::new().with_beta(2.0).init(); + let input = Tensor::<2>::from_data(TensorData::from([[0.0, 1.0]]), &device); + let out = model.forward(input); + // softplus(0, beta=2) = (1/2) * log(1 + exp(0)) = 0.5 * log(2) ≈ 0.3466 + // softplus(1, beta=2) = (1/2) * log(1 + exp(2)) = 0.5 * log(8.389) ≈ 1.0635 + let expected = TensorData::from([[0.3466, 1.0635]]); + out.to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn display() { + let config = SoftplusConfig::new().init(); + assert_eq!(alloc::format!("{config}"), "Softplus {beta: 1}"); + } +} diff --git a/crates/burn-nn/src/activation/softsign.rs b/crates/burn-nn/src/activation/softsign.rs new file mode 100644 index 0000000..3a9ed3d --- /dev/null +++ b/crates/burn-nn/src/activation/softsign.rs @@ -0,0 +1,37 @@ +use burn_core as burn; + +use burn::module::Module; +use burn::tensor::Tensor; + +/// Applies the softsign function element-wise +/// See also [softsign](burn::tensor::activation::softsign) +#[derive(Module, Debug, Default)] +pub struct Softsign; + +impl Softsign { + /// Create the module. + pub fn new() -> Self { + Self {} + } + /// Applies the forward pass on the input tensor. + /// + /// # Shapes + /// + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + burn::tensor::activation::softsign(input) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn display() { + let layer = Softsign::new(); + + assert_eq!(alloc::format!("{layer}"), "Softsign"); + } +} diff --git a/crates/burn-nn/src/activation/swiglu.rs b/crates/burn-nn/src/activation/swiglu.rs new file mode 100644 index 0000000..4f7b8b8 --- /dev/null +++ b/crates/burn-nn/src/activation/swiglu.rs @@ -0,0 +1,150 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Initializer, Module, ModuleDisplay}; +use burn::tensor::activation::silu; +use burn::tensor::{Device, Tensor}; + +use crate::{Linear, LinearConfig, LinearLayout}; + +/// Configuration to create a [SwiGlu](SwiGlu) activation layer using the [init function](SwiGluConfig::init). +#[derive(Config, Debug)] +pub struct SwiGluConfig { + /// The size of the input features. + pub d_input: usize, + /// The size of the output features. + pub d_output: usize, + /// If a bias should be applied during the linear transformation. Default behaviour is False + /// for SwiGLU activation implementations. + #[config(default = false)] + pub bias: bool, + /// The type of function used to initialize the linear layer parameters + #[config( + default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0), fan_out_only:false}" + )] + pub initializer: Initializer, + /// The layout in which the linear parameters are stored. + #[config(default = "LinearLayout::Row")] + pub layout: LinearLayout, +} + +/// Applies the SwiGLU or Swish Gated Linear Unit to the input tensor. +/// The SwiGLU activation function is defined as: +/// `SwiGLU(x) = Swish(W_inner * x + b_inner) * (W_outer * x + b_outer)` +/// +/// Should be created with [SwiGluConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct SwiGlu { + /// The inner linear layer for Swish activation function + /// with `d_input` input features and `d_output` output features. + pub linear_inner: Linear, + /// The outer linear layer for element wise multiplication + /// with `d_input` input features and `d_output` output features. + pub linear_outer: Linear, +} + +impl ModuleDisplay for SwiGlu { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let [d_input, d_output] = self.linear_inner.weight.shape().dims(); + content + .add("d_input", &d_input) + .add("d_output", &d_output) + .add("bias", &self.linear_inner.bias.is_some()) + .optional() + } +} + +impl SwiGluConfig { + /// Initialize a new [SwiGLU](SwiGlu) activation layer. + pub fn init(&self, device: &Device) -> SwiGlu { + SwiGlu { + linear_inner: LinearConfig::new(self.d_input, self.d_output) + .with_bias(self.bias) + .with_initializer(self.initializer.clone()) + .with_layout(self.layout) + .init(device), + linear_outer: LinearConfig::new(self.d_input, self.d_output) + .with_bias(self.bias) + .with_initializer(self.initializer.clone()) + .with_layout(self.layout) + .init(device), + } + } +} + +impl SwiGlu { + /// Applies the Swish Gated Linear Unit to the input tensor. + /// + /// # Shapes + /// + /// - input: `[batch_size, seq_length, d_input]` + /// - output: `[batch_size, seq_length, d_output]` + pub fn forward(&self, input: Tensor) -> Tensor { + let x = self.linear_inner.forward(input.clone()); + let x = silu(x); + x.mul(self.linear_outer.forward(input)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn test_swiglu_forward_no_bias() { + let device = Device::default(); + device.seed(0); + + let config = SwiGluConfig::new(3, 3).with_initializer(Initializer::Constant { value: 0.5 }); + let swiglu = config.init(&device); + let input = Tensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + let output = swiglu.forward(input); + let expected_output = Tensor::<2>::from_data( + [[8.5732, 8.5732, 8.5732], [56.2189, 56.2189, 56.2189]], + &device, + ); + output + .to_data() + .assert_approx_eq::(&expected_output.to_data(), Tolerance::default()); + } + + #[test] + fn test_swiglu_forward_with_bias() { + let device = Device::default(); + device.seed(0); + + let config = SwiGluConfig::new(3, 3) + .with_bias(true) + .with_initializer(Initializer::Constant { value: 0.5 }); + let swiglu = config.init(&device); + let input = Tensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + let output = swiglu.forward(input); + let expected_output = Tensor::<2>::from_data( + [[11.8909, 11.8909, 11.8909], [63.9785, 63.9785, 63.9785]], + &device, + ); + output + .to_data() + .assert_approx_eq::(&expected_output.to_data(), Tolerance::default()); + } + + #[test] + fn display() { + let config = SwiGluConfig::new(3, 5); + let swiglu = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{swiglu}"), + "SwiGlu {d_input: 3, d_output: 5, bias: false, params: 30}" + ); + } +} diff --git a/crates/burn-nn/src/activation/tanh.rs b/crates/burn-nn/src/activation/tanh.rs new file mode 100644 index 0000000..cee074a --- /dev/null +++ b/crates/burn-nn/src/activation/tanh.rs @@ -0,0 +1,37 @@ +use burn_core as burn; + +use burn::module::Module; +use burn::tensor::Tensor; + +/// Applies the tanh activation function element-wise +/// See also [tanh](burn::tensor::activation::tanh) +#[derive(Module, Debug, Default)] +pub struct Tanh; + +impl Tanh { + /// Create the module. + pub fn new() -> Self { + Self {} + } + /// Applies the forward pass on the input tensor. + /// + /// # Shapes + /// + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + burn::tensor::activation::tanh(input) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn display() { + let layer = Tanh::new(); + + assert_eq!(alloc::format!("{layer}"), "Tanh"); + } +} diff --git a/crates/burn-nn/src/activation/thresholded_relu.rs b/crates/burn-nn/src/activation/thresholded_relu.rs new file mode 100644 index 0000000..6e6f31a --- /dev/null +++ b/crates/burn-nn/src/activation/thresholded_relu.rs @@ -0,0 +1,79 @@ +use burn::config::Config; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; +use burn_core as burn; + +use burn::tensor::activation::thresholded_relu; + +/// Thresholded ReLU layer. +/// +/// Should be created with [ThresholdedReluConfig](ThresholdedReluConfig). +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct ThresholdedRelu { + /// The alpha threshold. + pub alpha: f64, +} + +/// Configuration to create a [ThresholdedRelu](ThresholdedRelu) layer using the [init function](ThresholdedReluConfig::init). +#[derive(Config, Debug)] +pub struct ThresholdedReluConfig { + /// The alpha threshold. Default is 1.0 + #[config(default = "1.0")] + pub alpha: f64, +} + +impl ThresholdedReluConfig { + /// Initialize a new [ThresholdedRelu](ThresholdedRelu) layer. + pub fn init(&self) -> ThresholdedRelu { + ThresholdedRelu { alpha: self.alpha } + } +} + +impl ModuleDisplay for ThresholdedRelu { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content.add("alpha", &self.alpha).optional() + } +} + +impl ThresholdedRelu { + /// Forward pass for the Thresholded ReLU layer. + /// + /// See [thresholded_relu](burn::tensor::activation::thresholded_relu) for more information. + /// + /// # Shapes + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + thresholded_relu(input, self.alpha) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + + #[test] + fn test_thresholded_relu_forward() { + let device = Default::default(); + let model = ThresholdedReluConfig::new().init(); + let input = Tensor::<2>::from_data(TensorData::from([[0.5, 1.5, -0.2]]), &device); + let out = model.forward(input); + let expected = TensorData::from([[0.0, 1.5, 0.0]]); + out.to_data().assert_eq(&expected, false); + } + + #[test] + fn display() { + let config = ThresholdedReluConfig::new().init(); + assert_eq!(alloc::format!("{config}"), "ThresholdedRelu {alpha: 1}"); + } +} diff --git a/crates/burn-nn/src/lib.rs b/crates/burn-nn/src/lib.rs new file mode 100644 index 0000000..d248e12 --- /dev/null +++ b/crates/burn-nn/src/lib.rs @@ -0,0 +1,29 @@ +#![cfg_attr(not(feature = "std"), no_std)] +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] +#![recursion_limit = "256"] + +//! Burn neural network module. + +/// Loss module +pub mod loss; + +/// Neural network modules implementations. +pub mod modules; +pub use modules::*; + +pub mod activation; +pub use activation::{ + celu::*, elu::*, gelu::*, glu::*, hard_shrink::*, hard_sigmoid::*, leaky_relu::*, prelu::*, + relu::*, selu::*, shrink::*, sigmoid::*, soft_shrink::*, softplus::*, softsign::*, swiglu::*, + tanh::*, thresholded_relu::*, +}; + +mod padding; + +pub use padding::*; + +// For backward compat, `burn::nn::Initializer` +pub use burn_core::module::Initializer; + +extern crate alloc; diff --git a/crates/burn-nn/src/loss/binary_cross_entropy.rs b/crates/burn-nn/src/loss/binary_cross_entropy.rs new file mode 100644 index 0000000..40ca7ff --- /dev/null +++ b/crates/burn-nn/src/loss/binary_cross_entropy.rs @@ -0,0 +1,409 @@ +use burn_core as burn; + +use alloc::vec::Vec; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::activation::log_sigmoid; +use burn::tensor::{Device, Int, Tensor}; +use burn::{config::Config, module::Module}; + +/// Configuration to create a [Binary Cross-entropy loss](BinaryCrossEntropyLoss) using the [init function](BinaryCrossEntropyLossConfig::init). +#[derive(Config, Debug)] +pub struct BinaryCrossEntropyLossConfig { + /// Create weighted binary cross-entropy with a weight for each class. + /// + /// The loss of a specific sample will simply be multiplied by its label weight. + pub weights: Option>, + + /// Create binary cross-entropy with label smoothing according to [When Does Label Smoothing Help?](https://arxiv.org/abs/1906.02629). + /// + /// Hard labels {0, 1} will be changed to `y_smoothed = y(1 - a) + a / num_classes`. + /// Alpha = 0 would be the same as default. + pub smoothing: Option, + + /// Treat the inputs as logits, applying a sigmoid activation when computing the loss. + #[config(default = false)] + pub logits: bool, +} + +impl BinaryCrossEntropyLossConfig { + /// Initialize [Binary Cross-entropy loss](BinaryCrossEntropyLoss). + pub fn init(&self, device: &Device) -> BinaryCrossEntropyLoss { + self.assertions(); + BinaryCrossEntropyLoss { + weights: self + .weights + .as_ref() + .map(|e| Tensor::<1>::from_floats(e.as_slice(), device)), + smoothing: self.smoothing, + logits: self.logits, + } + } + + fn assertions(&self) { + if let Some(alpha) = self.smoothing { + assert!( + (0.0..=1.).contains(&alpha), + "Alpha of Cross-entropy loss with smoothed labels should be in interval [0, 1]. Got {alpha}" + ); + }; + if let Some(weights) = self.weights.as_ref() { + assert!( + weights.iter().all(|e| e > &0.), + "Weights of cross-entropy have to be positive." + ); + } + } +} + +/// Calculate the binary cross entropy loss from the input logits and the targets. +/// +/// Should be created using [BinaryCrossEntropyLossConfig] +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct BinaryCrossEntropyLoss { + /// Weights for cross-entropy. + pub weights: Option>, + /// Label smoothing alpha. + pub smoothing: Option, + /// Treat the inputs as logits + pub logits: bool, +} + +impl ModuleDisplay for BinaryCrossEntropyLoss { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("weights", &self.weights) + .add("smoothing", &self.smoothing) + .add("logits", &self.logits) + .optional() + } +} + +impl BinaryCrossEntropyLoss { + /// Compute the criterion on the input tensor. + /// + /// # Shapes + /// + /// Binary: + /// - logits: `[batch_size]` + /// - targets: `[batch_size]` + /// + /// Multi-label: + /// - logits: `[batch_size, num_classes]` + /// - targets: `[batch_size, num_classes]` + pub fn forward(&self, logits: Tensor, targets: Tensor) -> Tensor<1> { + self.assertions(&logits, &targets); + + let mut targets_float = targets.clone().float(); + let shape = targets.dims(); + + if let Some(alpha) = self.smoothing { + let num_classes = if D > 1 { shape[D - 1] } else { 2 }; + targets_float = targets_float * (1. - alpha) + alpha / num_classes as f32; + } + + let mut loss = if self.logits { + // Numerically stable by combining `log(sigmoid(x))` with `log_sigmoid(x)` + (targets_float.neg() + 1.) * logits.clone() - log_sigmoid(logits) + } else { + // - (target * log(input) + (1 - target) * log(1 - input)) + // https://github.com/tracel-ai/burn/issues/2739: clamp at -100.0 to avoid undefined values + (targets_float.clone() - 1) * logits.clone().neg().log1p().clamp_min(-100.0) + - targets_float * logits.log().clamp_min(-100.0) + }; + + if let Some(weights) = &self.weights { + let weights = if D > 1 { + weights.clone().expand(shape) + } else { + // Flatten targets and expand resulting weights to make it compatible with + // Tensor for binary 1-D case + weights + .clone() + .gather(0, targets.flatten(0, 0)) + .expand(shape) + }; + loss = loss * weights; + } + + loss.mean() + } + + fn assertions(&self, logits: &Tensor, targets: &Tensor) { + let logits_dims = logits.dims(); + let targets_dims = targets.dims(); + assert!( + logits_dims == targets_dims, + "Shape of targets ({targets_dims:?}) should correspond to outer shape of logits ({logits_dims:?})." + ); + + if let Some(weights) = &self.weights + && D > 1 + { + let targets_classes = targets_dims[D - 1]; + let weights_classes = weights.dims()[0]; + assert!( + weights_classes == targets_classes, + "The number of classes ({weights_classes}) does not match the weights provided ({targets_classes})." + ); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Tolerance; + use burn::tensor::{TensorData, activation::sigmoid}; + type FT = f32; + + #[test] + fn test_binary_cross_entropy_preds_all_correct() { + let device = Default::default(); + let preds = Tensor::<1>::from_floats([1.0, 0.0, 1.0, 0.0], &device); + let targets = Tensor::<1, Int>::from_data(TensorData::from([1, 0, 1, 0]), &device); + + let loss_actual = BinaryCrossEntropyLossConfig::new() + .init(&device) + .forward(preds, targets) + .into_data(); + + let loss_expected = TensorData::from([0.000]); + loss_actual.assert_approx_eq::(&loss_expected, Tolerance::default()); + } + + #[test] + fn test_binary_cross_entropy_preds_all_incorrect() { + let device = Default::default(); + let preds = Tensor::<1>::from_floats([0.0, 1.0, 0.0, 1.0], &device); + let targets = Tensor::<1, Int>::from_data(TensorData::from([1, 0, 1, 0]), &device); + + let loss_actual = BinaryCrossEntropyLossConfig::new() + .init(&device) + .forward(preds, targets) + .into_data(); + + let loss_expected = TensorData::from([100.000]); // clamped value + loss_actual.assert_approx_eq::(&loss_expected, Tolerance::default()); + } + + #[test] + fn test_binary_cross_entropy() { + // import torch + // from torch import nn + // input = torch.tensor([0.8271, 0.9626, 0.3796, 0.2355]) + // target = torch.tensor([0., 1., 0., 1.]) + // loss = nn.BCELoss() + // sigmoid = nn.Sigmoid() + // out = loss(sigmoid(input), target) # tensor(0.7491) + + let device = Default::default(); + let logits = Tensor::<1>::from_floats([0.8271, 0.9626, 0.3796, 0.2355], &device); + let targets = Tensor::<1, Int>::from_data(TensorData::from([0, 1, 0, 1]), &device); + + let loss_actual = BinaryCrossEntropyLossConfig::new() + .init(&device) + .forward(sigmoid(logits), targets) + .into_data(); + + let loss_expected = TensorData::from([0.7491]); + loss_actual.assert_approx_eq::(&loss_expected, Tolerance::relative(1e-4)); + } + + #[test] + fn test_binary_cross_entropy_with_logits() { + let device = Default::default(); + let logits = Tensor::<1>::from_floats([0.8271, 0.9626, 0.3796, 0.2355], &device); + let targets = Tensor::<1, Int>::from_data(TensorData::from([0, 1, 0, 1]), &device); + + let loss_actual = BinaryCrossEntropyLossConfig::new() + .with_logits(true) + .init(&device) + .forward(logits, targets) + .into_data(); + + let loss_expected = TensorData::from([0.7491]); + loss_actual.assert_approx_eq::(&loss_expected, Tolerance::relative(1e-4)); + } + + #[test] + fn test_binary_cross_entropy_with_weights() { + // import torch + // from torch import nn + // input = torch.tensor([0.8271, 0.9626, 0.3796, 0.2355]) + // target = torch.tensor([0, 1, 0, 1]) + // weights = torch.tensor([3., 7.]).gather(0, target) + // loss = nn.BCELoss(weights) + // sigmoid = nn.Sigmoid() + // out = loss(sigmoid(input), target.float()) # tensor(3.1531) + + let device = Default::default(); + let logits = Tensor::<1>::from_floats([0.8271, 0.9626, 0.3796, 0.2355], &device); + let targets = Tensor::<1, Int>::from_data(TensorData::from([0, 1, 0, 1]), &device); + let weights = [3., 7.]; + + let loss_actual = BinaryCrossEntropyLossConfig::new() + .with_weights(Some(weights.to_vec())) + .init(&device) + .forward(sigmoid(logits), targets) + .into_data(); + + let loss_expected = TensorData::from([3.1531]); + loss_actual.assert_approx_eq::(&loss_expected, Tolerance::relative(1e-4)); + } + + #[test] + fn test_binary_cross_entropy_with_smoothing() { + // import torch + // from torch import nn + // input = torch.tensor([0.8271, 0.9626, 0.3796, 0.2355]) + // target = torch.tensor([0., 1., 0., 1.]) + // target_smooth = target * (1 - 0.1) + (0.1 / 2) + // loss = nn.BCELoss() + // sigmoid = nn.Sigmoid() + // out = loss(sigmoid(input), target_smooth) # tensor(0.7490) + + let device = Default::default(); + let logits = Tensor::<1>::from_floats([0.8271, 0.9626, 0.3796, 0.2355], &device); + let targets = Tensor::<1, Int>::from_data(TensorData::from([0, 1, 0, 1]), &device); + + let loss_actual = BinaryCrossEntropyLossConfig::new() + .with_smoothing(Some(0.1)) + .init(&device) + .forward(sigmoid(logits), targets) + .into_data(); + + let loss_expected = TensorData::from([0.7490]); + loss_actual.assert_approx_eq::(&loss_expected, Tolerance::relative(1e-4)); + } + + #[test] + fn test_binary_cross_entropy_multilabel() { + // import torch + // from torch import nn + // input = torch.tensor([[0.5150, 0.3097, 0.7556], [0.4974, 0.9879, 0.1564]]) + // target = torch.tensor([[1., 0., 1.], [1., 0., 0.]]) + // weights = torch.tensor([3., 7., 0.9]) + // loss = nn.BCEWithLogitsLoss() + // out = loss(input, target) # tensor(0.7112) + + let device = Default::default(); + let logits = Tensor::<2>::from_floats( + [[0.5150, 0.3097, 0.7556], [0.4974, 0.9879, 0.1564]], + &device, + ); + let targets = + Tensor::<2, Int>::from_data(TensorData::from([[1, 0, 1], [1, 0, 0]]), &device); + + let loss_actual = BinaryCrossEntropyLossConfig::new() + .with_logits(true) + .init(&device) + .forward(logits, targets) + .into_data(); + + let loss_expected = TensorData::from([0.7112]); + loss_actual.assert_approx_eq::(&loss_expected, Tolerance::relative(1e-4)); + } + + #[test] + fn test_binary_cross_entropy_multilabel_with_weights() { + // import torch + // from torch import nn + // input = torch.tensor([[0.5150, 0.3097, 0.7556], [0.4974, 0.9879, 0.1564]]) + // target = torch.tensor([[1., 0., 1.], [1., 0., 0.]]) + // loss = nn.BCEWithLogitsLoss() + // out = loss(input, target) # tensor(3.1708) + + let device = Default::default(); + let logits = Tensor::<2>::from_floats( + [[0.5150, 0.3097, 0.7556], [0.4974, 0.9879, 0.1564]], + &device, + ); + let targets = + Tensor::<2, Int>::from_data(TensorData::from([[1, 0, 1], [1, 0, 0]]), &device); + let weights = [3., 7., 0.9]; + + let loss_actual = BinaryCrossEntropyLossConfig::new() + .with_logits(true) + .with_weights(Some(weights.to_vec())) + .init(&device) + .forward(logits, targets) + .into_data(); + + let loss_expected = TensorData::from([3.1708]); + loss_actual.assert_approx_eq::(&loss_expected, Tolerance::default()); + } + + #[test] + fn test_binary_cross_entropy_multilabel_with_smoothing() { + // import torch + // from torch import nn + // input = torch.tensor([[0.5150, 0.3097, 0.7556], [0.4974, 0.9879, 0.1564]]) + // target = torch.tensor([[1., 0., 1.], [1., 0., 0.]]) + // target_smooth = target * (1 - 0.1) + (0.1 / 3) + // loss = nn.BCELoss() + // sigmoid = nn.Sigmoid() + // out = loss(sigmoid(input), target_smooth) # tensor(0.7228) + + let device = Default::default(); + let logits = Tensor::<2>::from_floats( + [[0.5150, 0.3097, 0.7556], [0.4974, 0.9879, 0.1564]], + &device, + ); + let targets = + Tensor::<2, Int>::from_data(TensorData::from([[1, 0, 1], [1, 0, 0]]), &device); + + let loss_actual = BinaryCrossEntropyLossConfig::new() + .with_smoothing(Some(0.1)) + .init(&device) + .forward(sigmoid(logits), targets) + .into_data(); + + let loss_expected = TensorData::from([0.7228]); + loss_actual.assert_approx_eq::(&loss_expected, Tolerance::default()); + } + + #[test] + #[should_panic = "The number of classes"] + fn multilabel_weights_should_match_target() { + // import torch + // from torch import nn + // input = torch.tensor([[0.5150, 0.3097, 0.7556], [0.4974, 0.9879, 0.1564]]) + // target = torch.tensor([[1., 0., 1.], [1., 0., 0.]]) + // loss = nn.BCEWithLogitsLoss() + // out = loss(input, target) # tensor(3.1708) + + let device = Default::default(); + let logits = Tensor::<2>::from_floats( + [[0.5150, 0.3097, 0.7556], [0.4974, 0.9879, 0.1564]], + &device, + ); + let targets = + Tensor::<2, Int>::from_data(TensorData::from([[1, 0, 1], [1, 0, 0]]), &device); + let weights = [3., 7.]; + + let _loss = BinaryCrossEntropyLossConfig::new() + .with_logits(true) + .with_weights(Some(weights.to_vec())) + .init(&device) + .forward(logits, targets); + } + + #[test] + fn display() { + let config = + BinaryCrossEntropyLossConfig::new().with_weights(Some(alloc::vec![3., 7., 0.9])); + let loss = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{loss}"), + "BinaryCrossEntropyLoss {weights: Tensor {rank: 1, shape: [3]}, smoothing: None, logits: false}" + ); + } +} diff --git a/crates/burn-nn/src/loss/cosine_embedding.rs b/crates/burn-nn/src/loss/cosine_embedding.rs new file mode 100644 index 0000000..0d94811 --- /dev/null +++ b/crates/burn-nn/src/loss/cosine_embedding.rs @@ -0,0 +1,294 @@ +use alloc::format; + +use burn::tensor::linalg::cosine_similarity; + +use burn_core as burn; + +use crate::loss::reduction::Reduction; +use burn::config::Config; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::{Int, Tensor, activation::relu}; + +/// Configuration for CosineEmbeddingLoss. +#[derive(Config, Debug)] +pub struct CosineEmbeddingLossConfig { + /// Margin for negative samples. + #[config(default = 0.0)] + pub margin: f32, + + /// Specifies the reduction to apply to the output. + #[config(default = "Reduction::Mean")] + pub reduction: Reduction, +} + +impl CosineEmbeddingLossConfig { + /// Initialize CosineEmbeddingLoss. + pub fn init(&self) -> CosineEmbeddingLoss { + CosineEmbeddingLoss { + margin: self.margin, + reduction: self.reduction.clone(), + } + } +} + +/// Cosine embedding loss between two tensors. +/// +/// Measures cosine distance between tensors. +/// Used for learning embeddings or similarity. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct CosineEmbeddingLoss { + /// Margin value. Default: 0.0 + pub margin: f32, + + /// Reduction method + #[module(skip)] + pub reduction: Reduction, +} + +impl Default for CosineEmbeddingLoss { + fn default() -> Self { + CosineEmbeddingLossConfig::new().init() + } +} + +impl ModuleDisplay for CosineEmbeddingLoss { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("margin", &self.margin) + .add("reduction", format!("{:?}", self.reduction).as_str()) + .optional() + } +} + +impl CosineEmbeddingLoss { + /// Creates a new instance + pub fn new() -> Self { + CosineEmbeddingLossConfig::new().init() + } + + /// Compute loss with reduction. + /// + /// # Shapes + /// + /// - input1: ``[batch_size, embedding_dim]`` + /// - input2: ``[batch_size, embedding_dim]`` + /// - target: ``[batch_size]`` with values 1 or -1 + /// + /// # Returns + /// + /// Loss tensor of shape ``[1]`` + pub fn forward( + &self, + input1: Tensor<2>, + input2: Tensor<2>, + target: Tensor<1, Int>, + ) -> Tensor<1> { + let tensor = self.forward_no_reduction(input1, input2, target); + match &self.reduction { + Reduction::Mean | Reduction::Auto => tensor.mean(), + Reduction::Sum => tensor.sum(), + other => panic!("{other:?} reduction is not supported"), + } + } + + /// Compute loss without applying reduction. + /// + /// # Arguments + /// + /// * `input1` - First input tensor of shape ``[batch_size, embedding_dim]`` + /// * `input2` - Second input tensor of shape ``[batch_size, embedding_dim]`` + /// * `target` - Target tensor of shape ``[batch_size]`` with values 1 or -1 + /// + /// # Returns + /// + /// Tensor of per-element losses with shape ``[batch_size]`` + pub fn forward_no_reduction( + &self, + input1: Tensor<2>, + input2: Tensor<2>, + target: Tensor<1, Int>, + ) -> Tensor<1> { + self.assertions(&input1, &input2, &target); + + // cos_sim shape: [batch_size, 1] + let cos_sim = cosine_similarity(input1, input2, 1, None); + // cos_sim shape: [batch_size] + let cos_sim: Tensor<1> = cos_sim.squeeze_dim(1); + + let mut loss = cos_sim.zeros_like(); + + // Similar pairs (target == 1) - Formula: L = 1 - cos_sim + let similar_mask = target.clone().equal_elem(1); + let similar_loss = cos_sim.clone().neg().add_scalar(1); + loss = loss.mask_where(similar_mask, similar_loss); + + // Dissimilar pairs (target == -1) - Formula: L = max(0, cos_sim - margin) + let dissimilar_mask = target.equal_elem(-1); + let dissimilar_loss = relu(cos_sim.clone().sub_scalar(self.margin)); + loss = loss.mask_where(dissimilar_mask, dissimilar_loss); + + // return loss shape: [batch_size] + loss + } + + fn assertions(&self, input1: &Tensor<2>, input2: &Tensor<2>, target: &Tensor<1, Int>) { + let [batch_size1, dim1] = input1.dims(); + let [batch_size2, dim2] = input2.dims(); + let [batch_size_target] = target.dims(); + + assert_eq!( + batch_size1, batch_size2, + "Batch size of input1 ({batch_size1}) must match batch size of input2 ({batch_size2})" + ); + + assert_eq!( + dim1, dim2, + "Embedding dimension of input1 ({dim1}) must match embedding dimension of input2 ({dim2})" + ); + + assert_eq!( + batch_size1, batch_size_target, + "Batch size of inputs ({batch_size1}) must match batch size of target ({batch_size_target})" + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn cosine_embedding_loss_positive_target() { + let device = Default::default(); + + // Two identical vectors should have cosine similarity of 1 + let input1 = Tensor::<2>::from_data(TensorData::from([[1.0, 0.0], [0.0, 1.0]]), &device); + + let input2 = Tensor::<2>::from_data(TensorData::from([[1.0, 0.0], [0.0, 1.0]]), &device); + + // Target 1 means that inputs should be similar + let target = Tensor::<1, Int>::from_data(TensorData::from([1, 1]), &device); + + let loss = CosineEmbeddingLossConfig::new().init(); + let loss_no_reduction = + loss.forward_no_reduction(input1.clone(), input2.clone(), target.clone()); + let loss_mean = loss.forward(input1.clone(), input2.clone(), target.clone()); + + let loss_sum = loss.forward(input1, input2, target); + + // For identical vectors, 1 - cos_sim = 1 - 1 = 0 + let expected_no_reduction = TensorData::from([0.0, 0.0]); + loss_no_reduction + .into_data() + .assert_approx_eq::(&expected_no_reduction, Tolerance::default()); + + let expected_mean = TensorData::from([0.0]); + loss_mean + .into_data() + .assert_approx_eq::(&expected_mean, Tolerance::default()); + + let expected_sum = TensorData::from([0.0]); + loss_sum + .into_data() + .assert_approx_eq::(&expected_sum, Tolerance::default()); + } + + #[test] + fn cosine_embedding_loss_negative_target() { + let device = Default::default(); + + // Two identical vectors should have cosine similarity of 1 + let input1 = Tensor::<2>::from_data(TensorData::from([[1.0, 0.0], [0.0, 1.0]]), &device); + + let input2 = Tensor::<2>::from_data(TensorData::from([[1.0, 0.0], [0.0, 1.0]]), &device); + + // Target -1 means that inputs should be dissimilar + let target = Tensor::<1, Int>::from_data(TensorData::from([-1, -1]), &device); + + // With margin 0.0, max(0, cos_sim - margin) = max(0, 1 - 0) = 1 + let loss = CosineEmbeddingLossConfig::new().init(); + let loss_no_reduction = + loss.forward_no_reduction(input1.clone(), input2.clone(), target.clone()); + let loss_mean = loss.forward(input1.clone(), input2.clone(), target.clone()); + + // Create a loss with Sum reduction for testing + let loss_sum_config = CosineEmbeddingLossConfig::new().with_reduction(Reduction::Sum); + let loss_sum = + loss_sum_config + .init() + .forward(input1.clone(), input2.clone(), target.clone()); + + let expected_no_reduction = TensorData::from([1.0, 1.0]); + loss_no_reduction + .into_data() + .assert_approx_eq::(&expected_no_reduction, Tolerance::default()); + + let expected_mean = TensorData::from([1.0]); + loss_mean + .into_data() + .assert_approx_eq::(&expected_mean, Tolerance::default()); + + let expected_sum = TensorData::from([2.0]); + loss_sum + .into_data() + .assert_approx_eq::(&expected_sum, Tolerance::default()); + + // With margin 0.5, max(0, cos_sim - margin) = max(0, 1 - 0.5) = 0.5 + let loss_with_margin = CosineEmbeddingLossConfig::new().with_margin(0.5).init(); + let loss_with_margin = loss_with_margin.forward(input1, input2, target); + + let expected = TensorData::from([0.5]); + loss_with_margin + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn cosine_embedding_loss_mixed_targets() { + let device = Default::default(); + + let input1 = Tensor::<2>::from_data(TensorData::from([[1.0, 0.0], [0.0, 1.0]]), &device); + + let input2 = Tensor::<2>::from_data(TensorData::from([[1.0, 0.0], [0.0, 1.0]]), &device); + + // Mixed targets + let target = Tensor::<1, Int>::from_data(TensorData::from([1, -1]), &device); + + let loss = CosineEmbeddingLossConfig::new().init(); + let loss_no_reduction = + loss.forward_no_reduction(input1.clone(), input2.clone(), target.clone()); + let loss_mean = loss.forward(input1, input2, target); + + let expected_no_reduction = TensorData::from([0.0, 1.0]); + loss_no_reduction + .into_data() + .assert_approx_eq::(&expected_no_reduction, Tolerance::default()); + + let expected_mean = TensorData::from([0.5]); + loss_mean + .into_data() + .assert_approx_eq::(&expected_mean, Tolerance::default()); + } + + #[test] + fn display() { + let config = CosineEmbeddingLossConfig::new().with_margin(0.5); + let loss = config.init(); + + assert_eq!( + alloc::format!("{loss}"), + "CosineEmbeddingLoss {margin: 0.5, reduction: Mean}" + ); + } +} diff --git a/crates/burn-nn/src/loss/cross_entropy.rs b/crates/burn-nn/src/loss/cross_entropy.rs new file mode 100644 index 0000000..a34ad41 --- /dev/null +++ b/crates/burn-nn/src/loss/cross_entropy.rs @@ -0,0 +1,514 @@ +use burn_core as burn; +use burn_core::tensor::IndexingUpdateOp; + +use alloc::string::ToString; +use alloc::vec; +use alloc::vec::Vec; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::activation::log_softmax; +use burn::tensor::{Bool, Device, Int, Tensor}; +use burn::{config::Config, module::Module}; + +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float; + +/// Configuration to create a [Cross-entropy loss](CrossEntropyLoss) using the [init function](CrossEntropyLossConfig::init). +#[derive(Config, Debug)] +pub struct CrossEntropyLossConfig { + /// Create padded cross entropy. + /// + /// Prevents pad tokens from impacting loss calculation. + pub pad_tokens: Option>, + + /// Create weighted cross-entropy. + /// + /// The loss of a specific sample will simply be given by: weight * log(p(x)) * 1, + /// + /// # Pre-conditions + /// - The order of the weight vector should correspond to the label integer assignment. + /// - Targets assigned negative Int's will not be allowed. + pub weights: Option>, + + /// Create cross-entropy with label smoothing. + /// + /// Hard labels {0, 1} will be changed to y_smoothed = y(1 - a) + a / nr_classes. + /// Alpha = 0 would be the same as default. + pub smoothing: Option, + + /// Create cross-entropy with probabilities as input instead of logits. + /// + #[config(default = true)] + pub logits: bool, +} + +impl CrossEntropyLossConfig { + /// Initialize [Cross-entropy loss](CrossEntropyLoss). + pub fn init(&self, device: &Device) -> CrossEntropyLoss { + self.assertions(); + CrossEntropyLoss { + pad_tokens: self.pad_tokens.clone(), + weights: self + .weights + .as_ref() + .map(|e| Tensor::<1>::from_floats(e.as_slice(), device)), + smoothing: self.smoothing, + logits: self.logits, + } + } + + fn assertions(&self) { + if let Some(alpha) = self.smoothing { + assert!( + (0.0..=1.).contains(&alpha), + "Alpha of Cross-entropy loss with smoothed labels should be in interval [0, 1]. Got {alpha}" + ); + }; + if let Some(weights) = self.weights.as_ref() { + assert!( + weights.iter().all(|e| e > &0.), + "Weights of cross-entropy have to be positive." + ); + } + } +} + +/// Calculate the cross entropy loss from the input logits and the targets. +/// +/// Should be created using [CrossEntropyLossConfig] +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct CrossEntropyLoss { + /// Pad tokens to ignore in the loss calculation. + pub pad_tokens: Option>, + /// Weights for cross-entropy. + pub weights: Option>, + /// Label smoothing factor. + pub smoothing: Option, + /// Use logits as input. + pub logits: bool, +} + +impl ModuleDisplay for CrossEntropyLoss { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let pad_tokens = if let Some(pad_tokens) = &self.pad_tokens { + alloc::format!("Vec<0..{}>", pad_tokens.len()) + } else { + "None".to_string() + }; + + content + .add("pad_tokens", &pad_tokens) + .add("weights", &self.weights) + .add("smoothing", &self.smoothing) + .add("logits", &self.logits) + .optional() + } +} + +impl CrossEntropyLoss { + /// For backward compatibility. + pub fn new(pad_index: Option, device: &Device) -> Self { + CrossEntropyLossConfig::new() + .with_pad_tokens(pad_index.map(|e| vec![e])) + .init(device) + } + + /// Compute the criterion on the input tensor. + /// + /// # Shapes + /// + /// - logits: `[batch_size, num_targets]` + /// - targets: `[batch_size]` + pub fn forward(&self, logits: Tensor<2>, targets: Tensor<1, Int>) -> Tensor<1> { + Self::assertions(logits.clone(), targets.clone()); + match self.smoothing { + Some(alpha) => self.forward_smoothed(logits, targets, alpha), + _ => self.forward_default(logits, targets), + } + } + + fn forward_smoothed( + &self, + logits: Tensor<2>, + targets: Tensor<1, Int>, + alpha: f32, + ) -> Tensor<1> { + let mask = self.padding_mask(&targets); + let tensor = if self.logits { + log_softmax(logits, 1) + } else { + logits.log() + }; + let [batch_size, nr_classes] = tensor.dims(); + let tensor = tensor + * Self::compute_smoothed_targets([batch_size, nr_classes], targets.clone(), alpha); + + match &self.weights { + Some(weights) => { + let tensor = tensor + * weights + .clone() + .reshape([1, nr_classes]) + .repeat_dim(0, batch_size); + let weights = weights.clone().gather(0, targets); + let tensor = Self::apply_mask_2d(tensor, mask); + tensor.sum().neg() / weights.sum() + } + None => { + let tensor = Self::apply_mask_2d(tensor, mask); + tensor.sum_dim(1).mean().neg() + } + } + } + + fn forward_default(&self, logits: Tensor<2>, targets: Tensor<1, Int>) -> Tensor<1> { + let [batch_size] = targets.dims(); + + let mask = self.padding_mask(&targets); + let target_indices = targets.clone().reshape([batch_size, 1]); + let tensor = if self.logits { + log_softmax(logits, 1).gather(1, target_indices) + } else { + // TODO: finfo stable eps + let finfo = logits.dtype().finfo().unwrap(); + let eps = finfo.min_positive.sqrt(); + logits.clamp_min(eps).gather(1, target_indices).log() + }; + + match &self.weights { + Some(weights) => { + let weights = weights.clone().gather(0, targets); + let tensor = tensor.reshape([batch_size]) * weights.clone(); + let tensor = Self::apply_mask_1d(tensor, mask); + tensor.sum().neg() / weights.sum() + } + None => { + let tensor = Self::apply_mask_1d(tensor.reshape([batch_size]), mask); + tensor.mean().neg() + } + } + } + + fn compute_smoothed_targets( + shape: [usize; 2], + targets: Tensor<1, Int>, + alpha: f32, + ) -> Tensor<2> { + let [batch_size, nr_classes] = shape; + let device = &targets.device(); + let targets_matrix = Tensor::<2>::zeros(shape, device).scatter( + 1, + targets.reshape([batch_size, 1]), + Tensor::ones([batch_size, 1], device), + IndexingUpdateOp::Add, + ); + targets_matrix * (1. - alpha) + alpha / nr_classes as f32 + } + + fn padding_mask(&self, targets: &Tensor<1, Int>) -> Option> { + let mut mask = None; + if let Some(pad_tokens) = &self.pad_tokens { + let mut res = targets.clone().equal_elem(pad_tokens[0] as i64).int(); + for x in pad_tokens { + res = res + targets.clone().equal_elem(*x as i64).int(); + } + mask = Some(res.greater_elem(0)); + } + + mask + } + + fn apply_mask_1d(mut tensor: Tensor<1>, mask: Option>) -> Tensor<1> { + if let Some(mask) = mask { + tensor = tensor.mask_fill(mask, 0); + } + + tensor + } + + fn apply_mask_2d(mut tensor: Tensor<2>, mask: Option>) -> Tensor<2> { + if let Some(mask) = mask { + let [batch_size, nr_classes] = tensor.dims(); + tensor = tensor.mask_fill(mask.reshape([batch_size, 1]).repeat_dim(1, nr_classes), 0); + } + + tensor + } + + fn assertions(logits: Tensor<2>, targets: Tensor<1, Int>) { + let [logits_height, _] = logits.dims(); + let [targets_height] = targets.dims(); + assert!( + logits_height == targets_height, + "Shape of targets ({targets_height}) should correspond to outer shape of logits ({logits_height})." + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Tolerance; + use burn::tensor::{Distribution, TensorData, loss::cross_entropy_with_logits}; + type FT = f32; + + macro_rules! setup { + () => {{ + let [batch_size, num_targets] = [4, 5]; + let device = Default::default(); + let logits = Tensor::<2>::random( + [batch_size, num_targets], + Distribution::Normal(0., 1.0), + &device, + ); + let targets = Tensor::<1, Int>::from_data(TensorData::from([2, 0, 4, 1]), &device); + let targets_logits = Tensor::<2>::from_data( + TensorData::from([ + [0.0, 0.0, 1.0, 0.0, 0.0], + [1.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 1.0], + [0.0, 1.0, 0.0, 0.0, 0.0], + ]), + &device, + ); + (logits, targets, targets_logits) + }}; + } + + macro_rules! setup_padded { + () => {{ + let [batch_size, num_targets, pad_index] = [4, 5, 1]; + let device = Default::default(); + let logits = Tensor::<2>::random( + [batch_size, num_targets], + Distribution::Normal(0., 1.0), + &device, + ); + let targets = + Tensor::<1, Int>::from_data(TensorData::from([2, 0, 4, pad_index as i64]), &device); + let targets_logits = Tensor::<2>::from_data( + TensorData::from([ + [0.0, 0.0, 0.0, 0.0, 0.0], + [1.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 1.0], + [0.0, 0.0, 0.0, 0.0, 0.0], + ]), + &device, + ); + (logits, targets, targets_logits) + }}; + } + + #[test] + fn test_cross_entropy_loss_with_weights() { + let (logits, targets, targets_logits) = setup!(); + let weights = vec![1.0, 2., 3., 4., 5.]; + let device = Default::default(); + let loss_1 = CrossEntropyLossConfig::new() + .with_weights(Some(weights.clone())) + .init(&device) + .forward(logits.clone(), targets); + let tensor = log_softmax(logits, 1); + let loss_2 = tensor + * targets_logits + * Tensor::<1>::from_floats(weights.as_slice(), &device) + .unsqueeze() + .repeat_dim(0, 4); + let loss_2 = loss_2.sum().neg() / (1. + 2. + 3. + 5.); + loss_1 + .into_data() + .assert_approx_eq::(&loss_2.into_data(), Tolerance::default()); + } + + #[test] + fn test_label_smoothing_with_weights_and_alpha_zero() { + let (logits, targets, _) = setup!(); + let device = Default::default(); + let weights = vec![1.0, 2., 3., 4., 5.]; + let loss_1 = CrossEntropyLossConfig::new() + .with_weights(Some(weights.clone())) + .init(&device) + .forward(logits.clone(), targets.clone()); + let loss_2 = CrossEntropyLossConfig::new() + .with_weights(Some(weights.clone())) + .with_smoothing(Some(0.)) + .init(&device) + .forward(logits.clone(), targets); + loss_1 + .into_data() + .assert_approx_eq::(&loss_2.into_data(), Tolerance::default()); + } + + #[test] + fn test_cross_entropy_loss() { + let (logits, targets, targets_logits) = setup!(); + let device = Default::default(); + let loss_1 = CrossEntropyLossConfig::new() + .init(&device) + .forward(logits.clone(), targets); + let loss_2 = cross_entropy_with_logits(logits, targets_logits); + + loss_1 + .into_data() + .assert_approx_eq::(&loss_2.into_data(), Tolerance::default()); + } + + #[test] + fn test_label_smoothing_alpha_equal_zero() { + let (logits, targets, _) = setup!(); + let device = Default::default(); + let loss_1 = CrossEntropyLossConfig::new() + .init(&device) + .forward(logits.clone(), targets.clone()); + let loss_2 = CrossEntropyLossConfig::new() + .with_smoothing(Some(0.)) + .init(&device) + .forward(logits, targets); + + loss_1 + .into_data() + .assert_approx_eq::(&loss_2.into_data(), Tolerance::default()); + } + + #[test] + fn test_cross_entropy_loss_with_pad_token() { + let (logits, targets, targets_logits) = setup_padded!(); + let pad_index = 1; + + let loss_1 = CrossEntropyLossConfig::new() + .with_pad_tokens(Some(vec![pad_index, 2])) + .init(&logits.device()) + .forward(logits.clone(), targets); + let loss_2 = cross_entropy_with_logits(logits, targets_logits); + + loss_1 + .into_data() + .assert_approx_eq::(&loss_2.into_data(), Tolerance::default()); + } + + #[test] + fn test_label_smoothing_with_zero_alpha_and_pad_token() { + let (logits, targets, _) = setup_padded!(); + let pad_index = 1; + + let loss_1 = CrossEntropyLossConfig::new() + .with_pad_tokens(Some(vec![pad_index, 2])) + .init(&logits.device()) + .forward(logits.clone(), targets.clone()); + let loss_2 = CrossEntropyLossConfig::new() + .with_pad_tokens(Some(vec![pad_index, 2])) + .with_smoothing(Some(0.)) + .init(&logits.device()) + .forward(logits.clone(), targets); + + loss_1 + .into_data() + .assert_approx_eq::(&loss_2.into_data(), Tolerance::default()); + } + + #[test] + fn test_label_smoothing_target_conversion() { + let (logits, targets, _) = setup!(); + let smoothed_targets = + CrossEntropyLoss::compute_smoothed_targets(logits.dims(), targets, 0.05); + let targets_logits = Tensor::<2>::from_data( + TensorData::from([ + [0.01, 0.01, 0.96, 0.01, 0.01], + [0.96, 0.01, 0.01, 0.01, 0.01], + [0.01, 0.01, 0.01, 0.01, 0.96], + [0.01, 0.96, 0.01, 0.01, 0.01], + ]), + &Default::default(), + ); + smoothed_targets + .into_data() + .assert_approx_eq::(&targets_logits.into_data(), Tolerance::default()); + } + + #[test] + fn test_label_smoothing() { + let (logits, targets, _) = setup!(); + let device = Default::default(); + let loss_1 = CrossEntropyLossConfig::new() + .with_smoothing(Some(0.05)) + .init(&device) + .forward(logits.clone(), targets); + let targets_logits = Tensor::<2>::from_data( + TensorData::from([ + [0.01, 0.01, 0.96, 0.01, 0.01], + [0.96, 0.01, 0.01, 0.01, 0.01], + [0.01, 0.01, 0.01, 0.01, 0.96], + [0.01, 0.96, 0.01, 0.01, 0.01], + ]), + &device, + ); + + let x = log_softmax(logits, 1); + let loss_2 = (x * targets_logits).sum_dim(1).mean().neg(); + + loss_1 + .into_data() + .assert_approx_eq::(&loss_2.into_data(), Tolerance::default()); + } + + #[test] + fn test_logits_flag_affects_output() { + let device = Default::default(); + + let probs = Tensor::<2>::from_data( + TensorData::from([ + [0.1, 0.2, 0.7, 0.0, 0.0], + [0.7, 0.1, 0.1, 0.1, 0.0], + [0.2, 0.2, 0.2, 0.2, 0.2], + [0.0, 0.3, 0.3, 0.2, 0.2], + ]), + &device, + ); + + let targets = Tensor::<1, Int>::from_data(TensorData::from([2, 0, 4, 1]), &device); + + let loss_logits = CrossEntropyLossConfig::new() + .init(&device) + .forward(probs.clone(), targets.clone()); + + let loss_probs = CrossEntropyLossConfig::new() + .with_logits(false) + .init(&device) + .forward(probs, targets); + + // They must differ if logits flag is implemented correctly + let loss_logits = loss_logits.into_data(); + let loss_probs = loss_probs.into_data(); + + loss_logits.assert_approx_eq::(&TensorData::from([1.354197]), Tolerance::default()); + loss_probs.assert_approx_eq::(&TensorData::from([0.88169014]), Tolerance::default()); + + assert_ne!( + loss_logits.as_slice::().unwrap(), + loss_probs.as_slice::().unwrap(), + "logits flag should change computation (log_softmax vs log)" + ); + } + + #[test] + fn display() { + let config = CrossEntropyLossConfig::new() + .with_weights(Some(alloc::vec![3., 7., 0.9])) + .with_smoothing(Some(0.5)); + let loss = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{loss}"), + "CrossEntropyLoss {pad_tokens: None, weights: Tensor {rank: 1, shape: [3]}, smoothing: 0.5, logits: true}" + ); + } + + // TODO: backward tests +} diff --git a/crates/burn-nn/src/loss/ctc.rs b/crates/burn-nn/src/loss/ctc.rs new file mode 100644 index 0000000..a130795 --- /dev/null +++ b/crates/burn-nn/src/loss/ctc.rs @@ -0,0 +1,1417 @@ +#![allow(clippy::excessive_precision)] + +use super::Reduction; +use burn::config::Config; +use burn::module::Module; +use burn::tensor::{Int, Tensor}; +use burn_core as burn; + +/// Configuration for the [CTC Loss](CTCLoss) module. +#[derive(Config, Debug)] +pub struct CTCLossConfig { + /// The index number used to represent the blank label. Default value is `0`. + #[config(default = 0)] + pub blank: usize, + /// Whether to zero infinite losses and the associated gradients. Default value is `false`. + #[config(default = false)] + pub zero_infinity: bool, +} + +impl CTCLossConfig { + /// Initialize a new [CTC Loss](CTCLoss) module + pub fn init(&self) -> CTCLoss { + CTCLoss { + blank: self.blank, + zero_infinity: self.zero_infinity, + } + } +} + +/// Computes the Connectionist Temporal Classification (CTC) loss. +/// +/// Calculates the loss between a continuous (unsegmented) time series and a target sequence. +/// CTC sums over the probability of all possible alignments of the input to the target, +/// producing a loss value that is differentiable with respect to each input node. +/// +/// The input to this loss is expected to be **log-probabilities** (e.g,, via `log_softmax`), +/// not raw logits. +/// +/// # References +/// +/// - [Connectionist Temporal Classification: Labelling Unsegmented Sequence Data with Recurrent Neural Networks](https://www.cs.toronto.edu/~graves/icml_2006.pdf) +/// +/// # Example +/// +/// ```rust,ignore +/// use burn::tensor::{Tensor, Int}; +/// use burn::tensor::activation::log_softmax; +/// use burn::nn::loss::{CTCLossConfig, CTCLoss}; +/// +/// let device = Default::default(); +/// +/// // Initialize CTC Loss with default configuration +/// let ctc_loss = CTCLossConfig::new().init(); +/// +/// // Initialize CTC Loss with custom configuration +/// let ctc_loss = CTCLossConfig::new() +/// .with_blank(1) +/// .with_zero_infinity(true) +/// .init(); +/// +/// // Prepare inputs (Logits shape: [Time, Batch, Class]) +/// // In your actual code, the logits would be the output of your model +/// let logits = Tensor::<3>::ones([10, 2, 5], &device); +/// let log_probs = log_softmax(logits, 2); +/// +/// // Targets shape: [Batch, Max_Target_Len] +/// // Note: Targets should not contain the blank index (1). +/// let targets = Tensor::<2, Int>::from_data([[0, 2], [3, 4]], &device); +/// +/// // Lengths shape: [Batch] +/// let input_lengths = Tensor::<1, Int>::from_data([10, 8], &device); +/// let target_lengths = Tensor::<1, Int>::from_data([2, 2], &device); +/// +/// // Compute loss +/// let loss = ctc_loss.forward(log_probs, targets, input_lengths, target_lengths); +/// ``` +#[derive(Module, Debug)] +pub struct CTCLoss { + blank: usize, + zero_infinity: bool, +} + +impl CTCLoss { + /// Computes the CTC loss for the input log-probabilities and targets with no reduction applied. + /// + /// # Arguments + /// + /// - `log_probs`: The log-probabilities of the outputs (e.g., from `log_softmax`). + /// - `targets`: A 2D tensor containing the target class indices. These indices should not + /// include the blank index used in CTC loss. The targets are padded to the length of the longest sequence. + /// - `input_lengths`: A 1D tensor containing the actual length of the input sequence for each batch. This + /// allows retrieving the actual sequence of log-probabilities from `log_probs` if the batch contains + /// sequences of varying lengths. + /// - `target_lengths`: A 1D tensor containing the actual length of the target sequence for each target + /// sequence in `targets`. + /// + /// # Returns + /// + /// - A 1D tensor of shape `[batch_size]` containing the loss for each sample. + /// + /// # Shapes + /// + /// - `log_probs`: `[time_steps, batch_size, num_classes]` where `num_classes` includes blank. + /// - `targets`: `[batch_size, max_target_length]` + /// - `input_lengths`: `[batch_size]` + /// - `target_lengths`: `[batch_size]` + pub fn forward( + &self, + log_probs: Tensor<3>, + targets: Tensor<2, Int>, + input_lengths: Tensor<1, Int>, + target_lengths: Tensor<1, Int>, + ) -> Tensor<1> { + let [max_input_length, batch_size, num_classes] = log_probs.dims(); + let max_target_len = targets.dims()[1]; + let input_lengths_len = input_lengths.dims()[0]; + let target_lengths_len = target_lengths.dims()[0]; + self.assertions( + batch_size, + num_classes, + targets.clone(), + input_lengths_len, + target_lengths_len, + ); + self.length_assertions( + input_lengths.clone(), + target_lengths.clone(), + max_target_len, + max_input_length, + ); + + let mut loss = burn::tensor::module::ctc_loss( + log_probs, + targets, + input_lengths, + target_lengths, + self.blank, + ); + + if self.zero_infinity { + let inf_mask = loss.clone().is_inf(); + loss = loss.clone().mask_where(inf_mask, loss.clone().zeros_like()); + } + + loss + } + + /// Computes the CTC loss for the input log-probabilities and targets with reduction. + /// + /// # Arguments + /// + /// - `log_probs`: The log-probabilities of the outputs (e.g., from `log_softmax`). + /// - `targets`: A 2D tensor containing the target class indices. These indices should not + /// include the blank index used in CTC loss. The targets are padded to the length of the longest sequence. + /// - `input_lengths`: A 1D tensor containing the actual length of the input sequence for each batch. This + /// allows retrieving the actual sequence of log-probabilities from `log_probs` if the batch contains + /// sequences of varying lengths. + /// - `target_lengths`: A 1D tensor containing the actual length of the target sequence for each target + /// sequence in `targets`. + /// - `reduction`: The reduction stratey to apply to the loss tensor containing the CTC loss values for + /// each sample (e.g., mean, sum). For the mean reduction strategy, the output losses will be divided + /// by the target lengths and then the mean over the batch is taken. This follows PyTorch's behavior. + /// + /// # Returns + /// + /// - A 1D tensor of shape `[1]` containing the reduced loss value. + /// + /// # Shapes + /// + /// - `log_probs`: `[time_steps, batch_size, num_classes]` where `num_classes` includes blank. + /// - `targets`: `[batch_size, max_target_length]` + /// - `input_lengths`: `[batch_size]` + /// - `target_lengths`: `[batch_size]` + /// + /// # Panics + /// - If `reduction` is not one of `Reduction::Auto`, `Reduction::Mean`, and `Reduction::Sum`. + /// - If `blank` index is greater than or equal to `num_classes`. + /// - If the batch dimension of `log_probs`, `targets`, `input_lengths`, and `target_lengths` do not match. + pub fn forward_with_reduction( + &self, + log_probs: Tensor<3>, + targets: Tensor<2, Int>, + input_lengths: Tensor<1, Int>, + target_lengths: Tensor<1, Int>, + reduction: Reduction, + ) -> Tensor<1> { + let ctc_loss_tensor = + self.forward(log_probs, targets, input_lengths, target_lengths.clone()); + + match reduction { + Reduction::Auto | Reduction::Mean => { + // Following PyTorch's behavior where the output losses are divided + // by the target lengths and then the mean over the batch is taken + let target_lengths_float = target_lengths.float(); + ctc_loss_tensor.div(target_lengths_float).mean() + } + Reduction::Sum => ctc_loss_tensor.sum(), + other => panic!("{other:?} reduction is not supported"), + } + } + + /// Checks the per-element length invariants required by the alpha + /// recursion. These require reading the length tensors from the device, + /// so the checks are gated behind `cfg(debug_assertions)` to avoid the + /// device-to-host sync in release builds. + /// + /// Validated: + /// - `target_lengths[i] >= 0` + /// - `target_lengths[i] <= max_target_len` + /// - `input_lengths[i] >= target_lengths[i]` + /// - `input_lengths[i] <= max_input_length` + #[allow(unused_variables)] + fn length_assertions( + &self, + input_lengths: Tensor<1, Int>, + target_lengths: Tensor<1, Int>, + max_target_len: usize, + max_input_length: usize, + ) { + #[cfg(debug_assertions)] + { + let target_lengths_data = target_lengths.into_data(); + let input_lengths_data = input_lengths.into_data(); + let target_iter = target_lengths_data.iter::(); + let input_iter = input_lengths_data.iter::(); + for (i, (tl, il)) in target_iter.zip(input_iter).enumerate() { + assert!(tl >= 0, "target_lengths[{i}] = {tl} must be non-negative"); + assert!( + tl as usize <= max_target_len, + "target_lengths[{i}] = {tl} exceeds the targets tensor width {max_target_len}" + ); + assert!( + il >= tl, + "input_lengths[{i}] = {il} must be >= target_lengths[{i}] = {tl} \ + (no valid CTC alignment otherwise)" + ); + assert!( + il as usize <= max_input_length, + "input_lengths[{i}] = {il} exceeds the log_probs time dimension \ + {max_input_length}" + ); + } + } + } + + fn assertions( + &self, + batch_size: usize, + num_classes: usize, + targets: Tensor<2, Int>, + input_lengths_len: usize, + target_lengths_len: usize, + ) { + assert!( + self.blank < num_classes, + "blank index {} must be less than num_classes {}", + self.blank, + num_classes + ); + assert_eq!( + targets.dims()[0], + batch_size, + "targets batch dimension {} must equal batch_size {}", + targets.dims()[0], + batch_size + ); + assert_eq!( + input_lengths_len, batch_size, + "input_lengths length {} must equal batch_size {}", + input_lengths_len, batch_size + ); + assert_eq!( + target_lengths_len, batch_size, + "target_lengths length {} must equal batch_size {}", + target_lengths_len, batch_size + ); + } +} + +#[cfg(test)] +mod tests { + use burn::tensor::{TensorData, Tolerance}; + + use super::*; + + // --------------------------------------------------------------- + // Assertions + // --------------------------------------------------------------- + + #[test] + #[should_panic(expected = "blank index")] + fn test_ctc_loss_panics_invalid_blank_index() { + let device = Default::default(); + // blank=5 is out of bounds for num_classes=3 + let ctc = CTCLossConfig::new().with_blank(5).init(); + + let log_probs = Tensor::<3>::zeros([2, 1, 3], &device); + let targets = Tensor::<2, Int>::from_data([[1]], &device); + let input_lengths = Tensor::<1, Int>::from_data([2], &device); + let target_lengths = Tensor::<1, Int>::from_data([1], &device); + + ctc.forward(log_probs, targets, input_lengths, target_lengths); + } + + #[test] + #[should_panic(expected = "must equal batch_size")] + fn test_ctc_loss_panics_mismatched_batch_size() { + let device = Default::default(); + let ctc = CTCLossConfig::new().init(); + + // Logits batch size = 2 + let log_probs = Tensor::<3>::zeros([2, 2, 3], &device); + // Targets batch size = 1 (Mismatch) + let targets = Tensor::<2, Int>::from_data([[1]], &device); + let input_lengths = Tensor::<1, Int>::from_data([2, 2], &device); + let target_lengths = Tensor::<1, Int>::from_data([1, 1], &device); + + ctc.forward(log_probs, targets, input_lengths, target_lengths); + } + + #[test] + #[should_panic(expected = "input_lengths length")] + fn test_ctc_loss_panics_input_lengths_mismatch() { + let device = Default::default(); + let ctc = CTCLossConfig::new().init(); + + // Logits batch size = 2 + let log_probs = Tensor::<3>::zeros([2, 2, 3], &device); + let targets = Tensor::<2, Int>::from_data([[1], [2]], &device); + + // Input lengths size = 1 (Mismatch) + let input_lengths = Tensor::<1, Int>::from_data([2], &device); + let target_lengths = Tensor::<1, Int>::from_data([1, 1], &device); + + ctc.forward(log_probs, targets, input_lengths, target_lengths); + } + + #[test] + #[should_panic(expected = "target_lengths length")] + fn test_ctc_loss_panics_target_lengths_mismatch() { + let device = Default::default(); + let ctc = CTCLossConfig::new().init(); + + // Logits batch size = 2 + let log_probs = Tensor::<3>::zeros([2, 2, 3], &device); + let targets = Tensor::<2, Int>::from_data([[1], [2]], &device); + let input_lengths = Tensor::<1, Int>::from_data([2, 2], &device); + + // Target lengths size = 1 (Mismatch) + let target_lengths = Tensor::<1, Int>::from_data([1], &device); + + ctc.forward(log_probs, targets, input_lengths, target_lengths); + } + + // --------------------------------------------------------------- + // Edge Case & Config Tests + // --------------------------------------------------------------- + + #[test] + fn test_ctc_loss_repeated_labels_minimum_input_length() { + // T=3, N=1, C=2, blank=0, target=[1, 1], uniform P = 1/2. + // + // The minimum T for target [1, 1] is 3: the only valid path is (1, 0, 1). + // prob = (1/2)^3 = 1/8 + // Loss = -ln(1/8) = 3 * ln(2) + let device = Default::default(); + let ctc = CTCLossConfig::new().init(); + + let log_probs = Tensor::<3>::full([3, 1, 2], 0.5_f32.ln(), &device); + let targets = Tensor::<2, Int>::from_data([[1_i32, 1]], &device); + let input_lengths = Tensor::<1, Int>::from_data([3_i32], &device); + let target_lengths = Tensor::<1, Int>::from_data([2_i32], &device); + + let loss = ctc.forward(log_probs, targets, input_lengths, target_lengths); + loss.into_data().assert_approx_eq::( + &TensorData::from([3.0 * 2.0_f32.ln()]), + Tolerance::rel_abs(1e-3, 1e-3), + ); + } + + #[test] + fn test_ctc_loss_custom_blank_uniform() { + // T=3, N=1, C=3, blank=2, target=[0, 1], uniform P = 1/3. + // + // Two distinct labels, 3 classes, 3 time steps, just with + // blank=2 instead of 0. + // 5 valid paths → total = 5/27 + // Loss = -ln(5/27) + let device = Default::default(); + let ctc = CTCLossConfig::new().with_blank(2).init(); + + let log_probs = Tensor::<3>::full([3, 1, 3], (1.0_f32 / 3.0).ln(), &device); + let targets = Tensor::<2, Int>::from_data([[0_i32, 1]], &device); + let input_lengths = Tensor::<1, Int>::from_data([3_i32], &device); + let target_lengths = Tensor::<1, Int>::from_data([2_i32], &device); + + let loss = ctc.forward(log_probs, targets, input_lengths, target_lengths); + + loss.into_data().assert_approx_eq::( + &TensorData::from([-(5.0_f32 / 27.0).ln()]), + Tolerance::default(), + ); + } + + // --------------------------------------------------------------- + // zero_infinity tests + // --------------------------------------------------------------- + + #[test] + fn test_ctc_loss_zero_infinity_produces_inf_when_disabled() { + // T=2, N=1, C=3, blank=0, target=[1, 1], input_length=2 + // Target [1, 1] requires at least 3 time steps → no valid paths → loss = +inf + let device = Default::default(); + let ctc = CTCLossConfig::new().with_zero_infinity(false).init(); + + let log_probs = Tensor::<3>::full([2, 1, 3], (1.0_f32 / 3.0).ln(), &device); + let targets = Tensor::<2, Int>::from_data([[1_i32, 1]], &device); + let input_lengths = Tensor::<1, Int>::from_data([2_i32], &device); + let target_lengths = Tensor::<1, Int>::from_data([2_i32], &device); + + let loss = ctc.forward(log_probs, targets, input_lengths, target_lengths); + let loss_data = loss.into_data().to_vec::().unwrap(); + assert!( + loss_data[0].is_infinite() && loss_data[0] > 0.0, + "Expected +inf, got {}", + loss_data[0] + ); + } + + #[test] + fn test_ctc_loss_zero_infinity_masks_inf_when_enabled() { + // Same inputs as above, but zero_infinity=true → loss should be 0.0 + let device = Default::default(); + let ctc = CTCLossConfig::new().with_zero_infinity(true).init(); + + let log_probs = Tensor::<3>::full([2, 1, 3], (1.0_f32 / 3.0).ln(), &device); + let targets = Tensor::<2, Int>::from_data([[1_i32, 1]], &device); + let input_lengths = Tensor::<1, Int>::from_data([2_i32], &device); + let target_lengths = Tensor::<1, Int>::from_data([2_i32], &device); + + let loss = ctc.forward(log_probs, targets, input_lengths, target_lengths); + + loss.into_data() + .assert_approx_eq::(&TensorData::from([0.0]), Tolerance::default()); + } + + #[test] + fn test_ctc_loss_zero_infinity_does_not_affect_finite_loss() { + // Verify that zero_infinity=true does not change a finite loss value. + let device = Default::default(); + let ctc = CTCLossConfig::new().with_zero_infinity(true).init(); + + let log_probs = Tensor::<3>::full([2, 1, 2], 0.5_f32.ln(), &device); + let targets = Tensor::<2, Int>::from_data([[1_i32]], &device); + let input_lengths = Tensor::<1, Int>::from_data([2_i32], &device); + let target_lengths = Tensor::<1, Int>::from_data([1_i32], &device); + + let loss = ctc.forward(log_probs, targets, input_lengths, target_lengths); + + loss.into_data() + .assert_approx_eq::(&TensorData::from([-(0.75_f32).ln()]), Tolerance::default()); + } +} + +#[cfg(test)] +mod pytorch_comparison_tests { + use super::*; + use burn::tensor::activation::log_softmax; + use burn_core::tensor::{Device, TensorData, Tolerance}; + + /// Deterministic logits: sin((t*7 + n*13 + c*3) * 0.1). + fn generate_logits(t_size: usize, n_size: usize, c_size: usize, device: &Device) -> Tensor<3> { + let mut data = Vec::with_capacity(t_size * n_size * c_size); + for t in 0..t_size { + for n in 0..n_size { + for c in 0..c_size { + data.push(((t * 7 + n * 13 + c * 3) as f32 * 0.1).sin()); + } + } + } + Tensor::<3>::from_data(TensorData::new(data, [t_size, n_size, c_size]), device) + } + + /// Runs a CTC forward + backward test and asserts against expected values from PyTorch. + /// + /// This helper performs the following steps: + /// 1. Generates deterministic logits using a sine-wave formula. + /// 2. Computes the CTC loss (forward pass). + /// 3. Asserts the computed loss matches `expected_losses`. + /// 4. Backpropagates the sum of the loss. + /// 5. Asserts the resulting gradients w.r.t. logits match `expected_grad_flat`. + /// + /// # Arguments + /// + /// - `expected_losses`: per-sample loss values from PyTorch (reduction='none'). + /// - `expected_grad_flat`: flattened gradient of sum(loss) w.r.t. logits. + #[allow(clippy::too_many_arguments)] + fn run_comparison( + label: &str, + t_size: usize, + n_size: usize, + c_size: usize, + targets_flat: Vec, + target_shape: [usize; 2], + input_lengths: Vec, + target_lengths: Vec, + blank: usize, + expected_losses: &[f32], + expected_grad_flat: &[f32], + ) { + let device = Device::default().autodiff(); + let ctc = CTCLossConfig::new().with_blank(blank).init(); + + let logits = generate_logits(t_size, n_size, c_size, &device).require_grad(); + let log_probs = log_softmax(logits.clone(), 2); + + let targets = + Tensor::<2, Int>::from_data(TensorData::new(targets_flat, target_shape), &device); + let input_lengths = + Tensor::<1, Int>::from_data(TensorData::new(input_lengths, [n_size]), &device); + let target_lengths = + Tensor::<1, Int>::from_data(TensorData::new(target_lengths, [n_size]), &device); + + let loss = ctc.forward(log_probs, targets, input_lengths, target_lengths); + + println!("=== {} ===", label); + println!(" Loss: {loss}"); + + let tolerance = Tolerance::rel_abs(1e-3, 1e-3); + loss.to_data() + .assert_approx_eq::(&TensorData::from(expected_losses), tolerance); + + let loss_sum = loss.sum(); + let grads = loss_sum.backward(); + let logits_grad = logits.grad(&grads).unwrap(); + + logits_grad + .reshape::<1, _>([-1]) + .into_data() + .assert_approx_eq::(&TensorData::from(expected_grad_flat), tolerance); + } + + #[test] + fn test_ctc_loss_uniform_input_lengths() { + // T=5, N=3, C=4, all input_lengths = 5 + // Expected losses and gradient from PyTorch + let expected_losses = [3.5236570835113525_f32, 3.495313882827759, 4.262677192687988]; + let expected_grad_flat = [ + -0.1679008007_f32, + -0.4595540464, + 0.2795598209, + 0.3478950262, + -0.3913056254, + -0.0832268298, + 0.2535884976, + 0.2209439576, + -0.0502742566, + 0.2766197622, + 0.2054125518, + -0.4317580462, + -0.0544800088, + -0.3144550920, + 0.0847885981, + 0.2841464877, + -0.1844545156, + -0.2063435912, + 0.2222184092, + 0.1685796976, + 0.0278018005, + 0.2657383382, + -0.0336986706, + -0.2598414719, + -0.0482986756, + -0.0098767160, + -0.1533526182, + 0.2115280181, + -0.1380317956, + -0.2198686600, + 0.2042596638, + 0.1536407918, + 0.0534787849, + 0.1819230020, + -0.2805589139, + 0.0451571345, + -0.0895631388, + 0.1996460557, + -0.2741115987, + 0.1640286744, + -0.2200077325, + -0.1693530381, + 0.2101601064, + 0.1792006642, + 0.0398471877, + -0.1131042913, + -0.2363226712, + 0.3095797896, + -0.2163617164, + 0.2740726173, + -0.2124865055, + 0.1547756046, + -0.4312027395, + -0.0446923785, + 0.2330704331, + 0.2428246588, + -0.0050083841, + -0.6256869435, + 0.2689785957, + 0.3617166877, + ]; + run_comparison( + "T=5, N=3, C=4 (uniform input lengths)", + 5, + 3, + 4, + vec![1, 2, 0, 1, 0, 0, 3, 2, 1], + [3, 3], + vec![5, 5, 5], + vec![2, 1, 3], + 0, + &expected_losses, + &expected_grad_flat, + ); + } + + #[test] + fn test_ctc_loss_repeated_labels() { + // T=8, N=4, C=6, includes consecutive repeated label [1,1,2] + // Expected losses and gradient from PyTorch + let expected_losses = [ + 8.84203052520752_f32, + 9.023029327392578, + 9.398024559020996, + 9.008068084716797, + ]; + let expected_grad_flat = [ + -0.2766432464, + -0.5202965736, + 0.1523768753, + 0.1896236390, + 0.2200277001, + 0.2349116206, + -0.1854365915, + 0.2031330466, + -0.4260218740, + 0.1678018719, + 0.1360142529, + 0.1045092493, + -0.6603536606, + 0.2278252542, + 0.1691786796, + 0.1262856424, + 0.0972681716, + 0.0397959016, + -0.0894432291, + -0.5457318425, + 0.1490373611, + 0.1462858170, + 0.1569476575, + 0.1829041988, + -0.2842915654, + -0.4220107496, + 0.1822281033, + 0.1889107376, + 0.1791101843, + 0.1560532600, + -0.1155678406, + 0.2295538932, + -0.2645366490, + -0.0288553704, + 0.1027252972, + 0.0766806602, + -0.5448347330, + 0.2031028718, + 0.1589304954, + 0.1322451383, + 0.1189499870, + -0.0683937520, + -0.0873993114, + -0.3051757514, + -0.2355299890, + 0.1586059481, + 0.2018169016, + 0.2676822543, + -0.3225219846, + -0.2611543834, + 0.1922984123, + 0.1632783115, + 0.1297036558, + 0.0983960181, + -0.1507159024, + 0.2256962359, + -0.1040333956, + -0.1514528394, + 0.0985243544, + 0.0819815546, + -0.2940836251, + 0.1586865336, + 0.1468491107, + 0.1485087872, + 0.1639631987, + -0.3239239752, + -0.0767390430, + -0.0434846729, + -0.4023587406, + -0.0052628326, + 0.2273432612, + 0.3005020022, + -0.2598774135, + -0.2188862711, + 0.1678501070, + 0.1352078766, + 0.1002781317, + 0.0754275694, + -0.1502914876, + 0.1930875033, + -0.0709601715, + -0.2219523191, + 0.1243555173, + 0.1257609427, + -0.0574148744, + 0.1152269915, + 0.1307857931, + 0.1599020809, + 0.2068412602, + -0.5553412437, + -0.0536844917, + 0.0758557543, + -0.2106334567, + -0.2509877980, + 0.1757438034, + 0.2637061775, + -0.1759711355, + -0.2431350052, + 0.1071053818, + 0.1259848624, + 0.1004033238, + 0.0856125653, + -0.1173698306, + 0.1213828772, + -0.1768893301, + -0.2070008069, + 0.1709136516, + 0.2089634240, + 0.0153109450, + 0.0967332721, + 0.1268781722, + 0.1706230640, + 0.2291058898, + -0.6386513710, + -0.0536664203, + 0.1378114969, + 0.0360041447, + -0.2989685237, + -0.0084722806, + 0.1872915775, + -0.1523490399, + -0.2111770809, + -0.0390694551, + 0.1366800815, + 0.1302325875, + 0.1356829405, + -0.0982905105, + -0.0127884001, + -0.3586881459, + -0.0259541404, + 0.2114149332, + 0.2843062580, + -0.0324133746, + 0.1084750593, + 0.1447229236, + 0.1862253845, + 0.2259712219, + -0.6329812407, + -0.1173689738, + 0.1914442331, + 0.1654772907, + -0.1376858056, + -0.2194855511, + 0.1176188141, + -0.1529908478, + -0.0606661662, + -0.3384291232, + 0.1524862647, + 0.1777049750, + 0.2218948901, + -0.0923086405, + -0.2855934799, + -0.3215619624, + 0.1726681292, + 0.2303666323, + 0.2964293361, + -0.2508065701, + 0.1479703039, + 0.1753441393, + 0.1917535067, + 0.1919818372, + -0.4562432170, + -0.2350299209, + 0.2257601619, + 0.1863904297, + 0.0388212129, + -0.2966264784, + 0.0806845874, + -0.1992894858, + 0.1068909168, + -0.5761897564, + 0.1624972969, + 0.2155302167, + 0.2905607820, + -0.1168124676, + -0.6870660186, + 0.1488010883, + 0.1881926507, + 0.2230074406, + 0.2438773215, + -0.5771554708, + 0.1980127096, + 0.1924194694, + 0.1714663208, + 0.1415647417, + -0.1263078004, + -0.3408652246, + 0.2292248607, + 0.1707807332, + 0.1269564927, + -0.2634142637, + 0.0773174241, + ]; + run_comparison( + "T=8, N=4, C=6 (repeated labels)", + 8, + 4, + 6, + vec![1, 1, 2, 0, 2, 3, 2, 1, 5, 0, 0, 0, 1, 2, 3, 4], + [4, 4], + vec![8, 8, 8, 8], + vec![3, 4, 1, 4], + 0, + &expected_losses, + &expected_grad_flat, + ); + } + + #[test] + fn test_ctc_loss_long_sequence() { + // T=10, N=2, C=8 + // Expected losses and gradient from PyTorch + let expected_losses = [12.629399299621582, 12.298524856567383]; + let expected_grad_flat = [ + -0.2570972741, + -0.6013792753, + 0.1061997041, + 0.1321590245, + 0.1533492655, + 0.1637226790, + 0.1598964781, + 0.1431493312, + -0.2540431321, + 0.1788398325, + -0.4038805366, + 0.1477340311, + 0.1197479516, + 0.0920107216, + 0.0686140805, + 0.0509770736, + -0.1364373565, + -0.3724762201, + 0.1489177048, + -0.0966964588, + 0.1463697106, + 0.1275274903, + 0.1033692732, + 0.0794258416, + -0.1771971881, + 0.2073454857, + -0.3109439015, + 0.1249521226, + -0.0101635465, + 0.0692621097, + 0.0533472970, + 0.0433975980, + -0.1398337185, + -0.0874802172, + 0.1705365479, + -0.2174201906, + 0.1150254831, + 0.0460043959, + 0.0647982135, + 0.0483694859, + -0.2332949787, + 0.1969220787, + -0.1270586401, + 0.1098557115, + -0.1364655048, + 0.0715296715, + 0.0553609394, + 0.0631506816, + -0.2169117928, + 0.0929956511, + 0.1624538749, + -0.2009791434, + 0.0904926360, + -0.0248185843, + 0.0532633252, + 0.0435040221, + -0.2313277274, + 0.1497355998, + -0.0024202778, + 0.1029939279, + -0.2776987851, + 0.0963881761, + 0.0351882279, + 0.1271408647, + -0.2590557337, + 0.1577988416, + 0.1429322213, + -0.1401246637, + 0.0866033062, + -0.1151762009, + 0.0683368817, + 0.0586853735, + -0.1322475076, + 0.0806737095, + 0.0528722852, + 0.0920089707, + -0.3037962914, + 0.1280544847, + -0.1391123086, + 0.2215466499, + -0.1918463260, + 0.1376975775, + 0.1160097718, + -0.0549413785, + 0.0970225409, + -0.2708687484, + 0.1147320047, + 0.0521945432, + -0.0504456684, + -0.0012221609, + 0.0644332916, + 0.0818370953, + -0.1036835983, + 0.1512031406, + -0.4072600305, + 0.2651379406, + -0.0681083873, + 0.0860663429, + 0.0810486302, + 0.0434282124, + 0.1056238264, + -0.2994530201, + 0.1729898751, + -0.1215954795, + -0.0481944978, + -0.1697723418, + 0.0725984722, + 0.0692019314, + 0.0859903544, + 0.1680216491, + -0.4071443677, + 0.2292988002, + -0.0205532499, + 0.0566616580, + 0.0326749459, + 0.0861379728, + 0.1142501161, + -0.0448331088, + 0.2054910213, + -0.4298293889, + -0.0647637174, + -0.4240962267, + 0.1013666242, + -0.0110451467, + 0.1519176364, + 0.1661346704, + -0.0719586164, + 0.1524447650, + -0.0496110357, + 0.0562372655, + -0.1889088154, + 0.1013496071, + 0.1339637935, + 0.1694275290, + 0.2007708699, + -0.4232292175, + -0.0401752405, + -0.2951072752, + 0.1443216652, + -0.2857291698, + 0.1489982456, + 0.1327733696, + 0.1096193567, + 0.0852990299, + -0.0413062274, + 0.0820900649, + -0.7903561592, + 0.1329460591, + 0.1535883099, + 0.1631743014, + 0.1585651338, + 0.1412984729, + -0.1033771932, + 0.1799504310, + 0.1697744429, + -0.5749052763, + 0.1189445183, + 0.0911802500, + 0.0679325759, + 0.0505003072, + ]; + run_comparison( + "T=10, N=2, C=8", + 10, + 2, + 8, + vec![1, 3, 5, 7, 2, 2, 4, 6, 1, 3], + [2, 5], + vec![10, 10], + vec![5, 5], + 0, + &expected_losses, + &expected_grad_flat, + ); + } + + #[test] + fn test_ctc_loss_mixed_input_lengths() { + // T=12, N=3, C=5, input_lengths=[12, 7, 10] + // Expected losses and gradient from PyTorch + let expected_losses = [10.595505714416504, 6.8078508377075195, 7.705057144165039]; + let expected_grad_flat = [ + -0.4790987670, + -0.2554937005, + 0.1991624236, + 0.2478453964, + 0.2875846624, + -0.3495813310, + 0.2268397957, + 0.2150714993, + -0.2442178279, + 0.1518878639, + -0.2764556706, + 0.2474014312, + -0.2137086987, + 0.1371368915, + 0.1056260392, + -0.2729502618, + -0.3609606028, + 0.2159237266, + 0.2238420397, + 0.1941450834, + -0.2953839302, + 0.1920599341, + 0.1974952668, + -0.2054278404, + 0.1112565696, + -0.1719199270, + 0.2299505472, + -0.2864859998, + 0.1497263014, + 0.0787290633, + -0.2035763413, + -0.3042884767, + 0.2126964629, + 0.1810975969, + 0.1140707731, + -0.2759391963, + 0.0975771844, + 0.1823379993, + -0.1112988219, + 0.1073228419, + -0.1336459517, + 0.1869296581, + -0.1996247321, + 0.1846873760, + -0.0383463502, + -0.2254105806, + -0.1834360659, + 0.1925925612, + 0.1462381780, + 0.0700158924, + -0.2259973884, + -0.0393539183, + 0.1802661419, + -0.0571591072, + 0.1422442794, + -0.0609069727, + 0.1089282706, + -0.0313654318, + 0.2186669111, + -0.2353227735, + -0.2840364873, + -0.0632198900, + 0.1755636632, + 0.1377806067, + 0.0339120962, + -0.1904856712, + -0.2139032930, + 0.1827126741, + 0.0056131603, + 0.2160631120, + -0.0243270602, + -0.0070458520, + 0.1070247591, + 0.2239368409, + -0.2995886803, + -0.2955487072, + 0.0309870224, + 0.1654911339, + 0.1581364125, + -0.0590658709, + -0.2191396207, + -0.3791662455, + 0.1803640425, + 0.1225430891, + 0.2953987718, + -0.0436352938, + -0.1575258970, + 0.1785279512, + 0.1756918877, + -0.1530586481, + -0.1834939867, + 0.0909025446, + 0.1423641294, + 0.1959712654, + -0.2457439601, + -0.3619639874, + -0.3929221630, + 0.1820438206, + 0.2454170734, + 0.3274252713, + -0.0628800318, + -0.2567180395, + 0.2112283260, + 0.0507859327, + 0.0575838275, + -0.0587697029, + 0.1174769849, + 0.0783569664, + 0.2290501744, + -0.3661144078, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + -0.0725664943, + -0.1532069892, + 0.2162397504, + -0.1248963475, + 0.1344300956, + -0.0362483934, + 0.1295878887, + -0.0502482466, + 0.2470482886, + -0.2901395261, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + -0.1349253207, + 0.0867646411, + 0.1998746395, + -0.2658679783, + 0.1141540110, + -0.0705668628, + 0.1519546807, + -0.2509805560, + 0.2475892603, + -0.0779965296, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + -0.2338010073, + 0.2471641302, + 0.1834627241, + -0.3026831448, + 0.1058573127, + -0.1155209392, + 0.1921830922, + -0.4129956067, + 0.2229512781, + 0.1133821756, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + -0.2636392713, + 0.2323469073, + -0.2913427949, + 0.1800564528, + 0.1425786912, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + 0.0000000000, + ]; + run_comparison( + "T=12, N=3, C=5 (mixed input lengths)", + 12, + 3, + 5, + vec![1, 4, 2, 0, 3, 1, 0, 0, 2, 4, 1, 3], + [3, 4], + vec![12, 7, 10], + vec![3, 2, 4], + 0, + &expected_losses, + &expected_grad_flat, + ); + } + + #[test] + fn test_ctc_loss_sum_reduction() { + // Same inputs as comparison_uniform_input_lengths, sum reduction + let device = Device::default().autodiff(); + let ctc = CTCLossConfig::new().init(); + + let logits = generate_logits(5, 3, 4, &device).require_grad(); + let log_probs = log_softmax(logits.clone(), 2); + let targets = Tensor::<2, Int>::from_data( + TensorData::new(vec![1_i32, 2, 0, 1, 0, 0, 3, 2, 1], [3, 3]), + &device, + ); + let il = Tensor::<1, Int>::from_data([5_i32, 5, 5], &device); + let tl = Tensor::<1, Int>::from_data([2_i32, 1, 3], &device); + + let loss = ctc.forward_with_reduction(log_probs, targets, il, tl, Reduction::Sum); + + let expected_sum = 11.2816486359_f32; // Expected value from PyTorch + loss.to_data() + .assert_approx_eq::(&TensorData::from([expected_sum]), Tolerance::default()); + + let grads = loss.backward(); + let logits_grad = logits.grad(&grads).unwrap(); + // Expected gradient from PyTorch + let expected_grad = [ + -0.1679008007_f32, + -0.4595540464, + 0.2795598209, + 0.3478950262, + -0.3913056254, + -0.0832268298, + 0.2535884976, + 0.2209439576, + -0.0502742566, + 0.2766197622, + 0.2054125518, + -0.4317580462, + -0.0544800088, + -0.3144550920, + 0.0847885981, + 0.2841464877, + -0.1844545156, + -0.2063435912, + 0.2222184092, + 0.1685796976, + 0.0278018005, + 0.2657383382, + -0.0336986706, + -0.2598414719, + -0.0482986756, + -0.0098767160, + -0.1533526182, + 0.2115280181, + -0.1380317956, + -0.2198686600, + 0.2042596638, + 0.1536407918, + 0.0534787849, + 0.1819230020, + -0.2805589139, + 0.0451571345, + -0.0895631388, + 0.1996460557, + -0.2741115987, + 0.1640286744, + -0.2200077325, + -0.1693530381, + 0.2101601064, + 0.1792006642, + 0.0398471877, + -0.1131042913, + -0.2363226712, + 0.3095797896, + -0.2163617164, + 0.2740726173, + -0.2124865055, + 0.1547756046, + -0.4312027395, + -0.0446923785, + 0.2330704331, + 0.2428246588, + -0.0050083841, + -0.6256869435, + 0.2689785957, + 0.3617166877, + ]; + logits_grad + .reshape::<1, _>([-1]) + .into_data() + .assert_approx_eq::(&TensorData::from(expected_grad), Tolerance::default()); + } + + #[test] + fn test_ctc_loss_mean_reduction() { + let device = Device::default().autodiff(); + let ctc = CTCLossConfig::new().init(); + + let logits = generate_logits(5, 3, 4, &device).require_grad(); + let log_probs = log_softmax(logits.clone(), 2); + let targets = Tensor::<2, Int>::from_data( + TensorData::new(vec![1_i32, 2, 0, 1, 0, 0, 3, 2, 1], [3, 3]), + &device, + ); + let il = Tensor::<1, Int>::from_data([5_i32, 5, 5], &device); + let tl = Tensor::<1, Int>::from_data([2_i32, 1, 3], &device); + + println!( + "{:?} {:?} {:?} {:?}", + log_probs.device(), + targets.device(), + il.device(), + tl.device() + ); + let loss = ctc.forward_with_reduction(log_probs, targets, il, tl, Reduction::Mean); + + let expected_mean = 2.2260115147_f32; // Expected value from PyTorch + loss.to_data() + .assert_approx_eq::(&TensorData::from([expected_mean]), Tolerance::default()); + + let grads = loss.backward(); + let logits_grad = logits.grad(&grads).unwrap(); + // Expected gradient from PyTorch + let expected_grad = [ + -0.0279834662_f32, + -0.0765923411, + 0.0465933047, + 0.0579825081, + -0.1304352134, + -0.0277422778, + 0.0845294967, + 0.0736479908, + -0.0055860290, + 0.0307355281, + 0.0228236169, + -0.0479731150, + -0.0090800021, + -0.0524091832, + 0.0141314333, + 0.0473577492, + -0.0614848398, + -0.0687812045, + 0.0740728080, + 0.0561932363, + 0.0030890885, + 0.0295264814, + -0.0037442972, + -0.0288712755, + -0.0080497796, + -0.0016461194, + -0.0255587716, + 0.0352546684, + -0.0460105985, + -0.0732895583, + 0.0680865571, + 0.0512135960, + 0.0059420872, + 0.0202136654, + -0.0311732125, + 0.0050174589, + -0.0149271907, + 0.0332743451, + -0.0456852652, + 0.0273381118, + -0.0733359158, + -0.0564510152, + 0.0700533763, + 0.0597335547, + 0.0044274656, + -0.0125671430, + -0.0262580756, + 0.0343977548, + -0.0360602848, + 0.0456787720, + -0.0354144201, + 0.0257959347, + -0.1437342465, + -0.0148974592, + 0.0776901469, + 0.0809415579, + -0.0005564869, + -0.0695207715, + 0.0298865121, + 0.0401907414, + ]; + logits_grad + .reshape::<1, _>([-1]) + .into_data() + .assert_approx_eq::(&TensorData::from(expected_grad), Tolerance::default()); + } +} diff --git a/crates/burn-nn/src/loss/huber.rs b/crates/burn-nn/src/loss/huber.rs new file mode 100644 index 0000000..f4fc5f5 --- /dev/null +++ b/crates/burn-nn/src/loss/huber.rs @@ -0,0 +1,208 @@ +use burn_core as burn; + +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::{config::Config, module::Module}; + +use super::Reduction; + +/// Configuration to create a [Huber loss](HuberLoss). +#[derive(Config, Debug)] +pub struct HuberLossConfig { + /// The bound where the Huber loss function changes from quadratic to linear behaviour. + pub delta: f32, +} + +impl HuberLossConfig { + /// Initialize [Huber loss](HuberLoss). + pub fn init(&self) -> HuberLoss { + self.assertions(); + HuberLoss { + delta: self.delta, + lin_bias: self.delta * self.delta * 0.5, + } + } + + fn assertions(&self) { + assert!( + self.delta >= 0., // This also tests for normality + "Delta for Huber loss must be a non-negative number." + ); + } +} + +/// Calculate the Huber loss between the inputs and the target. +/// +/// The loss for each element of the residuals `r = targets - predictions` is given by +/// +/// ```text +/// L(r) = 0.5 * r^2 if |r| <= d +/// L(r) = 0.5 * d^2 + d * (|r| - d) if |r| > d +/// ``` +/// +/// where `d` is the configured `delta`. In particular, this is equal to the +/// [L2 Loss](super::MseLoss) for residuals with magnitude smaller than `delta`, +/// but behaves linearly instead of quadratically for large residuals. +/// +/// This loss function is less sensitive to outliers than the mean squared error loss. +/// +/// See also: +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct HuberLoss { + /// The bound where the Huber loss function changes from quadratic to linear behaviour. + pub delta: f32, + /// Precomputed value for the linear bias. + pub lin_bias: f32, // delta * delta * 0.5 precomputed +} + +impl ModuleDisplay for HuberLoss { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("delta", &self.delta) + .add("lin_bias", &self.lin_bias) + .optional() + } +} + +impl HuberLoss { + /// Compute the loss element-wise for the predictions and targets, then reduce + /// to a single loss value. + /// + /// `Reduction::Auto` behaves as `Reduction::Mean`. + /// + /// # Shapes + /// + /// - predictions: \[...dims\] + /// - targets: \[...dims\] + /// - output: \[1\] + pub fn forward( + &self, + predictions: Tensor, + targets: Tensor, + reduction: Reduction, + ) -> Tensor<1> { + let loss = self.forward_no_reduction(predictions, targets); + match reduction { + Reduction::Mean | Reduction::Auto => loss.mean(), + Reduction::Sum => loss.sum(), + other => panic!("{other:?} reduction is not supported"), + } + } + /// Compute the loss element-wise for the predictions and targets. + /// + /// # Shapes + /// + /// - predictions: [...dims] + /// - targets: [...dims] + /// - output: [...dims] + pub fn forward_no_reduction( + &self, + predictions: Tensor, + targets: Tensor, + ) -> Tensor { + let residuals = targets - predictions; + self.forward_residuals(residuals) + } + /// Compute the loss element-wise for the given residuals. + /// + /// # Shapes + /// + /// - residuals: [...dims] + /// - output: [...dims] + pub fn forward_residuals(&self, residuals: Tensor) -> Tensor { + let is_large = residuals.clone().abs().greater_elem(self.delta); + // We are interested in `sign(r)` when `abs(r) > self.delta`. Note that the + // `sign()` function, in general, suffers from a jump at 0. + // Instead the following tensor implements `delta * sign(r)` for values outside + // the bound: + let softsign = residuals.clone().clamp(-self.delta, self.delta); + + // 0.5 * d^2 + d * (|r| - d) = + // d * |r| - 0.5 * d^2 + // Moreover |r| = sign(r) * r + let outside = softsign.mul(residuals.clone()).sub_scalar(self.lin_bias); + + let inside = residuals.square().mul_scalar(0.5); + inside.mask_where(is_large, outside) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn test_huber_loss() { + let predict = TensorData::from([-2., -0.5, 0., 0.3, 1.]); + let targets = TensorData::from([0., 0., 0., 0., 0.]); + + let device = Default::default(); + + let predict = Tensor::<1>::from_data(predict, &device); + let targets = Tensor::<1>::from_data(targets, &device); + + let huber = HuberLossConfig::new(0.5).init(); + + let loss_sum = huber.forward(predict.clone(), targets.clone(), Reduction::Sum); + let loss = huber.forward(predict.clone(), targets.clone(), Reduction::Auto); + let loss_no_reduction = huber.forward_no_reduction(predict, targets); + + let expected = TensorData::from([0.875, 0.125, 0., 0.045, 0.375]); + loss_no_reduction + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([0.284]); + loss.into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([1.42]); + loss_sum + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[cfg(feature = "std")] + #[test] + fn test_huber_ad_loss() { + use burn::tensor::Device; + let device = Device::default().autodiff(); + let predict = TensorData::from([-2., -0.5, 0., 0.3, 1.]); + let targets = TensorData::from([0., 0., 0., 0., 0.]); + + let predict = Tensor::<1>::from_data(predict, &device).require_grad(); + let targets = Tensor::<1>::from_data(targets, &device); + + let loss = HuberLossConfig::new(0.5).init(); + let loss = loss.forward_no_reduction(predict.clone(), targets); + + let grads = loss.backward(); + let grads_predict = predict.grad(&grads).unwrap(); + + let expected = TensorData::from([-0.5, -0.5, 0., 0.3, 0.5]); + grads_predict + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn display() { + let config = HuberLossConfig::new(0.5); + let loss = config.init(); + + assert_eq!( + alloc::format!("{loss}"), + "HuberLoss {delta: 0.5, lin_bias: 0.125}" + ); + } +} diff --git a/crates/burn-nn/src/loss/kldiv.rs b/crates/burn-nn/src/loss/kldiv.rs new file mode 100644 index 0000000..324be47 --- /dev/null +++ b/crates/burn-nn/src/loss/kldiv.rs @@ -0,0 +1,200 @@ +use burn_core as burn; + +use super::Reduction; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::{config::Config, module::Module}; + +/// Configuration to create a [KLDiv loss](KLDivLoss). +#[derive(Config, Debug)] +pub struct KLDivLossConfig { + /// Specifies whether target is the log space. Default: False. + #[config(default = false)] + pub log_target: bool, +} + +impl KLDivLossConfig { + /// Initialize [KLDiv Loss](KLDivLoss). + pub fn init(&self) -> KLDivLoss { + KLDivLoss { + log_target: self.log_target, + } + } +} + +/// Kullback-Leibler Divergence Loss +/// +/// KL Divergence shows the difference between two probability distributions by measuring information loss +/// +/// KLDivLoss = +/// ```tex +/// y_{true} \cdot (\log{y_{true}} - \log{y_{pred}}) +/// ``` +/// By default, the loss expects the input in the log-space. +/// The targets may also be provided in the log-space if `log_target` is true. +/// +/// See +/// - [Kullback–Leibler divergence](https://en.wikipedia.org/wiki/Kullback-Leibler_divergence) +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct KLDivLoss { + /// Specifies whether target is the log space. Default: False. + pub log_target: bool, +} + +impl ModuleDisplay for KLDivLoss { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content.add("log_target", &self.log_target).optional() + } +} + +impl KLDivLoss { + /// Compute the criterion on the input tensor. + /// + /// `Reduction::Auto` behaves as `Reduction::BatchMean`,`Reduction::Mean` dose not align with the math definition. + /// + /// # Shapes + /// + /// - predictions: \[batch_size,num_targets\] + /// - targets: \[batch_size,num_targets\] + /// - output: \[1\] + pub fn forward( + &self, + predictions: Tensor, + targets: Tensor, + reduction: Reduction, + ) -> Tensor<1> { + let loss = self.forward_no_reduction(predictions, targets); + match reduction { + Reduction::BatchMean | Reduction::Auto => { + let batch_size = loss.dims()[0] as f32; + loss.sum().div_scalar(batch_size) + } + Reduction::Mean => loss.mean(), + Reduction::Sum => loss.sum(), + } + } + /// Compute the criterion on the input tensor without reducing. + pub fn forward_no_reduction( + &self, + predictions: Tensor, + targets: Tensor, + ) -> Tensor { + match self.log_target { + true => targets.clone().exp().mul(targets.sub(predictions)), + false => { + let epsilon = targets + .dtype() + .finfo() + .unwrap_or(burn::tensor::FloatDType::F32.finfo()) + .min_positive; + let log_target = targets.clone().clamp(epsilon, 1.0).log(); + targets.mul(log_target.sub(predictions)) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn test_kl_div_loss() { + let predict = TensorData::from([[-1.0, -0.5], [-2.0, -0.2]]); + let targets = TensorData::from([[0.4, 0.6], [0.1, 0.9]]); + + let device = Default::default(); + let predict = Tensor::<2>::from_data(predict, &device); + let targets = Tensor::<2>::from_data(targets, &device); + + let kl_loss = KLDivLossConfig { log_target: false }.init(); + + let loss_sum = kl_loss.forward(predict.clone(), targets.clone(), Reduction::Sum); + let loss_batch_mean = + kl_loss.forward(predict.clone(), targets.clone(), Reduction::BatchMean); + let loss_no_reduction = kl_loss.forward_no_reduction(predict, targets); + + let expected_no_reduction = + TensorData::from([[0.0334837139, -0.0064953566], [-0.0302585065, 0.0851755068]]); + loss_no_reduction + .into_data() + .assert_approx_eq::(&expected_no_reduction, Tolerance::absolute(1e-5)); + + let expected_sum = TensorData::from([0.08191]); + loss_sum + .into_data() + .assert_approx_eq::(&expected_sum, Tolerance::absolute(1e-5)); + + let expected_batch_mean = TensorData::from([0.04095]); + loss_batch_mean + .into_data() + .assert_approx_eq::(&expected_batch_mean, Tolerance::absolute(1e-5)); + } + + #[test] + fn test_kl_div_loss_log_target() { + let device = Default::default(); + let predict = Tensor::<1>::from_data([-1.0, -2.0], &device); + let targets = Tensor::<1>::from_data([-0.5, -1.5], &device); + + let kl_loss = KLDivLossConfig { log_target: true }.init(); + + let loss_no_reduction = kl_loss.forward_no_reduction(predict.clone(), targets.clone()); + let expected_none = TensorData::from([0.3032653299, 0.1115650801]); + loss_no_reduction + .into_data() + .assert_approx_eq::(&expected_none, Tolerance::absolute(1e-5)); + + let loss_batch_mean = + kl_loss.forward(predict.clone(), targets.clone(), Reduction::BatchMean); + let expected_bm = TensorData::from([0.207415204965]); + loss_batch_mean + .into_data() + .assert_approx_eq::(&expected_bm, Tolerance::absolute(1e-5)); + + let loss_sum = kl_loss.forward(predict, targets, Reduction::Sum); + let expected_sum = TensorData::from([0.414830409931]); + loss_sum + .into_data() + .assert_approx_eq::(&expected_sum, Tolerance::absolute(1e-5)); + } + + #[cfg(feature = "std")] + #[test] + fn test_kl_div_ad_loss() { + use burn::tensor::Device; + let device = Device::default().autodiff(); + let predict = Tensor::<2>::from_data([[-1.0, -0.5]], &device).require_grad(); + let targets = Tensor::<2>::from_data([[0.4, 0.6]], &device); + + let kl_loss = KLDivLossConfig { log_target: false }.init(); + let loss = kl_loss.forward(predict.clone(), targets, Reduction::Sum); + + let grads = loss.backward(); + let grads_predict = predict.grad(&grads).unwrap(); + + // d/d_pred [target * (log_target - pred)] = -target + let expected = TensorData::from([[-0.4, -0.6]]); + grads_predict + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn display() { + let config = KLDivLossConfig { log_target: true }; + let loss = config.init(); + + assert_eq!(alloc::format!("{loss}"), "KLDivLoss {log_target: true}"); + } +} diff --git a/crates/burn-nn/src/loss/lp_loss.rs b/crates/burn-nn/src/loss/lp_loss.rs new file mode 100644 index 0000000..b03357c --- /dev/null +++ b/crates/burn-nn/src/loss/lp_loss.rs @@ -0,0 +1,636 @@ +use super::Reduction; +use burn::config::Config; +use burn::module::Module; +use burn::tensor::Tensor; +use burn_core as burn; + +/// Configuration for the [Lp Loss](LpLoss) module. +/// +/// # Example +/// +/// ```ignore +/// use burn_nn::loss::{LpLossConfig, Reduction}; +/// +/// // Create L1 loss (MAE when using mean reduction) +/// let l1_loss = LpLossConfig::l1(); +/// +/// // Create L2 loss (MSE when using mean reduction) +/// let l2_loss = LpLossConfig::l2(); +/// +/// // Create custom Lp loss with p=3 +/// let l3_loss = LpLossConfig::new(3.0).init(); +/// ``` +#[derive(Config, Debug)] +pub struct LpLossConfig { + /// The exponent `p` determining the type of error measurement. + /// + /// Common values: + /// - `p = 1.0`: L1 loss (MAE with mean reduction) - robust to outliers + /// - `p = 2.0`: L2 loss (MSE with mean reduction) - standard choice, differentiable everywhere + /// - `p > 2.0`: Increasingly sensitive to large errors (outliers) + /// - `0 < p < 1`: More robust to outliers than L1 (quasi-norm) + pub p: f64, +} + +impl LpLossConfig { + /// Initializes a [Lp Loss](LpLoss) module. + /// + /// # Panics + /// + /// Panics if `p <= 0`. + pub fn init(&self) -> LpLoss { + self.assertions(); + LpLoss { p: self.p } + } + + /// Creates L1 loss (p=1). + /// + /// When used with `Reduction::Mean`, this computes Mean Absolute Error (MAE). + /// When used with `Reduction::Sum`, this computes Sum of Absolute Errors (SAE). + pub fn l1() -> LpLoss { + LpLoss { p: 1.0 } + } + + /// Creates L2 loss (p=2). + /// + /// When used with `Reduction::Mean`, this computes Mean Squared Error (MSE). + /// When used with `Reduction::Sum`, this computes Sum of Squared Errors (SSE). + pub fn l2() -> LpLoss { + LpLoss { p: 2.0 } + } + + fn assertions(&self) { + assert!(self.p > 0.0, "The order of the norm p must be positive.") + } +} + +/// Computes the Lp Loss between predictions and targets. +/// +/// This loss function computes the element-wise p-th power of absolute errors, +/// then reduces them via mean or sum. +/// +/// # Mathematical Definition +/// +/// For predictions `ŷ` and targets `y`, the element-wise loss is: +/// +/// ```text +/// Lᵢ = |ŷᵢ - yᵢ|ᵖ +/// ``` +/// +/// With mean reduction (default), the final loss is: +/// +/// ```text +/// L = (1/n) × Σᵢ |ŷᵢ - yᵢ|ᵖ +/// ``` +/// +/// # Notes +/// +/// - This implementation computes `|error|^p`, **not** the Lp norm `(Σ|error|^p)^(1/p)`. +/// - The `p = 1` case uses an optimized `abs()` operation. +/// - The `p = 2` case uses an optimized computation `error * error` instead of `powf`. +/// +/// # Example +/// +/// ```ignore +/// use burn_nn::loss::{LpLossConfig, Reduction}; +/// use burn::tensor::Tensor; +/// +/// // Create L2 loss +/// let l2_loss = LpLossConfig::l2(); +/// +/// let predictions: Tensor = /* model output */; +/// let targets: Tensor = /* ground truth */; +/// +/// // Compute loss with mean reduction (MSE) +/// let mse = l2_loss.forward(predictions.clone(), targets.clone(), Reduction::Mean); +/// +/// // Compute loss with sum reduction (SSE) +/// let sse = l2_loss.forward(predictions.clone(), targets.clone(), Reduction::Sum); +/// +/// // Compute loss with no reduction +/// let unreduced_l2_loss = l2_loss.forward_no_reduction(predictions, targets); +/// ``` +#[derive(Module, Debug)] +pub struct LpLoss { + /// The order of the norm (e.g., 1 for L1, 2 for L2). + /// Equivalently, the exponent `p` for computing `|error|^p`. + pub p: f64, +} + +impl LpLoss { + /// Computes the element-wise loss `|error|^p` with reduction. + /// + /// # Arguments + /// + /// * `predictions` - The model's predicted values. + /// * `targets` - The ground truth target values. + /// * `reduction` - Specifies how to reduce the element-wise losses: + /// - `Reduction::Mean` or `Reduction::Auto`: Returns the mean of all element-wise losses. + /// - `Reduction::Sum`: Returns the sum of all element-wise losses. + /// + /// # Returns + /// + /// A scalar tensor containing the reduced loss value. + /// + /// # Shapes + /// + /// - predictions: `[...dims]` - Any shape + /// - targets: `[...dims]` - Must match predictions shape + /// - output: `[1]` - Scalar loss value + pub fn forward( + &self, + predictions: Tensor, + targets: Tensor, + reduction: Reduction, + ) -> Tensor<1> { + let unreduced_loss = self.forward_no_reduction(predictions, targets); + + match reduction { + Reduction::Mean | Reduction::Auto => unreduced_loss.mean(), + Reduction::Sum => unreduced_loss.sum(), + other => panic!("{other:?} reduction is not supported"), + } + } + + /// Computes the element-wise loss `|error|^p` without reduction. + /// + /// # Arguments + /// + /// * `predictions` - The model's predicted values. + /// * `targets` - The ground truth target values. + /// + /// # Returns + /// + /// A tensor of the same shape as the inputs, containing `|prediction - target|^p` + /// for each element. + /// + /// # Shapes + /// + /// - predictions: `[...dims]` - Any shape + /// - targets: `[...dims]` - Must match predictions shape + /// - output: `[...dims]` - Same shape as inputs + pub fn forward_no_reduction( + &self, + predictions: Tensor, + targets: Tensor, + ) -> Tensor { + let error = predictions.sub(targets); + + // Use simplified/optimized expressions for common cases (p = 1, p = 2) + if self.p == 1.0 { + // L1 loss + error.abs() + } else if self.p == 2.0 { + // L2 loss + error.clone().mul(error) + } else { + error.abs().powf_scalar(self.p) + } + } + + /// Computes the element-wise loss `|error|^p` with reduction over specified dimensions. + /// + /// Calculates element-wise `|predictions - targets|^p`, then takes the mean + /// over the specified dimensions. Useful for per-sample or per-channel losses (e.g., when + /// working with images). + /// + /// Dimensions can be provided in any order. They are sorted internally and + /// reduced from highest to lowest to ensure indices remain valid. + /// + /// # Arguments + /// + /// * `predictions` - The model's predicted values. + /// * `targets` - The ground truth target values. + /// * `dims` - Dimensions to reduce over. + /// + /// # Returns + /// + /// A tensor with the specified dimensions reduced to size 1. + /// + /// # Example + /// + /// ```ignore + /// // Image tensor: [batch, C, H, W] + /// let l2_loss = LpLossConfig::l2(); + /// + /// // Per-image MSE for PSNR: reduce over C, H, W → [batch, 1, 1, 1] + /// let mse_per_image = l2_loss.forward_reduce_dims(predictions, targets, &[1, 2, 3]); + /// ``` + pub fn forward_reduce_dims( + &self, + predictions: Tensor, + targets: Tensor, + dims: &[usize], + ) -> Tensor { + let error = self.forward_no_reduction(predictions, targets); + + // Sort the dimensions to ascending order + let mut sorted_dims = dims.to_vec(); + sorted_dims.sort(); + + // Reduce over specified dimensions + error.mean_dims(sorted_dims.as_slice()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn test_lp_loss_l1_constructor() { + let loss_func_l1 = LpLossConfig::l1(); + let loss_func_p1 = LpLossConfig::new(1.0).init(); + assert_eq!(loss_func_l1.p, 1.0); + assert_eq!(loss_func_l1.p, loss_func_p1.p); + } + + #[test] + fn test_lp_loss_l2_constructor() { + let loss_func_l2 = LpLossConfig::l2(); + let loss_func_p2 = LpLossConfig::new(2.0).init(); + assert_eq!(loss_func_l2.p, 2.0); + assert_eq!(loss_func_l2.p, loss_func_p2.p); + } + + #[test] + fn test_lp_loss_l1() { + let device = Default::default(); + let predictions = + Tensor::<2>::from_data(TensorData::from([[1.0, 2.0], [3.0, 4.0]]), &device); + + let targets = Tensor::<2>::from_data(TensorData::from([[2.0, 1.0], [3.0, 2.0]]), &device); + + let loss_func = LpLossConfig::l1(); + let loss_no_reduction = + loss_func.forward_no_reduction(predictions.clone(), targets.clone()); + let loss_auto = loss_func.forward(predictions.clone(), targets.clone(), Reduction::Auto); + let loss_sum = loss_func.forward(predictions, targets, Reduction::Sum); + + let expected = TensorData::from([[1.0, 1.0], [0.0, 2.0]]); + loss_no_reduction.into_data().assert_eq(&expected, false); + + let expected = TensorData::from([1.0]); + loss_auto.into_data().assert_eq(&expected, false); + + let expected = TensorData::from([4.0]); + loss_sum.into_data().assert_eq(&expected, false); + } + + #[test] + fn test_lp_loss_l2() { + let device = Default::default(); + let predictions = + Tensor::<2>::from_data(TensorData::from([[1.0, 2.0], [3.0, 4.0]]), &device); + + let targets = Tensor::<2>::from_data(TensorData::from([[2.0, 1.0], [3.0, 2.0]]), &device); + + let loss_func = LpLossConfig::l2(); + let loss_no_reduction = + loss_func.forward_no_reduction(predictions.clone(), targets.clone()); + let loss_auto = loss_func.forward(predictions.clone(), targets.clone(), Reduction::Auto); + let loss_sum = loss_func.forward(predictions, targets, Reduction::Sum); + + let expected = TensorData::from([[1.0, 1.0], [0.0, 4.0]]); + loss_no_reduction.into_data().assert_eq(&expected, false); + + let expected = TensorData::from([1.5]); + loss_auto.into_data().assert_eq(&expected, false); + + let expected = TensorData::from([6.0]); + loss_sum.into_data().assert_eq(&expected, false); + } + + #[test] + fn test_lp_loss_p_half() { + // L0.5 quasi-norm: more robust to outliers than L1 + let device = Default::default(); + let predictions = + Tensor::<2>::from_data(TensorData::from([[1.0, 2.0], [3.0, 4.0]]), &device); + + let targets = Tensor::<2>::from_data(TensorData::from([[2.0, 1.0], [3.0, 0.0]]), &device); + + let loss_func = LpLossConfig::new(0.5).init(); + let loss_no_reduction = + loss_func.forward_no_reduction(predictions.clone(), targets.clone()); + let loss_auto = loss_func.forward(predictions.clone(), targets.clone(), Reduction::Auto); + let loss_sum = loss_func.forward(predictions, targets, Reduction::Sum); + + // |1-2|^0.5 = 1, |2-1|^0.5 = 1, |3-3|^0.5 = 0, |4-0|^0.5 = 2 + let expected = TensorData::from([[1.0, 1.0], [0.0, 2.0]]); + loss_no_reduction.into_data().assert_eq(&expected, false); + + let expected = TensorData::from([1.0]); + loss_auto.into_data().assert_eq(&expected, false); + + let expected = TensorData::from([4.0]); + loss_sum.into_data().assert_eq(&expected, false); + } + + #[test] + fn test_lp_loss_p3() { + // L3 norm: more sensitive to outliers than L2 + let device = Default::default(); + let predictions = + Tensor::<2>::from_data(TensorData::from([[1.0, 2.0], [3.0, 4.0]]), &device); + + let targets = Tensor::<2>::from_data(TensorData::from([[2.0, 1.0], [3.0, 2.0]]), &device); + + let loss_func = LpLossConfig::new(3.0).init(); + let loss_no_reduction = + loss_func.forward_no_reduction(predictions.clone(), targets.clone()); + let loss_auto = loss_func.forward(predictions.clone(), targets.clone(), Reduction::Auto); + let loss_sum = loss_func.forward(predictions, targets, Reduction::Sum); + + // |1-2|^3 = 1, |2-1|^3 = 1, |3-3|^3 = 0, |4-2|^3 = 8 + let expected = TensorData::from([[1.0, 1.0], [0.0, 8.0]]); + loss_no_reduction.into_data().assert_eq(&expected, false); + + let expected = TensorData::from([2.5]); + loss_auto.into_data().assert_eq(&expected, false); + + let expected = TensorData::from([10.0]); + loss_sum.into_data().assert_eq(&expected, false); + } + + #[test] + fn test_lp_loss_zero_error() { + // Test when predictions exactly match targets + let device = Default::default(); + let predictions = + Tensor::<2>::from_data(TensorData::from([[1.0, 2.0], [3.0, 4.0]]), &device); + + let targets = predictions.clone(); + + let loss_func_l1 = LpLossConfig::l1(); + let loss_func_l2 = LpLossConfig::l2(); + + let l1_loss = loss_func_l1.forward(predictions.clone(), targets.clone(), Reduction::Auto); + let l2_loss = loss_func_l2.forward(predictions, targets, Reduction::Auto); + + let expected = TensorData::from([0.0]); + l1_loss.into_data().assert_eq(&expected, false); + l2_loss.into_data().assert_eq(&expected, false); + } + + #[test] + fn test_lp_loss_negative_errors() { + // Test that negative errors are handled correctly (absolute value) + let device = Default::default(); + let predictions = Tensor::<1>::from_data(TensorData::from([1.0, 2.0, 3.0]), &device); + let targets = Tensor::<1>::from_data(TensorData::from([3.0, 4.0, 5.0]), &device); + let loss_func_l1 = LpLossConfig::l1(); + let loss_func_p1 = LpLossConfig::new(1.0).init(); + + let loss_no_reduction_l1 = + loss_func_l1.forward_no_reduction(predictions.clone(), targets.clone()); + let loss_no_reduction_p1 = loss_func_p1.forward_no_reduction(predictions, targets); + + // All errors are negative: 1-3=-2, 2-4=-2, 3-5=-2, but |error| = 2 + let expected = TensorData::from([2.0, 2.0, 2.0]); + loss_no_reduction_l1.into_data().assert_eq(&expected, false); + loss_no_reduction_p1.into_data().assert_eq(&expected, false); + } + + #[test] + fn test_lp_loss_3d_tensor() { + let device = Default::default(); + let predictions = Tensor::<3>::from_data( + TensorData::from([[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]), + &device, + ); + let targets = Tensor::<3>::from_data( + TensorData::from([[[0.0, 2.0], [3.0, 5.0]], [[4.0, 6.0], [7.0, 10.0]]]), + &device, + ); + let loss_func_l2 = LpLossConfig::l2(); + let loss_func_p2 = LpLossConfig::new(2.0).init(); + + let loss_l2 = loss_func_l2.forward(predictions.clone(), targets.clone(), Reduction::Auto); + let loss_p2 = loss_func_p2.forward(predictions, targets, Reduction::Auto); + + // Errors: 1, 0, 0, -1, 1, 0, 0, -2 + // Squared: 1, 0, 0, 1, 1, 0, 0, 4 + // Mean: 7/8 = 0.875 + let expected = TensorData::from([0.875]); + loss_l2.into_data().assert_eq(&expected, false); + loss_p2.into_data().assert_eq(&expected, false); + } + + #[test] + #[should_panic(expected = "The order of the norm p must be positive.")] + fn test_lp_loss_negative_p_panics() { + let _ = LpLossConfig::new(-1.0).init(); + } + + #[test] + #[should_panic(expected = "The order of the norm p must be positive.")] + fn test_lp_loss_zero_p_panics() { + let _ = LpLossConfig::new(0.0).init(); + } + + #[test] + fn test_lp_loss_fractional_p() { + // Test p = 1.5 + let device = Default::default(); + let predictions = Tensor::<1>::from_data(TensorData::from([0.0, 4.0]), &device); + + let targets = Tensor::<1>::from_data(TensorData::from([1.0, 0.0]), &device); + + let loss_func = LpLossConfig::new(1.5).init(); + let loss_no_reduction = loss_func.forward_no_reduction(predictions, targets); + + // |0-1|^1.5 = 1, |4-0|^1.5 = 8 + let expected = TensorData::from([1.0, 8.0]); + loss_no_reduction.into_data().assert_eq(&expected, false); + } + + #[test] + fn test_forward_reduce_dims_single_dim() { + let device = Default::default(); + // Shape: [2, 3] + let predictions = Tensor::<2>::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), + &device, + ); + let targets = Tensor::<2>::from_data( + TensorData::from([[0.0, 2.0, 6.0], [1.0, 5.0, 6.0]]), + &device, + ); + let loss_func_l2 = LpLossConfig::l2(); + let loss_func_p2 = LpLossConfig::new(2.0).init(); + + // Reduce over dim 1 -> should give [2, 1] shape + let loss_l2 = loss_func_l2.forward_reduce_dims(predictions.clone(), targets.clone(), &[1]); + let loss_p2 = loss_func_p2.forward_reduce_dims(predictions, targets, &[1]); + + // Errors row 0: [1, 0, -3] -> squared: [1, 0, 9] -> mean: 10/3 + // Errors row 1: [3, 0, 0] -> squared: [9, 0, 0] -> mean: 3 + let expected = TensorData::from([[10.0 / 3.0], [3.0]]); + loss_l2 + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + loss_p2 + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn test_forward_reduce_dims_first_dim() { + let device = Default::default(); + // Shape: [2, 3] + let predictions = Tensor::<2>::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), + &device, + ); + let targets = Tensor::<2>::from_data( + TensorData::from([[0.0, 2.0, 6.0], [1.0, 5.0, 6.0]]), + &device, + ); + let loss_func = LpLossConfig::l2(); + + // Reduce over dim 0 -> should give [1, 3] shape + let loss = loss_func.forward_reduce_dims(predictions, targets, &[0]); + + // Squared errors: [[1, 0, 9], [9, 0, 0]] + // Mean over dim 0: [5, 0, 4.5] + let expected = TensorData::from([[5.0, 0.0, 4.5]]); + loss.into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn test_forward_reduce_dims_multiple_dims() { + let device = Default::default(); + // Shape: [2, 2, 2] + let predictions = Tensor::<3>::from_data( + TensorData::from([[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]), + &device, + ); + let targets = Tensor::<3>::from_data( + TensorData::from([[[0.0, 2.0], [3.0, 6.0]], [[4.0, 6.0], [7.0, 10.0]]]), + &device, + ); + let loss_func = LpLossConfig::l2(); + + // Reduce over dims 1 and 2 -> should give [2, 1, 1] shape + let loss = loss_func.forward_reduce_dims(predictions, targets, &[1, 2]); + + // Batch 0 errors: [[1, 0], [0, -2]] -> squared: [[1, 0], [0, 4]] -> mean: 5/4 = 1.25 + // Batch 1 errors: [[1, 0], [0, -2]] -> squared: [[1, 0], [0, 4]] -> mean: 5/4 = 1.25 + let expected = TensorData::from([[[1.25]], [[1.25]]]); + loss.into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn test_forward_reduce_dims_all_dims() { + let device = Default::default(); + // Shape: [2, 2] + let predictions = + Tensor::<2>::from_data(TensorData::from([[1.0, 2.0], [3.0, 4.0]]), &device); + let targets = Tensor::<2>::from_data(TensorData::from([[2.0, 1.0], [3.0, 2.0]]), &device); + let loss_func = LpLossConfig::l2(); + + // Reduce over all dims -> should give [1, 1] shape + let loss = loss_func.forward_reduce_dims(predictions, targets, &[0, 1]); + + // Errors: [[-1, 1], [0, 2]] -> squared: [[1, 1], [0, 4]] -> mean: 1.5 + let expected = TensorData::from([[1.5]]); + loss.into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn test_forward_reduce_dims_image_batch() { + // Simulate per-image loss for [batch, C, H, W] tensor (common use case for PSNR) + let device = Default::default(); + // Shape: [2, 1, 2, 2] (batch=2, C=1, H=2, W=2) + let predictions = Tensor::<4>::from_data( + TensorData::from([ + [[[1.0, 2.0], [3.0, 4.0]]], // Image 1 + [[[5.0, 6.0], [7.0, 8.0]]], // Image 2 + ]), + &device, + ); + let targets = Tensor::<4>::from_data( + TensorData::from([ + [[[0.0, 2.0], [3.0, 6.0]]], // Target 1 + [[[5.0, 5.0], [7.0, 7.0]]], // Target 2 + ]), + &device, + ); + let loss_func = LpLossConfig::l2(); + + // Reduce over C, H, W (dims 1, 2, 3) to get per-image MSE + let loss = loss_func.forward_reduce_dims(predictions, targets, &[1, 2, 3]); + + // Image 1 errors: [[1, 0], [0, -2]] -> squared: [[1, 0], [0, 4]] -> mean: 1.25 + // Image 2 errors: [[0, 1], [0, 1]] -> squared: [[0, 1], [0, 1]] -> mean: 0.5 + let expected = TensorData::from([[[[1.25]]], [[[0.5]]]]); + loss.into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn test_forward_reduce_dims_with_p1() { + let device = Default::default(); + // Shape: [2, 3] + let predictions = Tensor::<2>::from_data( + TensorData::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), + &device, + ); + let targets = Tensor::<2>::from_data( + TensorData::from([[0.0, 5.0, 3.0], [1.0, 5.0, 9.0]]), + &device, + ); + let loss_func = LpLossConfig::l1(); + + // Reduce over dim 1 -> should give [2, 1] shape + let loss = loss_func.forward_reduce_dims(predictions, targets, &[1]); + + // Abs errors row 0: [1, 3, 0] -> mean: 4/3 + // Abs errors row 1: [3, 0, 3] -> mean: 2 + let expected = TensorData::from([[4.0 / 3.0], [2.0]]); + loss.into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn test_forward_reduce_dims_empty_dims() { + // Reducing over no dimensions should return the unreduced loss + let device = Default::default(); + let predictions = + Tensor::<2>::from_data(TensorData::from([[1.0, 2.0], [3.0, 4.0]]), &device); + let targets = Tensor::<2>::from_data(TensorData::from([[0.0, 2.0], [3.0, 6.0]]), &device); + let loss_func = LpLossConfig::l2(); + let loss_reduce_dims = + loss_func.forward_reduce_dims(predictions.clone(), targets.clone(), &[]); + let loss_no_reduction = loss_func.forward_no_reduction(predictions, targets); + + // Should be equivalent + loss_reduce_dims + .into_data() + .assert_eq(&loss_no_reduction.into_data(), true); + } + + #[test] + fn test_forward_reduce_dims_zero_error() { + let device = Default::default(); + // Shape: [2, 2, 2] + let predictions = Tensor::<3>::from_data( + TensorData::from([[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]), + &device, + ); + let targets = predictions.clone(); + let loss_func = LpLossConfig::l2(); + let loss = loss_func.forward_reduce_dims(predictions, targets, &[1, 2]); + + // All zeros, reduced to shape: [2, 1, 1] + let expected = TensorData::from([[[0.0]], [[0.0]]]); + loss.into_data().assert_eq(&expected, false); + } +} diff --git a/crates/burn-nn/src/loss/mod.rs b/crates/burn-nn/src/loss/mod.rs new file mode 100644 index 0000000..69017a6 --- /dev/null +++ b/crates/burn-nn/src/loss/mod.rs @@ -0,0 +1,25 @@ +mod binary_cross_entropy; +mod cosine_embedding; +mod cross_entropy; +mod ctc; +mod huber; +mod kldiv; +mod lp_loss; +mod mse; +mod poisson; +mod reduction; +mod rnnt; +mod smooth_l1; + +pub use binary_cross_entropy::*; +pub use cosine_embedding::*; +pub use cross_entropy::*; +pub use ctc::*; +pub use huber::*; +pub use kldiv::*; +pub use lp_loss::*; +pub use mse::*; +pub use poisson::*; +pub use reduction::*; +pub use rnnt::*; +pub use smooth_l1::*; diff --git a/crates/burn-nn/src/loss/mse.rs b/crates/burn-nn/src/loss/mse.rs new file mode 100644 index 0000000..2a99e9a --- /dev/null +++ b/crates/burn-nn/src/loss/mse.rs @@ -0,0 +1,86 @@ +use burn_core as burn; + +use crate::loss::reduction::Reduction; + +use burn::module::Module; +use burn::tensor::Tensor; + +/// Calculate the mean squared error loss from the input logits and the targets. +#[derive(Module, Debug)] +pub struct MseLoss; + +impl Default for MseLoss { + fn default() -> Self { + Self::new() + } +} + +impl MseLoss { + /// Create the criterion. + pub fn new() -> Self { + Self + } + + /// Compute the criterion on the input tensor. + /// + /// # Shapes + /// + /// - logits: [batch_size, num_targets] + /// - targets: [batch_size, num_targets] + pub fn forward( + &self, + logits: Tensor, + targets: Tensor, + reduction: Reduction, + ) -> Tensor<1> { + let tensor = self.forward_no_reduction(logits, targets); + match reduction { + Reduction::Mean | Reduction::Auto => tensor.mean(), + Reduction::Sum => tensor.sum(), + other => panic!("{other:?} reduction is not supported"), + } + } + + /// Compute the criterion on the input tensor without reducing. + pub fn forward_no_reduction( + &self, + logits: Tensor, + targets: Tensor, + ) -> Tensor { + logits.sub(targets).square() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + + #[test] + fn test_mse_loss() { + let device = Default::default(); + let logits = Tensor::<2>::from_data(TensorData::from([[1.0, 2.0], [3.0, 4.0]]), &device); + + let targets = Tensor::<2>::from_data(TensorData::from([[2.0, 1.0], [3.0, 2.0]]), &device); + + let mse = MseLoss::new(); + let loss_no_reduction = mse.forward_no_reduction(logits.clone(), targets.clone()); + let loss = mse.forward(logits.clone(), targets.clone(), Reduction::Auto); + let loss_sum = mse.forward(logits, targets, Reduction::Sum); + + let expected = TensorData::from([[1.0, 1.0], [0.0, 4.0]]); + loss_no_reduction.into_data().assert_eq(&expected, false); + + let expected = TensorData::from([1.5]); + loss.into_data().assert_eq(&expected, false); + + let expected = TensorData::from([6.0]); + loss_sum.into_data().assert_eq(&expected, false); + } + + #[test] + fn display() { + let loss = MseLoss::new(); + assert_eq!(alloc::format!("{loss}"), "MseLoss"); + } +} diff --git a/crates/burn-nn/src/loss/poisson.rs b/crates/burn-nn/src/loss/poisson.rs new file mode 100644 index 0000000..e54c746 --- /dev/null +++ b/crates/burn-nn/src/loss/poisson.rs @@ -0,0 +1,405 @@ +use burn_core as burn; +use core::f32::consts::PI; + +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::{config::Config, module::Module}; + +use super::Reduction; + +/// Configuration for creating a [PoissonNllLoss](PoissonNllLoss) instance. +/// +/// This configuration allows customization of the Poisson Negative Log Likelihood (NLL) loss +/// behavior, such as whether the input is in log-space, whether to include the Stirling +/// approximation term, and a small epsilon value to avoid numerical instability. +#[derive(Config, Debug)] +pub struct PoissonNllLossConfig { + /// If `true`, the predictions are expected to be in log-space. + /// + /// When `log_input` is `true`, the loss is computed as: + /// ```text + /// L(predictions, target) = exp(predictions) - target * predictions + /// ``` + /// When `log_input` is `false`, the loss is computed as: + /// ```text + /// L(predictions, target) = predictions - target * log(predictions + eps) + /// ``` + #[config(default = true)] + pub log_input: bool, + /// Whether to compute the full loss, including the Stirling approximation term. + /// + /// When `full` is `true`, the Stirling approximation term is added to the loss: + /// ```text + /// target * log(target) - target + 0.5 * log(2 * PI * target) + /// ``` + #[config(default = false)] + pub full: bool, + /// A small value to avoid evaluation of `log(0)` when `log_input` is `false`. + /// + /// This epsilon value is added to the predictions to ensure numerical stability + /// when computing the logarithm. + #[config(default = 1e-8)] + pub eps: f64, +} + +impl PoissonNllLossConfig { + /// Initializes a [PoissonNllLoss](PoissonNllLoss) instance with the current configuration. + /// + /// # Panics + /// - Panics if `eps` is not a positive number. + pub fn init(&self) -> PoissonNllLoss { + self.assertions(); + PoissonNllLoss { + log_input: self.log_input, + full: self.full, + eps: self.eps, + } + } + + /// Validates the configuration parameters. + /// + /// # Panics + /// - Panics if `eps` is not a positive number. + fn assertions(&self) { + assert!( + self.eps > 0., + "eps for PoissonNllLoss must be a positive number." + ); + } +} + +/// Negative Log Likelihood (NLL) loss with a Poisson distribution assumption for the target. +/// +/// This loss function is used when the target values are assumed to follow a Poisson distribution. +/// The loss is defined as: +/// ```text +/// target ~ Poisson(input) +/// L(predictions, target) = predictions - target * log(predictions) + log(target!) +/// ``` +/// The last term (`log(target!)`) can be omitted or approximated using Stirling's formula. +/// The approximation is applied for `target > 1`, while for `target <= 1`, zeros are added to the loss. +/// +/// For more details, see: +/// +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct PoissonNllLoss { + /// If `true`, the predictions are expected to be in log-space. + pub log_input: bool, + /// Whether to compute the full loss, including the Stirling approximation term. + pub full: bool, + /// A small value to avoid evaluation of `log(0)` when `log_input` is `false`. + pub eps: f64, +} + +impl ModuleDisplay for PoissonNllLoss { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("log_input", &self.log_input) + .add("full", &self.full) + .add("eps", &self.eps) + .optional() + } +} + +impl PoissonNllLoss { + /// Computes the loss element-wise for the given predictions and targets, then reduces + /// the result to a single loss value. + /// + /// # Arguments + /// - `predictions`: The predicted values. + /// - `targets`: The target values. + /// - `reduction`: The reduction method to apply. `Reduction::Auto` behaves as `Reduction::Mean`. + /// + /// # Shapes + /// - `predictions`: `[...dims]` + /// - `targets`: `[...dims]` + /// - `output`: `[1]` + /// + /// # Panics + /// - Panics if the shapes of `predictions` and `targets` do not match. + /// - Panics if any target value is negative. + /// - Panics if `log_input` is `false` and any prediction value is negative. + pub fn forward( + &self, + predictions: Tensor, + targets: Tensor, + reduction: Reduction, + ) -> Tensor<1> { + let loss = self.forward_no_reduction(predictions, targets); + match reduction { + Reduction::Mean | Reduction::Auto => loss.mean(), + Reduction::Sum => loss.sum(), + other => panic!("{other:?} reduction is not supported"), + } + } + + /// Computes the loss element-wise for the given predictions and targets without reduction. + /// + /// # Arguments + /// - `predictions`: The predicted values. + /// - `targets`: The target values. + /// + /// # Shapes + /// - `predictions`: `[...dims]` + /// - `targets`: `[...dims]` + /// - `output`: `[...dims]` + /// + /// # Panics + /// - Panics if the shapes of `predictions` and `targets` do not match. + /// - Panics if any target value is negative. + /// - Panics if `log_input` is `false` and any prediction value is negative. + pub fn forward_no_reduction( + &self, + predictions: Tensor, + targets: Tensor, + ) -> Tensor { + self.assertions(&predictions, &targets); + let mut loss; + if self.log_input { + loss = predictions.clone().exp() - targets.clone() * predictions; + } else { + loss = predictions.clone() - targets.clone() * (predictions + self.eps).log(); + } + if self.full { + let log_stirling_term = targets.clone() * targets.clone().log() - targets.clone() + + (targets.clone() * 2. * PI).log() * 0.5; + loss = loss + + log_stirling_term + .mask_where(targets.clone().lower_equal_elem(1), targets.zeros_like()); + } + loss + } + + /// Validates the input tensors for the loss computation. + /// + /// # Panics + /// - Panics if the shapes of `predictions` and `targets` do not match. + /// - Panics if any target value is negative. + /// - Panics if `log_input` is `false` and any prediction value is negative. + fn assertions(&self, predictions: &Tensor, targets: &Tensor) { + let predictions_dims = predictions.dims(); + let targets_dims = targets.dims(); + assert!( + predictions_dims == targets_dims, + "Shape of targets ({targets_dims:?}) should correspond to outer shape of predictions ({predictions_dims:?})." + ); + assert!( + targets + .clone() + .greater_equal_elem(0.) + .all() + .into_scalar::(), + "All the values of `targets` must be non-negative." + ); + if !self.log_input { + assert!( + predictions + .clone() + .greater_equal_elem(0.) + .all() + .into_scalar::(), + "When `log_input` is `false`, all the values of `predictions` must be non-negative." + ); + } + } +} + +#[cfg(test)] +mod tests { + #![allow(clippy::approx_constant)] + + use super::*; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn test_poisson_nll_loss() { + let predictions = TensorData::from([0., 0., -40., 1., 2., 3.]); + let targets = TensorData::from([1., 4.5, 2.5, 0., 0., 2.]); + + let device = Default::default(); + + let predictions = Tensor::<1>::from_data(predictions, &device); + let targets = Tensor::<1>::from_data(targets, &device); + + let poisson = PoissonNllLossConfig::new().init(); + + let loss_sum = poisson.forward(predictions.clone(), targets.clone(), Reduction::Sum); + let loss = poisson.forward(predictions.clone(), targets.clone(), Reduction::Auto); + let loss_no_reduction = poisson.forward_no_reduction(predictions, targets); + + let expected = TensorData::from([1.0000, 1.0000, 100.0000, 2.7183, 7.3891, 14.0855]); + loss_no_reduction + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([21.0321]); + loss.into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([126.1929]); + loss_sum + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn test_poisson_nll_loss_no_log_input() { + let predictions = TensorData::from([0.0, 0.5, 1.0, 1.0, 2.71828, 7.38905, 20.0855]); + let targets = TensorData::from([2., 3., 1., 4.5, 0., 0., 2.]); + + let device = Default::default(); + + let predictions = Tensor::<1>::from_data(predictions, &device); + let targets = Tensor::<1>::from_data(targets, &device); + + let poisson = PoissonNllLossConfig::new().with_log_input(false).init(); + + let loss_no_reduction = poisson.forward_no_reduction(predictions.clone(), targets.clone()); + + let expected = TensorData::from([36.84136, 2.579441, 1.0, 1.0, 2.71828, 7.38905, 14.0855]); + loss_no_reduction + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn test_poisson_nll_loss_full() { + let predictions = TensorData::from([0., 0., -40., 1., 2., 3.]); + let targets = TensorData::from([1., 4.5, 2.5, 0., 0., 2.]); + + let device = Default::default(); + + let predictions = Tensor::<1>::from_data(predictions, &device); + let targets = Tensor::<1>::from_data(targets, &device); + + let poisson = PoissonNllLossConfig::new().with_full(true).init(); + + let loss_sum = poisson.forward(predictions.clone(), targets.clone(), Reduction::Sum); + let loss = poisson.forward(predictions.clone(), targets.clone(), Reduction::Auto); + let loss_no_reduction = poisson.forward_no_reduction(predictions, targets); + + let expected = TensorData::from([1.0000, 4.9393, 101.1678, 2.7183, 7.3891, 14.7373]); + loss_no_reduction + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([21.9920]); + loss.into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([131.9518]); + loss_sum + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[cfg(feature = "std")] + #[test] + fn test_poisson_nll_loss_gradients() { + use burn::tensor::Device; + let device = Device::default().autodiff(); + + let predictions = TensorData::from([0., 0., -40., 1., 2., 3.]); + let targets = TensorData::from([1., 4.5, 2.5, 0., 0., 2.]); + + let predictions1 = Tensor::<1>::from_data(predictions, &device).require_grad(); + let predictions2 = predictions1.clone(); + let targets = Tensor::<1>::from_data(targets, &device); + + let poisson = PoissonNllLossConfig::new().with_full(false).init(); + let poisson_full = PoissonNllLossConfig::new().with_full(true).init(); + + let loss_sum = poisson.forward(predictions1.clone(), targets.clone(), Reduction::Sum); + let loss_full_sum = + poisson_full.forward(predictions2.clone(), targets.clone(), Reduction::Sum); + + let grads = loss_sum.backward(); + let grads_full = loss_full_sum.backward(); + + let grads_predictions1 = predictions1.grad(&grads).unwrap(); + let grads_predictions2 = predictions2.grad(&grads_full).unwrap(); + + let expected = TensorData::from([0.0000, -3.5000, -2.5000, 2.7183, 7.3891, 18.0855]); + + grads_predictions1 + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + grads_predictions2 + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + #[should_panic = "eps for PoissonNllLoss must be a positive number."] + fn test_negative_eps() { + let _poisson = PoissonNllLossConfig::new().with_eps(0.).init(); + } + + #[test] + #[should_panic = "All the values of `targets` must be non-negative."] + fn test_targets_with_negative_values() { + let predictions = TensorData::from([0., 0., -40., 1., 2., 3., 4.]); + let targets = TensorData::from([1., 4.5, 2.5, 0., 0., 2., -0.42]); + + let device = Default::default(); + + let predictions = Tensor::<1>::from_data(predictions, &device); + let targets = Tensor::<1>::from_data(targets, &device); + + let poisson = PoissonNllLossConfig::new().init(); + + let _loss = poisson.forward(predictions.clone(), targets.clone(), Reduction::Auto); + } + + #[test] + #[should_panic = "Shape of targets"] + fn test_shape_tensors() { + let predictions = TensorData::from([0., 1., 2.]); + let targets = TensorData::from([0., 1.]); + + let device = Default::default(); + + let predictions = Tensor::<1>::from_data(predictions, &device); + let targets = Tensor::<1>::from_data(targets, &device); + + let poisson = PoissonNllLossConfig::new().init(); + + let _loss = poisson.forward_no_reduction(predictions.clone(), targets.clone()); + } + + #[test] + #[should_panic = "When `log_input` is `false`, all the values of `predictions` must be non-negative."] + fn test_exp_predictions_non_negative() { + let predictions = TensorData::from([0.3, -0.1, 0.4]); + let targets = TensorData::from([0., 1., 0.]); + + let device = Default::default(); + + let predictions = Tensor::<1>::from_data(predictions, &device); + let targets = Tensor::<1>::from_data(targets, &device); + + let poisson = PoissonNllLossConfig::new().with_log_input(false).init(); + + let _loss = poisson.forward_no_reduction(predictions.clone(), targets.clone()); + } + + #[test] + fn display() { + let config = PoissonNllLossConfig::new(); + let loss = config.init(); + + assert_eq!( + alloc::format!("{loss}"), + "PoissonNllLoss {log_input: true, full: false, eps: 0.00000001}" + ); + } +} diff --git a/crates/burn-nn/src/loss/reduction.rs b/crates/burn-nn/src/loss/reduction.rs new file mode 100644 index 0000000..ddb15fe --- /dev/null +++ b/crates/burn-nn/src/loss/reduction.rs @@ -0,0 +1,19 @@ +use burn_core as burn; + +use burn::config::Config; + +/// The reduction type for the loss. +#[derive(Config, Debug)] +pub enum Reduction { + /// The mean of the losses will be returned. + Mean, + + /// The sum of the losses will be returned. + Sum, + + /// The sum of the losses divided by the batch_size will be returned. + BatchMean, + + /// The mean of the losses will be returned. + Auto, +} diff --git a/crates/burn-nn/src/loss/rnnt.rs b/crates/burn-nn/src/loss/rnnt.rs new file mode 100644 index 0000000..3440baa --- /dev/null +++ b/crates/burn-nn/src/loss/rnnt.rs @@ -0,0 +1,766 @@ +use super::Reduction; +use alloc::vec; +use burn::config::Config; +use burn::module::Module; +use burn::tensor::{Bool, Device, Int, Tensor, s}; +use burn_core as burn; +use core::f32; + +/// Configuration for [RNNTLoss](RNNTLoss). +#[derive(Config, Debug)] +pub struct RNNTLossConfig { + /// Index of the blank label in the vocabulary. Default: `0`. + #[config(default = 0)] + pub blank: usize, + /// Treat the inputs as logits, applying a log-softmax on the last dimension internally. + /// If `false`, the input must already be log-probabilities. Default: `true`. + #[config(default = true)] + pub logits: bool, +} + +impl RNNTLossConfig { + /// Initializes a [RNNTLoss](RNNTLoss) module. + pub fn init(&self) -> RNNTLoss { + RNNTLoss { + blank: self.blank, + logits: self.logits, + } + } +} + +/// RNN Transducer (RNNT) loss, as described in +/// [Sequence Transduction with Recurrent Neural Networks](https://arxiv.org/abs/1211.3711). +/// +/// Computes the negative log-likelihood over a 2D lattice of encoder time steps (T) +/// and output labels (U), marginalizing over all valid alignments. +/// +/// # Example +/// +/// ```rust,ignore +/// let rnnt = RNNTLossConfig::new().init(); +/// +/// // logits: [B, T, U+1, V] from the joiner network +/// let loss = rnnt.forward(logits, targets, logit_lengths, target_lengths); +/// ``` +#[derive(Module, Debug)] +pub struct RNNTLoss { + blank: usize, + logits: bool, +} + +impl RNNTLoss { + /// Computes per-sample RNNT loss (no reduction). Returns shape `[B]`. + /// + /// - `logits`: `[B, T, U+1, V]` — joiner output (raw logits or log-probs) + /// - `targets`: `[B, U]` — target label indices (must not contain blank) + /// - `logit_lengths`: `[B]` — actual encoder lengths per sample + /// - `target_lengths`: `[B]` — actual target lengths per sample + pub fn forward( + &self, + logits: Tensor<4>, + targets: Tensor<2, Int>, + logit_lengths: Tensor<1, Int>, + target_lengths: Tensor<1, Int>, + ) -> Tensor<1> { + let device = logits.device(); + let [b, max_t, max_up1, v] = logits.dims(); + let max_u = max_up1 - 1; + + self.check_inputs(b, v, &targets, &logit_lengths, &target_lengths, max_u); + + let log_probs = if self.logits { + let vocab_dim = 3; // last dim of [B, T, U+1, V] + burn::tensor::activation::log_softmax(logits, vocab_dim) + } else { + logits + }; + + let (lpb, lpl) = self.extract_log_probs(log_probs, targets); + let u_mask = self.create_u_mask(&target_lengths, b, max_up1, &device); + let neg_inf = Tensor::<2>::full([b, max_up1], f32::NEG_INFINITY, &device); + + // Forward pass: compute log_alpha across the (T, U) lattice + let mut alpha = self.init_alpha(&lpl, b, max_up1, &device); + alpha = neg_inf.clone().mask_where(u_mask.clone(), alpha); + + let logit_lengths_exp = logit_lengths.clone().reshape([b, 1]).expand([b, max_up1]); + + for t in 1..max_t { + let new = self.step_alpha(&alpha, &lpb, &lpl, t); + let new = neg_inf.clone().mask_where(u_mask.clone(), new); + + // Only update alpha for samples where t < logit_lengths[b] + let valid = logit_lengths_exp.clone().greater_elem(t as i64); + alpha = alpha.mask_where(valid, new); + } + + self.gather_loss(alpha, &lpb, logit_lengths, target_lengths, b) + } + + /// Computes RNNT loss with the given reduction. Returns shape `[1]`. + pub fn forward_with_reduction( + &self, + logits: Tensor<4>, + targets: Tensor<2, Int>, + logit_lengths: Tensor<1, Int>, + target_lengths: Tensor<1, Int>, + reduction: Reduction, + ) -> Tensor<1> { + let loss = self.forward(logits, targets, logit_lengths, target_lengths); + match reduction { + Reduction::Auto | Reduction::Mean => loss.mean(), + Reduction::Sum => loss.sum(), + other => panic!("{other:?} reduction is not supported"), + } + } + + /// Gathers `log_prob_blank[B, T, U+1]` and `log_prob_label[B, T, U]` from the full + /// log-probability tensor by indexing into the vocab dimension. + fn extract_log_probs( + &self, + log_probs: Tensor<4>, + targets: Tensor<2, Int>, + ) -> (Tensor<3>, Tensor<3>) { + let [b, max_t, max_up1, v] = log_probs.dims(); + let max_u = max_up1 - 1; + let vocab_dim = 3; + + // Blank probabilities: slice log_probs in vocab dim using the blank index + let lpb = log_probs + .clone() + .slice_dim(vocab_dim, self.blank) + .squeeze_dim::<3>(vocab_dim); + + // Label probabilities: gather target labels across vocab dim (only first U positions) + let tgt = targets + .reshape([b, 1, max_u, 1]) + .expand([b, max_t, max_u, 1]); + let lpl = log_probs + .slice(s![.., .., 0..max_u, 0..v]) + .gather(vocab_dim, tgt) + .squeeze_dim::<3>(vocab_dim); + + (lpb, lpl) + } + + /// Sets up log_alpha at t=0: `alpha(0,0) = 0`, then cumsum of label probs along u. + fn init_alpha(&self, lpl: &Tensor<3>, b: usize, max_up1: usize, device: &Device) -> Tensor<2> { + // Label probs at t=0 + let lpl_0 = lpl.clone().slice(s![.., 0..1, ..]).squeeze_dim::<2>(1); + let zero_col = Tensor::<2>::zeros([b, 1], device); + let prefix = Tensor::cat(vec![zero_col, lpl_0.slice(s![.., 0..(max_up1 - 1)])], 1); + + prefix.cumsum(1) + } + + /// Boolean mask `[B, U+1]` that is true where `u <= target_lengths[b]`. + fn create_u_mask( + &self, + target_lengths: &Tensor<1, Int>, + b: usize, + max_up1: usize, + device: &Device, + ) -> Tensor<2, Bool> { + let indices = Tensor::<1, Int>::arange(0..max_up1 as i64, device) + .reshape([1, max_up1]) + .expand([b, max_up1]); + let lengths = target_lengths.clone().reshape([b, 1]).expand([b, max_up1]); + indices.lower_equal(lengths) + } + + /// One time step of the forward recurrence: + /// + /// alpha(t, u) = logaddexp( + /// alpha(t-1, u) + blank(t-1, u), + /// alpha(t, u-1) + label(t, u-1), + /// ) + fn step_alpha( + &self, + alpha: &Tensor<2>, + lpb: &Tensor<3>, + lpl: &Tensor<3>, + t: usize, + ) -> Tensor<2> { + let [b, max_up1] = alpha.dims(); + let device = alpha.device(); + + // Blank transition: alpha(t-1, :) + blank_prob(t-1, :) + let blank_prob = lpb + .clone() + .slice(s![.., (t - 1)..t, ..]) + .squeeze_dim::<2>(1); + let from_blank = alpha.clone().add(blank_prob); + + let mut new = Tensor::<2>::full([b, max_up1], f32::NEG_INFINITY, &device); + new = new.slice_assign(s![.., 0..1], from_blank.clone().slice(s![.., 0..1])); + + // Label probs at time t + let label_prob = lpl + .clone() + .slice(s![.., t..(t + 1), ..]) + .squeeze_dim::<2>(1); + + for u in 1..max_up1 { + let via_blank = from_blank.clone().slice(s![.., u..(u + 1)]); + let via_label = new + .clone() + .slice(s![.., (u - 1)..u]) + .add(label_prob.clone().slice(s![.., (u - 1)..u])); + new = new.slice_assign(s![.., u..(u + 1)], self.log_sum_exp(via_blank, via_label)); + } + new + } + + /// Extracts `-(alpha(T_b, U_b) + blank(T_b, U_b))` for each sample in the batch. + fn gather_loss( + &self, + alpha: Tensor<2>, + lpb: &Tensor<3>, + logit_lengths: Tensor<1, Int>, + target_lengths: Tensor<1, Int>, + b: usize, + ) -> Tensor<1> { + let device = alpha.device(); + // Anchor the index dtype on `u_idx` so all three coordinate tensors share a + // common int dtype before stacking - `cat` panics on dtype mismatch and the + // caller's lengths may not use the device's default IntElem. + let u_idx = target_lengths; + let int_dtype = u_idx.dtype(); + let t_idx = logit_lengths.sub_scalar(1).cast(int_dtype); + let b_idx = Tensor::<1, Int>::arange(0..b as i64, (&device, int_dtype)); + + let alpha_tu: Tensor<1> = + alpha.gather_nd(Tensor::stack::<2>(vec![b_idx.clone(), u_idx.clone()], 1)); + let lpb_tu: Tensor<1> = lpb + .clone() + .gather_nd(Tensor::stack::<2>(vec![b_idx, t_idx, u_idx], 1)); + + alpha_tu.add(lpb_tu).neg() + } + + fn check_inputs( + &self, + b: usize, + v: usize, + targets: &Tensor<2, Int>, + logit_lengths: &Tensor<1, Int>, + target_lengths: &Tensor<1, Int>, + max_u: usize, + ) { + assert!( + self.blank < v, + "blank index {} must be less than vocab_size {}", + self.blank, + v + ); + assert_eq!( + targets.dims()[0], + b, + "targets batch dimension {} must equal batch_size {}", + targets.dims()[0], + b + ); + assert_eq!( + targets.dims()[1], + max_u, + "targets length dimension {} must equal max_target_len (max_u) {}", + targets.dims()[1], + max_u + ); + assert_eq!( + logit_lengths.dims()[0], + b, + "logit_lengths length {} must equal batch_size {}", + logit_lengths.dims()[0], + b + ); + assert_eq!( + target_lengths.dims()[0], + b, + "target_lengths length {} must equal batch_size {}", + target_lengths.dims()[0], + b + ); + } + + /// Numerically stable `log(exp(a) + exp(b))`, handling `-inf` inputs. + fn log_sum_exp(&self, a: Tensor, b: Tensor) -> Tensor { + let a_inf = a.clone().equal_elem(f32::NEG_INFINITY); + let b_inf = b.clone().equal_elem(f32::NEG_INFINITY); + + // Replace -inf with 0 to prevent NaN in the subtraction (masked out below) + let a_safe = a.clone().mask_fill(a_inf.clone(), 0.0); + let b_safe = b.clone().mask_fill(b_inf.clone(), 0.0); + + // log(exp(a) + exp(b)) = max(a,b) + log(1 + exp(-|a-b|)) + let max = a_safe.clone().max_pair(b_safe.clone()); + let result = max.add(a_safe.sub(b_safe).abs().neg().exp().add_scalar(1.0).log()); + + // If a=-inf, result is b; if b=-inf, result is a; if both -inf, stays -inf + let result = result.mask_where(a_inf, b); + result.mask_where(b_inf, a) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::{TensorData, Tolerance}; + const NUM_LABELS: usize = 2; // vocab size for simple unit tests + + #[test] + fn config_defaults() { + let cfg = RNNTLossConfig::new(); + assert_eq!(cfg.blank, 0); + assert!(cfg.logits); + } + + #[test] + #[should_panic(expected = "blank index")] + fn panics_on_invalid_blank() { + let dev = Default::default(); + let rnnt = RNNTLossConfig::new().with_blank(5).init(); + rnnt.forward( + Tensor::<4>::zeros([1, 2, 2, 3], &dev), + Tensor::<2, Int>::from_data([[1_i32]], &dev), + Tensor::<1, Int>::from_data([2], &dev), + Tensor::<1, Int>::from_data([1], &dev), + ); + } + + #[test] + #[should_panic(expected = "must equal batch_size")] + fn panics_on_batch_mismatch() { + let dev = Default::default(); + let rnnt = RNNTLossConfig::new().init(); + rnnt.forward( + Tensor::<4>::zeros([2, 3, 2, 3], &dev), + Tensor::<2, Int>::from_data([[1_i32]], &dev), + Tensor::<1, Int>::from_data([3, 3], &dev), + Tensor::<1, Int>::from_data([1, 1], &dev), + ); + } + + #[test] + #[should_panic(expected = "logit_lengths length")] + fn panics_on_logit_lengths_mismatch() { + let dev = Default::default(); + let rnnt = RNNTLossConfig::new().init(); + rnnt.forward( + Tensor::<4>::zeros([2, 3, 2, 3], &dev), + Tensor::<2, Int>::from_data([[1_i32], [2]], &dev), + Tensor::<1, Int>::from_data([3], &dev), + Tensor::<1, Int>::from_data([1, 1], &dev), + ); + } + + #[test] + #[should_panic(expected = "target_lengths length")] + fn panics_on_target_lengths_mismatch() { + let dev = Default::default(); + let rnnt = RNNTLossConfig::new().init(); + rnnt.forward( + Tensor::<4>::zeros([2, 3, 2, 3], &dev), + Tensor::<2, Int>::from_data([[1_i32], [2]], &dev), + Tensor::<1, Int>::from_data([3, 3], &dev), + Tensor::<1, Int>::from_data([1], &dev), + ); + } + + #[test] + fn single_token_uniform_probs() { + // B=1, T=2, U=1, V=2, uniform probs: P(blank) = P(label) = 1/V + // + // Two alignment paths (label emitted at t=0 or t=1), each with T+U emissions: + // total_prob = T * (1/V)^(T+1) = 2 * (1/2)^3 = 1/4 + // loss = -ln(1/4) = 2*ln(2) + let dev = Default::default(); + let rnnt = RNNTLossConfig::new().with_logits(false).init(); + let time_steps = 2; + let target_len = 1; + let v = NUM_LABELS as f32; + let log_uniform = (1.0 / v).ln(); + + let loss = rnnt.forward( + Tensor::<4>::full( + [1, time_steps, target_len + 1, NUM_LABELS], + log_uniform, + &dev, + ), + Tensor::<2, Int>::from_data([[1_i32]], &dev), + Tensor::<1, Int>::from_data([time_steps as i64], &dev), + Tensor::<1, Int>::from_data([target_len as i64], &dev), + ); + // Each path: T-1 blanks + U labels + 1 final blank = T + U emissions + let num_paths = time_steps as f32; + let emissions_per_path = (time_steps + target_len) as f32; + let total_prob = num_paths * v.powf(-emissions_per_path); + let expected_loss = -total_prob.ln(); + loss.into_data().assert_approx_eq::( + &TensorData::from([expected_loss]), + Tolerance::absolute(1e-4), + ); + } + + #[test] + fn empty_target() { + // B=1, T=3, U=0, V=2, uniform probs: only the all-blanks path exists. + // + // Single path with T emissions (T-1 blanks + 1 final blank, all at u=0): + // total_prob = (1/V)^T = (1/2)^3 = 1/8 + // loss = T*ln(V) = 3*ln(2) + let dev = Default::default(); + let rnnt = RNNTLossConfig::new().with_logits(false).init(); + let time_steps = 3; + let target_len = 0; + let v = NUM_LABELS as f32; + let log_uniform = (1.0 / v).ln(); + + let loss = rnnt.forward( + Tensor::<4>::full([1, time_steps, 2, NUM_LABELS], log_uniform, &dev), + Tensor::<2, Int>::from_data([[1_i32]], &dev), + Tensor::<1, Int>::from_data([time_steps as i64], &dev), + Tensor::<1, Int>::from_data([target_len as i64], &dev), + ); + // T + U = T emissions total for U=0 + let expected_loss = -v.powf(-((time_steps + target_len) as f32)).ln(); + loss.into_data().assert_approx_eq::( + &TensorData::from([expected_loss]), + Tolerance::absolute(1e-4), + ); + } + + #[test] + fn logits_equivalence() { + // Verify that logits=true (internal log_softmax on raw logits) + // gives the same loss as logits=false with external log_softmax. + let dev = Default::default(); + let [bs, time_steps, up1, vocab] = [1, 2, 3, 4]; + let num_elements = bs * time_steps * up1 * vocab; + let target_len = up1 - 1; + + let data: Vec = (0..num_elements).map(|i| (i as f32 * 0.3).sin()).collect(); + let logits = Tensor::<4>::from_data( + burn_core::tensor::TensorData::new(data, [bs, time_steps, up1, vocab]), + &dev, + ); + let targets = Tensor::<2, Int>::from_data([[1_i32, 2]], &dev); + let logit_lengths = Tensor::<1, Int>::from_data([time_steps as i64], &dev); + let target_lengths = Tensor::<1, Int>::from_data([target_len as i64], &dev); + + let vocab_dim = 3; + let fused = RNNTLossConfig::new().with_logits(true).init().forward( + logits.clone(), + targets.clone(), + logit_lengths.clone(), + target_lengths.clone(), + ); + + let log_probs = burn::tensor::activation::log_softmax(logits, vocab_dim); + let manual = RNNTLossConfig::new().with_logits(false).init().forward( + log_probs, + targets, + logit_lengths, + target_lengths, + ); + + fused + .into_data() + .assert_approx_eq::(&manual.into_data(), Tolerance::absolute(1e-4)); + } +} + +/// Tests comparing forward loss and backward gradients against torchaudio.functional.rnnt_loss. +/// +/// Logits are generated deterministically via sin((b*11+t*7+u*13+v*3)*0.1) so the same +/// values can be reproduced in a Python script for cross-checking. +#[cfg(test)] +#[allow(clippy::identity_op, clippy::too_many_arguments)] +mod pytorch_comparison_tests { + use super::*; + use burn::tensor::{TensorData, Tolerance}; + fn tol() -> Tolerance { + Tolerance::absolute(1e-3) + } + + /// Deterministic logits matching the Python reference generator. + /// Uses coprime coefficients to avoid repeating patterns across dimensions. + fn make_logits(bs: usize, t: usize, u: usize, v: usize, dev: &Device) -> Tensor<4> { + let mut data = Vec::with_capacity(bs * t * u * v); + for bi in 0..bs { + for ti in 0..t { + for ui in 0..u { + for vi in 0..v { + let idx = bi * 11 + ti * 7 + ui * 13 + vi * 3; + data.push((idx as f32 * 0.1).sin()); + } + } + } + } + Tensor::from_data(TensorData::new(data, [bs, t, u, v]), dev) + } + + /// Checks that gradients along the vocab dim sum to ~0 at every (b, t, u) position. + /// This must hold because log_softmax is applied on the last dim, + /// and the Jacobian of log_softmax has the property that each row sums to zero. + fn check_vocab_grad_sums(grad: &[f32], bs: usize, t: usize, up1: usize, v: usize) { + for bi in 0..bs { + for ti in 0..t { + for ui in 0..up1 { + let base = ((bi * t + ti) * up1 + ui) * v; + let sum: f32 = (0..v).map(|vi| grad[base + vi]).sum(); + TensorData::from([sum]) + .assert_approx_eq::(&TensorData::from([0.0f32]), tol()); + } + } + } + } + + /// Returns the V-sized gradient slice at position (b, t, u) in a flattened [B, T, U+1, V] grad. + fn grad_at( + grad: &[f32], + b: usize, + t: usize, + u: usize, + max_t: usize, + up1: usize, + v: usize, + ) -> &[f32] { + let base = ((b * max_t + t) * up1 + u) * v; + &grad[base..base + v] + } + + /// Asserts that a gradient slice at position (b, t, u) matches expected values. + fn assert_grad( + grad: &[f32], + b: usize, + t: usize, + u: usize, + max_t: usize, + up1: usize, + v: usize, + expected: &[f32], + ) { + TensorData::from(grad_at(grad, b, t, u, max_t, up1, v)) + .assert_approx_eq::(&TensorData::from(expected), tol()); + } + + #[test] + fn basic_b1() { + // B=1, T=4, U+1=3, V=3, targets=[1,2] + let dev = Device::default().autodiff(); + let rnnt = RNNTLossConfig::new().init(); + let logits = make_logits(1, 4, 3, 3, &dev).require_grad(); + + let loss = rnnt.forward( + logits.clone(), + Tensor::<2, Int>::from_data([[1_i32, 2]], &dev), + Tensor::<1, Int>::from_data([4_i32], &dev), + Tensor::<1, Int>::from_data([2_i32], &dev), + ); + loss.clone() + .into_data() + .assert_approx_eq::(&TensorData::from([4.4491f32]), tol()); + + let grads = loss.sum().backward(); + let grad = logits + .grad(&grads) + .unwrap() + .into_data() + .to_vec::() + .unwrap(); + + // Spot-check first, middle, and last (t, u) positions against torchaudio + assert_grad(&grad, 0, 0, 0, 4, 3, 3, &[-0.2041, -0.2246, 0.4287]); + assert_grad(&grad, 0, 2, 0, 4, 3, 3, &[0.0079, -0.0640, 0.0561]); + assert_grad(&grad, 0, 3, 2, 4, 3, 3, &[-0.6899, 0.3231, 0.3667]); + check_vocab_grad_sums(&grad, 1, 4, 3, 3); + } + + #[test] + fn batched_b2() { + // B=2, T=5, U+1=4, V=4, targets=[[1,2,3],[2,1,3]] + let dev = Device::default().autodiff(); + let rnnt = RNNTLossConfig::new().init(); + let logits = make_logits(2, 5, 4, 4, &dev).require_grad(); + + let loss = rnnt.forward( + logits.clone(), + Tensor::<2, Int>::from_data(TensorData::new(vec![1_i32, 2, 3, 2, 1, 3], [2, 3]), &dev), + Tensor::<1, Int>::from_data([5_i32, 5], &dev), + Tensor::<1, Int>::from_data([3_i32, 3], &dev), + ); + loss.clone() + .into_data() + .assert_approx_eq::(&TensorData::from([7.9356f32, 7.2033]), tol()); + + let grads = loss.sum().backward(); + let grad = logits + .grad(&grads) + .unwrap() + .into_data() + .to_vec::() + .unwrap(); + + // Spot-check: first position of each sample, and last position + assert_grad(&grad, 0, 0, 0, 5, 4, 4, &[-0.3161, -0.3113, 0.2796, 0.3479]); + assert_grad(&grad, 1, 0, 0, 5, 4, 4, &[-0.2766, 0.2602, -0.2248, 0.2411]); + assert_grad(&grad, 0, 4, 3, 5, 4, 4, &[-0.8216, 0.2296, 0.2786, 0.3133]); + assert_grad(&grad, 1, 4, 3, 5, 4, 4, &[-0.7185, 0.2735, 0.2437, 0.2012]); + check_vocab_grad_sums(&grad, 2, 5, 4, 4); + } + + #[test] + fn variable_lengths_b3() { + // B=3, T=6, U+1=4, V=5 + // logit_lengths=[6,4,5], target_lengths=[3,2,1] + // Tests that masking works correctly for variable-length sequences. + let dev = Device::default().autodiff(); + let rnnt = RNNTLossConfig::new().init(); + let logits = make_logits(3, 6, 4, 5, &dev).require_grad(); + + let loss = rnnt.forward( + logits.clone(), + Tensor::<2, Int>::from_data( + TensorData::new(vec![1_i32, 2, 3, 4, 1, 0, 2, 0, 0], [3, 3]), + &dev, + ), + Tensor::<1, Int>::from_data([6_i32, 4, 5], &dev), + Tensor::<1, Int>::from_data([3_i32, 2, 1], &dev), + ); + loss.clone() + .into_data() + .assert_approx_eq::(&TensorData::from([10.7458f32, 8.0196, 8.3316]), tol()); + + let grads = loss.sum().backward(); + let grad = logits + .grad(&grads) + .unwrap() + .into_data() + .to_vec::() + .unwrap(); + let stride = 4 * 5; // U+1 * V per time step + let zeros = vec![0.0f32; 5]; + + // Sample 0 (full length=6): spot-check first and last active positions + assert_grad( + &grad, + 0, + 0, + 0, + 6, + 4, + 5, + &[-0.4232, -0.3114, 0.1992, 0.2478, 0.2876], + ); + assert_grad( + &grad, + 0, + 5, + 3, + 6, + 4, + 5, + &[-0.8016, 0.2170, 0.2172, 0.1991, 0.1683], + ); + + // Sample 1 (logit_length=4): gradients beyond t=3 should be zero + assert_grad( + &grad, + 1, + 0, + 0, + 6, + 4, + 5, + &[-0.2502, 0.2160, 0.2173, 0.2002, -0.3833], + ); + let sample1_t4_start = 1 * 6 * stride + 4 * stride; + for i in 0..(2 * stride) { + // t=4 and t=5 should all be zero + assert!( + grad[sample1_t4_start + i].abs() < 1e-3, + "sample 1, t>=4: grad[{}] = {} (expected 0)", + i, + grad[sample1_t4_start + i] + ); + } + + // Sample 1 (target_length=2): u=3 positions should be zero within active time steps + for ti in 0..4 { + assert_grad(&grad, 1, ti, 3, 6, 4, 5, &zeros); + } + + // Sample 2 (logit_length=5): t=5 should be zero + let sample2_t5_start = 2 * 6 * stride + 5 * stride; + for i in 0..stride { + assert!( + grad[sample2_t5_start + i].abs() < 1e-3, + "sample 2, t=5: grad[{}] = {} (expected 0)", + i, + grad[sample2_t5_start + i] + ); + } + + check_vocab_grad_sums(&grad, 3, 6, 4, 5); + } + + #[test] + fn sum_reduction() { + let dev = Device::default().autodiff(); + let rnnt = RNNTLossConfig::new().init(); + let logits = make_logits(2, 5, 4, 4, &dev).require_grad(); + let tgt = + Tensor::<2, Int>::from_data(TensorData::new(vec![1_i32, 2, 3, 2, 1, 3], [2, 3]), &dev); + let il = Tensor::<1, Int>::from_data([5_i32, 5], &dev); + let tl = Tensor::<1, Int>::from_data([3_i32, 3], &dev); + + let loss = rnnt.forward_with_reduction(logits.clone(), tgt, il, tl, Reduction::Sum); + // 7.9356 + 7.2033 = 15.1389 + loss.clone() + .into_data() + .assert_approx_eq::(&TensorData::from([15.1389f32]), tol()); + + let grads = loss.backward(); + let g = logits + .grad(&grads) + .unwrap() + .into_data() + .to_vec::() + .unwrap(); + TensorData::from(&g[..4]).assert_approx_eq::( + &TensorData::from([-0.3161f32, -0.3113, 0.2796, 0.3479]), + tol(), + ); + } + + #[test] + fn mean_reduction() { + let dev = Device::default().autodiff(); + let rnnt = RNNTLossConfig::new().init(); + let logits = make_logits(2, 5, 4, 4, &dev).require_grad(); + let tgt = + Tensor::<2, Int>::from_data(TensorData::new(vec![1_i32, 2, 3, 2, 1, 3], [2, 3]), &dev); + let il = Tensor::<1, Int>::from_data([5_i32, 5], &dev); + let tl = Tensor::<1, Int>::from_data([3_i32, 3], &dev); + + let loss = rnnt.forward_with_reduction(logits.clone(), tgt, il, tl, Reduction::Mean); + // 15.1389 / 2 = 7.5694 + loss.clone() + .into_data() + .assert_approx_eq::(&TensorData::from([7.5694f32]), tol()); + + // Gradients should be half the sum-reduction gradients (mean over batch of 2) + let grads = loss.backward(); + let g = logits + .grad(&grads) + .unwrap() + .into_data() + .to_vec::() + .unwrap(); + TensorData::from(&g[..4]).assert_approx_eq::( + &TensorData::from([-0.1581f32, -0.1557, 0.1398, 0.1739]), + tol(), + ); + } +} diff --git a/crates/burn-nn/src/loss/smooth_l1.rs b/crates/burn-nn/src/loss/smooth_l1.rs new file mode 100644 index 0000000..529e203 --- /dev/null +++ b/crates/burn-nn/src/loss/smooth_l1.rs @@ -0,0 +1,501 @@ +use super::Reduction; +use burn::config::Config; +use burn::module::Module; +use burn::tensor::Tensor; +use burn_core as burn; + +/// Configuration for the [SmoothL1Loss](SmoothL1Loss) module. +/// +/// Smooth L1 loss combines L1 and L2 loss, using L2 loss for small errors (below beta) +/// and L1 loss for large errors (above beta). This makes it less sensitive to outliers +/// than MSE while maintaining smooth gradients near zero. +/// +/// # Example +/// +/// ```ignore +/// use burn_nn::loss::{SmoothL1LossConfig, Reduction}; +/// +/// // Create Smooth L1 loss with default beta=1.0 +/// let smooth_l1 = SmoothL1LossConfig::new().init(); +/// +/// // Create with custom beta +/// let smooth_l1_custom = SmoothL1LossConfig::new().with_beta(0.5).init(); +/// ``` +#[derive(Config, Debug)] +pub struct SmoothL1LossConfig { + /// Specifies the threshold at which to change between L1 and L2 loss. + /// The value must be positive. Default: 1.0 + #[config(default = 1.0)] + pub beta: f32, +} + +impl SmoothL1LossConfig { + /// Initializes a [Smooth L1 Loss](SmoothL1Loss) module. + /// + /// # Panics + /// + /// Panics if `beta <= 0`. + pub fn init(&self) -> SmoothL1Loss { + self.assertions(); + SmoothL1Loss { beta: self.beta } + } + + fn assertions(&self) { + assert!(self.beta > 0.0, "The parameter beta must be positive.") + } +} + +/// Computes the Smooth L1 Loss between predictions and targets. +/// +/// This loss function uses L2 loss for small errors (below beta) and L1 loss for +/// large errors (above beta), providing robustness to outliers while maintaining +/// smooth gradients near |x - y| = 0. +/// +/// # Mathematical Definition +/// +/// For predictions `x` and targets `y`, the element-wise loss is: +/// +/// - L_i = 0.5 * (x_i - y_i)² / beta , if |x_i - y_i| < beta +/// - L_i = |x_i - y_i| - 0.5 * beta , otherwise +/// +/// # Notes +/// +/// Smooth L1 loss is closely related to HuberLoss since it is equivalent to HuberLoss +/// scaled by `1/beta`: +/// `SmoothL1(x, y, beta) = Huber(x, y, beta) / beta` +/// +/// This leads to the following differences: +/// +/// - As beta approaches 0, Smooth L1 loss converges to L1Loss, while HuberLoss converges to 0. +/// When beta = 0, Smooth L1 loss is equivalent to L1 loss. Thus, the `beta` +/// parameter in Burn must be positive. L1Loss should be used for beta = 0. +/// - As beta approaches positive infinity, Smooth L1 loss converges to a constant 0 loss, while +/// HuberLoss converges to L2Loss. +/// +/// # Example +/// +/// ```rust,ignore +/// use burn_nn::loss::{SmoothL1LossConfig, Reduction}; +/// use burn::tensor::Tensor; +/// +/// // Create Smooth L1 loss with the default beta=1.0 +/// let smooth_l1 = SmoothL1LossConfig::new().init(); +/// +/// let predictions: Tensor = /* model output */; +/// let targets: Tensor = /* ground truth */; +/// +/// // Compute element-wise loss without reduction +/// let element_wise = smooth_l1.forward(predictions.clone(), targets.clone()); +/// +/// // Compute loss with mean reduction +/// let loss = smooth_l1.forward_with_reduction(predictions.clone(), targets.clone(), Reduction::Mean); +/// +/// // Per-image loss: reduce over C, H, W → [batch, 1, 1, 1] +/// let loss_per_image = smooth_l1.forward_reduce_dims(predictions, targets, &[1, 2, 3]); +/// ``` +#[derive(Module, Debug)] +pub struct SmoothL1Loss { + /// Specifies the threshold at which to change between L1 and L2 loss. + /// The value must be positive. Default: 1.0 + pub beta: f32, +} + +impl SmoothL1Loss { + /// Computes the element-wise smooth L1 loss without reduction. + /// + /// # Arguments + /// + /// - `predictions` - The model's predicted values. + /// - `targets` - The ground truth target values. + /// + /// # Returns + /// + /// A tensor of the same shape as the inputs, containing the smooth L1 loss + /// for each element. + /// + /// # Shapes + /// + /// - predictions: `[...dims]` - Any shape + /// - targets: `[...dims]` - Must match predictions shape + /// - output: `[...dims]` - Same shape as inputs + pub fn forward(&self, predictions: Tensor, targets: Tensor) -> Tensor { + let error = predictions.sub(targets); + let abs_error = error.clone().abs(); + + // The L1 case: |error| - 0.5 * beta (when |error| >= beta) + let l1_loss = abs_error.clone().sub_scalar(0.5 * self.beta); + + // The L2 case: 0.5 * (error)^2 / beta (when |error| < beta) + let l2_loss = error.square().mul_scalar(0.5).div_scalar(self.beta); + + let l2_mask = abs_error.lower_elem(self.beta); + l1_loss.mask_where(l2_mask, l2_loss) + } + + /// Computes the smooth L1 loss with reduction. + /// + /// # Arguments + /// + /// - `predictions` - The model's predicted values. + /// - `targets` - The ground truth target values. + /// - `reduction` - Specifies how to reduce the element-wise losses: + /// - `Reduction::Mean` or `Reduction::Auto`: Returns the mean of all element-wise losses. + /// - `Reduction::Sum`: Returns the sum of all element-wise losses. + /// + /// # Returns + /// + /// A scalar tensor containing the reduced loss value. + /// + /// # Shapes + /// + /// - predictions: `[...dims]` - Any shape + /// - targets: `[...dims]` - Must match predictions shape + /// - output: `[1]` - Scalar loss value + pub fn forward_with_reduction( + &self, + predictions: Tensor, + targets: Tensor, + reduction: Reduction, + ) -> Tensor<1> { + let unreduced_loss = self.forward(predictions, targets); + + match reduction { + Reduction::Mean | Reduction::Auto => unreduced_loss.mean(), + Reduction::Sum => unreduced_loss.sum(), + other => panic!("{other:?} reduction is not supported"), + } + } + + /// Computes the smooth L1 loss with reduction over specified dimensions. + /// + /// Calculates element-wise smooth L1 loss, then takes the mean + /// over the specified dimensions. Useful for per-sample or per-channel losses. + /// + /// Dimensions can be provided in any order. They are sorted internally and + /// reduced from highest to lowest to ensure indices remain valid. + /// + /// # Arguments + /// + /// - `predictions` - The model's predicted values. + /// - `targets` - The ground truth target values. + /// - `dims` - Dimensions to reduce over. + /// + /// # Returns + /// + /// A tensor with the specified dimensions reduced to size 1. + /// + /// # Example + /// + /// ```ignore + /// // Consider image tensor with shape [batch, C, H, W] + /// let smooth_l1 = SmoothL1LossConfig::new().init(); + /// + /// // Per-image loss: reduce over C, H, W → [batch, 1, 1, 1] + /// let loss_per_image = smooth_l1.forward_reduce_dims(predictions, targets, &[1, 2, 3]); + /// ``` + pub fn forward_reduce_dims( + &self, + predictions: Tensor, + targets: Tensor, + dims: &[usize], + ) -> Tensor { + let error = self.forward(predictions, targets); + + // Sort the dimensions to ascending order + let mut sorted_dims = dims.to_vec(); + sorted_dims.sort(); + + // Reduce over specified dimensions + error.mean_dims(sorted_dims.as_slice()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + + type FT = f32; + + // ========================================================================= + // Configuration Tests + // ========================================================================= + + #[test] + fn test_smooth_l1_config_default_beta() { + let loss = SmoothL1LossConfig::new().init(); + assert_eq!(loss.beta, 1.0); + } + + #[test] + fn test_smooth_l1_config_custom_beta() { + let loss = SmoothL1LossConfig::new().with_beta(2.5).init(); + assert_eq!(loss.beta, 2.5); + } + + #[test] + #[should_panic(expected = "The parameter beta must be positive")] + fn test_smooth_l1_config_beta_zero_panics() { + SmoothL1LossConfig::new().with_beta(0.0).init(); + } + + #[test] + #[should_panic(expected = "The parameter beta must be positive")] + fn test_smooth_l1_config_beta_negative_panics() { + SmoothL1LossConfig::new().with_beta(-1.0).init(); + } + + // ========================================================================= + // Forward Pass (Element-wise) Tests + // ========================================================================= + + #[test] + fn test_smooth_l1_forward_l2_region() { + // Beta = 1.0, errors = 0.0 and 0.5 (both < beta, use L2 formula) + // L2 formula: 0.5 * error^2 / beta + // error = 0.0 -> loss = 0.5 * 0.0 / 1.0 = 0.0 + // error = 0.5 -> loss = 0.5 * 0.25 / 1.0 = 0.125 + let device = Default::default(); + let loss = SmoothL1LossConfig::new().init(); + + let predictions = Tensor::<2>::from_data(TensorData::from([[0.0_f32, 0.5]]), &device); + let targets = Tensor::<2>::from_data(TensorData::from([[0.0_f32, 0.0]]), &device); + + let output = loss.forward(predictions, targets); + let expected = TensorData::from([[0.0_f32, 0.125]]); + output.into_data().assert_eq(&expected, false); + } + + #[test] + fn test_smooth_l1_forward_l1_region() { + // Beta = 1.0, errors = 0.0 and 2.0 (2.0 >= beta, use L1 formula) + // L1 formula: |error| - 0.5 * beta + // L2 formula: 0.5 * (error)^2 / beta + // error = 0.0 -> loss = 0.0 + // error = 2.0 -> loss = 2.0 - 0.5 = 1.5 + let device = Default::default(); + let loss = SmoothL1LossConfig::new().init(); + + let predictions = Tensor::<2>::from_data(TensorData::from([[0.0_f32, 2.0]]), &device); + let targets = Tensor::<2>::from_data(TensorData::from([[0.0_f32, 0.0]]), &device); + + let output = loss.forward(predictions, targets); + let expected = TensorData::from([[0.0_f32, 1.5]]); + output.into_data().assert_eq(&expected, false); + } + + #[test] + fn test_smooth_l1_forward_zero_error() { + let device = Default::default(); + let loss = SmoothL1LossConfig::new().init(); + + let predictions = Tensor::<2>::from_data(TensorData::from([[1.0_f32, 2.0, 3.0]]), &device); + let targets = predictions.clone(); + + let output = loss.forward(predictions, targets); + let expected = TensorData::from([[0.0_f32, 0.0, 0.0]]); + output.into_data().assert_eq(&expected, false); + } + + #[test] + fn test_smooth_l1_forward_negative_errors() { + // Ensure absolute value is used correctly + // L1 formula: |error| - 0.5 * beta + // L2 formula: 0.5 * (error)^2 / beta + // Beta = 1.0, error = -3.0 (L1: 3.0 - 0.5 = 2.5) + let device = Default::default(); + let loss = SmoothL1LossConfig::new().init(); + + let predictions = Tensor::<1>::from_data(TensorData::from([-3.0_f32]), &device); + let targets = Tensor::<1>::zeros([1], &device); + + let output = loss.forward(predictions, targets); + let expected = TensorData::from([2.5_f32]); + output.into_data().assert_eq(&expected, false); + } + + #[test] + fn test_smooth_l1_forward_mixed_regions() { + // Test with errors in both L1 and L2 regions + // Beta = 1.0 + // L1 formula: |error| - 0.5 * beta + // L2 formula: 0.5 * (error)^2 / beta + // error = 0.5 -> L2: 0.5 * 0.25 / 1 = 0.125 + // error = 1.5 -> L1: 1.5 - 0.5 = 1.0 + // error = 3.0 -> L1: 3.0 - 0.5 = 2.5 + let device = Default::default(); + let loss = SmoothL1LossConfig::new().init(); + + let predictions = Tensor::<1>::from_data(TensorData::from([0.5_f32, 1.5, 3.0]), &device); + let targets = Tensor::<1>::zeros([3], &device); + + let output = loss.forward(predictions, targets); + let expected = TensorData::from([0.125_f32, 1.0, 2.5]); + output.into_data().assert_eq(&expected, false); + } + + #[test] + fn test_smooth_l1_custom_beta_values() { + // Test with beta = 0.5 + // error = 0.25 (< beta): L2 = 0.5 * 0.0625 / 0.5 = 0.0625 + // error = 1.0 (>= beta): L1 = 1.0 - 0.25 = 0.75 + let device = Default::default(); + let loss = SmoothL1LossConfig::new().with_beta(0.5).init(); + + let predictions = Tensor::<1>::from_data(TensorData::from([0.25_f32, 1.0]), &device); + let targets = Tensor::<1>::zeros([2], &device); + + let output = loss.forward(predictions, targets); + let expected = TensorData::from([0.0625_f32, 0.75]); + output.into_data().assert_eq(&expected, false); + } + + // ========================================================================= + // forward_with_reduction Tests + // ========================================================================= + + #[test] + fn test_smooth_l1_reduction_mean() { + // Errors: 0.5 (L2: 0.125), 2.0 (L1: 1.5) + // Mean: (0.125 + 1.5) / 2 = 0.8125 + let device = Default::default(); + let loss = SmoothL1LossConfig::new().init(); + + let predictions = Tensor::<2>::from_data(TensorData::from([[0.5_f32, 2.0]]), &device); + let targets = Tensor::<2>::from_data(TensorData::from([[0.0_f32, 0.0]]), &device); + + let output = loss.forward_with_reduction(predictions, targets, Reduction::Mean); + let expected = TensorData::from([0.8125_f32]); + output.into_data().assert_eq(&expected, false); + } + + #[test] + fn test_smooth_l1_reduction_sum() { + // Errors: 0.5 (L2: 0.125), 2.0 (L1: 1.5) + // Sum: 1.625 + let device = Default::default(); + let loss = SmoothL1LossConfig::new().init(); + + let predictions = Tensor::<2>::from_data(TensorData::from([[0.5_f32, 2.0]]), &device); + let targets = Tensor::<2>::from_data(TensorData::from([[0.0_f32, 0.0]]), &device); + + let output = loss.forward_with_reduction(predictions, targets, Reduction::Sum); + let expected = TensorData::from([1.625_f32]); + output.into_data().assert_eq(&expected, false); + } + + #[test] + fn test_smooth_l1_reduction_auto_equals_mean() { + let device = Default::default(); + let loss = SmoothL1LossConfig::new().init(); + + let predictions = Tensor::<1>::from_data(TensorData::from([2.0_f32]), &device); + let targets = Tensor::<1>::zeros([1], &device); + + let mean_out = + loss.forward_with_reduction(predictions.clone(), targets.clone(), Reduction::Mean); + let auto_out = loss.forward_with_reduction(predictions, targets, Reduction::Auto); + + mean_out.into_data().assert_eq(&auto_out.into_data(), false); + } + + // ========================================================================= + // Dimension Reduction Tests + // ========================================================================= + + #[test] + fn test_smooth_l1_forward_reduce_dims_single_dim() { + // Beta = 2.0 + // L1 formula: |error| - 0.5 * beta + // L2 formula: 0.5 * (error)^2 / beta + // Row 0: errors [0.0, 1.0, 4.0] + // error = 0.0 -> L2: 0.0 + // error = 1.0 -> L2: 0.5 * 1.0 / 2.0 = 0.25 + // error = 4.0 -> L1: 4.0 - 1.0 = 3.0 + // Mean = 3.25 / 3 = 1.083333... + // Row 1: errors [0.0, 0.0, 0.0] -> Mean = 0.0 + let device = Default::default(); + let loss = SmoothL1LossConfig::new().with_beta(2.0).init(); + + let predictions = Tensor::<2>::from_data( + TensorData::from([[0.0_f32, 1.0, 4.0], [5.0_f32, 5.0, 5.0]]), + &device, + ); + let targets = Tensor::<2>::from_data( + TensorData::from([[0.0_f32, 0.0, 0.0], [5.0_f32, 5.0, 5.0]]), + &device, + ); + + let output = loss.forward_reduce_dims(predictions, targets, &[1]); + let expected = TensorData::from([[3.25_f32 / 3.0], [0.0]]); // 3.25/3 = 1.0833... + output + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn test_smooth_l1_forward_reduce_dims_image_batch() { + // Simulate per-image Smooth L1 loss for [batch, C, H, W] tensor + // (common in object detection like Fast R-CNN) + let device = Default::default(); + let loss = SmoothL1LossConfig::new().init(); // beta = 1.0 + + // Shape: [2, 1, 2, 2] (batch=2, C=1, H=2, W=2) + let predictions = Tensor::<4>::from_data( + TensorData::from([ + [[[0.5_f32, 2.0], [0.0, 3.0]]], // Image 1 + [[[1.0_f32, 0.0], [0.5, 1.5]]], // Image 2 + ]), + &device, + ); + let targets = Tensor::<4>::zeros([2, 1, 2, 2], &device); + + // Reduce over C, H, W (dims 1, 2, 3) to get per-image loss + let output = loss.forward_reduce_dims(predictions, targets, &[1, 2, 3]); + + // Image 1: losses [[0.125, 1.5], [0.0, 2.5]] -> mean: 4.125 / 4 = 1.03125 + // Image 2: losses [[0.5, 0.0], [0.125, 1.0]] -> mean: 1.625 / 4 = 0.40625 + let expected = TensorData::from([[[[1.03125_f32]]], [[[0.40625_f32]]]]); + output.into_data().assert_eq(&expected, false); + } + + #[test] + fn test_smooth_l1_forward_reduce_dims_unsorted() { + // Test that unsorted dimensions are handled correctly (sorted internally) + let device = Default::default(); + let loss = SmoothL1LossConfig::new().init(); + + let predictions = Tensor::<3>::from_data( + TensorData::from([[[1.0_f32, 2.0], [3.0, 4.0]], [[5.0_f32, 6.0], [7.0, 8.0]]]), + &device, + ); + let targets = Tensor::<3>::zeros([2, 2, 2], &device); + + // Pass dims in reverse order + let output = loss.forward_reduce_dims(predictions.clone(), targets.clone(), &[2, 1]); + let expected_output = loss.forward_reduce_dims(predictions, targets, &[1, 2]); + + output + .into_data() + .assert_eq(&expected_output.into_data(), false); + } + + #[test] + fn test_smooth_l1_forward_reduce_dims_empty_dims() { + // Reducing over no dimensions should return the unreduced loss + let device = Default::default(); + let loss = SmoothL1LossConfig::new().init(); + + let predictions = + Tensor::<2>::from_data(TensorData::from([[0.5_f32, 2.0], [0.0, 3.0]]), &device); + let targets = Tensor::<2>::zeros([2, 2], &device); + + let loss_reduce_dims = loss.forward_reduce_dims(predictions.clone(), targets.clone(), &[]); + let loss_no_reduction = loss.forward(predictions, targets); + + loss_reduce_dims + .into_data() + .assert_eq(&loss_no_reduction.into_data(), false); + } +} diff --git a/crates/burn-nn/src/modules/attention/cross_attention.rs b/crates/burn-nn/src/modules/attention/cross_attention.rs new file mode 100644 index 0000000..d555588 --- /dev/null +++ b/crates/burn-nn/src/modules/attention/cross_attention.rs @@ -0,0 +1,618 @@ +//! Cross-Attention Module for Burn +//! +//! Features: +//! - Asymmetric Input Shapes (Query vs Context) +//! - Grouped Query Attention (GQA) & Multi-Query Attention (MQA) support +//! - Quantization-Safe Masking (min_float) +//! - Sparse-Ready (quiet_softmax) +//! - KV Caching for Streaming Inference + +use crate::cache::TensorCache; +use crate::modules::{Linear, LinearConfig}; +use crate::{Dropout, DropoutConfig}; +use burn_core as burn; + +use burn::{ + config::Config, + module::{Initializer, Module}, + tensor::{ + Bool, Device, Tensor, + activation::{quiet_softmax, softmax}, + }, +}; + +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float as _; + +#[derive(Config, Debug)] +/// Configuration to create a [CrossAttention](CrossAttention) layer using the [init function](CrossAttentionConfig::init). +pub struct CrossAttentionConfig { + /// Dimension of the Query (e.g., Decoder state). + pub d_model: usize, + /// Dimension of the Context (e.g., Encoder audio embeddings). + pub d_context: usize, + /// Number of heads for the Query. + pub n_heads: usize, + /// Number of heads for Key/Value (Set to 1 for MQA, set to n_heads for MHA). + pub n_heads_kv: usize, + /// Dimension of a single head. + pub d_head: usize, + /// Dropout rate. + #[config(default = 0.1)] + pub dropout: f64, + /// Masking value. Use -1.0e4 for f16/bf16 safety. + #[config(default = -1.0e4)] + pub min_float: f64, + /// Use quiet_softmax to allow zero-attention (good for sparse/quantized models). + #[config(default = false)] + pub quiet_softmax: bool, +} + +#[derive(Module, Debug)] +/// The Cross attention module +/// +/// # Params +/// +/// - `query`: [`Linear`] layer with `d_model` input and output features. +/// - `key`: [`Linear`] layer with `d_model` input and output features. +/// - `value`: [`Linear`] layer with `d_model` input and output features. +/// - `output`: [`Linear`] layer with `d_model` input and output features. +/// +/// Should be created with [CrossAttentionConfig]. +pub struct CrossAttention { + query: Linear, + key: Linear, + value: Linear, + output: Linear, + dropout: Dropout, + + n_heads: usize, + n_heads_kv: usize, + d_head: usize, + scale: f64, + min_float: f64, + quiet_softmax: bool, +} + +/// Cache for the [Cross Attention](CrossAttention) layer. +/// +/// To be used during inference when context is constant. +pub struct CrossAttentionCache { + /// Cached key tensor. + pub k: TensorCache<4>, + /// Cached value tensor. + pub v: TensorCache<4>, +} + +impl CrossAttentionCache { + /// Create a new empty cache. + pub fn new() -> Self { + Self { + k: TensorCache::empty(), + v: TensorCache::empty(), + } + } +} + +impl Default for CrossAttentionCache { + fn default() -> Self { + Self::new() + } +} + +impl CrossAttentionConfig { + /// Initializes a new cross-attention module. + /// + /// # Arguments + /// + /// * `device` - The device on which to initialize the module. + /// + /// # Returns + /// + /// A new [CrossAttention] module. + pub fn init(&self, device: &Device) -> CrossAttention { + // Safety Rail for GQA + assert_eq!( + self.n_heads % self.n_heads_kv, + 0, + "Query heads must be divisible by KV heads" + ); + + let init_linear = |in_dim, out_dim| { + LinearConfig::new(in_dim, out_dim) + .with_initializer(Initializer::KaimingUniform { + gain: 1.0 / (self.d_head as f64).sqrt(), + fan_out_only: false, + }) + .init(device) + }; + + CrossAttention { + // ADVICE: Asymmetric Projections + query: init_linear(self.d_model, self.n_heads * self.d_head), + key: init_linear(self.d_context, self.n_heads_kv * self.d_head), + value: init_linear(self.d_context, self.n_heads_kv * self.d_head), + output: init_linear(self.n_heads * self.d_head, self.d_model), + + dropout: DropoutConfig::new(self.dropout).init(), + n_heads: self.n_heads, + n_heads_kv: self.n_heads_kv, + d_head: self.d_head, + scale: (self.d_head as f64).sqrt().recip(), + min_float: self.min_float, + quiet_softmax: self.quiet_softmax, + } + } +} + +impl CrossAttention { + /// Applies cross-attention to query using context as key and value. + /// + /// # Arguments + /// + /// * `query` - Query tensor of shape `[batch, seq_len_query, d_model]`. + /// * `context` - Context tensor of shape `[batch, seq_len_context, d_context]`. + /// * `mask` - Optional attention mask of shape `[batch, seq_len_context]` where `true` indicates positions to mask. + /// + /// # Returns + /// + /// Output tensor of shape `[batch, seq_len_query, d_model]`. + pub fn forward( + &self, + query: Tensor<3>, + context: Tensor<3>, + mask: Option>, + ) -> Tensor<3> { + let [batch, l_q, _] = query.dims(); + let [_, l_k, _] = context.dims(); + + // 1. Projections + let q = self.query.forward(query); + let k = self.key.forward(context.clone()); + let v = self.value.forward(context); + + // 2. Reshape Heads + // Q: [Batch, Heads, L_q, D_head] + let q = q + .reshape([batch, l_q, self.n_heads, self.d_head]) + .swap_dims(1, 2); + + // K, V: [Batch, Heads_KV, L_k, D_head] + let k = k + .reshape([batch, l_k, self.n_heads_kv, self.d_head]) + .swap_dims(1, 2); + let v = v + .reshape([batch, l_k, self.n_heads_kv, self.d_head]) + .swap_dims(1, 2); + + // 3. GQA Expansion + // ADVICE: Handle GQA by repeating KV heads to match Query heads + let (k, v) = if self.n_heads != self.n_heads_kv { + let n_rep = self.n_heads / self.n_heads_kv; + (self.repeat_kv(k, n_rep), self.repeat_kv(v, n_rep)) + } else { + (k, v) + }; + + // 4. Score Calculation + let scores = q.matmul(k.transpose()) * self.scale; + + // 5. Masking + // ADVICE: Use min_float for F16/FP8 safety + let scores = if let Some(mask) = mask { + let mask = mask.reshape([batch, 1, 1, l_k]); + scores.mask_fill(mask, self.min_float) + } else { + scores + }; + + // 6. Softmax + // ADVICE: Optional Quiet Softmax for sparse networks + let weights = if self.quiet_softmax { + quiet_softmax(scores, 3) + } else { + softmax(scores, 3) + }; + + let weights = self.dropout.forward(weights); + + // 7. Aggregate & Output + let output = weights.matmul(v); + let output = output + .swap_dims(1, 2) + .reshape([batch, l_q, self.n_heads * self.d_head]); + + self.output.forward(output) + } + + /// Applies cross-attention to query using context as key and value. + /// + /// This method uses a cache to avoid recomputing key and value tensors when the context is the same. + /// + /// # Arguments + /// + /// * `query` - Query tensor of shape `[batch, seq_len_query, d_model]`. + /// * `context` - Context tensor of shape `[batch, seq_len_context, d_context]`. + /// * `mask` - Optional attention mask of shape `[batch, seq_len_context]` where `true` indicates positions to mask. + /// * `cache` - The cache to use. + /// + /// # Returns + /// + /// Output tensor of shape `[batch, seq_len_query, d_model]`. + pub fn forward_cache( + &self, + query: Tensor<3>, + context: Tensor<3>, + mask: Option>, + cache: &mut CrossAttentionCache, + ) -> Tensor<3> { + let [batch, l_q, _] = query.dims(); + + // 1. Projections + let q = self.query.forward(query); + + let k_compute = |context: Tensor<3>| { + let [batch, l_k, _] = context.dims(); + self.key + .forward(context) + .reshape([batch, l_k, self.n_heads_kv, self.d_head]) + .swap_dims(1, 2) + }; + let v_compute = |context: Tensor<3>| { + let [batch, l_k, _] = context.dims(); + self.value + .forward(context) + .reshape([batch, l_k, self.n_heads_kv, self.d_head]) + .swap_dims(1, 2) + }; + + let k = cache.k.forward_full(context.clone(), k_compute); + let v = cache.v.forward_full(context, v_compute); + + let [_, _, l_k, _] = k.dims(); + + // 2. Reshape Heads + // Q: [Batch, Heads, L_q, D_head] + let q = q + .reshape([batch, l_q, self.n_heads, self.d_head]) + .swap_dims(1, 2); + + // K, V are already in their correct shape from k_compute and v_compute + + // 3. GQA Expansion + // ADVICE: Handle GQA by repeating KV heads to match Query heads + let (k, v) = if self.n_heads != self.n_heads_kv { + let n_rep = self.n_heads / self.n_heads_kv; + (self.repeat_kv(k, n_rep), self.repeat_kv(v, n_rep)) + } else { + (k, v) + }; + + // 4. Score Calculation + let scores = q.matmul(k.transpose()) * self.scale; + + // 5. Masking + // ADVICE: Use min_float for F16/FP8 safety + let scores = if let Some(mask) = mask { + let mask = mask.reshape([batch, 1, 1, l_k]); + scores.mask_fill(mask, self.min_float) + } else { + scores + }; + + // 6. Softmax + // ADVICE: Optional Quiet Softmax for sparse networks + let weights = if self.quiet_softmax { + quiet_softmax(scores, 3) + } else { + softmax(scores, 3) + }; + + let weights = self.dropout.forward(weights); + + // 7. Aggregate & Output + let output = weights.matmul(v); + let output = output + .swap_dims(1, 2) + .reshape([batch, l_q, self.n_heads * self.d_head]); + + self.output.forward(output) + } + + /// Helper for Grouped Query Attention + fn repeat_kv(&self, x: Tensor<4>, n_rep: usize) -> Tensor<4> { + let [b, h, l, d] = x.dims(); + x.reshape([b, h, 1, l, d]) + .expand([b, h, n_rep, l, d]) + .reshape([b, h * n_rep, l, d]) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::{Distribution, Int, Shape, Tensor, Tolerance}; + + #[test] + fn test_cross_attention_mha_shapes() { + let [ + batch_size, + seq_len_query, + seq_len_context, + d_model, + d_context, + n_heads, + d_head, + ] = [7, 13, 15, 32, 40, 4, 8]; + let device = Default::default(); + let config = CrossAttentionConfig { + d_model, + d_context, + n_heads, + n_heads_kv: n_heads, // MHA case + d_head, + dropout: 0.1, + min_float: -1.0e4, + quiet_softmax: false, + }; + let cross_attn = config.init(&device); + + let query = Tensor::random( + [batch_size, seq_len_query, d_model], + Distribution::Default, + &device, + ); + let context = Tensor::random( + [batch_size, seq_len_context, d_context], + Distribution::Default, + &device, + ); + + let output = cross_attn.forward(query, context, None); + + assert_eq!( + output.shape(), + Shape::new([batch_size, seq_len_query, d_model]), + "Output should have the correct shape", + ); + } + + #[test] + fn test_cross_attention_gqa_shapes() { + let [ + batch_size, + seq_len_query, + seq_len_context, + d_model, + d_context, + n_heads, + n_heads_kv, + d_head, + ] = [7, 13, 15, 32, 40, 4, 2, 8]; + let device = Default::default(); + let config = CrossAttentionConfig { + d_model, + d_context, + n_heads, + n_heads_kv, // GQA case + d_head, + dropout: 0.1, + min_float: -1.0e4, + quiet_softmax: false, + }; + let cross_attn = config.init(&device); + + let query = Tensor::random( + [batch_size, seq_len_query, d_model], + Distribution::Default, + &device, + ); + let context = Tensor::random( + [batch_size, seq_len_context, d_context], + Distribution::Default, + &device, + ); + + let output = cross_attn.forward(query, context, None); + + assert_eq!( + output.shape(), + Shape::new([batch_size, seq_len_query, d_model]), + "Output should have the correct shape", + ); + } + + #[test] + fn test_cross_attention_mqa_shapes() { + let [ + batch_size, + seq_len_query, + seq_len_context, + d_model, + d_context, + n_heads, + d_head, + ] = [7, 13, 15, 32, 40, 4, 8]; + let device = Default::default(); + let config = CrossAttentionConfig { + d_model, + d_context, + n_heads, + n_heads_kv: 1, // MQA case + d_head, + dropout: 0.1, + min_float: -1.0e4, + quiet_softmax: false, + }; + let cross_attn = config.init(&device); + + let query = Tensor::random( + [batch_size, seq_len_query, d_model], + Distribution::Default, + &device, + ); + let context = Tensor::random( + [batch_size, seq_len_context, d_context], + Distribution::Default, + &device, + ); + + let output = cross_attn.forward(query, context, None); + + assert_eq!( + output.shape(), + Shape::new([batch_size, seq_len_query, d_model]), + "Output should have the correct shape", + ); + } + + #[test] + fn test_cross_attention_mask() { + let [ + batch_size, + seq_len_query, + seq_len_context, + d_model, + d_context, + n_heads, + d_head, + ] = [3, 6, 8, 12, 16, 4, 3]; + let num_padded = 2; + let device = Default::default(); + let config = CrossAttentionConfig { + d_model, + d_context, + n_heads, + n_heads_kv: n_heads, + d_head, + dropout: 0.0, // No dropout for deterministic test + min_float: -1.0e4, + quiet_softmax: false, + }; + let cross_attn = config.init(&device); + + // Create a padding mask for the context + let mut mask: Tensor<2, Int> = Tensor::zeros([batch_size, seq_len_context], &device); + mask = mask.slice_assign( + [0..batch_size, seq_len_context - num_padded..seq_len_context], + Tensor::ones([batch_size, num_padded], &device), + ); + let mask_bool = mask.equal_elem(1); + + let query = Tensor::<3>::random( + [batch_size, seq_len_query, d_model], + Distribution::Default, + &device, + ); + + let context_1 = Tensor::<3>::random( + [batch_size, seq_len_context, d_context], + Distribution::Default, + &device, + ); + + // Change the padded part of the context tensor + let context_2 = context_1.clone().slice_assign( + [ + 0..batch_size, + seq_len_context - num_padded..seq_len_context, + 0..d_context, + ], + Tensor::random( + [batch_size, num_padded, d_context], + Distribution::Default, + &device, + ), + ); + + // The outputs should be the same since the changed part is masked. + let output_1 = cross_attn.forward(query.clone(), context_1, Some(mask_bool.clone())); + let output_2 = cross_attn.forward(query, context_2, Some(mask_bool)); + + output_1 + .into_data() + .assert_approx_eq(&output_2.into_data(), Tolerance::::default()); + } + + #[test] + #[should_panic] + fn test_gqa_panic_if_n_heads_not_divisible_by_n_heads_kv() { + let device = Default::default(); + let config = CrossAttentionConfig { + d_model: 32, + d_context: 32, + n_heads: 5, + n_heads_kv: 2, + d_head: 8, + dropout: 0.1, + min_float: -1.0e4, + quiet_softmax: false, + }; + config.init(&device); + } + + #[test] + fn test_cross_attention_cache() { + let [ + batch_size, + seq_len_query, + seq_len_context, + d_model, + d_context, + n_heads, + d_head, + ] = [3, 6, 8, 12, 16, 4, 3]; + let device = Default::default(); + let config = CrossAttentionConfig { + d_model, + d_context, + n_heads, + n_heads_kv: n_heads, + d_head, + dropout: 0.0, // No dropout for deterministic test + min_float: -1.0e4, + quiet_softmax: false, + }; + let cross_attn = config.init(&device); + + let query1 = Tensor::<3>::random( + [batch_size, seq_len_query, d_model], + Distribution::Default, + &device, + ); + let context = Tensor::<3>::random( + [batch_size, seq_len_context, d_context], + Distribution::Default, + &device, + ); + + // First forward pass, no cache + let output1 = cross_attn.forward(query1.clone(), context.clone(), None); + + // Second forward pass with cache + let mut cache = CrossAttentionCache::new(); + let output2 = cross_attn.forward_cache(query1.clone(), context.clone(), None, &mut cache); + + // The two outputs should be identical + output1 + .into_data() + .assert_approx_eq(&output2.into_data(), Tolerance::::default()); + + // Third forward pass with different query, but same context and cache + let query2 = Tensor::<3>::random( + [batch_size, seq_len_query, d_model], + Distribution::Default, + &device, + ); + let output3 = cross_attn.forward_cache(query2.clone(), context.clone(), None, &mut cache); + + // For control, do a forward pass without cache with query2 + let output4 = cross_attn.forward(query2.clone(), context.clone(), None); + + // output3 and output4 should be identical + output3 + .into_data() + .assert_approx_eq(&output4.into_data(), Tolerance::::default()); + } +} diff --git a/crates/burn-nn/src/modules/attention/mask.rs b/crates/burn-nn/src/modules/attention/mask.rs new file mode 100644 index 0000000..f123027 --- /dev/null +++ b/crates/burn-nn/src/modules/attention/mask.rs @@ -0,0 +1,155 @@ +use burn_core as burn; +use burn_core::config::Config; + +use alloc::vec::Vec; + +use burn::tensor::{Bool, Device, Int, Shape, Tensor, TensorData}; + +/// Generate an autoregressive attention mask. +/// +/// The mask can be used in Transformer modules to train models to generate tensors sequentially. +pub fn generate_autoregressive_mask( + batch_size: usize, + seq_length: usize, + device: &Device, +) -> Tensor<3, Bool> { + let mask = Tensor::<2, Bool>::tril_mask([seq_length, seq_length], 0, device); + mask.expand([batch_size, seq_length, seq_length]) +} + +/// Generate a padding attention mask. +pub struct GeneratePaddingMask { + /// The generated tensor. + pub tensor: Tensor<2, Int>, + + /// The generated mask. + pub mask: Tensor<2, Bool>, +} + +/// Defines an enumeration to specify sequence length options for padding +#[derive(Config, Debug, Copy)] +pub enum SeqLengthOption { + /// No maximum length; use the longest sequence + NoMax, + /// Maximum length specified, truncate if necessary + Max(usize), + /// Fixed length, pad or truncate to this exact length + Fixed(usize), +} + +impl From> for SeqLengthOption { + fn from(val: Option) -> Self { + match val { + Some(max) => SeqLengthOption::Max(max), + None => SeqLengthOption::NoMax, + } + } +} + +/// Generates a padding attention mask for a batch of token sequences. +/// +/// # Arguments +/// +/// * `pad_token` - The token ID used for padding +/// * `tokens_list` - Vector of token sequences (each sequence is a vector of token IDs) +/// * `seq_length` - Sequence length option (NoMax, Max, or Fixed) +/// * `device` - The device for tensor operations +/// +/// # Returns +/// +/// A `GeneratePaddingMask` containing the padded tensor and corresponding mask +pub fn generate_padding_mask( + pad_token: usize, + tokens_list: Vec>, + seq_length: impl Into, + device: &Device, +) -> GeneratePaddingMask { + let tokens_max = || { + tokens_list + .iter() + .map(|tokens| tokens.len()) + .max() + .unwrap_or(1) + }; + + let size = match seq_length.into() { + SeqLengthOption::NoMax => tokens_max(), + SeqLengthOption::Max(max) => usize::min(tokens_max(), max), + SeqLengthOption::Fixed(limit) => limit, + }; + let batch_size = tokens_list.len(); + + let mut tensor = Tensor::zeros([batch_size, size], device); + tensor = tensor.add_scalar(pad_token as i64); + + for (index, tokens) in tokens_list.into_iter().enumerate() { + let seq_length = tokens.len().min(size); + tensor = tensor.slice_assign( + [index..index + 1, 0..seq_length], + Tensor::from_data( + TensorData::new( + tokens.into_iter().take(size).map(|e| e as i64).collect(), + Shape::new([1, seq_length]), + ), + device, + ), + ); + } + + let mask = tensor.clone().equal_elem(pad_token as i64); + + GeneratePaddingMask { tensor, mask } +} + +#[cfg(test)] +mod tests { + use super::*; + use alloc::vec; + use burn::tensor::TensorData; + + #[test] + fn test_generate_autoregressive_mask() { + let device = Default::default(); + + let mask = generate_autoregressive_mask(2, 3, &device); + + mask.into_data().assert_eq( + &TensorData::from([ + [ + [false, true, true], + [false, false, true], + [false, false, false], + ], + [ + [false, true, true], + [false, false, true], + [false, false, false], + ], + ]), + false, + ); + } + + #[test] + fn test_generate_padding_mask() { + let device = Default::default(); + let tokens = vec![ + vec![3, 3, 3], + vec![3, 3, 3], + vec![3, 3, 3, 4], + vec![3, 3, 3, 4, 10, 15], + ]; + + let mask = generate_padding_mask(0, tokens, None, &device); + + mask.mask.into_data().assert_eq( + &TensorData::from([ + [false, false, false, true, true, true], + [false, false, false, true, true, true], + [false, false, false, false, true, true], + [false, false, false, false, false, false], + ]), + false, + ); + } +} diff --git a/crates/burn-nn/src/modules/attention/mha.rs b/crates/burn-nn/src/modules/attention/mha.rs new file mode 100644 index 0000000..51fc2ee --- /dev/null +++ b/crates/burn-nn/src/modules/attention/mha.rs @@ -0,0 +1,578 @@ +use burn_core as burn; + +use crate::activation::Gelu; +use crate::cache::TensorCache; +use crate::{Dropout, DropoutConfig, Linear, LinearConfig}; +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Initializer, Module, ModuleDisplay}; +use burn::tensor::{Bool, Device, Tensor}; + +use burn::tensor::activation::{quiet_softmax, softmax}; +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float as _; + +/// Configuration to create a [Multi Head Attention](MultiHeadAttention) layer using the [init function](MultiHeadAttentionConfig::init). +#[derive(Config, Debug)] +pub struct MultiHeadAttentionConfig { + /// The size of each linear layer. + pub d_model: usize, + /// The number of heads. + pub n_heads: usize, + /// The dropout rate. Default: 0.1 + #[config(default = 0.1)] + pub dropout: f64, + /// The minimum value a float can take. Default: -1.0e4 + /// This is used to mask attention scores before calculating attention weights. + /// A value too low might result in NaN. + #[config(default = -1.0e4)] + pub min_float: f64, + /// Use "quiet softmax" instead of regular softmax. + /// + /// - Usage may improve performance by allowing attention heads to deposit no information (if the sequence contains no information relevant to that head). + /// - Usage may reduce the entropy of weights in the model, enhancing quantization and compression. + /// + /// Reference: + #[config(default = false)] + pub quiet_softmax: bool, + /// The type of function used to initialize neural network parameters + #[config( + default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0), fan_out_only:false}" + )] + pub initializer: Initializer, + + /// Enable bias for the key [Linear] layer. + #[config(default = true)] + pub key_bias: bool, + + /// Enable bias for the value [Linear] layer. + #[config(default = true)] + pub value_bias: bool, + + /// Enable bias for the query [Linear] layer. + #[config(default = true)] + pub query_bias: bool, + + /// Enable bias for the output [Linear] layer. + #[config(default = true)] + pub output_bias: bool, +} + +/// The multihead attention module as describe in the paper [Attention Is All You Need](https://arxiv.org/abs/1706.03762). +/// +/// # Params +/// +/// - `query`: [`Linear`] layer with `d_model` input and output features. +/// - `key`: [`Linear`] layer with `d_model` input and output features. +/// - `value`: [`Linear`] layer with `d_model` input and output features. +/// - `output`: [`Linear`] layer with `d_model` input and output features. +/// +/// Should be created with [MultiHeadAttentionConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct MultiHeadAttention { + /// Linear layer to transform the input features into the query space. + pub query: Linear, + /// Linear layer to transform the input features into the key space. + pub key: Linear, + /// Linear layer to transform the input features into the value space. + pub value: Linear, + /// Linear layer to transform the output features back to the original space. + pub output: Linear, + /// Dropout layer. + pub dropout: Dropout, + /// Activation function. + pub activation: Gelu, + /// The size of each linear layer. + pub d_model: usize, + /// The number of heads. + pub n_heads: usize, + /// Size of the key and query vectors. + pub d_k: usize, + /// Minimum value a float can take. + pub min_float: f64, + /// Use "quiet softmax" instead of regular softmax. + pub quiet_softmax: bool, +} + +impl ModuleDisplay for MultiHeadAttention { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("d_model", &self.d_model) + .add("n_heads", &self.n_heads) + .add("d_k", &self.d_k) + .add("dropout", &self.dropout.prob) + .add("min_float", &self.min_float) + .add("quiet_softmax", &self.quiet_softmax) + .optional() + } +} + +/// [Multihead attention](MultiHeadAttention) forward pass input argument. +#[derive(Debug, Clone)] +pub struct MhaInput { + /// Shape `[batch_size, seq_length_1, d_model]` + query: Tensor<3>, + /// Shape `[batch_size, seq_length_2, d_model]` + key: Tensor<3>, + /// Shape `[batch_size, seq_length_2, d_model]` + value: Tensor<3>, + mask_pad: Option>, + mask_attn: Option>, +} + +impl MultiHeadAttentionConfig { + /// Initialize a new [multihead attention](MultiHeadAttention) module. + pub fn init(&self, device: &Device) -> MultiHeadAttention { + let linear_cfg = LinearConfig::new(self.d_model, self.d_model) + .with_initializer(self.initializer.clone()); + + MultiHeadAttention { + query: linear_cfg.clone().with_bias(self.query_bias).init(device), + key: linear_cfg.clone().with_bias(self.key_bias).init(device), + value: linear_cfg.clone().with_bias(self.value_bias).init(device), + output: linear_cfg.clone().with_bias(self.output_bias).init(device), + dropout: DropoutConfig::new(self.dropout).init(), + activation: Gelu::new(), + n_heads: self.n_heads, + d_k: self.d_model / self.n_heads, + min_float: self.min_float, + quiet_softmax: self.quiet_softmax, + d_model: self.d_model, + } + } +} + +impl MhaInput { + /// Create a [multihead attention](MultiHeadAttention) input argument + /// by setting the query, key and value to the given tensor. + /// + /// # Shape + /// - tensor: `[batch_size, seq_length, d_model]` + pub fn self_attn(tensor: Tensor<3>) -> Self { + Self { + query: tensor.clone(), + key: tensor.clone(), + value: tensor, + mask_pad: None, + mask_attn: None, + } + } + + /// Create a [multihead attention](MultiHeadAttention) input argument. + pub fn new(query: Tensor<3>, key: Tensor<3>, value: Tensor<3>) -> Self { + Self { + query, + key, + value, + mask_pad: None, + mask_attn: None, + } + } + + /// Register the padding mask. + pub fn mask_pad(mut self, mask_pad: Tensor<2, Bool>) -> Self { + self.mask_pad = Some(mask_pad); + self + } + + /// Register the attention mask. + pub fn mask_attn(mut self, mask_attn: Tensor<3, Bool>) -> Self { + self.mask_attn = Some(mask_attn); + self + } +} + +/// [Multihead attention](MultiHeadAttention) outputs. +#[derive(Debug, Clone)] +pub struct MhaOutput { + /// The attention weights `[batch_size, n_heads, seq_length_1, seq_length_2]`. + pub weights: Tensor<4>, + /// The context tensor `[batch_size, seq_length_1, d_model]`. + pub context: Tensor<3>, +} + +impl MultiHeadAttention { + /// Applies the forward pass on the input tensors. + /// + /// See [MultiHeadAttention](MultiHeadAttention) for more information. + /// + /// # Shapes + /// + /// - query: `[batch_size, seq_length_1, d_model]` + /// - key: `[batch_size, seq_length_2, d_model]` + /// - value: `[batch_size, seq_length_2, d_model]` + /// - output: `[batch_size, seq_length_1, d_model]` + pub fn forward(&self, input: MhaInput) -> MhaOutput { + let [batch_size, seq_length_1, d_model] = input.query.dims(); + + let query = self.attention_linear(input.query, &self.query); + let key = self.attention_linear(input.key, &self.key); + let value = self.attention_linear(input.value, &self.value); + + let attn_scores = self.attn_scores(query, key); + let weights = self.attn_weights(attn_scores, input.mask_pad, input.mask_attn); + + let context = weights.clone().matmul(value); + let context = context + .swap_dims(1, 2) + .reshape([batch_size, seq_length_1, d_model]); + let context = self.output.forward(context); + + MhaOutput { weights, context } + } + + /// Applies the forward pass using a cache. + /// + /// # Shapes + /// + /// - query: `[batch_size, seq_length_1, d_model]` + /// - key: `[batch_size, seq_length_2, d_model]` + /// - value: `[batch_size, seq_length_2, d_model]` + /// - output: `[batch_size, seq_length_1, d_model]` + pub fn forward_cache(&self, input: MhaInput, cache: &mut MhaCache) -> MhaOutput { + let [batch_size, seq_length_1, d_model] = input.query.dims(); + + let query = cache + .query + .forward(input.query, |t| self.attention_linear(t, &self.query)); + let key = cache + .key + .forward(input.key, |t| self.attention_linear(t, &self.key)); + let value = cache + .value + .forward(input.value, |t| self.attention_linear(t, &self.value)); + + let attn_scores = self.attn_scores(query, key); + let weights = self.attn_weights(attn_scores, input.mask_pad, input.mask_attn); + + let context = weights.clone().matmul(value); + let context = context + .swap_dims(1, 2) + .reshape([batch_size, seq_length_1, d_model]); + + let context = cache.output.forward(context, |t| self.output.forward(t)); + + MhaOutput { weights, context } + } + + fn attn_scores(&self, query: Tensor<4>, key: Tensor<4>) -> Tensor<4> { + let attn_scores = query + .matmul(key.transpose()) + .div_scalar((self.d_k as f32).sqrt()); + + self.dropout.forward(attn_scores) + } + + fn attn_weights( + &self, + mut attn_scores: Tensor<4>, + mask_pad: Option>, + mask_attn: Option>, + ) -> Tensor<4> { + if let Some(mask_pad) = mask_pad { + let [batch_size, seq_length] = mask_pad.dims(); + + attn_scores = attn_scores.mask_fill( + mask_pad.reshape([batch_size, 1, 1, seq_length]), + self.min_float, + ); + } + + if let Some(mask_attn) = mask_attn { + let [batch_size, seq_length_1, seq_length_2] = mask_attn.dims(); + + attn_scores = attn_scores.mask_fill( + mask_attn.reshape([batch_size, 1, seq_length_1, seq_length_2]), + self.min_float, + ); + } + + if self.quiet_softmax { + quiet_softmax(attn_scores, 3) + } else { + softmax(attn_scores, 3) + } + } + + fn attention_linear(&self, x: Tensor<3>, linear: &Linear) -> Tensor<4> { + let [batch_size, seq_length, _d_model] = x.dims(); + linear + .forward(x) + .reshape([batch_size, seq_length, self.n_heads, self.d_k]) + .swap_dims(1, 2) + } +} + +/// Cache for the [Multi Head Attention](MultiHeadAttention) layer. +/// +/// To be used during inference when decoding tokens. +pub struct MhaCache { + query: MhaLinearCache<4>, + key: MhaLinearCache<4>, + value: MhaLinearCache<4>, + output: MhaLinearCache<3>, +} + +enum MhaLinearCache { + Autoregressive(TensorCache, usize), + Full(TensorCache), +} + +impl MhaCache { + /// Initialize a cache for autoregressive inference. + pub fn autoregressive() -> Self { + Self { + query: MhaLinearCache::Autoregressive(TensorCache::empty(), 2), + key: MhaLinearCache::Autoregressive(TensorCache::empty(), 2), + value: MhaLinearCache::Autoregressive(TensorCache::empty(), 2), + output: MhaLinearCache::Autoregressive(TensorCache::empty(), 1), + } + } + + /// Initialize a cache for autoregressive inference, but with a fixed memory used for keys and + /// values (cross-attention). + pub fn autoregressive_cross_attention() -> Self { + Self { + query: MhaLinearCache::Autoregressive(TensorCache::empty(), 2), + key: MhaLinearCache::Full(TensorCache::empty()), + value: MhaLinearCache::Full(TensorCache::empty()), + output: MhaLinearCache::Autoregressive(TensorCache::empty(), 1), + } + } +} + +impl MhaLinearCache { + pub fn forward) -> Tensor>( + &mut self, + tensor: Tensor<3>, + func: F, + ) -> Tensor { + match self { + MhaLinearCache::Autoregressive(cache, dim) => { + cache.forward_autoregressive(tensor, *dim, func) + } + MhaLinearCache::Full(cache) => cache.forward_full(tensor, func), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::attention::generate_autoregressive_mask; + use alloc::vec::Vec; + use burn::tensor::Int; + use burn::tensor::Tolerance; + use burn::tensor::{Distribution, Shape}; + + #[test] + fn test_enable_bias() { + let d_model = 32; + let n_heads = 4; + let device = Default::default(); + + let base_cfg = MultiHeadAttentionConfig::new(d_model, n_heads); + + // Default, everything is enabled. + let mha = base_cfg.init(&device); + assert!(mha.query.bias.is_some()); + assert!(mha.key.bias.is_some()); + assert!(mha.value.bias.is_some()); + assert!(mha.output.bias.is_some()); + + let mha = base_cfg.clone().with_query_bias(false).init(&device); + assert!(mha.query.bias.is_none()); + assert!(mha.key.bias.is_some()); + assert!(mha.value.bias.is_some()); + assert!(mha.output.bias.is_some()); + + let mha = base_cfg.clone().with_key_bias(false).init(&device); + assert!(mha.query.bias.is_some()); + assert!(mha.key.bias.is_none()); + assert!(mha.value.bias.is_some()); + assert!(mha.output.bias.is_some()); + + let mha = base_cfg.clone().with_value_bias(false).init(&device); + assert!(mha.query.bias.is_some()); + assert!(mha.key.bias.is_some()); + assert!(mha.value.bias.is_none()); + assert!(mha.output.bias.is_some()); + + let mha = base_cfg.clone().with_output_bias(false).init(&device); + assert!(mha.query.bias.is_some()); + assert!(mha.key.bias.is_some()); + assert!(mha.value.bias.is_some()); + assert!(mha.output.bias.is_none()); + } + + #[test] + fn test_self_attention_shapes() { + let [batch_size, seq_length, d_model, n_heads] = [7, 13, 32, 4]; + let device = Default::default(); + let mha = MultiHeadAttentionConfig::new(d_model, n_heads).init(&device); + let input = MhaInput::self_attn(Tensor::random( + [batch_size, seq_length, d_model], + Distribution::Default, + &device, + )); + + let output = mha.forward(input); + + assert_eq!( + output.context.shape(), + Shape::new([batch_size, seq_length, d_model]), + "Context should have the correct shape", + ); + assert_eq!( + output.weights.shape(), + Shape::new([batch_size, n_heads, seq_length, seq_length]), + "Weights should have the correct shape", + ); + } + + #[test] + fn test_generic_mha_shapes() { + let [batch_size, seq_length_1, seq_length_2, d_model, n_heads] = [7, 13, 15, 32, 4]; + let mha = MultiHeadAttentionConfig::new(d_model, n_heads).init(&Default::default()); + let device = Default::default(); + let input = MhaInput::new( + Tensor::random( + [batch_size, seq_length_1, d_model], + Distribution::Default, + &device, + ), + Tensor::random( + [batch_size, seq_length_2, d_model], + Distribution::Default, + &device, + ), + Tensor::random( + [batch_size, seq_length_2, d_model], + Distribution::Default, + &device, + ), + ); + + let output = mha.forward(input); + + assert_eq!( + output.context.shape(), + Shape::new([batch_size, seq_length_1, d_model]), + "Context should have the correct shape", + ); + assert_eq!( + output.weights.shape(), + Shape::new([batch_size, n_heads, seq_length_1, seq_length_2]), + "Weights should have the correct shape", + ); + } + + #[test] + fn test_self_attention_mask_pad() { + let [batch_size, seq_length, d_model, n_heads, num_padded] = [3, 6, 32, 2, 2]; + let device = Default::default(); + let mha = MultiHeadAttentionConfig::new(d_model, n_heads).init(&device); + + // Create a padding mask + let mask_pad: Tensor<2, Int> = Tensor::zeros([batch_size, seq_length], &device); + let mask_pad = mask_pad.slice_assign( + [0..batch_size, seq_length - num_padded..seq_length], + Tensor::ones([batch_size, num_padded], &device), + ); + let mask_pad = mask_pad.equal_elem(1).to_device(&device); + + let tensor_1 = Tensor::<3>::random( + [batch_size, seq_length, d_model], + Distribution::Default, + &device, + ); + // Change the end of the tensor + let tensor_2 = tensor_1.clone().slice_assign( + [ + 0..batch_size, + seq_length - num_padded..seq_length, + 0..d_model, + ], + Tensor::random( + [batch_size, num_padded, d_model], + Distribution::Default, + &device, + ), + ); + + let input_1 = MhaInput::self_attn(tensor_1).mask_pad(mask_pad.clone()); + let input_2 = MhaInput::self_attn(tensor_2).mask_pad(mask_pad); + + let output_1 = mha.forward(input_1); + let output_2 = mha.forward(input_2); + + // Check that the beginning of each tensor is the same + output_1 + .context + .slice([0..batch_size, 0..seq_length - num_padded, 0..d_model]) + .into_data() + .assert_approx_eq( + &output_2 + .context + .slice([0..batch_size, 0..seq_length - num_padded, 0..d_model]) + .into_data(), + Tolerance::::default(), + ); + } + + #[test] + fn test_autoregressive_mask_should_have_same_output_as_autoregressive_decoding() { + let [batch_size, seq_length, d_model, n_heads] = [3, 4, 12, 2]; + let device = Default::default(); + let mha = MultiHeadAttentionConfig::new(d_model, n_heads).init(&device); + + let tensor = Tensor::<3>::random( + [batch_size, seq_length, d_model], + Distribution::Default, + &device, + ); + let mask_attn = generate_autoregressive_mask(batch_size, seq_length, &tensor.device()); + let input = MhaInput::self_attn(tensor.clone()).mask_attn(mask_attn); + + let output_1 = mha.forward(input); + let mut output_2 = Vec::new(); + let mut cache = MhaCache::autoregressive(); + + for i in 1..seq_length + 1 { + let tensor = tensor.clone().slice([0..batch_size, 0..i, 0..d_model]); + let input = MhaInput::self_attn(tensor); + let next_tok = mha.forward_cache(input, &mut cache).context.slice([ + 0..batch_size, + i - 1..i, + 0..d_model, + ]); + output_2.push(next_tok); + } + + let output_2 = Tensor::cat(output_2, 1); + + output_1 + .context + .into_data() + .assert_approx_eq::(&output_2.into_data(), Tolerance::default()); + } + + #[test] + fn display() { + let config = MultiHeadAttentionConfig::new(2, 4); + let mha = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{mha}"), + "MultiHeadAttention {d_model: 2, n_heads: 4, d_k: 0, \ + dropout: 0.1, min_float: -10000, quiet_softmax: false, params: 24}" + ); + } +} diff --git a/crates/burn-nn/src/modules/attention/mod.rs b/crates/burn-nn/src/modules/attention/mod.rs new file mode 100644 index 0000000..460e473 --- /dev/null +++ b/crates/burn-nn/src/modules/attention/mod.rs @@ -0,0 +1,7 @@ +mod cross_attention; +mod mask; +mod mha; + +pub use cross_attention::*; +pub use mask::*; +pub use mha::*; diff --git a/crates/burn-nn/src/modules/cache/autoregressive.rs b/crates/burn-nn/src/modules/cache/autoregressive.rs new file mode 100644 index 0000000..1769741 --- /dev/null +++ b/crates/burn-nn/src/modules/cache/autoregressive.rs @@ -0,0 +1,51 @@ +use alloc::vec; +use burn_core as burn; + +use super::{CacheState, TensorCache}; +use burn::tensor::Tensor; + +impl TensorCache { + pub(crate) fn forward_autoregressive( + &mut self, + tensor: Tensor<3>, + dim_cat: usize, + func: F, + ) -> Tensor + where + F: Fn(Tensor<3>) -> Tensor, + { + let mut tensor_old = CacheState::Empty; + core::mem::swap(&mut self.state, &mut tensor_old); + + let tensor_new = match tensor_old { + CacheState::Value(tensor_old) => { + let [batch_size, seq_length, d_model] = tensor.dims(); + let next_seq_token = + tensor.slice([0..batch_size, (seq_length - 1)..seq_length, 0..d_model]); + let next_seq_token = func(next_seq_token); + + Tensor::cat(vec![tensor_old, next_seq_token], dim_cat) + } + _ => func(tensor), + }; + + self.state = CacheState::Value(tensor_new.clone()); + tensor_new + } + + pub(crate) fn forward_full(&mut self, tensor: Tensor<3>, func: F) -> Tensor + where + F: Fn(Tensor<3>) -> Tensor, + { + let mut tensor_old = CacheState::Empty; + core::mem::swap(&mut self.state, &mut tensor_old); + + let tensor_new = match tensor_old { + CacheState::Value(tensor_old) => tensor_old, + _ => func(tensor), + }; + + self.state = CacheState::Value(tensor_new.clone()); + tensor_new + } +} diff --git a/crates/burn-nn/src/modules/cache/base.rs b/crates/burn-nn/src/modules/cache/base.rs new file mode 100644 index 0000000..b0d547a --- /dev/null +++ b/crates/burn-nn/src/modules/cache/base.rs @@ -0,0 +1,26 @@ +use burn_core as burn; + +use burn::tensor::Tensor; + +pub(crate) enum CacheState { + Value(T), + Empty, +} + +/// A cache for a tensor. +pub struct TensorCache { + pub(crate) state: CacheState>, +} + +impl TensorCache { + /// Creates a new empty cache. + /// + /// # Returns + /// + /// The empty cache. + pub fn empty() -> Self { + Self { + state: CacheState::Empty, + } + } +} diff --git a/crates/burn-nn/src/modules/cache/mod.rs b/crates/burn-nn/src/modules/cache/mod.rs new file mode 100644 index 0000000..8050cb4 --- /dev/null +++ b/crates/burn-nn/src/modules/cache/mod.rs @@ -0,0 +1,4 @@ +mod autoregressive; +mod base; + +pub use base::*; diff --git a/crates/burn-nn/src/modules/conv/checks.rs b/crates/burn-nn/src/modules/conv/checks.rs new file mode 100644 index 0000000..0e24c27 --- /dev/null +++ b/crates/burn-nn/src/modules/conv/checks.rs @@ -0,0 +1,22 @@ +pub(crate) fn checks_channels_div_groups(channels_in: usize, channels_out: usize, groups: usize) { + let channels_in_div_by_group = channels_in.is_multiple_of(groups); + let channels_out_div_by_group = channels_out.is_multiple_of(groups); + + if !channels_in_div_by_group || !channels_out_div_by_group { + panic!( + "Both channels must be divisible by the number of groups. Got \ + channels_in={channels_in}, channels_out={channels_out}, groups={groups}" + ); + } +} + +// https://github.com/tracel-ai/burn/issues/2676 +/// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel +/// size is not supported as it will not produce the same output size. +pub(crate) fn check_same_padding_support(kernel_size: &[usize]) { + for k in kernel_size.iter() { + if k % 2 == 0 { + unimplemented!("Same padding with an even kernel size is not supported"); + } + } +} diff --git a/crates/burn-nn/src/modules/conv/conv1d.rs b/crates/burn-nn/src/modules/conv/conv1d.rs new file mode 100644 index 0000000..bbbcf02 --- /dev/null +++ b/crates/burn-nn/src/modules/conv/conv1d.rs @@ -0,0 +1,295 @@ +use alloc::format; + +use burn_core as burn; + +use crate::{PaddingConfig1d, conv::checks}; +use burn::tensor::{Device, Tensor, module::conv1d, ops::PaddedConvOptions}; +use burn::{ + config::Config, + module::{Content, DisplaySettings, Initializer, Module, ModuleDisplay, Param}, +}; + +/// Configuration to create a [1D convolution](Conv1d) layer using the [init function](Conv1dConfig::init). +#[derive(Config, Debug)] +pub struct Conv1dConfig { + /// The number of input channels. + pub channels_in: usize, + /// The number of output channels. + pub channels_out: usize, + /// The size of the kernel. + pub kernel_size: usize, + /// The stride of the convolution. + #[config(default = "1")] + pub stride: usize, + /// Spacing between kernel elements. + #[config(default = "1")] + pub dilation: usize, + /// Controls the connections between input and output channels. + #[config(default = "1")] + pub groups: usize, + /// The padding configuration. + /// + /// Supports symmetric and asymmetric padding. `Same` padding with even kernel sizes + /// will automatically use asymmetric padding to preserve input dimensions. + #[config(default = "PaddingConfig1d::Valid")] + pub padding: PaddingConfig1d, + /// If bias should be added to the output. + #[config(default = true)] + pub bias: bool, + /// The type of function used to initialize neural network parameters + #[config( + default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0),fan_out_only:false}" + )] + pub initializer: Initializer, +} + +/// Applies a 1D convolution over input tensors. +/// +/// Should be created with [Conv1dConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Conv1d { + /// Tensor of shape `[channels_out, channels_in / groups, kernel_size]` + pub weight: Param>, + /// Tensor of shape `[channels_out]` + pub bias: Option>>, + /// Stride of the convolution. + pub stride: usize, + /// Size of the kernel. + pub kernel_size: usize, + /// Spacing between kernel elements. + pub dilation: usize, + /// Controls the connections between input and output channels. + pub groups: usize, + /// Padding configuration. + #[module(skip)] + pub padding: PaddingConfig1d, +} + +impl ModuleDisplay for Conv1d { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + // Format stride/dilation as strings + let stride = format!("{:?}", self.stride); + let kernel_size = format!("{:?}", self.kernel_size); + let dilation = format!("{:?}", self.dilation); + + // Extract channels in/out from weight dims + let [channels_out, group_channels_in, _] = self.weight.dims(); + let channels_in = group_channels_in * self.groups; + let ch_out = format!("{:?}", channels_out); + let ch_in = format!("{:?}", channels_in); + + content + .add("ch_in", &ch_in) + .add("ch_out", &ch_out) + .add("stride", &stride) + .add("kernel_size", &kernel_size) + .add("dilation", &dilation) + .add("groups", &self.groups) + .add_debug_attribute("padding", &self.padding) + .optional() + } +} +impl Conv1dConfig { + /// Initialize a new [conv1d](Conv1d) module. + pub fn init(&self, device: &Device) -> Conv1d { + checks::checks_channels_div_groups(self.channels_in, self.channels_out, self.groups); + + let shape = [ + self.channels_out, + self.channels_in / self.groups, + self.kernel_size, + ]; + + let fan_in: usize = self.channels_in / self.groups * self.kernel_size; + let fan_out = self.channels_out / self.groups * self.kernel_size; + + let weight = self + .initializer + .init_with(shape, Some(fan_in), Some(fan_out), device); + let mut bias = None; + + if self.bias { + bias = + Some( + self.initializer + .init_with([self.channels_out], Some(fan_in), None, device), + ); + } + + Conv1d { + weight, + bias, + stride: self.stride, + kernel_size: self.kernel_size, + padding: self.padding.clone(), + dilation: self.dilation, + groups: self.groups, + } + } +} + +impl Conv1d { + /// Applies the forward pass on the input tensor. + /// + /// See [conv1d](burn::tensor::module::conv1d) for more information. + /// + /// # Shapes + /// + /// - input: `[batch_size, channels_in, length_in]` + /// - output: `[batch_size, channels_out, length_out]` + pub fn forward(&self, input: Tensor<3>) -> Tensor<3> { + let length = input.dims()[2]; + + // Calculate padding as pair - handles Same, Valid, and Explicit uniformly + let (left, right) = + self.padding + .calculate_padding_1d_pair(length, self.kernel_size, self.stride); + + let options = PaddedConvOptions::asymmetric( + [self.stride], + [left], + [right], + [self.dilation], + self.groups, + ); + + conv1d( + input, + self.weight.val(), + self.bias.as_ref().map(|bias| bias.val()), + options, + ) + } +} + +#[cfg(test)] +mod tests { + use burn::tensor::ElementConversion; + type FT = f32; + + use super::*; + use burn::tensor::TensorData; + + #[test] + fn initializer_default() { + let device = Device::default(); + device.seed(0); + + let config = Conv1dConfig::new(5, 5, 5); + let k = (config.channels_in * config.kernel_size) as f64; + let k = (config.groups as f64 / k).sqrt().elem::(); + let conv = config.init(&device); + + conv.weight.to_data().assert_within_range(-k..k); + } + + #[test] + fn initializer_zeros() { + let device = Device::default(); + device.seed(0); + + let config = Conv1dConfig::new(5, 5, 5).with_initializer(Initializer::Zeros); + let conv = config.init(&device); + + assert_eq!(config.initializer, Initializer::Zeros); + conv.weight + .to_data() + .assert_eq(&TensorData::zeros::(conv.weight.shape()), false); + } + + #[test] + fn same_with_even_kernel_uses_asymmetric_padding() { + let device = Default::default(); + let config = Conv1dConfig::new(4, 4, 2) + .with_padding(PaddingConfig1d::Same) + .with_initializer(Initializer::Constant { value: 1.0 }) + .with_bias(false); + let conv = config.init(&device); + + // Input: [batch=1, channels=4, length=5] + let input = Tensor::<3>::ones([1, 4, 5], &device); + let output = conv.forward(input); + + // Same padding should preserve spatial dimensions + assert_eq!(output.dims(), [1, 4, 5]); + } + + #[test] + fn display() { + let config = Conv1dConfig::new(5, 5, 5); + let conv = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{conv}"), + "Conv1d {ch_in: 5, ch_out: 5, stride: 1, kernel_size: 5, dilation: 1, groups: 1, padding: Valid, params: 130}" + ); + } + + #[test] + #[should_panic = "Number of channels in input tensor and input channels of convolution must be equal. got: 4, expected: 5"] + fn input_channels_mismatch() { + let config = Conv1dConfig::new(5, 3, 3); + let conv = config.init(&Default::default()); + + let input = Tensor::<3>::zeros([1, 4, 10], &Default::default()); + let _ = conv.forward(input); + } + + #[test] + fn asymmetric_padding_forward() { + let device = Default::default(); + // Create conv with asymmetric padding: left=1, right=2 + let config = Conv1dConfig::new(2, 3, 3) + .with_padding(PaddingConfig1d::Explicit(1, 2)) + .with_initializer(Initializer::Constant { value: 1.0 }) + .with_bias(false); + let conv = config.init(&device); + + // Input: [batch=1, channels=2, length=4] + let input = Tensor::<3>::ones([1, 2, 4], &device); + let output = conv.forward(input); + + // With asymmetric padding (1, 2), input length 4 becomes 4+1+2=7 + // Output length = (7 - 3) / 1 + 1 = 5 + assert_eq!(output.dims(), [1, 3, 5]); + } + + #[test] + fn symmetric_explicit_padding_forward() { + let device = Default::default(); + // Create conv with symmetric explicit padding: left=2, right=2 + let config = Conv1dConfig::new(2, 3, 3) + .with_padding(PaddingConfig1d::Explicit(2, 2)) + .with_initializer(Initializer::Constant { value: 1.0 }) + .with_bias(false); + let conv = config.init(&device); + + // Input: [batch=1, channels=2, length=4] + let input = Tensor::<3>::ones([1, 2, 4], &device); + let output = conv.forward(input); + + // With symmetric padding (2, 2), input length 4 becomes 4+2+2=8 + // Output length = (8 - 3) / 1 + 1 = 6 + assert_eq!(output.dims(), [1, 3, 6]); + } + + #[test] + fn initializer_fan_out() { + let device = Device::default(); + device.seed(0); + + let init = Initializer::XavierUniform { gain: 1.0 }; + + let config = Conv1dConfig::new(2, 2, 1).with_initializer(init.clone()); + let c = config.init(&device); + + let _ = c.weight.val(); + } +} diff --git a/crates/burn-nn/src/modules/conv/conv2d.rs b/crates/burn-nn/src/modules/conv/conv2d.rs new file mode 100644 index 0000000..661c36e --- /dev/null +++ b/crates/burn-nn/src/modules/conv/conv2d.rs @@ -0,0 +1,346 @@ +use alloc::format; + +use burn_core as burn; + +use crate::PaddingConfig2d; +use burn::config::Config; +use burn::module::Initializer; +use burn::module::{Content, DisplaySettings, Module, ModuleDisplay, Param}; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn::tensor::module::conv2d; +use burn::tensor::ops::PaddedConvOptions; + +use crate::conv::checks; + +/// Configuration to create a [2D convolution](Conv2d) layer, using the [init function](Conv2dConfig::init). +#[derive(Config, Debug)] +pub struct Conv2dConfig { + /// The number of channels. + pub channels: [usize; 2], + /// The size of the kernel. + pub kernel_size: [usize; 2], + /// The stride of the convolution. + #[config(default = "[1, 1]")] + pub stride: [usize; 2], + /// Spacing between kernel elements. + #[config(default = "[1, 1]")] + pub dilation: [usize; 2], + /// Controls the connections between input and output channels. + #[config(default = "1")] + pub groups: usize, + /// The padding configuration. + /// + /// Supports symmetric and asymmetric padding. `Same` padding with even kernel sizes + /// will automatically use asymmetric padding to preserve input dimensions. + #[config(default = "PaddingConfig2d::Valid")] + pub padding: PaddingConfig2d, + /// If bias should be added to the output. + #[config(default = true)] + pub bias: bool, + /// The type of function used to initialize neural network parameters + #[config( + default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0),fan_out_only:false}" + )] + pub initializer: Initializer, +} + +/// Applies a 2D convolution over input tensors. +/// +/// Should be created with [Conv2dConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Conv2d { + /// Tensor of shape `[channels_out, channels_in / groups, kernel_size_1, kernel_size_2]` + pub weight: Param>, + /// Tensor of shape `[channels_out]` + pub bias: Option>>, + /// Stride of the convolution. + pub stride: [usize; 2], + /// Size of the kernel. + pub kernel_size: [usize; 2], + /// Spacing between kernel elements. + pub dilation: [usize; 2], + /// Controls the connections between input and output channels. + pub groups: usize, + /// The padding configuration. + #[module(skip)] + pub padding: PaddingConfig2d, +} + +impl Conv2dConfig { + /// Initialize a new [conv2d](Conv2d) module. + pub fn init(&self, device: &Device) -> Conv2d { + checks::checks_channels_div_groups(self.channels[0], self.channels[1], self.groups); + + let shape = [ + self.channels[1], + self.channels[0] / self.groups, + self.kernel_size[0], + self.kernel_size[1], + ]; + + let k = self.kernel_size.iter().product::(); + let fan_in = self.channels[0] / self.groups * k; + let fan_out = self.channels[1] / self.groups * k; + + let weight = self + .initializer + .init_with(shape, Some(fan_in), Some(fan_out), device); + let mut bias = None; + + if self.bias { + bias = Some(self.initializer.init_with( + [self.channels[1]], + Some(fan_in), + Some(fan_out), + device, + )); + } + + Conv2d { + weight, + bias, + stride: self.stride, + kernel_size: self.kernel_size, + dilation: self.dilation, + padding: self.padding.clone(), + groups: self.groups, + } + } +} + +impl ModuleDisplay for Conv2d { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + // Format the stride, kernel_size and dilation as strings, formatted as arrays instead of indexed. + let stride = format!("{:?}", self.stride); + let kernel_size = format!("{:?}", self.kernel_size); + let dilation = format!("{:?}", self.dilation); + let [channels_out, group_channels_in, _, _] = self.weight.dims(); + let channels_in = group_channels_in * self.groups; + let ch_out = format!("{:?}", channels_out); + let ch_in = format!("{:?}", channels_in); + content + .add("ch_in", &ch_in) + .add("ch_out", &ch_out) + .add("stride", &stride) + .add("kernel_size", &kernel_size) + .add("dilation", &dilation) + .add("groups", &self.groups) + .add_debug_attribute("padding", &self.padding) + .optional() + } +} + +impl Conv2d { + /// Applies the forward pass on the input tensor. + /// + /// See [conv2d](burn::tensor::module::conv2d) for more information. + /// + /// # Shapes + /// - `input`: `[batch_size, channels_in, height_in, width_in]` + /// - `output`: `[batch_size, channels_out, height_out, width_out]` + /// + /// # Example + /// ```rust,ignore + /// use burn::nn::conv::Conv2dConfig; + /// use burn::tensor::Tensor; + /// + /// // Assuming backend type alias `B` + /// let device = Default::default(); + /// let conv = Conv2dConfig::new([3, 8], [3, 3]).init(&device); + /// + /// let x = Tensor::::zeros([1, 3, 28, 28], &device); + /// let y = conv.forward(x); + /// + /// println!("{:?}", y.dims()); // [1, 8, 26, 26] + /// ``` + pub fn forward(&self, input: Tensor<4>) -> Tensor<4> { + let [_batch_size, _channels_in, height_in, width_in] = input.dims(); + + // Calculate padding as pairs - handles Same, Valid, and Explicit uniformly + let ((top, bottom), (left, right)) = self.padding.calculate_padding_2d_pairs( + height_in, + width_in, + &self.kernel_size, + &self.stride, + ); + + let options = PaddedConvOptions::asymmetric( + self.stride, + [top, left], + [bottom, right], + self.dilation, + self.groups, + ); + + conv2d( + input, + self.weight.val(), + self.bias.as_ref().map(|bias| bias.val()), + options, + ) + } +} + +#[cfg(test)] +mod tests { + use burn::tensor::{ElementConversion, Tolerance}; + + use super::*; + use burn::tensor::TensorData; + type FT = f32; + + #[test] + fn initializer_default() { + let device = Device::default(); + device.seed(0); + + let config = Conv2dConfig::new([5, 1], [5, 5]); + let k = (config.channels[0] * config.kernel_size[0] * config.kernel_size[1]) as f64; + let k = (config.groups as f64 / k).sqrt().elem::(); + let conv = config.init(&device); + + conv.weight.to_data().assert_within_range(-k..k); + } + + #[test] + fn initializer_zeros() { + let device = Device::default(); + device.seed(0); + + let config = Conv2dConfig::new([5, 2], [5, 5]).with_initializer(Initializer::Zeros); + let conv = config.init(&device); + + assert_eq!(config.initializer, Initializer::Zeros); + conv.weight.to_data().assert_approx_eq::( + &TensorData::zeros::(conv.weight.shape()), + Tolerance::default(), + ); + } + + #[test] + fn initializer_fan_out() { + let device = Device::default(); + device.seed(0); + + let init = Initializer::KaimingUniform { + gain: 1.0 / 3.0f64.sqrt(), + fan_out_only: true, // test that fan_out is passed to `init_with()` + }; + + let config = Conv2dConfig::new([5, 1], [5, 5]).with_initializer(init.clone()); + let c = config.init(&device); + let _ = c.weight.val(); // initializes param + + assert_eq!(config.initializer, init); + } + + #[test] + fn initializer_fan_with_groups_is_valid() { + let device = Device::default(); + device.seed(0); + + let init = Initializer::KaimingUniform { + gain: 1.0 / 3.0f64.sqrt(), + fan_out_only: true, + }; + + let config = Conv2dConfig::new([4, 4], [1, 1]) + .with_initializer(init.clone()) + .with_groups(4); + let _ = config.init(&device); + + assert_eq!(config.initializer, init); + } + + #[test] + #[should_panic = "Both channels must be divisible by the number of groups."] + fn channels_with_groups_is_invalid() { + let device = Default::default(); + let config = Conv2dConfig::new([1, 4], [1, 1]).with_groups(4); + let _ = config.init(&device); + } + + #[test] + fn same_with_even_kernel_uses_asymmetric_padding() { + let device = Default::default(); + let config = Conv2dConfig::new([4, 4], [2, 2]) + .with_padding(PaddingConfig2d::Same) + .with_initializer(Initializer::Constant { value: 1.0 }) + .with_bias(false); + let conv = config.init(&device); + + // Input: [batch=1, channels=4, height=5, width=5] + let input = Tensor::<4>::ones([1, 4, 5, 5], &device); + let output = conv.forward(input); + + // Same padding should preserve spatial dimensions + assert_eq!(output.dims(), [1, 4, 5, 5]); + } + + #[test] + fn display() { + let config = Conv2dConfig::new([5, 1], [5, 5]); + let conv = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{conv}"), + "Conv2d {ch_in: 5, ch_out: 1, stride: [1, 1], kernel_size: [5, 5], dilation: [1, 1], groups: 1, padding: Valid, params: 126}" + ); + } + + #[test] + #[should_panic = "Number of channels in input tensor and input channels of convolution must be equal. got: 4, expected: 5"] + fn input_channels_mismatch() { + let config = Conv2dConfig::new([5, 3], [3, 3]); + let conv = config.init(&Default::default()); + + let input = Tensor::<4>::zeros([1, 4, 10, 10], &Default::default()); + let _ = conv.forward(input); + } + + #[test] + fn asymmetric_padding_forward() { + let device = Default::default(); + // Create conv with asymmetric padding: top=1, left=2, bottom=3, right=4 + let config = Conv2dConfig::new([2, 3], [3, 3]) + .with_padding(PaddingConfig2d::Explicit(1, 2, 3, 4)) + .with_initializer(Initializer::Constant { value: 1.0 }) + .with_bias(false); + let conv = config.init(&device); + + // Input: [batch=1, channels=2, height=4, width=5] + let input = Tensor::<4>::ones([1, 2, 4, 5], &device); + let output = conv.forward(input); + + // Height: 4 + 1 + 3 = 8, output = (8 - 3) / 1 + 1 = 6 + // Width: 5 + 2 + 4 = 11, output = (11 - 3) / 1 + 1 = 9 + assert_eq!(output.dims(), [1, 3, 6, 9]); + } + + #[test] + fn symmetric_explicit_padding_forward() { + let device = Default::default(); + // Create conv with symmetric explicit padding: top=2, left=2, bottom=2, right=2 + let config = Conv2dConfig::new([2, 3], [3, 3]) + .with_padding(PaddingConfig2d::Explicit(2, 2, 2, 2)) + .with_initializer(Initializer::Constant { value: 1.0 }) + .with_bias(false); + let conv = config.init(&device); + + // Input: [batch=1, channels=2, height=4, width=5] + let input = Tensor::<4>::ones([1, 2, 4, 5], &device); + let output = conv.forward(input); + + // Height: 4 + 2 + 2 = 8, output = (8 - 3) / 1 + 1 = 6 + // Width: 5 + 2 + 2 = 9, output = (9 - 3) / 1 + 1 = 7 + assert_eq!(output.dims(), [1, 3, 6, 7]); + } +} diff --git a/crates/burn-nn/src/modules/conv/conv3d.rs b/crates/burn-nn/src/modules/conv/conv3d.rs new file mode 100644 index 0000000..cf540b4 --- /dev/null +++ b/crates/burn-nn/src/modules/conv/conv3d.rs @@ -0,0 +1,273 @@ +use alloc::format; + +use burn_core as burn; + +use crate::PaddingConfig3d; +use burn::config::Config; +use burn::module::Initializer; +use burn::module::{Content, DisplaySettings, Module, ModuleDisplay, Param}; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn::tensor::module::conv3d; +use burn::tensor::ops::ConvOptions; + +use crate::conv::checks; + +/// Configuration to create a [3D convolution](Conv3d) layer, using the [init function](Conv3dConfig::init). +#[derive(Config, Debug)] +pub struct Conv3dConfig { + /// The number of channels. + pub channels: [usize; 2], + /// The size of the kernel. + pub kernel_size: [usize; 3], + /// The stride of the convolution. + #[config(default = "[1, 1, 1]")] + pub stride: [usize; 3], + /// Spacing between kernel elements. + #[config(default = "[1, 1, 1]")] + pub dilation: [usize; 3], + /// Controls the connections between input and output channels. + #[config(default = "1")] + pub groups: usize, + /// The padding configuration. + #[config(default = "PaddingConfig3d::Valid")] + pub padding: PaddingConfig3d, + /// If bias should be added to the output. + #[config(default = true)] + pub bias: bool, + /// The type of function used to initialize neural network parameters + #[config( + default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0),fan_out_only:false}" + )] + pub initializer: Initializer, +} + +/// Applies a 3D convolution over input tensors. +/// +/// Should be created with [Conv3dConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Conv3d { + /// Tensor of shape `[channels_out, channels_in / groups, kernel_size_1, kernel_size_2, kernel_size_3]` + pub weight: Param>, + /// Tensor of shape `[channels_out]` + pub bias: Option>>, + /// Stride of the convolution. + pub stride: [usize; 3], + /// Size of the kernel. + pub kernel_size: [usize; 3], + /// Spacing between kernel elements. + pub dilation: [usize; 3], + /// Controls the connections between input and output channels. + pub groups: usize, + /// The padding configuration. + #[module(skip)] + pub padding: PaddingConfig3d, +} + +impl Conv3dConfig { + /// Initialize a new [conv3d](Conv3d) module. + pub fn init(&self, device: &Device) -> Conv3d { + checks::checks_channels_div_groups(self.channels[0], self.channels[1], self.groups); + if self.padding == PaddingConfig3d::Same { + checks::check_same_padding_support(&self.kernel_size); + } + + let shape = [ + self.channels[1], + self.channels[0] / self.groups, + self.kernel_size[0], + self.kernel_size[1], + self.kernel_size[2], + ]; + + let k = self.kernel_size.iter().product::(); + let fan_in = self.channels[0] / self.groups * k; + let fan_out = self.channels[1] / self.groups * k; + + let weight = self + .initializer + .init_with(shape, Some(fan_in), Some(fan_out), device); + let mut bias = None; + + if self.bias { + bias = Some(self.initializer.init_with( + [self.channels[1]], + Some(fan_in), + Some(fan_out), + device, + )); + } + + Conv3d { + weight, + bias, + stride: self.stride, + kernel_size: self.kernel_size, + dilation: self.dilation, + padding: self.padding.clone(), + groups: self.groups, + } + } +} + +impl ModuleDisplay for Conv3d { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + // Format arrays as strings (consistent with Conv2d/Conv1d). + let stride = format!("{:?}", self.stride); + let kernel_size = format!("{:?}", self.kernel_size); + let dilation = format!("{:?}", self.dilation); + + // Weight dims: [channels_out, channels_in/groups, k1, k2, k3] + let [channels_out, group_channels_in, _, _, _] = self.weight.dims(); + let channels_in = group_channels_in * self.groups; + let ch_out = format!("{:?}", channels_out); + let ch_in = format!("{:?}", channels_in); + + content + .add("ch_in", &ch_in) + .add("ch_out", &ch_out) + .add("stride", &stride) + .add("kernel_size", &kernel_size) + .add("dilation", &dilation) + .add("groups", &self.groups) + .add_debug_attribute("padding", &self.padding) + .optional() + } +} + +impl Conv3d { + /// Applies the forward pass on the input tensor. + /// + /// See [conv3d](burn::tensor::module::conv3d) for more information. + /// + /// # Shapes + /// + /// - input: `[batch_size, channels_in, depth_in, height_in, width_in]` + /// - output: `[batch_size, channels_out, depth_out, height_out, width_out]` + pub fn forward(&self, input: Tensor<5>) -> Tensor<5> { + let [_batch_size, _channels_in, depth_in, height_in, width_in] = input.dims(); + let padding = self.padding.calculate_padding_3d( + depth_in, + height_in, + width_in, + &self.kernel_size, + &self.stride, + ); + conv3d( + input, + self.weight.val(), + self.bias.as_ref().map(|bias| bias.val()), + ConvOptions::new(self.stride, padding, self.dilation, self.groups), + ) + } +} + +#[cfg(test)] +mod tests { + use burn::tensor::{ElementConversion, Tolerance}; + type FT = f32; + + use super::*; + use burn::tensor::TensorData; + + #[test] + fn initializer_default() { + let device = Device::default(); + device.seed(0); + + let config = Conv3dConfig::new([5, 1], [5, 5, 5]); + let k = (config.channels[0] + * config.kernel_size[0] + * config.kernel_size[1] + * config.kernel_size[2]) as f64; + let k = (config.groups as f64 / k).sqrt().elem::(); + let conv = config.init(&device); + + conv.weight.to_data().assert_within_range(-k..k); + } + + #[test] + fn initializer_zeros() { + let device = Device::default(); + device.seed(0); + + let config = Conv3dConfig::new([5, 2], [5, 5, 5]).with_initializer(Initializer::Zeros); + let conv = config.init(&device); + + assert_eq!(config.initializer, Initializer::Zeros); + conv.weight.to_data().assert_approx_eq::( + &TensorData::zeros::(conv.weight.shape()), + Tolerance::default(), + ); + } + + #[test] + fn initializer_fan_out() { + let device = Device::default(); + device.seed(0); + + let init = Initializer::KaimingUniform { + gain: 1.0 / 3.0f64.sqrt(), + fan_out_only: true, // test that fan_out is passed to `init_with()` + }; + let config = Conv3dConfig::new([5, 1], [5, 5, 5]).with_initializer(init.clone()); + let c = config.init(&device); + let _ = c.weight.val(); // initializes param + + assert_eq!(config.initializer, init); + } + + #[test] + fn initializer_fan_with_groups_is_valid() { + let device = Device::default(); + device.seed(0); + + let init = Initializer::KaimingUniform { + gain: 1.0 / 3.0f64.sqrt(), + fan_out_only: true, + }; + + let config = Conv3dConfig::new([4, 4], [1, 1, 1]) + .with_initializer(init.clone()) + .with_groups(4); + let _ = config.init(&device); + + assert_eq!(config.initializer, init); + } + + #[test] + #[should_panic = "Same padding with an even kernel size is not supported"] + fn same_with_even_kernel_is_invalid() { + let device = Default::default(); + let config = Conv3dConfig::new([4, 4], [2, 2, 2]).with_padding(PaddingConfig3d::Same); + let _ = config.init(&device); + } + + #[test] + fn display() { + let config = Conv3dConfig::new([5, 1], [5, 5, 5]); + let conv = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{conv}"), + "Conv3d {ch_in: 5, ch_out: 1, stride: [1, 1, 1], kernel_size: [5, 5, 5], dilation: [1, 1, 1], groups: 1, padding: Valid, params: 626}" + ); + } + + #[test] + #[should_panic = "Number of channels in input tensor and input channels of convolution must be equal. got: 4, expected: 5"] + fn input_channels_mismatch() { + let config = Conv3dConfig::new([5, 3], [3, 3, 3]); + let conv = config.init(&Default::default()); + + let input = Tensor::<5>::zeros([1, 4, 10, 10, 10], &Default::default()); + let _ = conv.forward(input); + } +} diff --git a/crates/burn-nn/src/modules/conv/conv_transpose1d.rs b/crates/burn-nn/src/modules/conv/conv_transpose1d.rs new file mode 100644 index 0000000..037f954 --- /dev/null +++ b/crates/burn-nn/src/modules/conv/conv_transpose1d.rs @@ -0,0 +1,214 @@ +use alloc::format; + +use burn_core as burn; + +use crate::conv::checks; +use burn::config::Config; +use burn::module::Content; +use burn::module::DisplaySettings; +use burn::module::Initializer; +use burn::module::Module; +use burn::module::ModuleDisplay; +use burn::module::Param; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn::tensor::module::conv_transpose1d; +use burn::tensor::ops::ConvTransposeOptions; + +/// Configuration to create an [1D transposed convolution](ConvTranspose1d) layer +/// using the [init function](ConvTranspose1dConfig::init). +#[derive(Config, Debug)] +pub struct ConvTranspose1dConfig { + /// The number of channels. + pub channels: [usize; 2], + /// The size of the kernel. + pub kernel_size: usize, + /// The stride of the convolution. + #[config(default = "1")] + pub stride: usize, + /// Spacing between kernel elements. + #[config(default = "1")] + pub dilation: usize, + /// Controls the connections between input and output channels. + #[config(default = "1")] + pub groups: usize, + /// The padding configuration. + #[config(default = "0")] + pub padding: usize, + /// The padding output configuration. + #[config(default = "0")] + pub padding_out: usize, + /// If bias should be added to the output. + #[config(default = true)] + pub bias: bool, + /// The type of function used to initialize neural network parameters + #[config( + default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0),fan_out_only:false}" + )] + pub initializer: Initializer, +} + +/// Applies a 1D transposed convolution over input tensors. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct ConvTranspose1d { + /// Tensor of shape `[channels_in, channels_out / groups, kernel_size]` + pub weight: Param>, + /// Tensor of shape `[channels_out]` + pub bias: Option>>, + /// Stride of the convolution. + pub stride: usize, + /// Size of the kernel. + pub kernel_size: usize, + /// Spacing between kernel elements. + pub dilation: usize, + /// Controls the connections between input and output channels. + pub groups: usize, + /// The padding configuration. + pub padding: usize, + /// The padding output configuration. + pub padding_out: usize, + /// The number of channels. + pub channels: [usize; 2], +} + +impl ModuleDisplay for ConvTranspose1d { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("channels", &format!("{:?}", self.channels)) + .add("stride", &self.stride) + .add("kernel_size", &self.kernel_size) + .add("dilation", &self.dilation) + .add("groups", &self.groups) + .add("padding", &self.padding) + .add("padding_out", &self.padding_out) + .optional() + } +} + +impl ConvTranspose1dConfig { + /// Initialize a new [conv transpose 1d](ConvTranspose1d) module. + pub fn init(&self, device: &Device) -> ConvTranspose1d { + checks::checks_channels_div_groups(self.channels[0], self.channels[1], self.groups); + + let shape = [ + self.channels[0], + self.channels[1] / self.groups, + self.kernel_size, + ]; + + let fan_in = self.channels[1] / self.groups * self.kernel_size; + let weight = self + .initializer + .init_with(shape, Some(fan_in), None, device); + let mut bias = None; + + if self.bias { + bias = Some( + self.initializer + .init_with([self.channels[1]], Some(fan_in), None, device), + ); + } + + ConvTranspose1d { + weight, + bias, + stride: self.stride, + kernel_size: self.kernel_size, + dilation: self.dilation, + groups: self.groups, + padding: self.padding, + padding_out: self.padding_out, + channels: self.channels, + } + } +} + +impl ConvTranspose1d { + /// Applies the forward pass on the input tensor. + /// + /// See also [conv_transpose1d](burn::tensor::module::conv_transpose1d). + /// + /// # Shapes + /// + /// - input: `[batch_size, channels_in, length_in]` + /// - output: `[batch_size, channels_out, length_out]` + pub fn forward(&self, input: Tensor<3>) -> Tensor<3> { + conv_transpose1d( + input, + self.weight.val(), + self.bias.as_ref().map(|bias| bias.val()), + ConvTransposeOptions::new( + [self.stride], + [self.padding], + [self.padding_out], + [self.dilation], + self.groups, + ), + ) + } +} + +#[cfg(test)] +mod tests { + use burn::tensor::{ElementConversion, Tolerance}; + + use super::*; + use burn::tensor::TensorData; + type FT = f32; + + #[test] + fn initializer_default() { + let device = Device::default(); + device.seed(0); + + let config = ConvTranspose1dConfig::new([5, 1], 5); + let k = (config.channels[1] * config.kernel_size) as f64; + let k = (config.groups as f64 / k).sqrt().elem::(); + let conv = config.init(&device); + + conv.weight.to_data().assert_within_range(-k..k); + } + + #[test] + fn initializer_zeros() { + let device = Device::default(); + device.seed(0); + + let config = ConvTranspose1dConfig::new([5, 2], 5).with_initializer(Initializer::Zeros); + let conv = config.init(&device); + + assert_eq!(config.initializer, Initializer::Zeros); + conv.weight.to_data().assert_approx_eq::( + &TensorData::zeros::(conv.weight.shape()), + Tolerance::default(), + ); + } + + #[test] + fn display() { + let config = ConvTranspose1dConfig::new([5, 2], 5); + let conv = config.init(&Default::default()); + + assert_eq!( + format!("{conv}"), + "ConvTranspose1d {channels: [5, 2], stride: 1, kernel_size: 5, dilation: 1, groups: 1, padding: 0, padding_out: 0, params: 52}" + ); + } + + #[test] + #[should_panic = "Number of channels in input tensor and input channels of convolution must be equal. got: 4, expected: 5"] + fn input_channels_mismatch() { + let config = ConvTranspose1dConfig::new([5, 3], 3); + let conv = config.init(&Default::default()); + + let input = Tensor::<3>::zeros([1, 4, 10], &Default::default()); + let _ = conv.forward(input); + } +} diff --git a/crates/burn-nn/src/modules/conv/conv_transpose2d.rs b/crates/burn-nn/src/modules/conv/conv_transpose2d.rs new file mode 100644 index 0000000..bd92476 --- /dev/null +++ b/crates/burn-nn/src/modules/conv/conv_transpose2d.rs @@ -0,0 +1,215 @@ +use alloc::format; + +use burn_core as burn; + +use crate::conv::checks; +use burn::config::Config; +use burn::module::Content; +use burn::module::DisplaySettings; +use burn::module::Initializer; +use burn::module::Module; +use burn::module::ModuleDisplay; +use burn::module::Param; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn::tensor::module::conv_transpose2d; +use burn::tensor::ops::ConvTransposeOptions; + +/// Configuration to create an [2D transposed convolution](ConvTranspose2d) layer +/// using the [init function](ConvTranspose2dConfig::init). +#[derive(Config, Debug)] +pub struct ConvTranspose2dConfig { + /// The number of channels. + pub channels: [usize; 2], + /// The size of the kernel. + pub kernel_size: [usize; 2], + /// The stride of the convolution. + #[config(default = "[1, 1]")] + pub stride: [usize; 2], + /// Spacing between kernel elements. + #[config(default = "[1, 1]")] + pub dilation: [usize; 2], + /// Controls the connections between input and output channels. + #[config(default = "1")] + pub groups: usize, + /// The padding configuration. + #[config(default = "[0, 0]")] + pub padding: [usize; 2], + /// The padding output configuration. + #[config(default = "[0, 0]")] + pub padding_out: [usize; 2], + /// If bias should be added to the output. + #[config(default = true)] + pub bias: bool, + /// The type of function used to initialize neural network parameters + #[config( + default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0),fan_out_only:false}" + )] + pub initializer: Initializer, +} + +/// Applies a 2D transposed convolution over input tensors. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct ConvTranspose2d { + /// Tensor of shape `[channels_in, channels_out / groups, kernel_size_1, kernel_size_2]` + pub weight: Param>, + /// Tensor of shape `[channels_out]` + pub bias: Option>>, + /// Stride of the convolution. + pub stride: [usize; 2], + /// Size of the kernel. + pub kernel_size: [usize; 2], + /// Spacing between kernel elements. + pub dilation: [usize; 2], + /// Controls the connections between input and output channels. + pub groups: usize, + /// Padding configuration. + pub padding: [usize; 2], + /// Padding output configuration. + pub padding_out: [usize; 2], + /// Number of channels. + pub channels: [usize; 2], +} + +impl ModuleDisplay for ConvTranspose2d { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("channels", &format!("{:?}", self.channels)) + .add("stride", &format!("{:?}", self.stride)) + .add("kernel_size", &format!("{:?}", self.kernel_size)) + .add("dilation", &format!("{:?}", self.dilation)) + .add("groups", &self.groups) + .add("padding", &format!("{:?}", self.padding)) + .add("padding_out", &format!("{:?}", self.padding_out)) + .optional() + } +} + +impl ConvTranspose2dConfig { + /// Initialize a new [conv transpose 2d](ConvTranspose2d) module. + pub fn init(&self, device: &Device) -> ConvTranspose2d { + checks::checks_channels_div_groups(self.channels[0], self.channels[1], self.groups); + + let shape = [ + self.channels[0], + self.channels[1] / self.groups, + self.kernel_size[0], + self.kernel_size[1], + ]; + + let fan_in = self.channels[1] / self.groups * self.kernel_size.iter().product::(); + let weight = self + .initializer + .init_with(shape, Some(fan_in), None, device); + let mut bias = None; + + if self.bias { + bias = Some( + self.initializer + .init_with([self.channels[1]], Some(fan_in), None, device), + ); + } + + ConvTranspose2d { + weight, + bias, + stride: self.stride, + kernel_size: self.kernel_size, + dilation: self.dilation, + groups: self.groups, + padding: self.padding, + padding_out: self.padding_out, + channels: self.channels, + } + } +} + +impl ConvTranspose2d { + /// Applies the forward pass on the input tensor. + /// + /// See also [conv_transpose2d](burn::tensor::module::conv_transpose2d). + /// + /// # Shapes + /// + /// - input: `[batch_size, channels_in, height_in, width_in]` + /// - output: `[batch_size, channels_out, height_out, width_out]` + pub fn forward(&self, input: Tensor<4>) -> Tensor<4> { + conv_transpose2d( + input, + self.weight.val(), + self.bias.as_ref().map(|bias| bias.val()), + ConvTransposeOptions::new( + self.stride, + self.padding, + self.padding_out, + self.dilation, + self.groups, + ), + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + use burn::tensor::{ElementConversion, Tolerance}; + type FT = f32; + + #[test] + fn initializer_default() { + let device = Device::default(); + device.seed(0); + + let config = ConvTranspose2dConfig::new([5, 1], [5, 5]); + let k = (config.channels[1] * config.kernel_size[0] * config.kernel_size[1]) as f64; + let k = (config.groups as f64 / k).sqrt().elem::(); + let conv = config.init(&device); + + conv.weight.to_data().assert_within_range(-k..k); + } + + #[test] + fn initializer_zeros() { + let device = Device::default(); + device.seed(0); + + let config = + ConvTranspose2dConfig::new([5, 2], [5, 5]).with_initializer(Initializer::Zeros); + let conv = config.init(&device); + + assert_eq!(config.initializer, Initializer::Zeros); + conv.weight.to_data().assert_approx_eq::( + &TensorData::zeros::(conv.weight.shape()), + Tolerance::default(), + ); + } + + #[test] + fn display() { + let config = ConvTranspose2dConfig::new([5, 2], [5, 5]); + let conv = config.init(&Default::default()); + + assert_eq!( + format!("{conv}"), + "ConvTranspose2d {channels: [5, 2], stride: [1, 1], kernel_size: [5, 5], dilation: [1, 1], groups: 1, padding: [0, 0], padding_out: [0, 0], params: 252}" + ); + } + + #[test] + #[should_panic = "Number of channels in input tensor and input channels of convolution must be equal. got: 4, expected: 5"] + fn input_channels_mismatch() { + let config = ConvTranspose2dConfig::new([5, 3], [3, 3]); + let conv = config.init(&Default::default()); + + let input = Tensor::<4>::zeros([1, 4, 10, 10], &Default::default()); + let _ = conv.forward(input); + } +} diff --git a/crates/burn-nn/src/modules/conv/conv_transpose3d.rs b/crates/burn-nn/src/modules/conv/conv_transpose3d.rs new file mode 100644 index 0000000..d9d0c37 --- /dev/null +++ b/crates/burn-nn/src/modules/conv/conv_transpose3d.rs @@ -0,0 +1,220 @@ +use alloc::format; + +use burn_core as burn; + +use crate::conv::checks; +use burn::config::Config; +use burn::module::Content; +use burn::module::DisplaySettings; +use burn::module::Initializer; +use burn::module::Module; +use burn::module::ModuleDisplay; +use burn::module::Param; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn::tensor::module::conv_transpose3d; +use burn::tensor::ops::ConvTransposeOptions; + +/// Configuration to create an [3D transposed convolution](ConvTranspose3d) layer +/// using the [init function](ConvTranspose3dConfig::init). +#[derive(Config, Debug)] +pub struct ConvTranspose3dConfig { + /// The number of channels. + pub channels: [usize; 2], + /// The size of the kernel. + pub kernel_size: [usize; 3], + /// The stride of the convolution. + #[config(default = "[1, 1, 1]")] + pub stride: [usize; 3], + /// Spacing between kernel elements. + #[config(default = "[1, 1, 1]")] + pub dilation: [usize; 3], + /// Controls the connections between input and output channels. + #[config(default = "1")] + pub groups: usize, + /// The padding configuration. + #[config(default = "[0, 0, 0]")] + pub padding: [usize; 3], + /// The padding output configuration. + #[config(default = "[0, 0, 0]")] + pub padding_out: [usize; 3], + /// If bias should be added to the output. + #[config(default = true)] + pub bias: bool, + /// The type of function used to initialize neural network parameters + #[config( + default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0),fan_out_only:false}" + )] + pub initializer: Initializer, +} + +/// Applies a 3D transposed convolution over input tensors. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct ConvTranspose3d { + /// Tensor of shape `[channels_in, channels_out / groups, kernel_size_1, kernel_size_2, kernel_size_3]` + pub weight: Param>, + /// Tensor of shape `[channels_out]` + pub bias: Option>>, + /// Stride of the convolution. + pub stride: [usize; 3], + /// Size of the kernel. + pub kernel_size: [usize; 3], + /// Spacing between kernel elements. + pub dilation: [usize; 3], + /// Controls the connections between input and output channels. + pub groups: usize, + /// Padding configuration. + pub padding: [usize; 3], + /// Padding output configuration. + pub padding_out: [usize; 3], + /// Number of channels. + pub channels: [usize; 2], +} + +impl ModuleDisplay for ConvTranspose3d { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("channels", &format!("{:?}", self.channels)) + .add("stride", &format!("{:?}", self.stride)) + .add("kernel_size", &format!("{:?}", self.kernel_size)) + .add("dilation", &format!("{:?}", self.dilation)) + .add("groups", &self.groups) + .add("padding", &format!("{:?}", self.padding)) + .add("padding_out", &format!("{:?}", self.padding_out)) + .optional() + } +} + +impl ConvTranspose3dConfig { + /// Initialize a new [conv transpose 2d](ConvTranspose3d) module. + pub fn init(&self, device: &Device) -> ConvTranspose3d { + checks::checks_channels_div_groups(self.channels[0], self.channels[1], self.groups); + + let shape = [ + self.channels[0], + self.channels[1] / self.groups, + self.kernel_size[0], + self.kernel_size[1], + self.kernel_size[2], + ]; + + let fan_in = self.channels[1] / self.groups * self.kernel_size.iter().product::(); + let weight = self + .initializer + .init_with(shape, Some(fan_in), None, device); + let mut bias = None; + + if self.bias { + bias = Some( + self.initializer + .init_with([self.channels[1]], Some(fan_in), None, device), + ); + } + + ConvTranspose3d { + weight, + bias, + stride: self.stride, + kernel_size: self.kernel_size, + dilation: self.dilation, + groups: self.groups, + padding: self.padding, + padding_out: self.padding_out, + channels: self.channels, + } + } +} + +impl ConvTranspose3d { + /// Applies the forward pass on the input tensor. + /// + /// See also [conv_transpose3d](burn::tensor::module::conv_transpose3d). + /// + /// # Shapes + /// + /// - input: `[batch_size, channels_in, depth_in, height_in, width_in]` + /// - output: `[batch_size, channels_out, depth_out, height_out, width_out]` + pub fn forward(&self, input: Tensor<5>) -> Tensor<5> { + conv_transpose3d( + input, + self.weight.val(), + self.bias.as_ref().map(|bias| bias.val()), + ConvTransposeOptions::new( + self.stride, + self.padding, + self.padding_out, + self.dilation, + self.groups, + ), + ) + } +} + +#[cfg(test)] +mod tests { + use burn::tensor::{ElementConversion, Tolerance}; + type FT = f32; + + use super::*; + use burn::tensor::TensorData; + + #[test] + fn initializer_default() { + let device = Device::default(); + device.seed(0); + + let config = ConvTranspose3dConfig::new([5, 1], [5, 5, 5]); + let k = (config.channels[1] + * config.kernel_size[0] + * config.kernel_size[1] + * config.kernel_size[2]) as f64; + let k = (config.groups as f64 / k).sqrt().elem::(); + let conv = config.init(&device); + + conv.weight.to_data().assert_within_range(-k..k); + } + + #[test] + fn initializer_zeros() { + let device = Device::default(); + device.seed(0); + + let config = + ConvTranspose3dConfig::new([5, 2], [5, 5, 5]).with_initializer(Initializer::Zeros); + let conv = config.init(&device); + + assert_eq!(config.initializer, Initializer::Zeros); + conv.weight.to_data().assert_approx_eq::( + &TensorData::zeros::(conv.weight.shape()), + Tolerance::default(), + ); + } + + #[test] + fn display() { + let config = ConvTranspose3dConfig::new([5, 2], [5, 5, 5]); + let conv = config.init(&Default::default()); + + assert_eq!( + format!("{conv}"), + "ConvTranspose3d {channels: [5, 2], stride: [1, 1, 1], kernel_size: [5, 5, 5], dilation: [1, 1, 1], groups: 1, padding: [0, 0, 0], padding_out: [0, 0, 0], params: 1252}" + ); + } + + #[test] + #[should_panic = "Number of channels in input tensor and input channels of convolution must be equal. got: 4, expected: 5"] + fn input_channels_mismatch() { + let config = ConvTranspose3dConfig::new([5, 3], [3, 3, 3]); + let conv = config.init(&Default::default()); + + let input = Tensor::<5>::zeros([1, 4, 10, 10, 10], &Default::default()); + let _ = conv.forward(input); + } +} diff --git a/crates/burn-nn/src/modules/conv/deform_conv2d.rs b/crates/burn-nn/src/modules/conv/deform_conv2d.rs new file mode 100644 index 0000000..5306e1a --- /dev/null +++ b/crates/burn-nn/src/modules/conv/deform_conv2d.rs @@ -0,0 +1,293 @@ +use alloc::format; +use burn::tensor::ops::DeformConvOptions; + +use burn_core as burn; + +use crate::PaddingConfig2d; +use burn::config::Config; +use burn::module::Initializer; +use burn::module::{Content, DisplaySettings, Module, ModuleDisplay, Param}; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn::tensor::module::deform_conv2d; + +use crate::conv::checks; + +/// Configuration to create a [deformable 2D convolution](DeformConv2d) layer, using the [init function](DeformConv2dConfig::init). +#[derive(Config, Debug)] +pub struct DeformConv2dConfig { + /// The number of channels. + pub channels: [usize; 2], + /// The size of the kernel. + pub kernel_size: [usize; 2], + /// The stride of the convolution. + #[config(default = "[1, 1]")] + pub stride: [usize; 2], + /// Spacing between kernel elements. + #[config(default = "[1, 1]")] + pub dilation: [usize; 2], + /// Controls the connections between input and output channels. + #[config(default = "1")] + pub weight_groups: usize, + /// Offset groups. + #[config(default = "1")] + pub offset_groups: usize, + /// The padding configuration. + /// + /// ### Warning + /// Only symmetric padding is currently supported. As such, using `Same` padding with an even kernel + /// size is not supported as it will not produce the same output size. + #[config(default = "PaddingConfig2d::Valid")] + pub padding: PaddingConfig2d, + /// If bias should be added to the output. + #[config(default = true)] + pub bias: bool, + /// The type of function used to initialize neural network parameters + #[config( + default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0),fan_out_only:false}" + )] + pub initializer: Initializer, +} + +/// Applies a deformable 2D convolution over input tensors. +/// +/// Should be created with [DeformConv2dConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct DeformConv2d { + /// Tensor of shape `[channels_out, channels_in / groups, kernel_size_1, kernel_size_2]` + pub weight: Param>, + /// Tensor of shape `[channels_out]` + pub bias: Option>>, + /// Stride of the convolution. + pub stride: [usize; 2], + /// Size of the kernel. + pub kernel_size: [usize; 2], + /// Spacing between kernel elements. + pub dilation: [usize; 2], + /// Controls the connections between input and output channels. + pub weight_groups: usize, + /// Offset groups. + pub offset_groups: usize, + /// The padding configuration. + #[module(skip)] + pub padding: PaddingConfig2d, +} + +impl DeformConv2dConfig { + /// Initialize a new [DeformConv2d](DeformConv2d) module. + pub fn init(&self, device: &Device) -> DeformConv2d { + checks::checks_channels_div_groups(self.channels[0], self.channels[1], self.weight_groups); + if self.padding == PaddingConfig2d::Same { + checks::check_same_padding_support(&self.kernel_size); + } + + let shape = [ + self.channels[1], + self.channels[0] / self.weight_groups, + self.kernel_size[0], + self.kernel_size[1], + ]; + + let k = self.kernel_size.iter().product::(); + let fan_in = self.channels[0] / self.weight_groups * k; + let fan_out = self.channels[1] / self.weight_groups * k; + + let weight = self + .initializer + .init_with(shape, Some(fan_in), Some(fan_out), device); + let mut bias = None; + + if self.bias { + bias = Some(self.initializer.init_with( + [self.channels[1]], + Some(fan_in), + Some(fan_out), + device, + )); + } + + DeformConv2d { + weight, + bias, + stride: self.stride, + kernel_size: self.kernel_size, + dilation: self.dilation, + padding: self.padding.clone(), + weight_groups: self.weight_groups, + offset_groups: self.weight_groups, + } + } +} + +impl ModuleDisplay for DeformConv2d { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + // Format the stride, kernel_size and dilation as strings, formatted as arrays instead of indexed. + let stride = format!("{:?}", self.stride); + let kernel_size = format!("{:?}", self.kernel_size); + let dilation = format!("{:?}", self.dilation); + + content + .add("stride", &stride) + .add("kernel_size", &kernel_size) + .add("dilation", &dilation) + .add("weight_groups", &self.weight_groups) + .add("offset_groups", &self.offset_groups) + .add_debug_attribute("padding", &self.padding) + .optional() + } +} + +impl DeformConv2d { + /// Applies the forward pass on the input tensor. + /// + /// See [deform_conv2d](burn::tensor::module::deform_conv2d) for more information. + /// + /// # Shapes + /// + /// - input: `[batch_size, channels_in, height_in, width_in]` + /// - offset: `[batch_size, 2 * offset_groups * kernel_height * kernel_width, height_out, width_out]` + /// - mask: `[batch_size, offset_groups * kernel_height * kernel_width, height_out, width_out]` + /// - output: `[batch_size, channels_out, height_out, width_out]` + pub fn forward( + &self, + input: Tensor<4>, + offset: Tensor<4>, + mask: Option>, + ) -> Tensor<4> { + let [_batch_size, _channels_in, height_in, width_in] = input.dims(); + let padding = + self.padding + .calculate_padding_2d(height_in, width_in, &self.kernel_size, &self.stride); + deform_conv2d( + input, + offset, + self.weight.val(), + mask, + self.bias.as_ref().map(|bias| bias.val()), + DeformConvOptions::new( + self.stride, + padding, + self.dilation, + self.weight_groups, + self.offset_groups, + ), + ) + } +} + +#[cfg(test)] +mod tests { + use burn::tensor::{ElementConversion, Tolerance}; + type FT = f32; + + use super::*; + use burn::tensor::TensorData; + + #[test] + fn initializer_default() { + let device = Device::default(); + device.seed(0); + + let config = DeformConv2dConfig::new([5, 1], [5, 5]); + let k = (config.channels[0] * config.kernel_size[0] * config.kernel_size[1]) as f64; + let k = (config.offset_groups as f64 / k).sqrt().elem::(); + let conv = config.init(&device); + + conv.weight.to_data().assert_within_range(-k..k); + } + + #[test] + fn initializer_zeros() { + let device = Device::default(); + device.seed(0); + + let config = DeformConv2dConfig::new([5, 2], [5, 5]).with_initializer(Initializer::Zeros); + let conv = config.init(&device); + + assert_eq!(config.initializer, Initializer::Zeros); + conv.weight.to_data().assert_approx_eq::( + &TensorData::zeros::(conv.weight.shape()), + Tolerance::default(), + ); + } + + #[test] + fn initializer_fan_out() { + let device = Device::default(); + device.seed(0); + + let init = Initializer::KaimingUniform { + gain: 1.0 / 3.0f64.sqrt(), + fan_out_only: true, // test that fan_out is passed to `init_with()` + }; + + let config = DeformConv2dConfig::new([5, 1], [5, 5]).with_initializer(init.clone()); + let c = config.init(&device); + let _ = c.weight.val(); // initializes param + + assert_eq!(config.initializer, init); + } + + #[test] + fn initializer_fan_with_groups_is_valid() { + let device = Device::default(); + device.seed(0); + + let init = Initializer::KaimingUniform { + gain: 1.0 / 3.0f64.sqrt(), + fan_out_only: true, + }; + + let config = DeformConv2dConfig::new([4, 4], [1, 1]) + .with_initializer(init.clone()) + .with_weight_groups(4); + let _ = config.init(&device); + + assert_eq!(config.initializer, init); + } + + #[test] + #[should_panic = "Both channels must be divisible by the number of groups."] + fn channels_with_groups_is_invalid() { + let device = Default::default(); + let config = DeformConv2dConfig::new([1, 4], [1, 1]).with_weight_groups(4); + let _ = config.init(&device); + } + + #[test] + #[should_panic = "Same padding with an even kernel size is not supported"] + fn same_with_even_kernel_is_invalid() { + let device = Default::default(); + let config = DeformConv2dConfig::new([4, 4], [2, 2]).with_padding(PaddingConfig2d::Same); + let _ = config.init(&device); + } + + #[test] + fn display() { + let config = DeformConv2dConfig::new([5, 1], [5, 5]); + let conv = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{conv}"), + "DeformConv2d {stride: [1, 1], kernel_size: [5, 5], dilation: [1, 1], weight_groups: 1, offset_groups: 1, padding: Valid, params: 126}" + ); + } + + #[test] + #[should_panic = "Number of channels in input tensor and input channels of convolution must be equal. got: 4, expected: 5"] + fn input_channels_mismatch() { + let config = DeformConv2dConfig::new([5, 3], [3, 3]); + let conv = config.init(&Default::default()); + + let input = Tensor::<4>::zeros([1, 4, 10, 10], &Default::default()); + let offset = Tensor::<4>::zeros([1, 2 * 3 * 3, 10, 10], &Default::default()); + let _ = conv.forward(input, offset, None); + } +} diff --git a/crates/burn-nn/src/modules/conv/mod.rs b/crates/burn-nn/src/modules/conv/mod.rs new file mode 100644 index 0000000..9c83141 --- /dev/null +++ b/crates/burn-nn/src/modules/conv/mod.rs @@ -0,0 +1,17 @@ +mod conv1d; +mod conv2d; +mod conv3d; +mod conv_transpose1d; +mod conv_transpose2d; +mod conv_transpose3d; +mod deform_conv2d; + +pub(crate) mod checks; + +pub use conv_transpose1d::*; +pub use conv_transpose2d::*; +pub use conv_transpose3d::*; +pub use conv1d::*; +pub use conv2d::*; +pub use conv3d::*; +pub use deform_conv2d::*; diff --git a/crates/burn-nn/src/modules/dropout.rs b/crates/burn-nn/src/modules/dropout.rs new file mode 100644 index 0000000..ef9be1f --- /dev/null +++ b/crates/burn-nn/src/modules/dropout.rs @@ -0,0 +1,118 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Module, ModuleDisplay}; +use burn::tensor::{Distribution, Tensor}; + +/// Configuration to create a [Dropout](Dropout) layer using the [init function](DropoutConfig::init). +#[derive(Config, Debug)] +pub struct DropoutConfig { + /// The probability of randomly zeroes some elements of the input tensor during training. + pub prob: f64, +} + +/// Set at random some elements of the input tensor to zero during training. +/// +/// This is an effective regularization technique as describe in the paper +/// [Improving neural networks by preventing co-adaptation of feature detectors](https://arxiv.org/abs/1207.0580). +/// +/// The input is also scaled during training to `1 / (1 - prob_keep)`. +/// +/// Should be created with [DropoutConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Dropout { + /// The probability of randomly zeroes some elements of the input tensor during training. + pub prob: f64, +} + +impl DropoutConfig { + /// Initialize a new [dropout](Dropout) module. + pub fn init(&self) -> Dropout { + if self.prob < 0.0 || self.prob > 1.0 { + panic!( + "Dropout probability should be between 0 and 1, but got {}", + self.prob + ); + } + Dropout { prob: self.prob } + } +} + +impl Dropout { + /// Applies the forward pass on the input tensor. + /// + /// See [Dropout](Dropout) for more information. + /// + /// # Shapes + /// + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + if !input.device().is_autodiff() || self.prob == 0.0 { + return input; + } + + let prob_keep = 1.0 - self.prob; + let random = input.random_like(Distribution::Bernoulli(prob_keep)); + let x = input * random; + + x * (1.0 / prob_keep) + } +} + +impl ModuleDisplay for Dropout { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content.add("prob", &self.prob).optional() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Shape; + + #[cfg(feature = "std")] + #[test] + fn with_ad_backend_should_mark_input() { + use burn::tensor::Device; + let device = Device::default().autodiff(); + let tensor = Tensor::<2>::ones(Shape::new([100, 100]), &device); + let dropout = DropoutConfig::new(0.5).init(); + + let output = dropout.forward(tensor.clone()); + + assert_ne!(tensor.to_data(), output.to_data()); + } + + #[test] + fn without_ad_backend_should_not_change_input() { + let tensor = Tensor::<2>::ones(Shape::new([100, 100]), &Default::default()); + let dropout = DropoutConfig::new(0.5).init(); + + let output = dropout.forward(tensor.clone()); + + assert_eq!(tensor.to_data(), output.to_data()); + } + + #[test] + fn display() { + let config = DropoutConfig::new(0.5); + let layer = config.init(); + + assert_eq!(alloc::format!("{layer}"), "Dropout {prob: 0.5}"); + } + + #[test] + #[should_panic = "Dropout probability should be between 0 and 1,"] + fn dropout_prob_invalid() { + let config = DropoutConfig::new(-10.); + let _layer = config.init(); + } +} diff --git a/crates/burn-nn/src/modules/embedding.rs b/crates/burn-nn/src/modules/embedding.rs new file mode 100644 index 0000000..1fc5481 --- /dev/null +++ b/crates/burn-nn/src/modules/embedding.rs @@ -0,0 +1,109 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::Initializer; +use burn::module::Module; +use burn::module::Param; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Int; +use burn::tensor::{Device, Tensor}; + +use burn::tensor::module::embedding; + +/// Configuration to create an [Embedding](Embedding) layer using the [init function](EmbeddingConfig::init). +#[derive(Config, Debug)] +pub struct EmbeddingConfig { + /// The number of embedding vectors. + pub n_embedding: usize, + /// The size of each vector. + pub d_model: usize, + /// The type of function used to initialize neural network parameters + #[config(default = "Initializer::Normal{mean:0.0, std:1.0}")] + pub initializer: Initializer, +} + +/// Lookup table to store a fix number of vectors. +/// +/// Should be created with [EmbeddingConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Embedding { + /// The learnable weights of the module of shape `[n_embedding, d_model]` initialized + /// from a normal distribution `N(0, 1)`. + pub weight: Param>, +} + +impl ModuleDisplay for Embedding { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let [n_embedding, d_model] = self.weight.shape().dims(); + content + .add("n_embedding", &n_embedding) + .add("d_model", &d_model) + .optional() + } +} + +impl EmbeddingConfig { + /// Initialize a new [embedding](Embedding) module. + pub fn init(&self, device: &Device) -> Embedding { + let weight = self + .initializer + .init([self.n_embedding, self.d_model], device); + + Embedding { weight } + } +} + +impl Embedding { + /// Applies the forward pass on the input tensor. + /// + /// See also [embedding](burn::tensor::module::embedding). + /// + /// # Shapes + /// + /// - input: `[batch_size, seq_length]` + /// - output: `[batch_size, seq_length, d_model]` + pub fn forward(&self, input: Tensor<2, Int>) -> Tensor<3> { + embedding(self.weight.val(), input) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn initializer_zeros() { + let device = Device::default(); + device.seed(0); + + let config = EmbeddingConfig::new(5, 5).with_initializer(Initializer::Zeros); + let embed = config.init(&Default::default()); + + assert_eq!(config.initializer, Initializer::Zeros); + embed.weight.to_data().assert_approx_eq::( + &TensorData::zeros::(embed.weight.shape()), + Tolerance::default(), + ); + } + + #[test] + fn display() { + let config = EmbeddingConfig::new(100, 10); + let embed = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{embed}"), + "Embedding {n_embedding: 100, d_model: 10, params: 1000}" + ); + } +} diff --git a/crates/burn-nn/src/modules/identity.rs b/crates/burn-nn/src/modules/identity.rs new file mode 100644 index 0000000..718880c --- /dev/null +++ b/crates/burn-nn/src/modules/identity.rs @@ -0,0 +1,32 @@ +use burn_core as burn; + +use burn::Tensor; +use burn::module::Module; + +/// Identity Module. +#[derive(Module, Default, Debug)] +pub struct Identity; + +impl Identity { + /// Create the module. + pub fn new() -> Self { + Self {} + } + + /// Forward pass, returns the input tensor. + pub fn forward(&self, input: Tensor) -> Tensor { + input + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn display() { + let layer = Identity::new(); + + assert_eq!(alloc::format!("{layer}"), "Identity"); + } +} diff --git a/crates/burn-nn/src/modules/interpolate/interpolate1d.rs b/crates/burn-nn/src/modules/interpolate/interpolate1d.rs new file mode 100644 index 0000000..5e3acea --- /dev/null +++ b/crates/burn-nn/src/modules/interpolate/interpolate1d.rs @@ -0,0 +1,257 @@ +use alloc::format; + +use burn::tensor::module::interpolate; + +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Module, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::tensor::ops::InterpolateOptions; + +use super::InterpolateMode; + +/// Configuration for the 1D interpolation module. +/// +/// This struct defines the configuration options for the 1D interpolation operation. +/// It allows specifying the output size, scale factor, and interpolation mode. +#[derive(Config, Debug)] +pub struct Interpolate1dConfig { + /// Output size of the interpolated tensor. + /// If specified, this takes precedence over `scale_factor`. + #[config(default = "None")] + pub output_size: Option, + + /// Scale factor for resizing the input tensor. + /// This is used when `output_size` is not specified. + #[config(default = "None")] + pub scale_factor: Option, + + /// Interpolation mode to use for resizing. + /// Determines how the output values are calculated. + #[config(default = "InterpolateMode::Nearest")] + pub mode: InterpolateMode, + + /// If `true`, the input and output tensors are aligned by their corner pixels. + /// If `false`, half-pixel coordinate mapping is used instead. + #[config(default = true)] + pub align_corners: bool, +} + +/// Interpolate module for resizing 1D tensors with shape [N, C, L]. +/// +/// This struct represents a 1D interpolation module that can resize tensors +/// using various interpolation methods. It provides flexibility in specifying +/// either an output size or a scale factor for resizing, along with options +/// for the interpolation mode. +/// +/// The module can be used to upsample or downsample 1D tensors, preserving the +/// number of channels and batch size while adjusting the length dimension. +/// +/// The module can be created using the [Interpolate1dConfig] struct and the +/// `init` method, which returns an instance of the [Interpolate1d] struct. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Interpolate1d { + /// Output size of the interpolated tensor + pub output_size: Option, + + /// Scale factor for resizing the input tensor + pub scale_factor: Option, + + /// Interpolation mode used for resizing + #[module(skip)] + pub mode: InterpolateMode, + + /// Whether to align corner pixels + pub align_corners: bool, +} + +impl Interpolate1dConfig { + /// Initialize the interpolation module + pub fn init(self) -> Interpolate1d { + Interpolate1d { + output_size: self.output_size, + scale_factor: self.scale_factor, + mode: self.mode, + align_corners: self.align_corners, + } + } +} + +impl Interpolate1d { + /// Performs the forward pass of the 1D interpolation module + /// + /// # Arguments + /// + /// * `input` - Input tensor with shape [N, C, L] + /// + /// # Returns + /// + /// Resized tensor with shape [N, C, L'], where L' is determined by + /// the output_size or scale_factor specified in the module configuration + /// + /// # Example + /// + /// ```ignore + /// let input = Tensor::::random([1, 3, 64], Distribution::Uniform(0.0, 1.0), &device); + /// let interpolate = Interpolate1dConfig::new() + /// .with_output_size(Some(128)) + /// .init(); + /// let output = interpolate.forward(input); + /// assert_eq!(output.dims(), [1, 3, 128]); + /// ``` + pub fn forward(&self, input: Tensor<3>) -> Tensor<3> { + let output_size = calculate_output_size(input.dims(), self.output_size, self.scale_factor); + + // Use the interpolate operation to resize the temporal input tensor + // by adding a new dimension for the interpolation axis + let input = input.unsqueeze_dim(2); + + let result = interpolate( + input, + [1, output_size], + InterpolateOptions::new(self.mode.clone().into()) + .with_align_corners(self.align_corners), + ); + + result.squeeze_dims(&[2]) + } +} + +/// Calculate output size based on input dimensions, output size, and scale factor +/// +/// # Arguments +/// +/// * `input_dims` - Input dimensions of the tensor +/// * `output_size` - Output size for the interpolated tensor +/// * `scale_factor` - Scale factor for resizing the tensor +/// +/// # Returns +/// +/// Output size for the interpolated tensor +/// +/// # Panics +/// +/// Panics if neither output_size nor scale_factor is provided +/// or if the scale factor is too large +fn calculate_output_size( + input_dims: [usize; 3], + output_size: Option, + scale_factor: Option, +) -> usize { + match (output_size, scale_factor) { + (Some(output_size), None) => { + // Use provided + output_size + } + (None, Some(scale_factor)) => { + // Calculate output size based on scale factor + let [_, _, l] = input_dims; + + let new_dim = (l as f64) * (scale_factor as f64); + + if new_dim > usize::MAX as f64 { + panic!("Scale factor is too large"); + } + + new_dim as usize + } + _ => panic!("Either output_size or scale_factor must be provided"), + } +} + +impl ModuleDisplay for Interpolate1d { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add_debug_attribute("mode", &self.mode) + .add("output_size", &format!("{:?}", self.output_size)) + .add("scale_factor", &self.scale_factor) + .optional() + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::Distribution; + + use super::*; + #[test] + fn test_calculate_output_size() { + let input_dims = [1, 1, 4]; + + let output_size = calculate_output_size(input_dims, Some(2), None); + assert_eq!(output_size, 2); + + let output_size = calculate_output_size(input_dims, None, Some(2.0)); + assert_eq!(output_size, 8); + + let output_size = calculate_output_size(input_dims, None, Some(0.5)); + assert_eq!(output_size, 2); + + let output_size = calculate_output_size(input_dims, None, Some(1.5)); + assert_eq!(output_size, 6); + } + + #[test] + #[should_panic(expected = "Either output_size or scale_factor must be provided")] + fn test_panic() { + let input_dims = [1, 1, 4]; + calculate_output_size(input_dims, None, None); + } + + #[test] + #[should_panic(expected = "Scale factor is too large")] + fn test_large_scale_factor() { + let input_dims = [1, 1, usize::MAX - 1]; + calculate_output_size(input_dims, None, Some(2.0)); + } + + #[test] + fn test_module() { + let input = Tensor::<3>::random( + [2, 3, 4], + Distribution::Uniform(0.0, 1.0), + &Default::default(), + ); + + // Test with output_size + let config = Interpolate1dConfig::new().with_output_size(Some(8)); + let interpolate = config.init(); + let output = interpolate.forward(input.clone()); + assert_eq!(output.dims(), [2, 3, 8]); + + // Test with scale_factor + let config = Interpolate1dConfig::new().with_scale_factor(Some(0.5)); + let interpolate = config.init(); + let output = interpolate.forward(input.clone()); + assert_eq!(output.dims(), [2, 3, 2]); + + // Test with different interpolation mode + let config = Interpolate1dConfig::new() + .with_output_size(Some(6)) + .with_mode(InterpolateMode::Linear); + let interpolate = config.init(); + let output = interpolate.forward(input); + assert_eq!(output.dims(), [2, 3, 6]); + } + + #[test] + fn display() { + let config = Interpolate1dConfig::new().with_output_size(Some(20)); + let layer = config.init(); + + assert_eq!( + alloc::format!("{layer}"), + "Interpolate1d {mode: Nearest, output_size: Some(20), \ + scale_factor: None}" + ); + } +} diff --git a/crates/burn-nn/src/modules/interpolate/interpolate2d.rs b/crates/burn-nn/src/modules/interpolate/interpolate2d.rs new file mode 100644 index 0000000..e99fb40 --- /dev/null +++ b/crates/burn-nn/src/modules/interpolate/interpolate2d.rs @@ -0,0 +1,259 @@ +use alloc::format; + +use burn::tensor::module::interpolate; + +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Module, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::tensor::ops::InterpolateOptions; + +use super::InterpolateMode; + +/// Configuration for the 2D interpolation module. +/// +/// This struct defines the configuration options for the 2D interpolation operation. +/// It allows specifying the output size, scale factor, and interpolation mode. +#[derive(Config, Debug)] +pub struct Interpolate2dConfig { + /// Output size of the interpolated tensor. + /// If specified, this takes precedence over `scale_factor`. + #[config(default = "None")] + pub output_size: Option<[usize; 2]>, + + /// Scale factor for resizing the input tensor. + /// This is used when `output_size` is not specified. + #[config(default = "None")] + pub scale_factor: Option<[f32; 2]>, + + /// Interpolation mode to use for resizing. + /// Determines how the output values are calculated. + #[config(default = "InterpolateMode::Nearest")] + pub mode: InterpolateMode, + + /// If `true`, the input and output tensors are aligned by their corner pixels. + /// If `false`, half-pixel coordinate mapping is used instead. + #[config(default = true)] + pub align_corners: bool, +} + +/// Interpolate module for resizing tensors with shape [N, C, H, W]. +/// +/// This struct represents an interpolation module that can resize tensors +/// using various interpolation methods. It provides flexibility in specifying +/// either an output size or a scale factor for resizing, along with options +/// for the interpolation mode. +/// +/// The module can be used to upsample or downsample tensors, preserving the +/// number of channels and batch size while adjusting the height and width +/// dimensions. +/// +/// The module can be created using the [Interpolate2dConfig] struct and the +/// `init` method, which returns an instance of the [Interpolate2d] struct. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Interpolate2d { + /// Output size of the interpolated tensor + pub output_size: Option<[usize; 2]>, + + /// Scale factor for resizing the input tensor + pub scale_factor: Option<[f32; 2]>, + + /// Interpolation mode used for resizing + #[module(skip)] + pub mode: InterpolateMode, + + /// Whether to align corner pixels + pub align_corners: bool, +} + +impl Interpolate2dConfig { + /// Initialize the interpolation module + pub fn init(self) -> Interpolate2d { + Interpolate2d { + output_size: self.output_size, + scale_factor: self.scale_factor, + mode: self.mode, + align_corners: self.align_corners, + } + } +} +impl Interpolate2d { + /// Performs the forward pass of the interpolation module + /// + /// # Arguments + /// + /// * `input` - Input tensor with shape [N, C, H, W] + /// + /// # Returns + /// + /// Resized tensor with shape [N, C, H', W'], where H' and W' are determined by + /// the output_size or scale_factor specified in the module configuration + /// + /// # Example + /// + /// ```ignore + /// let input = Tensor::::random([1, 3, 64, 64], Distribution::Uniform(0.0, 1.0), &device); + /// let interpolate = Interpolate2dConfig::new() + /// .with_output_size(Some([128, 128])) + /// .init(); + /// let output = interpolate.forward(input); + /// assert_eq!(output.dims(), [1, 3, 128, 128]); + /// ``` + pub fn forward(&self, input: Tensor<4>) -> Tensor<4> { + let output_size = calculate_output_size(input.dims(), self.output_size, self.scale_factor); + interpolate( + input, + output_size, + InterpolateOptions::new(self.mode.clone().into()) + .with_align_corners(self.align_corners), + ) + } +} + +/// Calculates the output size for tensor interpolation. +/// +/// # Arguments +/// +/// * `input_dims` - The dimensions of the input tensor [N, C, H, W]. +/// * `output_size` - Optional desired output size [H', W']. +/// * `scale_factor` - Optional scale factor for height and width [scale_h, scale_w]. +/// +/// # Returns +/// +/// A tuple [H', W'] representing the calculated output size. +/// +/// # Panics +/// +/// Panics if neither `output_size` nor `scale_factor` is provided, +/// or if the scale factor results in dimensions exceeding usize::MAX. +fn calculate_output_size( + input_dims: [usize; 4], + output_size: Option<[usize; 2]>, + scale_factor: Option<[f32; 2]>, +) -> [usize; 2] { + match (output_size, scale_factor) { + (Some(output_size), None) => { + // Use provided + output_size + } + (None, Some(scale_factor)) => { + // Calculate output size based on scale factor + let [_, _, h, w] = input_dims; + + let new_dim_h = (h as f64) * (scale_factor[0] as f64); + + if new_dim_h > usize::MAX as f64 { + panic!("Scale factor for height is too large"); + } + + let new_dim_w = (w as f64) * (scale_factor[1] as f64); + + if new_dim_w > usize::MAX as f64 { + panic!("Scale factor for width is too large"); + } + + [new_dim_h as usize, new_dim_w as usize] + } + _ => panic!("Either output_size or scale_factor must be provided"), + } +} + +impl ModuleDisplay for Interpolate2d { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add_debug_attribute("mode", &self.mode) + .add("output_size", &format!("{:?}", self.output_size)) + .add("scale_factor", &self.scale_factor) + .optional() + } +} +#[cfg(test)] +mod tests { + use burn::tensor::Distribution; + + use super::*; + + #[test] + fn test_calculate_output_size() { + let input_dims = [1, 1, 4, 4]; + + let output_size = calculate_output_size(input_dims, Some([2, 2]), None); + assert_eq!(output_size, [2, 2]); + + let output_size = calculate_output_size(input_dims, None, Some([2.0, 2.0])); + assert_eq!(output_size, [8, 8]); + + let output_size = calculate_output_size([1, 1, 4, 4], None, Some([0.5, 0.5])); + assert_eq!(output_size, [2, 2]); + + let output_size = calculate_output_size([1, 1, 4, 4], None, Some([2.0, 1.5])); + assert_eq!(output_size, [8, 6]); + } + + #[test] + #[should_panic(expected = "Either output_size or scale_factor must be provided")] + fn test_missing_params() { + calculate_output_size([1, 1, 4, 4], None, None); + } + + #[test] + #[should_panic(expected = "Scale factor for height is too large")] + fn test_infinite_height() { + calculate_output_size([1, 1, usize::MAX - 1, 4], None, Some([2.0, 1.0])); + } + + #[test] + #[should_panic(expected = "Scale factor for width is too large")] + fn test_infinite_width() { + calculate_output_size([1, 1, 4, usize::MAX - 1], None, Some([1.0, 2.0])); + } + + #[test] + fn test_module() { + let input = Tensor::<4>::random( + [2, 3, 4, 4], + Distribution::Uniform(0.0, 1.0), + &Default::default(), + ); + + // Test with output_size + let config = Interpolate2dConfig::new().with_output_size(Some([8, 8])); + let interpolate = config.init(); + let output = interpolate.forward(input.clone()); + assert_eq!(output.dims(), [2, 3, 8, 8]); + + // Test with scale_factor + let config = Interpolate2dConfig::new().with_scale_factor(Some([0.5, 0.5])); + let interpolate = config.init(); + let output = interpolate.forward(input.clone()); + assert_eq!(output.dims(), [2, 3, 2, 2]); + + // Test with different interpolation mode + let config = Interpolate2dConfig::new() + .with_output_size(Some([6, 6])) + .with_mode(InterpolateMode::Linear); + let interpolate = config.init(); + let output = interpolate.forward(input); + assert_eq!(output.dims(), [2, 3, 6, 6]); + } + + #[test] + fn display() { + let config = Interpolate2dConfig::new().with_output_size(Some([20, 20])); + let layer = config.init(); + + assert_eq!( + alloc::format!("{layer}"), + "Interpolate2d {mode: Nearest, output_size: Some([20, 20]), \ + scale_factor: None}" + ); + } +} diff --git a/crates/burn-nn/src/modules/interpolate/mod.rs b/crates/burn-nn/src/modules/interpolate/mod.rs new file mode 100644 index 0000000..ce56965 --- /dev/null +++ b/crates/burn-nn/src/modules/interpolate/mod.rs @@ -0,0 +1,64 @@ +mod interpolate1d; +mod interpolate2d; + +pub use interpolate1d::*; +pub use interpolate2d::*; + +use burn_core as burn; + +use burn::config::Config; +use burn::tensor::ops::InterpolateMode as OpsInterpolateMode; + +/// Algorithm used for downsampling and upsampling +/// +/// This enum defines different interpolation modes for resampling data. +#[derive(Config, Debug)] +pub enum InterpolateMode { + /// Nearest-neighbor floor interpolation + /// + /// This mode selects the value with floor mapping to a sample point for each output pixel. + /// It is applicable for both temporal and spatial data. + Nearest, + + /// Nearest-neighbor exact interpolation + /// + /// This mode selects the value of the nearest sample point for each output pixel. + /// It is applicable for both temporal and spatial data. + NearestExact, + + /// Linear interpolation + /// + /// This mode calculates the output value using linear + /// interpolation between nearby sample points. + /// + /// It is applicable for both temporal and spatial data. + Linear, + + /// Cubic interpolation + /// + /// This mode uses cubic interpolation to calculate the output value + /// based on surrounding sample points. + /// + /// It is applicable for both temporal and spatial data and generally + /// provides smoother results than linear interpolation. + Cubic, + + /// Lanczos3 interpolation + /// + /// This mode uses a 6-tap sinc-based Lanczos filter (a=3) to calculate + /// the output value. It generally provides high-quality results, + /// especially for downsampling. + Lanczos, +} + +impl From for OpsInterpolateMode { + fn from(mode: InterpolateMode) -> Self { + match mode { + InterpolateMode::Nearest => OpsInterpolateMode::Nearest, + InterpolateMode::NearestExact => OpsInterpolateMode::NearestExact, + InterpolateMode::Linear => OpsInterpolateMode::Bilinear, + InterpolateMode::Cubic => OpsInterpolateMode::Bicubic, + InterpolateMode::Lanczos => OpsInterpolateMode::Lanczos3, + } + } +} diff --git a/crates/burn-nn/src/modules/linear.rs b/crates/burn-nn/src/modules/linear.rs new file mode 100644 index 0000000..1c765b6 --- /dev/null +++ b/crates/burn-nn/src/modules/linear.rs @@ -0,0 +1,317 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::Param; +use burn::module::{Content, DisplaySettings, Initializer, Module, ModuleDisplay}; +use burn::tensor::module::linear; +use burn::tensor::{Device, Tensor}; + +/// Configuration to create a [`Linear`] layer using the [init function](LinearConfig::init). +#[derive(Config, Debug)] +pub struct LinearConfig { + /// The size of the input features. + pub d_input: usize, + /// The size of the output features. + pub d_output: usize, + /// If a bias should be applied during the linear transformation. + #[config(default = true)] + pub bias: bool, + /// The type of function used to initialize neural network parameters + #[config( + default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0), fan_out_only:false}" + )] + pub initializer: Initializer, + /// The layout in which the linear parameters are stored. + #[config(default = "LinearLayout::Row")] + pub layout: LinearLayout, +} + +#[derive(Config, Debug, Copy)] +/// The layout in which the linear parameters are stored. +/// +/// This can have performance impacts. +pub enum LinearLayout { + /// Parameters are stored in Row major. + Row, + /// Parameters are stored in Col major. + Col, +} + +/// Applies a linear transformation to the input tensor. +/// +/// Should be created with [LinearConfig] +/// +/// `O = IW + b` +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Linear { + /// Matrix of shape `[d_input, d_output]` initialized from a uniform distribution: + /// `U(-k, k)`, where `k = sqrt(1 / d_input)` + pub weight: Param>, + /// Vector of size `d_output` initialized from a uniform distribution: + /// `U(-k, k)`, where `k = sqrt(1 / d_input)` + pub bias: Option>>, +} + +impl LinearConfig { + /// Initialize a new [`Linear`] module. + pub fn init(&self, device: &Device) -> Linear { + let weight = match self.layout { + LinearLayout::Row => { + let shape = [self.d_input, self.d_output]; + self.initializer + .init_with(shape, Some(self.d_input), Some(self.d_output), device) + } + LinearLayout::Col => { + let shape = [self.d_output, self.d_input]; + + self.initializer + .init_with(shape, Some(self.d_output), Some(self.d_input), device) + // The param is already transposed when init. We re-transpose to have + // [d_output, d_input] while saving. + .save_mapper(move |tensor| { + let device = tensor.device(); + device.sync().unwrap(); + let tensor = tensor.transpose(); + device.sync().unwrap(); + tensor + }) + // When loading from record we have to transpose. + .load_mapper(move |tensor| { + let device = tensor.device(); + device.sync().unwrap(); + let tensor = tensor.transpose(); + device.sync().unwrap(); + + tensor + }) + // When loading from initialization, we have to transpose. + .init_mapper(|tensor| { + let device = tensor.device(); + device.sync().unwrap(); + let tensor = tensor.transpose(); + device.sync().unwrap(); + tensor + }) + } + }; + let bias = if self.bias { + Some(self.initializer.init_with( + [self.d_output], + Some(self.d_input), + Some(self.d_output), + device, + )) + } else { + None + }; + + Linear { weight, bias } + } +} + +impl Linear { + /// Applies the forward pass on the input tensor. + /// + /// # Arguments + /// + /// - `input` - The input tensor of shape `[..., d_input]`. + /// + /// # Shapes + /// + /// - input: `[..., d_input]` + /// - output: `[..., d_output]` + /// + /// # Returns + /// + /// The transformed tensor of shape `[..., d_output]`. + pub fn forward(&self, input: Tensor) -> Tensor { + linear( + input, + self.weight.val(), + self.bias.as_ref().map(|b| b.val()), + ) + } +} + +impl ModuleDisplay for Linear { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let [d_input, d_output] = self.weight.shape().dims(); + content + .add("d_input", &d_input) + .add("d_output", &d_output) + .add("bias", &self.bias.is_some()) + .optional() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::module::ParamId; + use burn::store::ModuleRecord; + use burn::tensor::ElementConversion; + use burn::tensor::Tolerance; + use burn::tensor::{Shape, TensorData}; + type FT = f32; + + #[test] + fn initializer_default() { + let device = Device::default(); + device.seed(0); + + let config = LinearConfig::new(5, 5); + let k = (1.0 / config.d_input as f64).sqrt().elem::(); + let linear = config.init(&device); + + assert_eq!( + config.initializer, + Initializer::KaimingUniform { + gain: 1.0 / 3.0f64.sqrt(), + fan_out_only: false + } + ); + linear.weight.to_data().assert_within_range(-k..k); + } + + #[test] + fn initializer_zeros() { + let device = Device::default(); + device.seed(0); + + let config = LinearConfig::new(5, 5).with_initializer(Initializer::Zeros); + let linear = config.init(&device); + + assert_eq!(config.initializer, Initializer::Zeros); + linear.weight.to_data().assert_approx_eq::( + &TensorData::zeros::(linear.weight.shape()), + Tolerance::default(), + ); + } + + #[test] + fn test_linear_forward_no_bias() { + let device = Device::default(); + device.seed(0); + + let value = 2.; + let config = LinearConfig::new(2, 3) + .with_initializer(Initializer::Constant { value }) + .with_bias(false); + let linear = config.init(&device); + + let input = Tensor::<2>::ones(Shape::new([1, 2]), &device); + let result = linear.forward(input); + let expected_result = Tensor::<2>::from_data([[4., 4., 4.]], &device); + + assert_eq!(result.into_data(), expected_result.into_data()); + } + + #[test] + fn test_linear_forward_with_bias() { + let device = Device::default(); + device.seed(0); + + let device = Device::default(); + + let value = 2.; + let config = LinearConfig::new(2, 3).with_initializer(Initializer::Constant { value }); + let linear = config.init(&device); + + let input = Tensor::<2>::ones(Shape::new([1, 2]), &device); + let result = linear.forward(input); + let expected_result = Tensor::<2>::from_data([[6., 6., 6.]], &device); + + assert_eq!(result.into_data(), expected_result.into_data()); + } + + #[test] + fn test_linear_1d() { + let device = Device::default(); + device.seed(0); + + let value = 2.; + let config = LinearConfig::new(2, 3).with_initializer(Initializer::Constant { value }); + let linear = config.init(&device); + + let input_1d = Tensor::<1>::ones(Shape::new([2]), &device); + let input_2d = Tensor::<2>::ones(Shape::new([1, 2]), &device); + + let result_1d = linear.forward(input_1d).unsqueeze::<2>(); + let result_2d = linear.forward(input_2d); + + assert_eq!(result_1d.into_data(), result_2d.into_data()); + } + + #[test] + fn display() { + let config = LinearConfig::new(3, 5); + let linear = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{linear}"), + "Linear {d_input: 3, d_output: 5, bias: true, params: 20}" + ); + } + + #[test] + fn layout() { + let device = Default::default(); + // The `Col` layout exposes the weight matrix as `[d_input, d_output]`. + let linear = LinearConfig::new(6, 12) + .with_layout(LinearLayout::Col) + .init(&device); + assert_eq!(linear.weight.dims(), [6, 12], "Shape is as configured"); + } + + #[test] + fn round_trip_burnpack() { + let device = Default::default(); + let linear = LinearConfig::new(6, 12).init(&device); + + let weight_before = linear.weight.val().to_data(); + let data = linear.into_record().into_bytes().unwrap(); + + let linear = LinearConfig::new(6, 12) + .init(&device) + .load_record(ModuleRecord::from_bytes(data).unwrap()); + + linear + .weight + .val() + .to_data() + .assert_eq(&weight_before, true); + } + + #[test] + fn col_row_same_result() { + let device = Default::default(); + let config_col = LinearConfig::new(6, 12).with_layout(LinearLayout::Col); + let linear_col = config_col.init(&device); + let signal = Tensor::<2>::random([8, 6], burn::tensor::Distribution::Default, &device); + let value = linear_col.forward(signal.clone()); + + let data_1 = value.into_data(); + + let weights = linear_col.weight.val().into_data(); + let weights = Tensor::from_data(weights, &device); + + let linear = Linear { + weight: Param::initialized(ParamId::new(), weights), + bias: linear_col + .bias + .map(|b| Param::initialized(ParamId::new(), b.val())), + }; + + let value = linear.forward(signal); + let data_2 = value.into_data(); + + data_1.assert_approx_eq::(&data_2, Default::default()); + } +} diff --git a/crates/burn-nn/src/modules/mod.rs b/crates/burn-nn/src/modules/mod.rs new file mode 100644 index 0000000..d08d91f --- /dev/null +++ b/crates/burn-nn/src/modules/mod.rs @@ -0,0 +1,41 @@ +/// Attention module +pub mod attention; + +/// Cache module +pub mod cache; + +/// Convolution module +pub mod conv; + +/// Pooling module +pub mod pool; + +/// Transformer module +pub mod transformer; + +mod identity; +/// Interpolate module +pub mod interpolate; + +mod dropout; +mod embedding; +mod linear; +mod noise; +mod pos_encoding; +mod rnn; +mod rope_encoding; +mod unfold; + +pub mod norm; + +pub use norm::{batch::*, group::*, instance::*, layer::*, local_response::*, rms::*}; + +pub use dropout::*; +pub use embedding::*; +pub use identity::*; +pub use linear::*; +pub use noise::*; +pub use pos_encoding::*; +pub use rnn::*; +pub use rope_encoding::*; +pub use unfold::*; diff --git a/crates/burn-nn/src/modules/noise.rs b/crates/burn-nn/src/modules/noise.rs new file mode 100644 index 0000000..aead989 --- /dev/null +++ b/crates/burn-nn/src/modules/noise.rs @@ -0,0 +1,116 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Module, ModuleDisplay}; +use burn::tensor::{Distribution, Tensor}; + +/// Configuration to create a [GaussianNoise](GaussianNoise) layer using the [init function](GaussianNoiseConfig::init). +#[derive(Config, Debug)] +pub struct GaussianNoiseConfig { + /// Standard deviation of the normal noise distribution. + pub std: f64, +} + +/// Add pseudorandom Gaussian noise to an arbitrarily shaped tensor. +/// +/// This is an effective regularization technique that also contributes to data augmentation. +/// Please keep in mind that the value of [std](GaussianNoise::std) should be chosen with care in order to avoid +/// distortion. +/// +/// Should be created with [GaussianNoiseConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct GaussianNoise { + /// Standard deviation of the normal noise distribution. + pub std: f64, +} + +impl GaussianNoiseConfig { + /// Initialize a new [Gaussian noise](GaussianNoise) module. + pub fn init(&self) -> GaussianNoise { + if self.std.is_sign_negative() { + panic!( + "Standard deviation is required to be non-negative, but got {}", + self.std + ); + } + GaussianNoise { std: self.std } + } +} + +impl GaussianNoise { + /// Applies the forward pass on the input tensor. + /// + /// See [GaussianNoise](GaussianNoise) for more information. + /// + /// # Shapes + /// + /// - input: `[..., any]` + /// - output: `[..., any]` + pub fn forward(&self, input: Tensor) -> Tensor { + if input.device().is_autodiff() && self.std != 0.0 { + let noise = Tensor::random( + input.shape(), + Distribution::Normal(0.0, self.std), + &input.device(), + ); + input + noise + } else { + input + } + } +} + +impl ModuleDisplay for GaussianNoise { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content.add("std", &self.std).optional() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::{Device, Shape}; + + #[cfg(feature = "std")] + #[test] + fn with_ad_backend_should_mark_input() { + let device = Device::default().autodiff(); + let tensor = Tensor::<2>::ones(Shape::new([100, 100]), &device); + let noise = GaussianNoiseConfig::new(0.5).init(); + + let output = noise.forward(tensor.clone()); + + assert_ne!(tensor.to_data(), output.to_data()); + } + + #[test] + fn without_ad_backend_should_not_change_input() { + let tensor = Tensor::<2>::ones(Shape::new([100, 100]), &Default::default()); + let noise = GaussianNoiseConfig::new(0.5).init(); + + let output = noise.forward(tensor.clone()); + + assert_eq!(tensor.to_data(), output.to_data()); + } + + #[test] + #[should_panic(expected = "Standard deviation is required to be non-negative")] + fn negative_std_should_panic() { + GaussianNoiseConfig { std: -0.5 }.init(); + } + + #[test] + fn display() { + let config = GaussianNoiseConfig::new(0.5); + let layer = config.init(); + + assert_eq!(alloc::format!("{layer}"), "GaussianNoise {std: 0.5}"); + } +} diff --git a/crates/burn-nn/src/modules/norm/batch.rs b/crates/burn-nn/src/modules/norm/batch.rs new file mode 100644 index 0000000..8cb51ed --- /dev/null +++ b/crates/burn-nn/src/modules/norm/batch.rs @@ -0,0 +1,485 @@ +use burn_core as burn; + +use burn::module::Initializer; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::{Device, Tensor}; +use burn::{ + config::Config, + module::{Module, Param, RunningState}, +}; + +/// [`BatchNorm`] Configuration. +/// +/// Used to create a [`BatchNorm`] layer using the [`BatchNormConfig::init`]. +#[derive(Config, Debug)] +pub struct BatchNormConfig { + /// The number of features. + pub num_features: usize, + /// A value required for numerical stability. Default: 1e-5 + #[config(default = 1e-5)] + pub epsilon: f64, + /// Momentum used to update the metrics. Default: 0.1 + #[config(default = 0.1)] + pub momentum: f64, +} + +/// Applies Batch Normalization over a tensor. +/// +/// Based upon the paper [Batch Normalization](https://arxiv.org/abs/1502.03167). +/// +/// Assumes input tensor is of shape ``[batch_size, channels, ...]``. +/// +/// `Y = norm(X) * γ + β` +/// +/// Where: +/// - `X` is the input tensor +/// - `Y` is the output tensor +/// - `norm` is the normalization function +/// - `γ` is the learnable weight +/// - `β` is the learnable bias +/// +/// Should be created using [`BatchNormConfig`]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct BatchNorm { + /// The learnable weight gamma. + pub gamma: Param>, + /// The learnable weight beta. + pub beta: Param>, + /// The running mean. + pub running_mean: RunningState>, + /// The running variance. + pub running_var: RunningState>, + /// Momentum used to update the metrics. + pub momentum: f64, + /// A value required for numerical stability. + pub epsilon: f64, +} + +impl BatchNormConfig { + /// Initializes a new [batch norm](BatchNorm) module. + pub fn init(&self, device: &Device) -> BatchNorm { + let gamma = Initializer::Ones.init([self.num_features], device); + let beta = Initializer::Zeros.init([self.num_features], device); + + let running_mean = Tensor::zeros([self.num_features], device); + let running_var = Tensor::ones([self.num_features], device); + + BatchNorm { + gamma, + beta, + running_mean: RunningState::new(running_mean), + running_var: RunningState::new(running_var), + momentum: self.momentum, + epsilon: self.epsilon, + } + } +} + +impl BatchNorm { + /// Applies the forward pass on the input tensor. + /// + /// See [`BatchNorm`] for more information. + /// + /// # Shapes + /// + /// - `input`: ``[batch_size, channels, ...]`` + /// - `output`: ``[batch_size, channels, ...]`` + /// + /// # Panics + /// + /// This function will panic if the input tensor has rank < 2. + pub fn forward(&self, input: Tensor) -> Tensor { + // Should be move to a compilation error when const generic support that kind of + // validation. https://github.com/rust-lang/rust/issues/76560 + if D < 2 { + panic!( + "BatchNorm can only be applied on tensors of rank >= 2 with the following shape \ + [batch_size, channels, ...], received {}D tensor", + D + ); + } + + match input.device().is_autodiff() { + true => self.forward_train(input), + false => self.forward_inference(input), + } + } + + fn forward_inference(&self, input: Tensor) -> Tensor { + let device = input.device(); + let channels = input.dims()[1]; + let mean = self.running_mean.value().to_device(&device); + let var = self.running_var.value().to_device(&device); + + let mut shape = [1; D]; + shape[1] = channels; + + self.forward_shared(input, mean.reshape(shape), var.reshape(shape)) + } + + fn forward_train(&self, input: Tensor) -> Tensor { + let device = input.device(); + let dims = input.dims(); + let batch_size = dims[0]; + let channels = dims[1]; + + let mut shape_unsqueeze = [1; D]; + let mut flatten_size = batch_size; + shape_unsqueeze[1] = channels; + + for dim in dims.iter().take(D).skip(2) { + flatten_size *= dim; + } + + let mean = input + .clone() + .swap_dims(0, 1) + .reshape([channels, flatten_size]) + .mean_dim(1) + .reshape(shape_unsqueeze); + + let var = input + .clone() + .sub(mean.clone()) + .square() + .swap_dims(0, 1) + .reshape([channels, flatten_size]) + .mean_dim(1) + .reshape(shape_unsqueeze); + + let running_mean = self.running_mean.value_sync().to_device(&device); + let running_var = self.running_var.value_sync().to_device(&device); + + let running_mean = running_mean.mul_scalar(1.0 - self.momentum).add( + mean.clone() + .detach() + .mul_scalar(self.momentum) + .reshape([channels]), + ); + let running_var = running_var.mul_scalar(1.0 - self.momentum).add( + var.clone() + .detach() + .mul_scalar(self.momentum) + .reshape([channels]), + ); + + self.running_mean.update(running_mean.detach()); + self.running_var.update(running_var.detach()); + + self.forward_shared(input, mean, var) + } + + fn forward_shared( + &self, + x: Tensor, + mean: Tensor, + var: Tensor, + ) -> Tensor { + let channels = x.dims()[1]; + let mut shape = [1; D]; + shape[1] = channels; + + let std = var.add_scalar(self.epsilon).sqrt(); + + let x = x.sub(mean); + let x = x.div(std); + + let x = x.mul(self.gamma.val().reshape(shape)); + + x.add(self.beta.val().reshape(shape)) + } +} + +impl ModuleDisplay for BatchNorm { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let [num_features] = self.beta.shape().dims(); + + content + .add("num_features", &num_features) + .add("momentum", &self.momentum) + .add("epsilon", &self.epsilon) + .optional() + } +} + +#[cfg(feature = "std")] +#[cfg(test)] +mod tests_1d { + use super::*; + use burn::module::AutodiffModule; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn batch_norm_forward_train() { + let device = Device::default().autodiff(); + let module = BatchNormConfig::new(3).init(&device); + + let output = module.forward(input_tensor(&device)); + + output + .to_data() + .assert_approx_eq::(&expected_train(), Tolerance::rel_abs(0.1, 0.001)); + } + + #[test] + fn batch_norm_forward_inference() { + let device = Device::default(); + let device_autodiff = device.clone().autodiff(); + let module = BatchNormConfig::new(3).init(&device_autodiff); + + module.forward(input_tensor(&device_autodiff)); + let module = module.valid(); + let output = module.forward(input_tensor(&device)); + + output + .to_data() + .assert_approx_eq::(&expected_valid(), Tolerance::default()); + } + + fn expected_valid() -> TensorData { + TensorData::from([ + [[0.9409, 0.6976], [0.5892, 0.8774], [0.9106, 0.6844]], + [[0.6012, 0.0782], [-0.0394, 0.9270], [0.6181, 0.5492]], + ]) + } + + fn expected_train() -> TensorData { + TensorData::from([ + [ + [1.1483e+00, 3.7521e-01], + [1.6272e-03, 7.5067e-01], + [1.6204e+00, -4.5168e-02], + ], + [ + [6.8856e-02, -1.5923e+00], + [-1.6318e+00, 8.7949e-01], + [-5.3368e-01, -1.0416e+00], + ], + ]) + } + + fn input_tensor(device: &Device) -> Tensor<3> { + Tensor::<3>::from_floats( + [ + [[0.9601, 0.7277], [0.6272, 0.9034], [0.9378, 0.7230]], + [[0.6356, 0.1362], [0.0249, 0.9509], [0.6600, 0.5945]], + ], + device, + ) + } + + #[test] + fn batch_norm_forward_train_inference() { + let device = Device::default(); + let device_autodiff = device.clone().autodiff(); + let module = BatchNormConfig::new(3).init(&device_autodiff); + + module.forward(input_tensor(&device_autodiff)); + let module = module.valid(); + let output = module.forward(input_tensor(&device)); + + output + .to_data() + .assert_approx_eq::(&expected_valid(), Tolerance::default()); + + let module = module.train(); + let output = module.forward(input_tensor(&device_autodiff)); + output + .to_data() + .assert_approx_eq::(&expected_train(), Tolerance::default()); + } +} + +#[cfg(feature = "std")] +#[cfg(test)] +mod tests_2d { + use super::*; + use burn::module::AutodiffModule; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn batch_norm_forward_train() { + let device = Device::default().autodiff(); + let module = BatchNormConfig::new(3).init(&device); + + let output = module.forward(input_tensor(&device)); + + let expected = TensorData::from([ + [ + [[1.5136, 0.7506], [-1.2216, 0.1477]], + [[0.3135, 1.2252], [-0.4150, 0.6130]], + [[1.4186, 0.3372], [-1.5183, 1.5262]], + ], + [ + [[0.4483, -1.1914], [-1.2010, 0.7537]], + [[-1.6752, 1.3822], [-0.5058, -0.9381]], + [[0.0200, -0.3097], [-0.5715, -0.9026]], + ], + ]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(0.1, 0.001)); + } + + #[test] + fn batch_norm_forward_inference() { + let device = Device::default(); + let device_autodiff = device.clone().autodiff(); + let module = BatchNormConfig::new(3).init(&device_autodiff); + + module.forward(input_tensor(&device_autodiff)); + let module = module.valid(); + let output = module.forward(input_tensor(&device)); + + let expected = TensorData::from([ + [ + [[0.9538, 0.7103], [0.0808, 0.5179]], + [[0.6015, 0.8910], [0.3703, 0.6966]], + [[0.9171, 0.6912], [0.3037, 0.9395]], + ], + [ + [[0.6138, 0.0904], [0.0874, 0.7113]], + [[-0.0297, 0.9408], [0.3415, 0.2042]], + [[0.6250, 0.5561], [0.5013, 0.4323]], + ], + ]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn batch_norm_running_mean() { + let device = Device::default().autodiff(); + let module = BatchNormConfig::new(3).init(&device); + + let _output = module.forward(input_tensor(&device)); + + let running_mean = module.running_mean.value_sync(); + + let expected = TensorData::from([0.0499, 0.0532, 0.0656]); + running_mean + .reshape([3]) + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn batch_norm_running_var() { + let device = Device::default().autodiff(); + let module = BatchNormConfig::new(3).init(&device); + + let _output = module.forward(input_tensor(&device)); + + let running_var = module.running_var.value_sync(); + + let expected = TensorData::from([0.9106, 0.9105, 0.9045]); + running_var + .reshape([3]) + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn batch_norm_running_mean_inner_module() { + let device = Device::default().autodiff(); + let module = BatchNormConfig::new(3).init(&device); + + let _output = module.forward(input_tensor(&device)); + + let module_valid = module.valid(); + let running_mean = module_valid.running_mean.value(); + let running_mean_after = module.running_mean.value(); + + running_mean_after + .into_data() + .assert_approx_eq::(&running_mean.into_data(), Tolerance::default()); + } + + #[test] + fn batch_norm_grads() { + let device = Device::default().autodiff(); + let module = BatchNormConfig::new(3).init(&device); + let input = input_tensor(&device).require_grad(); + + let output = module.forward(input.clone()); + + let grads = output.backward(); + + let tolerance = Tolerance::rel_abs(0.1, 0.001); + let expected = TensorData::from([0.0000e+00, -5.9035e-07, -6.0011e-07]); + module + .gamma + .grad(&grads) + .unwrap() + .reshape([3]) + .into_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([8., 8., 8.]); + module + .beta + .grad(&grads) + .unwrap() + .reshape([3]) + .into_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([ + [ + [[0.0000e+00, 0.0000e+00], [0.0000e+00, 0.0000e+00]], + [[7.6400e-08, 2.9848e-07], [-1.0110e-07, 1.4933e-07]], + [[5.3570e-07, 1.2732e-07], [-5.7336e-07, 5.7632e-07]], + ], + [ + [[0.0000e+00, 0.0000e+00], [0.0000e+00, 0.0000e+00]], + [[-4.0807e-07, 3.3673e-07], [-1.2323e-07, -2.2854e-07]], + [[7.5642e-09, -1.1695e-07], [-2.1582e-07, -3.4078e-07]], + ], + ]); + input + .grad(&grads) + .unwrap() + .into_data() + .assert_approx_eq::(&expected, tolerance); + } + + fn input_tensor(device: &Device) -> Tensor<4> { + Tensor::<4>::from_floats( + [ + [ + [[0.9601, 0.7277], [0.1270, 0.5441]], + [[0.6272, 0.9034], [0.4066, 0.7179]], + [[0.9378, 0.7230], [0.3544, 0.9591]], + ], + [ + [[0.6356, 0.1362], [0.1333, 0.7287]], + [[0.0249, 0.9509], [0.3791, 0.2481]], + [[0.6600, 0.5945], [0.5424, 0.4767]], + ], + ], + device, + ) + } + + #[test] + fn display() { + let batch_norm = BatchNormConfig::new(3).init(&Default::default()); + + assert_eq!( + format!("{batch_norm}"), + "BatchNorm {num_features: 3, momentum: 0.1, epsilon: 0.00001, params: 12}" + ); + } +} diff --git a/crates/burn-nn/src/modules/norm/group.rs b/crates/burn-nn/src/modules/norm/group.rs new file mode 100644 index 0000000..0f53761 --- /dev/null +++ b/crates/burn-nn/src/modules/norm/group.rs @@ -0,0 +1,331 @@ +use burn::module::Initializer; +use burn_core as burn; + +use burn::config::Config; +use burn::module::Module; +use burn::module::Param; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Device; +use burn::tensor::Tensor; + +/// Configuration to create a [GroupNorm](GroupNorm) layer using the [init function](GroupNormConfig::init). +#[derive(Debug, Config)] +pub struct GroupNormConfig { + /// The number of groups to separate the channels into + pub num_groups: usize, + /// The number of channels expected in the input + pub num_channels: usize, + /// A value required for numerical stability. Default: 1e-5 + #[config(default = 1e-5)] + pub epsilon: f64, + /// A boolean value that when set to `true`, this module has learnable + /// per-channel affine parameters initialized to ones (for weights) + /// and zeros (for biases). Default: `true` + #[config(default = true)] + pub affine: bool, +} + +/// Applies Group Normalization over a mini-batch of inputs as described in the paper [Group Normalization](https://arxiv.org/abs/1803.08494). +/// +/// `Y = groupnorm(X) * γ + β` +/// +/// Where: +/// - `X` is the input tensor +/// - `Y` is the output tensor +/// - `γ` is the learnable weight +/// - `β` is the learnable bias +/// +/// Should be created using [GroupNormConfig](GroupNormConfig). +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct GroupNorm { + /// The learnable weight + pub gamma: Option>>, + /// The learnable bias + pub beta: Option>>, + /// The number of groups to separate the channels into + pub num_groups: usize, + /// The number of channels expected in the input + pub num_channels: usize, + /// A value required for numerical stability + pub epsilon: f64, + /// A boolean value that when set to `true`, this module has learnable + pub affine: bool, +} + +impl ModuleDisplay for GroupNorm { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("num_groups", &self.num_groups) + .add("num_channels", &self.num_channels) + .add("epsilon", &self.epsilon) + .add("affine", &self.affine) + .optional() + } +} + +impl GroupNormConfig { + /// Initialize a new [group norm](GroupNorm) module. + pub fn init(&self, device: &Device) -> GroupNorm { + assert_eq!( + self.num_channels % self.num_groups, + 0, + "The number of channels must be divisible by the number of groups" + ); + + let (gamma, beta) = if self.affine { + let gamma = Initializer::Ones.init([self.num_channels], device); + let beta = Initializer::Zeros.init([self.num_channels], device); + + (Some(gamma), Some(beta)) + } else { + (None, None) + }; + + GroupNorm { + num_groups: self.num_groups, + num_channels: self.num_channels, + gamma, + beta, + epsilon: self.epsilon, + affine: self.affine, + } + } +} + +impl GroupNorm { + /// Applies the forward pass on the input tensor. + /// + /// See [GroupNorm](GroupNorm) for more information. + /// + /// # Shapes + /// + /// - input: `[batch_size, num_channels, *]` + /// - output: `[batch_size, num_channels, *]` + pub fn forward(&self, input: Tensor) -> Tensor { + if input.shape()[1] != self.num_channels { + panic!( + "The number of channels in the input tensor should be equal to the number of channels in the GroupNorm module. Expected {}, got {}", + self.num_channels, + input.shape()[1] + ); + } + + let gamma = self.gamma.as_ref().map(|x| x.val()); + let beta = self.beta.as_ref().map(|x| x.val()); + + group_norm( + input, + gamma, + beta, + self.num_groups, + self.epsilon, + self.affine, + ) + } +} + +/// Applies Group Normalization over a mini-batch of inputs as described in the paper [Group Normalization](https://arxiv.org/abs/1803.08494). +/// +/// `Y = groupnorm(X) * γ + β` +/// +/// Where: +/// - `X` is the input tensor +/// - `Y` is the output tensor +/// - `γ` is the learnable weight +/// - `β` is the learnable bias +/// +pub(crate) fn group_norm( + input: Tensor, + gamma: Option>, + beta: Option>, + num_groups: usize, + epsilon: f64, + affine: bool, +) -> Tensor { + if (beta.is_none() || gamma.is_none()) && affine { + panic!("Affine is set to true, but gamma or beta is None"); + } + + let shape = input.shape(); + if shape.num_elements() <= 2 { + panic!( + "input rank for GroupNorm should be at least 3, but got {}", + shape.num_elements() + ); + } + + let batch_size = shape[0]; + let num_channels = shape[1]; + + let hidden_size = shape[2..].iter().product::() * num_channels / num_groups; + let input = input.reshape([batch_size, num_groups, hidden_size]); + + let mean = input.clone().sum_dim(2) / hidden_size as f64; + let input = input.sub(mean); + + let var = input.clone().square().sum_dim(2) / hidden_size as f64; + let input_normalized = input.div(var.add_scalar(epsilon).sqrt()); + + if affine { + let mut affine_shape = [1; D]; + affine_shape[1] = num_channels; + + input_normalized + .reshape(shape) + .mul(gamma.clone().unwrap().reshape(affine_shape)) + .add(beta.clone().unwrap().reshape(affine_shape)) + } else { + input_normalized.reshape(shape) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use alloc::format; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn group_norm_forward_affine_false() { + let device = Default::default(); + let module = GroupNormConfig::new(2, 6).with_affine(false).init(&device); + + assert!(module.gamma.is_none()); + assert!(module.beta.is_none()); + + let input = Tensor::<3>::from_data( + TensorData::from([ + [ + [-0.3034, 0.2726, -0.9659], + [-1.1845, -1.3236, 0.0172], + [1.9507, 1.2554, -0.8625], + [1.0682, 0.3604, 0.3985], + [-0.4957, -0.4461, -0.9721], + [1.5157, -0.1546, -0.5596], + ], + [ + [-1.6698, -0.4040, -0.7927], + [0.3736, -0.0975, -0.1351], + [-0.9461, 0.5461, -0.6334], + [-1.0919, -0.1158, 0.1213], + [-0.9535, 0.1281, 0.4372], + [-0.2845, 0.3488, 0.5641], + ], + ]), + &device, + ); + + let output = module.forward(input); + + let expected = TensorData::from([ + [ + [-0.1653, 0.3748, -0.7866], + [-0.9916, -1.1220, 0.1353], + [1.9485, 1.2965, -0.6896], + [1.2769, 0.3628, 0.4120], + [-0.7427, -0.6786, -1.3578], + [1.8547, -0.3022, -0.8252], + ], + [ + [-1.9342, 0.0211, -0.5793], + [1.2223, 0.4945, 0.4365], + [-0.8163, 1.4887, -0.3333], + [-1.7960, -0.0392, 0.3875], + [-1.5469, 0.3998, 0.9561], + [-0.3428, 0.7970, 1.1845], + ], + ]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn group_norm_forward_affine_true() { + let device = Default::default(); + let module = GroupNormConfig::new(3, 6).with_affine(true).init(&device); + + let tolerance = Tolerance::permissive(); + module + .gamma + .as_ref() + .expect("gamma should not be None") + .val() + .to_data() + .assert_approx_eq::(&TensorData::ones::([6]), tolerance); + + module + .beta + .as_ref() + .expect("beta should not be None") + .val() + .to_data() + .assert_approx_eq::(&TensorData::zeros::([6]), tolerance); + + let input = Tensor::<3>::from_data( + TensorData::from([ + [ + [0.3345, 0.4429, 0.6639], + [0.5041, 0.4175, 0.8437], + [0.6159, 0.3758, 0.4071], + [0.5417, 0.5785, 0.7671], + [0.3837, 0.9883, 0.0420], + [0.4808, 0.8989, 0.6144], + ], + [ + [0.3930, 0.2098, 0.0602], + [0.2298, 0.9425, 0.0333], + [0.7409, 0.8172, 0.8879], + [0.4846, 0.0486, 0.2029], + [0.6741, 0.9765, 0.6864], + [0.2827, 0.5534, 0.2125], + ], + ]), + &device, + ); + + let output = module.forward(input); + + let expected = TensorData::from([ + [ + [-1.1694, -0.5353, 0.7572], + [-0.1775, -0.6838, 1.8087], + [0.5205, -1.3107, -1.0723], + [-0.0459, 0.2351, 1.6734], + [-0.5796, 1.3218, -1.6544], + [-0.2744, 1.0406, 0.1459], + ], + [ + [0.2665, -0.3320, -0.8205], + [-0.2667, 2.0612, -0.9085], + [0.6681, 0.9102, 1.1345], + [-0.1453, -1.5287, -1.0389], + [0.4253, 1.5962, 0.4731], + [-1.0903, -0.0419, -1.3623], + ], + ]); + output + .to_data() + .assert_approx_eq::(&expected, tolerance); + } + + #[test] + fn display() { + let config = GroupNormConfig::new(3, 6); + let group_norm = config.init(&Default::default()); + + assert_eq!( + format!("{group_norm}"), + "GroupNorm {num_groups: 3, num_channels: 6, epsilon: 0.00001, affine: true, params: 12}" + ); + } +} diff --git a/crates/burn-nn/src/modules/norm/instance.rs b/crates/burn-nn/src/modules/norm/instance.rs new file mode 100644 index 0000000..a64c87f --- /dev/null +++ b/crates/burn-nn/src/modules/norm/instance.rs @@ -0,0 +1,223 @@ +use burn_core as burn; + +use crate::norm::group_norm; +use burn::config::Config; +use burn::module::Initializer; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::module::{Module, Param}; +use burn::tensor::{Device, Tensor}; + +/// Configuration to create a [InstanceNorm](InstanceNorm) layer using the [init function](InstanceNormConfig::init). +#[derive(Debug, Config)] +pub struct InstanceNormConfig { + /// The number of channels expected in the input + pub num_channels: usize, + /// A value required for numerical stability. Default: 1e-5 + #[config(default = 1e-5)] + pub epsilon: f64, + /// A boolean value that when set to `true`, this module has learnable + /// per-channel affine parameters initialized to ones (for weights) + /// and zeros (for biases). Default: `true` + #[config(default = true)] + pub affine: bool, +} + +/// Applies Instance Normalization over a tensor as described in the paper [Instance Normalization](https://arxiv.org/abs/1607.08022) +/// +/// Should be created using [InstanceNormConfig](InstanceNormConfig). +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct InstanceNorm { + /// The learnable weight + pub gamma: Option>>, + /// The learnable bias + pub beta: Option>>, + /// The number of channels expected in the input + pub num_channels: usize, + /// A value required for numerical stability + pub epsilon: f64, + /// A boolean value that when set to `true`, this module has learnable + pub affine: bool, +} + +impl ModuleDisplay for InstanceNorm { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("num_channels", &self.num_channels) + .add("epsilon", &self.epsilon) + .add("affine", &self.affine) + .optional() + } +} + +impl InstanceNormConfig { + /// Initialize a new [instance norm](InstanceNorm) module. + pub fn init(&self, device: &Device) -> InstanceNorm { + let (gamma, beta) = if self.affine { + let gamma = Initializer::Ones.init([self.num_channels], device); + let beta = Initializer::Zeros.init([self.num_channels], device); + + (Some(gamma), Some(beta)) + } else { + (None, None) + }; + + InstanceNorm { + gamma, + beta, + num_channels: self.num_channels, + epsilon: self.epsilon, + affine: self.affine, + } + } +} + +impl InstanceNorm { + /// Applies the forward pass on the input tensor. + /// + /// See also [InstanceNormConfig](InstanceNormConfig) for more information. + /// + /// # Shapes + /// + /// - input: `[batch_size, num_channels, *]` + /// - output: `[batch_size, num_channels, *]` + pub fn forward(&self, input: Tensor) -> Tensor { + // Instance norm is equivalent to group norm when the number of groups is equal to the number of channels. + let num_groups = self.num_channels; + + let gamma = self.gamma.as_ref().map(|x| x.val()); + let beta = self.beta.as_ref().map(|x| x.val()); + + group_norm(input, gamma, beta, num_groups, self.epsilon, self.affine) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use alloc::format; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn instance_norm_forward_affine_false() { + let device = Default::default(); + let module = InstanceNormConfig::new(6).with_affine(false).init(&device); + + let input = Tensor::<3>::from_data( + TensorData::from([ + [ + [-0.3034, 0.2726, -0.9659], + [-1.1845, 1.4078, 0.9774], + [0.3963, -1.3738, 1.4125], + [1.0682, 0.3604, 0.3985], + [-0.4957, -0.4461, -0.9721], + [1.5157, -0.1546, -0.5596], + ], + [ + [-1.6698, -0.4040, -0.7927], + [0.3736, -0.0975, -0.1351], + [-0.9461, 0.5461, -0.6334], + [-1.0919, -0.1158, 0.1213], + [-0.9535, 0.1281, 0.4372], + [-0.2845, 0.3488, 0.5641], + ], + ]), + &device, + ); + + let output = module.forward(input); + + let expected = TensorData::from([ + [ + [0.0569, 1.1952, -1.2522], + [-1.3971, 0.8883, 0.5088], + [0.2183, -1.3192, 1.1009], + [1.4126, -0.7649, -0.6477], + [0.5999, 0.8091, -1.409], + [1.39, -0.4696, -0.9205], + ], + [ + [-1.3492, 1.0417, 0.3075], + [1.411, -0.6243, -0.7867], + [-0.9363, 1.386, -0.4497], + [-1.3899, 0.4692, 0.9208], + [-1.3822, 0.4319, 0.9503], + [-1.3714, 0.3868, 0.9846], + ], + ]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn instance_norm_forward_affine_true() { + let device = Default::default(); + let module = InstanceNormConfig::new(6).with_affine(true).init(&device); + + let input = Tensor::<3>::from_data( + TensorData::from([ + [ + [0.3345, 0.4429, 0.6639], + [0.5041, 0.4175, 0.8437], + [0.6159, 0.3758, 0.4071], + [0.5417, 0.5785, 0.7671], + [0.3837, 0.9883, 0.0420], + [0.4808, 0.8989, 0.6144], + ], + [ + [0.3930, 0.2098, 0.0602], + [0.2298, 0.9425, 0.0333], + [0.7409, 0.8172, 0.8879], + [0.4846, 0.0486, 0.2029], + [0.6741, 0.9765, 0.6864], + [0.2827, 0.5534, 0.2125], + ], + ]), + &device, + ); + + let output = module.forward(input); + + let expected = TensorData::from([ + [ + [-1.06458, -0.2738, 1.33838], + [-0.45848, -0.92929, 1.38777], + [1.40388, -0.84877, -0.55511], + [-0.88515, -0.51245, 1.3976], + [-0.22397, 1.32124, -1.09727], + [-1.05468, 1.34316, -0.28848], + ], + [ + [1.26372, -0.08229, -1.18144], + [-0.44049, 1.38403, -0.94354], + [-1.23828, 0.03109, 1.2072], + [1.32524, -1.08999, -0.23524], + [-0.75061, 1.4132, -0.66259], + [-0.45469, 1.38697, -0.93228], + ], + ]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn display() { + let config = InstanceNormConfig::new(6); + let instance_norm = config.init(&Default::default()); + + assert_eq!( + format!("{instance_norm}"), + "InstanceNorm {num_channels: 6, epsilon: 0.00001, affine: true, params: 12}" + ); + } +} diff --git a/crates/burn-nn/src/modules/norm/layer.rs b/crates/burn-nn/src/modules/norm/layer.rs new file mode 100644 index 0000000..b781e2e --- /dev/null +++ b/crates/burn-nn/src/modules/norm/layer.rs @@ -0,0 +1,239 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::Content; +use burn::module::DisplaySettings; +use burn::module::Initializer; +use burn::module::Module; +use burn::module::ModuleDisplay; +use burn::module::Param; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn::tensor::module::layer_norm; + +/// Configuration to create a [LayerNorm](LayerNorm) layer using the [init function](LayerNormConfig::init). +#[derive(Debug, Config)] +pub struct LayerNormConfig { + /// The size of the input features. + pub d_model: usize, + /// A value required for numerical stability. Default: 1e-5 + #[config(default = 1e-5)] + pub epsilon: f64, + /// If a bias (beta) should be applied during the normalization. Default: true + #[config(default = true)] + pub bias: bool, +} + +/// Applies Layer Normalization over an input tensor as described in the paper [Layer Normalization](https://arxiv.org/abs/1607.06450). +/// +/// `Y = norm(X) * γ + β` +/// +/// Where: +/// - `X` is the input tensor +/// - `Y` is the output tensor +/// - `γ` is the learnable weight (scale) +/// - `β` is the learnable bias (optional) +/// +/// Should be created using [LayerNormConfig](LayerNormConfig). +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct LayerNorm { + /// The learnable weight (scale). + pub gamma: Param>, + /// The learnable bias (optional). + pub beta: Option>>, + /// A value required for numerical stability. + epsilon: f64, +} + +impl LayerNormConfig { + /// Initialize a new [layer norm](LayerNorm) module. + pub fn init(&self, device: &Device) -> LayerNorm { + let gamma = Initializer::Ones.init([self.d_model], device); + let beta = if self.bias { + Some(Initializer::Zeros.init([self.d_model], device)) + } else { + None + }; + + LayerNorm { + gamma, + beta, + epsilon: self.epsilon, + } + } +} + +impl LayerNorm { + /// Applies the forward pass on the input tensor. + /// + /// See the [LayerNorm](LayerNorm) documentation for more information. + /// + /// # Shapes + /// + /// - input: `[..., any, d_model]` + /// - output: `[..., any, d_model]` + pub fn forward(&self, input: Tensor) -> Tensor { + layer_norm( + input, + self.gamma.val(), + self.beta.as_ref().map(|b| b.val()), + self.epsilon, + ) + } +} + +impl ModuleDisplay for LayerNorm { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let [d_model] = self.gamma.shape().dims(); + content + .add("d_model", &d_model) + .add("epsilon", &self.epsilon) + .add("bias", &self.beta.is_some()) + .optional() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use alloc::format; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn layer_norm_forward() { + let device = Default::default(); + let module = LayerNormConfig::new(10).init(&device); + let input = Tensor::<2>::from_data( + TensorData::from([[ + -0.6897, -2.7106, 2.2222, -1.0330, -0.8933, 1.1765, 0.0601, 1.5252, -0.3630, 0.6728, + ]]), + &device, + ); + + let output = module.forward(input); + + let expected = TensorData::from([[ + -0.4990, -1.9680, 1.6178, -0.7486, -0.6470, 0.8576, 0.0461, 1.1111, -0.2614, 0.4915, + ]]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn layer_norm_forward_large_epsilon() { + let device = Default::default(); + let module = LayerNormConfig::new(10).with_epsilon(1e-1).init(&device); + let input = Tensor::<2>::from_data( + TensorData::from([[ + -0.6897, -2.7106, 2.2222, -1.0330, -0.8933, 1.1765, 0.0601, 1.5252, -0.3630, 0.6728, + ]]), + &device, + ); + + let output = module.forward(input); + + let expected = TensorData::from([[ + -0.4863, -1.9180, 1.5766, -0.7295, -0.6305, 0.8358, 0.0449, 1.0828, -0.2548, 0.4790, + ]]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn layer_norm_forward_no_bias() { + let device = Default::default(); + let module = LayerNormConfig::new(10).with_bias(false).init(&device); + let input = Tensor::<2>::from_data( + TensorData::from([[ + -0.6897, -2.7106, 2.2222, -1.0330, -0.8933, 1.1765, 0.0601, 1.5252, -0.3630, 0.6728, + ]]), + &device, + ); + + let output = module.forward(input); + + // With bias=false, output matches the bias=true case (beta is zero-initialized + // by default), confirming the `None` branch in the backend hook produces the + // pre-beta result. + let expected = TensorData::from([[ + -0.4990, -1.9680, 1.6178, -0.7486, -0.6470, 0.8576, 0.0461, 1.1111, -0.2614, 0.4915, + ]]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[cfg(feature = "std")] + #[test] + fn layer_norm_backward() { + let device = Device::default().autodiff(); + let module = LayerNormConfig::new(2).init(&device); + let tensor_1 = Tensor::<2>::from_data(TensorData::from([[0.0, 1.0], [3.0, 4.0]]), &device) + .require_grad(); + let tensor_2 = Tensor::<2>::from_data(TensorData::from([[6.0, 7.0], [9.0, 10.0]]), &device) + .require_grad(); + + let x = tensor_1.clone().matmul(tensor_2.clone()); + + let output = module.forward(x); + let grads = output.backward(); + + let tensor_1_grad = tensor_1.grad(&grads).unwrap(); + let tensor_2_grad = tensor_2.grad(&grads).unwrap(); + let gamma_grad = module.gamma.grad(&grads).unwrap(); + let beta_grad = module.beta.as_ref().unwrap().grad(&grads).unwrap(); + + let expected = TensorData::from([-2.0, 2.0]); + gamma_grad + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::from([2.0, 2.0]); + beta_grad + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::zeros::(tensor_1_grad.shape()); + tensor_1_grad + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + + let expected = TensorData::zeros::(tensor_2_grad.shape()); + tensor_2_grad + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn display() { + let config = LayerNormConfig::new(6); + let layer_norm = config.init(&Default::default()); + + assert_eq!( + format!("{layer_norm}"), + "LayerNorm {d_model: 6, epsilon: 0.00001, bias: true, params: 12}" + ); + } + + #[test] + fn display_no_bias() { + let config = LayerNormConfig::new(6).with_bias(false); + let layer_norm = config.init(&Default::default()); + + assert_eq!( + format!("{layer_norm}"), + "LayerNorm {d_model: 6, epsilon: 0.00001, bias: false, params: 6}" + ); + } +} diff --git a/crates/burn-nn/src/modules/norm/local_response.rs b/crates/burn-nn/src/modules/norm/local_response.rs new file mode 100644 index 0000000..1048bab --- /dev/null +++ b/crates/burn-nn/src/modules/norm/local_response.rs @@ -0,0 +1,678 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Module, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::tensor::module::avg_pool1d; +use burn::tensor::ops::PadMode; + +/// Configuration to create a [LocalResponseNorm](LocalResponseNorm) layer +/// using the [init function](LocalResponseNormConfig::init). +#[derive(Config, Debug)] +pub struct LocalResponseNormConfig { + /// Number of channels in the sliding normalization window. + pub size: usize, + /// Scaling parameter. Default: 0.0001 + #[config(default = 0.0001)] + pub alpha: f64, + /// Exponent. Default: 0.75 + #[config(default = 0.75)] + pub beta: f64, + /// Bias constant (called `bias` in ONNX). Default: 1.0 + #[config(default = 1.0)] + pub k: f64, +} + +impl LocalResponseNormConfig { + /// Initialize a new [LocalResponseNorm](LocalResponseNorm) module. + /// + /// # Panics + /// + /// Panics if `size` is 0. + pub fn init(&self) -> LocalResponseNorm { + assert!(self.size > 0, "size must be greater than 0."); + + LocalResponseNorm { + size: self.size, + alpha: self.alpha, + beta: self.beta, + k: self.k, + } + } +} + +/// Applies Local Response Normalization as described in +/// [ImageNet Classification with Deep Convolutional Neural Networks](https://papers.nips.cc/paper/2012/hash/c399862d3b9d6b76c8436e924a68c45b-Abstract.html). +/// +/// `Y = X / (k + (alpha / size) * sum(X^2))^beta` +/// +/// Where the sum is computed over a sliding window of `size` channels. +/// +/// For odd `size`, the window is centered on each channel position. +/// For even `size`, the window uses asymmetric padding and includes the current +/// channel plus one extra channel on the positive side. +/// +/// Should be created using [LocalResponseNormConfig](LocalResponseNormConfig). +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct LocalResponseNorm { + /// Number of channels in the sliding window. + size: usize, + /// Scaling parameter. + alpha: f64, + /// Exponent. + beta: f64, + /// Bias constant. + k: f64, +} + +impl LocalResponseNorm { + /// Applies Local Response Normalization on the input tensor. + /// + /// # Shapes + /// + /// - input: `[N, C, D1, D2, ..., Dk]` (rank >= 3) + /// - output: same shape as input + /// + /// # Panics + /// + /// Panics if the input tensor rank is less than 3. + pub fn forward(&self, input: Tensor) -> Tensor { + assert!( + D >= 3, + "LocalResponseNorm requires input rank >= 3, got {D}" + ); + + let shape = input.dims(); + let n = shape[0]; + let c = shape[1]; + let d_flat: usize = shape[2..].iter().product(); + + // Square the input + let squared = input.clone().square(); + + // Flatten spatial dims: [N, C, D1..Dk] -> [N, C, D_flat] + let flat: Tensor<3> = squared.reshape([n, c, d_flat]); + + // Move channel to last dim: [N, D_flat, C] + let transposed = flat.swap_dims(1, 2); + + // Batch all spatial positions: [N*D_flat, 1, C] + let batched: Tensor<3> = transposed.reshape([n * d_flat, 1, c]); + + let pad_left = (self.size - 1) / 2; + let pad_right = self.size / 2; + let square_avg = if pad_left != pad_right { + let padded = batched.pad((pad_left, pad_right, 0, 0), PadMode::Constant(0.0)); + avg_pool1d(padded, self.size, 1, 0, true, false) + } else { + avg_pool1d(batched, self.size, 1, pad_left, true, false) + }; + + // Restore shape: [N*D_flat, 1, C] -> [N, D_flat, C] -> [N, C, D_flat] -> original + let unbatched: Tensor<3> = square_avg.reshape([n, d_flat, c]); + let untransposed = unbatched.swap_dims(1, 2); + let square_avg_restored: Tensor = untransposed.reshape(shape); + + // Apply LRN formula: output = input / (k + alpha * avg(x^2))^beta + input + / square_avg_restored + .mul_scalar(self.alpha) + .add_scalar(self.k) + .powf_scalar(self.beta) + } +} + +impl ModuleDisplay for LocalResponseNorm { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("size", &self.size) + .add("alpha", &self.alpha) + .add("beta", &self.beta) + .add("k", &self.k) + .optional() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use alloc::format; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + + type FT = f32; + + // --- Correctness tests (values from PyTorch, torch.manual_seed(42)) --- + + #[test] + fn forward_default_params() { + // size=5, alpha=0.0001, beta=0.75, k=1.0, input [1,3,4,4] + let device = Default::default(); + let module = LocalResponseNormConfig::new(5).init(); + let input = Tensor::<4>::from_data( + TensorData::from([[ + [ + [1.9269, 1.4873, 0.9007, -2.1055], + [0.6784, -1.2345, -0.0431, -1.6047], + [-0.7521, 1.6487, -0.3925, -1.4036], + [-0.7279, -0.5594, -0.7688, 0.7624], + ], + [ + [1.6423, -0.1596, -0.4974, 0.4396], + [-0.7581, 1.0783, 0.8008, 1.6806], + [1.2791, 1.2964, 0.6105, 1.3347], + [-0.2316, 0.0418, -0.2516, 0.8599], + ], + [ + [-1.3847, -0.8712, -0.2234, 1.7174], + [0.3189, -0.4245, 0.3057, -0.7746], + [-1.5576, 0.9956, -0.8798, -0.6011], + [-1.2742, 2.1228, -1.2347, -0.4879], + ], + ]]), + &device, + ); + + let output = module.forward(input); + + let expected = TensorData::from([[ + [ + [1.9267, 1.4872, 0.9007, -2.1053], + [0.6784, -1.2345, -0.0431, -1.6045], + [-0.7521, 1.6486, -0.3925, -1.4035], + [-0.7279, -0.5594, -0.7688, 0.7624], + ], + [ + [1.6421, -0.1596, -0.4974, 0.4395], + [-0.7581, 1.0783, 0.8008, 1.6805], + [1.2790, 1.2963, 0.6105, 1.3347], + [-0.2316, 0.0418, -0.2516, 0.8598], + ], + [ + [-1.3845, -0.8712, -0.2234, 1.7172], + [0.3189, -0.4245, 0.3057, -0.7745], + [-1.5575, 0.9956, -0.8798, -0.6011], + [-1.2741, 2.1226, -1.2346, -0.4879], + ], + ]]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(5e-3, 1e-4)); + } + + #[test] + fn forward_custom_params() { + // size=3, alpha=0.001, beta=0.5, k=2.0, input [1,4,3,3] + let device = Default::default(); + let module = LocalResponseNormConfig::new(3) + .with_alpha(0.001) + .with_beta(0.5) + .with_k(2.0) + .init(); + let input = Tensor::<4>::from_data( + TensorData::from([[ + [ + [1.9269, 1.4873, 0.9007], + [-2.1055, 0.6784, -1.2345], + [-0.0431, -1.6047, -0.7521], + ], + [ + [1.6487, -0.3925, -1.4036], + [-0.7279, -0.5594, -0.7688], + [0.7624, 1.6423, -0.1596], + ], + [ + [-0.4974, 0.4396, 0.3189], + [-0.4245, 0.3057, -0.7746], + [0.0349, 0.3211, 1.5736], + ], + [ + [-0.8455, -1.2742, 2.1228], + [-1.2347, -0.4879, -1.4181], + [0.8963, 0.0499, 2.2667], + ], + ]]), + &device, + ); + + let output = module.forward(input); + + let expected = TensorData::from([[ + [ + [1.3618, 1.0515, 0.6368], + [-1.4882, 0.4797, -0.8728], + [-0.0305, -1.1342, -0.5318], + ], + [ + [1.1652, -0.2775, -0.9923], + [-0.5145, -0.3955, -0.5435], + [0.5391, 1.1608, -0.1128], + ], + [ + [-0.3516, 0.3108, 0.2254], + [-0.3001, 0.2162, -0.5476], + [0.0247, 0.2270, 1.1120], + ], + [ + [-0.5978, -0.9008, 1.5005], + [-0.8729, -0.3450, -1.0025], + [0.6337, 0.0353, 1.6018], + ], + ]]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(5e-3, 1e-4)); + } + + #[test] + fn forward_even_size() { + // size=2, alpha=0.0001, beta=0.75, k=1.0, input [1,3,2,2] + let device = Default::default(); + let module = LocalResponseNormConfig::new(2).init(); + let input = Tensor::<4>::from_data( + TensorData::from([[ + [[0.3367, 0.1288], [0.2345, 0.2303]], + [[-1.1229, -0.1863], [2.2082, -0.6380]], + [[0.4617, 0.2674], [0.5349, 0.8094]], + ]]), + &device, + ); + + let output = module.forward(input); + + let expected = TensorData::from([[ + [[0.3367, 0.1288], [0.2345, 0.2303]], + [[-1.1228, -0.1863], [2.2078, -0.6380]], + [[0.4616, 0.2673], [0.5348, 0.8093]], + ]]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(5e-3, 1e-4)); + } + + #[test] + fn forward_even_size_uses_asymmetric_positive_side_window() { + let device = Default::default(); + let module = LocalResponseNormConfig::new(2) + .with_alpha(1.0) + .with_beta(1.0) + .with_k(0.0) + .init(); + let input = Tensor::<3>::from_data(TensorData::from([[[1.0], [2.0], [4.0]]]), &device); + + let output = module.forward(input); + + // For size=2, the implementation pads asymmetrically and uses: + // c0 -> avg([1^2, 2^2]) = 2.5 + // c1 -> avg([2^2, 4^2]) = 10.0 + // c2 -> avg([4^2, 0]) = 8.0 + let expected = TensorData::from([[[0.4], [0.2], [0.5]]]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(1e-5, 1e-6)); + } + + #[test] + fn forward_3d() { + // size=3, input [1,4,6] + let device = Default::default(); + let module = LocalResponseNormConfig::new(3).init(); + let input = Tensor::<3>::from_data( + TensorData::from([[ + [1.9269, 1.4873, 0.9007, -2.1055, 0.6784, -1.2345], + [-0.0431, -1.6047, 0.3559, -0.6866, -0.4934, 0.2415], + [-1.1109, 0.0915, -2.3169, -0.2168, -0.3097, -0.3957], + [0.8034, -0.6216, -0.5920, -0.0631, -0.8286, 0.3309], + ]]), + &device, + ); + + let output = module.forward(input); + + let expected = TensorData::from([[ + [1.9267, 1.4871, 0.9007, -2.1053, 0.6784, -1.2345], + [-0.0431, -1.6045, 0.3558, -0.6865, -0.4933, 0.2415], + [-1.1109, 0.0915, -2.3166, -0.2168, -0.3097, -0.3957], + [0.8034, -0.6216, -0.5919, -0.0631, -0.8285, 0.3309], + ]]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(5e-3, 1e-4)); + } + + #[test] + fn forward_5d() { + // size=3, input [1,3,2,2,2] + let device = Default::default(); + let module = LocalResponseNormConfig::new(3).init(); + let input = Tensor::<5>::from_data( + TensorData::from([[ + [ + [[1.9269, 1.4873], [0.9007, -2.1055]], + [[0.6784, -1.2345], [-0.0431, -1.6047]], + ], + [ + [[0.3559, -0.6866], [-0.4934, 0.2415]], + [[-1.1109, 0.0915], [-2.3169, -0.2168]], + ], + [ + [[-0.3097, -0.3957], [0.8034, -0.6216]], + [[-0.5920, -0.0631], [-0.8286, 0.3309]], + ], + ]]), + &device, + ); + + let output = module.forward(input); + + let expected = TensorData::from([[ + [ + [[1.9267, 1.4872], [0.9007, -2.1053]], + [[0.6784, -1.2345], [-0.0431, -1.6046]], + ], + [ + [[0.3558, -0.6866], [-0.4933, 0.2415]], + [[-1.1108, 0.0915], [-2.3166, -0.2168]], + ], + [ + [[-0.3097, -0.3957], [0.8034, -0.6216]], + [[-0.5920, -0.0631], [-0.8284, 0.3309]], + ], + ]]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(5e-3, 1e-4)); + } + + // --- Edge case tests --- + + #[test] + fn forward_size_1() { + // size=1: window covers only self-channel, input [1,3,3,3] + let device = Default::default(); + let module = LocalResponseNormConfig::new(1).init(); + let input = Tensor::<4>::from_data( + TensorData::from([[ + [ + [1.9269, 1.4873, 0.9007], + [-2.1055, 0.6784, -1.2345], + [-0.0431, -1.6047, -0.7521], + ], + [ + [1.6487, -0.3925, 0.2415], + [-1.1109, 0.0915, -2.3169], + [-0.2168, -1.3847, -0.8712], + ], + [ + [-0.2234, -0.6216, -0.5920], + [-0.0631, -0.8286, 0.3309], + [-1.5576, 0.9956, -0.8798], + ], + ]]), + &device, + ); + + let output = module.forward(input); + + let expected = TensorData::from([[ + [ + [1.9264, 1.4870, 0.9007], + [-2.1048, 0.6784, -1.2344], + [-0.0431, -1.6044, -0.7521], + ], + [ + [1.6484, -0.3925, 0.2415], + [-1.1108, 0.0915, -2.3160], + [-0.2168, -1.3845, -0.8712], + ], + [ + [-0.2234, -0.6216, -0.5920], + [-0.0631, -0.8285, 0.3309], + [-1.5573, 0.9956, -0.8797], + ], + ]]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(5e-3, 1e-4)); + } + + #[test] + fn forward_c_less_than_size() { + // C=2 < size=5, input [1,2,3,3] + let device = Default::default(); + let module = LocalResponseNormConfig::new(5).init(); + let input = Tensor::<4>::from_data( + TensorData::from([[ + [ + [1.9269, 1.4873, -0.4974], + [0.4396, -0.7581, 1.0783], + [0.8008, 1.6806, 0.3559], + ], + [ + [-0.6866, 0.6105, 1.3347], + [-0.2316, 0.0418, -0.2516], + [0.8599, -0.3097, -0.3957], + ], + ]]), + &device, + ); + + let output = module.forward(input); + + let expected = TensorData::from([[ + [ + [1.9268, 1.4872, -0.4974], + [0.4396, -0.7581, 1.0783], + [0.8008, 1.6805, 0.3559], + ], + [ + [-0.6866, 0.6104, 1.3347], + [-0.2316, 0.0418, -0.2516], + [0.8598, -0.3097, -0.3957], + ], + ]]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::rel_abs(5e-3, 1e-4)); + } + + #[test] + fn forward_multi_batch() { + // N=2, size=5, input [2,3,4,4] + let device = Default::default(); + let module = LocalResponseNormConfig::new(5).init(); + let input = Tensor::<4>::from_data( + TensorData::from([ + [ + [ + [1.9269, 1.4873, 0.9007, -2.1055], + [0.6784, -1.2345, -0.0431, -1.6047], + [-0.7521, 1.6487, -0.3925, -1.4036], + [-0.7279, -0.5594, -0.7688, 0.7624], + ], + [ + [1.6423, -0.1596, -0.4974, 0.4396], + [-0.7581, 1.0783, 0.8008, 1.6806], + [1.2791, 1.2964, 0.6105, 1.3347], + [-0.2316, 0.0418, -0.2516, 0.8599], + ], + [ + [-1.3847, -0.8712, -0.2234, 1.7174], + [0.3189, -0.4245, 0.3057, -0.7746], + [-1.5576, 0.9956, -0.8798, -0.6011], + [-1.2742, 2.1228, -1.2347, -0.4879], + ], + ], + [ + [ + [-0.9138, -0.6581, 0.0780, 0.5258], + [-0.4880, 1.1914, -0.8140, -0.7360], + [-1.4032, 0.0360, -0.0635, 0.6756], + [-0.0978, 1.8446, -1.1845, 1.3835], + ], + [ + [1.4451, 0.8564, 2.2181, 0.5232], + [0.3466, -0.1973, -1.0546, 1.2780], + [-0.1722, 0.5238, 0.0566, 0.4263], + [0.5750, -0.6417, -2.2064, -0.7508], + ], + [ + [0.0109, -0.3387, -1.3407, -0.5854], + [0.5362, 0.5246, 1.1412, 0.0516], + [0.7440, -0.4816, -1.0495, 0.6039], + [-1.7223, -0.8278, 1.3347, 0.4835], + ], + ], + ]), + &device, + ); + + let output = module.forward(input); + + let out_data = output.to_data(); + assert_eq!(out_data.shape, [2, 3, 4, 4].into()); + let expected_full = TensorData::from([ + [ + [ + [1.9267, 1.4872, 0.9007, -2.1053], + [0.6784, -1.2345, -0.0431, -1.6045], + [-0.7521, 1.6486, -0.3925, -1.4035], + [-0.7279, -0.5594, -0.7688, 0.7624], + ], + [ + [1.6421, -0.1596, -0.4974, 0.4395], + [-0.7581, 1.0783, 0.8008, 1.6805], + [1.2790, 1.2963, 0.6105, 1.3347], + [-0.2316, 0.0418, -0.2516, 0.8598], + ], + [ + [-1.3845, -0.8712, -0.2234, 1.7172], + [0.3189, -0.4245, 0.3057, -0.7745], + [-1.5575, 0.9956, -0.8798, -0.6011], + [-1.2741, 2.1226, -1.2346, -0.4879], + ], + ], + [ + [ + [-0.9138, -0.6581, 0.0780, 0.5258], + [-0.4880, 1.1913, -0.8140, -0.7360], + [-1.4032, 0.0360, -0.0635, 0.6756], + [-0.0978, 1.8445, -1.1844, 1.3835], + ], + [ + [1.4451, 0.8564, 2.2179, 0.5232], + [0.3466, -0.1973, -1.0545, 1.2780], + [-0.1722, 0.5238, 0.0566, 0.4263], + [0.5750, -0.6417, -2.2061, -0.7508], + ], + [ + [0.0109, -0.3387, -1.3405, -0.5854], + [0.5362, 0.5246, 1.1411, 0.0516], + [0.7439, -0.4816, -1.0494, 0.6039], + [-1.7222, -0.8277, 1.3345, 0.4835], + ], + ], + ]); + out_data.assert_approx_eq::(&expected_full, Tolerance::rel_abs(5e-3, 1e-4)); + } + + // --- Validation / panic tests --- + + #[test] + #[should_panic(expected = "size must be greater than 0")] + fn config_size_zero_panics() { + LocalResponseNormConfig::new(0).init(); + } + + #[test] + #[should_panic(expected = "LocalResponseNorm requires input rank >= 3")] + fn forward_rank_2_panics() { + let module = LocalResponseNormConfig::new(3).init(); + let input = Tensor::<2>::zeros([2, 4], &Default::default()); + module.forward(input); + } + + // --- Autodiff --- + + #[cfg(feature = "std")] + #[test] + fn backward() { + use burn_core::tensor::Device; + + let device = Device::default().autodiff(); + let module = LocalResponseNormConfig::new(5).init(); + let input = Tensor::<4>::from_data( + TensorData::from([[ + [ + [1.9269, 1.4873, 0.9007, -2.1055], + [0.6784, -1.2345, -0.0431, -1.6047], + [-0.7521, 1.6487, -0.3925, -1.4036], + [-0.7279, -0.5594, -0.7688, 0.7624], + ], + [ + [1.6423, -0.1596, -0.4974, 0.4396], + [-0.7581, 1.0783, 0.8008, 1.6806], + [1.2791, 1.2964, 0.6105, 1.3347], + [-0.2316, 0.0418, -0.2516, 0.8599], + ], + [ + [-1.3847, -0.8712, -0.2234, 1.7174], + [0.3189, -0.4245, 0.3057, -0.7746], + [-1.5576, 0.9956, -0.8798, -0.6011], + [-1.2742, 2.1228, -1.2347, -0.4879], + ], + ]]), + &device, + ) + .require_grad(); + + let output = module.forward(input.clone()); + let grads = output.sum().backward(); + let input_grad = input.grad(&grads).unwrap(); + + assert_eq!(input_grad.dims(), [1, 3, 4, 4]); + + let expected_grad = TensorData::from([[ + [ + [0.9997, 0.9999, 1.0000, 0.9999], + [1.0000, 0.9999, 1.0000, 0.9999], + [0.9999, 0.9997, 1.0000, 0.9999], + [0.9999, 1.0000, 0.9999, 1.0000], + ], + [ + [0.9998, 1.0000, 1.0000, 0.9999], + [1.0000, 1.0000, 1.0000, 0.9999], + [1.0000, 0.9998, 1.0000, 1.0000], + [1.0000, 0.9999, 1.0000, 0.9999], + ], + [ + [1.0000, 1.0000, 1.0000, 0.9999], + [1.0000, 1.0000, 1.0000, 0.9999], + [0.9999, 0.9998, 1.0000, 0.9999], + [0.9999, 0.9998, 0.9999, 1.0000], + ], + ]]); + input_grad + .to_data() + .assert_approx_eq::(&expected_grad, Tolerance::rel_abs(5e-3, 1e-4)); + } + + // --- Display --- + + #[test] + fn display() { + let config = LocalResponseNormConfig::new(5); + let module = config.init(); + assert_eq!( + format!("{module}"), + "LocalResponseNorm {size: 5, alpha: 0.0001, beta: 0.75, k: 1}" + ); + } +} diff --git a/crates/burn-nn/src/modules/norm/mod.rs b/crates/burn-nn/src/modules/norm/mod.rs new file mode 100644 index 0000000..48e0777 --- /dev/null +++ b/crates/burn-nn/src/modules/norm/mod.rs @@ -0,0 +1,30 @@ +//! # Normalization Layers +//! +//! Users who wish to provide an abstraction over swappable normalization +//! layers can use the [`Normalization`] wrapper, with support for: +//! * [`Normalization::Batch`] - [`BatchNorm`] +//! * [`Normalization::Group`] - [`GroupNorm`] +//! * [`Normalization::Instance`] - [`InstanceNorm`] +//! * [`Normalization::Layer`] - [`LayerNorm`] +//! * [`Normalization::Rms`] - [`RmsNorm`] +//! +//! [`NormalizationConfig`] can be used as a generic normalization policy: +//! * Construct a config with arbitrary input features (we suggest `0`). +//! * Clone and match that config to the target input layer, +//! using the [`NormalizationConfig::with_num_features()`] method. +pub(crate) mod batch; +pub(crate) mod group; +pub(crate) mod instance; +pub(crate) mod layer; +pub(crate) mod local_response; +pub(crate) mod rms; + +mod normalization_wrapper; + +pub use batch::*; +pub use group::*; +pub use instance::*; +pub use layer::*; +pub use local_response::*; +pub use normalization_wrapper::*; +pub use rms::*; diff --git a/crates/burn-nn/src/modules/norm/normalization_wrapper.rs b/crates/burn-nn/src/modules/norm/normalization_wrapper.rs new file mode 100644 index 0000000..e6450b2 --- /dev/null +++ b/crates/burn-nn/src/modules/norm/normalization_wrapper.rs @@ -0,0 +1,397 @@ +use burn_core as burn; + +use crate::{ + BatchNorm, BatchNormConfig, GroupNorm, GroupNormConfig, Identity, InstanceNorm, + InstanceNormConfig, LayerNorm, LayerNormConfig, RmsNorm, RmsNormConfig, +}; +use burn::prelude::{Config, Module}; +use burn::tensor::Device; +use burn::tensor::Tensor; + +/// ['Normalization'] Configuration. +/// +/// The enum is non-exhaustive to prepare for future additions. +/// +/// Can be used as a generic configuration for normalization layers: +/// * Construct a config with arbitrary input features (we suggest `0`). +/// * Clone and match that config to the target input layer, +/// using the [`NormalizationConfig::with_num_features()`] method. +#[derive(Config, Debug)] +#[non_exhaustive] +pub enum NormalizationConfig { + /// ['Identity'] Configuration. + Identity, + + /// ['BatchNorm'] Configuration. + Batch(BatchNormConfig), + + /// ['GroupNorm'] Configuration. + Group(GroupNormConfig), + + /// ['InstanceNorm'] Configuration. + Instance(InstanceNormConfig), + + /// ['LayerNorm'] Configuration. + Layer(LayerNormConfig), + + /// ['RmsNorm'] Configuration. + Rms(RmsNormConfig), +} + +impl From for NormalizationConfig { + fn from(config: BatchNormConfig) -> Self { + Self::Batch(config) + } +} + +impl From for NormalizationConfig { + fn from(config: GroupNormConfig) -> Self { + Self::Group(config) + } +} + +impl From for NormalizationConfig { + fn from(config: InstanceNormConfig) -> Self { + Self::Instance(config) + } +} + +impl From for NormalizationConfig { + fn from(config: LayerNormConfig) -> Self { + Self::Layer(config) + } +} + +impl From for NormalizationConfig { + fn from(config: RmsNormConfig) -> Self { + Self::Rms(config) + } +} + +impl NormalizationConfig { + /// Initialize a ['Norm'] layer. + pub fn init(&self, device: &Device) -> Normalization { + match self { + NormalizationConfig::Identity => Identity::new().into(), + NormalizationConfig::Batch(config) => config.init(device).into(), + NormalizationConfig::Group(config) => config.init(device).into(), + NormalizationConfig::Instance(config) => config.init(device).into(), + NormalizationConfig::Layer(config) => config.init(device).into(), + NormalizationConfig::Rms(config) => config.init(device).into(), + } + } + + /// Set the number of features. + pub fn with_num_features(self, num_features: usize) -> Self { + match self { + NormalizationConfig::Identity => self, + NormalizationConfig::Batch(config) => BatchNormConfig { + num_features, + ..config + } + .into(), + NormalizationConfig::Group(config) => GroupNormConfig { + num_channels: num_features, + ..config + } + .into(), + NormalizationConfig::Instance(config) => InstanceNormConfig { + num_channels: num_features, + ..config + } + .into(), + NormalizationConfig::Layer(config) => LayerNormConfig { + d_model: num_features, + ..config + } + .into(), + NormalizationConfig::Rms(config) => RmsNormConfig { + d_model: num_features, + ..config + } + .into(), + } + } + + /// Get the number of features. + pub fn num_features(&self) -> usize { + match self { + NormalizationConfig::Identity => 0, + NormalizationConfig::Batch(config) => config.num_features, + NormalizationConfig::Group(config) => config.num_channels, + NormalizationConfig::Instance(config) => config.num_channels, + NormalizationConfig::Layer(config) => config.d_model, + NormalizationConfig::Rms(config) => config.d_model, + } + } +} + +/// Normalization Layer Wrapper +/// +/// Provides support for built-in ``burn::nn::norm`` norm layers: +/// * [`Normalization::Batch`] - [`BatchNorm`] +/// * [`Normalization::Group`] - [`GroupNorm`] +/// * [`Normalization::Instance`] - [`InstanceNorm`] +/// * [`Normalization::Layer`] - [`LayerNorm`] +/// * [`Normalization::Rms`] - [`RmsNorm`] +/// +/// The enum is non-exhaustive, to prepare for future additions. +#[derive(Module, Debug)] +#[non_exhaustive] +pub enum Normalization { + /// ['Identity'] layer. + Identity(Identity), + + /// [`BatchNorm`] layer. + Batch(BatchNorm), + + /// [`GroupNorm`] layer. + Group(GroupNorm), + + /// ['InstanceNorm'] layer. + Instance(InstanceNorm), + + /// [`LayerNorm`] layer. + Layer(LayerNorm), + + /// ['RmsNorm'] layer. + Rms(RmsNorm), +} + +impl From for Normalization { + fn from(layer: Identity) -> Self { + Self::Identity(layer) + } +} + +impl From for Normalization { + fn from(layer: BatchNorm) -> Self { + Self::Batch(layer) + } +} + +impl From for Normalization { + fn from(layer: GroupNorm) -> Self { + Self::Group(layer) + } +} + +impl From for Normalization { + fn from(layer: InstanceNorm) -> Self { + Self::Instance(layer) + } +} + +impl From for Normalization { + fn from(layer: LayerNorm) -> Self { + Self::Layer(layer) + } +} + +impl From for Normalization { + fn from(layer: RmsNorm) -> Self { + Self::Rms(layer) + } +} + +impl Normalization { + /// Applies normalization to a tensor. + /// + /// The normalization contract depends upon the wrapped norm layer; + /// but all norm layers assume an input of at least rank 2; + /// and produce an output of the same rank and shape. + pub fn forward(&self, input: Tensor) -> Tensor { + match self { + Normalization::Identity(norm) => norm.forward(input), + Normalization::Batch(norm) => norm.forward(input), + Normalization::Group(norm) => norm.forward(input), + Normalization::Instance(norm) => norm.forward(input), + Normalization::Layer(norm) => norm.forward(input), + Normalization::Rms(norm) => norm.forward(input), + } + } + + /// Get the number of features. + pub fn num_features(&self) -> usize { + match self { + Normalization::Identity(_) => 0, + Normalization::Batch(norm) => norm.gamma.shape()[0], + Normalization::Group(norm) => norm.num_channels, + Normalization::Instance(norm) => norm.num_channels, + Normalization::Layer(norm) => norm.gamma.shape()[0], + Normalization::Rms(norm) => norm.gamma.shape()[0], + } + } +} + +#[cfg(feature = "std")] +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn test_match_feature_size() { + let config: NormalizationConfig = BatchNormConfig::new(0).into(); + assert_eq!(config.num_features(), 0); + let config = config.with_num_features(12); + assert_eq!(config.num_features(), 12); + + let config: NormalizationConfig = GroupNormConfig::new(4, 0).into(); + assert_eq!(config.num_features(), 0); + let config = config.with_num_features(12); + assert_eq!(config.num_features(), 12); + + let config: NormalizationConfig = InstanceNormConfig::new(0).into(); + assert_eq!(config.num_features(), 0); + let config = config.with_num_features(12); + assert_eq!(config.num_features(), 12); + + let config: NormalizationConfig = LayerNormConfig::new(0).into(); + assert_eq!(config.num_features(), 0); + let config = config.with_num_features(12); + assert_eq!(config.num_features(), 12); + + let config: NormalizationConfig = RmsNormConfig::new(0).into(); + assert_eq!(config.num_features(), 0); + let config = config.with_num_features(12); + assert_eq!(config.num_features(), 12); + } + + #[test] + fn test_identity_norm() { + let device = Device::default().autodiff(); + + let num_features = 12; + let input: Tensor<4> = Tensor::ones([2, num_features, 3, 4], &device); + + let config: NormalizationConfig = NormalizationConfig::Identity; + + let layer = config.init(&device); + assert_eq!(layer.num_features(), 0); + + let expected = input.clone(); + let output = layer.forward(input); + + output.to_data().assert_eq(&expected.to_data(), true); + } + + #[test] + fn test_batch_norm() { + let device = Device::default().autodiff(); + + let num_features = 12; + let input: Tensor<4> = Tensor::ones([2, num_features, 3, 4], &device); + + let config: NormalizationConfig = BatchNormConfig::new(12).into(); + + let layer = config.init(&device); + assert_eq!(layer.num_features(), 12); + + let expected = match &layer { + Normalization::Batch(inner) => inner.forward(input.clone()), + _ => panic!("Unexpected layer type"), + }; + + let output = layer.forward(input); + + output.to_data().assert_eq(&expected.to_data(), true); + } + + #[test] + fn test_group_norm() { + let device = Device::default().autodiff(); + + let num_features = 12; + let input: Tensor<4> = Tensor::ones([2, num_features, 3, 4], &device); + + let config: NormalizationConfig = GroupNormConfig::new(3, num_features).into(); + + let layer = config.init(&device); + assert_eq!(layer.num_features(), 12); + + let expected = match &layer { + Normalization::Group(inner) => inner.forward(input.clone()), + _ => panic!("Unexpected layer type"), + }; + + let output = layer.forward(input); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::default()); + } + + #[test] + fn test_instance_norm() { + let device = Device::default().autodiff(); + + let num_features = 12; + let input: Tensor<4> = Tensor::ones([2, num_features, 3, 4], &device); + + let config: NormalizationConfig = InstanceNormConfig::new(num_features).into(); + + let layer = config.init(&device); + assert_eq!(layer.num_features(), 12); + + let expected = match &layer { + Normalization::Instance(inner) => inner.forward(input.clone()), + _ => panic!("Unexpected layer type"), + }; + + let output = layer.forward(input); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::default()); + } + + #[test] + fn test_layer_norm() { + let device = Device::default().autodiff(); + + let num_features = 12; + let input: Tensor<4> = Tensor::ones([2, 3, 4, num_features], &device); + + let config: NormalizationConfig = LayerNormConfig::new(num_features).into(); + + let layer = config.init(&device); + assert_eq!(layer.num_features(), 12); + + let expected = match &layer { + Normalization::Layer(inner) => inner.forward(input.clone()), + _ => panic!("Unexpected layer type"), + }; + + let output = layer.forward(input); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::default()); + } + + #[test] + fn test_rms_norm() { + let device = Device::default().autodiff(); + + let num_features = 12; + let input: Tensor<4> = Tensor::ones([2, 3, 4, num_features], &device); + + let config: NormalizationConfig = RmsNormConfig::new(num_features).into(); + + let layer = config.init(&device); + assert_eq!(layer.num_features(), 12); + + let expected = match &layer { + Normalization::Rms(inner) => inner.forward(input.clone()), + _ => panic!("Unexpected layer type"), + }; + + let output = layer.forward(input); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::default()); + } +} diff --git a/crates/burn-nn/src/modules/norm/rms.rs b/crates/burn-nn/src/modules/norm/rms.rs new file mode 100644 index 0000000..462c378 --- /dev/null +++ b/crates/burn-nn/src/modules/norm/rms.rs @@ -0,0 +1,131 @@ +use burn::tensor::DType; + +use burn_core as burn; + +use burn::config::Config; +use burn::module::Initializer; +use burn::module::Module; +use burn::module::Param; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Device; +use burn::tensor::Tensor; + +/// Configuration to create a [RMS Norm](RmsNorm) layer using the [init function](RmsNormConfig::init). +#[derive(Config, Debug)] +pub struct RmsNormConfig { + /// The size of the input features. + pub d_model: usize, + /// A value required for numerical stability. Default: 1e-5 + #[config(default = 1e-5)] + pub epsilon: f64, +} + +impl RmsNormConfig { + /// Initialize a new [RMS Norm](RmsNorm) module. + /// + /// # Panics + /// + /// Panics if `epsilon` is not positive. + pub fn init(&self, device: &Device) -> RmsNorm { + assert!(self.epsilon > 0.0, "epsilon must be positive."); + + let gamma = Initializer::Ones.init([self.d_model], device); + + RmsNorm { + gamma, + epsilon: self.epsilon, + } + } +} + +/// Applies RMS Normalization over an input tensor along the last dimension. +/// +/// `Y = X / sqrt(mean(X^2) + eps) * gamma` +/// +/// Where: +/// - `X` is the input tensor +/// - `Y` is the output tensor +/// - `gamma` is the learnable weight +/// - `mean` is the mean operation +/// - `eps` is a small value to avoid division by zero. +/// +/// Should be created using the [RmsNormConfig](RmsNormConfig) configuration. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct RmsNorm { + /// The learnable parameter to scale the normalized tensor + pub gamma: Param>, + /// A value required for numerical stability + pub epsilon: f64, +} + +impl RmsNorm { + /// Applies the forward pass on the input tensor. + /// + /// See the [RmsNorm](RmsNorm) documentation for more information. + /// + /// # Shapes + /// + /// - input: `[..., any, d_model]` + /// - output: `[..., any, d_model]` + pub fn forward(&self, x: Tensor) -> Tensor { + // Calculate the root-mean-square norm of the input tensor along the last dimension + let dtype = x.dtype(); + let rms = (x.clone().cast(DType::F32).square().mean_dim(D - 1) + self.epsilon).sqrt(); + (x / rms.cast(dtype)) * self.gamma.val().unsqueeze() + } +} + +impl ModuleDisplay for RmsNorm { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let [d_model] = self.gamma.shape().dims(); + content + .add("d_model", &d_model) + .add("epsilon", &self.epsilon) + .optional() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use alloc::format; + use burn::tensor::TensorData; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn rms_norm_forward() { + let device = Default::default(); + let module = RmsNormConfig::new(3).with_epsilon(1e-5).init(&device); + + let input = Tensor::arange(0..9, &device).float().reshape([3, 3]); + let output = module.forward(input); + + let expected = TensorData::from([ + [0.0000, 0.7746, 1.5492], + [0.7348, 0.9798, 1.2247], + [0.8514, 0.9933, 1.1352], + ]); + output + .to_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn display() { + let config = RmsNormConfig::new(6); + let layer_norm = config.init(&Default::default()); + + assert_eq!( + format!("{layer_norm}"), + "RmsNorm {d_model: 6, epsilon: 0.00001, params: 6}" + ); + } +} diff --git a/crates/burn-nn/src/modules/pool/adaptive_avg_pool1d.rs b/crates/burn-nn/src/modules/pool/adaptive_avg_pool1d.rs new file mode 100644 index 0000000..e8f4864 --- /dev/null +++ b/crates/burn-nn/src/modules/pool/adaptive_avg_pool1d.rs @@ -0,0 +1,76 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; + +use burn::tensor::module::adaptive_avg_pool1d; + +/// Configuration to create a [1D adaptive avg pooling](AdaptiveAvgPool1d) layer using the [init function](AdaptiveAvgPool1dConfig::init). +#[derive(Config, Debug)] +pub struct AdaptiveAvgPool1dConfig { + /// The size of the output. + pub output_size: usize, +} + +/// Applies a 1D adaptive avg pooling over input tensors. +/// +/// Should be created with [AdaptiveAvgPool1dConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct AdaptiveAvgPool1d { + /// The size of the output. + pub output_size: usize, +} + +impl ModuleDisplay for AdaptiveAvgPool1d { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content.add("output_size", &self.output_size).optional() + } +} + +impl AdaptiveAvgPool1dConfig { + /// Initialize a new [adaptive avg pool 1d](AdaptiveAvgPool1d) module. + pub fn init(&self) -> AdaptiveAvgPool1d { + AdaptiveAvgPool1d { + output_size: self.output_size, + } + } +} + +impl AdaptiveAvgPool1d { + /// Applies the forward pass on the input tensor. + /// + /// See [adaptive_avg_pool1d](burn::tensor::module::adaptive_avg_pool1d) for more information. + /// + /// # Shapes + /// + /// - input: `[batch_size, channels, length]` + /// - output: `[batch_size, channels, length_out]` + pub fn forward(&self, input: Tensor<3>) -> Tensor<3> { + adaptive_avg_pool1d(input, self.output_size) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn display() { + let config = AdaptiveAvgPool1dConfig::new(3); + let layer = config.init(); + + assert_eq!( + alloc::format!("{layer}"), + "AdaptiveAvgPool1d {output_size: 3}" + ); + } +} diff --git a/crates/burn-nn/src/modules/pool/adaptive_avg_pool2d.rs b/crates/burn-nn/src/modules/pool/adaptive_avg_pool2d.rs new file mode 100644 index 0000000..ba70b9b --- /dev/null +++ b/crates/burn-nn/src/modules/pool/adaptive_avg_pool2d.rs @@ -0,0 +1,78 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; + +use burn::tensor::module::adaptive_avg_pool2d; + +/// Configuration to create a [2D adaptive avg pooling](AdaptiveAvgPool2d) layer using the [init function](AdaptiveAvgPool2dConfig::init). +#[derive(Config, Debug)] +pub struct AdaptiveAvgPool2dConfig { + /// The size of the output. + pub output_size: [usize; 2], +} + +/// Applies a 2D adaptive avg pooling over input tensors. +/// +/// Should be created with [AdaptiveAvgPool2dConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct AdaptiveAvgPool2d { + /// The size of the output. + pub output_size: [usize; 2], +} + +impl ModuleDisplay for AdaptiveAvgPool2d { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let output_size = alloc::format!("{:?}", self.output_size); + + content.add("output_size", &output_size).optional() + } +} + +impl AdaptiveAvgPool2dConfig { + /// Initialize a new [adaptive avg pool 2d](AdaptiveAvgPool2d) module. + pub fn init(&self) -> AdaptiveAvgPool2d { + AdaptiveAvgPool2d { + output_size: self.output_size, + } + } +} + +impl AdaptiveAvgPool2d { + /// Applies the forward pass on the input tensor. + /// + /// See [adaptive_avg_pool2d](burn::tensor::module::adaptive_avg_pool2d) for more information. + /// + /// # Shapes + /// + /// - input: `[batch_size, channels, height_in, width_in]` + /// - output: `[batch_size, channels, height_out, width_out]` + pub fn forward(&self, input: Tensor<4>) -> Tensor<4> { + adaptive_avg_pool2d(input, self.output_size) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn display() { + let config = AdaptiveAvgPool2dConfig::new([3, 3]); + let layer = config.init(); + + assert_eq!( + alloc::format!("{layer}"), + "AdaptiveAvgPool2d {output_size: [3, 3]}" + ); + } +} diff --git a/crates/burn-nn/src/modules/pool/avg_pool1d.rs b/crates/burn-nn/src/modules/pool/avg_pool1d.rs new file mode 100644 index 0000000..06f771d --- /dev/null +++ b/crates/burn-nn/src/modules/pool/avg_pool1d.rs @@ -0,0 +1,219 @@ +use burn_core as burn; + +use crate::PaddingConfig1d; +use burn::config::Config; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::tensor::ops::PadMode; + +use burn::tensor::module::avg_pool1d; + +/// Configuration to create a [1D avg pooling](AvgPool1d) layer using the [init function](AvgPool1dConfig::init). +#[derive(Config, Debug)] +pub struct AvgPool1dConfig { + /// The size of the kernel. + pub kernel_size: usize, + /// The stride. + #[config(default = "kernel_size")] + pub stride: usize, + /// The padding configuration. + /// + /// Supports symmetric and asymmetric padding. `Same` padding with even kernel sizes + /// will automatically use asymmetric padding to preserve input dimensions. + #[config(default = "PaddingConfig1d::Valid")] + pub padding: PaddingConfig1d, + /// If the padding is counted in the denominator when computing the average. + #[config(default = "true")] + pub count_include_pad: bool, + /// If true, use ceiling instead of floor for output size calculation. + #[config(default = "false")] + pub ceil_mode: bool, +} + +/// Applies a 1D avg pooling over input tensors. +/// +/// Should be created with [AvgPool1dConfig](AvgPool1dConfig). +/// +/// # Remarks +/// +/// The zero-padding values will be included in the calculation +/// of the average. This means that the zeros are counted as +/// legitimate values, and they contribute to the denominator +/// when calculating the average. This is equivalent to +/// `torch.nn.AvgPool2d` with `count_include_pad=True`. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct AvgPool1d { + /// The stride. + pub stride: usize, + /// The size of the kernel. + pub kernel_size: usize, + /// The padding configuration. + #[module(skip)] + pub padding: PaddingConfig1d, + /// If the padding is counted in the denominator when computing the average. + pub count_include_pad: bool, + /// If true, use ceiling instead of floor for output size calculation. + pub ceil_mode: bool, +} + +impl ModuleDisplay for AvgPool1d { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("kernel_size", &self.kernel_size) + .add("stride", &self.stride) + .add_debug_attribute("padding", &self.padding) + .add("count_include_pad", &self.count_include_pad) + .add("ceil_mode", &self.ceil_mode) + .optional() + } +} + +impl AvgPool1dConfig { + /// Initialize a new [avg pool 1d](AvgPool1d) module. + pub fn init(&self) -> AvgPool1d { + AvgPool1d { + stride: self.stride, + kernel_size: self.kernel_size, + padding: self.padding.clone(), + count_include_pad: self.count_include_pad, + ceil_mode: self.ceil_mode, + } + } +} + +impl AvgPool1d { + /// Applies the forward pass on the input tensor. + /// + /// See [avg_pool1d](burn::tensor::module::avg_pool1d) for more information. + /// + /// # Shapes + /// + /// - input: `[batch_size, channels, length_in]` + /// - output: `[batch_size, channels, length_out]` + pub fn forward(&self, input: Tensor<3>) -> Tensor<3> { + let [_batch_size, _channels, length] = input.dims(); + + // Calculate padding as pair - handles Same, Valid, and Explicit uniformly + let (left, right) = + self.padding + .calculate_padding_1d_pair(length, self.kernel_size, self.stride); + + // TODO: Move asymmetric padding to functional level via PoolOptions + // See: https://github.com/tracel-ai/burn/issues/4362 + // Handle asymmetric padding by applying explicit pad operation first + if left != right { + // Burn's pad takes (left, right, top, bottom) for the last two dimensions + // For 1D (NCL format), we only pad L (last dim), so top/bottom = 0 + let padded = input.pad((left, right, 0, 0), PadMode::Constant(0.0)); + // Use zero padding for the pool operation since we already padded + avg_pool1d( + padded, + self.kernel_size, + self.stride, + 0, + self.count_include_pad, + self.ceil_mode, + ) + } else { + // Symmetric padding + avg_pool1d( + input, + self.kernel_size, + self.stride, + left, + self.count_include_pad, + self.ceil_mode, + ) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[test] + fn same_with_even_kernel_uses_asymmetric_padding() { + let device = Default::default(); + let config = AvgPool1dConfig::new(2) + .with_stride(1) + .with_padding(PaddingConfig1d::Same); + let pool = config.init(); + + // Input: [batch=1, channels=2, length=5] + let input = Tensor::<3>::ones([1, 2, 5], &device); + let output = pool.forward(input); + + // Same padding should preserve spatial dimensions + assert_eq!(output.dims(), [1, 2, 5]); + } + + #[test] + fn display() { + let config = AvgPool1dConfig::new(3); + let layer = config.init(); + + assert_eq!( + alloc::format!("{layer}"), + "AvgPool1d {kernel_size: 3, stride: 3, padding: Valid, count_include_pad: true, ceil_mode: false}" + ); + } + + #[rstest] + #[case(1)] + #[case(2)] + fn default_strides_match_kernel_size(#[case] kernel_size: usize) { + let config = AvgPool1dConfig::new(kernel_size); + + assert_eq!( + config.stride, kernel_size, + "Expected stride ({:?}) to match kernel size ({:?}) in default AvgPool1dConfig::new constructor", + config.stride, config.kernel_size + ); + } + + #[test] + fn asymmetric_padding_forward() { + let device = Default::default(); + // Create avg pool with asymmetric padding: left=1, right=2 + let config = AvgPool1dConfig::new(3) + .with_stride(1) + .with_padding(PaddingConfig1d::Explicit(1, 2)); + let pool = config.init(); + + // Input: [batch=1, channels=2, length=4] + let input = Tensor::<3>::ones([1, 2, 4], &device); + let output = pool.forward(input); + + // With asymmetric padding (1, 2), input length 4 becomes 4+1+2=7 + // Output length = (7 - 3) / 1 + 1 = 5 + assert_eq!(output.dims(), [1, 2, 5]); + } + + #[test] + fn symmetric_explicit_padding_forward() { + let device = Default::default(); + // Create avg pool with symmetric explicit padding: left=2, right=2 + let config = AvgPool1dConfig::new(3) + .with_stride(1) + .with_padding(PaddingConfig1d::Explicit(2, 2)); + let pool = config.init(); + + // Input: [batch=1, channels=2, length=4] + let input = Tensor::<3>::ones([1, 2, 4], &device); + let output = pool.forward(input); + + // With symmetric padding (2, 2), input length 4 becomes 4+2+2=8 + // Output length = (8 - 3) / 1 + 1 = 6 + assert_eq!(output.dims(), [1, 2, 6]); + } +} diff --git a/crates/burn-nn/src/modules/pool/avg_pool2d.rs b/crates/burn-nn/src/modules/pool/avg_pool2d.rs new file mode 100644 index 0000000..b3d5acb --- /dev/null +++ b/crates/burn-nn/src/modules/pool/avg_pool2d.rs @@ -0,0 +1,222 @@ +use burn_core as burn; + +use crate::PaddingConfig2d; +use burn::config::Config; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::tensor::ops::PadMode; + +use burn::tensor::module::avg_pool2d; + +/// Configuration to create a [2D avg pooling](AvgPool2d) layer using the [init function](AvgPool2dConfig::init). +#[derive(Config, Debug)] +pub struct AvgPool2dConfig { + /// The size of the kernel. + pub kernel_size: [usize; 2], + /// The strides. + #[config(default = "kernel_size")] + pub strides: [usize; 2], + /// The padding configuration. + /// + /// Supports symmetric and asymmetric padding. `Same` padding with even kernel sizes + /// will automatically use asymmetric padding to preserve input dimensions. + #[config(default = "PaddingConfig2d::Valid")] + pub padding: PaddingConfig2d, + /// If the padding is counted in the denominator when computing the average. + #[config(default = "true")] + pub count_include_pad: bool, + /// If true, use ceiling instead of floor for output size calculation. + #[config(default = "false")] + pub ceil_mode: bool, +} + +/// Applies a 2D avg pooling over input tensors. +/// +/// Should be created with [AvgPool2dConfig](AvgPool2dConfig). +/// +/// # Remarks +/// +/// The zero-padding values will be included in the calculation +/// of the average. This means that the zeros are counted as +/// legitimate values, and they contribute to the denominator +/// when calculating the average. This is equivalent to +/// `torch.nn.AvgPool2d` with `count_include_pad=True`. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct AvgPool2d { + /// Stride of the pooling. + pub stride: [usize; 2], + /// Size of the kernel. + pub kernel_size: [usize; 2], + /// Padding configuration. + #[module(skip)] + pub padding: PaddingConfig2d, + /// If the padding is counted in the denominator when computing the average. + pub count_include_pad: bool, + /// If true, use ceiling instead of floor for output size calculation. + pub ceil_mode: bool, +} + +impl ModuleDisplay for AvgPool2d { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("kernel_size", &alloc::format!("{:?}", self.kernel_size)) + .add("stride", &alloc::format!("{:?}", self.stride)) + .add_debug_attribute("padding", &self.padding) + .add("count_include_pad", &self.count_include_pad) + .add("ceil_mode", &self.ceil_mode) + .optional() + } +} + +impl AvgPool2dConfig { + /// Initialize a new [avg pool 2d](AvgPool2d) module. + pub fn init(&self) -> AvgPool2d { + AvgPool2d { + stride: self.strides, + kernel_size: self.kernel_size, + padding: self.padding.clone(), + count_include_pad: self.count_include_pad, + ceil_mode: self.ceil_mode, + } + } +} + +impl AvgPool2d { + /// Applies the forward pass on the input tensor. + /// + /// See [avg_pool2d](burn::tensor::module::avg_pool2d) for more information. + /// + /// # Shapes + /// + /// - input: `[batch_size, channels, height_in, width_in]` + /// - output: `[batch_size, channels, height_out, width_out]` + pub fn forward(&self, input: Tensor<4>) -> Tensor<4> { + let [_batch_size, _channels_in, height_in, width_in] = input.dims(); + + // Calculate padding as pairs - handles Same, Valid, and Explicit uniformly + let ((top, bottom), (left, right)) = self.padding.calculate_padding_2d_pairs( + height_in, + width_in, + &self.kernel_size, + &self.stride, + ); + + // TODO: Move asymmetric padding to functional level via PoolOptions + // See: https://github.com/tracel-ai/burn/issues/4362 + // Handle asymmetric padding by applying explicit pad operation first + if top != bottom || left != right { + // Burn's pad takes (left, right, top, bottom) for the last two dimensions + let padded = input.pad((left, right, top, bottom), PadMode::Constant(0.0)); + // Use zero padding for the pool operation since we already padded + avg_pool2d( + padded, + self.kernel_size, + self.stride, + [0, 0], + self.count_include_pad, + self.ceil_mode, + ) + } else { + // Symmetric padding + avg_pool2d( + input, + self.kernel_size, + self.stride, + [top, left], + self.count_include_pad, + self.ceil_mode, + ) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[test] + fn same_with_even_kernel_uses_asymmetric_padding() { + let device = Default::default(); + let config = AvgPool2dConfig::new([2, 2]) + .with_strides([1, 1]) + .with_padding(PaddingConfig2d::Same); + let pool = config.init(); + + // Input: [batch=1, channels=2, height=5, width=5] + let input = Tensor::<4>::ones([1, 2, 5, 5], &device); + let output = pool.forward(input); + + // Same padding should preserve spatial dimensions + assert_eq!(output.dims(), [1, 2, 5, 5]); + } + + #[test] + fn display() { + let config = AvgPool2dConfig::new([3, 3]); + + let layer = config.init(); + + assert_eq!( + alloc::format!("{layer}"), + "AvgPool2d {kernel_size: [3, 3], stride: [3, 3], padding: Valid, count_include_pad: true, ceil_mode: false}" + ); + } + + #[rstest] + #[case([2, 2])] + #[case([1, 2])] + fn default_strides_match_kernel_size(#[case] kernel_size: [usize; 2]) { + let config = AvgPool2dConfig::new(kernel_size); + + assert_eq!( + config.strides, kernel_size, + "Expected strides ({:?}) to match kernel size ({:?}) in default AvgPool2dConfig::new constructor", + config.strides, config.kernel_size + ); + } + + #[test] + fn asymmetric_padding_forward() { + let device = Default::default(); + // Create avg pool with asymmetric padding: top=1, left=2, bottom=3, right=4 + let config = AvgPool2dConfig::new([3, 3]) + .with_strides([1, 1]) + .with_padding(PaddingConfig2d::Explicit(1, 2, 3, 4)); + let pool = config.init(); + + // Input: [batch=1, channels=2, height=4, width=5] + let input = Tensor::<4>::ones([1, 2, 4, 5], &device); + let output = pool.forward(input); + + // Height: 4 + 1 + 3 = 8, output = (8 - 3) / 1 + 1 = 6 + // Width: 5 + 2 + 4 = 11, output = (11 - 3) / 1 + 1 = 9 + assert_eq!(output.dims(), [1, 2, 6, 9]); + } + + #[test] + fn symmetric_explicit_padding_forward() { + let device = Default::default(); + // Create avg pool with symmetric explicit padding: top=2, left=2, bottom=2, right=2 + let config = AvgPool2dConfig::new([3, 3]) + .with_strides([1, 1]) + .with_padding(PaddingConfig2d::Explicit(2, 2, 2, 2)); + let pool = config.init(); + + // Input: [batch=1, channels=2, height=4, width=5] + let input = Tensor::<4>::ones([1, 2, 4, 5], &device); + let output = pool.forward(input); + + // Height: 4 + 2 + 2 = 8, output = (8 - 3) / 1 + 1 = 6 + // Width: 5 + 2 + 2 = 9, output = (9 - 3) / 1 + 1 = 7 + assert_eq!(output.dims(), [1, 2, 6, 7]); + } +} diff --git a/crates/burn-nn/src/modules/pool/max_pool1d.rs b/crates/burn-nn/src/modules/pool/max_pool1d.rs new file mode 100644 index 0000000..f2df065 --- /dev/null +++ b/crates/burn-nn/src/modules/pool/max_pool1d.rs @@ -0,0 +1,213 @@ +use burn_core as burn; + +use crate::PaddingConfig1d; +use burn::config::Config; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::tensor::ops::PadMode; + +use burn::tensor::module::max_pool1d; + +/// Configuration to create a [1D max pooling](MaxPool1d) layer using the [init function](MaxPool1dConfig::init). +#[derive(Config, Debug)] +pub struct MaxPool1dConfig { + /// The size of the kernel. + pub kernel_size: usize, + /// The stride. + #[config(default = "kernel_size")] + pub stride: usize, + /// The padding configuration. + /// + /// Supports symmetric and asymmetric padding. `Same` padding with even kernel sizes + /// will automatically use asymmetric padding to preserve input dimensions. + #[config(default = "PaddingConfig1d::Valid")] + pub padding: PaddingConfig1d, + /// The dilation. + #[config(default = "1")] + pub dilation: usize, + /// If true, use ceiling instead of floor for output size calculation. + #[config(default = "false")] + pub ceil_mode: bool, +} + +/// Applies a 1D max pooling over input tensors. +/// +/// Should be created with [MaxPool1dConfig](MaxPool1dConfig). +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct MaxPool1d { + /// The stride. + pub stride: usize, + /// The size of the kernel. + pub kernel_size: usize, + /// The padding configuration. + #[module(skip)] + pub padding: PaddingConfig1d, + /// The dilation. + pub dilation: usize, + /// If true, use ceiling instead of floor for output size calculation. + pub ceil_mode: bool, +} + +impl ModuleDisplay for MaxPool1d { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("kernel_size", &self.kernel_size) + .add("stride", &self.stride) + .add_debug_attribute("padding", &self.padding) + .add("dilation", &self.dilation) + .add("ceil_mode", &self.ceil_mode) + .optional() + } +} + +impl MaxPool1dConfig { + /// Initialize a new [max pool 1d](MaxPool1d) module. + pub fn init(&self) -> MaxPool1d { + MaxPool1d { + stride: self.stride, + kernel_size: self.kernel_size, + padding: self.padding.clone(), + dilation: self.dilation, + ceil_mode: self.ceil_mode, + } + } +} + +impl MaxPool1d { + /// Applies the forward pass on the input tensor. + /// + /// See [max_pool1d](burn::tensor::module::max_pool1d) for more information. + /// + /// # Shapes + /// + /// - input: `[batch_size, channels, length_in]` + /// - output: `[batch_size, channels, length_out]` + pub fn forward(&self, input: Tensor<3>) -> Tensor<3> { + let [_batch_size, _channels, length] = input.dims(); + + // Calculate padding as pair - handles Same, Valid, and Explicit uniformly + let (left, right) = + self.padding + .calculate_padding_1d_pair(length, self.kernel_size, self.stride); + + // TODO: Move asymmetric padding to functional level via PoolOptions + // See: https://github.com/tracel-ai/burn/issues/4362 + // Handle asymmetric padding by applying explicit pad operation first + if left != right { + // For 1D (NCL format), pad the length dimension with (left, right) + // and no padding for channel dimension (top=0, bottom=0) + // Use -inf for max pooling so padded values don't affect the max + let padded = input.pad((left, right, 0, 0), PadMode::Constant(f32::NEG_INFINITY)); + // Use zero padding for the pool operation since we already padded + max_pool1d( + padded, + self.kernel_size, + self.stride, + 0, + self.dilation, + self.ceil_mode, + ) + } else { + // Symmetric padding + max_pool1d( + input, + self.kernel_size, + self.stride, + left, + self.dilation, + self.ceil_mode, + ) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[test] + fn same_with_even_kernel_uses_asymmetric_padding() { + let device = Default::default(); + let config = MaxPool1dConfig::new(2) + .with_stride(1) + .with_padding(PaddingConfig1d::Same); + let pool = config.init(); + + // Input: [batch=1, channels=2, length=5] + let input = Tensor::<3>::ones([1, 2, 5], &device); + let output = pool.forward(input); + + // Same padding should preserve spatial dimensions + assert_eq!(output.dims(), [1, 2, 5]); + } + + #[test] + fn display() { + let config = MaxPool1dConfig::new(3); + + let layer = config.init(); + + assert_eq!( + alloc::format!("{layer}"), + "MaxPool1d {kernel_size: 3, stride: 3, padding: Valid, dilation: 1, ceil_mode: false}" + ); + } + + #[rstest] + #[case(1)] + #[case(2)] + fn default_strides_match_kernel_size(#[case] kernel_size: usize) { + let config = MaxPool1dConfig::new(kernel_size); + + assert_eq!( + config.stride, kernel_size, + "Expected stride ({:?}) to match kernel size ({:?}) in default MaxPool1dConfig::new constructor", + config.stride, config.kernel_size + ); + } + + #[test] + fn asymmetric_padding_forward() { + let device = Default::default(); + // Create max pool with asymmetric padding: left=1, right=2 + let config = MaxPool1dConfig::new(3) + .with_stride(1) + .with_padding(PaddingConfig1d::Explicit(1, 2)); + let pool = config.init(); + + // Input: [batch=1, channels=2, length=4] + let input = Tensor::<3>::ones([1, 2, 4], &device); + let output = pool.forward(input); + + // With asymmetric padding (1, 2), input length 4 becomes 4+1+2=7 + // Output length = (7 - 3) / 1 + 1 = 5 + assert_eq!(output.dims(), [1, 2, 5]); + } + + #[test] + fn symmetric_explicit_padding_forward() { + let device = Default::default(); + // Create max pool with symmetric explicit padding: left=2, right=2 + let config = MaxPool1dConfig::new(3) + .with_stride(1) + .with_padding(PaddingConfig1d::Explicit(2, 2)); + let pool = config.init(); + + // Input: [batch=1, channels=2, length=4] + let input = Tensor::<3>::ones([1, 2, 4], &device); + let output = pool.forward(input); + + // With symmetric padding (2, 2), input length 4 becomes 4+2+2=8 + // Output length = (8 - 3) / 1 + 1 = 6 + assert_eq!(output.dims(), [1, 2, 6]); + } +} diff --git a/crates/burn-nn/src/modules/pool/max_pool2d.rs b/crates/burn-nn/src/modules/pool/max_pool2d.rs new file mode 100644 index 0000000..6061093 --- /dev/null +++ b/crates/burn-nn/src/modules/pool/max_pool2d.rs @@ -0,0 +1,218 @@ +use burn_core as burn; + +use crate::PaddingConfig2d; +use burn::config::Config; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Tensor; +use burn::tensor::ops::PadMode; + +use burn::tensor::module::max_pool2d; + +/// Configuration to create a [2D max pooling](MaxPool2d) layer using the [init function](MaxPool2dConfig::init). +#[derive(Debug, Config)] +pub struct MaxPool2dConfig { + /// The size of the kernel. + pub kernel_size: [usize; 2], + /// The strides. + #[config(default = "kernel_size")] + pub strides: [usize; 2], + /// The padding configuration. + /// + /// Supports symmetric and asymmetric padding. `Same` padding with even kernel sizes + /// will automatically use asymmetric padding to preserve input dimensions. + #[config(default = "PaddingConfig2d::Valid")] + pub padding: PaddingConfig2d, + /// The dilation. + #[config(default = "[1, 1]")] + pub dilation: [usize; 2], + /// If true, use ceiling instead of floor for output size calculation. + #[config(default = "false")] + pub ceil_mode: bool, +} + +/// Applies a 2D max pooling over input tensors. +/// +/// Should be created with [MaxPool2dConfig](MaxPool2dConfig). +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct MaxPool2d { + /// The strides. + pub stride: [usize; 2], + /// The size of the kernel. + pub kernel_size: [usize; 2], + /// The padding configuration. + #[module(skip)] + pub padding: PaddingConfig2d, + /// The dilation. + pub dilation: [usize; 2], + /// If true, use ceiling instead of floor for output size calculation. + pub ceil_mode: bool, +} + +impl ModuleDisplay for MaxPool2d { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("kernel_size", &alloc::format!("{:?}", self.kernel_size)) + .add("stride", &alloc::format!("{:?}", self.stride)) + .add_debug_attribute("padding", &self.padding) + .add("dilation", &alloc::format!("{:?}", self.dilation)) + .add("ceil_mode", &self.ceil_mode) + .optional() + } +} + +impl MaxPool2dConfig { + /// Initialize a new [max pool 2d](MaxPool2d) module. + pub fn init(&self) -> MaxPool2d { + MaxPool2d { + stride: self.strides, + kernel_size: self.kernel_size, + padding: self.padding.clone(), + dilation: self.dilation, + ceil_mode: self.ceil_mode, + } + } +} + +impl MaxPool2d { + /// Applies the forward pass on the input tensor. + /// + /// See [max_pool2d](burn::tensor::module::max_pool2d) for more information. + /// + /// # Shapes + /// + /// - input: `[batch_size, channels, height_in, width_in]` + /// - output: `[batch_size, channels, height_out, width_out]` + pub fn forward(&self, input: Tensor<4>) -> Tensor<4> { + let [_batch_size, _channels_in, height_in, width_in] = input.dims(); + + // Calculate padding as pairs - handles Same, Valid, and Explicit uniformly + let ((top, bottom), (left, right)) = self.padding.calculate_padding_2d_pairs( + height_in, + width_in, + &self.kernel_size, + &self.stride, + ); + + // TODO: Move asymmetric padding to functional level via PoolOptions + // See: https://github.com/tracel-ai/burn/issues/4362 + // Handle asymmetric padding by applying explicit pad operation first + if top != bottom || left != right { + // Burn's pad takes (left, right, top, bottom) for the last two dimensions + // Use -inf for max pooling so padded values don't affect the max + let padded = input.pad( + (left, right, top, bottom), + PadMode::Constant(f32::NEG_INFINITY), + ); + // Use zero padding for the pool operation since we already padded + max_pool2d( + padded, + self.kernel_size, + self.stride, + [0, 0], + self.dilation, + self.ceil_mode, + ) + } else { + // Symmetric padding + max_pool2d( + input, + self.kernel_size, + self.stride, + [top, left], + self.dilation, + self.ceil_mode, + ) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[test] + fn same_with_even_kernel_uses_asymmetric_padding() { + let device = Default::default(); + let config = MaxPool2dConfig::new([2, 2]) + .with_strides([1, 1]) + .with_padding(PaddingConfig2d::Same); + let pool = config.init(); + + // Input: [batch=1, channels=2, height=5, width=5] + let input = Tensor::<4>::ones([1, 2, 5, 5], &device); + let output = pool.forward(input); + + // Same padding should preserve spatial dimensions + assert_eq!(output.dims(), [1, 2, 5, 5]); + } + + #[test] + fn display() { + let config = MaxPool2dConfig::new([3, 3]); + + let layer = config.init(); + + assert_eq!( + alloc::format!("{layer}"), + "MaxPool2d {kernel_size: [3, 3], stride: [3, 3], padding: Valid, dilation: [1, 1], ceil_mode: false}" + ); + } + + #[rstest] + #[case([2, 2])] + #[case([1, 2])] + fn default_strides_match_kernel_size(#[case] kernel_size: [usize; 2]) { + let config = MaxPool2dConfig::new(kernel_size); + + assert_eq!( + config.strides, kernel_size, + "Expected strides ({:?}) to match kernel size ({:?}) in default MaxPool2dConfig::new constructor", + config.strides, config.kernel_size + ); + } + + #[test] + fn asymmetric_padding_forward() { + let device = Default::default(); + // Create max pool with asymmetric padding: top=1, left=2, bottom=3, right=4 + let config = MaxPool2dConfig::new([3, 3]) + .with_strides([1, 1]) + .with_padding(PaddingConfig2d::Explicit(1, 2, 3, 4)); + let pool = config.init(); + + // Input: [batch=1, channels=2, height=4, width=5] + let input = Tensor::<4>::ones([1, 2, 4, 5], &device); + let output = pool.forward(input); + + // Height: 4 + 1 + 3 = 8, output = (8 - 3) / 1 + 1 = 6 + // Width: 5 + 2 + 4 = 11, output = (11 - 3) / 1 + 1 = 9 + assert_eq!(output.dims(), [1, 2, 6, 9]); + } + + #[test] + fn symmetric_explicit_padding_forward() { + let device = Default::default(); + // Create max pool with symmetric explicit padding: top=2, left=2, bottom=2, right=2 + let config = MaxPool2dConfig::new([3, 3]) + .with_strides([1, 1]) + .with_padding(PaddingConfig2d::Explicit(2, 2, 2, 2)); + let pool = config.init(); + + // Input: [batch=1, channels=2, height=4, width=5] + let input = Tensor::<4>::ones([1, 2, 4, 5], &device); + let output = pool.forward(input); + + // Height: 4 + 2 + 2 = 8, output = (8 - 3) / 1 + 1 = 6 + // Width: 5 + 2 + 2 = 9, output = (9 - 3) / 1 + 1 = 7 + assert_eq!(output.dims(), [1, 2, 6, 7]); + } +} diff --git a/crates/burn-nn/src/modules/pool/mod.rs b/crates/burn-nn/src/modules/pool/mod.rs new file mode 100644 index 0000000..622a4b6 --- /dev/null +++ b/crates/burn-nn/src/modules/pool/mod.rs @@ -0,0 +1,13 @@ +mod adaptive_avg_pool1d; +mod adaptive_avg_pool2d; +mod avg_pool1d; +mod avg_pool2d; +mod max_pool1d; +mod max_pool2d; + +pub use adaptive_avg_pool1d::*; +pub use adaptive_avg_pool2d::*; +pub use avg_pool1d::*; +pub use avg_pool2d::*; +pub use max_pool1d::*; +pub use max_pool2d::*; diff --git a/crates/burn-nn/src/modules/pos_encoding.rs b/crates/burn-nn/src/modules/pos_encoding.rs new file mode 100644 index 0000000..409491b --- /dev/null +++ b/crates/burn-nn/src/modules/pos_encoding.rs @@ -0,0 +1,289 @@ +use burn_core as burn; + +use alloc::vec::Vec; +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Module, ModuleDisplay}; + +use burn::tensor::TensorData; +use burn::tensor::{Device, Tensor}; + +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float as _; + +/// Configuration to create a [PositionalEncoding](PositionalEncoding) layer using the [init function](PositionalEncodingConfig::init). +#[derive(Config, Debug)] +pub struct PositionalEncodingConfig { + /// Maximum sequence size to use. + #[config(default = "5_000")] + pub max_sequence_size: usize, + + /// The size of each vector. + pub d_model: usize, + + /// Max time scale to use. + #[config(default = "10_000")] + pub max_timescale: usize, +} + +/// Positional encoding layer for transformer models. +/// +/// This layer adds positional information to the input embeddings, allowing the transformer model +/// to take into account the order of the sequence. The positional encoding is added to the input +/// embeddings by computing a set of sinusoidal functions with different frequencies and phases. +/// +/// Sinusoids are used for positional embedding introduced in +/// [Attention is all you need](https://arxiv.org/abs/1706.03762). +/// +/// The reference implementation can be found here: +/// [LANGUAGE MODELING WITH NN.TRANSFORMER AND TORCHTEXT +/// ](https://pytorch.org/tutorials/beginner/transformer_tutorial.html) +/// +/// Should be created using [PositionalEncodingConfig] +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct PositionalEncoding { + /// The sinusoids used to add positional information to the input embeddings. + pub sinusoids: Tensor<3>, + /// The maximum sequence size to use. + pub max_sequence_size: usize, + /// Max time scale to use. + pub max_timescale: usize, +} + +impl ModuleDisplay for PositionalEncoding { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let [_, _, d_model] = self.sinusoids.shape().dims(); + content + .add("d_model", &d_model) + .add("max_sequence_size", &self.max_sequence_size) + .add("max_timescale", &self.max_timescale) + .optional() + } +} + +impl PositionalEncodingConfig { + /// Initialize a new [PositionalEncoding](PositionalEncoding) module. + pub fn init(&self, device: &Device) -> PositionalEncoding { + let sinusoids = generate_sinusoids( + self.max_sequence_size, + self.d_model, + self.max_timescale, + device, + ) + .unsqueeze::<3>(); + + PositionalEncoding { + sinusoids, + max_sequence_size: self.max_sequence_size, + max_timescale: self.max_timescale, + } + } +} + +impl PositionalEncoding { + /// Applies the forward pass on the input tensor by adding the sinusoids to the input. + /// + /// # Shapes + /// + /// * input: [batch_size, seq_length, d_model] + /// * output: [batch_size, seq_length, d_model] + /// + /// + /// # Panics + /// + /// * Panics if the input sequence length is greater than the maximum sequence size. + /// * Panics if the input d_model is not equal to the d_model of the sinusoids. + pub fn forward(&self, input: Tensor<3>) -> Tensor<3> { + let [_, seq_length, d_model_input] = input.dims(); + + let [batch_size, max_sequence_size, d_model] = self.sinusoids.dims(); + + assert!( + max_sequence_size >= seq_length, + "max_sequence_size({max_sequence_size}) must be greater or equal than length({seq_length})" + ); + + assert!( + d_model_input == d_model, + "d_model({d_model_input}) of the input must be equal to d_model of encoding({d_model})" + ); + + let slices = [0..batch_size, 0..seq_length, 0..d_model]; + + input.add(self.sinusoids.clone().slice(slices)) + } +} + +/// Returns sinusoids for positional embedding introduced in +/// [Attention is all you need](https://arxiv.org/abs/1706.03762). +/// +/// The reference implementation can be found here: +/// [LANGUAGE MODELING WITH NN.TRANSFORMER AND TORCHTEXT +/// ](https://pytorch.org/tutorials/beginner/transformer_tutorial.html) +/// +/// # Arguments +/// +/// * `length` - The length of the sequence. +/// * `d_model` - The size of each vector. +/// * `max_timescale` - The maximum time scale to use. +/// +/// # Returns +/// +/// A tensor of shape [length, d_model] containing the sinusoids. +pub fn generate_sinusoids( + length: usize, + d_model: usize, + max_timescale: usize, + device: &Device, +) -> Tensor<2> { + assert!(d_model.is_multiple_of(2), "d_model must be even"); + assert!( + max_timescale >= length, + "max_timescale must be greater than or equal to length" + ); + + // Calculate the increment for the logarithmic timescale + let log_timescale_increment = -(max_timescale as f32).ln() / d_model as f32; + + // Create a vector to hold the sinusoids + let mut scaled_time_sin_cos = Vec::with_capacity(length); + + // Loop over each position in the sequence + for i in 0..length { + // Create a vector to hold the sinusoids for this position + let mut row = Vec::with_capacity(d_model / 2); + // Loop over each dimension of the sinusoids + for k in (0..d_model).step_by(2) { + // Calculate the division term for this dimension + let div_term = (k as f32 * log_timescale_increment).exp(); + // Calculate the sine and cosine values for this dimension and position + row.push((div_term * i as f32).sin()); + row.push((div_term * i as f32).cos()); + } + + // Add the sinusoids for this position to the vector + scaled_time_sin_cos.push(row); + } + + // Convert the sinusoids to a tensor and return it + let data = TensorData::new( + scaled_time_sin_cos.into_iter().flatten().collect(), + [length, d_model], + ); + + Tensor::<2>::from_data(data, device) +} + +#[cfg(test)] +mod tests { + + use super::*; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn test_module() { + let d_model = 6; + let length = 3; + + // expected to broadcast + let batch_size = 2; + + let device = Default::default(); + let pe = PositionalEncodingConfig::new(d_model).init(&device); + + // Use a tensor of zeros as input for easy verification of the output + // The output should be the sinusoids broadcasted to the input shape + let tensor = Tensor::zeros([batch_size, length, d_model], &device); + + let output = pe.forward(tensor); + + assert_eq!(&*output.shape(), [batch_size, length, d_model]); + + let expected = Tensor::<3>::from_floats( + [ + [ + [0.00000, 1.00000, 0.00000, 1.00000, 0.00000, 1.00000], + [0.84147, 0.54030, 0.04640, 0.99892, 0.00215, 1.00000], + [0.90930, -0.41615, 0.09270, 0.99569, 0.00431, 0.99999], + ], + [ + [0.00000, 1.00000, 0.00000, 1.00000, 0.00000, 1.00000], + [0.84147, 0.54030, 0.04640, 0.99892, 0.00215, 1.00000], + [0.90930, -0.41615, 0.09270, 0.99569, 0.00431, 0.99999], + ], + ], + &device, + ); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::default()); + } + + #[test] + fn test_generate_sinusoids() { + let device = Default::default(); + let sinusoids = generate_sinusoids(12, 6, 10_000, &device); + + // The values are taken from the pytorch reference implementation + let expected = Tensor::<2>::from_floats( + [ + [0.00000, 1.00000, 0.00000, 1.00000, 0.00000, 1.00000], + [0.84147, 0.54030, 0.04640, 0.99892, 0.00215, 1.00000], + [0.90930, -0.41615, 0.09270, 0.99569, 0.00431, 0.99999], + [0.14112, -0.98999, 0.13880, 0.99032, 0.00646, 0.99998], + [-0.75680, -0.65364, 0.18460, 0.98281, 0.00862, 0.99996], + [-0.95892, 0.28366, 0.23000, 0.97319, 0.01077, 0.99994], + [-0.27942, 0.96017, 0.27491, 0.96147, 0.01293, 0.99992], + [0.65699, 0.75390, 0.31922, 0.94768, 0.01508, 0.99989], + [0.98936, -0.14550, 0.36285, 0.93185, 0.01723, 0.99985], + [0.41212, -0.91113, 0.40570, 0.91401, 0.01939, 0.99981], + [-0.54402, -0.83907, 0.44767, 0.89420, 0.02154, 0.99977], + [-0.99999, 0.00443, 0.48868, 0.87246, 0.02370, 0.99972], + ], + &device, + ); + sinusoids + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::default()); + } + + #[test] + #[should_panic] + fn d_model_input_should_match() { + let d_model = 8; + let device = Default::default(); + let pe = PositionalEncodingConfig::new(d_model).init(&device); + let input = Tensor::zeros([1, 5, 10], &device); + let _output = pe.forward(input); + } + + #[test] + #[should_panic] + fn input_length_should_be_less_than_max_len() { + let d_model = 8; + let device = Default::default(); + let pe = PositionalEncodingConfig::new(d_model).init(&device); + let input = Tensor::zeros([1, 6_000, d_model], &device); + let _output = pe.forward(input); + } + + #[test] + fn display() { + let config = PositionalEncodingConfig::new(4); + let pe = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{pe}"), + "PositionalEncoding {d_model: 4, max_sequence_size: 5000, max_timescale: 10000}" + ); + } +} diff --git a/crates/burn-nn/src/modules/rnn/basic.rs b/crates/burn-nn/src/modules/rnn/basic.rs new file mode 100644 index 0000000..d063d8d --- /dev/null +++ b/crates/burn-nn/src/modules/rnn/basic.rs @@ -0,0 +1,698 @@ +use burn_core as burn; + +use crate::GateController; +use crate::activation::{Activation, ActivationConfig}; +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Initializer, Module, ModuleDisplay}; +use burn::tensor::Device; +use burn::tensor::Tensor; + +/// A RnnState is used to store hidden state in RNN. +pub struct RnnState { + /// The hidden state. + pub hidden: Tensor, +} + +impl RnnState { + /// Initialize a new [RNN State](RnnState). + pub fn new(hidden: Tensor) -> Self { + Self { hidden } + } +} + +/// Configuration to create a [Rnn](Rnn) module using the [init function](RnnConfig::init). +#[derive(Config, Debug)] +pub struct RnnConfig { + /// The size of the input features. + pub d_input: usize, + /// The size of the hidden state. + pub d_hidden: usize, + /// If a bias should be applied during the Rnn transformation. + pub bias: bool, + /// Rnn initializer + #[config(default = "Initializer::XavierNormal{gain:1.0}")] + pub initializer: Initializer, + /// If true, the input tensor is expected to be `[batch_size, seq_length, input_size]`. + /// If false, the input tensor is expected to be `[seq_length, batch_size, input_size]`. + #[config(default = true)] + pub batch_first: bool, + /// If true, process the sequence in reverse order. + /// This is useful for implementing reverse-direction RNNs (e.g., ONNX reverse direction). + #[config(default = false)] + pub reverse: bool, + /// Optional hidden state clip threshold. If provided, hidden state values are clipped + /// to the range `[-clip, +clip]` after each timestep. This can help prevent + /// exploding values during inference. + pub clip: Option, + /// Activation function applied to the hidden state before computing hidden output. + /// Default is Tanh, which is standard for Rnn. + #[config(default = "ActivationConfig::Tanh")] + pub hidden_activation: ActivationConfig, +} + +/// The Rnn module. This implementation is for a unidirectional, stateless, Rnn. +/// Should be created with [RnnConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Rnn { + /// gate controller for Rnn (has single gate). + pub gate: GateController, + /// The hidden state of the Rnn. + pub d_hidden: usize, + /// If true, input is `[batch_size, seq_length, input_size]`. + /// If false, input is `[seq_length, batch_size, input_size]`. + pub batch_first: bool, + /// If true, process the sequence in reverse order. + pub reverse: bool, + /// Optional hidden state clip threshold. + pub clip: Option, + /// Activation function for hidden output. + pub hidden_activation: Activation, +} + +impl ModuleDisplay for Rnn { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let [d_input, _] = self.gate.input_transform.weight.shape().dims(); + let bias = self.gate.input_transform.bias.is_some(); + + content + .add("d_input", &d_input) + .add("d_hidden", &self.d_hidden) + .add("bias", &bias) + .optional() + } +} + +impl RnnConfig { + /// Initialize a new [Rnn](Rnn) module. + pub fn init(&self, device: &Device) -> Rnn { + let d_output = self.d_hidden; + + let new_gate = || { + GateController::new( + self.d_input, + d_output, + self.bias, + self.initializer.clone(), + device, + ) + }; + + Rnn { + gate: new_gate(), + d_hidden: self.d_hidden, + batch_first: self.batch_first, + reverse: self.reverse, + clip: self.clip, + hidden_activation: self.hidden_activation.init(device), + } + } +} + +impl Rnn { + /// Applies the forward pass on the input tensor. This RNN implementation + /// returns the state for each element in a sequence (i.e., across seq_length) and a final state. + /// + /// ## Parameters: + /// - batched_input: The input tensor of shape: + /// - `[batch_size, sequence_length, input_size]` if `batch_first` is true (default) + /// - `[sequence_length, batch_size, input_size]` if `batch_first` is false + /// - state: An optional `RnnState` representing the initial hidden state. + /// The state tensor has shape `[batch_size, hidden_size]`. + /// If no initial state is provided, these tensors are initialized to zeros. + /// + /// ## Returns: + /// - output: A tensor represents the output features of Rnn. Shape: + /// - `[batch_size, sequence_length, hidden_size]` if `batch_first` is true + /// - `[sequence_length, batch_size, hidden_size]` if `batch_first` is false + /// - state: A `RnnState` represents the final hidden state. The hidden state tensor has the shape + /// `[batch_size, hidden_size]`. + pub fn forward( + &self, + batched_input: Tensor<3>, + state: Option>, + ) -> (Tensor<3>, RnnState<2>) { + // Convert to batch-first layout internally if needed + let batched_input = if self.batch_first { + batched_input + } else { + batched_input.swap_dims(0, 1) + }; + + let device = batched_input.device(); + let [batch_size, seq_length, _] = batched_input.dims(); + + // Process sequence in forward or reverse order based on config + let (output, state) = if self.reverse { + self.forward_iter( + batched_input.iter_dim(1).rev().zip((0..seq_length).rev()), + state, + batch_size, + seq_length, + &device, + ) + } else { + self.forward_iter( + batched_input.iter_dim(1).zip(0..seq_length), + state, + batch_size, + seq_length, + &device, + ) + }; + + // Convert output back to seq-first layout if needed + let output = if self.batch_first { + output + } else { + output.swap_dims(0, 1) + }; + + (output, state) + } + + fn forward_iter, usize)>>( + &self, + input_timestep_iter: I, + state: Option>, + batch_size: usize, + seq_length: usize, + device: &Device, + ) -> (Tensor<3>, RnnState<2>) { + let mut batched_hidden_state = + Tensor::empty([batch_size, seq_length, self.d_hidden], device); + + let mut hidden_state = match state { + Some(state) => state.hidden, + None => Tensor::zeros([batch_size, self.d_hidden], device), + }; + + for (input_t, t) in input_timestep_iter { + let input_t = input_t.squeeze_dim(1); + + // Compute gate output: h_t = activation(W_i @ x_t + W_h @ h_{t-1} + b) + let biased_gate_sum = self + .gate + .gate_product(input_t.clone(), hidden_state.clone()); + + let output_values = self.hidden_activation.forward(biased_gate_sum); + + // Update hidden state + hidden_state = output_values; + + // Apply hidden state clipping if configured + if let Some(clip) = self.clip { + hidden_state = hidden_state.clamp(-clip, clip); + } + + let unsqueezed_hidden_state = hidden_state.clone().unsqueeze_dim(1); + + // store the hidden state for this timestep + batched_hidden_state = batched_hidden_state.slice_assign( + [0..batch_size, t..(t + 1), 0..self.d_hidden], + unsqueezed_hidden_state.clone(), + ); + } + + (batched_hidden_state, RnnState::new(hidden_state)) + } +} + +/// Configuration to create a [BiRnn](BiRnn) module using the [init function](BiRnnConfig::init). +#[derive(Config, Debug)] +pub struct BiRnnConfig { + /// The size of the input features. + pub d_input: usize, + /// The size of the hidden state. + pub d_hidden: usize, + /// If a bias should be applied during the BiRnn transformation. + pub bias: bool, + /// BiRnn initializer + #[config(default = "Initializer::XavierNormal{gain:1.0}")] + pub initializer: Initializer, + /// If true, the input tensor is expected to be `[batch_size, seq_length, input_size]`. + /// If false, the input tensor is expected to be `[seq_length, batch_size, input_size]`. + #[config(default = true)] + pub batch_first: bool, + /// Optional hidden state clip threshold. + pub clip: Option, + /// Activation function applied to the hidden state before computing hidden output. + #[config(default = "ActivationConfig::Tanh")] + pub hidden_activation: ActivationConfig, +} + +/// The BiRnn module. This implementation is for Bidirectional RNN. +/// Should be created with [BiRnnConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct BiRnn { + /// RNN for the forward direction. + pub forward: Rnn, + /// RNN for the reverse direction. + pub reverse: Rnn, + /// The size of the hidden state. + pub d_hidden: usize, + /// If true, input is `[batch_size, seq_length, input_size]`. + /// If false, input is `[seq_length, batch_size, input_size]`. + pub batch_first: bool, +} + +impl ModuleDisplay for BiRnn { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let [d_input, _] = self.forward.gate.input_transform.weight.shape().dims(); + let bias = self.forward.gate.input_transform.bias.is_some(); + + content + .add("d_input", &d_input) + .add("d_hidden", &self.d_hidden) + .add("bias", &bias) + .optional() + } +} + +impl BiRnnConfig { + /// Initialize a new [Bidirectional RNN](BiRnn) module. + pub fn init(&self, device: &Device) -> BiRnn { + // Internal RNNs always use batch_first=true; BiRnn handles layout conversion + let base_config = RnnConfig::new(self.d_input, self.d_hidden, self.bias) + .with_initializer(self.initializer.clone()) + .with_batch_first(true) + .with_clip(self.clip) + .with_hidden_activation(self.hidden_activation.clone()); + + BiRnn { + forward: base_config.clone().init(device), + reverse: base_config.init(device), + d_hidden: self.d_hidden, + batch_first: self.batch_first, + } + } +} + +impl BiRnn { + /// Applies the forward pass on the input tensor. This Bidirectional RNN implementation + /// returns the state for each element in a sequence (i.e., across seq_length) and a final state. + /// + /// ## Parameters: + /// - batched_input: The input tensor of shape: + /// - `[batch_size, sequence_length, input_size]` if `batch_first` is true (default) + /// - `[sequence_length, batch_size, input_size]` if `batch_first` is false + /// - state: An optional `RnnState` representing the hidden state. + /// Each state tensor has shape `[2, batch_size, hidden_size]`. + /// If no initial state is provided, these tensors are initialized to zeros. + /// + /// ## Returns: + /// - output: A tensor represents the output features of RNN. Shape: + /// - `[batch_size, sequence_length, hidden_size * 2]` if `batch_first` is true + /// - `[sequence_length, batch_size, hidden_size * 2]` if `batch_first` is false + /// - state: A `RnnState` represents the final forward and reverse states. + /// The `state.hidden` have the shape `[2, batch_size, hidden_size]`. + pub fn forward( + &self, + batched_input: Tensor<3>, + state: Option>, + ) -> (Tensor<3>, RnnState<3>) { + // Convert to batch-first layout internally if needed + let batched_input = if self.batch_first { + batched_input + } else { + batched_input.swap_dims(0, 1) + }; + + let device = batched_input.clone().device(); + let [batch_size, seq_length, _] = batched_input.shape().dims(); + + let [init_state_forward, init_state_reverse] = match state { + Some(state) => { + let hidden_state_forward = state + .hidden + .clone() + .slice([0..1, 0..batch_size, 0..self.d_hidden]) + .squeeze_dim(0); + let hidden_state_reverse = state + .hidden + .slice([1..2, 0..batch_size, 0..self.d_hidden]) + .squeeze_dim(0); + + [ + Some(RnnState::new(hidden_state_forward)), + Some(RnnState::new(hidden_state_reverse)), + ] + } + None => [None, None], + }; + + // forward direction + let (batched_hidden_state_forward, final_state_forward) = self + .forward + .forward(batched_input.clone(), init_state_forward); + + // reverse direction + let (batched_hidden_state_reverse, final_state_reverse) = self.reverse.forward_iter( + batched_input.iter_dim(1).rev().zip((0..seq_length).rev()), + init_state_reverse, + batch_size, + seq_length, + &device, + ); + + let output = Tensor::cat( + [batched_hidden_state_forward, batched_hidden_state_reverse].to_vec(), + 2, + ); + + // Convert output back to seq-first layout if needed + let output = if self.batch_first { + output + } else { + output.swap_dims(0, 1) + }; + + let state = RnnState::new(Tensor::stack( + [final_state_forward.hidden, final_state_reverse.hidden].to_vec(), + 0, + )); + + (output, state) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Linear; + use burn::module::Param; + use burn::tensor::{Distribution, TensorData}; + use burn::tensor::{ElementConversion, Tolerance}; + type FT = f32; + + fn create_single_feature_gate_controller( + weights: f32, + biases: f32, + device: &Device, + ) -> GateController { + let record_1 = Linear { + weight: Param::from_data(TensorData::from([[weights]]), device), + bias: Some(Param::from_data(TensorData::from([biases]), device)), + }; + let record_2 = Linear { + weight: Param::from_data(TensorData::from([[weights]]), device), + bias: Some(Param::from_data(TensorData::from([biases]), device)), + }; + GateController::create_with_weights(record_1, record_2) + } + + #[test] + fn test_with_uniform_initializer() { + let device = Device::default(); + device.seed(0); + + let config = RnnConfig::new(5, 5, false) + .with_initializer(Initializer::Uniform { min: 0.0, max: 1.0 }); + let rnn = config.init(&device); + + let gate_to_data = |gate: GateController| gate.input_transform.weight.val().to_data(); + + gate_to_data(rnn.gate).assert_within_range::(0.elem()..1.elem()); + } + + /// Test forward pass with simple input vector. + /// + /// Simple RNN: h_t = tanh(W_input @ x_t + W_hidden @ h_{t-1} + b) + /// With input=0.1, weight_input=0.5, bias=0.0, h_0=0.0, weight_hidden=0.5 + /// h_t = tanh(0.5*0.1 + 0.5*0) = tanh(0.05) = 0.04995 + #[test] + fn test_forward_single_input_single_feature() { + let device = Device::default(); + device.seed(0); + + let config = RnnConfig::new(1, 1, false); + let device = Default::default(); + let mut rnn = config.init(&device); + + rnn.gate = create_single_feature_gate_controller(0.5, 0.0, &device); + + // single timestep with single feature + let input = Tensor::<3>::from_data(TensorData::from([[[0.1]]]), &device); + + let (output, state) = rnn.forward(input, None); + + let tolerance = Tolerance::default(); + let expected = TensorData::from([[0.04995]]); + state + .hidden + .to_data() + .assert_approx_eq::(&expected, tolerance); + + output + .select(0, Tensor::arange(0..1, &device)) + .squeeze_dim::<2>(0) + .to_data() + .assert_approx_eq::(&state.hidden.to_data(), tolerance); + } + + #[test] + fn test_batched_forward_pass_batch_of_one() { + let device = Default::default(); + let rnn = RnnConfig::new(64, 1024, true).init(&device); + let batched_input = Tensor::<3>::random([1, 2, 64], Distribution::Default, &device); + + let (output, state) = rnn.forward(batched_input, None); + assert_eq!(output.dims(), [1, 2, 1024]); + assert_eq!(state.hidden.dims(), [1, 1024]); + } + + #[test] + #[cfg(feature = "std")] + fn test_batched_backward_pass() { + use burn::tensor::Shape; + let device = Device::default().autodiff(); + let rnn = RnnConfig::new(64, 32, true).init(&device); + let shape: Shape = [8, 10, 64].into(); + let batched_input = Tensor::<3>::random(shape, Distribution::Default, &device); + + let (output, _) = rnn.forward(batched_input.clone(), None); + let fake_loss = output; + let grads = fake_loss.backward(); + + let some_gradient = rnn.gate.hidden_transform.weight.grad(&grads).unwrap(); + + // Asserts that the gradients exist and are non-zero + assert_ne!( + some_gradient + .any() + .into_data() + .iter::() + .next() + .unwrap(), + 0.0 + ); + } + + #[test] + fn test_bidirectional() { + let device = Device::default(); + device.seed(0); + + let config = BiRnnConfig::new(2, 3, true); + let mut rnn = config.init(&device); + + fn create_gate_controller( + input_weights: [[f32; D1]; D2], + input_biases: [f32; D1], + hidden_weights: [[f32; D1]; D1], + hidden_biases: [f32; D1], + device: &Device, + ) -> GateController { + let input_record = Linear { + weight: Param::from_data(TensorData::from(input_weights), device), + bias: Some(Param::from_data(TensorData::from(input_biases), device)), + }; + let hidden_record = Linear { + weight: Param::from_data(TensorData::from(hidden_weights), device), + bias: Some(Param::from_data(TensorData::from(hidden_biases), device)), + }; + GateController::create_with_weights(input_record, hidden_record) + } + + // [batch_size=1, seq_length=4, input_size=2] + let input = Tensor::<3>::from_data( + TensorData::from([[ + [0.949, -0.861], + [0.892, 0.927], + [-0.173, -0.301], + [-0.081, 0.992], + ]]), + &device, + ); + + // [2, batch_size=1, hidden_size=3] + let h0 = Tensor::<3>::from_data( + TensorData::from([[[0.280, 0.360, -1.242]], [[-0.588, 0.729, -0.788]]]), + &device, + ); + + rnn.forward.gate = create_gate_controller( + // input_weights: [input_size=2, hidden_size=3] + [[0.367, 0.091, 0.342], [0.322, 0.533, 0.059]], + // input_biases: [hidden_size=3] + [-0.196, 0.354, 0.209], + // hidden_weights: [hidden_size=3, hidden_size=3] + [ + [-0.320, 0.232, -0.165], + [0.093, -0.572, -0.315], + [-0.467, 0.325, 0.046], + ], + // hidden_biases: [hidden_size=3] + [0.181, -0.190, -0.245], + &device, + ); + + rnn.reverse.gate = create_gate_controller( + [[-0.055, 0.506, 0.247], [-0.369, 0.178, -0.258]], + [0.540, -0.164, 0.033], + [ + [0.159, 0.180, -0.037], + [-0.443, 0.485, -0.488], + [0.098, -0.085, -0.140], + ], + [-0.510, 0.105, 0.114], + &device, + ); + + // [batch_size=1, sequence_length=4, hidden_size * 2 = 6] + // The expected output values were computed from PyTorch + let expected_output_with_init_state = TensorData::from([[ + [0.5226, -0.6370, 0.0210, 0.0685, 0.3867, 0.3602], + [0.3580, 0.8431, 0.4129, -0.3175, 0.4374, 0.1766], + [-0.3837, -0.2703, -0.3957, -0.1542, -0.1122, 0.0725], + [0.5059, 0.5527, 0.1244, -0.6779, 0.3725, -0.3387], + ]]); + let expected_output_without_init_state = TensorData::from([[ + [0.0560, -0.2056, 0.2334, 0.0892, 0.3912, 0.3607], + [0.4340, 0.7378, 0.3714, -0.2394, 0.4235, 0.2002], + [-0.3962, -0.2097, -0.3798, 0.0532, -0.2067, 0.1727], + [0.5075, 0.5298, 0.1083, -0.3200, 0.0764, -0.1282], + ]]); + + //`[2, batch_size=1, hidden_size=3]` + let expected_hn_with_init_state = + TensorData::from([[[0.5059, 0.5527, 0.1244]], [[0.0685, 0.3867, 0.3602]]]); + let expected_hn_without_init_state = + TensorData::from([[[0.5075, 0.5298, 0.1083]], [[0.0892, 0.3912, 0.3607]]]); + + let (output_with_init_state, state_with_init_state) = + rnn.forward(input.clone(), Some(RnnState::new(h0))); + let (output_without_init_state, state_without_init_state) = rnn.forward(input, None); + + let tolerance = Tolerance::permissive(); + output_with_init_state + .to_data() + .assert_approx_eq::(&expected_output_with_init_state, tolerance); + output_without_init_state + .to_data() + .assert_approx_eq::(&expected_output_without_init_state, tolerance); + state_with_init_state + .hidden + .to_data() + .assert_approx_eq::(&expected_hn_with_init_state, tolerance); + state_without_init_state + .hidden + .to_data() + .assert_approx_eq::(&expected_hn_without_init_state, tolerance); + } + + #[test] + fn display_rnn() { + let config = RnnConfig::new(2, 3, true); + + let layer = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{layer}"), + "Rnn {d_input: 2, d_hidden: 3, bias: true, params: 21}" + ); + } + + #[test] + fn display_birnn() { + let config = BiRnnConfig::new(2, 3, true); + + let layer = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{layer}"), + "BiRnn {d_input: 2, d_hidden: 3, bias: true, params: 42}" + ); + } + + #[test] + fn test_rnn_clipping() { + let device = Default::default(); + + // Create Rnn with clipping enabled + let clip_value = 0.3; + let config = RnnConfig::new(4, 8, true).with_clip(Some(clip_value)); + let rnn = config.init(&device); + + let input = Tensor::<3>::random([2, 5, 4], Distribution::Default, &device); + let (_, state) = rnn.forward(input, None); + + // Verify output values are within the clip range + let hidden_state: Vec = state.hidden.to_data().to_vec().unwrap(); + for val in hidden_state { + assert!( + val >= -clip_value as f32 && val <= clip_value as f32, + "Value {} is outside clip range [-{}, {}]", + val, + clip_value, + clip_value + ); + } + } + + #[test] + fn test_forward_reverse_sequence() { + let device = Device::default(); + device.seed(0); + + // Create RNN with reverse=true to process sequence in reverse order + let config = RnnConfig::new(1, 1, false).with_reverse(true); + let mut rnn = config.init(&device); + + rnn.gate = create_single_feature_gate_controller(0.5, 0.0, &device); + + // Create input with 3 timesteps: [0.1, 0.2, 0.3] + // Shape: [batch_size=1, seq_length=3, input_features=1] + let input = Tensor::<3>::from_data(TensorData::from([[[0.1], [0.2], [0.3]]]), &device); + + let (output, state) = rnn.forward(input, None); + + // With reverse=true and weight=0.5, sequence is processed in reverse: + // t=2 (last): h = tanh(0.5*0.3 + 0.5*0) = tanh(0.15) ≈ 0.1488850 + // t=1 (mid): h = tanh(0.5*0.2 + 0.5*0.1488850) ≈ 0.17269433 + // t=0 (first): h = tanh(0.5*0.1 + 0.5*0.17269433) ≈ 0.135508 + let expected_final_hidden = TensorData::from([[0.135508]]); + + let tolerance = Tolerance::default(); + state + .hidden + .to_data() + .assert_approx_eq::(&expected_final_hidden, tolerance); + + // Verify output tensor has correct shape and matches state at final timestep + assert_eq!(output.dims(), [1, 3, 1]); + } +} diff --git a/crates/burn-nn/src/modules/rnn/gate_controller.rs b/crates/burn-nn/src/modules/rnn/gate_controller.rs new file mode 100644 index 0000000..cc64c38 --- /dev/null +++ b/crates/burn-nn/src/modules/rnn/gate_controller.rs @@ -0,0 +1,75 @@ +use burn_core as burn; + +use crate::{Linear, LinearConfig, LinearLayout}; +use burn::module::{Initializer, Module}; +use burn::tensor::{Device, Tensor}; + +/// A GateController represents a gate in an LSTM cell. An +/// LSTM cell generally contains three gates: an input gate, +/// forget gate, and output gate. Additionally, cell gate +/// is just used to compute the cell state. +/// +/// An Lstm gate is modeled as two linear transformations. +/// The results of these transformations are used to calculate +/// the gate's output. +#[derive(Module, Debug)] +pub struct GateController { + /// Represents the affine transformation applied to input vector + pub input_transform: Linear, + /// Represents the affine transformation applied to the hidden state + pub hidden_transform: Linear, +} + +impl GateController { + /// Initialize a new [gate_controller](GateController) module. + pub fn new( + d_input: usize, + d_output: usize, + bias: bool, + initializer: Initializer, + device: &Device, + ) -> Self { + Self { + input_transform: LinearConfig { + d_input, + d_output, + bias, + initializer: initializer.clone(), + layout: LinearLayout::Row, + } + .init(device), + hidden_transform: LinearConfig { + d_input: d_output, + d_output, + bias, + initializer, + layout: LinearLayout::Row, + } + .init(device), + } + } + + /// Helper function for performing weighted matrix product for a gate and adds + /// bias, if any. + /// + /// Mathematically, performs `Wx*X + Wh*H + b`, where: + /// Wx = weight matrix for the connection to input vector X + /// Wh = weight matrix for the connection to hidden state H + /// X = input vector + /// H = hidden state + /// b = bias terms + pub fn gate_product(&self, input: Tensor<2>, hidden: Tensor<2>) -> Tensor<2> { + self.input_transform.forward(input) + self.hidden_transform.forward(hidden) + } + + /// Used to initialize a gate controller with known weight layers, + /// allowing for predictable behavior. Used only for testing in + /// lstm. + #[cfg(test)] + pub fn create_with_weights(input_transform: Linear, hidden_transform: Linear) -> Self { + Self { + input_transform, + hidden_transform, + } + } +} diff --git a/crates/burn-nn/src/modules/rnn/gru.rs b/crates/burn-nn/src/modules/rnn/gru.rs new file mode 100644 index 0000000..bba73e8 --- /dev/null +++ b/crates/burn-nn/src/modules/rnn/gru.rs @@ -0,0 +1,998 @@ +use burn_core as burn; + +use super::gate_controller::GateController; +use crate::activation::{Activation, ActivationConfig}; +use burn::config::Config; +use burn::module::Initializer; +use burn::module::Module; +use burn::module::{Content, DisplaySettings, ModuleDisplay}; +use burn::tensor::Device; +use burn::tensor::Tensor; + +/// Configuration to create a [gru](Gru) module using the [init function](GruConfig::init). +#[derive(Config, Debug)] +pub struct GruConfig { + /// The size of the input features. + pub d_input: usize, + /// The size of the hidden state. + pub d_hidden: usize, + /// If a bias should be applied during the Gru transformation. + pub bias: bool, + /// If reset gate should be applied after weight multiplication. + /// + /// This configuration option controls how the reset gate is applied to the hidden state. + /// * `true` - (Default) Match the initial arXiv version of the paper [Learning Phrase Representations using RNN Encoder-Decoder for + /// Statistical Machine Translation (v1)](https://arxiv.org/abs/1406.1078v1) and apply the reset gate after multiplication by + /// the weights. This matches the behavior of [PyTorch GRU](https://pytorch.org/docs/stable/generated/torch.nn.GRU.html#torch.nn.GRU). + /// * `false` - Match the most recent revision of [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine + /// Translation (v3)](https://arxiv.org/abs/1406.1078) and apply the reset gate before the weight multiplication. + /// + /// The differing implementations can give slightly different numerical results and have different efficiencies. For more + /// motivation for why the `true` can be more efficient see [Optimizing RNNs with Differentiable Graphs](https://svail.github.io/diff_graphs). + /// + /// To set this field to `false` use [`with_reset_after`](`GruConfig::with_reset_after`). + #[config(default = "true")] + pub reset_after: bool, + /// Gru initializer + #[config(default = "Initializer::XavierNormal{gain:1.0}")] + pub initializer: Initializer, + /// Activation function for the update and reset gates. + /// Default is Sigmoid, which is standard for GRU gates. + #[config(default = "ActivationConfig::Sigmoid")] + pub gate_activation: ActivationConfig, + /// Activation function for the new/candidate gate. + /// Default is Tanh, which is standard for GRU. + #[config(default = "ActivationConfig::Tanh")] + pub hidden_activation: ActivationConfig, + /// Optional hidden state clip threshold. If provided, hidden state values are clipped + /// to the range `[-clip, +clip]` after each timestep. This can help prevent + /// exploding values during inference. + pub clip: Option, +} + +/// The Gru (Gated recurrent unit) module. This implementation is for a unidirectional, stateless, Gru. +/// +/// Introduced in the paper: [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://arxiv.org/abs/1406.1078). +/// +/// Should be created with [GruConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Gru { + /// The update gate controller. + pub update_gate: GateController, + /// The reset gate controller. + pub reset_gate: GateController, + /// The new gate controller. + pub new_gate: GateController, + /// The size of the hidden state. + pub d_hidden: usize, + /// If reset gate should be applied after weight multiplication. + pub reset_after: bool, + /// Activation function for gates (update, reset). + pub gate_activation: Activation, + /// Activation function for new/candidate gate. + pub hidden_activation: Activation, + /// Optional hidden state clip threshold. + pub clip: Option, +} + +impl ModuleDisplay for Gru { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let [d_input, _] = self.update_gate.input_transform.weight.shape().dims(); + let bias = self.update_gate.input_transform.bias.is_some(); + + content + .add("d_input", &d_input) + .add("d_hidden", &self.d_hidden) + .add("bias", &bias) + .add("reset_after", &self.reset_after) + .optional() + } +} + +impl GruConfig { + /// Initialize a new [gru](Gru) module. + pub fn init(&self, device: &Device) -> Gru { + let d_output = self.d_hidden; + + let update_gate = GateController::new( + self.d_input, + d_output, + self.bias, + self.initializer.clone(), + device, + ); + let reset_gate = GateController::new( + self.d_input, + d_output, + self.bias, + self.initializer.clone(), + device, + ); + let new_gate = GateController::new( + self.d_input, + d_output, + self.bias, + self.initializer.clone(), + device, + ); + + Gru { + update_gate, + reset_gate, + new_gate, + d_hidden: self.d_hidden, + reset_after: self.reset_after, + gate_activation: self.gate_activation.init(device), + hidden_activation: self.hidden_activation.init(device), + clip: self.clip, + } + } +} + +impl Gru { + /// Applies the forward pass on the input tensor. This GRU implementation + /// returns a state tensor with dimensions `[batch_size, sequence_length, hidden_size]`. + /// + /// # Parameters + /// - batched_input: `[batch_size, sequence_length, input_size]`. + /// - state: An optional tensor representing an initial cell state with dimensions + /// `[batch_size, hidden_size]`. If none is provided, an empty state will be used. + /// + /// # Returns + /// - output: `[batch_size, sequence_length, hidden_size]` + pub fn forward(&self, batched_input: Tensor<3>, state: Option>) -> Tensor<3> { + let device = batched_input.device(); + let [batch_size, seq_length, _] = batched_input.shape().dims(); + + self.forward_iter( + batched_input.iter_dim(1).zip(0..seq_length), + state, + batch_size, + seq_length, + &device, + ) + .0 + } + + /// Forward pass variant that accepts an iterator over timesteps. + /// Used by BiGru to process sequences in either direction. + /// + /// # Parameters + /// - input_timestep_iter: Iterator yielding (input_tensor, timestep_index) pairs. + /// The timestep_index determines where in the output tensor to store results. + /// - state: Optional initial hidden state with shape `[batch_size, hidden_size]`. + /// - batch_size: Batch size of the input. + /// - seq_length: Sequence length of the input. + /// - device: Device to create tensors on. + /// + /// # Returns + /// - output: `[batch_size, sequence_length, hidden_size]` + /// - final_hidden: Final hidden state `[batch_size, hidden_size]` + pub(crate) fn forward_iter, usize)>>( + &self, + input_timestep_iter: I, + state: Option>, + batch_size: usize, + seq_length: usize, + device: &Device, + ) -> (Tensor<3>, Tensor<2>) { + let mut batched_hidden_state = + Tensor::empty([batch_size, seq_length, self.d_hidden], device); + + let mut hidden_t = match state { + Some(state) => state, + None => Tensor::zeros([batch_size, self.d_hidden], device), + }; + + for (input_t, t) in input_timestep_iter { + let input_t = input_t.squeeze_dim(1); + + // u(pdate)g(ate) tensors + let biased_ug_input_sum = + self.gate_product(&input_t, &hidden_t, None, &self.update_gate); + let update_values = self.gate_activation.forward(biased_ug_input_sum); + + // r(eset)g(ate) tensors + let biased_rg_input_sum = + self.gate_product(&input_t, &hidden_t, None, &self.reset_gate); + let reset_values = self.gate_activation.forward(biased_rg_input_sum); + + // n(ew)g(ate) tensor + let biased_ng_input_sum = if self.reset_after { + self.gate_product(&input_t, &hidden_t, Some(&reset_values), &self.new_gate) + } else { + let reset_t = hidden_t.clone().mul(reset_values); + self.gate_product(&input_t, &reset_t, None, &self.new_gate) + }; + let candidate_state = self.hidden_activation.forward(biased_ng_input_sum); + + // calculate linear interpolation between previous hidden state and candidate state: + // h_t = (1 - z_t) * g_t + z_t * h_{t-1} + let one_minus_z = update_values.clone().neg().add_scalar(1.0); + hidden_t = candidate_state.mul(one_minus_z) + update_values.mul(hidden_t); + + // Apply hidden state clipping if configured + if let Some(clip) = self.clip { + hidden_t = hidden_t.clamp(-clip, clip); + } + + let unsqueezed_hidden_state = hidden_t.clone().unsqueeze_dim(1); + + batched_hidden_state = batched_hidden_state.slice_assign( + [0..batch_size, t..(t + 1), 0..self.d_hidden], + unsqueezed_hidden_state, + ); + } + + (batched_hidden_state, hidden_t) + } + + /// Helper function for performing weighted matrix product for a gate and adds + /// bias, if any, and optionally applies reset to hidden state. + /// + /// Mathematically, performs `Wx*X + r .* (Wh*H + b)`, where: + /// Wx = weight matrix for the connection to input vector X + /// Wh = weight matrix for the connection to hidden state H + /// X = input vector + /// H = hidden state + /// b = bias terms + /// r = reset state + fn gate_product( + &self, + input: &Tensor<2>, + hidden: &Tensor<2>, + reset: Option<&Tensor<2>>, + gate: &GateController, + ) -> Tensor<2> { + let input_product = input.clone().matmul(gate.input_transform.weight.val()); + let hidden_product = hidden.clone().matmul(gate.hidden_transform.weight.val()); + + let input_part = match &gate.input_transform.bias { + Some(bias) => input_product + bias.val().unsqueeze(), + None => input_product, + }; + + let hidden_part = match &gate.hidden_transform.bias { + Some(bias) => hidden_product + bias.val().unsqueeze(), + None => hidden_product, + }; + + match reset { + Some(r) => input_part + r.clone().mul(hidden_part), + None => input_part + hidden_part, + } + } +} + +/// Configuration to create a [BiGru](BiGru) module using the [init function](BiGruConfig::init). +#[derive(Config, Debug)] +pub struct BiGruConfig { + /// The size of the input features. + pub d_input: usize, + /// The size of the hidden state. + pub d_hidden: usize, + /// If a bias should be applied during the BiGru transformation. + pub bias: bool, + /// If reset gate should be applied after weight multiplication. + #[config(default = "true")] + pub reset_after: bool, + /// BiGru initializer + #[config(default = "Initializer::XavierNormal{gain:1.0}")] + pub initializer: Initializer, + /// If true, the input tensor is expected to be `[batch_size, seq_length, input_size]`. + /// If false, the input tensor is expected to be `[seq_length, batch_size, input_size]`. + #[config(default = true)] + pub batch_first: bool, + /// Activation function for the update and reset gates. + #[config(default = "ActivationConfig::Sigmoid")] + pub gate_activation: ActivationConfig, + /// Activation function for the new/candidate gate. + #[config(default = "ActivationConfig::Tanh")] + pub hidden_activation: ActivationConfig, + /// Optional hidden state clip threshold. + pub clip: Option, +} + +/// The BiGru module. This implementation is for Bidirectional GRU. +/// +/// Based on the paper: [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://arxiv.org/abs/1406.1078). +/// +/// Should be created with [BiGruConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct BiGru { + /// GRU for the forward direction. + pub forward: Gru, + /// GRU for the reverse direction. + pub reverse: Gru, + /// The size of the hidden state. + pub d_hidden: usize, + /// If true, input is `[batch_size, seq_length, input_size]`. + /// If false, input is `[seq_length, batch_size, input_size]`. + pub batch_first: bool, +} + +impl ModuleDisplay for BiGru { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let [d_input, _] = self + .forward + .update_gate + .input_transform + .weight + .shape() + .dims(); + let bias = self.forward.update_gate.input_transform.bias.is_some(); + + content + .add("d_input", &d_input) + .add("d_hidden", &self.d_hidden) + .add("bias", &bias) + .optional() + } +} + +impl BiGruConfig { + /// Initialize a new [Bidirectional GRU](BiGru) module. + pub fn init(&self, device: &Device) -> BiGru { + // Internal GRUs always use batch_first=true; BiGru handles layout conversion + let base_config = GruConfig::new(self.d_input, self.d_hidden, self.bias) + .with_initializer(self.initializer.clone()) + .with_reset_after(self.reset_after) + .with_gate_activation(self.gate_activation.clone()) + .with_hidden_activation(self.hidden_activation.clone()) + .with_clip(self.clip); + + BiGru { + forward: base_config.clone().init(device), + reverse: base_config.init(device), + d_hidden: self.d_hidden, + batch_first: self.batch_first, + } + } +} + +impl BiGru { + /// Applies the forward pass on the input tensor. This Bidirectional GRU implementation + /// returns the state for each element in a sequence (i.e., across seq_length) and a final state. + /// + /// ## Parameters: + /// - batched_input: The input tensor of shape: + /// - `[batch_size, sequence_length, input_size]` if `batch_first` is true (default) + /// - `[sequence_length, batch_size, input_size]` if `batch_first` is false + /// - state: An optional tensor representing the initial hidden state with shape + /// `[2, batch_size, hidden_size]`. If no initial state is provided, it is initialized to zeros. + /// + /// ## Returns: + /// - output: A tensor representing the output features. Shape: + /// - `[batch_size, sequence_length, hidden_size * 2]` if `batch_first` is true + /// - `[sequence_length, batch_size, hidden_size * 2]` if `batch_first` is false + /// - state: The final forward and reverse hidden states stacked along dimension 0 + /// with shape `[2, batch_size, hidden_size]`. + pub fn forward( + &self, + batched_input: Tensor<3>, + state: Option>, + ) -> (Tensor<3>, Tensor<3>) { + // Convert to batch-first layout internally if needed + let batched_input = if self.batch_first { + batched_input + } else { + batched_input.swap_dims(0, 1) + }; + + let device = batched_input.clone().device(); + let [batch_size, seq_length, _] = batched_input.shape().dims(); + + let [init_state_forward, init_state_reverse] = match state { + Some(state) => { + let hidden_state_forward = state + .clone() + .slice([0..1, 0..batch_size, 0..self.d_hidden]) + .squeeze_dim(0); + let hidden_state_reverse = state + .slice([1..2, 0..batch_size, 0..self.d_hidden]) + .squeeze_dim(0); + + [Some(hidden_state_forward), Some(hidden_state_reverse)] + } + None => [None, None], + }; + + // forward direction + let (batched_hidden_state_forward, final_state_forward) = self.forward.forward_iter( + batched_input.clone().iter_dim(1).zip(0..seq_length), + init_state_forward, + batch_size, + seq_length, + &device, + ); + + // reverse direction + let (batched_hidden_state_reverse, final_state_reverse) = self.reverse.forward_iter( + batched_input.iter_dim(1).rev().zip((0..seq_length).rev()), + init_state_reverse, + batch_size, + seq_length, + &device, + ); + + let output = Tensor::cat( + [batched_hidden_state_forward, batched_hidden_state_reverse].to_vec(), + 2, + ); + + // Convert output back to seq-first layout if needed + let output = if self.batch_first { + output + } else { + output.swap_dims(0, 1) + }; + + let state = Tensor::stack([final_state_forward, final_state_reverse].to_vec(), 0); + + (output, state) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Linear; + use burn::module::Param; + use burn::tensor::Tolerance; + use burn::tensor::{Distribution, TensorData}; + + type FT = f32; + + fn init_gru(reset_after: bool, device: &Device) -> Gru { + fn create_gate_controller(weights: f32, biases: f32, device: &Device) -> GateController { + let record_1 = Linear { + weight: Param::from_data(TensorData::from([[weights]]), device), + bias: Some(Param::from_data(TensorData::from([biases]), device)), + }; + let record_2 = Linear { + weight: Param::from_data(TensorData::from([[weights]]), device), + bias: Some(Param::from_data(TensorData::from([biases]), device)), + }; + GateController::create_with_weights(record_1, record_2) + } + + let config = GruConfig::new(1, 1, false).with_reset_after(reset_after); + let mut gru = config.init(device); + + gru.update_gate = create_gate_controller(0.5, 0.0, device); + gru.reset_gate = create_gate_controller(0.6, 0.0, device); + gru.new_gate = create_gate_controller(0.7, 0.0, device); + gru + } + + /// Test forward pass with simple input vector. + /// + /// z_t = sigmoid(0.5*0.1 + 0.5*0) = 0.5125 + /// r_t = sigmoid(0.6*0.1 + 0.*0) = 0.5150 + /// g_t = tanh(0.7*0.1 + 0.7*0) = 0.0699 + /// + /// h_t = z_t * h' + (1 - z_t) * g_t = 0.0341 + #[test] + fn tests_forward_single_input_single_feature() { + let device = Device::default(); + device.seed(0); + + let mut gru = init_gru(false, &device); + + let input = Tensor::<3>::from_data(TensorData::from([[[0.1]]]), &device); + let expected = TensorData::from([[0.034]]); + + // Reset gate applied to hidden state before the matrix multiplication + let state = gru.forward(input.clone(), None); + + let output = state + .select(0, Tensor::arange(0..1, &device)) + .squeeze_dim::<2>(0); + + let tolerance = Tolerance::default(); + output + .to_data() + .assert_approx_eq::(&expected, tolerance); + + // Reset gate applied to hidden state after the matrix multiplication + gru.reset_after = true; // override forward behavior + let state = gru.forward(input, None); + + let output = state + .select(0, Tensor::arange(0..1, &device)) + .squeeze_dim::<2>(0); + + output + .to_data() + .assert_approx_eq::(&expected, tolerance); + } + + #[test] + fn tests_forward_seq_len_3() { + let device = Device::default(); + device.seed(0); + let mut gru = init_gru(true, &device); + + let input = Tensor::<3>::from_data(TensorData::from([[[0.1], [0.2], [0.3]]]), &device); + let expected = TensorData::from([[0.0341], [0.0894], [0.1575]]); + + let result = gru.forward(input.clone(), None); + let output = result + .select(0, Tensor::arange(0..1, &device)) + .squeeze_dim::<2>(0); + + let tolerance = Tolerance::default(); + output + .to_data() + .assert_approx_eq::(&expected, tolerance); + + // Reset gate applied to hidden state before the matrix multiplication + gru.reset_after = false; // override forward behavior + let state = gru.forward(input, None); + + let output = state + .select(0, Tensor::arange(0..1, &device)) + .squeeze_dim::<2>(0); + + output + .to_data() + .assert_approx_eq::(&expected, tolerance); + } + + #[test] + fn test_batched_forward_pass() { + let device = Default::default(); + let gru = GruConfig::new(64, 1024, true).init(&device); + let batched_input = Tensor::<3>::random([8, 10, 64], Distribution::Default, &device); + + let hidden_state = gru.forward(batched_input, None); + + assert_eq!(&*hidden_state.shape(), [8, 10, 1024]); + } + + #[test] + fn display() { + let config = GruConfig::new(2, 8, true); + + let layer = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{layer}"), + "Gru {d_input: 2, d_hidden: 8, bias: true, reset_after: true, params: 288}" + ); + } + + #[test] + fn test_bigru_batched_forward_pass() { + let device = Default::default(); + let bigru = BiGruConfig::new(64, 1024, true).init(&device); + let batched_input = Tensor::<3>::random([8, 10, 64], Distribution::Default, &device); + + let (output, state) = bigru.forward(batched_input, None); + + // Output should have hidden_size * 2 features (forward + reverse concatenated) + assert_eq!(&*output.shape(), [8, 10, 2048]); + // State should have shape [2, batch_size, hidden_size] + assert_eq!(&*state.shape(), [2, 8, 1024]); + } + + #[test] + fn test_bigru_with_initial_state() { + let device = Default::default(); + let bigru = BiGruConfig::new(32, 64, true).init(&device); + let batched_input = Tensor::<3>::random([4, 5, 32], Distribution::Default, &device); + let initial_state = Tensor::<3>::random([2, 4, 64], Distribution::Default, &device); + + let (output, state) = bigru.forward(batched_input, Some(initial_state)); + + assert_eq!(&*output.shape(), [4, 5, 128]); + assert_eq!(&*state.shape(), [2, 4, 64]); + } + + #[test] + fn test_bigru_seq_first() { + let device = Default::default(); + let bigru = BiGruConfig::new(32, 64, true) + .with_batch_first(false) + .init(&device); + // Input shape: [seq_length, batch_size, input_size] when batch_first=false + let batched_input = Tensor::<3>::random([5, 4, 32], Distribution::Default, &device); + + let (output, state) = bigru.forward(batched_input, None); + + // Output shape: [seq_length, batch_size, hidden_size * 2] + assert_eq!(&*output.shape(), [5, 4, 128]); + assert_eq!(&*state.shape(), [2, 4, 64]); + } + + /// Test BiGru against PyTorch reference implementation. + /// Expected values computed with PyTorch nn.GRU(bidirectional=True). + #[test] + fn test_bigru_against_pytorch() { + let device = Device::default(); + device.seed(0); + + let config = BiGruConfig::new(2, 3, true); + let mut bigru = config.init(&device); + + fn create_gate_controller( + input_weights: [[f32; D1]; D2], + input_biases: [f32; D1], + hidden_weights: [[f32; D1]; D1], + hidden_biases: [f32; D1], + device: &Device, + ) -> GateController { + let input_record = Linear { + weight: Param::from_data(TensorData::from(input_weights), device), + bias: Some(Param::from_data(TensorData::from(input_biases), device)), + }; + let hidden_record = Linear { + weight: Param::from_data(TensorData::from(hidden_weights), device), + bias: Some(Param::from_data(TensorData::from(hidden_biases), device)), + }; + GateController::create_with_weights(input_record, hidden_record) + } + + let input = Tensor::<3>::from_data( + TensorData::from([[ + [0.949, -0.861], + [0.892, 0.927], + [-0.173, -0.301], + [-0.081, 0.992], + ]]), + &device, + ); + let h0 = Tensor::<3>::from_data( + TensorData::from([[[0.280, 0.360, -1.242]], [[-0.588, 0.729, -0.788]]]), + &device, + ); + + // Forward GRU gates (weights from PyTorch with seed 42, transposed for burn) + bigru.forward.update_gate = create_gate_controller( + [[-0.2811, 0.5090, 0.5018], [0.3391, -0.4236, 0.1081]], + [0.2932, -0.3519, -0.5715], + [ + [-0.3471, 0.5214, 0.0961], + [0.0545, -0.4904, -0.1875], + [-0.5702, 0.4457, 0.3568], + ], + [-0.0100, 0.4518, -0.4102], + &device, + ); + + bigru.forward.reset_gate = create_gate_controller( + [[0.4414, -0.1353, -0.1265], [0.4792, 0.5304, 0.1165]], + [-0.2524, 0.3333, 0.1033], + [ + [-0.2695, -0.0677, -0.4557], + [0.1472, -0.2345, -0.2662], + [-0.2660, 0.3830, -0.1630], + ], + [0.1663, 0.2391, 0.1826], + &device, + ); + + bigru.forward.new_gate = create_gate_controller( + [[0.4266, 0.2784, 0.4451], [0.0782, -0.0815, 0.0853]], + [-0.2231, -0.4428, 0.4737], + [ + [0.0900, -0.1821, 0.2430], + [0.4665, 0.1551, 0.5155], + [0.0631, -0.1566, 0.3337], + ], + [0.0364, -0.3941, 0.1780], + &device, + ); + + // Reverse GRU gates + bigru.reverse.update_gate = create_gate_controller( + [[-0.3444, 0.1924, -0.4765], [0.5193, 0.5556, -0.5727]], + [0.1090, 0.1779, -0.5385], + [ + [0.1221, 0.3925, 0.5287], + [-0.1472, -0.4187, -0.1948], + [0.3441, -0.3082, -0.2047], + ], + [0.0016, -0.2148, -0.0400], + &device, + ); + + bigru.reverse.reset_gate = create_gate_controller( + [[-0.1988, -0.1203, -0.3422], [0.1769, 0.4788, -0.3443]], + [-0.5053, -0.3676, 0.5771], + [ + [-0.3936, 0.3504, -0.4486], + [0.3063, -0.1370, -0.2914], + [-0.2334, 0.3303, 0.1760], + ], + [-0.5080, -0.2488, -0.3456], + &device, + ); + + bigru.reverse.new_gate = create_gate_controller( + [[-0.4517, 0.2339, 0.4797], [-0.3884, 0.2067, -0.2982]], + [-0.3792, -0.1922, 0.0903], + [ + [-0.5586, -0.0762, -0.3944], + [-0.3306, -0.4191, -0.4898], + [0.1442, 0.0135, -0.3179], + ], + [-0.3912, -0.3963, -0.3368], + &device, + ); + + // Expected values from PyTorch + let expected_output_with_init = TensorData::from([[ + [0.24537, 0.14018, 0.19449, -0.49777, -0.15647, 0.48392], + [0.27468, -0.14514, 0.56205, -0.60381, -0.04986, 0.15683], + [-0.04062, -0.33486, 0.52330, -0.42244, -0.12644, -0.12034], + [-0.11743, -0.53873, 0.54429, -0.64943, 0.30127, -0.41943], + ]]); + + let expected_hn_with_init = TensorData::from([ + [[-0.11743, -0.53873, 0.54429]], + [[-0.49777, -0.15647, 0.48392]], + ]); + + let expected_output_without_init = TensorData::from([[ + [0.07452, -0.08247, 0.46677, -0.46770, -0.18086, 0.47519], + [0.15843, -0.27144, 0.65781, -0.50286, -0.12806, 0.14884], + [-0.10704, -0.41573, 0.53954, -0.24794, -0.24003, -0.10294], + [-0.16505, -0.57952, 0.53565, -0.23598, -0.07137, -0.28937], + ]]); + + let expected_hn_without_init = TensorData::from([ + [[-0.16505, -0.57952, 0.53565]], + [[-0.46770, -0.18086, 0.47519]], + ]); + + let (output_with_init, hn_with_init) = bigru.forward(input.clone(), Some(h0)); + let (output_without_init, hn_without_init) = bigru.forward(input, None); + + let tolerance = Tolerance::permissive(); + output_with_init + .to_data() + .assert_approx_eq::(&expected_output_with_init, tolerance); + output_without_init + .to_data() + .assert_approx_eq::(&expected_output_without_init, tolerance); + hn_with_init + .to_data() + .assert_approx_eq::(&expected_hn_with_init, tolerance); + hn_without_init + .to_data() + .assert_approx_eq::(&expected_hn_without_init, tolerance); + } + + #[test] + fn bigru_display() { + let config = BiGruConfig::new(2, 8, true); + + let layer = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{layer}"), + "BiGru {d_input: 2, d_hidden: 8, bias: true, params: 576}" + ); + } + + #[test] + fn test_gru_custom_activations() { + let device = Default::default(); + + // Create GRU with custom activations (ReLU instead of Sigmoid/Tanh) + let config = GruConfig::new(4, 8, true) + .with_gate_activation(ActivationConfig::Relu) + .with_hidden_activation(ActivationConfig::Relu); + let gru = config.init(&device); + + let input = Tensor::<3>::random([2, 3, 4], Distribution::Default, &device); + + // Should run without panicking and produce valid output + let output = gru.forward(input, None); + assert_eq!(&*output.shape(), [2, 3, 8]); + } + + #[test] + fn test_bigru_custom_activations() { + let device = Default::default(); + + // Create BiGRU with custom activations + let config = BiGruConfig::new(4, 8, true) + .with_gate_activation(ActivationConfig::Relu) + .with_hidden_activation(ActivationConfig::Relu); + let bigru = config.init(&device); + + let input = Tensor::<3>::random([2, 3, 4], Distribution::Default, &device); + + let (output, state) = bigru.forward(input, None); + assert_eq!(&*output.shape(), [2, 3, 16]); // hidden_size * 2 + assert_eq!(&*state.shape(), [2, 2, 8]); + } + + #[test] + fn test_gru_clipping() { + let device = Default::default(); + + // Create GRU with clipping enabled + let clip_value = 0.5; + let config = GruConfig::new(4, 8, true).with_clip(Some(clip_value)); + let gru = config.init(&device); + + let input = Tensor::<3>::random([2, 5, 4], Distribution::Default, &device); + + let output = gru.forward(input, None); + + // Verify output values are within the clip range + let output_data: Vec = output.to_data().to_vec().unwrap(); + for val in output_data { + assert!( + val >= -clip_value as f32 && val <= clip_value as f32, + "Value {} is outside clip range [-{}, {}]", + val, + clip_value, + clip_value + ); + } + } + + #[test] + fn test_bigru_clipping() { + let device = Default::default(); + + // Create BiGRU with clipping enabled + let clip_value = 0.3; + let config = BiGruConfig::new(4, 8, true).with_clip(Some(clip_value)); + let bigru = config.init(&device); + + let input = Tensor::<3>::random([2, 5, 4], Distribution::Default, &device); + + let (output, state) = bigru.forward(input, None); + + // Verify output values are within the clip range + let output_data: Vec = output.to_data().to_vec().unwrap(); + for val in output_data { + assert!( + val >= -clip_value as f32 && val <= clip_value as f32, + "Output value {} is outside clip range [-{}, {}]", + val, + clip_value, + clip_value + ); + } + + // Verify state values are within the clip range + let state_data: Vec = state.to_data().to_vec().unwrap(); + for val in state_data { + assert!( + val >= -clip_value as f32 && val <= clip_value as f32, + "State value {} is outside clip range [-{}, {}]", + val, + clip_value, + clip_value + ); + } + } + + /// Test Gru against PyTorch reference implementation. + /// Expected values computed with PyTorch nn.GRU (seed=42 for weights, seed=123 for input). + #[test] + fn test_gru_against_pytorch() { + let device = Device::default(); + device.seed(0); + + let config = GruConfig::new(2, 3, true); + let mut gru = config.init(&device); + + fn create_gate_controller( + input_weights: [[f32; D1]; D2], + input_biases: [f32; D1], + hidden_weights: [[f32; D1]; D1], + hidden_biases: [f32; D1], + device: &Device, + ) -> GateController { + let input_record = Linear { + weight: Param::from_data(TensorData::from(input_weights), device), + bias: Some(Param::from_data(TensorData::from(input_biases), device)), + }; + let hidden_record = Linear { + weight: Param::from_data(TensorData::from(hidden_weights), device), + bias: Some(Param::from_data(TensorData::from(hidden_biases), device)), + }; + GateController::create_with_weights(input_record, hidden_record) + } + + // Input: [batch=1, seq=4, input=2] + let input = Tensor::<3>::from_data( + TensorData::from([[ + [-0.11147, 0.12036], + [-0.36963, -0.24042], + [-1.19692, 0.20927], + [-0.97236, -0.75505], + ]]), + &device, + ); + + // Initial hidden state: [batch=1, hidden=3] + let h0 = Tensor::<2>::from_data(TensorData::from([[0.3239, -0.10852, 0.21033]]), &device); + + // Update gate (z) - weights from PyTorch, transposed for Burn's Row layout + gru.update_gate = create_gate_controller( + [[-0.2811, 0.5090, 0.5018], [0.3391, -0.4236, 0.1081]], + [0.2932, -0.3519, -0.5715], + [ + [-0.3471, 0.5214, 0.0961], + [0.0545, -0.4904, -0.1875], + [-0.5702, 0.4457, 0.3568], + ], + [-0.0100, 0.4518, -0.4102], + &device, + ); + + // Reset gate (r) + gru.reset_gate = create_gate_controller( + [[0.4414, -0.1353, -0.1265], [0.4792, 0.5304, 0.1165]], + [-0.2524, 0.3333, 0.1033], + [ + [-0.2695, -0.0677, -0.4557], + [0.1472, -0.2345, -0.2662], + [-0.2660, 0.3830, -0.1630], + ], + [0.1663, 0.2391, 0.1826], + &device, + ); + + // New gate (n) + gru.new_gate = create_gate_controller( + [[0.4266, 0.2784, 0.4451], [0.0782, -0.0815, 0.0853]], + [-0.2231, -0.4428, 0.4737], + [ + [0.0900, -0.1821, 0.2430], + [0.4665, 0.1551, 0.5155], + [0.0631, -0.1566, 0.3337], + ], + [0.0364, -0.3941, 0.1780], + &device, + ); + + // Expected values from PyTorch + let expected_output_with_h0 = TensorData::from([[ + [0.05665, -0.34932, 0.43267], + [-0.1737, -0.49246, 0.38099], + [-0.35401, -0.68099, 0.05061], + [-0.47854, -0.70427, -0.13648], + ]]); + + let expected_output_no_h0 = TensorData::from([[ + [-0.0985, -0.31661, 0.36126], + [-0.24563, -0.47784, 0.34609], + [-0.39497, -0.67659, 0.03083], + [-0.50146, -0.70066, -0.14894], + ]]); + + let output_with_h0 = gru.forward(input.clone(), Some(h0)); + let output_no_h0 = gru.forward(input, None); + + let tolerance = Tolerance::permissive(); + output_with_h0 + .to_data() + .assert_approx_eq::(&expected_output_with_h0, tolerance); + output_no_h0 + .to_data() + .assert_approx_eq::(&expected_output_no_h0, tolerance); + } +} diff --git a/crates/burn-nn/src/modules/rnn/lstm.rs b/crates/burn-nn/src/modules/rnn/lstm.rs new file mode 100644 index 0000000..16302db --- /dev/null +++ b/crates/burn-nn/src/modules/rnn/lstm.rs @@ -0,0 +1,856 @@ +use burn_core as burn; + +use crate::GateController; +use crate::activation::{Activation, ActivationConfig}; +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Initializer, Module, ModuleDisplay}; +use burn::tensor::Device; +use burn::tensor::Tensor; + +/// A LstmState is used to store cell state and hidden state in LSTM. +pub struct LstmState { + /// The cell state. + pub cell: Tensor, + /// The hidden state. + pub hidden: Tensor, +} + +impl LstmState { + /// Initialize a new [LSTM State](LstmState). + pub fn new(cell: Tensor, hidden: Tensor) -> Self { + Self { cell, hidden } + } +} + +/// Configuration to create a [Lstm](Lstm) module using the [init function](LstmConfig::init). +#[derive(Config, Debug)] +pub struct LstmConfig { + /// The size of the input features. + pub d_input: usize, + /// The size of the hidden state. + pub d_hidden: usize, + /// If a bias should be applied during the Lstm transformation. + pub bias: bool, + /// Lstm initializer + #[config(default = "Initializer::XavierNormal{gain:1.0}")] + pub initializer: Initializer, + /// If true, the input tensor is expected to be `[batch_size, seq_length, input_size]`. + /// If false, the input tensor is expected to be `[seq_length, batch_size, input_size]`. + #[config(default = true)] + pub batch_first: bool, + /// If true, process the sequence in reverse order. + /// This is useful for implementing reverse-direction LSTMs (e.g., ONNX reverse direction). + #[config(default = false)] + pub reverse: bool, + /// Optional cell state clip threshold. If provided, cell state values are clipped + /// to the range `[-clip, +clip]` after each timestep. This can help prevent + /// exploding values during inference. + pub clip: Option, + /// If true, couples the input and forget gates: `f_t = 1 - i_t`. + /// This reduces the number of parameters and is based on GRU-style simplification. + #[config(default = false)] + pub input_forget: bool, + /// Activation function for the input, forget, and output gates. + /// Default is Sigmoid, which is standard for LSTM gates. + #[config(default = "ActivationConfig::Sigmoid")] + pub gate_activation: ActivationConfig, + /// Activation function for the cell gate (candidate cell state). + /// Default is Tanh, which is standard for LSTM. + #[config(default = "ActivationConfig::Tanh")] + pub cell_activation: ActivationConfig, + /// Activation function applied to the cell state before computing hidden output. + /// Default is Tanh, which is standard for LSTM. + #[config(default = "ActivationConfig::Tanh")] + pub hidden_activation: ActivationConfig, +} + +/// The Lstm module. This implementation is for a unidirectional, stateless, Lstm. +/// +/// Introduced in the paper: [Long Short-Term Memory](https://www.researchgate.net/publication/13853244). +/// +/// Should be created with [LstmConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Lstm { + /// The input gate regulates which information to update and store in the cell state at each time step. + pub input_gate: GateController, + /// The forget gate is used to control which information to discard or keep in the memory cell at each time step. + /// Note: When `input_forget` is true, this gate is not used (forget = 1 - input). + pub forget_gate: GateController, + /// The output gate determines which information from the cell state to output at each time step. + pub output_gate: GateController, + /// The cell gate is used to compute the cell state that stores and carries information through time. + pub cell_gate: GateController, + /// The hidden state of the LSTM. + pub d_hidden: usize, + /// If true, input is `[batch_size, seq_length, input_size]`. + /// If false, input is `[seq_length, batch_size, input_size]`. + pub batch_first: bool, + /// If true, process the sequence in reverse order. + pub reverse: bool, + /// Optional cell state clip threshold. + pub clip: Option, + /// If true, couples input and forget gates: f_t = 1 - i_t. + pub input_forget: bool, + /// Activation function for gates (input, forget, output). + pub gate_activation: Activation, + /// Activation function for cell gate (candidate cell state). + pub cell_activation: Activation, + /// Activation function for hidden output. + pub hidden_activation: Activation, +} + +impl ModuleDisplay for Lstm { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let [d_input, _] = self.input_gate.input_transform.weight.shape().dims(); + let bias = self.input_gate.input_transform.bias.is_some(); + + content + .add("d_input", &d_input) + .add("d_hidden", &self.d_hidden) + .add("bias", &bias) + .optional() + } +} + +impl LstmConfig { + /// Initialize a new [lstm](Lstm) module. + pub fn init(&self, device: &Device) -> Lstm { + let d_output = self.d_hidden; + + let new_gate = || { + GateController::new( + self.d_input, + d_output, + self.bias, + self.initializer.clone(), + device, + ) + }; + + Lstm { + input_gate: new_gate(), + forget_gate: new_gate(), + output_gate: new_gate(), + cell_gate: new_gate(), + d_hidden: self.d_hidden, + batch_first: self.batch_first, + reverse: self.reverse, + clip: self.clip, + input_forget: self.input_forget, + gate_activation: self.gate_activation.init(device), + cell_activation: self.cell_activation.init(device), + hidden_activation: self.hidden_activation.init(device), + } + } +} + +impl Lstm { + /// Applies the forward pass on the input tensor. This LSTM implementation + /// returns the state for each element in a sequence (i.e., across seq_length) and a final state. + /// + /// ## Parameters: + /// - batched_input: The input tensor of shape: + /// - `[batch_size, sequence_length, input_size]` if `batch_first` is true (default) + /// - `[sequence_length, batch_size, input_size]` if `batch_first` is false + /// - state: An optional `LstmState` representing the initial cell state and hidden state. + /// Each state tensor has shape `[batch_size, hidden_size]`. + /// If no initial state is provided, these tensors are initialized to zeros. + /// + /// ## Returns: + /// - output: A tensor represents the output features of LSTM. Shape: + /// - `[batch_size, sequence_length, hidden_size]` if `batch_first` is true + /// - `[sequence_length, batch_size, hidden_size]` if `batch_first` is false + /// - state: A `LstmState` represents the final states. Both `state.cell` and `state.hidden` have the shape + /// `[batch_size, hidden_size]`. + pub fn forward( + &self, + batched_input: Tensor<3>, + state: Option>, + ) -> (Tensor<3>, LstmState<2>) { + // Convert to batch-first layout internally if needed + let batched_input = if self.batch_first { + batched_input + } else { + batched_input.swap_dims(0, 1) + }; + + let device = batched_input.device(); + let [batch_size, seq_length, _] = batched_input.dims(); + + // Process sequence in forward or reverse order based on config + let (output, state) = if self.reverse { + self.forward_iter( + batched_input.iter_dim(1).rev().zip((0..seq_length).rev()), + state, + batch_size, + seq_length, + &device, + ) + } else { + self.forward_iter( + batched_input.iter_dim(1).zip(0..seq_length), + state, + batch_size, + seq_length, + &device, + ) + }; + + // Convert output back to seq-first layout if needed + let output = if self.batch_first { + output + } else { + output.swap_dims(0, 1) + }; + + (output, state) + } + + fn forward_iter, usize)>>( + &self, + input_timestep_iter: I, + state: Option>, + batch_size: usize, + seq_length: usize, + device: &Device, + ) -> (Tensor<3>, LstmState<2>) { + let mut batched_hidden_state = + Tensor::empty([batch_size, seq_length, self.d_hidden], device); + + let (mut cell_state, mut hidden_state) = match state { + Some(state) => (state.cell, state.hidden), + None => ( + Tensor::zeros([batch_size, self.d_hidden], device), + Tensor::zeros([batch_size, self.d_hidden], device), + ), + }; + + for (input_t, t) in input_timestep_iter { + let input_t = input_t.squeeze_dim(1); + + // i(nput)g(ate) tensors + let biased_ig_input_sum = self + .input_gate + .gate_product(input_t.clone(), hidden_state.clone()); + let input_values = self.gate_activation.forward(biased_ig_input_sum); + + // f(orget)g(ate) tensors - either computed or coupled to input gate + let forget_values = if self.input_forget { + // Coupled mode: f_t = 1 - i_t + input_values.clone().neg().add_scalar(1.0) + } else { + let biased_fg_input_sum = self + .forget_gate + .gate_product(input_t.clone(), hidden_state.clone()); + self.gate_activation.forward(biased_fg_input_sum) + }; + + // o(output)g(ate) tensors + let biased_og_input_sum = self + .output_gate + .gate_product(input_t.clone(), hidden_state.clone()); + let output_values = self.gate_activation.forward(biased_og_input_sum); + + // c(ell)g(ate) tensors + let biased_cg_input_sum = self + .cell_gate + .gate_product(input_t.clone(), hidden_state.clone()); + let candidate_cell_values = self.cell_activation.forward(biased_cg_input_sum); + + cell_state = forget_values * cell_state.clone() + input_values * candidate_cell_values; + + // Apply cell state clipping if configured + if let Some(clip) = self.clip { + cell_state = cell_state.clamp(-clip, clip); + } + + hidden_state = output_values * self.hidden_activation.forward(cell_state.clone()); + + let unsqueezed_hidden_state = hidden_state.clone().unsqueeze_dim(1); + + // store the hidden state for this timestep + batched_hidden_state = batched_hidden_state.slice_assign( + [0..batch_size, t..(t + 1), 0..self.d_hidden], + unsqueezed_hidden_state.clone(), + ); + } + + ( + batched_hidden_state, + LstmState::new(cell_state, hidden_state), + ) + } +} + +/// Configuration to create a [BiLstm](BiLstm) module using the [init function](BiLstmConfig::init). +#[derive(Config, Debug)] +pub struct BiLstmConfig { + /// The size of the input features. + pub d_input: usize, + /// The size of the hidden state. + pub d_hidden: usize, + /// If a bias should be applied during the BiLstm transformation. + pub bias: bool, + /// BiLstm initializer + #[config(default = "Initializer::XavierNormal{gain:1.0}")] + pub initializer: Initializer, + /// If true, the input tensor is expected to be `[batch_size, seq_length, input_size]`. + /// If false, the input tensor is expected to be `[seq_length, batch_size, input_size]`. + #[config(default = true)] + pub batch_first: bool, + /// Optional cell state clip threshold. + pub clip: Option, + /// If true, couples the input and forget gates. + #[config(default = false)] + pub input_forget: bool, + /// Activation function for the input, forget, and output gates. + #[config(default = "ActivationConfig::Sigmoid")] + pub gate_activation: ActivationConfig, + /// Activation function for the cell gate (candidate cell state). + #[config(default = "ActivationConfig::Tanh")] + pub cell_activation: ActivationConfig, + /// Activation function applied to the cell state before computing hidden output. + #[config(default = "ActivationConfig::Tanh")] + pub hidden_activation: ActivationConfig, +} + +/// The BiLstm module. This implementation is for Bidirectional LSTM. +/// +/// Introduced in the paper: [Framewise phoneme classification with bidirectional LSTM and other neural network architectures](https://www.cs.toronto.edu/~graves/ijcnn_2005.pdf). +/// +/// Should be created with [BiLstmConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct BiLstm { + /// LSTM for the forward direction. + pub forward: Lstm, + /// LSTM for the reverse direction. + pub reverse: Lstm, + /// The size of the hidden state. + pub d_hidden: usize, + /// If true, input is `[batch_size, seq_length, input_size]`. + /// If false, input is `[seq_length, batch_size, input_size]`. + pub batch_first: bool, +} + +impl ModuleDisplay for BiLstm { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let [d_input, _] = self + .forward + .input_gate + .input_transform + .weight + .shape() + .dims(); + let bias = self.forward.input_gate.input_transform.bias.is_some(); + + content + .add("d_input", &d_input) + .add("d_hidden", &self.d_hidden) + .add("bias", &bias) + .optional() + } +} + +impl BiLstmConfig { + /// Initialize a new [Bidirectional LSTM](BiLstm) module. + pub fn init(&self, device: &Device) -> BiLstm { + // Internal LSTMs always use batch_first=true; BiLstm handles layout conversion + let base_config = LstmConfig::new(self.d_input, self.d_hidden, self.bias) + .with_initializer(self.initializer.clone()) + .with_batch_first(true) + .with_clip(self.clip) + .with_input_forget(self.input_forget) + .with_gate_activation(self.gate_activation.clone()) + .with_cell_activation(self.cell_activation.clone()) + .with_hidden_activation(self.hidden_activation.clone()); + + BiLstm { + forward: base_config.clone().init(device), + reverse: base_config.init(device), + d_hidden: self.d_hidden, + batch_first: self.batch_first, + } + } +} + +impl BiLstm { + /// Applies the forward pass on the input tensor. This Bidirectional LSTM implementation + /// returns the state for each element in a sequence (i.e., across seq_length) and a final state. + /// + /// ## Parameters: + /// - batched_input: The input tensor of shape: + /// - `[batch_size, sequence_length, input_size]` if `batch_first` is true (default) + /// - `[sequence_length, batch_size, input_size]` if `batch_first` is false + /// - state: An optional `LstmState` representing the initial cell state and hidden state. + /// Each state tensor has shape `[2, batch_size, hidden_size]`. + /// If no initial state is provided, these tensors are initialized to zeros. + /// + /// ## Returns: + /// - output: A tensor represents the output features of LSTM. Shape: + /// - `[batch_size, sequence_length, hidden_size * 2]` if `batch_first` is true + /// - `[sequence_length, batch_size, hidden_size * 2]` if `batch_first` is false + /// - state: A `LstmState` represents the final forward and reverse states. Both `state.cell` and + /// `state.hidden` have the shape `[2, batch_size, hidden_size]`. + pub fn forward( + &self, + batched_input: Tensor<3>, + state: Option>, + ) -> (Tensor<3>, LstmState<3>) { + // Convert to batch-first layout internally if needed + let batched_input = if self.batch_first { + batched_input + } else { + batched_input.swap_dims(0, 1) + }; + + let device = batched_input.clone().device(); + let [batch_size, seq_length, _] = batched_input.shape().dims(); + + let [init_state_forward, init_state_reverse] = match state { + Some(state) => { + let cell_state_forward = state + .cell + .clone() + .slice([0..1, 0..batch_size, 0..self.d_hidden]) + .squeeze_dim(0); + let hidden_state_forward = state + .hidden + .clone() + .slice([0..1, 0..batch_size, 0..self.d_hidden]) + .squeeze_dim(0); + let cell_state_reverse = state + .cell + .slice([1..2, 0..batch_size, 0..self.d_hidden]) + .squeeze_dim(0); + let hidden_state_reverse = state + .hidden + .slice([1..2, 0..batch_size, 0..self.d_hidden]) + .squeeze_dim(0); + + [ + Some(LstmState::new(cell_state_forward, hidden_state_forward)), + Some(LstmState::new(cell_state_reverse, hidden_state_reverse)), + ] + } + None => [None, None], + }; + + // forward direction + let (batched_hidden_state_forward, final_state_forward) = self + .forward + .forward(batched_input.clone(), init_state_forward); + + // reverse direction + let (batched_hidden_state_reverse, final_state_reverse) = self.reverse.forward_iter( + batched_input.iter_dim(1).rev().zip((0..seq_length).rev()), + init_state_reverse, + batch_size, + seq_length, + &device, + ); + + let output = Tensor::cat( + [batched_hidden_state_forward, batched_hidden_state_reverse].to_vec(), + 2, + ); + + // Convert output back to seq-first layout if needed + let output = if self.batch_first { + output + } else { + output.swap_dims(0, 1) + }; + + let state = LstmState::new( + Tensor::stack( + [final_state_forward.cell, final_state_reverse.cell].to_vec(), + 0, + ), + Tensor::stack( + [final_state_forward.hidden, final_state_reverse.hidden].to_vec(), + 0, + ), + ); + + (output, state) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Linear; + use burn::module::Param; + use burn::tensor::{Distribution, TensorData}; + use burn::tensor::{ElementConversion, Tolerance}; + type FT = f32; + + #[test] + fn test_with_uniform_initializer() { + let device = Device::default(); + device.seed(0); + + let config = LstmConfig::new(5, 5, false) + .with_initializer(Initializer::Uniform { min: 0.0, max: 1.0 }); + let lstm = config.init(&Default::default()); + + let gate_to_data = |gate: GateController| gate.input_transform.weight.val().to_data(); + + gate_to_data(lstm.input_gate).assert_within_range::(0.elem()..1.elem()); + gate_to_data(lstm.forget_gate).assert_within_range::(0.elem()..1.elem()); + gate_to_data(lstm.output_gate).assert_within_range::(0.elem()..1.elem()); + gate_to_data(lstm.cell_gate).assert_within_range::(0.elem()..1.elem()); + } + + /// Test forward pass with simple input vector. + /// + /// f_t = sigmoid(0.7*0.1 + 0.7*0) = sigmoid(0.07) = 0.5173928 + /// i_t = sigmoid(0.5*0.1 + 0.5*0) = sigmoid(0.05) = 0.5123725 + /// o_t = sigmoid(1.1*0.1 + 1.1*0) = sigmoid(0.11) = 0.5274723 + /// c_t = tanh(0.9*0.1 + 0.9*0) = tanh(0.09) = 0.0892937 + /// C_t = f_t * 0 + i_t * c_t = 0 + 0.5123725 * 0.0892937 = 0.04575243 + /// h_t = o_t * tanh(C_t) = 0.5274723 * tanh(0.04575243) = 0.5274723 * 0.04568173 = 0.024083648 + #[test] + fn test_forward_single_input_single_feature() { + let device = Device::default(); + device.seed(0); + + let config = LstmConfig::new(1, 1, false); + let mut lstm = config.init(&device); + + fn create_gate_controller(weights: f32, biases: f32, device: &Device) -> GateController { + let record_1 = Linear { + weight: Param::from_data(TensorData::from([[weights]]), device), + bias: Some(Param::from_data(TensorData::from([biases]), device)), + }; + let record_2 = Linear { + weight: Param::from_data(TensorData::from([[weights]]), device), + bias: Some(Param::from_data(TensorData::from([biases]), device)), + }; + GateController::create_with_weights(record_1, record_2) + } + + lstm.input_gate = create_gate_controller(0.5, 0.0, &device); + lstm.forget_gate = create_gate_controller(0.7, 0.0, &device); + lstm.cell_gate = create_gate_controller(0.9, 0.0, &device); + lstm.output_gate = create_gate_controller(1.1, 0.0, &device); + + // single timestep with single feature + let input = Tensor::<3>::from_data(TensorData::from([[[0.1]]]), &device); + + let (output, state) = lstm.forward(input, None); + + let expected = TensorData::from([[0.046]]); + let tolerance = Tolerance::default(); + state + .cell + .to_data() + .assert_approx_eq::(&expected, tolerance); + + let expected = TensorData::from([[0.0242]]); + state + .hidden + .to_data() + .assert_approx_eq::(&expected, tolerance); + + output + .select(0, Tensor::arange(0..1, &device)) + .squeeze_dim::<2>(0) + .to_data() + .assert_approx_eq::(&state.hidden.to_data(), tolerance); + } + + #[test] + fn test_batched_forward_pass() { + let device = Default::default(); + let lstm = LstmConfig::new(64, 1024, true).init(&device); + let batched_input = Tensor::<3>::random([8, 10, 64], Distribution::Default, &device); + + let (output, state) = lstm.forward(batched_input, None); + + assert_eq!(output.dims(), [8, 10, 1024]); + assert_eq!(state.cell.dims(), [8, 1024]); + assert_eq!(state.hidden.dims(), [8, 1024]); + } + + #[test] + fn test_batched_forward_pass_batch_of_one() { + let device = Default::default(); + let lstm = LstmConfig::new(64, 1024, true).init(&device); + let batched_input = Tensor::<3>::random([1, 2, 64], Distribution::Default, &device); + + let (output, state) = lstm.forward(batched_input, None); + + assert_eq!(output.dims(), [1, 2, 1024]); + assert_eq!(state.cell.dims(), [1, 1024]); + assert_eq!(state.hidden.dims(), [1, 1024]); + } + + #[test] + #[cfg(feature = "std")] + fn test_batched_backward_pass() { + use burn::tensor::Shape; + let device = Device::default().autodiff(); + let lstm = LstmConfig::new(64, 32, true).init(&device); + let shape: Shape = [8, 10, 64].into(); + let batched_input = Tensor::<3>::random(shape, Distribution::Default, &device); + + let (output, _) = lstm.forward(batched_input.clone(), None); + let fake_loss = output; + let grads = fake_loss.backward(); + + let some_gradient = lstm + .output_gate + .hidden_transform + .weight + .grad(&grads) + .unwrap(); + + // Asserts that the gradients exist and are non-zero + assert_ne!( + some_gradient + .any() + .into_data() + .iter::() + .next() + .unwrap(), + 0.0 + ); + } + + #[test] + fn test_bidirectional() { + let device = Device::default(); + device.seed(0); + + let config = BiLstmConfig::new(2, 3, true); + let mut lstm = config.init(&device); + + fn create_gate_controller( + input_weights: [[f32; D1]; D2], + input_biases: [f32; D1], + hidden_weights: [[f32; D1]; D1], + hidden_biases: [f32; D1], + device: &Device, + ) -> GateController { + let input_record = Linear { + weight: Param::from_data(TensorData::from(input_weights), device), + bias: Some(Param::from_data(TensorData::from(input_biases), device)), + }; + let hidden_record = Linear { + weight: Param::from_data(TensorData::from(hidden_weights), device), + bias: Some(Param::from_data(TensorData::from(hidden_biases), device)), + }; + GateController::create_with_weights(input_record, hidden_record) + } + + let input = Tensor::<3>::from_data( + TensorData::from([[ + [0.949, -0.861], + [0.892, 0.927], + [-0.173, -0.301], + [-0.081, 0.992], + ]]), + &device, + ); + let h0 = Tensor::<3>::from_data( + TensorData::from([[[0.280, 0.360, -1.242]], [[-0.588, 0.729, -0.788]]]), + &device, + ); + let c0 = Tensor::<3>::from_data( + TensorData::from([[[0.723, 0.397, -0.262]], [[0.471, 0.613, 1.885]]]), + &device, + ); + + lstm.forward.input_gate = create_gate_controller( + [[0.367, 0.091, 0.342], [0.322, 0.533, 0.059]], + [-0.196, 0.354, 0.209], + [ + [-0.320, 0.232, -0.165], + [0.093, -0.572, -0.315], + [-0.467, 0.325, 0.046], + ], + [0.181, -0.190, -0.245], + &device, + ); + + lstm.forward.forget_gate = create_gate_controller( + [[-0.342, -0.084, -0.420], [-0.432, 0.119, 0.191]], + [0.315, -0.413, -0.041], + [ + [0.453, 0.063, 0.561], + [0.211, 0.149, 0.213], + [-0.499, -0.158, 0.068], + ], + [-0.431, -0.535, 0.125], + &device, + ); + + lstm.forward.cell_gate = create_gate_controller( + [[-0.046, -0.382, 0.321], [-0.533, 0.558, 0.004]], + [-0.358, 0.282, -0.078], + [ + [-0.358, 0.109, 0.139], + [-0.345, 0.091, -0.368], + [-0.508, 0.221, -0.507], + ], + [0.502, -0.509, -0.247], + &device, + ); + + lstm.forward.output_gate = create_gate_controller( + [[-0.577, -0.359, 0.216], [-0.550, 0.268, 0.243]], + [-0.227, -0.274, 0.039], + [ + [-0.383, 0.449, 0.222], + [-0.357, -0.093, 0.449], + [-0.106, 0.236, 0.360], + ], + [-0.361, -0.209, -0.454], + &device, + ); + + lstm.reverse.input_gate = create_gate_controller( + [[-0.055, 0.506, 0.247], [-0.369, 0.178, -0.258]], + [0.540, -0.164, 0.033], + [ + [0.159, 0.180, -0.037], + [-0.443, 0.485, -0.488], + [0.098, -0.085, -0.140], + ], + [-0.510, 0.105, 0.114], + &device, + ); + + lstm.reverse.forget_gate = create_gate_controller( + [[-0.154, -0.432, -0.547], [-0.369, -0.310, -0.175]], + [0.141, 0.004, 0.055], + [ + [-0.005, -0.277, -0.515], + [-0.011, -0.101, -0.365], + [0.426, 0.379, 0.337], + ], + [-0.382, 0.331, -0.176], + &device, + ); + + lstm.reverse.cell_gate = create_gate_controller( + [[-0.571, 0.228, -0.287], [-0.331, 0.110, 0.219]], + [-0.206, -0.546, 0.462], + [ + [0.449, -0.240, 0.071], + [-0.045, 0.131, 0.124], + [0.138, -0.201, 0.191], + ], + [-0.030, 0.211, -0.352], + &device, + ); + + lstm.reverse.output_gate = create_gate_controller( + [[0.491, -0.442, 0.333], [0.313, -0.121, -0.070]], + [-0.387, -0.250, 0.066], + [ + [-0.030, 0.268, 0.299], + [-0.019, -0.280, -0.314], + [0.466, -0.365, -0.248], + ], + [-0.398, -0.199, -0.566], + &device, + ); + + let expected_output_with_init_state = TensorData::from([[ + [0.23764, -0.03442, 0.04414, -0.15635, -0.03366, -0.05798], + [0.00473, -0.02254, 0.02988, -0.16510, -0.00306, 0.08742], + [0.06210, -0.06509, -0.05339, -0.01710, 0.02091, 0.16012], + [-0.03420, 0.07774, -0.09774, -0.02604, 0.12584, 0.20872], + ]]); + let expected_output_without_init_state = TensorData::from([[ + [0.08679, -0.08776, -0.00528, -0.15969, -0.05322, -0.08863], + [-0.02577, -0.05057, 0.00033, -0.17558, -0.03679, 0.03142], + [0.02942, -0.07411, -0.06044, -0.03601, -0.09998, 0.04846], + [-0.04026, 0.07178, -0.10189, -0.07349, -0.04576, 0.05550], + ]]); + let expected_hn_with_init_state = TensorData::from([ + [[-0.03420, 0.07774, -0.09774]], + [[-0.15635, -0.03366, -0.05798]], + ]); + let expected_cn_with_init_state = TensorData::from([ + [[-0.13593, 0.17125, -0.22395]], + [[-0.45425, -0.11206, -0.12908]], + ]); + let expected_hn_without_init_state = TensorData::from([ + [[-0.04026, 0.07178, -0.10189]], + [[-0.15969, -0.05322, -0.08863]], + ]); + let expected_cn_without_init_state = TensorData::from([ + [[-0.15839, 0.15923, -0.23569]], + [[-0.47407, -0.17493, -0.19643]], + ]); + + let (output_with_init_state, state_with_init_state) = + lstm.forward(input.clone(), Some(LstmState::new(c0, h0))); + let (output_without_init_state, state_without_init_state) = lstm.forward(input, None); + + let tolerance = Tolerance::permissive(); + output_with_init_state + .to_data() + .assert_approx_eq::(&expected_output_with_init_state, tolerance); + output_without_init_state + .to_data() + .assert_approx_eq::(&expected_output_without_init_state, tolerance); + state_with_init_state + .hidden + .to_data() + .assert_approx_eq::(&expected_hn_with_init_state, tolerance); + state_with_init_state + .cell + .to_data() + .assert_approx_eq::(&expected_cn_with_init_state, tolerance); + state_without_init_state + .hidden + .to_data() + .assert_approx_eq::(&expected_hn_without_init_state, tolerance); + state_without_init_state + .cell + .to_data() + .assert_approx_eq::(&expected_cn_without_init_state, tolerance); + } + + #[test] + fn display_lstm() { + let config = LstmConfig::new(2, 3, true); + + let layer = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{layer}"), + "Lstm {d_input: 2, d_hidden: 3, bias: true, params: 84}" + ); + } + + #[test] + fn display_bilstm() { + let config = BiLstmConfig::new(2, 3, true); + + let layer = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{layer}"), + "BiLstm {d_input: 2, d_hidden: 3, bias: true, params: 168}" + ); + } +} diff --git a/crates/burn-nn/src/modules/rnn/mod.rs b/crates/burn-nn/src/modules/rnn/mod.rs new file mode 100644 index 0000000..b651874 --- /dev/null +++ b/crates/burn-nn/src/modules/rnn/mod.rs @@ -0,0 +1,15 @@ +mod gate_controller; + +/// Basic RNN. +pub mod basic; + +/// Gated Recurrent Unit module. +pub mod gru; + +/// Long Short-Term Memory module. +pub mod lstm; + +pub use basic::*; +pub use gate_controller::*; +pub use gru::*; +pub use lstm::*; diff --git a/crates/burn-nn/src/modules/rope_encoding.rs b/crates/burn-nn/src/modules/rope_encoding.rs new file mode 100644 index 0000000..48784a8 --- /dev/null +++ b/crates/burn-nn/src/modules/rope_encoding.rs @@ -0,0 +1,627 @@ +use burn_core as burn; + +use alloc::vec; +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Module, ModuleDisplay}; +use burn::tensor::Int; +use burn::tensor::{Device, Tensor}; +use core::ops::Range; + +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float as _; + +/// Configuration to create a [RotaryEncoding](RotaryEncoding) layer using the [init function](RotaryEncodingConfig::init). +#[derive(Config, Debug)] +pub struct RotaryEncodingConfig { + /// Maximum sequence length of input + pub max_sequence_length: usize, + + /// Size of the input embedding or hidden dimension + pub d_model: usize, + + /// Scaling factor for frequency computation. Defaults to 10000.0 + #[config(default = "10000.0")] + pub theta: f32, +} + +impl RotaryEncodingConfig { + /// Initialize a new [RotaryEncoding](RotaryEncoding) module. + /// + /// # Panics + /// + /// Panics if the size of input embedding dimension is not even. + /// Panics if the theta parameter is not positive. + pub fn init(&self, device: &Device) -> RotaryEncoding { + self.initialize(|x| x, device) + } + + /// Initialize a new [RotaryEncoding](RotaryEncoding) module with a custom frequency scaling function. + /// This is useful to apply different RoPE extensions. + /// + /// # Panics + /// + /// Panics if the size of input embedding dimension is not even. + /// Panics if the theta parameter is not positive. + pub fn init_with_frequency_scaling( + &self, + scaling: impl Fn(Tensor<1>) -> Tensor<1>, + device: &Device, + ) -> RotaryEncoding { + self.initialize(scaling, device) + } + + /// Initialize a new [RotaryEncoding](RotaryEncoding) module. + /// + /// # Panics + /// + /// Panics if the size of input embedding dimension is not even. + /// Panics if the theta parameter is not positive. + fn initialize( + &self, + scaling: impl Fn(Tensor<1>) -> Tensor<1>, + device: &Device, + ) -> RotaryEncoding { + assert_eq!( + self.d_model % 2, + 0, + "The input embedding dimension must be even" + ); + assert!( + self.theta > 0.0, + "Theta parameter must be positive (default: 10000)." + ); + + // Calculate the rotation frequencies for positional embeddings based on the formula + // `theta = 1 / (theta ^ (2i / d_model)) for i in [0..d_model/2]` + let exponent = Tensor::<1, Int>::arange_step(0..self.d_model as i64, 2, device) + .float() + .div_scalar(self.d_model as f32); + + // Calculate (10000 ^ (2i / d_model)) by using the log base property `exp(log(10000) * (2i / d_model))` + // This is done since burn doesn't support exponentiation of scalar to tensor + let theta = exponent.mul_scalar(self.theta.ln()).exp().recip(); + + let theta = scaling(theta); + + let freq_complex = + RotaryEncoding::compute_rotary_frequencies(0..self.max_sequence_length, theta.clone()); + + RotaryEncoding { + freq_complex, + theta, + start_offset: 0, + } + } +} + +/// A module that applies rotary positional encoding to a tensor. +/// Rotary Position Encoding or Embedding (RoPE), is a type of position embedding which encodes +/// absolute positional information with rotation matrix and naturally incorporates +/// explicit relative position dependency in self-attention formulation. +/// +/// Introduced in the paper: [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://arxiv.org/abs/2104.09864) +/// +/// Should be created using [RotaryEncodingConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct RotaryEncoding { + /// Complex frequency tensor of shape (max_sequence_length, d_model, 2) with real and imaginary components + // Essentially a cache of pre-computed RoPE values. + pub freq_complex: Tensor<3>, + /// Frequency vector used to compute/apply the complex rotations. + pub theta: Tensor<1>, + start_offset: usize, +} + +impl ModuleDisplay for RotaryEncoding { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let [max_sequence_length, d_model, _] = self.freq_complex.shape().dims(); + content + .add("d_model", &d_model) + .add("max_sequence_length", &max_sequence_length) + .optional() + } +} + +#[allow(clippy::single_range_in_vec_init)] +impl RotaryEncoding { + /// Applies rotary positional encoding to a tensor of dimensions (..., seq_len, d_model) + /// + /// # Arguments: + /// * `x` - Input tensor of shape (..., seq_len, d_model). Accommodate both 3D and 4D tensors + /// for (batch size, seq_len, hidden_dim) or (batch size, num_heads, seq_len, hidden_dim) + /// respectively. + /// + /// # Returns: + /// Output tensor with the same shape as input tensor after applying rotary encoding. + /// + /// # Panics + /// If the input tensor does not have at least 2 dimensions for sequence length and hidden dimension. + pub fn forward(&self, x: Tensor) -> Tensor { + self.apply(x, 0) + } + + /// Applies rotary positional encoding to a tensor of dimensions (..., seq_len, d_model) + /// + /// # Arguments: + /// * `x` - Input tensor of shape (..., seq_len, d_model). Accommodate both 3D and 4D tensors + /// for (batch size, seq_len, hidden_dim) or (batch size, num_heads, seq_len, hidden_dim) + /// respectively. + /// * `start` - Sequence start position index. + /// + /// # Returns: + /// Output tensor with the same shape as input tensor after applying rotary encoding. + /// + /// # Panics + /// If the input tensor does not have at least 2 dimensions for sequence length and hidden dimension. + pub fn apply(&self, x: Tensor, start: usize) -> Tensor { + assert!( + D >= 2, + "Input tensor must have at least 2 dimensions for sequence length and hidden dimension" + ); + + let device = x.device(); + let input_shape = x.shape(); + + // Extract the sequence length and embedding dimension, other dimensions are kept generic + // to allow both 3D and 4D tensors i.e. batch_size or (batch_size, num_heads) + let (seq_len, d_model) = (x.dims()[D - 2], x.dims()[D - 1]); + let dummy_dim_size = input_shape.num_elements() / (seq_len * d_model); + + // Create a dummy tensor with signed ones based on the 2D rotation matrix + // [[cos, -sin], [sin, cos]] + let sign_tensor = + Tensor::<2>::from_floats([[1.0, 0.0, 0.0, 1.0], [0.0, -1.0, 1.0, 0.0]], &device); + + // Rotate input using the frequency tensor. Slice the frequencies till input sequence length + let out: Tensor<4> = x + .reshape([dummy_dim_size, seq_len, d_model / 2, 2]) + .matmul(sign_tensor.unsqueeze()) + .reshape([dummy_dim_size, seq_len, d_model, 2]) + * self + .freq_complex + .clone() + .slice([start..start + seq_len]) + .unsqueeze(); + + // Sum the real and imaginary components to get output tensor and reshape to original shape + out.sum_dim(-1).reshape(input_shape) + } + + /// Shifts the pre-computed rotary frequency to cover a new range of positions. + /// + /// This method updates the internal frequency tensor `freq_complex` to store + /// the rotary positional encodings for a new window of positions starting at `start`. + pub fn shift(&mut self, start: usize) { + let max_seq_len = self.freq_complex.dims()[0]; + assert!( + start > self.start_offset, + "Shift start position must be monotonically increasing" + ); + + let current_end = self.start_offset + max_seq_len; + + if start >= current_end { + // Overwrite the whole buffer + let new_freqs = + Self::compute_rotary_frequencies(start..start + max_seq_len, self.theta.clone()); + self.freq_complex + .inplace(|freqs| freqs.slice_assign([0..max_seq_len], new_freqs)); + } else { + // Shift the tail + let num_keep = current_end - start; + let start_rel = start - self.start_offset; + let tail_freqs = self.freq_complex.clone().slice([start_rel..max_seq_len]); + self.freq_complex + .inplace(|freqs| freqs.slice_assign([0..num_keep], tail_freqs)); + // Compute the rest and assign + let new_freqs = Self::compute_rotary_frequencies( + current_end..start + max_seq_len, + self.theta.clone(), + ); + self.freq_complex + .inplace(|freqs| freqs.slice_assign([num_keep..max_seq_len], new_freqs)); + } + self.start_offset = start; + } + + /// Resets the pre-computed rotary frequency window to start at position 0. + pub fn reset(&mut self) { + if self.start_offset == 0 { + return; + } + + let max_seq_len = self.freq_complex.dims()[0]; + let freqs = Self::compute_rotary_frequencies(0..max_seq_len, self.theta.clone()); + self.freq_complex + .inplace(|freq_complex| freq_complex.slice_assign([0..max_seq_len], freqs)); + self.start_offset = 0; + } + + /// Computes the positional rotation frequencies (cosine and sine values) used in RoPE. + /// + /// # Arguments + /// - `range`: Range of position indices `[start, end)`. + /// - `theta`: 1D tensor of shape `(d_model / 2)` containing base angular frequencies. + /// + /// # Returns + /// Tensor of shape `(range.len(), d_model, 2)` containing `[cos, sin]` pairs for each position and frequency. + fn compute_rotary_frequencies(range: Range, theta: Tensor<1>) -> Tensor<3> { + let d_model = theta.dims()[0] * 2; + let num_positions = range.end - range.start; + + // Generate frequency values for positional embeddings + let frequencies: Tensor<2> = + Tensor::<1, Int>::arange(range.start as i64..range.end as i64, &theta.device()) + .float() + .unsqueeze() + .transpose() + .repeat_dim(1, d_model / 2) + * theta.unsqueeze(); + + // Convert frequency values to complex numbers (polar form) + let p_cos = frequencies.clone().cos(); + let p_sin = frequencies.sin(); + + Tensor::cat(vec![p_cos, p_sin], 1) + .reshape([num_positions, 2, d_model / 2]) + .transpose() + .unsqueeze_dim::<4>(2) + .repeat_dim(2, 2) + .reshape([num_positions, d_model, 2]) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn test_rotary_encoding_forward() { + let device = Default::default(); + let rotary_encoding = RotaryEncodingConfig::new(10, 4).init(&device); + + let input = Tensor::<3>::from_floats( + [ + [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], + [[9.0, 10.0, 11.0, 12.0], [13.0, 14.0, 15.0, 16.0]], + ], + &device, + ); + + // Input = [Batch size, Num of heads, Seq_len, d_model] + let input = input.unsqueeze::<4>(); + + let output = rotary_encoding.forward(input); + let expected_output = Tensor::<3>::from_floats( + [ + [ + [1.0000, 2.0000, 3.0000, 4.0000], + [-2.3473, 7.4492, 6.9197, 8.0696], + ], + [ + [9.0000, 10.0000, 11.0000, 12.0000], + [-4.7567, 18.5034, 14.8393, 16.1492], + ], + ], + &device, + ); + + output + .squeeze_dim::<3>(0) + .to_data() + .assert_approx_eq::(&expected_output.to_data(), Tolerance::default()); + } + + #[test] + fn test_rotary_encoding_3d() { + let device = Default::default(); + let rotary_encoding = RotaryEncodingConfig::new(10, 4).init(&device); + + let input = Tensor::<3>::from_floats( + [ + [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], + [[9.0, 10.0, 11.0, 12.0], [13.0, 14.0, 15.0, 16.0]], + ], + &device, + ); + + // Input = [Batch size, Num of heads, Seq_len, d_model] + // let input = input.unsqueeze::<4>(); + + let output = rotary_encoding.forward(input); + let expected_output = Tensor::<3>::from_floats( + [ + [ + [1.0000, 2.0000, 3.0000, 4.0000], + [-2.3473, 7.4492, 6.9197, 8.0696], + ], + [ + [9.0000, 10.0000, 11.0000, 12.0000], + [-4.7567, 18.5034, 14.8393, 16.1492], + ], + ], + &device, + ); + + output + .to_data() + .assert_approx_eq::(&expected_output.to_data(), Tolerance::default()); + } + + #[test] + fn test_zero_input_rotary_encoding_forward() { + let device = Default::default(); + let rotary_encoding = RotaryEncodingConfig::new(10, 4).init(&device); + + // Use a tensor of exact zeros as input. The output rotary embedding should be zeros as well + let input = Tensor::<4>::zeros([1, 2, 2, 4], &device); + + let output = rotary_encoding.forward(input); + let expected_output = Tensor::<3>::from_floats( + [ + [ + [0.0000, 0.0000, 0.0000, 0.0000], + [0.0000, 0.0000, 0.0000, 0.0000], + ], + [ + [0.0000, 0.0000, 0.0000, 0.0000], + [0.0000, 0.0000, 0.0000, 0.0000], + ], + ], + &device, + ); + + output + .squeeze_dim::<3>(0) + .to_data() + .assert_approx_eq::(&expected_output.to_data(), Tolerance::default()); + } + + #[test] + #[should_panic] + fn test_valid_input_hidden_dim() { + // Hidden dimension must be even to be able to split into real and imaginary components + // for rotation + let d_model = 15; + let device = Default::default(); + let pe = RotaryEncodingConfig::new(10, d_model).init(&device); + let input = Tensor::<3>::zeros([1, 5, d_model], &device); + let _output = pe.forward(input); + } + + #[test] + fn test_rotary_encoding_frequencies() { + let device = Default::default(); + let rotary_encoding = RotaryEncodingConfig::new(2, 8).init(&device); + + let expected_freqs = Tensor::<3>::from_floats( + [ + [ + [1.0000, 0.0000], + [1.0000, 0.0000], + [1.0000, 0.0000], + [1.0000, 0.0000], + ], + [ + [5.4030e-01, 8.4147e-01], + [9.9500e-01, 9.9833e-02], + [9.9995e-01, 9.9998e-03], + [9.9999e-01, 9.9999e-04], + ], + ], + &device, + ) + .unsqueeze_dim::<4>(2) + .repeat_dim(2, 2) + .reshape([2, 8, 2]); + + rotary_encoding + .freq_complex + .to_data() + .assert_approx_eq::(&expected_freqs.to_data(), Tolerance::default()); + } + + fn apply_freq_scaling_by_parts(freqs: Tensor<1>) -> Tensor<1> { + // Adapted from: https://github.com/meta-llama/llama-models/blob/main/models/llama3/reference_impl/model.py#L45 + let scale_factor = 8.; + let low_freq_factor = 1.; + let high_freq_factor = 4.; + let old_context_len = 8192.; + + let low_freq_wavelen = old_context_len / low_freq_factor; + let high_freq_wavelen = old_context_len / high_freq_factor; + + let wavelen = freqs.clone().recip().mul_scalar(2. * core::f32::consts::PI); + + // if wavelen >= high_freq_wavelen + let cond = wavelen.clone().greater_equal_elem(high_freq_wavelen); + let smooth = wavelen + .clone() + .recip() + .mul_scalar(old_context_len) + .sub_scalar(low_freq_factor) + .div_scalar(high_freq_factor - low_freq_factor); + // (1 - smooth) * freq / scale_factor + smooth * freq + let new_freqs = smooth + .clone() + .neg() + .add_scalar(1.) + .mul(freqs.clone().div_scalar(scale_factor)) + .add(smooth.clone().mul(freqs.clone())); + let new_freqs = freqs.clone().mask_where(cond, new_freqs); + + // if wavelen > low_freq_wavelen + let cond = wavelen.clone().greater_elem(low_freq_wavelen); + let new_freqs = new_freqs.mask_where(cond, freqs.clone().div_scalar(scale_factor)); + + // if wavelen < high_freq_wavelen + let cond = wavelen.lower_elem(high_freq_wavelen); + new_freqs.mask_where(cond, freqs) + } + + #[test] + fn test_rotary_encoding_with_frequency_scaling() { + let device = Default::default(); + let rotary_encoding = RotaryEncodingConfig::new(2, 8) + .init_with_frequency_scaling(apply_freq_scaling_by_parts, &device); + + let expected_freqs = Tensor::<3>::from_floats( + [ + [ + [1.0000, 0.0000], + [1.0000, 0.0000], + [1.0000, 0.0000], + [1.0000, 0.0000], + ], + [ + [5.4030e-01, 8.4148e-01], + [9.9500e-01, 9.9833e-02], + [9.9995e-01, 9.9998e-03], + [1.0000, 2.1361e-04], + ], + ], + &device, + ) + .unsqueeze_dim::<4>(2) + .repeat_dim(2, 2) + .reshape([2, 8, 2]); + + rotary_encoding + .freq_complex + .to_data() + .assert_approx_eq::(&expected_freqs.to_data(), Tolerance::default()); + } + + #[test] + fn test_rotary_encoding_shift_full() { + let device = Default::default(); + let rotary_encoding = RotaryEncodingConfig::new(10, 4).init(&device); + + // Input = [Batch size, Num of heads, Seq_len, d_model] + let input = Tensor::<3>::from_floats( + [ + [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], + [[9.0, 10.0, 11.0, 12.0], [13.0, 14.0, 15.0, 16.0]], + ], + &device, + ) + .unsqueeze::<4>(); + + // Initializing for a bigger cache (e.g., max_seq_len = 10) should give the same result + // as using a smaller cache of pre-computed RoPE frequencies that are shifted to the same + // initial position + let expected_output = rotary_encoding.apply(input.clone(), 6); + + let mut rotary_encoding = RotaryEncodingConfig::new(4, 4).init(&device); + rotary_encoding.shift(6); // start > 4 will perform a full re-compute + + let output = rotary_encoding.apply(input, 0); + + output + .into_data() + .assert_approx_eq::(&expected_output.into_data(), Tolerance::default()); + } + + #[test] + fn test_rotary_encoding_shift() { + let device = Default::default(); + let rotary_encoding = RotaryEncodingConfig::new(10, 4).init(&device); + + // Input = [Batch size, Num of heads, Seq_len, d_model] + let input = Tensor::<3>::from_floats( + [ + [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], + [[9.0, 10.0, 11.0, 12.0], [13.0, 14.0, 15.0, 16.0]], + ], + &device, + ) + .unsqueeze::<4>(); + + // Initializing for a bigger cache (e.g., max_seq_len = 10) should give the same result + // as using a smaller cache of pre-computed RoPE frequencies that are shifted to the same + // initial position + let expected_output = rotary_encoding.apply(input.clone(), 2); + + let mut rotary_encoding = RotaryEncodingConfig::new(4, 4).init(&device); + rotary_encoding.shift(2); // start < 4 will shift the (current_end - start) freqs and compute the rest + + let output = rotary_encoding.apply(input, 0); + + output + .into_data() + .assert_approx_eq::(&expected_output.into_data(), Tolerance::default()); + } + + #[test] + fn test_rotary_encoding_shift_multiple() { + let device = Default::default(); + let mut rotary_encoding = RotaryEncodingConfig::new(4, 4).init(&device); + rotary_encoding.shift(2); + rotary_encoding.shift(5); + } + + #[test] + #[should_panic = "Shift start position must be monotonically increasing"] + fn test_rotary_encoding_shift_should_increase() { + let device = Default::default(); + let mut rotary_encoding = RotaryEncodingConfig::new(4, 4).init(&device); + rotary_encoding.shift(6); + rotary_encoding.shift(4); // should be monotonically increasing + } + + #[test] + fn test_rotary_encoding_reset_after_shift() { + let device = Default::default(); + let rotary_encoding = RotaryEncodingConfig::new(10, 4).init(&device); + + let input = Tensor::<3>::from_floats( + [ + [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], + [[9.0, 10.0, 11.0, 12.0], [13.0, 14.0, 15.0, 16.0]], + ], + &device, + ) + .unsqueeze::<4>(); + + let expected_output = rotary_encoding.apply(input.clone(), 0); + let expected_shifted_output = rotary_encoding.apply(input.clone(), 2); + + let mut rotary_encoding = RotaryEncodingConfig::new(4, 4).init(&device); + rotary_encoding.shift(2); + rotary_encoding.reset(); + + assert_eq!(rotary_encoding.start_offset, 0); + let output = rotary_encoding.apply(input.clone(), 0); + output + .into_data() + .assert_approx_eq::(&expected_output.into_data(), Tolerance::default()); + + // Reset must also allow subsequent shifts from the beginning again. + rotary_encoding.shift(2); + let output = rotary_encoding.apply(input, 0); + output + .into_data() + .assert_approx_eq::(&expected_shifted_output.into_data(), Tolerance::default()); + } + + #[test] + fn display() { + let config = RotaryEncodingConfig::new(10, 4); + let pe = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{pe}"), + "RotaryEncoding {d_model: 4, max_sequence_length: 10}" + ); + } +} diff --git a/crates/burn-nn/src/modules/transformer/decoder.rs b/crates/burn-nn/src/modules/transformer/decoder.rs new file mode 100644 index 0000000..0b261df --- /dev/null +++ b/crates/burn-nn/src/modules/transformer/decoder.rs @@ -0,0 +1,571 @@ +use burn_core as burn; + +use alloc::vec::Vec; + +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Initializer, Module, ModuleDisplay}; +use burn::tensor::{Bool, Device, Tensor}; + +use crate::activation::ActivationConfig; +use crate::cache::TensorCache; +use crate::{ + Dropout, DropoutConfig, LayerNorm, LayerNormConfig, + attention::{MhaCache, MhaInput, MultiHeadAttention, MultiHeadAttentionConfig}, +}; + +use super::{PositionWiseFeedForward, PositionWiseFeedForwardConfig}; + +/// Configuration to create a [Transformer Decoder](TransformerDecoder) layer using the [init function](TransformerDecoderConfig::init). +#[derive(Config, Debug)] +pub struct TransformerDecoderConfig { + /// The size of the model. + pub d_model: usize, + /// The size of the position-wise feed-forward network. + pub d_ff: usize, + /// The number of attention heads. + pub n_heads: usize, + /// The number of layers. + pub n_layers: usize, + /// The dropout rate. Default: 0.1 + #[config(default = 0.1)] + pub dropout: f64, + /// Layer norm will be applied first instead of after the other modules. + #[config(default = false)] + pub norm_first: bool, + /// Use "quiet softmax" instead of regular softmax. + /// + /// - Usage may improve performance by allowing attention heads to deposit no information (if the sequence contains no information relevant to that head). + /// - Usage may reduce the entropy of weights in the model, enhancing quantization and compression. + /// + /// Reference: + #[config(default = false)] + pub quiet_softmax: bool, + /// The type of function used to initialize neural network parameters + #[config( + default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0), fan_out_only:false}" + )] + pub initializer: Initializer, + /// The activation function used in the position-wise feed-forward network. Default: Gelu + #[config(default = "ActivationConfig::Gelu")] + pub activation: ActivationConfig, + /// The epsilon value for layer normalization. Default: 1e-5 + #[config(default = 1e-5)] + pub layer_norm_eps: f64, +} + +/// The transformer decoder module as describe in the paper [Attention Is All You Need](https://arxiv.org/abs/1706.03762). +/// +/// # Params +/// +/// - layers: transformer decoder layers with `d_model` input and output features. +/// +/// Should be created using [TransformerDecoderConfig] +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct TransformerDecoder { + /// Transformer decoder layers. + pub layers: Vec, + + /// The size of the model. + pub d_model: usize, + + /// The size of the position-wise feed-forward network. + pub d_ff: usize, + + /// The number of attention heads. + pub n_heads: usize, + + /// The number of layers. + pub n_layers: usize, + + /// The dropout rate. Default: 0.1 + pub dropout: f64, + + /// Layer norm will be applied first instead of after the other modules. + pub norm_first: bool, + + /// Use "quiet softmax" instead of regular softmax. + pub quiet_softmax: bool, +} + +impl ModuleDisplay for TransformerDecoder { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("d_model", &self.d_model) + .add("d_ff", &self.d_ff) + .add("n_heads", &self.n_heads) + .add("n_layers", &self.n_layers) + .add("dropout", &self.dropout) + .add("norm_first", &self.norm_first) + .add("quiet_softmax", &self.quiet_softmax) + .optional() + } +} + +impl TransformerDecoderConfig { + /// Initialize a new [Transformer Decoder](TransformerDecoder) module. + pub fn init(&self, device: &Device) -> TransformerDecoder { + let layers = (0..self.n_layers) + .map(|_| TransformerDecoderLayer::new(self, device)) + .collect::>(); + + TransformerDecoder { + layers, + d_model: self.d_model, + d_ff: self.d_ff, + n_heads: self.n_heads, + n_layers: self.n_layers, + dropout: self.dropout, + norm_first: self.norm_first, + quiet_softmax: self.quiet_softmax, + } + } +} + +/// [Transformer Decoder](TransformerDecoder) forward pass input argument. +#[derive(Debug)] +pub struct TransformerDecoderInput { + target: Tensor<3>, + target_mask_pad: Option>, + target_mask_attn: Option>, + memory: Tensor<3>, + memory_mask_pad: Option>, + memory_mask_attn: Option>, +} + +impl TransformerDecoderInput { + /// Create a [transformer decoder](TransformerDecoder) input argument. + pub fn new(target: Tensor<3>, memory: Tensor<3>) -> Self { + Self { + target, + target_mask_pad: None, + target_mask_attn: None, + memory, + memory_mask_pad: None, + memory_mask_attn: None, + } + } + + /// Register the memory padding mask. + pub fn memory_mask_pad(mut self, mask_pad: Tensor<2, Bool>) -> Self { + self.memory_mask_pad = Some(mask_pad); + self + } + + /// Register the memory attention mask. + pub fn memory_mask_attn(mut self, mask_attn: Tensor<3, Bool>) -> Self { + self.memory_mask_attn = Some(mask_attn); + self + } + + /// Register the target padding mask. + pub fn target_mask_pad(mut self, mask_pad: Tensor<2, Bool>) -> Self { + self.target_mask_pad = Some(mask_pad); + self + } + + /// Register the target attention mask. + pub fn target_mask_attn(mut self, mask_attn: Tensor<3, Bool>) -> Self { + self.target_mask_attn = Some(mask_attn); + self + } +} + +/// [Transformer Decoder](TransformerDecoder) layer module. +#[derive(Module, Debug)] +pub struct TransformerDecoderLayer { + /// Cross-attention module. + pub cross_attn: MultiHeadAttention, + /// Self-attention module. + pub self_attn: MultiHeadAttention, + /// Position-wise feed-forward module. + pub pwff: PositionWiseFeedForward, + /// First layer norm. + pub norm_1: LayerNorm, + /// Second layer norm. + pub norm_2: LayerNorm, + /// Third layer norm. + pub norm_3: LayerNorm, + /// Dropout. + pub dropout: Dropout, + /// Whether to apply norm first. + pub norm_first: bool, +} + +/// Autoregressive cache for a single [Transformer Decoder Layer](TransformerDecoderLayer). +pub struct TransformerDecoderLayerAutoregressiveCache { + /// Cross-attention cache. + pub cross_attn: MhaCache, + /// Self-attention cache. + pub self_attn: MhaCache, + /// Position-wise feed-forward cache. + pub pwff: TensorCache<3>, + /// First layer norm cache. + pub norm_1: TensorCache<3>, + /// Second layer norm cache. + pub norm_2: TensorCache<3>, + /// Third layer norm cache. + pub norm_3: TensorCache<3>, +} + +impl TransformerDecoderLayerAutoregressiveCache { + /// Create an empty cache. + pub fn empty() -> Self { + Self { + cross_attn: MhaCache::autoregressive_cross_attention(), + self_attn: MhaCache::autoregressive(), + pwff: TensorCache::empty(), + norm_1: TensorCache::empty(), + norm_2: TensorCache::empty(), + norm_3: TensorCache::empty(), + } + } +} + +/// Autoregressive cache for the [Transformer Decoder](TransformerDecoder) layer. +/// +/// To be used during inference when decoding tokens. +pub struct TransformerDecoderAutoregressiveCache { + layers: Vec, +} + +impl TransformerDecoderAutoregressiveCache { + fn empty(num_layers: usize) -> Self { + Self { + layers: (0..num_layers) + .map(|_| TransformerDecoderLayerAutoregressiveCache::empty()) + .collect(), + } + } +} + +impl TransformerDecoderLayer { + /// Create a new [TransformerDecoderLayer](TransformerDecoderLayer). + pub fn new(config: &TransformerDecoderConfig, device: &Device) -> Self { + let self_attn = MultiHeadAttentionConfig::new(config.d_model, config.n_heads) + .with_initializer(config.initializer.clone()) + .with_dropout(config.dropout) + .with_quiet_softmax(config.quiet_softmax) + .init(device); + + let cross_attn = MultiHeadAttentionConfig::new(config.d_model, config.n_heads) + .with_initializer(config.initializer.clone()) + .with_dropout(config.dropout) + .with_quiet_softmax(config.quiet_softmax) + .init(device); + let norm_1 = LayerNormConfig::new(config.d_model) + .with_epsilon(config.layer_norm_eps) + .init(device); + let norm_2 = LayerNormConfig::new(config.d_model) + .with_epsilon(config.layer_norm_eps) + .init(device); + let norm_3 = LayerNormConfig::new(config.d_model) + .with_epsilon(config.layer_norm_eps) + .init(device); + let dropout = DropoutConfig::new(config.dropout).init(); + let pwff = PositionWiseFeedForwardConfig::new(config.d_model, config.d_ff) + .with_initializer(config.initializer.clone()) + .with_dropout(config.dropout) + .with_activation(config.activation.clone()) + .init(device); + + Self { + cross_attn, + self_attn, + norm_1, + norm_2, + norm_3, + pwff, + dropout, + norm_first: config.norm_first, + } + } + + /// Applies the TransformerDecoder forward pass to the input tensor. + pub fn forward(&self, mut input: TransformerDecoderInput) -> TransformerDecoderInput { + // Self attention residual path. + let x = input.target; + let mut residual_path = x.clone(); + + // Normalize. + if self.norm_first { + residual_path = self.norm_3.forward(residual_path); + } + + // Self attention. + let mut self_attn_input = MhaInput::self_attn(residual_path); + if let Some(mask_pad) = &input.target_mask_pad { + self_attn_input = self_attn_input.mask_pad(mask_pad.clone()); + } + if let Some(mask_attn) = &input.target_mask_attn { + self_attn_input = self_attn_input.mask_attn(mask_attn.clone()); + } + let residual_path = self.self_attn.forward(self_attn_input).context; + + let residual_path = self.dropout.forward(residual_path); + let mut x = x + residual_path; + + // Cross attention residual path. + // Normalize. + let residual_path = if self.norm_first { + self.norm_1.forward(x.clone()) + } else { + x = self.norm_1.forward(x); + x.clone() + }; + + // Cross attention. + let mut cross_attn_input = + MhaInput::new(residual_path, input.memory.clone(), input.memory.clone()); + if let Some(mask_pad) = &input.memory_mask_pad { + cross_attn_input = cross_attn_input.mask_pad(mask_pad.clone()); + } + if let Some(mask_attn) = &input.memory_mask_attn { + cross_attn_input = cross_attn_input.mask_attn(mask_attn.clone()); + } + let residual_path = self.cross_attn.forward(cross_attn_input).context; + + let residual_path = self.dropout.forward(residual_path); + let mut x = x + residual_path; + + // Feed forward residual path. + // Normalize. + let residual_path = if self.norm_first { + self.norm_2.forward(x.clone()) + } else { + x = self.norm_2.forward(x); + x.clone() + }; + + let residual_path = self.pwff.forward(residual_path); + let residual_path = self.dropout.forward(residual_path); + let mut x = x + residual_path; + + // Main path. + // Normalize. + if !self.norm_first { + x = self.norm_3.forward(x) + } + + input.target = x; + input + } + + /// Applies the forward pass using an autoregressive cache. + pub fn forward_autoregressive_inference( + &self, + mut input: TransformerDecoderInput, + cache: &mut TransformerDecoderLayerAutoregressiveCache, + ) -> TransformerDecoderInput { + // Self attention residual path. + let x = input.target; + let mut residual_path = x.clone(); + + // Normalize. + if self.norm_first { + residual_path = cache + .norm_3 + .forward_autoregressive(residual_path, 1, |x| self.norm_3.forward(x)); + } + + // Self attention. + let mut self_attn_input = MhaInput::self_attn(residual_path); + if let Some(mask_pad) = &input.target_mask_pad { + self_attn_input = self_attn_input.mask_pad(mask_pad.clone()); + } + if let Some(mask_attn) = &input.target_mask_attn { + self_attn_input = self_attn_input.mask_attn(mask_attn.clone()); + } + let residual_path = self + .self_attn + .forward_cache(self_attn_input, &mut cache.self_attn) + .context; + + let residual_path = self.dropout.forward(residual_path); + let mut x = x + residual_path; + + // Cross attention residual path. + // Normalize. + let residual_path = if self.norm_first { + cache + .norm_1 + .forward_autoregressive(x.clone(), 1, |x| self.norm_1.forward(x)) + } else { + x = cache + .norm_1 + .forward_autoregressive(x, 1, |x| self.norm_1.forward(x)); + x.clone() + }; + + // Cross attention. + let mut cross_attn_input = + MhaInput::new(residual_path, input.memory.clone(), input.memory.clone()); + if let Some(mask_pad) = &input.memory_mask_pad { + cross_attn_input = cross_attn_input.mask_pad(mask_pad.clone()); + } + if let Some(mask_attn) = &input.memory_mask_attn { + cross_attn_input = cross_attn_input.mask_attn(mask_attn.clone()); + } + let residual_path = self + .cross_attn + .forward_cache(cross_attn_input, &mut cache.cross_attn) + .context; + + let residual_path = self.dropout.forward(residual_path); + let mut x = x + residual_path; + + // Feed forward residual path. + // Normalize. + let residual_path = if self.norm_first { + cache + .norm_2 + .forward_autoregressive(x.clone(), 1, |x| self.norm_2.forward(x)) + } else { + x = cache + .norm_2 + .forward_autoregressive(x, 1, |x| self.norm_2.forward(x)); + x.clone() + }; + + let residual_path = cache + .pwff + .forward_autoregressive(residual_path, 1, |x| self.pwff.forward(x)); + let residual_path = self.dropout.forward(residual_path); + let mut x = x + residual_path; + + // Main path. + // Normalize. + if !self.norm_first { + x = cache + .norm_3 + .forward_autoregressive(x, 1, |x| self.norm_3.forward(x)) + } + + input.target = x; + input + } +} + +impl TransformerDecoder { + /// Applies the forward pass. + pub fn forward(&self, mut input: TransformerDecoderInput) -> Tensor<3> { + for layer in self.layers.iter() { + input = layer.forward(input); + } + + input.target + } + + /// Applies the forward pass on the input using autoregressive cache. + pub fn forward_autoregressive_inference( + &self, + mut input: TransformerDecoderInput, + cache: &mut TransformerDecoderAutoregressiveCache, + ) -> Tensor<3> { + for i in 0..self.layers.len() { + let layer = self.layers.get(i).unwrap(); + let cache = cache.layers.get_mut(i).unwrap(); + + input = layer.forward_autoregressive_inference(input, cache); + } + + input.target + } + /// Create an empty autoregressive cache. + pub fn new_autoregressive_cache(&self) -> TransformerDecoderAutoregressiveCache { + TransformerDecoderAutoregressiveCache::empty(self.layers.len()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::attention::generate_autoregressive_mask; + + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn test_autoregressive_norm_last() { + let [d_model, d_ff, n_heads, num_layers] = [12, 24, 2, 3]; + let device = Device::default(); + device.seed(0); + + test_autoregressive( + TransformerDecoderConfig::new(d_model, d_ff, n_heads, num_layers) + .with_norm_first(false), + ) + } + + #[test] + fn test_autoregressive_norm_first() { + let [d_model, d_ff, n_heads, num_layers] = [12, 24, 2, 3]; + let device = Device::default(); + device.seed(0); + + test_autoregressive( + TransformerDecoderConfig::new(d_model, d_ff, n_heads, num_layers).with_norm_first(true), + ) + } + + fn test_autoregressive(config: TransformerDecoderConfig) { + let device = Default::default(); + let [batch_size, seq_length, d_model] = [3, 4, config.d_model]; + let transformer = config.init(&device); + + let memory = Tensor::arange(0..(batch_size * seq_length * d_model) as i64, &device) + .float() + .reshape([batch_size, seq_length, d_model]); + let target = Tensor::arange(0..(batch_size * seq_length * d_model) as i64, &device) + .float() + .reshape([batch_size, seq_length, d_model]); + let mask_attn = generate_autoregressive_mask(batch_size, seq_length, &target.device()); + let input = TransformerDecoderInput::new(target.clone(), memory.clone()) + .target_mask_attn(mask_attn); + + // Normal forward using masking. + let output_1 = transformer.forward(input); + + // Forward using the autoregressive cache. + let mut output_2 = Vec::new(); + let mut cache = transformer.new_autoregressive_cache(); + + for i in 1..seq_length + 1 { + let target = target.clone().slice([0..batch_size, 0..i, 0..d_model]); + + let mask_attn = generate_autoregressive_mask(batch_size, i, &target.device()); + let input = TransformerDecoderInput::new(target.clone(), memory.clone()) + .target_mask_attn(mask_attn); + let next_tok = transformer // Greedy sampling + .forward_autoregressive_inference(input, &mut cache) + .slice([0..batch_size, i - 1..i, 0..d_model]); + output_2.push(next_tok); + } + + let output_2 = Tensor::cat(output_2, 1); + + // Should produce the same tokens. + let tolerance = Tolerance::rel_abs(5e-3, 1e-4); + output_1 + .into_data() + .assert_approx_eq::(&output_2.into_data(), tolerance); + } + + #[test] + fn display() { + let config = TransformerDecoderConfig::new(2, 4, 2, 3); + let transformer = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{transformer}"), + "TransformerDecoder {d_model: 2, d_ff: 4, n_heads: 2, n_layers: 3, \ + dropout: 0.1, norm_first: false, quiet_softmax: false, params: 246}" + ); + } +} diff --git a/crates/burn-nn/src/modules/transformer/encoder.rs b/crates/burn-nn/src/modules/transformer/encoder.rs new file mode 100644 index 0000000..2b4bae6 --- /dev/null +++ b/crates/burn-nn/src/modules/transformer/encoder.rs @@ -0,0 +1,489 @@ +use burn_core as burn; + +use alloc::vec::Vec; + +use super::{PositionWiseFeedForward, PositionWiseFeedForwardConfig}; +use crate::{ + Dropout, DropoutConfig, LayerNorm, LayerNormConfig, + activation::ActivationConfig, + attention::{MhaCache, MhaInput, MultiHeadAttention, MultiHeadAttentionConfig}, + cache::TensorCache, +}; +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Initializer, Module, ModuleDisplay}; +use burn::tensor::{Bool, Device, Tensor}; + +/// Configuration to create a [Transformer Encoder](TransformerEncoder) layer using the [init function](TransformerEncoderConfig::init). +#[derive(Config, Debug)] +pub struct TransformerEncoderConfig { + /// The size of the model. + pub d_model: usize, + /// The size of the position-wise feed-forward network. + pub d_ff: usize, + /// The number of attention heads. + pub n_heads: usize, + /// The number of layers. + pub n_layers: usize, + /// The dropout rate. Default: 0.1 + #[config(default = 0.1)] + pub dropout: f64, + /// Layer norm will be applied first instead of after the other modules. + #[config(default = false)] + pub norm_first: bool, + /// Use "quiet softmax" instead of regular softmax. + /// + /// - Usage may improve performance by allowing attention heads to deposit no information (if the sequence contains no information relevant to that head). + /// - Usage may reduce the entropy of weights in the model, enhancing quantization and compression. + /// + /// Reference: + #[config(default = false)] + pub quiet_softmax: bool, + /// The type of function used to initialize neural network parameters + #[config( + default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0), fan_out_only:false}" + )] + pub initializer: Initializer, + /// The activation function used in the position-wise feed-forward network. Default: Gelu + #[config(default = "ActivationConfig::Gelu")] + pub activation: ActivationConfig, + /// The epsilon value for layer normalization. Default: 1e-5 + #[config(default = 1e-5)] + pub layer_norm_eps: f64, +} + +/// The transformer encoder module as describe in the paper [Attention Is All You Need](https://arxiv.org/abs/1706.03762). +/// +/// # Params +/// +/// - layers: transformer encoder layers with `d_model` input and output features. +/// +/// Should be created using [TransformerEncoderConfig] +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct TransformerEncoder { + /// The transformer encoder layers. + pub layers: Vec, + + /// The size of the model. + pub d_model: usize, + + /// The size of the position-wise feed-forward network. + pub d_ff: usize, + + /// The number of attention heads. + pub n_heads: usize, + + /// The number of layers. + pub n_layers: usize, + + /// The dropout rate. Default: 0.1 + pub dropout: f64, + + /// Layer norm will be applied first instead of after the other modules. + pub norm_first: bool, + + /// Use "quiet softmax" instead of regular softmax. + pub quiet_softmax: bool, +} + +impl ModuleDisplay for TransformerEncoder { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("d_model", &self.d_model) + .add("d_ff", &self.d_ff) + .add("n_heads", &self.n_heads) + .add("n_layers", &self.n_layers) + .add("dropout", &self.dropout) + .add("norm_first", &self.norm_first) + .add("quiet_softmax", &self.quiet_softmax) + .optional() + } +} + +/// [Transformer Encoder](TransformerEncoder) forward pass input argument. +#[derive(Debug)] +pub struct TransformerEncoderInput { + tensor: Tensor<3>, + mask_pad: Option>, + mask_attn: Option>, +} + +impl TransformerEncoderInput { + /// Create a [transformer encoder](TransformerEncoder) input argument. + pub fn new(tensor: Tensor<3>) -> Self { + Self { + tensor, + mask_pad: None, + mask_attn: None, + } + } + + /// Register the padding mask. + pub fn mask_pad(mut self, mask_pad: Tensor<2, Bool>) -> Self { + self.mask_pad = Some(mask_pad); + self + } + + /// Register the attention mask. + pub fn mask_attn(mut self, mask_attn: Tensor<3, Bool>) -> Self { + self.mask_attn = Some(mask_attn); + self + } +} +impl TransformerEncoderConfig { + /// Initialize a new [transformer encoder](TransformerEncoder) module. + pub fn init(&self, device: &Device) -> TransformerEncoder { + let layers = (0..self.n_layers) + .map(|_| TransformerEncoderLayer::new(self, device)) + .collect::>(); + + TransformerEncoder { + layers, + d_model: self.d_model, + d_ff: self.d_ff, + n_heads: self.n_heads, + n_layers: self.n_layers, + dropout: self.dropout, + norm_first: self.norm_first, + quiet_softmax: self.quiet_softmax, + } + } +} + +impl TransformerEncoder { + /// Applies the forward pass on the input tensor. + /// + /// # Shapes + /// + /// - tensor: `[batch_size, seq_length, d_model]` + /// - output: `[batch_size, seq_length, d_model]` + pub fn forward(&self, input: TransformerEncoderInput) -> Tensor<3> { + let mut x = input.tensor; + + for layer in self.layers.iter() { + x = layer.forward(x, input.mask_pad.clone(), input.mask_attn.clone()); + } + + x + } + /// Applies the forward pass on the input tensor using autoregressive cache. + /// + /// # Shapes + /// + /// - tensor: `[batch_size, seq_length, d_model]` + /// - output: `[batch_size, seq_length, d_model]` + pub fn forward_autoregressive_inference( + &self, + input: TransformerEncoderInput, + cache: &mut TransformerEncoderAutoregressiveCache, + ) -> Tensor<3> { + let mut x = input.tensor; + + for i in 0..self.layers.len() { + let layer = self.layers.get(i).unwrap(); + let cache = cache.layers.get_mut(i).unwrap(); + + x = layer.forward_autoregressive_inference( + x, + input.mask_pad.clone(), + input.mask_attn.clone(), + cache, + ); + } + + x + } + + /// Create an empty autoregressive cache. + pub fn new_autoregressive_cache(&self) -> TransformerEncoderAutoregressiveCache { + TransformerEncoderAutoregressiveCache::empty(self.layers.len()) + } +} + +/// Transformer encoder layer module. +#[derive(Module, Debug)] +pub struct TransformerEncoderLayer { + /// Multi-head self-attention sub-layer. + pub mha: MultiHeadAttention, + /// Position-wise feed-forward sub-layer. + pub pwff: PositionWiseFeedForward, + /// Layer normalization applied around the feed-forward sub-layer. + pub norm_1: LayerNorm, + /// Layer normalization applied around the attention sub-layer. + pub norm_2: LayerNorm, + /// Dropout module applied to residual connections. + pub dropout: Dropout, + /// If `true`, apply layer normalization before sub-layers (pre-norm), + /// otherwise apply it after (post-norm). + pub norm_first: bool, +} + +impl TransformerEncoderLayer { + /// Create a new transformer encoder layer from the given configuration. + pub fn new(config: &TransformerEncoderConfig, device: &Device) -> Self { + let mha = MultiHeadAttentionConfig::new(config.d_model, config.n_heads) + .with_initializer(config.initializer.clone()) + .with_dropout(config.dropout) + .with_quiet_softmax(config.quiet_softmax) + .init(device); + let norm_1 = LayerNormConfig::new(config.d_model) + .with_epsilon(config.layer_norm_eps) + .init(device); + let norm_2 = LayerNormConfig::new(config.d_model) + .with_epsilon(config.layer_norm_eps) + .init(device); + let dropout = DropoutConfig::new(config.dropout).init(); + let pwff = PositionWiseFeedForwardConfig::new(config.d_model, config.d_ff) + .with_initializer(config.initializer.clone()) + .with_dropout(config.dropout) + .with_activation(config.activation.clone()) + .init(device); + + Self { + mha, + norm_1, + norm_2, + pwff, + dropout, + norm_first: config.norm_first, + } + } + + /// Applies the forward pass on the input tensor. + /// + /// # Shapes + /// + /// - input: `[batch_size, seq_length, d_model]` + /// - output: `[batch_size, seq_length, d_model]` + pub fn forward( + &self, + input: Tensor<3>, + mask_pad: Option>, + mask_attn: Option>, + ) -> Tensor<3> { + // Multi-head attention residual path. + let x = input; + let mut residual_path = x.clone(); + + // Normalize. + if self.norm_first { + residual_path = self.norm_2.forward(residual_path) + } + + // Multi-head attention. + let mut input_mhs = MhaInput::self_attn(residual_path); + if let Some(mask_pad) = mask_pad { + input_mhs = input_mhs.mask_pad(mask_pad); + } + if let Some(mask_attn) = mask_attn { + input_mhs = input_mhs.mask_attn(mask_attn); + } + let residual_path = self.mha.forward(input_mhs).context; + + let residual_path = self.dropout.forward(residual_path); + let mut x = x + residual_path; + + // Feed forward residual path. + // Normalize. + let residual_path = if self.norm_first { + self.norm_1.forward(x.clone()) + } else { + x = self.norm_1.forward(x); + x.clone() + }; + + // Feed forward. + let residual_path = self.pwff.forward(residual_path); + let residual_path = self.dropout.forward(residual_path); + let mut x = x + residual_path; + + // Main path. + // Normalize. + if !self.norm_first { + x = self.norm_2.forward(x) + } + + x + } + + /// Applies the forward pass using an autoregressive cache. + pub fn forward_autoregressive_inference( + &self, + input: Tensor<3>, + mask_pad: Option>, + mask_attn: Option>, + cache: &mut TransformerEncoderLayerAutoregressiveCache, + ) -> Tensor<3> { + // Multi-head attention residual path. + let x = input; + let mut residual_path = x.clone(); + + // Normalize. + if self.norm_first { + residual_path = cache + .norm_2 + .forward_autoregressive(residual_path, 1, |x| self.norm_2.forward(x)) + } + + // Multi-head attention. + let mut input_mhs = MhaInput::self_attn(residual_path); + if let Some(mask_pad) = mask_pad { + input_mhs = input_mhs.mask_pad(mask_pad); + } + if let Some(mask_attn) = mask_attn { + input_mhs = input_mhs.mask_attn(mask_attn); + } + let residual_path = self.mha.forward_cache(input_mhs, &mut cache.mha).context; + + let residual_path = self.dropout.forward(residual_path); + let mut x = x + residual_path; + + // Feed forward residual path. + // Normalize. + let residual_path = if self.norm_first { + cache + .norm_1 + .forward_autoregressive(x.clone(), 1, |x| self.norm_1.forward(x)) + } else { + x = cache + .norm_1 + .forward_autoregressive(x, 1, |x| self.norm_1.forward(x)); + x.clone() + }; + + // Feed forward. + let residual_path = cache + .pwff + .forward_autoregressive(residual_path, 1, |x| self.pwff.forward(x)); + let residual_path = self.dropout.forward(residual_path); + let mut x = x + residual_path; + + // Main path. + // Normalize. + if !self.norm_first { + x = cache + .norm_2 + .forward_autoregressive(x, 1, |x| self.norm_2.forward(x)) + } + + x + } +} + +/// Autoregressive cache for a single [Transformer Encoder Layer](TransformerEncoderLayer). +pub struct TransformerEncoderLayerAutoregressiveCache { + /// Multi-head attention cache. + pub mha: MhaCache, + /// Position-wise feed-forward cache. + pub pwff: TensorCache<3>, + /// First layer norm cache. + pub norm_1: TensorCache<3>, + /// Second layer norm cache. + pub norm_2: TensorCache<3>, +} + +impl TransformerEncoderLayerAutoregressiveCache { + /// Create an empty cache. + pub fn empty() -> Self { + Self { + mha: MhaCache::autoregressive(), + pwff: TensorCache::empty(), + norm_1: TensorCache::empty(), + norm_2: TensorCache::empty(), + } + } +} + +/// Autoregressive cache for the [Transformer Encoder](TransformerEncoder) layer. +/// +/// To be used during inference when decoding tokens. +pub struct TransformerEncoderAutoregressiveCache { + layers: Vec, +} + +impl TransformerEncoderAutoregressiveCache { + fn empty(num_layers: usize) -> Self { + Self { + layers: (0..num_layers) + .map(|_| TransformerEncoderLayerAutoregressiveCache::empty()) + .collect(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::attention::generate_autoregressive_mask; + use burn::tensor::Distribution; + use burn::tensor::Tolerance; + type FT = f32; + + #[test] + fn test_autoregressive_norm_last() { + let [d_model, d_ff, n_heads, num_layers] = [12, 24, 2, 3]; + test_autoregressive( + TransformerEncoderConfig::new(d_model, d_ff, n_heads, num_layers) + .with_norm_first(false), + ) + } + + #[test] + fn test_autoregressive_norm_first() { + let [d_model, d_ff, n_heads, num_layers] = [12, 24, 2, 3]; + test_autoregressive( + TransformerEncoderConfig::new(d_model, d_ff, n_heads, num_layers).with_norm_first(true), + ) + } + + fn test_autoregressive(config: TransformerEncoderConfig) { + let [batch_size, seq_length, d_model] = [3, 4, config.d_model]; + let device = Default::default(); + let transformer = config.init(&device); + + let tensor = Tensor::<3>::random( + [batch_size, seq_length, d_model], + Distribution::Default, + &device, + ); + let mask_attn = generate_autoregressive_mask(batch_size, seq_length, &tensor.device()); + let input = TransformerEncoderInput::new(tensor.clone()).mask_attn(mask_attn); + + let output_1 = transformer.forward(input); + let mut output_2 = Vec::new(); + let mut cache = transformer.new_autoregressive_cache(); + + for i in 1..seq_length + 1 { + let tensor = tensor.clone().slice([0..batch_size, 0..i, 0..d_model]); + let input = TransformerEncoderInput::new(tensor.clone()); + let next_tok = transformer + .forward_autoregressive_inference(input, &mut cache) + .slice([0..batch_size, i - 1..i, 0..d_model]); + output_2.push(next_tok); + } + + let output_2 = Tensor::cat(output_2, 1); + + output_1 + .into_data() + .assert_approx_eq::(&output_2.into_data(), Tolerance::permissive()); + } + + #[test] + fn display() { + let config = TransformerEncoderConfig::new(2, 4, 2, 3); + let transformer = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{transformer}"), + "TransformerEncoder {d_model: 2, d_ff: 4, n_heads: 2, \ + n_layers: 3, dropout: 0.1, norm_first: false, quiet_softmax: false, params: 162}" + ); + } +} diff --git a/crates/burn-nn/src/modules/transformer/mod.rs b/crates/burn-nn/src/modules/transformer/mod.rs new file mode 100644 index 0000000..54397c6 --- /dev/null +++ b/crates/burn-nn/src/modules/transformer/mod.rs @@ -0,0 +1,7 @@ +mod decoder; +mod encoder; +mod pwff; + +pub use decoder::*; +pub use encoder::*; +pub use pwff::*; diff --git a/crates/burn-nn/src/modules/transformer/pwff.rs b/crates/burn-nn/src/modules/transformer/pwff.rs new file mode 100644 index 0000000..6c4d051 --- /dev/null +++ b/crates/burn-nn/src/modules/transformer/pwff.rs @@ -0,0 +1,128 @@ +use burn_core as burn; + +use crate::activation::{Activation, ActivationConfig}; +use crate::{Dropout, DropoutConfig, Linear, LinearConfig}; +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Initializer, Module, ModuleDisplay}; +use burn::tensor::{Device, Tensor}; + +/// Configuration to create a [position-wise feed-forward](PositionWiseFeedForward) layer using the [init function](PositionWiseFeedForwardConfig::init). +#[derive(Config, Debug)] +pub struct PositionWiseFeedForwardConfig { + /// The size of the input and output features. + pub d_model: usize, + /// The size of the hidden inner features. + pub d_ff: usize, + /// The dropout rate. Default: 0.1 + #[config(default = 0.1)] + pub dropout: f64, + /// The type of function used to initialize neural network parameters + #[config( + default = "Initializer::KaimingUniform{gain:1.0/num_traits::Float::sqrt(3.0), fan_out_only:false}" + )] + pub initializer: Initializer, + /// The activation function used between the two linear layers. Default: Gelu + #[config(default = "ActivationConfig::Gelu")] + pub activation: ActivationConfig, +} + +/// Applies the position-wise feed-forward network to the input tensor from the paper [Attention Is All You Need](https://arxiv.org/pdf/1706.03762v7). +/// +/// # Params +/// +/// - linear inner: Linear layer with `d_model` input features and `d_ff` output features. +/// - linear outer: Linear layer with `d_ff` input features and `d_model` output features. +/// +/// `FFN(x) = max(0, xW1 + b1)W2 + b2` +/// +/// Should be created using [PositionWiseFeedForwardConfig] +/// +/// # Notes +/// +/// The `activation` field is currently marked `#[module(skip)]` for backward +/// compatibility with records saved before this field was introduced (when +/// the activation was always `Gelu` and had no state). This means activation +/// state is **not persisted** when saving or loading records. +/// +/// For stateless activations (GELU, ReLU, etc.) this has no effect. +/// **If you are using `SwiGLU`, its learnable parameters will not be saved or +/// loaded correctly.** +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct PositionWiseFeedForward { + /// Linear layer with `d_model` input features and `d_ff` output features. + pub linear_inner: Linear, + /// Linear layer with `d_ff` input features and `d_model` output features. + pub linear_outer: Linear, + /// Dropout layer. + pub dropout: Dropout, + /// Activation function. + #[module(skip)] // for backward compatibility with previous `gelu` field name + pub activation: Activation, +} + +impl ModuleDisplay for PositionWiseFeedForward { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let [d_model, dff] = self.linear_inner.weight.shape().dims(); + + content + .add("d_model", &d_model) + .add("d_ff", &dff) + .add("prob", &self.dropout.prob) + .optional() + } +} + +impl PositionWiseFeedForwardConfig { + /// Initialize a new [position-wise feed-forward](PositionWiseFeedForward) module. + pub fn init(&self, device: &Device) -> PositionWiseFeedForward { + PositionWiseFeedForward { + linear_inner: LinearConfig::new(self.d_model, self.d_ff) + .with_initializer(self.initializer.clone()) + .init(device), + linear_outer: LinearConfig::new(self.d_ff, self.d_model) + .with_initializer(self.initializer.clone()) + .init(device), + dropout: DropoutConfig::new(self.dropout).init(), + activation: self.activation.init(device), + } + } +} + +impl PositionWiseFeedForward { + /// Applies the forward pass on the input tensor. + /// + /// # Shapes + /// + /// - tensor: `[batch_size, seq_length, d_model]` + /// - output: `[batch_size, seq_length, d_model]` + pub fn forward(&self, input: Tensor) -> Tensor { + let x = self.linear_inner.forward(input); + let x = self.activation.forward(x); + let x = self.dropout.forward(x); + + self.linear_outer.forward(x) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn display() { + let config = PositionWiseFeedForwardConfig::new(2, 4); + let pwff = config.init(&Default::default()); + + assert_eq!( + alloc::format!("{pwff}"), + "PositionWiseFeedForward {d_model: 2, d_ff: 4, prob: 0.1, params: 22}" + ); + } +} diff --git a/crates/burn-nn/src/modules/unfold.rs b/crates/burn-nn/src/modules/unfold.rs new file mode 100644 index 0000000..c16fe9c --- /dev/null +++ b/crates/burn-nn/src/modules/unfold.rs @@ -0,0 +1,103 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Module, ModuleDisplay}; + +use burn::tensor::Tensor; +use burn::tensor::module::unfold4d; +use burn::tensor::ops::UnfoldOptions; + +/// Configuration to create an [unfold 4d](Unfold4d) layer using the [init function](Unfold4dConfig::init). +#[derive(Config, Debug)] +pub struct Unfold4dConfig { + /// The size of the kernel. + pub kernel_size: [usize; 2], + /// The stride of the convolution. + #[config(default = "[1, 1]")] + pub stride: [usize; 2], + /// Spacing between kernel elements. + #[config(default = "[1, 1]")] + pub dilation: [usize; 2], + /// The padding configuration. + #[config(default = "[0, 0]")] + pub padding: [usize; 2], +} + +/// Four-dimensional unfolding. +/// +/// Should be created with [Unfold4dConfig]. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Unfold4d { + /// The size of the kernel. + pub kernel_size: [usize; 2], + /// The stride of the convolution. + pub stride: [usize; 2], + /// Spacing between kernel elements. + pub dilation: [usize; 2], + /// The padding configuration. + pub padding: [usize; 2], +} + +impl ModuleDisplay for Unfold4d { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("kernel_size", &alloc::format!("{:?}", self.kernel_size)) + .add("stride", &alloc::format!("{:?}", self.stride)) + .add("dilation", &alloc::format!("{:?}", self.dilation)) + .add("padding", &alloc::format!("{:?}", self.padding)) + .optional() + } +} + +impl Unfold4dConfig { + /// Initializes a new [Unfold4d] module. + pub fn init(&self) -> Unfold4d { + Unfold4d { + kernel_size: self.kernel_size, + stride: self.stride, + dilation: self.dilation, + padding: self.padding, + } + } +} + +impl Unfold4d { + /// Applies the forward pass on the input tensor. + /// + /// See [unfold4d](burn::tensor::module::unfold4d) for more information. + /// + /// # Shapes + /// + /// input: `[batch_size, channels_in, height, width]` + /// returns: `[batch_size, channels_in * kernel_size_1 * kernel_size_2, number of blocks]` + pub fn forward(&self, input: Tensor<4>) -> Tensor<3> { + unfold4d( + input, + self.kernel_size, + UnfoldOptions::new(self.stride, self.padding, self.dilation), + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn display() { + let config = Unfold4dConfig::new([3, 3]); + let unfold = config.init(); + + assert_eq!( + alloc::format!("{unfold}"), + "Unfold4d {kernel_size: [3, 3], stride: [1, 1], dilation: [1, 1], padding: [0, 0]}" + ); + } +} diff --git a/crates/burn-nn/src/padding.rs b/crates/burn-nn/src/padding.rs new file mode 100644 index 0000000..cfacfc6 --- /dev/null +++ b/crates/burn-nn/src/padding.rs @@ -0,0 +1,247 @@ +use burn_core as burn; + +use burn::config::Config; + +/// Calculate asymmetric padding for "same" convolution. +/// Returns (start_padding, end_padding) where start is applied first (top/left). +/// For odd total padding, the extra pad goes to the end (bottom/right) following ONNX convention. +fn calculate_same_padding(kernel_size: usize, stride: usize, size_in: usize) -> (usize, usize) { + let size_out = size_in.div_ceil(stride); // ceil division for same padding + let total_padding = if size_out > 0 { + let needed = (size_out - 1) * stride + kernel_size; + needed.saturating_sub(size_in) + } else { + 0 + }; + let pad_start = total_padding / 2; + let pad_end = total_padding - pad_start; + (pad_start, pad_end) +} + +/// Padding configuration for 1D operators. +#[derive(Config, Debug, PartialEq)] +pub enum PaddingConfig1d { + /// Dynamically calculates padding to ensure output size matches input size. + Same, + /// No padding applied. + Valid, + /// Applies explicit padding values. + /// Format: (left, right) + /// For symmetric padding, use the same value for both (e.g., `Explicit(1, 1)`). + Explicit(usize, usize), +} + +impl PaddingConfig1d { + /// Calculate padding as (left, right) pair for 1D operations. + /// For `Same` padding, this computes the actual asymmetric padding if needed. + pub(crate) fn calculate_padding_1d_pair( + &self, + length: usize, + kernel_size: usize, + stride: usize, + ) -> (usize, usize) { + match self { + Self::Valid => (0, 0), + Self::Same => calculate_same_padding(kernel_size, stride, length), + Self::Explicit(left, right) => (*left, *right), + } + } +} + +/// Padding configuration for 2D operators. +#[derive(Config, Debug, PartialEq)] +pub enum PaddingConfig2d { + /// Dynamically calculates padding to preserve input dimensions in output. + Same, + /// No padding applied. + Valid, + /// Applies explicit padding values. + /// Format: (top, left, bottom, right) + /// For symmetric padding, use matching values (e.g., `Explicit(1, 1, 1, 1)`). + Explicit(usize, usize, usize, usize), +} + +impl PaddingConfig2d { + /// Calculate padding as ((top, bottom), (left, right)) pairs for 2D operations. + /// For `Same` padding, this computes the actual asymmetric padding if needed. + pub(crate) fn calculate_padding_2d_pairs( + &self, + height: usize, + width: usize, + kernel_size: &[usize; 2], + stride: &[usize; 2], + ) -> ((usize, usize), (usize, usize)) { + match self { + Self::Valid => ((0, 0), (0, 0)), + Self::Same => { + let (top, bottom) = calculate_same_padding(kernel_size[0], stride[0], height); + let (left, right) = calculate_same_padding(kernel_size[1], stride[1], width); + ((top, bottom), (left, right)) + } + Self::Explicit(top, left, bottom, right) => ((*top, *bottom), (*left, *right)), + } + } + + /// Calculate symmetric padding for 2D operations. + /// Returns padding values [height, width] (same for both sides). + /// Panics if asymmetric padding is detected. + pub(crate) fn calculate_padding_2d( + &self, + height: usize, + width: usize, + kernel_size: &[usize; 2], + stride: &[usize; 2], + ) -> [usize; 2] { + let ((top, bottom), (left, right)) = + self.calculate_padding_2d_pairs(height, width, kernel_size, stride); + if top != bottom || left != right { + panic!("Asymmetric padding should be handled via calculate_padding_2d_pairs()") + } + [top, left] + } +} + +/// Padding configuration for 3D operators. +#[derive(Config, Debug, PartialEq)] +pub enum PaddingConfig3d { + /// Dynamically calculates padding to preserve input dimensions in output. + Same, + /// No padding applied. + Valid, + /// Applies explicit symmetric padding values. + /// Format: (depth, height, width) — same padding on both sides of each dimension. + Explicit(usize, usize, usize), +} + +impl PaddingConfig3d { + /// Calculate symmetric padding for 3D operations. + /// Returns padding values [depth, height, width] (same for both sides). + pub(crate) fn calculate_padding_3d( + &self, + depth: usize, + height: usize, + width: usize, + kernel_size: &[usize; 3], + stride: &[usize; 3], + ) -> [usize; 3] { + match self { + Self::Valid => [0, 0, 0], + Self::Same => { + let (front, back) = calculate_same_padding(kernel_size[0], stride[0], depth); + let (top, bottom) = calculate_same_padding(kernel_size[1], stride[1], height); + let (left, right) = calculate_same_padding(kernel_size[2], stride[2], width); + if front != back || top != bottom || left != right { + panic!( + "Asymmetric 3D 'Same' padding is not supported. \ + Use odd kernel sizes for symmetric padding." + ) + } + [front, top, left] + } + Self::Explicit(depth, height, width) => [*depth, *height, *width], + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // ==================== PaddingConfig1d Tests ==================== + + #[test] + fn test_padding_config_1d_calculate_pair_valid() { + let padding = PaddingConfig1d::Valid; + assert_eq!(padding.calculate_padding_1d_pair(10, 3, 1), (0, 0)); + } + + #[test] + fn test_padding_config_1d_calculate_pair_explicit() { + let padding = PaddingConfig1d::Explicit(1, 2); + assert_eq!(padding.calculate_padding_1d_pair(10, 3, 1), (1, 2)); + } + + #[test] + fn test_padding_config_1d_calculate_pair_same() { + let padding = PaddingConfig1d::Same; + // kernel=3, stride=1, length=10: total=2, start=1, end=1 + assert_eq!(padding.calculate_padding_1d_pair(10, 3, 1), (1, 1)); + } + + // ==================== PaddingConfig2d Tests ==================== + + #[test] + fn test_padding_config_2d_calculate_pairs_valid() { + let padding = PaddingConfig2d::Valid; + assert_eq!( + padding.calculate_padding_2d_pairs(10, 10, &[3, 3], &[1, 1]), + ((0, 0), (0, 0)) + ); + } + + #[test] + fn test_padding_config_2d_calculate_pairs_explicit() { + let padding = PaddingConfig2d::Explicit(1, 2, 3, 4); + assert_eq!( + padding.calculate_padding_2d_pairs(10, 10, &[3, 3], &[1, 1]), + ((1, 3), (2, 4)) + ); + } + + #[test] + fn test_padding_config_2d_calculate_symmetric_valid() { + let padding = PaddingConfig2d::Valid; + assert_eq!( + padding.calculate_padding_2d(10, 10, &[3, 3], &[1, 1]), + [0, 0] + ); + } + + #[test] + fn test_padding_config_2d_calculate_symmetric_explicit() { + let padding = PaddingConfig2d::Explicit(2, 3, 2, 3); + assert_eq!( + padding.calculate_padding_2d(10, 10, &[3, 3], &[1, 1]), + [2, 3] + ); + } + + #[test] + #[should_panic( + expected = "Asymmetric padding should be handled via calculate_padding_2d_pairs" + )] + fn test_padding_config_2d_calculate_symmetric_asymmetric_panics() { + let padding = PaddingConfig2d::Explicit(1, 2, 3, 4); + let _ = padding.calculate_padding_2d(10, 10, &[3, 3], &[1, 1]); + } + + // ==================== PaddingConfig3d Tests ==================== + + #[test] + fn test_padding_config_3d_calculate_valid() { + let padding = PaddingConfig3d::Valid; + assert_eq!( + padding.calculate_padding_3d(10, 10, 10, &[3, 3, 3], &[1, 1, 1]), + [0, 0, 0] + ); + } + + #[test] + fn test_padding_config_3d_calculate_explicit() { + let padding = PaddingConfig3d::Explicit(1, 2, 3); + assert_eq!( + padding.calculate_padding_3d(10, 10, 10, &[3, 3, 3], &[1, 1, 1]), + [1, 2, 3] + ); + } + + #[test] + fn test_padding_config_3d_calculate_same_odd_kernel() { + let padding = PaddingConfig3d::Same; + // kernel=3, stride=1: total=2, symmetric (1,1) per dim + assert_eq!( + padding.calculate_padding_3d(10, 10, 10, &[3, 3, 3], &[1, 1, 1]), + [1, 1, 1] + ); + } +} diff --git a/crates/burn-nn/tests/quantize.rs b/crates/burn-nn/tests/quantize.rs new file mode 100644 index 0000000..1cde05a --- /dev/null +++ b/crates/burn-nn/tests/quantize.rs @@ -0,0 +1,162 @@ +use burn_core as burn; + +use burn::module::{Module, Quantizer}; +use burn::tensor::{ + Distribution, Tensor, Tolerance, + quantization::{Calibration, QuantLevel, QuantParam, QuantScheme, QuantValue}, +}; +use burn_nn::{ + LinearConfig, + transformer::{TransformerEncoder, TransformerEncoderConfig, TransformerEncoderInput}, +}; + +fn should_quantize_module Tensor>( + module: M, + scheme: QuantScheme, + func: F, + tolerance: Tolerance, +) { + let result = func(&module); + + let calibration = Calibration::MinMax; + let mut quantizer = Quantizer { + calibration, + scheme, + }; + let q_module = module.quantize_weights(&mut quantizer); + let q_result = func(&q_module); + + result + .into_data() + .assert_approx_eq::(&q_result.into_data(), tolerance); +} + +#[test] +fn should_quantize_transformer() { + let device = Default::default(); + let transformer: TransformerEncoder = + TransformerEncoderConfig::new(128, 256, 2, 2).init(&device); + let signal = Tensor::random([2, 32, 128], Distribution::Default, &device); + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::block([32])) + .with_param(QuantParam::F32); + + should_quantize_module( + transformer, + scheme, + |tr| tr.forward(TransformerEncoderInput::new(signal.clone())), + Tolerance::rel_abs(1e-2, 2e-2), // slightly higher abs tolerance (permissive: 1e-2) + ); +} + +#[test] +fn should_quantize_linear_128_256() { + let device = Default::default(); + let transformer = LinearConfig::new(128, 256).with_bias(false).init(&device); + let signal = Tensor::<2>::random([1, 128], Distribution::Default, &device); + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::Tensor) + .with_param(QuantParam::F32); + + should_quantize_module( + transformer, + scheme, + |tr| tr.forward(signal.clone()), + Tolerance::permissive(), + ); +} + +#[test] +fn should_quantize_linear() { + let device = Default::default(); + let transformer = LinearConfig::new(32, 32).with_bias(false).init(&device); + let signal = Tensor::<2>::random([1, 32], Distribution::Default, &device); + // Default scheme should select supported QuantStore default + // TODO: set native if dtype is supported by the test backend + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::Tensor) + // .with_store(QuantStore::Native) + .with_param(QuantParam::F32); + + should_quantize_module( + transformer, + scheme, + |tr| tr.forward(signal.clone()), + Tolerance::permissive(), + ); +} + +#[test] +fn should_quantize_linear_weights() { + let device = Default::default(); + let transformer = LinearConfig::new(32, 32).with_bias(false).init(&device); + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::Tensor) + .with_param(QuantParam::F32); + + should_quantize_module( + transformer, + scheme, + |tr| tr.weight.val().dequantize(), + Tolerance::permissive(), + ); +} + +#[test] +fn should_quantize_linear_blocks() { + let device = Default::default(); + let transformer = LinearConfig::new(32, 32).with_bias(false).init(&device); + let signal = Tensor::<2>::random([1, 32], Distribution::Default, &device); + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::block([16])) + // .with_store(QuantStore::Native) + .with_param(QuantParam::F32); + + should_quantize_module( + transformer, + scheme, + |tr| tr.forward(signal.clone()), + Tolerance::permissive(), + ); +} + +#[test] +fn should_quantize_linear_weights_blocks() { + let device = Default::default(); + let transformer = LinearConfig::new(32, 32).with_bias(false).init(&device); + let scheme = device + .settings() + .quantization + .scheme + .with_value(QuantValue::Q8S) + .with_level(QuantLevel::block([16])) + // .with_store(QuantStore::Native) + .with_param(QuantParam::F32); + + should_quantize_module( + transformer, + scheme, + |tr| tr.weight.val().dequantize(), + Tolerance::permissive(), + ); +} diff --git a/crates/burn-no-std-tests/Cargo.toml b/crates/burn-no-std-tests/Cargo.toml new file mode 100644 index 0000000..d09e247 --- /dev/null +++ b/crates/burn-no-std-tests/Cargo.toml @@ -0,0 +1,26 @@ +[package] +authors = [ + "nathanielsimard ", + "Dilshod Tadjibaev (@antimora)", +] +edition.workspace = true +license.workspace = true +name = "burn-no-std-tests" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-no-std-tests" +version.workspace = true + +[lints] +workspace = true + +[features] +default = [] + +[dependencies] + +# ** Please make sure all dependencies support no_std ** + +burn = { workspace = true } +burn-flex = { workspace = true } + +burn-store = { workspace = true, features = ["safetensors", "burnpack"] } diff --git a/crates/burn-no-std-tests/README.md b/crates/burn-no-std-tests/README.md new file mode 100644 index 0000000..0dae942 --- /dev/null +++ b/crates/burn-no-std-tests/README.md @@ -0,0 +1,29 @@ +The `burn-no-std-tests` contains integration tests aimed to check `no_std` compatibility of `burn`, `burn-core`, `burn-tensor` and `burn-flex` packages. + +Currently there is only a minimal test that checks if mnist model can be built with `no_std`. More tests should be added to check completeness. + +The continuous integration (CI) should build with additional targets: + + * `wasm32-unknown-unknown` - WebAssembly + * `thumbv7m-none-eabi` - ARM Cortex-M3 + * `thumbv6m-none-eabi` - ARM Cortex-M0+ + +Shell commands to build and test the package: + +```sh + +# install the new targets if not installed previously +rustup target add thumbv6m-none-eabi +rustup target add thumbv7m-none-eabi +rustup target add wasm32-unknown-unknown + +# build for various targets +cargo build # regular build +cargo build --target thumbv7m-none-eabi +cargo build --target wasm32-unknown-unknown +RUSTFLAGS="--cfg portable_atomic_unsafe_assume_single_core" cargo build --target thumbv6m-none-eabi + +# test +cargo test + + ``` \ No newline at end of file diff --git a/crates/burn-no-std-tests/src/burnpack.rs b/crates/burn-no-std-tests/src/burnpack.rs new file mode 100644 index 0000000..76fb2dd --- /dev/null +++ b/crates/burn-no-std-tests/src/burnpack.rs @@ -0,0 +1,158 @@ +// Test Burnpack storage in no-std environment + +use burn::{ + module::Module, + nn, + tensor::{Device, Tensor}, +}; + +use burn_store::{BurnpackStore, ModuleSnapshot, PathFilter}; + +/// Simple model for testing Burnpack storage +#[derive(Module, Debug)] +pub struct TestModel { + linear1: nn::Linear, + linear2: nn::Linear, + batch_norm: nn::BatchNorm, +} + +impl TestModel { + pub fn new(device: &Device) -> Self { + Self { + linear1: nn::LinearConfig::new(10, 20).init(device), + linear2: nn::LinearConfig::new(20, 10).init(device), + batch_norm: nn::BatchNormConfig::new(10).init(device), + } + } + + pub fn forward(&self, x: Tensor<2>) -> Tensor<2> { + let x = self.linear1.forward(x); + let x = self.linear2.forward(x); + // Apply batch norm (expand to 3D, apply, then squeeze back) + let x: Tensor<3> = x.unsqueeze_dim(2); + let x = self.batch_norm.forward(x); + x.squeeze_dim(2) + } +} + +/// Test basic Burnpack save and load in no-std +pub fn test_burnpack_basic(device: &Device) { + // Create a model + let model = TestModel::new(device); + + // Save to bytes (no file I/O in no-std) + let mut save_store = BurnpackStore::from_bytes(None); + model + .save_into(&mut save_store) + .expect("Failed to save model"); + + // Get the serialized bytes + let bytes = save_store.get_bytes().expect("Failed to get bytes"); + + // Load from bytes + let mut load_store = BurnpackStore::from_bytes(Some(bytes)); + let mut loaded_model = TestModel::new(device); + let result = loaded_model + .load_from(&mut load_store) + .expect("Failed to load model"); + + // Verify all tensors were loaded + assert!(result.is_success(), "Should have no errors"); + assert!(!result.applied.is_empty(), "Should have loaded tensors"); + + // Test that the model still works + let input = Tensor::<2>::ones([2, 10], device); + let _output = loaded_model.forward(input); +} + +/// Test Burnpack with filtering in no-std +pub fn test_burnpack_filtering(device: &Device) { + let model = TestModel::new(device); + + // Save only linear1 weights + let filter = PathFilter::new() + .with_full_path("linear1.weight") + .with_full_path("linear1.bias"); + let mut save_store = BurnpackStore::from_bytes(None).with_filter(filter); + model + .save_into(&mut save_store) + .expect("Failed to save filtered model"); + + let bytes = save_store.get_bytes().expect("Failed to get bytes"); + + // Load with partial loading allowed + let mut load_store = BurnpackStore::from_bytes(Some(bytes)).allow_partial(true); + let mut partial_model = TestModel::new(device); + let result = partial_model + .load_from(&mut load_store) + .expect("Failed to load partial model"); + + // Verify that only linear1 was loaded + assert_eq!(result.applied.len(), 2, "Should have loaded 2 tensors"); + assert!(!result.missing.is_empty(), "Should have missing tensors"); +} + +/// Test Burnpack with metadata in no-std +pub fn test_burnpack_metadata(device: &Device) { + let model = TestModel::new(device); + + // Save with metadata + let mut save_store = BurnpackStore::from_bytes(None) + .metadata("version", "1.0.0") + .metadata("environment", "no-std") + .metadata("model_type", "test"); + model + .save_into(&mut save_store) + .expect("Failed to save model with metadata"); + + let bytes = save_store.get_bytes().expect("Failed to get bytes"); + + // Load and verify it works + let mut load_store = BurnpackStore::from_bytes(Some(bytes)); + let mut loaded_model = TestModel::new(device); + let result = loaded_model + .load_from(&mut load_store) + .expect("Failed to load model with metadata"); + + assert!(result.is_success(), "Should load successfully"); +} + +// Note: Key remapping test is omitted as KeyRemapper requires std feature + +// Note: Regex filtering test is omitted as with_regex requires std feature + +/// Test Burnpack with match_all in no-std +pub fn test_burnpack_match_all(device: &Device) { + let model = TestModel::new(device); + + // Save with match_all (should save everything) + let mut save_store = BurnpackStore::from_bytes(None).match_all(); + model + .save_into(&mut save_store) + .expect("Failed to save model"); + + let bytes = save_store.get_bytes().expect("Failed to get bytes"); + + // Load everything + let mut load_store = BurnpackStore::from_bytes(Some(bytes)); + let mut loaded_model = TestModel::new(device); + let result = loaded_model + .load_from(&mut load_store) + .expect("Failed to load model"); + + assert!(result.is_success(), "Should load successfully"); + // linear1 (weight, bias) + linear2 (weight, bias) + batch_norm (4 params) + assert_eq!(result.applied.len(), 8, "Should load all 8 tensors"); + assert!(result.missing.is_empty(), "Should have no missing tensors"); + assert!(result.unused.is_empty(), "Should have no unused tensors"); +} + +/// Run all Burnpack no-std tests +pub fn run_all_tests(device: &Device) { + test_burnpack_basic(device); + test_burnpack_filtering(device); + test_burnpack_metadata(device); + // test_burnpack_remapping requires KeyRemapper which needs std + // test_burnpack_regex_filter requires with_regex which needs std + test_burnpack_match_all(device); +} diff --git a/crates/burn-no-std-tests/src/conv.rs b/crates/burn-no-std-tests/src/conv.rs new file mode 100644 index 0000000..5262af8 --- /dev/null +++ b/crates/burn-no-std-tests/src/conv.rs @@ -0,0 +1,49 @@ +// Originally copied from the burn/examples/mnist package + +use burn::{ + config::Config, + module::Module, + nn, + tensor::{Device, Tensor}, +}; + +#[derive(Module, Debug)] +pub struct ConvBlock { + conv: nn::conv::Conv2d, + pool: nn::pool::MaxPool2d, + activation: nn::Gelu, +} + +#[derive(Config, Debug)] +pub struct ConvBlockConfig { + channels: [usize; 2], + #[config(default = "[3, 3]")] + kernel_size: [usize; 2], +} + +impl ConvBlock { + pub fn new(config: &ConvBlockConfig, device: &Device) -> Self { + let conv = nn::conv::Conv2dConfig::new(config.channels, config.kernel_size) + .with_padding(nn::PaddingConfig2d::Same) + .init(device); + let pool = nn::pool::MaxPool2dConfig::new(config.kernel_size) + .with_strides([1, 1]) + .with_padding(nn::PaddingConfig2d::Same) + .init(); + let activation = nn::Gelu::new(); + + Self { + conv, + pool, + activation, + } + } + + pub fn forward(&self, input: Tensor<4>) -> Tensor<4> { + let x = self.conv.forward(input.clone()); + let x = self.pool.forward(x); + let x = self.activation.forward(x); + + (x + input) / 2.0 + } +} diff --git a/crates/burn-no-std-tests/src/lib.rs b/crates/burn-no-std-tests/src/lib.rs new file mode 100644 index 0000000..58de35e --- /dev/null +++ b/crates/burn-no-std-tests/src/lib.rs @@ -0,0 +1,9 @@ +#![no_std] + +pub mod burnpack; +pub mod conv; +pub mod mlp; +pub mod model; +pub mod safetensors; + +extern crate alloc; diff --git a/crates/burn-no-std-tests/src/mlp.rs b/crates/burn-no-std-tests/src/mlp.rs new file mode 100644 index 0000000..17f78ce --- /dev/null +++ b/crates/burn-no-std-tests/src/mlp.rs @@ -0,0 +1,67 @@ +// Originally copied from the burn/examples/mnist package + +use alloc::vec::Vec; + +use burn::{ + config::Config, + module::Module, + nn, + tensor::{Device, Tensor}, +}; + +/// Configuration to create a [Multilayer Perceptron](Mlp) layer. +#[derive(Config, Debug)] +pub struct MlpConfig { + /// The number of layers. + #[config(default = 3)] + pub num_layers: usize, + /// The dropout rate. + #[config(default = 0.5)] + pub dropout: f64, + /// The size of each layer. + #[config(default = 256)] + pub d_model: usize, +} + +/// Multilayer Perceptron module. +#[derive(Module, Debug)] +pub struct Mlp { + linears: Vec, + dropout: nn::Dropout, + activation: nn::Relu, +} + +impl Mlp { + /// Create the module from the given configuration. + pub fn new(config: &MlpConfig, device: &Device) -> Self { + let mut linears = Vec::with_capacity(config.num_layers); + + for _ in 0..config.num_layers { + linears.push(nn::LinearConfig::new(config.d_model, config.d_model).init(device)); + } + + Self { + linears, + dropout: nn::DropoutConfig::new(0.3).init(), + activation: nn::Relu::new(), + } + } + + /// Applies the forward pass on the input tensor. + /// + /// # Shapes + /// + /// - input: `[batch_size, d_model]` + /// - output: `[batch_size, d_model]` + pub fn forward(&self, input: Tensor<2>) -> Tensor<2> { + let mut x = input; + + for linear in self.linears.iter() { + x = linear.forward(x); + x = self.dropout.forward(x); + x = self.activation.forward(x); + } + + x + } +} diff --git a/crates/burn-no-std-tests/src/model.rs b/crates/burn-no-std-tests/src/model.rs new file mode 100644 index 0000000..ecf3b78 --- /dev/null +++ b/crates/burn-no-std-tests/src/model.rs @@ -0,0 +1,66 @@ +// Originally copied from the burn/examples/mnist package + +use crate::{ + conv::{ConvBlock, ConvBlockConfig}, + mlp::{Mlp, MlpConfig}, +}; + +use burn::{ + config::Config, + module::Module, + nn, + tensor::{Device, Tensor}, +}; + +#[derive(Config, Debug)] +pub struct MnistConfig { + #[config(default = 42)] + pub seed: u64, + + pub mlp: MlpConfig, + + #[config(default = 784)] + pub input_size: usize, + + #[config(default = 10)] + pub output_size: usize, +} + +#[derive(Module, Debug)] +pub struct Model { + mlp: Mlp, + conv: ConvBlock, + input: nn::Linear, + output: nn::Linear, + num_classes: usize, +} + +impl Model { + pub fn new(config: &MnistConfig, device: &Device) -> Self { + let mlp = Mlp::new(&config.mlp, device); + let input = nn::LinearConfig::new(config.input_size, config.mlp.d_model).init(device); + let output = nn::LinearConfig::new(config.mlp.d_model, config.output_size).init(device); + let conv = ConvBlock::new(&ConvBlockConfig::new([1, 1]), device); + + Self { + mlp, + conv, + output, + input, + num_classes: config.output_size, + } + } + + pub fn forward(&self, input: Tensor<3>) -> Tensor<2> { + let [batch_size, height, width] = input.dims(); + + let x = input.reshape([batch_size, 1, height, width]).detach(); + let x = self.conv.forward(x); + let x = x.reshape([batch_size, height * width]); + + let x = self.input.forward(x); + let x = self.mlp.forward(x); + + self.output.forward(x) + } +} diff --git a/crates/burn-no-std-tests/src/safetensors.rs b/crates/burn-no-std-tests/src/safetensors.rs new file mode 100644 index 0000000..06bc554 --- /dev/null +++ b/crates/burn-no-std-tests/src/safetensors.rs @@ -0,0 +1,111 @@ +// Test SafeTensors storage in no-std environment + +use burn::{ + module::Module, + nn, + tensor::{Device, Tensor}, +}; + +use burn_store::{ModuleSnapshot, SafetensorsStore}; + +/// Simple model for testing SafeTensors storage +#[derive(Module, Debug)] +pub struct TestModel { + linear1: nn::Linear, + linear2: nn::Linear, +} + +impl TestModel { + pub fn new(device: &Device) -> Self { + Self { + linear1: nn::LinearConfig::new(10, 20).init(device), + linear2: nn::LinearConfig::new(20, 10).init(device), + } + } + + pub fn forward(&self, x: Tensor<2>) -> Tensor<2> { + let x = self.linear1.forward(x); + self.linear2.forward(x) + } +} + +/// Test basic SafeTensors save and load in no-std +pub fn test_safetensors_basic(device: &Device) { + // Create a model + let model = TestModel::new(device); + + // Save to bytes (no file I/O in no-std) + let mut save_store = SafetensorsStore::from_bytes(None); + model + .save_into(&mut save_store) + .expect("Failed to save model"); + + // Get the serialized bytes + let bytes = save_store.get_bytes().expect("Failed to get bytes"); + + // Load from bytes + let mut load_store = SafetensorsStore::from_bytes(Some(bytes)); + let mut loaded_model = TestModel::new(device); + loaded_model + .load_from(&mut load_store) + .expect("Failed to load model"); + + // Test that the model still works + let input = Tensor::<2>::ones([2, 10], device); + let _output = loaded_model.forward(input); +} + +/// Test SafeTensors with filtering in no-std +pub fn test_safetensors_filtering(device: &Device) { + let model = TestModel::new(device); + + // Save only linear1 weights + let mut save_store = SafetensorsStore::from_bytes(None) + .with_full_path("linear1.weight") + .with_full_path("linear1.bias"); + model + .save_into(&mut save_store) + .expect("Failed to save filtered model"); + + let bytes = save_store.get_bytes().expect("Failed to get bytes"); + + // Load with partial loading allowed + let mut load_store = SafetensorsStore::from_bytes(Some(bytes)).allow_partial(true); + let mut partial_model = TestModel::new(device); + let result = partial_model + .load_from(&mut load_store) + .expect("Failed to load partial model"); + + // Verify that only linear1 was loaded + assert_eq!(result.applied.len(), 2, "Should have loaded 2 tensors"); + assert!(!result.missing.is_empty(), "Should have missing tensors"); +} + +/// Test SafeTensors with metadata in no-std +pub fn test_safetensors_metadata(device: &Device) { + let model = TestModel::new(device); + + // Save with metadata + let mut save_store = SafetensorsStore::from_bytes(None) + .metadata("version", "1.0.0") + .metadata("environment", "no-std"); + model + .save_into(&mut save_store) + .expect("Failed to save model with metadata"); + + let bytes = save_store.get_bytes().expect("Failed to get bytes"); + + // Load and verify it works + let mut load_store = SafetensorsStore::from_bytes(Some(bytes)); + let mut loaded_model = TestModel::new(device); + loaded_model + .load_from(&mut load_store) + .expect("Failed to load model with metadata"); +} + +/// Run all SafeTensors no-std tests +pub fn run_all_tests(device: &Device) { + test_safetensors_basic(device); + test_safetensors_filtering(device); + test_safetensors_metadata(device); +} diff --git a/crates/burn-no-std-tests/tests/burnpack_tests.rs b/crates/burn-no-std-tests/tests/burnpack_tests.rs new file mode 100644 index 0000000..3e174f2 --- /dev/null +++ b/crates/burn-no-std-tests/tests/burnpack_tests.rs @@ -0,0 +1,10 @@ +extern crate alloc; + +#[test] +fn test_burnpack_no_std() { + use burn_no_std_tests::burnpack; + let device = Default::default(); + + // Run all Burnpack tests + burnpack::run_all_tests(&device); +} diff --git a/crates/burn-no-std-tests/tests/safetensors_tests.rs b/crates/burn-no-std-tests/tests/safetensors_tests.rs new file mode 100644 index 0000000..3f3484c --- /dev/null +++ b/crates/burn-no-std-tests/tests/safetensors_tests.rs @@ -0,0 +1,10 @@ +extern crate alloc; + +#[test] +fn test_safetensors_no_std() { + use burn_no_std_tests::safetensors; + let device = Default::default(); + + // Run all SafeTensors tests + safetensors::run_all_tests(&device); +} diff --git a/crates/burn-no-std-tests/tests/test_integration.rs b/crates/burn-no-std-tests/tests/test_integration.rs new file mode 100644 index 0000000..9839a75 --- /dev/null +++ b/crates/burn-no-std-tests/tests/test_integration.rs @@ -0,0 +1,28 @@ +#![no_std] // Must keep it for testing + +use burn_no_std_tests::mlp::*; +use burn_no_std_tests::model::*; + +use burn::tensor::{Distribution, Tensor}; + +#[test] +fn test_mnist_model_with_random_input() { + // Model configurations + let device = Default::default(); + let mlp_config = MlpConfig::new(); + let mnist_config = MnistConfig::new(mlp_config); + let mnist_model = Model::new(&mnist_config, &device); + + // Pass a fixed seed for random, otherwise a build generated random seed is used + device.seed(mnist_config.seed); + + // Some random input + let input_shape = [1, 28, 28]; + let input = Tensor::<3>::random(input_shape, Distribution::Default, &device); + + // Run through the model + let output = mnist_model.forward(input); + + assert_eq!(&*output.shape(), [1, 10]); + assert!(output.to_data().iter::().all(|x| x <= 1.0)); +} diff --git a/crates/burn-optim/Cargo.toml b/crates/burn-optim/Cargo.toml new file mode 100644 index 0000000..68bdf71 --- /dev/null +++ b/crates/burn-optim/Cargo.toml @@ -0,0 +1,52 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science", "no-std", "embedded", "wasm"] +description = "Optimizer building blocks for the Burn deep learning framework" +documentation = "https://docs.rs/burn-optim" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "tensor", "pytorch", "ndarray"] +license.workspace = true +name = "burn-optim" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-optim" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["std", "burn-core/default"] +doc = [ + "std", + # Doc features + "burn-core/doc", +] +std = ["burn-core/std", "burn-pack/std", "num-traits/std", "serde/std", "log"] +tracing = ["burn-core/tracing"] + +# collective = ["burn-collective"] + +[dependencies] + +# ** Please make sure all dependencies support no_std when std is disabled ** +burn-core = { workspace = true, features = ["autodiff", "extension"] } +burn-derive = { workspace = true } +burn-pack = { workspace = true, default-features = false } + +num-traits = { workspace = true } +derive-new = { workspace = true } +log = { workspace = true, optional = true } +serde = { workspace = true, features = ["derive"] } + +# The same implementation of HashMap in std but with no_std support (only alloc crate is needed) +hashbrown = { workspace = true, features = ["serde"] } # no_std compatible + +[dev-dependencies] +burn-nn = { workspace = true, features = ["default"] } +burn-flex = { workspace = true, features = ["default"] } +burn-autodiff = { workspace = true, features = ["default"] } +rstest = { workspace = true } + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-optim/README.md b/crates/burn-optim/README.md new file mode 100644 index 0000000..84e45eb --- /dev/null +++ b/crates/burn-optim/README.md @@ -0,0 +1,3 @@ +# Burn Optimizers + +Core building blocks for Burn optimizers. \ No newline at end of file diff --git a/crates/burn-optim/src/grad_clipping/base.rs b/crates/burn-optim/src/grad_clipping/base.rs new file mode 100644 index 0000000..e76fbd5 --- /dev/null +++ b/crates/burn-optim/src/grad_clipping/base.rs @@ -0,0 +1,139 @@ +use burn_core as burn; + +use burn::{config::Config, tensor::Tensor}; + +/// Gradient Clipping provides a way to mitigate exploding gradients +#[derive(Config, Debug)] +pub enum GradientClippingConfig { + /// Clip the gradient by value. + Value(f32), + + /// Clip the gradient by norm. + Norm(f32), +} + +impl GradientClippingConfig { + /// Initialize the gradient clipping. + /// + /// # Returns + /// + /// The gradient clipping. + pub fn init(&self) -> GradientClipping { + match self { + GradientClippingConfig::Value(val) => GradientClipping::Value(*val), + GradientClippingConfig::Norm(val) => GradientClipping::Norm(*val), + } + } +} + +/// Gradient Clipping provides a way to mitigate exploding gradients +/// by clipping every component of the gradient by value or by norm during +/// backpropagation. +#[derive(Clone)] +pub enum GradientClipping { + /// Clip the gradient by value. + Value(f32), + + /// Clip the gradient by norm. + Norm(f32), +} + +impl GradientClipping { + /// Clip the gradient. + /// + /// # Arguments + /// + /// * `grad` - The gradient to clip. + /// + /// # Returns + /// + /// The clipped gradient. + pub fn clip_gradient(&self, grad: Tensor) -> Tensor { + match self { + GradientClipping::Value(threshold) => self.clip_by_value(grad, *threshold), + GradientClipping::Norm(max_norm) => self.clip_by_norm(grad, *max_norm), + } + } + + fn clip_by_value(&self, grad: Tensor, threshold: f32) -> Tensor { + let greater_mask = grad.clone().greater_elem(threshold); + let lower_mask = grad.clone().lower_elem(-threshold); + + let clipped_grad = grad.mask_fill(greater_mask, threshold); + + clipped_grad.mask_fill(lower_mask, -threshold) + } + + fn clip_by_norm(&self, grad: Tensor, threshold: f32) -> Tensor { + let norm = Self::l2_norm(grad.clone()); + let min_positive = grad + .dtype() + .finfo() + .unwrap_or(burn::tensor::FloatDType::F32.finfo()) + .min_positive; + let clip_coef = threshold / norm.add_scalar(min_positive); + let clip_coef_clamped = clip_coef.clamp_max(1.0); + grad.mul(clip_coef_clamped.unsqueeze()) + } + + fn l2_norm(tensor: Tensor) -> Tensor<1> { + let squared = tensor.square(); + let sum = squared.sum(); + sum.sqrt() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Tensor; + + #[test] + fn test_clip_by_value() { + let gradient: Tensor<2> = Tensor::from_floats( + [ + [0.6294, 0.0940, 0.8176, 0.8824, 0.5228, 0.4310], + [0.7152, 0.9559, 0.7893, 0.5684, 0.5939, 0.8883], + ], + &Default::default(), + ); + + let clipped_gradient = GradientClipping::Value(0.5).clip_gradient(gradient); + let clipped_gradient_data = clipped_gradient.into_data(); + + for value in clipped_gradient_data.iter::() { + assert!(value <= 0.5); + } + } + + #[test] + fn test_clip_by_norm() { + let gradient: Tensor<2> = Tensor::from_floats( + [ + [0.6294, 0.0940, 0.8176, 0.8824, 0.5228, 0.4310], + [0.7152, 0.9559, 0.7893, 0.5684, 0.5939, 0.8883], + ], + &Default::default(), + ); + + let clipped_gradient = GradientClipping::Norm(2.2).clip_gradient(gradient); + let clipped_gradient_data = clipped_gradient.into_data(); + + for value in clipped_gradient_data.iter::() { + assert!(value <= 0.88); + } + } + #[test] + fn test_clip_by_norm_no_clipping() { + let gradient: Tensor<2> = Tensor::from_floats( + [[0.3, 0.4, 0.5, 0.2], [0.1, 0.6, 0.3, 0.4]], + &Default::default(), + ); + + let clipped_gradient = GradientClipping::Norm(2.2).clip_gradient(gradient.clone()); + + clipped_gradient + .into_data() + .assert_eq(&gradient.into_data(), true); + } +} diff --git a/crates/burn-optim/src/grad_clipping/mod.rs b/crates/burn-optim/src/grad_clipping/mod.rs new file mode 100644 index 0000000..096c94e --- /dev/null +++ b/crates/burn-optim/src/grad_clipping/mod.rs @@ -0,0 +1,2 @@ +mod base; +pub use base::*; diff --git a/crates/burn-optim/src/lib.rs b/crates/burn-optim/src/lib.rs new file mode 100644 index 0000000..9a24cb3 --- /dev/null +++ b/crates/burn-optim/src/lib.rs @@ -0,0 +1,28 @@ +#![cfg_attr(not(feature = "std"), no_std)] +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] +#![recursion_limit = "256"] + +//! Burn optimizers. + +#[macro_use] +extern crate derive_new; + +extern crate alloc; + +/// Optimizer module. +mod optim; +pub use optim::*; + +/// Gradient clipping module. +pub mod grad_clipping; + +/// Learning rate scheduler module. +#[cfg(feature = "std")] +pub mod lr_scheduler; + +/// Type alias for the learning rate. +/// +/// LearningRate also implements [learning rate scheduler](crate::lr_scheduler::LrScheduler) so it +/// can be used for constant learning rate. +pub type LearningRate = f64; // We could potentially change the type. diff --git a/crates/burn-optim/src/lr_scheduler/base.rs b/crates/burn-optim/src/lr_scheduler/base.rs new file mode 100644 index 0000000..b2a9152 --- /dev/null +++ b/crates/burn-optim/src/lr_scheduler/base.rs @@ -0,0 +1,312 @@ +use std::fmt::Debug; + +use alloc::collections::BTreeMap; +pub(super) use alloc::string::String; +use alloc::vec::Vec; +use burn_core as burn; +use burn_core::config::Config; + +use crate::lr_scheduler::composed::ComposedLrSchedulerConfig; +use crate::lr_scheduler::cosine::CosineAnnealingLrSchedulerConfig; +use crate::lr_scheduler::exponential::ExponentialLrSchedulerConfig; +use crate::lr_scheduler::linear::LinearLrSchedulerConfig; +use crate::lr_scheduler::noam::NoamLrSchedulerConfig; +use crate::lr_scheduler::step::StepLrSchedulerConfig; +use crate::{RecordState, StateSink, StateSource, join_path}; +use burn::store::RecordError; +use burn::tensor::{Bytes, Device}; +use burn_pack::{Reader, Scalar, Writer}; + +use crate::LearningRate; + +macro_rules! impl_from_for_scheduler { + ($($variant:ident($config:ident)),* $(,)?) => { + $( + impl From<$config> for LrSchedulerConfig { + fn from(config: $config) -> Self { + LrSchedulerConfig::$variant(config) + } + } + )* + }; +} + +/// Learning rate scheduler defines how the learning rate will evolve during training. +pub trait LrScheduler: LrSchedulerClone + Send + Sync { + /// Perform the scheduler step, potentially updating its state, and returning the effective + /// learning rate. + fn step(&mut self) -> LearningRate; + + /// Get the current state of the scheduler as a [record](LrSchedulerRecord). + fn to_record(&self) -> LrSchedulerRecord; + + /// Load the state of the scheduler from a [record](LrSchedulerRecord). + fn load_record(&mut self, record: LrSchedulerRecord); +} + +/// Implements the clone of a boxed [`LrScheduler`]. +pub trait LrSchedulerClone { + /// Clones a boxed [`LrScheduler`]. + fn clone_box(&self) -> Box; +} + +impl LrSchedulerClone for T +where + T: 'static + LrScheduler + Clone, +{ + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } +} + +impl Clone for Box { + fn clone(&self) -> Box { + self.as_ref().clone_box() + } +} + +/// A wrapper over a dynamic [`LrScheduler`]. +#[derive(Clone)] +pub struct DynLrScheduler { + scheduler: Box, +} + +impl DynLrScheduler { + /// Perform the scheduler step, potentially updating its state, and returning the effective + /// learning rate. + pub fn step(&mut self) -> LearningRate { + self.scheduler.step() + } + + /// Get the current state of the scheduler as a [record](LrSchedulerRecord). + pub fn to_record(&self) -> LrSchedulerRecord { + self.scheduler.to_record() + } + + /// Load the state of the scheduler from a [record](LrSchedulerRecord). + pub fn load_record(mut self, record: LrSchedulerRecord) -> Self { + self.scheduler.load_record(record); + self + } +} + +impl From for DynLrScheduler +where + S: LrScheduler + 'static, +{ + fn from(scheduler: S) -> Self { + Self { + scheduler: Box::new(scheduler), + } + } +} + +/// The serialized state of a [learning rate scheduler](LrScheduler), stored in the +/// [burnpack](burn_pack) format. +/// +/// Scheduler state is just a handful of scalars (step counters, current learning rate), so the +/// record holds named typed scalars and no tensors. Composed schedulers nest their children's +/// records under an index prefix via [`with_record`](Self::with_record) / [`record`](Self::record). +#[derive(Default, Clone, Debug)] +pub struct LrSchedulerRecord { + scalars: BTreeMap, +} + +impl LrSchedulerRecord { + /// Create an empty record. + pub fn new() -> Self { + Self::default() + } + + /// Whether the record holds no scalars. + pub fn is_empty(&self) -> bool { + self.scalars.is_empty() + } + + /// Store a scalar under `key`. + pub fn with_scalar>(mut self, key: &str, value: V) -> Self { + self.scalars.insert(String::from(key), value.into()); + self + } + + /// Read the scalar stored under `key`, if present and of a compatible type. + pub fn scalar>(&self, key: &str) -> Option { + self.scalars + .get(key) + .copied() + .and_then(|scalar| V::try_from(scalar).ok()) + } + + /// Merge a child `record`'s scalars under `prefix` (used to compose schedulers). + pub fn with_record(mut self, prefix: &str, record: LrSchedulerRecord) -> Self { + for (key, value) in record.scalars { + self.scalars.insert(join_path(prefix, &key), value); + } + self + } + + /// Extract the child record previously merged under `prefix`. + pub fn record(&self, prefix: &str) -> LrSchedulerRecord { + let head = join_path(prefix, ""); + let scalars = self + .scalars + .iter() + .filter_map(|(key, value)| { + key.strip_prefix(&head) + .map(|stripped| (String::from(stripped), *value)) + }) + .collect(); + LrSchedulerRecord { scalars } + } + + /// Build a record from a scheduler [state](RecordState). + /// + /// Scheduler state is scalar-only, so this reuses the same [`RecordState`] decomposition as + /// optimizer states (it panics in debug builds if a tensor leaf is produced). + pub fn from_state(state: &S) -> Self { + let mut sink = StateSink::default(); + state.state_flatten("", &mut sink); + debug_assert!( + sink.tensors.is_empty(), + "learning rate scheduler state is expected to be scalar-only" + ); + Self { + scalars: sink.scalars.into_iter().collect(), + } + } + + /// Reconstruct a scheduler [state](RecordState) from this record. + /// + /// Uses the default device; scheduler state is scalar-only so no tensor is ever allocated. + pub fn into_state(&self) -> Option { + let mut source = StateSource::new(self.scalars.clone()); + S::state_unflatten("", &mut source, &Device::default()) + } + + /// Serialize the record to an in-memory burnpack byte buffer. + pub fn into_bytes(self) -> Result { + Ok(self.into_writer().into_bytes()?) + } + + /// Reconstruct a record from an in-memory burnpack byte buffer. + pub fn from_bytes(bytes: Bytes) -> Result { + let reader = Reader::from_bytes(bytes)?; + Ok(Self { + scalars: reader.scalars().clone(), + }) + } + + /// Save the record to a burnpack file on disk. + #[cfg(feature = "std")] + pub fn save>(self, path: P) -> Result<(), RecordError> { + self.into_writer().write_to_file(path)?; + Ok(()) + } + + /// Load the record from a burnpack file on disk. + #[cfg(feature = "std")] + pub fn load>(path: P) -> Result { + let reader = Reader::from_file(path)?; + Ok(Self { + scalars: reader.scalars().clone(), + }) + } + + fn into_writer(self) -> Writer { + let mut writer = Writer::new(Vec::new()); + for (key, value) in &self.scalars { + writer = writer.with_scalar(key, *value); + } + writer + } +} + +/// An enum for possible learning rate scheduler configs. +#[derive(Config, Debug)] +pub enum LrSchedulerConfig { + /// A constant learning rate. + Constant(LearningRate), + /// A [`LinearLrSchedulerConfig`] + Linear(LinearLrSchedulerConfig), + /// A [`CosineAnnealingLrSchedulerConfig`] + Cosine(CosineAnnealingLrSchedulerConfig), + /// A [`ExponentialLrSchedulerConfig`] + Exponential(ExponentialLrSchedulerConfig), + /// A [`NoamLrSchedulerConfig`] + Noam(NoamLrSchedulerConfig), + /// A [`StepLrSchedulerConfig`] + Step(StepLrSchedulerConfig), + /// A [`ComposedLrSchedulerConfig`] + Composed(ComposedLrSchedulerConfig), +} + +impl_from_for_scheduler!( + Constant(LearningRate), + Linear(LinearLrSchedulerConfig), + Cosine(CosineAnnealingLrSchedulerConfig), + Exponential(ExponentialLrSchedulerConfig), + Noam(NoamLrSchedulerConfig), + Step(StepLrSchedulerConfig), + Composed(ComposedLrSchedulerConfig), +); + +#[cfg(test)] +pub(super) mod test_utils { + use super::*; + + // A small tolerance for learning rate comparisons. Depending on how learning rates are + // computed, floating-point arithmetic error might exceed f64::EPSILON, so a larger value is + // used here. + const LOOSE_EPSILON: LearningRate = 1e-10; + + pub fn check_lr_sequence(mut scheduler: S, expected_lrs: I) + where + I: IntoIterator, + S: LrScheduler, + { + expected_lrs + .into_iter() + .enumerate() + .for_each(|(i, expected)| { + let lr = scheduler.step(); + assert!( + (lr - expected).abs() < LOOSE_EPSILON, + "Scheduled learning rate {lr} is not approximately equal to the expected value \ + {expected} at step {i}", + ); + }); + } + + // save_at_step is the number of steps to run the scheduler before saving and loading back its + // state. + pub fn check_save_load(mut scheduler: S, save_at_step: usize) + where + S: Clone + LrScheduler, + { + let mut truth = scheduler.clone(); + // Consume some steps before saving and loading back + (0..save_at_step).for_each(|_| { + truth.step(); + scheduler.step(); + }); + let rec = scheduler.to_record(); + scheduler.load_record(rec); + + // Validate that the scheduler resumes from where it left off. + compare_steps(&mut scheduler, &mut truth, save_at_step); + } + + // Check if two schedulers produce the same learning rate sequences over the specified number of + // steps. + pub fn compare_steps(a: &mut S, b: &mut S, num_steps: usize) { + (0..num_steps).for_each(|i| { + let lr_a = a.step(); + let lr_b = b.step(); + assert!( + (lr_a - lr_b).abs() < LOOSE_EPSILON, + "The two learning rates ({lr_a}, {lr_b}) at position {i} in the remaining \ + sequences are not approximately equal", + ); + }); + } +} diff --git a/crates/burn-optim/src/lr_scheduler/composed.rs b/crates/burn-optim/src/lr_scheduler/composed.rs new file mode 100644 index 0000000..b843e79 --- /dev/null +++ b/crates/burn-optim/src/lr_scheduler/composed.rs @@ -0,0 +1,162 @@ +use burn_core::{self as burn}; + +use super::cosine::CosineAnnealingLrSchedulerConfig; +use super::exponential::ExponentialLrSchedulerConfig; +use super::linear::LinearLrSchedulerConfig; +use super::noam::NoamLrSchedulerConfig; +use super::{LrScheduler, LrSchedulerRecord, String}; +use crate::LearningRate; +use crate::lr_scheduler::module_lr_scheduler::ModuleLrScheduler; +use crate::lr_scheduler::step::StepLrSchedulerConfig; +use crate::lr_scheduler::{DynLrScheduler, LrSchedulerConfig}; + +use burn::config::Config; + +/// Compose multiple [learning rate schedulers](LrScheduler) together. +#[derive(Config, Debug)] +pub struct ComposedLrSchedulerConfig { + #[config(default = "Vec::new()")] + schedulers: Vec, + #[config(default = "SchedulerReduction::Prod")] + reduction: SchedulerReduction, +} + +/// Compose multiple [learning rate schedulers](LrScheduler) together. +#[derive(Clone)] +pub struct ComposedLrScheduler { + schedulers: Vec, + reduction: SchedulerReduction, +} + +/// Defines how the learning rates generated by the schedulers are combined. +#[derive(Config, Debug, Copy)] +pub enum SchedulerReduction { + /// All learning rates are averaged. + Avg, + /// All learning rates are summed. + Sum, + /// All learning rates are multiplied. + Prod, +} + +impl ComposedLrSchedulerConfig { + /// Initialize the learning rate scheduler. + pub(crate) fn build(&self) -> Result { + let mut schedulers: Vec = Vec::with_capacity(self.schedulers.len()); + for config in self.schedulers.iter() { + let config = match config { + LrSchedulerConfig::Constant(lr) => (*lr).into(), + LrSchedulerConfig::Linear(config) => config.build()?.into(), + LrSchedulerConfig::Cosine(config) => config.build()?.into(), + LrSchedulerConfig::Exponential(config) => config.build()?.into(), + LrSchedulerConfig::Noam(config) => config.build()?.into(), + LrSchedulerConfig::Step(config) => config.build()?.into(), + LrSchedulerConfig::Composed(config) => config.build()?.into(), + }; + schedulers.push(config); + } + + Ok(ComposedLrScheduler { + schedulers, + reduction: self.reduction, + }) + } + + /// Initializes a [module learning rate scheduler](ModuleLrScheduler). + pub fn init(&self) -> Result { + self.build().map(|s| s.into()) + } + + /// Appends a [constant learning rate](crate::lr_scheduler::constant::ConstantLr). + pub fn constant(mut self, lr: LearningRate) -> Self { + self.schedulers.push(LrSchedulerConfig::Constant(lr)); + self + } + + /// Appends a [linear scheduler](crate::lr_scheduler::linear::LinearLrScheduler). + pub fn linear(mut self, config: LinearLrSchedulerConfig) -> Self { + self.schedulers.push(LrSchedulerConfig::Linear(config)); + self + } + + /// Appends a [cosine scheduler](ComposedLrSchedulerConfig). + pub fn cosine(mut self, config: CosineAnnealingLrSchedulerConfig) -> Self { + self.schedulers.push(LrSchedulerConfig::Cosine(config)); + self + } + + /// Appends an [exponential scheduler](crate::lr_scheduler::exponential::ExponentialLrScheduler). + pub fn exponential(mut self, config: ExponentialLrSchedulerConfig) -> Self { + self.schedulers.push(LrSchedulerConfig::Exponential(config)); + self + } + + /// Appends a [noam scheduler](crate::lr_scheduler::noam::NoamLrScheduler). + pub fn noam(mut self, config: NoamLrSchedulerConfig) -> Self { + self.schedulers.push(LrSchedulerConfig::Noam(config)); + self + } + + /// Appends a [step scheduler](crate::lr_scheduler::step::StepLrScheduler). + pub fn step(mut self, config: StepLrSchedulerConfig) -> Self { + self.schedulers.push(LrSchedulerConfig::Step(config)); + self + } + + /// Appends a [composed scheduler](ComposedLrScheduler). + pub fn composed(mut self, config: Self) -> Self { + self.schedulers.push(LrSchedulerConfig::Composed(config)); + self + } +} + +impl ComposedLrScheduler { + /// Add a custom learning rate scheduler to existing ones. + pub fn with_custom_scheduler(mut self, scheduler: S) -> Self { + self.schedulers.push(scheduler.into()); + self + } +} + +impl LrScheduler for ComposedLrScheduler { + fn step(&mut self) -> LearningRate { + let mut step = match self.reduction { + SchedulerReduction::Avg => 0.0, + SchedulerReduction::Sum => 0.0, + SchedulerReduction::Prod => 1.0, + }; + let num_scheduler = self.schedulers.len() as f64; + + for lr in self.schedulers.iter_mut().map(|s| s.step()) { + step = match self.reduction { + SchedulerReduction::Avg => step + (lr / num_scheduler), + SchedulerReduction::Sum => step + lr, + SchedulerReduction::Prod => step * lr, + } + } + + step + } + + fn to_record(&self) -> LrSchedulerRecord { + let mut record = LrSchedulerRecord::new(); + for (index, item) in self.schedulers.iter().enumerate() { + let sub = item.to_record(); + record = record.with_record(&index.to_string(), sub); + } + record + } + + fn load_record(&mut self, record: LrSchedulerRecord) { + self.schedulers = self + .schedulers + .clone() + .iter_mut() + .enumerate() + .map(|(index, item)| { + let sub = record.record(&index.to_string()); + item.clone().load_record(sub) + }) + .collect(); + } +} diff --git a/crates/burn-optim/src/lr_scheduler/constant.rs b/crates/burn-optim/src/lr_scheduler/constant.rs new file mode 100644 index 0000000..3e87836 --- /dev/null +++ b/crates/burn-optim/src/lr_scheduler/constant.rs @@ -0,0 +1,42 @@ +use super::{LrScheduler, LrSchedulerRecord}; +use crate::LearningRate; + +/// Constant learning rate implementing [learning rate scheduler](LrScheduler). +/// +/// # Notes +/// +/// You can also use [learning rate](LearningRate) which the same effect. +#[derive(new, Clone, Debug)] +pub struct ConstantLr { + lr: LearningRate, +} + +impl From for ConstantLr { + fn from(lr: LearningRate) -> Self { + Self { lr } + } +} + +impl LrScheduler for ConstantLr { + fn step(&mut self) -> LearningRate { + self.lr + } + + fn to_record(&self) -> LrSchedulerRecord { + LrSchedulerRecord::new() + } + + fn load_record(&mut self, _record: LrSchedulerRecord) {} +} + +impl LrScheduler for LearningRate { + fn step(&mut self) -> LearningRate { + *self + } + + fn to_record(&self) -> LrSchedulerRecord { + LrSchedulerRecord::new() + } + + fn load_record(&mut self, _record: LrSchedulerRecord) {} +} diff --git a/crates/burn-optim/src/lr_scheduler/cosine.rs b/crates/burn-optim/src/lr_scheduler/cosine.rs new file mode 100644 index 0000000..76c9b00 --- /dev/null +++ b/crates/burn-optim/src/lr_scheduler/cosine.rs @@ -0,0 +1,204 @@ +use burn_core as burn; + +use super::{LrScheduler, LrSchedulerRecord, String}; +use crate::LearningRate; +use crate::RecordState; +use crate::lr_scheduler::module_lr_scheduler::ModuleLrScheduler; +use burn::config::Config; + +/// The configuration for creating a [Cosine Annealing learning rate scheduler with warm +/// restarts](CosineAnnealingLrScheduler). +/// +/// This scheduler returns the learning rate `initial_lr` at the first step, then changes it by +/// following a cosine function. After `num_iters` iterations, the learning rate is reset to +/// `initial_lr`. +#[derive(Config, Debug)] +pub struct CosineAnnealingLrSchedulerConfig { + // The initial learning rate. + initial_lr: LearningRate, + // The final learning rate. + #[config(default = 0.0)] + min_lr: LearningRate, + // The number of iterations between two restarts. The two restart iterations themselves are not + // included. + num_iters: usize, +} + +impl CosineAnnealingLrSchedulerConfig { + /// Initializes a [Cosine learning rate scheduler](CosineAnnealingLrScheduler). + pub(crate) fn build(&self) -> Result { + if self.initial_lr <= 0. || self.initial_lr > 1. { + return Err("Initial learning rate must be greater than 0 and at most 1".into()); + } + if self.min_lr < 0.0 || self.min_lr > self.initial_lr { + return Err( + "Minimum learning rate must be at least 0 and at most equal to the initial \ + learning rate" + .into(), + ); + } + if self.num_iters == 0 { + return Err("Number of iterations must be at least 1".into()); + } + + Ok(CosineAnnealingLrScheduler { + min_lr: self.min_lr, + max_lr: self.initial_lr, + num_iters: self.num_iters, + current_iter: usize::MAX, + }) + } + + /// Initializes a [module learning rate scheduler](ModuleLrScheduler). + /// + /// # Errors + /// + /// An error will be returned if any of the following conditions is true: + /// + /// * `initial_lr` is out of range (0.0, 1.0] + /// * `min_lr` is out of range [0.0, `initial_lr`] + /// * `num_iters` is 0 + pub fn init(&self) -> Result { + self.build().map(|s| s.into()) + } +} + +/// A Cosine Annealing learning rate scheduler. +/// +/// This scheduler is described in [SGDR: Stochastic Gradient Descent with Warm +/// Restarts](https://arxiv.org/abs/1608.03983). See [CosineAnnealingLrSchedulerConfig] for more +/// information. +#[derive(Clone, Copy, Debug)] +pub struct CosineAnnealingLrScheduler { + min_lr: LearningRate, + max_lr: LearningRate, + num_iters: usize, + current_iter: usize, +} + +impl LrScheduler for CosineAnnealingLrScheduler { + fn step(&mut self) -> LearningRate { + // Make current_iter overflow from usize::MAX to 0 to get the initial learning rate on the + // first call. We could've used i64 with an initial value -1, but keeping it in usize saves + // us from some type casting here. + self.current_iter = self.current_iter.wrapping_add(1) % (self.num_iters + 1); + self.min_lr + + 0.5 + * (self.max_lr - self.min_lr) + * (1.0 + + (self.current_iter as f64 / self.num_iters as f64 * std::f64::consts::PI) + .cos()) + } + + fn to_record(&self) -> LrSchedulerRecord { + LrSchedulerRecord::from_state(&CosineAnnealingLrSchedulerState { + current_iter: self.current_iter, + }) + } + + fn load_record(&mut self, record: LrSchedulerRecord) { + if let Some(state) = record.into_state::() { + self.current_iter = state.current_iter; + } + } +} + +/// The serializable state of a [cosine annealing scheduler](CosineAnnealingLrScheduler). +#[derive(RecordState, Clone, Debug)] +pub struct CosineAnnealingLrSchedulerState { + current_iter: usize, +} + +#[cfg(test)] +mod tests { + use super::super::test_utils; + use super::*; + + #[test] + fn config_initial_lr_too_low() { + let r = CosineAnnealingLrSchedulerConfig::new(0., 10).build(); + assert!(r.is_err(), "Should return an error"); + assert_eq!( + r.unwrap_err(), + "Initial learning rate must be greater than 0 and at most 1", + "Error messages should match", + ); + } + + #[test] + fn config_initial_lr_too_high() { + let r = CosineAnnealingLrSchedulerConfig::new(1.5, 10).build(); + assert!(r.is_err(), "Should return an error"); + assert_eq!( + r.unwrap_err(), + "Initial learning rate must be greater than 0 and at most 1", + "Error messages should match", + ); + } + + #[test] + fn config_min_lr_too_low() { + let r = CosineAnnealingLrSchedulerConfig::new(0.5, 10) + .with_min_lr(-0.1) + .build(); + assert!(r.is_err(), "Should return an error"); + assert_eq!( + r.unwrap_err(), + "Minimum learning rate must be at least 0 and at most equal to the initial learning \ + rate", + "Error messages should match", + ); + } + + #[test] + fn config_min_lr_too_high() { + let r = CosineAnnealingLrSchedulerConfig::new(0.5, 10) + .with_min_lr(0.6) + .build(); + assert!(r.is_err(), "Should return an error"); + assert_eq!( + r.unwrap_err(), + "Minimum learning rate must be at least 0 and at most equal to the initial learning \ + rate", + "Error messages should match", + ); + } + + #[test] + fn config_num_iters_too_low() { + let r = CosineAnnealingLrSchedulerConfig::new(0.5, 0).build(); + assert!(r.is_err(), "Should return an error"); + assert_eq!( + r.unwrap_err(), + "Number of iterations must be at least 1", + "Error messages should match", + ); + } + + #[test] + fn test_lr_change() { + const INITIAL_LR: LearningRate = 0.5; + const MIN_LR: LearningRate = 0.1; + + let scheduler = CosineAnnealingLrSchedulerConfig::new(INITIAL_LR, 2) + .with_min_lr(MIN_LR) + .build() + .unwrap(); + let expected_lrs = [ + INITIAL_LR, // cos(0) + (INITIAL_LR + MIN_LR) * 0.5, // cos(PI/2) + MIN_LR, // cos(PI) + INITIAL_LR, // restart + ]; + test_utils::check_lr_sequence(scheduler, expected_lrs); + } + + #[test] + fn test_save_and_load() { + const NUM_ITERS: usize = 9; + let scheduler = CosineAnnealingLrSchedulerConfig::new(1.0, NUM_ITERS) + .build() + .unwrap(); + test_utils::check_save_load(scheduler, NUM_ITERS / 3 * 2); + } +} diff --git a/crates/burn-optim/src/lr_scheduler/exponential.rs b/crates/burn-optim/src/lr_scheduler/exponential.rs new file mode 100644 index 0000000..2c457ae --- /dev/null +++ b/crates/burn-optim/src/lr_scheduler/exponential.rs @@ -0,0 +1,153 @@ +use burn_core as burn; + +use super::{LrScheduler, LrSchedulerRecord, String}; +use crate::LearningRate; +use crate::RecordState; +use crate::lr_scheduler::module_lr_scheduler::ModuleLrScheduler; +use burn::config::Config; + +/// The configuration for creating an [exponential learning rate scheduler](ExponentialLrScheduler). +/// +/// This scheduler returns the learning rate `initial_lr` at the first step, then multiplies it by +/// a constant `gamma` at every iteration. At any iteration `i` (which starts from 0), the learning +/// rate is given by `initial_lr * gamma^i`. +#[derive(Config, Debug)] +pub struct ExponentialLrSchedulerConfig { + // The initial learning rate. + initial_lr: LearningRate, + // The constant that the learning rate is multiplied by on each iteration. + gamma: f64, +} + +impl ExponentialLrSchedulerConfig { + /// Initializes a [exponential learning rate scheduler](ExponentialLrScheduler). + pub(crate) fn build(&self) -> Result { + if self.initial_lr <= 0. || self.initial_lr > 1. { + return Err("Initial learning rate must be greater than 0 and at most 1".into()); + } + if self.gamma <= 0. || self.gamma > 1. { + return Err("Gamma must be greater than 0 and at most 1".into()); + } + + Ok(ExponentialLrScheduler { + // Such an initial value eliminates the need for special-case handling of the first + // learning rate. + previous_lr: self.initial_lr / self.gamma, + gamma: self.gamma, + }) + } + + /// Initializes a [module learning rate scheduler](ModuleLrScheduler). + /// + /// # Errors + /// + /// An error will be returned if any of the following conditions is true: + /// + /// * `initial_lr` is out of range (0.0, 1.0] + /// * `gamma` is out of range (0.0, 1.0] + pub fn init(&self) -> Result { + self.build().map(|s| s.into()) + } +} + +/// A exponential learning rate scheduler. +/// +/// See [ExponentialLrSchedulerConfig] for more information. +#[derive(Clone, Copy, Debug)] +pub struct ExponentialLrScheduler { + // The previous iteration's learning rate. + previous_lr: LearningRate, + // The constant that the learning rate is multiplied by on each iteration. + gamma: f64, +} + +impl LrScheduler for ExponentialLrScheduler { + fn step(&mut self) -> LearningRate { + self.previous_lr *= self.gamma; + self.previous_lr + } + + fn to_record(&self) -> LrSchedulerRecord { + LrSchedulerRecord::from_state(&ExponentialLrSchedulerState { + previous_lr: self.previous_lr, + }) + } + + fn load_record(&mut self, record: LrSchedulerRecord) { + if let Some(state) = record.into_state::() { + self.previous_lr = state.previous_lr; + } + } +} + +/// The serializable state of an [exponential scheduler](ExponentialLrScheduler). +#[derive(RecordState, Clone, Debug)] +pub struct ExponentialLrSchedulerState { + // `f64` (not the `LearningRate` alias) so the derive recognizes it as a scalar leaf. + previous_lr: f64, +} + +#[cfg(test)] +mod tests { + use super::super::test_utils; + use super::*; + + #[test] + fn config_initial_lr_too_low() { + let r = ExponentialLrSchedulerConfig::new(0., 0.5).build(); + assert!(r.is_err(), "Should return an error"); + assert_eq!( + r.unwrap_err(), + "Initial learning rate must be greater than 0 and at most 1", + "Error messages should match", + ); + } + + #[test] + fn config_initial_lr_too_high() { + let r = ExponentialLrSchedulerConfig::new(1.5, 0.5).build(); + assert!(r.is_err(), "Should return an error"); + assert_eq!( + r.unwrap_err(), + "Initial learning rate must be greater than 0 and at most 1", + "Error messages should match", + ); + } + + #[test] + fn config_gamma_too_low() { + let r = ExponentialLrSchedulerConfig::new(0.5, 0.0).build(); + assert!(r.is_err(), "Should return an error"); + assert_eq!( + r.unwrap_err(), + "Gamma must be greater than 0 and at most 1", + "Error messages should match", + ); + } + + #[test] + fn config_gamma_too_high() { + let r = ExponentialLrSchedulerConfig::new(0.5, 1.5).build(); + assert!(r.is_err(), "Should return an error"); + assert_eq!( + r.unwrap_err(), + "Gamma must be greater than 0 and at most 1", + "Error messages should match", + ); + } + + #[test] + fn test_lr_change() { + let scheduler = ExponentialLrSchedulerConfig::new(0.8, 0.1).build().unwrap(); + let expected_lrs = [0.8, 0.08, 0.008, 0.0008, 0.00008]; + test_utils::check_lr_sequence(scheduler, expected_lrs); + } + + #[test] + fn test_save_and_load() { + let scheduler = ExponentialLrSchedulerConfig::new(0.083, 0.3) + .build() + .unwrap(); + test_utils::check_save_load(scheduler, 7); + } +} diff --git a/crates/burn-optim/src/lr_scheduler/linear.rs b/crates/burn-optim/src/lr_scheduler/linear.rs new file mode 100644 index 0000000..08affd2 --- /dev/null +++ b/crates/burn-optim/src/lr_scheduler/linear.rs @@ -0,0 +1,186 @@ +use burn_core as burn; + +use super::{LrScheduler, LrSchedulerRecord, String}; +use crate::LearningRate; +use crate::RecordState; +use crate::lr_scheduler::module_lr_scheduler::ModuleLrScheduler; +use burn::config::Config; + +/// The configuration for creating a [linear learning rate scheduler](LinearLrScheduler). +/// +/// This scheduler returns the learning rate `initial_lr` at the first step, then changes it by a +/// constant amount on each iteration until reaching a final learning rate `final_lr`. The +/// `num_iters` parameter controls how many iterations are needed to go from `initial_lr` to +/// `final_lr`. +#[derive(Config, Debug)] +pub struct LinearLrSchedulerConfig { + // The initial learning rate. + initial_lr: LearningRate, + // The final learning rate. + final_lr: LearningRate, + // The number of iterations before reaching the final learning rate. + num_iters: usize, +} + +impl LinearLrSchedulerConfig { + /// Initializes a [linear learning rate scheduler](LinearLrScheduler). + pub(crate) fn build(&self) -> Result { + if self.initial_lr <= 0. || self.initial_lr > 1. { + return Err("Initial learning rate must be greater than 0 and at most 1".into()); + } + if self.final_lr < 0. || self.final_lr > 1. { + return Err("Final learning rate must be at least 0 and at most 1".into()); + } + if self.num_iters == 0 { + return Err("Number of iterations must be at least 1".into()); + } + + Ok(LinearLrScheduler { + final_lr: self.final_lr, + step_size: (self.final_lr - self.initial_lr) / self.num_iters as f64, + remaining_iters: self.num_iters + 1, + }) + } + + /// Initializes a [module learning rate scheduler](ModuleLrScheduler). + /// + /// # Errors + /// + /// An error will be returned if any of the following conditions is true: + /// + /// * `initial_lr` is out of range (0.0, 1.0] + /// * `final_lr` is out of range [0.0, 1.0] + /// * `num_iters` is 0 + pub fn init(&self) -> Result { + self.build().map(|s| s.into()) + } +} + +/// A linear learning rate scheduler. +/// +/// See [LinearLrSchedulerConfig] for more information. +#[derive(Clone, Copy, Debug)] +pub struct LinearLrScheduler { + // The final learning rate after the linear changing process stops. + final_lr: LearningRate, + // The amount that the learning rate changes by on each iteration. + step_size: f64, + // The number of iterations left before reaching the final learning rate. + remaining_iters: usize, +} + +impl LrScheduler for LinearLrScheduler { + fn step(&mut self) -> LearningRate { + self.remaining_iters -= (self.remaining_iters != 0) as usize; + self.final_lr - self.step_size * self.remaining_iters as f64 + } + + fn to_record(&self) -> LrSchedulerRecord { + LrSchedulerRecord::from_state(&LinearLrSchedulerState { + remaining_iters: self.remaining_iters, + }) + } + + fn load_record(&mut self, record: LrSchedulerRecord) { + if let Some(state) = record.into_state::() { + self.remaining_iters = state.remaining_iters; + } + } +} + +/// The serializable state of a [linear scheduler](LinearLrScheduler). +#[derive(RecordState, Clone, Debug)] +pub struct LinearLrSchedulerState { + remaining_iters: usize, +} + +#[cfg(test)] +mod tests { + use super::super::test_utils; + use super::*; + + #[test] + fn config_initial_lr_too_low() { + let r = LinearLrSchedulerConfig::new(0., 0.5, 100).build(); + assert!(r.is_err(), "Should return an error"); + assert_eq!( + r.unwrap_err(), + "Initial learning rate must be greater than 0 and at most 1", + "Error messages should match", + ); + } + + #[test] + fn config_initial_lr_too_high() { + let r = LinearLrSchedulerConfig::new(1.5, 0.5, 100).build(); + assert!(r.is_err(), "Should return an error"); + assert_eq!( + r.unwrap_err(), + "Initial learning rate must be greater than 0 and at most 1", + "Error messages should match", + ); + } + + #[test] + fn config_final_lr_too_low() { + let r = LinearLrSchedulerConfig::new(0.5, -0.5, 100).build(); + assert!(r.is_err(), "Should return an error"); + assert_eq!( + r.unwrap_err(), + "Final learning rate must be at least 0 and at most 1", + "Error messages should match", + ); + } + + #[test] + fn config_final_lr_too_high() { + let r = LinearLrSchedulerConfig::new(0.5, 1.5, 100).build(); + assert!(r.is_err(), "Should return an error"); + assert_eq!( + r.unwrap_err(), + "Final learning rate must be at least 0 and at most 1", + "Error messages should match", + ); + } + + #[test] + fn config_num_iters_too_low() { + let r = LinearLrSchedulerConfig::new(0.9, 0.1, 0).build(); + assert!(r.is_err(), "Should return an error"); + assert_eq!( + r.unwrap_err(), + "Number of iterations must be at least 1", + "Error messages should match", + ); + } + + #[test] + fn test_lr_decreasing() { + let scheduler = LinearLrSchedulerConfig::new(0.9, 0.5, 4).build().unwrap(); + let expected_lrs = [0.9, 0.8, 0.7, 0.6, 0.5, 0.5]; + test_utils::check_lr_sequence(scheduler, expected_lrs); + } + + #[test] + fn test_lr_increasing() { + let scheduler = LinearLrSchedulerConfig::new(0.01, 0.04, 3).build().unwrap(); + let expected_lrs = [0.01, 0.02, 0.03, 0.04, 0.04]; + test_utils::check_lr_sequence(scheduler, expected_lrs); + } + + #[test] + fn test_lr_unchanging() { + let scheduler = LinearLrSchedulerConfig::new(0.3, 0.3, 2).build().unwrap(); + let expected_lrs = [0.3, 0.3, 0.3, 0.3]; + test_utils::check_lr_sequence(scheduler, expected_lrs); + } + + #[test] + fn test_save_and_load() { + const NUM_ITERS: usize = 6; + let scheduler = LinearLrSchedulerConfig::new(1.0, 0.01, NUM_ITERS) + .build() + .unwrap(); + test_utils::check_save_load(scheduler, NUM_ITERS / 3 * 2); + } +} diff --git a/crates/burn-optim/src/lr_scheduler/mod.rs b/crates/burn-optim/src/lr_scheduler/mod.rs new file mode 100644 index 0000000..9c50451 --- /dev/null +++ b/crates/burn-optim/src/lr_scheduler/mod.rs @@ -0,0 +1,27 @@ +/// Constant learning rate scheduler +pub mod constant; + +/// Composed learning rate scheduler +pub mod composed; + +/// Linear learning rate scheduler +pub mod linear; + +/// Learning rate policies +pub mod module_lr_scheduler; + +/// Noam learning rate scheduler +pub mod noam; + +/// Exponential learning rate scheduler +pub mod exponential; + +/// Cosine learning rate scheduler +pub mod cosine; + +/// Step learning rate scheduler +pub mod step; + +mod base; + +pub use base::*; diff --git a/crates/burn-optim/src/lr_scheduler/module_lr_scheduler.rs b/crates/burn-optim/src/lr_scheduler/module_lr_scheduler.rs new file mode 100644 index 0000000..4a7aa7f --- /dev/null +++ b/crates/burn-optim/src/lr_scheduler/module_lr_scheduler.rs @@ -0,0 +1,381 @@ +use burn_core::{self as burn, module::ParamId}; + +use burn::config::Config; +use burn_core::module::ParamGroup; + +use crate::{ + LearningRate, + lr_scheduler::{DynLrScheduler, LrScheduler, LrSchedulerConfig, LrSchedulerRecord}, +}; + +#[derive(Clone)] +struct LrGroup { + group: ParamGroup, + lr: f64, +} + +/// Determines what learning rate to use for a given trainable parameter. +#[derive(Clone, Default)] +pub struct ModuleLearningRate { + groups: Vec, +} + +impl From for ModuleLearningRate { + fn from(value: LearningRate) -> Self { + Self { + groups: vec![LrGroup { + group: ParamGroup::all(), + lr: value, + }], + } + } +} + +impl ModuleLearningRate { + /// Get the effective learning rate for the given parameter. + pub fn lr_from_param(&self, id: ParamId, path: Option<&str>) -> LearningRate { + self.groups + .iter() + .filter_map(|val| val.group.matches(&id, path).then_some(val.lr)) + .next_back() + .expect("Should match at least one parameter group.") + } + + /// Get the base learning rate value which's group matches all parameters. + pub fn base(&self) -> LearningRate { + self.groups + .first() + .expect("Should have at least one learning rate.") + .lr + } +} + +#[derive(Config, Debug)] +struct LrSchedulerGroupConfig { + group: ParamGroup, + scheduler: LrSchedulerConfig, +} + +#[derive(new, Clone)] +struct LrSchedulerGroup { + group: ParamGroup, + scheduler: DynLrScheduler, +} + +/// Configuration for a [ModuleLrScheduler]. +#[derive(Config, Debug)] +pub struct ModuleLrSchedulerConfig { + base: LrSchedulerConfig, + #[config(default = "Vec::new()")] + scheduler_groups: Vec, +} + +/// A learning rate scheduler that maps specific parameter groups to dedicated sub-schedulers. +/// +/// This allows heterogeneous learning rate schedules across different layers of a model +/// (e.g., discriminative layer training or fine-tuning). +#[derive(Clone)] +pub struct ModuleLrScheduler { + groups: Vec, +} + +impl ModuleLrSchedulerConfig { + /// Initialize a new learning rate policy scheduler. + pub fn init(&self) -> Result { + let mut groups = Vec::with_capacity(self.scheduler_groups.len()); + + let base = match &self.base { + LrSchedulerConfig::Constant(lr) => (*lr).into(), + LrSchedulerConfig::Linear(config) => config.build()?.into(), + LrSchedulerConfig::Cosine(config) => config.build()?.into(), + LrSchedulerConfig::Exponential(config) => config.build()?.into(), + LrSchedulerConfig::Noam(config) => config.build()?.into(), + LrSchedulerConfig::Step(config) => config.build()?.into(), + LrSchedulerConfig::Composed(config) => config.build()?.into(), + }; + groups.push(LrSchedulerGroup { + group: ParamGroup::all(), + scheduler: base, + }); + + for group in self.scheduler_groups.iter() { + let scheduler = match &group.scheduler { + LrSchedulerConfig::Constant(lr) => (*lr).into(), + LrSchedulerConfig::Linear(config) => config.build()?.into(), + LrSchedulerConfig::Cosine(config) => config.build()?.into(), + LrSchedulerConfig::Exponential(config) => config.build()?.into(), + LrSchedulerConfig::Noam(config) => config.build()?.into(), + LrSchedulerConfig::Step(config) => config.build()?.into(), + LrSchedulerConfig::Composed(config) => config.build()?.into(), + }; + groups.push(LrSchedulerGroup::new(group.group.clone(), scheduler)); + } + + Ok(ModuleLrScheduler { groups }) + } + + /// Add a new parameter group to the scheduler's policy. + pub fn with_group( + mut self, + group: ParamGroup, + scheduler: impl Into, + ) -> Self { + self.scheduler_groups.push(LrSchedulerGroupConfig { + group, + scheduler: scheduler.into(), + }); + self + } +} + +impl ModuleLrScheduler { + /// Create a [ModuleLrScheduler]. + /// + /// # Arguments + /// + /// * `scheduler` - The policy's default learning rate scheduler. + pub fn new(scheduler: S) -> Self { + Self { + groups: vec![LrSchedulerGroup { + group: ParamGroup::all(), + scheduler: scheduler.into(), + }], + } + } + + /// Perform the scheduler step of every scheduler and returns the effective learning rate policy. + pub fn step(&mut self) -> ModuleLearningRate { + let groups = self + .groups + .iter_mut() + .map(|s| { + let lr = s.scheduler.step(); + + LrGroup { + group: s.group.clone(), + lr, + } + }) + .collect(); + + ModuleLearningRate { groups } + } + + /// Get the current state of the schedulers as a [record](LrSchedulerRecord). + pub fn to_record(&self) -> super::LrSchedulerRecord { + let mut record = LrSchedulerRecord::new(); + for (index, item) in self.groups.iter().enumerate() { + let sub = item.scheduler.to_record(); + record = record.with_record(&index.to_string(), sub); + } + + record + } + + /// Load the state of the schedulers from a [record](LrSchedulerRecord). + pub fn load_record(mut self, record: super::LrSchedulerRecord) -> Self { + self.groups = self + .groups + .into_iter() + .enumerate() + .map(|(index, item)| { + let sub = record.record(&index.to_string()); + let scheduler = item.scheduler.load_record(sub); + + LrSchedulerGroup { + group: item.group, + scheduler, + } + }) + .collect(); + + self + } + + /// Add a new parameter group to the scheduler's policy. + pub fn with_group(mut self, group: ParamGroup, scheduler: impl Into) -> Self { + self.groups.push(LrSchedulerGroup { + group, + scheduler: scheduler.into(), + }); + self + } +} + +impl From for ModuleLrScheduler +where + S: LrScheduler + 'static, +{ + fn from(value: S) -> Self { + Self::new(value) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::lr_scheduler::linear::LinearLrSchedulerConfig; + use burn_core::module::ParamGroup; + + const EPSILON: f64 = 1e-10; + + fn check_approx(actual: f64, expected: f64) { + assert!( + (actual - expected).abs() < EPSILON, + "expected {expected}, got {actual}", + ); + } + + #[test] + fn step_yields_constant_default_lr() { + let mut scheduler = ModuleLrScheduler::new(0.01_f64); + for _ in 0..3 { + check_approx(scheduler.step().base(), 0.01); + } + } + + #[test] + fn step_advances_linear_default_scheduler() { + let linear = LinearLrSchedulerConfig::new(0.9, 0.5, 4).build().unwrap(); + let mut scheduler = ModuleLrScheduler::new(linear); + + let expected = [0.9, 0.8, 0.7, 0.6, 0.5, 0.5]; + for expected_lr in expected { + check_approx(scheduler.step().base(), expected_lr); + } + } + + #[test] + fn save_load_preserves_default_scheduler_state() { + let make = + || ModuleLrScheduler::new(LinearLrSchedulerConfig::new(1.0, 0.1, 9).build().unwrap()); + + let mut original = make(); + let mut truth = make(); + + for _ in 0..5 { + original.step(); + truth.step(); + } + + let record = original.to_record(); + let mut restored = make().load_record(record); + + for _ in 0..4 { + check_approx(restored.step().base(), truth.step().base()); + } + } + + #[test] + fn group_param_gets_group_lr() { + let id_group = ParamId::new(); + let id_default = ParamId::new(); + + let mut scheduler = ModuleLrSchedulerConfig::new(0.001.into()) + .with_group(ParamGroup::from_ids(vec![id_group.clone()]), 0.1) + .init() + .unwrap(); + + let policy = scheduler.step(); + // id_group is in the explicit group = group LR + check_approx(policy.lr_from_param(id_group, None), 0.1); + // id_default is not in any group = default LR + check_approx(policy.lr_from_param(id_default, None), 0.001); + } + + #[test] + fn path_group_matches_param_by_path_substring() { + let mut scheduler = ModuleLrSchedulerConfig::new(0.001.into()) + .with_group(ParamGroup::from_predicate("backbone"), 0.1) + .init() + .unwrap(); + + let policy = scheduler.step(); + let id = ParamId::new(); + + check_approx( + policy.lr_from_param(id.clone(), Some("model.backbone.layer.weight")), + 0.1, + ); + check_approx( + policy.lr_from_param(id, Some("model.head.layer.weight")), + 0.001, + ); + } + + #[test] + fn multiple_groups_are_independent() { + let id_a = ParamId::new(); + let id_b = ParamId::new(); + let id_default = ParamId::new(); + + let mut scheduler = + ModuleLrSchedulerConfig::new(LinearLrSchedulerConfig::new(0.001, 0.0001, 4).into()) + .with_group( + ParamGroup::from_ids(vec![id_a.clone()]), + LinearLrSchedulerConfig::new(0.1, 0.01, 4), + ) + .with_group( + ParamGroup::from_ids(vec![id_b.clone()]), + LinearLrSchedulerConfig::new(0.5, 0.05, 4), + ) + .init() + .unwrap(); + + // Each group returns its own initial LR + let policy = scheduler.step(); + check_approx(policy.lr_from_param(id_a.clone(), None), 0.1); + check_approx(policy.lr_from_param(id_b.clone(), None), 0.5); + check_approx(policy.lr_from_param(id_default.clone(), None), 0.001); + + // All three schedulers advanced; LRs are strictly between initial and final + let policy = scheduler.step(); + let lr_a = policy.lr_from_param(id_a, None); + let lr_b = policy.lr_from_param(id_b, None); + let lr_default = policy.lr_from_param(id_default, None); + assert!( + lr_a < 0.1 && lr_a > 0.01, + "group-a LR should have decayed: {lr_a}" + ); + assert!( + lr_b < 0.5 && lr_b > 0.05, + "group-b LR should have decayed: {lr_b}" + ); + assert!( + lr_default < 0.001 && lr_default > 0.0001, + "default LR should have decayed: {lr_default}" + ); + } + + #[test] + fn save_load_with_groups_preserves_state() { + let id_group = ParamId::new(); + + let make = || { + ModuleLrSchedulerConfig::new(LinearLrSchedulerConfig::new(0.01, 0.001, 9).into()) + .with_group( + ParamGroup::from_ids(vec![id_group.clone()]), + LinearLrSchedulerConfig::new(0.1, 0.01, 9), + ) + .init() + .unwrap() + }; + + let mut original = make(); + let mut truth = make(); + + for _ in 0..5 { + original.step(); + truth.step(); + } + + let record = original.to_record(); + let mut restored = make().load_record(record); + + for _ in 0..4 { + let lr_restored = restored.step().lr_from_param(id_group.clone(), None); + let lr_truth = truth.step().lr_from_param(id_group.clone(), None); + check_approx(lr_restored, lr_truth); + } + } +} diff --git a/crates/burn-optim/src/lr_scheduler/noam.rs b/crates/burn-optim/src/lr_scheduler/noam.rs new file mode 100644 index 0000000..f0c72d4 --- /dev/null +++ b/crates/burn-optim/src/lr_scheduler/noam.rs @@ -0,0 +1,147 @@ +use burn_core as burn; + +use burn::config::Config; + +use super::{LrScheduler, LrSchedulerRecord, String}; +use crate::LearningRate; +use crate::RecordState; +use crate::lr_scheduler::module_lr_scheduler::ModuleLrScheduler; + +/// Configuration to create a [noam](NoamLrScheduler) learning rate scheduler. +#[derive(Config, Debug)] +pub struct NoamLrSchedulerConfig { + /// The overall scale factor for the learning rate decay. + factor: f64, + /// The number of steps before the exponential decay stats. + #[config(default = 4000)] + warmup_steps: usize, + /// The size of the model. + #[config(default = 512)] + model_size: usize, +} + +/// Noam learning rate scheduler as described in [Attention Is All You Need](https://arxiv.org/abs/1706.03762). +#[derive(Clone, Debug)] +pub struct NoamLrScheduler { + warmup_steps: f64, + embedding_size: f64, + factor: f64, + step: f64, +} + +impl NoamLrSchedulerConfig { + /// Initialize a new [noam](NoamLrScheduler) learning rate scheduler. + pub(crate) fn build(&self) -> Result { + if self.warmup_steps == 0 { + return Err( + "Number of steps before exponential decay starts must be greater than 0".into(), + ); + } + if self.model_size == 0 { + return Err("Model size must be greater than 0".into()); + } + + Ok(NoamLrScheduler { + warmup_steps: self.warmup_steps as f64, + embedding_size: self.model_size as f64, + factor: self.factor, + step: 0.0, + }) + } + + /// Initializes a [module learning rate scheduler](ModuleLrScheduler). + /// + /// # Errors + /// + /// An error will be returned if any of the following conditions is true: + /// + /// * `warmup_steps` is 0 + /// * `model_size` is 0 + pub fn init(&self) -> Result { + self.build().map(|s| s.into()) + } +} + +impl LrScheduler for NoamLrScheduler { + fn step(&mut self) -> LearningRate { + self.step += 1.0; + + let arg1 = self.step.powf(-0.5); + let arg2 = self.step * self.warmup_steps.powf(-1.5); + + self.factor * self.embedding_size.powf(-0.5) * f64::min(arg1, arg2) + } + + fn to_record(&self) -> LrSchedulerRecord { + LrSchedulerRecord::from_state(&NoamLrSchedulerState { step: self.step }) + } + + fn load_record(&mut self, record: LrSchedulerRecord) { + if let Some(state) = record.into_state::() { + self.step = state.step; + } + } +} + +/// The serializable state of a [noam scheduler](NoamLrScheduler). +#[derive(RecordState, Clone, Debug)] +pub struct NoamLrSchedulerState { + step: f64, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_config_warmup_steps_invalid() { + let r = NoamLrSchedulerConfig::new(0.1).with_warmup_steps(0).build(); + assert!(r.is_err(), "Should return an error"); + } + + #[test] + fn test_config_warmup_steps_valid() { + let r = NoamLrSchedulerConfig::new(0.1).with_warmup_steps(1).build(); + assert!(r.is_ok(), "Should return a success value"); + } + + #[test] + fn test_config_model_size_invalid() { + let r = NoamLrSchedulerConfig::new(0.1).with_model_size(0).build(); + assert!(r.is_err(), "Should return an error"); + } + + #[test] + fn test_config_model_size_valid() { + let r = NoamLrSchedulerConfig::new(0.1).with_model_size(1).build(); + assert!(r.is_ok(), "Should return a success value"); + } + + #[test] + fn test_function_increase_and_decrease() { + let warmup_steps = 100; + let mut scheduler = NoamLrSchedulerConfig::new(10.0) + .with_warmup_steps(warmup_steps) + .build() + .unwrap(); + let mut lr_current = 0.0; + + for _ in 0..warmup_steps { + let lr = scheduler.step(); + assert!( + lr > lr_current, + "Learning rate should increase before the warmup_steps is reached." + ); + lr_current = lr; + } + + for _ in 0..warmup_steps { + let lr = scheduler.step(); + assert!( + lr < lr_current, + "Learning rate should decrease after the warmup_steps is reached." + ); + lr_current = lr; + } + } +} diff --git a/crates/burn-optim/src/lr_scheduler/step.rs b/crates/burn-optim/src/lr_scheduler/step.rs new file mode 100644 index 0000000..ebf4443 --- /dev/null +++ b/crates/burn-optim/src/lr_scheduler/step.rs @@ -0,0 +1,233 @@ +use burn_core as burn; + +use burn::config::Config; + +use super::{LrScheduler, LrSchedulerRecord, String}; +use crate::lr_scheduler::module_lr_scheduler::ModuleLrScheduler; +use crate::{LearningRate, RecordState}; + +/// The configuration for create a [step learning rate scheduler](StepLrScheduler). +/// +/// This scheduler returns the learning rate `initial_lr` from the start, and keeps doing so until +/// the same value has been given for `step_size` times. Then it multiplies the learning rate by +/// `gamma` before repeating the process. +/// +/// Gamma values out of range (0.0, 1.0) and non-positive initial learning rates are acceptable, but +/// a warning log will be output for such a value in case of mistyping. +/// +/// ## Notes +/// +/// The [step](StepLrScheduler::step) method of the scheduler panics if it is called more than +/// `i32::MAX + 1` times. +#[derive(Config, Debug)] +pub struct StepLrSchedulerConfig { + // The learning rate at the initial step. + initial_lr: LearningRate, + // The number of iterations over which the learning rate remains unchanged before the next + // update. + step_size: usize, + /// The factor by which the learning rate is multiplied with each update. Default: 0.1. + #[config(default = 0.1)] + gamma: f64, +} + +impl StepLrSchedulerConfig { + /// Initializes a [step learning rate scheduler](StepLrScheduler). + pub(crate) fn build(&self) -> Result { + if self.step_size == 0 { + return Err("Step size must be greater than 0".into()); + } + + // Atypical values of `initial_lr` and `gamma` are not rejected because they might be useful + // in some cases like debugging (e.g., https://datascience.stackexchange.com/q/89518). + if self.initial_lr <= 0.0 { + log::warn!( + "Initial learning rate value of {} is not a positive number. Ignore this warning \ + if it is intended.", + self.initial_lr + ); + } + if self.gamma <= 0.0 || self.gamma >= 1.0 { + log::warn!( + "Gamma value of {} is out of range (0.0, 1.0). Ignore this warning if it is \ + intended.", + self.gamma + ); + } + + Ok(StepLrScheduler { + init_lr: self.initial_lr, + step_size: self.step_size, + gamma: self.gamma, + iter_idx: -1, + }) + } + + /// Initializes a [module learning rate scheduler](ModuleLrScheduler). + /// + /// # Errors + /// + /// An error will be returned if `step_size` is 0. + pub fn init(&self) -> Result { + self.build().map(|s| s.into()) + } +} + +/// Step learning rate scheduler. +#[derive(Clone, Debug)] +pub struct StepLrScheduler { + init_lr: LearningRate, + step_size: usize, + gamma: f64, + // The index of the current iteration. + // `i32` is used for avoiding truncating the exponent when taking powers of `gamma`. + iter_idx: i32, +} + +impl LrScheduler for StepLrScheduler { + fn step(&mut self) -> LearningRate { + self.iter_idx = self + .iter_idx + .checked_add(1) + .expect("`.step()` should be called no more than `i32::MAX + 1` times"); + // Type casting below causes no truncation, as all the values fall within the ranges. + self.init_lr + * self + .gamma + .powi((self.iter_idx as usize / self.step_size) as i32) + } + + fn to_record(&self) -> LrSchedulerRecord { + LrSchedulerRecord::from_state(&StepLrSchedulerState { + iter_idx: self.iter_idx, + }) + } + + fn load_record(&mut self, record: LrSchedulerRecord) { + if let Some(state) = record.into_state::() { + self.iter_idx = state.iter_idx; + } + } +} + +/// The serializable state of a [step learning rate scheduler](StepLrScheduler). +#[derive(RecordState, Clone, Debug)] +pub struct StepLrSchedulerState { + iter_idx: i32, +} + +#[cfg(test)] +mod tests { + use super::super::test_utils; + use super::*; + + // Warning logs for initial LR and gamma are not tested because there seems no straightforward + // way to do it. + // + // Creating a mock logger that collects logs into `String` for later examination seems a possible + // solution, but unit tests run in the same process in parallel, where the single logger would + // be shared by multiple tests, so logs from different tests would be mixed up with no easy way + // to separate them. + // Using "--test-threads=1" could prevent mixup, but whether the ability to test logging is + // worth the slowdown would be a question. Also, using a primitive provided by `std` to + // synchronize the logger across tests is not an option since we need to support `no-std`. + // Maybe the mocking approach can be reconsidered after we are given an option to run tests in + // separate processes like what the issue below is proposing: + // https://github.com/rust-lang/rust/issues/47506 + // + // As a side note, a helper crate exists for the exact purpose: + // https://crates.io/crates/testing_logger + // but the crate has been unmaintained and using it would introduce another dependency. + + #[test] + fn test_config_step_size_zero() { + let r = StepLrSchedulerConfig::new(1.0, 0).build(); + assert!(r.is_err(), "Should return an error"); + } + + #[test] + fn test_config_step_size_nonzero() { + let r = StepLrSchedulerConfig::new(1.0, 1).build(); + assert!(r.is_ok(), "Should return a success value"); + } + + #[test] + fn test_config_default_gamma() { + const INIT_LR: LearningRate = 0.4; + const STEP_SIZE: usize = 2; + + let mut default = StepLrSchedulerConfig::new(INIT_LR, STEP_SIZE) + .build() + .unwrap(); + let mut explicit = StepLrSchedulerConfig::new(INIT_LR, STEP_SIZE) + .with_gamma(0.1) + .build() + .unwrap(); + test_utils::compare_steps(&mut default, &mut explicit, 3 * STEP_SIZE); + } + + #[test] + fn test_lr_decreasing() { + let scheduler = StepLrSchedulerConfig::new(0.5, 3) + .with_gamma(0.1) + .build() + .unwrap(); + let expected_lrs = [0.5, 0.5, 0.5, 0.05, 0.05, 0.05, 0.005, 0.005, 0.005]; + test_utils::check_lr_sequence(scheduler, expected_lrs); + } + + #[test] + fn test_lr_increasing() { + let scheduler = StepLrSchedulerConfig::new(0.1, 2) + .with_gamma(2.0) + .build() + .unwrap(); + let expected_lrs = [0.1, 0.1, 0.2, 0.2, 0.4, 0.4]; + test_utils::check_lr_sequence(scheduler, expected_lrs); + } + + #[test] + fn test_lr_unchanging() { + let scheduler = StepLrSchedulerConfig::new(3.1, 1) + .with_gamma(1.0) + .build() + .unwrap(); + let expected_lrs = [3.1, 3.1, 3.1]; + test_utils::check_lr_sequence(scheduler, expected_lrs); + } + + #[test] + fn test_save_and_load() { + const STEP_SIZE: usize = 10; + + let scheduler = StepLrSchedulerConfig::new(0.007, STEP_SIZE) + .with_gamma(0.03) + .build() + .unwrap(); + test_utils::check_save_load(scheduler, 3 * STEP_SIZE / 2); + } + + // It's too time consuming to actually run a scheduler `i32::MAX` steps, so an approach that + // depends on private fields is used to implement the test. + #[test] + fn test_number_of_calls_within_limit() { + // Create a scheduler that has already run `i32::MAX` steps + let mut scheduler = StepLrSchedulerConfig::new(0.1, 2).build().unwrap(); + scheduler.load_record(LrSchedulerRecord::from_state(&StepLrSchedulerState { + iter_idx: i32::MAX - 1, + })); + scheduler.step(); + } + + #[test] + #[should_panic = "i32::MAX"] + fn test_number_of_calls_over_limit() { + // Create a scheduler that has already run `i32::MAX` steps + let mut scheduler = StepLrSchedulerConfig::new(0.1, 2).build().unwrap(); + scheduler.load_record(LrSchedulerRecord::from_state(&StepLrSchedulerState { + iter_idx: i32::MAX - 1, + })); + scheduler.step(); + scheduler.step(); + } +} diff --git a/crates/burn-optim/src/optim/adagrad.rs b/crates/burn-optim/src/optim/adagrad.rs new file mode 100644 index 0000000..509168c --- /dev/null +++ b/crates/burn-optim/src/optim/adagrad.rs @@ -0,0 +1,286 @@ +use burn_core as burn; + +use crate::RecordState; + +use burn::config::Config; +use burn::tensor::Device; +use burn::tensor::Tensor; + +use super::{ + Optimizer, + decay::{WeightDecay, WeightDecayConfig}, + module_optimizer::ModuleOptimizer, +}; +use crate::{LearningRate, grad_clipping::GradientClippingConfig}; + +/// AdaGrad configuration. +#[derive(Config, Debug)] +pub struct AdaGradConfig { + #[config(default = 0.)] + lr_decay: f64, + #[config(default = 1e-5)] + epsilon: f32, + /// [Weight decay](WeightDecayConfig) config. + weight_decay: Option, + /// [Gradient Clipping](GradientClippingConfig) config. + grad_clipping: Option, +} + +/// AdaGrad optimizer +#[derive(Clone)] +pub struct AdaGrad { + lr_decay: LrDecay, + weight_decay: Option, +} + +/// AdaGrad state. +#[derive(RecordState, Clone, new)] +pub struct AdaGradState { + lr_decay: LrDecayState, +} + +impl Optimizer for AdaGrad { + type State = AdaGradState; + + fn step( + &self, + lr: LearningRate, + tensor: Tensor, + mut grad: Tensor, + state: Option>, + ) -> (Tensor, Option>) { + let mut state_lr_decay = None; + + if let Some(state) = state { + state_lr_decay = Some(state.lr_decay); + } + + if let Some(weight_decay) = &self.weight_decay { + grad = weight_decay.transform(grad, tensor.clone()); + } + + let (grad, state_lr_decay) = self.lr_decay.transform(grad, lr, state_lr_decay); + + let state = AdaGradState::new(state_lr_decay); + + (tensor - grad, Some(state)) + } + + fn to_device(mut state: Self::State, device: &Device) -> Self::State { + state.lr_decay = state.lr_decay.to_device(device); + state + } +} + +impl AdaGradConfig { + /// Build an [`AdaGrad`] from the config. + pub(crate) fn build(&self) -> AdaGrad { + AdaGrad { + lr_decay: LrDecay { + lr_decay: self.lr_decay, + epsilon: self.epsilon, + }, + weight_decay: self.weight_decay.as_ref().map(WeightDecay::new), + } + } + + /// Initialize AdaGrad optimizer. + /// + /// # Returns + /// + /// Returns an optimizer that can be used to optimize a module. + pub fn init(&self) -> ModuleOptimizer { + let mut optim = ModuleOptimizer::from(self.build()); + if let Some(config) = &self.grad_clipping { + optim = optim.with_grad_clipping(config.init()); + } + optim + } +} + +/// Learning rate decay state (also includes sum state). +#[derive(RecordState, new, Clone)] +pub struct LrDecayState { + time: usize, + sum: Tensor, +} + +#[derive(Clone)] +struct LrDecay { + lr_decay: f64, + epsilon: f32, +} + +impl LrDecay { + pub fn transform( + &self, + grad: Tensor, + lr: LearningRate, + lr_decay_state: Option>, + ) -> (Tensor, LrDecayState) { + let state = if let Some(mut state) = lr_decay_state { + state.sum = state.sum.add(grad.clone().square()); + state.time += 1; + state + } else { + LrDecayState::new(1, grad.clone().square()) + }; + + let new_lr = lr / (1. + (state.time as f64 - 1.) * self.lr_decay); + + let grad = grad + .div(state.sum.clone().sqrt().add_scalar(self.epsilon)) + .mul_scalar(new_lr); + + (grad, state) + } +} + +impl LrDecayState { + /// Move state to device. + /// + /// # Arguments + /// + /// * `device` - Device to move state to. + /// + /// # Returns + /// + /// Returns state moved to device. + pub fn to_device(mut self, device: &Device) -> Self { + self.sum = self.sum.to_device(device); + self + } +} + +#[cfg(test)] +mod tests { + use burn::tensor::Tolerance; + + use super::*; + use crate::GradientsParams; + use burn::module::Param; + use burn::tensor::{Distribution, Tensor, TensorData}; + use burn_nn::{Linear, LinearConfig}; + + const LEARNING_RATE: LearningRate = 0.01; + + #[test] + fn test_adagrad_optimizer_save_load_state() { + let device = Device::default().autodiff(); + let linear = LinearConfig::new(6, 6).init(&device); + let x = Tensor::<2>::random([2, 6], Distribution::Default, &device); + let mut optimizer = create_adagrad(); + let grads = linear.forward(x).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let _linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let bytes = optimizer.into_bytes().unwrap(); + assert!(!bytes.is_empty()); + + #[cfg(feature = "std")] + optimizer + .save(std::env::temp_dir().as_path().join("test_optim_adagrad")) + .unwrap(); + + let state_optim_before = optimizer.to_record(); + let optimizer = create_adagrad().from_bytes(bytes).unwrap(); + let state_optim_after = optimizer.to_record(); + + assert_eq!(state_optim_before.len(), state_optim_after.len()); + } + + #[test] + fn test_adagrad_optimizer_with_numbers() { + let device = Device::default().autodiff(); + let linear = given_linear_layer( + TensorData::from([ + [-0.3206, 0.1374, 0.4043, 0.3200, 0.0859, 0.0671], + [0.0777, -0.0185, -0.3667, 0.2550, 0.1955, -0.2922], + [-0.0190, 0.0346, -0.2962, 0.2484, -0.2780, 0.3130], + [-0.2980, -0.2214, -0.3715, -0.2981, -0.0761, 0.1626], + [0.3300, -0.2182, 0.3717, -0.1729, 0.3796, -0.0304], + [-0.0159, -0.0120, 0.1258, 0.1921, 0.0293, 0.3833], + ]), + TensorData::from([-0.3905, 0.0884, -0.0970, 0.1176, 0.1366, 0.0130]), + &device, + ); + let x_1 = Tensor::<2>::from_floats( + [ + [0.6294, 0.0940, 0.8176, 0.8824, 0.5228, 0.4310], + [0.7152, 0.9559, 0.7893, 0.5684, 0.5939, 0.8883], + ], + &device, + ) + .require_grad(); + let x_2 = Tensor::<2>::from_floats( + [ + [0.8491, 0.2108, 0.8939, 0.4433, 0.5527, 0.2528], + [0.3270, 0.0412, 0.5538, 0.9605, 0.3195, 0.9085], + ], + &device, + ) + .require_grad(); + + let mut optimizer = AdaGradConfig::new() + .with_epsilon(1e-8) + .with_lr_decay(0.5) + .init(); + + let grads = linear.forward(x_1).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let grads = linear.forward(x_2).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let state_updated = linear; + let weights_expected = TensorData::from([ + [-0.334989, 0.123011, 0.389911, 0.305611, 0.071511, 0.052711], + [ + 0.066144, -0.030056, -0.378256, 0.243444, 0.183944, -0.303756, + ], + [ + -0.033462, 0.020138, -0.310662, 0.233938, -0.292462, 0.298538, + ], + [ + -0.312636, -0.236036, -0.386136, -0.312736, -0.090736, 0.147964, + ], + [ + 0.315896, -0.232304, 0.357596, -0.187004, 0.365496, -0.044504, + ], + [-0.030305, -0.026405, 0.111395, 0.177695, 0.014895, 0.368895], + ]); + let bias_expected = TensorData::from([ + -0.405214, 0.073686, -0.111714, 0.102886, 0.121886, -0.001714, + ]); + + let (weight_updated, bias_updated) = ( + state_updated.weight.val().into_data(), + state_updated.bias.unwrap().val().into_data(), + ); + + let tolerance = Tolerance::absolute(1e-6); + bias_updated.assert_approx_eq::(&bias_expected, tolerance); + weight_updated.assert_approx_eq::(&weights_expected, tolerance); + } + + fn given_linear_layer(weight: TensorData, bias: TensorData, device: &Device) -> Linear { + Linear { + weight: Param::from_data(weight, device), + bias: Some(Param::from_data(bias, device)), + } + } + + fn create_adagrad() -> ModuleOptimizer { + let config = AdaGradConfig::new(); + AdaGrad { + lr_decay: LrDecay { + lr_decay: config.lr_decay, + epsilon: config.epsilon, + }, + weight_decay: config.weight_decay.as_ref().map(WeightDecay::new), + } + .into() + } +} diff --git a/crates/burn-optim/src/optim/adam.rs b/crates/burn-optim/src/optim/adam.rs new file mode 100644 index 0000000..3e65c07 --- /dev/null +++ b/crates/burn-optim/src/optim/adam.rs @@ -0,0 +1,551 @@ +use burn_core as burn; + +use crate::RecordState; + +use burn::config::Config; +use burn::tensor::Device; +use burn::tensor::Tensor; + +use super::{ + Optimizer, + decay::{WeightDecay, WeightDecayConfig}, + module_optimizer::ModuleOptimizer, +}; +use crate::{LearningRate, grad_clipping::GradientClippingConfig}; + +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float as _; + +/// Adam configuration. +#[derive(Config, Debug)] +pub struct AdamConfig { + /// Parameter for Adam. + #[config(default = 0.9)] + beta_1: f32, + /// Parameter for Adam. + #[config(default = 0.999)] + beta_2: f32, + /// A value required for numerical stability. + #[config(default = 1e-5)] + epsilon: f32, + /// Whether to use AMSGrad algorithm + #[config(default = false)] + amsgrad: bool, + /// [Weight decay](WeightDecayConfig) config. + weight_decay: Option, + /// [Gradient Clipping](GradientClippingConfig) config. + grad_clipping: Option, +} + +/// Adam optimizer. +/// +/// See: +/// - [Adam: A Method for Stochastic Optimization](https://arxiv.org/pdf/1412.6980.pdf). +/// - [On the Convergence of Adam and Beyond](https://openreview.net/forum?id=ryQu7f-RZ) +#[derive(Clone)] +pub struct Adam { + momentum: AdaptiveMomentum, + weight_decay: Option, +} + +/// Adam state. +#[derive(RecordState, Clone, new)] +pub struct AdamState { + /// The current adaptive momentum. + pub momentum: AdaptiveMomentumState, +} + +impl Optimizer for Adam { + type State = AdamState; + + fn step( + &self, + lr: LearningRate, + tensor: Tensor, + mut grad: Tensor, + state: Option>, + ) -> (Tensor, Option>) { + let mut state_momentum = None; + + if let Some(state) = state { + state_momentum = Some(state.momentum); + } + + if let Some(weight_decay) = &self.weight_decay { + grad = weight_decay.transform(grad, tensor.clone()); + } + + let (grad, state_momentum) = self.momentum.transform(grad, state_momentum); + + let state = AdamState::new(state_momentum); + let delta = grad.mul_scalar(lr); + + (tensor - delta, Some(state)) + } + + fn to_device(mut state: Self::State, device: &Device) -> Self::State { + state.momentum = state.momentum.to_device(device); + state + } +} + +impl AdamConfig { + /// Build an [`Adam`] from the config. + pub(crate) fn build(&self) -> Adam { + Adam { + momentum: AdaptiveMomentum { + beta_1: self.beta_1, + beta_2: self.beta_2, + epsilon: self.epsilon, + amsgrad: self.amsgrad, + }, + weight_decay: self.weight_decay.as_ref().map(WeightDecay::new), + } + } + + /// Initialize Adam optimizer. + /// + /// # Returns + /// + /// Returns an optimizer that can be used to optimize a module. + pub fn init(&self) -> ModuleOptimizer { + let mut optim = ModuleOptimizer::from(self.build()); + if let Some(config) = &self.grad_clipping { + optim = optim.with_grad_clipping(config.init()); + } + optim + } +} + +/// Adaptive momentum state. +#[derive(RecordState, new, Clone)] +pub struct AdaptiveMomentumState { + /// The number of iterations aggregated. + pub time: usize, + /// The first order momentum. + pub moment_1: Tensor, + /// The second order momentum. + pub moment_2: Tensor, + /// Max of second order momentum (for AMSGrad) + #[new(default)] + pub max_moment_2: Option>, +} + +#[derive(Clone)] +struct AdaptiveMomentum { + beta_1: f32, + beta_2: f32, + epsilon: f32, + amsgrad: bool, +} + +impl AdaptiveMomentum { + pub fn transform( + &self, + grad: Tensor, + momentum_state: Option>, + ) -> (Tensor, AdaptiveMomentumState) { + let state = if let Some(mut state) = momentum_state { + let factor = 1.0 - self.beta_1; + state.moment_1 = state + .moment_1 + .mul_scalar(self.beta_1) + .add(grad.clone().mul_scalar(factor)); + + let factor = 1.0 - self.beta_2; + state.moment_2 = state + .moment_2 + .mul_scalar(self.beta_2) + .add(grad.square().mul_scalar(factor)); + if self.amsgrad { + let max_v = state + .max_moment_2 + .take() + .unwrap_or_else(|| state.moment_2.clone()); + + let new_max = max_v.max_pair(state.moment_2.clone()); + state.max_moment_2 = Some(new_max); + } + + state.time += 1; + + state + } else { + let factor = 1.0 - self.beta_1; + let moment_1 = grad.clone().mul_scalar(factor); + + let factor = 1.0 - self.beta_2; + let moment_2 = grad.square().mul_scalar(factor); + let max_moment_2 = self.amsgrad.then(|| moment_2.clone()); + AdaptiveMomentumState { + time: 1, + moment_1, + moment_2, + max_moment_2, + } + }; + + let time = state.time as i32; + let bias_correction2_sqrt = (1.0 - self.beta_2.powi(time)).sqrt(); + let combined_factor = bias_correction2_sqrt / (1.0 - self.beta_1.powi(time)); + + let v_to_use = if self.amsgrad { + state.max_moment_2.as_ref().unwrap_or(&state.moment_2) + } else { + &state.moment_2 + }; + + let grad = state.moment_1.clone().mul_scalar(combined_factor).div( + v_to_use + .clone() + .sqrt() + .add_scalar(self.epsilon * bias_correction2_sqrt), + ); + (grad, state) + } +} + +impl AdaptiveMomentumState { + /// Move state to device. + /// + /// # Arguments + /// + /// * `device` - Device to move state to. + /// + /// # Returns + /// + /// Returns state moved to device. + pub fn to_device(mut self, device: &Device) -> Self { + self.moment_1 = self.moment_1.to_device(device); + self.moment_2 = self.moment_2.to_device(device); + self.max_moment_2 = self.max_moment_2.map(|tensor| tensor.to_device(device)); + self + } +} + +#[cfg(test)] +mod tests { + use burn::tensor::Tolerance; + + use super::*; + use crate::GradientsParams; + use burn::module::Param; + use burn::tensor::{Distribution, Tensor, TensorData}; + use burn_nn::{Linear, LinearConfig}; + + const LEARNING_RATE: LearningRate = 0.01; + + #[test] + fn test_adam_optimizer_save_load_state() { + let device = Device::default().autodiff(); + let linear = LinearConfig::new(6, 6).init(&device); + let x = Tensor::<2>::random([2, 6], Distribution::Default, &device); + let mut optimizer = create_adam(); + let grads = linear.forward(x).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let _linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let bytes = optimizer.into_bytes().unwrap(); + assert!(!bytes.is_empty()); + + #[cfg(feature = "std")] + optimizer + .save(std::env::temp_dir().as_path().join("test_optim_adam")) + .unwrap(); + + let state_optim_before = optimizer.to_record(); + let optimizer = create_adam().from_bytes(bytes).unwrap(); + let state_optim_after = optimizer.to_record(); + + assert_eq!(state_optim_before.len(), state_optim_after.len()); + } + + /// A burnpack round-trip must restore the full state — the moment tensors, amsgrad's optional + /// `max_moment_2`, and the `time` step counter — so that a subsequent step produces identical + /// parameters whether taken on the original or the reloaded optimizer. + #[test] + fn test_adam_state_survives_burnpack_round_trip() { + let device = Device::default().autodiff(); + let mut linear = LinearConfig::new(6, 6).init(&device); + let mut optimizer = AdamConfig::new().with_amsgrad(true).init(); + + // Warm up the optimizer state over a few steps. + for i in 1..=3 { + let x = Tensor::<2>::ones([2, 6], &device) + .mul_scalar(i as f32 * 0.1) + .require_grad(); + let grads = linear.forward(x).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + } + + // Round-trip the optimizer state through the burnpack format. No device is needed on load: + // each parameter's state is migrated to that parameter's device on the next step. + let bytes = optimizer.into_bytes().unwrap(); + let mut reloaded = AdamConfig::new() + .with_amsgrad(true) + .init() + .from_bytes(bytes) + .unwrap(); + + // One more identical step on each optimizer must yield identical parameters. + let x = Tensor::<2>::ones([2, 6], &device) + .mul_scalar(0.4) + .require_grad(); + let grads_original = + GradientsParams::from_grads(linear.forward(x.clone()).backward(), &linear); + let grads_reloaded = GradientsParams::from_grads(linear.forward(x).backward(), &linear); + + let from_original = optimizer.step(LEARNING_RATE.into(), linear.clone(), grads_original); + let from_reloaded = reloaded.step(LEARNING_RATE.into(), linear, grads_reloaded); + + let weight_original = from_original.weight.to_data(); + let weight_reloaded = from_reloaded.weight.to_data(); + weight_original.assert_approx_eq::(&weight_reloaded, Tolerance::absolute(1e-6)); + } + + #[test] + fn test_adam_optimizer_with_amsgrad_50_steps() { + let device = Device::default().autodiff(); + let mut linear = given_linear_layer( + TensorData::from([ + [-0.3206, 0.1374, 0.4043, 0.3200, 0.0859, 0.0671], + [0.0777, -0.0185, -0.3667, 0.2550, 0.1955, -0.2922], + [-0.0190, 0.0346, -0.2962, 0.2484, -0.2780, 0.3130], + [-0.2980, -0.2214, -0.3715, -0.2981, -0.0761, 0.1626], + [0.3300, -0.2182, 0.3717, -0.1729, 0.3796, -0.0304], + [-0.0159, -0.0120, 0.1258, 0.1921, 0.0293, 0.3833], + ]), + TensorData::from([-0.3905, 0.0884, -0.0970, 0.1176, 0.1366, 0.0130]), + &device, + ); + + let mut optimizer = AdamConfig::new() + .with_epsilon(1e-8) + .with_beta_1(0.9) + .with_beta_2(0.999) + .with_amsgrad(true) + .with_weight_decay(Some(WeightDecayConfig::new(0.5))) + .init(); + + for i in 1..=50 { + let x = Tensor::<2>::ones([2, 6], &device) + .mul_scalar(i as f32 * 0.1) + .require_grad(); + + let grads = linear.forward(x).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + } + + let state_updated = linear; + let weight_updated = state_updated.weight.to_data(); + let bias_updated = state_updated.bias.unwrap().to_data(); + + let weights_expected = TensorData::from([ + [ + -0.9125810265541077, + -0.45855265855789185, + -0.1915993094444275, + -0.2759990692138672, + -0.5099529027938843, + -0.5287043452262878, + ], + [ + -0.5181325674057007, + -0.6139854788780212, + -0.9574727416038513, + -0.34102925658226013, + -0.400514155626297, + -0.8847861886024475, + ], + [ + -0.614483118057251, + -0.5611032247543335, + -0.8887064456939697, + -0.34762972593307495, + -0.8708556890487671, + -0.2830044627189636, + ], + [ + -0.8904699683189392, + -0.8151527643203735, + -0.9621278643608093, + -0.8905676603317261, + -0.671261191368103, + -0.4333854615688324, + ], + [ + -0.26599061489105225, + -0.8119961023330688, + -0.22424538433551788, + -0.7672406435012817, + -0.2163349837064743, + -0.6258266568183899, + ], + [ + -0.611397922039032, + -0.6075160503387451, + -0.4701341986656189, + -0.4039117991924286, + -0.5663845539093018, + -0.21262989938259125, + ], + ]); + let bias_expected = TensorData::from([ + -0.8817203044891357, + -0.4038999378681183, + -0.5889149308204651, + -0.37475723028182983, + -0.3557940721511841, + -0.47914788126945496, + ]); + + let tolerance = Tolerance::absolute(1e-5); + weight_updated.assert_approx_eq::(&weights_expected, tolerance); + bias_updated.assert_approx_eq::(&bias_expected, tolerance); + } + #[test] + fn test_adam_optimizer_with_numbers() { + let device = Device::default().autodiff(); + let linear = given_linear_layer( + TensorData::from([ + [-0.3206, 0.1374, 0.4043, 0.3200, 0.0859, 0.0671], + [0.0777, -0.0185, -0.3667, 0.2550, 0.1955, -0.2922], + [-0.0190, 0.0346, -0.2962, 0.2484, -0.2780, 0.3130], + [-0.2980, -0.2214, -0.3715, -0.2981, -0.0761, 0.1626], + [0.3300, -0.2182, 0.3717, -0.1729, 0.3796, -0.0304], + [-0.0159, -0.0120, 0.1258, 0.1921, 0.0293, 0.3833], + ]), + TensorData::from([-0.3905, 0.0884, -0.0970, 0.1176, 0.1366, 0.0130]), + &device, + ); + let x_1 = Tensor::<2>::from_floats( + [ + [0.6294, 0.0940, 0.8176, 0.8824, 0.5228, 0.4310], + [0.7152, 0.9559, 0.7893, 0.5684, 0.5939, 0.8883], + ], + &device, + ) + .require_grad(); + let x_2 = Tensor::<2>::from_floats( + [ + [0.8491, 0.2108, 0.8939, 0.4433, 0.5527, 0.2528], + [0.3270, 0.0412, 0.5538, 0.9605, 0.3195, 0.9085], + ], + &device, + ) + .require_grad(); + + let mut optimizer = AdamConfig::new() + .with_epsilon(1e-8) + .with_beta_1(0.9) + .with_beta_2(0.999) + .with_weight_decay(Some(WeightDecayConfig::new(0.5))) + .init(); + + let grads = linear.forward(x_1).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let grads = linear.forward(x_2).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let state_updated = linear; + let weights_expected = TensorData::from([ + [-0.340528, 0.118929, 0.384336, 0.300010, 0.066034, 0.047154], + [ + 0.057757, -0.036690, -0.386649, 0.235010, 0.175624, -0.312133, + ], + [ + -0.038940, 0.016306, -0.316151, 0.228410, -0.297819, 0.293047, + ], + [ + -0.317929, -0.239100, -0.391449, -0.318087, -0.095948, 0.142651, + ], + [ + 0.310050, -0.235909, 0.351736, -0.192888, 0.359710, -0.050343, + ], + [-0.035840, -0.030203, 0.105840, 0.172110, 0.009440, 0.363346], + ]); + let bias_expected = TensorData::from([ + -0.410499, 0.068401, -0.116999, 0.097601, 0.116601, -0.006999, + ]); + + let (weight_updated, bias_updated) = ( + state_updated.weight.to_data(), + state_updated.bias.unwrap().to_data(), + ); + + let tolerance = Tolerance::absolute(1e-2); + bias_updated.assert_approx_eq::(&bias_expected, tolerance); + weight_updated.assert_approx_eq::(&weights_expected, tolerance); + } + + #[test] + fn test_adam_optimizer_no_nan() { + let device = Device::default().autodiff(); + let linear = given_linear_layer( + TensorData::from([ + [-0.3206, 0.1374, 0.4043, 0.3200, 0.0859, 0.0671], + [0.0777, -0.0185, -0.3667, 0.2550, 0.1955, -0.2922], + [-0.0190, 0.0346, -0.2962, 0.2484, -0.2780, 0.3130], + [-0.2980, -0.2214, -0.3715, -0.2981, -0.0761, 0.1626], + [0.3300, -0.2182, 0.3717, -0.1729, 0.3796, -0.0304], + [-0.0159, -0.0120, 0.1258, 0.1921, 0.0293, 0.3833], + ]), + TensorData::from([-0.3905, 0.0884, -0.0970, 0.1176, 0.1366, 0.0130]), + &device, + ); + + let x = Tensor::<2>::from_floats( + [ + [0.8491, 0.2108, 0.8939, 0.4433, 0.5527, 0.2528], + [0.3270, 0.0412, 0.5538, 0.9605, 0.3195, 0.9085], + ], + &device, + ) + .require_grad(); + + let mut optimizer = AdamConfig::new() + .with_epsilon(1e-8) + .with_beta_1(0.9) + .with_beta_2(0.999) + .with_weight_decay(Some(WeightDecayConfig::new(0.5))) + .init(); + + let grads = linear.forward(x.clone()).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let grads = linear.forward(x).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let state_updated = linear; + assert!(!state_updated.weight.to_data().as_slice::().unwrap()[0].is_nan()); + } + + fn given_linear_layer(weight: TensorData, bias: TensorData, device: &Device) -> Linear { + Linear { + weight: Param::from_data(weight, device), + bias: Some(Param::from_data(bias, device)), + } + } + + fn create_adam() -> ModuleOptimizer { + let config = AdamConfig::new(); + Adam { + momentum: AdaptiveMomentum { + beta_1: config.beta_1, + beta_2: config.beta_2, + epsilon: config.epsilon, + amsgrad: config.amsgrad, + }, + weight_decay: config.weight_decay.as_ref().map(WeightDecay::new), + } + .into() + } +} diff --git a/crates/burn-optim/src/optim/adamw.rs b/crates/burn-optim/src/optim/adamw.rs new file mode 100644 index 0000000..7195f94 --- /dev/null +++ b/crates/burn-optim/src/optim/adamw.rs @@ -0,0 +1,586 @@ +use burn_core as burn; + +use crate::RecordState; +use burn::config::Config; +use burn::tensor::Device; +use burn::tensor::Tensor; + +use super::{AdaptiveMomentumState, Optimizer, module_optimizer::ModuleOptimizer}; +use crate::{LearningRate, grad_clipping::GradientClippingConfig}; + +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float as _; + +/// [`AdamW`] Configuration. +#[derive(Config, Debug)] +pub struct AdamWConfig { + /// Parameter for AdamW. + #[config(default = 0.9)] + beta_1: f32, + /// Parameter for AdamW. + #[config(default = 0.999)] + beta_2: f32, + /// A value required for numerical stability. + #[config(default = 1e-5)] + epsilon: f32, + /// Weight decay config. + #[config(default = 1e-4)] + weight_decay: f32, + + /// Cautious weight decay config. + /// + /// See: + #[config(default = false)] + cautious_weight_decay: bool, + + /// Whether to use AMSGrad algorithm + #[config(default = false)] + amsgrad: bool, + /// [Gradient Clipping](GradientClippingConfig) config. + grad_clipping: Option, +} + +/// AdamW optimizer. +/// +/// See: +/// - [Decoupled Weight Decay Regularization, Loshchilov and Hutter, 2019](https://arxiv.org/abs/1711.05101). +/// - [Cautious Weight Decay, 2025](https://arxiv.org/abs/2510.12402) +/// - [On the Convergence of Adam and Beyond](https://openreview.net/forum?id=ryQu7f-RZ) +/// +/// Configured by [`AdamWConfig`]. +#[derive(Clone)] +pub struct AdamW { + momentum: AdaptiveMomentumW, + weight_decay: f32, + cautious_weight_decay: bool, +} + +/// AdamW state. +#[derive(RecordState, Clone, new)] +pub struct AdamWState { + /// Th current adaptive momentum state. + pub momentum: AdaptiveMomentumState, +} + +impl Optimizer for AdamW { + type State = AdamWState; + + /// A single optimization step for any tensor that represents the parameters of a model. + fn step( + &self, + // Learning rate. + lr: LearningRate, + // Any tensor that represents the parameters of a model. + tensor: Tensor, + // Gradient of the loss w.r.t. the parameters. + grad: Tensor, + // State of the optimizer. + state: Option>, + ) -> (Tensor, Option>) { + let (raw_delta, momentum_state) = self.momentum.transform(grad, state.map(|s| s.momentum)); + + let decay_rate = lr * (self.weight_decay as f64); + + let decayed_tensor = if decay_rate == 0.0 { + tensor.clone() + } else if self.cautious_weight_decay { + // Cautious weight decay. + // See: https://arxiv.org/abs/2510.12402 + let tensor_pos = tensor.clone().greater_equal_elem(0.0); + let grad_pos = momentum_state.moment_1.clone().greater_equal_elem(0.0); + let differ = tensor_pos.not_equal(grad_pos); + + // Zero out the decay where the decay is counter to the update direction. + tensor.clone() - tensor.mul_scalar(decay_rate).mask_fill(differ, 0.0) + } else { + tensor.clone().mul_scalar(1.0 - decay_rate) + }; + + let tensor_updated = decayed_tensor - raw_delta.mul_scalar(lr); + + let state = AdamWState { + momentum: momentum_state, + }; + + (tensor_updated, Some(state)) + } + + fn to_device(mut state: Self::State, device: &Device) -> Self::State { + state.momentum = state.momentum.to_device(device); + state + } +} + +impl AdamWConfig { + /// Build an [`AdamW`] from the config. + pub(crate) fn build(&self) -> AdamW { + AdamW { + momentum: AdaptiveMomentumW { + beta_1: self.beta_1, + beta_2: self.beta_2, + epsilon: self.epsilon, + amsgrad: self.amsgrad, + }, + weight_decay: self.weight_decay, + cautious_weight_decay: self.cautious_weight_decay, + } + } + + /// Initialize AdamW optimizer. + /// + /// # Returns + /// + /// Returns an optimizer that can be used to optimize a module. + pub fn init(&self) -> ModuleOptimizer { + let mut optim = ModuleOptimizer::from(self.build()); + if let Some(config) = &self.grad_clipping { + optim = optim.with_grad_clipping(config.init()); + } + optim + } +} + +#[derive(Clone)] +struct AdaptiveMomentumW { + beta_1: f32, + beta_2: f32, + epsilon: f32, + amsgrad: bool, +} + +impl AdaptiveMomentumW { + pub fn transform( + &self, + grad: Tensor, + state: Option>, + ) -> (Tensor, AdaptiveMomentumState) { + let factor_1 = 1.0 - self.beta_1; + let factor_2 = 1.0 - self.beta_2; + + let state = if let Some(mut state) = state { + // Update first moment estimate. + state.moment_1 = state + .moment_1 + .mul_scalar(self.beta_1) + .add(grad.clone().mul_scalar(factor_1)); + + // Update second moment estimate. + state.moment_2 = state + .moment_2 + .mul_scalar(self.beta_2) + .add(grad.square().mul_scalar(factor_2)); + + if self.amsgrad { + let max_v = state + .max_moment_2 + .take() + .unwrap_or_else(|| state.moment_2.clone()); + state.max_moment_2 = Some(max_v.max_pair(state.moment_2.clone())); + } + + // Update time. + state.time += 1; + + state + } else { + // Initialize first moment estimate. + let moment_1 = grad.clone().mul_scalar(factor_1); + + // Initialize second moment estimate. + let moment_2 = grad.square().mul_scalar(factor_2); + let max_moment_2 = self.amsgrad.then(|| moment_2.clone()); + AdaptiveMomentumState { + time: 1, + moment_1, + moment_2, + max_moment_2, + } + }; + + let time: i32 = state.time as i32; + + // Compute bias-corrected first and second moment estimates. + let moment_1_corrected = state + .moment_1 + .clone() + .div_scalar(1f32 - self.beta_1.powi(time)); + + let v_to_use = if self.amsgrad { + state.max_moment_2.as_ref().unwrap_or(&state.moment_2) + } else { + &state.moment_2 + }; + + let moment_2_corrected = v_to_use.clone().div_scalar(1f32 - self.beta_2.powi(time)); + + let update_delta = + moment_1_corrected.div(moment_2_corrected.sqrt().add_scalar(self.epsilon)); + + (update_delta, state) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::GradientsParams; + use burn::module::Param; + use burn::tensor::Tolerance; + use burn::tensor::{Distribution, Tensor, TensorData}; + use burn_nn::{Linear, LinearConfig}; + + type FT = f32; + + const LEARNING_RATE: LearningRate = 0.01; + + #[test] + fn test_adamw_optimizer_save_load_state() { + let device = Device::default().autodiff(); + let linear = LinearConfig::new(6, 6).init(&device); + let x = Tensor::<2>::random([2, 6], Distribution::Default, &device); + let mut optimizer = create_adamw(); + let grads = linear.forward(x).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let _linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let bytes = optimizer.into_bytes().unwrap(); + assert!(!bytes.is_empty()); + + #[cfg(feature = "std")] + optimizer + .save(std::env::temp_dir().as_path().join("test_optim_adamw")) + .unwrap(); + + let state_optim_before = optimizer.to_record(); + let optimizer = create_adamw().from_bytes(bytes).unwrap(); + let state_optim_after = optimizer.to_record(); + + assert_eq!(state_optim_before.len(), state_optim_after.len()); + } + #[test] + fn test_adamw_optimizer_with_amsgrad_50_steps() { + let device = Device::default().autodiff(); + let mut linear = given_linear_layer( + TensorData::from([ + [-0.3206, 0.1374, 0.4043, 0.3200, 0.0859, 0.0671], + [0.0777, -0.0185, -0.3667, 0.2550, 0.1955, -0.2922], + [-0.0190, 0.0346, -0.2962, 0.2484, -0.2780, 0.3130], + [-0.2980, -0.2214, -0.3715, -0.2981, -0.0761, 0.1626], + [0.3300, -0.2182, 0.3717, -0.1729, 0.3796, -0.0304], + [-0.0159, -0.0120, 0.1258, 0.1921, 0.0293, 0.3833], + ]), + TensorData::from([-0.3905, 0.0884, -0.0970, 0.1176, 0.1366, 0.0130]), + &device, + ); + + let mut optimizer = AdamWConfig::new() + .with_epsilon(1e-8) + .with_beta_1(0.9) + .with_beta_2(0.999) + .with_amsgrad(true) + .with_weight_decay(0.5) + .init(); + + for i in 1..=50 { + let x = Tensor::<2>::ones([2, 6], &device) + .mul_scalar(i as f32 * 0.1) + .require_grad(); + + let grads = linear.forward(x).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + } + + let state_updated = linear; + let weight_updated = state_updated.weight.to_data(); + let bias_updated = state_updated.bias.unwrap().to_data(); + + let weights_expected = TensorData::from([ + [ + -0.7822558283805847, + -0.42578864097595215, + -0.21805696189403534, + -0.28366872668266296, + -0.46587175130844116, + -0.4805040955543518, + ], + [ + -0.4722539782524109, + -0.5471276640892029, + -0.8181359767913818, + -0.33425918221473694, + -0.3805687427520752, + -0.7601516842842102, + ], + [ + -0.5475167632102966, + -0.5057991743087769, + -0.763265073299408, + -0.3393959403038025, + -0.7490996718406677, + -0.28911691904067993, + ], + [ + -0.7646660208702087, + -0.7050473093986511, + -0.8218720555305481, + -0.7647438049316406, + -0.5919585227966309, + -0.40617525577545166, + ], + [ + -0.27588561177253723, + -0.7025567889213562, + -0.24343004822731018, + -0.6672990918159485, + -0.23728127777576447, + -0.556389570236206, + ], + [ + -0.5451040267944336, + -0.5420684814453125, + -0.4348171353340149, + -0.3832150399684906, + -0.5099242925643921, + -0.23440153896808624, + ], + ]); + let bias_expected = TensorData::from([ + -0.7473056316375732, + -0.3745720386505127, + -0.5188710689544678, + -0.35184532403945923, + -0.33705732226371765, + -0.4332566559314728, + ]); + + let tolerance = Tolerance::absolute(1e-5); + weight_updated.assert_approx_eq::(&weights_expected, tolerance); + bias_updated.assert_approx_eq::(&bias_expected, tolerance); + } + #[test] + fn test_adamw_optimizer_with_numbers() { + let device = Device::default().autodiff(); + let linear = given_linear_layer( + TensorData::from([ + [-0.3206, 0.1374, 0.4043, 0.3200, 0.0859, 0.0671], + [0.0777, -0.0185, -0.3667, 0.2550, 0.1955, -0.2922], + [-0.0190, 0.0346, -0.2962, 0.2484, -0.2780, 0.3130], + [-0.2980, -0.2214, -0.3715, -0.2981, -0.0761, 0.1626], + [0.3300, -0.2182, 0.3717, -0.1729, 0.3796, -0.0304], + [-0.0159, -0.0120, 0.1258, 0.1921, 0.0293, 0.3833], + ]), + TensorData::from([-0.3905, 0.0884, -0.0970, 0.1176, 0.1366, 0.0130]), + &device, + ); + let x_1 = Tensor::<2>::from_floats( + [ + [0.6294, 0.0940, 0.8176, 0.8824, 0.5228, 0.4310], + [0.7152, 0.9559, 0.7893, 0.5684, 0.5939, 0.8883], + ], + &device, + ) + .require_grad(); + let x_2 = Tensor::<2>::from_floats( + [ + [0.8491, 0.2108, 0.8939, 0.4433, 0.5527, 0.2528], + [0.3270, 0.0412, 0.5538, 0.9605, 0.3195, 0.9085], + ], + &device, + ) + .require_grad(); + + let mut optimizer = AdamWConfig::new() + .with_epsilon(1e-8) + .with_beta_1(0.9) + .with_beta_2(0.999) + .with_weight_decay(0.5) + .init(); + + let grads = linear.forward(x_1).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let grads = linear.forward(x_2).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let state_updated = linear; + let weights_expected = TensorData::from([ + [-0.337295, 0.117827, 0.380358, 0.296868, 0.065232, 0.046534], + [ + 0.057032, -0.036518, -0.382951, 0.232516, 0.173738, -0.309182, + ], + [ + -0.038703, 0.016052, -0.313155, 0.225982, -0.295039, 0.289981, + ], + [ + -0.314920, -0.237394, -0.387704, -0.315067, -0.095153, 0.141081, + ], + [ + 0.306815, -0.234226, 0.348083, -0.191115, 0.356002, -0.049993, + ], + [-0.035634, -0.030083, 0.104636, 0.170244, 0.009196, 0.359580], + ]); + let bias_expected = TensorData::from([ + -0.406555, 0.067568, -0.115982, 0.096477, 0.115287, -0.007080, + ]); + + let (weight_updated, bias_updated) = ( + state_updated.weight.to_data(), + state_updated.bias.unwrap().to_data(), + ); + + let tolerance = Tolerance::absolute(1e-2); + bias_updated.assert_approx_eq::(&bias_expected, tolerance); + weight_updated.assert_approx_eq::(&weights_expected, tolerance); + } + + #[test] + fn test_adamw_optimizer_with_numbers_cautious() { + let device = Device::default().autodiff(); + let linear = given_linear_layer( + TensorData::from([ + [-0.3206, 0.1374, 0.4043, 0.3200, 0.0859, 0.0671], + [0.0777, -0.0185, -0.3667, 0.2550, 0.1955, -0.2922], + [-0.0190, 0.0346, -0.2962, 0.2484, -0.2780, 0.3130], + [-0.2980, -0.2214, -0.3715, -0.2981, -0.0761, 0.1626], + [0.3300, -0.2182, 0.3717, -0.1729, 0.3796, -0.0304], + [-0.0159, -0.0120, 0.1258, 0.1921, 0.0293, 0.3833], + ]), + TensorData::from([-0.3905, 0.0884, -0.0970, 0.1176, 0.1366, 0.0130]), + &device, + ); + let x_1 = Tensor::<2>::from_floats( + [ + [0.6294, 0.0940, 0.8176, 0.8824, 0.5228, 0.4310], + [0.7152, 0.9559, 0.7893, 0.5684, 0.5939, 0.8883], + ], + &device, + ) + .require_grad(); + let x_2 = Tensor::<2>::from_floats( + [ + [0.8491, 0.2108, 0.8939, 0.4433, 0.5527, 0.2528], + [0.3270, 0.0412, 0.5538, 0.9605, 0.3195, -0.9085], + ], + &device, + ) + .require_grad(); + + let mut optimizer = AdamWConfig::new() + .with_cautious_weight_decay(true) + .with_epsilon(1e-8) + .with_beta_1(0.9) + .with_beta_2(0.999) + .with_weight_decay(0.5) + .init(); + + let grads = linear.forward(x_1).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let grads = linear.forward(x_2).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let state_updated = linear; + let weights_expected = TensorData::from([ + [-0.337295, 0.117827, 0.380358, 0.296868, 0.065232, 0.046534], + [ + 0.057032, -0.036518, -0.382951, 0.232516, 0.173738, -0.309182, + ], + [ + -0.038703, 0.016052, -0.313155, 0.225982, -0.295039, 0.289981, + ], + [ + -0.314920, -0.237394, -0.387704, -0.315067, -0.095153, 0.141081, + ], + [ + 0.306815, -0.234226, 0.348083, -0.191115, 0.356002, -0.049993, + ], + [ + -0.035634, -0.030083, 0.104636, 0.170244, 0.009196, 0.37061332, + ], + ]); + let bias_expected = TensorData::from([ + -0.406555, 0.067568, -0.115982, 0.096477, 0.115287, -0.007080, + ]); + + let (weight_updated, bias_updated) = ( + state_updated.weight.to_data(), + state_updated.bias.unwrap().to_data(), + ); + + let tolerance = Tolerance::absolute(1e-2); + bias_updated.assert_approx_eq::(&bias_expected, tolerance); + weight_updated.assert_approx_eq::(&weights_expected, tolerance); + } + + #[test] + fn test_adam_optimizer_no_nan() { + let device = Device::default().autodiff(); + let linear = given_linear_layer( + TensorData::from([ + [-0.3206, 0.1374, 0.4043, 0.3200, 0.0859, 0.0671], + [0.0777, -0.0185, -0.3667, 0.2550, 0.1955, -0.2922], + [-0.0190, 0.0346, -0.2962, 0.2484, -0.2780, 0.3130], + [-0.2980, -0.2214, -0.3715, -0.2981, -0.0761, 0.1626], + [0.3300, -0.2182, 0.3717, -0.1729, 0.3796, -0.0304], + [-0.0159, -0.0120, 0.1258, 0.1921, 0.0293, 0.3833], + ]), + TensorData::from([-0.3905, 0.0884, -0.0970, 0.1176, 0.1366, 0.0130]), + &device, + ); + + let x = Tensor::<2>::from_floats( + [ + [0.8491, 0.2108, 0.8939, 0.4433, 0.5527, 0.2528], + [0.3270, 0.0412, 0.5538, 0.9605, 0.3195, 0.9085], + ], + &device, + ) + .require_grad(); + + let mut optimizer = AdamWConfig::new() + .with_epsilon(1e-8) + .with_beta_1(0.9) + .with_beta_2(0.999) + .with_weight_decay(0.5) + .init(); + + let grads = linear.forward(x.clone()).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let grads = linear.forward(x).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let state_updated = linear; + assert!(!state_updated.weight.to_data().as_slice::().unwrap()[0].is_nan()); + } + + fn given_linear_layer(weight: TensorData, bias: TensorData, device: &Device) -> Linear { + Linear { + weight: Param::from_data(weight, device), + bias: Some(Param::from_data(bias, device)), + } + } + + fn create_adamw() -> ModuleOptimizer { + let config = AdamWConfig::new(); + AdamW { + momentum: AdaptiveMomentumW { + beta_1: config.beta_1, + beta_2: config.beta_2, + epsilon: config.epsilon, + amsgrad: config.amsgrad, + }, + weight_decay: config.weight_decay, + cautious_weight_decay: false, + } + .into() + } +} diff --git a/crates/burn-optim/src/optim/adan.rs b/crates/burn-optim/src/optim/adan.rs new file mode 100644 index 0000000..fae141f --- /dev/null +++ b/crates/burn-optim/src/optim/adan.rs @@ -0,0 +1,448 @@ +use burn_core as burn; + +use crate::RecordState; +use burn::config::Config; +use burn::tensor::Device; +use burn::tensor::Tensor; + +use super::{Optimizer, module_optimizer::ModuleOptimizer}; +use crate::{LearningRate, grad_clipping::GradientClippingConfig}; + +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float as _; + +/// [`Adan`] Configuration. +/// +/// See: +/// - [Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models](https://arxiv.org/abs/2208.06677). +#[derive(Config, Debug)] +pub struct AdanConfig { + /// Parameter for the first moment. + #[config(default = 0.98)] + beta_1: f32, + /// Parameter for the gradient-difference momentum. + #[config(default = 0.92)] + beta_2: f32, + /// Parameter for the second moment. + #[config(default = 0.99)] + beta_3: f32, + /// A value required for numerical stability. + #[config(default = 1e-8)] + epsilon: f32, + /// Weight decay factor. + #[config(default = 0.0)] + weight_decay: f32, + /// Disable proximal weight decay and use the decoupled update instead. + #[config(default = false)] + no_prox: bool, + /// [Gradient Clipping](GradientClippingConfig) config. + grad_clipping: Option, +} + +/// Adan optimizer. +/// +/// See: +/// - [Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models](https://arxiv.org/abs/2208.06677). +/// +/// Configured by [`AdanConfig`]. +#[derive(Clone)] +pub struct Adan { + momentum: AdaptiveNesterovMomentum, + weight_decay: f32, + no_prox: bool, +} + +/// Adan state. +#[derive(RecordState, Clone, new)] +pub struct AdanState { + /// The current adaptive Nesterov momentum state. + pub momentum: AdaptiveNesterovMomentumState, +} + +impl Optimizer for Adan { + type State = AdanState; + + fn step( + &self, + lr: LearningRate, + tensor: Tensor, + grad: Tensor, + state: Option>, + ) -> (Tensor, Option>) { + let (raw_delta, momentum_state) = self.momentum.transform(grad, state.map(|s| s.momentum)); + + let decay_rate = lr * (self.weight_decay as f64); + let delta = raw_delta.mul_scalar(lr); + + let tensor_updated = if self.no_prox { + if decay_rate == 0.0 { + tensor - delta + } else { + tensor.mul_scalar(1.0 - decay_rate) - delta + } + } else { + let updated = tensor - delta; + if decay_rate == 0.0 { + updated + } else { + updated.div_scalar(1.0 + decay_rate) + } + }; + + (tensor_updated, Some(AdanState::new(momentum_state))) + } + + fn to_device(mut state: Self::State, device: &Device) -> Self::State { + state.momentum = state.momentum.to_device(device); + state + } +} + +impl AdanConfig { + /// Build an [`Adan`] from the config. + pub(crate) fn build(&self) -> Adan { + Adan { + momentum: AdaptiveNesterovMomentum { + beta_1: self.beta_1, + beta_2: self.beta_2, + beta_3: self.beta_3, + epsilon: self.epsilon, + }, + weight_decay: self.weight_decay, + no_prox: self.no_prox, + } + } + + /// Initialize Adan optimizer. + /// + /// # Returns + /// + /// Returns an optimizer that can be used to optimize a module. + pub fn init(&self) -> ModuleOptimizer { + let mut optim = ModuleOptimizer::from(self.build()); + if let Some(config) = &self.grad_clipping { + optim = optim.with_grad_clipping(config.init()); + } + optim + } +} + +/// Adaptive Nesterov momentum state. +#[derive(RecordState, Clone, new)] +pub struct AdaptiveNesterovMomentumState { + /// The number of iterations aggregated. + pub time: usize, + /// The first order momentum. + pub exp_avg: Tensor, + /// The gradient-difference weighted second order momentum. + pub exp_avg_sq: Tensor, + /// The gradient-difference momentum. + pub exp_avg_diff: Tensor, + /// The negated previous gradient. + pub neg_pre_grad: Tensor, +} + +#[derive(Clone)] +struct AdaptiveNesterovMomentum { + beta_1: f32, + beta_2: f32, + beta_3: f32, + epsilon: f32, +} + +impl AdaptiveNesterovMomentum { + pub fn transform( + &self, + grad: Tensor, + state: Option>, + ) -> (Tensor, AdaptiveNesterovMomentumState) { + let state = if let Some(mut state) = state { + let grad_diff = state.neg_pre_grad.clone().add(grad.clone()); + let grad_diff_sq = grad_diff + .clone() + .mul_scalar(self.beta_2) + .add(grad.clone()) + .square(); + + state.exp_avg = state + .exp_avg + .mul_scalar(self.beta_1) + .add(grad.clone().mul_scalar(1.0 - self.beta_1)); + state.exp_avg_diff = state + .exp_avg_diff + .mul_scalar(self.beta_2) + .add(grad_diff.mul_scalar(1.0 - self.beta_2)); + state.exp_avg_sq = state + .exp_avg_sq + .mul_scalar(self.beta_3) + .add(grad_diff_sq.mul_scalar(1.0 - self.beta_3)); + state.neg_pre_grad = grad.mul_scalar(-1.0); + state.time += 1; + state + } else { + AdaptiveNesterovMomentumState::new( + 1, + grad.clone().mul_scalar(1.0 - self.beta_1), + grad.clone().square().mul_scalar(1.0 - self.beta_3), + grad.zeros_like(), + grad.clone().mul_scalar(-1.0), + ) + }; + + let time = state.time as i32; + let denom = state + .exp_avg_sq + .clone() + .sqrt() + .div_scalar((1.0 - self.beta_3.powi(time)).sqrt()) + .add_scalar(self.epsilon); + let update = state + .exp_avg + .clone() + .div_scalar(1.0 - self.beta_1.powi(time)) + .div(denom.clone()) + .add( + state + .exp_avg_diff + .clone() + .mul_scalar(self.beta_2) + .div_scalar(1.0 - self.beta_2.powi(time)) + .div(denom), + ); + + (update, state) + } +} + +impl AdaptiveNesterovMomentumState { + #[allow(clippy::wrong_self_convention)] + fn to_device(mut self, device: &Device) -> Self { + self.exp_avg = self.exp_avg.to_device(device); + self.exp_avg_sq = self.exp_avg_sq.to_device(device); + self.exp_avg_diff = self.exp_avg_diff.to_device(device); + self.neg_pre_grad = self.neg_pre_grad.to_device(device); + self + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::GradientsParams; + use burn::module::Param; + use burn::tensor::Tolerance; + use burn::tensor::{Distribution, Tensor, TensorData}; + use burn_nn::{Linear, LinearConfig}; + + type FT = f32; + + const LEARNING_RATE: LearningRate = 0.01; + + #[test] + fn test_adan_optimizer_save_load_state() { + let device = Device::default().autodiff(); + let linear = LinearConfig::new(6, 6).init(&device); + let x = Tensor::<2>::random([2, 6], Distribution::Default, &device); + let mut optimizer = create_adan(); + let grads = linear.forward(x).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let _linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let bytes = optimizer.into_bytes().unwrap(); + assert!(!bytes.is_empty()); + + #[cfg(feature = "std")] + optimizer + .save(std::env::temp_dir().as_path().join("test_optim_adan")) + .unwrap(); + + let state_optim_before = optimizer.to_record(); + let optimizer = create_adan().from_bytes(bytes).unwrap(); + let state_optim_after = optimizer.to_record(); + + assert_eq!(state_optim_before.len(), state_optim_after.len()); + } + + #[test] + fn test_adan_optimizer_with_numbers() { + let device = Device::default().autodiff(); + let linear = given_linear_layer( + TensorData::from([ + [-0.3206, 0.1374, 0.4043, 0.3200, 0.0859, 0.0671], + [0.0777, -0.0185, -0.3667, 0.2550, 0.1955, -0.2922], + [-0.0190, 0.0346, -0.2962, 0.2484, -0.2780, 0.3130], + [-0.2980, -0.2214, -0.3715, -0.2981, -0.0761, 0.1626], + [0.3300, -0.2182, 0.3717, -0.1729, 0.3796, -0.0304], + [-0.0159, -0.0120, 0.1258, 0.1921, 0.0293, 0.3833], + ]), + TensorData::from([-0.3905, 0.0884, -0.0970, 0.1176, 0.1366, 0.0130]), + &device, + ); + let x_1 = Tensor::<2>::from_floats( + [ + [0.6294, 0.0940, 0.8176, 0.8824, 0.5228, 0.4310], + [0.7152, 0.9559, 0.7893, 0.5684, 0.5939, 0.8883], + ], + &device, + ) + .require_grad(); + let x_2 = Tensor::<2>::from_floats( + [ + [0.8491, 0.2108, 0.8939, 0.4433, 0.5527, 0.2528], + [0.3270, 0.0412, 0.5538, 0.9605, 0.3195, 0.9085], + ], + &device, + ) + .require_grad(); + + let mut optimizer = AdanConfig::new() + .with_beta_1(0.98) + .with_beta_2(0.92) + .with_beta_3(0.99) + .with_epsilon(1e-8) + .with_weight_decay(0.02) + .init(); + + let grads = linear.forward(x_1).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let grads = linear.forward(x_2).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let state_updated = linear; + let weights_expected = TensorData::from([ + [ + -0.34034607, + 0.11747075, + 0.38426402, + 0.29999772, + 0.06599136, + 0.04719888, + ], + [ + 0.0644293, + -0.031732224, + -0.37979296, + 0.24165839, + 0.18218218, + -0.30532277, + ], + [ + -0.038910445, + 0.01466812, + -0.31599957, + 0.2283826, + -0.29780683, + 0.2929568, + ], + [ + -0.3178632, + -0.24129382, + -0.39133376, + -0.31796312, + -0.09605193, + 0.14255258, + ], + [ + 0.31026322, + -0.23771758, + 0.3519465, + -0.19243571, + 0.35984334, + -0.049992695, + ], + [ + -0.03577819, + -0.031879753, + 0.10586514, + 0.17213862, + 0.009403733, + 0.36326218, + ], + ]); + let bias_expected = TensorData::from([ + -0.4103378, + 0.06837065, + -0.116955206, + 0.097558975, + 0.11655137, + -0.006999196, + ]); + + let (weight_updated, bias_updated) = ( + state_updated.weight.to_data(), + state_updated.bias.unwrap().to_data(), + ); + + let tolerance = Tolerance::absolute(1e-5); + bias_updated.assert_approx_eq::(&bias_expected, tolerance); + weight_updated.assert_approx_eq::(&weights_expected, tolerance); + } + + #[test] + fn test_adan_optimizer_no_nan() { + let device = Device::default().autodiff(); + let linear = given_linear_layer( + TensorData::from([ + [-0.3206, 0.1374, 0.4043, 0.3200, 0.0859, 0.0671], + [0.0777, -0.0185, -0.3667, 0.2550, 0.1955, -0.2922], + [-0.0190, 0.0346, -0.2962, 0.2484, -0.2780, 0.3130], + [-0.2980, -0.2214, -0.3715, -0.2981, -0.0761, 0.1626], + [0.3300, -0.2182, 0.3717, -0.1729, 0.3796, -0.0304], + [-0.0159, -0.0120, 0.1258, 0.1921, 0.0293, 0.3833], + ]), + TensorData::from([-0.3905, 0.0884, -0.0970, 0.1176, 0.1366, 0.0130]), + &device, + ); + + let x = Tensor::<2>::from_floats( + [ + [0.8491, 0.2108, 0.8939, 0.4433, 0.5527, 0.2528], + [0.3270, 0.0412, 0.5538, 0.9605, 0.3195, 0.9085], + ], + &device, + ) + .require_grad(); + + let mut optimizer = AdanConfig::new() + .with_epsilon(1e-8) + .with_weight_decay(0.02) + .init(); + + let grads = linear.forward(x.clone()).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let grads = linear.forward(x).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let state_updated = linear; + assert!(!state_updated.weight.to_data().as_slice::().unwrap()[0].is_nan()); + } + + fn given_linear_layer(weight: TensorData, bias: TensorData, device: &Device) -> Linear { + Linear { + weight: Param::from_data(weight, device), + bias: Some(Param::from_data(bias, device)), + } + } + + fn create_adan() -> ModuleOptimizer { + let config = AdanConfig::new(); + Adan { + momentum: AdaptiveNesterovMomentum { + beta_1: config.beta_1, + beta_2: config.beta_2, + beta_3: config.beta_3, + epsilon: config.epsilon, + }, + weight_decay: config.weight_decay, + no_prox: config.no_prox, + } + .into() + } +} diff --git a/crates/burn-optim/src/optim/base.rs b/crates/burn-optim/src/optim/base.rs new file mode 100644 index 0000000..abb41c2 --- /dev/null +++ b/crates/burn-optim/src/optim/base.rs @@ -0,0 +1,51 @@ +use burn_core::Tensor; + +use burn_core::module::ParamId; +use burn_core::tensor::Device; + +use super::GradientsParams; +use alloc::vec::Vec; + +#[derive(Default)] +/// Exposes multiple gradients for each parameter. +pub struct MultiGradientsParams { + /// Each [GradientsParams] has its associated [Device]. + pub grads: Vec<(GradientsParams, Device)>, +} + +impl MultiGradientsParams { + /// Removes the gradients for the given [parameter id](ParamId). + /// + /// Potentially accumulates the gradients from multiple sources using a device associated with + /// a parameter id. The same parameter will be accumulated using the same device during + /// all training. + pub fn remove(&mut self, id: ParamId) -> Option<(Tensor, Device)> { + let (mut tensor, device, index) = self.select(id)?; + + for (i, (grads, _)) in self.grads.iter_mut().enumerate() { + if i == index { + continue; + } + + if let Some(grad) = grads.remove::(id) { + tensor = tensor + grad.to_device(&device); + } + } + + Some((tensor, device)) + } + + fn select(&mut self, id: ParamId) -> Option<(Tensor, Device, usize)> { + let id_val = id.val() as usize; + for i in 0..self.grads.len() { + let selected_device_index = (id_val + i) % self.grads.len(); + + if let Some(acc) = self.grads[selected_device_index].0.remove::(id) { + let device = &self.grads[selected_device_index].1; + return Some((acc.to_device(device), device.clone(), selected_device_index)); + } + } + + None + } +} diff --git a/crates/burn-optim/src/optim/decay.rs b/crates/burn-optim/src/optim/decay.rs new file mode 100644 index 0000000..332fb7c --- /dev/null +++ b/crates/burn-optim/src/optim/decay.rs @@ -0,0 +1,64 @@ +use burn_core as burn; + +use crate::RecordState; +use burn::config::Config; +use burn::tensor::Device; +use burn::tensor::Tensor; + +/// Configuration to create [weight decay](WeightDecay). +#[derive(Config, Debug)] +pub struct WeightDecayConfig { + /// L2 penalty. + pub penalty: f32, +} + +/// State of [weight decay](WeightDecay). +#[derive(RecordState, Clone, new)] +pub struct WeightDecayState { + pub(crate) grad_last_step: Tensor, +} + +/// Weight decay implementation that transforms gradients. +#[derive(Clone)] +pub struct WeightDecay { + penalty: f32, +} + +impl WeightDecay { + /// Creates a new [weight decay](WeightDecay) from a [config](WeightDecayConfig). + pub fn new(config: &WeightDecayConfig) -> Self { + Self { + penalty: config.penalty, + } + } + + /// Transforms a gradient. + /// + /// # Arguments + /// + /// * `grad` - Gradient to transform. + /// * `tensor` - Tensor param of the last iteration. + /// + /// # Returns + /// + /// * `grad` - Transformed gradient. + pub fn transform(&self, grad: Tensor, tensor: Tensor) -> Tensor { + tensor.mul_scalar(self.penalty).add(grad) + } +} + +impl WeightDecayState { + /// Moves the state to a device. + /// + /// # Arguments + /// + /// * `device` - Device to move the state to. + /// + /// # Returns + /// + /// * `self` - Moved state. + pub fn to_device(mut self, device: &Device) -> Self { + self.grad_last_step = self.grad_last_step.to_device(device); + self + } +} diff --git a/crates/burn-optim/src/optim/grad_accum.rs b/crates/burn-optim/src/optim/grad_accum.rs new file mode 100644 index 0000000..6fdce36 --- /dev/null +++ b/crates/burn-optim/src/optim/grad_accum.rs @@ -0,0 +1,119 @@ +use burn_core as burn; + +use core::marker::PhantomData; + +use burn::module::{AutodiffModule, ModuleVisitor, Param}; +use burn::tensor::Tensor; + +use super::GradientsParams; + +/// Accumulate gradients into a single [GradientsParams] object. +pub struct GradientsAccumulator { + grads: GradientsParams, + phantom: PhantomData, +} + +impl Default for GradientsAccumulator { + fn default() -> Self { + Self::new() + } +} + +impl GradientsAccumulator { + /// Create a new gradients accumulator. + pub fn new() -> Self { + Self { + grads: GradientsParams::new(), + phantom: PhantomData, + } + } +} + +impl GradientsAccumulator { + /// Accumulate the given gradients for each parameter in the given module. + pub fn accumulate(&mut self, module: &M, grads: GradientsParams) + where + M: AutodiffModule, + { + let mut visitor = ModuleGradsAccumulator::::new(&mut self.grads, grads); + module.visit(&mut visitor); + } + + /// Return the accumulated gradients and reset the accumulator state. + pub fn grads(&mut self) -> GradientsParams { + let mut grads = GradientsParams::new(); + core::mem::swap(&mut self.grads, &mut grads); + + grads + } +} + +#[derive(new)] +struct ModuleGradsAccumulator<'a, M> { + grads: &'a mut GradientsParams, + grads_new: GradientsParams, + phantom: PhantomData, +} + +impl ModuleVisitor for ModuleGradsAccumulator<'_, M> { + fn visit_float(&mut self, param: &Param>) { + let grad_updated = match self.grads_new.remove::(param.id) { + Some(new) => match self.grads.remove::(param.id) { + Some(grad) => grad.add(new), + None => new, + }, + None => match self.grads.remove::(param.id) { + Some(grad) => grad, + None => return, + }, + }; + + self.grads.register::(param.id, grad_updated); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::{Device, Distribution}; + use burn_nn::{Linear, LinearConfig}; + + #[test] + fn test_accumulate_gradients_one_step() { + let device = Device::default().autodiff(); + let mut accumulator = GradientsAccumulator::new(); + let layer = layer(&device); + let loss = layer.forward(random_tensor(&device)); + let grads = GradientsParams::from_grads(loss.backward(), &layer); + + accumulator.accumulate(&layer, grads); + + let grads = accumulator.grads(); + assert!(!grads.is_empty()) + } + + #[test] + fn test_accumulate_gradients_two_steps() { + let device = Device::default().autodiff(); + let mut accumulator = GradientsAccumulator::new(); + let layer = layer(&device); + let loss_1 = layer.forward(random_tensor(&device)); + let loss_2 = layer.forward(random_tensor(&device)); + let grads_1 = GradientsParams::from_grads(loss_1.backward(), &layer); + let grads_2 = GradientsParams::from_grads(loss_2.backward(), &layer); + + accumulator.accumulate(&layer, grads_1); + accumulator.accumulate(&layer, grads_2); + + let grads = accumulator.grads(); + assert_eq!(grads.len(), 2) + } + + fn layer(device: &Device) -> Linear { + LinearConfig::new(20, 20).init(device) + } + + fn random_tensor(device: &Device) -> Tensor<2> { + Tensor::<2>::random([2, 20], Distribution::Default, device) + } +} diff --git a/crates/burn-optim/src/optim/grads.rs b/crates/burn-optim/src/optim/grads.rs new file mode 100644 index 0000000..5ff60be --- /dev/null +++ b/crates/burn-optim/src/optim/grads.rs @@ -0,0 +1,130 @@ +use burn_core as burn; + +use burn::{ + Tensor, + tensor::{Device, Gradients, container::TensorContainer}, +}; + +use burn::module::{AutodiffModule, ParamId}; + +use super::visitor::{GradientsParamsChangeDevice, GradientsParamsConverter}; + +/// Data type that contains gradients for parameters. +#[derive(Default, Debug)] +pub struct GradientsParams { + container: TensorContainer, +} + +impl GradientsParams { + /// Creates a new [GradientsParams](GradientsParams). + pub fn new() -> Self { + Self::default() + } + + /// Extract each tensor gradients for the given [module](AutodiffModule). + /// + /// Note: This consumes the gradients. See ['from_module'] to extract gradients only for + /// a specific module. + pub fn from_grads(grads: Gradients, module: &M) -> Self { + let mut grads = grads; + Self::from_module(&mut grads, module) + } + + /// Extract each tensor gradients for the given [module](AutodiffModule). + pub fn from_module(grads: &mut Gradients, module: &M) -> Self { + let mut grads_params = GradientsParams::new(); + let mut visitor = GradientsParamsConverter::::new(grads, &mut grads_params, None); + module.visit(&mut visitor); + grads_params + } + + /// Extract tensor gradients for the given [module](AutodiffModule) and given parameters. + pub fn from_params( + grads: &mut Gradients, + module: &M, + params: &[ParamId], + ) -> Self { + let mut grads_params = GradientsParams::new(); + let mut visitor = + GradientsParamsConverter::::new(grads, &mut grads_params, Some(params.to_vec())); + module.visit(&mut visitor); + grads_params + } + + /// Get the gradients for the given [parameter id](ParamId). + /// + /// # Notes + /// + /// You should use [remove](GradientsParams::remove) if you want to get the gradients + /// only one time. + pub fn get(&self, id: ParamId) -> Option> { + self.container.get(&id) + } + + /// Remove the gradients for the given [parameter id](ParamId). + pub fn remove(&mut self, id: ParamId) -> Option> { + self.container.remove(&id) + } + + /// Register a gradients tensor for the given [parameter id](ParamId). + /// + /// # Notes + /// + /// If a tensor is already registered for the given [parameter id](ParamId), it will be replaced. + pub fn register(&mut self, id: ParamId, value: Tensor) { + // TODO: always call value.inner() to make sure? + self.container.register(id, value) + } + + /// The number of gradients tensors registered. + pub fn len(&self) -> usize { + self.container.len() + } + + /// If any tensor is contained. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Change the device of each tensor gradients registered for the given [module](AutodiffModule). + pub fn to_device(mut self, device: &Device, module: &M) -> Self { + let mut visitor = GradientsParamsChangeDevice::::new(device, &mut self); + module.visit(&mut visitor); + self + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::module::{Module, list_param_ids}; + use burn::tensor::Distribution; + use burn_nn::{Linear, LinearConfig}; + + #[test] + fn test_convert_grads() { + let device = Device::default().autodiff(); + let layer_1 = layer(&device); + let mut layer_2 = layer_1.clone(); + layer_2 = layer_2.fork(&device); + let loss_1 = layer_1.forward(random_tensor(&device)); + let loss_2 = layer_2.forward(random_tensor(&device)); + let grads_1 = GradientsParams::from_grads(loss_1.backward(), &layer_1); + let grads_2 = GradientsParams::from_grads(loss_2.backward(), &layer_2); + + let param_ids_1 = list_param_ids(&layer_1); + let param_ids_2 = list_param_ids(&layer_2); + + assert_eq!(param_ids_1, param_ids_2); + assert_eq!(grads_1.len(), param_ids_1.len()); + assert_eq!(grads_2.len(), param_ids_2.len()); + } + + fn layer(device: &Device) -> Linear { + LinearConfig::new(20, 20).init(device) + } + + fn random_tensor(device: &Device) -> Tensor<2> { + Tensor::<2>::random([2, 20], Distribution::Default, device) + } +} diff --git a/crates/burn-optim/src/optim/lbfgs.rs b/crates/burn-optim/src/optim/lbfgs.rs new file mode 100644 index 0000000..b2fd261 --- /dev/null +++ b/crates/burn-optim/src/optim/lbfgs.rs @@ -0,0 +1,1075 @@ +#![allow(clippy::excessive_precision)] + +use burn_core as burn; + +use super::GradientsParams; +use crate::{LearningRate, OptimizerRecord}; +use crate::{RecordState, StateSink, StateSource}; +use burn::config::Config; +use burn::module::{AutodiffModule, Module, ModuleMapper, ModuleVisitor, Param}; +use burn::store::RecordError; +use burn::tensor::{Bytes, Device, Tensor, TensorData}; +use serde::{Deserialize, Serialize}; + +use alloc::vec; +use alloc::vec::Vec; +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float as _; + +/// Cubic Interpolate +/// +/// Uses two points (x1, f1), (x2, f2) and their first derivatives g1,g2 to construct +/// a cubic interpolant and return its minimum within the given bounds. +fn cubic_interpolate( + x1: f64, + f1: f64, + g1: f64, + x2: f64, + f2: f64, + g2: f64, + bounds: Option<(f64, f64)>, +) -> f64 { + // Compute bounds of interpolation area + let (min_bound, max_bound) = bounds.unwrap_or(if x1 <= x2 { (x1, x2) } else { (x2, x1) }); + // Code for most common case: cubic interpolation of 2 points + // with function and derivative values for both + // Solution in this case (where x2 is the farthest point) + // d1 = g1 + g2 - 3*(f1 - f2) / (x1-x2); + // d2 = sqrt(d1^2 - g1 * g2); + // min_pos = x2 - (x2 - x1)*((g2 + d2 - d1)/(g2 - g1 + 2*d2)); + // t_new = min(max(min_pos,min_bound), max_bound); + let d1 = g1 + g2 - 3.0 * (f1 - f2) / (x1 - x2); + let d2_square = d1 * d1 - g1 * g2; + + if d2_square >= 0.0 { + let d2 = d2_square.sqrt(); + let min_pos = if x1 <= x2 { + x2 - (x2 - x1) * ((g2 + d2 - d1) / (g2 - g1 + 2.0 * d2)) + } else { + x1 - (x1 - x2) * ((g1 + d2 - d1) / (g1 - g2 + 2.0 * d2)) + }; + min_pos.max(min_bound).min(max_bound) + } else { + (min_bound + max_bound) / 2.0 + } +} +/// Auxiliary Struct For Strong_Wolfe +struct LineSearchSample { + // step size + t: f64, + // loss + f: f64, + // gradient + g: Tensor<1>, + // directional derivative + gtd: f64, +} + +#[allow(clippy::too_many_arguments)] +fn strong_wolfe( + // obj_func(x,step size,direction) -> (loss,grad) + obj_func: &mut F, + x: &Tensor<1>, + // initial step size + mut t: f64, + d: &Tensor<1>, + f: f64, + g: Tensor<1>, + gtd: f64, + c1: f64, + c2: f64, + tolerance_change: f64, + max_ls: usize, +) -> (f64, Tensor<1>, f64, usize) +where + F: FnMut(&Tensor<1>, f64, &Tensor<1>) -> (f64, Tensor<1>), +{ + let d_norm: f64 = d.clone().abs().max().into_scalar(); + + // evaluate objective and gradient using initial step + let (mut f_new, mut g_new) = obj_func(x, t, d); + let mut ls_func_evals = 1; + let mut gtd_new = g_new.clone().dot(d.clone()).into_scalar(); + + // bracket an interval [t_prev,t] containing a point satisfying the Wolfe criteria + let (mut t_prev, mut f_prev, mut g_prev, mut gtd_prev) = (0.0, f, g.clone(), gtd); + let mut done = false; + let mut ls_iter = 0; + + // the interval [low,high] using for Zoom phase + let mut bracket: Option<[LineSearchSample; 2]> = None; + // point which satisfy the wolfe condition + let mut wolfe_bracket: Option = None; + while ls_iter < max_ls { + // Checking Conditions. + + // Checking the Armijo Condition and function value increasing condition. + // Armijo: f(x+t*d) <= f(x) + c_1 t gtd + if f_new > (f + c1 * t * gtd) || (ls_iter > 1 && f_new >= f_prev) { + bracket = Some([ + LineSearchSample { + t: t_prev, + f: f_prev, + g: g_prev, + gtd: gtd_prev, + }, + LineSearchSample { + t, + f: f_new, + g: g_new.clone(), + gtd: gtd_new, + }, + ]); + break; + } + + // Checking Strong Wolfe Condition + // |gtd_new| <= -c_2 gtd + if gtd_new.abs() <= -c2 * gtd { + wolfe_bracket = Some(LineSearchSample { + t, + f: f_new, + g: g_new.clone(), + gtd: gtd_new, + }); + done = true; + break; + } + + // gtd_new >=0 , there must be a local minimum in the interval. + if gtd_new >= 0.0 { + bracket = Some([ + LineSearchSample { + t: t_prev, + f: f_prev, + g: g_prev, + gtd: gtd_prev, + }, + LineSearchSample { + t, + f: f_new, + g: g_new.clone(), + gtd: gtd_new, + }, + ]); + break; + } + + // interpolate + let min_step = t + 0.01 * (t - t_prev); + let max_step = t * 10.0; + let t_next = cubic_interpolate( + t_prev, + f_prev, + gtd_prev, + t, + f_new, + gtd_new, + Some((min_step, max_step)), + ); + t_prev = t; + f_prev = f_new; + g_prev = g_new; + gtd_prev = gtd_new; + + // next step + t = t_next; + (f_new, g_new) = obj_func(x, t, d); + ls_func_evals += 1; + gtd_new = g_new.clone().dot(d.clone()).into_scalar(); + ls_iter += 1; + } + if let Some(sample) = wolfe_bracket { + return (sample.f, sample.g, sample.t, ls_func_evals); + } + + let mut bracket = bracket.unwrap_or_else(|| { + [ + LineSearchSample { + t: 0.0, + f, + g: g.clone(), + gtd, + }, + LineSearchSample { + t, + f: f_new, + g: g_new.clone(), + gtd: gtd_new, + }, + ] + }); + + // zoom phase + let mut insuf_progress = false; + + // find high and low points in bracket + let (mut low_idx, mut high_idx) = if bracket[0].f <= bracket[1].f { + (0, 1) + } else { + (1, 0) + }; + + while !done && ls_iter < max_ls { + let diff = (bracket[1].t - bracket[0].t).abs(); + // line-search bracket is so small + if diff * d_norm < tolerance_change { + break; + } + + // compute new trial value + t = cubic_interpolate( + bracket[0].t, + bracket[0].f, + bracket[0].gtd, + bracket[1].t, + bracket[1].f, + bracket[1].gtd, + None, + ); + + let b_min = bracket[0].t.min(bracket[1].t); + let b_max = bracket[0].t.max(bracket[1].t); + let eps = 0.1 * (b_max - b_min); + + if (b_max - t).min(t - b_min) < eps { + // interpolation close to boundary + if insuf_progress || t >= b_max || t <= b_min { + t = if (t - b_max).abs() < (t - b_min).abs() { + b_max - eps + } else { + b_min + eps + }; + insuf_progress = false; + } else { + insuf_progress = true; + } + } else { + insuf_progress = false; + } + + // Evaluate new point + (f_new, g_new) = obj_func(x, t, d); + + ls_func_evals += 1; + gtd_new = g_new.clone().dot(d.clone()).into_scalar(); + ls_iter += 1; + + let armijo_holds = f_new <= (f + c1 * t * gtd) && f_new < bracket[low_idx].f; + + if !armijo_holds { + bracket[high_idx] = LineSearchSample { + t, + f: f_new, + g: g_new, + gtd: gtd_new, + }; + } else { + if gtd_new.abs() <= -c2 * gtd { + return (f_new, g_new, t, ls_func_evals); + } + + if gtd_new * (bracket[high_idx].t - bracket[low_idx].t) >= 0.0 { + bracket[high_idx] = LineSearchSample { + t: bracket[low_idx].t, + f: bracket[low_idx].f, + g: bracket[low_idx].g.clone(), + gtd: bracket[low_idx].gtd, + }; + } + bracket[low_idx] = LineSearchSample { + t, + f: f_new, + g: g_new, + gtd: gtd_new, + }; + } + + if bracket[0].f <= bracket[1].f { + low_idx = 0; + high_idx = 1; + } else { + low_idx = 1; + high_idx = 0; + } + } + // return stuff + ( + bracket[low_idx].f, + bracket[low_idx].g.clone(), + bracket[low_idx].t, + ls_func_evals, + ) +} + +/// Strategy for the line search optimization phase +#[derive(Clone, Default, Debug, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum LineSearchFn { + /// No line search performed + #[default] + None, + /// strong wolfe conditions + /// + /// See: + StrongWolfe, +} + +/// LBFGS Configuration. +#[derive(Config, Debug)] +pub struct LBFGSConfig { + /// Maximal number of iterations per optimization step (default: 20) + #[config(default = 20)] + pub max_iter: usize, + /// Update history size (default: 100). + #[config(default = 100)] + pub history_size: usize, + /// Termination tolerance on first order optimality (default: 1e-7). + #[config(default = 1e-7)] + pub tolerance_grad: f64, + /// Termination tolerance on function value/parameter changes (default: 1e-9). + #[config(default = 1e-9)] + pub tolerance_change: f64, + /// Maximal number of function evaluations per optimization step (default: max_iter * 1.25). + #[config(default = "None")] + pub max_eval: Option, + /// Either ‘strong_wolfe’ or None (default: None). + #[config(default = "LineSearchFn::None")] + pub line_search_fn: LineSearchFn, +} + +impl LBFGSConfig { + /// Initialize AdamW optimizer + /// + /// # Returns + /// + /// Returns an optimizer that can be used to optimize a module + pub fn init(&self) -> LBFGS { + // by default max_eval = max_iter * 5/4 + let max_eval = self.max_eval.unwrap_or(self.max_iter * 5 / 4); + LBFGS { + config: LBFGSConfig { + max_iter: self.max_iter, + history_size: self.history_size, + tolerance_grad: self.tolerance_grad, + tolerance_change: self.tolerance_change, + max_eval: Some(max_eval), + line_search_fn: self.line_search_fn, + }, + state: Default::default(), + } + } +} + +/// Collects gradients in module visit order. +struct FlattenGradsVisitorInner<'a> { + grads: &'a GradientsParams, + tensors: &'a mut Vec>, +} + +impl ModuleVisitor for FlattenGradsVisitorInner<'_> { + fn visit_float(&mut self, param: &Param>) { + if let Some(g) = self.grads.get::(param.id) { + let numel = g.shape().num_elements(); + self.tensors.push(g.reshape([numel])); + } + } +} + +/// Flatten params to inner backend 1D tensor. +fn flatten_params_inner(module: &M) -> Tensor<1> { + let mut tensors = Vec::new(); + let mut visitor = FlattenParamsVisitorInner { + tensors: &mut tensors, + }; + module.visit(&mut visitor); + if tensors.is_empty() { + return Tensor::empty([0], &module.devices()[0].clone().inner()); + } + Tensor::cat(tensors, 0) +} + +struct FlattenParamsVisitorInner<'a> { + tensors: &'a mut Vec>, +} + +impl ModuleVisitor for FlattenParamsVisitorInner<'_> { + fn visit_float(&mut self, param: &Param>) { + let t = param.val().inner(); + let numel = t.shape().num_elements(); + self.tensors.push(t.reshape([numel])); + } +} + +/// Flatten gradients for a module. +fn flatten_grads_inner(module: &M, grads: &GradientsParams) -> Tensor<1> { + let mut tensors = Vec::new(); + let mut visitor = FlattenGradsVisitorInner { + grads, + tensors: &mut tensors, + }; + module.visit(&mut visitor); + if tensors.is_empty() { + return Tensor::empty([0], &module.devices()[0].clone().inner()); + } + Tensor::cat(tensors, 0) +} + +/// Mapper that assigns each float param from a flat inner-backend 1D tensor. +struct ParamsFromFlatMapperInner<'a> { + flat: &'a Tensor<1>, + offset: &'a mut usize, +} + +impl ParamsFromFlatMapperInner<'_> { + fn take_slice(&mut self, numel: usize) -> Tensor<1> { + let start = *self.offset; + *self.offset += numel; + self.flat.clone().slice(start..*self.offset) + } +} + +impl ModuleMapper for ParamsFromFlatMapperInner<'_> { + fn map_float(&mut self, param: Param>) -> Param> { + let (id, tensor, mapper) = param.consume(); + let numel = tensor.shape().num_elements(); + let slice_1d = self.take_slice(numel); + let new_inner = slice_1d.reshape(tensor.shape()); + let new_tensor = Tensor::from_inner(new_inner).require_grad(); + Param::from_mapped_value(id, new_tensor, mapper) + } +} + +/// Overwrite module parameters from a flat inner-backend 1D tensor +fn set_params_from_flat_inner(module: M, flat: Tensor<1>) -> M { + let mut offset = 0; + let mut mapper = ParamsFromFlatMapperInner { + flat: &flat, + offset: &mut offset, + }; + module.map(&mut mapper) +} + +/// L-BFGS optimizer state +#[derive(Clone, RecordState)] +pub struct LBFGSState { + /// Historical displacement vectors + pub history_s: Vec>, + /// Historical gradient difference vectors + pub history_y: Vec>, + /// Search direction + pub d: Option>, + /// Step size from the previous iteration + pub t: Option, + /// Flattened gradient from the previous iteration + pub prev_flat_grad: Option>, + /// Loss value from the previous iteration + pub prev_loss: Option, + /// Global iteration count + pub g_iter: usize, +} + +impl LBFGSState { + /// The device of the state's tensors, if any have been populated. + fn current_device(&self) -> Option { + self.prev_flat_grad + .as_ref() + .or(self.d.as_ref()) + .or(self.history_s.first()) + .map(|t| t.device()) + } + + /// Moves all historical tensors to the target device. + pub fn to_device(self, device: &Device) -> Self { + Self { + history_s: self + .history_s + .into_iter() + .map(|t| t.to_device(device)) + .collect(), + history_y: self + .history_y + .into_iter() + .map(|t| t.to_device(device)) + .collect(), + d: self.d.map(|t| t.to_device(device)), + t: self.t, + prev_flat_grad: self.prev_flat_grad.map(|t| t.to_device(device)), + prev_loss: self.prev_loss, + g_iter: self.g_iter, + } + } +} +impl Default for LBFGSState { + fn default() -> Self { + Self { + history_s: Vec::new(), + history_y: Vec::new(), + d: None, + t: Some(1.0), + prev_flat_grad: None, + prev_loss: None, + g_iter: 0, + } + } +} + +/// L-BFGS optimizer. +/// +/// Ported from [pytorch](https://github.com/pytorch/pytorch/torch/optim/lbfgs.py). Heavily inspired by [miniFunc](https://www.cs.ubc.ca/~schmidtm/Software/minFunc.html) +/// +/// See also: +/// - [L-BFGS](https://en.wikipedia.org/wiki/Limited-memory_BFGS) +/// +/// # Note +/// This optimizer is memory intensive +#[derive(Clone)] +pub struct LBFGS { + config: LBFGSConfig, + state: LBFGSState, +} + +impl LBFGS { + /// Decompose the optimizer state into a serializable [`OptimizerRecord`] (burnpack format). + /// + /// L-BFGS keeps a single global state rather than per-parameter state, so its tensors are + /// named directly (e.g. `history_s.0`) and carry no parameter id. + pub fn to_record(&self) -> OptimizerRecord { + let mut sink = StateSink::default(); + RecordState::state_flatten(&self.state, "", &mut sink); + + let tensors = sink + .tensors + .into_iter() + .map(|(name, data)| { + burn_pack::Tensor::new(name, data.dtype, data.shape, None, data.bytes) + }) + .collect(); + let scalars = sink.scalars.into_iter().collect(); + + OptimizerRecord { + tensors, + scalars, + paths: Default::default(), + } + } + + /// Load the optimizer state from an [`OptimizerRecord`]. + /// + /// State tensors are materialized on the default device; the state is migrated to the gradient + /// device on the next [`step`](LBFGS::step), so no device argument is needed. + pub fn load_record(mut self, record: OptimizerRecord) -> Self { + let device = Device::default(); + let mut source = StateSource::new(record.scalars); + for tensor in record.tensors { + let data = TensorData::from_bytes(tensor.bytes, tensor.shape, tensor.dtype); + source.insert_tensor(tensor.name, data); + } + if let Some(state) = LBFGSState::state_unflatten("", &mut source, &device) { + self.state = state; + } + self + } + + /// Serialize the optimizer state to an in-memory burnpack byte buffer. + pub fn into_bytes(&self) -> Result { + self.to_record().into_bytes() + } + + /// Load the optimizer state from an in-memory burnpack byte buffer. + pub fn from_bytes(self, bytes: Bytes) -> Result { + Ok(self.load_record(OptimizerRecord::from_bytes(bytes)?)) + } + + /// Save the optimizer state to a burnpack file on disk. + #[cfg(feature = "std")] + pub fn save>(&self, path: P) -> Result<(), RecordError> { + self.to_record().save(path) + } + + /// Load the optimizer state from a burnpack file on disk. + #[cfg(feature = "std")] + pub fn load>(self, path: P) -> Result { + Ok(self.load_record(OptimizerRecord::load(path)?)) + } + + /// A single optimization step for any tensor that represents the parameters of a model. + pub fn step(&mut self, lr: LearningRate, mut module: M, mut closure: F) -> (M, f64) + where + M: AutodiffModule + Clone, + F: FnMut(M) -> (f64, GradientsParams), + { + // evaluate initial f(x) and df/dx + let (mut loss, grads) = closure(module.clone()); + let mut current_evals = 1; + + let mut flat_grad = flatten_grads_inner::(&module, &grads); + let mut x_flat = flatten_params_inner::(&module); + + // Migrate the state to the gradient's device when they differ (e.g. just after loading a + // record on the default device). This is a no-op once the state is built from gradients. + let device = flat_grad.device(); + if self.state.current_device().is_some_and(|d| d != device) { + self.state = core::mem::take(&mut self.state).to_device(&device); + } + + let opt_cond = + flat_grad.clone().abs().max().into_scalar::() <= self.config.tolerance_grad; + // optimal condition + if opt_cond { + return (module, loss); + } + + // tensors cached in state + let mut d = self + .state + .d + .take() + .unwrap_or_else(|| flat_grad.clone().neg()); + let mut t = self.state.t.unwrap_or(lr); + let mut prev_flat_grad = self.state.prev_flat_grad.take(); + + let mut n_iter = 0; + + // optimize for a max of max_iter iterations + while n_iter < self.config.max_iter { + // keep track of nb of iterations + n_iter += 1; + self.state.g_iter += 1; + + // compute gradient descent direction + if self.state.g_iter == 1 { + d = flat_grad.clone().neg(); + self.state.history_s.clear(); + self.state.history_y.clear(); + } else { + // do lbfgs update (update memory) + if let Some(pg) = prev_flat_grad.as_ref() { + let y = flat_grad.clone().sub(pg.clone()); + let s = d.clone().mul_scalar(t); + + let ys: f64 = y.clone().dot(s.clone()).into_scalar(); + + if ys > 1e-10 { + // updating memory + if self.state.history_s.len() >= self.config.history_size { + // shift history by one (limited-memory) + self.state.history_s.remove(0); + self.state.history_y.remove(0); + } + self.state.history_s.push(s); + self.state.history_y.push(y); + } + } + + // compute the approximate (L-BFGS) inverse Hessian + // multiplied by the gradient + let num_old = self.state.history_s.len(); + let mut q = flat_grad.clone().neg(); + let mut alphas: Vec> = + vec![Tensor::zeros([1], &flat_grad.device().inner()); num_old]; + + if num_old > 0 { + // multiply by initial Hessian + // r/d is the final direction + for i in (0..num_old).rev() { + let s = &self.state.history_s[i]; + let y = &self.state.history_y[i]; + let rho = y.clone().dot(s.clone()).powf_scalar(-1.0); + let alpha = rho.clone().mul(s.clone().dot(q.clone())); + alphas[i] = alpha.clone(); + q = q.sub(y.clone().mul(alpha)); + } + + let last_s = &self.state.history_s[num_old - 1]; + let last_y = &self.state.history_y[num_old - 1]; + let ys = last_y.clone().dot(last_s.clone()); + let yy = last_y.clone().dot(last_y.clone()); + let h_diag = ys.div(yy); + + let mut r = q.mul(h_diag); + + for ((s, y), alpha) in self + .state + .history_s + .iter() + .zip(self.state.history_y.iter()) + .zip(alphas) + .take(num_old) + { + let rho = y.clone().dot(s.clone()).powf_scalar(-1.0); + + let beta = rho.mul(y.clone().dot(r.clone())); + + r = r.add(s.clone().mul(alpha.sub(beta))); + } + d = r; + } else { + d = q; + } + } + + prev_flat_grad = Some(flat_grad.clone()); + let prev_loss_iter = loss; + + // compute step len + if self.state.g_iter == 1 { + let grad_l1: f64 = flat_grad.clone().abs().sum().into_scalar(); + t = (1.0f64 / grad_l1).min(1.0) * lr; + } else { + t = lr; + } + + // directional derivative + let gtd = flat_grad.clone().dot(d.clone()).into_scalar(); + + if gtd > -self.config.tolerance_change { + break; + } + + let ls_func_evals; + + if let LineSearchFn::StrongWolfe = self.config.line_search_fn { + // perform line search, using user function + let mut obj_func = |current_x: &Tensor<1>, step: f64, dir: &Tensor<1>| { + let update = dir.clone().mul_scalar(step); + let new_x = current_x.clone().add(update); + let tmp_module = set_params_from_flat_inner::(module.clone(), new_x); + let (l, g) = closure(tmp_module); + (l, flatten_grads_inner::(&module, &g)) + }; + + let (ls_f, ls_g, ls_t, evals) = strong_wolfe( + &mut obj_func, + &x_flat, + t, + &d, + loss, + flat_grad.clone(), + gtd, + 1e-4, + 0.9, + self.config.tolerance_change, + self.config.max_eval.unwrap() - current_evals, + ); + + loss = ls_f; + flat_grad = ls_g; + t = ls_t; + ls_func_evals = evals; + + x_flat = x_flat.add(d.clone().mul_scalar(t)); + module = set_params_from_flat_inner::(module, x_flat.clone()); + } else { + // no line search, simply move with fixed-step + let step_vec = d.clone().mul_scalar(t); + x_flat = x_flat.add(step_vec); + module = set_params_from_flat_inner::(module, x_flat.clone()); + // re-evaluate function only if not in last iteration + // the reason we do this: in a stochastic setting, + // no use to re-evaluate that function here + let (new_loss, new_grads) = closure(module.clone()); + loss = new_loss; + flat_grad = flatten_grads_inner::(&module, &new_grads); + ls_func_evals = 1; + } + + // update func eval + current_evals += ls_func_evals; + + // check conditions + + if current_evals >= self.config.max_eval.unwrap() { + break; + } + + if flat_grad.clone().abs().max().into_scalar::() <= self.config.tolerance_grad { + break; + } + + if d.clone().mul_scalar(t).abs().max().into_scalar::() + <= self.config.tolerance_change + { + break; + } + + if (loss - prev_loss_iter).abs() < self.config.tolerance_change { + break; + } + } + self.state.d = Some(d); + self.state.t = Some(t); + self.state.prev_flat_grad = prev_flat_grad; + self.state.prev_loss = Some(loss); + (module, loss) + } + /// Moves the optimizer state to the specified device. + pub fn to_device(self, device: &Device) -> Self { + Self { + config: self.config, + // History tensors reside in InnerBackend, so we convert the device accordingly + state: self.state.to_device(device), + } + } +} + +#[cfg(test)] +mod tests { + + use super::*; + use crate::GradientsParams; + use burn::module::Param; + use burn::tensor::{Tensor, TensorData}; + use burn_nn::Linear; + + fn given_linear_layer(weight: TensorData, bias: TensorData, device: &Device) -> Linear { + Linear { + weight: Param::from_data(weight, device), + bias: Some(Param::from_data(bias, device)), + } + } + #[test] + fn test_cubic_interpolate() { + let tolerance = 1e-8; + + // basic + let (x1, f1, g1, x2, f2, g2) = (-1.0, 1.0, -2.0, 1.0, 1.0, 2.0); + let result = cubic_interpolate(x1, f1, g1, x2, f2, g2, None); + assert!( + (result - 0.00000).abs() < tolerance, + "Basic: Result {} should be close to 0.0", + result + ); + + // bound + let (x1, f1, g1, x2, f2, g2) = (0.0, 0.25, -1.0, 1.0, 0.25, 1.0); + let bounds = Some((0.6, 1.0)); + let result = cubic_interpolate(x1, f1, g1, x2, f2, g2, bounds); + assert!( + (result - 0.6000000000).abs() < tolerance, + "Bound: Result {} should be clamped to 0.6", + result + ); + + // d2_square < 0,should return mid value + let (x1, f1, g1, x2, f2, g2) = (0.0, 0.0, 10.0, 1.0, 5.0, 10.0); + let result = cubic_interpolate(x1, f1, g1, x2, f2, g2, Some((0.0, 1.0))); + assert!( + (result - 0.5000000).abs() < tolerance, + "Fallback: Result {} should be midpoint 0.5", + result + ); + + // asymmetric + let (x1, f1, g1, x2, f2, g2) = (0.0, 1.0, -5.0, 1.0, 0.5, 1.0); + let result = cubic_interpolate(x1, f1, g1, x2, f2, g2, None); + assert!( + (result - 0.4606553370833684).abs() < tolerance, + "Asymmetric: Result {} should be 0.4606553370833684", + result + ); + + // not good value + let (x1, f1, g1, x2, f2, g2) = ( + 1.231232145, + -0.12567458754, + 9.1231243007, + 8.239105015, + -100.9012398021, + 123201321.0293982, + ); + let result_1 = cubic_interpolate(x1, f1, g1, x2, f2, g2, None); + let result_2 = cubic_interpolate(x1, f1, g1, x2, f2, g2, Some((-4.4, 4.4))); + assert!( + (result_1 - 5.9031480234724434).abs() < tolerance, + "not good value 1: Result {} should be 5.9031480234724434", + result + ); + assert!( + (result_2 - 4.4000000000000004).abs() < tolerance, + "not good value 2: Result {} should be 4.4000000000000004", + result + ); + } + #[test] + fn test_strong_wolfe_direct_comparison() { + let device = Device::default().autodiff(); + let tol = 1e-6; + + { + let x = Tensor::<1>::from_floats([2.1321912957_f64], &device); + let d = Tensor::<1>::from_floats([0.91312321_f64], &device); + let t_initial = 1.213132_f64; + fn func(x_base: &Tensor<1>, t_val: f64, d_vec: &Tensor<1>) -> (f64, Tensor<1>) { + let curr_x = x_base.clone().add(d_vec.clone().mul_scalar(t_val)); + let x2 = curr_x.clone().mul(curr_x.clone()); + let x3 = x2.clone().mul(curr_x.clone()); + let x4 = x2.clone().mul(x2.clone()); + + // f(x) = x^4 - 2*x^2 + x + let f_elements = x4 - x2.mul_scalar(2.0) + curr_x.clone(); + + let f_val = f_elements.sum().into_scalar(); + + // g(x) = 4*x^3 - 4*x + 1 + let g = x3.mul_scalar(4.0) - curr_x.clone().mul_scalar(4.0) + + Tensor::ones_like(&curr_x); + + (f_val, g) + } + let (f_init, g_init) = func(&x, 0.0, &d); + let gtd_init = g_init.clone().dot(d.clone()).into_scalar::(); + println!("Initial State: f={},gtd = {}", f_init, gtd_init); + assert!((f_init - 13.7080059052).abs() < tol); + assert!((gtd_init - 28.5305728912).abs() < tol); + let mut obj_func = |xb: &Tensor<1>, tv: f64, dv: &Tensor<1>| func(xb, tv, dv); + + let (f_final, _g_final, t_final, evals) = strong_wolfe( + &mut obj_func, + &x, + t_initial, + &d, + f_init, + g_init, + gtd_init, + 1e-4, // c1 + 0.9, // c2 + 1e-9, // tolerance_change + 10, // max_ls + ); + let g_f = _g_final.into_scalar::(); + println!( + "f_final:{:?},_g_final:{:?},t_final:{:?},evals:{:?}", + f_final, g_f, t_final, evals + ); + assert!((f_final - 13.708005905151367).abs() < tol); + assert!((g_f - 31.2450428009).abs() < tol); + assert!((t_final - 0.0).abs() < tol); + assert!((evals == 11)); + } + } + #[test] + fn test_lbfgs_strong_wolfe_comparison() { + let device = Device::default().autodiff(); + let tol = 1e-5; + let x_data = Tensor::<2>::from_data([[1.0], [2.0], [3.0]], &device); + let y_true = Tensor::<2>::from_data([[3.0], [5.0], [7.0]], &device); + let weight = TensorData::from([[0.5f64]]); + let bias = TensorData::from([0.1f64]); + let module = given_linear_layer(weight, bias, &device); + + let mut optimizer = LBFGSConfig::new() + .with_line_search_fn(LineSearchFn::StrongWolfe) + .init(); + let mut closure = |mod_in: Linear| { + let output = mod_in.forward(x_data.clone()); + let loss = burn_nn::loss::MseLoss::new().forward( + output, + y_true.clone(), + burn_nn::loss::Reduction::Sum, + ); + + let grads = loss.backward(); + let grads_params = GradientsParams::from_grads(grads, &mod_in); + + (loss.into_scalar::(), grads_params) + }; + let initial_loss = closure(module.clone()).0; + assert!((initial_loss - 50.1300048828).abs() < tol); + let (updated_module, final_loss) = optimizer.step(0.001, module, &mut closure); + assert!((final_loss - 0.0234732367).abs() < tol); + let optimized_data: f64 = updated_module.weight.val().into_scalar(); + let optimized_bias: f64 = updated_module.bias.as_ref().unwrap().val().into_scalar(); + assert!((optimized_data - 2.0570652485).abs() < tol); + assert!((optimized_bias - 0.8106800914).abs() < tol); + } + + // A burnpack round-trip of the L-BFGS state (which holds `Vec` history buffers, optional + // tensors and optional scalars) must restore enough that a further step agrees with the original. + #[test] + fn test_lbfgs_burnpack_round_trip() { + let device = Device::default().autodiff(); + let tol = 1e-6; + let x_data = Tensor::<2>::from_data([[1.0], [2.0], [3.0]], &device); + let y_true = Tensor::<2>::from_data([[3.0], [5.0], [7.0]], &device); + let module = given_linear_layer( + TensorData::from([[0.5f64]]), + TensorData::from([0.1f64]), + &device, + ); + + let make_closure = || { + let x = x_data.clone(); + let y = y_true.clone(); + move |mod_in: Linear| { + let output = mod_in.forward(x.clone()); + let loss = burn_nn::loss::MseLoss::new().forward( + output, + y.clone(), + burn_nn::loss::Reduction::Sum, + ); + let grads = loss.backward(); + let grads_params = GradientsParams::from_grads(grads, &mod_in); + (loss.into_scalar::(), grads_params) + } + }; + + let mut optimizer = LBFGSConfig::new() + .with_line_search_fn(LineSearchFn::StrongWolfe) + .init(); + let (module, _) = optimizer.step(0.001, module, &mut make_closure()); + + // Round-trip the optimizer state. State tensors live on the inner (non-autodiff) backend. + let bytes = optimizer.into_bytes().unwrap(); + let mut reloaded = LBFGSConfig::new() + .with_line_search_fn(LineSearchFn::StrongWolfe) + .init() + .from_bytes(bytes) + .unwrap(); + + // A further identical step on each optimizer must agree — exercising the restored history. + let (_, loss_original) = optimizer.step(0.001, module.clone(), &mut make_closure()); + let (_, loss_reloaded) = reloaded.step(0.001, module, &mut make_closure()); + assert!( + (loss_original - loss_reloaded).abs() < tol, + "losses differ after burnpack round-trip: {loss_original} vs {loss_reloaded}" + ); + } + + #[test] + fn test_lbfgs_no_strong_wolfe_comparison() { + let device = Device::default().autodiff(); + let tol = 1e-5; + let x_data = Tensor::<2>::from_data([[1.0], [2.0], [3.0]], &device); + let y_true = Tensor::<2>::from_data([[3.0], [5.0], [7.0]], &device); + let weight = TensorData::from([[0.5f64]]); + let bias = TensorData::from([0.1f64]); + let module = given_linear_layer(weight, bias, &device); + + let mut optimizer = LBFGSConfig::new() + .with_line_search_fn(LineSearchFn::None) + .init(); + let mut closure = |mod_in: Linear| { + let output = mod_in.forward(x_data.clone()); + let loss = burn_nn::loss::MseLoss::new().forward( + output, + y_true.clone(), + burn_nn::loss::Reduction::Sum, + ); + + let grads = loss.backward(); + let grads_params = GradientsParams::from_grads(grads, &mod_in); + + (loss.into_scalar::(), grads_params) + }; + let initial_loss = closure(module.clone()).0; + assert!((initial_loss - 50.1300048828).abs() < tol); + let (updated_module, final_loss) = optimizer.step(0.001, module, &mut closure); + assert!((final_loss - 48.2181930542).abs() < tol); + let optimized_data: f64 = updated_module.weight.val().into_scalar(); + let optimized_bias: f64 = updated_module.bias.as_ref().unwrap().val().into_scalar(); + + assert!((optimized_data - 0.5302446192).abs() < tol); + assert!((optimized_bias - 0.1142520783).abs() < tol); + } +} diff --git a/crates/burn-optim/src/optim/mod.rs b/crates/burn-optim/src/optim/mod.rs new file mode 100644 index 0000000..68631f0 --- /dev/null +++ b/crates/burn-optim/src/optim/mod.rs @@ -0,0 +1,34 @@ +/// Weight decay module for optimizers. +pub mod decay; + +/// Momentum module for optimizers. +pub mod momentum; + +mod adagrad; +mod adam; +mod adamw; +mod adan; +mod base; +mod grad_accum; +mod grads; +mod lbfgs; +mod module; +mod muon; +mod rmsprop; +mod sgd; +mod state; +mod visitor; + +pub use adagrad::*; +pub use adam::*; +pub use adamw::*; +pub use adan::*; +pub use base::*; +pub use grad_accum::*; +pub use grads::*; +pub use lbfgs::*; +pub use module::*; +pub use muon::*; +pub use rmsprop::*; +pub use sgd::*; +pub use state::*; diff --git a/crates/burn-optim/src/optim/module/base.rs b/crates/burn-optim/src/optim/module/base.rs new file mode 100644 index 0000000..b9cbab2 --- /dev/null +++ b/crates/burn-optim/src/optim/module/base.rs @@ -0,0 +1,210 @@ +use alloc::sync::Arc; +use core::any::Any; + +use crate::{RecordState, StateSink, StateSource}; +use burn_core as burn; +use burn_core::tensor::kind::BridgeTensor; + +use crate::LearningRate; +use burn::tensor::{Device, Tensor}; + +/// An opinionated trait to simplify the process of implementing an optimizer. +/// +/// Implementations don't have to handle missing gradients, loading and exporting records, +/// navigate the module parameter structure, handle tracked and untracked tensors, and the likes. +/// Wrap one in a [`ModuleOptimizer`](crate::optim::ModuleOptimizer) to optimize a whole module. +pub trait Optimizer: Send + Sync + Clone + 'static { + /// The state of the optimizer for a single parameter of rank `D`. + /// + /// It implements [`RecordState`] (which itself requires `Send + Sync + 'static`) so it can be + /// decomposed into named tensors and scalars for the burnpack format. + type State: Clone + RecordState; + + /// The optimizer step is performed for one tensor at a time with its gradient and state. + /// + /// Note that the state is passed as parameter, so implementations don't have to handle + /// the saving and loading of recorded states. + fn step( + &self, + lr: LearningRate, + tensor: Tensor, + grad: Tensor, + state: Option>, + ) -> (Tensor, Option>); + + /// Change the device of the state. + /// + /// This function will be called accordingly to have the state on the same device as the + /// gradient and the tensor when the [step](Optimizer::step) function is called. + fn to_device(state: Self::State, device: &Device) -> Self::State; +} + +/// A type-erased optimizer state for a single parameter. +/// +/// It wraps a concrete `O::State` together with its rank `D` so that the rank can be recovered +/// when the state is later interpreted by the originating [`Optimizer`] (during a step, a device +/// transfer or serialization). +#[derive(Clone)] +pub struct DynState { + state: Arc, + rank: usize, +} + +impl DynState { + /// Erase a concrete optimizer state of rank `rank`. + pub fn create(state: T, rank: usize) -> Self { + Self { + state: Arc::new(state), + rank, + } + } + + /// Recover the concrete state by value, moving it out when this is the only handle and + /// cloning only when the underlying state is still shared. Panics if `T` does not match the + /// stored type. + pub fn downcast(self) -> T { + let state = self + .state + .downcast::() + .expect("The dynamic optimizer state should match the optimizer state type."); + Arc::try_unwrap(state).unwrap_or_else(|state| (*state).clone()) + } + + /// Borrow the concrete state without cloning. Panics if `T` does not match the stored type. + pub fn downcast_ref(&self) -> &T { + self.state + .downcast_ref::() + .expect("The dynamic optimizer state should match the optimizer state type.") + } + + /// The rank of the parameter this state belongs to. + pub fn rank(&self) -> usize { + self.rank + } +} + +/// Dispatch a runtime `rank` to a body parameterized by a `const D: usize`. +macro_rules! dispatch_rank { + ($rank:expr, $d:ident => $body:block) => { + match $rank { + 0 => { + const $d: usize = 0; + $body + } + 1 => { + const $d: usize = 1; + $body + } + 2 => { + const $d: usize = 2; + $body + } + 3 => { + const $d: usize = 3; + $body + } + 4 => { + const $d: usize = 4; + $body + } + 5 => { + const $d: usize = 5; + $body + } + 6 => { + const $d: usize = 6; + $body + } + 7 => { + const $d: usize = 7; + $body + } + 8 => { + const $d: usize = 8; + $body + } + other => panic!("Unsupported tensor rank for optimizer state: {other}"), + } + }; +} + +/// Object-safe view over an [`Optimizer`], allowing [`ModuleOptimizer`](crate::optim::ModuleOptimizer) +/// to stay non-generic. Rank-generic operations are dispatched on a runtime rank. +pub trait DynOptimizer: Send + Sync { + /// Perform an optimizer step for a single parameter of the given `rank`. + fn step_dyn( + &self, + rank: usize, + lr: LearningRate, + tensor: BridgeTensor, + grad: BridgeTensor, + state: Option, + ) -> (BridgeTensor, Option); + + /// Move a state to the given device. + fn to_device_dyn(&self, state: DynState, device: &Device) -> DynState; + + /// Decompose a state into named tensors and scalars under `prefix`. + fn state_flatten(&self, prefix: &str, state: &DynState, out: &mut StateSink); + + /// Rebuild a state of the given `rank` from named tensors and scalars under `prefix`. + /// + /// Returns `None` when the record does not contain a reconstructable state for this parameter + /// (e.g. a truncated or foreign file); the caller leaves that parameter without state, so it is + /// re-initialized on the next step. + fn state_unflatten( + &self, + rank: usize, + prefix: &str, + src: &mut StateSource, + device: &Device, + ) -> Option; +} + +impl DynOptimizer for O { + fn step_dyn( + &self, + rank: usize, + lr: LearningRate, + tensor: BridgeTensor, + grad: BridgeTensor, + state: Option, + ) -> (BridgeTensor, Option) { + dispatch_rank!(rank, D => { + let (tensor, state) = self.step( + lr, + Tensor::::from_bridge(tensor), + Tensor::::from_bridge(grad), + state.map(|state| state.downcast::>()), + ); + + (tensor.into_bridge(), state.map(|state| DynState::create(state, D))) + }) + } + + fn to_device_dyn(&self, state: DynState, device: &Device) -> DynState { + dispatch_rank!(state.rank(), D => { + let state = O::to_device::(state.downcast::>(), device); + DynState::create(state, D) + }) + } + + fn state_flatten(&self, prefix: &str, state: &DynState, out: &mut StateSink) { + dispatch_rank!(state.rank(), D => { + RecordState::state_flatten(state.downcast_ref::>(), prefix, out); + }) + } + + fn state_unflatten( + &self, + rank: usize, + prefix: &str, + src: &mut StateSource, + device: &Device, + ) -> Option { + dispatch_rank!(rank, D => { + let state = as RecordState>::state_unflatten(prefix, src, device)?; + Some(DynState::create(state, D)) + }) + } +} diff --git a/crates/burn-optim/src/optim/module/mod.rs b/crates/burn-optim/src/optim/module/mod.rs new file mode 100644 index 0000000..4219067 --- /dev/null +++ b/crates/burn-optim/src/optim/module/mod.rs @@ -0,0 +1,10 @@ +mod base; +pub use base::*; + +/// Adaptor module for optimizers. +pub mod module_optimizer; +pub use module_optimizer::*; + +/// Record module for optimizers. +pub mod record; +pub use record::*; diff --git a/crates/burn-optim/src/optim/module/module_optimizer.rs b/crates/burn-optim/src/optim/module/module_optimizer.rs new file mode 100644 index 0000000..d98e6e5 --- /dev/null +++ b/crates/burn-optim/src/optim/module/module_optimizer.rs @@ -0,0 +1,749 @@ +use burn_core as burn; +use burn_core::module::ParamGroup; + +use super::Optimizer; +use crate::lr_scheduler::module_lr_scheduler::ModuleLearningRate; +use crate::{ + DynOptimizer, DynState, MultiGradientsParams, OptimizerRecord, StateSink, StateSource, + grad_clipping::GradientClipping, optim::GradientsParams, optim::state::join_path, +}; + +use alloc::collections::BTreeMap; +use alloc::string::ToString; +use alloc::sync::Arc; +use alloc::vec::Vec; +use burn::module::{AutodiffModule, ModuleMapper, Param, ParamId}; +use burn::store::RecordError; +use burn::tensor::{Bytes, Device, Tensor, TensorData}; +use hashbrown::HashMap; + +/// Scalar key (per parameter) under which the parameter's state rank is persisted. +/// +/// Reserved: a custom [`Optimizer::State`](crate::Optimizer::State) must not have a top-level +/// scalar field named `__rank`, as it would collide with this key in the record. +const RANK_KEY: &str = "__rank"; + +#[derive(Clone)] +struct OptimizerGroup { + group: ParamGroup, + optim: Arc, + grad_clipping: Option, +} + +impl OptimizerGroup { + pub(crate) fn set_gradient_clipping(&mut self, gradient_clipping: GradientClipping) { + self.grad_clipping = Some(gradient_clipping) + } +} + +/// Keep a reference to the optimizer to avoid matching every step. +#[derive(Clone)] +struct OptimizationContext { + optim: Arc, + grad_clipping: Option, + path: Option, + state: DynState, +} + +/// Optimizes a whole module by applying a per-parameter [`Optimizer`] to each of its parameters. +/// +/// It is non-generic over the module and optimizer: any `O: Optimizer` is type-erased behind a +/// dynamic optimizer, and per-parameter states are kept as type-erased states keyed by +/// [`ParamId`](burn::module::ParamId). Build one with `optimizer.into()` or +/// `OptimizerConfig::init()`. +/// +/// It is possible to use different optimizers for different parameters. To do so, use the +/// [ModuleOptimizer::with_group] function to add an optimizer for all parameters matching the +/// provided group. +#[derive(Clone)] +pub struct ModuleOptimizer { + optimizers: Vec, + param_context: HashMap, +} + +impl From for ModuleOptimizer +where + O: Optimizer, +{ + fn from(optim: O) -> Self { + Self { + param_context: HashMap::new(), + optimizers: vec![OptimizerGroup { + group: ParamGroup::all(), + optim: Arc::new(optim), + grad_clipping: None, + }], + } + } +} + +impl ModuleOptimizer { + /// Check if the optimizer has gradient clipping. + /// If there are multiple optimizers, checks if any group has gradient clipping. + pub fn has_gradient_clipping(&self) -> bool { + self.optimizers.iter().any(|g| g.grad_clipping.is_some()) + } + + /// Access the gradient clipping. + /// If there are multiple optimizers, returns the first optimizer's [GradientClipping]. + pub fn grad_clipping(&self) -> Option<&GradientClipping> { + self.optimizers + .first() + .expect("Should have at least one optimizer") + .grad_clipping + .as_ref() + } + + /// Sets the gradient clipping. + /// If there are multiple optimizers, assigns it to the first one. + /// + /// # Arguments + /// + /// * `gradient_clipping` - The gradient clipping. + /// + /// # Returns + /// + /// The optimizer. + pub fn with_grad_clipping(mut self, gradient_clipping: GradientClipping) -> Self { + self.optimizers + .first_mut() + .expect("Should have at least one optimizer") + .set_gradient_clipping(gradient_clipping); + self + } + + fn step_common( + &mut self, + lr_policy: ModuleLearningRate, + module: M, + mut grads: GradAdaptor, + ) -> M { + module.map(&mut ModuleOptimizerMapper::new( + self.optimizers.iter().collect(), + &mut self.param_context, + &mut grads, + lr_policy, + )) + } + + /// Adds an optimizer specific to a parameter group. + /// + /// Parameters matching this group will be optimized using the provided optimizer + /// and gradient clipping configuration. + /// + /// ### Matching Rules + /// * **Precedence:** If a parameter matches multiple groups, the *last* group added takes precedence. + /// * **Fallback:** The first optimizer added must match all parameters to act as a global fallback. + /// + /// ### Side Effects + /// * **State Reset:** Adding a new group will reset any existing optimizer states for parameters + /// that match the new group. + pub fn with_group( + mut self, + group: ParamGroup, + optim: O, + grad_clipping: Option, + ) -> Self + where + O: DynOptimizer + 'static, + { + self.optimizers.push(OptimizerGroup { + group: group.clone(), + optim: Arc::new(optim), + grad_clipping, + }); + self.param_context + .retain(|id, param_state| !group.matches(id, param_state.path.as_deref())); + self + } +} + +impl ModuleOptimizer { + /// Update the `module` parameters with the given `gradients`, advancing the optimizer state. + pub fn step( + &mut self, + lr_module: ModuleLearningRate, + module: M, + grads: GradientsParams, + ) -> M { + self.step_common(lr_module, module, grads.into()) + } + + /// Like [`step`](Self::step), but accumulating gradients sourced from multiple devices. + pub fn step_multi( + &mut self, + lr_module: ModuleLearningRate, + module: M, + grads: MultiGradientsParams, + ) -> M { + self.step_common(lr_module, module, grads.into()) + } + + fn optim_from_param( + &self, + id: ParamId, + path: Option<&str>, + ) -> (&'_ Arc, Option) { + self.optimizers + .iter() + .filter_map(|val| { + val.group + .matches(&id, path) + .then_some((&val.optim, val.grad_clipping.clone())) + }) + .next_back() + .expect("Should match at least one parameter group.") + } + + /// Decompose the optimizer state into a serializable [`OptimizerRecord`]. + pub fn to_record(&self) -> OptimizerRecord { + let mut tensors = Vec::new(); + let mut scalars = BTreeMap::new(); + let mut paths = BTreeMap::new(); + + for (id, param_state) in self.param_context.iter() { + let prefix = id.val().to_string(); + let mut sink = StateSink::default(); + param_state + .optim + .state_flatten(&prefix, ¶m_state.state, &mut sink); + + // Persist the parameter rank explicitly so the state can be reconstructed even when it + // carries no tensors, and without inferring the rank from tensor shapes. + scalars.insert( + join_path(&prefix, RANK_KEY), + burn_pack::Scalar::from(param_state.state.rank()), + ); + // Save parameter path to be able to match to the right group when loading. + if let Some(path) = ¶m_state.path { + paths.insert(prefix, path.clone()); + } + + for (name, data) in sink.tensors { + tensors.push(burn_pack::Tensor::new( + name, + data.dtype, + data.shape, + Some(id.val()), + data.bytes, + )); + } + for (name, value) in sink.scalars { + scalars.insert(name, value); + } + } + + OptimizerRecord { + tensors, + scalars, + paths, + } + } + + /// Load the optimizer state from an [`OptimizerRecord`]. + /// + /// State tensors are materialized on the default device; no device argument is needed because + /// each parameter's state is migrated to that parameter's (gradient's) device on the next + /// [`step`](ModuleOptimizer::step) — see the `to_device` call in the step path. The load device + /// is therefore irrelevant to correctness. + pub fn load_record(mut self, record: OptimizerRecord) -> Self { + let device = Device::default(); + let mut ranks: BTreeMap = BTreeMap::new(); + let mut paths: BTreeMap = BTreeMap::new(); + + // Recover each parameter's rank from its persisted `__rank` scalar (authoritative). Keys + // are `"{param_id}.__rank"`, so strip the dotted suffix to recover the id. + let suffix = alloc::format!(".{RANK_KEY}"); + for (name, value) in record.scalars.iter() { + if let Some(id_str) = name.strip_suffix(&suffix) + && let (Ok(id), Ok(rank)) = (id_str.parse::(), usize::try_from(*value)) + { + ranks.insert(id, rank); + } + } + + for (name, path) in record.paths.iter() { + if let Ok(id) = name.parse::() { + paths.insert(id, path.to_string()); + } + } + + let mut source = StateSource::new(record.scalars); + + for tensor in record.tensors { + let id = tensor + .param_id + .expect("Optimizer record tensors should carry a parameter id."); + let name = tensor.name; + let data = TensorData::from_bytes(tensor.bytes, tensor.shape, tensor.dtype); + // Fall back to inferring rank from a tensor shape if no `__rank` scalar was present. + ranks.entry(id).or_insert(data.shape.len()); + source.insert_tensor(name, data); + } + + let mut states = HashMap::new(); + for (id, rank) in ranks { + let prefix = id.to_string(); + let path = paths.get(&id); + let (optim, grad_clipping) = + self.optim_from_param(id.into(), path.map(|path| path.as_str())); + // Skip parameters whose state can't be reconstructed (truncated/foreign record); they + // are re-initialized lazily on the next step rather than aborting the load. + if let Some(state) = optim.state_unflatten(rank, &prefix, &mut source, &device) { + states.insert( + ParamId::from(id), + OptimizationContext { + optim: optim.clone(), + path: path.cloned(), + state, + grad_clipping, + }, + ); + } + } + + self.param_context = states; + self + } + + /// Serialize the optimizer state to an in-memory burnpack byte buffer. + pub fn into_bytes(&self) -> Result { + self.to_record().into_bytes() + } + + /// Load the optimizer state from an in-memory burnpack byte buffer. + pub fn from_bytes(self, bytes: Bytes) -> Result { + Ok(self.load_record(OptimizerRecord::from_bytes(bytes)?)) + } + + /// Save the optimizer state to a burnpack file on disk. + #[cfg(feature = "std")] + pub fn save>(&self, path: P) -> Result<(), RecordError> { + self.to_record().save(path) + } + + /// Load the optimizer state from a burnpack file on disk. + #[cfg(feature = "std")] + pub fn load>(self, path: P) -> Result { + Ok(self.load_record(OptimizerRecord::load(path)?)) + } +} + +/// Wrapper to unify the `remove` method for [GradientsParams] and [MultiGradientsParams]. +pub enum GradAdaptor { + /// Wrapper for [`GradientsParams`]. + Single(GradientsParams), + + /// Wrapper for [`MultiGradientsParams`]. + Multi(MultiGradientsParams), +} + +impl From for GradAdaptor { + fn from(grads: GradientsParams) -> Self { + Self::Single(grads) + } +} + +impl From for GradAdaptor { + fn from(grads: MultiGradientsParams) -> Self { + Self::Multi(grads) + } +} + +impl GradAdaptor { + /// Remove a gradient parameter by ID. + /// + /// # Returns + /// Maybe the (tensor, device) pair. + pub fn remove(&mut self, id: ParamId) -> Option<(Tensor, Device)> { + match self { + GradAdaptor::Single(grads) => grads.remove(id).map(|t| { + let device = t.device(); + (t, device) + }), + GradAdaptor::Multi(grads) => grads.remove(id), + } + } +} + +struct ModuleOptimizerMapper<'a> { + path: Vec, + optimizer_groups: Vec<&'a OptimizerGroup>, + states: &'a mut HashMap, + grads: &'a mut GradAdaptor, + lr_module: ModuleLearningRate, +} + +impl<'a> ModuleOptimizerMapper<'a> { + pub(crate) fn new( + optimizer_groups: Vec<&'a OptimizerGroup>, + states: &'a mut HashMap, + grads: &'a mut GradAdaptor, + lr_module: ModuleLearningRate, + ) -> Self { + Self { + path: vec![], + optimizer_groups, + states, + grads, + lr_module, + } + } + + fn optimizer_from_param( + &self, + id: ParamId, + path: Option<&str>, + ) -> (Arc, Option) { + self.optimizer_groups + .iter() + .filter_map(|val| { + val.group + .matches(&id, path) + .then_some((val.optim.clone(), val.grad_clipping.clone())) + }) + .next_back() + .expect("Should match at least one parameter group.") + } +} + +impl ModuleMapper for ModuleOptimizerMapper<'_> { + fn enter_module(&mut self, name: &str, _container_type: &str) { + self.path.push(name.to_string()); + } + + fn exit_module(&mut self, _name: &str, _container_type: &str) { + self.path.pop(); + } + + fn map_float(&mut self, param: Param>) -> Param> { + let (id, tensor, mapper) = param.consume(); + let grad = self.grads.remove(id); + + let tensor = if let Some((grad, device)) = grad { + let is_require_grad = tensor.is_require_grad(); + #[cfg(feature = "std")] + let is_distributed = tensor.is_distributed(); + + let entry = self.states.remove_entry(&id); + let key = entry.as_ref().map(|(k, _)| *k); + let tensor = if tensor.device() != device { + tensor.to_device(&device) + } else { + tensor + }; + + let path = self.path.join("."); + let (optim, grad_clipping, existing_dyn_state) = match entry.map(|(_, s)| s) { + Some(OptimizationContext { + optim, + grad_clipping, + state, + .. + }) => (optim, grad_clipping, Some(state)), + None => { + let (optim, grad_clipping) = + self.optimizer_from_param(id, Some(path.as_str())).clone(); + (optim, grad_clipping, None) + } + }; + + debug_assert_eq!( + grad.device(), + device, + "The gradient is on the provided device" + ); + let clipped_grad: Tensor = if let Some(g_clipping) = grad_clipping.as_ref() { + g_clipping.clip_gradient(grad) + } else { + grad + }; + + debug_assert_eq!( + tensor.device(), + device, + "Tensor and gradients are on the same device." + ); + + let lr = self.lr_module.lr_from_param(id, Some(path.as_str())); + let (tensor, state) = optim.step_dyn( + D, + lr, + tensor.inner().into_bridge(), + clipped_grad.into_bridge(), + existing_dyn_state.map(|s| optim.to_device_dyn(s, &device)), + ); + + if let Some(state) = state { + self.states.insert( + key.unwrap_or(id), + OptimizationContext { + optim, + path: Some(path), + state, + grad_clipping, + }, + ); + } + + let mut tensor = Tensor::from_inner(Tensor::from_bridge(tensor)); + + if is_require_grad { + tensor = tensor.require_grad(); + } + #[cfg(feature = "std")] + if is_distributed { + tensor = tensor.set_distributed(id) + } + + tensor + } else { + tensor + }; + + Param::from_mapped_value(id, tensor, mapper) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + AdamConfig, GradientsParams, SgdConfig, + lr_scheduler::module_lr_scheduler::ModuleLearningRate, + }; + use burn::module::ParamGroup; + use burn::tensor::{Distribution, Tensor, Tolerance}; + use burn_derive::Module; + use burn_nn::{Linear, LinearConfig}; + + #[derive(Module, Debug)] + struct TwoLayerModel { + layer_a: Linear, + layer_b: Linear, + } + + fn make_model(device: &Device) -> TwoLayerModel { + TwoLayerModel { + layer_a: LinearConfig::new(4, 4).init(device), + layer_b: LinearConfig::new(4, 4).init(device), + } + } + + fn make_grads(model: &TwoLayerModel, x: Tensor<2>) -> GradientsParams { + let out = model.layer_a.forward(x.clone()) + model.layer_b.forward(x); + GradientsParams::from_grads(out.mean().backward(), model) + } + + fn lr() -> ModuleLearningRate { + ModuleLearningRate::from(0.01_f64) + } + + fn sgd() -> ModuleOptimizer { + ModuleOptimizer::from(SgdConfig::new().init()) + } + + /// to_record / load_record must fully preserve a stateful optimizer's internal state so that + /// a step on the restored optimizer is numerically identical to one taken on the original. + #[test] + fn default_optimizer_state_survives_round_trip() { + let device = Device::default().autodiff(); + let mut model = make_model(&device); + let mut optim: ModuleOptimizer = AdamConfig::new().init(); + + for _ in 0..3 { + let x = Tensor::<2>::random([2, 4], Distribution::Default, &device); + model = optim.step(lr(), model.clone(), make_grads(&model, x)); + } + + let record = optim.to_record(); + let mut reloaded: ModuleOptimizer = AdamConfig::new().init().load_record(record); + + let x = Tensor::<2>::random([2, 4], Distribution::Default, &device); + let grads_a = make_grads(&model, x.clone()); + let grads_b = make_grads(&model, x); + let from_orig = optim.step(lr(), model.clone(), grads_a); + let from_reload = reloaded.step(lr(), model, grads_b); + + from_orig + .layer_a + .weight + .val() + .into_data() + .assert_approx_eq::( + &from_reload.layer_a.weight.val().into_data(), + Tolerance::absolute(1e-6), + ); + } + + /// The paths saved in OptimizerRecord enable load_record to route each parameter to the + /// correct group optimizer. A step on the restored optimizer must be numerically identical + /// to one on the original — for both the group (Adam) and the default (SGD) optimizer. + #[test] + fn group_optimizer_routes_correctly_after_record_round_trip() { + let device = Device::default().autodiff(); + let mut model = make_model(&device); + + let make_optim = || { + sgd().with_group( + ParamGroup::from_predicate("layer_a"), + AdamConfig::new().build(), + None, + ) + }; + let mut optim = make_optim(); + + for _ in 0..3 { + let x = Tensor::<2>::random([2, 4], Distribution::Default, &device); + model = optim.step(lr(), model.clone(), make_grads(&model, x)); + } + + let record = optim.to_record(); + let mut reloaded = make_optim().load_record(record); + + let x = Tensor::<2>::random([2, 4], Distribution::Default, &device); + let grads_a = make_grads(&model, x.clone()); + let grads_b = make_grads(&model, x); + let from_orig = optim.step(lr(), model.clone(), grads_a); + let from_reload = reloaded.step(lr(), model, grads_b); + + from_orig + .layer_a + .weight + .val() + .into_data() + .assert_approx_eq::( + &from_reload.layer_a.weight.val().into_data(), + Tolerance::absolute(1e-6), + ); + from_orig + .layer_b + .weight + .val() + .into_data() + .assert_approx_eq::( + &from_reload.layer_b.weight.val().into_data(), + Tolerance::absolute(1e-6), + ); + } + + /// Adding a group after training must clear accumulated state for matching params + /// while leaving state for non-matching params untouched. + #[test] + fn with_group_clears_state_for_matching_params() { + let device = Device::default().autodiff(); + let mut model = make_model(&device); + let mut optim: ModuleOptimizer = AdamConfig::new().init(); + + for _ in 0..2 { + let x = Tensor::<2>::random([2, 4], Distribution::Default, &device); + model = optim.step(lr(), model.clone(), make_grads(&model, x)); + } + + // Switch layer_a to SGD. Its Adam state must be cleared. + optim = optim.with_group( + ParamGroup::from_predicate("layer_a"), + SgdConfig::new().build(), + None, + ); + + let x = Tensor::<2>::random([2, 4], Distribution::Default, &device); + _ = optim.step(lr(), model.clone(), make_grads(&model, x)); + + let record = optim.to_record(); + + let time_key_count = record.scalars.keys().filter(|k| k.contains("time")).count(); + assert_eq!( + time_key_count, 2, + "only layer_b params should carry Adam's time scalar after group switch" + ); + } + + #[test] + fn group_grad_clipping_applies_only_to_matching_params() { + let device = Device::default().autodiff(); + let mut model = make_model(&device); + + let lr_value = 0.01; + let threshold = 0.01_f32; + let bound = (lr_value * threshold as f64) as f32; + + let mut optim = sgd().with_group( + ParamGroup::from_predicate("layer_a"), + SgdConfig::new().build(), + Some(GradientClipping::Value(threshold)), + ); + + // Large inputs produce gradients that clearly exceed the clipping threshold. + let x = Tensor::<2>::random([2, 4], Distribution::Uniform(100.0, 200.0), &device); + let weight_a_before = model.layer_a.weight.val(); + let weight_b_before = model.layer_b.weight.val(); + + model = optim.step( + ModuleLearningRate::from(lr_value), + model.clone(), + make_grads(&model, x), + ); + + let diff_a = weight_a_before - model.layer_a.weight.val(); + let diff_b = weight_b_before - model.layer_b.weight.val(); + + let eps = 1e-6_f32; + for value in diff_a.into_data().iter::() { + assert!( + value.abs() <= bound + eps, + "layer_a's group grad_clipping should keep every update within lr * threshold" + ); + } + assert!( + diff_b + .into_data() + .iter::() + .any(|value| value.abs() > bound + eps), + "layer_b uses the default optimizer and should not be clipped" + ); + } + + /// An OptimizerRecord with empty paths map must load cleanly. + /// Parameters default to the default optimizer. + #[test] + fn record_without_paths_loads_without_panic() { + let device = Device::default().autodiff(); + let mut model = make_model(&device); + let mut optim: ModuleOptimizer = AdamConfig::new().init(); + + let x = Tensor::<2>::random([2, 4], Distribution::Default, &device); + model = optim.step(lr(), model.clone(), make_grads(&model, x)); + + // Simulate a record with empty paths. + let mut record = optim.to_record(); + record.paths.clear(); + + let mut optim_loaded: ModuleOptimizer = AdamConfig::new().init().load_record(record); + + let x = Tensor::<2>::random([2, 4], Distribution::Default, &device); + let grads_a = make_grads(&model, x.clone()); + let grads_b = make_grads(&model, x); + let from_orig = optim.step(lr(), model.clone(), grads_a); + let from_reload = optim_loaded.step(lr(), model, grads_b); + + from_orig + .layer_a + .weight + .val() + .into_data() + .assert_approx_eq::( + &from_reload.layer_a.weight.val().into_data(), + Tolerance::absolute(1e-6), + ); + from_orig + .layer_b + .weight + .val() + .into_data() + .assert_approx_eq::( + &from_reload.layer_b.weight.val().into_data(), + Tolerance::absolute(1e-6), + ); + } +} diff --git a/crates/burn-optim/src/optim/module/record/mod.rs b/crates/burn-optim/src/optim/module/record/mod.rs new file mode 100644 index 0000000..4f9bbdd --- /dev/null +++ b/crates/burn-optim/src/optim/module/record/mod.rs @@ -0,0 +1,92 @@ +use alloc::collections::BTreeMap; +use alloc::string::String; +use alloc::vec::Vec; + +use burn::store::RecordError; +use burn::tensor::Bytes; +use burn_core as burn; + +use burn_pack::{Reader, Scalar, Writer}; + +/// A serialized optimizer state, stored in the [burnpack](burn_pack) format. +/// +/// Unlike a module record (keyed by module path), an optimizer record is keyed per parameter: +/// each parameter's state is decomposed into tensors named `"{param_id}.{field}"` (carrying the +/// originating `param_id`) plus a few typed scalar entries kept in the burnpack scalar map. +/// +/// Obtain one from a [`ModuleOptimizer`](crate::optim::ModuleOptimizer) with +/// [`to_record`](crate::optim::ModuleOptimizer::to_record), then save it +/// ([`save`](Self::save) / [`into_bytes`](Self::into_bytes)) or apply it back with +/// [`load_record`](crate::optim::ModuleOptimizer::load_record). +#[derive(Default)] +pub struct OptimizerRecord { + pub(crate) tensors: Vec, + pub(crate) scalars: BTreeMap, + pub(crate) paths: BTreeMap, +} + +impl core::fmt::Debug for OptimizerRecord { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("OptimizerRecord") + .field("num_tensors", &self.tensors.len()) + .field("num_scalars", &self.scalars.len()) + .finish() + } +} + +impl OptimizerRecord { + /// The number of tensors in the record. + pub fn len(&self) -> usize { + self.tensors.len() + } + + /// Whether the record holds no tensors. + pub fn is_empty(&self) -> bool { + self.tensors.is_empty() + } + + /// Serialize the record to an in-memory burnpack byte buffer. + pub fn into_bytes(self) -> Result { + Ok(self.into_writer().into_bytes()?) + } + + /// Reconstruct a record from an in-memory burnpack byte buffer. + pub fn from_bytes(bytes: Bytes) -> Result { + Self::from_reader(Reader::from_bytes(bytes)?) + } + + /// Save the record to a burnpack file on disk. + #[cfg(feature = "std")] + pub fn save>(self, path: P) -> Result<(), RecordError> { + self.into_writer().write_to_file(path)?; + Ok(()) + } + + /// Load a record from a burnpack file on disk. + #[cfg(feature = "std")] + pub fn load>(path: P) -> Result { + Self::from_reader(Reader::from_file(path)?) + } + + fn into_writer(self) -> Writer { + let mut writer = Writer::new(self.tensors); + for (key, value) in &self.scalars { + writer = writer.with_scalar(key, *value); + } + for (key, value) in &self.paths { + writer = writer.with_metadata(key, value); + } + writer + } + + fn from_reader(reader: Reader) -> Result { + let scalars = reader.scalars().clone(); + let paths = reader.metadata().clone(); + let tensors = reader.into_tensors()?; + Ok(Self { + tensors, + scalars, + paths, + }) + } +} diff --git a/crates/burn-optim/src/optim/momentum.rs b/crates/burn-optim/src/optim/momentum.rs new file mode 100644 index 0000000..ac0163f --- /dev/null +++ b/crates/burn-optim/src/optim/momentum.rs @@ -0,0 +1,94 @@ +use burn_core as burn; + +use crate::RecordState; +use burn::config::Config; +use burn::tensor::Device; +use burn::tensor::{ElementConversion, Tensor}; + +/// Configuration to create [momentum](Momentum). +#[derive(Config, Debug)] +pub struct MomentumConfig { + /// Momentum factor + #[config(default = 0.9)] + pub momentum: f64, + /// Dampening factor. + #[config(default = 0.1)] + pub dampening: f64, + /// Enables Nesterov momentum, see [On the importance of initialization and + /// momentum in deep learning](http://www.cs.toronto.edu/~hinton/absps/momentum.pdf). + #[config(default = false)] + pub nesterov: bool, +} + +/// State of [momentum](Momentum). +#[derive(RecordState, Clone, new)] +pub struct MomentumState { + velocity: Tensor, +} + +/// Momentum implementation that transforms gradients. +#[derive(Clone)] +pub struct Momentum { + momentum: f32, + dampening: f64, + nesterov: bool, +} + +impl Momentum { + /// Creates a new [momentum](Momentum) from a [config](MomentumConfig). + pub fn new(config: &MomentumConfig) -> Self { + Self { + momentum: config.momentum.elem(), + dampening: config.dampening, + nesterov: config.nesterov, + } + } + + /// Transforms a gradient. + /// + /// # Arguments + /// + /// * `grad` - Gradient to transform. + /// * `state` - State of the optimizer. + /// + /// # Returns + /// + /// * `grad` - Transformed gradient. + /// * `state` - State of the optimizer. + pub fn transform( + &self, + grad: Tensor, + state: Option>, + ) -> (Tensor, MomentumState) { + let velocity = if let Some(state) = state { + grad.clone() + .mul_scalar(1.0 - self.dampening) + .add(state.velocity.mul_scalar(self.momentum)) + } else { + grad.clone() + }; + + let grad = match self.nesterov { + true => velocity.clone().mul_scalar(self.momentum).add(grad), + false => velocity.clone(), + }; + + (grad, MomentumState::new(velocity)) + } +} + +impl MomentumState { + /// Moves the state to a device. + /// + /// # Arguments + /// + /// * `device` - Device to move the state to. + /// + /// # Returns + /// + /// * `self` - Moved state. + pub fn to_device(mut self, device: &Device) -> Self { + self.velocity = self.velocity.to_device(device); + self + } +} diff --git a/crates/burn-optim/src/optim/muon.rs b/crates/burn-optim/src/optim/muon.rs new file mode 100644 index 0000000..65ec271 --- /dev/null +++ b/crates/burn-optim/src/optim/muon.rs @@ -0,0 +1,762 @@ +use burn_core as burn; + +use crate::RecordState; + +use burn::config::Config; +use burn::tensor::Device; +use burn::tensor::Tensor; +use serde::{Deserialize, Serialize}; + +use super::{ + Optimizer, + decay::WeightDecayConfig, + module_optimizer::ModuleOptimizer, + momentum::{Momentum, MomentumConfig, MomentumState}, +}; +use crate::LearningRate; + +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float as _; + +/// Learning rate adjustment method for Muon optimizer. +/// +/// Muon adjusts the learning rate based on parameter shape to maintain consistent +/// RMS across rectangular matrices. +/// +/// # References +/// +/// - Original: [Muon: An optimizer for hidden layers](https://kellerjordan.github.io/posts/muon/) +/// - Moonshot: [Muon is Scalable for LLM Training](https://arxiv.org/pdf/2502.16982) +#[derive(Clone, Default, Debug, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum AdjustLrFn { + /// Keller Jordan's original method: `lr * sqrt(max(1, A/B))` + /// + /// This scales the learning rate based on the aspect ratio of the weight matrix, + /// ensuring that tall matrices (more rows than columns) get proportionally larger + /// learning rates. + /// + /// # Example + /// + /// For a [1024, 512] matrix: `lr * sqrt(1024/512) = lr * 1.414` + #[default] + Original, + + /// Moonshot's method: `lr * 0.2 * sqrt(max(A, B))` + /// + /// This method is designed to match AdamW's RMS, allowing Muon to directly reuse + /// learning rates and weight decay values tuned for AdamW without retuning. + /// + /// # Example + /// + /// For a [1024, 512] matrix: `lr * 0.2 * sqrt(1024) = lr * 6.4` + MatchRmsAdamW, +} + +impl AdjustLrFn { + /// Calculate the learning rate adjustment ratio for a given parameter shape. + /// + /// # Arguments + /// + /// * `shape` - Parameter shape (uses first two dimensions) + /// + /// # Returns + /// + /// Adjustment ratio to multiply with the base learning rate + fn adjustment_ratio(&self, shape: &[usize]) -> f64 { + if shape.len() < 2 { + return 1.0; + } + + let a = shape[0] as f64; + let b = shape[1] as f64; + + match self { + Self::Original => { + // sqrt(max(1, A/B)) + let ratio = a / b; + ratio.max(1.0).sqrt() + } + Self::MatchRmsAdamW => { + // 0.2 * sqrt(max(A, B)) + 0.2 * a.max(b).sqrt() + } + } + } +} + +/// Muon configuration. +/// +/// Muon is an optimizer specifically designed for 2D parameters of neural network +/// hidden layers (weight matrices). Other parameters such as biases and embeddings +/// should be optimized using a standard method such as AdamW. +/// +/// # Learning Rate Adjustment +/// +/// Muon adjusts the learning rate based on parameter shape to maintain consistent +/// RMS across rectangular matrices. Two methods are available: +/// +/// - **Original**: Uses `sqrt(max(1, A/B))` where A and B are the first two dimensions. +/// This is Keller Jordan's method and is the default. +/// +/// - **MatchRmsAdamW**: Uses `0.2 * sqrt(max(A, B))`. This is Moonshot's method +/// designed to match AdamW's RMS, allowing direct reuse of AdamW hyperparameters. +/// +/// # Example +/// +/// ```ignore +/// use burn_optim::{MuonConfig, AdjustLrFn}; +/// +/// // Using default (Original) method +/// let optimizer = MuonConfig::new().init(); +/// +/// // Using MatchRmsAdamW for AdamW-compatible hyperparameters +/// let optimizer = MuonConfig::new() +/// .with_adjust_lr_fn(AdjustLrFn::MatchRmsAdamW) +/// .init(); +/// ``` +/// +/// # References +/// +/// - [Muon: An optimizer for hidden layers in neural networks](https://kellerjordan.github.io/posts/muon/) +/// - [Muon is Scalable for LLM Training](https://arxiv.org/pdf/2502.16982) +/// - [PyTorch Implementation](https://github.com/pytorch/pytorch/blob/main/torch/optim/muon.py) +/// - [Original Implementation](https://github.com/KellerJordan/Muon) +#[derive(Config, Debug)] +pub struct MuonConfig { + /// [Weight decay](WeightDecayConfig) config. + weight_decay: Option, + + /// [Momentum](MomentumConfig) config. + /// + /// Muon always uses momentum. Default configuration: + /// - momentum: 0.95 + /// - dampening: 0.0 + /// - nesterov: true + #[config(default = "MomentumConfig { momentum: 0.95, dampening: 0.0, nesterov: true }")] + momentum: MomentumConfig, + + /// Newton-Schulz iteration coefficients (a, b, c). + /// + /// These coefficients are selected to maximize the slope at zero for the + /// quintic iteration. Default values are from Keller Jordan's implementation. + #[config(default = "(3.4445, -4.775, 2.0315)")] + ns_coefficients: (f32, f32, f32), + + /// Epsilon for numerical stability. + #[config(default = 1e-7)] + epsilon: f32, + + /// Number of Newton-Schulz iteration steps. + #[config(default = 5)] + ns_steps: usize, + + /// Learning rate adjustment method. + /// + /// Controls how the learning rate is adjusted based on parameter shape. + /// See [`AdjustLrFn`] for available methods. + #[config(default = "AdjustLrFn::Original")] + adjust_lr_fn: AdjustLrFn, +} + +impl MuonConfig { + /// Build a [`Muon`] from the config. + pub(crate) fn build(&self) -> Muon { + let momentum = Momentum::new(&self.momentum); + let weight_decay_penalty = self.weight_decay.as_ref().map(|wd| wd.penalty); + + Muon { + momentum, + ns_params: NewtonSchulzParams::new(self.ns_coefficients, self.ns_steps), + weight_decay_penalty, + epsilon: self.epsilon, + adjust_lr_fn: self.adjust_lr_fn, + } + } + + /// Initialize Muon optimizer. + /// + /// # Returns + /// + /// Returns an optimizer adaptor that can be used to optimize a module. + /// + /// # Example + /// + /// ```ignore + /// use burn_optim::{MuonConfig, AdjustLrFn, decay::WeightDecayConfig}; + /// + /// // Basic configuration with default (Original) LR adjustment + /// let optimizer = MuonConfig::new() + /// .with_weight_decay(Some(WeightDecayConfig::new(0.01))) + /// .init(); + /// + /// // With AdamW-compatible settings using MatchRmsAdamW + /// let optimizer = MuonConfig::new() + /// .with_adjust_lr_fn(AdjustLrFn::MatchRmsAdamW) + /// .with_weight_decay(Some(WeightDecayConfig::new(0.1))) + /// .init(); + /// + /// // Custom momentum and NS settings + /// let optimizer = MuonConfig::new() + /// .with_momentum(MomentumConfig { + /// momentum: 0.9, + /// dampening: 0.1, + /// nesterov: false, + /// }) + /// .with_ns_steps(7) + /// .init(); + /// ``` + pub fn init(&self) -> ModuleOptimizer { + ModuleOptimizer::from(self.build()) + } +} + +/// Parameters for Newton-Schulz orthogonalization. +#[derive(Clone, Copy)] +struct NewtonSchulzParams { + a: f32, + b: f32, + c: f32, + steps: usize, +} + +impl NewtonSchulzParams { + fn new(coefficients: (f32, f32, f32), steps: usize) -> Self { + Self { + a: coefficients.0, + b: coefficients.1, + c: coefficients.2, + steps, + } + } +} + +/// Muon optimizer. +/// +/// Muon internally runs standard SGD-momentum, and then performs an orthogonalization +/// post-processing step, in which each 2D parameter's update is replaced with the +/// nearest orthogonal matrix. For efficient orthogonalization we use a Newton-Schulz +/// iteration, which has the advantage that it can be stably run in bfloat16 on the GPU. +/// +/// # Important Notes +/// +/// 1. **Only for 2D+ parameters**: Muon is designed for weight matrices. Use AdamW +/// or SGD for biases, embeddings, and layer norms. +/// +/// 2. **Learning rate adjustment**: Muon automatically adjusts the learning rate based +/// on parameter shape. See [`AdjustLrFn`] for details. +/// +/// 3. **Weight decay timing**: Unlike typical optimizers, Muon applies weight decay +/// AFTER orthogonalization but uses the original (unadjusted) learning rate for it. +#[derive(Clone)] +pub struct Muon { + momentum: Momentum, + ns_params: NewtonSchulzParams, + weight_decay_penalty: Option, + epsilon: f32, + adjust_lr_fn: AdjustLrFn, +} + +impl Muon { + /// Adjust learning rate based on parameter shape. + /// + /// # Arguments + /// + /// * `lr` - Base learning rate + /// * `shape` - Parameter shape (uses first two dimensions) + /// + /// # Returns + /// + /// Adjusted learning rate + /// + /// ```ignore + /// // For a [1024, 512] weight matrix with lr=0.01: + /// // Original: 0.01 * sqrt(1024/512) = 0.01 * 1.414 = 0.01414 + /// // MatchRmsAdamW: 0.01 * 0.2 * sqrt(1024) = 0.01 * 0.2 * 32 = 0.064 + /// ``` + fn adjust_lr(&self, lr: LearningRate, shape: &[usize]) -> LearningRate { + lr * self.adjust_lr_fn.adjustment_ratio(shape) + } + + /// Perform Newton-Schulz orthogonalization on a gradient tensor. + /// + /// This computes the zeroth power (orthogonalization) of the input matrix G + /// using a quintic Newton-Schulz iteration. + /// + /// # Algorithm + /// + /// 1. Transpose if tall matrix (A > B) + /// 2. Normalize: X = X / ||X|| + /// 3. For k steps: + /// - A = X @ X^T + /// - B = b*A + c*A^2 + /// - X = a*X + B@X + /// 4. Transpose back if needed + /// + /// # References + /// + /// - Original: https://github.com/KellerJordan/Muon/blob/master/muon.py + /// - PyTorch: https://github.com/pytorch/pytorch/blob/main/torch/optim/muon.py + fn zeropower_via_newtonschulz(&self, g: Tensor) -> Tensor { + let shape = g.shape(); + let dim_m2 = shape[D - 2]; + let dim_m1 = shape[D - 1]; + + // Step 1: Transpose if tall matrix (more rows than columns) + let (mut x, needs_transpose) = if dim_m2 > dim_m1 { + (g.swap_dims(D - 2, D - 1), true) + } else { + (g, false) + }; + + // Step 2: Normalize by Frobenius norm + // X = X / (||X|| + epsilon) + let norm = x + .clone() + .powf_scalar(2.0) + .sum() + .sqrt() + .clamp_min(self.epsilon) + .unsqueeze(); + + x = x.div(norm); + + // Step 3: Newton-Schulz iteration + // This is the quintic iteration with coefficients (a, b, c) + let NewtonSchulzParams { a, b, c, steps } = self.ns_params; + + for _ in 0..steps { + // A = X @ X^T + let x_t = x.clone().swap_dims(D - 2, D - 1); + let a_matrix = x.clone().matmul(x_t); + + // B = b*A + c*A@A + let a_squared = a_matrix.clone().matmul(a_matrix.clone()); + let b_matrix = a_matrix.mul_scalar(b).add(a_squared.mul_scalar(c)); + + // X = a*X + B@X + x = x.clone().mul_scalar(a).add(b_matrix.matmul(x.clone())); + } + + // Step 4: Restore transpose if it was a tall matrix + if needs_transpose { + x = x.swap_dims(D - 2, D - 1); + } + + x + } +} + +/// Muon state. +#[derive(RecordState, Clone, new)] +pub struct MuonState { + /// Current momentum state + pub momentum: MomentumState, +} + +impl Optimizer for Muon { + type State = MuonState; + + /// Perform a single Muon optimization step. + /// + /// # Algorithm + /// + /// 1. Apply momentum to gradient + /// 2. Orthogonalize update via Newton-Schulz + /// 3. Adjust learning rate based on parameter shape + /// 4. Apply weight decay (using original lr) + /// 5. Update parameter (using adjusted lr) + /// + /// # Notes + /// + /// Unlike typical optimizers, the weight decay and parameter update use + /// different learning rates: + /// - Weight decay uses the original `lr` + /// - Parameter update uses the shape-adjusted `lr` + /// + /// # Panics + /// This function will panic if the input tensors are not 2D. + fn step( + &self, + lr: LearningRate, + tensor: Tensor, + grad: Tensor, + state: Option>, + ) -> (Tensor, Option>) { + assert!( + D == 2, + "Newton-Schulz iteration requires 2D tensors, got {}D", + D + ); + + // Step 1: Apply momentum + let state_momentum = state.map(|s| s.momentum); + let (grad, new_momentum_state) = self.momentum.transform(grad, state_momentum); + + // Step 2: Orthogonalize via Newton-Schulz + let update = self.zeropower_via_newtonschulz(grad); + + // Step 3: Adjust learning rate based on parameter shape + let adjusted_lr = self.adjust_lr(lr, &tensor.shape()); + + // Step 4: Apply weight decay (using ORIGINAL lr, not adjusted) + // Muon applies weight decay AFTER orthogonalization + let tensor = if let Some(penalty) = self.weight_decay_penalty { + let decay_factor = 1.0 - lr * penalty as f64; + tensor.mul_scalar(decay_factor) + } else { + tensor + }; + + // Step 5: Update parameter (using ADJUSTED lr) + let delta = update.mul_scalar(adjusted_lr); + let new_state = MuonState::new(new_momentum_state); + + (tensor - delta, Some(new_state)) + } + + fn to_device(mut state: Self::State, device: &Device) -> Self::State { + state.momentum = state.momentum.to_device(device); + state + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{GradientsParams, Optimizer}; + use burn::module::Param; + use burn::tensor::{Distribution, Tensor, TensorData}; + use burn_nn::{Linear, LinearConfig}; + + const TOLERANCE: f64 = 1e-8; + + fn given_linear_layer_no_bias(weight: TensorData) -> Linear { + let device = Device::default().autodiff(); + Linear { + weight: Param::from_data(weight, &device), + bias: None, // No bias for Muon optimizer + } + } + + #[test] + fn test_adjust_lr_fn_original() { + let method = AdjustLrFn::Original; + + // Square matrix [512, 512] -> sqrt(1) = 1.0 + let ratio = method.adjustment_ratio(&[512, 512]); + assert!((ratio - 1.0).abs() < TOLERANCE); + + // Tall matrix [1024, 512] -> sqrt(2) ≈ 1.414 + let ratio = method.adjustment_ratio(&[1024, 512]); + let expected = (2.0f64).sqrt(); + assert!((ratio - expected).abs() < TOLERANCE); + + // Wide matrix [512, 1024] -> max(1, 0.5) = 1.0 + let ratio = method.adjustment_ratio(&[512, 1024]); + assert!((ratio - 1.0).abs() < TOLERANCE); + } + + #[test] + fn test_adjust_lr_fn_match_rms_adamw() { + let method = AdjustLrFn::MatchRmsAdamW; + + // [1024, 512] -> 0.2 * sqrt(1024) = 6.4 + let ratio = method.adjustment_ratio(&[1024, 512]); + let expected = 0.2 * 1024.0f64.sqrt(); + assert!((ratio - expected).abs() < TOLERANCE); + + // [512, 512] -> 0.2 * sqrt(512) ≈ 4.525 + let ratio = method.adjustment_ratio(&[512, 512]); + let expected = 0.2 * 512.0f64.sqrt(); + assert!((ratio - expected).abs() < TOLERANCE); + } + + #[test] + #[should_panic(expected = "Newton-Schulz iteration requires 2D tensors, got 1D")] + fn test_1d_tensor_panics() { + let device = Default::default(); + let config = MuonConfig::new(); + let optim = Muon { + momentum: Momentum::new(&config.momentum), + ns_params: NewtonSchulzParams::new(config.ns_coefficients, config.ns_steps), + weight_decay_penalty: None, + epsilon: config.epsilon, + adjust_lr_fn: config.adjust_lr_fn, + }; + + let tensor_1d = Tensor::<1>::zeros([512], &device); + let grad_1d = Tensor::<1>::ones([512], &device); + + let _ = optim.step(0.01, tensor_1d, grad_1d, None); + } + + #[test] + fn test_muon_optimizer_save_load_state() { + let device = Device::default().autodiff(); + // Use Linear layer WITHOUT bias for Muon optimizer + let linear = LinearConfig::new(6, 6) + .with_bias(false) // No bias - only 2D weight matrix + .init(&device); + + let x = Tensor::<2>::random([2, 6], Distribution::Default, &device); + + let mut optimizer = MuonConfig::new().init(); + let grads = linear.forward(x).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let _linear = optimizer.step(0.01.into(), linear, grads); + + let state_before = optimizer.to_record(); + let bytes = optimizer.into_bytes().unwrap(); + + let optimizer_loaded = MuonConfig::new().init().from_bytes(bytes).unwrap(); + let state_after = optimizer_loaded.to_record(); + + assert_eq!(state_before.len(), state_after.len()); + } + + #[test] + fn test_muon_with_weight_decay() { + let device = Device::default().autodiff(); + // Create Linear layer WITHOUT bias for Muon + let linear = given_linear_layer_no_bias(TensorData::from([ + [1.0, 1.0, 1.0, 1.0], + [1.0, 1.0, 1.0, 1.0], + [1.0, 1.0, 1.0, 1.0], + [1.0, 1.0, 1.0, 1.0], + ])); + + let x = Tensor::<2>::from_floats([[0.5, 0.5, 0.5, 0.5], [0.5, 0.5, 0.5, 0.5]], &device) + .require_grad(); + + let mut optimizer = MuonConfig::new() + .with_weight_decay(Some(WeightDecayConfig::new(0.01))) + .init(); + + let grads = linear.forward(x).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(0.01.into(), linear, grads); + + let state = linear; + let weight = state.weight.to_data(); + + for val in weight.as_slice::().unwrap() { + assert!( + *val < 1.0, + "Weight should be reduced by weight decay, got {}", + val + ); + } + } + + #[test] + fn test_newton_schulz_orthogonalization() { + let device = Default::default(); + let matrix = Tensor::<2>::from_floats([[1.0, 0.5], [0.5, 1.0]], &device); + + let config = MuonConfig::new(); + let muon = Muon { + momentum: Momentum::new(&config.momentum), + ns_params: NewtonSchulzParams::new(config.ns_coefficients, config.ns_steps), + weight_decay_penalty: None, + epsilon: config.epsilon, + adjust_lr_fn: config.adjust_lr_fn, + }; + + let orthogonalized = muon.zeropower_via_newtonschulz(matrix); + let o_t = orthogonalized.clone().transpose(); + let product = orthogonalized.matmul(o_t); + + let data = product.into_data(); + let values = data.as_slice::().unwrap(); + + assert!( + (values[0] - 1.0).abs() < 0.1, + "Product[0,0] should be ~1.0, got {}", + values[0] + ); + assert!( + (values[3] - 1.0).abs() < 0.1, + "Product[1,1] should be ~1.0, got {}", + values[3] + ); + } + + #[test] + fn test_tall_matrix_transpose() { + // Test that tall matrices (A > B) are transposed during Newton-Schulz iteration + // and then transposed back + let device = Default::default(); + + // Create a tall matrix: [8, 4] (more rows than columns) + let tall_matrix = Tensor::<2>::from_floats( + [ + [1.0, 0.5, 0.3, 0.2], + [0.5, 1.0, 0.4, 0.1], + [0.3, 0.4, 1.0, 0.5], + [0.2, 0.1, 0.5, 1.0], + [0.1, 0.2, 0.3, 0.4], + [0.4, 0.3, 0.2, 0.1], + [0.2, 0.4, 0.1, 0.3], + [0.3, 0.1, 0.4, 0.2], + ], + &device, + ); + + let config = MuonConfig::new(); + let muon = Muon { + momentum: Momentum::new(&config.momentum), + ns_params: NewtonSchulzParams::new(config.ns_coefficients, config.ns_steps), + weight_decay_penalty: None, + epsilon: config.epsilon, + adjust_lr_fn: config.adjust_lr_fn, + }; + + // Perform Newton-Schulz orthogonalization + let orthogonalized = muon.zeropower_via_newtonschulz(tall_matrix.clone()); + + // Verify shape is preserved (should be transposed internally but returned in original shape) + let original_shape = tall_matrix.shape(); + let result_shape = orthogonalized.shape(); + assert_eq!( + original_shape.dims::<2>(), + result_shape.dims::<2>(), + "Shape should be preserved: [8, 4]" + ); + + // Verify output is different from input (orthogonalization happened) + let original_data = tall_matrix.into_data(); + let result_data = orthogonalized.into_data(); + assert_ne!( + original_data.as_slice::().unwrap(), + result_data.as_slice::().unwrap(), + "Orthogonalized matrix should differ from input" + ); + + // For comparison, test a wide matrix [4, 8] should NOT be transposed + let wide_matrix = Tensor::<2>::from_floats( + [ + [1.0, 0.5, 0.3, 0.2, 0.1, 0.4, 0.2, 0.3], + [0.5, 1.0, 0.4, 0.1, 0.2, 0.3, 0.4, 0.1], + [0.3, 0.4, 1.0, 0.5, 0.3, 0.2, 0.1, 0.4], + [0.2, 0.1, 0.5, 1.0, 0.4, 0.1, 0.3, 0.2], + ], + &device, + ); + + let orthogonalized_wide = muon.zeropower_via_newtonschulz(wide_matrix.clone()); + + // Verify wide matrix shape is also preserved + let wide_original_shape = wide_matrix.shape(); + let wide_result_shape = orthogonalized_wide.shape(); + assert_eq!( + wide_original_shape.dims::<2>(), + wide_result_shape.dims::<2>(), + "Wide matrix shape should be preserved: [4, 8]" + ); + } + + #[test] + fn test_zero_gradient() { + // Test that Muon handles zero gradients gracefully + let device = Default::default(); + + let tensor = Tensor::<2>::from_floats( + [ + [1.0, 0.5, 0.3, 0.2], + [0.5, 1.0, 0.4, 0.1], + [0.3, 0.4, 1.0, 0.5], + [0.2, 0.1, 0.5, 1.0], + ], + &device, + ); + + // Zero gradient - all zeros + let zero_grad = Tensor::<2>::zeros([4, 4], &device); + + let config = MuonConfig::new(); + let muon = Muon { + momentum: Momentum::new(&config.momentum), + ns_params: NewtonSchulzParams::new(config.ns_coefficients, config.ns_steps), + weight_decay_penalty: None, + epsilon: config.epsilon, + adjust_lr_fn: config.adjust_lr_fn, + }; + + // Should not panic or produce NaN + let (updated_tensor, state) = muon.step(0.01, tensor.clone(), zero_grad, None); + + // Verify state was created + assert!(state.is_some()); + + // With zero gradient and no weight decay, tensor should remain unchanged + let original_data = tensor.into_data(); + let updated_data = updated_tensor.clone().into_data(); + + let original_vals = original_data.as_slice::().unwrap(); + let updated_vals = updated_data.as_slice::().unwrap(); + + for (orig, upd) in original_vals.iter().zip(updated_vals.iter()) { + assert!( + (orig - upd).abs() < 1e-6, + "With zero gradient, tensor should remain unchanged (or very close)" + ); + } + + // Verify no NaN values + for val in updated_vals { + assert!( + !val.is_nan(), + "Result should not contain NaN values with zero gradient" + ); + } + + // Test with weight decay - should still work + let muon_with_decay = Muon { + momentum: Momentum::new(&config.momentum), + ns_params: NewtonSchulzParams::new(config.ns_coefficients, config.ns_steps), + weight_decay_penalty: Some(0.01), + epsilon: config.epsilon, + adjust_lr_fn: config.adjust_lr_fn, + }; + + let tensor2 = Tensor::<2>::from_floats( + [ + [1.0, 0.5, 0.3, 0.2], + [0.5, 1.0, 0.4, 0.1], + [0.3, 0.4, 1.0, 0.5], + [0.2, 0.1, 0.5, 1.0], + ], + &device, + ); + let zero_grad2 = Tensor::<2>::zeros([4, 4], &device); + + let (updated_tensor_decay, _) = + muon_with_decay.step(0.01, tensor2.clone(), zero_grad2, None); + + // With zero gradient but with weight decay, tensor should be slightly reduced + let updated_decay_data = updated_tensor_decay.into_data(); + let updated_decay_vals = updated_decay_data.as_slice::().unwrap(); + + for val in updated_decay_vals { + assert!( + !val.is_nan(), + "Result should not contain NaN with zero gradient and weight decay" + ); + } + + // With weight decay, values should be slightly smaller than original + let original_vals2 = tensor2.into_data().as_slice::().unwrap().to_vec(); + for (orig, upd) in original_vals2.iter().zip(updated_decay_vals.iter()) { + if orig.abs() > 1e-6 { + // Non-zero values should be reduced by weight decay + assert!( + upd.abs() < orig.abs(), + "Weight decay should reduce magnitude: original={}, updated={}", + orig, + upd + ); + } + } + } +} diff --git a/crates/burn-optim/src/optim/rmsprop.rs b/crates/burn-optim/src/optim/rmsprop.rs new file mode 100644 index 0000000..2ace0e4 --- /dev/null +++ b/crates/burn-optim/src/optim/rmsprop.rs @@ -0,0 +1,534 @@ +use burn_core as burn; + +use crate::RecordState; + +use super::{ + Optimizer, + decay::{WeightDecay, WeightDecayConfig}, + module_optimizer::ModuleOptimizer, +}; +use crate::{LearningRate, grad_clipping::GradientClippingConfig}; + +use burn::config::Config; +use burn::tensor::{Device, Tensor}; + +/// Configuration to create the [RmsProp](RmsProp) optimizer. +#[derive(Config, Debug)] +pub struct RmsPropConfig { + /// Smoothing constant. + #[config(default = 0.99)] + alpha: f32, + /// momentum for RmsProp. + #[config(default = 0.9)] + momentum: f32, + /// A value required for numerical stability. + #[config(default = 1e-5)] + epsilon: f32, + /// if True, compute the centered RmsProp, the gradient is normalized by an estimation of its variance + #[config(default = false)] + centered: bool, + /// [Weight decay](WeightDecayConfig) config. + weight_decay: Option, + /// [Gradient Clipping](GradientClippingConfig) config. + grad_clipping: Option, +} + +impl RmsPropConfig { + /// Build a [`RmsProp`] from the config. + pub(crate) fn build(&self) -> RmsProp { + let weight_decay = self.weight_decay.as_ref().map(WeightDecay::new); + RmsProp { + alpha: self.alpha, + centered: self.centered, + weight_decay, + momentum: RmsPropMomentum { + momentum: self.momentum, + epsilon: self.epsilon, + }, + } + } + + /// Initialize RmsProp optimizer. + /// + /// # Returns + /// + /// Returns an optimizer that can be used to optimize a module. + pub fn init(&self) -> ModuleOptimizer { + let mut optim = ModuleOptimizer::from(self.build()); + if let Some(config) = &self.grad_clipping { + optim = optim.with_grad_clipping(config.init()); + } + + optim + } +} + +/// Optimizer that implements stochastic gradient descent with momentum. +/// The optimizer can be configured with [RmsPropConfig](RmsPropConfig). +#[derive(Clone)] +pub struct RmsProp { + alpha: f32, + // epsilon: f32, + centered: bool, + // momentum: Option, + momentum: RmsPropMomentum, + weight_decay: Option, +} + +impl Optimizer for RmsProp { + type State = RmsPropState; + + fn step( + &self, + lr: LearningRate, + tensor: Tensor, + mut grad: Tensor, + state: Option>, + ) -> (Tensor, Option>) { + // fetch state for params + let mut state_square_avg = None; + let mut state_centered = None; + let mut state_momentum = None; + if let Some(state) = state { + state_square_avg = Some(state.square_avg); + state_centered = Some(state.centered); + state_momentum = state.momentum; + } + + // weight_decay transform + if let Some(weight_decay) = &self.weight_decay { + grad = weight_decay.transform(grad, tensor.clone()); + } + + // square_avg transform + let (grad, state_square_avg) = + SquareAvgState::transform(self.alpha, grad, state_square_avg); + + // centered transform + let (grad, state_square_avg, state_centered) = CenteredState::transform( + self.alpha, + self.centered, + grad, + state_square_avg, + state_centered, + ); + + // momentum transform + let (grad, state_centered, state_momentum) = + self.momentum + .transform(grad, state_centered, state_momentum); + + // transition state + let state = RmsPropState::new(state_square_avg, state_centered, state_momentum); + + // tensor param transform + let delta = grad.mul_scalar(lr); + (tensor - delta, Some(state)) + } + + fn to_device(mut state: Self::State, device: &Device) -> Self::State { + state.square_avg = state.square_avg.to_device(device); + state.centered = state.centered.to_device(device); + state.momentum = state.momentum.map(|momentum| momentum.to_device(device)); + state + } +} + +/// State of [RmsProp](RmsProp) +#[derive(RecordState, Clone, new)] +pub struct RmsPropState { + /// Current squared average state. + pub square_avg: SquareAvgState, + /// Current centered state + pub centered: CenteredState, + /// Current gradient momentum, if any. + pub momentum: Option>, +} + +/// [SquareAvgState](SquareAvgState) is to store and pass optimizer step params. +#[derive(RecordState, Clone, new)] +pub struct SquareAvgState { + /// Current squared average. + pub square_avg: Tensor, +} + +impl SquareAvgState { + /// transform [SquareAvgState] to the next step + fn transform(alpha: f32, grad: Tensor, state: Option) -> (Tensor, Self) { + match state { + Some(state) => { + let square_avg = state + .square_avg + .mul_scalar(alpha) + .add(grad.clone().square().mul_scalar(1. - alpha)); + (grad, Self { square_avg }) + } + _ => { + let square_avg = grad.clone().square().mul_scalar(1. - alpha); + (grad, Self { square_avg }) + } + } + } + + /// Moves the state to a device. + /// + /// # Arguments + /// + /// * `device` - Device to move the state to. + /// + /// # Returns + /// + /// * `self` - Moved state. + pub fn to_device(mut self, device: &Device) -> Self { + self.square_avg = self.square_avg.to_device(device); + self + } +} + +/// [CenteredState](CenteredState) is to store and pass optimizer step params. +#[derive(RecordState, Clone, new)] +pub struct CenteredState { + /// The averaged gradient to calculate the centered gradient, if available. + pub grad_avg: Option>, + /// The current average value. + pub avg: Tensor, +} + +impl CenteredState { + /// transform [CenteredState] to the next step + fn transform( + alpha: f32, + centered: bool, + grad: Tensor, + square_avg_state: SquareAvgState, + centered_state: Option, + ) -> (Tensor, SquareAvgState, Self) { + if centered { + let grad_avg_constant = grad.clone().mul_scalar(1. - alpha); + let grad_avg = match centered_state { + Some(state) => state + .grad_avg + .map_or(grad_avg_constant.clone(), move |grad_avg| { + grad_avg.mul_scalar(alpha).add(grad_avg_constant) + }), + _ => grad_avg_constant, + }; + let avg = square_avg_state + .square_avg + .clone() + .sub(grad_avg.clone().square()); + + ( + grad, + square_avg_state, + Self { + grad_avg: Some(grad_avg), + avg, + }, + ) + } else { + ( + grad, + square_avg_state.clone(), + Self { + grad_avg: None, + avg: square_avg_state.square_avg, + }, + ) + } + } + + /// Moves the state to a device. + /// + /// # Arguments + /// + /// * `device` - Device to move the state to. + /// + /// # Returns + /// + /// * `self` - Moved state. + pub fn to_device(mut self, device: &Device) -> Self { + self.grad_avg = self.grad_avg.map(|grad_avg| grad_avg.to_device(device)); + self.avg = self.avg.to_device(device); + self + } +} + +/// [RmsPropMomentum](RmsPropMomentum) is to store config status for optimizer. +/// (, which is stored in [optimizer](RmsProp) itself and not passed in during `step()` calculation) +#[derive(Clone)] +pub struct RmsPropMomentum { + momentum: f32, + epsilon: f32, +} + +impl RmsPropMomentum { + /// transform [grad](Tensor) and [RmsPropMomentumState] to the next step + fn transform( + &self, + grad: Tensor, + centered_state: CenteredState, + momentum_state: Option>, + ) -> (Tensor, CenteredState, Option>) { + let grad = grad.div(centered_state.avg.clone().sqrt().add_scalar(self.epsilon)); + + if self.momentum > 0. { + let buf = match momentum_state { + Some(state) => state.buf.mul_scalar(self.momentum).add(grad), + _ => grad, + }; + ( + buf.clone(), + centered_state, + Some(RmsPropMomentumState { buf }), + ) + } else { + (grad, centered_state, None) + } + } +} + +/// [RmsPropMomentumState](RmsPropMomentumState) is to store and pass optimizer step params. +#[derive(RecordState, Clone, new)] +pub struct RmsPropMomentumState { + buf: Tensor, +} + +impl RmsPropMomentumState { + /// Moves the state to a device. + /// + /// # Arguments + /// + /// * `device` - Device to move the state to. + /// + /// # Returns + /// + /// * `self` - Moved state. + pub fn to_device(mut self, device: &Device) -> Self { + self.buf = self.buf.to_device(device); + self + } +} + +#[cfg(test)] +mod tests { + use burn::tensor::Tolerance; + + use super::*; + use crate::optim::GradientsParams; + use burn::module::Param; + use burn::tensor::{Distribution, Tensor, TensorData}; + use burn_nn::{Linear, LinearConfig}; + + type FT = f32; + + const LEARNING_RATE: LearningRate = 0.01; + + #[test] + fn test_rmsprop_optimizer_save_load_state() { + let device = Device::default().autodiff(); + let linear = LinearConfig::new(6, 6).init(&device); + let x = Tensor::<2>::random([2, 6], Distribution::Default, &device); + let mut optimizer = create_rmsprop(); + let grads = linear.forward(x).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let _linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let bytes = optimizer.into_bytes().unwrap(); + assert!(!bytes.is_empty()); + + #[cfg(feature = "std")] + optimizer + .save(std::env::temp_dir().as_path().join("test_optim_rmsprop")) + .unwrap(); + + let state_optim_before = optimizer.to_record(); + let optimizer = create_rmsprop().from_bytes(bytes).unwrap(); + let state_optim_after = optimizer.to_record(); + + assert_eq!(state_optim_before.len(), state_optim_after.len()); + } + + /// used for test differences and debug + #[test] + fn test_rmsprop_optimizer_with_numbers_basic() { + let device = Device::default().autodiff(); + let linear = given_linear_layer( + TensorData::from([ + [1., 1., 1., 1., 1., 1.], + [1., 1., 1., 1., 1., 1.], + [1., 1., 1., 1., 1., 1.], + [1., 1., 1., 1., 1., 1.], + [1., 1., 1., 1., 1., 1.], + [1., 1., 1., 1., 1., 1.], + ]), + TensorData::from([0.5, 0.5, 0.5, 0.5, 0.5, 0.5]), + &device, + ); + let x_1 = Tensor::<2>::from_floats( + [ + [0.6294, 0.0940, 0.8176, 0.8824, 0.5228, 0.4310], + [0.7152, 0.9559, 0.7893, 0.5684, 0.5939, 0.8883], + ], + &device, + ) + .require_grad(); + let x_2 = Tensor::<2>::from_floats( + [ + [0.8491, 0.2108, 0.8939, 0.4433, 0.5527, 0.2528], + [0.3270, 0.0412, 0.5538, 0.9605, 0.3195, 0.9085], + ], + &device, + ) + .require_grad(); + + let mut optimizer = RmsPropConfig::new() + .with_alpha(0.99) + .with_epsilon(1e-8) + .with_weight_decay(WeightDecayConfig::new(0.05).into()) + .with_momentum(0.9) + .with_centered(false) + .init(); + + // println!("linear is {:?}", linear); + let grads = linear.forward(x_1).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + // println!("linear is {:?}", linear); + let grads = linear.forward(x_2).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + // println!("linear is {:?}", linear); + let state_updated = linear; + + let (weight_updated, bias_updated) = ( + state_updated.weight.to_data(), + state_updated.bias.unwrap().to_data(), + ); + + // println!("\nweight_updated\n{:?}", weight_updated); + // println!("\nbias_updated\n{:?}", bias_updated); + + let weights_expected = TensorData::from([ + [0.743937, 0.743937, 0.743937, 0.743937, 0.743937, 0.743937], + [0.783809, 0.783809, 0.783809, 0.783809, 0.783809, 0.783809], + [0.742881, 0.742881, 0.742881, 0.742881, 0.742881, 0.742881], + [0.740366, 0.740366, 0.740366, 0.740366, 0.740366, 0.740366], + [0.748005, 0.748005, 0.748005, 0.748005, 0.748005, 0.748005], + [0.743710, 0.743710, 0.743710, 0.743710, 0.743710, 0.743710], + ]); + let bias_expected = + TensorData::from([0.239199, 0.239199, 0.239199, 0.239199, 0.239199, 0.239199]); + + let tolerance = Tolerance::absolute(1e-6); + bias_updated.assert_approx_eq::(&bias_expected, tolerance); + weight_updated.assert_approx_eq::(&weights_expected, tolerance); + } + + #[test] + fn test_rmsprop_optimizer_with_numbers() { + let device = Device::default().autodiff(); + let linear = given_linear_layer( + TensorData::from([ + [-0.3206, 0.1374, 0.4043, 0.3200, 0.0859, 0.0671], + [0.0777, -0.0185, -0.3667, 0.2550, 0.1955, -0.2922], + [-0.0190, 0.0346, -0.2962, 0.2484, -0.2780, 0.3130], + [-0.2980, -0.2214, -0.3715, -0.2981, -0.0761, 0.1626], + [0.3300, -0.2182, 0.3717, -0.1729, 0.3796, -0.0304], + [-0.0159, -0.0120, 0.1258, 0.1921, 0.0293, 0.3833], + ]), + TensorData::from([-0.3905, 0.0884, -0.0970, 0.1176, 0.1366, 0.0130]), + &device, + ); + let x_1 = Tensor::<2>::from_floats( + [ + [0.6294, 0.0940, 0.8176, 0.8824, 0.5228, 0.4310], + [0.7152, 0.9559, 0.7893, 0.5684, 0.5939, 0.8883], + ], + &device, + ) + .require_grad(); + let x_2 = Tensor::<2>::from_floats( + [ + [0.8491, 0.2108, 0.8939, 0.4433, 0.5527, 0.2528], + [0.3270, 0.0412, 0.5538, 0.9605, 0.3195, 0.9085], + ], + &device, + ) + .require_grad(); + + let mut optimizer = RmsPropConfig::new() + .with_alpha(0.99) + .with_epsilon(1e-8) + .with_weight_decay(WeightDecayConfig::new(0.05).into()) + .with_momentum(0.9) + .with_centered(false) + .init(); + + let grads = linear.forward(x_1).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let grads = linear.forward(x_2).backward(); + let grads = GradientsParams::from_grads(grads, &linear); + let linear = optimizer.step(LEARNING_RATE.into(), linear, grads); + + let state_updated = linear; + let weights_expected = TensorData::from([ + [ + -0.576399, -0.118494, 0.148353, 0.064070, -0.169983, -0.188779, + ], + [ + -0.135571, -0.231448, -0.578445, 0.041143, -0.018162, -0.504207, + ], + [ + -0.275990, -0.222397, -0.553153, -0.008625, -0.534956, 0.055967, + ], + [ + -0.557575, -0.480979, -0.631072, -0.557675, -0.335686, -0.096997, + ], + [ + 0.078313, -0.469618, 0.119993, -0.424341, 0.127890, -0.281912, + ], + [ + -0.271996, -0.268097, -0.130324, -0.064037, -0.226805, 0.127126, + ], + ]); + let bias_expected = TensorData::from([ + -0.651299, -0.172400, -0.357800, -0.143200, -0.124200, -0.247800, + ]); + + let (weight_updated, bias_updated) = ( + state_updated.weight.to_data(), + state_updated.bias.unwrap().to_data(), + ); + + // println!("\nweight_updated\n{:?}", weight_updated); + // println!("\nbias_updated\n{:?}", bias_updated); + + let tolerance = Tolerance::absolute(1e-6); + bias_updated.assert_approx_eq::(&bias_expected, tolerance); + weight_updated.assert_approx_eq::(&weights_expected, tolerance); + } + + fn given_linear_layer(weight: TensorData, bias: TensorData, device: &Device) -> Linear { + Linear { + weight: Param::from_data(weight, device), + bias: Some(Param::from_data(bias, device)), + } + } + + fn create_rmsprop() -> ModuleOptimizer { + RmsPropConfig { + alpha: 0.99, + epsilon: 1e-9, + centered: false, + weight_decay: Some(WeightDecayConfig { penalty: 0.05 }), + momentum: 0.9, + grad_clipping: None, + } + .init() + } +} diff --git a/crates/burn-optim/src/optim/sgd.rs b/crates/burn-optim/src/optim/sgd.rs new file mode 100644 index 0000000..82f63b9 --- /dev/null +++ b/crates/burn-optim/src/optim/sgd.rs @@ -0,0 +1,176 @@ +use burn_core as burn; + +use super::Optimizer; +use super::decay::{WeightDecay, WeightDecayConfig}; +use super::module_optimizer::ModuleOptimizer; +use super::momentum::{Momentum, MomentumConfig, MomentumState}; +use crate::LearningRate; +use crate::RecordState; +use crate::grad_clipping::GradientClippingConfig; +use burn::config::Config; +use burn::tensor::Device; +use burn::tensor::Tensor; + +/// Configuration to create the [Sgd](Sgd) optimizer. +#[derive(Config, Debug)] +pub struct SgdConfig { + /// [Weight decay](WeightDecayConfig) config. + weight_decay: Option, + /// [Momentum](MomentumConfig) config. + momentum: Option, + /// [Gradient Clipping](GradientClippingConfig) config. + gradient_clipping: Option, +} + +/// Optimizer that implements stochastic gradient descent with momentum. +/// +/// The optimizer can be configured with [SgdConfig](SgdConfig). +#[derive(Clone)] +pub struct Sgd { + momentum: Option, + weight_decay: Option, +} + +/// State of [Sgd](Sgd). +#[derive(RecordState, Clone, new)] +pub struct SgdState { + /// The current state of the momentum (if any). + pub momentum: Option>, +} + +impl SgdConfig { + /// Build a [`Sgd`] from the config. + pub(crate) fn build(&self) -> Sgd { + Sgd { + momentum: self.momentum.as_ref().map(Momentum::new), + weight_decay: self.weight_decay.as_ref().map(WeightDecay::new), + } + } + + /// Initializes the SGD optimizer from the configuration. + pub fn init(&self) -> ModuleOptimizer { + let mut optim = ModuleOptimizer::from(self.build()); + if let Some(config) = &self.gradient_clipping { + optim = optim.with_grad_clipping(config.init()); + } + optim + } +} + +impl Optimizer for Sgd { + type State = SgdState; + + fn step( + &self, + lr: LearningRate, + tensor: Tensor, + mut grad: Tensor, + state: Option>, + ) -> (Tensor, Option>) { + let mut state_momentum = None; + + if let Some(state) = state { + state_momentum = state.momentum; + } + + if let Some(weight_decay) = &self.weight_decay { + grad = weight_decay.transform(grad, tensor.clone()); + } + + if let Some(momentum) = &self.momentum { + let (grad_out, state) = momentum.transform(grad, state_momentum); + state_momentum = Some(state); + grad = grad_out; + } + + let state = SgdState::new(state_momentum); + let delta = grad.mul_scalar(lr); + + (tensor - delta, Some(state)) + } + + fn to_device(mut state: Self::State, device: &Device) -> Self::State { + state.momentum = state.momentum.map(|state| state.to_device(device)); + state + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{grad_clipping::GradientClipping, optim::GradientsParams}; + use burn::tensor::{Distribution, Shape}; + use burn_nn::{Linear, LinearConfig}; + + const LEARNING_RATE: LearningRate = 0.02; + + #[test] + fn with_updated_params_should_have_state() { + let device = Device::default().autodiff(); + let layer = layer(&device); + let mut optim = sgd_with_all(); + let loss = layer.forward(random_tensor(&device)); + let grads = loss.backward(); + let grads = GradientsParams::from_grads(grads, &layer); + let _layer = optim.step(LEARNING_RATE.into(), layer, grads); + + let record = optim.to_record(); + + assert!(!record.is_empty()); + } + + #[test] + fn without_updated_params_should_not_have_state() { + let optim = sgd_with_all(); + let record = optim.to_record(); + assert!(record.is_empty()); + } + + #[test] + fn can_attach_gradient_clipping() { + let optim = sgd_with_all().with_grad_clipping(GradientClipping::Value(0.5)); + assert!(optim.has_gradient_clipping()); + } + + #[test] + fn should_load_state() { + let device = Device::default().autodiff(); + let layer = layer(&device); + let mut optim = sgd_with_all(); + let loss = layer.forward(random_tensor(&device)); + let grads = loss.backward(); + let grads = GradientsParams::from_grads(grads, &layer); + let _layer = optim.step(LEARNING_RATE.into(), layer, grads); + + let record = optim.to_record(); + let bytes = optim.into_bytes().unwrap(); + let optim_new = sgd_with_all(); + let record_new = optim_new.to_record(); + let optim_new = optim_new.from_bytes(bytes).unwrap(); + let state_restored = optim_new.to_record(); + + assert_ne!(record.len(), record_new.len()); + assert_eq!(record.len(), state_restored.len()); + } + + fn random_tensor(device: &Device) -> Tensor<2> { + Tensor::<2>::random(Shape::new([2, 20]), Distribution::Default, device) + } + + fn layer(device: &Device) -> Linear { + LinearConfig::new(20, 20).init(device) + } + + fn sgd_with_all() -> ModuleOptimizer { + SgdConfig { + weight_decay: Some(WeightDecayConfig { penalty: 0.05 }), + momentum: Some(MomentumConfig { + momentum: 0.9, + dampening: 0.1, + nesterov: true, + }), + gradient_clipping: None, + } + .init() + } +} diff --git a/crates/burn-optim/src/optim/state.rs b/crates/burn-optim/src/optim/state.rs new file mode 100644 index 0000000..788b898 --- /dev/null +++ b/crates/burn-optim/src/optim/state.rs @@ -0,0 +1,282 @@ +//! Decomposition of a state struct into named tensors and scalars for the burnpack format. +//! +//! Used for both optimizer state — keyed per-[`ParamId`](burn_core::module::ParamId) rather than +//! per module path — and learning-rate scheduler state. Each state is usually a small struct +//! holding a few tensors plus scalar bookkeeping (e.g. a step counter). [`RecordState`] flattens +//! such a struct into a flat list of named tensors plus a few typed [scalars](burn_pack::Scalar), +//! and reconstructs it on load. +//! +//! Implementations are generated with `#[derive(RecordState)]`; the derive supports the field +//! shapes `Tensor`, `Option>`, `Vec>`, scalars, optional scalars, nested +//! states and optional nested states. Scalar fields rely on `burn_pack`'s [`From`]/[`TryFrom`] +//! conversions to [`Scalar`](burn_pack::Scalar). + +use alloc::collections::BTreeMap; +use alloc::format; +use alloc::string::String; +use alloc::vec::Vec; + +use burn_core::tensor::{Device, TensorData}; +use burn_pack::Scalar; + +pub use burn_derive::RecordState; + +/// Join a `prefix` and a `leaf` into a dot-separated path (`"prefix.leaf"`). +/// +/// An empty prefix yields the leaf unchanged, so a top-level call can pass `""`. +pub fn join_path(prefix: &str, leaf: &str) -> String { + if prefix.is_empty() { + return String::from(leaf); + } + let mut path = String::with_capacity(prefix.len() + 1 + leaf.len()); + path.push_str(prefix); + path.push('.'); + path.push_str(leaf); + path +} + +/// Join a `prefix` and a numeric `index` into a dot-separated path (`"prefix.3"`). +/// +/// Used by the derive to name the elements of a `Vec` field. +pub fn join_index(prefix: &str, index: usize) -> String { + format!("{prefix}.{index}") +} + +/// Accumulates the named tensors and scalars produced while flattening a [`RecordState`]. +#[derive(Default, Debug)] +pub struct StateSink { + /// The collected `(name, data)` tensor leaves. + pub tensors: Vec<(String, TensorData)>, + /// The collected `(name, value)` scalar leaves. + pub scalars: Vec<(String, Scalar)>, +} + +impl StateSink { + /// Record a tensor leaf named `{prefix}.{leaf}`. + pub fn push_tensor(&mut self, prefix: &str, leaf: &str, data: TensorData) { + self.tensors.push((join_path(prefix, leaf), data)); + } + + /// Record a scalar leaf named `{prefix}.{leaf}`. + pub fn push_scalar(&mut self, prefix: &str, leaf: &str, value: Scalar) { + self.scalars.push((join_path(prefix, leaf), value)); + } +} + +/// Provides the named tensors and scalars consumed while reconstructing an [`RecordState`]. +/// +/// Tensors are taken (removed) by name so the same source can feed several parameters in turn; +/// scalars are looked up by name and left in place. +#[derive(Default, Debug)] +pub struct StateSource { + tensors: BTreeMap, + scalars: BTreeMap, +} + +impl StateSource { + /// Create a source from an existing scalar map (e.g. the burnpack scalars). + pub fn new(scalars: BTreeMap) -> Self { + Self { + tensors: BTreeMap::new(), + scalars, + } + } + + /// Register a tensor under its full `name`. + pub fn insert_tensor(&mut self, name: String, data: TensorData) { + self.tensors.insert(name, data); + } + + /// Take the tensor named `{prefix}.{leaf}`, if present. + pub fn take_tensor(&mut self, prefix: &str, leaf: &str) -> Option { + self.tensors.remove(&join_path(prefix, leaf)) + } + + /// Read the scalar named `{prefix}.{leaf}`, if present. + pub fn take_scalar(&mut self, prefix: &str, leaf: &str) -> Option { + self.scalars.get(&join_path(prefix, leaf)).copied() + } + + /// Whether any tensor or scalar leaf was recorded under `prefix` (a key beginning with + /// `"{prefix}."`). + /// + /// Used by the derive to tell an absent optional nested state (nothing recorded) apart from a + /// present one whose leaves all happen to be optional. + pub fn has_under(&self, prefix: &str) -> bool { + let pat = join_path(prefix, ""); + self.tensors.keys().any(|k| k.starts_with(&pat)) + || self.scalars.keys().any(|k| k.starts_with(&pat)) + } +} + +/// A type that can be flattened into named tensors and scalars and rebuilt from them. +/// +/// Generated with `#[derive(RecordState)]`. The `prefix` threads the parameter identity (and any +/// nested field path) through the recursion; leaves are named `{prefix}.{field}`. +pub trait RecordState: Sized + Send + Sync + 'static { + /// Flatten `self` into `out`, naming every leaf under `prefix`. + fn state_flatten(&self, prefix: &str, out: &mut StateSink); + + /// Rebuild a value from `src`, reading leaves named under `prefix`. + /// + /// Returns `None` when a required leaf is absent. The derive uses [`StateSource::has_under`] to + /// presence-test optional nested states, so an `Option` field is `None` exactly when + /// nothing was recorded under its path. + /// + /// A scalar leaf that is present but holds an incompatible [`Scalar`] variant (e.g. a + /// hand-edited or forward-version file) is treated the same as absent. + fn state_unflatten(prefix: &str, src: &mut StateSource, device: &Device) -> Option; +} + +/// The empty state — for stateless values (e.g. a constant learning rate scheduler). +impl RecordState for () { + fn state_flatten(&self, _prefix: &str, _out: &mut StateSink) {} + + fn state_unflatten(_prefix: &str, _src: &mut StateSource, _device: &Device) -> Option { + Some(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Tensor; + use burn_core as burn; + + /// Flatten `state` then rebuild it from the produced tensors and scalars. + fn round_trip(state: &T) -> Option { + let mut sink = StateSink::default(); + state.state_flatten("p", &mut sink); + + let scalars: BTreeMap = sink.scalars.into_iter().collect(); + let mut source = StateSource::new(scalars); + for (name, data) in sink.tensors { + source.insert_tensor(name, data); + } + + T::state_unflatten("p", &mut source, &Device::default()) + } + + fn tensor(values: &[f32]) -> Tensor<1> { + Tensor::from_data(TensorData::from(values), &Device::default()) + } + + fn data(t: &Tensor<1>) -> Vec { + t.clone().into_data().to_vec().unwrap() + } + + #[derive(RecordState, Clone, Debug)] + struct Inner { + weight: Tensor, + step: i64, + } + + #[derive(RecordState, Clone, Debug)] + struct Full { + t: Tensor, + opt_tensor: Option>, + history: Vec>, + count: usize, + opt_scalar: Option, + nested: Inner, + opt_nested: Option>, + } + + #[test] + fn all_field_kinds_round_trip() { + let state = Full::<1> { + t: tensor(&[1.0, 2.0]), + opt_tensor: Some(tensor(&[3.0])), + history: vec![tensor(&[4.0]), tensor(&[5.0, 6.0])], + count: 7, + opt_scalar: Some(8.5), + nested: Inner { + weight: tensor(&[9.0]), + step: 10, + }, + opt_nested: Some(Inner { + weight: tensor(&[11.0]), + step: 12, + }), + }; + + let out = round_trip(&state).unwrap(); + + assert_eq!(data(&out.t), vec![1.0, 2.0]); + assert_eq!(data(&out.opt_tensor.unwrap()), vec![3.0]); + assert_eq!(out.history.len(), 2); + assert_eq!(data(&out.history[0]), vec![4.0]); + assert_eq!(data(&out.history[1]), vec![5.0, 6.0]); + assert_eq!(out.count, 7); + assert_eq!(out.opt_scalar, Some(8.5)); + assert_eq!(data(&out.nested.weight), vec![9.0]); + assert_eq!(out.nested.step, 10); + let opt_nested = out.opt_nested.unwrap(); + assert_eq!(data(&opt_nested.weight), vec![11.0]); + assert_eq!(opt_nested.step, 12); + } + + #[test] + fn absent_optionals_round_trip_to_none() { + let state = Full::<1> { + t: tensor(&[1.0]), + opt_tensor: None, + history: vec![], + count: 0, + opt_scalar: None, + nested: Inner { + weight: tensor(&[2.0]), + step: 0, + }, + opt_nested: None, + }; + + let out = round_trip(&state).unwrap(); + + assert!(out.opt_tensor.is_none()); + assert!(out.history.is_empty()); + assert!(out.opt_scalar.is_none()); + assert!(out.opt_nested.is_none()); + } + + /// Regression: an optional nested state whose fields are all optional must come back `None` + /// when nothing was recorded under its path — it must not be resurrected as `Some`. + #[derive(RecordState, Clone, Debug)] + struct AllOptional { + x: Option>, + y: Option, + } + + #[derive(RecordState, Clone, Debug)] + struct OuterOpt { + inner: Option>, + } + + #[test] + fn optional_all_optional_nested_stays_none() { + let state = OuterOpt::<1> { inner: None }; + let out = round_trip(&state).unwrap(); + assert!(out.inner.is_none()); + } + + #[test] + fn optional_all_optional_nested_present_with_content() { + let state = OuterOpt::<1> { + inner: Some(AllOptional { + x: Some(tensor(&[1.0, 2.0])), + y: None, + }), + }; + let out = round_trip(&state).unwrap(); + let inner = out.inner.expect("present because a leaf was recorded"); + assert_eq!(data(&inner.x.unwrap()), vec![1.0, 2.0]); + assert!(inner.y.is_none()); + } + + #[test] + fn missing_required_tensor_yields_none() { + // A source with nothing recorded cannot rebuild a struct that has a required tensor leaf. + let mut source = StateSource::new(BTreeMap::new()); + assert!(Inner::<1>::state_unflatten("p", &mut source, &Device::default()).is_none()); + } +} diff --git a/crates/burn-optim/src/optim/visitor.rs b/crates/burn-optim/src/optim/visitor.rs new file mode 100644 index 0000000..739ac41 --- /dev/null +++ b/crates/burn-optim/src/optim/visitor.rs @@ -0,0 +1,57 @@ +use burn_core as burn; + +use super::GradientsParams; +use burn::module::{AutodiffModule, ModuleVisitor, Param, ParamId}; +use burn::tensor::{Device, Gradients, Tensor}; +use core::marker::PhantomData; + +#[cfg(not(feature = "std"))] +use alloc::vec::Vec; + +#[derive(new)] +pub struct GradientsParamsConverter<'a, M: AutodiffModule> { + grads: &'a mut Gradients, + grads_params: &'a mut GradientsParams, + phatom: PhantomData, + filter: Option>, +} + +#[derive(new)] +pub struct GradientsParamsChangeDevice<'a, M: AutodiffModule> { + device: &'a Device, + grads: &'a mut GradientsParams, + phatom: PhantomData, +} + +impl ModuleVisitor for GradientsParamsConverter<'_, M> +where + M: AutodiffModule, +{ + fn visit_float(&mut self, param: &Param>) { + if let Some(filter) = self.filter.as_ref() + && !filter.contains(¶m.id) + { + return; + } + + let Some(grad) = param.val().grad_remove(self.grads) else { + return; + }; + + self.grads_params.register(param.id, grad); + } +} + +impl ModuleVisitor for GradientsParamsChangeDevice<'_, M> +where + M: AutodiffModule, +{ + fn visit_float(&mut self, param: &Param>) { + let Some(grad) = self.grads.remove::(param.id) else { + return; + }; + + self.grads + .register::(param.id, grad.to_device(self.device)); + } +} diff --git a/crates/burn-pack/Cargo.toml b/crates/burn-pack/Cargo.toml new file mode 100644 index 0000000..54eb2b8 --- /dev/null +++ b/crates/burn-pack/Cargo.toml @@ -0,0 +1,36 @@ +[package] +authors = ["Dilshod Tadjibaev (@antimora)"] +categories = ["science", "no-std", "embedded", "wasm"] +description = "The burnpack binary serialization format for Burn" +documentation = "https://docs.rs/burn-pack" +edition.workspace = true +keywords = [ + "deep-learning", + "machine-learning", + "tensor", + "storage", + "serialization", +] +license.workspace = true +name = "burn-pack" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-pack" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["std"] +std = ["burn-std/std", "byteorder/std", "ciborium/std"] + +[dependencies] +burn-std = { workspace = true } + +# External dependencies +byteorder = { workspace = true, default-features = false } +ciborium = { workspace = true } +serde = { workspace = true } + +[dev-dependencies] +tempfile = { workspace = true } diff --git a/crates/burn-pack/README.md b/crates/burn-pack/README.md new file mode 100644 index 0000000..de9acf8 --- /dev/null +++ b/crates/burn-pack/README.md @@ -0,0 +1,38 @@ +# Burn Pack + +> The **burnpack** binary serialization format for the Burn deep learning framework + +[![Current Crates.io Version](https://img.shields.io/crates/v/burn-pack.svg)](https://crates.io/crates/burn-pack) +[![Documentation](https://docs.rs/burn-pack/badge.svg)](https://docs.rs/burn-pack) + +`burn-pack` reads and writes the burnpack container format. It is tensor-library-agnostic and +dependency-light: it depends only on [`burn-std`](https://crates.io/crates/burn-std) (for `DType` +/ `Bytes`), `serde`, and a CBOR codec — it knows the on-disk format but has no notion of Burn +modules. Tensor data is `Bytes`-native and read lazily from files (256-byte aligned for +zero-copy mmap and efficient GPU transfers), and the reader is hardened against malformed input. + +If you just want to save and load Burn models, use the higher-level +[`burn-core`](https://crates.io/crates/burn-core) record API or +[`burn-store`](https://crates.io/crates/burn-store) (which adds PyTorch/SafeTensors interop). +Reach for `burn-pack` directly only to produce or consume the raw format. + +## Usage + +```rust +use burn_pack::{Bytes, DType, Reader, Tensor, Writer}; + +let raw: Vec = [1.0f32, 2.0, 3.0, 4.0].iter().flat_map(|v| v.to_le_bytes()).collect(); +let tensor = Tensor::new("weight".to_string(), DType::F32, vec![2, 2], None, Bytes::from_bytes_vec(raw)); + +let packed = Writer::new(vec![tensor]).into_bytes().unwrap(); +let reader = Reader::from_bytes(packed).unwrap(); +assert_eq!(reader.into_tensors().unwrap()[0].shape.to_vec(), vec![2, 2]); +``` + +Use `Writer::write_to_file` / `Reader::from_file` for disk I/O (the default `std` feature; disable +it for no-std targets). See the [docs](https://docs.rs/burn-pack) for the format layout and the +full API. + +## License + +This project is dual-licensed under MIT and Apache-2.0. diff --git a/crates/burn-pack/src/base.rs b/crates/burn-pack/src/base.rs new file mode 100644 index 0000000..6b525d0 --- /dev/null +++ b/crates/burn-pack/src/base.rs @@ -0,0 +1,390 @@ +//! Core types and constants for the Burnpack file format. +//! +//! See the [parent module](crate::burnpack) for the complete file format specification. + +use alloc::collections::BTreeMap; +use alloc::string::String; +use alloc::vec::Vec; +use burn_std::DType; +use byteorder::{ByteOrder, LittleEndian}; +use serde::{Deserialize, Serialize}; + +/// Magic number identifying a Burnpack file: "BURN" in ASCII (0x4255524E) +/// When written to file in little-endian format, appears as "NRUB" bytes +pub const MAGIC_NUMBER: u32 = 0x4255524E; + +/// Current format version +pub const FORMAT_VERSION: u16 = 0x0001; + +/// Size of the magic number in bytes +pub const MAGIC_SIZE: usize = 4; + +/// Size of the format version in bytes +pub const VERSION_SIZE: usize = 2; + +/// Size of the metadata size field in bytes +pub const METADATA_SIZE_FIELD_SIZE: usize = 4; + +/// Total header size (computed from components) +pub const HEADER_SIZE: usize = MAGIC_SIZE + VERSION_SIZE + METADATA_SIZE_FIELD_SIZE; + +/// Alignment for tensor data in bytes. +/// +/// All tensor data is aligned to 256-byte boundaries to enable efficient +/// memory-mapped (mmap) zero-copy loading. This alignment ensures: +/// - Proper pointer alignment for all tensor element types (f64 requires 8-byte alignment) +/// - Cache-line friendly access (most CPUs use 64-byte cache lines) +/// - GPU memory alignment (CUDA prefers 256-byte for coalesced access) +/// - Future-proofing for wider SIMD (AVX-512 = 64 bytes, future AVX-1024 = 128 bytes) +/// +/// Industry alignment choices: +/// - 256-byte: GGUF, MLX, ncnn, MNN, TNN, vLLM-AWQ, Marlin (15+ formats) +/// - 64-byte: SafeTensors (minimum for AVX-512) +/// - 4096-byte: Core ML +/// +/// 256-byte alignment has negligible overhead for typical tensor sizes while +/// providing maximum compatibility with current and future hardware. +pub const TENSOR_ALIGNMENT: u64 = 256; + +/// Calculate the byte offset where the tensor data section starts. +/// +/// The data section is padded to start at a 256-byte aligned position +/// so that all tensor offsets (which are relative to data section) result +/// in properly aligned absolute file positions for mmap zero-copy access. +/// +/// This function must be used consistently by both writer and reader. +#[inline] +pub fn aligned_data_section_start(metadata_size: usize) -> usize { + let unaligned_start = (HEADER_SIZE + metadata_size) as u64; + // Keep multiplication in u64 space to avoid overflow on 32-bit systems + (unaligned_start.div_ceil(TENSOR_ALIGNMENT) * TENSOR_ALIGNMENT) as usize +} + +// Security limits to prevent DoS attacks via resource exhaustion +// These can be adjusted based on your use case + +/// Maximum allowed metadata size (100 MB) +/// Prevents memory exhaustion attacks via oversized metadata claims +pub const MAX_METADATA_SIZE: u32 = 100 * 1024 * 1024; + +/// Maximum allowed tensor size per tensor +/// Prevents memory exhaustion attacks via oversized tensor claims +/// 32-bit platforms: 2 GB limit (to fit within usize range) +/// 64-bit platforms: 10 GB limit +#[cfg(target_pointer_width = "32")] +pub const MAX_TENSOR_SIZE: usize = 2 * 1024 * 1024 * 1024; +#[cfg(not(target_pointer_width = "32"))] +pub const MAX_TENSOR_SIZE: usize = 10 * 1024 * 1024 * 1024; + +/// Maximum allowed number of tensors (100,000) +/// Prevents resource exhaustion via excessive tensor counts +pub const MAX_TENSOR_COUNT: usize = 100_000; + +/// Maximum CBOR deserialization recursion depth (128 levels) +/// Prevents stack overflow attacks via deeply nested CBOR structures +pub const MAX_CBOR_RECURSION_DEPTH: usize = 128; + +/// Maximum allowed file size (100 GB) +/// Prevents resource exhaustion from extremely large files +/// This limit applies to file-based loading (mmap and buffered) +#[cfg(feature = "std")] +pub const MAX_FILE_SIZE: u64 = 100 * 1024 * 1024 * 1024; + +/// Byte range for magic number in header +pub const fn magic_range() -> core::ops::Range { + let start = 0; + let end = start + MAGIC_SIZE; + start..end +} + +/// Byte range for format version in header +pub const fn version_range() -> core::ops::Range { + let start = MAGIC_SIZE; + let end = start + VERSION_SIZE; + start..end +} + +/// Byte range for metadata size field in header +pub const fn metadata_size_range() -> core::ops::Range { + let start = MAGIC_SIZE + VERSION_SIZE; + let end = start + METADATA_SIZE_FIELD_SIZE; + start..end +} + +// Compile-time validation that ranges are correct +const _: () = assert!(MAGIC_SIZE + VERSION_SIZE + METADATA_SIZE_FIELD_SIZE == HEADER_SIZE); + +/// Header structure for Burnpack files +#[derive(Debug, Clone, Copy)] +pub struct Header { + /// Magic number (4 bytes): 0x4255524E ("BURN") + pub magic: u32, + /// Format version (2 bytes) + pub version: u16, + /// Size of CBOR metadata in bytes (4 bytes) + pub metadata_size: u32, +} + +impl Header { + /// Create a new header with the given metadata size + pub fn new(metadata_size: u32) -> Self { + Self { + magic: MAGIC_NUMBER, + version: FORMAT_VERSION, + metadata_size, + } + } + + /// Serialize header into bytes + pub fn into_bytes(self) -> [u8; HEADER_SIZE] { + let mut bytes = [0u8; HEADER_SIZE]; + LittleEndian::write_u32(&mut bytes[magic_range()], self.magic); + LittleEndian::write_u16(&mut bytes[version_range()], self.version); + LittleEndian::write_u32(&mut bytes[metadata_size_range()], self.metadata_size); + bytes + } + + /// Deserialize header from bytes + pub fn from_bytes(bytes: &[u8]) -> Result { + if bytes.len() < HEADER_SIZE { + return Err(Error::InvalidHeader); + } + + let magic = LittleEndian::read_u32(&bytes[magic_range()]); + if magic != MAGIC_NUMBER { + return Err(Error::InvalidMagicNumber); + } + + let version = LittleEndian::read_u16(&bytes[version_range()]); + let metadata_size = LittleEndian::read_u32(&bytes[metadata_size_range()]); + + Ok(Self { + magic, + version, + metadata_size, + }) + } +} + +/// A typed scalar value stored alongside tensors in a burnpack container. +/// +/// Scalars are kept in the CBOR metadata section (not the tensor data section), so they carry +/// no alignment cost. The field is optional in the format: files written before scalar support +/// simply omit it, and readers default it to empty. +/// +/// Convert to/from the primitive numeric and boolean types with [`From`] / [`TryFrom`] +/// (e.g. `Scalar::from(3usize)`, `u32::try_from(scalar)`). +#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)] +pub enum Scalar { + /// A signed integer. + Int(i64), + /// An unsigned integer. + UInt(u64), + /// A floating-point number. + Float(f64), + /// A boolean. + Bool(bool), +} + +/// Error returned when a [`Scalar`] cannot be converted to a requested primitive type +/// (wrong variant or out of range). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ScalarConversionError; + +impl core::fmt::Display for ScalarConversionError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + write!(f, "scalar value does not fit the requested type") + } +} + +impl core::error::Error for ScalarConversionError {} + +macro_rules! impl_scalar_int { + ($($t:ty => $variant:ident),* $(,)?) => { + $( + impl From<$t> for Scalar { + fn from(value: $t) -> Self { + Scalar::$variant(value as _) + } + } + + impl TryFrom for $t { + type Error = ScalarConversionError; + fn try_from(scalar: Scalar) -> Result { + match scalar { + Scalar::Int(v) => v.try_into().map_err(|_| ScalarConversionError), + Scalar::UInt(v) => v.try_into().map_err(|_| ScalarConversionError), + _ => Err(ScalarConversionError), + } + } + } + )* + }; +} + +impl_scalar_int!( + i8 => Int, i16 => Int, i32 => Int, i64 => Int, isize => Int, + u8 => UInt, u16 => UInt, u32 => UInt, u64 => UInt, usize => UInt, +); + +impl From for Scalar { + fn from(value: f64) -> Self { + Scalar::Float(value) + } +} + +impl From for Scalar { + fn from(value: f32) -> Self { + Scalar::Float(value as f64) + } +} + +impl From for Scalar { + fn from(value: bool) -> Self { + Scalar::Bool(value) + } +} + +impl TryFrom for f64 { + type Error = ScalarConversionError; + fn try_from(scalar: Scalar) -> Result { + match scalar { + Scalar::Float(v) => Ok(v), + Scalar::Int(v) => Ok(v as f64), + Scalar::UInt(v) => Ok(v as f64), + _ => Err(ScalarConversionError), + } + } +} + +impl TryFrom for f32 { + type Error = ScalarConversionError; + fn try_from(scalar: Scalar) -> Result { + // Mirror `f64`'s acceptance of integer variants; float reads may be lossy (documented). + match scalar { + Scalar::Float(v) => Ok(v as f32), + Scalar::Int(v) => Ok(v as f32), + Scalar::UInt(v) => Ok(v as f32), + _ => Err(ScalarConversionError), + } + } +} + +impl TryFrom for bool { + type Error = ScalarConversionError; + fn try_from(scalar: Scalar) -> Result { + match scalar { + Scalar::Bool(v) => Ok(v), + _ => Err(ScalarConversionError), + } + } +} + +/// Metadata structure serialized with CBOR +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct Metadata { + /// Tensor descriptors mapped by name for efficient lookup + pub tensors: BTreeMap, + /// Optional additional metadata + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + pub metadata: BTreeMap, + /// Optional typed scalars mapped by name. + /// + /// Defaulted on read for backward compatibility with files written before scalar support. + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + pub scalars: BTreeMap, +} + +/// Individual tensor descriptor +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct TensorDescriptor { + /// Data type of the tensor + pub dtype: DType, + /// Tensor shape dimensions + pub shape: Vec, + /// Byte offsets in data section (start, end) + pub data_offsets: (u64, u64), + /// Parameter ID for training state persistence matching. + /// Generated automatically if not present during loading. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub param_id: Option, +} + +/// Error types for Burnpack operations +#[derive(Debug)] +pub enum Error { + InvalidHeader, + InvalidMagicNumber, + InvalidVersion, + MetadataSerializationError(String), + MetadataDeserializationError(String), + IoError(String), + TensorNotFound(String), + TensorBytesSizeMismatch(String), + ValidationError(String), +} + +impl core::fmt::Display for Error { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + Error::InvalidHeader => write!(f, "Invalid header: insufficient bytes"), + Error::InvalidMagicNumber => write!(f, "Invalid magic number"), + Error::InvalidVersion => write!(f, "Unsupported version"), + Error::MetadataSerializationError(e) => { + write!(f, "Metadata serialization error: {}", e) + } + Error::MetadataDeserializationError(e) => { + write!(f, "Metadata deserialization error: {}", e) + } + Error::IoError(e) => write!(f, "I/O error: {}", e), + Error::TensorNotFound(name) => write!(f, "Tensor not found: {}", name), + Error::TensorBytesSizeMismatch(e) => { + write!(f, "Tensor bytes size mismatch: {}", e) + } + Error::ValidationError(e) => write!(f, "Validation error: {}", e), + } + } +} + +impl core::error::Error for Error {} + +#[cfg(test)] +mod scalar_tests { + use super::*; + + #[test] + fn int_round_trips_through_checked_conversion() { + assert_eq!(i32::try_from(Scalar::from(-5i32)).unwrap(), -5); + assert_eq!(u8::try_from(Scalar::from(200u8)).unwrap(), 200); + assert_eq!(usize::try_from(Scalar::from(42usize)).unwrap(), 42); + } + + #[test] + fn out_of_range_int_conversion_is_rejected() { + // u64 value beyond i32::MAX cannot become i32. + let big = Scalar::from(5_000_000_000u64); + assert!(i32::try_from(big).is_err()); + // Negative value cannot become u32. + assert!(u32::try_from(Scalar::from(-1i32)).is_err()); + // 300 does not fit in u8. + assert!(u8::try_from(Scalar::from(300u32)).is_err()); + } + + #[test] + fn float_and_bool_variant_mismatches_are_rejected() { + // An integer field must not read a stored float. + assert!(i64::try_from(Scalar::Float(1.5)).is_err()); + // A bool field must not read a stored int. + assert!(bool::try_from(Scalar::Int(1)).is_err()); + // A float field must not read a stored bool. + assert!(f64::try_from(Scalar::Bool(true)).is_err()); + } + + #[test] + fn float_accepts_int_variants_symmetrically() { + assert_eq!(f64::try_from(Scalar::Int(3)).unwrap(), 3.0); + assert_eq!(f32::try_from(Scalar::Int(3)).unwrap(), 3.0); + assert_eq!(f64::try_from(Scalar::Float(2.5)).unwrap(), 2.5); + assert_eq!(f32::try_from(Scalar::Float(2.5)).unwrap(), 2.5); + } +} diff --git a/crates/burn-pack/src/lib.rs b/crates/burn-pack/src/lib.rs new file mode 100644 index 0000000..a4a030b --- /dev/null +++ b/crates/burn-pack/src/lib.rs @@ -0,0 +1,130 @@ +#![cfg_attr(not(feature = "std"), no_std)] + +//! # Burn Pack +//! +//! The **burnpack** binary serialization format for the Burn deep learning framework. +//! +//! `burn-pack` is intentionally minimal and tensor-library-agnostic: it depends only on +//! `burn-std` (for [`DType`] / [`Bytes`]), `serde`, and a CBOR codec. It knows how to read +//! and write the burnpack container format but has no notion of Burn modules or tensors. +//! Higher layers (e.g. `burn-core`) bridge between [`Tensor`] entries and their own +//! tensor/snapshot types. +//! +//! Write a pack with [`Writer`], read one with [`Reader`]; both operate on [`Tensor`] +//! entries that carry the format-level metadata plus a lazy provider of the raw +//! little-endian bytes. +//! +//! ``` +//! use burn_pack::{Bytes, DType, Reader, Tensor, Writer}; +//! +//! // A 2x2 f32 tensor, as raw little-endian bytes. +//! let raw: Vec = [1.0f32, 2.0, 3.0, 4.0] +//! .iter() +//! .flat_map(|v| v.to_le_bytes()) +//! .collect(); +//! let tensor = Tensor::new( +//! "weight".to_string(), +//! DType::F32, +//! vec![2, 2], +//! Some(42), // optional param id +//! Bytes::from_bytes_vec(raw), +//! ); +//! +//! // Write to an in-memory buffer ... +//! let packed = Writer::new(vec![tensor]) +//! .with_metadata("producer", "burn-pack docs") +//! .into_bytes() +//! .unwrap(); +//! +//! // ... and read it back. +//! let reader = Reader::from_bytes(packed).unwrap(); +//! assert_eq!(reader.metadata()["producer"], "burn-pack docs"); +//! // Consume the reader to get the tensors (zero-copy views into the source). +//! let tensors = reader.into_tensors().unwrap(); +//! assert_eq!(tensors.len(), 1); +//! assert_eq!(tensors[0].name, "weight"); +//! assert_eq!(tensors[0].shape.to_vec(), vec![2, 2]); +//! assert_eq!(tensors[0].param_id, Some(42)); +//! ``` +//! +//! # File format +//! +//! A burnpack file has three parts: a fixed-size header, a CBOR metadata blob, and a +//! 256-byte-aligned tensor data section. All multi-byte integers are little-endian. +//! +//! ```text +//! ┌──────────────────────────────────────────────────────────────┐ +//! │ Header — 10 bytes ([`HEADER_SIZE`]) │ +//! │ magic : u32 — 0x4255524E "BURN" ([`MAGIC_NUMBER`]) │ +//! │ version : u16 — format version ([`FORMAT_VERSION`]) │ +//! │ metadata_size : u32 — byte length of the CBOR metadata │ +//! ├──────────────────────────────────────────────────────────────┤ +//! │ Metadata — CBOR, `metadata_size` bytes │ +//! │ tensors : map │ +//! │ dtype : [`DType`] │ +//! │ shape : list │ +//! │ data_offsets : (start, end) relative to the data section │ +//! │ param_id : optional u64 (training-state identity) │ +//! │ metadata : map user key/value pairs │ +//! ├──────────────────────────────────────────────────────────────┤ +//! │ Padding to the next 256-byte boundary │ +//! │ ([`aligned_data_section_start`]) │ +//! ├──────────────────────────────────────────────────────────────┤ +//! │ Tensor data section │ +//! │ each tensor's bytes start on a 256-byte boundary │ +//! │ ([`TENSOR_ALIGNMENT`]) for aligned, lazy file-backed │ +//! │ loading (see [`Bytes::from_file`]). │ +//! │ tensors sliced zero-copy. │ +//! └──────────────────────────────────────────────────────────────┘ +//! ``` +//! +//! ## Why 256-byte alignment +//! +//! Aligning every tensor to a 256-byte boundary ([`TENSOR_ALIGNMENT`]) lets a reader +//! memory-map the file and hand out tensor slices without copying, while satisfying the +//! alignment requirements of every element type (including 8-byte `f64`), cache lines, +//! and GPU coalesced access. 256 bytes matches the choice made by GGUF, MLX, ncnn, and +//! other major formats. +//! +//! ## Safety limits +//! +//! Reading is hardened against malicious or corrupt inputs. The reader rejects files +//! that exceed any of the following before allocating for them: +//! +//! - [`MAX_METADATA_SIZE`] — largest CBOR metadata blob +//! - [`MAX_TENSOR_COUNT`] — largest number of tensors +//! - [`MAX_TENSOR_SIZE`] — largest single tensor +//! - [`MAX_CBOR_RECURSION_DEPTH`] — deepest CBOR nesting (stack-overflow guard) +//! - [`MAX_FILE_SIZE`] — largest file accepted by the file loaders (std only) +//! +//! It also validates that the file is large enough to contain every tensor it claims, +//! returning [`Error::ValidationError`] otherwise. +//! +//! ## Feature Flags +//! +//! - `std`: Enables file I/O ([`Reader::from_file`] / [`Writer::write_to_file`]) (default) + +extern crate alloc; + +mod base; +mod reader; +mod tensor; +mod writer; + +#[cfg(feature = "std")] +pub use base::MAX_FILE_SIZE; +pub use base::{ + Error, FORMAT_VERSION, HEADER_SIZE, Header, MAGIC_NUMBER, MAX_CBOR_RECURSION_DEPTH, + MAX_METADATA_SIZE, MAX_TENSOR_COUNT, MAX_TENSOR_SIZE, Scalar, ScalarConversionError, + TENSOR_ALIGNMENT, aligned_data_section_start, +}; +pub use reader::Reader; +pub use tensor::Tensor; +pub use writer::Writer; + +/// The canonical file extension for burnpack files (without the leading dot). +pub const EXTENSION: &str = "bpk"; + +// Re-export the core types so callers can build [`Tensor`] entries and inspect descriptors +// without depending on `burn-std` directly. +pub use burn_std::{Bytes, DType, Shape}; diff --git a/crates/burn-pack/src/reader.rs b/crates/burn-pack/src/reader.rs new file mode 100644 index 0000000..cc252ec --- /dev/null +++ b/crates/burn-pack/src/reader.rs @@ -0,0 +1,382 @@ +use super::base::{ + Error, FORMAT_VERSION, HEADER_SIZE, Header, MAX_CBOR_RECURSION_DEPTH, MAX_METADATA_SIZE, + MAX_TENSOR_COUNT, MAX_TENSOR_SIZE, Metadata, TensorDescriptor, aligned_data_section_start, +}; +use super::tensor::Tensor; +use alloc::format; +use alloc::string::{String, ToString}; +use alloc::vec::Vec; +use burn_std::{Bytes, Shape}; + +#[cfg(feature = "std")] +use super::base::MAX_FILE_SIZE; +#[cfg(feature = "std")] +use alloc::vec; +#[cfg(feature = "std")] +use std::fs::File; +#[cfg(feature = "std")] +use std::io::Read; +#[cfg(feature = "std")] +use std::path::Path; + +/// Reader for loading burnpack containers. +pub struct Reader { + metadata: Metadata, + source: Source, + /// Absolute byte offset where the (256-byte aligned) tensor data section starts. + data_offset: usize, +} + +impl Reader { + /// Load a pack from an in-memory [`Bytes`] buffer. + /// + /// Loading is lazy: only the header and metadata are parsed here. The buffer is kept as-is + /// (no copy, no share) and is only turned into zero-copy [`Bytes::view`] windows when you + /// consume the reader with [`into_tensors`](Self::into_tensors). For large models, prefer + /// [`from_file`](Self::from_file), which keeps tensor data file-backed and lazy. + pub fn from_bytes(bytes: Bytes) -> Result { + let header = read_header(&bytes)?; + let metadata_end = HEADER_SIZE + .checked_add(header.metadata_size as usize) + .ok_or(Error::InvalidHeader)?; + if bytes.len() < metadata_end { + return Err(Error::InvalidHeader); + } + let metadata = parse_metadata(&bytes[HEADER_SIZE..metadata_end])?; + + let available = bytes.len(); + Self::assemble(&header, metadata, Source::Memory(bytes), available) + } + + /// Load a pack from a file. + /// + /// Only the header and metadata are read up front; the whole file is wrapped in a single + /// lazy [`Bytes::from_file`] source, and each tensor's data is a [`Bytes::view`] window into + /// it, read from disk only when accessed. This integrates with the Burn ecosystem's file + /// allocation (pinned-memory staging) for fast file-to-GPU transfers. + /// + /// If `path` has no extension and does not exist as given, the canonical + /// [`crate::EXTENSION`] (`.bpk`) is appended. + #[cfg(feature = "std")] + pub fn from_file>(path: P) -> Result { + let path = path.as_ref(); + let path = if path.extension().is_none() && !path.exists() { + path.with_extension(crate::EXTENSION) + } else { + path.to_path_buf() + }; + + let mut file = File::open(&path).map_err(io_err)?; + + let file_size = file.metadata().map_err(io_err)?.len(); + if file_size > MAX_FILE_SIZE { + return Err(Error::ValidationError(format!( + "File size {file_size} bytes exceeds maximum allowed size of {MAX_FILE_SIZE} bytes" + ))); + } + + let mut header_bytes = [0u8; HEADER_SIZE]; + file.read_exact(&mut header_bytes).map_err(io_err)?; + let header = read_header(&header_bytes)?; + + let mut metadata_bytes = vec![0u8; header.metadata_size as usize]; + file.read_exact(&mut metadata_bytes).map_err(io_err)?; + let metadata = parse_metadata(&metadata_bytes)?; + + let source = Source::File(Bytes::from_file(path.as_path(), file_size, 0)); + Self::assemble(&header, metadata, source, file_size as usize) + } + + /// Finish construction once the header, metadata, and data source are known. + /// + /// Centralizes the truncation check and the aligned data-section offset so both + /// [`from_bytes`](Self::from_bytes) and [`from_file`](Self::from_file) stay in sync. + /// `available` is the number of bytes the source can actually supply. + fn assemble( + header: &Header, + metadata: Metadata, + source: Source, + available: usize, + ) -> Result { + let metadata_end = HEADER_SIZE + header.metadata_size as usize; + validate_total_size(&metadata, metadata_end, available)?; + + Ok(Self { + metadata, + source, + data_offset: aligned_data_section_start(header.metadata_size as usize), + }) + } + + /// Consume the reader, returning all tensors in sorted (alphabetical) name order. + /// + /// Each tensor's bytes are a zero-copy [`Bytes::view`] window into the source — no per-tensor + /// copy. For an in-memory source the buffer is [shared](Bytes::shared) once here (a cheap + /// `Arc` move, never a data copy) to make those views available; for a file source the windows + /// are file-backed and read lazily on access. Consuming `self` lets us hand the source's + /// ownership to the views directly, so loading never reads or copies tensor data eagerly. + pub fn into_tensors(self) -> Result, Error> { + let Reader { + metadata, + source, + data_offset, + } = self; + + // Make the source view-capable: a plain in-memory buffer has no zero-copy window until + // it's shared behind an `Arc`, whereas a file-backed source already windows lazily (and + // must NOT be shared, or every view would materialize the whole file). + let source = match source { + Source::Memory(bytes) => bytes.shared(), + #[cfg(feature = "std")] + Source::File(bytes) => bytes, + }; + + let mut tensors = Vec::with_capacity(metadata.tensors.len()); + for (name, descriptor) in &metadata.tensors { + let (start, end) = tensor_range(data_offset, name, descriptor)?; + let bytes = source.view(start, end).map_err(|_| { + Error::ValidationError(format!( + "Tensor '{name}' data range {start}..{end} could not be viewed (source is {} bytes)", + source.len() + )) + })?; + tensors.push(make_tensor(name, descriptor, bytes)?); + } + Ok(tensors) + } + + /// The user-supplied key/value metadata stored alongside the tensors. + /// + /// For per-tensor info (dtype/shape/param id), use [`into_tensors`](Self::into_tensors) — for a + /// file-backed reader that does not read any tensor data until a tensor's bytes are accessed. + pub fn metadata(&self) -> &alloc::collections::BTreeMap { + &self.metadata.metadata + } + + /// The typed scalars stored alongside the tensors. + /// + /// Empty for files written before scalar support. + pub fn scalars(&self) -> &alloc::collections::BTreeMap { + &self.metadata.scalars + } + + /// The names of all tensors in the pack, in sorted (alphabetical) order. + pub fn tensor_names(&self) -> Vec<&str> { + self.metadata.tensors.keys().map(|n| n.as_str()).collect() + } + + /// Read a single tensor's raw little-endian bytes by name (always copies). + /// + /// Returns [`Error::TensorNotFound`] if no tensor with that name exists. + pub fn tensor_data(&self, name: &str) -> Result, Error> { + let descriptor = self + .metadata + .tensors + .get(name) + .ok_or_else(|| Error::TensorNotFound(name.to_string()))?; + let (start, end) = tensor_range(self.data_offset, name, descriptor)?; + + match &self.source { + #[cfg(feature = "std")] + Source::File(bytes) => { + // A file-backed view reads just this tensor's range from disk. + let view = bytes.view(start, end).map_err(|_| { + Error::ValidationError(format!( + "Tensor '{name}' data range {start}..{end} could not be viewed" + )) + })?; + let slice: &[u8] = &view; + Ok(slice.to_vec()) + } + Source::Memory(bytes) => Ok(memory_chunk(bytes, start, end)?.to_vec()), + } + } +} + +/// Compute and validate the absolute `[start, end)` byte range of a tensor. +fn tensor_range( + data_offset: usize, + name: &str, + descriptor: &TensorDescriptor, +) -> Result<(usize, usize), Error> { + let to_usize = |offset: u64| -> Result { + offset.try_into().map_err(|_| { + Error::ValidationError(format!( + "Tensor '{name}' has corrupted offset data: offset {offset} exceeds platform maximum" + )) + }) + }; + let overflow = || { + Error::ValidationError(format!( + "Tensor '{name}' has corrupted offset data: overflow" + )) + }; + + let start = data_offset + .checked_add(to_usize(descriptor.data_offsets.0)?) + .ok_or_else(overflow)?; + let end = data_offset + .checked_add(to_usize(descriptor.data_offsets.1)?) + .ok_or_else(overflow)?; + + if end < start { + return Err(Error::ValidationError(format!( + "Tensor '{name}' has corrupted offset data: end {end} < start {start}" + ))); + } + if end - start > MAX_TENSOR_SIZE { + return Err(Error::ValidationError(format!( + "Tensor '{name}' size {} exceeds maximum allowed size of {MAX_TENSOR_SIZE} bytes (potential DoS attack)", + end - start + ))); + } + Ok((start, end)) +} + +/// Parse and validate a header from a buffer that starts with it. +fn read_header(buf: &[u8]) -> Result { + if buf.len() < HEADER_SIZE { + return Err(Error::InvalidHeader); + } + let header = Header::from_bytes(&buf[..HEADER_SIZE])?; + if header.version > FORMAT_VERSION { + return Err(Error::InvalidVersion); + } + if header.metadata_size > MAX_METADATA_SIZE { + return Err(Error::ValidationError(format!( + "Metadata size {} exceeds maximum allowed size of {MAX_METADATA_SIZE} bytes (potential DoS attack)", + header.metadata_size + ))); + } + Ok(header) +} + +/// Deserialize the CBOR metadata and validate the tensor count. +fn parse_metadata(bytes: &[u8]) -> Result { + let metadata: Metadata = + ciborium::de::from_reader_with_recursion_limit(bytes, MAX_CBOR_RECURSION_DEPTH) + .map_err(|e| Error::MetadataDeserializationError(e.to_string()))?; + if metadata.tensors.len() > MAX_TENSOR_COUNT { + return Err(Error::ValidationError(format!( + "File contains {} tensors, exceeding maximum of {MAX_TENSOR_COUNT} (potential DoS attack)", + metadata.tensors.len() + ))); + } + Ok(metadata) +} + +/// Ensure the available bytes can hold every tensor the metadata claims. +fn validate_total_size( + metadata: &Metadata, + metadata_end: usize, + available: usize, +) -> Result<(), Error> { + if metadata.tensors.is_empty() { + return Ok(()); + } + let max_offset = metadata + .tensors + .values() + .map(|t| t.data_offsets.1) + .max() + .unwrap_or(0); + let max_offset: usize = max_offset.try_into().map_err(|_| { + Error::ValidationError(format!("Data offset {max_offset} exceeds platform maximum")) + })?; + let min_size = metadata_end + .checked_add(max_offset) + .ok_or_else(|| Error::ValidationError("File size calculation overflow".into()))?; + if available < min_size { + return Err(Error::ValidationError(format!( + "File truncated: expected at least {min_size} bytes, got {available} bytes" + ))); + } + Ok(()) +} + +/// Borrow a tensor's `[start, end)` range out of an in-memory buffer, bounds-checked. +fn memory_chunk(source: &Bytes, start: usize, end: usize) -> Result<&[u8], Error> { + let data: &[u8] = source; + data.get(start..end).ok_or_else(|| { + Error::ValidationError(format!( + "Tensor data range {start}..{end} is out of bounds (buffer is {} bytes)", + data.len() + )) + }) +} + +/// Build a [`Tensor`] entry from a descriptor + its data bytes. +fn make_tensor(name: &str, descriptor: &TensorDescriptor, bytes: Bytes) -> Result { + let shape = descriptor + .shape + .iter() + .map(|&s| { + s.try_into().map_err(|_| { + Error::ValidationError(format!( + "Tensor '{name}' has corrupted shape data: dimension {s} exceeds platform maximum" + )) + }) + }) + .collect::, Error>>()?; + + Ok(Tensor::new( + name.to_string(), + descriptor.dtype, + Shape::from(shape), + descriptor.param_id, + bytes, + )) +} + +/// Where a [`Reader`] gets its tensor data from. +/// +/// Both variants hold a single [`Bytes`] spanning the whole container, and tensors are carved +/// out of it with zero-copy [`Bytes::view`] windows. +enum Source { + /// The whole pack lives in memory. + Memory(Bytes), + /// The pack lives in a file; tensor data is read lazily via file-backed [`Bytes::view`]. + #[cfg(feature = "std")] + File(Bytes), +} + +#[cfg(feature = "std")] +fn io_err(e: std::io::Error) -> Error { + Error::IoError(e.to_string()) +} + +// Verifies the on-disk layout invariant (256-byte tensor alignment), which needs access to the +// internal `TensorDescriptor` offsets. Public round-trip/error tests live in `tests/`. +#[cfg(all(test, feature = "std"))] +mod tests { + use super::*; + use crate::{TENSOR_ALIGNMENT, Tensor, Writer}; + use burn_std::DType; + + fn tensor(name: &str, elems: usize) -> Tensor { + Tensor::new( + name.to_string(), + DType::F32, + alloc::vec![elems], + None, + Bytes::from_bytes_vec(alloc::vec![0u8; elems * 4]), + ) + } + + #[test] + fn tensor_offsets_are_256_aligned() { + // Odd sizes force the writer to insert padding between tensors. + let packed = Writer::new(vec![tensor("a", 3), tensor("b", 1), tensor("c", 2)]) + .into_bytes() + .unwrap(); + let reader = Reader::from_bytes(packed).unwrap(); + + for (name, descriptor) in &reader.metadata.tensors { + assert_eq!( + descriptor.data_offsets.0 % TENSOR_ALIGNMENT, + 0, + "tensor '{name}' start offset is not 256-aligned" + ); + } + } +} diff --git a/crates/burn-pack/src/tensor.rs b/crates/burn-pack/src/tensor.rs new file mode 100644 index 0000000..eac913f --- /dev/null +++ b/crates/burn-pack/src/tensor.rs @@ -0,0 +1,56 @@ +//! Tensor-library-agnostic tensor entry for the burnpack format. +//! +//! The burnpack reader and writer operate on [`Tensor`] values, which carry only the +//! format-level information (name, dtype, shape, optional param id) plus the raw +//! little-endian [`Bytes`]. Keeping the bytes as [`Bytes`] (rather than a custom buffer +//! type) integrates with the rest of the Burn ecosystem: a reader can hand out +//! file-backed bytes ([`Bytes::from_file`]) for fast, lazy file-to-GPU transfers, while a +//! writer simply consumes already-materialized bytes. + +use alloc::string::String; + +use burn_std::{Bytes, DType, Shape}; + +/// A single tensor in a burnpack container, decoupled from any tensor library. +/// +/// The [`bytes`](Self::bytes) field holds the tensor's data in little-endian layout, +/// matching the element count implied by [`shape`](Self::shape) and [`dtype`](Self::dtype). +/// When produced by a [`Reader`](crate::Reader) loading from a file, the bytes are +/// file-backed and only read from disk when accessed. +#[derive(Clone)] +pub struct Tensor { + /// Fully-qualified tensor name (e.g. `"encoder.layer1.weight"`). + pub name: String, + /// Data type of the tensor. + pub dtype: DType, + /// Tensor shape. + pub shape: Shape, + /// Optional parameter id, used to preserve identities for stateful training. + pub param_id: Option, + /// The tensor's raw little-endian bytes. + pub bytes: Bytes, +} + +impl Tensor { + /// Create a tensor entry from its metadata and raw bytes. + pub fn new( + name: String, + dtype: DType, + shape: impl Into, + param_id: Option, + bytes: Bytes, + ) -> Self { + Self { + name, + dtype, + shape: shape.into(), + param_id, + bytes, + } + } + + /// Number of raw bytes the tensor occupies. + pub fn byte_len(&self) -> usize { + self.bytes.len() + } +} diff --git a/crates/burn-pack/src/writer.rs b/crates/burn-pack/src/writer.rs new file mode 100644 index 0000000..c2a1120 --- /dev/null +++ b/crates/burn-pack/src/writer.rs @@ -0,0 +1,411 @@ +use super::base::{ + Error, FORMAT_VERSION, HEADER_SIZE, Header, MAGIC_NUMBER, Metadata, Scalar, TENSOR_ALIGNMENT, + TensorDescriptor, aligned_data_section_start, +}; +use super::tensor::Tensor; +use alloc::collections::BTreeMap; +use alloc::format; +use alloc::string::{String, ToString}; +use alloc::vec; +use alloc::vec::Vec; +use burn_std::Bytes; + +#[cfg(feature = "std")] +use std::fs::File; +#[cfg(feature = "std")] +use std::io::{Read, Write}; +#[cfg(feature = "std")] +use std::path::Path; + +/// Align an offset to the specified alignment boundary. +/// +/// Returns the smallest value >= `offset` that is a multiple of `alignment`. +#[inline] +const fn align_offset(offset: u64, alignment: u64) -> u64 { + offset.div_ceil(alignment) * alignment +} + +/// Maximum number of bytes materialized from a single tensor at a time while +/// streaming its data into a [`Sink`]. +/// +/// Large device-resident tensors are read back to host memory lazily, one +/// [`Bytes::view`] window at a time, instead of all at once. This keeps the +/// transient (often pinned) host staging buffer bounded by this size regardless +/// of how large the tensor is. The value is a multiple of [`TENSOR_ALIGNMENT`] +/// so each window starts on an aligned device offset. +const WRITE_CHUNK_SIZE: usize = 8 * 1024 * 1024; + +/// Writer for creating Burnpack files +pub struct Writer { + /// Tensors to write + pub(crate) tensors: Vec, + /// Metadata key-value pairs + pub(crate) metadata: BTreeMap, + /// Typed scalars keyed by name + pub(crate) scalars: BTreeMap, +} + +impl Writer { + /// Create a new writer + pub fn new(tensors: Vec) -> Self { + Self { + tensors, + metadata: BTreeMap::new(), + scalars: BTreeMap::new(), + } + } + + /// Builder pattern: add metadata and return self + pub fn with_metadata(mut self, key: &str, value: &str) -> Self { + self.metadata.insert(key.to_string(), value.to_string()); + self + } + + /// Builder pattern: add a typed scalar and return self. + pub fn with_scalar(mut self, key: &str, value: Scalar) -> Self { + self.scalars.insert(key.to_string(), value); + self + } + + /// Calculate the total size needed for the burnpack data. + /// + /// This is useful when you want to pre-allocate a buffer for `write_into()`. + /// The size includes padding bytes for both metadata alignment and tensor alignment. + pub fn size(&self) -> Result { + Ok(self.plan()?.total_size()) + } + + /// Write burnpack data into a caller-provided buffer. + /// + /// The buffer must be large enough to hold all data. Use `size()` to determine + /// the required buffer size. If the buffer is too small, this will return an error. + /// + /// This allows the caller to control buffer allocation, enabling optimizations like: + /// - Buffer reuse across multiple writes + /// - Custom allocators + /// - Pinned memory for GPU transfers + /// + /// # Arguments + /// + /// * `buffer` - Mutable slice to write data into. Must be at least `size()` bytes. + pub fn write_into(self, buffer: &mut [u8]) -> Result<(), Error> { + let layout = self.plan()?; + let total_size = layout.total_size(); + + if buffer.len() < total_size { + return Err(Error::IoError(format!( + "Buffer too small: need {} bytes, got {} bytes", + total_size, + buffer.len() + ))); + } + + let mut sink = BufferSink { buffer, offset: 0 }; + self.write_container(&layout, &mut sink) + } + + /// Write to a byte buffer (convenience method). + /// + /// This allocates a buffer internally and writes the burnpack data. + /// For more control over buffer allocation, use `size()` + `write_into()`. + pub fn into_bytes(self) -> Result { + let layout = self.plan()?; + let mut buffer = vec![0u8; layout.total_size()]; + + let mut sink = BufferSink { + buffer: &mut buffer, + offset: 0, + }; + self.write_container(&layout, &mut sink)?; + + Ok(Bytes::from_bytes_vec(buffer)) + } + + /// Write directly to a file (more memory efficient for large models). + /// + /// If `path` has no extension, the canonical [`crate::EXTENSION`] (`.bpk`) is appended. + #[cfg(feature = "std")] + pub fn write_to_file>(self, path: P) -> Result<(), Error> { + let path = path.as_ref(); + let path = if path.extension().is_none() { + path.with_extension(crate::EXTENSION) + } else { + path.to_path_buf() + }; + + let layout = self.plan()?; + let file = File::create(path).map_err(|e| Error::IoError(e.to_string()))?; + + let mut sink = FileSink { file }; + self.write_container(&layout, &mut sink)?; + + sink.file.flush().map_err(|e| Error::IoError(e.to_string())) + } + + /// Build the complete on-disk layout: header, serialized metadata, and the + /// position and size of the (aligned) tensor data section. + fn plan(&self) -> Result { + let (metadata, metadata_bytes, data_size) = self.build_metadata()?; + + let metadata_size: u32 = metadata_bytes.len().try_into().map_err(|_| { + Error::IoError(format!( + "Metadata size {} exceeds maximum of {} bytes", + metadata_bytes.len(), + u32::MAX + )) + })?; + + let header = Header { + magic: MAGIC_NUMBER, + version: FORMAT_VERSION, + metadata_size, + }; + + let data_section_start = aligned_data_section_start(metadata_bytes.len()); + + Ok(Layout { + metadata, + metadata_bytes, + header, + data_section_start, + data_size, + }) + } + + /// Serialize the metadata structure (tensor descriptors + key-value pairs) to CBOR. + /// + /// Also returns the size of the tensor data section, computed while assigning offsets. + fn build_metadata(&self) -> Result<(Metadata, Vec, usize), Error> { + let (tensors, data_size) = self.build_descriptors()?; + let metadata = Metadata { + tensors, + metadata: self.metadata.clone(), + scalars: self.scalars.clone(), + }; + + let mut metadata_bytes = Vec::new(); + ciborium::ser::into_writer(&metadata, &mut metadata_bytes) + .map_err(|e| Error::MetadataSerializationError(e.to_string()))?; + + Ok((metadata, metadata_bytes, data_size)) + } + + /// Build tensor descriptors, assigning each tensor an aligned offset within + /// the data section so that absolute file positions are mmap-friendly. + /// + /// Returns the descriptors plus the total data-section size — the running offset after the + /// last tensor. Offsets only grow, so this is also the highest descriptor end offset. + fn build_descriptors(&self) -> Result<(BTreeMap, usize), Error> { + let mut tensors = BTreeMap::new(); + let mut current_offset = 0u64; + + for tensor in &self.tensors { + let data_len = tensor.bytes.len() as u64; + + // Align the start offset for mmap zero-copy support. + let aligned_start = align_offset(current_offset, TENSOR_ALIGNMENT); + let end = aligned_start.checked_add(data_len).ok_or_else(|| { + Error::IoError(format!( + "Tensor offset overflow: {} + {} exceeds maximum", + aligned_start, data_len + )) + })?; + + // Descriptors are keyed by name, but the tensor data is written from the + // (ordered) `self.tensors` list. A duplicate name would collapse to a single + // descriptor while still writing two data blocks, corrupting the container. + if tensors + .insert( + tensor.name.clone(), + TensorDescriptor { + dtype: tensor.dtype, + shape: tensor.shape.iter().map(|&s| s as u64).collect(), + data_offsets: (aligned_start, end), + param_id: tensor.param_id, + }, + ) + .is_some() + { + return Err(Error::ValidationError(format!( + "Duplicate tensor name '{}'", + tensor.name + ))); + } + + current_offset = end; + } + + Ok((tensors, current_offset as usize)) + } + + /// Emit the full container — header, metadata, alignment padding, then tensor data + /// — into `sink`, which decides where the bytes ultimately land. + fn write_container(self, layout: &Layout, sink: &mut impl Sink) -> Result<(), Error> { + sink.write(&layout.header.into_bytes())?; + sink.write(&layout.metadata_bytes)?; + + // Pad so the data section starts at its aligned position. + let unaligned_data_start = HEADER_SIZE + layout.metadata_bytes.len(); + if layout.data_section_start > unaligned_data_start { + sink.pad(layout.data_section_start - unaligned_data_start)?; + } + + self.write_tensors(&layout.metadata, sink) + } + + /// Write each tensor's data into `sink`, inserting alignment padding between + /// tensors so every tensor lands at its descriptor's aligned offset. + fn write_tensors(self, metadata: &Metadata, sink: &mut impl Sink) -> Result<(), Error> { + // Position within the data section (relative to its aligned start). + let mut data_offset = 0usize; + + for tensor in self.tensors.into_iter() { + let (aligned_offset, data) = Self::resolve_tensor(tensor, metadata)?; + + if aligned_offset > data_offset { + sink.pad(aligned_offset - data_offset)?; + data_offset = aligned_offset; + } + + Self::write_tensor_data(&data, sink)?; + data_offset += data.len(); + } + + Ok(()) + } + + /// Stream a single tensor's bytes into `sink`, materializing at most + /// [`WRITE_CHUNK_SIZE`] bytes at a time. + /// + /// When the backing supports zero-copy windows — device-resident + /// ([lazy](burn_std::Bytes) device readback), file, or shared buffers — each + /// chunk is taken as a [`Bytes::view`] and read just-in-time, then dropped + /// before the next one. A large device tensor is therefore copied to host in + /// bounded pieces rather than through one big (pinned) staging buffer, so the + /// whole tensor never has to be resident at once. + /// + /// Backings without a zero-copy window (e.g. a plain heap `Vec`) are already + /// host-resident, so [`Bytes::view`] reports it can't window them and the + /// remaining bytes are written in a single pass. + fn write_tensor_data(data: &Bytes, sink: &mut impl Sink) -> Result<(), Error> { + let len = data.len(); + let mut offset = 0; + + while offset < len { + let end = (offset + WRITE_CHUNK_SIZE).min(len); + match data.view(offset, end) { + Ok(chunk) => { + sink.write(&chunk)?; + offset = end; + } + // No zero-copy window available (already host-resident): write + // whatever remains in one shot. View support is a property of the + // backing, so this only ever happens on the first iteration. + Err(_) => { + sink.write(&data[offset..])?; + break; + } + } + } + + Ok(()) + } + + /// Look up a tensor's aligned offset from the metadata and validate that its + /// bytes match the length the descriptor reserved for it. + fn resolve_tensor(tensor: Tensor, metadata: &Metadata) -> Result<(usize, Bytes), Error> { + let descriptor = metadata.tensors.get(&tensor.name).ok_or_else(|| { + Error::IoError(format!( + "Internal error: tensor '{}' not found in metadata", + tensor.name + )) + })?; + + let (start, end) = descriptor.data_offsets; + let declared_len = (end - start) as usize; + let actual_len = tensor.bytes.len(); + if actual_len != declared_len { + return Err(Error::TensorBytesSizeMismatch(format!( + "tensor '{}' has inconsistent length (expected {}, got {})", + tensor.name, declared_len, actual_len + ))); + } + + Ok((start as usize, tensor.bytes)) + } +} + +/// The computed on-disk layout of a burnpack container. +/// +/// Captures everything needed to emit the bytes: the serialized metadata, the +/// header, where the aligned data section begins, and how large it is. Built once +/// via [`Writer::plan`] and shared by `size`, `write_into`, `to_bytes`, and +/// `write_to_file`. +struct Layout { + metadata: Metadata, + metadata_bytes: Vec, + header: Header, + data_section_start: usize, + data_size: usize, +} + +impl Layout { + /// Total number of bytes the container occupies. + fn total_size(&self) -> usize { + self.data_section_start + self.data_size + } +} + +/// A sequential destination for the bytes of a burnpack container. +/// +/// Padding and data are written in order; each implementation advances its own +/// cursor, letting the writer stay agnostic about whether bytes land in a buffer +/// or a file. +trait Sink { + /// Write `count` zero bytes of alignment padding. + fn pad(&mut self, count: usize) -> Result<(), Error>; + /// Write `data` verbatim. + fn write(&mut self, data: &[u8]) -> Result<(), Error>; +} + +/// Sink that copies into a caller-provided buffer. +struct BufferSink<'a> { + buffer: &'a mut [u8], + offset: usize, +} + +impl Sink for BufferSink<'_> { + fn pad(&mut self, count: usize) -> Result<(), Error> { + self.buffer[self.offset..self.offset + count].fill(0); + self.offset += count; + Ok(()) + } + + fn write(&mut self, data: &[u8]) -> Result<(), Error> { + self.buffer[self.offset..self.offset + data.len()].copy_from_slice(data); + self.offset += data.len(); + Ok(()) + } +} + +/// Sink that streams directly to a file. +#[cfg(feature = "std")] +struct FileSink { + file: File, +} + +#[cfg(feature = "std")] +impl Sink for FileSink { + fn pad(&mut self, count: usize) -> Result<(), Error> { + // Stream zeros without allocating a `count`-sized buffer per call. + std::io::copy(&mut std::io::repeat(0).take(count as u64), &mut self.file) + .map(|_| ()) + .map_err(|e| Error::IoError(e.to_string())) + } + + fn write(&mut self, data: &[u8]) -> Result<(), Error> { + self.file + .write_all(data) + .map_err(|e| Error::IoError(e.to_string())) + } +} diff --git a/crates/burn-pack/tests/alignment.rs b/crates/burn-pack/tests/alignment.rs new file mode 100644 index 0000000..3ed6e20 --- /dev/null +++ b/crates/burn-pack/tests/alignment.rs @@ -0,0 +1,61 @@ +//! Alignment guarantees and buffer-sizing (`size`/`write_into`). + +mod common; + +use burn_pack::{HEADER_SIZE, Reader, TENSOR_ALIGNMENT, Writer, aligned_data_section_start}; +use common::{f32_tensor, read_f32}; + +#[test] +fn data_section_start_is_aligned() { + let align = TENSOR_ALIGNMENT as usize; + for metadata_size in [0usize, 1, 10, 245, 246, 247, 4096] { + let start = aligned_data_section_start(metadata_size); + assert_eq!(start % align, 0, "data section start must be 256-aligned"); + assert!( + start >= HEADER_SIZE + metadata_size, + "must clear header+metadata" + ); + assert!( + start < HEADER_SIZE + metadata_size + align, + "minimal padding" + ); + } +} + +// Per-tensor offset alignment is verified by an in-crate unit test (it needs the internal +// tensor descriptors); see `src/reader.rs`. + +#[test] +fn size_matches_to_bytes_length() { + let writer = Writer::new(vec![ + f32_tensor("a", &[1.0, 2.0, 3.0], &[3], None), + f32_tensor("b", &[4.0], &[1], None), + ]); + + let size = writer.size().unwrap(); + let bytes = writer.into_bytes().unwrap(); + assert_eq!(size, bytes.len()); +} + +#[test] +fn write_into_matches_to_bytes_and_round_trips() { + // `write_into` and `into_bytes` each consume the writer, so build one per call. + let make_writer = || Writer::new(vec![f32_tensor("w", &[1.0, 2.0, 3.0, 4.0], &[2, 2], None)]); + + let mut buffer = vec![0u8; make_writer().size().unwrap()]; + make_writer().write_into(&mut buffer).unwrap(); + + let from_to_bytes = make_writer().into_bytes().unwrap(); + assert_eq!(&buffer[..], &from_to_bytes[..]); + + let reader = Reader::from_bytes(burn_pack::Bytes::from_bytes_vec(buffer)).unwrap(); + let tensors = reader.into_tensors().unwrap(); + assert_eq!(read_f32(&tensors[0]), vec![1.0, 2.0, 3.0, 4.0]); +} + +#[test] +fn write_into_rejects_too_small_buffer() { + let writer = Writer::new(vec![f32_tensor("w", &[1.0, 2.0], &[2], None)]); + let mut buffer = vec![0u8; writer.size().unwrap() - 1]; + assert!(writer.write_into(&mut buffer).is_err()); +} diff --git a/crates/burn-pack/tests/common/mod.rs b/crates/burn-pack/tests/common/mod.rs new file mode 100644 index 0000000..bc917f3 --- /dev/null +++ b/crates/burn-pack/tests/common/mod.rs @@ -0,0 +1,44 @@ +//! Shared helpers for the burn-pack format integration tests. +//! +//! Included by multiple test binaries; not every helper is used by every one. +#![allow(dead_code)] + +use burn_pack::{Bytes, DType, Tensor}; + +/// Build an f32 [`Tensor`] entry from values + shape. +pub fn f32_tensor(name: &str, values: &[f32], shape: &[usize], param_id: Option) -> Tensor { + let raw: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); + Tensor::new( + name.to_string(), + DType::F32, + shape.to_vec(), + param_id, + Bytes::from_bytes_vec(raw), + ) +} + +/// Build a [`Tensor`] entry from raw little-endian bytes with an explicit dtype. +pub fn raw_tensor( + name: &str, + dtype: DType, + shape: &[usize], + bytes: Vec, + param_id: Option, +) -> Tensor { + Tensor::new( + name.to_string(), + dtype, + shape.to_vec(), + param_id, + Bytes::from_bytes_vec(bytes), + ) +} + +/// Decode a tensor entry's bytes as f32 values. +pub fn read_f32(tensor: &Tensor) -> Vec { + let slice: &[u8] = &tensor.bytes; + slice + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect() +} diff --git a/crates/burn-pack/tests/errors.rs b/crates/burn-pack/tests/errors.rs new file mode 100644 index 0000000..7ccbd1c --- /dev/null +++ b/crates/burn-pack/tests/errors.rs @@ -0,0 +1,95 @@ +//! Malformed / malicious input handling. + +mod common; + +use burn_pack::{ + Bytes, Error, FORMAT_VERSION, Header, MAGIC_NUMBER, MAX_METADATA_SIZE, Reader, Writer, +}; +use common::f32_tensor; + +fn header_bytes(version: u16, metadata_size: u32) -> Bytes { + let header = Header { + magic: MAGIC_NUMBER, + version, + metadata_size, + }; + Bytes::from_bytes_vec(header.into_bytes().to_vec()) +} + +#[test] +fn rejects_too_short_input() { + assert!(matches!( + Reader::from_bytes(Bytes::from_bytes_vec(vec![0u8; 4])), + Err(Error::InvalidHeader) + )); +} + +#[test] +fn rejects_bad_magic() { + let mut bytes = vec![0u8; 10]; + bytes[..4].copy_from_slice(&0xDEAD_BEEFu32.to_le_bytes()); + assert!(matches!( + Reader::from_bytes(Bytes::from_bytes_vec(bytes)), + Err(Error::InvalidMagicNumber) + )); +} + +#[test] +fn rejects_future_version() { + let bytes = header_bytes(FORMAT_VERSION + 1, 0); + assert!(matches!( + Reader::from_bytes(bytes), + Err(Error::InvalidVersion) + )); +} + +#[test] +fn rejects_oversized_metadata_claim() { + // The reader bails out on the metadata-size claim before allocating for it. + let bytes = header_bytes(FORMAT_VERSION, MAX_METADATA_SIZE + 1); + assert!(matches!( + Reader::from_bytes(bytes), + Err(Error::ValidationError(_)) + )); +} + +#[test] +fn rejects_metadata_size_past_eof() { + // Header claims more metadata than the buffer actually contains. + let bytes = header_bytes(FORMAT_VERSION, 4096); + assert!(Reader::from_bytes(bytes).is_err()); +} + +#[test] +fn rejects_duplicate_tensor_names() { + // Descriptors are keyed by name but data is written from the tensor list: a duplicate + // name must be rejected up front, not silently corrupt the container. + let writer = Writer::new(vec![ + f32_tensor("w", &[1.0, 2.0], &[2], None), + f32_tensor("w", &[3.0, 4.0], &[2], None), + ]); + assert!(matches!( + writer.into_bytes(), + Err(Error::ValidationError(_)) + )); +} + +#[test] +fn rejects_truncated_data_section() { + // A valid pack, truncated well into its data section, must be rejected (not silently + // read). Use a large tensor and drop half the file so we are unambiguously below the + // size the metadata claims. + let values: Vec = (0..512).map(|i| i as f32).collect(); + let packed = Writer::new(vec![f32_tensor("w", &values, &[512], None)]) + .into_bytes() + .unwrap(); + + let slice: &[u8] = &packed; + let mut bytes = slice.to_vec(); + bytes.truncate(bytes.len() / 2); + + assert!(matches!( + Reader::from_bytes(Bytes::from_bytes_vec(bytes)), + Err(Error::ValidationError(_)) + )); +} diff --git a/crates/burn-pack/tests/header.rs b/crates/burn-pack/tests/header.rs new file mode 100644 index 0000000..a69f869 --- /dev/null +++ b/crates/burn-pack/tests/header.rs @@ -0,0 +1,44 @@ +//! Header encoding / decoding and format constants. + +use burn_pack::{Error, FORMAT_VERSION, HEADER_SIZE, Header, MAGIC_NUMBER}; + +#[test] +fn header_constants() { + assert_eq!(HEADER_SIZE, 10); + // "BURN" in ASCII. + assert_eq!(MAGIC_NUMBER, 0x4255524E); + assert_eq!(MAGIC_NUMBER.to_le_bytes(), *b"NRUB"); +} + +#[test] +fn header_round_trip() { + let header = Header::new(1234); + let bytes = header.into_bytes(); + assert_eq!(bytes.len(), HEADER_SIZE); + + let decoded = Header::from_bytes(&bytes).unwrap(); + assert_eq!(decoded.magic, MAGIC_NUMBER); + assert_eq!(decoded.version, FORMAT_VERSION); + assert_eq!(decoded.metadata_size, 1234); +} + +#[test] +fn header_rejects_bad_magic() { + let bad = Header { + magic: 0xDEAD_BEEF, + version: FORMAT_VERSION, + metadata_size: 0, + }; + assert!(matches!( + Header::from_bytes(&bad.into_bytes()), + Err(Error::InvalidMagicNumber) + )); +} + +#[test] +fn header_rejects_short_input() { + assert!(matches!( + Header::from_bytes(&[0u8; 4]), + Err(Error::InvalidHeader) + )); +} diff --git a/crates/burn-pack/tests/round_trip.rs b/crates/burn-pack/tests/round_trip.rs new file mode 100644 index 0000000..a6705e4 --- /dev/null +++ b/crates/burn-pack/tests/round_trip.rs @@ -0,0 +1,308 @@ +//! Write → read round-trip coverage for the burnpack format. + +mod common; + +use burn_pack::{Bytes, DType, Reader, Scalar, Tensor, Writer}; +use common::{f32_tensor, raw_tensor, read_f32}; + +#[test] +fn single_tensor_round_trip() { + let tensor = f32_tensor("weight", &[1.0, 2.0, 3.0, 4.0], &[2, 2], Some(7)); + + let packed = Writer::new(vec![tensor]).into_bytes().unwrap(); + let reader = Reader::from_bytes(packed).unwrap(); + let tensors = reader.into_tensors().unwrap(); + + assert_eq!(tensors.len(), 1); + let t = &tensors[0]; + assert_eq!(t.name, "weight"); + assert_eq!(t.dtype, DType::F32); + assert_eq!(t.shape.to_vec(), vec![2, 2]); + assert_eq!(t.param_id, Some(7)); + assert_eq!(t.byte_len(), 16); + assert_eq!(read_f32(t), vec![1.0, 2.0, 3.0, 4.0]); +} + +#[test] +fn multiple_tensors_returned_sorted_by_name() { + // Insert out of order; the reader yields them sorted (BTreeMap ordering). + let packed = Writer::new(vec![ + f32_tensor("zebra", &[9.0], &[1], None), + f32_tensor("alpha", &[1.0], &[1], None), + f32_tensor("mango", &[5.0], &[1], None), + ]) + .into_bytes() + .unwrap(); + + let reader = Reader::from_bytes(packed).unwrap(); + // Read names before consuming the reader, then consume it for the tensors. + assert_eq!(reader.tensor_names(), vec!["alpha", "mango", "zebra"]); + let names: Vec<_> = reader + .into_tensors() + .unwrap() + .iter() + .map(|t| t.name.clone()) + .collect(); + + assert_eq!(names, vec!["alpha", "mango", "zebra"]); +} + +#[test] +fn tensors_with_varied_sizes_map_to_correct_names() { + // On-disk data order = write order (z, a, m), which differs from the name-sorted order the + // reader returns. The odd sizes force different alignment gaps between tensors, so this + // exercises the gap/offset arithmetic of the view-based loader: each tensor's zero-copy + // window must still come back attached to the right name. + let z: Vec = (0..5).map(|i| 100.0 + i as f32).collect(); + let a: Vec = (0..3).map(|i| i as f32).collect(); + let m: Vec = (0..7).map(|i| 50.0 + i as f32).collect(); + + let packed = Writer::new(vec![ + f32_tensor("zebra", &z, &[5], None), + f32_tensor("alpha", &a, &[3], None), + f32_tensor("mango", &m, &[7], None), + ]) + .into_bytes() + .unwrap(); + + let reader = Reader::from_bytes(packed).unwrap(); + let tensors = reader.into_tensors().unwrap(); + + let names: Vec<_> = tensors.iter().map(|t| t.name.clone()).collect(); + assert_eq!(names, vec!["alpha", "mango", "zebra"]); + assert_eq!(read_f32(&tensors[0]), a); + assert_eq!(read_f32(&tensors[1]), m); + assert_eq!(read_f32(&tensors[2]), z); +} + +#[test] +fn user_metadata_round_trip() { + let packed = Writer::new(vec![f32_tensor("w", &[1.0], &[1], None)]) + .with_metadata("producer", "burn-pack") + .with_metadata("format", "burnpack") + .into_bytes() + .unwrap(); + + let reader = Reader::from_bytes(packed).unwrap(); + assert_eq!(reader.metadata()["producer"], "burn-pack"); + assert_eq!(reader.metadata()["format"], "burnpack"); + // Per-tensor info comes from the tensors themselves. + let tensors = reader.into_tensors().unwrap(); + assert_eq!(tensors[0].dtype, DType::F32); + assert_eq!(tensors[0].shape.to_vec(), vec![1]); +} + +#[test] +fn param_id_present_and_absent() { + let packed = Writer::new(vec![ + f32_tensor("with_id", &[1.0], &[1], Some(123)), + f32_tensor("without_id", &[2.0], &[1], None), + ]) + .into_bytes() + .unwrap(); + + let reader = Reader::from_bytes(packed).unwrap(); + let tensors = reader.into_tensors().unwrap(); + let with = tensors.iter().find(|t| t.name == "with_id").unwrap(); + let without = tensors.iter().find(|t| t.name == "without_id").unwrap(); + + assert_eq!(with.param_id, Some(123)); + assert_eq!(without.param_id, None); +} + +#[test] +fn empty_pack() { + let packed = Writer::new(vec![]).into_bytes().unwrap(); + let reader = Reader::from_bytes(packed).unwrap(); + assert!(reader.tensor_names().is_empty()); + assert!(reader.into_tensors().unwrap().is_empty()); +} + +#[test] +fn dtype_and_byte_len_preserved() { + // (dtype, bytes-per-element) + let cases = [ + (DType::F32, 4usize), + (DType::F64, 8), + (DType::I64, 8), + (DType::I32, 4), + (DType::I8, 1), + (DType::U8, 1), + (DType::BF16, 2), + (DType::F16, 2), + ]; + + let n = 3; + let tensors = cases + .iter() + .enumerate() + .map(|(i, (dtype, elem))| { + let bytes = vec![i as u8 + 1; n * elem]; + raw_tensor(&format!("t{i}"), *dtype, &[n], bytes, None) + }) + .collect(); + + let packed = Writer::new(tensors).into_bytes().unwrap(); + let reader = Reader::from_bytes(packed).unwrap(); + let read = reader.into_tensors().unwrap(); + + for (i, (dtype, elem)) in cases.iter().enumerate() { + let t = read.iter().find(|t| t.name == format!("t{i}")).unwrap(); + assert_eq!(t.dtype, *dtype, "dtype preserved for t{i}"); + assert_eq!(t.shape.to_vec(), vec![n]); + assert_eq!(t.byte_len(), n * elem, "byte_len preserved for t{i}"); + let materialized: &[u8] = &t.bytes; + assert_eq!(materialized, &vec![i as u8 + 1; n * elem][..]); + } +} + +#[test] +fn in_memory_round_trip() { + let packed = Writer::new(vec![f32_tensor("w", &[1.5, 2.5, 3.5], &[3], None)]) + .into_bytes() + .unwrap(); + + let reader = Reader::from_bytes(packed).unwrap(); + let tensors = reader.into_tensors().unwrap(); + assert_eq!(read_f32(&tensors[0]), vec![1.5, 2.5, 3.5]); +} + +// Reading from a file backs each tensor with a file-backed `Bytes::view` window, read lazily on +// access. +#[cfg(feature = "std")] +#[test] +fn file_backed_tensors() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("zc.bpk"); + Writer::new(vec![f32_tensor("w", &[1.5, 2.5, 3.5], &[3], None)]) + .write_to_file(&path) + .unwrap(); + + let reader = Reader::from_file(&path).unwrap(); + let tensors = reader.into_tensors().unwrap(); + assert_eq!(read_f32(&tensors[0]), vec![1.5, 2.5, 3.5]); +} + +// A `shared()` tensor exposes zero-copy `view` windows, so the writer streams it +// in bounded chunks rather than materializing it whole. Data larger than the +// writer's internal chunk size exercises the multi-chunk path (several full +// windows plus a partial tail); the concatenated windows must reproduce the +// original bytes exactly. +#[test] +fn shared_tensor_chunked_write_round_trip() { + // Comfortably larger than the writer's 8 MiB chunk size, and not a whole + // multiple of it, so the final chunk is partial. + let len = 8 * 1024 * 1024 + 4096; + let data: Vec = (0..len).map(|i| (i % 251) as u8).collect(); + + let bytes = Bytes::from_bytes_vec(data.clone()).shared(); + let tensor = Tensor::new("big".to_string(), DType::U8, vec![len], None, bytes); + + let packed = Writer::new(vec![tensor]).into_bytes().unwrap(); + let reader = Reader::from_bytes(packed).unwrap(); + let read = reader.into_tensors().unwrap(); + + assert_eq!(read.len(), 1); + let materialized: &[u8] = &read[0].bytes; + assert_eq!(materialized.len(), len); + assert_eq!(materialized, &data[..]); +} + +#[test] +fn read_single_tensor_data_by_name() { + let packed = Writer::new(vec![ + f32_tensor("a", &[1.0, 2.0], &[2], None), + f32_tensor("b", &[3.0], &[1], None), + ]) + .into_bytes() + .unwrap(); + + let reader = Reader::from_bytes(packed).unwrap(); + let raw = reader.tensor_data("a").unwrap(); + let values: Vec = raw + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(); + assert_eq!(values, vec![1.0, 2.0]); + + assert!(reader.tensor_data("missing").is_err()); +} + +#[cfg(feature = "std")] +#[test] +fn file_round_trip() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("model.bpk"); + + Writer::new(vec![f32_tensor( + "weight", + &[1.0, 2.0, 3.0, 4.0], + &[2, 2], + Some(1), + )]) + .with_metadata("producer", "burn-pack") + .write_to_file(&path) + .unwrap(); + + let reader = Reader::from_file(&path).unwrap(); + assert_eq!(reader.metadata()["producer"], "burn-pack"); + let tensors = reader.into_tensors().unwrap(); + assert_eq!(tensors.len(), 1); + assert_eq!(read_f32(&tensors[0]), vec![1.0, 2.0, 3.0, 4.0]); + assert_eq!(tensors[0].param_id, Some(1)); +} + +#[test] +fn extensionless_path_appends_bpk() { + let dir = tempfile::tempdir().unwrap(); + // No extension on the path: the writer should append `.bpk`, and the reader should find it + // when given the same extension-less path. + let path = dir.path().join("model"); + + Writer::new(vec![f32_tensor("weight", &[1.0, 2.0], &[2], Some(7))]) + .write_to_file(&path) + .unwrap(); + + assert!( + path.with_extension("bpk").exists(), + "writer should have created `model.bpk`" + ); + assert!( + !path.exists(), + "no extension-less `model` file should exist" + ); + + let reader = Reader::from_file(&path).unwrap(); + let tensors = reader.into_tensors().unwrap(); + assert_eq!(read_f32(&tensors[0]), vec![1.0, 2.0]); + assert_eq!(tensors[0].param_id, Some(7)); +} + +#[test] +fn typed_scalars_round_trip() { + let packed = Writer::new(vec![f32_tensor("w", &[1.0], &[1], None)]) + .with_scalar("step", Scalar::UInt(42)) + .with_scalar("lr", Scalar::Float(0.001)) + .with_scalar("flag", Scalar::Bool(true)) + .with_scalar("offset", Scalar::Int(-3)) + .into_bytes() + .unwrap(); + + let reader = Reader::from_bytes(packed).unwrap(); + let scalars = reader.scalars(); + assert_eq!(scalars.get("step"), Some(&Scalar::UInt(42))); + assert_eq!(scalars.get("lr"), Some(&Scalar::Float(0.001))); + assert_eq!(scalars.get("flag"), Some(&Scalar::Bool(true))); + assert_eq!(scalars.get("offset"), Some(&Scalar::Int(-3))); +} + +#[test] +fn scalars_absent_by_default() { + // A pack written without scalars exposes an empty scalar map (backward compatible: files + // predating scalar support simply omit the field and default to empty on read). + let packed = Writer::new(vec![f32_tensor("w", &[1.0], &[1], None)]) + .into_bytes() + .unwrap(); + let reader = Reader::from_bytes(packed).unwrap(); + assert!(reader.scalars().is_empty()); +} diff --git a/crates/burn-remote/Cargo.toml b/crates/burn-remote/Cargo.toml new file mode 100644 index 0000000..0dac4a8 --- /dev/null +++ b/crates/burn-remote/Cargo.toml @@ -0,0 +1,81 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science"] +description = "Backend router decorator over the network." +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "data"] +license.workspace = true +name = "burn-remote" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-router-remote" +documentation = "https://docs.rs/burn-router-remote" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["client", "server", "iroh", "websocket"] +doc = [] +tracing = [ + "burn-communication?/tracing", + "burn-ir/tracing", + "burn-router/tracing", + "burn-std/tracing", + "burn-backend/tracing", +] + +client = ["tokio/sync"] +# Cache recurring op-groups on the server and invoke them by id. Wraps `RemoteBackend` in `Fusion`. +fusion = ["client", "dep:burn-fusion", "burn-router/fusion"] +server = ["tokio/sync"] +iroh = ["dep:iroh", "tokio/io-util"] +websocket = [ + "server", + "dep:burn-communication", + "burn-communication/websocket", + "burn-communication/data-service", + "dep:tokio-util", +] + + +[dependencies] +burn-ir = { workspace = true, features = ["default"] } +burn-backend = { workspace = true, features = [ + "default", + "serde", # For DTypeUsageSet to be de/serialized +] } +burn-std = { workspace = true, features = ["default"] } +burn-router = { workspace = true, features = ["default"] } +burn-fusion = { workspace = true, optional = true } +# Native WebSocket transport; optional so the Iroh client stays wasm-compatible. +burn-communication = { workspace = true, optional = true } + +bytes = { workspace = true } +log = { workspace = true } +rand = { workspace = true } + +tokio = { workspace = true, features = ["sync", "io-util"] } +serde = { workspace = true, features = ["derive"] } +serde_bytes = { workspace = true } +rmp-serde = { workspace = true } +iroh = { workspace = true, optional = true } + +tokio-util = { workspace = true, optional = true } + +# Native server runtime: multi-thread, timers, signals. The wasm build uses the JS event loop/timers. +[target.'cfg(not(target_family = "wasm"))'.dependencies] +tokio = { workspace = true, features = ["rt-multi-thread", "time", "signal"] } + +[target.'cfg(all(target_family = "wasm", target_os = "unknown"))'.dependencies] +wasm-bindgen-futures = { workspace = true } +gloo-timers = { version = "0.3", features = ["futures"] } +futures-util = { workspace = true } + +[dev-dependencies] +burn-flex = { workspace = true, features = ["default"] } +burn-tensor = { workspace = true, features = ["default", "remote-websocket"] } + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-remote/README.md b/crates/burn-remote/README.md new file mode 100644 index 0000000..6a6d4be --- /dev/null +++ b/crates/burn-remote/README.md @@ -0,0 +1,84 @@ +# Burn Remote + +Burn Remote executes tensor operations on compute peers reached through +[Iroh](https://iroh.computer/). Iroh is the primary transport: peers are identified by +cryptographic endpoint IDs, direct connections are preferred, and configured relays are used +when NAT traversal cannot establish a direct path. + +## Client + +Applications own the Iroh endpoint configuration. This keeps identity persistence, relay policy, +address lookup, and fleet discovery outside Burn: + +```rust,ignore +use burn::{Device, Tensor}; +use burn::backend::remote::{Endpoint, RemoteNode, endpoint::presets}; + +let endpoint = Endpoint::builder(presets::N0) + // .secret_key(persistent_secret_key) + // .relay_mode(custom_relay_mode) + // .address_lookup(platform_lookup) + .bind() + .await?; +let node = RemoteNode::from_endpoint(endpoint); + +// Supplied by your platform, invitation, or other discovery mechanism. +let compute_peer = fleet.lookup("gpu-worker-7").await?; +let device = Device::remote_iroh(&node, compute_peer, 0); + +let output = Tensor::<1>::from_floats([1.0, 2.0], &device) * 2.0; +``` + +`RemoteNode` is process-level. Clone it rather than creating one per device: clones share one +Iroh endpoint, peer connection pool, and multiplexed QUIC connections. + +Platforms can issue a `RemoteTicket` containing an `EndpointAddr` and opaque authorization bytes. +Burn passes the credential to the compute peer's `PeerAuthorizer`; signature format, expiry, +tenant policy, and fleet membership remain application concerns. + +## Compute peer + +```rust,ignore +use burn::{Device, server::{self, Channel}}; +use burn::backend::remote::RemoteNode; + +let node = RemoteNode::bind().await?; +println!("compute peer: {}", node.endpoint().addr()); + +server::start_async( + Device::cuda(0), + Channel::Iroh { node }, +).await; +``` + +For an endpoint shared with other Iroh protocols, register Burn's composable handler in the +application router: + +```rust,ignore +use burn_remote::BURN_REMOTE_ALPN; +use iroh::protocol::Router; + +let burn = node + .protocol::(devices) + .with_authorizer(|request| platform.verify(request.peer, request.credential)); + +let router = Router::builder(node.endpoint().clone()) + .accept(BURN_REMOTE_ALPN, burn) + .accept(MY_OTHER_ALPN, other_protocol) + .spawn(); +``` + +## Tensor movement + +Moving a tensor between different Iroh compute peers does not route the payload through the +client. The destination peer opens an authenticated stream directly to the source peer. Each +transfer uses a random, short-lived capability bound to the destination's authenticated endpoint +identity and limited to the number of downloads requested by the operation. + +Multiple devices hosted by the same compute peer retain the in-process fast path. +Tensor movement between an Iroh peer and a legacy WebSocket peer is not supported. + +## WebSocket compatibility + +The `websocket` feature preserves `Device::remote("ws://host:port", index)` and +`Channel::WebSocket`. It is intended for compatibility; new integrations should use Iroh. diff --git a/crates/burn-remote/src/client/base.rs b/crates/burn-remote/src/client/base.rs new file mode 100644 index 0000000..cf59770 --- /dev/null +++ b/crates/burn-remote/src/client/base.rs @@ -0,0 +1,68 @@ +use super::{RemoteDevice, service::RemoteService}; +use burn_backend::{DeviceHandle, backend::Device}; + +/// A thin handle to a `RemoteService` running on its own device-runner thread. +/// +/// Every `RouterClient` method delegates to `handle.submit` / `submit_blocking`, so all +/// connection state, the tokio runtime, the response-demux task, and the op batch buffer +/// live on the service side. +pub struct RemoteClient { + pub(crate) device: RemoteDevice, + pub(crate) handle: DeviceHandle, +} + +impl RemoteClient { + pub fn init(device: RemoteDevice) -> Self { + // `DeviceHandle::new` initializes the service the first time it's called for a given + // device id. `RemoteService::init` is deliberately cheap — it records the endpoint but + // does NOT connect, because cubecl holds a process-global lock across it; the actual + // connect + handshake happens lazily on first use (or via `ensure_connected`). + // Subsequent calls return a handle to the existing service. + let handle = DeviceHandle::::new(device.to_id()); + Self { device, handle } + } + + /// Force the lazily-established connection to be opened now, populating the device's + /// settings/device-count cells. Used by the settings path (`RemoteDevice::defaults` / + /// `enumerate`), which needs the handshake reply before any op has flushed. Runs the + /// connect on the service's runner thread, so it can't sit under cubecl's global lock. + pub(crate) fn ensure_connected(&self) { + self.handle + .submit_blocking(|s| s.ensure_connected()) + .expect("Service call failed"); + } + + /// Establish the session asynchronously, the way the browser requires. + /// + /// The connect + handshake cannot block the single browser thread, so it runs off the device + /// handle: the service hands back the connection parameters, the network round-trip happens + /// with `.await`, and the opened session is installed back into the service. A no-op once the + /// session is up. + #[cfg(target_family = "wasm")] + pub(crate) async fn connect_async(&self) { + use crate::client::service::wasm_connect; + + let Some(plan) = self + .handle + .submit_blocking(|s| s.wasm_connect_plan()) + .expect("Service call failed") + else { + return; + }; + + let connected = wasm_connect(plan).await; + + self.handle + .submit_blocking(move |s| s.wasm_install(connected)) + .expect("Service call failed"); + } +} + +impl Clone for RemoteClient { + fn clone(&self) -> Self { + Self { + device: self.device.clone(), + handle: self.handle.clone(), + } + } +} diff --git a/crates/burn-remote/src/client/channel.rs b/crates/burn-remote/src/client/channel.rs new file mode 100644 index 0000000..d45ac8f --- /dev/null +++ b/crates/burn-remote/src/client/channel.rs @@ -0,0 +1,70 @@ +use burn_backend::Shape; +use burn_ir::TensorIr; +use burn_router::{RouterChannel, RouterTensor, get_client}; + +use super::{ + RemoteClient, + runner::{RemoteBridge, RemoteDevice, RemoteTensorHandle}, +}; + +/// A local channel with direct connection to the backend runner clients. +pub struct RemoteChannel; + +impl RouterChannel for RemoteChannel { + type Device = RemoteDevice; + type Bridge = RemoteBridge; + type Client = RemoteClient; + + fn name(device: &Self::Device) -> String { + format!("remote-{device:?}") + } + + fn init_client(device: &Self::Device) -> Self::Client { + RemoteClient::init(device.clone()) + } + + fn get_tensor_handle(tensor: &TensorIr, client: &Self::Client) -> RemoteTensorHandle { + RemoteTensorHandle { + client: client.clone(), + tensor: tensor.clone(), + } + } + + fn register_tensor( + _client: &Self::Client, + _handle: RemoteTensorHandle, + _shape: Shape, + _dtype: burn_backend::DType, + ) -> RouterTensor { + // This function is normally only used to move a tensor from a device to another. + // + // In other words, to change the client. + panic!("Can't register manually a tensor on a remote channel."); + } + + fn change_client_backend( + tensor: RouterTensor, + target_device: &Self::Device, // target device + ) -> RouterTensor { + // Get tensor handle from current client + let original_client = tensor.client.clone(); + let desc = tensor.into_ir(); + let handle = Self::get_tensor_handle(&desc, &original_client); + + let handle = handle.change_backend(target_device); + + let id = handle.tensor.id; + + let target_client = get_client::(target_device); + let router_tensor: RouterTensor = + RouterTensor::new(id, handle.tensor.shape, handle.tensor.dtype, target_client); + + router_tensor + } +} + +impl Clone for RemoteChannel { + fn clone(&self) -> Self { + RemoteChannel + } +} diff --git a/crates/burn-remote/src/client/custom_op.rs b/crates/burn-remote/src/client/custom_op.rs new file mode 100644 index 0000000..e3d07bc --- /dev/null +++ b/crates/burn-remote/src/client/custom_op.rs @@ -0,0 +1,94 @@ +//! Client-side helper for registering [custom operations](CustomOpIr) on the remote backend. +//! +//! A backend extension ships its op to the server as `OperationIr::Custom`, where a registered +//! [`CustomOpRegistry`](burn_router::CustomOpRegistry) handler executes it. How the op is registered +//! on the client differs by whether the `fusion` feature is enabled — the remote backend is a plain +//! [`BackendRouter`](burn_router::BackendRouter) without it and a +//! [`Fusion`](burn_fusion::Fusion)-wrapped one with it — and so does the tensor primitive it returns +//! (`RouterTensor` vs `FusionTensor`). [`CustomOpClient`] hides that difference: the same code builds +//! and registers a custom op, and gets back `FloatTensor` either way. + +use burn_backend::tensor::FloatTensor; +use burn_ir::{CustomOpIr, OperationIr, TensorId}; + +use crate::client::RemoteChannel; +use crate::{RemoteBackend, RemoteDevice}; + +/// The router channel used by the remote backend. Transport-erased, so it carries custom ops over +/// whichever transport the device was opened on. +type Channel = RemoteChannel; + +/// A client for registering [custom operations](CustomOpIr) on the remote backend, transparently +/// across the `fusion` feature. +/// +/// Drop-in for the lower-level client a backend extension would otherwise reach for: allocate output +/// ids with [`create_empty_handle`](Self::create_empty_handle), build a [`CustomOpIr`], then +/// [`register`](Self::register) it. With `fusion` enabled the op joins the cached op-graph (via the +/// fusion client); without it the op streams through the router client. Either way the op reaches the +/// server, and the returned tensors are the matching `FloatTensor`. +pub struct CustomOpClient { + #[cfg(not(feature = "fusion"))] + inner: ::Client, + #[cfg(feature = "fusion")] + inner: burn_fusion::client::GlobalFusionClient>, + /// Kept so the fusion path can build the op's unfused handler, which needs the device to reach + /// the router client when the op isn't fused into a graph. + #[cfg(feature = "fusion")] + device: RemoteDevice, +} + +impl CustomOpClient { + /// Create a client bound to the given remote device. + pub fn new(device: &RemoteDevice) -> Self { + #[cfg(not(feature = "fusion"))] + let inner = burn_router::get_client::(device); + #[cfg(feature = "fusion")] + let inner = burn_fusion::get_client::>(device); + + Self { + inner, + #[cfg(feature = "fusion")] + device: device.clone(), + } + } + + /// Allocate a fresh, uninitialized tensor id for a custom op output. + /// + /// Use it to build the output [`TensorIr`](burn_ir::TensorIr)s of the [`CustomOpIr`] passed to + /// [`register`](Self::register). The id is allocated by the same client that registers the op, so + /// it is consistent with how that client tracks tensors (fusion ids under `fusion`, router ids + /// otherwise). + pub fn create_empty_handle(&self) -> TensorId { + #[cfg(not(feature = "fusion"))] + { + use burn_router::RouterClient; + self.inner.create_empty_handle() + } + #[cfg(feature = "fusion")] + { + self.inner.create_empty_handle() + } + } + + /// Register a custom op and return its (uninitialized) output tensors. + pub fn register(&self, op: CustomOpIr) -> Vec> { + #[cfg(not(feature = "fusion"))] + { + use burn_router::RouterClient; + self.inner.register(OperationIr::Custom(op)) + } + #[cfg(feature = "fusion")] + { + use burn_fusion::stream::StreamId; + use burn_router::CustomOperation; + // The op runs on the server through the `CustomOpRegistry`. When it's fused into a + // cached op-graph it travels as part of the graph; when it isn't (a one-op segment), + // `CustomOperation` is the unfused fallback that ships it through the router client. + self.inner.register( + StreamId::current(), + OperationIr::Custom(op.clone()), + CustomOperation::::new(op, self.device.clone()), + ) + } + } +} diff --git a/crates/burn-remote/src/client/mod.rs b/crates/burn-remote/src/client/mod.rs new file mode 100644 index 0000000..fd13df2 --- /dev/null +++ b/crates/burn-remote/src/client/mod.rs @@ -0,0 +1,11 @@ +mod base; +mod channel; +mod custom_op; +mod runner; +pub(crate) mod runtime; +pub(crate) mod service; + +pub use base::*; +pub use channel::*; +pub use custom_op::CustomOpClient; +pub use runner::RemoteDevice; diff --git a/crates/burn-remote/src/client/runner.rs b/crates/burn-remote/src/client/runner.rs new file mode 100644 index 0000000..413d720 --- /dev/null +++ b/crates/burn-remote/src/client/runner.rs @@ -0,0 +1,516 @@ +use super::{RemoteChannel, RemoteClient, service}; +use crate::shared::{LocalTransferId, TaskResponseContent, TensorRemote, TransferCapability}; +use crate::{PeerAddr, PeerId}; +use burn_backend::{DeviceId, DeviceOps, ExecutionError, StreamId, TensorData}; +use burn_ir::TensorIr; +use burn_router::{MultiBackendBridge, RouterClient, RouterTensor, get_client}; +use burn_std::DeviceSettings; +use burn_std::{backtrace::BackTrace, future::DynFut}; +#[cfg(feature = "websocket")] +use std::sync::Arc; +use std::sync::Mutex; + +use super::runtime::Executor; +use service::RemoteEndpoint; + +// It is very important to block on any request made via the service, since ordering is +// crucial when registering operations or creating tensors. The `DeviceHandle` queue +// preserves submission order, so `submit` is sufficient for cheap fire-and-forget ops; we +// only `submit_blocking` for paths that need to read the service's response. +impl RouterClient for RemoteClient { + type Device = RemoteDevice; + + fn register_op(&self, op: burn_ir::OperationIr) { + let stream_id = StreamId::current(); + // Device ids in the op's payload are *client* remote device ids; rewrite them to + // server-local device indices the server can resolve to its own backend devices. Applies + // to every op — only ops that actually carry device ids are rewritten. + let op = self.resolve_devices(op); + self.handle.submit(move |s| s.register_op(stream_id, op)); + } + + fn read_tensor_async( + &self, + tensor: burn_ir::TensorIr, + ) -> DynFut> { + // Issue the request synchronously so ordering is preserved relative to subsequent + // submissions; the returned future just awaits the server's response. + let stream_id = StreamId::current(); + let rx = self + .handle + .submit_blocking(move |s| s.read_tensor(stream_id, tensor)) + .expect("Service call failed"); + + Box::pin(async move { + match rx.await { + Ok(TaskResponseContent::ReadTensor(res)) => res, + Ok(_) => panic!("Invalid response type for ReadTensor"), + Err(e) => Err(ExecutionError::Generic { + reason: format!("Failed to read tensor: {e:?}"), + backtrace: BackTrace::capture(), + }), + } + }) + } + + fn register_tensor_data(&self, data: TensorData) -> RouterTensor { + let shape = data.shape.clone(); + let dtype = data.dtype; + let id = service::new_tensor_id(); + + // Fire-and-forget: the outgoing batch flushes itself once buffered data bytes (or the task + // count) cross their threshold — see `OutgoingBatch` — so no explicit flush is needed here. + let stream_id = StreamId::current(); + self.handle + .submit(move |s| s.register_tensor(stream_id, id, data)); + + RouterTensor::new(id, shape, dtype, self.clone()) + } + + fn device(&self) -> Self::Device { + self.device.clone() + } + + fn sync(&self) -> Result<(), ExecutionError> { + let stream_id = StreamId::current(); + self.handle + .submit_blocking(|s| s.sync(stream_id)) + .expect("Service call failed") + } + + fn seed(&self, seed: u64) { + self.handle.submit(move |s| s.seed(seed)); + } + + fn create_empty_handle(&self) -> burn_ir::TensorId { + service::new_tensor_id() + } + + fn register_alias(&self, new_id: burn_ir::TensorId, src_id: burn_ir::TensorId) { + let stream_id = StreamId::current(); + self.handle + .submit(move |s| s.register_alias(stream_id, new_id, src_id)); + } + + fn dtype_usage(&self, dtype: burn_std::DType) -> burn_backend::DTypeUsageSet { + self.handle + .submit_blocking(move |s| s.dtype_usage(dtype)) + .expect("Service call failed") + } + + fn flush(&self) { + self.handle.submit_blocking(|s| s.flush()).unwrap(); + } + + fn register_and_execute_graph( + &self, + graph_id: burn_ir::GraphId, + relative_graph: Vec, + bindings: burn_ir::GraphBindings, + ) { + let stream_id = StreamId::current(); + self.handle.submit(move |s| { + s.register_and_execute_graph(stream_id, graph_id, relative_graph, bindings) + }); + } + + fn execute_graph(&self, graph_id: burn_ir::GraphId, bindings: burn_ir::GraphBindings) { + let stream_id = StreamId::current(); + self.handle + .submit(move |s| s.execute_graph(stream_id, graph_id, bindings)); + } +} + +impl RemoteClient { + /// Rewrite the device ids carried by an op so the server can resolve them. + /// + /// This runs for every op, but only ops that carry device ids (currently the collective ops) + /// are affected On the client, the participating devices are identified by their *remote* + /// device ids (`type_id = 0`, `index_id = ` the local-registry index that encodes + /// `address`+device index). The server can't reverse that registry hash, so we translate each + /// id to the plain server-local device index (kept in `index_id`, `type_id` left 0). The + /// server then maps each index to its own backend device id before executing — see + /// `RemoteServer::resolve_devices`. + /// + /// Only same-server collectives are supported for now: every participating device must live + /// on the same address as the tensor's device. A cross-server group panics with a clear + /// message rather than silently reducing the wrong devices. + fn resolve_devices(&self, mut op: burn_ir::OperationIr) -> burn_ir::OperationIr { + use burn_ir::{DistributedOperationIr, OperationIr}; + + if let OperationIr::Distributed(DistributedOperationIr::AllReduce(desc)) = &mut op { + let local_peer = self.device.peer_id(); + for id in desc.device_ids.iter_mut() { + let (endpoint, device_index) = service::endpoint_for(id.index_id as u32).expect( + "an all_reduce device must be a registered remote device on this process", + ); + assert_eq!( + endpoint.peer_id(), + local_peer, + "cross-peer all_reduce is not supported yet: the tensor is on `{local_peer}` \ + but the collective includes a device on `{}`", + endpoint.peer_id(), + ); + id.type_id = 0; + id.index_id = device_index as u16; + } + log::trace!("All-reduce on {:?} ({local_peer}): {desc:?}", self.device); + } + + op + } +} + +#[derive(Clone, Debug)] +/// A remote compute device identified by its endpoint and device index. +/// +/// Two RemoteDevices with the same endpoint but different indices point at distinct devices on +/// the same peer; each gets its own registry id and service connection, and transfers between +/// them take the same-peer fast path. +pub struct RemoteDevice { + pub(crate) endpoint: RemoteEndpoint, + /// Device index on the remote peer. + pub(crate) device_index: u32, + /// Local registry id for this device. + pub(crate) id: u32, +} + +impl RemoteDevice { + /// Create a legacy WebSocket device from a URL. + #[cfg(feature = "websocket")] + pub fn websocket(address: &str, device_index: usize) -> Self { + let endpoint = RemoteEndpoint::WebSocket { + address: burn_communication::Address::from(address), + authorization: Arc::from([]), + }; + let device_index = device_index as u32; + let id = service::register_endpoint(endpoint.clone(), Executor::capture(), device_index); + Self { + endpoint, + device_index, + id, + } + } + + /// Create an Iroh remote device dialing `peer` from `endpoint`. + /// + /// The application owns the Iroh endpoint; Burn dials the compute peer from it. + #[cfg(feature = "iroh")] + pub fn iroh(endpoint: &iroh::Endpoint, peer: iroh::EndpointAddr, device_index: usize) -> Self { + Self::iroh_authorized(endpoint, peer, device_index, Vec::new()) + } + + /// Like [`iroh`](Self::iroh), but carries an authorization credential the server's PeerAuthorizer will check. + #[cfg(feature = "iroh")] + pub fn iroh_authorized( + endpoint: &iroh::Endpoint, + peer: iroh::EndpointAddr, + device_index: usize, + authorization: Vec, + ) -> Self { + let node = crate::transport::iroh::node::RemoteNode::from_endpoint(endpoint.clone()); + let endpoint = RemoteEndpoint::Iroh { + node, + peer, + authorization: authorization.into(), + }; + let device_index = device_index as u32; + let id = service::register_endpoint(endpoint.clone(), Executor::capture(), device_index); + Self { + endpoint, + device_index, + id, + } + } + + /// Forces the client connection to be established immediately using the default protocol. + /// This is a no-op if the connection is already up for this device. + pub fn connect(&self) { + // `get_client` initializes the (lazy) service if needed; `ensure_connected` then opens + // the sockets and runs the handshake on the runner thread, so the settings/device-count + // cells are populated by the time we return. + get_client::(self).ensure_connected(); + } + + /// Establish the session asynchronously. Browser entry point: wasm cannot block to connect, + /// so call and await this once before using the device. No-op if already connected. + #[cfg(target_family = "wasm")] + pub async fn connect_async(&self) { + get_client::(self).connect_async().await; + } + + /// Initialize the client for this device using a custom protocol channel. + /// + /// Only creates the lazy service; the socket and handshake open on first use. Call + /// [`connect`](Self::connect) when the connection and device settings are needed immediately. + pub fn connect_with_channel>(&self) { + // `get_client` forces service initialization if the client doesn't exist yet; + // `RemoteService::init` records the endpoint but defers the connect to first use. + get_client::(self); + } + + /// The stable identity of the compute peer. + pub fn peer_id(&self) -> PeerId { + self.endpoint.peer_id() + } + + /// The peer identity plus its current dialing hints. + pub fn peer_addr(&self) -> PeerAddr { + self.endpoint.peer_addr() + } + + /// The peer address as a string. Prefer `peer_addr` for typed access. + pub fn address(&self) -> String { + self.peer_addr().to_string() + } + + /// The index of this device on its server. + pub fn device_index(&self) -> usize { + self.device_index as usize + } + + /// List every device hosted by the WebSocket server at `address`. + /// + /// Connects to index 0 to read the device count from the init handshake, then returns one + /// RemoteDevice per index. Remaining indices connect lazily on first use, matching the + /// behavior of [`Device::enumerate`](burn_backend::tensor::Device) for local backends. + #[cfg(feature = "websocket")] + pub fn enumerate_websocket(address: &str) -> Vec { + // Device 0 always exists (a server must host at least one device); connecting to it + // populates the device-count cell for its registry id. + let device = Self::websocket(address, 0); + device.connect(); + + let count = service::device_count_for(device.id) + .expect("Device count populated by the init handshake during connect"); + + (0..count as usize) + .map(|index| Self::websocket(address, index)) + .collect() + } + + /// List every device hosted by an Iroh peer. + #[cfg(feature = "iroh")] + pub fn enumerate_iroh(endpoint: &iroh::Endpoint, peer: iroh::EndpointAddr) -> Vec { + let device = Self::iroh(endpoint, peer.clone(), 0); + device.connect(); + let count = service::device_count_for(device.id) + .expect("Device count populated by the init handshake during connect"); + (0..count as usize) + .map(|index| Self::iroh(endpoint, peer.clone(), index)) + .collect() + } +} + +impl PartialEq for RemoteDevice { + fn eq(&self, other: &Self) -> bool { + self.id == other.id + } +} + +impl Eq for RemoteDevice {} + +impl Default for RemoteDevice { + fn default() -> Self { + #[cfg(feature = "websocket")] + { + let address = match std::env::var("BURN_REMOTE_ADDRESS") { + Ok(address) => address, + Err(_) => String::from("ws://127.0.0.1:3000"), + }; + + Self::websocket(&address, 0) + } + #[cfg(not(feature = "websocket"))] + panic!( + "RemoteDevice::default requires the `websocket` compatibility feature; construct an Iroh device through RemoteNode::device" + ) + } +} + +impl burn_std::device::Device for RemoteDevice { + fn from_id(device_id: DeviceId) -> Self { + if device_id.type_id != 0 { + panic!("Invalid device id: {device_id} (expected type 0)"); + } + let (endpoint, device_index) = service::endpoint_for(device_id.index_id as u32) + .unwrap_or_else(|| panic!("Invalid device id: {device_id}")); + Self { + endpoint, + device_index, + id: device_id.index_id as u32, + } + } + + fn to_id(&self) -> DeviceId { + DeviceId { + type_id: 0, + index_id: self.id as u16, + } + } +} + +impl DeviceOps for RemoteDevice { + fn defaults(&self) -> DeviceSettings { + // Lazy-connect on first access. Callers like `Device::configure` or + // `Device::default()`-driven dispatch can hit `defaults` before the user has + // triggered any op, so we need to establish the session here. `connect` is + // idempotent — a no-op once the client has been initialized for this device. + if !service::has_settings(self.id) { + self.connect(); + } + service::settings_for(self.id) + } +} + +pub struct RemoteBridge; + +pub struct RemoteTensorHandle { + pub(crate) client: RemoteClient, + pub(crate) tensor: TensorIr, +} + +static TRANSFER_COUNTER: Mutex> = Mutex::new(None); + +/// Allocate the next process-unique [`LocalTransferId`] for a same-peer transfer. +/// +/// The id keys the server's transfer rendezvous (`local_comm` / `external_comm`), so two +/// transfers that are ever in flight at the same time MUST get distinct ids; otherwise a `take` +/// can pick up the wrong (or an overwritten) exposed primitive and its peer hangs forever. The +/// counter is incremented **in place** in the static; `LocalTransferId` is `Copy`, so the +/// earlier `transfer_counter.unwrap()` copied the value out and incremented a throwaway local, +/// leaving every transfer after the first sharing id 1 (harmless sequentially, a deadlock under +/// concurrency). +fn get_next_transfer_id() -> LocalTransferId { + let mut transfer_counter = TRANSFER_COUNTER.lock().unwrap(); + match transfer_counter.as_mut() { + Some(id) => { + id.next(); + *id + } + None => { + let id = LocalTransferId::from(0); + *transfer_counter = Some(id); + id + } + } +} + +impl RemoteTensorHandle { + /// Move the tensor to `target_device`, picking the cheapest path. + /// + /// When the source and target live on the **same** server (same address, different device + /// index), the data never leaves the process — see [`change_backend_local`]. Otherwise we + /// fall back to the cross-server path that streams the data server-to-server without the + /// client ever seeing it. + pub(crate) fn change_backend(self, target_device: &RemoteDevice) -> Self { + if self.client.device.peer_id() == target_device.peer_id() { + self.change_backend_local(target_device) + } else { + assert_eq!( + self.client.device.peer_addr().is_iroh(), + target_device.peer_addr().is_iroh(), + "Moving a tensor between Iroh and legacy WebSocket compute peers is not supported" + ); + self.change_backend_remote(target_device) + } + } + + /// Same-host transfer: hand the device-resident primitive from the source session to the + /// target session on the same server, which moves it with the inner backend's `to_device`. + /// No host round-trip. + fn change_backend_local(mut self, target_device: &RemoteDevice) -> Self { + // The stream of the calling user thread: the source reads the tensor back on the + // stream that produced it, and the target registers the result on the stream that + // will consume it — same contract as the regular op/tensor registration paths. + let stream_id = StreamId::current(); + let transfer_id = get_next_transfer_id(); + let tensor = self.tensor.clone(); + self.client.handle.submit(move |s| { + s.expose_tensor_local(stream_id, tensor, transfer_id); + }); + // Force the expose (and the ops producing the tensor) onto the wire now — same reason + // as the cross-server path below. + self.client.handle.flush_queue(); + + let target_client = get_client::(target_device); + let new_id = service::new_tensor_id(); + target_client.handle.submit(move |s| { + s.register_tensor_local(stream_id, transfer_id, new_id); + }); + target_client.handle.flush_queue(); + + self.tensor.id = new_id; + self.client = target_client; + + self + } + + /// Changes the backend of the tensor via a dWebSocket. + /// We ask the original server to expose the tensor, then ask the target server to fetch + /// the tensor. The target server will open a new network connection to the original server + /// to download the data. + /// This way the client never sees the tensor's data, and we avoid a bottleneck. + fn change_backend_remote(mut self, target_device: &RemoteDevice) -> Self { + // See `change_backend_local`: carry the calling thread's stream so the source readback + // and the target registration land on the client streams, not arbitrary server threads. + let stream_id = StreamId::current(); + let capability = TransferCapability::random(); + let tensor = self.tensor.clone(); + let target = target_device.peer_id(); + self.client.handle.submit(move |s| { + s.expose_tensor_remote(stream_id, tensor, 1, capability, target); + }); + // `submit` only enqueues the closure on the device-runner queue; the runner + // wouldn't drain it until 32 ops accumulated. `flush_queue` forces the runner + // thread to run the closure now, and `expose_tensor_remote` itself flushes the + // service batch onto the wire — so the source server receives the expose before + // the target server starts trying to download. + self.client.handle.flush_queue(); + + let target_client = get_client::(target_device); + + let peer = self.client.device.peer_addr(); + let new_id = service::new_tensor_id(); + target_client.handle.submit(move |s| { + s.register_tensor_remote(stream_id, TensorRemote { capability, peer }, new_id); + }); + // Same as the source side: drain the closure queue so it runs now and + // `register_tensor_remote` flushes the registration onto the target's wire. + target_client.handle.flush_queue(); + + self.tensor.id = new_id; + self.client = target_client; + + self + } +} + +impl MultiBackendBridge for RemoteBridge { + type TensorHandle = RemoteTensorHandle; + type Device = RemoteDevice; + + fn change_backend_float( + tensor: Self::TensorHandle, + _shape: burn_backend::Shape, + target_device: &Self::Device, + ) -> Self::TensorHandle { + tensor.change_backend(target_device) + } + + fn change_backend_int( + tensor: Self::TensorHandle, + _shape: burn_backend::Shape, + target_device: &Self::Device, + ) -> Self::TensorHandle { + tensor.change_backend(target_device) + } + + fn change_backend_bool( + tensor: Self::TensorHandle, + _shape: burn_backend::Shape, + target_device: &Self::Device, + ) -> Self::TensorHandle { + tensor.change_backend(target_device) + } +} diff --git a/crates/burn-remote/src/client/runtime.rs b/crates/burn-remote/src/client/runtime.rs new file mode 100644 index 0000000..d1ea429 --- /dev/null +++ b/crates/burn-remote/src/client/runtime.rs @@ -0,0 +1,104 @@ +//! The client session runtime. + +/// Process-global Tokio runtime for sessions opened outside an ambient runtime. +/// +/// Fallback for synchronous callers (scripts, REPLs, notebooks) and the legacy WebSocket path. A +/// native Iroh node also binds on this runtime so its endpoint and session tasks share one executor. +/// When a device is built inside an existing runtime, that one is used instead and this fallback is +/// never created. +#[cfg(not(target_family = "wasm"))] +pub(crate) fn blocking_runtime() -> &'static tokio::runtime::Runtime { + use std::sync::OnceLock; + static RUNTIME: OnceLock = OnceLock::new(); + RUNTIME.get_or_init(|| { + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("Can build the Burn Remote blocking runtime") + }) +} + +/// Executor for a remote session's writer and response-demux tasks. +/// +/// Captured once at device construction and carried by the device registry, so the session reuses +/// whatever runtime owns its transport. On native this wraps a Tokio runtime handle: the ambient +/// runtime if one is active, otherwise the shared [`blocking_runtime`]. In the browser Iroh runs on +/// the JS event loop; tasks are spawned with `spawn_local` and blocking calls are unavailable. +#[derive(Clone, Debug)] +pub(crate) enum Executor { + #[cfg(not(target_family = "wasm"))] + Tokio(tokio::runtime::Handle), + #[cfg(target_family = "wasm")] + WasmLocal, +} + +/// Handle to a spawned session task. Joinable on native; a no-op in the browser where tasks +/// run on the event loop and cannot be awaited. +pub(crate) struct SpawnHandle { + #[cfg(not(target_family = "wasm"))] + inner: tokio::task::JoinHandle<()>, +} + +impl Executor { + /// Capture the executor for a new session at device-construction time. + /// + /// Native: the ambient Tokio runtime if one is active, otherwise the shared [`blocking_runtime`]. + /// Browser: the JS event loop. The result is stored in the device registry. + #[cfg(not(target_family = "wasm"))] + pub(crate) fn capture() -> Self { + match tokio::runtime::Handle::try_current() { + Ok(handle) => Self::Tokio(handle), + Err(_) => Self::Tokio(blocking_runtime().handle().clone()), + } + } + + #[cfg(target_family = "wasm")] + pub(crate) fn capture() -> Self { + Self::WasmLocal + } + + pub(crate) fn block_on(&self, future: F) -> F::Output { + match self { + #[cfg(not(target_family = "wasm"))] + Self::Tokio(handle) => handle.block_on(future), + #[cfg(target_family = "wasm")] + Self::WasmLocal => { + core::mem::drop(future); + panic!( + "Blocking remote calls are not supported on wasm. Establish the session with \ + `RemoteDevice::connect_async(...).await` and read tensors with \ + `into_data_async().await`." + ) + } + } + } + + #[cfg(not(target_family = "wasm"))] + pub(crate) fn spawn(&self, future: F) -> SpawnHandle + where + F: core::future::Future + Send + 'static, + { + match self { + Self::Tokio(handle) => SpawnHandle { + inner: handle.spawn(future), + }, + } + } + + /// Spawn a session task on the browser event loop. The Iroh streams these tasks own are not + /// `Send`, which is why the wasm path uses `spawn_local` rather than the native `spawn`. + #[cfg(target_family = "wasm")] + pub(crate) fn spawn(&self, future: F) -> SpawnHandle + where + F: core::future::Future + 'static, + { + wasm_bindgen_futures::spawn_local(future); + SpawnHandle {} + } + + /// Wait for a spawned task to finish. No-op in the browser. + #[cfg(not(target_family = "wasm"))] + pub(crate) fn join(&self, handle: SpawnHandle) { + let _ = self.block_on(handle.inner); + } +} diff --git a/crates/burn-remote/src/client/service.rs b/crates/burn-remote/src/client/service.rs new file mode 100644 index 0000000..0238105 --- /dev/null +++ b/crates/burn-remote/src/client/service.rs @@ -0,0 +1,687 @@ +use crate::metrics::{MetricSide, TelemetryLogger, logger_task}; +use crate::shared::{ + LocalTransferId, PROTOCOL_VERSION, RemoteMessage, RequestId, SessionId, SessionInfo, + SessionInit, Task, TaskResponse, TaskResponseContent, TensorRemote, TransferCapability, +}; +use crate::telemetry::{CHANNEL_CAPACITY, TelemetryEvent, TelemetryProbe, serialized_len}; +use burn_backend::{ + DTypeUsageSet, ExecutionError, TensorData, + backend::{DeviceId, DeviceService, ServerUtilitiesHandle}, +}; +use burn_ir::{OperationIr, TensorId, TensorIr}; +use burn_std::{DType, DeviceSettings, id::StreamId}; +// Only the native `sync` path captures a backtrace; the wasm path returns without blocking. +#[cfg(not(target_family = "wasm"))] +use burn_std::backtrace::BackTrace; +use std::sync::{Arc, OnceLock}; +use tokio::sync::oneshot; + +mod batch; +mod conn; +mod pending; +mod registry; +mod writer; + +use batch::OutgoingBatch; +use conn::{ResponseChannel, open_channels}; +use pending::{PendingResponses, Responder}; +use writer::SubmitWriter; + +use super::runtime::Executor; +pub(crate) use conn::{RemoteEndpoint, SubmitChannel}; +use registry::{device_count_cell, executor_for, settings_cell}; +pub(crate) use registry::{device_count_for, has_settings, new_tensor_id, settings_for}; +pub(crate) use registry::{endpoint_for, register_endpoint}; + +/// All the state owned by the device-runner thread for a single remote device. +/// +/// `RemoteService` lives behind a [`DeviceHandle`](burn_backend::DeviceHandle); every call +/// from the `RouterClient` shim hops onto the runner thread via the device handle's +/// `submit` / `submit_blocking`, so the service has exclusive access to the connection, +/// the callback map, and the outgoing task buffer without any locking on its own state. +/// +/// The service mirrors that submit-style API internally: fire-and-forget calls go +/// through [`submit`](Self::submit) (push onto the [`batch`](Self::batch), flush once it +/// reaches the configured flush threshold), and response-producing calls go through +/// [`submit_blocking`](Self::submit_blocking) (push, then flush right away so the request +/// is enqueued before we await the oneshot). +/// +/// The moving parts are split into focused types: [`OutgoingBatch`] buffers tasks and +/// decides when to flush, [`SubmitWriter`] owns the socket and serializes frames off the +/// runner thread, and [`PendingResponses`] correlates response-producing requests with the +/// caller awaiting each reply. +/// +/// All tokio work (connecting, the writer task, awaiting responses, the response-demux +/// task) happens on the [`Executor`] captured from the device's endpoint. The caller never +/// sees a runtime handle. +pub struct RemoteService { + executor: Executor, + /// Where to connect on first use. The connection is established lazily (see + /// [`ensure_connected`](Self::ensure_connected)) rather than in [`init`](Self::init), + /// because cubecl holds a process-global device-registry lock across `init` — opening the + /// sockets there would serialize every remote device's setup behind that lock. + endpoint: RemoteEndpoint, + device_index: u32, + /// Owns the request socket and serializes outgoing frames off the runner thread. + /// `None` until the first task (or a settings read) triggers the lazy connect. + writer: Option, + /// Buffers outgoing tasks and signals when a batch is ready for the wire. + batch: OutgoingBatch, + /// Request-id allocation + the callbacks awaiting response-producing tasks. + pending: PendingResponses, + /// Emits this device's telemetry (the ops and graphs it sends). + probe: TelemetryProbe, + /// Shared cell populated from the init handshake (read by `RemoteDevice::defaults`). + settings: Arc>, + /// Shared cell populated from the init handshake (read by `RemoteDevice::enumerate`). + device_count: Arc>, + session_id: SessionId, + closed: bool, +} + +impl DeviceService for RemoteService { + fn init(device_id: DeviceId) -> Self { + let (id, endpoint, device_index) = Self::resolve_endpoint(device_id); + // The executor was captured at device-construction time (in the runtime that owns the + // transport) and stored in the registry alongside the endpoint; the service just reuses it. + let executor = executor_for(id).expect("device registered with a captured executor"); + let session_id = SessionId::new(); + + let probe = if TelemetryLogger::enabled() { + TelemetryProbe::new(CHANNEL_CAPACITY) + } else { + TelemetryProbe::disabled() + }; + if let Some(task) = logger_task(&probe, MetricSide::Client) { + executor.spawn(task); + } + + // Lazy connect: `init` must return promptly. cubecl holds a process-global + // device-registry lock across this call (to make device-handle creation atomic), so + // doing the blocking network connect + handshake here would serialize every remote + // device's setup behind that lock — N devices would connect strictly one at a time. + // Instead we record the endpoint and open the sockets on the first real use, off the + // lock and on the device-runner thread (see `ensure_connected`). + Self { + executor, + endpoint, + device_index, + writer: None, + batch: { + let cfg = burn_std::config::config(); + let remote = cfg.remote(); + OutgoingBatch::new(remote.flush_threshold, remote.flush_bytes_threshold) + }, + pending: PendingResponses::new(), + probe, + settings: settings_cell(id), + device_count: device_count_cell(id), + session_id, + closed: false, + } + } + + fn utilities(&self) -> ServerUtilitiesHandle { + // The remote backend surfaces device settings through the endpoint registry + // (`RemoteDevice::defaults` → `settings_for`), not through this handle, so the device + // layer never reads what we return here. Handing back an empty handle keeps `init` — + // and the cubecl lock it runs under — free of the blocking network handshake. + Arc::new(()) + } +} + +/// Construction helpers for [`RemoteService::init`], one per step of bringing a connection +/// up. Kept separate from the public submit-style API below. +impl RemoteService { + /// Resolve a device id to its registry index, parsed network [`Address`], and the device + /// index to select on the server. + fn resolve_endpoint(device_id: DeviceId) -> (u32, RemoteEndpoint, u32) { + let id = device_id.index_id as u32; + let (endpoint, device_index) = endpoint_for(id) + .unwrap_or_else(|| panic!("No endpoint registered for device id {device_id}")); + (id, endpoint, device_index) + } + + /// Native synchronous wrapper over [`open_channels`](conn::open_channels): blocks the runner + /// thread until the streams are open. + #[cfg(not(target_family = "wasm"))] + fn connect_streams( + executor: &Executor, + endpoint: &RemoteEndpoint, + ) -> (SubmitChannel, ResponseChannel) { + executor + .block_on(open_channels(endpoint)) + .unwrap_or_else(|err: String| panic!("{err}")) + } + + /// Send the session-init handshake on both streams and wait for the device settings the + /// server replies with on the response stream. Both streams carry the same `Vec` + /// wire format; the handshake is just a single-element batch. + async fn handshake_async( + request: &mut SubmitChannel, + response: &mut ResponseChannel, + endpoint: &RemoteEndpoint, + session_id: SessionId, + device_index: u32, + ) -> (DeviceSettings, u32) { + let init_bytes: bytes::Bytes = rmp_serde::to_vec(&vec![RemoteMessage::Init( + SessionInit::new(session_id, device_index, endpoint.authorization().to_vec()), + )]) + .expect("Can serialize RemoteMessage::Init") + .into(); + + let result: Result<(DeviceSettings, u32), String> = async { + request.send(init_bytes).await?; + + let msg = response + .recv() + .await? + .expect("Server disconnected during initialization"); + let reply: TaskResponse = + rmp_serde::from_slice(&msg).expect("Can deserialize init handshake payload"); + + match reply.content { + TaskResponseContent::Init(SessionInfo { + version, + settings, + device_count, + .. + }) => { + if version != PROTOCOL_VERSION { + panic!( + "Server uses Burn Remote protocol version {version}, expected {PROTOCOL_VERSION}" + ); + } + Ok((settings, device_count)) + } + other => panic!("Expected Init response, got {other:?}"), + } + } + .await; + + result.unwrap_or_else(|err| { + panic!( + "Failed to initialize remote session at {}: {err}", + endpoint.peer_addr() + ) + }) + } + + /// Native synchronous wrapper over [`handshake_async`](Self::handshake_async). + #[cfg(not(target_family = "wasm"))] + fn handshake( + executor: &Executor, + request: &mut SubmitChannel, + response: &mut ResponseChannel, + endpoint: &RemoteEndpoint, + session_id: SessionId, + device_index: u32, + ) -> (DeviceSettings, u32) { + executor.block_on(Self::handshake_async( + request, + response, + endpoint, + session_id, + device_index, + )) + } + + /// Spawn the response-demux task: route each [`TaskResponse`] to its pending callback by + /// [`RequestId`] via the [`Responder`]. Lives on the service runtime; exits when the + /// response stream closes. + fn spawn_response_demux( + executor: &Executor, + mut response: ResponseChannel, + responder: Responder, + ) { + // Detached: the task owns the response stream and runs until it closes. + let _demux = executor.spawn(async move { + loop { + match response.recv().await { + Ok(Some(msg)) => { + let reply: TaskResponse = match rmp_serde::from_slice(&msg) { + Ok(r) => r, + Err(err) => { + log::error!("Failed to deserialize remote response: {err:?}"); + continue; + } + }; + if !responder.complete(reply.id, reply.content) { + log::warn!("No pending callback for response id {:?}", reply.id); + } + } + Ok(None) => { + log::warn!("Remote response stream closed"); + break; + } + Err(err) => { + log::warn!("Remote response stream error: {err:?}"); + break; + } + } + } + + // The response stream is gone (clean close or error): the server will never answer + // any in-flight or future request on this connection. Fail every waiting caller and + // gate new ones so they error out instead of blocking forever on a dead server. + responder.disconnect(); + }); + } +} + +/// Session state captured from a [`RemoteService`] to drive an asynchronous (wasm) connect. +#[cfg(target_family = "wasm")] +pub(crate) struct WasmConnectPlan { + endpoint: RemoteEndpoint, + session_id: SessionId, + device_index: u32, + responder: Responder, +} + +/// A session opened by [`wasm_connect`], ready to be installed back into the service. +#[cfg(target_family = "wasm")] +pub(crate) struct WasmConnected { + writer: SubmitWriter, + settings: DeviceSettings, + device_count: u32, +} + +/// Open and hand-shake a session on the browser event loop. +/// +/// This is the async counterpart of [`RemoteService::ensure_connected`]: it runs the parts that +/// would block (connecting the Iroh streams, the init round-trip) with `.await`, then spawns the +/// response-demux and writer tasks with `spawn_local`. The returned [`WasmConnected`] is `Send`, +/// so the caller can install it back into the service through the device handle. +#[cfg(target_family = "wasm")] +pub(crate) async fn wasm_connect(plan: WasmConnectPlan) -> WasmConnected { + let executor = Executor::WasmLocal; + + let (mut request, mut response) = open_channels(&plan.endpoint) + .await + .unwrap_or_else(|err| panic!("{err}")); + let (settings, device_count) = RemoteService::handshake_async( + &mut request, + &mut response, + &plan.endpoint, + plan.session_id, + plan.device_index, + ) + .await; + + RemoteService::spawn_response_demux(&executor, response, plan.responder); + let writer = SubmitWriter::spawn(&executor, request); + + WasmConnected { + writer, + settings, + device_count, + } +} + +impl RemoteService { + /// Buffer a fire-and-forget op. The buffer is flushed automatically once it reaches the + /// configured flush threshold. + pub fn register_op(&mut self, stream_id: StreamId, op: OperationIr) { + // An op streamed individually (not part of a cached graph) is an unfused op. + self.probe + .emit(|| TelemetryEvent::unfused_op(self.session_id, stream_id, &op)); + self.submit_task(Task::RegisterOperation(stream_id, op)); + } + + /// Buffer a fire-and-forget "register a reusable op-graph and run its first invocation" task. + pub fn register_and_execute_graph( + &mut self, + stream_id: StreamId, + graph_id: burn_ir::GraphId, + relative_graph: Vec, + bindings: burn_ir::GraphBindings, + ) { + // The first invocation both registers (one-time graph cost) and executes (a replay). + self.probe.emit(|| { + TelemetryEvent::graph_registered( + self.session_id, + graph_id, + &relative_graph, + serialized_len(&relative_graph), + ) + }); + self.probe.emit(|| { + TelemetryEvent::graph_executed( + self.session_id, + graph_id, + stream_id, + serialized_len(&bindings), + ) + }); + self.submit_task(Task::RegisterAndExecuteGraph { + stream_id, + graph_id, + relative_graph, + bindings, + }); + } + + /// Buffer a fire-and-forget "execute a registered graph" task. + pub fn execute_graph( + &mut self, + stream_id: StreamId, + graph_id: burn_ir::GraphId, + bindings: burn_ir::GraphBindings, + ) { + self.probe.emit(|| { + TelemetryEvent::graph_executed( + self.session_id, + graph_id, + stream_id, + serialized_len(&bindings), + ) + }); + self.submit_task(Task::ExecuteGraph { + stream_id, + graph_id, + bindings, + }); + } + + pub fn register_tensor(&mut self, stream_id: StreamId, id: TensorId, data: TensorData) { + self.submit_task(Task::RegisterTensor(stream_id, id, data)); + self.flush(); + } + + /// Buffer a fire-and-forget "alias `src_id` under `new_id`" task. FIFO submission keeps it + /// after whatever materialized `src_id`, so the server has the source handle when it runs. + pub fn register_alias(&mut self, stream_id: StreamId, new_id: TensorId, src_id: TensorId) { + self.submit_task(Task::RegisterAlias { + stream_id, + new_id, + src_id, + }); + } + + /// Register a tensor produced by a cross-server transfer, flushing immediately. + /// + /// This is a coordination point with no following barrier on this client — the caller + /// (`change_backend`) switches to the target client right after submitting it — so we + /// flush right away instead of relying on a later op to piggy-back the flush. Otherwise + /// the task would sit buffered below the flush threshold and the target server would + /// never start the download. + pub fn register_tensor_remote( + &mut self, + stream_id: StreamId, + tensor: TensorRemote, + new_id: TensorId, + ) { + self.submit_task(Task::RegisterTensorRemote(stream_id, tensor, new_id)); + self.flush(); + } + + /// Expose a tensor for a cross-server transfer, flushing immediately. + /// + /// Flushed for the same reason as [`register_tensor_remote`](Self::register_tensor_remote): + /// the source server must actually receive the expose so the target server's download can + /// complete, and the caller has no later barrier on this client to carry the flush. + pub fn expose_tensor_remote( + &mut self, + stream_id: StreamId, + tensor: TensorIr, + count: u32, + capability: TransferCapability, + target: crate::PeerId, + ) { + self.submit_task(Task::ExposeTensorRemote { + stream_id, + tensor, + count, + capability, + target, + }); + self.flush(); + } + + /// Expose a tensor for a same-host transfer, flushing immediately. + /// + /// Source side of the local path; flushed for the same reason as + /// [`expose_tensor_remote`](Self::expose_tensor_remote). + pub fn expose_tensor_local( + &mut self, + stream_id: StreamId, + tensor: TensorIr, + transfer_id: LocalTransferId, + ) { + self.submit_task(Task::ExposeTensorLocal { + stream_id, + tensor, + transfer_id, + }); + self.flush(); + } + + /// Register a tensor produced by a same-host transfer, flushing immediately. + /// + /// Target side of the local path; flushed for the same reason as + /// [`register_tensor_remote`](Self::register_tensor_remote). + pub fn register_tensor_local( + &mut self, + stream_id: StreamId, + transfer_id: LocalTransferId, + new_id: TensorId, + ) { + self.submit_task(Task::RegisterTensorLocal { + stream_id, + transfer_id, + new_id, + }); + self.flush(); + } + + pub fn seed(&mut self, seed: u64) { + self.submit_task(Task::Seed(seed)); + } + + /// Initiate a tensor read. The returned receiver resolves when the server response + /// arrives. + /// + /// The request id rides on the task itself; the server echoes it back so the + /// response-demux task can hand the response to the right pending callback. + pub fn read_tensor( + &mut self, + stream_id: StreamId, + tensor: TensorIr, + ) -> oneshot::Receiver { + self.submit_request(|id| Task::ReadTensor(id, stream_id, tensor)) + } + + pub fn sync(&mut self, stream_id: StreamId) -> Result<(), ExecutionError> { + #[cfg(not(target_family = "wasm"))] + { + let rx = self.submit_request(|id| Task::SyncBackend(id, stream_id)); + match self.executor.block_on(rx) { + Ok(TaskResponseContent::SyncBackend(res)) => res, + Ok(other) => panic!("Invalid response for SyncBackend: {other:?}"), + Err(_) => Err(ExecutionError::Generic { + reason: "Remote response channel closed before sync completed".into(), + backtrace: BackTrace::capture(), + }), + } + } + #[cfg(target_family = "wasm")] + { + // The browser thread cannot block to await completion. Flushing forces every pending + // operation onto the wire; the writer's FIFO ordering means anything submitted after + // this `sync` still executes after it on the peer, and results are observed through + // async reads. There is no host-side barrier to wait on, so we return immediately. + let _ = stream_id; + self.flush(); + Ok(()) + } + } + + pub fn dtype_usage(&mut self, dtype: DType) -> DTypeUsageSet { + let rx = self.submit_request(|id| Task::DTypeUsage(id, dtype)); + match self.executor.block_on(rx) { + Ok(TaskResponseContent::DTypeUsage(set)) => set, + Ok(other) => panic!("Invalid response for DTypeUsage: {other:?}"), + Err(_) => panic!("Remote response channel closed before dtype_usage completed"), + } + } + + /// Buffer a fire-and-forget compute task. Thin wrapper over [`submit`](Self::submit) + /// that wraps the task in [`RemoteMessage::Task`]. + fn submit_task(&mut self, task: Task) { + self.submit(RemoteMessage::Task(task)); + } + + /// Issue a response-producing compute task: allocate its [`RequestId`], register the + /// pending callback, and flush immediately so it's enqueued before the caller awaits + /// the returned receiver. `make_task` builds the task from the freshly allocated id. + fn submit_request( + &mut self, + make_task: impl FnOnce(RequestId) -> Task, + ) -> oneshot::Receiver { + let request_id = self.pending.next_id(); + let rx = self.pending.register(request_id); + self.submit_blocking(RemoteMessage::Task(make_task(request_id))); + rx + } + + /// Append a task to the outgoing buffer; flush only once it hits the threshold. + /// + /// Use this for fire-and-forget tasks. The runner thread is single-threaded, so + /// pushing into the buffer is just a `Vec::push` — no locking, no tokio hop, no + /// network send. + fn submit(&mut self, task: RemoteMessage) { + if self.batch.push(task) { + self.flush(); + } + } + + /// Append a task and flush the buffer immediately. + /// + /// Use this when the caller is about to await a response: flushing enqueues the batch + /// to the writer task, and the FIFO writer guarantees it reaches the wire before any + /// later frame — so the response can't be missed. We no longer block on the actual + /// send; correctness comes from FIFO enqueue order plus the oneshot await. + fn submit_blocking(&mut self, task: RemoteMessage) { + self.batch.push(task); + self.flush(); + } + + /// Open the session socket and run the init handshake — exactly once. + /// + /// Called lazily from the device-runner thread: from [`flush`](Self::flush) on the first + /// task, or via [`RemoteClient::ensure_connected`](super::RemoteClient) on the settings + /// path. Never from [`init`](Self::init), so the blocking connect + handshake runs off + /// cubecl's global device-registry lock and devices connect in parallel. Always runs on + /// the single runner thread, so the idempotent check needs no locking. + pub fn ensure_connected(&mut self) { + if self.writer.is_some() { + return; + } + + #[cfg(not(target_family = "wasm"))] + { + log::debug!( + "Connecting to {} (device {}) ...", + self.endpoint.peer_addr(), + self.device_index + ); + let (mut request, mut response) = Self::connect_streams(&self.executor, &self.endpoint); + let (settings, device_count) = Self::handshake( + &self.executor, + &mut request, + &mut response, + &self.endpoint, + self.session_id, + self.device_index, + ); + + // Publish to the shared cells so `RemoteDevice::defaults`/`enumerate` can read them. + let _ = self.settings.set(settings); + let _ = self.device_count.set(device_count); + + Self::spawn_response_demux(&self.executor, response, self.pending.responder()); + self.writer = Some(SubmitWriter::spawn(&self.executor, request)); + } + + #[cfg(target_family = "wasm")] + panic!( + "Remote session to {} is not connected. On wasm, establish it with \ + `RemoteDevice::connect_async(...).await` before running tensor operations.", + self.endpoint.peer_addr() + ); + } + + /// Capture everything needed to open the session asynchronously, or `None` if it is already + /// connected. The browser cannot block the runner thread, so the connect + handshake runs + /// off the device handle (see [`wasm_connect`]) and the result is handed back through + /// [`wasm_install`](Self::wasm_install). Every captured value is `Send`, so it can cross the + /// device handle even though the Iroh streams it later owns are not. + #[cfg(target_family = "wasm")] + pub(crate) fn wasm_connect_plan(&mut self) -> Option { + if self.writer.is_some() { + return None; + } + Some(WasmConnectPlan { + endpoint: self.endpoint.clone(), + session_id: self.session_id, + device_index: self.device_index, + responder: self.pending.responder(), + }) + } + + /// Install a connection opened by [`wasm_connect`] and publish its handshake results. + #[cfg(target_family = "wasm")] + pub(crate) fn wasm_install(&mut self, connected: WasmConnected) { + if self.writer.is_some() { + // A concurrent connect already installed a session; drop this one rather than + // leaking two writers for the same device. + return; + } + let _ = self.settings.set(connected.settings); + let _ = self.device_count.set(connected.device_count); + self.writer = Some(connected.writer); + } + + /// Hand whatever's currently buffered to the writer task as one batch (the writer + /// serializes it off the runner thread). No-op when the buffer is empty; otherwise opens + /// the connection first if it isn't already up. + pub fn flush(&mut self) { + if self.batch.is_empty() { + return; + } + self.ensure_connected(); + log::trace!("Flush session: {}", self.session_id); + let batch = self.batch.take(); + let writer = self + .writer + .as_ref() + .expect("writer is set by ensure_connected"); + writer.send(&self.executor, batch); + } +} + +impl Drop for RemoteService { + fn drop(&mut self) { + if self.closed { + return; + } + self.closed = true; + + // If we never connected, there's no server-side session to close and no writer to + // drain — whatever was buffered never had a connection to go out on, so just drop it. + if self.writer.is_none() { + return; + } + + // Best-effort teardown: append Close to whatever's still buffered and let the + // writer drain + flush it before we join the task, so the runtime isn't torn down + // mid-send. Serialization happens in the writer task now, so Drop can't panic on it. + self.batch.push(RemoteMessage::Close(self.session_id)); + let batch = self.batch.take(); + let writer = self + .writer + .as_mut() + .expect("writer present (checked above)"); + writer.shutdown(&self.executor, Some(batch)); + } +} diff --git a/crates/burn-remote/src/client/service/batch.rs b/crates/burn-remote/src/client/service/batch.rs new file mode 100644 index 0000000..fd6aa89 --- /dev/null +++ b/crates/burn-remote/src/client/service/batch.rs @@ -0,0 +1,142 @@ +//! Outgoing task buffering. + +use crate::shared::{RemoteMessage, Task}; + +/// Accumulates outgoing [`RemoteMessage`]s on the runner thread and decides when a batch is ready +/// for the wire. +/// +/// Pure buffering logic — no runtime, no socket. [`push`](Self::push) reports when a flush +/// threshold is reached so the caller can flush; [`take`](Self::take) drains the buffer into a +/// frame, leaving it empty for the next batch. +/// +/// Two thresholds drive a flush, whichever hits first: +/// - **task count** — bounds latency for streams of many tiny ops; +/// - **buffered data bytes** — bounds how much tensor data sits unsent, so a big upload (or several +/// moderate ones) goes out promptly instead of waiting on the count threshold. Tracking a running +/// sum, rather than flushing on each individual large message, lets small tensors keep batching +/// while still capping the in-flight data. +/// +/// Batching is wire-level only: every task keeps its own `StreamId`/`RequestId`, so the +/// server sees the same per-task semantics whether tasks arrive batched or one per frame. +pub(crate) struct OutgoingBatch { + tasks: Vec, + threshold: usize, + /// Running sum of [`data_len`] over the buffered tasks; reset by [`take`](Self::take). + bytes: usize, + bytes_threshold: usize, +} + +/// Size of the bulk tensor data a message carries, in bytes (0 for metadata-only messages). Drives +/// the byte-based flush threshold. +fn data_len(msg: &RemoteMessage) -> usize { + match msg { + RemoteMessage::Task(Task::RegisterTensor(_, _, data)) => data.bytes.len(), + _ => 0, + } +} + +impl OutgoingBatch { + /// Create a buffer that signals a flush once `threshold` tasks *or* `bytes_threshold` buffered + /// data bytes accumulate. + pub(crate) fn new(threshold: usize, bytes_threshold: usize) -> Self { + Self { + tasks: Vec::with_capacity(threshold), + threshold, + bytes: 0, + bytes_threshold, + } + } + + /// Append a task. Returns `true` once the buffer has reached either flush threshold, i.e. + /// the caller should [`take`](Self::take) and send. + pub(crate) fn push(&mut self, task: RemoteMessage) -> bool { + self.bytes += data_len(&task); + self.tasks.push(task); + self.tasks.len() >= self.threshold || self.bytes >= self.bytes_threshold + } + + pub(crate) fn is_empty(&self) -> bool { + self.tasks.is_empty() + } + + /// Drain the accumulated tasks, leaving the buffer empty for the next batch. + pub(crate) fn take(&mut self) -> Vec { + self.bytes = 0; + std::mem::take(&mut self.tasks) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::shared::Task; + use burn_backend::{StreamId, TensorData}; + use burn_ir::TensorId; + + /// A metadata-only task (contributes 0 to the byte counter). + fn task() -> RemoteMessage { + RemoteMessage::Task(Task::Seed(0)) + } + + /// A `RegisterTensor` carrying `n` bytes of data. + fn data_task(n: usize) -> RemoteMessage { + RemoteMessage::Task(Task::RegisterTensor( + StreamId::current(), + TensorId::new(0), + TensorData::new(vec![0u8; n], [n]), + )) + } + + // A large byte threshold so these count-driven tests aren't affected by the byte trigger. + const NO_BYTES: usize = usize::MAX; + + #[test] + fn push_signals_flush_only_at_threshold() { + let mut batch = OutgoingBatch::new(3, NO_BYTES); + assert!(!batch.push(task())); // 1 + assert!(!batch.push(task())); // 2 + assert!(batch.push(task())); // 3 == threshold + } + + #[test] + fn threshold_of_one_flushes_every_push() { + let mut batch = OutgoingBatch::new(1, NO_BYTES); + assert!(batch.push(task())); + assert!(batch.push(task())); + } + + #[test] + fn flushes_once_buffered_bytes_reach_threshold() { + // Count threshold high enough that only the byte threshold can trigger. + let mut batch = OutgoingBatch::new(100, 1000); + assert!(!batch.push(data_task(400))); // 400 buffered + assert!(!batch.push(data_task(400))); // 800 buffered + assert!(batch.push(data_task(400))); // 1200 >= 1000 -> flush + } + + #[test] + fn take_resets_the_byte_counter() { + let mut batch = OutgoingBatch::new(100, 1000); + assert!(!batch.push(data_task(900))); + batch.take(); // drains tasks and resets the byte sum + // Counter restarted from zero, so 900 again does not trip the 1000-byte threshold. + assert!(!batch.push(data_task(900))); + } + + #[test] + fn take_drains_and_resets() { + let mut batch = OutgoingBatch::new(4, NO_BYTES); + assert!(batch.is_empty()); + + batch.push(task()); + batch.push(task()); + assert!(!batch.is_empty()); + + let drained = batch.take(); + assert_eq!(drained.len(), 2); + + // Taking resets the buffer, so a subsequent take yields nothing. + assert!(batch.is_empty()); + assert_eq!(batch.take().len(), 0); + } +} diff --git a/crates/burn-remote/src/client/service/conn.rs b/crates/burn-remote/src/client/service/conn.rs new file mode 100644 index 0000000..f93743d --- /dev/null +++ b/crates/burn-remote/src/client/service/conn.rs @@ -0,0 +1,190 @@ +//! The client's transport connection layer. +//! +//! All `cfg(feature = ...)` transport selection on the client lives here: how to reach a peer +//! ([`RemoteEndpoint`]), the two halves of an opened session +//! ([`SubmitChannel`] / [`ResponseChannel`]), and the connect logic ([`open_channels`]). The +//! service and the registry are written against these and stay transport-agnostic. + +use std::sync::Arc; + +use crate::transport::link::{FrameSink, FrameSource}; +use crate::{PeerAddr, PeerId}; + +#[cfg(feature = "iroh")] +use crate::transport::iroh::node::RemoteNode; +#[cfg(feature = "websocket")] +use burn_communication::{Address, ProtocolClient}; + +/// Everything needed to establish a session with a remote compute peer. +#[derive(Clone, Debug)] +pub(crate) enum RemoteEndpoint { + #[cfg(feature = "iroh")] + Iroh { + node: RemoteNode, + peer: iroh::EndpointAddr, + authorization: Arc<[u8]>, + }, + #[cfg(feature = "websocket")] + WebSocket { + address: Address, + authorization: Arc<[u8]>, + }, +} + +impl RemoteEndpoint { + pub(crate) fn peer_addr(&self) -> PeerAddr { + match self { + #[cfg(feature = "iroh")] + Self::Iroh { peer, .. } => PeerAddr::Iroh(peer.clone()), + #[cfg(feature = "websocket")] + Self::WebSocket { address, .. } => PeerAddr::WebSocket(address.clone()), + } + } + + pub(crate) fn peer_id(&self) -> PeerId { + self.peer_addr().id() + } + + pub(crate) fn authorization(&self) -> &[u8] { + match self { + #[cfg(feature = "iroh")] + Self::Iroh { authorization, .. } => authorization, + #[cfg(feature = "websocket")] + Self::WebSocket { authorization, .. } => authorization, + } + } + + /// The stable registry key for this endpoint (identity + authorization, no mutable dialing + /// hints), so the same compute peer reuses one device id across reconnects. + pub(crate) fn key(&self) -> EndpointKey { + match self { + #[cfg(feature = "iroh")] + Self::Iroh { + node, + peer, + authorization, + .. + } => EndpointKey::Iroh { + local: node.id(), + remote: peer.id, + authorization: authorization.clone(), + }, + #[cfg(feature = "websocket")] + Self::WebSocket { + address, + authorization, + .. + } => EndpointKey::WebSocket { + address: address.clone(), + authorization: authorization.clone(), + }, + } + } +} + +/// Stable identity used to deduplicate endpoints in the device registry. +#[derive(Clone, PartialEq, Eq, Hash)] +pub(crate) enum EndpointKey { + #[cfg(feature = "iroh")] + Iroh { + local: iroh::EndpointId, + remote: iroh::EndpointId, + authorization: Arc<[u8]>, + }, + #[cfg(feature = "websocket")] + WebSocket { + address: Address, + authorization: Arc<[u8]>, + }, +} + +/// Outgoing half of an opened session (the submit stream). +pub(crate) enum SubmitChannel { + #[cfg(feature = "iroh")] + Iroh(iroh::endpoint::SendStream), + #[cfg(feature = "websocket")] + WebSocket(Box), +} + +impl SubmitChannel { + pub(crate) async fn send(&mut self, bytes: bytes::Bytes) -> Result<(), String> { + match self { + #[cfg(feature = "iroh")] + Self::Iroh(stream) => FrameSink::send(stream, bytes).await, + #[cfg(feature = "websocket")] + Self::WebSocket(sink) => FrameSink::send(sink.as_mut(), bytes).await, + } + } + + pub(crate) async fn close(&mut self) -> Result<(), String> { + match self { + #[cfg(feature = "iroh")] + Self::Iroh(stream) => FrameSink::close(stream).await, + #[cfg(feature = "websocket")] + Self::WebSocket(sink) => FrameSink::close(sink.as_mut()).await, + } + } +} + +/// Incoming half of an opened session (the response stream). +pub(crate) enum ResponseChannel { + #[cfg(feature = "iroh")] + Iroh(iroh::endpoint::RecvStream), + #[cfg(feature = "websocket")] + WebSocket(Box), +} + +impl ResponseChannel { + pub(crate) async fn recv(&mut self) -> Result, String> { + match self { + #[cfg(feature = "iroh")] + Self::Iroh(stream) => FrameSource::recv(stream).await, + #[cfg(feature = "websocket")] + Self::WebSocket(stream) => FrameSource::recv(stream.as_mut()).await, + } + } +} + +/// Open the session for `endpoint`, returning its submit + response halves. +/// +/// Done up front so a missing server surfaces here rather than on the first op, and the demux / +/// writer tasks can be spawned on already-open streams. +pub(crate) async fn open_channels( + endpoint: &RemoteEndpoint, +) -> Result<(SubmitChannel, ResponseChannel), String> { + match endpoint { + #[cfg(feature = "iroh")] + RemoteEndpoint::Iroh { node, peer, .. } => { + let (send, recv) = node + .open_stream( + &PeerAddr::Iroh(peer.clone()), + crate::transport::iroh::node::StreamKind::Session, + ) + .await?; + Ok((SubmitChannel::Iroh(send), ResponseChannel::Iroh(recv))) + } + #[cfg(feature = "websocket")] + RemoteEndpoint::WebSocket { address, .. } => { + use burn_communication::websocket::WsClient; + // One full-duplex socket per session, split into the submit (sink) + response (source) + // halves — matching the Iroh single-stream model. + let channel = WsClient::connect(address.clone(), "session") + .await + .map_err(|err| connect_error("session", &endpoint.peer_addr(), &err))?; + let (sink, source) = channel.split(); + Ok(( + SubmitChannel::WebSocket(Box::new(sink)), + ResponseChannel::WebSocket(Box::new(source)), + )) + } + } +} + +/// Actionable panic message for a failed channel connect. +#[cfg(feature = "websocket")] +fn connect_error(route: &str, peer: &PeerAddr, err: &E) -> String { + format!( + "Failed to open remote '{route}' channel to {peer}: {err:?}. \ + Is a `burn-remote` compute node running at that peer?" + ) +} diff --git a/crates/burn-remote/src/client/service/pending.rs b/crates/burn-remote/src/client/service/pending.rs new file mode 100644 index 0000000..ff4745f --- /dev/null +++ b/crates/burn-remote/src/client/service/pending.rs @@ -0,0 +1,201 @@ +//! Request/response correlation. + +use crate::shared::{RequestId, TaskResponseContent}; +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, +}; +use tokio::sync::oneshot; + +/// The callbacks awaiting replies, plus whether the connection is still live. +/// +/// Both live under one mutex so that registering a callback and tearing every callback down on +/// disconnect are mutually exclusive: a registration can't slip into the map *after* the +/// response-demux task has already drained it on disconnect and left it to wait for a reply that +/// will never come. +struct State { + callbacks: HashMap>, + /// Cleared by [`Responder::disconnect`] once the response stream is gone; gates new + /// registrations so a post-disconnect request fails fast instead of parking forever. + connected: bool, +} + +type SharedState = Arc>; + +/// Correlates response-producing requests with the caller awaiting each one. +/// +/// Each response-producing task ([`ReadTensor`](crate::shared::Task::ReadTensor), +/// `SyncBackend`, `DTypeUsage`) carries a [`RequestId`]; the server echoes it on the +/// response. The runner thread [`register`](Self::register)s a [`oneshot`] callback before +/// sending the task, and the response-demux task delivers the reply through a [`Responder`]. +/// +/// The state is guarded by a plain [`std::sync::Mutex`]: the lock is only ever held for a single +/// insert/remove/drain and never across an `.await`, so neither the runner thread nor the demux +/// task needs the tokio runtime to touch it. +pub(crate) struct PendingResponses { + state: SharedState, + next_id: RequestId, +} + +impl PendingResponses { + pub(crate) fn new() -> Self { + Self { + state: Arc::new(Mutex::new(State { + callbacks: HashMap::new(), + connected: true, + })), + next_id: 0, + } + } + + /// Allocate the next monotonic [`RequestId`]. + pub(crate) fn next_id(&mut self) -> RequestId { + let id = self.next_id; + self.next_id += 1; + id + } + + /// Register a callback for `id`, returning the receiver the caller awaits for the reply. + /// + /// If the connection has already dropped (the demux task called [`Responder::disconnect`]), + /// the sender is dropped immediately, so the returned receiver resolves to a `RecvError` right + /// away rather than blocking on a response the dead server will never send. + pub(crate) fn register(&self, id: RequestId) -> oneshot::Receiver { + let (tx, rx) = oneshot::channel(); + let mut state = self.state.lock().unwrap(); + if state.connected { + state.callbacks.insert(id, tx); + } + // Disconnected: drop `tx` here, leaving `rx` already-closed. + rx + } + + /// A cheap, cloneable handle the response-demux task uses to deliver replies. + pub(crate) fn responder(&self) -> Responder { + Responder { + state: self.state.clone(), + } + } +} + +/// Delivers responses to the callbacks registered in [`PendingResponses`]. Held by the +/// response-demux task, decoupled from the [`PendingResponses`] the runner thread owns. +pub(crate) struct Responder { + state: SharedState, +} + +impl Responder { + /// Deliver `content` to the caller waiting on `id`. Returns `false` if no callback is + /// registered (unknown id, or the caller dropped its receiver), in which case the + /// response is discarded. + pub(crate) fn complete(&self, id: RequestId, content: TaskResponseContent) -> bool { + match self.state.lock().unwrap().callbacks.remove(&id) { + Some(tx) => { + // Receiver dropped is fine (caller no longer cares). + let _ = tx.send(content); + true + } + None => false, + } + } + + /// Mark the connection dropped and fail every pending caller. + /// + /// Called by the response-demux task when the response stream closes or errors (server down, + /// connection reset). Clearing the map drops every registered sender, so each waiting receiver + /// resolves to a `RecvError` that the callers (`sync`, `read_tensor`, `dtype_usage`) translate + /// into an error instead of blocking forever. Clearing `connected` makes any later request + /// fail fast in [`PendingResponses::register`] rather than registering a doomed callback. + pub(crate) fn disconnect(&self) { + let mut state = self.state.lock().unwrap(); + state.connected = false; + state.callbacks.clear(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn content() -> TaskResponseContent { + TaskResponseContent::SyncBackend(Ok(())) + } + + #[test] + fn next_id_is_monotonic() { + let mut pending = PendingResponses::new(); + assert_eq!(pending.next_id(), 0); + assert_eq!(pending.next_id(), 1); + assert_eq!(pending.next_id(), 2); + } + + #[test] + fn register_then_complete_delivers_to_receiver() { + let mut pending = PendingResponses::new(); + let id = pending.next_id(); + let mut rx = pending.register(id); + + assert!(pending.responder().complete(id, content())); + // `try_recv` resolves synchronously once the sender has fired — no runtime needed. + assert!(matches!( + rx.try_recv(), + Ok(TaskResponseContent::SyncBackend(Ok(()))) + )); + } + + #[test] + fn complete_unknown_id_returns_false() { + let pending = PendingResponses::new(); + assert!(!pending.responder().complete(42, content())); + } + + #[test] + fn complete_consumes_the_callback() { + let mut pending = PendingResponses::new(); + let id = pending.next_id(); + let _rx = pending.register(id); + + assert!(pending.responder().complete(id, content())); + // The callback was removed on first delivery; a duplicate response finds nothing. + assert!(!pending.responder().complete(id, content())); + } + + #[test] + fn disconnect_fails_every_pending_caller() { + let mut pending = PendingResponses::new(); + let id0 = pending.next_id(); + let id1 = pending.next_id(); + let mut rx0 = pending.register(id0); + let mut rx1 = pending.register(id1); + + // The response stream died with both requests still in flight. + pending.responder().disconnect(); + + // Both receivers resolve immediately with an error instead of hanging. + assert!(matches!( + rx0.try_recv(), + Err(oneshot::error::TryRecvError::Closed) + )); + assert!(matches!( + rx1.try_recv(), + Err(oneshot::error::TryRecvError::Closed) + )); + } + + #[test] + fn register_after_disconnect_returns_a_closed_receiver() { + let mut pending = PendingResponses::new(); + pending.responder().disconnect(); + + // A request issued after the connection dropped must not park forever. + let id = pending.next_id(); + let mut rx = pending.register(id); + assert!(matches!( + rx.try_recv(), + Err(oneshot::error::TryRecvError::Closed) + )); + + // And it was never inserted, so a late response for it finds nothing. + assert!(!pending.responder().complete(id, content())); + } +} diff --git a/crates/burn-remote/src/client/service/registry.rs b/crates/burn-remote/src/client/service/registry.rs new file mode 100644 index 0000000..4dc2875 --- /dev/null +++ b/crates/burn-remote/src/client/service/registry.rs @@ -0,0 +1,144 @@ +//! Process-global registry connecting Burn's compact `DeviceId` to rich remote endpoints. + +use burn_ir::TensorId; +use burn_std::DeviceSettings; +use std::{ + collections::HashMap, + sync::{ + Arc, Mutex, OnceLock, + atomic::{AtomicU64, Ordering}, + }, +}; + +use super::conn::{EndpointKey, RemoteEndpoint}; +use crate::client::runtime::Executor; + +static TENSOR_ID_COUNTER: AtomicU64 = AtomicU64::new(0); + +pub(crate) fn new_tensor_id() -> TensorId { + TensorId::new(TENSOR_ID_COUNTER.fetch_add(1, Ordering::Relaxed)) +} + +struct EndpointRegistry { + next_index: u32, + by_endpoint: HashMap<(EndpointKey, u32), u32>, + by_index: HashMap, +} + +#[derive(Clone)] +struct EndpointEntry { + endpoint: RemoteEndpoint, + executor: Executor, + device_index: u32, + settings: Arc>, + device_count: Arc>, +} + +static REGISTRY: OnceLock> = OnceLock::new(); + +fn registry() -> &'static Mutex { + REGISTRY.get_or_init(|| { + Mutex::new(EndpointRegistry { + next_index: 0, + by_endpoint: HashMap::new(), + by_index: HashMap::new(), + }) + }) +} + +pub(crate) fn register_endpoint( + endpoint: RemoteEndpoint, + executor: Executor, + device_index: u32, +) -> u32 { + let key = (endpoint.key(), device_index); + let mut registry = registry().lock().unwrap(); + if let Some(id) = registry.by_endpoint.get(&key).copied() { + // Refresh mutable dialing hints + the captured runtime while preserving the stable device + // id and settings cells. + let entry = registry.by_index.get_mut(&id).unwrap(); + entry.endpoint = endpoint; + entry.executor = executor; + return id; + } + + let id = registry.next_index; + registry.next_index += 1; + registry.by_endpoint.insert(key, id); + registry.by_index.insert( + id, + EndpointEntry { + endpoint, + executor, + device_index, + settings: Arc::new(OnceLock::new()), + device_count: Arc::new(OnceLock::new()), + }, + ); + id +} + +pub(crate) fn endpoint_for(id: u32) -> Option<(RemoteEndpoint, u32)> { + registry() + .lock() + .unwrap() + .by_index + .get(&id) + .map(|entry| (entry.endpoint.clone(), entry.device_index)) +} + +/// The runtime captured for `id`'s device at construction, used to drive its session tasks. +pub(crate) fn executor_for(id: u32) -> Option { + registry() + .lock() + .unwrap() + .by_index + .get(&id) + .map(|entry| entry.executor.clone()) +} + +pub(crate) fn settings_for(id: u32) -> DeviceSettings { + *settings_cell(id) + .get() + .expect("Remote service has not connected to this device yet") +} + +pub(crate) fn has_settings(id: u32) -> bool { + registry() + .lock() + .unwrap() + .by_index + .get(&id) + .is_some_and(|entry| entry.settings.get().is_some()) +} + +pub(crate) fn settings_cell(id: u32) -> Arc> { + registry() + .lock() + .unwrap() + .by_index + .get(&id) + .expect("Device id not registered") + .settings + .clone() +} + +pub(crate) fn device_count_cell(id: u32) -> Arc> { + registry() + .lock() + .unwrap() + .by_index + .get(&id) + .expect("Device id not registered") + .device_count + .clone() +} + +pub(crate) fn device_count_for(id: u32) -> Option { + registry() + .lock() + .unwrap() + .by_index + .get(&id) + .and_then(|entry| entry.device_count.get().copied()) +} diff --git a/crates/burn-remote/src/client/service/writer.rs b/crates/burn-remote/src/client/service/writer.rs new file mode 100644 index 0000000..2d3419f --- /dev/null +++ b/crates/burn-remote/src/client/service/writer.rs @@ -0,0 +1,137 @@ +//! Outgoing-frame writer task. + +use crate::client::runtime::{Executor, SpawnHandle}; +use crate::client::service::SubmitChannel; +use crate::shared::RemoteMessage; +use tokio::sync::mpsc; + +/// Bound on task batches queued for the writer task on native targets. +/// +/// [`send`](SubmitWriter::send) enqueues here instead of touching the socket directly. +/// Bounded so a stalled socket surfaces as backpressure on the runner thread (a `send` +/// blocks only when the writer has fallen this many batches behind) rather than as unbounded +/// memory growth. The browser cannot block the runner, so the wasm build uses an unbounded +/// channel and never applies backpressure. +#[cfg(not(target_family = "wasm"))] +const WRITE_QUEUE_CAP: usize = 16; + +#[cfg(not(target_family = "wasm"))] +type BatchSender = mpsc::Sender>; +#[cfg(target_family = "wasm")] +type BatchSender = mpsc::UnboundedSender>; + +/// Owns the submit channel on the service runtime and turns task batches into wire frames. +/// +/// The runner thread hands raw [`RemoteMessage`] batches to [`send`](Self::send) over a channel; +/// the writer task serializes and `await`s each socket send fully before pulling the next, so +/// frames reach the wire in FIFO order without ever parking the runner thread on serialization or +/// the network. Serializing here (rather than on the runner thread) lets encoding one frame — +/// which for `RegisterTensor` carries full tensor payloads — overlap with the runner registering +/// the next op. That single-task FIFO drain is also what guarantees a frame is fully flushed +/// before the next begins — the socket sink itself offers no such queue. +pub(crate) struct SubmitWriter { + /// `Option` so [`shutdown`](Self::shutdown) can drop the sender to signal the task to + /// finish once it has drained. + tx: Option, + /// Joined on shutdown on native; on wasm the detached task drains on the event loop. + handle: Option, +} + +impl SubmitWriter { + /// Spawn the writer task on `runtime`, taking ownership of the submit `channel`. + pub(crate) fn spawn(runtime: &Executor, mut channel: SubmitChannel) -> Self { + #[cfg(not(target_family = "wasm"))] + let (tx, mut rx) = mpsc::channel::>(WRITE_QUEUE_CAP); + #[cfg(target_family = "wasm")] + let (tx, mut rx) = mpsc::unbounded_channel::>(); + + let handle = runtime.spawn(async move { + while let Some(batch) = rx.recv().await { + let bytes: bytes::Bytes = match rmp_serde::to_vec(&batch) { + Ok(b) => b.into(), + Err(err) => { + log::error!("Failed to serialize outgoing task batch: {err:?}; dropping"); + continue; + } + }; + if let Err(err) = channel.send(bytes).await { + log::warn!("Remote submit writer send failed: {err:?}; closing writer"); + return; + } + } + let _ = channel.close().await; + }); + Self { + tx: Some(tx), + handle: Some(handle), + } + } + + /// Enqueue a task batch for serialization + the wire. + /// + /// Native: a non-blocking [`try_send`](mpsc::Sender::try_send) fast path; only when the writer + /// has fallen [`WRITE_QUEUE_CAP`] batches behind (socket backpressure) do we block the runner + /// thread waiting for room. Wasm: an unbounded, always non-blocking send. + #[cfg(not(target_family = "wasm"))] + pub(crate) fn send(&self, runtime: &Executor, batch: Vec) { + let Some(tx) = self.tx.as_ref() else { + log::warn!("Remote submit writer already shut down; dropping outgoing batch"); + return; + }; + match tx.try_send(batch) { + Ok(()) => {} + Err(mpsc::error::TrySendError::Full(batch)) => { + // Backpressure: wait for the writer to drain a slot. FIFO order is preserved + // because this runs on the single runner thread, after the fast-path sends. + if runtime.block_on(tx.send(batch)).is_err() { + log::warn!( + "Remote submit writer task has exited (server disconnected?); dropping outgoing batch" + ); + } + } + // The writer task exits if a socket send fails (server gone / connection reset). + // Drop the batch with a warning rather than panicking on the runner thread, which + // would crash the caller's thread instead of surfacing a recoverable disconnect. + Err(mpsc::error::TrySendError::Closed(_)) => { + log::warn!( + "Remote submit writer task has exited (server disconnected?); dropping outgoing batch" + ); + } + } + } + + #[cfg(target_family = "wasm")] + pub(crate) fn send(&self, _runtime: &Executor, batch: Vec) { + let Some(tx) = self.tx.as_ref() else { + log::warn!("Remote submit writer already shut down; dropping outgoing batch"); + return; + }; + if tx.send(batch).is_err() { + log::warn!( + "Remote submit writer task has exited (server disconnected?); dropping outgoing batch" + ); + } + } + + /// Best-effort teardown: enqueue an optional final batch, drop the sender so the writer + /// drains and exits, then (on native) join it so the runtime isn't torn down mid-send. In the + /// browser the task is detached and finishes draining on the event loop after the sender drops. + pub(crate) fn shutdown(&mut self, runtime: &Executor, final_batch: Option>) { + if let (Some(batch), Some(tx)) = (final_batch, self.tx.as_ref()) { + #[cfg(not(target_family = "wasm"))] + let _ = runtime.block_on(tx.send(batch)); + #[cfg(target_family = "wasm")] + let _ = tx.send(batch); + } + // Drop the sender so the writer's `rx.recv()` returns `None` once it has drained. + self.tx.take(); + if let Some(handle) = self.handle.take() { + #[cfg(not(target_family = "wasm"))] + runtime.join(handle); + #[cfg(target_family = "wasm")] + let _ = handle; + } + #[cfg(target_family = "wasm")] + let _ = runtime; + } +} diff --git a/crates/burn-remote/src/lib.rs b/crates/burn-remote/src/lib.rs new file mode 100644 index 0000000..a42b2eb --- /dev/null +++ b/crates/burn-remote/src/lib.rs @@ -0,0 +1,1004 @@ +//! Peer-to-peer remote tensor execution for Burn. +//! +//! Iroh is the primary transport. Applications own an Iroh [`Endpoint`] and build remote devices +//! from it; a server hosts compute on its own endpoint. Compute sessions use bidirectional QUIC +//! streams, while cross-peer tensor movement uses independent authenticated streams without +//! routing payloads through the controlling client. +//! +//! The optional `websocket` feature retains the legacy address-and-port transport. + +#[cfg(feature = "client")] +mod client; + +#[cfg(feature = "server")] +pub mod server; + +pub(crate) mod shared; +pub mod telemetry; +mod transport; + +pub use burn_ir as ir; +pub use burn_router::RouterClient; + +/// Network-traffic savings metric for op-graph caching, shared by the client device service and the +/// server session worker. +#[cfg(any(feature = "client", feature = "server"))] +pub(crate) mod metrics; + +#[cfg(feature = "iroh")] +pub use iroh::{Endpoint, EndpointAddr, EndpointId}; +#[cfg(feature = "iroh")] +pub use transport::iroh::RemoteSecret; +#[cfg(feature = "iroh")] +pub use transport::iroh::node::BURN_REMOTE_ALPN; +pub use transport::{PeerAddr, PeerId}; + +#[cfg(feature = "client")] +mod __client { + use super::*; + + use burn_router::BackendRouter; + + /// The remote backend allows you to run computation on a remote device. + /// + /// Iroh is the primary transport. Applications own an Iroh [`Endpoint`], resolve a compute peer + /// through their own discovery/control plane, and construct devices from the endpoint and the + /// peer's address with [`RemoteDevice::iroh`] (or the `Device::remote_iroh` facade). + /// + /// ```rust, ignore + /// let endpoint = Endpoint::builder(presets::N0).bind().await?; + /// let remote = RemoteDevice::iroh(&endpoint, compute_peer, 0); + /// ``` + /// + /// For backends that aren't part of `DispatchDevice` but implement + /// `BackendIr`, build a [`server::RemoteServerBuilder`] directly with the + /// concrete backend type parameter — that is also how custom operations + /// (backend extensions) are hosted, via + /// [`custom_op`](server::RemoteServerBuilder::custom_op). + #[cfg(not(feature = "fusion"))] + pub type RemoteBackend = BackendRouter; + + /// With the `fusion` feature enabled, the remote backend is wrapped in + /// [`Fusion`](burn_fusion::Fusion) — exactly like the CubeCL backends — so recurring groups of + /// operations are cached on the server and invoked by id, sending a repeated computation + /// (e.g. a model block per step) over the network once instead of every step. + #[cfg(feature = "fusion")] + pub type RemoteBackend = burn_fusion::Fusion>; + + pub use client::{CustomOpClient, RemoteChannel, RemoteDevice}; +} +#[cfg(feature = "client")] +pub use __client::*; + +#[cfg(all(test, feature = "client", feature = "server"))] +mod tests { + use burn_flex::Flex; + use burn_tensor::{Device, DeviceType, Distribution, Tensor}; + + /// Run `body` on a worker thread and fail the test if it doesn't finish within `timeout`. + /// + /// A deadlock in the remote backend manifests as a hung worker, so without a watchdog the + /// test would block the whole suite forever. We can't forcibly kill the hung thread (it's + /// parked on a blocking recv deep in the backend), so on timeout we panic from the test + /// thread and let the process exit carry the stuck worker away. + fn with_deadlock_watchdog(timeout: std::time::Duration, body: impl FnOnce() + Send + 'static) { + let (tx, rx) = std::sync::mpsc::channel(); + let handle = std::thread::spawn(move || { + body(); + let _ = tx.send(()); + }); + match rx.recv_timeout(timeout) { + Ok(()) => { + handle.join().expect("worker thread panicked"); + } + Err(_) => panic!( + "Deadlock: the remote multi-device workload did not finish within {timeout:?}" + ), + } + } + + #[test] + pub fn test_to_device_over_websocket() { + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_io() + .build() + .unwrap(); + + rt.spawn( + crate::server::RemoteServerBuilder::::new(vec![Default::default()]) + .port(3000) + .start_async(), + ); + rt.spawn( + crate::server::RemoteServerBuilder::::new(vec![Default::default()]) + .port(3010) + .start_async(), + ); + + // Give the servers a moment to bind before clients try to connect. + std::thread::sleep(std::time::Duration::from_millis(500)); + + let device_1 = Device::remote_websocket("ws://localhost:3000", 0); + let device_2 = Device::remote_websocket("ws://localhost:3010", 0); + + // Some random input on device 1. + let input_shape = [1, 28, 28]; + let input = Tensor::<3>::random(input_shape, Distribution::Default, &device_1); + let numbers_expected: Vec = input.to_data().to_vec().unwrap(); + + // Move tensor to device 2. + let input = input.to_device(&device_2); + let numbers: Vec = input.to_data().to_vec().unwrap(); + assert_eq!(numbers, numbers_expected); + + // Move tensor back to device 1. + let input = input.to_device(&device_1); + let numbers: Vec = input.to_data().to_vec().unwrap(); + assert_eq!(numbers, numbers_expected); + + rt.shutdown_background(); + } + + /// End-to-end backend extension over the wire: the client ships a custom op as + /// `OperationIr::Custom`, and the server executes it through a handler registered on the + /// builder. Mirrors how a backend extension hosts its ops — the user hand-writes the client + /// side (here, building the `CustomOpIr`) and registers the server handler. + /// + /// Only runs without `fusion`, since it drives the router client (`RemoteBackend`) directly. + #[test] + #[cfg(not(feature = "fusion"))] + pub fn test_custom_op_over_websocket() { + use crate::{RemoteBackend, RemoteDevice}; + use burn_backend::{Scalar, TensorData, TensorMetadata, ops::FloatTensorOps}; + use burn_ir::{CustomOpIr, OperationIr, ScalarIr, TensorIr}; + use burn_router::RouterClient; + + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_io() + .build() + .unwrap(); + + // Host a "scale" custom op: multiply the input float tensor by a scalar argument. + rt.spawn( + crate::server::RemoteServerBuilder::::new(vec![Default::default()]) + .port(3200) + .custom_op("scale", |handles, ir, _device| { + let input = handles.get_float_tensor::(&ir.inputs[0]); + let factor: Scalar = ir.scalars[0].into(); + let output = Flex::float_mul_scalar(input, factor); + handles.register_float_tensor::(&ir.outputs[0].id, output); + }) + .start_async(), + ); + + // Give the server a moment to bind before the client connects. + std::thread::sleep(std::time::Duration::from_millis(500)); + + // Drive the remote backend directly (no autodiff/dispatch glue). A real backend extension + // would wrap this in a hand-written `impl MyExt for RemoteBackend`. + let device = RemoteDevice::websocket("ws://localhost:3200", 0); + let input = >::float_from_data( + TensorData::from([2.0f32, 4.0, 6.0]), + &device, + ); + + // Client side: build the custom op (input tensor + the scale factor as a scalar) and ship + // it through the remote client as `OperationIr::Custom`. + let client = input.client.clone(); + let shape = input.shape(); + let dtype = input.dtype(); + let out_ir = TensorIr::uninit(client.create_empty_handle(), shape, dtype); + let desc = CustomOpIr::with_scalars( + "scale", + &[input.into_ir()], + &[out_ir], + vec![ScalarIr::Float(3.0)], + ); + let out = client.register(OperationIr::Custom(desc)).remove(0); + + let data = rt + .block_on(>::float_into_data(out)) + .unwrap(); + let values: Vec = data.to_vec().unwrap(); + assert_eq!(values, vec![6.0, 12.0, 18.0]); + + rt.shutdown_background(); + } + + /// The Iroh counterpart of [`test_custom_op_over_websocket`]: the server hosts a "scale" handler + /// registered on the protocol's custom-op registry, and the client ships the op as + /// `OperationIr::Custom`. Confirms custom ops travel the merged session path over Iroh just as + /// they do over WebSocket. + /// + /// Only runs without `fusion`, since it drives the router client (`RemoteBackend`) directly. + #[test] + #[cfg(not(feature = "fusion"))] + pub fn test_custom_op_over_iroh() { + use crate::{ + BURN_REMOTE_ALPN, RemoteBackend, RemoteDevice, + server::{AllowAll, CustomOpRegistry, IrohRemoteProtocol}, + telemetry::TelemetryProbe, + }; + use burn_backend::{Scalar, TensorData, TensorMetadata, ops::FloatTensorOps}; + use burn_ir::{CustomOpIr, OperationIr, ScalarIr, TensorIr}; + use burn_router::RouterClient; + use iroh::{Endpoint, RelayMode, endpoint::presets, protocol::Router}; + + async fn local_endpoint() -> Endpoint { + Endpoint::builder(presets::Minimal) + .relay_mode(RelayMode::Disabled) + .clear_ip_transports() + .bind_addr("127.0.0.1:0") + .unwrap() + .bind() + .await + .unwrap() + } + + // Server on its own runtime, hosting a "scale" custom op (multiply the input by a scalar). + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap(); + + let server = rt.block_on(local_endpoint()); + let server_addr = server.addr(); + + let mut custom_ops = CustomOpRegistry::::default(); + custom_ops.register("scale", |handles, ir, _device| { + let input = handles.get_float_tensor::(&ir.inputs[0]); + let factor: Scalar = ir.scalars[0].into(); + let output = Flex::float_mul_scalar(input, factor); + handles.register_float_tensor::(&ir.outputs[0].id, output); + }); + + let protocol = IrohRemoteProtocol::::new( + server.clone(), + vec![Default::default()], + std::sync::Arc::new(AllowAll), + TelemetryProbe::disabled(), + custom_ops, + ); + + // spawn() must run in a runtime context so iroh can schedule its tasks. + let router = { + let _guard = rt.enter(); + Router::builder(server) + .accept(BURN_REMOTE_ALPN, protocol) + .spawn() + }; + + let client = rt.block_on(local_endpoint()); + let device = RemoteDevice::iroh(&client, server_addr, 0); + + // Client side: build the custom op and ship it as `OperationIr::Custom`, exactly as the + // WebSocket test does -- only the transport differs. + let input = >::float_from_data( + TensorData::from([2.0f32, 4.0, 6.0]), + &device, + ); + let remote_client = input.client.clone(); + let shape = input.shape(); + let dtype = input.dtype(); + let out_ir = TensorIr::uninit(remote_client.create_empty_handle(), shape, dtype); + let desc = CustomOpIr::with_scalars( + "scale", + &[input.into_ir()], + &[out_ir], + vec![ScalarIr::Float(3.0)], + ); + let out = remote_client.register(OperationIr::Custom(desc)).remove(0); + + let data = rt + .block_on(>::float_into_data(out)) + .unwrap(); + let values: Vec = data.to_vec().unwrap(); + assert_eq!(values, vec![6.0, 12.0, 18.0]); + + rt.block_on(router.shutdown()).unwrap(); + } + + /// A single server hosting multiple devices: two indices on the same address resolve to + /// two distinct sessions (distinct interpreters/runner threads). Moving a tensor between + /// them exercises the multi-device path within one host. + #[test] + pub fn test_multi_device_single_server() { + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_io() + .build() + .unwrap(); + + // One server, two devices. + rt.spawn( + crate::server::RemoteServerBuilder::::new(vec![ + Default::default(), + Default::default(), + ]) + .port(3030) + .start_async(), + ); + + std::thread::sleep(std::time::Duration::from_millis(500)); + + let device_0 = Device::remote_websocket("ws://localhost:3030", 0); + let device_1 = Device::remote_websocket("ws://localhost:3030", 1); + + // Distinct indices on the same address must be distinct devices. + assert_ne!(device_0, device_1); + + let input_shape = [1, 28, 28]; + let input = Tensor::<3>::random(input_shape, Distribution::Default, &device_0); + let numbers_expected: Vec = input.to_data().to_vec().unwrap(); + + // Move tensor to the second device on the same host and back. + let input = input.to_device(&device_1); + let numbers: Vec = input.to_data().to_vec().unwrap(); + assert_eq!(numbers, numbers_expected); + + let input = input.to_device(&device_0); + let numbers: Vec = input.to_data().to_vec().unwrap(); + assert_eq!(numbers, numbers_expected); + + rt.shutdown_background(); + } + + /// Concurrent multi-device regression (DDP-style): two user threads, each pinned to a device, + /// running simultaneously and each iteration moving a tensor to the *other* device and back. + /// + /// This used to deadlock because the client's transfer-id counter never persisted its + /// increment (`LocalTransferId` is `Copy`, so the increment landed on a throwaway local). + /// Every same-host transfer after the first reused the same id; sequentially that's harmless + /// (each expose is taken before the next), but two transfers in flight at once then collided + /// in the server's `local_comm` rendezvous and one `take` hung forever. Single-threaded + /// workloads never tripped it — `into_data` serializes each iteration — so the bug only shows + /// up under genuine concurrency. + #[test] + fn test_multi_device_concurrent_to_device_deadlock() { + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_io() + .build() + .unwrap(); + + rt.spawn( + crate::server::RemoteServerBuilder::::new(vec![ + Default::default(), + Default::default(), + ]) + .port(3060) + .start_async(), + ); + + std::thread::sleep(std::time::Duration::from_millis(500)); + + with_deadlock_watchdog(std::time::Duration::from_secs(30), || { + let device0 = Device::remote_websocket("ws://localhost:3060", 0); + let device1 = Device::remote_websocket("ws://localhost:3060", 1); + + let run = |home: Device, away: Device| { + move || { + for _ in 0..100 { + let t = Tensor::<2>::random([8, 8], Distribution::Default, &home); + // home -> away, op there, away -> home. + let t = t.to_device(&away); + let t = t * 2.0; + let t = t.to_device(&home); + let _ = t.sum().into_data(); + } + } + }; + + let h0 = std::thread::spawn(run(device0.clone(), device1.clone())); + let h1 = std::thread::spawn(run(device1, device0)); + h0.join().unwrap(); + h1.join().unwrap(); + }); + + rt.shutdown_background(); + } + + /// `Device::enumerate(DeviceType::remote_websocket(addr))` lists every device the server hosts, by + /// connecting once and reading the device count off the init handshake. + #[test] + pub fn test_enumerate_remote_devices() { + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_io() + .build() + .unwrap(); + + // One server hosting three devices. + rt.spawn( + crate::server::RemoteServerBuilder::::new(vec![ + Default::default(), + Default::default(), + Default::default(), + ]) + .port(3040) + .start_async(), + ); + + std::thread::sleep(std::time::Duration::from_millis(500)); + + let devices = + Device::enumerate(DeviceType::remote_websocket("ws://localhost:3040")).into_vec(); + + // The server reports its three devices, in index order. + assert_eq!(devices.len(), 3); + assert_eq!( + devices[0], + Device::remote_websocket("ws://localhost:3040", 0) + ); + assert_eq!( + devices[1], + Device::remote_websocket("ws://localhost:3040", 1) + ); + assert_eq!( + devices[2], + Device::remote_websocket("ws://localhost:3040", 2) + ); + + // Distinct indices are distinct devices. + assert_ne!(devices[0], devices[1]); + assert_ne!(devices[1], devices[2]); + + // The enumerated devices are usable: run an op on the last one. + let input = Tensor::<2>::from_floats([[1.0, 2.0], [3.0, 4.0]], &devices[2]); + let numbers: Vec = (input * 2.0).to_data().to_vec().unwrap(); + assert_eq!(numbers, vec![2.0, 4.0, 6.0, 8.0]); + + rt.shutdown_background(); + } + + /// Exercises the cross-backend transfer body: local tensor → remote (data round-trip + /// through `TensorData`), an op on the remote, then remote → local. + /// Run `body` on a worker thread and report whether it finished within `timeout`. + /// + /// Unlike [`with_deadlock_watchdog`], a panic inside `body` counts as "finished": these error + /// tests assert that a failure *surfaces* (as an `Err` or a panic) instead of hanging, so all + /// that matters is the thread came back. Returns `false` if it was still running at the + /// deadline — i.e. the call hung. + fn finishes_within(timeout: std::time::Duration, body: impl FnOnce() + Send + 'static) -> bool { + let (tx, rx) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + // Swallow panics: a disconnected read panicking on the error path is an acceptable + // "didn't hang" outcome and must not abort the whole test process. + let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(body)); + let _ = tx.send(()); + }); + rx.recv_timeout(timeout).is_ok() + } + + /// When the server goes down, a client call that awaits a response (here a tensor read) must + /// fail promptly instead of blocking forever. Before the fix the response-demux task just + /// exited on the closed stream, leaving every pending callback — and any later request — + /// parked on a oneshot that would never be completed. + #[test] + fn test_server_down_does_not_hang_client() { + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_io() + .build() + .unwrap(); + + rt.spawn( + crate::server::RemoteServerBuilder::::new(vec![Default::default()]) + .port(3070) + .start_async(), + ); + + std::thread::sleep(std::time::Duration::from_millis(500)); + + let device = Device::remote_websocket("ws://localhost:3070", 0); + + // One successful round-trip so the sockets are actually up and the demux task is running. + let input = Tensor::<2>::from_floats([[1.0, 2.0], [3.0, 4.0]], &device); + let warmup: Vec = (input * 2.0).to_data().to_vec().unwrap(); + assert_eq!(warmup, vec![2.0, 4.0, 6.0, 8.0]); + + // Kill the server: dropping its runtime closes the listener and both client sockets. + rt.shutdown_timeout(std::time::Duration::from_millis(100)); + // Let the client's response-demux observe the closed stream and fail pending callers. + std::thread::sleep(std::time::Duration::from_millis(500)); + + // A read now has no server to answer it. It must error out (which `to_data` surfaces as a + // panic), not hang. + let finished = finishes_within(std::time::Duration::from_secs(10), move || { + let t = Tensor::<2>::from_floats([[5.0, 6.0], [7.0, 8.0]], &device); + let _ = (t * 2.0).to_data(); + }); + assert!( + finished, + "client hung waiting for a response after the server went down" + ); + } + + /// A client that disconnects abruptly mid-session (socket dropped, no `Close`) must not wedge + /// the server: it should clean the session up and keep serving everyone else. We drive a raw + /// connection here because a `Device`'s client is process-cached and never dropped mid-test. + #[test] + fn test_client_disconnect_handled_cleanly_by_server() { + use crate::shared::{RemoteMessage, SessionId, Task}; + use burn_communication::{CommunicationChannel, Message, ProtocolClient}; + use std::str::FromStr; + + type Client = burn_communication::websocket::WsClient; + + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_io() + .build() + .unwrap(); + + rt.spawn( + crate::server::RemoteServerBuilder::::new(vec![Default::default()]) + .port(3090) + .start_async(), + ); + + std::thread::sleep(std::time::Duration::from_millis(500)); + + // Raw client: connect a submit stream, init a session and send one task so the server + // spawns the session worker, then drop the socket without a `Close` to mimic a crash. + { + let rtc = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let address = burn_communication::Address::from_str("ws://localhost:3090").unwrap(); + let session_id = SessionId::new(); + + rtc.block_on(async { + let mut submit = Client::connect(address, "session") + .await + .expect("raw session connect"); + + let frame = |msgs: Vec| -> Message { + Message::new(rmp_serde::to_vec(&msgs).unwrap().into()) + }; + + submit + .send(frame(vec![RemoteMessage::Init( + crate::shared::SessionInit::new(session_id, 0, vec![]), + )])) + .await + .expect("send init"); + submit + .send(frame(vec![RemoteMessage::Task(Task::Seed(0))])) + .await + .expect("send task"); + // Drop `submit` here (end of block): the server sees the stream end without a + // `Close` and must run the cleanup path. + }); + } + + // Give the server a moment to tear the abandoned session down. + std::thread::sleep(std::time::Duration::from_millis(500)); + + // The server must have survived: a fresh, normal client on the same server still works. + let finished = finishes_within(std::time::Duration::from_secs(10), || { + let device = Device::remote_websocket("ws://localhost:3090", 0); + let input = Tensor::<2>::from_floats([[10.0, 20.0]], &device); + let numbers: Vec = (input * 3.0).to_data().to_vec().unwrap(); + assert_eq!(numbers, vec![30.0, 60.0]); + }); + assert!( + finished, + "server stopped serving after a client disconnected abruptly" + ); + + rt.shutdown_timeout(std::time::Duration::from_millis(100)); + } + + #[test] + pub fn test_to_device_local_to_remote() { + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_io() + .build() + .unwrap(); + + rt.spawn( + crate::server::RemoteServerBuilder::::new(vec![Default::default()]) + .port(3020) + .start_async(), + ); + + std::thread::sleep(std::time::Duration::from_millis(500)); + + let local = Device::default(); + let remote = Device::remote_websocket("ws://localhost:3020", 0); + + // Create on local, move to remote. + let input = Tensor::<2>::from_floats([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &local); + let on_remote = input.clone().to_device(&remote); + + // Run an op while on the remote. + let doubled = on_remote * 2.0; + + // Move back to local and verify. + let back = doubled.to_device(&local); + let numbers: Vec = back.to_data().to_vec().unwrap(); + assert_eq!(numbers, vec![2.0, 4.0, 6.0, 8.0, 10.0, 12.0]); + + rt.shutdown_background(); + } +} + +#[cfg(all(test, feature = "fusion", feature = "server"))] +mod fusion_tests { + use crate::{RemoteBackend, RemoteDevice, client::RemoteChannel}; + use burn_backend::{Backend, Shape, TensorData}; + use burn_router::BackendRouter; + + // `RemoteBackend` is `Fusion` under the `fusion` feature; `PlainRemote` is the + // unwrapped router backend, used as the reference to compare against. + type PlainRemote = BackendRouter; + + fn input() -> TensorData { + TensorData::from([[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]) + } + + /// Run the same small multi-op graph `iters` times on backend `B`, reading the result each + /// iteration. Every iteration has an identical op structure, so a fusion backend registers the + /// optimization once and replays it by id on the later iterations. + /// + /// The graph deliberately includes a reshape, so one of the intermediate tensors (`c`) has a + /// *different* shape than the inputs/outputs — exercising the server's reconstruction of + /// intermediate shapes from the shape-dim map (rather than them being sent per replay). + fn run>( + device: &RemoteDevice, + iters: usize, + ) -> Vec> { + let mut out = Vec::new(); + for _ in 0..iters { + let a = B::float_from_data(input(), device); // [2, 3] + let b = B::float_exp(a); // [2, 3] + let c = B::float_reshape(b, Shape::from([3, 2])); // [3, 2] intermediate (distinct shape) + let d = B::float_log(c); // [3, 2] + let data = burn_std::reader::try_read_sync(B::float_into_data(d)) + .expect("remote read should resolve synchronously") + .expect("read should succeed"); + out.push(data.to_vec::().unwrap()); + } + out + } + + /// The fusion-enabled remote backend must produce exactly the same results as the plain remote + /// backend across a repeated computation that exercises register-once + replay-by-id. + #[test] + fn fusion_matches_plain_remote() { + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_io() + .build() + .unwrap(); + + rt.spawn( + crate::server::RemoteServerBuilder::::new(vec![Default::default()]) + .port(3100) + .start_async(), + ); + rt.spawn( + crate::server::RemoteServerBuilder::::new(vec![Default::default()]) + .port(3110) + .start_async(), + ); + std::thread::sleep(std::time::Duration::from_millis(500)); + + let plain_device = RemoteDevice::websocket("ws://localhost:3100", 0); + let fused_device = RemoteDevice::websocket("ws://localhost:3110", 0); + + let iters = 5; + let expected = run::(&plain_device, iters); + let actual = run::(&fused_device, iters); + + assert_eq!(actual.len(), iters); + assert_eq!(expected.len(), iters); + for (a, e) in actual.iter().zip(expected.iter()) { + assert_eq!(a.len(), e.len()); + for (av, ev) in a.iter().zip(e.iter()) { + assert!( + (av - ev).abs() < 1e-5, + "fusion result {av} differs from plain remote {ev}" + ); + } + } + + rt.shutdown_background(); + } + + /// A *source* custom op (no tensor inputs — it builds a tensor from scalars on the server) whose + /// output is then consumed by a follow-up op, read back, repeated to exercise register-once + + /// replay. This mirrors the server-side data-loader extension pattern and isolates it from the + /// training stack — if the fusion graph mishandles a source custom op's outputs, it surfaces here + /// as a "Should have handle for tensor ..." panic on the server. + #[test] + fn fusion_custom_source_op_then_followup() { + use crate::client::CustomOpClient; + use burn_backend::DType; + use burn_backend::ops::FloatTensorOps; + use burn_flex::Flex; + use burn_ir::{CustomOpIr, OperationOutput, ScalarIr, TensorIr}; + + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_io() + .build() + .unwrap(); + + rt.spawn( + crate::server::RemoteServerBuilder::::new(vec![Default::default()]) + .port(3120) + .custom_op("make_floats", |handles, ir, device| { + // Build a 1-D float tensor from the op's scalars — a pure source (no inputs). + let values: Vec = ir.scalars.iter().map(|s| s.elem::()).collect(); + let n = values.len(); + let tensor = Flex::float_from_data(TensorData::new(values, [n]), device); + handles.register_float_tensor::(&ir.outputs[0].id, tensor); + }) + .start_async(), + ); + std::thread::sleep(std::time::Duration::from_millis(500)); + + let device = RemoteDevice::websocket("ws://localhost:3120", 0); + + for i in 0..5 { + let client = CustomOpClient::new(&device); + let out_ir = + TensorIr::uninit(client.create_empty_handle(), Shape::from([3]), DType::F32); + let made = client + .register(CustomOpIr::with_scalars( + "make_floats", + &[], + &[out_ir], + vec![ + ScalarIr::Float(1.0), + ScalarIr::Float(2.0), + ScalarIr::Float(3.0), + ], + )) + .output(); + + // Follow-up op consuming the source output — forces a graph that references the custom + // op's output as an input, the scenario that breaks during training. + let doubled = >::float_exp(made); + let data = burn_std::reader::try_read_sync(>::float_into_data(doubled)) + .expect("remote read should resolve synchronously") + .expect("read should succeed"); + + let values = data.to_vec::().unwrap(); + let expected = [1.0f32.exp(), 2.0f32.exp(), 3.0f32.exp()]; + for (a, e) in values.iter().zip(expected.iter()) { + assert!((a - e).abs() < 1e-4, "iter {i}: {a} vs {e}"); + } + } + + rt.shutdown_background(); + } + + /// Read a source custom op's output *directly* (it is the boundary output, with no follow-up op + /// consuming it). This is what the data loader does — `batch.tokens.to_data()` — and the case the + /// other two tests don't cover (they always feed the output into another op first). + #[test] + fn fusion_read_source_output_directly() { + use crate::client::CustomOpClient; + use burn_backend::DType; + use burn_backend::ops::FloatTensorOps; + use burn_flex::Flex; + use burn_ir::{CustomOpIr, OperationOutput, ScalarIr, TensorIr}; + + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_io() + .build() + .unwrap(); + + rt.spawn( + crate::server::RemoteServerBuilder::::new(vec![Default::default()]) + .port(3140) + .custom_op("make_floats", |handles, ir, device| { + let values: Vec = ir.scalars.iter().map(|s| s.elem::()).collect(); + let n = values.len(); + let tensor = Flex::float_from_data(TensorData::new(values, [n]), device); + handles.register_float_tensor::(&ir.outputs[0].id, tensor); + }) + .start_async(), + ); + std::thread::sleep(std::time::Duration::from_millis(500)); + + let device = RemoteDevice::websocket("ws://localhost:3140", 0); + + for i in 0..5 { + let client = CustomOpClient::new(&device); + let out_ir = + TensorIr::uninit(client.create_empty_handle(), Shape::from([3]), DType::F32); + let made = client + .register(CustomOpIr::with_scalars( + "make_floats", + &[], + &[out_ir], + vec![ + ScalarIr::Float(1.0), + ScalarIr::Float(2.0), + ScalarIr::Float(3.0), + ], + )) + .output(); + + // Read the source output directly — no follow-up op. + let data = burn_std::reader::try_read_sync(>::float_into_data(made)) + .expect("remote read should resolve synchronously") + .expect("read should succeed"); + let values = data.to_vec::().unwrap(); + assert_eq!(values, vec![1.0, 2.0, 3.0], "iter {i}"); + } + + rt.shutdown_background(); + } + + /// Closer to the data-loader: a source custom op with *three* outputs that are consumed at + /// *different depths* of the following graph (one immediately, one mid-graph, one only at the + /// end — like `tokens`/`mask`/`labels`). Exercises a source op's outputs surviving across many + /// ops before being bound as inputs, under register-once + replay. + #[test] + fn fusion_custom_source_multi_output_long_lived() { + use crate::client::CustomOpClient; + use burn_backend::DType; + use burn_backend::ops::FloatTensorOps; + use burn_flex::Flex; + use burn_ir::{CustomOpIr, OperationOutput, ScalarIr, TensorIr}; + + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_io() + .build() + .unwrap(); + + rt.spawn( + crate::server::RemoteServerBuilder::::new(vec![Default::default()]) + .port(3130) + .custom_op("make3", |handles, ir, device| { + // 9 scalars → three [3] outputs (chunks of 3). + let values: Vec = ir.scalars.iter().map(|s| s.elem::()).collect(); + for (i, out) in ir.outputs.iter().enumerate() { + let chunk = values[i * 3..(i + 1) * 3].to_vec(); + let tensor = Flex::float_from_data(TensorData::new(chunk, [3]), device); + handles.register_float_tensor::(&out.id, tensor); + } + }) + .start_async(), + ); + std::thread::sleep(std::time::Duration::from_millis(500)); + + let device = RemoteDevice::websocket("ws://localhost:3130", 0); + + for i in 0..5 { + let client = CustomOpClient::new(&device); + let mk = |client: &CustomOpClient| { + TensorIr::uninit(client.create_empty_handle(), Shape::from([3]), DType::F32) + }; + let [a, b, c] = client + .register(CustomOpIr::with_scalars( + "make3", + &[], + &[mk(&client), mk(&client), mk(&client)], + (1..=9).map(|v| ScalarIr::Float(v as f64)).collect(), + )) + .outputs::<3>(); + + // a consumed immediately, b mid-graph, c only at the end — so b and c are source-op + // outputs that survive across several ops before being bound as inputs. + type B = RemoteBackend; + let t = >::float_exp(a); + let t = >::float_add(t, b); + let t = >::float_log(t); + let t = >::float_add(t, c); + let data = + burn_std::reader::try_read_sync(>::float_into_data(t)) + .expect("remote read should resolve synchronously") + .expect("read should succeed"); + + let values = data.to_vec::().unwrap(); + // a=[1,2,3], b=[4,5,6], c=[7,8,9]; t = log(exp(a)+b) + c + let expected: Vec = (0..3) + .map(|k| { + let a = (k + 1) as f32; + let b = (k + 4) as f32; + let c = (k + 7) as f32; + (a.exp() + b).ln() + c + }) + .collect(); + for (g, e) in values.iter().zip(expected.iter()) { + assert!((g - e).abs() < 1e-3, "iter {i}: {g} vs {e}"); + } + } + + rt.shutdown_background(); + } + + /// Regression guard for the `free_handle` drop-suppression override + /// (`RouterFusionRuntime::free_handle`): a *second live reference* to a source op's output is + /// held across the graph drain that consumes the first one. The override removes the drained + /// block's container entry and bumps the handle refcount so `RouterTensor::drop` doesn't + /// re-register a redundant server `Drop`; if that bookkeeping mishandles a surviving clone, the + /// retained tensor's id is freed too early and the *next* graph that uses it panics on the + /// server with "Should have handle for tensor ..." (or reads back garbage). Looping exercises + /// register-once + replay so the bug would surface on a later iteration even if the first slips + /// through. + #[test] + fn fusion_custom_source_output_clone_survives_drain() { + use crate::client::CustomOpClient; + use burn_backend::DType; + use burn_backend::ops::FloatTensorOps; + use burn_flex::Flex; + use burn_ir::{CustomOpIr, OperationOutput, ScalarIr, TensorIr}; + + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_io() + .build() + .unwrap(); + + rt.spawn( + crate::server::RemoteServerBuilder::::new(vec![Default::default()]) + .port(3150) + .custom_op("make_floats", |handles, ir, device| { + let values: Vec = ir.scalars.iter().map(|s| s.elem::()).collect(); + let n = values.len(); + let tensor = Flex::float_from_data(TensorData::new(values, [n]), device); + handles.register_float_tensor::(&ir.outputs[0].id, tensor); + }) + .start_async(), + ); + std::thread::sleep(std::time::Duration::from_millis(500)); + + let device = RemoteDevice::websocket("ws://localhost:3150", 0); + + type B = RemoteBackend; + for i in 0..5 { + let client = CustomOpClient::new(&device); + let out_ir = + TensorIr::uninit(client.create_empty_handle(), Shape::from([3]), DType::F32); + let made = client + .register(CustomOpIr::with_scalars( + "make_floats", + &[], + &[out_ir], + vec![ + ScalarIr::Float(1.0), + ScalarIr::Float(2.0), + ScalarIr::Float(3.0), + ], + )) + .output(); + + // A second live reference to the same source output. Kept alive across the drain that + // consumes `made` below — this is the scenario the override's refcount bump must not + // free. + let kept = made.clone(); + + // Consume `made` in a graph and read it back, forcing a drain that frees the block's + // handles while `kept` still references the source output's id. + let exp = >::float_exp(made); + let exp_data = + burn_std::reader::try_read_sync(>::float_into_data(exp)) + .expect("remote read should resolve synchronously") + .expect("read should succeed"); + let exp_values = exp_data.to_vec::().unwrap(); + let exp_expected = [1.0f32.exp(), 2.0f32.exp(), 3.0f32.exp()]; + for (g, e) in exp_values.iter().zip(exp_expected.iter()) { + assert!((g - e).abs() < 1e-3, "iter {i} (exp): {g} vs {e}"); + } + + // Now use the retained clone in a *new* graph. If the drain above freed the id out from + // under it, this read fails server-side or returns garbage. + let log = >::float_log(kept); + let log_data = + burn_std::reader::try_read_sync(>::float_into_data(log)) + .expect("remote read should resolve synchronously") + .expect("read should succeed"); + let log_values = log_data.to_vec::().unwrap(); + let log_expected = [1.0f32.ln(), 2.0f32.ln(), 3.0f32.ln()]; + for (g, e) in log_values.iter().zip(log_expected.iter()) { + assert!((g - e).abs() < 1e-3, "iter {i} (log): {g} vs {e}"); + } + } + + rt.shutdown_background(); + } +} diff --git a/crates/burn-remote/src/metrics.rs b/crates/burn-remote/src/metrics.rs new file mode 100644 index 0000000..75a2eaa --- /dev/null +++ b/crates/burn-remote/src/metrics.rs @@ -0,0 +1,148 @@ +//! Logs op-graph caching (fusion) savings as `[remote ...]` lines, gated behind the runtime log +//! level (`[remote]` in `burn.toml`, or `BURN_REMOTE_LOG`). A best-effort [`TelemetryProbe`] +//! subscriber; client and server each run their own, distinguished by a [`MetricSide`] label. +//! +//! [`TelemetryProbe`]: crate::telemetry::TelemetryProbe + +use core::fmt; + +use burn_ir::GraphId; +use burn_std::config::config; +use burn_std::config::log_remote; +use burn_std::config::remote::RemoteLogLevel; + +use crate::telemetry::{ + OpClass, TelemetryEvent, TelemetryProbe, TelemetrySubscription, TrafficAggregator, +}; + +const MIB: u64 = 1024 * 1024; + +/// Which side of the connection a logger reports, used to label its log lines. +#[derive(Clone, Copy, Debug)] +pub(crate) enum MetricSide { + /// The client: bytes it would have sent vs. bytes it actually sent. + Client, + /// The server: bytes it would have received vs. bytes it actually received. + Server, +} + +impl fmt::Display for MetricSide { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + MetricSide::Client => f.write_str("client"), + MetricSide::Server => f.write_str("server"), + } + } +} + +/// A best-effort telemetry subscriber that logs op-graph caching savings. +pub(crate) struct TelemetryLogger { + side: MetricSide, + aggregator: TrafficAggregator, + /// Highest whole-MiB savings already reported, so `Basic` logging emits at most one line per + /// additional mebibyte saved instead of one per replay. + logged_mib: u64, +} + +impl TelemetryLogger { + pub(crate) fn new(side: MetricSide) -> Self { + Self { + side, + aggregator: TrafficAggregator::default(), + logged_mib: 0, + } + } + + /// Whether remote logging is configured. + pub(crate) fn enabled() -> bool { + level() != RemoteLogLevel::Disabled + } + + /// Drain and log until the probe's senders are dropped. + pub(crate) async fn run(mut self, mut events: TelemetrySubscription) { + while let Some(event) = events.recv().await { + self.aggregator.apply(&event); + match event.as_ref() { + TelemetryEvent::GraphRegistered { + graph, ops, bytes, .. + } => self.log_registration(*graph, ops.len(), *bytes), + TelemetryEvent::GraphExecuted { .. } => self.log_progress(), + _ => {} + } + } + } + + fn log_registration(&self, graph: GraphId, ops: usize, bytes: usize) { + let side = self.side; + log_remote(RemoteLogLevel::Full, || { + format!( + "[remote {side}] registered graph {graph:?}: {ops} ops, {bytes} bytes (sent once)" + ) + }); + } + + fn log_progress(&mut self) { + let snapshot = self.aggregator.snapshot(); + let saved = snapshot.saved(); + let saved_pct = percentage(saved, snapshot.baseline); + let side = self.side; + + if level() >= RemoteLogLevel::Full { + let fused_pct = percentage(snapshot.fused_ops, snapshot.total_ops()); + let breakdown = format_unfused_by_kind(&self.aggregator.unfused_by_kind()); + log_remote(RemoteLogLevel::Full, || { + format!( + "[remote {side}] op-graph caching saved {saved} bytes ({saved_pct:.1}% of \ + baseline); fused {fused_pct:.1}% of {} ops total; unfused by kind: [{breakdown}]", + snapshot.total_ops() + ) + }); + } else { + // Basic: self-throttle to one summary per additional mebibyte saved. + let mib = saved / MIB; + if mib > self.logged_mib { + self.logged_mib = mib; + log_remote(RemoteLogLevel::Basic, || { + format!( + "[remote {side}] op-graph caching has saved ~{mib} MiB ({saved} bytes, \ + {saved_pct:.1}% of baseline) of network traffic" + ) + }); + } + } + } +} + +/// The logger's drain loop for `probe`, or `None` when remote logging is off. Spawn it detached. +pub(crate) fn logger_task( + probe: &TelemetryProbe, + side: MetricSide, +) -> Option + Send + 'static> { + if !TelemetryLogger::enabled() { + return None; + } + let subscription = probe.subscribe()?; + Some(TelemetryLogger::new(side).run(subscription)) +} + +fn level() -> RemoteLogLevel { + config().remote().logger.level +} + +fn format_unfused_by_kind(entries: &[(OpClass, u64)]) -> String { + entries + .iter() + .map(|(kind, count)| format!("{}={count}", kind.label())) + .collect::>() + .join(", ") +} + +/// `part` as a percentage of `whole`. Zero when `whole` is 0 (nothing measured yet), so the first +/// log line can't divide by zero. +fn percentage(part: u64, whole: u64) -> f64 { + if whole == 0 { + 0.0 + } else { + part as f64 / whole as f64 * 100.0 + } +} diff --git a/crates/burn-remote/src/server/builder.rs b/crates/burn-remote/src/server/builder.rs new file mode 100644 index 0000000..8808fb9 --- /dev/null +++ b/crates/burn-remote/src/server/builder.rs @@ -0,0 +1,163 @@ +use burn_backend::tensor::Device; +use burn_ir::{BackendIr, CustomOpIr, HandleContainer}; +use burn_router::CustomOpRegistry; + +/// Transport used to serve remote clients. +/// +/// Selects the protocol the [`RemoteServerBuilder`] serves on. +#[derive(Clone)] +pub enum Channel { + /// WebSocket server bound to `0.0.0.0:port`. + #[cfg(feature = "websocket")] + WebSocket { + /// Port to bind on. + port: u16, + }, + /// Iroh peer-to-peer transport. The server's address is `secret.id()`; clients dial that. + #[cfg(feature = "iroh")] + Iroh { + /// The server's stable identity, its address knob (like a port for WebSocket). + secret: Box, + }, +} + +/// Default port used when none is configured on the builder. +#[cfg(feature = "websocket")] +const DEFAULT_PORT: u16 = 3000; + +impl core::fmt::Debug for Channel { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + #[cfg(feature = "websocket")] + Channel::WebSocket { port } => f.debug_struct("WebSocket").field("port", port).finish(), + // Show the public identity, never the secret key material. + #[cfg(feature = "iroh")] + Channel::Iroh { secret } => f.debug_struct("Iroh").field("id", &secret.id()).finish(), + } + } +} + +impl Default for Channel { + fn default() -> Self { + #[cfg(feature = "websocket")] + return Channel::WebSocket { port: DEFAULT_PORT }; + // Without WebSocket the default is Iroh on a fresh random identity; a host that wants a + // dialable address sets its own secret with [`Channel::Iroh`]. + #[cfg(all(feature = "iroh", not(feature = "websocket")))] + return Channel::Iroh { + secret: Box::new(crate::RemoteSecret::random()), + }; + } +} + +/// Builder for a remote-execution server. +/// +/// Configures the transport ([`channel`](Self::channel) / [`port`](Self::port)) and the custom +/// operation handlers ([`custom_op`](Self::custom_op) / [`custom_ops`](Self::custom_ops)), then +/// starts the server with [`start`](Self::start) (blocking) or [`start_async`](Self::start_async). +/// +/// The builder is generic over the concrete backend `B`: custom ops are typed by `B`, since their +/// handlers call into `B`'s primitives. A backend extension hosts its ops here — the server-side +/// counterpart of the client building `OperationIr::Custom`. Custom ops are served the same way over +/// either transport. +/// +/// ```rust,ignore +/// RemoteServerBuilder::new(devices) +/// .port(3000) +/// .custom_op("fused_matmul_add_relu", |handles, ir, _device| { +/// let ([lhs, rhs, bias], [out]) = ir.as_fixed::<3, 1>(); +/// let lhs = handles.get_float_tensor::(lhs); +/// let rhs = handles.get_float_tensor::(rhs); +/// let bias = handles.get_float_tensor::(bias); +/// let result = ::fused_matmul_add_relu(lhs, rhs, bias); +/// handles.register_float_tensor::(&out.id, result); +/// }) +/// .start(); +/// ``` +pub struct RemoteServerBuilder { + devices: Vec>, + channel: Channel, + custom_ops: CustomOpRegistry, +} + +impl RemoteServerBuilder { + /// Create a new builder hosting the given devices. + /// + /// `devices` is indexed by the device index a client selects at session init; `devices[0]` is + /// the default device. Must be non-empty. Defaults to WebSocket on port `3000` (or Iroh when + /// WebSocket is not compiled in) with no custom ops. + pub fn new(devices: Vec>) -> Self { + Self { + devices, + channel: Channel::default(), + custom_ops: CustomOpRegistry::default(), + } + } + + /// Select the transport (protocol and its configuration). + pub fn channel(mut self, channel: Channel) -> Self { + self.channel = channel; + self + } + + /// Set the port, keeping the WebSocket transport. Convenience over [`channel`](Self::channel). + #[cfg(feature = "websocket")] + pub fn port(mut self, port: u16) -> Self { + self.channel = Channel::WebSocket { port }; + self + } + + /// Register a handler for a [custom operation](burn_ir::OperationIr::Custom), keyed by `id`. + /// + /// The `id` must match the one the client puts in its [`CustomOpIr`]. Chainable; registering + /// the same id twice keeps the last handler. + pub fn custom_op(mut self, id: &str, handler: F) -> Self + where + F: Fn(&mut HandleContainer, &CustomOpIr, &B::Device) + Send + Sync + 'static, + { + self.custom_ops.register(id, handler); + self + } + + /// Replace the custom-op handlers with a prebuilt [`CustomOpRegistry`]. + pub fn custom_ops(mut self, custom_ops: CustomOpRegistry) -> Self { + self.custom_ops = custom_ops; + self + } + + /// Start the server on the caller's async runtime, serving until shutdown. + #[cfg(not(target_family = "wasm"))] + pub async fn start_async(self) { + // The backend is hosted on an async runtime: tensor readbacks must materialize + // eagerly rather than deferring a blocking device→host copy onto an executor worker. + burn_std::set_runtime_kind(burn_std::RuntimeKind::Async); + + match self.channel { + #[cfg(feature = "websocket")] + Channel::WebSocket { port } => { + crate::transport::websocket::start_websocket_async::( + self.devices, + port, + self.custom_ops, + ) + .await; + } + #[cfg(feature = "iroh")] + Channel::Iroh { secret } => { + crate::transport::iroh::server::start_iroh_async::( + *secret, + self.devices, + self.custom_ops, + ) + .await; + } + } + } + + /// Start the server, blocking the current thread until shutdown. + #[cfg(not(target_family = "wasm"))] + pub fn start(self) { + let runtime = tokio::runtime::Runtime::new().expect("Failed to build the tokio runtime"); + runtime.block_on(self.start_async()); + } +} diff --git a/crates/burn-remote/src/server/local_comm.rs b/crates/burn-remote/src/server/local_comm.rs new file mode 100644 index 0000000..82c058f --- /dev/null +++ b/crates/burn-remote/src/server/local_comm.rs @@ -0,0 +1,93 @@ +//! In-process communication between two sessions hosted by the same server. +//! +//! Whenever sessions on the **same** server (same address, different device index) need to +//! hand device-resident data to each other, there is no reason to round-trip through the host: +//! both sessions' [`TensorInterpreter`](burn_router::TensorInterpreter)s live in the same +//! process, so the source can hand its primitive straight to the target, which moves it onto +//! its own device with the inner backend's `to_device`. This is the same-host counterpart of the +//! cross-server tensor transfer, which moves data *between* servers over the network. +//! +//! Today this backs cross-device `to_device`, but the mechanism is general: any same-host +//! collective (e.g. all-reduce) that needs ranks to exchange primitives uses the same +//! expose/take rendezvous. +//! +//! This registry is the rendezvous point. The source session exposes a primitive under a +//! [`LocalTransferId`]; the target session takes it. Either may arrive first — the source op +//! and the target op travel on separate connections with no cross-connection ordering — so the +//! taker waits on a [`Notify`] until the primitive shows up. + +use crate::shared::{LocalTransferId, SessionId}; +use burn_ir::{BackendIr, HandleKind}; +use std::collections::HashMap; +use std::sync::Arc; +use tokio::sync::{Mutex, Notify}; + +/// A primitive parked in the registry, tagged with the session that exposed it so the entry can +/// be reclaimed if that session goes away before its target takes it. +struct Pending { + source: SessionId, + tensor: HandleKind, +} + +/// Rendezvous registry for same-host communication, keyed by [`LocalTransferId`]. +pub(crate) struct LocalCommService { + /// Primitives exposed by source sessions, waiting to be taken by their target. + pending: Mutex>>, + /// Wakes takers whenever a new primitive is exposed. + notify: Arc, +} + +impl LocalCommService { + pub fn new() -> Self { + Self { + pending: Mutex::new(HashMap::new()), + notify: Arc::new(Notify::new()), + } + } + + /// Expose `tensor` under `transfer_id`, waking any session already waiting for it. `source` is + /// the session exposing it, so a stranded entry can be purged if that session closes before its + /// target takes the tensor (see [`purge_session`](Self::purge_session)). + pub async fn expose( + &self, + source: SessionId, + transfer_id: LocalTransferId, + tensor: HandleKind, + ) { + self.pending + .lock() + .await + .insert(transfer_id, Pending { source, tensor }); + self.notify.notify_waiters(); + } + + /// Take the primitive exposed under `transfer_id`, waiting until it is exposed. + pub async fn take(&self, transfer_id: LocalTransferId) -> HandleKind { + loop { + // Register interest *before* checking the map so an `expose` that lands between + // the check and the await can't slip through the gap and leave us parked forever. + let notified = self.notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + + if let Some(pending) = self.pending.lock().await.remove(&transfer_id) { + return pending.tensor; + } + + notified.await; + } + } + + /// Drop every primitive `session` exposed that no target ever took. + /// + /// Called when a session closes. In normal operation every exposed primitive is taken (and so + /// removed) by its target, but if the source session tears down first — a crash, a dropped + /// connection, a transfer whose target never runs — its entry would otherwise sit in the map + /// for the server's lifetime, pinning device memory. Purging by source session reclaims it. + pub async fn purge_session(&self, session: SessionId) { + self.pending + .lock() + .await + .retain(|_, pending| pending.source != session); + } +} diff --git a/crates/burn-remote/src/server/mod.rs b/crates/burn-remote/src/server/mod.rs new file mode 100644 index 0000000..1452181 --- /dev/null +++ b/crates/burn-remote/src/server/mod.rs @@ -0,0 +1,17 @@ +pub(crate) mod local_comm; +pub(crate) mod pump; +pub(crate) mod service; +pub(crate) mod session; +pub(crate) mod spawn; +pub(crate) mod transfer; +pub(crate) mod worker; + +mod builder; + +pub use builder::{Channel, RemoteServerBuilder}; +pub use burn_router::{CustomOpHandler, CustomOpRegistry}; + +#[cfg(feature = "iroh")] +pub use crate::transport::iroh::protocol::{ + AllowAll, AuthorizationRequest, IrohRemoteProtocol, PeerAuthorizer, RemoteProtocol, +}; diff --git a/crates/burn-remote/src/server/pump.rs b/crates/burn-remote/src/server/pump.rs new file mode 100644 index 0000000..4c64fae --- /dev/null +++ b/crates/burn-remote/src/server/pump.rs @@ -0,0 +1,143 @@ +//! The transport-agnostic session pump. +//! +//! One session is one duplex link. [`drive_session`] reads the init handshake, authorizes the peer, +//! binds the session, replies with the device settings, then concurrently drains task responses to +//! the sink (a detached writer) while forwarding submitted task batches from the source to the +//! session worker. This is the single implementation both transports (iroh, websocket) drive — the +//! per-transport modules only build the [`FrameSource`]/[`FrameSink`] halves and the authorizer. + +use std::sync::Arc; + +use crate::PeerId; +use crate::server::service::{SessionService, parse_init_handshake}; +use crate::server::spawn::spawn_detached; +use crate::shared::{ + PROTOCOL_VERSION, RemoteMessage, SessionInfo, SessionInit, TaskResponse, TaskResponseContent, +}; +use crate::transport::link::{FrameSink, FrameSource}; + +/// Drive one session to completion over a duplex link. +/// +/// `authorize` runs once, after the init handshake is parsed and before the session is bound — it +/// is where a transport with an authenticated peer identity (iroh) enforces its policy; transports +/// without one (websocket) pass an allow-all closure. `server_peer_id` is echoed to the client in +/// the handshake response (the server's own identity, or `None` for websocket). +/// +/// Returns `Err` on a protocol violation or a transport error; the caller logs it. A clean client +/// `Close` (or stream end) returns `Ok(())`. +pub(crate) async fn drive_session( + mut source: Src, + mut sink: Snk, + service: Arc, + server_peer_id: Option, + authorize: A, +) -> Result<(), String> +where + Src: FrameSource, + Snk: FrameSink, + S: SessionService, + A: FnOnce(&SessionInit) -> Result<(), String>, +{ + // The session stream opens with exactly one `Init` frame. + let handshake = source + .recv() + .await? + .ok_or_else(|| "Session stream closed before initialization".to_string())?; + let init = parse_init_handshake(&handshake)?; + + // Authorize before any session state is created. + authorize(&init)?; + + // Bind the session (creating it + its worker on demand) and claim its response receiver. + let task_sender = service + .session_task_sender(init.session_id, init.device_index) + .await; + let mut responses = service + .take_response_receiver(init.session_id, init.device_index) + .await?; + + // Reply with the selected device's settings + this server's identity, so the client can fill in + // `RemoteDevice::defaults`/`enumerate` without an extra round-trip. + let info = TaskResponse { + id: 0, + content: TaskResponseContent::Init(SessionInfo { + version: PROTOCOL_VERSION, + settings: service.device_settings(init.device_index), + device_count: service.device_count(), + peer_id: server_peer_id, + }), + }; + let info = rmp_serde::to_vec(&info) + .map_err(|err| format!("Failed to encode session handshake response: {err}"))?; + sink.send(info.into()).await?; + + // Detached writer: drain the session's responses onto the sink until the queue closes (every + // sender — the worker and any in-flight readback task — has dropped). + let (writer_done, writer_result) = tokio::sync::oneshot::channel(); + spawn_detached(async move { + let result = async { + while let Some(response) = responses.recv().await { + let bytes = rmp_serde::to_vec(&response) + .map_err(|err| format!("Failed to encode task response: {err}"))?; + sink.send(bytes.into()).await?; + } + sink.close().await + } + .await; + let _ = writer_done.send(result); + }); + + // Reader loop: forward each submitted task batch to the session worker in arrival order. + let result = loop { + let Some(frame) = source.recv().await? else { + break Ok(()); + }; + let messages: Vec = rmp_serde::from_slice(&frame) + .map_err(|err| format!("Invalid remote task batch: {err}"))?; + let mut close = false; + let mut protocol_error = None; + for message in messages { + match message { + RemoteMessage::Task(task) => { + task_sender + .send(task) + .await + .map_err(|_| "Session worker stopped".to_string())?; + } + RemoteMessage::Close(id) if id == init.session_id => { + close = true; + break; + } + RemoteMessage::Close(id) => { + protocol_error = Some(format!( + "Session {} attempted to close unrelated session {id}", + init.session_id + )); + break; + } + RemoteMessage::Init(_) => { + protocol_error = Some("A session stream cannot be initialized twice".into()); + break; + } + } + } + if let Some(err) = protocol_error { + break Err(err); + } + if close { + break Ok(()); + } + }; + + // Teardown: drop our task sender and close the session so its worker drains and exits, which + // closes the response queue and ends the writer; then await the writer so we don't tear the + // runtime down mid-send. + drop(task_sender); + service.close(init.session_id).await; + match writer_result.await { + Ok(Ok(())) => {} + Ok(Err(err)) => log::warn!("Session response writer failed: {err}"), + Err(_) => log::warn!("Session response writer stopped before finishing"), + } + result +} diff --git a/crates/burn-remote/src/server/service/base.rs b/crates/burn-remote/src/server/service/base.rs new file mode 100644 index 0000000..4da3b4f --- /dev/null +++ b/crates/burn-remote/src/server/service/base.rs @@ -0,0 +1,72 @@ +//! Shared helpers for the connection handlers. + +use crate::shared::{PROTOCOL_VERSION, RemoteMessage, SessionInit}; + +/// Decode the single `RemoteMessage::Init` a fresh submit (or fetch) socket opens with. +/// +/// The handshake frame must hold exactly one message and it must be an `Init`; anything else is +/// a protocol error. +pub(crate) fn parse_init_handshake(bytes: &[u8]) -> Result { + let mut messages = rmp_serde::from_slice::>(bytes) + .map_err(|err| format!("Failed to decode init handshake: {err:?}"))?; + + match messages.pop() { + Some(RemoteMessage::Init(init)) if messages.is_empty() => { + if init.version != PROTOCOL_VERSION { + return Err(format!( + "Unsupported Burn Remote protocol version {} (expected {PROTOCOL_VERSION})", + init.version + )); + } + Ok(init) + } + other => Err(format!( + "Init handshake expected a single RemoteMessage::Init, got {other:?}" + )), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::shared::SessionId; + + fn encode(messages: &[RemoteMessage]) -> Vec { + rmp_serde::to_vec(messages).unwrap() + } + + #[test] + fn handshake_accepts_a_single_init() { + let id = SessionId::new(); + let bytes = encode(&[RemoteMessage::Init(SessionInit::new(id, 3, vec![]))]); + let init = parse_init_handshake(&bytes).unwrap(); + assert_eq!(init.session_id, id); + assert_eq!(init.device_index, 3); + } + + #[test] + fn handshake_rejects_an_empty_batch() { + assert!(parse_init_handshake(&encode(&[])).is_err()); + } + + #[test] + fn handshake_rejects_more_than_one_message() { + let id = SessionId::new(); + let bytes = encode(&[ + RemoteMessage::Init(SessionInit::new(id, 0, vec![])), + RemoteMessage::Close(id), + ]); + assert!(parse_init_handshake(&bytes).is_err()); + } + + #[test] + fn handshake_rejects_a_non_init() { + let id = SessionId::new(); + assert!(parse_init_handshake(&encode(&[RemoteMessage::Close(id)])).is_err()); + } + + #[test] + fn handshake_rejects_garbage() { + assert!(parse_init_handshake(&[0xff, 0x00, 0x13]).is_err()); + } +} diff --git a/crates/burn-remote/src/server/service/mod.rs b/crates/burn-remote/src/server/service/mod.rs new file mode 100644 index 0000000..91b1017 --- /dev/null +++ b/crates/burn-remote/src/server/service/mod.rs @@ -0,0 +1,52 @@ +//! The session-layer service trait, plus the init-handshake helper shared by the session pump. +//! +//! The pump ([`drive_session`](crate::server::pump::drive_session)) is written against +//! [`SessionService`] rather than a concrete type, so the submit-loop teardown logic stays testable +//! against a fake service with no backend and no live socket. + +mod base; + +pub(crate) use base::parse_init_handshake; + +use std::future::Future; + +use burn_std::DeviceSettings; +use tokio::sync::mpsc; + +use crate::shared::{SessionId, Task, TaskResponse}; + +/// What the session pump needs from the session layer: bind a session to its worker channel, claim +/// its response receiver, read the device metadata for the handshake, and tear a session down. +/// +/// Async methods return `impl Future + Send` so a session future built on them stays `Send` and can +/// be spawned by the server. The production implementation is +/// [`SessionManager`](super::session::SessionManager). +pub(crate) trait SessionService: Send + Sync + 'static { + /// The channel forwarding tasks to `session_id`'s worker, creating the session (and spawning + /// its worker) on demand. + fn session_task_sender( + &self, + session_id: SessionId, + device_index: u32, + ) -> impl Future> + Send; + + /// Claim the session's result receiver. Errors if a receiver was already taken — the protocol + /// allows only one session stream per session. + fn take_response_receiver( + &self, + session_id: SessionId, + device_index: u32, + ) -> impl Future, String>> + Send; + + /// The default settings of the device at `device_index`, returned on the handshake so the + /// client can resolve op dtypes without an extra round-trip. + fn device_settings(&self, device_index: u32) -> DeviceSettings; + + /// The number of devices this server hosts, returned on the handshake so the client can + /// enumerate every device behind the address. + fn device_count(&self) -> u32; + + /// Drop the session, letting its worker drain and exit. A `close` for an unknown session is a + /// no-op. + fn close(&self, session_id: SessionId) -> impl Future + Send; +} diff --git a/crates/burn-remote/src/server/session.rs b/crates/burn-remote/src/server/session.rs new file mode 100644 index 0000000..9c8311d --- /dev/null +++ b/crates/burn-remote/src/server/session.rs @@ -0,0 +1,222 @@ +use burn_backend::tensor::Device; +use burn_ir::BackendIr; +use burn_router::{CustomOpRegistry, TensorInterpreter}; +use std::{ + collections::HashMap, + sync::{Arc, Once}, +}; +use tokio::sync::{Mutex, mpsc}; + +use crate::metrics::{MetricSide, logger_task}; +use crate::server::local_comm::LocalCommService; +use crate::server::service::SessionService; +use crate::server::spawn::spawn_detached; +use crate::server::transfer::TensorTransfer; +use crate::server::worker::SessionHandler; +use crate::shared::{SessionId, Task, TaskResponse}; +use crate::telemetry::{TelemetryEvent, TelemetryProbe}; + +/// Capacity for the per-session response queue. +/// +/// Sized larger than the typical in-flight read/sync count so that request processing +/// doesn't block on backpressure during a burst, but small enough that a stuck response +/// writer surfaces as a backpressure stall rather than memory growth. +const RESPONSE_CHANNEL_CAPACITY: usize = 64; + +/// Coordinates per-session state. +/// +/// Each [`Session`] owns a dedicated [`SessionHandler`] that holds the session's +/// [`TensorInterpreter`] with its own [`HandleContainer`](burn_ir::HandleContainer) — different +/// sessions never share tensor handles, so concurrent sessions can't race on each other's backend +/// state. Cross-session tensor transfers go through `external_comm` (cross-server) or `local_comm` +/// (same-host), each of which has its own rendezvous. +/// +/// Tasks run on the handler's worker threads, not the submit handler: the latter only decodes the +/// incoming batch and forwards each [`Task`] to the session over a bounded channel. Inside the +/// session a dispatcher routes each task to a per-stream worker thread, so per-stream ordering is +/// preserved while independent streams — and other sessions — keep making progress even when one +/// stream is parked on a blocking op (a same-host transfer rendezvous or an all-reduce barrier). +pub struct SessionManager +where + B: BackendIr, + T: TensorTransfer, +{ + /// All devices this server hosts, indexed by the device index the client selects at + /// session init. `devices[0]` is the default device (`DeviceIndex::Default`). + devices: Vec>, + pub(crate) transfer: Arc, + /// Rendezvous registry for same-host tensor transfers between this server's sessions. + pub(crate) local_comm: Arc>, + /// Custom-op handlers shared (read-only) with every session's interpreter. + custom_ops: CustomOpRegistry, + sessions: Mutex>, + probe: TelemetryProbe, + /// Spawns the telemetry logger once, on the first session. + logger: Once, +} + +struct Session { + /// Inbound channel to the session's dispatcher thread; cloned once per submit connection. + task_sender: mpsc::Sender, + receiver: Option>, +} + +impl SessionManager +where + B: BackendIr, + T: TensorTransfer, +{ + pub fn new(devices: Vec>, transfer: Arc) -> Self { + assert!( + !devices.is_empty(), + "A remote server must host at least one device" + ); + Self { + devices, + transfer, + local_comm: Arc::new(LocalCommService::new()), + custom_ops: CustomOpRegistry::default(), + sessions: Mutex::new(HashMap::new()), + probe: TelemetryProbe::disabled(), + logger: Once::new(), + } + } + + /// Register custom-op handlers, shared read-only with every session's interpreter. + pub fn with_custom_ops(mut self, custom_ops: CustomOpRegistry) -> Self { + self.custom_ops = custom_ops; + self + } + + /// Emit telemetry into `probe`. + pub fn with_telemetry(mut self, probe: TelemetryProbe) -> Self { + self.probe = probe; + self + } + + fn ensure_logger(&self) { + self.logger.call_once(|| { + if let Some(task) = logger_task(&self.probe, MetricSide::Server) { + spawn_detached(task); + } + }); + } + + /// Resolve the device at `device_index`. + /// + /// The index is validated against the server's device count on the client init handshake, so + /// an out-of-range index here is a protocol/configuration error (e.g. a client enumerating + /// more devices than this server hosts). Fail loudly rather than silently collapsing onto + /// device 0 — for a collective that would reduce a device against itself and silently corrupt + /// the result instead of producing a clear failure. + pub(crate) fn device(&self, device_index: u32) -> Device { + self.devices + .get(device_index as usize) + .cloned() + .unwrap_or_else(|| { + panic!( + "Requested device index {device_index} but server hosts only {} device(s)", + self.devices.len() + ) + }) + } + + async fn with_session( + &self, + session_id: SessionId, + device_index: u32, + f: impl FnOnce(&mut Session) -> R, + ) -> R { + self.ensure_logger(); + let mut sessions = self.sessions.lock().await; + let entry = sessions.entry(session_id).or_insert_with(|| { + let (sender, receiver) = mpsc::channel(RESPONSE_CHANNEL_CAPACITY); + // The session is pinned to its device for its whole lifetime; the first handler + // (request or response) to touch it fixes the device index. Spawn the handler that + // owns the runner — this runs inside the tokio runtime, so the handler can capture + // the runtime handle its worker threads need for the async parts of a task. + let runner = TensorInterpreter::with_custom_ops( + self.device(device_index), + self.custom_ops.clone(), + ); + let task_sender = SessionHandler::spawn( + session_id, + runner, + sender, + self.transfer.clone(), + self.local_comm.clone(), + self.probe.clone(), + ); + self.probe.emit(|| TelemetryEvent::SessionOpened { + session: session_id, + device: device_index, + }); + Session { + task_sender, + receiver: Some(receiver), + } + }); + f(entry) + } +} + +impl SessionService for SessionManager +where + B: BackendIr, + T: TensorTransfer, +{ + /// Resolve the channel used to forward [`Task`]s to `session_id`'s dispatcher thread, + /// creating the session (and spawning its handler) on demand. The pump resolves this once and + /// reuses it for every task, instead of re-locking the sessions map per task. + async fn session_task_sender( + &self, + session_id: SessionId, + device_index: u32, + ) -> mpsc::Sender { + self.with_session(session_id, device_index, |s| s.task_sender.clone()) + .await + } + + /// Take the response receiver for `session_id`. + /// + /// Returns `Err` if a responder has already been registered for this session — the protocol + /// allows only one session stream per session. + async fn take_response_receiver( + &self, + session_id: SessionId, + device_index: u32, + ) -> Result, String> { + self.with_session(session_id, device_index, |s| { + s.receiver + .take() + .ok_or_else(|| format!("Response receiver already taken for session {session_id}")) + }) + .await + } + + /// The device settings for `device_index`, used by the handshake before any session-specific + /// runner is needed. + fn device_settings(&self, device_index: u32) -> burn_std::DeviceSettings { + use burn_backend::backend::DeviceOps; + self.device(device_index).defaults() + } + + /// The total number of devices this server hosts. Sent to the client on the init handshake so + /// it can enumerate every device behind the address (see [`RemoteDevice::enumerate`]). + fn device_count(&self) -> u32 { + self.devices.len() as u32 + } + + /// Drop the session, detaching its handler. Removing the map entry drops the inbound task + /// sender it holds; once the pump also drops its cloned task sender, the dispatcher's channel + /// closes, so the per-stream workers drain and exit and the handler flushes its runner, + /// releasing any backend state held only by this session's tensors. + async fn close(&self, session_id: SessionId) { + let mut sessions = self.sessions.lock().await; + if sessions.remove(&session_id).is_some() { + self.probe.emit(|| TelemetryEvent::SessionClosed { + session: session_id, + }); + } + } +} diff --git a/crates/burn-remote/src/server/spawn.rs b/crates/burn-remote/src/server/spawn.rs new file mode 100644 index 0000000..c819d0f --- /dev/null +++ b/crates/burn-remote/src/server/spawn.rs @@ -0,0 +1,47 @@ +use core::future::Future; + +#[cfg(not(target_family = "wasm"))] +pub(crate) fn spawn_detached(future: F) +where + F: Future + Send + 'static, +{ + tokio::spawn(future); +} + +#[cfg(target_family = "wasm")] +pub(crate) fn spawn_detached(future: F) +where + F: Future + 'static, +{ + wasm_bindgen_futures::spawn_local(future); +} + +/// Resolve when the process is asked to stop (Ctrl+C, or `SIGTERM` on Unix). +/// +/// The single shutdown trigger shared by the turnkey WebSocket and Iroh server entry points. +#[cfg(all( + not(target_family = "wasm"), + any(feature = "websocket", feature = "iroh") +))] +pub(crate) async fn os_shutdown_signal() { + let ctrl_c = async { + tokio::signal::ctrl_c() + .await + .expect("failed to install Ctrl+C handler"); + }; + + #[cfg(unix)] + let terminate = async { + tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .expect("failed to install signal handler") + .recv() + .await; + }; + #[cfg(not(unix))] + let terminate = std::future::pending::<()>(); + + tokio::select! { + _ = ctrl_c => {}, + _ = terminate => {}, + } +} diff --git a/crates/burn-remote/src/server/transfer.rs b/crates/burn-remote/src/server/transfer.rs new file mode 100644 index 0000000..8671c7c --- /dev/null +++ b/crates/burn-remote/src/server/transfer.rs @@ -0,0 +1,40 @@ +//! The tensor-transfer seam between the session worker and a transport. +//! +//! Server-to-server tensor movement is independent from the compute-session transport: a tensor is +//! *exposed* on the source server and *downloaded* by the target server without routing the data +//! through the controlling client. Each transport implements this its own way — +//! [`IrohTransfer`](crate::transport::iroh::IrohTransfer) over authenticated Iroh streams, +//! [`WebSocketTransfer`](crate::transport::websocket::WebSocketTransfer) over the legacy data +//! service — and the session worker drives it through this trait, knowing nothing about the wire. + +use std::future::Future; + +use burn_backend::TensorData; +use burn_ir::BackendIr; + +use crate::shared::TransferCapability; +use crate::{PeerAddr, PeerId}; + +/// Server-to-server tensor movement, independent from the compute-session transport. +pub(crate) trait TensorTransfer: Send + Sync + 'static { + fn expose_data( + &self, + data: TensorData, + max_downloads: u32, + capability: TransferCapability, + target: PeerId, + ) -> impl Future + Send; + + fn download_tensor( + &self, + remote: PeerAddr, + capability: TransferCapability, + ) -> impl Future> + Send; + + fn fail( + &self, + capability: TransferCapability, + target: PeerId, + reason: String, + ) -> impl Future + Send; +} diff --git a/crates/burn-remote/src/server/worker.rs b/crates/burn-remote/src/server/worker.rs new file mode 100644 index 0000000..6479e4d --- /dev/null +++ b/crates/burn-remote/src/server/worker.rs @@ -0,0 +1,457 @@ +//! Per-session handler and its worker thread. +//! +//! A [`SessionHandler`] owns everything that is constant for the lifetime of a session — the +//! session's [`TensorInterpreter`] (with its own [`HandleContainer`](burn_ir::HandleContainer)), +//! the cross-server / same-host comm services, the result sender, the runtime handle and the +//! session id — and exposes a single [`process_task`](SessionHandler::process_task) method that +//! runs one task against that state. +//! +//! The submit handler does no work itself: it decodes the incoming message batch and forwards +//! each [`Task`] to the session's worker over a bounded channel. The worker drives the session's +//! tasks in **global submission order** (a single FIFO), preserving every ordering the protocol +//! relies on — including *cross-stream* ones. A tensor produced on one client stream (say a +//! dataloader thread) and consumed on another (the main thread) is only safe because the producer +//! task is applied before the consumer task; the interpreter looks up input handles eagerly and +//! panics if one is missing, so the order the client submitted in must be preserved verbatim. +//! Per-stream parallelism would break exactly this, so the stream id on a task is used only to set +//! the backend's thread-local stream (via [`StreamId::executes`]), not to reorder work. +//! +//! ## Why a dedicated OS thread, not a tokio task +//! +//! Some tasks are **synchronously blocking** — a same-host transfer's +//! [`RegisterTensorLocal`](crate::shared::Task::RegisterTensorLocal) waits on the rendezvous +//! (`local_comm.take`) for its source session to expose the primitive, and a collective op +//! (all-reduce / sync-collective) parks until *every* participating device reaches the barrier. +//! Running those on a shared tokio worker would tie up a runtime thread; once more devices are +//! blocked than the runtime has workers, the remaining devices can never be scheduled to reach the +//! barrier and it deadlocks. Giving each session its own OS thread keeps that blocking off the +//! shared runtime, so a barrier (or rendezvous) on one session can't stall another session's +//! worker or a runtime thread. +//! +//! In the browser the worker instead runs on the JS event loop via `spawn_local`, so a browser peer +//! cannot host co-located collective / same-host-transfer participants (the single thread would +//! stall on the blocking tasks above). Ordinary compute is unaffected. + +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +use burn_ir::{BackendIr, GraphBindings, GraphId}; +use burn_router::{Graph, TensorInterpreter}; +use burn_std::id::StreamId; +#[cfg(not(target_family = "wasm"))] +use tokio::runtime::Handle; +use tokio::sync::mpsc; + +use crate::server::local_comm::LocalCommService; +use crate::server::spawn::spawn_detached; +use crate::server::transfer::TensorTransfer; +use crate::shared::{RequestId, SessionId, Task, TaskResponse, TaskResponseContent}; +use crate::telemetry::{ + TelemetryEvent, TelemetryProbe, TransferPhase, TransferScope, serialized_len, +}; +use burn_ir::OperationIr; + +/// Capacity of the per-session task channel feeding the worker thread. +/// +/// The submit handler forwards with an `await`ing send, so a full channel applies async +/// backpressure (the submit task yields and stops reading the socket) rather than blocking an OS +/// thread or growing memory without bound. Sized so a burst of fire-and-forget ops doesn't stall +/// the submit loop while the worker is mid-task. +const TASK_CHANNEL_CAPACITY: usize = 64; + +/// Everything constant for the lifetime of a session. +/// +/// Owned by the session's worker thread, which runs every task against this one interpreter and +/// these comm services. Held behind an [`Arc`] only so detached readback tasks can be spawned with +/// a clone of the result sender; the interpreter itself is single-owner here. +pub(crate) struct SessionHandler +where + B: BackendIr, + T: TensorTransfer, +{ + session_id: SessionId, + runner: TensorInterpreter, + response_sender: mpsc::Sender, + transfer: Arc, + local_comm: Arc>, + /// Cache of client-registered reusable op-graphs, keyed by id (see + /// [`Task::RegisterAndExecuteGraph`]). + /// + /// `Mutex` only for interior mutability under `process_task(&self)`. Each graph is held behind + /// an [`Arc`] so a replay clones the handle out under a short lock and runs `replay_graph` + /// *after* releasing it — the lock guards only the map, never the backend dispatch. That keeps + /// the cache from becoming a serialization point if graph execution is ever driven from more + /// than the current single-FIFO worker. + graphs: Mutex>, + probe: TelemetryProbe, +} + +impl SessionHandler +where + B: BackendIr, + T: TensorTransfer, +{ + /// Start the session worker; the worker runs until every clone of the returned sender is + /// dropped, then flushes and exits. + pub(crate) fn spawn( + session_id: SessionId, + runner: TensorInterpreter, + response_sender: mpsc::Sender, + transfer: Arc, + local_comm: Arc>, + probe: TelemetryProbe, + ) -> mpsc::Sender { + let handler = SessionHandler { + session_id, + runner, + response_sender, + transfer, + local_comm, + graphs: Mutex::new(HashMap::new()), + probe, + }; + let (sender, receiver) = mpsc::channel(TASK_CHANNEL_CAPACITY); + handler.drive(receiver); + sender + } + + // Must be called from within the runtime that owns the endpoint: it captures `Handle::current`. + #[cfg(not(target_family = "wasm"))] + fn drive(self, receiver: mpsc::Receiver) { + let handle = Handle::current(); + let session_id = self.session_id; + std::thread::Builder::new() + .name(format!("burn-remote-session-{session_id}")) + .spawn(move || handle.block_on(self.run(receiver))) + .expect("Failed to spawn session worker thread"); + } + + #[cfg(target_family = "wasm")] + fn drive(self, receiver: mpsc::Receiver) { + spawn_detached(self.run(receiver)); + } + + /// Drain the task channel, running each task to completion in arrival order, then tear the + /// session down in an order that releases its memory. + async fn run(mut self, mut receiver: mpsc::Receiver) { + let session_id = self.session_id; + + log::debug!("Session {session_id} worker started"); + while let Some(task) = receiver.recv().await { + if let Err(err) = self.process_task(task).await { + // One task failing doesn't tear down the session: read/sync/dtype failures surface + // to the client through their response, fire-and-forget failures are logged here, + // and the worker keeps processing subsequent tasks. + log::error!("Task on session {session_id} failed: {err}"); + } + } + + // Reclaim any same-host transfers this session exposed that no target ever took, so a + // half-finished transfer doesn't strand device memory in the shared registry. + self.local_comm.purge_session(session_id).await; + + // The task channel closed: every submit connection bound to this session has gone away + // (clean `Close` or disconnect). + log::debug!("Session {session_id} worker draining and exiting"); + + // Destructure so we control drop order explicitly. The `..` fields (response sender, comm + // services) drop here: dropping the response sender closes the fetch writer's queue, ending + // its task. + let SessionHandler { runner, .. } = self; + let device = runner.device(); + + // Flush outstanding backend work before dropping the runner so the session's tensors aren't + // freed with GPU work still queued against them. + if let Err(err) = runner.sync() { + log::warn!("runner.sync() at session {session_id} close failed: {err:?}"); + } + if let Err(err) = B::sync(&device) { + log::warn!("B::sync(device) at session {session_id} close failed: {err:?}"); + } + + // Drop the runner: this frees every tensor handle the session held back to the backend's + // allocator. Must happen before `memory_cleanup`, otherwise the memory is still live and + // nothing is reclaimable. + drop(runner); + + // Hand the freed memory back instead of leaving it parked in the allocator's pool for a + // session that no longer exists. A no-op on backends that don't pool (the `Backend` + // default), but on cubecl backends (wgpu/cuda) this returns the session's device memory so + // a long-lived server doesn't accumulate it across session churn. + B::memory_cleanup(&device); + } + + /// Execute a single [`Task`] against this session's state. + /// + /// Sync work is wrapped in [`StreamId::executes`](burn_std::id::StreamId::executes) so the + /// runner's thread-local stream id matches the one the client assigned to this op. + /// Response-producing tasks carry their own [`RequestId`] for routing the response back to the + /// right pending callback on the client. Async work (data-service transfers, + /// `read_tensor_async`) runs without a stream context — the relevant stream id is captured into + /// the future at construction time via `executes`. + async fn process_task(&mut self, task: Task) -> Result<(), String> { + match task { + Task::RegisterOperation(stream_id, op) => { + // An op received individually (not as part of a cached graph) is an unfused op. + self.emit_op(stream_id, &op); + stream_id.executes(|| self.runner.register_op(op)); + Ok(()) + } + Task::RegisterAndExecuteGraph { + stream_id, + graph_id, + relative_graph, + bindings, + } => { + self.probe.emit(|| { + TelemetryEvent::graph_registered( + self.session_id, + graph_id, + &relative_graph, + serialized_len(&relative_graph), + ) + }); + self.emit_graph_executed(graph_id, stream_id, &bindings); + stream_id.executes(|| { + let graph = Graph::new(relative_graph); + self.graphs.lock().unwrap().insert(graph_id, graph.clone()); + graph.replay(&mut self.runner, bindings); + }); + Ok(()) + } + Task::ExecuteGraph { + stream_id, + graph_id, + bindings, + } => { + let graph = { + let cache = self.graphs.lock().unwrap(); + cache + .get(&graph_id) + .cloned() + .ok_or_else(|| format!("Execute of unknown graph {graph_id:?}"))? + }; + self.emit_graph_executed(graph_id, stream_id, &bindings); + stream_id.executes(|| graph.replay(&mut self.runner, bindings)); + Ok(()) + } + Task::RegisterTensor(stream_id, id, data) => { + stream_id.executes(|| self.runner.register_tensor_data_id(id, data)); + Ok(()) + } + Task::RegisterAlias { + stream_id, + new_id, + src_id, + } => { + stream_id.executes(|| self.runner.register_alias(new_id, src_id)); + Ok(()) + } + Task::RegisterTensorRemote(stream_id, remote, new_id) => { + log::trace!( + "Registering remote tensor (transfer {:?} from {:?})", + remote.capability, + remote.peer, + ); + let peer = Some(format!("{:?}", remote.peer)); + self.emit_transfer(peer.clone(), TransferScope::Remote, TransferPhase::Started); + let data = match self + .transfer + .download_tensor(remote.peer.clone(), remote.capability) + .await + { + Some(data) => { + self.emit_transfer(peer, TransferScope::Remote, TransferPhase::Completed); + data + } + None => { + self.emit_transfer(peer, TransferScope::Remote, TransferPhase::Failed); + return Err(format!( + "Failed to download tensor for transfer {:?} from {:?}", + remote.capability, remote.peer, + )); + } + }; + // Register on the client stream that will consume `new_id`, carried over the + // wire — not the arbitrary tokio worker running this task. + stream_id.executes(|| self.runner.register_tensor_data_id(new_id, data)); + Ok(()) + } + Task::ExposeTensorLocal { + stream_id, + tensor, + transfer_id, + } => { + // Source side of a same-host transfer. Grab the device-resident primitive + // (no host readback) and park it in the registry for the target session to + // pick up. Runs in order on this session's worker, so it is ordered after the op + // that produced `tensor` — the handle is guaranteed present. Read it back on the + // client stream that produced it, carried over the wire. + let kind = stream_id.executes(|| self.runner.get_tensor(&tensor)); + self.local_comm + .expose(self.session_id, transfer_id, kind) + .await; + Ok(()) + } + Task::RegisterTensorLocal { + stream_id, + transfer_id, + new_id, + } => { + // Target side of a same-host transfer. Wait for the source to expose the + // primitive, then move it onto this session's device and register it. Awaited + // in order so subsequent ops on this session that consume `new_id` see it + // registered first — same ordering contract as `RegisterTensorRemote`. The wait + // blocks only this session's worker, not the source session's. + let kind = self.local_comm.take(transfer_id).await; + stream_id.executes(|| self.runner.register_tensor_to_device(new_id, kind)); + self.emit_transfer(None, TransferScope::Local, TransferPhase::Completed); + Ok(()) + } + Task::ExposeTensorRemote { + stream_id, + tensor, + count, + capability, + target, + } => { + log::trace!("Exposing tensor (transfer {capability:?})"); + // Same shape as `ReadTensor`: the sync part of `read_tensor_async` runs in order + // to preserve stream ordering, but the readback + expose are detached so a + // cross-server hand-off doesn't stall this session's op registration on a + // GPU→host copy. A target that downloads before the expose lands simply blocks on + // the data service's `new_tensor_notify`, so there is no race. + let fut = stream_id.executes(|| self.runner.read_tensor_async(tensor)); + let transfer = self.transfer.clone(); + let probe = self.probe.clone(); + let session_id = self.session_id; + let peer = Some(format!("{target:?}")); + probe.emit(|| TelemetryEvent::Transfer { + session: session_id, + peer: peer.clone(), + scope: TransferScope::Remote, + phase: TransferPhase::Started, + }); + spawn_detached(async move { + match fut.await { + Ok(data) => { + transfer.expose_data(data, count, capability, target).await; + probe.emit(|| TelemetryEvent::Transfer { + session: session_id, + peer, + scope: TransferScope::Remote, + phase: TransferPhase::Completed, + }); + } + Err(e) => { + let reason = format!( + "read_tensor_async for transfer {capability:?} failed: {e:?}" + ); + log::error!( + "read_tensor_async for transfer {capability:?} failed: {e:?}" + ); + transfer.fail(capability, target, reason).await; + probe.emit(|| TelemetryEvent::Transfer { + session: session_id, + peer, + scope: TransferScope::Remote, + phase: TransferPhase::Failed, + }); + } + } + }); + Ok(()) + } + Task::Seed(seed) => { + self.runner.seed(seed); + Ok(()) + } + Task::ReadTensor(request_id, stream_id, tensor) => { + // `read_tensor_async` is sync at construction — it locks the context and + // captures the tensor's position in the command stream — and returns a future + // for the actual host readback. Run the sync part in order (so ordering vs. later + // ops is preserved), then detach the readback await onto its own task. Awaiting it + // here would stall the worker on the GPU→host copy and stop us registering + // subsequent ops, draining the device queue into a bubble. The client demuxes + // responses by request id, so out-of-order completion is fine. + self.probe.emit(|| TelemetryEvent::Read { + session: self.session_id, + request: request_id, + }); + let fut = stream_id.executes(|| self.runner.read_tensor_async(tensor)); + let sender = self.response_sender.clone(); + spawn_detached(async move { + // Under an async runtime the backend reads eagerly (see `RuntimeKind::Async`), + // so `data` is already host-resident here — the fetch handler only serializes + // resident bytes and no blocking device→host copy runs on the shared runtime. + let data = fut.await; + if sender + .send(TaskResponse { + content: TaskResponseContent::ReadTensor(data), + id: request_id, + }) + .await + .is_err() + { + log::warn!( + "Response receiver dropped before read for request {request_id} could be sent" + ); + } + }); + Ok(()) + } + Task::SyncBackend(request_id, stream_id) => { + self.probe.emit(|| TelemetryEvent::Sync { + session: self.session_id, + request: request_id, + }); + let res = stream_id.executes(|| self.runner.sync()); + self.send_response(request_id, TaskResponseContent::SyncBackend(res)) + .await + } + Task::DTypeUsage(request_id, dtype) => { + let res = self.runner.dtype_usage(dtype); + self.send_response(request_id, TaskResponseContent::DTypeUsage(res)) + .await + } + } + } + + async fn send_response( + &self, + request_id: RequestId, + content: TaskResponseContent, + ) -> Result<(), String> { + self.response_sender + .send(TaskResponse { + content, + id: request_id, + }) + .await + .map_err(|_| { + format!( + "Response receiver dropped before result for request {request_id} could be sent" + ) + }) + } + + fn emit_op(&self, stream: StreamId, op: &OperationIr) { + self.probe + .emit(|| TelemetryEvent::unfused_op(self.session_id, stream, op)); + } + + fn emit_transfer(&self, peer: Option, scope: TransferScope, phase: TransferPhase) { + self.probe.emit(|| TelemetryEvent::Transfer { + session: self.session_id, + peer, + scope, + phase, + }); + } + + fn emit_graph_executed(&self, graph: GraphId, stream: StreamId, bindings: &GraphBindings) { + self.probe.emit(|| { + TelemetryEvent::graph_executed(self.session_id, graph, stream, serialized_len(bindings)) + }); + } +} diff --git a/crates/burn-remote/src/shared/mod.rs b/crates/burn-remote/src/shared/mod.rs new file mode 100644 index 0000000..6e4ad92 --- /dev/null +++ b/crates/burn-remote/src/shared/mod.rs @@ -0,0 +1,4 @@ +mod task; + +#[allow(unused_imports)] +pub(crate) use task::*; diff --git a/crates/burn-remote/src/shared/task.rs b/crates/burn-remote/src/shared/task.rs new file mode 100644 index 0000000..4b36ac6 --- /dev/null +++ b/crates/burn-remote/src/shared/task.rs @@ -0,0 +1,236 @@ +use burn_backend::{DTypeUsageSet, ExecutionError, TensorData}; +use burn_ir::{GraphBindings, GraphId, OperationIr, TensorId, TensorIr}; +use burn_std::{ + DType, DeviceSettings, + id::{IdGenerator, StreamId}, +}; +use serde::{Deserialize, Serialize}; +use std::fmt::Display; + +use crate::{PeerAddr, PeerId}; + +/// Current Burn Remote application-protocol version. +pub const PROTOCOL_VERSION: u16 = 1; + +/// Routing id for a task whose result is fetched back. +/// +/// Only the result-producing tasks ([`Task::ReadTensor`], [`Task::SyncBackend`], +/// [`Task::DTypeUsage`]) carry a `RequestId`; the server echoes it on its [`TaskResponse`] so +/// the client demultiplexes results back to the right pending callback. Fire-and-forget tasks +/// have no id because no result ever comes back. Collective ops (all-reduce, sync-collective) +/// are plain fire-and-forget [`OperationIr`]s carried by [`Task::RegisterOperation`]. +pub type RequestId = u64; + +/// Process-local rendezvous id for moving tensors between devices on one compute peer. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct LocalTransferId(u64); + +impl LocalTransferId { + #[cfg(feature = "client")] + pub fn next(&mut self) { + self.0 += 1; + } +} + +impl From for LocalTransferId { + fn from(value: u64) -> Self { + Self(value) + } +} + +/// Unforgeable, bounded-use capability authorizing a peer-to-peer tensor download. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct TransferCapability([u8; 32]); + +impl TransferCapability { + #[cfg(feature = "client")] + pub fn random() -> Self { + Self(rand::random()) + } + + /// Deterministic compatibility key for the legacy 64-bit WebSocket transfer service. + #[cfg(feature = "websocket")] + pub(crate) fn legacy_id(self) -> u64 { + u64::from_le_bytes( + self.0[..8] + .try_into() + .expect("slice has exactly eight bytes"), + ) + } +} + +/// Unique identifier that can represent a session. +#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash, Serialize, Deserialize, PartialOrd, Ord)] +pub struct SessionId { + id: u64, +} + +impl Display for SessionId { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "SessionId({})", self.id) + } +} + +impl SessionId { + /// Create a new [session id](SessionId). + #[allow(dead_code)] + pub fn new() -> Self { + Self { + id: IdGenerator::generate(), + } + } + + /// The underlying numeric identifier. + pub fn value(&self) -> u64 { + self.id + } +} + +/// A single message on a session's stream: either a session-lifecycle signal (`Init`/`Close`) or a +/// [`Task`] to run. +#[allow(missing_docs, clippy::large_enum_variant)] +#[derive(Serialize, Deserialize, Debug)] +pub enum RemoteMessage { + /// A unit of work to run within the bound session. + Task(Task), + /// Open a session bound to the device at the given index on the server. + Init(SessionInit), + Close(SessionId), +} + +/// Client-side session handshake. +#[allow(missing_docs)] +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct SessionInit { + pub version: u16, + pub session_id: SessionId, + pub device_index: u32, + /// Opaque application credential interpreted by the compute node's authorizer. + #[serde(with = "serde_bytes")] + pub authorization: Vec, +} + +impl SessionInit { + #[cfg(any(feature = "client", test))] + pub fn new(session_id: SessionId, device_index: u32, authorization: Vec) -> Self { + Self { + version: PROTOCOL_VERSION, + session_id, + device_index, + authorization, + } + } +} + +/// Server-side session handshake response. +#[allow(missing_docs)] +#[derive(Serialize, Deserialize, Debug)] +pub struct SessionInfo { + pub version: u16, + pub settings: DeviceSettings, + pub device_count: u32, + /// Authenticated identity of the compute node, when the transport provides one. + pub peer_id: Option, +} + +#[allow(missing_docs)] +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct TensorRemote { + pub capability: TransferCapability, + pub peer: PeerAddr, +} + +#[allow(missing_docs, clippy::large_enum_variant)] +#[derive(Serialize, Deserialize, Debug)] +pub enum Task { + Seed(u64), + /// A single [`OperationIr`] tagged with the stream it was issued on. + /// + /// Each op is sent independently — batching across streams would lose stream + /// identity, and the server-side backend (fusion, etc.) handles any batching of + /// its own. + RegisterOperation(StreamId, OperationIr), + /// Register a reusable group of operations under `graph_id` (relative form) *and* immediately + /// execute it with `bindings`. + /// + /// Registering a new graph always coincides with its first execution, so the two are combined + /// into a single task to avoid a redundant round-trip on a cache miss. The server caches the + /// graph (per session) and replays it on every subsequent [`ExecuteGraph`](Task::ExecuteGraph). + /// This is how the fusion-enabled client avoids re-sending a recurring op sequence (e.g. a + /// model block) on every step. + RegisterAndExecuteGraph { + stream_id: StreamId, + graph_id: GraphId, + relative_graph: Vec, + bindings: GraphBindings, + }, + /// Replay an already-registered graph with the given concrete bindings. + ExecuteGraph { + stream_id: StreamId, + graph_id: GraphId, + bindings: GraphBindings, + }, + RegisterTensor(StreamId, TensorId, TensorData), + /// Register `new_id` as an alias of `src_id` — a second server handle over the same buffer. + /// + /// Emitted by the fusion layer's cross-stream sharing so a tensor used on multiple streams + /// gets one server id per stream view, each freeable independently. Ordered after the task + /// that materializes `src_id`, like every other op on the session's FIFO worker. + RegisterAlias { + stream_id: StreamId, + new_id: TensorId, + src_id: TensorId, + }, + RegisterTensorRemote(StreamId, TensorRemote, TensorId), + ExposeTensorRemote { + stream_id: StreamId, + tensor: TensorIr, + count: u32, + capability: TransferCapability, + target: PeerId, + }, + /// Source side of a same-host transfer: hand the device-resident primitive for `tensor` + /// to the server's local comm registry under `transfer_id`. No host readback — the + /// counterpart [`RegisterTensorLocal`](Task::RegisterTensorLocal), running on the + /// target session of the same server, moves it onto the target device via the inner + /// backend's `to_device`. + /// + /// `stream_id` is the client stream that produced `tensor`, so the server reads it back on + /// the same stream the producing ops ran on rather than an arbitrary server thread. + ExposeTensorLocal { + stream_id: StreamId, + tensor: TensorIr, + transfer_id: LocalTransferId, + }, + /// Target side of a same-host transfer: wait for `transfer_id` to be exposed, move the + /// primitive onto this session's device, and register it under `new_id`. + /// + /// `stream_id` is the client stream that will consume `new_id`, so the registration lands + /// on the same stream as the ops that use the transferred tensor. + RegisterTensorLocal { + stream_id: StreamId, + transfer_id: LocalTransferId, + new_id: TensorId, + }, + ReadTensor(RequestId, StreamId, TensorIr), + SyncBackend(RequestId, StreamId), + DTypeUsage(RequestId, DType), +} + +#[allow(missing_docs)] +#[derive(Serialize, Deserialize, Debug)] +pub struct TaskResponse { + pub content: TaskResponseContent, + pub id: RequestId, +} + +#[allow(missing_docs)] +#[derive(Serialize, Deserialize, Debug)] +pub enum TaskResponseContent { + /// Server responds with the selected device's settings plus the total number of devices + /// it hosts (so the client can enumerate them, see [`RemoteDevice::enumerate`]). + Init(SessionInfo), + ReadTensor(Result), + SyncBackend(Result<(), ExecutionError>), + DTypeUsage(DTypeUsageSet), +} diff --git a/crates/burn-remote/src/telemetry.rs b/crates/burn-remote/src/telemetry.rs new file mode 100644 index 0000000..6fb0aa4 --- /dev/null +++ b/crates/burn-remote/src/telemetry.rs @@ -0,0 +1,531 @@ +//! Best-effort server telemetry. +//! +//! Session workers emit [`TelemetryEvent`]s into a [`TelemetryProbe`]; a monitoring view holds a +//! [`TelemetrySubscription`] and renders them. The probe never blocks compute: events are built +//! only while a subscriber is attached and are dropped on lag, so a slow or idle viewer cannot +//! apply backpressure to a worker. + +use std::collections::HashMap; +use std::sync::Arc; + +use burn_backend::{DType, Shape}; +use burn_ir::{GraphId, NumericOperationIr, OperationIr, TensorId}; +use burn_std::id::StreamId; +use serde::{Deserialize, Serialize}; +use tokio::sync::broadcast; + +use crate::shared::{RequestId, SessionId}; + +/// A tensor as it appears in the dataflow graph: identity plus the shape and dtype it was +/// produced with. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TensorRef { + pub id: TensorId, + pub shape: Shape, + pub dtype: DType, +} + +/// One operation inside a cached graph, in relative form (positional tensor ids). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct GraphOp { + pub kind: OpClass, + pub inputs: Vec, + pub outputs: Vec, +} + +impl GraphOp { + pub(crate) fn summarize(op: &OperationIr) -> Self { + Self { + kind: classify(op), + inputs: op.inputs().map(|tensor| tensor.id).collect(), + outputs: op + .outputs() + .map(|tensor| TensorRef { + id: tensor.id, + shape: tensor.shape.clone(), + dtype: tensor.dtype, + }) + .collect(), + } + } +} + +/// Coarse operation category, used to colour and aggregate the op stream. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub enum OpClass { + Matmul, + Conv, + Linear, + Activation, + Reduction, + Elementwise, + Compare, + Cast, + Index, + Reshape, + Init, + Random, + Drop, + Custom, + Distributed, + Other, +} + +impl OpClass { + /// Every variant, in a stable order for layout and iteration. + pub const ALL: [OpClass; 16] = [ + OpClass::Matmul, + OpClass::Conv, + OpClass::Linear, + OpClass::Activation, + OpClass::Reduction, + OpClass::Elementwise, + OpClass::Compare, + OpClass::Cast, + OpClass::Index, + OpClass::Reshape, + OpClass::Init, + OpClass::Random, + OpClass::Drop, + OpClass::Custom, + OpClass::Distributed, + OpClass::Other, + ]; + + /// Stable lowercase label for display and aggregation. + pub fn label(self) -> &'static str { + match self { + OpClass::Matmul => "matmul", + OpClass::Conv => "conv", + OpClass::Linear => "linear", + OpClass::Activation => "activation", + OpClass::Reduction => "reduction", + OpClass::Elementwise => "elementwise", + OpClass::Compare => "compare", + OpClass::Cast => "cast", + OpClass::Index => "index", + OpClass::Reshape => "reshape", + OpClass::Init => "init", + OpClass::Random => "random", + OpClass::Drop => "drop", + OpClass::Custom => "custom", + OpClass::Distributed => "distributed", + OpClass::Other => "other", + } + } +} + +/// Classify an [`OperationIr`] into a coarse [`OpClass`]. +pub fn classify(op: &OperationIr) -> OpClass { + use OperationIr as O; + match op { + O::NumericFloat(_, n) | O::NumericInt(_, n) => classify_numeric(n), + O::Float(_, f) => classify_float(f), + O::Module(m) => classify_module(m), + O::Activation(_) => OpClass::Activation, + O::BaseFloat(b) | O::BaseInt(b) | O::BaseBool(b) => classify_base(b), + O::Init(_) => OpClass::Init, + O::Custom(_) => OpClass::Custom, + O::Drop(_) => OpClass::Drop, + O::Distributed(_) => OpClass::Distributed, + O::Bool(_) | O::Int(_) => OpClass::Elementwise, + } +} + +fn classify_numeric(op: &NumericOperationIr) -> OpClass { + use NumericOperationIr as N; + match op { + N::Mean(_) + | N::MeanDim(_) + | N::Sum(_) + | N::SumDim(_) + | N::Prod(_) + | N::ProdDim(_) + | N::Max(_) + | N::MaxDim(_) + | N::MaxDimWithIndices(_) + | N::Min(_) + | N::MinDim(_) + | N::MinDimWithIndices(_) + | N::MaxAbs(_) + | N::MaxAbsDim(_) + | N::ArgMax(_) + | N::ArgMin(_) + | N::ArgTopK(_) + | N::TopK(_) + | N::CumSum(_) + | N::CumProd(_) + | N::CumMin(_) + | N::CumMax(_) + | N::Sort(_) + | N::SortWithIndices(_) + | N::ArgSort(_) => OpClass::Reduction, + N::Greater(_) + | N::GreaterElem(_) + | N::GreaterEqual(_) + | N::GreaterEqualElem(_) + | N::Lower(_) + | N::LowerElem(_) + | N::LowerEqual(_) + | N::LowerEqualElem(_) => OpClass::Compare, + N::Full(_) | N::IntRandom(_) => OpClass::Init, + _ => OpClass::Elementwise, + } +} + +fn classify_float(op: &burn_ir::FloatOperationIr) -> OpClass { + use burn_ir::FloatOperationIr as F; + match op { + F::Matmul(_) => OpClass::Matmul, + F::Random(_) => OpClass::Random, + F::IntoInt(_) | F::Quantize(_) | F::Dequantize(_) => OpClass::Cast, + _ => OpClass::Elementwise, + } +} + +fn classify_module(op: &burn_ir::ModuleOperationIr) -> OpClass { + let name = format!("{op:?}"); + if name.starts_with("Linear") || name.starts_with("Embedding") { + OpClass::Linear + } else if name.starts_with("Conv") || name.starts_with("DeformableConv") { + OpClass::Conv + } else { + OpClass::Other + } +} + +fn classify_base(op: &burn_ir::BaseOperationIr) -> OpClass { + use burn_ir::BaseOperationIr as B; + match op { + B::Reshape(_) + | B::SwapDims(_) + | B::Permute(_) + | B::Flip(_) + | B::Expand(_) + | B::RepeatDim(_) + | B::Cat(_) + | B::Unfold(_) => OpClass::Reshape, + B::Cast(_) => OpClass::Cast, + B::Empty(_) | B::Ones(_) | B::Zeros(_) => OpClass::Init, + B::Equal(_) + | B::EqualElem(_) + | B::NotEqual(_) + | B::NotEqualElem(_) + | B::All(_) + | B::Any(_) + | B::AllDim(_) + | B::AnyDim(_) => OpClass::Compare, + _ => OpClass::Index, + } +} + +/// Whether a transfer crosses hosts or stays on this machine. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum TransferScope { + Local, + Remote, +} + +/// Lifecycle point of a tensor transfer. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum TransferPhase { + Started, + Completed, + Failed, +} + +/// A single observation from a session worker. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum TelemetryEvent { + SessionOpened { + session: SessionId, + device: u32, + }, + SessionClosed { + session: SessionId, + }, + Op { + session: SessionId, + stream: StreamId, + kind: OpClass, + inputs: Vec, + outputs: Vec, + }, + TensorDropped { + session: SessionId, + tensor: TensorId, + }, + Transfer { + session: SessionId, + peer: Option, + scope: TransferScope, + phase: TransferPhase, + }, + Read { + session: SessionId, + request: RequestId, + }, + Sync { + session: SessionId, + request: RequestId, + }, + GraphRegistered { + session: SessionId, + graph: GraphId, + ops: Vec, + bytes: usize, + }, + GraphExecuted { + session: SessionId, + graph: GraphId, + stream: StreamId, + bindings_bytes: usize, + }, +} + +impl TelemetryEvent { + pub(crate) fn unfused_op(session: SessionId, stream: StreamId, op: &OperationIr) -> Self { + match op { + OperationIr::Drop(tensor) => TelemetryEvent::TensorDropped { + session, + tensor: tensor.id, + }, + _ => { + let GraphOp { + kind, + inputs, + outputs, + } = GraphOp::summarize(op); + TelemetryEvent::Op { + session, + stream, + kind, + inputs, + outputs, + } + } + } + } + + pub(crate) fn graph_registered( + session: SessionId, + graph: GraphId, + ops: &[OperationIr], + bytes: usize, + ) -> Self { + TelemetryEvent::GraphRegistered { + session, + graph, + ops: ops.iter().map(GraphOp::summarize).collect(), + bytes, + } + } + + pub(crate) fn graph_executed( + session: SessionId, + graph: GraphId, + stream: StreamId, + bindings_bytes: usize, + ) -> Self { + TelemetryEvent::GraphExecuted { + session, + graph, + stream, + bindings_bytes, + } + } +} + +pub(crate) const CHANNEL_CAPACITY: usize = 4096; + +/// Cloneable handle a worker emits into. Inert until a [`TelemetrySubscription`] is attached. +#[derive(Clone)] +pub struct TelemetryProbe { + tx: Option>>, +} + +impl TelemetryProbe { + /// A probe that drops everything. Worker emit points compile to a cheap no-op against it. + pub fn disabled() -> Self { + Self { tx: None } + } + + /// Create an active probe with no initial subscription. Subscribers attach later with + /// [`subscribe`](Self::subscribe); the probe stays inert until at least one is listening. + pub fn new(capacity: usize) -> Self { + let (tx, _rx) = broadcast::channel(capacity); + Self { tx: Some(tx) } + } + + /// Create an active probe and its first subscription. `capacity` bounds the per-subscriber + /// backlog; older events are dropped once a subscriber falls that far behind. + pub fn channel(capacity: usize) -> (Self, TelemetrySubscription) { + let (tx, rx) = broadcast::channel(capacity); + (Self { tx: Some(tx) }, TelemetrySubscription { rx }) + } + + /// Attach another subscription, or `None` if this probe is disabled. + pub fn subscribe(&self) -> Option { + self.tx + .as_ref() + .map(|tx| TelemetrySubscription { rx: tx.subscribe() }) + } + + /// Whether building and emitting an event is worthwhile right now. + #[inline] + pub fn is_active(&self) -> bool { + self.tx.as_ref().is_some_and(|tx| tx.receiver_count() > 0) + } + + /// Emit an event, building it only when a subscriber is listening. + #[inline] + pub fn emit(&self, event: impl FnOnce() -> TelemetryEvent) { + if let Some(tx) = &self.tx + && tx.receiver_count() > 0 + { + let _ = tx.send(Arc::new(event())); + } + } +} + +/// Receiving end of a [`TelemetryProbe`]. +pub struct TelemetrySubscription { + rx: broadcast::Receiver>, +} + +/// Outcome of a non-blocking drain. +pub enum DrainStatus { + /// The channel is still open; `lagged` counts events dropped before this drain. + Open { lagged: u64 }, + /// Every sender has been dropped. + Closed, +} + +impl TelemetrySubscription { + /// Drain all currently buffered events into `sink` without blocking, reporting whether the + /// channel is still open and how many events were dropped to lag. Intended to be called once + /// per UI frame. + pub fn drain_into(&mut self, sink: &mut Vec>) -> DrainStatus { + use broadcast::error::TryRecvError; + let mut lagged = 0; + loop { + match self.rx.try_recv() { + Ok(event) => sink.push(event), + Err(TryRecvError::Empty) => return DrainStatus::Open { lagged }, + Err(TryRecvError::Lagged(n)) => lagged += n, + Err(TryRecvError::Closed) => return DrainStatus::Closed, + } + } + } + + /// Await the next event, skipping lag gaps. Returns `None` once every sender is dropped. + pub async fn recv(&mut self) -> Option> { + use broadcast::error::RecvError; + loop { + match self.rx.recv().await { + Ok(event) => return Some(event), + Err(RecvError::Lagged(_)) => continue, + Err(RecvError::Closed) => return None, + } + } + } +} + +struct GraphCost { + bytes: u64, + ops: u64, +} + +/// A running fold of the op-graph caching economics over a telemetry stream. +#[derive(Default)] +pub struct TrafficAggregator { + graphs: HashMap, + baseline: u64, + actual: u64, + fused_ops: u64, + unfused_ops: u64, + unfused_by_kind: HashMap, +} + +impl TrafficAggregator { + pub fn apply(&mut self, event: &TelemetryEvent) { + match event { + TelemetryEvent::GraphRegistered { + graph, ops, bytes, .. + } => { + self.actual += *bytes as u64; + self.graphs.insert( + *graph, + GraphCost { + bytes: *bytes as u64, + ops: ops.len() as u64, + }, + ); + } + TelemetryEvent::GraphExecuted { + graph, + bindings_bytes, + .. + } => { + if let Some(cost) = self.graphs.get(graph) { + self.baseline += cost.bytes; + self.fused_ops += cost.ops; + } + self.actual += *bindings_bytes as u64; + } + TelemetryEvent::Op { kind, .. } => { + self.unfused_ops += 1; + *self.unfused_by_kind.entry(*kind).or_default() += 1; + } + TelemetryEvent::TensorDropped { .. } => { + self.unfused_ops += 1; + *self.unfused_by_kind.entry(OpClass::Drop).or_default() += 1; + } + _ => {} + } + } + + pub fn snapshot(&self) -> FusionSnapshot { + FusionSnapshot { + fused_ops: self.fused_ops, + unfused_ops: self.unfused_ops, + baseline: self.baseline, + actual: self.actual, + } + } + + /// Unfused-op tally by class, highest first. + pub fn unfused_by_kind(&self) -> Vec<(OpClass, u64)> { + let mut entries: Vec<_> = self + .unfused_by_kind + .iter() + .map(|(kind, count)| (*kind, *count)) + .collect(); + entries.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.label().cmp(b.0.label()))); + entries + } +} + +pub struct FusionSnapshot { + pub fused_ops: u64, + pub unfused_ops: u64, + pub baseline: u64, + pub actual: u64, +} + +impl FusionSnapshot { + /// Baseline a non-fusion peer would have moved, minus what fusion actually moved. + pub fn saved(&self) -> u64 { + self.baseline.saturating_sub(self.actual) + } + + pub fn total_ops(&self) -> u64 { + self.fused_ops + self.unfused_ops + } +} + +pub(crate) fn serialized_len(value: &T) -> usize { + rmp_serde::to_vec(value) + .map(|bytes| bytes.len()) + .unwrap_or(0) +} diff --git a/crates/burn-remote/src/transport/identity.rs b/crates/burn-remote/src/transport/identity.rs new file mode 100644 index 0000000..39a058c --- /dev/null +++ b/crates/burn-remote/src/transport/identity.rs @@ -0,0 +1,109 @@ +use core::fmt; + +#[cfg(feature = "websocket")] +use burn_communication::Address; +use serde::{Deserialize, Serialize}; + +/// Stable identity of a Burn Remote peer. +/// +/// Network paths are deliberately not part of the identity. An Iroh peer keeps the same +/// identity while moving between direct addresses and relays. +#[derive(Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Debug)] +pub enum PeerId { + /// An Iroh endpoint, authenticated by its public key. + #[cfg(feature = "iroh")] + Iroh(iroh::EndpointId), + /// A legacy WebSocket endpoint. + #[cfg(feature = "websocket")] + WebSocket(Address), +} + +impl fmt::Display for PeerId { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + #[cfg(feature = "iroh")] + Self::Iroh(id) => write!(f, "iroh://{id}"), + #[cfg(feature = "websocket")] + Self::WebSocket(address) => address.fmt(f), + } + } +} + +// Only the Iroh server resolves a `PeerId` back to an endpoint id (to address tensor transfers). +#[cfg(all(feature = "iroh", feature = "server"))] +impl PeerId { + /// The Iroh endpoint id, or `None` for a non-Iroh peer. + pub(crate) fn into_iroh_id(self) -> Option { + #[cfg(feature = "websocket")] + return match self { + Self::Iroh(id) => Some(id), + Self::WebSocket(_) => None, + }; + #[cfg(not(feature = "websocket"))] + { + let Self::Iroh(id) = self; + Some(id) + } + } +} + +/// A peer identity plus the mutable network paths that may be used to reach it. +#[derive(Clone, PartialEq, Eq, Hash, Serialize, Deserialize, Debug)] +pub enum PeerAddr { + /// An Iroh endpoint address. It may contain direct and relay paths, or only an endpoint id + /// when the configured Iroh address lookup can resolve it. + #[cfg(feature = "iroh")] + Iroh(iroh::EndpointAddr), + /// A legacy WebSocket address. + #[cfg(feature = "websocket")] + WebSocket(Address), +} + +impl PeerAddr { + /// Return the stable peer identity, excluding dialing hints. + pub fn id(&self) -> PeerId { + match self { + #[cfg(feature = "iroh")] + Self::Iroh(address) => PeerId::Iroh(address.id), + #[cfg(feature = "websocket")] + Self::WebSocket(address) => PeerId::WebSocket(address.clone()), + } + } + + /// Return true when this is an Iroh peer. + pub fn is_iroh(&self) -> bool { + match self { + #[cfg(feature = "iroh")] + Self::Iroh(_) => true, + #[cfg(feature = "websocket")] + Self::WebSocket(_) => false, + } + } +} + +impl fmt::Display for PeerAddr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.id().fmt(f) + } +} + +#[cfg(feature = "iroh")] +impl From for PeerAddr { + fn from(value: iroh::EndpointAddr) -> Self { + Self::Iroh(value) + } +} + +#[cfg(feature = "iroh")] +impl From for PeerAddr { + fn from(value: iroh::EndpointId) -> Self { + Self::Iroh(value.into()) + } +} + +#[cfg(feature = "websocket")] +impl From

for PeerAddr { + fn from(value: Address) -> Self { + Self::WebSocket(value) + } +} diff --git a/crates/burn-remote/src/transport/iroh/link.rs b/crates/burn-remote/src/transport/iroh/link.rs new file mode 100644 index 0000000..8e6b413 --- /dev/null +++ b/crates/burn-remote/src/transport/iroh/link.rs @@ -0,0 +1,27 @@ +//! Iroh implementations of the session-link frame traits. +//! +//! An Iroh session is one bidirectional QUIC stream; its two halves are already separate owned +//! values (`SendStream` / `RecvStream`), so they map directly onto [`FrameSink`] / [`FrameSource`]. + +use bytes::Bytes; +use iroh::endpoint::{RecvStream, SendStream}; + +use super::node::{recv_frame, send_frame}; +use crate::transport::link::{FrameSink, FrameSource}; + +impl FrameSink for SendStream { + async fn send(&mut self, frame: Bytes) -> Result<(), String> { + send_frame(self, &frame).await + } + + async fn close(&mut self) -> Result<(), String> { + self.finish() + .map_err(|err| format!("Failed to finish Iroh stream: {err}")) + } +} + +impl FrameSource for RecvStream { + async fn recv(&mut self) -> Result, String> { + recv_frame(self).await.map(|frame| frame.map(Bytes::from)) + } +} diff --git a/crates/burn-remote/src/transport/iroh/mod.rs b/crates/burn-remote/src/transport/iroh/mod.rs new file mode 100644 index 0000000..4b415ff --- /dev/null +++ b/crates/burn-remote/src/transport/iroh/mod.rs @@ -0,0 +1,21 @@ +//! Iroh transport implementation. +//! +//! Self-contained: everything `cfg(feature = "iroh")`-specific that the session, transfer, and +//! client layers depend on lives under this module. + +mod secret; +pub use secret::RemoteSecret; + +mod link; +pub mod node; + +#[cfg(feature = "server")] +pub mod protocol; +#[cfg(all(feature = "server", not(target_family = "wasm")))] +pub mod server; +#[cfg(feature = "server")] +mod time; +#[cfg(feature = "server")] +mod transfer; +#[cfg(feature = "server")] +pub(crate) use transfer::IrohTransfer; diff --git a/crates/burn-remote/src/transport/iroh/node.rs b/crates/burn-remote/src/transport/iroh/node.rs new file mode 100644 index 0000000..b18d167 --- /dev/null +++ b/crates/burn-remote/src/transport/iroh/node.rs @@ -0,0 +1,247 @@ +//! Process-level Iroh endpoint used by Burn Remote clients and compute nodes. + +use std::{collections::HashMap, sync::Arc}; + +use iroh::{ + Endpoint, EndpointAddr, EndpointId, + endpoint::{Connection, RecvStream, SendStream}, +}; +use tokio::sync::Mutex; +use tokio::sync::OnceCell; + +use crate::{PeerAddr, PeerId}; + +/// ALPN used by the version-one Burn Remote protocol. +pub const BURN_REMOTE_ALPN: &[u8] = b"burn/remote/1"; + +/// Identifies the purpose of a bidirectional stream inside a shared Iroh connection. +#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize)] +pub(crate) enum StreamKind { + Session, + TensorTransfer, +} + +#[derive(Debug, serde::Serialize, serde::Deserialize)] +struct StreamHeader { + version: u16, + kind: StreamKind, +} + +const STREAM_VERSION: u16 = 1; +const MAX_FRAME_SIZE: usize = 1024 * 1024 * 1024; + +struct RemoteNodeInner { + endpoint: Endpoint, + connections: Mutex>>>, +} + +/// A process-level Burn Remote networking node. +/// +/// Clone this handle freely. Every clone shares one Iroh [`Endpoint`] and one connection pool, +/// so all remote devices in the process multiplex their sessions over the same peer connection. +#[derive(Clone)] +pub struct RemoteNode { + inner: Arc, +} + +impl core::fmt::Debug for RemoteNode { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("RemoteNode") + .field("endpoint_id", &self.id()) + .finish_non_exhaustive() + } +} + +impl RemoteNode { + /// Use an application-configured Iroh endpoint. + /// + /// Applications serving Burn Remote must include [`BURN_REMOTE_ALPN`] in the endpoint's + /// accepted ALPN list, or route that ALPN to Burn's protocol handler. + /// + /// On native client builds, the runtime that drives a device's session is captured when the + /// device is created (see [`RemoteNode::device`]), not here — so create devices from the + /// Tokio runtime that owns this endpoint. + pub fn from_endpoint(endpoint: Endpoint) -> Self { + Self { + inner: Arc::new(RemoteNodeInner { + endpoint, + connections: Mutex::new(HashMap::new()), + }), + } + } + + /// The cryptographic identity of this node. + pub fn id(&self) -> EndpointId { + self.inner.endpoint.id() + } + + /// Access the underlying endpoint for relay, discovery, router, and observability setup. + pub fn endpoint(&self) -> &Endpoint { + &self.inner.endpoint + } + + pub(crate) async fn open_stream( + &self, + peer: &PeerAddr, + kind: StreamKind, + ) -> Result<(SendStream, RecvStream), String> { + // Only the Iroh variant remains when the websocket transport is compiled out. + #[cfg_attr( + not(feature = "websocket"), + allow(clippy::infallible_destructuring_match) + )] + let peer = match peer { + PeerAddr::Iroh(peer) => peer, + #[cfg(feature = "websocket")] + PeerAddr::WebSocket(_) => { + return Err("Iroh node cannot open a stream to a non-Iroh peer".into()); + } + }; + let connection = self.connection(peer.clone()).await?; + let (mut send, recv) = connection + .open_bi() + .await + .map_err(|err| format!("Failed to open Iroh stream to {}: {err}", peer.id))?; + let header = rmp_serde::to_vec(&StreamHeader { + version: STREAM_VERSION, + kind, + }) + .map_err(|err| format!("Failed to encode Iroh stream header: {err}"))?; + send_frame(&mut send, &header).await?; + Ok((send, recv)) + } + + #[cfg(feature = "server")] + pub(crate) async fn accept_stream( + connection: &Connection, + ) -> Result, String> { + let (send, mut recv) = match connection.accept_bi().await { + Ok(stream) => stream, + Err(err) => { + if connection.close_reason().is_some() { + return Ok(None); + } + return Err(format!("Failed to accept Iroh stream: {err}")); + } + }; + let Some(frame) = recv_frame(&mut recv).await? else { + return Ok(None); + }; + let header: StreamHeader = rmp_serde::from_slice(&frame) + .map_err(|err| format!("Invalid Iroh stream header: {err}"))?; + if header.version != STREAM_VERSION { + return Err(format!( + "Unsupported Burn Remote stream version {} (expected {STREAM_VERSION})", + header.version + )); + } + Ok(Some((header.kind, send, recv))) + } + + async fn connection(&self, peer: EndpointAddr) -> Result { + loop { + let cell = { + let mut connections = self.inner.connections.lock().await; + connections + .entry(peer.id) + .or_insert_with(|| Arc::new(OnceCell::new())) + .clone() + }; + + if let Some(connection) = cell.get() + && connection.close_reason().is_some() + { + let mut connections = self.inner.connections.lock().await; + if connections + .get(&peer.id) + .is_some_and(|current| Arc::ptr_eq(current, &cell)) + { + connections.remove(&peer.id); + } + continue; + } + + let endpoint = self.inner.endpoint.clone(); + let peer_for_connect = peer.clone(); + let connection = cell + .get_or_try_init(|| async move { + endpoint + .connect(peer_for_connect.clone(), BURN_REMOTE_ALPN) + .await + .map_err(|err| { + format!( + "Failed to connect to Iroh peer {}: {err}", + peer_for_connect.id + ) + }) + }) + .await?; + return Ok(connection.clone()); + } + } + + #[cfg(feature = "server")] + pub(crate) async fn remember_connection(&self, connection: Connection) { + let remote = connection.remote_id(); + let cell = { + let mut connections = self.inner.connections.lock().await; + match connections.get(&remote) { + Some(cell) + if cell + .get() + .is_some_and(|existing| existing.close_reason().is_none()) => + { + return; + } + _ => { + let cell = Arc::new(OnceCell::new()); + connections.insert(remote, cell.clone()); + cell + } + } + }; + let _ = cell.set(connection); + } +} + +pub(crate) async fn send_frame(send: &mut SendStream, bytes: &[u8]) -> Result<(), String> { + if bytes.len() > MAX_FRAME_SIZE { + return Err(format!( + "Burn Remote frame is too large: {} bytes (max {MAX_FRAME_SIZE})", + bytes.len() + )); + } + send.write_all(&(bytes.len() as u64).to_le_bytes()) + .await + .map_err(|err| format!("Failed to write Iroh frame length: {err}"))?; + send.write_all(bytes) + .await + .map_err(|err| format!("Failed to write Iroh frame: {err}"))?; + Ok(()) +} + +pub(crate) async fn recv_frame(recv: &mut RecvStream) -> Result>, String> { + let mut length = [0u8; 8]; + match recv.read_exact(&mut length).await { + Ok(_) => {} + Err(iroh::endpoint::ReadExactError::FinishedEarly(0)) => return Ok(None), + Err(err) => return Err(format!("Failed to read Iroh frame length: {err}")), + } + let length = u64::from_le_bytes(length) as usize; + if length > MAX_FRAME_SIZE { + return Err(format!( + "Peer sent an oversized Burn Remote frame: {length} bytes (max {MAX_FRAME_SIZE})" + )); + } + let mut bytes = vec![0; length]; + recv.read_exact(&mut bytes) + .await + .map_err(|err| format!("Failed to read Iroh frame: {err}"))?; + Ok(Some(bytes)) +} + +impl From<&RemoteNode> for PeerId { + fn from(value: &RemoteNode) -> Self { + PeerId::Iroh(value.id()) + } +} diff --git a/crates/burn-remote/src/transport/iroh/protocol.rs b/crates/burn-remote/src/transport/iroh/protocol.rs new file mode 100644 index 0000000..0c7ca2a --- /dev/null +++ b/crates/burn-remote/src/transport/iroh/protocol.rs @@ -0,0 +1,201 @@ +//! Iroh protocol handler for Burn Remote compute and tensor-transfer streams. + +use std::{fmt, sync::Arc}; + +use burn_backend::tensor::Device; +use burn_ir::BackendIr; +use burn_router::CustomOpRegistry; +use iroh::{ + Endpoint, EndpointId, + endpoint::{Connection, RecvStream, SendStream}, + protocol::{AcceptError, DynProtocolHandler, ProtocolHandler}, +}; + +use crate::{ + PeerId, + server::{pump::drive_session, session::SessionManager, spawn::spawn_detached}, + telemetry::TelemetryProbe, +}; + +use super::{ + IrohTransfer, + node::{RemoteNode, StreamKind}, +}; + +/// Information presented to a compute node before a remote session is accepted. +pub struct AuthorizationRequest<'a> { + /// Authenticated Iroh identity of the connecting peer. + pub peer: EndpointId, + /// Compute-device index requested by the peer. + pub device_index: u32, + /// Opaque credential supplied by the application when creating the remote device. + pub credential: &'a [u8], +} + +/// Application authorization policy for incoming compute sessions. +pub trait PeerAuthorizer: Send + Sync + 'static { + /// Return `Ok(())` to allow the session, or a user-facing rejection reason. + fn authorize(&self, request: AuthorizationRequest<'_>) -> Result<(), String>; +} + +impl PeerAuthorizer for F +where + F: Fn(AuthorizationRequest<'_>) -> Result<(), String> + Send + Sync + 'static, +{ + fn authorize(&self, request: AuthorizationRequest<'_>) -> Result<(), String> { + self(request) + } +} + +#[derive(Debug, Default)] +pub struct AllowAll; + +impl PeerAuthorizer for AllowAll { + fn authorize(&self, _request: AuthorizationRequest<'_>) -> Result<(), String> { + Ok(()) + } +} + +/// Iroh protocol handler for Burn Remote compute and tensor-transfer streams. +/// +/// Register this handler in an existing Iroh `Router` to compose Burn with other application +/// protocols on the same endpoint. +pub struct IrohRemoteProtocol { + node: RemoteNode, + sessions: Arc>>, + transfer: Arc>, + authorizer: Arc, +} + +impl fmt::Debug for IrohRemoteProtocol { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("IrohRemoteProtocol") + .field("endpoint_id", &self.node.id()) + .finish_non_exhaustive() + } +} + +impl IrohRemoteProtocol { + /// Create a handler hosting `devices` on `node`. + pub fn new( + endpoint: Endpoint, + devices: Vec>, + authorizer: Arc, + probe: TelemetryProbe, + custom_ops: CustomOpRegistry, + ) -> Self { + let node = RemoteNode::from_endpoint(endpoint); + let transfer = Arc::new(IrohTransfer::new(node.clone())); + + let sessions = Arc::new( + SessionManager::new(devices.to_vec(), transfer.clone()) + .with_telemetry(probe.clone()) + .with_custom_ops(custom_ops.clone()), + ); + Self { + node, + sessions, + transfer, + authorizer, + } + } + + /// Drive an accepted session stream through the shared [`drive_session`] pump. + /// + /// The Iroh-specific parts are just the authenticated peer identity (`remote_id`, checked by the + /// application's [`PeerAuthorizer`]) and this server's own id, echoed to the client. + async fn handle_session( + sessions: Arc>>, + authorizer: Arc, + server_id: EndpointId, + remote_id: EndpointId, + send: SendStream, + recv: RecvStream, + ) -> Result<(), String> { + drive_session( + recv, + send, + sessions, + Some(PeerId::Iroh(server_id)), + |init| { + authorizer.authorize(AuthorizationRequest { + peer: remote_id, + device_index: init.device_index, + credential: &init.authorization, + }) + }, + ) + .await + } +} + +/// A backend-erased Burn Remote protocol handler. +/// +/// The dispatch layer resolves a `Device` to a concrete backend and builds an +/// [`IrohRemoteProtocol`]; this wraps it as a single non-generic type, so an application can +/// register Burn on its own Iroh `Router` without naming a backend. Hand it directly to +/// `RouterBuilder::accept` under [`BURN_REMOTE_ALPN`](super::node::BURN_REMOTE_ALPN). +pub struct RemoteProtocol(Box); + +impl RemoteProtocol { + /// Erase a concrete protocol handler behind this non-generic type. + pub fn new(handler: impl ProtocolHandler) -> Self { + Self(handler.into()) + } +} + +impl fmt::Debug for RemoteProtocol { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("RemoteProtocol").finish_non_exhaustive() + } +} + +impl From for Box { + fn from(protocol: RemoteProtocol) -> Self { + protocol.0 + } +} + +impl ProtocolHandler for IrohRemoteProtocol { + async fn accept(&self, connection: Connection) -> Result<(), AcceptError> { + let remote_id = connection.remote_id(); + self.node.remember_connection(connection.clone()).await; + loop { + let Some((kind, send, recv)) = RemoteNode::accept_stream(&connection) + .await + .map_err(user_error)? + else { + return Ok(()); + }; + + match kind { + StreamKind::Session => { + let sessions = self.sessions.clone(); + let authorizer = self.authorizer.clone(); + let server_id = self.node.id(); + spawn_detached(async move { + if let Err(err) = Self::handle_session( + sessions, authorizer, server_id, remote_id, send, recv, + ) + .await + { + log::warn!("Rejected or failed Iroh remote session: {err}"); + } + }); + } + StreamKind::TensorTransfer => { + let transfer = self.transfer.clone(); + spawn_detached(async move { + if let Err(err) = transfer.handle_stream(remote_id, send, recv).await { + log::warn!("Iroh tensor-transfer stream failed: {err}"); + } + }); + } + } + } + } +} + +fn user_error(reason: String) -> AcceptError { + AcceptError::from_err(std::io::Error::other(reason)) +} diff --git a/crates/burn-remote/src/transport/iroh/secret.rs b/crates/burn-remote/src/transport/iroh/secret.rs new file mode 100644 index 0000000..df80211 --- /dev/null +++ b/crates/burn-remote/src/transport/iroh/secret.rs @@ -0,0 +1,35 @@ +//! Server identity for the Iroh transport. + +/// A compute server's stable identity: the secret stays on the server, and the public +/// [`id`](Self::id) it yields is the address clients dial. Generate one with [`random`](Self::random) +/// and persist [`to_bytes`](Self::to_bytes) for a stable address across restarts, or derive it from a +/// seed with [`from_bytes`](Self::from_bytes). +#[derive(Clone)] +pub struct RemoteSecret(iroh::SecretKey); + +impl RemoteSecret { + /// A fresh random identity. Persist [`to_bytes`](Self::to_bytes) to reuse the same address later. + pub fn random() -> Self { + Self(iroh::SecretKey::generate()) + } + + /// A deterministic identity from 32 seed bytes (e.g. a hash of an application name). + pub fn from_bytes(bytes: [u8; 32]) -> Self { + Self(iroh::SecretKey::from_bytes(&bytes)) + } + + /// The raw 32 bytes, to persist and reload a stable identity. + pub fn to_bytes(&self) -> [u8; 32] { + self.0.to_bytes() + } + + /// The public identity clients dial. + pub fn id(&self) -> iroh::EndpointId { + self.0.public() + } + + #[cfg(feature = "server")] + pub(crate) fn secret_key(&self) -> iroh::SecretKey { + self.0.clone() + } +} diff --git a/crates/burn-remote/src/transport/iroh/server.rs b/crates/burn-remote/src/transport/iroh/server.rs new file mode 100644 index 0000000..c6bf312 --- /dev/null +++ b/crates/burn-remote/src/transport/iroh/server.rs @@ -0,0 +1,54 @@ +use crate::server::spawn::os_shutdown_signal; +use crate::telemetry::TelemetryProbe; +use crate::transport::iroh::node::BURN_REMOTE_ALPN; +use crate::transport::iroh::protocol::{AllowAll, IrohRemoteProtocol}; +#[cfg(not(target_family = "wasm"))] +use burn_backend::tensor::Device; +#[cfg(not(target_family = "wasm"))] +use burn_ir::BackendIr; +#[cfg(not(target_family = "wasm"))] +use burn_router::CustomOpRegistry; +use iroh::{Endpoint, endpoint::presets, protocol::Router}; +use std::sync::Arc; + +/// Serve Burn Remote over Iroh until the process receives its shutdown signal. +/// +/// Binds a server endpoint with the stable identity carried by `secret` and hosts `devices` as the +/// sole protocol on it. Reached through [`RemoteServerBuilder`](super::RemoteServerBuilder) (the +/// single turnkey entry point); use [`RemoteNode::protocol`] for composition with other protocols. +#[cfg(not(target_family = "wasm"))] +pub(crate) async fn start_iroh_async( + secret: crate::RemoteSecret, + devices: Vec>, + custom_ops: CustomOpRegistry, +) { + let endpoint = Endpoint::builder(presets::N0) + .secret_key(secret.secret_key()) + .alpns(vec![BURN_REMOTE_ALPN.to_vec()]) + .bind() + .await + .expect("Can bind the Burn Remote server endpoint"); + + let probe = if crate::metrics::TelemetryLogger::enabled() { + TelemetryProbe::new(crate::telemetry::CHANNEL_CAPACITY) + } else { + TelemetryProbe::disabled() + }; + + let protocol = IrohRemoteProtocol::new( + endpoint.clone(), + devices, + Arc::new(AllowAll), + probe, + custom_ops, + ); + + let router = Router::builder(endpoint) + .accept(BURN_REMOTE_ALPN, protocol) + .spawn(); + + os_shutdown_signal().await; + if let Err(err) = router.shutdown().await { + log::warn!("Burn Remote Iroh router shutdown failed: {err}"); + } +} diff --git a/crates/burn-remote/src/transport/iroh/time.rs b/crates/burn-remote/src/transport/iroh/time.rs new file mode 100644 index 0000000..e6f931a --- /dev/null +++ b/crates/burn-remote/src/transport/iroh/time.rs @@ -0,0 +1,29 @@ +use core::{future::Future, time::Duration}; + +#[cfg(not(target_family = "wasm"))] +pub(crate) async fn sleep(duration: Duration) { + tokio::time::sleep(duration).await; +} + +#[cfg(target_family = "wasm")] +pub(crate) async fn sleep(duration: Duration) { + gloo_timers::future::sleep(duration).await; +} + +/// `Err(())` if `duration` elapses before `future` completes. +#[cfg(not(target_family = "wasm"))] +pub(crate) async fn timeout(duration: Duration, future: F) -> Result { + tokio::time::timeout(duration, future).await.map_err(|_| ()) +} + +#[cfg(target_family = "wasm")] +pub(crate) async fn timeout(duration: Duration, future: F) -> Result { + use futures_util::future::{Either, select}; + + let timer = gloo_timers::future::sleep(duration); + futures_util::pin_mut!(future, timer); + match select(future, timer).await { + Either::Left((output, _)) => Ok(output), + Either::Right(((), _)) => Err(()), + } +} diff --git a/crates/burn-remote/src/transport/iroh/transfer.rs b/crates/burn-remote/src/transport/iroh/transfer.rs new file mode 100644 index 0000000..ba2f209 --- /dev/null +++ b/crates/burn-remote/src/transport/iroh/transfer.rs @@ -0,0 +1,241 @@ +//! Authenticated tensor transfer over independent Iroh streams. +//! +//! Server-to-server tensor movement rides its own bidirectional stream +//! ([`StreamKind::TensorTransfer`]), separate from the compute session: a tensor is exposed under a +//! [`TransferCapability`] for a specific target peer, and only that peer can download it. + +use std::collections::HashMap; +use std::sync::Arc; + +use burn_backend::TensorData; +use burn_ir::BackendIr; +use tokio::sync::{Mutex, Notify}; + +use super::node::{RemoteNode, StreamKind, recv_frame, send_frame}; +use crate::server::transfer::TensorTransfer; +use crate::shared::TransferCapability; +use crate::{PeerAddr, PeerId}; + +#[derive(Debug, serde::Serialize, serde::Deserialize)] +enum TransferMessage { + Request(TransferCapability), + Tensor(TensorData), + Denied(String), +} + +struct ExposedTensor { + bytes: bytes::Bytes, + target: iroh::EndpointId, + downloads: u32, + max_downloads: u32, +} + +/// Authenticated tensor transfer service carried on independent Iroh streams. +pub(crate) struct IrohTransfer { + node: RemoteNode, + exposed: Arc>>, + exposed_notify: Notify, + _backend: core::marker::PhantomData, +} + +impl IrohTransfer { + pub(crate) fn new(node: RemoteNode) -> Self { + Self { + node, + exposed: Arc::new(Mutex::new(HashMap::new())), + exposed_notify: Notify::new(), + _backend: core::marker::PhantomData, + } + } + + pub(crate) async fn handle_stream( + &self, + remote: iroh::EndpointId, + mut send: iroh::endpoint::SendStream, + mut recv: iroh::endpoint::RecvStream, + ) -> Result<(), String> { + let request = recv_frame(&mut recv) + .await? + .ok_or_else(|| "Tensor-transfer stream closed before its request".to_string())?; + let TransferMessage::Request(capability) = rmp_serde::from_slice(&request) + .map_err(|err| format!("Invalid tensor-transfer request: {err}"))? + else { + return Err("Expected a tensor-transfer request".into()); + }; + + let response = match self.take(capability, remote).await { + Ok(bytes) => bytes, + Err(reason) => rmp_serde::to_vec(&TransferMessage::Denied(reason)) + .map(bytes::Bytes::from) + .map_err(|err| format!("Failed to encode tensor-transfer denial: {err}"))?, + }; + send_frame(&mut send, &response).await?; + send.finish() + .map_err(|err| format!("Failed to finish tensor-transfer stream: {err}"))?; + Ok(()) + } + + async fn take( + &self, + capability: TransferCapability, + remote: iroh::EndpointId, + ) -> Result { + super::time::timeout(TRANSFER_WAIT_TIMEOUT, async { + loop { + let notified = self.exposed_notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + { + let mut exposed = self.exposed.lock().await; + if let Some(mut tensor) = exposed.remove(&capability) { + if tensor.target != remote { + exposed.insert(capability, tensor); + return Err(format!( + "Transfer capability is not authorized for peer {remote}" + )); + } + tensor.downloads += 1; + let bytes = tensor.bytes.clone(); + if tensor.downloads < tensor.max_downloads { + exposed.insert(capability, tensor); + } + return Ok(bytes); + } + } + notified.as_mut().await; + } + }) + .await + .map_err(|_| format!("Timed out waiting for tensor transfer {capability:?}"))? + } + + async fn expose_response( + &self, + bytes: bytes::Bytes, + max_downloads: u32, + capability: TransferCapability, + target: iroh::EndpointId, + ) { + self.exposed.lock().await.insert( + capability, + ExposedTensor { + bytes, + target, + downloads: 0, + max_downloads, + }, + ); + self.exposed_notify.notify_waiters(); + + let exposed = self.exposed.clone(); + crate::server::spawn::spawn_detached(async move { + super::time::sleep(TRANSFER_CAPABILITY_TTL).await; + exposed.lock().await.remove(&capability); + }); + } +} + +const TRANSFER_WAIT_TIMEOUT: core::time::Duration = core::time::Duration::from_secs(300); +const TRANSFER_CAPABILITY_TTL: core::time::Duration = core::time::Duration::from_secs(300); + +impl TensorTransfer for IrohTransfer { + async fn expose_data( + &self, + data: TensorData, + max_downloads: u32, + capability: TransferCapability, + target: PeerId, + ) { + let Some(target) = target.into_iroh_id() else { + log::error!("An Iroh tensor transfer cannot target a non-Iroh peer"); + return; + }; + let bytes = match rmp_serde::to_vec(&TransferMessage::Tensor(data)) { + Ok(bytes) => bytes::Bytes::from(bytes), + Err(err) => { + log::error!("Failed to encode tensor transfer {capability:?}: {err}"); + return; + } + }; + self.expose_response(bytes, max_downloads, capability, target) + .await; + } + + async fn download_tensor( + &self, + remote: PeerAddr, + capability: TransferCapability, + ) -> Option { + match &remote { + PeerAddr::Iroh(_) => {} + #[cfg(feature = "websocket")] + PeerAddr::WebSocket(_) => { + log::error!("An Iroh compute node cannot download from a non-Iroh peer"); + return None; + } + } + let (mut send, mut recv) = match self + .node + .open_stream(&remote, StreamKind::TensorTransfer) + .await + { + Ok(streams) => streams, + Err(err) => { + log::error!("{err}"); + return None; + } + }; + let request = match rmp_serde::to_vec(&TransferMessage::Request(capability)) { + Ok(request) => request, + Err(err) => { + log::error!("Failed to encode tensor-transfer request: {err}"); + return None; + } + }; + if let Err(err) = send_frame(&mut send, &request).await { + log::error!("{err}"); + return None; + } + let _ = send.finish(); + let response = match recv_frame(&mut recv).await { + Ok(Some(response)) => response, + Ok(None) => { + log::error!("Tensor-transfer peer closed without a response"); + return None; + } + Err(err) => { + log::error!("{err}"); + return None; + } + }; + match rmp_serde::from_slice(&response) { + Ok(TransferMessage::Tensor(data)) => Some(data), + Ok(TransferMessage::Denied(reason)) => { + log::error!("Tensor transfer denied: {reason}"); + None + } + Ok(TransferMessage::Request(_)) => { + log::error!("Tensor-transfer peer returned a request instead of tensor data"); + None + } + Err(err) => { + log::error!("Invalid tensor-transfer response: {err}"); + None + } + } + } + + async fn fail(&self, capability: TransferCapability, target: PeerId, reason: String) { + let Some(target) = target.into_iroh_id() else { + return; + }; + let bytes = match rmp_serde::to_vec(&TransferMessage::Denied(reason)) { + Ok(bytes) => bytes.into(), + Err(err) => { + log::error!("Failed to encode tensor-transfer failure: {err}"); + return; + } + }; + self.expose_response(bytes, 1, capability, target).await; + } +} diff --git a/crates/burn-remote/src/transport/link.rs b/crates/burn-remote/src/transport/link.rs new file mode 100644 index 0000000..175b9e7 --- /dev/null +++ b/crates/burn-remote/src/transport/link.rs @@ -0,0 +1,44 @@ +//! The session-link abstraction. +//! +//! A session is a duplex, length-framed message channel: the client submits a stream of +//! [`RemoteMessage`](crate::shared::RemoteMessage)s and the server returns a stream of +//! [`TaskResponse`](crate::shared::TaskResponse)s. Every transport realizes this as one +//! bidirectional stream, split into an outgoing [`FrameSink`] and an incoming [`FrameSource`] so +//! the response-writer task can own the sink while the request-reader loop owns the source. +//! +//! Frames are opaque `Bytes` here; encoding/decoding to the protocol types lives in the session +//! pump (server) and client service, so the transport layer only moves bytes. + +use bytes::Bytes; +use core::future::Future; + +/// `Send` on native targets, unconstrained in the browser. +/// +/// Iroh streams are `!Send` on wasm (they live on the JS event loop), so the link traits cannot +/// require `Send` unconditionally. The real `Send` requirement is applied where session tasks are +/// spawned, via the cfg'd [`spawn_detached`](crate::server::spawn::spawn_detached) / +/// [`Executor`](crate::client::service::Executor) helpers — exactly as the concrete channel enums +/// did before this abstraction existed. +#[cfg(not(target_family = "wasm"))] +pub(crate) trait MaybeSend: Send {} +#[cfg(not(target_family = "wasm"))] +impl MaybeSend for T {} +#[cfg(target_family = "wasm")] +pub(crate) trait MaybeSend {} +#[cfg(target_family = "wasm")] +impl MaybeSend for T {} + +/// The outgoing half of a session link: writes length-framed messages to the peer. +pub(crate) trait FrameSink: MaybeSend + 'static { + /// Send one already-encoded frame. + fn send(&mut self, frame: Bytes) -> impl Future> + MaybeSend; + + /// Finish the stream; no more frames will be sent. + fn close(&mut self) -> impl Future> + MaybeSend; +} + +/// The incoming half of a session link: reads length-framed messages from the peer. +pub(crate) trait FrameSource: MaybeSend + 'static { + /// Receive the next frame, or `None` when the peer closes the stream cleanly. + fn recv(&mut self) -> impl Future, String>> + MaybeSend; +} diff --git a/crates/burn-remote/src/transport/mod.rs b/crates/burn-remote/src/transport/mod.rs new file mode 100644 index 0000000..29c6ae1 --- /dev/null +++ b/crates/burn-remote/src/transport/mod.rs @@ -0,0 +1,18 @@ +//! Transport layer for Burn Remote. +//! +//! `burn-remote` owns the transport seam: identity (`PeerId`/`PeerAddr`) is transport-agnostic and +//! lives here, and each concrete transport (iroh, websocket) is a self-contained submodule that +//! plugs into the session and transfer layers. The rest of the crate is written against these +//! abstractions and stays `cfg`-free; the `cfg(feature = ...)` selection between transports is +//! contained to this module. + +mod identity; +pub use identity::{PeerAddr, PeerId}; + +pub(crate) mod link; + +#[cfg(feature = "iroh")] +pub mod iroh; + +#[cfg(feature = "websocket")] +pub(crate) mod websocket; diff --git a/crates/burn-remote/src/transport/websocket/link.rs b/crates/burn-remote/src/transport/websocket/link.rs new file mode 100644 index 0000000..a85b6d9 --- /dev/null +++ b/crates/burn-remote/src/transport/websocket/link.rs @@ -0,0 +1,59 @@ +//! WebSocket implementations of the session-link frame traits. +//! +//! A WebSocket session is one full-duplex socket; `burn_communication` splits it into independent +//! send/receive halves, which map directly onto [`FrameSink`] / [`FrameSource`]. The inherent +//! `send`/`recv`/`close` on each half do the binary framing; here we only adapt the message type +//! (`Message` ↔ `Bytes`) and the error type (`String`). + +use bytes::Bytes; + +use burn_communication::Message; +use burn_communication::websocket::{WsClientSink, WsClientStream, WsServerSink, WsServerStream}; + +use crate::transport::link::{FrameSink, FrameSource}; + +impl FrameSink for WsServerSink { + async fn send(&mut self, frame: Bytes) -> Result<(), String> { + WsServerSink::send(self, Message::new(frame)) + .await + .map_err(|err| err.to_string()) + } + + async fn close(&mut self) -> Result<(), String> { + WsServerSink::close(self) + .await + .map_err(|err| err.to_string()) + } +} + +impl FrameSource for WsServerStream { + async fn recv(&mut self) -> Result, String> { + WsServerStream::recv(self) + .await + .map(|message| message.map(|message| message.data)) + .map_err(|err| err.to_string()) + } +} + +impl FrameSink for WsClientSink { + async fn send(&mut self, frame: Bytes) -> Result<(), String> { + WsClientSink::send(self, Message::new(frame)) + .await + .map_err(|err| err.to_string()) + } + + async fn close(&mut self) -> Result<(), String> { + WsClientSink::close(self) + .await + .map_err(|err| err.to_string()) + } +} + +impl FrameSource for WsClientStream { + async fn recv(&mut self) -> Result, String> { + WsClientStream::recv(self) + .await + .map(|message| message.map(|message| message.data)) + .map_err(|err| err.to_string()) + } +} diff --git a/crates/burn-remote/src/transport/websocket/mod.rs b/crates/burn-remote/src/transport/websocket/mod.rs new file mode 100644 index 0000000..795df86 --- /dev/null +++ b/crates/burn-remote/src/transport/websocket/mod.rs @@ -0,0 +1,17 @@ +//! WebSocket transport implementation — the legacy address-and-port transport. +//! +//! Self-contained: the session uses one full-duplex socket (split into [`FrameSink`]/[`FrameSource`] +//! halves in [`link`]) driven by the shared session pump, and the turnkey server lives in [`server`]. +//! +//! [`FrameSink`]: crate::transport::link::FrameSink +//! [`FrameSource`]: crate::transport::link::FrameSource + +mod link; + +#[cfg(feature = "server")] +mod transfer; + +#[cfg(not(target_family = "wasm"))] +mod server; +#[cfg(not(target_family = "wasm"))] +pub(crate) use server::start_websocket_async; diff --git a/crates/burn-remote/src/transport/websocket/server.rs b/crates/burn-remote/src/transport/websocket/server.rs new file mode 100644 index 0000000..7eb5c24 --- /dev/null +++ b/crates/burn-remote/src/transport/websocket/server.rs @@ -0,0 +1,69 @@ +//! The turnkey WebSocket compute server. + +use std::sync::Arc; + +use burn_backend::tensor::Device; +use burn_ir::BackendIr; +use burn_router::CustomOpRegistry; +use tokio_util::sync::CancellationToken; + +use burn_communication::{ + ProtocolServer, + external_comm::{ExternalCommServer, ExternalCommService}, + websocket::{WebSocket, WsServer, WsServerChannel}, +}; + +use super::transfer::WebSocketTransfer; +use crate::server::{pump::drive_session, session::SessionManager, spawn::os_shutdown_signal}; + +/// Serve a WebSocket compute node on the given port, until shutdown. +/// +/// The session protocol is a single full-duplex `/session` socket per session (split into a sink + +/// source and driven by the shared [`drive_session`] pump); cross-server tensor transfers ride the +/// same server via [`route_external_comm`](ExternalCommServer::route_external_comm). Driven through +/// [`RemoteServerBuilder`](crate::server::RemoteServerBuilder) rather than called directly. +#[cfg(not(target_family = "wasm"))] +pub(crate) async fn start_websocket_async( + devices: Vec>, + port: u16, + custom_ops: CustomOpRegistry, +) { + let cancel_token = CancellationToken::new(); + let external = Arc::new(ExternalCommService::::new(cancel_token)); + let transfer = Arc::new(WebSocketTransfer { + inner: external.clone(), + }); + let probe = if crate::metrics::TelemetryLogger::enabled() { + crate::telemetry::TelemetryProbe::new(crate::telemetry::CHANNEL_CAPACITY) + } else { + crate::telemetry::TelemetryProbe::disabled() + }; + let sessions = Arc::new( + SessionManager::new(devices, transfer) + .with_custom_ops(custom_ops) + .with_telemetry(probe), + ); + + let server = WsServer::new(port) + .route("/session", { + let sessions = sessions.clone(); + move |channel: WsServerChannel| { + let sessions = sessions.clone(); + async move { + let (sink, source) = channel.split(); + // WebSocket has no authenticated peer identity, so the server presents none + // (`peer_id: None`) and authorizes every session. + if let Err(err) = + drive_session(source, sink, sessions, None, |_init| Ok(())).await + { + log::warn!("WebSocket remote session failed: {err}"); + } + } + } + }) + .route_external_comm(external); + + if let Err(err) = server.serve(os_shutdown_signal()).await { + log::error!("Burn Remote WebSocket server stopped: {err:?}"); + } +} diff --git a/crates/burn-remote/src/transport/websocket/transfer.rs b/crates/burn-remote/src/transport/websocket/transfer.rs new file mode 100644 index 0000000..65929a7 --- /dev/null +++ b/crates/burn-remote/src/transport/websocket/transfer.rs @@ -0,0 +1,74 @@ +//! Legacy WebSocket tensor transfer, carried over the `burn_communication` data service. + +use std::sync::Arc; + +use burn_backend::TensorData; +use burn_ir::BackendIr; + +use crate::server::transfer::TensorTransfer; +use crate::shared::TransferCapability; +use crate::{PeerAddr, PeerId}; + +pub(crate) struct WebSocketTransfer { + pub(crate) inner: Arc< + burn_communication::external_comm::ExternalCommService< + B, + burn_communication::websocket::WebSocket, + >, + >, +} + +impl Clone for WebSocketTransfer { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + } + } +} + +impl TensorTransfer for WebSocketTransfer { + async fn expose_data( + &self, + data: TensorData, + max_downloads: u32, + capability: TransferCapability, + _target: PeerId, + ) { + self.inner + .expose_data(data, max_downloads, capability_to_legacy_id(capability)) + .await; + } + + async fn download_tensor( + &self, + remote: PeerAddr, + capability: TransferCapability, + ) -> Option { + // Only the WebSocket arm remains when the Iroh transport is compiled out. + #[cfg_attr(not(feature = "iroh"), allow(clippy::infallible_destructuring_match))] + let address = match remote { + PeerAddr::WebSocket(address) => address, + #[cfg(feature = "iroh")] + PeerAddr::Iroh(_) => { + log::error!("A WebSocket compute node cannot download from a non-WebSocket peer"); + return None; + } + }; + self.inner + .download_tensor(address, capability_to_legacy_id(capability)) + .await + } + + async fn fail(&self, _capability: TransferCapability, _target: PeerId, reason: String) { + log::error!("Legacy WebSocket tensor transfer failed before exposure: {reason}"); + } +} + +fn capability_to_legacy_id( + capability: TransferCapability, +) -> burn_communication::external_comm::TensorTransferId { + // WebSocket is a compatibility transport without authenticated peer identity. Preserve its old + // transfer service while deriving a collision-resistant-enough rendezvous key from the + // capability. Iroh uses the complete capability and enforces the destination identity. + capability.legacy_id().into() +} diff --git a/crates/burn-remote/tests/iroh.rs b/crates/burn-remote/tests/iroh.rs new file mode 100644 index 0000000..9ff8eea --- /dev/null +++ b/crates/burn-remote/tests/iroh.rs @@ -0,0 +1,210 @@ +#![cfg(all(feature = "client", feature = "server", feature = "iroh"))] + +use burn_flex::Flex; +use burn_ir::BackendIr; +use burn_remote::{ + BURN_REMOTE_ALPN, RemoteDevice, + server::{AllowAll, IrohRemoteProtocol}, + telemetry::TelemetryProbe, +}; +use burn_tensor::{Device, Tensor}; +use iroh::{Endpoint, RelayMode, endpoint::presets, protocol::Router}; + +async fn local_endpoint() -> Endpoint { + Endpoint::builder(presets::Minimal) + .relay_mode(RelayMode::Disabled) + .clear_ip_transports() + .bind_addr("127.0.0.1:0") + .unwrap() + .bind() + .await + .unwrap() +} + +fn spawn_router( + endpoint: Endpoint, + authorizer: impl burn_remote::server::PeerAuthorizer, + probe: TelemetryProbe, +) -> Router { + let protocol = IrohRemoteProtocol::::new( + endpoint.clone(), + vec![Default::default()], + std::sync::Arc::new(authorizer), + probe, + burn_remote::server::CustomOpRegistry::default(), + ); + Router::builder(endpoint) + .accept(BURN_REMOTE_ALPN, protocol) + .spawn() +} + +#[tokio::test(flavor = "multi_thread")] +async fn executes_over_iroh_session_stream() { + let server = local_endpoint().await; + let client = local_endpoint().await; + let router = spawn_router::(server.clone(), AllowAll, TelemetryProbe::disabled()); + + let remote = RemoteDevice::iroh(&client, server.addr(), 0); + remote.connect(); + let device = Device::new(remote); + + let output = Tensor::<1>::from_floats([1.0, 2.0, 3.0], &device) * 2.0; + assert_eq!( + output.to_data().to_vec::().unwrap(), + vec![2.0, 4.0, 6.0] + ); + + router.shutdown().await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn transfers_tensor_directly_between_iroh_compute_peers() { + let source_server = local_endpoint().await; + let target_server = local_endpoint().await; + let client = local_endpoint().await; + + let source_router = + spawn_router::(source_server.clone(), AllowAll, TelemetryProbe::disabled()); + let target_router = + spawn_router::(target_server.clone(), AllowAll, TelemetryProbe::disabled()); + + let source_remote = RemoteDevice::iroh(&client, source_server.addr(), 0); + let target_remote = RemoteDevice::iroh(&client, target_server.addr(), 0); + source_remote.connect(); + target_remote.connect(); + let source = Device::new(source_remote); + let target = Device::new(target_remote); + + let tensor = Tensor::<1>::from_floats([3.0, 5.0, 7.0], &source); + let transferred = tensor.to_device(&target); + assert_eq!( + transferred.to_data().to_vec::().unwrap(), + vec![3.0, 5.0, 7.0] + ); + + source_router.shutdown().await.unwrap(); + target_router.shutdown().await.unwrap(); +} + +/// The synchronous client path used by scripts, REPLs and Rust notebooks: no `async`, no ambient +/// runtime in the calling code. The device is created on the client's runtime (so the session +/// reuses it, the way [`RemoteNode::bind_blocking`] does internally) and every operation is then +/// driven synchronously off it. +#[test] +fn synchronous_client_round_trip() { + // Server on its own local runtime, kept alive for the duration of the test. + let server_runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap(); + let server_guard = server_runtime.enter(); + let server = server_runtime.block_on(local_endpoint()); + let router = spawn_router::(server.clone(), AllowAll, TelemetryProbe::disabled()); + let server_addr = server.addr(); + drop(server_guard); + + // Client on a node that owns its runtime, used entirely synchronously from this (non-runtime) + // thread, exactly what a notebook cell does. + let client_runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap(); + let client_endpoint = client_runtime.block_on(local_endpoint()); + + // Create the device on the client's runtime so the session captures it; the round-trip below + // then runs from this non-runtime thread, exactly what a notebook cell does. + let remote = { + let _guard = client_runtime.enter(); + RemoteDevice::iroh(&client_endpoint, server_addr, 0) + }; + remote.connect(); + let device = Device::new(remote); + + let output = Tensor::<1>::from_floats([1.0, 2.0, 3.0], &device) * 2.0; + assert_eq!( + output.to_data().to_vec::().unwrap(), + vec![2.0, 4.0, 6.0] + ); + + server_runtime.block_on(router.shutdown()).unwrap(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn passes_application_credentials_to_the_peer_authorizer() { + let server = local_endpoint().await; + let client = local_endpoint().await; + let router = spawn_router::( + server.clone(), + |request: burn_remote::server::AuthorizationRequest<'_>| { + (request.credential == b"fleet-ticket") + .then_some(()) + .ok_or_else(|| "invalid fleet ticket".to_string()) + }, + TelemetryProbe::disabled(), + ); + let remote = RemoteDevice::iroh_authorized(&client, server.addr(), 0, b"fleet-ticket".to_vec()); + remote.connect(); + let device = Device::new(remote); + let data = Tensor::<1>::from_floats([4.0], &device).to_data(); + assert_eq!(data.to_vec::().unwrap(), vec![4.0]); + + router.shutdown().await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread")] +#[cfg(feature = "fusion")] +async fn fused_compute_surfaces_as_graph_telemetry() { + use burn_remote::telemetry::{TelemetryEvent, TelemetryProbe, TrafficAggregator}; + use std::time::Duration; + + let server = local_endpoint().await; + let client = local_endpoint().await; + + let (probe, mut events) = TelemetryProbe::channel(4096); + let router = spawn_router::(server.clone(), AllowAll, probe); + let remote = RemoteDevice::iroh(&client, server.addr(), 0); + remote.connect(); + let device = Device::new(remote); + + // A multi-op float expression fuses into a cached graph; running it twice forces a replay, and + // each read flushes the fusion stream so the server actually executes the graph. + for _ in 0..2 { + let x = Tensor::<1>::from_floats([1.0, 2.0, 3.0], &device); + let y = ((x * 2.0) + 1.0).exp().log(); + let _ = y.to_data(); + } + + // Fold the stream the same way a logger or dashboard would, and check the derived economics. + let mut aggregator = TrafficAggregator::default(); + let (mut saw_registered, mut saw_executed) = (false, false); + let collect = async { + while !(saw_registered && saw_executed && aggregator.snapshot().fused_ops > 0) { + let Some(event) = events.recv().await else { + break; + }; + aggregator.apply(&event); + match event.as_ref() { + TelemetryEvent::GraphRegistered { ops, bytes, .. } => { + saw_registered = !ops.is_empty() && *bytes > 0 + } + TelemetryEvent::GraphExecuted { .. } => saw_executed = true, + _ => {} + } + } + }; + tokio::time::timeout(Duration::from_secs(10), collect) + .await + .expect("fused-path telemetry did not arrive in time"); + + assert!( + saw_registered, + "expected a GraphRegistered event carrying the graph's ops and size" + ); + assert!(saw_executed, "expected a GraphExecuted replay heartbeat"); + assert!( + aggregator.snapshot().fused_ops > 0, + "the aggregator should price the replayed graph's ops as fused" + ); + + router.shutdown().await.unwrap(); +} diff --git a/crates/burn-rl/Cargo.toml b/crates/burn-rl/Cargo.toml new file mode 100644 index 0000000..9458a20 --- /dev/null +++ b/crates/burn-rl/Cargo.toml @@ -0,0 +1,29 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science"] +description = "RL crate for the Burn framework" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "tensor", "pytorch", "ndarray"] +license.workspace = true +name = "burn-rl" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-rl" +documentation = "https://docs.rs/burn-rl" +version.workspace = true + +[dependencies] +burn-core = { workspace = true, features = [ + "dataset", + "std", +] } + +derive-new.workspace = true +log = { workspace = true } +rand.workspace = true + +[dev-dependencies] +# Test backend +burn-core = { workspace = true, features = ["default", "flex"]} + +[lints] +workspace = true diff --git a/crates/burn-rl/LICENSE-APACHE b/crates/burn-rl/LICENSE-APACHE new file mode 120000 index 0000000..1cd601d --- /dev/null +++ b/crates/burn-rl/LICENSE-APACHE @@ -0,0 +1 @@ +../../LICENSE-APACHE \ No newline at end of file diff --git a/crates/burn-rl/LICENSE-MIT b/crates/burn-rl/LICENSE-MIT new file mode 120000 index 0000000..b2cfbdc --- /dev/null +++ b/crates/burn-rl/LICENSE-MIT @@ -0,0 +1 @@ +../../LICENSE-MIT \ No newline at end of file diff --git a/crates/burn-rl/README.md b/crates/burn-rl/README.md new file mode 100644 index 0000000..0b0221b --- /dev/null +++ b/crates/burn-rl/README.md @@ -0,0 +1,6 @@ +# Burn RL + + + + diff --git a/crates/burn-rl/src/environment/base.rs b/crates/burn-rl/src/environment/base.rs new file mode 100644 index 0000000..16af869 --- /dev/null +++ b/crates/burn-rl/src/environment/base.rs @@ -0,0 +1,46 @@ +/// The result of taking a step in an environment. +pub struct StepResult { + /// The updated state. + pub next_state: S, + /// The reward. + pub reward: f64, + /// If the environment reached a terminal state. + pub done: bool, + /// If the environment reached its max length. + pub truncated: bool, +} + +/// Trait to be implemented for a RL environment. +pub trait Environment { + /// The type of the state. + type State; + /// The type of actions. + type Action; + + /// The maximum number of step for one episode. + const MAX_STEPS: usize; + + /// Returns the current state. + fn state(&self) -> Self::State; + /// Take a step in the environment given an action. + fn step(&mut self, action: Self::Action) -> StepResult; + /// Reset the environment to an initial state. + fn reset(&mut self); +} + +/// Trait to define how to initialize an environment. +/// By default, any function returning an environment implements it. +pub trait EnvironmentInit: Clone { + /// Initialize the environment. + fn init(&self) -> E; +} + +impl EnvironmentInit for F +where + F: Fn() -> E + Clone, + E: Environment, +{ + fn init(&self) -> E { + (self)() + } +} diff --git a/crates/burn-rl/src/environment/mod.rs b/crates/burn-rl/src/environment/mod.rs new file mode 100644 index 0000000..cbcb6ac --- /dev/null +++ b/crates/burn-rl/src/environment/mod.rs @@ -0,0 +1,3 @@ +mod base; + +pub use base::*; diff --git a/crates/burn-rl/src/lib.rs b/crates/burn-rl/src/lib.rs new file mode 100644 index 0000000..a02018d --- /dev/null +++ b/crates/burn-rl/src/lib.rs @@ -0,0 +1,163 @@ +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] + +//! A library for training reinforcement learning agents. + +/// Module for implementing an environment. +pub mod environment; +/// Module for implementing a policy. +pub mod policy; +/// Transition buffer. +pub mod transition_buffer; + +pub use environment::*; +pub use policy::*; +pub use transition_buffer::*; + +#[cfg(test)] +pub(crate) mod tests { + use crate::{Batchable, Policy, PolicyState}; + + use burn_core::tensor::Device; + + /// Mock policy for testing + /// + /// Calling `forward()` with a [MockObservation](MockObservation) (list of f32) returns a [MockActionDistribution](MockActionDistribution) + /// containing a list of 0s of the same length as the observation. + /// + /// Calling `action()` with a [MockObservation](MockObservation) (list of f32) returns a [MockAction](MockAction) with a list of actions of the same length as the observation. + /// The actions are all 1 if the call is requested as deterministic, or else 0. + #[derive(Clone)] + pub(crate) struct MockPolicy {} + + impl MockPolicy { + pub fn new() -> Self { + Self {} + } + } + + impl Policy for MockPolicy { + type Observation = MockObservation; + type ActionDistribution = MockActionDistribution; + type Action = MockAction; + type ActionContext = MockActionContext; + type PolicyState = MockPolicyState; + + fn forward(&mut self, obs: Self::Observation) -> Self::ActionDistribution { + let mut dists = vec![]; + + for _ in obs.0 { + dists.push(MockActionDistribution(vec![0.])); + } + MockActionDistribution::batch(dists) + } + + fn action( + &mut self, + obs: Self::Observation, + deterministic: bool, + ) -> (Self::Action, Vec) { + let mut actions = vec![]; + let mut contexts = vec![]; + + for _ in obs.0 { + if deterministic { + actions.push(MockAction(vec![1])); + } else { + actions.push(MockAction(vec![0])); + } + contexts.push(MockActionContext); + } + + (MockAction::batch(actions), contexts) + } + + fn update(&mut self, _update: Self::PolicyState) {} + + fn state(&self) -> Self::PolicyState { + MockPolicyState + } + + fn load_record(self, _record: ::Record) -> Self { + self + } + + fn to_device(self, _device: &Device) -> Self { + self + } + } + + /// Mock observation for testing represented as a vector of f32. Can call `batch()` and `unbatch` on it. + #[derive(Clone)] + pub(crate) struct MockObservation(pub Vec); + + /// Mock action for testing represented as a vector of i32. Can call `batch()` and `unbatch` on it. + #[derive(Clone)] + pub(crate) struct MockAction(pub Vec); + + /// Mock action distribution for testing represented as a vector of i32. Can call `batch()` and `unbatch` on it. + #[derive(Clone)] + pub(crate) struct MockActionDistribution(Vec); + + #[derive(Clone)] + pub(crate) struct MockActionContext; + + /// Mock policy state for testing represented as an arbitrary `usize` that has no effect on the policy. + #[derive(Clone)] + pub(crate) struct MockPolicyState; + + #[derive(Clone)] + pub(crate) struct MockRecord { + _item: usize, + } + + impl PolicyState for MockPolicyState { + type Record = MockRecord; + + fn into_record(self) -> Self::Record { + MockRecord { _item: 0 } + } + + fn load_record(&self, _record: Self::Record) -> Self { + self.clone() + } + } + + impl Batchable for MockObservation { + fn batch(items: Vec) -> Self { + MockObservation(items.iter().flat_map(|m| m.0.clone()).collect()) + } + + fn unbatch(self) -> Vec { + vec![MockObservation(self.0)] + } + } + + impl Batchable for MockAction { + fn batch(items: Vec) -> Self { + MockAction(items.iter().flat_map(|m| m.0.clone()).collect()) + } + + fn unbatch(self) -> Vec { + let mut actions = vec![]; + for a in self.0 { + actions.push(MockAction(vec![a])); + } + actions + } + } + + impl Batchable for MockActionDistribution { + fn batch(items: Vec) -> Self { + MockActionDistribution(items.iter().flat_map(|m| m.0.clone()).collect()) + } + + fn unbatch(self) -> Vec { + let mut dists = vec![]; + for _ in self.0 { + dists.push(MockActionDistribution(vec![0.])); + } + dists + } + } +} diff --git a/crates/burn-rl/src/policy/async_policy.rs b/crates/burn-rl/src/policy/async_policy.rs new file mode 100644 index 0000000..f6a4ec3 --- /dev/null +++ b/crates/burn-rl/src/policy/async_policy.rs @@ -0,0 +1,486 @@ +use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + mpsc::{self, Sender}, + }, + thread::spawn, +}; + +use burn_core::tensor::Device; + +use crate::{ActionContext, Batchable, Policy, PolicyState}; + +#[derive(Clone)] +struct PolicyInferenceServer { + // `num_agents` used to make sure autobatching doesn't block the agents if they are less than the autobatch size. + num_agents: Arc, + max_autobatch_size: usize, + inner_policy: P, + batch_action: Vec>, + batch_logits: Vec>, +} + +impl

PolicyInferenceServer

+where + P: Policy, + P::Observation: Clone + Batchable, + P::ActionDistribution: Clone + Batchable, + P::Action: Clone + Batchable, + P::ActionContext: Clone, +{ + pub fn new(max_autobatch_size: usize, inner_policy: P) -> Self { + Self { + num_agents: Arc::new(AtomicUsize::new(0)), + max_autobatch_size, + inner_policy, + batch_action: vec![], + batch_logits: vec![], + } + } + + pub fn push_action(&mut self, item: ActionItem) { + self.batch_action.push(item); + if self.len_actions() + >= self + .num_agents + .load(Ordering::Relaxed) + .min(self.max_autobatch_size) + { + self.flush_actions(); + } + } + + pub fn push_logits(&mut self, item: ForwardItem) { + self.batch_logits.push(item); + if self.len_logits() + >= self + .num_agents + .load(Ordering::Relaxed) + .min(self.max_autobatch_size) + { + self.flush_logits(); + } + } + + pub fn len_actions(&self) -> usize { + self.batch_action.len() + } + + pub fn len_logits(&self) -> usize { + self.batch_logits.len() + } + + pub fn flush_actions(&mut self) { + if self.len_actions() == 0 { + return; + } + let input: Vec<_> = self + .batch_action + .iter() + .map(|m| m.inference_state.clone()) + .collect(); + // Only deterministic if all actions are requested as deterministic. + let deterministic = self.batch_action.iter().all(|item| item.deterministic); + let (actions, context) = self + .inner_policy + .action(P::Observation::batch(input), deterministic); + let actions: Vec<_> = actions.unbatch(); + + for (i, item) in self.batch_action.iter().enumerate() { + item.sender + .send(ActionContext { + context: vec![context[i].clone()], + action: actions[i].clone(), + }) + .expect("Autobatcher should be able to send resulting actions."); + } + self.batch_action.clear(); + } + + pub fn flush_logits(&mut self) { + if self.len_logits() == 0 { + return; + } + let input: Vec<_> = self + .batch_logits + .iter() + .map(|m| m.inference_state.clone()) + .collect(); + let output = self.inner_policy.forward(P::Observation::batch(input)); + let logits: Vec<_> = output.unbatch(); + for (i, item) in self.batch_logits.iter().enumerate() { + item.sender + .send(logits[i].clone()) + .expect("Autobatcher should be able to send resulting probabilities."); + } + self.batch_logits.clear(); + } + + pub fn update_policy(&mut self, policy_update: P::PolicyState) { + if self.len_actions() > 0 { + self.flush_actions(); + } + if self.len_logits() > 0 { + self.flush_logits(); + } + self.inner_policy.update(policy_update); + } + + pub fn policy_to_device(&mut self, device: &Device) { + self.inner_policy = self.inner_policy.clone().to_device(device); + } + + pub fn state(&self) -> P::PolicyState { + self.inner_policy.state() + } + + pub fn increment_agents(&mut self, num: usize) { + self.num_agents.fetch_add(num, Ordering::Relaxed); + } + + pub fn decrement_agents(&mut self, num: usize) { + self.num_agents.fetch_sub(num, Ordering::Relaxed); + if self.len_actions() + >= self + .num_agents + .load(Ordering::Relaxed) + .min(self.max_autobatch_size) + { + self.flush_actions(); + } + if self.len_logits() + >= self + .num_agents + .load(Ordering::Relaxed) + .min(self.max_autobatch_size) + { + self.flush_logits(); + } + } +} + +enum InferenceMessage { + ActionMessage(ActionItem), + ForwardMessage(ForwardItem), + PolicyUpdate(P::PolicyState), + ToDevice(Device), + PolicyRequest(Sender), + IncrementAgents(usize), + DecrementAgents(usize), +} + +#[derive(Clone)] +struct ActionItem { + sender: Sender>>, + inference_state: S, + deterministic: bool, +} + +#[derive(Clone)] +struct ForwardItem { + sender: Sender, + inference_state: S, +} + +/// An asynchronous policy using an inference server with autobatching. +#[derive(Clone)] +pub struct AsyncPolicy { + inference_state_sender: Sender>, +} + +impl

AsyncPolicy

+where + P: Policy + Clone + Send + 'static, + P::ActionContext: Clone + Send, + P::PolicyState: Send, + P::Observation: Clone + Send + Batchable, + P::ActionDistribution: Clone + Send + Batchable, + P::Action: Clone + Send + Batchable, +{ + /// Create the policy. + /// + /// # Arguments + /// + /// * `autobatch_size` - Number of observations to accumulate before running a pass of inference. + /// * `inner_policy` - The policy used to take actions. + pub fn new(autobatch_size: usize, inner_policy: P) -> Self { + let (sender, receiver) = std::sync::mpsc::channel(); + let mut autobatcher = PolicyInferenceServer::new(autobatch_size, inner_policy.clone()); + spawn(move || { + loop { + match receiver.recv() { + Ok(msg) => match msg { + InferenceMessage::ActionMessage(item) => autobatcher.push_action(item), + InferenceMessage::ForwardMessage(item) => autobatcher.push_logits(item), + InferenceMessage::PolicyUpdate(update) => autobatcher.update_policy(update), + InferenceMessage::ToDevice(device) => autobatcher.policy_to_device(&device), + InferenceMessage::PolicyRequest(sender) => sender + .send(autobatcher.state()) + .expect("Autobatcher should be able to send current policy state."), + InferenceMessage::IncrementAgents(num) => autobatcher.increment_agents(num), + InferenceMessage::DecrementAgents(num) => autobatcher.decrement_agents(num), + }, + Err(err) => { + log::error!("Error in AsyncPolicy : {}", err); + break; + } + } + } + }); + + Self { + inference_state_sender: sender, + } + } + + /// Increment the number of agents using the inference server. + pub fn increment_agents(&self, num: usize) { + self.inference_state_sender + .send(InferenceMessage::IncrementAgents(num)) + .expect("Can send message to autobatcher.") + } + + /// Decrement the number of agents using the inference server. + pub fn decrement_agents(&self, num: usize) { + self.inference_state_sender + .send(InferenceMessage::DecrementAgents(num)) + .expect("Can send message to autobatcher.") + } +} + +impl

Policy for AsyncPolicy

+where + P: Policy + Send + 'static, +{ + type ActionContext = P::ActionContext; + type PolicyState = P::PolicyState; + + type Observation = P::Observation; + type ActionDistribution = P::ActionDistribution; + type Action = P::Action; + + fn forward(&mut self, states: Self::Observation) -> Self::ActionDistribution { + let (action_sender, action_receiver) = std::sync::mpsc::channel(); + let item = ForwardItem { + sender: action_sender, + inference_state: states, + }; + self.inference_state_sender + .send(InferenceMessage::ForwardMessage(item)) + .expect("Should be able to send message to inference_server"); + action_receiver + .recv() + .expect("AsyncPolicy should receive queued probabilities.") + } + + fn action( + &mut self, + states: Self::Observation, + deterministic: bool, + ) -> (Self::Action, Vec) { + let (action_sender, action_receiver) = std::sync::mpsc::channel(); + let item = ActionItem { + sender: action_sender, + inference_state: states, + deterministic, + }; + self.inference_state_sender + .send(InferenceMessage::ActionMessage(item)) + .expect("should be able to send message to inference_server."); + let action = action_receiver + .recv() + .expect("AsyncPolicy should receive queued actions."); + (action.action, action.context) + } + + fn update(&mut self, update: Self::PolicyState) { + self.inference_state_sender + .send(InferenceMessage::PolicyUpdate(update)) + .expect("AsyncPolicy should be able to send policy state.") + } + + fn state(&self) -> Self::PolicyState { + let (sender, receiver) = mpsc::channel(); + self.inference_state_sender + .send(InferenceMessage::PolicyRequest(sender)) + .expect("should be able to send message to inference_server."); + receiver + .recv() + .expect("AsyncPolicy should be able to receive policy state.") + } + + fn to_device(self, device: &Device) -> Self { + self.inference_state_sender + .send(InferenceMessage::ToDevice(device.clone())) + .expect("AsyncPolicy should be able to send policy state."); + self + } + + fn load_record(self, _record: ::Record) -> Self { + unimplemented!( + "Not implemented yet. Please load the record on the inner policy before creating an async policy." + ) + } +} + +#[cfg(test)] +#[allow(clippy::needless_range_loop)] +mod tests { + use std::thread::JoinHandle; + use std::time::Duration; + + use crate::tests::{MockAction, MockObservation, MockPolicy}; + + use super::*; + + #[test] + fn test_multiple_actions_before_flush() { + fn launch_thread(policy: &AsyncPolicy, handles: &mut Vec>) { + let mut thread_policy = policy.clone(); + let handle = spawn(move || { + thread_policy.action(MockObservation(vec![0.]), false); + }); + handles.push(handle); + } + + let policy = AsyncPolicy::new(8, MockPolicy::new()); + policy.increment_agents(1000); + + let mut handles = vec![]; + launch_thread(&policy, &mut handles); + std::thread::sleep(Duration::from_millis(10)); + assert!(!handles[0].is_finished()); + + for _ in 0..6 { + launch_thread(&policy, &mut handles); + } + std::thread::sleep(Duration::from_millis(10)); + for i in 0..7 { + assert!(!handles[i].is_finished()); + } + + launch_thread(&policy, &mut handles); + std::thread::sleep(Duration::from_millis(10)); + for i in 0..8 { + assert!(handles[i].is_finished()); + } + + let mut handles = vec![]; + launch_thread(&policy, &mut handles); + std::thread::sleep(Duration::from_millis(10)); + assert!(!handles[0].is_finished()); + } + + #[test] + fn test_multiple_forward_before_flush() { + fn launch_thread(policy: &AsyncPolicy, handles: &mut Vec>) { + let mut thread_policy = policy.clone(); + let handle = spawn(move || { + thread_policy.forward(MockObservation(vec![0.])); + }); + handles.push(handle); + } + + let policy = AsyncPolicy::new(8, MockPolicy::new()); + policy.increment_agents(1000); + + let mut handles = vec![]; + launch_thread(&policy, &mut handles); + std::thread::sleep(Duration::from_millis(10)); + assert!(!handles[0].is_finished()); + + for _ in 0..6 { + launch_thread(&policy, &mut handles); + } + std::thread::sleep(Duration::from_millis(10)); + for i in 0..7 { + assert!(!handles[i].is_finished()); + } + + launch_thread(&policy, &mut handles); + std::thread::sleep(Duration::from_millis(10)); + for i in 0..8 { + assert!(handles[i].is_finished()); + } + + let mut handles = vec![]; + launch_thread(&policy, &mut handles); + std::thread::sleep(Duration::from_millis(10)); + assert!(!handles[0].is_finished()); + } + + #[test] + fn test_async_policy_deterministic_behaviour() { + fn launch_thread( + policy: &AsyncPolicy, + handles: &mut Vec>, + deterministic: bool, + ) { + let mut thread_policy = policy.clone(); + let handle = spawn(move || { + let (action, _) = thread_policy.action(MockObservation(vec![0.]), deterministic); + action + }); + handles.push(handle); + } + + let policy = AsyncPolicy::new(2, MockPolicy::new()); + policy.increment_agents(1000); + + let mut handles = vec![]; + launch_thread(&policy, &mut handles, true); + launch_thread(&policy, &mut handles, false); + for _ in 0..2 { + let action = handles.pop().unwrap().join().unwrap(); + assert_eq!(action.0, vec![0]); + } + + let mut handles = vec![]; + launch_thread(&policy, &mut handles, true); + launch_thread(&policy, &mut handles, true); + for _ in 0..2 { + let action = handles.pop().unwrap().join().unwrap(); + assert_eq!(action.0, vec![1]); + } + } + + #[test] + fn flush_when_running_agents_smaller_than_autobatch_size() { + fn launch_thread(policy: &AsyncPolicy, handles: &mut Vec>) { + let mut thread_policy = policy.clone(); + let handle = spawn(move || { + thread_policy.action(MockObservation(vec![0.]), false); + }); + handles.push(handle); + } + + let policy = AsyncPolicy::new(8, MockPolicy::new()); + policy.increment_agents(3); + + let mut handles = vec![]; + launch_thread(&policy, &mut handles); + launch_thread(&policy, &mut handles); + std::thread::sleep(Duration::from_millis(10)); + assert!(!handles[0].is_finished()); + assert!(!handles[1].is_finished()); + + launch_thread(&policy, &mut handles); + std::thread::sleep(Duration::from_millis(10)); + for i in 0..3 { + assert!(handles[i].is_finished()); + } + + let mut handles = vec![]; + launch_thread(&policy, &mut handles); + launch_thread(&policy, &mut handles); + std::thread::sleep(Duration::from_millis(10)); + assert!(!handles[0].is_finished()); + assert!(!handles[1].is_finished()); + + policy.decrement_agents(1); + std::thread::sleep(Duration::from_millis(10)); + assert!(handles[0].is_finished()); + assert!(handles[1].is_finished()); + } +} diff --git a/crates/burn-rl/src/policy/base.rs b/crates/burn-rl/src/policy/base.rs new file mode 100644 index 0000000..dd16c08 --- /dev/null +++ b/crates/burn-rl/src/policy/base.rs @@ -0,0 +1,122 @@ +use derive_new::new; + +use burn_core::tensor::Device; + +use crate::TransitionBatch; + +/// An action along with additional context about the decision. +#[derive(Clone, new)] +pub struct ActionContext { + /// The context. + pub context: C, + /// The action. + pub action: A, +} + +/// The state of a policy. +pub trait PolicyState { + /// The type of the record. + type Record; + + /// Convert the state to a record. + fn into_record(self) -> Self::Record; + /// Load the state from a record. + fn load_record(&self, record: Self::Record) -> Self; +} + +/// Defines how an environment's state is converted to a policy's observation. +pub trait ToObservation { + /// Convert an environment's state to a policy's observation, moving it to the given device if needed. + fn to_observation(&self, device: &Device) -> O; +} + +/// Defines how an environment's action is converted to a policy's action. +pub trait ToAction { + /// Convert an environment's action to a policy's action, moving it to the given device if needed. + fn to_action(&self, device: &Device) -> A; +} + +/// Trait for a RL policy. +pub trait Policy: Clone { + /// The observation given as input to the policy. + type Observation; + /// The action distribution parameters defining how the action will be sampled. + type ActionDistribution; + /// The action. + type Action; + + /// Additional context on the policy's decision. + type ActionContext; + /// The current parameterization of the policy. + type PolicyState: PolicyState; + + /// Produces the action distribution from a batch of observations. + fn forward(&mut self, obs: Self::Observation) -> Self::ActionDistribution; + /// Gives the action from a batch of observations. + fn action( + &mut self, + obs: Self::Observation, + deterministic: bool, + ) -> (Self::Action, Vec); + + /// Update the policy's parameters. + fn update(&mut self, update: Self::PolicyState); + /// Returns the current parameterization. + fn state(&self) -> Self::PolicyState; + + /// Loads the policy on the given device. + fn to_device(self, device: &Device) -> Self; + /// Loads the policy parameters from a record. + fn load_record(self, record: ::Record) -> Self; +} + +/// Trait for a type that can be batched and unbatched (split). +pub trait Batchable: Sized { + /// Create a batch from a list of items. + fn batch(value: Vec) -> Self; + /// Create a list from batched items. + fn unbatch(self) -> Vec; +} + +/// A training output. +pub struct RLTrainOutput { + /// The policy. + pub policy: P, + /// The item. + pub item: TO, +} + +/// Batched transitions for a PolicyLearner. +pub type LearnerTransitionBatch

= + TransitionBatch<

::Observation,

::Action>; + +/// Learner for a policy. +pub trait PolicyLearner +where + ::Observation: Clone + Batchable, + ::ActionDistribution: Clone + Batchable, + ::Action: Clone + Batchable, +{ + /// Additional context of a training step. + type TrainContext; + /// The policy to train. + type InnerPolicy: Policy; + /// The record of the learner. + type Record; + + /// Execute a training step on the policy. + fn train( + &mut self, + input: LearnerTransitionBatch, + ) -> RLTrainOutput::PolicyState>; + /// Returns the learner's current policy for validation. + fn policy(&self) -> Self::InnerPolicy; + /// Update the learner's policy. + fn update_policy(&mut self, update: Self::InnerPolicy); + /// Convert the learner's state into a record. + fn record(&self) -> Self::Record; + /// Load the learner's state from a record. + fn load_record(self, record: Self::Record) -> Self; + /// Returns the device used for training. + fn device(&self) -> Device; +} diff --git a/crates/burn-rl/src/policy/mod.rs b/crates/burn-rl/src/policy/mod.rs new file mode 100644 index 0000000..fa7e967 --- /dev/null +++ b/crates/burn-rl/src/policy/mod.rs @@ -0,0 +1,5 @@ +mod async_policy; +mod base; + +pub use async_policy::*; +pub use base::*; diff --git a/crates/burn-rl/src/transition_buffer/base.rs b/crates/burn-rl/src/transition_buffer/base.rs new file mode 100644 index 0000000..01a386e --- /dev/null +++ b/crates/burn-rl/src/transition_buffer/base.rs @@ -0,0 +1,239 @@ +use burn_core::{Tensor, prelude::Device, tensor::Distribution}; +use derive_new::new; + +use super::SliceAccess; + +/// A state transition in an environment. +#[derive(Clone, new)] +pub struct Transition { + /// The initial state. + pub state: S, + /// The state after the step was taken. + pub next_state: S, + /// The action taken in the step. + pub action: A, + /// The reward. + pub reward: Tensor<1>, + /// If the environment has reached a terminal state. + pub done: Tensor<1>, +} + +/// A batch of transitions. +pub struct TransitionBatch { + /// Batched initial states. + pub states: SB, + /// Batched resulting states. + pub next_states: SB, + /// Batched actions. + pub actions: AB, + /// Batched rewards. + pub rewards: Tensor<2>, + /// Batched flags for terminal states. + pub dones: Tensor<2>, +} + +/// A tensor-backed circular buffer for transitions. +/// +/// Uses [`SliceAccess`] to store state and action batches in contiguous +/// tensor storage, enabling efficient random sampling via `select`. +/// The buffer lazily initializes its storage on the first `push` call. +pub struct TransitionBuffer { + states: Option, + next_states: Option, + actions: Option, + rewards: Option>, + dones: Option>, + capacity: usize, + write_head: usize, + len: usize, + device: Device, +} + +impl TransitionBuffer { + /// Creates a new buffer. Storage is lazily allocated on the first `push`. + pub fn new(capacity: usize, device: &Device) -> Self { + Self { + states: None, + next_states: None, + actions: None, + rewards: None, + dones: None, + capacity, + write_head: 0, + len: 0, + device: device.clone(), + } + } + + fn ensure_init(&mut self, state: &SB, next_state: &SB, action: &AB) { + if self.states.is_none() { + self.states = Some(SB::zeros_like(state, self.capacity, &self.device)); + self.next_states = Some(SB::zeros_like(next_state, self.capacity, &self.device)); + self.actions = Some(AB::zeros_like(action, self.capacity, &self.device)); + self.rewards = Some(Tensor::zeros([self.capacity, 1], &self.device)); + self.dones = Some(Tensor::zeros([self.capacity, 1], &self.device)); + } + } + + /// Add a transition, overwriting the oldest if full. + pub fn push(&mut self, state: SB, next_state: SB, action: AB, reward: f32, done: bool) { + self.ensure_init(&state, &next_state, &action); + + let idx = self.write_head % self.capacity; + + self.states + .as_mut() + .unwrap() + .slice_assign_inplace(idx, state); + self.next_states + .as_mut() + .unwrap() + .slice_assign_inplace(idx, next_state); + self.actions + .as_mut() + .unwrap() + .slice_assign_inplace(idx, action); + + let reward = Tensor::from_data([[reward]], &self.device); + self.rewards + .as_mut() + .unwrap() + .inplace(|r| r.slice_assign(idx..idx + 1, reward)); + + let done_val = if done { 1.0f32 } else { 0.0 }; + let done = Tensor::from_data([[done_val]], &self.device); + self.dones + .as_mut() + .unwrap() + .inplace(|d| d.slice_assign(idx..idx + 1, done)); + + self.write_head += 1; + if self.len < self.capacity { + self.len += 1; + } + } + + /// Sample a random batch of transitions. + pub fn sample(&self, batch_size: usize) -> TransitionBatch { + assert!(batch_size <= self.len, "batch_size exceeds buffer length"); + + let indices = Tensor::<1>::random( + [batch_size], + Distribution::Uniform(0.0, self.len as f64), + &self.device, + ) + .int(); + + TransitionBatch { + states: self + .states + .as_ref() + .unwrap() + .clone() + .select(0, indices.clone()), + next_states: self + .next_states + .as_ref() + .unwrap() + .clone() + .select(0, indices.clone()), + actions: self + .actions + .as_ref() + .unwrap() + .clone() + .select(0, indices.clone()), + rewards: self + .rewards + .as_ref() + .unwrap() + .clone() + .select(0, indices.clone()), + dones: self.dones.as_ref().unwrap().clone().select(0, indices), + } + } + + /// Current number of stored transitions. + pub fn len(&self) -> usize { + self.len + } + + /// Whether the buffer is empty. + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// Buffer capacity. + pub fn capacity(&self) -> usize { + self.capacity + } +} + +#[cfg(test)] +mod tests { + use super::*; + + type TB = Tensor<2>; + + fn push_transition(buffer: &mut TransitionBuffer, device: &Device, val: f32) { + let state = Tensor::<2>::from_data([[val, val]], device); + let next_state = Tensor::<2>::from_data([[val + 1.0, val + 1.0]], device); + let action = Tensor::<2>::from_data([[val]], device); + buffer.push(state, next_state, action, val, false); + } + + #[test] + fn push_increment_len() { + let device = Default::default(); + let mut buffer = TransitionBuffer::::new(5, &device); + + assert_eq!(buffer.len(), 0); + assert!(buffer.is_empty()); + + push_transition(&mut buffer, &device, 1.0); + assert_eq!(buffer.len(), 1); + + push_transition(&mut buffer, &device, 2.0); + assert_eq!(buffer.len(), 2); + } + + #[test] + fn push_overwrites_when_full() { + let device = Default::default(); + let mut buffer = TransitionBuffer::::new(3, &device); + + for i in 0..5 { + push_transition(&mut buffer, &device, i as f32); + } + + assert_eq!(buffer.len(), 3); + assert_eq!(buffer.capacity(), 3); + } + + #[test] + fn sample_returns_correct_shapes() { + let device = Default::default(); + let mut buffer = TransitionBuffer::::new(10, &device); + + for i in 0..5 { + push_transition(&mut buffer, &device, i as f32); + } + + let batch = buffer.sample(3); + assert_eq!(batch.states.dims(), [3, 2]); + assert_eq!(batch.next_states.dims(), [3, 2]); + assert_eq!(batch.actions.dims(), [3, 1]); + assert_eq!(batch.rewards.dims(), [3, 1]); + assert_eq!(batch.dones.dims(), [3, 1]); + } + + #[test] + #[should_panic(expected = "batch_size exceeds buffer length")] + fn sample_panics_when_batch_too_large() { + let device = Default::default(); + let mut buffer = TransitionBuffer::::new(5, &device); + + push_transition(&mut buffer, &device, 1.0); + buffer.sample(5); + } +} diff --git a/crates/burn-rl/src/transition_buffer/mod.rs b/crates/burn-rl/src/transition_buffer/mod.rs new file mode 100644 index 0000000..806834d --- /dev/null +++ b/crates/burn-rl/src/transition_buffer/mod.rs @@ -0,0 +1,5 @@ +mod base; +mod slice_access; + +pub use base::*; +pub use slice_access::*; diff --git a/crates/burn-rl/src/transition_buffer/slice_access.rs b/crates/burn-rl/src/transition_buffer/slice_access.rs new file mode 100644 index 0000000..2d3b0a8 --- /dev/null +++ b/crates/burn-rl/src/transition_buffer/slice_access.rs @@ -0,0 +1,36 @@ +use burn_core::prelude::*; + +/// Trait for types that support tensor-like slice operations, +/// enabling storage in a [`TransitionBuffer`](super::TransitionBuffer). +/// +/// Implement this trait for any type that wraps tensors and can be stored +/// in a replay buffer. The buffer uses these operations for: +/// - Pre-allocating storage (`zeros_like`) +/// - Writing transitions (`slice_assign_inplace`) +/// - Sampling batches (`select`) +pub trait SliceAccess: Clone + Sized { + /// Create zeroed storage matching the shape of `sample` but with `capacity` rows + /// along the first dimension. + fn zeros_like(sample: &Self, capacity: usize, device: &Device) -> Self; + + /// Select rows at the given indices along the specified dimension. + fn select(self, dim: usize, indices: Tensor<1, Int>) -> Self; + + /// Assign `value` at row `index` along the first dimension, in place. + fn slice_assign_inplace(&mut self, index: usize, value: Self); +} + +impl SliceAccess for Tensor<2> { + fn zeros_like(sample: &Self, capacity: usize, device: &Device) -> Self { + let feature_dim = sample.dims()[1]; + Tensor::zeros([capacity, feature_dim], device) + } + + fn select(self, dim: usize, indices: Tensor<1, Int>) -> Self { + Tensor::select(self, dim, indices) + } + + fn slice_assign_inplace(&mut self, index: usize, value: Self) { + self.inplace(|t| t.slice_assign(index..index + 1, value)); + } +} diff --git a/crates/burn-rocm/Cargo.toml b/crates/burn-rocm/Cargo.toml new file mode 100644 index 0000000..ebd44a4 --- /dev/null +++ b/crates/burn-rocm/Cargo.toml @@ -0,0 +1,44 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science"] +description = "ROCm HIP backend for the Burn framework" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "gpu", "rocm", "hip"] +license.workspace = true +name = "burn-rocm" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-rocm" +documentation = "https://docs.rs/burn-rocm" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["fusion", "burn-cubecl/default", "cubecl/default"] +doc = ["burn-cubecl/doc"] +std = ["burn-cubecl/std", "cubecl/std"] + +tracing = [ + "cubecl/tracing", + "burn-cubecl/tracing", + "burn-backend/tracing", + "burn-fusion?/tracing", +] + +fusion = ["burn-fusion", "burn-cubecl/fusion"] +autotune = ["burn-cubecl/autotune"] +autotune-checks = ["burn-cubecl/autotune-checks"] + +[dependencies] +cubecl = { workspace = true, features = ["hip"] } +burn-cubecl = { workspace = true, features = ["default"] } +burn-backend = { workspace = true, features = [ + "default", + "cubecl-hip", +] } +burn-fusion = { workspace = true, optional = true, features = ["default"] } + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-rocm/README.md b/crates/burn-rocm/README.md new file mode 100644 index 0000000..583bbec --- /dev/null +++ b/crates/burn-rocm/README.md @@ -0,0 +1,7 @@ +# burn-rocm + +Backend using ROCm HIP runtime. + +To execute the tests for this backend set an environment variable called `ROCM_PATH` or `CUBECL_ROCM_PATH` to the installation path of ROCm. It is often `/opt/rocm`. + +For now this backend requires the version `6.2.2` of ROCm or a compatible version. diff --git a/crates/burn-rocm/src/lib.rs b/crates/burn-rocm/src/lib.rs new file mode 100644 index 0000000..2b7e562 --- /dev/null +++ b/crates/burn-rocm/src/lib.rs @@ -0,0 +1,14 @@ +#![cfg_attr(docsrs, feature(doc_cfg))] +extern crate alloc; + +use burn_cubecl::CubeBackend; + +pub use cubecl::hip::AmdDevice as RocmDevice; + +use cubecl::hip::HipRuntime; + +#[cfg(not(feature = "fusion"))] +pub type Rocm = CubeBackend; + +#[cfg(feature = "fusion")] +pub type Rocm = burn_fusion::Fusion>; diff --git a/crates/burn-router/Cargo.toml b/crates/burn-router/Cargo.toml new file mode 100644 index 0000000..94b95c4 --- /dev/null +++ b/crates/burn-router/Cargo.toml @@ -0,0 +1,53 @@ +[package] +authors = [ + "laggui ", + "nathanielsimard ", +] +categories = ["science"] +description = "Multi-backend router decorator for the Burn framework" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "data"] +license.workspace = true +name = "burn-router" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-router" +documentation = "https://docs.rs/burn-router" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["std"] +std = ["burn-backend/std", "burn-std/std", "burn-ir/std"] +# Client-side operation fusion / caching for router backends (e.g. the remote backend). Requires +# std (the fusion engine uses threads). +fusion = ["std", "dep:burn-fusion", "dep:serde"] +doc = ["default"] +tracing = [ + "burn-backend/tracing", + "burn-ir/tracing", + "burn-std/tracing", +] + +[dependencies] +burn-ir = { workspace = true } +burn-backend = { workspace = true } +burn-fusion = { workspace = true, optional = true } +burn-std = { workspace = true } +hashbrown = { workspace = true } +spin = { workspace = true } +log = { workspace = true } +serde = { workspace = true, optional = true } + +[dev-dependencies] +burn-tensor = { workspace = true } +burn-flex = { workspace = true, features = ["default"] } +burn-wgpu = { workspace = true, features = [ + "std", +] } + + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-router/README.md b/crates/burn-router/README.md new file mode 100644 index 0000000..be3e69a --- /dev/null +++ b/crates/burn-router/README.md @@ -0,0 +1,3 @@ +# Burn Router + +A multi-backend extension that forwards the tensor operations to the appropriate backend. diff --git a/crates/burn-router/src/backend.rs b/crates/burn-router/src/backend.rs new file mode 100644 index 0000000..cbba4a6 --- /dev/null +++ b/crates/burn-router/src/backend.rs @@ -0,0 +1,69 @@ +use super::{RouterChannel, RouterClient, RouterTensor, get_client}; +use alloc::{format, string::String}; +use burn_backend::{Backend, BackendTypes, DType, ExecutionError}; +use core::marker::PhantomData; + +/// A backend that forwards the tensor operations to the appropriate backend (given multiple backends). +pub struct BackendRouter { + r: PhantomData, +} + +impl core::fmt::Debug for BackendRouter { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_fmt(format_args!("router")) + } +} + +impl Clone for BackendRouter { + fn clone(&self) -> Self { + Self { r: PhantomData } + } +} + +impl Default for BackendRouter { + fn default() -> Self { + Self { r: PhantomData } + } +} + +impl BackendTypes for BackendRouter { + type Device = R::Device; + + type FloatTensorPrimitive = RouterTensor; + type IntTensorPrimitive = RouterTensor; + type BoolTensorPrimitive = RouterTensor; + type QuantizedTensorPrimitive = RouterTensor; + + type GraphPrimitive = burn_backend::GraphUnsupported; +} + +impl Backend for BackendRouter { + fn name(device: &Self::Device) -> String { + format!("router<{}>", R::name(device)) + } + + fn seed(device: &Self::Device, seed: u64) { + let client = get_client::(device); + client.seed(seed); + } + + fn sync(device: &Self::Device) -> Result<(), ExecutionError> { + let client = get_client::(device); + client.sync() + } + + fn dtype_usage(device: &Self::Device, dtype: DType) -> burn_backend::DTypeUsageSet { + let client = get_client::(device); + client.dtype_usage(dtype) + } + + fn device_count(_: u16) -> usize { + // This is what was there before, not sure if it's actually correct + 1 + } + + fn flush(device: &Self::Device) { + let client = get_client::(device); + client.flush(); + } +} diff --git a/crates/burn-router/src/bridge/base.rs b/crates/burn-router/src/bridge/base.rs new file mode 100644 index 0000000..8e61854 --- /dev/null +++ b/crates/burn-router/src/bridge/base.rs @@ -0,0 +1,32 @@ +use burn_backend::{Shape, backend::DeviceOps}; + +/// Allows tensors to be transferred between multiple backends. +pub trait MultiBackendBridge: Send + Sync + 'static { + /// The type that can be used to point to a tensor of any kind. + type TensorHandle; + /// Device type used by the backends. + type Device: DeviceOps; + + /// Change the backend of the given float tensor. + fn change_backend_float( + tensor: Self::TensorHandle, + shape: Shape, + target_device: &Self::Device, + ) -> Self::TensorHandle; + + /// Change the backend of the given int tensor. + fn change_backend_int( + tensor: Self::TensorHandle, + shape: Shape, + target_device: &Self::Device, + ) -> Self::TensorHandle; + + /// Change the backend of the given bool tensor. + fn change_backend_bool( + tensor: Self::TensorHandle, + shape: Shape, + target_device: &Self::Device, + ) -> Self::TensorHandle; + + // TODO: change_backend_quantized +} diff --git a/crates/burn-router/src/bridge/mod.rs b/crates/burn-router/src/bridge/mod.rs new file mode 100644 index 0000000..cbcb6ac --- /dev/null +++ b/crates/burn-router/src/bridge/mod.rs @@ -0,0 +1,3 @@ +mod base; + +pub use base::*; diff --git a/crates/burn-router/src/channel/base.rs b/crates/burn-router/src/channel/base.rs new file mode 100644 index 0000000..c206919 --- /dev/null +++ b/crates/burn-router/src/channel/base.rs @@ -0,0 +1,60 @@ +use alloc::string::String; +use burn_backend::{DType, Shape, backend::DeviceOps}; +use burn_ir::TensorIr; + +use crate::{MultiBackendBridge, RouterClient, RouterTensor, get_client}; + +/// Type alias for `
::TensorHandle`. +pub type TensorHandle
=
::TensorHandle; + +/// Defines the connection channel and operations for a setup with multiple backend router clients. +pub trait RouterChannel: Clone + Send + Sync + 'static + Sized { + /// Device type. + type Device: DeviceOps; + /// A bridge that can transfer tensors between multiple backends. + type Bridge: MultiBackendBridge; + /// Client type. + type Client: RouterClient; + + /// Name of the channel. + fn name(device: &Self::Device) -> String; + + /// Initialize a new client for the given device. + fn init_client(device: &Self::Device) -> Self::Client; + + /// Get the tensor handle corresponding to the [tensor representation](TensorIr). + fn get_tensor_handle(tensor: &TensorIr, client: &Self::Client) -> TensorHandle; + + /// Create a tensor with the given handle and shape. + fn register_tensor( + client: &Self::Client, + handle: TensorHandle, + shape: Shape, + dtype: DType, + ) -> RouterTensor; + + /// Change the tensor to a different client backend. + fn change_client_backend( + tensor: RouterTensor, + device: &Self::Device, // target device + ) -> RouterTensor { + // Get tensor handle from current client + let original_client = tensor.client.clone(); + let desc = tensor.into_ir(); + let mut handle = Self::get_tensor_handle(&desc, &original_client); + + if desc.dtype.is_float() { + handle = Self::Bridge::change_backend_float(handle, desc.shape.clone(), device); + } else if desc.dtype.is_int() { + handle = Self::Bridge::change_backend_int(handle, desc.shape.clone(), device); + } else if desc.dtype.is_bool() { + handle = Self::Bridge::change_backend_bool(handle, desc.shape.clone(), device); + } else { + unimplemented!() + } + + // Register tensor handle on target client + let target_client = get_client::(device); + Self::register_tensor(&target_client, handle, desc.shape, desc.dtype) + } +} diff --git a/crates/burn-router/src/channel/mod.rs b/crates/burn-router/src/channel/mod.rs new file mode 100644 index 0000000..cbcb6ac --- /dev/null +++ b/crates/burn-router/src/channel/mod.rs @@ -0,0 +1,3 @@ +mod base; + +pub use base::*; diff --git a/crates/burn-router/src/client/base.rs b/crates/burn-router/src/client/base.rs new file mode 100644 index 0000000..122bf1d --- /dev/null +++ b/crates/burn-router/src/client/base.rs @@ -0,0 +1,161 @@ +use crate::{RouterChannel, RouterTensor}; +use alloc::boxed::Box; +use alloc::vec::Vec; +use burn_backend::{ + DType, TensorData, + backend::{DeviceId, DeviceOps, ExecutionError}, +}; +use burn_ir::{GraphBindings, GraphId, OperationIr, TensorId, TensorIr}; +use burn_std::future::DynFut; +use core::ops::DerefMut; +use hashbrown::HashMap; +use spin::Mutex; + +/// Type alias for `::Client`. +pub type Client = ::Client; +pub(crate) static CLIENTS: RouterClientLocator = RouterClientLocator::new(); + +type Key = (core::any::TypeId, DeviceId); + +/// Define how to interact with the interpreter. +pub trait RouterClient: Clone + Send + Sync + Sized { + /// Device type. + type Device: DeviceOps; + + /// Register a new tensor operation to be executed by the interpreter (server). + fn register_op(&self, op: OperationIr); + /// Register a new tensor operation to be executed by the interpreter (server). + /// + /// Returns the new (uninitialized) output tensor(s) generated by the registered operation. + fn register(&self, op: OperationIr) -> Vec> { + let out = op + .outputs() + .map(|output| { + RouterTensor::new(output.id, output.shape.clone(), output.dtype, self.clone()) + }) + .collect(); + self.register_op(op); + + out + } + /// Read the values contained by a tensor. + fn read_tensor_async(&self, tensor: TensorIr) -> DynFut>; + /// Sync the interpreter, ensure that all computations are finished. + fn sync(&self) -> Result<(), ExecutionError>; + /// Eagerly submit the operations registered so far without waiting for them to complete. + /// + /// Unlike [`sync`](Self::sync), this does not block on results — it only ensures the + /// buffered operations are handed off for execution (and, for the remote backend, sent to + /// the server) instead of sitting in a local buffer. + fn flush(&self); + /// Create a new (uninitialized) empty tensor and returns its corresponding [tensor id](TensorId). + fn create_empty_handle(&self) -> TensorId; + /// Create a new [RouterTensor] from the tensor data. + fn register_tensor_data(&self, data: TensorData) -> RouterTensor; + /// Get the current device used by all operations handled by this client. + fn device(&self) -> Self::Device; + /// Seed the interpreter. + fn seed(&self, seed: u64); + /// Returns the supported data type usage set + fn dtype_usage(&self, dtype: DType) -> burn_backend::DTypeUsageSet; + + /// Register a reusable group of operations (in relative form) under `graph_id` *and* run its + /// first invocation with `bindings`, so it can later be replayed by id with + /// [`execute_graph`](Self::execute_graph). + /// + /// Registration always coincides with the first execution, so they're combined to save a + /// round-trip on a cache miss. Used by the fusion layer to avoid re-sending a recurring + /// op-graph. + fn register_and_execute_graph( + &self, + graph_id: GraphId, + relative_graph: Vec, + bindings: GraphBindings, + ); + + /// Replay a previously [registered](Self::register_and_execute_graph) graph with the given + /// concrete bindings. + fn execute_graph(&self, graph_id: GraphId, bindings: GraphBindings); + + /// Register `new_id` as an alias of `src_id` — a second handle over the same backing buffer. + /// + /// Used by the fusion layer's cross-stream sharing (see + /// [`FusionRuntime::alias_handle`](burn_fusion::FusionRuntime::alias_handle)): when a tensor is + /// shared to another stream, that stream's view needs its own id so that consuming it (a + /// `ReadWrite` last-use) frees only this alias, leaving the original handle valid. The server + /// clones the source handle (an `Arc`-style refcount on the device buffer) under `new_id`. + fn register_alias(&self, new_id: TensorId, src_id: TensorId); +} + +pub(crate) struct RouterClientLocator { + clients: Mutex>>>, +} + +/// Get the client for the given device +pub fn get_client(device: &R::Device) -> Client { + CLIENTS.client::(device) +} + +/// Initialize a new client for the given device. +/// +/// If a (global) seed was previously set, the client seed is set. +fn new_client(device: &R::Device) -> Client { + R::init_client(device) +} + +impl RouterClientLocator { + /// Create a new client locator. + pub const fn new() -> Self { + Self { + clients: Mutex::new(None), + } + } + + /// Get the router client for the given device. + /// + /// If a client isn't already initialized, it is created. + pub fn client(&self, device: &R::Device) -> Client { + let device_id = device.id(); + let client_id = (core::any::TypeId::of::(), device_id); + let mut clients = self.clients.lock(); + + if clients.is_none() { + let client = new_client::(device); + Self::register_inner::(client_id, client, &mut clients); + } + + match clients.deref_mut() { + Some(clients) => match clients.get(&client_id) { + Some(client) => { + let client: &Client = client.downcast_ref().unwrap(); + client.clone() + } + None => { + let client = new_client::(device); + let any = Box::new(client.clone()); + clients.insert(client_id, any); + client + } + }, + _ => unreachable!(), + } + } + + fn register_inner( + key: Key, + client: Client, + clients: &mut Option>>, + ) { + if clients.is_none() { + *clients = Some(HashMap::new()); + } + + if let Some(clients) = clients { + if clients.contains_key(&key) { + panic!("Client already created for device {key:?}"); + } + + clients.insert(key, Box::new(client)); + } + } +} diff --git a/crates/burn-router/src/client/mod.rs b/crates/burn-router/src/client/mod.rs new file mode 100644 index 0000000..cbcb6ac --- /dev/null +++ b/crates/burn-router/src/client/mod.rs @@ -0,0 +1,3 @@ +mod base; + +pub use base::*; diff --git a/crates/burn-router/src/custom_op.rs b/crates/burn-router/src/custom_op.rs new file mode 100644 index 0000000..0c76b42 --- /dev/null +++ b/crates/burn-router/src/custom_op.rs @@ -0,0 +1,152 @@ +use alloc::string::{String, ToString}; +use alloc::sync::Arc; +use burn_backend::backend::BackendTypes; +use burn_ir::{BackendIr, CustomOpIr, HandleContainer}; +use hashbrown::HashMap; + +/// A handler that executes a single custom operation against a backend's tensor handles. +/// +/// It reads its input tensors (and any scalar arguments) out of `handles`, runs the actual backend +/// computation on the interpreter's `device`, and registers the resulting output tensor(s) back +/// into `handles` under the ids declared by the [`CustomOpIr`]. The `device` lets a handler create +/// brand-new tensors (e.g. a custom op with no tensor inputs, like a server-side data loader). This +/// is the server-side counterpart of the client building an `OperationIr::Custom` — see +/// [`CustomOpRegistry`]. +pub type CustomOpHandler = Arc< + dyn Fn( + &mut HandleContainer<::Handle>, + &CustomOpIr, + &::Device, + ) + Send + + Sync, +>; + +/// A set of [custom operation handlers](CustomOpHandler), keyed by the custom op's id. +/// +/// The remote backend ships custom ops to the server as `OperationIr::Custom(CustomOpIr)`; the +/// server's [`TensorInterpreter`](crate::TensorInterpreter) can't know how to execute an arbitrary +/// custom op on its own, so it looks the id up in this registry and calls the matching handler. +/// +/// The registry is *owned* state (cheap to clone — handlers live behind an [`Arc`]), built once +/// before the server starts and shared read-only across every session's interpreter. There is no +/// global/static registry to manage. +#[derive(Clone)] +pub struct CustomOpRegistry { + handlers: HashMap>, +} + +impl Default for CustomOpRegistry { + fn default() -> Self { + Self { + handlers: HashMap::new(), + } + } +} + +impl CustomOpRegistry { + /// Create a new, empty registry. + pub fn new() -> Self { + Self::default() + } + + /// Register the handler executed when a custom op with the given `id` reaches the interpreter. + /// + /// The `id` must match the one the client puts in its [`CustomOpIr`]. Registering the same id + /// twice replaces the previous handler. + pub fn register(&mut self, id: &str, handler: F) + where + F: Fn(&mut HandleContainer, &CustomOpIr, &B::Device) + Send + Sync + 'static, + { + self.handlers.insert(id.to_string(), Arc::new(handler)); + } + + /// Get the handler registered for `id`, if any. + pub(crate) fn get(&self, id: &str) -> Option<&CustomOpHandler> { + self.handlers.get(id) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::TensorInterpreter; + use burn_backend::{DType, Scalar, Shape, TensorData, ops::FloatTensorOps}; + use burn_flex::Flex; + use burn_ir::{OperationIr, ScalarIr, TensorId, TensorIr}; + use std::sync::Mutex; + + #[test] + fn custom_op_is_dispatched_with_scalars() { + // A handler that scales a float tensor by a scalar argument, and records what it received so + // the test can assert the interpreter routed the op (and its scalars) to it. + let seen_scalars = Arc::new(Mutex::new(None)); + let seen_scalars_handler = seen_scalars.clone(); + + let mut registry = CustomOpRegistry::::new(); + registry.register("scale", move |handles, ir, _device| { + let input = handles.get_float_tensor::(&ir.inputs[0]); + let factor: Scalar = ir.scalars[0].into(); + let output = Flex::float_mul_scalar(input, factor); + handles.register_float_tensor::(&ir.outputs[0].id, output); + *seen_scalars_handler.lock().unwrap() = Some(ir.scalars.clone()); + }); + + let mut interp = TensorInterpreter::::with_custom_ops(Default::default(), registry); + let input = interp.register_tensor_data_desc(TensorData::from([2.0f32, 4.0])); + let output = TensorIr::uninit(TensorId::new(1_000_000), Shape::from([2]), DType::F32); + + let desc = + CustomOpIr::with_scalars("scale", &[input], &[output], vec![ScalarIr::Float(3.0)]); + interp.register_op(OperationIr::Custom(desc)); + + // The interpreter routed the op to our handler, carrying the scalar we shipped. (The handler + // also read its input and registered an output; either failing would panic before here.) + assert_eq!( + seen_scalars.lock().unwrap().clone(), + Some(vec![ScalarIr::Float(3.0)]) + ); + } + + #[test] + fn custom_op_can_create_tensors_from_the_device() { + // A no-input handler that builds a tensor from scratch on the interpreter's device — the + // shape of a server-side data loader. Exercises the `&device` passed to the handler. + use burn_backend::ops::IntTensorOps; + + let mut registry = CustomOpRegistry::::new(); + registry.register("load", |handles, ir, device| { + let values: Vec = ir.scalars.iter().map(|s| s.elem::()).collect(); + let n = values.len(); + let tensor = Flex::int_from_data(TensorData::new(values, [n]), device); + handles.register_int_tensor::(&ir.outputs[0].id, tensor); + }); + + let mut interp = TensorInterpreter::::with_custom_ops(Default::default(), registry); + let out = TensorIr::uninit(TensorId::new(2_000_000), Shape::from([3]), DType::I64); + let desc = CustomOpIr::with_scalars( + "load", + &[], + &[out.clone()], + vec![ScalarIr::Int(10), ScalarIr::Int(20), ScalarIr::Int(30)], + ); + interp.register_op(OperationIr::Custom(desc)); + + // The handler created the tensor on the device and registered it under the output id. + let data = interp.get_tensor(&TensorIr { + status: burn_ir::TensorStatus::ReadOnly, + ..out + }); + match data { + burn_ir::HandleKind::Int(_) => {} + _ => panic!("expected an int tensor"), + } + } + + #[test] + #[should_panic(expected = "No custom-op handler registered")] + fn unregistered_custom_op_panics() { + let mut interp = TensorInterpreter::::new(Default::default()); + let desc = CustomOpIr::new("missing", &[], &[]); + interp.register_op(OperationIr::Custom(desc)); + } +} diff --git a/crates/burn-router/src/fusion.rs b/crates/burn-router/src/fusion.rs new file mode 100644 index 0000000..fb903a1 --- /dev/null +++ b/crates/burn-router/src/fusion.rs @@ -0,0 +1,576 @@ +//! Generic client-side graph caching for any [router backend](crate::BackendRouter). +//! +//! Wrapping a router backend as [`Fusion`](burn_fusion::Fusion) turns recurring groups of tensor +//! operations into reusable, client-cached graphs. A single greedy [`RouterFuser`] accumulates +//! every operation and is drained only at a sync point, so one graph covers each connected block +//! between syncs. On execution the group is registered once on the backend (via the +//! [`RouterClient`]) and thereafter invoked by id with only the changing bindings — for the remote +//! backend this means a recurring computation (e.g. a model block) crosses the network once +//! instead of every step. +//! +//! `burn-fusion` names its hook `Optimization`; here that hook *is* a cached graph execution +//! ([`RouterGraphExecution`]), to distinguish it from a compute backend's kernel fusion. + +use std::collections::HashSet; +use std::marker::PhantomData; +use std::sync::atomic::{AtomicU64, Ordering}; + +use burn_backend::DType; +use burn_backend::ops::FloatTensorOps; +use burn_fusion::stream::{Context, Operation, OrderedExecution}; +use burn_fusion::{ + FuserProperties, FuserStatus, FusionBackend, FusionRuntime, NumOperations, OperationFuser, + Optimization, +}; +use burn_ir::{ + BackendIr, CustomOpIr, GraphBindings, GraphId, Handle, HandleContainer, OperationIr, ScalarIr, + TensorHandle, TensorId, TensorIr, TensorStatus, +}; +use serde::{Deserialize, Serialize}; + +use burn_std::config::config; + +use crate::{BackendRouter, RouterChannel, RouterClient, RouterTensor, get_client}; + +static GRAPH_ID_COUNTER: AtomicU64 = AtomicU64::new(0); + +fn next_graph_id() -> GraphId { + GraphId(GRAPH_ID_COUNTER.fetch_add(1, Ordering::Relaxed)) +} + +// The router backend already implements `Backend`; these two impls add the `BackendIr` + +// `FusionBackend` glue so it can be wrapped as `Fusion>`. All four tensor +// primitives are `RouterTensor`, so the handle conversions are the identity (mirroring the +// `impl BackendIr for Fusion` in burn-fusion). +impl BackendIr for BackendRouter { + type Handle = RouterTensor; + + fn float_tensor(handle: TensorHandle) -> burn_backend::tensor::FloatTensor { + handle.handle + } + + fn int_tensor(handle: TensorHandle) -> burn_backend::tensor::IntTensor { + handle.handle + } + + fn bool_tensor(handle: TensorHandle) -> burn_backend::tensor::BoolTensor { + handle.handle + } + + fn quantized_tensor( + handle: TensorHandle, + ) -> burn_backend::tensor::QuantizedTensor { + handle.handle + } + + fn float_tensor_handle(tensor: burn_backend::tensor::FloatTensor) -> Self::Handle { + tensor + } + + fn int_tensor_handle(tensor: burn_backend::tensor::IntTensor) -> Self::Handle { + tensor + } + + fn bool_tensor_handle(tensor: burn_backend::tensor::BoolTensor) -> Self::Handle { + tensor + } + + fn quantized_tensor_handle( + tensor: burn_backend::tensor::QuantizedTensor, + ) -> Self::Handle { + tensor + } +} + +impl FusionBackend for BackendRouter { + type FusionRuntime = RouterFusionRuntime; + type FullPrecisionBackend = Self; + + fn cast_float(tensor: burn_backend::tensor::FloatTensor, dtype: DType) -> Self::Handle { + Self::float_cast(tensor, dtype.into()) + } +} + +/// The [fusion runtime](FusionRuntime) for a [router backend](BackendRouter). +/// +/// Its [handle](FusionRuntime::FusionHandle) is a [`RouterTensor`] — a lightweight reference to a +/// backend-resident tensor plus the client to reach it — and its "optimization" is a backend-cached +/// op-graph (see [`RouterGraphExecution`]). +pub struct RouterFusionRuntime { + _p: PhantomData, +} + +impl core::fmt::Debug for RouterFusionRuntime { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_str("RouterFusionRuntime") + } +} + +impl FusionRuntime for RouterFusionRuntime { + type OptimizationState = RouterGraphExecutionState; + type Optimization = RouterGraphExecution; + type FusionHandle = RouterTensor; + type FusionDevice = R::Device; + + fn fusers(device: R::Device) -> Vec>> { + vec![Box::new(RouterFuser::::new(device))] + } + + fn alias_handle(handle: &RouterTensor) -> RouterTensor { + // The router handle is a thin id into a server-side tensor, so a bare `clone()` would keep + // the same server id — every cross-stream alias would collapse onto one server handle, and + // the first stream to consume it (`ReadWrite`) would free it for the others. Instead mint a + // fresh id and have the server register it as an alias of the same buffer (a refcounted + // clone), so each stream's view frees independently. Mirrors local backends, where the + // default `clone()` already yields an independent `HandleContainer` entry over a shared + // `Arc` buffer. + let id = handle.client.create_empty_handle(); + handle.client.register_alias(id, handle.id); + RouterTensor::new( + id, + handle.shape.clone(), + handle.dtype, + handle.client.clone(), + ) + } + + fn free_handle(handles: &mut HandleContainer>, tensor: &TensorIr) { + // Only `ReadWrite` (last-use) nodes are freed here, matching `HandleContainer::free`. + if tensor.status != TensorStatus::ReadWrite { + return; + } + // The drained block already freed this tensor server-side: every consumed op (a `Drop`, or + // a `ReadWrite` input of a compute op) is replayed on the server and pops its handle. So + // remove the client entry WITHOUT letting `RouterTensor::drop` register a *second*, + // redundant `Drop` for the same id. Bumping the refcount makes that drop a no-op. + if let Some(Handle::Existing(handle)) = handles.remove_handle(tensor.id) { + handle + .count + .fetch_add(1, core::sync::atomic::Ordering::Relaxed); + } + } +} + +/// A greedy operation fuser that records every operation and never closes itself. +/// +/// Because [`status`](OperationFuser::status) always reports [`FuserStatus::Open`], the fusion +/// engine keeps deferring in lazy mode and only drains the queue at a sync point (a read, `sync`, +/// or `flush`). At that point [`finish`](OperationFuser::finish) yields a single +/// [`RouterGraphExecution`] covering the whole accumulated (connected) block. +pub struct RouterFuser { + device: R::Device, + ops: Vec, + score: u64, + score_max: u64, + num_since_max_unchanged: usize, + /// Close the graph once it reaches this many ops, if set (`FusionConfig::max_graph_size`). + max_graph_size: Option, + /// Close the graph once the score hasn't reached a new max for this many consecutive ops + /// (`FusionConfig::growth_patience`). + growth_patience: usize, +} + +impl RouterFuser { + fn new(device: R::Device) -> Self { + let cfg = config(); + let fusion = cfg.fusion(); + Self { + device, + ops: Vec::new(), + score: 0, + score_max: 0, + num_since_max_unchanged: 0, + max_graph_size: fusion.max_graph_size, + growth_patience: fusion.growth_patience, + } + } + + /// Value-based fusion score for the currently accumulated ops. + /// + /// Benefit: estimated % of serialized bytes caching saves per replay (the relative graph vs the + /// per-replay bindings — see [`estimate_saved_pct`]), weighted by `FACTOR_SAVED`. Cost: a + /// penalty that grows with op count past `FREE_OPS`, so the score *peaks* at a "right-sized" + /// graph and then decays once the per-op overhead outweighs the savings. + /// + /// Floored at 1 for any non-empty graph: a score of 0 makes `find_best_optimization_index` + /// treat the graph as "don't fuse" (streaming every op unfused → the cache never replays). The + /// `+1` doesn't move the argmax, so it doesn't change which size wins. + fn score(&self) -> u64 { + const FACTOR_SAVED: u64 = 100; // weight per % point saved → benefit in 0..=10_000 + const FREE_OPS: usize = 64; // ops below this are not penalized + const PENALTY_PER_OP: u64 = 0; // penalty per op beyond FREE_OPS + + let benefit = estimate_saved_pct(&self.ops) * FACTOR_SAVED; + let penalty = (self.ops.len().saturating_sub(FREE_OPS) as u64) * PENALTY_PER_OP; + benefit.saturating_sub(penalty) + 1 + } +} + +impl Clone for RouterFuser { + fn clone(&self) -> Self { + Self { + device: self.device.clone(), + ops: self.ops.clone(), + score: self.score, + score_max: self.score_max, + num_since_max_unchanged: self.num_since_max_unchanged, + max_graph_size: self.max_graph_size, + growth_patience: self.growth_patience, + } + } +} + +impl OperationFuser> for RouterFuser { + fn fuse(&mut self, operation: &OperationIr) { + self.ops.push(operation.clone()); + + self.score = self.score(); + if self.score > self.score_max { + self.score_max = self.score; + self.num_since_max_unchanged = 0; + } else { + self.num_since_max_unchanged += 1; + } + } + + fn finish(&mut self) -> RouterGraphExecution { + let ops = core::mem::take(&mut self.ops); + RouterGraphExecution::new(ops, self.device.clone()) + } + + fn reset(&mut self) { + self.ops.clear(); + } + + fn status(&self) -> FuserStatus { + let over_max = self.max_graph_size.is_some_and(|max| self.len() > max); + if self.num_since_max_unchanged >= self.growth_patience || over_max { + FuserStatus::Closed + } else { + FuserStatus::Open + } + } + + fn properties(&self) -> FuserProperties { + FuserProperties { + score: self.score, + ready: self.ops.len() > 1, + } + } + + fn len(&self) -> usize { + self.ops.len() + } + + fn clone_dyn(&self) -> Box>> { + Box::new(self.clone()) + } +} + +/// Unfused execution of a single [custom op](OperationIr::Custom) on a router backend. +/// +/// Every built-in op reaches the backend through a `BackendRouter` method, so when the fuser leaves +/// it unfused (a one-op segment — e.g. a source op whose output is read right away) `burn-fusion` +/// runs that method op by op. A custom op has no such method, so it needs this handler: it is the +/// [`Operation`] registered alongside the op's IR, and runs only on the unfused path. When the op is +/// instead fused into a [`RouterGraphExecution`] it ships as part of the graph and this never runs. +/// +/// It mirrors what a built-in op does on that path: resolve the fused input handles to their +/// backend tensor ids, ship the op through the [`RouterClient`], and bind the outputs so downstream +/// ops find them. (Without it the op would silently do nothing and the backend would never create +/// the output handle — "Should have handle for tensor ..." on the next access.) +pub struct CustomOperation { + ir: CustomOpIr, + device: R::Device, +} + +impl core::fmt::Debug for CustomOperation { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("CustomOperation") + .field("id", &self.ir.id) + .finish() + } +} + +impl CustomOperation { + /// Create the unfused handler for `ir`, executed on `device`. + pub fn new(ir: CustomOpIr, device: R::Device) -> Self { + Self { ir, device } + } +} + +impl Operation> for CustomOperation { + fn execute(&self, handles: &mut HandleContainer>) { + let client = get_client::(&self.device); + + // Map each fused input handle to its backend tensor id. `into_ir` carries the + // refcount/free semantics, so a last-use (`ReadWrite`) input still frees correctly. + let inputs: Vec = self + .ir + .inputs + .iter() + .map(|input| handles.get_handle(&input.id, &input.status).into_ir()) + .collect(); + + // Mint fresh backend ids for the outputs. + let outputs: Vec> = self + .ir + .outputs + .iter() + .map(|out| { + RouterTensor::new( + client.create_empty_handle(), + out.shape.clone(), + out.dtype, + client.clone(), + ) + }) + .collect(); + + // Ship the op to the backend with the translated ids; scalars travel unchanged. + client.register_op(OperationIr::Custom(CustomOpIr { + id: self.ir.id.clone(), + inputs, + outputs: outputs.iter().map(|out| out.to_ir_out()).collect(), + scalars: self.ir.scalars.clone(), + })); + + // Bind the outputs under their fused ids so any downstream op resolves them. + for (out, tensor) in self.ir.outputs.iter().zip(outputs) { + handles.register_handle(out.id, tensor); + } + } +} + +/// A reusable group of operations, registered once on the backend and thereafter invoked by id. +/// +/// The recorded [`graph`](Vec) is in *relative* form (positional tensor ids, relative +/// shape-dim ids, scalar placeholders), which is invariant across invocations — that is what lets +/// the backend cache it and the client reuse it. On each [`execute`](Optimization::execute) only +/// the concrete bindings are computed from the [`Context`] and sent; the (large) graph itself +/// travels only on the first invocation. +pub struct RouterGraphExecution { + graph: Vec, + device: R::Device, + /// Relative ids of the graph's boundary, precomputed once from the (static) graph so each + /// replay is O(boundary) instead of re-scanning every tensor. Inputs are tensors not produced + /// by a compute op (external data / prior results, incl. `Init`/`from_data` outputs); outputs + /// are compute-produced tensors that survive (neither consumed in place nor dropped here). + input_ids: Vec, + output_ids: Vec, + /// Backend-side id, assigned and registered on the first execution and reused afterwards. + graph_id: Option, + _p: PhantomData, +} + +/// Serializable state for a [`RouterGraphExecution`]. +/// +/// The backend id is intentionally not serialized — it is per-connection state, so a deserialized +/// graph re-registers itself on first use. +#[derive(Serialize, Deserialize)] +pub struct RouterGraphExecutionState { + graph: Vec, +} + +impl RouterGraphExecution { + fn new(graph: Vec, device: R::Device) -> Self { + let (input_ids, output_ids) = classify_boundary(&graph); + Self { + graph, + device, + input_ids, + output_ids, + graph_id: None, + _p: PhantomData, + } + } +} + +/// Precompute the graph's boundary as `(input relative ids, surviving-output relative ids)`. +/// +/// - **Inputs** are tensors that aren't produced by a *compute* op: external data and prior-block +/// results read by the graph, plus `Init`/`from_data` outputs (whose handle is registered +/// out-of-band, so they're really inputs even though an `Init` op "produces" them). +/// - **Outputs** are compute-produced tensors that survive — neither consumed in place +/// (`ReadWrite` anywhere) nor dropped within the graph. This mirrors the fusion engine's own +/// `drain_queue` freeing logic, keeping client-side handle state consistent with the unfused path. +/// +/// Intermediate tensors (compute-produced and consumed/dropped here) are in neither list — the +/// replay owns their ids — so a replay only ever touches the boundary. +fn classify_boundary(graph: &[OperationIr]) -> (Vec, Vec) { + let mut referenced: HashSet = HashSet::new(); + let mut compute_produced: HashSet = HashSet::new(); + let mut consumed: HashSet = HashSet::new(); + for op in graph { + if let OperationIr::Drop(tensor) = op { + consumed.insert(tensor.id); + } + if !matches!(op, OperationIr::Init(_)) { + for tensor in op.outputs() { + compute_produced.insert(tensor.id); + } + } + for tensor in op.nodes() { + referenced.insert(tensor.id); + if tensor.status == TensorStatus::ReadWrite { + consumed.insert(tensor.id); + } + } + } + + let inputs = referenced + .iter() + .filter(|id| !compute_produced.contains(id)) + .copied() + .collect(); + let outputs = compute_produced + .iter() + .filter(|id| !consumed.contains(id)) + .copied() + .collect(); + (inputs, outputs) +} + +/// Estimate the % of serialized bytes that caching this op-graph saves per replay. +/// +/// Dependency-free structural estimate (reuses [`classify_boundary`]): the baseline is the whole +/// relative graph's bytes (op overhead + each tensor's id/dtype/status + its dims); the per-replay +/// bindings are just the boundary `(relative id, concrete id)` pairs plus the distinct dim table. +/// Scalars/ranges travel in both and cancel in the ratio. Returns 0..=100. +fn estimate_saved_pct(ops: &[OperationIr]) -> u64 { + if ops.is_empty() { + return 0; + } + const TENSOR: u64 = 10; // id (≈8) + dtype + status + const PER_DIM: u64 = 8; + const OP_OVERHEAD: u64 = 8; // variant tag + small bookkeeping + const PAIR: u64 = 16; // a (relative id, concrete id) binding entry + + let mut baseline = 0u64; + let mut dims: HashSet = HashSet::new(); + for op in ops { + baseline += OP_OVERHEAD; + for tensor in op.nodes().into_iter().chain(op.outputs()) { + baseline += TENSOR; + for dim in tensor.shape.iter() { + baseline += PER_DIM; + dims.insert(*dim); + } + } + } + + let (inputs, outputs) = classify_boundary(ops); + let bindings = (inputs.len() + outputs.len()) as u64 * PAIR + dims.len() as u64 * PER_DIM; + + let saved = baseline.saturating_sub(bindings); + (saved * 100 / baseline).min(100) +} + +impl core::fmt::Debug for RouterGraphExecution { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("RouterGraphExecution") + .field("len", &self.graph.len()) + .finish() + } +} + +impl NumOperations for RouterGraphExecution { + fn len(&self) -> usize { + self.graph.len() + } + + fn name(&self) -> &'static str { + "RouterGraphExecution" + } +} + +impl Optimization> for RouterGraphExecution { + fn execute( + &mut self, + context: &mut Context>, + _execution: &OrderedExecution>, + ) { + let client = get_client::(&self.device); + + // Walk only the precomputed boundary — never every tensor — so a replay is O(boundary). + // Inputs reuse their resident concrete id; surviving outputs get a fresh id and a handle. + // Intermediates are never touched here: the replay owns their ids and derives their shapes + // from the shape table below. + let mut tensors: Vec<(TensorId, TensorId)> = + Vec::with_capacity(self.input_ids.len() + self.output_ids.len()); + + for &input_id in &self.input_ids { + let global_id = match context.tensors.get(&input_id) { + Some(global) => global.id, + None => continue, + }; + if let Some(handle) = context.handles.get_handle_ref(&global_id) { + tensors.push((input_id, handle.id())); + } + } + + for &output_id in &self.output_ids { + // Extract owned metadata first so the `context.tensors` borrow ends before we touch + // `context.handles`. + let output = context + .tensors + .get(&output_id) + .map(|global| (global.id, global.shape.clone(), global.dtype)); + if let Some((fusion_id, shape, dtype)) = output { + let concrete_id = client.create_empty_handle(); + tensors.push((output_id, concrete_id)); + let handle = RouterTensor::new(concrete_id, shape, dtype, client.clone()); + context.handles.register_handle(fusion_id, handle); + } + } + + // Dense shape-dim table indexed by relative dim id (dim ids are dense `0..N`). + let mut shapes = vec![0usize; context.shapes_relative2global.len()]; + for (relative, concrete) in context.shapes_relative2global.iter() { + if *relative < shapes.len() { + shapes[*relative] = *concrete; + } + } + + // Concrete scalar values, indexed by their placeholder id. + let mut scalars = vec![ScalarIr::UInt(0); context.scalars.len()]; + for (scalar_id, value) in context.scalars.iter() { + let idx = scalar_id.value as usize; + if idx < scalars.len() { + scalars[idx] = *value; + } + } + + // Concrete slice ranges, indexed by their placeholder id (carried in a relative range's + // `start`). Cheap to clone — a few `Slice`s per slice op — and, like scalars, they can + // change between invocations of the same cached graph, so they travel every time. + let ranges = context.ranges.clone(); + + let bindings = GraphBindings { + tensors, + shapes, + scalars, + ranges, + }; + match self.graph_id { + // Already registered: replay by id, sending only the bindings. + Some(id) => client.execute_graph(id, bindings), + // First invocation: register the relative graph and execute it in a single round-trip. + None => { + let id = next_graph_id(); + self.graph_id = Some(id); + client.register_and_execute_graph(id, self.graph.clone(), bindings); + } + }; + } + + fn to_state(&self) -> RouterGraphExecutionState { + RouterGraphExecutionState { + graph: self.graph.clone(), + } + } + + fn from_state(device: &R::Device, state: RouterGraphExecutionState) -> Self { + Self::new(state.graph, device.clone()) + } +} diff --git a/crates/burn-router/src/graph.rs b/crates/burn-router/src/graph.rs new file mode 100644 index 0000000..a2c44a0 --- /dev/null +++ b/crates/burn-router/src/graph.rs @@ -0,0 +1,137 @@ +//! A cached relative op-graph and its replay against a [`TensorInterpreter`]. +//! +//! A router server registers a recurring op sequence once as a [`Graph`] (in *relative* form: +//! positional tensor ids, relative shape dims, placeholder scalars/ranges), then replays it by id +//! with only the per-invocation [`GraphBindings`]. This is the server-side counterpart of the +//! client's cached optimization — it turns a recurring computation (e.g. a model block per step) +//! into one registration plus cheap replays. + +use core::sync::atomic::{AtomicU64, Ordering}; + +use alloc::sync::Arc; +use alloc::vec::Vec; + +use burn_backend::Slice; +use burn_ir::{BackendIr, GraphBindings, IrVisitorMut, OperationIr, ScalarIr, TensorId, TensorIr}; +use hashbrown::HashMap; + +use crate::TensorInterpreter; + +/// Server-allocated ids for a replay's intermediate tensors carry this high bit so they can never +/// collide with client-allocated ids (whose monotonic counter never reaches `1 << 63`). The bit is +/// purely server-internal: intermediates are produced and freed within a single replay and are +/// never referenced by the client. +const INTERMEDIATE_ID_BIT: u64 = 1 << 63; +static INTERMEDIATE_ID_COUNTER: AtomicU64 = AtomicU64::new(0); + +fn alloc_intermediate_id() -> TensorId { + let value = INTERMEDIATE_ID_COUNTER.fetch_add(1, Ordering::Relaxed); + TensorId::new(value | INTERMEDIATE_ID_BIT) +} + +/// A cached relative op-graph, registered once and [replayed](Graph::replay) by id with +/// per-invocation [`GraphBindings`]. +/// +/// Cheap to clone — the op list is shared behind an [`Arc`]. That lets a server keep its graph +/// cache behind a short-lived lock: clone the [`Graph`] handle out under the lock, then replay +/// after releasing it, so the lock never spans the backend dispatch. +#[derive(Clone, Debug)] +pub struct Graph { + ops: Arc>, +} + +impl Graph { + /// Wrap a relative op-graph so it can be replayed. + pub fn new(ops: Vec) -> Self { + Self { ops: Arc::new(ops) } + } + + /// Number of operations in the graph. + pub fn len(&self) -> usize { + self.ops.len() + } + + /// Whether the graph has no operations. + pub fn is_empty(&self) -> bool { + self.ops.is_empty() + } + + /// Replay the graph against `interpreter`, rebinding the relative form to concrete tensors. + /// + /// The graph is in relative form: its tensor ids are positional, shape dims are relative ids, + /// and scalars/ranges are placeholders. `bindings` arrive packaged the way the replay uses + /// them — the boundary `tensors` map is moved straight into the working id table and grown with + /// intermediate ids on demand, and `shapes` is a dense table indexed by relative dim id. For + /// each op we: + /// - resolve every tensor id to its boundary binding, or a freshly allocated intermediate id + /// (memoized so all references to one intermediate agree); + /// - rewrite every tensor's shape dims in place via the shape table (so intermediates get + /// correct shapes too, without being sent); + /// - substitute scalar placeholders and restore concrete slice ranges; + /// + /// then hand the rebound op to the unchanged [`TensorInterpreter::register_op`], reproducing the + /// exact sequence of global ops the client would have streamed op-by-op. + pub fn replay( + &self, + interpreter: &mut TensorInterpreter, + bindings: GraphBindings, + ) { + let GraphBindings { + tensors, + shapes, + scalars, + ranges, + } = bindings; + // The boundary map *is* the working id table — seeded here, intermediates added on demand. + let mut ids: HashMap = tensors.into_iter().collect(); + for op in self.ops.iter() { + let mut op = op.clone(); + let mut visitor = ReplayVisitor { + ids: &mut ids, + shapes: &shapes, + scalars: &scalars, + ranges: &ranges, + }; + op.visit_mut(&mut visitor); + interpreter.register_op(op); + } + } +} + +/// Rebinds a relative op's tensors, scalars, and ranges to their concrete values during replay. +struct ReplayVisitor<'a> { + /// The working id table; intermediates are allocated on demand and memoized here so all + /// references to one intermediate agree. Persists across ops within a replay. + ids: &'a mut HashMap, + shapes: &'a [usize], + scalars: &'a [ScalarIr], + ranges: &'a [Slice], +} + +impl IrVisitorMut for ReplayVisitor<'_> { + fn visit_tensor_mut(&mut self, tensor: &mut TensorIr) { + tensor.id = *self + .ids + .entry(tensor.id) + .or_insert_with(alloc_intermediate_id); + for dim in tensor.shape.iter_mut() { + *dim = self.shapes.get(*dim).copied().unwrap_or(*dim); + } + } + + fn visit_scalar_mut(&mut self, scalar: &mut ScalarIr) { + if let ScalarIr::UInt(placeholder) = *scalar + && (placeholder as usize) < self.scalars.len() + { + *scalar = self.scalars[placeholder as usize]; + } + } + + fn visit_range_mut(&mut self, range: &mut Slice) { + // Restore concrete slice bounds: relativization replaced each range with a placeholder + // whose `start` is the binding id (see `OperationConverter::relative_range`). + if let Some(concrete) = self.ranges.get(range.start as usize) { + *range = *concrete; + } + } +} diff --git a/crates/burn-router/src/interpreter.rs b/crates/burn-router/src/interpreter.rs new file mode 100644 index 0000000..94c7ba2 --- /dev/null +++ b/crates/burn-router/src/interpreter.rs @@ -0,0 +1,2170 @@ +use core::sync::atomic::{AtomicU64, Ordering}; + +use super::{RouterClient, RouterTensor}; +use crate::CustomOpRegistry; +use crate::{ + binary_bool_ops, binary_float_cmp_ops, binary_float_ops, binary_int_cmp_ops, binary_int_ops, + reduce_float_dim_ops, reduce_float2int_dim_ops, reduce_int_dim_ops, scalar_float_cmp_ops, + scalar_float_ops, scalar_int_cmp_ops, scalar_int_ops, unary_float_ops, unary_int_ops, +}; +use alloc::boxed::Box; +use burn_backend::{ + Backend, DType, DeviceOps, ExecutionError, Shape, TensorData, distributed::DistributedOps, + tensor::IndexingUpdateOp, +}; +use burn_ir::{ + ActivationOperationIr, BackendIr, BaseOperationIr, BoolOperationIr, FloatOperationIr, + HandleContainer, HandleKind, IntOperationIr, ModuleOperationIr, NumericOperationIr, + OperationIr, TensorId, TensorIr, TensorStatus, +}; +use burn_std::{DeviceSettings, future::DynFut}; + +/// An interpreter's context contains a [handle container](HandleContainer) to manage +/// (i.e., fetch and update) existing tensors. +pub struct InterpreterContext { + /// Handle container to retrieve tensors based on their intermediate representation. + handles: HandleContainer, +} + +static COUNTER: AtomicU64 = AtomicU64::new(0); + +impl InterpreterContext { + /// Create a new (uninitialized) empty tensor and returns its corresponding [tensor id](TensorId). + fn create_empty_handle(&mut self) -> TensorId { + let value = COUNTER.fetch_add(1, Ordering::Relaxed); + TensorId::new(value) + } +} + +/// A tensor interpreter is responsible for executing tensor operations for a given [intermediate backend](BackendIr). +/// +/// Single-owner and **not** `Clone`: the interpreter is driven by exactly one thread (the remote +/// server's per-session worker), so the handle container is owned directly rather than behind a +/// lock — every op takes `&mut self`, with no per-op `Mutex` on the hot path. +pub struct TensorInterpreter { + context: InterpreterContext, + device: B::Device, + /// Handlers for [custom operations](OperationIr::Custom), keyed by id. Shared read-only across + /// every session, so executing a custom op is a map lookup plus a call. + custom_ops: CustomOpRegistry, +} + +impl core::fmt::Debug for TensorInterpreter { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("TensorInterpreter") + .field("device", &self.device) + .finish() + } +} + +impl TensorInterpreter { + /// Create a new interpreter without any custom operation handlers. + pub fn new(device: B::Device) -> Self { + Self::with_custom_ops(device, CustomOpRegistry::default()) + } + + /// Create a new interpreter with the given [custom operation handlers](CustomOpRegistry). + pub fn with_custom_ops(device: B::Device, custom_ops: CustomOpRegistry) -> Self { + Self { + context: InterpreterContext { + handles: HandleContainer::new(), + }, + device, + custom_ops, + } + } + + /// Get the tensor handle for the given [tensor representation](TensorIr). + pub fn get_tensor_handle(&mut self, tensor: &TensorIr) -> B::Handle { + let handles = &mut self.context.handles; + handles.get_tensor_handle(tensor).handle + } + + /// Take the typed backend primitive for `tensor`, dispatching on its dtype. + /// + /// Unlike [`get_tensor_handle`](Self::get_tensor_handle) (which returns the opaque + /// `B::Handle`), this returns the concrete float/int/bool primitive so the caller can hand + /// it to `B::*_to_device`. Used by the same-host transfer path, which moves a tensor between + /// two interpreters living in the same server process without a host round-trip. + pub fn get_tensor(&mut self, tensor: &TensorIr) -> HandleKind { + let handles = &mut self.context.handles; + let dtype = tensor.dtype; + if dtype.is_float() { + HandleKind::Float(handles.get_float_tensor::(tensor)) + } else if dtype.is_int() { + HandleKind::Int(handles.get_int_tensor::(tensor)) + } else if dtype.is_bool() { + HandleKind::Bool(handles.get_bool_tensor::(tensor)) + } else { + todo!("Local transfer of {dtype:?} tensors is not supported yet"); + } + } + + /// Move a primitive produced on another interpreter's device onto this interpreter's device + /// and register it under `id`. + /// + /// The counterpart of [`get_tensor`](Self::get_tensor): the source interpreter hands over its + /// primitive, and the destination calls `B::*_to_device` onto its own device. When both + /// interpreters share the same device, the backend's `to_device` is a cheap no-op. + pub fn register_tensor_to_device(&mut self, id: TensorId, tensor: HandleKind) { + let ctx = &mut self.context; + match tensor { + HandleKind::Float(tensor) => { + let tensor = B::float_to_device(tensor, &self.device); + ctx.handles.register_float_tensor::(&id, tensor); + } + HandleKind::Int(tensor) => { + let tensor = B::int_to_device(tensor, &self.device); + ctx.handles.register_int_tensor::(&id, tensor); + } + HandleKind::Bool(tensor) => { + let tensor = B::bool_to_device(tensor, &self.device); + ctx.handles.register_bool_tensor::(&id, tensor); + } + HandleKind::Quantized(_) => { + todo!("Local transfer of quantized tensors is not supported yet"); + } + } + } + + /// Create a tensor with the given handle and shape. + pub fn register_tensor( + &mut self, + handle: B::Handle, + shape: Shape, + dtype: DType, + client: C, + ) -> RouterTensor { + let ctx = &mut self.context; + let id = ctx.create_empty_handle(); + + ctx.handles.register_handle(id, handle); + + RouterTensor::new(id, shape, dtype, client) + } + + /// Register `new_id` as an alias of `src_id`: a second handle over the same backing buffer. + /// + /// Used for the fusion layer's cross-stream sharing (see + /// [`RouterClient::register_alias`](crate::RouterClient::register_alias)). Cloning the backend + /// handle is a cheap `Arc`-style refcount bump, so the buffer survives until *both* ids are + /// freed — consuming the alias on one stream can't pull it out from under the other. + pub fn register_alias(&mut self, new_id: TensorId, src_id: TensorId) { + let ctx = &mut self.context; + let handle = ctx + .handles + .get_handle_ref(&src_id) + .expect("alias source tensor must be materialized before it is aliased") + .clone(); + ctx.handles.register_handle(new_id, handle); + } + + /// Register a tensor from its data and id. + pub fn register_tensor_data_id(&mut self, id: TensorId, data: TensorData) { + let ctx = &mut self.context; + let dtype = data.dtype; + + if dtype.is_float() { + let tensor = B::float_from_data(data, &self.device); + ctx.handles.register_float_tensor::(&id, tensor) + } else if dtype.is_int() { + let tensor = B::int_from_data(data, &self.device); + ctx.handles.register_int_tensor::(&id, tensor) + } else if dtype.is_bool() { + let tensor = B::bool_from_data(data, &self.device); + ctx.handles.register_bool_tensor::(&id, tensor) + } else if let DType::QFloat(_) = dtype { + todo!(); + } + } + + /// Register a tensor and returns its intermediate representation. + pub fn register_tensor_data_desc(&mut self, data: TensorData) -> TensorIr { + let ctx = &mut self.context; + let id = ctx.create_empty_handle(); + let shape = data.shape.clone(); + let dtype = data.dtype; + + if dtype.is_float() { + let tensor = B::float_from_data(data, &self.device); + ctx.handles.register_float_tensor::(&id, tensor) + } else if dtype.is_int() { + let tensor = B::int_from_data(data, &self.device); + ctx.handles.register_int_tensor::(&id, tensor) + } else if dtype.is_bool() { + let tensor = B::bool_from_data(data, &self.device); + ctx.handles.register_bool_tensor::(&id, tensor) + } else if let DType::QFloat(_) = dtype { + todo!(); + } + + TensorIr { + id, + shape, + status: TensorStatus::ReadWrite, + dtype, + } + } + + /// Get the device default settings. + pub fn device_settings(&self) -> DeviceSettings { + self.device.defaults() + } +} + +impl TensorInterpreter { + /// Execute a tensor operation. + pub fn register_op(&mut self, op: OperationIr) { + // Remove unused tensor handles + let ctx = &mut self.context; + + let handles = &mut ctx.handles; + match &op { + // For every op: get the input(s), execute the operation and register the output(s) + OperationIr::BaseFloat(op) => match op { + BaseOperationIr::Reshape(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + + let output = B::float_reshape(tensor, desc.out.shape.clone()); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::SwapDims(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + + let output = B::float_swap_dims(tensor, desc.dim1, desc.dim2); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::Permute(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + + let output = B::float_permute(tensor, &desc.axes); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::Flip(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + + let output = B::float_flip(tensor, &desc.axes); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::Expand(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + + let output = B::float_expand(tensor, desc.out.shape.clone()); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::Unfold(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + + let output = B::float_unfold(tensor, desc.dim, desc.size, desc.step); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::Slice(desc) => { + let tensor = handles.get_float_tensor::(&desc.tensor); + + let output = B::float_slice(tensor, &desc.ranges); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::SliceAssign(desc) => { + let tensor = handles.get_float_tensor::(&desc.tensor); + let value = handles.get_float_tensor::(&desc.value); + + let output = B::float_slice_assign(tensor, &desc.ranges, value); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::Gather(desc) => { + let tensor = handles.get_float_tensor::(&desc.tensor); + let indices = handles.get_int_tensor::(&desc.indices); + + let output = B::float_gather(desc.dim, tensor, indices); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::Scatter(desc) => { + let tensor = handles.get_float_tensor::(&desc.tensor); + let indices = handles.get_int_tensor::(&desc.indices); + let value = handles.get_float_tensor::(&desc.value); + + let output = match desc.update { + IndexingUpdateOp::Add => { + B::float_scatter_add(desc.dim, tensor, indices, value) + } + _ => unimplemented!(), + }; + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::ScatterNd(desc) => { + let data = handles.get_float_tensor::(&desc.data); + let indices = handles.get_int_tensor::(&desc.indices); + let values = handles.get_float_tensor::(&desc.values); + + let output = B::float_scatter_nd(data, indices, values, desc.reduction); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::GatherNd(desc) => { + let data = handles.get_float_tensor::(&desc.data); + let indices = handles.get_int_tensor::(&desc.indices); + + let output = B::float_gather_nd(data, indices); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::Select(desc) => { + let tensor = handles.get_float_tensor::(&desc.tensor); + let indices = handles.get_int_tensor::(&desc.indices); + + let output = B::float_select(tensor, desc.dim, indices); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::SelectAssign(desc) => { + let tensor = handles.get_float_tensor::(&desc.tensor); + let indices = handles.get_int_tensor::(&desc.indices); + let value = handles.get_float_tensor::(&desc.value); + + let output = match desc.update { + IndexingUpdateOp::Add => { + B::float_select_add(tensor, desc.dim, indices, value) + } + _ => unimplemented!(), + }; + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::MaskWhere(desc) => { + let tensor = handles.get_float_tensor::(&desc.tensor); + let mask = handles.get_bool_tensor::(&desc.mask); + let value = handles.get_float_tensor::(&desc.value); + + let output = B::float_mask_where(tensor, mask, value); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::MaskFill(desc) => { + let tensor = handles.get_float_tensor::(&desc.tensor); + let mask = handles.get_bool_tensor::(&desc.mask); + + let output = B::float_mask_fill(tensor, mask, desc.value.into()); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::Equal(desc) => { + binary_float_cmp_ops!(handles, desc, B::float_equal) + } + BaseOperationIr::EqualElem(desc) => { + scalar_float_cmp_ops!(handles, desc, B::float_equal_elem) + } + BaseOperationIr::RepeatDim(desc) => { + let tensor = handles.get_float_tensor::(&desc.tensor); + + let output = B::float_repeat_dim(tensor, desc.dim, desc.times); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::Cat(desc) => { + let tensors = desc + .tensors + .iter() + .map(|tensor| handles.get_float_tensor::(tensor)) + .collect(); + + let output = B::float_cat(tensors, desc.dim); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::Cast(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + let output = B::float_cast(tensor, desc.out.dtype.into()); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::Empty(desc) => { + let shape = desc.out.shape.clone(); + let output = B::float_empty(shape, &self.device, desc.out.dtype.into()); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::Ones(desc) => { + let shape = desc.out.shape.clone(); + let output = B::float_ones(shape, &self.device, desc.out.dtype.into()); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::Zeros(desc) => { + let shape = desc.out.shape.clone(); + let output = B::float_zeros(shape, &self.device, desc.out.dtype.into()); + handles.register_float_tensor::(&desc.out.id, output); + } + BaseOperationIr::NotEqual(desc) => { + binary_float_cmp_ops!(handles, desc, B::float_not_equal) + } + BaseOperationIr::NotEqualElem(desc) => { + scalar_float_cmp_ops!(handles, desc, B::float_not_equal_elem) + } + BaseOperationIr::All(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + let output = B::float_all(tensor, desc.out.dtype.into()); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::Any(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + let output = B::float_any(tensor, desc.out.dtype.into()); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::AllDim(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + let output = B::float_all_dim(tensor, desc.axis, desc.out.dtype.into()); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::AnyDim(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + let output = B::float_any_dim(tensor, desc.axis, desc.out.dtype.into()); + handles.register_bool_tensor::(&desc.out.id, output); + } + }, + OperationIr::BaseInt(op) => match op { + BaseOperationIr::Reshape(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + + let output = B::int_reshape(tensor, desc.out.shape.clone()); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::SwapDims(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + + let output = B::int_swap_dims(tensor, desc.dim1, desc.dim2); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::Permute(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + + let output = B::int_permute(tensor, &desc.axes); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::Flip(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + + let output = B::int_flip(tensor, &desc.axes); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::Expand(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + + let output = B::int_expand(tensor, desc.out.shape.clone()); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::Unfold(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + + let output = B::int_unfold(tensor, desc.dim, desc.size, desc.step); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::Slice(desc) => { + let tensor = handles.get_int_tensor::(&desc.tensor); + + let output = B::int_slice(tensor, &desc.ranges); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::SliceAssign(desc) => { + let tensor = handles.get_int_tensor::(&desc.tensor); + let value = handles.get_int_tensor::(&desc.value); + + let output = B::int_slice_assign(tensor, &desc.ranges, value); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::Gather(desc) => { + let tensor = handles.get_int_tensor::(&desc.tensor); + let indices = handles.get_int_tensor::(&desc.indices); + + let output = B::int_gather(desc.dim, tensor, indices); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::Scatter(desc) => { + let tensor = handles.get_int_tensor::(&desc.tensor); + let indices = handles.get_int_tensor::(&desc.indices); + let value = handles.get_int_tensor::(&desc.value); + + let output = match desc.update { + IndexingUpdateOp::Add => { + B::int_scatter_add(desc.dim, tensor, indices, value) + } + _ => unimplemented!(), + }; + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::ScatterNd(desc) => { + let data = handles.get_int_tensor::(&desc.data); + let indices = handles.get_int_tensor::(&desc.indices); + let values = handles.get_int_tensor::(&desc.values); + + let output = B::int_scatter_nd(data, indices, values, desc.reduction); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::GatherNd(desc) => { + let data = handles.get_int_tensor::(&desc.data); + let indices = handles.get_int_tensor::(&desc.indices); + + let output = B::int_gather_nd(data, indices); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::Select(desc) => { + let tensor = handles.get_int_tensor::(&desc.tensor); + let indices = handles.get_int_tensor::(&desc.indices); + + let output = B::int_select(tensor, desc.dim, indices); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::SelectAssign(desc) => { + let tensor = handles.get_int_tensor::(&desc.tensor); + let indices = handles.get_int_tensor::(&desc.indices); + let value = handles.get_int_tensor::(&desc.value); + + let output = match desc.update { + IndexingUpdateOp::Add => { + B::int_select_add(tensor, desc.dim, indices, value) + } + _ => unimplemented!(), + }; + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::MaskWhere(desc) => { + let tensor = handles.get_int_tensor::(&desc.tensor); + let mask = handles.get_bool_tensor::(&desc.mask); + let value = handles.get_int_tensor::(&desc.value); + + let output = B::int_mask_where(tensor, mask, value); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::MaskFill(desc) => { + let tensor = handles.get_int_tensor::(&desc.tensor); + let mask = handles.get_bool_tensor::(&desc.mask); + + let output = B::int_mask_fill(tensor, mask, desc.value.into()); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::Equal(desc) => { + binary_int_cmp_ops!(handles, desc, B::int_equal) + } + BaseOperationIr::EqualElem(desc) => { + scalar_int_cmp_ops!(handles, desc, B::int_equal_elem) + } + BaseOperationIr::RepeatDim(desc) => { + let tensor = handles.get_int_tensor::(&desc.tensor); + + let output = B::int_repeat_dim(tensor, desc.dim, desc.times); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::Cat(desc) => { + let tensors = desc + .tensors + .iter() + .map(|tensor| handles.get_int_tensor::(tensor)) + .collect(); + + let output = B::int_cat(tensors, desc.dim); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::Cast(_) => unreachable!(), + BaseOperationIr::Empty(desc) => { + let shape = desc.out.shape.clone(); + let output = B::int_empty(shape, &self.device, desc.out.dtype.into()); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::Ones(desc) => { + let shape = desc.out.shape.clone(); + let output = B::int_ones(shape, &self.device, desc.out.dtype.into()); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::Zeros(desc) => { + let shape = desc.out.shape.clone(); + let output = B::int_zeros(shape, &self.device, desc.out.dtype.into()); + handles.register_int_tensor::(&desc.out.id, output); + } + BaseOperationIr::NotEqual(desc) => { + binary_int_cmp_ops!(handles, desc, B::int_not_equal) + } + BaseOperationIr::NotEqualElem(desc) => { + scalar_int_cmp_ops!(handles, desc, B::int_not_equal_elem) + } + BaseOperationIr::All(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + let output = B::int_all(tensor, desc.out.dtype.into()); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::Any(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + let output = B::int_any(tensor, desc.out.dtype.into()); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::AllDim(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + let output = B::int_all_dim(tensor, desc.axis, desc.out.dtype.into()); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::AnyDim(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + let output = B::int_any_dim(tensor, desc.axis, desc.out.dtype.into()); + handles.register_bool_tensor::(&desc.out.id, output); + } + }, + OperationIr::BaseBool(op) => match op { + BaseOperationIr::Reshape(desc) => { + let tensor = handles.get_bool_tensor::(&desc.input); + + let output = B::bool_reshape(tensor, desc.out.shape.clone()); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::SwapDims(desc) => { + let tensor = handles.get_bool_tensor::(&desc.input); + + let output = B::bool_swap_dims(tensor, desc.dim1, desc.dim2); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::Permute(desc) => { + let tensor = handles.get_bool_tensor::(&desc.input); + + let output = B::bool_permute(tensor, &desc.axes); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::Flip(desc) => { + let tensor = handles.get_bool_tensor::(&desc.input); + + let output = B::bool_flip(tensor, &desc.axes); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::Expand(desc) => { + let tensor = handles.get_bool_tensor::(&desc.input); + + let output = B::bool_expand(tensor, desc.out.shape.clone()); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::Unfold(desc) => { + let tensor = handles.get_bool_tensor::(&desc.input); + + let output = B::bool_unfold(tensor, desc.dim, desc.size, desc.step); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::Slice(desc) => { + let tensor = handles.get_bool_tensor::(&desc.tensor); + + let output = B::bool_slice(tensor, &desc.ranges); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::SliceAssign(desc) => { + let tensor = handles.get_bool_tensor::(&desc.tensor); + let value = handles.get_bool_tensor::(&desc.value); + + let output = B::bool_slice_assign(tensor, &desc.ranges, value); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::Gather(desc) => { + let tensor = handles.get_bool_tensor::(&desc.tensor); + let indices = handles.get_int_tensor::(&desc.indices); + + let output = B::bool_gather(desc.dim, tensor, indices); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::Scatter(desc) => { + let tensor = handles.get_bool_tensor::(&desc.tensor); + let indices = handles.get_int_tensor::(&desc.indices); + let value = handles.get_bool_tensor::(&desc.value); + + let output = match desc.update { + IndexingUpdateOp::Add => { + B::bool_scatter_or(desc.dim, tensor, indices, value) + } + _ => unimplemented!(), + }; + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::ScatterNd(_) => { + unreachable!("scatter_nd not supported for bool tensors") + } + BaseOperationIr::GatherNd(_) => { + unreachable!("gather_nd not supported for bool tensors") + } + BaseOperationIr::Select(desc) => { + let tensor = handles.get_bool_tensor::(&desc.tensor); + let indices = handles.get_int_tensor::(&desc.indices); + + let output = B::bool_select(tensor, desc.dim, indices); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::SelectAssign(desc) => { + let tensor = handles.get_bool_tensor::(&desc.tensor); + let indices = handles.get_int_tensor::(&desc.indices); + let value = handles.get_bool_tensor::(&desc.value); + + let output = match desc.update { + IndexingUpdateOp::Add => { + B::bool_select_or(tensor, desc.dim, indices, value) + } + _ => unimplemented!(), + }; + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::MaskWhere(desc) => { + let tensor = handles.get_bool_tensor::(&desc.tensor); + let mask = handles.get_bool_tensor::(&desc.mask); + let value = handles.get_bool_tensor::(&desc.value); + + let output = B::bool_mask_where(tensor, mask, value); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::MaskFill(desc) => { + let tensor = handles.get_bool_tensor::(&desc.tensor); + let mask = handles.get_bool_tensor::(&desc.mask); + + let output = B::bool_mask_fill(tensor, mask, desc.value.into()); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::Equal(desc) => { + let lhs = handles.get_bool_tensor::(&desc.lhs); + let rhs = handles.get_bool_tensor::(&desc.rhs); + + let output = B::bool_equal(lhs, rhs); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::EqualElem(desc) => { + let lhs = handles.get_bool_tensor::(&desc.lhs); + + let output = B::bool_equal_elem(lhs, desc.rhs.into()); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::RepeatDim(desc) => { + let tensor = handles.get_bool_tensor::(&desc.tensor); + + let output = B::bool_repeat_dim(tensor, desc.dim, desc.times); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::Cat(desc) => { + let tensors = desc + .tensors + .iter() + .map(|tensor| handles.get_bool_tensor::(tensor)) + .collect(); + + let output = B::bool_cat(tensors, desc.dim); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::Cast(_) => unreachable!(), + BaseOperationIr::Empty(desc) => { + let shape = desc.out.shape.clone(); + let output = B::bool_empty(shape, &self.device, desc.out.dtype.into()); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::Zeros(desc) => { + let shape = desc.out.shape.clone(); + let output = B::bool_zeros(shape, &self.device, desc.out.dtype.into()); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::Ones(desc) => { + let shape = desc.out.shape.clone(); + let output = B::bool_ones(shape, &self.device, desc.out.dtype.into()); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::NotEqual(desc) => { + let lhs = handles.get_bool_tensor::(&desc.lhs); + let rhs = handles.get_bool_tensor::(&desc.rhs); + let output = B::bool_not_equal(lhs, rhs); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::NotEqualElem(desc) => { + let lhs = handles.get_bool_tensor::(&desc.lhs); + let output = B::bool_not_equal_elem(lhs, desc.rhs.into()); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::All(desc) => { + let tensor = handles.get_bool_tensor::(&desc.input); + let output = B::bool_all(tensor); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::Any(desc) => { + let tensor = handles.get_bool_tensor::(&desc.input); + let output = B::bool_any(tensor); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::AllDim(desc) => { + let tensor = handles.get_bool_tensor::(&desc.input); + let output = B::bool_all_dim(tensor, desc.axis); + handles.register_bool_tensor::(&desc.out.id, output); + } + BaseOperationIr::AnyDim(desc) => { + let tensor = handles.get_bool_tensor::(&desc.input); + let output = B::bool_any_dim(tensor, desc.axis); + handles.register_bool_tensor::(&desc.out.id, output); + } + }, + OperationIr::NumericFloat(_dtype, op) => match op { + NumericOperationIr::Add(desc) => { + binary_float_ops!(handles, desc, B::float_add) + } + NumericOperationIr::AddScalar(desc) => { + scalar_float_ops!(handles, desc, B::float_add_scalar) + } + NumericOperationIr::Sub(desc) => { + binary_float_ops!(handles, desc, B::float_sub) + } + NumericOperationIr::SubScalar(desc) => { + scalar_float_ops!(handles, desc, B::float_sub_scalar) + } + NumericOperationIr::Div(desc) => { + binary_float_ops!(handles, desc, B::float_div) + } + NumericOperationIr::DivScalar(desc) => { + scalar_float_ops!(handles, desc, B::float_div_scalar) + } + NumericOperationIr::Rem(desc) => { + binary_float_ops!(handles, desc, B::float_remainder) + } + NumericOperationIr::RemScalar(desc) => { + scalar_float_ops!(handles, desc, B::float_remainder_scalar) + } + NumericOperationIr::Mul(desc) => { + binary_float_ops!(handles, desc, B::float_mul) + } + NumericOperationIr::MulScalar(desc) => { + scalar_float_ops!(handles, desc, B::float_mul_scalar) + } + NumericOperationIr::Abs(desc) => { + unary_float_ops!(handles, desc, B::float_abs) + } + NumericOperationIr::Full(desc) => { + let shape = desc.out.shape.clone(); + let output = B::float_full( + shape, + desc.value.into(), + &self.device, + desc.out.dtype.into(), + ); + handles.register_float_tensor::(&desc.out.id, output); + } + NumericOperationIr::MeanDim(desc) => { + reduce_float_dim_ops!(handles, desc, |tensor, axis, _| B::float_mean_dim( + tensor, axis + )) + } + NumericOperationIr::Mean(desc) => { + unary_float_ops!(handles, desc, B::float_mean) + } + NumericOperationIr::Sum(desc) => { + unary_float_ops!(handles, desc, B::float_sum) + } + NumericOperationIr::SumDim(desc) => { + reduce_float_dim_ops!(handles, desc, |tensor, axis, _| B::float_sum_dim( + tensor, axis + )) + } + NumericOperationIr::Prod(desc) => { + unary_float_ops!(handles, desc, B::float_prod) + } + NumericOperationIr::ProdDim(desc) => { + reduce_float_dim_ops!(handles, desc, |tensor, axis, _| B::float_prod_dim( + tensor, axis + )) + } + NumericOperationIr::Greater(desc) => { + binary_float_cmp_ops!(handles, desc, B::float_greater) + } + NumericOperationIr::GreaterElem(desc) => { + scalar_float_cmp_ops!(handles, desc, B::float_greater_elem) + } + NumericOperationIr::GreaterEqual(desc) => { + binary_float_cmp_ops!(handles, desc, B::float_greater_equal) + } + NumericOperationIr::GreaterEqualElem(desc) => { + scalar_float_cmp_ops!(handles, desc, B::float_greater_equal_elem) + } + NumericOperationIr::Lower(desc) => { + binary_float_cmp_ops!(handles, desc, B::float_lower) + } + NumericOperationIr::LowerElem(desc) => { + scalar_float_cmp_ops!(handles, desc, B::float_lower_elem) + } + NumericOperationIr::LowerEqual(desc) => { + binary_float_cmp_ops!(handles, desc, B::float_lower_equal) + } + NumericOperationIr::LowerEqualElem(desc) => { + scalar_float_cmp_ops!(handles, desc, B::float_lower_equal_elem) + } + NumericOperationIr::ArgMax(desc) => { + reduce_float2int_dim_ops!(handles, desc, |tensor, axis, _, dtype| { + B::float_argmax(tensor, axis, dtype) + }) + } + NumericOperationIr::ArgTopK(desc) => { + reduce_float2int_dim_ops!(handles, desc, B::float_argtopk) + } + NumericOperationIr::ArgMin(desc) => { + reduce_float2int_dim_ops!(handles, desc, |tensor, axis, _, dtype| { + B::float_argmin(tensor, axis, dtype) + }) + } + NumericOperationIr::Max(desc) => { + unary_float_ops!(handles, desc, B::float_max) + } + NumericOperationIr::MaxDimWithIndices(desc) => { + let tensor = handles.get_float_tensor::(&desc.tensor); + + let (output, output_idx) = B::float_max_dim_with_indices( + tensor, + desc.dim, + desc.out_indices.dtype.into(), + ); + handles.register_float_tensor::(&desc.out.id, output); + handles.register_int_tensor::(&desc.out_indices.id, output_idx); + } + NumericOperationIr::MinDimWithIndices(desc) => { + let tensor = handles.get_float_tensor::(&desc.tensor); + + let (output, output_idx) = B::float_min_dim_with_indices( + tensor, + desc.dim, + desc.out_indices.dtype.into(), + ); + handles.register_float_tensor::(&desc.out.id, output); + handles.register_int_tensor::(&desc.out_indices.id, output_idx); + } + NumericOperationIr::Min(desc) => { + unary_float_ops!(handles, desc, B::float_min) + } + NumericOperationIr::MaxDim(desc) => { + reduce_float_dim_ops!(handles, desc, |tensor, axis, _| B::float_max_dim( + tensor, axis + )) + } + NumericOperationIr::TopK(desc) => { + reduce_float_dim_ops!(handles, desc, B::float_topk) + } + NumericOperationIr::MinDim(desc) => { + reduce_float_dim_ops!(handles, desc, |tensor, axis, _| B::float_min_dim( + tensor, axis + )) + } + NumericOperationIr::MaxAbs(desc) => { + unary_float_ops!(handles, desc, B::float_max_abs) + } + NumericOperationIr::MaxAbsDim(desc) => { + reduce_float_dim_ops!(handles, desc, |tensor, axis, _| B::float_max_abs_dim( + tensor, axis + )) + } + NumericOperationIr::Clamp(desc) => { + let tensor = handles.get_float_tensor::(&desc.tensor); + + let output = B::float_clamp(tensor, desc.min.into(), desc.max.into()); + handles.register_float_tensor::(&desc.out.id, output); + } + NumericOperationIr::IntRandom(_) => unreachable!(), + NumericOperationIr::Powi(desc) => { + let lhs = handles.get_float_tensor::(&desc.lhs); + let rhs = handles.get_int_tensor::(&desc.rhs); + let output = (B::float_powi)(lhs, rhs); + handles.register_float_tensor::(&desc.out.id, output); + } + NumericOperationIr::PowiScalar(desc) => { + scalar_float_ops!(handles, desc, B::float_powi_scalar) + } + NumericOperationIr::CumSum(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + let output = B::float_cumsum(tensor, desc.axis); + handles.register_float_tensor::(&desc.out.id, output); + } + NumericOperationIr::CumProd(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + let output = B::float_cumprod(tensor, desc.axis); + handles.register_float_tensor::(&desc.out.id, output); + } + NumericOperationIr::CumMin(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + let output = B::float_cummin(tensor, desc.axis); + handles.register_float_tensor::(&desc.out.id, output); + } + NumericOperationIr::CumMax(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + let output = B::float_cummax(tensor, desc.axis); + handles.register_float_tensor::(&desc.out.id, output); + } + NumericOperationIr::Neg(desc) => { + unary_float_ops!(handles, desc, B::float_neg) + } + NumericOperationIr::Sign(desc) => { + unary_float_ops!(handles, desc, B::float_sign) + } + NumericOperationIr::ClampMin(desc) => { + scalar_float_ops!(handles, desc, B::float_clamp_min) + } + NumericOperationIr::ClampMax(desc) => { + scalar_float_ops!(handles, desc, B::float_clamp_max) + } + NumericOperationIr::Sort(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + let output = B::float_sort(tensor, desc.dim, desc.descending); + handles.register_float_tensor::(&desc.out.id, output); + } + NumericOperationIr::SortWithIndices(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + let (values, indices) = B::float_sort_with_indices( + tensor, + desc.dim, + desc.descending, + desc.out_indices.dtype.into(), + ); + handles.register_float_tensor::(&desc.out.id, values); + handles.register_int_tensor::(&desc.out_indices.id, indices); + } + NumericOperationIr::ArgSort(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + let output = + B::float_argsort(tensor, desc.dim, desc.descending, desc.out.dtype.into()); + handles.register_int_tensor::(&desc.out.id, output); + } + }, + OperationIr::NumericInt(_dtype, op) => match op { + NumericOperationIr::Add(desc) => { + binary_int_ops!(handles, desc, B::int_add) + } + NumericOperationIr::AddScalar(desc) => { + scalar_int_ops!(handles, desc, B::int_add_scalar) + } + NumericOperationIr::Sub(desc) => { + binary_int_ops!(handles, desc, B::int_sub) + } + NumericOperationIr::SubScalar(desc) => { + scalar_int_ops!(handles, desc, B::int_sub_scalar) + } + NumericOperationIr::Div(desc) => { + binary_int_ops!(handles, desc, B::int_div) + } + NumericOperationIr::DivScalar(desc) => { + scalar_int_ops!(handles, desc, B::int_div_scalar) + } + NumericOperationIr::Rem(desc) => { + binary_int_ops!(handles, desc, B::int_remainder) + } + NumericOperationIr::RemScalar(desc) => { + scalar_int_ops!(handles, desc, B::int_remainder_scalar) + } + NumericOperationIr::Mul(desc) => { + binary_int_ops!(handles, desc, B::int_mul) + } + NumericOperationIr::MulScalar(desc) => { + scalar_int_ops!(handles, desc, B::int_mul_scalar) + } + NumericOperationIr::Abs(desc) => { + unary_int_ops!(handles, desc, B::int_abs) + } + NumericOperationIr::Full(desc) => { + let shape = desc.out.shape.clone(); + let output = B::int_full( + shape, + desc.value.into(), + &self.device, + desc.out.dtype.into(), + ); + handles.register_int_tensor::(&desc.out.id, output); + } + NumericOperationIr::MeanDim(desc) => { + reduce_int_dim_ops!(handles, desc, |tensor, axis, _| B::int_mean_dim( + tensor, axis + )) + } + NumericOperationIr::Mean(desc) => { + unary_int_ops!(handles, desc, B::int_mean) + } + NumericOperationIr::Sum(desc) => { + unary_int_ops!(handles, desc, B::int_sum) + } + NumericOperationIr::SumDim(desc) => { + reduce_int_dim_ops!(handles, desc, |tensor, axis, _| B::int_sum_dim( + tensor, axis + )) + } + NumericOperationIr::Prod(desc) => { + unary_int_ops!(handles, desc, B::int_prod) + } + NumericOperationIr::ProdDim(desc) => { + reduce_int_dim_ops!(handles, desc, |tensor, axis, _| B::int_prod_dim( + tensor, axis + )) + } + NumericOperationIr::Greater(desc) => { + binary_int_cmp_ops!(handles, desc, B::int_greater) + } + NumericOperationIr::GreaterElem(desc) => { + scalar_int_cmp_ops!(handles, desc, B::int_greater_elem) + } + NumericOperationIr::GreaterEqual(desc) => { + binary_int_cmp_ops!(handles, desc, B::int_greater_equal) + } + NumericOperationIr::GreaterEqualElem(desc) => { + scalar_int_cmp_ops!(handles, desc, B::int_greater_equal_elem) + } + NumericOperationIr::Lower(desc) => { + binary_int_cmp_ops!(handles, desc, B::int_lower) + } + NumericOperationIr::LowerElem(desc) => { + scalar_int_cmp_ops!(handles, desc, B::int_lower_elem) + } + NumericOperationIr::LowerEqual(desc) => { + binary_int_cmp_ops!(handles, desc, B::int_lower_equal) + } + NumericOperationIr::LowerEqualElem(desc) => { + scalar_int_cmp_ops!(handles, desc, B::int_lower_equal_elem) + } + NumericOperationIr::ArgMax(desc) => { + reduce_int_dim_ops!(handles, desc, |tensor, axis, _| B::int_argmax( + tensor, axis + )) + } + NumericOperationIr::ArgTopK(desc) => { + reduce_int_dim_ops!(handles, desc, B::int_argtopk) + } + NumericOperationIr::ArgMin(desc) => { + reduce_int_dim_ops!(handles, desc, |tensor, axis, _| B::int_argmin( + tensor, axis + )) + } + NumericOperationIr::Max(desc) => { + unary_int_ops!(handles, desc, B::int_max) + } + NumericOperationIr::MaxDimWithIndices(desc) => { + let tensor = handles.get_int_tensor::(&desc.tensor); + + let (output, output_idx) = B::int_max_dim_with_indices(tensor, desc.dim); + handles.register_int_tensor::(&desc.out.id, output); + handles.register_int_tensor::(&desc.out_indices.id, output_idx); + } + NumericOperationIr::MinDimWithIndices(desc) => { + let tensor = handles.get_int_tensor::(&desc.tensor); + + let (output, output_idx) = B::int_min_dim_with_indices(tensor, desc.dim); + handles.register_int_tensor::(&desc.out.id, output); + handles.register_int_tensor::(&desc.out_indices.id, output_idx); + } + NumericOperationIr::Min(desc) => { + unary_int_ops!(handles, desc, B::int_min) + } + NumericOperationIr::MaxDim(desc) => { + reduce_int_dim_ops!(handles, desc, |tensor, axis, _| B::int_max_dim( + tensor, axis + )) + } + NumericOperationIr::TopK(desc) => { + reduce_int_dim_ops!(handles, desc, B::int_topk) + } + NumericOperationIr::MinDim(desc) => { + reduce_int_dim_ops!(handles, desc, |tensor, axis, _| B::int_min_dim( + tensor, axis + )) + } + NumericOperationIr::MaxAbs(desc) => { + unary_int_ops!(handles, desc, B::int_max_abs) + } + NumericOperationIr::MaxAbsDim(desc) => { + reduce_int_dim_ops!(handles, desc, |tensor, axis, _| B::int_max_abs_dim( + tensor, axis + )) + } + NumericOperationIr::Clamp(desc) => { + let tensor = handles.get_int_tensor::(&desc.tensor); + + let output = B::int_clamp(tensor, desc.min.into(), desc.max.into()); + handles.register_int_tensor::(&desc.out.id, output); + } + NumericOperationIr::IntRandom(desc) => { + let shape = desc.out.shape.clone(); + + let output = B::int_random( + shape, + desc.distribution, + &self.device, + desc.out.dtype.into(), + ); + handles.register_int_tensor::(&desc.out.id, output); + } + NumericOperationIr::Powi(desc) => { + let lhs = handles.get_int_tensor::(&desc.lhs); + let rhs = handles.get_int_tensor::(&desc.rhs); + + let output = B::int_powi(lhs, rhs); + handles.register_int_tensor::(&desc.out.id, output); + } + NumericOperationIr::PowiScalar(desc) => { + scalar_int_ops!(handles, desc, B::int_powi_scalar) + } + NumericOperationIr::CumSum(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + let output = B::int_cumsum(tensor, desc.axis); + handles.register_int_tensor::(&desc.out.id, output); + } + NumericOperationIr::CumProd(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + let output = B::int_cumprod(tensor, desc.axis); + handles.register_int_tensor::(&desc.out.id, output); + } + NumericOperationIr::CumMin(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + let output = B::int_cummin(tensor, desc.axis); + handles.register_int_tensor::(&desc.out.id, output); + } + NumericOperationIr::CumMax(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + let output = B::int_cummax(tensor, desc.axis); + handles.register_int_tensor::(&desc.out.id, output); + } + NumericOperationIr::Neg(desc) => { + unary_int_ops!(handles, desc, B::int_neg) + } + NumericOperationIr::Sign(desc) => { + unary_int_ops!(handles, desc, B::int_sign) + } + NumericOperationIr::ClampMin(desc) => { + scalar_int_ops!(handles, desc, B::int_clamp_min) + } + NumericOperationIr::ClampMax(desc) => { + scalar_int_ops!(handles, desc, B::int_clamp_max) + } + NumericOperationIr::Sort(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + let output = B::int_sort(tensor, desc.dim, desc.descending); + handles.register_int_tensor::(&desc.out.id, output); + } + NumericOperationIr::SortWithIndices(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + let (values, indices) = + B::int_sort_with_indices(tensor, desc.dim, desc.descending); + handles.register_int_tensor::(&desc.out.id, values); + handles.register_int_tensor::(&desc.out_indices.id, indices); + } + NumericOperationIr::ArgSort(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + let output = B::int_argsort(tensor, desc.dim, desc.descending); + handles.register_int_tensor::(&desc.out.id, output); + } + }, + OperationIr::Bool(op) => match op { + BoolOperationIr::IntoFloat(desc) => { + let tensor = handles.get_bool_tensor::(&desc.input); + + let output = B::bool_into_float(tensor, desc.out.dtype.into()); + handles.register_float_tensor::(&desc.out.id, output); + } + BoolOperationIr::IntoInt(desc) => { + let tensor = handles.get_bool_tensor::(&desc.input); + + let output = B::bool_into_int(tensor, desc.out.dtype.into()); + handles.register_int_tensor::(&desc.out.id, output); + } + BoolOperationIr::Not(desc) => { + let tensor = handles.get_bool_tensor::(&desc.input); + + let output = B::bool_not(tensor); + handles.register_bool_tensor::(&desc.out.id, output); + } + BoolOperationIr::And(desc) => { + binary_bool_ops!(handles, desc, B::bool_and) + } + BoolOperationIr::Or(desc) => { + binary_bool_ops!(handles, desc, B::bool_or) + } + BoolOperationIr::Xor(desc) => { + binary_bool_ops!(handles, desc, B::bool_xor) + } + }, + OperationIr::Int(op) => match op { + IntOperationIr::IntoFloat(desc) => { + let tensor = handles.get_int_tensor::(&desc.input); + + let output = B::int_into_float(tensor, desc.out.dtype.into()); + handles.register_float_tensor::(&desc.out.id, output); + } + IntOperationIr::Matmul(desc) => { + binary_int_ops!(handles, desc, B::int_matmul) + } + IntOperationIr::BitwiseAnd(desc) => { + binary_int_ops!(handles, desc, B::bitwise_and) + } + IntOperationIr::BitwiseAndScalar(desc) => { + scalar_int_ops!(handles, desc, B::bitwise_and_scalar) + } + IntOperationIr::BitwiseOr(desc) => { + binary_int_ops!(handles, desc, B::bitwise_or) + } + IntOperationIr::BitwiseOrScalar(desc) => { + scalar_int_ops!(handles, desc, B::bitwise_or_scalar) + } + IntOperationIr::BitwiseXor(desc) => { + binary_int_ops!(handles, desc, B::bitwise_xor) + } + IntOperationIr::BitwiseXorScalar(desc) => { + scalar_int_ops!(handles, desc, B::bitwise_xor_scalar) + } + IntOperationIr::BitwiseNot(desc) => { + unary_int_ops!(handles, desc, B::bitwise_not) + } + IntOperationIr::BitwiseLeftShift(desc) => { + binary_int_ops!(handles, desc, B::bitwise_left_shift) + } + IntOperationIr::BitwiseRightShift(desc) => { + binary_int_ops!(handles, desc, B::bitwise_right_shift) + } + IntOperationIr::BitwiseLeftShiftScalar(desc) => { + scalar_int_ops!(handles, desc, B::bitwise_left_shift_scalar) + } + IntOperationIr::BitwiseRightShiftScalar(desc) => { + scalar_int_ops!(handles, desc, B::bitwise_right_shift_scalar) + } + }, + OperationIr::Float(_dtype, op) => match op { + FloatOperationIr::Exp(desc) => { + unary_float_ops!(handles, desc, B::float_exp) + } + FloatOperationIr::Powf(desc) => { + binary_float_ops!(handles, desc, B::float_powf) + } + FloatOperationIr::Log(desc) => { + unary_float_ops!(handles, desc, B::float_log) + } + FloatOperationIr::Log1p(desc) => { + unary_float_ops!(handles, desc, B::float_log1p) + } + FloatOperationIr::Erf(desc) => { + unary_float_ops!(handles, desc, B::float_erf) + } + FloatOperationIr::PowfScalar(desc) => { + scalar_float_ops!(handles, desc, B::float_powf_scalar) + } + FloatOperationIr::Sqrt(desc) => { + unary_float_ops!(handles, desc, B::float_sqrt) + } + FloatOperationIr::Cos(desc) => { + unary_float_ops!(handles, desc, B::float_cos) + } + FloatOperationIr::Sin(desc) => { + unary_float_ops!(handles, desc, B::float_sin) + } + FloatOperationIr::Tanh(desc) => { + unary_float_ops!(handles, desc, B::float_tanh) + } + FloatOperationIr::Tan(desc) => unary_float_ops!(handles, desc, B::float_tan), + FloatOperationIr::Cosh(desc) => unary_float_ops!(handles, desc, B::float_cosh), + FloatOperationIr::Sinh(desc) => unary_float_ops!(handles, desc, B::float_sinh), + FloatOperationIr::ArcCos(desc) => unary_float_ops!(handles, desc, B::float_acos), + FloatOperationIr::ArcCosh(desc) => unary_float_ops!(handles, desc, B::float_acosh), + FloatOperationIr::ArcSin(desc) => unary_float_ops!(handles, desc, B::float_asin), + FloatOperationIr::ArcSinh(desc) => unary_float_ops!(handles, desc, B::float_asinh), + FloatOperationIr::ArcTan(desc) => unary_float_ops!(handles, desc, B::float_atan), + FloatOperationIr::ArcTanh(desc) => unary_float_ops!(handles, desc, B::float_atanh), + FloatOperationIr::ArcTan2(desc) => binary_float_ops!(handles, desc, B::float_atan2), + FloatOperationIr::Hypot(desc) => binary_float_ops!(handles, desc, B::float_hypot), + FloatOperationIr::Round(desc) => { + unary_float_ops!(handles, desc, B::float_round) + } + FloatOperationIr::Floor(desc) => { + unary_float_ops!(handles, desc, B::float_floor) + } + FloatOperationIr::Ceil(desc) => { + unary_float_ops!(handles, desc, B::float_ceil) + } + FloatOperationIr::Trunc(desc) => { + unary_float_ops!(handles, desc, B::float_trunc) + } + FloatOperationIr::IntoInt(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + + let output = B::float_into_int(tensor, desc.out.dtype.into()); + handles.register_int_tensor::(&desc.out.id, output); + } + FloatOperationIr::Matmul(desc) => { + binary_float_ops!(handles, desc, B::float_matmul) + } + FloatOperationIr::Cross(desc) => { + let lhs = handles.get_float_tensor::(&desc.lhs); + let rhs = handles.get_float_tensor::(&desc.rhs); + let output = B::float_cross(lhs, rhs, desc.dim); + handles.register_float_tensor::(&desc.out.id, output); + } + FloatOperationIr::Random(desc) => { + let shape = desc.out.shape.clone(); + + let output = B::float_random( + shape, + desc.distribution, + &self.device, + desc.out.dtype.into(), + ); + handles.register_float_tensor::(&desc.out.id, output); + } + FloatOperationIr::Recip(desc) => { + unary_float_ops!(handles, desc, B::float_recip) + } + FloatOperationIr::Quantize(_) => todo!(), + FloatOperationIr::Dequantize(_) => todo!(), + FloatOperationIr::IsNan(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + + let output = B::float_is_nan(tensor, desc.out.dtype.into()); + handles.register_bool_tensor::(&desc.out.id, output); + } + FloatOperationIr::IsInf(desc) => { + let tensor = handles.get_float_tensor::(&desc.input); + + let output = B::float_is_inf(tensor, desc.out.dtype.into()); + handles.register_bool_tensor::(&desc.out.id, output); + } + FloatOperationIr::GridSample2d(desc) => { + let tensor = handles.get_float_tensor::(&desc.tensor); + let grid = handles.get_float_tensor::(&desc.grid); + + let output = B::float_grid_sample_2d(tensor, grid, desc.options.clone().into()); + handles.register_float_tensor::(&desc.out.id, output); + } + }, + OperationIr::Module(op) => match op { + ModuleOperationIr::Embedding(desc) => { + let weights = handles.get_float_tensor::(&desc.weights); + let indices = handles.get_int_tensor::(&desc.indices); + + let output = B::embedding(weights, indices); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::EmbeddingBackward(desc) => { + let weights = handles.get_float_tensor::(&desc.weights); + let indices = handles.get_int_tensor::(&desc.indices); + let output_grad = handles.get_float_tensor::(&desc.out_grad); + + let output = B::embedding_backward(weights, output_grad, indices); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::Linear(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let bias = desc + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + + let output = B::linear(x, weight, bias); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::LinearXBackward(desc) => { + let weight = handles.get_float_tensor::(&desc.weight); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + + let output = B::linear_x_backward(weight, output_grad); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::LinearWeightBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + + let output = B::linear_weight_backward(x, output_grad); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::LinearBiasBackward(desc) => { + let output_grad = handles.get_float_tensor::(&desc.output_grad); + + let output = B::linear_bias_backward(output_grad); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::Conv1d(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let bias = desc + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + + let output = B::conv1d(x, weight, bias, desc.clone().options.into()); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::Conv1dXBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + + let output = + B::conv1d_x_backward(x, weight, output_grad, desc.clone().options.into()); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::Conv1dWeightBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + + let output = B::conv1d_weight_backward( + x, + weight, + output_grad, + desc.clone().options.into(), + ); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::Conv1dBiasBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let bias = handles.get_float_tensor::(&desc.bias); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + + let output = B::conv1d_bias_backward(x, bias, output_grad); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::Conv2d(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let bias = desc + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + + let output = B::conv2d(x, weight, bias, desc.clone().options.into()); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::Conv2dXBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + + let output = + B::conv2d_x_backward(x, weight, output_grad, desc.clone().options.into()); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::Conv2dWeightBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + + let output = B::conv2d_weight_backward( + x, + weight, + output_grad, + desc.clone().options.into(), + ); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::Conv2dBiasBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let bias = handles.get_float_tensor::(&desc.bias); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + + let output = B::conv2d_bias_backward(x, bias, output_grad); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::Conv3d(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let bias = desc + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + + let output = B::conv3d(x, weight, bias, desc.options.clone().into()); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::Conv3dXBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + + let output = + B::conv3d_x_backward(x, weight, output_grad, desc.clone().options.into()); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::Conv3dWeightBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + + let output = B::conv3d_weight_backward( + x, + weight, + output_grad, + desc.clone().options.into(), + ); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::Conv3dBiasBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let bias = handles.get_float_tensor::(&desc.bias); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + + let output = B::conv3d_bias_backward(x, bias, output_grad); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::DeformableConv2d(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let offset = handles.get_float_tensor::(&desc.offset); + let mask = desc + .mask + .as_ref() + .map(|mask| handles.get_float_tensor::(mask)); + let weight = handles.get_float_tensor::(&desc.weight); + let bias = desc + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + + let output = B::deform_conv2d( + x, + offset, + weight, + mask, + bias, + desc.options.clone().into(), + ); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::DeformableConv2dBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let offset = handles.get_float_tensor::(&desc.offset); + let mask = desc + .mask + .as_ref() + .map(|mask| handles.get_float_tensor::(mask)); + let weight = handles.get_float_tensor::(&desc.weight); + let bias = desc + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + let output_grad = handles.get_float_tensor::(&desc.out_grad); + + let output = B::deform_conv2d_backward( + x, + offset, + weight, + mask, + bias, + output_grad, + desc.options.clone().into(), + ); + + handles.register_float_tensor::(&desc.input_grad.id, output.x_grad); + handles.register_float_tensor::(&desc.offset_grad.id, output.offset_grad); + handles.register_float_tensor::(&desc.weight_grad.id, output.weight_grad); + if let Some((mask_grad, field)) = output.mask_grad.zip(desc.mask_grad.as_ref()) + { + handles.register_float_tensor::(&field.id, mask_grad); + } + if let Some((bias_grad, field)) = output.bias_grad.zip(desc.bias_grad.as_ref()) + { + handles.register_float_tensor::(&field.id, bias_grad); + } + } + ModuleOperationIr::ConvTranspose1d(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let bias = desc + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + + let output = B::conv_transpose1d(x, weight, bias, desc.options.clone().into()); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::ConvTranspose2d(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let bias = desc + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + + let output = B::conv_transpose2d(x, weight, bias, desc.options.clone().into()); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::ConvTranspose3d(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let bias = desc + .bias + .as_ref() + .map(|bias| handles.get_float_tensor::(bias)); + + let output = B::conv_transpose3d(x, weight, bias, desc.options.clone().into()); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::AvgPool1d(desc) => { + let x = handles.get_float_tensor::(&desc.x); + + let output = B::avg_pool1d( + x, + desc.kernel_size, + desc.stride, + desc.padding, + desc.count_include_pad, + desc.ceil_mode, + ); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::AvgPool2d(desc) => { + let x = handles.get_float_tensor::(&desc.x); + + let output = B::avg_pool2d( + x, + desc.kernel_size, + desc.stride, + desc.padding, + desc.count_include_pad, + desc.ceil_mode, + ); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::AvgPool1dBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let grad = handles.get_float_tensor::(&desc.grad); + + let output = B::avg_pool1d_backward( + x, + grad, + desc.kernel_size, + desc.stride, + desc.padding, + desc.count_include_pad, + desc.ceil_mode, + ); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::AvgPool2dBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let grad = handles.get_float_tensor::(&desc.grad); + + let output = B::avg_pool2d_backward( + x, + grad, + desc.kernel_size, + desc.stride, + desc.padding, + desc.count_include_pad, + desc.ceil_mode, + ); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::AdaptiveAvgPool1d(desc) => { + let x = handles.get_float_tensor::(&desc.x); + + let output = B::adaptive_avg_pool1d(x, desc.output_size); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::AdaptiveAvgPool2d(desc) => { + let x = handles.get_float_tensor::(&desc.x); + + let output = B::adaptive_avg_pool2d(x, desc.output_size); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::AdaptiveAvgPool1dBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let grad = handles.get_float_tensor::(&desc.grad); + + let output = B::adaptive_avg_pool1d_backward(x, grad); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::AdaptiveAvgPool2dBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let grad = handles.get_float_tensor::(&desc.grad); + + let output = B::adaptive_avg_pool2d_backward(x, grad); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::MaxPool1d(desc) => { + let x = handles.get_float_tensor::(&desc.x); + + let output = B::max_pool1d( + x, + desc.kernel_size, + desc.stride, + desc.padding, + desc.dilation, + desc.ceil_mode, + ); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::MaxPool1dWithIndices(desc) => { + let x = handles.get_float_tensor::(&desc.x); + + let output = B::max_pool1d_with_indices( + x, + desc.kernel_size, + desc.stride, + desc.padding, + desc.dilation, + desc.ceil_mode, + desc.out_indices.dtype.into(), + ); + handles.register_float_tensor::(&desc.out.id, output.output); + handles.register_int_tensor::(&desc.out_indices.id, output.indices); + } + ModuleOperationIr::MaxPool1dWithIndicesBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let output_grad = handles.get_float_tensor::(&desc.grad); + let indices = handles.get_int_tensor::(&desc.indices); + + let output = B::max_pool1d_with_indices_backward( + x, + desc.kernel_size, + desc.stride, + desc.padding, + desc.dilation, + desc.ceil_mode, + output_grad, + indices, + ); + handles.register_float_tensor::(&desc.out.id, output.x_grad); + } + ModuleOperationIr::MaxPool2d(desc) => { + let x = handles.get_float_tensor::(&desc.x); + + let output = B::max_pool2d( + x, + desc.kernel_size, + desc.stride, + desc.padding, + desc.dilation, + desc.ceil_mode, + ); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::MaxPool2dWithIndices(desc) => { + let x = handles.get_float_tensor::(&desc.x); + + let output = B::max_pool2d_with_indices( + x, + desc.kernel_size, + desc.stride, + desc.padding, + desc.dilation, + desc.ceil_mode, + desc.out_indices.dtype.into(), + ); + handles.register_float_tensor::(&desc.out.id, output.output); + handles.register_int_tensor::(&desc.out_indices.id, output.indices); + } + ModuleOperationIr::MaxPool2dWithIndicesBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let output_grad = handles.get_float_tensor::(&desc.grad); + let indices = handles.get_int_tensor::(&desc.indices); + + let output = B::max_pool2d_with_indices_backward( + x, + desc.kernel_size, + desc.stride, + desc.padding, + desc.dilation, + desc.ceil_mode, + output_grad, + indices, + ); + handles.register_float_tensor::(&desc.out.id, output.x_grad); + } + ModuleOperationIr::Interpolate(desc) => { + let x = handles.get_float_tensor::(&desc.x); + + let output = B::interpolate(x, desc.output_size, desc.options.clone().into()); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::InterpolateBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let grad = handles.get_float_tensor::(&desc.grad); + + let output = B::interpolate_backward( + x, + grad, + desc.output_size, + desc.options.clone().into(), + ); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::Rfft(desc) => { + let signal = handles.get_float_tensor::(&desc.signal); + let (out_re, out_im) = B::rfft(signal, desc.dim, desc.n); + + handles.register_float_tensor::(&desc.out_re.id, out_re); + handles.register_float_tensor::(&desc.out_im.id, out_im); + } + ModuleOperationIr::IRfft(desc) => { + let spectrum_re = handles.get_float_tensor::(&desc.input_re); + let spectrum_im = handles.get_float_tensor::(&desc.input_im); + let signal = B::irfft(spectrum_re, spectrum_im, desc.dim, desc.n); + + handles.register_float_tensor::(&desc.out_signal.id, signal); + } + ModuleOperationIr::Attention(desc) => { + let query = handles.get_float_tensor::(&desc.query); + let key = handles.get_float_tensor::(&desc.key); + let value = handles.get_float_tensor::(&desc.value); + let mask = desc.mask.as_ref().map(|m| handles.get_bool_tensor::(m)); + let attn_bias = desc + .attn_bias + .as_ref() + .map(|ab| handles.get_float_tensor::(ab)); + + let output = B::attention( + query, + key, + value, + mask, + attn_bias, + desc.options.clone().into(), + ); + + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::CtcLoss(desc) => { + let log_probs = handles.get_float_tensor::(&desc.log_probs); + let targets = handles.get_int_tensor::(&desc.targets); + let input_lengths = handles.get_int_tensor::(&desc.input_lengths); + let target_lengths = handles.get_int_tensor::(&desc.target_lengths); + + let output = B::ctc_loss( + log_probs, + targets, + input_lengths, + target_lengths, + desc.blank, + ); + + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::CtcLossBackward(desc) => { + let log_probs = handles.get_float_tensor::(&desc.log_probs); + let targets = handles.get_int_tensor::(&desc.targets); + let input_lengths = handles.get_int_tensor::(&desc.input_lengths); + let target_lengths = handles.get_int_tensor::(&desc.target_lengths); + let grad_loss = handles.get_float_tensor::(&desc.grad_loss); + + let output = B::ctc_loss_backward( + log_probs, + targets, + input_lengths, + target_lengths, + grad_loss, + desc.blank, + ); + + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::LayerNorm(desc) => { + let input = handles.get_float_tensor::(&desc.input); + let gamma = handles.get_float_tensor::(&desc.gamma); + let beta = desc.beta.as_ref().map(|b| handles.get_float_tensor::(b)); + let output = B::layer_norm(input, gamma, beta, desc.epsilon.elem()); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::Unfold4d(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let options = burn_backend::ops::UnfoldOptions::new( + desc.options.stride, + desc.options.padding, + desc.options.dilation, + ); + let output = B::unfold4d(x, desc.kernel_size, options); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::ConvTranspose1dWeightBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + let output = B::conv_transpose1d_weight_backward( + x, + weight, + output_grad, + desc.options.clone().into(), + ); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::ConvTranspose1dBiasBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let bias = handles.get_float_tensor::(&desc.bias); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + let output = B::conv_transpose1d_bias_backward(x, bias, output_grad); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::ConvTranspose2dWeightBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + let output = B::conv_transpose2d_weight_backward( + x, + weight, + output_grad, + desc.options.clone().into(), + ); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::ConvTranspose2dBiasBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let bias = handles.get_float_tensor::(&desc.bias); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + let output = B::conv_transpose2d_bias_backward(x, bias, output_grad); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::ConvTranspose3dWeightBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let weight = handles.get_float_tensor::(&desc.weight); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + let output = B::conv_transpose3d_weight_backward( + x, + weight, + output_grad, + desc.options.clone().into(), + ); + handles.register_float_tensor::(&desc.out.id, output); + } + ModuleOperationIr::ConvTranspose3dBiasBackward(desc) => { + let x = handles.get_float_tensor::(&desc.x); + let bias = handles.get_float_tensor::(&desc.bias); + let output_grad = handles.get_float_tensor::(&desc.output_grad); + let output = B::conv_transpose3d_bias_backward(x, bias, output_grad); + handles.register_float_tensor::(&desc.out.id, output); + } + }, + OperationIr::Activation(op) => match op { + ActivationOperationIr::Relu(desc) => { + let input = handles.get_float_tensor::(&desc.input); + let output = B::relu(input); + handles.register_float_tensor::(&desc.out.id, output); + } + ActivationOperationIr::ReluBackward(desc) => { + let output = handles.get_float_tensor::(&desc.lhs); + let grad = handles.get_float_tensor::(&desc.rhs); + let result = B::relu_backward(output, grad); + handles.register_float_tensor::(&desc.out.id, result); + } + ActivationOperationIr::LeakyRelu(desc) => { + let input = handles.get_float_tensor::(&desc.lhs); + let result = B::leaky_relu(input, desc.rhs.into()); + handles.register_float_tensor::(&desc.out.id, result); + } + ActivationOperationIr::PRelu(desc) => { + let input = handles.get_float_tensor::(&desc.lhs); + let alpha = handles.get_float_tensor::(&desc.rhs); + let result = B::prelu(input, alpha); + handles.register_float_tensor::(&desc.out.id, result); + } + ActivationOperationIr::Gelu(desc) => { + let input = handles.get_float_tensor::(&desc.input); + let result = B::gelu(input); + handles.register_float_tensor::(&desc.out.id, result); + } + ActivationOperationIr::GeluBackward(desc) => { + let x = handles.get_float_tensor::(&desc.lhs); + let grad = handles.get_float_tensor::(&desc.rhs); + let result = B::gelu_backward(x, grad); + handles.register_float_tensor::(&desc.out.id, result); + } + ActivationOperationIr::Sigmoid(desc) => { + let input = handles.get_float_tensor::(&desc.input); + let result = B::sigmoid(input); + handles.register_float_tensor::(&desc.out.id, result); + } + ActivationOperationIr::SigmoidBackward(desc) => { + let output = handles.get_float_tensor::(&desc.lhs); + let grad = handles.get_float_tensor::(&desc.rhs); + let result = B::sigmoid_backward(output, grad); + handles.register_float_tensor::(&desc.out.id, result); + } + ActivationOperationIr::HardSigmoid(desc) => { + let input = handles.get_float_tensor::(&desc.tensor); + let result = B::hard_sigmoid(input, desc.alpha.into(), desc.beta.into()); + handles.register_float_tensor::(&desc.out.id, result); + } + ActivationOperationIr::LogSigmoid(desc) => { + let input = handles.get_float_tensor::(&desc.input); + let result = B::log_sigmoid(input); + handles.register_float_tensor::(&desc.out.id, result); + } + ActivationOperationIr::LogSigmoidBackward(desc) => { + let x = handles.get_float_tensor::(&desc.lhs); + let grad = handles.get_float_tensor::(&desc.rhs); + let result = B::log_sigmoid_backward(x, grad); + handles.register_float_tensor::(&desc.out.id, result); + } + ActivationOperationIr::Softmax(desc) => { + let input = handles.get_float_tensor::(&desc.input); + let result = B::softmax(input, desc.axis); + handles.register_float_tensor::(&desc.out.id, result); + } + ActivationOperationIr::LogSoftmax(desc) => { + let input = handles.get_float_tensor::(&desc.input); + let result = B::log_softmax(input, desc.axis); + handles.register_float_tensor::(&desc.out.id, result); + } + ActivationOperationIr::Softmin(desc) => { + let input = handles.get_float_tensor::(&desc.input); + let result = B::softmin(input, desc.axis); + handles.register_float_tensor::(&desc.out.id, result); + } + }, + OperationIr::Custom(desc) => match self.custom_ops.get(&desc.id) { + Some(handler) => handler(handles, desc, &self.device), + None => panic!( + "No custom-op handler registered for `{}`. Register one on the server via \ + `CustomOpRegistry`/the server builder before starting it.", + desc.id + ), + }, + OperationIr::Init(_) => { + // Nothing to do. + } + OperationIr::Drop(repr) => { + handles.remove_handle(repr.id); + } + OperationIr::Distributed(op) => match op { + burn_ir::DistributedOperationIr::AllReduce(desc) => { + let tensor = handles.get_float_tensor::(&desc.tensor); + let device_ids = desc.device_ids.iter().map(|id| (*id).into()).collect(); + + let output = >::all_reduce(tensor, desc.op, device_ids); + // Safety: the collective tensor is resolved through the normal op stream + // (a `SyncCollective` op follows), so the handle is valid once that runs. + let output = unsafe { output.assume_resolved() }; + B::flush(&self.device); + handles.register_float_tensor::(&desc.out.id, output); + } + burn_ir::DistributedOperationIr::SyncCollective => B::sync_collective(&self.device), + }, + } + } + + /// Read a tensor's data, returning a future for the host readback. + pub fn read_tensor_async( + &mut self, + tensor: TensorIr, + ) -> DynFut> { + let ctx = &mut self.context; + + enum Output { + Float(B::FloatTensorPrimitive), + Int(B::IntTensorPrimitive), + Bool(B::BoolTensorPrimitive), + } + + let tensor = if tensor.dtype.is_float() { + let tensor = ctx.handles.get_float_tensor::(&tensor); + Output::::Float(tensor) + } else if tensor.dtype.is_int() { + let tensor = ctx.handles.get_int_tensor::(&tensor); + Output::Int(tensor) + } else if tensor.dtype.is_bool() { + let tensor = ctx.handles.get_bool_tensor::(&tensor); + Output::Bool(tensor) + } else if let DType::QFloat(_) = tensor.dtype { + todo!() + } else { + unimplemented!() + }; + + match tensor { + Output::Float(val) => Box::pin(B::float_into_data(val)), + Output::Int(val) => Box::pin(B::int_into_data(val)), + Output::Bool(val) => Box::pin(B::bool_into_data(val)), + } + } + + /// The device this interpreter executes on. + pub fn device(&self) -> B::Device { + self.device.clone() + } + + /// Block until all queued backend work has completed. + pub fn sync(&self) -> Result<(), ExecutionError> { + B::sync(&self.device) + } + + /// Seed the backend's RNG. + pub fn seed(&self, seed: u64) { + B::seed(&self.device, seed) + } + + /// The set of supported usages for `dtype` on this backend. + pub fn dtype_usage(&self, dtype: DType) -> burn_backend::DTypeUsageSet { + B::dtype_usage(&self.device, dtype) + } +} diff --git a/crates/burn-router/src/lib.rs b/crates/burn-router/src/lib.rs new file mode 100644 index 0000000..13a4d3d --- /dev/null +++ b/crates/burn-router/src/lib.rs @@ -0,0 +1,31 @@ +#![cfg_attr(not(feature = "std"), no_std)] +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] +#![recursion_limit = "138"] + +//! Burn multi-backend router. + +mod backend; +mod bridge; +mod channel; +mod client; +mod custom_op; +#[cfg(feature = "fusion")] +mod fusion; +mod graph; +mod interpreter; +mod ops; +mod tensor; + +pub use backend::*; +pub use bridge::*; +pub use channel::*; +pub use client::*; +pub use custom_op::*; +#[cfg(feature = "fusion")] +pub use fusion::*; +pub use graph::*; +pub use interpreter::*; +pub use tensor::*; + +extern crate alloc; diff --git a/crates/burn-router/src/ops/activation.rs b/crates/burn-router/src/ops/activation.rs new file mode 100644 index 0000000..bcfa04f --- /dev/null +++ b/crates/burn-router/src/ops/activation.rs @@ -0,0 +1,166 @@ +use crate::{BackendRouter, RouterChannel, RouterClient}; +use burn_backend::{Scalar, ops::ActivationOps, tensor::FloatTensor}; +use burn_ir::{ + ActivationOperationIr, BinaryOpIr, DimOpIr, HardSigmoidOpIr, OperationIr, OperationOutput, + ScalarOpIr, UnaryOpIr, +}; + +impl ActivationOps for BackendRouter { + fn relu(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Activation(ActivationOperationIr::Relu(desc))) + .output() + } + + fn relu_backward(output: FloatTensor, grad: FloatTensor) -> FloatTensor { + let client = output.client.clone(); + let desc = BinaryOpIr::create(output.into_ir(), grad.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Activation( + ActivationOperationIr::ReluBackward(desc), + )) + .output() + } + + fn leaky_relu(tensor: FloatTensor, negative_slope: Scalar) -> FloatTensor { + let client = tensor.client.clone(); + let desc = ScalarOpIr::create(tensor.into_ir(), negative_slope.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Activation(ActivationOperationIr::LeakyRelu( + desc, + ))) + .output() + } + + fn prelu(tensor: FloatTensor, alpha: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = BinaryOpIr::create(tensor.into_ir(), alpha.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Activation(ActivationOperationIr::PRelu(desc))) + .output() + } + + fn gelu(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Activation(ActivationOperationIr::Gelu(desc))) + .output() + } + + fn gelu_backward(x: FloatTensor, grad: FloatTensor) -> FloatTensor { + let client = x.client.clone(); + let desc = BinaryOpIr::create(x.into_ir(), grad.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Activation( + ActivationOperationIr::GeluBackward(desc), + )) + .output() + } + + fn sigmoid(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Activation(ActivationOperationIr::Sigmoid( + desc, + ))) + .output() + } + + fn sigmoid_backward(output: FloatTensor, grad: FloatTensor) -> FloatTensor { + let client = output.client.clone(); + let desc = BinaryOpIr::create(output.into_ir(), grad.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Activation( + ActivationOperationIr::SigmoidBackward(desc), + )) + .output() + } + + fn hard_sigmoid(tensor: FloatTensor, alpha: Scalar, beta: Scalar) -> FloatTensor { + let client = tensor.client.clone(); + let desc = HardSigmoidOpIr::create(tensor.into_ir(), alpha.into(), beta.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Activation(ActivationOperationIr::HardSigmoid( + desc, + ))) + .output() + } + + fn log_sigmoid(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Activation(ActivationOperationIr::LogSigmoid( + desc, + ))) + .output() + } + + fn log_sigmoid_backward(x: FloatTensor, grad: FloatTensor) -> FloatTensor { + let client = x.client.clone(); + let desc = BinaryOpIr::create(x.into_ir(), grad.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Activation( + ActivationOperationIr::LogSigmoidBackward(desc), + )) + .output() + } + + fn softmax(tensor: FloatTensor, dim: usize) -> FloatTensor { + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register(OperationIr::Activation(ActivationOperationIr::Softmax( + desc, + ))) + .output() + } + + fn log_softmax(tensor: FloatTensor, dim: usize) -> FloatTensor { + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register(OperationIr::Activation(ActivationOperationIr::LogSoftmax( + desc, + ))) + .output() + } + + fn softmin(tensor: FloatTensor, dim: usize) -> FloatTensor { + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register(OperationIr::Activation(ActivationOperationIr::Softmin( + desc, + ))) + .output() + } +} diff --git a/crates/burn-router/src/ops/binary.rs b/crates/burn-router/src/ops/binary.rs new file mode 100644 index 0000000..b00e36a --- /dev/null +++ b/crates/burn-router/src/ops/binary.rs @@ -0,0 +1,69 @@ +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! binary_float_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let lhs = $handles.get_float_tensor::(&$desc.lhs); + let rhs = $handles.get_float_tensor::(&$desc.rhs); + let output = $ops(lhs, rhs); + + $handles.register_float_tensor::(&$desc.out.id, output); + }}; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! binary_float_cmp_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let lhs = $handles.get_float_tensor::(&$desc.lhs); + let rhs = $handles.get_float_tensor::(&$desc.rhs); + let output = $ops(lhs, rhs, $desc.out.dtype.into()); + + $handles.register_bool_tensor::(&$desc.out.id, output); + }}; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! binary_int_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let lhs = $handles.get_int_tensor::(&$desc.lhs); + let rhs = $handles.get_int_tensor::(&$desc.rhs); + let output = $ops(lhs, rhs); + + $handles.register_int_tensor::(&$desc.out.id, output); + }}; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! binary_int_cmp_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let lhs = $handles.get_int_tensor::(&$desc.lhs); + let rhs = $handles.get_int_tensor::(&$desc.rhs); + let output = $ops(lhs, rhs, $desc.out.dtype.into()); + + $handles.register_bool_tensor::(&$desc.out.id, output); + }}; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! binary_bool_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let lhs = $handles.get_bool_tensor::(&$desc.lhs); + let rhs = $handles.get_bool_tensor::(&$desc.rhs); + let output = $ops(lhs, rhs); + + $handles.register_bool_tensor::(&$desc.out.id, output); + }}; +} diff --git a/crates/burn-router/src/ops/bool_tensor.rs b/crates/burn-router/src/ops/bool_tensor.rs new file mode 100644 index 0000000..92139d7 --- /dev/null +++ b/crates/burn-router/src/ops/bool_tensor.rs @@ -0,0 +1,433 @@ +use alloc::vec::Vec; +use burn_backend::backend::ExecutionError; +use burn_std::{BoolDType, FloatDType, IntDType}; + +use crate::{BackendRouter, RouterChannel, RouterClient, get_client}; +use burn_backend::ops::BoolTensorOps; +use burn_backend::tensor::{BoolTensor, Device, FloatTensor, IndexingUpdateOp, IntTensor}; +use burn_backend::{Scalar, Shape, Slice, TensorData}; +use burn_ir::{ + BaseOperationIr, BinaryOpIr, BoolOperationIr, CastOpIr, CatOpIr, CreationOpIr, FlipOpIr, + GatherOpIr, InitOperationIr, MaskFillOpIr, MaskWhereOpIr, OperationIr, OperationOutput, + PermuteOpIr, ReduceDimOpIr, ReduceOpIr, RepeatDimOpIr, ScalarOpIr, ScatterOpIr, + SelectAssignOpIr, SelectOpIr, ShapeOpIr, SliceAssignOpIr, SliceOpIr, SwapDimsOpIr, UnaryOpIr, + UnfoldOpIr, +}; + +impl BoolTensorOps for BackendRouter { + fn bool_empty(shape: Shape, device: &Device, dtype: BoolDType) -> BoolTensor { + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register(OperationIr::BaseBool(BaseOperationIr::Empty(desc))) + .output() + } + + fn bool_zeros(shape: Shape, device: &Device, dtype: BoolDType) -> BoolTensor { + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register(OperationIr::BaseBool(BaseOperationIr::Zeros(desc))) + .output() + } + + fn bool_ones(shape: Shape, device: &Device, dtype: BoolDType) -> BoolTensor { + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register(OperationIr::BaseBool(BaseOperationIr::Ones(desc))) + .output() + } + + async fn bool_into_data(tensor: BoolTensor) -> Result { + tensor.into_data().await + } + + fn bool_from_data(data: TensorData, device: &Device) -> BoolTensor { + let client = get_client::(device); + let out = client.register_tensor_data(data); + let desc = InitOperationIr { + out: out.to_ir_out(), + }; + + // Call register op when output is already initialized + client.register_op(OperationIr::Init(desc)); + + out + } + + fn bool_into_int(tensor: BoolTensor, out_dtype: IntDType) -> IntTensor { + let client = tensor.client.clone(); + let desc = CastOpIr::create(tensor.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Bool(BoolOperationIr::IntoInt(desc))) + .output() + } + + fn bool_into_float(tensor: BoolTensor, out_dtype: FloatDType) -> FloatTensor { + let client = tensor.client.clone(); + let desc = CastOpIr::create(tensor.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Bool(BoolOperationIr::IntoFloat(desc))) + .output() + } + + fn bool_to_device(tensor: BoolTensor, device: &Device) -> BoolTensor { + if &tensor.client.device() == device { + return tensor; + } + R::change_client_backend(tensor, device) + } + + fn bool_reshape(tensor: BoolTensor, shape: Shape) -> BoolTensor { + let client = tensor.client.clone(); + let desc = ShapeOpIr::reshape(tensor.into_ir(), shape, || client.create_empty_handle()); + + client + .register(OperationIr::BaseBool(BaseOperationIr::Reshape(desc))) + .output() + } + + fn bool_slice(tensor: BoolTensor, slices: &[Slice]) -> BoolTensor { + let client = tensor.client.clone(); + let desc = SliceOpIr::create(tensor.into_ir(), slices.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseBool(BaseOperationIr::Slice(desc))) + .output() + } + + fn bool_slice_assign( + tensor: BoolTensor, + slices: &[burn_backend::Slice], + value: BoolTensor, + ) -> BoolTensor { + let client = tensor.client.clone(); + let desc = + SliceAssignOpIr::create(tensor.into_ir(), slices.into(), value.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseBool(BaseOperationIr::SliceAssign(desc))) + .output() + } + + fn bool_equal(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseBool(BaseOperationIr::Equal(desc))) + .output() + } + + fn bool_not(tensor: BoolTensor) -> BoolTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Bool(BoolOperationIr::Not(desc))) + .output() + } + + fn bool_and(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Bool(BoolOperationIr::And(desc))) + .output() + } + + fn bool_or(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Bool(BoolOperationIr::Or(desc))) + .output() + } + + fn bool_xor(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::Bool(BoolOperationIr::Xor(desc))) + .output() + } + + fn bool_not_equal(lhs: BoolTensor, rhs: BoolTensor) -> BoolTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::BaseBool(BaseOperationIr::NotEqual(desc))) + .output() + } + + fn bool_not_equal_elem(lhs: BoolTensor, rhs: Scalar) -> BoolTensor { + let client = lhs.client.clone(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs.into(), || client.create_empty_handle()); + client + .register(OperationIr::BaseBool(BaseOperationIr::NotEqualElem(desc))) + .output() + } + + fn bool_all(tensor: BoolTensor) -> BoolTensor { + let client = tensor.client.clone(); + let dtype = tensor.dtype; + let desc = + ReduceOpIr::create_bool(tensor.into_ir(), dtype, || client.create_empty_handle()); + client + .register(OperationIr::BaseBool(BaseOperationIr::All(desc))) + .output() + } + + fn bool_any(tensor: BoolTensor) -> BoolTensor { + let client = tensor.client.clone(); + let dtype = tensor.dtype; + let desc = + ReduceOpIr::create_bool(tensor.into_ir(), dtype, || client.create_empty_handle()); + client + .register(OperationIr::BaseBool(BaseOperationIr::Any(desc))) + .output() + } + + fn bool_all_dim(tensor: BoolTensor, dim: usize) -> BoolTensor { + let client = tensor.client.clone(); + let dtype = tensor.dtype; + let desc = ReduceDimOpIr::create_bool(tensor.into_ir(), dim, 1, dtype, || { + client.create_empty_handle() + }); + client + .register(OperationIr::BaseBool(BaseOperationIr::AllDim(desc))) + .output() + } + + fn bool_any_dim(tensor: BoolTensor, dim: usize) -> BoolTensor { + let client = tensor.client.clone(); + let dtype = tensor.dtype; + let desc = ReduceDimOpIr::create_bool(tensor.into_ir(), dim, 1, dtype, || { + client.create_empty_handle() + }); + client + .register(OperationIr::BaseBool(BaseOperationIr::AnyDim(desc))) + .output() + } + + fn bool_swap_dims(tensor: BoolTensor, dim1: usize, dim2: usize) -> BoolTensor { + let client = tensor.client.clone(); + let desc = SwapDimsOpIr::create(tensor.into_ir(), dim1, dim2, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseBool(BaseOperationIr::SwapDims(desc))) + .output() + } + + fn bool_permute(tensor: BoolTensor, axes: &[usize]) -> BoolTensor { + let client = tensor.client.clone(); + let desc = PermuteOpIr::create(tensor.into_ir(), axes.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseBool(BaseOperationIr::Permute(desc))) + .output() + } + + fn bool_flip(tensor: BoolTensor, axes: &[usize]) -> BoolTensor { + let client = tensor.client.clone(); + let desc = FlipOpIr::create(tensor.into_ir(), axes.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseBool(BaseOperationIr::Flip(desc))) + .output() + } + + fn bool_expand(tensor: BoolTensor, shape: Shape) -> BoolTensor { + let client = tensor.client.clone(); + let desc = ShapeOpIr::expand(tensor.into_ir(), shape, || client.create_empty_handle()); + + client + .register(OperationIr::BaseBool(BaseOperationIr::Expand(desc))) + .output() + } + + fn bool_cat(tensors: Vec>, dim: usize) -> BoolTensor { + let client = tensors.first().unwrap().client.clone(); + let tensors = tensors.into_iter().map(|t| t.into_ir()).collect(); + let desc = CatOpIr::create(tensors, dim, || client.create_empty_handle()); + + client + .register(OperationIr::BaseBool(BaseOperationIr::Cat(desc))) + .output() + } + + fn bool_repeat_dim(tensor: BoolTensor, dim: usize, times: usize) -> BoolTensor { + let client = tensor.client.clone(); + let desc = RepeatDimOpIr::create(tensor.into_ir(), dim, times, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseBool(BaseOperationIr::RepeatDim(desc))) + .output() + } + + fn bool_unfold( + tensor: BoolTensor, + dim: usize, + size: usize, + step: usize, + ) -> BoolTensor { + let client = tensor.client.clone(); + let desc = UnfoldOpIr::create(tensor.into_ir(), dim, size, step, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseBool(BaseOperationIr::Unfold(desc))) + .output() + } + + fn bool_mask_where( + tensor: BoolTensor, + mask: BoolTensor, + value: BoolTensor, + ) -> BoolTensor { + let client = tensor.client.clone(); + let desc = MaskWhereOpIr::create(tensor.into_ir(), mask.into_ir(), value.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseBool(BaseOperationIr::MaskWhere(desc))) + .output() + } + + fn bool_mask_fill( + tensor: BoolTensor, + mask: BoolTensor, + value: Scalar, + ) -> BoolTensor { + let client = tensor.client.clone(); + let value = value.into(); + let desc = MaskFillOpIr::create(tensor.into_ir(), mask.into_ir(), value, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseBool(BaseOperationIr::MaskFill(desc))) + .output() + } + + fn bool_gather( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + ) -> BoolTensor { + let client = tensor.client.clone(); + let desc = GatherOpIr::create(tensor.into_ir(), dim, indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseBool(BaseOperationIr::Gather(desc))) + .output() + } + + fn bool_scatter_or( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor { + let client = tensor.client.clone(); + let desc = ScatterOpIr::create( + tensor.into_ir(), + dim, + indices.into_ir(), + value.into_ir(), + IndexingUpdateOp::Add, + || client.create_empty_handle(), + ); + + client + .register(OperationIr::BaseBool(BaseOperationIr::Scatter(desc))) + .output() + } + + fn bool_equal_elem(lhs: BoolTensor, rhs: Scalar) -> BoolTensor { + let dtype = lhs.dtype; + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, dtype, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseBool(BaseOperationIr::EqualElem(desc))) + .output() + } + + fn bool_select( + tensor: BoolTensor, + dim: usize, + indices: IntTensor, + ) -> BoolTensor { + let client = tensor.client.clone(); + let desc = SelectOpIr::create(tensor.into_ir(), dim, indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseBool(BaseOperationIr::Select(desc))) + .output() + } + + fn bool_select_or( + tensor: BoolTensor, + dim: usize, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor { + let client = tensor.client.clone(); + let desc = SelectAssignOpIr::create( + tensor.into_ir(), + dim, + indices.into_ir(), + value.into_ir(), + IndexingUpdateOp::Add, + || client.create_empty_handle(), + ); + + client + .register(OperationIr::BaseBool(BaseOperationIr::SelectAssign(desc))) + .output() + } +} diff --git a/crates/burn-router/src/ops/distributed.rs b/crates/burn-router/src/ops/distributed.rs new file mode 100644 index 0000000..19a1316 --- /dev/null +++ b/crates/burn-router/src/ops/distributed.rs @@ -0,0 +1,42 @@ +use alloc::vec::Vec; + +use burn_backend::{ + DeviceId, + distributed::{CollectiveTensor, DistributedOps, ReduceOperation}, + tensor::{Device, FloatTensor}, +}; +use burn_ir::{AllReduceOpIr, DeviceIdIr, DistributedOperationIr, OperationIr, OperationOutput}; + +use crate::{BackendRouter, RouterChannel, RouterClient, get_client}; + +impl DistributedOps for BackendRouter { + fn all_reduce( + tensor: FloatTensor, + op: ReduceOperation, + device_ids: Vec, + ) -> CollectiveTensor { + let client = tensor.client.clone(); + let device_ids = device_ids.into_iter().map(DeviceIdIr::from).collect(); + let desc = AllReduceOpIr::create(tensor.into_ir(), op, device_ids, || { + client.create_empty_handle() + }); + + let output = client + .register(OperationIr::Distributed(DistributedOperationIr::AllReduce( + desc, + ))) + .output(); + + CollectiveTensor::new(output) + } + + fn sync_collective(device: &Device) { + // Fire-and-forget, like `all_reduce`: register a `SyncCollective` op on the device's + // stream instead of a blocking call. The interpreter resolves it through the normal op + // stream like any other distributed op. + let client = get_client::(device); + client.register_op(OperationIr::Distributed( + DistributedOperationIr::SyncCollective, + )); + } +} diff --git a/crates/burn-router/src/ops/int_tensor.rs b/crates/burn-router/src/ops/int_tensor.rs new file mode 100644 index 0000000..1ca348d --- /dev/null +++ b/crates/burn-router/src/ops/int_tensor.rs @@ -0,0 +1,1312 @@ +use alloc::vec::Vec; +use burn_backend::backend::ExecutionError; +use burn_std::{BoolDType, FloatDType}; + +use crate::{BackendRouter, RouterChannel, RouterClient, get_client}; +use burn_backend::tensor::{BoolTensor, Device, FloatTensor, IndexingUpdateOp, IntTensor}; +use burn_backend::{Distribution, IntDType, Scalar, Shape, Slice, TensorData, ops::IntTensorOps}; +use burn_ir::{ + BaseOperationIr, BinaryOpIr, CastOpIr, CatOpIr, ClampOpIr, CreationOpIr, DimOpIr, FlipOpIr, + FullOpIr, GatherNdOpIr, GatherOpIr, InitOperationIr, IntOperationIr, MaskFillOpIr, + MaskWhereOpIr, MatmulOpIr, NumericOperationIr, OperationIr, OperationOutput, PermuteOpIr, + RandomOpIr, ReduceDimOpIr, ReduceDimWithIndicesOpIr, ReduceOpIr, RepeatDimOpIr, ScalarOpIr, + ScatterNdOpIr, ScatterOpIr, SelectAssignOpIr, SelectOpIr, ShapeOpIr, SliceAssignOpIr, + SliceOpIr, SortOpIr, SortWithIndicesOpIr, SwapDimsOpIr, UnaryOpIr, UnfoldOpIr, +}; + +impl IntTensorOps for BackendRouter { + fn int_empty(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register(OperationIr::BaseInt(BaseOperationIr::Empty(desc))) + .output() + } + + async fn int_into_data(tensor: IntTensor) -> Result { + tensor.into_data().await + } + + fn int_from_data(data: TensorData, device: &Device) -> IntTensor { + let client = get_client::(device); + let out = client.register_tensor_data(data); + let desc = InitOperationIr { + out: out.to_ir_out(), + }; + + // Call register op when output is already initialized + client.register_op(OperationIr::Init(desc)); + + out + } + + fn int_to_device(tensor: IntTensor, device: &Device) -> IntTensor { + if &tensor.client.device() == device { + return tensor; + } + R::change_client_backend(tensor, device) + } + + fn int_reshape(tensor: IntTensor, shape: Shape) -> IntTensor { + let client = tensor.client.clone(); + let desc = ShapeOpIr::reshape(tensor.into_ir(), shape, || client.create_empty_handle()); + + client + .register(OperationIr::BaseInt(BaseOperationIr::Reshape(desc))) + .output() + } + + fn int_slice(tensor: IntTensor, slices: &[Slice]) -> IntTensor { + let client = tensor.client.clone(); + let desc = SliceOpIr::create(tensor.into_ir(), slices.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseInt(BaseOperationIr::Slice(desc))) + .output() + } + + fn int_slice_assign( + tensor: IntTensor, + slices: &[burn_backend::Slice], + value: IntTensor, + ) -> IntTensor { + let client = tensor.client.clone(); + let desc = + SliceAssignOpIr::create(tensor.into_ir(), slices.into(), value.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseInt(BaseOperationIr::SliceAssign(desc))) + .output() + } + + fn int_matmul(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + let client = lhs.client.clone(); + let desc = MatmulOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Int(IntOperationIr::Matmul(desc))) + .output() + } + + fn int_mask_where( + tensor: IntTensor, + mask: BoolTensor, + value: IntTensor, + ) -> IntTensor { + let client = tensor.client.clone(); + let desc = MaskWhereOpIr::create(tensor.into_ir(), mask.into_ir(), value.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseInt(BaseOperationIr::MaskWhere(desc))) + .output() + } + + fn int_mask_fill( + tensor: IntTensor, + mask: BoolTensor, + value: Scalar, + ) -> IntTensor { + let client = tensor.client.clone(); + let value = value.into(); + let desc = MaskFillOpIr::create(tensor.into_ir(), mask.into_ir(), value, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseInt(BaseOperationIr::MaskFill(desc))) + .output() + } + + fn int_gather( + dim: usize, + tensor: IntTensor, + indices: IntTensor, + ) -> IntTensor { + let client = tensor.client.clone(); + let desc = GatherOpIr::create(tensor.into_ir(), dim, indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseInt(BaseOperationIr::Gather(desc))) + .output() + } + + fn int_scatter_add( + dim: usize, + tensor: IntTensor, + indices: IntTensor, + value: IntTensor, + ) -> IntTensor { + let client = tensor.client.clone(); + let desc = ScatterOpIr::create( + tensor.into_ir(), + dim, + indices.into_ir(), + value.into_ir(), + IndexingUpdateOp::Add, + || client.create_empty_handle(), + ); + + client + .register(OperationIr::BaseInt(BaseOperationIr::Scatter(desc))) + .output() + } + + fn int_scatter_nd( + data: IntTensor, + indices: IntTensor, + values: IntTensor, + reduction: IndexingUpdateOp, + ) -> IntTensor { + let client = data.client.clone(); + let desc = ScatterNdOpIr::create( + data.into_ir(), + indices.into_ir(), + values.into_ir(), + reduction, + || client.create_empty_handle(), + ); + + client + .register(OperationIr::BaseInt(BaseOperationIr::ScatterNd(desc))) + .output() + } + + fn int_gather_nd(data: IntTensor, indices: IntTensor) -> IntTensor { + let client = data.client.clone(); + let desc = GatherNdOpIr::create(data.into_ir(), indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseInt(BaseOperationIr::GatherNd(desc))) + .output() + } + + fn int_select( + tensor: IntTensor, + dim: usize, + indices: IntTensor, + ) -> IntTensor { + let client = tensor.client.clone(); + let desc = SelectOpIr::create(tensor.into_ir(), dim, indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseInt(BaseOperationIr::Select(desc))) + .output() + } + + fn int_select_add( + tensor: IntTensor, + dim: usize, + indices: IntTensor, + value: IntTensor, + ) -> IntTensor { + let client = tensor.client.clone(); + let desc = SelectAssignOpIr::create( + tensor.into_ir(), + dim, + indices.into_ir(), + value.into_ir(), + IndexingUpdateOp::Add, + || client.create_empty_handle(), + ); + + client + .register(OperationIr::BaseInt(BaseOperationIr::SelectAssign(desc))) + .output() + } + + fn int_cat(tensors: Vec>, dim: usize) -> IntTensor { + let client = tensors.first().unwrap().client.clone(); + let tensors = tensors.into_iter().map(|t| t.into_ir()).collect(); + let desc = CatOpIr::create(tensors, dim, || client.create_empty_handle()); + + client + .register(OperationIr::BaseInt(BaseOperationIr::Cat(desc))) + .output() + } + + fn int_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseInt(BaseOperationIr::Equal(desc))) + .output() + } + + fn int_equal_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseInt(BaseOperationIr::EqualElem(desc))) + .output() + } + + fn int_greater( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericInt( + desc.lhs.dtype, + NumericOperationIr::Greater(desc), + )) + .output() + } + + fn int_greater_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericInt( + desc.lhs.dtype, + NumericOperationIr::GreaterElem(desc), + )) + .output() + } + + fn int_greater_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericInt( + desc.lhs.dtype, + NumericOperationIr::GreaterEqual(desc), + )) + .output() + } + + fn int_greater_equal_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericInt( + desc.lhs.dtype, + NumericOperationIr::GreaterEqualElem(desc), + )) + .output() + } + + fn int_lower( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericInt( + desc.lhs.dtype, + NumericOperationIr::Lower(desc), + )) + .output() + } + + fn int_lower_elem(lhs: IntTensor, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericInt( + desc.lhs.dtype, + NumericOperationIr::LowerElem(desc), + )) + .output() + } + + fn int_lower_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericInt( + desc.lhs.dtype, + NumericOperationIr::LowerEqual(desc), + )) + .output() + } + + fn int_lower_equal_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericInt( + desc.lhs.dtype, + NumericOperationIr::LowerEqualElem(desc), + )) + .output() + } + + fn int_add(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Add(desc), + )) + .output() + } + + fn int_add_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::AddScalar(desc), + )) + .output() + } + + fn int_sub(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Sub(desc), + )) + .output() + } + + fn int_sub_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::SubScalar(desc), + )) + .output() + } + + fn int_mul(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Mul(desc), + )) + .output() + } + + fn int_mul_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::MulScalar(desc), + )) + .output() + } + + fn int_div(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Div(desc), + )) + .output() + } + + fn int_div_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::DivScalar(desc), + )) + .output() + } + + fn int_remainder(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Rem(desc), + )) + .output() + } + + fn int_remainder_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::RemScalar(desc), + )) + .output() + } + + fn int_zeros(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register(OperationIr::BaseInt(BaseOperationIr::Zeros(desc))) + .output() + } + + fn int_ones(shape: Shape, device: &Device, dtype: IntDType) -> IntTensor { + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register(OperationIr::BaseInt(BaseOperationIr::Ones(desc))) + .output() + } + + fn int_full( + shape: Shape, + fill_value: Scalar, + device: &Device, + dtype: IntDType, + ) -> IntTensor { + let client = get_client::(device); + let dtype = dtype.into(); + let value = fill_value.into(); + let desc = FullOpIr::create(shape, dtype, value, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Full(desc), + )) + .output() + } + + fn int_sum(tensor: IntTensor) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Sum(desc), + )) + .output() + } + + fn int_sum_dim(tensor: IntTensor, axis: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = + ReduceDimOpIr::create(tensor.into_ir(), axis, 1, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::SumDim(desc), + )) + .output() + } + + fn int_prod(tensor: IntTensor) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Prod(desc), + )) + .output() + } + + fn int_prod_dim(tensor: IntTensor, dim: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::ProdDim(desc), + )) + .output() + } + + fn int_mean(tensor: IntTensor) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Mean(desc), + )) + .output() + } + + fn int_mean_dim(tensor: IntTensor, dim: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::MeanDim(desc), + )) + .output() + } + + fn int_cumsum(tensor: IntTensor, dim: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::CumSum(desc), + )) + .output() + } + + fn int_cumprod(tensor: IntTensor, dim: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::CumProd(desc), + )) + .output() + } + + fn int_cummin(tensor: IntTensor, dim: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::CumMin(desc), + )) + .output() + } + + fn int_cummax(tensor: IntTensor, dim: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::CumMax(desc), + )) + .output() + } + + fn int_argmax(tensor: IntTensor, dim: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::ArgMax(desc), + )) + .output() + } + + fn int_argtopk(tensor: IntTensor, dim: usize, k: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, k, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::ArgTopK(desc), + )) + .output() + } + + fn int_topk(tensor: IntTensor, dim: usize, k: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, k, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::TopK(desc), + )) + .output() + } + + fn int_argmin(tensor: IntTensor, dim: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::ArgMin(desc), + )) + .output() + } + + fn int_clamp(tensor: IntTensor, min: Scalar, max: Scalar) -> IntTensor { + let client = tensor.client.clone(); + let min = min.into(); + let max = max.into(); + let desc = ClampOpIr::create(tensor.into_ir(), min, max, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Clamp(desc), + )) + .output() + } + + fn int_abs(tensor: IntTensor) -> IntTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Abs(desc), + )) + .output() + } + + fn int_into_float(tensor: IntTensor, out_dtype: FloatDType) -> FloatTensor { + let client = tensor.client.clone(); + let desc = CastOpIr::create(tensor.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Int(IntOperationIr::IntoFloat(desc))) + .output() + } + + fn int_swap_dims(tensor: IntTensor, dim1: usize, dim2: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = SwapDimsOpIr::create(tensor.into_ir(), dim1, dim2, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseInt(BaseOperationIr::SwapDims(desc))) + .output() + } + + fn int_max(tensor: IntTensor) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Max(desc), + )) + .output() + } + + fn int_max_dim(tensor: IntTensor, dim: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::MaxDim(desc), + )) + .output() + } + + fn int_max_dim_with_indices( + tensor: IntTensor, + dim: usize, + ) -> (IntTensor, IntTensor) { + let dtype = tensor.dtype; + let client = tensor.client.clone(); + let desc = ReduceDimWithIndicesOpIr::create(tensor.into_ir(), dim, dtype, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericInt( + desc.tensor.dtype, + NumericOperationIr::MaxDimWithIndices(desc), + )) + .outputs() + .into() + } + + fn int_max_abs(tensor: IntTensor) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::MaxAbs(desc), + )) + .output() + } + + fn int_max_abs_dim(tensor: IntTensor, dim: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::MaxAbsDim(desc), + )) + .output() + } + + fn int_min(tensor: IntTensor) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Min(desc), + )) + .output() + } + + fn int_min_dim(tensor: IntTensor, dim: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::MinDim(desc), + )) + .output() + } + + fn int_min_dim_with_indices( + tensor: IntTensor, + dim: usize, + ) -> (IntTensor, IntTensor) { + let dtype = tensor.dtype; + let client = tensor.client.clone(); + let desc = ReduceDimWithIndicesOpIr::create(tensor.into_ir(), dim, dtype, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::MinDimWithIndices(desc), + )) + .outputs() + .into() + } + + fn int_random( + shape: Shape, + distribution: Distribution, + device: &Device, + dtype: IntDType, + ) -> IntTensor { + let client = get_client::(device); + let dtype = dtype.into(); + let desc = RandomOpIr::create(shape, dtype, distribution, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + dtype, + NumericOperationIr::IntRandom(desc), + )) + .output() + } + + fn int_permute(tensor: IntTensor, axes: &[usize]) -> IntTensor { + let client = tensor.client.clone(); + let desc = PermuteOpIr::create(tensor.into_ir(), axes.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseInt(BaseOperationIr::Permute(desc))) + .output() + } + + fn int_expand(tensor: IntTensor, shape: Shape) -> IntTensor { + let client = tensor.client.clone(); + let desc = ShapeOpIr::expand(tensor.into_ir(), shape, || client.create_empty_handle()); + + client + .register(OperationIr::BaseInt(BaseOperationIr::Expand(desc))) + .output() + } + + fn int_flip(tensor: IntTensor, axes: &[usize]) -> IntTensor { + let client = tensor.client.clone(); + let desc = FlipOpIr::create(tensor.into_ir(), axes.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseInt(BaseOperationIr::Flip(desc))) + .output() + } + + fn int_repeat_dim(tensor: IntTensor, dim: usize, times: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = RepeatDimOpIr::create(tensor.into_ir(), dim, times, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseInt(BaseOperationIr::RepeatDim(desc))) + .output() + } + + fn bitwise_and(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Int(IntOperationIr::BitwiseAnd(desc))) + .output() + } + + fn bitwise_or(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Int(IntOperationIr::BitwiseOr(desc))) + .output() + } + + fn bitwise_xor(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Int(IntOperationIr::BitwiseXor(desc))) + .output() + } + + fn bitwise_not(tensor: IntTensor) -> IntTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Int(IntOperationIr::BitwiseNot(desc))) + .output() + } + + fn bitwise_and_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::Int(IntOperationIr::BitwiseAndScalar(desc))) + .output() + } + + fn bitwise_or_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::Int(IntOperationIr::BitwiseOrScalar(desc))) + .output() + } + + fn bitwise_xor_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::Int(IntOperationIr::BitwiseXorScalar(desc))) + .output() + } + + fn bitwise_left_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Int(IntOperationIr::BitwiseLeftShift(desc))) + .output() + } + + fn bitwise_left_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::Int(IntOperationIr::BitwiseLeftShiftScalar( + desc, + ))) + .output() + } + + fn bitwise_right_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Int(IntOperationIr::BitwiseRightShift(desc))) + .output() + } + + fn bitwise_right_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::Int(IntOperationIr::BitwiseRightShiftScalar( + desc, + ))) + .output() + } + + fn int_cast(tensor: IntTensor, dtype: burn_backend::IntDType) -> IntTensor { + let client = tensor.client.clone(); + let desc = CastOpIr::create(tensor.into_ir(), dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseInt(BaseOperationIr::Cast(desc))) + .output() + } + + fn int_unfold( + tensor: IntTensor, + dim: usize, + size: usize, + step: usize, + ) -> IntTensor { + let client = tensor.client.clone(); + let desc = UnfoldOpIr::create(tensor.into_ir(), dim, size, step, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseInt(BaseOperationIr::Unfold(desc))) + .output() + } + + fn int_powi(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Powi(desc), + )) + .output() + } + + fn int_neg(tensor: IntTensor) -> IntTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Neg(desc), + )) + .output() + } + + fn int_sign(tensor: IntTensor) -> IntTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Sign(desc), + )) + .output() + } + + fn int_clamp_min(tensor: IntTensor, min: Scalar) -> IntTensor { + let client = tensor.client.clone(); + let desc = ScalarOpIr::create(tensor.into_ir(), min.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::ClampMin(desc), + )) + .output() + } + + fn int_clamp_max(tensor: IntTensor, max: Scalar) -> IntTensor { + let client = tensor.client.clone(); + let desc = ScalarOpIr::create(tensor.into_ir(), max.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::ClampMax(desc), + )) + .output() + } + + fn int_not_equal( + lhs: IntTensor, + rhs: IntTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::BaseInt(BaseOperationIr::NotEqual(desc))) + .output() + } + + fn int_not_equal_elem( + lhs: IntTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let desc = + ScalarOpIr::create_comparison(lhs.into_ir(), rhs.into(), out_dtype.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::BaseInt(BaseOperationIr::NotEqualElem(desc))) + .output() + } + + fn int_all(tensor: IntTensor, out_dtype: BoolDType) -> BoolTensor { + let client = tensor.client.clone(); + let desc = ReduceOpIr::create_bool(tensor.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::BaseInt(BaseOperationIr::All(desc))) + .output() + } + + fn int_any(tensor: IntTensor, out_dtype: BoolDType) -> BoolTensor { + let client = tensor.client.clone(); + let desc = ReduceOpIr::create_bool(tensor.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::BaseInt(BaseOperationIr::Any(desc))) + .output() + } + + fn int_all_dim(tensor: IntTensor, dim: usize, out_dtype: BoolDType) -> BoolTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create_bool(tensor.into_ir(), dim, 1, out_dtype.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::BaseInt(BaseOperationIr::AllDim(desc))) + .output() + } + + fn int_any_dim(tensor: IntTensor, dim: usize, out_dtype: BoolDType) -> BoolTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create_bool(tensor.into_ir(), dim, 1, out_dtype.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::BaseInt(BaseOperationIr::AnyDim(desc))) + .output() + } + + fn int_sort(tensor: IntTensor, dim: usize, descending: bool) -> IntTensor { + let client = tensor.client.clone(); + let desc = SortOpIr::create(tensor.into_ir(), dim, descending, || { + client.create_empty_handle() + }); + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::Sort(desc), + )) + .output() + } + + fn int_sort_with_indices( + tensor: IntTensor, + dim: usize, + descending: bool, + ) -> (IntTensor, IntTensor) { + let client = tensor.client.clone(); + let dtype = tensor.dtype; + let desc = SortWithIndicesOpIr::create(tensor.into_ir(), dim, descending, dtype, || { + client.create_empty_handle() + }); + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::SortWithIndices(desc), + )) + .outputs() + .into() + } + + fn int_argsort(tensor: IntTensor, dim: usize, descending: bool) -> IntTensor { + let client = tensor.client.clone(); + let dtype = tensor.dtype; + let desc = SortOpIr::create_arg(tensor.into_ir(), dim, descending, dtype, || { + client.create_empty_handle() + }); + client + .register(OperationIr::NumericInt( + dtype, + NumericOperationIr::ArgSort(desc), + )) + .output() + } + + fn int_powi_scalar_impl(lhs: IntTensor, rhs: Scalar) -> IntTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::NumericInt( + desc.out.dtype, + NumericOperationIr::PowiScalar(desc), + )) + .output() + } +} diff --git a/crates/burn-router/src/ops/mod.rs b/crates/burn-router/src/ops/mod.rs new file mode 100644 index 0000000..6d8c4be --- /dev/null +++ b/crates/burn-router/src/ops/mod.rs @@ -0,0 +1,10 @@ +mod activation; +mod binary; +mod bool_tensor; +mod distributed; +mod int_tensor; +mod module; +mod qtensor; +mod tensor; +mod transaction; +mod unary; diff --git a/crates/burn-router/src/ops/module.rs b/crates/burn-router/src/ops/module.rs new file mode 100644 index 0000000..c73d5ab --- /dev/null +++ b/crates/burn-router/src/ops/module.rs @@ -0,0 +1,1068 @@ +use alloc::boxed::Box; + +use burn_backend::ops::{ + AttentionModuleOptions, ConvOptions, ConvTransposeOptions, DeformConv2dBackward, + DeformConvOptions, InterpolateOptions, MaxPool1dBackward, MaxPool1dWithIndices, + MaxPool2dBackward, MaxPool2dWithIndices, ModuleOps, +}; +use burn_backend::tensor::{BoolTensor, FloatTensor, IntTensor}; +use burn_ir::*; +use burn_std::IntDType; + +use crate::{BackendRouter, RouterChannel, RouterClient}; + +impl ModuleOps for BackendRouter { + fn embedding(weights: FloatTensor, indices: IntTensor) -> FloatTensor { + let client = weights.client.clone(); + let desc = EmbeddingOpIr::create(weights.into_ir(), indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Module(ModuleOperationIr::Embedding(desc))) + .output() + } + + fn embedding_backward( + weights: FloatTensor, + output_grad: FloatTensor, + indices: IntTensor, + ) -> FloatTensor { + let client = weights.client.clone(); + let desc = EmbeddingBackwardOpIr::create( + weights.into_ir(), + output_grad.into_ir(), + indices.into_ir(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::EmbeddingBackward( + desc, + ))) + .output() + } + + fn linear( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = LinearOpIr::create( + x.into_ir(), + weight.into_ir(), + bias.map(|bias| bias.into_ir()), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::Linear(desc))) + .output() + } + + fn linear_x_backward( + weight: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + let client = weight.client.clone(); + let desc = LinearXBackwardOpIr::create(weight.into_ir(), output_grad.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Module(ModuleOperationIr::LinearXBackward( + desc, + ))) + .output() + } + + fn linear_weight_backward( + x: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = LinearWeightBackwardOpIr::create(x.into_ir(), output_grad.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Module( + ModuleOperationIr::LinearWeightBackward(desc), + )) + .output() + } + + fn linear_bias_backward(output_grad: FloatTensor) -> FloatTensor { + let client = output_grad.client.clone(); + let desc = + LinearBiasBackwardOpIr::create(output_grad.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Module(ModuleOperationIr::LinearBiasBackward( + desc, + ))) + .output() + } + + fn layer_norm( + tensor: FloatTensor, + gamma: FloatTensor, + beta: Option>, + epsilon: f64, + ) -> FloatTensor { + let client = tensor.client.clone(); + let desc = LayerNormOpIr::create( + tensor.into_ir(), + gamma.into_ir(), + beta.map(|b| b.into_ir()), + epsilon, + || client.create_empty_handle(), + ); + client + .register(OperationIr::Module(ModuleOperationIr::LayerNorm(desc))) + .output() + } + + fn unfold4d( + x: FloatTensor, + kernel_size: [usize; 2], + options: burn_backend::ops::UnfoldOptions, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = Unfold4dOpIr::create( + x.into_ir(), + kernel_size, + Unfold4dOptionsIr { + stride: options.stride, + padding: options.padding, + dilation: options.dilation, + }, + || client.create_empty_handle(), + ); + client + .register(OperationIr::Module(ModuleOperationIr::Unfold4d(desc))) + .output() + } + + fn conv_transpose1d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<1>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = ConvTranspose1dWeightBackwardOpIr::create( + x.into_ir(), + weight.into_ir(), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + client + .register(OperationIr::Module( + ModuleOperationIr::ConvTranspose1dWeightBackward(desc), + )) + .output() + } + + fn conv_transpose1d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = ConvTranspose1dBiasBackwardOpIr::create( + x.into_ir(), + bias.into_ir(), + output_grad.into_ir(), + || client.create_empty_handle(), + ); + client + .register(OperationIr::Module( + ModuleOperationIr::ConvTranspose1dBiasBackward(desc), + )) + .output() + } + + fn conv_transpose2d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<2>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = ConvTranspose2dWeightBackwardOpIr::create( + x.into_ir(), + weight.into_ir(), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + client + .register(OperationIr::Module( + ModuleOperationIr::ConvTranspose2dWeightBackward(desc), + )) + .output() + } + + fn conv_transpose2d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = ConvTranspose2dBiasBackwardOpIr::create( + x.into_ir(), + bias.into_ir(), + output_grad.into_ir(), + || client.create_empty_handle(), + ); + client + .register(OperationIr::Module( + ModuleOperationIr::ConvTranspose2dBiasBackward(desc), + )) + .output() + } + + fn conv_transpose3d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvTransposeOptions<3>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = ConvTranspose3dWeightBackwardOpIr::create( + x.into_ir(), + weight.into_ir(), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + client + .register(OperationIr::Module( + ModuleOperationIr::ConvTranspose3dWeightBackward(desc), + )) + .output() + } + + fn conv_transpose3d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = ConvTranspose3dBiasBackwardOpIr::create( + x.into_ir(), + bias.into_ir(), + output_grad.into_ir(), + || client.create_empty_handle(), + ); + client + .register(OperationIr::Module( + ModuleOperationIr::ConvTranspose3dBiasBackward(desc), + )) + .output() + } + + fn conv1d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<1>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = Conv1dOpIr::create( + x.into_ir(), + weight.into_ir(), + bias.map(|bias| bias.into_ir()), + options.into(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::Conv1d(desc))) + .output() + } + + fn conv1d_x_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<1>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = Conv1dXBackwardOpIr::create( + x.into_ir(), + weight.into_ir(), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::Conv1dXBackward( + desc, + ))) + .output() + } + + fn conv1d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<1>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = Conv1dWeightBackwardOpIr::create( + x.into_ir(), + weight.into_ir(), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module( + ModuleOperationIr::Conv1dWeightBackward(desc), + )) + .output() + } + + fn conv1d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = Conv1dBiasBackwardOpIr::create( + x.into_ir(), + bias.into_ir(), + output_grad.into_ir(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::Conv1dBiasBackward( + desc, + ))) + .output() + } + + fn conv2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<2>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = Conv2dOpIr::create( + x.into_ir(), + weight.into_ir(), + bias.map(|bias| bias.into_ir()), + options.into(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::Conv2d(desc))) + .output() + } + + fn conv2d_x_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<2>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = Conv2dXBackwardOpIr::create( + x.into_ir(), + weight.into_ir(), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::Conv2dXBackward( + desc, + ))) + .output() + } + + fn conv2d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<2>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = Conv2dWeightBackwardOpIr::create( + x.into_ir(), + weight.into_ir(), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module( + ModuleOperationIr::Conv2dWeightBackward(desc), + )) + .output() + } + + fn conv2d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = Conv2dBiasBackwardOpIr::create( + x.into_ir(), + bias.into_ir(), + output_grad.into_ir(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::Conv2dBiasBackward( + desc, + ))) + .output() + } + + fn conv3d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvOptions<3>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = Conv3dOpIr::create( + x.into_ir(), + weight.into_ir(), + bias.map(|bias| bias.into_ir()), + options.into(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::Conv3d(desc))) + .output() + } + + fn conv3d_x_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<3>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = Conv3dXBackwardOpIr::create( + x.into_ir(), + weight.into_ir(), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::Conv3dXBackward( + desc, + ))) + .output() + } + + fn conv3d_weight_backward( + x: FloatTensor, + weight: FloatTensor, + output_grad: FloatTensor, + options: ConvOptions<3>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = Conv3dWeightBackwardOpIr::create( + x.into_ir(), + weight.into_ir(), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module( + ModuleOperationIr::Conv3dWeightBackward(desc), + )) + .output() + } + + fn conv3d_bias_backward( + x: FloatTensor, + bias: FloatTensor, + output_grad: FloatTensor, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = Conv3dBiasBackwardOpIr::create( + x.into_ir(), + bias.into_ir(), + output_grad.into_ir(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::Conv3dBiasBackward( + desc, + ))) + .output() + } + + fn conv_transpose1d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<1>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = ConvTranspose1dOpIr::create( + x.into_ir(), + weight.into_ir(), + bias.map(|bias| bias.into_ir()), + options.into(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::ConvTranspose1d( + desc, + ))) + .output() + } + + fn conv_transpose2d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<2>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = ConvTranspose2dOpIr::create( + x.into_ir(), + weight.into_ir(), + bias.map(|bias| bias.into_ir()), + options.into(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::ConvTranspose2d( + desc, + ))) + .output() + } + + fn conv_transpose3d( + x: FloatTensor, + weight: FloatTensor, + bias: Option>, + options: ConvTransposeOptions<3>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = ConvTranspose3dOpIr::create( + x.into_ir(), + weight.into_ir(), + bias.map(|bias| bias.into_ir()), + options.into(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::ConvTranspose3d( + desc, + ))) + .output() + } + + fn avg_pool1d( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = AvgPool1dOpIr::create( + x.into_ir(), + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::AvgPool1d(desc))) + .output() + } + + fn avg_pool2d( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = AvgPool2dOpIr::create( + x.into_ir(), + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::AvgPool2d(desc))) + .output() + } + + fn avg_pool1d_backward( + x: FloatTensor, + grad: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = AvgPool1dBackwardOpIr::create( + x.into_ir(), + grad.into_ir(), + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::AvgPool1dBackward( + desc, + ))) + .output() + } + + fn avg_pool2d_backward( + x: FloatTensor, + grad: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = AvgPool2dBackwardOpIr::create( + x.into_ir(), + grad.into_ir(), + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::AvgPool2dBackward( + desc, + ))) + .output() + } + + fn max_pool1d( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = MaxPool1dOpIr::create( + x.into_ir(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::MaxPool1d(desc))) + .output() + } + + fn max_pool2d( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = MaxPool2dOpIr::create( + x.into_ir(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::MaxPool2d(desc))) + .output() + } + + fn max_pool1d_with_indices( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + indices_dtype: IntDType, + ) -> MaxPool1dWithIndices { + let client = x.client.clone(); + let desc = MaxPool1dWithIndicesOpIr::create( + x.into_ir(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + indices_dtype.into(), + || client.create_empty_handle(), + ); + + let [out, out_indices] = client + .register(OperationIr::Module( + ModuleOperationIr::MaxPool1dWithIndices(desc), + )) + .outputs(); + + MaxPool1dWithIndices::new(out, out_indices) + } + + fn max_pool2d_with_indices( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + indices_dtype: IntDType, + ) -> MaxPool2dWithIndices { + let client = x.client.clone(); + let desc = MaxPool2dWithIndicesOpIr::create( + x.into_ir(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + indices_dtype.into(), + || client.create_empty_handle(), + ); + + let [out, out_indices] = client + .register(OperationIr::Module( + ModuleOperationIr::MaxPool2dWithIndices(desc), + )) + .outputs(); + + MaxPool2dWithIndices::new(out, out_indices) + } + + fn max_pool1d_with_indices_backward( + x: FloatTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + output_grad: FloatTensor, + indices: IntTensor, + ) -> MaxPool1dBackward { + let client = x.client.clone(); + + let desc = MaxPool1dWithIndicesBackwardOpIr::create( + x.into_ir(), + output_grad.into_ir(), + indices.into_ir(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + || client.create_empty_handle(), + ); + + let out = client + .register(OperationIr::Module( + ModuleOperationIr::MaxPool1dWithIndicesBackward(desc), + )) + .output(); + + MaxPool1dBackward::new(out) + } + + fn max_pool2d_with_indices_backward( + x: FloatTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + output_grad: FloatTensor, + indices: IntTensor, + ) -> MaxPool2dBackward { + let client = x.client.clone(); + + let desc = MaxPool2dWithIndicesBackwardOpIr::create( + x.into_ir(), + output_grad.into_ir(), + indices.into_ir(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + || client.create_empty_handle(), + ); + + let out = client + .register(OperationIr::Module( + ModuleOperationIr::MaxPool2dWithIndicesBackward(desc), + )) + .output(); + + MaxPool2dBackward::new(out) + } + + fn adaptive_avg_pool1d(x: FloatTensor, output_size: usize) -> FloatTensor { + let client = x.client.clone(); + + let desc = AdaptiveAvgPool1dOpIr::create(x.into_ir(), output_size, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Module(ModuleOperationIr::AdaptiveAvgPool1d( + desc, + ))) + .output() + } + + fn adaptive_avg_pool2d(x: FloatTensor, output_size: [usize; 2]) -> FloatTensor { + let client = x.client.clone(); + + let desc = AdaptiveAvgPool2dOpIr::create(x.into_ir(), output_size, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Module(ModuleOperationIr::AdaptiveAvgPool2d( + desc, + ))) + .output() + } + + fn adaptive_avg_pool1d_backward( + x: FloatTensor, + grad: FloatTensor, + ) -> FloatTensor { + let client = x.client.clone(); + + let desc = AdaptiveAvgPool1dBackwardOpIr::create(x.into_ir(), grad.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Module( + ModuleOperationIr::AdaptiveAvgPool1dBackward(desc), + )) + .output() + } + + fn adaptive_avg_pool2d_backward( + x: FloatTensor, + grad: FloatTensor, + ) -> FloatTensor { + let client = x.client.clone(); + + let desc = AdaptiveAvgPool2dBackwardOpIr::create(x.into_ir(), grad.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Module( + ModuleOperationIr::AdaptiveAvgPool2dBackward(desc), + )) + .output() + } + + fn interpolate( + x: FloatTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = InterpolateOpIr::create(x.into_ir(), output_size, options.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Module(ModuleOperationIr::Interpolate(desc))) + .output() + } + + fn interpolate_backward( + x: FloatTensor, + grad: FloatTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = InterpolateBackwardOpIr::create( + x.into_ir(), + grad.into_ir(), + output_size, + options.into(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::InterpolateBackward( + desc, + ))) + .output() + } + + fn deform_conv2d( + x: FloatTensor, + offset: FloatTensor, + weight: FloatTensor, + mask: Option>, + bias: Option>, + options: DeformConvOptions<2>, + ) -> FloatTensor { + let client = x.client.clone(); + let desc = DeformConv2dOpIr::create( + x.into_ir(), + offset.into_ir(), + weight.into_ir(), + mask.map(|mask| mask.into_ir()), + bias.map(|bias| bias.into_ir()), + options.into(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::DeformableConv2d( + Box::new(desc), + ))) + .output() + } + + fn deform_conv2d_backward( + x: FloatTensor, + offset: FloatTensor, + weight: FloatTensor, + mask: Option>, + bias: Option>, + output_grad: FloatTensor, + options: DeformConvOptions<2>, + ) -> DeformConv2dBackward { + let client = x.client.clone(); + let has_bias = bias.is_some(); + let has_mask = mask.is_some(); + + let desc = DeformConv2dBackwardOpIr::create( + x.into_ir(), + offset.into_ir(), + weight.into_ir(), + mask.map(|mask| mask.into_ir()), + bias.map(|bias| bias.into_ir()), + output_grad.into_ir(), + options.into(), + || client.create_empty_handle(), + ); + let mut outputs = client + .register(OperationIr::Module( + ModuleOperationIr::DeformableConv2dBackward(Box::new(desc)), + )) + .into_iter(); + + // When the number of outputs is variable, the order is important + let input_grad = outputs.next().unwrap(); + let offset_grad = outputs.next().unwrap(); + let weight_grad = outputs.next().unwrap(); + let mask_grad = has_mask.then(|| outputs.next().unwrap()); + let bias_grad = has_bias.then(|| outputs.next().unwrap()); + + DeformConv2dBackward::new(input_grad, offset_grad, weight_grad, mask_grad, bias_grad) + } + + fn attention( + query: FloatTensor, + key: FloatTensor, + value: FloatTensor, + mask: Option>, + attn_bias: Option>, + options: AttentionModuleOptions, + ) -> FloatTensor { + let client = query.client.clone(); + let desc = AttentionOpIr::create( + query.into_ir(), + key.into_ir(), + value.into_ir(), + mask.map(|m: BoolTensor| m.into_ir()), + attn_bias.map(|ab| ab.into_ir()), + options.into(), + || client.create_empty_handle(), + ); + + client + .register(OperationIr::Module(ModuleOperationIr::Attention(desc))) + .output() + } + + fn rfft( + _signal: FloatTensor, + _dim: usize, + _n: Option, + ) -> (FloatTensor, FloatTensor) { + todo!("rfft is not supported for backend-router") + } + + fn irfft( + _spectrum_re: FloatTensor, + _spectrum_im: FloatTensor, + _dim: usize, + _n: Option, + ) -> FloatTensor { + todo!("irfft is not supported for backend-router") + } +} diff --git a/crates/burn-router/src/ops/qtensor.rs b/crates/burn-router/src/ops/qtensor.rs new file mode 100644 index 0000000..6740611 --- /dev/null +++ b/crates/burn-router/src/ops/qtensor.rs @@ -0,0 +1,64 @@ +use burn_backend::{ + ExecutionError, FloatDType, Shape, TensorData, + ops::QTensorOps, + quantization::{QuantScheme, QuantizationParametersPrimitive}, + tensor::{Device, FloatTensor, QuantizedTensor}, +}; + +use crate::{BackendRouter, RouterChannel}; + +impl QTensorOps for BackendRouter { + fn q_from_data(_data: TensorData, _device: &Device) -> QuantizedTensor { + unimplemented!() + } + + fn quantize( + _tensor: FloatTensor, + _scheme: &QuantScheme, + _qparams: QuantizationParametersPrimitive, + ) -> QuantizedTensor { + unimplemented!() + } + + fn quantize_dynamic( + _tensor: FloatTensor, + _scheme: &QuantScheme, + ) -> QuantizedTensor { + unimplemented!() + } + + fn dequantize(_tensor: QuantizedTensor, _dtype: FloatDType) -> FloatTensor { + unimplemented!() + } + + fn q_to_device( + _tensor: QuantizedTensor, + _device: &Device, + ) -> QuantizedTensor { + unimplemented!() + } + + fn q_reshape(_tensor: QuantizedTensor, _shape: Shape) -> QuantizedTensor { + unimplemented!() + } + + async fn q_into_data(_tensor: QuantizedTensor) -> Result { + unimplemented!() + } + + fn q_swap_dims( + _tensor: QuantizedTensor, + _dim1: usize, + _dim2: usize, + ) -> QuantizedTensor { + unimplemented!() + } + + fn q_permute(_tensor: QuantizedTensor, _axes: &[usize]) -> QuantizedTensor { + unimplemented!() + } + + fn q_flip(_tensor: QuantizedTensor, _axes: &[usize]) -> QuantizedTensor { + unimplemented!() + } +} diff --git a/crates/burn-router/src/ops/tensor.rs b/crates/burn-router/src/ops/tensor.rs new file mode 100644 index 0000000..3243721 --- /dev/null +++ b/crates/burn-router/src/ops/tensor.rs @@ -0,0 +1,1593 @@ +use alloc::vec::Vec; +use burn_backend::Scalar; +use burn_backend::backend::ExecutionError; +use burn_std::{BoolDType, IntDType}; + +use crate::{BackendRouter, RouterChannel, RouterClient, get_client}; +use burn_backend::tensor::{BoolTensor, Device, FloatTensor, IndexingUpdateOp, IntTensor}; +use burn_backend::{Distribution, FloatDType, Shape, Slice, TensorData, ops::FloatTensorOps}; +use burn_ir::{ + BaseOperationIr, BinaryOpIr, CastOpIr, CatOpIr, ClampOpIr, CreationOpIr, CrossOpIr, DimOpIr, + FlipOpIr, FloatOperationIr, FullOpIr, GatherNdOpIr, GatherOpIr, GridSample2dOpIr, + InitOperationIr, MaskFillOpIr, MaskWhereOpIr, MatmulOpIr, NumericOperationIr, OperationIr, + OperationOutput, PermuteOpIr, RandomOpIr, ReduceDimOpIr, ReduceDimWithIndicesOpIr, ReduceOpIr, + RepeatDimOpIr, ScalarOpIr, ScatterNdOpIr, ScatterOpIr, SelectAssignOpIr, SelectOpIr, ShapeOpIr, + SliceAssignOpIr, SliceOpIr, SortOpIr, SortWithIndicesOpIr, SwapDimsOpIr, UnaryOpIr, UnfoldOpIr, +}; + +impl FloatTensorOps for BackendRouter { + fn float_from_data(data: TensorData, device: &Device) -> FloatTensor { + let client = get_client::(device); + let out = client.register_tensor_data(data); + let desc = InitOperationIr { + out: out.to_ir_out(), + }; + + // Call register op when output is already initialized + client.register_op(OperationIr::Init(desc)); + + out + } + + fn float_random( + shape: Shape, + distribution: Distribution, + device: &Device, + dtype: FloatDType, + ) -> FloatTensor { + let client = get_client::(device); + let dtype = dtype.into(); + let desc = RandomOpIr::create(shape, dtype, distribution, || client.create_empty_handle()); + + client + .register(OperationIr::Float(dtype, FloatOperationIr::Random(desc))) + .output() + } + + fn float_zeros(shape: Shape, device: &Device, dtype: FloatDType) -> FloatTensor { + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::Zeros(desc))) + .output() + } + + fn float_ones(shape: Shape, device: &Device, dtype: FloatDType) -> FloatTensor { + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::Ones(desc))) + .output() + } + + fn float_full( + shape: Shape, + fill_value: Scalar, + device: &Device, + dtype: FloatDType, + ) -> FloatTensor { + let client = get_client::(device); + let dtype = dtype.into(); + let value = fill_value.into(); + let desc = FullOpIr::create(shape, dtype, value, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Full(desc), + )) + .output() + } + + async fn float_into_data(tensor: FloatTensor) -> Result { + tensor.into_data().await + } + + fn float_to_device(tensor: FloatTensor, device: &Device) -> FloatTensor { + if &tensor.client.device() == device { + return tensor; + } + R::change_client_backend(tensor, device) + } + + fn float_into_int(tensor: FloatTensor, out_dtype: IntDType) -> IntTensor { + let client = tensor.client.clone(); + let desc = CastOpIr::create(tensor.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Float( + desc.input.dtype, + FloatOperationIr::IntoInt(desc), + )) + .output() + } + + fn float_empty(shape: Shape, device: &Device, dtype: FloatDType) -> FloatTensor { + let client = get_client::(device); + let desc = CreationOpIr::create(shape, dtype.into(), || client.create_empty_handle()); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::Empty(desc))) + .output() + } + + fn float_add(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Add(desc), + )) + .output() + } + + fn float_add_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::AddScalar(desc), + )) + .output() + } + + fn float_clamp(tensor: FloatTensor, min: Scalar, max: Scalar) -> FloatTensor { + let client = tensor.client.clone(); + let min = min.into(); + let max = max.into(); + let desc = ClampOpIr::create(tensor.into_ir(), min, max, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Clamp(desc), + )) + .output() + } + + fn float_sub(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Sub(desc), + )) + .output() + } + + fn float_sub_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::SubScalar(desc), + )) + .output() + } + + fn float_mul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Mul(desc), + )) + .output() + } + + fn float_mul_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::MulScalar(desc), + )) + .output() + } + + fn float_div(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Div(desc), + )) + .output() + } + + fn float_div_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::DivScalar(desc), + )) + .output() + } + + fn float_remainder(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Rem(desc), + )) + .output() + } + + fn float_remainder_scalar(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::RemScalar(desc), + )) + .output() + } + + fn float_matmul(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + let client = lhs.client.clone(); + let desc = MatmulOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Matmul(desc), + )) + .output() + } + + fn float_cross( + lhs: FloatTensor, + rhs: FloatTensor, + dim: usize, + ) -> FloatTensor { + let client = lhs.client.clone(); + let desc = CrossOpIr::create(lhs.into_ir(), rhs.into_ir(), dim, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Cross(desc), + )) + .output() + } + + fn float_grid_sample_2d( + tensor: FloatTensor, + grid: FloatTensor, + options: burn_backend::ops::GridSampleOptions, + ) -> FloatTensor { + let client = tensor.client.clone(); + let desc = + GridSample2dOpIr::create(tensor.into_ir(), grid.into_ir(), options.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::GridSample2d(desc), + )) + .output() + } + + fn float_swap_dims(tensor: FloatTensor, dim1: usize, dim2: usize) -> FloatTensor { + let client = tensor.client.clone(); + let desc = SwapDimsOpIr::create(tensor.into_ir(), dim1, dim2, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::SwapDims(desc))) + .output() + } + + fn float_reshape(tensor: FloatTensor, shape: Shape) -> FloatTensor { + let client = tensor.client.clone(); + let desc = ShapeOpIr::reshape(tensor.into_ir(), shape, || client.create_empty_handle()); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::Reshape(desc))) + .output() + } + + fn float_gather( + dim: usize, + tensor: FloatTensor, + indices: IntTensor, + ) -> FloatTensor { + let client = tensor.client.clone(); + let desc = GatherOpIr::create(tensor.into_ir(), dim, indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::Gather(desc))) + .output() + } + + fn float_scatter_add( + dim: usize, + tensor: FloatTensor, + indices: IntTensor, + value: FloatTensor, + ) -> FloatTensor { + let client = tensor.client.clone(); + let desc = ScatterOpIr::create( + tensor.into_ir(), + dim, + indices.into_ir(), + value.into_ir(), + IndexingUpdateOp::Add, + || client.create_empty_handle(), + ); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::Scatter(desc))) + .output() + } + + fn float_scatter_nd( + data: FloatTensor, + indices: IntTensor, + values: FloatTensor, + reduction: IndexingUpdateOp, + ) -> FloatTensor { + let client = data.client.clone(); + let desc = ScatterNdOpIr::create( + data.into_ir(), + indices.into_ir(), + values.into_ir(), + reduction, + || client.create_empty_handle(), + ); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::ScatterNd(desc))) + .output() + } + + fn float_gather_nd(data: FloatTensor, indices: IntTensor) -> FloatTensor { + let client = data.client.clone(); + let desc = GatherNdOpIr::create(data.into_ir(), indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::GatherNd(desc))) + .output() + } + + fn float_select( + tensor: FloatTensor, + dim: usize, + indices: IntTensor, + ) -> FloatTensor { + let client = tensor.client.clone(); + let desc = SelectOpIr::create(tensor.into_ir(), dim, indices.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::Select(desc))) + .output() + } + + fn float_select_add( + tensor: FloatTensor, + dim: usize, + indices: IntTensor, + value: FloatTensor, + ) -> FloatTensor { + let client = tensor.client.clone(); + let desc = SelectAssignOpIr::create( + tensor.into_ir(), + dim, + indices.into_ir(), + value.into_ir(), + IndexingUpdateOp::Add, + || client.create_empty_handle(), + ); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::SelectAssign(desc))) + .output() + } + + fn float_slice(tensor: FloatTensor, slices: &[Slice]) -> FloatTensor { + let client = tensor.client.clone(); + let desc = SliceOpIr::create(tensor.into_ir(), slices.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::Slice(desc))) + .output() + } + + fn float_slice_assign( + tensor: FloatTensor, + slices: &[burn_backend::Slice], + value: FloatTensor, + ) -> FloatTensor { + let client = tensor.client.clone(); + let desc = + SliceAssignOpIr::create(tensor.into_ir(), slices.into(), value.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::SliceAssign(desc))) + .output() + } + + fn float_mask_where( + tensor: FloatTensor, + mask: BoolTensor, + value: FloatTensor, + ) -> FloatTensor { + let client = tensor.client.clone(); + let desc = MaskWhereOpIr::create(tensor.into_ir(), mask.into_ir(), value.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::MaskWhere(desc))) + .output() + } + + fn float_mask_fill( + tensor: FloatTensor, + mask: BoolTensor, + value: Scalar, + ) -> FloatTensor { + let client = tensor.client.clone(); + let value = value.into(); + let desc = MaskFillOpIr::create(tensor.into_ir(), mask.into_ir(), value, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::MaskFill(desc))) + .output() + } + + fn float_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::Equal(desc))) + .output() + } + + fn float_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::EqualElem(desc))) + .output() + } + + fn float_greater( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.lhs.dtype, + NumericOperationIr::Greater(desc), + )) + .output() + } + + fn float_greater_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.lhs.dtype, + NumericOperationIr::GreaterElem(desc), + )) + .output() + } + + fn float_greater_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.lhs.dtype, + NumericOperationIr::GreaterEqual(desc), + )) + .output() + } + + fn float_greater_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.lhs.dtype, + NumericOperationIr::GreaterEqualElem(desc), + )) + .output() + } + + fn float_lower( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.lhs.dtype, + NumericOperationIr::Lower(desc), + )) + .output() + } + + fn float_lower_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.lhs.dtype, + NumericOperationIr::LowerElem(desc), + )) + .output() + } + + fn float_lower_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.lhs.dtype, + NumericOperationIr::LowerEqual(desc), + )) + .output() + } + + fn float_lower_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create_comparison(lhs.into_ir(), rhs, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.lhs.dtype, + NumericOperationIr::LowerEqualElem(desc), + )) + .output() + } + + fn float_sum(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Sum(desc), + )) + .output() + } + + fn float_sum_dim(tensor: FloatTensor, axis: usize) -> FloatTensor { + let client = tensor.client.clone(); + let desc = + ReduceDimOpIr::create(tensor.into_ir(), axis, 1, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::SumDim(desc), + )) + .output() + } + + fn float_prod(tensor: IntTensor) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Prod(desc), + )) + .output() + } + + fn float_prod_dim(tensor: IntTensor, dim: usize) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::ProdDim(desc), + )) + .output() + } + + fn float_mean(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Mean(desc), + )) + .output() + } + + fn float_mean_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::MeanDim(desc), + )) + .output() + } + + fn float_cumsum(tensor: FloatTensor, dim: usize) -> FloatTensor { + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::CumSum(desc), + )) + .output() + } + + fn float_cumprod(tensor: FloatTensor, dim: usize) -> FloatTensor { + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::CumProd(desc), + )) + .output() + } + + fn float_cummin(tensor: FloatTensor, dim: usize) -> FloatTensor { + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::CumMin(desc), + )) + .output() + } + + fn float_cummax(tensor: FloatTensor, dim: usize) -> FloatTensor { + let client = tensor.client.clone(); + let desc = DimOpIr::create(tensor.into_ir(), dim, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::CumMax(desc), + )) + .output() + } + + fn float_exp(lhs: FloatTensor) -> FloatTensor { + let client = lhs.client.clone(); + let desc = UnaryOpIr::create(lhs.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Exp(desc), + )) + .output() + } + + fn float_log(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Log(desc), + )) + .output() + } + + fn float_log1p(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Log1p(desc), + )) + .output() + } + + fn float_powf_scalar_impl(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::PowfScalar(desc), + )) + .output() + } + + fn float_sqrt(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Sqrt(desc), + )) + .output() + } + + fn float_abs(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Abs(desc), + )) + .output() + } + + fn float_cos(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Cos(desc), + )) + .output() + } + + fn float_cosh(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Cosh(desc), + )) + .output() + } + + fn float_sin(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Sin(desc), + )) + .output() + } + + fn float_sinh(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Sinh(desc), + )) + .output() + } + + fn float_tan(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Tan(desc), + )) + .output() + } + + fn float_tanh(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Tanh(desc), + )) + .output() + } + + fn float_acos(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::ArcCos(desc), + )) + .output() + } + + fn float_acosh(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::ArcCosh(desc), + )) + .output() + } + + fn float_asin(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::ArcSin(desc), + )) + .output() + } + + fn float_asinh(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::ArcSinh(desc), + )) + .output() + } + + fn float_atan(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::ArcTan(desc), + )) + .output() + } + + fn float_atanh(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::ArcTanh(desc), + )) + .output() + } + + fn float_atan2(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::ArcTan2(desc), + )) + .output() + } + + fn float_hypot(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Hypot(desc), + )) + .output() + } + + fn float_round(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Round(desc), + )) + .output() + } + + fn float_floor(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Floor(desc), + )) + .output() + } + + fn float_ceil(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Ceil(desc), + )) + .output() + } + + fn float_trunc(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Trunc(desc), + )) + .output() + } + + fn float_recip(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Recip(desc), + )) + .output() + } + + fn float_erf(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Erf(desc), + )) + .output() + } + + fn float_cat(tensors: Vec>, dim: usize) -> FloatTensor { + let client = tensors.first().unwrap().client.clone(); + let tensors = tensors.into_iter().map(|t| t.into_ir()).collect(); + let desc = CatOpIr::create(tensors, dim, || client.create_empty_handle()); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::Cat(desc))) + .output() + } + + fn float_argmax(tensor: FloatTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create_arg(tensor.into_ir(), dim, 1, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.input.dtype, + NumericOperationIr::ArgMax(desc), + )) + .output() + } + + fn float_argtopk( + tensor: FloatTensor, + dim: usize, + k: usize, + out_dtype: IntDType, + ) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create_arg(tensor.into_ir(), dim, k, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.input.dtype, + NumericOperationIr::ArgTopK(desc), + )) + .output() + } + + fn float_topk(tensor: FloatTensor, dim: usize, k: usize) -> FloatTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, k, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.input.dtype, + NumericOperationIr::TopK(desc), + )) + .output() + } + + fn float_repeat_dim(tensor: FloatTensor, dim: usize, times: usize) -> FloatTensor { + let client = tensor.client.clone(); + let desc = RepeatDimOpIr::create(tensor.into_ir(), dim, times, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::RepeatDim(desc))) + .output() + } + + fn float_argmin(tensor: FloatTensor, dim: usize, out_dtype: IntDType) -> IntTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create_arg(tensor.into_ir(), dim, 1, out_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.input.dtype, + NumericOperationIr::ArgMin(desc), + )) + .output() + } + + fn float_max(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Max(desc), + )) + .output() + } + + fn float_max_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::MaxDim(desc), + )) + .output() + } + + fn float_max_dim_with_indices( + tensor: FloatTensor, + dim: usize, + indices_dtype: IntDType, + ) -> (FloatTensor, IntTensor) { + let client = tensor.client.clone(); + let desc = + ReduceDimWithIndicesOpIr::create(tensor.into_ir(), dim, indices_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.tensor.dtype, + NumericOperationIr::MaxDimWithIndices(desc), + )) + .outputs() + .into() + } + + fn float_max_abs(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::MaxAbs(desc), + )) + .output() + } + + fn float_max_abs_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::MaxAbsDim(desc), + )) + .output() + } + + fn float_min(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = ReduceOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Min(desc), + )) + .output() + } + + fn float_min_dim(tensor: FloatTensor, dim: usize) -> FloatTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create(tensor.into_ir(), dim, 1, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::MinDim(desc), + )) + .output() + } + + fn float_min_dim_with_indices( + tensor: FloatTensor, + dim: usize, + indices_dtype: IntDType, + ) -> (FloatTensor, IntTensor) { + let client = tensor.client.clone(); + let desc = + ReduceDimWithIndicesOpIr::create(tensor.into_ir(), dim, indices_dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.tensor.dtype, + NumericOperationIr::MinDimWithIndices(desc), + )) + .outputs() + .into() + } + + fn float_powf(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::Float( + desc.out.dtype, + FloatOperationIr::Powf(desc), + )) + .output() + } + + fn float_neg(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Neg(desc), + )) + .output() + } + + fn float_sign(tensor: FloatTensor) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnaryOpIr::create(tensor.into_ir(), || client.create_empty_handle()); + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Sign(desc), + )) + .output() + } + + fn float_clamp_min(tensor: FloatTensor, min: Scalar) -> FloatTensor { + let client = tensor.client.clone(); + let desc = ScalarOpIr::create(tensor.into_ir(), min.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::ClampMin(desc), + )) + .output() + } + + fn float_clamp_max(tensor: FloatTensor, max: Scalar) -> FloatTensor { + let client = tensor.client.clone(); + let desc = ScalarOpIr::create(tensor.into_ir(), max.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::ClampMax(desc), + )) + .output() + } + + fn float_not_equal( + lhs: FloatTensor, + rhs: FloatTensor, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let desc = + BinaryOpIr::create_comparison(lhs.into_ir(), rhs.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::BaseFloat(BaseOperationIr::NotEqual(desc))) + .output() + } + + fn float_not_equal_elem( + lhs: FloatTensor, + rhs: Scalar, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + let client = lhs.client.clone(); + let desc = + ScalarOpIr::create_comparison(lhs.into_ir(), rhs.into(), out_dtype.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::BaseFloat(BaseOperationIr::NotEqualElem(desc))) + .output() + } + + fn float_all(tensor: FloatTensor, out_dtype: burn_std::BoolDType) -> BoolTensor { + let client = tensor.client.clone(); + let desc = ReduceOpIr::create_bool(tensor.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::BaseFloat(BaseOperationIr::All(desc))) + .output() + } + + fn float_any(tensor: FloatTensor, out_dtype: burn_std::BoolDType) -> BoolTensor { + let client = tensor.client.clone(); + let desc = ReduceOpIr::create_bool(tensor.into_ir(), out_dtype.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::BaseFloat(BaseOperationIr::Any(desc))) + .output() + } + + fn float_all_dim( + tensor: FloatTensor, + dim: usize, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create_bool(tensor.into_ir(), dim, 1, out_dtype.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::BaseFloat(BaseOperationIr::AllDim(desc))) + .output() + } + + fn float_any_dim( + tensor: FloatTensor, + dim: usize, + out_dtype: burn_std::BoolDType, + ) -> BoolTensor { + let client = tensor.client.clone(); + let desc = ReduceDimOpIr::create_bool(tensor.into_ir(), dim, 1, out_dtype.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::BaseFloat(BaseOperationIr::AnyDim(desc))) + .output() + } + + fn float_sort(tensor: FloatTensor, dim: usize, descending: bool) -> FloatTensor { + let client = tensor.client.clone(); + let desc = SortOpIr::create(tensor.into_ir(), dim, descending, || { + client.create_empty_handle() + }); + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Sort(desc), + )) + .output() + } + + fn float_sort_with_indices( + tensor: FloatTensor, + dim: usize, + descending: bool, + indices_dtype: IntDType, + ) -> (FloatTensor, IntTensor) { + let client = tensor.client.clone(); + let desc = SortWithIndicesOpIr::create( + tensor.into_ir(), + dim, + descending, + indices_dtype.into(), + || client.create_empty_handle(), + ); + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::SortWithIndices(desc), + )) + .outputs() + .into() + } + + fn float_argsort( + tensor: FloatTensor, + dim: usize, + descending: bool, + out_dtype: IntDType, + ) -> IntTensor { + let client = tensor.client.clone(); + let dtype = tensor.dtype; + let desc = + SortOpIr::create_arg(tensor.into_ir(), dim, descending, out_dtype.into(), || { + client.create_empty_handle() + }); + client + .register(OperationIr::NumericFloat( + dtype, + NumericOperationIr::ArgSort(desc), + )) + .output() + } + + fn float_powi(lhs: FloatTensor, rhs: IntTensor) -> FloatTensor { + let client = lhs.client.clone(); + let desc = BinaryOpIr::create(lhs.into_ir(), rhs.into_ir(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::Powi(desc), + )) + .output() + } + + fn float_powi_scalar_impl(lhs: FloatTensor, rhs: Scalar) -> FloatTensor { + let client = lhs.client.clone(); + let rhs = rhs.into(); + let desc = ScalarOpIr::create(lhs.into_ir(), rhs, || client.create_empty_handle()); + + client + .register(OperationIr::NumericFloat( + desc.out.dtype, + NumericOperationIr::PowiScalar(desc), + )) + .output() + } + + fn float_permute(tensor: FloatTensor, axes: &[usize]) -> FloatTensor { + let client = tensor.client.clone(); + let desc = PermuteOpIr::create(tensor.into_ir(), axes.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::Permute(desc))) + .output() + } + + fn float_expand(tensor: FloatTensor, shape: Shape) -> FloatTensor { + let client = tensor.client.clone(); + let desc = ShapeOpIr::expand(tensor.into_ir(), shape, || client.create_empty_handle()); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::Expand(desc))) + .output() + } + + fn float_flip(tensor: FloatTensor, axes: &[usize]) -> FloatTensor { + let client = tensor.client.clone(); + let desc = FlipOpIr::create(tensor.into_ir(), axes.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::Flip(desc))) + .output() + } + + fn float_cast(tensor: FloatTensor, dtype: burn_backend::FloatDType) -> FloatTensor { + let client = tensor.client.clone(); + let desc = CastOpIr::create(tensor.into_ir(), dtype.into(), || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::Cast(desc))) + .output() + } + + fn float_unfold( + tensor: FloatTensor, + dim: usize, + size: usize, + step: usize, + ) -> FloatTensor { + let client = tensor.client.clone(); + let desc = UnfoldOpIr::create(tensor.into_ir(), dim, size, step, || { + client.create_empty_handle() + }); + + client + .register(OperationIr::BaseFloat(BaseOperationIr::Unfold(desc))) + .output() + } +} diff --git a/crates/burn-router/src/ops/transaction.rs b/crates/burn-router/src/ops/transaction.rs new file mode 100644 index 0000000..ce7b9e0 --- /dev/null +++ b/crates/burn-router/src/ops/transaction.rs @@ -0,0 +1,5 @@ +use burn_backend::ops::TransactionOps; + +use crate::{BackendRouter, RouterChannel}; + +impl TransactionOps for BackendRouter {} diff --git a/crates/burn-router/src/ops/unary.rs b/crates/burn-router/src/ops/unary.rs new file mode 100644 index 0000000..0e6cb01 --- /dev/null +++ b/crates/burn-router/src/ops/unary.rs @@ -0,0 +1,160 @@ +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! scalar_float_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let lhs = $handles.get_float_tensor::(&$desc.lhs); + let output = $ops(lhs, $desc.rhs.into()); + + $handles.register_float_tensor::(&$desc.out.id, output); + }}; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! scalar_float_dim_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let lhs = $handles.get_float_tensor::(&$desc.lhs); + let output = $ops(lhs, $desc.rhs); + + $handles.register_float_tensor::(&$desc.out.id, output); + }}; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! reduce_float_dim_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let input = $handles.get_float_tensor::(&$desc.input); + let output = $ops(input, $desc.axis, $desc.accumulator_len); + + $handles.register_float_tensor::(&$desc.out.id, output); + }}; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! reduce_float2int_dim_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let input = $handles.get_float_tensor::(&$desc.input); + let output = $ops( + input, + $desc.axis, + $desc.accumulator_len, + $desc.out.dtype.into(), + ); + + $handles.register_int_tensor::(&$desc.out.id, output); + }}; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! reduce_int_dim_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let input = $handles.get_int_tensor::(&$desc.input); + let output = $ops(input, $desc.axis, $desc.accumulator_len); + + $handles.register_int_tensor::(&$desc.out.id, output); + }}; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! scalar_float2int_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let lhs = $handles.get_float_tensor::(&$desc.lhs); + let output = $ops(lhs, $desc.rhs); + + $handles.register_int_tensor::(&$desc.out.id, output); + }}; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! scalar_float_cmp_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let lhs = $handles.get_float_tensor::(&$desc.lhs); + let output = $ops(lhs, $desc.rhs.into(), $desc.out.dtype.into()); + + $handles.register_bool_tensor::(&$desc.out.id, output); + }}; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! unary_float_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let lhs = $handles.get_float_tensor::(&$desc.input); + let output = $ops(lhs); + + $handles.register_float_tensor::(&$desc.out.id, output); + }}; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! scalar_int_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let lhs = $handles.get_int_tensor::(&$desc.lhs); + let output = $ops(lhs, $desc.rhs.into()); + + $handles.register_int_tensor::(&$desc.out.id, output); + }}; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! scalar_int_dim_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let lhs = $handles.get_int_tensor::(&$desc.lhs); + let output = $ops(lhs, $desc.rhs); + + $handles.register_int_tensor::(&$desc.out.id, output); + }}; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! scalar_int_cmp_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let lhs = $handles.get_int_tensor::(&$desc.lhs); + let output = $ops(lhs, $desc.rhs.into(), $desc.out.dtype.into()); + + $handles.register_bool_tensor::(&$desc.out.id, output); + }}; +} + +#[allow(missing_docs)] +#[macro_export(local_inner_macros)] +macro_rules! unary_int_ops { + ( + $handles:expr, $desc:expr, $ops:expr + ) => {{ + let lhs = $handles.get_int_tensor::(&$desc.input); + let output = $ops(lhs); + + $handles.register_int_tensor::(&$desc.out.id, output); + }}; +} diff --git a/crates/burn-router/src/tensor.rs b/crates/burn-router/src/tensor.rs new file mode 100644 index 0000000..2b6bb60 --- /dev/null +++ b/crates/burn-router/src/tensor.rs @@ -0,0 +1,158 @@ +use core::sync::atomic::{AtomicU32, Ordering}; + +use alloc::format; +use alloc::{sync::Arc, vec::Vec}; + +use super::RouterClient; +use burn_backend::{DType, Shape, TensorData, TensorMetadata, backend::ExecutionError}; +use burn_ir::{TensorId, TensorIr, TensorStatus}; + +/// Tensor primitive for the [router backend](crate::BackendRouter). +pub struct RouterTensor { + pub(crate) id: TensorId, + pub(crate) shape: Shape, + pub(crate) dtype: DType, + /// The client that has this tensor + pub client: C, + pub(crate) count: Arc, +} + +impl TensorMetadata for RouterTensor { + type Device = C::Device; + fn dtype(&self) -> DType { + self.dtype + } + + fn shape(&self) -> Shape { + self.shape.clone() + } + + fn rank(&self) -> usize { + self.shape.num_dims() + } + + fn device(&self) -> Self::Device { + self.client.device() + } + + fn can_mut(&self) -> bool { + // Same sharing rule as `into_ir`: a handle cloned on its stream + // (count > 1) is read-only for the runner, a unique one is read-write. + self.count.load(Ordering::Acquire) <= 1 + } +} + +impl RouterTensor { + /// The id identifying this tensor on its [client](RouterClient). + pub fn id(&self) -> TensorId { + self.id + } + + /// Create a new router tensor. + pub fn new(id: TensorId, shape: Shape, dtype: DType, client: C) -> Self { + Self { + id, + shape, + dtype, + client, + count: Arc::new(AtomicU32::new(1)), + } + } + + pub(crate) async fn into_data(self) -> Result { + self.client.clone().read_tensor_async(self.into_ir()).await + } + + /// Get the ir for this tensor + pub fn into_ir(mut self) -> TensorIr { + let count = self.count.load(Ordering::Relaxed); + let status = self.status(count); + let mut shape_out = Shape::from(Vec::::new()); + core::mem::swap(&mut self.shape, &mut shape_out); + + if let TensorStatus::ReadWrite = status { + // Avoids an unwanted drop on the same thread. + // + // Since `drop` is called after `into_ir`, we must not register a drop if the tensor + // was consumed with a `ReadWrite` status. + self.count.fetch_add(1, Ordering::Relaxed); + } + + TensorIr { + status, + shape: shape_out, + id: self.id, + dtype: self.dtype, + } + } + + pub(crate) fn to_ir_out(&self) -> TensorIr { + TensorIr { + status: TensorStatus::NotInit, + shape: self.shape.clone(), + id: self.id, + dtype: self.dtype, + } + } + + pub(crate) fn status(&self, count: u32) -> TensorStatus { + if count <= 1 { + TensorStatus::ReadWrite + } else { + TensorStatus::ReadOnly + } + } +} + +impl core::fmt::Debug for RouterTensor { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_str( + format!( + "{{ id: {:?}, shape: {:?}, dtype: {:?}, device: {:?} }}", + self.id, + self.shape, + self.dtype, + self.client.device(), + ) + .as_str(), + ) + } +} + +impl Clone for RouterTensor { + fn clone(&self) -> Self { + self.count.fetch_add(1, Ordering::Relaxed); + + Self { + id: self.id, + shape: self.shape.clone(), + client: self.client.clone(), + dtype: self.dtype, + count: self.count.clone(), + } + } +} + +impl Drop for RouterTensor { + fn drop(&mut self) { + let count = self.count.fetch_sub(1, Ordering::Relaxed); + + match self.status(count) { + TensorStatus::ReadWrite => { + let id = self.id; + let mut shape = Shape::from(Vec::::new()); + core::mem::swap(&mut shape, &mut self.shape); + + let ir = TensorIr { + id, + shape, + status: TensorStatus::ReadWrite, + dtype: self.dtype, + }; + self.client.register_op(burn_ir::OperationIr::Drop(ir)); + } + TensorStatus::ReadOnly => {} + TensorStatus::NotInit => {} + } + } +} diff --git a/crates/burn-std/Cargo.toml b/crates/burn-std/Cargo.toml new file mode 100644 index 0000000..ddcde75 --- /dev/null +++ b/crates/burn-std/Cargo.toml @@ -0,0 +1,73 @@ +[package] +authors = ["Dilshod Tadjibaev (@antimora)"] +categories = [] +description = "Core types and utilities shared across the Burn ecosystem." +documentation = "https://docs.rs/burn-std" +edition.workspace = true +keywords = [] +license.workspace = true +name = "burn-std" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-std" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["std", "cubecl-common/default"] +doc = ["default"] +std = ["cubecl-common/std", "num-traits/std", "rand/std"] +tracing = ["cubecl-common/tracing"] + +network = ["dep:indicatif", "dep:reqwest", "dep:tokio"] + +[dependencies] +ahash = { workspace = true } +hashbrown = { workspace = true } +bytemuck = { workspace = true, features = ["extern_crate_alloc"] } +derive-new = { workspace = true } +enumset = { workspace = true } +half = { workspace = true, features = ["bytemuck"] } +num-traits = { workspace = true } +rand = { workspace = true, default-features = false } +rand_distr = { workspace = true } +serde = { workspace = true } +smallvec = { workspace = true, features = ["serde"] } +spin = { workspace = true } +thiserror = { workspace = true } + +# ParamId +data-encoding = { workspace = true } +uuid = { workspace = true } + +cubecl-common = { workspace = true, default-features = false, features = [ + "serde", + "shared-bytes", +] } +cubecl-zspace = { workspace = true, default-features = false } +# Enable extra-platforms for portable-atomic support on targets without native atomics (e.g., thumbv6m) +# This is needed because cubecl-common's shared-bytes feature pulls in bytes +bytes = { workspace = true } + +# Network downloader +indicatif = { workspace = true, optional = true } +reqwest = { workspace = true, optional = true } +tokio = { workspace = true, optional = true } + +[dev-dependencies] +dashmap = { workspace = true } +paste = { workspace = true } +rand = { workspace = true, features = ["thread_rng"] } +serde_json = { workspace = true, features = ["alloc"] } +serial_test = { workspace = true } +tempfile = { workspace = true } + +# Enable extra-platforms for bytes on targets without native atomics (e.g., thumbv6m-none-eabi) +[target.'cfg(not(target_has_atomic = "ptr"))'.dependencies] +bytes = { workspace = true, features = ["extra-platforms"] } +portable-atomic-util = { workspace = true } + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-std/LICENSE-APACHE b/crates/burn-std/LICENSE-APACHE new file mode 120000 index 0000000..1cd601d --- /dev/null +++ b/crates/burn-std/LICENSE-APACHE @@ -0,0 +1 @@ +../../LICENSE-APACHE \ No newline at end of file diff --git a/crates/burn-std/LICENSE-MIT b/crates/burn-std/LICENSE-MIT new file mode 120000 index 0000000..b2cfbdc --- /dev/null +++ b/crates/burn-std/LICENSE-MIT @@ -0,0 +1 @@ +../../LICENSE-MIT \ No newline at end of file diff --git a/crates/burn-std/README.md b/crates/burn-std/README.md new file mode 100644 index 0000000..6969baf --- /dev/null +++ b/crates/burn-std/README.md @@ -0,0 +1,7 @@ +# Burn Standard Library + +`burn-std` provides the core types and utilities shared across the Burn ecosystem. +It includes foundational definitions for shapes, indexing, and data types. + +This crate supports both `std` and `no_std` environments and must compile with +`cargo build --no-default-features` as well. diff --git a/crates/burn-std/src/config/autodiff.rs b/crates/burn-std/src/config/autodiff.rs new file mode 100644 index 0000000..1038a2f --- /dev/null +++ b/crates/burn-std/src/config/autodiff.rs @@ -0,0 +1,43 @@ +use cubecl_common::config::logger::{LogLevel, LoggerConfig}; + +/// Configuration for autodiff in Burn. +#[derive(Default, Clone, Debug, serde::Serialize, serde::Deserialize)] +pub struct AutodiffConfig { + /// Logger configuration for autodiff logs. + #[serde(default)] + pub logger: LoggerConfig, +} + +/// Log levels for autodiff logging. +#[derive( + Default, + Clone, + Copy, + Debug, + PartialEq, + Eq, + PartialOrd, + Ord, + serde::Serialize, + serde::Deserialize, +)] +pub enum AutodiffLogLevel { + /// Autodiff logging is disabled. + #[default] + #[serde(rename = "disabled")] + Disabled, + + /// Log backward graph size and the checkpoint strategy applied per forward pass. + #[serde(rename = "basic")] + Basic, + + /// Additionally log each tensor that gets checkpointed or recomputed. + #[serde(rename = "medium")] + Medium, + + /// Log every graph node traversal and recomputation event. + #[serde(rename = "full")] + Full, +} + +impl LogLevel for AutodiffLogLevel {} diff --git a/crates/burn-std/src/config/base.rs b/crates/burn-std/src/config/base.rs new file mode 100644 index 0000000..b668e69 --- /dev/null +++ b/crates/burn-std/src/config/base.rs @@ -0,0 +1,104 @@ +use cubecl_common::config::RuntimeConfig; +use cubecl_common::stub::Arc; + +use super::autodiff::AutodiffConfig; +use super::fusion::FusionConfig; +use super::remote::RemoteConfig; + +/// Static mutex holding the global Burn configuration, initialized as `None`. +static BURN_GLOBAL_CONFIG: spin::Mutex>> = spin::Mutex::new(None); + +/// Represents the global configuration for Burn. +#[derive(Default, Clone, Debug, serde::Serialize, serde::Deserialize)] +pub struct BurnConfig { + /// Configuration for operation fusion. + #[serde(default)] + fusion: FusionConfig, + + /// Configuration for autodiff. + #[serde(default)] + autodiff: AutodiffConfig, + + /// Configuration for the remote backend. + #[serde(default)] + remote: RemoteConfig, +} + +impl BurnConfig { + /// Returns a reference to the operation-fusion configuration. + pub fn fusion(&self) -> &FusionConfig { + &self.fusion + } + + /// Returns a reference to the autodiff configuration. + pub fn autodiff(&self) -> &AutodiffConfig { + &self.autodiff + } + + /// Returns a reference to the remote-backend configuration. + pub fn remote(&self) -> &RemoteConfig { + &self.remote + } +} + +impl RuntimeConfig for BurnConfig { + fn storage() -> &'static spin::Mutex>> { + &BURN_GLOBAL_CONFIG + } + + fn file_names() -> &'static [&'static str] { + &["burn.toml", "Burn.toml"] + } + + // Match cubecl-common's `std_io` cfg: only available on platforms where + // the trait method exists. See cubecl-common's build.rs. + #[cfg(all( + feature = "std", + any( + target_os = "windows", + target_os = "linux", + target_os = "macos", + target_os = "android" + ) + ))] + fn override_from_env(mut self) -> Self { + use super::fusion::FusionLogLevel; + use super::remote::RemoteLogLevel; + + if let Ok(val) = std::env::var("BURN_FUSION_LOG") { + let level = match val.to_ascii_lowercase().as_str() { + "disabled" | "off" | "0" => FusionLogLevel::Disabled, + "basic" => FusionLogLevel::Basic, + "medium" => FusionLogLevel::Medium, + "full" | "1" => FusionLogLevel::Full, + _ => self.fusion.logger.level, + }; + self.fusion.logger.level = level; + // Default to stderr so tests can see the output via `cargo test -- --nocapture`. + if level != FusionLogLevel::Disabled { + self.fusion.logger.stderr = true; + } + } + + if let Ok(val) = std::env::var("BURN_FUSION_MAX_EXPLORATIONS") + && let Ok(n) = val.parse::() + { + self.fusion.beam_search.max_explorations = Some(n); + } + + if let Ok(val) = std::env::var("BURN_REMOTE_LOG") { + let level = match val.to_ascii_lowercase().as_str() { + "disabled" | "off" | "0" => RemoteLogLevel::Disabled, + "basic" | "1" => RemoteLogLevel::Basic, + "full" | "2" => RemoteLogLevel::Full, + _ => self.remote.logger.level, + }; + self.remote.logger.level = level; + if level != RemoteLogLevel::Disabled { + self.remote.logger.stderr = true; + } + } + + self + } +} diff --git a/crates/burn-std/src/config/fusion.rs b/crates/burn-std/src/config/fusion.rs new file mode 100644 index 0000000..52fd606 --- /dev/null +++ b/crates/burn-std/src/config/fusion.rs @@ -0,0 +1,117 @@ +use cubecl_common::config::logger::{LogLevel, LoggerConfig}; + +/// Configuration for operation fusion in Burn. +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +pub struct FusionConfig { + /// Logger configuration for fusion logs. + #[serde(default)] + pub logger: LoggerConfig, + + /// Beam search configuration used when exploring fusion opportunities. + #[serde(default)] + pub beam_search: BeamSearchConfig, + + /// Maximum number of operations in a single client-cached graph (router graph caching). + /// + /// The greedy graph fuser keeps accumulating ops into one cached graph; once it reaches this + /// many ops the graph is closed and dispatched. `None` (the default) sets no size cap, leaving + /// [`growth_patience`](Self::growth_patience) to decide when a graph stops being worth growing; + /// set a value to also bound the per-graph memory and the cost of building and replaying any + /// single graph. + #[serde(default)] + pub max_graph_size: Option, + + /// Close the current graph once its fusion score hasn't reached a new maximum for this many + /// consecutive ops (router graph caching). + /// + /// The fuser scores the accumulated ops as they grow and tracks the best score so far; once + /// this many ops have been added without beating it, the graph has stopped getting more worth + /// caching and is closed. Higher values keep growing the graph longer in search of a better + /// score. + #[serde(default = "default_growth_patience")] + pub growth_patience: usize, +} + +impl Default for FusionConfig { + fn default() -> Self { + Self { + logger: LoggerConfig::default(), + beam_search: BeamSearchConfig::default(), + max_graph_size: None, + growth_patience: default_growth_patience(), + } + } +} + +fn default_growth_patience() -> usize { + 32 +} + +/// Beam search configuration controlling how the fusion optimizer explores independent blocks +/// of operations. +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +pub struct BeamSearchConfig { + /// Maximum number of independent blocks explored during the fusion search. + /// + /// Higher values can find better fusion opportunities at the cost of more cache misses + /// in the fusion cache. + #[serde(default = "default_max_blocks")] + pub max_blocks: usize, + + /// Stop exploring new optimizations after this many explorations (per stream). + /// + /// Each cache miss runs the (relatively expensive) optimizer to build a new optimization. A + /// graph whose relative form changes every step never caches, so it re-explores forever. Once + /// this cap is reached, cache-missing segments execute *unfused* (no optimizer work, nothing + /// added to the cache) instead. Cache *hits* are unaffected, so already-cached stable graphs + /// keep replaying. `None` (the default) never disables exploration. + #[serde(default)] + pub max_explorations: Option, +} + +impl Default for BeamSearchConfig { + fn default() -> Self { + Self { + max_blocks: default_max_blocks(), + max_explorations: None, + } + } +} + +fn default_max_blocks() -> usize { + 5 +} + +/// Log levels for fusion logging. +#[derive( + Default, + Clone, + Copy, + Debug, + PartialEq, + Eq, + PartialOrd, + Ord, + serde::Serialize, + serde::Deserialize, +)] +pub enum FusionLogLevel { + /// Fusion logging is disabled. + #[default] + #[serde(rename = "disabled")] + Disabled, + + /// Log the final execution strategy selected per stream (single vs composed). + #[serde(rename = "basic")] + Basic, + + /// Log block merge/split decisions and cache hit/miss events. + #[serde(rename = "medium")] + Medium, + + /// Log every registration, rejection and scoring decision. + #[serde(rename = "full")] + Full, +} + +impl LogLevel for FusionLogLevel {} diff --git a/crates/burn-std/src/config/logger.rs b/crates/burn-std/src/config/logger.rs new file mode 100644 index 0000000..11a5884 --- /dev/null +++ b/crates/burn-std/src/config/logger.rs @@ -0,0 +1,186 @@ +use alloc::{string::String, vec::Vec}; +use core::fmt::Display; +use cubecl_common::{ + config::{ + RuntimeConfig, + logger::{LogLevel, LoggerConfig, LoggerSinks}, + }, + stub::Arc, +}; + +use super::{ + autodiff::AutodiffLogLevel, base::BurnConfig, fusion::FusionLogLevel, remote::RemoteLogLevel, +}; + +static BURN_LOGGER: spin::Mutex> = spin::Mutex::new(None); + +#[cfg(feature = "std")] +std::thread_local! { + static LOCAL_CONFIG: std::cell::OnceCell> = + const { std::cell::OnceCell::new() }; +} + +/// Returns the current [`BurnConfig`], cached in thread-local storage on native targets. +/// +/// On the first call from a given thread this fetches the global config via +/// [`BurnConfig::get`] (which locks a spin mutex) and caches the `Arc` thread-locally. +/// Subsequent calls on the same thread only pay an `Arc` clone. On `no_std` builds this +/// is equivalent to [`BurnConfig::get`]. +/// +/// Safe because [`BurnConfig::set`] panics after the first read, so the cached snapshot +/// matches the global singleton for the whole program lifetime. +pub fn config() -> Arc { + #[cfg(feature = "std")] + { + LOCAL_CONFIG.with(|cell| cell.get_or_init(BurnConfig::get).clone()) + } + #[cfg(not(feature = "std"))] + { + BurnConfig::get() + } +} + +/// Central logging utility for Burn, managing one sink registry shared across subsystems. +#[derive(Debug)] +pub struct Logger { + sinks: LoggerSinks, + fusion_index: Vec, + autodiff_index: Vec, + remote_index: Vec, + /// The configuration snapshot the logger was initialized with. + pub config: Arc, +} + +impl Default for Logger { + fn default() -> Self { + Self::new() + } +} + +impl Logger { + /// Creates a new `Logger` from the current global `BurnConfig`. + /// + /// Note that creating a logger is somewhat expensive because it opens file handles for any + /// sink configured with a file path. + pub fn new() -> Self { + let config = BurnConfig::get(); + let mut sinks = LoggerSinks::new(); + + let fusion_index = register_enabled( + &mut sinks, + &config.fusion().logger, + config.fusion().logger.level != FusionLogLevel::Disabled, + ); + let autodiff_index = register_enabled( + &mut sinks, + &config.autodiff().logger, + config.autodiff().logger.level != AutodiffLogLevel::Disabled, + ); + let remote_index = register_enabled( + &mut sinks, + &config.remote().logger, + config.remote().logger.level != RemoteLogLevel::Disabled, + ); + + Self { + sinks, + fusion_index, + autodiff_index, + remote_index, + config, + } + } + + /// Writes `msg` to all configured fusion sinks. + pub fn log_fusion(&mut self, msg: &S) { + self.sinks.log(&self.fusion_index, msg); + } + + /// Writes `msg` to all configured autodiff sinks. + pub fn log_autodiff(&mut self, msg: &S) { + self.sinks.log(&self.autodiff_index, msg); + } + + /// Writes `msg` to all configured remote-backend sinks. + pub fn log_remote(&mut self, msg: &S) { + self.sinks.log(&self.remote_index, msg); + } + + /// Returns the current fusion log level. + pub fn log_level_fusion(&self) -> FusionLogLevel { + self.config.fusion().logger.level + } + + /// Returns the current remote-backend log level. + pub fn log_level_remote(&self) -> RemoteLogLevel { + self.config.remote().logger.level + } + + /// Returns the current autodiff log level. + pub fn log_level_autodiff(&self) -> AutodiffLogLevel { + self.config.autodiff().logger.level + } +} + +fn register_enabled( + sinks: &mut LoggerSinks, + config: &LoggerConfig, + enabled: bool, +) -> Vec { + if enabled { + sinks.register(config) + } else { + Vec::new() + } +} + +/// Emit a fusion log message when the configured level is at least `level`. +/// +/// The message is only constructed when logging is enabled. +pub fn log_fusion(level: FusionLogLevel, f: F) +where + F: FnOnce() -> String, +{ + let current = config().fusion().logger.level; + if current < level { + return; + } + let msg = f(); + let mut guard = BURN_LOGGER.lock(); + let logger = guard.get_or_insert_with(Logger::new); + logger.log_fusion(&msg); +} + +/// Emit an autodiff log message when the configured level is at least `level`. +/// +/// The message is only constructed when logging is enabled. +pub fn log_autodiff(level: AutodiffLogLevel, f: F) +where + F: FnOnce() -> String, +{ + let current = config().autodiff().logger.level; + if current < level { + return; + } + let msg = f(); + let mut guard = BURN_LOGGER.lock(); + let logger = guard.get_or_insert_with(Logger::new); + logger.log_autodiff(&msg); +} + +/// Emit a remote-backend log message when the configured level is at least `level`. +/// +/// The message is only constructed when logging is enabled. +pub fn log_remote(level: RemoteLogLevel, f: F) +where + F: FnOnce() -> String, +{ + let current = config().remote().logger.level; + if current < level { + return; + } + let msg = f(); + let mut guard = BURN_LOGGER.lock(); + let logger = guard.get_or_insert_with(Logger::new); + logger.log_remote(&msg); +} diff --git a/crates/burn-std/src/config/mod.rs b/crates/burn-std/src/config/mod.rs new file mode 100644 index 0000000..478be3d --- /dev/null +++ b/crates/burn-std/src/config/mod.rs @@ -0,0 +1,14 @@ +/// Autodiff config module. +pub mod autodiff; +/// Fusion config module. +pub mod fusion; +/// Remote backend config module. +pub mod remote; + +mod base; +mod logger; + +pub use base::*; +pub use cubecl_common::config::RuntimeConfig; +pub use cubecl_common::config::logger::{LogCrateLevel, LogLevel, LoggerConfig, LoggerSinks}; +pub use logger::*; diff --git a/crates/burn-std/src/config/remote.rs b/crates/burn-std/src/config/remote.rs new file mode 100644 index 0000000..90546be --- /dev/null +++ b/crates/burn-std/src/config/remote.rs @@ -0,0 +1,75 @@ +use cubecl_common::config::logger::{LogLevel, LoggerConfig}; + +/// Configuration for the remote backend. +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +pub struct RemoteConfig { + /// Logger configuration for remote-backend logs (e.g. the bytes-saved metric of the + /// fusion / op-graph caching feature). + #[serde(default)] + pub logger: LoggerConfig, + + /// Flush the outgoing task buffer once this many tasks have accumulated. + /// + /// Wire-level batching only — every task keeps its own stream/request id, so the server sees + /// the same per-task semantics. Larger values batch more aggressively (fewer, bigger frames); + /// smaller values cut latency for chains of fire-and-forget submits. + #[serde(default = "default_flush_threshold")] + pub flush_threshold: usize, + + /// Flush once this many bytes of buffered tensor data accumulate, independent of + /// [`flush_threshold`](Self::flush_threshold). + /// + /// Bounds how much tensor data sits unsent so large uploads go out promptly while small ops keep + /// batching. Larger batches fewer/bigger frames; smaller cuts latency for data-heavy streams. + #[serde(default = "default_flush_bytes_threshold")] + pub flush_bytes_threshold: usize, +} + +impl Default for RemoteConfig { + fn default() -> Self { + Self { + logger: LoggerConfig::default(), + flush_threshold: default_flush_threshold(), + flush_bytes_threshold: default_flush_bytes_threshold(), + } + } +} + +fn default_flush_threshold() -> usize { + 4 +} + +fn default_flush_bytes_threshold() -> usize { + 1024 * 1024 +} + +/// Log levels for remote-backend logging. +#[derive( + Default, + Clone, + Copy, + Debug, + PartialEq, + Eq, + PartialOrd, + Ord, + serde::Serialize, + serde::Deserialize, +)] +pub enum RemoteLogLevel { + /// Remote logging is disabled. + #[default] + #[serde(rename = "disabled")] + Disabled, + + /// Log periodic summaries — e.g. the cumulative network bytes saved by op-graph caching, + /// emitted once per additional mebibyte saved. + #[serde(rename = "basic")] + Basic, + + /// Log every optimization registration and replay, with per-message sizes. + #[serde(rename = "full")] + Full, +} + +impl LogLevel for RemoteLogLevel {} diff --git a/crates/burn-std/src/data/compare.rs b/crates/burn-std/src/data/compare.rs new file mode 100644 index 0000000..090f491 --- /dev/null +++ b/crates/burn-std/src/data/compare.rs @@ -0,0 +1,428 @@ +use alloc::format; +use alloc::string::String; +use num_traits::{Float, ToPrimitive}; + +use super::TensorData; +use crate::{BoolStore, DType, Element, ElementOrdered, bf16, f16}; + +/// The tolerance used to compare to floating point numbers. +/// +/// Generally, two numbers `x` and `y` are approximately equal if +/// +/// ```text +/// |x - y| < max(R * (|x + y|), A) +/// ``` +/// +/// where `R` is the relative tolerance and `A` is the absolute tolerance. +/// +/// +/// The most common way to initialize this struct is to use `Tolerance::::default()`. +/// In that case, the relative and absolute tolerances are computed using an heuristic based +/// on the EPSILON and MIN_POSITIVE values of the given floating point type `F`. +/// +/// Another common initialization is `Tolerance::::rel_abs(1e-4, 1e-5).set_half_precision_relative(1e-2)`. +/// This will use a sane default to manage values too close to 0.0 and +/// use different relative tolerances depending on the floating point precision. +#[derive(Debug, Clone, Copy)] +pub struct Tolerance { + relative: F, + absolute: F, +} + +impl Default for Tolerance { + fn default() -> Self { + Self::balanced() + } +} + +impl Tolerance { + /// Create a tolerance with strict precision setting. + pub fn strict() -> Self { + Self { + relative: F::from(0.00).unwrap(), + absolute: F::from(64).unwrap() * F::min_positive_value(), + } + } + /// Create a tolerance with balanced precision setting. + pub fn balanced() -> Self { + Self { + relative: F::from(0.005).unwrap(), // 0.5% + absolute: F::from(1e-5).unwrap(), + } + } + + /// Create a tolerance with permissive precision setting. + pub fn permissive() -> Self { + Self { + relative: F::from(0.01).unwrap(), // 1.0% + absolute: F::from(0.01).unwrap(), + } + } + /// When comparing two numbers, this uses both the relative and absolute differences. + /// + /// That is, `x` and `y` are approximately equal if + /// + /// ```text + /// |x - y| < max(R * (|x + y|), A) + /// ``` + /// + /// where `R` is the `relative` tolerance and `A` is the `absolute` tolerance. + pub fn rel_abs(relative: FF, absolute: FF) -> Self { + let relative = Self::check_relative(relative); + let absolute = Self::check_absolute(absolute); + + Self { relative, absolute } + } + + /// When comparing two numbers, this uses only the relative difference. + /// + /// That is, `x` and `y` are approximately equal if + /// + /// ```text + /// |x - y| < R * max(|x|, |y|) + /// ``` + /// + /// where `R` is the relative `tolerance`. + pub fn relative(tolerance: FF) -> Self { + let relative = Self::check_relative(tolerance); + + Self { + relative, + absolute: F::from(0.0).unwrap(), + } + } + + /// When comparing two numbers, this uses only the absolute difference. + /// + /// That is, `x` and `y` are approximately equal if + /// + /// ```text + /// |x - y| < A + /// ``` + /// + /// where `A` is the absolute `tolerance`. + pub fn absolute(tolerance: FF) -> Self { + let absolute = Self::check_absolute(tolerance); + + Self { + relative: F::from(0.0).unwrap(), + absolute, + } + } + + /// Change the relative tolerance to the given one. + pub fn set_relative(mut self, tolerance: FF) -> Self { + self.relative = Self::check_relative(tolerance); + self + } + + /// Change the relative tolerance to the given one only if `F` is half precision. + pub fn set_half_precision_relative(mut self, tolerance: FF) -> Self { + if core::mem::size_of::() == 2 { + self.relative = Self::check_relative(tolerance); + } + self + } + + /// Change the relative tolerance to the given one only if `F` is single precision. + pub fn set_single_precision_relative(mut self, tolerance: FF) -> Self { + if core::mem::size_of::() == 4 { + self.relative = Self::check_relative(tolerance); + } + self + } + + /// Change the relative tolerance to the given one only if `F` is double precision. + pub fn set_double_precision_relative(mut self, tolerance: FF) -> Self { + if core::mem::size_of::() == 8 { + self.relative = Self::check_relative(tolerance); + } + self + } + + /// Change the absolute tolerance to the given one. + pub fn set_absolute(mut self, tolerance: FF) -> Self { + self.absolute = Self::check_absolute(tolerance); + self + } + + /// Change the absolute tolerance to the given one only if `F` is half precision. + pub fn set_half_precision_absolute(mut self, tolerance: FF) -> Self { + if core::mem::size_of::() == 2 { + self.absolute = Self::check_absolute(tolerance); + } + self + } + + /// Change the absolute tolerance to the given one only if `F` is single precision. + pub fn set_single_precision_absolute(mut self, tolerance: FF) -> Self { + if core::mem::size_of::() == 4 { + self.absolute = Self::check_absolute(tolerance); + } + self + } + + /// Change the absolute tolerance to the given one only if `F` is double precision. + pub fn set_double_precision_absolute(mut self, tolerance: FF) -> Self { + if core::mem::size_of::() == 8 { + self.absolute = Self::check_absolute(tolerance); + } + self + } + + /// Checks if `x` and `y` are approximately equal given the tolerance. + pub fn approx_eq(&self, x: F, y: F) -> bool { + // See the accepted answer here + // https://stackoverflow.com/questions/4915462/how-should-i-do-floating-point-comparison + + // This also handles the case where both a and b are infinity so that we don't need + // to manage it in the rest of the function. + if x == y { + return true; + } + + let diff = (x - y).abs(); + let max = F::max(x.abs(), y.abs()); + + diff < self.absolute.max(self.relative * max) + } + + fn check_relative(tolerance: FF) -> F { + let tolerance = F::from(tolerance).unwrap(); + assert!(tolerance <= F::one()); + tolerance + } + + fn check_absolute(tolerance: FF) -> F { + let tolerance = F::from(tolerance).unwrap(); + assert!(tolerance >= F::zero()); + tolerance + } +} + +impl TensorData { + /// Asserts the data is equal to another data. + /// + /// # Arguments + /// + /// * `other` - The other data. + /// * `strict` - If true, the data types must the be same. + /// Otherwise, the comparison is done in the current data type. + /// + /// # Panics + /// + /// Panics if the data is not equal. + #[track_caller] + pub fn assert_eq(&self, other: &Self, strict: bool) { + if strict { + assert_eq!( + self.dtype, other.dtype, + "Data types differ ({:?} != {:?})", + self.dtype, other.dtype + ); + } + + match self.dtype { + DType::F64 => self.assert_eq_elem::(other), + DType::F32 | DType::Flex32 => self.assert_eq_elem::(other), + DType::F16 => self.assert_eq_elem::(other), + DType::BF16 => self.assert_eq_elem::(other), + DType::I64 => self.assert_eq_elem::(other), + DType::I32 => self.assert_eq_elem::(other), + DType::I16 => self.assert_eq_elem::(other), + DType::I8 => self.assert_eq_elem::(other), + DType::U64 => self.assert_eq_elem::(other), + DType::U32 => self.assert_eq_elem::(other), + DType::U16 => self.assert_eq_elem::(other), + DType::U8 => self.assert_eq_elem::(other), + DType::Bool(BoolStore::Native) => self.assert_eq_elem::(other), + DType::Bool(BoolStore::U8) => self.assert_eq_elem::(other), + DType::Bool(BoolStore::U32) => self.assert_eq_elem::(other), + DType::QFloat(q) => { + // Strict or not, it doesn't make sense to compare quantized data to not quantized data for equality + let q_other = if let DType::QFloat(q_other) = other.dtype { + q_other + } else { + panic!("Quantized data differs from other not quantized data") + }; + + // Data equality mostly depends on input quantization type, but we also check level + if q.value == q_other.value && q.level == q_other.level { + self.assert_eq_elem::(other) + } else { + panic!("Quantization schemes differ ({q:?} != {q_other:?})") + } + } + } + } + + #[track_caller] + fn assert_eq_elem(&self, other: &Self) { + let mut message = String::new(); + if self.shape != other.shape { + message += format!( + "\n => Shape is different: {:?} != {:?}", + self.shape, other.shape + ) + .as_str(); + } + + let mut num_diff = 0; + let max_num_diff = 5; + for (i, (a, b)) in self.iter::().zip(other.iter::()).enumerate() { + if !a.eq(&b) { + // Only print the first 5 different values. + if num_diff < max_num_diff { + message += format!("\n => Position {i}: {a} != {b}").as_str(); + } + num_diff += 1; + } + } + + if num_diff >= max_num_diff { + message += format!("\n{} more errors...", num_diff - max_num_diff).as_str(); + } + + if !message.is_empty() { + panic!("Tensors are not eq:{message}"); + } + } + + /// Asserts the data is approximately equal to another data. + /// + /// # Arguments + /// + /// * `other` - The other data. + /// * `tolerance` - The tolerance of the comparison. + /// + /// # Panics + /// + /// Panics if the data is not approximately equal. + #[track_caller] + pub fn assert_approx_eq(&self, other: &Self, tolerance: Tolerance) { + let mut message = String::new(); + if self.shape != other.shape { + message += format!( + "\n => Shape is different: {:?} != {:?}", + self.shape, other.shape + ) + .as_str(); + } + + let iter = self.iter::().zip(other.iter::()); + + let mut num_diff = 0; + let max_num_diff = 5; + + for (i, (a, b)) in iter.enumerate() { + //if they are both nan, then they are equally nan + let both_nan = a.is_nan() && b.is_nan(); + //this works for both infinities + let both_inf = + a.is_infinite() && b.is_infinite() && ((a > F::zero()) == (b > F::zero())); + + if both_nan || both_inf { + continue; + } + + if !tolerance.approx_eq(F::from(a).unwrap(), F::from(b).unwrap()) { + // Only print the first 5 different values. + if num_diff < max_num_diff { + let diff_abs = ToPrimitive::to_f64(&(a - b).abs()).unwrap(); + let max = F::max(a.abs(), b.abs()); + let diff_rel = diff_abs / ToPrimitive::to_f64(&max).unwrap(); + + let tol_rel = ToPrimitive::to_f64(&tolerance.relative).unwrap(); + let tol_abs = ToPrimitive::to_f64(&tolerance.absolute).unwrap(); + + message += format!( + "\n => Position {i}: {a} != {b}\n diff (rel = {diff_rel:+.2e}, abs = {diff_abs:+.2e}), tol (rel = {tol_rel:+.2e}, abs = {tol_abs:+.2e})" + ) + .as_str(); + } + num_diff += 1; + } + } + + if num_diff >= max_num_diff { + message += format!("\n{} more errors...", num_diff - 5).as_str(); + } + + if !message.is_empty() { + panic!("Tensors are not approx eq:{message}"); + } + } + + /// Asserts each value is within a given range. + /// + /// # Arguments + /// + /// * `range` - The range. + /// + /// # Panics + /// + /// If any value is not within the half-open range bounded inclusively below + /// and exclusively above (`start..end`). + pub fn assert_within_range(&self, range: core::ops::Range) { + for elem in self.iter::() { + if elem.cmp(&range.start).is_lt() || elem.cmp(&range.end).is_ge() { + panic!("Element ({elem:?}) is not within range {range:?}"); + } + } + } + + /// Asserts each value is within a given inclusive range. + /// + /// # Arguments + /// + /// * `range` - The range. + /// + /// # Panics + /// + /// If any value is not within the half-open range bounded inclusively (`start..=end`). + pub fn assert_within_range_inclusive( + &self, + range: core::ops::RangeInclusive, + ) { + let start = range.start(); + let end = range.end(); + + for elem in self.iter::() { + if elem.cmp(start).is_lt() || elem.cmp(end).is_gt() { + panic!("Element ({elem:?}) is not within range {range:?}"); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn should_assert_appox_eq_limit() { + let data1 = TensorData::from([[3.0, 5.0, 6.0]]); + let data2 = TensorData::from([[3.03, 5.0, 6.0]]); + + data1.assert_approx_eq::(&data2, Tolerance::absolute(3e-2)); + data1.assert_approx_eq::(&data2, Tolerance::absolute(3e-2)); + } + + #[test] + #[should_panic] + fn should_assert_approx_eq_above_limit() { + let data1 = TensorData::from([[3.0, 5.0, 6.0]]); + let data2 = TensorData::from([[3.031, 5.0, 6.0]]); + + data1.assert_approx_eq::(&data2, Tolerance::absolute(1e-2)); + } + + #[test] + #[should_panic] + fn should_assert_approx_eq_check_shape() { + let data1 = TensorData::from([[3.0, 5.0, 6.0, 7.0]]); + let data2 = TensorData::from([[3.0, 5.0, 6.0]]); + + data1.assert_approx_eq::(&data2, Tolerance::absolute(1e-2)); + } +} diff --git a/crates/burn-std/src/data/mod.rs b/crates/burn-std/src/data/mod.rs new file mode 100644 index 0000000..cf5d2dc --- /dev/null +++ b/crates/burn-std/src/data/mod.rs @@ -0,0 +1,5 @@ +mod compare; +mod tensor; + +pub use compare::*; +pub use tensor::*; diff --git a/crates/burn-std/src/data/tensor.rs b/crates/burn-std/src/data/tensor.rs new file mode 100644 index 0000000..2aa3b66 --- /dev/null +++ b/crates/burn-std/src/data/tensor.rs @@ -0,0 +1,936 @@ +use core::f32; + +use alloc::boxed::Box; +use alloc::format; +use alloc::string::String; +use alloc::vec::Vec; +use bytemuck::{AnyBitPattern, CheckedBitPattern, Zeroable, cast_mut, checked::CheckedCastError}; +use rand::Rng; +use thiserror::Error; + +use crate::Scalar; +use crate::distribution::Distribution; +use crate::element::{Element, ElementConversion}; +use crate::tensor::DType; +use crate::{ + BoolStore, Bytes, QuantLevel, QuantMode, QuantScheme, QuantValue, QuantizedBytes, Shape, bf16, + f16, +}; + +use serde::{Deserialize, Serialize}; + +/// Data structure for tensors. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct TensorData { + /// The values of the tensor (as bytes). + pub bytes: Bytes, + + /// The shape of the tensor. + #[serde(with = "shape_inner")] + pub shape: Shape, + + /// The data type of the tensor. + pub dtype: DType, +} + +// For backward compatibility with shape `Vec` +mod shape_inner { + use crate::SmallVec; + + use super::*; + + pub fn serialize( + shape: &Shape, + serializer: S, + ) -> Result { + shape.as_slice().serialize(serializer) + } + + pub fn deserialize<'de, D: serde::Deserializer<'de>>( + deserializer: D, + ) -> Result { + let dims = SmallVec::<[usize; _]>::deserialize(deserializer)?; + Ok(Shape::new_raw(dims)) + } +} + +impl TensorData { + /// Creates a new tensor data structure. + pub fn new>(value: Vec, shape: S) -> Self { + // Ensure shape is valid + let shape = shape.into(); + Self::check_data_len(&value, &shape); + + Self { + bytes: Bytes::from_elems(value), + shape, + dtype: E::dtype(), + } + } + + /// Creates a new quantized tensor data structure. + pub fn quantized>( + value: Vec, + shape: S, + scheme: QuantScheme, + qparams: &[f32], + ) -> Self { + let shape = shape.into(); + Self::check_data_len(&value, &shape); + + let q_bytes = QuantizedBytes::new(value, scheme, qparams); + + Self { + bytes: q_bytes.bytes, + shape, + dtype: DType::QFloat(q_bytes.scheme), + } + } + + /// Creates a new tensor data structure from raw bytes. + pub fn from_bytes>(bytes: Bytes, shape: S, dtype: DType) -> Self { + Self { + bytes, + shape: shape.into(), + dtype, + } + } + + /// Creates a new tensor data structure from raw bytes stored in a vector. + /// + /// Prefer [`TensorData::new`] or [`TensorData::quantized`] over this method unless you are + /// certain that the bytes representation is valid. + pub fn from_bytes_vec>(bytes: Vec, shape: S, dtype: DType) -> Self { + Self { + bytes: Bytes::from_bytes_vec(bytes), + shape: shape.into(), + dtype, + } + } + + // Check that the input vector contains a correct number of elements + fn check_data_len(data: &[E], shape: &Shape) { + let expected_data_len = Self::numel(shape); + let num_data = data.len(); + assert_eq!( + expected_data_len, num_data, + "Shape {shape:?} is invalid for input of size {num_data:?}", + ); + } + + /// Returns the immutable slice view of the tensor data. + pub fn as_slice(&self) -> Result<&[E], DataError> { + if self.matches_target_dtype::() { + match E::dtype() { + // The only way to create a bool `TensorData` with invalid values is by unsafely modifying + // the dtype. This should be considered unsafe to begin with, so we unsafely cast bool + // to u8 to skip bit validation. Validation iterates through the entire vector, so it's slow. + DType::Bool(BoolStore::Native) => { + let slice = bytemuck::checked::try_cast_slice::<_, u8>(&self.bytes) + .map_err(DataError::CastError)?; + Ok(unsafe { core::mem::transmute::<&[u8], &[E]>(slice) }) + } + _ => bytemuck::checked::try_cast_slice(&self.bytes).map_err(DataError::CastError), + } + } else { + Err(DataError::TypeMismatch(format!( + "Invalid target element type (expected {:?}, got {:?})", + self.dtype, + E::dtype() + ))) + } + } + + /// Returns the mutable slice view of the tensor data. + /// + /// # Panics + /// If the target element type is different from the stored element type. + pub fn as_mut_slice(&mut self) -> Result<&mut [E], DataError> { + if self.matches_target_dtype::() { + match E::dtype() { + // The only way to create a bool `TensorData` with invalid values is by unsafely modifying + // the dtype. This should be considered unsafe to begin with, so we unsafely cast bool + // to u8 to skip bit validation. Validation iterates through the entire vector, so it's slow. + DType::Bool(BoolStore::Native) => { + let slice = bytemuck::checked::try_cast_slice_mut::<_, u8>(&mut self.bytes) + .map_err(DataError::CastError)?; + Ok(unsafe { core::mem::transmute::<&mut [u8], &mut [E]>(slice) }) + } + _ => bytemuck::checked::try_cast_slice_mut(&mut self.bytes) + .map_err(DataError::CastError), + } + } else { + Err(DataError::TypeMismatch(format!( + "Invalid target element type (expected {:?}, got {:?})", + self.dtype, + E::dtype() + ))) + } + } + + /// Returns the tensor data as a vector of scalar values. + pub fn to_vec(&self) -> Result, DataError> { + Ok(self.as_slice()?.to_vec()) + } + + /// Returns the tensor data as a vector of scalar values. + pub fn into_vec(self) -> Result, DataError> { + // This means we cannot call `into_vec` for QFloat + if !self.matches_target_dtype::() { + return Err(DataError::TypeMismatch(format!( + "Invalid target element type (expected {:?}, got {:?})", + self.dtype, + E::dtype() + ))); + } + + match E::dtype() { + // The only way to create a bool `TensorData` with invalid values is by unsafely modifying + // the dtype. This should be considered unsafe to begin with, so we unsafely cast bool + // to u8 to skip bit validation. Validation iterates through the entire vector, so it's slow. + DType::Bool(BoolStore::Native) => { + let vec = self.into_vec_unchecked::()?; + Ok(unsafe { core::mem::transmute::, Vec>(vec) }) + } + _ => self.into_vec_unchecked(), + } + } + + /// Returns the tensor data as a vector of scalar values. Does not check dtype. + fn into_vec_unchecked(self) -> Result, DataError> { + let mut me = self; + me.bytes = match me.bytes.try_into_vec::() { + Ok(elems) => return Ok(elems), + Err(bytes) => bytes, + }; + + // The bytes might have been deserialized and allocated with a different align. + // In that case, we have to memcopy the data into a new vector, more suitably allocated + Ok(bytemuck::checked::try_cast_slice(me.as_bytes()) + .map_err(DataError::CastError)? + .to_vec()) + } + + fn matches_target_dtype(&self) -> bool { + let target_dtype = E::dtype(); + match self.dtype { + DType::Bool(BoolStore::U8) => { + matches!(target_dtype, DType::U8 | DType::Bool(BoolStore::U8)) + } + DType::Bool(BoolStore::U32) => { + matches!(target_dtype, DType::U32 | DType::Bool(BoolStore::U32)) + } + dtype => dtype == target_dtype, + } + } + + /// Returns an iterator over the values of the tensor data. + pub fn iter(&self) -> Box + '_> { + if E::dtype() == self.dtype { + Box::new(bytemuck::checked::cast_slice(&self.bytes).iter().copied()) + } else { + match self.dtype { + DType::I8 => Box::new( + bytemuck::checked::cast_slice(&self.bytes) + .iter() + .map(|e: &i8| e.elem::()), + ), + DType::I16 => Box::new( + bytemuck::checked::cast_slice(&self.bytes) + .iter() + .map(|e: &i16| e.elem::()), + ), + DType::I32 => Box::new( + bytemuck::checked::cast_slice(&self.bytes) + .iter() + .map(|e: &i32| e.elem::()), + ), + DType::I64 => Box::new( + bytemuck::checked::cast_slice(&self.bytes) + .iter() + .map(|e: &i64| e.elem::()), + ), + DType::U8 => Box::new(self.bytes.iter().map(|e| e.elem::())), + DType::U16 => Box::new( + bytemuck::checked::cast_slice(&self.bytes) + .iter() + .map(|e: &u16| e.elem::()), + ), + DType::U32 => Box::new( + bytemuck::checked::cast_slice(&self.bytes) + .iter() + .map(|e: &u32| e.elem::()), + ), + DType::U64 => Box::new( + bytemuck::checked::cast_slice(&self.bytes) + .iter() + .map(|e: &u64| e.elem::()), + ), + DType::BF16 => Box::new( + bytemuck::checked::cast_slice(&self.bytes) + .iter() + .map(|e: &bf16| e.elem::()), + ), + DType::F16 => Box::new( + bytemuck::checked::cast_slice(&self.bytes) + .iter() + .map(|e: &f16| e.elem::()), + ), + DType::F32 | DType::Flex32 => Box::new( + bytemuck::checked::cast_slice(&self.bytes) + .iter() + .map(|e: &f32| e.elem::()), + ), + DType::F64 => Box::new( + bytemuck::checked::cast_slice(&self.bytes) + .iter() + .map(|e: &f64| e.elem::()), + ), + // bool is a byte value equal to either 0 or 1 + DType::Bool(BoolStore::Native) | DType::Bool(BoolStore::U8) => { + Box::new(self.bytes.iter().map(|e| e.elem::())) + } + DType::Bool(BoolStore::U32) => Box::new( + bytemuck::checked::cast_slice(&self.bytes) + .iter() + .map(|e: &u32| e.elem::()), + ), + DType::QFloat(scheme) => match scheme { + QuantScheme { + level: QuantLevel::Tensor | QuantLevel::Block(_), + mode: QuantMode::Symmetric, + value: + QuantValue::Q8F + | QuantValue::Q8S + // Represent sub-byte values as i8 + | QuantValue::Q4F + | QuantValue::Q4S + | QuantValue::Q2F + | QuantValue::Q2S, + .. + } => { + // Quantized int8 values + let q_bytes = QuantizedBytes { + bytes: self.bytes.clone(), + scheme, + num_elements: self.num_elements(), + }; + let (values, _) = q_bytes.into_vec_i8(); + + Box::new( + values + .iter() + .map(|e: &i8| e.elem::()) + .collect::>() + .into_iter(), + ) + } + QuantScheme { + level: QuantLevel::Tensor | QuantLevel::Block(_), + mode: QuantMode::Symmetric, + value: + QuantValue::E4M3 | QuantValue::E5M2 | QuantValue::E2M1, + .. + } => { + unimplemented!("Not yet implemented for iteration"); + } + }, + } + } + } + + /// Returns the rank (the number of dimensions). + pub fn rank(&self) -> usize { + self.shape.len() + } + + /// Returns the total number of elements of the tensor data. + pub fn num_elements(&self) -> usize { + Self::numel(&self.shape) + } + + fn numel(shape: &[usize]) -> usize { + shape.iter().product() + } + + /// Populates the data with random values. + pub fn random>( + shape: S, + distribution: Distribution, + rng: &mut R, + ) -> Self { + let shape = shape.into(); + let num_elements = Self::numel(&shape); + let mut data = Vec::with_capacity(num_elements); + + for _ in 0..num_elements { + data.push(E::random(distribution, rng)); + } + + TensorData::new(data, shape) + } + + /// Populates the data with zeros. + pub fn zeros>(shape: S) -> TensorData { + let shape = shape.into(); + let num_elements = Self::numel(&shape); + let mut data = Vec::::with_capacity(num_elements); + + for _ in 0..num_elements { + data.push(0.elem()); + } + + TensorData::new(data, shape) + } + + /// Populates the data with ones. + pub fn ones>(shape: S) -> TensorData { + let shape = shape.into(); + let num_elements = Self::numel(&shape); + let mut data = Vec::::with_capacity(num_elements); + + for _ in 0..num_elements { + data.push(1.elem()); + } + + TensorData::new(data, shape) + } + + /// Populates the data with the given value + pub fn full>(shape: S, fill_value: E) -> TensorData { + let shape = shape.into(); + let num_elements = Self::numel(&shape); + let mut data = Vec::::with_capacity(num_elements); + for _ in 0..num_elements { + data.push(fill_value) + } + + TensorData::new(data, shape) + } + + /// Populates the data with the given value + pub fn full_dtype, S: Into>( + shape: S, + fill_value: E, + dtype: DType, + ) -> TensorData { + let fill_value = fill_value.into(); + match dtype { + DType::F64 => Self::full::(shape, fill_value.elem()), + DType::F32 | DType::Flex32 => Self::full::(shape, fill_value.elem()), + DType::F16 => Self::full::(shape, fill_value.elem()), + DType::BF16 => Self::full::(shape, fill_value.elem()), + DType::I64 => Self::full::(shape, fill_value.elem()), + DType::I32 => Self::full::(shape, fill_value.elem()), + DType::I16 => Self::full::(shape, fill_value.elem()), + DType::I8 => Self::full::(shape, fill_value.elem()), + DType::U64 => Self::full::(shape, fill_value.elem()), + DType::U32 => Self::full::(shape, fill_value.elem()), + DType::U16 => Self::full::(shape, fill_value.elem()), + DType::U8 => Self::full::(shape, fill_value.elem()), + DType::Bool(BoolStore::Native) => Self::full::(shape, fill_value.elem()), + DType::Bool(BoolStore::U8) => { + Self::full::(shape, fill_value.elem()).into_bool_u8() + } + DType::Bool(BoolStore::U32) => { + Self::full::(shape, fill_value.elem()).into_bool_u32() + } + DType::QFloat(_) => unreachable!(), + } + } + + // Unchecked, used to overwrite the dtype + fn into_bool_u8(mut self) -> Self { + self.dtype = DType::Bool(BoolStore::U8); + self + } + + // Unchecked, used to overwrite the dtype + fn into_bool_u32(mut self) -> Self { + self.dtype = DType::Bool(BoolStore::U32); + self + } + + /// Converts the data to a different element type. + pub fn convert(self) -> Self { + self.convert_dtype(E::dtype()) + } + + /// Converts the data to a different element type. + pub fn convert_dtype(self, dtype: DType) -> Self { + if dtype == self.dtype { + self + } else if dtype.size() == self.dtype.size() + && !matches!( + self.dtype, + DType::Bool(BoolStore::Native) | DType::QFloat(_) + ) + && !matches!(dtype, DType::Bool(BoolStore::Native) | DType::QFloat(_)) + { + match self.dtype { + DType::F64 => self.convert_inplace_dtype::(dtype), + DType::F32 | DType::Flex32 => self.convert_inplace_dtype::(dtype), + DType::F16 => self.convert_inplace_dtype::(dtype), + DType::BF16 => self.convert_inplace_dtype::(dtype), + DType::I64 => self.convert_inplace_dtype::(dtype), + DType::I32 => self.convert_inplace_dtype::(dtype), + DType::I16 => self.convert_inplace_dtype::(dtype), + DType::I8 => self.convert_inplace_dtype::(dtype), + DType::U64 => self.convert_inplace_dtype::(dtype), + DType::U32 => self.convert_inplace_dtype::(dtype), + DType::U16 => self.convert_inplace_dtype::(dtype), + DType::U8 => self.convert_inplace_dtype::(dtype), + DType::Bool(BoolStore::U8) => self.convert_inplace_dtype::(dtype), + DType::Bool(BoolStore::U32) => self.convert_inplace_dtype::(dtype), + DType::Bool(BoolStore::Native) | DType::QFloat(_) => unreachable!(), + } + } else { + match self.dtype { + DType::F64 => self.convert_clone_dtype::(dtype), + DType::F32 | DType::Flex32 => self.convert_clone_dtype::(dtype), + DType::F16 => self.convert_clone_dtype::(dtype), + DType::BF16 => self.convert_clone_dtype::(dtype), + DType::I64 => self.convert_clone_dtype::(dtype), + DType::I32 => self.convert_clone_dtype::(dtype), + DType::I16 => self.convert_clone_dtype::(dtype), + DType::I8 => self.convert_clone_dtype::(dtype), + DType::U64 => self.convert_clone_dtype::(dtype), + DType::U32 => self.convert_clone_dtype::(dtype), + DType::U16 => self.convert_clone_dtype::(dtype), + DType::U8 => self.convert_clone_dtype::(dtype), + DType::Bool(BoolStore::Native) => self.convert_clone_dtype::(dtype), + DType::Bool(BoolStore::U8) => self.convert_clone_dtype::(dtype), + DType::Bool(BoolStore::U32) => self.convert_clone_dtype::(dtype), + DType::QFloat(_) => unreachable!(), + } + } + } + + fn convert_inplace_dtype(self, dtype: DType) -> Self { + match dtype { + DType::F64 => self.convert_inplace::(), + DType::F32 | DType::Flex32 => self.convert_inplace::(), + DType::F16 => self.convert_inplace::(), + DType::BF16 => self.convert_inplace::(), + DType::I64 => self.convert_inplace::(), + DType::I32 => self.convert_inplace::(), + DType::I16 => self.convert_inplace::(), + DType::I8 => self.convert_inplace::(), + DType::U64 => self.convert_inplace::(), + DType::U32 => self.convert_inplace::(), + DType::U16 => self.convert_inplace::(), + DType::U8 => self.convert_inplace::(), + DType::Bool(BoolStore::U8) => self.convert_inplace::().into_bool_u8(), + DType::Bool(BoolStore::U32) => self.convert_inplace::().into_bool_u32(), + DType::Bool(BoolStore::Native) | DType::QFloat(_) => unreachable!(), + } + } + + fn convert_inplace( + mut self, + ) -> Self { + for x in bytemuck::cast_slice_mut::<_, Current>(&mut self.bytes) { + let t: Target = x.elem(); + let x = cast_mut::<_, Target>(x); + *x = t; + } + + self.dtype = Target::dtype(); + + self + } + + fn convert_clone_dtype(self, dtype: DType) -> Self { + match dtype { + DType::F64 => self.convert_clone::(), + DType::F32 | DType::Flex32 => self.convert_clone::(), + DType::F16 => self.convert_clone::(), + DType::BF16 => self.convert_clone::(), + DType::I64 => self.convert_clone::(), + DType::I32 => self.convert_clone::(), + DType::I16 => self.convert_clone::(), + DType::I8 => self.convert_clone::(), + DType::U64 => self.convert_clone::(), + DType::U32 => self.convert_clone::(), + DType::U16 => self.convert_clone::(), + DType::U8 => self.convert_clone::(), + DType::Bool(BoolStore::Native) => self.convert_clone::(), + DType::Bool(BoolStore::U8) => self.convert_clone::().into_bool_u8(), + DType::Bool(BoolStore::U32) => self.convert_clone::().into_bool_u32(), + DType::QFloat(_) => unreachable!(), + } + } + + fn convert_clone( + self, + ) -> Self { + let this = bytemuck::checked::cast_slice::<_, Current>(&self.bytes); + let mut out: Vec = ::alloc::vec![Zeroable::zeroed(); self.num_elements()]; + + for (x, out) in this.iter().zip(&mut out) { + *out = x.elem(); + } + + Self::new(out, self.shape) + } + + /// Returns the data as a slice of bytes. + pub fn as_bytes(&self) -> &[u8] { + &self.bytes + } + + /// Returns the bytes representation of the data. + pub fn into_bytes(self) -> Bytes { + self.bytes + } +} + +impl From<[E; A]> for TensorData { + fn from(elems: [E; A]) -> Self { + TensorData::new(elems.to_vec(), [A]) + } +} + +impl From<[usize; A]> for TensorData { + fn from(elems: [usize; A]) -> Self { + TensorData::new(elems.iter().map(|&e| e as i64).collect(), [A]) + } +} + +impl From<&[usize]> for TensorData { + fn from(elems: &[usize]) -> Self { + let mut data = Vec::with_capacity(elems.len()); + for elem in elems.iter() { + data.push(*elem as i64); + } + + TensorData::new(data, [elems.len()]) + } +} + +impl From<&[E]> for TensorData { + fn from(elems: &[E]) -> Self { + let mut data = Vec::with_capacity(elems.len()); + for elem in elems.iter() { + data.push(*elem); + } + + TensorData::new(data, [elems.len()]) + } +} + +impl From<[[E; B]; A]> for TensorData { + fn from(elems: [[E; B]; A]) -> Self { + let mut data = Vec::with_capacity(A * B); + for elem in elems.into_iter().take(A) { + for elem in elem.into_iter().take(B) { + data.push(elem); + } + } + + TensorData::new(data, [A, B]) + } +} + +impl From<[[[E; C]; B]; A]> + for TensorData +{ + fn from(elems: [[[E; C]; B]; A]) -> Self { + let mut data = Vec::with_capacity(A * B * C); + + for elem in elems.into_iter().take(A) { + for elem in elem.into_iter().take(B) { + for elem in elem.into_iter().take(C) { + data.push(elem); + } + } + } + + TensorData::new(data, [A, B, C]) + } +} + +impl + From<[[[[E; D]; C]; B]; A]> for TensorData +{ + fn from(elems: [[[[E; D]; C]; B]; A]) -> Self { + let mut data = Vec::with_capacity(A * B * C * D); + + for elem in elems.into_iter().take(A) { + for elem in elem.into_iter().take(B) { + for elem in elem.into_iter().take(C) { + for elem in elem.into_iter().take(D) { + data.push(elem); + } + } + } + } + + TensorData::new(data, [A, B, C, D]) + } +} + +impl + From<[[[[[Elem; E]; D]; C]; B]; A]> for TensorData +{ + fn from(elems: [[[[[Elem; E]; D]; C]; B]; A]) -> Self { + let mut data = Vec::with_capacity(A * B * C * D * E); + + for elem in elems.into_iter().take(A) { + for elem in elem.into_iter().take(B) { + for elem in elem.into_iter().take(C) { + for elem in elem.into_iter().take(D) { + for elem in elem.into_iter().take(E) { + data.push(elem); + } + } + } + } + } + + TensorData::new(data, [A, B, C, D, E]) + } +} +impl core::fmt::Display for TensorData { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let fmt = match self.dtype { + DType::F64 => format!("{:?}", self.as_slice::().unwrap()), + DType::F32 | DType::Flex32 => format!("{:?}", self.as_slice::().unwrap()), + DType::F16 => format!("{:?}", self.as_slice::().unwrap()), + DType::BF16 => format!("{:?}", self.as_slice::().unwrap()), + DType::I64 => format!("{:?}", self.as_slice::().unwrap()), + DType::I32 => format!("{:?}", self.as_slice::().unwrap()), + DType::I16 => format!("{:?}", self.as_slice::().unwrap()), + DType::I8 => format!("{:?}", self.as_slice::().unwrap()), + DType::U64 => format!("{:?}", self.as_slice::().unwrap()), + DType::U32 => format!("{:?}", self.as_slice::().unwrap()), + DType::U16 => format!("{:?}", self.as_slice::().unwrap()), + DType::U8 => format!("{:?}", self.as_slice::().unwrap()), + DType::Bool(BoolStore::Native) => format!("{:?}", self.as_slice::().unwrap()), + DType::Bool(BoolStore::U8) => format!("{:?}", self.as_slice::().unwrap()), + DType::Bool(BoolStore::U32) => format!("{:?}", self.as_slice::().unwrap()), + DType::QFloat(scheme) => match scheme { + QuantScheme { + level: QuantLevel::Tensor | QuantLevel::Block(_), + mode: QuantMode::Symmetric, + value: + QuantValue::Q8F + | QuantValue::Q8S + // Display sub-byte values as i8 + | QuantValue::Q4F + | QuantValue::Q4S + | QuantValue::Q2F + | QuantValue::Q2S, + .. + } => { + format!("{:?} {scheme:?}", self.iter::().collect::>()) + }, + QuantScheme { + level: QuantLevel::Tensor | QuantLevel::Block(_), + mode: QuantMode::Symmetric, + value: + QuantValue::E4M3 | QuantValue::E5M2 | QuantValue::E2M1, + .. + } => { + unimplemented!("Can't format yet"); + } + }, + }; + f.write_str(fmt.as_str()) + } +} + +/// The things that can go wrong when manipulating tensor data. +#[derive(Debug, Error)] +pub enum DataError { + /// Failed to cast the values to a specified element type. + #[error("Failed to cast values to the specified element type.\nError:\n {0}")] + CastError(CheckedCastError), + /// Invalid target element type. + #[error("{0}")] + TypeMismatch(String), +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::shape; + use alloc::vec; + use rand::{ + SeedableRng, + rngs::{StdRng, SysRng}, + }; + + #[test] + fn should_have_rank() { + let shape = [3, 5, 6]; + let data = TensorData::random::( + shape, + Distribution::Default, + &mut StdRng::try_from_rng(&mut SysRng).unwrap(), + ); + + assert_eq!(data.rank(), 3); + } + + #[test] + fn into_vec_should_yield_same_value_as_iter() { + let shape = [3, 5, 6]; + let data = TensorData::random::( + shape, + Distribution::Default, + &mut StdRng::try_from_rng(&mut SysRng).unwrap(), + ); + + let expected = data.iter::().collect::>(); + let actual = data.into_vec::().unwrap(); + + assert_eq!(expected, actual); + } + + #[test] + #[should_panic] + fn into_vec_should_assert_wrong_dtype() { + let shape = [3, 5, 6]; + let data = TensorData::random::( + shape, + Distribution::Default, + &mut StdRng::try_from_rng(&mut SysRng).unwrap(), + ); + + data.into_vec::().unwrap(); + } + + #[test] + fn should_have_right_num_elements() { + let shape = [3, 5, 6]; + let num_elements: usize = shape.iter().product(); + let data = TensorData::random::( + shape, + Distribution::Default, + &mut StdRng::try_from_rng(&mut SysRng).unwrap(), + ); + + assert_eq!(num_elements, data.bytes.len() / 4); // f32 stored as u8s + assert_eq!(num_elements, data.as_slice::().unwrap().len()); + } + + #[test] + fn should_have_right_shape() { + let data = TensorData::from([[3.0, 5.0, 6.0]]); + assert_eq!(data.shape, shape![1, 3]); + + let data = TensorData::from([[4.0, 5.0, 8.0], [3.0, 5.0, 6.0]]); + assert_eq!(data.shape, shape![2, 3]); + + let data = TensorData::from([3.0, 5.0, 6.0]); + assert_eq!(data.shape, shape![3]); + } + + #[test] + fn should_convert_bytes_correctly() { + let mut vector: Vec = Vec::with_capacity(5); + vector.push(2.0); + vector.push(3.0); + let data1 = TensorData::new(vector, vec![2]); + + let factor = core::mem::size_of::() / core::mem::size_of::(); + assert_eq!(data1.bytes.len(), 2 * factor); + assert_eq!(data1.bytes.capacity(), 5 * factor); + } + + #[test] + fn should_convert_bytes_correctly_inplace() { + fn test_precision() { + let data = TensorData::new((0..32).collect(), [32]); + for (i, val) in data + .clone() + .convert::() + .into_vec::() + .unwrap() + .into_iter() + .enumerate() + { + assert_eq!(i as u32, val.elem::()) + } + } + test_precision::(); + test_precision::(); + test_precision::(); + test_precision::(); + } + + macro_rules! test_dtypes { + ($test_name:ident, $($dtype:ty),*) => { + $( + paste::paste! { + #[test] + fn [<$test_name _ $dtype:snake>]() { + let full_dtype = TensorData::full_dtype([2, 16], 4, <$dtype>::dtype()); + let full = TensorData::full::<$dtype, _>([2, 16], 4.elem()); + assert_eq!(full_dtype, full); + } + } + )* + }; +} + + test_dtypes!( + should_create_with_dtype, + bool, + i8, + i16, + i32, + i64, + u8, + u16, + u32, + u64, + f16, + bf16, + f32, + f64 + ); + + #[test] + fn should_serialize_deserialize_tensor_data() { + let data = TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3]); + assert_eq!( + data.as_bytes(), + [ + 0, 0, 128, 63, 0, 0, 0, 64, 0, 0, 64, 64, 0, 0, 128, 64, 0, 0, 160, 64, 0, 0, 192, + 64 + ] + ); + let serialized = serde_json::to_string(&data).unwrap(); + let deserialized: TensorData = serde_json::from_str(&serialized).unwrap(); + assert_eq!(data, deserialized); + } + + #[test] + fn should_deserialize_tensor_data_with_shape_inner() { + // TensorData `shape` was previously a Vec. + let serialized = r#"{ + "bytes": [0, 0, 128, 63, 0, 0, 0, 64, 0, 0, 64, 64, 0, 0, 128, 64, 0, 0, 160, 64, 0, 0, 192, 64], + "shape": [2, 3], + "dtype": "F32" + }"#; + + let data: TensorData = serde_json::from_str(serialized).unwrap(); + assert_eq!(data.shape, shape![2, 3]); + assert_eq!( + data.as_slice::().unwrap(), + &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + ); + } + + #[test] + fn should_serialize_shape_as_flat_array() { + // Ensure the new Shape serializes identically to how Vec used to, + // i.e. as a flat JSON array, not as an object like `{"dims": [2, 3]}`. + let data = TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3]); + let serialized = serde_json::to_string(&data).unwrap(); + let json: serde_json::Value = serde_json::from_str(&serialized).unwrap(); + assert_eq!(json["shape"], serde_json::json!([2, 3])); + } +} diff --git a/crates/burn-std/src/device_settings.rs b/crates/burn-std/src/device_settings.rs new file mode 100644 index 0000000..0047b62 --- /dev/null +++ b/crates/burn-std/src/device_settings.rs @@ -0,0 +1,125 @@ +//! Device data types: settings and errors. + +use alloc::format; +use alloc::string::String; +use cubecl_common::backtrace::BackTrace; +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +use crate::{BoolDType, DType, FloatDType, IntDType, QuantConfig}; + +/// Settings controlling the default data types for a specific device. +/// +/// These settings are managed in a global registry that enforces strict initialization semantics: +/// +/// 1. Manual Initialization: You can set these once at the start of your program using `set_default_dtypes`. +/// 2. Default Initialization: If an operation (like creating a tensor) occurs before manual initialization, +/// the settings are permanently locked to their default values. +/// 3. Immutability: Once initialized, settings cannot be changed. This ensures consistent behavior across +/// all threads and operations. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub struct DeviceSettings { + /// Default floating-point data type. + pub float_dtype: FloatDType, + /// Default integer data type. + pub int_dtype: IntDType, + /// Default bool data type. + pub bool_dtype: BoolDType, + /// Quantization configuration. + pub quantization: QuantConfig, +} + +impl DeviceSettings { + /// Creates a new [`DeviceSettings`] from any types convertible into the dtype kinds and the [quantization config](QuantConfig). + pub fn new( + float_dtype: impl Into, + int_dtype: impl Into, + bool_dtype: impl Into, + quantization: QuantConfig, + ) -> Self { + Self { + float_dtype: float_dtype.into(), + int_dtype: int_dtype.into(), + bool_dtype: bool_dtype.into(), + quantization, + } + } + + /// Creates a new [`DeviceSettings`] from any types convertible into the dtype kinds. + pub fn with_dtypes( + float_dtype: impl Into, + int_dtype: impl Into, + bool_dtype: impl Into, + ) -> Self { + Self::new(float_dtype, int_dtype, bool_dtype, Default::default()) + } +} + +/// Errors returned by device-related operations. +/// +/// Examples include attempting to use an unsupported data type on a device or initialization +/// errors like attempting to change a settings in an invalid context. +#[derive(Debug, Error)] +pub enum DeviceError { + /// Unsupported data type by the device. + #[error("Device {device} does not support the requested data type {dtype:?}")] + UnsupportedDType { + /// The string representation of the device. + device: String, + /// The data type that caused the error. + dtype: DType, + }, + /// Device settings have already been initialized. + #[error("Device {device} settings have already been initialized")] + AlreadyInitialized { + /// The string representation of the device. + device: String, + }, +} + +impl DeviceError { + /// Helper to create a [`DeviceError::UnsupportedDType`] from any device. + pub fn unsupported_dtype(device: &D, dtype: DType) -> Self { + Self::UnsupportedDType { + device: format!("{device:?}"), + dtype, + } + } + + /// Helper to create a [`DeviceError::AlreadyInitialized`] from any device. + pub fn already_initialized(device: &D) -> Self { + Self::AlreadyInitialized { + device: format!("{device:?}"), + } + } +} + +/// An error that can happen when syncing a device. +#[derive(Error, Serialize, Deserialize)] +pub enum ExecutionError { + /// A generic error happened during execution. + /// + /// The backtrace and context information should be included in the reason string. + #[error("An error happened during execution\nCaused by:\n {reason}")] + WithContext { + /// The reason of the error. + reason: String, + }, + /// A generic error happened during execution thrown in the Burn project. + /// + /// The full context isn't captured by the string alone. + #[error("An error happened during execution\nCaused by:\n {reason}")] + Generic { + /// The reason of the error. + reason: String, + /// The backtrace. + #[serde(skip)] + backtrace: BackTrace, + }, +} + +impl core::fmt::Debug for ExecutionError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_fmt(format_args!("{self}")) + } +} diff --git a/crates/burn-std/src/distributed.rs b/crates/burn-std/src/distributed.rs new file mode 100644 index 0000000..2fb765e --- /dev/null +++ b/crates/burn-std/src/distributed.rs @@ -0,0 +1,17 @@ +use serde::{Deserialize, Serialize}; + +/// The different ways to execute the reduce operation. +#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy, Serialize, Deserialize)] +pub enum ReduceOperation { + /// The sum of the values. + Sum, + /// The mean of the values. + Mean, +} + +/// Parameter struct for setting up and getting parameters for distributed operations. +#[derive(Clone, Debug)] +pub struct DistributedConfig { + /// How to execute the all_reduce operation. + pub all_reduce_op: ReduceOperation, +} diff --git a/crates/burn-std/src/distribution.rs b/crates/burn-std/src/distribution.rs new file mode 100644 index 0000000..d16ebc1 --- /dev/null +++ b/crates/burn-std/src/distribution.rs @@ -0,0 +1,125 @@ +//! Random value distributions used to initialize and populate tensor data. + +use rand::{Rng, RngExt, distr::StandardUniform}; + +use super::element::{Element, ElementConversion}; + +/// Distribution for random value of a tensor. +#[derive(Debug, Default, Clone, Copy, PartialEq, serde::Serialize, serde::Deserialize)] +pub enum Distribution { + /// Uniform distribution from 0 (inclusive) to 1 (exclusive). + #[default] + Default, + + /// Bernoulli distribution with the given probability. + Bernoulli(f64), + + /// Uniform distribution `[low, high)`. + Uniform(f64, f64), + + /// Normal distribution with the given mean and standard deviation. + Normal(f64, f64), +} + +/// Distribution sampler for random value of a tensor. +#[derive(new)] +pub struct DistributionSampler<'a, E, R> +where + StandardUniform: rand::distr::Distribution, + E: rand::distr::uniform::SampleUniform, + R: Rng, +{ + kind: DistributionSamplerKind, + rng: &'a mut R, +} + +/// Distribution sampler kind for random value of a tensor. +pub enum DistributionSamplerKind +where + StandardUniform: rand::distr::Distribution, + E: rand::distr::uniform::SampleUniform, +{ + /// Standard distribution. + Standard(rand::distr::StandardUniform), + + /// Uniform distribution. + Uniform(rand::distr::Uniform), + + /// Bernoulli distribution. + Bernoulli(rand::distr::Bernoulli), + + /// Normal distribution. + Normal(rand_distr::Normal), +} + +impl DistributionSampler<'_, E, R> +where + StandardUniform: rand::distr::Distribution, + E: rand::distr::uniform::SampleUniform, + E: Element, + R: Rng, +{ + /// Sames a random value from the distribution. + pub fn sample(&mut self) -> E { + match &self.kind { + DistributionSamplerKind::Standard(distribution) => self.rng.sample(distribution), + DistributionSamplerKind::Uniform(distribution) => self.rng.sample(distribution), + DistributionSamplerKind::Bernoulli(distribution) => { + if self.rng.sample(distribution) { + 1.elem() + } else { + 0.elem() + } + } + DistributionSamplerKind::Normal(distribution) => self.rng.sample(distribution).elem(), + } + } +} + +impl Distribution { + /// Creates a new distribution sampler. + /// + /// # Arguments + /// + /// * `rng` - The random number generator. + /// + /// # Returns + /// + /// The distribution sampler. + pub fn sampler(self, rng: &'_ mut R) -> DistributionSampler<'_, E, R> + where + R: Rng, + E: Element + rand::distr::uniform::SampleUniform, + StandardUniform: rand::distr::Distribution, + { + let kind = match self { + Distribution::Default => { + DistributionSamplerKind::Standard(rand::distr::StandardUniform {}) + } + Distribution::Uniform(low, high) => DistributionSamplerKind::Uniform( + rand::distr::Uniform::new(low.elem::(), high.elem::()).unwrap(), + ), + Distribution::Bernoulli(prob) => { + DistributionSamplerKind::Bernoulli(rand::distr::Bernoulli::new(prob).unwrap()) + } + Distribution::Normal(mean, std) => { + DistributionSamplerKind::Normal(rand_distr::Normal::new(mean, std).unwrap()) + } + }; + + DistributionSampler::new(kind, rng) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_distribution_default() { + let dist: Distribution = Default::default(); + + assert_eq!(dist, Distribution::Default); + assert_eq!(Distribution::default(), Distribution::Default); + } +} diff --git a/crates/burn-std/src/element/base.rs b/crates/burn-std/src/element/base.rs new file mode 100644 index 0000000..ee22ff9 --- /dev/null +++ b/crates/burn-std/src/element/base.rs @@ -0,0 +1,290 @@ +use core::cmp::Ordering; +use rand::Rng; + +use crate::distribution::Distribution; +use crate::{BoolStore, DType, bf16, f16, flex32}; + +use super::cast::ToElement; + +/// Core element trait for tensor values. +/// +/// This trait defines the minimal set of capabilities required for a type to be +/// stored and manipulated as a tensor element across all backends. +pub trait Element: + ToElement + + ElementRandom + + ElementConversion + + ElementEq + + bytemuck::CheckedBitPattern + + bytemuck::NoUninit + + bytemuck::Zeroable + + core::fmt::Debug + + core::fmt::Display + + Default + + Send + + Sync + + Copy + + 'static +{ + /// The dtype of the element. + fn dtype() -> DType; +} + +/// Ordered element trait for tensor values. +/// +/// This trait extends [`Element`] with ordering semantics, enabling comparison +/// and order-dependent operations in generic Rust implementations. +/// +/// Backends that implement these operations entirely at the device level do +/// not rely on this trait. It only constrains the scalar type for generic Rust code. +pub trait ElementOrdered: Element + ElementComparison + ElementLimits {} + +/// Element conversion trait for tensor. +pub trait ElementConversion { + /// Converts an element to another element. + /// + /// # Arguments + /// + /// * `elem` - The element to convert. + /// + /// # Returns + /// + /// The converted element. + fn from_elem(elem: E) -> Self; + + /// Converts and returns the converted element. + fn elem(self) -> E; +} + +/// Element trait for random value of a tensor. +pub trait ElementRandom { + /// Returns a random value for the given distribution. + /// + /// # Arguments + /// + /// * `distribution` - The distribution to sample from. + /// * `rng` - The random number generator. + /// + /// # Returns + /// + /// The random value. + fn random(distribution: Distribution, rng: &mut R) -> Self; +} + +/// Element trait for equality of a tensor. +pub trait ElementEq { + /// Returns whether `self` and `other` are equal. + fn eq(&self, other: &Self) -> bool; +} + +/// Element ordering trait. +pub trait ElementComparison { + /// Returns and [Ordering] between `self` and `other`. + fn cmp(&self, other: &Self) -> Ordering; +} + +/// Element limits trait. +pub trait ElementLimits { + /// The minimum representable value + const MIN: Self; + /// The maximum representable value + const MAX: Self; +} + +/// Macro to implement the element trait for a type. +#[macro_export] +macro_rules! make_element { + ( + ty $type:ident, + convert $convert:expr, + random $random:expr, + cmp $cmp:expr, + dtype $dtype:expr + ) => { + make_element!(ty $type, convert $convert, random $random, cmp $cmp, dtype $dtype, min $type::MIN, max $type::MAX); + }; + ( + ty $type:ident, + convert $convert:expr, + random $random:expr, + cmp $cmp:expr, + dtype $dtype:expr, + min $min:expr, + max $max:expr + ) => { + impl Element for $type { + #[inline(always)] + fn dtype() -> $crate::DType { + $dtype + } + } + impl ElementEq for $type { + fn eq(&self, other: &Self) -> bool { + self == other + } + } + + impl ElementConversion for $type { + #[inline(always)] + fn from_elem(elem: E) -> Self { + #[allow(clippy::redundant_closure_call)] + $convert(&elem) + } + #[inline(always)] + fn elem(self) -> E { + E::from_elem(self) + } + } + + impl ElementRandom for $type { + fn random(distribution: Distribution, rng: &mut R) -> Self { + #[allow(clippy::redundant_closure_call)] + $random(distribution, rng) + } + } + + impl ElementComparison for $type { + fn cmp(&self, other: &Self) -> Ordering { + let a = self.elem::<$type>(); + let b = other.elem::<$type>(); + #[allow(clippy::redundant_closure_call)] + $cmp(&a, &b) + } + } + + impl ElementLimits for $type { + const MIN: Self = $min; + const MAX: Self = $max; + } + + impl ElementOrdered for $type {} + + }; +} + +make_element!( + ty f64, + convert ToElement::to_f64, + random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(), + cmp |a: &f64, b: &f64| a.total_cmp(b), + dtype DType::F64 +); + +make_element!( + ty f32, + convert ToElement::to_f32, + random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(), + cmp |a: &f32, b: &f32| a.total_cmp(b), + dtype DType::F32 +); + +make_element!( + ty i64, + convert ToElement::to_i64, + random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(), + cmp |a: &i64, b: &i64| Ord::cmp(a, b), + dtype DType::I64 +); + +make_element!( + ty u64, + convert ToElement::to_u64, + random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(), + cmp |a: &u64, b: &u64| Ord::cmp(a, b), + dtype DType::U64 +); + +make_element!( + ty i32, + convert ToElement::to_i32, + random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(), + cmp |a: &i32, b: &i32| Ord::cmp(a, b), + dtype DType::I32 +); + +make_element!( + ty u32, + convert ToElement::to_u32, + random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(), + cmp |a: &u32, b: &u32| Ord::cmp(a, b), + dtype DType::U32 +); + +make_element!( + ty i16, + convert ToElement::to_i16, + random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(), + cmp |a: &i16, b: &i16| Ord::cmp(a, b), + dtype DType::I16 +); + +make_element!( + ty u16, + convert ToElement::to_u16, + random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(), + cmp |a: &u16, b: &u16| Ord::cmp(a, b), + dtype DType::U16 +); + +make_element!( + ty i8, + convert ToElement::to_i8, + random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(), + cmp |a: &i8, b: &i8| Ord::cmp(a, b), + dtype DType::I8 +); + +make_element!( + ty u8, + convert ToElement::to_u8, + random |distribution: Distribution, rng: &mut R| distribution.sampler(rng).sample(), + cmp |a: &u8, b: &u8| Ord::cmp(a, b), + dtype DType::U8 +); + +make_element!( + ty f16, + convert ToElement::to_f16, + random |distribution: Distribution, rng: &mut R| { + let sample: f32 = distribution.sampler(rng).sample(); + f16::from_elem(sample) + }, + cmp |a: &f16, b: &f16| a.total_cmp(b), + dtype DType::F16 +); +make_element!( + ty bf16, + convert ToElement::to_bf16, + random |distribution: Distribution, rng: &mut R| { + let sample: f32 = distribution.sampler(rng).sample(); + bf16::from_elem(sample) + }, + cmp |a: &bf16, b: &bf16| a.total_cmp(b), + dtype DType::BF16 +); + +make_element!( + ty flex32, + convert |elem: &dyn ToElement| flex32::from_f32(elem.to_f32()), + random |distribution: Distribution, rng: &mut R| { + let sample: f32 = distribution.sampler(rng).sample(); + flex32::from_elem(sample) + }, + cmp |a: &flex32, b: &flex32| a.total_cmp(b), + dtype DType::Flex32, + min flex32::from_f32(f16::MIN.to_f32_const()), + max flex32::from_f32(f16::MAX.to_f32_const()) +); + +make_element!( + ty bool, + convert ToElement::to_bool, + random |distribution: Distribution, rng: &mut R| { + let sample: u8 = distribution.sampler(rng).sample(); + bool::from_elem(sample) + }, + cmp |a: &bool, b: &bool| Ord::cmp(a, b), + dtype DType::Bool(BoolStore::Native), + min false, + max true +); diff --git a/crates/burn-std/src/element/cast.rs b/crates/burn-std/src/element/cast.rs new file mode 100644 index 0000000..792adeb --- /dev/null +++ b/crates/burn-std/src/element/cast.rs @@ -0,0 +1,705 @@ +use core::mem::size_of; + +use crate::{bf16, f16}; + +/// A generic trait for converting a value to a number. +/// Adapted from num_traits::ToPrimitive to support [bool]. +/// +/// A value can be represented by the target type when it lies within +/// the range of scalars supported by the target type. +/// For example, a negative integer cannot be represented by an unsigned +/// integer type, and an `i64` with a very high magnitude might not be +/// convertible to an `i32`. +/// On the other hand, conversions with possible precision loss or truncation +/// are admitted, like an `f32` with a decimal part to an integer type, or +/// even a large `f64` saturating to `f32` infinity. +/// +/// The methods *panic* when the value cannot be represented by the target type. +pub trait ToElement { + /// Converts the value of `self` to an `isize`. + #[inline] + fn to_isize(&self) -> isize { + ToElement::to_isize(&self.to_i64()) + } + + /// Converts the value of `self` to an `i8`. + #[inline] + fn to_i8(&self) -> i8 { + ToElement::to_i8(&self.to_i64()) + } + + /// Converts the value of `self` to an `i16`. + #[inline] + fn to_i16(&self) -> i16 { + ToElement::to_i16(&self.to_i64()) + } + + /// Converts the value of `self` to an `i32`. + #[inline] + fn to_i32(&self) -> i32 { + ToElement::to_i32(&self.to_i64()) + } + + /// Converts the value of `self` to an `i64`. + fn to_i64(&self) -> i64; + + /// Converts the value of `self` to an `i128`. + /// + /// The default implementation converts through `to_i64()`. Types implementing + /// this trait should override this method if they can represent a greater range. + #[inline] + fn to_i128(&self) -> i128 { + i128::from(self.to_i64()) + } + + /// Converts the value of `self` to a `usize`. + #[inline] + fn to_usize(&self) -> usize { + ToElement::to_usize(&self.to_u64()) + } + + /// Converts the value of `self` to a `u8`. + #[inline] + fn to_u8(&self) -> u8 { + ToElement::to_u8(&self.to_u64()) + } + + /// Converts the value of `self` to a `u16`. + #[inline] + fn to_u16(&self) -> u16 { + ToElement::to_u16(&self.to_u64()) + } + + /// Converts the value of `self` to a `u32`. + #[inline] + fn to_u32(&self) -> u32 { + ToElement::to_u32(&self.to_u64()) + } + + /// Converts the value of `self` to a `u64`. + fn to_u64(&self) -> u64; + + /// Converts the value of `self` to a `u128`. + /// + /// The default implementation converts through `to_u64()`. Types implementing + /// this trait should override this method if they can represent a greater range. + #[inline] + fn to_u128(&self) -> u128 { + u128::from(self.to_u64()) + } + + /// Converts the value of `self` to an `f16`. Overflows may map to positive + /// or negative infinity. + #[inline] + fn to_f16(&self) -> f16 { + f16::from_f32(self.to_f32()) + } + + /// Converts the value of `self` to an `bf16`. Overflows may map to positive + /// or negative infinity. + #[inline] + fn to_bf16(&self) -> bf16 { + bf16::from_f32(self.to_f32()) + } + + /// Converts the value of `self` to an `f32`. Overflows may map to positive + /// or negative infinity. + #[inline] + fn to_f32(&self) -> f32 { + ToElement::to_f32(&self.to_f64()) + } + + /// Converts the value of `self` to an `f64`. Overflows may map to positive + /// or negative infinity. + /// + /// The default implementation tries to convert through `to_i64()`, and + /// failing that through `to_u64()`. Types implementing this trait should + /// override this method if they can represent a greater range. + #[inline] + fn to_f64(&self) -> f64 { + ToElement::to_f64(&self.to_u64()) + } + + /// Converts the value of `self` to a bool. + /// Rust only considers 0 and 1 to be valid booleans, but for compatibility, C semantics are + /// adopted (anything that's not 0 is true). + /// + /// The default implementation tries to convert through `to_i64()`, and + /// failing that through `to_u64()`. Types implementing this trait should + /// override this method if they can represent a greater range. + #[inline] + fn to_bool(&self) -> bool { + ToElement::to_bool(&self.to_u64()) + } +} + +macro_rules! impl_to_element_int_to_int { + ($SrcT:ident : $( $(#[$cfg:meta])* fn $method:ident -> $DstT:ident ; )*) => {$( + #[inline] + $(#[$cfg])* + fn $method(&self) -> $DstT { + let min = $DstT::MIN as $SrcT; + let max = $DstT::MAX as $SrcT; + if size_of::<$SrcT>() <= size_of::<$DstT>() || (min <= *self && *self <= max) { + *self as $DstT + } else { + panic!( + "Element cannot be represented in the target type: {:?}({:?}) => {:?}", + core::any::type_name::<$SrcT>(), + self, + core::any::type_name::<$DstT>(), + ) + } + } + )*} +} + +macro_rules! impl_to_element_int_to_uint { + ($SrcT:ident : $( $(#[$cfg:meta])* fn $method:ident -> $DstT:ident ; )*) => {$( + #[inline] + $(#[$cfg])* + fn $method(&self) -> $DstT { + let max = $DstT::MAX as $SrcT; + if 0 <= *self && (size_of::<$SrcT>() <= size_of::<$DstT>() || *self <= max) { + *self as $DstT + } else { + panic!( + "Element cannot be represented in the target type: {:?}({:?}) => {:?}", + core::any::type_name::<$SrcT>(), + self, + core::any::type_name::<$DstT>(), + ) + } + } + )*} +} + +macro_rules! impl_to_element_int { + ($T:ident) => { + impl ToElement for $T { + impl_to_element_int_to_int! { $T: + fn to_isize -> isize; + fn to_i8 -> i8; + fn to_i16 -> i16; + fn to_i32 -> i32; + fn to_i64 -> i64; + fn to_i128 -> i128; + } + + impl_to_element_int_to_uint! { $T: + fn to_usize -> usize; + fn to_u8 -> u8; + fn to_u16 -> u16; + fn to_u32 -> u32; + fn to_u64 -> u64; + fn to_u128 -> u128; + } + + #[inline] + fn to_f32(&self) -> f32 { + *self as f32 + } + #[inline] + fn to_f64(&self) -> f64 { + *self as f64 + } + #[inline] + fn to_bool(&self) -> bool { + *self != 0 + } + } + }; +} + +impl_to_element_int!(isize); +impl_to_element_int!(i8); +impl_to_element_int!(i16); +impl_to_element_int!(i32); +impl_to_element_int!(i64); +impl_to_element_int!(i128); + +macro_rules! impl_to_element_uint_to_int { + ($SrcT:ident : $( $(#[$cfg:meta])* fn $method:ident -> $DstT:ident ; )*) => {$( + #[inline] + $(#[$cfg])* + fn $method(&self) -> $DstT { + let max = $DstT::MAX as $SrcT; + if size_of::<$SrcT>() < size_of::<$DstT>() || *self <= max { + *self as $DstT + } else { + panic!( + "Element cannot be represented in the target type: {:?}({:?}) => {:?}", + core::any::type_name::<$SrcT>(), + self, + core::any::type_name::<$DstT>(), + ) + } + } + )*} +} + +macro_rules! impl_to_element_uint_to_uint { + ($SrcT:ident : $( $(#[$cfg:meta])* fn $method:ident -> $DstT:ident ; )*) => {$( + #[inline] + $(#[$cfg])* + fn $method(&self) -> $DstT { + let max = $DstT::MAX as $SrcT; + if size_of::<$SrcT>() <= size_of::<$DstT>() || *self <= max { + *self as $DstT + } else { + panic!( + "Element cannot be represented in the target type: {:?}({:?}) => {:?}", + core::any::type_name::<$SrcT>(), + self, + core::any::type_name::<$DstT>(), + ) + } + } + )*} +} + +macro_rules! impl_to_element_uint { + ($T:ident) => { + impl ToElement for $T { + impl_to_element_uint_to_int! { $T: + fn to_isize -> isize; + fn to_i8 -> i8; + fn to_i16 -> i16; + fn to_i32 -> i32; + fn to_i64 -> i64; + fn to_i128 -> i128; + } + + impl_to_element_uint_to_uint! { $T: + fn to_usize -> usize; + fn to_u8 -> u8; + fn to_u16 -> u16; + fn to_u32 -> u32; + fn to_u64 -> u64; + fn to_u128 -> u128; + } + + #[inline] + fn to_f32(&self) -> f32 { + *self as f32 + } + #[inline] + fn to_f64(&self) -> f64 { + *self as f64 + } + #[inline] + fn to_bool(&self) -> bool { + *self != 0 + } + } + }; +} + +impl_to_element_uint!(usize); +impl_to_element_uint!(u8); +impl_to_element_uint!(u16); +impl_to_element_uint!(u32); +impl_to_element_uint!(u64); +impl_to_element_uint!(u128); + +macro_rules! impl_to_element_float_to_float { + ($SrcT:ident : $( fn $method:ident -> $DstT:ident ; )*) => {$( + #[inline] + fn $method(&self) -> $DstT { + // We can safely cast all values, whether NaN, +-inf, or finite. + // Finite values that are reducing size may saturate to +-inf. + *self as $DstT + } + )*} +} + +macro_rules! float_to_int_unchecked { + // SAFETY: Must not be NaN or infinite; must be representable as the integer after truncating. + // We already checked that the float is in the exclusive range `(MIN-1, MAX+1)`. + ($float:expr => $int:ty) => { + unsafe { $float.to_int_unchecked::<$int>() } + }; +} + +macro_rules! impl_to_element_float_to_signed_int { + ($f:ident : $( $(#[$cfg:meta])* fn $method:ident -> $i:ident ; )*) => {$( + #[inline] + $(#[$cfg])* + fn $method(&self) -> $i { + // Float as int truncates toward zero, so we want to allow values + // in the exclusive range `(MIN-1, MAX+1)`. + if size_of::<$f>() > size_of::<$i>() { + // With a larger size, we can represent the range exactly. + const MIN_M1: $f = $i::MIN as $f - 1.0; + const MAX_P1: $f = $i::MAX as $f + 1.0; + if *self > MIN_M1 && *self < MAX_P1 { + return float_to_int_unchecked!(*self => $i); + } + } else { + // We can't represent `MIN-1` exactly, but there's no fractional part + // at this magnitude, so we can just use a `MIN` inclusive boundary. + const MIN: $f = $i::MIN as $f; + // We can't represent `MAX` exactly, but it will round up to exactly + // `MAX+1` (a power of two) when we cast it. + const MAX_P1: $f = $i::MAX as $f; + if *self >= MIN && *self < MAX_P1 { + return float_to_int_unchecked!(*self => $i); + } + } + panic!("Float cannot be represented in the target signed int type") + } + )*} +} + +macro_rules! impl_to_element_float_to_unsigned_int { + ($f:ident : $( $(#[$cfg:meta])* fn $method:ident -> $u:ident ; )*) => {$( + #[inline] + $(#[$cfg])* + fn $method(&self) -> $u { + // Float as int truncates toward zero, so we want to allow values + // in the exclusive range `(-1, MAX+1)`. + if size_of::<$f>() > size_of::<$u>() { + // With a larger size, we can represent the range exactly. + const MAX_P1: $f = $u::MAX as $f + 1.0; + if *self > -1.0 && *self < MAX_P1 { + return float_to_int_unchecked!(*self => $u); + } + } else { + // We can't represent `MAX` exactly, but it will round up to exactly + // `MAX+1` (a power of two) when we cast it. + // (`u128::MAX as f32` is infinity, but this is still ok.) + const MAX_P1: $f = $u::MAX as $f; + if *self > -1.0 && *self < MAX_P1 { + return float_to_int_unchecked!(*self => $u); + } + } + panic!("Float cannot be represented in the target unsigned int type") + } + )*} +} + +macro_rules! impl_to_element_float { + ($T:ident) => { + impl ToElement for $T { + impl_to_element_float_to_signed_int! { $T: + fn to_isize -> isize; + fn to_i8 -> i8; + fn to_i16 -> i16; + fn to_i32 -> i32; + fn to_i64 -> i64; + fn to_i128 -> i128; + } + + impl_to_element_float_to_unsigned_int! { $T: + fn to_usize -> usize; + fn to_u8 -> u8; + fn to_u16 -> u16; + fn to_u32 -> u32; + fn to_u64 -> u64; + fn to_u128 -> u128; + } + + impl_to_element_float_to_float! { $T: + fn to_f32 -> f32; + fn to_f64 -> f64; + } + + #[inline] + fn to_bool(&self) -> bool { + *self != 0.0 + } + } + }; +} + +impl_to_element_float!(f32); +impl_to_element_float!(f64); + +impl ToElement for f16 { + #[inline] + fn to_i64(&self) -> i64 { + Self::to_f32(*self).to_i64() + } + #[inline] + fn to_u64(&self) -> u64 { + Self::to_f32(*self).to_u64() + } + #[inline] + fn to_i8(&self) -> i8 { + Self::to_f32(*self).to_i8() + } + #[inline] + fn to_u8(&self) -> u8 { + Self::to_f32(*self).to_u8() + } + #[inline] + fn to_i16(&self) -> i16 { + Self::to_f32(*self).to_i16() + } + #[inline] + fn to_u16(&self) -> u16 { + Self::to_f32(*self).to_u16() + } + #[inline] + fn to_i32(&self) -> i32 { + Self::to_f32(*self).to_i32() + } + #[inline] + fn to_u32(&self) -> u32 { + Self::to_f32(*self).to_u32() + } + #[inline] + fn to_f16(&self) -> f16 { + *self + } + #[inline] + fn to_f32(&self) -> f32 { + Self::to_f32(*self) + } + #[inline] + fn to_f64(&self) -> f64 { + Self::to_f64(*self) + } + #[inline] + fn to_bool(&self) -> bool { + *self != f16::from_f32_const(0.0) + } +} + +impl ToElement for bf16 { + #[inline] + fn to_i64(&self) -> i64 { + Self::to_f32(*self).to_i64() + } + #[inline] + fn to_u64(&self) -> u64 { + Self::to_f32(*self).to_u64() + } + #[inline] + fn to_i8(&self) -> i8 { + Self::to_f32(*self).to_i8() + } + #[inline] + fn to_u8(&self) -> u8 { + Self::to_f32(*self).to_u8() + } + #[inline] + fn to_i16(&self) -> i16 { + Self::to_f32(*self).to_i16() + } + #[inline] + fn to_u16(&self) -> u16 { + Self::to_f32(*self).to_u16() + } + #[inline] + fn to_i32(&self) -> i32 { + Self::to_f32(*self).to_i32() + } + #[inline] + fn to_u32(&self) -> u32 { + Self::to_f32(*self).to_u32() + } + #[inline] + fn to_bf16(&self) -> bf16 { + *self + } + #[inline] + fn to_f32(&self) -> f32 { + Self::to_f32(*self) + } + #[inline] + fn to_f64(&self) -> f64 { + Self::to_f64(*self) + } + #[inline] + fn to_bool(&self) -> bool { + *self != bf16::from_f32_const(0.0) + } +} + +impl ToElement for crate::flex32 { + #[inline] + fn to_i64(&self) -> i64 { + Self::to_f32(*self).to_i64() + } + #[inline] + fn to_u64(&self) -> u64 { + Self::to_f32(*self).to_u64() + } + #[inline] + fn to_i8(&self) -> i8 { + Self::to_f32(*self).to_i8() + } + #[inline] + fn to_u8(&self) -> u8 { + Self::to_f32(*self).to_u8() + } + #[inline] + fn to_i16(&self) -> i16 { + Self::to_f32(*self).to_i16() + } + #[inline] + fn to_u16(&self) -> u16 { + Self::to_f32(*self).to_u16() + } + #[inline] + fn to_i32(&self) -> i32 { + Self::to_f32(*self).to_i32() + } + #[inline] + fn to_u32(&self) -> u32 { + Self::to_f32(*self).to_u32() + } + #[inline] + fn to_f32(&self) -> f32 { + Self::to_f32(*self) + } + #[inline] + fn to_f64(&self) -> f64 { + Self::to_f64(*self) + } + #[inline] + fn to_bool(&self) -> bool { + *self != crate::flex32::from_f32(0.0) + } +} + +impl ToElement for bool { + #[inline] + fn to_i64(&self) -> i64 { + *self as i64 + } + #[inline] + fn to_u64(&self) -> u64 { + *self as u64 + } + #[inline] + fn to_i8(&self) -> i8 { + *self as i8 + } + #[inline] + fn to_u8(&self) -> u8 { + *self as u8 + } + #[inline] + fn to_i16(&self) -> i16 { + *self as i16 + } + #[inline] + fn to_u16(&self) -> u16 { + *self as u16 + } + #[inline] + fn to_i32(&self) -> i32 { + *self as i32 + } + #[inline] + fn to_u32(&self) -> u32 { + *self as u32 + } + #[inline] + fn to_f32(&self) -> f32 { + self.to_u8() as f32 + } + #[inline] + fn to_f64(&self) -> f64 { + self.to_u8() as f64 + } + #[inline] + fn to_bool(&self) -> bool { + *self + } +} + +mod tests { + #[allow(unused_imports)] + use super::*; + + #[test] + fn to_element_float() { + let f32_toolarge = 1e39f64; + assert_eq!(f32_toolarge.to_f32(), f32::INFINITY); + assert_eq!((-f32_toolarge).to_f32(), f32::NEG_INFINITY); + assert_eq!((f32::MAX as f64).to_f32(), f32::MAX); + assert_eq!((-f32::MAX as f64).to_f32(), -f32::MAX); + assert_eq!(f64::INFINITY.to_f32(), f32::INFINITY); + assert_eq!((f64::NEG_INFINITY).to_f32(), f32::NEG_INFINITY); + assert!((f64::NAN).to_f32().is_nan()); + } + + #[test] + #[should_panic] + fn to_element_signed_to_u8_underflow() { + let _x = (-1i8).to_u8(); + } + + #[test] + #[should_panic] + fn to_element_signed_to_u16_underflow() { + let _x = (-1i8).to_u16(); + } + + #[test] + #[should_panic] + fn to_element_signed_to_u32_underflow() { + let _x = (-1i8).to_u32(); + } + + #[test] + #[should_panic] + fn to_element_signed_to_u64_underflow() { + let _x = (-1i8).to_u64(); + } + + #[test] + #[should_panic] + fn to_element_signed_to_u128_underflow() { + let _x = (-1i8).to_u128(); + } + + #[test] + #[should_panic] + fn to_element_signed_to_usize_underflow() { + let _x = (-1i8).to_usize(); + } + + #[test] + #[should_panic] + fn to_element_unsigned_to_u8_overflow() { + let _x = 256.to_u8(); + } + + #[test] + #[should_panic] + fn to_element_unsigned_to_u16_overflow() { + let _x = 65_536.to_u16(); + } + + #[test] + #[should_panic] + fn to_element_unsigned_to_u32_overflow() { + let _x = 4_294_967_296u64.to_u32(); + } + + #[test] + #[should_panic] + fn to_element_unsigned_to_u64_overflow() { + let _x = 18_446_744_073_709_551_616u128.to_u64(); + } + + #[test] + fn to_element_int_to_float() { + assert_eq!((-1).to_f32(), -1.0); + assert_eq!((-1).to_f64(), -1.0); + assert_eq!(255.to_f32(), 255.0); + assert_eq!(65_535.to_f64(), 65_535.0); + } + + #[test] + fn to_element_float_to_int() { + assert_eq!((-1.0).to_i8(), -1); + assert_eq!(1.0.to_u8(), 1); + assert_eq!(1.8.to_u16(), 1); + assert_eq!(123.456.to_u32(), 123); + } +} diff --git a/crates/burn-std/src/element/mod.rs b/crates/burn-std/src/element/mod.rs new file mode 100644 index 0000000..c1f7884 --- /dev/null +++ b/crates/burn-std/src/element/mod.rs @@ -0,0 +1,10 @@ +//! Traits and helpers for working with element types and conversions. + +mod base; +mod scalar; + +/// Tensor element casting. +pub mod cast; + +pub use base::*; +pub use scalar::*; diff --git a/crates/burn-std/src/element/scalar.rs b/crates/burn-std/src/element/scalar.rs new file mode 100644 index 0000000..709102d --- /dev/null +++ b/crates/burn-std/src/element/scalar.rs @@ -0,0 +1,111 @@ +use crate::{BoolStore, DType, bf16, f16}; +use num_traits::ToPrimitive; + +#[cfg(not(feature = "std"))] +#[allow(unused_imports)] +use num_traits::Float; + +use crate::{Element, ElementConversion}; + +/// A scalar element. +#[derive(Clone, Copy, Debug)] +#[allow(missing_docs)] +pub enum Scalar { + Float(f64), + Int(i64), + UInt(u64), + Bool(bool), +} + +impl Scalar { + /// Creates a scalar with the specified data type. + /// + /// # Note + /// [`QFloat`](DType::QFloat) scalars are represented as float for element-wise operations. + pub fn new(value: E, dtype: &DType) -> Self { + if dtype.is_float() | matches!(dtype, &DType::QFloat(_)) { + Self::Float(value.elem()) + } else if dtype.is_int() { + Self::Int(value.elem()) + } else if dtype.is_uint() { + Self::UInt(value.elem()) + } else if dtype.is_bool() { + match dtype { + DType::Bool(BoolStore::Native) => Self::Bool(value.elem()), + DType::Bool(BoolStore::U8) | DType::Bool(BoolStore::U32) => { + Self::UInt(value.elem()) + } + _ => unreachable!(), + } + } else { + unimplemented!("Scalar not supported for {dtype:?}") + } + } + + /// Converts and returns the converted element. + pub fn elem(self) -> E { + match self { + Self::Float(x) => x.elem(), + Self::Int(x) => x.elem(), + Self::UInt(x) => x.elem(), + Self::Bool(x) => x.elem(), + } + } + + /// Returns the exact integer value, if valid. + pub fn try_as_integer(&self) -> Option { + match self { + Scalar::Float(x) => (x.floor() == *x).then(|| Self::Int(x.to_i64().unwrap())), + Scalar::Int(_) | Scalar::UInt(_) => Some(*self), + Scalar::Bool(x) => Some(Scalar::Int(*x as i64)), + } + } +} + +macro_rules! impl_from_scalar { + ($($ty:ty => $variant:ident),+ $(,)?) => { + $( + impl From<$ty> for Scalar { + fn from(value: $ty) -> Self { + Scalar::$variant(value.elem()) + } + } + )+ + }; +} + +impl_from_scalar! { + f64 => Float, f32 => Float, f16 => Float, bf16 => Float, + i64 => Int, i32 => Int, i16 => Int, i8 => Int, + u64 => UInt, u32 => UInt, u16 => UInt, u8 => UInt, bool => Bool, +} + +// CubeCL requirement +impl ToPrimitive for Scalar { + fn to_i64(&self) -> Option { + match self { + Scalar::Float(x) => x.to_i64(), + Scalar::UInt(x) => x.to_i64(), + Scalar::Int(x) => Some(*x), + Scalar::Bool(x) => Some(*x as i64), + } + } + + fn to_u64(&self) -> Option { + match self { + Scalar::Float(x) => x.to_u64(), + Scalar::UInt(x) => Some(*x), + Scalar::Int(x) => x.to_u64(), + Scalar::Bool(x) => Some(*x as u64), + } + } + + fn to_f64(&self) -> Option { + match self { + Scalar::Float(x) => Some(*x), + Scalar::UInt(x) => x.to_f64(), + Scalar::Int(x) => x.to_f64(), + Scalar::Bool(x) => (*x as u8).to_f64(), + } + } +} diff --git a/crates/burn-std/src/id.rs b/crates/burn-std/src/id.rs new file mode 100644 index 0000000..4533792 --- /dev/null +++ b/crates/burn-std/src/id.rs @@ -0,0 +1,193 @@ +//! # Unique Identifiers +use crate::rand::gen_random; + +/// Simple ID generator. +pub struct IdGenerator {} + +impl IdGenerator { + /// Generates a new ID. + pub fn generate() -> u64 { + // Generate a random u64 (18,446,744,073,709,551,615 combinations) + let random_bytes: [u8; 8] = gen_random(); + u64::from_le_bytes(random_bytes) + } +} + +pub use cubecl_common::stream_id::StreamId; + +use core::hash::{BuildHasher, Hasher}; + +use alloc::str::FromStr; +use data_encoding::BASE32_DNSSEC; +use serde::{Deserialize, Serialize}; + +// Hashbrown changed its default hasher in 0.15, but there are some issues +// https://github.com/rust-lang/hashbrown/issues/577 +// Also, `param_serde_deserialize_legacy_uuid` doesn't pass with the default hasher. +type DefaultHashBuilder = core::hash::BuildHasherDefault; + +/// Unique ID for a parameter of a module. +#[derive(Debug, Hash, PartialEq, Eq, Clone, Copy, PartialOrd, Ord, Serialize, Deserialize)] +pub struct ParamId { + value: u64, +} + +impl From for ParamId { + fn from(value: u64) -> Self { + Self { value } + } +} + +impl Default for ParamId { + fn default() -> Self { + Self::new() + } +} + +impl ParamId { + /// Create a new parameter ID. + pub fn new() -> Self { + Self { + value: IdGenerator::generate(), + } + } + + /// Gets the internal value of the id. + pub fn val(&self) -> u64 { + self.value + } +} + +impl FromStr for ParamId { + type Err = &'static str; // Or a custom error type if preferred + + /// Construct a param id from str. + /// + /// Preserves compatibility with previous formats (6 bytes, 16-byte uuid). + /// + /// # Returns + /// A `Result` containing the `ParamId` when valid. + fn from_str(encoded: &str) -> Result { + let u64_id: Option = match BASE32_DNSSEC.decode(encoded.as_bytes()) { + Ok(bytes) => { + let mut buffer = [0u8; 8]; + buffer[..bytes.len()].copy_from_slice(&bytes); + Some(u64::from_le_bytes(buffer)) + } + Err(_) => match uuid::Uuid::try_parse(encoded) { + // Backward compatibility with uuid parameter identifiers + Ok(id) => { + // Hash the 128-bit uuid to 64-bit + // Though not *theoretically* unique, the probability of a collision should be extremely low + let mut hasher = DefaultHashBuilder::default().build_hasher(); + // let mut hasher = DefaultHasher::new(); + hasher.write(id.as_bytes()); + Some(hasher.finish()) + } + Err(_) => None, + }, + }; + u64_id.map(Self::from).ok_or("Invalid id.") + } +} + +impl core::fmt::Display for ParamId { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let encoded = BASE32_DNSSEC.encode(&self.value.to_le_bytes()); + f.write_str(&encoded) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use alloc::collections::BTreeSet; + use alloc::string::ToString; + + #[cfg(feature = "std")] + use dashmap::DashSet; //Concurrent HashMap + #[cfg(feature = "std")] + use std::{sync::Arc, thread}; + + #[test] + fn uniqueness_test() { + const IDS_CNT: usize = 10_000; + + let mut set: BTreeSet = BTreeSet::new(); + + for _i in 0..IDS_CNT { + assert!(set.insert(IdGenerator::generate())); + } + + assert_eq!(set.len(), IDS_CNT); + } + + #[cfg(feature = "std")] + #[test] + fn thread_safety_test() { + const NUM_THREADS: usize = 10; + const NUM_REPEATS: usize = 1_000; + const EXPECTED_TOTAL_IDS: usize = NUM_THREADS * NUM_REPEATS; + + let set: Arc> = Arc::new(DashSet::new()); + + let mut handles = vec![]; + + for _ in 0..NUM_THREADS { + let set = set.clone(); + + let handle = thread::spawn(move || { + for _i in 0..NUM_REPEATS { + assert!(set.insert(IdGenerator::generate())); + } + }); + handles.push(handle); + } + + for handle in handles { + handle.join().unwrap(); + } + assert_eq!(set.len(), EXPECTED_TOTAL_IDS); + } + + #[test] + fn param_serde_try_deserialize() { + let val = ParamId::from(123456u64); + let deserialized = ParamId::from_str(&val.to_string()).unwrap(); + assert_eq!(val, deserialized); + + assert_eq!(ParamId::from_str("invalid_id"), Err("Invalid id.")); + } + + #[test] + fn param_serde_deserialize() { + let val = ParamId::from(123456u64); + let deserialized = ParamId::from_str(&val.to_string()).unwrap(); + assert_eq!(val, deserialized); + } + + #[test] + fn param_serde_deserialize_legacy() { + let legacy_val = [45u8; 6]; + let param_id = ParamId::from_str(&BASE32_DNSSEC.encode(&legacy_val)).unwrap(); + assert_eq!(param_id.val().to_le_bytes()[0..6], legacy_val); + assert_eq!(param_id.val().to_le_bytes()[6..], [0, 0]); + } + + #[test] + fn param_serde_deserialize_legacy_uuid() { + // Ensure support for legacy uuid deserialization and make sure it results in the same output + let legacy_id = "30b82c23-788d-4d63-a743-ada258d5f13c"; + let param_id1 = ParamId::from_str(legacy_id).unwrap(); + let param_id2 = ParamId::from_str(legacy_id).unwrap(); + assert_eq!(param_id1, param_id2); + } + + #[test] + #[should_panic = "Invalid id."] + fn param_serde_deserialize_invalid_id() { + let invalid_uuid = "30b82c23-788d-4d63-ada258d5f13c"; + let _ = ParamId::from_str(invalid_uuid).unwrap(); + } +} diff --git a/crates/burn-std/src/lib.rs b/crates/burn-std/src/lib.rs new file mode 100644 index 0000000..f05178f --- /dev/null +++ b/crates/burn-std/src/lib.rs @@ -0,0 +1,88 @@ +#![cfg_attr(not(feature = "std"), no_std)] +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] + +//! # Burn Standard Library +//! +//! This library contains core types and utilities shared across Burn, including shapes, indexing, +//! and data types. + +#[macro_use] +extern crate derive_new; + +extern crate alloc; + +/// Id module contains types for unique identifiers. +pub mod id; + +/// Tensor utilities. +pub mod tensor; +pub use tensor::*; + +/// Tensor data representation and helpers. +pub mod data; +pub use data::*; + +/// Random value distributions. +pub mod distribution; +pub use distribution::*; + +/// Traits for tensor element types and conversions. +pub mod element; +pub use element::*; + +mod device_settings; +pub use device_settings::*; + +/// Runtime kind of the host program (async / sync / no-std). +pub mod runtime_kind; +pub use runtime_kind::*; + +/// Distributed configurations. +pub mod distributed; + +/// Configuration types for tensor operations (conv, pool, interpolate, pad, etc). +pub mod ops; +pub use ops::*; + +/// Burn runtime configurations. +pub mod config; + +/// Common Errors. +pub use cubecl_zspace::errors::{self, *}; + +/// Network utilities. +#[cfg(feature = "network")] +pub mod network; + +/// An ID unique to any unordered combination of devices, used by collective / +/// communication primitives (distributed training etc.). +/// +/// Mirrors `cubecl_runtime::server::CommunicationId` so that the +/// `burn_fusion::FusionUtilities::initialized_comms` set (and other consumers) +/// can be reused without depending on cubecl directly. +#[derive(Clone, Debug, Hash, Eq, PartialEq)] +pub struct CommunicationId { + /// Stable hash of the (sorted) set of device ids that participate. + pub id: u64, +} + +impl From> for CommunicationId { + fn from(mut value: alloc::vec::Vec) -> Self { + use core::hash::{Hash, Hasher}; + // Sort so any permutation of the same devices yields the same id. + value.sort(); + let mut hasher = ahash::AHasher::default(); + value.hash(&mut hasher); + CommunicationId { + id: hasher.finish(), + } + } +} + +pub use cubecl_common::bytes::*; +pub use cubecl_common::device_handle::DeviceHandle; +pub use cubecl_common::*; +pub use half::{bf16, f16}; + +pub use cubecl_common::flex32; diff --git a/crates/burn-std/src/network.rs b/crates/burn-std/src/network.rs new file mode 100644 index 0000000..621cc10 --- /dev/null +++ b/crates/burn-std/src/network.rs @@ -0,0 +1,57 @@ +//! # Common Network Utilities + +/// Network download utilities. +pub mod downloader { + use indicatif::{ProgressBar, ProgressState, ProgressStyle}; + use reqwest::Client; + use std::io::Write; + + /// Download the file at the specified url. + /// File download progress is reported with the help of a [progress bar](indicatif). + /// + /// # Arguments + /// + /// * `url` - The file URL to download. + /// * `message` - The message to display on the progress bar during download. + /// + /// # Returns + /// + /// A vector of bytes containing the downloaded file data. + #[tokio::main(flavor = "current_thread")] + pub async fn download_file_as_bytes(url: &str, message: &str) -> Vec { + // Get file from web + let mut response = Client::new().get(url).send().await.unwrap(); + let total_size = response.content_length().unwrap(); + + // Pretty progress bar + let pb = ProgressBar::new(total_size); + let msg = message.to_owned(); + pb.set_style( + ProgressStyle::with_template( + "{msg}\n {wide_bar:.cyan/blue} {bytes}/{total_bytes} ({eta})", + ) + .unwrap() + .with_key( + "eta", + |state: &ProgressState, w: &mut dyn std::fmt::Write| { + write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap() + }, + ) + .progress_chars("▬ "), + ); + pb.set_message(msg.clone()); + + // Read stream into bytes + let mut downloaded: u64 = 0; + let mut bytes: Vec = Vec::with_capacity(total_size as usize); + while let Some(chunk) = response.chunk().await.unwrap() { + let num_bytes = bytes.write(&chunk).unwrap(); + let new = std::cmp::min(downloaded + (num_bytes as u64), total_size); + downloaded = new; + pb.set_position(new); + } + pb.finish_with_message(msg); + + bytes + } +} diff --git a/crates/burn-std/src/ops.rs b/crates/burn-std/src/ops.rs new file mode 100644 index 0000000..40d4a48 --- /dev/null +++ b/crates/burn-std/src/ops.rs @@ -0,0 +1,480 @@ +//! Configuration types for tensor operations. + +use crate::ElementConversion; +use core::num::NonZeroUsize; + +/// Check that the parameter value is non-zero. +// NOTE: for now we keep usize but we could refactor the parameters to hold `NonZeroUsize`. +pub(crate) fn check_nonzero(value: usize, msg: &str) -> usize { + NonZeroUsize::new(value).expect(msg); + value +} + +/// Convolution options. +#[derive(Debug, Clone, Hash, PartialEq, Eq)] +pub struct ConvOptions { + /// Stride (non-zero). + pub stride: [usize; N], + + /// Padding. + pub padding: [usize; N], + + /// Dilation (non-zero). + pub dilation: [usize; N], + + /// Groups (non-zero). + pub groups: usize, +} + +impl ConvOptions { + /// Constructs a new `ConvOptions`. + pub fn new( + stride: [usize; N], + padding: [usize; N], + dilation: [usize; N], + groups: usize, + ) -> Self { + Self { + stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")), + padding, + dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")), + groups: check_nonzero(groups, "groups must be non-zero"), + } + } +} + +/// Convolution options with support for asymmetric padding. +/// +/// Wraps [`ConvOptions`] (which represents symmetric padding for the backend op) +/// and adds optional asymmetric padding. When asymmetric padding is specified, +/// the functional convolution layer applies an explicit pad operation before +/// dispatching to the backend. +/// +/// Implements `From>` for backward compatibility. +#[derive(Debug, Clone)] +pub struct PaddedConvOptions { + /// The underlying convolution options for the backend. + pub options: ConvOptions, + /// Padding at the end of each dimension (e.g., bottom/right for 2D). + /// If `None`, padding is symmetric (same as `options.padding`). + /// If `Some`, specifies different end-padding per dimension. + pub padding_end: Option<[usize; N]>, +} + +impl PaddedConvOptions { + /// Creates options with asymmetric padding. + /// + /// `padding_start` is stored in `ConvOptions::padding`. + /// `padding_end` specifies the end padding per dimension. + pub fn asymmetric( + stride: [usize; N], + padding_start: [usize; N], + padding_end: [usize; N], + dilation: [usize; N], + groups: usize, + ) -> Self { + let options = ConvOptions::new(stride, padding_start, dilation, groups); + if padding_start == padding_end { + Self { + options, + padding_end: None, + } + } else { + Self { + options, + padding_end: Some(padding_end), + } + } + } + + /// Returns true if padding is asymmetric. + pub fn is_asymmetric(&self) -> bool { + self.padding_end.is_some() + } +} + +impl From> for PaddedConvOptions { + fn from(options: ConvOptions) -> Self { + Self { + options, + padding_end: None, + } + } +} + +/// Deformable convolution options. +#[derive(Debug, Clone, Hash, PartialEq, Eq)] +pub struct DeformConvOptions { + /// Stride (non-zero). + pub stride: [usize; N], + + /// Padding. + pub padding: [usize; N], + + /// Dilation (non-zero). + pub dilation: [usize; N], + + /// Weight Groups (non-zero). + pub weight_groups: usize, + + /// Offset Groups (non-zero). + pub offset_groups: usize, +} + +impl DeformConvOptions { + /// Constructs a new `DeformConvOptions`. + pub fn new( + stride: [usize; N], + padding: [usize; N], + dilation: [usize; N], + weight_groups: usize, + offset_groups: usize, + ) -> Self { + Self { + stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")), + padding, + dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")), + weight_groups: check_nonzero(weight_groups, "weight groups must be non-zero"), + offset_groups: check_nonzero(offset_groups, "offset groups must be non-zero"), + } + } +} + +/// Transposed convolution options. +#[derive(Debug, Clone, Hash, PartialEq, Eq)] +pub struct ConvTransposeOptions { + /// Stride (non-zero). + pub stride: [usize; N], + + /// Padding. + pub padding: [usize; N], + + /// Padding out. + pub padding_out: [usize; N], + + /// Dilation (non-zero). + pub dilation: [usize; N], + + /// Groups (non-zero). + pub groups: usize, +} + +impl ConvTransposeOptions { + /// Constructs a new `ConvTransposeOptions`. + pub fn new( + stride: [usize; N], + padding: [usize; N], + padding_out: [usize; N], + dilation: [usize; N], + groups: usize, + ) -> Self { + Self { + stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")), + padding, + padding_out, + dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")), + groups: check_nonzero(groups, "groups must be non-zero"), + } + } +} + +/// Unfold operation options. +#[derive(Debug, Clone)] +pub struct UnfoldOptions { + /// The number of positions to slide over the input tensor in each dimension. + /// A stride of `[1, 1]` will slide the kernel one pixel at a time. + pub stride: [usize; 2], + + /// The number of zero-padding pixels added to each side of the input tensor in each dimension. + pub padding: [usize; 2], + + /// The spacing between the blocks (patches) in the original input tensor. + pub dilation: [usize; 2], +} + +impl UnfoldOptions { + /// Constructs a new `UnfoldOptions`. + pub fn new(stride: [usize; 2], padding: [usize; 2], dilation: [usize; 2]) -> Self { + Self { + stride: stride.map(|s| check_nonzero(s, "stride must be non-zero")), + padding, + dilation: dilation.map(|d| check_nonzero(d, "dilation must be non-zero")), + } + } +} + +/// Algorithm used. +#[derive(new, Debug, Clone, serde::Deserialize, serde::Serialize)] +pub enum InterpolateMode { + /// Nearest-neighbor floor interpolation. + /// Matches the legacy behavior of OpenCV’s INTER_NEAREST. It results in a bottom-right shift when resizing. + /// + Nearest, + + /// Nearest-neighbor exact interpolation. + /// + NearestExact, + + /// Bilinear interpolation. + /// + Bilinear, + + /// Bicubic interpolation. + /// + Bicubic, + + /// Lanczos3 interpolation (6-tap sinc-based filter). + /// + Lanczos3, +} + +/// Interpolation options. +#[derive(Debug, Clone)] +pub struct InterpolateOptions { + /// Algorithm used. + pub mode: InterpolateMode, + /// If `true`, the input and output tensors are aligned by their corner pixels. + /// If `false`, half-pixel coordinate mapping is used instead. + pub align_corners: bool, +} + +impl InterpolateOptions { + /// Create new interpolate options with the given mode. + /// Defaults to `align_corners = true`. + pub fn new(mode: InterpolateMode) -> Self { + Self { + mode, + align_corners: true, + } + } + + /// Set align_corners. + pub fn with_align_corners(mut self, align_corners: bool) -> Self { + self.align_corners = align_corners; + self + } +} + +/// Padding mode for grid sampling when coordinates are out of bounds. +/// +/// Matches PyTorch's `padding_mode` parameter in `grid_sample`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Deserialize, serde::Serialize)] +pub enum GridSamplePaddingMode { + /// Fill with zeros for out-of-bounds coordinates. + #[default] + Zeros, + /// Clamp coordinates to the border (use nearest edge value). + Border, + /// Reflect coordinates at the boundary. + Reflection, +} + +/// Options for grid sampling operations. +#[derive(Debug, Clone)] +pub struct GridSampleOptions { + /// Interpolation mode (bilinear, nearest, or bicubic). + pub mode: InterpolateMode, + /// Padding mode for out-of-bounds coordinates. + pub padding_mode: GridSamplePaddingMode, + /// If `true`, grid values of -1 and 1 correspond to the corner pixels. + /// If `false`, they correspond to the corner points of the corner pixels + /// (i.e., -1 maps to -0.5 and 1 maps to size - 0.5 in pixel coordinates). + pub align_corners: bool, +} + +impl Default for GridSampleOptions { + fn default() -> Self { + Self { + mode: InterpolateMode::Bilinear, + padding_mode: GridSamplePaddingMode::Zeros, + align_corners: false, + } + } +} + +impl From for GridSampleOptions { + fn from(value: InterpolateMode) -> Self { + GridSampleOptions::new(value) + } +} + +impl GridSampleOptions { + /// Create new grid sample options with the given interpolation mode. + /// + /// Uses default values for padding_mode (Zeros) and align_corners (false). + pub fn new(mode: InterpolateMode) -> Self { + Self { + mode, + ..Default::default() + } + } + + /// Set the padding mode. + pub fn with_padding_mode(mut self, padding_mode: GridSamplePaddingMode) -> Self { + self.padding_mode = padding_mode; + self + } + + /// Set align_corners. + pub fn with_align_corners(mut self, align_corners: bool) -> Self { + self.align_corners = align_corners; + self + } +} + +/// Padding mode for tensor pad operations. +/// +/// Defines how values are filled when padding a tensor beyond its original boundaries. +/// Padding can be applied to any dimension of a tensor. +/// +/// # Modes +/// +/// - [`Constant`](PadMode::Constant): Fill with a specified value (default: 0.0) +/// - [`Reflect`](PadMode::Reflect): Mirror values at boundary, excluding edge (requires padding < dim_size) +/// - [`Edge`](PadMode::Edge): Replicate boundary values +#[derive(Debug, Clone, Copy, PartialEq, serde::Deserialize, serde::Serialize)] +pub enum PadMode { + /// Fill padded regions with a constant value. + /// + /// # Example + /// For tensor `[1, 2, 3]` with padding 2 on the left and value 0: + /// Result: `[0, 0, 1, 2, 3]` + Constant(f32), + + /// Reflect values at the boundary, excluding the edge value. + /// + /// Padding must be less than the dimension size (i.e., `padding < dim_size`). + /// + /// # Example + /// For tensor `[1, 2, 3, 4]` with padding 2 on the left: + /// Result: `[3, 2, 1, 2, 3, 4]` (reflects from index 1, not 0) + Reflect, + + /// Replicate the edge values. + /// + /// # Example + /// For tensor `[1, 2, 3, 4]` with padding 2 on the left: + /// Result: `[1, 1, 1, 2, 3, 4]` + Edge, +} + +impl Default for PadMode { + fn default() -> Self { + PadMode::Constant(0.0) + } +} + +impl From for PadMode { + fn from(value: E) -> Self { + PadMode::Constant(value.elem()) + } +} + +/// Options for the attention module. +#[derive(Debug, Clone, Copy, Default, PartialEq, serde::Deserialize, serde::Serialize)] +pub struct AttentionModuleOptions { + /// Custom scale factor applied to QK^T. When `None`, defaults to `1/sqrt(head_dim)`. + pub scale: Option, + + /// Soft capping applied before softmax: `softcap * tanh(scores / softcap)`. + /// Used by Gemma-2 and similar models. Must be positive when set. + pub softcap: Option, + + /// When `true`, applies causal (autoregressive) masking so that each query position + /// can only attend to key positions at or before it. This is more efficient than + /// passing an explicit lower-triangular bool mask because backends can use optimized + /// kernel paths (e.g. flash attention with causal mode). + pub is_causal: bool, +} + +/// Computation to be used to update the existing values in indexed assignment operations (scatter/select). +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub enum IndexingUpdateOp { + /// Overwrite existing values. + Assign, + /// Performs an addition. + Add, + /// Multiply existing values. + Mul, + /// Take element-wise minimum. + Min, + /// Take element-wise maximum. + Max, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + #[should_panic = "stride must be non-zero"] + fn conv_options_stride_zero() { + let _opt = ConvOptions::new([0, 1], [0, 0], [1, 1], 1); + } + + #[test] + #[should_panic = "dilation must be non-zero"] + fn conv_options_dilation_zero() { + let _opt = ConvOptions::new([1, 1], [0, 0], [0, 0], 1); + } + + #[test] + #[should_panic = "groups must be non-zero"] + fn conv_options_groups_zero() { + let _opt = ConvOptions::new([1, 1], [0, 0], [1, 1], 0); + } + + #[test] + #[should_panic = "stride must be non-zero"] + fn conv_transpose_options_stride_zero() { + let _opt = ConvTransposeOptions::new([0, 1], [0, 0], [0, 0], [1, 1], 1); + } + + #[test] + #[should_panic = "dilation must be non-zero"] + fn conv_transpose_options_dilation_zero() { + let _opt = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [0, 0], 1); + } + + #[test] + #[should_panic = "groups must be non-zero"] + fn conv_transpose_options_groups_zero() { + let _opt = ConvTransposeOptions::new([1, 1], [0, 0], [0, 0], [1, 1], 0); + } + + #[test] + #[should_panic = "stride must be non-zero"] + fn deform_conv_options_stride_zero() { + let _opt = DeformConvOptions::new([0, 1], [0, 0], [1, 1], 1, 1); + } + + #[test] + #[should_panic = "dilation must be non-zero"] + fn deform_conv_options_dilation_zero() { + let _opt = DeformConvOptions::new([1, 1], [0, 0], [0, 0], 1, 1); + } + + #[test] + #[should_panic = "weight groups must be non-zero"] + fn deform_conv_options_weights_groups_zero() { + let _opt = DeformConvOptions::new([1, 1], [0, 0], [1, 1], 0, 1); + } + + #[test] + #[should_panic = "offset groups must be non-zero"] + fn deform_conv_options_offset_groups_zero() { + let _opt = DeformConvOptions::new([1, 1], [0, 0], [1, 1], 1, 0); + } + + #[test] + #[should_panic = "stride must be non-zero"] + fn unfold_options_stride_zero() { + let _opt = UnfoldOptions::new([0, 1], [0, 0], [1, 1]); + } + + #[test] + #[should_panic = "dilation must be non-zero"] + fn unfold_options_dilation_zero() { + let _opt = UnfoldOptions::new([1, 1], [0, 0], [0, 0]); + } +} diff --git a/crates/burn-std/src/runtime_kind.rs b/crates/burn-std/src/runtime_kind.rs new file mode 100644 index 0000000..79e7f0d --- /dev/null +++ b/crates/burn-std/src/runtime_kind.rs @@ -0,0 +1,62 @@ +//! Runtime kind of the host program. +//! +//! Some backend decisions depend not on the device but on *how the host program itself is +//! being driven* — whether `main` runs on an asynchronous runtime, a synchronous +//! thread-based runtime, or a restricted no-std environment. This module stores that kind +//! in a process global so backends can read it and adapt their behavior. +//! +//! The canonical example is tensor readback (`into_data`): deferring the device→host copy +//! lazily is fine under a sync/threaded runtime (a later blocking read just parks a thread), +//! but under an async runtime the same blocking read parks an executor worker and starves +//! the runtime, so the read must materialize eagerly instead. + +use core::sync::atomic::{AtomicU8, Ordering}; + +/// How the host program is being driven. +/// +/// Set once near program start with [`set_runtime_kind`]; read anywhere with +/// [`runtime_kind`]. Defaults to [`RuntimeKind::Sync`]. +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum RuntimeKind { + /// Synchronous, thread-based runtime (the default). + /// + /// Blocking readbacks are fine here, so tensor reads may defer the device→host copy + /// lazily and materialize it on first access. + #[default] + Sync = 0, + /// Asynchronous runtime (e.g. tokio). + /// + /// Blocking a runtime worker starves the executor, so tensor reads must materialize + /// eagerly inside the awaited future rather than deferring a blocking copy. + Async = 1, + /// Restricted no-std environment. + NoStd = 2, +} + +impl RuntimeKind { + fn from_u8(value: u8) -> Self { + match value { + 1 => RuntimeKind::Async, + 2 => RuntimeKind::NoStd, + _ => RuntimeKind::Sync, + } + } +} + +static RUNTIME_KIND: AtomicU8 = AtomicU8::new(RuntimeKind::Sync as u8); + +/// Declare the [kind](RuntimeKind) of runtime the host program is running on. +/// +/// Intended to be called once near program start (e.g. when a remote server declares that +/// it hosts the backend on an async runtime). Backends read it via [`runtime_kind`]. +pub fn set_runtime_kind(kind: RuntimeKind) { + RUNTIME_KIND.store(kind as u8, Ordering::Relaxed); +} + +/// Return the currently declared [kind](RuntimeKind) of runtime the host program runs on. +/// +/// Defaults to [`RuntimeKind::Sync`] until [`set_runtime_kind`] is called. +pub fn runtime_kind() -> RuntimeKind { + RuntimeKind::from_u8(RUNTIME_KIND.load(Ordering::Relaxed)) +} diff --git a/crates/burn-std/src/tensor/container.rs b/crates/burn-std/src/tensor/container.rs new file mode 100644 index 0000000..b092e28 --- /dev/null +++ b/crates/burn-std/src/tensor/container.rs @@ -0,0 +1,86 @@ +use alloc::boxed::Box; +use core::any::Any; + +#[cfg(not(feature = "std"))] +use alloc::vec::Vec; +#[cfg(not(feature = "std"))] +use hashbrown::HashMap; + +#[cfg(feature = "std")] +use std::collections::HashMap; + +/// Contains tensor of arbitrary dimension. +#[derive(Debug)] +pub struct TensorContainer { + tensors: HashMap>, +} + +impl Default for TensorContainer +where + ID: core::hash::Hash + PartialEq + Eq + core::fmt::Debug, +{ + fn default() -> Self { + Self::new() + } +} + +impl TensorContainer +where + ID: core::hash::Hash + PartialEq + Eq + core::fmt::Debug, +{ + /// Create an empty container. + pub fn new() -> Self { + Self { + tensors: HashMap::new(), + } + } + + /// Get a tensor with the given ID. + pub fn get(&self, id: &ID) -> Option { + let grad = self.tensors.get(id)?; + + let tensor = grad.downcast_ref::().unwrap(); + + Some(tensor.clone()) + } + + /// Get a mutable reference to the tensor with the given ID. + pub fn get_mut_ref(&mut self, id: &ID) -> Option<&mut T> { + let grad = self.tensors.get_mut(id)?; + + let tensor = grad.downcast_mut::().unwrap(); + + Some(tensor) + } + + /// Register a new tensor for the given ID. + /// + /// # Notes + /// + /// If a tensor is already registered for the given ID, it will be replaced. + pub fn register(&mut self, id: ID, value: T) { + self.tensors.insert(id, Box::new(value)); + } + + /// Remove a tensor for the given ID and returns it. + pub fn remove(&mut self, id: &ID) -> Option { + self.tensors + .remove(id) + .map(|item| *item.downcast::().unwrap()) + } + + /// The number of tensors registered. + pub fn len(&self) -> usize { + self.tensors.len() + } + + /// If any tensor is contained. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Get id of every tensor in the container + pub fn ids(&self) -> Vec<&ID> { + self.tensors.keys().collect() + } +} diff --git a/crates/burn-std/src/tensor/dtype.rs b/crates/burn-std/src/tensor/dtype.rs new file mode 100644 index 0000000..c3fa02b --- /dev/null +++ b/crates/burn-std/src/tensor/dtype.rs @@ -0,0 +1,408 @@ +//! Tensor data type. + +use serde::{Deserialize, Serialize}; + +use crate::tensor::quantization::{QuantScheme, QuantStore, QuantValue}; +use crate::{bf16, f16}; + +#[allow(missing_docs)] +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Serialize, Deserialize)] +pub enum DType { + F64, + F32, + Flex32, + F16, + BF16, + I64, + I32, + I16, + I8, + U64, + U32, + U16, + U8, + Bool(BoolStore), + QFloat(QuantScheme), +} + +// Conversions between `DType` and cubecl's `ElemType` / `StorageType` are +// intentionally not implemented here so that `burn-std` does not depend on +// `cubecl`. Backend code that needs them should call the named functions in +// `burn_backend::cubecl` (e.g. `elem_type_to_dtype`, `dtype_to_elem_type`, +// `dtype_to_storage_type`). + +impl DType { + /// Returns the size of a type in bytes. + pub const fn size(&self) -> usize { + match self { + DType::F64 => core::mem::size_of::(), + DType::F32 => core::mem::size_of::(), + DType::Flex32 => core::mem::size_of::(), + DType::F16 => core::mem::size_of::(), + DType::BF16 => core::mem::size_of::(), + DType::I64 => core::mem::size_of::(), + DType::I32 => core::mem::size_of::(), + DType::I16 => core::mem::size_of::(), + DType::I8 => core::mem::size_of::(), + DType::U64 => core::mem::size_of::(), + DType::U32 => core::mem::size_of::(), + DType::U16 => core::mem::size_of::(), + DType::U8 => core::mem::size_of::(), + DType::Bool(store) => match store { + BoolStore::Native => core::mem::size_of::(), + BoolStore::U8 => core::mem::size_of::(), + BoolStore::U32 => core::mem::size_of::(), + }, + DType::QFloat(scheme) => match scheme.store { + QuantStore::Native => match scheme.value { + QuantValue::Q8F | QuantValue::Q8S => core::mem::size_of::(), + // e2m1 native is automatically packed by the kernels, so the actual storage is + // 8 bits wide. + QuantValue::E4M3 | QuantValue::E5M2 | QuantValue::E2M1 => { + core::mem::size_of::() + } + QuantValue::Q4F | QuantValue::Q4S | QuantValue::Q2F | QuantValue::Q2S => { + // Sub-byte values have fractional size + 0 + } + }, + QuantStore::PackedU32(_) => core::mem::size_of::(), + QuantStore::PackedNative(_) => match scheme.value { + QuantValue::E2M1 => core::mem::size_of::(), + _ => 0, + }, + }, + } + } + /// Returns true if the data type is a floating point type. + pub fn is_float(&self) -> bool { + matches!( + self, + DType::F64 | DType::F32 | DType::Flex32 | DType::F16 | DType::BF16 + ) + } + /// Returns true if the data type is a signed integer type. + pub fn is_int(&self) -> bool { + matches!(self, DType::I64 | DType::I32 | DType::I16 | DType::I8) + } + /// Returns true if the data type is an unsigned integer type. + pub fn is_uint(&self) -> bool { + matches!(self, DType::U64 | DType::U32 | DType::U16 | DType::U8) + } + + /// Returns true if the data type is a boolean type + pub fn is_bool(&self) -> bool { + matches!(self, DType::Bool(_)) + } + + /// Returns float precision info if this is a float dtype, `None` otherwise. + /// + /// Analogous to `torch.finfo(dtype)` or `numpy.finfo(dtype)`. + pub const fn finfo(&self) -> Option { + match self { + DType::F64 => Some(FloatDType::F64.finfo()), + DType::F32 => Some(FloatDType::F32.finfo()), + DType::Flex32 => Some(FloatDType::Flex32.finfo()), + DType::F16 => Some(FloatDType::F16.finfo()), + DType::BF16 => Some(FloatDType::BF16.finfo()), + _ => None, + } + } + + /// Returns the data type name. + pub fn name(&self) -> &'static str { + match self { + DType::F64 => "f64", + DType::F32 => "f32", + DType::Flex32 => "flex32", + DType::F16 => "f16", + DType::BF16 => "bf16", + DType::I64 => "i64", + DType::I32 => "i32", + DType::I16 => "i16", + DType::I8 => "i8", + DType::U64 => "u64", + DType::U32 => "u32", + DType::U16 => "u16", + DType::U8 => "u8", + DType::Bool(store) => match store { + BoolStore::Native => "bool", + BoolStore::U8 => "bool(u8)", + BoolStore::U32 => "bool(u32)", + }, + DType::QFloat(_) => "qfloat", + } + } +} + +#[allow(missing_docs)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)] +pub enum FloatDType { + F64, + F32, + Flex32, + F16, + BF16, +} + +/// Numerical precision properties for a floating-point dtype. +/// +/// Equivalent to NumPy's `finfo` / PyTorch's `torch.finfo`. All values are +/// widened to `f64` so they can be inspected without knowing the concrete +/// element type at compile time. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct FloatInfo { + /// Machine epsilon: smallest value such that `1.0 + epsilon != 1.0`. + pub epsilon: f64, + /// Largest representable finite value. + pub max: f64, + /// Most negative representable finite value. + pub min: f64, + /// Smallest positive normal value. + pub min_positive: f64, +} + +impl FloatDType { + /// Returns numerical precision properties for this float dtype. + /// + /// Analogous to `torch.finfo(dtype)` or `numpy.finfo(dtype)`. + pub const fn finfo(self) -> FloatInfo { + match self { + FloatDType::F64 => FloatInfo { + epsilon: f64::EPSILON, + max: f64::MAX, + min: f64::MIN, + min_positive: f64::MIN_POSITIVE, // ~2.225e-308 + }, + FloatDType::F32 => FloatInfo { + epsilon: f32::EPSILON as f64, + max: f32::MAX as f64, + min: f32::MIN as f64, + min_positive: f32::MIN_POSITIVE as f64, // ~1.175e-38 + }, + // Flex32 stores as f32 but computes at reduced (f16-like) precision. + // Use f16 precision limits so stability code stays safe. + FloatDType::Flex32 => FloatInfo { + epsilon: f16::EPSILON.to_f64_const(), + max: f16::MAX.to_f64_const(), + min: f16::MIN.to_f64_const(), + min_positive: f16::MIN_POSITIVE.to_f64_const(), // ~6.104e-5 + }, + FloatDType::F16 => FloatInfo { + epsilon: f16::EPSILON.to_f64_const(), + max: f16::MAX.to_f64_const(), + min: f16::MIN.to_f64_const(), + min_positive: f16::MIN_POSITIVE.to_f64_const(), // ~6.104e-5 + }, + FloatDType::BF16 => FloatInfo { + epsilon: bf16::EPSILON.to_f64_const(), + max: bf16::MAX.to_f64_const(), + min: bf16::MIN.to_f64_const(), + min_positive: bf16::MIN_POSITIVE.to_f64_const(), // ~1.175e-38 + }, + } + } +} + +impl From for FloatDType { + fn from(value: DType) -> Self { + match value { + DType::F64 => FloatDType::F64, + DType::F32 => FloatDType::F32, + DType::Flex32 => FloatDType::Flex32, + DType::F16 => FloatDType::F16, + DType::BF16 => FloatDType::BF16, + _ => panic!("Expected float data type, got {value:?}"), + } + } +} + +impl From for DType { + fn from(value: FloatDType) -> Self { + match value { + FloatDType::F64 => DType::F64, + FloatDType::F32 => DType::F32, + FloatDType::Flex32 => DType::Flex32, + FloatDType::F16 => DType::F16, + FloatDType::BF16 => DType::BF16, + } + } +} + +#[allow(missing_docs)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)] +pub enum IntDType { + I64, + I32, + I16, + I8, + U64, + U32, + U16, + U8, +} + +impl From for IntDType { + fn from(value: DType) -> Self { + match value { + DType::I64 => IntDType::I64, + DType::I32 => IntDType::I32, + DType::I16 => IntDType::I16, + DType::I8 => IntDType::I8, + DType::U64 => IntDType::U64, + DType::U32 => IntDType::U32, + DType::U16 => IntDType::U16, + DType::U8 => IntDType::U8, + _ => panic!("Expected int data type, got {value:?}"), + } + } +} + +impl From for DType { + fn from(value: IntDType) -> Self { + match value { + IntDType::I64 => DType::I64, + IntDType::I32 => DType::I32, + IntDType::I16 => DType::I16, + IntDType::I8 => DType::I8, + IntDType::U64 => DType::U64, + IntDType::U32 => DType::U32, + IntDType::U16 => DType::U16, + IntDType::U8 => DType::U8, + } + } +} + +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Serialize, Deserialize)] +/// Data type used to store boolean values. +pub enum BoolStore { + /// Stored as native boolean type (e.g. `bool`). + Native, + /// Stored as 8-bit unsigned integer. + U8, + /// Stored as 32-bit unsigned integer. + U32, +} + +/// Boolean dtype. +/// +/// This is currently an alias to [`BoolStore`], since it only varies by the storage representation. +pub type BoolDType = BoolStore; + +#[allow(deprecated)] +impl From for BoolDType { + fn from(value: DType) -> Self { + match value { + DType::Bool(store) => match store { + BoolStore::Native => BoolDType::Native, + BoolStore::U8 => BoolDType::U8, + BoolStore::U32 => BoolDType::U32, + }, + // For compat BoolElem associated type + DType::U8 => BoolDType::U8, + DType::U32 => BoolDType::U32, + _ => panic!("Expected bool data type, got {value:?}"), + } + } +} + +impl From for DType { + fn from(value: BoolDType) -> Self { + match value { + BoolDType::Native => DType::Bool(BoolStore::Native), + BoolDType::U8 => DType::Bool(BoolStore::U8), + BoolDType::U32 => DType::Bool(BoolStore::U32), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn finfo_f32() { + let info = FloatDType::F32.finfo(); + assert_eq!(info.epsilon, f32::EPSILON as f64); + assert_eq!(info.max, f32::MAX as f64); + assert_eq!(info.min, f32::MIN as f64); + assert_eq!(info.min_positive, f32::MIN_POSITIVE as f64); + } + + #[test] + fn finfo_f64() { + let info = FloatDType::F64.finfo(); + assert_eq!(info.epsilon, f64::EPSILON); + assert_eq!(info.max, f64::MAX); + assert_eq!(info.min, f64::MIN); + assert_eq!(info.min_positive, f64::MIN_POSITIVE); + } + + #[test] + fn finfo_f16() { + let info = FloatDType::F16.finfo(); + assert_eq!(info.epsilon, f16::EPSILON.to_f64_const()); + assert!(info.epsilon > 0.0); + assert!(info.min_positive > 0.0); + // f16 epsilon is much larger than f32 + assert!(info.epsilon > FloatDType::F32.finfo().epsilon); + } + + #[test] + fn finfo_bf16() { + let info = FloatDType::BF16.finfo(); + assert_eq!(info.epsilon, bf16::EPSILON.to_f64_const()); + assert!(info.epsilon > 0.0); + assert!(info.min_positive > 0.0); + // bf16 epsilon is larger than f32 (fewer mantissa bits) + assert!(info.epsilon > FloatDType::F32.finfo().epsilon); + } + + #[test] + fn finfo_flex32_uses_f16_limits() { + let flex = FloatDType::Flex32.finfo(); + let f16_info = FloatDType::F16.finfo(); + assert_eq!(flex.epsilon, f16_info.epsilon); + assert_eq!(flex.min_positive, f16_info.min_positive); + } + + #[test] + fn dtype_finfo_delegates_to_float_dtype() { + assert_eq!(DType::F32.finfo(), Some(FloatDType::F32.finfo())); + assert_eq!(DType::F64.finfo(), Some(FloatDType::F64.finfo())); + assert_eq!(DType::F16.finfo(), Some(FloatDType::F16.finfo())); + assert_eq!(DType::BF16.finfo(), Some(FloatDType::BF16.finfo())); + assert_eq!(DType::Flex32.finfo(), Some(FloatDType::Flex32.finfo())); + } + + #[test] + fn dtype_finfo_returns_none_for_non_float() { + assert!(DType::I32.finfo().is_none()); + assert!(DType::U8.finfo().is_none()); + assert!(DType::Bool(BoolStore::Native).finfo().is_none()); + } + + #[test] + fn finfo_invariants() { + for dtype in [ + FloatDType::F64, + FloatDType::F32, + FloatDType::F16, + FloatDType::BF16, + FloatDType::Flex32, + ] { + let info = dtype.finfo(); + assert!(info.epsilon > 0.0, "{dtype:?}: epsilon must be positive"); + assert!( + info.min_positive > 0.0, + "{dtype:?}: min_positive must be positive" + ); + assert!(info.max > 0.0, "{dtype:?}: max must be positive"); + assert!(info.min < 0.0, "{dtype:?}: min must be negative"); + assert!( + info.max > info.min_positive, + "{dtype:?}: max > min_positive" + ); + } + } +} diff --git a/crates/burn-std/src/tensor/matmul.rs b/crates/burn-std/src/tensor/matmul.rs new file mode 100644 index 0000000..a517f28 --- /dev/null +++ b/crates/burn-std/src/tensor/matmul.rs @@ -0,0 +1,352 @@ +use crate::tensor::{Metadata, Shape, Strides}; + +/// Accelerated matmul tiles consume rows in chunks of this size; a row count +/// that is a multiple of it already tiles cleanly. +const ROW_TILE: usize = 32; + +/// The action to take for a batched matmul with a broadcast rhs. +#[derive(Debug, PartialEq, Eq, Clone, Copy)] +pub enum MatmulTransformAction { + /// Launch the matmul as-is. + Keep, + /// Fold every lhs/out batch dim into the rows: `[.., b, m, k] @ [.., 1, k, n]` + /// runs as one `[b*m, k] @ [k, n]` matmul against the shared rhs, instead of + /// `b` matmuls that each re-read it. + MergeBatches { + /// The merged row count (batches × rows). + rows: usize, + }, +} + +impl MatmulTransformAction { + /// Apply the action to one merged operand (lhs or out): fold the batch dims + /// into the rows, in place. A pure metadata rewrite — the buffer is shared. + /// + /// Requires batch-contiguous rows (see [MatmulTransformAnalysis]); the merge + /// keeps the rank, with every batch dim set to 1. + pub fn apply(&self, meta: &mut Metadata) { + let rows = match self { + MatmulTransformAction::Keep => return, + MatmulTransformAction::MergeBatches { rows } => *rows, + }; + + let rank = meta.rank(); + let stride_rows = merged_row_stride(meta.shape(), meta.strides()) + .expect("The action requires batch-contiguous rows"); + + for i in 0..rank - 2 { + meta.shape_mut()[i] = 1; + meta.strides_mut()[i] = rows * stride_rows; + } + meta.shape_mut()[rank - 2] = rows; + meta.strides_mut()[rank - 2] = stride_rows; + } +} + +/// The facts of a batched matmul `lhs @ rhs` relevant to folding its batches +/// into its rows. +#[derive(Debug, PartialEq, Eq, Clone, Copy)] +pub struct MatmulTransformAnalysis { + /// Product of the lhs batch dims. + batches: usize, + /// Rows of one batch (`m`). + rows: usize, + /// Columns of the output (`n`). + cols: usize, + /// The rhs is shared by every batch, and each operand's rows advance with a + /// single stride across batch boundaries (no holes), so the fold is a pure + /// view. + mergeable: bool, +} + +impl MatmulTransformAnalysis { + /// Analyze from shapes alone: the operands are assumed batch-contiguous, as + /// holds for freshly materialized tensors. Use [Self::from_metadata] when + /// the layouts are known. + /// + /// The rhs may have a lower rank than the lhs (implicit broadcast). + pub fn from_shapes(lhs: &Shape, rhs: &Shape) -> Self { + Self::new(lhs, rhs, true) + } + + /// Analyze from full metadata: on top of the shape facts, the lhs and out + /// rows must advance with a single stride across batch boundaries. + pub fn from_metadata(lhs: &Metadata, rhs: &Metadata, out: &Metadata) -> Self { + let mergeable = merged_row_stride(lhs.shape(), lhs.strides()).is_some() + && merged_row_stride(out.shape(), out.strides()).is_some(); + + Self::new(lhs.shape(), rhs.shape(), mergeable) + } + + fn new(lhs: &Shape, rhs: &Shape, rows_contiguous: bool) -> Self { + let rank = lhs.num_dims(); + let rank_rhs = rhs.num_dims(); + + let batches = lhs[..rank - 2].iter().product(); + let rows = lhs[rank - 2]; + let cols = rhs[rank_rhs - 1]; + + // Every batch dim the rhs actually has must be broadcast. + let rhs_shared = rhs[..rank_rhs - 2].iter().all(|&dim| dim == 1); + + Self { + batches, + rows, + cols, + mergeable: rhs_shared && rows_contiguous, + } + } +} + +/// Decides the [action](MatmulTransformAction) to take for a batched matmul, +/// given its [analysis](MatmulTransformAnalysis). +#[derive(Debug, Default, Clone, Copy)] +pub enum MatmulTransformPolicy { + /// Merge the batches into the rows when the fold is a pure view and the + /// merged problem tiles better: per-batch rows that already fill row tiles + /// stay batched (a batched matmul beats one big matmul at that scale), tiny + /// row counts merge when that brings the output closer to square. + #[default] + BetterTiling, + /// Never transform. + Never, +} + +impl MatmulTransformPolicy { + /// The action to take for a matmul with the given analysis. + pub fn action(&self, analysis: &MatmulTransformAnalysis) -> MatmulTransformAction { + match self { + MatmulTransformPolicy::Never => MatmulTransformAction::Keep, + MatmulTransformPolicy::BetterTiling => { + if !analysis.mergeable || analysis.batches == 1 { + return MatmulTransformAction::Keep; + } + + // Rows already fill the tiles: keep the batched form. + if analysis.rows.is_multiple_of(ROW_TILE) { + return MatmulTransformAction::Keep; + } + + let rows = analysis.batches * analysis.rows; + + // The merge must bring the output closer to square. + if !squarer(rows, analysis.rows, analysis.cols) { + return MatmulTransformAction::Keep; + } + + MatmulTransformAction::MergeBatches { rows } + } + } + } +} + +/// Whether an output of `rows_new` x `cols` is closer to square than +/// `rows` x `cols`. +fn squarer(rows_new: usize, rows: usize, cols: usize) -> bool { + // min(a, c) / max(a, c) > min(b, c) / max(b, c), cross-multiplied to stay + // in integers. + rows_new.min(cols) * rows.max(cols) > rows.min(cols) * rows_new.max(cols) +} + +/// The single stride with which rows advance across batch boundaries, when the +/// batch and row dims form a hole-free chain — the requirement for folding the +/// batches into the rows as a pure view. Dims of size 1 carry arbitrary strides +/// and never anchor the chain. +fn merged_row_stride(shape: &Shape, strides: &Strides) -> Option { + let rank = shape.num_dims(); + let mut chained: Option<(usize, usize)> = None; + + for i in (0..rank - 1).rev() { + if shape[i] == 1 { + continue; + } + match chained { + None => chained = Some((strides[i], strides[i] * shape[i])), + Some((row_stride, expected)) => { + if strides[i] != expected { + return None; + } + chained = Some((row_stride, strides[i] * shape[i])); + } + } + } + + match chained { + Some((row_stride, _)) => Some(row_stride), + // Every merged dim is a unit dim: a single row, any stride works. + None => Some(shape[rank - 1] * strides[rank - 1]), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::shape; + + fn analysis(lhs: &[usize], rhs: &[usize]) -> MatmulTransformAnalysis { + MatmulTransformAnalysis::from_shapes(&Shape::from(lhs), &Shape::from(rhs)) + } + + fn action(analysis: &MatmulTransformAnalysis) -> MatmulTransformAction { + MatmulTransformPolicy::default().action(analysis) + } + + #[test] + fn merges_batched_vec_mat() { + // A decode step: 16 broadcast vec-mats fold into one 16-row matmul. + let analysis = analysis(&[16, 1, 4096], &[4096, 14336]); + + assert_eq!( + action(&analysis), + MatmulTransformAction::MergeBatches { rows: 16 } + ); + } + + #[test] + fn keeps_row_tiled_batches() { + // A prefill step: 512 rows per batch already fill the row tiles; the + // batched matmul beats one big matmul. + let analysis = analysis(&[8, 512, 4096], &[4096, 14336]); + + assert_eq!(action(&analysis), MatmulTransformAction::Keep); + } + + #[test] + fn keeps_single_batch() { + let analysis = analysis(&[1, 1, 4096], &[4096, 14336]); + + assert_eq!(action(&analysis), MatmulTransformAction::Keep); + } + + #[test] + fn keeps_matrix() { + let analysis = analysis(&[16, 4096], &[4096, 14336]); + + assert_eq!(action(&analysis), MatmulTransformAction::Keep); + } + + #[test] + fn keeps_batched_rhs() { + // The rhs differs per batch: nothing is shared, nothing to fold. + let analysis = analysis(&[8, 1, 64], &[8, 64, 32]); + + assert_eq!(action(&analysis), MatmulTransformAction::Keep); + } + + #[test] + fn merges_with_broadcast_rhs_rank() { + let analysis = analysis(&[8, 1, 64], &[1, 64, 32]); + + assert_eq!( + action(&analysis), + MatmulTransformAction::MergeBatches { rows: 8 } + ); + } + + #[test] + fn keeps_when_merge_leaves_square() { + // 1000 rows tower over 64 cols; folding to 4000 only makes the output + // less square. + let analysis = analysis(&[4, 1000, 64], &[64, 64]); + + assert_eq!(action(&analysis), MatmulTransformAction::Keep); + } + + #[test] + fn merges_multiple_batch_dims() { + let analysis = analysis(&[2, 8, 1, 4096], &[4096, 14336]); + + assert_eq!( + action(&analysis), + MatmulTransformAction::MergeBatches { rows: 16 } + ); + } + + #[test] + fn never_policy_keeps() { + let analysis = analysis(&[16, 1, 4096], &[4096, 14336]); + + assert_eq!( + MatmulTransformPolicy::Never.action(&analysis), + MatmulTransformAction::Keep + ); + } + + #[test] + fn metadata_merges_pitched_rows() { + // Padded batches with one row each still fold: the merged matrix reads + // its rows with the single (pitched) stride 8192. + let lhs = Metadata::new(shape![16, 1, 4096], crate::strides![8192, 4096, 1]); + let rhs = Metadata::new(shape![1, 4096, 14336], crate::strides![0, 14336, 1]); + let out = Metadata::new(shape![16, 1, 14336], crate::strides![14336, 14336, 1]); + + let analysis = MatmulTransformAnalysis::from_metadata(&lhs, &rhs, &out); + + assert_eq!( + action(&analysis), + MatmulTransformAction::MergeBatches { rows: 16 } + ); + } + + #[test] + fn metadata_keeps_batch_holes() { + // Rows are contiguous within a batch but batches leave a gap: no single + // row stride can express the fold. + let lhs = Metadata::new(shape![4, 3, 8], crate::strides![48, 8, 1]); + let rhs = Metadata::new(shape![1, 8, 32], crate::strides![0, 32, 1]); + let out = Metadata::new(shape![4, 3, 32], crate::strides![96, 32, 1]); + + let analysis = MatmulTransformAnalysis::from_metadata(&lhs, &rhs, &out); + + assert_eq!(action(&analysis), MatmulTransformAction::Keep); + } + + #[test] + fn metadata_merges_contiguous() { + let lhs = Metadata::new(shape![16, 1, 4096], crate::strides![4096, 4096, 1]); + let rhs = Metadata::new(shape![1, 4096, 14336], crate::strides![0, 14336, 1]); + let out = Metadata::new(shape![16, 1, 14336], crate::strides![14336, 14336, 1]); + + let analysis = MatmulTransformAnalysis::from_metadata(&lhs, &rhs, &out); + + assert_eq!( + action(&analysis), + MatmulTransformAction::MergeBatches { rows: 16 } + ); + } + + #[test] + fn metadata_ignores_unit_dim_strides() { + // Size-1 dims carry arbitrary strides; they must not anchor the chain. + let lhs = Metadata::new(shape![16, 1, 4096], crate::strides![4096, 12345, 1]); + let rhs = Metadata::new(shape![1, 4096, 14336], crate::strides![0, 14336, 1]); + let out = Metadata::new(shape![16, 1, 14336], crate::strides![14336, 99999, 1]); + + let analysis = MatmulTransformAnalysis::from_metadata(&lhs, &rhs, &out); + + assert_eq!( + action(&analysis), + MatmulTransformAction::MergeBatches { rows: 16 } + ); + } + + #[test] + fn apply_folds_batches_in_place() { + let mut lhs = Metadata::new(shape![16, 1, 4096], crate::strides![4096, 4096, 1]); + + MatmulTransformAction::MergeBatches { rows: 16 }.apply(&mut lhs); + + assert_eq!(lhs.shape(), &shape![1, 16, 4096]); + assert_eq!(lhs.strides()[1], 4096); + assert_eq!(lhs.strides()[2], 1); + } + + #[test] + fn apply_keep_is_noop() { + let mut lhs = Metadata::new(shape![16, 1, 4096], crate::strides![4096, 4096, 1]); + + MatmulTransformAction::Keep.apply(&mut lhs); + + assert_eq!(lhs.shape(), &shape![16, 1, 4096]); + } +} diff --git a/crates/burn-std/src/tensor/mod.rs b/crates/burn-std/src/tensor/mod.rs new file mode 100644 index 0000000..a52718d --- /dev/null +++ b/crates/burn-std/src/tensor/mod.rs @@ -0,0 +1,336 @@ +/// Generic container for storing tensors keyed by an id. +pub mod container; +/// Tensor data type definitions. +pub mod dtype; +/// Batched matmul transformation utilities. +pub mod matmul; +/// Quantization data representation. +pub mod quantization; +/// Tensor shape utilities. +pub mod shape; +/// Tensor slicing utilities. +pub mod slice; + +pub use dtype::*; +pub use matmul::*; +pub use quantization::*; +pub use shape::*; +pub use slice::*; + +pub use cubecl_zspace::indexing::{self, *}; +pub use cubecl_zspace::{Strides, metadata::Metadata, strides}; + +/// Check if the current tensor is contiguous. +/// +/// A tensor is considered contiguous if its elements are stored in memory +/// such that the stride at position `k` is equal to the product of the shapes +/// of all dimensions greater than `k`. +/// +/// This means that strides increase as you move from the rightmost to the leftmost dimension. +pub fn is_contiguous(shape: &[usize], strides: &[usize]) -> bool { + if shape.is_empty() { + return true; + } + + for (&expected, &stride) in contiguous_strides(shape).iter().zip(strides) { + if expected != stride { + return false; + } + } + + true +} + +/// Computes the strides for a contiguous tensor with the given shape. +/// +/// In a contiguous row-major tensor, the stride for each dimension +/// equals the product of all dimension sizes to its right. +pub fn contiguous_strides(shape: &[usize]) -> Strides { + let mut strides = strides![0; shape.len()]; + let mut current = 1; + + for (i, &dim) in shape.iter().enumerate().rev() { + strides[i] = current; + current *= dim; + } + + strides +} + +/// The action to take for a reshape operation. +#[derive(Debug)] +pub enum ReshapeAction { + /// Updating the strides is sufficient to handle the reshape. + UpdateStrides { + /// The new strides. + strides: Strides, + }, + /// The strides are not compatible, we should recompute the buffer. + Recompute, + /// The strides are already correct. + NoChange, +} + +/// The reshape kind. +#[derive(Debug, PartialEq)] +pub enum ReshapeAnalysis { + /// Original tensor is contiguous, can update the strides. + IsContiguous, + /// Original tensor is highly permuted, can't update the strides. + HighlyPermuted, + /// Only batch dimensions are added, can update the strides. + Broadcasted, + /// Dimensions are only split, can update the strides. + Split, + /// Original tensor is bigger than output shape. + SmallerRank, + /// New shape is the same. + NoChange, +} + +impl ReshapeAnalysis { + /// Returns the proper action to take for the current analysis. + pub fn action(&self, shape: &[usize], strides: &[usize], shape_new: &[usize]) -> ReshapeAction { + match self { + ReshapeAnalysis::IsContiguous => ReshapeAction::UpdateStrides { + strides: contiguous_strides(shape_new), + }, + ReshapeAnalysis::NoChange => ReshapeAction::NoChange, + ReshapeAnalysis::HighlyPermuted | ReshapeAnalysis::SmallerRank => { + ReshapeAction::Recompute + } + ReshapeAnalysis::Broadcasted => { + let shape_rank = shape.len(); + let shape_new_rank = shape_new.len(); + let n_new_batch = shape_new_rank - shape_rank; + let num_elems = shape.iter().product::(); + let strides_new = broadcast_strides(n_new_batch, shape_rank, num_elems, strides); + + ReshapeAction::UpdateStrides { + strides: strides_new, + } + } + ReshapeAnalysis::Split => { + let strides_new = split_strides(shape, strides, shape_new); + + ReshapeAction::UpdateStrides { + strides: strides_new, + } + } + } + } +} + +/// Returns the proper action to take when reshaping a tensor. +pub fn reshape_action(shape: &Shape, strides: &Strides, shape_new: &Shape) -> ReshapeAction { + reshape_analysis(shape, Some(strides), shape_new).action(shape, strides, shape_new) +} + +/// Calculate the new strides given added batch dimensions. +pub fn broadcast_strides( + n_new_batch: usize, + rank_prev: usize, + num_elems: usize, + strides: &[usize], +) -> Strides { + let mut strides_new = strides![num_elems; rank_prev + n_new_batch]; + + for (i, s) in strides.iter().enumerate() { + strides_new[i + n_new_batch] = *s; + } + + strides_new +} + +/// Calculate the new strides given added split dimensions. +pub fn split_strides(shape: &[usize], strides: &[usize], shape_new: &[usize]) -> Strides { + let mut strides_new = strides![1; shape_new.len()]; + + // Unit dims in the old shape never anchor a group of new dims, and their + // stride can be arbitrary (0 for broadcast views, pitched values, ...). + // Skip them so a real new dim never inherits a unit dim's stride — + // e.g. reshaping [26, 1] with strides [1, 0] (a `repeat_dim` broadcast + // view) to [26, 1, 1] must keep stride 1 on dim 0; propagating the 0 + // would make every index along dim 0 alias the first element. + let skip_unit_dims = |mut idx: usize| { + while idx > 0 && shape[idx] == 1 { + idx -= 1; + } + idx + }; + + let mut old_idx = skip_unit_dims(shape.len() - 1); + let mut current_stride = strides[old_idx]; + let mut dim_prod = 1; + + for (i, dim) in shape_new.iter().enumerate().rev() { + dim_prod *= *dim; + strides_new[i] = current_stride; + if *dim == 1 { + continue; + } else if dim_prod == shape[old_idx] { + old_idx = skip_unit_dims(old_idx.saturating_sub(1)); + current_stride = strides[old_idx]; + dim_prod = 1; + } else { + current_stride *= *dim; + } + } + + strides_new +} + +/// Returns the analysis of a reshape operation. +pub fn reshape_analysis( + shape: &Shape, + strides: Option<&Strides>, + shape_new: &Shape, +) -> ReshapeAnalysis { + let shape_rank = shape.len(); + let shape_new_rank = shape_new.len(); + + let is_contiguous = match strides { + Some(strides) => is_contiguous(shape, strides), + None => false, + }; + + if is_contiguous { + return ReshapeAnalysis::IsContiguous; + } + + if shape_new_rank < shape_rank { + return ReshapeAnalysis::SmallerRank; + } + + let n_new_batch = shape_new_rank - shape_rank; + + match n_new_batch > 0 { + true => { + if shape.as_ref() == &shape_new[n_new_batch..shape_new_rank] + && shape_new[0..n_new_batch].iter().all(|it| *it == 1) + { + return ReshapeAnalysis::Broadcasted; + } else { + let mut dim_prod = 1; + let mut old_idx = 0; + for dim in shape_new.iter() { + dim_prod *= *dim; + + // We need to ignore unit dims because they don't affect analysis and break + // things because they match the default `dim_prod`. If we don't do this, + // reshapes like [2, 3] to [2, 3, 1] will panic from out of bounds access. + if *dim == 1 { + continue; + } else if dim_prod == shape[old_idx] { + dim_prod = 1; + old_idx += 1; + } else if dim_prod > shape[old_idx] { + return ReshapeAnalysis::HighlyPermuted; + } + } + return ReshapeAnalysis::Split; + } + } + + false => { + if shape == shape_new { + return ReshapeAnalysis::NoChange; + } + } + }; + + ReshapeAnalysis::HighlyPermuted +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_reshape_analysis_is_contiguous() { + let analysis = reshape_analysis( + &[32, 1, 1, 1].into(), + Some(&[1, 1, 1, 1].into()), + &[1, 1, 32, 1, 1, 1].into(), + ); + + assert_eq!(analysis, ReshapeAnalysis::IsContiguous) + } + + #[test] + fn test_reshape_analysis_is_contiguous_2() { + let analysis = reshape_analysis( + &[32, 1, 1, 8].into(), + Some(&[8, 8, 8, 1].into()), + &[1, 1, 32, 1, 1, 8].into(), + ); + + assert_eq!(analysis, ReshapeAnalysis::IsContiguous) + } + + #[test] + fn test_reshape_analysis_broadcasted_batch() { + let analysis = reshape_analysis( + &[32, 1, 1, 1].into(), + Some(&[1, 32, 32, 32].into()), + &[1, 1, 32, 1, 1, 1].into(), + ); + + assert_eq!(analysis, ReshapeAnalysis::Broadcasted) + } + + #[test] + fn test_reshape_analysis_unsqueeze_split() { + // Unsqueeze + let analysis = reshape_analysis( + &[32, 1, 1, 1].into(), + Some(&[1, 32, 32, 32].into()), + &[32, 1, 1, 1, 1].into(), + ); + + assert_eq!(analysis, ReshapeAnalysis::Split) + } + + #[test] + fn test_reshape_analysis_split() { + let analysis = reshape_analysis( + &[32, 1, 1, 1].into(), + Some(&[1, 32, 32, 32].into()), + &[4, 8, 1, 1, 1].into(), + ); + + assert_eq!(analysis, ReshapeAnalysis::Split) + } + + #[test] + fn test_split_strides_trailing_unit_dim_broadcast_view() { + // A `repeat_dim` broadcast view: [26, 1] with strides [1, 0], + // unsqueezed to [26, 1, 1]. Dim 0 must keep stride 1 — propagating + // the unit dim's stride 0 makes every row alias row 0 (this breaks + // e.g. `scatter` index tensors built via unsqueeze + repeat). + let strides = split_strides(&[26, 1], &[1, 0], &[26, 1, 1]); + assert_eq!(strides.as_ref(), &[1, 1, 1]); + } + + #[test] + fn test_split_strides_trailing_unit_dims_arbitrary_strides() { + // Unit dims can carry arbitrary strides (broadcast 0, pitched + // values, ...). They must not anchor the stride walk. + let strides = split_strides(&[32, 1, 1, 1], &[1, 32, 32, 32], &[32, 1, 1, 1, 1]); + assert_eq!(strides.as_ref(), &[1, 1, 1, 1, 1]); + } + + #[test] + fn test_split_strides_split_of_broadcast_dim_keeps_zero() { + // Splitting a real broadcast (stride 0) dim keeps 0 on the split + // parts; the leading real dim keeps its stride. + let strides = split_strides(&[26, 16], &[1, 0], &[26, 4, 4]); + assert_eq!(strides.as_ref(), &[1, 0, 0]); + } + + #[test] + fn test_split_strides_plain_unsqueeze() { + let strides = split_strides(&[26, 16], &[16, 1], &[26, 16, 1]); + assert_eq!(strides.as_ref(), &[16, 1, 1]); + } +} diff --git a/crates/burn-std/src/tensor/quantization.rs b/crates/burn-std/src/tensor/quantization.rs new file mode 100644 index 0000000..23fb308 --- /dev/null +++ b/crates/burn-std/src/tensor/quantization.rs @@ -0,0 +1,432 @@ +//! Quantization data representation. + +// Re-exported types +pub use cubecl_common::quant::scheme::{ + BlockSize, QuantLevel, QuantMode, QuantParam, QuantScheme, QuantStore, QuantValue, +}; + +/// Alignment (in bytes) for quantization parameters in serialized tensor data. +/// +/// NOTE: This is currently f32-based since scales were originally always f32. +/// With `QuantParam` now supporting different precisions (F16, BF16, etc.), +/// this alignment may need to be revisited in the future. +pub const QPARAM_ALIGN: usize = core::mem::align_of::(); + +use alloc::vec::Vec; +use core::any::TypeId; +use num_traits::PrimInt; +use serde::{Deserialize, Serialize}; + +use crate::{DType, Metadata, Shape, bytes::Bytes}; + +/// Configuration for a device quantization behavior. +/// +/// This configuration determines how tensors are quantized and how quantization rules +/// propagate through operations on a given device. It is applied once during device +/// initialization. See also the [device settings](crate::DeviceSettings). +#[derive(new, Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)] +pub struct QuantConfig { + /// Defines how a tensor is quantized. + pub scheme: QuantScheme, + /// How quantization is propagated during computation. + pub propagation: QuantPropagation, + // NOTE: accumulation is currently unused, only scheme and propagation have an impact + // /// The precision used for the accumulation in various kernels. + // pub acc: QuantAcc, +} + +#[derive( + Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Default, +)] +/// The precision of accumulating elements. +pub enum QuantAcc { + /// Full precision. + #[default] + F32, + /// Half precision. + F16, + /// bfloat16 precision. + BF16, +} + +/// Calibration method used to compute the quantization range mapping. +pub enum Calibration { + /// Computes quantization range mapping based on the min and max values. + MinMax, + /// Absolute-mean calibration for BitNet b1.58-style `{-1, 0, +1}` weight quantization. + /// + /// The range is `[-γ, +γ]` where γ = `mean(|W|)` per tensor or per block (BitNet b1.58 + /// §3.1). Use with `QuantValue::Q2S` and `QuantStore::PackedU32` for 2-bit packed storage. + AbsMean, +} + +/// Specify if the output of an operation is quantized using the scheme of the input +/// or returned unquantized. +#[derive( + Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Default, +)] +pub enum QuantPropagation { + /// The output is quantized using the scheme of the input. + Propagate, + /// The output is not quantized. + #[default] + Inhibit, +} + +/// The quantization tensor data parameters. +#[derive(Clone, Debug)] +pub struct QParams { + /// The scaling factor. + pub scales: S, +} + +/// A quantization parameter tensor descriptor. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct QParamTensor { + /// Start of the tensor in the buffer + pub offset_start: usize, + /// Offset of tensor end from the end of the buffer + pub offset_end: usize, + /// Metadata of the tensor + pub metadata: Metadata, + /// Data type of the tensor + pub dtype: DType, +} + +/// Calculate the shape of the quantization parameters for a given tensor and level +pub fn params_shape(data_shape: &Shape, level: QuantLevel) -> Shape { + match level { + QuantLevel::Tensor => Shape::new([1]), + QuantLevel::Block(block_size) => { + let mut params_shape = data_shape.clone(); + let block_size = block_size.to_dim_vec(data_shape.num_dims()); + + for (shape, block_size) in params_shape.iter_mut().zip(block_size) { + *shape = (*shape).div_ceil(block_size as usize); + } + + params_shape + } + } +} + +/// Quantized data bytes representation. +/// +/// # Notes +/// 1) The quantized values are packed into 32-bit unsigned integers. For example, int8 +/// quantized values pack 4 grouped values into a single `u32`. When unpacking these values, +/// we make sure to retrieve only the meaningful values (and ignore the alignment padding). +/// 2) Quantization parameters are appended to the tensor data. +/// As such, the last bytes always correspond to the scale parameter. +/// If the quantization scheme includes an offset (zero-point) parameter, it is next to last. +pub struct QuantizedBytes { + /// The quantized values and quantization parameters represented as bytes. + pub bytes: Bytes, + /// The quantization scheme. + pub scheme: QuantScheme, + /// The number of quantized elements. + pub num_elements: usize, +} + +impl QuantizedBytes { + /// Creates a new quantized bytes representation. + pub fn new( + value: Vec, + scheme: QuantScheme, + scales: &[f32], + ) -> Self { + let num_elements = value.len(); + // Only used for 8-bit quantization data comparison in tests + if TypeId::of::() != TypeId::of::() { + panic!("Invalid quantized type"); + } + + // Re-interpret `Vec` as `Vec` with `Vec::from_raw_parts` + let i8s: Vec = bytemuck::allocation::cast_vec(value); + let mut bytes = Bytes::from_elems(i8s); + + let scales = match scheme.level { + QuantLevel::Tensor => &scales[..1], + QuantLevel::Block(_block_size) => scales, + }; + let scale_bytes = encode_scales(scales, scheme.param); + bytes.extend_from_byte_slice_aligned(scale_bytes.as_slice(), QPARAM_ALIGN); + + Self { + bytes, + scheme, + num_elements, + } + } + + /// Returns the int8 quantized values with the quantization parameters. + pub fn into_vec_i8(self) -> (Vec, QParams>) { + let param = self.scheme.param; + let (values, (qparams, num_params)) = self.split_values_off(); + + // Quantization parameters are added at the end of the tensor data. + // As such, the last bytes always correspond to the scale parameter(s), + // stored at the scheme's param dtype. For example, per-block + // quantization can have multiple parameters for a single tensor: + // [scale, scale, scale, ...] + let scales_size = scale_size(param) * num_params; + let scales = decode_scales(&qparams[qparams.len() - scales_size..], param); + + (values, QParams { scales }) + } + + fn split_i8_values(self, scale_bytes: usize) -> (Vec, Vec) { + let mut values = read_bytes_to_i8(self.bytes); + + let values_end = values.len() - scale_bytes; + let qparams = values.split_off(values_end); + + (values, bytemuck::cast_vec(qparams)) + } + + /// Splits the quantized values of the tensor from the quantization parameters. + /// + /// Returns the values in i8 and a newly allocated vector containing the + /// quantization parameter bytes. + fn split_values_off(self) -> (Vec, (Vec, usize)) { + let num_params = match self.scheme.level { + QuantLevel::Tensor => 1, + QuantLevel::Block(block_size) => self.num_elements / block_size.num_elements(), + }; + let scale_bytes = scale_size(self.scheme.param) * num_params; + + if let QuantStore::PackedU32(packed_dim) = self.scheme.store { + assert_eq!( + packed_dim, 0, + "Packing must be on innermost dimension for splitting off values" + ); + } + + let (values, qparams) = match self.scheme.store { + QuantStore::Native => self.split_i8_values(scale_bytes), + QuantStore::PackedU32(_) => match self.scheme.value { + QuantValue::Q8F | QuantValue::Q8S => self.split_i8_values(scale_bytes), + QuantValue::Q4F | QuantValue::Q4S | QuantValue::Q2F | QuantValue::Q2S => { + let split_at = self.bytes.len() - scale_bytes; + let qparams = self.bytes[split_at..].to_vec(); + let values = bytemuck::cast_slice::<_, u32>(&self.bytes[..split_at]); + // Sub-byte values are unpacked as i8s for value equality tests + let values = unpack_q_to_i8s(values, self.num_elements, &self.scheme.value); + (values, qparams) + } + QuantValue::E4M3 | QuantValue::E5M2 | QuantValue::E2M1 => { + unimplemented!("Not yet supported") + } + }, + QuantStore::PackedNative(_) => unimplemented!("Not yet supported"), + }; + + (values, (qparams, num_params)) + } +} + +/// Bytes per stored scale entry for the given param dtype. +fn scale_size(param: QuantParam) -> usize { + match param { + QuantParam::F32 => 4, + QuantParam::F16 | QuantParam::BF16 => 2, + QuantParam::UE8M0 | QuantParam::UE4M3 => 1, + } +} + +/// Decode stored scale entries into f32. +fn decode_scales(bytes: &[u8], param: QuantParam) -> Vec { + match param { + QuantParam::F32 => bytes + .chunks_exact(4) + .map(|c| f32::from_ne_bytes([c[0], c[1], c[2], c[3]])) + .collect(), + QuantParam::F16 => bytes + .chunks_exact(2) + .map(|c| crate::f16::from_ne_bytes([c[0], c[1]]).to_f32()) + .collect(), + QuantParam::BF16 => bytes + .chunks_exact(2) + .map(|c| crate::bf16::from_ne_bytes([c[0], c[1]]).to_f32()) + .collect(), + QuantParam::UE8M0 | QuantParam::UE4M3 => unimplemented!("Not yet supported"), + } +} + +/// Encode f32 scales at the param dtype for serialization. +fn encode_scales(scales: &[f32], param: QuantParam) -> Vec { + match param { + QuantParam::F32 => scales.iter().flat_map(|s| s.to_ne_bytes()).collect(), + QuantParam::F16 => scales + .iter() + .flat_map(|s| crate::f16::from_f32(*s).to_ne_bytes()) + .collect(), + QuantParam::BF16 => scales + .iter() + .flat_map(|s| crate::bf16::from_f32(*s).to_ne_bytes()) + .collect(), + QuantParam::UE8M0 | QuantParam::UE4M3 => unimplemented!("Not yet supported"), + } +} + +fn read_bytes_to_i8(bytes: Bytes) -> Vec { + match bytes.try_into_vec::() { + Ok(val) => val, + // Safety, + // + // `Vec` can be Re-interpreted as `Vec` since they share the same alignment. + Err(bytes) => unsafe { core::mem::transmute::, Vec>(bytes.to_vec()) }, + } +} + +/// Pack signed 8-bit integer values into a sequence of unsigned 32-bit integers. +pub fn pack_i8s_to_u32s(values: Vec) -> Vec { + // Shift and combine groups of four 8-bit values into a u32. + // Same as doing this: + // let result = (d_u8 & 0xFF) << 24 | (c_u8 & 0xFF) << 16 | (b_u8 & 0xFF) << 8 | (a_u8 & 0xFF); + #[cfg(target_endian = "big")] + { + values + .chunks(4) + .map(|x| { + x.iter() + .enumerate() + .fold(0u32, |acc, (i, x)| acc | (*x as u32 & 0xFF) << (i * 8)) + }) + .collect() + } + + // The order of bytes in little endian matches the above description, we just need to + // handle padding when the number of values is not a factor of 4 + #[cfg(target_endian = "little")] + { + let mut values = values; + let remainder = values.len() % 4; + if remainder != 0 { + // Pad with zeros + values.extend(core::iter::repeat_n(0, 4 - remainder)); + } + + let len = values.len() / 4; + let capacity = values.capacity() / 4; + + // Pre-forget the old vec and re-interpret as u32 + let mut values = core::mem::ManuallyDrop::new(values); + let ptr = values.as_mut_ptr() as *mut u32; + + unsafe { Vec::from_raw_parts(ptr, len, capacity) } + } +} + +/// Unpack integer values into a sequence of signed 8-bit integers. +pub(crate) fn unpack_q_to_i8s( + values: &[Q], + numel: usize, + value: &QuantValue, +) -> Vec { + let size_store = size_of::() * 8; + let size_quant = value.size_bits(); + let num_quants = size_store / size_quant; + let mask = Q::from((1 << size_quant) - 1).unwrap(); + let sign_shift = 8 - size_quant; // sign extension for sub-byte values + values + .iter() + .enumerate() + .flat_map(|(i, &packed)| { + // A single u32 could contain less than four 8-bit values... + let n = core::cmp::min(num_quants, numel - i * num_quants); + // Extract each 8-bit segment from u32 and cast back to i8 + // Same as doing this (when 4 values are fully packed): + // let a = (packed & 0xFF) as i8; + // let b = ((packed >> 8) & 0xFF) as i8; + // let c = ((packed >> 16) & 0xFF) as i8; + // let d = ((packed >> 24) & 0xFF) as i8; + (0..n).map(move |i| { + let raw = (packed >> (i * size_quant) & mask).to_u8().unwrap(); + ((raw << sign_shift) as i8) >> sign_shift + }) + }) + .collect() +} + +#[cfg(test)] +mod tests { + + use super::*; + use alloc::vec; + + #[test] + fn should_pack_i8s_to_u32() { + let packed = pack_i8s_to_u32s(vec![-128, 2, -3, 127]); + + assert_eq!(packed, vec![2147287680]); + } + + #[test] + fn should_pack_i8s_to_u32_padded() { + let packed = pack_i8s_to_u32s(vec![-128, 2, -3, 127, 55]); + let packed_padded = pack_i8s_to_u32s(vec![-128, 2, -3, 127, 55, 0, 0, 0]); + + assert_eq!(packed, vec![2147287680, 55]); + assert_eq!(packed, packed_padded); + } + + #[test] + fn should_unpack_u32s_to_i8s() { + let unpacked = unpack_q_to_i8s(&[2147287680u32], 4, &QuantValue::Q8S); + + assert_eq!(unpacked, vec![-128, 2, -3, 127]); + } + + #[test] + fn should_unpack_u32s_to_i8s_padded() { + let unpacked = unpack_q_to_i8s(&[55u32], 1, &QuantValue::Q8S); + + assert_eq!(unpacked, vec![55]); + } + + #[test] + fn should_unpack_u32s_to_i8s_arange() { + let unpacked = unpack_q_to_i8s( + &[ + 0u32, 286331136, 286331153, 572657937, 572662306, 857874978, 858993459, 858993459, + 1145324612, 1145324612, 1431655748, 1431655765, 1717982549, 1717986918, 2003199590, + 2004318071, + ], + 128, + &QuantValue::Q4S, + ); + + assert_eq!( + unpacked, + vec![ + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, + 3, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 5, 5, + 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, + 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7 + ] + ); + } + + #[test] + fn should_pack_unpack_quantization_parameters_per_tensor_symmetric() { + // Quantized [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]] + let scale = 0.03937008; + let values = vec![0i8, 25, 51, 76, 102, 127]; + + let q_bytes = QuantizedBytes::new( + values.clone(), + QuantScheme::default() + .with_value(QuantValue::Q8S) + .with_store(QuantStore::Native), + &[scale], + ); + + let (q_values, qparams) = q_bytes.into_vec_i8(); + + assert_eq!(qparams.scales, vec![scale]); + + assert_eq!(q_values, values); + } +} diff --git a/crates/burn-std/src/tensor/shape.rs b/crates/burn-std/src/tensor/shape.rs new file mode 100644 index 0000000..12313a9 --- /dev/null +++ b/crates/burn-std/src/tensor/shape.rs @@ -0,0 +1,271 @@ +//! Tensor shape definition. + +use super::{Slice, SliceArg}; +use alloc::vec::Vec; +use core::ops::Range; + +pub use crate::errors::ExpressionError; + +pub use cubecl_zspace::{MetadataError, Shape, SmallVec, calculate_matmul_output, shape}; + +/// Slice-related ops on [`Shape`] +pub trait SliceOps: Sized { + /// Convert shape dimensions to full covering ranges (0..dim) for each dimension. + fn into_ranges(self) -> Vec>; + /// Converts slice arguments into an array of slice specifications for the shape. + /// + /// This method returns an array of `Slice` objects that can be used for slicing operations. + /// The slices are clamped to the shape's dimensions. Similar to `into_ranges()`, but + /// allows custom slice specifications instead of full ranges. + /// For creating complex slice specifications, use the [`s!`] macro. + /// + /// # Arguments + /// + /// * `slices` - An array of slice specifications, where each element can be: + /// - A range (e.g., `2..5`) + /// - An index + /// - A `Slice` object + /// - The output of the [`s!`] macro for advanced slicing + /// + /// # Behavior + /// + /// - Supports partial and full slicing in any number of dimensions. + /// - Missing ranges are treated as full slices if D > D2. + /// - Handles negative indices by wrapping around from the end of the dimension. + /// - Clamps ranges to the shape's dimensions if they exceed the bounds. + /// + /// # Returns + /// + /// An array of `Slice` objects corresponding to the provided slice specifications, + /// clamped to the shape's actual dimensions. + /// + /// # Examples + /// + /// ```rust + /// use burn_std::{Shape, Slice, s, SliceOps}; + /// + /// fn example() { + /// // 1D slicing + /// let slices = Shape::new([4]).into_slices(1..4); + /// assert_eq!(slices[0].to_range(4), 1..3); + /// + /// // 2D slicing + /// let slices = Shape::new([3, 4]).into_slices(s![1..4, 0..2]); + /// assert_eq!(slices[0].to_range(3), 1..3); + /// assert_eq!(slices[1].to_range(4), 0..2); + /// + /// // Using negative indices + /// let slices = Shape::new([3]).into_slices(..-2); + /// assert_eq!(slices[0].to_range(3), 0..1); + /// + /// // Using the slice macro to select different ranges + /// let slices = Shape::new([2, 3, 4]).into_slices(s![.., 1..-1]); + /// assert_eq!(slices[0].to_range(2), 0..2); + /// assert_eq!(slices[1].to_range(3), 1..2); + /// } + /// ``` + /// + /// # See Also + /// + /// - [`s!`] - The recommended macro for creating slice specifications + /// - [`Shape::into_ranges`] - Convert to full covering ranges + /// + /// [`s!`]: crate::s! + fn into_slices(self, slices: S) -> Vec + where + S: SliceArg; + /// Compute the output shape from the given slices. + fn slice(self, slices: &[Slice]) -> Result; +} + +impl SliceOps for Shape { + fn into_ranges(self) -> Vec> { + self.iter().map(|&d| 0..d).collect() + } + + fn into_slices(self, slices: S) -> Vec + where + S: SliceArg, + { + slices.into_slices(&self) + } + + fn slice(mut self, slices: &[Slice]) -> Result { + if slices.len() > self.rank() { + return Err(MetadataError::RankMismatch { + left: self.rank(), + right: slices.len(), + }); + } + + slices + .iter() + .zip(self.iter_mut()) + .for_each(|(slice, dim_size)| *dim_size = slice.output_size(*dim_size)); + + Ok(self) + } +} + +#[cfg(test)] +#[allow(clippy::identity_op, reason = "useful for clarity")] +mod tests { + use super::*; + use crate::s; + use alloc::vec; + + #[test] + fn test_into_ranges() { + let dims = [2, 3, 4, 5]; + let shape = Shape::new(dims); + assert_eq!(shape.into_ranges(), vec![0..2, 0..3, 0..4, 0..5]); + } + + #[allow(clippy::single_range_in_vec_init)] + #[test] + fn test_into_slices() { + let slices = Shape::new([3]).into_slices(1..4); + assert_eq!(slices[0].to_range(3), 1..3); + + let slices = Shape::new([3, 4]).into_slices(s![1..4, 0..2]); + assert_eq!(slices[0].to_range(3), 1..3); + assert_eq!(slices[1].to_range(4), 0..2); + + let slices = Shape::new([3]).into_slices(..-2); + assert_eq!(slices[0].to_range(3), 0..1); + + let slices = Shape::new([2, 3, 4]).into_slices(s![.., 1..-1]); + assert_eq!(slices[0].to_range(2), 0..2); + assert_eq!(slices[1].to_range(3), 1..2); + + let slices = Shape::new([2, 3, 4]).into_slices(s![..20, 2]); + assert_eq!(slices[0].to_range(2), 0..2); + assert_eq!(slices[1].to_range(3), 2..3); + } + + #[test] + fn test_shape_as_slice() { + let dims = [2, 3, 4, 5]; + let shape = Shape::new(dims); + + assert_eq!(shape.as_slice(), dims.as_slice()); + + // Deref coercion + let shape_slice: &[usize] = &shape; + assert_eq!(shape_slice, *&[2, 3, 4, 5]); + } + + #[test] + fn test_shape_as_mut_slice() { + let mut dims = [2, 3, 4, 5]; + let mut shape = Shape::new(dims); + + let shape_mut = shape.as_mut_slice(); + assert_eq!(shape_mut, dims.as_mut_slice()); + shape_mut[1] = 6; + + assert_eq!(shape_mut, &[2, 6, 4, 5]); + + let mut shape = Shape::new(dims); + let shape = &mut shape[..]; + shape[1] = 6; + + assert_eq!(shape, shape_mut) + } + + #[test] + fn test_shape_slice_output_shape_basic() { + // Test basic slicing with step=1 + let slices = [ + Slice::new(0, Some(5), 1), // 5 elements + Slice::new(2, Some(8), 1), // 6 elements + ]; + let original_shape = Shape::new([10, 10, 10]); + let result = original_shape.slice(&slices).unwrap(); + assert_eq!(result, Shape::new([5, 6, 10])); + } + + #[test] + fn test_shape_slice_output_shape_with_positive_steps() { + // Test slicing with various positive steps + let slices = [ + Slice::new(0, Some(10), 2), // [0,2,4,6,8] -> 5 elements + Slice::new(1, Some(9), 3), // [1,4,7] -> 3 elements + Slice::new(0, Some(7), 4), // [0,4] -> 2 elements + ]; + let original_shape = Shape::new([20, 20, 20, 30]); + let result = original_shape.slice(&slices).unwrap(); + assert_eq!(result, Shape::new([5, 3, 2, 30])); + } + + #[test] + fn test_shape_slice_output_shape_with_negative_steps() { + // Test slicing with negative steps (backward iteration) + let slices = [ + Slice::new(0, Some(10), -1), // 10 elements traversed backward + Slice::new(2, Some(8), -2), // [7,5,3] -> 3 elements + ]; + let original_shape = Shape::new([20, 20, 20]); + let result = original_shape.slice(&slices).unwrap(); + assert_eq!(result, Shape::new([10, 3, 20])); + } + + #[test] + fn test_shape_slice_output_shape_mixed_steps() { + // Test with a mix of positive, negative, and unit steps + let slices = [ + Slice::from_range_stepped(1..6, 1), // 5 elements + Slice::from_range_stepped(0..10, -3), // [9,6,3,0] -> 4 elements + Slice::from_range_stepped(2..14, 4), // [2,6,10] -> 3 elements + ]; + let original_shape = Shape::new([20, 20, 20]); + let result = original_shape.slice(&slices).unwrap(); + assert_eq!(result, Shape::new([5, 4, 3])); + } + + #[test] + fn test_shape_slice_output_shape_partial_dims() { + // Test when slices has fewer dimensions than original shape + let slices = [ + Slice::from_range_stepped(2..7, 2), // [2,4,6] -> 3 elements + ]; + let original_shape = Shape::new([10, 20, 30, 40]); + let result = original_shape.slice(&slices).unwrap(); + assert_eq!(result, Shape::new([3, 20, 30, 40])); + } + + #[test] + fn test_shape_slice_output_shape_edge_cases() { + // Test edge cases with small ranges and large steps + let slices = [ + Slice::from_range_stepped(0..1, 1), // Single element + Slice::from_range_stepped(0..10, 100), // Step larger than range -> 1 element + Slice::from_range_stepped(5..5, 1), // Empty range -> 0 elements + ]; + let original_shape = Shape::new([10, 20, 30]); + let result = original_shape.slice(&slices).unwrap(); + assert_eq!(result, Shape::new([1, 1, 0])); + } + + #[test] + fn test_shape_slice_output_shape_empty() { + // Test with no slice infos (should return original shape) + let slices = []; + let original_shape = Shape::new([10, 20, 30]); + let result = original_shape.slice(&slices).unwrap(); + assert_eq!(result, Shape::new([10, 20, 30])); + } + + #[test] + fn test_shape_slice_output_shape_uneven_division() { + // Test cases where range size doesn't divide evenly by step + let slices = [ + Slice::from_range_stepped(0..7, 3), // ceil(7/3) = 3 elements: [0,3,6] + Slice::from_range_stepped(0..11, 4), // ceil(11/4) = 3 elements: [0,4,8] + Slice::from_range_stepped(1..10, 5), // ceil(9/5) = 2 elements: [1,6] + ]; + let original_shape = Shape::new([20, 20, 20]); + let result = original_shape.slice(&slices).unwrap(); + assert_eq!(result, Shape::new([3, 3, 2])); + } +} diff --git a/crates/burn-std/src/tensor/slice.rs b/crates/burn-std/src/tensor/slice.rs new file mode 100644 index 0000000..7a90e44 --- /dev/null +++ b/crates/burn-std/src/tensor/slice.rs @@ -0,0 +1,937 @@ +//! Tensor slice utilities. + +use crate::Shape; +use crate::indexing::AsIndex; +use alloc::format; +use alloc::vec::Vec; +use core::fmt::{Display, Formatter}; +use core::ops::{Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive}; +use core::str::FromStr; + +/// Trait for slice arguments that can be converted into an array of slices. +/// This allows the `slice` method to accept both single slices (from `s![..]`) +/// and arrays of slices (from `s![.., ..]` or `[0..5, 1..3]`). +pub trait SliceArg { + /// Convert to an vec of slices with clamping to shape dimensions. + /// + /// Returns a [Slice] for each dimension in `shape`. + fn into_slices(self, shape: &Shape) -> Vec; +} + +impl + Clone> SliceArg for &[S] { + fn into_slices(self, shape: &Shape) -> Vec { + assert!( + self.len() <= shape.num_dims(), + "Too many slices provided for shape, got {} but expected at most {}", + self.len(), + shape.num_dims() + ); + + shape + .iter() + .enumerate() + .map(|(i, dim_size)| { + let slice = if i >= self.len() { + Slice::full() + } else { + self[i].clone().into() + }; + // Apply shape clamping by converting to range and back + let clamped_range = slice.to_range(*dim_size); + Slice::new( + clamped_range.start as isize, + Some(clamped_range.end as isize), + slice.step(), + ) + }) + .collect::>() + } +} + +impl SliceArg for &Vec { + fn into_slices(self, shape: &Shape) -> Vec { + self.as_slice().into_slices(shape) + } +} + +impl SliceArg for [T; R] +where + T: Into + Clone, +{ + fn into_slices(self, shape: &Shape) -> Vec { + self.as_slice().into_slices(shape) + } +} + +impl SliceArg for T +where + T: Into, +{ + fn into_slices(self, shape: &Shape) -> Vec { + let slice: Slice = self.into(); + [slice].as_slice().into_slices(shape) + } +} + +/// Slice argument constructor for tensor indexing. +/// +/// The `s![]` macro is used to create multi-dimensional slice specifications for tensors. +/// It converts various range syntax forms into a `&[Slice]` that can be used with +/// `tensor.slice()` and `tensor.slice_assign()` operations. +/// +/// # Syntax Overview +/// +/// ## Basic Forms +/// +/// * **`s![index]`** - Index a single element (produces a subview with that axis removed) +/// * **`s![range]`** - Slice a range of elements +/// * **`s![range;step]`** - Slice a range with a custom step +/// * **`s![dim1, dim2, ...]`** - Multiple dimensions, each can be any of the above forms +/// +/// ## Range Types +/// +/// All standard Rust range types are supported: +/// * **`a..b`** - From `a` (inclusive) to `b` (exclusive) +/// * **`a..=b`** - From `a` to `b` (both inclusive) +/// * **`a..`** - From `a` to the end +/// * **`..b`** - From the beginning to `b` (exclusive) +/// * **`..=b`** - From the beginning to `b` (inclusive) +/// * **`..`** - The full range (all elements) +/// +/// ## Negative Indices +/// +/// Negative indices count from the end of the axis: +/// * **`-1`** refers to the last element +/// * **`-2`** refers to the second-to-last element +/// * And so on... +/// +/// This works in all range forms: `s![-3..-1]`, `s![-2..]`, `s![..-1]` +/// +/// ## Step Syntax +/// +/// Steps control the stride between selected elements: +/// * **`;step`** after a range specifies the step +/// * **Positive steps** select every nth element going forward +/// * **Negative steps** select every nth element going backward +/// * Default step is `1` when not specified +/// * Step cannot be `0` +/// +/// ### Negative Step Behavior +/// +/// With negative steps, the range bounds still specify *which* elements to include, +/// but the traversal order is reversed: +/// +/// * `s![0..5;-1]` selects indices `[4, 3, 2, 1, 0]` (not `[0, 1, 2, 3, 4]`) +/// * `s![2..8;-2]` selects indices `[7, 5, 3]` (starting from 7, going backward by 2) +/// * `s![..;-1]` reverses the entire axis +/// +/// This matches the semantics of NumPy and the ndarray crate. +/// +/// # Examples +/// +/// ## Basic Slicing +/// +/// ```rust,ignore +/// use burn_tensor::{Tensor, s}; +/// +/// # fn example(tensor: Tensor) { +/// // Select rows 0-5 (exclusive) +/// let subset = tensor.slice(s![0..5, .., ..]); +/// +/// // Select the last row +/// let last_row = tensor.slice(s![-1, .., ..]); +/// +/// // Select columns 2, 3, 4 +/// let cols = tensor.slice(s![.., 2..5, ..]); +/// +/// // Select a single element at position [1, 2, 3] +/// let element = tensor.slice(s![1, 2, 3]); +/// # } +/// ``` +/// +/// ## Slicing with Steps +/// +/// ```rust,ignore +/// use burn_tensor::{Tensor, s}; +/// +/// # fn example(tensor: Tensor) { +/// // Select every 2nd row +/// let even_rows = tensor.slice(s![0..10;2, ..]); +/// +/// // Select every 3rd column +/// let cols = tensor.slice(s![.., 0..9;3]); +/// +/// // Select every 2nd element in reverse order +/// let reversed_even = tensor.slice(s![10..0;-2, ..]); +/// # } +/// ``` +/// +/// ## Reversing Dimensions +/// +/// ```rust,ignore +/// use burn_tensor::{Tensor, s}; +/// +/// # fn example(tensor: Tensor) { +/// // Reverse the first dimension +/// let reversed = tensor.slice(s![..;-1, ..]); +/// +/// // Reverse both dimensions +/// let fully_reversed = tensor.slice(s![..;-1, ..;-1]); +/// +/// // Reverse a specific range +/// let range_reversed = tensor.slice(s![2..8;-1, ..]); +/// # } +/// ``` +/// +/// ## Complex Multi-dimensional Slicing +/// +/// ```rust,ignore +/// use burn_tensor::{Tensor, s}; +/// +/// # fn example(tensor: Tensor) { +/// // Mix of different slice types +/// let complex = tensor.slice(s![ +/// 0..10;2, // Every 2nd element from 0 to 10 +/// .., // All elements in dimension 1 +/// 5..15;-3, // Every 3rd element from 14 down to 5 +/// -1 // Last element in dimension 3 +/// ]); +/// +/// // Using inclusive ranges +/// let inclusive = tensor.slice(s![2..=5, 1..=3, .., ..]); +/// +/// // Negative indices with steps +/// let from_end = tensor.slice(s![-5..-1;2, .., .., ..]); +/// # } +/// ``` +/// +/// ## Slice Assignment +/// +/// ```rust,ignore +/// use burn_tensor::{Tensor, s}; +/// +/// # fn example(tensor: Tensor, values: Tensor) { +/// // Assign to every 2nd row +/// let tensor = tensor.slice_assign(s![0..10;2, ..], values); +/// +/// // Assign to a reversed slice +/// let tensor = tensor.slice_assign(s![..;-1, 0..5], values); +/// # } +/// ``` +#[macro_export] +macro_rules! s { + // Empty - should not happen + [] => { + compile_error!("Empty slice specification") + }; + + // Single expression with step + [$range:expr; $step:expr] => { + { + #[allow(clippy::reversed_empty_ranges)] + { + $crate::tensor::Slice::from_range_stepped($range, $step) + } + } + }; + + // Single expression without step (no comma after) + [$range:expr] => { + { + #[allow(clippy::reversed_empty_ranges)] + { + $crate::tensor::Slice::from($range) + } + } + }; + + // Two or more expressions with first having step + [$range:expr; $step:expr, $($rest:tt)*] => { + { + #[allow(clippy::reversed_empty_ranges)] + { + $crate::s!(@internal [$crate::tensor::Slice::from_range_stepped($range, $step)] $($rest)*) + } + } + }; + + // Two or more expressions with first not having step + [$range:expr, $($rest:tt)*] => { + { + #[allow(clippy::reversed_empty_ranges)] + { + $crate::s!(@internal [$crate::tensor::Slice::from($range)] $($rest)*) + } + } + }; + + // Internal: finished parsing + (@internal [$($acc:expr),*]) => { + [$($acc),*] + }; + + // Internal: parse range with step followed by comma + (@internal [$($acc:expr),*] $range:expr; $step:expr, $($rest:tt)*) => { + $crate::s!(@internal [$($acc,)* $crate::tensor::Slice::from_range_stepped($range, $step as isize)] $($rest)*) + }; + + // Internal: parse range with step at end + (@internal [$($acc:expr),*] $range:expr; $step:expr) => { + $crate::s!(@internal [$($acc,)* $crate::tensor::Slice::from_range_stepped($range, $step as isize)]) + }; + + // Internal: parse range without step followed by comma + (@internal [$($acc:expr),*] $range:expr, $($rest:tt)*) => { + $crate::s!(@internal [$($acc,)* $crate::tensor::Slice::from($range)] $($rest)*) + }; + + // Internal: parse range without step at end + (@internal [$($acc:expr),*] $range:expr) => { + $crate::s!(@internal [$($acc,)* $crate::tensor::Slice::from($range)]) + }; +} + +/// A slice specification for a single tensor dimension. +/// +/// This struct represents a range with an optional step, used for advanced indexing +/// operations on tensors. It is typically created using the [`s!`] macro rather than +/// constructed directly. +/// +/// # Fields +/// +/// * `start` - The starting index (inclusive). Negative values count from the end. +/// * `end` - The ending index (exclusive). `None` means to the end of the dimension. +/// * `step` - The stride between elements. Must be non-zero. +/// +/// # Index Interpretation +/// +/// - **Positive indices**: Count from the beginning (0-based) +/// - **Negative indices**: Count from the end (-1 is the last element) +/// - **Bounds checking**: Indices are clamped to valid ranges +/// +/// # Step Behavior +/// +/// - **Positive step**: Traverse forward through the range +/// - **Negative step**: Traverse backward through the range +/// - **Step size**: Determines how many elements to skip +/// +/// # Examples +/// +/// While you typically use the [`s!`] macro, you can also construct slices directly: +/// +/// ```rust,ignore +/// use burn_tensor::Slice; +/// +/// // Equivalent to s![2..8] +/// let slice1 = Slice::new(2, Some(8), 1); +/// +/// // Equivalent to s![0..10;2] +/// let slice2 = Slice::new(0, Some(10), 2); +/// +/// // Equivalent to s![..;-1] (reverse) +/// let slice3 = Slice::new(0, None, -1); +/// ``` +/// +/// See also the [`s!`] macro for the preferred way to create slices. +#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct Slice { + /// Slice start index. + pub start: isize, + /// Slice end index (exclusive). + pub end: Option, + /// Step between elements (default: 1). + pub step: isize, +} + +/// Defines an [`Iterator`] over a [`Slice`]. +#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct SliceIter { + slice: Slice, + current: isize, +} + +impl Iterator for SliceIter { + type Item = isize; + + fn next(&mut self) -> Option { + let next = self.current; + self.current += self.slice.step; + + if let Some(end) = self.slice.end { + if self.slice.is_reversed() { + if next <= end { + return None; + } + } else if next >= end { + return None; + } + } + + Some(next) + } +} + +/// Note: Unbounded [`Slice`]s produce infinite iterators. +impl IntoIterator for Slice { + type Item = isize; + type IntoIter = SliceIter; + + fn into_iter(self) -> Self::IntoIter { + SliceIter { + slice: self, + current: self.start, + } + } +} + +impl Default for Slice { + fn default() -> Self { + Self::full() + } +} + +impl Slice { + /// Creates a new slice with start, end, and step + pub const fn new(start: isize, end: Option, step: isize) -> Self { + assert!(step != 0, "Step cannot be zero"); + Self { start, end, step } + } + + /// Creates a slice that represents the full range. + pub const fn full() -> Self { + Self::new(0, None, 1) + } + + /// Creates a slice that represents a single index + pub fn index(idx: isize) -> Self { + Self { + start: idx, + end: handle_signed_inclusive_end(idx), + step: 1, + } + } + + /// Converts the slice to a vector. + pub fn into_vec(self) -> Vec { + assert!( + self.end.is_some(), + "Slice must have an end to convert to a vector: {self:?}" + ); + self.into_iter().collect() + } + + /// Clips the slice to a maximum size. + /// + /// # Example + /// + /// ```rust,ignore + /// assert_eq!( + /// Slice::new(0, None, 1).bound_to(10), + /// Slice::new(0, Some(10), 1)); + /// assert_eq!( + /// Slice::new(0, Some(5), 1).bound_to(10), + /// Slice::new(0, Some(5), 1)); + /// assert_eq!( + /// Slice::new(0, None, -1).bound_to(10), + /// Slice::new(0, Some(-11), -1)); + /// assert_eq!( + /// Slice::new(0, Some(-5), -1).bound_to(10), + /// Slice::new(0, Some(-5), -1)); + /// ``` + pub fn bound_to(self, size: usize) -> Self { + let mut bounds = size as isize; + + if let Some(end) = self.end { + if end > 0 { + bounds = end.min(bounds); + } else { + bounds = end.max(-(bounds + 1)); + } + } else if self.is_reversed() { + bounds = -(bounds + 1); + } + + Self { + end: Some(bounds), + ..self + } + } + + /// Creates a slice with a custom step + pub fn with_step(start: isize, end: Option, step: isize) -> Self { + assert!(step != 0, "Step cannot be zero"); + Self { start, end, step } + } + + /// Creates a slice from a range with a specified step + pub fn from_range_stepped>(range: R, step: isize) -> Self { + assert!(step != 0, "Step cannot be zero"); + let mut slice = range.into(); + slice.step = step; + slice + } + + /// Returns the step of the slice + pub fn step(&self) -> isize { + self.step + } + + /// Returns the range for this slice given a dimension size + pub fn range(&self, size: usize) -> Range { + self.to_range(size) + } + + /// Convert this slice to a range for a dimension of the given size. + /// + /// # Arguments + /// + /// * `size` - The size of the dimension to slice. + /// + /// # Returns + /// + /// A `Range` representing the slice bounds. + pub fn to_range(&self, size: usize) -> Range { + // Always return a valid range with start <= end + // The step information will be handled separately + let start = convert_signed_index(self.start, size); + let end = match self.end { + Some(end) => convert_signed_index(end, size), + None => size, + }; + start..end + } + + /// Converts the slice into a range and step tuple + pub fn to_range_and_step(&self, size: usize) -> (Range, isize) { + let range = self.to_range(size); + (range, self.step) + } + + /// Returns true if the step is negative + pub fn is_reversed(&self) -> bool { + self.step < 0 + } + + /// Calculates the output size for this slice operation + pub fn output_size(&self, dim_size: usize) -> usize { + let range = self.to_range(dim_size); + // Handle empty slices (start >= end) + if range.start >= range.end { + return 0; + } + let len = range.end - range.start; + if self.step.unsigned_abs() == 1 { + len + } else { + len.div_ceil(self.step.unsigned_abs()) + } + } +} + +fn convert_signed_index(index: isize, size: usize) -> usize { + if index < 0 { + (size as isize + index).max(0) as usize + } else { + (index as usize).min(size) + } +} + +fn handle_signed_inclusive_end(end: isize) -> Option { + match end { + -1 => None, + end => Some(end + 1), + } +} + +impl From> for Slice { + fn from(r: Range) -> Self { + Self { + start: r.start.as_index(), + end: Some(r.end.as_index()), + step: 1, + } + } +} + +impl From> for Slice { + fn from(r: RangeInclusive) -> Self { + Self { + start: r.start().as_index(), + end: handle_signed_inclusive_end(r.end().as_index()), + step: 1, + } + } +} + +impl From> for Slice { + fn from(r: RangeFrom) -> Self { + Self { + start: r.start.as_index(), + end: None, + step: 1, + } + } +} + +impl From> for Slice { + fn from(r: RangeTo) -> Self { + Self { + start: 0, + end: Some(r.end.as_index()), + step: 1, + } + } +} + +impl From> for Slice { + fn from(r: RangeToInclusive) -> Self { + Self { + start: 0, + end: handle_signed_inclusive_end(r.end.as_index()), + step: 1, + } + } +} + +impl From for Slice { + fn from(_: RangeFull) -> Self { + Self { + start: 0, + end: None, + step: 1, + } + } +} + +impl From for Slice { + fn from(i: usize) -> Self { + Slice::index(i as isize) + } +} + +impl From for Slice { + fn from(i: isize) -> Self { + Slice::index(i) + } +} + +impl From for Slice { + fn from(i: i32) -> Self { + Slice::index(i as isize) + } +} + +impl From for Slice { + fn from(i: i64) -> Self { + Slice::index(i as isize) + } +} + +impl Display for Slice { + fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result { + if self.step == 1 + && let Some(end) = self.end + && self.start == end - 1 + { + f.write_fmt(format_args!("{}", self.start)) + } else { + if self.start != 0 { + f.write_fmt(format_args!("{}", self.start))?; + } + f.write_str("..")?; + if let Some(end) = self.end { + f.write_fmt(format_args!("{}", end))?; + } + if self.step != 1 { + f.write_fmt(format_args!(";{}", self.step))?; + } + Ok(()) + } + } +} + +impl FromStr for Slice { + type Err = crate::ExpressionError; + + fn from_str(source: &str) -> Result { + let mut s = source.trim(); + + let parse_int = |v: &str| -> Result { + v.parse::().map_err(|e| { + crate::ExpressionError::parse_error( + format!("Invalid integer: '{v}': {}", e), + source, + ) + }) + }; + + let mut start: isize = 0; + let mut end: Option = None; + let mut step: isize = 1; + + if let Some((head, tail)) = s.split_once(";") { + step = parse_int(tail)?; + s = head; + } + + if s.is_empty() { + return Err(crate::ExpressionError::parse_error( + "Empty expression", + source, + )); + } + + if let Some((start_s, end_s)) = s.split_once("..") { + if !start_s.is_empty() { + start = parse_int(start_s)?; + } + if !end_s.is_empty() { + if let Some(end_s) = end_s.strip_prefix('=') { + end = Some(parse_int(end_s)? + 1); + } else { + end = Some(parse_int(end_s)?); + } + } + } else { + start = parse_int(s)?; + end = Some(start + 1); + } + + if step == 0 { + return Err(crate::ExpressionError::invalid_expression( + "Step cannot be zero", + source, + )); + } + + Ok(Slice::new(start, end, step)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use alloc::string::ToString; + use alloc::vec; + + #[test] + fn test_slice_to_str() { + assert_eq!(Slice::new(0, None, 1).to_string(), ".."); + + assert_eq!(Slice::new(0, Some(1), 1).to_string(), "0"); + + assert_eq!(Slice::new(0, Some(10), 1).to_string(), "..10"); + assert_eq!(Slice::new(1, Some(10), 1).to_string(), "1..10"); + + assert_eq!(Slice::new(-3, Some(10), -2).to_string(), "-3..10;-2"); + } + + #[test] + fn test_slice_from_str() { + assert_eq!("1".parse::(), Ok(Slice::new(1, Some(2), 1))); + assert_eq!("..".parse::(), Ok(Slice::new(0, None, 1))); + assert_eq!("..3".parse::(), Ok(Slice::new(0, Some(3), 1))); + assert_eq!("..=3".parse::(), Ok(Slice::new(0, Some(4), 1))); + + assert_eq!("-12..3".parse::(), Ok(Slice::new(-12, Some(3), 1))); + assert_eq!("..;-1".parse::(), Ok(Slice::new(0, None, -1))); + + assert_eq!("..=3;-2".parse::(), Ok(Slice::new(0, Some(4), -2))); + + assert_eq!( + "..;0".parse::(), + Err(crate::ExpressionError::invalid_expression( + "Step cannot be zero", + "..;0" + )) + ); + + assert_eq!( + "".parse::(), + Err(crate::ExpressionError::parse_error("Empty expression", "")) + ); + assert_eq!( + "a".parse::(), + Err(crate::ExpressionError::parse_error( + "Invalid integer: 'a': invalid digit found in string", + "a" + )) + ); + assert_eq!( + "..a".parse::(), + Err(crate::ExpressionError::parse_error( + "Invalid integer: 'a': invalid digit found in string", + "..a" + )) + ); + assert_eq!( + "a:b:c".parse::(), + Err(crate::ExpressionError::parse_error( + "Invalid integer: 'a:b:c': invalid digit found in string", + "a:b:c" + )) + ); + } + + #[test] + fn test_slice_output_size() { + // Test the output_size method directly + assert_eq!(Slice::new(0, Some(10), 1).output_size(10), 10); + assert_eq!(Slice::new(0, Some(10), 2).output_size(10), 5); + assert_eq!(Slice::new(0, Some(10), 3).output_size(10), 4); // ceil(10/3) + assert_eq!(Slice::new(0, Some(10), -1).output_size(10), 10); + assert_eq!(Slice::new(0, Some(10), -2).output_size(10), 5); + assert_eq!(Slice::new(2, Some(8), -3).output_size(10), 2); // ceil(6/3) + assert_eq!(Slice::new(5, Some(5), 1).output_size(10), 0); // empty range + } + + #[test] + fn test_bound_to() { + assert_eq!( + Slice::new(0, None, 1).bound_to(10), + Slice::new(0, Some(10), 1) + ); + assert_eq!( + Slice::new(0, Some(5), 1).bound_to(10), + Slice::new(0, Some(5), 1) + ); + + assert_eq!( + Slice::new(0, None, -1).bound_to(10), + Slice::new(0, Some(-11), -1) + ); + assert_eq!( + Slice::new(0, Some(-5), -1).bound_to(10), + Slice::new(0, Some(-5), -1) + ); + } + + #[test] + fn test_slice_iter() { + assert_eq!( + Slice::new(2, Some(3), 1).into_iter().collect::>(), + vec![2] + ); + assert_eq!( + Slice::new(3, Some(-1), -1).into_iter().collect::>(), + vec![3, 2, 1, 0] + ); + + assert_eq!(Slice::new(3, Some(-1), -1).into_vec(), vec![3, 2, 1, 0]); + + assert_eq!( + Slice::new(3, None, 2) + .into_iter() + .take(3) + .collect::>(), + vec![3, 5, 7] + ); + assert_eq!( + Slice::new(3, None, 2) + .bound_to(8) + .into_iter() + .collect::>(), + vec![3, 5, 7] + ); + } + + #[test] + #[should_panic( + expected = "Slice must have an end to convert to a vector: Slice { start: 0, end: None, step: 1 }" + )] + fn test_unbound_slice_into_vec() { + Slice::new(0, None, 1).into_vec(); + } + + #[test] + fn into_slices_should_return_for_all_shape_dims() { + let slice = s![1]; + let shape = Shape::new([2, 3, 1]); + + let slices = slice.into_slices(&shape); + + assert_eq!(slices.len(), shape.len()); + + assert_eq!(slices[0], Slice::new(1, Some(2), 1)); + assert_eq!(slices[1], Slice::new(0, Some(3), 1)); + assert_eq!(slices[2], Slice::new(0, Some(1), 1)); + + let slice = s![1, 0..2]; + let slices = slice.into_slices(&shape); + + assert_eq!(slices.len(), shape.len()); + + assert_eq!(slices[0], Slice::new(1, Some(2), 1)); + assert_eq!(slices[1], Slice::new(0, Some(2), 1)); + assert_eq!(slices[2], Slice::new(0, Some(1), 1)); + + let slice = s![..]; + let slices = slice.into_slices(&shape); + + assert_eq!(slices.len(), shape.len()); + + assert_eq!(slices[0], Slice::new(0, Some(2), 1)); + assert_eq!(slices[1], Slice::new(0, Some(3), 1)); + assert_eq!(slices[2], Slice::new(0, Some(1), 1)); + } + + #[test] + fn into_slices_all_dimensions() { + let slice = s![1, ..2, ..]; + let shape = Shape::new([2, 3, 1]); + + let slices = slice.into_slices(&shape); + + assert_eq!(slices.len(), shape.len()); + + assert_eq!(slices[0], Slice::new(1, Some(2), 1)); + assert_eq!(slices[1], Slice::new(0, Some(2), 1)); + assert_eq!(slices[2], Slice::new(0, Some(1), 1)); + } + + #[test] + fn into_slices_supports_empty_dimensions() { + let slice = s![.., 1, ..]; + let shape = Shape::new([0, 3, 1]); + + let slices = slice.into_slices(&shape); + + assert_eq!(slices.len(), shape.len()); + + assert_eq!(slices[0], Slice::new(0, Some(0), 1)); + assert_eq!(slices[1], Slice::new(1, Some(2), 1)); + assert_eq!(slices[2], Slice::new(0, Some(1), 1)); + } + + #[test] + #[should_panic = "Too many slices provided for shape"] + fn into_slices_should_match_shape_rank() { + let slice = s![.., 1, ..]; + let shape = Shape::new([3, 1]); + + let _ = slice.into_slices(&shape); + } + + #[test] + fn should_support_const_and_full() { + static SLICES: [Slice; 2] = [Slice::full(), Slice::new(2, None, 1)]; + assert_eq!(SLICES[0], Slice::new(0, None, 1)); + assert_eq!(SLICES[1], Slice::new(2, None, 1)); + } + + #[test] + fn should_support_default() { + assert_eq!(Slice::default(), Slice::new(0, None, 1)); + } + + #[test] + fn should_support_copy() { + let mut slice = Slice::new(1, Some(3), 2); + let slice_copy = slice; + + slice.end = Some(4); + + assert_eq!(slice, Slice::new(1, Some(4), 2)); + assert_eq!(slice_copy, Slice::new(1, Some(3), 2)); + } +} diff --git a/crates/burn-std/tests/config.rs b/crates/burn-std/tests/config.rs new file mode 100644 index 0000000..2d919b7 --- /dev/null +++ b/crates/burn-std/tests/config.rs @@ -0,0 +1,80 @@ +/// Tests for loading `burn.toml` from disk. +/// +/// The `cubecl` feature toggles whether `BurnConfig` has a `cubecl` sub-config. We still +/// want `burn.toml` files that contain `[cubecl.*]` sections to load cleanly either way — +/// serde's default behaviour is to ignore unknown top-level fields, and these tests lock +/// that in. +#[cfg(feature = "std")] +mod config_loading { + use burn_std::config::autodiff::AutodiffLogLevel; + use burn_std::config::fusion::FusionLogLevel; + use burn_std::config::{BurnConfig, RuntimeConfig}; + use std::io::Write; + use tempfile::NamedTempFile; + + const EMPTY_TOML: &str = ""; + + const MINIMAL_TOML: &str = r#" +[fusion.beam_search] +max_blocks = 3 +"#; + + const FULL_TOML: &str = r#" +[fusion.beam_search] +max_blocks = 10 + +[fusion.logger] +level = "medium" +stdout = true + +[autodiff.logger] +level = "basic" + +[cubecl.autotune] +level = "full" + +[cubecl.compilation.logger] +level = "full" +stdout = true +"#; + + fn write_toml(content: &str) -> NamedTempFile { + let mut file = NamedTempFile::new().expect("create temp toml"); + file.write_all(content.as_bytes()).expect("write toml"); + file.flush().expect("flush toml"); + file + } + + #[test] + fn load_empty_toml_uses_defaults() { + let file = write_toml(EMPTY_TOML); + let config = BurnConfig::from_file_path(file.path()).expect("parse empty toml"); + + assert_eq!(config.fusion().beam_search.max_blocks, 5); + assert_eq!(config.fusion().logger.level, FusionLogLevel::Disabled); + assert_eq!(config.autodiff().logger.level, AutodiffLogLevel::Disabled); + } + + #[test] + fn load_minimal_toml_overrides_defaults() { + let file = write_toml(MINIMAL_TOML); + let config = BurnConfig::from_file_path(file.path()).expect("parse minimal toml"); + + assert_eq!(config.fusion().beam_search.max_blocks, 3); + // Untouched sections keep defaults. + assert_eq!(config.autodiff().logger.level, AutodiffLogLevel::Disabled); + } + + /// The important case: a single `burn.toml` with both Burn and CubeCL sections must + /// load whether or not the `cubecl` feature is enabled on `burn-std`. + #[test] + fn load_full_toml_ignores_cubecl_when_feature_off() { + let file = write_toml(FULL_TOML); + let config = BurnConfig::from_file_path(file.path()).expect("parse full toml"); + + assert_eq!(config.fusion().beam_search.max_blocks, 10); + assert_eq!(config.fusion().logger.level, FusionLogLevel::Medium); + assert!(config.fusion().logger.stdout); + assert_eq!(config.autodiff().logger.level, AutodiffLogLevel::Basic); + } +} diff --git a/crates/burn-store/Cargo.toml b/crates/burn-store/Cargo.toml new file mode 100644 index 0000000..10a8af0 --- /dev/null +++ b/crates/burn-store/Cargo.toml @@ -0,0 +1,97 @@ +[package] +authors = ["Dilshod Tadjibaev (@antimora)"] +categories = ["science", "no-std", "embedded", "wasm"] +description = "Storage and serialization infrastructure for Burn" +documentation = "https://docs.rs/burn-store" +edition.workspace = true +keywords = [ + "deep-learning", + "machine-learning", + "tensor", + "storage", + "serialization", +] +license.workspace = true +name = "burn-store" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-store" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["std", "pytorch", "safetensors", "burnpack", "memmap"] +memmap = ["std", "dep:memmap2"] +std = [ + "dep:memmap2", + "safetensors/std", + "burn-core/std", + "burn-pack?/std", + "dep:regex", + "byteorder/std", +] + +# The burnpack format lives in the `burn-pack` crate; this enables the BurnpackStore over it. +burnpack = ["dep:burn-pack"] +cuda = ["burn-core/cuda"] +metal = ["wgpu", "burn-core/metal"] +tch = ["burn-core/tch"] +wgpu = ["burn-core/wgpu"] + +safetensors = ["dep:safetensors"] + +pytorch = ["std", "zip", "serde", "tar", "dep:thiserror", "dep:num-traits", "dep:burn-tensor"] + +[dependencies] +burn-core = { workspace = true } +burn-pack = { workspace = true, optional = true } +burn-tensor = { workspace = true, optional = true } + +# External dependencies +byteorder = { workspace = true, default-features = false } +bytes = { workspace = true } +ciborium = { workspace = true, optional = true } +half = { workspace = true } +hashbrown = { workspace = true, features = ["serde"] } +memmap2 = { workspace = true, optional = true } +num-traits = { workspace = true, optional = true } +regex = { workspace = true, optional = true } +serde = { workspace = true, optional = true } +textdistance = { workspace = true } +thiserror = { workspace = true, optional = true } +zip = { workspace = true, optional = true } +tar = { workspace = true, optional = true } + +# Workaround to force broken minor version to update +lzma-rust2 = { workspace = true, optional = true } + +safetensors = { workspace = true, optional = true } + +[dev-dependencies] +# burn-import = { path = "../burn-import", version = "0.22.0" } # disabled (circular dep in publish, only for bench) +burn-core = { workspace = true, features = ["flex"] } +burn-nn = { workspace = true } + +divan = "0.1" +tempfile = { workspace = true } + +[[bench]] +harness = false +name = "resnet18_loading" + +[[bench]] +harness = false +name = "unified_loading" + +[[bench]] +harness = false +name = "unified_saving" + +[[bench]] +harness = false +name = "zero_copy_loading" + +# Enable extra-platforms for bytes on targets without native atomics (e.g., thumbv6m-none-eabi) +[target.'cfg(not(target_has_atomic = "ptr"))'.dependencies] +bytes = { workspace = true, features = ["extra-platforms"] } diff --git a/crates/burn-store/MIGRATION.md b/crates/burn-store/MIGRATION.md new file mode 100644 index 0000000..1d50964 --- /dev/null +++ b/crates/burn-store/MIGRATION.md @@ -0,0 +1,325 @@ +# Migration Guide: burn-import to burn-store + +This guide helps you migrate from the deprecated `burn-import` recorders (`PyTorchFileRecorder`, +`SafetensorsFileRecorder`) to the new `burn-store` API (`PytorchStore`, `SafetensorsStore`). + +## Overview + +The new `burn-store` API provides: + +- **Simpler API**: Load directly into models instead of records +- **Fluent builder pattern**: Chain configuration methods +- **Better error handling**: Detailed load results with applied/missing/errors info +- **Bidirectional support**: Both load and save operations +- **More features**: Filtering, partial loading, metadata, zero-copy loading + +## Quick Migration + +### PyTorch Files (.pt/.pth) + +**Before (burn-import):** + +```rust +use burn::record::{FullPrecisionSettings, Recorder}; +use burn_import::pytorch::{LoadArgs, PyTorchFileRecorder}; + +// Load into a record, then create model from record +let record: ModelRecord = PyTorchFileRecorder::::default() + .load("model.pt".into(), &device) + .expect("Failed to load"); + +let model = Model::init(&device).load_record(record); +``` + +**After (burn-store):** + +```rust +use burn_store::{ModuleSnapshot, PytorchStore}; + +// Initialize model, then load weights directly +let mut model = Model::init(&device); +let mut store = PytorchStore::from_file("model.pt"); +model.load_from(&mut store).expect("Failed to load"); +``` + +### SafeTensors Files (.safetensors) + +**Before (burn-import):** + +```rust +use burn::record::{FullPrecisionSettings, Recorder}; +use burn_import::safetensors::{AdapterType, LoadArgs, SafetensorsFileRecorder}; + +let record: ModelRecord = SafetensorsFileRecorder::::default() + .load("model.safetensors".into(), &device) + .expect("Failed to load"); + +let model = Model::init(&device).load_record(record); +``` + +**After (burn-store):** + +```rust +use burn_store::{ModuleSnapshot, PyTorchToBurnAdapter, SafetensorsStore}; + +let mut model = Model::init(&device); + +// For SafeTensors exported from PyTorch, use the adapter +let mut store = SafetensorsStore::from_file("model.safetensors") + .with_from_adapter(PyTorchToBurnAdapter); +model.load_from(&mut store).expect("Failed to load"); + +// For native Burn SafeTensors, no adapter needed +let mut store = SafetensorsStore::from_file("model.safetensors"); +model.load_from(&mut store).expect("Failed to load"); +``` + +## API Mapping + +### PyTorchFileRecorder Options + +| burn-import | burn-store | +| ---------------------------------------------- | ------------------------------------------- | +| `LoadArgs::new(path)` | `PytorchStore::from_file(path)` | +| `.with_key_remap(pattern, replacement)` | `.with_key_remapping(pattern, replacement)` | +| `.with_top_level_key(key)` | `.with_top_level_key(key)` | +| `.with_debug_print()` | _(use tracing/logging instead)_ | +| `PyTorchFileRecorder::` | _(precision handled automatically)_ | + +### SafetensorsFileRecorder Options + +| burn-import | burn-store | +| -------------------------------------------------- | ------------------------------------------- | +| `LoadArgs::new(path)` | `SafetensorsStore::from_file(path)` | +| `.with_key_remap(pattern, replacement)` | `.with_key_remapping(pattern, replacement)` | +| `.with_adapter_type(AdapterType::PyTorch)` | `.with_from_adapter(PyTorchToBurnAdapter)` | +| `.with_adapter_type(AdapterType::NoAdapter)` | _(default, no adapter)_ | +| `.with_debug_print()` | _(use tracing/logging instead)_ | +| `SafetensorsFileRecorder::` | _(precision handled automatically)_ | + +## Detailed Examples + +### Key Remapping + +**Before:** + +```rust +let args = LoadArgs::new("model.pt".into()) + .with_key_remap("conv\\.(.*)", "$1") + .with_key_remap("^old_prefix\\.", "new_prefix."); + +let record: ModelRecord = PyTorchFileRecorder::::default() + .load(args, &device)?; +``` + +**After:** + +```rust +let mut store = PytorchStore::from_file("model.pt") + .with_key_remapping("conv\\.(.*)", "$1") + .with_key_remapping("^old_prefix\\.", "new_prefix."); + +model.load_from(&mut store)?; +``` + +### Top-Level Key Access + +**Before:** + +```rust +let args = LoadArgs::new("checkpoint.pt".into()) + .with_top_level_key("state_dict"); + +let record: ModelRecord = PyTorchFileRecorder::::default() + .load(args, &device)?; +``` + +**After:** + +```rust +let mut store = PytorchStore::from_file("checkpoint.pt") + .with_top_level_key("state_dict"); + +model.load_from(&mut store)?; +``` + +### PyTorch Adapter for SafeTensors + +**Before:** + +```rust +use burn_import::safetensors::{AdapterType, LoadArgs}; + +let args = LoadArgs::new("pytorch_model.safetensors".into()) + .with_adapter_type(AdapterType::PyTorch); + +let record: ModelRecord = SafetensorsFileRecorder::::default() + .load(args, &device)?; +``` + +**After:** + +```rust +use burn_store::{PyTorchToBurnAdapter, SafetensorsStore}; + +let mut store = SafetensorsStore::from_file("pytorch_model.safetensors") + .with_from_adapter(PyTorchToBurnAdapter); + +model.load_from(&mut store)?; +``` + +## New Features in burn-store + +### Partial Loading + +Handle missing tensors gracefully: + +```rust +let mut store = PytorchStore::from_file("model.pt") + .allow_partial(true); + +let result = model.load_from(&mut store)?; +println!("Loaded: {:?}", result.applied); +println!("Missing: {:?}", result.missing); +``` + +### Filtering + +Load only specific tensors: + +```rust +let mut store = SafetensorsStore::from_file("model.safetensors") + .with_regex(r"^encoder\..*") // Only encoder layers + .allow_partial(true); + +model.load_from(&mut store)?; +``` + +### Saving Models + +Save models (not supported by old recorders): + +```rust +// Save to SafeTensors +let mut store = SafetensorsStore::from_file("output.safetensors") + .metadata("version", "1.0"); +model.save_into(&mut store)?; + +// Save to Burnpack (native format) +let mut store = BurnpackStore::from_file("output.bpk"); +model.save_into(&mut store)?; +``` + +### Load Results + +Get detailed information about loading: + +```rust +let result = model.load_from(&mut store)?; + +// Print the full result for debugging - shows applied, skipped, missing, and errors +println!("{}", result); + +// Or access individual fields +println!("Applied: {} tensors", result.applied.len()); +println!("Skipped: {} tensors", result.skipped.len()); +println!("Missing: {:?}", result.missing); +println!("Errors: {:?}", result.errors); + +// Check if fully successful +if result.is_success() { + println!("All tensors loaded successfully"); +} +``` + +The `LoadResult` implements `Display`, so printing it shows a formatted summary with suggestions for +common issues (e.g., using `allow_partial(true)` for missing tensors). + +## Updating Cargo.toml + +**Before:** + +```toml +[dependencies] +burn-import = { version = "0.x", features = ["pytorch", "safetensors"] } +``` + +**After:** + +```toml +[dependencies] +burn-store = { version = "0.x", features = ["pytorch", "safetensors"] } +``` + +## Common Migration Issues + +### 1. Model vs Record + +The new API loads directly into models, not records. Update your model initialization: + +```rust +// Before: Create record, then model from record +let record = recorder.load(...)?; +let model = Model::init(&device).load_record(record); + +// After: Create model, then load into it +let mut model = Model::init(&device); +model.load_from(&mut store)?; +``` + +### 2. Inference Functions + +If you had functions that took `ModelRecord`, update them to take `Model`: + +```rust +// Before +fn infer(record: ModelRecord) { + let model = Model::init(&device).load_record(record); + // ... +} + +// After +fn infer(model: Model) { + // Model already has weights loaded + // ... +} +``` + +### 3. Precision Settings + +The old API required explicit precision settings. The new API handles this automatically: + +```rust +// Before: Had to specify FullPrecisionSettings or HalfPrecisionSettings +PyTorchFileRecorder::::default() + +// After: Precision handled automatically based on tensor dtype +PytorchStore::from_file("model.pt") +``` + +### 4. Error Handling + +The new API provides richer error information: + +```rust +// Before: Simple Result +let record = recorder.load(args, &device)?; + +// After: LoadResult with detailed info +let result = model.load_from(&mut store)?; + +// Print the result to see a helpful summary with suggestions +println!("{}", result); + +// Or handle specific issues programmatically +if !result.errors.is_empty() { + for (path, error) in &result.errors { + eprintln!("Error loading {}: {}", path, error); + } +} +``` + +## See Also + +- [burn-store README](README.md) - Full documentation +- [import-model-weights example](../../examples/import-model-weights/) - Working example diff --git a/crates/burn-store/README.md b/crates/burn-store/README.md new file mode 100644 index 0000000..e36d7d8 --- /dev/null +++ b/crates/burn-store/README.md @@ -0,0 +1,90 @@ +# Burn Store + +> Advanced model storage and serialization for the Burn deep learning framework + +[![Current Crates.io Version](https://img.shields.io/crates/v/burn-store.svg)](https://crates.io/crates/burn-store) +[![Documentation](https://docs.rs/burn-store/badge.svg)](https://docs.rs/burn-store) + +A comprehensive storage library for Burn that enables efficient model serialization, cross-framework +interoperability, and advanced tensor management. + +> **Migrating from burn-import?** See the [Migration Guide](MIGRATION.md) for help moving from +> `PyTorchFileRecorder`/`SafetensorsFileRecorder` to the new Store API. + +## Features + +- **Burnpack Format** - Native Burn format with CBOR metadata, memory-mapped loading, ParamId + persistence for stateful training, and no-std support +- **SafeTensors Format** - Industry-standard format for secure and efficient tensor serialization +- **PyTorch Support** - Direct loading of PyTorch .pth/.pt files with automatic weight + transformation +- **Zero-Copy Loading** - Memory-mapped files and lazy tensor materialization for optimal + performance +- **Flexible Filtering** - Load/save specific model subsets with regex, exact paths, or custom + predicates +- **Tensor Remapping** - Rename tensors during load/save for framework compatibility +- **Half-Precision Storage** - Automatic F32/F16 conversion with smart defaults for reduced model + file size +- **No-std Support** - Burnpack and SafeTensors formats available in embedded and WASM environments + +## Quick Start + +```rust +use burn_store::{ModuleSnapshot, PytorchStore, SafetensorsStore, BurnpackStore, HalfPrecisionAdapter}; + +// Load from PyTorch +let mut store = PytorchStore::from_file("model.pt"); +model.load_from(&mut store)?; + +// Load from SafeTensors (with PyTorch adapter) +let mut store = SafetensorsStore::from_file("model.safetensors") + .with_from_adapter(PyTorchToBurnAdapter); +model.load_from(&mut store)?; + +// Save to Burnpack +let mut store = BurnpackStore::from_file("model.bpk"); +model.save_into(&mut store)?; + +// Save with half-precision (F32 -> F16, ~50% smaller files) +let adapter = HalfPrecisionAdapter::new(); +let mut store = BurnpackStore::from_file("model_f16.bpk") + .with_to_adapter(adapter.clone()); +model.save_into(&mut store)?; + +// Load half-precision back (F16 -> F32, same adapter) +let mut store = BurnpackStore::from_file("model_f16.bpk") + .with_from_adapter(adapter); +model.load_from(&mut store)?; +``` + +## Documentation + +For comprehensive documentation including: + +- Exporting weights from PyTorch +- Loading weights into Burn models +- Saving models to various formats +- Advanced features (filtering, remapping, partial loading, zero-copy) +- API reference and troubleshooting + +See the **[Burn Book - Saving and Loading](../../burn-book/src/saving-and-loading.md)** chapter. + +## Running Benchmarks + +```bash +# Generate model files (one-time setup) +uv run benches/generate_unified_models.py + +# Run loading benchmarks +cargo bench --bench unified_loading + +# Run saving benchmarks +cargo bench --bench unified_saving + +# With specific backend +cargo bench --bench unified_loading --features metal +``` + +## License + +This project is dual-licensed under MIT and Apache-2.0. diff --git a/crates/burn-store/benches/download_resnet18.py b/crates/burn-store/benches/download_resnet18.py new file mode 100644 index 0000000..6393465 --- /dev/null +++ b/crates/burn-store/benches/download_resnet18.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.8" +# dependencies = [ +# "torch", +# "torchvision", +# ] +# /// +""" +Download ResNet18 PyTorch model for benchmarking. +This script downloads a pre-trained ResNet18 model from PyTorch Hub +and saves it in a format suitable for benchmarking. +""" + +import os +import sys +import tempfile +from pathlib import Path + +import torch +import torchvision.models as models + +def download_resnet18(): + """Download ResNet18 model and save to temp directory.""" + + # Create a temporary directory for the model + temp_dir = Path(tempfile.gettempdir()) / "burn_resnet18_benchmark" + temp_dir.mkdir(parents=True, exist_ok=True) + + output_path = temp_dir / "resnet18.pth" + + # Check if already downloaded + if output_path.exists(): + file_size_mb = output_path.stat().st_size / (1024 * 1024) + print(f"✅ ResNet18 already exists at: {output_path}") + print(f" Size: {file_size_mb:.1f} MB") + return str(output_path) + + print("📥 Downloading ResNet18 model...") + + try: + # Download pre-trained ResNet18 model + model = models.resnet18(pretrained=True) + + # Save the model state dict (this is what burn-store reads) + # Using the legacy format for compatibility + torch.save(model.state_dict(), output_path, _use_new_zipfile_serialization=False) + + file_size_mb = output_path.stat().st_size / (1024 * 1024) + print(f"✅ Successfully downloaded ResNet18 to: {output_path}") + print(f" Size: {file_size_mb:.1f} MB") + print(f" Format: PyTorch legacy format") + + # Verify it's readable + state_dict = torch.load(output_path, map_location='cpu') + print(f" Tensors: {len(state_dict)} tensors") + + # Print a few tensor names and shapes for verification + print("\n Sample tensors:") + for i, (name, tensor) in enumerate(state_dict.items()): + if i < 3: + print(f" - {name}: {list(tensor.shape)}") + + return str(output_path) + + except Exception as e: + print(f"❌ Failed to download ResNet18: {e}") + sys.exit(1) + +def main(): + """Main entry point.""" + path = download_resnet18() + + # Write the path to a file that the benchmark can read + bench_config = Path(tempfile.gettempdir()) / "burn_resnet18_benchmark" / "path.txt" + bench_config.write_text(path) + + print(f"\n💡 Model ready for benchmarking") + print(f" Run: cargo bench --bench resnet18_loading") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/crates/burn-store/benches/generate_unified_models.py b/crates/burn-store/benches/generate_unified_models.py new file mode 100644 index 0000000..6216d22 --- /dev/null +++ b/crates/burn-store/benches/generate_unified_models.py @@ -0,0 +1,175 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.8" +# dependencies = [ +# "torch", +# "safetensors", +# "packaging", +# "numpy", +# ] +# /// +""" +Generate a large model (~312MB) in both PyTorch and SafeTensors formats for unified benchmarking. + +Usage: + uv run benches/generate_unified_models.py + +The script will create model files in /tmp/simple_bench_models/ directory. +""" + +import torch +import torch.nn as nn +import os +from pathlib import Path +import tempfile +from safetensors.torch import save_file + +def get_temp_dir(): + """Get the appropriate temp directory.""" + temp_dir = Path(tempfile.gettempdir()) / "simple_bench_models" + temp_dir.mkdir(parents=True, exist_ok=True) + return temp_dir + +class LargeModel(nn.Module): + """Large model with 20 layers to match Rust benchmark.""" + def __init__(self): + super().__init__() + self.layers = nn.ModuleList() + + # Create a model with 20 layers matching the Rust LargeModel + for i in range(20): + in_size = 1024 if i == 0 else 2048 + out_size = 2048 + self.layers.append(nn.Linear(in_size, out_size)) + + print(f"Created model with {len(self.layers)} layers") + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + +def calculate_model_size(model): + """Calculate the size of the model in MB.""" + total_params = sum(p.numel() for p in model.parameters()) + size_mb = (total_params * 4) / (1024 * 1024) # 4 bytes per float32 + return total_params, size_mb + +def initialize_weights(model): + """Initialize model weights with random values.""" + for param in model.parameters(): + if param.dim() > 1: + nn.init.xavier_uniform_(param) + else: + nn.init.zeros_(param) + +def save_pytorch_format(model, output_dir): + """Save model in PyTorch format.""" + pt_path = output_dir / "large_model.pt" + + # Save as checkpoint with model_state_dict (common format) + checkpoint = { + 'model_state_dict': model.state_dict(), + 'metadata': { + 'model_type': 'large_benchmark_model', + 'num_layers': len(model.layers), + } + } + torch.save(checkpoint, pt_path) + + return pt_path + +def save_safetensors_format(model, output_dir): + """Save model in SafeTensors format.""" + st_path = output_dir / "large_model.safetensors" + + # Convert state dict to safetensors format + state_dict = model.state_dict() + # Ensure all tensors are contiguous and on CPU + state_dict = {k: v.contiguous().cpu() for k, v in state_dict.items()} + + # Save with metadata + metadata = { + 'model_type': 'large_benchmark_model', + 'num_layers': str(len(model.layers)), + } + save_file(state_dict, st_path, metadata=metadata) + + return st_path + +def verify_files(pt_path, st_path): + """Verify the saved files can be loaded.""" + # Verify PyTorch file + checkpoint = torch.load(pt_path, map_location='cpu') + pt_keys = set(checkpoint['model_state_dict'].keys()) + print(f" PyTorch file: {len(pt_keys)} tensors") + + # Verify SafeTensors file + from safetensors import safe_open + with safe_open(st_path, framework="pt", device="cpu") as f: + st_keys = set(f.keys()) + print(f" SafeTensors file: {len(st_keys)} tensors") + + # Check keys match + if pt_keys != st_keys: + print(" ⚠️ Warning: Keys don't match between formats!") + else: + print(" ✓ Keys match between formats") + +def main(): + print("🔧 Generating unified benchmark model files...") + print("") + + output_dir = get_temp_dir() + print(f"📁 Output directory: {output_dir}") + print("") + + # Set random seed for reproducibility + torch.manual_seed(42) + + # Create the large model + print("📝 Creating large model...") + model = LargeModel() + + # Calculate and display model size + total_params, size_mb = calculate_model_size(model) + print(f" Total parameters: {total_params:,}") + print(f" Model size: {size_mb:.2f} MB") + print("") + + # Initialize weights + print("🎲 Initializing weights...") + initialize_weights(model) + + # Save in PyTorch format + print("💾 Saving PyTorch format...") + pt_path = save_pytorch_format(model, output_dir) + pt_size_mb = pt_path.stat().st_size / (1024 * 1024) + print(f" Saved: {pt_path}") + print(f" File size: {pt_size_mb:.2f} MB") + print("") + + # Save in SafeTensors format + print("💾 Saving SafeTensors format...") + st_path = save_safetensors_format(model, output_dir) + st_size_mb = st_path.stat().st_size / (1024 * 1024) + print(f" Saved: {st_path}") + print(f" File size: {st_size_mb:.2f} MB") + print("") + + # Verify files + print("🔍 Verifying saved files...") + verify_files(pt_path, st_path) + print("") + + print(f"✅ Model files generated successfully!") + print("") + print("📊 Summary:") + print(f" PyTorch file: {pt_path.name} ({pt_size_mb:.2f} MB)") + print(f" SafeTensors file: {st_path.name} ({st_size_mb:.2f} MB)") + print("") + print("💡 To run the unified benchmark:") + print(" cargo bench --bench unified_loading") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/crates/burn-store/benches/resnet18_loading.rs b/crates/burn-store/benches/resnet18_loading.rs new file mode 100644 index 0000000..df8604e --- /dev/null +++ b/crates/burn-store/benches/resnet18_loading.rs @@ -0,0 +1,213 @@ +//! Benchmark for ResNet18 loading to verify lazy loading memory usage. +//! +//! resnet18.pth is pytorch's legacy file format. +//! +//! This benchmark loads a ResNet18 model and materializes all tensors +//! to ensure memory usage stays reasonable with lazy loading. +//! +//! Run the benchmark: +//! ```bash +//! cargo bench --bench resnet18_loading +//! ``` + +use burn_store::pytorch::PytorchReader; +use divan::{AllocProfiler, Bencher}; +use std::path::PathBuf; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +#[allow(clippy::manual_range_contains)] +fn main() { + // Check if ResNet18 file exists + let path = resnet18_path(); + if !path.exists() { + eprintln!("❌ ResNet18 model not found!"); + eprintln!(); + eprintln!("Please download it first by running:"); + eprintln!(" python benches/download_resnet18.py"); + eprintln!(); + eprintln!("Or if you don't have Python/PyTorch installed:"); + eprintln!(" uv run benches/download_resnet18.py"); + eprintln!(); + eprintln!("Expected location: {}", path.display()); + std::process::exit(1); + } + + // Verify file size is reasonable + let metadata = std::fs::metadata(&path).expect("Failed to read file metadata"); + let size_mb = metadata.len() as f64 / 1_048_576.0; + + if size_mb < 40.0 || size_mb > 50.0 { + eprintln!( + "⚠️ Warning: ResNet18 file size ({:.1} MB) seems unusual", + size_mb + ); + eprintln!("Expected size is around 45 MB"); + } + + println!("✅ Found ResNet18 model at: {}", path.display()); + println!("📦 File size: {:.1} MB", size_mb); + println!("📊 Running ResNet18 loading benchmarks...\n"); + + // Run divan benchmarks + divan::main(); +} + +/// Get the path to ResNet18 model file +fn resnet18_path() -> PathBuf { + // First try to read from the path file created by download script + let temp_dir = std::env::temp_dir(); + let config_file = temp_dir.join("burn_resnet18_benchmark").join("path.txt"); + + if config_file.exists() + && let Ok(path_str) = std::fs::read_to_string(&config_file) + { + let path = PathBuf::from(path_str.trim()); + if path.exists() { + return path; + } + } + + // Fallback to default location + temp_dir + .join("burn_resnet18_benchmark") + .join("resnet18.pth") +} + +#[divan::bench(sample_count = 10)] +fn load_resnet18_metadata(bencher: Bencher) { + let path = resnet18_path(); + + bencher.bench_local(|| { + let reader = PytorchReader::new(&path).expect("Failed to load ResNet18"); + let metadata = reader.metadata(); + + // Just access metadata without materializing tensors + assert_eq!(metadata.tensor_count, 122); + }); +} + +#[divan::bench(sample_count = 5)] +fn load_resnet18_materialize_all(bencher: Bencher) { + let path = resnet18_path(); + + bencher.bench_local(|| { + let reader = PytorchReader::new(&path).expect("Failed to load ResNet18"); + let keys = reader.keys(); + + let mut total_bytes = 0usize; + + // Materialize all tensors one by one + for key in &keys { + let tensor = reader.get(key).expect("Failed to get tensor"); + // Materialize the tensor data + let _data = tensor.to_data().expect("Failed to materialize tensor data"); + total_bytes += tensor.data_len(); + } + + // Verify we processed all the data + assert!(total_bytes > 40_000_000); // Should be ~45MB + }); +} + +#[divan::bench(sample_count = 5)] +fn load_resnet18_materialize_sequential(bencher: Bencher) { + let path = resnet18_path(); + + bencher.bench_local(|| { + let reader = PytorchReader::new(&path).expect("Failed to load ResNet18"); + let keys = reader.keys(); + + // Materialize tensors one at a time, letting previous ones be dropped + // This simulates processing tensors sequentially without keeping all in memory + for key in &keys { + let tensor = reader.get(key).expect("Failed to get tensor"); + let data = tensor.to_data().expect("Failed to materialize tensor data"); + + // Do minimal work with the data to prevent optimization + let sum = match data.dtype { + burn_core::tensor::DType::F32 => data + .as_slice::() + .map(|s| s.iter().sum::()) + .unwrap_or(0.0) as f64, + burn_core::tensor::DType::F64 => data + .as_slice::() + .map(|s| s.iter().sum::()) + .unwrap_or(0.0), + _ => 0.0, + }; + + // Use the sum to prevent dead code elimination + std::hint::black_box(sum); + } + }); +} + +#[divan::bench(sample_count = 10)] +fn load_resnet18_largest_tensor(bencher: Bencher) { + let path = resnet18_path(); + + bencher.bench_local(|| { + let reader = PytorchReader::new(&path).expect("Failed to load ResNet18"); + + // Find and materialize only the largest tensor + // This tests peak memory for a single tensor operation + let keys = reader.keys(); + let mut largest_key = String::new(); + let mut largest_size = 0usize; + + for key in &keys { + let tensor = reader.get(key).expect("Failed to get tensor"); + let size = tensor.data_len(); + if size > largest_size { + largest_size = size; + largest_key = key.clone(); + } + } + + // Materialize the largest tensor + let tensor = reader + .get(&largest_key) + .expect("Failed to get largest tensor"); + let _data = tensor.to_data().expect("Failed to materialize tensor data"); + + assert!(largest_size > 9_000_000); // Should be ~9MB for layer4.0.conv2.weight + }); +} + +#[divan::bench(sample_count = 10)] +fn load_resnet18_memory_profile(bencher: Bencher) { + let path = resnet18_path(); + + bencher + .with_inputs(|| path.clone()) + .bench_local_values(|path| { + let reader = PytorchReader::new(&path).expect("Failed to load ResNet18"); + let keys = reader.keys(); + + let mut peak_single_tensor = 0usize; + let mut total_data = 0usize; + + // Process each tensor and track memory + for key in &keys { + let tensor = reader.get(key).expect("Failed to get tensor"); + let tensor_size = tensor.data_len(); + + // Track largest single tensor + if tensor_size > peak_single_tensor { + peak_single_tensor = tensor_size; + } + + // Materialize the tensor + let data = tensor.to_data().expect("Failed to materialize tensor data"); + total_data += tensor_size; + + // Drop data immediately to test lazy loading memory efficiency + drop(data); + } + + // Return stats for verification + (peak_single_tensor, total_data) + }); +} diff --git a/crates/burn-store/benches/unified_loading.rs b/crates/burn-store/benches/unified_loading.rs new file mode 100644 index 0000000..60b8efa --- /dev/null +++ b/crates/burn-store/benches/unified_loading.rs @@ -0,0 +1,321 @@ +#![recursion_limit = "256"] + +//! Unified benchmark comparing all loading methods: +//! - BurnpackStore (new native format) +//! - NamedMpkFileRecorder (old native format) +//! - SafetensorsStore (new) +//! - SafetensorsFileRecorder (old) +//! - PytorchStore (new) +//! - PyTorchFileRecorder (old) +//! +//! Before running this benchmark, generate the model files: +//! ```bash +//! cd crates/burn-store +//! uv run benches/generate_unified_models.py +//! ``` +//! +//! Then run the benchmark: +//! ```bash +//! cargo bench --bench unified_loading +//! ``` + +use burn_core as burn; + +use burn_core::module::Module; +use burn_core::prelude::*; +use burn_core::store::ModuleRecord; +// use burn_import::pytorch::{LoadArgs, PyTorchFileRecorder}; +// use burn_import::safetensors::SafetensorsFileRecorder; +use burn_nn as nn; +use burn_store::{ + BurnpackStore, ModuleSnapshot, PyTorchToBurnAdapter, PytorchStore, SafetensorsStore, +}; +use divan::{AllocProfiler, Bencher}; +use std::fs; +use std::path::{Path, PathBuf}; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +// Use the same LargeModel as other benchmarks for fair comparison +#[derive(Module, Debug)] +struct LargeModel { + layers: Vec, +} + +impl LargeModel { + fn new(device: &Device) -> Self { + let mut layers = Vec::new(); + // Create a model with 20 layers - same as safetensor_loading benchmark + for i in 0..20 { + let in_size = if i == 0 { 1024 } else { 2048 }; + layers.push(nn::LinearConfig::new(in_size, 2048).init(device)); + } + Self { layers } + } +} + +/// Get the path to the model files +fn get_model_dir() -> PathBuf { + std::env::temp_dir().join("simple_bench_models") +} + +/// Generate Burnpack and NamedMpk files from existing SafeTensors file +fn generate_burn_formats(st_path: &Path, bp_path: &Path, mpk_path: &Path) { + let device = Device::flex(); + + // Load the model from SafeTensors + let mut model = LargeModel::new(&device); + let mut store = SafetensorsStore::from_file(st_path).with_from_adapter(PyTorchToBurnAdapter); + model + .load_from(&mut store) + .expect("Failed to load from SafeTensors"); + + // Save as Burnpack + if !bp_path.exists() { + println!(" Creating Burnpack file..."); + let mut burnpack_store = BurnpackStore::from_file(bp_path); + model + .save_into(&mut burnpack_store) + .expect("Failed to save as Burnpack"); + } + + // Save as NamedMpk + if !mpk_path.exists() { + println!(" Creating NamedMpk file..."); + model + .save_file(mpk_path) + .expect("Failed to save as NamedMpk"); + } +} + +/// Get paths to the model files +fn get_model_paths() -> (PathBuf, PathBuf, PathBuf, PathBuf) { + let dir = get_model_dir(); + ( + dir.join("large_model.bpk"), + dir.join("large_model.mpk"), + dir.join("large_model.safetensors"), + dir.join("large_model.pt"), + ) +} + +/// Check if model files exist +fn check_model_files() -> Result<(), String> { + let (_, _, st_path, pt_path) = get_model_paths(); + + // For now, only check safetensors and pytorch files (will generate burnpack/mpk later) + if !st_path.exists() || !pt_path.exists() { + return Err(format!( + "\n❌ Model files not found!\n\ + \n\ + Please generate the model files first by running:\n\ + \n\ + cd crates/burn-store\n\ + uv run benches/generate_unified_models.py\n\ + \n\ + Expected files:\n\ + - {}\n\ + - {}\n", + st_path.display(), + pt_path.display() + )); + } + + Ok(()) +} + +fn main() { + // Check if model files exist before running benchmarks + match check_model_files() { + Ok(()) => { + let (bp_path, mpk_path, st_path, pt_path) = get_model_paths(); + + // First, generate Burnpack and MPK files if they don't exist + if !bp_path.exists() || !mpk_path.exists() { + println!("⏳ Generating Burnpack and NamedMpk files from SafeTensors..."); + generate_burn_formats(&st_path, &bp_path, &mpk_path); + } + + let bp_size = fs::metadata(&bp_path) + .ok() + .map(|m| m.len() as f64 / 1_048_576.0); + let mpk_size = fs::metadata(&mpk_path) + .ok() + .map(|m| m.len() as f64 / 1_048_576.0); + let st_size = fs::metadata(&st_path).unwrap().len() as f64 / 1_048_576.0; + let pt_size = fs::metadata(&pt_path).unwrap().len() as f64 / 1_048_576.0; + + println!("✅ Found model files:"); + if let Some(size) = bp_size { + println!(" Burnpack: {} ({:.1} MB)", bp_path.display(), size); + } + if let Some(size) = mpk_size { + println!(" NamedMpk: {} ({:.1} MB)", mpk_path.display(), size); + } + println!(" SafeTensors: {} ({:.1} MB)", st_path.display(), st_size); + println!(" PyTorch: {} ({:.1} MB)", pt_path.display(), pt_size); + println!(); + println!("🚀 Running unified loading benchmarks..."); + println!(); + println!("Comparing 6 loading methods:"); + println!(" 1. BurnpackStore (new native format - lazy loading)"); + println!(" 2. NamedMpkFileRecorder (old native format - loads all to memory)"); + println!(" 3. SafetensorsStore (new)"); + println!(" 4. SafetensorsFileRecorder (old)"); + println!(" 5. PytorchStore (new)"); + println!(" 6. PyTorchFileRecorder (old)"); + println!(); + println!("Available backends:"); + println!(" - Flex (CPU)"); + #[cfg(feature = "wgpu")] + println!(" - WGPU (GPU)"); + #[cfg(feature = "cuda")] + println!(" - CUDA (NVIDIA GPU)"); + #[cfg(feature = "tch")] + println!(" - LibTorch"); + #[cfg(feature = "metal")] + println!(" - Metal (Apple GPU)"); + println!(); + + divan::main(); + } + Err(msg) => { + eprintln!("{}", msg); + std::process::exit(1); + } + } +} + +// Macro to generate benchmarks for each backend +macro_rules! bench_backend { + ($device:expr, $mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name, sample_count = 10)] + mod $mod_name { + use super::*; + + #[divan::bench] + fn burnpack_store(bencher: Bencher) { + let (bp_path, _, _, _) = get_model_paths(); + let file_size = fs::metadata(&bp_path).unwrap().len(); + + bencher + .counter(divan::counter::BytesCount::new(file_size)) + .bench(|| { + let device = $device; + let mut model = LargeModel::new(&device); + let mut store = BurnpackStore::from_file(bp_path.clone()); + model.load_from(&mut store).expect("Failed to load"); + }); + } + + #[divan::bench] + fn namedmpk_recorder(bencher: Bencher) { + let (_, mpk_path, _, _) = get_model_paths(); + let file_size = fs::metadata(&mpk_path).unwrap().len(); + + bencher + .counter(divan::counter::BytesCount::new(file_size)) + .bench(|| { + let device = $device; + let model = LargeModel::new(&device); + model.load_record(ModuleRecord::load(&mpk_path).expect("Failed to load")); + }); + } + + #[divan::bench] + fn safetensors_store(bencher: Bencher) { + let (_, _, st_path, _) = get_model_paths(); + let file_size = fs::metadata(&st_path).unwrap().len(); + + bencher + .counter(divan::counter::BytesCount::new(file_size)) + .bench(|| { + let device = $device; + let mut model = LargeModel::new(&device); + let mut store = SafetensorsStore::from_file(st_path.clone()) + .with_from_adapter(PyTorchToBurnAdapter); + model.load_from(&mut store).expect("Failed to load"); + }); + } + + // #[divan::bench] + // fn safetensors_recorder(bencher: Bencher) { + // let (_, _, st_path, _) = get_model_paths(); + // let file_size = fs::metadata(&st_path).unwrap().len(); + + // bencher + // .counter(divan::counter::BytesCount::new(file_size)) + // .bench(|| { + // let device: Device = $device.into(); + // let recorder = SafetensorsFileRecorder::::default(); + // let record = recorder + // .load(st_path.clone().into(), &device) + // .expect("Failed to load"); + // let _model = LargeModel::new(&device).load_record(record); + // }); + // } + + #[divan::bench] + fn pytorch_store(bencher: Bencher) { + let (_, _, _, pt_path) = get_model_paths(); + let file_size = fs::metadata(&pt_path).unwrap().len(); + + bencher + .counter(divan::counter::BytesCount::new(file_size)) + .bench(|| { + let device: Device = $device.into(); + let mut model = LargeModel::new(&device); + let mut store = PytorchStore::from_file(pt_path.clone()) + .with_top_level_key("model_state_dict") + .allow_partial(true); + model.load_from(&mut store).expect("Failed to load"); + }); + } + + // #[divan::bench] + // fn pytorch_recorder(bencher: Bencher) { + // let (_, _, _, pt_path) = get_model_paths(); + // let file_size = fs::metadata(&pt_path).unwrap().len(); + + // bencher + // .counter(divan::counter::BytesCount::new(file_size)) + // .bench(|| { + // let device: Device = $device.into(); + // let recorder = PyTorchFileRecorder::::default(); + // let load_args = + // LoadArgs::new(pt_path.clone()).with_top_level_key("model_state_dict"); + // let record = recorder.load(load_args, &device).expect("Failed to load"); + // let _model = LargeModel::new(&device).load_record(record); + // }); + // } + } + }; +} + +// Generate benchmarks for each backend +bench_backend!(Device::flex(), ndarray_backend, "NdArray Backend (CPU)"); + +#[cfg(feature = "wgpu")] +bench_backend!( + Device::wgpu(Default::default()), + wgpu_backend, + "WGPU Backend (GPU)" +); + +#[cfg(feature = "cuda")] +bench_backend!( + CudaDevice::default(), + cuda_backend, + "CUDA Backend (NVIDIA GPU)" +); + +#[cfg(feature = "tch")] +bench_backend!(Device::libtorch(), tch_backend, "LibTorch Backend"); + +#[cfg(feature = "metal")] +bench_backend!( + Device::wgpu(Default::default()), + metal_backend, + "Metal Backend (Apple GPU)" +); diff --git a/crates/burn-store/benches/unified_saving.rs b/crates/burn-store/benches/unified_saving.rs new file mode 100644 index 0000000..a4c31f5 --- /dev/null +++ b/crates/burn-store/benches/unified_saving.rs @@ -0,0 +1,175 @@ +#![recursion_limit = "256"] + +//! Unified benchmark comparing all saving methods: +//! - BurnpackStore (new native format) +//! - NamedMpkFileRecorder (old native format) +//! - SafetensorsStore (new) +//! +//! Before running this benchmark, ensure the directory exists: +//! ```bash +//! mkdir -p /tmp/simple_bench_models +//! ``` +//! +//! Then run the benchmark: +//! ```bash +//! cargo bench --bench unified_saving +//! ``` +use burn_core as burn; + +use burn_core::module::Module; +use burn_core::prelude::*; +use burn_nn as nn; +use burn_store::{BurnpackStore, ModuleSnapshot, SafetensorsStore}; +use divan::{AllocProfiler, Bencher}; +use std::fs; +use std::path::PathBuf; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +// Use the same LargeModel as other benchmarks for fair comparison +#[derive(Module, Debug)] +struct LargeModel { + layers: Vec, +} + +impl LargeModel { + fn new(device: &Device) -> Self { + let mut layers = Vec::new(); + // Create a model with 20 layers - same as loading benchmarks + for i in 0..20 { + let in_size = if i == 0 { 1024 } else { 2048 }; + layers.push(nn::LinearConfig::new(in_size, 2048).init(device)); + } + Self { layers } + } +} + +/// Get the path to the output directory +fn get_output_dir() -> PathBuf { + std::env::temp_dir().join("simple_bench_models_saving") +} + +/// Ensure output directory exists +fn ensure_output_dir() -> Result<(), String> { + let dir = get_output_dir(); + if !dir.exists() { + fs::create_dir_all(&dir) + .map_err(|e| format!("Failed to create output directory: {}", e))?; + } + Ok(()) +} + +fn main() { + match ensure_output_dir() { + Ok(()) => { + println!("✅ Output directory ready: {}", get_output_dir().display()); + println!(); + println!("🚀 Running unified saving benchmarks..."); + println!(); + println!("Comparing 3 saving methods:"); + println!(" 1. BurnpackStore (new native format)"); + println!(" 2. NamedMpkFileRecorder (old native format)"); + println!(" 3. SafetensorsStore (new)"); + println!(); + println!("Available backends:"); + println!(" - Flex (CPU)"); + #[cfg(feature = "wgpu")] + println!(" - WGPU (GPU)"); + #[cfg(feature = "cuda")] + println!(" - CUDA (NVIDIA GPU)"); + #[cfg(feature = "tch")] + println!(" - LibTorch"); + #[cfg(feature = "metal")] + println!(" - Metal (Apple GPU)"); + println!(); + + divan::main(); + } + Err(msg) => { + eprintln!("❌ {}", msg); + std::process::exit(1); + } + } +} + +// Macro to generate benchmarks for each backend +macro_rules! bench_backend { + ($device:expr, $mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name, sample_count = 10)] + mod $mod_name { + use super::*; + + #[divan::bench] + fn burnpack_store(bencher: Bencher) { + bencher.bench(|| { + let device = $device; + let model = LargeModel::new(&device); + let output_path = get_output_dir().join("test_burnpack.bpk"); + let mut store = BurnpackStore::from_file(output_path.clone()).overwrite(true); + model + .save_into(&mut store) + .expect("Failed to save with BurnpackStore"); + // Clean up + let _ = fs::remove_file(output_path); + }); + } + + #[divan::bench] + fn namedmpk_recorder(bencher: Bencher) { + bencher.bench(|| { + let device = $device; + let model = LargeModel::new(&device); + let output_path = get_output_dir().join("test_namedmpk.mpk"); + model + .save_file(output_path.clone()) + .expect("Failed to save with NamedMpkFileRecorder"); + // Clean up + let _ = fs::remove_file(output_path); + }); + } + + #[divan::bench] + fn safetensors_store(bencher: Bencher) { + bencher.bench(|| { + let device = $device; + let model = LargeModel::new(&device); + let output_path = get_output_dir().join("test_safetensors_store.safetensors"); + let mut store = SafetensorsStore::from_file(output_path.clone()); + model + .save_into(&mut store) + .expect("Failed to save with SafetensorsStore"); + // Clean up + let _ = fs::remove_file(output_path); + }); + } + } + }; +} + +// Generate benchmarks for each backend +bench_backend!(Device::flex(), ndarray_backend, "NdArray Backend (CPU)"); + +#[cfg(feature = "wgpu")] +bench_backend!( + Device::wgpu(Default::default()), + wgpu_backend, + "WGPU Backend (GPU)" +); + +#[cfg(feature = "cuda")] +bench_backend!( + CudaDevice::default(), + cuda_backend, + "CUDA Backend (NVIDIA GPU)" +); + +#[cfg(feature = "tch")] +bench_backend!(Device::libtorch(), tch_backend, "LibTorch Backend"); + +#[cfg(feature = "metal")] +bench_backend!( + Device::wgpu(Default::default()), + metal_backend, + "Metal Backend (Apple GPU)" +); diff --git a/crates/burn-store/benches/zero_copy_loading.rs b/crates/burn-store/benches/zero_copy_loading.rs new file mode 100644 index 0000000..59cb54c --- /dev/null +++ b/crates/burn-store/benches/zero_copy_loading.rs @@ -0,0 +1,483 @@ +#![recursion_limit = "256"] + +//! Benchmark comparing different loading modes for BurnpackStore. +//! +//! This benchmark measures the performance difference between: +//! - `from_file()` - File-based loading (reader handles lazy/mmap internally) +//! - `from_static()` - Zero-copy mode from static bytes (stays in .rodata) +//! - `from_bytes()` - Loading from owned bytes +//! - `from_bytes()` with shared `Bytes` - Loading from shared/arc bytes +//! +//! ## Understanding the Results +//! +//! The key difference is how tensor data reaches the backend: +//! - **from_file**: File → reader (lazy/mmap) → backend +//! - **from_bytes**: Owned bytes → copy to backend +//! - **from_static**: Static .rodata → zero-copy slice → backend +//! - **shared bytes**: Arc'd bytes → zero-copy clone → backend +//! +//! GPU backends that can consume `Bytes` directly will show larger benefits +//! from zero-copy paths. +//! +//! ## Running the benchmark +//! +//! Before running this benchmark, generate the model files: +//! ```bash +//! cd crates/burn-store +//! uv run benches/generate_unified_models.py +//! ``` +//! +//! Then run the benchmark: +//! ```bash +//! cargo bench --bench zero_copy_loading +//! ``` + +use burn_core as burn; + +use burn_core::module::Module; +use burn_core::prelude::*; +use burn_core::tensor::{AllocationProperty, Bytes}; +use burn_nn as nn; +use burn_store::{ + BurnpackStore, ModuleSnapshot, ModuleStore, PyTorchToBurnAdapter, SafetensorsStore, +}; +use divan::{AllocProfiler, Bencher}; +use std::fs; +use std::path::PathBuf; +use std::sync::OnceLock; + +#[global_allocator] +static ALLOC: AllocProfiler = AllocProfiler::system(); + +// Static storage for embedded model bytes (simulating include_bytes!) +static STATIC_MODEL_BYTES: OnceLock<&'static [u8]> = OnceLock::new(); + +// Use the same LargeModel as other benchmarks for fair comparison +#[derive(Module, Debug)] +struct LargeModel { + layers: Vec, +} + +impl LargeModel { + fn new(device: &Device) -> Self { + let mut layers = Vec::new(); + // Create a model with 20 layers - same as unified_loading benchmark + for i in 0..20 { + let in_size = if i == 0 { 1024 } else { 2048 }; + layers.push(nn::LinearConfig::new(in_size, 2048).init(device)); + } + Self { layers } + } +} + +/// Get the path to the model files +fn get_model_dir() -> PathBuf { + std::env::temp_dir().join("simple_bench_models") +} + +/// Get path to Burnpack model file +fn get_burnpack_path() -> PathBuf { + get_model_dir().join("large_model.bpk") +} + +/// Generate Burnpack file from existing SafeTensors file if needed +fn ensure_burnpack_file() { + let bp_path = get_burnpack_path(); + let st_path = get_model_dir().join("large_model.safetensors"); + + if bp_path.exists() { + return; + } + + if !st_path.exists() { + panic!( + "\n❌ SafeTensors model file not found!\n\ + \n\ + Please generate the model files first by running:\n\ + \n\ + cd crates/burn-store\n\ + uv run benches/generate_unified_models.py\n\ + \n\ + Expected file: {}\n", + st_path.display() + ); + } + + println!("⏳ Generating Burnpack file from SafeTensors..."); + + let device = Device::flex(); + + // Load from SafeTensors + let mut model = LargeModel::new(&device); + let mut store = SafetensorsStore::from_file(&st_path).with_from_adapter(PyTorchToBurnAdapter); + model + .load_from(&mut store) + .expect("Failed to load from SafeTensors"); + + // Save as Burnpack + let mut burnpack_store = BurnpackStore::from_file(&bp_path); + model + .save_into(&mut burnpack_store) + .expect("Failed to save as Burnpack"); + + println!("✅ Created Burnpack file: {}", bp_path.display()); +} + +/// Initialize static model bytes (simulating include_bytes! at runtime for benchmarks) +fn get_static_model_bytes() -> &'static [u8] { + STATIC_MODEL_BYTES.get_or_init(|| { + let bp_path = get_burnpack_path(); + let bytes = fs::read(&bp_path).expect("Failed to read Burnpack file"); + // Leak the bytes to get a 'static lifetime (acceptable for benchmarks) + Box::leak(bytes.into_boxed_slice()) + }) +} + +fn main() { + // Ensure Burnpack file exists + ensure_burnpack_file(); + + let bp_path = get_burnpack_path(); + let file_size = fs::metadata(&bp_path).unwrap().len() as f64 / 1_048_576.0; + + println!("✅ Found Burnpack model file:"); + println!(" Path: {}", bp_path.display()); + println!(" Size: {:.1} MB", file_size); + println!(); + println!("🚀 Running loading mode benchmarks..."); + println!(); + println!("Comparing loading modes:"); + println!(" 1. file - from_file() - lazy/mmap file loading"); + println!(" 2. static_bytes - from_bytes() with Vec copy from static"); + println!(" 3. static_zero_copy - from_static() - zero-copy from static"); + println!(" 4. memory_shared - from_bytes() with shared Bytes (cheap Arc clone)"); + println!(); + println!("Available backends:"); + println!(" - Flex (CPU)"); + #[cfg(feature = "wgpu")] + println!(" - WGPU (GPU)"); + #[cfg(feature = "cuda")] + println!(" - CUDA (NVIDIA GPU)"); + #[cfg(feature = "tch")] + println!(" - LibTorch"); + #[cfg(feature = "metal")] + println!(" - Metal (Apple GPU)"); + println!(); + + // Pre-initialize static bytes before benchmarks + let _ = get_static_model_bytes(); + + divan::main(); +} + +// Macro to generate benchmarks for each backend +macro_rules! bench_backend { + ($device:expr, $mod_name:ident, $backend_name:literal) => { + #[divan::bench_group(name = $backend_name, sample_count = 10)] + mod $mod_name { + use super::*; + + /// File-based loading (reader handles lazy/mmap internally) + #[divan::bench] + fn file(bencher: Bencher) { + let bp_path = get_burnpack_path(); + let file_size = fs::metadata(&bp_path).unwrap().len(); + + bencher + .counter(divan::counter::BytesCount::new(file_size)) + .bench(|| { + let device = $device; + let mut model = LargeModel::new(&device); + let mut store = BurnpackStore::from_file(&bp_path); + model.load_from(&mut store).expect("Failed to load"); + }); + } + + /// Static bytes with copy (simulating old behavior: copy to Vec first) + #[divan::bench] + fn static_bytes(bencher: Bencher) { + let static_bytes = get_static_model_bytes(); + let file_size = static_bytes.len() as u64; + + bencher + .counter(divan::counter::BytesCount::new(file_size)) + .bench(|| { + let device = $device; + let mut model = LargeModel::new(&device); + + let bytes = Bytes::from_bytes_vec(static_bytes.to_vec()); + let mut store = BurnpackStore::from_bytes(Some(bytes)); + model.load_from(&mut store).expect("Failed to load"); + }); + } + + /// Static bytes with zero-copy (from_static keeps data in .rodata) + #[divan::bench] + fn static_zero_copy(bencher: Bencher) { + let static_bytes = get_static_model_bytes(); + let file_size = static_bytes.len() as u64; + + bencher + .counter(divan::counter::BytesCount::new(file_size)) + .bench(|| { + let device = $device; + let mut model = LargeModel::new(&device); + + let mut store = BurnpackStore::from_static(static_bytes); + model.load_from(&mut store).expect("Failed to load"); + }); + } + + /// In-memory shared bytes (cheap Arc clone, zero-copy) + #[divan::bench] + fn memory_shared(bencher: Bencher) { + let static_bytes = get_static_model_bytes(); + let file_size = static_bytes.len() as u64; + + // Pre-create shared bytes outside the benchmark loop + let shared = bytes::Bytes::from_static(static_bytes); + + bencher + .counter(divan::counter::BytesCount::new(file_size)) + .bench(|| { + let device = $device; + let mut model = LargeModel::new(&device); + + let bytes = Bytes::from_shared(shared.clone(), AllocationProperty::Other); + let mut store = BurnpackStore::from_bytes(Some(bytes)); + model.load_from(&mut store).expect("Failed to load"); + }); + } + } + }; +} + +// ============================================================================= +// Zero-copy verification (proves operations use static region data) +// ============================================================================= + +/// Verify that zero-copy loading actually uses data from the static region. +/// This runs once at startup to prove correctness before benchmarking. +#[divan::bench_group(name = "Zero-Copy Verification", sample_count = 1)] +mod verification { + use super::*; + + /// Verify data is readable and correct using sum().into_scalar(). + /// Note: sum() triggers COW copy, so this shows ops work correctly on zero-copy data. + #[divan::bench] + fn verify_ops_produce_correct_results() { + let static_bytes = get_static_model_bytes(); + + let device = Device::flex(); + let mut model = LargeModel::new(&device); + let mut store = BurnpackStore::from_static(static_bytes); + model.load_from(&mut store).expect("Failed to load"); + + // Compute sum of first layer weight - proves data is valid + let weight = model.layers[0].weight.val(); + let sum: f32 = weight.sum().into_scalar(); + + assert!(sum.is_finite(), "Sum should be finite"); + println!("✅ Verified: Operations on zero-copy data produce valid results"); + println!(" - First layer sum: {:.4}", sum); + } + + /// Verify operations produce correct results on zero-copy data + #[divan::bench] + fn verify_operations_on_static_data() { + let static_bytes = get_static_model_bytes(); + + // Load model with zero-copy + let device = Device::flex(); + let mut model = LargeModel::new(&device); + let mut store = BurnpackStore::from_static(static_bytes); + model.load_from(&mut store).expect("Failed to load"); + + // Perform operations on the loaded weights + let weight = model.layers[0].weight.val(); + let shape = weight.shape(); + + // Test 1: Sum should be finite (not NaN or Inf) + let sum: f32 = weight.clone().sum().to_data().to_vec().unwrap()[0]; + assert!( + sum.is_finite(), + "Operation failed: sum is not finite ({})", + sum + ); + + // Test 2: Matrix multiply with itself transposed (W @ W.T) + let transposed = weight.clone().transpose(); + let matmul_result = weight.clone().matmul(transposed); + let matmul_sum: f32 = matmul_result.sum().to_data().to_vec().unwrap()[0]; + assert!( + matmul_sum.is_finite(), + "Matmul failed: result sum is not finite ({})", + matmul_sum + ); + + // Test 3: Element-wise operations + let doubled = weight.clone() * 2.0; + let doubled_sum: f32 = doubled.sum().to_data().to_vec().unwrap()[0]; + assert!( + (doubled_sum - sum * 2.0).abs() < 1e-3, + "Element-wise op failed: doubled_sum ({}) != sum*2 ({})", + doubled_sum, + sum * 2.0 + ); + + println!("✅ Verified: Operations on zero-copy data produce correct results"); + println!(" - Weight shape: {:?}", shape.as_slice()); + println!(" - Sum: {:.4}", sum); + println!(" - Matmul result sum: {:.4}", matmul_sum); + } + + /// Compare from_static vs from_bytes: verify both produce identical results + #[divan::bench] + fn verify_static_vs_bytes_equality() { + let static_bytes = get_static_model_bytes(); + let device = Device::flex(); + + // Load with from_static (zero-copy) + let mut model_zc = LargeModel::new(&device); + let mut store_zc = BurnpackStore::from_static(static_bytes); + model_zc + .load_from(&mut store_zc) + .expect("Failed to load zero-copy"); + + // Load with from_bytes (copy) + let mut model_copy = LargeModel::new(&device); + let bytes = Bytes::from_bytes_vec(static_bytes.to_vec()); + let mut store_copy = BurnpackStore::from_bytes(Some(bytes)); + model_copy + .load_from(&mut store_copy) + .expect("Failed to load copy"); + + // Compare weights from both models + for (i, (layer_zc, layer_copy)) in model_zc + .layers + .iter() + .zip(model_copy.layers.iter()) + .enumerate() + { + let weight_zc = layer_zc.weight.val(); + let weight_copy = layer_copy.weight.val(); + + // Check shapes match + assert_eq!( + weight_zc.shape(), + weight_copy.shape(), + "Layer {} weight shapes don't match", + i + ); + + // Check values match (using sum as a proxy) + let sum_zc: f32 = weight_zc.clone().sum().to_data().to_vec().unwrap()[0]; + let sum_copy: f32 = weight_copy.clone().sum().to_data().to_vec().unwrap()[0]; + assert!( + (sum_zc - sum_copy).abs() < 1e-6, + "Layer {} weight sums don't match: zero-copy={}, copy={}", + i, + sum_zc, + sum_copy + ); + } + + println!( + "✅ Verified: Zero-copy and copy loading produce identical results for all {} layers", + model_zc.layers.len() + ); + } +} + +// ============================================================================= +// Store-only benchmarks (no backend allocation overhead) +// These show the TRUE zero-copy benefit at the store level +// ============================================================================= + +#[divan::bench_group(name = "Store Only (no backend)", sample_count = 10)] +mod store_only { + use super::*; + + /// File-based store - measures store overhead only + #[divan::bench] + fn file(bencher: Bencher) { + let bp_path = get_burnpack_path(); + let file_size = fs::metadata(&bp_path).unwrap().len(); + + bencher + .counter(divan::counter::BytesCount::new(file_size)) + .bench(|| { + let mut store = BurnpackStore::from_file(&bp_path); + let snapshots = store.get_all_snapshots().expect("Failed to get snapshots"); + for snapshot in snapshots.values() { + let _data = snapshot.to_data().expect("Failed to get tensor data"); + } + }); + } + + /// Static bytes with copy - measures store overhead only + #[divan::bench] + fn static_bytes(bencher: Bencher) { + let static_bytes = get_static_model_bytes(); + let file_size = static_bytes.len() as u64; + + bencher + .counter(divan::counter::BytesCount::new(file_size)) + .bench(|| { + let bytes = Bytes::from_bytes_vec(static_bytes.to_vec()); + let mut store = BurnpackStore::from_bytes(Some(bytes)); + let snapshots = store.get_all_snapshots().expect("Failed to get snapshots"); + for snapshot in snapshots.values() { + let _data = snapshot.to_data().expect("Failed to get tensor data"); + } + }); + } + + /// Static bytes with zero-copy - measures store overhead only + #[divan::bench] + fn static_zero_copy(bencher: Bencher) { + let static_bytes = get_static_model_bytes(); + let file_size = static_bytes.len() as u64; + + bencher + .counter(divan::counter::BytesCount::new(file_size)) + .bench(|| { + let mut store = BurnpackStore::from_static(static_bytes); + let snapshots = store.get_all_snapshots().expect("Failed to get snapshots"); + for snapshot in snapshots.values() { + let _data = snapshot.to_data().expect("Failed to get tensor data"); + } + }); + } +} + +// ============================================================================= +// Full model loading benchmarks (includes backend allocation) +// ============================================================================= + +// Generate benchmarks for each backend +bench_backend!(Device::flex(), ndarray_backend, "NdArray Backend (CPU)"); + +#[cfg(feature = "wgpu")] +bench_backend!( + Device::wgpu(Default::default()), + wgpu_backend, + "WGPU Backend (GPU)" +); + +#[cfg(feature = "cuda")] +bench_backend!( + Device::cuda(Default::default()), + cuda_backend, + "CUDA Backend (NVIDIA GPU)" +); + +#[cfg(feature = "tch")] +bench_backend!(LibTorchDevice::default(), tch_backend, "LibTorch Backend"); + +#[cfg(feature = "metal")] +bench_backend!( + Device::wgpu(Default::default()), + metal_backend, + "Metal Backend (Apple GPU)" +); diff --git a/crates/burn-store/examples/burnpack_inspect.rs b/crates/burn-store/examples/burnpack_inspect.rs new file mode 100644 index 0000000..0f02cc4 --- /dev/null +++ b/crates/burn-store/examples/burnpack_inspect.rs @@ -0,0 +1,144 @@ +//! Example: Generate a Burnpack file for inspection +//! +//! This example creates a simple Burnpack file that you can examine to understand the format. +//! +//! Usage: +//! cargo run --example burnpack-inspect [output_path] +//! +//! Example: +//! cargo run --example burnpack-inspect sample.bpk +//! cargo run --example burnpack-inspect /tmp/test.bpk +//! +//! After generating the file, examine it with: +//! hexdump -C sample.bpk | head -100 +//! xxd sample.bpk | head -100 +//! hexyl sample.bpk +use burn_core::{self as burn, tensor::Device}; + +use burn_core::module::Module; +use burn_nn::{Linear, LinearConfig}; +use burn_store::{BurnpackStore, ModuleSnapshot}; +use std::env; + +// Simple model with a few layers +#[derive(Module, Debug)] +struct SampleModel { + linear1: Linear, + linear2: Linear, + linear3: Linear, +} + +impl SampleModel { + fn new(device: &Device) -> Self { + Self { + linear1: LinearConfig::new(128, 64).init(device), + linear2: LinearConfig::new(64, 32).init(device), + linear3: LinearConfig::new(32, 10).init(device), + } + } +} + +fn main() { + // Get output path from command line or use default + let output_path = env::args() + .nth(1) + .unwrap_or_else(|| "sample.bpk".to_string()); + + println!("Creating sample Burnpack file: {}", output_path); + println!(); + + // Create a simple model + let device = Default::default(); + let model = SampleModel::new(&device); + + // Save to Burnpack format with metadata + let mut store = BurnpackStore::from_file(&output_path) + .overwrite(true) + .metadata("format", "burnpack") + .metadata("description", "Sample file for examining Burnpack format") + .metadata("version", env!("CARGO_PKG_VERSION")) + .metadata("author", "Burn Example"); + + model.save_into(&mut store).expect("Failed to save model"); + + println!("✅ Successfully created: {}", output_path); + println!(); + println!("📋 File Structure:"); + println!(" ┌─────────────────────────────────────┐"); + println!(" │ Header (10 bytes) │"); + println!(" ├─────────────────────────────────────┤"); + println!(" │ - Magic: 0x4E525542 (BURN in LE) │"); + println!(" │ - Version: 0x0001 (2 bytes) │"); + println!(" │ - Metadata size: (4 bytes, u32 LE) │"); + println!(" ├─────────────────────────────────────┤"); + println!(" │ Metadata (CBOR format) │"); + println!(" ├─────────────────────────────────────┤"); + println!(" │ - Tensor descriptors │"); + println!(" │ * name, dtype, shape, offsets │"); + println!(" │ - User metadata │"); + println!(" ├─────────────────────────────────────┤"); + println!(" │ Tensor Data (raw bytes, LE) │"); + println!(" ├─────────────────────────────────────┤"); + println!(" │ - linear1.weight [64, 128] │"); + println!(" │ - linear1.bias [64] │"); + println!(" │ - linear2.weight [32, 64] │"); + println!(" │ - linear2.bias [32] │"); + println!(" │ - linear3.weight [10, 32] │"); + println!(" │ - linear3.bias [10] │"); + println!(" └─────────────────────────────────────┘"); + println!(); + println!("📊 Model Contents:"); + println!(" - linear1.weight: [64, 128] = 8,192 params → 32,768 bytes"); + println!(" - linear1.bias: [64] = 64 params → 256 bytes"); + println!(" - linear2.weight: [32, 64] = 2,048 params → 8,192 bytes"); + println!(" - linear2.bias: [32] = 32 params → 128 bytes"); + println!(" - linear3.weight: [10, 32] = 320 params → 1,280 bytes"); + println!(" - linear3.bias: [10] = 10 params → 40 bytes"); + println!(" ───────────────────────────────────────────────────────"); + + let total_params = 8192 + 64 + 2048 + 32 + 320 + 10; + let total_bytes = total_params * 4; + println!( + " Total: {} parameters = {} KB", + total_params, + total_bytes / 1024 + ); + println!(); + + // Get actual file size + if let Ok(metadata) = std::fs::metadata(&output_path) { + let file_size = metadata.len(); + println!( + "📦 File size: {} bytes ({:.2} KB)", + file_size, + file_size as f64 / 1024.0 + ); + } + + println!(); + println!("🔍 Inspection Commands:"); + println!(); + println!(" # View first 100 bytes in hex:"); + println!(" hexdump -C {} | head -20", output_path); + println!(); + println!(" # View header only (10 bytes):"); + println!(" head -c 10 {} | hexdump -C", output_path); + println!(); + println!(" # View with prettier hex viewer (if installed):"); + println!(" hexyl {} | head -50", output_path); + println!(); + println!(" # View in binary format:"); + println!(" xxd -b {} | head -20", output_path); + println!(); + println!(" # Extract and examine header:"); + println!(" # Magic (bytes 0-3): Should be 42 55 52 4E (BURN)"); + println!(" # Version (bytes 4-5): Should be 01 00"); + println!(" # Metadata size (bytes 6-9): u32 little-endian"); + println!(); + println!(" # Load back the model:"); + println!( + " # let mut store = BurnpackStore::from_file(\"{}\");", + output_path + ); + println!(" # model.load_from(&mut store)?;"); +} diff --git a/crates/burn-store/examples/half_precision.rs b/crates/burn-store/examples/half_precision.rs new file mode 100644 index 0000000..edb325b --- /dev/null +++ b/crates/burn-store/examples/half_precision.rs @@ -0,0 +1,86 @@ +//! Example: Save and load a model with half-precision (F32 <-> F16) +//! +//! Demonstrates using HalfPrecisionAdapter to automatically convert between +//! F32 and F16 during saving/loading. The same adapter instance handles both +//! directions. +//! +//! Usage: +//! cargo run -p burn-store --example half_precision + +use burn_core::module::Module; +use burn_core::{self as burn, tensor::Device}; +use burn_nn::{LayerNorm, LayerNormConfig, Linear, LinearConfig}; +use burn_store::{BurnpackStore, HalfPrecisionAdapter, ModuleSnapshot}; + +// A model with mixed layer types to show selective conversion +#[derive(Module, Debug)] +struct DemoModel { + linear1: Linear, + norm: LayerNorm, + linear2: Linear, +} + +impl DemoModel { + fn new(device: &Device) -> Self { + Self { + linear1: LinearConfig::new(128, 64).init(device), + norm: LayerNormConfig::new(64).init(device), + linear2: LinearConfig::new(64, 10).init(device), + } + } +} + +fn main() { + let device = Default::default(); + let model = DemoModel::new(&device); + + // 1) Save at full F32 precision (baseline) + let dir = tempfile::tempdir().expect("Failed to create temp dir"); + let path_f32 = dir.path().join("model_f32"); + let path_f16 = dir.path().join("model_f16"); + let path_mixed = dir.path().join("model_mixed"); + + let mut store = BurnpackStore::from_file(path_f32.to_str().unwrap()).overwrite(true); + model.save_into(&mut store).expect("Failed to save F32"); + let size_f32 = std::fs::metadata(format!("{}.bpk", path_f32.display())) + .map(|m| m.len()) + .unwrap_or(0); + + // 2) Save with default half-precision (all default modules get F16) + let adapter = HalfPrecisionAdapter::new(); + let mut store = BurnpackStore::from_file(path_f16.to_str().unwrap()) + .overwrite(true) + .with_to_adapter(adapter.clone()); + model.save_into(&mut store).expect("Failed to save F16"); + let size_f16 = std::fs::metadata(format!("{}.bpk", path_f16.display())) + .map(|m| m.len()) + .unwrap_or(0); + + // 3) Save with without_module: keep LayerNorm at F32 + let adapter_no_norm = HalfPrecisionAdapter::new().without_module("LayerNorm"); + let mut store = BurnpackStore::from_file(path_mixed.to_str().unwrap()) + .overwrite(true) + .with_to_adapter(adapter_no_norm); + model.save_into(&mut store).expect("Failed to save mixed"); + let size_mixed = std::fs::metadata(format!("{}.bpk", path_mixed.display())) + .map(|m| m.len()) + .unwrap_or(0); + + println!("F32 (full precision): {} bytes", size_f32); + println!("F16 (default modules): {} bytes", size_f16); + println!("Mixed (norm stays F32): {} bytes", size_mixed); + println!( + "F16 savings: {:.1}%", + (1.0 - size_f16 as f64 / size_f32 as f64) * 100.0 + ); + + // 4) Round-trip: load the F16 file back to F32 with the same adapter + let mut load_store = + BurnpackStore::from_file(path_f16.to_str().unwrap()).with_from_adapter(adapter); + let mut model2 = DemoModel::new(&device); + let result = model2.load_from(&mut load_store).expect("Failed to load"); + println!( + "\nRound-trip loaded {} tensors successfully", + result.applied.len() + ); +} diff --git a/crates/burn-store/pytorch-tests/Cargo.toml b/crates/burn-store/pytorch-tests/Cargo.toml new file mode 100644 index 0000000..0ef49f9 --- /dev/null +++ b/crates/burn-store/pytorch-tests/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "pytorch-tests" +version.workspace = true +edition.workspace = true +license.workspace = true + +[dev-dependencies] +burn = { workspace = true, features = ["default", "flex"] } +burn-store = { workspace = true, features = ["default", "std", "pytorch"] } +serde = { workspace = true } +float-cmp = { workspace = true } diff --git a/crates/burn-store/pytorch-tests/src/lib.rs b/crates/burn-store/pytorch-tests/src/lib.rs new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/crates/burn-store/pytorch-tests/src/lib.rs @@ -0,0 +1 @@ + diff --git a/crates/burn-store/pytorch-tests/tests/batch_norm/batch_norm2d.pt b/crates/burn-store/pytorch-tests/tests/batch_norm/batch_norm2d.pt new file mode 100644 index 0000000..abb1088 Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/batch_norm/batch_norm2d.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/batch_norm/export_weights.py b/crates/burn-store/pytorch-tests/tests/batch_norm/export_weights.py new file mode 100755 index 0000000..24a65c9 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/batch_norm/export_weights.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.norm1 = nn.BatchNorm2d(5) + + def forward(self, x): + x = self.norm1(x) + return x + + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(1) + + model = Model().to(torch.device("cpu")) + + # Condition batch norm (each forward will affect the running stats) + x1 = torch.ones(1, 5, 2, 2) - 0.5 + _ = model(x1) + model.eval() # Set to eval mode to freeze running stats + # Save the model after the first forward + torch.save(model.state_dict(), "batch_norm2d.pt") + + x2 = torch.ones(1, 5, 2, 2) - 0.3 + print("Input shape: {}", x2.shape) + output = model(x2) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + + + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/batch_norm/mod.rs b/crates/burn-store/pytorch-tests/tests/batch_norm/mod.rs new file mode 100644 index 0000000..2d5540e --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/batch_norm/mod.rs @@ -0,0 +1,61 @@ +use burn::{ + module::Module, + nn::{BatchNorm, BatchNormConfig}, + tensor::{Device, Tensor}, +}; + +#[derive(Module, Debug)] +pub struct Net { + norm1: BatchNorm, +} + +impl Net { + pub fn new(device: &Device) -> Self { + Self { + norm1: BatchNormConfig::new(5).init(device), // Python model uses BatchNorm2d(5) + } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + self.norm1.forward(x) + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::Tolerance; + use burn_store::{ModuleSnapshot, PytorchStore}; + + use super::*; + + #[test] + fn batch_norm2d() { + let device = Default::default(); + let mut model = Net::new(&device); + let mut store = PytorchStore::from_file("tests/batch_norm/batch_norm2d.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + let input = Tensor::<4>::ones([1, 5, 2, 2], &device) - 0.3; + + let output = model.forward(input); + + let expected = Tensor::<4>::from_data( + [[ + [[0.68515635, 0.68515635], [0.68515635, 0.68515635]], + [[0.68515635, 0.68515635], [0.68515635, 0.68515635]], + [[0.68515635, 0.68515635], [0.68515635, 0.68515635]], + [[0.68515635, 0.68515635], [0.68515635, 0.68515635]], + [[0.68515635, 0.68515635], [0.68515635, 0.68515635]], + ]], + &device, + ); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::default()); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/boolean/boolean.pt b/crates/burn-store/pytorch-tests/tests/boolean/boolean.pt new file mode 100644 index 0000000..f75eb2a Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/boolean/boolean.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/boolean/export_weights.py b/crates/burn-store/pytorch-tests/tests/boolean/export_weights.py new file mode 100755 index 0000000..1b2e020 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/boolean/export_weights.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + buffer = torch.tensor([True, False, True]) + self.register_buffer("buffer", buffer, persistent=True) + + def forward(self, x): + x = self.buffer + return x + + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(1) + + model = Model().to(torch.device("cpu")) + + torch.save(model.state_dict(), "boolean.pt") + + input = torch.ones(3, 3) + print("Input shape: {}", input.shape) + print("Input: {}", input) + output = model(input) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + + + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/boolean/mod.rs b/crates/burn-store/pytorch-tests/tests/boolean/mod.rs new file mode 100644 index 0000000..7387a63 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/boolean/mod.rs @@ -0,0 +1,53 @@ +use burn::{ + module::{Module, Param, ParamId}, + tensor::{Bool, Device, Tensor, TensorData}, +}; + +#[derive(Module, Debug)] +pub struct Net { + buffer: Param>, +} + +impl Net { + /// Create a new model with placeholder values. + pub fn init(device: &Device) -> Self { + Self { + buffer: Param::initialized( + ParamId::new(), + Tensor::from_bool(TensorData::from([false, false, false]), device), + ), + } + } + + /// Forward pass of the model. + pub fn forward(&self, _x: Tensor<2>) -> Tensor<1, Bool> { + self.buffer.val() + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::TensorData; + use burn_store::{ModuleSnapshot, PytorchStore}; + + use super::*; + + #[test] + fn boolean() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/boolean/boolean.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + let input = Tensor::<2>::ones([3, 3], &device); + + let output = model.forward(input); + + let expected = Tensor::<1, Bool>::from_bool(TensorData::from([true, false, true]), &device); + + assert_eq!(output.to_data(), expected.to_data()); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/buffer/buffer.pt b/crates/burn-store/pytorch-tests/tests/buffer/buffer.pt new file mode 100644 index 0000000..5a80651 Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/buffer/buffer.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/buffer/export_weights.py b/crates/burn-store/pytorch-tests/tests/buffer/export_weights.py new file mode 100755 index 0000000..ecd678c --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/buffer/export_weights.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + buffer = torch.ones(3, 3) + self.register_buffer("buffer", buffer, persistent=True) + + def forward(self, x): + x = self.buffer + x + return x + + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(1) + + model = Model().to(torch.device("cpu")) + + torch.save(model.state_dict(), "buffer.pt") + + input = torch.ones(3, 3) + print("Input shape: {}", input.shape) + print("Input: {}", input) + output = model(input) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + + + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/buffer/mod.rs b/crates/burn-store/pytorch-tests/tests/buffer/mod.rs new file mode 100644 index 0000000..8e10afe --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/buffer/mod.rs @@ -0,0 +1,52 @@ +use burn::{ + module::{Module, Param}, + tensor::{Device, Tensor}, +}; + +#[derive(Module, Debug)] +pub struct Net { + buffer: Param>, +} + +impl Net { + /// Create a new model with placeholder values. + pub fn init(device: &Device) -> Self { + Self { + buffer: Param::from_tensor(Tensor::zeros([3, 3], device)), + } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<2>) -> Tensor<2> { + self.buffer.val() + x + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::Tolerance; + use burn_store::{ModuleSnapshot, PytorchStore}; + + use super::*; + + #[test] + fn buffer() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/buffer/buffer.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + let input = Tensor::<2>::ones([3, 3], &device); + + let output = model.forward(input); + + let expected = Tensor::<2>::ones([3, 3], &device) * 2.0; + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::default()); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/complex_nested/complex_nested.pt b/crates/burn-store/pytorch-tests/tests/complex_nested/complex_nested.pt new file mode 100644 index 0000000..1fca425 Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/complex_nested/complex_nested.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/complex_nested/export_weights.py b/crates/burn-store/pytorch-tests/tests/complex_nested/export_weights.py new file mode 100755 index 0000000..9bae3bf --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/complex_nested/export_weights.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class ConvBlock(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size): + super(ConvBlock, self).__init__() + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size) + self.norm = nn.BatchNorm2d(out_channels) + + def forward(self, x): + x = self.conv(x) + x = self.norm(x) + return x + +class Net(nn.Module): + def __init__(self): + super(Net, self).__init__() + + self.conv_blocks = nn.Sequential( + ConvBlock(2, 4, (3, 2)), + ConvBlock(4, 6, (3, 2)), + ) + self.norm1 = nn.BatchNorm2d(6) + + self.fc1 = nn.Linear(120, 12) + self.fc2 = nn.Linear(12, 10) + + def forward(self, x): + x = self.conv_blocks(x) + x = self.norm1(x) + x = torch.flatten(x, 1) + x = self.fc1(x) + x = F.relu(x) + x = self.fc2(x) + x = F.log_softmax(x, dim=1) + return x + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(2) + + + model = Net().to(torch.device("cpu")) + + # Condition the model (batch norm requires a forward pass to compute the mean and variance) + x1 = torch.ones(1, 2, 9, 6) - 0.1 + x2 = torch.ones(1, 2, 9, 6) - 0.3 + output = model(x1) + output = model(x2) + model.eval() # set to eval mode + + torch.save(model.state_dict(), "complex_nested.pt") + + # feed test data + x = torch.ones(1, 2, 9, 6) - 0.5 + output = model(x) + print("Input shape: {}", x.shape) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + + + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/complex_nested/mod.rs b/crates/burn-store/pytorch-tests/tests/complex_nested/mod.rs new file mode 100644 index 0000000..01b67b9 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/complex_nested/mod.rs @@ -0,0 +1,234 @@ +use burn::tensor::Tolerance; +use burn::{ + module::Module, + nn::{ + BatchNorm, BatchNormConfig, Linear, LinearConfig, + conv::{Conv2d, Conv2dConfig}, + }, + tensor::{ + Device, Tensor, + activation::{log_softmax, relu}, + }, +}; +use burn_store::{ModuleSnapshot, PytorchStore}; + +#[derive(Module, Debug)] +pub struct ConvBlock { + conv: Conv2d, + norm: BatchNorm, +} + +#[derive(Module, Debug)] +pub struct Net { + conv_blocks: Vec, + norm1: BatchNorm, + fc1: Linear, + fc2: Linear, +} + +impl Net { + pub fn init(device: &Device) -> Self { + let conv_blocks = vec![ + ConvBlock { + conv: Conv2dConfig::new([2, 4], [3, 2]).init(device), + norm: BatchNormConfig::new(4).init(device), // matches conv output channels + }, + ConvBlock { + conv: Conv2dConfig::new([4, 6], [3, 2]).init(device), + norm: BatchNormConfig::new(6).init(device), // matches conv output channels + }, + ]; + let norm1 = BatchNormConfig::new(6).init(device); + let fc1 = LinearConfig::new(120, 12).init(device); + let fc2 = LinearConfig::new(12, 10).init(device); + + Self { + conv_blocks, + norm1, + fc1, + fc2, + } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<4>) -> Tensor<2> { + let x = self.conv_blocks[0].forward(x); + let x = self.conv_blocks[1].forward(x); + let x = self.norm1.forward(x); + let x = x.reshape([0, -1]); + let x = self.fc1.forward(x); + let x = relu(x); + let x = self.fc2.forward(x); + + log_softmax(x, 1) + } +} + +impl ConvBlock { + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + let x = self.conv.forward(x); + + self.norm.forward(x) + } +} + +/// Partial model to test loading of partial records. +#[derive(Module, Debug)] +pub struct PartialNet { + conv1: ConvBlock, +} + +impl PartialNet { + /// Create a new model from the given record. + pub fn init(device: &Device) -> Self { + let conv1 = ConvBlock { + conv: Conv2dConfig::new([2, 4], [3, 2]).init(device), + norm: BatchNormConfig::new(4).init(device), // matches conv output channels + }; + Self { conv1 } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + self.conv1.forward(x) + } +} + +/// Model with extra fields to test loading of records (e.g. from a different model). +#[derive(Module, Debug)] +pub struct PartialWithExtraNet { + conv1: ConvBlock, + extra_field: bool, // This field is not present in the pytorch model +} + +impl PartialWithExtraNet { + /// Create a new model from the given record. + pub fn init(device: &Device) -> Self { + let conv1 = ConvBlock { + conv: Conv2dConfig::new([2, 4], [3, 2]).init(device), + norm: BatchNormConfig::new(4).init(device), // matches conv output channels + }; + + Self { + conv1, + extra_field: true, + } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + self.conv1.forward(x) + } +} + +fn model_test(model: Net, precision: f32) { + let device = Default::default(); + + let input = Tensor::<4>::ones([1, 2, 9, 6], &device) - 0.5; + + let output = model.forward(input); + + let expected = Tensor::<2>::from_data( + [[ + -2.306_613, + -2.058_945_4, + -2.298_372_7, + -2.358_294, + -2.296_395_5, + -2.416_090_5, + -2.107_669, + -2.428_420_8, + -2.526_469, + -2.319_918_6, + ]], + &device, + ); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::absolute(precision)); +} + +#[test] +fn full_record() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/complex_nested/complex_nested.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + model_test(model, 1e-8); +} + +#[test] +fn full_record_autodiff() { + let device = Device::default().autodiff(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/complex_nested/complex_nested.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); +} + +#[test] +fn half_record() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/complex_nested/complex_nested.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + model_test(model, 1e-4); +} + +#[test] +fn partial_model_loading() { + let device = Default::default(); + let mut model = PartialNet::init(&device); + + // Load the full model but rename "conv_blocks.0.*" to "conv1.*" + let mut store = PytorchStore::from_file("tests/complex_nested/complex_nested.pt") + .with_key_remapping("conv_blocks\\.0\\.(.*)", "conv1.$1") + .allow_partial(true); + + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + let input = Tensor::<4>::ones([1, 2, 9, 6], &device) - 0.5; + + let output = model.forward(input); + + // get the sum of all elements in the output tensor for quick check + let sum = output.sum(); + + assert!((sum.into_scalar::() - 4.871538).abs() < 0.000002); +} + +#[test] +fn extra_field_model_loading() { + let device = Default::default(); + let mut model = PartialWithExtraNet::init(&device); + + // Load the full model but rename "conv_blocks.0.*" to "conv1.*" + let mut store = PytorchStore::from_file("tests/complex_nested/complex_nested.pt") + .with_key_remapping("conv_blocks\\.0\\.(.*)", "conv1.$1") + .allow_partial(true); + + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + let input = Tensor::<4>::ones([1, 2, 9, 6], &device) - 0.5; + + let output = model.forward(input); + + // get the sum of all elements in the output tensor for quick check + let sum = output.sum(); + + assert!((sum.into_scalar::() - 4.871538).abs() < 0.000002); + + assert!(model.extra_field); +} diff --git a/crates/burn-store/pytorch-tests/tests/config/export_weights.py b/crates/burn-store/pytorch-tests/tests/config/export_weights.py new file mode 100755 index 0000000..8dd194b --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/config/export_weights.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.fc1 = nn.Linear(2, 3) + self.fc2 = nn.Linear(3, 4, bias=False) + + def forward(self, x): + x = self.fc1(x) + x = F.relu(x) # Add relu so that PyTorch optimizer does not combine fc1 and fc2 + x = self.fc2(x) + + return x + +CONFIG = { + "n_head": 2, + "n_layer": 3, + "d_model": 512, + "some_float": 0.1, + "some_int": 1, + "some_bool": True, + "some_str": "hello", + "some_list_int": [1, 2, 3], + "some_list_str": ["hello", "world"], + "some_list_float": [0.1, 0.2, 0.3], + "some_dict": { + "some_key": "some_value" + } +} + +class ModelWithBias(nn.Module): + def __init__(self): + super(ModelWithBias, self).__init__() + self.fc1 = nn.Linear(2, 3) + + def forward(self, x): + x = self.fc1(x) + + return x + + +def main(): + + model = Model().to(torch.device("cpu")) + + weights_with_config = { + "my_model": model.state_dict(), + "my_config": CONFIG + } + + torch.save(weights_with_config, "weights_with_config.pt") + + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/config/mod.rs b/crates/burn-store/pytorch-tests/tests/config/mod.rs new file mode 100644 index 0000000..6c8deb3 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/config/mod.rs @@ -0,0 +1,53 @@ +#![allow(clippy::too_many_arguments)] // To mute derive Config warning +use std::collections::HashMap; + +use burn::config::Config; + +#[allow(clippy::too_many_arguments)] +#[derive(Debug, PartialEq, Config)] +struct NetConfig { + n_head: usize, + n_layer: usize, + d_model: usize, + some_float: f64, + some_int: i32, + some_bool: bool, + some_str: String, + some_list_int: Vec, + some_list_str: Vec, + some_list_float: Vec, + some_dict: HashMap, +} + +#[cfg(test)] +mod tests { + use burn_store::pytorch::PytorchReader; + + use super::*; + + #[test] + fn test_net_config() { + let config_expected = NetConfig { + n_head: 2, + n_layer: 3, + d_model: 512, + some_float: 0.1, + some_int: 1, + some_bool: true, + some_str: "hello".to_string(), + some_list_int: vec![1, 2, 3], + some_list_str: vec!["hello".to_string(), "world".to_string()], + some_list_float: vec![0.1, 0.2, 0.3], + some_dict: { + let mut map = HashMap::new(); + map.insert("some_key".to_string(), "some_value".to_string()); + map + }, + }; + let path = "tests/config/weights_with_config.pt"; + let top_level_key = Some("my_config"); + let config: NetConfig = PytorchReader::load_config(path, top_level_key).unwrap(); + + assert_eq!(config, config_expected); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/config/weights_with_config.pt b/crates/burn-store/pytorch-tests/tests/config/weights_with_config.pt new file mode 100644 index 0000000..3c19e80 Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/config/weights_with_config.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/conv1d/conv1d.pt b/crates/burn-store/pytorch-tests/tests/conv1d/conv1d.pt new file mode 100644 index 0000000..b1c43ee Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/conv1d/conv1d.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/conv1d/export_weights.py b/crates/burn-store/pytorch-tests/tests/conv1d/export_weights.py new file mode 100755 index 0000000..bda180e --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/conv1d/export_weights.py @@ -0,0 +1,39 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.conv1 = nn.Conv1d(2, 2, 2) + self.conv2 = nn.Conv1d(2, 2, 2, bias=False) + + def forward(self, x): + x = self.conv1(x) + x = self.conv2(x) + return x + + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(1) + + model = Model().to(torch.device("cpu")) + + torch.save(model.state_dict(), "conv1d.pt") + + input = torch.rand(1, 2, 6) + print("Input shape: {}", input.shape) + print("Input: {}", input) + output = model(input) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + + + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/conv1d/mod.rs b/crates/burn-store/pytorch-tests/tests/conv1d/mod.rs new file mode 100644 index 0000000..236bd31 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/conv1d/mod.rs @@ -0,0 +1,97 @@ +use burn::{ + module::Module, + nn::conv::{Conv1d, Conv1dConfig}, + tensor::{Device, Tensor}, +}; + +#[derive(Module, Debug)] +pub struct Net { + conv1: Conv1d, + conv2: Conv1d, +} + +impl Net { + /// Create a new model from the given record. + pub fn init(device: &Device) -> Self { + let conv1 = Conv1dConfig::new(2, 2, 2).init(device); + let conv2 = Conv1dConfig::new(2, 2, 2).with_bias(false).init(device); + + Self { conv1, conv2 } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<3>) -> Tensor<3> { + let x = self.conv1.forward(x); + + self.conv2.forward(x) + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::Tolerance; + use burn_store::{ModuleSnapshot, PytorchStore}; + type FT = f32; + + use super::*; + + fn conv1d(model: Net, precision: f32) { + let device = Default::default(); + + let input = Tensor::<3>::from_data( + [[ + [ + 0.93708336, 0.65559506, 0.31379688, 0.19801933, 0.41619217, 0.28432965, + ], + [ + 0.33977574, + 0.523_940_8, + 0.798_063_9, + 0.77176833, + 0.01122457, + 0.80996025, + ], + ]], + &device, + ); + + let output = model.forward(input); + + let expected = Tensor::<3>::from_data( + [[ + [0.02987457, 0.03134188, 0.04234261, -0.02437721], + [-0.03788019, -0.02972012, -0.00806090, -0.01981254], + ]], + &device, + ); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::absolute(precision)); + } + + #[test] + fn conv1d_full_precision() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/conv1d/conv1d.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + conv1d(model, 1e-7); + } + + #[test] + fn conv1d_half_precision() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/conv1d/conv1d.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + conv1d(model, 1e-4); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/conv2d/conv2d.pt b/crates/burn-store/pytorch-tests/tests/conv2d/conv2d.pt new file mode 100644 index 0000000..3f19ad2 Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/conv2d/conv2d.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/conv2d/export_weights.py b/crates/burn-store/pytorch-tests/tests/conv2d/export_weights.py new file mode 100755 index 0000000..cdc3f76 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/conv2d/export_weights.py @@ -0,0 +1,39 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.conv1 = nn.Conv2d(2, 2, (2,2)) + self.conv2 = nn.Conv2d(2, 2, (2,2), bias=False) + + def forward(self, x): + x = self.conv1(x) + x = self.conv2(x) + return x + + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(1) + + model = Model().to(torch.device("cpu")) + + torch.save(model.state_dict(), "conv2d.pt") + + input = torch.rand(1, 2, 5, 5) + print("Input shape: {}", input.shape) + print("Input: {}", input) + output = model(input) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + + + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/conv2d/mod.rs b/crates/burn-store/pytorch-tests/tests/conv2d/mod.rs new file mode 100644 index 0000000..fca78a7 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/conv2d/mod.rs @@ -0,0 +1,133 @@ +use burn::{ + module::Module, + nn::conv::{Conv2d, Conv2dConfig}, + tensor::{Device, Tensor}, +}; + +#[derive(Module, Debug)] +pub struct Net { + conv1: Conv2d, + conv2: Conv2d, +} + +impl Net { + /// Create a new model from the given record. + pub fn init(device: &Device) -> Self { + let conv1 = Conv2dConfig::new([2, 2], [2, 2]).init(device); + let conv2 = Conv2dConfig::new([2, 2], [2, 2]) + .with_bias(false) + .init(device); + + Self { conv1, conv2 } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + let x = self.conv1.forward(x); + + self.conv2.forward(x) + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::Tolerance; + use burn_store::{ModuleSnapshot, PytorchStore}; + + use super::*; + + fn conv2d(model: Net, precision: f32) { + let device = Default::default(); + + let input = Tensor::<4>::from_data( + [[ + [ + [ + 0.024_595_8, + 0.25883394, + 0.93905586, + 0.416_715_5, + 0.713_979_7, + ], + [0.267_644_3, 0.990_609, 0.28845078, 0.874_962_4, 0.505_920_8], + [0.23659128, 0.757_007_4, 0.23458993, 0.64705235, 0.355_621_4], + [0.445_182_8, 0.01930594, 0.26160914, 0.771_317, 0.37846136], + [ + 0.99802476, + 0.900_794_2, + 0.476_588_2, + 0.16625845, + 0.804_481_1, + ], + ], + [ + [ + 0.65517855, + 0.17679012, + 0.824_772_3, + 0.803_550_9, + 0.943_447_5, + ], + [0.21972018, 0.417_697, 0.49031407, 0.57302874, 0.12054086], + [0.14518881, 0.772_002_3, 0.38275403, 0.744_236_7, 0.52850497], + [0.664_172_4, 0.60994434, 0.681_799_7, 0.74785537, 0.03694397], + [ + 0.751_675_7, + 0.148_438_4, + 0.12274551, + 0.530_407_2, + 0.414_796_4, + ], + ], + ]], + &device, + ); + + let output = model.forward(input); + + let expected = Tensor::<4>::from_data( + [[ + [ + [-0.02502128, 0.00250649, 0.04841233], + [0.04589614, -0.00296854, 0.01991477], + [0.02920526, 0.059_497_3, 0.04326791], + ], + [ + [-0.04825336, 0.080_190_9, -0.02375088], + [0.02885434, 0.09638263, -0.07460806], + [0.02004079, 0.06244051, 0.035_887_1], + ], + ]], + &device, + ); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::absolute(precision)); + } + + #[test] + fn conv2d_full_precision() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/conv2d/conv2d.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + conv2d(model, 1e-7); + } + + #[test] + fn conv2d_half_precision() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/conv2d/conv2d.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + conv2d(model, 1e-4); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/conv_transpose1d/conv_transpose1d.pt b/crates/burn-store/pytorch-tests/tests/conv_transpose1d/conv_transpose1d.pt new file mode 100644 index 0000000..7929d6e Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/conv_transpose1d/conv_transpose1d.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/conv_transpose1d/export_weights.py b/crates/burn-store/pytorch-tests/tests/conv_transpose1d/export_weights.py new file mode 100755 index 0000000..a15250c --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/conv_transpose1d/export_weights.py @@ -0,0 +1,39 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.conv1 = nn.ConvTranspose1d(2, 2, 2) + self.conv2 = nn.ConvTranspose1d(2, 2, 2, bias=False) + + def forward(self, x): + x = self.conv1(x) + x = self.conv2(x) + return x + + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(1) + + model = Model().to(torch.device("cpu")) + + torch.save(model.state_dict(), "conv_transpose1d.pt") + + input = torch.rand(1, 2, 2) + print("Input shape: {}", input.shape) + print("Input: {}", input) + output = model(input) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + + + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/conv_transpose1d/mod.rs b/crates/burn-store/pytorch-tests/tests/conv_transpose1d/mod.rs new file mode 100644 index 0000000..27ef048 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/conv_transpose1d/mod.rs @@ -0,0 +1,86 @@ +use burn::{ + module::Module, + nn::conv::{ConvTranspose1d, ConvTranspose1dConfig}, + tensor::{Device, Tensor}, +}; + +#[derive(Module, Debug)] +pub struct Net { + conv1: ConvTranspose1d, + conv2: ConvTranspose1d, +} + +impl Net { + /// Create a new model from the given record. + pub fn init(device: &Device) -> Self { + let conv1 = ConvTranspose1dConfig::new([2, 2], 2).init(device); + let conv2 = ConvTranspose1dConfig::new([2, 2], 2) + .with_bias(false) + .init(device); + + Self { conv1, conv2 } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<3>) -> Tensor<3> { + let x = self.conv1.forward(x); + + self.conv2.forward(x) + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::Tolerance; + use burn_store::{ModuleSnapshot, PytorchStore}; + + use super::*; + + fn conv_transpose1d(model: Net, precision: f32) { + let device = Default::default(); + + let input = Tensor::<3>::from_data( + [[[0.93708336, 0.65559506], [0.31379688, 0.19801933]]], + &device, + ); + + let output = model.forward(input); + + let expected = Tensor::<3>::from_data( + [[ + [0.02935525, 0.01119324, -0.01356167, -0.00682688], + [0.01644749, -0.01429807, 0.00083987, 0.00279229], + ]], + &device, + ); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::absolute(precision)); + } + + #[test] + fn conv_transpose1d_full() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/conv_transpose1d/conv_transpose1d.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + conv_transpose1d(model, 1e-8); + } + + #[test] + fn conv_transpose1d_half() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/conv_transpose1d/conv_transpose1d.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + conv_transpose1d(model, 1e-4); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/conv_transpose2d/conv_transpose2d.pt b/crates/burn-store/pytorch-tests/tests/conv_transpose2d/conv_transpose2d.pt new file mode 100644 index 0000000..07b0fbe Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/conv_transpose2d/conv_transpose2d.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/conv_transpose2d/export_weights.py b/crates/burn-store/pytorch-tests/tests/conv_transpose2d/export_weights.py new file mode 100755 index 0000000..8dc5ff8 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/conv_transpose2d/export_weights.py @@ -0,0 +1,39 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.conv1 = nn.ConvTranspose2d(2, 2, (2, 2)) + self.conv2 = nn.ConvTranspose2d(2, 2, (2, 2), bias=False) + + def forward(self, x): + x = self.conv1(x) + x = self.conv2(x) + return x + + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(1) + + model = Model().to(torch.device("cpu")) + + torch.save(model.state_dict(), "conv_transpose2d.pt") + + input = torch.rand(1, 2, 2, 2) + print("Input shape: {}", input.shape) + print("Input: {}", input) + output = model(input) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + + + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/conv_transpose2d/mod.rs b/crates/burn-store/pytorch-tests/tests/conv_transpose2d/mod.rs new file mode 100644 index 0000000..6c2050e --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/conv_transpose2d/mod.rs @@ -0,0 +1,98 @@ +use burn::{ + module::Module, + nn::conv::{ConvTranspose2d, ConvTranspose2dConfig}, + tensor::{Device, Tensor}, +}; + +#[derive(Module, Debug)] +pub struct Net { + conv1: ConvTranspose2d, + conv2: ConvTranspose2d, +} + +impl Net { + /// Create a new model from the given record. + pub fn init(device: &Device) -> Self { + let conv1 = ConvTranspose2dConfig::new([2, 2], [2, 2]).init(device); + let conv2 = ConvTranspose2dConfig::new([2, 2], [2, 2]) + .with_bias(false) + .init(device); + + Self { conv1, conv2 } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + let x = self.conv1.forward(x); + + self.conv2.forward(x) + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::Tolerance; + use burn_store::{ModuleSnapshot, PytorchStore}; + + use super::*; + + fn conv_transpose2d(model: Net, precision: f32) { + let device = Default::default(); + + let input = Tensor::<4>::from_data( + [[ + [[0.024_595_8, 0.25883394], [0.93905586, 0.416_715_5]], + [[0.713_979_7, 0.267_644_3], [0.990_609, 0.28845078]], + ]], + &device, + ); + + let output = model.forward(input); + + let expected = Tensor::<4>::from_data( + [[ + [ + [0.04547675, 0.01879685, -0.01636661, 0.00310803], + [0.02090115, 0.01192738, -0.048_240_2, 0.02252235], + [0.03249975, -0.00460748, 0.05003899, 0.04029131], + [0.02185687, -0.10226749, -0.06508022, -0.01267705], + ], + [ + [0.00277598, -0.00513832, -0.059_048_3, 0.00567626], + [-0.03149522, -0.195_757_4, 0.03474613, 0.01997269], + [-0.10096474, 0.00679589, 0.041_919_7, -0.02464108], + [-0.03174751, 0.02963913, -0.02703723, -0.01860938], + ], + ]], + &device, + ); + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::absolute(precision)); + } + + #[test] + fn conv_transpose2d_full() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/conv_transpose2d/conv_transpose2d.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + conv_transpose2d(model, 1e-7); + } + + #[test] + fn conv_transpose2d_half() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/conv_transpose2d/conv_transpose2d.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + conv_transpose2d(model, 1e-4); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/debug_test.pt b/crates/burn-store/pytorch-tests/tests/debug_test.pt new file mode 100644 index 0000000..da686a3 Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/debug_test.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/embedding/embedding.pt b/crates/burn-store/pytorch-tests/tests/embedding/embedding.pt new file mode 100644 index 0000000..71a52ca Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/embedding/embedding.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/embedding/export_weights.py b/crates/burn-store/pytorch-tests/tests/embedding/export_weights.py new file mode 100755 index 0000000..9510e3d --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/embedding/export_weights.py @@ -0,0 +1,37 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.embed = nn.Embedding(10, 3) + + def forward(self, x): + x = self.embed(x) + return x + + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(1) + + model = Model().to(torch.device("cpu")) + + torch.save(model.state_dict(), "embedding.pt") + + input = torch.LongTensor([[1, 2, 4, 5], [4, 3, 2, 9]]) + print("Input shape: {}", input.shape) + print("Input: {}", input) + output = model(input) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + + + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/embedding/mod.rs b/crates/burn-store/pytorch-tests/tests/embedding/mod.rs new file mode 100644 index 0000000..9164b3a --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/embedding/mod.rs @@ -0,0 +1,86 @@ +use burn::{ + module::Module, + nn::{Embedding, EmbeddingConfig}, + tensor::{Device, Int, Tensor}, +}; + +#[derive(Module, Debug)] +pub struct Net { + embed: Embedding, +} + +impl Net { + /// Create a new model. + pub fn init(device: &Device) -> Self { + let embed = EmbeddingConfig::new(10, 3).init(device); + Self { embed } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<2, Int>) -> Tensor<3> { + self.embed.forward(x) + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::Tolerance; + use burn_store::{ModuleSnapshot, PytorchStore}; + + use super::*; + + fn embedding(model: Net, precision: f32) { + let device = Default::default(); + + let input = Tensor::<2, Int>::from_data([[1, 2, 4, 5], [4, 3, 2, 9]], &device); + + let output = model.forward(input); + + let expected = Tensor::<3>::from_data( + [ + [ + [-1.609_484_9, -0.10016718, -0.609_188_9], + [-0.97977227, -1.609_096_3, -0.712_144_6], + [-0.22227049, 1.687_113_4, -0.32062083], + [-0.29934573, 1.879_345_7, -0.07213178], + ], + [ + [-0.22227049, 1.687_113_4, -0.32062083], + [0.303_722, -0.777_314_3, -0.25145486], + [-0.97977227, -1.609_096_3, -0.712_144_6], + [-0.02878714, 2.357_111, -1.037_338_7], + ], + ], + &device, + ); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::absolute(precision)); + } + + #[test] + fn embedding_full_precision() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/embedding/embedding.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + embedding(model, 1e-3); + } + + #[test] + fn embedding_half_precision() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/embedding/embedding.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + embedding(model, 1e-3); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/enum_module/enum_depthwise_false.pt b/crates/burn-store/pytorch-tests/tests/enum_module/enum_depthwise_false.pt new file mode 100644 index 0000000..27e8935 Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/enum_module/enum_depthwise_false.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/enum_module/enum_depthwise_true.pt b/crates/burn-store/pytorch-tests/tests/enum_module/enum_depthwise_true.pt new file mode 100644 index 0000000..e9488a6 Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/enum_module/enum_depthwise_true.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/enum_module/export_weights.py b/crates/burn-store/pytorch-tests/tests/enum_module/export_weights.py new file mode 100755 index 0000000..7a22623 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/enum_module/export_weights.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python3 +import torch +from torch import nn, Tensor + +class DwsConv(nn.Module): + """Depthwise separable convolution.""" + + def __init__(self, in_channels: int, out_channels: int, kernel_size: int) -> None: + super().__init__() + # Depthwise conv + self.dconv = nn.Conv2d(in_channels, in_channels, kernel_size, groups=in_channels) + # Pointwise conv + self.pconv = nn.Conv2d(in_channels, out_channels, kernel_size=1, groups=1) + + def forward(self, x: Tensor) -> Tensor: + x = self.dconv(x) + return self.pconv(x) + + +class Model(nn.Module): + def __init__(self, depthwise: bool = False) -> None: + super().__init__() + self.conv = DwsConv(2, 2, 3) if depthwise else nn.Conv2d(2, 2, 3) + + def forward(self, x: Tensor) -> Tensor: + return self.conv(x) + + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(1) + + model = Model().to(torch.device("cpu")) + + torch.save(model.state_dict(), "enum_depthwise_false.pt") + + input = torch.rand(1, 2, 5, 5) + + print("Depthwise is False") + print("Input shape: {}", input.shape) + print("Input: {}", input) + output = model(input) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + + print("Depthwise is True") + model = Model(depthwise=True).to(torch.device("cpu")) + torch.save(model.state_dict(), "enum_depthwise_true.pt") + + print("Input shape: {}", input.shape) + print("Input: {}", input) + output = model(input) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/enum_module/mod.rs b/crates/burn-store/pytorch-tests/tests/enum_module/mod.rs new file mode 100644 index 0000000..539980c --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/enum_module/mod.rs @@ -0,0 +1,196 @@ +use burn::{ + module::Module, + nn::conv::{Conv2d, Conv2dConfig}, + tensor::{Device, Tensor}, +}; + +#[derive(Module, Debug)] +#[allow(clippy::large_enum_variant)] +pub enum Conv { + DwsConv(DwsConv), + Conv(Conv2d), +} + +#[derive(Module, Debug)] +pub struct DwsConv { + dconv: Conv2d, + pconv: Conv2d, +} + +#[derive(Module, Debug)] +pub struct Net { + conv: Conv, +} + +impl Net { + /// Create a new model with DwsConv variant. + pub fn init_dws_conv(device: &Device) -> Self { + let dconv = Conv2dConfig::new([2, 2], [3, 3]) + .with_groups(2) + .init(device); + let pconv = Conv2dConfig::new([2, 2], [1, 1]) + .with_groups(1) + .init(device); + Net { + conv: Conv::DwsConv(DwsConv { dconv, pconv }), + } + } + + /// Create a new model with Conv variant. + pub fn init_conv(device: &Device) -> Self { + let conv2d_config = Conv2dConfig::new([2, 2], [3, 3]); + Net { + conv: Conv::Conv(conv2d_config.init(device)), + } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + match &self.conv { + Conv::DwsConv(dws_conv) => { + let x = dws_conv.dconv.forward(x); + dws_conv.pconv.forward(x) + } + Conv::Conv(conv) => conv.forward(x), + } + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::Tolerance; + use burn_store::{ModuleSnapshot, PytorchStore}; + type FT = f32; + + use super::*; + + #[test] + fn depthwise_false() { + let device = Default::default(); + let mut model = Net::init_conv(&device); + let mut store = PytorchStore::from_file("tests/enum_module/enum_depthwise_false.pt"); + + model + .load_from(&mut store) + .expect("Should decode state successfully"); + let input = Tensor::<4>::from_data( + [[ + [ + [0.713_979_7, 0.267_644_3, 0.990_609, 0.28845078, 0.874_962_4], + [0.505_920_8, 0.23659128, 0.757_007_4, 0.23458993, 0.64705235], + [0.355_621_4, 0.445_182_8, 0.01930594, 0.26160914, 0.771_317], + [0.37846136, 0.99802476, 0.900_794_2, 0.476_588_2, 0.16625845], + [ + 0.804_481_1, + 0.65517855, + 0.17679012, + 0.824_772_3, + 0.803_550_9, + ], + ], + [ + [0.943_447_5, 0.21972018, 0.417_697, 0.49031407, 0.57302874], + [0.12054086, 0.14518881, 0.772_002_3, 0.38275403, 0.744_236_7], + [0.52850497, 0.664_172_4, 0.60994434, 0.681_799_7, 0.74785537], + [ + 0.03694397, + 0.751_675_7, + 0.148_438_4, + 0.12274551, + 0.530_407_2, + ], + [0.414_796_4, 0.793_662, 0.21043217, 0.05550903, 0.863_884_4], + ], + ]], + &device, + ); + + let output = model.forward(input); + + let expected = Tensor::<4>::from_data( + [[ + [ + [0.35449377, -0.02832414, 0.490_976_1], + [0.29709217, 0.332_586_3, 0.30594018], + [0.18101373, 0.30932188, 0.30558896], + ], + [ + [-0.17683622, -0.13244139, -0.05608707], + [0.23467252, -0.07038684, 0.255_044_1], + [-0.241_931_3, -0.20476191, -0.14468731], + ], + ]], + &device, + ); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::default()); + } + + #[test] + fn depthwise_true() { + let device = Default::default(); + let mut model = Net::init_dws_conv(&device); + let mut store = PytorchStore::from_file("tests/enum_module/enum_depthwise_true.pt"); + + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + let input = Tensor::<4>::from_data( + [[ + [ + [0.713_979_7, 0.267_644_3, 0.990_609, 0.28845078, 0.874_962_4], + [0.505_920_8, 0.23659128, 0.757_007_4, 0.23458993, 0.64705235], + [0.355_621_4, 0.445_182_8, 0.01930594, 0.26160914, 0.771_317], + [0.37846136, 0.99802476, 0.900_794_2, 0.476_588_2, 0.16625845], + [ + 0.804_481_1, + 0.65517855, + 0.17679012, + 0.824_772_3, + 0.803_550_9, + ], + ], + [ + [0.943_447_5, 0.21972018, 0.417_697, 0.49031407, 0.57302874], + [0.12054086, 0.14518881, 0.772_002_3, 0.38275403, 0.744_236_7], + [0.52850497, 0.664_172_4, 0.60994434, 0.681_799_7, 0.74785537], + [ + 0.03694397, + 0.751_675_7, + 0.148_438_4, + 0.12274551, + 0.530_407_2, + ], + [0.414_796_4, 0.793_662, 0.21043217, 0.05550903, 0.863_884_4], + ], + ]], + &device, + ); + + let output = model.forward(input); + + let expected = Tensor::<4>::from_data( + [[ + [ + [0.77874625, 0.859_017_6, 0.834_283_5], + [0.773_056_4, 0.73817325, 0.78292674], + [0.710_775_2, 0.747_187_2, 0.733_264_4], + ], + [ + [-0.44891885, -0.49027523, -0.394_170_7], + [-0.43836114, -0.33961445, -0.387_311_5], + [-0.581_134_3, -0.34197026, -0.535_035_7], + ], + ]], + &device, + ); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::default()); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/group_norm/export_weights.py b/crates/burn-store/pytorch-tests/tests/group_norm/export_weights.py new file mode 100755 index 0000000..9e44950 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/group_norm/export_weights.py @@ -0,0 +1,37 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.norm1 = nn.GroupNorm(2, 6) + + def forward(self, x): + x = self.norm1(x) + return x + + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(1) + + model = Model().to(torch.device("cpu")) + + torch.save(model.state_dict(), "group_norm.pt") + + x2 = torch.rand(1, 6, 2, 2) + print("Input shape: {}", x2.shape) + print("Input: {}", x2) + output = model(x2) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + + + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/group_norm/group_norm.pt b/crates/burn-store/pytorch-tests/tests/group_norm/group_norm.pt new file mode 100644 index 0000000..aa8b439 Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/group_norm/group_norm.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/group_norm/mod.rs b/crates/burn-store/pytorch-tests/tests/group_norm/mod.rs new file mode 100644 index 0000000..a02cd42 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/group_norm/mod.rs @@ -0,0 +1,90 @@ +use burn::{ + module::Module, + nn::{GroupNorm, GroupNormConfig}, + tensor::{Device, Tensor}, +}; + +#[derive(Module, Debug)] +pub struct Net { + norm1: GroupNorm, +} + +impl Net { + /// Create a new model from the given record. + pub fn init(device: &Device) -> Self { + let norm1 = GroupNormConfig::new(2, 6).init(device); + Self { norm1 } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + self.norm1.forward(x) + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::Tolerance; + use burn_store::{ModuleSnapshot, PytorchStore}; + + use super::*; + + fn group_norm(model: Net, precision: f32) { + let device = Default::default(); + + let input = Tensor::<4>::from_data( + [[ + [[0.757_631_6, 0.27931088], [0.40306926, 0.73468447]], + [[0.02928156, 0.799_858_6], [0.39713734, 0.75437194]], + [[0.569_508_5, 0.43877792], [0.63868046, 0.524_665_9]], + [[0.682_614_1, 0.305_149_5], [0.46354562, 0.45498633]], + [[0.572_472, 0.498_002_6], [0.93708336, 0.65559506]], + [[0.31379688, 0.19801933], [0.41619217, 0.28432965]], + ]], + &device, + ); + + let output = model.forward(input); + + let expected = Tensor::<4>::from_data( + [[ + [[1.042_578_5, -1.122_016_7], [-0.56195974, 0.938_733_6]], + [[-2.253_500_7, 1.233_672_9], [-0.588_804_1, 1.027_827_3]], + [[0.19124532, -0.40036356], [0.504_276_5, -0.01168585]], + [[1.013_829_2, -0.891_984_6], [-0.09224463, -0.13546038]], + [[0.45772314, 0.08172822], [2.298_641_4, 0.877_410_4]], + [[-0.84832406, -1.432_883_4], [-0.331_331_5, -0.997_103_7]], + ]], + &device, + ); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::absolute(precision)); + } + + #[test] + fn group_norm_full() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/group_norm/group_norm.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + group_norm(model, 1e-3); + } + + #[test] + fn group_norm_half() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/group_norm/group_norm.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + group_norm(model, 1e-3); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/integer/export_weights.py b/crates/burn-store/pytorch-tests/tests/integer/export_weights.py new file mode 100755 index 0000000..151dda3 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/integer/export_weights.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + buffer = torch.tensor([1, 2, 3]) + self.register_buffer("buffer", buffer, persistent=True) + + def forward(self, x): + x = self.buffer + return x + + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(1) + + model = Model().to(torch.device("cpu")) + + torch.save(model.state_dict(), "integer.pt") + + input = torch.ones(3, 3) + print("Input shape: {}", input.shape) + print("Input: {}", input) + output = model(input) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + + + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/integer/integer.pt b/crates/burn-store/pytorch-tests/tests/integer/integer.pt new file mode 100644 index 0000000..bcf22cd Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/integer/integer.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/integer/mod.rs b/crates/burn-store/pytorch-tests/tests/integer/mod.rs new file mode 100644 index 0000000..8a6b470 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/integer/mod.rs @@ -0,0 +1,76 @@ +use burn::{ + module::{Module, Param, ParamId}, + tensor::{Device, Int, Tensor, TensorData}, +}; + +#[derive(Module, Debug)] +pub struct Net { + buffer: Param>, +} + +impl Net { + /// Create a new model with placeholder values. + pub fn init(device: &Device) -> Self { + Self { + buffer: Param::initialized( + ParamId::new(), + Tensor::<1, Int>::from_data(TensorData::from([0, 0, 0]), device), + ), + } + } + + /// Forward pass of the model. + pub fn forward(&self, _x: Tensor<2>) -> Tensor<1, Int> { + self.buffer.val() + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::DType; + use burn_store::{ModuleSnapshot, PytorchStore}; + + use super::*; + + fn integer(model: Net) { + let device = Default::default(); + + let input = Tensor::<2>::ones([3, 3], &device); + + let output = model.forward(input); + let data = output.to_data(); + + // The .pt file stores int64 (PyTorch's default int dtype); we pin + // that here to catch a regression where the loader silently casts + // to the backend's native IntElem (i32 for Flex). + assert_eq!(data.dtype, DType::I64); + + let values = data.iter::().collect::>(); + assert_eq!(values, vec![1i64, 2, 3]); + } + + #[test] + fn integer_full_precision() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/integer/integer.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + integer(model); + } + + #[test] + fn integer_half_precision() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/integer/integer.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + integer(model); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/key_remap/export_weights.py b/crates/burn-store/pytorch-tests/tests/key_remap/export_weights.py new file mode 100755 index 0000000..1b7490b --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/key_remap/export_weights.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class ConvModule(nn.Module): + def __init__(self): + super(ConvModule, self).__init__() + self.conv1 = nn.Conv2d(2, 2, (2,2)) + self.conv2 = nn.Conv2d(2, 2, (2,2), bias=False) + + def forward(self, x): + x = self.conv1(x) + x = self.conv2(x) + return x + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.conv = ConvModule() + + def forward(self, x): + x = self.conv(x) + return x + + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(1) + + model = Model().to(torch.device("cpu")) + + torch.save(model.state_dict(), "key_remap.pt") + + input = torch.rand(1, 2, 5, 5) + print("Input shape: {}", input.shape) + print("Input: {}", input) + output = model(input) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + + + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/key_remap/key_remap.pt b/crates/burn-store/pytorch-tests/tests/key_remap/key_remap.pt new file mode 100644 index 0000000..088a943 Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/key_remap/key_remap.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/key_remap/mod.rs b/crates/burn-store/pytorch-tests/tests/key_remap/mod.rs new file mode 100644 index 0000000..bb8c422 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/key_remap/mod.rs @@ -0,0 +1,117 @@ +use burn::{ + module::Module, + nn::conv::{Conv2d, Conv2dConfig}, + tensor::{Device, Tensor}, +}; + +#[derive(Module, Debug)] +pub struct Net { + conv1: Conv2d, + conv2: Conv2d, +} + +impl Net { + /// Create a new model. + pub fn init(device: &Device) -> Self { + let conv1 = Conv2dConfig::new([2, 2], [2, 2]).init(device); + let conv2 = Conv2dConfig::new([2, 2], [2, 2]) + .with_bias(false) + .init(device); + Self { conv1, conv2 } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + let x = self.conv1.forward(x); + + self.conv2.forward(x) + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::Tolerance; + use burn_store::{ModuleSnapshot, PytorchStore}; + type FT = f32; + + use super::*; + + #[test] + fn key_remap() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/key_remap/key_remap.pt") + .with_key_remapping("conv\\.(.*)", "$1"); // Remove "conv" prefix, e.g. "conv.conv1" -> "conv1" + + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + let input = Tensor::<4>::from_data( + [[ + [ + [ + 0.024_595_8, + 0.25883394, + 0.93905586, + 0.416_715_5, + 0.713_979_7, + ], + [0.267_644_3, 0.990_609, 0.28845078, 0.874_962_4, 0.505_920_8], + [0.23659128, 0.757_007_4, 0.23458993, 0.64705235, 0.355_621_4], + [0.445_182_8, 0.01930594, 0.26160914, 0.771_317, 0.37846136], + [ + 0.99802476, + 0.900_794_2, + 0.476_588_2, + 0.16625845, + 0.804_481_1, + ], + ], + [ + [ + 0.65517855, + 0.17679012, + 0.824_772_3, + 0.803_550_9, + 0.943_447_5, + ], + [0.21972018, 0.417_697, 0.49031407, 0.57302874, 0.12054086], + [0.14518881, 0.772_002_3, 0.38275403, 0.744_236_7, 0.52850497], + [0.664_172_4, 0.60994434, 0.681_799_7, 0.74785537, 0.03694397], + [ + 0.751_675_7, + 0.148_438_4, + 0.12274551, + 0.530_407_2, + 0.414_796_4, + ], + ], + ]], + &device, + ); + + let output = model.forward(input); + + let expected = Tensor::<4>::from_data( + [[ + [ + [-0.02502128, 0.00250649, 0.04841233], + [0.04589614, -0.00296854, 0.01991477], + [0.02920526, 0.059_497_3, 0.04326791], + ], + [ + [-0.04825336, 0.080_190_9, -0.02375088], + [0.02885434, 0.09638263, -0.07460806], + [0.02004079, 0.06244051, 0.035_887_1], + ], + ]], + &device, + ); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::default()); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/key_remap_chained/export_weights.py b/crates/burn-store/pytorch-tests/tests/key_remap_chained/export_weights.py new file mode 100755 index 0000000..95da5a4 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/key_remap_chained/export_weights.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python3 + +import torch +from torch import nn, Tensor + + +class ConvBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.block = nn.Sequential( + nn.Conv2d(in_channels, out_channels, 1, bias=False), + nn.BatchNorm2d(out_channels), + ) + + def forward(self, x: Tensor) -> Tensor: + return self.block(x) + + +class Model(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(3, 6, 3, bias=False) + self.bn = nn.BatchNorm2d(6) + self.layer = nn.Sequential(ConvBlock(6, 6), ConvBlock(6, 6)) + + def forward(self, x: Tensor) -> Tensor: + x = self.conv(x) + x = self.bn(x) + x = self.layer(x) + + return x + + +def main(): + torch.set_printoptions(precision=8) + torch.manual_seed(42) + + model = Model() + + input = torch.rand(1, 3, 4, 4) + model(input) # condition batch norm + model.eval() + + with torch.no_grad(): + print(f"Input shape: {input.shape}") + print("Input type: {}", input.dtype) + print(f"Input: {input}") + output = model(input) + + print(f"Output: {output}") + print(f"Output Shape: {output.shape}") + + torch.save(model.state_dict(), "key_remap.pt") + + +if __name__ == "__main__": + main() diff --git a/crates/burn-store/pytorch-tests/tests/key_remap_chained/key_remap.pt b/crates/burn-store/pytorch-tests/tests/key_remap_chained/key_remap.pt new file mode 100644 index 0000000..8c2b11b Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/key_remap_chained/key_remap.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/key_remap_chained/mod.rs b/crates/burn-store/pytorch-tests/tests/key_remap_chained/mod.rs new file mode 100644 index 0000000..31dc80e --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/key_remap_chained/mod.rs @@ -0,0 +1,172 @@ +use burn::{ + module::Module, + nn::{ + BatchNorm, BatchNormConfig, + conv::{Conv2d, Conv2dConfig}, + }, + tensor::{Device, Tensor}, +}; + +/// Some module that implements a specific method so it can be used in a sequential block. +pub trait ForwardModule { + fn forward(&self, input: Tensor<4>) -> Tensor<4>; +} + +/// Conv2d + BatchNorm block. +#[derive(Module, Debug)] +pub struct ConvBlock { + conv: Conv2d, + bn: BatchNorm, +} + +impl ForwardModule for ConvBlock { + fn forward(&self, input: Tensor<4>) -> Tensor<4> { + let out = self.conv.forward(input); + self.bn.forward(out) + } +} + +impl ConvBlock { + pub fn new(in_channels: usize, out_channels: usize, device: &Device) -> Self { + let conv = Conv2dConfig::new([in_channels, out_channels], [1, 1]) + .with_bias(false) + .init(device); + let bn = BatchNormConfig::new(out_channels).init(device); + + Self { conv, bn } + } +} + +/// Collection of sequential blocks. +#[derive(Module, Debug)] +pub struct ModuleBlock { + blocks: Vec, +} + +impl ModuleBlock { + pub fn forward(&self, input: Tensor<4>) -> Tensor<4> { + let mut out = input; + for block in &self.blocks { + out = block.forward(out); + } + out + } +} + +impl ModuleBlock { + pub fn new(device: &Device) -> Self { + let blocks = vec![ConvBlock::new(6, 6, device), ConvBlock::new(6, 6, device)]; + + Self { blocks } + } +} + +#[derive(Module, Debug)] +pub struct Model { + conv: Conv2d, + bn: BatchNorm, + layer: ModuleBlock, +} + +impl Model { + pub fn new(device: &Device) -> Self { + let conv = Conv2dConfig::new([3, 6], [3, 3]) + .with_bias(false) + .init(device); + let bn = BatchNormConfig::new(6).init(device); + + let layer = ModuleBlock::new(device); + + Self { conv, bn, layer } + } + + pub fn forward(&self, input: Tensor<4>) -> Tensor<4> { + let out = self.conv.forward(input); + let out = self.bn.forward(out); + self.layer.forward(out) + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::Tolerance; + use burn_store::{ModuleSnapshot, PytorchStore}; + type FT = f32; + + use super::*; + + #[test] + #[should_panic] + fn key_remap_chained_missing_pattern() { + // Loading record should fail due to missing pattern to map the layer.blocks + let device = Default::default(); + let mut model: Model<_> = Model::new(&device); + let mut store = PytorchStore::from_file("tests/key_remap_chained/key_remap.pt") + // Map *.block.0.* -> *.conv.* + .with_key_remapping("(.+)\\.block\\.0\\.(.+)", "$1.conv.$2") + // Map *.block.1.* -> *.bn.* + .with_key_remapping("(.+)\\.block\\.1\\.(.+)", "$1.bn.$2"); + + model + .load_from(&mut store) + .expect("Should decode state successfully"); + } + + #[test] + fn key_remap_chained() { + let device = Default::default(); + let mut model: Model<_> = Model::new(&device); + let mut store = PytorchStore::from_file("tests/key_remap_chained/key_remap.pt") + // Map *.block.0.* -> *.conv.* + .with_key_remapping("(.+)\\.block\\.0\\.(.+)", "$1.conv.$2") + // Map *.block.1.* -> *.bn.* + .with_key_remapping("(.+)\\.block\\.1\\.(.+)", "$1.bn.$2") + // Map layer.[i].* -> layer.blocks.[i].* + .with_key_remapping("layer\\.([0-9])\\.(.+)", "layer.blocks.$1.$2"); + + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + let input = Tensor::<4>::from_data( + [[ + [ + [0.76193494, 0.626_546_1, 0.49510366, 0.11974698], + [0.07161391, 0.03232569, 0.704_681, 0.254_516], + [0.399_373_7, 0.21224737, 0.40888822, 0.14808255], + [0.17329216, 0.665_855_4, 0.351_401_8, 0.808_671_6], + ], + [ + [0.33959562, 0.13321638, 0.41178054, 0.257_626_3], + [0.347_029_2, 0.02400219, 0.77974546, 0.15189773], + [0.75130886, 0.726_892_1, 0.85721636, 0.11647397], + [0.859_598_4, 0.263_624_2, 0.685_534_6, 0.96955734], + ], + [ + [0.42948407, 0.49613327, 0.38488472, 0.08250773], + [0.73995143, 0.00364107, 0.81039995, 0.87411255], + [0.972_853_2, 0.38206023, 0.08917904, 0.61241513], + [0.77621365, 0.00234562, 0.38650817, 0.20027226], + ], + ]], + &device, + ); + let expected = Tensor::<4>::from_data( + [[ + [[0.198_967_1, 0.17847246], [0.06883702, 0.20012866]], + [[0.17582723, 0.11344293], [0.05444185, 0.13307181]], + [[0.192_229_5, 0.20391327], [0.06150475, 0.22688155]], + [[0.00230906, -0.02177845], [0.01129148, 0.00925517]], + [[0.14751078, 0.14433631], [0.05498439, 0.29049855]], + [[0.16868964, 0.133_269_3], [0.06917118, 0.35094324]], + ]], + &device, + ); + + let output = model.forward(input); + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::default()); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/layer_norm/export_weights.py b/crates/burn-store/pytorch-tests/tests/layer_norm/export_weights.py new file mode 100755 index 0000000..04b4234 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/layer_norm/export_weights.py @@ -0,0 +1,37 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.norm1 = nn.LayerNorm(2) + + def forward(self, x): + x = self.norm1(x) + return x + + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(1) + + model = Model().to(torch.device("cpu")) + + torch.save(model.state_dict(), "layer_norm.pt") + + x2 = torch.rand(1, 2, 2, 2) + print("Input shape: {}", x2.shape) + print("Input: {}", x2) + output = model(x2) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + + + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/layer_norm/layer_norm.pt b/crates/burn-store/pytorch-tests/tests/layer_norm/layer_norm.pt new file mode 100644 index 0000000..f792f58 Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/layer_norm/layer_norm.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/layer_norm/mod.rs b/crates/burn-store/pytorch-tests/tests/layer_norm/mod.rs new file mode 100644 index 0000000..a0def47 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/layer_norm/mod.rs @@ -0,0 +1,81 @@ +use burn::{ + module::Module, + nn::{LayerNorm, LayerNormConfig}, + tensor::{Device, Tensor}, +}; + +#[derive(Module, Debug)] +pub struct Net { + norm1: LayerNorm, +} + +impl Net { + /// Create a new model. + pub fn init(device: &Device) -> Self { + let norm1 = LayerNormConfig::new(2).init(device); + Self { norm1 } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + self.norm1.forward(x) + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::Tolerance; + use burn_store::{ModuleSnapshot, PytorchStore}; + type FT = f32; + + use super::*; + + fn layer_norm(model: Net, precision: f32) { + let device = Default::default(); + + let input = Tensor::<4>::from_data( + [[ + [[0.757_631_6, 0.27931088], [0.40306926, 0.73468447]], + [[0.02928156, 0.799_858_6], [0.39713734, 0.75437194]], + ]], + &device, + ); + + let output = model.forward(input); + + let expected = Tensor::<4>::from_data( + [[ + [[0.99991274, -0.999_912_5], [-0.999_818_3, 0.999_818_3]], + [[-0.999_966_2, 0.99996626], [-0.99984336, 0.99984336]], + ]], + &device, + ); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::absolute(precision)); + } + + #[test] + fn layer_norm_full() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/layer_norm/layer_norm.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + layer_norm(model, 1e-3); + } + + #[test] + fn layer_norm_half() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/layer_norm/layer_norm.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + layer_norm(model, 1e-3); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/linear/export_weights.py b/crates/burn-store/pytorch-tests/tests/linear/export_weights.py new file mode 100755 index 0000000..5d76638 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/linear/export_weights.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.fc1 = nn.Linear(2, 3) + self.fc2 = nn.Linear(3, 4, bias=False) + + def forward(self, x): + x = self.fc1(x) + x = F.relu(x) # Add relu so that PyTorch optimizer does not combine fc1 and fc2 + x = self.fc2(x) + + return x + + +class ModelWithBias(nn.Module): + def __init__(self): + super(ModelWithBias, self).__init__() + self.fc1 = nn.Linear(2, 3) + + def forward(self, x): + x = self.fc1(x) + + return x + + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(1) + + model = Model().to(torch.device("cpu")) + model_with_bias = ModelWithBias().to(torch.device("cpu")) + + torch.save(model.state_dict(), "linear.pt") + torch.save(model_with_bias.state_dict(), "linear_with_bias.pt") + + input = torch.rand(1, 2, 2, 2) + print("Input shape: {}", input.shape) + print("Input: {}", input) + + output = model(input) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + print("Model with bias") + output = model_with_bias(input) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + + + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/linear/linear.pt b/crates/burn-store/pytorch-tests/tests/linear/linear.pt new file mode 100644 index 0000000..fd5f4c9 Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/linear/linear.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/linear/linear_with_bias.pt b/crates/burn-store/pytorch-tests/tests/linear/linear_with_bias.pt new file mode 100644 index 0000000..dc4c68f Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/linear/linear_with_bias.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/linear/mod.rs b/crates/burn-store/pytorch-tests/tests/linear/mod.rs new file mode 100644 index 0000000..8546ad7 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/linear/mod.rs @@ -0,0 +1,153 @@ +use burn::{ + module::Module, + nn::{Linear, LinearConfig, Relu}, + tensor::{Device, Tensor}, +}; + +#[derive(Module, Debug)] +pub struct Net { + fc1: Linear, + fc2: Linear, + relu: Relu, +} + +impl Net { + /// Create a new model. + pub fn init(device: &Device) -> Self { + let fc1 = LinearConfig::new(2, 3).init(device); + let fc2 = LinearConfig::new(3, 4).with_bias(false).init(device); + let relu = Relu; + + Self { fc1, fc2, relu } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + let x = self.fc1.forward(x); + let x = self.relu.forward(x); + + self.fc2.forward(x) + } +} + +#[derive(Module, Debug)] +struct NetWithBias { + fc1: Linear, +} + +impl NetWithBias { + /// Create a new model. + pub fn init(device: &Device) -> Self { + let fc1 = LinearConfig::new(2, 3).init(device); + + Self { fc1 } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + self.fc1.forward(x) + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::Tolerance; + use burn_store::{ModuleSnapshot, PytorchStore}; + type FT = f32; + + use super::*; + + fn linear_test(model: Net, precision: f32) { + let device = Default::default(); + + let input = Tensor::<4>::from_data( + [[ + [[0.63968194, 0.97427773], [0.830_029_9, 0.04443115]], + [[0.024_595_8, 0.25883394], [0.93905586, 0.416_715_5]], + ]], + &device, + ); + + let output = model.forward(input); + let expected = Tensor::<4>::from_data( + [[ + [ + [0.09778349, -0.13756673, 0.04962806, 0.08856435], + [0.03163241, -0.02848549, 0.01437942, 0.11905234], + ], + [ + [0.07628226, -0.10757702, 0.03656857, 0.03824598], + [0.05443089, -0.06904714, 0.02744314, 0.09997337], + ], + ]], + &device, + ); + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::absolute(precision)); + } + + #[test] + fn linear_full_precision() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/linear/linear.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + linear_test(model, 1e-7); + } + + #[test] + fn linear_half_precision() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/linear/linear.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + linear_test(model, 1e-4); + } + + #[test] + fn linear_with_bias() { + let device = Default::default(); + + let mut model = NetWithBias::init(&device); + let mut store = PytorchStore::from_file("tests/linear/linear_with_bias.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + let input = Tensor::<4>::from_data( + [[ + [[0.63968194, 0.97427773], [0.830_029_9, 0.04443115]], + [[0.024_595_8, 0.25883394], [0.93905586, 0.416_715_5]], + ]], + &device, + ); + + let output = model.forward(input); + + let expected = Tensor::<4>::from_data( + [[ + [ + [-0.00432095, -1.107_101_2, 0.870_691_4], + [0.024_595_5, -0.954_462_9, 0.48518157], + ], + [ + [0.34315687, -0.757_384_2, 0.548_288], + [-0.06608963, -1.072_072_7, 0.645_800_5], + ], + ]], + &device, + ); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::default()); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/missing_module_field/export_weights.py b/crates/burn-store/pytorch-tests/tests/missing_module_field/export_weights.py new file mode 100755 index 0000000..e65585c --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/missing_module_field/export_weights.py @@ -0,0 +1,24 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.conv1 = nn.Conv2d(2, 2, (2,2)) + + def forward(self, x): + x = self.conv1(x) + return x + + +def main(): + torch.set_printoptions(precision=8) + torch.manual_seed(1) + model = Model().to(torch.device("cpu")) + torch.save(model.state_dict(), "missing_module_field.pt") + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/missing_module_field/missing_module_field.pt b/crates/burn-store/pytorch-tests/tests/missing_module_field/missing_module_field.pt new file mode 100644 index 0000000..8162e99 Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/missing_module_field/missing_module_field.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/missing_module_field/mod.rs b/crates/burn-store/pytorch-tests/tests/missing_module_field/mod.rs new file mode 100644 index 0000000..a213b6c --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/missing_module_field/mod.rs @@ -0,0 +1,36 @@ +use burn::{module::Module, nn::conv::Conv2d, tensor::Device}; + +#[derive(Module, Debug)] +#[allow(unused)] +pub struct Net { + do_not_exist_in_pt: Conv2d, +} + +#[cfg(test)] +mod tests { + + use burn::nn::conv::Conv2dConfig; + use burn_store::{ModuleSnapshot, PytorchStore}; + + use super::*; + + impl Net { + pub fn init(device: &Device) -> Self { + Self { + do_not_exist_in_pt: Conv2dConfig::new([2, 2], [2, 2]).init(device), + } + } + } + + #[test] + #[should_panic(expected = "do_not_exist_in_pt")] + fn should_fail_if_struct_field_is_missing() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = + PytorchStore::from_file("tests/missing_module_field/missing_module_field.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/non_contiguous_indexes/export_weights.py b/crates/burn-store/pytorch-tests/tests/non_contiguous_indexes/export_weights.py new file mode 100755 index 0000000..9a76930 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/non_contiguous_indexes/export_weights.py @@ -0,0 +1,42 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + num_layers = 5 # Number of repeated convolutional layers + + # Create a list to store the layers + layers = [] + for _ in range(num_layers): + layers.append(nn.Conv2d(2, 2, kernel_size=3, padding=1, bias=True)) + layers.append(nn.ReLU(inplace=True)) + + # Use nn.Sequential to create a single module from the layers + self.fc = nn.Sequential(*layers) + + def forward(self, x): + x = self.fc(x) + return x + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(1) + + model = Model().to(torch.device("cpu")) + + torch.save(model.state_dict(), "non_contiguous_indexes.pt") + + input = torch.rand(1, 2, 5, 5) + print("Input shape: {}", input.shape) + print("Input: {}", input) + output = model(input) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/non_contiguous_indexes/mod.rs b/crates/burn-store/pytorch-tests/tests/non_contiguous_indexes/mod.rs new file mode 100644 index 0000000..5b2753e --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/non_contiguous_indexes/mod.rs @@ -0,0 +1,109 @@ +use burn::{ + module::Module, + nn::{ + PaddingConfig2d, + conv::{Conv2d, Conv2dConfig}, + }, + tensor::{Device, Tensor, activation::relu}, +}; + +#[derive(Module, Debug)] +pub struct Net { + fc: Vec, +} + +impl Net { + /// Create a new model with placeholder values. + pub fn init(device: &Device) -> Self { + let conv2d_config = Conv2dConfig::new([2, 2], [3, 3]).with_padding(PaddingConfig2d::Same); + // The PyTorch file has 5 Conv2d layers at non-contiguous indices (0, 2, 4, 6, 8) + // in the Sequential (alternating with ReLU layers) + let fc = vec![ + conv2d_config.init(device), + conv2d_config.init(device), + conv2d_config.init(device), + conv2d_config.init(device), + conv2d_config.init(device), + ]; + Net { fc } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + self.fc.iter().fold(x, |x_i, conv| relu(conv.forward(x_i))) + } +} + +#[cfg(test)] +mod tests { + + use burn::tensor::Tolerance; + use burn_store::{ModuleSnapshot, PytorchStore}; + type FT = f32; + + use super::*; + + #[test] + fn non_contiguous_indexes() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = + PytorchStore::from_file("tests/non_contiguous_indexes/non_contiguous_indexes.pt"); + + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + let input = Tensor::<4>::from_data( + [[ + [ + [ + 0.67890584, + 0.307_537_2, + 0.265_156_2, + 0.528_318_8, + 0.86194897, + ], + [0.14828813, 0.73480314, 0.821_220_7, 0.989_098_6, 0.15003455], + [0.62109494, 0.13028657, 0.926_875_1, 0.30604684, 0.80117637], + [0.514_885_7, 0.46105868, 0.484_046_1, 0.58499724, 0.73569804], + [0.58018994, 0.65252745, 0.05023766, 0.864_268_7, 0.935_932], + ], + [ + [0.913_302_9, 0.869_611_3, 0.139_184_3, 0.314_65, 0.94086266], + [0.11917073, 0.953_610_6, 0.10675198, 0.14779574, 0.744_439], + [0.14075547, 0.38544965, 0.863_745_9, 0.89604443, 0.97287786], + [0.39854127, 0.11136961, 0.99230546, 0.39348692, 0.29428244], + [0.621_886_9, 0.15033776, 0.828_640_1, 0.81336635, 0.10325938], + ], + ]], + &device, + ); + + let output = model.forward(input); + + let expected = Tensor::<4>::from_data( + [[ + [ + [0.00000000, 0.00000000, 0.00000000, 0.00000000, 0.00000000], + [0.00000000, 0.00000000, 0.00000000, 0.00000000, 0.00000000], + [0.00000000, 0.00000000, 0.00000000, 0.00000000, 0.00000000], + [0.00000000, 0.00000000, 0.00000000, 0.00000000, 0.00000000], + [0.04485746, 0.03582812, 0.03432692, 0.02892298, 0.013_844_3], + ], + [ + [0.00000000, 0.00000000, 0.00000000, 0.00000000, 0.00000000], + [0.00000000, 0.00000000, 0.00000000, 0.00000000, 0.00000000], + [0.00000000, 0.00000000, 0.00000000, 0.00000000, 0.00000000], + [0.00000000, 0.00000000, 0.00000000, 0.00000000, 0.00000000], + [0.00000000, 0.00000000, 0.00000000, 0.00000000, 0.00000000], + ], + ]], + &device, + ); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::absolute(1e-7)); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/non_contiguous_indexes/non_contiguous_indexes.pt b/crates/burn-store/pytorch-tests/tests/non_contiguous_indexes/non_contiguous_indexes.pt new file mode 100644 index 0000000..ff4b515 Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/non_contiguous_indexes/non_contiguous_indexes.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/test_int.pt b/crates/burn-store/pytorch-tests/tests/test_int.pt new file mode 100644 index 0000000..dfb11c4 Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/test_int.pt differ diff --git a/crates/burn-store/pytorch-tests/tests/test_mod.rs b/crates/burn-store/pytorch-tests/tests/test_mod.rs new file mode 100644 index 0000000..b414a55 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/test_mod.rs @@ -0,0 +1,20 @@ +mod batch_norm; +mod boolean; +mod buffer; +mod complex_nested; +mod config; +mod conv1d; +mod conv2d; +mod conv_transpose1d; +mod conv_transpose2d; +mod embedding; +mod enum_module; +mod group_norm; +mod integer; +mod key_remap; +mod key_remap_chained; +mod layer_norm; +mod linear; +mod missing_module_field; +mod non_contiguous_indexes; +mod top_level_key; diff --git a/crates/burn-store/pytorch-tests/tests/top_level_key/export_weights.py b/crates/burn-store/pytorch-tests/tests/top_level_key/export_weights.py new file mode 100755 index 0000000..4045629 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/top_level_key/export_weights.py @@ -0,0 +1,24 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.conv1 = nn.Conv2d(2, 2, (2,2)) + + def forward(self, x): + x = self.conv1(x) + return x + + +def main(): + torch.set_printoptions(precision=8) + torch.manual_seed(1) + model = Model().to(torch.device("cpu")) + torch.save({"my_state_dict": model.state_dict()}, "top_level_key.pt") + +if __name__ == '__main__': + main() diff --git a/crates/burn-store/pytorch-tests/tests/top_level_key/mod.rs b/crates/burn-store/pytorch-tests/tests/top_level_key/mod.rs new file mode 100644 index 0000000..0b8dd10 --- /dev/null +++ b/crates/burn-store/pytorch-tests/tests/top_level_key/mod.rs @@ -0,0 +1,47 @@ +use burn::{module::Module, nn::conv::Conv2d, tensor::Device}; + +#[derive(Module, Debug)] +#[allow(unused)] +pub struct Net { + conv1: Conv2d, +} + +#[cfg(test)] +mod tests { + + use burn::nn::conv::Conv2dConfig; + use burn_store::{ModuleSnapshot, PytorchStore}; + + use super::*; + + impl Net { + pub fn init(device: &Device) -> Self { + Self { + conv1: Conv2dConfig::new([2, 2], [2, 2]).init(device), + } + } + } + + #[test] + #[should_panic] + fn should_fail_if_not_found() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/top_level_key/top_level_key.pt"); + model + .load_from(&mut store) + .expect("Should decode state successfully"); + } + + #[test] + fn should_load() { + let device = Default::default(); + let mut model = Net::init(&device); + let mut store = PytorchStore::from_file("tests/top_level_key/top_level_key.pt") + .with_top_level_key("my_state_dict"); + + model + .load_from(&mut store) + .expect("Should decode state successfully"); + } +} diff --git a/crates/burn-store/pytorch-tests/tests/top_level_key/top_level_key.pt b/crates/burn-store/pytorch-tests/tests/top_level_key/top_level_key.pt new file mode 100644 index 0000000..2658e07 Binary files /dev/null and b/crates/burn-store/pytorch-tests/tests/top_level_key/top_level_key.pt differ diff --git a/crates/burn-store/safetensors-tests/Cargo.toml b/crates/burn-store/safetensors-tests/Cargo.toml new file mode 100644 index 0000000..86abed7 --- /dev/null +++ b/crates/burn-store/safetensors-tests/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "safetensors-tests" +version.workspace = true +edition.workspace = true +license.workspace = true + +[dev-dependencies] +burn = { workspace = true, features = ["default"] } +burn-flex = { workspace = true, features = ["default"] } +burn-autodiff = { workspace = true, features = ["default"] } +burn-store = { workspace = true, features = ["default", "std", "safetensors"] } +serde = { workspace = true } +float-cmp = { workspace = true } diff --git a/crates/burn-store/safetensors-tests/src/lib.rs b/crates/burn-store/safetensors-tests/src/lib.rs new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/crates/burn-store/safetensors-tests/src/lib.rs @@ -0,0 +1 @@ + diff --git a/crates/burn-store/safetensors-tests/tests/multi_layer/mod.rs b/crates/burn-store/safetensors-tests/tests/multi_layer/mod.rs new file mode 100644 index 0000000..45d38d8 --- /dev/null +++ b/crates/burn-store/safetensors-tests/tests/multi_layer/mod.rs @@ -0,0 +1,90 @@ +use burn::{ + module::Module, + nn::{ + BatchNorm, BatchNormConfig, Linear, LinearConfig, PaddingConfig2d, Relu, + conv::{Conv2d, Conv2dConfig}, + }, + tensor::{Device, Tensor}, +}; + +#[derive(Module, Debug)] +pub struct Net { + conv1: Conv2d, + norm1: BatchNorm, + fc1: Linear, + relu: Relu, +} + +impl Net { + pub fn new(device: &Device) -> Self { + Self { + conv1: Conv2dConfig::new([3, 4], [3, 3]) + .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1)) + .init(device), + norm1: BatchNormConfig::new(4).init(device), + fc1: LinearConfig::new(4 * 8 * 8, 16).init(device), + relu: Relu::new(), + } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<4>) -> Tensor<2> { + let x = self.conv1.forward(x); + let x = self.norm1.forward(x); + let x = self.relu.forward(x); + // Flatten all dimensions except the batch dimension + let x = x.flatten(1, 3); + self.fc1.forward(x) + } +} + +#[cfg(test)] +mod tests { + use burn::tensor::Tolerance; + use burn_store::{ModuleSnapshot, PyTorchToBurnAdapter, SafetensorsStore}; + + use super::*; + + #[test] + fn multi_layer_model() { + let device = Default::default(); + let mut model = Net::new(&device); + let mut store = SafetensorsStore::from_file("tests/multi_layer/multi_layer.safetensors") + .with_from_adapter(PyTorchToBurnAdapter); + + model + .load_from(&mut store) + .expect("Should decode state successfully"); + + let input = Tensor::<4>::ones([1, 3, 8, 8], &device); + + let output = model.forward(input); + + // Note: Expected values should be updated based on the actual output from the PyTorch model + let expected = Tensor::<2>::from_data( + [[ + 0.04971555, + -0.16849735, + 0.05182848, + -0.18032673, + 0.23138367, + 0.05041867, + 0.13005908, + -0.32202929, + -0.07915690, + -0.03232457, + -0.19790289, + -0.17476529, + -0.19627589, + -0.21757686, + -0.31376451, + 0.08377837, + ]], + &device, + ); + + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::default()); + } +} diff --git a/crates/burn-store/safetensors-tests/tests/multi_layer/multi_layer.py b/crates/burn-store/safetensors-tests/tests/multi_layer/multi_layer.py new file mode 100755 index 0000000..487462d --- /dev/null +++ b/crates/burn-store/safetensors-tests/tests/multi_layer/multi_layer.py @@ -0,0 +1,49 @@ +#!/usr/bin/env python3 + +import torch +import torch.nn as nn +import torch.nn.functional as F +from safetensors.torch import save_file + + +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + self.conv1 = nn.Conv2d(3, 4, kernel_size=3, padding=1) + self.norm1 = nn.BatchNorm2d(4) + self.flatten = nn.Flatten() + self.fc1 = nn.Linear(4 * 8 * 8, 16) # Changed for smaller input size + + def forward(self, x): + x = self.conv1(x) + x = self.norm1(x) + x = F.relu(x) + x = self.flatten(x) + x = self.fc1(x) + return x + + +def main(): + + torch.set_printoptions(precision=8) + torch.manual_seed(1) + + model = Model().to(torch.device("cpu")) + + # Use a smaller input size + # 1 batch, 3 channels (RGB), 8x8 image (small input) + x1 = torch.ones(1, 3, 8, 8) + _ = model(x1) + model.eval() # Set to eval mode to freeze running stats + # Save the model to safetensors after the first forward + save_file(model.state_dict(), "multi_layer.safetensors") + + x2 = torch.ones(1, 3, 8, 8) + print("Input shape: {}", x2.shape) + output = model(x2) + print("Output: {}", output) + print("Output Shape: {}", output.shape) + + +if __name__ == "__main__": + main() diff --git a/crates/burn-store/safetensors-tests/tests/multi_layer/multi_layer.safetensors b/crates/burn-store/safetensors-tests/tests/multi_layer/multi_layer.safetensors new file mode 100644 index 0000000..ce34270 Binary files /dev/null and b/crates/burn-store/safetensors-tests/tests/multi_layer/multi_layer.safetensors differ diff --git a/crates/burn-store/safetensors-tests/tests/test_mod.rs b/crates/burn-store/safetensors-tests/tests/test_mod.rs new file mode 100644 index 0000000..4a39d58 --- /dev/null +++ b/crates/burn-store/safetensors-tests/tests/test_mod.rs @@ -0,0 +1 @@ +mod multi_layer; diff --git a/crates/burn-store/src/adapter.rs b/crates/burn-store/src/adapter.rs new file mode 100644 index 0000000..102601a --- /dev/null +++ b/crates/burn-store/src/adapter.rs @@ -0,0 +1,1022 @@ +//! Module adapters for transforming tensor snapshots during save/load +//! +//! This module provides adapters for: +//! - PyTorch/Burn format conversion (weight transposition, parameter renaming) +//! - Mixed-precision storage (F32/F16 dtype casting via [`HalfPrecisionAdapter`]) +//! - Adapter chaining for composing multiple transformations + +use crate::TensorSnapshot; + +use alloc::boxed::Box; +use alloc::format; +use alloc::rc::Rc; +use alloc::string::String; +use alloc::string::ToString; +use alloc::vec; + +use burn_core::tensor::shape; +use burn_core::tensor::{DType, TensorData}; +use hashbrown::HashSet; + +// Module type names as they appear in the container_type field +// These come from the Module derive macro which uses stringify! on the struct name +// Format: "Struct:TypeName" for user-defined structs +mod module_names { + // The actual string constants that match what the Module derive macro produces + pub const LINEAR: &str = "Struct:Linear"; + pub const BATCH_NORM: &str = "Struct:BatchNorm"; + pub const LAYER_NORM: &str = "Struct:LayerNorm"; + pub const GROUP_NORM: &str = "Struct:GroupNorm"; + pub const EMBEDDING: &str = "Struct:Embedding"; + pub const CONV1D: &str = "Struct:Conv1d"; + pub const CONV2D: &str = "Struct:Conv2d"; + pub const CONV3D: &str = "Struct:Conv3d"; + pub const CONV_TRANSPOSE1D: &str = "Struct:ConvTranspose1d"; + pub const CONV_TRANSPOSE2D: &str = "Struct:ConvTranspose2d"; + pub const CONV_TRANSPOSE3D: &str = "Struct:ConvTranspose3d"; + pub const DEFORM_CONV2D: &str = "Struct:DeformConv2d"; + pub const INSTANCE_NORM: &str = "Struct:InstanceNorm"; + pub const RMS_NORM: &str = "Struct:RmsNorm"; + pub const PRELU: &str = "Struct:PRelu"; +} + +/// Trait for adapting tensor snapshots between different module formats +pub trait ModuleAdapter: Send + Sync { + /// Adapt a tensor snapshot based on its container type and parameter name + fn adapt(&self, snapshot: &TensorSnapshot) -> TensorSnapshot; + + /// Get alternative parameter name to try during matching + /// + /// When looking for a parameter in a module, this method provides an alternative + /// name to try if the direct name doesn't match. This enables matching parameters + /// with different naming conventions (e.g., PyTorch's "weight" vs Burn's "gamma"). + /// + /// # Arguments + /// * `param_name` - The parameter name we're looking for + /// * `container_type` - The type of container module (e.g., "BatchNorm") + /// + /// # Returns + /// Alternative parameter name to try, or None if no alternative exists + fn get_alternative_param_name( + &self, + _param_name: &str, + _container_type: &str, + ) -> Option { + None + } + + /// Clone the adapter into a boxed trait object + fn clone_box(&self) -> Box; + + /// Chain adapters together, applying `self` first and then `next`. + /// + /// This is useful when multiple transformations are required when importing model weights + /// (e.g. PyTorch -> Burn layout conversion, then dtype casting, then custom remapping). + /// + /// The semantics follow a simple pipeline: + /// - `adapt`: `next.adapt(&self.adapt(snapshot))` + /// - `get_alternative_param_name`: try `self` first; if it returns an alternative name, + /// try `next` with that name, otherwise return the first alternative name. + fn chain(self, next: A) -> ChainAdapter + where + Self: Sized + 'static, + A: ModuleAdapter + 'static, + { + ChainAdapter::new(self, next) + } +} + +impl Clone for Box { + fn clone(&self) -> Self { + self.clone_box() + } +} + +/// Adapter that applies two adapters in sequence. +/// +/// This allows composing smaller adapters instead of creating one large monolithic adapter. +#[derive(Clone)] +pub struct ChainAdapter { + first: Box, + second: Box, +} + +impl ChainAdapter { + /// Create a new adapter chain. + pub fn new(first: A, second: B) -> Self + where + A: ModuleAdapter + 'static, + B: ModuleAdapter + 'static, + { + Self { + first: Box::new(first), + second: Box::new(second), + } + } +} + +impl ModuleAdapter for ChainAdapter { + fn adapt(&self, snapshot: &TensorSnapshot) -> TensorSnapshot { + let snapshot = self.first.adapt(snapshot); + self.second.adapt(&snapshot) + } + + fn get_alternative_param_name(&self, param_name: &str, container_type: &str) -> Option { + if let Some(name) = self + .first + .get_alternative_param_name(param_name, container_type) + { + self.second + .get_alternative_param_name(&name, container_type) + .or(Some(name)) + } else { + self.second + .get_alternative_param_name(param_name, container_type) + } + } + + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } +} + +/// Identity adapter that passes tensors through unchanged +#[derive(Debug, Clone, Default)] +pub struct IdentityAdapter; + +impl ModuleAdapter for IdentityAdapter { + fn adapt(&self, snapshot: &TensorSnapshot) -> TensorSnapshot { + snapshot.clone() + } + + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } +} + +/// Returns the default set of module types that `HalfPrecisionAdapter` converts. +/// +/// Includes: Linear, Embedding, all Conv variants, LayerNorm, GroupNorm, +/// InstanceNorm, RmsNorm, PRelu. +/// +/// Excludes BatchNorm by default because `running_var` underflows in F16. +fn default_half_precision_modules() -> HashSet { + let modules = [ + module_names::LINEAR, + module_names::EMBEDDING, + module_names::CONV1D, + module_names::CONV2D, + module_names::CONV3D, + module_names::CONV_TRANSPOSE1D, + module_names::CONV_TRANSPOSE2D, + module_names::CONV_TRANSPOSE3D, + module_names::DEFORM_CONV2D, + module_names::LAYER_NORM, + module_names::GROUP_NORM, + module_names::INSTANCE_NORM, + module_names::RMS_NORM, + module_names::PRELU, + ]; + modules.iter().map(|s| s.to_string()).collect() +} + +/// Adapter for mixed-precision (F32/F16) model storage. +/// +/// Auto-detects conversion direction from the snapshot's dtype: +/// - F32 source -> cast to F16 (typical for saving) +/// - F16 source -> cast to F32 (typical for loading) +/// - Other dtypes -> passed through unchanged +/// +/// The same instance works for both `with_to_adapter` (save) and `with_from_adapter` (load). +/// +/// By default, converts weights in: Linear, Embedding, Conv*, LayerNorm, GroupNorm, +/// InstanceNorm, RmsNorm, PRelu. BatchNorm is excluded because `running_var` underflows in F16. +/// +/// # Examples +/// +/// Default usage (same adapter for save and load): +/// ```rust +/// # use burn_store::HalfPrecisionAdapter; +/// let adapter = HalfPrecisionAdapter::new(); +/// // store.with_to_adapter(adapter.clone()); // F32 -> F16 on save +/// // store.with_from_adapter(adapter); // F16 -> F32 on load +/// ``` +/// +/// Exclude a module type: +/// ```rust +/// # use burn_store::HalfPrecisionAdapter; +/// let adapter = HalfPrecisionAdapter::new() +/// .without_module("LayerNorm"); +/// ``` +/// +/// Add a custom module type: +/// ```rust +/// # use burn_store::HalfPrecisionAdapter; +/// let adapter = HalfPrecisionAdapter::new() +/// .with_module("CustomLayer"); +/// ``` +#[derive(Debug, Clone)] +pub struct HalfPrecisionAdapter { + modules: HashSet, +} + +impl HalfPrecisionAdapter { + /// Create a new adapter with the default set of modules. + pub fn new() -> Self { + Self { + modules: default_half_precision_modules(), + } + } + + /// Add a module type to convert. Accepts both short (`"MyLayer"`) and + /// qualified (`"Struct:MyLayer"`) forms. + /// + /// Note: short names are mapped to `"Struct:Name"`. If you have an Enum-based + /// module, use the qualified form `"Enum:MyModule"` explicitly. + pub fn with_module(mut self, module_type: impl Into) -> Self { + let name = module_type.into(); + if name.contains(':') { + self.modules.insert(name); + } else { + self.modules.insert(format!("Struct:{}", name)); + } + self + } + + /// Remove a module type from conversion. Accepts both short and qualified forms. + pub fn without_module(mut self, module_type: impl Into) -> Self { + let name = module_type.into(); + let key = if name.contains(':') { + name + } else { + format!("Struct:{}", name) + }; + assert!( + self.modules.contains(&key), + "without_module called with '{}' which is not in the module set", + key + ); + self.modules.remove(&key); + self + } + + /// Check whether the tensor belongs to a module that should be converted. + fn should_convert(&self, snapshot: &TensorSnapshot) -> bool { + snapshot + .module_type() + .is_some_and(|mt| self.modules.contains(&mt)) + } +} + +impl Default for HalfPrecisionAdapter { + fn default() -> Self { + Self::new() + } +} + +impl ModuleAdapter for HalfPrecisionAdapter { + fn adapt(&self, snapshot: &TensorSnapshot) -> TensorSnapshot { + // Determine target dtype from source: F32 -> F16, F16 -> F32, anything else -> skip + let target_dtype = match snapshot.dtype { + DType::F32 => DType::F16, + DType::F16 => DType::F32, + _ => return snapshot.clone(), + }; + + if !self.should_convert(snapshot) { + return snapshot.clone(); + } + + let original_data_fn = snapshot.clone_data_fn(); + + let cast_data_fn = Rc::new(move || { + let data = original_data_fn()?; + Ok(data.convert_dtype(target_dtype)) + }); + + TensorSnapshot::from_closure( + cast_data_fn, + target_dtype, + snapshot.shape.clone(), + snapshot.path_stack.clone().unwrap_or_default(), + snapshot.container_stack.clone().unwrap_or_default(), + snapshot.tensor_id.unwrap_or_default(), + ) + } + + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } +} + +/// Adapter for converting from PyTorch format to Burn format +/// +/// Handles: +/// - Linear layer weight transposition (PyTorch: [out, in] → Burn: [in, out]) +/// - Normalization parameter renaming (weight → gamma, bias → beta) +#[derive(Debug, Clone, Default)] +pub struct PyTorchToBurnAdapter; + +impl ModuleAdapter for PyTorchToBurnAdapter { + fn adapt(&self, snapshot: &TensorSnapshot) -> TensorSnapshot { + adapt_pytorch_tensor(snapshot, PyTorchConversionDirection::PyTorchToBurn) + } + + fn get_alternative_param_name(&self, param_name: &str, container_type: &str) -> Option { + // For PyTorch->Burn: When looking for Burn names (gamma/beta), try PyTorch names (weight/bias) + if is_normalization_layer(container_type) { + burn_norm_param_to_pytorch(param_name).map(|s| s.to_string()) + } else { + None + } + } + + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } +} + +/// Adapter for converting from Burn format to PyTorch format +/// +/// Handles: +/// - Linear layer weight transposition (Burn: [in, out] → PyTorch: [out, in]) +/// - Normalization parameter renaming (gamma → weight, beta → bias) +#[derive(Debug, Clone, Default)] +pub struct BurnToPyTorchAdapter; + +impl ModuleAdapter for BurnToPyTorchAdapter { + fn adapt(&self, snapshot: &TensorSnapshot) -> TensorSnapshot { + adapt_pytorch_tensor(snapshot, PyTorchConversionDirection::BurnToPyTorch) + } + + fn get_alternative_param_name(&self, param_name: &str, container_type: &str) -> Option { + // For Burn->PyTorch: When looking for PyTorch names (weight/bias), try Burn names (gamma/beta) + if is_normalization_layer(container_type) { + pytorch_norm_param_to_burn(param_name).map(|s| s.to_string()) + } else { + None + } + } + + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } +} + +/// Direction of PyTorch conversion for parameter naming +#[derive(Debug, Clone, Copy)] +enum PyTorchConversionDirection { + PyTorchToBurn, + BurnToPyTorch, +} + +/// Check if container type is a normalization layer +fn is_normalization_layer(container_type: &str) -> bool { + matches!( + container_type, + module_names::BATCH_NORM + | module_names::LAYER_NORM + | module_names::GROUP_NORM + | module_names::RMS_NORM + ) +} + +/// Map PyTorch normalization parameter name to Burn +fn pytorch_norm_param_to_burn(param_name: &str) -> Option<&'static str> { + match param_name { + "weight" => Some("gamma"), + "bias" => Some("beta"), + _ => None, + } +} + +/// Map Burn normalization parameter name to PyTorch +fn burn_norm_param_to_pytorch(param_name: &str) -> Option<&'static str> { + match param_name { + "gamma" => Some("weight"), + "beta" => Some("bias"), + _ => None, + } +} + +/// Core tensor adaptation logic for PyTorch format conversions +fn adapt_pytorch_tensor( + snapshot: &TensorSnapshot, + direction: PyTorchConversionDirection, +) -> TensorSnapshot { + // Extract path and parameter name + let (path_stack, param_name) = match get_path_and_param(snapshot) { + Some(result) => result, + None => return snapshot.clone(), + }; + + // Get module type for matching (ignores Vec/Array wrappers) + let module_type = match snapshot.module_type() { + Some(mt) => mt, + None => return snapshot.clone(), // No user-defined module found + }; + + // Linear: transpose weight (bidirectional - same operation both ways) + if module_type == module_names::LINEAR && param_name == "weight" && snapshot.shape.len() == 2 { + return transpose_2d_tensor(snapshot); + } + + // Normalization layers: rename parameters based on direction + if is_normalization_layer(&module_type) { + let new_name = match direction { + PyTorchConversionDirection::PyTorchToBurn => pytorch_norm_param_to_burn(param_name), + PyTorchConversionDirection::BurnToPyTorch => burn_norm_param_to_pytorch(param_name), + }; + + if let Some(new_name) = new_name { + return rename_parameter(snapshot, path_stack, new_name); + } + } + + snapshot.clone() +} + +/// Extract path stack and parameter name from snapshot +fn get_path_and_param(snapshot: &TensorSnapshot) -> Option<(&[String], &str)> { + let path_stack = snapshot.path_stack.as_ref()?; + let param_name = path_stack.last()?.as_str(); + Some((path_stack.as_slice(), param_name)) +} + +/// Rename a parameter in the snapshot +fn rename_parameter( + snapshot: &TensorSnapshot, + path_stack: &[String], + new_name: &str, +) -> TensorSnapshot { + let mut new_path = path_stack.to_vec(); + *new_path.last_mut().unwrap() = new_name.to_string(); + + TensorSnapshot::from_closure( + snapshot.clone_data_fn(), + snapshot.dtype, + snapshot.shape.clone(), + new_path, + snapshot.container_stack.clone().unwrap_or_default(), + snapshot.tensor_id.unwrap_or_default(), + ) +} + +/// Transpose a 2D tensor +fn transpose_2d_tensor(snapshot: &TensorSnapshot) -> TensorSnapshot { + if snapshot.shape.len() != 2 { + return snapshot.clone(); + } + + let original_data_fn = snapshot.clone_data_fn(); + let dtype = snapshot.dtype; + let transposed_shape = shape![snapshot.shape[1], snapshot.shape[0]]; + + // Create a lazy closure that transposes when called + let transposed_data_fn = Rc::new(move || { + let data = original_data_fn()?; + Ok(transpose_tensor_data(data)) + }); + + TensorSnapshot::from_closure( + transposed_data_fn, + dtype, + transposed_shape, + snapshot.path_stack.clone().unwrap_or_default(), + snapshot.container_stack.clone().unwrap_or_default(), + snapshot.tensor_id.unwrap_or_default(), + ) +} + +/// Transpose tensor data (assumes 2D shape is already validated) +fn transpose_tensor_data(data: TensorData) -> TensorData { + let shape = &data.shape; + let rows = shape[0]; + let cols = shape[1]; + let transposed_shape = vec![cols, rows]; + + // Get the raw bytes and element size + let bytes = data.as_bytes(); + let element_size = data.dtype.size(); + + // Create a new buffer for transposed data + let mut transposed_bytes = vec![0u8; bytes.len()]; + + // Transpose at the byte level - works for any data type + for i in 0..rows { + for j in 0..cols { + let src_idx = (i * cols + j) * element_size; + let dst_idx = (j * rows + i) * element_size; + + // Copy the bytes for this element + transposed_bytes[dst_idx..dst_idx + element_size] + .copy_from_slice(&bytes[src_idx..src_idx + element_size]); + } + } + + // Create new TensorData from transposed bytes + TensorData::from_bytes_vec(transposed_bytes, transposed_shape, data.dtype) +} + +#[cfg(test)] +mod tests { + use super::*; + use alloc::rc::Rc; + use alloc::sync::Arc; + use burn_core::tensor::{DType, Shape, TensorData}; + use core::sync::atomic::{AtomicUsize, Ordering}; + + #[test] + fn test_module_names_match_burn_nn() { + // If these types are renamed or moved in `burn-nn`, this test will fail to compile. + #[allow(unused_imports)] + use burn_nn::{ + BatchNorm, Embedding, GroupNorm, InstanceNorm, LayerNorm, Linear, PRelu, RmsNorm, + conv::{ + Conv1d, Conv2d, Conv3d, ConvTranspose1d, ConvTranspose2d, ConvTranspose3d, + DeformConv2d, + }, + }; + + assert_eq!(module_names::LINEAR, "Struct:Linear"); + assert_eq!(module_names::BATCH_NORM, "Struct:BatchNorm"); + assert_eq!(module_names::LAYER_NORM, "Struct:LayerNorm"); + assert_eq!(module_names::GROUP_NORM, "Struct:GroupNorm"); + assert_eq!(module_names::EMBEDDING, "Struct:Embedding"); + assert_eq!(module_names::CONV1D, "Struct:Conv1d"); + assert_eq!(module_names::CONV2D, "Struct:Conv2d"); + assert_eq!(module_names::CONV3D, "Struct:Conv3d"); + assert_eq!(module_names::CONV_TRANSPOSE1D, "Struct:ConvTranspose1d"); + assert_eq!(module_names::CONV_TRANSPOSE2D, "Struct:ConvTranspose2d"); + assert_eq!(module_names::CONV_TRANSPOSE3D, "Struct:ConvTranspose3d"); + assert_eq!(module_names::DEFORM_CONV2D, "Struct:DeformConv2d"); + assert_eq!(module_names::INSTANCE_NORM, "Struct:InstanceNorm"); + assert_eq!(module_names::RMS_NORM, "Struct:RmsNorm"); + assert_eq!(module_names::PRELU, "Struct:PRelu"); + } + + fn create_test_snapshot(path: &str, shape: Shape, container_type: &str) -> TensorSnapshot { + let path_parts: Vec = path.split('.').map(|s| s.to_string()).collect(); + let values = vec![1.0f32; shape.iter().product()]; + let data = TensorData::new(values, shape.clone()); + + TensorSnapshot::from_closure( + Rc::new(move || Ok(data.clone())), + DType::F32, + shape, + path_parts, + vec![container_type.to_string()], + burn_core::module::ParamId::new(), + ) + } + + #[test] + fn test_pytorch_to_burn_linear_weight() { + let adapter = PyTorchToBurnAdapter; + + // Linear layer weight should be transposed + let snapshot = create_test_snapshot("fc.weight", shape![10, 5], module_names::LINEAR); + let adapted = adapter.adapt(&snapshot); + assert_eq!(adapted.shape, shape![5, 10]); + + // Linear layer bias should not be transposed + let snapshot = create_test_snapshot("fc.bias", shape![10], module_names::LINEAR); + let adapted = adapter.adapt(&snapshot); + assert_eq!(adapted.shape, shape![10]); + } + + #[test] + fn test_pytorch_to_burn_norm_params() { + let adapter = PyTorchToBurnAdapter; + + // BatchNorm weight -> gamma + let snapshot = create_test_snapshot("norm.weight", shape![10], module_names::BATCH_NORM); + let adapted = adapter.adapt(&snapshot); + assert_eq!(adapted.full_path(), "norm.gamma"); + + // BatchNorm bias -> beta + let snapshot = create_test_snapshot("norm.bias", shape![10], module_names::BATCH_NORM); + let adapted = adapter.adapt(&snapshot); + assert_eq!(adapted.full_path(), "norm.beta"); + } + + #[test] + fn test_burn_to_pytorch_linear_weight() { + let adapter = BurnToPyTorchAdapter; + + // Linear layer weight should be transposed + let snapshot = create_test_snapshot("fc.weight", shape![5, 10], module_names::LINEAR); + let adapted = adapter.adapt(&snapshot); + assert_eq!(adapted.shape, shape![10, 5]); + } + + #[test] + fn test_burn_to_pytorch_norm_params() { + let adapter = BurnToPyTorchAdapter; + + // BatchNorm gamma -> weight + let snapshot = create_test_snapshot("norm.gamma", shape![10], module_names::BATCH_NORM); + let adapted = adapter.adapt(&snapshot); + assert_eq!(adapted.full_path(), "norm.weight"); + + // BatchNorm beta -> bias + let snapshot = create_test_snapshot("norm.beta", shape![10], module_names::BATCH_NORM); + let adapted = adapter.adapt(&snapshot); + assert_eq!(adapted.full_path(), "norm.bias"); + } + + #[test] + fn test_transpose_different_dtypes() { + // Test that transpose works for different data types + + // Test with F32 + let f32_data = TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3]); + let transposed = transpose_tensor_data(f32_data); + assert_eq!(transposed.shape, shape![3, 2]); + let values = transposed.to_vec::().unwrap(); + assert_eq!(values, vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0]); + + // Test with I32 + let i32_data = TensorData::new(vec![1i32, 2, 3, 4, 5, 6], [2, 3]); + let transposed = transpose_tensor_data(i32_data); + assert_eq!(transposed.shape, shape![3, 2]); + let values = transposed.to_vec::().unwrap(); + assert_eq!(values, vec![1, 4, 2, 5, 3, 6]); + + // Test with F64 + let f64_data = TensorData::new(vec![1.0f64, 2.0, 3.0, 4.0], [2, 2]); + let transposed = transpose_tensor_data(f64_data); + assert_eq!(transposed.shape, shape![2, 2]); + let values = transposed.to_vec::().unwrap(); + assert_eq!(values, vec![1.0, 3.0, 2.0, 4.0]); + } + + #[test] + fn test_no_container_info() { + let adapter = PyTorchToBurnAdapter; + + // Without container info, adapter returns unchanged for non-norm parameters + let mut snapshot = create_test_snapshot("fc.weight", shape![10, 5], module_names::LINEAR); + snapshot.container_stack = None; + + // Without container info, no transformation occurs for linear layers + let adapted = adapter.adapt(&snapshot); + assert_eq!(adapted.shape, shape![10, 5]); // No transposition without container info + + // Test a non-linear, non-norm parameter - should pass through unchanged + let mut snapshot2 = create_test_snapshot("other.weight", shape![10, 5], "Struct:Other"); + snapshot2.container_stack = None; + let adapted2 = adapter.adapt(&snapshot2); + assert_eq!(adapted2.shape, shape![10, 5]); // No transposition + } + + #[derive(Clone)] + struct RenameParamAdapter { + from: &'static str, + to: &'static str, + called: Arc, + } + + impl ModuleAdapter for RenameParamAdapter { + fn adapt(&self, snapshot: &TensorSnapshot) -> TensorSnapshot { + self.called.fetch_add(1, Ordering::Relaxed); + + let path_stack = match snapshot.path_stack.as_ref() { + Some(stack) => stack, + None => return snapshot.clone(), + }; + let param = match path_stack.last() { + Some(p) => p.as_str(), + None => return snapshot.clone(), + }; + if param != self.from { + return snapshot.clone(); + } + + let mut new_path = path_stack.to_vec(); + *new_path.last_mut().unwrap() = self.to.to_string(); + + TensorSnapshot::from_closure( + snapshot.clone_data_fn(), + snapshot.dtype, + snapshot.shape.clone(), + new_path, + snapshot.container_stack.clone().unwrap_or_default(), + snapshot.tensor_id.unwrap_or_default(), + ) + } + + fn get_alternative_param_name( + &self, + _param_name: &str, + _container_type: &str, + ) -> Option { + None + } + + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } + } + + #[derive(Clone)] + struct AltNameAdapter { + from: &'static str, + to: &'static str, + called: Arc, + } + + impl ModuleAdapter for AltNameAdapter { + fn adapt(&self, snapshot: &TensorSnapshot) -> TensorSnapshot { + TensorSnapshot::from_closure( + snapshot.clone_data_fn(), + snapshot.dtype, + snapshot.shape.clone(), + snapshot.path_stack.clone().unwrap_or_default(), + snapshot.container_stack.clone().unwrap_or_default(), + snapshot.tensor_id.unwrap_or_default(), + ) + } + + fn get_alternative_param_name( + &self, + param_name: &str, + _container_type: &str, + ) -> Option { + self.called.fetch_add(1, Ordering::Relaxed); + if param_name == self.from { + Some(self.to.to_string()) + } else { + None + } + } + + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } + } + + #[test] + fn test_chain_adapter_pipes_adapt() { + let called1 = Arc::new(AtomicUsize::new(0)); + let called2 = Arc::new(AtomicUsize::new(0)); + + let a = RenameParamAdapter { + from: "weight", + to: "a", + called: called1.clone(), + }; + let b = RenameParamAdapter { + from: "a", + to: "b", + called: called2.clone(), + }; + + let chain = a.chain(b); + let snapshot = create_test_snapshot("fc.weight", shape![2, 2], module_names::LINEAR); + let adapted = chain.adapt(&snapshot); + + assert_eq!(adapted.full_path(), "fc.b"); + assert_eq!(called1.load(Ordering::Relaxed), 1); + assert_eq!(called2.load(Ordering::Relaxed), 1); + } + + #[test] + fn test_chain_adapter_alternative_name_pipes_and_fallbacks() { + let called1 = Arc::new(AtomicUsize::new(0)); + let called2 = Arc::new(AtomicUsize::new(0)); + + let a = AltNameAdapter { + from: "gamma", + to: "weight", + called: called1.clone(), + }; + let b = AltNameAdapter { + from: "weight", + to: "scale", + called: called2.clone(), + }; + + let chain = a.chain(b); + let alt = chain.get_alternative_param_name("gamma", module_names::LAYER_NORM); + assert_eq!(alt.as_deref(), Some("scale")); + assert_eq!(called1.load(Ordering::Relaxed), 1); + assert_eq!(called2.load(Ordering::Relaxed), 1); + + // If the second adapter doesn't have a mapping for the first alternative, + // fall back to the first alternative name. + let called1 = Arc::new(AtomicUsize::new(0)); + let called2 = Arc::new(AtomicUsize::new(0)); + let a = AltNameAdapter { + from: "gamma", + to: "weight", + called: called1.clone(), + }; + let b = AltNameAdapter { + from: "something-else", + to: "unused", + called: called2.clone(), + }; + let chain = a.chain(b); + let alt = chain.get_alternative_param_name("gamma", module_names::LAYER_NORM); + assert_eq!(alt.as_deref(), Some("weight")); + assert_eq!(called1.load(Ordering::Relaxed), 1); + assert_eq!(called2.load(Ordering::Relaxed), 1); + + // If the first adapter doesn't provide an alternative, try the second with the original name. + let called1 = Arc::new(AtomicUsize::new(0)); + let called2 = Arc::new(AtomicUsize::new(0)); + let a = AltNameAdapter { + from: "something-else", + to: "unused", + called: called1.clone(), + }; + let b = AltNameAdapter { + from: "gamma", + to: "weight", + called: called2.clone(), + }; + let chain = a.chain(b); + let alt = chain.get_alternative_param_name("gamma", module_names::LAYER_NORM); + assert_eq!(alt.as_deref(), Some("weight")); + assert_eq!(called1.load(Ordering::Relaxed), 1); + assert_eq!(called2.load(Ordering::Relaxed), 1); + + // clone_box must preserve behavior. + let boxed = chain.clone_box(); + let alt = boxed.get_alternative_param_name("gamma", module_names::LAYER_NORM); + assert_eq!(alt.as_deref(), Some("weight")); + } + + #[test] + fn test_half_precision_f32_to_f16() { + let adapter = HalfPrecisionAdapter::new(); + let snapshot = create_test_snapshot("fc.weight", shape![2, 3], module_names::LINEAR); + + let adapted = adapter.adapt(&snapshot); + assert_eq!(adapted.dtype, DType::F16); + assert_eq!(adapted.shape, shape![2, 3]); + + let data = adapted.to_data().unwrap(); + assert_eq!(data.dtype, DType::F16); + } + + #[test] + fn test_half_precision_f16_to_f32() { + let adapter = HalfPrecisionAdapter::new(); + + // Create an F16 snapshot + let values = vec![1.0f32; 6]; + let data = TensorData::new(values, shape![2, 3]).convert_dtype(DType::F16); + let path_parts = vec!["fc".to_string(), "weight".to_string()]; + let snapshot = TensorSnapshot::from_closure( + Rc::new(move || Ok(data.clone())), + DType::F16, + shape![2, 3], + path_parts, + vec![module_names::LINEAR.to_string()], + burn_core::module::ParamId::new(), + ); + + let adapted = adapter.adapt(&snapshot); + assert_eq!(adapted.dtype, DType::F32); + } + + #[test] + fn test_half_precision_skips_batch_norm() { + let adapter = HalfPrecisionAdapter::new(); + + // BatchNorm is excluded by default + let snapshot = create_test_snapshot("norm.weight", shape![10], module_names::BATCH_NORM); + let adapted = adapter.adapt(&snapshot); + assert_eq!(adapted.dtype, DType::F32); // unchanged + } + + #[test] + fn test_half_precision_converts_default_modules() { + let adapter = HalfPrecisionAdapter::new(); + + // Linear + let snapshot = create_test_snapshot("fc.weight", shape![2, 3], module_names::LINEAR); + assert_eq!(adapter.adapt(&snapshot).dtype, DType::F16); + + // Embedding + let snapshot = create_test_snapshot("emb.weight", shape![100, 64], module_names::EMBEDDING); + assert_eq!(adapter.adapt(&snapshot).dtype, DType::F16); + + // Conv2d + let snapshot = + create_test_snapshot("conv.weight", shape![3, 3, 3, 3], module_names::CONV2D); + assert_eq!(adapter.adapt(&snapshot).dtype, DType::F16); + + // LayerNorm (included by default) + let snapshot = create_test_snapshot("norm.gamma", shape![10], module_names::LAYER_NORM); + assert_eq!(adapter.adapt(&snapshot).dtype, DType::F16); + + // GroupNorm + let snapshot = create_test_snapshot("gn.gamma", shape![10], module_names::GROUP_NORM); + assert_eq!(adapter.adapt(&snapshot).dtype, DType::F16); + + // RmsNorm + let snapshot = create_test_snapshot("rms.weight", shape![10], module_names::RMS_NORM); + assert_eq!(adapter.adapt(&snapshot).dtype, DType::F16); + } + + #[test] + fn test_half_precision_without_module() { + let adapter = HalfPrecisionAdapter::new().without_module("LayerNorm"); + + // LayerNorm removed from conversion set + let snapshot = create_test_snapshot("norm.gamma", shape![10], module_names::LAYER_NORM); + assert_eq!(adapter.adapt(&snapshot).dtype, DType::F32); + + // Linear still converted + let snapshot = create_test_snapshot("fc.weight", shape![2, 3], module_names::LINEAR); + assert_eq!(adapter.adapt(&snapshot).dtype, DType::F16); + } + + #[test] + fn test_half_precision_with_module() { + let adapter = HalfPrecisionAdapter::new().with_module("CustomLayer"); + + // Custom module should now be converted + let snapshot = create_test_snapshot("custom.weight", shape![5], "Struct:CustomLayer"); + assert_eq!(adapter.adapt(&snapshot).dtype, DType::F16); + } + + #[test] + fn test_half_precision_with_qualified_name() { + let adapter = HalfPrecisionAdapter::new().with_module("Struct:CustomLayer"); + + let snapshot = create_test_snapshot("custom.weight", shape![5], "Struct:CustomLayer"); + assert_eq!(adapter.adapt(&snapshot).dtype, DType::F16); + } + + #[test] + fn test_half_precision_chain() { + let adapter = PyTorchToBurnAdapter.chain(HalfPrecisionAdapter::new()); + + let snapshot = create_test_snapshot("fc.weight", shape![10, 5], module_names::LINEAR); + let adapted = adapter.adapt(&snapshot); + + // Should be both transposed and cast + assert_eq!(adapted.shape, shape![5, 10]); + assert_eq!(adapted.dtype, DType::F16); + } + + #[test] + fn test_half_precision_skips_no_container() { + let adapter = HalfPrecisionAdapter::new(); + let mut snapshot = create_test_snapshot("fc.weight", shape![2, 3], module_names::LINEAR); + snapshot.container_stack = None; + + // No module type info: skip + let adapted = adapter.adapt(&snapshot); + assert_eq!(adapted.dtype, DType::F32); + } + + #[test] + fn test_half_precision_skips_non_float() { + use burn_core::tensor::quantization::QuantScheme; + + let adapter = HalfPrecisionAdapter::new(); + + // QFloat source: skip + let qfloat_dtype = DType::QFloat(QuantScheme::default()); + let snapshot = create_test_snapshot("fc.weight", shape![2, 3], module_names::LINEAR); + let qfloat_snapshot = TensorSnapshot::from_closure( + snapshot.clone_data_fn(), + qfloat_dtype, + snapshot.shape.clone(), + snapshot.path_stack.clone().unwrap_or_default(), + snapshot.container_stack.clone().unwrap_or_default(), + snapshot.tensor_id.unwrap_or_default(), + ); + let adapted = adapter.adapt(&qfloat_snapshot); + assert_eq!(adapted.dtype, qfloat_dtype); + } + + #[test] + fn test_half_precision_default_module_count() { + let adapter = HalfPrecisionAdapter::new(); + // 14 modules: Linear, Embedding, Conv1d-3d, ConvTranspose1d-3d, + // DeformConv2d, LayerNorm, GroupNorm, InstanceNorm, RmsNorm, PRelu + assert_eq!(adapter.modules.len(), 14); + } + + #[test] + fn test_half_precision_without_module_qualified() { + let adapter = HalfPrecisionAdapter::new().without_module("Struct:LayerNorm"); + + let snapshot = create_test_snapshot("norm.gamma", shape![10], module_names::LAYER_NORM); + assert_eq!(adapter.adapt(&snapshot).dtype, DType::F32); + } + + #[test] + fn test_half_precision_with_module_batch_norm_opt_in() { + let adapter = HalfPrecisionAdapter::new().with_module("BatchNorm"); + + let snapshot = create_test_snapshot("bn.weight", shape![10], module_names::BATCH_NORM); + assert_eq!(adapter.adapt(&snapshot).dtype, DType::F16); + } +} diff --git a/crates/burn-store/src/applier.rs b/crates/burn-store/src/applier.rs new file mode 100644 index 0000000..648c095 --- /dev/null +++ b/crates/burn-store/src/applier.rs @@ -0,0 +1,585 @@ +//! Applier that correctly applies tensor snapshots with adapter support + +use alloc::boxed::Box; +use alloc::format; +use alloc::string::{String, ToString}; +use alloc::vec::Vec; + +use hashbrown::{HashMap, HashSet}; + +use burn_core::module::{ModuleMapper, Param}; +use burn_core::tensor::{Bool, Device, Int, Shape, Tensor}; + +use crate::apply_result::{ApplyError, ApplyResult}; +use crate::{ModuleAdapter, PathFilter, TensorSnapshot}; + +/// Applier that applies tensor snapshots to module parameters +/// with proper adapter support using container type information +pub struct Applier { + /// Map of tensor paths to their snapshots + snapshots: HashMap, + /// Current path in the module hierarchy + path_stack: Vec, + /// Current container type stack in the module hierarchy + container_stack: Vec, + /// Optional filter for selective application + filter: Option, + /// Optional adapter to transform tensors based on container types + adapter: Option>, + /// Successfully applied tensor paths + applied: Vec, + /// Skipped tensor paths + skipped: HashSet, + /// Errors encountered during application + errors: Vec, + /// Track visited paths with their container stacks (in dot notation) to find missing tensors + visited_paths: HashMap, + /// Skip enum variant names when matching paths + /// When true, "feature.BaseConv.weight" will also try to match "feature.weight" + skip_enum_variants: bool, +} + +impl Applier { + /// Create a new applier with snapshots, optional filter, and optional adapter + /// + /// # Arguments + /// + /// * `views` - A vector of TensorSnapshot objects to apply + /// * `filter` - An optional [`PathFilter`] to determine which tensors to apply. + /// When `None`, all available tensors are applied. + /// * `adapter` - Optional adapter to transform tensors based on container types + /// * `skip_enum_variants` - Skip enum variant names when matching paths + pub fn new( + views: Vec, + filter: Option, + adapter: Option>, + skip_enum_variants: bool, + ) -> Self { + let views_map: HashMap = views + .into_iter() + .map(|view| (view.full_path(), view)) + .collect(); + + Self { + snapshots: views_map, + path_stack: Vec::new(), + container_stack: Vec::new(), + filter, + adapter, + applied: Vec::new(), + skipped: HashSet::new(), + errors: Vec::new(), + visited_paths: HashMap::new(), + skip_enum_variants, + } + } + + /// Get the current path in the module hierarchy + fn current_path(&self) -> String { + self.path_stack.join(".") + } + + /// Get the current module type (last Struct/Enum in container stack) + fn current_module_type(&self) -> Option<&str> { + self.container_stack + .iter() + .rev() + .find(|ct| ct.starts_with("Struct:") || ct.starts_with("Enum:")) + .map(|s| s.as_str()) + } + + /// Check if a tensor should be applied based on filter + fn should_apply(&self) -> bool { + match &self.filter { + None => true, + Some(f) => f.matches_with_container_path(&self.path_stack, &self.container_stack), + } + } + + /// Convert the applier into a result + pub fn into_result(self) -> ApplyResult { + let mut unused: Vec = self + .snapshots + .keys() + .filter(|path| !self.visited_paths.contains_key(*path) && !self.skipped.contains(*path)) + .cloned() + .collect(); + // Sort for stable output order + unused.sort(); + + // Create a set of successfully applied paths for efficient lookup + let applied_set: HashSet = self.applied.iter().cloned().collect(); + + // Extract paths that have errors - these are not "missing", they were found but had issues + let errored_paths: HashSet = self + .errors + .iter() + .map(|e| match e { + ApplyError::ShapeMismatch { path, .. } => path.clone(), + ApplyError::DTypeMismatch { path, .. } => path.clone(), + ApplyError::AdapterError { path, .. } => path.clone(), + ApplyError::LoadError { path, .. } => path.clone(), + }) + .collect(); + + // A path is missing if it was visited but not successfully applied, not skipped, and didn't have an error + // Store both the path and its container stack (in dot notation) + let mut missing: Vec<(String, String)> = self + .visited_paths + .into_iter() + .filter(|(p, _)| { + !applied_set.contains(p) && !self.skipped.contains(p) && !errored_paths.contains(p) + }) + .collect(); + // Sort for stable output order (by path) + missing.sort_by(|a, b| a.0.cmp(&b.0)); + + // Convert skipped HashSet to sorted Vec for stable output + let mut skipped: Vec = self.skipped.into_iter().collect(); + skipped.sort(); + + ApplyResult { + applied: self.applied, + skipped, + missing, + unused, + errors: self.errors, + } + } + + /// Apply a tensor snapshot with shape validation and optional adapter transformation + /// Returns None if snapshot not found, filtered, or validation fails + fn apply_tensor( + &mut self, + target_device: &Device, + target_shape: Shape, + ) -> Option> + where + K: burn_core::tensor::kind::Basic, + { + let path = self.current_path(); + let container_stack_str = self.container_stack.join("."); + self.visited_paths.insert(path.clone(), container_stack_str); + + // Try to get snapshot with original path first + let mut snapshot = self.snapshots.get(&path).cloned(); + + // If not found and we have an adapter, try alternative parameter names + if snapshot.is_none() + && let Some(ref adapter) = self.adapter + && let Some(module_type) = self.current_module_type() + { + // Get alternative name based on current module type (user-defined module only) + let param_name = self.path_stack.last()?; + + if let Some(alt_name) = adapter.get_alternative_param_name(param_name, module_type) { + // Build alternative path with parameter name substitution + let mut alt_path_stack = self.path_stack.clone(); + *alt_path_stack.last_mut().unwrap() = alt_name.clone(); + let alt_path = alt_path_stack.join("."); + + // Try to get snapshot with alternative name + snapshot = self.snapshots.get(&alt_path).cloned(); + + // Don't mark the alternative path as visited - only the original Burn path + // should be tracked. The alternative path is just for lookup. + } + } + + let mut snapshot = snapshot?; + + // Apply adapter transformation using current container_stack context (for data transformation like transpose) + if let Some(ref adapter) = self.adapter { + // Create a temporary snapshot with current context for adaptation + let snapshot_with_context = TensorSnapshot::from_closure( + snapshot.clone_data_fn(), + snapshot.dtype, + snapshot.shape.clone(), + self.path_stack.clone(), + self.container_stack.clone(), + snapshot.tensor_id.unwrap_or_default(), + ); + + // Transform using adapter (handles transpose) + snapshot = adapter.adapt(&snapshot_with_context); + } + + // Check if we should apply based on filter + if !self.should_apply() { + self.skipped.insert(path.clone()); + return None; + } + + // Load tensor data + let data = match snapshot.to_data() { + Ok(data) => data, + Err(e) => { + self.errors.push(ApplyError::LoadError { + path: path.clone(), + message: format!("Failed to load tensor data: {:?}", e), + }); + return None; // Signal caller to fall back to initialization + } + }; + + // Validate shape + if data.shape != target_shape { + self.errors.push(ApplyError::ShapeMismatch { + path: path.clone(), + expected: target_shape, + found: data.shape, + }); + return None; // Signal caller to fall back to initialization + } + + self.applied.push(path); + Some(Tensor::from_data(data, (target_device, snapshot.dtype))) + } +} + +impl ModuleMapper for Applier { + fn enter_module(&mut self, name: &str, container_type: &str) { + // Always track the container type for proper module type detection + self.container_stack.push(container_type.to_string()); + + // Only add to path if it's not an enum variant (when skip_enum_variants is enabled) + // This ensures paths are built without enum variant names from the start + if !self.skip_enum_variants || !container_type.starts_with("Enum:") { + self.path_stack.push(name.to_string()); + } + } + + fn exit_module(&mut self, _name: &str, container_type: &str) { + self.container_stack.pop(); + + // Only pop from path if we added it (not an enum variant when skip_enum_variants is enabled) + if !self.skip_enum_variants || !container_type.starts_with("Enum:") { + self.path_stack.pop(); + } + } + + fn map_float(&mut self, param: Param>) -> Param> { + let param_id = param.id; + let target_device = param.lazy_device(); + let target_shape = param.lazy_shape(); + + // Try to apply snapshot with shape validation + match self.apply_tensor(&target_device, target_shape) { + Some(tensor) => { + // We have a tensor to apply - load it + param.transform_for_load(tensor, param_id) + } + None => { + // No snapshot, filtered, or validation failed - return param unchanged + param + } + } + } + + fn map_int(&mut self, param: Param>) -> Param> { + let param_id = param.id; + let target_device = param.lazy_device(); + let target_shape = param.lazy_shape(); + + // Try to apply snapshot with shape validation + match self.apply_tensor(&target_device, target_shape) { + Some(tensor) => { + // We have a tensor to apply - load it + param.transform_for_load(tensor, param_id) + } + None => { + // No snapshot, filtered, or validation failed - return param unchanged + param + } + } + } + + fn map_bool( + &mut self, + param: Param>, + ) -> Param> { + let param_id = param.id; + let target_device = param.lazy_device(); + let target_shape = param.lazy_shape(); + + // Try to apply snapshot with shape validation + match self.apply_tensor(&target_device, target_shape) { + Some(tensor) => { + // We have a tensor to apply - load it + param.transform_for_load(tensor, param_id) + } + None => { + // No snapshot, filtered, or validation failed - return param unchanged + param + } + } + } +} + +#[cfg(all(test, feature = "std", target_has_atomic = "ptr"))] +mod tests { + use super::*; + use burn_core::module::{ModuleMapper, Param, ParamId}; + use burn_core::tensor::{DType, Tensor, TensorData}; + + #[test] + fn root_level_parameters() { + let device = Default::default(); + + // Create root-level parameters (not inside any module) + let weight = Param::>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + let bias = Param::>::from_data([5.0, 6.0], &device); + + // Create snapshots with root-level paths (single-element path, no nested modules) + let weight_snapshot = crate::TensorSnapshot::from_data( + weight.val().to_data(), + vec!["weight".to_string()], // root-level parameter name + vec![], // no container + ParamId::new(), + ); + + let bias_snapshot = crate::TensorSnapshot::from_data( + bias.val().to_data(), + vec!["bias".to_string()], // root-level parameter name + vec![], // no container + ParamId::new(), + ); + + // Create applier with root-level snapshots + let mut applier = Applier::new(vec![weight_snapshot, bias_snapshot], None, None, false); + + // Create new params to load into + let weight_target = Param::initialized(ParamId::new(), Tensor::<2>::zeros([2, 2], &device)); + let bias_target = Param::initialized(ParamId::new(), Tensor::<1>::zeros([2], &device)); + + // Apply using the ModuleMapper interface - simulate module traversal + // Enter "weight" path (as if we're visiting a field named "weight") + applier.enter_module("weight", ""); + let weight_loaded = applier.map_float(weight_target); + applier.exit_module("weight", ""); + + // Enter "bias" path (as if we're visiting a field named "bias") + applier.enter_module("bias", ""); + let bias_loaded = applier.map_float(bias_target); + applier.exit_module("bias", ""); + + // Verify values were loaded + let weight_data = weight_loaded.val().to_data().to_vec::().unwrap(); + let bias_data = bias_loaded.val().to_data().to_vec::().unwrap(); + + assert_eq!(weight_data, vec![1.0, 2.0, 3.0, 4.0]); + assert_eq!(bias_data, vec![5.0, 6.0]); + + // Verify applier result + let result = applier.into_result(); + assert_eq!(result.applied.len(), 2); + assert_eq!(result.errors.len(), 0); + } + + /// Test that the applier preserves dtype when loading tensor data. + /// This is a regression test for the bug where F16 tensors were being + /// loaded as F32 because `Tensor::from_data` was used instead of + /// without specifying the target dtype + #[test] + fn dtype_preservation_f64() { + let device = Default::default(); + + // Create TensorData with F64 dtype explicitly + let f64_data = TensorData::new(vec![1.0f64, 2.0, 3.0, 4.0], [2, 2]); + assert_eq!(f64_data.dtype, DType::F64, "Test setup: data should be F64"); + + // Create a snapshot with F64 data + let snapshot = crate::TensorSnapshot::from_data( + f64_data.clone(), + vec!["weight".to_string()], + vec![], + ParamId::new(), + ); + assert_eq!( + snapshot.dtype, + DType::F64, + "Snapshot should preserve F64 dtype" + ); + + // Create applier with the F64 snapshot + let mut applier = Applier::new(vec![snapshot], None, None, false); + + // Create target parameter + let target = Param::initialized( + ParamId::new(), + Tensor::<2>::zeros([2, 2], (&device, DType::F64)), + ); + + // Apply the snapshot + applier.enter_module("weight", ""); + let loaded = applier.map_float(target); + applier.exit_module("weight", ""); + + // Verify dtype is preserved - this would fail before the fix + // because the data would be converted to the backend's default FloatElem + assert_eq!( + loaded.val().dtype(), + DType::F64, + "Loaded tensor should have F64 dtype" + ); + + // Verify data values are correct + let loaded_data = loaded.val().to_data().to_vec::().unwrap(); + assert_eq!(loaded_data, vec![1.0, 2.0, 3.0, 4.0]); + + // Verify applier result + let result = applier.into_result(); + assert_eq!(result.applied.len(), 1); + assert_eq!(result.errors.len(), 0); + } + + /// Test that F32 dtype is preserved when loading (verifies we didn't break F32 handling) + #[test] + fn dtype_preservation_f32() { + let device = Default::default(); + + // Create TensorData with F32 dtype + let f32_data = TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0], [2, 2]); + assert_eq!(f32_data.dtype, DType::F32); + + // Create a snapshot with F32 data + let snapshot = crate::TensorSnapshot::from_data( + f32_data.clone(), + vec!["weight".to_string()], + vec![], + ParamId::new(), + ); + assert_eq!(snapshot.dtype, DType::F32); + + // Create applier with the F32 snapshot + let mut applier = Applier::new(vec![snapshot], None, None, false); + + // Create target parameter + let target = Param::initialized(ParamId::new(), Tensor::<2>::zeros([2, 2], &device)); + + // Apply the snapshot + applier.enter_module("weight", ""); + let loaded = applier.map_float(target); + applier.exit_module("weight", ""); + + // Verify dtype is F32 + assert_eq!(loaded.val().dtype(), DType::F32); + + // Verify data values + let loaded_data = loaded.val().to_data().to_vec::().unwrap(); + assert_eq!(loaded_data, vec![1.0, 2.0, 3.0, 4.0]); + } + + /// Test that F16 dtype is correctly preserved in TensorSnapshot. + /// + /// Verifies the snapshot layer preserves the F16 dtype tag through a + /// round-trip; does not materialize the data through a backend. + #[test] + fn dtype_preservation_f16_snapshot() { + use half::f16; + + // Create TensorData with F16 dtype using the half crate + let f16_values: Vec = vec![ + f16::from_f32(1.0), + f16::from_f32(2.0), + f16::from_f32(3.0), + f16::from_f32(4.0), + ]; + let f16_data = TensorData::new(f16_values.clone(), [2, 2]); + assert_eq!( + f16_data.dtype, + DType::F16, + "TensorData should have F16 dtype" + ); + + // Create a snapshot with F16 data + let snapshot = crate::TensorSnapshot::from_data( + f16_data.clone(), + vec!["weight".to_string()], + vec![], + ParamId::new(), + ); + + // Verify snapshot preserves F16 dtype + assert_eq!( + snapshot.dtype, + DType::F16, + "TensorSnapshot should preserve F16 dtype" + ); + + // Verify the data can be retrieved with correct dtype + let retrieved_data = snapshot.to_data().expect("Should be able to retrieve data"); + assert_eq!( + retrieved_data.dtype, + DType::F16, + "Retrieved data should have F16 dtype" + ); + + // Verify the actual values are preserved + let retrieved_values: Vec = retrieved_data + .to_vec() + .expect("Should be able to convert to f16 vec"); + assert_eq!( + retrieved_values, f16_values, + "F16 values should be preserved" + ); + + // Note: To fully test F16 tensor creation, you would need a backend + // that supports F16 (like CUDA or WebGPU). The applier fix ensures + // that `Tensor::from_data(data, (device, snapshot.dtype))` is + // called with DType::F16, which will correctly create an F16 tensor + // on backends that support it. + } + + /// Test that BF16 dtype is correctly preserved in TensorSnapshot. + #[test] + fn dtype_preservation_bf16_snapshot() { + use half::bf16; + + // Create TensorData with BF16 dtype + let bf16_values: Vec = vec![ + bf16::from_f32(1.0), + bf16::from_f32(2.0), + bf16::from_f32(3.0), + bf16::from_f32(4.0), + ]; + let bf16_data = TensorData::new(bf16_values.clone(), [2, 2]); + assert_eq!( + bf16_data.dtype, + DType::BF16, + "TensorData should have BF16 dtype" + ); + + // Create a snapshot with BF16 data + let snapshot = crate::TensorSnapshot::from_data( + bf16_data.clone(), + vec!["weight".to_string()], + vec![], + ParamId::new(), + ); + + // Verify snapshot preserves BF16 dtype + assert_eq!( + snapshot.dtype, + DType::BF16, + "TensorSnapshot should preserve BF16 dtype" + ); + + // Verify the data can be retrieved with correct dtype + let retrieved_data = snapshot.to_data().expect("Should be able to retrieve data"); + assert_eq!( + retrieved_data.dtype, + DType::BF16, + "Retrieved data should have BF16 dtype" + ); + + // Verify the actual values are preserved + let retrieved_values: Vec = retrieved_data + .to_vec() + .expect("Should be able to convert to bf16 vec"); + assert_eq!( + retrieved_values, bf16_values, + "BF16 values should be preserved" + ); + } +} diff --git a/crates/burn-store/src/apply_result.rs b/crates/burn-store/src/apply_result.rs new file mode 100644 index 0000000..310a4e7 --- /dev/null +++ b/crates/burn-store/src/apply_result.rs @@ -0,0 +1,300 @@ +//! Result types and diagnostics for tensor application operations + +use alloc::string::String; +use alloc::vec; +use alloc::vec::Vec; + +use burn_core::tensor::{DType, Shape}; + +/// Error types that can occur during tensor application +#[derive(Debug, Clone)] +pub enum ApplyError { + /// Shape mismatch between expected and actual tensor + ShapeMismatch { + /// Path of the tensor + path: String, + /// Expected shape + expected: Shape, + /// Found shape + found: Shape, + }, + /// Data type mismatch between expected and actual tensor + DTypeMismatch { + /// Path of the tensor + path: String, + /// Expected data type + expected: DType, + /// Found data type + found: DType, + }, + /// Error from adapter transformation + AdapterError { + /// Path of the tensor + path: String, + /// Error message + message: String, + }, + /// Error loading tensor data + LoadError { + /// Path of the tensor + path: String, + /// Error message + message: String, + }, +} + +impl core::fmt::Display for ApplyError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + Self::ShapeMismatch { + path, + expected, + found, + } => { + write!( + f, + "Shape mismatch for '{}': expected {:?}, found {:?}", + path, expected, found + ) + } + Self::DTypeMismatch { + path, + expected, + found, + } => { + write!( + f, + "DType mismatch for '{}': expected {:?}, found {:?}", + path, expected, found + ) + } + Self::AdapterError { path, message } => { + write!(f, "Adapter error for '{}': {}", path, message) + } + Self::LoadError { path, message } => { + write!(f, "Load error for '{}': {}", path, message) + } + } + } +} + +impl core::error::Error for ApplyError {} + +/// Result of applying tensor snapshots to a module +#[derive(Clone)] +pub struct ApplyResult { + /// Successfully applied tensor paths + pub applied: Vec, + /// Skipped tensor paths (due to filter) + pub skipped: Vec, + /// Missing tensor paths with their container stacks in dot notation (path, container_stack) + /// Container stack shows the hierarchy: "Struct:Model.Struct:Linear" or "Struct:Model.Enum:ConvType.Struct:Linear" + pub missing: Vec<(String, String)>, + /// Unused tensor paths (in snapshots but not in module) + pub unused: Vec, + /// Errors encountered during application + pub errors: Vec, +} + +impl ApplyResult { + /// Try to strip enum variant from a path + /// e.g., "field.BaseConv.weight" -> "field.weight" + fn strip_enum_variant(path: &str) -> Option { + let segments: Vec<&str> = path.split('.').collect(); + + // Find segments that look like enum variants (CamelCase in middle of path) + let variant_indices: Vec = segments + .iter() + .enumerate() + .filter(|(i, segment)| { + *i > 0 && *i < segments.len() - 1 // Not first or last + && !segment.is_empty() + && segment.chars().next().map(|c| c.is_uppercase()).unwrap_or(false) + && segment.len() > 1 + && segment.chars().skip(1).any(|c| c.is_lowercase()) + }) + .map(|(i, _)| i) + .collect(); + + if variant_indices.is_empty() { + return None; + } + + // Remove the first found variant and return the modified path + let mut result_segments = segments.clone(); + result_segments.remove(variant_indices[0]); + Some(result_segments.join(".")) + } + + /// Find similar paths for a given missing path (for "Did you mean?" suggestions) + fn find_similar_paths(&self, missing_path: &str, max_suggestions: usize) -> Vec { + // First, try exact match with enum variant stripped + if let Some(stripped) = Self::strip_enum_variant(missing_path) + && self.unused.contains(&stripped) + { + return vec![stripped]; + } + + // Fall back to Jaro similarity (used by Elixir for "did you mean?" suggestions) + // Jaro gives higher weight to matching prefixes, ideal for hierarchical tensor paths + let mut similarities: Vec<(String, f64)> = self + .unused + .iter() + .map(|available| { + let similarity = textdistance::nstr::jaro(missing_path, available); + (available.clone(), similarity) + }) + .collect(); + + // Sort by similarity (higher = more similar) + similarities + .sort_by(|(_, a), (_, b)| b.partial_cmp(a).unwrap_or(core::cmp::Ordering::Equal)); + + // Only suggest paths with >= 70% similarity + const SIMILARITY_THRESHOLD: f64 = 0.7; + similarities + .into_iter() + .filter(|(_, sim)| *sim >= SIMILARITY_THRESHOLD) + .take(max_suggestions) + .map(|(path, _)| path) + .collect() + } +} + +impl ApplyResult { + /// Check if the apply operation was successful (no errors) + /// Note: Missing tensors are not considered errors when allow_partial is true + pub fn is_success(&self) -> bool { + self.errors.is_empty() + } +} + +impl core::fmt::Debug for ApplyResult { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + // Delegate to Display for comprehensive output + core::fmt::Display::fmt(self, f) + } +} + +impl core::fmt::Display for ApplyResult { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + writeln!(f, "┌─ Tensor Loading Summary ─────────────────────────")?; + writeln!(f, "│")?; + writeln!( + f, + "│ ✓ Successfully applied: {} tensors", + self.applied.len() + )?; + writeln!(f, "│ ⊘ Skipped (filtered): {} tensors", self.skipped.len())?; + writeln!( + f, + "│ ✗ Missing in source: {} tensors", + self.missing.len() + )?; + writeln!(f, "│ ? Unused in target: {} tensors", self.unused.len())?; + writeln!(f, "│ ! Errors: {} errors", self.errors.len())?; + + if !self.missing.is_empty() { + writeln!(f, "│")?; + writeln!( + f, + "├─ Missing Tensors (requested by model but not found in source)" + )?; + writeln!(f, "│")?; + + // Use actual container stack data to detect enum variants + // Count how many missing paths have "Enum:" in their container stack + let enum_variant_missing: Vec<_> = self + .missing + .iter() + .filter(|(_, stack)| stack.contains("Enum:")) + .collect(); + + if !enum_variant_missing.is_empty() { + writeln!( + f, + "│ ⚠️ {} paths contain enum variants (detected from container stack)", + enum_variant_missing.len() + )?; + writeln!( + f, + "│ Burn includes enum variant names in paths, but PyTorch doesn't." + )?; + writeln!( + f, + "│ Example: Burn has 'field.BaseConv.weight', PyTorch has 'field.weight'" + )?; + writeln!(f, "│")?; + writeln!( + f, + "│ 💡 Solution 1: Enable skip_enum_variants flag (simplest):" + )?; + writeln!(f, "│")?; + writeln!( + f, + "│ let mut store = PytorchStore::from_file(\"model.pth\")" + )?; + writeln!(f, "│ .skip_enum_variants(true); // ← Add this")?; + writeln!(f, "│")?; + writeln!( + f, + "│ 💡 Solution 2: Remap enum keys in source (most precise):" + )?; + writeln!(f, "│")?; + writeln!( + f, + "│ let mut store = SafetensorsStore::from_file(\"model.safetensors\")" + )?; + writeln!( + f, + "│ .with_key_remapping(r\"field\\.(\\w+)\", \"field.BaseConv.$1\");" + )?; + writeln!(f, "│")?; + } + + writeln!(f, "│ First 10 missing tensors:")?; + for (path, _) in self.missing.iter().take(10) { + writeln!(f, "│ • {}", path)?; + + // Show "Did you mean?" suggestions for this path + let suggestions = self.find_similar_paths(path, 1); + if !suggestions.is_empty() { + writeln!(f, "│ Did you mean: '{}'?", suggestions[0])?; + } + } + if self.missing.len() > 10 { + writeln!(f, "│ ... and {} more", self.missing.len() - 10)?; + } + } + + if !self.unused.is_empty() { + writeln!(f, "│")?; + writeln!(f, "├─ Unused Tensors (in source but not used by model)")?; + writeln!(f, "│")?; + writeln!(f, "│ First 10 unused tensors:")?; + for path in self.unused.iter().take(10) { + writeln!(f, "│ • {}", path)?; + } + if self.unused.len() > 10 { + writeln!(f, "│ ... and {} more", self.unused.len() - 10)?; + } + } + + if !self.errors.is_empty() { + writeln!(f, "│")?; + writeln!(f, "├─ Errors")?; + writeln!(f, "│")?; + for error in self.errors.iter().take(10) { + writeln!(f, "│ ⚠️ {}", error)?; + } + if self.errors.len() > 10 { + writeln!(f, "│ ... and {} more", self.errors.len() - 10)?; + } + } + + writeln!(f, "│")?; + write!(f, "└───────────────────────────────────────────────────")?; + + Ok(()) + } +} diff --git a/crates/burn-store/src/bridge.rs b/crates/burn-store/src/bridge.rs new file mode 100644 index 0000000..2184c54 --- /dev/null +++ b/crates/burn-store/src/bridge.rs @@ -0,0 +1,51 @@ +//! Bridge between [`TensorSnapshot`] (burn-core) and [`burn_pack::Tensor`] +//! (the tensor-agnostic burnpack format entry), used by [`BurnpackStore`](crate::BurnpackStore). + +use alloc::format; +use alloc::rc::Rc; +use alloc::string::{String, ToString}; +use alloc::vec; +use alloc::vec::Vec; + +use burn_pack::{Error as PackError, Tensor as PackTensor}; + +use super::TensorSnapshot; +use burn_core::module::ParamId; +use burn_core::tensor::TensorData; + +/// Convert a [`TensorSnapshot`] into a [`PackTensor`] entry, materializing its data. +pub fn snapshot_to_tensor(snapshot: &TensorSnapshot) -> Result { + let data = snapshot + .to_data() + .map_err(|e| PackError::IoError(format!("{e:?}")))?; + Ok(PackTensor::new( + snapshot.full_path(), + snapshot.dtype, + snapshot.shape.clone(), + snapshot.tensor_id.map(|id| id.val()), + data.bytes, + )) +} + +/// Convert a [`PackTensor`] entry into a lazy [`TensorSnapshot`]. +/// +/// The tensor's [`Bytes`](burn_pack::Bytes) may be file-backed (from [`Reader::from_file`](burn_pack::Reader::from_file)), +/// in which case the data is only read from disk when the snapshot is materialized. +pub fn tensor_to_snapshot(tensor: PackTensor) -> TensorSnapshot { + let dtype = tensor.dtype; + let shape = tensor.shape.clone(); + let path_stack: Vec = tensor.name.split('.').map(|s| s.to_string()).collect(); + let tensor_id = tensor.param_id.map(ParamId::from).unwrap_or_default(); + + let bytes = tensor.bytes; + let shape_for_closure = shape.clone(); + let data_fn = Rc::new(move || { + Ok(TensorData::from_bytes( + bytes.clone(), + shape_for_closure.clone(), + dtype, + )) + }); + + TensorSnapshot::from_closure(data_fn, dtype, shape, path_stack, vec![], tensor_id) +} diff --git a/crates/burn-store/src/burnpack.rs b/crates/burn-store/src/burnpack.rs new file mode 100644 index 0000000..d85f50f --- /dev/null +++ b/crates/burn-store/src/burnpack.rs @@ -0,0 +1,512 @@ +#[cfg(feature = "std")] +use std::path::PathBuf; + +#[cfg(feature = "std")] +use crate::KeyRemapper; +use crate::bridge; +use crate::{ + IdentityAdapter, ModuleAdapter, ModuleSnapshot, ModuleStore, PathFilter, TensorSnapshot, +}; +use alloc::boxed::Box; +use alloc::collections::BTreeMap; +use alloc::format; +use alloc::string::String; +use alloc::vec::Vec; +use burn_core::tensor::Bytes; +use burn_pack::{Error as PackError, Reader, Tensor as PackTensor, Writer}; + +/// Store mode for BurnpackStore +enum StoreMode { + #[cfg(feature = "std")] + File(PathBuf), + Bytes(Option), +} + +/// BurnpackStore - A Burn-specific file format store using CBOR for metadata +pub struct BurnpackStore { + /// Store mode - either file path or bytes + mode: StoreMode, + /// Optional filter for selective loading/saving + filter: Option, + /// Additional metadata + metadata: BTreeMap, + /// Allow partial loading (ignore missing tensors) + allow_partial: bool, + /// Validate tensors during loading (check shapes and dtypes) + validate: bool, + /// Allow overwriting existing files (default: false) + overwrite: bool, + /// Automatically append .bpk extension if not present (default: true) + #[cfg(feature = "std")] + auto_extension: bool, + /// Key remapper for tensor name transformations + #[cfg(feature = "std")] + remapper: KeyRemapper, + /// Adapter applied when loading (source -> Burn) + from_adapter: Box, + /// Adapter applied when saving (Burn -> target) + to_adapter: Box, + /// Reader for loading + reader: Option, + /// Cached tensor snapshots (parsed once, reused) + snapshots_cache: Option>, +} + +impl BurnpackStore { + /// Get the default metadata that includes Burn framework information. + /// + /// This includes: + /// - `format`: "burnpack" + /// - `producer`: "burn" + /// - `version`: The version of burn-store crate (from CARGO_PKG_VERSION) + /// + /// These metadata fields are automatically added to all saved models. + pub fn default_metadata() -> BTreeMap { + let mut metadata = BTreeMap::new(); + metadata.insert("format".into(), "burnpack".into()); + metadata.insert("producer".into(), "burn".into()); + metadata.insert("version".into(), env!("CARGO_PKG_VERSION").into()); + metadata + } + /// Create a new store from a file path + /// + /// By default, automatically appends `.bpk` extension if the path doesn't have one. + /// Use `.auto_extension(false)` to disable this behavior. + /// + /// # Examples + /// + /// ```no_run + /// # use burn_store::BurnpackStore; + /// // Automatically appends .bpk + /// let store = BurnpackStore::from_file("model"); // creates "model.bpk" + /// + /// // Already has extension, no append + /// let store = BurnpackStore::from_file("model.bpk"); // uses "model.bpk" + /// let store = BurnpackStore::from_file("model.myext"); // uses "model.myext" + /// + /// // Disable auto-extension + /// let store = BurnpackStore::from_file("model").auto_extension(false); // uses "model" + /// ``` + #[cfg(feature = "std")] + pub fn from_file>(path: P) -> Self { + Self { + mode: StoreMode::File(path.as_ref().to_path_buf()), + filter: None, + metadata: Self::default_metadata(), + allow_partial: false, + validate: true, + overwrite: false, + #[cfg(feature = "std")] + auto_extension: true, + #[cfg(feature = "std")] + remapper: KeyRemapper::new(), + from_adapter: Box::new(IdentityAdapter), + to_adapter: Box::new(IdentityAdapter), + reader: None, + snapshots_cache: None, + } + } + + /// Create a new store from bytes (for reading) or empty (for writing) + pub fn from_bytes(bytes: Option) -> Self { + Self { + mode: StoreMode::Bytes(bytes), + filter: None, + metadata: Self::default_metadata(), + allow_partial: false, + validate: true, + overwrite: false, + #[cfg(feature = "std")] + auto_extension: false, // Not used for bytes mode + #[cfg(feature = "std")] + remapper: KeyRemapper::new(), + from_adapter: Box::new(IdentityAdapter), + to_adapter: Box::new(IdentityAdapter), + reader: None, + snapshots_cache: None, + } + } + + /// Create a new store from static bytes with zero-copy loading enabled. + /// + /// This is optimized for embedded model weights where the data lives in the + /// binary's `.rodata` section. Tensor data is sliced without copying, keeping + /// the static reference alive. + /// + /// # Example + /// + /// ```ignore + /// static MODEL_DATA: &[u8] = include_bytes!("model.bpk"); + /// let store = BurnpackStore::from_static(MODEL_DATA); + /// ``` + pub fn from_static(data: &'static [u8]) -> Self { + use burn_core::tensor::AllocationProperty; + + // Create bytes::Bytes from static data (zero-copy, stays in .rodata) + let shared = bytes::Bytes::from_static(data); + + // Wrap in cubecl Bytes with shared-bytes allocation controller + let bytes = Bytes::from_shared(shared, AllocationProperty::Other); + + Self { + mode: StoreMode::Bytes(Some(bytes)), + filter: None, + metadata: Self::default_metadata(), + allow_partial: false, + validate: true, + overwrite: false, + #[cfg(feature = "std")] + auto_extension: false, + #[cfg(feature = "std")] + remapper: KeyRemapper::new(), + from_adapter: Box::new(IdentityAdapter), + to_adapter: Box::new(IdentityAdapter), + reader: None, + snapshots_cache: None, + } + } + + /// Add metadata key-value pair + pub fn metadata(mut self, key: impl Into, value: impl Into) -> Self { + self.metadata.insert(key.into(), value.into()); + self + } + + /// Clear all metadata (including defaults) + /// + /// This removes all metadata including the default format, producer, and version fields. + /// Use with caution as some tools may expect these fields to be present. + pub fn clear_metadata(mut self) -> Self { + self.metadata.clear(); + self + } + + /// Allow partial loading (ignore missing tensors) + /// + /// When set to `true`, the store will not fail if some tensors are missing + /// during loading. This is useful when loading a subset of a model's parameters. + /// + /// Default: `false` + pub fn allow_partial(mut self, allow: bool) -> Self { + self.allow_partial = allow; + self + } + + /// Enable or disable validation during loading + /// + /// When validation is enabled, the store will check that loaded tensors + /// match the expected shapes and data types. Disabling validation can + /// improve performance but may lead to runtime errors if data is corrupted. + /// + /// Default: `true` + pub fn validate(mut self, validate: bool) -> Self { + self.validate = validate; + self + } + + /// Allow overwriting existing files when saving + /// + /// When set to `false`, attempting to save to an existing file will result in an error. + /// When set to `true`, existing files will be overwritten without warning. + /// + /// Default: `false` + pub fn overwrite(mut self, overwrite: bool) -> Self { + self.overwrite = overwrite; + self + } + + /// Enable or disable automatic .bpk extension appending + /// + /// When enabled (default), automatically appends `.bpk` to the file path + /// if no extension is detected. If an extension is already present, it is preserved. + /// + /// When disabled, uses the exact path provided without modification. + /// + /// Default: `true` + /// + /// # Examples + /// + /// ```no_run + /// # use burn_store::BurnpackStore; + /// // With auto_extension enabled (default) + /// let store = BurnpackStore::from_file("model"); // -> "model.bpk" + /// + /// // With auto_extension disabled + /// let store = BurnpackStore::from_file("model") + /// .auto_extension(false); // -> "model" + /// ``` + #[cfg(feature = "std")] + pub fn auto_extension(mut self, enable: bool) -> Self { + self.auto_extension = enable; + self + } + + /// Set the adapter for loading tensors (converting from source format to Burn). + pub fn with_from_adapter(mut self, adapter: impl ModuleAdapter + 'static) -> Self { + self.from_adapter = Box::new(adapter); + self + } + + /// Set the adapter for saving tensors (converting from Burn to target format). + pub fn with_to_adapter(mut self, adapter: impl ModuleAdapter + 'static) -> Self { + self.to_adapter = Box::new(adapter); + self + } + + /// Set path filter for selective loading/saving + pub fn with_filter(mut self, filter: PathFilter) -> Self { + self.filter = Some(filter); + self + } + + /// Add regex pattern to filter + #[cfg(feature = "std")] + pub fn with_regex(mut self, pattern: &str) -> Self { + let filter = self.filter.unwrap_or_default(); + self.filter = Some(filter.with_regex(pattern)); + self + } + + /// Add exact path to filter + pub fn with_full_path(mut self, path: impl Into) -> Self { + let filter = self.filter.unwrap_or_default(); + self.filter = Some(filter.with_full_path(path)); + self + } + + /// Match all tensors (no filtering) + pub fn match_all(mut self) -> Self { + self.filter = Some(PathFilter::new().match_all()); + self + } + + /// Set key remapper for tensor name transformations during loading + #[cfg(feature = "std")] + pub fn remap(mut self, remapper: KeyRemapper) -> Self { + self.remapper = remapper; + self + } + + /// Add a single regex pattern for key remapping + #[cfg(feature = "std")] + pub fn with_remap_pattern(mut self, from: S1, to: S2) -> Self + where + S1: AsRef, + S2: Into, + { + self.remapper = self + .remapper + .add_pattern(from.as_ref(), to.into()) + .expect("Invalid regex pattern"); + self + } + + /// Set the path filter + pub fn filter(mut self, filter: PathFilter) -> Self { + self.filter = Some(filter); + self + } + + /// Get the bytes after writing (only valid for bytes mode after collecting) + pub fn get_bytes(&self) -> Result { + // `collect_from` caches the written bytes back into `self.mode` for bytes mode. + match &self.mode { + StoreMode::Bytes(Some(bytes)) => Ok(bytes.clone()), + _ => Err(PackError::IoError("No bytes available".into())), + } + } + + /// Process the file path with auto-extension logic + #[cfg(feature = "std")] + fn process_path(&self, path: &std::path::Path) -> PathBuf { + if !self.auto_extension { + return path.to_path_buf(); + } + + // Check if path already has an extension + if path.extension().is_some() { + // Has extension, use as-is + return path.to_path_buf(); + } + + // No extension, append .bpk + let mut new_path = path.to_path_buf(); + new_path.set_extension("bpk"); + new_path + } + + /// Ensure the reader is initialized, loading from storage if needed + fn ensure_reader(&mut self) -> Result<&Reader, PackError> { + if self.reader.is_none() { + let reader = match &self.mode { + #[cfg(feature = "std")] + StoreMode::File(path) => { + let final_path = self.process_path(path); + Reader::from_file(&final_path)? + } + StoreMode::Bytes(Some(bytes)) => Reader::from_bytes(bytes.clone())?, + StoreMode::Bytes(None) => { + return Err(PackError::IoError("No bytes to read from".into())); + } + }; + self.reader = Some(reader); + } + + self.reader + .as_ref() + .ok_or_else(|| PackError::IoError("Reader not initialized".into())) + } +} + +impl ModuleStore for BurnpackStore { + type Error = PackError; + + fn collect_from(&mut self, module: &M) -> Result<(), Self::Error> { + // Invalidate cache since we're writing new data + self.snapshots_cache = None; + self.reader = None; + + // Collect snapshots from module with adapter + let snapshots = module.collect(self.filter.clone(), Some(self.to_adapter.clone()), false); + + // Bridge snapshots to tensor-agnostic burnpack entries (materializing their data) + let tensors: Vec = snapshots + .iter() + .map(bridge::snapshot_to_tensor) + .collect::>()?; + + // Initialize writer with tensors + let mut writer = Writer::new(tensors); + + // Add metadata using builder pattern + for (key, value) in &self.metadata { + writer = writer.with_metadata(key.as_str(), value.as_str()); + } + + // Write to storage based on mode, consuming the writer. For bytes mode the generated + // bytes are cached back into `self.mode` so `get_bytes` can return them afterwards. + match &self.mode { + #[cfg(feature = "std")] + StoreMode::File(path) => { + // Process path with auto-extension logic + let final_path = self.process_path(path); + + // Check if file exists and overwrite is disabled + if final_path.exists() && !self.overwrite { + return Err(PackError::IoError(format!( + "File already exists: {}. Use .overwrite(true) to overwrite.", + final_path.display() + ))); + } + writer.write_to_file(&final_path)?; + } + StoreMode::Bytes(_) => { + // Generate and store the bytes + let bytes_data = writer.into_bytes()?; + // Update mode with bytes - this pattern is irrefutable in no-std mode + #[cfg_attr(not(feature = "std"), allow(irrefutable_let_patterns))] + let StoreMode::Bytes(bytes_ref) = &mut self.mode else { + unreachable!("We just matched Bytes variant"); + }; + *bytes_ref = Some(bytes_data); + } + } + + Ok(()) + } + + fn apply_to( + &mut self, + module: &mut M, + ) -> Result { + // Get all snapshots using the cached method + let snapshots: Vec = self.get_all_snapshots()?.values().cloned().collect(); + + // Apply all snapshots at once to the module + // Burnpack is Burn's native format, so no enum variant skipping needed + // Filter is applied here during apply, not during cache population + let result = module.apply( + snapshots, + self.filter.clone(), + Some(self.from_adapter.clone()), + false, + ); + + // Validate if needed + if self.validate && !result.errors.is_empty() { + return Err(PackError::ValidationError(format!( + "Import errors: {:?}", + result.errors + ))); + } + + // Check for missing tensors if partial loading is not allowed + if !self.allow_partial && !result.missing.is_empty() { + return Err(PackError::ValidationError(format!( + "Missing tensors: {:?}", + result.missing + ))); + } + + Ok(result) + } + + fn get_snapshot(&mut self, name: &str) -> Result, Self::Error> { + // Ensure cache is populated + self.ensure_snapshots_cache()?; + Ok(self.snapshots_cache.as_ref().unwrap().get(name)) + } + + fn get_all_snapshots(&mut self) -> Result<&BTreeMap, Self::Error> { + // Ensure cache is populated + self.ensure_snapshots_cache()?; + Ok(self.snapshots_cache.as_ref().unwrap()) + } + + fn keys(&mut self) -> Result, Self::Error> { + // Always use the cache to ensure remapping is applied consistently + Ok(self.get_all_snapshots()?.keys().cloned().collect()) + } +} + +impl BurnpackStore { + /// Ensure the snapshots cache is populated + fn ensure_snapshots_cache(&mut self) -> Result<(), PackError> { + if self.snapshots_cache.is_some() { + return Ok(()); + } + + // Ensure reader is loaded + self.ensure_reader()?; + + // Consume the reader, bridging its tensors to snapshots. File-backed readers keep the + // tensor data lazy (read on materialization); in-memory shared sources are zero-copy. + // Taking the reader hands the source's ownership to the tensor views, so nothing is read + // or copied eagerly. The snapshots cache below is what's reused on later calls. + let reader = self + .reader + .take() + .expect("reader initialized by ensure_reader"); + let snapshots: Vec = reader + .into_tensors()? + .into_iter() + .map(bridge::tensor_to_snapshot) + .collect(); + + // Apply remapping if configured (but NOT filtering - that's done at apply time) + #[cfg(feature = "std")] + let snapshots = if !self.remapper.patterns.is_empty() { + let (remapped, _remapped_names) = self.remapper.remap(snapshots); + remapped + } else { + snapshots + }; + + // Build the cache as BTreeMap + let cache: BTreeMap = + snapshots.into_iter().map(|s| (s.full_path(), s)).collect(); + + self.snapshots_cache = Some(cache); + Ok(()) + } +} diff --git a/crates/burn-store/src/collector.rs b/crates/burn-store/src/collector.rs new file mode 100644 index 0000000..db2c74b --- /dev/null +++ b/crates/burn-store/src/collector.rs @@ -0,0 +1,1137 @@ +use alloc::boxed::Box; +use alloc::string::{String, ToString}; +use alloc::vec::Vec; + +use burn_core::tensor::{Bool, Int, Tensor}; + +use crate::{ModuleAdapter, PathFilter, TensorSnapshot}; +use burn_core::module::{ModuleVisitor, Param, ParamId}; + +/// Collects tensor views from modules without copying data. +/// +/// This collector traverses a module hierarchy and creates lightweight views +/// of tensors that can be materialized to `TensorData` on demand. +/// +/// # Examples +/// +/// ## Collect all tensors +/// ```rust,no_run +/// # use burn_store::Collector; +/// let collector = Collector::new(None, None, false); +/// // Use with module.visit(&mut collector); +/// let all_tensors = collector.tensors; +/// ``` +/// +/// ## Filter with single pattern +/// ```rust,no_run +/// # use burn_store::{Collector, PathFilter}; +/// let filter = PathFilter::new().with_regex(r"^encoder\..*"); +/// let collector = Collector::new(Some(filter), None, false); +/// // Use with module.visit(&mut collector); +/// // Only collects tensors starting with "encoder." +/// ``` +/// +/// ## Filter with multiple patterns (OR union) +/// ```rust,no_run +/// # use burn_store::{Collector, PathFilter}; +/// let filter = PathFilter::new() +/// .with_regex(r"^encoder\..*") // Match all encoder tensors +/// .with_regex(r".*\.bias$"); // OR match any bias tensors +/// let collector = Collector::new(Some(filter), None, false); +/// // Use with module.visit(&mut collector); +/// // Collects tensors matching ANY of the patterns +/// ``` +pub struct Collector { + /// Collection of tensor views + pub tensors: Vec, + path_stack: Vec, + container_stack: Vec, + filter: Option, + adapter: Option>, + /// Skip enum variant names when building paths + /// When true, enum variant names are not included in tensor paths + skip_enum_variants: bool, +} + +impl Default for Collector { + fn default() -> Self { + Self::new(None, None, false) + } +} + +impl Collector { + /// Create a new tensor view collector with an optional filter and adapter. + /// + /// # Arguments + /// + /// * `filter` - An optional [`PathFilter`] to determine which tensors to collect. + /// When `None`, all tensors are collected. + /// * `adapter` - Optional adapter to transform tensors based on container types. + /// Applied to all collected tensors before returning. + /// * `skip_enum_variants` - Skip enum variant names when building paths. + /// When true, paths will not include enum variant names (e.g., "feature.weight" + /// instead of "feature.BaseConv.weight"). Useful when exporting to formats + /// like PyTorch that don't use enum variants. + /// + /// # Examples + /// + /// ```rust,no_run + /// # use burn_store::{Collector, PathFilter}; + /// // Collect all tensors without adapter + /// let collector = Collector::new(None, None, false); + /// + /// // Use PathFilter builder + /// let filter = PathFilter::new() + /// .with_regex(r"^encoder\..*") + /// .with_full_path("decoder.weight"); + /// let collector = Collector::new(Some(filter), None, false); + /// + /// // Skip enum variants for PyTorch export + /// let collector = Collector::new(None, None, true); + /// ``` + pub fn new( + filter: Option, + adapter: Option>, + skip_enum_variants: bool, + ) -> Self { + Self { + tensors: Vec::new(), + path_stack: Vec::new(), + container_stack: Vec::new(), + filter, + adapter, + skip_enum_variants, + } + } + + /// Apply the adapter to collected tensors and return the result. + pub fn into_tensors(self) -> Vec { + if let Some(adapter) = self.adapter { + self.tensors + .into_iter() + .map(|snapshot| adapter.adapt(&snapshot)) + .collect() + } else { + self.tensors + } + } + + fn should_collect(&self, path: &[String], container_stack: &[String]) -> bool { + // If filter is present, use it; otherwise collect all + match &self.filter { + None => true, + Some(f) => f.matches_with_container_path(path, container_stack), + } + } +} + +impl ModuleVisitor for Collector { + fn enter_module(&mut self, name: &str, container_type: &str) { + // Always track the container type for proper filtering and module type detection + self.container_stack.push(container_type.to_string()); + + // Only add to path if it's not an enum variant (when skip_enum_variants is enabled) + // This ensures paths are built without enum variant names from the start + if !self.skip_enum_variants || !container_type.starts_with("Enum:") { + self.path_stack.push(name.to_string()); + } + } + + fn exit_module(&mut self, _name: &str, container_type: &str) { + self.container_stack.pop(); + + // Only pop from path if we added it (not an enum variant when skip_enum_variants is enabled) + if !self.skip_enum_variants || !container_type.starts_with("Enum:") { + self.path_stack.pop(); + } + } + + fn visit_float(&mut self, param: &Param>) { + if self.should_collect(&self.path_stack, &self.container_stack) { + self.tensors.push(TensorSnapshot::from_float( + ¶m.transform_for_save().val(), + self.path_stack.clone(), + self.container_stack.clone(), + param.id, + )); + } + } + + fn visit_int(&mut self, param: &Param>) { + if self.should_collect(&self.path_stack, &self.container_stack) { + self.tensors.push(TensorSnapshot::from_int( + ¶m.transform_for_save().val(), + self.path_stack.clone(), + self.container_stack.clone(), + param.id, + )); + } + } + + fn visit_bool(&mut self, param: &Param>) { + if self.should_collect(&self.path_stack, &self.container_stack) { + self.tensors.push(TensorSnapshot::from_bool( + ¶m.transform_for_save().val(), + self.path_stack.clone(), + self.container_stack.clone(), + param.id, + )); + } + } + + fn visit_float_with_path( + &mut self, + path: &[String], + id: ParamId, + tensor: &Tensor, + ) { + // For path-based visits, we use the current container stack for filtering + if self.should_collect(path, &self.container_stack) { + self.tensors.push(TensorSnapshot::from_float( + tensor, + path.to_vec(), + self.container_stack.clone(), + id, + )); + } + } + + fn visit_int_with_path( + &mut self, + path: &[String], + id: ParamId, + tensor: &Tensor, + ) { + if self.should_collect(path, &self.container_stack) { + self.tensors.push(TensorSnapshot::from_int( + tensor, + path.to_vec(), + self.container_stack.clone(), + id, + )); + } + } + + fn visit_bool_with_path( + &mut self, + path: &[String], + id: ParamId, + tensor: &Tensor, + ) { + if self.should_collect(path, &self.container_stack) { + self.tensors.push(TensorSnapshot::from_bool( + tensor, + path.to_vec(), + self.container_stack.clone(), + id, + )); + } + } +} + +#[cfg(all(test, feature = "std"))] +mod tests { + use super::*; + + use burn_core as burn; + + use alloc::collections::BTreeMap; + use alloc::string::String; + use burn_core::module::{Module, Param}; + use burn_core::tensor::{Device, shape}; + use burn_nn::LinearConfig; + + #[test] + fn tensor_snapshot_collector() { + let device = Default::default(); + let tensor = Tensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + + let mut collector = Collector::new(None, None, false); + let id = ParamId::new(); + + // Collect a tensor + collector.visit_float_with_path(&["model".to_string(), "weight".to_string()], id, &tensor); + + assert_eq!(collector.tensors.len(), 1); + assert_eq!(collector.tensors[0].full_path(), "model.weight"); + + // Verify the tensor can be converted to data + let view = &collector.tensors[0]; + let data = view.to_data().unwrap(); + assert_eq!(data.shape, shape![2, 2]); + } + + #[test] + fn root_level_parameters() { + use burn_core::module::ModuleVisitor; + + let device = Default::default(); + + // Create root-level parameters (single-element path, not nested in modules) + let weight = Param::>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + let bias = Param::>::from_data([5.0, 6.0], &device); + + let mut collector = Collector::new(None, None, false); + + // Simulate module traversal for root-level parameters + // Enter "weight" path (as if we're visiting a field named "weight") + ModuleVisitor::enter_module(&mut collector, "weight", ""); + ModuleVisitor::visit_float(&mut collector, &weight); + ModuleVisitor::exit_module(&mut collector, "weight", ""); + + // Enter "bias" path (as if we're visiting a field named "bias") + ModuleVisitor::enter_module(&mut collector, "bias", ""); + ModuleVisitor::visit_float(&mut collector, &bias); + ModuleVisitor::exit_module(&mut collector, "bias", ""); + + // Verify both parameters were collected + assert_eq!(collector.tensors.len(), 2); + + // Verify paths are correct (single-element paths) + assert_eq!(collector.tensors[0].full_path(), "weight"); + assert_eq!(collector.tensors[1].full_path(), "bias"); + + // Verify data is correct + let weight_data = collector.tensors[0] + .to_data() + .unwrap() + .to_vec::() + .unwrap(); + let bias_data = collector.tensors[1] + .to_data() + .unwrap() + .to_vec::() + .unwrap(); + + assert_eq!(weight_data, vec![1.0, 2.0, 3.0, 4.0]); + assert_eq!(bias_data, vec![5.0, 6.0]); + } + + #[test] + #[cfg(target_has_atomic = "ptr")] + fn tensor_snapshot_collector_with_filter() { + let device = Default::default(); + let tensor = Tensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + + let filter = PathFilter::new().with_regex(r"^encoder\..*"); + let mut collector = Collector::new(Some(filter), None, false); + let id = ParamId::new(); + + // This should be collected + collector.visit_float_with_path( + &["encoder".to_string(), "weight".to_string()], + id, + &tensor, + ); + // This should NOT be collected + collector.visit_float_with_path( + &["decoder".to_string(), "weight".to_string()], + id, + &tensor, + ); + + assert_eq!(collector.tensors.len(), 1); + assert_eq!(collector.tensors[0].full_path(), "encoder.weight"); + } + + #[test] + #[cfg(target_has_atomic = "ptr")] + fn tensor_snapshot_collector_with_multiple_filters() { + let device = Default::default(); + let tensor = Tensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + + // Multiple patterns - collect if matches ANY (OR union) + let filter = PathFilter::new() + .with_regex(r"^encoder\..*") // Match encoder.* + .with_regex(r".*\.bias$"); // Match *.bias + let mut collector = Collector::new(Some(filter), None, false); + let id = ParamId::new(); + + // These should be collected + collector.visit_float_with_path( + &["encoder".to_string(), "weight".to_string()], + id, + &tensor, + ); // matches first pattern + collector.visit_float_with_path(&["decoder".to_string(), "bias".to_string()], id, &tensor); // matches second pattern + collector.visit_float_with_path(&["encoder".to_string(), "bias".to_string()], id, &tensor); // matches both patterns + + // This should NOT be collected + collector.visit_float_with_path( + &["decoder".to_string(), "weight".to_string()], + id, + &tensor, + ); // matches neither + + assert_eq!(collector.tensors.len(), 3); + let paths: Vec = collector.tensors.iter().map(|v| v.full_path()).collect(); + assert!(paths.contains(&"encoder.weight".to_string())); + assert!(paths.contains(&"decoder.bias".to_string())); + assert!(paths.contains(&"encoder.bias".to_string())); + assert!(!paths.contains(&"decoder.weight".to_string())); + } + + #[test] + fn tensor_snapshot_collector_with_predicate() { + let device = Default::default(); + let tensor = Tensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + + // Use predicate function for filtering + fn filter_fn(path: &str, _container_path: &str) -> bool { + path.starts_with("encoder.") || path == "decoder.bias" + } + let filter = PathFilter::new().with_predicate(filter_fn); + let mut collector = Collector::new(Some(filter), None, false); + let id = ParamId::new(); + + // These should be collected + collector.visit_float_with_path( + &["encoder".to_string(), "weight".to_string()], + id, + &tensor, + ); + collector.visit_float_with_path(&["encoder".to_string(), "bias".to_string()], id, &tensor); + collector.visit_float_with_path(&["decoder".to_string(), "bias".to_string()], id, &tensor); + + // This should NOT be collected + collector.visit_float_with_path( + &["decoder".to_string(), "weight".to_string()], + id, + &tensor, + ); + collector.visit_float_with_path(&["other".to_string(), "tensor".to_string()], id, &tensor); + + assert_eq!(collector.tensors.len(), 3); + let paths: Vec = collector.tensors.iter().map(|v| v.full_path()).collect(); + assert!(paths.contains(&"encoder.weight".to_string())); + assert!(paths.contains(&"encoder.bias".to_string())); + assert!(paths.contains(&"decoder.bias".to_string())); + assert!(!paths.contains(&"decoder.weight".to_string())); + assert!(!paths.contains(&"other.tensor".to_string())); + } + + #[test] + fn tensor_snapshot_collector_predicate_with_complex_logic() { + let device = Default::default(); + let tensor = Tensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + + // Complex predicate with multiple conditions + fn complex_filter(path: &str, _container_path: &str) -> bool { + let parts: Vec<&str> = path.split('.').collect(); + if parts.len() != 3 { + return false; + } + // Only collect if it's layer1 or layer2, and it's a weight tensor + (parts[1] == "layer1" || parts[1] == "layer2") && parts[2] == "weight" + } + let filter = PathFilter::new().with_predicate(complex_filter); + let mut collector = Collector::new(Some(filter), None, false); + let id = ParamId::new(); + + // These should be collected + collector.visit_float_with_path( + &[ + "model".to_string(), + "layer1".to_string(), + "weight".to_string(), + ], + id, + &tensor, + ); + collector.visit_float_with_path( + &[ + "model".to_string(), + "layer2".to_string(), + "weight".to_string(), + ], + id, + &tensor, + ); + + // These should NOT be collected + collector.visit_float_with_path( + &[ + "model".to_string(), + "layer1".to_string(), + "bias".to_string(), + ], + id, + &tensor, + ); + collector.visit_float_with_path( + &[ + "model".to_string(), + "layer3".to_string(), + "weight".to_string(), + ], + id, + &tensor, + ); + collector.visit_float_with_path( + &["encoder".to_string(), "weight".to_string()], + id, + &tensor, + ); // wrong structure + + assert_eq!(collector.tensors.len(), 2); + let paths: Vec = collector.tensors.iter().map(|v| v.full_path()).collect(); + assert!(paths.contains(&"model.layer1.weight".to_string())); + assert!(paths.contains(&"model.layer2.weight".to_string())); + assert!(!paths.contains(&"model.layer1.bias".to_string())); + assert!(!paths.contains(&"model.layer3.weight".to_string())); + assert!(!paths.contains(&"encoder.weight".to_string())); + } + + // Test visitor that collects tensor paths + struct TensorPathCollector { + pub paths: BTreeMap)>, + path_stack: Vec, + } + + impl TensorPathCollector { + fn new() -> Self { + Self { + paths: BTreeMap::new(), + path_stack: Vec::new(), + } + } + + fn current_path(&self) -> String { + self.path_stack.join(".") + } + } + + impl ModuleVisitor for TensorPathCollector { + fn enter_module(&mut self, name: &str, _container_type: &str) { + self.path_stack.push(name.to_string()); + } + + fn exit_module(&mut self, _name: &str, _container_type: &str) { + self.path_stack.pop(); + } + + fn visit_float(&mut self, param: &Param>) { + let path = self.current_path(); + if !path.is_empty() { + self.paths.insert( + path, + (param.id, param.transform_for_save().val().shape().to_vec()), + ); + } + } + + fn visit_int(&mut self, param: &Param>) { + let path = self.current_path(); + if !path.is_empty() { + self.paths.insert( + path, + (param.id, param.transform_for_save().val().shape().to_vec()), + ); + } + } + + fn visit_bool(&mut self, param: &Param>) { + let path = self.current_path(); + if !path.is_empty() { + self.paths.insert( + path, + (param.id, param.transform_for_save().val().shape().to_vec()), + ); + } + } + } + + // Simple nested module for testing + #[derive(Module, Debug)] + struct InnerModule { + weight: Param>, + bias: Param>, + } + + #[derive(Module, Debug)] + struct OuterModule { + layer1: InnerModule, + layer2: InnerModule, + } + + impl InnerModule { + fn new(device: &Device) -> Self { + Self { + weight: Param::from_data([[1.0, 2.0], [3.0, 4.0]], device), + bias: Param::from_data([5.0, 6.0], device), + } + } + } + + impl OuterModule { + fn new(device: &Device) -> Self { + Self { + layer1: InnerModule::new(device), + layer2: InnerModule::new(device), + } + } + } + + #[test] + fn nested_module_path_tracking() { + let device = Default::default(); + let module = OuterModule::new(&device); + + let mut collector = TensorPathCollector::new(); + module.visit(&mut collector); + + let paths = collector.paths; + + // Verify we have the expected paths + // Note: Param fields are themselves modules, so we get an extra level + assert!(paths.contains_key("layer1.weight"), "Missing layer1.weight"); + assert!(paths.contains_key("layer1.bias"), "Missing layer1.bias"); + assert!(paths.contains_key("layer2.weight"), "Missing layer2.weight"); + assert!(paths.contains_key("layer2.bias"), "Missing layer2.bias"); + + // Verify the shapes are correct + assert_eq!(paths.get("layer1.weight").unwrap().1, vec![2, 2]); + assert_eq!(paths.get("layer1.bias").unwrap().1, vec![2]); + assert_eq!(paths.get("layer2.weight").unwrap().1, vec![2, 2]); + assert_eq!(paths.get("layer2.bias").unwrap().1, vec![2]); + } + + #[test] + fn linear_module_paths() { + let device = Default::default(); + let config = LinearConfig::new(10, 20).with_bias(true); + let linear = config.init(&device); + + let mut collector = TensorPathCollector::new(); + linear.visit(&mut collector); + + let paths = collector.paths; + + // Linear module has weight and optional bias + assert!(paths.contains_key("weight")); + assert!(paths.contains_key("bias")); + + // Check dimensions + assert_eq!(paths.get("weight").unwrap().1, vec![10, 20]); + assert_eq!(paths.get("bias").unwrap().1, vec![20]); + } + + // Deep nesting test structures (4+ levels) + #[derive(Module, Debug)] + struct Level4Module { + weight: Param>, + bias: Param>, + } + + #[derive(Module, Debug)] + struct Level3Module { + layer: Level4Module, + extra: Level4Module, + } + + #[derive(Module, Debug)] + struct Level2Module { + block1: Level3Module, + block2: Level3Module, + } + + #[derive(Module, Debug)] + struct Level1Module { + encoder: Level2Module, + decoder: Level2Module, + } + + #[derive(Module, Debug)] + struct DeepModel { + backbone: Level1Module, + head: Level4Module, + } + + impl Level4Module { + fn new(device: &Device) -> Self { + Self { + weight: Param::from_data([[1.0, 2.0], [3.0, 4.0]], device), + bias: Param::from_data([5.0, 6.0], device), + } + } + } + + impl Level3Module { + fn new(device: &Device) -> Self { + Self { + layer: Level4Module::new(device), + extra: Level4Module::new(device), + } + } + } + + impl Level2Module { + fn new(device: &Device) -> Self { + Self { + block1: Level3Module::new(device), + block2: Level3Module::new(device), + } + } + } + + impl Level1Module { + fn new(device: &Device) -> Self { + Self { + encoder: Level2Module::new(device), + decoder: Level2Module::new(device), + } + } + } + + impl DeepModel { + fn new(device: &Device) -> Self { + Self { + backbone: Level1Module::new(device), + head: Level4Module::new(device), + } + } + } + + #[test] + fn deep_module_path_tracking() { + let device = Default::default(); + let model = DeepModel::new(&device); + + let mut collector = Collector::new(None, None, false); + model.visit(&mut collector); + + let views = collector.tensors; + let paths: Vec = views.iter().map(|v| v.full_path()).collect(); + + // Test 5-level deep paths + assert!(paths.contains(&"backbone.encoder.block1.layer.weight".to_string())); + assert!(paths.contains(&"backbone.encoder.block1.layer.bias".to_string())); + assert!(paths.contains(&"backbone.encoder.block1.extra.weight".to_string())); + assert!(paths.contains(&"backbone.encoder.block1.extra.bias".to_string())); + + assert!(paths.contains(&"backbone.encoder.block2.layer.weight".to_string())); + assert!(paths.contains(&"backbone.encoder.block2.layer.bias".to_string())); + assert!(paths.contains(&"backbone.encoder.block2.extra.weight".to_string())); + assert!(paths.contains(&"backbone.encoder.block2.extra.bias".to_string())); + + assert!(paths.contains(&"backbone.decoder.block1.layer.weight".to_string())); + assert!(paths.contains(&"backbone.decoder.block1.layer.bias".to_string())); + assert!(paths.contains(&"backbone.decoder.block1.extra.weight".to_string())); + assert!(paths.contains(&"backbone.decoder.block1.extra.bias".to_string())); + + assert!(paths.contains(&"backbone.decoder.block2.layer.weight".to_string())); + assert!(paths.contains(&"backbone.decoder.block2.layer.bias".to_string())); + assert!(paths.contains(&"backbone.decoder.block2.extra.weight".to_string())); + assert!(paths.contains(&"backbone.decoder.block2.extra.bias".to_string())); + + // Test 2-level paths + assert!(paths.contains(&"head.weight".to_string())); + assert!(paths.contains(&"head.bias".to_string())); + + // Total should be 18 tensors (16 from backbone + 2 from head) + assert_eq!(views.len(), 18); + + // Verify data can be materialized + let view = views + .iter() + .find(|v| v.full_path() == "backbone.encoder.block1.layer.weight") + .unwrap(); + let data = view.to_data().unwrap(); + assert_eq!(data.shape, shape![2, 2]); + } + + #[test] + fn deep_module_filtered_export() { + let device = Default::default(); + let model = DeepModel::new(&device); + + // Test filtering at different depths + #[cfg(target_has_atomic = "ptr")] + { + let filter = PathFilter::new().with_regex(r"^backbone\.encoder\..*"); + let mut collector = Collector::new(Some(filter), None, false); + model.visit(&mut collector); + assert_eq!(collector.tensors.len(), 8); // Only encoder tensors + } + + // Test filtering specific blocks + #[cfg(target_has_atomic = "ptr")] + { + let filter = PathFilter::new().with_regex(r".*\.block1\..*"); + let mut collector = Collector::new(Some(filter), None, false); + model.visit(&mut collector); + assert_eq!(collector.tensors.len(), 8); // block1 in both encoder and decoder + } + + // Test filtering by tensor type at any depth + #[cfg(target_has_atomic = "ptr")] + { + let filter = PathFilter::new().with_regex(r".*\.weight$"); + let mut collector = Collector::new(Some(filter), None, false); + model.visit(&mut collector); + assert_eq!(collector.tensors.len(), 9); // All weight tensors + } + + // Test complex multi-pattern filtering + #[cfg(target_has_atomic = "ptr")] + { + let filter = PathFilter::new() + .with_regex(r"^backbone\.encoder\.block1\..*") // All encoder.block1 tensors + .with_regex(r"^backbone\.decoder\..*\.bias$") // All decoder biases + .with_regex(r"^head\.weight$"); // Head weight only + let mut collector = Collector::new(Some(filter), None, false); + model.visit(&mut collector); + + // Should have: + // - 4 from encoder.block1 (2 weights + 2 biases) + // - 4 decoder biases + // - 1 head weight + assert_eq!(collector.tensors.len(), 9); + + let paths: Vec = collector.tensors.iter().map(|v| v.full_path()).collect(); + assert!(paths.contains(&"backbone.encoder.block1.layer.weight".to_string())); + assert!(paths.contains(&"backbone.decoder.block1.layer.bias".to_string())); + assert!(paths.contains(&"head.weight".to_string())); + assert!(!paths.contains(&"head.bias".to_string())); // Not included + } + } + + use crate::traits::ModuleSnapshot; + use burn_nn::Linear; + use hashbrown::HashMap; + + // Test module with Option fields + #[derive(Module, Debug)] + struct OptionalFieldModule { + required: Param>, + optional: Option>>, + } + + impl OptionalFieldModule { + fn new_with_optional(device: &Device) -> Self { + Self { + required: Param::from_data([[1.0, 2.0], [3.0, 4.0]], device), + optional: Some(Param::from_data([5.0, 6.0], device)), + } + } + + fn new_without_optional(device: &Device) -> Self { + Self { + required: Param::from_data([[1.0, 2.0], [3.0, 4.0]], device), + optional: None, + } + } + } + + #[test] + fn optional_field_module_with_value() { + let device = Default::default(); + let module = OptionalFieldModule::new_with_optional(&device); + + let views: HashMap = module + .collect(None, None, false) + .into_iter() + .map(|v| (v.full_path(), v)) + .collect(); + + assert_eq!(views.len(), 2); + assert!(views.contains_key("required")); + assert!(views.contains_key("optional")); + } + + #[test] + fn optional_field_module_without_value() { + let device = Default::default(); + let module = OptionalFieldModule::new_without_optional(&device); + + let views: HashMap = module + .collect(None, None, false) + .into_iter() + .map(|v| (v.full_path(), v)) + .collect(); + + assert_eq!(views.len(), 1); + assert!(views.contains_key("required")); + assert!(!views.contains_key("optional")); + } + + // Test Vec of modules + #[derive(Module, Debug)] + struct VecModule { + layers: Vec, + } + + impl VecModule { + fn new(device: &Device, num_layers: usize) -> Self { + Self { + layers: (0..num_layers) + .map(|_| LinearConfig::new(10, 10).init(device)) + .collect(), + } + } + } + + // Test tuple of modules + #[derive(Module, Debug)] + struct TupleModule { + layers: (Linear, Linear, Linear), + } + + impl TupleModule { + fn new(device: &Device) -> Self { + Self { + layers: ( + LinearConfig::new(10, 10).init(device), + LinearConfig::new(10, 10).init(device), + LinearConfig::new(10, 10).init(device), + ), + } + } + } + + #[test] + fn vec_module_collect() { + let device = Default::default(); + let module = VecModule::new(&device, 3); + + let views: HashMap = module + .collect(None, None, false) + .into_iter() + .map(|v| (v.full_path(), v)) + .collect(); + + // With the fix, all Vec items should now be properly indexed and visited + assert_eq!(views.len(), 6); // 3 layers × 2 tensors each = 6 tensors + + // Check that all indexed paths exist + assert!(views.contains_key("layers.0.weight")); + assert!(views.contains_key("layers.0.bias")); + assert!(views.contains_key("layers.1.weight")); + assert!(views.contains_key("layers.1.bias")); + assert!(views.contains_key("layers.2.weight")); + assert!(views.contains_key("layers.2.bias")); + } + + #[test] + fn tuple_module_collect() { + let device = Default::default(); + let module = TupleModule::new(&device); + + let snapshots = module.collect(None, None, false); + assert_eq!(snapshots.len(), 6); + + let views: HashMap = + snapshots.into_iter().map(|v| (v.full_path(), v)).collect(); + + assert_eq!(views.len(), 6); + + assert!(views.contains_key("layers.0.weight")); + assert!(views.contains_key("layers.0.bias")); + assert!(views.contains_key("layers.1.weight")); + assert!(views.contains_key("layers.1.bias")); + assert!(views.contains_key("layers.2.weight")); + assert!(views.contains_key("layers.2.bias")); + } + + // Test array of modules + #[derive(Module, Debug)] + struct ArrayModule { + layers: [Linear; 3], + } + + impl ArrayModule { + fn new(device: &Device) -> Self { + Self { + layers: [ + LinearConfig::new(10, 10).init(device), + LinearConfig::new(10, 10).init(device), + LinearConfig::new(10, 10).init(device), + ], + } + } + } + + #[test] + fn array_module_collect() { + let device = Default::default(); + let module = ArrayModule::new(&device); + + let views: HashMap = module + .collect(None, None, false) + .into_iter() + .map(|v| (v.full_path(), v)) + .collect(); + + // All array items should be properly indexed + assert_eq!(views.len(), 6); // 3 layers × 2 tensors each = 6 tensors + + // Check indexed paths + for i in 0..3 { + assert!(views.contains_key(&format!("layers.{}.weight", i))); + assert!(views.contains_key(&format!("layers.{}.bias", i))); + } + } + + // Test enum modules + #[derive(Module, Debug)] + enum EnumModule { + LayerA(Linear), + LayerB(Linear), + LayerC(Linear), + } + + #[test] + fn enum_module_collect() { + let device = Default::default(); + + // Test variant A + let module_a = EnumModule::LayerA(LinearConfig::new(10, 20).init(&device)); + let views_a: HashMap = module_a + .collect(None, None, false) + .into_iter() + .map(|v| (v.full_path(), v)) + .collect(); + + // Should have the variant name in the path + assert_eq!(views_a.len(), 2); + assert!(views_a.contains_key("LayerA.weight")); + assert!(views_a.contains_key("LayerA.bias")); + + // Test variant B + let module_b = EnumModule::LayerB(LinearConfig::new(10, 20).init(&device)); + let views_b: HashMap = module_b + .collect(None, None, false) + .into_iter() + .map(|v| (v.full_path(), v)) + .collect(); + + assert_eq!(views_b.len(), 2); + assert!(views_b.contains_key("LayerB.weight")); + assert!(views_b.contains_key("LayerB.bias")); + } + + // Container type tracking tests + #[test] + fn linear_container_type() { + let device = Default::default(); + + #[derive(Module, Debug)] + struct ModelWithLinear { + linear: Linear, + } + + impl ModelWithLinear { + fn new(device: &Device) -> Self { + Self { + linear: LinearConfig::new(10, 20).init(device), + } + } + } + + let model = ModelWithLinear::new(&device); + + let views: HashMap = model + .collect(None, None, false) + .into_iter() + .map(|v| (v.full_path(), v)) + .collect(); + + // Check that tensors inside Linear layers have "Struct:Linear" as their module type + for (path, view) in views.iter() { + if path == "linear.weight" || path == "linear.bias" { + assert_eq!( + view.module_type(), + Some("Struct:Linear".to_string()), + "Tensor '{}' should have module type 'Struct:Linear'", + path + ); + } + } + } + + #[test] + fn complex_model_container_types() { + let device = Default::default(); + + #[derive(Module, Debug)] + struct ComplexModel { + linear_layers: [Linear; 2], + vec_layers: Vec, + single_linear: Linear, + } + + impl ComplexModel { + fn new(device: &Device) -> Self { + Self { + linear_layers: [ + LinearConfig::new(100, 50).init(device), + LinearConfig::new(50, 10).init(device), + ], + vec_layers: vec![ + LinearConfig::new(10, 10).init(device), + LinearConfig::new(10, 10).init(device), + ], + single_linear: LinearConfig::new(10, 1).init(device), + } + } + } + + let model = ComplexModel::new(&device); + + let views: HashMap = model + .collect(None, None, false) + .into_iter() + .map(|v| (v.full_path(), v)) + .collect(); + + // Should have 10 tensors total + assert_eq!(views.len(), 10); + + // Verify different module types + for (_path, view) in views.iter() { + assert_eq!(view.module_type(), Some("Struct:Linear".to_string())); + } + } + + #[test] + fn collect_with_container_filter() { + let device = Default::default(); + + #[derive(Module, Debug)] + struct FilterTestModel { + layers: Vec, + } + + impl FilterTestModel { + fn new(device: &Device) -> Self { + Self { + layers: vec![ + LinearConfig::new(10, 10).init(device), + LinearConfig::new(10, 10).init(device), + ], + } + } + } + + let model = FilterTestModel::new(&device); + + // Filter to only collect tensors from Linear modules + let filter = PathFilter::new().with_predicate(|_path, container_path| { + container_path.split('.').next_back() == Some("Struct:Linear") + }); + + let linear_views: Vec = model.collect(Some(filter), None, false); + + // All collected tensors should be from Linear modules + for view in linear_views.iter() { + assert_eq!( + view.module_type(), + Some("Struct:Linear".to_string()), + "All tensors should be from Linear modules" + ); + } + + // Should have collected all Linear tensors + assert_eq!(linear_views.len(), 4); + } +} diff --git a/crates/burn-store/src/filter.rs b/crates/burn-store/src/filter.rs new file mode 100644 index 0000000..1d995b6 --- /dev/null +++ b/crates/burn-store/src/filter.rs @@ -0,0 +1,625 @@ +use alloc::format; +use alloc::string::String; +use alloc::vec::Vec; +use core::fmt; + +#[cfg(feature = "std")] +use regex::Regex; + +/// A sophisticated path filter that supports multiple matching strategies. +/// +/// The filter uses an OR logic - a path is included if it matches ANY of the configured criteria. +/// This allows for flexible and powerful filtering configurations. +/// +/// # Examples +/// +/// ```rust,no_run +/// # use burn_store::PathFilter; +/// // Create a filter that matches encoder paths or any weight path +/// let filter = PathFilter::new() +/// .with_regex(r"^encoder\..*") +/// .with_regex(r".*\.weight$") +/// .with_full_path("special_tensor"); +/// +/// // Check if a path should be included +/// if filter.matches("encoder.layer1.weight") { +/// // This will match due to both regex patterns +/// } +/// ``` +#[derive(Debug, Clone, Default)] +pub struct PathFilter { + /// Compiled regex patterns for matching paths + #[cfg(feature = "std")] + regex_patterns: Vec, + + /// Exact full paths to match + exact_paths: Vec, + + /// Predicate functions for custom matching logic based on path and container path + /// Note: These cannot be cloned, so we store them separately + predicates: Vec bool>, + + /// If true, matches all paths (overrides other filters) + match_all: bool, +} + +impl PathFilter { + /// Create a new empty filter (matches nothing by default) + pub fn new() -> Self { + Self::default() + } + + /// Create a filter that matches all paths + pub fn all() -> Self { + Self { + match_all: true, + ..Default::default() + } + } + + /// Create a filter that matches nothing + pub fn none() -> Self { + Self::default() + } + + /// Add a regex pattern for matching paths + #[cfg(feature = "std")] + pub fn with_regex>(mut self, pattern: S) -> Self { + if let Ok(regex) = Regex::new(pattern.as_ref()) { + self.regex_patterns.push(regex); + } + // TODO: Consider returning Result to handle regex compilation errors + self + } + + /// Add multiple regex patterns + #[cfg(feature = "std")] + pub fn with_regexes(mut self, patterns: I) -> Self + where + I: IntoIterator, + S: AsRef, + { + for pattern in patterns { + if let Ok(regex) = Regex::new(pattern.as_ref()) { + self.regex_patterns.push(regex); + } + } + self + } + + /// Add an exact full path to match + pub fn with_full_path>(mut self, path: S) -> Self { + self.exact_paths.push(path.into()); + self + } + + /// Add multiple exact full paths + pub fn with_full_paths(mut self, paths: I) -> Self + where + I: IntoIterator, + S: Into, + { + self.exact_paths.extend(paths.into_iter().map(|p| p.into())); + self + } + + /// Add a predicate function for custom matching based on path and container path + pub fn with_predicate(mut self, predicate: fn(&str, &str) -> bool) -> Self { + self.predicates.push(predicate); + self + } + + /// Add multiple predicates + pub fn with_predicates(mut self, predicates: I) -> Self + where + I: IntoIterator bool>, + { + self.predicates.extend(predicates); + self + } + + /// Set to match all paths + pub fn match_all(mut self) -> Self { + self.match_all = true; + self + } + + /// Check if a path matches this filter (assumes empty container path for backward compatibility) + pub fn matches(&self, path: &str) -> bool { + self.matches_with_container_path_str(path, "") + } + + /// Check if a path and container type match this filter (for backward compatibility) + pub fn matches_with_container(&self, path: &str, container_type: &str) -> bool { + // For backward compatibility, treat single container type as the full path + self.matches_with_container_path_str(path, container_type) + } + + /// Check if a path and container path match this filter + pub fn matches_with_container_path(&self, path: &[String], container_stack: &[String]) -> bool { + let path_str = path.join("."); + let container_path = container_stack.join("."); + self.matches_with_container_path_str(&path_str, &container_path) + } + + /// Check if a path and container path (dot-notated strings) match this filter + pub fn matches_with_container_path_str(&self, path: &str, container_path: &str) -> bool { + // If match_all is set, always return true + if self.match_all { + return true; + } + + // If no filters are configured, match nothing + if self.is_empty() { + return false; + } + + // Check exact path matches + if self.exact_paths.iter().any(|p| p == path) { + return true; + } + + // Check regex patterns (on the path) + #[cfg(feature = "std")] + { + for regex in &self.regex_patterns { + if regex.is_match(path) { + return true; + } + } + } + + // Check predicates with container path + if self + .predicates + .iter() + .any(|pred| pred(path, container_path)) + { + return true; + } + + false + } + + /// Check if the filter is empty (matches nothing) + pub fn is_empty(&self) -> bool { + if self.match_all { + return false; + } + + #[cfg(feature = "std")] + let regex_empty = self.regex_patterns.is_empty(); + #[cfg(not(feature = "std"))] + let regex_empty = true; + + self.exact_paths.is_empty() && self.predicates.is_empty() && regex_empty + } + + /// Get the number of filter criteria configured + pub fn criteria_count(&self) -> usize { + if self.match_all { + return 1; + } + + #[allow(unused_mut)] + let mut count = self.exact_paths.len() + self.predicates.len(); + + #[cfg(feature = "std")] + { + count += self.regex_patterns.len(); + } + + count + } + + /// Clear all regex patterns + #[cfg(feature = "std")] + pub fn clear_regex(&mut self) -> &mut Self { + self.regex_patterns.clear(); + self + } + + /// Clear all exact paths + pub fn clear_paths(&mut self) -> &mut Self { + self.exact_paths.clear(); + self + } + + /// Clear all predicates + pub fn clear_predicates(&mut self) -> &mut Self { + self.predicates.clear(); + self + } + + /// Clear all filters + pub fn clear(&mut self) -> &mut Self { + #[cfg(feature = "std")] + self.clear_regex(); + + self.clear_paths().clear_predicates(); + self.match_all = false; + self + } + + /// Create a filter from regex patterns only + #[cfg(feature = "std")] + pub fn from_regex_patterns(patterns: I) -> Self + where + I: IntoIterator, + S: AsRef, + { + Self::new().with_regexes(patterns) + } + + /// Create a filter from exact paths only + pub fn from_paths(paths: I) -> Self + where + I: IntoIterator, + S: Into, + { + Self::new().with_full_paths(paths) + } + + /// Create a filter from a single predicate + pub fn from_predicate(predicate: fn(&str, &str) -> bool) -> Self { + Self::new().with_predicate(predicate) + } + + /// Combine with another filter using OR logic + pub fn or(mut self, other: Self) -> Self { + if self.match_all || other.match_all { + return Self::all(); + } + + #[cfg(feature = "std")] + { + self.regex_patterns.extend(other.regex_patterns); + } + + self.exact_paths.extend(other.exact_paths); + self.predicates.extend(other.predicates); + + self + } +} + +impl fmt::Display for PathFilter { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + if self.match_all { + return write!(f, "PathFilter::all()"); + } + + if self.is_empty() { + return write!(f, "PathFilter::none()"); + } + + write!(f, "PathFilter[")?; + + let mut parts = Vec::new(); + + #[cfg(feature = "std")] + if !self.regex_patterns.is_empty() { + parts.push(format!("regex: {:?}", self.regex_patterns)); + } + + if !self.exact_paths.is_empty() { + parts.push(format!("paths: {:?}", self.exact_paths)); + } + + if !self.predicates.is_empty() { + parts.push(format!("predicates: {}", self.predicates.len())); + } + + write!(f, "{}]", parts.join(", ")) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn empty_filter() { + let filter = PathFilter::new(); + assert!(filter.is_empty()); + assert!(!filter.matches("encoder.weight")); + assert!(!filter.matches("decoder.bias")); + } + + #[test] + fn match_all() { + let filter = PathFilter::all(); + assert!(!filter.is_empty()); + assert!(filter.matches("encoder.weight")); + assert!(filter.matches("decoder.bias")); + assert!(filter.matches("anything")); + } + + #[test] + fn exact_paths() { + let filter = PathFilter::new() + .with_full_path("encoder.weight") + .with_full_path("decoder.bias"); + + assert!(filter.matches("encoder.weight")); + assert!(filter.matches("decoder.bias")); + assert!(!filter.matches("encoder.bias")); + assert!(!filter.matches("decoder.weight")); + } + + #[test] + #[cfg(feature = "std")] + fn regex_patterns() { + let filter = PathFilter::new() + .with_regex(r"^encoder\..*") + .with_regex(r".*\.weight$"); + + assert!(filter.matches("encoder.layer1.bias")); + assert!(filter.matches("decoder.weight")); + assert!(filter.matches("encoder.weight")); + assert!(!filter.matches("decoder.bias")); + } + + #[test] + fn predicates() { + fn contains_norm(path: &str, _container_path: &str) -> bool { + path.contains("norm") + } + + fn is_short(path: &str, _container_path: &str) -> bool { + path.len() < 10 + } + + let filter = PathFilter::new() + .with_predicate(contains_norm) + .with_predicate(is_short); + + assert!(filter.matches("norm.weight")); + assert!(filter.matches("layer.norm.bias")); + assert!(filter.matches("bias")); + assert!(!filter.matches("encoder.decoder.weight.long.name")); + } + + #[test] + fn combined_filters() { + let filter = PathFilter::new() + .with_full_path("special.tensor") + .with_predicate(|path, _container_path| path.contains("attention")); + + #[cfg(feature = "std")] + let filter = filter.with_regex(r"^encoder\..*"); + + assert!(filter.matches("special.tensor")); + assert!(filter.matches("self_attention.query")); + + #[cfg(feature = "std")] + assert!(filter.matches("encoder.anything")); + + assert!(!filter.matches("decoder.weight")); + } + + #[test] + fn or_combination() { + let encoder_filter = PathFilter::new().with_full_path("encoder.weight"); + let decoder_filter = PathFilter::new().with_full_path("decoder.bias"); + + let combined = encoder_filter.or(decoder_filter); + + assert!(combined.matches("encoder.weight")); + assert!(combined.matches("decoder.bias")); + assert!(!combined.matches("model.head.weight")); + } + + #[test] + #[cfg(feature = "std")] + fn common_patterns() { + // Test encoder pattern + let encoder = PathFilter::new().with_regex(r"^encoder\..*"); + assert!(encoder.matches("encoder.weight")); + assert!(!encoder.matches("decoder.weight")); + + // Test weights-only pattern + let weights = PathFilter::new().with_regex(r".*\.weight$"); + assert!(weights.matches("encoder.weight")); + assert!(weights.matches("decoder.weight")); + assert!(!weights.matches("encoder.bias")); + + // Test layer-specific patterns + let layers = PathFilter::new() + .with_regex(r"(^|.*\.)layers\.0\.") + .with_regex(r"(^|.*\.)layers\.2\.") + .with_regex(r"(^|.*\.)layers\.4\."); + assert!(layers.matches("model.layers.0.weight")); + assert!(layers.matches("layers.2.bias")); + assert!(!layers.matches("layers.1.weight")); + } + + #[test] + fn criteria_count() { + let filter = PathFilter::new() + .with_full_path("path1") + .with_full_path("path2") + .with_predicate(|_, _| true); + + #[cfg(feature = "std")] + let filter = filter.with_regex(".*"); + + #[cfg(feature = "std")] + assert_eq!(filter.criteria_count(), 4); + + #[cfg(not(feature = "std"))] + assert_eq!(filter.criteria_count(), 3); + } + + #[test] + fn clear_operations() { + let mut filter = PathFilter::new().with_full_path("test"); + + filter.clear_paths(); + assert!(!filter.matches("test")); + + filter.clear(); + assert!(filter.is_empty()); + } + + #[test] + fn container_predicates() { + // Filter that matches only Linear module weights + let linear_weights = PathFilter::new().with_predicate(|path, container_path| { + container_path.split('.').next_back() == Some("Linear") && path.ends_with(".weight") + }); + + assert!(linear_weights.matches_with_container("layer1.weight", "Linear")); + assert!(!linear_weights.matches_with_container("layer1.weight", "Conv2d")); + assert!(!linear_weights.matches_with_container("layer1.bias", "Linear")); + + // Filter for specific container types + let conv_only = PathFilter::new().with_predicate(|_path, container_path| { + let last = container_path.split('.').next_back(); + last == Some("Conv2d") || last == Some("ConvTranspose2d") + }); + + assert!(conv_only.matches_with_container("encoder.weight", "Conv2d")); + assert!(conv_only.matches_with_container("decoder.weight", "ConvTranspose2d")); + assert!(!conv_only.matches_with_container("fc.weight", "Linear")); + + // Combine path and container predicates + let combined = PathFilter::new() + .with_predicate(|path, _container_path| path.starts_with("encoder.")) + .with_predicate(|_path, container_path| { + container_path.split('.').next_back() == Some("BatchNorm2d") + }); + + // Should match either condition (OR logic) + assert!(combined.matches_with_container("encoder.layer1", "Linear")); + assert!(combined.matches_with_container("decoder.bn", "BatchNorm2d")); + assert!(!combined.matches_with_container("decoder.layer", "Linear")); + } + + #[test] + fn container_predicate_with_regex() { + // Combine regex patterns with container predicates + #[cfg(feature = "std")] + { + let filter = PathFilter::new() + .with_regex(r"^encoder\..*") + .with_predicate(|path, container_path| { + container_path.split('.').next_back() == Some("Linear") + && path.contains(".bias") + }); + + // Matches due to regex + assert!(filter.matches_with_container("encoder.layer1.weight", "Conv2d")); + // Matches due to container predicate + assert!(filter.matches_with_container("decoder.fc.bias", "Linear")); + // Doesn't match either + assert!(!filter.matches_with_container("decoder.conv.weight", "Conv2d")); + } + } + + #[test] + fn container_stack_predicates() { + // Filter using full container path - only tensors nested in a specific hierarchy + let nested_filter = PathFilter::new().with_predicate(|_path, container_path| { + // Check if tensor is nested within: Model -> TransformerBlock -> Linear + let parts: Vec<&str> = container_path.split('.').collect(); + parts.len() >= 3 + && parts[0] == "Model" + && parts[1] == "TransformerBlock" + && parts[2] == "Linear" + }); + + assert!(nested_filter.matches_with_container_path_str( + "encoder.weight", + "Model.TransformerBlock.Linear.Param" + )); + assert!( + !nested_filter + .matches_with_container_path_str("decoder.weight", "Model.Decoder.Linear.Param") + ); + assert!(!nested_filter.matches_with_container_path_str( + "encoder.weight", + "Model.TransformerBlock.Conv2d.Param" + )); + + // Filter that checks for specific depth in hierarchy + let depth_filter = PathFilter::new().with_predicate(|_path, container_path| { + let parts: Vec<&str> = container_path.split('.').collect(); + parts.len() == 4 && parts.get(2) == Some(&"Linear") + }); + + assert!(depth_filter.matches_with_container_path_str( + "model.layer.weight", + "Model.TransformerBlock.Linear.Param" + )); + assert!( + !depth_filter + .matches_with_container_path_str("model.weight", "Model.TransformerBlock.Conv2d") + ); // Too shallow + + // Filter that checks any Linear in the path (not just the last) + let any_linear = PathFilter::new() + .with_predicate(|_path, container_path| container_path.contains("Linear")); + + assert!( + any_linear.matches_with_container_path_str( + "some.path", + "Model.TransformerBlock.Linear.Param" + ) + ); + assert!( + any_linear.matches_with_container_path_str("other.path", "Model.Decoder.Linear.Param") + ); + assert!( + !any_linear.matches_with_container_path_str( + "conv.path", + "Model.TransformerBlock.Conv2d.Param" + ) + ); + } + + #[test] + fn container_path_dot_notation() { + // Filter using dot-notated container path + let dot_filter = PathFilter::new().with_predicate(|_path, container_path| { + container_path.starts_with("Model.TransformerBlock") + }); + + // Test with matches_with_container_path + assert!( + dot_filter.matches_with_container_path_str("weight", "Model.TransformerBlock.Linear") + ); + assert!(!dot_filter.matches_with_container_path_str("weight", "Model.Decoder.Linear")); + + // Filter that checks for specific patterns in container path + let pattern_filter = PathFilter::new().with_predicate(|_path, container_path| { + // Match any path that has Linear after a block + container_path.contains("Block.Linear") || container_path.contains("Block.Conv") + }); + + assert!( + pattern_filter + .matches_with_container_path_str("weight", "Model.TransformerBlock.Linear") + ); + assert!(pattern_filter.matches_with_container_path_str("weight", "Model.ResBlock.Conv2d")); + assert!(!pattern_filter.matches_with_container_path_str("weight", "Model.Linear.Param")); + + // Filter combining path and container path patterns + let combined = PathFilter::new().with_predicate(|path, container_path| { + // Only weights in Linear layers that are inside blocks + path.ends_with(".weight") + && container_path.contains("Block") + && container_path.split('.').next_back() == Some("Linear") + }); + + assert!( + combined + .matches_with_container_path_str("layer.weight", "Model.TransformerBlock.Linear") + ); + assert!( + !combined + .matches_with_container_path_str("layer.bias", "Model.TransformerBlock.Linear") + ); + assert!(!combined.matches_with_container_path_str("layer.weight", "Model.Decoder.Linear")); + } +} diff --git a/crates/burn-store/src/keyremapper.rs b/crates/burn-store/src/keyremapper.rs new file mode 100644 index 0000000..f7bebd9 --- /dev/null +++ b/crates/burn-store/src/keyremapper.rs @@ -0,0 +1,674 @@ +use alloc::collections::BTreeMap; +use alloc::string::{String, ToString}; +use alloc::vec::Vec; + +use regex::{self, Regex}; + +use crate::TensorSnapshot; + +/// Key remapper for transforming tensor names. +/// +/// This allows mapping tensor names from one naming convention to another, +/// which is useful for loading models from different frameworks or versions. +/// +/// # Examples +/// +/// ```rust +/// # use burn_store::KeyRemapper; +/// // Create a key remapper +/// let remapper = KeyRemapper::new() +/// .add_pattern(r"^pytorch\.(.*)", "burn.$1").expect("valid regex") // pytorch.layer -> burn.layer +/// .add_pattern(r"\.gamma$", ".weight").expect("valid regex"); // layer.gamma -> layer.weight +/// +/// // Use remapper with stores +/// // store.remap(remapper) +/// ``` +#[derive(Debug, Clone, Default)] +pub struct KeyRemapper { + /// Pattern-based remapping rules (regex pattern, replacement string) + pub patterns: Vec<(Regex, String)>, +} + +impl KeyRemapper { + /// Create a new empty key remapper + pub fn new() -> Self { + Self::default() + } + + /// Add a remapping pattern (compiles regex) + /// + /// # Arguments + /// + /// * `from` - Source pattern (regex string) + /// * `to` - Replacement string (can include capture groups like `$1`) + /// + /// # Returns + /// + /// * `Ok(Self)` - Updated remapping configuration + /// * `Err(regex::Error)` - If regex compilation fails + pub fn add_pattern(mut self, from: S1, to: S2) -> Result + where + S1: AsRef, + S2: Into, + { + let regex = Regex::new(from.as_ref())?; + self.patterns.push((regex, to.into())); + Ok(self) + } + + /// Create from a list of compiled regex patterns + pub fn from_compiled_patterns(patterns: Vec<(Regex, String)>) -> Self { + Self { patterns } + } + + /// Create from string patterns (will compile to regex) + /// + /// # Arguments + /// + /// * `patterns` - Vector of (pattern, replacement) tuples + /// + /// # Returns + /// + /// * `Ok(Self)` - New remapping configuration + /// * `Err(regex::Error)` - If any regex compilation fails + pub fn from_patterns(patterns: Vec<(S1, S2)>) -> Result + where + S1: AsRef, + S2: Into, + { + let mut compiled_patterns = Vec::new(); + for (pattern, replacement) in patterns { + let regex = Regex::new(pattern.as_ref())?; + compiled_patterns.push((regex, replacement.into())); + } + Ok(Self { + patterns: compiled_patterns, + }) + } + + /// Create from an iterator of patterns + /// + /// # Arguments + /// + /// * `iter` - Iterator yielding (pattern, replacement) tuples + /// + /// # Returns + /// + /// * `Ok(Self)` - New remapping configuration + /// * `Err(regex::Error)` - If any regex compilation fails + pub fn from_pattern_iter(iter: I) -> Result + where + I: IntoIterator, + S1: AsRef, + S2: Into, + { + let patterns: Result, _> = iter + .into_iter() + .map(|(from, to)| Ok((Regex::new(from.as_ref())?, to.into()))) + .collect(); + Ok(Self { + patterns: patterns?, + }) + } + + /// Check if the remapping is empty + pub fn is_empty(&self) -> bool { + self.patterns.is_empty() + } + + /// Convert to the format expected by remap_tensor_paths_with_patterns + pub fn to_regex_pairs(&self) -> Vec<(Regex, String)> { + self.patterns.clone() + } + + /// Remap tensor paths using the configured patterns. + /// + /// # Arguments + /// + /// * `tensors` - Vec of TensorSnapshots to remap + /// + /// # Returns + /// + /// A tuple containing: + /// * The remapped Vec of TensorSnapshots with updated paths + /// * A vector of (new_path, original_path) showing the transformations + pub fn remap( + &self, + mut tensors: Vec, + ) -> (Vec, Vec<(String, String)>) { + if self.patterns.is_empty() { + let remapped_names = tensors + .iter() + .map(|v| { + let path = v.full_path(); + (path.clone(), path) + }) + .collect(); + return (tensors, remapped_names); + } + + let mut remapped_snapshots = Vec::new(); + let mut remapped_names = Vec::new(); + + for mut snapshot in tensors.drain(..) { + let original_path = snapshot.full_path(); + let mut new_path = original_path.clone(); + + // Apply all patterns to get the new path + for (pattern, replacement) in &self.patterns { + if pattern.is_match(&new_path) { + new_path = pattern + .replace_all(&new_path, replacement.as_str()) + .to_string(); + } + } + + // Update the snapshot's internal path_stack if the path changed + if new_path != original_path + && let Some(ref mut path_stack) = snapshot.path_stack + { + *path_stack = new_path.split('.').map(|s| s.to_string()).collect(); + } + + remapped_names.push((new_path.clone(), original_path)); + remapped_snapshots.push(snapshot); + } + + (remapped_snapshots, remapped_names) + } +} + +/// Map tensor paths to have contiguous numeric indices. +/// +/// This function detects numeric indices in tensor paths and renumbers them +/// to be contiguous (0, 1, 2, ...) while preserving their relative order. +/// It handles nested sequential structures by processing ALL numeric indices +/// in each path independently based on their position context. +/// +/// This is useful when loading PyTorch models that have gaps in layer numbering, +/// such as when using `nn.Sequential` with mixed layer types (e.g., Conv2d + ReLU +/// where only Conv2d has parameters). +/// +/// # Example +/// +/// Simple case - input paths: +/// - `fc.0.weight`, `fc.0.bias` +/// - `fc.2.weight`, `fc.2.bias` +/// - `fc.4.weight`, `fc.4.bias` +/// +/// Output paths: +/// - `fc.0.weight`, `fc.0.bias` +/// - `fc.1.weight`, `fc.1.bias` +/// - `fc.2.weight`, `fc.2.bias` +/// +/// Nested case - input paths: +/// - `feature.layers.0.conv_block.0.weight` +/// - `feature.layers.0.conv_block.2.weight` +/// - `feature.layers.2.conv_block.0.weight` +/// - `feature.layers.2.conv_block.2.weight` +/// +/// Output paths: +/// - `feature.layers.0.conv_block.0.weight` +/// - `feature.layers.0.conv_block.1.weight` +/// - `feature.layers.1.conv_block.0.weight` +/// - `feature.layers.1.conv_block.1.weight` +/// +/// # Arguments +/// +/// * `tensors` - Vec of TensorSnapshots to map +/// +/// # Returns +/// +/// A tuple containing: +/// * The mapped Vec of TensorSnapshots with updated paths +/// * A vector of (new_path, original_path) showing the transformations +pub fn map_indices_contiguous( + mut tensors: Vec, +) -> (Vec, Vec<(String, String)>) { + if tensors.is_empty() { + return (tensors, Vec::new()); + } + + // Step 1: Collect all paths and find all index positions + // For each index position (identified by prefix using ORIGINAL indices), + // collect all indices seen at that position. + // + // Key: prefix using original path (e.g., "feature.layers." or "feature.layers.0.conv_block.") + // Value: BTreeMap of original_index -> new_index + let mut index_maps: BTreeMap> = BTreeMap::new(); + + // First pass: collect all indices at each position using original prefixes + for snapshot in &tensors { + let path = snapshot.full_path(); + let parts: Vec<&str> = path.split('.').collect(); + + // Check each part for numeric indices + for (i, part) in parts.iter().enumerate() { + if let Ok(index) = part.parse::() { + // The prefix is everything before this index (using original path) + let prefix = if i > 0 { + format!("{}.", parts[..i].join(".")) + } else { + String::new() + }; + + index_maps + .entry(prefix) + .or_default() + .entry(index) + .or_insert(usize::MAX); // Placeholder + } + } + } + + // Second pass: assign contiguous indices for each position + for indices in index_maps.values_mut() { + let mut sorted_indices: Vec = indices.keys().cloned().collect(); + sorted_indices.sort(); + + for (new_idx, old_idx) in sorted_indices.into_iter().enumerate() { + indices.insert(old_idx, new_idx); + } + } + + // Third pass: apply the remapping to all tensors + // We use original prefixes for lookup since that's how we collected indices + let mut mapped_snapshots = Vec::new(); + let mut transformations = Vec::new(); + + for mut snapshot in tensors.drain(..) { + let original_path = snapshot.full_path(); + let new_path = remap_all_indices_with_original_prefix(&original_path, &index_maps); + + // Update the snapshot's internal path_stack if the path changed + if new_path != original_path + && let Some(ref mut path_stack) = snapshot.path_stack + { + *path_stack = new_path.split('.').map(|s| s.to_string()).collect(); + } + + transformations.push((new_path, original_path)); + mapped_snapshots.push(snapshot); + } + + (mapped_snapshots, transformations) +} + +/// Remap all numeric indices in a path using the provided index maps. +/// Uses original path prefixes for lookup. +fn remap_all_indices_with_original_prefix( + path: &str, + index_maps: &BTreeMap>, +) -> String { + let parts: Vec<&str> = path.split('.').collect(); + let mut result_parts: Vec = Vec::with_capacity(parts.len()); + + for (i, part) in parts.iter().enumerate() { + if let Ok(index) = part.parse::() { + // Build the prefix from ORIGINAL parts (not remapped) + let prefix = if i > 0 { + format!("{}.", parts[..i].join(".")) + } else { + String::new() + }; + + // Look up the new index using original prefix + if let Some(index_map) = index_maps.get(&prefix) + && let Some(&new_index) = index_map.get(&index) + { + result_parts.push(new_index.to_string()); + continue; + } + } + // Not a numeric index or no mapping found, keep as-is + result_parts.push((*part).to_string()); + } + + result_parts.join(".") +} + +#[cfg(all(test, feature = "std"))] +mod tests { + use super::*; + use burn_core::module::ParamId; + use burn_core::tensor::{Bytes, DType, TensorData, shape}; + + fn create_test_tensor_snapshot(name: &str) -> TensorSnapshot { + let data = TensorData { + bytes: Bytes::from_bytes_vec(vec![1, 2, 3, 4]), + shape: shape![2, 2], + dtype: DType::F32, + }; + let path_parts: Vec = name.split('.').map(|s| s.to_string()).collect(); + TensorSnapshot::from_data(data, path_parts, vec!["Test".to_string()], ParamId::new()) + } + + #[test] + fn test_key_remapper_basic() { + let remapper = KeyRemapper::new() + .add_pattern(r"^encoder\.", "transformer.encoder.") + .expect("valid regex"); + + let tensors = vec![ + create_test_tensor_snapshot("encoder.layer1.weight"), + create_test_tensor_snapshot("decoder.layer1.weight"), + ]; + + let (remapped, transformations) = remapper.remap(tensors); + + // Check that remapped views exist with correct paths + assert!( + remapped + .iter() + .any(|v| v.full_path() == "transformer.encoder.layer1.weight") + ); + assert!( + remapped + .iter() + .any(|v| v.full_path() == "decoder.layer1.weight") + ); + assert_eq!(remapped.len(), 2); + + // Check transformations + let encoder_transform = transformations + .iter() + .find(|(_new, old)| old == "encoder.layer1.weight") + .expect("should find encoder transformation"); + assert_eq!(encoder_transform.0, "transformer.encoder.layer1.weight"); + } + + #[test] + fn test_key_remapper_multiple_patterns() { + let remapper = KeyRemapper::new() + .add_pattern(r"^encoder\.", "transformer.encoder.") + .expect("valid regex") + .add_pattern(r"\.gamma$", ".weight") + .expect("valid regex"); + + let tensors = vec![create_test_tensor_snapshot("encoder.layer1.gamma")]; + + let (remapped, _) = remapper.remap(tensors); + + assert!( + remapped + .iter() + .any(|v| v.full_path() == "transformer.encoder.layer1.weight") + ); + assert_eq!(remapped.len(), 1); + } + + #[test] + fn test_key_remapper_from_patterns() { + let patterns = vec![(r"^pytorch\.", "burn."), (r"\.bias$", ".bias_param")]; + let remapper = KeyRemapper::from_patterns(patterns).expect("valid patterns"); + + let tensors = vec![create_test_tensor_snapshot("pytorch.linear.bias")]; + + let (remapped, _) = remapper.remap(tensors); + + assert!( + remapped + .iter() + .any(|v| v.full_path() == "burn.linear.bias_param") + ); + } + + #[test] + fn test_key_remapper_empty() { + let remapper = KeyRemapper::new(); + assert!(remapper.is_empty()); + + let tensors = vec![create_test_tensor_snapshot("test.weight")]; + + let (remapped, transformations) = remapper.remap(tensors); + + assert!(remapped.iter().any(|v| v.full_path() == "test.weight")); + assert_eq!(remapped.len(), 1); + assert_eq!(transformations.len(), 1); + assert_eq!( + transformations[0], + ("test.weight".to_string(), "test.weight".to_string()) + ); + } + + #[test] + fn test_map_indices_contiguous_basic() { + // Simulate PyTorch nn.Sequential with Conv2d (0, 2, 4) and ReLU (1, 3, 5) + // Only Conv2d layers have parameters + let tensors = vec![ + create_test_tensor_snapshot("fc.0.weight"), + create_test_tensor_snapshot("fc.0.bias"), + create_test_tensor_snapshot("fc.2.weight"), + create_test_tensor_snapshot("fc.2.bias"), + create_test_tensor_snapshot("fc.4.weight"), + create_test_tensor_snapshot("fc.4.bias"), + ]; + + let (reindexed, transformations) = map_indices_contiguous(tensors); + + // Check that indices are now contiguous + assert!(reindexed.iter().any(|v| v.full_path() == "fc.0.weight")); + assert!(reindexed.iter().any(|v| v.full_path() == "fc.0.bias")); + assert!(reindexed.iter().any(|v| v.full_path() == "fc.1.weight")); + assert!(reindexed.iter().any(|v| v.full_path() == "fc.1.bias")); + assert!(reindexed.iter().any(|v| v.full_path() == "fc.2.weight")); + assert!(reindexed.iter().any(|v| v.full_path() == "fc.2.bias")); + assert_eq!(reindexed.len(), 6); + + // Check transformations + let transform_2_to_1 = transformations + .iter() + .find(|(_, old)| old == "fc.2.weight") + .expect("should find fc.2.weight transformation"); + assert_eq!(transform_2_to_1.0, "fc.1.weight"); + + let transform_4_to_2 = transformations + .iter() + .find(|(_, old)| old == "fc.4.weight") + .expect("should find fc.4.weight transformation"); + assert_eq!(transform_4_to_2.0, "fc.2.weight"); + } + + #[test] + fn test_map_indices_contiguous_already_contiguous() { + // Already contiguous indices should remain unchanged + let tensors = vec![ + create_test_tensor_snapshot("fc.0.weight"), + create_test_tensor_snapshot("fc.1.weight"), + create_test_tensor_snapshot("fc.2.weight"), + ]; + + let (reindexed, transformations) = map_indices_contiguous(tensors); + + assert!(reindexed.iter().any(|v| v.full_path() == "fc.0.weight")); + assert!(reindexed.iter().any(|v| v.full_path() == "fc.1.weight")); + assert!(reindexed.iter().any(|v| v.full_path() == "fc.2.weight")); + assert_eq!(reindexed.len(), 3); + + // All transformations should have same old and new paths + for (new, old) in &transformations { + assert_eq!(new, old); + } + } + + #[test] + fn test_map_indices_contiguous_multiple_prefixes() { + // Different prefixes should be mapped independently + let tensors = vec![ + create_test_tensor_snapshot("encoder.0.weight"), + create_test_tensor_snapshot("encoder.2.weight"), + create_test_tensor_snapshot("decoder.1.weight"), + create_test_tensor_snapshot("decoder.5.weight"), + ]; + + let (reindexed, _) = map_indices_contiguous(tensors); + + // encoder: 0, 2 -> 0, 1 + assert!( + reindexed + .iter() + .any(|v| v.full_path() == "encoder.0.weight") + ); + assert!( + reindexed + .iter() + .any(|v| v.full_path() == "encoder.1.weight") + ); + + // decoder: 1, 5 -> 0, 1 + assert!( + reindexed + .iter() + .any(|v| v.full_path() == "decoder.0.weight") + ); + assert!( + reindexed + .iter() + .any(|v| v.full_path() == "decoder.1.weight") + ); + } + + #[test] + fn test_map_indices_contiguous_no_indices() { + // Paths without indices should remain unchanged + let tensors = vec![ + create_test_tensor_snapshot("encoder.weight"), + create_test_tensor_snapshot("decoder.bias"), + ]; + + let (reindexed, transformations) = map_indices_contiguous(tensors); + + assert!(reindexed.iter().any(|v| v.full_path() == "encoder.weight")); + assert!(reindexed.iter().any(|v| v.full_path() == "decoder.bias")); + + for (new, old) in &transformations { + assert_eq!(new, old); + } + } + + #[test] + fn test_map_indices_contiguous_empty() { + let tensors: Vec = vec![]; + let (reindexed, transformations) = map_indices_contiguous(tensors); + + assert!(reindexed.is_empty()); + assert!(transformations.is_empty()); + } + + #[test] + fn test_map_indices_contiguous_mixed_indexed_and_non_indexed() { + // Mix of indexed and non-indexed paths + let tensors = vec![ + create_test_tensor_snapshot("fc.0.weight"), + create_test_tensor_snapshot("fc.2.weight"), + create_test_tensor_snapshot("output.weight"), // no index + ]; + + let (reindexed, _) = map_indices_contiguous(tensors); + + assert!(reindexed.iter().any(|v| v.full_path() == "fc.0.weight")); + assert!(reindexed.iter().any(|v| v.full_path() == "fc.1.weight")); // 2 -> 1 + assert!(reindexed.iter().any(|v| v.full_path() == "output.weight")); // unchanged + } + + #[test] + fn test_map_indices_contiguous_nested_sequential() { + // Test nested sequential structures like: + // feature = nn.Sequential(ConvBlock, ReLU, ConvBlock, ReLU, ConvBlock) + // where ConvBlock = nn.Sequential(Conv2d, ReLU, Conv2d) + // + // This produces paths like: + // feature.layers.0.conv_block.0.weight (layer 0, conv 0) + // feature.layers.0.conv_block.2.weight (layer 0, conv 2 - skipping ReLU at 1) + // feature.layers.2.conv_block.0.weight (layer 2 - skipping ReLU at 1, conv 0) + // feature.layers.2.conv_block.2.weight (layer 2, conv 2) + let tensors = vec![ + create_test_tensor_snapshot("feature.layers.0.conv_block.0.weight"), + create_test_tensor_snapshot("feature.layers.0.conv_block.2.weight"), + create_test_tensor_snapshot("feature.layers.2.conv_block.0.weight"), + create_test_tensor_snapshot("feature.layers.2.conv_block.2.weight"), + ]; + + let (mapped, transformations) = map_indices_contiguous(tensors); + + // Expected mapping: + // feature.layers: 0, 2 -> 0, 1 + // feature.layers.0.conv_block: 0, 2 -> 0, 1 + // feature.layers.2.conv_block: 0, 2 -> 0, 1 + // + // Result: + // feature.layers.0.conv_block.0.weight -> feature.layers.0.conv_block.0.weight + // feature.layers.0.conv_block.2.weight -> feature.layers.0.conv_block.1.weight + // feature.layers.2.conv_block.0.weight -> feature.layers.1.conv_block.0.weight + // feature.layers.2.conv_block.2.weight -> feature.layers.1.conv_block.1.weight + + assert!( + mapped + .iter() + .any(|v| v.full_path() == "feature.layers.0.conv_block.0.weight"), + "0.0 should stay as 0.0" + ); + assert!( + mapped + .iter() + .any(|v| v.full_path() == "feature.layers.0.conv_block.1.weight"), + "0.2 should become 0.1" + ); + assert!( + mapped + .iter() + .any(|v| v.full_path() == "feature.layers.1.conv_block.0.weight"), + "2.0 should become 1.0" + ); + assert!( + mapped + .iter() + .any(|v| v.full_path() == "feature.layers.1.conv_block.1.weight"), + "2.2 should become 1.1" + ); + + // Verify specific transformations + let t1 = transformations + .iter() + .find(|(_, old)| old == "feature.layers.2.conv_block.2.weight"); + assert_eq!( + t1.map(|(new, _)| new.as_str()), + Some("feature.layers.1.conv_block.1.weight"), + "2.2 should map to 1.1" + ); + } + + #[test] + fn test_map_indices_contiguous_deeply_nested() { + // Test with three levels of nesting + let tensors = vec![ + create_test_tensor_snapshot("a.0.b.0.c.0.weight"), + create_test_tensor_snapshot("a.0.b.0.c.2.weight"), + create_test_tensor_snapshot("a.0.b.2.c.0.weight"), + create_test_tensor_snapshot("a.2.b.0.c.0.weight"), + ]; + + let (mapped, _) = map_indices_contiguous(tensors); + + // a: 0, 2 -> 0, 1 + // a.0.b: 0, 2 -> 0, 1 + // a.2.b: 0 -> 0 + // a.0.b.0.c: 0, 2 -> 0, 1 + // a.0.b.2.c: 0 -> 0 + // a.2.b.0.c: 0 -> 0 + + assert!(mapped.iter().any(|v| v.full_path() == "a.0.b.0.c.0.weight")); + assert!( + mapped.iter().any(|v| v.full_path() == "a.0.b.0.c.1.weight"), + "a.0.b.0.c.2 should become a.0.b.0.c.1" + ); + assert!( + mapped.iter().any(|v| v.full_path() == "a.0.b.1.c.0.weight"), + "a.0.b.2.c.0 should become a.0.b.1.c.0" + ); + assert!( + mapped.iter().any(|v| v.full_path() == "a.1.b.0.c.0.weight"), + "a.2.b.0.c.0 should become a.1.b.0.c.0" + ); + } +} diff --git a/crates/burn-store/src/lib.rs b/crates/burn-store/src/lib.rs new file mode 100644 index 0000000..eea7f57 --- /dev/null +++ b/crates/burn-store/src/lib.rs @@ -0,0 +1,124 @@ +#![cfg_attr(not(feature = "std"), no_std)] + +//! # Burn Store +//! +//! Advanced model storage and serialization infrastructure for the Burn deep learning framework. +//! +//! This crate provides comprehensive functionality for storing and loading Burn modules +//! and their tensor data, with support for cross-framework interoperability, flexible filtering, +//! and efficient memory management through lazy materialization. +//! +//! ## Key Features +//! +//! - **Burnpack Format**: Native Burn format with CBOR metadata, ParamId persistence for stateful training, and no-std support +//! - **SafeTensors Format**: Industry-standard format for secure and efficient tensor serialization +//! - **PyTorch Compatibility**: Load PyTorch models directly into Burn with automatic weight transformation +//! - **Zero-Copy Loading**: Memory-mapped files and lazy tensor materialization for optimal performance +//! - **Flexible Filtering**: Load/save specific model subsets using regex, exact paths, or custom predicates +//! - **Tensor Remapping**: Rename tensors during load/save operations for framework compatibility +//! - **No-std Support**: Core functionality available in embedded and WASM environments +//! +//! ## Quick Start +//! +//! ### Basic Save and Load +//! +//! ```rust,ignore +//! use burn_store::{ModuleSnapshot, SafetensorsStore}; +//! +//! // Save a model +//! let mut store = SafetensorsStore::from_file("model.safetensors"); +//! model.save_into(&mut store)?; +//! +//! // Load a model +//! let mut store = SafetensorsStore::from_file("model.safetensors"); +//! model.load_from(&mut store)?; +//! ``` +//! +//! ### Loading PyTorch Models +//! +//! ```rust,ignore +//! use burn_store::PytorchStore; +//! +//! // Load PyTorch model (automatic weight transformation via PyTorchToBurnAdapter) +//! let mut store = PytorchStore::from_file("pytorch_model.pth") +//! .with_top_level_key("state_dict") // Access nested state dict if needed +//! .allow_partial(true); // Skip unknown tensors +//! +//! model.load_from(&mut store)?; +//! ``` +//! +//! ### Filtering and Remapping +//! +//! ```rust,no_run +//! # use burn_store::SafetensorsStore; +//! // Save only specific layers with renaming +//! let mut store = SafetensorsStore::from_file("encoder.safetensors") +//! .with_regex(r"^encoder\..*") // Filter: only encoder layers +//! .with_key_remapping(r"^encoder\.", "transformer.") // Rename: encoder.X -> transformer.X +//! .metadata("subset", "encoder_only"); +//! +//! // Use store with model.save_into(&mut store)?; +//! ``` +//! +//! ## Core Components +//! +//! - [`ModuleSnapshot`]: Extension trait for Burn modules providing `collect()` and `apply()` methods +//! - [`BurnpackStore`]: Native Burn format with ParamId persistence for stateful training workflows +//! - [`SafetensorsStore`]: Primary storage implementation supporting the SafeTensors format +//! - [`PytorchStore`]: PyTorch model loader supporting .pth and .pt files +//! - [`PathFilter`]: Flexible filtering system for selective tensor loading/saving +//! - [`KeyRemapper`]: Advanced tensor name remapping with regex patterns +//! - [`ModuleAdapter`]: Framework adapters for cross-framework compatibility +//! +//! ## Feature Flags +//! +//! - `std`: Enables file I/O and other std-only features (default) +//! - `safetensors`: Enables SafeTensors format support (default) + +extern crate alloc; + +mod adapter; +mod applier; +mod apply_result; +mod collector; +mod filter; +mod tensor_snapshot; +mod traits; + +pub use adapter::{ + BurnToPyTorchAdapter, ChainAdapter, HalfPrecisionAdapter, IdentityAdapter, ModuleAdapter, + PyTorchToBurnAdapter, +}; +pub use applier::Applier; +pub use apply_result::{ApplyError, ApplyResult}; +pub use collector::Collector; +pub use filter::PathFilter; +pub use tensor_snapshot::{TensorSnapshot, TensorSnapshotError}; +pub use traits::{ModuleSnapshot, ModuleStore}; + +#[cfg(feature = "std")] +mod keyremapper; +#[cfg(feature = "std")] +pub use keyremapper::{KeyRemapper, map_indices_contiguous}; + +/// Serde-based deserialization of nested values, used for importing model weights from external +/// formats (e.g. PyTorch's pickle `.pt` files). +#[cfg(feature = "pytorch")] +pub mod nested; + +#[cfg(feature = "pytorch")] +pub mod pytorch; +#[cfg(feature = "pytorch")] +pub use pytorch::{PytorchStore, PytorchStoreError}; + +#[cfg(feature = "safetensors")] +mod safetensors; +#[cfg(feature = "safetensors")] +pub use safetensors::{SafetensorsStore, SafetensorsStoreError}; + +#[cfg(feature = "burnpack")] +mod bridge; +#[cfg(feature = "burnpack")] +mod burnpack; +#[cfg(feature = "burnpack")] +pub use burnpack::BurnpackStore; diff --git a/crates/burn-store/src/nested/adapter.rs b/crates/burn-store/src/nested/adapter.rs new file mode 100644 index 0000000..b7b3e4f --- /dev/null +++ b/crates/burn-store/src/nested/adapter.rs @@ -0,0 +1,83 @@ +use super::data::NestedValue; + +/// A trait that defines the adapter for a Burn module. +/// +/// This is used to adapt an incoming module to a Burn module. +pub trait BurnModuleAdapter: Sized { + /// Adapts a module. + fn adapt(name: &str, data: NestedValue) -> NestedValue { + match name { + "BatchNorm" => Self::adapt_batch_norm(data), + "Conv1d" => Self::adapt_conv1d(data), + "Conv2d" => Self::adapt_conv2d(data), + "Conv3d" => Self::adapt_conv3d(data), + "ConvTranspose1d" => Self::adapt_conv_transpose_1d(data), + "ConvTranspose2d" => Self::adapt_conv_transpose_2d(data), + "ConvTranspose3d" => Self::adapt_conv_transpose_3d(data), + "Embedding" => Self::adapt_embedding(data), + "GroupNorm" => Self::adapt_group_norm(data), + "LayerNorm" => Self::adapt_layer_norm(data), + "Linear" => Self::adapt_linear(data), + _ => data, + } + } + + /// Adapts a linear module. + fn adapt_linear(data: NestedValue) -> NestedValue { + data + } + + /// Adapts a Convolution 1D module. + fn adapt_conv1d(data: NestedValue) -> NestedValue { + data + } + + /// Adapts a Convolution 2D module. + fn adapt_conv2d(data: NestedValue) -> NestedValue { + data + } + + /// Adapts a Convolution 3D module. + fn adapt_conv3d(data: NestedValue) -> NestedValue { + data + } + + /// Adapts convolution transpose 1D module. + fn adapt_conv_transpose_1d(data: NestedValue) -> NestedValue { + data + } + + /// Adapts convolution transpose 2D module. + fn adapt_conv_transpose_2d(data: NestedValue) -> NestedValue { + data + } + + /// Adapts convolution transpose 2D module. + fn adapt_conv_transpose_3d(data: NestedValue) -> NestedValue { + data + } + + /// Adapts embedding module. + fn adapt_embedding(data: NestedValue) -> NestedValue { + data + } + + /// Adapts group normalization module. + fn adapt_group_norm(data: NestedValue) -> NestedValue { + data + } + + /// Adapts layer normalization module. + fn adapt_layer_norm(data: NestedValue) -> NestedValue { + data + } + + /// Adapts batch normalization module. + fn adapt_batch_norm(data: NestedValue) -> NestedValue { + data + } +} + +/// Default adapter that takes no action. +pub struct DefaultAdapter; +impl BurnModuleAdapter for DefaultAdapter {} diff --git a/crates/burn-store/src/nested/data.rs b/crates/burn-store/src/nested/data.rs new file mode 100644 index 0000000..0067d3a --- /dev/null +++ b/crates/burn-store/src/nested/data.rs @@ -0,0 +1,235 @@ +use std::collections::HashMap; + +use burn_tensor::Bytes; +use core::fmt; +use num_traits::cast::ToPrimitive; + +/// The main data structure used for deserialization. +/// +/// It can hold tree-like structures of nested maps and vectors. +#[derive(Clone)] +pub enum NestedValue { + /// The default value, which actually does not hold any value and it is used to indicate that + /// the value should be populated with the default value. It contains an optional string with + /// the originator field name. + Default(Option), + + /// A boolean value. + Bool(bool), + + /// A string value. + String(String), + + /// Floating point 32-bit value. + F32(f32), + + /// Floating point 64-bit value. + F64(f64), + + /// Signed 16-bit integer value. + I16(i16), + + /// Signed 32-bit integer value. + I32(i32), + + /// Signed 64-bit integer value. + I64(i64), + + /// Unsigned 8-bit integer value. + U8(u8), + + /// Unsigned 16-bit integer value used for bf16 and f16 serialization + U16(u16), + + /// Unsigned 64-bit integer value. + U64(u64), + + /// A map of nested values (typically used for structs) + Map(HashMap), + + /// A vector of nested values (typically used for vector of structs or numbers) + Vec(Vec), + + /// A vector of 8-bit unsigned integer values. + U8s(Vec), + + /// A vector of 16-bit unsigned integer values. + U16s(Vec), + + /// A vector of 32-bit floating point values. + F32s(Vec), + + /// An opaque vector of bytes, with alignment. + Bytes(Bytes), +} + +impl NestedValue { + /// Get the nested value as a map. + pub fn as_map(self) -> Option> { + match self { + NestedValue::Map(map) => Some(map), + _ => None, + } + } + + /// Get the nested value as a boolean. + pub fn as_bool(self) -> Option { + match self { + NestedValue::Bool(bool) => Some(bool), + _ => None, + } + } + + /// Get the nested value as a string. + pub fn as_string(self) -> Option { + match self { + NestedValue::String(string) => Some(string), + _ => None, + } + } + + /// Get the nested value as a f32. + pub fn as_f32(self) -> Option { + match self { + NestedValue::F32(f32) => Some(f32), + NestedValue::F64(f) => f.to_f32(), + _ => None, + } + } + + /// Get the nested value as a f64. + pub fn as_f64(self) -> Option { + match self { + NestedValue::F64(f64) => Some(f64), + NestedValue::F32(f) => f.to_f64(), + _ => None, + } + } + + /// Get the nested value as an i16. + pub fn as_i16(self) -> Option { + match self { + NestedValue::I16(i16) => Some(i16), + NestedValue::I32(i) => i.to_i16(), + NestedValue::I64(i) => i.to_i16(), + NestedValue::U16(u) => u.to_i16(), + NestedValue::U64(u) => u.to_i16(), + _ => None, + } + } + + /// Get the nested value as an i32. + pub fn as_i32(self) -> Option { + match self { + NestedValue::I32(i32) => Some(i32), + NestedValue::I16(i) => i.to_i32(), + NestedValue::I64(i) => i.to_i32(), + NestedValue::U16(u) => u.to_i32(), + NestedValue::U64(u) => u.to_i32(), + _ => None, + } + } + + /// Get the nested value as an i64. + pub fn as_i64(self) -> Option { + match self { + NestedValue::I64(i64) => Some(i64), + NestedValue::I16(i) => i.to_i64(), + NestedValue::I32(i) => i.to_i64(), + NestedValue::U16(u) => u.to_i64(), + NestedValue::U64(u) => u.to_i64(), + _ => None, + } + } + + /// Get the nested value as a u8. + pub fn as_u8(self) -> Option { + match self { + NestedValue::U8(u8) => Some(u8), + NestedValue::I16(i) => i.to_u8(), + NestedValue::I32(i) => i.to_u8(), + NestedValue::I64(i) => i.to_u8(), + NestedValue::U16(u) => u.to_u8(), + NestedValue::U64(u) => u.to_u8(), + _ => None, + } + } + + /// Get the nested value as a u16. + pub fn as_u16(self) -> Option { + match self { + NestedValue::U16(u16) => Some(u16), + NestedValue::I16(i) => i.to_u16(), + NestedValue::I32(i) => i.to_u16(), + NestedValue::I64(i) => i.to_u16(), + NestedValue::U64(u) => u.to_u16(), + _ => None, + } + } + + /// Get the nested value as a u64. + pub fn as_u64(self) -> Option { + match self { + NestedValue::U64(u64) => Some(u64), + NestedValue::I16(i) => i.to_u64(), + NestedValue::I32(i) => i.to_u64(), + NestedValue::I64(i) => i.to_u64(), + NestedValue::U16(u) => u.to_u64(), + _ => None, + } + } + + /// Get the nested value as a vector of bytes. + pub fn as_bytes(self) -> Option { + match self { + NestedValue::Bytes(u) => Some(u), + NestedValue::U8s(u) => Some(Bytes::from_elems(u)), + _ => None, + } + } +} + +fn write_vec_truncated( + vec: &[T], + f: &mut core::fmt::Formatter, +) -> fmt::Result { + write!(f, "Vec([")?; + for (i, v) in vec.iter().take(3).enumerate() { + if i > 0 { + write!(f, ", ")?; + } + write!(f, "{v:?}")?; + } + write!(f, ", ...] len={})", vec.len()) +} + +impl fmt::Debug for NestedValue { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + // Truncate values for vector + NestedValue::Vec(vec) if vec.len() > 3 => write_vec_truncated(vec, f), + NestedValue::U8s(vec) if vec.len() > 3 => write_vec_truncated(vec, f), + NestedValue::U16s(vec) if vec.len() > 3 => write_vec_truncated(vec, f), + NestedValue::F32s(vec) if vec.len() > 3 => write_vec_truncated(vec, f), + NestedValue::Bytes(bytes) if bytes.len() > 3 => write_vec_truncated(bytes, f), + // Handle other variants as usual + NestedValue::Default(origin) => f.debug_tuple("Default").field(origin).finish(), + NestedValue::Bool(b) => f.debug_tuple("Bool").field(b).finish(), + NestedValue::String(s) => f.debug_tuple("String").field(s).finish(), + NestedValue::F32(val) => f.debug_tuple("F32").field(val).finish(), + NestedValue::F64(val) => f.debug_tuple("F64").field(val).finish(), + NestedValue::I16(val) => f.debug_tuple("I16").field(val).finish(), + NestedValue::I32(val) => f.debug_tuple("I32").field(val).finish(), + NestedValue::I64(val) => f.debug_tuple("I64").field(val).finish(), + NestedValue::U8(val) => f.debug_tuple("U8").field(val).finish(), + NestedValue::U16(val) => f.debug_tuple("U16").field(val).finish(), + NestedValue::U64(val) => f.debug_tuple("U64").field(val).finish(), + NestedValue::Map(map) => f.debug_map().entries(map.iter()).finish(), + NestedValue::Vec(vec) => f.debug_list().entries(vec.iter()).finish(), + NestedValue::U8s(vec) => f.debug_list().entries(vec.iter()).finish(), + NestedValue::U16s(vec) => f.debug_list().entries(vec.iter()).finish(), + NestedValue::F32s(vec) => f.debug_list().entries(vec.iter()).finish(), + NestedValue::Bytes(bytes) => f.debug_list().entries(bytes.iter()).finish(), + } + } +} diff --git a/crates/burn-store/src/nested/de.rs b/crates/burn-store/src/nested/de.rs new file mode 100644 index 0000000..383b1f0 --- /dev/null +++ b/crates/burn-store/src/nested/de.rs @@ -0,0 +1,1006 @@ +use core::ptr; +use std::collections::HashMap; + +use super::data::NestedValue; +use super::{adapter::BurnModuleAdapter, error::Error}; + +use serde::de::{EnumAccess, VariantAccess}; +use serde::{ + de::{self, DeserializeSeed, IntoDeserializer, MapAccess, SeqAccess, Visitor}, + forward_to_deserialize_any, +}; + +const RECORD_ITEM_SUFFIX: &str = "RecordItem"; + +/// A deserializer for the nested value data structure. +pub struct Deserializer { + // This string starts with the input data and characters are truncated off + // the beginning as data is parsed. + value: Option, + default_for_missing_fields: bool, + phantom: std::marker::PhantomData, +} + +impl Deserializer { + /// Creates a new deserializer with the given nested value. + /// + /// # Arguments + /// + /// * `value` - A nested value. + /// * `default_for_missing_fields` - A boolean indicating whether to add missing fields with default value. + pub fn new(value: NestedValue, default_for_missing_fields: bool) -> Self { + Self { + value: Some(value), + default_for_missing_fields, + phantom: std::marker::PhantomData, + } + } +} + +impl<'de, A: BurnModuleAdapter> serde::Deserializer<'de> for Deserializer { + type Error = Error; + + fn deserialize_any(self, _visitor: V) -> Result + where + V: Visitor<'de>, + { + unimplemented!("deserialize_any is not implemented") + } + + fn deserialize_struct( + self, + name: &'static str, + fields: &'static [&'static str], + visitor: V, + ) -> Result + where + V: Visitor<'de>, + { + let value = match self.value { + Some(value) => { + // Adapt modules + if let Some(name) = name.strip_suffix(RECORD_ITEM_SUFFIX) { + A::adapt(name, value) + } else { + value + } + } + None => { + return Err(de::Error::custom(format!( + "Expected some value but got {:?}", + self.value + ))); + } + }; + + match value { + NestedValue::Map(map) => { + // Add missing fields into the map with default value if needed. + let map = if self.default_for_missing_fields { + let mut map = map; + for field in fields.iter().map(|s| s.to_string()) { + map.entry(field.clone()) + .or_insert(NestedValue::Default(Some(field))); + } + map + } else { + map + }; + + visitor.visit_map(HashMapAccess::::new( + map, + self.default_for_missing_fields, + )) + } + + _ => Err(de::Error::custom(format!( + "Expected struct but got {value:?}" + ))), + } + } + + fn deserialize_string(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_string(self.value.unwrap().as_string().unwrap().to_string()) + } + + fn deserialize_ignored_any(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_unit() + } + + fn deserialize_map(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + match self.value { + Some(NestedValue::Map(map)) => visitor.visit_map(HashMapAccess::::new( + map, + self.default_for_missing_fields, + )), + + _ => Err(de::Error::custom(format!( + "Expected map value but got {:?}", + self.value + ))), + } + } + + fn deserialize_bool(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_bool(self.value.unwrap().as_bool().unwrap()) + } + + fn deserialize_i8(self, _visitor: V) -> Result + where + V: Visitor<'de>, + { + unimplemented!("deserialize_i8 is not implemented") + } + + fn deserialize_i16(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_i16(self.value.unwrap().as_i16().unwrap().to_owned()) + } + + fn deserialize_i32(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_i32(self.value.unwrap().as_i32().unwrap().to_owned()) + } + + fn deserialize_i64(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_i64(self.value.unwrap().as_i64().unwrap().to_owned()) + } + + fn deserialize_u8(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_u8(self.value.unwrap().as_u8().unwrap().to_owned()) + } + + fn deserialize_u16(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_u16(self.value.unwrap().as_u16().unwrap().to_owned()) + } + + fn deserialize_u32(self, _visitor: V) -> Result + where + V: Visitor<'de>, + { + unimplemented!("deserialize_u32 is not implemented") + } + + fn deserialize_u64(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_u64(self.value.unwrap().as_u64().unwrap().to_owned()) + } + + fn deserialize_f32(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_f32(self.value.unwrap().as_f32().unwrap().to_owned()) + } + + fn deserialize_f64(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_f64(self.value.unwrap().as_f64().unwrap().to_owned()) + } + + fn deserialize_char(self, _visitor: V) -> Result + where + V: Visitor<'de>, + { + unimplemented!("deserialize_char is not implemented") + } + + fn deserialize_str(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_str(self.value.unwrap().as_string().unwrap().as_ref()) + } + + fn deserialize_bytes(self, _visitor: V) -> Result + where + V: Visitor<'de>, + { + unimplemented!("deserialize_bytes is not implemented") + } + + fn deserialize_byte_buf(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + let bytes = self.value.unwrap().as_bytes().unwrap(); + match bytes.try_into_vec::() { + Ok(bytes) => visitor.visit_byte_buf(bytes), + Err(bytes) => visitor.visit_bytes(&bytes), + } + } + + fn deserialize_option(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + if let Some(value) = self.value { + visitor.visit_some(Deserializer::::new( + value, + self.default_for_missing_fields, + )) + } else { + visitor.visit_none() + } + } + + fn deserialize_unit(self, _visitor: V) -> Result + where + V: Visitor<'de>, + { + unimplemented!("deserialize_unit is not implemented") + } + + fn deserialize_unit_struct( + self, + _name: &'static str, + _visitor: V, + ) -> Result + where + V: Visitor<'de>, + { + unimplemented!("deserialize_unit_struct is not implemented") + } + + fn deserialize_newtype_struct( + self, + _name: &'static str, + visitor: V, + ) -> Result + where + V: Visitor<'de>, + { + visitor.visit_newtype_struct(Deserializer::::new( + self.value.unwrap(), + self.default_for_missing_fields, + )) + } + + fn deserialize_seq(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + if let Some(value) = self.value { + match value { + NestedValue::Vec(_) => visitor.visit_seq(VecSeqAccess::::new( + value, + self.default_for_missing_fields, + )), + NestedValue::U8s(_) => visitor.visit_seq(VecSeqAccess::::new( + value, + self.default_for_missing_fields, + )), + NestedValue::U16s(_) => visitor.visit_seq(VecSeqAccess::::new( + value, + self.default_for_missing_fields, + )), + NestedValue::F32s(_) => visitor.visit_seq(VecSeqAccess::::new( + value, + self.default_for_missing_fields, + )), + _ => Err(de::Error::custom(format!("Expected Vec but got {value:?}"))), + } + } else { + Err(de::Error::custom("Expected Vec but got None")) + } + } + + fn deserialize_tuple(self, _len: usize, _visitor: V) -> Result + where + V: Visitor<'de>, + { + unimplemented!("deserialize_tuple is not implemented") + } + + fn deserialize_tuple_struct( + self, + _name: &'static str, + _len: usize, + _visitor: V, + ) -> Result + where + V: Visitor<'de>, + { + unimplemented!("deserialize_tuple_struct is not implemented") + } + + /// Deserializes an enum by attempting to match its variants against the provided data. + /// + /// This function attempts to deserialize an enum by iterating over its possible variants + /// and trying to deserialize the data into each until one succeeds. We need to do this + /// because we don't have a way to know which variant to deserialize from the data. + /// + /// This is similar to Serde's + /// [untagged enum deserialization](https://serde.rs/enum-representations.html#untagged), + /// but it's on the deserializer side. Using `#[serde(untagged)]` on the enum will force + /// using `deserialize_any`, which is not what we want because we want to use methods, such + /// as `visit_struct`. Also we do not wish to use auto generate code for Deserialize just + /// for enums because it will affect other serialization and deserialization, such + /// as JSON and Bincode. + /// + /// # Safety + /// The function uses an unsafe block to clone the `visitor`. This is necessary because + /// the `Visitor` trait does not have a `Clone` implementation, and we need to clone it + /// as we are going to use it multiple times. The Visitor is a code generated unit struct + /// with no states or mutations, so it is safe to clone it in this case. We mainly care + /// about the `visit_enum` method, which is the only method that will be called on the + /// cloned visitor. + fn deserialize_enum( + self, + _name: &'static str, + variants: &'static [&'static str], + visitor: V, + ) -> Result + where + V: Visitor<'de>, + { + fn clone_unsafely(thing: &T) -> T { + unsafe { + // Allocate memory for the clone. + let mut clone = std::mem::MaybeUninit::::uninit(); + // Get a mutable pointer to the allocated memory. + let clone_ptr = clone.as_mut_ptr(); + // Copy the memory + ptr::copy_nonoverlapping(thing as *const T, clone_ptr, 1); + // Assume the cloned data is initialized and convert it to an owned instance of T. + clone.assume_init() + } + } + + // Try each variant in order + for &variant in variants { + // clone visitor to avoid moving it + let cloned_visitor = clone_unsafely(&visitor); + let result = cloned_visitor.visit_enum(ProbeEnumAccess::::new( + self.value.clone().unwrap(), + variant.to_owned(), + self.default_for_missing_fields, + )); + + if result.is_ok() { + return result; + } + } + + Err(de::Error::custom("No variant match")) + } + + fn deserialize_identifier(self, _visitor: V) -> Result + where + V: Visitor<'de>, + { + unimplemented!("deserialize_identifier is not implemented") + } +} + +/// A sequence access for a vector in the nested value data structure. +struct VecSeqAccess { + iter: Box>, + default_for_missing_fields: bool, + phantom: std::marker::PhantomData, +} + +// Concrete implementation for `Vec` +impl VecSeqAccess { + fn new(vec: NestedValue, default_for_missing_fields: bool) -> Self { + match vec { + NestedValue::Vec(v) => VecSeqAccess { + iter: Box::new(v.into_iter()), + default_for_missing_fields, + phantom: std::marker::PhantomData, + }, + _ => panic!("Invalid vec sequence"), + } + } +} + +// Concrete implementation for `Vec` +impl VecSeqAccess { + fn new(vec: NestedValue, default_for_missing_fields: bool) -> Self { + match vec { + NestedValue::U8s(v) => VecSeqAccess { + iter: Box::new(v.into_iter()), + default_for_missing_fields, + phantom: std::marker::PhantomData, + }, + _ => panic!("Invalid vec sequence"), + } + } +} + +// Concrete implementation for `Vec` +impl VecSeqAccess { + fn new(vec: NestedValue, default_for_missing_fields: bool) -> Self { + match vec { + NestedValue::U16s(v) => VecSeqAccess { + iter: Box::new(v.into_iter()), + default_for_missing_fields, + phantom: std::marker::PhantomData, + }, + _ => panic!("Invalid vec sequence"), + } + } +} + +// Concrete implementation for `Vec` +impl VecSeqAccess { + fn new(vec: NestedValue, default_for_missing_fields: bool) -> Self { + match vec { + NestedValue::F32s(v) => VecSeqAccess { + iter: Box::new(v.into_iter()), + default_for_missing_fields, + phantom: std::marker::PhantomData, + }, + _ => panic!("Invalid vec sequence"), + } + } +} + +// Concrete implementation for `Vec` +impl<'de, A> SeqAccess<'de> for VecSeqAccess +where + NestedValueWrapper: IntoDeserializer<'de, Error>, + A: BurnModuleAdapter, +{ + type Error = Error; + + fn next_element_seed(&mut self, seed: T) -> Result, Self::Error> + where + T: DeserializeSeed<'de>, + { + let item = match self.iter.next() { + Some(v) => v, + None => return Ok(None), + }; + + seed.deserialize( + NestedValueWrapper::::new(item, self.default_for_missing_fields).into_deserializer(), + ) + .map(Some) + } +} + +// Concrete implementation for `Vec` +impl<'de, A> SeqAccess<'de> for VecSeqAccess +where + NestedValueWrapper: IntoDeserializer<'de, Error>, + A: BurnModuleAdapter, +{ + type Error = Error; + + fn next_element_seed(&mut self, seed: T) -> Result, Self::Error> + where + T: DeserializeSeed<'de>, + { + let item = match self.iter.next() { + Some(v) => v, + None => return Ok(None), + }; + + seed.deserialize( + NestedValueWrapper::::new(NestedValue::U8(item), self.default_for_missing_fields) + .into_deserializer(), + ) + .map(Some) + } +} + +// Concrete implementation for `Vec` +impl<'de, A> SeqAccess<'de> for VecSeqAccess +where + NestedValueWrapper: IntoDeserializer<'de, Error>, + A: BurnModuleAdapter, +{ + type Error = Error; + + fn next_element_seed(&mut self, seed: T) -> Result, Self::Error> + where + T: DeserializeSeed<'de>, + { + let item = match self.iter.next() { + Some(v) => v, + None => return Ok(None), + }; + + seed.deserialize( + NestedValueWrapper::::new(NestedValue::U16(item), self.default_for_missing_fields) + .into_deserializer(), + ) + .map(Some) + } +} + +// Concrete implementation for `Vec` +impl<'de, A> SeqAccess<'de> for VecSeqAccess +where + NestedValueWrapper: IntoDeserializer<'de, Error>, + A: BurnModuleAdapter, +{ + type Error = Error; + + fn next_element_seed(&mut self, seed: T) -> Result, Self::Error> + where + T: DeserializeSeed<'de>, + { + let item = match self.iter.next() { + Some(v) => v, + None => return Ok(None), + }; + + seed.deserialize( + NestedValueWrapper::::new(NestedValue::F32(item), self.default_for_missing_fields) + .into_deserializer(), + ) + .map(Some) + } +} + +/// A map access for a map in the nested value data structure. +struct HashMapAccess { + iter: std::collections::hash_map::IntoIter, + next_value: Option, + default_for_missing_fields: bool, + phantom: std::marker::PhantomData, +} + +impl HashMapAccess { + fn new(map: HashMap, default_for_missing_fields: bool) -> Self { + HashMapAccess { + iter: map.into_iter(), + next_value: None, + default_for_missing_fields, + phantom: std::marker::PhantomData, + } + } +} + +impl<'de, A> MapAccess<'de> for HashMapAccess +where + String: IntoDeserializer<'de, Error>, + NestedValueWrapper: IntoDeserializer<'de, Error>, + A: BurnModuleAdapter, +{ + type Error = Error; + + fn next_key_seed(&mut self, seed: T) -> Result, Self::Error> + where + T: DeserializeSeed<'de>, + { + match self.iter.next() { + Some((k, v)) => { + // Keep the value for the next call to next_value_seed. + self.next_value = Some(v); + // Deserialize the key. + seed.deserialize(k.into_deserializer()).map(Some) + } + None => Ok(None), + } + } + + fn next_value_seed(&mut self, seed: T) -> Result + where + T: DeserializeSeed<'de>, + { + match self.next_value.take() { + Some(NestedValue::Default(originator)) => { + seed.deserialize(DefaultDeserializer::new(originator)) + } + Some(v) => seed.deserialize( + NestedValueWrapper::new(v, self.default_for_missing_fields).into_deserializer(), + ), + None => seed.deserialize(DefaultDeserializer::new(None)), + } + } +} + +struct ProbeEnumAccess { + value: NestedValue, + current_variant: String, + default_for_missing_fields: bool, + phantom: std::marker::PhantomData, +} + +impl ProbeEnumAccess { + fn new(value: NestedValue, current_variant: String, default_for_missing_fields: bool) -> Self { + ProbeEnumAccess { + value, + current_variant, + default_for_missing_fields, + phantom: std::marker::PhantomData, + } + } +} + +impl<'de, A> EnumAccess<'de> for ProbeEnumAccess +where + A: BurnModuleAdapter, +{ + type Error = Error; + type Variant = Self; + + fn variant_seed(self, seed: V) -> Result<(V::Value, Self::Variant), Self::Error> + where + V: DeserializeSeed<'de>, + { + seed.deserialize(self.current_variant.clone().into_deserializer()) + .map(|v| (v, self)) + } +} + +impl<'de, A> VariantAccess<'de> for ProbeEnumAccess +where + A: BurnModuleAdapter, +{ + type Error = Error; + + fn newtype_variant_seed(self, seed: T) -> Result + where + T: DeserializeSeed<'de>, + { + let value = seed.deserialize( + NestedValueWrapper::::new(self.value, self.default_for_missing_fields) + .into_deserializer(), + )?; + Ok(value) + } + + fn unit_variant(self) -> Result<(), Self::Error> { + // Support tensor `DType` deserialization + match self.value { + NestedValue::Map(value) if value.contains_key("DType") => { + match value.get("DType") { + Some(NestedValue::String(variant)) => { + if *variant == self.current_variant { + Ok(()) + } else { + Err(Error::Other("Wrong variant".to_string())) // wrong match + } + } + _ => panic!("expected DType variant as string"), + } + } + _ => unimplemented!( + "unit variant is not implemented because it is not used in the burn module" + ), + } + } + + fn tuple_variant(self, _len: usize, _visitor: V) -> Result + where + V: Visitor<'de>, + { + unimplemented!("tuple variant is not implemented because it is not used in the burn module") + } + + fn struct_variant( + self, + _fields: &'static [&'static str], + _visitor: V, + ) -> Result + where + V: Visitor<'de>, + { + unimplemented!( + "struct variant is not implemented because it is not used in the burn module" + ) + } +} + +/// A wrapper for the nested value data structure with a burn module adapter. +struct NestedValueWrapper { + value: NestedValue, + default_for_missing_fields: bool, + phantom: std::marker::PhantomData, +} + +impl NestedValueWrapper { + fn new(value: NestedValue, default_for_missing_fields: bool) -> Self { + Self { + value, + default_for_missing_fields, + phantom: std::marker::PhantomData, + } + } +} + +impl IntoDeserializer<'_, Error> for NestedValueWrapper { + type Deserializer = Deserializer; + + fn into_deserializer(self) -> Self::Deserializer { + Deserializer::::new(self.value, self.default_for_missing_fields) + } +} + +/// A default deserializer that always returns the default value. +struct DefaultDeserializer { + /// The originator field name (the top-level missing field name) + originator_field_name: Option, +} + +impl DefaultDeserializer { + fn new(originator_field_name: Option) -> Self { + Self { + originator_field_name, + } + } +} + +impl<'de> serde::Deserializer<'de> for DefaultDeserializer { + type Error = Error; + + fn deserialize_any(self, _visitor: V) -> Result + where + V: Visitor<'de>, + { + unimplemented!() + } + + fn deserialize_i32(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_i32(Default::default()) + } + + fn deserialize_f32(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_f32(Default::default()) + } + + fn deserialize_i16(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_i16(Default::default()) + } + + fn deserialize_i64(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_i64(Default::default()) + } + + fn deserialize_u16(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_u16(Default::default()) + } + + fn deserialize_u64(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_u64(Default::default()) + } + + fn deserialize_f64(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_f64(Default::default()) + } + + fn deserialize_bool(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_bool(Default::default()) + } + + fn deserialize_char(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_char(Default::default()) + } + + fn deserialize_str(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_str(Default::default()) + } + + fn deserialize_i8(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_i8(Default::default()) + } + + fn deserialize_u8(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_u8(Default::default()) + } + + fn deserialize_u32(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_u32(Default::default()) + } + + fn deserialize_option(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_none() + } + + fn deserialize_seq(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_seq(DefaultSeqAccess::new(None)) + } + + fn deserialize_string(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_string(Default::default()) + } + + fn deserialize_struct( + self, + name: &'static str, + _fields: &'static [&'static str], + _visitor: V, + ) -> Result + where + V: Visitor<'de>, + { + // Return an error if the originator field name is not set + Err(Error::Other(format!( + "Missing source values for the '{}' field of type '{}'. Please verify the source data and ensure the field name is correct", + self.originator_field_name.unwrap_or("UNKNOWN".to_string()), + name, + ))) + } + + fn deserialize_tuple_struct( + self, + _name: &'static str, + len: usize, + visitor: V, + ) -> Result + where + V: Visitor<'de>, + { + visitor.visit_seq(DefaultSeqAccess::new(Some(len))) + } + + fn deserialize_tuple(self, len: usize, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_seq(DefaultSeqAccess::new(Some(len))) + } + + fn deserialize_map(self, visitor: V) -> Result + where + V: Visitor<'de>, + { + visitor.visit_map(DefaultMapAccess::new()) + } + + forward_to_deserialize_any! { + u128 bytes byte_buf unit unit_struct newtype_struct + enum identifier ignored_any + } +} + +/// A default sequence access that always returns None (empty sequence). +pub struct DefaultSeqAccess { + size: Option, +} + +impl Default for DefaultSeqAccess { + fn default() -> Self { + Self::new(None) + } +} + +impl DefaultSeqAccess { + /// Creates a new default sequence access with the given size hint. + pub fn new(size: Option) -> Self { + DefaultSeqAccess { size } + } +} + +impl<'de> SeqAccess<'de> for DefaultSeqAccess { + type Error = Error; + + fn next_element_seed(&mut self, seed: T) -> Result, Self::Error> + where + T: DeserializeSeed<'de>, + { + match self.size { + Some(0) => Ok(None), + Some(ref mut size) => { + *size -= 1; + seed.deserialize(DefaultDeserializer::new(None)).map(Some) + } + None => Ok(None), + } + } + + fn size_hint(&self) -> Option { + self.size + } +} + +/// A default map access that always returns None (empty map). +pub struct DefaultMapAccess; + +impl Default for DefaultMapAccess { + fn default() -> Self { + Self::new() + } +} + +impl DefaultMapAccess { + /// Creates a new default map access. + pub fn new() -> Self { + DefaultMapAccess + } +} + +impl<'de> MapAccess<'de> for DefaultMapAccess { + type Error = Error; + + fn next_key_seed(&mut self, _seed: T) -> Result, Self::Error> + where + T: DeserializeSeed<'de>, + { + // Since this is a default implementation, we'll just return None. + Ok(None) + } + + fn next_value_seed(&mut self, _seed: T) -> Result + where + T: DeserializeSeed<'de>, + { + unimplemented!("This should never be called since next_key_seed always returns None") + } + + fn size_hint(&self) -> Option { + // Since this is a default implementation, we'll just return None. + None + } +} diff --git a/crates/burn-store/src/nested/error.rs b/crates/burn-store/src/nested/error.rs new file mode 100644 index 0000000..6b645f3 --- /dev/null +++ b/crates/burn-store/src/nested/error.rs @@ -0,0 +1,31 @@ +/// The error type for nested-value serde. +#[derive(thiserror::Error, Debug)] +pub enum Error { + /// Failed to deserialize. + #[error("failed to deserialize: {0}")] + Deserialize(#[from] serde::de::value::Error), + + /// Failed to serialize. + #[error("failed to serialize")] + Serialize(String), + + /// Encountered an invalid state. + #[error("invalid state")] + InvalidState, + + /// Other error. + #[error("other error: {0}")] + Other(String), +} + +impl serde::de::Error for Error { + fn custom(msg: T) -> Self { + Error::Deserialize(serde::de::value::Error::custom(msg.to_string())) + } +} + +impl serde::ser::Error for Error { + fn custom(msg: T) -> Self { + Error::Serialize(msg.to_string()) + } +} diff --git a/crates/burn-store/src/nested/mod.rs b/crates/burn-store/src/nested/mod.rs new file mode 100644 index 0000000..de363d1 --- /dev/null +++ b/crates/burn-store/src/nested/mod.rs @@ -0,0 +1,14 @@ +//! Serde-based deserialization of nested values, used for importing model weights from external +//! formats (e.g. PyTorch's pickle `.pt` files) into Burn modules via `burn-store`. + +/// The adapter trait that is used to convert the nested value to the module type. +pub mod adapter; + +/// The main data structure used for deserialization. +pub mod data; + +/// The deserializer that converts a nested value into a typed item. +pub mod de; + +/// Error types. +pub mod error; diff --git a/crates/burn-store/src/pytorch/lazy_data.rs b/crates/burn-store/src/pytorch/lazy_data.rs new file mode 100644 index 0000000..30a2260 --- /dev/null +++ b/crates/burn-store/src/pytorch/lazy_data.rs @@ -0,0 +1,570 @@ +//! Lazy data loading support for PyTorch files. +//! +//! This module provides abstractions for lazy loading of tensor data from PyTorch files, +//! avoiding the need to load all data into memory upfront. + +use alloc::string::String; +use alloc::vec::Vec; +use std::collections::HashMap; +use std::fs::File; +use std::io::{BufReader, Read, Seek}; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, Mutex, RwLock}; +use zip::ZipArchive; + +/// A data source that can lazily load tensor data. +#[derive(Clone)] +pub enum LazyDataSource { + /// ZIP archive with lazy loading + Zip(Arc>), + /// TAR archive format (older torchvision models) + Tar(Arc>), + /// Legacy format with multiple storages in single blob + LegacyMultiStorage(Arc>), +} + +/// ZIP archive source for lazy loading +pub struct ZipSource { + path: PathBuf, + // Cache the file list to avoid reopening archive repeatedly + file_list: Vec<(String, u64, u64)>, // (name, offset, compressed_size) +} + +/// TAR archive source for lazy loading (older torchvision models like AlexNet, SqueezeNet) +/// +/// Older PyTorch/torchvision models (pre-1.6) use TAR format instead of ZIP. +/// The TAR archive contains: +/// - `sys_info`: System info pickle (endianness, type sizes) +/// - `pickle`: OrderedDict mapping tensor names to storage keys +/// - `tensors`: Tensor metadata pickles (unused, metadata is embedded in pickle) +/// - `storages`: Storage count + sequential (metadata pickle, element count, raw data) +pub struct TarSource { + /// Cached storage map: storage_key -> (offset_in_storages, size_bytes) + storage_map: HashMap, + /// The raw storages data (kept in memory for TAR format) + storages_data: Vec, +} + +/// Legacy multi-storage source for old PyTorch format (0.1.10 - 1.5) +/// +/// Legacy format stores tensor data as concatenated raw binary without explicit +/// storage boundaries. This source tracks storage usage during tensor parsing +/// to build a storage map for lazy loading. +/// +/// ## Storage Layout +/// - Pickle metadata with tensor definitions +/// - List of storage keys (determines concatenation order) +/// - Raw binary blob with all storages concatenated +pub struct LegacyMultiStorageSource { + path: PathBuf, + data_offset: u64, + #[allow(dead_code)] + data_size: u64, + // Map of storage_key -> (offset_in_blob, size) + storage_map: RwLock>>, + // Storage keys in order (for boundary calculation) + storage_keys: RwLock>>, + // Track storage usage as tensors are accessed + storage_usage: RwLock>, // key -> max_bytes_needed +} + +impl ZipSource { + /// Create a new ZIP source + pub fn new(path: PathBuf) -> std::io::Result { + let file = File::open(&path)?; + let reader = BufReader::new(file); + let mut archive = ZipArchive::new(reader)?; + + // Cache file metadata + let mut file_list = Vec::new(); + for i in 0..archive.len() { + let file = archive.by_index(i)?; + let name = file.name().to_string(); + let offset = file.data_start(); + let compressed_size = file.compressed_size(); + file_list.push(( + name, + offset.expect("should have an offset"), + compressed_size, + )); + } + + Ok(Self { path, file_list }) + } + + /// Check if a file exists in the archive + pub fn contains(&self, name: &str) -> bool { + self.file_list.iter().any(|(n, _, _)| n == name) + } + + /// Get list of data files (excluding pickle files) + pub fn data_files(&self) -> Vec { + self.file_list + .iter() + .filter(|(name, _, _)| name.starts_with("data/") || name.contains("/data/")) + .filter(|(name, _, _)| !name.ends_with(".pkl") && !name.ends_with("/")) + .map(|(name, _, _)| name.clone()) + .collect() + } + + /// Read a specific file from the archive + pub fn read_file(&self, name: &str) -> std::io::Result> { + let file = File::open(&self.path)?; + let reader = BufReader::new(file); + let mut archive = ZipArchive::new(reader)?; + + let mut file = archive.by_name(name)?; + let mut contents = Vec::with_capacity(file.size() as usize); + file.read_to_end(&mut contents)?; + Ok(contents) + } + + /// Read a portion of a file + pub fn read_file_range( + &self, + name: &str, + offset: usize, + length: usize, + ) -> std::io::Result> { + let file = File::open(&self.path)?; + let reader = BufReader::new(file); + let mut archive = ZipArchive::new(reader)?; + + let mut file = archive.by_name(name)?; + let mut buffer = vec![0u8; length]; + + // Skip to offset + let mut skip_buffer = vec![0u8; offset.min(8192)]; + let mut skipped = 0; + while skipped < offset { + let to_skip = (offset - skipped).min(skip_buffer.len()); + file.read_exact(&mut skip_buffer[..to_skip])?; + skipped += to_skip; + } + + // Read the requested data + file.read_exact(&mut buffer)?; + Ok(buffer) + } +} + +impl LegacyMultiStorageSource { + /// Create a new legacy multi-storage source + pub fn new(path: PathBuf, data_offset: u64, data_size: u64) -> Self { + Self { + path, + data_offset, + data_size, + storage_map: RwLock::new(None), + storage_keys: RwLock::new(None), + storage_usage: RwLock::new(HashMap::new()), + } + } + + /// Set the ordered storage keys from the pickle + pub fn set_storage_keys(&self, keys: Vec) { + let mut storage_keys = self + .storage_keys + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + *storage_keys = Some(keys); + } + + /// Track storage usage from tensor access + /// This is called from within tensor loading closures + pub fn track_storage_usage(&self, storage_key: &str, offset: usize, size: usize) { + let mut usage = self + .storage_usage + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let max_extent = offset + size; + usage + .entry(storage_key.to_string()) + .and_modify(|current| *current = (*current).max(max_extent)) + .or_insert(max_extent); + + // Try to build storage map if we have enough information + drop(usage); + self.try_build_storage_map(); + } + + /// Try to build the storage map from tracked usage + fn try_build_storage_map(&self) { + // Only build if we don't already have a map + if self + .storage_map + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .is_some() + { + return; + } + + // Check if we have storage keys + let keys_guard = self + .storage_keys + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(ref keys) = *keys_guard { + let usage = self + .storage_usage + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + + // Only build if we have usage info for all storages + if keys.iter().all(|k| usage.contains_key(k)) { + let mut map = HashMap::new(); + let mut current_offset = 0u64; + + for key in keys { + if let Some(&size) = usage.get(key) { + // Each storage in the binary section is preceded by an 8 byte + // element count (u64 little-endian). Skip it by adding 8 to the + // data offset, and include it in the stride to the next storage. + map.insert(key.clone(), (current_offset + 8, size as u64)); + current_offset += 8 + size as u64; + } + } + + // Set the storage map + drop(keys_guard); + drop(usage); + let mut storage_map = self + .storage_map + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + *storage_map = Some(map); + } + } + } + + /// Read data for a specific storage key + /// Only loads the specific storage portion, never the entire blob + pub fn read(&self, key: &str) -> std::io::Result> { + // Extract numeric key from paths like "data/0" or just "0" + let storage_key = key.split('/').next_back().unwrap_or(key); + + // Get storage map - must be available for lazy loading to work + let storage_map = self + .storage_map + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + + if let Some(ref map) = *storage_map + && let Some(&(offset, size)) = map.get(storage_key) + { + // Load only this specific storage + let mut file = File::open(&self.path)?; + file.seek(std::io::SeekFrom::Start(self.data_offset + offset))?; + + let mut buffer = vec![0u8; size as usize]; + file.read_exact(&mut buffer)?; + return Ok(buffer); + } + + // NO FALLBACK! If we don't have storage boundaries, we cannot load data lazily + // The storage map MUST be built from tensor metadata for lazy loading to work + Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "Storage boundaries not available for key '{}'. Cannot perform lazy loading.", + storage_key + ), + )) + } +} + +impl TarSource { + /// Create a new TAR source by parsing storages data. + /// + /// # Arguments + /// * `storages_data` - Raw storages blob with structure: + /// - Count pickle (number of storages) + /// - For each storage: metadata pickle + u64 num_elements + raw binary data + pub fn new(storages_data: Vec) -> std::io::Result { + use super::pickle_reader::{read_pickle, storage_type_to_element_size}; + use std::io::Cursor; + + let mut storage_map = HashMap::new(); + let mut pos = 0usize; + + // First, read the count of storages + let mut cursor = Cursor::new(&storages_data[pos..]); + let storage_count = + if let Ok(super::pickle_reader::Object::Int(count)) = read_pickle(&mut cursor) { + pos += cursor.position() as usize; + count as usize + } else { + 0 + }; + + // Parse each storage entry + for _i in 0..storage_count { + if pos >= storages_data.len() { + break; + } + + // Read the storage metadata pickle: (storage_key, device, storage_type) + let mut cursor = Cursor::new(&storages_data[pos..]); + if let Ok(obj) = read_pickle(&mut cursor) { + let pickle_size = cursor.position() as usize; + pos += pickle_size; + + // Extract storage info from pickle tuple + let (storage_key, storage_type) = match obj { + super::pickle_reader::Object::Tuple(tuple) if tuple.len() >= 3 => { + let key = match &tuple[0] { + super::pickle_reader::Object::Int(i) => i.to_string(), + super::pickle_reader::Object::String(s) => s.clone(), + _ => continue, + }; + // tuple[1] is device (e.g., "cpu") + // tuple[2] is storage type class + let stype = match &tuple[2] { + super::pickle_reader::Object::Class { name, .. } => name.clone(), + other => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("Expected Class for storage type, got {:?}", other), + )); + } + }; + (key, stype) + } + _ => continue, + }; + + // Read the number of elements (u64 little-endian) + if pos + 8 > storages_data.len() { + break; + } + let num_elements = u64::from_le_bytes([ + storages_data[pos], + storages_data[pos + 1], + storages_data[pos + 2], + storages_data[pos + 3], + storages_data[pos + 4], + storages_data[pos + 5], + storages_data[pos + 6], + storages_data[pos + 7], + ]) as usize; + pos += 8; + + // Determine element size from storage type + let element_size = storage_type_to_element_size(&storage_type) + .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; + + let data_size = num_elements * element_size; + + // Store the offset to raw data and its size + storage_map.insert(storage_key, (pos, data_size)); + + // Skip the raw binary data + pos += data_size; + } else { + break; + } + } + + Ok(Self { + storage_map, + storages_data, + }) + } + + /// Read data for a specific storage key + pub fn read_file(&self, key: &str) -> std::io::Result> { + // Extract the storage key from paths like "data/0" + let storage_key = key.split('/').next_back().unwrap_or(key); + + if let Some(&(offset, size)) = self.storage_map.get(storage_key) + && offset + size <= self.storages_data.len() + { + return Ok(self.storages_data[offset..offset + size].to_vec()); + } + + Err(std::io::Error::new( + std::io::ErrorKind::NotFound, + format!("Storage key '{}' not found in TAR archive", storage_key), + )) + } + + /// Read a range of data for a specific storage key (avoids double allocation) + pub fn read_file_range( + &self, + key: &str, + offset: usize, + length: usize, + ) -> std::io::Result> { + let storage_key = key.split('/').next_back().unwrap_or(key); + + if let Some(&(storage_offset, storage_size)) = self.storage_map.get(storage_key) + && storage_offset + storage_size <= self.storages_data.len() + { + let start = storage_offset + offset; + let end = (storage_offset + offset + length).min(storage_offset + storage_size); + return Ok(self.storages_data[start..end].to_vec()); + } + + Err(std::io::Error::new( + std::io::ErrorKind::NotFound, + format!("Storage key '{}' not found in TAR archive", storage_key), + )) + } + + /// Check if a storage key exists + pub fn contains(&self, key: &str) -> bool { + let storage_key = key.split('/').next_back().unwrap_or(key); + self.storage_map.contains_key(storage_key) + } + + /// Get list of storage keys + pub fn keys(&self) -> Vec { + self.storage_map.keys().cloned().collect() + } +} + +impl LazyDataSource { + /// Create from a ZIP file + pub fn from_zip(path: impl AsRef) -> std::io::Result { + Ok(Self::Zip(Arc::new(Mutex::new(ZipSource::new( + path.as_ref().to_path_buf(), + )?)))) + } + + /// Create from a TAR archive's storages data + pub fn from_tar(storages_data: &[u8]) -> std::io::Result { + Ok(Self::Tar(Arc::new(Mutex::new(TarSource::new( + storages_data.to_vec(), + )?)))) + } + + /// Create from a legacy multi-storage file + pub fn from_legacy_multi_storage( + path: impl AsRef, + data_offset: u64, + data_size: u64, + ) -> Self { + Self::LegacyMultiStorage(Arc::new(Mutex::new(LegacyMultiStorageSource::new( + path.as_ref().to_path_buf(), + data_offset, + data_size, + )))) + } + + /// Read data for a specific key + pub fn read(&self, key: &str) -> std::io::Result> { + match self { + Self::Zip(source) => { + let source = source + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + source.read_file(key) + } + Self::Tar(source) => { + let source = source + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + source.read_file(key) + } + Self::LegacyMultiStorage(source) => { + let source = source + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + source.read(key) + } + } + } + + /// Read a portion of data for a specific key + pub fn read_range(&self, key: &str, offset: usize, length: usize) -> std::io::Result> { + match self { + Self::Zip(source) => { + let source = source + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + source.read_file_range(key, offset, length) + } + Self::Tar(source) => { + let source = source + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + source.read_file_range(key, offset, length) + } + Self::LegacyMultiStorage(source) => { + // For legacy format, read only the requested range + let storage_key = key.split('/').next_back().unwrap_or(key); + let source = source + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + + // Get storage boundaries + let storage_map = source + .storage_map + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(ref map) = *storage_map + && let Some(&(storage_offset, storage_size)) = map.get(storage_key) + { + // Calculate actual file position + let file_offset = source.data_offset + storage_offset + offset as u64; + let read_length = length.min((storage_size as usize).saturating_sub(offset)); + + // Read only the requested range + let mut file = File::open(&source.path)?; + file.seek(std::io::SeekFrom::Start(file_offset))?; + + let mut buffer = vec![0u8; read_length]; + file.read_exact(&mut buffer)?; + Ok(buffer) + } else { + Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "Storage boundaries not available for key '{}'. Cannot perform lazy loading.", + storage_key + ), + )) + } + } + } + } + + /// Check if a key exists + pub fn contains(&self, key: &str) -> bool { + match self { + Self::Zip(source) => { + let source = source + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + source.contains(key) + } + Self::Tar(source) => { + let source = source + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + source.contains(key) + } + Self::LegacyMultiStorage(_) => true, // Legacy format has all data + } + } + + /// Get list of available keys (for ZIP sources) + pub fn keys(&self) -> Vec { + match self { + Self::Zip(source) => { + let source = source + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + source.data_files() + } + Self::Tar(source) => { + let source = source + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + source.keys() + } + Self::LegacyMultiStorage(_) => vec![], // Legacy format doesn't have distinct keys + } + } +} diff --git a/crates/burn-store/src/pytorch/mod.rs b/crates/burn-store/src/pytorch/mod.rs new file mode 100644 index 0000000..5eecf6b --- /dev/null +++ b/crates/burn-store/src/pytorch/mod.rs @@ -0,0 +1,48 @@ +//! PyTorch format support for burn-store. +//! +//! This module provides comprehensive support for loading PyTorch model files (.pth, .pt) +//! into Burn, with automatic weight transformation and flexible configuration options. +//! +//! ## Features +//! +//! - **Direct .pth/.pt file loading**: Load PyTorch checkpoint and state dict files +//! - **Automatic weight transformation**: `PyTorchToBurnAdapter` is applied by default: +//! - Linear layer weights are automatically transposed +//! - Normalization parameters are renamed (gamma → weight, beta → bias) +//! - Conv2d weights maintain their format +//! - **Flexible filtering**: Load only specific layers or parameters +//! - **Key remapping**: Rename tensors during loading to match your model structure +//! - **Partial loading**: Continue even when some tensors are missing +//! +//! ## Example +//! +//! ```rust,ignore +//! use burn_store::PytorchStore; +//! +//! // Load a PyTorch model (PyTorchToBurnAdapter is applied automatically) +//! let mut store = PytorchStore::from_file("model.pth") +//! .with_top_level_key("state_dict") // Access nested state dict +//! .with_regex(r"^encoder\..*") // Only load encoder layers +//! .with_key_remapping(r"^fc\.", "linear.") // Rename fc -> linear +//! .allow_partial(true); // Skip missing tensors +//! +//! let mut model = MyModel::new(&device); +//! let result = model.load_from(&mut store)?; +//! +//! println!("Loaded {} tensors", result.applied.len()); +//! if !result.missing.is_empty() { +//! println!("Missing tensors: {:?}", result.missing); +//! } +//! ``` + +pub mod lazy_data; +pub mod pickle_reader; +pub mod reader; +pub mod store; + +#[cfg(test)] +pub mod tests; + +// Main public interface +pub use reader::{PytorchError, PytorchReader}; +pub use store::{PytorchStore, PytorchStoreError}; diff --git a/crates/burn-store/src/pytorch/pickle_reader.rs b/crates/burn-store/src/pytorch/pickle_reader.rs new file mode 100644 index 0000000..5362859 --- /dev/null +++ b/crates/burn-store/src/pytorch/pickle_reader.rs @@ -0,0 +1,1538 @@ +//! Just enough pickle support to be able to read PyTorch checkpoints. +//! +//! This implementation is based on the candle project's pickle loader with significant +//! modifications for improved separation of concerns and extended PyTorch compatibility. +//! +//! Original source: +//! +//! Modifications include: +//! - Lazy tensor data loading for memory efficiency +//! - Extended PyTorch version compatibility (0.1.10 - 2.x) +//! - Better separation of pickle parsing and tensor extraction +//! - Support for both legacy and modern PyTorch formats +use crate::TensorSnapshot; +use crate::pytorch::lazy_data::LazyDataSource; +use alloc::rc::Rc; +use alloc::string::{String, ToString}; +use alloc::vec::Vec; +use burn_core::module::ParamId; +use burn_core::tensor::{BoolStore, DType, TensorData}; +use byteorder::{LittleEndian, ReadBytesExt}; +use half::{bf16, f16}; +use std::collections::HashMap; +use std::io::{self, BufRead}; +use std::sync::Arc; + +/// Error type for pickle operations +#[derive(Debug)] +pub enum PickleError { + Io(io::Error), + InvalidOpCode(u8), + InvalidProtocol(u8), + UnexpectedOpCode(OpCode), + UnsupportedType(String), + InvalidData(String), + StackUnderflow, + MemoNotFound(u32), + InvalidShapeOrType, +} + +impl From for PickleError { + fn from(e: io::Error) -> Self { + PickleError::Io(e) + } +} + +impl std::fmt::Display for PickleError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + PickleError::Io(e) => write!(f, "IO error: {}", e), + PickleError::InvalidOpCode(code) => write!( + f, + "Invalid pickle opcode: 0x{:02x}. The file may be corrupted or use an unsupported pickle protocol.", + code + ), + PickleError::InvalidProtocol(proto) => write!( + f, + "Invalid or unsupported pickle protocol version: {}. Supported versions are 2-5.", + proto + ), + PickleError::UnexpectedOpCode(op) => { + write!(f, "Unexpected pickle opcode {:?} in current context", op) + } + PickleError::UnsupportedType(ty) => write!( + f, + "Unsupported Python type '{}'. This may indicate a full model save rather than a state_dict.", + ty + ), + PickleError::InvalidData(msg) => write!(f, "Invalid data in pickle file: {}", msg), + PickleError::StackUnderflow => { + write!(f, "Pickle stack underflow - the file may be corrupted") + } + PickleError::MemoNotFound(idx) => write!( + f, + "Pickle memo reference {} not found - the file may be corrupted", + idx + ), + PickleError::InvalidShapeOrType => { + write!(f, "Invalid tensor shape or data type in PyTorch file") + } + } + } +} + +impl std::error::Error for PickleError {} + +type Result = std::result::Result; + +/// Convert PyTorch storage type name to element size in bytes. +/// +/// This is used to calculate storage sizes for lazy loading. +/// The storage type names follow PyTorch's naming convention (e.g., "FloatStorage", "BFloat16Storage"). +/// +/// Returns an error for unknown storage types to avoid silently loading garbage data. +pub fn storage_type_to_element_size(storage_type: &str) -> std::result::Result { + match storage_type { + "DoubleStorage" | "LongStorage" | "ComplexFloatStorage" => Ok(8), + "FloatStorage" | "IntStorage" | "ComplexHalfStorage" => Ok(4), + "HalfStorage" | "BFloat16Storage" | "ShortStorage" => Ok(2), + "ByteStorage" | "CharStorage" | "BoolStorage" => Ok(1), + _ => Err(format!("Unknown storage type: {}", storage_type)), + } +} + +// https://docs.juliahub.com/Pickle/LAUNc/0.1.0/opcode/ +#[repr(u8)] +#[derive(Debug, Eq, PartialEq, Clone)] +pub enum OpCode { + // https://github.com/python/cpython/blob/ed25f097160b5cbb0c9a1f9a746d2f1bbc96515a/Lib/pickletools.py#L2123 + Proto = 0x80, + Global = b'c', + BinPut = b'q', + LongBinPut = b'r', + EmptyTuple = b')', + Reduce = b'R', + Mark = b'(', + BinUnicode = b'X', + ShortBinString = b'U', + BinInt = b'J', + Int = b'I', + Tuple = b't', + BinPersId = b'Q', + BinInt1 = b'K', + BinInt2 = b'M', + Tuple1 = 0x85, + Tuple2 = 0x86, + Tuple3 = 0x87, + NewTrue = 0x88, + NewFalse = 0x89, + None = b'N', + BinGet = b'h', + LongBinGet = b'j', + SetItem = b's', + SetItems = b'u', + EmptyDict = b'}', + Dict = b'd', + Build = b'b', + Stop = b'.', + NewObj = 0x81, + EmptyList = b']', + List = b'l', + BinFloat = b'G', + Append = b'a', + Appends = b'e', + Long1 = 0x8a, + Memoize = 0x94, +} + +// Avoid using FromPrimitive so as not to drag another dependency. +impl TryFrom for OpCode { + type Error = u8; + fn try_from(value: u8) -> std::result::Result { + match value { + 0x80 => Ok(Self::Proto), + b'c' => Ok(Self::Global), + b'q' => Ok(Self::BinPut), + b'r' => Ok(Self::LongBinPut), + b')' => Ok(Self::EmptyTuple), + b'R' => Ok(Self::Reduce), + b'(' => Ok(Self::Mark), + b'X' => Ok(Self::BinUnicode), + b'U' => Ok(Self::ShortBinString), + b'J' => Ok(Self::BinInt), + b'I' => Ok(Self::Int), + b't' => Ok(Self::Tuple), + b'Q' => Ok(Self::BinPersId), + b'K' => Ok(Self::BinInt1), + b'M' => Ok(Self::BinInt2), + b'N' => Ok(Self::None), + 0x85 => Ok(Self::Tuple1), + 0x86 => Ok(Self::Tuple2), + 0x87 => Ok(Self::Tuple3), + 0x88 => Ok(Self::NewTrue), + 0x89 => Ok(Self::NewFalse), + b'h' => Ok(Self::BinGet), + b'j' => Ok(Self::LongBinGet), + b's' => Ok(Self::SetItem), + b'u' => Ok(Self::SetItems), + b'}' => Ok(Self::EmptyDict), + b'd' => Ok(Self::Dict), + b'b' => Ok(Self::Build), + b'.' => Ok(Self::Stop), + 0x81 => Ok(Self::NewObj), + b']' => Ok(Self::EmptyList), + b'l' => Ok(Self::List), + b'G' => Ok(Self::BinFloat), + b'a' => Ok(Self::Append), + b'e' => Ok(Self::Appends), + 0x8a => Ok(Self::Long1), + 0x94 => Ok(Self::Memoize), + value => Err(value), + } + } +} + +fn read_to_newline(r: &mut R) -> Result> { + let mut data: Vec = Vec::with_capacity(32); + r.read_until(b'\n', &mut data)?; + data.pop(); + if data.last() == Some(&b'\r') { + data.pop(); + } + Ok(data) +} + +fn buf_to_str(buf: &[u8]) -> Result { + String::from_utf8(buf.to_vec()) + .map_err(|e| PickleError::InvalidData(format!("Invalid UTF-8: {}", e))) +} + +#[derive(Debug, Clone)] +pub enum Object { + Class { + module_name: String, + name: String, + }, + String(String), + Int(i64), + Float(f64), + Bool(bool), + None, + Tuple(Vec), + List(Vec), + Dict(HashMap), + Persistent(Vec), + PersistentTuple(Vec), + Reduce { + callable: Box, + args: Box, + }, + Build { + callable: Box, + args: Box, + }, + TorchParam(TensorSnapshot), +} + +fn rebuild_from_type_v2( + o: Object, + memo: &mut HashMap, + data_source: &Option>, +) -> Result { + let args = if let Object::Tuple(args) = o { + if args.is_empty() { + return Err(PickleError::InvalidData( + "rebuild_from_type_v2: empty args".to_string(), + )); + } + args + } else { + return Err(PickleError::InvalidData(format!( + "rebuild_from_type_v2: expected tuple got {:?}", + o + ))); + }; + let func = &args[0]; + match func { + Object::Class { module_name, name } => { + let module_name = module_name.as_str(); + let name = name.as_str(); + // For rebuild_tensor_v2, the args might already be in a tuple + let actual_args = if args.len() == 2 && matches!(&args[1], Object::Tuple(_)) { + // If there's only one arg and it's a tuple, use it directly + args[1].clone() + } else { + // Otherwise, wrap the remaining args in a tuple + Object::Tuple(args[1..].to_vec()) + }; + if module_name == "torch._utils" && name == "_rebuild_tensor_v2" { + rebuild_tensor_v2(actual_args, memo, data_source) + } else if module_name == "torch._utils" && name == "_rebuild_tensor" { + // Legacy _rebuild_tensor (PyTorch < 1.6) + // Same as v2 but with fewer arguments: (storage, storage_offset, size, stride) + rebuild_tensor(actual_args, memo, data_source) + } else if module_name == "torch._tensor" && name == "_rebuild_from_type_v2" { + rebuild_from_type_v2(actual_args, memo, data_source) + } else if module_name == "torch._utils" && name == "_rebuild_parameter" { + rebuild_parameter(actual_args, memo, data_source) + } else if module_name == "collections" && name == "OrderedDict" { + // OrderedDict is treated as a regular Dict in our implementation + Ok(Object::Dict(HashMap::new())) + } else { + Err(PickleError::UnsupportedType(format!( + "{}.{}", + module_name, name + ))) + } + } + _ => Err(PickleError::InvalidData(format!( + "rebuild_from_type_v2: expected class got {:?}", + func + ))), + } +} + +fn rebuild_parameter( + args: Object, + memo: &mut HashMap, + data_source: &Option>, +) -> Result { + let args = if let Object::Tuple(args) = args { + if args.is_empty() { + return Err(PickleError::InvalidData( + "rebuild_parameter: empty args".to_string(), + )); + } + args + } else { + return Err(PickleError::InvalidData(format!( + "rebuild_parameter: expected tuple got {:?}", + args + ))); + }; + let data = &args[0]; + let tensor = match data { + Object::Reduce { + callable: _, + args: _, + } => rebuild_from_type_v2(data.clone(), memo, data_source)?, + _ => data.clone(), + }; + Ok(tensor) +} + +/// Parse storage argument and extract storage info and tuple. +fn parse_storage_arg(arg: &Object, fn_name: &str) -> Result<(Vec, Option>)> { + match arg { + Object::Persistent(data) => Ok((data.clone(), None)), + Object::PersistentTuple(tuple) => Ok((vec![], Some(tuple.clone()))), + // Also accept regular Tuple for TAR format compatibility + Object::Tuple(tuple) => Ok((vec![], Some(tuple.clone()))), + _ => Err(PickleError::InvalidData(format!( + "{}: expected persistent id got {:?}", + fn_name, arg + ))), + } +} + +/// Parse shape argument. +fn parse_shape_arg(arg: &Object, fn_name: &str) -> Result> { + match arg { + Object::Tuple(shape) => shape + .iter() + .map(|x| match x { + Object::Int(i) => Ok(*i as usize), + _ => Err(PickleError::InvalidData( + "shape must contain ints".to_string(), + )), + }) + .collect::>>(), + _ => Err(PickleError::InvalidData(format!( + "{}: expected shape tuple got {:?}", + fn_name, arg + ))), + } +} + +/// Legacy _rebuild_tensor function for PyTorch < 1.6. +/// Thin wrapper that parses 4 arguments and calls rebuild_tensor_impl. +fn rebuild_tensor( + args: Object, + _memo: &mut HashMap, + data_source: &Option>, +) -> Result { + let args = if let Object::Tuple(args) = args { + args + } else { + return Err(PickleError::InvalidData(format!( + "rebuild_tensor: expected tuple got {:?}", + args + ))); + }; + + if args.len() < 4 { + return Err(PickleError::InvalidData(format!( + "rebuild_tensor: expected at least 4 args, got {}", + args.len() + ))); + } + + let (storage_info, storage_tuple) = parse_storage_arg(&args[0], "rebuild_tensor")?; + let storage_offset = match &args[1] { + Object::Int(offset) => *offset as usize, + _ => 0, + }; + let shape = parse_shape_arg(&args[2], "rebuild_tensor")?; + + rebuild_tensor_impl( + storage_info, + storage_tuple, + storage_offset, + shape, + data_source, + ) +} + +/// Modern _rebuild_tensor_v2 function for PyTorch >= 1.6. +/// Thin wrapper that parses 5+ arguments and calls rebuild_tensor_impl. +fn rebuild_tensor_v2( + args: Object, + _memo: &mut HashMap, + data_source: &Option>, +) -> Result { + let args = if let Object::Tuple(args) = args { + args + } else { + return Err(PickleError::InvalidData(format!( + "rebuild_tensor_v2: expected tuple got {:?}", + args + ))); + }; + + if args.len() < 5 { + return Err(PickleError::InvalidData(format!( + "rebuild_tensor_v2: expected at least 5 args, got {}", + args.len() + ))); + } + + let (storage_info, storage_tuple) = parse_storage_arg(&args[0], "rebuild_tensor_v2")?; + let storage_offset = match &args[1] { + Object::Int(offset) => *offset as usize, + _ => 0, + }; + let shape = parse_shape_arg(&args[2], "rebuild_tensor_v2")?; + // args[3] is stride (unused) + // args[4] is requires_grad (unused) + // args[5] is backward_hooks (unused) + + rebuild_tensor_impl( + storage_info, + storage_tuple, + storage_offset, + shape, + data_source, + ) +} + +/// Helper to convert storage type name to DType. +fn storage_type_to_dtype(storage_type: &str) -> Result { + match storage_type { + "FloatStorage" => Ok(DType::F32), + "DoubleStorage" => Ok(DType::F64), + "HalfStorage" => Ok(DType::F16), + "BFloat16Storage" => Ok(DType::BF16), + "LongStorage" => Ok(DType::I64), + "IntStorage" => Ok(DType::I32), + "ShortStorage" => Ok(DType::I16), + "CharStorage" => Ok(DType::I8), + "ByteStorage" => Ok(DType::U8), + "BoolStorage" => Ok(DType::Bool(BoolStore::Native)), + _ => Err(PickleError::InvalidData(format!( + "Unknown storage type: {}", + storage_type + ))), + } +} + +/// Core implementation for rebuilding tensors. +/// Shared by both rebuild_tensor (legacy) and rebuild_tensor_v2 (modern). +fn rebuild_tensor_impl( + storage_info: Vec, + storage_tuple: Option>, + storage_offset: usize, + shape: Vec, + data_source: &Option>, +) -> Result { + // Parse the storage info to extract dtype and storage key + // The persistent ID is typically a tuple like: ('storage', 'FloatStorage', '0', 'cpu', 4) + let (dtype, storage_key, storage_total_elements) = if let Some(tuple) = storage_tuple { + // Direct tuple access + if tuple.len() >= 3 { + let storage_type = match &tuple[1] { + Object::String(s) => s.as_str(), + Object::Class { + module_name: _, + name, + } => name.as_str(), + other => { + return Err(PickleError::InvalidData(format!( + "Expected storage type as String or Class, got {:?}", + other + ))); + } + }; + let dtype = storage_type_to_dtype(storage_type)?; + let key = match &tuple[2] { + Object::String(s) => s.clone(), + other => { + return Err(PickleError::InvalidData(format!( + "Expected storage key as String, got {:?}", + other + ))); + } + }; + // Extract total element count from index 4 + let total_elements = match tuple.get(4) { + Some(Object::Int(n)) => Some(*n as usize), + _ => None, + }; + (dtype, key, total_elements) + } else { + return Err(PickleError::InvalidData(format!( + "Storage tuple too short, expected at least 3 elements, got {}", + tuple.len() + ))); + } + } else if !storage_info.is_empty() { + // Legacy string-based parsing + let storage_str = String::from_utf8_lossy(&storage_info); + if storage_str.starts_with("Tuple(") { + // Parse from the debug representation we stored + let parts: Vec<&str> = storage_str + .trim_start_matches("Tuple(") + .trim_end_matches(")") + .split(", ") + .map(|s| { + let trimmed = s.trim_matches('"'); + if let Some(inner) = trimmed + .strip_prefix("Object::String(\"") + .and_then(|s| s.strip_suffix("\")")) + { + inner + } else { + trimmed + } + }) + .collect(); + + if parts.len() >= 3 { + let dtype = storage_type_to_dtype(parts[1])?; + (dtype, parts[2].to_string(), None) + } else { + return Err(PickleError::InvalidData(format!( + "Storage info tuple too short, expected at least 3 parts, got {}", + parts.len() + ))); + } + } else { + return Err(PickleError::InvalidData(format!( + "Invalid storage info format: {}", + storage_str + ))); + } + } else { + return Err(PickleError::InvalidData( + "No storage information available".to_string(), + )); + }; + + // If no data source, we can't load tensor data + let data_source = match data_source { + Some(ds) => ds.clone(), + None => { + return Err(PickleError::InvalidData( + "Cannot load tensor data without a data source".to_string(), + )); + } + }; + + // Create clones for the closure + let data_source_clone = data_source.clone(); + let shape_clone = shape.clone(); + + // Find the correct data file key + let data_file_key = { + let exact_key = format!("data/{}", storage_key); + if data_source.contains(&exact_key) { + exact_key + } else { + // Try other patterns + data_source + .keys() + .into_iter() + .find(|key| { + key.ends_with(&format!("/data/{}", storage_key)) + || (key.contains("/data/") && key.rsplit('/').next() == Some(&storage_key)) + }) + .unwrap_or_else(|| format!("data/{}", storage_key)) + } + }; + + // There is no guarantee that the storage size equals the python object size. + // The size of the object can be smaller than the storage (e.g., uncloned views + // of a tensor are saved to the .pt/.pth files). Thus, the actual storage size + // should be used, which can be computed using the total number of elements + // in the storage. + if let LazyDataSource::LegacyMultiStorage(ref source) = *data_source { + let source = source + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let bytes_needed = match storage_total_elements { + Some(total) => total * dtype.size(), + None => (storage_offset + shape.iter().product::()) * dtype.size(), + }; + source.track_storage_usage(&storage_key, 0, bytes_needed); + } + + // Create a TensorSnapshot with a closure that loads the actual data on-demand + Ok(Object::TorchParam(TensorSnapshot::from_closure( + Rc::new(move || { + // Load data only when needed + if let Ok(data) = data_source_clone.read(&data_file_key) { + // Parse the binary data based on dtype + let num_elements = shape_clone.iter().product::().max(1); + + // Use dtype.size() to get element size in bytes + let element_size = dtype.size(); + + // Apply storage offset + let offset_bytes = storage_offset * element_size; + if offset_bytes >= data.len() { + return Ok(TensorData::new( + vec![0.0f32; num_elements], + shape_clone.clone(), + )); + } + + let data_slice = &data[offset_bytes..]; + let available_elements = data_slice.len() / element_size; + let elements_to_read = num_elements.min(available_elements); + + // Convert bytes to the appropriate type + match dtype { + DType::F32 => { + let mut values = Vec::with_capacity(num_elements); + for i in 0..elements_to_read { + let bytes = [ + data_slice[i * element_size], + data_slice[i * element_size + 1], + data_slice[i * element_size + 2], + data_slice[i * element_size + 3], + ]; + values.push(f32::from_le_bytes(bytes)); + } + // Pad with zeros if needed + values.resize(num_elements, 0.0); + Ok(TensorData::new(values, shape_clone.clone())) + } + DType::F64 => { + let mut values = Vec::with_capacity(num_elements); + for i in 0..elements_to_read { + let mut bytes = [0u8; 8]; + bytes.copy_from_slice( + &data_slice[i * element_size..(i + 1) * element_size], + ); + values.push(f64::from_le_bytes(bytes)); + } + values.resize(num_elements, 0.0); + Ok(TensorData::new(values, shape_clone.clone())) + } + DType::I64 => { + let mut values = Vec::with_capacity(num_elements); + for i in 0..elements_to_read { + let mut bytes = [0u8; 8]; + bytes.copy_from_slice( + &data_slice[i * element_size..(i + 1) * element_size], + ); + values.push(i64::from_le_bytes(bytes)); + } + values.resize(num_elements, 0); + Ok(TensorData::new(values, shape_clone.clone())) + } + DType::I32 => { + let mut values = Vec::with_capacity(num_elements); + for i in 0..elements_to_read { + let mut bytes = [0u8; 4]; + bytes.copy_from_slice( + &data_slice[i * element_size..(i + 1) * element_size], + ); + values.push(i32::from_le_bytes(bytes)); + } + values.resize(num_elements, 0); + Ok(TensorData::new(values, shape_clone.clone())) + } + DType::I16 => { + let mut values = Vec::with_capacity(num_elements); + for i in 0..elements_to_read { + let mut bytes = [0u8; 2]; + bytes.copy_from_slice( + &data_slice[i * element_size..(i + 1) * element_size], + ); + values.push(i16::from_le_bytes(bytes)); + } + values.resize(num_elements, 0); + Ok(TensorData::new(values, shape_clone.clone())) + } + DType::I8 => { + let mut values = Vec::with_capacity(num_elements); + for &byte in data_slice.iter().take(elements_to_read) { + values.push(byte as i8); + } + values.resize(num_elements, 0); + Ok(TensorData::new(values, shape_clone.clone())) + } + DType::Bool(BoolStore::Native) => { + let mut values = Vec::with_capacity(num_elements); + for &byte in data_slice.iter().take(elements_to_read) { + values.push(byte != 0); + } + values.resize(num_elements, false); + Ok(TensorData::new(values, shape_clone.clone())) + } + DType::F16 => { + let mut values = Vec::with_capacity(num_elements); + for i in 0..elements_to_read { + let mut bytes = [0u8; 2]; + bytes.copy_from_slice( + &data_slice[i * element_size..(i + 1) * element_size], + ); + values.push(f16::from_le_bytes(bytes)); + } + values.resize(num_elements, f16::ZERO); + Ok(TensorData::new(values, shape_clone.clone())) + } + DType::BF16 => { + let mut values = Vec::with_capacity(num_elements); + for i in 0..elements_to_read { + let mut bytes = [0u8; 2]; + bytes.copy_from_slice( + &data_slice[i * element_size..(i + 1) * element_size], + ); + values.push(bf16::from_le_bytes(bytes)); + } + values.resize(num_elements, bf16::ZERO); + Ok(TensorData::new(values, shape_clone.clone())) + } + DType::U8 => { + let mut values = Vec::with_capacity(num_elements); + for &byte in data_slice.iter().take(elements_to_read) { + values.push(byte); + } + values.resize(num_elements, 0); + Ok(TensorData::new(values, shape_clone.clone())) + } + DType::U16 => { + let mut values = Vec::with_capacity(num_elements); + for i in 0..elements_to_read { + let mut bytes = [0u8; 2]; + bytes.copy_from_slice( + &data_slice[i * element_size..(i + 1) * element_size], + ); + values.push(u16::from_le_bytes(bytes)); + } + values.resize(num_elements, 0); + Ok(TensorData::new(values, shape_clone.clone())) + } + DType::U32 => { + let mut values = Vec::with_capacity(num_elements); + for i in 0..elements_to_read { + let mut bytes = [0u8; 4]; + bytes.copy_from_slice( + &data_slice[i * element_size..(i + 1) * element_size], + ); + values.push(u32::from_le_bytes(bytes)); + } + values.resize(num_elements, 0); + Ok(TensorData::new(values, shape_clone.clone())) + } + DType::U64 => { + let mut values = Vec::with_capacity(num_elements); + for i in 0..elements_to_read { + let mut bytes = [0u8; 8]; + bytes.copy_from_slice( + &data_slice[i * element_size..(i + 1) * element_size], + ); + values.push(u64::from_le_bytes(bytes)); + } + values.resize(num_elements, 0); + Ok(TensorData::new(values, shape_clone.clone())) + } + _ => { + // For any remaining unsupported types, return an error + Err(crate::TensorSnapshotError::DataError(format!( + "Unsupported dtype for tensor data reading: {:?}", + dtype + ))) + } + } + } else { + // If no data file found, return zeros of the appropriate type + let num_elements = shape_clone.iter().product::().max(1); + match dtype { + DType::F32 => Ok(TensorData::new( + vec![0.0f32; num_elements], + shape_clone.clone(), + )), + DType::F64 => Ok(TensorData::new( + vec![0.0f64; num_elements], + shape_clone.clone(), + )), + DType::F16 => Ok(TensorData::new( + vec![f16::ZERO; num_elements], + shape_clone.clone(), + )), + DType::BF16 => Ok(TensorData::new( + vec![bf16::ZERO; num_elements], + shape_clone.clone(), + )), + DType::I64 => Ok(TensorData::new( + vec![0i64; num_elements], + shape_clone.clone(), + )), + DType::I32 => Ok(TensorData::new( + vec![0i32; num_elements], + shape_clone.clone(), + )), + DType::I16 => Ok(TensorData::new( + vec![0i16; num_elements], + shape_clone.clone(), + )), + DType::I8 => Ok(TensorData::new( + vec![0i8; num_elements], + shape_clone.clone(), + )), + DType::U8 => Ok(TensorData::new( + vec![0u8; num_elements], + shape_clone.clone(), + )), + DType::U16 => Ok(TensorData::new( + vec![0u16; num_elements], + shape_clone.clone(), + )), + DType::U32 => Ok(TensorData::new( + vec![0u32; num_elements], + shape_clone.clone(), + )), + DType::U64 => Ok(TensorData::new( + vec![0u64; num_elements], + shape_clone.clone(), + )), + DType::Bool(BoolStore::Native) => Ok(TensorData::new( + vec![false; num_elements], + shape_clone.clone(), + )), + _ => { + // For any remaining unsupported types, return an error + Err(crate::TensorSnapshotError::DataError(format!( + "Unsupported dtype for tensor data reading: {:?}", + dtype + ))) + } + } + } + }), + dtype, + shape.into(), + vec![], // path_stack + vec![], // container_stack + ParamId::new(), // tensor_id + ))) +} + +pub struct Stack { + stack: Vec, + memo: HashMap, + data_source: Option>, +} + +impl Default for Stack { + fn default() -> Self { + Self::new() + } +} + +impl Stack { + pub fn new() -> Self { + // For cases where no data source is needed (pure pickle without tensor data) + Self { + stack: Vec::new(), + memo: HashMap::new(), + data_source: None, + } + } + + pub fn with_data_source(data_source: Arc) -> Self { + Self { + stack: Vec::new(), + memo: HashMap::new(), + data_source: Some(data_source), + } + } + + fn push(&mut self, o: Object) { + self.stack.push(o) + } + + fn pop(&mut self) -> Result { + match self.stack.pop() { + None => Err(PickleError::StackUnderflow), + Some(o) => Ok(o), + } + } + + fn top(&self) -> Result { + match self.stack.last() { + None => Err(PickleError::StackUnderflow), + Some(o) => Ok(o.clone()), + } + } + + fn pop_to_marker(&mut self) -> Result> { + let marker_pos = self + .stack + .iter() + .rposition(|o| { + matches!(o, Object::Class { module_name, name } + if module_name == "mark" && name == "mark") + }) + .ok_or(PickleError::InvalidData("marker not found".to_string()))?; + + let result = self.stack.split_off(marker_pos + 1); + self.stack.pop(); // Remove the marker + Ok(result) + } + + fn last_mut(&mut self) -> Result<&mut Object> { + match self.stack.last_mut() { + None => Err(PickleError::StackUnderflow), + Some(o) => Ok(o), + } + } + + fn push_mark(&mut self) { + self.stack.push(Object::Class { + module_name: "mark".to_string(), + name: "mark".to_string(), + }); + } + + fn memo_get(&self, idx: u32) -> Result { + self.memo + .get(&idx) + .cloned() + .ok_or(PickleError::MemoNotFound(idx)) + } + + fn memo_put(&mut self, idx: u32, obj: Object) { + self.memo.insert(idx, obj); + } + + fn memo_len(&self) -> usize { + self.memo.len() + } +} + +fn read_global(r: &mut R, stack: &mut Stack) -> Result<()> { + let module_name = buf_to_str(&read_to_newline(r)?)?; + let name = buf_to_str(&read_to_newline(r)?)?; + stack.push(Object::Class { module_name, name }); + Ok(()) +} + +fn read_long1(r: &mut R, stack: &mut Stack) -> Result<()> { + let len = r.read_u8()? as usize; + let mut data = vec![0u8; len]; + r.read_exact(&mut data)?; + // Handle little-endian signed integer + let mut value = 0i64; + for (i, &byte) in data.iter().enumerate().take(8) { + // Only process up to 8 bytes for i64, and use wrapping to avoid overflow + value |= (byte as i64).wrapping_shl((i as u32) * 8); + } + // Handle sign extension for negative numbers + if len < 8 && data.last().is_some_and(|&b| b & 0x80 != 0) { + // Sign extend + for i in len..8 { + value |= 0xffi64.wrapping_shl((i as u32) * 8); + } + } + stack.push(Object::Int(value)); + Ok(()) +} + +fn read_string(r: &mut R, stack: &mut Stack, len: usize) -> Result<()> { + let mut data = vec![0u8; len]; + r.read_exact(&mut data)?; + let s = buf_to_str(&data)?; + stack.push(Object::String(s)); + Ok(()) +} + +fn read_bin_int(r: &mut R, stack: &mut Stack) -> Result<()> { + let v = r.read_i32::()?; + stack.push(Object::Int(v as i64)); + Ok(()) +} + +fn read_int(r: &mut R, stack: &mut Stack) -> Result<()> { + // INT opcode reads an integer as ASCII string followed by newline + let line = read_to_newline(r)?; + let s = buf_to_str(&line)?; + let v = s + .parse::() + .map_err(|e| PickleError::InvalidData(format!("Invalid INT value '{}': {}", s, e)))?; + stack.push(Object::Int(v)); + Ok(()) +} + +fn read_bin_int1(r: &mut R, stack: &mut Stack) -> Result<()> { + let v = r.read_u8()?; + stack.push(Object::Int(v as i64)); + Ok(()) +} + +fn read_bin_int2(r: &mut R, stack: &mut Stack) -> Result<()> { + let v = r.read_u16::()?; + stack.push(Object::Int(v as i64)); + Ok(()) +} + +fn read_bin_float(r: &mut R, stack: &mut Stack) -> Result<()> { + // Python's BINFLOAT uses big-endian encoding + let v = r.read_f64::()?; + stack.push(Object::Float(v)); + Ok(()) +} + +pub fn read_pickle(r: &mut R) -> Result { + // For pure pickle without tensor data, no data source is needed + read_pickle_with_optional_data(r, None) +} + +/// Skip over a pickle without parsing it fully +/// This is useful for legacy format where we need to skip the main object +/// that contains tensors but we don't have a data source yet +pub fn skip_pickle(r: &mut R) -> Result<()> { + // Read the protocol marker if present + let mut first_byte = [0u8; 1]; + r.read_exact(&mut first_byte)?; + + if first_byte[0] == 0x80 { + // PROTO marker - read protocol version + let mut proto_version = [0u8; 1]; + r.read_exact(&mut proto_version)?; + } + // If not PROTO, the first byte is an opcode - continue to main loop + + // Helper to skip until newline + fn skip_line(r: &mut R) -> Result<()> { + let mut buf = Vec::new(); + r.read_until(b'\n', &mut buf)?; + Ok(()) + } + + // Helper to skip length-prefixed data + fn skip_length_prefixed(r: &mut R, length: usize) -> Result<()> { + let mut skip_buf = vec![0u8; length.min(8192)]; + let mut skipped = 0; + while skipped < length { + let to_skip = (length - skipped).min(skip_buf.len()); + r.read_exact(&mut skip_buf[..to_skip])?; + skipped += to_skip; + } + Ok(()) + } + + // Process first byte if it wasn't PROTO + let mut pending_byte = if first_byte[0] != 0x80 { + Some(first_byte[0]) + } else { + None + }; + + // Scan until we find STOP (0x2e) opcode + loop { + let byte = if let Some(b) = pending_byte.take() { + b + } else { + let mut byte = [0u8; 1]; + r.read_exact(&mut byte)?; + byte[0] + }; + + match byte { + 0x2e => { + // STOP - end of pickle + break; + } + // === Newline-terminated string opcodes === + 0x63 => { + // GLOBAL - two newline-terminated strings (module\nname\n) + skip_line(r)?; + skip_line(r)?; + } + 0x69 => { + // INST - two newline-terminated strings + skip_line(r)?; + skip_line(r)?; + } + 0x53 => { + // STRING - quoted string ending with newline + skip_line(r)?; + } + 0x46 | 0x49 | 0x4c => { + // FLOAT, INT, LONG - newline-terminated ASCII + skip_line(r)?; + } + 0x50 => { + // PERSID - newline-terminated persistent ID + skip_line(r)?; + } + // === Length-prefixed binary opcodes === + 0x58 | 0x42 | 0x43 | 0x54 | 0x55 | 0x56 | 0x8c | 0x8d | 0x8e => { + // String/bytes opcodes with length prefixes + let length = match byte { + 0x43 | 0x55 | 0x8c => { + // SHORT versions - 1 byte length + let mut len_byte = [0u8; 1]; + r.read_exact(&mut len_byte)?; + len_byte[0] as usize + } + 0x42 | 0x54 | 0x58 | 0x56 => { + // Regular versions - 4 byte length + let mut len_bytes = [0u8; 4]; + r.read_exact(&mut len_bytes)?; + u32::from_le_bytes(len_bytes) as usize + } + 0x8d | 0x8e => { + // 8-byte length versions + let mut len_bytes = [0u8; 8]; + r.read_exact(&mut len_bytes)?; + u64::from_le_bytes(len_bytes) as usize + } + _ => 0, + }; + skip_length_prefixed(r, length)?; + } + // === Fixed-size integer opcodes === + 0x4b => { + // BININT1 - 1 byte + let mut buf = [0u8; 1]; + r.read_exact(&mut buf)?; + } + 0x4d => { + // BININT2 - 2 bytes + let mut buf = [0u8; 2]; + r.read_exact(&mut buf)?; + } + 0x4a => { + // BININT - 4 bytes (signed int) + let mut buf = [0u8; 4]; + r.read_exact(&mut buf)?; + } + 0x47 => { + // BINFLOAT - 8 bytes + let mut buf = [0u8; 8]; + r.read_exact(&mut buf)?; + } + // === Variable-length integer opcodes === + 0x8a => { + // LONG1 - 1 byte length, then that many bytes + let mut len_byte = [0u8; 1]; + r.read_exact(&mut len_byte)?; + let length = len_byte[0] as usize; + skip_length_prefixed(r, length)?; + } + 0x8b => { + // LONG4 - 4 byte length, then that many bytes + let mut len_bytes = [0u8; 4]; + r.read_exact(&mut len_bytes)?; + let length = u32::from_le_bytes(len_bytes) as usize; + skip_length_prefixed(r, length)?; + } + // === Memo opcodes === + 0x71 | 0x68 => { + // BINPUT, BINGET - 1 byte index + let mut buf = [0u8; 1]; + r.read_exact(&mut buf)?; + } + 0x72 | 0x6a => { + // LONG_BINPUT, LONG_BINGET - 4 byte index + let mut buf = [0u8; 4]; + r.read_exact(&mut buf)?; + } + 0x67 | 0x70 => { + // GET, PUT - newline-terminated decimal index + skip_line(r)?; + } + // === Extension opcodes === + 0x82 => { + // EXT1 - 1 byte code + let mut buf = [0u8; 1]; + r.read_exact(&mut buf)?; + } + 0x83 => { + // EXT2 - 2 byte code + let mut buf = [0u8; 2]; + r.read_exact(&mut buf)?; + } + 0x84 => { + // EXT4 - 4 byte code + let mut buf = [0u8; 4]; + r.read_exact(&mut buf)?; + } + // === Frame opcode (protocol 4+) === + 0x95 => { + // FRAME - 8 byte frame size (we don't actually use framing, just skip the size) + let mut buf = [0u8; 8]; + r.read_exact(&mut buf)?; + } + // === Opcodes with no additional data === + // These just manipulate the stack or are markers + 0x28 | 0x29 | 0x30 | 0x31 | 0x32 | // MARK, TUPLE, POP, POP_MARK, DUP + 0x4e | 0x52 | 0x5d | 0x5b | 0x7d | // NONE, REDUCE, LIST, EMPTY_LIST, EMPTY_DICT + 0x61 | 0x62 | 0x64 | 0x65 | 0x73 | // APPEND, BUILD, DICT, APPENDS, SETITEM + 0x74 | 0x75 | 0x85 | 0x86 | 0x87 | // TUPLE, SETITEMS, TUPLE1, TUPLE2, TUPLE3 + 0x88 | 0x89 | 0x8f | 0x90 | 0x91 | // NEWTRUE, NEWFALSE, STACK_GLOBAL, MEMOIZE, EMPTY_SET + 0x92 | 0x93 | 0x94 | 0x51 | 0x81 => { // ADDITEMS, FROZENSET, NEWOBJ, BINPERSID, NEWOBJ_EX + // No additional data to skip + } + _ => { + // Unknown opcode - assume no additional data + // This is a best-effort approach + } + } + } + + Ok(()) +} + +pub fn read_pickle_with_data( + r: &mut R, + data_source: Arc, +) -> Result { + read_pickle_with_optional_data(r, Some(data_source)) +} + +fn get_dict_key(obj: Object) -> Result { + match obj { + Object::String(s) => Ok(s), + Object::Int(i) => Ok(i.to_string()), + _ => Err(PickleError::InvalidData(format!( + "dict key must be a valid type, got {obj:?}" + ))), + } +} + +pub fn read_pickle_with_optional_data( + r: &mut R, + data_source: Option>, +) -> Result { + let mut stack = match data_source { + Some(ds) => Stack::with_data_source(ds), + None => Stack::new(), + }; + loop { + let op_code = r.read_u8()?; + let op_code = OpCode::try_from(op_code).map_err(PickleError::InvalidOpCode)?; + match op_code { + OpCode::Proto => { + let version = r.read_u8()?; + if version > 5 { + return Err(PickleError::InvalidProtocol(version)); + } + } + OpCode::Global => read_global(r, &mut stack)?, + OpCode::BinInt => read_bin_int(r, &mut stack)?, + OpCode::Int => read_int(r, &mut stack)?, + OpCode::BinInt1 => read_bin_int1(r, &mut stack)?, + OpCode::BinInt2 => read_bin_int2(r, &mut stack)?, + OpCode::BinFloat => read_bin_float(r, &mut stack)?, + OpCode::BinUnicode => { + let len = r.read_u32::()? as usize; + read_string(r, &mut stack, len)? + } + OpCode::ShortBinString => { + let len = r.read_u8()? as usize; + read_string(r, &mut stack, len)? + } + OpCode::Long1 => read_long1(r, &mut stack)?, + OpCode::None => stack.push(Object::None), + OpCode::NewTrue => stack.push(Object::Bool(true)), + OpCode::NewFalse => stack.push(Object::Bool(false)), + OpCode::EmptyTuple => stack.push(Object::Tuple(Vec::new())), + OpCode::EmptyList => stack.push(Object::List(Vec::new())), + OpCode::EmptyDict => stack.push(Object::Dict(HashMap::new())), + OpCode::Tuple => { + let objs = stack.pop_to_marker()?; + stack.push(Object::Tuple(objs)) + } + OpCode::Tuple1 => { + let obj = stack.pop()?; + stack.push(Object::Tuple(vec![obj])) + } + OpCode::Tuple2 => { + let obj2 = stack.pop()?; + let obj1 = stack.pop()?; + stack.push(Object::Tuple(vec![obj1, obj2])) + } + OpCode::Tuple3 => { + let obj3 = stack.pop()?; + let obj2 = stack.pop()?; + let obj1 = stack.pop()?; + stack.push(Object::Tuple(vec![obj1, obj2, obj3])) + } + OpCode::Append => { + let value = stack.pop()?; + match stack.last_mut()? { + Object::List(list) => list.push(value), + _ => return Err(PickleError::UnexpectedOpCode(op_code)), + } + } + OpCode::Appends => { + let objs = stack.pop_to_marker()?; + match stack.last_mut()? { + Object::List(list) => list.extend(objs), + _ => return Err(PickleError::UnexpectedOpCode(op_code)), + } + } + OpCode::SetItem => { + let value = stack.pop()?; + let key = stack.pop()?; + match stack.last_mut()? { + Object::Dict(dict) => { + if let Object::String(key) = key { + dict.insert(key, value); + } else { + return Err(PickleError::InvalidData( + "dict key must be a string".to_string(), + )); + } + } + _ => return Err(PickleError::UnexpectedOpCode(op_code)), + } + } + OpCode::SetItems => { + let mut objs = stack.pop_to_marker()?; + if objs.len() % 2 != 0 { + return Err(PickleError::InvalidData( + "setitems requires even number of objects".to_string(), + )); + } + match stack.last_mut()? { + Object::Dict(dict) => { + while !objs.is_empty() { + let key = objs.remove(0); + let value = objs.remove(0); + let key = get_dict_key(key)?; + dict.insert(key, value); + } + } + _ => return Err(PickleError::UnexpectedOpCode(op_code)), + } + } + OpCode::BinPut => { + let idx = r.read_u8()? as u32; + let obj = stack.top()?; + stack.memo_put(idx, obj); + } + OpCode::LongBinPut => { + let idx = r.read_u32::()?; + let obj = stack.top()?; + stack.memo_put(idx, obj); + } + OpCode::BinGet => { + let idx = r.read_u8()? as u32; + let obj = stack.memo_get(idx)?; + stack.push(obj); + } + OpCode::LongBinGet => { + let idx = r.read_u32::()?; + let obj = stack.memo_get(idx)?; + stack.push(obj); + } + OpCode::Mark => stack.push_mark(), + OpCode::BinPersId => { + let pid = stack.pop()?; + match pid { + Object::String(s) => { + stack.push(Object::Persistent(s.into_bytes())); + } + Object::Tuple(tuple) => { + // The persistent ID is a tuple (e.g., ('storage', 'FloatStorage', '0', 'cpu', 4)) + // Store it as a PersistentTuple for proper handling + stack.push(Object::PersistentTuple(tuple)); + } + _ => { + return Err(PickleError::InvalidData(format!( + "persistent id must be a string or tuple, got {:?}", + pid + ))); + } + } + } + OpCode::Reduce => { + let args = stack.pop()?; + let callable = stack.pop()?; + + // Check if this is an OrderedDict + if let Object::Class { module_name, name } = &callable { + if module_name == "collections" && name == "OrderedDict" { + // OrderedDict can be created with items: OrderedDict([(key1, val1), ...]) + // The args is typically a tuple containing a list of [key, value] pairs + let mut dict = HashMap::new(); + + // Extract items from args + let items = match &args { + Object::Tuple(tuple) if !tuple.is_empty() => { + // Args is a tuple, get the first element (the list of items) + match &tuple[0] { + Object::List(list) => Some(list.clone()), + _ => None, + } + } + Object::List(list) => Some(list.clone()), + _ => None, + }; + + if let Some(items) = items { + for item in items { + // Each item is a list/tuple of [key, value] + match item { + Object::List(pair) | Object::Tuple(pair) if pair.len() >= 2 => { + if let Object::String(key) = &pair[0] { + dict.insert(key.clone(), pair[1].clone()); + } + } + _ => {} + } + } + } + + stack.push(Object::Dict(dict)); + } else { + let _obj = Object::Reduce { + callable: Box::new(callable.clone()), + args: Box::new(args.clone()), + }; + let obj = rebuild_from_type_v2( + Object::Tuple(vec![callable, args]), + &mut stack.memo, + &stack.data_source, + )?; + stack.push(obj); + } + } else { + let _obj = Object::Reduce { + callable: Box::new(callable.clone()), + args: Box::new(args.clone()), + }; + let obj = rebuild_from_type_v2( + Object::Tuple(vec![callable, args]), + &mut stack.memo, + &stack.data_source, + )?; + stack.push(obj); + } + } + OpCode::Build => { + let args = stack.pop()?; + let obj = stack.pop()?; + match obj { + Object::Dict(mut dict) => { + // For dicts, BUILD updates with the args + if let Object::Dict(update) = args { + dict.extend(update); + } + stack.push(Object::Dict(dict)); + } + _ => { + stack.push(Object::Build { + callable: Box::new(obj), + args: Box::new(args), + }); + } + } + } + OpCode::NewObj => { + let args = stack.pop()?; + let cls = stack.pop()?; + stack.push(Object::Reduce { + callable: Box::new(cls), + args: Box::new(args), + }); + } + OpCode::Dict => { + let objs = stack.pop_to_marker()?; + let mut dict = HashMap::new(); + if objs.len() % 2 != 0 { + return Err(PickleError::InvalidData( + "dict requires even number of objects".to_string(), + )); + } + for chunk in objs.chunks(2) { + let key = get_dict_key(chunk[0].clone())?; + dict.insert(key, chunk[1].clone()); + } + stack.push(Object::Dict(dict)); + } + OpCode::List => { + let objs = stack.pop_to_marker()?; + stack.push(Object::List(objs)); + } + OpCode::Memoize => { + // Store top of stack in memo without popping + // The memo index is the current number of items in the memo + let obj = stack.top()?; + let idx = stack.memo_len() as u32; + stack.memo_put(idx, obj); + } + OpCode::Stop => break, + } + } + stack.pop() +} + +/// Load tensors from a pickle file (PyTorch checkpoint format) +pub fn read_pickle_tensors(reader: &mut R) -> Result> { + let obj = read_pickle(reader)?; + + // Extract tensors from the loaded object + let mut tensors = HashMap::new(); + let mut path = Vec::new(); + extract_tensors(&obj, &mut path, &mut tensors); + + Ok(tensors) +} + +fn extract_tensors<'a>( + obj: &'a Object, + path: &mut Vec<&'a str>, + tensors: &mut HashMap, +) { + match obj { + Object::Dict(dict) => { + for (key, value) in dict { + path.push(key); + extract_tensors(value, path, tensors); + path.pop(); + } + } + Object::TorchParam(snapshot) => { + // Only allocate the string here when we actually insert + tensors.insert(path.join("."), snapshot.clone()); + } + _ => {} + } +} diff --git a/crates/burn-store/src/pytorch/reader.rs b/crates/burn-store/src/pytorch/reader.rs new file mode 100644 index 0000000..8b3f4d1 --- /dev/null +++ b/crates/burn-store/src/pytorch/reader.rs @@ -0,0 +1,1125 @@ +//! PyTorch file reader implementation. +//! +//! This module provides support for reading PyTorch checkpoint files (.pt/.pth). +//! +//! # Supported Formats +//! +//! ## 1. Modern ZIP Format (PyTorch 1.6+) +//! Files are ZIP archives containing: +//! - `data.pkl` or `archive/data.pkl`: Pickled tensor metadata +//! - `data/` directory: Binary tensor data files +//! +//! ## 2. TAR Format (older torchvision models like AlexNet, SqueezeNet) +//! TAR archives containing: +//! - `sys_info`: System info pickle (endianness, type sizes) +//! - `pickle`: OrderedDict mapping tensor names to storage keys +//! - `tensors`: Tensor metadata (unused, metadata is in pickle) +//! - `storages`: Count pickle + sequential (metadata, num_elements, raw data) +//! +//! ## 3. Legacy Pickle Format (PyTorch 0.1.10 - 1.5) +//! Sequential pickle streams with the structure: +//! - Magic number pickle (0x1950a86a20f9469cfc6c) +//! - Protocol version pickle (e.g., 1001) +//! - System info pickle (endianness, type sizes) +//! - Model data pickle (state_dict or full model) +//! +//! ## 4. Simple Pickle Format +//! Direct pickle file with a dictionary at the root, commonly used for +//! manually saved state_dicts. +//! +//! # Compatibility +//! +//! The reader handles backward compatibility by detecting the file format +//! automatically. Files from PyTorch 0.1.10 through current versions are +//! supported, though full model saves (vs state_dict) may have limitations +//! as they contain Python code references. + +use crate::TensorSnapshot; +use crate::nested::{adapter::DefaultAdapter, data::NestedValue, de::Deserializer}; +use alloc::string::{String, ToString}; +use alloc::vec::Vec; +use serde::de::DeserializeOwned; +use std::collections::HashMap; +use std::fs::File; +use std::io::{BufReader, Read, Seek, SeekFrom}; +use std::path::Path; + +use super::lazy_data::LazyDataSource; +use super::pickle_reader::{Object, PickleError, read_pickle, read_pickle_with_data}; +use std::sync::Arc; + +/// Error type for PyTorch file operations +#[derive(Debug)] +pub enum PytorchError { + /// IO error + Io(std::io::Error), + /// Pickle parsing error + Pickle(PickleError), + /// Zip archive error + Zip(zip::result::ZipError), + /// TAR archive error + Tar(std::io::Error), + /// Invalid file format + InvalidFormat(String), + /// Key not found + KeyNotFound(String), + /// Serde deserialization error + Serde(crate::nested::error::Error), +} + +impl From for PytorchError { + fn from(e: std::io::Error) -> Self { + PytorchError::Io(e) + } +} + +impl From for PytorchError { + fn from(e: PickleError) -> Self { + PytorchError::Pickle(e) + } +} + +impl From for PytorchError { + fn from(e: zip::result::ZipError) -> Self { + PytorchError::Zip(e) + } +} + +impl From for PytorchError { + fn from(e: crate::nested::error::Error) -> Self { + PytorchError::Serde(e) + } +} + +impl std::fmt::Display for PytorchError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + PytorchError::Io(e) => write!(f, "IO error: {}", e), + PytorchError::Pickle(e) => write!( + f, + "Pickle parsing error: {}. This may indicate an unsupported PyTorch file format or corrupted file.", + e + ), + PytorchError::Zip(e) => write!(f, "Zip archive error: {}", e), + PytorchError::Tar(e) => write!(f, "TAR archive error: {}", e), + PytorchError::InvalidFormat(msg) => write!(f, "Invalid PyTorch file format: {}", msg), + PytorchError::KeyNotFound(key) => write!( + f, + "Key '{}' not found in PyTorch file. Available keys may be listed with the keys() method.", + key + ), + PytorchError::Serde(e) => write!(f, "Serde deserialization error: {}", e), + } + } +} + +impl std::error::Error for PytorchError {} + +type Result = std::result::Result; + +/// Metadata about a PyTorch file +/// +/// Contains information about the file format, version, and other properties +/// that can be useful for debugging or compatibility checking. +#[derive(Debug, Clone)] +pub struct PytorchMetadata { + /// Format version (e.g., "1.0" for modern ZIP format) + pub format_version: Option, + /// File format type (ZIP, Legacy, or Pickle) + pub format_type: FileFormat, + /// Byte order (endianness) - currently only LittleEndian is supported + pub byte_order: ByteOrder, + /// Whether the file has storage alignment information + pub has_storage_alignment: bool, + /// PyTorch version that saved the file (if available) + pub pytorch_version: Option, + /// Number of tensors in the file + pub tensor_count: usize, + /// Total size of tensor data in bytes (if available) + pub total_data_size: Option, +} + +impl PytorchMetadata { + /// Check if this is a modern format file (ZIP-based, PyTorch 1.6+) + pub fn is_modern_format(&self) -> bool { + matches!(self.format_type, FileFormat::Zip) + } + + /// Check if this is a legacy format file (PyTorch 0.1.10 - 1.5) + pub fn is_legacy_format(&self) -> bool { + matches!(self.format_type, FileFormat::Legacy) + } +} + +/// File format type +#[derive(Debug, Clone, PartialEq)] +pub enum FileFormat { + /// ZIP-based format (PyTorch 1.6+) + Zip, + /// TAR-based format (older torchvision models) + Tar, + /// Legacy format (PyTorch 0.1.10 - 1.5) + Legacy, + /// Simple pickle file + Pickle, +} + +/// Byte order (endianness) +#[derive(Debug, Clone, PartialEq)] +pub enum ByteOrder { + LittleEndian, + BigEndian, +} + +/// PyTorch checkpoint reader +/// +/// This is the main interface for reading PyTorch checkpoint files (.pt/.pth). +/// It supports multiple PyTorch formats including modern ZIP-based format (1.6+), +/// legacy format (0.1.10-1.5), and simple pickle files. +/// +/// # Example +/// ```rust,no_run +/// # use burn_store::pytorch::PytorchReader; +/// # fn example() -> Result<(), Box> { +/// // Load a checkpoint file +/// let reader = PytorchReader::new("model.pt")?; +/// +/// // Get tensor names +/// let keys = reader.keys(); +/// +/// // Access a specific tensor +/// if let Some(tensor) = reader.get("conv1.weight") { +/// let data = tensor.to_data(); // Materializes the tensor +/// } +/// +/// // Check file metadata +/// println!("Format: {:?}", reader.metadata().format_type); +/// println!("Tensor count: {}", reader.metadata().tensor_count); +/// # Ok(()) +/// # } +/// ``` +pub struct PytorchReader { + tensors: HashMap, + metadata: PytorchMetadata, +} + +impl PytorchReader { + /// Load a PyTorch checkpoint file + /// + /// # Arguments + /// * `path` - Path to the PyTorch file (.pt or .pth) + /// + /// # Returns + /// A `PytorchReader` with lazy-loaded tensors and metadata + pub fn new>(path: P) -> Result { + let (tensors, metadata) = load_pytorch_file_with_metadata(path.as_ref(), None)?; + Ok(Self { tensors, metadata }) + } + + /// Load a PyTorch checkpoint with a specific top-level key + /// + /// Many PyTorch checkpoints store the model weights under a specific key + /// like "state_dict", "model", or "model_state_dict". + /// + /// # Arguments + /// * `path` - Path to the PyTorch file + /// * `key` - Top-level key to extract (e.g., "state_dict") + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::pytorch::PytorchReader; + /// # fn example() -> Result<(), Box> { + /// let reader = PytorchReader::with_top_level_key("checkpoint.pt", "state_dict")?; + /// # Ok(()) + /// # } + /// ``` + pub fn with_top_level_key>(path: P, key: &str) -> Result { + let (tensors, metadata) = load_pytorch_file_with_metadata(path.as_ref(), Some(key))?; + Ok(Self { tensors, metadata }) + } + + /// Load from a reader + /// + /// This method is useful when loading from non-file sources like memory buffers. + /// Note: Metadata detection is limited when loading from a reader. + /// + /// # Arguments + /// * `reader` - Any type implementing `Read` + /// * `top_level_key` - Optional key to extract + pub fn from_reader(reader: R, top_level_key: Option<&str>) -> Result { + // For reader-based loading, we don't have full metadata access + let tensors = load_from_reader(reader, top_level_key)?; + let metadata = PytorchMetadata { + format_version: None, + format_type: FileFormat::Pickle, // Default assumption + byte_order: ByteOrder::LittleEndian, + has_storage_alignment: false, + pytorch_version: None, + tensor_count: tensors.len(), + total_data_size: None, + }; + Ok(Self { tensors, metadata }) + } + + /// Get all tensor names + pub fn keys(&self) -> Vec { + self.tensors.keys().cloned().collect() + } + + /// Get a tensor by name + pub fn get(&self, name: &str) -> Option<&TensorSnapshot> { + self.tensors.get(name) + } + + /// Get all tensors + pub fn tensors(&self) -> &HashMap { + &self.tensors + } + + /// Take ownership of all tensors + pub fn into_tensors(self) -> HashMap { + self.tensors + } + + /// Get metadata about the loaded file + /// + /// Provides information about the file format, version, endianness, etc. + pub fn metadata(&self) -> &PytorchMetadata { + &self.metadata + } + + /// Get the number of tensors in the file + pub fn len(&self) -> usize { + self.tensors.len() + } + + /// Check if the file contains no tensors + pub fn is_empty(&self) -> bool { + self.tensors.is_empty() + } + + /// Read raw pickle data from a PyTorch file + /// + /// This is useful for extracting configuration or metadata that isn't tensor data. + /// Returns a simplified JSON-like structure that can be easily converted to other formats. + /// + /// # Arguments + /// * `path` - Path to the PyTorch file + /// * `top_level_key` - Optional key to extract from the top-level dictionary + /// + /// # Returns + /// A `PickleValue` representing the pickle data structure + pub fn read_pickle_data>( + path: P, + top_level_key: Option<&str>, + ) -> Result { + read_pickle_as_value(path.as_ref(), top_level_key) + } + + /// Load and deserialize configuration data from a PyTorch file + /// + /// This method reads configuration or metadata stored in PyTorch checkpoint files + /// and deserializes it into the specified type. It's particularly useful for + /// extracting model configurations that might be saved alongside model weights. + /// + /// # Arguments + /// * `path` - Path to the PyTorch file (.pt or .pth) + /// * `top_level_key` - Optional key to extract specific data within the pickle file. + /// If `None`, the entire content is deserialized. + /// + /// # Type Parameters + /// * `D` - The target type to deserialize into. Must implement `DeserializeOwned`. + /// + /// # Returns + /// A `Result` containing the deserialized configuration data, or an `Error` if + /// reading or deserialization fails. + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::pytorch::PytorchReader; + /// # use serde::Deserialize; + /// # fn example() -> Result<(), Box> { + /// #[derive(Debug, Deserialize)] + /// struct ModelConfig { + /// hidden_size: usize, + /// num_layers: usize, + /// } + /// + /// let config: ModelConfig = PytorchReader::load_config("model.pth", Some("config"))?; + /// # Ok(()) + /// # } + /// ``` + pub fn load_config(path: P, top_level_key: Option<&str>) -> Result + where + D: DeserializeOwned, + P: AsRef, + { + // Read the PyTorch file and extract the pickle data + let pickle_value = Self::read_pickle_data(path, top_level_key)?; + + // Convert PickleValue to NestedValue + let nested_value = convert_pickle_to_nested_value(pickle_value)?; + + // Create a deserializer with the default adapter + let deserializer = Deserializer::::new(nested_value, false); + + // Deserialize the nested value into the target type + let value = D::deserialize(deserializer)?; + Ok(value) + } +} + +/// Simplified representation of pickle data +/// +/// This enum provides a JSON-like structure that's easier to work with +/// than the internal pickle Object type. +#[derive(Debug, Clone, PartialEq)] +pub enum PickleValue { + /// None/null value + None, + /// Boolean value + Bool(bool), + /// Integer value + Int(i64), + /// Floating point value + Float(f64), + /// String value + String(String), + /// List/array of values + List(Vec), + /// Dictionary/map of string keys to values + Dict(HashMap), + /// Binary data + Bytes(Vec), +} + +/// Internal function to load a PyTorch file with metadata +fn load_pytorch_file_with_metadata( + path: &Path, + top_level_key: Option<&str>, +) -> Result<(HashMap, PytorchMetadata)> { + // First, try to read as a zip file + if let Ok(file) = File::open(path) + && let Ok(mut archive) = zip::ZipArchive::new(BufReader::new(file)) + { + // PyTorch saves the main data in various locations within the zip + let mut pickle_data = Vec::new(); + let mut pickle_found = false; + + // Try different common pickle file locations + let possible_pickle_paths = [ + "data.pkl", + "archive/data.pkl", + // Look for any .pkl file in the root or first-level directories + ]; + + for pickle_path in &possible_pickle_paths { + if archive.by_name(pickle_path).is_ok() { + let mut pickle_file = archive.by_name(pickle_path)?; + pickle_file.read_to_end(&mut pickle_data)?; + pickle_found = true; + break; + } + } + + // If not found in standard locations, search for any .pkl file + if !pickle_found { + for i in 0..archive.len() { + let file = archive.by_index(i)?; + let name = file.name().to_string(); + drop(file); // Release the borrow + + if name.ends_with("data.pkl") { + let mut file = archive.by_index(i)?; + file.read_to_end(&mut pickle_data)?; + pickle_found = true; + break; + } + } + } + + if !pickle_found { + return Err(PytorchError::InvalidFormat( + "No data.pkl file found in ZIP archive. Expected PyTorch 1.6+ format with data.pkl or archive/data.pkl".to_string(), + )); + } + + // Check for format version (optional) + let format_version = if let Ok(mut version_file) = archive.by_name(".format_version") { + let mut version_data = Vec::new(); + version_file.read_to_end(&mut version_data)?; + let version_str = String::from_utf8_lossy(&version_data); + let version = version_str.trim().to_string(); + Some(version) + } else { + None + }; + + // Check for byteorder file to detect endianness + let is_big_endian = if let Ok(mut byteorder_file) = archive.by_name("byteorder") { + let mut byteorder_data = Vec::new(); + byteorder_file.read_to_end(&mut byteorder_data)?; + let byteorder_str = String::from_utf8_lossy(&byteorder_data); + byteorder_str.trim() == "big" + } else { + false // Default to little-endian if no byteorder file + }; + + if is_big_endian { + // Big-endian files are not yet supported as they require different byte order conversion + // TODO: To support big-endian files, we need to: + // 1. Pass endianness info through to pickle_reader + // 2. Use from_be_bytes instead of from_le_bytes for tensor data + // 3. Handle byte swapping for all numeric types (f32, f64, i32, etc.) + return Err(PytorchError::InvalidFormat( + "Big-endian PyTorch files are not yet supported. The file was saved on a big-endian system and requires byte order conversion.".to_string() + )); + } + + // Check for storage alignment file + let has_storage_alignment = archive.by_name(".storage_alignment").is_ok(); + + // Check for PyTorch version (if saved) + let pytorch_version = if let Ok(mut version_file) = archive.by_name("version") { + let mut version_data = Vec::new(); + version_file.read_to_end(&mut version_data)?; + Some(String::from_utf8_lossy(&version_data).trim().to_string()) + } else { + None + }; + + // Create a lazy data source instead of loading all data upfront + let data_source = Arc::new(LazyDataSource::from_zip(path)?); + + // Calculate total data size without loading + let mut total_data_size = 0usize; + for i in 0..archive.len() { + let file = archive.by_index(i)?; + let name = file.name(); + + // Look for data files - they can be in various locations + let is_data_file = (name.contains("/data/") + || name.starts_with("data/") + || name.starts_with("archive/data/")) + && !name.ends_with(".pkl") + && !name.ends_with("/"); + + if is_data_file { + total_data_size += file.size() as usize; + } + } + + // Parse the pickle data with lazy data source + let mut pickle_reader = BufReader::new(pickle_data.as_slice()); + let obj = read_pickle_with_data(&mut pickle_reader, data_source)?; + + // Extract tensors with their data + let tensors = extract_tensors_with_data(obj, top_level_key)?; + + // Create metadata + let metadata = PytorchMetadata { + format_version, + format_type: FileFormat::Zip, + byte_order: if is_big_endian { + ByteOrder::BigEndian + } else { + ByteOrder::LittleEndian + }, + has_storage_alignment, + pytorch_version, + tensor_count: tensors.len(), + total_data_size: Some(total_data_size), + }; + + return Ok((tensors, metadata)); + } + + // If not a zip or zip reading failed, try TAR format + if is_tar_file(path) { + return load_tar_pytorch_file_with_metadata(path, top_level_key); + } + + // Try reading as a plain pickle file + let mut file = File::open(path)?; + + // Check for PyTorch legacy format (starts with magic number as pickled integer) + let mut header = [0u8; 15]; + // Use read() instead of read_exact() to handle files smaller than 15 bytes + let bytes_read = file.read(&mut header)?; + file.seek(std::io::SeekFrom::Start(0))?; + + // Only check for legacy format if we have enough bytes + // PyTorch legacy format detection (PyTorch 0.1.10 - 1.3) + // Reference: https://github.com/pytorch/pytorch/blob/main/torch/serialization.py#L65 + // + // These files use sequential pickle streams with metadata before the actual data. + // Format structure: + // 1. Magic number (0x1950a86a20f9469cfc6c) stored as LONG1 pickle + // 2. Protocol version (e.g., 1001) + // 3. System info dict (protocol_version, little_endian, type_sizes) + // 4. Actual model data (state_dict or full model) + // 5. Storage keys list (pickle) + // 6. Raw binary data for each storage + // + // The pattern is: 0x80 0x02 0x8a 0x0a (PROTO 2, LONG1 with 10 bytes) + // followed by 10 bytes of magic number (little-endian), then 0x2e (STOP) + let is_legacy_format = bytes_read >= 15 + && header[0] == 0x80 // PROTO opcode + && header[1] == 0x02 // Protocol version 2 + && header[2] == 0x8a // LONG1 opcode + && header[3] == 0x0a // 10 bytes follow + // Magic number 0x1950a86a20f9469cfc6c in little-endian + && header[4] == 0x6c + && header[5] == 0xfc + && header[6] == 0x9c + && header[7] == 0x46 + && header[8] == 0xf9 + && header[9] == 0x20 + && header[10] == 0x6a + && header[11] == 0xa8 + && header[12] == 0x50 + && header[13] == 0x19 + && header[14] == 0x2e; // STOP opcode + + if is_legacy_format { + return load_legacy_pytorch_file_with_metadata(path, top_level_key); + } + + // Standard pickle file + // This might be a pickle with tensor references, so we need to handle that case + // For plain pickle files without a separate data section, we can't use lazy loading + // so we'll just create empty placeholder tensors for the structure + let file = File::open(path)?; + let mut reader = BufReader::new(file); + + // Try reading without data source first + match read_pickle(&mut reader) { + Ok(obj) => { + let tensors = extract_tensors_with_data(obj, top_level_key)?; + let tensor_count = tensors.len(); + Ok(( + tensors, + PytorchMetadata { + format_version: None, + format_type: FileFormat::Pickle, + byte_order: ByteOrder::LittleEndian, + has_storage_alignment: false, + pytorch_version: None, + tensor_count, + total_data_size: None, + }, + )) + } + Err(e) + if e.to_string() + .contains("Cannot load tensor data without a data source") => + { + // This pickle file contains tensor data but we're trying to read it without + // providing a data source. This shouldn't happen in normal usage as PyTorch + // files with actual tensor data should be in ZIP or legacy format. + Err(PytorchError::InvalidFormat( + "Pickle file contains tensor data but no data source is available. This file should be loaded as ZIP or legacy format.".to_string() + )) + } + Err(e) => Err(PytorchError::Pickle(e)), + } +} + +/// Load from a reader +fn load_from_reader( + reader: R, + top_level_key: Option<&str>, +) -> Result> { + let mut buf_reader = BufReader::new(reader); + + // Try reading without data source + match read_pickle(&mut buf_reader) { + Ok(obj) => extract_tensors_with_data(obj, top_level_key), + Err(e) + if e.to_string() + .contains("Cannot load tensor data without a data source") => + { + // This reader contains tensor data but we can't load it without a file path + Err(PytorchError::InvalidFormat( + "Reader contains tensor data but no data source is available. Use file-based loading instead.".to_string() + )) + } + Err(e) => Err(PytorchError::Pickle(e)), + } +} + +/// Extract tensors from a parsed pickle object +fn extract_tensors_with_data( + obj: Object, + top_level_key: Option<&str>, +) -> Result> { + let dict = match obj { + Object::Dict(dict) => { + if let Some(key) = top_level_key { + // Extract the nested dictionary if a top-level key is specified + match dict.get(key) { + Some(Object::Dict(nested)) => nested.clone(), + _ => { + return Err(PytorchError::KeyNotFound(format!( + "Top-level key '{}' not found or is not a dictionary. Available top-level keys in file: {:?}", + key, + dict.keys().collect::>() + ))); + } + } + } else { + dict + } + } + _ => { + return Err(PytorchError::InvalidFormat( + "Expected a dictionary at the root of the PyTorch file, but found a different type. The file may be a full model save rather than a state_dict.".to_string(), + )); + } + }; + + let mut tensors = HashMap::new(); + let mut path = Vec::new(); + extract_tensors_recursive(&Object::Dict(dict), &mut path, &mut tensors); + Ok(tensors) +} + +/// Recursively extract tensors from an object +fn extract_tensors_recursive<'a>( + obj: &'a Object, + path: &mut Vec<&'a str>, + tensors: &mut HashMap, +) { + match obj { + Object::Dict(dict) => { + for (key, value) in dict { + path.push(key); + extract_tensors_recursive(value, path, tensors); + path.pop(); + } + } + Object::TorchParam(snapshot) => { + // The TensorSnapshot already contains the data loading closure + // Only allocate the string here when we actually insert + tensors.insert(path.join("."), snapshot.clone()); + } + _ => {} + } +} + +/// Load a legacy PyTorch file with metadata +fn load_legacy_pytorch_file_with_metadata( + path: &Path, + top_level_key: Option<&str>, +) -> Result<(HashMap, PytorchMetadata)> { + let file = File::open(path)?; + let mut reader = BufReader::new(file); + + // Skip metadata pickles + // 1. Magic number + let _ = read_pickle(&mut reader).map_err(|e| { + PytorchError::InvalidFormat(format!( + "Failed to read magic number from legacy format: {}", + e + )) + })?; + + // 2. Protocol version + let _ = read_pickle(&mut reader).map_err(|e| { + PytorchError::InvalidFormat(format!( + "Failed to read protocol version from legacy format: {}", + e + )) + })?; + + // 3. System info + let _ = read_pickle(&mut reader).map_err(|e| { + PytorchError::InvalidFormat(format!( + "Failed to read system info from legacy format: {}", + e + )) + })?; + + // Save position before main pickle + let main_pickle_pos = reader.stream_position()?; + + // 4. Skip main object - it might contain tensors so we can't parse it yet + // We'll re-read it with a data source later + use crate::pytorch::pickle_reader::skip_pickle; + skip_pickle(&mut reader).map_err(|e| { + PytorchError::InvalidFormat(format!( + "Failed to skip main object in legacy format: {}", + e + )) + })?; + + // 5. Storage keys list (sorted keys as written by PyTorch) + let storage_keys = match read_pickle(&mut reader) { + Ok(Object::List(keys)) => keys + .into_iter() + .filter_map(|obj| match obj { + Object::String(s) => Some(s), + _ => None, + }) + .collect::>(), + _ => vec![], + }; + + // 6. Raw binary data starts here + let data_start_pos = reader.stream_position()?; + let file_size = reader.seek(SeekFrom::End(0))?; + let data_size = file_size - data_start_pos; + + // Create a lazy data source for legacy multi-storage format + let data_source = Arc::new(LazyDataSource::from_legacy_multi_storage( + path, + data_start_pos, + data_size, + )); + + // Set storage keys BEFORE parsing the main pickle + // This is critical because track_storage_usage() is called during parsing + // and it needs storage_keys to build the storage map + if let LazyDataSource::LegacyMultiStorage(ref source) = *data_source + && !storage_keys.is_empty() + { + let source = source + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + source.set_storage_keys(storage_keys.clone()); + } + + // Now re-read the main pickle with lazy data source + reader.seek(SeekFrom::Start(main_pickle_pos))?; + let main_obj = read_pickle_with_data(&mut reader, data_source.clone())?; + + // Extract tensors normally + let tensors = extract_tensors_with_data(main_obj, top_level_key)?; + + // Create metadata for legacy format + let metadata = PytorchMetadata { + format_version: None, // Legacy format doesn't have version files + format_type: FileFormat::Legacy, + byte_order: ByteOrder::LittleEndian, // Legacy format is little-endian + has_storage_alignment: false, + pytorch_version: None, // Could parse from protocol version, but not reliable + tensor_count: tensors.len(), + total_data_size: Some(data_size as usize), + }; + + Ok((tensors, metadata)) +} + +/// Check if a file is a TAR archive +fn is_tar_file(path: &Path) -> bool { + if let Ok(mut file) = File::open(path) { + // TAR files have "ustar" magic at offset 257 + let mut header = [0u8; 263]; + if file.read_exact(&mut header).is_ok() { + // Check for "ustar" magic at offset 257 + return &header[257..262] == b"ustar"; + } + } + false +} + +/// Load a TAR format PyTorch file with metadata +fn load_tar_pytorch_file_with_metadata( + path: &Path, + top_level_key: Option<&str>, +) -> Result<(HashMap, PytorchMetadata)> { + use tar::Archive; + + let file = File::open(path)?; + let mut archive = Archive::new(BufReader::new(file)); + + // Extract the main entries from the TAR archive + let mut sys_info_data: Option> = None; + let mut pickle_data: Option> = None; + let mut storages_data: Option> = None; + + for entry in archive.entries().map_err(PytorchError::Tar)? { + let mut entry = entry.map_err(PytorchError::Tar)?; + let entry_path = entry + .path() + .map_err(PytorchError::Tar)? + .to_string_lossy() + .to_string(); + + // Skip PAX headers + if entry_path.contains("@PaxHeader") { + continue; + } + + // Normalize path (remove ./ prefix if present) + let normalized = entry_path.trim_start_matches("./"); + + match normalized { + "sys_info" => { + let mut data = Vec::new(); + entry.read_to_end(&mut data).map_err(PytorchError::Tar)?; + sys_info_data = Some(data); + } + "pickle" => { + let mut data = Vec::new(); + entry.read_to_end(&mut data).map_err(PytorchError::Tar)?; + pickle_data = Some(data); + } + "storages" => { + let mut data = Vec::new(); + entry.read_to_end(&mut data).map_err(PytorchError::Tar)?; + storages_data = Some(data); + } + _ => {} + } + } + + // Validate required entries + let pickle_data = pickle_data.ok_or_else(|| { + PytorchError::InvalidFormat("TAR file missing 'pickle' entry".to_string()) + })?; + let storages_data = storages_data.ok_or_else(|| { + PytorchError::InvalidFormat("TAR file missing 'storages' entry".to_string()) + })?; + + // Parse sys_info to check endianness + let is_little_endian = if let Some(ref data) = sys_info_data { + parse_tar_sys_info(data)? + } else { + true // Default to little-endian + }; + + if !is_little_endian { + return Err(PytorchError::InvalidFormat( + "Big-endian TAR PyTorch files are not supported".to_string(), + )); + } + + // Create TarSource for lazy loading + let data_source = Arc::new(LazyDataSource::from_tar(&storages_data)?); + + // Parse the pickle (OrderedDict of name -> storage_key) + let mut pickle_reader = BufReader::new(pickle_data.as_slice()); + let obj = read_pickle_with_data(&mut pickle_reader, data_source)?; + + // Extract tensors + let tensors = extract_tensors_with_data(obj, top_level_key)?; + + let metadata = PytorchMetadata { + format_version: None, + format_type: FileFormat::Tar, + byte_order: ByteOrder::LittleEndian, + has_storage_alignment: false, + pytorch_version: None, + tensor_count: tensors.len(), + total_data_size: Some(storages_data.len()), + }; + + Ok((tensors, metadata)) +} + +/// Parse sys_info pickle from TAR format to extract endianness +fn parse_tar_sys_info(data: &[u8]) -> Result { + let mut reader = BufReader::new(data); + let obj = read_pickle(&mut reader)?; + + if let Object::Dict(dict) = obj + && let Some(Object::Bool(little_endian)) = dict.get("little_endian") + { + return Ok(*little_endian); + } + + Ok(true) // Default assumption +} + +/// Read pickle data from a PyTorch file as a simplified value +fn read_pickle_as_value(path: &Path, top_level_key: Option<&str>) -> Result { + use crate::pytorch::lazy_data::LazyDataSource; + use crate::pytorch::pickle_reader::{read_pickle, read_pickle_with_data}; + use std::sync::Arc; + + // Try to open as ZIP first + if let Ok(file) = File::open(path) + && let Ok(mut archive) = zip::ZipArchive::new(BufReader::new(file)) + { + // Read pickle data from ZIP + let mut pickle_data = Vec::new(); + + // Try standard locations + for pickle_path in &["data.pkl", "archive/data.pkl"] { + if let Ok(mut pickle_file) = archive.by_name(pickle_path) { + pickle_file.read_to_end(&mut pickle_data)?; + break; + } + } + + // If not found, search for any .pkl file + if pickle_data.is_empty() { + for i in 0..archive.len() { + let file = archive.by_index(i)?; + let name = file.name().to_string(); + drop(file); + + if name.ends_with("data.pkl") { + let mut file = archive.by_index(i)?; + file.read_to_end(&mut pickle_data)?; + break; + } + } + } + + if !pickle_data.is_empty() { + // Create a data source for the ZIP file + let data_source = LazyDataSource::from_zip(path)?; + let data_source_arc = Arc::new(data_source); + + let mut reader = BufReader::new(pickle_data.as_slice()); + let obj = read_pickle_with_data(&mut reader, data_source_arc)?; + return convert_object_to_value(obj, top_level_key); + } + } + + // Try as plain pickle file + // First attempt without data source (for pure metadata files) + let file = File::open(path)?; + let mut reader = BufReader::new(file); + + match read_pickle(&mut reader) { + Ok(obj) => convert_object_to_value(obj, top_level_key), + Err(e) + if e.to_string() + .contains("Cannot load tensor data without a data source") => + { + // File contains tensors, need to use full PytorchReader + // Use the regular reader to get proper tensor handling + let reader = PytorchReader::new(path)?; + + // Convert tensors to PickleValue structure + let mut result = std::collections::HashMap::new(); + for key in reader.keys() { + // For pickle value extraction, we just need the structure, not the actual data + result.insert( + key.clone(), + PickleValue::String(format!("", key)), + ); + } + + if let Some(key) = top_level_key { + Ok(PickleValue::Dict( + [(key.to_string(), PickleValue::Dict(result))] + .into_iter() + .collect(), + )) + } else { + Ok(PickleValue::Dict(result)) + } + } + Err(e) => Err(PytorchError::Pickle(e)), + } +} + +/// Convert internal Object to public PickleValue +fn convert_object_to_value(obj: Object, top_level_key: Option<&str>) -> Result { + use crate::pytorch::pickle_reader::Object; + + // If a top-level key is specified, extract it first + if let Some(key) = top_level_key + && let Object::Dict(dict) = obj + { + if let Some(value) = dict.get(key) { + return object_to_pickle_value(value.clone()); + } else { + return Err(PytorchError::KeyNotFound(format!( + "Key '{}' not found in pickle data", + key + ))); + } + } + + object_to_pickle_value(obj) +} + +/// Convert Object to PickleValue +fn object_to_pickle_value(obj: Object) -> Result { + use crate::pytorch::pickle_reader::Object; + + Ok(match obj { + Object::None => PickleValue::None, + Object::Bool(b) => PickleValue::Bool(b), + Object::Int(i) => PickleValue::Int(i), + Object::Float(f) => PickleValue::Float(f), + Object::String(s) => PickleValue::String(s), + Object::Persistent(data) => { + // Persistent data is raw bytes + PickleValue::Bytes(data) + } + Object::PersistentTuple(tuple) => { + // Convert persistent tuples to lists + let mut values = Vec::new(); + for item in tuple { + values.push(object_to_pickle_value(item)?); + } + PickleValue::List(values) + } + Object::List(list) => { + let mut values = Vec::new(); + for item in list { + values.push(object_to_pickle_value(item)?); + } + PickleValue::List(values) + } + Object::Dict(dict) => { + let mut map = HashMap::new(); + for (k, v) in dict { + map.insert(k, object_to_pickle_value(v)?); + } + PickleValue::Dict(map) + } + Object::Tuple(tuple) => { + // Convert tuples to lists in the public API + let mut values = Vec::new(); + for item in tuple { + values.push(object_to_pickle_value(item)?); + } + PickleValue::List(values) + } + Object::TorchParam(_) => { + // Skip tensor parameters in config reading + PickleValue::None + } + Object::Class { .. } | Object::Build { .. } | Object::Reduce { .. } => { + // Complex objects are represented as None for simplicity + PickleValue::None + } + }) +} + +/// Convert PickleValue to NestedValue for deserialization +fn convert_pickle_to_nested_value(value: PickleValue) -> Result { + Ok(match value { + PickleValue::None => NestedValue::Default(None), + PickleValue::Bool(b) => NestedValue::Bool(b), + PickleValue::Int(i) => NestedValue::I64(i), + PickleValue::Float(f) => NestedValue::F64(f), + PickleValue::String(s) => NestedValue::String(s), + PickleValue::List(list) => { + let mut vec = Vec::new(); + for item in list { + vec.push(convert_pickle_to_nested_value(item)?); + } + NestedValue::Vec(vec) + } + PickleValue::Dict(dict) => { + let mut map = HashMap::new(); + for (k, v) in dict { + map.insert(k, convert_pickle_to_nested_value(v)?); + } + NestedValue::Map(map) + } + PickleValue::Bytes(data) => { + // Convert bytes to a list of u8 values + let vec: Vec = data.into_iter().map(NestedValue::U8).collect(); + NestedValue::Vec(vec) + } + }) +} diff --git a/crates/burn-store/src/pytorch/store.rs b/crates/burn-store/src/pytorch/store.rs new file mode 100644 index 0000000..4cb2068 --- /dev/null +++ b/crates/burn-store/src/pytorch/store.rs @@ -0,0 +1,439 @@ +//! PyTorch store implementation for saving and loading models in PyTorch format. + +use crate::{ + ApplyResult, KeyRemapper, ModuleSnapshot, ModuleStore, PathFilter, PyTorchToBurnAdapter, + TensorSnapshot, map_indices_contiguous, +}; + +use alloc::collections::BTreeMap; + +use alloc::format; +use alloc::string::{String, ToString}; +use alloc::vec::Vec; +use core::fmt; +use std::path::PathBuf; + +use super::reader::{PytorchError as ReaderError, PytorchReader}; + +/// Errors that can occur during PyTorch operations. +#[derive(Debug)] +pub enum PytorchStoreError { + /// Reader error. + Reader(ReaderError), + + /// I/O error. + Io(std::io::Error), + + /// Tensor not found. + TensorNotFound(String), + + /// Validation failed. + ValidationFailed(String), + + /// Other error. + Other(String), +} + +impl fmt::Display for PytorchStoreError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Reader(e) => write!(f, "PyTorch reader error: {}", e), + Self::Io(e) => write!(f, "I/O error: {}", e), + Self::TensorNotFound(name) => write!(f, "Tensor not found: {}", name), + Self::ValidationFailed(msg) => write!(f, "Validation failed: {}", msg), + Self::Other(msg) => write!(f, "{}", msg), + } + } +} + +impl std::error::Error for PytorchStoreError {} + +impl From for PytorchStoreError { + fn from(e: ReaderError) -> Self { + PytorchStoreError::Reader(e) + } +} + +impl From for PytorchStoreError { + fn from(e: std::io::Error) -> Self { + PytorchStoreError::Io(e) + } +} + +/// PyTorch store for file-based storage only. +/// +/// This store allows loading models from PyTorch checkpoint files (.pt/.pth) +/// with automatic weight transformation using `PyTorchToBurnAdapter`. +/// Linear weights are automatically transposed and normalization parameters +/// are renamed (gamma -> weight, beta -> bias). +/// +/// Note that saving to PyTorch format is not yet supported. +pub struct PytorchStore { + pub(crate) path: PathBuf, + pub(crate) filter: PathFilter, + pub(crate) remapper: KeyRemapper, + pub(crate) validate: bool, + pub(crate) allow_partial: bool, + pub(crate) top_level_key: Option, + pub(crate) skip_enum_variants: bool, + /// Enable contiguous mapping of layer indices (default: true) + pub(crate) map_indices_contiguous: bool, + /// Cached tensor snapshots (parsed once, reused) + snapshots_cache: Option>, +} + +impl PytorchStore { + /// Create a store for loading from a PyTorch file. + /// + /// # Arguments + /// * `path` - Path to the PyTorch checkpoint file (.pt or .pth) + /// + /// # Example + /// ```rust,no_run + /// use burn_store::PytorchStore; + /// + /// let store = PytorchStore::from_file("model.pth"); + /// ``` + pub fn from_file(path: impl Into) -> Self { + Self { + path: path.into(), + filter: PathFilter::new(), + remapper: KeyRemapper::new(), + validate: true, + allow_partial: false, + top_level_key: None, + // PyTorch models never include enum variant names in paths + skip_enum_variants: true, + // Enable contiguous index mapping by default for PyTorch files + // This handles nn.Sequential models with gaps in layer indices + map_indices_contiguous: true, + snapshots_cache: None, + } + } + + /// Set a top-level key to extract tensors from. + /// + /// PyTorch files often contain nested dictionaries. Use this to extract + /// tensors from a specific top-level key like "state_dict" or "model_state_dict". + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::PytorchStore; + /// let store = PytorchStore::from_file("checkpoint.pth") + /// .with_top_level_key("model_state_dict"); + /// ``` + pub fn with_top_level_key(mut self, key: impl Into) -> Self { + self.top_level_key = Some(key.into()); + self + } + + /// Filter which tensors to load. + pub fn filter(mut self, filter: PathFilter) -> Self { + self.filter = filter; + self + } + + /// Add a regex pattern to filter tensors. + /// + /// Multiple patterns can be added and they work with OR logic. + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::PytorchStore; + /// let store = PytorchStore::from_file("model.pth") + /// .with_regex(r"^encoder\..*") // Match all encoder tensors + /// .with_regex(r".*\.weight$"); // OR match any weight tensors + /// ``` + pub fn with_regex>(mut self, pattern: S) -> Self { + self.filter = self.filter.with_regex(pattern); + self + } + + /// Add multiple regex patterns to filter tensors. + pub fn with_regexes(mut self, patterns: I) -> Self + where + I: IntoIterator, + S: AsRef, + { + self.filter = self.filter.with_regexes(patterns); + self + } + + /// Add an exact full path to match. + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::PytorchStore; + /// let store = PytorchStore::from_file("model.pth") + /// .with_full_path("encoder.layer1.weight") + /// .with_full_path("decoder.output.bias"); + /// ``` + pub fn with_full_path>(mut self, path: S) -> Self { + self.filter = self.filter.with_full_path(path); + self + } + + /// Add multiple exact full paths to match. + pub fn with_full_paths(mut self, paths: I) -> Self + where + I: IntoIterator, + S: Into, + { + self.filter = self.filter.with_full_paths(paths); + self + } + + /// Add a predicate function for custom filtering logic. + /// + /// The predicate receives the tensor path and container path. + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::PytorchStore; + /// let store = PytorchStore::from_file("model.pth") + /// .with_predicate(|path, _| path.starts_with("encoder.") || path.ends_with(".bias")); + /// ``` + pub fn with_predicate(mut self, predicate: fn(&str, &str) -> bool) -> Self { + self.filter = self.filter.with_predicate(predicate); + self + } + + /// Add multiple predicate functions. + pub fn with_predicates(mut self, predicates: I) -> Self + where + I: IntoIterator bool>, + { + self.filter = self.filter.with_predicates(predicates); + self + } + + /// Set the filter to match all paths (disables filtering). + pub fn match_all(mut self) -> Self { + self.filter = self.filter.match_all(); + self + } + + /// Remap tensor names during load. + pub fn remap(mut self, remapper: KeyRemapper) -> Self { + self.remapper = remapper; + self + } + + /// Add a regex pattern to remap tensor names during load. + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::PytorchStore; + /// let store = PytorchStore::from_file("model.pth") + /// .with_key_remapping(r"^encoder\.", "transformer.encoder.") // encoder.X -> transformer.encoder.X + /// .with_key_remapping(r"\.gamma$", ".weight"); // X.gamma -> X.weight + /// ``` + pub fn with_key_remapping( + mut self, + from_pattern: impl AsRef, + to_pattern: impl Into, + ) -> Self { + self.remapper = self + .remapper + .add_pattern(from_pattern, to_pattern) + .expect("Invalid regex pattern"); + self + } + + /// Set whether to validate tensors during loading (default: true). + pub fn validate(mut self, validate: bool) -> Self { + self.validate = validate; + self + } + + /// Allow partial loading of tensors (continue even if some tensors are missing). + pub fn allow_partial(mut self, allow: bool) -> Self { + self.allow_partial = allow; + self + } + + /// Skip enum variant names when matching tensor paths (default: true). + /// + /// When enabled, tensor paths from PyTorch that don't include enum variants + /// can be matched against Burn module paths that do include them. + /// For example, PyTorch path "feature.weight" can match Burn path "feature.BaseConv.weight". + /// + /// This defaults to `true` for PytorchStore since PyTorch models never include + /// enum variant names in their parameter paths. + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::PytorchStore; + /// // Disable enum variant skipping (not typical) + /// let store = PytorchStore::from_file("model.pth") + /// .skip_enum_variants(false); + /// ``` + pub fn skip_enum_variants(mut self, skip: bool) -> Self { + self.skip_enum_variants = skip; + self + } + + /// Enable or disable automatic contiguous mapping of layer indices (default: true). + /// + /// When enabled, non-contiguous numeric indices in tensor paths are renumbered + /// to be contiguous. This is useful when loading PyTorch models that have gaps + /// in layer numbering, such as when using `nn.Sequential` with mixed layer types + /// (e.g., Conv2d layers at indices 0, 2, 4 with ReLU layers at 1, 3, 5). + /// + /// # Example + /// + /// With index mapping enabled (default): + /// - `fc.0.weight` → `fc.0.weight` + /// - `fc.2.weight` → `fc.1.weight` (gap filled) + /// - `fc.4.weight` → `fc.2.weight` (gap filled) + /// + /// # Arguments + /// + /// * `map` - `true` to enable contiguous index mapping, `false` to disable + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::PytorchStore; + /// // Disable contiguous index mapping if your model already has contiguous indices + /// let store = PytorchStore::from_file("model.pth") + /// .map_indices_contiguous(false); + /// ``` + pub fn map_indices_contiguous(mut self, map: bool) -> Self { + self.map_indices_contiguous = map; + self + } + + /// Apply remapping to tensor snapshots. + fn apply_remapping(&self, snapshots: Vec) -> Vec { + if self.remapper.is_empty() { + return snapshots; + } + + let (remapped, _) = self.remapper.remap(snapshots); + remapped + } + + /// Create a PytorchReader for the configured path and options. + fn create_reader(&self) -> Result { + let reader = if let Some(ref key) = self.top_level_key { + PytorchReader::with_top_level_key(&self.path, key)? + } else { + PytorchReader::new(&self.path)? + }; + Ok(reader) + } +} + +impl ModuleStore for PytorchStore { + type Error = PytorchStoreError; + + fn collect_from(&mut self, _module: &M) -> Result<(), Self::Error> { + // Saving to PyTorch format is not yet supported + Err(PytorchStoreError::Other( + "Saving to PyTorch format is not yet supported. Use other formats for saving." + .to_string(), + )) + } + + fn apply_to(&mut self, module: &mut M) -> Result { + // Get snapshots from cache + let snapshots: Vec = self.get_all_snapshots()?.values().cloned().collect(); + + // Get filter (convert to Option for apply) + let filter_opt = if self.filter.is_empty() { + None + } else { + Some(self.filter.clone()) + }; + + // Apply to module with PyTorchToBurnAdapter (always used for PyTorch files) + // This adapter handles: + // - Transposing linear weights from PyTorch format to Burn format + // - Renaming normalization parameters (gamma -> weight, beta -> bias) + // Filter is applied here during apply, not during cache population + let result = module.apply( + snapshots, + filter_opt, + Some(Box::new(PyTorchToBurnAdapter)), + self.skip_enum_variants, + ); + + // Validate if needed + if self.validate && !result.errors.is_empty() { + return Err(PytorchStoreError::ValidationFailed(format!( + "Import errors:\n{}", + result + ))); + } + + if !self.allow_partial && !result.missing.is_empty() { + return Err(PytorchStoreError::TensorNotFound(format!( + "\n{}\n\nHint: Use `.allow_partial(true)` on the store to get an `ApplyResult` \ + with structured missing/error info instead of a hard failure.", + result + ))); + } + + Ok(result) + } + + fn get_snapshot(&mut self, name: &str) -> Result, Self::Error> { + self.ensure_snapshots_cache()?; + Ok(self.snapshots_cache.as_ref().unwrap().get(name)) + } + + fn get_all_snapshots(&mut self) -> Result<&BTreeMap, Self::Error> { + self.ensure_snapshots_cache()?; + Ok(self.snapshots_cache.as_ref().unwrap()) + } + + fn keys(&mut self) -> Result, Self::Error> { + // Always use the cache to ensure remapping is applied consistently + Ok(self.get_all_snapshots()?.keys().cloned().collect()) + } +} + +impl PytorchStore { + /// Ensure the snapshots cache is populated + fn ensure_snapshots_cache(&mut self) -> Result<(), PytorchStoreError> { + if self.snapshots_cache.is_some() { + return Ok(()); + } + + let reader = self.create_reader()?; + + // Convert to tensor snapshots + let mut snapshots: Vec = reader + .into_tensors() + .into_iter() + .map(|(key, mut snapshot)| { + // Parse the key into path parts (split by '.') + let path_parts: Vec = key.split('.').map(|s| s.to_string()).collect(); + + // Set the path stack from the key + snapshot.path_stack = Some(path_parts); + snapshot.container_stack = None; + snapshot.tensor_id = None; + + snapshot + }) + .collect(); + + // Apply remapping (but NOT filtering - that's done at apply time) + snapshots = self.apply_remapping(snapshots); + + // Apply contiguous index mapping if enabled + // This must be done after remapping so that remapped paths are mapped + if self.map_indices_contiguous { + let (mapped, _) = map_indices_contiguous(snapshots); + snapshots = mapped; + } + + // Build cache as BTreeMap + let cache: BTreeMap = + snapshots.into_iter().map(|s| (s.full_path(), s)).collect(); + + self.snapshots_cache = Some(cache); + Ok(()) + } +} diff --git a/crates/burn-store/src/pytorch/tests/mod.rs b/crates/burn-store/src/pytorch/tests/mod.rs new file mode 100644 index 0000000..b9d78f0 --- /dev/null +++ b/crates/burn-store/src/pytorch/tests/mod.rs @@ -0,0 +1,2 @@ +pub mod reader; +pub mod store; diff --git a/crates/burn-store/src/pytorch/tests/reader/create_legacy_with_offsets.py b/crates/burn-store/src/pytorch/tests/reader/create_legacy_with_offsets.py new file mode 100644 index 0000000..4e690b0 --- /dev/null +++ b/crates/burn-store/src/pytorch/tests/reader/create_legacy_with_offsets.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 +# /// script +# dependencies = ["torch"] +# /// +"""Create a legacy format PyTorch file with specific storage offsets to test offset handling.""" + +import torch + +# Create tensors with known values at specific storage offsets +# This will help us verify we're reading from the correct location + +# Create a state dict with tensors that share storage +# This is common in PyTorch models (e.g., weight and transposed weight views) +state_dict = {} + +# Create a base tensor with known pattern +base_data = torch.arange(100, dtype=torch.float32) + +# tensor1: uses elements 10-19 (offset 10*4 = 40 bytes) +tensor1 = base_data[10:20].clone() +tensor1[:] = torch.arange(1.0, 1.1, 0.01)[:10] # 1.00, 1.01, 1.02, ... + +# tensor2: uses elements 30-35 (offset 30*4 = 120 bytes) +tensor2 = base_data[30:35].clone() +tensor2[:] = torch.arange(2.0, 2.5, 0.1)[:5] # 2.0, 2.1, 2.2, 2.3, 2.4 + +# tensor3: starts at beginning (offset 0) +tensor3 = base_data[:5].clone() +tensor3[:] = torch.arange(3.0, 3.5, 0.1)[:5] # 3.0, 3.1, 3.2, 3.3, 3.4 + +state_dict['tensor1'] = tensor1 +state_dict['tensor2'] = tensor2 +state_dict['tensor3'] = tensor3 + +# Save in legacy format +output_file = 'test_data/legacy_with_offsets.pt' +torch.save(state_dict, output_file, _use_new_zipfile_serialization=False) + +print(f"Created {output_file}") + +# Verify by loading +loaded = torch.load(output_file, weights_only=False) +print("\nVerification - expected values:") +for key, tensor in loaded.items(): + print(f" {key}: {tensor.tolist()}") + print(f" Storage offset: {tensor.storage_offset()}") + print(f" Storage size: {len(tensor.storage())}") + +# Also create a test with multiple tensors sharing the same storage +# This is important for proper offset handling +shared_storage = torch.randn(1000) + +# Create views into the same storage at different offsets +view1 = shared_storage[100:110] # offset 100 +view2 = shared_storage[500:520] # offset 500 +view3 = shared_storage[0:10] # offset 0 + +# Need to save these properly - PyTorch will handle the storage sharing +shared_dict = { + 'view1': view1.clone(), # Clone to avoid view issues + 'view2': view2.clone(), + 'view3': view3.clone(), +} + +output_file2 = 'test_data/legacy_shared_storage.pt' +torch.save(shared_dict, output_file2, _use_new_zipfile_serialization=False) +print(f"\nCreated {output_file2}") + +# Print exact values for test verification +print("\nExact test values for legacy_with_offsets.pt:") +print("tensor1 (10 elements starting at 1.0):") +print(" First 3 values: [1.00, 1.01, 1.02]") +print("tensor2 (5 elements starting at 2.0):") +print(" All values: [2.0, 2.1, 2.2, 2.3, 2.4]") +print("tensor3 (5 elements starting at 3.0):") +print(" All values: [3.0, 3.1, 3.2, 3.3, 3.4]") \ No newline at end of file diff --git a/crates/burn-store/src/pytorch/tests/reader/create_tar_format.py b/crates/burn-store/src/pytorch/tests/reader/create_tar_format.py new file mode 100644 index 0000000..28549f7 --- /dev/null +++ b/crates/burn-store/src/pytorch/tests/reader/create_tar_format.py @@ -0,0 +1,361 @@ +#!/usr/bin/env python3 +""" +Create TAR format test fixtures for burn-store integration tests. + +The TAR format was used by very early versions of PyTorch (pre 0.1.10). +Modern torch.save cannot create this format, so we construct it manually. + +TAR format structure: + - sys_info: pickle with {protocol_version, little_endian, type_sizes} + - pickle: pickle with OrderedDict containing _rebuild_tensor_v2 REDUCE calls + - storages: count_pickle + for each storage: (key, device, class) pickle + u64 num_elements + raw data +""" + +import io +import pickle +import struct +import tarfile +import os +from collections import OrderedDict + + +def create_sys_info(): + """Create sys_info pickle data.""" + sys_info = { + "protocol_version": 1000, + "little_endian": True, + "type_sizes": { + "short": 2, + "int": 4, + "long": 8, + }, + } + return pickle.dumps(sys_info, protocol=2) + + +def encode_tensor_data(values: list, storage_type: str) -> tuple: + """Encode tensor values to bytes and return (bytes, element_size).""" + fmt_map = { + "FloatStorage": (" bytes: + """ + Create the storages binary blob manually. + + Args: + tensors: List of (key, storage_type, element_size, data_bytes) tuples + """ + buffer = io.BytesIO() + + # Write storage count as pickle (simple integer) + pickle.dump(len(tensors), buffer, protocol=2) + + for key, storage_type, element_size, data_bytes in tensors: + # Manually construct the tuple pickle with GLOBAL class reference + # Format: (key, "cpu", ) + + tuple_buffer = io.BytesIO() + # Protocol 2 header + tuple_buffer.write(b'\x80\x02') + + # Build tuple with MARK + items + TUPLE + tuple_buffer.write(b'(') # MARK + + # First item: storage key (string) + write_string(tuple_buffer, key) + + # Second item: device "cpu" + tuple_buffer.write(b'U\x03cpu') + + # Third item: class reference using GLOBAL + tuple_buffer.write(b'c') # GLOBAL opcode + tuple_buffer.write(b'torch\n') # module + tuple_buffer.write(storage_type.encode('ascii') + b'\n') # name + + # End tuple + tuple_buffer.write(b't') # TUPLE + tuple_buffer.write(b'.') # STOP + + buffer.write(tuple_buffer.getvalue()) + + # Write num_elements as u64 little-endian + num_elements = len(data_bytes) // element_size + buffer.write(struct.pack(" bytes: + """ + Create the main pickle containing _rebuild_tensor_v2 REDUCE calls. + + For each tensor, we need: + - GLOBAL torch._utils _rebuild_tensor_v2 + - MARK + - args tuple: (persistent_id, offset, shape, stride, requires_grad, hooks) + - TUPLE + - REDUCE + + The persistent_id is a PersistentTuple: ('storage', , key, device, num_elements) + """ + buffer = io.BytesIO() + + # Protocol 2 header + buffer.write(b'\x80\x02') + + # Build OrderedDict: GLOBAL + EMPTY_LIST + items + TUPLE + REDUCE + # OrderedDict([('name1', tensor1), ('name2', tensor2)]) + + # GLOBAL collections OrderedDict + buffer.write(b'ccollections\nOrderedDict\n') + + # Start list for items + buffer.write(b'(') # MARK + buffer.write(b']') # EMPTY_LIST + + # For each tensor, add (name, rebuilt_tensor) to the list + for name, storage_key, storage_type, shape, num_elements in tensors_info: + # Calculate stride for row-major (C) order + stride = [] + s = 1 + for dim in reversed(shape): + stride.insert(0, s) + s *= dim + + # Build inner tuple: (name, tensor_value) + buffer.write(b'(') # MARK for (name, value) tuple + + # Write name + write_string(buffer, name) + + # Now build the tensor using _rebuild_tensor_v2 REDUCE + # GLOBAL torch._utils _rebuild_tensor_v2 + buffer.write(b'ctorch._utils\n_rebuild_tensor_v2\n') + + # Build args tuple for _rebuild_tensor_v2 + # (persistent_id, offset, shape, stride, requires_grad, backward_hooks) + buffer.write(b'(') # MARK for args tuple + + # arg 0: persistent_id tuple: ('storage', class, key, device, num_elements) + # This will be converted to PersistentTuple by the reader + buffer.write(b'(') # MARK for persistent_id + + write_string(buffer, 'storage') + + # Class reference - GLOBAL torch FloatStorage + buffer.write(b'c') + buffer.write(b'torch\n') + buffer.write(storage_type.encode('ascii') + b'\n') + + # Storage key + write_string(buffer, storage_key) + + # Device + buffer.write(b'U\x03cpu') + + # num_elements + write_int(buffer, num_elements) + + buffer.write(b't') # TUPLE - end persistent_id + + # arg 1: storage offset (0) + buffer.write(b'K\x00') + + # arg 2: shape tuple + buffer.write(b'(') + for dim in shape: + write_int(buffer, dim) + buffer.write(b't') + + # arg 3: stride tuple + buffer.write(b'(') + for s_val in stride: + write_int(buffer, s_val) + buffer.write(b't') + + # arg 4: requires_grad (False) + buffer.write(b'\x89') # NEWFALSE + + # arg 5: backward_hooks (empty OrderedDict) + buffer.write(b'ccollections\nOrderedDict\n') + buffer.write(b'(') + buffer.write(b']') + buffer.write(b't') + buffer.write(b'R') # REDUCE to create empty OrderedDict + + buffer.write(b't') # TUPLE - end args tuple + + buffer.write(b'R') # REDUCE - call _rebuild_tensor_v2 with args + + buffer.write(b't') # TUPLE - end (name, tensor) tuple + + buffer.write(b'a') # APPEND to list + + buffer.write(b't') # TUPLE - wrap list in tuple for REDUCE + buffer.write(b'R') # REDUCE - call OrderedDict with the list + buffer.write(b'.') # STOP + + return buffer.getvalue() + + +def create_tar_pytorch_file(filename: str, tensors: dict, dtypes: dict): + """ + Create a TAR format PyTorch file. + + Args: + filename: Output file path + tensors: Dict of tensor_name -> (values_list, shape) + dtypes: Dict of tensor_name -> storage_type + """ + # Prepare storage data + storage_list = [] # (key, storage_type, element_size, data_bytes) + tensors_info = [] # (name, storage_key, storage_type, shape, num_elements) + + for idx, (name, (values, shape)) in enumerate(tensors.items()): + storage_key = str(idx) + storage_type = dtypes[name] + data_bytes, element_size = encode_tensor_data(values, storage_type) + num_elements = len(values) + + storage_list.append((storage_key, storage_type, element_size, data_bytes)) + tensors_info.append((name, storage_key, storage_type, shape, num_elements)) + + # Create the three main entries + sys_info_data = create_sys_info() + pickle_data = create_main_pickle_manual(tensors_info) + storages_data = create_storages_blob_manual(storage_list) + + # Write TAR archive + os.makedirs(os.path.dirname(filename) or ".", exist_ok=True) + + with tarfile.open(filename, "w") as tar: + # Add sys_info + tarinfo = tarfile.TarInfo(name="sys_info") + tarinfo.size = len(sys_info_data) + tar.addfile(tarinfo, io.BytesIO(sys_info_data)) + + # Add pickle + tarinfo = tarfile.TarInfo(name="pickle") + tarinfo.size = len(pickle_data) + tar.addfile(tarinfo, io.BytesIO(pickle_data)) + + # Add storages + tarinfo = tarfile.TarInfo(name="storages") + tarinfo.size = len(storages_data) + tar.addfile(tarinfo, io.BytesIO(storages_data)) + + size = os.path.getsize(filename) + print(f"Created {filename} ({size} bytes)") + print(f" Tensors: {list(tensors.keys())}") + + +def main(): + # Create test_data directory + os.makedirs("test_data", exist_ok=True) + + # Test 1: Single float32 tensor + create_tar_pytorch_file( + "test_data/tar_float32.tar", + {"tensor": ([1.0, 2.5, -3.7, 0.0], [4])}, + {"tensor": "FloatStorage"}, + ) + + # Test 2: Single float64 tensor + create_tar_pytorch_file( + "test_data/tar_float64.tar", + {"tensor": ([1.1, 2.2, 3.3], [3])}, + {"tensor": "DoubleStorage"}, + ) + + # Test 3: Single int64 tensor + create_tar_pytorch_file( + "test_data/tar_int64.tar", + {"tensor": ([100, -200, 300, 0], [4])}, + {"tensor": "LongStorage"}, + ) + + # Test 4: Multiple tensors (weight + bias) + create_tar_pytorch_file( + "test_data/tar_weight_bias.tar", + { + "weight": ([0.1, 0.2, 0.3, 0.4, 0.5, 0.6], [2, 3]), + "bias": ([0.01, 0.02], [2]), + }, + { + "weight": "FloatStorage", + "bias": "FloatStorage", + }, + ) + + # Test 5: Different dtypes in one file + create_tar_pytorch_file( + "test_data/tar_multi_dtype.tar", + { + "float_tensor": ([1.5, 2.5, 3.5], [3]), + "double_tensor": ([1.111, 2.222], [2]), + "int_tensor": ([10, 20, 30, 40], [4]), + }, + { + "float_tensor": "FloatStorage", + "double_tensor": "DoubleStorage", + "int_tensor": "LongStorage", + }, + ) + + # Test 6: 2D tensor for shape verification + create_tar_pytorch_file( + "test_data/tar_2d_tensor.tar", + { + "matrix": ([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0], [3, 4]), + }, + {"matrix": "FloatStorage"}, + ) + + print("\nAll TAR format test files created!") + print("\nTo run tests: cargo test -p burn-store --features pytorch test_tar") + + +if __name__ == "__main__": + main() diff --git a/crates/burn-store/src/pytorch/tests/reader/mod.rs b/crates/burn-store/src/pytorch/tests/reader/mod.rs new file mode 100644 index 0000000..0786e35 --- /dev/null +++ b/crates/burn-store/src/pytorch/tests/reader/mod.rs @@ -0,0 +1,1278 @@ +//! Tests for PyTorch file reader functionality +//! +//! Floating-point comparison tolerances: +//! - F16/BF16: 1e-2 (~3 decimal digits precision) +//! - F32: 1e-6 (~7 decimal digits precision) +//! - F64: 1e-10 (~16 decimal digits precision) + +#![allow(clippy::needless_range_loop)] + +use crate::pytorch::PytorchReader; +// Import internal types for testing only +use crate::pytorch::reader::{ByteOrder, FileFormat}; +use burn_core::tensor::{BoolStore, DType, TensorData, Tolerance, shape}; +use std::path::PathBuf; + +fn test_data_path(filename: &str) -> PathBuf { + // Get the path relative to the crate root + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("src") + .join("pytorch") + .join("tests") + .join("reader") + .join("test_data") + .join(filename) +} + +#[test] +fn test_float32_tensor() { + let path = test_data_path("float32.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load float32.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::F32); + assert_eq!(tensor.shape, shape![4]); + + let data = tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 4); + assert!((values[0] - 1.0).abs() < 1e-6); + assert!((values[1] - 2.5).abs() < 1e-6); + assert!((values[2] - (-3.7)).abs() < 1e-6); + assert!((values[3] - 0.0).abs() < 1e-6); +} + +#[test] +fn test_float64_tensor() { + let path = test_data_path("float64.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load float64.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::F64); + assert_eq!(tensor.shape, shape![3]); + + let data = tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 3); + assert!((values[0] - 1.1).abs() < 1e-10); + assert!((values[1] - 2.2).abs() < 1e-10); + assert!((values[2] - 3.3).abs() < 1e-10); +} + +#[test] +fn test_int64_tensor() { + let path = test_data_path("int64.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load int64.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::I64); + assert_eq!(tensor.shape, shape![4]); + + let data = tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values, &[100, -200, 300, 0]); +} + +#[test] +fn test_int32_tensor() { + let path = test_data_path("int32.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load int32.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::I32); + assert_eq!(tensor.shape, shape![3]); + + let data = tensor.to_data().unwrap(); + // Convert to the appropriate element type + let data_converted = data.convert::(); + let values = data_converted.as_slice::().unwrap(); + assert_eq!(values, &[10, 20, -30]); +} + +#[test] +fn test_int16_tensor() { + let path = test_data_path("int16.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load int16.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::I16); + assert_eq!(tensor.shape, shape![3]); + + let data = tensor.to_data().unwrap(); + let data_converted = data.convert::(); + let values = data_converted.as_slice::().unwrap(); + assert_eq!(values, &[1000, -2000, 3000]); +} + +#[test] +fn test_int8_tensor() { + let path = test_data_path("int8.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load int8.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::I8); + assert_eq!(tensor.shape, shape![4]); + + let data = tensor.to_data().unwrap(); + let data_converted = data.convert::(); + let values = data_converted.as_slice::().unwrap(); + assert_eq!(values, &[127, -128, 0, 50]); +} + +#[test] +fn test_bool_tensor() { + let path = test_data_path("bool.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load bool.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::Bool(BoolStore::Native)); + assert_eq!(tensor.shape, shape![5]); + + let data = tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values, &[true, false, true, true, false]); +} + +#[test] +fn test_uint8_tensor() { + let path = test_data_path("uint8.pt"); + + let reader = PytorchReader::new(&path).expect("Failed to load uint8.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::U8); + assert_eq!(tensor.shape, shape![4]); + + // Verify actual U8 values [0, 128, 255, 42] from test_data.py + let data = tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values, &[0, 128, 255, 42]); +} + +#[test] +fn test_float16_tensor() { + use half::f16; + + let path = test_data_path("float16.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load float16.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::F16); + assert_eq!(tensor.shape, shape![3]); + + // Verify actual F16 values [1.5, -2.25, 3.125] from test_data.py + let data = tensor.to_data().unwrap(); + assert_eq!(data.shape, shape![3]); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 3); + assert!((values[0].to_f32() - 1.5).abs() < 1e-2); + assert!((values[1].to_f32() - (-2.25)).abs() < 1e-2); + assert!((values[2].to_f32() - 3.125).abs() < 1e-2); +} + +#[test] +fn test_bfloat16_tensor() { + use half::bf16; + + let path = test_data_path("bfloat16.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load bfloat16.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::BF16); + assert_eq!(tensor.shape, shape![3]); + + // Verify actual BF16 values [1.5, -2.5, 3.5] from test_data.py + let data = tensor.to_data().unwrap(); + assert_eq!(data.shape, shape![3]); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 3); + assert!((values[0].to_f32() - 1.5).abs() < 1e-2); + assert!((values[1].to_f32() - (-2.5)).abs() < 1e-2); + assert!((values[2].to_f32() - 3.5).abs() < 1e-2); +} + +#[test] +fn test_2d_tensor() { + let path = test_data_path("tensor_2d.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load tensor_2d.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::F32); + assert_eq!(tensor.shape, shape![3, 2]); + + let data = tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 6); + // Check flattened values [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + for (i, expected) in [1.0, 2.0, 3.0, 4.0, 5.0, 6.0].iter().enumerate() { + assert!((values[i] - expected).abs() < 1e-6); + } +} + +#[test] +fn test_3d_tensor() { + let path = test_data_path("tensor_3d.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load tensor_3d.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::F32); + assert_eq!(tensor.shape, shape![2, 3, 4]); + + let data = tensor.to_data().unwrap(); + assert_eq!(data.shape, shape![2, 3, 4]); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 24); +} + +#[test] +fn test_4d_tensor() { + let path = test_data_path("tensor_4d.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load tensor_4d.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::F32); + assert_eq!(tensor.shape, shape![2, 3, 2, 2]); + + let data = tensor.to_data().unwrap(); + assert_eq!(data.shape, shape![2, 3, 2, 2]); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 24); +} + +#[test] +fn test_state_dict() { + let path = test_data_path("state_dict.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load state_dict.pt"); + let keys = reader.keys(); + + assert_eq!(keys.len(), 4); + assert!(keys.contains(&"weight".to_string())); + assert!(keys.contains(&"bias".to_string())); + assert!(keys.contains(&"running_mean".to_string())); + assert!(keys.contains(&"running_var".to_string())); + + // Check weight tensor + let weight = reader.get("weight").unwrap(); + assert_eq!(weight.shape, shape![3, 4]); + assert_eq!(weight.dtype, DType::F32); + + // Check bias tensor + let bias = reader.get("bias").unwrap(); + assert_eq!(bias.shape, shape![3]); + assert_eq!(bias.dtype, DType::F32); + + // Check running_mean (should be zeros) + let running_mean = reader.get("running_mean").unwrap(); + assert_eq!(running_mean.shape, shape![3]); + let mean_data = running_mean.to_data().unwrap(); + let mean_values = mean_data.as_slice::().unwrap(); + assert!(mean_values.iter().all(|&v| v.abs() < 1e-6)); + + // Check running_var (should be ones) + let running_var = reader.get("running_var").unwrap(); + assert_eq!(running_var.shape, shape![3]); + let var_data = running_var.to_data().unwrap(); + let var_values = var_data.as_slice::().unwrap(); + assert!(var_values.iter().all(|&v| (v - 1.0).abs() < 1e-6)); +} + +#[test] +fn test_nested_dict() { + let path = test_data_path("nested_dict.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load nested_dict.pt"); + let keys = reader.keys(); + + assert_eq!(keys.len(), 4); + assert!(keys.contains(&"layer1.weight".to_string())); + assert!(keys.contains(&"layer1.bias".to_string())); + assert!(keys.contains(&"layer2.weight".to_string())); + assert!(keys.contains(&"layer2.bias".to_string())); + + // Check layer1.weight and load data + let layer1_weight = reader.get("layer1.weight").unwrap(); + assert_eq!(layer1_weight.shape, shape![2, 3]); + assert_eq!(layer1_weight.dtype, DType::F32); + let data = layer1_weight.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 6); // 2x3 = 6 elements + + // Check layer2.weight and load data + let layer2_weight = reader.get("layer2.weight").unwrap(); + assert_eq!(layer2_weight.shape, shape![4, 2]); + assert_eq!(layer2_weight.dtype, DType::F32); + let data = layer2_weight.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 8); // 4x2 = 8 elements +} + +#[test] +fn test_checkpoint() { + let path = test_data_path("checkpoint.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load checkpoint.pt"); + let keys = reader.keys(); + + // Should have model_state_dict entries and optimizer entries + assert!(keys.contains(&"model_state_dict.fc1.weight".to_string())); + assert!(keys.contains(&"model_state_dict.fc1.bias".to_string())); + assert!(keys.contains(&"model_state_dict.fc2.weight".to_string())); + assert!(keys.contains(&"model_state_dict.fc2.bias".to_string())); + + // Check fc1.weight dimensions and load data + let fc1_weight = reader.get("model_state_dict.fc1.weight").unwrap(); + assert_eq!(fc1_weight.shape, shape![10, 5]); + let data = fc1_weight.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 50); // 10x5 = 50 elements + + // Check fc2.weight dimensions and load data + let fc2_weight = reader.get("model_state_dict.fc2.weight").unwrap(); + assert_eq!(fc2_weight.shape, shape![3, 10]); + let data = fc2_weight.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 30); // 3x10 = 30 elements +} + +#[test] +fn test_empty_tensor() { + let path = test_data_path("empty.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load empty.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.shape, shape![0]); // Empty tensor has shape [0] + assert_eq!(tensor.dtype, DType::F32); + + // Note: Empty tensors cannot be loaded with to_data() due to TensorData validation + // We verify the metadata is correct, which confirms the .pt file is being read +} + +#[test] +fn test_scalar_tensor() { + let path = test_data_path("scalar.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load scalar.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.shape, shape![]); // Scalar has empty shape + assert_eq!(tensor.dtype, DType::F32); + + let data = tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 1); + assert!((values[0] - 42.0).abs() < 1e-6); +} + +#[test] +fn test_large_shape() { + let path = test_data_path("large_shape.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load large_shape.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.shape, shape![100, 100]); + assert_eq!(tensor.dtype, DType::F32); + + let data = tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 10000); + + // Check specific non-zero values + assert!((values[0] - 1.0).abs() < 1e-6); // [0, 0] = 1.0 + assert!((values[5050] - 2.0).abs() < 1e-6); // [50, 50] = 2.0 + assert!((values[9999] - 3.0).abs() < 1e-6); // [99, 99] = 3.0 +} + +#[test] +fn test_mixed_types() { + let path = test_data_path("mixed_types.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load mixed_types.pt"); + let tensors = reader.tensors(); + + assert_eq!(tensors.len(), 4); + + // Check float32 tensor [1.0, 2.0] from test_data.py + let float32 = reader.get("float32").unwrap(); + assert_eq!(float32.dtype, DType::F32); + assert_eq!(float32.shape, shape![2]); + let data = float32.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert!((values[0] - 1.0).abs() < 1e-6); + assert!((values[1] - 2.0).abs() < 1e-6); + + // Check int64 tensor [100, 200] from test_data.py + let int64 = reader.get("int64").unwrap(); + assert_eq!(int64.dtype, DType::I64); + assert_eq!(int64.shape, shape![2]); + let data = int64.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values, &[100, 200]); + + // Check bool tensor [True, False] from test_data.py + let bool_tensor = reader.get("bool").unwrap(); + assert_eq!(bool_tensor.dtype, DType::Bool(BoolStore::Native)); + assert_eq!(bool_tensor.shape, shape![2]); + let data = bool_tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values, &[true, false]); + + // Check float64 tensor [1.1, 2.2] from test_data.py + let float64 = reader.get("float64").unwrap(); + assert_eq!(float64.dtype, DType::F64); + assert_eq!(float64.shape, shape![2]); + let data = float64.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert!((values[0] - 1.1).abs() < 1e-10); + assert!((values[1] - 2.2).abs() < 1e-10); +} + +#[test] +fn test_special_values() { + let path = test_data_path("special_values.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load special_values.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::F32); + assert_eq!(tensor.shape, shape![5]); + + let data = tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 5); + + // Check for special values + assert!(values[0].is_nan()); + assert!(values[1].is_infinite() && values[1] > 0.0); + assert!(values[2].is_infinite() && values[2] < 0.0); + assert!((values[3] - 0.0).abs() < 1e-6); + assert!((values[4] - 1.0).abs() < 1e-6); +} + +#[test] +fn test_extreme_values() { + let path = test_data_path("extreme_values.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load extreme_values.pt"); + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::F32); + assert_eq!(tensor.shape, shape![4]); + + let data = tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 4); + + // Very small positive + assert!(values[0] > 0.0 && values[0] < 1e-20); + // Very large positive + assert!(values[1] > 1e20); + // Very small negative + assert!(values[2] < 0.0 && values[2] > -1e-20); + // Very large negative + assert!(values[3] < -1e20); +} + +#[test] +fn test_parameter() { + let path = test_data_path("parameter.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load parameter.pt"); + let tensors = reader.tensors(); + + // nn.Parameter is typically saved as a regular tensor + assert_eq!(tensors.len(), 1); + let param = reader.get("param").unwrap(); + assert_eq!(param.shape, shape![3, 3]); + assert_eq!(param.dtype, DType::F32); + + let data = param.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 9); +} + +#[test] +fn test_buffers() { + let path = test_data_path("buffers.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load buffers.pt"); + let tensors = reader.tensors(); + + assert_eq!(tensors.len(), 2); + + // Check buffer1 (int32) + let buffer1 = reader.get("buffer1").unwrap(); + assert_eq!(buffer1.dtype, DType::I32); + assert_eq!(buffer1.shape, shape![3]); + let data1 = buffer1.to_data().unwrap(); + let data1_converted = data1.convert::(); + let values1 = data1_converted.as_slice::().unwrap(); + assert_eq!(values1, &[1, 2, 3]); + + // Check buffer2 (bool) + let buffer2 = reader.get("buffer2").unwrap(); + assert_eq!(buffer2.dtype, DType::Bool(BoolStore::Native)); + assert_eq!(buffer2.shape, shape![2]); + let data2 = buffer2.to_data().unwrap(); + let values2 = data2.as_slice::().unwrap(); + assert_eq!(values2, &[true, false]); +} + +#[test] +fn test_complex_structure() { + let path = test_data_path("complex_structure.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load complex_structure.pt"); + let keys = reader.keys(); + + // Should have nested structure tensors + assert!(keys.contains(&"state.encoder.layer_0.weight".to_string())); + assert!(keys.contains(&"state.encoder.layer_0.bias".to_string())); + assert!(keys.contains(&"state.encoder.layer_1.weight".to_string())); + assert!(keys.contains(&"state.encoder.layer_1.bias".to_string())); + assert!(keys.contains(&"state.decoder.weight".to_string())); + assert!(keys.contains(&"state.decoder.bias".to_string())); + + // Check encoder layer_0 weight and load data + let layer0_weight = reader.get("state.encoder.layer_0.weight").unwrap(); + assert_eq!(layer0_weight.shape, shape![4, 3]); + let data = layer0_weight.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 12); // 4x3 = 12 elements + + // Check decoder weight and load data + let decoder_weight = reader.get("state.decoder.weight").unwrap(); + assert_eq!(decoder_weight.shape, shape![3, 2]); + let data = decoder_weight.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 6); // 3x2 = 6 elements +} + +#[test] +fn test_read_pytorch_tensors_convenience() { + // Test reading and materializing tensors into memory + let path = test_data_path("state_dict.pt"); + let reader = PytorchReader::new(&path).expect("Failed to read file"); + + let keys = reader.keys(); + assert_eq!(keys.len(), 4); + assert!(keys.contains(&"weight".to_string())); + assert!(keys.contains(&"bias".to_string())); + + // Check that data can be materialized + let weight = reader.get("weight").unwrap(); + let weight_data = weight.to_data().unwrap(); + assert_eq!(weight_data.shape, shape![3, 4]); + assert_eq!(weight_data.dtype, DType::F32); +} + +#[test] +fn test_with_top_level_key() { + // Test loading with a specific top-level key + let path = test_data_path("checkpoint.pt"); + + // Load only model_state_dict + let reader = PytorchReader::with_top_level_key(&path, "model_state_dict") + .expect("Failed to load with top-level key"); + + let keys = reader.keys(); + // Should only have model weights, not optimizer state + assert!(keys.contains(&"fc1.weight".to_string())); + assert!(keys.contains(&"fc1.bias".to_string())); + assert!(keys.contains(&"fc2.weight".to_string())); + assert!(keys.contains(&"fc2.bias".to_string())); + + // Should NOT have nested paths with model_state_dict prefix + assert!(!keys.contains(&"model_state_dict.fc1.weight".to_string())); +} + +#[test] +fn test_legacy_format() { + // Test loading PyTorch legacy format (pre-1.6) + let path = test_data_path("simple_legacy.pt"); + + // This file has the sequential pickle structure of legacy PyTorch format + let reader = PytorchReader::new(&path).expect("Failed to load legacy format"); + let keys = reader.keys(); + + // Should have the tensors from the state dict + assert!(keys.contains(&"weight".to_string()), "Missing 'weight' key"); + assert!(keys.contains(&"bias".to_string()), "Missing 'bias' key"); + assert!( + keys.contains(&"running_mean".to_string()), + "Missing 'running_mean' key" + ); + + // Check weight tensor + let weight = reader.get("weight").expect("weight not found"); + // Note: values for `weight` in simple_legacy.pt are randomly generated + assert_eq!(weight.shape, shape![2, 3]); + assert_eq!(weight.dtype, DType::F32); + + // Check bias tensor + let bias = reader.get("bias").expect("bias not found"); + assert_eq!(bias.shape, shape![2]); + assert_eq!(bias.dtype, DType::F32); + + // Verify bias values are all ones + let bias_data = bias.to_data().unwrap(); + let expected_bias_data = TensorData::new(vec![1.0_f32, 1.0], vec![2]); + bias_data.assert_approx_eq::(&expected_bias_data, Tolerance::default()); + + // Check running_mean tensor + let running_mean = reader.get("running_mean").expect("running_mean not found"); + assert_eq!(running_mean.shape, shape![2]); + assert_eq!(running_mean.dtype, DType::F32); + + // Verify running_mean values are accessible + let mean_data = running_mean.to_data().unwrap(); + let expected_mean_data = TensorData::new(vec![0.0_f32, 0.0], vec![2]); + mean_data.assert_approx_eq::(&expected_mean_data, Tolerance::default()); +} + +#[test] +fn test_legacy_uncloned_views() { + let path = test_data_path("legacy_uncloned_views.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load legacy format"); + let keys = reader.keys(); + + // Should have the tensors from the state dict + assert!( + keys.contains(&"tensor1".to_string()), + "Missing 'tensor1' key" + ); + assert!( + keys.contains(&"tensor2".to_string()), + "Missing 'tensor2' key" + ); + + // Check tensor1 + let tensor1 = reader.get("tensor1").expect("tensor1 not found"); + assert_eq!(tensor1.shape, shape![10]); + assert_eq!(tensor1.dtype, DType::F32); + let tensor1_data = tensor1.to_data().unwrap(); + let expected_tensor1_data = + TensorData::new(vec![10, 11, 12, 13, 14, 15, 16, 17, 18, 19], vec![10]); + tensor1_data.assert_approx_eq::(&expected_tensor1_data, Tolerance::default()); + + // Check tensor2 + let tensor2 = reader.get("tensor2").expect("tensor2 not found"); + assert_eq!(tensor2.shape, shape![10]); + assert_eq!(tensor2.dtype, DType::F32); + let tensor2_data = tensor2.to_data().unwrap(); + let expected_tensor2_data = + TensorData::new(vec![50, 51, 52, 53, 54, 55, 56, 57, 58, 59], vec![10]); + tensor2_data.assert_approx_eq::(&expected_tensor2_data, Tolerance::default()); +} + +#[test] +fn test_legacy_with_offsets() { + // Test with legacy format file that has storage offsets + let path = test_data_path("legacy_with_offsets.pt"); + let reader = PytorchReader::new(&path).expect("Should read legacy file with offsets"); + assert_eq!(reader.keys().len(), 3, "Should have 3 tensors"); + + let tensor1 = reader + .get("tensor1") + .expect("Legacy file should contain tensor1"); + assert_eq!(tensor1.shape, shape![10]); + let data1 = tensor1.to_data().unwrap(); + let expected_data1 = TensorData::new( + vec![ + 1.00_f32, 1.01, 1.02, 1.03, 1.04, 1.05, 1.06, 1.07, 1.08, 1.09, + ], + vec![10], + ); + data1.assert_approx_eq::(&expected_data1, Tolerance::default()); + + let tensor2 = reader + .get("tensor2") + .expect("Legacy file should contain tensor2"); + assert_eq!(tensor2.shape, shape![5]); + let data2 = tensor2.to_data().unwrap(); + let expected_data2 = TensorData::new(vec![2.0_f32, 2.1, 2.2, 2.3, 2.4], vec![5]); + data2.assert_approx_eq::(&expected_data2, Tolerance::default()); + + let tensor3 = reader + .get("tensor3") + .expect("Legacy file should contain tensor3"); + assert_eq!(tensor3.shape, shape![5]); + let data3 = tensor3.to_data().unwrap(); + let expected_data3 = TensorData::new(vec![3.0_f32, 3.1, 3.2, 3.3, 3.4], vec![5]); + data3.assert_approx_eq::(&expected_data3, Tolerance::default()); +} + +#[test] +fn test_legacy_shared_storage() { + // Test with legacy format file that has shared storage + let path = test_data_path("legacy_shared_storage.pt"); + let reader = PytorchReader::new(&path).expect("Should read legacy file with shared storage"); + + let keys = reader.keys(); + assert!(keys.len() >= 2, "Should have at least 2 tensors"); + + // Check that tensors exist and can be loaded + for key in &keys { + assert!(reader.get(key).is_some(), "Should have tensor: {}", key); + let tensor = reader.get(key).unwrap(); + let data = tensor.to_data().unwrap(); + + // Verify tensor data can be accessed + match tensor.dtype { + DType::F32 => { + let values = data.as_slice::().unwrap(); + assert!(!values.is_empty(), "Tensor {} should have data", key); + } + DType::I64 => { + let values = data.as_slice::().unwrap(); + assert!(!values.is_empty(), "Tensor {} should have data", key); + } + _ => { + // For other types, just verify we can convert to data + assert!(!data.shape.is_empty(), "Tensor {} should have shape", key); + } + } + } +} + +#[test] +fn test_metadata_zip_format() { + // Test that metadata is properly populated for ZIP format files + let path = test_data_path("float32.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load float32.pt"); + + // Check metadata + let metadata = reader.metadata(); + assert_eq!(metadata.format_type, FileFormat::Zip); + assert_eq!(metadata.byte_order, ByteOrder::LittleEndian); + assert_eq!(metadata.tensor_count, 1); + assert!(metadata.total_data_size.is_some()); + + // Check that metadata is accessible + assert!(metadata.is_modern_format()); + assert!(!metadata.is_legacy_format()); +} + +#[test] +fn test_metadata_legacy_format() { + // Test that metadata is properly populated for legacy format files + let path = test_data_path("simple_legacy.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load legacy file"); + + // Check metadata + let metadata = reader.metadata(); + assert_eq!(metadata.format_type, FileFormat::Legacy); + assert_eq!(metadata.byte_order, ByteOrder::LittleEndian); + assert_eq!(metadata.tensor_count, 3); // weight, bias, running_mean + assert!(metadata.total_data_size.is_some()); +} + +#[test] +fn test_legacy_metadata_detailed() { + // Detailed test to prove we load all metadata for legacy format files + let path = test_data_path("simple_legacy.pt"); + let reader = PytorchReader::new(&path).expect("Failed to load legacy file"); + + // Get and examine metadata + let metadata = reader.metadata(); + + // Verify the metadata is correct for legacy format + assert_eq!( + metadata.format_type, + FileFormat::Legacy, + "Should be Legacy format" + ); + assert_eq!( + metadata.byte_order, + ByteOrder::LittleEndian, + "Legacy format is little-endian" + ); + assert_eq!( + metadata.tensor_count, 3, + "Should have 3 tensors: weight, bias, running_mean" + ); + assert!( + metadata.total_data_size.is_some(), + "Should have total data size" + ); + assert!( + metadata.total_data_size.unwrap() > 0, + "Data size should be positive" + ); + + // Legacy format specifics + assert_eq!( + metadata.format_version, None, + "Legacy format doesn't have version file" + ); + assert_eq!( + metadata.pytorch_version, None, + "Legacy format doesn't store PyTorch version reliably" + ); + assert!( + !metadata.has_storage_alignment, + "Legacy format doesn't have storage alignment" + ); + + // Also verify we can access the tensors + let keys = reader.keys(); + assert!( + keys.contains(&"weight".to_string()), + "Should have weight tensor" + ); + assert!( + keys.contains(&"bias".to_string()), + "Should have bias tensor" + ); + assert!( + keys.contains(&"running_mean".to_string()), + "Should have running_mean tensor" + ); +} + +#[test] +fn test_small_invalid_file() { + // Test that we handle broken/invalid files gracefully + let path = test_data_path("broken.pt"); + + // Should fail gracefully with an appropriate error + let result = PytorchReader::new(&path); + assert!(result.is_err(), "Expected error for broken file"); + + // The error should be a pickle error since the file is too small to be valid + if let Err(e) = result { + let err_str = format!("{}", e); + assert!( + err_str.contains("Pickle") || err_str.contains("Invalid"), + "Error should mention pickle or invalid format: {}", + err_str + ); + } +} + +#[test] +fn test_read_pickle_data_basic() { + use crate::pytorch::reader::PickleValue; + + // Test reading pickle data from a checkpoint file + let path = test_data_path("checkpoint.pt"); + + // Read the entire pickle data + let data = PytorchReader::read_pickle_data(&path, None).expect("Failed to read pickle data"); + + // Should be a dictionary at the root + if let PickleValue::Dict(dict) = data { + // Check that expected keys exist + assert!(dict.contains_key("model_state_dict")); + assert!(dict.contains_key("optimizer_state_dict")); + assert!(dict.contains_key("epoch")); + assert!(dict.contains_key("loss")); + + // Check epoch value + if let Some(PickleValue::Int(epoch)) = dict.get("epoch") { + assert_eq!(*epoch, 42); + } else { + panic!("Expected epoch to be an integer"); + } + + // Check loss value + if let Some(PickleValue::Float(loss)) = dict.get("loss") { + assert!(*loss > 0.0 && *loss < 1.0, "Loss should be between 0 and 1"); + } else { + panic!("Expected loss to be a float"); + } + } else { + panic!("Expected root to be a dictionary"); + } +} + +#[test] +fn test_read_pickle_data_with_key() { + use crate::pytorch::reader::PickleValue; + + // Test reading specific key from checkpoint + let path = test_data_path("checkpoint.pt"); + + // Read only the model_state_dict + let data = PytorchReader::read_pickle_data(&path, Some("model_state_dict")) + .expect("Failed to read pickle data with key"); + + // Should get the model_state_dict directly + if let PickleValue::Dict(dict) = data { + // Should have model weights + assert!(dict.contains_key("fc1.weight")); + assert!(dict.contains_key("fc1.bias")); + assert!(dict.contains_key("fc2.weight")); + assert!(dict.contains_key("fc2.bias")); + + // Should NOT have optimizer keys + assert!(!dict.contains_key("optimizer_state_dict")); + assert!(!dict.contains_key("epoch")); + } else { + panic!("Expected model_state_dict to be a dictionary"); + } +} + +#[test] +fn test_read_pickle_data_nested_structure() { + use crate::pytorch::reader::PickleValue; + + // Test reading nested dictionary structure + let path = test_data_path("nested_dict.pt"); + + let data = + PytorchReader::read_pickle_data(&path, None).expect("Failed to read nested structure"); + + if let PickleValue::Dict(dict) = data { + // nested_dict.pt has a nested structure, not flat keys + // It should have layer1 and layer2 as nested dicts + assert!(!dict.is_empty(), "Dictionary should not be empty"); + + // The structure depends on how the file was saved + // It could be flat keys like "layer1.weight" or nested dicts + // Just verify it's a valid dict structure + for (_key, value) in dict.iter() { + // Values could be None (tensors), nested dicts, or other types + assert!( + matches!(value, PickleValue::None | PickleValue::Dict(_)), + "Values should be None or nested dicts" + ); + } + } else { + panic!("Expected nested_dict to be a dictionary"); + } +} + +#[test] +fn test_read_pickle_data_types() { + use crate::pytorch::reader::PickleValue; + + // Test various data types in mixed_types.pt + let path = test_data_path("mixed_types.pt"); + + let data = PytorchReader::read_pickle_data(&path, None).expect("Failed to read mixed types"); + + if let PickleValue::Dict(dict) = data { + // The file contains different tensor types + assert!(dict.len() >= 3, "Should have at least 3 tensor types"); + + // All tensor values should be None in pickle data + for (_key, value) in dict.iter() { + // All values should be None (tensors are not included in pickle data) + assert!( + matches!(value, PickleValue::None), + "Tensors should be None in pickle data" + ); + } + } else { + panic!("Expected mixed_types to be a dictionary"); + } +} + +#[test] +fn test_read_pickle_data_key_not_found() { + // Test error handling when key doesn't exist + let path = test_data_path("checkpoint.pt"); + + let result = PytorchReader::read_pickle_data(&path, Some("nonexistent_key")); + assert!(result.is_err()); + + if let Err(e) = result { + let err_str = format!("{}", e); + assert!( + err_str.contains("not found"), + "Error should mention key not found: {}", + err_str + ); + } +} + +#[test] +fn test_read_pickle_data_simple_pickle() { + use crate::pytorch::reader::PickleValue; + + // Test reading a simple pickle file (not ZIP) + // Note: simple_legacy.pt is a legacy format file, not a simple pickle + // It may return None because legacy format reading is different + let path = test_data_path("state_dict.pt"); // Use a proper simple pickle file + + let data = PytorchReader::read_pickle_data(&path, None).expect("Failed to read simple pickle"); + + // Should contain state dict entries + if let PickleValue::Dict(dict) = data { + // state_dict.pt has weight, bias, running_mean, running_var + assert!(dict.len() >= 3); + assert!(dict.contains_key("weight")); + assert!(dict.contains_key("bias")); + + // All tensor values should be None in pickle data + for (_key, value) in dict.iter() { + assert!(matches!(value, PickleValue::None)); + } + } else { + panic!("Expected state_dict to contain a dictionary"); + } +} + +#[test] +fn test_load_config_basic() { + let path = test_data_path("checkpoint.pt"); + + // Define a struct that matches part of the checkpoint data + #[derive(Debug, serde::Deserialize, PartialEq)] + struct CheckpointConfig { + epoch: i64, + loss: f64, + } + + // Load config + let config: CheckpointConfig = + PytorchReader::load_config(&path, None).expect("Failed to load config"); + + // Verify values - based on test_read_pickle_data_basic + assert_eq!(config.epoch, 42); + assert!((config.loss - 0.123).abs() < 1e-6); +} + +#[test] +fn test_load_config_with_top_level_key() { + // Test that we can extract a non-existent key and get an appropriate error + let path = test_data_path("checkpoint.pt"); + + #[derive(Debug, serde::Deserialize, PartialEq)] + struct DummyConfig { + field: String, + } + + // Try loading with a valid top-level key that exists but has wrong structure + let result: Result = PytorchReader::load_config(&path, Some("epoch")); + + // This should fail because epoch is an integer, not a struct with a field + assert!(result.is_err()); + + // Now test that we can load with a real key that has the right structure + // Since checkpoint.pt doesn't have nested configs, let's use nested_dict.pt + let path2 = test_data_path("nested_dict.pt"); + + // Try to extract a specific nested key if it exists + // Since nested_dict has complex structure, let's just verify we can read it + let data = PytorchReader::read_pickle_data(&path2, None).unwrap(); + + // Verify it's a dict + if let crate::pytorch::reader::PickleValue::Dict(dict) = data { + assert!(!dict.is_empty()); + } else { + panic!("Expected a dict"); + } +} + +#[test] +fn test_load_config_complex_types() { + // For this test, let's create a comprehensive test using checkpoint.pt + // which has both metadata and state_dict fields + let path = test_data_path("checkpoint.pt"); + + // Define a partial config that only captures metadata fields + #[derive(Debug, serde::Deserialize, PartialEq)] + struct PartialCheckpoint { + epoch: i64, + loss: f64, + // We skip model_state_dict and optimizer_state_dict + // as they contain tensor references that become None + } + + // Load partial config + let config: PartialCheckpoint = + PytorchReader::load_config(&path, None).expect("Failed to load config"); + + // Verify we can extract the metadata + assert_eq!(config.epoch, 42); + assert!((config.loss - 0.123).abs() < 1e-6); +} + +#[test] +fn test_load_config_key_not_found() { + let path = test_data_path("checkpoint.pt"); + + #[derive(Debug, serde::Deserialize)] + struct DummyConfig { + #[allow(dead_code)] + field: String, + } + + // Try to load with non-existent key + let result: Result = PytorchReader::load_config(&path, Some("nonexistent")); + + assert!(result.is_err()); + let error = result.unwrap_err(); + assert!(error.to_string().contains("not found") || error.to_string().contains("Key")); +} + +#[test] +fn test_pickle_value_conversion() { + use crate::pytorch::reader::PickleValue; + + // Test that PickleValue provides useful data structures + let path = test_data_path("checkpoint.pt"); + let data = PytorchReader::read_pickle_data(&path, None).unwrap(); + + // Test pattern matching and data extraction + match data { + PickleValue::Dict(dict) => { + // Extract epoch as integer + if let Some(PickleValue::Int(epoch)) = dict.get("epoch") { + assert!(*epoch >= 0); + } + + // Extract loss as float + if let Some(PickleValue::Float(loss)) = dict.get("loss") { + assert!(loss.is_finite()); + } + + // Test nested access + if let Some(PickleValue::Dict(model_dict)) = dict.get("model_state_dict") { + assert!(!model_dict.is_empty()); + } + } + _ => panic!("Unexpected root type"), + } +} + +// ============================================================================ +// TAR Format Tests +// ============================================================================ +// The TAR format was used by very early versions of PyTorch (pre 0.1.10). +// These tests verify that we can correctly load models saved in this format. + +#[test] +fn test_tar_format_detection() { + // Test that is_tar_file correctly detects TAR files + let tar_path = test_data_path("tar_float32.tar"); + let zip_path = test_data_path("float32.pt"); + + // TAR file should be detected as TAR + let reader = PytorchReader::new(&tar_path).expect("Failed to load TAR file"); + let metadata = reader.metadata(); + assert_eq!(metadata.format_type, FileFormat::Tar); + + // ZIP file should NOT be detected as TAR + let reader = PytorchReader::new(&zip_path).expect("Failed to load ZIP file"); + let metadata = reader.metadata(); + assert_ne!(metadata.format_type, FileFormat::Tar); +} + +#[test] +fn test_tar_float32_tensor() { + let path = test_data_path("tar_float32.tar"); + let reader = PytorchReader::new(&path).expect("Failed to load tar_float32.tar"); + + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::F32); + assert_eq!(tensor.shape, shape![4]); + + let data = tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 4); + assert!((values[0] - 1.0).abs() < 1e-6); + assert!((values[1] - 2.5).abs() < 1e-6); + assert!((values[2] - (-3.7)).abs() < 1e-6); + assert!((values[3] - 0.0).abs() < 1e-6); +} + +#[test] +fn test_tar_float64_tensor() { + let path = test_data_path("tar_float64.tar"); + let reader = PytorchReader::new(&path).expect("Failed to load tar_float64.tar"); + + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::F64); + assert_eq!(tensor.shape, shape![3]); + + let data = tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 3); + assert!((values[0] - 1.1).abs() < 1e-10); + assert!((values[1] - 2.2).abs() < 1e-10); + assert!((values[2] - 3.3).abs() < 1e-10); +} + +#[test] +fn test_tar_int64_tensor() { + let path = test_data_path("tar_int64.tar"); + let reader = PytorchReader::new(&path).expect("Failed to load tar_int64.tar"); + + let tensor = reader.get("tensor").expect("tensor key not found"); + assert_eq!(tensor.dtype, DType::I64); + assert_eq!(tensor.shape, shape![4]); + + let data = tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values, &[100, -200, 300, 0]); +} + +#[test] +fn test_tar_multiple_tensors() { + // Test loading multiple tensors (weight + bias) with correct shapes + let path = test_data_path("tar_weight_bias.tar"); + let reader = PytorchReader::new(&path).expect("Failed to load tar_weight_bias.tar"); + + // Check weight tensor (2x3 matrix) + let weight = reader.get("weight").expect("weight key not found"); + assert_eq!(weight.dtype, DType::F32); + assert_eq!(weight.shape, shape![2, 3]); + + let data = weight.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 6); + assert!((values[0] - 0.1).abs() < 1e-6); + assert!((values[1] - 0.2).abs() < 1e-6); + assert!((values[5] - 0.6).abs() < 1e-6); + + // Check bias tensor (2-element vector) + let bias = reader.get("bias").expect("bias key not found"); + assert_eq!(bias.dtype, DType::F32); + assert_eq!(bias.shape, shape![2]); + + let data = bias.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 2); + assert!((values[0] - 0.01).abs() < 1e-6); + assert!((values[1] - 0.02).abs() < 1e-6); +} + +#[test] +fn test_tar_multi_dtype() { + // Test loading different dtypes from the same TAR file + let path = test_data_path("tar_multi_dtype.tar"); + let reader = PytorchReader::new(&path).expect("Failed to load tar_multi_dtype.tar"); + + // Float32 tensor + let float_tensor = reader + .get("float_tensor") + .expect("float_tensor key not found"); + assert_eq!(float_tensor.dtype, DType::F32); + let data = float_tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert!((values[0] - 1.5).abs() < 1e-6); + + // Float64 tensor + let double_tensor = reader + .get("double_tensor") + .expect("double_tensor key not found"); + assert_eq!(double_tensor.dtype, DType::F64); + let data = double_tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert!((values[0] - 1.111).abs() < 1e-10); + + // Int64 tensor + let int_tensor = reader.get("int_tensor").expect("int_tensor key not found"); + assert_eq!(int_tensor.dtype, DType::I64); + let data = int_tensor.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values, &[10, 20, 30, 40]); +} + +#[test] +fn test_tar_2d_tensor_shape() { + // Test that 2D tensor shapes are correctly preserved + let path = test_data_path("tar_2d_tensor.tar"); + let reader = PytorchReader::new(&path).expect("Failed to load tar_2d_tensor.tar"); + + let matrix = reader.get("matrix").expect("matrix key not found"); + assert_eq!(matrix.dtype, DType::F32); + assert_eq!(matrix.shape, shape![3, 4]); // 3 rows, 4 columns + + let data = matrix.to_data().unwrap(); + let values = data.as_slice::().unwrap(); + assert_eq!(values.len(), 12); + + // Verify values in row-major order + for i in 0..12 { + assert!((values[i] - (i as f32 + 1.0)).abs() < 1e-6); + } +} + +#[test] +fn test_tar_metadata() { + // Test that TAR metadata is correctly populated + let path = test_data_path("tar_float32.tar"); + let reader = PytorchReader::new(&path).expect("Failed to load tar_float32.tar"); + + let metadata = reader.metadata(); + assert_eq!(metadata.format_type, FileFormat::Tar); + assert_eq!(metadata.byte_order, ByteOrder::LittleEndian); + assert_eq!(metadata.tensor_count, 1); + assert!(metadata.total_data_size.is_some()); +} diff --git a/crates/burn-store/src/pytorch/tests/reader/simple_legacy.py b/crates/burn-store/src/pytorch/tests/reader/simple_legacy.py new file mode 100644 index 0000000..57610dd --- /dev/null +++ b/crates/burn-store/src/pytorch/tests/reader/simple_legacy.py @@ -0,0 +1,39 @@ +#!/usr/bin/env python3 +# /// script +# dependencies = ["torch"] +# /// +"""Create a simple legacy format PyTorch file.""" + +import torch + +# Create a simple state dict +state_dict = { + 'weight': torch.randn(2, 3), + 'bias': torch.ones(2), + 'running_mean': torch.zeros(2), +} + +# Save without using zip format (legacy format) +torch.save(state_dict, 'test_data/simple_legacy.pt', _use_new_zipfile_serialization=False) + +print("\nCreated simple_legacy.pt") + +# Verify +loaded = torch.load('test_data/simple_legacy.pt', weights_only=False) +print(f"Loaded {len(loaded)} tensors") +for key, val in loaded.items(): + print(f" {key}: shape {val.shape}, dtype {val.dtype}") + +# The storage has 100 elements while each tensor only uses 10 +base = torch.arange(100, dtype=torch.float32) +tensor1 = base[10:20] +tensor2 = base[50:60] +torch.save({'tensor1': tensor1, 'tensor2': tensor2}, 'test_data/legacy_uncloned_views.pt', + _use_new_zipfile_serialization=False) + +# Verify +print("\nCreated legacy_uncloned_views.pt") +loaded2 = torch.load('test_data/legacy_uncloned_views.pt', weights_only=False) +print(f"Loaded {len(loaded2)} tensors") +for key, val in loaded2.items(): + print(f" {key}: shape {val.shape}, dtype {val.dtype}") \ No newline at end of file diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data.py b/crates/burn-store/src/pytorch/tests/reader/test_data.py new file mode 100644 index 0000000..6020590 --- /dev/null +++ b/crates/burn-store/src/pytorch/tests/reader/test_data.py @@ -0,0 +1,227 @@ +#!/usr/bin/env python3 +# /// script +# dependencies = ["torch", "numpy"] +# /// +""" +Generate test PyTorch .pt files for testing the burn-store PyTorch reader. +Run with: uv run test_files.py +""" + +import torch +import numpy as np +import os +from pathlib import Path + +# Create test directory +test_dir = Path(__file__).parent / "test_data" +test_dir.mkdir(exist_ok=True) + +def save_test_file(filename, data, description): + """Save a test file and print what was saved.""" + filepath = test_dir / filename + torch.save(data, filepath) + print(f"✓ {filename}: {description}") + return filepath + +# Test 1: Simple tensors of different types +print("\n=== Generating Basic Tensor Tests ===") + +# Float32 tensor (wrap in dict for compatibility) +float32_tensor = torch.tensor([1.0, 2.5, -3.7, 0.0], dtype=torch.float32) +save_test_file("float32.pt", {"tensor": float32_tensor}, "Float32 tensor [1.0, 2.5, -3.7, 0.0]") + +# Float64 tensor +float64_tensor = torch.tensor([1.1, 2.2, 3.3], dtype=torch.float64) +save_test_file("float64.pt", {"tensor": float64_tensor}, "Float64 tensor [1.1, 2.2, 3.3]") + +# Int64 tensor +int64_tensor = torch.tensor([100, -200, 300, 0], dtype=torch.int64) +save_test_file("int64.pt", {"tensor": int64_tensor}, "Int64 tensor [100, -200, 300, 0]") + +# Int32 tensor +int32_tensor = torch.tensor([10, 20, -30], dtype=torch.int32) +save_test_file("int32.pt", {"tensor": int32_tensor}, "Int32 tensor [10, 20, -30]") + +# Int16 tensor +int16_tensor = torch.tensor([1000, -2000, 3000], dtype=torch.int16) +save_test_file("int16.pt", {"tensor": int16_tensor}, "Int16 tensor [1000, -2000, 3000]") + +# Int8 tensor +int8_tensor = torch.tensor([127, -128, 0, 50], dtype=torch.int8) +save_test_file("int8.pt", {"tensor": int8_tensor}, "Int8 tensor [127, -128, 0, 50]") + +# Boolean tensor +bool_tensor = torch.tensor([True, False, True, True, False], dtype=torch.bool) +save_test_file("bool.pt", {"tensor": bool_tensor}, "Bool tensor [True, False, True, True, False]") + +# Float16 tensor (half precision) +float16_tensor = torch.tensor([1.5, -2.25, 3.125], dtype=torch.float16) +save_test_file("float16.pt", {"tensor": float16_tensor}, "Float16 tensor [1.5, -2.25, 3.125]") + +# BFloat16 tensor +bfloat16_tensor = torch.tensor([1.5, -2.5, 3.5], dtype=torch.bfloat16) +save_test_file("bfloat16.pt", {"tensor": bfloat16_tensor}, "BFloat16 tensor [1.5, -2.5, 3.5]") + +# UInt8 tensor +uint8_tensor = torch.tensor([0, 128, 255, 42], dtype=torch.uint8) +save_test_file("uint8.pt", {"tensor": uint8_tensor}, "UInt8 tensor [0, 128, 255, 42]") + +# Test 2: Multi-dimensional tensors +print("\n=== Generating Multi-dimensional Tensor Tests ===") + +# 2D tensor +tensor_2d = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=torch.float32) +save_test_file("tensor_2d.pt", {"tensor": tensor_2d}, "2D tensor shape (3, 2)") + +# 3D tensor +torch.manual_seed(42) +tensor_3d = torch.randn(2, 3, 4) * 10 +save_test_file("tensor_3d.pt", {"tensor": tensor_3d}, "3D tensor shape (2, 3, 4)") + +# 4D tensor (common for conv weights) +tensor_4d = torch.randn(2, 3, 2, 2) +save_test_file("tensor_4d.pt", {"tensor": tensor_4d}, "4D tensor shape (2, 3, 2, 2)") + +# Test 3: State dict (multiple tensors) +print("\n=== Generating State Dict Tests ===") + +state_dict = { + "weight": torch.randn(3, 4), + "bias": torch.randn(3), + "running_mean": torch.zeros(3), + "running_var": torch.ones(3), +} +save_test_file("state_dict.pt", state_dict, "State dict with 4 tensors") + +# Nested state dict +nested_dict = { + "layer1": { + "weight": torch.randn(2, 3), + "bias": torch.randn(2) + }, + "layer2": { + "weight": torch.randn(4, 2), + "bias": torch.randn(4) + } +} +save_test_file("nested_dict.pt", nested_dict, "Nested state dict") + +# Test 4: Model checkpoint format +print("\n=== Generating Model Checkpoint Tests ===") + +# Typical checkpoint format (use string keys for compatibility) +checkpoint = { + "model_state_dict": { + "fc1.weight": torch.randn(10, 5), + "fc1.bias": torch.randn(10), + "fc2.weight": torch.randn(3, 10), + "fc2.bias": torch.randn(3), + }, + "optimizer_state_dict": { + "state": { + "0": { # Use string key instead of integer + "momentum_buffer": torch.randn(10, 5) + } + } + }, + "epoch": 42, + "loss": 0.123 +} +save_test_file("checkpoint.pt", checkpoint, "Full checkpoint with model and optimizer state") + +# Test 5: Edge cases +print("\n=== Generating Edge Case Tests ===") + +# Empty tensor (1D with 0 elements) +empty_tensor = torch.zeros(0) +save_test_file("empty.pt", {"tensor": empty_tensor}, "Empty tensor") + +# Scalar tensor (0-dimensional) +scalar_tensor = torch.tensor(42.0) +save_test_file("scalar.pt", {"tensor": scalar_tensor}, "Scalar tensor (0-dim)") + +# Large shape but small data (testing shape vs actual data) +sparse_like = torch.zeros(100, 100) +sparse_like[0, 0] = 1.0 +sparse_like[50, 50] = 2.0 +sparse_like[99, 99] = 3.0 +save_test_file("large_shape.pt", {"tensor": sparse_like}, "Large shape (100, 100) mostly zeros") + +# Test 6: Mixed types in dict +print("\n=== Generating Mixed Type Tests ===") + +mixed_types = { + "float32": torch.tensor([1.0, 2.0], dtype=torch.float32), + "int64": torch.tensor([100, 200], dtype=torch.int64), + "bool": torch.tensor([True, False], dtype=torch.bool), + "float64": torch.tensor([1.1, 2.2], dtype=torch.float64), +} +save_test_file("mixed_types.pt", mixed_types, "Dict with mixed tensor types") + +# Test 7: Special values +print("\n=== Generating Special Value Tests ===") + +# NaN and Inf values +special_values = torch.tensor([float('nan'), float('inf'), float('-inf'), 0.0, 1.0]) +save_test_file("special_values.pt", {"tensor": special_values}, "Tensor with NaN and Inf") + +# Very small and very large values +extreme_values = torch.tensor([1e-30, 1e30, -1e-30, -1e30], dtype=torch.float32) +save_test_file("extreme_values.pt", {"tensor": extreme_values}, "Tensor with extreme values") + +# Test 8: Parameter wrapper (common in models) +print("\n=== Generating Parameter Tests ===") + +import torch.nn as nn +param = nn.Parameter(torch.randn(3, 3)) +param_dict = {"param": param} +save_test_file("parameter.pt", param_dict, "nn.Parameter wrapped tensor") + +# Test 9: Buffer-style tensors +print("\n=== Generating Buffer Tests ===") + +# Simulate model buffers +buffers = { + "buffer1": torch.tensor([1, 2, 3], dtype=torch.int32), + "buffer2": torch.tensor([True, False], dtype=torch.bool), +} +save_test_file("buffers.pt", buffers, "Model buffers") + +# Test 10: Complex nested structure +print("\n=== Generating Complex Structure Tests ===") + +complex_structure = { + "metadata": { + "version": 1, + "name": "test_model" + }, + "state": { + "encoder": { + "layer_0": { + "weight": torch.randn(4, 3), + "bias": torch.randn(4) + }, + "layer_1": { + "weight": torch.randn(2, 4), + "bias": torch.randn(2) + } + }, + "decoder": { + "weight": torch.randn(3, 2), + "bias": torch.randn(3) + } + }, + "config": { + "hidden_size": 4, + "num_layers": 2 + } +} +save_test_file("complex_structure.pt", complex_structure, "Complex nested structure") + +print(f"\n✅ Generated {len(list(test_dir.glob('*.pt')))} test files in {test_dir}") +print("\nTest files can be used to verify PyTorch reader functionality:") +print("- Different data types (float32, int64, bool, etc.)") +print("- Multi-dimensional tensors") +print("- State dicts and nested structures") +print("- Edge cases (empty, scalar, special values)") +print("- Model checkpoints and parameters") diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/bfloat16.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/bfloat16.pt new file mode 100644 index 0000000..6e5c973 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/bfloat16.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/bool.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/bool.pt new file mode 100644 index 0000000..1ca216c Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/bool.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/broken.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/broken.pt new file mode 100644 index 0000000..f2ba8f8 --- /dev/null +++ b/crates/burn-store/src/pytorch/tests/reader/test_data/broken.pt @@ -0,0 +1 @@ +abc \ No newline at end of file diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/buffers.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/buffers.pt new file mode 100644 index 0000000..75d39e8 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/buffers.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/checkpoint.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/checkpoint.pt new file mode 100644 index 0000000..5c73260 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/checkpoint.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/complex_structure.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/complex_structure.pt new file mode 100644 index 0000000..5c2087c Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/complex_structure.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/empty.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/empty.pt new file mode 100644 index 0000000..1985924 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/empty.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/extreme_values.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/extreme_values.pt new file mode 100644 index 0000000..dec83dc Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/extreme_values.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/float16.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/float16.pt new file mode 100644 index 0000000..4262198 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/float16.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/float32.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/float32.pt new file mode 100644 index 0000000..1749780 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/float32.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/float64.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/float64.pt new file mode 100644 index 0000000..d49f4af Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/float64.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/int16.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/int16.pt new file mode 100644 index 0000000..624380e Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/int16.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/int32.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/int32.pt new file mode 100644 index 0000000..fa33c8d Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/int32.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/int64.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/int64.pt new file mode 100644 index 0000000..99611be Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/int64.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/int8.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/int8.pt new file mode 100644 index 0000000..a02b42b Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/int8.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/large_shape.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/large_shape.pt new file mode 100644 index 0000000..f1e9417 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/large_shape.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/legacy_shared_storage.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/legacy_shared_storage.pt new file mode 100644 index 0000000..21d849c Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/legacy_shared_storage.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/legacy_uncloned_views.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/legacy_uncloned_views.pt new file mode 100644 index 0000000..2b1fb2e Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/legacy_uncloned_views.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/legacy_with_offsets.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/legacy_with_offsets.pt new file mode 100644 index 0000000..6cc448f Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/legacy_with_offsets.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/mixed_types.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/mixed_types.pt new file mode 100644 index 0000000..031f4f5 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/mixed_types.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/nested_dict.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/nested_dict.pt new file mode 100644 index 0000000..686286c Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/nested_dict.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/parameter.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/parameter.pt new file mode 100644 index 0000000..53e3b29 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/parameter.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/scalar.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/scalar.pt new file mode 100644 index 0000000..7b579f9 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/scalar.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/simple_legacy.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/simple_legacy.pt new file mode 100644 index 0000000..61fb5fc Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/simple_legacy.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/special_values.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/special_values.pt new file mode 100644 index 0000000..f33ad35 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/special_values.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/state_dict.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/state_dict.pt new file mode 100644 index 0000000..d1869c0 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/state_dict.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/tar_2d_tensor.tar b/crates/burn-store/src/pytorch/tests/reader/test_data/tar_2d_tensor.tar new file mode 100644 index 0000000..abe8459 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/tar_2d_tensor.tar differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/tar_float32.tar b/crates/burn-store/src/pytorch/tests/reader/test_data/tar_float32.tar new file mode 100644 index 0000000..1463dc0 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/tar_float32.tar differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/tar_float64.tar b/crates/burn-store/src/pytorch/tests/reader/test_data/tar_float64.tar new file mode 100644 index 0000000..b6acf66 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/tar_float64.tar differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/tar_int64.tar b/crates/burn-store/src/pytorch/tests/reader/test_data/tar_int64.tar new file mode 100644 index 0000000..abf626c Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/tar_int64.tar differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/tar_multi_dtype.tar b/crates/burn-store/src/pytorch/tests/reader/test_data/tar_multi_dtype.tar new file mode 100644 index 0000000..4524991 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/tar_multi_dtype.tar differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/tar_weight_bias.tar b/crates/burn-store/src/pytorch/tests/reader/test_data/tar_weight_bias.tar new file mode 100644 index 0000000..24961dc Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/tar_weight_bias.tar differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/tensor_2d.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/tensor_2d.pt new file mode 100644 index 0000000..d27c30a Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/tensor_2d.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/tensor_3d.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/tensor_3d.pt new file mode 100644 index 0000000..c77b3a3 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/tensor_3d.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/tensor_4d.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/tensor_4d.pt new file mode 100644 index 0000000..644a3f0 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/tensor_4d.pt differ diff --git a/crates/burn-store/src/pytorch/tests/reader/test_data/uint8.pt b/crates/burn-store/src/pytorch/tests/reader/test_data/uint8.pt new file mode 100644 index 0000000..97cfc77 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/reader/test_data/uint8.pt differ diff --git a/crates/burn-store/src/pytorch/tests/store/mod.rs b/crates/burn-store/src/pytorch/tests/store/mod.rs new file mode 100644 index 0000000..324e629 --- /dev/null +++ b/crates/burn-store/src/pytorch/tests/store/mod.rs @@ -0,0 +1,1202 @@ +//! Comprehensive tests for PytorchStore with real model application +use burn_core as burn; + +use std::path::PathBuf; + +use crate::ModuleStore; +use crate::pytorch::PytorchStore; +use burn_core::module::Module; +use burn_core::tensor::Tensor; +use burn_nn::conv::{Conv2d, Conv2dConfig}; +use burn_nn::{Linear, LinearConfig}; + +/// Path to pytorch test files (now under burn-store) +fn pytorch_test_path(subdir: &str, filename: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("pytorch-tests") + .join("tests") + .join(subdir) + .join(filename) +} + +/// Path to burn-store test data files +fn test_data_path(filename: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("src") + .join("pytorch") + .join("tests") + .join("reader") + .join("test_data") + .join(filename) +} + +/// Path to store test data files +fn store_test_data_path(filename: &str) -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("src") + .join("pytorch") + .join("tests") + .join("store") + .join("test_data") + .join(filename) +} + +#[cfg(test)] +mod basic_tests { + use super::*; + + #[test] + fn test_store_creation() { + let store = PytorchStore::from_file("model.pth"); + assert!(store.validate); + assert!(!store.allow_partial); + assert!(store.top_level_key.is_none()); + // Contiguous index mapping is enabled by default for PyTorch files + assert!(store.map_indices_contiguous); + } + + #[test] + fn test_store_map_indices_contiguous_default() { + // Verify that map_indices_contiguous is enabled by default + let store = PytorchStore::from_file("model.pth"); + assert!( + store.map_indices_contiguous, + "map_indices_contiguous should be enabled by default" + ); + } + + #[test] + fn test_store_map_indices_contiguous_disabled() { + // Verify that we can disable map_indices_contiguous + let store = PytorchStore::from_file("model.pth").map_indices_contiguous(false); + assert!( + !store.map_indices_contiguous, + "map_indices_contiguous should be disabled after explicit call" + ); + } + + #[test] + fn test_store_with_top_level_key() { + let store = PytorchStore::from_file("model.pth").with_top_level_key("state_dict"); + assert_eq!(store.top_level_key, Some("state_dict".to_string())); + } + + #[test] + fn test_store_configuration() { + let store = PytorchStore::from_file("model.pth") + .validate(false) + .allow_partial(true) + .with_regex(r"^encoder\.") + .with_full_path("decoder.weight"); + + assert!(!store.validate); + assert!(store.allow_partial); + assert!(!store.filter.is_empty()); + } + + #[test] + fn test_store_with_remapping() { + let store = PytorchStore::from_file("model.pth").with_key_remapping(r"^old\.", "new."); + + assert!(!store.remapper.is_empty()); + } + + #[test] + fn test_store_save_not_supported() { + // Currently, saving to PyTorch format is not implemented + // The collect_from method always returns an error + let store = PytorchStore::from_file("test.pth"); + + // Just verify that store creation works + assert!(store.validate); + + // Note: Actually testing save would require a proper Module implementation + // which is complex. The implementation guarantees it returns an error. + } +} + +#[cfg(test)] +mod linear_model_tests { + use burn_core::tensor::Device; + + use super::*; + + #[derive(Module, Debug)] + pub struct SimpleLinearModel { + fc1: Linear, + fc2: Linear, + } + + impl SimpleLinearModel { + pub fn new(device: &Device) -> Self { + Self { + fc1: LinearConfig::new(2, 3).init(device), + fc2: LinearConfig::new(3, 4).init(device), + } + } + + pub fn forward(&self, x: Tensor<2>) -> Tensor<2> { + let x = self.fc1.forward(x); + self.fc2.forward(x) + } + } + + #[test] + fn test_load_linear_model() { + let device = Default::default(); + let path = pytorch_test_path("linear", "linear.pt"); + + // Create a model and load weights from PyTorch + let mut model = SimpleLinearModel::new(&device); + let mut store = PytorchStore::from_file(path).allow_partial(true); + + // Apply the PyTorch weights to our model + let result = store.apply_to(&mut model); + + assert!( + result.is_ok(), + "Failed to load linear model: {:?}", + result.err() + ); + + let result = result.unwrap(); + assert!(!result.applied.is_empty(), "No tensors were applied"); + + // Test forward pass with loaded weights + let input = Tensor::<2>::ones([1, 2], &device); + let output = model.forward(input); + + // Verify output shape + assert_eq!(&*output.shape(), [1, 4]); + } + + #[test] + fn test_load_linear_with_bias() { + let device = Default::default(); + let path = pytorch_test_path("linear", "linear_with_bias.pt"); + + // Single linear layer with bias + #[derive(Module, Debug)] + struct LinearWithBias { + fc1: Linear, + } + + let mut model = LinearWithBias { + fc1: LinearConfig::new(2, 3).init(&device), + }; + + let mut store = PytorchStore::from_file(path).allow_partial(true); + + let result = store.apply_to(&mut model); + assert!(result.is_ok(), "Failed to load model with bias"); + + // Verify biases were loaded + let result = result.unwrap(); + let bias_loaded = result.applied.iter().any(|s| s.contains("bias")); + assert!(bias_loaded, "Bias parameters not loaded"); + } + + #[test] + fn test_filter_layers() { + let device = Default::default(); + let path = pytorch_test_path("linear", "linear.pt"); + + let mut model = SimpleLinearModel::new(&device); + + // Only load fc1 layers + let mut store = PytorchStore::from_file(path) + .with_regex(r"^fc1\.") + .allow_partial(true); + + let result = store.apply_to(&mut model).unwrap(); + + // Should only have fc1 tensors + for tensor in &result.applied { + assert!(tensor.contains("fc1")); + assert!(!tensor.contains("fc2")); + } + } + + #[test] + fn test_remap_layer_names() { + let device = Default::default(); + let path = pytorch_test_path("linear", "linear.pt"); + + // Model with different layer names + #[derive(Module, Debug)] + struct RemappedModel { + linear1: Linear, + linear2: Linear, + } + + let mut model = RemappedModel { + linear1: LinearConfig::new(2, 3).init(&device), + linear2: LinearConfig::new(3, 4).init(&device), + }; + + let mut store = PytorchStore::from_file(path) + .with_key_remapping(r"^fc1\.", "linear1.") + .with_key_remapping(r"^fc2\.", "linear2.") + .allow_partial(true); + + let result = store.apply_to(&mut model); + assert!(result.is_ok(), "Failed to load with remapped names"); + + let result = result.unwrap(); + // Verify remapped names were applied + let has_linear1 = result.applied.iter().any(|s| s.contains("linear1")); + assert!(has_linear1, "Remapped names not applied"); + } +} + +#[cfg(test)] +mod conv_model_tests { + use burn_core::tensor::Device; + + use super::*; + + #[derive(Module, Debug)] + struct SimpleConvModel { + conv1: Conv2d, + conv2: Conv2d, + } + + impl SimpleConvModel { + pub fn new(device: &Device) -> Self { + Self { + conv1: Conv2dConfig::new([3, 16], [3, 3]).init(device), + conv2: Conv2dConfig::new([16, 32], [3, 3]).init(device), + } + } + } + + #[test] + fn test_load_conv2d_model() { + let device = Default::default(); + let path = pytorch_test_path("conv2d", "conv2d.pt"); + + // Check if file exists, skip if not + if !path.exists() { + println!("Skipping conv2d test - file not found: {:?}", path); + return; + } + + let mut model = SimpleConvModel::new(&device); + let mut store = PytorchStore::from_file(path).allow_partial(true); + + let result = store.apply_to(&mut model); + + if let Ok(result) = result { + assert!(!result.applied.is_empty(), "No conv tensors applied"); + + // Check for conv weights + let has_conv_weights = result.applied.iter().any(|s| s.contains("weight")); + assert!(has_conv_weights, "Conv weights not loaded"); + } + } + + #[test] + fn test_load_conv1d_model() { + let path = pytorch_test_path("conv1d", "conv1d.pt"); + + if !path.exists() { + println!("Skipping conv1d test - file not found: {:?}", path); + return; + } + + // Just test that we can create a store for conv1d files + let store = PytorchStore::from_file(path).allow_partial(true); + + assert!(store.allow_partial); + } +} + +#[cfg(test)] +mod complex_model_tests { + use super::*; + + #[test] + fn test_load_with_top_level_key() { + let path = test_data_path("checkpoint.pt"); + + // Just verify that we can create a store with top-level key + let store = PytorchStore::from_file(path) + .with_top_level_key("model_state_dict") + .allow_partial(true); + + assert_eq!(store.top_level_key, Some("model_state_dict".to_string())); + } + + #[test] + fn test_load_nested_structure() { + let path = test_data_path("complex_structure.pt"); + + // Just verify that we can create a store for nested structure + let store = PytorchStore::from_file(path).allow_partial(true); + + assert!(store.allow_partial); + } + + #[test] + fn test_legacy_format() { + let path = test_data_path("simple_legacy.pt"); + + if !path.exists() { + println!("Skipping legacy format test - file not found: {:?}", path); + return; + } + + // Just verify that we can create a store for legacy format + let store = PytorchStore::from_file(path).allow_partial(true); + + assert!(store.allow_partial); + + // Could load into an actual model if we had legacy model structure + } + + #[test] + fn test_key_remap_chained() { + let path = pytorch_test_path("linear", "linear.pt"); + + if !path.exists() { + println!("Skipping key remap test - file not found: {:?}", path); + return; + } + + let device = Default::default(); + + // Model with different layer names that need remapping + #[derive(Module, Debug)] + struct RemappedChainModel { + convolution1: Linear, // Will be remapped from fc1 + linear2: Linear, // Will be remapped from fc2 + } + + let mut model = RemappedChainModel { + convolution1: LinearConfig::new(2, 3).init(&device), + linear2: LinearConfig::new(3, 4).init(&device), + }; + + // Chain multiple remappings + let mut store = PytorchStore::from_file(path) + .with_key_remapping(r"^fc1\.", "convolution1.") + .with_key_remapping(r"^fc2\.", "linear2.") + .allow_partial(true); + + let result = store.apply_to(&mut model); + + if let Ok(result) = result { + // Check that remapped names were applied + assert!( + !result.applied.is_empty(), + "No tensors were applied after remapping" + ); + } + } +} + +#[cfg(test)] +mod adapter_tests { + use burn_core::tensor::Device; + + use super::*; + + #[derive(Module, Debug)] + pub struct SimpleLinearModel { + fc1: Linear, + fc2: Linear, + } + + impl SimpleLinearModel { + pub fn new(device: &Device) -> Self { + Self { + fc1: LinearConfig::new(2, 3).init(device), + fc2: LinearConfig::new(3, 4).init(device), + } + } + } + + #[test] + fn test_pytorch_adapter_always_applied() { + // Test that PyTorchToBurnAdapter is always applied internally + let path = pytorch_test_path("linear", "linear.pt"); + + if !path.exists() { + println!("Skipping adapter test - file not found: {:?}", path); + return; + } + + let device = Default::default(); + let mut model = SimpleLinearModel::new(&device); + + let mut store = PytorchStore::from_file(path).allow_partial(true); + + let result = store.apply_to(&mut model); + + // PyTorchToBurnAdapter is always applied internally + assert!( + result.is_ok(), + "Failed to load with internal PyTorchToBurnAdapter: {:?}", + result.err() + ); + assert!(!result.unwrap().applied.is_empty()); + } + + #[test] + fn test_pytorch_adapter_with_filtering() { + // Test that PyTorchToBurnAdapter works with filtering + let path = pytorch_test_path("linear", "linear.pt"); + + if !path.exists() { + println!("Skipping filtering test - file not found: {:?}", path); + return; + } + + let device = Default::default(); + let mut model = SimpleLinearModel::new(&device); + + // Filter to exclude bias tensors + let mut store = PytorchStore::from_file(path) + .with_predicate(|path, _| !path.contains("bias")) + .allow_partial(true); + + let result = store.apply_to(&mut model).unwrap(); + + // Should not have any bias tensors due to filtering + for applied_path in &result.applied { + assert!( + !applied_path.contains("bias"), + "Bias tensor was not filtered: {}", + applied_path + ); + } + } +} + +#[cfg(test)] +mod error_handling_tests { + use burn_core::tensor::Device; + + use super::*; + + #[derive(Module, Debug)] + pub struct SimpleLinearModel { + fc1: Linear, + fc2: Linear, + } + + impl SimpleLinearModel { + pub fn new(device: &Device) -> Self { + Self { + fc1: LinearConfig::new(2, 3).init(device), + fc2: LinearConfig::new(3, 4).init(device), + } + } + } + + #[test] + fn test_missing_file() { + let device = Default::default(); + let mut model = SimpleLinearModel::new(&device); + let mut store = PytorchStore::from_file("nonexistent.pth"); + + let result = store.apply_to(&mut model); + + assert!(result.is_err()); + match result { + Err(crate::pytorch::PytorchStoreError::Reader(_)) => {} + _ => panic!("Expected reader error for missing file"), + } + } + + #[test] + fn test_invalid_top_level_key() { + let path = pytorch_test_path("linear", "linear.pt"); + + if !path.exists() { + println!( + "Skipping invalid top level key test - file not found: {:?}", + path + ); + return; + } + + let device = Default::default(); + let mut model = SimpleLinearModel::new(&device); + + let mut store = PytorchStore::from_file(path).with_top_level_key("nonexistent_key"); + + let result = store.apply_to(&mut model); + + assert!(result.is_err(), "Should fail with invalid top level key"); + } + + #[test] + fn test_strict_validation() { + let path = pytorch_test_path("linear", "linear.pt"); + + if !path.exists() { + println!( + "Skipping strict validation test - file not found: {:?}", + path + ); + return; + } + + let device = Default::default(); + let mut model = SimpleLinearModel::new(&device); + + // Apply very restrictive filter that matches nothing + let mut store = PytorchStore::from_file(path) + .with_regex(r"^this_will_never_match$") + .validate(true) + .allow_partial(false); + + let result = store.apply_to(&mut model); + + // Should fail because no tensors match and allow_partial is false + assert!( + result.is_err(), + "Should fail when no tensors match with allow_partial=false" + ); + } +} + +#[cfg(test)] +mod enum_variant_tests { + use burn_core::tensor::Device; + + use super::*; + use crate::ModuleSnapshot; + + /// Enum representing different convolution block types (similar to YOLOX architecture) + #[derive(Module, Debug)] + pub enum ConvBlock { + /// Base convolution block + BaseConv(Linear), + /// Depthwise separable convolution block + DwsConv(Linear), + } + + /// Model with enum field that will have variant names in Burn paths + #[derive(Module, Debug)] + pub struct ModelWithEnum { + /// Feature extractor with enum variants + feature: ConvBlock, + /// Output classifier + classifier: Linear, + } + + impl ModelWithEnum { + pub fn new(device: &Device) -> Self { + Self { + feature: ConvBlock::BaseConv(LinearConfig::new(3, 64).init(device)), + classifier: LinearConfig::new(64, 10).init(device), + } + } + } + + #[test] + fn test_enum_variant_path_mismatch() { + let device = Default::default(); + let mut model = ModelWithEnum::new(&device); + + // Load PyTorch model that was generated without enum variant names + // PyTorch paths: "feature.weight", "feature.bias", "classifier.weight", "classifier.bias" + // Burn paths: "feature.BaseConv.weight", "feature.BaseConv.bias", "classifier.weight", "classifier.bias" + // ^^^^^^^^ enum variant name is included in Burn but not PyTorch + + let pytorch_file = store_test_data_path("model_without_enum_variants.pt"); + + // Try to load from PyTorch format (without enum variants) + // Explicitly disable skip_enum_variants to demonstrate the mismatch problem + let mut store = PytorchStore::from_file(pytorch_file) + .skip_enum_variants(false) // Disable to show the mismatch + .allow_partial(true) // Allow partial to see what's missing + .validate(false); // Disable validation to get detailed missing info + + let result = store.apply_to(&mut model); + + // The load should succeed (allow_partial=true) but report missing tensors + match result { + Ok(apply_result) => { + // Verify we have missing tensors + assert!( + !apply_result.missing.is_empty(), + "Should have missing tensors due to enum variant path mismatch" + ); + + // Check that missing paths contain enum variants + let enum_missing: Vec<_> = apply_result + .missing + .iter() + .filter(|(_, container_stack)| container_stack.contains("Enum:")) + .collect(); + + assert!( + !enum_missing.is_empty(), + "Missing tensors should be detected as having enum containers" + ); + + // Verify the paths look like what we expect + let has_base_conv_path = apply_result + .missing + .iter() + .any(|(path, _)| path.contains("BaseConv")); + + assert!( + has_base_conv_path, + "Should have missing paths with 'BaseConv' enum variant. Missing: {:?}", + apply_result + .missing + .iter() + .map(|(p, _)| p) + .collect::>() + ); + + // Print the diagnostic output to show enum detection + println!("\n{}", apply_result); + + // Verify the diagnostic message mentions enum variants + let display_output = format!("{}", apply_result); + assert!( + display_output.contains("enum variant"), + "Display output should mention enum variants" + ); + } + Err(e) => panic!( + "Load should succeed with allow_partial=true, got error: {}", + e + ), + } + } + + #[test] + fn test_enum_variant_detection_in_container_stack() { + let device = Default::default(); + + // Create model with enum + let model = ModelWithEnum::new(&device); + + // Collect snapshots to inspect container stacks + let snapshots = model.collect(None, None, false); + + // Find a snapshot from inside the enum + let enum_snapshot = snapshots + .iter() + .find(|s| s.full_path().contains("feature")) + .expect("Should have feature snapshots"); + + // Verify container stack contains enum marker + if let Some(container_stack) = &enum_snapshot.container_stack { + let container_str = container_stack.join("."); + assert!( + container_str.contains("Enum:ConvBlock"), + "Container stack should contain Enum:ConvBlock marker. Got: {}", + container_str + ); + } else { + panic!("Snapshot should have container_stack"); + } + } + + #[test] + fn test_skip_enum_variants_feature() { + let device = Default::default(); + let mut model = ModelWithEnum::new(&device); + + // Load PyTorch model that was generated without enum variant names + // PyTorch paths: "feature.weight", "feature.bias", "classifier.weight", "classifier.bias" + // Burn paths: "feature.BaseConv.weight", "feature.BaseConv.bias", "classifier.weight", "classifier.bias" + + let pytorch_file = store_test_data_path("model_without_enum_variants.pt"); + + // Try to load with skip_enum_variants enabled + let mut store = PytorchStore::from_file(pytorch_file) + .skip_enum_variants(true) // Enable enum variant skipping + .allow_partial(true) + .validate(false); + + let result = store.apply_to(&mut model); + + // The load should succeed and all tensors should be loaded + match result { + Ok(apply_result) => { + println!("\n{}", apply_result); + + // With skip_enum_variants enabled, we should successfully load the feature tensors + let feature_applied = apply_result + .applied + .iter() + .filter(|path| path.contains("feature")) + .count(); + + assert!( + feature_applied > 0, + "Should have applied feature tensors with skip_enum_variants=true. Applied: {:?}", + apply_result.applied + ); + + // The feature tensors should NOT be in missing anymore + let feature_missing = apply_result + .missing + .iter() + .filter(|(path, _)| path.contains("feature")) + .count(); + + assert_eq!( + feature_missing, 0, + "Feature tensors should not be missing with skip_enum_variants=true. Missing: {:?}", + apply_result.missing + ); + } + Err(e) => panic!( + "Load with skip_enum_variants should succeed, got error: {}", + e + ), + } + } +} + +#[cfg(test)] +mod direct_access_tests { + use super::*; + + #[test] + fn test_get_all_snapshots() { + let path = pytorch_test_path("linear", "linear.pt"); + + if !path.exists() { + println!("Skipping test - file not found: {:?}", path); + return; + } + + let mut store = PytorchStore::from_file(path); + let snapshots = store.get_all_snapshots().unwrap(); + + // linear.pt should have fc1.weight, fc1.bias, fc2.weight, fc2.bias + assert!(!snapshots.is_empty(), "Should have snapshots"); + assert!( + snapshots.contains_key("fc1.weight"), + "Should contain fc1.weight" + ); + assert!( + snapshots.contains_key("fc1.bias"), + "Should contain fc1.bias" + ); + } + + #[test] + fn test_get_snapshot_existing() { + let path = pytorch_test_path("linear", "linear.pt"); + + if !path.exists() { + println!("Skipping test - file not found: {:?}", path); + return; + } + + let mut store = PytorchStore::from_file(path); + + // Get existing snapshot + let snapshot = store.get_snapshot("fc1.weight").unwrap(); + assert!(snapshot.is_some(), "Should find fc1.weight"); + + let snapshot = snapshot.unwrap(); + // Linear weight should be 2D + assert_eq!(snapshot.shape.len(), 2, "Weight should be 2D tensor"); + + // Verify we can load data + let data = snapshot.to_data().unwrap(); + assert!(!data.bytes.is_empty(), "Data should not be empty"); + } + + #[test] + fn test_get_snapshot_not_found() { + let path = pytorch_test_path("linear", "linear.pt"); + + if !path.exists() { + println!("Skipping test - file not found: {:?}", path); + return; + } + + let mut store = PytorchStore::from_file(path); + + // Get non-existent snapshot + let snapshot = store.get_snapshot("nonexistent.weight").unwrap(); + assert!(snapshot.is_none(), "Should not find nonexistent tensor"); + } + + #[test] + fn test_keys() { + let path = pytorch_test_path("linear", "linear.pt"); + + if !path.exists() { + println!("Skipping test - file not found: {:?}", path); + return; + } + + let mut store = PytorchStore::from_file(path); + let keys = store.keys().unwrap(); + + assert!(!keys.is_empty(), "Should have keys"); + assert!( + keys.contains(&"fc1.weight".to_string()), + "Keys should contain fc1.weight" + ); + assert!( + keys.contains(&"fc1.bias".to_string()), + "Keys should contain fc1.bias" + ); + } + + #[test] + fn test_keys_fast_path() { + let path = pytorch_test_path("linear", "linear.pt"); + + if !path.exists() { + println!("Skipping test - file not found: {:?}", path); + return; + } + + // Create fresh store - cache should be empty + let mut store = PytorchStore::from_file(&path); + + // keys() should work without populating the full cache (fast path) + let keys = store.keys().unwrap(); + assert!(!keys.is_empty(), "Should have keys via fast path"); + + // Now call get_all_snapshots to populate cache + let snapshots = store.get_all_snapshots().unwrap(); + assert!(!snapshots.is_empty(), "Should have snapshots"); + + // keys() should now use the cached data + let keys2 = store.keys().unwrap(); + assert_eq!(keys.len(), keys2.len(), "Keys count should match"); + } + + #[test] + fn test_caching_behavior() { + let path = pytorch_test_path("linear", "linear.pt"); + + if !path.exists() { + println!("Skipping test - file not found: {:?}", path); + return; + } + + let mut store = PytorchStore::from_file(path); + + // First call populates cache + let snapshots1 = store.get_all_snapshots().unwrap(); + let count1 = snapshots1.len(); + + // Second call uses cache + let snapshots2 = store.get_all_snapshots().unwrap(); + let count2 = snapshots2.len(); + + assert_eq!(count1, count2, "Cached results should match"); + } + + #[test] + fn test_get_all_snapshots_with_remapping() { + let path = pytorch_test_path("linear", "linear.pt"); + + if !path.exists() { + println!("Skipping test - file not found: {:?}", path); + return; + } + + // Create store with key remapping + let mut store = PytorchStore::from_file(path).with_key_remapping(r"^fc1\.", "linear1."); + + let snapshots = store.get_all_snapshots().unwrap(); + + // Should have remapped keys + assert!( + snapshots.contains_key("linear1.weight"), + "Should contain remapped key linear1.weight. Keys: {:?}", + snapshots.keys().collect::>() + ); + assert!( + snapshots.contains_key("linear1.bias"), + "Should contain remapped key linear1.bias" + ); + + // Original keys should not exist + assert!( + !snapshots.contains_key("fc1.weight"), + "Should not contain original key fc1.weight" + ); + } + + #[test] + fn test_get_snapshot_with_remapped_name() { + let path = pytorch_test_path("linear", "linear.pt"); + + if !path.exists() { + println!("Skipping test - file not found: {:?}", path); + return; + } + + // Create store with key remapping + let mut store = PytorchStore::from_file(path).with_key_remapping(r"^fc1\.", "linear1."); + + // Should find by remapped name + let snapshot = store.get_snapshot("linear1.weight").unwrap(); + assert!(snapshot.is_some(), "Should find tensor by remapped name"); + + // Should NOT find by original name + let snapshot_orig = store.get_snapshot("fc1.weight").unwrap(); + assert!( + snapshot_orig.is_none(), + "Should not find tensor by original name after remapping" + ); + } + + #[test] + fn test_get_all_snapshots_ignores_filter() { + let path = pytorch_test_path("linear", "linear.pt"); + + if !path.exists() { + println!("Skipping test - file not found: {:?}", path); + return; + } + + // Create store with filter that only matches fc1 + let mut store = PytorchStore::from_file(path).with_regex(r"^fc1\."); + + // get_all_snapshots should return ALL tensors regardless of filter + let snapshots = store.get_all_snapshots().unwrap(); + + // Should have both fc1 and fc2 tensors + assert!( + snapshots.contains_key("fc1.weight"), + "Should contain fc1.weight" + ); + assert!( + snapshots.contains_key("fc2.weight"), + "Should contain fc2.weight (filter not applied to get_all_snapshots)" + ); + } +} + +/// Tests for contiguous index mapping feature +#[cfg(test)] +mod map_indices_contiguous_tests { + use burn_core::tensor::Device; + + use super::*; + + /// Model with a Vec of Conv2d layers that expects contiguous indices + #[derive(Module, Debug)] + struct SequentialConvModel { + fc: Vec, + } + + impl SequentialConvModel { + pub fn new(device: &Device, num_layers: usize) -> Self { + Self { + fc: (0..num_layers) + .map(|_| { + Conv2dConfig::new([2, 2], [3, 3]) + .with_bias(true) + .init(device) + }) + .collect(), + } + } + } + + #[test] + fn test_load_non_contiguous_indexes_with_mapping() { + // This test uses the non_contiguous_indexes.pt file which has: + // fc.0.weight, fc.0.bias, fc.2.weight, fc.2.bias, fc.4.weight, ... (non-contiguous) + // The Burn model expects fc.0, fc.1, fc.2, ... (contiguous) + + let path = pytorch_test_path("non_contiguous_indexes", "non_contiguous_indexes.pt"); + + if !path.exists() { + println!("Skipping test - file not found: {:?}", path); + return; + } + + let device = Default::default(); + + // Create model with 5 conv layers (matching the PyTorch model) + let mut model = SequentialConvModel::new(&device, 5); + + // Load with contiguous index mapping enabled (default) + let mut store = PytorchStore::from_file(&path) + .map_indices_contiguous(true) + .allow_partial(true) + .validate(false); + + let result = store.apply_to(&mut model); + + match result { + Ok(apply_result) => { + println!("Applied tensors: {:?}", apply_result.applied); + println!("Missing tensors: {:?}", apply_result.missing); + println!("Unused tensors: {:?}", apply_result.unused); + + // All fc layers should be loaded successfully + assert!( + !apply_result.applied.is_empty(), + "Should have applied tensors" + ); + + // Verify we have tensors from all 5 layers + // With mapping: fc.0, fc.1, fc.2, fc.3, fc.4 + for i in 0..5 { + let has_weight = apply_result + .applied + .iter() + .any(|p| p.contains(&format!("fc.{}.weight", i))); + let has_bias = apply_result + .applied + .iter() + .any(|p| p.contains(&format!("fc.{}.bias", i))); + + assert!( + has_weight, + "Should have applied fc.{}.weight, applied: {:?}", + i, apply_result.applied + ); + assert!( + has_bias, + "Should have applied fc.{}.bias, applied: {:?}", + i, apply_result.applied + ); + } + + // There should be no missing tensors (assuming model matches) + let missing_fc: Vec<_> = apply_result + .missing + .iter() + .filter(|(p, _)| p.starts_with("fc.")) + .collect(); + assert!( + missing_fc.is_empty(), + "Should have no missing fc tensors with index mapping. Missing: {:?}", + missing_fc + ); + } + Err(e) => panic!("Failed to load with index mapping: {}", e), + } + } + + #[test] + fn test_load_non_contiguous_indexes_without_mapping() { + // This test verifies that loading fails or has missing tensors when + // map_indices_contiguous is disabled + + let path = pytorch_test_path("non_contiguous_indexes", "non_contiguous_indexes.pt"); + + if !path.exists() { + println!("Skipping test - file not found: {:?}", path); + return; + } + + let device = Default::default(); + + // Create model with 5 conv layers + let mut model = SequentialConvModel::new(&device, 5); + + // Load with contiguous index mapping DISABLED + let mut store = PytorchStore::from_file(&path) + .map_indices_contiguous(false) // Disable index mapping + .allow_partial(true) + .validate(false); + + let result = store.apply_to(&mut model); + + match result { + Ok(apply_result) => { + println!( + "Without index mapping - Applied tensors: {:?}", + apply_result.applied + ); + println!( + "Without index mapping - Missing tensors: {:?}", + apply_result.missing + ); + + // Without index mapping, we should have missing tensors for fc.1, fc.3 + // because the source has fc.0, fc.2, fc.4, fc.6, fc.8 but model expects fc.0-4 + let missing_fc: Vec<_> = apply_result + .missing + .iter() + .filter(|(p, _)| p.starts_with("fc.")) + .collect(); + + assert!( + !missing_fc.is_empty(), + "Should have missing fc tensors without index mapping (indices 1, 3 don't exist in file)" + ); + + // Specifically, fc.1 and fc.3 should be missing + let has_fc1_missing = apply_result + .missing + .iter() + .any(|(p, _)| p.starts_with("fc.1.")); + let has_fc3_missing = apply_result + .missing + .iter() + .any(|(p, _)| p.starts_with("fc.3.")); + + assert!( + has_fc1_missing || has_fc3_missing, + "Should have fc.1 or fc.3 missing. Missing: {:?}", + apply_result.missing + ); + } + Err(e) => panic!("Unexpected error: {}", e), + } + } + + #[test] + fn test_mapping_applied_to_keys() { + // Verify that the keys returned by the store are mapped + let path = pytorch_test_path("non_contiguous_indexes", "non_contiguous_indexes.pt"); + + if !path.exists() { + println!("Skipping test - file not found: {:?}", path); + return; + } + + // With index mapping enabled (default) + let mut store_mapped = PytorchStore::from_file(&path).map_indices_contiguous(true); + + let keys_mapped = store_mapped.keys().unwrap(); + println!("Keys with index mapping: {:?}", keys_mapped); + + // Should have contiguous keys: fc.0, fc.1, fc.2, fc.3, fc.4 + assert!( + keys_mapped.iter().any(|k| k.starts_with("fc.1.")), + "With index mapping, should have fc.1 (from fc.2)" + ); + assert!( + keys_mapped.iter().any(|k| k.starts_with("fc.2.")), + "With index mapping, should have fc.2 (from fc.4)" + ); + + // Without index mapping + let mut store_no_mapping = PytorchStore::from_file(&path).map_indices_contiguous(false); + + let keys_no_mapping = store_no_mapping.keys().unwrap(); + println!("Keys without index mapping: {:?}", keys_no_mapping); + + // Should have original non-contiguous keys: fc.0, fc.2, fc.4, fc.6, fc.8 + assert!( + keys_no_mapping.iter().any(|k| k.starts_with("fc.2.")), + "Without index mapping, should have original fc.2" + ); + assert!( + keys_no_mapping.iter().any(|k| k.starts_with("fc.4.")), + "Without index mapping, should have original fc.4" + ); + assert!( + !keys_no_mapping.iter().any(|k| k.starts_with("fc.1.")), + "Without index mapping, should NOT have fc.1 (not in original file)" + ); + } +} diff --git a/crates/burn-store/src/pytorch/tests/store/test_data/generate_enum_test.py b/crates/burn-store/src/pytorch/tests/store/test_data/generate_enum_test.py new file mode 100644 index 0000000..f735153 --- /dev/null +++ b/crates/burn-store/src/pytorch/tests/store/test_data/generate_enum_test.py @@ -0,0 +1,81 @@ +#!/usr/bin/env python3 +""" +Generate PyTorch test data for enum variant path mismatch testing. + +This script creates a PyTorch checkpoint that simulates how PyTorch models +export their state dicts WITHOUT enum variant names in the paths. + +Example: +- PyTorch path: "feature.weight" +- Burn path: "feature.BaseConv.weight" (includes enum variant "BaseConv") + +Run with: uv run generate_enum_test.py +""" + +import torch +import torch.nn as nn + + +class SimpleModel(nn.Module): + """ + Simple PyTorch model that represents what a Burn enum model would look like + WITHOUT the enum variant names in the path. + + In Burn, this would be: + struct ModelWithEnum { + feature: ConvBlock, // enum with BaseConv, DwsConv variants + classifier: Linear, + } + + But PyTorch exports it as flat paths without the enum variant names. + """ + def __init__(self): + super().__init__() + # This represents the "feature" field which is an enum in Burn + # PyTorch doesn't have enums, so it's just a Linear layer + # Path will be: "feature.weight" and "feature.bias" + self.feature = nn.Linear(3, 64) + + # This represents the "classifier" field + # Path will be: "classifier.weight" and "classifier.bias" + self.classifier = nn.Linear(64, 10) + + def forward(self, x): + x = self.feature(x) + x = torch.relu(x) + x = self.classifier(x) + return x + + +def generate_enum_variant_mismatch_test(): + """Generate test file demonstrating enum variant path mismatch.""" + model = SimpleModel() + + # Initialize with some deterministic weights for testing + torch.manual_seed(42) + for param in model.parameters(): + param.data.normal_(0, 0.1) + + # Save the state dict + # PyTorch paths: "feature.weight", "feature.bias", "classifier.weight", "classifier.bias" + # Burn paths: "feature.BaseConv.weight", "feature.BaseConv.bias", ... + # ^^^^^^^^ enum variant is missing in PyTorch + torch.save(model.state_dict(), "model_without_enum_variants.pt") + + print("Generated: model_without_enum_variants.pt") + print("\nPyTorch state dict keys:") + for key in model.state_dict().keys(): + shape = tuple(model.state_dict()[key].shape) + print(f" {key}: {shape}") + + print("\nExpected Burn paths (with enum variant):") + print(" feature.BaseConv.weight: (3, 64)") + print(" feature.BaseConv.bias: (64,)") + print(" classifier.weight: (64, 10)") + print(" classifier.bias: (10,)") + + print("\n⚠️ Notice: Burn includes 'BaseConv' enum variant, PyTorch doesn't!") + + +if __name__ == "__main__": + generate_enum_variant_mismatch_test() diff --git a/crates/burn-store/src/pytorch/tests/store/test_data/model_without_enum_variants.pt b/crates/burn-store/src/pytorch/tests/store/test_data/model_without_enum_variants.pt new file mode 100644 index 0000000..1846aa2 Binary files /dev/null and b/crates/burn-store/src/pytorch/tests/store/test_data/model_without_enum_variants.pt differ diff --git a/crates/burn-store/src/safetensors/mod.rs b/crates/burn-store/src/safetensors/mod.rs new file mode 100644 index 0000000..b077ed1 --- /dev/null +++ b/crates/burn-store/src/safetensors/mod.rs @@ -0,0 +1,322 @@ +//! SafeTensors format support for Burn deep learning framework. +//! +//! [SafeTensors](https://github.com/huggingface/safetensors) is a simple, safe, and efficient format +//! for storing and loading tensors. It provides fast zero-copy deserialization and strong safety +//! guarantees, making it ideal for production environments. +//! +//! # Features +//! +//! - **Fast Loading**: Zero-copy tensor access using safetensors' built-in mechanisms +//! - **Safety**: Prevents arbitrary code execution during model loading +//! - **Efficiency**: Memory-mapped files enable lazy loading without reading entire file +//! - **Filtering**: Load only specific tensors using path filters +//! - **Remapping**: Transform tensor names during load/save operations +//! - **Metadata**: Store and retrieve custom metadata alongside tensors (automatic `format`, `producer` and `version` metadata included) +//! - **Cross-Platform**: Works on all platforms including no-std environments +//! +//! # Usage Examples +//! +//! ## Basic Save and Load +//! +//! ```rust,ignore +//! use burn_store::{SafetensorsStore, ModuleSnapshot}; +//! +//! // Save a model to a file +//! let mut store = SafetensorsStore::from_file("model.safetensors"); +//! model.save_into(&mut store)?; +//! +//! // Load a model from a file +//! let mut store = SafetensorsStore::from_file("model.safetensors"); +//! let mut model = Model::new(&device); +//! model.load_from(&mut store)?; +//! ``` +//! +//! ## Memory-Based Operations +//! +//! ```rust,ignore +//! use burn_store::{SafetensorsStore, ModuleSnapshot}; +//! +//! // Save to memory buffer +//! let mut store = SafetensorsStore::from_bytes(None); +//! model.save_into(&mut store)?; +//! let bytes = store.get_bytes()?; +//! +//! // Load from memory buffer +//! let mut store = SafetensorsStore::from_bytes(Some(bytes)); +//! let mut model = Model::new(&device); +//! model.load_from(&mut store)?; +//! ``` +//! +//! ## Advanced Features +//! +//! ### Filter Configuration with Builder Pattern +//! +//! ```rust,no_run +//! # use burn_store::SafetensorsStore; +//! // Filter with regex patterns (OR logic - matches any pattern) +//! let mut store = SafetensorsStore::from_file("model.safetensors") +//! .with_regex(r"^encoder\..*") // Match all encoder tensors +//! .with_regex(r".*\.bias$"); // OR match any bias tensors +//! +//! // Filter with exact paths +//! let mut store = SafetensorsStore::from_file("model.safetensors") +//! .with_full_path("encoder.weight") +//! .with_full_path("encoder.bias") +//! .with_full_paths(vec!["decoder.scale", "decoder.norm"]); +//! +//! // Custom filter logic with predicate +//! let mut store = SafetensorsStore::from_file("model.safetensors") +//! .with_predicate(|path, _dtype| { +//! // Only save layer weights (not biases) +//! path.contains("layer") && path.ends_with("weight") +//! }); +//! +//! // Combine multiple filter methods +//! let mut store = SafetensorsStore::from_file("model.safetensors") +//! .with_regex(r"^encoder\..*") // All encoder tensors +//! .with_full_path("decoder.scale") // Plus specific decoder.scale +//! .with_predicate(|path, _| { // Plus any projection tensors +//! path.contains("projection") +//! }); +//! +//! // Save or load all tensors (no filtering) +//! let mut store = SafetensorsStore::from_file("model.safetensors") +//! .match_all(); +//! ``` +//! +//! ### Tensor Name Remapping +//! +//! Remap tensor names during load/save operations for compatibility between different frameworks: +//! +//! ```rust,no_run +//! # use burn_store::{SafetensorsStore, KeyRemapper}; +//! // Using builder pattern for common remapping patterns +//! let mut store = SafetensorsStore::from_file("model.safetensors") +//! .with_key_remapping(r"^encoder\.", "transformer.encoder.") // encoder.X -> transformer.encoder.X +//! .with_key_remapping(r"\.gamma$", ".weight") // X.gamma -> X.weight +//! .with_key_remapping(r"\.beta$", ".bias"); // X.beta -> X.bias +//! +//! // Or using a pre-configured KeyRemapper for complex transformations +//! let remapper = KeyRemapper::new() +//! .add_pattern(r"^pytorch\.(.*)", "burn.$1").expect("valid regex") // pytorch.layer -> burn.layer +//! .add_pattern(r"^(.*)\.running_mean$", "$1.mean").expect("valid regex") // layer.running_mean -> layer.mean +//! .add_pattern(r"^(.*)\.running_var$", "$1.variance").expect("valid regex"); // layer.running_var -> layer.variance +//! +//! let mut store = SafetensorsStore::from_file("model.safetensors") +//! .remap(remapper); +//! ``` +//! +//! ### Framework Adapters +//! +//! Use adapters for automatic framework-specific transformations: +//! +//! ```rust,ignore +//! use burn_store::{SafetensorsStore, ModuleSnapshot, PyTorchToBurnAdapter, BurnToPyTorchAdapter}; +//! +//! // Loading PyTorch model into Burn +//! let mut store = SafetensorsStore::from_file("pytorch_model.safetensors") +//! .with_from_adapter(PyTorchToBurnAdapter) // Transposes linear weights, renames norm params +//! .allow_partial(true); // PyTorch models may have extra tensors +//! +//! let mut burn_model = Model::new(&device); +//! burn_model.load_from(&mut store)?; +//! +//! // Saving Burn model for PyTorch +//! let mut store = SafetensorsStore::from_file("for_pytorch.safetensors") +//! .with_to_adapter(BurnToPyTorchAdapter); // Transposes weights back, renames for PyTorch +//! +//! burn_model.save_into(&mut store)?; +//! ``` +//! +//! ### Additional Configuration Options +//! +//! ```rust,ignore +//! use burn_store::{SafetensorsStore, ModuleSnapshot}; +//! +//! let mut store = SafetensorsStore::from_file("model.safetensors") +//! // Add custom metadata +//! .metadata("version", "1.0.0") +//! .metadata("producer", "burn") +//! // Allow partial loading (continue even if some tensors are missing) +//! .allow_partial(true) +//! // Disable validation for faster loading +//! .validate(false); +//! +//! // Use the configured store +//! model.save_into(&mut store)?; // For saving +//! // or +//! model.load_from(&mut store)?; // For loading +//! ``` +//! +//! # Efficient Loading with SafeTensors +//! +//! SafeTensors provides efficient tensor loading through its zero-copy design: +//! +//! ```rust,ignore +//! use burn_store::{SafetensorsStore, ModuleSnapshot}; +//! +//! let mut store = SafetensorsStore::from_file("large_model.safetensors"); +//! // Uses memory mapping (when available) for zero-copy access +//! // Falls back to buffered reading when mmap is not available +//! let mut model = Model::new(&device); +//! model.load_from(&mut store)?; +//! ``` +//! +//! The safetensors approach provides: +//! - Zero-copy views - tensors are accessed directly from the mapped file +//! - Lazy loading - only accessed tensors are materialized +//! - Efficient memory usage - no unnecessary data duplication +//! +//! # Lazy Loading and Inspection +//! +//! SafeTensors provides efficient inspection and selective loading through its +//! zero-copy design and built-in metadata handling: +//! +//! ```rust,ignore +//! use burn_store::SafetensorsStore; +//! +//! // Open a file - uses safetensors' efficient header reading +//! let store = SafetensorsStore::from_file("large_model.safetensors"); +//! +//! // List all tensor names from the metadata +//! let tensor_names = store.list_tensors()?; +//! println!("Model contains {} tensors", tensor_names.len()); +//! +//! // Get tensor metadata without loading tensor data +//! if let Some((shape, dtype)) = store.tensor_info("encoder.weight")? { +//! println!("Encoder weight shape: {:?}, dtype: {:?}", shape, dtype); +//! } +//! +//! // Selectively load tensors - safetensors handles efficient access +//! let encoder_tensors = store.load_tensors(&[ +//! "encoder.weight", +//! "encoder.bias", +//! "encoder.norm" +//! ])?; +//! +//! // Distributed loading: each worker loads only its assigned layers +//! // SafeTensors' zero-copy views ensure minimal memory usage +//! let worker_layers = match worker_id { +//! 0 => vec!["encoder.layer1", "encoder.layer2"], +//! 1 => vec!["encoder.layer3", "encoder.layer4"], +//! 2 => vec!["decoder.layer1", "decoder.layer2"], +//! _ => vec!["head.weight", "head.bias"], +//! }; +//! let worker_tensors = store.load_tensors(&worker_layers)?; +//! ``` +//! +//! # Builder Pattern API Reference +//! +//! The SafetensorsStore provides a fluent builder API for configuration: +//! +//! ## Filtering Methods +//! +//! - **`with_regex(pattern)`** - Add regex pattern to match tensor names (OR logic with multiple patterns) +//! - **`with_full_path(path)`** - Add exact tensor path to include +//! - **`with_full_paths(paths)`** - Add multiple exact tensor paths to include +//! - **`with_predicate(fn)`** - Add custom filter function `fn(&str, &str) -> bool` +//! - **`match_all()`** - Disable filtering, include all tensors +//! +//! ## Remapping Methods +//! +//! - **`with_key_remapping(from, to)`** - Add regex pattern to rename tensors +//! - **`remap(KeyRemapper)`** - Use a pre-configured KeyRemapper for complex transformations +//! +//! ## Adapter Methods +//! +//! - **`with_from_adapter(adapter)`** - Set adapter for loading (e.g., PyTorchToBurnAdapter) +//! - **`with_to_adapter(adapter)`** - Set adapter for saving (e.g., BurnToPyTorchAdapter) +//! +//! ## Configuration Methods +//! +//! - **`metadata(key, value)`** - Add custom metadata to saved files (in addition to automatic `format`, `producer` and `version`) +//! - **`allow_partial(bool)`** - Allow loading even if some tensors are missing +//! - **`validate(bool)`** - Enable/disable tensor validation during loading +//! +//! All methods return `Self` for chaining: +//! +//! ```rust,no_run +//! use burn_store::{SafetensorsStore, PyTorchToBurnAdapter}; +//! +//! let store = SafetensorsStore::from_file("model.safetensors") +//! .with_regex(r"^encoder\..*") +//! .with_key_remapping(r"\.gamma$", ".weight") +//! .with_from_adapter(PyTorchToBurnAdapter) +//! .allow_partial(true) +//! .metadata("version", "2.0"); +//! ``` +//! +//! # Working with Bytes +//! +//! For direct byte operations without files: +//! +//! ```rust,ignore +//! use burn_store::{SafetensorsStore, ModuleSnapshot}; +//! +//! // Save to bytes with filtering and remapping +//! let mut store = SafetensorsStore::from_bytes(None) +//! .with_regex(r"^encoder\..*") // Only save encoder tensors +//! .with_key_remapping(r"^encoder\.", "transformer.") // Rename encoder.X -> transformer.X +//! .metadata("subset", "encoder_only"); +//! model.save_into(&mut store)?; +//! let bytes = store.get_bytes()?; +//! +//! // Load from bytes (allow partial since we only saved encoder) +//! let mut store = SafetensorsStore::from_bytes(Some(bytes)) +//! .with_key_remapping(r"^transformer\.", "encoder.") // Rename back: transformer.X -> encoder.X +//! .allow_partial(true); +//! let mut model = Model::new(&device); +//! let result = model.load_from(&mut store)?; +//! println!("Applied {} tensors", result.applied.len()); +//! ``` +//! +//! # Complete Example: PyTorch Model Migration +//! +//! Migrating a PyTorch model to Burn with filtering, remapping, and adapters: +//! +//! ```rust,ignore +//! use burn_store::{SafetensorsStore, ModuleSnapshot, PyTorchToBurnAdapter}; +//! +//! // Load PyTorch model with all transformations +//! let mut store = SafetensorsStore::from_file("pytorch_model.safetensors") +//! // Use PyTorch adapter for automatic transformations +//! .with_from_adapter(PyTorchToBurnAdapter) +//! // Only load transformer layers +//! .with_regex(r"^transformer\..*") +//! // Rename old layer names to new structure +//! .with_key_remapping(r"^transformer\.h\.(\d+)\.", "transformer.layer$1.") +//! // Skip unexpected tensors from PyTorch +//! .allow_partial(true) +//! // Add metadata about the conversion +//! .metadata("source", "pytorch") +//! .metadata("converted_by", "burn-store"); +//! +//! let mut model = TransformerModel::new(&device); +//! let result = model.load_from(&mut store)?; +//! +//! println!("Successfully loaded {} tensors", result.applied.len()); +//! if !result.missing.is_empty() { +//! println!("Missing tensors: {:?}", result.missing); +//! } +//! ``` +//! +//! # Format Details +//! +//! SafeTensors uses a simple binary format: +//! - **8 bytes**: Header size (unsigned little-endian 64-bit integer) +//! - **N bytes**: JSON header with tensor metadata +//! - Contains: `{"tensor_name": {"dtype": "F32", "shape": [1, 2, 3], "data_offsets": [start, end]}, ...}` +//! - Special key `__metadata__` for user-defined string metadata +//! - **Rest**: Raw tensor data (referenced by offsets in header) +//! +//! The format enables: +//! - **Secure loading**: No code execution, just data +//! - **Efficient access**: Use offsets to read only needed tensors +//! - **Simple parsing**: Standard JSON header with fixed structure + +mod store; + +pub use store::{SafetensorsStore, SafetensorsStoreError}; + +#[cfg(test)] +mod tests; diff --git a/crates/burn-store/src/safetensors/store.rs b/crates/burn-store/src/safetensors/store.rs new file mode 100644 index 0000000..7b299c3 --- /dev/null +++ b/crates/burn-store/src/safetensors/store.rs @@ -0,0 +1,1087 @@ +//! SafeTensors store implementation using the official safetensors crate. + +use crate::{ + ApplyResult, IdentityAdapter, ModuleAdapter, ModuleSnapshot, ModuleStore, PathFilter, + TensorSnapshot, +}; + +#[cfg(feature = "std")] +use crate::{KeyRemapper, map_indices_contiguous}; +use alloc::boxed::Box; +use alloc::collections::BTreeMap; +use alloc::format; +use alloc::string::{String, ToString}; +use alloc::vec; +use alloc::vec::Vec; +use burn_core::module::ParamId; +use burn_core::tensor::{BoolStore, DType, TensorData}; +use core::fmt; +use core::ops::Deref; +use hashbrown::HashMap; + +// Arc is only available on targets with atomic pointers +#[cfg(target_has_atomic = "ptr")] +use alloc::sync::Arc; + +// For targets without atomic pointers, we use Box instead +#[cfg(not(target_has_atomic = "ptr"))] +type Arc = Box; + +/// Errors that can occur during SafeTensors operations. +#[derive(Debug)] +pub enum SafetensorsStoreError { + /// SafeTensors crate error. + Safetensors(safetensors::SafeTensorError), + + /// I/O error. + #[cfg(feature = "std")] + Io(std::io::Error), + + /// Tensor not found. + TensorNotFound(String), + + /// Validation failed. + ValidationFailed(String), + + /// Other error. + Other(String), +} + +impl fmt::Display for SafetensorsStoreError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Safetensors(e) => write!(f, "SafeTensors error: {}", e), + #[cfg(feature = "std")] + Self::Io(e) => write!(f, "I/O error: {}", e), + Self::TensorNotFound(name) => write!(f, "Tensor not found: {}", name), + Self::ValidationFailed(msg) => write!(f, "Validation failed: {}", msg), + Self::Other(msg) => write!(f, "{}", msg), + } + } +} + +impl core::error::Error for SafetensorsStoreError {} + +impl From for SafetensorsStoreError { + fn from(e: safetensors::SafeTensorError) -> Self { + SafetensorsStoreError::Safetensors(e) + } +} + +#[cfg(feature = "std")] +impl From for SafetensorsStoreError { + fn from(e: std::io::Error) -> Self { + SafetensorsStoreError::Io(e) + } +} + +/// SafeTensors store supporting both file and memory storage. +pub enum SafetensorsStore { + /// File-based storage. + #[cfg(feature = "std")] + File(FileStore), + + /// Memory-based storage. + Memory(MemoryStore), +} + +impl Default for SafetensorsStore { + /// Create a default memory-based store. + fn default() -> Self { + Self::from_bytes(None) + } +} + +impl SafetensorsStore { + /// Get the default metadata that includes Burn framework information. + /// + /// This includes: + /// - `format`: "safetensors" + /// - `producer`: "burn" + /// - `version`: The version of burn-store crate (from CARGO_PKG_VERSION) + /// + /// These metadata fields are automatically added to all saved models. + pub fn default_metadata() -> HashMap { + let mut metadata = HashMap::new(); + metadata.insert("format".to_string(), "safetensors".to_string()); + metadata.insert("producer".to_string(), "burn".to_string()); + metadata.insert("version".to_string(), env!("CARGO_PKG_VERSION").to_string()); + metadata + } + + /// Create a store for loading from or saving to a file. + #[cfg(feature = "std")] + pub fn from_file(path: impl Into) -> Self { + Self::File(FileStore { + path: path.into(), + filter: PathFilter::new(), + remapper: KeyRemapper::new(), + metadata: Self::default_metadata(), + validate: true, + allow_partial: false, + overwrite: false, + skip_enum_variants: false, + // Contiguous index mapping is off by default for SafeTensors + // (SafeTensors files typically have clean, contiguous indices) + map_indices_contiguous: false, + from_adapter: Box::new(IdentityAdapter), + to_adapter: Box::new(IdentityAdapter), + snapshots_cache: None, + }) + } + + /// Create a store for working with bytes in memory. + pub fn from_bytes(bytes: Option>) -> Self { + Self::Memory(MemoryStore { + data: bytes.map(Arc::new), + filter: PathFilter::new(), + #[cfg(feature = "std")] + remapper: KeyRemapper::new(), + metadata: Self::default_metadata(), + validate: true, + allow_partial: false, + skip_enum_variants: false, + // Contiguous index mapping is off by default for SafeTensors + #[cfg(feature = "std")] + map_indices_contiguous: false, + from_adapter: Box::new(IdentityAdapter), + to_adapter: Box::new(IdentityAdapter), + snapshots_cache: None, + }) + } + + /// Filter which tensors to load/save. + pub fn filter(mut self, filter: PathFilter) -> Self { + match &mut self { + #[cfg(feature = "std")] + Self::File(p) => p.filter = filter, + Self::Memory(p) => p.filter = filter, + } + self + } + + /// Add a regex pattern to filter tensors. + /// + /// Multiple patterns can be added and they work with OR logic. + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::SafetensorsStore; + /// let store = SafetensorsStore::from_file("model.safetensors") + /// .with_regex(r"^encoder\..*") // Match all encoder tensors + /// .with_regex(r".*\.weight$"); // OR match any weight tensors + /// ``` + #[cfg(feature = "std")] + pub fn with_regex>(mut self, pattern: S) -> Self { + match &mut self { + #[cfg(feature = "std")] + Self::File(p) => p.filter = p.filter.clone().with_regex(pattern), + Self::Memory(p) => p.filter = p.filter.clone().with_regex(pattern), + } + self + } + + /// Add multiple regex patterns to filter tensors. + #[cfg(feature = "std")] + pub fn with_regexes(mut self, patterns: I) -> Self + where + I: IntoIterator, + S: AsRef, + { + match &mut self { + #[cfg(feature = "std")] + Self::File(p) => p.filter = p.filter.clone().with_regexes(patterns), + Self::Memory(p) => p.filter = p.filter.clone().with_regexes(patterns), + } + self + } + + /// Add an exact full path to match. + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::SafetensorsStore; + /// let store = SafetensorsStore::from_file("model.safetensors") + /// .with_full_path("encoder.layer1.weight") + /// .with_full_path("decoder.output.bias"); + /// ``` + pub fn with_full_path>(mut self, path: S) -> Self { + match &mut self { + #[cfg(feature = "std")] + Self::File(p) => p.filter = p.filter.clone().with_full_path(path), + Self::Memory(p) => p.filter = p.filter.clone().with_full_path(path), + } + self + } + + /// Add multiple exact full paths to match. + pub fn with_full_paths(mut self, paths: I) -> Self + where + I: IntoIterator, + S: Into, + { + match &mut self { + #[cfg(feature = "std")] + Self::File(p) => p.filter = p.filter.clone().with_full_paths(paths), + Self::Memory(p) => p.filter = p.filter.clone().with_full_paths(paths), + } + self + } + + /// Add a predicate function for custom filtering logic. + /// + /// The predicate receives the tensor path and container path. + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::SafetensorsStore; + /// let store = SafetensorsStore::from_file("model.safetensors") + /// .with_predicate(|path, _| path.starts_with("encoder.") || path.ends_with(".bias")); + /// ``` + pub fn with_predicate(mut self, predicate: fn(&str, &str) -> bool) -> Self { + match &mut self { + #[cfg(feature = "std")] + Self::File(p) => p.filter = p.filter.clone().with_predicate(predicate), + Self::Memory(p) => p.filter = p.filter.clone().with_predicate(predicate), + } + self + } + + /// Add multiple predicate functions. + pub fn with_predicates(mut self, predicates: I) -> Self + where + I: IntoIterator bool>, + { + match &mut self { + #[cfg(feature = "std")] + Self::File(p) => p.filter = p.filter.clone().with_predicates(predicates), + Self::Memory(p) => p.filter = p.filter.clone().with_predicates(predicates), + } + self + } + + /// Set the filter to match all paths (disables filtering). + pub fn match_all(mut self) -> Self { + match &mut self { + #[cfg(feature = "std")] + Self::File(p) => p.filter = p.filter.clone().match_all(), + Self::Memory(p) => p.filter = p.filter.clone().match_all(), + } + self + } + + /// Remap tensor names during load/save. + #[cfg(feature = "std")] + pub fn remap(mut self, remapper: KeyRemapper) -> Self { + match &mut self { + Self::File(p) => p.remapper = remapper, + Self::Memory(p) => p.remapper = remapper, + } + self + } + + /// Add a regex pattern to remap tensor names during load/save. + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::SafetensorsStore; + /// let store = SafetensorsStore::from_file("model.safetensors") + /// .with_key_remapping(r"^encoder\.", "transformer.encoder.") // encoder.X -> transformer.encoder.X + /// .with_key_remapping(r"\.gamma$", ".weight"); // X.gamma -> X.weight + /// ``` + #[cfg(feature = "std")] + pub fn with_key_remapping( + mut self, + from_pattern: impl AsRef, + to_pattern: impl Into, + ) -> Self { + match &mut self { + Self::File(p) => { + p.remapper = p + .remapper + .clone() + .add_pattern(from_pattern, to_pattern) + .expect("Invalid regex pattern"); + } + Self::Memory(p) => { + p.remapper = p + .remapper + .clone() + .add_pattern(from_pattern, to_pattern) + .expect("Invalid regex pattern"); + } + } + self + } + + /// Add metadata to be saved with the tensors. + pub fn metadata(mut self, key: impl Into, value: impl Into) -> Self { + let key = key.into(); + let value = value.into(); + match &mut self { + #[cfg(feature = "std")] + Self::File(p) => { + p.metadata.insert(key, value); + } + Self::Memory(p) => { + p.metadata.insert(key, value); + } + } + self + } + + /// Clear all metadata including the default Burn framework metadata. + /// + /// This removes the automatic `format`, `producer` and `version` fields. + /// Use this when you need complete control over metadata or when + /// saving models for use with other frameworks. + pub fn clear_metadata(mut self) -> Self { + match &mut self { + #[cfg(feature = "std")] + Self::File(p) => { + p.metadata.clear(); + } + Self::Memory(p) => { + p.metadata.clear(); + } + } + self + } + + /// Set whether to validate tensors during loading (default: true). + pub fn validate(mut self, validate: bool) -> Self { + match &mut self { + #[cfg(feature = "std")] + Self::File(p) => p.validate = validate, + Self::Memory(p) => p.validate = validate, + } + self + } + + /// Allow partial loading of tensors (continue even if some tensors are missing). + pub fn allow_partial(mut self, allow: bool) -> Self { + match &mut self { + #[cfg(feature = "std")] + Self::File(p) => p.allow_partial = allow, + Self::Memory(p) => p.allow_partial = allow, + } + self + } + + /// Skip enum variant names when loading or saving tensor paths. + /// + /// When enabled during **loading**, tensor paths from the source that don't include enum variants + /// can be matched against Burn module paths that do include them. + /// For example, source path "feature.weight" can match Burn path "feature.BaseConv.weight". + /// + /// When enabled during **saving**, enum variant names are omitted from the exported tensor paths, + /// making them compatible with PyTorch naming conventions. + /// For example, "feature.BaseConv.weight" becomes "feature.weight" in the exported file. + /// + /// This is useful when working with models from/to formats that don't include enum variant + /// names in their parameter paths (like PyTorch models). + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::SafetensorsStore; + /// // For PyTorch compatibility + /// let store = SafetensorsStore::from_file("model.safetensors") + /// .skip_enum_variants(true); + /// ``` + pub fn skip_enum_variants(mut self, skip: bool) -> Self { + match &mut self { + #[cfg(feature = "std")] + Self::File(p) => p.skip_enum_variants = skip, + Self::Memory(p) => p.skip_enum_variants = skip, + } + self + } + + /// Enable or disable automatic contiguous mapping of layer indices (default: false). + /// + /// When enabled, non-contiguous numeric indices in tensor paths are renumbered + /// to be contiguous. This is useful when loading models that have gaps + /// in layer numbering, such as PyTorch models using `nn.Sequential` with mixed + /// layer types (e.g., Conv2d layers at indices 0, 2, 4 with ReLU layers at 1, 3, 5). + /// + /// # Example + /// + /// With index mapping enabled: + /// - `fc.0.weight` → `fc.0.weight` + /// - `fc.2.weight` → `fc.1.weight` (gap filled) + /// - `fc.4.weight` → `fc.2.weight` (gap filled) + /// + /// # Arguments + /// + /// * `map` - `true` to enable contiguous index mapping, `false` to disable + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::SafetensorsStore; + /// // Enable contiguous index mapping for PyTorch-exported safetensors + /// let store = SafetensorsStore::from_file("model.safetensors") + /// .map_indices_contiguous(true); + /// ``` + #[cfg(feature = "std")] + pub fn map_indices_contiguous(mut self, map: bool) -> Self { + match &mut self { + Self::File(p) => p.map_indices_contiguous = map, + Self::Memory(p) => p.map_indices_contiguous = map, + } + self + } + + /// Set whether to overwrite existing files when saving (default: false). + /// + /// When set to `false`, attempting to save to an existing file will result in an error. + /// When set to `true`, existing files will be overwritten without warning. + /// + /// This setting only applies to file-based stores. + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::SafetensorsStore; + /// let mut store = SafetensorsStore::from_file("model.safetensors") + /// .overwrite(true); + /// // Will overwrite if file exists when saving + /// ``` + #[cfg(feature = "std")] + pub fn overwrite(mut self, overwrite: bool) -> Self { + match &mut self { + Self::File(p) => p.overwrite = overwrite, + Self::Memory(_) => { + // Memory stores don't have overwrite semantics, ignore + } + } + self + } + + /// Set the adapter for loading tensors (converting from source format to Burn). + pub fn with_from_adapter(mut self, adapter: impl ModuleAdapter + 'static) -> Self { + match &mut self { + #[cfg(feature = "std")] + Self::File(p) => p.from_adapter = Box::new(adapter), + Self::Memory(p) => p.from_adapter = Box::new(adapter), + } + self + } + + /// Set the adapter for saving tensors (converting from Burn to target format). + pub fn with_to_adapter(mut self, adapter: impl ModuleAdapter + 'static) -> Self { + match &mut self { + #[cfg(feature = "std")] + Self::File(p) => p.to_adapter = Box::new(adapter), + Self::Memory(p) => p.to_adapter = Box::new(adapter), + } + self + } + + /// Get saved bytes from memory-based store. + /// + /// # Example + /// ```rust,no_run + /// # use burn_store::SafetensorsStore; + /// # fn example() -> Result<(), Box> { + /// let mut store = SafetensorsStore::from_bytes(None); + /// // After saving model with collect_to()... + /// let bytes = store.get_bytes()?; + /// # Ok(()) + /// # } + /// ``` + pub fn get_bytes(&self) -> Result, SafetensorsStoreError> { + match self { + #[cfg(feature = "std")] + Self::File(_) => Err(SafetensorsStoreError::Other( + "Cannot get bytes from file-based store".to_string(), + )), + Self::Memory(p) => p + .data() + .map(|arc| arc.as_ref().clone()) + .ok_or_else(|| SafetensorsStoreError::Other("No data available".to_string())), + } + } +} + +/// File-based store. +#[cfg(feature = "std")] +pub struct FileStore { + path: std::path::PathBuf, + filter: PathFilter, + remapper: KeyRemapper, + metadata: HashMap, + validate: bool, + allow_partial: bool, + overwrite: bool, + skip_enum_variants: bool, + /// Enable contiguous mapping of layer indices (default: false) + map_indices_contiguous: bool, + from_adapter: Box, + to_adapter: Box, + /// Cached tensor snapshots (parsed once, reused) + snapshots_cache: Option>, +} + +/// Memory-based store. +pub struct MemoryStore { + data: Option>>, + filter: PathFilter, + #[cfg(feature = "std")] + remapper: KeyRemapper, + metadata: HashMap, + validate: bool, + allow_partial: bool, + skip_enum_variants: bool, + /// Enable contiguous mapping of layer indices (default: false) + #[cfg(feature = "std")] + map_indices_contiguous: bool, + from_adapter: Box, + to_adapter: Box, + /// Cached tensor snapshots (parsed once, reused) + snapshots_cache: Option>, +} + +impl Default for MemoryStore { + fn default() -> Self { + Self { + data: None, + filter: PathFilter::new(), + #[cfg(feature = "std")] + remapper: KeyRemapper::new(), + metadata: HashMap::new(), + validate: true, + allow_partial: false, + skip_enum_variants: false, + #[cfg(feature = "std")] + map_indices_contiguous: false, + from_adapter: Box::new(IdentityAdapter), + to_adapter: Box::new(IdentityAdapter), + snapshots_cache: None, + } + } +} + +impl MemoryStore { + #[cfg(test)] + pub(crate) fn data(&self) -> Option>> { + self.data.clone() + } + + #[cfg(not(test))] + fn data(&self) -> Option>> { + self.data.clone() + } + + #[cfg(test)] + pub(crate) fn set_data(&mut self, data: Vec) { + self.data = Some(Arc::new(data)); + } +} + +// Adapter to use TensorSnapshot directly with safetensors +#[derive(Debug)] +struct TensorSnapshotAdapter(TensorSnapshot); + +impl safetensors::View for TensorSnapshotAdapter { + fn dtype(&self) -> safetensors::Dtype { + // Convert from burn dtype to safetensors dtype + dtype_to_safetensors(self.0.dtype).unwrap_or(safetensors::Dtype::F32) + } + + fn shape(&self) -> &[usize] { + &self.0.shape + } + + fn data(&self) -> alloc::borrow::Cow<'_, [u8]> { + // Only materialize data when actually needed for serialization + let data = self + .0 + .to_data() + .unwrap_or_else(|e| panic!("Failed to get tensor data: {:?}", e)); + alloc::borrow::Cow::Owned(data.bytes.deref().to_vec()) + } + + fn data_len(&self) -> usize { + // Use the efficient data_len method from TensorSnapshot + self.0.data_len() + } +} + +impl ModuleStore for SafetensorsStore { + type Error = SafetensorsStoreError; + + fn collect_from(&mut self, module: &M) -> Result<(), Self::Error> { + // Invalidate cache since we're writing new data + match self { + #[cfg(feature = "std")] + Self::File(p) => p.snapshots_cache = None, + Self::Memory(p) => p.snapshots_cache = None, + } + + // Collect tensor snapshots from module with adapter + // The to_adapter converts from Burn format to target format for saving + let to_adapter = match self { + #[cfg(feature = "std")] + Self::File(p) => p.to_adapter.clone(), + Self::Memory(p) => p.to_adapter.clone(), + }; + let mut snapshots = module.collect(None, Some(to_adapter), self.get_skip_enum_variants()); + + // Apply filtering + snapshots = apply_filter(snapshots, self.get_filter()); + + // Apply remapping + #[cfg(feature = "std")] + { + snapshots = apply_remapping(snapshots, self.get_remapper()); + } + + // Get metadata (already includes format, producer and version from default_metadata) + let metadata = self.get_metadata().clone(); + + #[cfg(feature = "std")] + let std_metadata: std::collections::HashMap = metadata + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + + // Write to storage + match self { + #[cfg(feature = "std")] + Self::File(p) => { + // Check if file exists and overwrite is disabled + if p.path.exists() && !p.overwrite { + return Err(SafetensorsStoreError::Other(format!( + "File already exists: {}. Use .overwrite(true) to overwrite.", + p.path.display() + ))); + } + + // Convert to safetensors format + let tensors = snapshots_to_safetensors(snapshots)?; + + // Use serialize_to_file which streams directly to disk + // This calls the lazy closures on-demand without buffering everything + safetensors::serialize_to_file(tensors, Some(std_metadata), &p.path)?; + Ok(()) + } + Self::Memory(p) => { + // For memory, we need to serialize to bytes + let tensors = snapshots_to_safetensors(snapshots)?; + // For no-std, serialize still needs std HashMap when std feature is enabled + #[cfg(feature = "std")] + let data = safetensors::serialize(tensors, Some(std_metadata))?; + + #[cfg(not(feature = "std"))] + let data = safetensors::serialize(tensors, Some(metadata))?; + p.data = Some(Arc::new(data)); + Ok(()) + } + } + } + + fn apply_to(&mut self, module: &mut M) -> Result { + // Get snapshots from cache + let snapshots: Vec = self.get_all_snapshots()?.values().cloned().collect(); + + // Get the adapter + let adapter: Box = match self { + #[cfg(feature = "std")] + Self::File(p) => p.from_adapter.clone(), + Self::Memory(p) => p.from_adapter.clone(), + }; + + // Get filter (cloned to Option for apply) + let filter = self.get_filter(); + let filter_opt = if filter.is_empty() { + None + } else { + Some(filter.clone()) + }; + + // Apply to module with adapter + // The adapter will be applied during module traversal with proper container info + // Filter is applied here during apply, not during cache population + let result = module.apply( + snapshots, + filter_opt, + Some(adapter), + self.get_skip_enum_variants(), + ); + + // Validate if needed + if self.get_validate() && !result.errors.is_empty() { + return Err(SafetensorsStoreError::ValidationFailed(format!( + "Import errors: {:?}", + result.errors + ))); + } + + if !self.get_allow_partial() && !result.missing.is_empty() { + return Err(SafetensorsStoreError::TensorNotFound(format!( + "\n{}\n\nHint: Use `.allow_partial(true)` on the store to get an `ApplyResult` \ + with structured missing/error info instead of a hard failure.", + result + ))); + } + + Ok(result) + } + + fn get_snapshot(&mut self, name: &str) -> Result, Self::Error> { + // Ensure cache is populated + self.ensure_snapshots_cache()?; + let cache = match self { + #[cfg(feature = "std")] + Self::File(p) => p.snapshots_cache.as_ref().unwrap(), + Self::Memory(p) => p.snapshots_cache.as_ref().unwrap(), + }; + Ok(cache.get(name)) + } + + fn get_all_snapshots(&mut self) -> Result<&BTreeMap, Self::Error> { + // Ensure cache is populated + self.ensure_snapshots_cache()?; + let cache = match self { + #[cfg(feature = "std")] + Self::File(p) => p.snapshots_cache.as_ref().unwrap(), + Self::Memory(p) => p.snapshots_cache.as_ref().unwrap(), + }; + Ok(cache) + } + + fn keys(&mut self) -> Result, Self::Error> { + // Always use the cache to ensure remapping is applied consistently + Ok(self.get_all_snapshots()?.keys().cloned().collect()) + } +} + +impl SafetensorsStore { + fn get_filter(&self) -> &PathFilter { + match self { + #[cfg(feature = "std")] + Self::File(p) => &p.filter, + Self::Memory(p) => &p.filter, + } + } + + #[cfg(feature = "std")] + fn get_remapper(&self) -> &KeyRemapper { + match self { + Self::File(p) => &p.remapper, + Self::Memory(p) => &p.remapper, + } + } + + fn get_metadata(&self) -> &HashMap { + match self { + #[cfg(feature = "std")] + Self::File(p) => &p.metadata, + Self::Memory(p) => &p.metadata, + } + } + + fn get_validate(&self) -> bool { + match self { + #[cfg(feature = "std")] + Self::File(p) => p.validate, + Self::Memory(p) => p.validate, + } + } + + fn get_allow_partial(&self) -> bool { + match self { + #[cfg(feature = "std")] + Self::File(p) => p.allow_partial, + Self::Memory(p) => p.allow_partial, + } + } + + fn get_skip_enum_variants(&self) -> bool { + match self { + #[cfg(feature = "std")] + Self::File(p) => p.skip_enum_variants, + Self::Memory(p) => p.skip_enum_variants, + } + } + + #[cfg(feature = "std")] + fn get_map_indices_contiguous(&self) -> bool { + match self { + Self::File(p) => p.map_indices_contiguous, + Self::Memory(p) => p.map_indices_contiguous, + } + } + + /// Ensure the snapshots cache is populated + fn ensure_snapshots_cache(&mut self) -> Result<(), SafetensorsStoreError> { + // Check if cache exists + let has_cache = match self { + #[cfg(feature = "std")] + Self::File(p) => p.snapshots_cache.is_some(), + Self::Memory(p) => p.snapshots_cache.is_some(), + }; + + if has_cache { + return Ok(()); + } + + // Load snapshots + #[allow(unused_mut)] + let mut snapshots = match self { + #[cfg(feature = "std")] + Self::File(p) => safetensors_to_snapshots_lazy_file(&p.path)?, + Self::Memory(p) => { + let data_arc = p + .data + .clone() + .ok_or_else(|| SafetensorsStoreError::Other("No data loaded".to_string()))?; + safetensors_to_snapshots_lazy(data_arc)? + } + }; + + // Apply remapping (but NOT filtering - that's done at apply time) + #[cfg(feature = "std")] + { + snapshots = match self { + Self::File(p) => apply_remapping(snapshots, &p.remapper), + Self::Memory(p) => apply_remapping(snapshots, &p.remapper), + }; + } + + // Apply contiguous index mapping if enabled + // This must be done after remapping so that remapped paths are mapped + #[cfg(feature = "std")] + if self.get_map_indices_contiguous() { + let (mapped, _) = map_indices_contiguous(snapshots); + snapshots = mapped; + } + + // Build cache as BTreeMap + let cache: BTreeMap = + snapshots.into_iter().map(|s| (s.full_path(), s)).collect(); + + // Store cache + match self { + #[cfg(feature = "std")] + Self::File(p) => p.snapshots_cache = Some(cache), + Self::Memory(p) => p.snapshots_cache = Some(cache), + } + + Ok(()) + } +} + +/// Apply filter to tensor snapshots. +fn apply_filter(mut snapshots: Vec, filter: &PathFilter) -> Vec { + if filter.is_empty() { + return snapshots; + } + + snapshots.retain(|snapshot| { + let path = snapshot.full_path(); + filter.matches(&path) + }); + + snapshots +} + +/// Apply remapping to tensor snapshots. +#[cfg(feature = "std")] +fn apply_remapping(snapshots: Vec, remapper: &KeyRemapper) -> Vec { + if remapper.is_empty() { + return snapshots; + } + + let (remapped, _) = remapper.remap(snapshots); + remapped +} + +/// Convert TensorSnapshots to safetensors format lazily. +fn snapshots_to_safetensors( + snapshots: Vec, +) -> Result, SafetensorsStoreError> { + let mut tensors = Vec::new(); + + for snapshot in snapshots { + let name = snapshot.full_path(); + // No need to materialize data - TensorSnapshot now has dtype and shape cached! + tensors.push((name, TensorSnapshotAdapter(snapshot))); + } + + Ok(tensors) +} + +/// Convert safetensors to TensorSnapshots with lazy loading. +fn safetensors_to_snapshots_lazy( + data_arc: Arc>, +) -> Result, SafetensorsStoreError> { + // Parse to get metadata + let tensors = safetensors::SafeTensors::deserialize(&data_arc)?; + let mut snapshots = Vec::new(); + + for (name, tensor_snapshot) in tensors.tensors() { + // Extract metadata without materializing data + let dtype = safetensor_dtype_to_burn(tensor_snapshot.dtype())?; + let shape = tensor_snapshot.shape(); + let path_parts: Vec = name.split('.').map(|s| s.to_string()).collect(); + + // Create a lazy closure that will deserialize only this tensor when needed + #[cfg(target_has_atomic = "ptr")] + let data_clone = Arc::clone(&data_arc); + #[cfg(not(target_has_atomic = "ptr"))] + let data_clone = data_arc.clone(); + let name_clone = name.to_string(); + let data_fn = alloc::rc::Rc::new(move || { + // Re-deserialize when needed (this is cheap, just parsing header) + let tensors = safetensors::SafeTensors::deserialize(&data_clone).map_err(|e| { + crate::TensorSnapshotError::IoError(format!( + "Failed to re-deserialize safetensors: {}", + e + )) + })?; + + // Find our specific tensor + let tensor = tensors.tensor(&name_clone).map_err(|e| { + crate::TensorSnapshotError::DataError(format!( + "Tensor '{}' not found: {}", + name_clone, e + )) + })?; + + // Now materialize just this tensor's data + let bytes = burn_core::tensor::Bytes::from_bytes_vec(tensor.data().to_vec()); + Ok(TensorData { + bytes, + shape: tensor.shape().into(), + dtype: safetensor_dtype_to_burn(tensor.dtype()) + .map_err(|_| crate::TensorSnapshotError::DataError("Invalid dtype".into()))?, + }) + }); + + let snapshot = TensorSnapshot::from_closure( + data_fn, + dtype, + shape.into(), + path_parts, + vec![], // Empty container_stack - will be filled during module traversal + ParamId::new(), + ); + snapshots.push(snapshot); + } + + Ok(snapshots) +} + +/// Convert safetensors to TensorSnapshots with true on-demand loading from file. +/// This reads only the header initially, then loads tensor data on demand. +#[cfg(feature = "std")] +fn safetensors_to_snapshots_lazy_file( + path: &std::path::Path, +) -> Result, SafetensorsStoreError> { + // Always use memory mapping for the most efficient access + use memmap2::MmapOptions; + + // Memory map the file for efficient access + let file = std::fs::File::open(path)?; + let mmap = unsafe { MmapOptions::new().map(&file)? }; + let mmap_arc = Arc::new(mmap); + + // Parse just to get metadata (safetensors won't copy data with mmap) + let tensors = safetensors::SafeTensors::deserialize(&mmap_arc)?; + let mut snapshots = Vec::new(); + + for (name, tensor_snapshot) in tensors.tensors() { + let dtype = safetensor_dtype_to_burn(tensor_snapshot.dtype())?; + let shape = tensor_snapshot.shape(); + let path_parts: Vec = name.split('.').map(|s| s.to_string()).collect(); + + // Create a lazy closure that accesses the mmap'd data + let mmap_clone = Arc::clone(&mmap_arc); + let name_clone = name.to_string(); + + let data_fn = alloc::rc::Rc::new(move || { + // Re-parse to get the tensor snapshot (this is cheap with mmap) + let tensors = safetensors::SafeTensors::deserialize(&mmap_clone).map_err(|e| { + crate::TensorSnapshotError::IoError(format!("Failed to deserialize: {}", e)) + })?; + let tensor = tensors.tensor(&name_clone).map_err(|e| { + crate::TensorSnapshotError::DataError(format!( + "Tensor '{}' not found: {}", + name_clone, e + )) + })?; + + // Only now do we actually copy the tensor data + Ok(TensorData { + bytes: burn_core::tensor::Bytes::from_bytes_vec(tensor.data().to_vec()), + shape: tensor.shape().into(), + dtype: safetensor_dtype_to_burn(tensor.dtype()) + .map_err(|_| crate::TensorSnapshotError::DataError("Invalid dtype".into()))?, + }) + }); + + let snapshot = TensorSnapshot::from_closure( + data_fn, + dtype, + shape.into(), + path_parts, + vec![], // Empty container_stack - will be filled during module traversal + ParamId::new(), + ); + snapshots.push(snapshot); + } + + Ok(snapshots) +} + +/// Helper to convert safetensors Dtype to burn DType. +fn safetensor_dtype_to_burn(dtype: safetensors::Dtype) -> Result { + use safetensors::Dtype; + + match dtype { + Dtype::F64 => Ok(DType::F64), + Dtype::F32 => Ok(DType::F32), + Dtype::F16 => Ok(DType::F16), + Dtype::BF16 => Ok(DType::BF16), + Dtype::I64 => Ok(DType::I64), + Dtype::I32 => Ok(DType::I32), + Dtype::I16 => Ok(DType::I16), + Dtype::I8 => Ok(DType::I8), + Dtype::U64 => Ok(DType::U64), + Dtype::U32 => Ok(DType::U32), + Dtype::U8 => Ok(DType::U8), + Dtype::BOOL => Ok(DType::Bool(BoolStore::Native)), + _ => Err(SafetensorsStoreError::Other(format!( + "Unsupported dtype: {:?}", + dtype + ))), + } +} + +/// Helper to convert DType to safetensors Dtype. +fn dtype_to_safetensors(dtype: DType) -> Result { + use safetensors::Dtype; + + match dtype { + DType::F64 => Ok(Dtype::F64), + DType::F32 | DType::Flex32 => Ok(Dtype::F32), // Flex32 is stored as F32 + DType::F16 => Ok(Dtype::F16), + DType::BF16 => Ok(Dtype::BF16), + DType::I64 => Ok(Dtype::I64), + DType::I32 => Ok(Dtype::I32), + DType::I16 => Ok(Dtype::I16), + DType::I8 => Ok(Dtype::I8), + DType::U64 => Ok(Dtype::U64), + DType::U32 => Ok(Dtype::U32), + DType::U16 => Err(SafetensorsStoreError::Other( + "U16 dtype not yet supported in safetensors".to_string(), + )), + DType::U8 => Ok(Dtype::U8), + DType::Bool(BoolStore::Native) => Ok(Dtype::BOOL), + DType::Bool(BoolStore::U32) => Ok(Dtype::U32), + DType::Bool(BoolStore::U8) => Ok(Dtype::U8), + DType::QFloat(_) => Err(SafetensorsStoreError::Other( + "Quantized tensors not yet supported in safetensors".to_string(), + )), + } +} diff --git a/crates/burn-store/src/safetensors/tests/adapter.rs b/crates/burn-store/src/safetensors/tests/adapter.rs new file mode 100644 index 0000000..5bd9d1f --- /dev/null +++ b/crates/burn-store/src/safetensors/tests/adapter.rs @@ -0,0 +1,346 @@ +use burn_core as burn; + +use crate::{ + BurnToPyTorchAdapter, ModuleSnapshot, ModuleStore, PyTorchToBurnAdapter, SafetensorsStore, +}; +use burn_core::module::{Module, Param}; +use burn_core::tensor::{Device, Tensor}; +use burn_nn::{Linear, LinearConfig}; + +#[derive(Module, Debug)] +struct TestModel { + linear: Linear, + norm_weight: Param>, + norm_bias: Param>, +} + +impl TestModel { + fn new(device: &Device) -> Self { + Self { + linear: LinearConfig::new(4, 2).with_bias(true).init(device), + norm_weight: Param::from_data([1.0, 1.0], device), + norm_bias: Param::from_data([0.0, 0.0], device), + } + } +} + +#[test] +fn pytorch_to_burn_adapter_linear_transpose() { + let device = Default::default(); + let model = TestModel::new(&device); + + // Save with BurnToPyTorch adapter (will transpose linear weights) + let mut save_store = SafetensorsStore::from_bytes(None).with_to_adapter(BurnToPyTorchAdapter); + model.save_into(&mut save_store).unwrap(); + + // Load with PyTorchToBurn adapter (will transpose back) + let mut load_store = SafetensorsStore::from_bytes(None).with_from_adapter(PyTorchToBurnAdapter); + if let SafetensorsStore::Memory(ref mut p) = load_store + && let SafetensorsStore::Memory(ref p_save) = save_store + { + p.set_data(p_save.data().unwrap().as_ref().clone()); + } + + let mut model2 = TestModel::new(&device); + let result = model2.load_from(&mut load_store).unwrap(); + + // Should successfully load all tensors + assert!(!result.applied.is_empty()); + + // Verify the linear weights are the same after round-trip + let weight1 = model.linear.weight.val().to_data(); + let weight2 = model2.linear.weight.val().to_data(); + + assert_eq!(weight1.shape, weight2.shape); + let data1 = weight1.to_vec::().unwrap(); + let data2 = weight2.to_vec::().unwrap(); + + for (a, b) in data1.iter().zip(data2.iter()) { + assert!( + (a - b).abs() < 1e-6, + "Weights differ after adapter round-trip" + ); + } +} + +#[test] +fn pytorch_to_burn_adapter_norm_rename() { + let device = Default::default(); + + // Create a model with norm-like naming + #[derive(Module, Debug)] + struct NormModel { + norm_gamma: Param>, + norm_beta: Param>, + } + + impl NormModel { + fn new(device: &Device) -> Self { + Self { + norm_gamma: Param::from_data([1.0, 2.0, 3.0], device), + norm_beta: Param::from_data([0.1, 0.2, 0.3], device), + } + } + } + + let model = NormModel::new(&device); + + // Save with BurnToPyTorch adapter (will rename gamma->weight, beta->bias) + let mut save_store = SafetensorsStore::from_bytes(None).with_to_adapter(BurnToPyTorchAdapter); + model.save_into(&mut save_store).unwrap(); + + // The saved data should have PyTorch naming convention + // We can't directly verify the internal names, but we can verify round-trip works + + // Load with PyTorchToBurn adapter (will rename weight->gamma, bias->beta) + let mut load_store = SafetensorsStore::from_bytes(None).with_from_adapter(PyTorchToBurnAdapter); + if let SafetensorsStore::Memory(ref mut p) = load_store + && let SafetensorsStore::Memory(ref p_save) = save_store + { + p.set_data(p_save.data().unwrap().as_ref().clone()); + } + + let mut model2 = NormModel::new(&device); + let result = model2.load_from(&mut load_store).unwrap(); + + // Should load successfully + assert!(!result.applied.is_empty()); + + // Verify data is preserved + let gamma1 = model.norm_gamma.val().to_data().to_vec::().unwrap(); + let gamma2 = model2.norm_gamma.val().to_data().to_vec::().unwrap(); + let beta1 = model.norm_beta.val().to_data().to_vec::().unwrap(); + let beta2 = model2.norm_beta.val().to_data().to_vec::().unwrap(); + + assert_eq!(gamma1, gamma2); + assert_eq!(beta1, beta2); +} + +#[test] +fn no_adapter_preserves_original() { + let device = Default::default(); + let model = TestModel::new(&device); + + // Save without adapter + let mut save_store = SafetensorsStore::from_bytes(None); + model.save_into(&mut save_store).unwrap(); + + // Load without adapter + let mut load_store = SafetensorsStore::from_bytes(None); + if let SafetensorsStore::Memory(ref mut p) = load_store + && let SafetensorsStore::Memory(ref p_save) = save_store + { + p.set_data(p_save.data().unwrap().as_ref().clone()); + } + + let mut model2 = TestModel::new(&device); + let result = model2.load_from(&mut load_store).unwrap(); + + assert!(result.is_success()); + assert!(!result.applied.is_empty()); + + // Verify data is exactly the same + let weight1 = model.linear.weight.val().to_data(); + let weight2 = model2.linear.weight.val().to_data(); + + assert_eq!(weight1.shape, weight2.shape); + assert_eq!( + weight1.to_vec::().unwrap(), + weight2.to_vec::().unwrap() + ); +} + +#[test] +#[cfg(all(feature = "std", target_has_atomic = "ptr"))] +fn adapter_with_pytorch_import() { + use crate::PyTorchToBurnAdapter; + + let device = Default::default(); + + // Reference the safetensors file from burn-store + let safetensors_path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/safetensors-tests/tests/multi_layer/multi_layer.safetensors" + ); + + // Simple test model that matches some of the PyTorch structure + #[derive(Module, Debug)] + struct SimpleNet { + fc1: Linear, + } + + impl SimpleNet { + fn new(device: &Device) -> Self { + Self { + fc1: LinearConfig::new(4 * 8 * 8, 16).init(device), + } + } + } + + // Load with PyTorchToBurn adapter + let mut store = SafetensorsStore::from_file(safetensors_path) + .with_from_adapter(PyTorchToBurnAdapter) + .validate(false) + .allow_partial(true); + + let mut model = SimpleNet::new(&device); + let result = model.load_from(&mut store).unwrap(); + + // Should load some tensors (fc1 if it exists in the file) + // This mainly tests that the adapter works with real PyTorch files + assert!(!result.applied.is_empty() || !result.missing.is_empty()); +} + +#[test] +fn half_precision_adapter_round_trip() { + use crate::HalfPrecisionAdapter; + use burn_core::tensor::DType; + + let device = Default::default(); + let model = TestModel::new(&device); + + // Save with HalfPrecisionAdapter (F32 -> F16) + let adapter = HalfPrecisionAdapter::new(); + let mut save_store = SafetensorsStore::from_bytes(None).with_to_adapter(adapter.clone()); + model.save_into(&mut save_store).unwrap(); + + // Verify Linear tensors are F16, raw params stay F32 (no recognized module type) + let save_bytes = match &save_store { + SafetensorsStore::Memory(p) => p.data().unwrap().as_ref().clone(), + _ => panic!("Expected memory store"), + }; + let mut inspect_store = SafetensorsStore::from_bytes(Some(save_bytes.clone())); + let snapshots = inspect_store.get_all_snapshots().unwrap(); + for (name, snapshot) in snapshots.iter() { + if name.starts_with("linear") { + assert_eq!( + snapshot.dtype, + DType::F16, + "Linear tensor '{}' should be F16", + name + ); + } else { + assert_eq!( + snapshot.dtype, + DType::F32, + "Raw param '{}' should stay F32", + name + ); + } + } + + // Load back with same adapter (F16 -> F32) + let mut load_store = SafetensorsStore::from_bytes(Some(save_bytes)).with_from_adapter(adapter); + + let mut model2 = TestModel::new(&device); + let result = model2.load_from(&mut load_store).unwrap(); + + assert!(!result.applied.is_empty()); + + // Verify values are close (F32 -> F16 -> F32 has rounding) + let w1 = model.linear.weight.val().to_data().to_vec::().unwrap(); + let w2 = model2 + .linear + .weight + .val() + .to_data() + .to_vec::() + .unwrap(); + for (a, b) in w1.iter().zip(w2.iter()) { + assert!( + (a - b).abs() < 0.01, + "Weight values differ too much after F16 round-trip: {} vs {}", + a, + b + ); + } +} + +#[test] +fn half_precision_adapter_without_module() { + use crate::HalfPrecisionAdapter; + use burn_core::tensor::DType; + use burn_nn::{LayerNorm, LayerNormConfig}; + + #[derive(Module, Debug)] + struct MixedModel { + linear: Linear, + norm: LayerNorm, + } + + let device = Default::default(); + let model = MixedModel { + linear: LinearConfig::new(4, 2).with_bias(true).init(&device), + norm: LayerNormConfig::new(2).init(&device), + }; + + // Save: exclude LayerNorm from half-precision conversion + let adapter = HalfPrecisionAdapter::new().without_module("LayerNorm"); + let mut save_store = SafetensorsStore::from_bytes(None).with_to_adapter(adapter); + model.save_into(&mut save_store).unwrap(); + + // Verify: Linear tensors are F16, LayerNorm tensors remain F32 + let save_bytes = match &save_store { + SafetensorsStore::Memory(p) => p.data().unwrap().as_ref().clone(), + _ => panic!("Expected memory store"), + }; + let mut inspect_store = SafetensorsStore::from_bytes(Some(save_bytes)); + let snapshots = inspect_store.get_all_snapshots().unwrap(); + for (name, snapshot) in snapshots { + if name.starts_with("linear") { + assert_eq!( + snapshot.dtype, + DType::F16, + "Linear tensor '{}' should be F16", + name + ); + } else if name.starts_with("norm") { + assert_eq!( + snapshot.dtype, + DType::F32, + "LayerNorm tensor '{}' should stay F32", + name + ); + } + } +} + +#[test] +fn half_precision_adapter_default_converts_layer_norm() { + use crate::HalfPrecisionAdapter; + use burn_core::tensor::DType; + use burn_nn::{LayerNorm, LayerNormConfig}; + + #[derive(Module, Debug)] + struct NormModel { + linear: Linear, + norm: LayerNorm, + } + + let device = Default::default(); + let model = NormModel { + linear: LinearConfig::new(4, 2).with_bias(true).init(&device), + norm: LayerNormConfig::new(2).init(&device), + }; + + // Default adapter converts LayerNorm + let adapter = HalfPrecisionAdapter::new(); + let mut save_store = SafetensorsStore::from_bytes(None).with_to_adapter(adapter); + model.save_into(&mut save_store).unwrap(); + + let save_bytes = match &save_store { + SafetensorsStore::Memory(p) => p.data().unwrap().as_ref().clone(), + _ => panic!("Expected memory store"), + }; + let mut inspect_store = SafetensorsStore::from_bytes(Some(save_bytes)); + let snapshots = inspect_store.get_all_snapshots().unwrap(); + for (name, snapshot) in snapshots { + assert_eq!( + snapshot.dtype, + DType::F16, + "All tensors should be F16 by default, but '{}' is {:?}", + name, + snapshot.dtype + ); + } +} diff --git a/crates/burn-store/src/safetensors/tests/direct_access.rs b/crates/burn-store/src/safetensors/tests/direct_access.rs new file mode 100644 index 0000000..3171dfc --- /dev/null +++ b/crates/burn-store/src/safetensors/tests/direct_access.rs @@ -0,0 +1,338 @@ +use burn_core as burn; + +use crate::{ModuleStore, SafetensorsStore}; +use burn_core::module::{Module, Param}; +use burn_core::tensor::{Device, Tensor, shape}; + +// Test module for direct access tests +#[derive(Module, Debug)] +struct DirectAccessTestModule { + weight: Param>, + bias: Param>, + nested: DirectAccessNestedModule, +} + +#[derive(Module, Debug)] +struct DirectAccessNestedModule { + gamma: Param>, + beta: Param>, +} + +impl DirectAccessTestModule { + fn new(device: &Device) -> Self { + Self { + weight: Param::from_data([[1.0, 2.0], [3.0, 4.0]], device), + bias: Param::from_data([0.1, 0.2], device), + nested: DirectAccessNestedModule { + gamma: Param::from_data([1.0, 2.0], device), + beta: Param::from_data([0.5, 0.5], device), + }, + } + } +} + +#[test] +fn test_memory_get_all_snapshots() { + let device = Default::default(); + let module = DirectAccessTestModule::new(&device); + + // Save module to memory + let mut save_store = SafetensorsStore::from_bytes(None); + save_store.collect_from(&module).unwrap(); + + // Get bytes and create load store + let bytes = save_store.get_bytes().unwrap(); + let mut load_store = SafetensorsStore::from_bytes(Some(bytes)); + + // Get all snapshots + let snapshots = load_store.get_all_snapshots().unwrap(); + + assert_eq!(snapshots.len(), 4); + assert!(snapshots.contains_key("weight")); + assert!(snapshots.contains_key("bias")); + assert!(snapshots.contains_key("nested.gamma")); + assert!(snapshots.contains_key("nested.beta")); +} + +#[test] +fn test_memory_get_snapshot_existing() { + let device = Default::default(); + let module = DirectAccessTestModule::new(&device); + + let mut save_store = SafetensorsStore::from_bytes(None); + save_store.collect_from(&module).unwrap(); + let bytes = save_store.get_bytes().unwrap(); + + let mut load_store = SafetensorsStore::from_bytes(Some(bytes)); + + // Get existing snapshot + let snapshot = load_store.get_snapshot("weight").unwrap(); + assert!(snapshot.is_some()); + + let snapshot = snapshot.unwrap(); + assert_eq!(snapshot.shape, shape![2, 2]); + + // Verify data + let data = snapshot.to_data().unwrap(); + let values: Vec = data.to_vec().unwrap(); + assert_eq!(values, vec![1.0, 2.0, 3.0, 4.0]); +} + +#[test] +fn test_memory_get_snapshot_nested() { + let device = Default::default(); + let module = DirectAccessTestModule::new(&device); + + let mut save_store = SafetensorsStore::from_bytes(None); + save_store.collect_from(&module).unwrap(); + let bytes = save_store.get_bytes().unwrap(); + + let mut load_store = SafetensorsStore::from_bytes(Some(bytes)); + + // Get nested snapshot + let snapshot = load_store.get_snapshot("nested.gamma").unwrap(); + assert!(snapshot.is_some()); + + let snapshot = snapshot.unwrap(); + let data = snapshot.to_data().unwrap(); + let values: Vec = data.to_vec().unwrap(); + assert_eq!(values, vec![1.0, 2.0]); +} + +#[test] +fn test_memory_get_snapshot_not_found() { + let device = Default::default(); + let module = DirectAccessTestModule::new(&device); + + let mut save_store = SafetensorsStore::from_bytes(None); + save_store.collect_from(&module).unwrap(); + let bytes = save_store.get_bytes().unwrap(); + + let mut load_store = SafetensorsStore::from_bytes(Some(bytes)); + + // Get non-existent snapshot + let snapshot = load_store.get_snapshot("nonexistent").unwrap(); + assert!(snapshot.is_none()); +} + +#[test] +fn test_memory_keys() { + let device = Default::default(); + let module = DirectAccessTestModule::new(&device); + + let mut save_store = SafetensorsStore::from_bytes(None); + save_store.collect_from(&module).unwrap(); + let bytes = save_store.get_bytes().unwrap(); + + let mut load_store = SafetensorsStore::from_bytes(Some(bytes)); + + let keys = load_store.keys().unwrap(); + assert_eq!(keys.len(), 4); + assert!(keys.contains(&"weight".to_string())); + assert!(keys.contains(&"bias".to_string())); + assert!(keys.contains(&"nested.gamma".to_string())); + assert!(keys.contains(&"nested.beta".to_string())); +} + +#[test] +fn test_memory_caching_behavior() { + let device = Default::default(); + let module = DirectAccessTestModule::new(&device); + + let mut save_store = SafetensorsStore::from_bytes(None); + save_store.collect_from(&module).unwrap(); + let bytes = save_store.get_bytes().unwrap(); + + let mut load_store = SafetensorsStore::from_bytes(Some(bytes)); + + // Call get_all_snapshots multiple times - should return same cached data + let snapshots1 = load_store.get_all_snapshots().unwrap(); + assert_eq!(snapshots1.len(), 4); + + let snapshots2 = load_store.get_all_snapshots().unwrap(); + assert_eq!(snapshots2.len(), 4); + + // Verify we can still access individual snapshots after caching + let snapshot = load_store.get_snapshot("bias").unwrap(); + assert!(snapshot.is_some()); +} + +// ============================================================================ +// Tests for FileStore variant +// ============================================================================ + +#[test] +#[cfg(feature = "std")] +fn test_file_get_all_snapshots() { + use tempfile::tempdir; + + let device = Default::default(); + let module = DirectAccessTestModule::new(&device); + + let temp_dir = tempdir().unwrap(); + let path = temp_dir.path().join("test_get_all_snapshots.safetensors"); + + let mut save_store = SafetensorsStore::from_file(&path); + save_store.collect_from(&module).unwrap(); + + let mut load_store = SafetensorsStore::from_file(&path); + let snapshots = load_store.get_all_snapshots().unwrap(); + + assert_eq!(snapshots.len(), 4); + assert!(snapshots.contains_key("weight")); + assert!(snapshots.contains_key("bias")); + assert!(snapshots.contains_key("nested.gamma")); + assert!(snapshots.contains_key("nested.beta")); +} + +#[test] +#[cfg(feature = "std")] +fn test_file_get_snapshot_existing() { + use tempfile::tempdir; + + let device = Default::default(); + let module = DirectAccessTestModule::new(&device); + + let temp_dir = tempdir().unwrap(); + let path = temp_dir.path().join("test_get_snapshot.safetensors"); + + let mut save_store = SafetensorsStore::from_file(&path); + save_store.collect_from(&module).unwrap(); + + let mut load_store = SafetensorsStore::from_file(&path); + + let snapshot = load_store.get_snapshot("weight").unwrap(); + assert!(snapshot.is_some()); + + let snapshot = snapshot.unwrap(); + assert_eq!(snapshot.shape, shape![2, 2]); + + let data = snapshot.to_data().unwrap(); + let values: Vec = data.to_vec().unwrap(); + assert_eq!(values, vec![1.0, 2.0, 3.0, 4.0]); +} + +#[test] +#[cfg(feature = "std")] +fn test_file_get_snapshot_not_found() { + use tempfile::tempdir; + + let device = Default::default(); + let module = DirectAccessTestModule::new(&device); + + let temp_dir = tempdir().unwrap(); + let path = temp_dir.path().join("test_not_found.safetensors"); + + let mut save_store = SafetensorsStore::from_file(&path); + save_store.collect_from(&module).unwrap(); + + let mut load_store = SafetensorsStore::from_file(&path); + + let snapshot = load_store.get_snapshot("nonexistent").unwrap(); + assert!(snapshot.is_none()); +} + +#[test] +#[cfg(feature = "std")] +fn test_file_keys() { + use tempfile::tempdir; + + let device = Default::default(); + let module = DirectAccessTestModule::new(&device); + + let temp_dir = tempdir().unwrap(); + let path = temp_dir.path().join("test_keys.safetensors"); + + let mut save_store = SafetensorsStore::from_file(&path); + save_store.collect_from(&module).unwrap(); + + let mut load_store = SafetensorsStore::from_file(&path); + + let keys = load_store.keys().unwrap(); + assert_eq!(keys.len(), 4); + assert!(keys.contains(&"weight".to_string())); + assert!(keys.contains(&"bias".to_string())); + assert!(keys.contains(&"nested.gamma".to_string())); + assert!(keys.contains(&"nested.beta".to_string())); +} + +#[test] +#[cfg(feature = "std")] +fn test_file_keys_fast_path() { + use tempfile::tempdir; + + let device = Default::default(); + let module = DirectAccessTestModule::new(&device); + + let temp_dir = tempdir().unwrap(); + let path = temp_dir.path().join("test_keys_fast.safetensors"); + + let mut save_store = SafetensorsStore::from_file(&path); + save_store.collect_from(&module).unwrap(); + + // Create fresh store - cache should be empty + let mut load_store = SafetensorsStore::from_file(&path); + + // keys() should work without populating the full cache (fast path) + let keys = load_store.keys().unwrap(); + assert_eq!(keys.len(), 4); + + // Now call get_all_snapshots to populate cache + let snapshots = load_store.get_all_snapshots().unwrap(); + assert_eq!(snapshots.len(), 4); + + // keys() should now use the cached data + let keys2 = load_store.keys().unwrap(); + assert_eq!(keys2.len(), 4); +} + +#[test] +#[cfg(feature = "std")] +fn test_file_caching_behavior() { + use tempfile::tempdir; + + let device = Default::default(); + let module = DirectAccessTestModule::new(&device); + + let temp_dir = tempdir().unwrap(); + let path = temp_dir.path().join("test_caching.safetensors"); + + let mut save_store = SafetensorsStore::from_file(&path); + save_store.collect_from(&module).unwrap(); + + let mut load_store = SafetensorsStore::from_file(&path); + + // First call populates cache + let snapshots1 = load_store.get_all_snapshots().unwrap(); + assert_eq!(snapshots1.len(), 4); + + // Second call uses cache + let snapshots2 = load_store.get_all_snapshots().unwrap(); + assert_eq!(snapshots2.len(), 4); +} + +#[test] +#[cfg(feature = "std")] +fn test_file_cache_invalidation_on_save() { + use tempfile::tempdir; + + let device = Default::default(); + let module = DirectAccessTestModule::new(&device); + + let temp_dir = tempdir().unwrap(); + let path = temp_dir.path().join("test_invalidation.safetensors"); + + // Create store, save, and populate cache + let mut store = SafetensorsStore::from_file(&path).overwrite(true); + store.collect_from(&module).unwrap(); + + let snapshots1 = store.get_all_snapshots().unwrap(); + assert_eq!(snapshots1.len(), 4); + + // Save again (this should invalidate cache) + store.collect_from(&module).unwrap(); + + // Cache should be repopulated with fresh data + let snapshots2 = store.get_all_snapshots().unwrap(); + assert_eq!(snapshots2.len(), 4); +} diff --git a/crates/burn-store/src/safetensors/tests/error_handling.rs b/crates/burn-store/src/safetensors/tests/error_handling.rs new file mode 100644 index 0000000..d70e76a --- /dev/null +++ b/crates/burn-store/src/safetensors/tests/error_handling.rs @@ -0,0 +1,45 @@ +use crate::{ModuleSnapshot, SafetensorsStore}; +use burn_nn::LinearConfig; + +#[test] +fn shape_mismatch_errors() { + let device = Default::default(); + + // Create a module + let module = LinearConfig::new(2, 2).with_bias(true).init(&device); + + // Save module + let mut save_store = SafetensorsStore::from_bytes(None); + module.save_into(&mut save_store).unwrap(); + + // Try to load into incompatible module (different dimensions) + let mut incompatible_module = LinearConfig::new(3, 3).with_bias(true).init(&device); + + // Load without validation - should return errors in the result + let mut load_store = SafetensorsStore::from_bytes(None).validate(false); // Disable validation to get errors in result + if let SafetensorsStore::Memory(ref mut p) = load_store + && let SafetensorsStore::Memory(ref p_save) = save_store + { + // Get Arc and extract data + let data_arc = p_save.data().unwrap(); + p.set_data(data_arc.as_ref().clone()); + } + + let result = incompatible_module.load_from(&mut load_store).unwrap(); + + // Should have errors due to shape mismatch + assert!(!result.errors.is_empty()); + + // Try again with validation enabled - should return Err + let mut load_store_with_validation = SafetensorsStore::from_bytes(None).validate(true); + if let SafetensorsStore::Memory(ref mut p) = load_store_with_validation + && let SafetensorsStore::Memory(ref p_save) = save_store + { + // Get Arc and extract data + let data_arc = p_save.data().unwrap(); + p.set_data(data_arc.as_ref().clone()); + } + + let validation_result = incompatible_module.load_from(&mut load_store_with_validation); + assert!(validation_result.is_err()); +} diff --git a/crates/burn-store/src/safetensors/tests/file_io.rs b/crates/burn-store/src/safetensors/tests/file_io.rs new file mode 100644 index 0000000..56dcb06 --- /dev/null +++ b/crates/burn-store/src/safetensors/tests/file_io.rs @@ -0,0 +1,273 @@ +use burn_core as burn; + +use crate::{ModuleSnapshot, ModuleStore, SafetensorsStore}; +use burn_core::module::{Module, Param}; +use burn_core::tensor::{Device, Tensor}; +use burn_nn::{Initializer, LinearConfig}; + +use tempfile::tempdir; + +// Define a test model with forward pass +#[derive(Module, Debug)] +struct ForwardTestModel { + linear1: burn_nn::Linear, + linear2: burn_nn::Linear, +} + +impl ForwardTestModel { + fn forward(&self, input: Tensor<2>) -> Tensor<2> { + let x = self.linear1.forward(input); + let x = burn::tensor::activation::gelu(x); + self.linear2.forward(x) + } +} + +// Define config for the model +#[derive(burn::config::Config, Debug)] +struct ForwardTestModelConfig { + input_size: usize, + hidden_size: usize, + output_size: usize, +} + +impl ForwardTestModelConfig { + fn init(&self, device: &Device) -> ForwardTestModel { + ForwardTestModel { + linear1: LinearConfig::new(self.input_size, self.hidden_size) + .with_bias(true) + .init(device), + linear2: LinearConfig::new(self.hidden_size, self.output_size) + .with_bias(true) + .init(device), + } + } +} + +#[derive(Module, Debug)] +pub struct ModuleBasic { + weight_basic: Param>, +} + +impl ModuleBasic { + fn new(device: &Device) -> Self { + Self { + weight_basic: Initializer::Normal { + std: 1.0, + mean: 0.0, + } + .init([20, 20], device), + } + } +} + +#[derive(Module, Debug)] +pub struct ModuleComposed { + weight: Param>, + basic: ModuleBasic, + tuple: (ModuleBasic, ModuleBasic), +} + +impl ModuleComposed { + fn new(device: &Device) -> Self { + let weight = Initializer::Normal { + std: 1.0, + mean: 0.0, + } + .init([20, 20], device); + + Self { + weight, + basic: ModuleBasic::new(device), + tuple: (ModuleBasic::new(device), ModuleBasic::new(device)), + } + } +} + +#[test] +fn file_based_loading() { + use std::fs; + + let device = Default::default(); + let module = LinearConfig::new(4, 2).with_bias(true).init(&device); + + // Create temp file path + let temp_dir = std::env::temp_dir(); + let file_path = temp_dir.join("test_safetensors.st"); + + // Save to file + let mut save_store = SafetensorsStore::from_file(&file_path).metadata("test", "file_loading"); + + module.save_into(&mut save_store).unwrap(); + + // Verify file exists + assert!(file_path.exists()); + + // Load from file (will use memory-mapped loading if available) + let mut load_store = SafetensorsStore::from_file(&file_path); + + let mut loaded_module = LinearConfig::new(4, 2).with_bias(true).init(&device); + + let result = loaded_module.load_from(&mut load_store).unwrap(); + + assert!(result.is_success()); + assert_eq!(result.applied.len(), 2); // weight and bias + + // Clean up + fs::remove_file(file_path).ok(); +} + +#[test] +fn test_store_overwrite_protection() { + use tempfile::tempdir; + + let device = Default::default(); + let module = LinearConfig::new(4, 2).with_bias(true).init(&device); + + // Create temp directory and file path (file doesn't exist yet) + let temp_dir = tempdir().unwrap(); + let path = temp_dir.path().join("test_model.safetensors"); + + // First save - should succeed + let mut save_store = SafetensorsStore::from_file(&path); + save_store.collect_from(&module).unwrap(); + assert!(path.exists()); + + // Second save without overwrite flag - should fail + let mut save_store2 = SafetensorsStore::from_file(&path); + let result = save_store2.collect_from(&module); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("File already exists") + ); + + // Third save with overwrite flag - should succeed + let mut save_store3 = SafetensorsStore::from_file(&path).overwrite(true); + save_store3.collect_from(&module).unwrap(); + + // Verify file still exists and is valid + let mut load_store = SafetensorsStore::from_file(&path); + let mut module2 = LinearConfig::new(4, 2).with_bias(true).init(&device); + let result = load_store.apply_to(&mut module2).unwrap(); + assert!(result.is_success()); +} + +#[test] +fn test_store_overwrite_with_metadata() { + use tempfile::tempdir; + + let device = Default::default(); + let module = LinearConfig::new(4, 2).with_bias(true).init(&device); + + // Create temp directory and file path + let temp_dir = tempdir().unwrap(); + let path = temp_dir.path().join("test_model_metadata.safetensors"); + + // First save with v1 metadata and overwrite enabled + let mut save_store = SafetensorsStore::from_file(&path) + .metadata("model_version", "v1") + .overwrite(true); + save_store.collect_from(&module).unwrap(); + + // Second save with v2 metadata and overwrite enabled + let mut save_store2 = SafetensorsStore::from_file(&path) + .metadata("model_version", "v2") + .overwrite(true); + save_store2.collect_from(&module).unwrap(); + + // Load and verify the metadata was updated to v2 + let mut load_store = SafetensorsStore::from_file(&path); + // Since we can't easily access metadata after loading, we just verify the file loads successfully + let mut module2 = LinearConfig::new(4, 2).with_bias(true).init(&device); + let result = module2.load_from(&mut load_store).unwrap(); + assert!(result.is_success()); +} + +#[test] +fn test_forward_pass_preservation_after_save_load() { + let device = Default::default(); + + // Create model config + let config = ForwardTestModelConfig { + input_size: 4, + hidden_size: 8, + output_size: 2, + }; + + // Initialize model1 with random weights + let model1 = config.init(&device); + + // Create random input + let input = Tensor::<2>::random( + [1, 4], + burn_core::tensor::Distribution::Uniform(-1.0, 1.0), + &device, + ); + + // Forward pass with model1 -> output1 + let output1 = model1.forward(input.clone()); + + // Save model1 weights + let temp_dir = tempdir().unwrap(); + let path = temp_dir.path().join("forward_test_model.safetensors"); + let mut save_store = SafetensorsStore::from_file(&path); + save_store.collect_from(&model1).unwrap(); + + // Initialize model2 with different random weights + let mut model2 = config.init(&device); + + // Forward pass with model2 -> output2 (should differ from output1) + let output2 = model2.forward(input.clone()); + + // Verify output2 differs from output1 (different random weights) + assert!( + !output1 + .clone() + .all_close(output2.clone(), Some(1e-6), Some(1e-6)), + "output2 should differ from output1 (different random initializations)" + ); + + // Load model1 weights into model2 + let mut load_store = SafetensorsStore::from_file(&path); + let result = load_store.apply_to(&mut model2).unwrap(); + assert!(result.is_success()); + assert_eq!(result.applied.len(), 4); // 2 weights + 2 biases + + // Forward pass with model2 (now has model1 weights) -> output3 + let output3 = model2.forward(input.clone()); + + // Verify output3 equals output1 (same weights) + assert!( + output1.all_close(output3, Some(1e-6), Some(1e-6)), + "output3 should equal output1 after loading weights" + ); +} + +#[test] +fn should_save_load_compose() { + let device = Device::default(); + let module_1 = ModuleComposed::new(&device); + let mut module_2 = ModuleComposed::new(&device); + assert_ne!(module_1.weight.to_data(), module_2.weight.to_data()); + assert_ne!( + module_1.basic.weight_basic.to_data(), + module_2.basic.weight_basic.to_data() + ); + + let temp_dir = tempdir().unwrap(); + let path = temp_dir.path().join("save_load_compose.safetensors"); + let mut store = SafetensorsStore::from_file(&path); + module_1.save_into(&mut store).unwrap(); + + let mut load_store = SafetensorsStore::from_file(&path); + let result = module_2.load_from(&mut load_store).unwrap(); + assert!(result.is_success()); + + assert_eq!(module_1.weight.to_data(), module_2.weight.to_data()); + assert_eq!( + module_1.basic.weight_basic.to_data(), + module_2.basic.weight_basic.to_data() + ); +} diff --git a/crates/burn-store/src/safetensors/tests/filtering.rs b/crates/burn-store/src/safetensors/tests/filtering.rs new file mode 100644 index 0000000..6f139bc --- /dev/null +++ b/crates/burn-store/src/safetensors/tests/filtering.rs @@ -0,0 +1,167 @@ +use crate::{ModuleSnapshot, SafetensorsStore}; + +use super::round_trip::ComplexModule; + +#[test] +#[cfg(target_has_atomic = "ptr")] +fn filtered_export_import() { + let device = Default::default(); + let module1 = ComplexModule::new(&device); + let mut module2 = ComplexModule::new_zeros(&device); + + // Export only encoder tensors using the builder pattern + let mut save_store = SafetensorsStore::from_bytes(None).with_regex(r"^encoder\..*"); + module1.save_into(&mut save_store).unwrap(); + + // Import filtered tensors - need to allow partial since we only saved encoder tensors + let mut load_store = SafetensorsStore::from_bytes(None).allow_partial(true); + if let SafetensorsStore::Memory(ref mut p) = load_store + && let SafetensorsStore::Memory(ref p_save) = save_store + { + // Get Arc and extract data + let data_arc = p_save.data().unwrap(); + p.set_data(data_arc.as_ref().clone()); + } + let result = module2.load_from(&mut load_store).unwrap(); + + assert!(result.is_success()); + assert_eq!(result.applied.len(), 3); // encoder.weight, encoder.bias, encoder.norm + assert!(!result.missing.is_empty()); // decoder and layers tensors are missing +} + +#[test] +#[cfg(target_has_atomic = "ptr")] +fn builder_pattern_filtering() { + let device = Default::default(); + let module = ComplexModule::new(&device); + + // Test with_regex - multiple patterns (OR logic) + let mut store = SafetensorsStore::from_bytes(None) + .with_regex(r"^encoder\..*") // Match encoder tensors + .with_regex(r".*\.bias$"); // OR match any bias tensors + + let views = module.collect(None, None, false); + let filtered_count = views + .iter() + .filter(|v| { + let path = v.full_path(); + path.starts_with("encoder.") || path.ends_with(".bias") + }) + .count(); + + module.save_into(&mut store).unwrap(); + + // Verify we saved the expected number of tensors + if let SafetensorsStore::Memory(ref p) = store { + let data = p.data().unwrap(); + let tensors = safetensors::SafeTensors::deserialize(&data).unwrap(); + assert_eq!(tensors.len(), filtered_count); + } +} + +#[test] +fn builder_pattern_exact_paths() { + let device = Default::default(); + let module = ComplexModule::new(&device); + + // Test with_full_path and with_full_paths + let paths = vec!["encoder.weight", "decoder.scale"]; + let mut store = SafetensorsStore::from_bytes(None) + .with_full_path("encoder.norm") + .with_full_paths(paths.clone()); + + module.save_into(&mut store).unwrap(); + + // Verify only specified tensors were saved + if let SafetensorsStore::Memory(ref p) = store { + let data = p.data().unwrap(); + let tensors = safetensors::SafeTensors::deserialize(&data).unwrap(); + assert_eq!(tensors.len(), 3); // encoder.norm + encoder.weight + decoder.scale + + for (name, _) in tensors.tensors() { + assert!(name == "encoder.norm" || name == "encoder.weight" || name == "decoder.scale"); + } + } +} + +#[test] +fn builder_pattern_with_predicate() { + let device = Default::default(); + let module = ComplexModule::new(&device); + + // Test with_predicate - custom logic + let mut store = SafetensorsStore::from_bytes(None).with_predicate(|path, _| { + // Only save tensors with "layer" in the path and ending with "weight" + path.contains("layer") && path.ends_with("weight") + }); + + module.save_into(&mut store).unwrap(); + + // Verify only layer weights were saved + if let SafetensorsStore::Memory(ref p) = store { + let data = p.data().unwrap(); + let tensors = safetensors::SafeTensors::deserialize(&data).unwrap(); + + for (name, _) in tensors.tensors() { + assert!(name.contains("layer")); + assert!(name.ends_with("weight")); + } + } +} + +#[test] +fn builder_pattern_combined() { + let device = Default::default(); + let module = ComplexModule::new(&device); + + // Combine multiple filter methods + #[cfg(target_has_atomic = "ptr")] + { + let mut store = SafetensorsStore::from_bytes(None) + .with_regex(r"^encoder\..*") // All encoder tensors + .with_full_path("decoder.scale") // Plus specific decoder.scale + .with_predicate(|path, _| { + // Plus any projection tensors + path.contains("projection") + }); + + module.save_into(&mut store).unwrap(); + + if let SafetensorsStore::Memory(ref p) = store { + let data = p.data().unwrap(); + let tensors = safetensors::SafeTensors::deserialize(&data).unwrap(); + + // Should have encoder.*, decoder.scale, and projection tensors + let mut names = Vec::new(); + for (name, _) in tensors.tensors() { + names.push(name); + } + assert!(names.iter().any(|n| n == "encoder.weight")); + assert!(names.iter().any(|n| n == "encoder.bias")); + assert!(names.iter().any(|n| n == "encoder.norm")); + assert!(names.iter().any(|n| n == "decoder.scale")); + // decoder.projection.* should also be included due to predicate + assert!(names.iter().any(|n| n.contains("projection"))); + } + } +} + +#[test] +fn builder_pattern_match_all() { + let device = Default::default(); + let module = ComplexModule::new(&device); + + let all_views = module.collect(None, None, false); + let total_count = all_views.len(); + + // Test match_all - should save everything + let mut store = SafetensorsStore::from_bytes(None).match_all(); + + module.save_into(&mut store).unwrap(); + + if let SafetensorsStore::Memory(ref p) = store { + let data = p.data().unwrap(); + let tensors = safetensors::SafeTensors::deserialize(&data).unwrap(); + assert_eq!(tensors.len(), total_count); + } +} diff --git a/crates/burn-store/src/safetensors/tests/integration.rs b/crates/burn-store/src/safetensors/tests/integration.rs new file mode 100644 index 0000000..ef31e7b --- /dev/null +++ b/crates/burn-store/src/safetensors/tests/integration.rs @@ -0,0 +1,169 @@ +use burn_core as burn; + +use crate::{ModuleSnapshot, SafetensorsStore}; +use burn_core::module::{Module, Param}; +use burn_core::tensor::{Device, Tensor}; + +// Integration tests demonstrating the SafeTensors store API +#[derive(Module, Debug)] +struct IntegrationTestModel { + encoder: IntegrationEncoderModule, + decoder: IntegrationDecoderModule, + head: IntegrationHeadModule, +} + +#[derive(Module, Debug)] +struct IntegrationEncoderModule { + layer1: IntegrationLinearLayer, + layer2: IntegrationLinearLayer, + norm: IntegrationNormLayer, +} + +#[derive(Module, Debug)] +struct IntegrationDecoderModule { + layer1: IntegrationLinearLayer, + layer2: IntegrationLinearLayer, + norm: IntegrationNormLayer, +} + +#[derive(Module, Debug)] +struct IntegrationHeadModule { + weight: Param>, + bias: Param>, +} + +#[derive(Module, Debug)] +struct IntegrationLinearLayer { + weight: Param>, + bias: Param>, +} + +#[derive(Module, Debug)] +struct IntegrationNormLayer { + scale: Param>, + shift: Param>, +} + +impl IntegrationTestModel { + fn new(device: &Device) -> Self { + Self { + encoder: IntegrationEncoderModule::new(device), + decoder: IntegrationDecoderModule::new(device), + head: IntegrationHeadModule::new(device), + } + } +} + +impl IntegrationEncoderModule { + fn new(device: &Device) -> Self { + Self { + layer1: IntegrationLinearLayer::new(device, 1), + layer2: IntegrationLinearLayer::new(device, 2), + norm: IntegrationNormLayer::new(device), + } + } +} + +impl IntegrationDecoderModule { + fn new(device: &Device) -> Self { + Self { + layer1: IntegrationLinearLayer::new(device, 3), + layer2: IntegrationLinearLayer::new(device, 4), + norm: IntegrationNormLayer::new(device), + } + } +} + +impl IntegrationHeadModule { + fn new(device: &Device) -> Self { + Self { + weight: Param::from_data([[5.0, 6.0], [7.0, 8.0]], device), + bias: Param::from_data([9.0, 10.0], device), + } + } +} + +impl IntegrationLinearLayer { + fn new(device: &Device, seed: i32) -> Self { + let weight_data = [ + [seed as f32, (seed + 1) as f32], + [(seed + 2) as f32, (seed + 3) as f32], + ]; + let bias_data = [(seed + 4) as f32, (seed + 5) as f32]; + + Self { + weight: Param::from_data(weight_data, device), + bias: Param::from_data(bias_data, device), + } + } +} + +impl IntegrationNormLayer { + fn new(device: &Device) -> Self { + Self { + scale: Param::from_data([1.0, 2.0], device), + shift: Param::from_data([0.1, 0.2], device), + } + } +} + +#[test] +fn basic_usage() { + let device = Default::default(); + let model = IntegrationTestModel::new(&device); + + // Save using new API (format, producer and version are automatically added) + let mut save_store = SafetensorsStore::from_bytes(None).metadata("model_name", "test_model"); + + // Use collect_to method + model.save_into(&mut save_store).unwrap(); + + // Load using new API + let mut load_store = SafetensorsStore::from_bytes(None); + if let SafetensorsStore::Memory(ref mut p) = load_store + && let SafetensorsStore::Memory(ref p_save) = save_store + { + p.set_data(p_save.data().unwrap().as_ref().clone()); + } + + let mut target_model = IntegrationTestModel::new(&device); + let result = target_model.load_from(&mut load_store).unwrap(); + + assert!(result.is_success()); + assert_eq!(result.applied.len(), 14); // All tensors should be applied + assert_eq!(result.errors.len(), 0); + assert_eq!(result.unused.len(), 0); +} + +#[test] +#[cfg(target_has_atomic = "ptr")] +fn with_filtering() { + let device = Default::default(); + let model = IntegrationTestModel::new(&device); + + // Save only encoder tensors using the builder pattern + let mut save_store = SafetensorsStore::from_bytes(None) + .with_regex(r"^encoder\..*") + .metadata("subset", "encoder_only"); + + model.save_into(&mut save_store).unwrap(); + + // Load into new model - need to allow partial loading since we only saved encoder tensors + let mut load_store = SafetensorsStore::from_bytes(None).allow_partial(true); + if let SafetensorsStore::Memory(ref mut p) = load_store + && let SafetensorsStore::Memory(ref p_save) = save_store + { + p.set_data(p_save.data().unwrap().as_ref().clone()); + } + + let mut target_model = IntegrationTestModel::new(&device); + let result = target_model.load_from(&mut load_store).unwrap(); + + // Only encoder tensors should be applied + assert_eq!(result.applied.len(), 6); // encoder has 6 tensors (2 layers × 2 + norm × 2) + + // Check that only encoder tensors were applied + for tensor_name in &result.applied { + assert!(tensor_name.starts_with("encoder.")); + } +} diff --git a/crates/burn-store/src/safetensors/tests/metadata.rs b/crates/burn-store/src/safetensors/tests/metadata.rs new file mode 100644 index 0000000..94e8b56 --- /dev/null +++ b/crates/burn-store/src/safetensors/tests/metadata.rs @@ -0,0 +1,101 @@ +use crate::{ModuleSnapshot, SafetensorsStore}; +use burn_nn::LinearConfig; + +#[test] +fn default_metadata_included() { + // Verify that default metadata is automatically included + let default_metadata = SafetensorsStore::default_metadata(); + + // Check that format, producer and version are present + assert_eq!(default_metadata.get("format").unwrap(), "safetensors"); + assert_eq!(default_metadata.get("producer").unwrap(), "burn"); + assert!(default_metadata.contains_key("version")); + + // The version should be the crate version + let version = default_metadata.get("version").unwrap(); + assert!(!version.is_empty()); +} + +#[test] +fn metadata_preservation() { + let device = Default::default(); + let module = LinearConfig::new(4, 2).with_bias(true).init(&device); + + // Write with metadata - note that format, producer and version are automatically added + let mut save_store = SafetensorsStore::from_bytes(None) + .metadata("model_type", "linear") + .metadata("custom_field", "test_value"); + + module.save_into(&mut save_store).unwrap(); + + // Verify metadata was saved (would need to add a method to check metadata) + // For now, just verify the round trip works + let mut load_store = SafetensorsStore::from_bytes(None); + if let SafetensorsStore::Memory(ref mut p) = load_store + && let SafetensorsStore::Memory(ref p_save) = save_store + { + // Get Arc and extract data + let data_arc = p_save.data().unwrap(); + p.set_data(data_arc.as_ref().clone()); + } + + let mut module2 = LinearConfig::new(4, 2).with_bias(true).init(&device); + let result = module2.load_from(&mut load_store).unwrap(); + + assert!(result.is_success()); +} + +#[test] +fn clear_metadata_removes_all() { + let device = Default::default(); + let module = LinearConfig::new(4, 2).with_bias(true).init(&device); + + // Create store with custom metadata, then clear all + let mut save_store = SafetensorsStore::from_bytes(None) + .metadata("model_type", "linear") + .metadata("custom_field", "test_value") + .clear_metadata(); // Should remove all metadata including defaults + + module.save_into(&mut save_store).unwrap(); + + // Load and verify the module still works (metadata is optional) + let mut load_store = SafetensorsStore::from_bytes(None); + if let SafetensorsStore::Memory(ref mut p) = load_store + && let SafetensorsStore::Memory(ref p_save) = save_store + { + let data_arc = p_save.data().unwrap(); + p.set_data(data_arc.as_ref().clone()); + } + + let mut module2 = LinearConfig::new(4, 2).with_bias(true).init(&device); + let result = module2.load_from(&mut load_store).unwrap(); + + assert!(result.is_success()); +} + +#[test] +fn clear_then_add_custom_metadata() { + let device = Default::default(); + let module = LinearConfig::new(4, 2).with_bias(true).init(&device); + + // Clear all metadata, then add only custom ones + let mut save_store = SafetensorsStore::from_bytes(None) + .clear_metadata() + .metadata("only_custom", "value"); + + module.save_into(&mut save_store).unwrap(); + + // Verify round-trip works + let mut load_store = SafetensorsStore::from_bytes(None); + if let SafetensorsStore::Memory(ref mut p) = load_store + && let SafetensorsStore::Memory(ref p_save) = save_store + { + let data_arc = p_save.data().unwrap(); + p.set_data(data_arc.as_ref().clone()); + } + + let mut module2 = LinearConfig::new(4, 2).with_bias(true).init(&device); + let result = module2.load_from(&mut load_store).unwrap(); + + assert!(result.is_success()); +} diff --git a/crates/burn-store/src/safetensors/tests/mixed_datatypes.rs b/crates/burn-store/src/safetensors/tests/mixed_datatypes.rs new file mode 100644 index 0000000..116d22b --- /dev/null +++ b/crates/burn-store/src/safetensors/tests/mixed_datatypes.rs @@ -0,0 +1,436 @@ +use burn_core as burn; + +use burn_core::module::{Module, Param, ParamId}; +use burn_core::tensor::{Bool, Device, Int, Tensor}; +use burn_nn as nn; + +use crate::{ModuleSnapshot, SafetensorsStore}; + +/// Simple model with different data types for testing +#[derive(Module, Debug)] +pub struct MixedDtypeModel { + // Standard neural network layers (float tensors) + linear: nn::Linear, + + // Direct tensor parameters of different types + float_tensor: Param>, + + int_tensor: Param>, + + bool_tensor: Param>, +} + +impl MixedDtypeModel { + pub fn new(device: &Device) -> Self { + Self { + linear: nn::LinearConfig::new(3, 3).init(device), + + // Simple float values + float_tensor: Param::from_tensor(Tensor::from_floats( + [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + device, + )), + + // Simple integer values + int_tensor: Param::initialized( + ParamId::new(), + Tensor::from_ints([[1, 2, 3], [4, 5, 6]], device), + ), + + // Simple boolean values + bool_tensor: Param::initialized( + ParamId::new(), + Tensor::from_bool( + burn::tensor::TensorData::new( + vec![true, false, true, false, true, false], + [2, 3], + ), + device, + ), + ), + } + } +} + +#[cfg(test)] +#[allow(clippy::excessive_precision)] +mod tests { + use burn_core::tensor::{BoolStore, DType}; + + use super::*; + + #[test] + fn test_mixed_dtypes_round_trip() { + let device = Default::default(); + + // Create model with mixed data types + let model = MixedDtypeModel::new(&device); + + // Save to bytes + let mut save_store = SafetensorsStore::from_bytes(None); + model.save_into(&mut save_store).expect("Failed to save"); + let bytes = save_store.get_bytes().expect("Failed to get bytes"); + + // Load into a new model + let mut load_store = SafetensorsStore::from_bytes(Some(bytes)); + let mut loaded_model = MixedDtypeModel::new(&device); + loaded_model + .load_from(&mut load_store) + .expect("Failed to load"); + + // Verify float tensor is preserved + let orig_float = model.float_tensor.val().into_data(); + let loaded_float = loaded_model.float_tensor.val().into_data(); + assert_eq!(orig_float, loaded_float, "Float tensor not preserved"); + + // Verify integer tensor is preserved + let orig_int = model.int_tensor.val().into_data(); + let loaded_int = loaded_model.int_tensor.val().into_data(); + assert_eq!(orig_int, loaded_int, "Integer tensor not preserved"); + + // Verify boolean tensor is preserved + let orig_bool = model.bool_tensor.val().into_data(); + let loaded_bool = loaded_model.bool_tensor.val().into_data(); + assert_eq!(orig_bool, loaded_bool, "Boolean tensor not preserved"); + } + + #[test] + fn test_dtype_detection() { + let device = Default::default(); + + let model = MixedDtypeModel::new(&device); + let snapshots = model.collect(None, None, false); + + for snapshot in snapshots { + let path = snapshot.full_path(); + let dtype = snapshot.dtype; + + if path.contains("float_tensor") || path.contains("linear") { + assert_eq!( + dtype, + burn::tensor::DType::F32, + "Float tensor {} should have F32 dtype", + path + ); + } else if path.contains("int_tensor") { + assert!( + matches!( + dtype, + burn::tensor::DType::I64 + | burn::tensor::DType::I32 + | burn::tensor::DType::I16 + | burn::tensor::DType::I8 + ), + "Integer tensor {} should have integer dtype, got {:?}", + path, + dtype + ); + } else if path.contains("bool_tensor") { + assert_eq!( + dtype, + burn::tensor::DType::Bool(BoolStore::Native), + "Boolean tensor {} should have Bool dtype", + path + ); + } + } + } + + #[test] + fn test_extreme_values() { + let device = Device::default(); + + #[derive(Module, Debug)] + struct ExtremeValueModel { + large_floats: Param>, + small_floats: Param>, + large_ints: Param>, + } + + impl ExtremeValueModel { + fn new(device: &Device) -> Self { + Self { + large_floats: Param::from_tensor(Tensor::from_floats( + [1e30, -1e30, f32::MAX, f32::MIN], + device, + )), + small_floats: Param::from_tensor(Tensor::from_floats( + [1e-30, -1e-30, f32::MIN_POSITIVE, f32::EPSILON], + device, + )), + large_ints: Param::initialized( + ParamId::new(), + Tensor::from_ints([i32::MAX, i32::MIN, 0, -1], device), + ), + } + } + } + + let model = ExtremeValueModel::new(&device); + + // Save and load + let mut save_store = SafetensorsStore::from_bytes(None); + model.save_into(&mut save_store).expect("Failed to save"); + let bytes = save_store.get_bytes().expect("Failed to get bytes"); + + let mut load_store = SafetensorsStore::from_bytes(Some(bytes)); + let mut loaded_model = ExtremeValueModel::new(&device); + loaded_model + .load_from(&mut load_store) + .expect("Failed to load"); + + // Check exact preservation + assert_eq!( + model.large_floats.val().into_data(), + loaded_model.large_floats.val().into_data(), + "Large floats not preserved" + ); + assert_eq!( + model.small_floats.val().into_data(), + loaded_model.small_floats.val().into_data(), + "Small floats not preserved" + ); + assert_eq!( + model.large_ints.val().into_data(), + loaded_model.large_ints.val().into_data(), + "Large integers not preserved" + ); + } + + #[test] + fn test_mixed_precision_floats() { + // Note: While SafeTensors format supports storing tensors with different precisions + // (F16, BF16, F32, F64, etc.) in the same file, Burn's backend architecture currently + // requires all tensors in a model instance to share the same floating-point precision. + // + // However, for storage purposes, SafeTensors can correctly save and load tensors + // with their original precision, preserving the data type information in the file format. + // This test demonstrates that different precision backends work correctly with SafeTensors. + + // Test with f32 backend + { + let device = Default::default(); + + let model = MixedDtypeModel::new(&device); + + // Save to bytes + let mut save_store = SafetensorsStore::from_bytes(None); + model.save_into(&mut save_store).expect("Failed to save"); + let bytes = save_store.get_bytes().expect("Failed to get bytes"); + + // Load and verify + let mut load_store = SafetensorsStore::from_bytes(Some(bytes)); + let mut loaded_model = MixedDtypeModel::new(&device); + loaded_model + .load_from(&mut load_store) + .expect("Failed to load"); + + assert_eq!( + model.float_tensor.val().into_data(), + loaded_model.float_tensor.val().into_data(), + "F32 float tensor not preserved" + ); + } + + // Test with f64 weights + { + let device = Default::default(); + + #[derive(Module, Debug)] + struct F64Model { + weight: Param>, + double_precision: Param>, + } + + let model = F64Model { + weight: Param::from_tensor(Tensor::ones([3, 3], (&device, DType::F64))), + double_precision: Param::from_tensor(Tensor::from_data( + [ + [1.234567890123456789, 2.345678901234567890], + [3.456789012345678901, 4.567890123456789012], + ], + (&device, DType::F64), + )), + }; + + // Save to bytes + let mut save_store = SafetensorsStore::from_bytes(None); + model.save_into(&mut save_store).expect("Failed to save"); + let bytes = save_store.get_bytes().expect("Failed to get bytes"); + + // Load and verify + let mut load_store = SafetensorsStore::from_bytes(Some(bytes)); + let mut loaded_model = F64Model { + weight: Param::from_tensor(Tensor::ones([3, 3], &device)), + double_precision: Param::from_tensor(Tensor::zeros([2, 2], &device)), + }; + loaded_model + .load_from(&mut load_store) + .expect("Failed to load"); + + let orig = model.double_precision.val().into_data(); + let loaded = loaded_model.double_precision.val().into_data(); + assert_eq!(orig, loaded, "F64 double precision not preserved"); + } + } + + #[test] + fn test_mixed_precision_integers() { + let device = Default::default(); + + #[derive(Module, Debug)] + struct MultiIntModel { + // Note: Burn's Tensor uses the backend's default int type + // We can't directly specify i8, i16, etc. in the type system + // But we can test with different values that would fit in different ranges + small_ints: Param>, // Values that fit in i8 + medium_ints: Param>, // Values that fit in i16 + large_ints: Param>, // Values that need i32/i64 + } + + let model = MultiIntModel { + small_ints: Param::initialized( + ParamId::new(), + Tensor::from_ints([127i32, -128, 0, 42], &device), + ), + medium_ints: Param::initialized( + ParamId::new(), + Tensor::from_ints([32767i32, -32768, 1000, -1000], &device), + ), + large_ints: Param::initialized( + ParamId::new(), + Tensor::from_ints([i32::MAX, i32::MIN, 1_000_000, -1_000_000], &device), + ), + }; + + // Save to bytes + let mut save_store = SafetensorsStore::from_bytes(None); + model.save_into(&mut save_store).expect("Failed to save"); + let bytes = save_store.get_bytes().expect("Failed to get bytes"); + + // Load and verify + let mut load_store = SafetensorsStore::from_bytes(Some(bytes)); + let mut loaded_model = MultiIntModel { + small_ints: Param::initialized(ParamId::new(), Tensor::zeros([4], &device)), + medium_ints: Param::initialized(ParamId::new(), Tensor::zeros([4], &device)), + large_ints: Param::initialized(ParamId::new(), Tensor::zeros([4], &device)), + }; + loaded_model + .load_from(&mut load_store) + .expect("Failed to load"); + + assert_eq!( + model.small_ints.val().into_data(), + loaded_model.small_ints.val().into_data(), + "Small ints (i8 range) not preserved" + ); + assert_eq!( + model.medium_ints.val().into_data(), + loaded_model.medium_ints.val().into_data(), + "Medium ints (i16 range) not preserved" + ); + assert_eq!( + model.large_ints.val().into_data(), + loaded_model.large_ints.val().into_data(), + "Large ints (i32 range) not preserved" + ); + } + + #[test] + fn test_comprehensive_mixed_types() { + let device = Default::default(); + + #[derive(Module, Debug)] + struct ComprehensiveModel { + // Neural network layers + linear1: nn::Linear, + conv2d: nn::conv::Conv2d, + + // Different tensor types + float32_weights: Param>, + integer_indices: Param>, + boolean_mask: Param>, + } + + let model = ComprehensiveModel { + linear1: nn::LinearConfig::new(4, 8).init(&device), + conv2d: nn::conv::Conv2dConfig::new([3, 16], [3, 3]).init(&device), + + float32_weights: Param::from_tensor(Tensor::from_floats( + [[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]], + &device, + )), + + integer_indices: Param::initialized( + ParamId::new(), + Tensor::from_ints( + [[0, 1, 2, 3], [10, 20, 30, 40], [100, 200, 300, 400]], + &device, + ), + ), + + boolean_mask: Param::initialized( + ParamId::new(), + Tensor::from_bool( + burn::tensor::TensorData::new( + vec![true, false, false, true, false, true, true, false], + [2, 4], + ), + &device, + ), + ), + }; + + // Collect all tensors + let snapshots = model.collect(None, None, false); + + // Verify we have all expected tensors + let paths: Vec = snapshots.iter().map(|s| s.full_path()).collect(); + assert!(paths.iter().any(|p| p.contains("linear1"))); + assert!(paths.iter().any(|p| p.contains("conv2d"))); + assert!(paths.iter().any(|p| p.contains("float32_weights"))); + assert!(paths.iter().any(|p| p.contains("integer_indices"))); + assert!(paths.iter().any(|p| p.contains("boolean_mask"))); + + // Save to bytes + let mut save_store = SafetensorsStore::from_bytes(None); + model.save_into(&mut save_store).expect("Failed to save"); + let bytes = save_store.get_bytes().expect("Failed to get bytes"); + + // Load into fresh model + let mut load_store = SafetensorsStore::from_bytes(Some(bytes)); + let mut loaded_model = ComprehensiveModel { + linear1: nn::LinearConfig::new(4, 8).init(&device), + conv2d: nn::conv::Conv2dConfig::new([3, 16], [3, 3]).init(&device), + float32_weights: Param::from_tensor(Tensor::zeros([2, 2, 2], &device)), + integer_indices: Param::initialized(ParamId::new(), Tensor::zeros([3, 4], &device)), + boolean_mask: Param::initialized( + ParamId::new(), + Tensor::from_bool( + burn::tensor::TensorData::new(vec![false; 8], [2, 4]), + &device, + ), + ), + }; + loaded_model + .load_from(&mut load_store) + .expect("Failed to load"); + + // Verify all data is preserved + assert_eq!( + model.float32_weights.val().into_data(), + loaded_model.float32_weights.val().into_data(), + "Float32 weights not preserved" + ); + assert_eq!( + model.integer_indices.val().into_data(), + loaded_model.integer_indices.val().into_data(), + "Integer indices not preserved" + ); + assert_eq!( + model.boolean_mask.val().into_data(), + loaded_model.boolean_mask.val().into_data(), + "Boolean mask not preserved" + ); + } +} diff --git a/crates/burn-store/src/safetensors/tests/mod.rs b/crates/burn-store/src/safetensors/tests/mod.rs new file mode 100644 index 0000000..1ac3037 --- /dev/null +++ b/crates/burn-store/src/safetensors/tests/mod.rs @@ -0,0 +1,12 @@ +mod adapter; +mod direct_access; +mod error_handling; +#[cfg(feature = "std")] +mod file_io; +mod filtering; +mod integration; +mod metadata; +mod mixed_datatypes; +mod multi_layer_verify; +mod pytorch_import; +mod round_trip; diff --git a/crates/burn-store/src/safetensors/tests/multi_layer_verify.rs b/crates/burn-store/src/safetensors/tests/multi_layer_verify.rs new file mode 100644 index 0000000..31a3d4b --- /dev/null +++ b/crates/burn-store/src/safetensors/tests/multi_layer_verify.rs @@ -0,0 +1,111 @@ +//! Tests for multi-layer model loading with SafeTensors format +use burn_core as burn; + +use burn_core::module::Module; +use burn_core::tensor::{Device, Tensor, Tolerance}; + +use burn_nn::{ + BatchNorm, BatchNormConfig, Linear, LinearConfig, PaddingConfig2d, Relu, + conv::{Conv2d, Conv2dConfig}, +}; + +/// Multi-layer neural network model for testing +#[derive(Module, Debug)] +pub struct Net { + conv1: Conv2d, + norm1: BatchNorm, + fc1: Linear, + relu: Relu, +} + +impl Net { + /// Create a new network instance + pub fn new(device: &Device) -> Self { + Self { + conv1: Conv2dConfig::new([3, 4], [3, 3]) + .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1)) + .init(device), + norm1: BatchNormConfig::new(4).init(device), + fc1: LinearConfig::new(4 * 8 * 8, 16).init(device), + relu: Relu::new(), + } + } + + /// Forward pass of the model + pub fn forward(&self, x: Tensor<4>) -> Tensor<2> { + let x = self.conv1.forward(x); + let x = self.norm1.forward(x); + let x = self.relu.forward(x); + // Flatten all dimensions except the batch dimension + let x = x.flatten(1, 3); + self.fc1.forward(x) + } +} + +use crate::{ModuleSnapshot, PyTorchToBurnAdapter, SafetensorsStore}; + +/// Path to the multi_layer.safetensors test file +fn get_safetensors_path() -> &'static str { + concat!( + env!("CARGO_MANIFEST_DIR"), + "/safetensors-tests/tests/multi_layer/multi_layer.safetensors" + ) +} + +#[test] +fn multi_layer_model() { + let device = Default::default(); + let safetensors_path = get_safetensors_path(); + + // Load model from SafeTensors file with PyTorch adapter + let mut store = SafetensorsStore::from_file(safetensors_path) + .with_from_adapter(PyTorchToBurnAdapter) + .validate(false) + .allow_partial(true); + + let mut model = Net::new(&device); + let result = model.load_from(&mut store).unwrap(); + + // Verify loading was successful + assert!( + !result.applied.is_empty(), + "Should have loaded some tensors" + ); + assert!( + result.errors.is_empty(), + "Should have no errors: {:?}", + result.errors + ); + + // Test forward pass + let input = Tensor::<4>::ones([1, 3, 8, 8], &device); + let output = model.forward(input); + + // Expected output values from PyTorch model + let expected = Tensor::<2>::from_data( + [[ + 0.04971555, + -0.16849735, + 0.05182848, + -0.18032673, + 0.23138367, + 0.05041867, + 0.13005908, + -0.32202929, + -0.07915690, + -0.03232457, + -0.19790289, + -0.17476529, + -0.19627589, + -0.21757686, + -0.31376451, + 0.08377837, + ]], + &device, + ); + + // Verify output matches expected values + output + .to_data() + .assert_approx_eq::(&expected.to_data(), Tolerance::default()); +} diff --git a/crates/burn-store/src/safetensors/tests/pytorch_import.rs b/crates/burn-store/src/safetensors/tests/pytorch_import.rs new file mode 100644 index 0000000..6e2c633 --- /dev/null +++ b/crates/burn-store/src/safetensors/tests/pytorch_import.rs @@ -0,0 +1,201 @@ +use burn_core as burn; + +use crate::{ModuleSnapshot, SafetensorsStore}; +use burn_core::module::Module; +use burn_core::tensor::{Device, Tensor}; +use burn_nn::{ + BatchNorm, BatchNormConfig, Linear, LinearConfig, PaddingConfig2d, Relu, + conv::{Conv2d, Conv2dConfig}, +}; + +#[derive(Module, Debug)] +pub struct Net { + conv1: Conv2d, + norm1: BatchNorm, + fc1: Linear, + relu: Relu, +} + +impl Net { + pub fn new(device: &Device) -> Self { + Self { + conv1: Conv2dConfig::new([3, 4], [3, 3]) + .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1)) + .init(device), + norm1: BatchNormConfig::new(4).init(device), + fc1: LinearConfig::new(4 * 8 * 8, 16).init(device), + relu: Relu::new(), + } + } + + /// Forward pass of the model. + pub fn forward(&self, x: Tensor<4>) -> Tensor<2> { + let x = self.conv1.forward(x); + let x = self.norm1.forward(x); + let x = self.relu.forward(x); + // Flatten all dimensions except the batch dimension + let x = x.flatten(1, 3); + self.fc1.forward(x) + } +} + +#[test] +#[cfg(all(feature = "std", target_has_atomic = "ptr"))] +fn multi_layer_model_import() { + let device = Default::default(); + + // Reference the safetensors file from burn-import + let safetensors_path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/safetensors-tests/tests/multi_layer/multi_layer.safetensors" + ); + + // Load the model using SafetensorsStore + // Note: PyTorch and Burn have different conventions for linear layer weights + // PyTorch stores as [out_features, in_features], Burn as [in_features, out_features] + // Also, tensor names may differ (e.g., PyTorch uses different names for BatchNorm params) + let mut store = SafetensorsStore::from_file(safetensors_path) + .with_from_adapter(crate::PyTorchToBurnAdapter) // Use adapter to handle PyTorch format + .allow_partial(true); // Allow partial loading due to naming differences + let mut model = Net::new(&device); + + let result = model.load_from(&mut store).unwrap(); + + // With the adapter, weights should load correctly + assert!(!result.applied.is_empty()); + assert!( + result.errors.is_empty(), + "Should have no errors with adapter: {:?}", + result.errors + ); + + // Test forward pass with the loaded weights + // Note: Due to shape mismatches (PyTorch vs Burn conventions for linear layers), + // we can't directly compare outputs with PyTorch model. + // This test mainly verifies that the loading mechanism works. + let input = Tensor::<4>::ones([1, 3, 8, 8], &device); + let _output = model.forward(input); + + // Verify that some tensors were loaded successfully + // Conv and BatchNorm layers should load correctly + assert!(result.applied.iter().any(|n| n.contains("conv1"))); + assert!(result.applied.iter().any(|n| n.contains("norm1"))); +} + +#[test] +#[cfg(all(feature = "std", target_has_atomic = "ptr"))] +fn safetensors_round_trip_with_pytorch_model() { + let device = Default::default(); + + // Reference the safetensors file from burn-import + let safetensors_path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/safetensors-tests/tests/multi_layer/multi_layer.safetensors" + ); + + // Load the model from PyTorch safetensors + let mut load_store = SafetensorsStore::from_file(safetensors_path) + .with_from_adapter(crate::PyTorchToBurnAdapter) // Use adapter to handle PyTorch format + .allow_partial(true); // Allow partial loading due to naming differences + let mut model = Net::new(&device); + let load_result = model.load_from(&mut load_store).unwrap(); + // With the adapter, weights should load correctly + assert!(!load_result.applied.is_empty()); + assert!( + load_result.errors.is_empty(), + "Should have no errors with adapter: {:?}", + load_result.errors + ); + + // Save the model to memory + // Note: format, producer and version are automatically added + let mut save_store = SafetensorsStore::from_bytes(None).metadata("source", "pytorch"); + model.save_into(&mut save_store).unwrap(); + + // Load into a new model + let mut model2 = Net::new(&device); + let mut load_store2 = SafetensorsStore::from_bytes(None); + if let SafetensorsStore::Memory(ref mut p) = load_store2 + && let SafetensorsStore::Memory(ref p_save) = save_store + { + p.set_data(p_save.data().unwrap().as_ref().clone()); + } + + let result = model2.load_from(&mut load_store2).unwrap(); + assert!(!result.applied.is_empty()); + + // Verify both models produce the same output + let input = Tensor::<4>::ones([1, 3, 8, 8], &device); + let output1 = model.forward(input.clone()); + let output2 = model2.forward(input); + + // Check outputs are identical + let output1_data = output1.to_data().to_vec::().unwrap(); + let output2_data = output2.to_data().to_vec::().unwrap(); + + for (a, b) in output1_data.iter().zip(output2_data.iter()) { + assert!((a - b).abs() < 1e-7, "Outputs differ after round trip"); + } +} + +#[test] +#[cfg(all(feature = "std", target_has_atomic = "ptr"))] +fn partial_load_from_pytorch_model() { + let device = Default::default(); + + // Reference the safetensors file from burn-import + let safetensors_path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/safetensors-tests/tests/multi_layer/multi_layer.safetensors" + ); + + // Load only conv1 and norm1 parameters (not fc1) + let mut store = SafetensorsStore::from_file(safetensors_path) + .validate(false) // Disable validation due to shape differences + .allow_partial(true); + + let mut model = Net::new(&device); + + // Save initial fc1 weights for comparison + let _initial_fc1_weight = model.fc1.weight.val().to_data(); + + let result = model.load_from(&mut store).unwrap(); + + // Should load available tensors (with some errors due to shape mismatch) + assert!(!result.applied.is_empty()); + + // fc1 weight should remain unchanged if not in the file + // or should be updated if it is in the file + // This test verifies that partial loading works correctly +} + +#[test] +#[cfg(all(feature = "std", target_has_atomic = "ptr"))] +fn verify_tensor_names_from_pytorch() { + let device = Default::default(); + + // Reference the safetensors file from burn-import + let safetensors_path = concat!( + env!("CARGO_MANIFEST_DIR"), + "/safetensors-tests/tests/multi_layer/multi_layer.safetensors" + ); + + // Create a model and load from PyTorch + let mut model = Net::new(&device); + let mut store = SafetensorsStore::from_file(safetensors_path) + .validate(false) // Disable validation due to shape differences + .allow_partial(true); // Allow partial loading due to naming differences + let result = model.load_from(&mut store).unwrap(); + + // Check that we loaded some tensors (with errors due to shape mismatch) + assert!(!result.applied.is_empty()); + + // Collect tensor names from the model + let views = model.collect(None, None, false); + let tensor_names: Vec = views.iter().map(|v| v.full_path()).collect(); + + // Verify expected tensor names are present + assert!(tensor_names.iter().any(|n| n.contains("conv1"))); + assert!(tensor_names.iter().any(|n| n.contains("norm1"))); + assert!(tensor_names.iter().any(|n| n.contains("fc1"))); +} diff --git a/crates/burn-store/src/safetensors/tests/round_trip.rs b/crates/burn-store/src/safetensors/tests/round_trip.rs new file mode 100644 index 0000000..4dace86 --- /dev/null +++ b/crates/burn-store/src/safetensors/tests/round_trip.rs @@ -0,0 +1,103 @@ +use burn_core as burn; + +use crate::{ModuleSnapshot, SafetensorsStore}; +use burn_core::module::{Module, Param}; +use burn_core::tensor::{Device, Tensor, shape}; +use burn_nn::{Linear, LinearConfig}; + +#[derive(Module, Debug)] +pub(super) struct ComplexModule { + pub encoder: EncoderModule, + pub decoder: DecoderModule, + pub layers: Vec, +} + +#[derive(Module, Debug)] +pub(super) struct EncoderModule { + pub weight: Param>, + pub bias: Param>, + pub norm: Param>, +} + +#[derive(Module, Debug)] +pub(super) struct DecoderModule { + pub projection: Linear, + pub scale: Param>, +} + +impl ComplexModule { + pub fn new(device: &Device) -> Self { + Self { + encoder: EncoderModule { + weight: Param::from_data( + [[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]], + device, + ), + bias: Param::from_data([0.1, 0.2, 0.3], device), + norm: Param::from_data([1.0, 1.0, 1.0], device), + }, + decoder: DecoderModule { + projection: LinearConfig::new(4, 2).with_bias(true).init(device), + scale: Param::from_data([[0.5, 0.5], [0.5, 0.5]], device), + }, + layers: vec![ + LinearConfig::new(3, 4).with_bias(false).init(device), + LinearConfig::new(4, 3).with_bias(true).init(device), + ], + } + } + + pub fn new_zeros(device: &Device) -> Self { + Self { + encoder: EncoderModule { + weight: Param::from_tensor(Tensor::zeros([2, 2, 2], device)), + bias: Param::from_tensor(Tensor::zeros([3], device)), + norm: Param::from_tensor(Tensor::zeros([3], device)), + }, + decoder: DecoderModule { + projection: LinearConfig::new(4, 2).with_bias(true).init(device), + scale: Param::from_tensor(Tensor::zeros([2, 2], device)), + }, + layers: vec![ + LinearConfig::new(3, 4).with_bias(false).init(device), + LinearConfig::new(4, 3).with_bias(true).init(device), + ], + } + } +} + +#[test] +fn complex_module_round_trip() { + let device = Default::default(); + let module1 = ComplexModule::new(&device); + let mut module2 = ComplexModule::new_zeros(&device); + + // Save module1 using new store API + let mut save_store = SafetensorsStore::from_bytes(None); + module1.save_into(&mut save_store).unwrap(); + + // Load into module2 + let mut load_store = SafetensorsStore::from_bytes(None); + if let SafetensorsStore::Memory(ref mut p) = load_store + && let SafetensorsStore::Memory(ref p_save) = save_store + { + // Get Arc and extract data + let data_arc = p_save.data().unwrap(); + p.set_data(data_arc.as_ref().clone()); + } + let result = module2.load_from(&mut load_store).unwrap(); + + assert!(result.is_success()); + assert!(result.applied.len() > 5); + assert_eq!(result.errors.len(), 0); + + // Verify data was imported correctly + let module2_views = module2.collect(None, None, false); + let encoder_weight = module2_views + .iter() + .find(|v| v.full_path() == "encoder.weight") + .unwrap() + .to_data() + .unwrap(); + assert_eq!(encoder_weight.shape, shape![2, 2, 2]); +} diff --git a/crates/burn-store/src/tensor_snapshot.rs b/crates/burn-store/src/tensor_snapshot.rs new file mode 100644 index 0000000..0817659 --- /dev/null +++ b/crates/burn-store/src/tensor_snapshot.rs @@ -0,0 +1,617 @@ +use alloc::rc::Rc; +use alloc::string::String; +use alloc::string::ToString; +use alloc::vec::Vec; +use burn_core::module::ParamId; +use burn_core::tensor::quantization::{QPARAM_ALIGN, QuantParam, params_shape}; +use burn_core::tensor::{Bool, DType, Int, Shape, Tensor, TensorData}; +use half::f16; + +/// Returns the byte size of a quantization parameter type. +// TODO: Add `size_bytes()` method to `QuantParam` in cubecl and use it here. +const fn quant_param_size(param: QuantParam) -> usize { + match param { + QuantParam::F32 => core::mem::size_of::(), + QuantParam::F16 | QuantParam::BF16 => core::mem::size_of::(), + QuantParam::UE8M0 | QuantParam::UE4M3 => core::mem::size_of::(), + } +} + +/// Error type for TensorSnapshot operations +#[derive(Debug, Clone)] +pub enum TensorSnapshotError { + /// I/O error occurred while loading tensor data + IoError(String), + /// Data corruption or invalid format + DataError(String), + /// Panic occurred while loading tensor data + PanicError(String), +} + +impl core::fmt::Display for TensorSnapshotError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + Self::IoError(e) => write!(f, "I/O error: {}", e), + Self::DataError(e) => write!(f, "Data error: {}", e), + Self::PanicError(e) => write!(f, "Panic error: {}", e), + } + } +} + +impl core::error::Error for TensorSnapshotError {} + +/// A lightweight snapshot of a tensor that can lazily produce TensorData. +/// +/// TensorSnapshot stores a cloned tensor internally (which is cheap due to reference counting) +/// and only materializes the actual data when `to_data()` is called. This allows +/// efficient inspection of module structure without the overhead of copying all tensor data. +/// +/// The dtype and shape are cached for efficient access without requiring data materialization, +/// which is particularly useful for serialization formats that need metadata upfront. +pub struct TensorSnapshot { + /// Function to get tensor data when needed (Rc allows cloning) + data_fn: Rc Result>, + /// Data type of the tensor (cached for efficient access) + pub dtype: DType, + /// Shape of the tensor (cached for efficient access) + pub shape: Shape, + /// Path stack representing the module hierarchy + pub path_stack: Option>, + /// Container stack representing the container types at each level + pub container_stack: Option>, + /// Unique identifier for the tensor parameter + pub tensor_id: Option, +} + +impl TensorSnapshot { + /// Create a new tensor snapshot from a float tensor + pub fn from_float( + tensor: &Tensor, + path_stack: Vec, + container_stack: Vec, + tensor_id: ParamId, + ) -> Self { + let dtype = tensor.dtype(); + let shape = tensor.shape(); + let tensor = tensor.clone(); // Clone is cheap (reference counted) + Self { + data_fn: Rc::new(move || Ok(tensor.to_data())), + dtype, + shape, + path_stack: Some(path_stack), + container_stack: Some(container_stack), + tensor_id: Some(tensor_id), + } + } + + /// Create a new tensor snapshot from an int tensor + pub fn from_int( + tensor: &Tensor, + path_stack: Vec, + container_stack: Vec, + tensor_id: ParamId, + ) -> Self { + let dtype = tensor.dtype(); + let shape = tensor.shape(); + let tensor = tensor.clone(); // Clone is cheap (reference counted) + Self { + data_fn: Rc::new(move || Ok(tensor.to_data())), + dtype, + shape, + path_stack: Some(path_stack), + container_stack: Some(container_stack), + tensor_id: Some(tensor_id), + } + } + + /// Create a new tensor snapshot from a bool tensor + pub fn from_bool( + tensor: &Tensor, + path_stack: Vec, + container_stack: Vec, + tensor_id: ParamId, + ) -> Self { + let dtype = tensor.dtype(); + let shape = tensor.shape(); + let tensor = tensor.clone(); // Clone is cheap (reference counted) + Self { + data_fn: Rc::new(move || Ok(tensor.to_data())), + dtype, + shape, + path_stack: Some(path_stack), + container_stack: Some(container_stack), + tensor_id: Some(tensor_id), + } + } + + /// Convert to TensorData (this is where actual data copy happens) + #[cfg(feature = "std")] + pub fn to_data(&self) -> Result { + // Use AssertUnwindSafe since we're working with Rc which is not UnwindSafe + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| (self.data_fn)())).unwrap_or_else( + |_| { + Err(TensorSnapshotError::PanicError( + "Panic occurred while loading tensor data".to_string(), + )) + }, + ) + } + + /// Convert to TensorData (this is where actual data copy happens) + #[cfg(not(feature = "std"))] + pub fn to_data(&self) -> Result { + (self.data_fn)() // Can't catch panics in no-std, do it when core::panic::AssertUnwindSafe is available + } + + /// Get the full path by joining the path stack + pub fn full_path(&self) -> String { + self.path_stack + .as_ref() + .map(|stack| stack.join(".")) + .unwrap_or_default() + } + + /// Get the full container path by joining the container stack + pub fn container_path(&self) -> String { + self.container_stack + .as_ref() + .map(|stack| stack.join(".")) + .unwrap_or_default() + } + + /// Get the module type (last Struct/Enum in the hierarchy) + /// + /// Returns the last user-defined module type, skipping primitive containers + /// like "Vec", "Array". This is useful for determining which user-defined + /// module a tensor belongs to. + /// + /// # Examples + /// - `Linear.weight` → `Some("Struct:Linear")` + /// - `Vec[0].weight` → `Some("Struct:Linear")` + /// - `Linear.bias` (Optional) → `Some("Struct:Linear")` + /// - `Vec[0]` (no module) → `None` + pub fn module_type(&self) -> Option { + self.container_stack.as_ref().and_then(|stack| { + // Find the last user-defined type (Struct: or Enum:) + stack + .iter() + .rev() + .find(|ct| ct.starts_with("Struct:") || ct.starts_with("Enum:")) + .cloned() + }) + } + + /// Get the immediate container type (last in the container stack) + /// + /// Returns the last element in the container stack, which could be a + /// user-defined type ("Struct:", "Enum:") or a collection type ("Vec", "Array"). + /// This is useful for understanding the full container hierarchy. + /// + /// # Examples + /// - `Linear.weight` → `"Struct:Linear"` + /// - `Vec[0].weight` → `"Struct:Linear"` (the Linear, not the Vec) + /// - `Vec[0]` → `"Vec"` + pub fn container_type(&self) -> String { + self.container_stack + .as_ref() + .and_then(|stack| stack.last()) + .cloned() + .unwrap_or_else(|| "Unknown".to_string()) + } + + /// Create a TensorSnapshot from a closure that produces TensorData + /// This is used internally for lazy loading + pub fn from_closure( + data_fn: Rc Result>, + dtype: DType, + shape: Shape, + path_stack: Vec, + container_stack: Vec, + tensor_id: ParamId, + ) -> Self { + Self { + data_fn, + dtype, + shape, + path_stack: Some(path_stack), + container_stack: Some(container_stack), + tensor_id: Some(tensor_id), + } + } + + /// Create a TensorSnapshot from TensorData directly + pub fn from_data( + data: TensorData, + path_stack: Vec, + container_stack: Vec, + tensor_id: ParamId, + ) -> Self { + let dtype = data.dtype; + let shape = data.shape.clone(); + Self { + data_fn: Rc::new(move || Ok(data.clone())), + dtype, + shape, + path_stack: Some(path_stack), + container_stack: Some(container_stack), + tensor_id: Some(tensor_id), + } + } + + /// Get the size of the tensor data in bytes without materializing it. + /// + /// For regular (non-quantized) types, this is simply `shape.product() * dtype.size()`. + /// + /// For quantized types (`QFloat`), this accounts for: + /// - The quantized values (packed according to the quantization scheme) + /// - Alignment padding (values are aligned to 4-byte boundary) + /// - Quantization parameters (scale values appended to the data) + pub fn data_len(&self) -> usize { + const BITS_PER_BYTE: usize = 8; + + let num_elements: usize = self.shape.iter().product(); + + match self.dtype { + DType::QFloat(scheme) => { + // Calculate value bytes using scheme's packing information + let num_storage_elements = num_elements.div_ceil(scheme.num_quants()); + let value_bytes = + num_storage_elements * (scheme.size_bits_stored() / BITS_PER_BYTE); + + // Calculate number of quantization parameters (scales) + let num_params = params_shape(&self.shape, scheme.level).num_elements(); + + let aligned_value_bytes = value_bytes.div_ceil(QPARAM_ALIGN) * QPARAM_ALIGN; + let scale_bytes = num_params * quant_param_size(scheme.param); + + aligned_value_bytes + scale_bytes + } + _ => num_elements * self.dtype.size(), + } + } + + /// Clone the data function for lazy composition + pub fn clone_data_fn(&self) -> Rc Result> { + self.data_fn.clone() + } +} + +impl Clone for TensorSnapshot { + fn clone(&self) -> Self { + // Clone lazily - keep the same data function + Self { + data_fn: self.data_fn.clone(), + dtype: self.dtype, + shape: self.shape.clone(), + path_stack: self.path_stack.clone(), + container_stack: self.container_stack.clone(), + tensor_id: self.tensor_id, + } + } +} + +impl core::fmt::Debug for TensorSnapshot { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("TensorSnapshot") + .field("dtype", &self.dtype) + .field("shape", &self.shape) + .field("path_stack", &self.path_stack) + .field("container_stack", &self.container_stack) + .field("tensor_id", &self.tensor_id) + .finish() + } +} + +#[cfg(all(test, feature = "std"))] +mod tests { + use super::*; + use alloc::string::ToString; + use burn_core::tensor::{BoolStore, Device, shape}; + + #[test] + fn tensor_view_float() { + let device = Default::default(); + let tensor = Tensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + + let snapshot = TensorSnapshot::from_float( + &tensor, + vec!["test".to_string(), "weight".to_string()], + vec!["TestModule".to_string(), "Param".to_string()], + ParamId::new(), + ); + + // Test metadata access without materialization + assert_eq!(snapshot.dtype, DType::F32); + assert_eq!(snapshot.shape, shape![2, 2]); + assert_eq!(snapshot.full_path(), "test.weight"); + assert_eq!(snapshot.container_path(), "TestModule.Param"); + + // Test data materialization + let data = snapshot.to_data().unwrap(); + assert_eq!(data.shape, shape![2, 2]); + assert_eq!(data.dtype, DType::F32); + } + + #[test] + fn tensor_view_int() { + let device = Default::default(); + let tensor = Tensor::<2, Int>::from_data([[1, 2], [3, 4]], &device); + + let snapshot = TensorSnapshot::from_int( + &tensor, + vec!["test".to_string(), "int".to_string()], + vec!["TestModule".to_string(), "Param".to_string()], + ParamId::new(), + ); + + // Test metadata access without materialization + let dtype: DType = device.settings().int_dtype.into(); + assert_eq!(snapshot.dtype, dtype); + assert_eq!(snapshot.shape, shape![2, 2]); + + let data = snapshot.to_data().unwrap(); + assert_eq!(data.shape, shape![2, 2]); + assert_eq!(data.dtype, dtype); + } + + #[test] + fn tensor_view_bool() { + let device = Default::default(); + let tensor = Tensor::<2, Bool>::from_data([[true, false], [false, true]], &device); + + let snapshot = TensorSnapshot::from_bool( + &tensor, + vec!["test".to_string(), "bool".to_string()], + vec!["TestModule".to_string(), "Param".to_string()], + ParamId::new(), + ); + + // Test metadata access without materialization + assert_eq!(snapshot.dtype, DType::Bool(BoolStore::Native)); + assert_eq!(snapshot.shape, shape![2, 2]); + + let data = snapshot.to_data().unwrap(); + assert_eq!(data.shape, shape![2, 2]); + assert_eq!(data.dtype, DType::Bool(BoolStore::Native)); + } + + #[test] + fn data_len() { + let device = Device::default(); + let settings = device.settings(); + + // Test F32 tensor (4 bytes per element) + let tensor_f32 = Tensor::<2>::from_data([[1.0, 2.0], [3.0, 4.0]], &device); + let view_f32 = TensorSnapshot::from_float( + &tensor_f32, + vec!["test".to_string()], + vec!["Module".to_string()], + ParamId::new(), + ); + let dtype: DType = settings.float_dtype.into(); + assert_eq!(view_f32.data_len(), 4 * dtype.size()); // 4 elements * 4 bytes + + // Test int tensor + let tensor_i64 = Tensor::<3, Int>::from_data([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], &device); + let view_i64 = TensorSnapshot::from_int( + &tensor_i64, + vec!["test".to_string()], + vec!["Module".to_string()], + ParamId::new(), + ); + let dtype: DType = settings.int_dtype.into(); + assert_eq!(view_i64.data_len(), 8 * dtype.size()); // 8 elements * 8 bytes (I64) + + // Test Bool tensor (1 byte per element) + let tensor_bool = Tensor::<2, Bool>::from_data([[true, false], [false, true]], &device); + let view_bool = TensorSnapshot::from_bool( + &tensor_bool, + vec!["test".to_string()], + vec!["Module".to_string()], + ParamId::new(), + ); + let dtype: DType = settings.bool_dtype.into(); + assert_eq!(view_bool.data_len(), 4 * dtype.size()); // 4 elements * 1 byte + } + + #[test] + fn from_closure() { + let data = TensorData::from([1.0f32, 2.0, 3.0, 4.0]); + let dtype = data.dtype; + let shape = data.shape.clone(); + + let snapshot = TensorSnapshot::from_closure( + Rc::new(move || Ok(data.clone())), + dtype, + shape.clone(), + vec!["model".to_string(), "layer".to_string()], + vec!["Model".to_string(), "Layer".to_string()], + ParamId::new(), + ); + + // Test metadata access + assert_eq!(snapshot.dtype, DType::F32); + assert_eq!(snapshot.shape, shape![4]); + assert_eq!(snapshot.full_path(), "model.layer"); + assert_eq!(snapshot.data_len(), 16); // 4 * 4 bytes + + // Test data materialization + let materialized = snapshot.to_data().unwrap(); + assert_eq!(materialized.shape, shape![4]); + } + + #[test] + fn from_data() { + let data = TensorData::from([1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]); + let original_dtype = data.dtype; + let original_shape = data.shape.clone(); + + let snapshot = TensorSnapshot::from_data( + data, + vec!["encoder".to_string(), "weight".to_string()], + vec!["Struct:Encoder".to_string(), "Struct:Dense".to_string()], + ParamId::new(), + ); + + // Test metadata + assert_eq!(snapshot.dtype, original_dtype); + assert_eq!(snapshot.shape, original_shape); + assert_eq!(snapshot.full_path(), "encoder.weight"); + assert_eq!(snapshot.container_type(), "Struct:Dense"); + assert_eq!(snapshot.data_len(), 24); // 6 * 4 bytes + + // Test data materialization + let materialized = snapshot.to_data().unwrap(); + assert_eq!(materialized.shape, original_shape); + } + + #[test] + #[cfg(feature = "std")] + fn panic_catching_in_to_data() { + use alloc::rc::Rc; + + // Create a TensorSnapshot with a closure that panics + let snapshot = TensorSnapshot { + data_fn: Rc::new(|| panic!("Test panic in data_fn")), + dtype: DType::F32, + shape: shape![2, 2], + path_stack: Some(vec!["test".to_string()]), + container_stack: Some(vec!["Test".to_string()]), + tensor_id: Some(ParamId::new()), + }; + + // When std is available, to_data should catch the panic and return an error + let result = snapshot.to_data(); + assert!(result.is_err()); + + match result { + Err(TensorSnapshotError::PanicError(msg)) => { + assert!(msg.contains("Panic occurred")); + } + _ => panic!("Expected PanicError with panic message"), + } + } + + #[test] + fn error_propagation_in_closure() { + use alloc::rc::Rc; + + // Create a snapshot with a closure that returns an error + let snapshot = TensorSnapshot::from_closure( + Rc::new(|| Err(TensorSnapshotError::IoError("Simulated IO error".into()))), + DType::F32, + shape![2, 2], + vec!["error_test".into()], + vec![], + ParamId::new(), + ); + + // Should return an error when trying to get data + let result = snapshot.to_data(); + assert!(result.is_err()); + match result { + Err(TensorSnapshotError::IoError(msg)) => { + assert!(msg.contains("Simulated IO error")); + } + _ => panic!("Expected IoError"), + } + } + + #[test] + fn container_type_extraction() { + let device = Default::default(); + let tensor = Tensor::<1>::from_data([1.0, 2.0, 3.0], &device); + + let snapshot = TensorSnapshot::from_float( + &tensor, + vec![ + "model".to_string(), + "layer1".to_string(), + "weight".to_string(), + ], + vec![ + "Struct:Model".to_string(), + "Struct:Conv2d".to_string(), + "Struct:Param".to_string(), + ], + ParamId::new(), + ); + + assert_eq!(snapshot.container_type(), "Struct:Param"); + assert_eq!(snapshot.module_type(), Some("Struct:Param".to_string())); + assert_eq!( + snapshot.container_path(), + "Struct:Model.Struct:Conv2d.Struct:Param" + ); + assert_eq!(snapshot.full_path(), "model.layer1.weight"); + } + + #[test] + fn container_type_vs_module_type() { + let device = Default::default(); + let tensor = Tensor::<1>::from_data([1.0, 2.0, 3.0], &device); + + // Test case 1: Tensor inside a Vec + // container_stack: ["Struct:Model", "Vec", "Struct:Linear"] + let snapshot = TensorSnapshot::from_float( + &tensor, + vec![ + "model".to_string(), + "layers".to_string(), + "0".to_string(), + "weight".to_string(), + ], + vec![ + "Struct:Model".to_string(), + "Vec".to_string(), + "Struct:Linear".to_string(), + ], + ParamId::new(), + ); + + // container_type() returns the last element (Struct:Linear in this case) + assert_eq!(snapshot.container_type(), "Struct:Linear"); + // module_type() also returns Some(Struct:Linear) (skipping Vec) + assert_eq!(snapshot.module_type(), Some("Struct:Linear".to_string())); + + // Test case 2: Tensor that's just in a Vec + // container_stack: ["Vec"] + let snapshot2 = TensorSnapshot::from_float( + &tensor, + vec!["data".to_string(), "0".to_string()], + vec!["Vec".to_string()], + ParamId::new(), + ); + + // container_type() returns Vec + assert_eq!(snapshot2.container_type(), "Vec"); + // module_type() returns None (no Struct/Enum found) + assert_eq!(snapshot2.module_type(), None); + + // Test case 3: Nested collections + // container_stack: ["Struct:Model", "Vec", "Array", "Struct:Linear"] + let snapshot3 = TensorSnapshot::from_float( + &tensor, + vec![ + "model".to_string(), + "layers".to_string(), + "0".to_string(), + "sublayers".to_string(), + "1".to_string(), + "weight".to_string(), + ], + vec![ + "Struct:Model".to_string(), + "Vec".to_string(), + "Array".to_string(), + "Struct:Linear".to_string(), + ], + ParamId::new(), + ); + + // container_type() returns the immediate container + assert_eq!(snapshot3.container_type(), "Struct:Linear"); + // module_type() returns the last Struct/Enum + assert_eq!(snapshot3.module_type(), Some("Struct:Linear".to_string())); + } +} diff --git a/crates/burn-store/src/traits.rs b/crates/burn-store/src/traits.rs new file mode 100644 index 0000000..bb8f11b --- /dev/null +++ b/crates/burn-store/src/traits.rs @@ -0,0 +1,281 @@ +use alloc::boxed::Box; +use alloc::collections::BTreeMap; +use alloc::string::String; +use alloc::vec::Vec; + +use super::applier::Applier; +use super::apply_result::ApplyResult; +use crate::collector::Collector; +use crate::{ModuleAdapter, PathFilter, TensorSnapshot}; +use burn_core::module::Module; + +/// Extension trait for modules that provides tensor storage functionality. +/// +/// This trait provides convenient methods to collect and apply tensor snapshots from any Burn module. +/// Collection operations create lightweight tensor snapshots without immediately copying data. +/// Apply operations apply tensor data from snapshots to the corresponding tensors in the module. +pub trait ModuleSnapshot: Module { + /// Collects tensor snapshots for inspection without copying data. + /// + /// Returns a vector of `TensorSnapshot` objects that can lazily materialize the tensor data. + /// Each `TensorSnapshot` contains the full path accessible via `snapshot.full_path()`. + /// + /// # Arguments + /// + /// * `filter` - An optional [`PathFilter`] to determine which tensors to collect. + /// When `None`, all tensors are collected. + /// * `adapter` - Optional adapter to transform tensors based on container types. + /// Applied to all collected tensors before returning. + /// * `skip_enum_variants` - Skip enum variant names when building paths. + /// When true, paths will not include enum variant names (e.g., "feature.weight" + /// instead of "feature.BaseConv.weight"). Useful when exporting to formats + /// like PyTorch/SafeTensors that don't use enum variants. + fn collect( + &self, + filter: Option, + adapter: Option>, + skip_enum_variants: bool, + ) -> Vec { + let mut collector = Collector::new(filter, adapter, skip_enum_variants); + self.visit(&mut collector); + collector.into_tensors() + } + + /// Applies tensor snapshots to the module. + /// + /// This is the primary apply method that applies tensor data from `TensorSnapshot`s + /// to the corresponding tensors in the module. The snapshots are typically obtained + /// from `collect()` or loaded from storage. + /// + /// # Arguments + /// + /// * `snapshots` - A vector of TensorSnapshot objects + /// * `filter` - An optional [`PathFilter`] to determine which tensors to apply. + /// When `None`, all available tensors are applied. + /// * `adapter` - Optional adapter to transform tensors based on container types + /// * `skip_enum_variants` - Skip enum variant names when matching tensor paths + /// + /// # Returns + /// + /// An [`ApplyResult`] containing information about applied, skipped, missing, + /// and unused tensors, as well as any errors encountered. + /// + /// # Examples + /// + /// ```rust,ignore + /// use burn_store::PathFilter; + /// + /// // Apply all tensors + /// let result = model.apply(snapshots, None, None, false); + /// + /// // Apply only encoder tensors + /// let filter = PathFilter::new().with_regex(r"^encoder\..*"); + /// let result = model.apply(snapshots, Some(filter), None, false); + /// + /// // Apply with complex filter + /// let filter = PathFilter::new() + /// .with_regex(r"^encoder\..*") + /// .with_regex(r"^decoder\..*") + /// .with_full_path("head.weight"); + /// let result = model.apply(snapshots, Some(filter), None, false); + /// + /// // Apply with enum variant skipping (for PyTorch models) + /// let result = model.apply(snapshots, None, None, true); + /// ``` + fn apply( + &mut self, + snapshots: Vec, + filter: Option, + adapter: Option>, + skip_enum_variants: bool, + ) -> ApplyResult + where + Self: Sized, + { + let mut applier = Applier::new(snapshots, filter, adapter, skip_enum_variants); + + // Use unsafe to avoid cloning the entire module, which would double the memory usage + // We read the module out, map it, then write it back + // See https://github.com/tracel-ai/burn/issues/3754 + unsafe { + // Read the module out of self (moves it, leaving self in undefined state) + let module = core::ptr::read(self as *const Self); + + // Map the module to create a new one with updated tensors + let new_module = module.map(&mut applier); + + // Write the new module back to self + core::ptr::write(self as *mut Self, new_module); + } + + applier.into_result() + } + + /// Saves tensor snapshots into a [`ModuleStore`]. + /// + /// This method allows using a `ModuleStore` implementation to handle the + /// collection and writing logic in a configurable way. + /// + /// # Arguments + /// + /// * `store` - A mutable reference to a [`ModuleStore`] that will collect and save the tensors + fn save_into

(&self, store: &mut P) -> Result<(), P::Error> + where + P: ModuleStore, + { + store.collect_from(self) + } + + /// Loads tensor data from a [`ModuleStore`]. + /// + /// This method allows using a `ModuleStore` implementation to handle the + /// loading and application logic in a configurable way. + /// + /// # Arguments + /// + /// * `store` - A mutable reference to a [`ModuleStore`] that will load and apply tensors + fn load_from

(&mut self, store: &mut P) -> Result + where + P: ModuleStore, + { + store.apply_to(self) + } +} + +/// A trait for handling module storage operations. +/// +/// `ModuleStore` provides a unified interface for saving and loading module +/// tensor data with support for various storage formats and advanced features like filtering, +/// remapping, and metadata handling. +pub trait ModuleStore { + /// The error type that can be returned during storage operations. + /// + /// This should be a format-specific error type that provides detailed + /// information about what went wrong (e.g., I/O errors, format violations, + /// unsupported tensor types). + type Error: core::fmt::Debug + core::fmt::Display; + + /// Collect tensor data from a module and store it to storage. + /// + /// This method traverses the module structure, collects all tensor data + /// according to the store's configuration (filters, remapping, etc.), + /// and writes it to the underlying storage. + /// + /// # Arguments + /// + /// * `module` - The module to collect tensor data from. The module must + /// implement `ModuleSnapshot` to provide tensor access. + /// + /// # Returns + /// + /// * `Ok(())` - If all tensors were successfully collected and stored + /// * `Err(Self::Error)` - If an error occurred during collection or writing + fn collect_from(&mut self, module: &M) -> Result<(), Self::Error>; + + /// Load stored tensor data and apply it to a module. + /// + /// This method reads tensor data from storage and applies it to the provided + /// module. The operation is flexible and can handle partial matches, missing + /// tensors, and extra tensors in the storage. + /// + /// # Arguments + /// + /// * `module` - The module to apply tensor data to. The module must + /// implement `ModuleSnapshot` to allow tensor updates. + /// + /// # Returns + /// + /// * `Ok(ApplyResult)` - Detailed information about the apply operation: + /// - `applied`: List of successfully applied tensor names + /// - `missing`: Tensors expected by the module but not found in storage + /// - `skipped`: Tensors in storage that were not applied (filtered or not needed) + /// - `errors`: Non-critical errors that occurred during apply + /// * `Err(Self::Error)` - If a critical error prevented the apply operation + fn apply_to(&mut self, module: &mut M) -> Result; + + /// Get a single tensor snapshot by name. + /// + /// This method provides direct access to individual tensors in storage without + /// requiring a module. The returned `TensorSnapshot` uses lazy loading - tensor + /// data is only materialized when `to_data()` is called. + /// + /// **Note:** Key remapping is applied, so use the remapped name if configured. + /// Filters are NOT applied - use `apply_to()` for filtered loading. + /// + /// Results are cached after the first call for efficient repeated access. + /// + /// # Arguments + /// + /// * `name` - The tensor name/path (e.g., "encoder.layer1.weight") + /// + /// # Returns + /// + /// * `Ok(Some(&TensorSnapshot))` - Reference to the tensor snapshot if found + /// * `Ok(None)` - If no tensor with that name exists + /// * `Err(Self::Error)` - If an error occurred accessing storage + /// + /// # Example + /// + /// ```rust,ignore + /// let mut store = BurnpackStore::from_file("model.bpk"); + /// if let Some(snapshot) = store.get_snapshot("encoder.weight")? { + /// println!("Shape: {:?}", snapshot.shape); + /// println!("Dtype: {:?}", snapshot.dtype); + /// let data = snapshot.to_data()?; // Lazy load + /// } + /// ``` + fn get_snapshot(&mut self, name: &str) -> Result, Self::Error>; + + /// Get all tensor snapshots from storage as an ordered map. + /// + /// This method returns all tensors in storage as lazy-loading snapshots, + /// organized in a `BTreeMap` for efficient lookup by name. The map preserves + /// alphabetical ordering of tensor names. + /// + /// **Note:** This returns ALL tensors in storage, regardless of any filter + /// settings. Filters are only applied during `apply_to()`. Key remapping + /// IS applied, so tensor names reflect any configured remapping. + /// + /// Results are cached after the first call for efficient repeated access. + /// + /// # Returns + /// + /// * `Ok(&BTreeMap)` - Reference to all tensor snapshots + /// * `Err(Self::Error)` - If an error occurred accessing storage + /// + /// # Example + /// + /// ```rust,ignore + /// let mut store = SafetensorsStore::from_file("model.safetensors"); + /// let snapshots = store.get_all_snapshots()?; + /// for (name, snapshot) in snapshots { + /// println!("{}: {:?}", name, snapshot.shape); + /// } + /// ``` + fn get_all_snapshots(&mut self) -> Result<&BTreeMap, Self::Error>; + + /// Get all tensor names/keys in storage. + /// + /// This method returns the names of all tensors in storage. + /// Useful for inspecting storage contents or checking if specific tensors exist. + /// + /// **Note:** Returns ALL tensor names regardless of filter settings. + /// Key remapping IS applied, so names reflect any configured remapping. + /// + /// # Returns + /// + /// * `Ok(Vec)` - All tensor names in storage + /// * `Err(Self::Error)` - If an error occurred accessing storage + /// + /// # Example + /// + /// ```rust,ignore + /// let mut store = PytorchStore::from_file("model.pth"); + /// let keys = store.keys()?; + /// println!("Tensors in file: {:?}", keys); + /// ``` + fn keys(&mut self) -> Result, Self::Error>; +} + +// Blanket implementation for all modules +impl ModuleSnapshot for M {} diff --git a/crates/burn-tch/Cargo.toml b/crates/burn-tch/Cargo.toml new file mode 100644 index 0000000..6257242 --- /dev/null +++ b/crates/burn-tch/Cargo.toml @@ -0,0 +1,38 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science"] +description = "LibTorch backend for the Burn framework using the tch bindings." +documentation = "https://docs.rs/burn-tch" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "data"] +license.workspace = true +name = "burn-tch" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-tch" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["std"] +std = ["burn-backend/std"] +doc = ["tch/doc-only"] +tracing = [ + "burn-backend/tracing", +] + +[dependencies] +burn-backend = { workspace = true } + +libc = { workspace = true } +log = { workspace = true } +tch = { workspace = true, features = ["download-libtorch"] } +torch-sys = { workspace = true } # for build script lib dir detection + +[build-dependencies] +cc = "1.2.61" + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-tch/LICENSE-APACHE b/crates/burn-tch/LICENSE-APACHE new file mode 120000 index 0000000..1cd601d --- /dev/null +++ b/crates/burn-tch/LICENSE-APACHE @@ -0,0 +1 @@ +../../LICENSE-APACHE \ No newline at end of file diff --git a/crates/burn-tch/LICENSE-MIT b/crates/burn-tch/LICENSE-MIT new file mode 120000 index 0000000..b2cfbdc --- /dev/null +++ b/crates/burn-tch/LICENSE-MIT @@ -0,0 +1 @@ +../../LICENSE-MIT \ No newline at end of file diff --git a/crates/burn-tch/README.md b/crates/burn-tch/README.md new file mode 100644 index 0000000..559a18e --- /dev/null +++ b/crates/burn-tch/README.md @@ -0,0 +1,246 @@ +# Burn Torch Backend + +[Burn](https://github.com/tracel-ai/burn) Torch backend + +[![Current Crates.io Version](https://img.shields.io/crates/v/burn-tch.svg)](https://crates.io/crates/burn-tch) +[![license](https://shields.io/badge/license-MIT%2FApache--2.0-blue)](https://github.com/tracel-ai/burn-tch/blob/master/README.md) + +This crate provides a Torch backend for [Burn](https://github.com/tracel-ai/burn) utilizing the +[`tch-rs`](https://github.com/LaurentMazare/tch-rs) crate, which offers a Rust interface to the +[PyTorch](https://pytorch.org/) C++ API. + +The backend supports CPU (multithreaded), [CUDA](https://pytorch.org/docs/stable/notes/cuda.html) +(multiple GPUs), and [MPS](https://pytorch.org/docs/stable/notes/mps.html) devices (MacOS). + +## Installation + +[`tch-rs`](https://github.com/LaurentMazare/tch-rs) requires the C++ PyTorch library (LibTorch) to +be available on your system. + +By default, the CPU distribution is installed for LibTorch v2.9.0 as required by `tch-rs`. + +

+CUDA + +To install the latest compatible CUDA distribution, set the `TORCH_CUDA_VERSION` environment +variable before the `tch-rs` dependency is retrieved with `cargo`. + +```shell +export TORCH_CUDA_VERSION=cu128 +``` + +On Windows: + +```powershell +$Env:TORCH_CUDA_VERSION = "cu128" +``` + +> Note: `tch` doesn't expose the downloaded libtorch directory on Windows when using the automatic +> download feature, so the `torch_cuda.dll` cannot be detected properly during build. In this case, +> you can set the `LIBTORCH` environment variable to point to the `libtorch/` folder in `torch-sys` +> `OUT_DIR` (or move the downloaded lib to a different folder and point to it). + +For example, running the validation sample for the first time could be done with the following +commands: + +```shell +export TORCH_CUDA_VERSION=cu128 +cargo run --bin cuda --release +``` + +**Important:** make sure your driver version is compatible with the selected CUDA version. A CUDA +Toolkit installation is not required since LibTorch ships with the appropriate CUDA runtimes. Having +the latest driver version is recommended, but you can always take a look at the +[toolkit driver version table](https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html#id4) +or +[minimum required driver version](https://docs.nvidia.com/deploy/cuda-compatibility/index.html#minor-version-compatibility) +(limited feature-set, might not work with all operations). + +

+ +Once your installation is complete, you should be able to build/run your project. You can also +validate your installation by running the appropriate `cpu`, `cuda` or `mps` sample as below. + +```shell +cargo run --bin cpu --release +cargo run --bin cuda --release +cargo run --bin mps --release +``` + +_Note: no MPS distribution is available for automatic download at this time, please check out the +[manual instructions](#metal-mps)._ + +### Manual Download + +To install `tch-rs` with a different LibTorch distribution, you will have to manually download the +desired LibTorch distribution. The instructions are detailed in the sections below for each +platform. + +| Compute Platform | CPU | GPU | Linux | MacOS | Windows | Android | iOS | WASM | +| :------------------------ | :----------------------------: | :-: | :---: | :---: | :-----: | :-----: | :-: | :--: | +| [CPU](#cpu) | Yes | No | Yes | Yes | Yes | Yes | Yes | No | +| [CUDA](#cuda) | Yes [[1]](#cpu-sup) | Yes | Yes | No | Yes | No | No | No | +| [Metal (MPS)](#metal-mps) | No | Yes | No | Yes | No | No | No | No | +| Vulkan | Yes | Yes | Yes | Yes | Yes | Yes | No | No | + +[1] The LibTorch CUDA distribution also comes with CPU support. + +#### CPU + +
+🐧 Linux + +First, download the LibTorch CPU distribution. + +```shell +wget -O libtorch.zip https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.9.0%2Bcpu.zip +unzip libtorch.zip +``` + +Then, point to that installation using the `LIBTORCH` and `LD_LIBRARY_PATH` environment variables +before building `burn-tch` or a crate which depends on it. + +```shell +export LIBTORCH=/absolute/path/to/libtorch/ +export LD_LIBRARY_PATH=/absolute/path/to/libtorch/lib:$LD_LIBRARY_PATH +``` + +

+ +
+🍎 Mac + +First, download the LibTorch CPU distribution. + +```shell +wget -O libtorch.zip https://download.pytorch.org/libtorch/cpu/libtorch-macos-arm64-2.9.0.zip +unzip libtorch.zip +``` + +Then, point to that installation using the `LIBTORCH` and `DYLD_LIBRARY_PATH` environment variables +before building `burn-tch` or a crate which depends on it. + +```shell +export LIBTORCH=/absolute/path/to/libtorch/ +export DYLD_LIBRARY_PATH=/absolute/path/to/libtorch/lib:$DYLD_LIBRARY_PATH +``` + +

+ +
+🪟 Windows + +First, download the LibTorch CPU distribution. + +```powershell +wget https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-2.9.0%2Bcpu.zip -OutFile libtorch.zip +Expand-Archive libtorch.zip +``` + +Then, set the `LIBTORCH` environment variable and append the library to your path as with the +PowerShell commands below before building `burn-tch` or a crate which depends on it. + +```powershell +$Env:LIBTORCH = "/absolute/path/to/libtorch/" +$Env:Path += ";/absolute/path/to/libtorch/" +``` + +

+ +#### CUDA + +LibTorch 2.9.0 currently includes binary distributions with CUDA 12.6, 12.8 or 13.0 runtimes. The +manual installation instructions are detailed below for CUDA 12.6, but can be applied to the other +CUDA versions by replacing `cu126` with the corresponding version string (e.g., `cu130`). + +
+🐧 Linux + +First, download the LibTorch CUDA 12.6 distribution. + +```shell +wget -O libtorch.zip https://download.pytorch.org/libtorch/cu126/libtorch-shared-with-deps-2.9.0%2Bcu126.zip +unzip libtorch.zip +``` + +Then, point to that installation using the `LIBTORCH` and `LD_LIBRARY_PATH` environment variables +before building `burn-tch` or a crate which depends on it. + +```shell +export LIBTORCH=/absolute/path/to/libtorch/ +export LD_LIBRARY_PATH=/absolute/path/to/libtorch/lib:$LD_LIBRARY_PATH +``` + +**Note:** make sure your CUDA installation is in your `PATH` and `LD_LIBRARY_PATH`. + +

+ +
+🪟 Windows + +First, download the LibTorch CUDA 12.6 distribution. + +```powershell +wget https://download.pytorch.org/libtorch/cu126/libtorch-win-shared-with-deps-2.9.0%2Bcu126.zip -OutFile libtorch.zip +Expand-Archive libtorch.zip +``` + +Then, set the `LIBTORCH` environment variable and append the library to your path as with the +PowerShell commands below before building `burn-tch` or a crate which depends on it. + +```powershell +$Env:LIBTORCH = "/absolute/path/to/libtorch/" +$Env:Path += ";/absolute/path/to/libtorch/" +``` + +

+ +#### Metal (MPS) + +There is no official LibTorch distribution with MPS support at this time, so the easiest alternative +is to use a PyTorch installation. This requires a Python installation. + +_Note: MPS acceleration is available on MacOS 12.3+._ + +```shell +pip install torch==2.9.0 numpy==1.26.4 setuptools +export LIBTORCH_USE_PYTORCH=1 +export DYLD_LIBRARY_PATH=/path/to/pytorch/lib:$DYLD_LIBRARY_PATH +``` + +**Note:** if `venv` is used, it should be activated during coding and building, or the compiler may +not work properly. + +## Example Usage + +For a simple example, check out any of the test programs in [`src/bin/`](./src/bin/). Each program +sets the device to use and performs a simple element-wise addition. + +For a more complete example using the `tch` backend, take a loot at the +[Burn mnist example](https://github.com/tracel-ai/burn/tree/main/examples/mnist). + +## Too many environment variables? + +Try `.cargo/config.toml` ([cargo book](https://doc.rust-lang.org/cargo/reference/config.html#env)). + +Instead of setting the environments in your shell, you can manually add them to your +`.cargo/config.toml`: + +```toml +[env] +LD_LIBRARY_PATH = "/absolute/path/to/libtorch/lib" +LIBTORCH = "/absolute/path/to/libtorch/libtorch" +``` + +Or use bash commands below: + +```bash +mkdir .cargo +cat < .cargo/config.toml +[env] +LD_LIBRARY_PATH = "/absolute/path/to/libtorch/lib:$LD_LIBRARY_PATH" +LIBTORCH = "/absolute/path/to/libtorch/libtorch" +EOF +``` + +This will automatically include the old `LD_LIBRARY_PATH` value in the new one. diff --git a/crates/burn-tch/build.rs b/crates/burn-tch/build.rs new file mode 100644 index 0000000..7fb432b --- /dev/null +++ b/crates/burn-tch/build.rs @@ -0,0 +1,243 @@ +// The LIBTORCH environment variable can be used to specify the directory +// where libtorch has been installed. +// When not specified this script downloads the cpu version for libtorch +// and extracts it in OUT_DIR. +// +// On Linux, the TORCH_CUDA_VERSION environment variable can be used, +// like 9.0, 90, or cu90 to specify the version of CUDA to use for libtorch. + +use std::path::{Path, PathBuf}; +use std::{env, fs}; + +const PYTHON_PRINT_PYTORCH_DETAILS: &str = r" +import torch +from torch.utils import cpp_extension +print('LIBTORCH_VERSION:', torch.__version__.split('+')[0]) +print('LIBTORCH_CXX11:', torch._C._GLIBCXX_USE_CXX11_ABI) +for include_path in cpp_extension.include_paths(): + print('LIBTORCH_INCLUDE:', include_path) +for library_path in cpp_extension.library_paths(): + print('LIBTORCH_LIB:', library_path) +"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Os { + Linux, + Macos, + Windows, +} + +#[allow(dead_code)] +#[derive(Debug, Clone)] +struct SystemInfo { + os: Os, + cxx11_abi: String, + libtorch_include_dirs: Vec, + libtorch_lib_dir: PathBuf, +} + +fn env_var_rerun(name: &str) -> Result { + println!("cargo:rerun-if-env-changed={name}"); + env::var(name) +} + +impl SystemInfo { + fn new() -> Option { + let os = match env::var("CARGO_CFG_TARGET_OS") + .expect("Unable to get TARGET_OS") + .as_str() + { + "linux" => Os::Linux, + "windows" => Os::Windows, + "macos" => Os::Macos, + os => panic!("unsupported TARGET_OS '{os}'"), + }; + // Locate the currently active Python binary, similar to: + // https://github.com/PyO3/maturin/blob/243b8ec91d07113f97a6fe74d9b2dcb88086e0eb/src/target.rs#L547 + let python_interpreter = match os { + Os::Windows => PathBuf::from("python.exe"), + Os::Linux | Os::Macos => { + if env::var_os("VIRTUAL_ENV").is_some() { + PathBuf::from("python") + } else { + PathBuf::from("python3") + } + } + }; + let mut libtorch_include_dirs = vec![]; + let mut libtorch_lib_dir = None; + let cxx11_abi = if env_var_rerun("LIBTORCH_USE_PYTORCH").is_ok() { + let output = std::process::Command::new(&python_interpreter) + .arg("-c") + .arg(PYTHON_PRINT_PYTORCH_DETAILS) + .output() + .expect("error running python interpreter"); + let mut cxx11_abi = None; + for line in String::from_utf8_lossy(&output.stdout).lines() { + match line.strip_prefix("LIBTORCH_CXX11: ") { + Some("True") => cxx11_abi = Some("1".to_owned()), + Some("False") => cxx11_abi = Some("0".to_owned()), + _ => {} + } + if let Some(path) = line.strip_prefix("LIBTORCH_INCLUDE: ") { + libtorch_include_dirs.push(PathBuf::from(path)) + } + if let Some(path) = line.strip_prefix("LIBTORCH_LIB: ") { + libtorch_lib_dir = Some(PathBuf::from(path)) + } + } + match cxx11_abi { + Some(cxx11_abi) => cxx11_abi, + None => panic!("no cxx11 abi returned by python {output:?}"), + } + } else { + let libtorch = Self::prepare_libtorch_dir(os)?; + let includes = env_var_rerun("LIBTORCH_INCLUDE") + .map(PathBuf::from) + .unwrap_or_else(|_| libtorch.clone()); + let lib = env_var_rerun("LIBTORCH_LIB") + .map(PathBuf::from) + .unwrap_or_else(|_| libtorch.clone()); + libtorch_include_dirs.push(includes.join("include")); + libtorch_include_dirs.push(includes.join("include/torch/csrc/api/include")); + if lib.ends_with("lib") { + // DEP_TCH_LIBTORCH_LIB might already point to /lib + libtorch_lib_dir = Some(lib); + } else { + libtorch_lib_dir = Some(lib.join("lib")); + } + env_var_rerun("LIBTORCH_CXX11_ABI").unwrap_or_else(|_| "1".to_owned()) + }; + let libtorch_lib_dir = libtorch_lib_dir?; + Some(Self { + os, + cxx11_abi, + libtorch_include_dirs, + libtorch_lib_dir, + }) + } + + fn check_system_location(os: Os) -> Option { + match os { + Os::Linux => Path::new("/usr/lib/libtorch.so") + .exists() + .then(|| PathBuf::from("/usr")), + _ => None, + } + } + + fn prepare_libtorch_dir(os: Os) -> Option { + if let Ok(libtorch) = env_var_rerun("DEP_TCH_LIBTORCH_LIB") { + Some(PathBuf::from(libtorch)) + } else if let Ok(libtorch) = env_var_rerun("LIBTORCH") { + Some(PathBuf::from(libtorch)) + } else if let Some(pathbuf) = Self::check_system_location(os) { + Some(pathbuf) + } else { + check_out_dir() + } + } + + fn make(&self, use_cuda: bool, use_hip: bool) { + let cuda_dependency = if use_cuda || use_hip { + "src/cuda_hack/dummy_cuda_dependency.cpp" + } else { + "src/cuda_hack/fake_cuda_dependency.cpp" + }; + println!("cargo:rerun-if-changed={cuda_dependency}"); + + match self.os { + Os::Linux | Os::Macos => { + cc::Build::new() + .cpp(true) + .pic(true) + .warnings(false) + .includes(&self.libtorch_include_dirs) + .flag(format!("-Wl,-rpath={}", self.libtorch_lib_dir.display())) + .flag("-std=c++17") + .flag(format!("-D_GLIBCXX_USE_CXX11_ABI={}", self.cxx11_abi)) + .files(&[cuda_dependency]) + .compile("burn-tch"); + } + Os::Windows => { + cc::Build::new() + .cpp(true) + .pic(true) + .warnings(false) + .includes(&self.libtorch_include_dirs) + .flag("/std:c++17") + .files(&[cuda_dependency]) + .compile("burn-tch"); + } + }; + } + + fn make_cpu() { + let cuda_dependency = "src/cuda_hack/fake_cuda_dependency.cpp"; + println!("cargo:rerun-if-changed={cuda_dependency}"); + + let os = env::var("CARGO_CFG_TARGET_OS").expect("Unable to get TARGET_OS"); + + match os.as_str() { + "windows" => { + cc::Build::new() + .cpp(true) + .pic(true) + .warnings(false) + .flag("/std:c++17") + .files(&[cuda_dependency]) + .compile("burn-tch"); + } + _ => { + cc::Build::new() + .cpp(true) + .pic(true) + .warnings(false) + .flag("-std=c++17") + .files(&[cuda_dependency]) + .compile("tch"); + } + }; + } +} + +fn check_out_dir() -> Option { + let out_dir = env_var_rerun("OUT_DIR").ok()?; + let libtorch_dir = PathBuf::from(out_dir).join("libtorch"); + libtorch_dir.exists().then_some(libtorch_dir) +} + +fn main() { + let system_info = SystemInfo::new(); + let out_dir = env_var_rerun("OUT_DIR").expect("Failed to get out dir"); + + let mut gpu_found = false; + let found_dir = system_info.is_some(); + if let Some(system_info) = &system_info { + let si_lib = &system_info.libtorch_lib_dir; + let use_cuda = + si_lib.join("libtorch_cuda.so").exists() || si_lib.join("torch_cuda.dll").exists(); + let use_hip = + si_lib.join("libtorch_hip.so").exists() || si_lib.join("torch_hip.dll").exists(); + + system_info.make(use_cuda, use_hip); + gpu_found = use_cuda || use_hip; + } else { + SystemInfo::make_cpu(); + } + let check_file = PathBuf::from(out_dir).join("tch_gpu_check.rs"); + if gpu_found { + fs::write(check_file, "#[allow(clippy::no_effect)]\n()").unwrap(); + } else { + let message = if !found_dir { + r#"Could not find libtorch dir. + + If you are trying to use the automatically downloaded version, the path is not directly available on Windows. Instead, try setting the `LIBTORCH` environment variable for the manual download instructions. + + If the library has already been downloaded in the torch-sys OUT_DIR, you can point the variable to this path (or move the downloaded lib and point to it)."# + } else { + "No libtorch_cuda or libtorch_hip found. Download the GPU version of libtorch to use a GPU device" + }; + fs::write(check_file, format!("panic!(\"{message}\")")).unwrap(); + } +} diff --git a/crates/burn-tch/src/backend.rs b/crates/burn-tch/src/backend.rs new file mode 100644 index 0000000..43d675e --- /dev/null +++ b/crates/burn-tch/src/backend.rs @@ -0,0 +1,184 @@ +use crate::IntoKind; + +use super::TchTensor; +use burn_backend::backend::{Backend, BackendTypes, DeviceId, DeviceOps, ExecutionError}; +use burn_backend::ops::IntTensorOps; +use burn_backend::{BoolStore, DType, DeviceSettings}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +/// The device struct when using the `tch` backend. +/// +/// Note that you need to provide the device index when using Cuda. +/// +/// # Example +/// +/// ```no_run +/// use burn_tch::LibTorchDevice; +/// +/// let device_gpu_1 = LibTorchDevice::Cuda(0); // First GPU +/// let device_gpu_2 = LibTorchDevice::Cuda(1); // Second GPU +/// let device_cpu = LibTorchDevice::Cpu; // CPU +/// let device_mps = LibTorchDevice::Mps; // Metal Performance Shaders +/// let device_vulkan = LibTorchDevice::Vulkan; // Vulkan +/// ``` +#[derive(Default)] +pub enum LibTorchDevice { + /// CPU device. + #[default] + Cpu, + + /// Cuda device with the given index. The index is the index of the Cuda device in the list of + /// all Cuda devices found on the system. + Cuda(usize), + + /// Metal Performance Shaders device. + Mps, + + /// Vulkan device. + Vulkan, +} + +impl From for tch::Device { + #[allow( + unreachable_code, + reason = "CUDA branch always panics if the library is missing" + )] + fn from(device: LibTorchDevice) -> Self { + match device { + LibTorchDevice::Cpu => tch::Device::Cpu, + LibTorchDevice::Cuda(_num) => { + include!(concat!(env!("OUT_DIR"), "/tch_gpu_check.rs")); + tch::Device::Cuda(_num) + } + LibTorchDevice::Mps => tch::Device::Mps, + LibTorchDevice::Vulkan => tch::Device::Vulkan, + } + } +} + +impl From for LibTorchDevice { + fn from(device: tch::Device) -> Self { + match device { + tch::Device::Cpu => LibTorchDevice::Cpu, + tch::Device::Cuda(num) => LibTorchDevice::Cuda(num), + tch::Device::Mps => LibTorchDevice::Mps, + tch::Device::Vulkan => LibTorchDevice::Vulkan, + } + } +} + +impl burn_backend::Device for LibTorchDevice { + fn from_id(device_id: DeviceId) -> Self { + match device_id.type_id { + 0 => Self::Cuda(device_id.index_id as usize), + 1 => Self::Mps, + 2 => Self::Cpu, + 3 => Self::Vulkan, + _ => LibTorchDevice::Cpu, + } + } + + fn to_id(&self) -> DeviceId { + match self { + LibTorchDevice::Cuda(index) => DeviceId::new(0, *index as u16), + LibTorchDevice::Mps => DeviceId::new(1, 0), + LibTorchDevice::Cpu => DeviceId::new(2, 0), + LibTorchDevice::Vulkan => DeviceId::new(3, 0), + } + } +} + +impl DeviceOps for LibTorchDevice { + fn defaults(&self) -> DeviceSettings { + DeviceSettings::new( + DType::F32, + DType::I64, + DType::Bool(BoolStore::Native), + Default::default(), + ) + } +} + +/// Tensor backend that uses `LibTorch` with the [tch] crate for executing tensor operations. +/// +/// This backend is compatible with a wide range of hardwares ranging from CPUs to GPUs, but +/// requires `LibTorch` to be installed correctly. The CPU version can be downloaded +/// automatically and the CUDA version as well by setting the `TORCH_CUDA_VERSION` environment +/// variable. For more complex configurations, check out the manual installation for +/// [burn-tch](https://github.com/tracel-ai/burn/tree/main/crates/burn-tch). +/// +/// Refer to the [tch] crate for more information. +#[derive(Clone, Copy, Default, Debug)] +pub struct LibTorch; + +impl BackendTypes for LibTorch { + type Device = LibTorchDevice; + + type FloatTensorPrimitive = TchTensor; + type IntTensorPrimitive = TchTensor; + type BoolTensorPrimitive = TchTensor; + type QuantizedTensorPrimitive = TchTensor; + + type GraphPrimitive = burn_backend::GraphUnsupported; +} + +impl Backend for LibTorch { + fn seed(_device: &Self::Device, seed: u64) { + tch::manual_seed(seed as i64); + } + + fn ad_enabled(_device: &Self::Device) -> bool { + false + } + + fn name(device: &Self::Device) -> String { + match device { + LibTorchDevice::Cpu => "libtorch", + LibTorchDevice::Cuda(_) => "libtorch", + LibTorchDevice::Mps => "libtorch", + LibTorchDevice::Vulkan => "libtorch", + } + .to_string() + } + + fn sync(device: &Self::Device) -> Result<(), ExecutionError> { + match device { + LibTorchDevice::Cpu => (), + LibTorchDevice::Cuda(index) => { + tch::Cuda::synchronize(*index as i64); + } + _ => { + let int_dtype = device.defaults().int_dtype; + // When there is no explicit way to synchronize, we write and read one value to sync + burn_backend::read_sync(Self::int_into_data(Self::int_zeros( + [1].into(), + device, + int_dtype, + ))) + .unwrap(); + } + }; + + Ok(()) + } + + fn dtype_usage( + _device: &Self::Device, + dtype: burn_backend::DType, + ) -> burn_backend::DTypeUsageSet { + if dtype.try_into_kind().is_ok() { + burn_backend::DTypeUsage::general() + } else { + burn_backend::DTypeUsageSet::empty() + } + } + + fn device_count(type_id: u16) -> usize { + match type_id { + 0 => tch::Cuda::device_count() as usize, + _ => 1, + } + } + + fn flush(_device: &Self::Device) {} +} diff --git a/crates/burn-tch/src/bin/cpu.rs b/crates/burn-tch/src/bin/cpu.rs new file mode 100644 index 0000000..15df92e --- /dev/null +++ b/crates/burn-tch/src/bin/cpu.rs @@ -0,0 +1,14 @@ +use burn_backend::{TensorMetadata, ops::FloatTensorOps}; +use burn_tch::{LibTorch, LibTorchDevice}; + +fn main() { + type B = LibTorch; + let device = LibTorchDevice::Cpu; + + // Creation of two tensors, the first with explicit values and the second one with ones, with the same shape as the first + let tensor_1 = B::float_from_data([[2f32, 3.], [4., 5.]].into(), &device); + let tensor_2 = B::float_ones(tensor_1.shape(), &device, tensor_1.dtype().into()); + + // Print the element-wise addition of the two tensors. + println!("{}", B::float_add(tensor_1, tensor_2)); +} diff --git a/crates/burn-tch/src/bin/cuda.rs b/crates/burn-tch/src/bin/cuda.rs new file mode 100644 index 0000000..8be1c39 --- /dev/null +++ b/crates/burn-tch/src/bin/cuda.rs @@ -0,0 +1,19 @@ +use burn_backend::{TensorMetadata, ops::FloatTensorOps}; +use burn_tch::{LibTorch, LibTorchDevice}; + +fn main() { + assert!( + tch::utils::has_cuda(), + "Could not detect valid CUDA configuration" + ); + + type B = LibTorch; + let device = LibTorchDevice::Cuda(0); + + // Creation of two tensors, the first with explicit values and the second one with ones, with the same shape as the first + let tensor_1 = B::float_from_data([[2f32, 3.], [4., 5.]].into(), &device); + let tensor_2 = B::float_ones(tensor_1.shape(), &device, tensor_1.dtype().into()); + + // Print the element-wise addition of the two tensors. + println!("{}", B::float_add(tensor_1, tensor_2)); +} diff --git a/crates/burn-tch/src/bin/mps.rs b/crates/burn-tch/src/bin/mps.rs new file mode 100644 index 0000000..2ecad4b --- /dev/null +++ b/crates/burn-tch/src/bin/mps.rs @@ -0,0 +1,16 @@ +use burn_backend::{TensorMetadata, ops::FloatTensorOps}; +use burn_tch::{LibTorch, LibTorchDevice}; + +fn main() { + assert!(tch::utils::has_mps(), "Could not detect MPS"); + + type B = LibTorch; + let device = LibTorchDevice::Mps; + + // Creation of two tensors, the first with explicit values and the second one with ones, with the same shape as the first + let tensor_1 = B::float_from_data([[2f32, 3.], [4., 5.]].into(), &device); + let tensor_2 = B::float_ones(tensor_1.shape(), &device, tensor_1.dtype().into()); + + // Print the element-wise addition of the two tensors. + println!("{}", B::float_add(tensor_1, tensor_2)); +} diff --git a/crates/burn-tch/src/cuda_hack/dummy_cuda_dependency.cpp b/crates/burn-tch/src/cuda_hack/dummy_cuda_dependency.cpp new file mode 100644 index 0000000..54c98e4 --- /dev/null +++ b/crates/burn-tch/src/cuda_hack/dummy_cuda_dependency.cpp @@ -0,0 +1,28 @@ +#include +#include +#include +#include +using namespace std; +extern "C" { +void dummy_cuda_dependency(); +} + +struct cublasContext; + +namespace at { +namespace cuda { +cublasContext *getCurrentCUDABlasHandle(); +int warp_size(); +} // namespace cuda +} // namespace at +char *magma_strerror(int err); +void dummy_cuda_dependency() { + try { + at::cuda::getCurrentCUDABlasHandle(); + at::cuda::warp_size(); + } catch (std::exception &e) { + if (getenv("TCH_PRINT_CUDA_INIT_ERROR") != nullptr) { + std::cerr << "error initializing cuda: " << e.what() << std::endl; + } + } +} diff --git a/crates/burn-tch/src/cuda_hack/fake_cuda_dependency.cpp b/crates/burn-tch/src/cuda_hack/fake_cuda_dependency.cpp new file mode 100644 index 0000000..57c077e --- /dev/null +++ b/crates/burn-tch/src/cuda_hack/fake_cuda_dependency.cpp @@ -0,0 +1,5 @@ +extern "C" { +void dummy_cuda_dependency(); +} + +void dummy_cuda_dependency() {} diff --git a/crates/burn-tch/src/element.rs b/crates/burn-tch/src/element.rs new file mode 100644 index 0000000..36db5ce --- /dev/null +++ b/crates/burn-tch/src/element.rs @@ -0,0 +1,51 @@ +use burn_backend::Element; +use burn_backend::{bf16, f16}; + +/// The element type for the tch backend. +pub trait TchElement: Element + tch::kind::Element { + /// Returns the associated tensor kind for [`tch::kind::Element`]. + fn kind() -> tch::Kind { + Self::KIND + } +} + +impl TchElement for f64 {} +impl TchElement for f32 {} +impl TchElement for f16 {} +impl TchElement for bf16 { + fn kind() -> tch::Kind { + let mut kind = ::KIND; + // Incorrect kind mapping in tch definitions, force bfloat16 + if matches!(Self::dtype(), burn_backend::DType::BF16) && kind == tch::Kind::Half { + kind = tch::Kind::BFloat16 + } + kind + } +} + +impl TchElement for i64 {} +impl TchElement for i32 {} +impl TchElement for i16 {} +impl TchElement for i8 {} + +impl TchElement for u8 {} + +impl TchElement for bool {} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_elem_kinds() { + assert_eq!(f64::kind(), tch::Kind::Double); + assert_eq!(f32::kind(), tch::Kind::Float); + assert_eq!(f16::kind(), tch::Kind::Half); + assert_eq!(bf16::kind(), tch::Kind::BFloat16); + assert_eq!(i64::kind(), tch::Kind::Int64); + assert_eq!(i32::kind(), tch::Kind::Int); + assert_eq!(i16::kind(), tch::Kind::Int16); + assert_eq!(i8::kind(), tch::Kind::Int8); + assert_eq!(bool::kind(), tch::Kind::Bool); + } +} diff --git a/crates/burn-tch/src/lib.rs b/crates/burn-tch/src/lib.rs new file mode 100644 index 0000000..4424a6b --- /dev/null +++ b/crates/burn-tch/src/lib.rs @@ -0,0 +1,14 @@ +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] +#![allow(clippy::single_range_in_vec_init)] + +//! Burn Tch Backend + +mod backend; +mod element; +mod ops; +mod tensor; + +pub use backend::*; +pub use element::*; +pub use tensor::*; diff --git a/crates/burn-tch/src/ops/activation.rs b/crates/burn-tch/src/ops/activation.rs new file mode 100644 index 0000000..bbd1e23 --- /dev/null +++ b/crates/burn-tch/src/ops/activation.rs @@ -0,0 +1,61 @@ +use crate::{LibTorch, TchTensor}; +use burn_backend::ops::ActivationOps; + +impl ActivationOps for LibTorch { + fn relu(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.relu_(), |tensor| tensor.relu()) + } + + fn gelu(tensor: TchTensor) -> TchTensor { + tensor.unary_ops( + |mut tensor| tensor.gelu_("none"), + |tensor| tensor.gelu("none"), + ) + } + + fn gelu_backward(tensor: TchTensor, grad: TchTensor) -> TchTensor { + let storage = tensor.storage.clone(); + let tensor = tensor.tensor.gelu_backward(&grad.tensor, "none"); + + TchTensor::from_existing(tensor, storage) + } + + fn sigmoid(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.sigmoid_(), |tensor| tensor.sigmoid()) + } + + fn log_sigmoid(tensor: TchTensor) -> TchTensor { + // NOTE: we don't override log_sigmoid_backward because Torch has a special backward + // formula that uses a buffer with computed values from the forward pass + + // no in-place log_sigmoid_ + let storage = tensor.storage.clone(); + let tensor = tensor.tensor.log_sigmoid(); + + TchTensor::from_existing(tensor, storage) + } + + fn softmax(tensor: TchTensor, dim: usize) -> TchTensor { + let storage = tensor.storage.clone(); + let tensor = tensor.tensor.softmax(dim as i64, None); + TchTensor::from_existing(tensor, storage) + } + + fn log_softmax(tensor: TchTensor, dim: usize) -> TchTensor { + let storage = tensor.storage.clone(); + let tensor = tensor.tensor.log_softmax(dim as i64, None); + TchTensor::from_existing(tensor, storage) + } + + fn softmin(tensor: TchTensor, dim: usize) -> TchTensor { + let storage = tensor.storage.clone(); + let tensor = tensor.tensor.neg().softmax(dim as i64, None); + TchTensor::from_existing(tensor, storage) + } + + fn prelu(tensor: TchTensor, alpha: TchTensor) -> TchTensor { + let storage = tensor.storage.clone(); + let tensor = tensor.tensor.prelu(&alpha.tensor); + TchTensor::from_existing(tensor, storage) + } +} diff --git a/crates/burn-tch/src/ops/base.rs b/crates/burn-tch/src/ops/base.rs new file mode 100644 index 0000000..fb1eb7b --- /dev/null +++ b/crates/burn-tch/src/ops/base.rs @@ -0,0 +1,837 @@ +use burn_backend::tensor::IndexingUpdateOp; +use burn_backend::{Shape, TensorMetadata}; +use tch::Scalar; + +use crate::{LibTorchDevice, TchShape, TchTensor}; + +pub struct TchOps { + // e: PhantomData, +} + +impl TchOps { + pub fn to_device(tensor: TchTensor, device: &LibTorchDevice) -> TchTensor { + let device = (*device).into(); + + // We have to manually check if the device is the same, since when it's the case, we need to keep + // the same storage reference and not create a new one. + if tensor.tensor.device() == device { + return tensor; + } + + TchTensor::new(tensor.tensor.to(device)) + } + + pub fn reshape(tensor: TchTensor, shape: Shape) -> TchTensor { + let shape_tch: TchShape = shape.into(); + + TchTensor::from_existing(tensor.tensor.reshape(shape_tch.dims), tensor.storage) + } + + pub fn repeat_dim(tensor: TchTensor, dim: usize, times: usize) -> TchTensor { + let mut dims = vec![1; tensor.shape().num_dims()]; + dims[dim] = times as i64; + let tensor = tch::Tensor::repeat(&tensor.tensor, dims); + TchTensor::new(tensor) + } + + pub fn slice_with_steps(tensor: TchTensor, slices: &[burn_backend::Slice]) -> TchTensor { + let storage = tensor.storage.clone(); + let mut tensor = tensor.tensor.shallow_clone(); + + for (dim, slice) in slices.iter().enumerate() { + let dim_i64 = dim as i64; + // Convert slice to range using a dummy size (we'll use tensor dimensions) + let dim_size = tensor.size()[dim]; + let range = slice.to_range(dim_size as usize); + let start = range.start as i64; + let end = range.end as i64; + let step = slice.step as i64; + + if step > 0 { + // Forward stepping - use native slice + tensor = tensor.slice(dim_i64, Some(start), Some(end), step); + } else { + // Negative stepping - we need to handle the semantics correctly + // For negative steps, we iterate backwards from end-1 + // PyTorch's negative step works differently than our semantics + // We need to reverse the selected range + + // First get the slice with positive step + tensor = tensor.slice(dim_i64, Some(start), Some(end), 1); + + // Then reverse it and apply the step + if step == -1 { + // Simple reversal + tensor = tensor.flip([dim_i64]); + } else { + // Reverse and then take every nth element + tensor = tensor.flip([dim_i64]); + let abs_step = step.abs(); + tensor = tensor.slice(dim_i64, None, None, abs_step); + } + } + } + + TchTensor::partial(tensor, storage) + } + + pub fn slice_assign( + tensor: TchTensor, + slices: &[burn_backend::Slice], + value: TchTensor, + ) -> TchTensor { + // PyTorch's narrow operation only supports contiguous slices (step=1) + // For non-unit steps, we use advanced indexing as a workaround + let all_unit_steps = slices.iter().all(|s| s.step == 1); + + if all_unit_steps { + // Fast path: use narrow and copy_ for unit steps + let tch_shape = TchShape::from(tensor.shape()); + + // Copy the input tensor if we can't mutate it + let tensor_original: TchTensor = + tensor.unary_ops(|tensor| tensor, |tensor| tensor.copy()); + let tensor_original = tensor_original.tensor; + + let mut tensor = tensor_original.view_(tch_shape.dims); + + for (i, slice) in slices.iter().enumerate().take(slices.len()) { + // Convert Slice to range for narrow operation + let dim_size = tensor.size()[i] as usize; + let range = slice.to_range(dim_size); + let start = range.start as i64; + let length = (range.end - range.start) as i64; + + tensor = tensor.narrow(i as i64, start, length); + } + + tensor.copy_(&value.tensor); + TchTensor::new(tensor_original) + } else { + // Workaround for non-unit steps: use PyTorch's index_put operation + // This generates explicit indices for the slice and uses advanced indexing + let tensor_shape = tensor.shape(); + let dims = tensor_shape.clone(); + + // Copy the tensor since we'll modify it + let result_tensor = tensor.tensor.shallow_clone(); + + // Use advanced indexing to set the values + Self::slice_assign_with_advanced_indexing(result_tensor, slices, value.tensor, &dims) + } + } + + /// Generate indices for a slice with potentially non-unit step. + /// For negative steps, generates indices in reverse order. + fn generate_slice_indices(slice: &burn_backend::Slice, dim_size: usize) -> Vec { + let step = slice.step; + let range = slice.to_range(dim_size); + + let mut indices = Vec::new(); + + if step > 0 { + let mut idx = range.start as i64; + while idx < range.end as i64 { + indices.push(idx); + idx += step as i64; + } + } else if step < 0 { + // For negative steps, iterate backwards through the range + let mut idx = (range.end - 1) as i64; + while idx >= range.start as i64 { + indices.push(idx); + idx += step as i64; // step is negative, so this decreases + } + } + + indices + } + + /// Implementation using advanced indexing for non-unit steps. + /// Uses PyTorch's index_put operation to assign values at specific indices. + fn slice_assign_with_advanced_indexing( + mut tensor: tch::Tensor, + slices: &[burn_backend::Slice], + value: tch::Tensor, + dims: &[usize], + ) -> TchTensor { + // Generate all index combinations for the sliced regions + let mut index_sets: Vec> = Vec::new(); + for (i, slice) in slices.iter().enumerate() { + let dim_size = if i < dims.len() { dims[i] } else { 1 }; + let indices = Self::generate_slice_indices(slice, dim_size); + index_sets.push(indices); + } + + // For unsliced dimensions, include all indices + for &dim_size in dims.iter().skip(slices.len()) { + let indices: Vec = (0..dim_size as i64).collect(); + index_sets.push(indices); + } + + // Convert index sets to tensors for index_put + let mut final_indices = Vec::new(); + let total_elements = index_sets.iter().map(|s| s.len()).product::(); + + // Build flattened index arrays for each dimension using cartesian product + // This creates the index tensors needed for PyTorch's index_put operation + for dim_idx in 0..index_sets.len() { + let mut dim_indices = Vec::with_capacity(total_elements); + let repeat = index_sets[dim_idx + 1..] + .iter() + .map(|s| s.len()) + .product::() + .max(1); + let tile = index_sets[..dim_idx] + .iter() + .map(|s| s.len()) + .product::() + .max(1); + + for _ in 0..tile { + for &idx in &index_sets[dim_idx] { + for _ in 0..repeat { + dim_indices.push(idx); + } + } + } + + let indices_tensor = tch::Tensor::from_slice(&dim_indices).to_device(tensor.device()); + final_indices.push(indices_tensor); + } + + // PyTorch's index_put handles assignment correctly for negative steps + // following NumPy semantics: values[i] goes to selected_indices[i] + let value_flat = value.view(-1); + + // Use index_put to assign values - convert to Option + let final_indices_opt: Vec> = + final_indices.into_iter().map(Some).collect(); + tensor = tensor.index_put(&final_indices_opt, &value_flat, false); + + TchTensor::new(tensor) + } + + pub fn gather(dim: usize, tensor: TchTensor, indices: TchTensor) -> TchTensor { + let storage = tensor.storage.clone(); + let tensor = tensor.tensor.gather(dim as i64, &indices.tensor, false); + + TchTensor::from_existing(tensor, storage) + } + + pub fn scatter( + dim: usize, + tensor: TchTensor, + indices: TchTensor, + value: TchTensor, + ) -> TchTensor { + let storage = tensor.storage.clone(); + let tensor = tensor + .tensor + .scatter_add(dim as i64, &indices.tensor, &value.tensor); + + TchTensor::from_existing(tensor, storage) + } + + /// Flatten K-dimensional index tuples into 1D linear offsets, suitable for + /// use with PyTorch's scatter/gather along dim 0 of a flattened tensor. + /// + /// `indices` has shape `[num_updates, k]`. + /// Returns a 1D tensor of shape `[num_updates * slice_size]`. + fn flatten_nd_indices( + indices: &tch::Tensor, + data_shape: &[i64], + k: usize, + slice_size: i64, + ) -> tch::Tensor { + // Compute per-dimension strides for the first K dims + let mut strides = vec![0i64; k]; + if k > 0 { + strides[k - 1] = slice_size; + for i in (0..k - 1).rev() { + strides[i] = strides[i + 1] * data_shape[i + 1]; + } + } + + // base_offsets: [num_updates] = sum(indices[:, j] * strides[j]) + let strides_tensor = tch::Tensor::from_slice(&strides).to_device(indices.device()); + let base_offsets = indices + .matmul(&strides_tensor.unsqueeze(-1)) + .squeeze_dim(-1); + + // Expand base_offsets to [num_updates, slice_size] and add intra-slice offsets + let slice_offsets = tch::Tensor::arange(slice_size, (tch::Kind::Int64, indices.device())); + let flat = base_offsets.unsqueeze(-1) + slice_offsets.unsqueeze(0); + flat.reshape([flat.size().iter().product::()]) + } + + pub fn scatter_nd( + tensor: TchTensor, + indices: TchTensor, + values: TchTensor, + reduction: IndexingUpdateOp, + ) -> TchTensor { + let data_shape: Vec = tensor.tensor.size(); + let idx_shape: Vec = indices.tensor.size(); + let m = idx_shape.len(); + let k = *idx_shape.last().unwrap() as usize; + let num_updates: i64 = idx_shape[..m - 1].iter().product(); + let total_elems: i64 = data_shape.iter().product(); + let slice_size: i64 = data_shape[k..].iter().product(); + + // Flatten indices [I0,..,I_{M-2}, K] -> [num_updates, K] + let flat_indices = indices.tensor.reshape([num_updates, k as i64]); + let linear_idx = Self::flatten_nd_indices(&flat_indices, &data_shape, k, slice_size); + + // Flatten data and values to 1D + let mut flat_data = tensor.tensor.reshape([total_elems]); + let flat_values = values.tensor.reshape([num_updates * slice_size]); + + let result = match reduction { + IndexingUpdateOp::Assign => flat_data.scatter_(0, &linear_idx, &flat_values), + IndexingUpdateOp::Add => flat_data.scatter_add(0, &linear_idx, &flat_values), + IndexingUpdateOp::Mul => flat_data.scatter_reduce(0, &linear_idx, &flat_values, "prod"), + IndexingUpdateOp::Min => flat_data.scatter_reduce(0, &linear_idx, &flat_values, "amin"), + IndexingUpdateOp::Max => flat_data.scatter_reduce(0, &linear_idx, &flat_values, "amax"), + }; + + let storage = tensor.storage.clone(); + TchTensor::from_existing(result.reshape(data_shape), storage) + } + + pub fn gather_nd(tensor: TchTensor, indices: TchTensor) -> TchTensor { + let data_shape: Vec = tensor.tensor.size(); + let idx_shape: Vec = indices.tensor.size(); + let m = idx_shape.len(); + let k = *idx_shape.last().unwrap() as usize; + let num_indices: i64 = idx_shape[..m - 1].iter().product(); + let total_elems: i64 = data_shape.iter().product(); + let slice_size: i64 = data_shape[k..].iter().product(); + + let flat_indices = indices.tensor.reshape([num_indices, k as i64]); + let linear_idx = Self::flatten_nd_indices(&flat_indices, &data_shape, k, slice_size); + + let flat_data = tensor.tensor.reshape([total_elems]); + let gathered = flat_data.index_select(0, &linear_idx); + + // Output shape: idx_shape[..m-1] ++ data_shape[k..] + let mut out_shape: Vec = idx_shape[..m - 1].to_vec(); + out_shape.extend_from_slice(&data_shape[k..]); + + let storage = tensor.storage.clone(); + TchTensor::from_existing(gathered.reshape(out_shape), storage) + } + + pub fn index_select_dim(tensor: TchTensor, dim: usize, indices: TchTensor) -> TchTensor { + let storage = tensor.storage.clone(); + let tensor = tensor.tensor.index_select(dim as i64, &indices.tensor); + + TchTensor::from_existing(tensor, storage) + } + + pub fn select_assign( + tensor: TchTensor, + dim: usize, + indices: TchTensor, + value: TchTensor, + ) -> TchTensor { + tensor.clone().unary_ops( + |mut tensor| tensor.index_add_(dim as i64, &indices.tensor, &value.tensor), + |tensor| tensor.index_add(dim as i64, &indices.tensor, &value.tensor), + ) + } + + pub fn cat(tensors: Vec, dim: usize) -> TchTensor { + let tensors: Vec = tensors + .into_iter() + .map(|t| t.tensor.shallow_clone()) + .collect(); + let tensor = tch::Tensor::cat(&tensors, dim as i64); + + TchTensor::new(tensor) + } + + pub fn equal(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.eq_tensor_(rhs).to_kind(tch::Kind::Bool), + |lhs, rhs| rhs.eq_tensor_(lhs).to_kind(tch::Kind::Bool), + |lhs, rhs| lhs.eq_tensor(rhs), + ) + } + + pub fn equal_elem + Clone>(lhs: TchTensor, rhs: S) -> TchTensor { + lhs.unary_ops( + |mut tensor| tensor.eq_(rhs.clone().into()).to_kind(tch::Kind::Bool), + |tensor| tensor.eq(rhs.clone().into()), + ) + } + + pub fn greater(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.greater_tensor_(rhs).to_kind(tch::Kind::Bool), + |lhs, rhs| rhs.less_tensor_(lhs).to_kind(tch::Kind::Bool), + |lhs, rhs| lhs.greater_tensor(rhs), + ) + } + + pub fn greater_elem + Clone>(lhs: TchTensor, rhs: S) -> TchTensor { + lhs.unary_ops( + |mut tensor| tensor.greater_(rhs.clone().into()).to_kind(tch::Kind::Bool), + |tensor| tensor.greater(rhs.clone().into()), + ) + } + + pub fn greater_equal(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.greater_equal_tensor_(rhs).to_kind(tch::Kind::Bool), + |lhs, rhs| rhs.less_equal_tensor_(lhs).to_kind(tch::Kind::Bool), + |lhs, rhs| lhs.greater_equal_tensor(rhs), + ) + } + + pub fn greater_equal_elem + Clone>(lhs: TchTensor, rhs: S) -> TchTensor { + lhs.unary_ops( + |mut tensor| { + tensor + .greater_equal_(rhs.clone().into()) + .to_kind(tch::Kind::Bool) + }, + |tensor| tensor.greater_equal(rhs.clone().into()), + ) + } + + pub fn lower(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.less_tensor_(rhs).to_kind(tch::Kind::Bool), + |lhs, rhs| rhs.greater_tensor_(lhs).to_kind(tch::Kind::Bool), + |lhs, rhs| lhs.less_tensor(rhs), + ) + } + + pub fn lower_elem + Clone>(lhs: TchTensor, rhs: S) -> TchTensor { + lhs.unary_ops( + |mut tensor| tensor.less_(rhs.clone().into()).to_kind(tch::Kind::Bool), + |tensor| tensor.less(rhs.clone().into()), + ) + } + + pub fn lower_equal(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.less_equal_tensor_(rhs).to_kind(tch::Kind::Bool), + |lhs, rhs| rhs.greater_equal_tensor_(lhs).to_kind(tch::Kind::Bool), + |lhs, rhs| lhs.less_equal_tensor(rhs), + ) + } + + pub fn lower_equal_elem + Clone>(lhs: TchTensor, rhs: S) -> TchTensor { + lhs.unary_ops( + |mut tensor| { + tensor + .less_equal_(rhs.clone().into()) + .to_kind(tch::Kind::Bool) + }, + |tensor| tensor.less_equal(rhs.clone().into()), + ) + } + + pub fn add(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.f_add_(rhs).unwrap(), + |lhs, rhs| rhs.f_add_(lhs).unwrap(), + |lhs, rhs| lhs.f_add(rhs).unwrap(), + ) + } + + pub fn sub(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.f_sub_(rhs).unwrap(), + |lhs, rhs| lhs.f_sub(rhs).unwrap(), + |lhs, rhs| lhs.f_sub(rhs).unwrap(), + ) + } + + pub fn mul(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.f_mul_(rhs).unwrap(), + |lhs, rhs| rhs.f_mul_(lhs).unwrap(), + |lhs, rhs| lhs.f_mul(rhs).unwrap(), + ) + } + + pub fn div(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.f_div_(rhs).unwrap(), + |lhs, rhs| lhs.f_div(rhs).unwrap(), + |lhs, rhs| lhs.f_div(rhs).unwrap(), + ) + } + + pub fn remainder(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.f_remainder_tensor_(rhs).unwrap(), + |lhs, rhs| lhs.f_remainder_tensor(rhs).unwrap(), + |lhs, rhs| lhs.f_remainder_tensor(rhs).unwrap(), + ) + } + + pub fn mean(tensor: TchTensor) -> TchTensor { + // view as 1d tensor + let tensor = tensor.tensor.mean(tensor.tensor.kind()).view(1); + TchTensor::new(tensor) + } + + pub fn mean_dim(tensor: TchTensor, dim: usize) -> TchTensor { + TchTensor::from_existing( + tensor + .tensor + .mean_dim(Some([dim as i64].as_slice()), true, tensor.tensor.kind()), + tensor.storage, + ) + } + + pub fn sum(tensor: TchTensor) -> TchTensor { + // view as 1d tensor + let tensor = tensor.tensor.sum(tensor.tensor.kind()).view(1); + TchTensor::new(tensor) + } + + pub fn sum_dim(tensor: TchTensor, dim: usize) -> TchTensor { + TchTensor::from_existing( + tensor.tensor.sum_dim_intlist( + Some([dim as i64].as_slice()), + true, + tensor.tensor.kind(), + ), + tensor.storage, + ) + } + + pub fn prod(tensor: TchTensor) -> TchTensor { + // view as 1d tensor + let tensor = tensor.tensor.prod(tensor.tensor.kind()).view(1); + TchTensor::new(tensor) + } + + pub fn prod_dim(tensor: TchTensor, dim: usize) -> TchTensor { + TchTensor::from_existing( + tensor + .tensor + .prod_dim_int(dim as i64, true, tensor.tensor.kind()), + tensor.storage, + ) + } + + pub fn topk(tensor: TchTensor, dim: usize, k: usize) -> TchTensor { + let (value, _indices) = tensor.tensor.topk(k as i64, dim as i64, true, true); + TchTensor::from_existing(value, tensor.storage) + } + + pub fn argtopk(tensor: TchTensor, dim: usize, k: usize) -> TchTensor { + let (_value, indices) = tensor.tensor.topk(k as i64, dim as i64, true, true); + TchTensor::from_existing(indices, tensor.storage) + } + + pub fn cumsum(tensor: TchTensor, dim: usize) -> TchTensor { + TchTensor::from_existing( + tensor.tensor.cumsum(dim as i64, tensor.tensor.kind()), + tensor.storage, + ) + } + + pub fn cumprod(tensor: TchTensor, dim: usize) -> TchTensor { + TchTensor::from_existing( + tensor.tensor.cumprod(dim as i64, tensor.tensor.kind()), + tensor.storage, + ) + } + + pub fn cummin(tensor: TchTensor, dim: usize) -> TchTensor { + let (values, _indices) = tensor.tensor.cummin(dim as i64); + TchTensor::from_existing(values, tensor.storage) + } + + pub fn cummax(tensor: TchTensor, dim: usize) -> TchTensor { + // cummax returns (values, indices) tuple in PyTorch, we only need values + let (values, _indices) = tensor.tensor.cummax(dim as i64); + TchTensor::from_existing(values, tensor.storage) + } + + pub fn argmax(tensor: TchTensor, dim: usize) -> TchTensor { + let storage = tensor.storage.clone(); + let tensor = tensor.tensor.argmax(dim as i64, true); + + TchTensor::from_existing(tensor, storage) + } + + pub fn argmin(tensor: TchTensor, dim: usize) -> TchTensor { + let storage = tensor.storage.clone(); + let tensor = tensor.tensor.argmin(dim as i64, true); + + TchTensor::from_existing(tensor, storage) + } + + pub fn max_dim(tensor: TchTensor, dim: usize) -> TchTensor { + let storage = tensor.storage.clone(); + let (tensor, _indices) = tensor.tensor.max_dim(dim as i64, true); + + TchTensor::from_existing(tensor, storage) + } + + pub fn max_dim_with_indices(tensor: TchTensor, dim: usize) -> (TchTensor, TchTensor) { + let storage = tensor.storage.clone(); + let (tensor, indices) = tensor.tensor.max_dim(dim as i64, true); + + let tensor = TchTensor::from_existing(tensor, storage); + let indices = TchTensor::new(indices); + + (tensor, indices) + } + + pub fn min_dim(tensor: TchTensor, dim: usize) -> TchTensor { + let storage = tensor.storage.clone(); + let (tensor, _indices) = tensor.tensor.min_dim(dim as i64, true); + + TchTensor::from_existing(tensor, storage) + } + + pub fn min_dim_with_indices(tensor: TchTensor, dim: usize) -> (TchTensor, TchTensor) { + let storage = tensor.storage.clone(); + let (tensor, indices) = tensor.tensor.min_dim(dim as i64, true); + + let tensor = TchTensor::from_existing(tensor, storage); + let indices = TchTensor::new(indices); + + (tensor, indices) + } + + pub fn clamp_min + Clone + Copy>(tensor: TchTensor, min: S) -> TchTensor { + tensor.unary_ops( + |mut tensor| tensor.clamp_min_(min), + |tensor| tensor.clamp_min(min), + ) + } + + pub fn clamp_max + Clone + Copy>(tensor: TchTensor, max: S) -> TchTensor { + tensor.unary_ops( + |mut tensor| tensor.clamp_max_(max), + |tensor| tensor.clamp_max(max), + ) + } + + pub fn clamp + Clone + Copy>( + tensor: TchTensor, + min: S, + max: S, + ) -> TchTensor { + tensor.unary_ops( + |mut tensor| tensor.clamp_(min, max), + |tensor| tensor.clamp(min, max), + ) + } + + pub fn swap_dims(tensor: TchTensor, dim1: usize, dim2: usize) -> TchTensor { + let tensor = tensor.tensor.transpose(dim1 as i64, dim2 as i64); + TchTensor::new(tensor) + } + + pub fn permute(tensor: TchTensor, axes: &[usize]) -> TchTensor { + let tensor = tensor + .tensor + .permute(axes.iter().map(|x| *x as i64).collect::>()); + TchTensor::new(tensor) + } + + pub fn flip(tensor: TchTensor, axes: &[usize]) -> TchTensor { + let dims = axes.iter().map(|x| *x as i64).collect::>(); + let tensor = tensor.tensor.flip(dims); + TchTensor::new(tensor) + } + + pub fn pow(tensor: TchTensor, exponent: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + tensor, + exponent, + |lhs, rhs| lhs.f_pow_tensor_(rhs).unwrap(), + |lhs, rhs| lhs.f_pow(rhs).unwrap(), + |lhs, rhs| lhs.f_pow(rhs).unwrap(), + ) + } + + pub fn sign(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.sign_(), |tensor| tensor.sign()) + } + + pub fn expand(tensor: TchTensor, shape: Shape) -> TchTensor { + let storage = tensor.storage.clone(); + let broadcasted_tensor = tensor.tensor.broadcast_to(TchShape::from(shape).dims); + TchTensor::from_existing(broadcasted_tensor, storage) + } + + pub fn unfold(tensor: TchTensor, dim: usize, size: usize, step: usize) -> TchTensor { + let storage = tensor.storage.clone(); + let uf_tensor = tensor.tensor.unfold(dim as i64, size as i64, step as i64); + + TchTensor::from_existing(uf_tensor, storage) + } + + pub fn sort(tensor: TchTensor, dim: usize, descending: bool) -> TchTensor { + TchTensor::new(tensor.tensor.sort(dim as i64, descending).0) + } + + pub fn sort_with_indices( + tensor: TchTensor, + dim: usize, + descending: bool, + ) -> (TchTensor, TchTensor) { + let sorted = tensor.tensor.sort(dim as i64, descending); + (TchTensor::new(sorted.0), TchTensor::new(sorted.1)) + } + + pub fn argsort(tensor: TchTensor, dim: usize, descending: bool) -> TchTensor { + TchTensor::new(tensor.tensor.argsort(dim as i64, descending)) + } + + pub fn bitwise_and(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.f_bitwise_and_tensor_(rhs).unwrap(), + |lhs, rhs| rhs.f_bitwise_and_tensor_(lhs).unwrap(), + |lhs, rhs| lhs.f_bitwise_and_tensor(rhs).unwrap(), + ) + } + + pub fn bitwise_and_scalar + Clone>(tensor: TchTensor, scalar: S) -> TchTensor { + tensor.unary_ops( + |mut tensor| tensor.f_bitwise_and_(scalar.clone().into()).unwrap(), + |tensor| tensor.f_bitwise_and(scalar.clone().into()).unwrap(), + ) + } + + pub fn bitwise_or(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.f_bitwise_or_tensor_(rhs).unwrap(), + |lhs, rhs| rhs.f_bitwise_or_tensor_(lhs).unwrap(), + |lhs, rhs| lhs.f_bitwise_or_tensor(rhs).unwrap(), + ) + } + + pub fn bitwise_or_scalar + Clone>(tensor: TchTensor, scalar: S) -> TchTensor { + tensor.unary_ops( + |mut tensor| tensor.f_bitwise_or_(scalar.clone().into()).unwrap(), + |tensor| tensor.f_bitwise_or(scalar.clone().into()).unwrap(), + ) + } + + pub fn bitwise_xor(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.f_bitwise_xor_tensor_(rhs).unwrap(), + |lhs, rhs| rhs.f_bitwise_xor_tensor_(lhs).unwrap(), + |lhs, rhs| lhs.f_bitwise_xor_tensor(rhs).unwrap(), + ) + } + + pub fn bitwise_xor_scalar + Clone>(tensor: TchTensor, scalar: S) -> TchTensor { + tensor.unary_ops( + |mut tensor| tensor.f_bitwise_xor_(scalar.clone().into()).unwrap(), + |tensor| tensor.f_bitwise_xor(scalar.clone().into()).unwrap(), + ) + } + + pub fn bitwise_not(tensor: TchTensor) -> TchTensor { + tensor.unary_ops( + |mut tensor| tensor.f_bitwise_not_().unwrap(), + |tensor| tensor.f_bitwise_not().unwrap(), + ) + } + + pub fn bitwise_left_shift(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.f_bitwise_left_shift_(rhs).unwrap(), + |lhs, rhs| lhs.f_bitwise_left_shift(rhs).unwrap(), + |lhs, rhs| lhs.f_bitwise_left_shift(rhs).unwrap(), + ) + } + + pub fn bitwise_left_shift_scalar + Clone>( + tensor: TchTensor, + scalar: S, + ) -> TchTensor { + tensor.unary_ops( + |mut tensor| { + tensor + .f_bitwise_left_shift_tensor_scalar_(scalar.clone().into()) + .unwrap() + }, + |tensor| { + tensor + .f_bitwise_left_shift_tensor_scalar(scalar.clone().into()) + .unwrap() + }, + ) + } + + pub fn bitwise_right_shift(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.f_bitwise_right_shift_(rhs).unwrap(), + |lhs, rhs| lhs.f_bitwise_right_shift(rhs).unwrap(), + |lhs, rhs| lhs.f_bitwise_right_shift(rhs).unwrap(), + ) + } + + pub fn bitwise_right_shift_scalar + Clone>( + tensor: TchTensor, + scalar: S, + ) -> TchTensor { + tensor.unary_ops( + |mut tensor| { + tensor + .f_bitwise_right_shift_tensor_scalar_(scalar.clone().into()) + .unwrap() + }, + |tensor| { + tensor + .f_bitwise_right_shift_tensor_scalar(scalar.clone().into()) + .unwrap() + }, + ) + } + + pub fn atan2(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.f_atan2_(rhs).unwrap(), + |lhs, rhs| lhs.f_atan2(rhs).unwrap(), + |lhs, rhs| lhs.f_atan2(rhs).unwrap(), + ) + } +} diff --git a/crates/burn-tch/src/ops/bool_tensor.rs b/crates/burn-tch/src/ops/bool_tensor.rs new file mode 100644 index 0000000..5ff05d8 --- /dev/null +++ b/crates/burn-tch/src/ops/bool_tensor.rs @@ -0,0 +1,221 @@ +use super::TchOps; +use crate::IntoKind; +use crate::{LibTorch, LibTorchDevice, TchShape, TchTensor}; +use burn_backend::BoolStore; +use burn_backend::ExecutionError; +use burn_backend::IntDType; +use burn_backend::Scalar; +use burn_backend::tensor::BoolTensor; +use burn_backend::tensor::IntTensor; +use burn_backend::{BoolDType, FloatDType}; +use burn_backend::{Shape, TensorData, TensorMetadata, ops::BoolTensorOps}; + +impl BoolTensorOps for LibTorch { + fn bool_from_data(data: TensorData, device: &LibTorchDevice) -> TchTensor { + match data.dtype { + burn_backend::DType::Bool(BoolStore::Native) => { + TchTensor::from_data::(data, (*device).into()) + } + _ => unimplemented!("Unsupported dtype for `bool_from_data`"), + } + } + + fn bool_repeat_dim(tensor: TchTensor, dim: usize, times: usize) -> TchTensor { + TchOps::repeat_dim(tensor, dim, times) + } + + async fn bool_into_data(tensor: TchTensor) -> Result { + let shape = tensor.shape(); + let tensor = Self::bool_reshape(tensor.clone(), Shape::new([shape.num_elements()])); + let values: Result, tch::TchError> = tensor.tensor.shallow_clone().try_into(); + Ok(TensorData::new(values.unwrap(), shape)) + } + + fn bool_to_device(tensor: TchTensor, device: &LibTorchDevice) -> TchTensor { + TchOps::to_device(tensor, device) + } + + fn bool_reshape(tensor: TchTensor, shape: Shape) -> TchTensor { + TchOps::reshape(tensor, shape) + } + + fn bool_empty(shape: Shape, device: &LibTorchDevice, _dtype: BoolDType) -> TchTensor { + let tensor = tch::Tensor::empty( + TchShape::from(shape).dims, + (tch::Kind::Bool, (*device).into()), + ); + + TchTensor::new(tensor) + } + + fn bool_zeros(shape: Shape, device: &LibTorchDevice, _dtype: BoolDType) -> TchTensor { + let tensor = tch::Tensor::zeros( + TchShape::from(shape).dims, + (tch::Kind::Bool, (*device).into()), + ); + + TchTensor::new(tensor) + } + + fn bool_ones(shape: Shape, device: &LibTorchDevice, _dtype: BoolDType) -> TchTensor { + let tensor = tch::Tensor::ones( + TchShape::from(shape).dims, + (tch::Kind::Bool, (*device).into()), + ); + + TchTensor::new(tensor) + } + + fn bool_slice(tensor: TchTensor, slices: &[burn_backend::Slice]) -> TchTensor { + TchOps::slice_with_steps(tensor, slices) + } + + fn bool_slice_assign( + tensor: TchTensor, + slices: &[burn_backend::Slice], + value: TchTensor, + ) -> TchTensor { + TchOps::slice_assign(tensor, slices, value) + } + + fn bool_cat(tensors: Vec, dim: usize) -> TchTensor { + TchOps::cat(tensors, dim) + } + + fn bool_equal(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchOps::equal(lhs, rhs) + } + + fn bool_not(tensor: TchTensor) -> TchTensor { + tensor.unary_ops( + |mut tensor| tensor.eq_(0).to_kind(tch::Kind::Bool), + |tensor| tensor.eq(0), + ) + } + + fn bool_and(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.logical_and_(rhs), + |lhs, rhs| rhs.logical_and_(lhs), + |lhs, rhs| lhs.logical_and(rhs), + ) + } + + fn bool_or(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + lhs, + rhs, + |lhs, rhs| lhs.logical_or_(rhs), + |lhs, rhs| rhs.logical_or_(lhs), + |lhs, rhs| lhs.logical_or(rhs), + ) + } + + fn bool_into_int(tensor: TchTensor, out_dtype: IntDType) -> TchTensor { + let tensor = tensor.tensor.to_kind(out_dtype.into_kind()); + TchTensor::new(tensor) + } + + fn bool_into_float(tensor: TchTensor, out_dtype: FloatDType) -> TchTensor { + let tensor = tensor.tensor.to_kind(out_dtype.into_kind()); + TchTensor::new(tensor) + } + + fn bool_swap_dims(tensor: TchTensor, dim1: usize, dim2: usize) -> TchTensor { + TchOps::swap_dims(tensor, dim1, dim2) + } + + fn bool_permute(tensor: TchTensor, axes: &[usize]) -> TchTensor { + TchOps::permute(tensor, axes) + } + + fn bool_flip(tensor: TchTensor, axes: &[usize]) -> TchTensor { + TchOps::flip(tensor, axes) + } + + async fn bool_argwhere(tensor: TchTensor, out_dtype: IntDType) -> TchTensor { + TchTensor::new(tensor.tensor.argwhere().to_kind(out_dtype.into_kind())) + } + + fn bool_select(tensor: TchTensor, dim: usize, indices: TchTensor) -> TchTensor { + TchOps::index_select_dim(tensor, dim, indices) + } + + fn bool_select_or( + tensor: TchTensor, + dim: usize, + indices: TchTensor, + value: TchTensor, + ) -> TchTensor { + TchOps::select_assign(tensor, dim, indices, value) + } + + fn bool_expand(tensor: TchTensor, shape: Shape) -> TchTensor { + TchOps::expand(tensor, shape) + } + + fn bool_unfold( + tensor: IntTensor, + dim: usize, + size: usize, + step: usize, + ) -> IntTensor { + TchOps::unfold(tensor, dim, size, step) + } + + fn bool_mask_where( + tensor: BoolTensor, + mask: BoolTensor, + value: BoolTensor, + ) -> BoolTensor { + TchTensor::binary_ops_tensor( + tensor, + value, + |tensor, source| source.f_where_self(&mask.tensor, tensor).unwrap(), + |tensor, source| source.f_where_self(&mask.tensor, tensor).unwrap(), + |tensor, source| source.f_where_self(&mask.tensor, tensor).unwrap(), + ) + } + + fn bool_mask_fill( + tensor: BoolTensor, + mask: BoolTensor, + value: Scalar, + ) -> BoolTensor { + tensor.unary_ops( + |mut tensor| { + tensor + .f_masked_fill_(&mask.tensor, value.elem::()) + .unwrap() + }, + |tensor| { + tensor + .f_masked_fill(&mask.tensor, value.elem::()) + .unwrap() + }, + ) + } + + fn bool_gather( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + ) -> BoolTensor { + TchOps::gather(dim, tensor, indices) + } + + fn bool_scatter_or( + dim: usize, + tensor: BoolTensor, + indices: IntTensor, + value: BoolTensor, + ) -> BoolTensor { + TchOps::scatter(dim, tensor, indices, value) + } + + fn bool_equal_elem(lhs: BoolTensor, rhs: Scalar) -> BoolTensor { + TchOps::equal_elem(lhs, rhs.elem::()) + } +} diff --git a/crates/burn-tch/src/ops/int_tensor.rs b/crates/burn-tch/src/ops/int_tensor.rs new file mode 100644 index 0000000..9922f7b --- /dev/null +++ b/crates/burn-tch/src/ops/int_tensor.rs @@ -0,0 +1,540 @@ +use std::ops::Range; + +use burn_backend::{ + BoolDType, Distribution, ExecutionError, FloatDType, IntDType, Scalar, Shape, TensorData, + TensorMetadata, + ops::{FloatTensorOps, IntTensorOps}, + tensor::IntTensor, +}; + +use crate::{IntoKind, LibTorch, LibTorchDevice, TchShape, TchTensor}; + +use super::TchOps; + +impl IntTensorOps for LibTorch { + fn int_from_data(data: TensorData, device: &LibTorchDevice) -> TchTensor { + match data.dtype { + burn_backend::DType::I64 => TchTensor::from_data::(data, (*device).into()), + burn_backend::DType::I32 => TchTensor::from_data::(data, (*device).into()), + burn_backend::DType::I16 => TchTensor::from_data::(data, (*device).into()), + burn_backend::DType::I8 => TchTensor::from_data::(data, (*device).into()), + burn_backend::DType::U8 => TchTensor::from_data::(data, (*device).into()), + _ => unimplemented!("Unsupported dtype for `int_from_data`: {:?}", data.dtype), + } + } + + fn int_repeat_dim(tensor: TchTensor, dim: usize, times: usize) -> TchTensor { + TchOps::repeat_dim(tensor, dim, times) + } + + async fn int_into_data(tensor: TchTensor) -> Result { + let shape = tensor.shape(); + let tensor = Self::int_reshape(tensor.clone(), Shape::new([shape.num_elements()])); + let values: Result, tch::TchError> = tensor.tensor.shallow_clone().try_into(); + Ok(TensorData::new(values.unwrap(), shape)) + } + + fn int_to_device(tensor: TchTensor, device: &LibTorchDevice) -> TchTensor { + TchOps::to_device(tensor, device) + } + + fn int_reshape(tensor: TchTensor, shape: Shape) -> TchTensor { + TchOps::reshape(tensor, shape) + } + + fn int_empty(shape: Shape, device: &LibTorchDevice, dtype: IntDType) -> TchTensor { + let tensor = tch::Tensor::empty( + TchShape::from(shape).dims, + (dtype.into_kind(), (*device).into()), + ); + + TchTensor::new(tensor) + } + + fn int_slice(tensor: TchTensor, slices: &[burn_backend::Slice]) -> TchTensor { + TchOps::slice_with_steps(tensor, slices) + } + + fn int_slice_assign( + tensor: TchTensor, + slices: &[burn_backend::Slice], + value: TchTensor, + ) -> TchTensor { + TchOps::slice_assign(tensor, slices, value) + } + + fn int_cat(tensors: Vec, dim: usize) -> TchTensor { + TchOps::cat(tensors, dim) + } + + fn int_matmul(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + let int_dtype = lhs.dtype(); + let lhs = Self::int_into_float(lhs, FloatDType::F32); + let rhs = Self::int_into_float(rhs, FloatDType::F32); + let out = lhs.tensor.f_matmul(&rhs.tensor).unwrap(); + Self::float_into_int(TchTensor::new(out), int_dtype.into()) + } + + fn int_equal(lhs: TchTensor, rhs: TchTensor, _out_dtype: BoolDType) -> TchTensor { + TchOps::equal(lhs, rhs) + } + + fn int_equal_elem(lhs: TchTensor, rhs: Scalar, _out_dtype: BoolDType) -> TchTensor { + TchOps::equal_elem(lhs, rhs.elem::()) + } + + fn int_greater(lhs: TchTensor, rhs: TchTensor, _out_dtype: BoolDType) -> TchTensor { + TchOps::greater(lhs, rhs) + } + + fn int_greater_elem(lhs: TchTensor, rhs: Scalar, _out_dtype: BoolDType) -> TchTensor { + TchOps::greater_elem(lhs, rhs.elem::()) + } + + fn int_greater_equal(lhs: TchTensor, rhs: TchTensor, _out_dtype: BoolDType) -> TchTensor { + TchOps::greater_equal(lhs, rhs) + } + + fn int_greater_equal_elem(lhs: TchTensor, rhs: Scalar, _out_dtype: BoolDType) -> TchTensor { + TchOps::greater_equal_elem(lhs, rhs.elem::()) + } + + fn int_lower(lhs: TchTensor, rhs: TchTensor, _out_dtype: BoolDType) -> TchTensor { + TchOps::lower(lhs, rhs) + } + + fn int_lower_elem(lhs: TchTensor, rhs: Scalar, _out_dtype: BoolDType) -> TchTensor { + TchOps::lower_elem(lhs, rhs.elem::()) + } + + fn int_lower_equal(lhs: TchTensor, rhs: TchTensor, _out_dtype: BoolDType) -> TchTensor { + TchOps::lower_equal(lhs, rhs) + } + + fn int_lower_equal_elem(lhs: TchTensor, rhs: Scalar, _out_dtype: BoolDType) -> TchTensor { + TchOps::lower_equal_elem(lhs, rhs.elem::()) + } + + fn int_add(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchOps::add(lhs, rhs) + } + + fn int_add_scalar(lhs: TchTensor, rhs: Scalar) -> TchTensor { + lhs.unary_ops( + |mut tensor| tensor.f_add_scalar_(rhs.elem::()).unwrap(), + |tensor| tensor.f_add_scalar(rhs.elem::()).unwrap(), + ) + } + + fn int_sub(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchOps::sub(lhs, rhs) + } + + fn int_sub_scalar(lhs: TchTensor, rhs: Scalar) -> TchTensor { + lhs.unary_ops( + |mut tensor| tensor.f_sub_scalar_(rhs.elem::()).unwrap(), + |tensor| tensor.f_sub_scalar(rhs.elem::()).unwrap(), + ) + } + + fn int_mul(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchOps::mul(lhs, rhs) + } + + fn int_mul_scalar(lhs: TchTensor, rhs: Scalar) -> TchTensor { + lhs.unary_ops( + |mut tensor| tensor.f_mul_scalar_(rhs.elem::()).unwrap(), + |tensor| tensor.f_mul_scalar(rhs.elem::()).unwrap(), + ) + } + + fn int_div(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + let dtype = lhs.tensor.kind(); + let copy = false; + let non_blocking = true; + let lhs: TchTensor = + TchTensor::new(lhs.tensor.to_dtype(tch::Kind::Float, non_blocking, copy)); + let rhs: TchTensor = + TchTensor::new(rhs.tensor.to_dtype(tch::Kind::Float, non_blocking, copy)); + + let out = TchOps::div(lhs, rhs); + + TchTensor::new(out.tensor.to_dtype(dtype, non_blocking, copy)) + } + + fn int_div_scalar(lhs: TchTensor, rhs: Scalar) -> TchTensor { + let dtype = lhs.tensor.kind(); + let copy = false; + let non_blocking = true; + let lhs: TchTensor = + TchTensor::new(lhs.tensor.to_dtype(tch::Kind::Float, non_blocking, copy)); + + let out: TchTensor = lhs.unary_ops( + |mut tensor| tensor.f_div_scalar_(rhs.elem::()).unwrap(), + |tensor| tensor.f_div_scalar(rhs.elem::()).unwrap(), + ); + + TchTensor::new(out.tensor.to_dtype(dtype, non_blocking, copy)) + } + + fn int_remainder(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + let dtype = lhs.tensor.kind(); + let copy = false; + let non_blocking = true; + let lhs: TchTensor = + TchTensor::new(lhs.tensor.to_dtype(tch::Kind::Float, non_blocking, copy)); + let rhs: TchTensor = + TchTensor::new(rhs.tensor.to_dtype(tch::Kind::Float, non_blocking, copy)); + + let out = TchOps::remainder(lhs, rhs); + + TchTensor::new(out.tensor.to_dtype(dtype, non_blocking, copy)) + } + + fn int_remainder_scalar(lhs: TchTensor, rhs: Scalar) -> TchTensor { + lhs.unary_ops( + |tensor| tensor.f_remainder(rhs.elem::()).unwrap(), + |tensor| tensor.f_remainder(rhs.elem::()).unwrap(), + ) + } + + fn int_zeros(shape: Shape, device: &LibTorchDevice, dtype: IntDType) -> TchTensor { + let shape = TchShape::from(shape); + let device: tch::Device = (*device).into(); + + TchTensor::new(tch::Tensor::zeros(shape.dims, (dtype.into_kind(), device))) + } + + fn int_ones(shape: Shape, device: &LibTorchDevice, dtype: IntDType) -> TchTensor { + let shape = TchShape::from(shape); + let device: tch::Device = (*device).into(); + + TchTensor::new(tch::Tensor::ones(shape.dims, (dtype.into_kind(), device))) + } + + fn int_full( + shape: Shape, + fill_value: Scalar, + device: &LibTorchDevice, + dtype: IntDType, + ) -> TchTensor { + let shape = TchShape::from(shape); + let device: tch::Device = (*device).into(); + + TchTensor::new(tch::Tensor::full( + shape.dims, + fill_value.elem::(), + (dtype.into_kind(), device), + )) + } + + fn int_sum(tensor: TchTensor) -> TchTensor { + TchOps::sum(tensor) + } + + fn int_sum_dim(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::sum_dim(tensor, dim) + } + + fn int_prod(tensor: TchTensor) -> TchTensor { + TchOps::prod(tensor) + } + + fn int_prod_dim(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::prod_dim(tensor, dim) + } + + fn int_mean(tensor: TchTensor) -> TchTensor { + let dtype = tensor.tensor.kind(); + let tensor: TchTensor = + TchTensor::new(tensor.tensor.to_dtype(tch::Kind::Float, true, false)); + let output: TchTensor = TchTensor::new(TchOps::mean(tensor).tensor); + + TchTensor::new(output.tensor.to_dtype(dtype, true, false)) + } + + fn int_mean_dim(tensor: TchTensor, dim: usize) -> TchTensor { + let dtype = tensor.tensor.kind(); + let tensor: TchTensor = + TchTensor::new(tensor.tensor.to_dtype(tch::Kind::Float, true, false)); + + let output: TchTensor = TchTensor::new(TchOps::mean_dim(tensor, dim).tensor); + + TchTensor::new(output.tensor.to_dtype(dtype, true, false)) + } + + fn int_cumsum(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::cumsum(tensor, dim) + } + + fn int_cumprod(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::cumprod(tensor, dim) + } + + fn int_cummin(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::cummin(tensor, dim) + } + + fn int_cummax(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::cummax(tensor, dim) + } + + fn int_gather(dim: usize, tensor: TchTensor, indices: TchTensor) -> TchTensor { + TchOps::gather(dim, tensor, indices) + } + + fn int_scatter_add( + dim: usize, + tensor: TchTensor, + indices: TchTensor, + value: TchTensor, + ) -> TchTensor { + TchOps::scatter(dim, tensor, indices, value) + } + + fn int_scatter_nd( + data: TchTensor, + indices: TchTensor, + values: TchTensor, + reduction: burn_backend::tensor::IndexingUpdateOp, + ) -> TchTensor { + TchOps::scatter_nd(data, indices, values, reduction) + } + + fn int_gather_nd(data: TchTensor, indices: TchTensor) -> TchTensor { + TchOps::gather_nd(data, indices) + } + + fn int_select(tensor: TchTensor, dim: usize, indices: TchTensor) -> TchTensor { + TchOps::index_select_dim(tensor, dim, indices) + } + + fn int_select_add( + tensor: TchTensor, + dim: usize, + indices: TchTensor, + value: TchTensor, + ) -> TchTensor { + TchOps::select_assign(tensor, dim, indices, value) + } + + fn int_mask_where(tensor: TchTensor, mask: TchTensor, source: TchTensor) -> TchTensor { + TchTensor::binary_ops_tensor( + tensor, + source, + |tensor, source| source.f_where_self(&mask.tensor, tensor).unwrap(), + |tensor, source| source.f_where_self(&mask.tensor, tensor).unwrap(), + |tensor, source| source.f_where_self(&mask.tensor, tensor).unwrap(), + ) + } + + fn int_mask_fill(tensor: TchTensor, mask: TchTensor, value: Scalar) -> TchTensor { + let value = value.elem::(); + tensor.unary_ops( + |mut tensor| tensor.f_masked_fill_(&mask.tensor, value).unwrap(), + |tensor| tensor.f_masked_fill(&mask.tensor, value).unwrap(), + ) + } + + fn int_argmax(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::argmax(tensor, dim) + } + + fn int_argtopk(_tensor: TchTensor, _dim: usize, _k: usize) -> TchTensor { + panic!("argtopk not implemented for torch") + } + + fn int_topk(tensor: TchTensor, dim: usize, k: usize) -> TchTensor { + TchOps::topk(tensor, dim, k) + } + + fn int_argmin(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::argmin(tensor, dim) + } + + fn int_max_dim(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::max_dim(tensor, dim) + } + + fn int_max_dim_with_indices(tensor: TchTensor, dim: usize) -> (TchTensor, TchTensor) { + TchOps::max_dim_with_indices(tensor, dim) + } + + fn int_min_dim(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::min_dim(tensor, dim) + } + + fn int_min_dim_with_indices(tensor: TchTensor, dim: usize) -> (TchTensor, TchTensor) { + TchOps::min_dim_with_indices(tensor, dim) + } + + fn int_clamp_min(tensor: TchTensor, min: Scalar) -> TchTensor { + TchOps::clamp_min(tensor, min.elem::()) + } + + fn int_clamp_max(tensor: TchTensor, max: Scalar) -> TchTensor { + TchOps::clamp_max(tensor, max.elem::()) + } + + fn int_clamp(tensor: TchTensor, min: Scalar, max: Scalar) -> TchTensor { + TchOps::clamp(tensor, min.elem::(), max.elem::()) + } + + fn int_abs(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.abs_(), |tensor| tensor.abs()) + } + + fn int_into_float(tensor: TchTensor, out_dtype: FloatDType) -> TchTensor { + let tensor = tensor.tensor.to_kind(out_dtype.into_kind()); + TchTensor::new(tensor) + } + + fn int_swap_dims(tensor: IntTensor, dim1: usize, dim2: usize) -> IntTensor { + TchOps::swap_dims(tensor, dim1, dim2) + } + + fn int_random( + shape: Shape, + distribution: Distribution, + device: &LibTorchDevice, + dtype: IntDType, + ) -> TchTensor { + match distribution { + Distribution::Default => TchTensor::new(tch::Tensor::randint_low( + 0, + 255, + shape.iter().map(|i| *i as i64).collect::>(), + (dtype.into_kind(), (*device).into()), + )), + Distribution::Bernoulli(prob) => { + let mut tensor = TchTensor::empty(shape, *device, dtype.into()); + tensor + .mut_ops(|tensor| tensor.f_bernoulli_float_(prob).unwrap()) + .unwrap() + } + Distribution::Uniform(from, to) => TchTensor::new(tch::Tensor::randint_low( + from as i64, + to as i64, + shape.iter().map(|i| *i as i64).collect::>(), + (dtype.into_kind(), (*device).into()), + )), + Distribution::Normal(mean, std) => { + let mut tensor = TchTensor::empty(shape, *device, dtype.into()); + tensor.mut_ops(|tensor| tensor.normal_(mean, std)).unwrap() + } + } + } + + fn int_arange(range: Range, device: &LibTorchDevice, dtype: IntDType) -> TchTensor { + let device: tch::Device = (*device).into(); + let mut tensor = tch::Tensor::arange(range.end - range.start, (dtype.into_kind(), device)); + + if range.start != 0 { + tensor = tensor.f_add_scalar_(range.start).unwrap(); + } + + TchTensor::new(tensor) + } + + fn int_permute(tensor: IntTensor, axes: &[usize]) -> IntTensor { + TchOps::permute(tensor, axes) + } + + fn int_flip(tensor: IntTensor, axes: &[usize]) -> IntTensor { + TchOps::flip(tensor, axes) + } + + fn int_sign(tensor: IntTensor) -> IntTensor { + TchOps::sign(tensor) + } + + fn int_expand(tensor: IntTensor, shape: Shape) -> IntTensor { + TchOps::expand(tensor, shape) + } + + fn int_sort(tensor: IntTensor, dim: usize, descending: bool) -> IntTensor { + TchOps::sort(tensor, dim, descending) + } + + fn int_argsort(tensor: IntTensor, dim: usize, descending: bool) -> IntTensor { + TchOps::argsort(tensor, dim, descending) + } + + fn bitwise_and(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + TchOps::bitwise_and(lhs, rhs) + } + + fn bitwise_or(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + TchOps::bitwise_or(lhs, rhs) + } + + fn bitwise_xor(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + TchOps::bitwise_xor(lhs, rhs) + } + + fn bitwise_not(tensor: IntTensor) -> IntTensor { + TchOps::bitwise_not(tensor) + } + + fn bitwise_and_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + TchOps::bitwise_and_scalar(lhs, rhs.elem::()) + } + + fn bitwise_or_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + TchOps::bitwise_or_scalar(lhs, rhs.elem::()) + } + + fn bitwise_xor_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + TchOps::bitwise_xor_scalar(lhs, rhs.elem::()) + } + + fn bitwise_left_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + TchOps::bitwise_left_shift(lhs, rhs) + } + + fn bitwise_right_shift(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + TchOps::bitwise_right_shift(lhs, rhs) + } + + fn bitwise_left_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + TchOps::bitwise_left_shift_scalar(lhs, rhs.elem::()) + } + + fn bitwise_right_shift_scalar(lhs: IntTensor, rhs: Scalar) -> IntTensor { + TchOps::bitwise_right_shift_scalar(lhs, rhs.elem::()) + } + + fn int_cast(tensor: IntTensor, dtype: IntDType) -> IntTensor { + // NOTE: when dtypes of inputs to an arithmetic operation differ, tch handles type + // promotion based on a set of rules: https://pytorch.org/docs/stable/tensor_attributes.html#type-promotion-doc + + // Type promotion is not automatic on all backends so this behavior might differ + let kind = dtype.into_kind(); + + if tensor.tensor.kind() == kind { + tensor + } else { + TchTensor::new(tensor.tensor.to_kind(kind)) + } + } + + fn int_unfold( + tensor: IntTensor, + dim: usize, + size: usize, + step: usize, + ) -> IntTensor { + TchOps::unfold(tensor, dim, size, step) + } + + fn int_powi(lhs: IntTensor, rhs: IntTensor) -> IntTensor { + TchOps::pow(lhs, rhs) + } + + fn int_powi_scalar_impl(lhs: IntTensor, rhs: Scalar) -> IntTensor { + lhs.unary_ops( + |mut tensor| tensor.f_pow_(rhs.elem::()).unwrap(), + |tensor| tensor.pow_tensor_scalar(rhs.elem::()), + ) + } +} diff --git a/crates/burn-tch/src/ops/mod.rs b/crates/burn-tch/src/ops/mod.rs new file mode 100644 index 0000000..cebe9d2 --- /dev/null +++ b/crates/burn-tch/src/ops/mod.rs @@ -0,0 +1,10 @@ +mod activation; +mod base; +mod bool_tensor; +mod int_tensor; +mod module; +mod qtensor; +mod tensor; +mod transaction; + +pub(crate) use base::*; diff --git a/crates/burn-tch/src/ops/module.rs b/crates/burn-tch/src/ops/module.rs new file mode 100644 index 0000000..c5f9e36 --- /dev/null +++ b/crates/burn-tch/src/ops/module.rs @@ -0,0 +1,642 @@ +use crate::{IntoKind, LibTorch, TchTensor}; +use burn_backend::{ + IntDType, TensorMetadata, + ops::{ + AttentionModuleOptions, ConvOptions, ConvTransposeOptions, DeformConv2dBackward, + DeformConvOptions, InterpolateMode, InterpolateOptions, MaxPool1dWithIndices, + MaxPool2dBackward, MaxPool2dWithIndices, ModuleOps, attention::attention_fallback, + }, + tensor::{FloatTensor, IntTensor}, +}; + +impl ModuleOps for LibTorch { + fn embedding(weights: TchTensor, indices: TchTensor) -> TchTensor { + // Workaround for MPS "Placeholder storage has not been allocated" error. + // See: https://github.com/pytorch/pytorch/issues/123995 + // MPS uses lazy allocation and the embedding operation (which uses index_select) + // can fail if the tensors haven't been materialized yet. + // We work around this by performing the embedding on CPU and transferring back to MPS. + if matches!(weights.tensor.device(), tch::Device::Mps) { + let cpu_weights = weights.tensor.to(tch::Device::Cpu); + let cpu_indices = indices.tensor.to(tch::Device::Cpu); + let result = tch::Tensor::embedding(&cpu_weights, &cpu_indices, -1, false, false) + .to(tch::Device::Mps); + return TchTensor::new(result); + } + + let tensor = tch::Tensor::embedding(&weights.tensor, &indices.tensor, -1, false, false); + TchTensor::new(tensor) + } + + fn embedding_backward(weights: TchTensor, output: TchTensor, indices: TchTensor) -> TchTensor { + let [n_embedding, _d_model] = weights.shape().dims(); + + // Workaround for MPS "Placeholder storage has not been allocated" error. + // See: https://github.com/pytorch/pytorch/issues/123995 + if matches!(output.tensor.device(), tch::Device::Mps) { + let cpu_output = output.tensor.to(tch::Device::Cpu); + let cpu_indices = indices.tensor.to(tch::Device::Cpu); + let result = tch::Tensor::embedding_backward( + &cpu_output, + &cpu_indices, + n_embedding as i64, + -1, + false, + false, + ) + .to(tch::Device::Mps); + return TchTensor::new(result); + } + + let tensor = tch::Tensor::embedding_backward( + &output.tensor, + &indices.tensor, + n_embedding as i64, + -1, + false, + false, + ); + + TchTensor::new(tensor) + } + + fn conv1d( + x: TchTensor, + weight: TchTensor, + bias: Option, + options: ConvOptions<1>, + ) -> TchTensor { + let tensor = tch::Tensor::conv1d( + &x.tensor, + &weight.tensor, + bias.map(|t| t.tensor), + options.stride.map(|i| i as i64), + options.padding.map(|i| i as i64), + options.dilation.map(|i| i as i64), + options.groups as i64, + ); + + TchTensor::new(tensor) + } + + fn conv2d( + x: TchTensor, + weight: TchTensor, + bias: Option, + options: ConvOptions<2>, + ) -> TchTensor { + let tensor = tch::Tensor::conv2d( + &x.tensor, + &weight.tensor, + bias.map(|t| t.tensor), + options.stride.map(|i| i as i64), + options.padding.map(|i| i as i64), + options.dilation.map(|i| i as i64), + options.groups as i64, + ); + + TchTensor::new(tensor) + } + + fn conv3d( + x: TchTensor, + weight: TchTensor, + bias: Option, + options: ConvOptions<3>, + ) -> TchTensor { + let tensor = tch::Tensor::conv3d( + &x.tensor, + &weight.tensor, + bias.map(|t| t.tensor), + options.stride.map(|i| i as i64), + options.padding.map(|i| i as i64), + options.dilation.map(|i| i as i64), + options.groups as i64, + ); + + TchTensor::new(tensor) + } + + fn deform_conv2d( + _x: TchTensor, + _offset: TchTensor, + _weight: TchTensor, + _mask: Option, + _bias: Option, + _options: DeformConvOptions<2>, + ) -> TchTensor { + unimplemented!("Torch bindings don't support deform_conv2d"); + } + + fn deform_conv2d_backward( + _x: TchTensor, + _offset: TchTensor, + _weight: TchTensor, + _mask: Option, + _bias: Option, + _out_grad: TchTensor, + _options: DeformConvOptions<2>, + ) -> DeformConv2dBackward { + unimplemented!("Torch bindings don't support deform_conv2d"); + } + + fn conv_transpose1d( + x: TchTensor, + weight: TchTensor, + bias: Option, + options: ConvTransposeOptions<1>, + ) -> TchTensor { + let tensor = tch::Tensor::conv_transpose1d( + &x.tensor, + &weight.tensor, + bias.map(|t| t.tensor), + options.stride.map(|i| i as i64), + options.padding.map(|i| i as i64), + options.padding_out.map(|i| i as i64), + options.groups as i64, + options.dilation.map(|i| i as i64), + ); + + TchTensor::new(tensor) + } + + fn conv_transpose2d( + x: TchTensor, + weight: TchTensor, + bias: Option, + options: ConvTransposeOptions<2>, + ) -> TchTensor { + let tensor = tch::Tensor::conv_transpose2d( + &x.tensor, + &weight.tensor, + bias.map(|t| t.tensor), + options.stride.map(|i| i as i64), + options.padding.map(|i| i as i64), + options.padding_out.map(|i| i as i64), + options.groups as i64, + options.dilation.map(|i| i as i64), + ); + + TchTensor::new(tensor) + } + + fn conv_transpose3d( + x: TchTensor, + weight: TchTensor, + bias: Option, + options: ConvTransposeOptions<3>, + ) -> TchTensor { + let tensor = tch::Tensor::conv_transpose3d( + &x.tensor, + &weight.tensor, + bias.map(|t| t.tensor), + options.stride.map(|i| i as i64), + options.padding.map(|i| i as i64), + options.padding_out.map(|i| i as i64), + options.groups as i64, + options.dilation.map(|i| i as i64), + ); + + TchTensor::new(tensor) + } + + fn avg_pool1d( + x: TchTensor, + kernel_size: usize, + stride: usize, + padding: usize, + count_include_pad: bool, + ceil_mode: bool, + ) -> TchTensor { + let tensor = tch::Tensor::avg_pool1d( + &x.tensor, + [kernel_size as i64], + [stride as i64], + [padding as i64], + ceil_mode, + count_include_pad, + ); + + TchTensor::new(tensor) + } + fn avg_pool2d( + x: TchTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> TchTensor { + let tensor = tch::Tensor::avg_pool2d( + &x.tensor, + [kernel_size[0] as i64, kernel_size[1] as i64], + [stride[0] as i64, stride[1] as i64], + [padding[0] as i64, padding[1] as i64], + ceil_mode, + count_include_pad, + None, + ); + + TchTensor::new(tensor) + } + + fn avg_pool2d_backward( + x: TchTensor, + grad: TchTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, + ) -> TchTensor { + let tensor = tch::Tensor::avg_pool2d_backward( + &x.tensor, + &grad.tensor, + [kernel_size[0] as i64, kernel_size[1] as i64], + [stride[0] as i64, stride[1] as i64], + [padding[0] as i64, padding[1] as i64], + ceil_mode, + count_include_pad, + None, + ); + + TchTensor::new(tensor) + } + + fn max_pool1d( + x: TchTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + ) -> TchTensor { + let tensor = tch::Tensor::max_pool1d( + &x.tensor, + kernel_size as i64, + stride as i64, + padding as i64, + dilation as i64, + ceil_mode, + ); + + TchTensor::new(tensor) + } + + fn max_pool1d_with_indices( + x: TchTensor, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, + indices_dtype: IntDType, + ) -> MaxPool1dWithIndices { + let (tensor, indices) = tch::Tensor::max_pool1d_with_indices( + &x.tensor, + kernel_size as i64, + stride as i64, + padding as i64, + dilation as i64, + ceil_mode, + ); + + MaxPool1dWithIndices::new( + TchTensor::new(tensor), + TchTensor::new(indices.to_kind(indices_dtype.into_kind())), + ) + } + + fn max_pool2d( + x: TchTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + ) -> TchTensor { + let tensor = tch::Tensor::max_pool2d( + &x.tensor, + [kernel_size[0] as i64, kernel_size[1] as i64], + [stride[0] as i64, stride[1] as i64], + [padding[0] as i64, padding[1] as i64], + [dilation[0] as i64, dilation[1] as i64], + ceil_mode, + ); + + TchTensor::new(tensor) + } + + fn max_pool2d_with_indices( + x: TchTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + indices_dtype: IntDType, + ) -> MaxPool2dWithIndices { + let (tensor, indices) = tch::Tensor::max_pool2d_with_indices( + &x.tensor, + [kernel_size[0] as i64, kernel_size[1] as i64], + [stride[0] as i64, stride[1] as i64], + [padding[0] as i64, padding[1] as i64], + [dilation[0] as i64, dilation[1] as i64], + ceil_mode, + ); + + MaxPool2dWithIndices::new( + TchTensor::new(tensor), + TchTensor::new(indices.to_kind(indices_dtype.into_kind())), + ) + } + + fn max_pool2d_with_indices_backward( + x: TchTensor, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + output_grad: TchTensor, + indices: TchTensor, + ) -> MaxPool2dBackward { + let grad = tch::Tensor::max_pool2d_with_indices_backward( + &x.tensor, + &output_grad.tensor, + [kernel_size[0] as i64, kernel_size[1] as i64], + [stride[0] as i64, stride[1] as i64], + [padding[0] as i64, padding[1] as i64], + [dilation[0] as i64, dilation[1] as i64], + ceil_mode, + &indices.tensor, + ); + + MaxPool2dBackward::new(TchTensor::new(grad)) + } + + fn adaptive_avg_pool2d(x: TchTensor, output_size: [usize; 2]) -> TchTensor { + let tensor = tch::Tensor::adaptive_avg_pool2d(&x.tensor, output_size.map(|e| e as i64)); + + TchTensor::new(tensor) + } + + fn adaptive_avg_pool2d_backward(x: TchTensor, grad: TchTensor) -> TchTensor { + let tensor = tch::Tensor::internal_adaptive_avg_pool2d_backward(&x.tensor, &grad.tensor); + + TchTensor::new(tensor) + } + + fn adaptive_avg_pool1d(x: TchTensor, output_size: usize) -> TchTensor { + let tensor = tch::Tensor::adaptive_avg_pool1d(&x.tensor, output_size as i64); + + TchTensor::new(tensor) + } + + fn interpolate( + x: TchTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> TchTensor { + let output_size = output_size.map(|e| e as i64); + + let align_corners = options.align_corners; + let tensor = match options.mode { + InterpolateMode::Nearest => { + tch::Tensor::upsample_nearest2d(&x.tensor, output_size, None, None) + } + InterpolateMode::NearestExact => { + panic!("nearest exact interpolation is not supported by PyTorch/tch backend") + } + InterpolateMode::Bilinear => { + tch::Tensor::upsample_bilinear2d(&x.tensor, output_size, align_corners, None, None) + } + InterpolateMode::Bicubic => { + tch::Tensor::upsample_bicubic2d(&x.tensor, output_size, align_corners, None, None) + } + InterpolateMode::Lanczos3 => { + panic!("lanczos3 interpolation is not supported by PyTorch/tch backend") + } + }; + + TchTensor::new(tensor) + } + + fn interpolate_backward( + x: TchTensor, + grad: TchTensor, + output_size: [usize; 2], + options: InterpolateOptions, + ) -> TchTensor { + let output_size = output_size.map(|e| e as i64); + let [n, c, h_in, w_in] = x.shape().dims(); + let input_size = [n as i64, c as i64, h_in as i64, w_in as i64]; + let align_corners = options.align_corners; + + let tensor = match options.mode { + InterpolateMode::Nearest => tch::Tensor::upsample_nearest2d_backward( + &grad.tensor, + output_size, + input_size, + None, + None, + ), + InterpolateMode::NearestExact => { + panic!( + "nearest exact interpolation backward is not supported by PyTorch/tch backend" + ) + } + InterpolateMode::Bilinear => tch::Tensor::upsample_bilinear2d_backward( + &grad.tensor, + output_size, + input_size, + align_corners, + None, + None, + ), + InterpolateMode::Bicubic => tch::Tensor::upsample_bicubic2d_backward( + &grad.tensor, + output_size, + input_size, + align_corners, + None, + None, + ), + InterpolateMode::Lanczos3 => { + panic!("lanczos3 interpolation backward is not supported by PyTorch/tch backend") + } + }; + + TchTensor::new(tensor) + } + + fn attention( + query: TchTensor, + key: TchTensor, + value: TchTensor, + mask: Option, + attn_bias: Option, + options: AttentionModuleOptions, + ) -> TchTensor { + if attn_bias.is_some() { + return attention_fallback::(query, key, value, mask, attn_bias, options); + } + + TchTensor::new(tch::Tensor::scaled_dot_product_attention( + &query.tensor, + &key.tensor, + &value.tensor, + mask.map(|m| m.tensor), + 0., + options.is_causal, + options.scale, + false, + )) + } + + fn layer_norm( + tensor: TchTensor, + gamma: TchTensor, + beta: Option, + epsilon: f64, + ) -> TchTensor { + let shape = tensor.shape(); + let last_dim = shape[shape.num_dims() - 1] as i64; + + let tensor = tensor.tensor.layer_norm( + [last_dim], + Some(&gamma.tensor), + beta.as_ref().map(|b| &b.tensor), + epsilon, + true, + ); + + TchTensor::new(tensor) + } + + fn has_ctc_loss_backward() -> bool { + true + } + + fn ctc_loss( + log_probs: FloatTensor, + targets: IntTensor, + input_lengths: IntTensor, + target_lengths: IntTensor, + blank: usize, + ) -> FloatTensor { + // PyTorch's CTC requires int64 for targets and length tensors. + let targets_i64 = targets.tensor.to_kind(tch::Kind::Int64); + let input_lengths_i64 = input_lengths.tensor.to_kind(tch::Kind::Int64); + let target_lengths_i64 = target_lengths.tensor.to_kind(tch::Kind::Int64); + + // Reduction::None returns per-sample losses [N], matching the trait contract. + let tensor = tch::Tensor::ctc_loss_tensor( + &log_probs.tensor, + &targets_i64, + &input_lengths_i64, + &target_lengths_i64, + blank as i64, + tch::Reduction::None, + false, + ); + + TchTensor::new(tensor) + } + + fn ctc_loss_backward( + log_probs: FloatTensor, + targets: IntTensor, + input_lengths: IntTensor, + target_lengths: IntTensor, + grad_loss: FloatTensor, + blank: usize, + ) -> FloatTensor { + let targets_i64 = targets.tensor.to_kind(tch::Kind::Int64); + let input_lengths_i64 = input_lengths.tensor.to_kind(tch::Kind::Int64); + let target_lengths_i64 = target_lengths.tensor.to_kind(tch::Kind::Int64); + + // Recompute forward to get neg_log_likelihood and log_alpha (LibTorch's + // backward needs both). PyTorch caches log_alpha during the autograd + // forward; our trait has no caching slot for it, so we redo the alpha + // recursion here. This is still a single-call into LibTorch's fused + // kernel and avoids the ~40T host-side dispatches. + let (neg_log_likelihood, log_alpha) = tch::Tensor::internal_ctc_loss_tensor( + &log_probs.tensor, + &targets_i64, + &input_lengths_i64, + &target_lengths_i64, + blank as i64, + false, + ); + + let grad = tch::Tensor::internal_ctc_loss_backward_tensor( + &grad_loss.tensor, + &log_probs.tensor, + &targets_i64, + &input_lengths_i64, + &target_lengths_i64, + &neg_log_likelihood, + &log_alpha, + blank as i64, + false, + ); + + TchTensor::new(grad) + } + + fn rfft( + signal: FloatTensor, + dim: usize, + n: Option, + ) -> (FloatTensor, FloatTensor) { + let complex = signal + .tensor + .fft_rfft(n.map(|v| v as i64), dim as i64, "backward"); + let re = TchTensor::new(complex.real().contiguous()); + let im = TchTensor::new(complex.imag().contiguous()); + (re, im) + } + + fn irfft( + spectrum_re: FloatTensor, + spectrum_im: FloatTensor, + dim: usize, + n: Option, + ) -> FloatTensor { + let complex = tch::Tensor::complex(&spectrum_re.tensor, &spectrum_im.tensor); + TchTensor::new(complex.fft_irfft(n.map(|v| v as i64), dim as i64, "backward")) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn_backend::{ + TensorData, Tolerance, + ops::{FloatTensorOps, IntTensorOps}, + read_sync, + }; + + type B = crate::LibTorch; + + #[test] + fn ctc_loss_uniform() { + // T=3, N=1, C=2, blank=0, target=[1, 1]. + // Only valid alignment is (1, 0, 1) with prob (1/2)^3. + // Loss = -ln(1/8) = 3 * ln(2) + let device = Default::default(); + let log_probs_data = vec![(0.5f32).ln(); 3 * 2]; + let log_probs = B::float_from_data(TensorData::new(log_probs_data, [3, 1, 2]), &device); + let targets = B::int_from_data(TensorData::from([[1i64, 1]]), &device); + let input_lengths = B::int_from_data(TensorData::from([3i64]), &device); + let target_lengths = B::int_from_data(TensorData::from([2i64]), &device); + + let loss = + >::ctc_loss(log_probs, targets, input_lengths, target_lengths, 0); + + let out = read_sync(B::float_into_data(loss)).unwrap(); + let expected = TensorData::from([3.0f32 * 2.0f32.ln()]); + out.assert_approx_eq::(&expected, Tolerance::rel_abs(1e-3, 1e-3)); + } +} diff --git a/crates/burn-tch/src/ops/qtensor.rs b/crates/burn-tch/src/ops/qtensor.rs new file mode 100644 index 0000000..201a236 --- /dev/null +++ b/crates/burn-tch/src/ops/qtensor.rs @@ -0,0 +1,129 @@ +use burn_backend::{ + ExecutionError, FloatDType, IntDType, Shape, TensorData, + ops::QTensorOps, + quantization::{QuantScheme, QuantizationParametersPrimitive}, + tensor::{Device, FloatTensor, IntTensor, QuantizedTensor}, +}; + +use crate::{LibTorch, LibTorchDevice}; + +impl QTensorOps for LibTorch { + fn q_from_data(_data: TensorData, _device: &LibTorchDevice) -> QuantizedTensor { + unimplemented!() + } + + fn quantize( + _tensor: FloatTensor, + _scheme: &QuantScheme, + _qparams: QuantizationParametersPrimitive, + ) -> QuantizedTensor { + unimplemented!() + } + + fn quantize_dynamic( + _tensor: FloatTensor, + _scheme: &QuantScheme, + ) -> QuantizedTensor { + unimplemented!() + } + + fn dequantize(_tensor: QuantizedTensor, _dtype: FloatDType) -> FloatTensor { + unimplemented!() + } + + fn q_to_device( + _tensor: QuantizedTensor, + _device: &Device, + ) -> QuantizedTensor { + unimplemented!() + } + + fn q_reshape(_tensor: QuantizedTensor, _shape: Shape) -> QuantizedTensor { + unimplemented!() + } + + async fn q_into_data(_tensor: QuantizedTensor) -> Result { + unimplemented!() + } + fn q_swap_dims( + _tensor: QuantizedTensor, + _dim1: usize, + _dim2: usize, + ) -> QuantizedTensor { + unimplemented!() + } + + fn q_permute(_tensor: QuantizedTensor, _axes: &[usize]) -> QuantizedTensor { + unimplemented!() + } + + fn q_flip(_tensor: QuantizedTensor, _axes: &[usize]) -> QuantizedTensor { + unimplemented!() + } + + fn q_argmax( + _tensor: QuantizedTensor, + _dim: usize, + _out_dtype: IntDType, + ) -> IntTensor { + unimplemented!() + } + + fn q_argmin( + _tensor: QuantizedTensor, + _dim: usize, + _out_dtype: IntDType, + ) -> IntTensor { + unimplemented!() + } + + fn q_max_dim_with_indices( + _tensor: QuantizedTensor, + _dim: usize, + _indices_dtype: IntDType, + ) -> (QuantizedTensor, IntTensor) { + unimplemented!() + } + + fn q_max_dim(_tensor: QuantizedTensor, _dim: usize) -> QuantizedTensor { + unimplemented!() + } + + fn q_min_dim(_tensor: QuantizedTensor, _dim: usize) -> QuantizedTensor { + unimplemented!() + } + + fn q_min_dim_with_indices( + _tensor: QuantizedTensor, + _dim: usize, + _indices_dtype: IntDType, + ) -> (QuantizedTensor, IntTensor) { + unimplemented!() + } + + fn q_sort( + _tensor: QuantizedTensor, + _dim: usize, + _descending: bool, + ) -> QuantizedTensor { + unimplemented!() + } + + fn q_sort_with_indices( + _tensor: QuantizedTensor, + _dim: usize, + _descending: bool, + _indices_dtype: IntDType, + ) -> (QuantizedTensor, IntTensor) { + unimplemented!() + } + + fn q_argsort( + _tensor: QuantizedTensor, + _dim: usize, + _descending: bool, + _out_dtype: IntDType, + ) -> IntTensor { + unimplemented!() + } +} diff --git a/crates/burn-tch/src/ops/tensor.rs b/crates/burn-tch/src/ops/tensor.rs new file mode 100644 index 0000000..2892c43 --- /dev/null +++ b/crates/burn-tch/src/ops/tensor.rs @@ -0,0 +1,576 @@ +use super::TchOps; +use crate::{IntoKind, LibTorch, LibTorchDevice, TchShape, TchTensor}; +use burn_backend::backend::ExecutionError; +use burn_backend::tensor::{BoolTensor, FloatTensor, IntTensor}; +use burn_backend::{BoolDType, IntDType, Scalar, bf16, f16}; +use burn_backend::{ + DType, Distribution, FloatDType, Shape, TensorData, TensorMetadata, ops::FloatTensorOps, +}; + +impl FloatTensorOps for LibTorch { + fn float_from_data(data: TensorData, device: &LibTorchDevice) -> TchTensor { + match data.dtype { + DType::F64 => TchTensor::from_data::(data, (*device).into()), + DType::F32 => TchTensor::from_data::(data, (*device).into()), + DType::F16 => TchTensor::from_data::(data, (*device).into()), + DType::BF16 => TchTensor::from_data::(data, (*device).into()), + _ => unimplemented!("Unsupported dtype for `float_from_data`"), + } + } + + fn float_random( + shape: Shape, + distribution: Distribution, + device: &LibTorchDevice, + dtype: FloatDType, + ) -> TchTensor { + match distribution { + Distribution::Default => { + let mut tensor = TchTensor::empty(shape, *device, dtype.into()); + tensor + .mut_ops(|tensor| tensor.rand_like_out(tensor)) + .unwrap() + } + Distribution::Bernoulli(prob) => { + let mut tensor = TchTensor::empty(shape, *device, dtype.into()); + tensor + .mut_ops(|tensor| tensor.f_bernoulli_float_(prob).unwrap()) + .unwrap() + } + Distribution::Uniform(from, to) => { + let mut tensor = TchTensor::empty(shape, *device, dtype.into()); + tensor.mut_ops(|tensor| tensor.uniform_(from, to)).unwrap() + } + Distribution::Normal(mean, std) => { + let mut tensor = TchTensor::empty(shape, *device, dtype.into()); + tensor.mut_ops(|tensor| tensor.normal_(mean, std)).unwrap() + } + } + } + + fn float_repeat_dim(tensor: TchTensor, dim: usize, times: usize) -> TchTensor { + TchOps::repeat_dim(tensor, dim, times) + } + + fn float_zeros(shape: Shape, device: &LibTorchDevice, dtype: FloatDType) -> TchTensor { + let shape = TchShape::from(shape); + let device: tch::Device = (*device).into(); + + TchTensor::new(tch::Tensor::zeros(shape.dims, (dtype.into_kind(), device))) + } + + fn float_ones(shape: Shape, device: &LibTorchDevice, dtype: FloatDType) -> TchTensor { + let shape = TchShape::from(shape); + let device: tch::Device = (*device).into(); + + TchTensor::new(tch::Tensor::ones(shape.dims, (dtype.into_kind(), device))) + } + + async fn float_into_data(tensor: TchTensor) -> Result { + let shape = tensor.shape(); + let tensor = Self::float_reshape(tensor.clone(), Shape::new([shape.num_elements()])); + Ok(match tensor.tensor.kind() { + tch::Kind::Half => { + let values = Vec::::try_from(&tensor).unwrap(); + TensorData::new(values, shape) + } + tch::Kind::Float => { + let values = Vec::::try_from(&tensor).unwrap(); + TensorData::new(values, shape) + } + tch::Kind::Double => { + let values = Vec::::try_from(&tensor).unwrap(); + TensorData::new(values, shape) + } + tch::Kind::BFloat16 => { + let values = Vec::::try_from(&tensor).unwrap(); + TensorData::new(values, shape) + } + _ => panic!("Not a valid float kind"), + }) + } + + fn float_to_device(tensor: TchTensor, device: &LibTorchDevice) -> TchTensor { + TchOps::to_device(tensor, device) + } + + fn float_empty(shape: Shape, device: &LibTorchDevice, dtype: FloatDType) -> TchTensor { + let tensor = tch::Tensor::empty( + TchShape::from(shape).dims, + (dtype.into_kind(), (*device).into()), + ); + + TchTensor::new(tensor) + } + + fn float_add(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchOps::add(lhs, rhs) + } + + fn float_add_scalar(lhs: TchTensor, rhs: Scalar) -> TchTensor { + let rhs: f64 = rhs.elem(); + + lhs.unary_ops( + |mut tensor| tensor.f_add_scalar_(rhs).unwrap(), + |tensor| tensor.f_add_scalar(rhs).unwrap(), + ) + } + + fn float_sub(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchOps::sub(lhs, rhs) + } + + fn float_sub_scalar(lhs: TchTensor, rhs: Scalar) -> TchTensor { + let rhs: f64 = rhs.elem(); + + lhs.unary_ops( + |mut tensor| tensor.f_sub_scalar_(rhs).unwrap(), + |tensor| tensor.f_sub_scalar(rhs).unwrap(), + ) + } + + fn float_mul(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchOps::mul(lhs, rhs) + } + + fn float_mul_scalar(lhs: TchTensor, rhs: Scalar) -> TchTensor { + let rhs: f64 = rhs.elem(); + + lhs.unary_ops( + |mut tensor| tensor.f_mul_scalar_(rhs).unwrap(), + |tensor| tensor.f_mul_scalar(rhs).unwrap(), + ) + } + + fn float_div(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchOps::div(lhs, rhs) + } + + fn float_div_scalar(lhs: TchTensor, rhs: Scalar) -> TchTensor { + let rhs: f64 = rhs.elem(); + + lhs.unary_ops( + |mut tensor| tensor.f_div_scalar_(rhs).unwrap(), + |tensor| tensor.f_div_scalar(rhs).unwrap(), + ) + } + + fn float_remainder(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchOps::remainder(lhs, rhs) + } + + fn float_remainder_scalar(lhs: TchTensor, rhs: Scalar) -> TchTensor { + let rhs: f64 = rhs.elem(); + + lhs.unary_ops( + |tensor| tensor.f_remainder(rhs).unwrap(), + |tensor| tensor.f_remainder(rhs).unwrap(), + ) + } + + fn float_matmul(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + let tensor = lhs.tensor.matmul(&rhs.tensor); + TchTensor::new(tensor) + } + + fn float_cross(lhs: TchTensor, rhs: TchTensor, dim: usize) -> TchTensor { + let tensor = lhs.tensor.cross(&rhs.tensor, dim as i64); + TchTensor::new(tensor) + } + + fn float_recip(tensor: TchTensor) -> TchTensor { + TchTensor::new(tensor.tensor.reciprocal()) + } + + fn float_swap_dims(tensor: TchTensor, dim1: usize, dim2: usize) -> TchTensor { + TchOps::swap_dims(tensor, dim1, dim2) + } + + fn float_reshape(tensor: TchTensor, shape: Shape) -> TchTensor { + TchOps::reshape(tensor, shape) + } + + fn float_gather(dim: usize, tensor: TchTensor, indices: TchTensor) -> TchTensor { + TchOps::gather(dim, tensor, indices) + } + + fn float_scatter_add( + dim: usize, + tensor: TchTensor, + indices: TchTensor, + value: TchTensor, + ) -> TchTensor { + TchOps::scatter(dim, tensor, indices, value) + } + + fn float_scatter_nd( + data: TchTensor, + indices: TchTensor, + values: TchTensor, + reduction: burn_backend::tensor::IndexingUpdateOp, + ) -> TchTensor { + TchOps::scatter_nd(data, indices, values, reduction) + } + + fn float_gather_nd(data: TchTensor, indices: TchTensor) -> TchTensor { + TchOps::gather_nd(data, indices) + } + + fn float_select(tensor: TchTensor, dim: usize, indices: TchTensor) -> TchTensor { + TchOps::index_select_dim(tensor, dim, indices) + } + + fn float_select_add( + tensor: TchTensor, + dim: usize, + indices: TchTensor, + value: TchTensor, + ) -> TchTensor { + TchOps::select_assign(tensor, dim, indices, value) + } + + fn float_slice(tensor: TchTensor, slices: &[burn_backend::Slice]) -> TchTensor { + TchOps::slice_with_steps(tensor, slices) + } + + fn float_slice_assign( + tensor: TchTensor, + slices: &[burn_backend::Slice], + value: TchTensor, + ) -> TchTensor { + TchOps::slice_assign(tensor, slices, value) + } + + fn float_mask_where(tensor: TchTensor, mask: TchTensor, value: TchTensor) -> TchTensor { + let output = value.tensor.where_self(&mask.tensor, &tensor.tensor); + + TchTensor::new(output) + } + + fn float_mask_fill(tensor: TchTensor, mask: TchTensor, value: Scalar) -> TchTensor { + let value: f64 = value.elem(); + + tensor.unary_ops( + |mut tensor| tensor.f_masked_fill_(&mask.tensor, value).unwrap(), + |tensor| tensor.f_masked_fill(&mask.tensor, value).unwrap(), + ) + } + + fn float_equal(lhs: TchTensor, rhs: TchTensor, _out_dtype: BoolDType) -> TchTensor { + TchOps::equal(lhs, rhs) + } + + fn float_equal_elem(lhs: TchTensor, rhs: Scalar, _out_dtype: BoolDType) -> TchTensor { + TchOps::equal_elem(lhs, rhs.elem::()) + } + + fn float_greater(lhs: TchTensor, rhs: TchTensor, _out_dtype: BoolDType) -> TchTensor { + TchOps::greater(lhs, rhs) + } + + fn float_greater_elem(lhs: TchTensor, rhs: Scalar, _out_dtype: BoolDType) -> TchTensor { + TchOps::greater_elem(lhs, rhs.elem::()) + } + + fn float_greater_equal(lhs: TchTensor, rhs: TchTensor, _out_dtype: BoolDType) -> TchTensor { + TchOps::greater_equal(lhs, rhs) + } + + fn float_greater_equal_elem(lhs: TchTensor, rhs: Scalar, _out_dtype: BoolDType) -> TchTensor { + TchOps::greater_equal_elem(lhs, rhs.elem::()) + } + + fn float_lower(lhs: TchTensor, rhs: TchTensor, _out_dtype: BoolDType) -> TchTensor { + TchOps::lower(lhs, rhs) + } + + fn float_lower_elem(lhs: TchTensor, rhs: Scalar, _out_dtype: BoolDType) -> TchTensor { + TchOps::lower_elem(lhs, rhs.elem::()) + } + + fn float_lower_equal(lhs: TchTensor, rhs: TchTensor, _out_dtype: BoolDType) -> TchTensor { + TchOps::lower_equal(lhs, rhs) + } + + fn float_lower_equal_elem(lhs: TchTensor, rhs: Scalar, _out_dtype: BoolDType) -> TchTensor { + TchOps::lower_equal_elem(lhs, rhs.elem::()) + } + + fn float_mean(tensor: TchTensor) -> TchTensor { + TchOps::mean(tensor) + } + + fn float_sum(tensor: TchTensor) -> TchTensor { + TchOps::sum(tensor) + } + + fn float_sum_dim(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::sum_dim(tensor, dim) + } + + fn float_mean_dim(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::mean_dim(tensor, dim) + } + + fn float_cumsum(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::cumsum(tensor, dim) + } + + fn float_cumprod(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::cumprod(tensor, dim) + } + + fn float_cummin(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::cummin(tensor, dim) + } + + fn float_cummax(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::cummax(tensor, dim) + } + + fn float_prod(tensor: TchTensor) -> TchTensor { + TchOps::prod(tensor) + } + + fn float_prod_dim(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::prod_dim(tensor, dim) + } + + fn float_argmax(tensor: TchTensor, dim: usize, _indices_dtype: IntDType) -> TchTensor { + TchOps::argmax(tensor, dim) + } + + fn float_argtopk( + tensor: TchTensor, + dim: usize, + k: usize, + _indices_dtype: IntDType, + ) -> TchTensor { + TchOps::argtopk(tensor, dim, k) + } + + fn float_topk(tensor: TchTensor, dim: usize, k: usize) -> TchTensor { + TchOps::topk(tensor, dim, k) + } + + fn float_argmin(tensor: TchTensor, dim: usize, _out_dtype: IntDType) -> TchTensor { + TchOps::argmin(tensor, dim) + } + + fn float_max_dim(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::max_dim(tensor, dim) + } + + fn float_max_dim_with_indices( + tensor: TchTensor, + dim: usize, + _indices_dtype: IntDType, + ) -> (TchTensor, TchTensor) { + TchOps::max_dim_with_indices(tensor, dim) + } + + fn float_min_dim(tensor: TchTensor, dim: usize) -> TchTensor { + TchOps::min_dim(tensor, dim) + } + + fn float_min_dim_with_indices( + tensor: TchTensor, + dim: usize, + _indices_dtype: IntDType, + ) -> (TchTensor, TchTensor) { + TchOps::min_dim_with_indices(tensor, dim) + } + + fn float_exp(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.exp_(), |tensor| tensor.exp()) + } + + fn float_log(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.log_(), |tensor| tensor.log()) + } + + fn float_log1p(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.log1p_(), |tensor| tensor.log1p()) + } + + fn float_powf_scalar_impl(tensor: TchTensor, value: Scalar) -> TchTensor { + tensor.unary_ops( + |mut tensor| tensor.f_pow_(value.elem::()).unwrap(), + |tensor| tensor.pow_tensor_scalar(value.elem::()), + ) + } + + fn float_sqrt(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.sqrt_(), |tensor| tensor.sqrt()) + } + + fn float_abs(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.abs_(), |tensor| tensor.abs()) + } + + fn float_cos(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.cos_(), |tensor| tensor.cos()) + } + + fn float_cosh(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.cosh_(), |tensor| tensor.cosh()) + } + + fn float_sin(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.sin_(), |tensor| tensor.sin()) + } + + fn float_sinh(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.sinh_(), |tensor| tensor.sinh()) + } + + fn float_tan(tensor: FloatTensor) -> FloatTensor { + tensor.unary_ops(|mut tensor| tensor.tan_(), |tensor| tensor.tan()) + } + + fn float_tanh(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.tanh_(), |tensor| tensor.tanh()) + } + + fn float_acos(tensor: FloatTensor) -> FloatTensor { + tensor.unary_ops(|mut tensor| tensor.acos_(), |tensor| tensor.acos()) + } + + fn float_acosh(tensor: FloatTensor) -> FloatTensor { + tensor.unary_ops(|mut tensor| tensor.acosh_(), |tensor| tensor.acosh()) + } + + fn float_asin(tensor: FloatTensor) -> FloatTensor { + tensor.unary_ops(|mut tensor| tensor.asin_(), |tensor| tensor.asin()) + } + + fn float_asinh(tensor: FloatTensor) -> FloatTensor { + tensor.unary_ops(|mut tensor| tensor.asinh_(), |tensor| tensor.asinh()) + } + + fn float_atan(tensor: FloatTensor) -> FloatTensor { + tensor.unary_ops(|mut tensor| tensor.atan_(), |tensor| tensor.atan()) + } + + fn float_atanh(tensor: FloatTensor) -> FloatTensor { + tensor.unary_ops(|mut tensor| tensor.atanh_(), |tensor| tensor.atanh()) + } + + fn float_atan2(lhs: FloatTensor, rhs: FloatTensor) -> FloatTensor { + TchOps::atan2(lhs, rhs) + } + + fn float_round(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.round_(), |tensor| tensor.round()) + } + + fn float_floor(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.floor_(), |tensor| tensor.floor()) + } + + fn float_ceil(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.ceil_(), |tensor| tensor.ceil()) + } + + fn float_trunc(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.trunc_(), |tensor| tensor.trunc()) + } + + fn float_erf(tensor: TchTensor) -> TchTensor { + tensor.unary_ops(|mut tensor| tensor.erf_(), |tensor| tensor.erf()) + } + + fn float_cat(tensors: Vec, dim: usize) -> TchTensor { + TchOps::cat(tensors, dim) + } + + fn float_clamp_min(tensor: TchTensor, min: Scalar) -> TchTensor { + TchOps::clamp_min(tensor, min.elem::()) + } + + fn float_clamp_max(tensor: TchTensor, max: Scalar) -> TchTensor { + TchOps::clamp_max(tensor, max.elem::()) + } + + fn float_clamp(tensor: TchTensor, min: Scalar, max: Scalar) -> TchTensor { + TchOps::clamp(tensor, min.elem::(), max.elem::()) + } + + fn float_into_int(tensor: TchTensor, _out_dtype: IntDType) -> TchTensor { + let tensor = tensor.tensor.to_kind(tch::Kind::Int64); + TchTensor::new(tensor) + } + + fn float_powf(lhs: TchTensor, rhs: TchTensor) -> TchTensor { + TchOps::pow(lhs, rhs) + } + + fn float_permute(tensor: TchTensor, axes: &[usize]) -> TchTensor { + TchOps::permute(tensor, axes) + } + + fn float_flip(tensor: TchTensor, axes: &[usize]) -> TchTensor { + TchOps::flip(tensor, axes) + } + + fn float_sign(tensor: TchTensor) -> TchTensor { + TchOps::sign(tensor) + } + + fn float_expand(tensor: TchTensor, shape: Shape) -> TchTensor { + TchOps::expand(tensor, shape) + } + + fn float_sort(tensor: TchTensor, dim: usize, descending: bool) -> TchTensor { + TchOps::sort(tensor, dim, descending) + } + + fn float_sort_with_indices( + tensor: TchTensor, + dim: usize, + descending: bool, + _indices_dtype: IntDType, + ) -> (TchTensor, TchTensor) { + TchOps::sort_with_indices(tensor, dim, descending) + } + + fn float_argsort( + tensor: TchTensor, + dim: usize, + descending: bool, + _out_dtype: IntDType, + ) -> IntTensor { + TchOps::argsort(tensor, dim, descending) + } + + fn float_cast(tensor: TchTensor, dtype: FloatDType) -> TchTensor { + // NOTE: when dtypes of inputs to an arithmetic operation differ, tch handles type + // promotion based on a set of rules: https://pytorch.org/docs/stable/tensor_attributes.html#type-promotion-doc + + // Type promotion is not automatic on all backends so this behavior might differ + let kind = dtype.into_kind(); + + if tensor.tensor.kind() == kind { + tensor + } else { + TchTensor::new(tensor.tensor.to_kind(kind)) + } + } + + fn float_unfold( + tensor: FloatTensor, + dim: usize, + size: usize, + step: usize, + ) -> FloatTensor { + TchOps::unfold(tensor, dim, size, step) + } + + fn float_is_nan(tensor: FloatTensor, _out_dtype: BoolDType) -> BoolTensor { + TchTensor::new(tensor.tensor.isnan()) + } + + fn float_is_inf(tensor: FloatTensor, _out_dtype: BoolDType) -> BoolTensor { + TchTensor::new(tensor.tensor.isinf()) + } +} diff --git a/crates/burn-tch/src/ops/transaction.rs b/crates/burn-tch/src/ops/transaction.rs new file mode 100644 index 0000000..7b5838b --- /dev/null +++ b/crates/burn-tch/src/ops/transaction.rs @@ -0,0 +1,9 @@ +use burn_backend::distributed::DistributedOps; +use burn_backend::ops::TransactionOps; + +use crate::LibTorch; + +impl TransactionOps for LibTorch {} + +// DistributedOps has default implementations; LibTorch does not support collective operations. +impl DistributedOps for LibTorch {} diff --git a/crates/burn-tch/src/tensor.rs b/crates/burn-tch/src/tensor.rs new file mode 100644 index 0000000..4f1f0c5 --- /dev/null +++ b/crates/burn-tch/src/tensor.rs @@ -0,0 +1,511 @@ +use crate::{LibTorchDevice, TchElement}; +use burn_backend::{BoolStore, DType, FloatDType, IntDType, Shape, TensorData, TensorMetadata}; +use libc::c_void; +use std::sync::Arc; + +/// A reference to a tensor storage. +/// +/// We manually implement `Sync` and `Send` unsafely, so even if we could use `Rc`, it isn't safe. +#[allow(clippy::arc_with_non_send_sync)] +pub type StorageRef = Arc<*mut c_void>; + +/// A reference to a tensor storage. +#[derive(PartialEq, Debug, Clone)] +pub enum Storage { + /// When a tensor is a partial view of another tensor. + View { + /// Storage reference for the whole buffer. + buffer_ref: StorageRef, + /// Storage reference for the partial buffer. + view_ref: StorageRef, + }, + /// When a tensor use all of its buffer. + Owned { + /// Storage reference for the whole buffer. + buffer_ref: StorageRef, + }, +} + +impl Storage { + /// Check if the storage can be used inplace. + pub fn can_mut(&self) -> bool { + match self { + Storage::View { + buffer_ref: start_ref, + view_ref, + } => Arc::strong_count(start_ref) == 1 && Arc::strong_count(view_ref) == 1, + Storage::Owned { + buffer_ref: start_ref, + } => Arc::strong_count(start_ref) == 1, + } + } + + /// Get the whole buffer reference. + pub fn buffer_ref(&self) -> &StorageRef { + match self { + Storage::View { + buffer_ref: start_ref, + view_ref: _, + } => start_ref, + Storage::Owned { + buffer_ref: start_ref, + } => start_ref, + } + } +} + +/// A tensor using the tch backend. +#[derive(Debug, PartialEq)] +pub struct TchTensor { + /// Handle to the tensor. Call methods on this field. + pub tensor: tch::Tensor, + + /// The tensor's storage + pub storage: Storage, +} + +impl TensorMetadata for TchTensor { + type Device = LibTorchDevice; + fn dtype(&self) -> DType { + match self.tensor.kind() { + tch::Kind::Uint8 => DType::U8, + tch::Kind::Int8 => DType::I8, + tch::Kind::Int16 => DType::I16, + tch::Kind::Int => DType::I32, + tch::Kind::Int64 => DType::I64, + tch::Kind::Half => DType::F16, + tch::Kind::Float => DType::F32, + tch::Kind::Double => DType::F64, + tch::Kind::Bool => DType::Bool(BoolStore::Native), + tch::Kind::BFloat16 => DType::BF16, + // Complex and quantization types are not valid/implemented. + _ => unimplemented!(), + } + } + + fn shape(&self) -> Shape { + Shape::from(self.tensor.size()) + } + + fn rank(&self) -> usize { + self.tensor.dim() + } + fn device(&self) -> Self::Device { + self.tensor.device().into() + } + + fn can_mut(&self) -> bool { + // The inherent method: unique storage and no broadcast stride. + TchTensor::can_mut(self) + } +} + +impl core::fmt::Display for TchTensor { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.tensor) + } +} + +pub(crate) trait IntoKind { + fn try_into_kind(self) -> Result; + fn into_kind(self) -> tch::Kind + where + Self: Sized, + { + self.try_into_kind().unwrap() + } +} + +impl IntoKind for IntDType { + fn try_into_kind(self) -> Result { + let dtype: DType = self.into(); + dtype.try_into_kind() + } +} + +impl IntoKind for FloatDType { + fn try_into_kind(self) -> Result { + let dtype: DType = self.into(); + dtype.try_into_kind() + } +} + +impl IntoKind for DType { + fn try_into_kind(self) -> Result { + match self { + DType::F64 => Ok(tch::Kind::Double), + DType::F32 => Ok(tch::Kind::Float), + DType::Flex32 => Ok(tch::Kind::Float), + DType::F16 => Ok(tch::Kind::Half), + DType::BF16 => Ok(tch::Kind::BFloat16), + DType::I64 => Ok(tch::Kind::Int64), + // LibTorch backend currently forces I64 int dtype + // DType::I32 => Ok(tch::Kind::Int), + // DType::I16 => Ok(tch::Kind::Int16), + // DType::I8 => Ok(tch::Kind::Int8), + // DType::U8 => Ok(tch::Kind::Uint8), + DType::Bool(BoolStore::Native) => Ok(tch::Kind::Bool), + other => Err(tch::TchError::Kind(format!("Unsupported dtype {other:?}"))), + } + } +} + +impl TchTensor { + /// Create a new tensor. + /// + /// Note that if the tensor was created from an operation that may reuse the same tensor + /// storage as the parent, you should use [from_existing](TchTensor::from_existing) + /// instead. + pub fn new(tensor: tch::Tensor) -> Self { + #[allow(clippy::arc_with_non_send_sync)] + let storage = Storage::Owned { + buffer_ref: Arc::new(tensor.data_ptr()), + }; + + Self { tensor, storage } + } + + /// Create a tensor that was created from an operation executed on a parent tensor. + /// + /// If the child tensor shared the same storage as its parent, it will be cloned, effectively + /// tracking how much tensors point to the same memory space. + pub fn from_existing(tensor: tch::Tensor, storage_parent: Storage) -> Self { + let storage_child = tensor.data_ptr(); + let mut is_a_new_tensor = true; + + match &storage_parent { + Storage::View { + buffer_ref: start_ref, + view_ref, + } => { + if storage_child == *start_ref.as_ref() || storage_child == *view_ref.as_ref() { + is_a_new_tensor = false; + } + } + Storage::Owned { + buffer_ref: start_ref, + } => { + if storage_child == *start_ref.as_ref() { + is_a_new_tensor = false; + } + } + }; + + let storage = match is_a_new_tensor { + true => Storage::Owned { + #[allow(clippy::arc_with_non_send_sync)] + buffer_ref: Arc::new(storage_child), + }, + false => storage_parent.clone(), + }; + + Self { tensor, storage } + } + + /// Create a tensor that uses a part of its parent tensor such as slice and narrow. + pub fn partial(tensor: tch::Tensor, storage_parent: Storage) -> Self { + let storage = Storage::View { + buffer_ref: storage_parent.buffer_ref().clone(), + #[allow(clippy::arc_with_non_send_sync)] + view_ref: Arc::new(tensor.data_ptr()), + }; + Self { tensor, storage } + } +} + +// This is safe since we don't use autodiff from LibTorch. +// Also, atomic reference counting is used to know if the tensor's data can be reused. +// If there are multiple reference on the same tensor, it becomes read only. +unsafe impl Send for TchTensor {} +unsafe impl Sync for TchTensor {} + +impl TchTensor { + /// Checks if the tensor can be mutated in-place. + /// + /// Returns `true` if the tensor's stride does not contain zero (no broadcasting) + /// and the storage can be mutated. + pub fn can_mut(&self) -> bool { + let stride_contains_zero = self.tensor.stride().contains(&0); + + !stride_contains_zero && self.storage.can_mut() + } + + /// Executes an operation on a tensor if the data can be reused. + pub fn mut_ops tch::Tensor>( + &mut self, + func: F, + ) -> Option { + if !self.can_mut() { + return None; + } + + let data = self.storage.clone(); + Some(TchTensor::from_existing(func(&mut self.tensor), data)) + } + + /// Executes a unary operation, reusing the tensor data if possible. + pub fn unary_ops(self, fown: FOwn, fref: FRef) -> TchTensor + where + FOwn: Fn(tch::Tensor) -> tch::Tensor, + FRef: Fn(&tch::Tensor) -> tch::Tensor, + { + if !self.can_mut() { + return TchTensor::from_existing(fref(&self.tensor), self.storage); + } + + TchTensor::from_existing(fown(self.tensor), self.storage) + } + + /// Executes a binary operation, reusing the tensor data if possible. + pub fn binary_ops_tensor( + mut lhs: Self, + mut rhs: Self, + flmut: FLMut, + frmut: FRMut, + fref: FRef, + ) -> TchTensor + where + FLMut: Fn(&mut tch::Tensor, &tch::Tensor) -> tch::Tensor, + FRMut: Fn(&tch::Tensor, &mut tch::Tensor) -> tch::Tensor, + FRef: Fn(&tch::Tensor, &tch::Tensor) -> tch::Tensor, + { + let lhs_shape = lhs.shape(); + let rhs_shape = rhs.shape(); + + // Both lhs and rhs are expected to have the same rank + let d_out = lhs_shape.num_dims(); + let mut out_shape = Shape::from(vec![1usize; d_out]); + + for i in 0..d_out { + out_shape[i] = usize::max(lhs_shape[i], rhs_shape[i]); + } + + let num_elements_out = out_shape.num_elements(); + + // Attempt to mutate lhs tensor + if lhs_shape.num_elements() == num_elements_out + && let Some(output) = lhs.mut_ops(|lhs| flmut(lhs, &rhs.tensor)) + { + return output; + } + + // Attempt to mutate rhs tensor + if rhs_shape.num_elements() == num_elements_out + && let Some(output) = rhs.mut_ops(|rhs| frmut(&lhs.tensor, rhs)) + { + return output; + } + + let storage = lhs.storage; + let tensor = fref(&lhs.tensor, &rhs.tensor); + + TchTensor::from_existing(tensor, storage) + } +} + +impl Clone for TchTensor { + fn clone(&self) -> Self { + Self { + tensor: self.tensor.shallow_clone(), + storage: self.storage.clone(), + } + } +} + +/// A shape that can be used by LibTorch. +#[derive(Debug)] +pub struct TchShape { + /// The shape's dimensions. + pub dims: Vec, +} + +impl From for TchShape { + fn from(shape: Shape) -> Self { + TchShape { + dims: shape.iter().map(|d| *d as i64).collect(), + } + } +} + +impl From<&[usize]> for TchShape { + fn from(shape: &[usize]) -> Self { + TchShape { + dims: shape.iter().map(|d| *d as i64).collect(), + } + } +} + +impl TchTensor { + /// Creates a new tensor from a shape and a device. + /// + /// # Arguments + /// + /// * `data` - The tensor's data. + /// * `device` - The device on which the tensor will be allocated. + /// + /// # Returns + /// + /// A new tensor. + pub fn from_data(data: TensorData, device: tch::Device) -> Self { + let shape_tch = TchShape::from(data.shape.as_slice()); + let tensor = + tch::Tensor::from_data_size(&data.bytes, &shape_tch.dims, E::kind()).to(device); + + Self::new(tensor) + } +} + +impl TchTensor { + /// Creates an empty tensor from a shape and a device. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `device` - The device to create the tensor on. + /// + /// # Returns + /// + /// A new empty tensor. + pub fn empty(shape: Shape, device: LibTorchDevice, dtype: DType) -> Self { + let shape_tch = TchShape::from(shape); + let tensor = tch::Tensor::empty(shape_tch.dims, (dtype.into_kind(), device.into())); + + Self::new(tensor) + } +} + +// Adapted from `tch` to use patched `T::kind()` instead of `T::KIND` which is incorrect for bf16. +// TODO: remove when fixed in `tch` release (https://github.com/LaurentMazare/tch-rs/pull/996). +impl TryFrom<&TchTensor> for Vec { + type Error = tch::TchError; + fn try_from(tensor: &TchTensor) -> Result { + let tensor = &tensor.tensor; + let size = tensor.size(); + if size.len() != 1 { + Err(tch::TchError::Convert(format!( + "Attempting to convert a Tensor with {} dimensions to flat vector", + size.len() + )))?; + } + let numel = size[0] as usize; + let mut vec = vec![T::ZERO; numel]; + // Adapted to use patched `T::kind()` instead + // TODO: tensor.f_to_kind(T::KIND)?.f_copy_data(&mut vec, numel)?; + f_copy_data(&mut tensor.f_to_kind(T::kind())?, &mut vec, numel)?; + Ok(vec) + } +} + +unsafe fn ptr_to_string(ptr: *mut libc::c_char) -> Option { + if !ptr.is_null() { + unsafe { + let str = std::ffi::CStr::from_ptr(ptr).to_string_lossy().into_owned(); + libc::free(ptr as *mut libc::c_void); + Some(str) + } + } else { + None + } +} + +/// Copies `numel` elements from `self` to `dst`. +fn f_copy_data( + tensor: &mut tch::Tensor, + dst: &mut [T], + numel: usize, +) -> Result<(), tch::TchError> { + if T::kind() != tensor.f_kind()? { + return Err(tch::TchError::Kind(format!( + "incoherent elt kind, {:?} != {:?}", + tensor.f_kind(), + T::kind() + ))); + } + if dst.len() < numel { + return Err(tch::TchError::Shape(format!("slice len < {numel}"))); + } + + unsafe { + torch_sys::at_copy_data( + tensor.as_mut_ptr(), + dst.as_mut_ptr() as *const c_void, + numel, + T::kind().elt_size_in_bytes(), + ); + match ptr_to_string(torch_sys::get_and_reset_last_err()) { + None => Ok(()), + Some(c_error) => Err(tch::TchError::Torch(c_error)), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn_backend::ops::FloatTensorOps; + use burn_backend::{Backend, quantization::QuantScheme, read_sync}; + + type B = crate::LibTorch; + + #[test] + fn should_have_bf16_kind() { + let data = TensorData::from([4.0, 4.0]); + let tensor_1: TchTensor = B::float_from_data(data, &Default::default()); + let tensor_2 = B::float_cast(tensor_1, DType::BF16.into()); + + assert_eq!(tensor_2.tensor.kind(), tch::Kind::BFloat16); + + let out = read_sync(B::float_into_data(tensor_2)).unwrap(); + + out.assert_eq(&TensorData::from([4.0, 4.0]), false); + } + + #[test] + fn should_support_dtypes() { + let device = Default::default(); + + assert!(B::supports_dtype(&device, DType::F64)); + assert!(B::supports_dtype(&device, DType::F32)); + assert!(B::supports_dtype(&device, DType::Flex32)); + assert!(B::supports_dtype(&device, DType::F16)); + assert!(B::supports_dtype(&device, DType::BF16)); + assert!(B::supports_dtype(&device, DType::I64)); + assert!(B::supports_dtype(&device, DType::I32)); + assert!(B::supports_dtype(&device, DType::I16)); + assert!(B::supports_dtype(&device, DType::I8)); + assert!(B::supports_dtype(&device, DType::U8)); + assert!(B::supports_dtype(&device, DType::Bool(BoolStore::Native))); + + assert!(!B::supports_dtype(&device, DType::U64)); + assert!(!B::supports_dtype(&device, DType::U32)); + assert!(!B::supports_dtype(&device, DType::U16)); + assert!(!B::supports_dtype( + &device, + DType::QFloat(QuantScheme::default()) + )); + } + + #[test] + fn should_support_from_bf16() { + let data = TensorData::from([[1.0], [1.]]).convert_dtype(DType::BF16); + let tensor_1: TchTensor = B::float_from_data(data, &Default::default()); + let data = TensorData::from([[2.0], [2.]]).convert_dtype(DType::BF16); + let tensor_2 = B::float_from_data(data, &Default::default()); + + let tensor_3 = B::float_add(tensor_1, tensor_2); + + assert_eq!(tensor_3.tensor.kind(), tch::Kind::BFloat16); + + let out = read_sync(B::float_into_data(tensor_3)).unwrap(); + + out.assert_eq(&TensorData::from([[3.0], [3.0]]), false); + } +} + +unsafe extern "C" { + /// Dummy function to get CUDA to link properly + pub fn dummy_cuda_dependency(); +} + +#[used] +static INIT_ARRAY: [unsafe extern "C" fn(); 1] = [dummy_cuda_dependency]; diff --git a/crates/burn-tensor/Cargo.toml b/crates/burn-tensor/Cargo.toml new file mode 100644 index 0000000..ed0416b --- /dev/null +++ b/crates/burn-tensor/Cargo.toml @@ -0,0 +1,80 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science", "no-std", "embedded", "wasm"] +description = "Tensor library with user-friendly APIs and automatic differentiation support" +documentation = "https://docs.rs/burn-tensor" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "tensor", "pytorch", "ndarray"] +license.workspace = true +name = "burn-tensor" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-tensor" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["std", "burn-dispatch/default"] +doc = ["default"] +std = [ + "num-traits/std", + "burn-std/std", + "burn-backend/std", + "burn-dispatch/std", + "colored", +] +tracing = ["burn-std/tracing", "burn-backend/tracing", "burn-dispatch/tracing"] + +cubecl = ["burn-backend/cubecl"] +cubecl-cuda = ["burn-backend/cubecl-cuda"] +cubecl-hip = ["burn-backend/cubecl-hip"] +cubecl-wgpu = ["burn-backend/cubecl-wgpu"] +cubecl-metal = ["burn-backend/cubecl-metal"] +cubecl-vulkan = ["burn-backend/cubecl-vulkan"] +cubecl-webgpu = ["burn-backend/cubecl-webgpu"] +cubecl-cpu = ["burn-backend/cubecl-cpu"] + +# Backends +cuda = ["burn-dispatch/cuda", "cubecl-cuda"] +rocm = ["burn-dispatch/rocm", "cubecl-hip"] +flex = ["burn-dispatch/flex"] +ndarray = ["burn-dispatch/ndarray"] +tch = ["burn-dispatch/tch"] +wgpu = ["burn-dispatch/wgpu", "cubecl-wgpu"] +vulkan = ["wgpu", "burn-dispatch/vulkan", "cubecl-vulkan"] +webgpu = ["wgpu", "burn-dispatch/webgpu", "cubecl-webgpu"] +metal = ["wgpu", "burn-dispatch/metal", "cubecl-metal"] +cpu = ["burn-dispatch/cpu", "cubecl-cpu"] +autodiff = ["burn-dispatch/autodiff"] +# Remote compute client over Iroh. +remote = ["std", "burn-dispatch/remote"] +# Host a remote compute server over Iroh. Wasm-compatible. +remote-server = ["remote", "burn-dispatch/remote-server"] +# Add the legacy WebSocket transport. +remote-websocket = ["remote", "burn-dispatch/remote-websocket"] + +# For backend extensions +extension = [] + +# Backend features +autotune = ["burn-dispatch/autotune"] +autotune-checks = ["burn-dispatch/autotune-checks"] +fusion = ["burn-dispatch/fusion"] + +[dependencies] +burn-std = { workspace = true } +burn-backend = { workspace = true } +burn-dispatch = { workspace = true } +burn-backend-extension = { workspace = true, optional = true } + +colored = { workspace = true, optional = true } +derive-new = { workspace = true } +num-traits = { workspace = true } + +# Serialization +serde = { workspace = true } + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs", "--html-in-header", "katex-header.html"] diff --git a/crates/burn-tensor/LICENSE-APACHE b/crates/burn-tensor/LICENSE-APACHE new file mode 120000 index 0000000..1cd601d --- /dev/null +++ b/crates/burn-tensor/LICENSE-APACHE @@ -0,0 +1 @@ +../../LICENSE-APACHE \ No newline at end of file diff --git a/crates/burn-tensor/LICENSE-MIT b/crates/burn-tensor/LICENSE-MIT new file mode 120000 index 0000000..b2cfbdc --- /dev/null +++ b/crates/burn-tensor/LICENSE-MIT @@ -0,0 +1 @@ +../../LICENSE-MIT \ No newline at end of file diff --git a/crates/burn-tensor/README.md b/crates/burn-tensor/README.md new file mode 100644 index 0000000..b97bffd --- /dev/null +++ b/crates/burn-tensor/README.md @@ -0,0 +1,12 @@ +# Burn Tensor + +> [Burn](https://github.com/tracel-ai/burn) Tensor Library + +[![Current Crates.io Version](https://img.shields.io/crates/v/burn-tensor.svg)](https://crates.io/crates/burn-tensor) +[![license](https://shields.io/badge/license-MIT%2FApache--2.0-blue)](https://github.com/tracel-ai/burn-tensor/blob/master/README.md) + +This library provides the core abstractions required to run tensor operations with Burn. + +`Tensor`s are generic over the backend to allow users to perform operations using different +`Backend` implementations. Burn's tensors also support auto-differentiation thanks to the +`AutodiffBackend` trait. diff --git a/crates/burn-tensor/katex-header.html b/crates/burn-tensor/katex-header.html new file mode 120000 index 0000000..4a9eb97 --- /dev/null +++ b/crates/burn-tensor/katex-header.html @@ -0,0 +1 @@ +../../docs/katex-header.html \ No newline at end of file diff --git a/crates/burn-tensor/src/bridge/kind.rs b/crates/burn-tensor/src/bridge/kind.rs new file mode 100644 index 0000000..dda581d --- /dev/null +++ b/crates/burn-tensor/src/bridge/kind.rs @@ -0,0 +1,365 @@ +use alloc::vec::Vec; +use burn_backend::{TensorMetadata, TensorPrimitive, get_device_settings}; +use burn_dispatch::{Dispatch, DispatchTensor}; +use burn_std::DeviceSettings; + +/// A type-level representation of the kind of a float tensor +#[derive(Clone, Debug)] +pub struct Float; + +/// A type-level representation of the kind of a int tensor. +#[derive(Clone, Debug)] +pub struct Int; + +/// A type-level representation of the kind of a bool tensor. +#[derive(Clone, Debug)] +pub struct Bool; + +mod sealed { + pub trait Sealed {} +} + +impl sealed::Sealed for Float {} +impl sealed::Sealed for Int {} +impl sealed::Sealed for Bool {} + +/// A type-level representation of the kind of a tensor. +/// Metadata access is lazy. +/// +/// # Notes +/// This trait is intentionally sealed to keep the set of tensor kinds closed. +/// +/// Although exposed publicly, tensor kinds are not meant to be extensible: +/// the backend dispatch system, `DType`, and all tensor ops assume a fixed, +/// closed set of tensor kinds (e.g. Float, Int, Bool), each mapping directly to a +/// corresponding backend implementation. +pub trait TensorKind: sealed::Sealed + Clone + Send + Sync + core::fmt::Debug { + /// The tensor kind identifier. + const KIND: Kind; + + /// The name of the tensor kind. + fn name() -> &'static str { + Self::KIND.as_str() + } +} + +impl TensorKind for Float { + const KIND: Kind = Kind::Float; +} + +impl TensorKind for Int { + const KIND: Kind = Kind::Int; +} + +impl TensorKind for Bool { + const KIND: Kind = Kind::Bool; +} + +/// Represents the kind of a [`Tensor`](crate::Tensor). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Kind { + /// A float tensor kind. + Float, + /// An integer tensor kind. + Int, + /// A boolean tensor kind. + Bool, +} + +impl Kind { + /// Get the string representation of the [`Kind`]. + pub fn as_str(&self) -> &'static str { + match self { + Kind::Float => "Float", + Kind::Int => "Int", + Kind::Bool => "Bool", + } + } +} + +/// A type-tagged tensor at the bridge layer between the high-level tensor API +/// and the dispatch system. +/// +/// `BridgeTensor` serves as the runtime representation for the public tensor +/// kinds (Float, Int, Bool) and internal variants like quantized floats, wrapping +/// the uniform [`DispatchTensor`] used by the underlying dispatch layer. This +/// separation keeps tensor kind tracking out of the backends while avoiding +/// exposure of backend-level primitives in the public API. +pub struct BridgeTensor { + blob: bridge_opaque::Opaque, +} + +// Aligned, type-erased storage for `BridgeTensorVariant`. See `crate::macros` +// for why this indirection exists (it keeps the dispatch type tree out of +// downstream MIR). +burn_std::obfuscate!( + type: BridgeTensorVariant, + module: bridge_opaque, + derives: [Send, Sync] +); + +impl core::fmt::Debug for BridgeTensor { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("BridgeTensor") + .field("kind", &self.kind()) + .finish() + } +} + +impl Clone for BridgeTensor { + fn clone(&self) -> Self { + Self::new(self.as_variant().clone()) + } +} + +impl BridgeTensor { + fn as_variant(&self) -> &BridgeTensorVariant { + self.blob.as_ref() + } + + fn into_variant(self) -> BridgeTensorVariant { + self.blob.into_inner() + } + + fn new(inner: BridgeTensorVariant) -> Self { + Self { + blob: bridge_opaque::Opaque::new(inner), + } + } +} + +#[derive(Clone, Debug)] +/// Private type obfucated by Blob. +enum BridgeTensorVariant { + /// A boolean tensor. + Bool(DispatchTensor), + /// An integer tensor. + Int(DispatchTensor), + /// A floating-point tensor. + Float(DispatchTensor), + /// A quantized floating-point tensor. + QFloat(DispatchTensor), +} + +/// Runtime tag identifying which variant a [`BridgeTensor`] wraps. +/// +/// Exposed so callers can dispatch on the variant without having to reach the +/// private [`BridgeTensorVariant`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum BridgeKind { + /// A boolean tensor. + Bool, + /// An integer tensor. + Int, + /// A floating-point tensor. + Float, + /// A quantized floating-point tensor. + QFloat, +} + +/// Switches visibility based on the `extension` feature: `pub` when enabled +/// (so backend-extension authors can call these), `pub(crate)` otherwise (so +/// burn-tensor's own ops keep working without leaking the dispatch types). +macro_rules! ext_fn { + ($(#[$meta:meta])* fn $($tt:tt)+) => { + #[cfg(feature = "extension")] + $(#[$meta])* + pub fn $($tt)+ + #[cfg(not(feature = "extension"))] + $(#[$meta])* + pub(crate) fn $($tt)+ + }; +} + +impl BridgeTensor { + ext_fn! { + /// Builds a bridge tensor that wraps a floating-point dispatch tensor. + /// + /// Available with the `extension` feature for backend-extension authors. + fn float(tensor: DispatchTensor) -> Self { + Self::new(BridgeTensorVariant::Float(tensor)) + } + } + + ext_fn! { + /// Builds a bridge tensor that wraps an integer dispatch tensor. + /// + /// Available with the `extension` feature for backend-extension authors. + fn int(tensor: DispatchTensor) -> Self { + Self::new(BridgeTensorVariant::Int(tensor)) + } + } + + ext_fn! { + /// Builds a bridge tensor that wraps a boolean dispatch tensor. + /// + /// Available with the `extension` feature for backend-extension authors. + fn bool(tensor: DispatchTensor) -> Self { + Self::new(BridgeTensorVariant::Bool(tensor)) + } + } + + ext_fn! { + /// Builds a bridge tensor that wraps a quantized floating-point dispatch tensor. + /// + /// Available with the `extension` feature for backend-extension authors. + fn qfloat(tensor: DispatchTensor) -> Self { + Self::new(BridgeTensorVariant::QFloat(tensor)) + } + } + + /// Returns the runtime tag identifying which variant this tensor wraps. + pub fn kind(&self) -> BridgeKind { + match self.as_variant() { + BridgeTensorVariant::Bool(_) => BridgeKind::Bool, + BridgeTensorVariant::Int(_) => BridgeKind::Int, + BridgeTensorVariant::Float(_) => BridgeKind::Float, + BridgeTensorVariant::QFloat(_) => BridgeKind::QFloat, + } + } + + /// Returns `true` if this tensor is the float variant. + pub fn is_float(&self) -> bool { + matches!(self.kind(), BridgeKind::Float) + } + + /// Returns `true` if this tensor is the int variant. + pub fn is_int(&self) -> bool { + matches!(self.kind(), BridgeKind::Int) + } + + /// Returns `true` if this tensor is the bool variant. + pub fn is_bool(&self) -> bool { + matches!(self.kind(), BridgeKind::Bool) + } + + /// Returns `true` if this tensor is the quantized float variant. + pub fn is_qfloat(&self) -> bool { + matches!(self.kind(), BridgeKind::QFloat) + } + + ext_fn! { + /// Consumes the bridge tensor and returns its variant tag together with + /// the underlying dispatch tensor. + /// + /// Available with the `extension` feature for backend-extension authors. + fn into_parts(self) -> (BridgeKind, DispatchTensor) { + match self.into_variant() { + BridgeTensorVariant::Bool(t) => (BridgeKind::Bool, t), + BridgeTensorVariant::Int(t) => (BridgeKind::Int, t), + BridgeTensorVariant::Float(t) => (BridgeKind::Float, t), + BridgeTensorVariant::QFloat(t) => (BridgeKind::QFloat, t), + } + } + } + + ext_fn! { + /// Borrows the bridge tensor as its variant tag together with a reference + /// to the underlying dispatch tensor. + /// + /// Available with the `extension` feature for backend-extension authors. + fn as_parts(&self) -> (BridgeKind, &DispatchTensor) { + match self.as_variant() { + BridgeTensorVariant::Bool(t) => (BridgeKind::Bool, t), + BridgeTensorVariant::Int(t) => (BridgeKind::Int, t), + BridgeTensorVariant::Float(t) => (BridgeKind::Float, t), + BridgeTensorVariant::QFloat(t) => (BridgeKind::QFloat, t), + } + } + } + + /// Returns the dtype of the tensor. + pub fn dtype(&self) -> burn_std::DType { + match self.as_variant() { + BridgeTensorVariant::Bool(tensor) => tensor.dtype(), + BridgeTensorVariant::Int(tensor) => tensor.dtype(), + BridgeTensorVariant::Float(tensor) => tensor.dtype(), + BridgeTensorVariant::QFloat(tensor) => tensor.dtype(), + } + } + + /// Whether the tensor's buffer can be mutated in place (see + /// [`TensorMetadata::can_mut`]). + pub fn can_mut(&self) -> bool { + match self.as_variant() { + BridgeTensorVariant::Bool(tensor) => tensor.can_mut(), + BridgeTensorVariant::Int(tensor) => tensor.can_mut(), + BridgeTensorVariant::Float(tensor) => tensor.can_mut(), + BridgeTensorVariant::QFloat(tensor) => tensor.can_mut(), + } + } + + /// Returns the shape of the tensor. + pub fn shape(&self) -> burn_std::Shape { + match self.as_variant() { + BridgeTensorVariant::Bool(tensor) => tensor.shape(), + BridgeTensorVariant::Int(tensor) => tensor.shape(), + BridgeTensorVariant::Float(tensor) => tensor.shape(), + BridgeTensorVariant::QFloat(tensor) => tensor.shape(), + } + } + + /// Returns the number of dimensions of the tensor. + pub fn rank(&self) -> usize { + match self.as_variant() { + BridgeTensorVariant::Bool(tensor) => tensor.rank(), + BridgeTensorVariant::Int(tensor) => tensor.rank(), + BridgeTensorVariant::Float(tensor) => tensor.rank(), + BridgeTensorVariant::QFloat(tensor) => tensor.rank(), + } + } + + pub(crate) fn as_dispatch(&self) -> &DispatchTensor { + match self.as_variant() { + BridgeTensorVariant::Bool(tensor) => tensor, + BridgeTensorVariant::Int(tensor) => tensor, + BridgeTensorVariant::Float(tensor) => tensor, + BridgeTensorVariant::QFloat(tensor) => tensor, + } + } + + #[cfg(feature = "autodiff")] + pub(crate) fn as_float(&self) -> &DispatchTensor { + match self.as_variant() { + BridgeTensorVariant::Float(tensor) => tensor, + _ => panic!("Should be Float primitive kind"), + } + } + + pub(crate) fn into_dispatch_vec(tensors: Vec) -> Vec { + tensors.into_iter().map(Into::into).collect() + } + + pub(crate) fn into_float(self) -> DispatchTensor { + match self.into_variant() { + BridgeTensorVariant::Float(tensor) => tensor, + // Returns the dequantized float tensor. + BridgeTensorVariant::QFloat(tensor) => { + TensorPrimitive::::QFloat(tensor).tensor() + } + _ => panic!("Should be Float primitive kind"), + } + } + + pub(crate) fn device_settings(&self) -> DeviceSettings { + let device = match self.as_variant() { + BridgeTensorVariant::Bool(tensor) => tensor.device(), + BridgeTensorVariant::Int(tensor) => tensor.device(), + BridgeTensorVariant::Float(tensor) => tensor.device(), + BridgeTensorVariant::QFloat(tensor) => tensor.device(), + }; + + get_device_settings::(&device) + } +} + +impl From for DispatchTensor { + fn from(value: BridgeTensor) -> Self { + match value.into_variant() { + BridgeTensorVariant::Bool(tensor) => tensor, + BridgeTensorVariant::Int(tensor) => tensor, + BridgeTensorVariant::Float(tensor) => tensor, + BridgeTensorVariant::QFloat(tensor) => tensor, + } + } +} diff --git a/crates/burn-tensor/src/bridge/mod.rs b/crates/burn-tensor/src/bridge/mod.rs new file mode 100644 index 0000000..1245fa5 --- /dev/null +++ b/crates/burn-tensor/src/bridge/mod.rs @@ -0,0 +1,5 @@ +mod kind; +mod ops; + +pub use kind::*; +pub(crate) use ops::*; diff --git a/crates/burn-tensor/src/bridge/ops/autodiff.rs b/crates/burn-tensor/src/bridge/ops/autodiff.rs new file mode 100644 index 0000000..75ac383 --- /dev/null +++ b/crates/burn-tensor/src/bridge/ops/autodiff.rs @@ -0,0 +1,32 @@ +use crate::{bridge::BasicOps, ops::BridgeTensor}; + +/// Trait that list all operations that can be applied on all tensors on an autodiff backend. +/// +/// # Warnings +/// +/// This is an internal trait, use the public API provided by the [`Tensor`](crate::Tensor) struct. +pub(crate) trait BasicAutodiffOps: BasicOps { + /// Returns the inner tensor without the autodiff information. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// Users should prefer the [`Tensor::inner`](crate::Tensor::inner) + /// function, which is more high-level and designed for public use. + fn inner(tensor: BridgeTensor) -> BridgeTensor; + + /// Convert a tensor to the autodiff backend. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// Users should prefer the [`Tensor::from_inner`](crate::Tensor::from_inner) + /// function, which is more high-level and designed for public use. + fn from_inner(inner: BridgeTensor) -> BridgeTensor; +} diff --git a/crates/burn-tensor/src/bridge/ops/base.rs b/crates/burn-tensor/src/bridge/ops/base.rs new file mode 100644 index 0000000..2eee9cc --- /dev/null +++ b/crates/burn-tensor/src/bridge/ops/base.rs @@ -0,0 +1,727 @@ +use alloc::vec::Vec; +use burn_backend::{Scalar, TensorData, ops::TransactionPrimitive}; +use burn_dispatch::Dispatch; +use burn_std::{DType, ExecutionError, IndexingUpdateOp, Shape, Slice}; + +use crate::{ + Device, + ops::{BridgeTensor, TensorKind}, +}; + +/// Trait for the one basic op that still requires Backend +/// +/// # Warnings +/// +/// This is an internal trait, use the public API provided by the [`Tensor`](crate::Tensor) struct. +pub(crate) trait TransactionOp: BasicOps { + /// Read the data from the tensor using a transaction. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + fn register_transaction(tr: &mut TransactionPrimitive, tensor: BridgeTensor); +} + +/// Trait that list all operations that can be applied on all tensors. +/// +/// # Warnings +/// +/// This is an internal trait, use the public API provided by the [`Tensor`](crate::Tensor) struct. +pub(crate) trait BasicOps: TensorKind { + /// Creates an empty tensor with the given shape. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `device` - The device on which the tensor will be allocated. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The empty tensor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For creating empty tensors, users should prefer the [`Tensor::empty`](crate::Tensor::empty) + /// function, which is more high-level and designed for public use. + fn empty(shape: Shape, device: &Device, dtype: DType) -> BridgeTensor; + + /// Creates a tensor filled with zeros. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `device` - The device on which the tensor will be allocated. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The tensor filled with zeros. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For creating a tensor filled with zeros, users should prefer the [`Tensor::zeros`](crate::Tensor::zeros) + /// function, which is more high-level and designed for public use. + fn zeros(shape: Shape, device: &Device, dtype: DType) -> BridgeTensor; + + /// Creates a tensor filled with ones. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `device` - The device on which the tensor will be allocated. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The tensor filled with ones. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For creating a tensor filled with ones, users should prefer the [`Tensor::ones`](crate::Tensor::ones) + /// function, which is more high-level and designed for public use. + fn ones(shape: Shape, device: &Device, dtype: DType) -> BridgeTensor; + + /// Creates a tensor of the given shape where each element is equal to the provided value. + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `fill_value` - The value with which to fill the tensor. + /// * `device` - The device on which the tensor will be allocated. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// The tensor filled with the specified value. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For creating full tensors, users should prefer the [`Tensor::full`](crate::Tensor::full) + /// function, which is more high-level and designed for public use. + fn full(shape: Shape, fill_value: Scalar, device: &Device, dtype: DType) -> BridgeTensor; + + /// Reshapes the tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `shape` - The new shape of the tensor. + /// + /// # Returns + /// + /// The reshaped tensor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For reshaping a tensor, users should prefer the [`Tensor::reshape`](crate::Tensor::reshape) + /// function, which is more high-level and designed for public use. + fn reshape(tensor: BridgeTensor, shape: Shape) -> BridgeTensor; + + /// Transposes a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to transpose. + /// + /// # Returns + /// + /// The transposed tensor. + fn transpose(tensor: BridgeTensor) -> BridgeTensor; + + /// Swaps two dimensions of a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to swap the dimensions of. + /// * `dim1` - The first dimension to swap. + /// * `dim2` - The second dimension to swap. + /// + /// # Returns + /// + /// The tensor with the dimensions swapped. + fn swap_dims(tensor: BridgeTensor, dim1: usize, dim2: usize) -> BridgeTensor; + + /// Permutes the dimensions of a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to permute the dimensions of. + /// * `axes` - The new order of the dimensions. + /// + /// # Returns + /// + /// The tensor with the dimensions permuted. + fn permute(tensor: BridgeTensor, axes: &[usize]) -> BridgeTensor; + + /// Flips the tensor along the given axes. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to flip. + /// * `axes` - The axes to flip the tensor along. + /// + /// # Returns + /// + /// The tensor with the axes flipped. + fn flip(tensor: BridgeTensor, axes: &[usize]) -> BridgeTensor; + + /// Select tensor elements corresponding to the given slices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `slices` - The slices specifying ranges and steps for each dimension. + /// + /// # Returns + /// + /// The selected elements. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For selecting elements of a tensor, users should prefer the [`Tensor::slice`](crate::Tensor::slice) + /// function, which is more high-level and designed for public use. + fn slice(tensor: BridgeTensor, slices: &[Slice]) -> BridgeTensor; + + /// Assigns the given value to the tensor elements corresponding to the given slices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `slices` - The slices specifying which elements to assign, including support for steps. + /// * `value` - The value to assign. + /// + /// # Returns + /// + /// The tensor with the assigned values. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For assigning values to elements of a tensor, users should prefer the [`Tensor::slice_assign`](crate::Tensor::slice_assign) + /// function, which is more high-level and designed for public use. + fn slice_assign(tensor: BridgeTensor, slices: &[Slice], value: BridgeTensor) -> BridgeTensor; + + /// Select tensor elements along the given dimension corresponding to the given indices. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to select from. + /// * `dim` - The dimension along which to select. + /// * `indices` - The indices of the elements to select. + /// + /// # Returns + /// + /// The selected tensor elements. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For selecting elements from a tensor along an axis, users should prefer the [`Tensor::select`](crate::Tensor::select) + /// function, which is more high-level and designed for public use. + fn select(tensor: BridgeTensor, dim: usize, indices: BridgeTensor) -> BridgeTensor; + + /// Assign the selected elements along the given dimension corresponding to the given indices + /// from the value tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to assign elements to. + /// * `dim` - The axis along which to assign elements. + /// * `indices` - The indices of the elements to assign. + /// * `values` - The values to assign to the tensor. + /// * `update` - The operation used to update the existing values at the indexed positions (e.g., add). + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor, where each element is taken from the + /// corresponding element of the input tensor at the corresponding index along the specified axis, + /// except for the elements at the specified indices, which are taken from the corresponding + /// element of the values tensor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For assigning elements to a tensor along an axis, users should prefer the [`Tensor::select_assign`](crate::Tensor::select_assign) + /// function, which is more high-level and designed for public use. + fn select_assign( + tensor: BridgeTensor, + dim: usize, + indices: BridgeTensor, + values: BridgeTensor, + update: IndexingUpdateOp, + ) -> BridgeTensor; + + /// Selects elements from a tensor based on a boolean mask. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to select elements from if the corresponding element of the mask is true. + /// * `mask` - The boolean mask to use for selecting elements. + /// * `source` - The tensor to select elements from when the corresponding element of the mask is false. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensors, where each element is taken from the + /// corresponding element of the left hand side tensor if the corresponding element of the mask + /// is true, and from the corresponding element of the right hand side tensor otherwise. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For selecting elements from a tensor based on a boolean mask, users should prefer the [`Tensor::mask_where`](crate::Tensor::mask_where) + /// function, which is more high-level and designed for public use. + fn mask_where(tensor: BridgeTensor, mask: BridgeTensor, source: BridgeTensor) -> BridgeTensor; + + /// Fills elements of a tensor based on a boolean mask. + /// + /// # Arguments + /// + /// * `tensor` - The tensor where will be overwritten with the value + /// when the corresponding element of the mask is true. + /// * `mask` - The boolean mask to use for filling elements. + /// * `value` - The value to fill elements with when the corresponding element of the mask is true. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensors, where each element is taken from the + /// corresponding element unmodified if the corresponding element of the mask is false, and + /// filled with the value otherwise. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For filling elements of a tensor based on a boolean mask, users should prefer the [`Tensor::mask_fill`](crate::Tensor::mask_fill) + /// function, which is more high-level and designed for public use. + fn mask_fill(tensor: BridgeTensor, mask: BridgeTensor, value: Scalar) -> BridgeTensor; + + /// Gathers elements from a tensor along an axis. + /// + /// # Arguments + /// + /// * `dim` - The axis along which to gather elements. + /// * `tensor` - The tensor to gather elements from. + /// * `indices` - The indices of the elements to gather. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor, where each element is taken from the + /// corresponding element of the input tensor at the corresponding index along the specified axis. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For gathering elements from a tensor along an axis, users should prefer the [`Tensor::gather`](crate::Tensor::gather) + /// function, which is more high-level and designed for public use. + fn gather(dim: usize, tensor: BridgeTensor, indices: BridgeTensor) -> BridgeTensor; + + /// Scatters elements into a tensor along an axis. + /// + /// # Arguments + /// + /// * `dim` - The axis along which to scatter elements. + /// * `tensor` - The tensor to scatter elements into. + /// * `indices` - The indices of the elements to scatter. + /// * `values` - The values to scatter into the tensor. + /// * `update` - The operation used to update the existing values at the indexed positions (e.g., add). + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor, where each element is taken from the + /// corresponding element of the input tensor at the corresponding index along the specified axis, + /// except for the elements at the specified indices, which are taken from the corresponding + /// element of the values tensor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For scattering elements into a tensor along an axis, users should prefer the [`Tensor::scatter`](crate::Tensor::scatter) + /// function, which is more high-level and designed for public use. + fn scatter( + dim: usize, + tensor: BridgeTensor, + indices: BridgeTensor, + values: BridgeTensor, + update: IndexingUpdateOp, + ) -> BridgeTensor; + + /// Multi-dimensional scatter: update `data` at multi-index locations specified by `indices`. + fn scatter_nd( + data: BridgeTensor, + indices: BridgeTensor, + values: BridgeTensor, + reduction: IndexingUpdateOp, + ) -> BridgeTensor; + + /// Multi-dimensional gather: collect slices from `data` at multi-index locations. + fn gather_nd(data: BridgeTensor, indices: BridgeTensor) -> BridgeTensor; + + /// Returns the device on which the tensor is allocated. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// + /// # Returns + /// + /// The device on which the tensor is allocated. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For getting the device of a tensor, users should prefer the [`Tensor::device`](crate::Tensor::device) + /// function, which is more high-level and designed for public use. + fn device(tensor: &BridgeTensor) -> Device; + + /// Moves the tensor to the given device. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `device` - The device on which the tensor will be moved. + /// + /// # Returns + /// + /// The tensor on the given device. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For moving a tensor to a device, users should prefer the [`Tensor::to_device`](crate::Tensor::to_device) + /// function, which is more high-level and designed for public use. + #[allow(clippy::wrong_self_convention)] + fn to_device(tensor: BridgeTensor, device: &Device) -> BridgeTensor; + + /// Extracts the data from the tensor asynchronously. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// + /// # Returns + /// + /// The data of the tensor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For extracting the data of a tensor, users should prefer the [`Tensor::into_data`](crate::Tensor::into_data) + /// function, which is more high-level and designed for public use. + #[allow(clippy::wrong_self_convention)] + fn into_data_async( + tensor: BridgeTensor, + ) -> impl Future> + Send; + + /// Creates a tensor from the given data enforcing the provided data type. + /// + /// # Arguments + /// + /// * `data` - The data of the tensor. + /// * `device` - The device on which the tensor will be allocated. + /// * `dtype` - The target data type. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For creating a tensor from data, users should prefer the [`Tensor::from_data`](crate::Tensor::from_data) + /// function, which is more high-level and designed for public use. + fn from_data(data: TensorData, device: &Device, dtype: DType) -> BridgeTensor; + + /// Repeat the tensor along the given dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// * `dim` - The dimension along which the tensor will be repeated. + /// * `times` - The number of times the tensor will be repeated. + /// + /// # Returns + /// + /// The repeated tensor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For repeating a tensor, users should prefer the [`Tensor::repeat_dim`](crate::Tensor::repeat_dim) + /// function, which is more high-level and designed for public use. + fn repeat_dim(tensor: BridgeTensor, dim: usize, times: usize) -> BridgeTensor; + + /// Concatenates the given tensors along the given dimension. + /// + /// # Arguments + /// + /// * `vectors` - The tensors to concatenate. + /// * `dim` - The dimension along which the tensors will be concatenated. + /// + /// # Returns + /// + /// The concatenated tensor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For concatenating tensors, users should prefer the [`Tensor::cat`](crate::Tensor::cat) + /// function, which is more high-level and designed for public use. + fn cat(vectors: Vec, dim: usize) -> BridgeTensor; + + /// Equates the given tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// The tensor of booleans indicating whether the corresponding elements are equal. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For equating tensors, users should prefer the [`Tensor::equal`](crate::Tensor::equal) + /// function, which is more high-level and designed for public use. + fn equal(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor; + + /// Element-wise equality between two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side scalar. + /// + /// # Returns + /// + /// A boolean tensor with the same shape as the input tensors, where each element is true if the + /// corresponding elements of the input tensors are equal, and false otherwise. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For element-wise equality between two tensors, users should prefer the [`Tensor::equal_elem`](crate::Tensor::equal_elem) + /// function, which is more high-level and designed for public use. + fn equal_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor; + + /// Applies element-wise non-equality comparison between the given tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// The tensor of booleans indicating whether the corresponding elements are equal. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For non-equality comparison of tensors, users should prefer the [`Tensor::not_equal`](crate::Tensor::not_equal) + /// function, which is more high-level and designed for public use. + fn not_equal(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor; + + /// Element-wise non-equality between two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side scalar. + /// + /// # Returns + /// + /// A boolean tensor with the same shape as the input tensors, where each element is true if the + /// corresponding elements of the input tensors are equal, and false otherwise. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For element-wise non-equality between two tensors, users should prefer the [`Tensor::not_equal_elem`](crate::Tensor::not_equal_elem) + /// function, which is more high-level and designed for public use. + fn not_equal_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor; + + /// Tests if any element in the `tensor` evaluates to True. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// + /// # Returns + /// + /// A boolean tensor with a single element, True if any element in the input tensor evaluates to True, False otherwise. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. Users should prefer the [`Tensor::any`](crate::Tensor::any) + /// function, which is more high-level and designed for public use. + fn any(tensor: BridgeTensor) -> BridgeTensor; + + /// Tests if any element in the tensor evaluates to True along a given dimension dim. + /// + /// # Arguments + /// + /// * tensor - The tensor to test. + /// * dim - The axis along which to test. + /// + /// # Returns + /// + /// A boolean tensor with the same size as input tensor, except in the dim axis where the size is 1. + /// Returns True if any element in the input tensor along the given dimension evaluates to True, False otherwise. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. Users should prefer the [`Tensor::any_dim`](crate::Tensor::any_dim) + /// function, which is more high-level and designed for public use. + fn any_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor; + + /// Tests if all elements in the `tensor` evaluate to True. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// + /// # Returns + /// + /// A boolean tensor with a single element, True if all elements in the input tensor evaluates to True, False otherwise. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. Users should prefer the [`Tensor::all`](crate::Tensor::all) + /// function, which is more high-level and designed for public use. + fn all(tensor: BridgeTensor) -> BridgeTensor; + + /// Tests if all elements in the `tensor` evaluate to True along a given dimension `dim`. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. + /// + /// # Returns + /// + /// A boolean tensor with the same size as input `tensor`, except in the `dim` axis where the size is 1. + /// Returns True if all elements in the input tensor along the given dimension evaluate to True, False otherwise. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. Users should prefer the [`Tensor::all_dim`](crate::Tensor::all_dim) + /// function, which is more high-level and designed for public use. + fn all_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor; + + /// Broadcasts the given tensor to the specified shape. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to broadcast. + /// * `shape` - The shape to broadcast to. + /// + /// # Returns + /// + /// The broadcasted tensor. + fn expand(tensor: BridgeTensor, shape: Shape) -> BridgeTensor; + + /// Unfold windows along a dimension. + /// + /// Returns a view of the tensor with all complete windows of size `size` in dimension `dim`; + /// where windows are advanced by `step` at each index. + /// + /// The number of windows is `max(0, (shape[dim] - size).ceil_div(step))`. + /// + /// # Warning + /// + /// For the `ndarray` and `candle` backends; this is not a view but a full copy. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor to unfold; of shape ``[pre=..., dim shape, post=...]`` + /// * `dim` - the dimension to unfold. + /// * `size` - the size of each unfolded window. + /// * `step` - the step between each window. + /// + /// # Returns + /// + /// A tensor view with shape ``[pre=..., windows, post=..., size]``. + fn unfold(tensor: BridgeTensor, dim: usize, size: usize, step: usize) -> BridgeTensor; +} diff --git a/crates/burn-tensor/src/bridge/ops/bool.rs b/crates/burn-tensor/src/bridge/ops/bool.rs new file mode 100644 index 0000000..1c65952 --- /dev/null +++ b/crates/burn-tensor/src/bridge/ops/bool.rs @@ -0,0 +1,256 @@ +use alloc::vec::Vec; +use burn_backend::{ + AutodiffBackend, Scalar, TensorData, TensorMetadata, + ops::{BoolTensorOps, TransactionPrimitive}, +}; +use burn_dispatch::Dispatch; +use burn_std::{DType, ExecutionError, IndexingUpdateOp, Shape, Slice}; + +use crate::{ + Bool, Device, + bridge::{BasicAutodiffOps, BasicOps, TransactionOp}, + ops::BridgeTensor, +}; + +impl TransactionOp for Bool { + fn register_transaction(tr: &mut TransactionPrimitive, tensor: BridgeTensor) { + tr.register_bool(tensor.into()); + } +} + +impl BasicOps for Bool { + fn empty(shape: Shape, device: &Device, dtype: DType) -> BridgeTensor { + if !dtype.is_bool() { + panic!("Expected bool data type, got {dtype:?}"); + } + BridgeTensor::bool(Dispatch::bool_empty( + shape, + device.as_dispatch(), + dtype.into(), + )) + } + + fn zeros(shape: Shape, device: &Device, dtype: DType) -> BridgeTensor { + if !dtype.is_bool() { + panic!("Expected bool data type, got {dtype:?}"); + } + BridgeTensor::bool(Dispatch::bool_zeros( + shape, + device.as_dispatch(), + dtype.into(), + )) + } + fn ones(shape: Shape, device: &Device, dtype: DType) -> BridgeTensor { + if !dtype.is_bool() { + panic!("Expected bool data type, got {dtype:?}"); + } + BridgeTensor::bool(Dispatch::bool_ones( + shape, + device.as_dispatch(), + dtype.into(), + )) + } + + fn full(shape: Shape, fill_value: Scalar, device: &Device, dtype: DType) -> BridgeTensor { + if !dtype.is_bool() { + panic!("Expected bool data type, got {dtype:?}"); + } + if fill_value.elem() { + BridgeTensor::bool(Dispatch::bool_ones( + shape, + device.as_dispatch(), + dtype.into(), + )) + } else { + BridgeTensor::bool(Dispatch::bool_zeros( + shape, + device.as_dispatch(), + dtype.into(), + )) + } + } + + fn reshape(tensor: BridgeTensor, shape: Shape) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_reshape(tensor.into(), shape)) + } + + fn transpose(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_transpose(tensor.into())) + } + + fn swap_dims(tensor: BridgeTensor, dim1: usize, dim2: usize) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_swap_dims(tensor.into(), dim1, dim2)) + } + + fn slice(tensor: BridgeTensor, slices: &[Slice]) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_slice(tensor.into(), slices)) + } + + fn slice_assign(tensor: BridgeTensor, slices: &[Slice], value: BridgeTensor) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_slice_assign( + tensor.into(), + slices, + value.into(), + )) + } + + fn select(tensor: BridgeTensor, dim: usize, indices: BridgeTensor) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_select(tensor.into(), dim, indices.into())) + } + + fn select_assign( + tensor: BridgeTensor, + dim: usize, + indices: BridgeTensor, + values: BridgeTensor, + update: IndexingUpdateOp, + ) -> BridgeTensor { + match update { + IndexingUpdateOp::Add => BridgeTensor::bool(Dispatch::bool_select_or( + tensor.into(), + dim, + indices.into(), + values.into(), + )), + _ => unimplemented!(), + } + } + + fn mask_where(tensor: BridgeTensor, mask: BridgeTensor, source: BridgeTensor) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_mask_where( + tensor.into(), + mask.into(), + source.into(), + )) + } + + fn mask_fill(tensor: BridgeTensor, mask: BridgeTensor, value: Scalar) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_mask_fill(tensor.into(), mask.into(), value)) + } + + fn gather(dim: usize, tensor: BridgeTensor, indices: BridgeTensor) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_gather(dim, tensor.into(), indices.into())) + } + + fn scatter( + dim: usize, + tensor: BridgeTensor, + indices: BridgeTensor, + values: BridgeTensor, + update: IndexingUpdateOp, + ) -> BridgeTensor { + match update { + IndexingUpdateOp::Add => BridgeTensor::bool(Dispatch::bool_scatter_or( + dim, + tensor.into(), + indices.into(), + values.into(), + )), + _ => unimplemented!(), + } + } + + fn scatter_nd( + _data: BridgeTensor, + _indices: BridgeTensor, + _values: BridgeTensor, + _reduction: IndexingUpdateOp, + ) -> BridgeTensor { + panic!("scatter_nd is not supported for bool tensors") + } + + fn gather_nd(_data: BridgeTensor, _indices: BridgeTensor) -> BridgeTensor { + panic!("gather_nd is not supported for bool tensors") + } + + fn device(tensor: &BridgeTensor) -> Device { + Device::new(tensor.as_dispatch().device()) + } + + fn to_device(tensor: BridgeTensor, device: &Device) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_to_device( + tensor.into(), + device.as_dispatch(), + )) + } + + async fn into_data_async(tensor: BridgeTensor) -> Result { + Dispatch::bool_into_data(tensor.into()).await + } + + fn from_data(data: TensorData, device: &Device, dtype: DType) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_from_data( + data.convert_dtype(dtype), + device.as_dispatch(), + )) + } + + fn repeat_dim(tensor: BridgeTensor, dim: usize, times: usize) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_repeat_dim(tensor.into(), dim, times)) + } + + fn equal(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_equal(lhs.into(), rhs.into())) + } + + fn not_equal(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_not_equal(lhs.into(), rhs.into())) + } + + fn equal_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_equal_elem(lhs.into(), rhs)) + } + + fn not_equal_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_not_equal_elem(lhs.into(), rhs)) + } + + fn cat(vectors: Vec, dim: usize) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_cat( + BridgeTensor::into_dispatch_vec(vectors), + dim, + )) + } + + fn any(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_any(tensor.into())) + } + + fn any_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_any_dim(tensor.into(), dim)) + } + + fn all(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_all(tensor.into())) + } + + fn all_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_all_dim(tensor.into(), dim)) + } + + fn permute(tensor: BridgeTensor, axes: &[usize]) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_permute(tensor.into(), axes)) + } + + fn expand(tensor: BridgeTensor, shape: Shape) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_expand(tensor.into(), shape)) + } + + fn flip(tensor: BridgeTensor, axes: &[usize]) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_flip(tensor.into(), axes)) + } + + fn unfold(tensor: BridgeTensor, dim: usize, size: usize, step: usize) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_unfold(tensor.into(), dim, size, step)) + } +} + +impl BasicAutodiffOps for Bool { + fn inner(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_inner(tensor.into())) + } + + fn from_inner(inner: BridgeTensor) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_from_inner(inner.into())) + } +} diff --git a/crates/burn-tensor/src/bridge/ops/float.rs b/crates/burn-tensor/src/bridge/ops/float.rs new file mode 100644 index 0000000..e20f66c --- /dev/null +++ b/crates/burn-tensor/src/bridge/ops/float.rs @@ -0,0 +1,1025 @@ +use alloc::vec::Vec; +use burn_backend::{ + AutodiffBackend, Distribution, Scalar, TensorData, TensorMetadata, TensorPrimitive, + ops::{FloatTensorOps, QTensorOps, TransactionPrimitive}, +}; +use burn_dispatch::Dispatch; +use burn_std::{DType, ExecutionError, IndexingUpdateOp, Shape, Slice}; + +use crate::{ + Device, Float, + bridge::{BasicAutodiffOps, BasicOps, FloatMathOps, Numeric, Ordered, TransactionOp}, + ops::{BridgeKind, BridgeTensor}, +}; + +fn from_q_primitive(prim: TensorPrimitive) -> BridgeTensor { + match prim { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + } +} + +macro_rules! q_bin_ops { + ($lhs:ident, $rhs:ident, $op:ident, $q_op:ident) => {{ + let (lkind, lhs) = $lhs.into_parts(); + let (rkind, rhs) = $rhs.into_parts(); + match (lkind, rkind) { + (BridgeKind::Float, BridgeKind::Float) => BridgeTensor::float(Dispatch::$op(lhs, rhs)), + (BridgeKind::QFloat, BridgeKind::QFloat) => from_q_primitive(Dispatch::$q_op(lhs, rhs)), + (BridgeKind::QFloat, BridgeKind::Float) => { + let dtype = rhs.dtype(); + BridgeTensor::float(Dispatch::$op(Dispatch::dequantize(lhs, dtype.into()), rhs)) + } + (BridgeKind::Float, BridgeKind::QFloat) => { + let dtype = lhs.dtype(); + BridgeTensor::float(Dispatch::$op(lhs, Dispatch::dequantize(rhs, dtype.into()))) + } + _ => panic!("Should be Float primitive kind"), + } + }}; +} +impl TransactionOp for Float { + fn register_transaction(tr: &mut TransactionPrimitive, tensor: BridgeTensor) { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => tr.register_float(TensorPrimitive::Float(tensor)), + _ => panic!("Should be Float primitive kind"), + } + } +} + +impl BasicOps for Float { + fn empty(shape: Shape, device: &Device, dtype: DType) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_empty( + shape, + device.as_dispatch(), + dtype.into(), + )) + } + + fn zeros(shape: Shape, device: &Device, dtype: DType) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_zeros( + shape, + device.as_dispatch(), + dtype.into(), + )) + } + fn ones(shape: Shape, device: &Device, dtype: DType) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_ones( + shape, + device.as_dispatch(), + dtype.into(), + )) + } + + fn full(shape: Shape, fill_value: Scalar, device: &Device, dtype: DType) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_full( + shape, + fill_value, + device.as_dispatch(), + dtype.into(), + )) + } + + fn reshape(tensor: BridgeTensor, shape: Shape) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_reshape(tensor, shape)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_reshape(tensor, shape)), + _ => panic!("Should be Float primitive kind"), + } + } + + fn transpose(tensor: BridgeTensor) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_transpose(tensor)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_transpose(tensor)), + _ => panic!("Should be Float primitive kind"), + } + } + + fn swap_dims(tensor: BridgeTensor, dim1: usize, dim2: usize) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_swap_dims(tensor, dim1, dim2)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_swap_dims(tensor, dim1, dim2)), + _ => panic!("Should be Float primitive kind"), + } + } + + fn slice(tensor: BridgeTensor, slices: &[Slice]) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_slice(tensor, slices)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_slice(tensor, slices)), + _ => panic!("Should be Float primitive kind"), + } + } + + fn slice_assign(tensor: BridgeTensor, slices: &[Slice], value: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_slice_assign( + tensor.into_float(), + slices, + value.into_float(), + )) + } + + fn select(tensor: BridgeTensor, dim: usize, indices: BridgeTensor) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => { + BridgeTensor::float(Dispatch::float_select(tensor, dim, indices.into())) + } + BridgeKind::QFloat => { + BridgeTensor::qfloat(Dispatch::q_select(tensor, dim, indices.into())) + } + _ => panic!("Should be Float primitive kind"), + } + } + + fn select_assign( + tensor: BridgeTensor, + dim: usize, + indices: BridgeTensor, + values: BridgeTensor, + update: IndexingUpdateOp, + ) -> BridgeTensor { + // Select assign is ambiguous for QFloat + match update { + IndexingUpdateOp::Add => BridgeTensor::float(Dispatch::float_select_add( + tensor.into_float(), + dim, + indices.into(), + values.into_float(), + )), + other => unimplemented!("Unsupported update op {other:?}"), + } + } + + fn mask_where(tensor: BridgeTensor, mask: BridgeTensor, source: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_mask_where( + tensor.into_float(), + mask.into(), + source.into_float(), + )) + } + + fn mask_fill(tensor: BridgeTensor, mask: BridgeTensor, value: Scalar) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_mask_fill( + tensor.into_float(), + mask.into(), + value, + )) + } + + fn gather(dim: usize, tensor: BridgeTensor, indices: BridgeTensor) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => { + BridgeTensor::float(Dispatch::float_gather(dim, tensor, indices.into())) + } + BridgeKind::QFloat => { + BridgeTensor::qfloat(Dispatch::q_gather(dim, tensor, indices.into())) + } + _ => panic!("Should be Float primitive kind"), + } + } + + fn scatter( + dim: usize, + tensor: BridgeTensor, + indices: BridgeTensor, + values: BridgeTensor, + update: IndexingUpdateOp, + ) -> BridgeTensor { + match update { + IndexingUpdateOp::Add => BridgeTensor::float(Dispatch::float_scatter_add( + dim, + tensor.into_float(), + indices.into(), + values.into_float(), + )), + other => unimplemented!("Unsupported update op {other:?}"), + } + } + + fn scatter_nd( + data: BridgeTensor, + indices: BridgeTensor, + values: BridgeTensor, + reduction: IndexingUpdateOp, + ) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_scatter_nd( + data.into_float(), + indices.into(), + values.into_float(), + reduction, + )) + } + + fn gather_nd(data: BridgeTensor, indices: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_gather_nd(data.into_float(), indices.into())) + } + + fn device(tensor: &BridgeTensor) -> Device { + let (kind, tensor) = tensor.as_parts(); + match kind { + BridgeKind::Float => Device::new(tensor.device()), + BridgeKind::QFloat => Device::new(tensor.device()), + _ => panic!("Should be Float primitive kind"), + } + } + + fn to_device(tensor: BridgeTensor, device: &Device) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => { + BridgeTensor::float(Dispatch::float_to_device(tensor, device.as_dispatch())) + } + BridgeKind::QFloat => { + BridgeTensor::qfloat(Dispatch::q_to_device(tensor, device.as_dispatch())) + } + _ => panic!("Should be Float primitive kind"), + } + } + + async fn into_data_async(tensor: BridgeTensor) -> Result { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => Dispatch::float_into_data(tensor).await, + BridgeKind::QFloat => Dispatch::q_into_data(tensor).await, + _ => panic!("Should be Float primitive kind"), + } + } + + fn from_data(data: TensorData, device: &Device, dtype: DType) -> BridgeTensor { + if matches!(data.dtype, DType::QFloat(_)) { + // When the source is QFloat, there is no conversion path possible. + BridgeTensor::qfloat(Dispatch::q_from_data(data, device.as_dispatch())) + } else if dtype.is_float() { + BridgeTensor::float(Dispatch::float_from_data( + data.convert_dtype(dtype), + device.as_dispatch(), + )) + } else { + panic!("Expected float dtype, got {dtype:?}") + } + } + + fn repeat_dim(tensor: BridgeTensor, dim: usize, times: usize) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => { + BridgeTensor::float(Dispatch::float_repeat_dim(tensor, dim, times)) + } + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_repeat_dim(tensor, dim, times)), + _ => panic!("Should be Float primitive kind"), + } + } + + fn cat(vectors: Vec, dim: usize) -> BridgeTensor { + match vectors.first().unwrap().kind() { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_cat( + BridgeTensor::into_dispatch_vec(vectors), + dim, + )), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_cat( + BridgeTensor::into_dispatch_vec(vectors), + dim, + )), + _ => panic!("Should be Float primitive kind"), + } + } + + fn equal(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + let lhs = lhs.into_float(); + BridgeTensor::bool(Dispatch::float_equal(lhs, rhs.into_float(), bool_dtype)) + } + + fn not_equal(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + let lhs = lhs.into_float(); + BridgeTensor::bool(Dispatch::float_not_equal(lhs, rhs.into_float(), bool_dtype)) + } + + fn equal_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + let lhs = lhs.into_float(); + BridgeTensor::bool(Dispatch::float_equal_elem(lhs, rhs, bool_dtype)) + } + + fn not_equal_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + let lhs = lhs.into_float(); + BridgeTensor::bool(Dispatch::float_not_equal_elem(lhs, rhs, bool_dtype)) + } + + fn any(tensor: BridgeTensor) -> BridgeTensor { + let bool_dtype = tensor.device_settings().bool_dtype; + let tensor = tensor.into_float(); + BridgeTensor::bool(Dispatch::float_any(tensor, bool_dtype)) + } + + fn any_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + let bool_dtype = tensor.device_settings().bool_dtype; + let tensor = tensor.into_float(); + BridgeTensor::bool(Dispatch::float_any_dim(tensor, dim, bool_dtype)) + } + + fn all(tensor: BridgeTensor) -> BridgeTensor { + let bool_dtype = tensor.device_settings().bool_dtype; + let tensor = tensor.into_float(); + BridgeTensor::bool(Dispatch::float_all(tensor, bool_dtype)) + } + + fn all_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + let bool_dtype = tensor.device_settings().bool_dtype; + let tensor = tensor.into_float(); + BridgeTensor::bool(Dispatch::float_all_dim(tensor, dim, bool_dtype)) + } + + fn permute(tensor: BridgeTensor, axes: &[usize]) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_permute(tensor, axes)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_permute(tensor, axes)), + _ => panic!("Should be Float primitive kind"), + } + } + + fn expand(tensor: BridgeTensor, shape: Shape) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_expand(tensor.into_float(), shape)) + } + + fn flip(tensor: BridgeTensor, axes: &[usize]) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_flip(tensor, axes)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_flip(tensor, axes)), + _ => panic!("Should be Float primitive kind"), + } + } + + fn unfold(tensor: BridgeTensor, dim: usize, size: usize, step: usize) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_unfold(tensor.into_float(), dim, size, step)) + } +} + +impl Numeric for Float { + fn add(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + q_bin_ops!(lhs, rhs, float_add, q_add) + } + + fn add_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let (kind, lhs) = lhs.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_add_scalar(lhs, rhs)), + BridgeKind::QFloat => match Dispatch::q_add_scalar(lhs, rhs) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn sub(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + q_bin_ops!(lhs, rhs, float_sub, q_sub) + } + + fn sub_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let (kind, lhs) = lhs.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_sub_scalar(lhs, rhs)), + BridgeKind::QFloat => match Dispatch::q_sub_scalar(lhs, rhs) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn div(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + q_bin_ops!(lhs, rhs, float_div, q_div) + } + + fn div_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let (kind, lhs) = lhs.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_div_scalar(lhs, rhs)), + BridgeKind::QFloat => match Dispatch::q_div_scalar(lhs, rhs) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + fn remainder(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_remainder( + lhs.into_float(), + rhs.into_float(), + )) + } + + fn remainder_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_remainder_scalar(lhs.into_float(), rhs)) + } + + fn mul(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + q_bin_ops!(lhs, rhs, float_mul, q_mul) + } + + fn mul_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let (kind, lhs) = lhs.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_mul_scalar(lhs, rhs)), + BridgeKind::QFloat => match Dispatch::q_mul_scalar(lhs, rhs) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + fn neg(tensor: BridgeTensor) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_neg(tensor)), + BridgeKind::QFloat => match Dispatch::q_neg(tensor) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn sum(tensor: BridgeTensor) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_sum(tensor)), + BridgeKind::QFloat => match Dispatch::q_sum(tensor) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn sum_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_sum_dim(tensor, dim)), + BridgeKind::QFloat => match Dispatch::q_sum_dim(tensor, dim) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn prod(tensor: BridgeTensor) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_prod(tensor)), + BridgeKind::QFloat => match Dispatch::q_prod(tensor) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn prod_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_prod_dim(tensor, dim)), + BridgeKind::QFloat => match Dispatch::q_prod_dim(tensor, dim) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn mean(tensor: BridgeTensor) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_mean(tensor)), + BridgeKind::QFloat => match Dispatch::q_mean(tensor) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn mean_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_mean_dim(tensor, dim)), + BridgeKind::QFloat => match Dispatch::q_mean_dim(tensor, dim) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn cumsum(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_cumsum(tensor, dim)), + BridgeKind::QFloat => match Dispatch::q_cumsum(tensor, dim) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn cumprod(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_cumprod(tensor, dim)), + BridgeKind::QFloat => match Dispatch::q_cumprod(tensor, dim) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn abs(tensor: BridgeTensor) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_abs(tensor)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_abs(tensor)), + _ => panic!("Should be Float primitive kind"), + } + } + + fn powi(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + q_bin_ops!(lhs, rhs, float_powf, q_powf) + } + + fn powi_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let (kind, lhs) = lhs.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_powi_scalar(lhs, rhs)), + BridgeKind::QFloat => match Dispatch::q_powi_scalar(lhs, rhs) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn random( + shape: Shape, + distribution: Distribution, + device: &Device, + dtype: DType, + ) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_random( + shape, + distribution, + device.as_dispatch(), + dtype.into(), + )) + } + + fn sign(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_sign(tensor.into_float())) + } + + /// Applies the matrix multiplication operation. + /// + /// `C = AB` + /// + /// # Panics + /// + /// If the two tensors don't have a compatible shape. + fn matmul(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + let (lkind, lhs) = lhs.into_parts(); + let (rkind, rhs) = rhs.into_parts(); + match (lkind, rkind) { + (BridgeKind::Float, BridgeKind::Float) => { + BridgeTensor::float(Dispatch::float_matmul(lhs, rhs)) + } + (BridgeKind::Float, BridgeKind::QFloat) => from_q_primitive(Dispatch::q_matmul( + TensorPrimitive::Float(lhs), + TensorPrimitive::QFloat(rhs), + )), + (BridgeKind::QFloat, BridgeKind::Float) => from_q_primitive(Dispatch::q_matmul( + TensorPrimitive::QFloat(lhs), + TensorPrimitive::Float(rhs), + )), + (BridgeKind::QFloat, BridgeKind::QFloat) => from_q_primitive(Dispatch::q_matmul( + TensorPrimitive::QFloat(lhs), + TensorPrimitive::QFloat(rhs), + )), + _ => panic!("Should be Float primitive kind"), + } + } +} +impl Ordered for Float { + fn sort(tensor: BridgeTensor, dim: usize, descending: bool) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_sort(tensor, dim, descending)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_sort(tensor, dim, descending)), + _ => panic!("Should be Float primitive kind"), + } + } + + fn sort_with_indices( + tensor: BridgeTensor, + dim: usize, + descending: bool, + ) -> (BridgeTensor, BridgeTensor) { + let settings = tensor.device_settings(); + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => { + let (values, indices) = + Dispatch::float_sort_with_indices(tensor, dim, descending, settings.int_dtype); + (BridgeTensor::float(values), BridgeTensor::int(indices)) + } + BridgeKind::QFloat => { + let (values, indices) = + Dispatch::q_sort_with_indices(tensor, dim, descending, settings.int_dtype); + (BridgeTensor::qfloat(values), BridgeTensor::int(indices)) + } + _ => panic!("Should be Float primitive kind"), + } + } + + fn argsort(tensor: BridgeTensor, dim: usize, descending: bool) -> BridgeTensor { + let settings = tensor.device_settings(); + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::int(Dispatch::float_argsort( + tensor, + dim, + descending, + settings.int_dtype, + )), + BridgeKind::QFloat => BridgeTensor::int(Dispatch::q_argsort( + tensor, + dim, + descending, + settings.int_dtype, + )), + _ => panic!("Should be Float primitive kind"), + } + } + + fn cummin(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_cummin(tensor, dim)), + BridgeKind::QFloat => match Dispatch::q_cummin(tensor, dim) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn cummax(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_cummax(tensor, dim)), + BridgeKind::QFloat => match Dispatch::q_cummax(tensor, dim) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn greater(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + let lhs = lhs.into_float(); + BridgeTensor::bool(Dispatch::float_greater(lhs, rhs.into_float(), bool_dtype)) + } + + fn greater_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + let lhs = lhs.into_float(); + BridgeTensor::bool(Dispatch::float_greater_elem(lhs, rhs, bool_dtype)) + } + + fn greater_equal(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + let lhs = lhs.into_float(); + BridgeTensor::bool(Dispatch::float_greater_equal( + lhs, + rhs.into_float(), + bool_dtype, + )) + } + + fn greater_equal_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + let lhs = lhs.into_float(); + BridgeTensor::bool(Dispatch::float_greater_equal_elem(lhs, rhs, bool_dtype)) + } + + fn lower(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + let lhs = lhs.into_float(); + BridgeTensor::bool(Dispatch::float_lower(lhs, rhs.into_float(), bool_dtype)) + } + + fn lower_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + let lhs = lhs.into_float(); + BridgeTensor::bool(Dispatch::float_lower_elem(lhs, rhs, bool_dtype)) + } + + fn lower_equal(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + let lhs = lhs.into_float(); + BridgeTensor::bool(Dispatch::float_lower_equal( + lhs, + rhs.into_float(), + bool_dtype, + )) + } + + fn lower_equal_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + let lhs = lhs.into_float(); + BridgeTensor::bool(Dispatch::float_lower_equal_elem(lhs, rhs, bool_dtype)) + } + + fn argmax(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + let settings = tensor.device_settings(); + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => { + BridgeTensor::int(Dispatch::float_argmax(tensor, dim, settings.int_dtype)) + } + BridgeKind::QFloat => { + BridgeTensor::int(Dispatch::q_argmax(tensor, dim, settings.int_dtype)) + } + _ => panic!("Should be Float primitive kind"), + } + } + + fn argtopk(tensor: BridgeTensor, dim: usize, k: usize) -> BridgeTensor { + let settings = tensor.device_settings(); + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => { + BridgeTensor::int(Dispatch::float_argtopk(tensor, dim, k, settings.int_dtype)) + } + BridgeKind::QFloat => { + BridgeTensor::int(Dispatch::q_argtopk(tensor, dim, k, settings.int_dtype)) + } + _ => panic!("Should be Float primitive kind"), + } + } + + fn argmin(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + let settings = tensor.device_settings(); + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => { + BridgeTensor::int(Dispatch::float_argmin(tensor, dim, settings.int_dtype)) + } + BridgeKind::QFloat => { + BridgeTensor::int(Dispatch::q_argmin(tensor, dim, settings.int_dtype)) + } + _ => panic!("Should be Float primitive kind"), + } + } + + fn max(tensor: BridgeTensor) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_max(tensor)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_max(tensor)), + _ => panic!("Should be Float primitive kind"), + } + } + + fn max_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_max_dim(tensor, dim)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_max_dim(tensor, dim)), + _ => panic!("Should be Float primitive kind"), + } + } + + fn topk(tensor: BridgeTensor, dim: usize, k: usize) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_topk(tensor, dim, k)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_topk(tensor, dim, k)), + _ => panic!("Should be Float primitive kind"), + } + } + + fn max_dim_with_indices(tensor: BridgeTensor, dim: usize) -> (BridgeTensor, BridgeTensor) { + let settings = tensor.device_settings(); + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => { + let (values, indices) = + Dispatch::float_max_dim_with_indices(tensor, dim, settings.int_dtype); + (BridgeTensor::float(values), BridgeTensor::int(indices)) + } + BridgeKind::QFloat => { + let (values, indices) = + Dispatch::q_max_dim_with_indices(tensor, dim, settings.int_dtype); + (BridgeTensor::qfloat(values), BridgeTensor::int(indices)) + } + _ => panic!("Should be Float primitive kind"), + } + } + + fn min(tensor: BridgeTensor) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_min(tensor)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_min(tensor)), + _ => panic!("Should be Float primitive kind"), + } + } + + fn min_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_min_dim(tensor, dim)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_min_dim(tensor, dim)), + _ => panic!("Should be Float primitive kind"), + } + } + + fn min_dim_with_indices(tensor: BridgeTensor, dim: usize) -> (BridgeTensor, BridgeTensor) { + let settings = tensor.device_settings(); + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => { + let (values, indices) = + Dispatch::float_min_dim_with_indices(tensor, dim, settings.int_dtype); + (BridgeTensor::float(values), BridgeTensor::int(indices)) + } + BridgeKind::QFloat => { + let (values, indices) = + Dispatch::q_min_dim_with_indices(tensor, dim, settings.int_dtype); + (BridgeTensor::qfloat(values), BridgeTensor::int(indices)) + } + _ => panic!("Should be Float primitive kind"), + } + } + + fn clamp(tensor: BridgeTensor, min: Scalar, max: Scalar) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_clamp(tensor, min, max)), + BridgeKind::QFloat => match Dispatch::q_clamp(tensor, min, max) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn clamp_min(tensor: BridgeTensor, min: Scalar) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_clamp_min(tensor, min)), + BridgeKind::QFloat => match Dispatch::q_clamp_min(tensor, min) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn clamp_max(tensor: BridgeTensor, max: Scalar) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_clamp_max(tensor, max)), + BridgeKind::QFloat => match Dispatch::q_clamp_max(tensor, max) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } + } + + fn max_abs(tensor: BridgeTensor) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_max_abs(tensor)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_max_abs(tensor)), + _ => panic!("Should be Float primitive kind"), + } + } + + fn max_abs_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_max_abs_dim(tensor, dim)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_max_abs_dim(tensor, dim)), + _ => panic!("Should be Float primitive kind"), + } + } +} + +impl FloatMathOps for Float { + fn square(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_powi_scalar(tensor.into_float(), 2.into())) + } + fn sqrt(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_sqrt(tensor.into_float())) + } + fn cos(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_cos(tensor.into_float())) + } + + fn sin(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_sin(tensor.into_float())) + } + + fn tan(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_tan(tensor.into_float())) + } + + fn cosh(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_cosh(tensor.into_float())) + } + + fn sinh(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_sinh(tensor.into_float())) + } + + fn tanh(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_tanh(tensor.into_float())) + } + + fn acos(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_acos(tensor.into_float())) + } + + fn acosh(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_acosh(tensor.into_float())) + } + + fn asin(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_asin(tensor.into_float())) + } + + fn asinh(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_asinh(tensor.into_float())) + } + + fn atan(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_atan(tensor.into_float())) + } + + fn atanh(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_atanh(tensor.into_float())) + } + + fn atan2(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_atan2(lhs.into_float(), rhs.into_float())) + } + + fn exp(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_exp(tensor.into_float())) + } + + fn log(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_log(tensor.into_float())) + } + + fn log1p(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_log1p(tensor.into_float())) + } +} + +impl BasicAutodiffOps for Float { + fn inner(tensor: BridgeTensor) -> BridgeTensor { + let (kind, tensor) = tensor.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::inner(tensor)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_inner(tensor)), + _ => panic!("Should be Float primitive kind"), + } + } + + fn from_inner(inner: BridgeTensor) -> BridgeTensor { + let (kind, tensor) = inner.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::from_inner(tensor)), + BridgeKind::QFloat => BridgeTensor::qfloat(Dispatch::q_from_inner(tensor)), + _ => panic!("Should be Float primitive kind"), + } + } +} diff --git a/crates/burn-tensor/src/bridge/ops/int.rs b/crates/burn-tensor/src/bridge/ops/int.rs new file mode 100644 index 0000000..77e141d --- /dev/null +++ b/crates/burn-tensor/src/bridge/ops/int.rs @@ -0,0 +1,493 @@ +use alloc::vec::Vec; +use burn_backend::{ + AutodiffBackend, Distribution, Scalar, TensorData, TensorMetadata, + ops::{IntTensorOps, TransactionPrimitive}, +}; +use burn_dispatch::Dispatch; +use burn_std::{DType, ExecutionError, IndexingUpdateOp, Shape, Slice}; + +use crate::{ + Device, Int, + bridge::{BasicAutodiffOps, BasicOps, Numeric, Ordered, TransactionOp}, + ops::BridgeTensor, +}; + +impl TransactionOp for Int { + fn register_transaction(tr: &mut TransactionPrimitive, tensor: BridgeTensor) { + tr.register_int(tensor.into()); + } +} +impl BasicOps for Int { + fn empty(shape: Shape, device: &Device, dtype: DType) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_empty( + shape, + device.as_dispatch(), + dtype.into(), + )) + } + + fn zeros(shape: Shape, device: &Device, dtype: DType) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_zeros( + shape, + device.as_dispatch(), + dtype.into(), + )) + } + fn ones(shape: Shape, device: &Device, dtype: DType) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_ones( + shape, + device.as_dispatch(), + dtype.into(), + )) + } + + fn full(shape: Shape, fill_value: Scalar, device: &Device, dtype: DType) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_full( + shape, + fill_value, + device.as_dispatch(), + dtype.into(), + )) + } + + fn reshape(tensor: BridgeTensor, shape: Shape) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_reshape(tensor.into(), shape)) + } + + fn transpose(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_transpose(tensor.into())) + } + + fn swap_dims(tensor: BridgeTensor, dim1: usize, dim2: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_swap_dims(tensor.into(), dim1, dim2)) + } + + fn slice(tensor: BridgeTensor, slices: &[Slice]) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_slice(tensor.into(), slices)) + } + + fn slice_assign(tensor: BridgeTensor, slices: &[Slice], value: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_slice_assign( + tensor.into(), + slices, + value.into(), + )) + } + + fn select(tensor: BridgeTensor, dim: usize, indices: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_select(tensor.into(), dim, indices.into())) + } + + fn select_assign( + tensor: BridgeTensor, + dim: usize, + indices: BridgeTensor, + values: BridgeTensor, + update: IndexingUpdateOp, + ) -> BridgeTensor { + match update { + IndexingUpdateOp::Add => BridgeTensor::int(Dispatch::int_select_add( + tensor.into(), + dim, + indices.into(), + values.into(), + )), + _ => unimplemented!(), + } + } + + fn mask_where(tensor: BridgeTensor, mask: BridgeTensor, source: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_mask_where( + tensor.into(), + mask.into(), + source.into(), + )) + } + + fn mask_fill(tensor: BridgeTensor, mask: BridgeTensor, value: Scalar) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_mask_fill(tensor.into(), mask.into(), value)) + } + + fn gather(dim: usize, tensor: BridgeTensor, indices: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_gather(dim, tensor.into(), indices.into())) + } + + fn scatter( + dim: usize, + tensor: BridgeTensor, + indices: BridgeTensor, + values: BridgeTensor, + update: IndexingUpdateOp, + ) -> BridgeTensor { + match update { + IndexingUpdateOp::Add => BridgeTensor::int(Dispatch::int_scatter_add( + dim, + tensor.into(), + indices.into(), + values.into(), + )), + _ => unimplemented!(), + } + } + + fn scatter_nd( + data: BridgeTensor, + indices: BridgeTensor, + values: BridgeTensor, + reduction: IndexingUpdateOp, + ) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_scatter_nd( + data.into(), + indices.into(), + values.into(), + reduction, + )) + } + + fn gather_nd(data: BridgeTensor, indices: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_gather_nd(data.into(), indices.into())) + } + + fn device(tensor: &BridgeTensor) -> Device { + Device::new(tensor.as_dispatch().device()) + } + + fn to_device(tensor: BridgeTensor, device: &Device) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_to_device(tensor.into(), device.as_dispatch())) + } + + async fn into_data_async(tensor: BridgeTensor) -> Result { + Dispatch::int_into_data(tensor.into()).await + } + + fn from_data(data: TensorData, device: &Device, dtype: DType) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_from_data( + data.convert_dtype(dtype), + device.as_dispatch(), + )) + } + + fn repeat_dim(tensor: BridgeTensor, dim: usize, times: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_repeat_dim(tensor.into(), dim, times)) + } + + fn equal(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + BridgeTensor::int(Dispatch::int_equal(lhs.into(), rhs.into(), bool_dtype)) + } + + fn not_equal(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + BridgeTensor::int(Dispatch::int_not_equal(lhs.into(), rhs.into(), bool_dtype)) + } + + fn equal_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + BridgeTensor::int(Dispatch::int_equal_elem(lhs.into(), rhs, bool_dtype)) + } + + fn not_equal_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + BridgeTensor::int(Dispatch::int_not_equal_elem(lhs.into(), rhs, bool_dtype)) + } + + fn cat(vectors: Vec, dim: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_cat( + BridgeTensor::into_dispatch_vec(vectors), + dim, + )) + } + + fn any(tensor: BridgeTensor) -> BridgeTensor { + let settings = tensor.device_settings(); + BridgeTensor::int(Dispatch::int_any(tensor.into(), settings.bool_dtype)) + } + + fn any_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + let bool_dtype = tensor.device_settings().bool_dtype; + BridgeTensor::int(Dispatch::int_any_dim(tensor.into(), dim, bool_dtype)) + } + + fn all(tensor: BridgeTensor) -> BridgeTensor { + let bool_dtype = tensor.device_settings().bool_dtype; + BridgeTensor::int(Dispatch::int_all(tensor.into(), bool_dtype)) + } + + fn all_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + let bool_dtype = tensor.device_settings().bool_dtype; + BridgeTensor::int(Dispatch::int_all_dim(tensor.into(), dim, bool_dtype)) + } + + fn permute(tensor: BridgeTensor, axes: &[usize]) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_permute(tensor.into(), axes)) + } + + fn expand(tensor: BridgeTensor, shape: Shape) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_expand(tensor.into(), shape)) + } + + fn flip(tensor: BridgeTensor, axes: &[usize]) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_flip(tensor.into(), axes)) + } + + fn unfold(tensor: BridgeTensor, dim: usize, size: usize, step: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_unfold(tensor.into(), dim, size, step)) + } +} + +impl Numeric for Int { + fn add(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_add(lhs.into(), rhs.into())) + } + fn add_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_add_scalar(lhs.into(), rhs)) + } + fn sub(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_sub(lhs.into(), rhs.into())) + } + fn sub_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_sub_scalar(lhs.into(), rhs)) + } + fn div(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_div(lhs.into(), rhs.into())) + } + fn div_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_div_scalar(lhs.into(), rhs)) + } + fn remainder(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_remainder(lhs.into(), rhs.into())) + } + fn remainder_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_remainder_scalar(lhs.into(), rhs)) + } + fn mul(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_mul(lhs.into(), rhs.into())) + } + fn mul_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_mul_scalar(lhs.into(), rhs)) + } + fn neg(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_neg(tensor.into())) + } + + fn sum(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_sum(tensor.into())) + } + + fn sum_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_sum_dim(tensor.into(), dim)) + } + + fn prod(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_prod(tensor.into())) + } + + fn prod_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_prod_dim(tensor.into(), dim)) + } + + fn mean(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_mean(tensor.into())) + } + fn mean_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_mean_dim(tensor.into(), dim)) + } + fn cumsum(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_cumsum(tensor.into(), dim)) + } + fn cumprod(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_cumprod(tensor.into(), dim)) + } + + fn abs(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_abs(tensor.into())) + } + + fn powi(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_powi(lhs.into(), rhs.into())) + } + + fn powi_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_powi_scalar(lhs.into(), rhs)) + } + + fn random( + shape: Shape, + distribution: Distribution, + device: &Device, + dtype: DType, + ) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_random( + shape, + distribution, + device.as_dispatch(), + dtype.into(), + )) + } + + fn sign(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_sign(tensor.into())) + } + + /// Applies the matrix multiplication operation. + /// + /// `C = AB` + /// + /// # Panics + /// + /// If the two tensors don't have a compatible shape. + fn matmul(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_matmul(lhs.into(), rhs.into())) + } +} + +impl Ordered for Int { + fn sort(tensor: BridgeTensor, dim: usize, descending: bool) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_sort(tensor.into(), dim, descending)) + } + + fn sort_with_indices( + tensor: BridgeTensor, + dim: usize, + descending: bool, + ) -> (BridgeTensor, BridgeTensor) { + let (values, indices) = Dispatch::int_sort_with_indices(tensor.into(), dim, descending); + (BridgeTensor::int(values), BridgeTensor::int(indices)) + } + + fn argsort(tensor: BridgeTensor, dim: usize, descending: bool) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_argsort(tensor.into(), dim, descending)) + } + + fn cummin(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_cummin(tensor.into(), dim)) + } + + fn cummax(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_cummax(tensor.into(), dim)) + } + + fn greater(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + BridgeTensor::int(Dispatch::int_greater(lhs.into(), rhs.into(), bool_dtype)) + } + + fn greater_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + BridgeTensor::bool(Dispatch::int_greater_elem(lhs.into(), rhs, bool_dtype)) + } + + fn greater_equal(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + BridgeTensor::bool(Dispatch::int_greater_equal( + lhs.into(), + rhs.into(), + bool_dtype, + )) + } + + fn greater_equal_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + BridgeTensor::bool(Dispatch::int_greater_equal_elem( + lhs.into(), + rhs, + bool_dtype, + )) + } + + fn lower(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + BridgeTensor::bool(Dispatch::int_lower(lhs.into(), rhs.into(), bool_dtype)) + } + + fn lower_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + BridgeTensor::bool(Dispatch::int_lower_elem(lhs.into(), rhs, bool_dtype)) + } + + fn lower_equal(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + BridgeTensor::bool(Dispatch::int_lower_equal( + lhs.into(), + rhs.into(), + bool_dtype, + )) + } + + fn lower_equal_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let bool_dtype = lhs.device_settings().bool_dtype; + BridgeTensor::bool(Dispatch::int_lower_equal_elem(lhs.into(), rhs, bool_dtype)) + } + + fn argmax(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_argmax(tensor.into(), dim)) + } + + fn argtopk(tensor: BridgeTensor, dim: usize, k: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_argtopk(tensor.into(), dim, k)) + } + + fn topk(tensor: BridgeTensor, dim: usize, k: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_topk(tensor.into(), dim, k)) + } + + fn argmin(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_argmin(tensor.into(), dim)) + } + + fn max(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_max(tensor.into())) + } + + fn max_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_max_dim(tensor.into(), dim)) + } + + fn max_dim_with_indices(tensor: BridgeTensor, dim: usize) -> (BridgeTensor, BridgeTensor) { + let (values, indices) = Dispatch::int_max_dim_with_indices(tensor.into(), dim); + (BridgeTensor::int(values), BridgeTensor::int(indices)) + } + + fn max_abs(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_max_abs(tensor.into())) + } + + fn max_abs_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_max_abs_dim(tensor.into(), dim)) + } + + fn min(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_min(tensor.into())) + } + + fn min_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_min_dim(tensor.into(), dim)) + } + + fn min_dim_with_indices(tensor: BridgeTensor, dim: usize) -> (BridgeTensor, BridgeTensor) { + let (values, indices) = Dispatch::int_min_dim_with_indices(tensor.into(), dim); + (BridgeTensor::int(values), BridgeTensor::int(indices)) + } + + fn clamp(tensor: BridgeTensor, min: Scalar, max: Scalar) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_clamp(tensor.into(), min, max)) + } + + fn clamp_min(tensor: BridgeTensor, min: Scalar) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_clamp_min(tensor.into(), min)) + } + + fn clamp_max(tensor: BridgeTensor, max: Scalar) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_clamp_max(tensor.into(), max)) + } +} + +impl BasicAutodiffOps for Int { + fn inner(tensor: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_inner(tensor.into())) + } + + fn from_inner(inner: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_from_inner(inner.into())) + } +} diff --git a/crates/burn-tensor/src/bridge/ops/math.rs b/crates/burn-tensor/src/bridge/ops/math.rs new file mode 100644 index 0000000..443d2bc --- /dev/null +++ b/crates/burn-tensor/src/bridge/ops/math.rs @@ -0,0 +1,298 @@ +use crate::{bridge::Numeric, ops::BridgeTensor}; + +/// Trait that lists some floating-point mathematical operations are common to all float-like dtypes. +/// +/// # Warnings +/// +/// This is an internal trait, use the public API provided by the [`Tensor`](crate::Tensor) struct. +pub(crate) trait FloatMathOps: Numeric { + /// Applies element wise square operation + /// + #[cfg_attr(doc, doc = "$y_i = x^{2}$")] + #[cfg_attr(not(doc), doc = "`y = x^2`")] + fn square(tensor: BridgeTensor) -> BridgeTensor; + + /// Applies element wise exponential operation. + /// + #[cfg_attr(doc, doc = "$y_i = e^{x_i}$")] + #[cfg_attr(not(doc), doc = "`y = e^x`")] + fn exp(tensor: BridgeTensor) -> BridgeTensor; + + /// Applies the natural logarithm of one plus the input tensor, element-wise. + /// + #[cfg_attr(doc, doc = r#"$y_i = \log_e\(x_i + 1\)$"#)] + #[cfg_attr(not(doc), doc = "`y_i = log(x_i + 1)`")] + fn log1p(tensor: BridgeTensor) -> BridgeTensor; + + /// Applies element wise natural log operation *ln*. + /// + #[cfg_attr(doc, doc = r#"$y_i = \log_e\(x_i\)$"#)] + #[cfg_attr(not(doc), doc = "`y_i = log(x_i)`")] + fn log(tensor: BridgeTensor) -> BridgeTensor; + + /// Applies element wise root square operation. + /// + #[cfg_attr(doc, doc = r#"$y_i = \sqrt{x_i}$"#)] + #[cfg_attr(not(doc), doc = "`y_i = sqrt(x_i)`")] + fn sqrt(tensor: BridgeTensor) -> BridgeTensor; + /// Returns a new tensor with cosine values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with cosine values. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For the cosine of a tensor, users should prefer the [`Tensor::cos`](crate::Tensor::cos) + /// function, which is more high-level and designed for public use. + fn cos(tensor: BridgeTensor) -> BridgeTensor; + + /// Returns a new tensor with sine values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with sine values. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For the sine of a tensor, users should prefer the [`Tensor::sin`](crate::Tensor::sin) + /// function, which is more high-level and designed for public use. + fn sin(tensor: BridgeTensor) -> BridgeTensor; + + /// Returns a new tensor with tangent values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with tangent values. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For the tangent of a tensor, users should prefer the [`Tensor::tan`](crate::Tensor::tan) + /// function, which is more high-level and designed for public use. + fn tan(tensor: BridgeTensor) -> BridgeTensor; + + /// Returns a new tensor with hyperbolic cosine values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with hyperbolic cosine values. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For the hyperbolic cosine of a tensor, users should prefer the [`Tensor::cosh`](crate::Tensor::cosh) + /// function, which is more high-level and designed for public use. + fn cosh(tensor: BridgeTensor) -> BridgeTensor; + + /// Returns a new tensor with hyperbolic sine values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with hyperbolic sine values. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For the hyperbolic sine of a tensor, users should prefer the [`Tensor::sinh`](crate::Tensor::sinh) + /// function, which is more high-level and designed for public use. + fn sinh(tensor: BridgeTensor) -> BridgeTensor; + + /// Returns a new tensor with hyperbolic tangent values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with hyperbolic tangent values. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For the hyperbolic tangent of a tensor, users should prefer the [`Tensor::tanh`](crate::Tensor::tanh) + /// function, which is more high-level and designed for public use. + fn tanh(tensor: BridgeTensor) -> BridgeTensor; + + /// Returns a new tensor with inverse cosine values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with inverse cosine values. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For the inverse cosine of a tensor, users should prefer the [`Tensor::acos`](crate::Tensor::acos) + /// function, which is more high-level and designed for public use. + fn acos(tensor: BridgeTensor) -> BridgeTensor; + + /// Returns a new tensor with inverse hyperbolic cosine values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with inverse hyperbolic cosine values. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For the inverse hyperbolic cosine of a tensor, users should prefer the [`Tensor::acosh`](crate::Tensor::acosh) + /// function, which is more high-level and designed for public use. + fn acosh(tensor: BridgeTensor) -> BridgeTensor; + + /// Returns a new tensor with inverse sine values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with inverse sine values. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For the inverse sine of a tensor, users should prefer the [`Tensor::asin`](crate::Tensor::asin) + /// function, which is more high-level and designed for public use. + fn asin(tensor: BridgeTensor) -> BridgeTensor; + + /// Returns a new tensor with inverse hyperbolic sine values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with inverse hyperbolic sine values. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For the inverse hyperbolic sine of a tensor, users should prefer the [`Tensor::asinh`](crate::Tensor::asinh) + /// function, which is more high-level and designed for public use. + fn asinh(tensor: BridgeTensor) -> BridgeTensor; + + /// Returns a new tensor with inverse tangent values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with inverse tangent values. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For the inverse tangent of a tensor, users should prefer the [`Tensor::atan`](crate::Tensor::atan) + /// function, which is more high-level and designed for public use. + fn atan(tensor: BridgeTensor) -> BridgeTensor; + + /// Returns a new tensor with inverse hyperbolic tangent values. + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// + /// # Returns + /// + /// A tensor with the same shape as `tensor` with inverse hyperbolic tangent values. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For the inverse hyperbolic tangent of a tensor, users should prefer the [`Tensor::atanh`](crate::Tensor::atanh) + /// function, which is more high-level and designed for public use. + fn atanh(tensor: BridgeTensor) -> BridgeTensor; + + /// Returns a tensor with the four-quadrant inverse tangent values of `y` and `x`. + /// + /// # Arguments + /// + /// * `lhs` - The tensor with y coordinates. + /// * `rhs` - The tensor with x coordinates. + /// + /// # Returns + /// + /// A tensor with the four-quadrant inverse tangent values. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For the four-quadrant inverse tangent of two tensors, users should prefer the [`Tensor::atan2`](crate::Tensor::atan2) + /// function, which is more high-level and designed for public use. + fn atan2(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor; +} diff --git a/crates/burn-tensor/src/bridge/ops/mod.rs b/crates/burn-tensor/src/bridge/ops/mod.rs new file mode 100644 index 0000000..801df0d --- /dev/null +++ b/crates/burn-tensor/src/bridge/ops/mod.rs @@ -0,0 +1,14 @@ +mod autodiff; +mod base; +mod bool; +mod float; +mod int; +mod math; +mod numeric; +mod ordered; + +pub(crate) use autodiff::*; +pub(crate) use base::*; +pub(crate) use math::*; +pub(crate) use numeric::*; +pub(crate) use ordered::*; diff --git a/crates/burn-tensor/src/bridge/ops/numeric.rs b/crates/burn-tensor/src/bridge/ops/numeric.rs new file mode 100644 index 0000000..a141469 --- /dev/null +++ b/crates/burn-tensor/src/bridge/ops/numeric.rs @@ -0,0 +1,499 @@ +use burn_backend::{Distribution, Scalar}; +use burn_std::{DType, Shape}; + +use crate::{Device, bridge::BasicOps, ops::BridgeTensor}; + +/// Trait that list all operations that can be applied on all numerical tensors. +/// +/// # Warnings +/// +/// This is an internal trait, use the public API provided by the [`Tensor`](crate::Tensor) struct. +pub(crate) trait Numeric: BasicOps { + /// Adds two tensors together. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// The sum of the two tensors. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For adding tensors, users should prefer the [`Tensor::add`](crate::Tensor::add) + /// function, which is more high-level and designed for public use. + fn add(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor; + + /// Adds a scalar to a tensor element-wise. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side scalar. + /// + /// # Returns + /// + /// The sum of the tensor and the scalar. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For adding a scalar to a tensor, users should prefer the [`Tensor::add_scalar`](crate::Tensor::add_scalar) + /// function, which is more high-level and designed for public use. + fn add_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor; + + /// Subtracts two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// The difference of the two tensors. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For subtracting tensors, users should prefer the [`Tensor::sub`](crate::Tensor::sub) + /// function, which is more high-level and designed for public use. + fn sub(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor; + + /// Subtracts a scalar from a tensor element-wise. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side scalar. + /// + /// # Returns + /// + /// The difference of the tensor and the scalar. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For subtracting a scalar from a tensor, users should prefer the [`Tensor::sub_scalar`](crate::Tensor::sub_scalar) + /// function, which is more high-level and designed for public use. + fn sub_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor; + + /// Divides two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// The quotient of the two tensors. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For dividing tensors, users should prefer the [`Tensor::div`](crate::Tensor::div) + /// function, which is more high-level and designed for public use. + fn div(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor; + + /// Divides a tensor by a scalar element-wise. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side scalar. + /// + /// # Returns + /// + /// The quotient of the tensor and the scalar. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For dividing a tensor by a scalar, users should prefer the [`Tensor::div_scalar`](crate::Tensor::div_scalar) + /// function, which is more high-level and designed for public use. + fn div_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor; + + /// Computes the modulo element-wise. The result is the *signed* remainder of the division and its absolute value is + /// less than that of the divisor. + /// + /// # Arguments + /// + /// * `lhs` - The dividend. + /// * `rhs` - The divisor. + /// + /// # Returns + /// + /// The modulo of the input tensor with the divisor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For performing the modulo operation, users should prefer the [`Tensor::remainder`](crate::Tensor::remainder) + /// function, which is more high-level and designed for public use. + fn remainder(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor; + + /// Computes the modulo element-wise. The result is the *signed* remainder of the division and its absolute value is + /// less than that of the divisor. + /// + /// # Arguments + /// + /// * `lhs` - The dividend. + /// * `rhs` - The divisor. + /// + /// # Returns + /// + /// The modulo of the input tensor with the divisor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For performing the modulo operation, users should prefer the [`Tensor::remainder_scalar`](crate::Tensor::remainder_scalar) + /// function, which is more high-level and designed for public use. + fn remainder_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor; + + /// Multiplies two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// The product of the two tensors. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For multiplying tensors, users should prefer the [`Tensor::mul`](crate::Tensor::mul) + /// function, which is more high-level and designed for public use. + fn mul(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor; + + /// Multiplies a tensor by a scalar element-wise. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side scalar. + /// + /// # Returns + /// + /// The product of the tensor and the scalar. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For multiplying a tensor by a scalar, users should prefer the [`Tensor::mul_scalar`](crate::Tensor::mul_scalar) + /// function, which is more high-level and designed for public use. + fn mul_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor; + + /// Negates a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to negate. + /// + /// # Returns + /// + /// The negated tensor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For negating a tensor, users should prefer the [`Tensor::neg`](crate::Tensor::neg) + /// function, which is more high-level and designed for public use. + fn neg(tensor: BridgeTensor) -> BridgeTensor; + + /// Returns the signs of the elements of a tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor. + /// + /// # Returns + /// + /// The signs of the elements of the tensor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For getting the signs of the elements of a tensor, users should prefer the [`Tensor::sign`](crate::Tensor::sign) + /// function, which is more high-level and designed for public use. + fn sign(tensor: BridgeTensor) -> BridgeTensor; + + /// Sums all the elements of the tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to sum. + /// + /// # Returns + /// + /// The sum of all the elements of the tensor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For summing all the elements of a tensor, users should prefer the [`Tensor::sum`](crate::Tensor::sum) + /// function, which is more high-level and designed for public use. + fn sum(tensor: BridgeTensor) -> BridgeTensor; + + /// Sums all the elements of the tensor along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to sum. + /// * `dim` - The dimension along which to sum. + /// + /// # Returns + /// + /// The sum of all the elements of the tensor along the specified dimension. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For summing all the elements of a tensor along a dimension, users should prefer the [`Tensor::sum_dim`](crate::Tensor::sum_dim) + /// function, which is more high-level and designed for public use. + fn sum_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor; + + /// Computes the product of all the elements of the tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the product of. + /// + /// # Returns + /// + /// The product of all the elements of the tensor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For computing the product of all the elements of a tensor, users should prefer the + /// [`Tensor::prod`](crate::Tensor::prod) function, which is more high-level and designed for public use. + fn prod(tensor: BridgeTensor) -> BridgeTensor; + + /// Computes the product of all the elements of the tensor along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the product of. + /// * `dim` - The dimension along which to compute the product. + /// + /// # Returns + /// + /// The product of all the elements of the tensor along the specified dimension. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For computing the product of all the elements of a tensor along a dimension, users should prefer the + /// [`Tensor::prod_dim`](crate::Tensor::prod_dim) function, which is more high-level and designed for public use. + fn prod_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor; + + /// Computes the mean of all the elements of the tensor. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the mean of. + /// + /// # Returns + /// + /// The mean of all the elements of the tensor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For computing the mean of all the elements of a tensor, users should prefer the [`Tensor::mean`](crate::Tensor::mean) + /// function, which is more high-level and designed for public use. + fn mean(tensor: BridgeTensor) -> BridgeTensor; + + /// Computes the mean of all the elements of the tensor along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the mean of. + /// * `dim` - The dimension along which to compute the mean. + /// + /// # Returns + /// + /// The mean of all the elements of the tensor along the specified dimension. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For computing the mean of all the elements of a tensor along a dimension, users should prefer the + /// [`Tensor::mean_dim`](crate::Tensor::mean_dim) function, which is more high-level and designed for public use. + fn mean_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor; + + /// Computes the cumulative sum of elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the cumulative sum of. + /// * `dim` - The dimension along which to compute the cumulative sum. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor, where each element is the cumulative sum + /// of all elements up to and including that position along the specified dimension. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For computing the cumulative sum of elements along a dimension, users should prefer the + /// [`Tensor::cumsum`](crate::Tensor::cumsum) function, which is more high-level and designed for public use. + fn cumsum(tensor: BridgeTensor, dim: usize) -> BridgeTensor; + + /// Computes the cumulative product of elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the cumulative product of. + /// * `dim` - The dimension along which to compute the cumulative product. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor, where each element is the cumulative product + /// of all elements up to and including that position along the specified dimension. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For computing the cumulative product of elements along a dimension, users should prefer the + /// [`Tensor::cumprod`](crate::Tensor::cumprod) function, which is more high-level and designed for public use. + fn cumprod(tensor: BridgeTensor, dim: usize) -> BridgeTensor; + + /// Calculate absolute value on all elements of a tensor + /// + /// # Arguments + /// + /// * `tensor` - The tensor to apply abs to. + /// + /// # Returns + /// + /// A tensor with absolute values. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For calculating abs of the elements of a tensor, users should prefer the [`Tensor::abs`](crate::Tensor::abs) + /// function, which is more high-level and designed for public use. + fn abs(tensor: BridgeTensor) -> BridgeTensor; + + /// Element-wise power of a tensor + /// + /// # Arguments + /// * `tensor` - The tensor to apply power to. + /// * `power` - The power to apply to the tensor. + fn powi(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor; + + /// Element-wise power of a tensor to a scalar int + /// + /// # Arguments + /// * `tensor` - The tensor to apply power to. + /// * `power` - The power to apply to the tensor. + fn powi_scalar(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor; + + /// Create a random tensor. + /// + /// # Arguments + /// + /// * `shape` - The shape of the output tensor. + /// * `distribution` - The distribution used to sample. + /// * `device` - The device to use. + /// * `dtype` - The target data type. + /// + /// # Returns + /// + /// A new tensor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// Users should prefer the [`Tensor::random`](crate::Tensor::random) + /// function, which is more high-level and designed for public use. + fn random( + shape: Shape, + distribution: Distribution, + device: &Device, + dtype: DType, + ) -> BridgeTensor; + + /// Applies the matrix multiplication operation. + /// + /// ```math + /// C = AB + /// ``` + fn matmul(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor; +} diff --git a/crates/burn-tensor/src/bridge/ops/ordered.rs b/crates/burn-tensor/src/bridge/ops/ordered.rs new file mode 100644 index 0000000..269f10e --- /dev/null +++ b/crates/burn-tensor/src/bridge/ops/ordered.rs @@ -0,0 +1,643 @@ +use burn_backend::Scalar; + +use crate::{bridge::Numeric, ops::BridgeTensor}; + +/// Trait that list all operations that can be applied on all numerical tensors +/// whose elements have a well-defined ordering. +/// +/// This includes operations such as comparisons, minimum/maximum reductions, +/// and other order-dependent computations that are not strictly valid for all numerical +/// types. +/// +/// # Warnings +/// +/// This is an internal trait, use the public API provided by the [`Tensor`](crate::Tensor) struct. +pub(crate) trait Ordered: Numeric { + /// Sort the elements of the input `tensor` by value along a given dimension. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// * `dim` - The axis along which to sort. + /// * `descending` - The sorting order. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor, where the elements are sorted by value. + /// + /// # Remarks + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// Users should prefer the [`Tensor::sort`](crate::Tensor::sort) + /// function, which is more high-level and designed for public use. + fn sort(tensor: BridgeTensor, dim: usize, descending: bool) -> BridgeTensor; + + /// Sort the elements of the input `tensor` by value along a given dimension. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// * `dim` - The axis along which to sort. + /// * `descending` - The sorting order. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor and corresponding indices, where + /// the elements are sorted by value and the indices map back to the original input tensor. + /// + /// # Remarks + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For sorting the elements of a tensor, users should prefer the [`Tensor::sort_with_indices`](crate::Tensor::sort_with_indices) + /// function, which is more high-level and designed for public use. + fn sort_with_indices( + tensor: BridgeTensor, + dim: usize, + descending: bool, + ) -> (BridgeTensor, BridgeTensor); + + /// Returns the indices that sort the elements of the input `tensor` by value along a given dimension. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `tensor` - The input tensor. + /// * `dim` - The axis along which to sort. + /// * `descending` - The sorting order. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor the indices map back to the original input tensor. + /// + /// # Remarks + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// Users should prefer the [`Tensor::argsort`](crate::Tensor::argsort) + /// function, which is more high-level and designed for public use. + fn argsort(tensor: BridgeTensor, dim: usize, descending: bool) -> BridgeTensor; + + /// Computes the cumulative minimum of elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the cumulative minimum of. + /// * `dim` - The dimension along which to compute the cumulative minimum. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor, where each element is the minimum + /// of all elements up to and including that position along the specified dimension. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For computing the cumulative minimum of elements along a dimension, users should prefer the + /// [`Tensor::cummin`](crate::Tensor::cummin) function, which is more high-level and designed for public use. + fn cummin(tensor: BridgeTensor, dim: usize) -> BridgeTensor; + + /// Computes the cumulative maximum of elements along a dimension. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to compute the cumulative maximum of. + /// * `dim` - The dimension along which to compute the cumulative maximum. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor, where each element is the maximum + /// of all elements up to and including that position along the specified dimension. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For computing the cumulative maximum of elements along a dimension, users should prefer the + /// [`Tensor::cummax`](crate::Tensor::cummax) function, which is more high-level and designed for public use. + fn cummax(tensor: BridgeTensor, dim: usize) -> BridgeTensor; + + /// Element-wise greater than comparison between two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// A boolean tensor with the same shape as the input tensors, where each element is true if the + /// corresponding element of the left hand side tensor is greater than the corresponding element + /// of the right hand side tensor, and false otherwise. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For element-wise greater than comparison between two tensors, users should prefer the + /// [`Tensor::greater`](crate::Tensor::greater) function, which is more high-level and designed for public use. + fn greater(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor; + + /// Element-wise greater than comparison between a tensor and a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side scalar. + /// + /// # Returns + /// + /// A boolean tensor with the same shape as the input tensor, where each element is true if the + /// corresponding element of the left hand side tensor is greater than the right hand side + /// scalar, and false otherwise. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For element-wise greater than comparison between a tensor and a scalar, users should prefer the + /// [`Tensor::greater_elem`](crate::Tensor::greater_elem) function, which is more high-level and designed for public use. + fn greater_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor; + + /// Element-wise greater than or equal comparison between two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// A boolean tensor with the same shape as the input tensors, where each element is true if the + /// corresponding element of the left hand side tensor is greater than or equal to the + /// corresponding element of the right hand side tensor, and false otherwise. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For element-wise greater than or equal comparison between two tensors, users should prefer the + /// [`Tensor::greater_equal`](crate::Tensor::greater_equal) function, which is more high-level and designed for public use. + fn greater_equal(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor; + + /// Element-wise greater than or equal comparison between a tensor and a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side scalar. + /// + /// # Returns + /// + /// A boolean tensor with the same shape as the input tensor, where each element is true if the + /// corresponding element of the left hand side tensor is greater than or equal to the right + /// hand side scalar, and false otherwise. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For element-wise greater than or equal comparison between a tensor and a scalar, users should prefer the + /// [`Tensor::greater_equal_elem`](crate::Tensor::greater_equal_elem) function, which is more high-level and designed for public use. + fn greater_equal_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor; + + /// Element-wise less than comparison between two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// A boolean tensor with the same shape as the input tensors, where each element is true if the + /// corresponding element of the left hand side tensor is less than the corresponding element of + /// the right hand side tensor, and false otherwise. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For element-wise less than comparison between two tensors, users should prefer the + /// [`Tensor::lower`](crate::Tensor::lower) function, which is more high-level and designed for public use. + fn lower(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor; + + /// Element-wise less than comparison between a tensor and a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side scalar. + /// + /// # Returns + /// + /// A boolean tensor with the same shape as the input tensor, where each element is true if the + /// corresponding element of the left hand side tensor is less than the right hand side scalar, + /// and false otherwise. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For element-wise less than comparison between a tensor and a scalar, users should prefer the + /// [`Tensor::lower_elem`](crate::Tensor::lower_elem) function, which is more high-level and designed for public use. + fn lower_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor; + + /// Element-wise less than or equal comparison between two tensors. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side tensor. + /// + /// # Returns + /// + /// A boolean tensor with the same shape as the input tensors, where each element is true if the + /// corresponding element of the left hand side tensor is less than or equal to the corresponding + /// element of the right hand side tensor, and false otherwise. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For element-wise less than or equal comparison between two tensors, users should prefer the + /// [`Tensor::lower_equal`](crate::Tensor::lower_equal) function, which is more high-level and designed for public use. + fn lower_equal(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor; + + /// Element-wise less than or equal comparison between a tensor and a scalar. + /// + /// # Arguments + /// + /// * `lhs` - The left hand side tensor. + /// * `rhs` - The right hand side scalar. + /// + /// # Returns + /// + /// A boolean tensor with the same shape as the input tensor, where each element is true if the + /// corresponding element of the left hand side tensor is less than or equal to the right hand + /// side scalar, and false otherwise. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For element-wise less than or equal comparison between a tensor and a scalar, users should prefer the + /// [`Tensor::lower_equal_elem`](crate::Tensor::lower_equal_elem) function, which is more high-level and designed for public use. + fn lower_equal_elem(lhs: BridgeTensor, rhs: Scalar) -> BridgeTensor; + + /// Gets the indices of the maximum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `dim` - The axis along which to get the indices of the maximum elements. + /// * `tensor` - The tensor to get the indices of the maximum elements from. + /// + /// # Returns + /// + /// A tensor where the dimension `dim` has size 1 and all other dimensions + /// are the same as the input tensor. Each element is the index of the maximum + /// value. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For getting the indices of the maximum elements of a tensor along an axis, users should prefer the + /// [`Tensor::argmax`](crate::Tensor::argmax) function, which is more high-level and designed for public use. + fn argmax(tensor: BridgeTensor, dim: usize) -> BridgeTensor; + + /// Gets the indices of the k maximum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `dim` - The axis along which to get the indices of the maximum elements. + /// * `tensor` - The tensor to get the indices of the maximum elements from. + /// * `k` - k maximum elements to get + /// + /// # Returns + /// + /// A tensor where the dimension `dim` has size `k` and all other dimensions + /// are the same as the input tensor. Each element is the index of one of the + /// `k` largest values along the specified axis. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For getting the indices of the k maximum elements of a tensor along an axis, users should prefer the + /// [`Tensor::argtopk`](crate::Tensor::argtopk) function, which is more high-level and designed for public use. + fn argtopk(tensor: BridgeTensor, dim: usize, k: usize) -> BridgeTensor; + + /// Gets the values of the k maximum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `dim` - The axis along which to get the values of the maximum elements. + /// * `tensor` - The tensor to get the values of the maximum elements from. + /// * `k` - k maximum elements to get + /// + /// # Returns + /// + /// A tensor where the dimension `dim` has size `k` and all other dimensions + /// are the same as the input tensor. Each element is the value of one of the + /// `k` largest values along the specified axis. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For getting the values of the k maximum elements of a tensor along an axis, users should prefer the + /// [`Tensor::topk`](crate::Tensor::topk) function, which is more high-level and designed for public use. + fn topk(tensor: BridgeTensor, dim: usize, k: usize) -> BridgeTensor; + + /// Gets the indices of the minimum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `dim` - The axis along which to get the indices of the minimum elements. + /// * `tensor` - The tensor to get the indices of the minimum elements from. + /// + /// # Returns + /// + /// A tensor where the dimension `dim` has size 1 and all other dimensions + /// are the same as the input tensor. Each element is the index of the minimum + /// value. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For getting the indices of the minimum elements of a tensor along an axis, users should prefer the + /// [`Tensor::argmin`](crate::Tensor::argmin) function, which is more high-level and designed for public use. + fn argmin(tensor: BridgeTensor, dim: usize) -> BridgeTensor; + + /// Gets the maximum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `dim` - The axis along which to get the maximum elements. + /// + /// # Returns + /// + /// A single-element tensor containing the maximum element of the input tensor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For getting the maximum elements of a tensor along an axis, users should prefer the + /// [`Tensor::max`](crate::Tensor::max) function, which is more high-level and designed for public use. + fn max(tensor: BridgeTensor) -> BridgeTensor; + + /// Gets the maximum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements from. + /// * `dim` - The axis along which to get the maximum elements. + /// + /// # Returns + /// + /// A tensor with the same rank as the input tensor, but the given dim set to a shape of 1. + /// Each element is the maximum element of the corresponding input dim. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For getting the maximum elements of a tensor along an axis, users should prefer the + /// [`Tensor::max_dim`](crate::Tensor::max_dim) function, which is more high-level and designed for public use. + fn max_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor; + + /// Gets the maximum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements from. + /// * `dim` - The axis along which to get the maximum elements. + /// + /// # Returns + /// + /// A tuple containing the maximum element of the input tensor, and a tensor with the same shape + /// as the input tensor, where each element is the index of the maximum element of the input tensor + /// at the corresponding index along the specified axis. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For getting the maximum elements of a tensor along an axis, users should prefer the + /// [`Tensor::max_dim_with_indices`](crate::Tensor::max_dim_with_indices) function, + /// which is more high-level and designed for public use. + fn max_dim_with_indices(tensor: BridgeTensor, dim: usize) -> (BridgeTensor, BridgeTensor); + + /// Gets the maximum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `dim` - The axis along which to get the maximum elements. + /// + /// # Returns + /// + /// A single-element tensor containing the maximum absolute element of the input tensor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For getting the maximum absolute elements of a tensor, users should prefer the + /// [`Tensor::max_abs`](crate::Tensor::max_abs) function, which is more high-level and designed for public use. + fn max_abs(tensor: BridgeTensor) -> BridgeTensor; + + /// Gets the maximum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the maximum elements from. + /// * `dim` - The axis along which to get the maximum elements. + /// + /// # Returns + /// + /// A tensor with the same rank as the input tensor, but the given dim set to a shape of 1. + /// Each element is the maximum absolute element of the corresponding input dim. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For getting the maximum elements of a tensor along an axis, users should prefer the + /// [`Tensor::max_abs_dim`](crate::Tensor::max_abs_dim) function, which is more high-level and designed for public use. + fn max_abs_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor; + + /// Gets the minimum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the minimum elements from. + /// + /// # Returns + /// + /// A single-element tensor containing the minimum element of the input tensor. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For getting the minimum elements of a tensor along an axis, users should prefer the + /// [`Tensor::min`](crate::Tensor::min) function, which is more high-level and designed for public use. + fn min(tensor: BridgeTensor) -> BridgeTensor; + + /// Gets the minimum elements of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the minimum elements from. + /// * `dim` - The axis along which to get the minimum elements. + /// + /// # Returns + /// + /// A tensor with the same rank as the input tensor, but the given dim set to a shape of 1. + /// Each element is the minimum element of the corresponding input dim. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For getting the minimum elements of a tensor along an axis, users should prefer the + /// [`Tensor::min_dim`](crate::Tensor::min_dim) function, which is more high-level and designed for public use. + fn min_dim(tensor: BridgeTensor, dim: usize) -> BridgeTensor; + + /// Gets the minimum elements and indices of a tensor along an axis. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to get the minimum elements from. + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensor and corresponding indices, where + /// each element is the minimum element of the input tensor at the corresponding index + /// along the specified axis. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users, and not recommended to import + /// or use this function directly. + /// + /// For getting the minimum elements of a tensor along an axis, users should prefer the + /// [`Tensor::min_dim_with_indices`](crate::Tensor::min_dim_with_indices) function, which is more high-level and designed for public use. + fn min_dim_with_indices(tensor: BridgeTensor, dim: usize) -> (BridgeTensor, BridgeTensor); + + /// Clamp the tensor between the given min and max values. + /// + /// # Arguments + /// + /// * `min` - The minimum value. + /// * `max` - The maximum value. + /// + /// # Returns + /// + /// A new tensor with the values clamped between the given min and max values. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users. + /// + /// For clamping a tensor between the given min and max values, users should prefer the + /// [`Tensor::clamp`](crate::Tensor::clamp) function, which is more high-level and designed for public use. + fn clamp(tensor: BridgeTensor, min: Scalar, max: Scalar) -> BridgeTensor; + + /// Clamps a tensor under a minimum value. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to clamp. + /// * `min` - The minimum value. + /// + /// # Returns + /// + /// A new tensor with the values clamped under the given min value. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users. + /// + /// For clamping a tensor under a minimum value, users should prefer the + /// [`Tensor::clamp_min`](crate::Tensor::clamp_min) function, which is more high-level and designed for public use. + fn clamp_min(tensor: BridgeTensor, min: Scalar) -> BridgeTensor; + + /// Clamps a tensor over a maximum value. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to clamp. + /// * `max` - The maximum value. + /// + /// # Returns + /// + /// A new tensor with the values clamped over the given max value. + /// + /// # Remarks + /// + /// This is a low-level function used internally by the library to call different backend functions + /// with static dispatch. It is not designed for direct usage by users. + /// + /// For clamping a tensor over a maximum value, users should prefer the + /// [`Tensor::clamp_max`](crate::Tensor::clamp_max) function, which is more high-level and designed for public use. + fn clamp_max(tensor: BridgeTensor, max: Scalar) -> BridgeTensor; +} diff --git a/crates/burn-tensor/src/device.rs b/crates/burn-tensor/src/device.rs new file mode 100644 index 0000000..5186068 --- /dev/null +++ b/crates/burn-tensor/src/device.rs @@ -0,0 +1,1060 @@ +pub use burn_std::{ + DeviceError, DeviceSettings, ExecutionError, backtrace::BackTrace, device::DeviceId, +}; + +use burn_backend::{Backend, DeviceOps}; +#[allow(unused)] +use burn_dispatch::DispatchDeviceId; +use burn_dispatch::{Dispatch, DispatchDevice}; +use burn_std::{BoolDType, FloatDType, IntDType, TensorData}; + +#[cfg(feature = "remote-websocket")] +use alloc::string::String; +use alloc::vec; +use alloc::vec::Vec; + +/// A high-level device handle for tensor operations. +/// +/// [`Device`] provides a unified interface to interact with the underlying compute backend. +/// +/// Autodiff support is a property of the device rather than a separate type parameter. +#[cfg_attr( + feature = "autodiff", + doc = "Wrap a device with [`.autodiff()`](Device::autodiff) to enable automatic differentiation with the device." +)] +#[cfg_attr( + not(feature = "autodiff"), + doc = "Enable the `autodiff` feature to add automatic differentiation support to devices." +)] +/// +/// # Backend selection +/// +/// Enable the desired backend via Cargo feature flags, then call the +/// corresponding factory method: +/// +/// ```rust,ignore +/// // Default CUDA device (requires the `cuda` feature). +/// let device = Device::cuda(DeviceIndex::Default); +/// +/// // CUDA device at hardware index 1. +/// let device = Device::cuda(1); +/// +/// // WGPU with explicit selector (requires `wgpu`/`vulkan`/`metal`/`webgpu`). +/// let device = Device::wgpu(DeviceKind::DiscreteGpu(0)); +/// +/// // Default device for whichever backend is enabled. +/// let device = Default::default(); +/// ``` +/// +/// Available factory methods (each gated by its matching Cargo feature): +/// `Device::cpu`, `Device::cuda` / `Device::rocm` / `Device::libtorch_cuda` +/// (take an integer index or a [`DeviceIndex`]), `Device::wgpu` / +/// `Device::vulkan` / `Device::metal` / `Device::webgpu` (take a +/// [`DeviceKind`]), `Device::flex`, `Device::ndarray`, `Device::libtorch`, +/// `Device::libtorch_mps`, `Device::libtorch_vulkan`. +/// +/// # Autodiff +/// +/// Requires `autodiff` feature. +/// +/// Gradient computation is opt-in for a device: +/// +/// ```rust,ignore +/// let device = Device::default().autodiff(); +/// +/// // Tensors created on this device will track gradients +/// let x = Tensor::<1>::from_floats([1.0, 2.0, 3.0], &device); +/// ``` +pub struct Device { + blob: device_opaque::Opaque, +} + +// Aligned, type-erased storage for `DispatchDevice`. See `crate::macros` for +// why this indirection exists (it keeps the dispatch type tree out of +// downstream MIR). +burn_std::obfuscate!( + type: DispatchDevice, + module: device_opaque, + derives: [Send, Sync] +); + +impl Clone for Device { + fn clone(&self) -> Self { + Self::new(self.as_dispatch().clone()) + } +} + +impl Default for Device { + fn default() -> Self { + Self::new(DispatchDevice::default()) + } +} + +impl core::fmt::Debug for Device { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + write!(f, "Device<{:?}>", self.as_dispatch()) + } +} + +// Manually implement both `eq` and `ne` to add documentation on equality. +#[allow(clippy::partialeq_ne_impl)] +impl PartialEq for Device { + /// Compares devices based on hardware identity. + /// + /// Returns `true` if both devices represent the same compute resource. + /// Note that this comparison ignores autodiff and checkpointing settings. + /// To check if two devices have identical capabilities, check [`Device::is_autodiff`]. + fn eq(&self, other: &Self) -> bool { + self.as_dispatch() == other.as_dispatch() + } + + /// Compares devices based on hardware identity. + /// + /// Returns `false` if both devices represent the same compute resource, + /// even if one has autodiff enabled and the other does not. + fn ne(&self, other: &Self) -> bool { + !self.eq(other) + } +} + +impl Eq for Device {} + +impl Device { + /// Wrap a backend-specific device in a unified [`Device`]. + /// + /// Used by: + /// - the backend-specific factory methods below (`Device::cuda`, etc.) + /// — these are the recommended entry points for downstream code; + /// - burn-tensor's bridge ops, which already hold a [`DispatchDevice`] + /// and just need to wrap it; + /// - direct callers (tests, type-erased helpers) that have a concrete + /// backend device type at hand. + /// + /// Anything convertible into [`DispatchDevice`] is accepted, including + /// `DispatchDevice` itself. + pub fn new(device: impl Into) -> Self { + Self { + blob: device_opaque::Opaque::new(device.into()), + } + } + + /// Borrow the underlying [`DispatchDevice`]. + /// + /// The inverse of [`Device::new`]. Useful to backend-extension authors who need to dispatch on + /// the concrete backend variant (e.g. matching `DispatchDevice::Remote(_)`). + pub fn as_dispatch(&self) -> &DispatchDevice { + self.blob.as_ref() + } + + /// Crate-internal owning extraction of the underlying dispatch device. + pub(crate) fn into_dispatch(self) -> DispatchDevice { + self.blob.into_inner() + } +} + +impl> From for Device { + fn from(device: D) -> Self { + Self::new(device) + } +} + +/// Selector for the hardware index of a backend whose devices are simply +/// indexed (e.g. CUDA, ROCm). +/// +/// Backend factory methods that take an index (`Device::cuda`, `Device::rocm`, +/// `Device::libtorch_cuda`) accept `impl Into`, so the common +/// shorthand is to pass a plain integer literal: +/// +/// ```rust,ignore +/// Device::cuda(0); // hardware index 0 +/// Device::cuda(DeviceIndex::Default); // backend-chosen default +/// ``` +#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Default)] +pub enum DeviceIndex { + /// Target a specific hardware device by its index. + Specified(usize), + /// Let the backend pick its default device (typically index `0`). + #[default] + Default, +} + +impl DeviceIndex { + /// Construct a [`DeviceIndex::Specified`] from anything convertible into + /// a `usize`. + pub fn new(index: impl Into) -> Self { + Self::Specified(index.into()) + } + + /// Resolve to a concrete hardware index, defaulting to `0` for + /// [`DeviceIndex::Default`]. Backend factory methods are each gated by a + /// Cargo feature, so this looks dead when none of them are enabled. + #[allow(dead_code)] + fn resolve(self) -> usize { + match self { + DeviceIndex::Specified(i) => i, + DeviceIndex::Default => 0, + } + } +} + +impl From for DeviceIndex { + fn from(i: usize) -> Self { + Self::Specified(i) + } +} + +impl From for DeviceIndex { + fn from(i: u32) -> Self { + Self::Specified(i as usize) + } +} + +impl From for DeviceIndex { + fn from(i: u64) -> Self { + Self::Specified(i as usize) + } +} + +impl From for DeviceIndex { + fn from(i: i32) -> Self { + Self::Specified(usize::try_from(i).expect("device index must be non-negative")) + } +} + +impl From for DeviceIndex { + fn from(i: i64) -> Self { + Self::Specified(usize::try_from(i).expect("device index must be non-negative")) + } +} + +/// Selector for the more flexible backends whose device handle is a tagged +/// enum (e.g. WGPU, which can target a discrete/integrated/virtual GPU, a CPU +/// adapter, an externally-created wgpu setup, or just "best available"). +/// +/// The variants mirror `WgpuDevice` from cubecl so the mapping is direct, but +/// it is kept as a burn-owned enum so callers don't have to depend on cubecl. +#[derive(Clone, Debug, Hash, PartialEq, Eq, Default)] +pub enum DeviceKind { + /// Discrete GPU with the given index. The index is the index of the discrete GPU in the list + /// of all discrete GPUs found on the system. + DiscreteGpu(usize), + + /// Integrated GPU with the given index. The index is the index of the integrated GPU in the + /// list of all integrated GPUs found on the system. + IntegratedGpu(usize), + + /// Virtual GPU with the given index. The index is the index of the virtual GPU in the list of + /// all virtual GPUs found on the system. + VirtualGpu(usize), + + /// CPU. + Cpu, + + /// The best available device found with the current graphics API. + /// + /// This will prioritize GPUs wgpu recognizes as "high power". Additionally, you can override this using + /// the `CUBECL_WGPU_DEFAULT_DEVICE` environment variable. This variable is spelled as if i was a `WgpuDevice`, + /// so for example `CUBECL_WGPU_DEFAULT_DEVICE=IntegratedGpu(1)` or `CUBECL_WGPU_DEFAULT_DEVICE=Cpu` + #[default] + DefaultDevice, + + /// Use an externally created, existing, wgpu setup. This is helpful when using `CubeCL` in conjunction + /// with some existing wgpu setup (eg. egui or bevy), as resources can be transferred in & out of `CubeCL`. + /// + /// # Notes + /// + /// This can be initialized with `init_device` from the wgpu runtime. + Existing(u32), +} + +impl Device { + /// Default CPU device backed by CubeCL's CPU backend. + #[cfg(feature = "cpu")] + pub fn cpu() -> Self { + Self::new(burn_dispatch::devices::CpuDevice::default()) + } + + /// CUDA device at the given hardware index. + /// + /// Accepts a plain integer (e.g. `Device::cuda(0)`) or a + /// [`DeviceIndex`] — use [`DeviceIndex::Default`] to let the backend + /// pick. + #[cfg(feature = "cuda")] + pub fn cuda(index: impl Into) -> Self { + Self::new(burn_dispatch::devices::CudaDevice::new( + index.into().resolve(), + )) + } + + /// ROCm/HIP device at the given hardware index. + /// + /// Same selector semantics as [`Device::cuda`]. + #[cfg(feature = "rocm")] + pub fn rocm(index: impl Into) -> Self { + Self::new(burn_dispatch::devices::RocmDevice::new( + index.into().resolve(), + )) + } + + /// Flex backend device. + #[cfg(feature = "flex")] + pub fn flex() -> Self { + Self::new(burn_dispatch::devices::FlexDevice) + } + + /// Default NdArray (CPU) device. + #[cfg(feature = "ndarray")] + pub fn ndarray() -> Self { + Self::new(burn_dispatch::devices::NdArrayDevice::default()) + } + + /// LibTorch CPU device. + #[cfg(feature = "tch")] + pub fn libtorch() -> Self { + Self::new(burn_dispatch::devices::LibTorchDevice::Cpu) + } + + /// LibTorch CUDA device at the given hardware index. + #[cfg(feature = "tch")] + pub fn libtorch_cuda(index: impl Into) -> Self { + Self::new(burn_dispatch::devices::LibTorchDevice::Cuda( + index.into().resolve(), + )) + } + + /// LibTorch Metal Performance Shaders (MPS) device. + #[cfg(feature = "tch")] + pub fn libtorch_mps() -> Self { + Self::new(burn_dispatch::devices::LibTorchDevice::Mps) + } + + /// LibTorch Vulkan device. + #[cfg(feature = "tch")] + pub fn libtorch_vulkan() -> Self { + Self::new(burn_dispatch::devices::LibTorchDevice::Vulkan) + } + + /// Legacy WebSocket remote device. New integrations should prefer [`Device::remote_iroh`]. + /// + /// Connects to a burn-remote WebSocket server at the given address. `index` selects which of + /// the server's devices to use; two devices with the same address but different indices target + /// distinct devices on the same host. + #[cfg(feature = "remote-websocket")] + pub fn remote_websocket(address: &str, index: impl Into) -> Self { + let index = index.into().resolve(); + let device = burn_dispatch::devices::RemoteDevice::websocket(address, index); + device.connect(); // initializes the connection (required to get the device default settings) + Self::new(device) + } + + /// Iroh peer-to-peer remote device. + /// + /// `endpoint` is the application-owned Iroh endpoint to dial from; `peer` is the compute + /// server's identity (from [`RemoteSecret::id`](burn_dispatch::backends::remote::RemoteSecret::id)), + /// optionally carrying direct/relay dialing hints. + /// On wasm, use [`remote_iroh_async`](Self::remote_iroh_async) instead since sessions cannot + /// be opened synchronously. + #[cfg(all(feature = "remote", not(target_family = "wasm")))] + pub fn remote_iroh( + endpoint: &burn_dispatch::backends::remote::Endpoint, + peer: impl Into, + index: impl Into, + ) -> Self { + let index = index.into().resolve(); + let device = + burn_dispatch::backends::remote::RemoteDevice::iroh(endpoint, peer.into(), index); + device.connect(); + Self::new(device) + } + + /// Browser counterpart of [`remote_iroh`](Self::remote_iroh). Wasm cannot block to connect, + /// so the session is established asynchronously before the device is returned. + #[cfg(all(feature = "remote", any(target_family = "wasm", doc)))] + pub async fn remote_iroh_async( + endpoint: &burn_dispatch::backends::remote::Endpoint, + peer: impl Into, + index: impl Into, + ) -> Self { + let index = index.into().resolve(); + let device = + burn_dispatch::backends::remote::RemoteDevice::iroh(endpoint, peer.into(), index); + device.connect_async().await; + Self::new(device) + } + + /// Like `remote_iroh`, but carries an authorization credential the server's PeerAuthorizer + /// will check. Use against servers that require a credential; open servers take `remote_iroh`. + #[cfg(all(feature = "remote", not(target_family = "wasm")))] + pub fn remote_iroh_authorized( + endpoint: &burn_dispatch::backends::remote::Endpoint, + peer: impl Into, + index: impl Into, + credential: Vec, + ) -> Self { + let index = index.into().resolve(); + let device = burn_dispatch::backends::remote::RemoteDevice::iroh_authorized( + endpoint, + peer.into(), + index, + credential, + ); + device.connect(); + Self::new(device) + } + + /// Browser counterpart of `remote_iroh_authorized`. Establishes the session asynchronously. + #[cfg(all(feature = "remote", any(target_family = "wasm", doc)))] + pub async fn remote_iroh_authorized_async( + endpoint: &burn_dispatch::backends::remote::Endpoint, + peer: impl Into, + index: impl Into, + credential: Vec, + ) -> Self { + let index = index.into().resolve(); + let device = burn_dispatch::backends::remote::RemoteDevice::iroh_authorized( + endpoint, + peer.into(), + index, + credential, + ); + device.connect_async().await; + Self::new(device) + } + + /// WGPU device, selected via [`DeviceKind`]. + /// + /// This variant uses the runtime [`AutoCompiler`](burn_dispatch::backends::wgpu::AutoCompiler) + /// to dispatch to the most appropriate shader language (WGSL, SPIR-V, or MSL) based on the + /// enabled features. + /// + /// For [`DeviceKind::DefaultDevice`], the adapter is picked by `wgpu`'s + /// selection heuristics (high-power GPU preferred, or overridden by + /// `CUBECL_WGPU_DEFAULT_DEVICE`). + /// + /// `Device::vulkan`, `Device::metal`, and `Device::webgpu` also use the Wgpu runtime, + /// but bypass runtime dispatch by pinning specific compilers at compile time. + #[cfg(feature = "wgpu")] + pub fn wgpu(device_kind: DeviceKind) -> Self { + Self::new(DispatchDevice::Wgpu(wgpu_device(device_kind))) + } + + #[cfg(all(feature = "wgpu", target_family = "wasm"))] + /// Asynchronously creates a WGPU device, initializing the client. + pub async fn wgpu_async(device_kind: DeviceKind) -> Self { + Self::new(DispatchDevice::Wgpu(wgpu_init_async(device_kind).await)) + } + + /// Vulkan-backed WGPU device, selected via [`DeviceKind`]. + /// + /// Pins the wgpu shader compiler to SPIR-V at compile time, avoiding + /// the runtime [`AutoCompiler`](burn_dispatch::backends::wgpu::AutoCompiler) dispatch. + #[cfg(feature = "vulkan")] + pub fn vulkan(device_kind: DeviceKind) -> Self { + Self::new(DispatchDevice::Vulkan(wgpu_device(device_kind))) + } + + /// Metal-backed WGPU device, selected via [`DeviceKind`]. + /// + /// Pins the wgpu shader compiler to MSL at compile time. + #[cfg(feature = "metal")] + pub fn metal(device_kind: DeviceKind) -> Self { + Self::new(DispatchDevice::Metal(wgpu_device(device_kind))) + } + + /// WebGPU-backed device, selected via [`DeviceKind`]. + /// + /// Pins the wgpu shader compiler to WGSL at compile time. + #[cfg(feature = "webgpu")] + pub fn webgpu(device_kind: DeviceKind) -> Self { + Self::new(DispatchDevice::WebGpu(wgpu_device(device_kind))) + } + + /// Enables autodiff on this device. + /// + /// Autodiff is a property of the device: tensors created on the returned device + /// will participate in the autodiff graph. + /// + /// Only first-order autodiff is supported. Calling this method on a device that + /// already has autodiff enabled will panic. + /// + /// # Example + /// + /// ```rust,ignore + /// let device = Device::default().autodiff(); + /// let x = Tensor::<1>::from_floats([1.0, 2.0, 3.0], &device); + /// // x.backward() is now available + /// ``` + /// + /// # Panics + /// + /// Panics if autodiff is already enabled on this device. + #[cfg(feature = "autodiff")] + pub fn autodiff(self) -> Self { + match self.into_dispatch() { + DispatchDevice::Autodiff(_) => unimplemented!("Only first-order autodiff is supported"), + other => Self::new(DispatchDevice::autodiff(other)), + } + } + + /// Enables gradient checkpointing on the autodiff device. + /// + /// Gradient checkpointing recomputes activations during backpropagation for operations + /// marked as memory-bound, while compute-bound operations still cache their + /// output. This reduces peak memory usage at the cost of additional computation + /// for memory-bound ops. + /// + /// # Example + /// + /// ```rust,ignore + /// let device = Device::default().autodiff().gradient_checkpointing(); + /// ``` + /// + /// # Panics + /// + /// Panics if autodiff is not enabled on this device. + #[cfg(feature = "autodiff")] + pub fn gradient_checkpointing(self) -> Self { + match self.into_dispatch() { + DispatchDevice::Autodiff(device) => { + use burn_dispatch::CheckpointingStrategy; + + Self::new(DispatchDevice::autodiff_checkpointed( + device.inner(), + CheckpointingStrategy::Balanced, + )) + } + _ => panic!("Autodiff is not enabled on this device"), + } + } + + /// Returns the underlying device, removing the autodiff capability if present. + /// + /// If autodiff is not enabled, this method returns the device as-is. + /// + /// # Example + /// + /// ```rust,ignore + /// let device = Device::default().autodiff(); + /// let inner_device = device.inner(); + /// + /// assert!(!inner_device.is_autodiff()); + /// ``` + pub fn inner(self) -> Self { + if self.is_autodiff() { + Self::new(self.into_dispatch().inner()) + } else { + self + } + } + + /// Synchronize the device, waiting for all pending operations to complete. + /// + /// # Errors + /// + /// Returns an [`ExecutionError`] if an operation failed to execute. + pub fn sync(&self) -> Result<(), ExecutionError> { + Dispatch::sync(self.as_dispatch()) + } + + /// Flush the device's pending operations, handing them off for execution without waiting + /// for them to complete. + /// + /// Backends that buffer work hold registered operations in a local queue until enough + /// accumulate: the fusion backend batches ops to build optimizations, and the remote backend + /// batches them before sending them over the network. `flush` forces that queue out now — the + /// fusion backend processes its pending optimizations and the remote backend sends its batch to + /// the server. + /// + /// Unlike [`sync`](Self::sync), this does not block on results — it only ensures buffered + /// operations are dispatched instead of sitting idle. Eager backends, which execute each + /// operation as it is registered, have nothing buffered and treat this as a no-op. + pub fn flush(&self) { + Dispatch::flush(self.as_dispatch()) + } + + /// Seeds the random number generator for this device. + /// + /// Seeding before tensor operations that involve randomness (e.g. [`Tensor::random`](crate::Tensor::random)) + /// makes those operations reproducible in a single-threaded program. + /// + /// # Note + /// + /// Depending on the backend, the seed may be applied globally rather than scoped + /// to this specific device. It is guaranteed that at least this device will be seeded. + /// + /// # Example + /// + /// ```rust,ignore + /// let device = Default::default(); + /// device.seed(42); + /// let t = Tensor::<1>::random([8], Distribution::Default, &device); + /// ``` + pub fn seed(&self, seed: u64) { + Dispatch::seed(self.as_dispatch(), seed) + } + + /// Returns `true` if autodiff (gradient tracking) is enabled on this device. + /// + /// # Example + /// + /// ```rust,ignore + /// let device = Default::default(); + /// assert!(!device.is_autodiff()); + /// + /// let ad_device = device.autodiff(); + /// assert!(ad_device.is_autodiff()); + /// ``` + pub fn is_autodiff(&self) -> bool { + Dispatch::ad_enabled(self.as_dispatch()) + } + + /// Sets the current allocation mode to persistent. + pub fn memory_persistent_allocations< + Output: Send, + Input: Send, + Func: Fn(Input) -> Output + Send, + >( + &self, + input: Input, + func: Func, + ) -> Output { + Dispatch::memory_persistent_allocations(self.as_dispatch(), input, func) + } + + /// Triggers a memory cleanup on this device. + /// + /// The amount of memory reclaimed depends on the allocator implementation. + /// Calling this method does not guarantee that any memory will be freed. + pub fn memory_cleanup(&self) { + Dispatch::memory_cleanup(self.as_dispatch()); + } + + /// Prepares the given data for transfer between the CPU and accelerator devices such as GPUs. + /// + /// Depending on the backend, the data may be transferred to pinned memory + /// or another transfer-optimized format to improve transfer performance. + pub fn staging<'a, Iter>(&self, data: Iter) + where + Iter: Iterator, + { + Dispatch::staging(data, self.as_dispatch()); + } + + /// Returns the [`DeviceSettings`] for this device. + /// + /// Settings include the default float and integer data types used when creating + /// tensors on this device. + /// + /// See [`configure`](Device::configure) to configure them. + pub fn settings(&self) -> DeviceSettings { + burn_backend::get_device_settings::(self.as_dispatch()) + } + + /// Configures the [settings](DeviceSettings) for this device. + /// + /// This configures the dtype used when no explicit type is specified at tensor + /// creation time. + /// + /// Settings can only be initialized once per device, and must happen before any + /// tensor is created on the device. The first tensor operation will lock the device + /// to its defaults, causing subsequent initializations attempt to return + /// [`DeviceError::AlreadyInitialized`]. + /// + /// # Errors + /// + /// Returns [`DeviceError::AlreadyInitialized`] if settings have already been set + /// for this device (either by a prior call or because a tensor operation has + /// already occurred). + /// + /// # Example + /// + /// ```rust,ignore + /// let device = Default::default(); + /// + /// device.configure((FloatDType::F16, IntDType::I32))? + /// + /// // Float tensors will now use F16 + /// let floats = Tensor::<2>::zeros([2, 3], &device); + /// // Int tensors will now use I32 + /// let ints = Tensor::<2, Int>::zeros([2, 3], &device); + /// ``` + pub fn configure(&mut self, config: impl Into) -> Result<(), DeviceError> { + let mut config = config.into(); + + let defaults = self.as_dispatch().defaults(); + + let float_dtype = config.float_dtype.take().unwrap_or(defaults.float_dtype); + let int_dtype = config.int_dtype.take().unwrap_or(defaults.int_dtype); + let bool_dtype = config.bool_dtype.take().unwrap_or(defaults.bool_dtype); + + burn_backend::set_default_dtypes::( + self.as_dispatch(), + float_dtype, + int_dtype, + bool_dtype, + ) + } + + /// Retrieves all available [`Device`]s that match the given [`DeviceType`] filter. + /// + /// Local backends (CPU, CUDA, WGPU, …) enumerate the hardware found on the host. The + /// [`Remote`](DeviceType::Remote) variant instead lists every device hosted by the + /// `burn-remote` server at the given address — it connects to the server to learn how + /// many devices it exposes: + /// + /// ```rust,ignore + /// // Every CUDA device on this machine. + /// let local = Device::enumerate(DeviceType::Cuda); + /// + /// // Every device hosted by a remote server. + /// let remote = Device::enumerate(DeviceType::remote_websocket("ws://host:3000")); + /// + /// // Filters combine with `|`. + /// let both = Device::enumerate(DeviceType::Cuda | DeviceType::remote_websocket("ws://host:3000")); + /// ``` + pub fn enumerate(filter: impl Into) -> Devices { + #[allow(unused)] + let mut devices = Vec::new(); + + #[allow(clippy::never_loop)] // at least one backend is expected to be enabled. + for device_type in filter.into() { + #[allow(unused)] + let type_id = match device_type { + #[cfg(feature = "cpu")] + DeviceType::Cpu => DispatchDeviceId::Cpu, + #[cfg(feature = "cuda")] + DeviceType::Cuda => DispatchDeviceId::Cuda, + #[cfg(feature = "rocm")] + DeviceType::Rocm => DispatchDeviceId::Rocm, + #[cfg(feature = "wgpu")] + DeviceType::Wgpu => DispatchDeviceId::Wgpu, + #[cfg(feature = "metal")] + DeviceType::Metal => DispatchDeviceId::Metal, + #[cfg(feature = "vulkan")] + DeviceType::Vulkan => DispatchDeviceId::Vulkan, + #[cfg(feature = "webgpu")] + DeviceType::WebGpu => DispatchDeviceId::WebGpu, + #[cfg(feature = "flex")] + DeviceType::Flex => DispatchDeviceId::Flex, + #[cfg(feature = "ndarray")] + DeviceType::NdArray => DispatchDeviceId::NdArray, + #[cfg(feature = "tch")] + DeviceType::LibTorch => DispatchDeviceId::LibTorch, + // Remote devices are keyed by address, not a backend type id, so they take a + // dedicated enumeration path (connecting to the server for its device count). + #[cfg(feature = "remote-websocket")] + DeviceType::Remote(address) => { + for device in Dispatch::enumerate_remote_websocket(&address) { + devices.push(Device::new(device)); + } + continue; + } + }; + + #[allow(unreachable_code)] // need to have one backend enabled, so it is reachable + for device in Dispatch::enumerate(type_id) { + devices.push(Device::new(device)) + } + } + + Devices(devices) + } +} + +/// Map our backend-agnostic [`DeviceKind`] onto cubecl's `WgpuDevice` enum. +/// +/// Shared by [`Device::wgpu`], [`Device::vulkan`], [`Device::metal`], and +/// [`Device::webgpu`], which differ only in which Cargo feature gates them. +#[cfg(feature = "wgpu")] +fn wgpu_device(device_kind: DeviceKind) -> burn_dispatch::devices::WgpuDevice { + use burn_dispatch::devices::WgpuDevice; + match device_kind { + DeviceKind::DiscreteGpu(i) => WgpuDevice::DiscreteGpu(i), + DeviceKind::IntegratedGpu(i) => WgpuDevice::IntegratedGpu(i), + DeviceKind::VirtualGpu(i) => WgpuDevice::VirtualGpu(i), + DeviceKind::Cpu => WgpuDevice::Cpu, + DeviceKind::DefaultDevice => WgpuDevice::DefaultDevice, + DeviceKind::Existing(id) => WgpuDevice::Existing(id), + } +} + +#[cfg(all(feature = "wgpu", target_family = "wasm"))] +// TODO: this is only helpful for the default graphics api and runtime options.. we'd have to expose other methods but that leaks the types +// so we might have to introduce some wrapper types. +async fn wgpu_init_async(device_kind: DeviceKind) -> burn_dispatch::devices::WgpuDevice { + use burn_dispatch::backends::wgpu::{graphics::AutoGraphicsApi, init_setup_async}; + + let device = wgpu_device(device_kind); + init_setup_async::(&device, Default::default()).await; + device +} + +/// Represents the devices that can be used. +/// +/// `DeviceType` is used to filter the available device types for [`Device::enumerate`]. Most +/// variants are fieldless and select a backend's local hardware; [`Remote`](Self::Remote) +/// carries the network address of a `burn-remote` server whose devices should be listed. +/// +/// Variants combine into a [`DeviceFilter`] with the `|` operator, so a single +/// [`Device::enumerate`] call can span several backends and remote hosts. +#[allow(missing_docs)] +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum DeviceType { + #[cfg(feature = "cpu")] + Cpu, + #[cfg(feature = "cuda")] + Cuda, + #[cfg(feature = "rocm")] + Rocm, + #[cfg(feature = "wgpu")] + Wgpu, + #[cfg(feature = "metal")] + Metal, + #[cfg(feature = "vulkan")] + Vulkan, + #[cfg(feature = "webgpu")] + WebGpu, + #[cfg(feature = "flex")] + Flex, + #[cfg(feature = "ndarray")] + NdArray, + #[cfg(feature = "tch")] + LibTorch, + /// Devices hosted by the `burn-remote` server at the given address + /// (e.g. `"ws://host:3000"`). Unlike the other variants this is resolved at runtime by + /// connecting to the server, which reports how many devices it exposes. + #[cfg(feature = "remote-websocket")] + Remote(String), +} + +#[cfg(feature = "remote-websocket")] +impl DeviceType { + /// Filter selecting every device hosted by the `burn-remote` server at `address` + /// (e.g. `"ws://host:3000"`). + /// + /// Convenience for [`DeviceType::Remote`] that accepts anything string-like. + pub fn remote_websocket(address: impl Into) -> Self { + DeviceType::Remote(address.into()) + } +} + +/// A set of [`DeviceType`]s passed to [`Device::enumerate`]. +/// +/// Built from a single [`DeviceType`], a `Vec`, or by combining variants with the +/// `|` operator (`DeviceType::Cuda | DeviceType::Cpu`). Because [`DeviceType::Remote`] carries +/// an address, this is a plain list rather than a bitset. +#[derive(Debug, Clone, Default)] +pub struct DeviceFilter(Vec); + +impl DeviceFilter { + /// Create an empty filter. + pub fn new() -> Self { + Self::default() + } + + /// Add a [`DeviceType`] to the filter. + pub fn with(mut self, device_type: DeviceType) -> Self { + self.0.push(device_type); + self + } +} + +impl From for DeviceFilter { + fn from(value: DeviceType) -> Self { + DeviceFilter(vec![value]) + } +} + +impl From> for DeviceFilter { + fn from(value: Vec) -> Self { + DeviceFilter(value) + } +} + +impl IntoIterator for DeviceFilter { + type Item = DeviceType; + type IntoIter = alloc::vec::IntoIter; + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } +} + +impl core::ops::BitOr for DeviceType { + type Output = DeviceFilter; + fn bitor(self, rhs: Self) -> DeviceFilter { + DeviceFilter(vec![self, rhs]) + } +} + +impl core::ops::BitOr for DeviceFilter { + type Output = DeviceFilter; + fn bitor(mut self, rhs: DeviceType) -> DeviceFilter { + self.0.push(rhs); + self + } +} + +/// Configuration options used to initialize a device. +/// +/// Unlike [`DeviceSettings`], this type represents partial user-provided +/// configuration and does not require all settings to be specified. +/// +/// Any unspecified options will be resolved to device-specific defaults +/// when the device is initialized. +/// +/// Use [`Device::configure`] to apply this configuration to a device. +#[derive(new, Debug, Clone, Default)] +pub struct DeviceConfig { + /// Default floating-point data type. + pub float_dtype: Option, + + /// Default integer data type. + pub int_dtype: Option, + + /// Default boolean data type. + pub bool_dtype: Option, + // TODO: maybe quantization, but for now we keep this as device defaults +} + +impl DeviceConfig { + /// Sets the default floating-point data type for tensors created on the device. + pub fn float_dtype(mut self, dtype: impl Into) -> Self { + self.float_dtype = Some(dtype.into()); + self + } + + /// Sets the default integer data type for tensors created on the device. + pub fn int_dtype(mut self, dtype: impl Into) -> Self { + self.int_dtype = Some(dtype.into()); + self + } + + /// Sets the default boolean data type storage precision for tensors created on the device. + pub fn bool_dtype(mut self, dtype: impl Into) -> Self { + self.bool_dtype = Some(dtype.into()); + self + } +} + +impl From for DeviceConfig { + fn from(value: FloatDType) -> Self { + DeviceConfig::new(Some(value), None, None) + } +} + +impl From for DeviceConfig { + fn from(value: IntDType) -> Self { + DeviceConfig::new(None, Some(value), None) + } +} + +impl From for DeviceConfig { + fn from(value: BoolDType) -> Self { + DeviceConfig::new(None, None, Some(value)) + } +} + +impl From<(FloatDType, IntDType)> for DeviceConfig { + fn from(value: (FloatDType, IntDType)) -> Self { + DeviceConfig::new(Some(value.0), Some(value.1), None) + } +} + +/// A collection of [`Device`]s returned by [`Device::enumerate`]. +/// +/// This type provides bulk operations and transformations over multiple +/// devices, such as enabling autodiff or configuring the device settings. +/// +/// # Example +/// +/// ```rust,ignore +/// let mut devices = Device::enumerate(DeviceType::Cuda) +/// .autodiff(); +/// +/// devices.configure( +/// DeviceConfig::default().float_dtype(FloatDType::F16), +/// )?; +/// ``` +/// +/// `Devices` dereferences to a slice of [`Device`], so it can be iterated, +/// indexed, and passed anywhere a `&[Device]` is expected. +pub struct Devices(Vec); + +impl Devices { + /// Enables autodiff across all contained devices. + /// + /// Only first-order autodiff is supported. Calling this method on a device that + /// already has autodiff enabled will panic. + /// + /// See [`Device::autodiff`]. + #[cfg(feature = "autodiff")] + pub fn autodiff(mut self) -> Self { + for device in &mut self.0 { + *device = core::mem::take(device).autodiff(); + } + + self + } + + /// Configures the [settings](DeviceSettings) for all devices. + /// + /// This configures the dtype used when no explicit type is specified at tensor + /// creation time. + /// + /// Settings can only be initialized once per device, and must happen before any + /// tensor is created on the device. The first tensor operation will lock the device + /// to its defaults, causing subsequent initializations attempt to return + /// [`DeviceError::AlreadyInitialized`]. + /// + /// See [`Device::configure`]. + pub fn configure(&mut self, config: impl Into) -> Result<(), DeviceError> { + let config = config.into(); + for device in &mut self.0 { + device.configure(config.clone())?; + } + Ok(()) + } + + /// Returns the `Vec` of [`Device`]s. + pub fn into_vec(self) -> Vec { + self.0 + } +} + +// Loop over `&Devices` or `Devices` seamlessly +impl IntoIterator for Devices { + type Item = Device; + type IntoIter = alloc::vec::IntoIter; + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } +} + +impl core::ops::Deref for Devices { + type Target = [Device]; + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +#[cfg(all(test, feature = "flex", feature = "autodiff"))] +mod autodiff_move_tests { + use crate::{Device, Tensor}; + + // A non-tracked float tensor (e.g. a gradient) can be moved onto an autodiff device; it + // lands on the underlying hardware and stays non-tracked. Regression test for a panic in + // `float_to_device` ("Cannot move between autodiff and non-autodiff instances"). + #[test] + fn move_non_autodiff_float_tensor_to_autodiff_device() { + let device = Device::default(); + let ad_device = device.clone().autodiff(); + + let t = Tensor::<2>::from_floats([[1.0, 2.0], [3.0, 4.0]], &device); + let moved = t.to_device(&ad_device); + + assert_eq!( + moved.to_data().to_vec::().unwrap(), + vec![1.0, 2.0, 3.0, 4.0] + ); + } +} diff --git a/crates/burn-tensor/src/lib.rs b/crates/burn-tensor/src/lib.rs new file mode 100644 index 0000000..3f5807f --- /dev/null +++ b/crates/burn-tensor/src/lib.rs @@ -0,0 +1,59 @@ +#![cfg_attr(not(feature = "std"), no_std)] +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] + +//! This library provides the core abstractions required to run tensor operations with Burn. +//! `Tensor`s are generic over the backend to allow users to perform operations using different `Backend` implementations. +//! Burn's tensors also support auto-differentiation thanks to the `AutodiffBackend` trait. +//! +//! # Note for contributors: `*_impl` helpers +//! +//! Throughout this crate (e.g. in `tensor::api::float`, `tensor::api::int`, +//! `tensor::api::bool`, `tensor::api::cast`, and `tensor::activation`), public +//! generic methods on `Tensor` that need to call into `burn_dispatch` are +//! routed through small non-generic helper functions named `*_impl`, grouped +//! together at the bottom of each file under a banner like: +//! +//! ```text +//! // ===================================================================== +//! // Non-generic implementation helpers (outlined from the generic API). +//! // ===================================================================== +//! ``` +//! +//! These helpers take and return only `BridgeTensor` (a type-erased blob — no +//! `DispatchTensor` or other `burn_dispatch` types appear in their +//! signatures). Because the helpers are not generic over `D`, they are +//! compiled once, and the MIR of the public generic methods does not mention +//! any `burn_dispatch` types. Downstream crates that monomorphize the public +//! generic API therefore never have to resolve the cubecl-backed type tree, +//! which drastically cuts compile times for user code. +//! +//! When adding a new public method that calls a `Dispatch::*` op, follow the +//! same pattern: keep the generic method body thin and forward to a +//! non-generic `*_impl` helper alongside the existing ones. + +#[macro_use] +extern crate derive_new; + +extern crate alloc; + +mod bridge; +mod tensor; + +pub(crate) use tensor::check::macros::check; +pub use tensor::*; + +// Re-exported types +pub use burn_std::{ + AllocationProperty, Bytes, bf16, f16, + reader::{read_sync, try_read_sync}, + stream_id::StreamId, +}; + +mod device; +pub use device::*; + +#[cfg(feature = "remote-server")] +pub mod server; + +pub(crate) use burn_backend::TensorPrimitive; diff --git a/crates/burn-tensor/src/server.rs b/crates/burn-tensor/src/server.rs new file mode 100644 index 0000000..9292552 --- /dev/null +++ b/crates/burn-tensor/src/server.rs @@ -0,0 +1,114 @@ +//! Remote-execution server entry points. +//! +//! Hosts a Burn server that executes tensor operations on behalf of remote clients. The backend is +//! selected by the Device passed in; the transport is selected by Channel. Iroh is the primary +//! transport; WebSocket is retained for compatibility. +//! +//! Two serving modes: +//! +//! - Turnkey (start / start_async): no Iroh exposure. Pick an identity with RemoteSecret and pass +//! it in a Channel::Iroh; clients dial its public id. +//! - Composed (protocol): for applications that own their own Iroh router. Burn hands back only +//! its protocol handler to register alongside the application's own protocols. +//! +//! User-defined backends that implement BackendIr but are not part of DispatchDevice use +//! burn_remote::server::RemoteServerBuilder directly; that is also how custom operations +//! (backend extensions) are hosted, over either transport. + +use std::sync::Arc; + +use crate::Device; +pub use burn_dispatch::backends::remote::server::{ + AllowAll, AuthorizationRequest, PeerAuthorizer, RemoteProtocol, +}; +pub use burn_dispatch::backends::remote::telemetry; +pub use burn_dispatch::backends::remote::{Endpoint, RemoteSecret}; +pub use burn_dispatch::devices::BURN_REMOTE_ALPN; + +use telemetry::TelemetryProbe; + +/// Transport used to serve remote clients. Re-exported from `burn-remote` (via `burn-dispatch`) so +/// the whole stack shares one definition. +pub use burn_dispatch::remote_server::Channel; + +/// Build Burn's protocol handler for `device`'s backend. +/// +/// Returns a builder to optionally attach telemetry and an authorizer before calling `build`. +/// Register the result on an application-owned Iroh router under BURN_REMOTE_ALPN. +pub fn protocol(device: Device, endpoint: &Endpoint) -> RemoteProtocolBuilder<'_> { + RemoteProtocolBuilder::new(device, endpoint) +} + +/// Configures optional telemetry and authorization for a RemoteProtocol handler. +pub struct RemoteProtocolBuilder<'a> { + device: Device, + endpoint: &'a Endpoint, + probe: Option, + authorizer: Option>, +} + +impl<'a> RemoteProtocolBuilder<'a> { + /// Create a new builder for `device`'s backend, to register on `endpoint`. + pub fn new(device: Device, endpoint: &'a Endpoint) -> Self { + Self { + device, + endpoint, + probe: None, + authorizer: None, + } + } + + /// Attach a telemetry probe for per-session monitoring. Pair with + /// [`telemetry::TelemetryProbe::channel`] to obtain a subscription a dashboard can drain. + pub fn with_telemetry(mut self, probe: TelemetryProbe) -> Self { + self.probe = Some(probe); + self + } + + /// Authorize or reject each incoming compute session. The policy receives the peer identity, + /// the requested device index, and the opaque credential carried by the client's ticket. + pub fn with_authorizer(mut self, authorizer: impl PeerAuthorizer) -> Self { + self.authorizer = Some(Arc::new(authorizer)); + self + } + + /// Build the backend-erased protocol handler. + pub fn build(self) -> RemoteProtocol { + burn_dispatch::remote_server::remote_protocol( + self.device.into_dispatch(), + self.endpoint, + self.probe.unwrap_or_else(TelemetryProbe::disabled), + self.authorizer.unwrap_or_else(|| Arc::new(AllowAll)), + ) + } +} + +impl<'a> From> for RemoteProtocol { + fn from(builder: RemoteProtocolBuilder<'a>) -> Self { + builder.build() + } +} + +/// Start a remote-execution server, blocking the current thread. +/// +/// The backend is determined by `device`: e.g. `Device::cuda(0)` runs ops on +/// CUDA, `Device::flex()` on the Flex CPU backend. Autodiff devices are +/// transparently stripped; the autodiff graph is a client-side concern. +/// +/// # Panics +/// +/// Panics if `device` selects a backend that doesn't support remote execution +/// (currently `LibTorch`, or a `Remote` device; hosting on a remote device +/// makes no sense). +#[cfg(not(target_family = "wasm"))] +pub fn start(device: Device, channel: Channel) { + burn_dispatch::remote_server::start(device.into_dispatch(), channel) +} + +/// Start a remote-execution server on the caller's async runtime. +/// +/// See [`start`] for backend-selection rules. +#[cfg(not(target_family = "wasm"))] +pub async fn start_async(device: Device, channel: Channel) { + burn_dispatch::remote_server::start_async(device.into_dispatch(), channel).await +} diff --git a/crates/burn-tensor/src/tensor/activation/base.rs b/crates/burn-tensor/src/tensor/activation/base.rs new file mode 100644 index 0000000..cf55806 --- /dev/null +++ b/crates/burn-tensor/src/tensor/activation/base.rs @@ -0,0 +1,658 @@ +use burn_backend::ops::ActivationOps; +use burn_dispatch::Dispatch; + +use crate::check::TensorCheck; +use crate::ops::BridgeTensor; +use crate::{Tensor, check, s}; + +/// Applies the rectified linear unit function element-wise +/// as described in the paper [Deep Learning using Rectified Linear Units (ReLU)](https://arxiv.org/pdf/1803.08375). +/// +#[cfg_attr(doc, doc = "$$\\text{ReLU}\\(x\\) = \\(x\\)^+ = \\max\\(0, x\\)$$")] +#[cfg_attr(not(doc), doc = "`ReLU(x) = max(0, x)`")] +pub fn relu(tensor: Tensor) -> Tensor { + tensor.relu() +} + +/// Applies the leaky rectified linear unit function element-wise. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{LeakyReLU}\(x\) = \max\(0,x\) + \text{negative\\_slope} \cdot \min\(0, x\) +$$ + +or + +$$ +\text{LeakyReLU}(x) = + \begin{cases} + x & \text{if } x \geq 0 \newline + \text{negative\\_slope} \cdot x & \text{otherwise} + \end{cases} +$$ +"# +)] +#[cfg_attr( + not(doc), + doc = "`f(x) =`\n- `x for x >= 0`\n- `negative_slope * x if x < 0`" +)] +pub fn leaky_relu(tensor: Tensor, negative_slope: f64) -> Tensor { + Tensor::new(leaky_relu_impl(tensor.primitive, negative_slope)) +} + +/// Applies the Gaussian Error Linear Units function as described in the paper +/// [Gaussian Error Linear Units (GELUs)](https://arxiv.org/pdf/1606.08415v3.pdf). +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{GELU}(x) += x \cdot \Phi(x) += x \cdot \frac{1}{2}\left(1 + \text{erf}\left(\frac{x}{\sqrt{2}}\right)\right) +$$ + +where $\Phi(x)$ is the cumulative distribution function for the Gaussian distribution. +"# +)] +#[cfg_attr( + not(doc), + doc = r#" +`GELU(x) = x * Φ(x) = x * 1/2 * (1 + erf(x / sqrt(2)))` + +where `Φ(x)` is the cumulative distribution function for the Gaussian distribution. +"# +)] +pub fn gelu(tensor: Tensor) -> Tensor { + Tensor::new(gelu_impl(tensor.primitive)) +} + +/// Applies the tanh-based approximate GELU function element-wise. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{GELU\_approx}(x) += \frac{x}{2}\left(1 + \tanh\left(\sqrt{\frac{2}{\pi}}\left(x + 0.044715\,x^3\right)\right)\right) +$$ +"# +)] +#[cfg_attr( + not(doc), + doc = "`GELU_approx(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))`" +)] +pub fn gelu_approximate(tensor: Tensor) -> Tensor { + /// sqrt(2/π) precomputed as FRAC_2_SQRT_PI * FRAC_1_SQRT_2 + const SQRT_2_OVER_PI: f64 = + core::f64::consts::FRAC_2_SQRT_PI * core::f64::consts::FRAC_1_SQRT_2; + + let x = tensor; + let inner = x.clone() + x.clone().powf_scalar(3.0) * 0.044715; + let inner = inner * SQRT_2_OVER_PI; + (x.clone() * (inner.tanh() + 1)) * 0.5 +} + +/// Applies Parametric ReLu activation function as described in the paper +/// [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://arxiv.org/pdf/1502.01852). +/// +/// - The tensor is assumed to be of shape `[batch_size, channels, ...]`. +/// - `alpha` is assumed to be of shape `[channels]` or `[1]`. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{PReLU}\(x\) = \max\(0,x\) + \alpha \cdot \min\(0, x\) +$$ + +or + +$$ +\text{PReLU}(x) = + \begin{cases} + x & \text{if } x \geq 0 \newline + \alpha x & \text{otherwise} + \end{cases} +$$ +"# +)] +#[cfg_attr(not(doc), doc = "`PReLu(x) = max(0,x) + alpha * min(0,x)`")] +pub fn prelu(tensor: Tensor, alpha: Tensor<1>) -> Tensor { + check!(TensorCheck::check_prelu_shape::( + &tensor.shape(), + &alpha.shape() + )); + + let weight = if alpha.dims()[0] == 1 { + // if there is only 1 weight, then reshape it to (1,1,1... D times) so that the rank is D + alpha.reshape([1; D]) + } else { + // D>=2 because the case where D==1 and num_weights >1 is handled by check function + // there is more than 1 weight and rank is more than 2 + let num_weights = alpha.dims()[0]; + let mut s = [1; D]; + s[1] = num_weights; + // reshape the weights to (1, channels,1 ...) + alpha.reshape(s) + }; + + Tensor::new(prelu_impl(tensor.primitive, weight.primitive)) +} + +/// Applies the softmax function on the input tensor along the given dimension. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{softmax}\(x_i\) = \frac{\exp\(x_i\)}{\sum_j \exp\(x_j\)} +$$ +"# +)] +#[cfg_attr(not(doc), doc = "`softmax(x_i) = exp(x_i) / sum_j(exp(x_j))`")] +/// +/// # Arguments +/// - `dim`: the dimension along which Softmax will be computed. +/// +/// # Panics +/// - If `dim` is outside [0, D) +pub fn softmax(tensor: Tensor, dim: usize) -> Tensor { + check!(TensorCheck::dim_ops::("softmax", dim)); + + Tensor::new(softmax_impl(tensor.primitive, dim)) +} + +/// Applies the softmin function on the input tensor along the given dimension. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{softmin}\(x_i\) = \frac{\exp\(-x_i\)}{\sum_j \exp\(-x_j\)} +$$ +"# +)] +#[cfg_attr(not(doc), doc = "`softmin(x_i) = exp(-x_i) / sum_j(exp(-x_j)`")] +/// +/// # Arguments +/// - `dim`: the dimension along which Softmax will be computed. +/// +/// # Panics +/// - If `dim` is outside [0, D) +pub fn softmin(tensor: Tensor, dim: usize) -> Tensor { + check!(TensorCheck::dim_ops::("softmin", dim)); + + Tensor::new(softmin_impl(tensor.primitive, dim)) +} + +/// Applies the SoftPlus function element-wise. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{softplus}\(x\) = \frac{1}{\beta}\log\(1 + \exp\(\beta x\)\) +$$ +"# +)] +#[cfg_attr(not(doc), doc = "`softplus(x_i) = log(1 + exp(beta * x_i)) / beta`")] +/// +/// The SoftPlus function is a smooth approximation of the ReLU function. +pub fn softplus(tensor: Tensor, beta: f64) -> Tensor { + let tensor = (tensor.mul_scalar(beta).exp() + 1).log(); + tensor.div_scalar(beta) +} + +/// Applies the "quiet softmax" function on the input tensor along the given dimension. +/// +/// Also referred to as [`softmax1`](https://www.evanmiller.org/attention-is-off-by-one.html). +/// +/// This function is similar to the softmax function, but it allows for "no selection" when +/// all the outputs are close to zero. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{quiet\\_softmax}\(x_i\) = \frac{\exp\(x_i\)}{1 + \sum_j \exp\(x_j\)} +$$ +"# +)] +#[cfg_attr( + not(doc), + doc = "`quiet_softmax(x_i) = exp(x_i) / [ 1 + sum_j(exp(x_j)) ]`" +)] +/// +/// # Arguments +/// - `dim`: the dimension along which Softmax will be computed. +/// +/// # Panics +/// - If `dim` is outside [0, D) +pub fn quiet_softmax(tensor: Tensor, dim: usize) -> Tensor { + check!(TensorCheck::dim_ops::("softmax", dim)); + + let max_vals = tensor.clone().detach().max_dim(dim); + let exp_x = (tensor - max_vals.clone()).exp(); + let sum_exp = exp_x.clone().sum_dim(dim); + + exp_x.div(sum_exp + max_vals.neg().exp()) +} + +/// Applies the log softmax function on the input tensor along the given dimension. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{log\\_softmax}\(x_i\) += \log\left(\text{softmax}\(x_i\)\right) += \log\left(\frac{\exp\(x_i\)}{\sum_j \exp\(x_j\)}\right) +$$ +"# +)] +#[cfg_attr( + not(doc), + doc = "`log_softmax(x_i) = log(softmax(x_i)) = log(exp(x_i) / sum_j(exp(x_j)))`" +)] +/// +/// # Arguments +/// - `dim`: the dimension along which Softmax will be computed. +/// +/// # Panics +/// - If `dim` is outside [0, D) +pub fn log_softmax(tensor: Tensor, dim: usize) -> Tensor { + check!(TensorCheck::dim_ops::("log softmax", dim)); + + Tensor::new(log_softmax_impl(tensor.primitive, dim)) +} + +/// Applies the sigmoid function element-wise. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{sigmoid}\(x\) += \sigma(x) += \frac{1}{1 + \exp(-x)} +$$ +"# +)] +#[cfg_attr(not(doc), doc = "`sigmoid(x) = 1 / (1 + exp(-x))`")] +pub fn sigmoid(tensor: Tensor) -> Tensor { + Tensor::new(sigmoid_impl(tensor.primitive)) +} + +/// Applies the hard sigmoid function element-wise. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{hard\\_sigmoid}\(x\) = \max(0, \min(1, \alpha \cdot x + \beta)) +$$ +"# +)] +#[cfg_attr(not(doc), doc = "`hard_sigmoid(x) = max(0, min(1, alpha * x + beta))`")] +pub fn hard_sigmoid(tensor: Tensor, alpha: f64, beta: f64) -> Tensor { + Tensor::new(hard_sigmoid_impl(tensor.primitive, alpha, beta)) +} + +/// Applies the log sigmoid function element-wise. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{log\\_sigmoid}\(x\) = \log\left(\frac{1}{1 + \exp(-x)}\right) +$$ +"# +)] +#[cfg_attr(not(doc), doc = "`log_sigmoid(x) = log(1 / (1 + exp(-x)))`")] +pub fn log_sigmoid(tensor: Tensor) -> Tensor { + Tensor::new(log_sigmoid_impl(tensor.primitive)) +} + +/// Applies the SiLU function (also known as the swish function) element-wise. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{SiLU}\(x\) = x \cdot \sigma(x) = \frac{x}{1 + \exp(-x)} +$$ +"# +)] +#[cfg_attr(not(doc), doc = "`SiLU(x) = x * sigmoid(x) = x / (1 + exp(-x))`")] +pub fn silu(tensor: Tensor) -> Tensor { + tensor.clone().mul(sigmoid(tensor)) +} + +/// Applies the hard swish function element-wise. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{hard\_swish}\(x\) = x \cdot \text{hard\_sigmoid}(x) = x \cdot \max(0, \min(1, \frac{x}{6} + 0.5)) +$$ +"# +)] +#[cfg_attr( + not(doc), + doc = "`hard_swish(x) = x * hard_sigmoid(x) = x * max(0, min(1, x/6 + 0.5))`" +)] +pub fn hard_swish(tensor: Tensor) -> Tensor { + tensor.clone().mul(hard_sigmoid(tensor, 1.0 / 6.0, 0.5)) +} + +/// Applies the Mish function as described in the paper in +/// [Mish: A Self Regularized Non-Monotonic Neural Activation Function](https://arxiv.org/abs/1908.08681). +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{Mish}\(x\) += x \cdot \tanh(\text{Softplus}(x)) += \tanh\left(\log\(1 + \exp\(x\)\)\right) +$$ +"# +)] +#[cfg_attr( + not(doc), + doc = "`mish(x) = x * tanh(softplus(x)) = tanh(log(1 + exp(x)))`" +)] +pub fn mish(tensor: Tensor) -> Tensor { + tensor.clone().mul(softplus(tensor, 1.0).tanh()) +} + +/// Applies the tanh function element-wise. +pub fn tanh(tensor: Tensor) -> Tensor { + tensor.tanh() +} + +/// Applies the Exponential Linear Unit function element-wise. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{ELU}\(x\) = + \begin{cases} + x & \text{if } x > 0 \newline + \alpha \cdot (\exp(x) - 1) & \text{if } x \leq 0 + \end{cases} +$$ +"# +)] +#[cfg_attr( + not(doc), + doc = "`f(x) =`\n- `x for x > 0`\n- `alpha * (exp(x) - 1) for x <= 0`" +)] +pub fn elu(tensor: Tensor, alpha: f64) -> Tensor { + let mask = tensor.clone().lower_equal_elem(0); + let scaled = tensor.clone().exp().sub_scalar(1).mul_scalar(alpha); + tensor.mask_where(mask, scaled) +} + +/// Applies the Continuously Differentiable Exponential Linear Unit function element-wise. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{CELU}(x) = + \begin{cases} + x & \text{if } x \geq 0 \newline + \alpha \cdot \left(\exp\left(\frac{x}{\alpha}\right) - 1\right) & \text{otherwise} + \end{cases} +$$ +"# +)] +#[cfg_attr( + not(doc), + doc = "`celu(x) = max(0, x) + min(0, alpha * (exp(x / alpha) - 1))`" +)] +/// +/// See also [CELU](https://pytorch.org/docs/stable/generated/torch.nn.CELU.html) +/// +/// # Arguments +/// - `alpha`: scaling parameter for the negative part. +pub fn celu(tensor: Tensor, alpha: f64) -> Tensor { + let mask = tensor.clone().lower_equal_elem(0); + let scaled = tensor + .clone() + .div_scalar(alpha) + .exp() + .sub_scalar(1) + .mul_scalar(alpha); + tensor.mask_where(mask, scaled) +} + +/// Applies the Scaled Exponential Linear Unit function element-wise +/// as described in the paper [Self-Normalizing Neural Networks](https://arxiv.org/abs/1706.02515). +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{SELU}\(x\) = \gamma \cdot + \begin{cases} + x & \text{if } x > 0 \newline + \alpha \cdot (\exp(x) - 1) & \text{if } x \leq 0 + \end{cases} +$$ + +where $\alpha \approx 1.6733$ and $\gamma \approx 1.0507$. +"# +)] +#[cfg_attr( + not(doc), + doc = "`selu(x) = gamma * x if x > 0, gamma * alpha * (exp(x) - 1) if x <= 0`" +)] +pub fn selu(tensor: Tensor) -> Tensor { + // Constants from the SELU paper / ONNX spec + const ALPHA: f64 = 1.6732632423543772848170429916717_f64; + const GAMMA: f64 = 1.0507009873554804934193349852946_f64; + + let mask = tensor.clone().greater_equal_elem(0.0); + let positive = tensor.clone().mul_scalar(GAMMA); + let negative = tensor.exp().sub_scalar(1.0).mul_scalar(ALPHA * GAMMA); + + negative.mask_where(mask, positive) +} + +/// Applies the thresholded rectified linear unit function element-wise. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{ThresholdedReLU}(x) = + \begin{cases} + x & \text{if } x > \alpha \newline + 0 & \text{otherwise} + \end{cases} +$$ +"# +)] +#[cfg_attr(not(doc), doc = "`f(x) =`\n- `x if x > alpha`\n- `0 otherwise`")] +/// +/// # Arguments +/// - `alpha`: threshold value (default in ONNX is 1.0). +pub fn thresholded_relu(tensor: Tensor, alpha: f64) -> Tensor { + let mask = tensor.clone().lower_equal_elem(alpha); + tensor.mask_fill(mask, 0) +} + +/// Applies the gated linear unit function. +/// +/// GLU(a,b)=a⊗σ(b) where `a` is the first half of the input matrices and `b` is the second half. +/// +/// **Note**: +/// * The size of the input tensor along `dim` must be divisible by 2. +/// +/// ### Arguments +/// * `tensor` - The input tensor. +/// +/// ### Returns +/// * A tensor with the same shape as the input, except the size along `dim` is halved. +pub fn glu(tensor: Tensor, dim: usize) -> Tensor { + // TODO: Handle negative indices with AsIndex for compatibility with Pytorch nn.GLU. + + assert!( + tensor.dims()[dim].is_multiple_of(2), + "Input tensor along dimension {dim} must have an even size. N is divisible by 2." + ); + let new_len = tensor.dims()[dim] / 2; + + let a = tensor.clone().slice_dim(dim, s![0..new_len]); + let b = tensor.slice_dim(dim, s![new_len..new_len * 2]); + + a.mul(sigmoid(b)) +} + +/// Applies the Softsign function element-wise. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{softsign}(x) = \frac{x}{1 + |x|} +$$ +"# +)] +#[cfg_attr(not(doc), doc = "`softsign(x_i) = x_i / (1 + |x_i|)`")] +pub fn softsign(tensor: Tensor) -> Tensor { + tensor.clone().div(tensor.abs() + 1) +} + +/// Applies the HardShrink function element-wise. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{hard\_shrink}(x) = + \begin{cases} + x & \text{if } x > \lambda \newline + x & \text{if } x < -\lambda \newline + 0 & \text{otherwise} + \end{cases} +$$ +"# +)] +#[cfg_attr( + not(doc), + doc = "`hard_shrink(x) = x if x > lambda, x if x < -lambda, 0 otherwise`" +)] +/// # Arguments +/// - `lambda`: the lambda value for the Hard Shrink formulation. Default is 0.5. +pub fn hard_shrink(tensor: Tensor, lambda: f64) -> Tensor { + let mask = tensor.clone().abs().lower_equal_elem(lambda); + tensor.mask_fill(mask, 0) +} + +/// Applies the SoftShrink function element-wise. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{soft\_shrink}(x) = + \begin{cases} + x - \lambda & \text{if } x > \lambda \newline + x + \lambda & \text{if } x < -\lambda \newline + 0 & \text{otherwise} + \end{cases} +$$ +"# +)] +#[cfg_attr( + not(doc), + doc = "`soft_shrink(x) = x - lambda if x > lambda, x + lambda if x < -lambda, 0 otherwise`" +)] +/// # Arguments +/// - `lambda`: the lambda value for the Soft Shrink formulation. Default is 0.5. +pub fn soft_shrink(tensor: Tensor, lambda: f64) -> Tensor { + shrink(tensor, lambda, lambda) +} + +/// Applies the Shrink function element-wise. +/// +#[cfg_attr( + doc, + doc = r#" +$$ +\text{shrink}(x) = + \begin{cases} + x - \text{bias} & \text{if } x > \lambda \newline + x + \text{bias} & \text{if } x < -\lambda \newline + 0 & \text{otherwise} + \end{cases} +$$ +"# +)] +#[cfg_attr( + not(doc), + doc = "`shrink(x) = x - bias if x > lambda, x + bias if x < -lambda, 0 otherwise`" +)] +/// # Arguments +/// - `lambda`: the lambda value for the Shrink formulation. +/// - `bias`: the bias value for the Shrink formulation. +pub fn shrink(tensor: Tensor, lambda: f64, bias: f64) -> Tensor { + let abs_tensor = tensor.clone().abs(); + let sign = tensor.clone().sign(); + let shrunk = tensor.sub(sign.mul_scalar(bias)); + let mask = abs_tensor.lower_equal_elem(lambda); + shrunk.mask_fill(mask, 0) +} + +// ========================================================================= +// Non-generic implementation helpers (outlined from the generic API). +// See the crate-level docs for the rationale behind this pattern. +// ========================================================================= + +fn leaky_relu_impl(p: BridgeTensor, negative_slope: f64) -> BridgeTensor { + BridgeTensor::float(Dispatch::leaky_relu(p.into_float(), negative_slope.into())) +} + +fn gelu_impl(p: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::gelu(p.into_float())) +} + +fn prelu_impl(p: BridgeTensor, weight: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::prelu(p.into_float(), weight.into_float())) +} + +fn softmax_impl(p: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::float(Dispatch::softmax(p.into_float(), dim)) +} + +fn softmin_impl(p: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::float(Dispatch::softmin(p.into_float(), dim)) +} + +fn log_softmax_impl(p: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::float(Dispatch::log_softmax(p.into_float(), dim)) +} + +fn sigmoid_impl(p: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::sigmoid(p.into_float())) +} + +fn hard_sigmoid_impl(p: BridgeTensor, alpha: f64, beta: f64) -> BridgeTensor { + BridgeTensor::float(Dispatch::hard_sigmoid( + p.into_float(), + alpha.into(), + beta.into(), + )) +} + +fn log_sigmoid_impl(p: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::log_sigmoid(p.into_float())) +} diff --git a/crates/burn-tensor/src/tensor/activation/mod.rs b/crates/burn-tensor/src/tensor/activation/mod.rs new file mode 100644 index 0000000..cbcb6ac --- /dev/null +++ b/crates/burn-tensor/src/tensor/activation/mod.rs @@ -0,0 +1,3 @@ +mod base; + +pub use base::*; diff --git a/crates/burn-tensor/src/tensor/api/autodiff.rs b/crates/burn-tensor/src/tensor/api/autodiff.rs new file mode 100644 index 0000000..f393f10 --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/autodiff.rs @@ -0,0 +1,117 @@ +use crate::{Tensor, kind::Autodiff}; + +#[cfg(feature = "autodiff")] +use crate::ops::BridgeTensor; +#[cfg(feature = "autodiff")] +use burn_backend::AutodiffBackend; +#[cfg(feature = "autodiff")] +use burn_dispatch::Dispatch; + +#[cfg(feature = "autodiff")] +type AutodiffGradients = ::Gradients; + +// Aligned, type-erased storage for `AutodiffGradients`. See `crate::macros` +// for why this indirection exists. +#[cfg(feature = "autodiff")] +burn_std::obfuscate!( + type: AutodiffGradients, + module: gradients_opaque, + derives: [Send] +); + +/// Gradients container used during the backward pass. +#[cfg(feature = "autodiff")] +pub struct Gradients { + blob: gradients_opaque::Opaque, +} + +#[cfg(feature = "autodiff")] +impl Gradients { + /// Crate-internal constructor wrapping the dispatch-level gradients. + pub(crate) fn from_inner(inner: AutodiffGradients) -> Self { + Self { + blob: gradients_opaque::Opaque::new(inner), + } + } + + /// Crate-internal borrow of the underlying gradients container. + pub(crate) fn as_inner(&self) -> &AutodiffGradients { + self.blob.as_ref() + } + + /// Crate-internal mutable borrow of the underlying gradients container. + pub(crate) fn as_inner_mut(&mut self) -> &mut AutodiffGradients { + self.blob.as_mut() + } +} + +#[cfg(feature = "autodiff")] +impl Tensor { + /// Backward pass of the tensor. + pub fn backward(&self) -> Gradients { + backward_impl(&self.primitive) + } + + /// Get the gradients of a tensor if it exist. + /// + /// Returns a new reference to the same tensor. Therefore the same grad tensor can + /// be accessed multiple times. If you only need to get the gradients one time, + /// consider using [grad_remove](Tensor::grad_remove) for better performance. + pub fn grad(&self, grads: &Gradients) -> Option> { + grad_impl(&self.primitive, grads).map(Tensor::new) + } + + /// Remove the grad tensor from the [grads](AutodiffBackend::Gradients) struct returning the result. + pub fn grad_remove(&self, grads: &mut Gradients) -> Option> { + grad_remove_impl(&self.primitive, grads).map(Tensor::new) + } + + /// Replace the grad tensor from the [grads](AutodiffBackend::Gradients) struct with the provided + /// gradient. + pub fn grad_replace(&self, grads: &mut Gradients, grad: Tensor) { + grad_replace_impl(&self.primitive, grads, grad.primitive) + } +} + +#[cfg(feature = "autodiff")] +fn backward_impl(p: &BridgeTensor) -> Gradients { + Gradients::from_inner(Dispatch::backward(p.clone().into_float())) +} + +#[cfg(feature = "autodiff")] +fn grad_impl(p: &BridgeTensor, grads: &Gradients) -> Option { + Dispatch::grad(p.as_float(), grads.as_inner()).map(BridgeTensor::float) +} + +#[cfg(feature = "autodiff")] +fn grad_remove_impl(p: &BridgeTensor, grads: &mut Gradients) -> Option { + Dispatch::grad_remove(p.as_float(), grads.as_inner_mut()).map(BridgeTensor::float) +} + +#[cfg(feature = "autodiff")] +fn grad_replace_impl(p: &BridgeTensor, grads: &mut Gradients, grad: BridgeTensor) { + Dispatch::grad_replace(p.as_float(), grads.as_inner_mut(), grad.into_float()) +} + +impl Tensor { + /// Returns the inner tensor without the autodiff information. + pub fn inner(self) -> Tensor { + Tensor::new(K::inner(self.primitive)) + } + + /// Convert a tensor to the autodiff backend. + /// + /// # Arguments + /// + /// * `inner` - The tensor to convert. + /// + /// # Returns + /// + /// The tensor converted to the autodiff backend. + pub fn from_inner(inner: Tensor) -> Self { + Self::new(K::from_inner(inner.primitive)) + } +} + +// TODO: a lot of the `tensor.inner` / `Tensor::from_inner(...)` are actually scoped to perform some operations +// so it might be cleaner and easier to manage the device etc. if we provide a method to scope the autodiff? diff --git a/crates/burn-tensor/src/tensor/api/base.rs b/crates/burn-tensor/src/tensor/api/base.rs new file mode 100644 index 0000000..fa124a2 --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/base.rs @@ -0,0 +1,3454 @@ +#![allow(clippy::single_range_in_vec_init)] +use crate::check::unwrap_shape_reshape; +use crate::kind::Basic; +use crate::ops::BridgeTensor; + +use burn_backend::Scalar; + +use alloc::vec::Vec; + +use alloc::format; +use alloc::string::String; +use alloc::vec; + +use burn_std::ExecutionError; +use burn_std::{SliceOps, stub::RwLock}; +use core::iter::repeat; +use core::marker::PhantomData; +use core::{fmt::Debug, ops::Range}; +use serde::{Deserialize, Deserializer}; + +use crate::{AsIndex, Device, Slice, SliceArg, wrap_index}; +use crate::{Bool, ElementConversion, Float, Int, Shape, TensorData, check}; +use crate::{DType, Element}; +use crate::{IndexingUpdateOp, TensorCreationOptions}; +use crate::{cast::ToElement, check::TensorCheck}; +use serde::{Serialize, Serializer}; + +/// A tensor with a given backend, shape and data type. +/// +/// # Indexing +/// Indexing a tensor can be done using [`slice`](Tensor::slice) for all tensor types +/// or [`select`](Tensor::select) for numeric types. +/// +/// ## Example +/// +/// ```rust +/// use burn_tensor::Tensor; +/// use burn_tensor::Int; +/// +/// fn example() { +/// let device = Default::default(); +/// +/// let tensor = Tensor::<2>::from_data( +/// [ +/// [3.0, 4.9, 2.0], +/// [2.0, 1.9, 3.0], +/// [6.0, 1.5, 7.0], +/// [3.0, 4.9, 9.0], +/// ], +/// &device, +/// ); +/// +/// // Slice the tensor to get the second and third rows: +/// // [[2.0, 1.9, 3.0], [6.0, 1.5, 7.0]] +/// // The resulting tensor will have dimensions [2, 3]. +/// let slice = tensor.clone().slice([1..3]); +/// println!("{slice}"); +/// +/// // Slice the tensor to get the first two rows and the first 2 columns: +/// // [[3.0, 4.9], [2.0, 1.9]] +/// // The resulting tensor will have dimensions [2, 2]. +/// let slice = tensor.clone().slice([0..2, 0..2]); +/// println!("{slice}"); +/// +/// // Index the tensor along the dimension 1 to get the elements 0 and 2: +/// // [[3.0, 2.0], [2.0, 3.0], [6.0, 7.0], [3.0, 9.0]] +/// // The resulting tensor will have dimensions [4, 2] +/// let indices = Tensor::<1, Int>::from_data([0, 2], &device); +/// let indexed = tensor.select(1, indices); +/// println!("{indexed}"); +/// } +/// ``` +#[derive(new, Clone, Debug)] +pub struct Tensor +where + K: Basic, +{ + pub(crate) primitive: BridgeTensor, + _kind: PhantomData, +} + +impl From for Tensor +where + K: Basic, + T: Into, +{ + fn from(value: T) -> Self { + Tensor::from_data(value.into(), &Default::default()) + } +} + +impl Tensor +where + K: Basic, +{ + /// Executes an operation on the tensor and modifies its value. + /// + /// # Notes + /// + /// This won't necessarily reuse the same tensor data/buffer, but it should if there is + /// no other reference pointing to the same tensor. + /// + /// Wrapping operations with inplace is not an optimization, it's mainly there if you + /// want to mutate a tensor by using owned operations. A plausible usage would be to + /// update the weights of a mutable model reference. + pub fn inplace Self>(&mut self, func: F) { + let mut tensor_owned = Tensor::empty([0; D], &self.device()); + core::mem::swap(&mut tensor_owned, self); + + let mut tensor_new = func(tensor_owned); + core::mem::swap(&mut tensor_new, self); + } + + /// Returns the number of dimensions of the tensor. + pub fn rank(&self) -> usize { + self.primitive.rank() + } + + /// Returns the tensor primitive data type. + /// + /// # Note + /// Some element types are encoded in different primitive types depending on the backend + /// (e.g., bool could be encoded as `u8` or `u32`). + pub fn dtype(&self) -> DType { + self.primitive.dtype() + } + + /// Whether this tensor's buffer can be mutated in place — i.e. this handle + /// uniquely owns the allocation, so an in-place op writes it directly + /// instead of copying first (see `TensorMetadata::can_mut`). + /// + /// Backends that track buffer ownership (cubecl, fusion, tch) answer + /// precisely from the handle reference count; others conservatively return + /// `false` — they may alias the buffer, so an in-place write can't be + /// assumed safe. Useful to assert a hot-path op (e.g. a KV-cache + /// `slice_assign`) stays in place rather than silently copying. + pub fn can_mut(&self) -> bool { + self.primitive.can_mut() + } + + /// Create an empty tensor of the given shape. + /// + /// # Arguments + /// + /// - `shape`: The shape of the tensor. + /// - `device`: The device where the tensor will be created. + /// + /// # Example + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create an empty tensor with dimensions [2, 3, 4]. + /// let tensor = Tensor::<3>::empty([2, 3, 4], &device); + /// } + /// ``` + pub fn empty>(shape: S, options: impl Into) -> Self { + let opt = options.into(); + let shape = shape.into(); + let dtype = opt.resolve_dtype::(); + check!(TensorCheck::creation_ops::("Empty", &shape)); + Self::new(K::empty(shape, &opt.device, dtype)) + } + + /// Create a tensor of the given shape where each element is zero. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::zeros(Shape::new([2, 3]), &device); + /// println!("{tensor}"); + /// // [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] + /// } + /// ``` + pub fn zeros>(shape: S, options: impl Into) -> Self { + let opt = options.into(); + let shape = shape.into(); + let dtype = opt.resolve_dtype::(); + check!(TensorCheck::creation_ops::("Zeros", &shape)); + Self::new(K::zeros(shape, &opt.device, dtype)) + } + + /// Returns a new tensor with the same shape, dtype, and device as the current tensor filled with zeros. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.zeros_like(); + /// println!("{tensor}"); + /// // [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] + /// } + /// ``` + pub fn zeros_like(&self) -> Self { + Self::new(K::zeros(self.shape(), &self.device(), self.dtype())) + } + + /// Create a tensor of the given shape where each element is one. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::ones(Shape::new([2, 3]), &device); + /// println!("{tensor}"); + /// // [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]] + /// } + /// ``` + pub fn ones>(shape: S, options: impl Into) -> Self { + let opt = options.into(); + let shape = shape.into(); + let dtype = opt.resolve_dtype::(); + check!(TensorCheck::creation_ops::("Ones", &shape)); + Self::new(K::ones(shape, &opt.device, dtype)) + } + + /// Returns a new tensor with the same shape, dtype, and device as the current tensor filled with ones. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.ones_like(); + /// println!("{tensor}"); + /// // [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]] + /// } + /// ``` + pub fn ones_like(&self) -> Self { + Self::new(K::ones(self.shape(), &self.device(), self.dtype())) + } + + /// Create a tensor of the given shape where each element is equal to the provided value. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::full(Shape::new([2, 3]), 5.0, &device); + /// println!("{tensor}"); + /// // [[5.0, 5.0, 5.0], [5.0, 5.0, 5.0]] + /// } + /// ``` + pub fn full, E: ElementConversion>( + shape: S, + fill_value: E, + options: impl Into, + ) -> Self { + let opt = options.into(); + let shape = shape.into(); + let dtype = opt.resolve_dtype::(); + check!(TensorCheck::creation_ops::("Full", &shape)); + Self::new(K::full( + shape, + Scalar::new(fill_value, &dtype), + &opt.device, + dtype, + )) + } + + /// Returns a new tensor with the same shape, dtype, and device as the current tensor, + /// filled with the provided value. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.full_like(5.0); + /// println!("{tensor}"); + /// // [[5.0, 5.0, 5.0], [5.0, 5.0, 5.0]] + /// } + /// ``` + pub fn full_like(&self, fill_value: E) -> Self { + let dtype = self.dtype(); + Self::new(K::full( + self.shape(), + Scalar::new(fill_value, &dtype), + &self.device(), + dtype, + )) + } + + /// Returns the dimensions of the current tensor. + /// + /// # Example + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<3>::ones([2, 3, 4], &device); + /// let dims = tensor.dims(); // [2, 3, 4] + /// println!("{dims:?}"); + /// } + /// ``` + pub fn dims(&self) -> [usize; D] { + Self::shape(self).dims() + } + + /// Returns the shape of the current tensor. + /// + /// # Example + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<3>::ones([2, 3, 4], &device); + /// // Shape { dims: [2, 3, 4] } + /// let shape = tensor.shape(); + /// } + /// ``` + pub fn shape(&self) -> Shape { + self.primitive.shape() + } + + /// Reshape the tensor to have the given shape. + /// + /// The tensor has the same data and number of elements as the input. + /// + /// A `-1` in the shape is used to infer the remaining dimensions, e.g.: `[2, -1]` + /// will reshape the tensor with [2, 3, 4] dimensions to [2, 12]. + /// + /// A `0` in the shape instructs to keep the current dimension from the original tensor, + /// e.g.: `[2, 0, 4]` will reshape the tensor with [2, 3, 4] dimensions to [2, 3, 4]. + /// This is useful when reshaping tensors with unknown dimensions and combining with `-1` + /// to infer the remaining dimensions, e.g. `[0, -1]` will reshape the tensor + /// with [1, 3, 4] dimensions to [1, 12]. + /// + /// # Arguments + /// - `shape`: The new shape of the tensor. + /// + /// # Panics + /// - If the tensor contains more than one `-1` in the shape. + /// - If the tensor contains values that are not positive (other than -1). + /// - If the shape does not match the number of elements of the original shape. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a tensor with dimensions [2, 3, 4] + /// let tensor = Tensor::<3>::ones([2, 3, 4], &device); + /// // Reshape it to [2, 12], where 12 is inferred from the number of elements. + /// let reshaped = tensor.reshape([2, -1]); + /// println!("{reshaped}"); + /// } + /// ``` + pub fn reshape>(self, shape: S) -> Tensor { + // Convert reshape args to shape + let shape = shape.into_shape::(self.shape()); + Tensor::new(K::reshape(self.primitive, shape)) + } + + /// Transpose the tensor. + /// + /// For a 2D tensor, this is the standard matrix transpose. For `D > 2`, the transpose is + /// applied on the last two dimensions. For example, the transpose of a tensor with shape + /// `[1, 2, 3, 4]` will have shape `[1, 2, 4, 3]`. + /// + /// See also [`permute`](Tensor::permute). + /// + /// # Arguments + /// + /// * `tensor` - The tensor to transpose. + /// + /// # Returns + /// + /// The transposed tensor. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 2D tensor of shape [2, 3] + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// + /// // Transpose the tensor: + /// // [[1.0, 5.0], [-2.0, 9.0], [3.0, 6.0]] + /// // The resulting tensor will have dimensions [3, 2]. + /// let transposed = tensor.transpose(); + /// println!("{transposed}"); + /// } + /// ``` + pub fn transpose(self) -> Tensor { + Tensor::new(K::transpose(self.primitive)) + } + + /// Alias for `transpose`. + #[inline(always)] + pub fn t(self) -> Tensor { + self.transpose() + } + + /// Swaps two dimensions of a tensor. + /// + /// This is a no-op when `dim1 == dim2`, assuming both are within bounds. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to swap the dimensions of. + /// * `dim1` - The first dimension to swap, supports negative indexing. + /// * `dim2` - The second dimension to swap, supports negative indexing. + /// + /// # Returns + /// + /// The tensor with the dimensions swapped. + /// + /// # Panics + /// + /// When dimensions are out of bounds. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 2D tensor of shape [2, 3] + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// + /// // Swap the dimensions 0 and -1 (equivalent to `tensor.transpose()`): + /// // [[1.0, 5.0], [-2.0, 9.0], [3.0, 6.0]] + /// // The resulting tensor will have dimensions [3, 2]. + /// let swapped = tensor.swap_dims(0, -1); + /// println!("{swapped}"); + /// } + /// ``` + pub fn swap_dims(self, dim1: Dim1, dim2: Dim2) -> Tensor + where + Dim1: AsIndex, + Dim2: AsIndex, + { + let dim1 = dim1.expect_dim_index(D); + let dim2 = dim2.expect_dim_index(D); + check!(TensorCheck::swap_dims::(dim1, dim2)); + if dim1 == dim2 { + self + } else { + Tensor::new(K::swap_dims(self.primitive, dim1, dim2)) + } + } + + /// Permute the dimensions of the tensor. + /// + /// This is a no-op when the resolved `axes` match the current order. + /// + /// # Arguments + /// + /// * `axes` - The new order of the dimensions. The length of the axes + /// must be equal to the number of dimensions of the tensor. + /// The values must be unique and in the range of the number of dimensions. + /// The values can be negative, in which case they are used as an offset from the end. + /// + /// # Returns + /// + /// The tensor with the dimensions permuted. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 2D tensor of shape [3, 2] + /// let tensor = Tensor::<2>::from_data([[1.0, 5.0], [-2.0, 9.0], [3.0, 6.0]], &device); + /// + /// // Permute the dimensions 1 and 0: + /// // [[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]] + /// // The resulting tensor will have dimensions [3, 2]. + /// let permuted = tensor.permute([1, 0]); + /// println!("{permuted}"); + /// } + /// ``` + pub fn permute(self, axes: [Dim; D]) -> Tensor + where + Dim: AsIndex, + { + let mut no_op = true; + let mut fixed_axes = [0; D]; + for (i, axis) in axes.into_iter().enumerate() { + let dim = axis.expect_dim_index(D); + no_op &= dim == i; + fixed_axes[i] = dim; + } + + if no_op { + self + } else { + check!(TensorCheck::permute(fixed_axes)); + Tensor::new(K::permute(self.primitive, &fixed_axes)) + } + } + + /// Moves the dimension(s) of input at the position(s) in source to the position(s) in destination. + /// + /// Other dimensions of input that are not explicitly moved remain in their original order and appear + /// at the positions not specified in destination. + /// + /// # Arguments + /// + /// * `src` - The dimension(s) to move. The values must be unique and in the range of the number of dimensions. + /// The values can be negative, in which case they are used as an offset from the end. + /// + /// * `dst` - Destination positions for each of the original dims. These must also be unique. + /// + /// # Panics + /// + /// - If the source and destination dimensions are not of the same length. + /// - If the source and destination vectors contain duplicate values. + /// - If the source and destination vectors contain values that are out of bounds. + /// + /// # Returns + /// + /// The tensor with the dimensions moved. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 3D tensor of shape [3, 2, 1] + /// let tensor = Tensor::<3>::from_data([[[1.0], [5.0]], [[-2.0], [9.0]], [[3.0], [6.0]]], &device); + /// + /// // Move the dimensions 0 and 1: + /// // [[[1.0], [-2.0], [3.0]], [[5.0], [9.0], [6.0]]] + /// // The resulting tensor will have dimensions [2, 3, 1]. + /// let moved = tensor.movedim(1, 0); + /// println!("{moved}"); + /// } + /// ``` + /// + /// # Note + /// + /// This is a syntactic sugar for `permute`. It is used widely enough, so we define a separate Op + /// for it + pub fn movedim(self, src: S1, dst: S2) -> Tensor { + let source_dims = src.into_dim_vec::(); + let destination_dims = dst.into_dim_vec::(); + + check!(TensorCheck::movedim_args_length( + &source_dims, + &destination_dims + )); + + let mut m = [-1; D]; + for (&d, &s) in destination_dims.iter().zip(source_dims.iter()) { + m[d] = s as isize; + } + let mut axes: [isize; D] = [0; D]; + let mut source_i = 0; + for (dest_i, item) in axes.iter_mut().enumerate().take(D) { + *item = if m[dest_i] != -1 { + m[dest_i] + } else { + while source_dims.contains(&source_i) { + source_i += 1; + } + let result = source_i as isize; + source_i += 1; + result + }; + } + + self.permute(axes) + } + + /// Reverse the order of elements in the tensor along the given dimensions. + /// + /// # Arguments + /// + /// * `axes` - The dimensions to reverse. The values must be unique and in the range of the number of dimensions. + /// The values can be negative, in which case they are used as an offset from the end. + /// + /// # Returns + /// + /// The tensor with the axes flipped. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 2D tensor with dimensions [4, 3] + /// let tensor = Tensor::<2>::from_data( + /// [ + /// [3.0, 4.9, 2.0], + /// [2.0, 1.9, 3.0], + /// [4.0, 5.9, 8.0], + /// [1.4, 5.8, 6.0], + /// ], + /// &device, + /// ); + /// + /// // Flip the elements in dimensions 0 and 1: + /// // [[6.0, 5.8, 1.4], + /// // [8.0, 5.9, 4.0], + /// // [3.0, 1.9, 2.0], + /// // [2.0, 4.9, 3.0]] + /// // The resulting tensor will have dimensions [4, 3]. + /// let flipped = tensor.flip([0, 1]); + /// println!("{flipped}"); + /// } + /// ``` + pub fn flip(self, axes: [isize; N]) -> Tensor { + // Convert the axes to usize and handle negative values without using vector + let mut transformed_axes: [usize; N] = [0; N]; + for (i, &x) in axes.iter().enumerate() { + transformed_axes[i] = if x < 0 { + (D as isize + x) as usize + } else { + x as usize + }; + } + + // Check if the axes are valid + check!(TensorCheck::flip(D, &transformed_axes)); + + Tensor::new(K::flip(self.primitive, &transformed_axes)) + } + + /// Flatten the tensor along a given range of dimensions. + /// + /// This function collapses the specified range of dimensions into a single dimension, + /// effectively flattening the tensor in that range. + /// + /// # Arguments + /// + /// - `start_dim`: The starting dimension of the range to be flattened, + /// supports negative indexing. + /// - `end_dim`: The ending dimension of the range to be flattened (inclusive), + /// supports negative indexing. + /// + /// # Type Parameters + /// + /// - `D2`: The resulting number of dimensions in the flattened tensor. + /// + /// # Returns + /// + /// A new `Tensor` instance with the specified range of dimensions flattened. + /// + /// # Example + /// + /// ```rust + /// + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 3D tensor with dimensions [2, 3, 4] + /// let tensor = Tensor::<3>::ones(Shape::new([2, 3, 4]), &device); + /// + /// // Flatten the tensor from dimensions 1 to 2 (inclusive). + /// // The resulting tensor will have dimensions [2, 12] + /// let flattened: Tensor<2> = tensor.flatten(1, 2); + /// println!("{flattened}"); + /// } + /// ``` + pub fn flatten( + self, + start_dim: impl AsIndex, + end_dim: impl AsIndex, + ) -> Tensor { + let start_dim = start_dim.expect_dim_index(D); + let end_dim = end_dim.expect_dim_index(D); + check!(TensorCheck::flatten::(start_dim, end_dim)); + let new_shape = self.shape().flatten_dims(start_dim, end_dim); + + Tensor::new(K::reshape(self.primitive, new_shape)) + } + + /// Squeeze the tensor along all dimensions, removing dimensions + /// of size one, and effectively reducing the rank of the tensor. + /// + /// # Type Parameters + /// + /// - `D2`: The resulting number of dimensions in the squeezed tensor. + /// + /// # Returns + /// + /// A new `Tensor` instance with the specified dimension removed. + /// + /// # Example + /// + /// ```rust + /// + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 4D tensor with dimensions [1, 3, 1, 3] + /// let tensor = Tensor::<4>::from_data( + /// [[[[3.0, 4.9, 2.0]], [[2.0, 1.9, 3.0]], [[4.0, 5.9, 8.0]]]], + /// &device, + /// ); + /// + /// // Squeeze the tensor dimensions. + /// // The resulting tensor will have dimensions [3, 3]. + /// let squeezed = tensor.squeeze::<2>(); + /// println!("{squeezed}"); + /// } + /// ``` + pub fn squeeze(self) -> Tensor { + let new_dims = self + .shape() + .iter() + .filter_map(|&dim| if dim == 1 { None } else { Some(dim) }) + .collect::>(); + check!(TensorCheck::squeeze_dims_len::(new_dims.len())); + + Tensor::new(K::reshape(self.primitive, new_dims.into())) + } + + /// Squeeze the tensor along the given dimension, removing the specified dimension + /// of size one, and effectively reducing the rank of the tensor by one. + /// + /// # Arguments + /// + /// - `dim`: The dimension to be squeezed. + /// + /// # Type Parameters + /// + /// - `D2`: The resulting number of dimensions in the squeezed tensor. + /// + /// # Panics + /// + /// If the size in the squeezed dimension is not 1. + /// + /// # Returns + /// + /// A new `Tensor` instance with the specified dimension removed. + /// + /// # Example + /// + /// ```rust + /// + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 3D tensor with dimensions [3, 1, 3] + /// let tensor = Tensor::<3>::from_data( + /// [[[3.0, 4.9, 2.0]], [[2.0, 1.9, 3.0]], [[4.0, 5.9, 8.0]]], + /// &device, + /// ); + /// + /// // Squeeze the dimension 1. + /// // The resulting tensor will have dimensions [3, 3]. + /// let squeezed = tensor.squeeze_dim::<2>(1); + /// println!("{squeezed}"); + /// } + /// ``` + pub fn squeeze_dim(self, dim: usize) -> Tensor { + check!(TensorCheck::squeeze::(dim, &self.shape())); + + let current_dims = self.shape(); + let mut new_dims: [usize; D2] = [0; D2]; + + new_dims[..dim].copy_from_slice(¤t_dims[..dim]); + new_dims[dim..].copy_from_slice(¤t_dims[dim + 1..]); + + check!(TensorCheck::squeeze_dims_len::(new_dims.len())); + Tensor::new(K::reshape(self.primitive, new_dims.into())) + } + + /// Removes specified dimensions of size 1 from a tensor's shape. This function takes a tensor and + /// an array of dimensions (`dims`) to be squeezed. If `dims` is provided, only the dimensions + /// specified in this array will be removed. Each dimension in `dims` should correspond to a size of 1 + /// in the tensor; otherwise, the dimension will not be squeezed. If `dims` is empty, all single-dimensional entries + /// in the tensor will be removed. If entries in `dims` are negative, then dimensions will be counted + /// from the back. + /// + /// # Arguments + /// + /// - `dims`: The dimension(s) to be squeezed. + /// + /// # Type Parameters + /// + /// - `D2`: The resulting number of dimensions in the squeezed tensor. + /// + /// # Returns + /// + /// A new `Tensor` instance with the specified dimensions removed. + /// + /// # Example + /// + /// ```rust + /// + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 4D tensor with dimensions [2, 1, 4, 1] + /// let tensor = Tensor::<4>::ones(Shape::new([2, 1, 4, 1]), &device); + /// + /// // Squeeze the dimensions 1 and 3. + /// // The resulting tensor will have dimensions [2, 4]. + /// let squeezed: Tensor<2> = tensor.squeeze_dims(&[1, 3]); + /// println!("{squeezed}"); + /// } + /// ``` + pub fn squeeze_dims(self, dims: &[isize]) -> Tensor { + let current_dims = self.shape(); + let mut dim_indices: Vec; + + // Check if dims is empty, if yes then assign dim_indices all single-dimensional entries + if dims.is_empty() { + dim_indices = current_dims + .iter() + .enumerate() + .filter_map(|(index, &dim)| if dim == 1 { Some(index) } else { None }) + .collect(); + } else { + // If negative dims, count from the back + dim_indices = dims + .iter() + .map(|&d| { + if d < 0 { + (current_dims.len() as isize + d) as usize + } else { + d as usize + } + }) + .collect(); + } + + // Sort indices and remove duplicates + dim_indices.sort_unstable(); + dim_indices.dedup(); + + // Make sure squeeze_dims doesn't result in a tensor with < 1 dimensions + check!(TensorCheck::squeeze_dims_input::( + &dim_indices, + ¤t_dims + )); + + // Calculate new dimensions + let mut new_dims = Vec::new(); + for (index, &dim_size) in current_dims.iter().enumerate() { + // Exclude the dimension if it's explicitly marked for squeezing + if dim_indices.contains(&index) { + check!(TensorCheck::squeeze::(index, ¤t_dims)); + continue; + } + new_dims.push(dim_size); + } + + // Check that after squeezing, we still respect the D2 size + check!(TensorCheck::squeeze_dims_len::(new_dims.len())); + + Tensor::new(K::reshape(self.primitive, new_dims.into())) + } + + /// Unsqueeze the current tensor. Create new leading dimensions to fit the given size. + /// + /// # Type Parameters + /// + /// - `D2`: The resulting number of dimensions in the unsqueezed tensor. + /// + /// # Panics + /// + /// If the output size `D2` is smaller than the current number of dimensions. + /// + /// # Returns + /// + /// A new `Tensor` instance with the specified dimensions added. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 2D tensor with dimensions [3, 3] + /// let tensor = Tensor::<2>::ones(Shape::new([3, 3]), &device); + /// // Unsqueeze the tensor up to 4 dimensions. + /// // The resulting tensor will have dimensions [1, 1, 3, 3]. + /// let unsqueezed = tensor.unsqueeze::<4>(); + /// println!("{unsqueezed}"); + /// } + /// ``` + pub fn unsqueeze(self) -> Tensor { + check!(TensorCheck::unsqueeze::()); + + let mut dims = [1; D2]; + let num_ones = D2 - D; + let shape = self.shape(); + + dims[num_ones..(D + num_ones)].copy_from_slice(&shape[..D]); + + let shape = Shape::new(dims); + self.reshape(shape) + } + + /// Creates a new tensor with a dimension of size one inserted at the specified position. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 2D tensor with dimensions [3, 3] + /// let tensor = Tensor::<2>::ones(Shape::new([3, 3]), &device); + /// // Unsqueeze the dimension 1. + /// // The resulting tensor will have dimensions [3, 1, 3]. + /// let unsqueezed: Tensor<3> = tensor.unsqueeze_dim(1); + /// println!("{unsqueezed}"); + /// } + /// ``` + pub fn unsqueeze_dim(self, dim: usize) -> Tensor { + check!(TensorCheck::unsqueeze_dim::(dim)); + + let mut dims = [1; D2]; + let shape = self.shape(); + + dims[0..dim].copy_from_slice(&shape[0..dim]); + + if dim < D { + dims[dim] = 1; + dims[(dim + 1)..].copy_from_slice(&shape[dim..]); + } else { + dims[dim] = 1; + } + + let shape = Shape::new(dims); + self.reshape(shape) + } + + /// Creates a new tensor with added dimensions of size one inserted at the specified indices. + /// The indices can be negative, in which case they are counted from the last to the first dimension. + /// the axes can contain duplicates, in which case the number of dimensions inserted at the index + /// is the number of duplicates. + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 3D tensor with dimensions [3, 4, 5] + /// let tensor = Tensor::<3>::ones(Shape::new([3, 4, 5]), &device); + /// // Unsqueeze the leading dimension (0) once and the trailing dimension (-1) twice. + /// // The resulting tensor will have dimensions [1, 3, 4, 5, 1, 1]. + /// let unsqueezed: Tensor<6> = tensor.unsqueeze_dims(&[0, -1, -1]); + /// println!("{unsqueezed}"); + /// } + /// ``` + pub fn unsqueeze_dims(self, axes: &[impl AsIndex]) -> Tensor { + let mut new_dims = [1; D2]; + let old_dims = self.shape(); + //for checking if the dimension is in the acceptable range + + //part 1: convert the negative indices to positive + let mut neg_offset = D2; + let mut dim_indices = axes + .iter() + .map(|d| { + let d = d.as_index(); + // check if the dimension is in the acceptable range + check!(TensorCheck::unsqueeze_dims::<{ D2 }>(d)); + (if d < 0 { + neg_offset -= 1; // handle multiple negative indices (decrease dim value in reverse) + d + neg_offset as isize + 1 + } else { + d + }) as usize + }) + .collect::>(); + + //sort the indices + dim_indices.sort_unstable(); + + // Per the documented semantics, duplicate axes mean "insert N dims at that index". + // After sorting, N insertions at position `i` logically occupy positions + // `i, i+1, ..., i+N-1` in the output, so bump each duplicate to the next slot. + // Example: sorted `[0, 0, 3]` becomes `[0, 1, 3]`, matching the intent of + // "two 1s starting at index 0, plus one 1 at index 3". + for i in 1..dim_indices.len() { + if dim_indices[i] <= dim_indices[i - 1] { + dim_indices[i] = dim_indices[i - 1] + 1; + } + } + + // Re-validate after normalization: bumping duplicates forward can push the + // last index past `D2 - 1` (e.g. `[2, 2]` targeting rank 3 normalizes to + // `[2, 3]`). The per-axis check above only runs on pre-normalization values, + // so we re-check here to surface a clear `TensorCheck` error instead of + // letting the copy loop panic on an out-of-bounds `old_dims` read. + for &dim_index in &dim_indices { + check!(TensorCheck::unsqueeze_dims::<{ D2 }>(dim_index as isize)); + } + + // Loop over the entries/indices of the `new_dims` array. + // When the current entry should be 1 from the unsqueeze operation, simply increment + // the index for `dims_indices` to account for "adding" its entry to `new_dims`. + // Otherwise, the dim from the current entry of `old_dims` should be copied to `new_dims`. + let mut dim_indices_curr_idx = 0; + let mut old_dims_curr_idx = 0; + for new_dims_curr_idx in 0..D2 { + // If all indices in `dim_indices` have been processed, then + // simply copy all the remaining dims from `old_dims` to `new_dims` + if dim_indices_curr_idx == dim_indices.len() { + new_dims[new_dims_curr_idx..].copy_from_slice(&old_dims[old_dims_curr_idx..]); + break; + } + + if new_dims_curr_idx == dim_indices[dim_indices_curr_idx] { + dim_indices_curr_idx += 1; + } else { + new_dims[new_dims_curr_idx] = old_dims[old_dims_curr_idx]; + old_dims_curr_idx += 1; + } + } + + //lastly, create the shape and reshape + let shape = Shape::new(new_dims); + self.reshape(shape) + } + + /// Roll operation along a specific dimension; wrapping around the elements. + /// + /// ## Parameters + /// + /// - `shift`: The roll extent; supports negative values and wraps around. + /// - `dim`: The dimension to roll; supports negative indexing. + /// + /// ## Returns + /// + /// A new tensor with the specified dimension rolled by the given shift amount. + pub fn roll_dim(self, shift: Shift, dim: Dim) -> Self + where + Shift: AsIndex, + Dim: AsIndex, + { + let dim = dim.expect_dim_index(D); + let size = self.shape()[dim]; + if size == 0 { + // If the dimension is empty, return the tensor as is. + return self; + } + + let shift = wrap_index(shift, size); + if shift == 0 { + // If the shift is zero, return the tensor as is. + return self; + } + + self.unchecked_roll_dim(shift, dim) + } + + /// Internal implementation of `roll_dim` that does not canonicalize dimensions or shifts. + /// + /// ## Parameters + /// + /// - `shift`: The number of positions to shift; must be (0 < shift < size). + /// - `dim`: The dimension to roll; must be a valid index for the tensor's shape. + /// + /// ## Returns + /// + /// A new tensor with the specified dimension rolled by the given shift amount. + #[inline(always)] + fn unchecked_roll_dim(self, shift: usize, dim: usize) -> Self { + #[cfg(debug_assertions)] + { + let size = self.shape()[dim]; + assert!( + 0 < shift && shift < size, + "Expected: 0 < shift < size: found shift={shift}, size={size}", + ); + assert!( + dim < self.shape().num_dims(), + "Expected: dim < num_dims: found dim={dim}, num_dims={size}", + ); + } + + Tensor::cat( + vec![ + self.clone().slice_dim(dim, shift..), + self.slice_dim(dim, ..shift), + ], + dim, + ) + } + + /// Roll operation. + /// + /// Note: unlike ``pytorch``, `dims` and `shifts` must have the same length. + /// + /// A given `dim` may be rolled multiple times, and the shifts will be applied sequentially. + /// + /// ## Parameters + /// + /// - `shifts`: A slice of shifts corresponding to each dimension; + /// supports negative values and wraps around. + /// - `dims`: A slice of dimensions to roll; supports negative indexing. + /// + /// ## Returns + /// + /// A new tensor with the specified dimensions rolled by the given shifts. + pub fn roll(self, shifts: &[Shift], dims: &[Dim]) -> Self + where + Shift: AsIndex, + Dim: AsIndex, + { + assert_eq!( + dims.len(), + shifts.len(), + "Dimensions and shifts must align; found dims={dims:#?}, shifts={shifts:#?}", + ); + + // This is a fair amount of complexity, which could be replaced + // by a simple canonicalization of `dims` and wrapping of `shifts`. + // The work is done here to ensure that any roll operation + // which could be a no-op is a no-op; simplifying the accounting + // needed by backend-specific implementations of the inner roll op. + + let item_count = dims.len(); + + let shape = self.shape(); + + // Accumulate the effective shifts for each dimension. + let mut accumulated_shifts: Vec = vec![0; shape.len()]; + for i in 0..item_count { + let dim = dims[i].expect_dim_index(D); + accumulated_shifts[dim] += shifts[i].as_index(); + } + + // Do this after we've checked the validity of `dims` and `shifts`. + if self.shape().num_elements() == 0 { + // If the tensor is empty, return it as is. + return self; + } + + // Wrap the accumulated shifts, and filter out empty dimensions. + let mut effective_dims: Vec = Vec::with_capacity(item_count); + let mut effective_shifts: Vec = Vec::with_capacity(item_count); + for dim in 0..shape.len() { + // `wrap_index` should inline, and has a fast-exit path for zero shifts. + let shift = wrap_index(accumulated_shifts[dim], shape[dim]); + if shift == 0 { + continue; + } + + effective_dims.push(dim); + effective_shifts.push(shift); + } + + // If no shifts are needed, return the original tensor. + if effective_shifts.is_empty() { + return self; + } + + // At this point: + // - `dims` contains the effective dimensions to roll, in index order, + // - `shifts` contains the effective usize shifts for each dimension. + // - Every shift is non-zero, and less than the size of the corresponding dimension. + self.unchecked_roll(&effective_shifts, &effective_dims) + } + + /// `roll` internal implementation. + /// + /// ## Parameters + /// + /// - `shifts`: A slice of shifts corresponding to each dimension; + /// must be non-empty, the same length as `dims`, and all ``1..``. + /// - `dims`: A slice of dimensions to roll; must be non-empty; + /// the same length as `shifts`, and must not contain repeats. + /// + /// ## Panics + /// + /// Panics if the shifts and dimensions do not align, or if dimensions contain repeats. + /// + /// ## Returns + /// + /// A new tensor with the specified dimensions rolled by the given shifts. + #[inline(always)] + fn unchecked_roll(self, shifts: &[usize], dims: &[usize]) -> Self { + #[cfg(debug_assertions)] + { + assert!(!shifts.is_empty()); + assert_eq!( + shifts.len(), + dims.len(), + "Shifts and dimensions must align; found {} shifts and {} dims", + shifts.len(), + dims.len() + ); + + let mut unique_dims = dims.to_vec(); + unique_dims.dedup(); + + assert_eq!( + unique_dims.len(), + dims.len(), + "Dimensions must not contain repeats; found {} unique dims and {} total dims", + unique_dims.len(), + dims.len() + ) + } + + let x = self.unchecked_roll_dim(shifts[0], dims[0]); + + if dims.len() == 1 { + x + } else { + x.unchecked_roll(&shifts[1..], &dims[1..]) + } + } + + /// Returns a tensor containing the elements selected from the given slices. + /// + /// This method provides flexible tensor slicing with support for various range types, + /// negative indices, and stepped slicing. The method accepts both single slices and + /// arrays of slices, with the [`s!`] macro providing convenient syntax for complex patterns. + /// + /// # Arguments + /// + /// * `slices` - Can be: + /// - A single range for 1D slicing (e.g., `0..5`, `..`, `2..`) + /// - An array of ranges (e.g., `[0..2, 1..4]`) + /// - The [`s!`] macro output for advanced slicing with steps + /// - a `&Vec` or `&[Slice]` + /// + /// # Behavior + /// + /// - Supports partial and full slicing in any number of dimensions + /// - Handles negative indices by wrapping from the end (-1 is the last element) + /// - Automatically clamps ranges that exceed tensor dimensions + /// - Supports stepped slicing for selecting every nth element + /// - Negative steps reverse the selection order + /// + /// # Panics + /// + /// - If the number of slices exceeds the tensor's dimensions + /// - If a range is descending (e.g., 2..1) or empty (e.g., 1..1) without negative step + /// - If a step is zero + /// + /// # Examples + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape, s}; + /// + /// fn example() { + /// let device = Default::default(); + /// + /// // Single dimension slicing - no brackets needed! + /// let tensor = Tensor::<1, burn_tensor::Int>::arange(0..10, &device); + /// let slice = tensor.clone().slice(2..8); // Simple range + /// assert_eq!(slice.into_data().to_vec::().unwrap(), vec![2, 3, 4, 5, 6, 7]); + /// + /// // Using s! macro for single dimension with step + /// let slice = tensor.clone().slice(s![0..10;2]); // Every 2nd element + /// assert_eq!(slice.into_data().to_vec::().unwrap(), vec![0, 2, 4, 6, 8]); + /// + /// // Reverse a dimension with negative step + /// let slice = tensor.slice(s![..;-1]); // Reverse entire tensor + /// assert_eq!(slice.into_data().to_vec::().unwrap(), vec![9, 8, 7, 6, 5, 4, 3, 2, 1, 0]); + /// + /// // Multi-dimensional slicing + /// let tensor = Tensor::<2>::ones(Shape::new([4, 6]), &device); + /// + /// // Array syntax for simple ranges + /// let slice = tensor.clone().slice([1..3, 2..5]); + /// assert_eq!(slice.dims(), [2, 3]); + /// + /// // Advanced multi-dimensional with s! macro + /// let slice = tensor.clone().slice(s![0..4;2, ..;-1]); // Every 2nd row, reverse columns + /// assert_eq!(slice.dims(), [2, 6]); + /// + /// // Complex 3D example with mixed slice types + /// let tensor = Tensor::<3>::ones(Shape::new([4, 6, 8]), &device); + /// let slice = tensor.slice(s![1..3, ..;2, -3..]); // Rows 1-2, every 2nd col, last 3 depth + /// assert_eq!(slice.dims(), [2, 3, 3]); + /// + /// // Using negative indices + /// let tensor = Tensor::<2>::ones(Shape::new([4, 6]), &device); + /// let slice = tensor.slice(s![-2.., ..-1]); // Last 2 rows, all but last column + /// assert_eq!(slice.dims(), [2, 5]); + /// } + /// ``` + /// + /// # See Also + /// + /// - [`s!`] - The recommended macro for creating complex slice specifications + /// - [`slice_assign`](Self::slice_assign) - Assign values to a slice + /// - [`slice_fill`](Self::slice_fill) - Fill a slice with a constant value + /// - [`slice_dim`](Self::slice_dim) - Slice a single dimension + /// + /// [`s!`]: crate::s! + pub fn slice(self, slices: S) -> Self + where + S: SliceArg, + { + let shape = self.shape(); + let slices = slices.into_slices(&shape); + + // Validate slices + check!(TensorCheck::slice::(&shape, &slices)); + + // Calculate output shape and check for empty slices + let mut output_dims = shape.clone(); + for (dim, slice) in slices.iter().enumerate() { + output_dims[dim] = slice.output_size(shape[dim]); + } + + // Return empty tensor if any dimension is 0 (empty slice) + if output_dims.contains(&0) { + return Self::empty(output_dims, &self.device()); + } + Self::new(K::slice(self.primitive, &slices)) + } + + /// Assigns values to a slice of the tensor and returns the updated tensor. + /// + /// This method supports advanced slicing with steps, including negative steps for reverse + /// assignment. Like `slice`, it accepts both single slices and arrays, with the [`s!`] macro + /// providing powerful syntax for complex patterns. + /// + /// # Arguments + /// + /// * `slices` - Slice specification (same format as `slice` method) + /// * `values` - Tensor with values to assign (must match slice dimensions) + /// + /// # Panics + /// + /// - If slices exceed tensor dimensions + /// - If values dimensions don't match the selected slice shape + /// - If a step is zero + /// + /// # Examples + /// + /// ```rust + /// use burn_tensor::{Tensor, s}; + /// + /// fn example() { + /// let device = Default::default(); + /// + /// // Simple assignment to a sub-region + /// let mut tensor = Tensor::<2>::zeros([4, 6], &device); + /// let values = Tensor::<2>::ones([2, 3], &device); + /// tensor = tensor.slice_assign([1..3, 2..5], values); + /// // Now tensor[1..3, 2..5] contains ones + /// + /// // Single dimension assignment with step + /// let mut tensor = Tensor::<1>::zeros([10], &device); + /// let values = Tensor::<1>::ones([5], &device); + /// tensor = tensor.slice_assign(s![0..10;2], values); + /// // Now every 2nd element is 1: [1, 0, 1, 0, 1, 0, 1, 0, 1, 0] + /// + /// // Reverse assignment with negative step + /// let mut tensor = Tensor::<1>::from_data([0.0, 1.0, 2.0, 3.0, 4.0], &device); + /// let values = Tensor::<1>::from_data([10.0, 11.0, 12.0, 13.0, 14.0], &device); + /// tensor = tensor.slice_assign(s![..;-1], values); + /// // Assigns in reverse: [14, 13, 12, 11, 10] + /// + /// // Complex multi-dimensional assignment + /// let mut tensor = Tensor::<3>::zeros([4, 6, 8], &device); + /// let values = Tensor::<3>::ones([2, 3, 3], &device); + /// tensor = tensor.slice_assign(s![0..4;2, ..;2, -3..], values); + /// // Assigns to every 2nd row, every 2nd column, last 3 in depth + /// + /// // Mixed syntax example + /// let mut tensor = Tensor::<2>::zeros([8, 8], &device); + /// let pattern = Tensor::<2>::ones([4, 4], &device); + /// tensor = tensor.slice_assign(s![..;2, ..;2], pattern); + /// // Creates a checkerboard pattern with ones + /// } + /// ``` + /// + /// # See Also + /// + /// - [`s!`] - The recommended macro for creating complex slice specifications + /// - [`slice`](Self::slice) - Extract a slice from a tensor + /// - [`slice_fill`](Self::slice_fill) - Fill a slice with a constant value + /// + /// [`s!`]: crate::s! + pub fn slice_assign(self, slices: S, values: Self) -> Self + where + S: SliceArg, + { + let shape = self.shape(); + let slices = slices.into_slices(&shape); + + // Check if any slice produces 0 elements (empty assignment). + // Empty assignments are no-ops and would cause issues in backend implementations. + let is_empty_assignment = slices + .iter() + .enumerate() + .any(|(i, slice)| slice.output_size(shape[i]) == 0); + + if is_empty_assignment { + return self; + } + + check!(TensorCheck::slice_assign::( + &shape, + &values.shape(), + &slices + )); + + Self::new(K::slice_assign(self.primitive, &slices, values.primitive)) + } + + /// Fills a slice of the tensor with a constant value and returns the updated tensor. + /// + /// Like other slice methods, accepts both single slices and arrays. However, this method + /// currently **does not support stepped slicing** - use [`slice_assign`](Self::slice_assign) + /// with a constant tensor for stepped patterns. + /// + /// # Arguments + /// + /// * `slices` - Slice specification (same format as `slice` method, but no steps) + /// * `value` - The value to fill the slice with + /// + /// # Panics + /// + /// - If slices exceed tensor dimensions + /// - If any slice has a step != 1 (not yet supported) + /// + /// # Examples + /// + /// ```rust + /// use burn_tensor::{Tensor, s}; + /// + /// fn example() { + /// let device = Default::default(); + /// + /// // Simple fill for a single dimension + /// let mut tensor = Tensor::<1>::zeros([10], &device); + /// tensor = tensor.slice_fill(2..5, 1.0); + /// // Now tensor is [0, 0, 1, 1, 1, 0, 0, 0, 0, 0] + /// + /// // Multi-dimensional fill + /// let mut tensor = Tensor::<2>::zeros([4, 6], &device); + /// tensor = tensor.slice_fill([1..3, 2..5], -1.0); + /// // Fills the rectangle at rows 1-2, columns 2-4 with -1 + /// + /// // Using negative indices + /// let mut tensor = Tensor::<1>::zeros([10], &device); + /// tensor = tensor.slice_fill(-3.., 2.0); + /// // Fills the last 3 elements with 2.0 + /// + /// // Complex multi-dimensional example + /// let mut tensor = Tensor::<3>::ones([4, 6, 8], &device); + /// tensor = tensor.slice_fill(s![1..3, .., -2..], 0.0); + /// // Sets rows 1-2, all columns, last 2 in depth to 0 + /// + /// // Stepped slicing is supported + /// let mut tensor = Tensor::<1>::zeros([10], &device); + /// tensor = tensor.slice_fill(s![0..10;2], 1.0); + /// // Now every 2nd element is 1: [1, 0, 1, 0, 1, 0, 1, 0, 1, 0] + /// } + /// ``` + /// + /// # See Also + /// + /// - [`s!`] - The macro for creating slice specifications with steps + /// - [`slice`](Self::slice) - Extract a slice from a tensor + /// - [`slice_assign`](Self::slice_assign) - Assign tensor values to a slice + /// + /// [`s!`]: crate::s! + pub fn slice_fill(self, slices: S, value: E) -> Self + where + S: SliceArg, + { + let shape = self.shape(); + let slices = slices.into_slices(&shape); + + check!(TensorCheck::slice::(&shape, &slices)); + + let slice_shape = shape.slice(&slices).unwrap(); + let value = Tensor::<1, K>::from_data([value], (&self.device(), self.dtype())); + let value = value.expand(slice_shape); + self.slice_assign(&slices, value) + } + + /// Returns a new tensor with the specified dimension sliced. + /// + /// # Arguments + /// + /// * `dim`: The dimension to slice. + /// * `slice`: The slice specification for the dimension. Can be a range (e.g., `2..5`), + /// slice with step (via `s!` macro, e.g., `s![0..10;2]`), or any type that implements `Into`. + /// + /// # Returns + /// + /// A new tensor with the specified dimension sliced. + /// + /// # Panics + /// + /// If the slice is out of bounds for the specified dimension. + /// + /// # Examples + /// + /// ```rust + /// # use burn_tensor::{Tensor, s}; + /// # + /// # fn example() { + /// # let device = Default::default(); + /// let tensor = Tensor::<3>::zeros([3, 4, 5], &device); + /// + /// // Simple range slicing + /// let sliced = tensor.clone().slice_dim(1, 1..3); + /// assert_eq!(sliced.shape().as_slice(), [3, 2, 5]); + /// + /// // Slicing with step - take every 2nd element + /// let sliced = tensor.clone().slice_dim(2, s![0..5;2]); + /// assert_eq!(sliced.shape().as_slice(), [3, 4, 3]); // Takes indices 0, 2, 4 + /// + /// // Reverse slicing with negative step + /// let sliced = tensor.clone().slice_dim(1, s![..;-1]); + /// assert_eq!(sliced.shape().as_slice(), [3, 4, 5]); // Reverses dimension 1 + /// + /// // Select from index 2 with step 3 + /// let sliced = tensor.clone().slice_dim(0, s![2..;3]); + /// assert_eq!(sliced.shape().as_slice(), [1, 4, 5]); // Takes only index 2 + /// + /// // Select single index (reduces dimension to size 1) + /// let sliced = tensor.slice_dim(0, 1); + /// assert_eq!(sliced.shape().as_slice(), [1, 4, 5]); + /// # } + /// ``` + /// + /// # See Also + /// + /// - [`slice`](Self::slice) - Slice multiple dimensions simultaneously + /// - [`s!`] - The macro for creating complex slice specifications + /// + /// [`s!`]: crate::s! + pub fn slice_dim(self, dim: usize, slice: S) -> Self + where + S: Into, + { + check!(TensorCheck::check_dim::(dim)); + let slice: Slice = slice.into(); + + let mut slices = vec![Slice::full(); D]; + slices[dim] = slice; + + self.slice(&slices) + } + + /// Returns the device of the current tensor. + pub fn device(&self) -> Device { + K::device(&self.primitive) + } + + /// Move the tensor to the given device. + pub fn to_device(self, device: &Device) -> Self { + Self::new(K::to_device(self.primitive, device)) + } + + /// Select tensor elements along the given dimension corresponding to the given indices. + /// + /// # Arguments + /// + /// * `dim` - The dimension to select from. Supports negative indexing. + /// * `indices` - The indices of the elements to select. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Int}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [4.0, 5.0, 6.0]], &device); + /// let indices = Tensor::<1, Int>::from_data([0], &device); + /// let tensor = tensor.select(0, indices); + /// println!("{tensor}"); + /// // [[1.0, -2.0, 3.0]] + /// } + /// ``` + pub fn select(self, dim: impl AsIndex, indices: Tensor<1, Int>) -> Self { + let dim = dim.expect_dim_index(D); + check!(TensorCheck::select::(dim)); + Self::new(K::select(self.primitive, dim, indices.primitive)) + } + + /// Assign the selected elements along the given dimension corresponding to the given indices + /// from the value tensor to the original tensor using sum reduction. + /// + /// # Note + /// For booleans, the sum operator is logical or. + /// + /// # Arguments + /// + /// * `dim` - The dimension along which to select. Supports negative indexing. + /// * `indices` - The indices to select from the tensor. + /// * `values` - The values to assign to the selected indices. + /// * `update` - The operation used to update the existing values at the indexed positions (e.g., add). + /// + /// # Example + /// + /// Example using a 3D tensor: + /// + /// `input[indices[i], j, k] += values[i, j, k]; // dim = 0` + /// `input[i, indices[j], k] += values[i, j, k]; // dim = 1` + /// `input[i, j, indices[k]] += values[i, j, k]; // dim = 2` + /// `input[i, j, indices[k]] += values[i, j, k]; // dim = -1 (same as dim = 2)` + /// + /// # Warning + /// + /// Not all backends have runtime bound checks for the indices, so make sure they are valid. + /// Otherwise, out of bounds indices could lead to unexpected results instead of panicking. + pub fn select_assign( + self, + dim: impl AsIndex, + indices: Tensor<1, Int>, + values: Tensor, + update: IndexingUpdateOp, + ) -> Self { + let dim = dim.expect_dim_index(D); + check!(TensorCheck::select_assign::( + dim, + &indices.shape(), + &values.shape() + )); + + Self::new(K::select_assign( + self.primitive, + dim, + indices.primitive, + values.primitive, + update, + )) + } + + /// Update the given tensor with the value tensor where the mask is true. + /// + /// This is similar to [mask_fill](Tensor::mask_fill), however the value is a tensor instead of + /// a scalar. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape, Bool}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let mask = Tensor::<2, Bool>::from_data([[true, false, true], [false, true, false]], &device); + /// let value = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device); + /// let tensor = tensor.mask_where(mask, value); + /// println!("{tensor}"); + /// // [[2.0, -2.0, 4.0], [5.0, 2.0, 6.0]] + /// } + /// ``` + pub fn mask_where(self, mask: Tensor, value: Self) -> Self { + Self::new(K::mask_where( + self.primitive, + mask.primitive, + value.primitive, + )) + } + + /// Update the given tensor with the value where the mask is true. + /// + /// This is similar to [mask_where](Tensor::mask_where), however the value is a scalar instead of + /// a tensor. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape, Bool}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let mask = Tensor::<2, Bool>::from_data([[true, false, true], [false, true, false]], &device); + /// let tensor = tensor.mask_fill(mask, 3.0); + /// println!("{tensor}"); + /// // [[3.0, -2.0, 3.0], [5.0, 3.0, 6.0]] + /// } + /// ``` + pub fn mask_fill(self, mask: Tensor, value: E) -> Self { + let value = Scalar::new(value, &self.dtype()); + Self::new(K::mask_fill(self.primitive, mask.primitive, value)) + } + + /// Gather tensor elements corresponding to the given indices from the specified dim. + /// + /// Example using a 3D tensor: + /// + /// `output[i, j, k] = input[indices[i, j, k], j, k]; // dim = 0` + /// `output[i, j, k] = input[i, indices[i, j, k], k]; // dim = 1` + /// `output[i, j, k] = input[i, j, indices[i, j, k]]; // dim = 2` + /// + /// # Notes + /// + /// The index tensor should have the same shape as the original tensor except for the dim + /// specified. + /// + /// # Warning + /// Not all backends have runtime bound checks for the indices, so make sure the they are valid. + /// Otherwise, out of bounds indices could lead to unexpected results instead of panicking. + pub fn gather(self, dim: usize, indices: Tensor) -> Self { + check!(TensorCheck::gather::( + dim, + &self.shape(), + &indices.shape() + )); + + Self::new(K::gather(dim, self.primitive, indices.primitive)) + } + + /// Assign the gathered elements corresponding to the given indices along the specified dimension + /// from the value tensor to the original tensor using sum reduction. + /// + /// Example using a 3D tensor: + /// + /// `input[indices[i, j, k], j, k] += values[i, j, k]; // dim = 0` + /// `input[i, indices[i, j, k], k] += values[i, j, k]; // dim = 1` + /// `input[i, j, indices[i, j, k]] += values[i, j, k]; // dim = 2` + /// + /// # Arguments + /// * `dim` - The axis along which to scatter elements. + /// * `indices` - The indices of the elements to scatter. + /// * `values` - The values to scatter into the tensor. + /// * `update` - The operation used to update the existing values at the indexed positions (e.g., add). + /// + /// # Notes + /// + /// The index tensor should have the same shape as the original tensor except for the specified + /// dimension. The value and index tensors should have the same shape. + /// + /// Other references to the input tensor will not be modified by this operation. + /// + /// # Warning + /// Not all backends have runtime bound checks for the indices, so make sure the they are valid. + /// Otherwise, out of bounds indices could lead to unexpected results instead of panicking. + /// + /// # Panics + /// If the `update` is not `IndexingUpdateOp::Add`. Other operations are currently not implemented. + pub fn scatter( + self, + dim: usize, + indices: Tensor, + values: Self, + update: IndexingUpdateOp, + ) -> Self { + check!(TensorCheck::scatter::( + dim, + &self.shape(), + &indices.shape(), + &values.shape() + )); + + Self::new(K::scatter( + dim, + self.primitive, + indices.primitive, + values.primitive, + update, + )) + } + + /// Multi-dimensional scatter: update the tensor at locations given by `indices` using the specified `update` operation. + /// + /// The size of `indices`'s last axis (call it `K`) indexes the leading `K` dims of `self`; + /// the batch shape `indices.shape[0..M-1]` is preserved. `values` has shape + /// `indices.shape[0..M-1] ++ self.shape[K..D]`. Constraints: `K <= D` and `M >= 1`. + /// + /// # Arguments + /// * `indices` - The indices of the elements to scatter. + /// * `values` - The values to scatter into the tensor. + /// * `update` - The operation used to update the existing values at the indexed positions (e.g., add). + /// + /// # Note + /// + /// When `indices` contains duplicate entries, behavior varies by operation: + /// - For `Add`, accumulation is supported, though results may be non-deterministic on GPU + /// backends. + /// - For other operations (`Assign`, `Mul`, `Min`, `Max`), duplicate indices result in + /// undefined behavior for both the forward result and the backward gradients. + /// + /// For deterministic results and correct gradient calculation across all operations, + /// `indices` should contain unique entries. + /// + /// # Warning + /// + /// Not all backends have runtime bound checks for the indices, so make sure they are valid. + /// Otherwise, out of bounds indices could lead to unexpected results instead of panicking. + pub fn scatter_nd( + self, + indices: Tensor, + values: Tensor, + update: IndexingUpdateOp, + ) -> Self { + check!(TensorCheck::scatter_nd::( + &self.shape(), + &indices.shape(), + &values.shape() + )); + Self::new(K::scatter_nd( + self.primitive, + indices.primitive, + values.primitive, + update, + )) + } + + /// Multi-dimensional gather: collect slices from `self` at multi-index locations + /// specified by `indices`. + /// + /// The size of `indices`'s last axis (call it `K`) indexes the leading `K` dims of `self`; + /// the batch shape `indices.shape[0..M-1]` is preserved. The output has shape + /// `indices.shape[0..M-1] ++ self.shape[K..D]`. Constraints: `K <= D` and `M >= 1`. + /// + /// # Warning + /// + /// Not all backends have runtime bound checks for the indices, so make sure they are valid. + /// Otherwise, out of bounds indices could lead to unexpected results instead of panicking. + pub fn gather_nd( + self, + indices: Tensor, + ) -> Tensor { + check!(TensorCheck::gather_nd::(&indices.shape())); + Tensor::new(K::gather_nd(self.primitive, indices.primitive)) + } + + /// Converts the data of the current tensor. + /// + /// # Note + /// + /// For better performance, prefer using a [Transaction](crate::Transaction) when reading multiple + /// tensors at once. This may improve laziness, especially if executed on a different + /// thread in native environments. + pub fn into_data(self) -> TensorData { + into_data_sync_impl(self.primitive, K::KIND) + } + + /// Converts the data of the current tensor and returns any error that might have occurred since the + /// last time the device was synchronized. + /// + /// # Note + /// + /// For better performance, prefer using a [Transaction](crate::Transaction) when reading multiple + /// tensors at once. This may improve laziness, especially if executed on a different + /// thread in native environments. + pub fn try_into_data(self) -> Result { + try_into_data_sync_impl(self.primitive, K::KIND) + } + + /// Converts the data of the current tensor. + /// + /// # Note + /// + /// For better performance, prefer using a [Transaction](crate::Transaction) when reading multiple + /// tensors at once. This may improve laziness, especially if executed on a different + /// thread in native environments. + pub fn to_data(&self) -> TensorData { + self.clone().into_data() + } + + /// Returns the data of the current tensor. + pub fn into_data_async( + self, + ) -> impl core::future::Future> + Send { + into_data_async_impl(self.primitive, K::KIND) + } + + /// Returns the data of the current tensor. + pub fn to_data_async( + &self, + ) -> impl core::future::Future> + Send { + into_data_async_impl(self.primitive.clone(), K::KIND) + } + + /// Create a tensor from the given data on the given device. + pub fn from_data(data: T, options: impl Into) -> Self + where + T: Into, + { + let data = data.into(); + check!(TensorCheck::creation_ops::( + "From Data", + data.shape.as_slice() + )); + + // Use the given dtype when provided, otherwise default device dtype + let opt = options.into(); + let dtype = opt.resolve_dtype::(); + + Self::new(K::from_data(data, &opt.device, dtype)) + } + + /// Repeat the tensor along the given dimension. + /// + /// The output tensor has the same shape, except along the given dimension. + /// + /// # Arguments + /// - `dim`: The dimension to repeat. + /// - `times`: The number of times to repeat the tensor along the given dimension in the new tensor. + /// + /// # Returns + /// + /// A new tensor with the given dimension repeated `times` times. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 2D tensor with dimensions [3, 2] + /// let tensor = Tensor::<2>::from_data([[3.0, 4.9], [2.0, 1.9], [4.0, 5.9]], &device); + /// + /// // Repeat the tensor along the dimension 0 twice. + /// // [[3.0, 4.9], [2.0, 1.9], [4.0, 5.9], [3.0, 4.9], [2.0, 1.9], [4.0, 5.9]] + /// // The resulting tensor will have dimensions [6, 2]. + /// let repeated = tensor.repeat_dim(0, 2); + /// println!("{repeated}"); + /// } + /// ``` + pub fn repeat_dim(self, dim: usize, times: usize) -> Self { + if times > 0 { + Self::new(K::repeat_dim(self.primitive, dim, times)) + } else { + let shape = self.shape().repeat(dim, times).unwrap(); + Self::empty(shape, &self.device()) + } + } + + /// Repeat the tensor along the given dimensions. + /// # Arguments + /// - `sizes`: Borrowed slice of the number of times to repeat each dimension. + /// + /// # Returns + /// + /// A new tensor with the given dimensions repeated `times` times. + /// + /// # Panics + /// + /// If `sizes` contains more elements than the number of dimensions. + /// + /// # Example + /// + /// ```rust + /// + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 2D tensor with dimensions [3, 2] + /// let tensor = Tensor::<2>::from_data([[3.0, 4.9], [2.0, 1.9], [4.0, 5.9]], &device); + /// + /// // Repeat the tensor along the dimension 0 twice and the dimension 0 once. + /// // [[3.0, 4.9], [2.0, 1.9], [4.0, 5.9], [3.0, 4.9], [2.0, 1.9], [4.0, 5.9]] + /// // The resulting tensor will have dimensions [6, 2]. + /// let repeated = tensor.repeat(&[2, 1]); + /// } + /// ``` + pub fn repeat(self, sizes: &[usize]) -> Self { + if sizes.contains(&0) { + let mut shape = self.shape(); + for (dim, ×) in sizes.iter().enumerate() { + shape = shape.repeat(dim, times).unwrap(); + } + + return Self::empty(shape, &self.device()); + } + + let mut tensor = self; + for (dim, ×) in sizes.iter().enumerate() { + if times > 1 { + tensor = tensor.repeat_dim(dim, times); + } + } + tensor + } + + /// Applies element-wise equal comparison. + /// + /// # Returns + /// A boolean tensor that is `true` where input is equal to `other` and `false` elsewhere. + /// + /// # Panics + /// + /// If the two tensors don't have the same shape. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// let t1 = Tensor::<2>::from_data([[2.0, 4.9], [2.0, 1.9], [4.0, 5.9]], &device); + /// let t2 = Tensor::<2>::from_data([[3.0, 4.9], [2.0, 1.9], [4.0, 5.9]], &device); + /// // Compare the elements of the two 2D tensors with dimensions [3, 2]. + /// // [[false, true], [true, true], [true, true]] + /// let equal = t1.equal(t2); + /// println!("{equal}"); + /// } + /// ``` + pub fn equal(self, other: Self) -> Tensor { + check!(TensorCheck::binary_ops_ew("Equal", &self, &other)); + Tensor::new(K::equal(self.primitive, other.primitive)) + } + + /// Applies element-wise non-equality comparison. + /// + /// # Returns + /// A boolean tensor that is `true` where input is not equal to `other` and `false` elsewhere. + /// + /// # Panics + /// + /// If the two tensors don't have the same shape. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// let t1 = Tensor::<2>::from_data([[2.0, 4.9], [2.0, 1.9], [4.0, 5.9]], &device); + /// let t2 = Tensor::<2>::from_data([[3.0, 4.9], [2.0, 1.9], [4.0, 5.9]], &device); + /// // Compare the elements of the two 2D tensors for inequality. + /// // [[true, false], [false, false], [false, false]] + /// let not_equal = t1.not_equal(t2); + /// println!("{not_equal}"); + /// } + /// ``` + pub fn not_equal(self, other: Self) -> Tensor { + check!(TensorCheck::binary_ops_ew("NotEqual", &self, &other)); + Tensor::new(K::not_equal(self.primitive, other.primitive)) + } + + /// Applies element wise equal comparison and returns a boolean tensor. + /// + /// # Arguments + /// + /// * `other` - The element to compare. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.equal_elem(3.0); + /// println!("{tensor}"); + /// // [[false, false, true], [false, false, false]] + /// } + /// ``` + pub fn equal_elem(self, other: E) -> Tensor { + let other = Scalar::new(other, &self.dtype()); + Tensor::new(K::equal_elem(self.primitive, other)) + } + + /// Applies element wise non-equality comparison and returns a boolean tensor. + /// + /// # Arguments + /// + /// * `other` - The element to compare. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.not_equal_elem(3.0); + /// println!("{tensor}"); + /// // [[true, true, false], [true, true, true]] + /// } + /// ``` + pub fn not_equal_elem(self, other: E) -> Tensor { + let other = Scalar::new(other, &self.dtype()); + Tensor::new(K::not_equal_elem(self.primitive, other)) + } + + /// Concatenates all tensors into a new one along the given dimension. + /// + /// # Panics + /// + /// - If `dim` is higher than the rank. + /// - If `tensors` is an empty vector. + /// - If all tensors don't have the same shape (the dimension `dim` is ignored). + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// let t1 = Tensor::<2>::from_data([[3.0, 4.9, 2.0, 1.0], [2.0, 1.9, 3.0, 1.0]], &device); + /// let t2 = Tensor::<2>::from_data([[4.0, 5.9, 8.0], [1.4, 5.8, 6.0]], &device); + /// + /// // Concatenate the two tensors with shapes [2, 4] and [2, 3] along the dimension 1. + /// // [[3.0, 4.9, 2.0, 1.0, 4.0, 5.9, 8.0], [2.0, 1.9, 3.0, 1.0, 1.4, 5.8, 6.0]] + /// // The resulting tensor will have shape [2, 7]. + /// let concat = Tensor::cat(vec![t1, t2], 1); + /// println!("{concat}"); + /// } + /// ``` + pub fn cat(tensors: Vec, dim: usize) -> Self { + check!(TensorCheck::cat(tensors.as_slice(), dim)); + + // Filter out tensors with size 0 along the concatenation dimension. + // Empty tensors don't contribute to the output and would cause issues + // in backend implementations (e.g., division by zero in slice_assign). + // Safety: TensorCheck::cat ensures tensors is non-empty + let first_tensor = tensors.first().unwrap(); + let device = first_tensor.device(); + let mut shape = first_tensor.shape(); + + let non_empty_primitives: Vec<_> = tensors + .into_iter() + .filter(|t| t.shape()[dim] > 0) + .map(|t| t.primitive) + .collect(); + + // If all tensors were empty, return an empty tensor with size 0 on concat dim + if non_empty_primitives.is_empty() { + shape[dim] = 0; + return Self::empty(shape, &device); + } + + Self::new(K::cat(non_empty_primitives, dim)) + } + + /// Concatenates all tensors into a new one along a new dimension. + /// + /// # Panics + /// + /// - If all tensors don't have the same shape. + /// - If given dimension is not with range of 0..D2 + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// let t1 = Tensor::<2>::from_data([[3.0, 4.9, 2.0], [2.0, 1.9, 3.0]], &device); + /// let t2 = Tensor::<2>::from_data([[4.0, 5.9, 8.0], [1.4, 5.8, 6.0]], &device); + /// let t3 = Tensor::<2>::from_data([[4.0, 5.9, 8.0], [1.4, 5.8, 6.0]], &device); + /// + /// // Concatenate the three tensors with shape [2, 3] along a new dimension, 0. + /// // [[[3.0, 4.9, 2.0], [2.0, 1.9, 3.0]], + /// // [[4.0, 5.9, 8.0], [1.4, 5.8, 6.0]], + /// // [[4.0, 5.9, 8.0], [1.4, 5.8, 6.0]]] + /// // The resulting tensor will have shape [3, 2, 3]. + /// let stacked= Tensor::stack::<3>(vec![t1, t2, t3], 0); + /// println!("{stacked}"); + /// } + /// ``` + pub fn stack(tensors: Vec>, dim: usize) -> Tensor { + check!(TensorCheck::stack::(tensors.as_slice(), dim)); + let tensors = tensors.into_iter().map(|t| t.unsqueeze_dim(dim)).collect(); + Tensor::::cat(tensors, dim) + } + + /// Iterate over slices of tensors alongside a given dimension. + /// + /// # Panics + /// + /// If given dimension is greater than or equal to tensor rank. + /// + /// # Returns + /// + /// A tensor iterator. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[3.0, 4.9, 2.0], [2.0, 1.9, 3.0]], &device); + /// // Given a 2D tensor with dimensions [2, 3], iterate over slices of tensors along the dimension 0. + /// let iter = tensor.iter_dim(0); + /// for (i,tensor) in iter.enumerate() { + /// println!("Tensor {}: {}", i, tensor); + /// // Tensor 0: Tensor { data: [[3.0, 4.9, 2.0]], ... } + /// // Tensor 1: Tensor { data: [[2.0, 1.9, 3.0]], ... } + /// } + /// } + /// ``` + pub fn iter_dim(self, dim: usize) -> DimIter { + check!(TensorCheck::dim_ops::("iter_dim", dim)); + DimIter::new(self, dim) + } + + /// Returns a new tensor with the given dimension narrowed to the given range. + /// + /// # Panics + /// + /// - If the dimension is greater than the number of dimensions of the tensor. + /// - If the given range exceeds the number of elements on the given dimension. + /// + /// # Returns + /// + /// A new tensor with the given dimension narrowed to the given range. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 2D tensor with dimensions [4, 3] + /// let tensor = Tensor::<2>::from_data( + /// [ + /// [3.0, 4.9, 2.0], + /// [2.0, 1.9, 3.0], + /// [6.0, 1.5, 7.0], + /// [3.0, 4.9, 9.0], + /// ], + /// &device, + /// ); + /// // Narrow the tensor along the dimension 0, keeping 3 elements starting from index 1. + /// // [[2.0, 1.9, 3.0], [6.0, 1.5, 7.0], [3.0, 4.9, 9.0]] + /// // The resulting tensor will have dimensions [3, 3]. + /// let narrowed = tensor.narrow(0, 1, 3); + /// println!("{narrowed}"); + /// } + /// ``` + pub fn narrow(self, dim: usize, start: usize, length: usize) -> Self { + check!(TensorCheck::dim_ops::("narrow", dim)); + check!(TensorCheck::narrow(&self, dim, start, length)); + let dims = self.dims(); + + let ranges: [Range; D] = dims + .iter() + .enumerate() + .map(|(i, d)| { + if i == dim { + start..(start + length) + } else { + 0..*d + } + }) + .collect::>() + .try_into() + .unwrap(); + + Self::slice(self, ranges) + } + + /// Attempts to split the tensor into a specified number of chunks along a given dimension. + /// May return less chunks than requested if the tensor size is not divisible by the number of chunks. + /// + /// When the given dimension is evenly divisible by the number of chunks, the chunks will be of equal size. + /// Otherwise all chunks will be of equal size except for the last one. + /// + /// # Panics + /// + /// If the dimension is greater than the number of dimensions of the tensor. + /// + /// # Returns + /// A vector of tensors. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 2D tensor with dimensions [4, 3] + /// let tensor = Tensor::<2>::from_data( + /// [ + /// [3.0, 4.9, 2.0], + /// [2.0, 1.9, 3.0], + /// [6.0, 1.5, 7.0], + /// [3.0, 4.9, 9.0], + /// ], + /// &device, + /// ); + /// // Split the tensor along the dimension 1 into 2 chunks. + /// // The first chuck will have shape [4, 2]: + /// // [[3.0, 4.9], [2.0, 1.9], [6.0, 1.5], [3.0, 4.9]] + /// // The second chunk will have shape [4, 1]: + /// // [[2.0], [3.0], [7.0], [9.0]] + /// let chunks = tensor.chunk(2, 1); + /// println!("{chunks:?}"); + /// } + /// ``` + pub fn chunk(self, chunks: usize, dim: usize) -> Vec { + check!(TensorCheck::dim_ops::("chunk", dim)); + let size = self.shape()[dim]; + if size < chunks { + return (0..size) + .map(|i| Self::narrow(self.clone(), dim, i, 1)) + .collect(); + } + + let mut tensors = Vec::with_capacity(chunks); + let mut sum_chunk_size = 0; + if size.is_multiple_of(chunks) { + let chunk_size = size / chunks; + for _ in 0..chunks { + tensors.push(Self::narrow(self.clone(), dim, sum_chunk_size, chunk_size)); + sum_chunk_size += chunk_size; + } + } else { + let chunk_size = (size / chunks) + 1; // assumes not divisible + for _ in 0..chunks - 1 { + tensors.push(Self::narrow(self.clone(), dim, sum_chunk_size, chunk_size)); + sum_chunk_size += chunk_size; + } + let remainder = size % chunk_size; + tensors.push(Self::narrow(self.clone(), dim, sum_chunk_size, remainder)); + } + + tensors + } + + /// Splits the tensor into chunks of a specified size along a given dimension. + /// Each chunk is a view of the original tensor. + /// + /// If the tensor size along the given dimension is not divisible by `split_size`, + /// then the last chunk will be smaller. + /// + /// # Panics + /// + /// If the specified dimension to split along is greater than the number of dimensions of the tensor. + /// + /// # Returns + /// + /// A vector of tensors. + /// + /// # Example + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 1D tensor with 5 elements + /// let tensor = Tensor::<1>::from_data([0.0, 1.0, 2.0, 3.0, 4.0], &device); + /// // Split the tensor into chunks of size 2 along dimension 0 + /// let chunks = tensor.split(2, 0); + /// // The result is a vector of tensors: + /// // [Tensor([0.0, 1.0]), Tensor([2.0, 3.0]), Tensor([4.0])] + /// println!("{:?}", chunks); + /// } + /// ``` + pub fn split(self, split_size: usize, dim: usize) -> Vec { + check!(TensorCheck::split::(&self.shape(), split_size, dim)); + let size = self.shape()[dim]; + let mut tensors = Vec::new(); + + let mut start = 0; + while start < size { + let length = usize::min(split_size, size - start); + tensors.push(Self::narrow(self.clone(), dim, start, length)); + start += length; + } + + tensors + } + + /// Splits the tensor into chunks with the specified sizes along a given dimension. + /// Each chunk is a view of the original tensor. + /// + /// The sizes of the chunks are specified in the `split_sizes` vector. The sum of the sizes + /// in `split_sizes` must equal the size of the tensor along the specified dimension. + /// + /// # Panics + /// + /// If the specified dimension to split along is greater than the number of dimensions of the tensor or + /// if the sum of `dim_sizes` does not equal the size of the tensor along `dim`. + /// + /// # Returns + /// + /// A vector of tensors. + /// + /// # Example + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 1D tensor with 5 elements + /// let tensor = Tensor::<1>::from_data([0.0, 1.0, 2.0, 3.0, 4.0], &device); + /// // Split the tensor into chunks with sizes [2, 3] along dimension 0 + /// let chunks = tensor.split_with_sizes(vec![2, 3], 0); + /// // The result is a vector of tensors: + /// // [Tensor([0.0, 1.0]), Tensor([2.0, 3.0, 4.0])] + /// println!("{:?}", chunks); + /// } + /// ``` + pub fn split_with_sizes(self, split_sizes: Vec, dim: usize) -> Vec { + check!(TensorCheck::split_with_sizes::( + &self.shape(), + &split_sizes, + dim + )); + let mut tensors = Vec::new(); + + let mut start = 0; + for length in split_sizes { + if length == 0 { + continue; + } + tensors.push(Self::narrow(self.clone(), dim, start, length)); + start += length; + } + + tensors + } + + /// Tests if any element in the `tensor` evaluates to True. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. All input tensor types (Float, Int, Bool) are supported. + /// + /// # Returns + /// + /// A boolean tensor `Tensor<1, Bool>` containing a single element, True if any element in the input tensor + /// evaluates to True, False otherwise. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2, Bool>::from_data([[true,false,true],[false,true,false]], &device); + /// let tensor_two = Tensor::<2, Bool>::from_data([[false,false,false],[false,false,false]], &device); + /// + /// // Given a 2D tensor with dimensions [2, 3], test if any element in the tensor evaluates to True. + /// let any_tensor = tensor.any(); + /// println!("{}", any_tensor); + /// // Tensor { data: [true], ... } + /// + /// // Given a 2D tensor with dimensions [2, 3], test if any element in the tensor evaluates to True. + /// let any_tensor_two = tensor_two.any(); + /// println!("{}", any_tensor_two); + /// // Tensor { data: [false], ... } + /// } + /// ``` + pub fn any(self) -> Tensor<1, Bool> { + Tensor::new(K::any(self.primitive)) + } + + /// Tests if any element in the `tensor` evaluates to True along a given dimension `dim`. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. All input tensor types (Float, Int, Bool) are supported. + /// * `dim` - The axis along which to test. + /// + /// # Returns + /// + /// A boolean tensor `Tensor` with the same shape as input `tensor`, except in the `dim` axis + /// where the size is 1. The elem in the `dim` axis is True if any element along this dim in the input + /// evaluates to True, False otherwise. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = + /// Tensor::<2, Bool>::from_data([[true, false, false], [false, true, false]], &device); + /// // Check if any element in the tensor evaluates to True along the dimension 1. + /// // [[true], [true]], + /// let any_dim = tensor.clone().any_dim(1); + /// println!("{any_dim}"); + /// } + /// ``` + pub fn any_dim(self, dim: usize) -> Tensor { + Tensor::new(K::any_dim(self.primitive, dim)) + } + + /// Tests if all elements in the `tensor` evaluate to True. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. All input tensor types (Float, Int, Bool) are supported. + /// + /// # Returns + /// + /// A boolean tensor `Tensor<1, Bool>` with a single element, True if all elements in the input tensor + /// evaluate to True, False otherwise. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = + /// Tensor::<2, Bool>::from_data([[true, false, true], [true, true, true]], &device); + /// // Check if all elements in the tensor evaluate to True (which is not the case). + /// // [false] + /// let all = tensor.all(); + /// println!("{all}"); + /// } + /// ``` + pub fn all(self) -> Tensor<1, Bool> { + Tensor::new(K::all(self.primitive)) + } + + /// Tests if all elements in the `tensor` evaluate to True along a given dimension `dim`. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to test. All input tensor types (Float, Int, Bool) are supported. + /// * `dim` - The axis along which to test. + /// + /// # Returns + /// + /// A boolean tensor `Tensor` with the same shape as input `tensor`, except in the `dim` axis + /// where the size is 1. The elem in the `dim` axis is True if all elements along this dim in the input + /// evaluates to True, False otherwise. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = + /// Tensor::<2, Bool>::from_data([[true, true, false], [true, true, true]], &device); + /// // Check if all elements in the tensor evaluate to True along the dimension 1. + /// // [[true, true, false]] + /// let all_dim = tensor.clone().all_dim(0); + /// println!("{all_dim}"); + /// } + /// ``` + pub fn all_dim(self, dim: usize) -> Tensor { + Tensor::new(K::all_dim(self.primitive, dim)) + } + + /// Convert the tensor into a scalar. + /// + /// # Panics + /// + /// - If the tensor doesn't have one element. + /// - If the backend fails to read the tensor data synchronously. + /// + /// # Returns + /// + /// The scalar value of the tensor. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[3.0]], &device); + /// // Convert the tensor with a single element into a scalar. + /// let scalar: f32 = tensor.into_scalar(); + /// println!("{scalar}"); + /// } + /// ``` + pub fn into_scalar(self) -> E { + check!(TensorCheck::into_scalar::(&self.shape())); + + let err_msg = + "Error while reading data: use `try_into_scalar` instead to catch the error at runtime"; + + let data = self.into_data(); + data.iter::().next().expect(err_msg) + } + + /// Convert the tensor into a scalar and returns any error that might have occurred since the + /// last time the device was synchronized. + /// + /// # Panics + /// + /// - If the tensor doesn't have one element. + /// - If the backend fails to read the tensor data synchronously. + /// + /// # Returns + /// + /// The scalar value of the tensor. + pub fn try_into_scalar(self) -> Result { + check!(TensorCheck::into_scalar::(&self.shape())); + + let err_msg = + "Error while reading data: use `try_into_scalar` instead to catch the error at runtime"; + + let data = self.try_into_data()?; + Ok(data.iter::().next().expect(err_msg)) + } + + /// Convert the tensor into a scalar. + /// + /// # Panics + /// + /// If the tensor doesn't have one element. + pub async fn into_scalar_async(self) -> Result { + check!(TensorCheck::into_scalar::(&self.shape())); + + let err_msg = + "Error while reading data: use `try_into_scalar` instead to catch the error at runtime"; + + let data = self.into_data_async().await?; + Ok(data.iter::().next().expect(err_msg)) + } + + /// Broadcast the tensor to the given shape. + /// + /// Only singleton dimensions can be expanded to a larger size. Other dimensions must have the same size + /// (which can be inferred with `-1`). + /// + /// # Arguments + /// + /// * `shape` - The shape to broadcast the tensor to. + /// Can contain -1 for dimensions that should be inferred. + /// The number of elements in the shape must be greater or equal as + /// the number of dimensions of the tensor. + /// + /// # Panics + /// + /// If the tensor cannot be broadcasted to the given shape. + /// + /// # Returns + /// + /// A new tensor with the given shape. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// // Create a 2D tensor with dimensions [3, 1] + /// let tensor = Tensor::<2>::from_data([[1.], [2.], [3.]], &device); + /// // Expand the tensor to a new shape [3, 4] + /// // [[1.0, 1.0, 1.0, 1.0], [2.0, 2.0, 2.0, 2.0], [3.0, 3.0, 3.0, 3.0]] + /// let expanded = tensor.expand([3, 4]); + /// println!("{}", expanded); + /// } + /// ``` + pub fn expand>(self, shape: S) -> Tensor { + let shape = shape.into_shape(&self.shape()); + check!(TensorCheck::expand::( + "expand", + &self.shape(), + &shape, + )); + + Tensor::::new(K::expand(self.primitive, shape)) + } + + /// Unfold windows along a dimension. + /// + /// Returns a view of the tensor with all complete windows of size `size` in dimension `dim`; + /// where windows are advanced by `step` at each index. + /// + /// The number of windows is `max(0, (shape[dim] - size).ceil_div(step))`. + /// + /// The new view will have the unfolded dimension replaced by two dimensions; + /// one in the position of the original dimension, with size equal to the number of windows, + /// and one appended to the right-most position, with size equal to `size`. + /// + /// # Warning + /// + /// For the `ndarray` backend; this is not a view but a copy + /// with duplicated data. + /// + /// # Arguments + /// + /// * `dim` - the dimension to unfold. + /// * `size` - the size of each unfolded window. + /// * `step` - the step between each window. + /// + /// # Returns + /// + /// A tensor view with the shape ``[pre=..., windows, post=..., size]``. + pub fn unfold( + self, + dim: I, + size: usize, + step: usize, + ) -> Tensor { + let dim = dim.expect_dim_index(D); + check!(TensorCheck::unfold::( + "unfold", + &self.shape(), + dim, + size, + step, + )); + Tensor::::new(K::unfold(self.primitive, dim, size, step)) + } +} + +/// Iterator given by (Tensor::iter_dim). +pub struct DimIter +where + K: Basic, +{ + start: usize, + end: usize, + dim: usize, + ranges: [Range; D], + tensor: Tensor, +} + +impl Iterator for DimIter { + type Item = Tensor; + + fn next(&mut self) -> Option { + if self.start >= self.end { + return None; + } + + let mut ranges = self.ranges.clone(); + ranges[self.dim] = self.start..(self.start + 1); + + let slice = self.tensor.clone().slice(ranges); + self.start += 1; + + Some(slice) + } +} + +impl DoubleEndedIterator for DimIter { + fn next_back(&mut self) -> Option { + if self.start >= self.end { + return None; + } + + let mut ranges = self.ranges.clone(); + ranges[self.dim] = (self.end - 1)..self.end; + + let slice = self.tensor.clone().slice(ranges); + self.end = self.end.saturating_sub(1); + + Some(slice) + } +} + +impl DimIter { + fn new(tensor: Tensor, dim: usize) -> Self { + let dims = tensor.dims(); + let ranges = dims + .iter() + .map(|&dim| 0..dim) + .collect::>>(); + let ranges: [Range; D] = ranges.try_into().unwrap(); + Self { + end: dims[dim], + ranges, + start: 0, + dim, + tensor, + } + } +} + +struct DataIterFmt { + data: TensorData, + precision: Option, +} + +fn fmt_float(elem: E, precision: Option) -> String { + match precision { + Some(p) => format!("{elem:.p$}"), + None => fmt_elem(elem), + } +} + +fn fmt_elem(elem: E) -> String { + format!("{elem:?}") +} + +// TODO: refactor display +impl DataIterFmt { + fn next(&self) -> String { + match self.data.dtype { + DType::F64 => fmt_float(self.next_elem::(), self.precision), + DType::F32 | DType::Flex32 => fmt_float(self.next_elem::(), self.precision), + DType::F16 => fmt_float(self.next_elem::(), self.precision), + DType::BF16 => fmt_float(self.next_elem::(), self.precision), + DType::I64 => fmt_elem(self.next_elem::()), + DType::I32 => fmt_elem(self.next_elem::()), + DType::I16 => fmt_elem(self.next_elem::()), + DType::I8 => fmt_elem(self.next_elem::()), + DType::U64 => fmt_elem(self.next_elem::()), + DType::U32 => fmt_elem(self.next_elem::()), + DType::U16 => fmt_elem(self.next_elem::()), + DType::U8 => fmt_elem(self.next_elem::()), + DType::Bool(store) => match store { + burn_std::BoolStore::Native => fmt_elem(self.next_elem::()), + burn_std::BoolStore::U8 => fmt_elem(self.next_elem::().to_bool()), + burn_std::BoolStore::U32 => fmt_elem(self.next_elem::().to_bool()), + }, + DType::QFloat(_) => todo!(), // unreachable but we should fix that + } + } + + fn next_elem(&self) -> E { + self.data.iter::().next().unwrap() + } +} + +// The Display-formatting recursion used to live as generic methods on +// `Tensor` here. It has been outlined to non-generic free functions +// (`display_fmt_*`, `slice_bridge_by_kind`, `push_newline_indent_impl`) below, +// so it is compiled exactly once inside `burn-tensor` instead of being +// re-monomorphized for every `(D, K)` in downstream crates. That outlining is +// the difference between a ~7s and a ~0.5s incremental release rebuild for a +// program that just calls `println!("{tensor}")`. + +#[derive(Clone, Debug)] +/// Options for Tensor pretty printing +pub struct PrintOptions { + /// number of elements to start summarizing tensor + pub threshold: usize, + + /// number of starting elements and ending elements to display + pub edge_items: usize, + + /// Precision for floating point numbers + pub precision: Option, +} + +static PRINT_OPTS: RwLock = RwLock::new(PrintOptions::const_default()); + +impl PrintOptions { + /// Print options with default values + pub const fn const_default() -> Self { + Self { + threshold: 1000, + edge_items: 3, + precision: None, + } + } +} + +impl Default for PrintOptions { + fn default() -> Self { + Self::const_default() + } +} + +/// Set print options +pub fn set_print_options(options: PrintOptions) { + let mut print_opts = PRINT_OPTS.write().unwrap(); + *print_opts = options; +} + +/// Pretty print tensors +impl core::fmt::Display for Tensor +where + K: Basic, +{ + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + display_fmt_impl(&self.primitive, K::KIND, K::name(), f) + } +} + +/// Trait used for movedim arguments +pub trait MovedimArgs { + /// Converts into a set of dimensions `Vec` for the `tensor.movedim()` function + fn into_dim_vec(self) -> Vec; +} + +impl MovedimArgs for Vec { + fn into_dim_vec(self) -> Vec { + let set = self + .iter() + .map(|&dim| { + if dim < 0 { + (D as i32 + dim) as usize + } else { + dim as usize + } + }) + .collect::>(); + check!(TensorCheck::movedim_args_vec::(&set)); + + set + } +} + +impl MovedimArgs for Vec { + fn into_dim_vec(self) -> Vec { + check!(TensorCheck::movedim_args_vec::(&self)); + self + } +} + +impl MovedimArgs for usize { + #[allow(clippy::vec_init_then_push)] + fn into_dim_vec(self) -> Vec { + check!(TensorCheck::movedim_args_usize::(self)); + + let mut set = Vec::with_capacity(1); + set.push(self); + + set + } +} + +impl MovedimArgs for i32 { + #[allow(clippy::vec_init_then_push)] + fn into_dim_vec(self) -> Vec { + check!(TensorCheck::movedim_args_i32::(self)); + + let dim = if self < 0 { + (D as i32 + self) as usize + } else { + self as usize + }; + + let mut set = Vec::with_capacity(1); + set.push(dim); + + set + } +} + +/// Trait used for reshape arguments. +pub trait ReshapeArgs: Debug { + /// Converts to a shape. + fn into_shape(self, source: Shape) -> Shape; +} + +impl ReshapeArgs for [I; D2] { + fn into_shape(self, source: Shape) -> Shape { + unwrap_shape_reshape(source.reshape(self)) + } +} + +impl ReshapeArgs for Shape { + fn into_shape(self, source: Shape) -> Shape { + unwrap_shape_reshape(source.reshape(self)) + } +} + +/// Trait used for broadcast arguments. +pub trait BroadcastArgs { + /// Converts to a shape. + fn into_shape(self, shape: &Shape) -> Shape; +} + +impl BroadcastArgs for Shape { + fn into_shape(self, _shape: &Shape) -> Shape { + self + } +} + +impl BroadcastArgs for [E; D2] { + // Passing -1 as the size for a dimension means not changing the size of that dimension. + fn into_shape(self, shape: &Shape) -> Shape { + if self.len() < shape.num_dims() { + panic!( + "Broadcast arguments must be greater than the number of dimensions! got {}, need at least {}", + self.len(), + shape.num_dims() + ); + } + + // Zip the two shapes in reverse order and replace -1 with the actual dimension value. + let new_shape: Vec<_> = self + .iter() + .rev() + .map(|x| { + let primitive = x.as_index(); + if primitive < -1 || primitive == 0 { + panic!( + "Broadcast arguments must be positive or -1! Got {}", + primitive + ); + } + primitive + }) + .zip(shape.iter().rev().chain(repeat(&0)).take(self.len())) // Pad the original shape with 0s + .map(|(x, &y)| if x == -1 { y } else { x as usize }) + .collect::>() + .into_iter() + .rev() + .collect(); + + if new_shape.contains(&0) { + panic!( + "Cannot substitute -1 for a non-existing dimension! Got {:?}", + new_shape + ); + } + + let new_shape: [usize; D2] = new_shape.try_into().unwrap(); + + Shape::from(new_shape) + } +} + +impl Serialize for Tensor +where + K: Basic, +{ + fn serialize(&self, serializer: S) -> Result { + let data = self.to_data(); + data.serialize(serializer) + } +} + +impl<'de, const D: usize, K> Deserialize<'de> for Tensor +where + K: Basic, +{ + fn deserialize>(deserializer: De) -> Result { + let tensor = Tensor::from_data(TensorData::deserialize(deserializer)?, &Device::default()); + Ok(tensor) + } +} + +/// Non-generic outline of `into_data_async`. The public method just calls this +/// helper, so its monomorphization (per `D`/`K`) is trivial — the heavy async +/// state-machine code lives here, compiled once inside `burn-tensor`. +async fn into_data_async_impl( + primitive: BridgeTensor, + kind: crate::ops::Kind, +) -> Result { + use crate::ops::{BasicOps, Kind}; + match kind { + Kind::Float => ::into_data_async(primitive).await, + Kind::Int => ::into_data_async(primitive).await, + Kind::Bool => ::into_data_async(primitive).await, + } +} + +fn slice_bridge_by_kind(p: BridgeTensor, slices: &[Slice], kind: crate::ops::Kind) -> BridgeTensor { + use crate::ops::{BasicOps, Kind}; + match kind { + Kind::Float => ::slice(p, slices), + Kind::Int => ::slice(p, slices), + Kind::Bool => ::slice(p, slices), + } +} + +#[allow(clippy::too_many_arguments)] +fn display_fmt_inner( + primitive: &BridgeTensor, + kind: crate::ops::Kind, + acc: &mut String, + depth: usize, + multi_index: &mut [usize], + range: (usize, usize), + precision: Option, + dims: &[usize], +) { + let (start, end) = range; + let rank = dims.len(); + for i in start..end { + if i > 0 { + acc.push_str(", "); + } + multi_index[depth] = i; + let slices: Vec = (0..rank) + .map(|d| Slice::from((multi_index[d] as i64)..((multi_index[d] + 1) as i64))) + .collect(); + let sliced = slice_bridge_by_kind(primitive.clone(), &slices, kind); + let data = burn_std::reader::try_read_sync(into_data_async_impl(sliced, kind)); + if let Some(Ok(data)) = data { + let elem = DataIterFmt { data, precision }.next(); + acc.push_str(&elem); + } else { + acc.push_str(""); + } + } +} + +fn push_newline_indent_impl(acc: &mut String, indent: usize) { + acc.push('\n'); + for _ in 0..indent { + acc.push(' '); + } +} + +#[allow(clippy::too_many_arguments)] +fn display_fmt_outer( + primitive: &BridgeTensor, + kind: crate::ops::Kind, + acc: &mut String, + depth: usize, + multi_index: &mut [usize], + print_options: &PrintOptions, + summarize: bool, + range: (usize, usize), + dims: &[usize], +) { + let (start, end) = range; + for i in start..end { + if i > start { + acc.push(','); + push_newline_indent_impl(acc, depth + 1); + } + acc.push('['); + multi_index[depth] = i; + display_fmt_recursive( + primitive, + kind, + acc, + depth + 1, + multi_index, + print_options, + summarize, + dims, + ); + acc.push(']'); + } +} + +#[allow(clippy::too_many_arguments)] +fn display_fmt_recursive( + primitive: &BridgeTensor, + kind: crate::ops::Kind, + acc: &mut String, + depth: usize, + multi_index: &mut [usize], + print_options: &PrintOptions, + summarize: bool, + dims: &[usize], +) { + let edge_items = print_options.edge_items; + + if depth == 0 { + acc.push('['); + } + + if depth == dims.len() - 1 { + if summarize && dims[depth] > 2 * edge_items { + display_fmt_inner( + primitive, + kind, + acc, + depth, + multi_index, + (0, edge_items), + print_options.precision, + dims, + ); + acc.push_str(", ..."); + display_fmt_inner( + primitive, + kind, + acc, + depth, + multi_index, + (dims[depth] - edge_items, dims[depth]), + print_options.precision, + dims, + ); + } else { + display_fmt_inner( + primitive, + kind, + acc, + depth, + multi_index, + (0, dims[depth]), + print_options.precision, + dims, + ); + } + } else if summarize && dims[depth] > 2 * edge_items { + display_fmt_outer( + primitive, + kind, + acc, + depth, + multi_index, + print_options, + summarize, + (0, edge_items), + dims, + ); + acc.push(','); + push_newline_indent_impl(acc, depth + 1); + acc.push_str("..."); + push_newline_indent_impl(acc, depth + 1); + display_fmt_outer( + primitive, + kind, + acc, + depth, + multi_index, + print_options, + summarize, + (dims[depth] - edge_items, dims[depth]), + dims, + ); + } else { + display_fmt_outer( + primitive, + kind, + acc, + depth, + multi_index, + print_options, + summarize, + (0, dims[depth]), + dims, + ); + } + + if depth == 0 { + acc.push(']'); + } +} + +fn display_fmt_impl( + primitive: &BridgeTensor, + kind: crate::ops::Kind, + kind_name: &str, + f: &mut core::fmt::Formatter<'_>, +) -> core::fmt::Result { + writeln!(f, "Tensor {{")?; + { + let mut po = { PRINT_OPTS.read().unwrap().clone() }; + if let Some(precision) = f.precision() { + po.precision = Some(precision); + } + let shape = primitive.shape(); + let dims: Vec = shape.iter().copied().collect(); + let mut acc = String::new(); + let mut multi_index = vec![0; dims.len()]; + let num_elements: usize = dims.iter().product(); + let summarize = num_elements > po.threshold; + display_fmt_recursive( + primitive, + kind, + &mut acc, + 0, + &mut multi_index, + &po, + summarize, + &dims, + ); + writeln!(f, " data:")?; + write!(f, "{acc}")?; + writeln!(f, ",")?; + } + writeln!(f, " shape: {},", primitive.shape())?; + let device = match kind { + crate::ops::Kind::Float => ::device(primitive), + crate::ops::Kind::Int => ::device(primitive), + crate::ops::Kind::Bool => ::device(primitive), + }; + writeln!(f, " device: {:?},", device)?; + writeln!(f, " kind: {:?},", kind_name)?; + let dtype = primitive.dtype(); + writeln!(f, " dtype: {:?},", dtype.name())?; + write!(f, "}}") +} + +fn try_into_data_sync_impl( + primitive: BridgeTensor, + kind: crate::ops::Kind, +) -> Result { + crate::try_read_sync(into_data_async_impl(primitive, kind)).expect( + "Failed to read tensor data synchronously. + This can happen on platforms that don't support blocking futures like WASM. + If possible, try using into_data_async instead.", + ) +} + +fn into_data_sync_impl(primitive: BridgeTensor, kind: crate::ops::Kind) -> TensorData { + try_into_data_sync_impl(primitive, kind).expect( + "Error while reading data: use `try_into_data` instead to catch the error at runtime", + ) +} + +#[cfg(test)] +mod tests { + use burn_std::SliceOps; + + use crate::{Shape, s}; + + #[test] + fn slice_range_single_dim_leading() { + let shape = Shape::new([8, 4]); + + // Half-open range + let slices = shape.clone().into_slices([0..5]); + assert_eq!(slices[0].to_range(8), 0..5); + let slices = shape.clone().into_slices([-3..-1]); + assert_eq!(slices[0].to_range(8), 5..7); + + // Inclusive range + let slices = shape.clone().into_slices([0..=4]); + assert_eq!(slices[0].to_range(8), 0..5); + let slices = shape.clone().into_slices([-2..=-1]); + assert_eq!(slices[0].to_range(8), 6..8); + + // Unbounded start + let slices = shape.clone().into_slices([..3]); + assert_eq!(slices[0].to_range(8), 0..3); + let slices = shape.clone().into_slices([..-5]); + assert_eq!(slices[0].to_range(8), 0..3); + + // Unbounded end + let slices = shape.clone().into_slices([5..]); + assert_eq!(slices[0].to_range(8), 5..8); + let slices = shape.clone().into_slices([-3..]); + assert_eq!(slices[0].to_range(8), 5..8); + + // Full range + let slices = shape.into_slices([..]); + assert_eq!(slices[0].to_range(8), 0..8); + } + + #[test] + fn test_negative_slice_indices() { + use crate::Slice; + + // Test negative indices conversion + let slice: Slice = (-3..-1).into(); + assert_eq!(slice.start, -3); + assert_eq!(slice.end, Some(-1)); + + // Test to_range conversion with size 8 + let range = slice.to_range(8); + assert_eq!(range, 5..7); + + // Test with shape slice + let shape = Shape::new([8, 4]); + let result = shape.clone().into_slices([-3..-1]); + assert_eq!(result[0].to_range(8), 5..7); + + // Test more negative index cases + let slice2: Slice = (-5..).into(); + assert_eq!(slice2.to_range(10), 5..10); + + let slice3: Slice = (..-2).into(); + assert_eq!(slice3.to_range(10), 0..8); + + // Test with s! macro - single dimension returns Slice directly + let slice4 = s![-3..-1]; + assert_eq!(slice4.start, -3); + assert_eq!(slice4.end, Some(-1)); + } + + #[test] + fn slice_range_multi_dim() { + let shape = Shape::new([8, 4]); + + // Multiple ways to provide ranges + let slices = shape.clone().into_slices([0..5, 0..4]); + assert_eq!(slices[0].to_range(8), 0..5); + assert_eq!(slices[1].to_range(4), 0..4); + + let slices = shape.clone().into_slices([0.., 0..]); + assert_eq!(slices[0].to_range(8), 0..8); + assert_eq!(slices[1].to_range(4), 0..4); + + let slices = shape.clone().into_slices([0..=7, 0..=3]); + assert_eq!(slices[0].to_range(8), 0..8); + assert_eq!(slices[1].to_range(4), 0..4); + + let slices = shape.clone().into_slices([0..5, 0..3]); + assert_eq!(slices[0].to_range(8), 0..5); + assert_eq!(slices[1].to_range(4), 0..3); + + let slices = shape.into_slices([0.., 0..]); + assert_eq!(slices[0].to_range(8), 0..8); + assert_eq!(slices[1].to_range(4), 0..4); + } + + #[test] + fn slice_range_multi_dim_index() { + let shape = Shape::new([8, 4]); + + // Indices (single integer) should also convert to correct range + let slices = shape.clone().into_slices([0, 2]); + assert_eq!(slices[0].to_range(8), 0..1); + assert_eq!(slices[1].to_range(4), 2..3); + + let slices = shape.into_slices([-1, -1]); + assert_eq!(slices[0].to_range(8), 7..8); + assert_eq!(slices[1].to_range(4), 3..4); + } + + #[test] + fn slice_range_multi_dim_heterogeneous() { + // Slice macro `s![]` can be used to provide different range types + let shape = Shape::new([8, 4, 2]); + let slice = s![0..5, .., -1]; + let slices = shape.into_slices(slice); + assert_eq!(slices[0].to_range(8), 0..5); + assert_eq!(slices[1].to_range(4), 0..4); + assert_eq!(slices[2].to_range(2), 1..2); + + let shape = Shape::new([8, 4, 2, 3]); + let slice = s![..=4, 0..=3, .., -2..]; + let slices = shape.into_slices(slice); + assert_eq!(slices[0].to_range(8), 0..5); + assert_eq!(slices[1].to_range(4), 0..4); + assert_eq!(slices[2].to_range(2), 0..2); + assert_eq!(slices[3].to_range(3), 1..3); + + let shape = Shape::new([3, 4]); + let slice = s![1..-1, ..]; + let slices = shape.into_slices(slice); + assert_eq!(slices[0].to_range(3), 1..2); + assert_eq!(slices[1].to_range(4), 0..4); + } +} diff --git a/crates/burn-tensor/src/tensor/api/bool.rs b/crates/burn-tensor/src/tensor/api/bool.rs new file mode 100644 index 0000000..2fbe26e --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/bool.rs @@ -0,0 +1,508 @@ +use crate::{Bool, Cast, Device, Int, Shape, Tensor, TensorData, ops::BridgeTensor}; +use alloc::{vec, vec::Vec}; +use burn_backend::ops::BoolTensorOps; +use burn_dispatch::Dispatch; + +use crate::try_read_sync; + +/// The part of the tensor to keep when creating a triangular mask. +enum TriPart { + /// Upper triangular part. + Upper, + + /// Lower triangular part. + Lower, + + /// Diagonal part. + Diagonal, +} + +impl Tensor { + /// Create a boolean tensor from data on the given device. + /// + /// # Arguments + /// + /// * `data` - The tensor data. + /// * `device` - The device on which the tensor will be allocated. + /// + /// # Returns + /// + /// A boolean tensor. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2, Bool>::from_bool([[true, false], [false, true]], &device); + /// println!("{tensor}"); + /// } + /// ``` + pub fn from_bool>(data: A, device: &Device) -> Self { + Self::from_data(data.into(), device) + } + + /// Convert the bool tensor into an int tensor. + /// + /// # Returns + /// + /// An integer tensor where `true` is converted to `1` and `false` to `0`. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool}; + /// + /// fn example() { + /// let device = Default::default(); + /// let bool_tensor = Tensor::<1, Bool>::from_bool([true, false, true], &device); + /// let int_tensor = bool_tensor.int(); + /// println!("{int_tensor}"); // [1, 0, 1] + /// } + /// ``` + pub fn int(self) -> Tensor { + let device = self.device(); + Tensor::new(bool_to_int_impl(self.primitive, device)) + } + + /// Convert the bool tensor into a float tensor. + /// + /// # Returns + /// + /// A float tensor where `true` is converted to `1.0` and `false` to `0.0`. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool}; + /// + /// fn example() { + /// let device = Default::default(); + /// let bool_tensor = Tensor::<1, Bool>::from_bool([true, false, true], &device); + /// let float_tensor = bool_tensor.float(); + /// println!("{float_tensor}"); // [1.0, 0.0, 1.0] + /// } + /// ``` + pub fn float(self) -> Tensor { + let device = self.device(); + Tensor::new(bool_to_float_impl(self.primitive, device)) + } + + /// Converts a bool tensor to the specified data type. + /// + /// Supports casting to [`IntDType`](crate::IntDType) (producing an int tensor) + /// or [`FloatDType`](crate::FloatDType) (producing a float tensor). + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool, IntDType, FloatDType}; + /// + /// fn example() { + /// let device = Default::default(); + /// let bool_tensor = Tensor::<1, Bool>::from_bool([true, false, true], &device); + /// + /// // Cast to int + /// let int_tensor = bool_tensor.clone().cast(IntDType::I64); + /// + /// // Cast to float + /// let float_tensor = bool_tensor.cast(FloatDType::F32); + /// } + /// ``` + #[must_use] + pub fn cast>(self, dtype: T) -> Tensor { + T::cast(self, dtype) + } + + /// Inverses boolean values. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2, Bool>::from_bool([[true, false], [false, true]], &device); + /// let inverted = tensor.bool_not(); + /// println!("{inverted}"); // [[false, true], [true, false]] + /// } + /// ``` + pub fn bool_not(self) -> Self { + Tensor::new(bool_not_impl(self.primitive)) + } + + /// Performs logical and (`&&`) on two boolean tensors. + /// + /// # Arguments + /// + /// * `rhs` - The right-hand side tensor for the AND operation. + /// + /// # Returns + /// + /// A boolean tensor where each element is the result of `self[i] && rhs[i]`. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool}; + /// + /// fn example() { + /// let device = Default::default(); + /// let a = Tensor::<2, Bool>::from_bool([[true, true], [false, false]], &device); + /// let b = Tensor::<2, Bool>::from_bool([[true, false], [true, false]], &device); + /// let result = a.bool_and(b); + /// println!("{result}"); // [[true, false], [false, false]] + /// } + /// ``` + pub fn bool_and(self, rhs: Tensor) -> Tensor { + Tensor::new(bool_and_impl(self.primitive, rhs.primitive)) + } + + /// Performs logical or (`||`) on two boolean tensors. + /// + /// # Arguments + /// + /// * `rhs` - The right-hand side tensor for the OR operation. + /// + /// # Returns + /// + /// A boolean tensor where each element is the result of `self[i] || rhs[i]`. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool}; + /// + /// fn example() { + /// let device = Default::default(); + /// let a = Tensor::<2, Bool>::from_bool([[true, true], [false, false]], &device); + /// let b = Tensor::<2, Bool>::from_bool([[true, false], [true, false]], &device); + /// let result = a.bool_or(b); + /// println!("{result}"); // [[true, true], [true, false]] + /// } + /// ``` + pub fn bool_or(self, rhs: Tensor) -> Tensor { + Tensor::new(bool_or_impl(self.primitive, rhs.primitive)) + } + + /// Performs logical xor (`^`) on two boolean tensors. + /// + /// # Arguments + /// + /// * `rhs` - The right-hand side tensor for the XOR operation. + /// + /// # Returns + /// + /// A boolean tensor where each element is the result of `self[i] ^ rhs[i]`. + /// Returns `true` when exactly one of the operands is `true`. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool}; + /// + /// fn example() { + /// let device = Default::default(); + /// let a = Tensor::<2, Bool>::from_bool([[true, true], [false, false]], &device); + /// let b = Tensor::<2, Bool>::from_bool([[true, false], [true, false]], &device); + /// let result = a.bool_xor(b); + /// println!("{result}"); // [[false, true], [true, false]] + /// } + /// ``` + pub fn bool_xor(self, rhs: Tensor) -> Tensor { + Tensor::new(bool_xor_impl(self.primitive, rhs.primitive)) + } + + /// Compute the indices of `true` elements in the tensor (i.e., non-zero for boolean tensors). + /// + /// # Returns + /// + /// A vector of tensors, one for each dimension of the given tensor, containing the indices of + /// the non-zero elements in that dimension. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2, Bool>::from_bool( + /// [[true, false, true], [false, true, false], [false, true, false]], + /// &device, + /// ); + /// let indices = tensor.nonzero(); + /// println!("{}", indices[0]); // [0, 0, 1, 2] + /// println!("{}", indices[1]); // [0, 2, 1, 1] + /// } + /// ``` + pub fn nonzero(self) -> Vec> { + try_read_sync(self.nonzero_async()) + .expect("Failed to read tensor data synchronously. Try using nonzero_async instead.") + } + + /// Compute the indices of `true` elements in the tensor (i.e., non-zero for boolean tensors). + /// + /// # Returns + /// + /// A vector of tensors, one for each dimension of the given tensor, containing the indices of + /// the non-zero elements in that dimension. + pub async fn nonzero_async(self) -> Vec> { + let indices = self.argwhere_async().await; + + if indices.shape().num_elements() == 0 { + // Return empty vec when all elements are zero + return vec![]; + } + + let dims = indices.shape(); + indices + .chunk(dims[1], 1) + .into_iter() + .map(|t| t.reshape(Shape::new([dims[0]]))) + .collect() + } + + /// Compute the indices of the elements that are true, grouped by element. + /// + /// # Returns + /// + /// A tensor containing the indices of all non-zero elements of the given tensor. Each row in the + /// result contains the indices of a non-zero element. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2, Bool>::from_bool( + /// [[true, false, true], [false, true, false], [false, true, false]], + /// &device, + /// ); + /// let indices = tensor.argwhere(); + /// println!("{indices}"); // [[0, 0], [0, 2], [1, 1], [2, 1]] + /// } + /// ``` + pub fn argwhere(self) -> Tensor<2, Int> { + try_read_sync(self.argwhere_async()) + .expect("Failed to read tensor data synchronously. Try using argwhere_async instead.") + } + + /// Compute the indices of the elements that are true, grouped by element. + /// + /// # Returns + /// + /// A tensor containing the indices of all non-zero elements of the given tensor. Each row in the + /// result contains the indices of a non-zero element. + pub async fn argwhere_async(self) -> Tensor<2, Int> { + let out_dtype = self.device().settings().int_dtype; + let inner = Dispatch::bool_argwhere(self.primitive.into(), out_dtype).await; + Tensor::new(BridgeTensor::int(inner)) + } + + /// Creates a mask for the upper, lower triangle, or diagonal of a matrix, which can be used to + /// fill the specified area with a value. + fn tri_mask>(shape: S, tri_part: TriPart, offset: i64, device: &Device) -> Self { + let shape: Shape = shape.into(); + let height = shape[D - 2]; + let width = shape[D - 1]; + + // Generate row and column index tensors. + let row_indices: Tensor<1, Int> = Tensor::arange(0..height as i64, device); + let col_indices: Tensor<1, Int> = Tensor::arange(0..width as i64, device); + + // Prepare shapes for broadcasting. + let mut row_shape = [1; D]; + row_shape[D - 2] = height; + let mut col_shape = [1; D]; + col_shape[D - 1] = width; + + // Reshape for broadcasting. + let row_broadcast: Tensor = row_indices.reshape(Shape::new(row_shape)); + let col_broadcast = col_indices.reshape(Shape::new(col_shape)); + + // Broadcasting trick to create a matrix that facilitates comparison for mask generation. + let matrix = row_broadcast.clone() - (col_broadcast.clone() - offset); + + // Select the appropriate comparison function based on `tri_part`. + let compare = match tri_part { + TriPart::Upper => Tensor::greater_elem, + TriPart::Lower => Tensor::lower_elem, + TriPart::Diagonal => Tensor::not_equal_elem, + }; + + // Generate and return the mask by applying the comparison to the matrix. + compare(matrix, 0).unsqueeze() + } + + /// Creates a mask for the upper triangle of a matrix, which can be used to fill the specified + /// area with a value. + /// + /// This function generates a boolean tensor representing the mask of the upper triangle of a matrix. + /// + /// # Arguments + /// + /// * `shape`: The shape of the matrix. + /// * `offset`: The offset from the diagonal, where 0 means the diagonal, and positive values shift + /// towards the upper triangle. + /// * `device`: The device on which the tensor will be allocated. + /// + /// # Returns + /// + /// Returns a boolean tensor where `false` indicates the elements of the matrix that are part of the + /// upper triangle taking into account the specified `offset`. All other elements are `true`. + /// + /// # Example + /// ```rust + /// use burn_tensor::{Tensor, Bool}; + /// + /// fn example() { + /// let mask = Tensor::<2, Bool>::triu_mask([3, 3], 0, &Default::default()); + /// println!("{mask}"); + /// // [[false, false, false], + /// // [true, false, false], + /// // [true, true, false]] + /// } + /// ``` + pub fn triu_mask>(shape: S, offset: i64, device: &Device) -> Self { + Self::tri_mask(shape, TriPart::Upper, offset, device) + } + + /// Creates a mask for the lower triangle of a matrix, which can be used to fill the specified + /// area with a value. + /// + /// This function generates a boolean tensor representing the mask of the lower triangle of a matrix. + /// + /// # Arguments + /// + /// * `shape`: The shape of the matrix. + /// * `offset`: The offset from the diagonal, where 0 means the diagonal, and negative values shift + /// towards the lower triangle. + /// * `device`: The device on which the tensor will be allocated. + /// + /// # Returns + /// + /// Returns a boolean tensor where `false` indicates the elements of the matrix that are part of the + /// lower triangle taking into account the specified `offset`. All other elements are `true`. + /// + /// # Example + /// ```rust + /// use burn_tensor::{Tensor, Bool}; + /// + /// fn example() { + /// let mask = Tensor::<2, Bool>::tril_mask([3, 3], 0, &Default::default()); + /// println!("{mask}"); + /// // [[false, true, true], + /// // [false, false, true], + /// // [false, false, false]] + /// } + /// ``` + pub fn tril_mask>(shape: S, offset: i64, device: &Device) -> Self { + Self::tri_mask(shape, TriPart::Lower, offset, device) + } + + /// Creates a mask for the diagonal of a matrix, which can be used to fill the specified + /// area with a value. + /// + /// This function generates a boolean tensor representing the mask of the diagonal of a matrix. + /// + /// # Arguments + /// + /// * `shape`: The shape of the matrix. + /// * `offset`: The offset from the diagonal, where 0 means the diagonal, and positive values shift + /// towards the upper triangle. + /// * `device`: The device on which the tensor will be allocated. + /// + /// # Returns + /// + /// Returns a boolean tensor where `false` indicates the elements of the matrix that are part of the + /// diagonal. All other elements are `true`. + /// + /// # Example + /// ```rust + /// use burn_tensor::{Tensor, Bool}; + /// + /// fn example() { + /// let mask = Tensor::<2, Bool>::diag_mask([3, 3], 0, &Default::default()); + /// println!("{mask}"); + /// // [[false, true, true], + /// // [true, false, true], + /// // [true, true, false]] + /// } + /// ``` + pub fn diag_mask>(shape: S, offset: i64, device: &Device) -> Self { + Self::tri_mask(shape, TriPart::Diagonal, offset, device) + } +} + +// !tensor (bool only) +impl core::ops::Not for Tensor { + type Output = Tensor; + + fn not(self) -> Self::Output { + self.bool_not() + } +} + +// tensor & tensor (bool only) +impl core::ops::BitAnd for Tensor { + type Output = Tensor; + + fn bitand(self, tensor: Tensor) -> Self::Output { + self.bool_and(tensor) + } +} + +// tensor | tensor (bool only) +impl core::ops::BitOr for Tensor { + type Output = Tensor; + + fn bitor(self, tensor: Tensor) -> Self::Output { + self.bool_or(tensor) + } +} + +// tensor ^ tensor (bool only) +impl core::ops::BitXor for Tensor { + type Output = Tensor; + + fn bitxor(self, tensor: Tensor) -> Self::Output { + self.bool_xor(tensor) + } +} + +// ========================================================================= +// Non-generic implementation helpers (outlined from the generic API). +// See the crate-level docs for the rationale behind this pattern. +// ========================================================================= + +fn bool_to_int_impl(p: BridgeTensor, device: Device) -> BridgeTensor { + let out_dtype = device.settings().int_dtype; + BridgeTensor::int(Dispatch::bool_into_int(p.into(), out_dtype)) +} + +fn bool_to_float_impl(p: BridgeTensor, device: Device) -> BridgeTensor { + let out_dtype = device.settings().float_dtype; + BridgeTensor::float(Dispatch::bool_into_float(p.into(), out_dtype)) +} + +fn bool_not_impl(p: BridgeTensor) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_not(p.into())) +} + +fn bool_and_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_and(lhs.into(), rhs.into())) +} + +fn bool_or_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_or(lhs.into(), rhs.into())) +} + +fn bool_xor_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_xor(lhs.into(), rhs.into())) +} diff --git a/crates/burn-tensor/src/tensor/api/cartesian_grid.rs b/crates/burn-tensor/src/tensor/api/cartesian_grid.rs new file mode 100644 index 0000000..1a9c32e --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/cartesian_grid.rs @@ -0,0 +1,56 @@ +use crate::{Device, Int, Shape, Tensor}; +use alloc::vec::Vec; + +/// Generates a cartesian grid for the given tensor shape on the specified device. +/// The generated tensor is of dimension `D2 = D + 1`, where each element at dimension D contains the cartesian grid coordinates for that element. +/// +/// # Arguments +/// +/// * `shape` - The shape specifying the dimensions of the tensor. +/// * `device` - The device to create the tensor on. +/// +/// # Panics +/// +/// Panics if `D2` is not equal to `D+1`. +/// +/// # Examples +/// +/// ```rust +/// use burn_tensor::Int; +/// use burn_tensor::{Shape, Tensor}; +/// fn example() { +/// let device = Default::default(); +/// let result: Tensor<3, _> = Tensor::<2, Int>::cartesian_grid([2, 3], &device); +/// println!("{}", result); +/// } +/// ``` +pub fn cartesian_grid, const D: usize, const D2: usize>( + shape: S, + device: &Device, +) -> Tensor { + if D2 != D + 1 { + panic!("D2 must equal D + 1 for Tensor::cartesian_grid") + } + + let dims = shape.into(); + let mut indices: Vec> = Vec::new(); + + for dim in 0..D { + let dim_range: Tensor<1, Int> = Tensor::arange(0..dims[dim] as i64, device); + + let mut shape = [1; D]; + shape[dim] = dims[dim]; + let mut dim_range = dim_range.reshape(shape); + + for (i, &item) in dims.iter().enumerate() { + if i == dim { + continue; + } + dim_range = dim_range.repeat_dim(i, item); + } + + indices.push(dim_range); + } + + Tensor::stack::(indices, D) +} diff --git a/crates/burn-tensor/src/tensor/api/cast.rs b/crates/burn-tensor/src/tensor/api/cast.rs new file mode 100644 index 0000000..62ac6a0 --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/cast.rs @@ -0,0 +1,144 @@ +use burn_backend::ops::{BoolTensorOps, FloatTensorOps, IntTensorOps}; +use burn_backend::{DType, FloatDType, IntDType}; +use burn_dispatch::Dispatch; + +use crate::kind::Basic; +use crate::ops::BridgeTensor; +use crate::{Bool, Float, Int, Tensor}; + +/// Trait for types that represent a valid cast target from a tensor of kind `K`. +/// +/// The generic parameter `K` is the *input* tensor kind ([`Float`], [`Int`], or [`Bool`]). +/// Implementors declare the output kind and provide the actual cast logic. +pub trait Cast { + /// The output tensor kind after casting. + type OutputKind: Basic; + + /// Cast a tensor primitive to the target dtype. + fn cast(tensor: Tensor, dtype: Self) -> Tensor; +} + +// --- Float input impls --- + +impl Cast for FloatDType { + type OutputKind = Float; + + fn cast(tensor: Tensor, dtype: Self) -> Tensor { + if tensor.primitive.is_float() { + let current: FloatDType = tensor.dtype().into(); + if current == dtype { + return tensor; + } + Tensor::new(float_cast_impl(tensor.primitive, dtype)) + } else { + panic!("Should be Float primitive kind"); + } + } +} + +impl Cast for IntDType { + type OutputKind = Int; + + fn cast(tensor: Tensor, dtype: Self) -> Tensor { + Tensor::new(float_to_int_impl(tensor.primitive, dtype)) + } +} + +/// Backward-compatible impl: only float `DType` variants are accepted. +/// +/// # Panics +/// +/// Panics if `dtype` is not a float variant (e.g., `DType::I32`). +/// Use [`IntDType`] directly for cross-kind casting to int. +impl Cast for DType { + type OutputKind = Float; + + fn cast(tensor: Tensor, dtype: Self) -> Tensor { + let float_dtype: FloatDType = dtype.into(); + >::cast(tensor, float_dtype) + } +} + +// --- Int input impls --- + +impl Cast for IntDType { + type OutputKind = Int; + + fn cast(tensor: Tensor, dtype: Self) -> Tensor { + let current: IntDType = tensor.primitive.dtype().into(); + if current == dtype { + return tensor; + } + Tensor::new(int_cast_impl(tensor.primitive, dtype)) + } +} + +impl Cast for FloatDType { + type OutputKind = Float; + + fn cast(tensor: Tensor, dtype: Self) -> Tensor { + Tensor::new(int_to_float_impl(tensor.primitive, dtype)) + } +} + +/// Backward-compatible impl: only int `DType` variants are accepted. +/// +/// # Panics +/// +/// Panics if `dtype` is not an int variant (e.g., `DType::F32`). +/// Use [`FloatDType`] directly for cross-kind casting to float. +impl Cast for DType { + type OutputKind = Int; + + fn cast(tensor: Tensor, dtype: Self) -> Tensor { + let int_dtype: IntDType = dtype.into(); + >::cast(tensor, int_dtype) + } +} + +// --- Bool input impls --- + +impl Cast for IntDType { + type OutputKind = Int; + + fn cast(tensor: Tensor, dtype: Self) -> Tensor { + Tensor::new(bool_cast_to_int_impl(tensor.primitive, dtype)) + } +} + +impl Cast for FloatDType { + type OutputKind = Float; + + fn cast(tensor: Tensor, dtype: Self) -> Tensor { + Tensor::new(bool_cast_to_float_impl(tensor.primitive, dtype)) + } +} + +// ========================================================================= +// Non-generic implementation helpers (outlined from the generic API). +// See the crate-level docs for the rationale behind this pattern. +// ========================================================================= + +fn float_cast_impl(p: BridgeTensor, dtype: FloatDType) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_cast(p.into_float(), dtype)) +} + +fn float_to_int_impl(p: BridgeTensor, dtype: IntDType) -> BridgeTensor { + BridgeTensor::int(Dispatch::float_into_int(p.into_float(), dtype)) +} + +fn int_cast_impl(p: BridgeTensor, dtype: IntDType) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_cast(p.into(), dtype)) +} + +fn int_to_float_impl(p: BridgeTensor, dtype: FloatDType) -> BridgeTensor { + BridgeTensor::float(Dispatch::int_into_float(p.into(), dtype)) +} + +fn bool_cast_to_int_impl(p: BridgeTensor, dtype: IntDType) -> BridgeTensor { + BridgeTensor::bool(Dispatch::bool_into_int(p.into(), dtype)) +} + +fn bool_cast_to_float_impl(p: BridgeTensor, dtype: FloatDType) -> BridgeTensor { + BridgeTensor::float(Dispatch::bool_into_float(p.into(), dtype)) +} diff --git a/crates/burn-tensor/src/tensor/api/check.rs b/crates/burn-tensor/src/tensor/api/check.rs new file mode 100644 index 0000000..7e81351 --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/check.rs @@ -0,0 +1,1753 @@ +use crate::bridge::{BasicOps, Ordered}; +use crate::{DType, Shape, Slice, Tensor, cast::ToElement}; +use alloc::format; +use alloc::string::{String, ToString}; +use alloc::vec; +use alloc::vec::Vec; + +/// The struct should always be used with the [check](crate::check) macro. +/// +/// This is a simple pub(crate) data structure that efficiently checks tensor operations and +/// formats clear error messages. It's crucial that the checks are really fast, but it doesn't matter +/// when a failed check is discovered since the program will panic. +/// +/// # Notes +/// +/// Failing tensor checks will always result in a panic. +/// As mentioned in [The Rust Programming Language book](https://doc.rust-lang.org/book/ch09-03-to-panic-or-not-to-panic.html), +/// when there is no way to recover, panic should be used instead of a result. +/// +/// Most users will unwrap the results anyway, which will worsen the clarity of the code. Almost +/// all checks highlight programming errors, which means invalid programs that should be fixed. +/// Checks are not the ideal way to help users write correct programs, but they are still better +/// than backend errors. Other forms of compile-time validation could be developed, such as named +/// tensors, but we have to carefully evaluate the ease of use of the Tensor API. Adding overly +/// complex type validation checks might drastically worsen the API and result in harder-to-maintain +/// programs. +/// +/// # Design +/// +/// Maybe the Backend API should return a result for each operation, which would allow handling +/// all checks, even the ones that can't be efficiently checked before performing an operation, +/// such as the `index_select` operation. The downside of that approach is that all backend +/// implementation might re-implement the same checks, which may result in unnecessary code +/// duplication. Maybe a combination of both strategies could help to cover all use cases. +pub(crate) enum TensorCheck { + Ok, + Failed(FailedTensorCheck), +} + +impl TensorCheck { + /// Checks device and shape compatibility for element wise binary operations. + pub(crate) fn binary_ops_ew( + ops: &str, + lhs: &Tensor, + rhs: &Tensor, + ) -> Self { + Self::Ok + .binary_ops_device(ops, &lhs.device(), &rhs.device()) + .binary_ops_ew_shape::(ops, &lhs.shape(), &rhs.shape()) + } + + pub(crate) fn into_scalar(shape: &Shape) -> Self { + let mut check = Self::Ok; + + if shape.num_elements() != 1 { + check = check.register( + "Into Scalar", + TensorError::new("Only tensors with 1 element can be converted into scalar.") + .details(format!( + "Current tensor has {} elements", + shape.num_elements() + )), + ); + } + + check + } + + pub(crate) fn dim_ops(ops: &str, dim: usize) -> Self { + let mut check = Self::Ok; + + if dim >= D { + check = check.register( + ops, + TensorError::new("Given dimension is higher than the tensor rank.") + .details(format!("Tensor rank: '{D}', given dimension: '{dim}'.")), + ); + } + + check + } + + pub(crate) fn creation_ops(ops: &str, dims: &[usize]) -> Self { + let mut check = Self::Ok; + + if D == 0 { + check = check.register( + ops, + TensorError::new("Tried to create a 0-dim tensor, which is invalid.") + .details(format!("Tensor rank: '{D}', given dimensions: '{dims:?}'.")), + ); + } + + if dims.len() != D { + check = check.register( + ops, + TensorError::new("Given dimensions differ from the tensor rank.") + .details(format!("Tensor rank: '{D}', given dimensions: '{dims:?}'.")), + ); + } + + check + } + + pub(crate) fn narrow( + tensor: &Tensor, + dim: usize, + start: usize, + length: usize, + ) -> Self { + let mut check = Self::Ok; + + if length == 0 { + check = check.register( + "Narrow", + TensorError::new(format!( + "Can't narrow at dimension {dim}, length must be greater than 0", + )), + ); + } + + if start >= tensor.shape()[dim] { + check = check.register( + "Narrow", + TensorError::new(format!( + "Can't narrow at dimension {dim}, start exceeds the size of the tensor along \ + this dimension (Size={})", + tensor.shape()[dim] + )), + ); + } + + if start + length > tensor.shape()[dim] { + check = check.register( + "Narrow", + TensorError::new(format!( + "Can't narrow at dimension {dim}, start + length exceeds the size of the tensor \ + along this dimension (Size={})", + tensor.shape()[dim] + )), + ); + } + + check + } + + pub(crate) fn movedim_args_usize(dim: usize) -> Self { + let mut check = Self::Ok; + + if dim >= D { + check = check.register( + "Movedim", + TensorError::new( + "The given dimension exceeds the number of dimensions of the current tensor.", + ) + .details(format!( + "Current tensor has {D} dimensions, but the given dimension is {dim}.", + )), + ); + } + + check + } + + pub(crate) fn movedim_args_i32(dim: i32) -> Self { + let mut check = Self::Ok; + + if dim < -(D as i32) || dim >= D as i32 { + check = check.register( + "Movedim", + TensorError::new( + "The given dimension is out of bounds for the current tensor dimensions.", + ) + .details(format!( + "Current tensor has {D} dimensions, but the given dimension is {dim}.", + )), + ); + } + + check + } + + pub(crate) fn movedim_args_vec(dims: &Vec) -> Self { + let mut check = Self::Ok; + + // Check out of bounds + if dims.iter().any(|&x| x >= D) { + check = check.register( + "Movedim", + TensorError::new("The given dimensions are out of bounds.").details(format!( + "Current tensor has {D} dimensions, but the given dimensions are {dims:?}.", + )), + ); + } + + // Check there are no duplicates + for (i, &dim_i) in dims.iter().enumerate() { + for &dim_j in dims.iter().skip(i + 1) { + if dim_i == dim_j { + check = check.register( + "Movedim", + TensorError::new("The given dimensions contain duplicates.").details( + format!( + "The dimension {dim_i} is duplicated in the given dimensions {dims:?}.", + ), + ), + ); + } + } + } + + check + } + + pub(crate) fn movedim_args_length( + source_dims: &Vec, + destination_dims: &Vec, + ) -> Self { + let mut check = Self::Ok; + + if source_dims.len() != destination_dims.len() { + check = check.register( + "Movedim", + TensorError::new( + "The number of dimensions in source and destination must be equal.", + ) + .details(format!( + "Source dimensions: {source_dims:?}, Destination dimensions: {destination_dims:?}.", + )), + ) + } + + check + } + + pub(crate) fn flatten( + start_dim: usize, + end_dim: usize, + ) -> Self { + let mut check = Self::Ok; + + if start_dim > end_dim { + check = check.register( + "Flatten", + TensorError::new(format!( + "The start dim ({start_dim}) must be smaller than or equal to the end dim ({end_dim})" + )), + ); + } + + if D2 > D1 { + check = check.register( + "Flatten", + TensorError::new(format!( + "Result dim ({D2}) must be smaller than or equal to ({D1})" + )), + ); + } + + if D1 < end_dim + 1 { + check = check.register( + "Flatten", + TensorError::new(format!( + "The end dim ({end_dim}) must be smaller than the tensor dim ({D1})" + )), + ); + } + + if (D2 as i32) < (D1 as i32 - (end_dim as i32 - start_dim as i32)) { + check = check.register( + "Flatten", + TensorError::new(format!( + "The destination dimension ({D2}) must be large enough to accommodate the \ + flattening operation." + )), + ); + } + + check + } + + pub(crate) fn tri() -> Self { + let mut check = Self::Ok; + + if D < 2 { + check = check.register( + "Tri", + TensorError::new(format!( + "The input tensor must have at least 2 dimensions, got {D}" + )), + ); + } + + check + } + + pub(crate) fn squeeze(dim: usize, tensor_dims: &[usize]) -> Self { + let mut check = Self::Ok; + // This should actually be to check that the dimension to squeeze + // has a size of 1 + if tensor_dims[dim] != 1 { + check = check.register( + "Squeeze", + TensorError::new(format!( + "Can't squeeze dimension {dim} because its size is not 1", + )), + ); + } + + if dim >= tensor_dims.len() { + check = check.register( + "Squeeze", + TensorError::new(format!( + "Dimension index {dim} is out of bounds for tensor dimensions {tensor_dims:?}.", + )), + ); + } + + check + } + + pub(crate) fn squeeze_dims_input( + dim_indices: &[usize], + current_dims: &[usize], + ) -> Self { + let mut check = Self::Ok; + if dim_indices.len() >= current_dims.len() { + check = check.register( + "Squeeze", + TensorError::new("Attempted to squeeze too many dimensions!").details(format!( + "Got {} dims, tensor has {}", + dim_indices.len(), + current_dims.len() + )), + ); + } + + check + } + + pub(crate) fn squeeze_dims_len(new_dims_len: usize) -> Self { + let mut check = Self::Ok; + if new_dims_len == 0 { + // 0-dim tensor not supported + check = check.register( + "Squeeze", + TensorError::new( + "Resulting dimensions cannot be zero. To remove specific singleton dimensions while preserving at least one, use `squeeze_dims` instead.".to_string() + ), + ); + } + + if new_dims_len != D2 { + check = check.register( + "Squeeze", + TensorError::new(format!( + "Resulting dimensions {new_dims_len} do not match the required D2 size {D2}.", + )), + ); + } + + check + } + + pub(crate) fn unsqueeze() -> Self { + let mut check = Self::Ok; + if D2 < D1 { + check = check.register( + "Unsqueeze", + TensorError::new(format!( + "Can't unsqueeze smaller tensor, got dim {D2}, expected > {D1}", + )), + ); + } + + check + } + + pub(crate) fn unsqueeze_dim(dim: usize) -> Self { + let mut check = Self::Ok; + if D2 <= D1 { + check = check.register( + "Unsqueeze", + TensorError::new(format!( + "The unsqueezed rank must be greater than the input rank (D={D1}; D2={D2})", + )), + ); + } + + if dim > D1 { + check = check.register( + "Unsqueeze", + TensorError::new(format!( + "Can't unsqueeze at dimension {dim}, exceeds tensor dimensions (D={D1})", + )), + ); + } + + if dim >= D2 { + check = check.register( + "Unsqueeze", + TensorError::new(format!( + "Can't unsqueeze at dimension {dim}, exceeds output tensor dimensions (D2={D2})", + )), + ); + } + + check + } + + pub(crate) fn unsqueeze_dims(dim: isize) -> Self { + let mut check = Self::Ok; + let output_rank = D as isize; + //contains is right exclusive, so this is to spec + if !(-output_rank..output_rank).contains(&dim) { + check = check.register( + "Unsqueeze", + TensorError::new(format!( + "unsqueeze arg {dim} is out of range for the output tensor of rank {output_rank}", + )), + ); + } + check + } + + pub(crate) fn one_hot_tensor( + index_tensor: Tensor, + num_classes: usize, + ) -> Self { + let mut check = Self::Ok; + if index_tensor + .clone() + .greater_equal_elem(num_classes as i32) + .any() + .into_scalar::() + .to_bool() + { + check = check.register( + "One Hot", + TensorError::new(format!( + "Can't create a one hot tensor from ({index_tensor:?}) containing indexes greater or equal to the number of classes ({num_classes})", + )), + ); + } else if num_classes <= 1 { + check = check.register( + "One Hot", + TensorError::new("Can't create a one hot tensor with less than 2 classes"), + ) + } + check + } + + pub(crate) fn one_hot_tensor_rank() -> Self { + let mut check = Self::Ok; + if D + 1 != D2 { + check = check.register( + "One Hot", + TensorError::new( + "The one-hot tensor rank must correspond to the rank of the tensor + 1", + ) + .details(format!("Expected D2={}, got {D2}", D + 1)), + ); + } + check + } + + pub(crate) fn swap_dims(dim1: usize, dim2: usize) -> Self { + let mut check = Self::Ok; + + if dim1 > D || dim2 > D { + check = check.register( + "Swap Dims", + TensorError::new("The swap dimensions must be smaller than the tensor dimension") + .details(format!( + "Swap dims ({dim1}, {dim2}) on tensor with ({D}) dimensions." + )), + ); + } + + check + } + + pub(crate) fn permute(axes: [usize; D]) -> Self { + let check = Self::Ok; + + // Check if the axes are within the tensor dimensions + if let Some(axis) = axes.iter().find(|&x| *x >= D) { + return check.register( + "permute", + TensorError::new("The axes must be smaller than the tensor dimension.") + .details(format!("The '{axis}' axis is greater than {D} dimensions.")), + ); + } + + // Check if the axes are unique + let mut seen = [false; D]; + axes.iter().for_each(|&x| seen[x] = true); + if seen.iter().any(|&x| !x) { + return check.register( + "permute", + TensorError::new("The axes must be unique.") + .details(format!("The axes '{axes:?}' are not unique.")), + ); + } + + check + } + + pub(crate) fn flip(rank: usize, axes: &[usize]) -> Self { + let check = Self::Ok; + + // Check if the axes are within the tensor dimensions + if let Some(axis) = axes.iter().find(|&x| *x >= rank) { + return check.register( + "flip", + TensorError::new("The axes must be smaller than the tensor dimension.").details( + format!("The '{axis}' axis is greater than {rank} dimensions."), + ), + ); + } + + // Check if the axes are unique + let mut dedup = axes.to_vec(); + dedup.sort_unstable(); + dedup.dedup(); + if dedup.len() != axes.len() { + return check.register( + "flip", + TensorError::new("The axes must be unique.") + .details(format!("The axes '{axes:?}' are not unique.")), + ); + } + + check + } + + pub(crate) fn matmul(lhs: &Tensor, rhs: &Tensor) -> Self + where + K: BasicOps, + { + let mut check = Self::Ok; + + check = check.binary_ops_device("Matmul", &lhs.device(), &rhs.device()); + + if D < 2 { + return check; + } + + let shape_lhs = lhs.shape(); + let shape_rhs = rhs.shape(); + + let dim_lhs = shape_lhs[D - 1]; + let dim_rhs = shape_rhs[D - 2]; + + if dim_lhs != dim_rhs { + check = check.register( + "Matmul", + TensorError::new(format!( + "The inner dimension of matmul should be the same, but got {dim_lhs} and \ + {dim_rhs}." + )) + .details(format!( + "Lhs shape {:?}, rhs shape {:?}.", + shape_lhs, shape_rhs + )), + ); + } + + check + } + + pub(crate) fn cross( + lhs: &Tensor, + rhs: &Tensor, + dim: usize, + ) -> Self + where + K: BasicOps, + { + let mut check = Self::Ok; + + check = check.binary_ops_device("Cross", &lhs.device(), &rhs.device()); + + let shape_lhs = lhs.shape(); + let shape_rhs = rhs.shape(); + + if dim >= D { + check = check.register( + "Cross", + TensorError::new(format!( + "Dimension {dim} is out of bounds for tensors with {D} dimensions." + )), + ); + return check; + } + + let dim_size_lhs = shape_lhs[dim]; + let dim_size_rhs = shape_rhs[dim]; + + if dim_size_lhs != 3 || dim_size_rhs != 3 { + check = check.register( + "Cross", + TensorError::new(format!( + "Cross product requires dimension {dim} to have size 3, but got {dim_size_lhs} and {dim_size_rhs}." + )), + ); + } + + // Check broadcastability of other dimensions + for i in 0..D { + if i != dim { + let l = shape_lhs[i]; + let r = shape_rhs[i]; + if l != r && l != 1 && r != 1 { + check = check.register( + "Cross", + TensorError::new(format!( + "Tensors are not broadcastable along dimension {i}: {l} and {r}." + )), + ); + } + } + } + + check + } + + pub(crate) fn stack( + tensors: &[Tensor], + dim: usize, + ) -> Self { + let mut check = Self::Ok; + + if dim > D1 { + check = check.register( + "Stack", + TensorError::new( + "Can't stack tensors on a dim that exceeds the tensors dimension (inclusive)", + ) + .details(format!( + "Trying to concatenate tensors with {D1} dimensions on axis {dim}." + )), + ); + } + + if D1 == D2 { + check = check.register( + "Stack", + TensorError::new(format!( + "Can't stack tensors on existing dimension {dim}, the input and output ranks are the same (D={D1}; D2={D2}).\ + If you want to concatenate the tensors along the specified dimension ({dim}), use `Tensor::cat` instead.", + )), + ); + } + + if tensors.is_empty() { + return check.register( + "Stack", + TensorError::new("Can't stack an empty list of tensors."), + ); + } + + let shape_reference = tensors.first().unwrap().shape(); + + for tensor in tensors { + let shape = tensor.shape(); + + if shape_reference != shape { + return check.register( + "Stack", + TensorError::new("Can't stack tensors with different shapes").details(format!( + "Provided dimension ({dim}), tensors shapes: {:?}", + tensors.iter().map(Tensor::shape).collect::>() + )), + ); + } + } + + check + } + + pub(crate) fn cat(tensors: &[Tensor], dim: usize) -> Self { + let mut check = Self::Ok; + + if dim >= D { + check = check.register( + "Cat", + TensorError::new( + "Can't concatenate tensors on a dim that exceeds the tensors dimension", + ) + .details(format!( + "Trying to concatenate tensors with {D} dimensions on axis {dim}." + )), + ); + } + + if tensors.is_empty() { + return check.register( + "Cat", + TensorError::new("Can't concatenate an empty list of tensors."), + ); + } + + let mut shape_reference = tensors.first().unwrap().shape(); + shape_reference[dim] = 1; // We want to check every dims except the one where the + // concatenation happens. + + for tensor in tensors { + let mut shape = tensor.shape(); + shape[dim] = 1; // Ignore the concatenate dim. + + if shape_reference != shape { + return check.register( + "Cat", + TensorError::new( + "Can't concatenate tensors with different shapes, except for the provided \ + dimension", + ) + .details(format!( + "Provided dimension ({dim}), tensors shapes: {:?}", + tensors.iter().map(Tensor::shape).collect::>() + )), + ); + } + } + + check + } + + pub(crate) fn slice(shape: &Shape, slices: &[Slice]) -> Self { + let mut check = Self::Ok; + let n_dims_tensor = R; + let n_dims_slices = slices.len(); + + if n_dims_tensor < n_dims_slices { + check = check.register( + "Slice", + TensorError::new( + "The provided slices array has a higher number of dimensions than the current \ + tensor.", + ) + .details(format!( + "The slices array must be smaller or equal to the tensor number of \ + dimensions. Tensor number of dimensions: {n_dims_tensor}, slices array \ + length {n_dims_slices}." + )), + ); + } + + for (i, slice) in slices.iter().enumerate().take(R) { + let d_tensor = shape[i]; + + // Check the raw end value before conversion + if let Some(end) = slice.end + && end > 0 + && end as usize > d_tensor + { + check = check.register( + "Slice", + TensorError::new( + "The provided slice has a range that exceeds the current tensor \ + size.", + ) + .details(format!( + "The slice end index {} exceeds the size of the tensor ({}) at dimension {}. \ + Tensor shape {:?}.", + end, d_tensor, i, shape, + )), + ); + } + + // Empty slices (start >= end) are allowed and produce a tensor with size 0 + // in that dimension. This matches PyTorch behavior and is required for ONNX + // compatibility where dynamic slice ranges may become empty at runtime. + + if slice.step() == 0 { + check = check.register( + "Slice", + TensorError::new("The provided slice has a step of 0.").details(format!( + "The slice at dimension '{i}' has a step of 0. Step must be non-zero.", + )), + ); + } + } + + check + } + + pub(crate) fn slice_assign( + shape: &Shape, + shape_value: &Shape, + slices: &[crate::Slice], + ) -> Self { + let mut check = Self::Ok; + let n_dims_slices = slices.len(); + + if R < n_dims_slices { + check = check.register( + "Slice Assign", + TensorError::new( + "The provided slices array has a higher number of dimensions than the current \ + tensor.", + ) + .details(format!( + "The slices array must be smaller or equal to the tensor number of \ + dimensions. Tensor number of dimensions: {R}, slices array length {n_dims_slices}." + )), + ); + } + + for (i, slice) in slices.iter().enumerate().take(usize::min(R, n_dims_slices)) { + let d_tensor = shape[i]; + let d_tensor_value = shape_value[i]; + let range = slice.to_range(d_tensor); + + if range.end > d_tensor { + check = check.register( + "Range Assign", + TensorError::new( + "The provided slice has a range that exceeds the current tensor \ + size.", + ) + .details(format!( + "The range ({}..{}) exceeds the size of the tensor ({}) at dimension {}. \ + Current tensor shape {:?}, value tensor shape {:?}.", + range.start, range.end, d_tensor, i, shape, shape_value, + )), + ); + } + + // Calculate the number of elements selected with the given step + let num_elements = slice.output_size(d_tensor); + + if num_elements != d_tensor_value { + check = check.register( + "Slice Assign", + TensorError::new( + "The value tensor must match the amount of elements selected with the \ + slices array", + ) + .details(format!( + "The slice with range ({}..{}) and step {} selects {} elements but the value \ + tensor has {} elements at dimension {}. Current tensor shape {:?}, value tensor \ + shape {:?}.", + range.start, + range.end, + slice.step, + num_elements, + d_tensor_value, + i, + shape, + shape_value, + )), + ); + } + + // Note: Empty slices (start >= end with positive step) are handled at the API level + // by returning the original tensor unchanged, so we don't check for them here. + } + + check + } + + pub(crate) fn check_is_power_of_two(shape: &Shape, dim: usize) -> Self { + let mut check = Self::Ok; + let dim_size = shape[dim]; + + if !dim_size.is_power_of_two() { + check = check.register( + "Check Is Power of two", + TensorError::new("The provided dimension size must be a power of two.") + .details(format!("The length of dimension {dim} is {dim_size}.")), + ); + } + + check + } + + pub(crate) fn check_dim(dim: usize) -> Self { + let mut check = Self::Ok; + + if dim >= D { + check = check.register( + "Check Dim", + TensorError::new("The provided dimension exceeds the tensor dimensions.").details( + format!("Tensor has {D} dimensions, but the provided dimension is {dim}."), + ), + ); + } + + check + } + + pub(crate) fn gather(dim: usize, shape: &Shape, shape_indices: &Shape) -> Self { + Self::check_gather_scatter_indices::(Self::Ok, "Gather", dim, shape, shape_indices) + } + + pub(crate) fn scatter( + dim: usize, + shape: &Shape, + shape_indices: &Shape, + shape_value: &Shape, + ) -> Self { + let ops = "Scatter"; + let mut check = + Self::check_gather_scatter_indices::(Self::Ok, ops, dim, shape, shape_indices); + + if shape_indices != shape_value { + check = check.register( + ops, + TensorError::new( + "Indices tensor shape should be the same as the value tensor shape." + .to_string(), + ) + .details(format!( + "The shape differs: {:?} != {:?}", + shape_indices, shape_value + )), + ); + } + + check + } + + pub(crate) fn scatter_nd( + data_shape: &Shape, + indices_shape: &Shape, + values_shape: &Shape, + ) -> Self { + let ops = "ScatterNd"; + let mut check = Self::Ok; + + if M == 0 { + return check.register( + ops, + TensorError::new("Indices tensor must have rank >= 1".to_string()), + ); + } + + if indices_shape.num_elements() == 0 { + return check.register( + ops, + TensorError::new("Indices tensor must not be empty".to_string()), + ); + } + + let k = indices_shape[M - 1]; + + if k > D { + return check.register( + ops, + TensorError::new(format!( + "Last dimension of indices (K={k}) must be <= data rank (D={D})" + )), + ); + } + + let expected_dv = M - 1 + D - k; + if DV != expected_dv { + check = check.register( + ops, + TensorError::new(format!( + "Values rank DV={DV} does not match expected M-1+D-K = {expected_dv}" + )), + ); + } + + // Batch dims: first M-1 dims of values must equal first M-1 dims of indices + for i in 0..(M - 1) { + if values_shape[i] != indices_shape[i] { + check = check.register( + ops, + TensorError::new(format!( + "Batch dimension {i} mismatch: values={} vs indices={}", + values_shape[i], indices_shape[i] + )), + ); + } + } + + // Slice dims: last D-K dims of values must equal last D-K dims of data + for i in 0..(D - k) { + let val_idx = M - 1 + i; + let data_idx = k + i; + if val_idx < DV && data_idx < D && values_shape[val_idx] != data_shape[data_idx] { + check = check.register( + ops, + TensorError::new(format!( + "Slice dimension mismatch at values[{val_idx}]={} vs data[{data_idx}]={}", + values_shape[val_idx], data_shape[data_idx] + )), + ); + } + } + + check + } + + pub(crate) fn gather_nd( + indices_shape: &Shape, + ) -> Self { + let ops = "GatherNd"; + let mut check = Self::Ok; + + if M == 0 { + return check.register( + ops, + TensorError::new("Indices tensor must have rank >= 1".to_string()), + ); + } + + if indices_shape.num_elements() == 0 { + return check.register( + ops, + TensorError::new("Indices tensor must not be empty".to_string()), + ); + } + + let k = indices_shape[M - 1]; + + if k > D { + return check.register( + ops, + TensorError::new(format!( + "Last dimension of indices (K={k}) must be <= data rank (D={D})" + )), + ); + } + + let expected_dv = M - 1 + D - k; + if DV != expected_dv { + check = check.register( + ops, + TensorError::new(format!( + "Output rank DV={DV} does not match expected M-1+D-K = {expected_dv}" + )), + ); + } + + check + } + + pub(crate) fn select(dim: usize) -> Self { + Self::check_select_basic::(Self::Ok, "select", dim) + } + + pub(crate) fn take(dim: usize) -> Self { + let mut check = Self::check_select_basic::(Self::Ok, "Take", dim); + + // Calculate expected output dimensions + // DO = D - 1 + DI (remove 1 dim, add DI dims) + let expected_do = D + DI - 1; + if DO != expected_do { + check = check.register( + "Take", + TensorError::new("Output dimension mismatch").details(format!( + "Expected output dimension {} (D={} + DI={} - 1) but got DO={}", + expected_do, D, DI, DO + )), + ); + } + + check + } + + pub(crate) fn diag() -> Self { + let mut check = Self::Ok; + + if D < 2 { + check = check.register( + "Diag", + TensorError::new( + "Diagonal operations require + tensors with at least 2 dimensions.", + ) + .details(format!( + "Got tensor with {D} dimensions, + expected at least 2" + )), + ); + } + + if DO != D - 1 { + check = check.register( + "Diag", + TensorError::new("Output rank must be input rank minus 1 for diagonal") + .details(format!("Expected output rank {}, got {DO}", D - 1)), + ); + } + + check + } + + pub(crate) fn select_assign( + dim: usize, + shape_indices: &Shape, + shape_value: &Shape, + ) -> Self { + let mut check = Self::check_select_basic::(Self::Ok, "Select Assign", dim); + + if shape_value[dim] != shape_indices[0] { + check = check.register( + "Select Assign", + TensorError::new( + format!( + "Number of indices ({}) should be equal to value tensor dimensions {:?} on axis (dim={dim})", + shape_indices[0], + shape_value + ), + ) + ); + } + + check + } + + fn check_select_basic(mut check: Self, ops: &str, dim: usize) -> Self { + if dim > D { + check = check.register( + ops, + TensorError::new(format!( + "Can't index a tensor with ({D}) dimensions on axis ({dim})" + )), + ); + } + + check + } + fn check_gather_scatter_indices( + mut check: Self, + ops: &str, + dim: usize, + shape: &Shape, + shape_indices: &Shape, + ) -> Self { + if dim > D { + check = check.register( + ops, + TensorError::new(format!( + "Can't index a tensor with ({D}) dimensions on axis ({dim})" + )), + ); + } + + for i in 0..D { + if i == dim { + continue; + } + + let tensor_dim_i = shape[i]; + let indices_dim_i = shape_indices[i]; + + if tensor_dim_i != indices_dim_i { + check = check.register( + ops, + TensorError::new( + "The tensor shape should be the same as the index tensor shape." + .to_string(), + ) + .details(format!( + "The shape differs at dimension {i}: {tensor_dim_i} != {indices_dim_i}" + )), + ); + } + } + + check + } + + pub(crate) fn check_prelu_shape( + shape_tensor: &Shape, + shape_weight: &Shape, + ) -> Self { + let mut check = Self::Ok; + if shape_weight[0] == 1 { + check + } else if D >= 2 { + let channels = shape_tensor[1]; + let num_weights = shape_weight[0]; + if channels != num_weights { + check = check.register( + "PReLu", + TensorError::new( + "Number of channels in input tensor and number of weights must be equal", + ) + .details(format!( + "Got no. of channels: {channels}, no. of weights: {num_weights}", + )), + ); + return check; + } + check + } else { + check = check.register( + "PReLu", + TensorError::new( + "Number of channels in input tensor and number of weights must be equal", + ) + .details(format!( + "Got no. of channels: 1, no. of weights: {}", + shape_weight[0] + )), + ); + check + } + } + + /// Checks aggregate dimension such as mean and sum. + pub(crate) fn aggregate_dim(ops: &str, dim: usize) -> Self { + let mut check = Self::Ok; + + if dim > D { + check = check.register( + ops, + TensorError::new(format!( + "Can't aggregate a tensor with ({D}) dimensions on axis ({dim})" + )), + ); + } + + check + } + + pub(crate) fn sort_dim(ops: &str, dim: usize) -> Self { + let mut check = Self::Ok; + + if dim > D { + check = check.register( + ops, + TensorError::new(format!( + "Can't sort a tensor with ({D}) dimensions on axis ({dim})" + )), + ); + } + + check + } + + pub(crate) fn split( + tensor_dims: &[usize], + split_size: usize, + dim: usize, + ) -> Self { + let mut check = Self::Ok; + let op = "split"; + let tensor_rank = tensor_dims.len(); + + if dim >= tensor_rank { + check = check.register( + op, + TensorError::new("Given dimension is greater than or equal to the tensor rank.") + .details(format!("Tensor rank: '{D}', given dimension: '{dim}'")), + ); + } else { + let tensor_size = tensor_dims[dim]; + if split_size == 0 && tensor_size != 0 { + check = check.register( + op, + TensorError::new("split_size must be greater than 0 unless the tensor size along the dimension is 0.") + .details(format!("split_size: '{split_size}', tensor size along dim '{dim}': '{tensor_size}'.")), + ); + } + } + + check + } + + pub(crate) fn split_with_sizes( + tensor_dims: &[usize], + split_sizes: &[usize], + dim: usize, + ) -> Self { + let mut check = Self::Ok; + let op = "split_with_sizes"; + let tensor_rank = tensor_dims.len(); + + if dim >= tensor_rank { + check = check.register( + op, + TensorError::new("Given dimension is greater than or equal to the tensor rank.") + .details(format!("Tensor rank: '{D}', given dimension: '{dim}'.")), + ); + } else { + // Validate split_sizes add up to size of dimension to split along + let tensor_size = tensor_dims[dim]; + let total_split_size: usize = split_sizes.iter().sum(); + if total_split_size != tensor_size { + check = check.register( + op, + TensorError::new("The sum of split_sizes must equal the tensor size along the specified dimension.") + .details(format!("Sum of split_sizes: '{total_split_size}', tensor size along dim '{dim}': '{tensor_size}'.")), + ); + } + } + + check + } + + /// The goal is to minimize the cost of checks when there are no error, but it's way less + /// important when an error occurred, crafting a comprehensive error message is more important + /// than optimizing string manipulation. + fn register(self, ops: &str, error: TensorError) -> Self { + let errors = match self { + Self::Ok => vec![error], + Self::Failed(mut failed) => { + failed.errors.push(error); + failed.errors + } + }; + + Self::Failed(FailedTensorCheck { + ops: ops.to_string(), + errors, + }) + } + + /// Checks if shapes are compatible for element wise operations supporting broadcasting. + pub(crate) fn binary_ops_ew_shape( + self, + ops: &str, + lhs: &Shape, + rhs: &Shape, + ) -> Self { + let mut check = self; + + for i in 0..D { + let d_lhs = lhs[i]; + let d_rhs = rhs[i]; + + if d_lhs != d_rhs { + let is_broadcast = d_lhs == 1 || d_rhs == 1; + + if is_broadcast { + continue; + } + + check = check.register( + ops, + TensorError::new("The provided tensors have incompatible shapes.").details( + format!( + "Incompatible size at dimension '{}' => '{} != {}', which can't be \ + broadcasted. Lhs tensor shape {:?}, Rhs tensor shape {:?}.", + i, d_lhs, d_rhs, lhs, rhs, + ), + ), + ); + } + } + + check + } + + /// Checks if tensor devices are equal. + fn binary_ops_device( + self, + ops: &str, + lhs: &Device, + rhs: &Device, + ) -> Self { + match lhs != rhs { + true => self.register( + ops, + TensorError::new("The provided tensors are not on the same device.").details( + format!("Lhs tensor device {lhs:?}, Rhs tensor device {rhs:?}.",), + ), + ), + false => self, + } + } + + /// Checks if expand operation is possible for the given shapes. + pub(crate) fn expand( + ops: &str, + shape: &Shape, + to: &Shape, + ) -> Self { + let mut check = TensorCheck::Ok; + let max_dims = core::cmp::max(D1, D2); + + // Calculate the starting indices for each shape array, ensuring alignment from the right. + let start_index_shape = max_dims.saturating_sub(D1); + let start_index_to = max_dims.saturating_sub(D2); + + for i in 0..max_dims { + // Use 1 as the default dimension size for dimensions beyond the tensor's rank. + let d_shape = if i >= start_index_shape { + shape[i - start_index_shape] + } else { + 1 + }; + let d_to = if i >= start_index_to { + to[i - start_index_to] + } else { + 1 + }; + + if d_shape != d_to && d_shape != 1 && d_to != 1 { + // Register an incompatibility error. + check = check.register( + ops, + TensorError::new( + "The provided tensor can't be broadcasted to the target shape.", + ) + .details(format!( + "Incompatible size at dimension '{}' => '{} != {}', which can't be \ + broadcasted. Tensor shape {:?}, Target shape {:?}.", + max_dims - i - 1, + d_shape, + d_to, + shape, + to, + )), + ); + break; // Incompatibility found, no need to check further. + } + } + + check + } + + /// Checks if unfold operation is possible for the given shapes. + pub(crate) fn unfold( + ops: &str, + _shape: &Shape, + _dim: usize, + _size: usize, + _step: usize, + ) -> Self { + let mut check = TensorCheck::Ok; + + if D2 != D1 + 1 { + check = check.register( + ops, + TensorError::new("The unfold rank is incompatible with the input tensor rank.") + .details(format!( + "The output rank '{D2}' != the input rank + 1 '{D1}'.", + )), + ); + } + + check + } + + /// Checks if input is compatible with convolution weights. + pub(crate) fn conv( + ops: &str, + x: [usize; D1], + weight: [usize; D2], + groups: usize, + ) -> Self { + let mut check = TensorCheck::Ok; + let channels = x[1]; + let expected = weight[1] * groups; + if channels != expected { + check = check.register( + ops, + TensorError::new("Number of channels in input tensor and input channels of convolution must be equal.") + .details(format!("got: {channels}, expected: {expected}")), + ); + } + check + } + + /// Checks if input is compatible with transposed convolution weights. + pub fn conv_transpose( + ops: &str, + x: [usize; D1], + weight: [usize; D2], + ) -> Self { + let mut check = TensorCheck::Ok; + let channels = x[1]; + let expected = weight[0]; + if channels != expected { + check = check.register( + ops, + TensorError::new("Number of channels in input tensor and input channels of convolution must be equal.") + .details(format!("got: {channels}, expected: {expected}")), + ); + } + check + } + + /// Check the generic parameters for lu decomposition is valid. + pub fn lu_generic_param(ops: &str) -> Self { + let mut check = TensorCheck::Ok; + if D - 1 != D1 { + check = check.register( + ops, + TensorError::new( + "D - 1 = D1 must hold for the generic parameters of LU decomposition.", + ) + .details(format!("Got generic parameters D = {} and D1 = {}", D, D1)), + ); + } + check + } + + /// Check the input tensor for lu decomposition is valid. + pub fn lu_input_tensor(ops: &str, dims: &[usize], dtype: DType) -> Self { + let mut check = TensorCheck::Ok; + + if matches!(dtype, DType::QFloat(_)) { + check = check.register( + ops, + TensorError::new("The input tensor must have a real float dtype") + .details("Got an input tensor with a quantized float dtype".to_string()), + ); + } + + let n_dims = dims.len(); + if n_dims < 2 { + check = check.register( + ops, + TensorError::new( + "The input tensor for LU decomposition must have at least two dimensions.", + ) + .details(format!("Got input tensor with {} dimensions", n_dims)), + ); + } + + check + } + + /// Check if input tensor and generic parameters of `linalg::det()` are valid. + pub fn det( + dims: [usize; D], + dtype: DType, + ) -> Self { + let mut check = TensorCheck::Ok; + + if matches!(dtype, DType::QFloat(_)) { + check = check.register( + "det", + TensorError::new("The input tensor must have a real float dtype.") + .details("Got an input tensor with a quantized float dtype".to_string()), + ); + } + + if D1 != D - 1 { + check = check.register( + "det", + TensorError::new( + "D - 1 = D1 must hold for the generic parameters of the linalg::det function.", + ) + .details(format!("Got generic parameters D = {D} and D1 = {D1}")), + ); + } + + if D2 != D - 2 { + check = check.register( + "det", + TensorError::new("The output tensor rank must be less than input tensor rank by 2") + .details(format!( + "Got input tensor rank {D} and output tensor rank {D2}" + )), + ); + } + + if D < 3 { + check = check.register( + "det", + TensorError::new(format!( + "The input tensor must have at least 3 dimensions, got {D}" + )), + ); + } + + if dims[D - 1] != dims[D - 2] { + check = check.register( + "det", + TensorError::new("The last two dimensions of the input tensor must be equal") + .details(format!("Got input tensor with shape {:?}", dims)), + ); + } + + check + } +} + +pub(crate) struct FailedTensorCheck { + ops: String, + errors: Vec, +} + +impl FailedTensorCheck { + /// Format all the checks into a single message ready to be printed by a [panic](core::panic). + pub(crate) fn format(self) -> String { + self.errors.into_iter().enumerate().fold( + format!( + "=== Tensor Operation Error ===\n Operation: '{}'\n Reason:", + self.ops + ), + |accum, (number, error)| accum + error.format(number + 1).as_str(), + ) + "\n" + } +} + +struct TensorError { + description: String, + details: Option, +} + +impl TensorError { + pub(crate) fn new>(description: S) -> Self { + TensorError { + description: description.into(), + details: None, + } + } + + pub(crate) fn details>(mut self, details: S) -> Self { + self.details = Some(details.into()); + self + } + + fn format(self, number: usize) -> String { + let mut message = format!("\n {number}. "); + message += self.description.as_str(); + message += " "; + + if let Some(details) = self.details { + message += details.as_str(); + message += " "; + } + + message + } +} + +/// Module where we defined macros that can be used only in the project. +pub(crate) mod macros { + /// We use a macro for all checks, since the panic message file and line number will match the + /// function that does the check instead of a generic error.rs crate private unrelated file + /// and line number. + macro_rules! check { + ($check:expr) => { + if let TensorCheck::Failed(check) = $check { + core::panic!("{}", check.format()); + } + }; + } + pub(crate) use check; +} + +pub(crate) fn unwrap_shape_reshape(result: Result) -> Shape { + match result { + Ok(shape) => shape, + // `shape.reshape(new_shape)` should only return `MetadataError::Invalid`. + Err(burn_std::MetadataError::Invalid { reason }) => { + macros::check!({ + TensorCheck::Ok.register("Reshape", crate::check::TensorError::new(reason)) + }); + unreachable!() + } + Err(e) => panic!("{e:?}"), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use macros::check; + + #[test] + #[should_panic] + fn index_range_exceed_dimension() { + let slices = vec![Slice::from(0..2), Slice::from(0..4), Slice::from(1..8)]; + check!(TensorCheck::slice::<3>(&Shape::new([3, 5, 7]), &slices)); + } + + #[test] + #[should_panic] + fn index_range_exceed_number_of_dimensions() { + let slices = vec![Slice::from(0..1), Slice::from(0..1), Slice::from(0..1)]; + check!(TensorCheck::slice::<2>(&Shape::new([3, 5]), &slices)); + } + + #[test] + #[should_panic] + fn binary_ops_shapes_no_broadcast() { + check!(TensorCheck::binary_ops_ew_shape::<2>( + TensorCheck::Ok, + "TestOps", + &Shape::new([3, 5]), + &Shape::new([3, 6]) + )); + } + + #[test] + fn binary_ops_shapes_with_broadcast() { + check!(TensorCheck::binary_ops_ew_shape::<2>( + TensorCheck::Ok, + "Test", + &Shape::new([3, 5]), + &Shape::new([1, 5]) + )); + } + + #[test] + #[should_panic] + fn binary_ops_devices() { + check!(TensorCheck::binary_ops_device( + TensorCheck::Ok, + "Test", + &5, // We can pass anything that implements PartialEq as device + &8 + )); + } + + #[test] + #[should_panic] + fn movedim_args_out_of_bounds() { + check!(TensorCheck::movedim_args_usize::<3>(5)); + } + + #[test] + fn movedim_args_i32() { + check!(TensorCheck::movedim_args_i32::<3>(-3)); + } + + #[test] + #[should_panic] + fn movedim_args_too_negative() { + check!(TensorCheck::movedim_args_i32::<3>(-4)); + } + + #[test] + #[should_panic] + fn movedim_args_vec_out_of_bounds() { + check!(TensorCheck::movedim_args_vec::<3>(&vec![0, 1, 3])); + } + + #[test] + #[should_panic] + fn movedim_args_vec_duplicates() { + check!(TensorCheck::movedim_args_vec::<3>(&vec![0, 1, 1])); + } + + #[test] + #[should_panic] + fn movedim_args_length() { + check!(TensorCheck::movedim_args_length( + &vec![0, 1], + &vec![0, 1, 2] + )); + } + + #[test] + #[should_panic] + fn unsqueeze_dim_same_rank() { + check!(TensorCheck::unsqueeze_dim::<3, 3>(2)); + } +} diff --git a/crates/burn-tensor/src/tensor/api/extension.rs b/crates/burn-tensor/src/tensor/api/extension.rs new file mode 100644 index 0000000..be550b9 --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/extension.rs @@ -0,0 +1,463 @@ +use burn_backend::{Backend, TensorMetadata}; +pub use burn_dispatch::DispatchTensor; +use burn_dispatch::{BackendTensor, DispatchKindConversion}; +use burn_std::DType; + +use crate::{ + Bool, Float, Int, Tensor, + kind::Basic, + ops::{BridgeKind, BridgeTensor, Kind}, +}; + +use alloc::{format, string::String}; + +impl Tensor +where + K: Basic, +{ + /// Converts the tensor into its bridge-layer representation. + /// + /// This is primarily intended for backend extensions, allowing custom operations + /// to be inserted between the high-level tensor API and the dispatch layer before + /// deferring to a concrete backend. + pub fn into_bridge(self) -> BridgeTensor { + self.primitive + } + + /// Reconstructs a tensor from its [`BridgeTensor`] bridge representation. + /// + /// This is the inverse of [`Tensor::into_bridge`] and is primarily intended + /// for backend extensions, to wrap the output of a custom operation back into + /// the high-level tensor API. + /// + /// # Panics + /// + /// Panics if the [`BridgeTensor`] variant does not match the tensor kind `K` + /// (e.g. passing an int [`BridgeTensor`] when `K` is [`Float`](crate::Float). + pub fn from_bridge(tensor: BridgeTensor) -> Self { + let dtype = tensor.dtype(); + match (tensor.kind(), K::KIND) { + (BridgeKind::Bool, Kind::Bool) if dtype.is_bool() => Self::new(tensor), + (BridgeKind::Int, Kind::Int) if dtype.is_int() || dtype.is_uint() => Self::new(tensor), + (BridgeKind::Float, Kind::Float) if dtype.is_float() => Self::new(tensor), + (BridgeKind::QFloat, Kind::Float) if matches!(dtype, DType::QFloat(_)) => { + Self::new(tensor) + } + (_, kind) => panic!("Expected kind {kind:?}, got dtype {dtype:?}"), + } + } + + /// Converts from a dispatch tensor into a tensor. + /// + /// # Panics + /// Panis if the dispatch dtype does not match the tensor kind `K`. + pub fn from_dispatch(tensor: DispatchTensor) -> Self { + match (tensor.dtype(), K::KIND) { + (DType::QFloat(_), Kind::Float) => Self::new(BridgeTensor::qfloat(tensor)), + (dtype, Kind::Float) if dtype.is_float() => Self::new(BridgeTensor::float(tensor)), + (dtype, Kind::Int) if dtype.is_int() || dtype.is_uint() => { + Self::new(BridgeTensor::int(tensor)) + } + (dtype, Kind::Bool) if dtype.is_bool() => Self::new(BridgeTensor::bool(tensor)), + (dtype, kind) => panic!("Expected kind {kind:?}, got dtype {dtype:?}"), + } + } + + /// Converts the tensor into a dispatch tensor. + pub fn into_dispatch(self) -> DispatchTensor { + self.primitive.into() + } + + /// Safely downcasts a [`Tensor`] into its low-level backend primitive. + /// + /// This is primarily intended for backend extensions where interfacing with the direct backend + /// primitive (e.g., `B::FloatTensorPrimitive` or `AutodiffTensor`) is necessary. + /// + /// # Examples + /// + /// ```rust,ignore + /// // Extract the underlying CubeCL tensor primitive on the `Wgpu` backend + /// let cube_tensor = tensor.try_into_primitive::()?; + /// + /// // For an autodiff tensor, we can get the wrapped CubeCL tensor primitive on the `Autodiff` backend + /// let ad_tensor = tensor.try_into_primitive::>()?; + ///``` + /// + /// # Errors + /// + /// Returns a [`PrimitiveConversionError`] if the tensor does not currently live on the requested + /// backend `B` (including `Autodiff` mismatch). + pub fn try_into_primitive( + self, + ) -> Result<>::Primitive, PrimitiveConversionError> + where + K: BackendPrimitive, + DispatchTensor: DispatchKindConversion, + { + let dispatch = self.primitive.into(); + // `tensor.try_into_backend::>()` returns a `BackendTensor::Float(AutodiffTensor)` + // so it is automatically handled by the `try_into_primitive` impl for `Float` + let tensor = >::try_into_backend(dispatch) + .map_err(PrimitiveConversionError::BackendMismatch)?; + >::try_into_primitive(tensor) + } + + /// Reconstructs a high-level [`Tensor`] from a concrete backend primitive. + /// + /// This is the inverse of [`Tensor::try_into_primitive`]. + /// + /// # Panics + /// Panis if the tensor kind `K` does not match the tensor underlying primitive kind. + pub fn from_primitive(primitive: >::Primitive) -> Self + where + K: BackendPrimitive, + DispatchTensor: DispatchKindConversion, + { + let tensor = >::from_primitive(primitive); + let dispatch = >::from_backend(tensor); + Self::from_dispatch(dispatch) + } +} + +/// Error returned when a [`DispatchTensor`] cannot be converted to the requested concrete primitive type. +#[derive(Debug, PartialEq)] +pub enum PrimitiveConversionError { + /// The dispatch tensor's backend variant does not match the requested backend. + /// + /// For example, extracting a `Wgpu` primitive from a `Cuda` dispatch tensor. + BackendMismatch(String), + /// The tensor kind does not match the requested primitive kind. + /// + /// For example, extracting a float primitive from an int tensor. + KindMismatch(String), +} + +/// Trait to safely extract and wrap backend-specific primitives based on the high-level tensor kind. +/// +/// This trait functions as a type-level map linking frontend kinds (`Float`, `Int`, `Bool`) +/// to their corresponding backend-associated types (`B::FloatTensorPrimitive`, `B::IntTensorPrimitive`, etc.). +pub trait BackendPrimitive { + /// The backend tensor primitive. + type Primitive; + + /// Attempts to unpack the type-erased [`BackendTensor`] enum wrapper into the concrete + /// variant expected by this kind. + /// + /// # Errors + /// + /// Returns an error if there is a variant type mismatch (e.g., attempting to extract an + /// `Int` primitive out of a `BackendTensor::Float` enum variant). + fn try_into_primitive( + tensor: BackendTensor, + ) -> Result; + + /// Wraps a backend tensor primitive into its corresponding `BackendTensor` enum variant. + fn from_primitive(primitive: Self::Primitive) -> BackendTensor; +} + +impl BackendPrimitive for Float { + // NOTE: not implemented for QFloat + type Primitive = B::FloatTensorPrimitive; + + fn try_into_primitive( + tensor: BackendTensor, + ) -> Result { + match tensor { + BackendTensor::Float(t) => Ok(t), + other => Err(PrimitiveConversionError::KindMismatch(format!( + "Expected Float primitive, got variant: {}", + other.name() + ))), + } + } + + fn from_primitive(primitive: Self::Primitive) -> BackendTensor { + BackendTensor::Float(primitive) + } +} + +impl BackendPrimitive for Int { + type Primitive = B::IntTensorPrimitive; + + fn try_into_primitive( + tensor: BackendTensor, + ) -> Result { + match tensor { + BackendTensor::Int(t) => Ok(t), + other => Err(PrimitiveConversionError::KindMismatch(format!( + "Expected Int primitive, got variant: {}", + other.name() + ))), + } + } + + fn from_primitive(primitive: Self::Primitive) -> BackendTensor { + BackendTensor::Int(primitive) + } +} + +impl BackendPrimitive for Bool { + type Primitive = B::BoolTensorPrimitive; + + fn try_into_primitive( + tensor: BackendTensor, + ) -> Result { + match tensor { + BackendTensor::Bool(t) => Ok(t), + other => Err(PrimitiveConversionError::KindMismatch(format!( + "Expected Bool primitive, got variant: {}", + other.name() + ))), + } + } + + fn from_primitive(primitive: Self::Primitive) -> BackendTensor { + BackendTensor::Bool(primitive) + } +} + +#[cfg(test)] +mod tests { + use crate::{Bool, Int}; + + use super::*; + + type TestBackend = burn_dispatch::backends::Flex; + type TestAutodiffBackend = burn_dispatch::backends::Autodiff; + + // -- into_bridge / from_bridge roundtrip -- + + #[test] + fn float_tensor_bridge_roundtrip() { + let tensor = Tensor::<2>::zeros([2, 3], &Default::default()); + let shape = tensor.shape(); + let bridge = tensor.into_bridge(); + assert!(bridge.is_float()); + let tensor = Tensor::<2>::from_bridge(bridge); + assert_eq!(tensor.shape(), shape); + } + + #[test] + fn int_tensor_bridge_roundtrip() { + let tensor = Tensor::<2, Int>::zeros([2, 3], &Default::default()); + let shape = tensor.shape(); + let bridge = tensor.into_bridge(); + assert!(bridge.is_int()); + let tensor = Tensor::<2, Int>::from_bridge(bridge); + assert_eq!(tensor.shape(), shape); + } + + #[test] + fn bool_tensor_bridge_roundtrip() { + let tensor = Tensor::<2, Bool>::empty([2, 3], &Default::default()); + let shape = tensor.shape(); + let bridge = tensor.into_bridge(); + assert!(bridge.is_bool()); + let tensor = Tensor::<2, Bool>::from_bridge(bridge); + assert_eq!(tensor.shape(), shape); + } + + // -- from_bridge panics on kind mismatch -- + + #[test] + #[should_panic(expected = "Expected kind Float")] + fn from_bridge_int_as_float_panics() { + let bridge = Tensor::<2, Int>::zeros([2, 3], &Default::default()).into_bridge(); + let _tensor = Tensor::<2>::from_bridge(bridge); + } + + #[test] + #[should_panic(expected = "Expected kind Float")] + fn from_bridge_bool_as_float_panics() { + let bridge = Tensor::<2, Bool>::empty([2, 3], &Default::default()).into_bridge(); + let _tensor = Tensor::<2>::from_bridge(bridge); + } + + #[test] + #[should_panic(expected = "Expected kind Int")] + fn from_bridge_float_as_int_panics() { + let bridge = Tensor::<2>::zeros([2, 3], &Default::default()).into_bridge(); + let _tensor = Tensor::<2, Int>::from_bridge(bridge); + } + + #[test] + #[should_panic(expected = "Expected kind Bool")] + fn from_bridge_int_as_bool_panics() { + let bridge = Tensor::<2, Int>::zeros([2, 3], &Default::default()).into_bridge(); + let _tensor = Tensor::<2, Bool>::from_bridge(bridge); + } + + #[test] + #[should_panic(expected = "Expected kind Float")] + fn from_bridge_qfloat_variant_with_int_dtype_panics() { + // Construct a QFloat bridge tensor wrapping a non-qfloat dispatch tensor: + // kind tag says Float but dtype says otherwise. + let inner = Tensor::<2, Int>::zeros([2, 3], &Default::default()).into_dispatch(); + let bridge = BridgeTensor::qfloat(inner); + let _tensor = Tensor::<2>::from_bridge(bridge); + } + + // -- into_dispatch / from_dispatch roundtrip -- + + #[test] + fn float_primitive_roundtrip() { + let tensor = Tensor::<2>::zeros([2, 3], &Default::default()); + let shape = tensor.shape(); + let primitive = tensor.into_dispatch(); + let tensor = Tensor::<2>::from_dispatch(primitive); + assert_eq!(tensor.shape(), shape); + } + + #[test] + fn int_primitive_roundtrip() { + let tensor = Tensor::<2, Int>::zeros([2, 3], &Default::default()); + let shape = tensor.shape(); + let primitive = tensor.into_dispatch(); + let tensor = Tensor::<2, Int>::from_dispatch(primitive); + assert_eq!(tensor.shape(), shape); + } + + #[test] + fn bool_primitive_roundtrip() { + let tensor = Tensor::<2, Bool>::empty([2, 3], &Default::default()); + let shape = tensor.shape(); + let primitive = tensor.into_dispatch(); + let tensor = Tensor::<2, Bool>::from_dispatch(primitive); + assert_eq!(tensor.shape(), shape); + } + + // -- from_dispatch panics on dtype/kind mismatch -- + + #[test] + #[should_panic(expected = "Expected kind Float")] + fn from_dispatch_int_dtype_as_float_panics() { + let primitive = Tensor::<2, Int>::zeros([2, 3], &Default::default()).into_dispatch(); + let _tensor = Tensor::<2>::from_dispatch(primitive); + } + + #[test] + #[should_panic(expected = "Expected kind Int")] + fn from_dispatch_float_dtype_as_int_panics() { + let primitive = Tensor::<2>::zeros([2, 3], &Default::default()).into_dispatch(); + let _tensor = Tensor::<2, Int>::from_dispatch(primitive); + } + + #[test] + #[should_panic(expected = "Expected kind Bool")] + fn from_dispatch_float_dtype_as_bool_panics() { + let primitive = Tensor::<2>::zeros([2, 3], &Default::default()).into_dispatch(); + let _tensor = Tensor::<2, Bool>::from_dispatch(primitive); + } + + // -- try_into_primitive / from_primitive roundtrip -- + + #[test] + fn float_backend_primitive_roundtrip() { + let device = Default::default(); + let tensor = Tensor::<2>::zeros([2, 3], &device); + let shape = tensor.shape(); + + let primitive = tensor + .try_into_primitive::() + .expect("Failed to extract Float primitive"); + let tensor = Tensor::<2>::from_primitive::(primitive); + + assert_eq!(tensor.shape(), shape); + } + + #[test] + fn int_backend_primitive_roundtrip() { + let device = Default::default(); + let tensor = Tensor::<2, Int>::zeros([2, 3], &device); + let shape = tensor.shape(); + + let primitive = tensor + .try_into_primitive::() + .expect("Failed to extract Int primitive"); + let tensor = Tensor::<2, Int>::from_primitive::(primitive); + + assert_eq!(tensor.shape(), shape); + } + + #[test] + fn bool_backend_primitive_roundtrip() { + let device = Default::default(); + let tensor = Tensor::<2, Bool>::empty([2, 3], &device); + let shape = tensor.shape(); + + let primitive = tensor + .try_into_primitive::() + .expect("Failed to extract Bool primitive"); + let tensor = Tensor::<2, Bool>::from_primitive::(primitive); + + assert_eq!(tensor.shape(), shape); + } + + #[cfg(feature = "autodiff")] + #[test] + fn autodiff_backend_primitive_roundtrip() { + let device = crate::Device::default().autodiff(); + let tensor = Tensor::<2, Float>::empty([2, 3], &device).require_grad(); + let require_grad = tensor.is_require_grad(); + + let primitive = tensor + .try_into_primitive::() + .expect("Failed to extract Autodiff primitive"); + let tensor = Tensor::<2, Float>::from_primitive::(primitive); + + assert_eq!(tensor.is_require_grad(), require_grad); + } + + // -- try_into_primitive panics on backend or kind mismatch -- + + #[cfg(feature = "autodiff")] + #[test] + fn try_into_primitive_backend_mismatch() { + let device = Default::default(); + let tensor = Tensor::<2>::zeros([2, 3], &device); + + let err = tensor + .try_into_primitive::() + .unwrap_err(); + + assert!(matches!(err, PrimitiveConversionError::BackendMismatch(_))); + assert!(format!("{err:?}").contains("Expected Autodiff tensor, got backend:")); + } + + #[test] + #[should_panic(expected = "Expected kind Float")] + fn try_into_primitive_kind_mismatch() { + let device = Default::default(); + let tensor = Tensor::<2, Int>::zeros([2, 3], &device); + + let primitive = tensor + .try_into_primitive::() + .expect("Failed to extract Int primitive"); + let _tensor = Tensor::<2, Float>::from_primitive::(primitive); + } + + // -- BackendPrimitive trait error handling on mismatched variants -- + + #[test] + fn try_into_primitive_float_with_int_variant_returns_err() { + let device = Default::default(); + + // Valid Int primitive + let int_tensor = Tensor::<2, Int>::zeros([2, 3], &device); + let int_primitive = int_tensor.try_into_primitive::().unwrap(); + + // Wrap it artificially into the Backend layer + let backend_tensor = BackendTensor::::Int(int_primitive); + + // Attempt to extract it as a Float primitive using the BackendPrimitive trait + let err = >::try_into_primitive(backend_tensor) + .unwrap_err(); + + assert_eq!( + err, + PrimitiveConversionError::KindMismatch( + "Expected Float primitive, got variant: Int".into() + ) + ); + } +} diff --git a/crates/burn-tensor/src/tensor/api/float.rs b/crates/burn-tensor/src/tensor/api/float.rs new file mode 100644 index 0000000..a6a7c41 --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/float.rs @@ -0,0 +1,1276 @@ +use crate::AsIndex; +use crate::Cast; +use crate::Device; +use crate::Tensor; +use crate::cast::ToElement; +use crate::check; +use crate::check::TensorCheck; +use crate::kind::FloatMath; +use crate::ops::{BridgeKind, BridgeTensor}; +use crate::quantization::{QuantScheme, QuantizationParameters}; +use crate::tensor::stats; +use crate::tensor::{Distribution, TensorData}; +use crate::{Bool, Float, Int, TensorPrimitive}; +#[cfg(feature = "std")] +use burn_backend::AutodiffBackend; +use burn_backend::ElementConversion; +use burn_backend::Scalar; +use burn_backend::TensorMetadata; +#[cfg(feature = "std")] +use burn_backend::distributed::DistributedParamId; +use burn_backend::ops::ActivationOps; +use burn_backend::ops::FloatTensorOps; +use burn_backend::ops::GridSampleOptions; +use burn_backend::ops::QTensorOps; +use burn_backend::quantization::QuantizationParametersPrimitive; +use burn_dispatch::Dispatch; +use core::f32; + +/// Default RTOL value for `is_close` and `all_close`. +pub const DEFAULT_RTOL: f64 = 1e-5; + +/// Default ATOL value for `is_close` and `all_close`. +pub const DEFAULT_ATOL: f64 = 1e-8; + +impl Tensor { + /// Applies the [error function](https://en.wikipedia.org/wiki/Error_function) element wise. + /// + #[cfg_attr( + doc, + doc = r#" +$y_i = \text{erf}\(x_i\)$ + +The error function is defined as: + +$$\text{erf}\(x\) = \frac{2}{\sqrt{\pi}} \int_0^x e^{-t^2} dt$$ +"# + )] + #[cfg_attr(not(doc), doc = "`y_i = erf(x_i)`")] + pub fn erf(self) -> Self { + Self::new(erf_impl(self.primitive)) + } + + /// Applies [hypotenuse operation](https://en.wikipedia.org/wiki/Hypotenuse) element wise. + /// + #[cfg_attr(doc, doc = r#"$y_i = \sqrt{x_i^2 + y_i^2}$"#)] + #[cfg_attr(not(doc), doc = "`y_i = sqrt(x_i^2 + y_i^2)`")] + pub fn hypot(self, other: Self) -> Self { + Self::new(hypot_impl(self.primitive, other.primitive)) + } + + /// Applies [reciprocal operation](https://en.wikipedia.org/wiki/Multiplicative_inverse) + /// (or multiplicative inverse) element wise. + /// + #[cfg_attr(doc, doc = r#"$y_i = \frac{1}{x_i}$"#)] + #[cfg_attr(not(doc), doc = "`y_i = 1/x_i`")] + pub fn recip(self) -> Self { + Self::new(recip_impl(self.primitive)) + } + + /// Converts each of the elements of the input tensor from angles in degrees to radians. + /// + /// # Example + /// ```ignore + /// let tensor_in_radians = tensor.deg2rad(); + /// ``` + pub fn deg2rad(self) -> Self { + self.mul_scalar(f32::consts::PI / 180.0) + } + + /// Converts each of the elements of the input tensor from angles in radians to degrees. + /// + /// # Example + /// ```ignore + /// let tensor_in_degrees = tensor.rad2deg(); + /// ``` + pub fn rad2deg(self) -> Self { + self.mul_scalar(180.0 / f32::consts::PI) + } + + /// Applies element wise round operation. + /// + /// This function implements the [round half to even](https://en.wikipedia.org/wiki/Rounding#Rounding_half_to_even) + /// strategy, with halfway cases rounded to the nearest even integer value. + pub fn round(self) -> Self { + Self::new(round_impl(self.primitive)) + } + + /// Applies element wise floor operation. + pub fn floor(self) -> Self { + Self::new(floor_impl(self.primitive)) + } + + /// Applies element wise ceil operation. + pub fn ceil(self) -> Self { + Self::new(ceil_impl(self.primitive)) + } + + /// Create a tensor from floats (f32) on a given device. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// let _ = Tensor::<1>::from_floats([1.0, 2.0], &device); + /// let _ = Tensor::<2>::from_floats([[1.0, 2.0], [3.0, 4.0]], &device); + /// } + /// ``` + pub fn from_floats>(floats: A, device: &Device) -> Self { + Self::from_data(floats.into().convert::(), device) + } + + /// Returns a new tensor with the same shape and device as the current tensor and the data + /// cast to Integer. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// let float_tensor = Tensor::<1>::from_floats([1.0, 2.0], &device); + /// let int_tensor = float_tensor.int(); + /// } + /// ``` + pub fn int(self) -> Tensor { + let device = self.device(); + Tensor::new(int_impl(self.primitive, device)) + } + + /// Returns a new tensor with the same shape, dtype, and device as the current tensor filled random + /// values sampled from the given distribution. + pub fn random_like(&self, distribution: Distribution) -> Self { + Self::new(random_like_impl(&self.primitive, distribution)) + } + + /// Calculate the variance along the given dimension. + pub fn var(self, dim: usize) -> Self { + stats::var(self, dim) + } + + /// Calculate the variance along the given dimension without applying the Bessel’s correction. + pub fn var_bias(self, dim: usize) -> Self { + stats::var_bias(self, dim) + } + + /// Calculate the variance along the given dimension and also returns the mean. + pub fn var_mean(self, dim: usize) -> (Self, Self) { + let mean = self.clone().mean_dim(dim); + let var = stats::var_with_mean(self, mean.clone(), dim); + (var, mean) + } + + /// Calculate the variance along the given dimension without applying the Bessel’s correction and also returns the mean. + pub fn var_mean_bias(self, dim: usize) -> (Self, Self) { + let mean = self.clone().mean_dim(dim); + let var = stats::var_with_mean_bias(self, mean.clone(), dim); + (var, mean) + } + + /// Returns the median value along the specified dimension. + /// + /// The median is not unique for input tensors with an even number of elements + /// in the reduced dimension. In this case, the lower of the two medians is returned, + /// following PyTorch's behavior. + /// + /// # Note + /// + /// The current implementation performs a full sort along the specified dimension, + /// which has O(nlog(n)) complexity. Additionally, most backends currently fall back + /// to CPU for the sort operation, which may result in slower performance compared + /// to native GPU operations. + /// + /// # Arguments + /// + /// - `dim` - The dimension along which to compute the median. + /// + /// # Returns + /// + /// - A tensor containing the median values along the specified dimension. + /// + /// # Example 1 + /// + /// ```ignore + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data( + /// [[1.0, 5.0, 3.0, 2.0], [8.0, 4.0, 6.0, 7.0]], + /// &device, + /// ); + /// + /// // Median along dimension 0: + /// // sorted columns are [1.0, 8.0], [4.0, 5.0], [3.0, 6.0], [2.0, 7.0] + /// let median = tensor.median(0); + /// // Result: [[1.0, 4.0, 3.0, 2.0]] + /// + /// // Median along dimension 1: + /// // sorted rows are [1.0, 2.0, 3.0, 5.0] and [4.0, 6.0, 7.0, 8.0] + /// let median = tensor.median(1); + /// // Result: [[2.0], [6.0]] + /// ``` + /// + /// # Example 2 + /// + /// The median across all elements can be calculated as follows: + /// + /// ```ignore + /// // D is the number of dimensions of the tensor + /// let flattened_tensor: Tensor<1> = tensor.flatten(0, D - 1); + /// + /// // Calculate median for dim 0 since the tensor has become 1 dimensional + /// let median = flattened_tensor.median(0); + /// // Result: [4.0] + /// ``` + pub fn median(self, dim: usize) -> Self { + // TODO: Allow backend specialization. Optimally, implement a median kernel for cubecl + // instead of leveraging a full sort to get the median. + stats::median(self, dim) + } + + /// Returns the median value along the specified dimension and its index. + /// + /// The median is not unique for input tensors with an even number of elements + /// in the reduced dimension. In this case, the lower of the two medians is returned, + /// following PyTorch's behavior. + /// + /// # Note + /// + /// The current implementation performs a full sort along the specified dimension, + /// which has O(nlog(n)) complexity. Additionally, most backends currently fall back + /// to CPU for the sort operation, which may result in slower performance compared + /// to native GPU operations. + /// + /// # Arguments + /// + /// - `dim` - The dimension along which to compute the median. + /// + /// # Returns + /// + /// A tuple containing: + /// - A tensor with the median values. + /// - A tensor with the indices of the median values in the original tensor. + /// + /// # Example + /// + /// ```ignore + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data( + /// [[1.0, 5.0, 3.0, 2.0], [8.0, 4.0, 6.0, 7.0]], + /// &device, + /// ); + /// + /// // Median along dimension 1: + /// // sorted rows are [1.0, 2.0, 3.0, 5.0] and [4.0, 6.0, 7.0, 8.0] + /// let (values, indices) = tensor.median_with_indices(1); + /// // values: [[2.0], [6.0]], indices: [[3], [2]] (position in the original tensor) + /// ``` + pub fn median_with_indices(self, dim: usize) -> (Self, Tensor) { + // TODO: Allow backend specialization. Optimally, implement a median kernel for cubecl + // instead of leveraging a full sort to get the median. + stats::median_with_indices(self, dim) + } + + /// Converts a tensor to the specified data type. + /// + /// Supports both within-kind casting (e.g., `FloatDType::F64`) and cross-kind casting + /// (e.g., `IntDType::I64` to produce an int tensor). + /// + /// This is a no-op when casting to the current dtype within the same kind. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, FloatDType, IntDType}; + /// + /// fn example() { + /// let device = Default::default(); + /// let float_tensor = Tensor::<1>::from_floats([1.0, 2.5], &device); + /// + /// // Within-kind cast (float to float) + /// let f64_tensor = float_tensor.clone().cast(FloatDType::F64); + /// + /// // Cross-kind cast (float to int) + /// let int_tensor = float_tensor.cast(IntDType::I64); + /// } + /// ``` + #[must_use] + pub fn cast>(self, dtype: T) -> Tensor { + T::cast(self, dtype) + } + + /// Detach the current tensor from the autodiff graph. + /// + /// This function does nothing when autodiff is not enabled. + /// This can be used in batchers or elsewhere to ensure that previous operations are not + /// considered in the autodiff graph. + pub fn detach(self) -> Self { + Self::new(detach_impl(self.primitive)) + } + + /// Mark the tensor to keep gradients during the backward pass. + /// + /// This function does nothing when autodiff is not enabled. + pub fn require_grad(self) -> Self { + self.set_require_grad(true) + } + + /// Returns true if the tensor requires gradients during the backward pass. + pub fn is_require_grad(&self) -> bool { + is_require_grad_impl(&self.primitive) + } + + /// Mark the tensor as tracked or untracked depending on the require_grad argument. + /// When tracked, the gradients will be available after the backward pass. + /// + /// This function does nothing when autodiff is not enabled. + pub fn set_require_grad(self, require_grad: bool) -> Self { + Self::new(set_require_grad_impl(self.primitive, require_grad)) + } + + /// Applies the relu function to the tensor. + pub(crate) fn relu(self) -> Self { + Self::new(relu_impl(self.primitive)) + } + + /// Calculate covaraince matrix between different entries alongside a given dimension. + /// + /// # Arguments + /// + /// * `size` - The size of the square matrix. + /// * `correction_factor` - Is usually 1 for samples and 0 for population. + pub fn cov(self, dim: usize, correction_factor: usize) -> Tensor { + let n = self.dims()[dim]; + let centered = (self.clone() - self.mean_dim(dim)).swap_dims(dim, 0); + centered + .clone() + .transpose() + .matmul(centered) + .div_scalar(n as f32 - correction_factor as f32) + } + + /// Convert the tensor to a lower precision data type based on the quantization scheme. + /// + /// # Arguments + /// + /// * `scheme` - The quantization scheme. + /// * `qparams` - The pre-computed quantization parameters. + /// + /// # Returns + /// + /// The quantized tensor. + pub fn quantize(self, scheme: &QuantScheme, qparams: QuantizationParameters) -> Tensor { + Tensor::new(quantize_impl( + self.primitive, + scheme, + qparams.scales.primitive, + )) + } + + /// Dynamically convert the tensor to a lower precision data type based on the quantization scheme. + /// + /// # Arguments + /// + /// * `scheme` - The quantization scheme. + /// + /// # Returns + /// + /// The quantized tensor. + /// + /// # Notes + /// This uses [min-max calibration](crate::quantization::Calibration::MinMax). + pub fn quantize_dynamic(self, scheme: &QuantScheme) -> Tensor { + Tensor::new(quantize_dynamic_impl(self.primitive, scheme)) + } + + /// Convert the tensor back to a higher precision data type. + /// + /// If the tensor is not quantized, its value is simply returned. + /// + /// # Returns + /// + /// The dequantized tensor. + pub fn dequantize(self) -> Tensor { + Tensor::new(dequantize_impl(self.primitive)) + } + + /// Checks element wise if the tensor is close to another tensor. + /// + /// The tolerance is defined by the following equation: + /// + /// ```text + /// abs(a - b) <= (atol + rtol * abs(b)) + /// + /// where `a` is the first tensor, `b` is the second tensor, `rtol` is the relative tolerance, + /// and `atol` is the absolute tolerance. + /// ``` + /// + /// # Arguments + /// + /// * `other` - The tensor to compare with. + /// * `rtol` - Optional relative tolerance. Default is 1e-5; see `DEFAULT_RTOL`. + /// * `atol` - Optional absolute tolerance. Default is 1e-8; see `DEFAULT_ATOL`. + /// + /// # Returns + /// + /// A boolean tensor with the same shape as the input tensors. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor2 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor1.is_close(tensor2, None, None); + /// println!("{tensor}"); + /// // [[true, true, true], [true, true, true]] + /// } + /// ``` + pub fn is_close(self, other: Self, rtol: Option, atol: Option) -> Tensor { + let rtol = rtol.unwrap_or(DEFAULT_RTOL); + let atol = atol.unwrap_or(DEFAULT_ATOL); + + // check finite difference is close + let is_close_finite_val = self + .clone() + .sub(other.clone()) + .abs() + .lower_equal(other.clone().abs().mul_scalar(rtol).add_scalar(atol)) + .bool_and(self.clone().is_finite()) + .bool_and(other.clone().is_finite()); + + // check if both are infinite and have same sign + let inf_same_sign = self + .clone() + .is_finite() + .bool_not() + .bool_and(other.clone().is_finite().bool_not()) + .bool_and(self.equal(other)); + + is_close_finite_val.bool_or(inf_same_sign) + } + + /// Checks if all elements are close to another tensor. + /// + /// The tolerance is defined by the following equation: + /// + /// ```text + /// + /// abs(a - b) <= (atol + rtol * abs(b)) + /// + /// where `a` is the first tensor, `b` is the second tensor, `rtol` is the relative tolerance, + /// and `atol` is the absolute tolerance. + /// + /// ``` + /// + /// # Arguments + /// + /// * `other` - The tensor to compare with. + /// * `rtol` - Optional relative tolerance. Default is 1e-5; see `DEFAULT_RTOL`. + /// * `atol` - Optional absolute tolerance. Default is 1e-8; see `DEFAULT_ATOL`. + /// + /// # Returns + /// + /// A boolean scalar. + /// + /// # Remarks + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor2 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let result = tensor1.all_close(tensor2, None, None); + /// println!("{}", result); + /// // true + /// } + /// ``` + pub fn all_close(self, other: Self, rtol: Option, atol: Option) -> bool { + self.is_close(other, rtol, atol) + .all() + .into_scalar::() + .to_bool() + } + + /// Returns a new tensor with boolean elements indicating whether each element of the input is NaN. + /// + /// # Returns + /// + /// A boolean tensor where `true` indicates NaN and `false` indicates a non-NaN value. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, f64::NAN, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.is_nan(); + /// println!("{tensor}"); + /// // [[false, true, false], [false, false, false]] + /// } + /// ``` + pub fn is_nan(self) -> Tensor { + Tensor::new(is_nan_impl(self.primitive)) + } + + /// Checks if the tensor contains any NaN values. + /// + /// # Returns + /// + /// A boolean tensor with a single element indicating whether the tensor contains any NaN values. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [f64::NAN, 9.0, 6.0]], &device); + /// let tensor = tensor.contains_nan(); + /// println!("{tensor}"); + /// // [true] + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.contains_nan(); + /// println!("{tensor}"); + /// // [false] + /// } + /// ``` + pub fn contains_nan(self) -> Tensor<1, Bool> { + self.is_nan().any() + } + + /// Returns a new tensor with boolean elements indicating whether each element of the input is infinite (either +INF or -INF). + /// + /// # Returns + /// + /// A boolean tensor where `true` indicates that the value is infinite + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, f64::INFINITY, 3.0], [f64::NAN, 9.0, 6.0]], &device); + /// let tensor = tensor.is_finite(); + /// println!("{tensor}"); + /// // [[false, true, false], [false, false, false]] + /// } + /// ``` + pub fn is_inf(self) -> Tensor { + Tensor::new(is_inf_impl(self.primitive)) + } + + /// Returns a new tensor with boolean elements indicating whether each element of the input is finite + /// + /// # Returns + /// + /// A boolean tensor where `true` indicates that the value is finite and `false` indicates + /// either INF, -INF or NAN + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Bool, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, f64::INFINITY, 3.0], [f64::NAN, 9.0, 6.0]], &device); + /// let tensor = tensor.is_finite(); + /// println!("{tensor}"); + /// // [[true, false, true], [false, true, true]] + /// } + /// ``` + pub fn is_finite(self) -> Tensor { + self.clone() + .is_nan() + .bool_not() + .bool_and(self.is_inf().bool_not()) + } + + /// Samples tensor as a two-dimensional spatial grid of (possibly multi-channel) values, + /// using the given locations in [-1, 1]. + /// + /// # Arguments + /// + /// * `grid` - A tensor of locations, with shape (N, H_out, W_out, 2). Values are [-1, 1]. + /// A [x = -1, y = -1] means top-left, and [x = 1, y = 1] means bottom-right + /// * `options` - Grid sampling options (mode, padding_mode, align_corners) + /// + /// # Returns + /// + /// A tensor with shape (N, C, H_out, W_out) + /// + /// # Example + /// + /// ```ignore + /// use burn_tensor::ops::{GridSampleOptions, GridSamplePaddingMode, InterpolateMode}; + /// + /// // Default options (bilinear, zeros padding, align_corners=false) + /// let output = tensor.grid_sample_2d(grid, GridSampleOptions::default()); + /// + /// // Custom options + /// let options = GridSampleOptions::new(InterpolateMode::Bilinear) + /// .with_padding_mode(GridSamplePaddingMode::Border) + /// .with_align_corners(true); + /// let output = tensor.grid_sample_2d(grid, options); + /// ``` + pub fn grid_sample_2d( + self, + grid: Tensor, + options: impl Into, + ) -> Tensor { + Tensor::new(grid_sample_2d_impl( + self.primitive, + grid.primitive, + options.into(), + )) + } + + /// Computes the cross product of `self` and another tensor along a given dimension. + /// + /// Both `self` and `other` **must have size 3** along the specified `dim`, + /// because the cross product is only defined in three-dimensional space. + /// + /// # Arguments + /// + /// * `other` - The other tensor to take the cross product with. + /// * `dim` - The dimension along which to compute the cross product. + /// + /// # Returns + /// + /// A tensor containing the cross product of `self` and `other` along `dim`. + pub fn cross(self, other: Tensor, dim: Dim) -> Tensor { + let dim = dim.expect_dim_index(D); + check!(TensorCheck::cross(&self, &other, dim)); + Tensor::new(cross_impl(self.primitive, other.primitive, dim)) + } + + /// Applies element wise power operation with a float Tensor + /// + /// # Arguments + /// + /// * `other` - The tensor to apply the power operation with. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device); + /// let tensor = tensor1.powf(tensor2); + /// println!("{tensor}"); + /// // [[1.0, 8.0, 81.0], [5.0, 81.0, 216.0]] + /// } + /// ``` + pub fn powf(self, other: Self) -> Self { + Tensor::new(powf_impl(self.primitive, other.primitive)) + } + + /// Applies element wise power operation with a float scalar + /// + /// # Arguments + /// + /// * `other` - The scalar to apply the power operation with. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.powf_scalar(2.0); + /// println!("{tensor}"); + /// // [[1.0, 4.0, 9.0], [25.0, 81.0, 36.0]] + /// } + /// ``` + pub fn powf_scalar(self, other: E) -> Self { + let rhs = Scalar::new(other, &self.dtype()); + Tensor::new(powf_scalar_impl(self.primitive, rhs)) + } +} + +impl Tensor { + /// Draws samples from a categorical distribution defined by the last dimension + /// of the input tensor. + /// + /// The last dimension is treated as a (possibly unnormalized) set of weights + /// defining a categorical distribution over categories. All leading dimensions + /// are treated as batch dimensions. The method returns integer indices of the + /// sampled categories. + /// + /// # Arguments + /// + /// * `num_samples` - Number of samples to draw per distribution. Must be >= 1. + /// + /// # Panics + /// + /// Panics if `num_samples` is 0. + /// + /// # Note + /// + /// Distributions with all-zero weights produce undefined (NaN-based) sampling + /// results. Callers should ensure each distribution has at least one positive + /// weight. + /// + /// # Returns + /// + /// An integer tensor with the same shape as the input, except the last dimension + /// is replaced by `num_samples`, containing sampled category indices in + /// `[0, num_categories)`. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// let probs = Tensor::<2>::from_floats( + /// [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], + /// &device, + /// ); + /// let samples = probs.categorical(4); + /// // First row always samples index 1, second row always samples index 2 + /// println!("{samples}"); + /// } + /// ``` + pub fn categorical(self, num_samples: usize) -> Tensor { + assert!(num_samples > 0, "categorical: num_samples must be >= 1"); + + let shape = self.shape(); + let num_categories = shape[D - 1]; + let batch_size = (shape.num_elements() / num_categories).max(1); + let device = self.device(); + let dtype = self.dtype(); + + // Flatten leading dimensions into a single batch dimension: [batch, categories] + let flat: Tensor<2> = self.reshape([batch_size, num_categories]); + + // Normalize weights to probabilities + let sum = flat.clone().sum_dim(1); // [batch, 1] + let probs = flat / sum; + + // Cumulative sum along categories dimension + let cumsum = probs.cumsum(1); // [batch, categories] + + // Uniform random values for each sample + let uniform = Tensor::<2>::random( + [batch_size, num_samples], + Distribution::Uniform(0.0, 1.0), + (&device, dtype), + ); // [batch, num_samples] + + // Expand dimensions for broadcasting: + // cumsum: [batch, categories, 1] + // uniform: [batch, 1, num_samples] + let cumsum_3d: Tensor<3> = cumsum.unsqueeze_dim(2); + let uniform_3d: Tensor<3> = uniform.unsqueeze_dim(1); + + // Count categories where cumsum < uniform (inverse CDF) + let mask: Tensor<3, Bool> = cumsum_3d.lower(uniform_3d); + let indices: Tensor<2, Int> = mask.int().sum_dim(1).squeeze_dim::<2>(1); + + // Clamp to valid range to guard against floating-point imprecision in cumsum + let indices = indices.clamp(0, num_categories as i64 - 1); + + // Reshape back to [...leading_dims, num_samples] + let mut out_shape = shape; + out_shape[D - 1] = num_samples; + indices.reshape(out_shape) + } +} + +#[cfg(feature = "std")] +impl Tensor { + /// Returns true if the tensor is marked as distributed. + pub fn is_distributed(&self) -> bool { + is_distributed_impl(&self.primitive) + } + + /// Mark the tensor as distributed. + /// + /// This function does nothing when autodiff or distributed is not enabled. + pub fn set_distributed(self, param_id: DistributedParamId) -> Self { + Self::new(set_distributed_impl(self.primitive, param_id)) + } +} + +impl Tensor +where + K: FloatMath, +{ + /// Applies element wise square operation. + /// + #[cfg_attr(doc, doc = r#"$y_i = x_i * x_i$"#)] + #[cfg_attr(not(doc), doc = "`y_i = x_i * x_i`")] + pub fn square(self) -> Self { + Self::new(K::square(self.primitive)) + } + + /// Applies element wise exponential operation. + /// + #[cfg_attr(doc, doc = r#"$y_i = e^{x_i}$"#)] + #[cfg_attr(not(doc), doc = "`y = e^x`")] + pub fn exp(self) -> Self { + Self::new(K::exp(self.primitive)) + } + + /// Applies element wise natural logarithm of one plus the input tensor. + /// + #[cfg_attr(doc, doc = r#"$y_i = \log_e\(x_i + 1\)$"#)] + #[cfg_attr(not(doc), doc = "`y_i = log1p(x_i)`")] + pub fn log1p(self) -> Self { + Self::new(K::log1p(self.primitive)) + } + + /// Applies element wise natural log operation *ln*. + /// + #[cfg_attr(doc, doc = r#"$y_i = \log_e\(x_i\)$"#)] + #[cfg_attr(not(doc), doc = "`y_i = log(x_i)`")] + pub fn log(self) -> Self { + Self::new(K::log(self.primitive)) + } + + /// Applies element wise square root operation. + /// + pub fn sqrt(self) -> Self { + Tensor::new(K::sqrt(self.primitive)) + } + /// Applies element wise cosine operation. + /// + #[cfg_attr(doc, doc = r#"$y_i = \cos\(x_i\)$"#)] + #[cfg_attr(not(doc), doc = "`y_i = cos(x_i)`")] + pub fn cos(self) -> Self { + Tensor::new(K::cos(self.primitive)) + } + + /// Applies element wise sine operation. + /// + #[cfg_attr(doc, doc = r#"$y_i = \sin\(x_i\)$"#)] + #[cfg_attr(not(doc), doc = "`y_i = sin(x_i)`")] + pub fn sin(self) -> Self { + Tensor::new(K::sin(self.primitive)) + } + + /// Applies element wise tangent operation. + /// + #[cfg_attr(doc, doc = r#"$y_i = \tan\(x_i\)$"#)] + #[cfg_attr(not(doc), doc = "`y_i = tan(x_i)`")] + pub fn tan(self) -> Self { + Tensor::new(K::tan(self.primitive)) + } + + /// Applies element wise hyperbolic cosine operation. + /// + #[cfg_attr(doc, doc = r#"$y_i = \cosh\(x_i\)$"#)] + #[cfg_attr(not(doc), doc = "`y_i = cosh(x_i)`")] + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// + /// let tensor = Tensor::<1>::from_data([0.0, -1.0, 2.0], &device); + /// println!("{}", tensor.cosh()); // [1.0, 1.5430, 3.7621] + /// } + /// ``` + pub fn cosh(self) -> Self { + Tensor::new(K::cosh(self.primitive)) + } + + /// Applies element wise hyperbolic sine operation. + /// + #[cfg_attr(doc, doc = r#"$y_i = \sinh\(x_i\)$"#)] + #[cfg_attr(not(doc), doc = "`y_i = sinh(x_i)`")] + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// + /// let tensor = Tensor::<1>::from_data([0.0, -1.0, 2.0], &device); + /// println!("{}", tensor.sinh()); // [0.0, -1.1752, 3.6269] + /// } + /// ``` + pub fn sinh(self) -> Self { + Tensor::new(K::sinh(self.primitive)) + } + + /// Applies element wise hyperbolic tangent operation. + /// + #[cfg_attr(doc, doc = r#"$y_i = \tanh\(x_i\)$"#)] + #[cfg_attr(not(doc), doc = "`y_i = tanh(x_i)`")] + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// + /// let tensor = Tensor::<1>::from_data([0.0, -1.0, 2.0], &device); + /// println!("{}", tensor.tanh()); // [0.0, -0.7616, 0.9640] + /// } + /// ``` + pub fn tanh(self) -> Self { + Tensor::new(K::tanh(self.primitive)) + } + + /// Applies element wise inverse cosine operation. + /// + #[cfg_attr(doc, doc = r#"$y_i = \acos\(x_i\)$"#)] + #[cfg_attr(not(doc), doc = "`y_i = acos(x_i)`")] + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// + /// let tensor = Tensor::<1>::from_data([0.0, -1.0, 1.0], &device); + /// println!("{}", tensor.acos()); // [1.5708, 3.1416, 0.0] + /// } + /// ``` + pub fn acos(self) -> Self { + Tensor::new(K::acos(self.primitive)) + } + + /// Applies element wise inverse hyperbolic cosine operation. + /// + #[cfg_attr(doc, doc = r#"$y_i = \acosh\(x_i\)$"#)] + #[cfg_attr(not(doc), doc = "`y_i = acosh(x_i)`")] + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// + /// let tensor = Tensor::<1>::from_data([1.0, 2.0, 3.0], &device); + /// println!("{}", tensor.acosh()); // [0.0000, 1.3170, 1.7627] + /// } + /// ``` + pub fn acosh(self) -> Self { + Tensor::new(K::acosh(self.primitive)) + } + + /// Applies element wise inverse sine operation. + /// + #[cfg_attr(doc, doc = r#"$y_i = \asin\(x_i\)$"#)] + #[cfg_attr(not(doc), doc = "`y_i = asin(x_i)`")] + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// + /// let tensor = Tensor::<1>::from_data([0.0, -1.0, 1.0], &device); + /// println!("{}", tensor.asin()); // [ 0.0000, -1.5708, 1.5708] + /// } + /// ``` + pub fn asin(self) -> Self { + Tensor::new(K::asin(self.primitive)) + } + + /// Applies element wise inverse hyperbolic sine operation. + /// + #[cfg_attr(doc, doc = r#"$y_i = \asinh\(x_i\)$"#)] + #[cfg_attr(not(doc), doc = "`y_i = asinh(x_i)`")] + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// + /// let tensor = Tensor::<1>::from_data([0.0, -1.0, 1.0], &device); + /// println!("{}", tensor.asinh()); // [ 0.0000, -0.8814, 0.8814] + /// } + /// ``` + pub fn asinh(self) -> Self { + Tensor::new(K::asinh(self.primitive)) + } + + /// Applies element wise inverse tangent operation. + /// + #[cfg_attr(doc, doc = r#"$y_i = \atan\(x_i\)$"#)] + #[cfg_attr(not(doc), doc = "`y_i = atan(x_i)`")] + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// + /// let tensor = Tensor::<1>::from_data([0.0, -1.0, 2.0], &device); + /// println!("{}", tensor.atan()); // [ 0.0, -0.7854, 1.1071] + /// } + /// ``` + pub fn atan(self) -> Self { + Tensor::new(K::atan(self.primitive)) + } + + /// Applies element wise inverse hyperbolic tangent operation. + /// + #[cfg_attr(doc, doc = r#"$y_i = \atanh\(x_i\)$"#)] + #[cfg_attr(not(doc), doc = "`y_i = atanh(x_i)`")] + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// + /// let tensor = Tensor::<1>::from_data([0.0, -0.5, 0.5], &device); + /// println!("{}", tensor.atanh()); // [ 0.0, -0.5493, 0.5493] + /// } + /// ``` + pub fn atanh(self) -> Self { + Tensor::new(K::atanh(self.primitive)) + } + + /// Applies element wise inverse tangent operation using the signs of arguments to determine the correct quadrant. + /// + #[cfg_attr(doc, doc = r#"$z_i = \atan2\(y_i, x_i\)$"#)] + #[cfg_attr(not(doc), doc = "`z_i = atan2(y_i, x_i)`")] + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// + /// let lhs = Tensor::<1>::from_data([-2.0, 2.0, -2.0], &device); + /// let rhs = Tensor::<1>::from_data([1.0, -1.0, -1.0], &device); + /// println!("{}", lhs.atan2(rhs)); // [-1.1071, 2.0344, -2.0344] + /// } + /// ``` + pub fn atan2(self, other: Self) -> Self { + Tensor::new(K::atan2(self.primitive, other.primitive)) + } +} + +// ========================================================================= +// Non-generic implementation helpers (outlined from the public generic API). +// See the crate-level docs for the rationale behind this pattern. +// ========================================================================= + +fn erf_impl(p: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_erf(p.into_float())) +} + +fn recip_impl(p: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_recip(p.into_float())) +} + +fn hypot_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_hypot(lhs.into_float(), rhs.into_float())) +} + +fn round_impl(p: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_round(p.into_float())) +} + +fn floor_impl(p: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_floor(p.into_float())) +} + +fn ceil_impl(p: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_ceil(p.into_float())) +} + +fn int_impl(p: BridgeTensor, device: Device) -> BridgeTensor { + let out_dtype = device.settings().int_dtype; + BridgeTensor::int(Dispatch::float_into_int(p.into_float(), out_dtype)) +} + +fn random_like_impl(p: &BridgeTensor, distribution: Distribution) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_random( + p.shape(), + distribution, + &p.as_dispatch().device(), + p.dtype().into(), + )) +} + +fn detach_impl(p: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_detach(p.into_float())) +} + +fn is_require_grad_impl(p: &BridgeTensor) -> bool { + let (kind, tensor) = p.as_parts(); + match kind { + BridgeKind::Float => Dispatch::float_is_require_grad(tensor), + BridgeKind::QFloat => Dispatch::q_is_require_grad(tensor), + _ => panic!("Should be Float primitive kind"), + } +} + +fn set_require_grad_impl(p: BridgeTensor, require_grad: bool) -> BridgeTensor { + let (kind, tensor) = p.into_parts(); + match kind { + BridgeKind::Float => { + BridgeTensor::float(Dispatch::float_set_require_grad(tensor, require_grad)) + } + BridgeKind::QFloat => { + BridgeTensor::qfloat(Dispatch::q_set_require_grad(tensor, require_grad)) + } + _ => panic!("Should be Float primitive kind"), + } +} + +fn relu_impl(p: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::relu(p.into_float())) +} + +fn quantize_impl(p: BridgeTensor, scheme: &QuantScheme, scales: BridgeTensor) -> BridgeTensor { + BridgeTensor::qfloat(Dispatch::quantize( + p.into_float(), + scheme, + QuantizationParametersPrimitive { + scales: scales.into_float(), + }, + )) +} + +fn quantize_dynamic_impl(p: BridgeTensor, scheme: &QuantScheme) -> BridgeTensor { + BridgeTensor::qfloat(Dispatch::quantize_dynamic(p.into_float(), scheme)) +} + +fn dequantize_impl(p: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(p.into_float()) +} + +fn is_nan_impl(p: BridgeTensor) -> BridgeTensor { + let bool_dtype = p.device_settings().bool_dtype; + BridgeTensor::bool(Dispatch::float_is_nan(p.into_float(), bool_dtype)) +} + +fn is_inf_impl(p: BridgeTensor) -> BridgeTensor { + let bool_dtype = p.device_settings().bool_dtype; + BridgeTensor::bool(Dispatch::float_is_inf(p.into_float(), bool_dtype)) +} + +fn grid_sample_2d_impl( + p: BridgeTensor, + grid: BridgeTensor, + options: GridSampleOptions, +) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_grid_sample_2d( + p.into_float(), + grid.into_float(), + options, + )) +} + +fn cross_impl(p: BridgeTensor, other: BridgeTensor, dim: usize) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_cross( + p.into_float(), + other.into_float(), + dim, + )) +} + +fn powf_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + let (lkind, lhs) = lhs.into_parts(); + let (rkind, rhs) = rhs.into_parts(); + match (lkind, rkind) { + (BridgeKind::Float, BridgeKind::Float) => { + BridgeTensor::float(Dispatch::float_powf(lhs, rhs)) + } + (BridgeKind::QFloat, BridgeKind::QFloat) => match Dispatch::q_powf(lhs, rhs) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + (BridgeKind::QFloat, BridgeKind::Float) => { + let dtype = rhs.dtype(); + BridgeTensor::float(Dispatch::float_powf( + Dispatch::dequantize(lhs, dtype.into()), + rhs, + )) + } + (BridgeKind::Float, BridgeKind::QFloat) => { + let dtype = lhs.dtype(); + BridgeTensor::float(Dispatch::float_powf( + lhs, + Dispatch::dequantize(rhs, dtype.into()), + )) + } + _ => panic!("Should be Float primitive kind"), + } +} + +fn powf_scalar_impl(p: BridgeTensor, rhs: Scalar) -> BridgeTensor { + let (kind, lhs) = p.into_parts(); + match kind { + BridgeKind::Float => BridgeTensor::float(Dispatch::float_powf_scalar(lhs, rhs)), + BridgeKind::QFloat => match Dispatch::q_powf_scalar(lhs, rhs) { + TensorPrimitive::Float(out) => BridgeTensor::float(out), + TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out), + }, + _ => panic!("Should be Float primitive kind"), + } +} + +#[cfg(feature = "std")] +fn is_distributed_impl(p: &BridgeTensor) -> bool { + let (kind, tensor) = p.as_parts(); + match kind { + BridgeKind::Float => Dispatch::is_distributed(tensor), + BridgeKind::QFloat => unimplemented!(), + _ => panic!("Should be Float primitive kind"), + } +} + +#[cfg(feature = "std")] +fn set_distributed_impl(p: BridgeTensor, param_id: DistributedParamId) -> BridgeTensor { + let (kind, tensor) = p.into_parts(); + match kind { + BridgeKind::Float => { + BridgeTensor::float(Dispatch::set_distributed_params(tensor, param_id)) + } + BridgeKind::QFloat => unimplemented!(), + _ => panic!("Should be Float primitive kind"), + } +} diff --git a/crates/burn-tensor/src/tensor/api/fmod.rs b/crates/burn-tensor/src/tensor/api/fmod.rs new file mode 100644 index 0000000..5efb44e --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/fmod.rs @@ -0,0 +1,106 @@ +use crate::{Float, Tensor}; + +impl Tensor { + /// Computes the floating-point remainder of dividing `self` by `other`. + /// + /// The result has the same sign as `self` and magnitude less than `other`. + /// This is equivalent to the IEEE 754 remainder operation. + /// + /// # Special Cases (IEEE 754 compliant) + /// + /// - If `self` is ±∞ and `other` is not NaN, NaN is returned + /// - If `other` is ±0 and `self` is not NaN, NaN is returned + /// - If `other` is ±∞ and `self` is finite, `self` is returned + /// - If either argument is NaN, NaN is returned + /// + /// # Arguments + /// + /// * `other` - The divisor tensor. Must have the same shape as `self`. + /// + /// # Returns + /// + /// A tensor with the same shape where each element is the floating-point remainder. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// let dividend = Tensor::<1>::from_data([5.3, -5.3, 5.3, -5.3], &device); + /// let divisor = Tensor::<1>::from_data([2.0, 2.0, -2.0, -2.0], &device); + /// let result = dividend.fmod(divisor); + /// + /// // Result: [1.3, -1.3, 1.3, -1.3] + /// } + /// ``` + pub fn fmod(self, other: Self) -> Self { + // Normal case: fmod(x, y) = x - y * trunc(x / y) + let quotient = self.clone().div(other.clone()); + let truncated = quotient.trunc(); + let product = other.clone() * truncated.clone(); + + // When divisor is infinity and dividend is finite: + // - quotient is 0, truncated is 0 + // - but 0 * infinity = NaN, which is wrong + // We need to handle this case by replacing NaN with 0 when appropriate + + // Check if the product is NaN due to 0 * inf + let is_zero_times_inf = truncated.equal_elem(0.0).bool_and(other.is_inf()); + let zero_tensor = self.clone().mul_scalar(0.0); + let corrected_product = product.mask_where(is_zero_times_inf, zero_tensor); + + self - corrected_product + } + + /// Computes the floating-point remainder of dividing `self` by a scalar. + /// + /// The result has the same sign as `self` and magnitude less than the scalar. + /// + /// # Special Cases (IEEE 754 compliant) + /// + /// - If `self` is ±∞ and scalar is not NaN, NaN is returned + /// - If scalar is ±0 and `self` is not NaN, NaN is returned + /// - If scalar is ±∞ and `self` is finite, `self` is returned + /// - If either argument is NaN, NaN is returned + /// + /// # Arguments + /// + /// * `scalar` - The scalar divisor. + /// + /// # Returns + /// + /// A tensor with the same shape where each element is the floating-point remainder. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<1>::from_data([5.3, -5.3, 7.5, -7.5], &device); + /// let result = tensor.fmod_scalar(2.0); + /// + /// // Result: [1.3, -1.3, 1.5, -1.5] + /// } + /// ``` + pub fn fmod_scalar(self, scalar: f32) -> Self { + // Normal case: fmod(x, y) = x - y * trunc(x / y) + let quotient = self.clone().div_scalar(scalar); + let truncated = quotient.trunc(); + let product = truncated.mul_scalar(scalar); + + // Handle the special case where scalar is infinity + // When scalar is ±∞ and self is finite, quotient is 0, truncated is 0 + // but 0 * infinity = NaN, which is wrong - it should be 0 + if scalar.is_infinite() { + // For finite values, fmod(x, ±∞) = x + // For infinite values, fmod(±∞, ±∞) = NaN (which is handled by arithmetic) + return self; + } + + self - product + } +} diff --git a/crates/burn-tensor/src/tensor/api/graph.rs b/crates/burn-tensor/src/tensor/api/graph.rs new file mode 100644 index 0000000..7670d1f --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/graph.rs @@ -0,0 +1,153 @@ +use crate::Device; +use burn_backend::{Backend, BackendGraph}; +use burn_dispatch::Dispatch; + +/// A captured computation. +/// +/// [`capture`] records a closure once; [`replay`](Graph::replay) then re-runs it +/// as a single dispatch on backends with hardware graph support (CUDA/HIP — +/// collapsing hundreds of kernel launches into one), or simply re-executes the +/// closure everywhere else. The API is identical either way; only the speed +/// differs. +/// +/// The recorded graph replays against the exact device buffers used during +/// capture, so the closure must read its inputs from, and write its outputs to, +/// **stable** buffers held across the loop. Refresh those inputs in place with +/// [`Tensor::inplace`](crate::Tensor::inplace)-style writes before each replay. +/// +/// # Safety +/// +/// The graph has no explicit input/state/output signature: whatever buffers the +/// closure touched during capture are what every replay reads and writes, with +/// nothing tracking them afterwards. [`replay`](Graph::replay) is therefore +/// `unsafe` — see its safety contract. Safe, structured APIs can be layered on +/// top of this mechanism for specific workloads (e.g. a decode step with pinned +/// KV-cache and token buffers). +pub struct Graph { + device: Device, + output: T, + closure: F, + /// The captured hardware graph, or `None` to re-run the closure (a backend + /// without graph support, or a capture that failed). + hardware: Option>, +} + +/// Capture `closure` for repeated replay (see [`Graph`]). +/// +/// Runs `closure` twice: once to warm up (trigger autotuning and allocate every +/// buffer, so the capture itself needs no fresh device allocation — which is +/// illegal mid-capture), then once under capture. On a backend without graph +/// support the warmup/capture collapses to a single run and replay just re-runs +/// the closure. +pub fn capture(device: &Device, mut closure: F) -> Graph +where + F: FnMut() -> T, +{ + let dispatch = device.as_dispatch(); + // Prepare the allocator for capture, then warm up: this triggers all + // autotuning and allocates every buffer the capture will reuse. + // + // Several warmup iterations, not one: with fusion the first run builds and + // autotunes the fused optimization (a different execution path, allocating + // different scratch/metadata buffers than steady state). Only from the + // second run on does the closure take the cached fused path whose buffers + // the capture must reuse. Warming up a few times lets those steady-state + // allocations populate the persistent pool before capture opens — otherwise + // the capture run allocates them fresh, which faults mid-capture. + const WARMUP_ITERS: usize = 3; + let _ = Dispatch::graph_prepare(dispatch); + for _ in 0..WARMUP_ITERS { + // Hold the output alive across the sync. A lazy backend (fusion) elides + // a computation whose result is dropped before it drains — so dropping + // the warmup output would skip the very kernels (and allocations) the + // capture must reuse, and the capture run would then be the first to + // allocate them, faulting mid-capture. Keeping `out` until after `sync` + // forces the closure to actually execute and populate the pool. + let out = closure(); + let _ = Dispatch::sync(dispatch); + drop(out); + } + + let (hardware, output) = match Dispatch::graph_start_capture(dispatch) { + // Record the closure's launches into a graph. + Ok(()) => { + let output = closure(); + // A failed stop still resets the stream out of capture mode (the + // backend guarantees this even on error), so falling back to `None` — + // re-running the closure on replay — stays correct; we lose the + // graph, not stream health. + (Dispatch::graph_stop_capture(dispatch).ok(), output) + } + // No hardware graph support: fall back to re-running the closure. + Err(_) => (None, closure()), + }; + + Graph { + device: device.clone(), + output, + closure, + hardware, + } +} + +impl Graph +where + F: FnMut() -> T, +{ + /// Re-run the captured computation and return its output. + /// + /// On a captured graph this is one dispatch replaying the recorded launches + /// against their original buffers — so write fresh inputs into those + /// buffers first. On the fallback path it re-executes the closure. The + /// returned reference is the output produced during capture (whose buffer + /// the replay just overwrote), or the fresh output on the fallback path. + /// + /// # Safety + /// + /// On the hardware path this dispatches the recorded kernels against the raw + /// device buffers the closure touched during capture, with nothing checking + /// they are still valid. The caller must guarantee, until the replay's work + /// completes (e.g. it is followed by a read of the output or a device sync): + /// + /// - **Liveness** — every tensor the captured closure read or wrote still + /// exists. Dropping one frees its buffer for reuse by other allocations, + /// and a later replay would read or overwrite whatever now lives there. + /// Tensors owned by the closure (or by `self`, like the output) are kept + /// alive automatically; tensors the closure only borrowed must outlive + /// the replays. + /// - **No concurrent use** — no other stream or thread reads or writes a + /// tensor shared with the graph while the replay executes; the replay is + /// only ordered against work on its own capture stream. + /// - **Same-stream refreshes** — input refreshes and output reads are + /// issued on the stream the graph was captured on (the same device + /// thread/client), so they order correctly against the replay rather + /// than racing it with stale or torn data. + /// + /// On the fallback path (no hardware graph) this simply re-runs the closure + /// and is trivially safe. + pub unsafe fn replay(&mut self) -> &T { + match &self.hardware { + Some(graph) => { + // Safety: forwarded verbatim from this method's own contract. + unsafe { Dispatch::graph_replay(self.device.as_dispatch(), graph) } + .expect("graph replay should succeed"); + } + None => { + self.output = (self.closure)(); + } + } + &self.output + } + + /// The output tensor(s) the graph writes to — stable across replays on the + /// hardware path (the same buffer is overwritten each time). + pub fn output(&self) -> &T { + &self.output + } + + /// Whether this graph replays as a hardware dispatch (`true`) or by + /// re-running the closure (`false`). + pub fn is_hardware(&self) -> bool { + self.hardware.is_some() + } +} diff --git a/crates/burn-tensor/src/tensor/api/int.rs b/crates/burn-tensor/src/tensor/api/int.rs new file mode 100644 index 0000000..97184cb --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/int.rs @@ -0,0 +1,261 @@ +use burn_backend::{ElementConversion, Scalar, ops::IntTensorOps}; +use burn_dispatch::Dispatch; + +use crate::{ + Cast, Device, Float, Int, Shape, Tensor, TensorCreationOptions, TensorData, cartesian_grid, + ops::BridgeTensor, +}; + +use core::ops::Range; + +impl Tensor<1, Int> { + /// Returns a new integer tensor on the specified device. + /// + /// # Arguments + /// + /// * `range` - The range of values to generate. + /// * `device` - The device to create the tensor on. + pub fn arange(range: Range, options: impl Into) -> Self { + let opt = options.into(); + let dtype = opt.resolve_dtype::(); + Tensor::new(arange_impl(range, opt.device, dtype)) + } + + /// Returns a new integer tensor on the specified device. + /// + /// # Arguments + /// + /// * `range` - The range of values to generate. + /// * `step` - The step between each value. + pub fn arange_step( + range: Range, + step: usize, + options: impl Into, + ) -> Self { + let opt = options.into(); + let dtype = opt.resolve_dtype::(); + Tensor::new(arange_step_impl(range, step, opt.device, dtype)) + } +} + +impl Tensor { + /// Create a tensor from integers (i32), placing it on a given device. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Int}; + /// + /// fn example() { + /// let device = Default::default(); + /// let _x: Tensor<1, Int> = Tensor::from_ints([1, 2], &device); + /// let _y: Tensor<2, Int> = Tensor::from_ints([[1, 2], [3, 4]], &device); + /// } + /// ``` + pub fn from_ints>(ints: A, device: &Device) -> Self { + Self::from_data(ints.into().convert::(), device) + } + + /// Returns a new tensor with the same shape and device as the current tensor and the data + /// cast to Float. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Int, Tensor}; + /// + /// fn example() { + /// let device = Default::default(); + /// let int_tensor = Tensor::<1, Int>::arange(0..5, &device); + /// let float_tensor = int_tensor.float(); + /// } + /// ``` + pub fn float(self) -> Tensor { + let device = self.device(); + Tensor::new(int_to_float_impl(self.primitive, device)) + } + + /// Generates a cartesian grid for the given tensor shape on the specified device. + /// The generated tensor is of dimension `D2 = D + 1`, where each element at dimension D contains the cartesian grid coordinates for that element. + /// + /// # Arguments + /// + /// * `shape` - The shape specifying the dimensions of the tensor. + /// * `device` - The device to create the tensor on. + /// + /// # Panics + /// + /// Panics if `D2` is not equal to `D+1`. + /// + /// # Examples + /// + /// ```rust + /// use burn_tensor::Int; + /// use burn_tensor::{Shape, Tensor}; + /// fn example() { + /// let device = Default::default(); + /// let result: Tensor<3, _> = Tensor::<2, Int>::cartesian_grid([2, 3], &device); + /// println!("{}", result); + /// } + /// ``` + pub fn cartesian_grid, const D2: usize>( + shape: S, + device: &Device, + ) -> Tensor { + cartesian_grid::(shape, device) + } + + /// Applies the bitwise logical and operation with each bit representing the integer. + pub fn bitwise_and(self, other: Self) -> Self { + Self::new(bitwise_and_impl(self.primitive, other.primitive)) + } + + /// Applies the bitwise logical or operation with another tensor. + pub fn bitwise_or(self, other: Self) -> Self { + Self::new(bitwise_or_impl(self.primitive, other.primitive)) + } + + /// Applies the bitwise logical xor operation with another tensor. + pub fn bitwise_xor(self, other: Self) -> Self { + Self::new(bitwise_xor_impl(self.primitive, other.primitive)) + } + + /// Applies the bitwise logical not operation. + pub fn bitwise_not(self) -> Self { + Self::new(bitwise_not_impl(self.primitive)) + } + + /// Applies the bitwise logical and operation with each bit in the scalar and the integers in the tensor. + pub fn bitwise_and_scalar(self, other: impl ElementConversion) -> Self { + let other = Scalar::new(other, &self.dtype()); + Self::new(bitwise_and_scalar_impl(self.primitive, other)) + } + + /// Applies the bitwise logical or operation with each bit in the scalar and the integers in the tensor. + pub fn bitwise_or_scalar(self, other: impl ElementConversion) -> Self { + let other = Scalar::new(other, &self.dtype()); + Self::new(bitwise_or_scalar_impl(self.primitive, other)) + } + + /// Applies bitwise logical xor operation with each bit in the scalar and the integers in the tensor. + pub fn bitwise_xor_scalar(self, other: impl ElementConversion) -> Self { + let other = Scalar::new(other, &self.dtype()); + Self::new(bitwise_xor_scalar_impl(self.primitive, other)) + } + + /// Applies the bitwise left shift operation with the integers in the tensor. + pub fn bitwise_left_shift(self, other: Self) -> Self { + Self::new(bitwise_left_shift_impl(self.primitive, other.primitive)) + } + + /// Applies the bitwise right shift operation with the integers in the tensor. + pub fn bitwise_right_shift(self, other: Self) -> Self { + Self::new(bitwise_right_shift_impl(self.primitive, other.primitive)) + } + + /// Applies the bitwise left shift operation with the scalar. + pub fn bitwise_left_shift_scalar(self, other: impl ElementConversion) -> Self { + let other = Scalar::new(other, &self.dtype()); + Self::new(bitwise_left_shift_scalar_impl(self.primitive, other)) + } + + /// Applies the bitwise right shift operation with the scalar. + pub fn bitwise_right_shift_scalar(self, other: impl ElementConversion) -> Self { + let other = Scalar::new(other, &self.dtype()); + Self::new(bitwise_right_shift_scalar_impl(self.primitive, other)) + } + + /// Converts a tensor to the specified data type. + /// + /// Supports both within-kind casting (e.g., `IntDType::I64`) and cross-kind casting + /// (e.g., `FloatDType::F32` to produce a float tensor). + /// + /// This is a no-op when casting to the current dtype within the same kind. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Int, IntDType, FloatDType}; + /// + /// fn example() { + /// let device = Default::default(); + /// let int_tensor = Tensor::<1, Int>::arange(0..5, &device); + /// + /// // Within-kind cast (int to int) + /// let i64_tensor = int_tensor.clone().cast(IntDType::I64); + /// + /// // Cross-kind cast (int to float) + /// let float_tensor = int_tensor.cast(FloatDType::F32); + /// } + /// ``` + #[must_use] + pub fn cast>(self, dtype: T) -> Tensor { + T::cast(self, dtype) + } +} + +// ========================================================================= +// Non-generic implementation helpers (outlined from the generic API). +// See the crate-level docs for the rationale behind this pattern. +// ========================================================================= + +fn arange_impl(range: Range, device: Device, dtype: burn_std::DType) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_arange( + range, + device.as_dispatch(), + dtype.into(), + )) +} + +fn arange_step_impl( + range: Range, + step: usize, + device: Device, + dtype: burn_std::DType, +) -> BridgeTensor { + BridgeTensor::int(Dispatch::int_arange_step( + range, + step, + device.as_dispatch(), + dtype.into(), + )) +} + +fn int_to_float_impl(p: BridgeTensor, device: Device) -> BridgeTensor { + let out_dtype = device.settings().float_dtype; + BridgeTensor::float(Dispatch::int_into_float(p.into(), out_dtype)) +} + +fn bitwise_and_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::bitwise_and(lhs.into(), rhs.into())) +} +fn bitwise_or_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::bitwise_or(lhs.into(), rhs.into())) +} +fn bitwise_xor_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::bitwise_xor(lhs.into(), rhs.into())) +} +fn bitwise_not_impl(p: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::bitwise_not(p.into())) +} +fn bitwise_and_scalar_impl(p: BridgeTensor, other: Scalar) -> BridgeTensor { + BridgeTensor::int(Dispatch::bitwise_and_scalar(p.into(), other)) +} +fn bitwise_or_scalar_impl(p: BridgeTensor, other: Scalar) -> BridgeTensor { + BridgeTensor::int(Dispatch::bitwise_or_scalar(p.into(), other)) +} +fn bitwise_xor_scalar_impl(p: BridgeTensor, other: Scalar) -> BridgeTensor { + BridgeTensor::int(Dispatch::bitwise_xor_scalar(p.into(), other)) +} +fn bitwise_left_shift_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::bitwise_left_shift(lhs.into(), rhs.into())) +} +fn bitwise_right_shift_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor { + BridgeTensor::int(Dispatch::bitwise_right_shift(lhs.into(), rhs.into())) +} +fn bitwise_left_shift_scalar_impl(p: BridgeTensor, other: Scalar) -> BridgeTensor { + BridgeTensor::int(Dispatch::bitwise_left_shift_scalar(p.into(), other)) +} +fn bitwise_right_shift_scalar_impl(p: BridgeTensor, other: Scalar) -> BridgeTensor { + BridgeTensor::int(Dispatch::bitwise_right_shift_scalar(p.into(), other)) +} diff --git a/crates/burn-tensor/src/tensor/api/mod.rs b/crates/burn-tensor/src/tensor/api/mod.rs new file mode 100644 index 0000000..d51528b --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/mod.rs @@ -0,0 +1,35 @@ +pub(crate) mod check; + +mod autodiff; +mod base; +mod bool; +mod cartesian_grid; +mod cast; +mod float; +mod fmod; +mod graph; +mod int; +mod numeric; +mod options; +mod orderable; +mod pad; +pub use pad::IntoPadding; +mod take; +mod transaction; + +mod trunc; + +#[cfg(feature = "autodiff")] +pub use autodiff::*; +pub use base::*; +pub use cartesian_grid::cartesian_grid; +pub use cast::*; +pub use float::{DEFAULT_ATOL, DEFAULT_RTOL}; +pub use graph::{Graph, capture}; +pub use options::*; +pub use transaction::*; + +#[cfg(feature = "extension")] +mod extension; +#[cfg(feature = "extension")] +pub use extension::*; diff --git a/crates/burn-tensor/src/tensor/api/numeric.rs b/crates/burn-tensor/src/tensor/api/numeric.rs new file mode 100644 index 0000000..d76ec3b --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/numeric.rs @@ -0,0 +1,1127 @@ +use burn_backend::Scalar; + +use crate::alloc::borrow::ToOwned; +use crate::kind::Numeric; +use alloc::vec::Vec; + +use crate::{ + AsIndex, Bool, Distribution, ElementConversion, Int, Shape, Tensor, check, check::TensorCheck, +}; +use crate::{Device, IndexingUpdateOp, TensorCreationOptions}; + +impl Tensor +where + K: Numeric, +{ + /// Applies element wise addition operation. + /// + /// `y = x2 + x1` + /// + /// # Arguments + /// + /// * `other` - The tensor to add. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device); + /// let tensor = tensor1 + tensor2; + /// println!("{tensor}"); + /// // [[3.0, 1.0, 7.0], [6.0, 11.0, 9.0]] + /// } + /// ``` + #[allow(clippy::should_implement_trait)] + pub fn add(self, other: Self) -> Self { + check!(TensorCheck::binary_ops_ew("Add", &self, &other)); + Self::new(K::add(self.primitive, other.primitive)) + } + + /// Applies element wise addition operation with a scalar. + /// + /// `y = x + s` + /// + /// # Arguments + /// + /// * `other` - The scalar to add, element wise. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let scalar = 2.0; + /// let tensor = tensor + scalar; + /// println!("{tensor}"); + /// // [[3.0, 0.0, 5.0], [7.0, 11.0, 8.0]] + /// } + /// ``` + pub fn add_scalar(self, other: E) -> Self { + let other = Scalar::new(other, &self.dtype()); + Self::new(K::add_scalar(self.primitive, other)) + } + + /// Applies element wise subtraction operation. + /// + /// `y = x2 - x1` + /// + /// # Arguments + /// + /// * `other` - The tensor to subtract. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device); + /// let tensor = tensor1 - tensor2; + /// println!("{tensor}"); + /// // [[-1.0, -5.0, -1.0], [4.0, 7.0, 3.0]] + /// } + /// ``` + #[allow(clippy::should_implement_trait)] + pub fn sub(self, other: Self) -> Self { + check!(TensorCheck::binary_ops_ew("Sub", &self, &other)); + Self::new(K::sub(self.primitive, other.primitive)) + } + + /// Applies element wise subtraction operation with a scalar. + /// + /// `y = x - s` + /// + /// # Arguments + /// + /// * `other` - The scalar to subtract, element wise. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let scalar = 2.0; + /// let tensor = tensor - scalar; + /// println!("{tensor}"); + /// // [[-1.0, -4.0, 1.0], [3.0, 7.0, 4.0]] + /// } + /// ``` + pub fn sub_scalar(self, other: E) -> Self { + let other = Scalar::new(other, &self.dtype()); + Self::new(K::sub_scalar(self.primitive, other)) + } + + /// Applies element wise division operation. + /// + /// `y = x2 / x1` + /// + /// # Arguments + /// + /// * `other` - The tensor to divide. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device); + /// let tensor = tensor1 / tensor2; + /// println!("{tensor}"); + /// // [[0.5, -0.6666667, 0.75], [5.0, 4.5, 2.0]] + /// } + /// ``` + #[allow(clippy::should_implement_trait)] + pub fn div(self, other: Self) -> Self { + check!(TensorCheck::binary_ops_ew("Div", &self, &other)); + Self::new(K::div(self.primitive, other.primitive)) + } + + /// Applies element wise division operation with a scalar. + /// + /// `y = x / s` + /// + /// # Arguments + /// + /// * `other` - The scalar to divide, element wise. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let scalar = 2.0; + /// let tensor = tensor / scalar; + /// println!("{tensor}"); + /// // [[0.5, -1.0, 1.5], [2.5, 4.5, 3.0]] + /// } + /// ``` + pub fn div_scalar(self, other: E) -> Self { + let other = Scalar::new(other, &self.dtype()); + Self::new(K::div_scalar(self.primitive, other)) + } + + /// Applies element wise the remainder operation with a scalar. + /// + /// `y = x2 % x1` + pub fn remainder(self, other: Self) -> Self { + Self::new(K::remainder(self.primitive, other.primitive)) + } + + /// Applies element wise the remainder operation with a scalar. + /// + /// `y = x % s` + /// + /// # Arguments + /// + /// * `other` - The scalar to divide, element wise. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let scalar = 2.0; + /// let tensor = tensor1 % scalar; + /// println!("{tensor}"); + /// // [[1.0, 0.0, 1.0], [1.0, 1.0, 0.0]] + /// } + /// ``` + pub fn remainder_scalar(self, other: E) -> Self { + let other = Scalar::new(other, &self.dtype()); + Self::new(K::remainder_scalar(self.primitive, other)) + } + + /// Applies element wise multiplication operation. + /// + /// `y = x2 * x1` + /// + /// # Arguments + /// + /// * `other` - The tensor to multiply. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device); + /// let tensor = tensor1 * tensor2; + /// println!("{tensor}"); + /// // [[2.0, -6.0, 12.0], [5.0, 18.0, 18.0]] + /// } + /// ``` + #[allow(clippy::should_implement_trait)] + pub fn mul(self, other: Self) -> Self { + check!(TensorCheck::binary_ops_ew("Mul", &self, &other)); + Self::new(K::mul(self.primitive, other.primitive)) + } + + /// Applies element wise multiplication operation with a scalar. + /// + /// `y = x * s` + /// + /// # Arguments + /// + /// * `other` - The scalar to multiply, element wise. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let scalar = 2.0; + /// let tensor = tensor * scalar; + /// println!("{tensor}"); + /// // [[2.0, -4.0, 6.0], [10.0, 18.0, 12.0]] + /// } + /// ``` + pub fn mul_scalar(self, other: E) -> Self { + let other = Scalar::new(other, &self.dtype()); + Self::new(K::mul_scalar(self.primitive, other)) + } + + /// Switch sign of each element in the tensor. + /// + /// `y = -x` + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = -tensor; + /// println!("{tensor}"); + /// // [[-1.0, 2.0, -3.0], [-5.0, -9.0, -6.0]] + /// } + /// ``` + #[allow(clippy::should_implement_trait)] + pub fn neg(self) -> Self { + Self::new(K::neg(self.primitive)) + } + + /// Returns the signs of the elements of the input tensor. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.sign(); + /// println!("{tensor}"); + /// // [[1.0, -1.0, 1.0], [1.0, 1.0, 1.0]] + /// } + /// ``` + pub fn sign(self) -> Self { + Self::new(K::sign(self.primitive)) + } + + /// Aggregate all elements in the tensor with the mean operation. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.mean(); + /// println!("{tensor}"); + /// // [3.6666667] + /// } + /// ``` + pub fn mean(self) -> Tensor<1, K> { + Tensor::new(K::mean(self.primitive)) + } + + /// Aggregate all elements in the tensor with the sum operation. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.sum(); + /// println!("{tensor}"); + /// // [22.0] + /// } + /// ``` + pub fn sum(self) -> Tensor<1, K> { + Tensor::new(K::sum(self.primitive)) + } + + /// Aggregate all elements along the given *dimension* or *axis* + /// in the tensor with the mean operation. + /// + /// # Arguments + /// + /// * `dim` - The dimension or axis along which to aggregate the elements; + /// supports negative indexing. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.clone().mean_dim(0); + /// println!("{tensor}"); + /// // [[3.0, 3.5, 4.5]] + /// let tensor = tensor.clone().mean_dim(1); + /// println!("{tensor}"); + /// // [[0.6666667], [6.6666665]] + /// } + /// ``` + pub fn mean_dim(self, dim: I) -> Self { + let dim = dim.expect_dim_index(D); + check!(TensorCheck::aggregate_dim::("Mean", dim)); + Self::new(K::mean_dim(self.primitive, dim)) + } + + /// Aggregate all elements along the given *axes* + /// in the tensor with the mean operation. + /// + /// # Arguments + /// + /// * `dims` - the dimensions to aggregate; supports negative indexing. + /// + /// # Returns + /// + /// The returned tensor will have the same rank, + /// but the aggregated dimensions will have size 1. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[2.0, 4.0], [6.0, -4.0]], &device); + /// let tensor = tensor.clone().mean_dims(&[0, 1]); + /// println!("{tensor}"); + /// // [[2.0]] + /// } + /// ``` + pub fn mean_dims(self, dims: &[I]) -> Self { + dims.iter().fold(self, |tensor, &dim| tensor.mean_dim(dim)) + } + + /// Aggregate all elements along the given *dimension* or *axis* + /// in the tensor with the sum operation. + /// + /// # Arguments + /// + /// * `dim` - The dimension or axis along which to aggregate the elements; + /// supports negative indexing. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.clone().sum_dim(0); + /// println!("{tensor}"); + /// // [[6.0, 7.0, 9.0]] + /// let tensor = tensor.clone().sum_dim(1); + /// println!("{tensor}"); + /// // [[2.0], [20.0]] + /// } + /// ``` + pub fn sum_dim(self, dim: I) -> Self { + let dim = dim.expect_dim_index(D); + check!(TensorCheck::aggregate_dim::("Sum", dim)); + Self::new(K::sum_dim(self.primitive, dim)) + } + + /// Aggregate all elements along the given *axes* + /// in the tensor with the sum operation. + /// + /// # Arguments + /// + /// * `dims` - the dimensions to aggregate; supports negative indexing. + /// + /// # Returns + /// + /// The returned tensor will have the same rank, + /// but the aggregated dimensions will have size 1. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.clone().sum_dims(&[0, 1]); + /// println!("{tensor}"); + /// // [[27]] + /// } + /// ``` + pub fn sum_dims(self, dims: &[I]) -> Self { + dims.iter().fold(self, |tensor, &dim| tensor.sum_dim(dim)) + } + + /// Aggregate and squeeze along the given dimensions. + /// + /// This is equivalent to ``tensor.sum_dims(dims).squeeze_dims(dims)`` + /// + /// # Arguments + /// + /// * `dims` - the dimensions to aggregate; supports negative indexing. + /// + /// # Returns + /// + /// The returned tensor will have the same rank, + /// but the aggregated dimensions will have size 1. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<3>::from_data([ + /// [[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], + /// [[9.0, 2.0, 5.0], [5.0, 7.0, 7.0]], + /// ], &device); + /// let tensor = tensor.clone().sum_dims_squeeze::<1, _>(&[0, 1]); + /// println!("{tensor}"); + /// // [20.0, 16.0, 21.0] + /// } + /// ``` + pub fn sum_dims_squeeze(self, dims: &[I]) -> Tensor { + // TODO: remove idims when squeeze_dims uses AsIndex. + let idims = dims + .iter() + .map(|&dim| (dim.expect_dim_index(D)) as isize) + .collect::>(); + self.sum_dims(dims).squeeze_dims::(&idims) + } + + /// Aggregate all elements in the tensor with the product operation. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.prod(); + /// println!("{tensor}"); + /// // [-1620.0] + /// } + /// ``` + pub fn prod(self) -> Tensor<1, K> { + Tensor::new(K::prod(self.primitive)) + } + + /// Aggregate all elements along the given *dimension* or *axis* + /// in the tensor with the product operation. + /// + /// # Arguments + /// + /// * `dim` - The dimension or axis along which to aggregate the elements, + /// supports negative indexing. + /// + /// # Returns + /// + /// The returned tensor will have the same rank, + /// but the aggregated dimension will have size 1. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.clone().prod_dim(0); + /// println!("{tensor}"); + /// // [[5.0, -18.0, 18.0]] + /// let tensor = tensor.clone().prod_dim(1); + /// println!("{tensor}"); + /// // [[-6.0], [270.0]] + /// } + /// ``` + pub fn prod_dim(self, dim: I) -> Self { + let dim = dim.expect_dim_index(D); + check!(TensorCheck::aggregate_dim::("Prod", dim)); + Self::new(K::prod_dim(self.primitive, dim)) + } + + /// Aggregate all elements along the given *axes* + /// in the tensor with the prod operation. + /// + /// # Arguments + /// + /// * `dims` - the dimensions to aggregate, supports negative indexing. + /// + /// # Returns + /// + /// The returned tensor will have the same rank, + /// but the aggregated dimensions will have size 1. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.clone().sum_dims(&[0, 1]); + /// println!("{tensor}"); + /// // [[-1620.0]] + /// } + /// ``` + pub fn prod_dims(self, dims: &[I]) -> Self { + dims.iter().fold(self, |tensor, &dim| tensor.prod_dim(dim)) + } + + /// Computes the cumulative sum of elements along the given *dimension* or *axis*. + /// + /// # Arguments + /// + /// * `dim` - The dimension or axis along which to compute the cumulative sum. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + /// let result = tensor.clone().cumsum(0); + /// println!("{result}"); + /// // [[1.0, 2.0, 3.0], [5.0, 7.0, 9.0]] + /// let result = tensor.cumsum(1); + /// println!("{result}"); + /// // [[1.0, 3.0, 6.0], [4.0, 9.0, 15.0]] + /// } + /// ``` + pub fn cumsum(self, dim: usize) -> Self { + check!(TensorCheck::aggregate_dim::("CumSum", dim)); + Self::new(K::cumsum(self.primitive, dim)) + } + + /// Computes the cumulative product of elements along the given *dimension* or *axis*. + /// + /// # Arguments + /// + /// * `dim` - The dimension or axis along which to compute the cumulative product. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + /// let result = tensor.clone().cumprod(0); + /// println!("{result}"); + /// // [[1.0, 2.0, 3.0], [4.0, 10.0, 18.0]] + /// let result = tensor.cumprod(1); + /// println!("{result}"); + /// // [[1.0, 2.0, 6.0], [4.0, 20.0, 120.0]] + /// } + /// ``` + pub fn cumprod(self, dim: usize) -> Self { + check!(TensorCheck::aggregate_dim::("CumProd", dim)); + Self::new(K::cumprod(self.primitive, dim)) + } + + /// Apply element wise absolute value operation. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Int, Tensor}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2, Int>::from_ints([[1, -2, 3], [4, -5, 6], [7, -8, 9]], &device); + /// let tensor = tensor.abs(); + /// println!("{tensor}"); + /// // [[1, 2, 3], [4, 5, 6], [7, 8, 9]] + /// } + /// ``` + /// + /// # Notes + /// + /// For signed integer dtypes, this operation uses two's-complement wraparound semantics, similar to + /// `x.wrapping_abs()`. For example, `abs(i64::MIN) == i64::MIN`. + pub fn abs(self) -> Self { + Self::new(K::abs(self.primitive)) + } + + /// Returns the upper triangular part of a matrix (2-D tensor) or batch of matrices input, + /// the other elements of the result tensor out are set to 0. + /// + /// See also [`triu_mask`](Tensor::triu_mask). + /// + /// # Arguments + /// + /// * `diagonal` - The offset from the diagonal, where 0 means the diagonal, and positive values shift + /// towards the upper triangle. + /// + /// # Example + /// ```rust + /// use burn_tensor::{Int, Tensor}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2, Int>::from_ints( + /// [ + /// [1, 2, 3], + /// [4, 5, 6], + /// [7, 8, 9] + /// ], + /// &device + /// ); + /// let tensor = tensor.triu(1); + /// println!("{tensor}"); + /// // [ + /// // [0, 2, 3], + /// // [0, 0, 6], + /// // [0, 0, 0] + /// // ] + /// } + /// ``` + pub fn triu(self, diagonal: i64) -> Self { + check!(TensorCheck::tri::<{ D }>()); + + // last two dimensions + let shape = &self.shape()[D - 2..].to_owned(); + + let mask = Tensor::<2, Bool>::triu_mask(shape, diagonal, &self.device()).unsqueeze(); + self.mask_fill(mask, 0) + } + + /// Returns the lower triangular part of a matrix (2-D tensor) or batch of matrices input, + /// the other elements of the result tensor out are set to 0. + /// + /// See also [`tril_mask`](Tensor::tril_mask). + /// + /// # Arguments + /// + /// * `diagonal` - The offset from the diagonal, where 0 means the diagonal, and positive values shift + /// towards the upper triangle. + /// + /// # Example + /// ```rust + /// use burn_tensor::{Int, Tensor}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2, Int>::from_ints( + /// [ + /// [1, 2, 3], + /// [4, 5, 6], + /// [7, 8, 9] + /// ], + /// &device + /// ); + /// + /// let tensor = tensor.tril(-1); + /// println!("{tensor}"); + /// // [ + /// // [0, 0, 0], + /// // [4, 0, 0], + /// // [7, 8, 0] + /// // ] + /// } + /// ``` + pub fn tril(self, diagonal: i64) -> Self { + check!(TensorCheck::tri::<{ D }>()); + + // last two dimensions + let shape = &self.shape()[D - 2..].to_owned(); + let mask = Tensor::<2, Bool>::tril_mask(shape, diagonal, &self.device()).unsqueeze(); + + self.mask_fill(mask, 0) + } + + /// Applies element wise power operation with a integer Tensor + /// + /// # Arguments + /// + /// * `other` - The tensor to apply the power operation with. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape, Int}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor1 = Tensor::<2, Int>::from_ints([[1, -2, 3], [5, 9, 6]], &device); + /// let tensor2 = Tensor::<2, Int>::from_ints([[2, 3, 4], [1, 2, 3]], &device); + /// let tensor = tensor1.powi(tensor2); + /// println!("{tensor}"); + /// // [[1, -8, 81], [5, 81, 216]] + /// } + /// ``` + pub fn powi(self, other: Self) -> Self { + Self::new(K::powi(self.primitive, other.primitive)) + } + + /// Applies element wise power operation with a integer scalar + /// + /// # Arguments + /// + /// * `other` - The scalar to apply the power operation with. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape, Int}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2, Int>::from_ints([[1, -2, 3], [5, 9, 6]], &device); + /// let tensor = tensor.powi_scalar(2); + /// println!("{tensor}"); + /// + /// // [[1, 4, 9], [25, 81, 36]] + /// let tensor = Tensor::<2>::from_data([[1.5, -2., 3.], [5., 9., 6.]], &device); + /// let tensor = tensor.powi_scalar(2); + /// println!("{tensor}"); + /// // [[2.25, 4., 9.], [25., 81., 36.]] + /// } + /// ``` + pub fn powi_scalar(self, other: E) -> Self { + let other = Scalar::new(other, &self.dtype()); + Self::new(K::powi_scalar(self.primitive, other)) + } + + /// Converts the tensor to a boolean tensor by checking if the elements are non-zero. + /// + /// # Returns + /// + /// A boolean tensor with the same shape as the input tensor. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [0.0, 9.0, 6.0]], &device); + /// let tensor = tensor.bool(); + /// println!("{tensor}"); + /// // [ + /// // [true, true, true], + /// // [false, true, true] + /// // ] + /// } + pub fn bool(self) -> Tensor { + self.not_equal_elem(0) + } + + /// Create a random tensor of the given shape on the given device where each element is + /// sampled from the given distribution. + /// + /// See also [`random_like`](Tensor::random_like). + /// + /// # Arguments + /// + /// * `shape` - The shape of the tensor. + /// * `distribution` - The distribution to sample from. + /// * `device` - The device to create the tensor on. + /// + /// # Returns + /// + /// A new tensor with the given shape and elements sampled from the given distribution. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape, Distribution}; + /// + /// fn example() { + /// let device = Default::default(); + /// let distribution = Distribution::Uniform(0.0, 1.0); // Any random value between 0.0 and 1.0 + /// let tensor = Tensor::<2>::random(Shape::new([2, 3]), distribution, &device); + /// println!("{tensor}"); + /// // [ + /// // [0.08347523, 0.70498955, 0.60332155], + /// // [0.08173251, 0.18028641, 0.97942924] + /// // ] + /// } + /// ``` + pub fn random>( + shape: S, + distribution: Distribution, + options: impl Into, + ) -> Self { + // Use the given dtype when provided, otherwise default device dtype + let opt = options.into(); + let dtype = opt.resolve_dtype::(); + Self::new(K::random(shape.into(), distribution, &opt.device, dtype)) + } + + /// Applies the matrix multiplication operation. + /// + /// ```math + /// C = AB + /// ``` + /// + /// Shapes of the form `[..., B, 1, K] @ [..., 1, K, N]` are reinterpreted as + /// `[..., 1, B, K] @ [..., 1, K, N]`, turning a batched vec-mat into a general + /// matmul, which is often faster. + pub fn matmul(self, other: Self) -> Self { + check!(TensorCheck::matmul(&self, &other)); + + if D >= 3 { + let batch_index = D - 3; + let vector_index = D - 2; + let lhs_dims = &self.shape()[batch_index..D]; + let rhs_dims = &other.shape()[batch_index..D]; + + if let ([_, 1, k1], [1, k2, _]) = (lhs_dims, rhs_dims) + && k1 == k2 + { + return Tensor::new(K::matmul( + self.swap_dims(batch_index, vector_index).primitive, + other.primitive, + )) + .swap_dims(batch_index, vector_index); + } + } + + Tensor::new(K::matmul(self.primitive, other.primitive)) + } +} + +impl Tensor<1, K> +where + K: Numeric, +{ + /// Calculates the dot product with another tensor. + /// + /// `y = x2.dot(x1)` + /// + /// # Arguments + /// + /// * `other` - The tensor to compute dot product with. + /// + /// # Notes + /// + /// Both tensors must have the same number of elements. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor1 = Tensor::<1>::from_data([1.0, 2.0], &device); + /// let tensor2 = Tensor::<1>::from_data([-2.0, 3.0], &device); + /// let tensor = tensor1.dot(tensor2); + /// println!("{tensor}"); + /// // [4] + /// } + /// ``` + pub fn dot(self, other: Self) -> Self { + self.mul(other).sum() + } +} + +impl Tensor<2, K> +where + K: Numeric, +{ + /// Creates a new 2D tensor with ones on the diagonal and zeros elsewhere. + /// + /// # Arguments + /// + /// * `size` - The size of the square matrix. + pub fn eye(size: usize, device: &Device) -> Self { + let indices = Tensor::<1, Int>::arange(0..size as i64, device).unsqueeze::<2>(); + let ones = Self::ones([1, size], device); + let zeros = Self::zeros([size, size], device); + + zeros.scatter(0, indices, ones, IndexingUpdateOp::Add) + } +} + +// Tensor + tensor +impl core::ops::Add for Tensor { + type Output = Self; + + fn add(self, rhs: Self) -> Self::Output { + Self::add(self, rhs) + } +} + +// Tensor + scalar +impl core::ops::Add for Tensor { + type Output = Self; + + fn add(self, other: E) -> Self::Output { + Tensor::add_scalar(self, other) + } +} + +// Scalar + tensor +macro_rules! impl_tensor_scalar_add { + ($($t:ty),*) => { + $( + impl core::ops::Add> for $t + { + type Output = Tensor; + + fn add(self, tensor: Tensor) -> Self::Output { + Tensor::add_scalar(tensor, self) + } + } + )* + } +} +impl_tensor_scalar_add!(f32, f64, i32, i64, u32, u64); + +// Tensor - tensor +impl core::ops::Sub for Tensor { + type Output = Self; + + fn sub(self, rhs: Self) -> Self::Output { + Tensor::sub(self, rhs) + } +} + +// Tensor - scalar +impl core::ops::Sub for Tensor { + type Output = Self; + + fn sub(self, other: E) -> Self::Output { + Tensor::sub_scalar(self, other) + } +} + +// Scalar - tensor +macro_rules! impl_tensor_scalar_sub { + ($($t:ty),*) => { + $( + impl core::ops::Sub> for $t + { + type Output = Tensor; + + fn sub(self, tensor: Tensor) -> Self::Output { + Tensor::add_scalar(Tensor::neg(tensor), self) + } + } + )* + } +} +impl_tensor_scalar_sub!(f32, f64, i32, i64, u32, u64); + +// Tensor / tensor +impl core::ops::Div for Tensor { + type Output = Self; + + fn div(self, rhs: Self) -> Self::Output { + Tensor::div(self, rhs) + } +} + +// Tensor / scalar +impl core::ops::Div for Tensor { + type Output = Self; + + fn div(self, other: E) -> Self::Output { + Tensor::div_scalar(self, other) + } +} + +// Scalar / tensor (float only) +macro_rules! impl_tensor_scalar_div { + ($($t:ty),*) => { + $( + impl core::ops::Div> for $t + { + type Output = Tensor; + + fn div(self, tensor: Tensor) -> Self::Output { + tensor.recip().mul_scalar(self) + } + } + )* + } +} + +impl_tensor_scalar_div!(f32, f64); + +// Tensor % tensor. +impl core::ops::Rem for Tensor { + type Output = Self; + + fn rem(self, rhs: Self) -> Self::Output { + Tensor::remainder(self, rhs) + } +} + +// Tensor % scalar. +impl core::ops::Rem for Tensor { + type Output = Self; + + fn rem(self, other: E) -> Self::Output { + Tensor::remainder_scalar(self, other) + } +} + +// Tensor * tensor. +impl core::ops::Mul for Tensor { + type Output = Self; + + fn mul(self, rhs: Self) -> Self::Output { + Tensor::mul(self, rhs) + } +} + +// Tensor * scalar. +impl core::ops::Mul for Tensor { + type Output = Self; + + fn mul(self, other: E) -> Self::Output { + Tensor::mul_scalar(self, other) + } +} + +macro_rules! impl_tensor_scalar_mul { + ($($t:ty),*) => { + $( + impl core::ops::Mul> for $t + { + type Output = Tensor; + + fn mul(self, other: Tensor) -> Self::Output { + Tensor::mul_scalar(other, self) + } + } + )* + } +} + +impl_tensor_scalar_mul!(f32, f64, i32, i64, u32, u64); + +impl core::ops::Neg for Tensor +where + K: Numeric, +{ + type Output = Self; + + fn neg(self) -> Self::Output { + Tensor::neg(self) + } +} diff --git a/crates/burn-tensor/src/tensor/api/options.rs b/crates/burn-tensor/src/tensor/api/options.rs new file mode 100644 index 0000000..21cdf94 --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/options.rs @@ -0,0 +1,98 @@ +use burn_std::DType; + +use crate::{Device, bridge::BasicOps, ops::Kind}; + +/// Options for tensor creation. +/// +/// This struct allows specifying the `device` and overriding the data type when creating a tensor. +/// When the `dtype` is not specified, the [device's default settings](crate::DeviceSettings) is used. +#[derive(Debug, Clone)] +pub struct TensorCreationOptions { + /// Device where the tensor will be created. + pub device: Device, + /// Optional data type. + /// If `None`, the dtype will be inferred on creation from the [device settings](crate::DeviceSettings). + pub dtype: Option, +} + +impl Default for TensorCreationOptions { + /// Returns new options with the backend's default device. + fn default() -> Self { + Self::new(Default::default()) + } +} + +impl TensorCreationOptions { + /// Create new options with a specific device. + /// + /// Data type will follow the [device settings](crate::DeviceSettings) on tensor creation. + pub fn new(device: Device) -> Self { + Self { + device, + dtype: None, + } + } + + /// Set the tensor creation data type. + pub fn with_dtype(mut self, dtype: DType) -> Self { + self.dtype = Some(dtype); + + self + } + + /// Set the tensor creation device. + pub fn with_device(mut self, device: Device) -> Self { + self.device = device; + + self + } + + /// Returns the tensor data type, or a provided default if not set. + /// + /// This is useful for cases where [`TensorCreationOptions`] may not have an explicit `dtype`. + pub fn dtype_or(&self, dtype: DType) -> DType { + self.dtype.unwrap_or(dtype) + } + + /// Returns the tensor data type, or the default from the [device settings](crate::DeviceSettings). + pub(crate) fn resolve_dtype(&self) -> DType { + let settings = self.device.settings(); + let default = match K::KIND { + Kind::Float => settings.float_dtype.into(), + Kind::Int => settings.int_dtype.into(), + Kind::Bool => settings.bool_dtype.into(), + }; + match K::KIND { + // `BoolStore` variants are backend-specific storage layouts rather than + // semantic types, and a requested variant (e.g. forced from a serialized + // snapshot by `burn-store`) may be unsupported on the target device, so + // always resolve bool tensors to the device's default bool storage. + Kind::Bool => default, + _ => self.dtype.unwrap_or(default), + } + } +} + +impl From<&Device> for TensorCreationOptions { + /// Convenience conversion from a reference to a device. + /// + /// Example: + /// ```rust + /// use burn_tensor::TensorCreationOptions; + /// use burn_tensor::Device; + /// + /// fn example(device: Device) { + /// let options: TensorCreationOptions = (&device).into(); + /// } + /// ``` + fn from(device: &Device) -> Self { + TensorCreationOptions::new(device.clone()) + } +} + +impl From<(&crate::Device, DType)> for TensorCreationOptions { + /// Convenience conversion for a specified `(&device, dtype)` tuple. + fn from(args: (&crate::Device, DType)) -> Self { + TensorCreationOptions::new(args.0.clone()).with_dtype(args.1) + } +} diff --git a/crates/burn-tensor/src/tensor/api/orderable.rs b/crates/burn-tensor/src/tensor/api/orderable.rs new file mode 100644 index 0000000..8c921e0 --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/orderable.rs @@ -0,0 +1,1132 @@ +use burn_backend::{ElementConversion, Scalar}; +use burn_std::{AsIndex, IndexingUpdateOp}; + +use crate::kind::Ordered; +use crate::{Bool, Int, check}; +use crate::{Tensor, check::TensorCheck}; + +impl Tensor +where + K: Ordered, +{ + /// Sort the elements by value in ascending order along a given dimension. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `dim` - The dimension to sort along. + /// + /// # Returns + /// + /// A new tensor with the elements sorted in ascending order along the given dimension. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device); + /// let tensor = tensor.sort(0); + /// println!("{tensor}"); + /// // [[5.0, -2.0, 3.0], [12.0, 3.0, 6.0]] + /// let tensor = tensor.sort(1); + /// println!("{tensor}"); + /// // [[-2.0, 3.0, 12.0], [3.0, 5.0, 6.0]] + /// } + /// ``` + pub fn sort(self, dim: usize) -> Self { + check!(TensorCheck::sort_dim::("Sort", dim)); + Tensor::new(K::sort(self.primitive, dim, /*descending*/ false)) + } + + /// Sort the elements by value in descending order along a given dimension. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `dim` - The dimension to sort along. + /// + /// # Returns + /// + /// A new tensor with the elements sorted in descending order along the given dimension. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device); + /// let tensor = tensor.sort_descending(0); + /// println!("{tensor}"); + /// // [[12.0, 3.0, 6.0], [5.0, -2.0, 3.0]] + /// let tensor = tensor.sort_descending(1); + /// println!("{tensor}"); + /// // [[12.0, 3.0, -2.0], [6.0, 5.0, 3.0]] + /// } + /// ``` + pub fn sort_descending(self, dim: usize) -> Self { + check!(TensorCheck::sort_dim::("Sort", dim)); + Tensor::new(K::sort(self.primitive, dim, /*descending*/ true)) + } + + /// Sort the elements by value in ascending order along a given dimension. + /// Also returns the indices. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `dim` - The dimension to sort along. + /// + /// # Returns + /// + /// A tuple containing the sorted tensor and the indices tensor. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device); + /// let (tensor, indices) = tensor.sort_with_indices(0); + /// println!("{tensor}"); + /// // [[5.0, -2.0, 3.0], [12.0, 3.0, 6.0]] + /// println!("{}", indices); + /// // [[1, 0, 0], [0, 1, 1]] + /// } + /// ``` + pub fn sort_with_indices(self, dim: usize) -> (Self, Tensor) { + check!(TensorCheck::sort_dim::("Sort_with_indices", dim)); + let (values, indices) = + K::sort_with_indices(self.primitive, dim, /*descending*/ false); + (Tensor::new(values), Tensor::new(indices)) + } + + /// Sort the elements by value in descending order along a given dimension. + /// Also returns the indices. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `dim` - The dimension to sort along. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device); + /// let (tensor, indices) = tensor.sort_descending_with_indices(0); + /// println!("{tensor}"); + /// // [[12.0, 3.0, 6.0], [5.0, -2.0, 3.0]] + /// println!("{}", indices); + /// // [[0, 1, 1], [1, 0, 0]] + /// } + /// ``` + pub fn sort_descending_with_indices(self, dim: usize) -> (Self, Tensor) { + check!(TensorCheck::sort_dim::("Sort_with_indices", dim)); + let (values, indices) = K::sort_with_indices(self.primitive, dim, /*descending*/ true); + (Tensor::new(values), Tensor::new(indices)) + } + + /// Returns the indices that sort the elements by value in ascending order along a given dimension. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `dim` - The dimension to sort along. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device); + /// let tensor = tensor.argsort(0); + /// println!("{tensor}"); + /// // [[1, 0, 0], [0, 1, 1]] + /// } + /// ``` + pub fn argsort(self, dim: usize) -> Tensor { + check!(TensorCheck::sort_dim::("Argsort", dim)); + Tensor::new(K::argsort(self.primitive, dim, /*descending*/ false)) + } + + /// Returns the indices that sort the elements by value in descending order along a given dimension. + /// + /// This sort is unstable (i.e., may reorder equal elements). + /// + /// # Arguments + /// + /// * `dim` - The dimension to sort along. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device); + /// let tensor = tensor.argsort_descending(0); + /// println!("{tensor}"); + /// // [[0, 1, 1], [1, 0, 0]] + /// let tensor = tensor.argsort_descending(1); + /// println!("{tensor}"); + /// // [[0, 2, 1], [2, 0, 1]] + /// } + /// ``` + pub fn argsort_descending(self, dim: usize) -> Tensor { + check!(TensorCheck::sort_dim::("Argsort", dim)); + Tensor::new(K::argsort(self.primitive, dim, /*descending*/ true)) + } + + /// Returns the `k` largest elements of the given input tensor along a given dimension. + /// + /// # Arguments + /// + /// * `k` - The number of elements to return. + /// + /// # Returns + /// + /// A new tensor with the `k` largest elements along the given dimension. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device); + /// let tensor = tensor.topk(2, 0); + /// println!("{tensor}"); + /// // [[12.0, 3.0, 6.0], [5.0, -2.0, 3.0]] + /// let tensor = tensor.topk(1, 1); + /// println!("{tensor}"); + /// // [[12.0], [6.0]] + /// } + /// ``` + pub fn topk(self, k: usize, dim: usize) -> Self { + assert!(self.shape()[dim] > k); + Tensor::new(K::topk(self.primitive, dim, k)) + } + + /// Returns the `k` largest elements of the given input tensor along a given dimension. + /// Also returns the indices. + /// + /// # Arguments + /// + /// * `k` - The number of elements to return. + /// * `dim` - The dimension to sort along. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device); + /// let (tensor, indices) = tensor.topk_with_indices(2, 0); + /// println!("{tensor}"); + /// // [[12.0, 3.0, 6.0], [5.0, -2.0, 3.0]] + /// println!("{}", indices); + /// // [[0, 1, 1], [1, 0, 0]] + /// let (tensor, indices) = tensor.topk_with_indices(1, 1); + /// println!("{tensor}"); + /// // [[12.0], [6.0]] + /// println!("{indices}"); + /// // [[0], [2]] + /// } + /// ``` + pub fn topk_with_indices(self, k: usize, dim: usize) -> (Self, Tensor) { + let k_indices = Tensor::arange(0..k as i64, &self.device()); + let (values, indices) = self.sort_descending_with_indices(dim); + ( + values.select(dim, k_indices.clone()), + indices.select(dim, k_indices), + ) + } + + /// Create a one hot tensor. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example(){ + /// let device = Default::default(); + /// let indices: Tensor<1> = Tensor::from_floats([0.0, 1.0, 2.0, 3.0], &device); + /// let one_hot: Tensor<2> = indices.one_hot(4); + /// println!("{}", one_hot.to_data()); + /// // [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]] + /// } + /// ``` + pub fn one_hot(self, num_classes: usize) -> Tensor { + check!(TensorCheck::one_hot_tensor(self.clone(), num_classes)); + self.one_hot_fill(num_classes, 1.0, 0.0, -1) + } + + /// Create a one-hot encoded tensor with configurable `num_classes`, `on_value`, `off_value`, and `axis` including high-ranked tensors. + /// + /// # Arguments + /// + /// * `num_classes`: The number of classes for the one-hot encoding, which defines the size of the one-hot dimension. + /// * `on_value`: The value to assign for active positions (corresponding to indices). + /// * `off_value`: The value to assign for inactive positions. + /// * `axis`: The axis along which the one-hot dimension is added. Supports negative indexing. + /// + /// # Returns + /// + /// A tensor with one additional dimension for the one-hot encoding, where active positions are filled with `on_value` and others with `off_value`. + /// + /// # Example + /// ```rust + /// use burn_tensor::{Tensor, Float}; + /// fn example() { + /// let device = Default::default(); + /// let indices: Tensor<2, Float> = Tensor::from_floats([[0., 2.], [1., -1.]], &device); + /// // One-hot encoding + /// let tensor: Tensor<3, Float> = indices.one_hot_fill(3, 5.0.into(), 0.0.into(), -1); + /// println!("{tensor}"); + /// // [[[5.0, 0.0, 0.0], + /// // [0.0, 0.0, 5.0]], + /// // [[0.0, 5.0, 0.0], + /// // [0.0, 0.0, 5.0]]] + /// } + /// ``` + pub fn one_hot_fill( + self, + num_classes: usize, + on_value: f32, + off_value: f32, + axis: i64, + ) -> Tensor { + check!(TensorCheck::one_hot_tensor_rank::()); + // Initialize shape from the current tensor dimensions and prepare for modification + let mut shape = self.shape(); + let device = self.device(); + let rank = self.dims().len(); + + // Adjust negative axis to a positive index + let axis = if axis < 0 { + axis + rank as i64 + 1 + } else { + axis + }; + + // Ensure axis is within valid range + if axis < 0 || axis > rank as i64 { + panic!("Axis out of range. Accepted range is [-r-1, r] where r = rank(indices)."); + } + // Convert the input tensor to integer indices + let indices: Tensor = Tensor::from_data(self.to_data().convert::(), &device); + // Insert the new dimension for the one-hot representation + shape.insert(axis as usize, num_classes); + // Adjust indices to valid range and handle invalid indices + let adjusted_indices = indices + .clone() + .mask_fill(self.clone().lower_elem(0), num_classes as i64) // Handle negative indices + .add(indices.clone().mask_fill(self.clone().greater_elem(0), 0)); // Handle positive indices + // Unsqueeze the indices tensor along the specified axis + let indices_unsqueezed: Tensor = adjusted_indices.unsqueeze_dim(axis as usize); + + // Initialize the output tensor with the off_value + let output = Tensor::full(shape.clone(), off_value, &device); + + // Prepare scatter tensor for on_value and off_value adjustments + let scatter_on_values = Tensor::full(indices_unsqueezed.shape(), on_value, &device) + - Tensor::full(indices_unsqueezed.shape(), off_value, &self.device()); + + // Scatter on_value at the appropriate indices to create the one-hot representation + output.scatter( + axis as usize, + indices_unsqueezed, + scatter_on_values, + IndexingUpdateOp::Add, + ) + } + + /// Applies element wise greater comparison and returns a boolean tensor. + /// + /// # Panics + /// + /// If the two tensors don't have the same shape. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor2 = Tensor::<2>::from_data([[1.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device); + /// let tensor = tensor1.greater(tensor2); + /// println!("{tensor}"); + /// // [[false, false, false], [true, true, true]] + /// } + /// ``` + pub fn greater(self, other: Self) -> Tensor { + check!(TensorCheck::binary_ops_ew("Greater", &self, &other)); + Tensor::new(K::greater(self.primitive, other.primitive)) + } + + /// Applies element wise greater-equal comparison and returns a boolean tensor. + /// + /// # Panics + /// + /// If the two tensors don't have the same shape. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor2 = Tensor::<2>::from_data([[1.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device); + /// let tensor = tensor1.greater_equal(tensor2); + /// println!("{tensor}"); + /// // [[true, false, false], [true, true, true]] + /// } + /// ``` + pub fn greater_equal(self, other: Self) -> Tensor { + check!(TensorCheck::binary_ops_ew("Greater_equal", &self, &other)); + Tensor::new(K::greater_equal(self.primitive, other.primitive)) + } + + /// Applies element wise lower comparison and returns a boolean tensor. + /// + /// # Panics + /// + /// If the two tensors don't have the same shape. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor2 = Tensor::<2>::from_data([[1.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device); + /// let tensor = tensor1.lower(tensor2); + /// println!("{tensor}"); + /// // [[false, true, true], [false, false, false]] + /// } + /// ``` + pub fn lower(self, other: Self) -> Tensor { + check!(TensorCheck::binary_ops_ew("Lower", &self, &other)); + Tensor::new(K::lower(self.primitive, other.primitive)) + } + + /// Applies element wise lower-equal comparison and returns a boolean tensor. + /// + /// # Panics + /// + /// If the two tensors don't have the same shape. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor2 = Tensor::<2>::from_data([[1.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device); + /// let tensor = tensor1.lower_equal(tensor2); + /// println!("{tensor}"); + /// // [[true, true, true], [false, false, false]] + /// } + /// ``` + pub fn lower_equal(self, other: Self) -> Tensor { + check!(TensorCheck::binary_ops_ew("Lower_equal", &self, &other)); + Tensor::new(K::lower_equal(self.primitive, other.primitive)) + } + + /// Applies greater than `other` comparison and returns a boolean tensor. + /// + /// # Arguments + /// + /// * `other` - The element to compare. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.greater_elem(3.0); + /// println!("{tensor}"); + /// // [[false, false, true], [true, true, true]] + /// } + /// ``` + pub fn greater_elem(self, other: E) -> Tensor { + let other = Scalar::new(other, &self.dtype()); + Tensor::new(K::greater_elem(self.primitive, other)) + } + + /// Applies greater-equal than `other` comparison and returns a boolean tensor. + /// + /// # Arguments + /// + /// * `other` - The element to compare. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.greater_equal_elem(3.0); + /// println!("{tensor}"); + /// // [[false, false, true], [true, true, true]] + /// } + /// ``` + pub fn greater_equal_elem(self, other: E) -> Tensor { + let other = Scalar::new(other, &self.dtype()); + Tensor::new(K::greater_equal_elem(self.primitive, other)) + } + + /// Applies lower than `other` comparison and returns a boolean tensor. + /// + /// # Arguments + /// + /// * `other` - The element to compare. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.lower_elem(3.0); + /// println!("{tensor}"); + /// // [[true, true, false], [false, false, false]] + /// } + /// ``` + pub fn lower_elem(self, other: E) -> Tensor { + let other = Scalar::new(other, &self.dtype()); + Tensor::new(K::lower_elem(self.primitive, other)) + } + + /// Applies lower-equal than `other` comparison and returns a boolean tensor. + /// + /// # Arguments + /// + /// * `other` - The element to compare. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.lower_equal_elem(3.0); + /// println!("{tensor}"); + /// // [[true, true, true], [false, false, false]] + /// } + /// ``` + pub fn lower_equal_elem(self, other: E) -> Tensor { + let other = Scalar::new(other, &self.dtype()); + Tensor::new(K::lower_equal_elem(self.primitive, other)) + } + + /// Applies the argmax function along the given dimension and returns an integer tensor. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<3>::ones(Shape::new([2, 3, 3]), &device); + /// let tensor = tensor.argmax(1); + /// println!("{:?}", tensor.shape()); + /// // Shape { dims: [2, 1, 3] } + /// } + /// ``` + pub fn argmax(self, dim: usize) -> Tensor { + Tensor::new(K::argmax(self.primitive, dim)) + } + + /// Applies the argtopk function along the given dimension and returns an integer tensor. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<3>::ones(Shape::new([2, 3, 3]), &device); + /// let tensor = tensor.argtopk(1, 2); + /// println!("{:?}", tensor.shape()); + /// } + /// ``` + pub fn argtopk(self, k: usize, dim: usize) -> Tensor { + assert!(self.shape()[dim] > k); + Tensor::new(K::argtopk(self.primitive, dim, k)) + } + + /// Find the maximum value. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.max(); + /// println!("{tensor}"); + /// // [9.0] + /// } + /// ``` + pub fn max(self) -> Tensor<1, K> { + Tensor::new(K::max(self.primitive)) + } + + /// Find the maximum value along the given dimension. + /// + /// Also returns the indices. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let (tensor, index) = tensor.max_dim_with_indices(0); + /// // [[5.0, 9.0, 6.0]] + /// println!("{tensor}"); + /// // [[1, 1, 1]] + /// println!("{index}"); + /// } + /// ``` + pub fn max_dim_with_indices(self, dim: I) -> (Self, Tensor) { + let dim = dim.expect_dim_index(D); + check!(TensorCheck::aggregate_dim::("Max", dim)); + + let (tensor, index) = K::max_dim_with_indices(self.primitive, dim); + + let tensor = Tensor::new(tensor); + let index = Tensor::new(index); + + (tensor, index) + } + + /// Find the maximum absolute value. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -7.0, 3.0], [5.0, -1.0, 6.0]], &device); + /// let tensor = tensor.max_abs(); + /// println!("{tensor}"); + /// // [7.0] + /// } + /// ``` + pub fn max_abs(self) -> Tensor<1, K> { + Tensor::new(K::max_abs(self.primitive)) + } + + /// Finds the maximum pair wise values with another tensor. + /// + /// # Arguments + /// + /// * `other` - Other tensor to find maximum elements with + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensors containing the maximum value found + /// in the input tensors. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device); + /// let tensor = tensor1.max_pair(tensor2); + /// println!("{tensor}"); + /// // [[2.0, 3.0, 4.0], [5.0, 9.0, 6.0]] + /// } + /// ``` + pub fn max_pair(self, other: Self) -> Self { + let mask = self.clone().lower(other.clone()); + self.mask_where(mask, other) + } + + /// Find the maximum absolute value along the given dimension. + /// + /// # Arguments + /// + /// * `dim` - The dimension or axis along which to aggregate the elements, + /// supports negative indexing. + /// + /// # Returns + /// + /// The returned tensor will have the same rank, + /// but the aggregated dimension will have size 1. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.max_dim(0); + /// println!("{tensor}"); + /// // [[5.0, 9.0, 6.0]] + /// } + /// ``` + pub fn max_abs_dim(self, dim: I) -> Self { + let dim = dim.expect_dim_index(D); + check!(TensorCheck::aggregate_dim::("MaxAbs", dim)); + + Tensor::new(K::max_abs_dim(self.primitive, dim)) + } + + /// Find the maximum absolute value along the given dimensions. + /// + /// # Arguments + /// + /// * `dims` - The dimensions or axes along which to aggregate the elements, + /// supports negative indexing. + /// + /// # Returns + /// + /// The returned tensor will have the same rank, + /// but the aggregated dimensions will have size 1. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.max_abs_dims(&[0, 1]); + /// println!("{tensor}"); + /// // [[9.0]] + /// } + /// ``` + pub fn max_abs_dims(self, dims: &[I]) -> Self { + dims.iter() + .fold(self, |tensor, &dim| tensor.max_abs_dim(dim)) + } + + /// Applies the argmin function along the given dimension and returns an integer tensor. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<3>::ones(Shape::new([2, 3, 3]), &device); + /// let tensor = tensor.argmin(1); + /// println!("{:?}", tensor.shape()); + /// // Shape { dims: [2, 1, 3] } + /// } + /// ``` + pub fn argmin(self, dim: usize) -> Tensor { + Tensor::new(K::argmin(self.primitive, dim)) + } + + /// Find the minimum value. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.min(); + /// println!("{tensor}"); + /// // [-2.0] + /// } + /// ``` + pub fn min(self) -> Tensor<1, K> { + Tensor::new(K::min(self.primitive)) + } + + /// Find the minimum value along the given dimension. + /// + /// # Arguments + /// + /// * `dim` - The dimension or axis along which to aggregate the elements; + /// supports negative indexing. + /// + /// # Returns + /// + /// The returned tensor will have the same rank, + /// but the aggregated dimension will have size 1. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.min_dim(0); + /// println!("{tensor}"); + /// // [[1.0, -2.0, 3.0]] + /// } + /// ``` + pub fn min_dim(self, dim: I) -> Self { + let dim = dim.expect_dim_index(D); + check!(TensorCheck::aggregate_dim::("Min", dim)); + Tensor::new(K::min_dim(self.primitive, dim)) + } + + /// Find the minimum value along the given dimensions. + /// + /// # Arguments + /// + /// * `dims` - The dimensions or axes along which to aggregate the elements; + /// supports negative indexing. + /// + /// # Returns + /// + /// The returned tensor will have the same rank, + /// but the aggregated dimensions will have size 1. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.min_dims(&[0, 1]); + /// println!("{tensor}"); + /// // [[-2.0]] + /// } + /// ``` + pub fn min_dims(self, dims: &[I]) -> Self { + dims.iter().fold(self, |tensor, &dim| tensor.min_dim(dim)) + } + + /// Find the minimum value along the given dimension. + /// + /// Also returns the indices. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[7.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let (tensor, index) = tensor.min_dim_with_indices(0); + /// println!("{tensor}"); + /// // [[5.0, -2.0, 3.0]] + /// println!("{}", index); + /// // [[1, 0, 0]] + /// } + /// ``` + pub fn min_dim_with_indices(self, dim: I) -> (Self, Tensor) { + let dim = dim.expect_dim_index(D); + check!(TensorCheck::aggregate_dim::("Min", dim)); + + let (tensor, index) = K::min_dim_with_indices(self.primitive, dim); + + let tensor = Tensor::new(tensor); + let index = Tensor::new(index); + + (tensor, index) + } + + /// Finds the minimum pair wise values with another tensor. + /// + /// # Arguments + /// + /// * `other` - Other tensor to find minimum elements with + /// + /// # Returns + /// + /// A tensor with the same shape as the input tensors containing the minimum value found + /// between each element of the two source tensors. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device); + /// let tensor = tensor1.min_pair(tensor2); + /// println!("{tensor}"); + /// // [[1.0, -2.0, 3.0], [1.0, 2.0, 3.0]] + /// } + pub fn min_pair(self, other: Self) -> Self { + let mask = other.clone().lower(self.clone()); + self.mask_where(mask, other) + } + + /// Clamp element wise between the given min and max values. + /// + /// # Arguments + /// + /// * `min` - The minimum value. + /// * `max` - The maximum value. + /// + /// # Returns + /// + /// A new tensor with the values clamped between the given min and max values. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Int, Tensor}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2, Int>::from_ints( + /// [ + /// [1, 2, 3], + /// [4, 5, 6], + /// [7, 8, 9] + /// ], + /// &device); + /// let tensor = tensor.clamp(2, 6); + /// println!("{tensor}"); + /// // [[2, 2, 3], [4, 5, 6], [6, 6, 6]] + /// } + /// ``` + pub fn clamp(self, min: E, max: E) -> Self { + let dtype = self.dtype(); + Self::new(K::clamp( + self.primitive, + Scalar::new(min, &dtype), + Scalar::new(max, &dtype), + )) + } + + /// Clamp element wise under a minimum value. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to clamp. + /// * `min` - The minimum value. + /// + /// # Returns + /// + /// A new tensor with the values clamped under the given min value. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Int, Tensor}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2, Int>::from_ints( + /// [[1, 2, 3], [4, 5, 6], [7, 8, 9]], + /// &device); + /// let tensor = tensor.clamp_min(4); + /// println!("{tensor}"); + /// // [[4, 4, 4], [4, 5, 6], [7, 8, 9]] + /// } + /// ``` + pub fn clamp_min(self, min: E) -> Self { + let min = Scalar::new(min, &self.dtype()); + Self::new(K::clamp_min(self.primitive, min)) + } + + /// Clamp element wise over a maximum value. + /// + /// # Arguments + /// + /// * `tensor` - The tensor to clamp. + /// * `max` - The maximum value. + /// + /// # Returns + /// + /// A new tensor with the values clamped over the given max value. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Int, Tensor}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2, Int>::from_ints( + /// [[1, 2, 3], [4, 5, 6], [7, 8, 9]], + /// &device); + /// let tensor = tensor.clamp_max(5); + /// println!("{tensor}"); + /// // [[1, 2, 3], [4, 5, 5], [5, 5, 5]] + /// } + /// ``` + pub fn clamp_max(self, max: E) -> Self { + let max = Scalar::new(max, &self.dtype()); + Self::new(K::clamp_max(self.primitive, max)) + } + + /// Computes the cumulative minimum of elements along the given *dimension* or *axis*. + /// + /// # Arguments + /// + /// * `dim` - The dimension or axis along which to compute the cumulative minimum. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[3.0, 5.0, 2.0], [4.0, 1.0, 6.0]], &device); + /// let result = tensor.clone().cummin(0); + /// println!("{result}"); + /// // [[3.0, 5.0, 2.0], [3.0, 1.0, 2.0]] + /// let result = tensor.cummin(1); + /// println!("{result}"); + /// // [[3.0, 3.0, 2.0], [4.0, 1.0, 1.0]] + /// } + /// ``` + pub fn cummin(self, dim: usize) -> Self { + check!(TensorCheck::aggregate_dim::("CumMin", dim)); + Self::new(K::cummin(self.primitive, dim)) + } + + /// Computes the cumulative maximum of elements along the given *dimension* or *axis*. + /// + /// # Arguments + /// + /// * `dim` - The dimension or axis along which to compute the cumulative maximum. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[3.0, 1.0, 2.0], [4.0, 5.0, 2.0]], &device); + /// let result = tensor.clone().cummax(0); + /// println!("{result}"); + /// // [[3.0, 1.0, 2.0], [4.0, 5.0, 2.0]] + /// let result = tensor.cummax(1); + /// println!("{result}"); + /// // [[3.0, 3.0, 3.0], [4.0, 5.0, 5.0]] + /// } + /// ``` + pub fn cummax(self, dim: usize) -> Self { + check!(TensorCheck::aggregate_dim::("CumMax", dim)); + Self::new(K::cummax(self.primitive, dim)) + } + /// Find the maximum value along the given dimension. + /// + /// # Arguments + /// + /// * `dim` - The dimension or axis along which to aggregate the elements; + /// supports negative indexing. + /// + /// # Returns + /// + /// The returned tensor will have the same rank, + /// but the aggregated dimension will have size 1. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.max_dim(0); + /// println!("{tensor}"); + /// // [[5.0, 9.0, 6.0]] + /// } + /// ``` + pub fn max_dim(self, dim: I) -> Self { + let dim = dim.expect_dim_index(D); + check!(TensorCheck::aggregate_dim::("Max", dim)); + Tensor::new(K::max_dim(self.primitive, dim)) + } + + /// Find the maximum value along the given dimensions. + /// + /// # Arguments + /// + /// * `dims` - The dimensions or axis along which to aggregate the elements; + /// supports negative indexing. + /// + /// # Returns + /// + /// The returned tensor will have the same rank, + /// but the aggregated dimensions will have size 1. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device); + /// let tensor = tensor.max_dims(&[0, 1]); + /// println!("{tensor}"); + /// // [[9.0]] + /// } + /// ``` + pub fn max_dims(self, dims: &[I]) -> Self { + dims.iter().fold(self, |tensor, &dim| tensor.max_dim(dim)) + } +} diff --git a/crates/burn-tensor/src/tensor/api/pad.rs b/crates/burn-tensor/src/tensor/api/pad.rs new file mode 100644 index 0000000..8ef2ea0 --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/pad.rs @@ -0,0 +1,345 @@ +use alloc::vec::Vec; +use core::ops::Range; + +use crate::{ElementConversion, Tensor, kind::Numeric, ops::PadMode}; + +/// Trait for types that can be used as padding specifications. +/// +/// Padding is specified as `(before, after)` pairs per dimension, returned as a +/// fixed-size array `[(usize, usize); D]`. If fewer pairs than dimensions are provided, +/// they apply to the **last** N dimensions (earlier dimensions are left unpadded). +pub trait IntoPadding { + /// Converts into a fixed-size array of `(before, after)` padding pairs. + fn into_padding(self) -> [(usize, usize); D]; +} + +impl IntoPadding for [(usize, usize); N] { + fn into_padding(self) -> [(usize, usize); D] { + assert!( + N <= D, + "Padding has {} pairs but tensor only has {} dimensions", + N, + D + ); + let mut result = [(0usize, 0usize); D]; + let offset = D - N; + for (i, pair) in self.into_iter().enumerate() { + result[offset + i] = pair; + } + result + } +} + +/// Backward-compatible: `(left, right, top, bottom)` maps to last 2 dimensions. +/// +/// Equivalent to `[(top, bottom), (left, right)]`. +impl IntoPadding for (usize, usize, usize, usize) { + fn into_padding(self) -> [(usize, usize); D] { + let (left, right, top, bottom) = self; + let mut result = [(0usize, 0usize); D]; + result[D - 2] = (top, bottom); + result[D - 1] = (left, right); + result + } +} + +impl IntoPadding for &[(usize, usize)] { + fn into_padding(self) -> [(usize, usize); D] { + assert!( + self.len() <= D, + "Padding has {} pairs but tensor only has {} dimensions", + self.len(), + D + ); + let mut result = [(0usize, 0usize); D]; + let offset = D - self.len(); + for (i, &pair) in self.iter().enumerate() { + result[offset + i] = pair; + } + result + } +} + +impl IntoPadding for Vec<(usize, usize)> { + fn into_padding(self) -> [(usize, usize); D] { + assert!( + self.len() <= D, + "Padding has {} pairs but tensor only has {} dimensions", + self.len(), + D + ); + let mut result = [(0usize, 0usize); D]; + let offset = D - self.len(); + for (i, pair) in self.into_iter().enumerate() { + result[offset + i] = pair; + } + result + } +} + +/// Helper to build a range array for slice_assign, selecting a portion of one dimension. +fn build_slice_ranges( + dims: [usize; D], + target_dim: usize, + start: usize, + len: usize, +) -> [Range; D] { + dims.iter() + .enumerate() + .map(|(i, &size)| { + if i == target_dim { + start..start + len + } else { + 0..size + } + }) + .collect::>>() + .try_into() + .unwrap() +} + +impl Tensor +where + K: Numeric, +{ + /// Pads the tensor using the specified padding mode. + /// + /// Padding is specified as `(before, after)` pairs. If fewer pairs than tensor dimensions + /// are provided, they apply to the **last** N dimensions (unspecified leading dimensions + /// are left unpadded). + /// + /// For backward compatibility, a `(left, right, top, bottom)` tuple is also accepted, + /// which pads the last two dimensions. + /// + /// # Arguments + /// + /// * `padding` - Padding specification. Accepts: + /// - `[(before, after); N]` fixed-size array of pairs (N <= D) + /// - `&[(before, after)]` slice of pairs per dimension + /// - `Vec<(before, after)>` vector of pairs + /// - `(left, right, top, bottom)` tuple for last-2-dim backward compatibility + /// * `mode` - The padding mode: `Constant(value)`, `Reflect`, or `Edge`. + /// + /// # Returns + /// + /// A new tensor with the specified padding applied. + /// + /// # Panics + /// + /// - Panics if more padding pairs are provided than tensor dimensions. + /// - `Reflect` mode panics if padding exceeds `dimension_size - 1`. + /// - `Edge` mode panics if padding is applied to a zero-sized dimension. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Shape}; + /// use burn_tensor::ops::PadMode; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device); + /// + /// // Constant padding with value 0.0 (backward-compatible tuple) + /// let padded = tensor.clone().pad((1, 1, 1, 1), PadMode::Constant(0.0)); + /// + /// // Pad arbitrary dimensions with slice of (before, after) pairs + /// let padded = tensor.clone().pad([(1, 1), (2, 2)], PadMode::Constant(0.0)); + /// + /// // Pad only the last dimension + /// let padded = tensor.pad([(1, 1)], PadMode::Reflect); + /// } + /// ``` + pub fn pad(self, padding: impl IntoPadding, mode: impl Into) -> Self { + let pairs = padding.into_padding(); + match mode.into() { + PadMode::Constant(value) => pad_constant(self, &pairs, value), + PadMode::Reflect => pad_reflect(self, &pairs), + PadMode::Edge => pad_edge(self, &pairs), + } + } +} + +/// Pad with a constant value. +fn pad_constant( + tensor: Tensor, + padding: &[(usize, usize); D], + value: E, +) -> Tensor +where + K: Numeric, + E: ElementConversion, +{ + let mut padded_dims: [usize; D] = tensor.dims(); + + for (i, &(before, after)) in padding.iter().enumerate() { + padded_dims[i] += before + after; + } + + let ranges: [Range; D] = padded_dims + .iter() + .enumerate() + .map(|(i, &dim)| { + let (before, after) = padding[i]; + before..dim - after + }) + .collect::>>() + .try_into() + .unwrap(); + + let padded_tensor = Tensor::full(padded_dims, value, &tensor.device()); + + padded_tensor.slice_assign(ranges, tensor) +} + +/// Pad using reflection at the boundaries (excluding edge values). +/// +/// For ONNX "reflect" mode: mirrors from index 1, not index 0. +/// Example: `[1, 2, 3, 4]` with left padding 2 becomes `[3, 2, 1, 2, 3, 4]` +fn pad_reflect( + tensor: Tensor, + padding: &[(usize, usize); D], +) -> Tensor +where + K: Numeric, +{ + let dims = tensor.dims(); + + for (i, &(before, after)) in padding.iter().enumerate() { + if before > 0 || after > 0 { + assert!( + before < dims[i] && after < dims[i], + "Reflect padding ({}, {}) must be less than dimension {} size ({})", + before, + after, + i, + dims[i] + ); + } + } + + let mut result = tensor; + + for (i, &(before, after)) in padding.iter().enumerate() { + if before > 0 || after > 0 { + result = pad_reflect_dim(result, i, before, after); + } + } + + result +} + +/// Helper to pad a single dimension using reflection. +fn pad_reflect_dim( + tensor: Tensor, + dim: usize, + pad_before: usize, + pad_after: usize, +) -> Tensor +where + K: Numeric, +{ + let dims = tensor.dims(); + let dim_size = dims[dim]; + + // Calculate output dimensions + let mut output_dims = dims; + output_dims[dim] += pad_before + pad_after; + + // Create output tensor and place original in the center + let output = Tensor::zeros(output_dims, &tensor.device()); + let original_range = build_slice_ranges(output_dims, dim, pad_before, dim_size); + let mut output = output.slice_assign(original_range, tensor.clone()); + + // Assign reflected "before" padding (e.g., top or left) + // Reflect excludes the edge, so we take indices [1..pad_before+1] and flip + if pad_before > 0 { + let before_slice = tensor.clone().narrow(dim, 1, pad_before); + let before_flipped = before_slice.flip([dim as isize]); + let before_range = build_slice_ranges(output_dims, dim, 0, pad_before); + output = output.slice_assign(before_range, before_flipped); + } + + // Assign reflected "after" padding (e.g., bottom or right) + // Take indices [dim_size - pad_after - 1..dim_size - 1] and flip + if pad_after > 0 { + let start = dim_size - pad_after - 1; + let after_slice = tensor.narrow(dim, start, pad_after); + let after_flipped = after_slice.flip([dim as isize]); + let after_range = build_slice_ranges(output_dims, dim, pad_before + dim_size, pad_after); + output = output.slice_assign(after_range, after_flipped); + } + + output +} + +/// Pad by replicating edge values. +/// +/// Example: `[1, 2, 3, 4]` with left padding 2 becomes `[1, 1, 1, 2, 3, 4]` +fn pad_edge(tensor: Tensor, padding: &[(usize, usize); D]) -> Tensor +where + K: Numeric, +{ + let dims = tensor.dims(); + + for (i, &(before, after)) in padding.iter().enumerate() { + if before > 0 || after > 0 { + assert!( + dims[i] > 0, + "Cannot apply edge padding to zero-sized dimension {}", + i + ); + } + } + + let mut result = tensor; + + for (i, &(before, after)) in padding.iter().enumerate() { + if before > 0 || after > 0 { + result = pad_edge_dim(result, i, before, after); + } + } + + result +} + +/// Helper to pad a single dimension by replicating edge values. +fn pad_edge_dim( + tensor: Tensor, + dim: usize, + pad_before: usize, + pad_after: usize, +) -> Tensor +where + K: Numeric, +{ + let dims = tensor.dims(); + let dim_size = dims[dim]; + + // Calculate output dimensions + let mut output_dims = dims; + output_dims[dim] += pad_before + pad_after; + + // Create output tensor and place original in the center + let output = Tensor::zeros(output_dims, &tensor.device()); + let original_range = build_slice_ranges(output_dims, dim, pad_before, dim_size); + let mut output = output.slice_assign(original_range, tensor.clone()); + + // Assign "before" padding by repeating the first element + if pad_before > 0 { + let first_slice = tensor.clone().narrow(dim, 0, 1); + let before_pad = first_slice.repeat_dim(dim, pad_before); + let before_range = build_slice_ranges(output_dims, dim, 0, pad_before); + output = output.slice_assign(before_range, before_pad); + } + + // Assign "after" padding by repeating the last element + if pad_after > 0 { + let last_slice = tensor.narrow(dim, dim_size - 1, 1); + let after_pad = last_slice.repeat_dim(dim, pad_after); + let after_range = build_slice_ranges(output_dims, dim, pad_before + dim_size, pad_after); + output = output.slice_assign(after_range, after_pad); + } + + output +} diff --git a/crates/burn-tensor/src/tensor/api/take.rs b/crates/burn-tensor/src/tensor/api/take.rs new file mode 100644 index 0000000..de33a4a --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/take.rs @@ -0,0 +1,96 @@ +use crate::{AsIndex, Int, Tensor, check, check::TensorCheck, kind::Basic}; +use alloc::vec::Vec; + +impl Tensor +where + K: Basic, +{ + /// Takes elements from the tensor along the given dimension using indices of any dimensionality. + /// + /// This behaves like numpy's take function. When indices is multi-dimensional, + /// the output shape will be: input.shape\[:dim\] + indices.shape + input.shape\[dim+1:\] + /// + /// # Arguments + /// + /// * `dim` - The dimension along which to select elements. Supports negative indexing. + /// * `indices` - The indices of elements to select. Can be any dimensionality. + /// Must be valid indices in the range [0, dim_size). + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::{Tensor, Int}; + /// + /// fn example() { + /// let device = Default::default(); + /// + /// // Example with 1D indices + /// let tensor = Tensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device); + /// let indices = Tensor::<1, Int>::from_data([2, 0, 1], &device); + /// let result: Tensor<2> = tensor.clone().take::<1, 2>(-1, indices); // -1 refers to last dimension + /// println!("{result}"); + /// // [[3.0, 1.0, 2.0], [6.0, 4.0, 5.0]] + /// + /// // Example with 2D indices - output will have +1 dimension (2D -> 3D) + /// let indices_2d = Tensor::<2, Int>::from_data([[0, 2], [1, 0]], &device); + /// let result: Tensor<3> = tensor.take::<2, 3>(1, indices_2d); + /// println!("{result}"); + /// // [[[1.0, 3.0], [2.0, 1.0]], [[4.0, 6.0], [5.0, 4.0]]] + /// } + /// ``` + pub fn take( + self, + dim: impl AsIndex, + indices: Tensor, + ) -> Tensor { + let dim = dim.expect_dim_index(D); + check!(TensorCheck::take::(dim)); + + // Store the indices shape for reshaping later + let indices_shape = indices.shape(); + let indices_dims = indices_shape.clone(); + + // Flatten indices to 1D for processing + let indices_flat = indices.reshape([indices_shape.num_elements()]); + + // Perform the selection with the flattened indices + let selected = self.select(dim, indices_flat); + + // Build the output shape + // Output shape = input.shape[:dim] + indices.shape + input.shape[dim+1:] + let selected_shape = selected.shape(); + let mut new_shape = Vec::with_capacity(DO); + + // Add dimensions before the selected dimension + for i in 0..dim { + new_shape.push(selected_shape[i]); + } + + // Add all indices dimensions + for &idx_dim in indices_dims.iter() { + new_shape.push(idx_dim); + } + + // Add dimensions after the selected dimension + for i in (dim + 1)..D { + new_shape.push(selected_shape[i]); + } + + // Verify we have the correct number of dimensions + assert_eq!( + new_shape.len(), + DO, + "Internal error: shape calculation resulted in {} dims but expected {}", + new_shape.len(), + DO + ); + + // Convert to fixed-size array for reshape + let mut shape_array = [0; DO]; + for (i, &s) in new_shape.iter().enumerate() { + shape_array[i] = s; + } + + selected.reshape(shape_array) + } +} diff --git a/crates/burn-tensor/src/tensor/api/transaction.rs b/crates/burn-tensor/src/tensor/api/transaction.rs new file mode 100644 index 0000000..ddb31bb --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/transaction.rs @@ -0,0 +1,86 @@ +use super::Tensor; +use crate::{ExecutionError, TensorData}; +use alloc::vec::Vec; +use burn_backend::ops::TransactionPrimitive; +use burn_dispatch::Dispatch; + +/// A transaction can [read](Self::register) multiple tensors at once with a single operation improving +/// compute utilization with optimized laziness. +/// +/// # Example +/// +/// ```rust,ignore +/// let [output_data, loss_data, targets_data] = Transaction::default() +/// .register(output) +/// .register(loss) +/// .register(targets) +/// .execute() +/// .try_into() +/// .expect("Correct amount of tensor data"); +/// ``` +pub struct Transaction { + opaque: transaction_opaque::Opaque, +} + +burn_std::obfuscate!( + type: TransactionPrimitive, + module: transaction_opaque, + derives: [Send, Sync], +); + +impl Default for Transaction { + fn default() -> Self { + Self::from_op(TransactionPrimitive::::default()) + } +} + +impl Transaction { + /// Crate-internal constructor wrapping a dispatch-level transaction. + pub(crate) fn from_op(op: TransactionPrimitive) -> Self { + Self { + opaque: transaction_opaque::Opaque::new(op), + } + } + + /// Crate-internal mutable borrow of the underlying transaction primitive. + pub(crate) fn as_op_mut(&mut self) -> &mut TransactionPrimitive { + self.opaque.as_mut() + } + + /// Crate-internal owning extraction of the underlying transaction primitive. + pub(crate) fn into_op(self) -> TransactionPrimitive { + self.opaque.into_inner() + } + + /// Add a [tensor](Tensor) to the transaction to be read. + pub fn register( + mut self, + tensor: Tensor, + ) -> Self { + K::register_transaction(self.as_op_mut(), tensor.primitive); + self + } + + /// Executes the transaction synchronously and returns the [data](TensorData) in the same order + /// in which they were [registered](Self::register). + pub fn execute(self) -> Vec { + burn_std::future::block_on(self.execute_async()) + .expect("Error while reading data: use `try_execute` to handle error at runtime") + } + + /// Executes the transaction synchronously and returns the [data](TensorData) in the same + /// order in which they were [registered](Self::register). + /// + /// # Returns + /// + /// Any error that might have occurred since the last time the device was synchronized. + pub fn try_execute(self) -> Result, ExecutionError> { + burn_std::future::block_on(self.execute_async()) + } + + /// Executes the transaction asynchronously and returns the [data](TensorData) in the same order + /// in which they were [registered](Self::register). + pub async fn execute_async(self) -> Result, ExecutionError> { + self.into_op().execute_async().await + } +} diff --git a/crates/burn-tensor/src/tensor/api/trunc.rs b/crates/burn-tensor/src/tensor/api/trunc.rs new file mode 100644 index 0000000..139577f --- /dev/null +++ b/crates/burn-tensor/src/tensor/api/trunc.rs @@ -0,0 +1,43 @@ +use burn_backend::ops::FloatTensorOps; +use burn_dispatch::Dispatch; + +use crate::{Float, Tensor, ops::BridgeTensor}; + +impl Tensor { + /// Truncates the tensor element-wise, rounding toward zero. + /// + /// This function returns a new tensor with the same shape as the input tensor, + /// where each element is truncated toward zero. For positive values, this is + /// equivalent to floor, and for negative values, it's equivalent to ceil. + /// + /// # Special Cases (IEEE 754 compliant) + /// + /// - `trunc(±0)` returns ±0 (preserves sign of zero) + /// - `trunc(±∞)` returns ±∞ + /// - `trunc(NaN)` returns NaN + /// + /// # Returns + /// + /// A tensor with the same shape where each element has been truncated toward zero. + /// + /// # Example + /// + /// ```rust + /// use burn_tensor::Tensor; + /// + /// fn example() { + /// let device = Default::default(); + /// let tensor = Tensor::<1>::from_data([2.3, -1.7, 0.5, -0.5, 3.9], &device); + /// let truncated = tensor.trunc(); + /// + /// // Result: [2.0, -1.0, 0.0, -0.0, 3.0] + /// } + /// ``` + pub fn trunc(self) -> Self { + Self::new(trunc_impl(self.primitive)) + } +} + +fn trunc_impl(p: BridgeTensor) -> BridgeTensor { + BridgeTensor::float(Dispatch::float_trunc(p.into_float())) +} diff --git a/crates/burn-tensor/src/tensor/distributed.rs b/crates/burn-tensor/src/tensor/distributed.rs new file mode 100644 index 0000000..6c91476 --- /dev/null +++ b/crates/burn-tensor/src/tensor/distributed.rs @@ -0,0 +1,92 @@ +//! Distributed execution utilities. +//! +//! The core component of this module is [`DistributedContext`], which manages +//! the lifecycle of distributed synchronization clients. + +use alloc::vec::Vec; +use burn_backend::TensorMetadata; +use burn_backend::{DeviceOps, distributed::DistributedOps}; +use burn_dispatch::{Dispatch, DispatchTensor}; +pub use burn_std::distributed::*; + +use crate::{Device, Tensor, ops::BridgeTensor}; + +/// This structure acts as a resource handle for multi-device synchronization. +/// +/// Spawning this context automatically initializes the underlying distributed communication +/// servers, while dropping it guarantees a clean and safe teardown of all network resources. +#[derive(Debug)] +pub struct DistributedContext { + devices: Vec, +} + +impl DistributedContext { + /// Starts a distributed communication server for the provided devices. + /// + /// # Arguments + /// + /// * `devices` - The collection of compute devices participating in the distributed operations. + /// * `config` - Parameter aggregation settings, such as global reduction strategies (`Mean`, `Sum`, etc.). + pub fn init(devices: Vec, config: DistributedConfig) -> Self { + let dispatch_devices = devices + .iter() + .map(|d| d.as_dispatch().clone()) + .collect::>(); + Dispatch::start_communication_server(&dispatch_devices, config); + + Self { devices } + } +} + +impl Drop for DistributedContext { + fn drop(&mut self) { + if !self.devices.is_empty() { + Dispatch::close_communication_server(self.devices[0].as_dispatch()); + } + } +} + +/// A tensor handle used for a collective operation, that is not yet valid for use. +/// We must ensure collective operations are completed before accessing the underlying data. +#[derive(new, Clone)] +pub struct CollectiveTensor { + handle: DispatchTensor, +} + +impl CollectiveTensor { + /// Synchronizes the collective operation and returns a valid tensor handle. + pub fn resolve(self) -> Tensor { + Dispatch::sync_collective(&self.handle.device()); + Tensor::new(BridgeTensor::float(self.handle)) + } + + /// Returns the tensor handle without synchronizing. + /// + /// # Safety + /// + /// The caller must ensure that `sync_collective()` is called before + /// the returned handle is used in any computation. + pub unsafe fn assume_resolved(self) -> Tensor { + Tensor::new(BridgeTensor::float(self.handle)) + } +} + +/// Performs an all_reduce operation on the input tensor. +/// +/// # Arguments +/// - `input`: The input tensor. +/// - `op`: The aggregation operation. +/// - `device_ids`: The list of all devices with which to `all_reduce` +/// +/// # Returns +/// A [CollectiveTensor] containing the handle of the result. +pub fn all_reduce( + input: Tensor, + op: ReduceOperation, + device_ids: Vec, +) -> CollectiveTensor { + let device_ids = device_ids.iter().map(|d| d.as_dispatch().id()).collect(); + let collective = Dispatch::all_reduce(input.primitive.into_float(), op, device_ids); + // Safety: we call `assume_resolved` only to wrap it in `burn_tensor`'s [CollectiveTensor]. + CollectiveTensor::new(unsafe { collective.assume_resolved() }) +} diff --git a/crates/burn-tensor/src/tensor/grid/affine_grid.rs b/crates/burn-tensor/src/tensor/grid/affine_grid.rs new file mode 100644 index 0000000..382da64 --- /dev/null +++ b/crates/burn-tensor/src/tensor/grid/affine_grid.rs @@ -0,0 +1,59 @@ +use crate::ElementConversion; +use crate::s; +use crate::tensor::{Int, Tensor}; +use alloc::vec; + +/// Generate a tensor with homogeonous coordinates of each element's +/// transformed location +/// +/// +/// See: +/// - [torch.nn.functional.affine_grid](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.affine_grid.html) +/// +/// * `transform` - Transformation with shape (batch_size, 2, 3) +/// * `dims` - dimensions as (batch_size, channels, height, width) +/// +/// # Returns +/// +/// Tensor with shape (batch_size, height, width, 2), where dim 2 is (x, y) +/// All coordinates are broadcast on the batch dim +pub fn affine_grid_2d(transform: Tensor<3>, dims: [usize; 4]) -> Tensor<4> { + let [batch_size, _c, height, width] = dims; + + let device = &transform.device(); + + let x = Tensor::<1, Int>::arange(0..width as i64, device) + .reshape([1, width]) + .expand([height, width]); + let y = Tensor::<1, Int>::arange(0..height as i64, device) + .reshape([height, 1]) + .expand([height, width]); + + // from ints (0..(width-1)) and (0..(height-1)), to (-1.0..1.0) + let x = x + .float() + .div_scalar(((width - 1) as f32 / 2.0).elem::()) + .sub_scalar((1_f32).elem::()); + let y = y + .float() + .div_scalar(((height - 1) as f32 / 2.0).elem::()) + .sub_scalar((1_f32).elem::()); + + // Broadcast to batch dimension + let x = x.unsqueeze_dim::<3>(0).expand([batch_size, height, width]); // [B, H, W] + let y = y.unsqueeze_dim::<3>(0).expand([batch_size, height, width]); // [B, H, W] + + // Apply affine transform + let a_11 = transform.clone().slice(s![.., 0, 0]); + let a_12 = transform.clone().slice(s![.., 0, 1]); + let trans_x = transform.clone().slice(s![.., 0, 2]); + + let a_21 = transform.clone().slice(s![.., 1, 0]); + let a_22 = transform.clone().slice(s![.., 1, 1]); + let trans_y = transform.slice(s![.., 1, 2]); + + let grid_x = a_11.mul(x.clone()).add(a_12.mul(y.clone())).add(trans_x); + let grid_y = a_21.mul(x).add(a_22.mul(y)).add(trans_y); + + Tensor::stack(vec![grid_x, grid_y], 3) +} diff --git a/crates/burn-tensor/src/tensor/grid/meshgrid.rs b/crates/burn-tensor/src/tensor/grid/meshgrid.rs new file mode 100644 index 0000000..affbb8e --- /dev/null +++ b/crates/burn-tensor/src/tensor/grid/meshgrid.rs @@ -0,0 +1,104 @@ +use crate::kind::Basic; +use crate::tensor::Tensor; +use crate::tensor::grid::{GridIndexing, GridOptions, GridSparsity, IndexPos}; +use alloc::vec::Vec; + +/// Return a collection of coordinate matrices for coordinate vectors. +/// +/// Takes N 1D tensors and returns N tensors where each tensor represents the coordinates +/// in one dimension across an N-dimensional grid. +/// +/// Based upon `options.sparse`, the generated coordinate tensors can either be `Sparse` or `Dense`: +/// * In `Sparse` mode, output tensors will have shape 1 everywhere except their cardinal dimension. +/// * In `Dense` mode, output tensors will be expanded to the full grid shape. +/// +/// Based upon `options.indexing`, the generated coordinate tensors will use either: +/// * `Matrix` indexing, where dimensions are in the same order as their cardinality. +/// * `Cartesian` indexing; where the first two dimensions are swapped. +/// +/// See: +/// - [numpy.meshgrid](https://numpy.org/doc/stable/reference/generated/numpy.meshgrid.html) +/// - [torch.meshgrid](https://pytorch.org/docs/stable/generated/torch.meshgrid.html) +/// +/// # Arguments +/// +/// * `tensors` - A slice of 1D tensors +/// * `options` - the options. +/// +/// # Returns +/// +/// A vector of N N-dimensional tensors representing the grid coordinates. +pub fn meshgrid(tensors: &[Tensor<1, K>; N], options: O) -> [Tensor; N] +where + K: Basic, + O: Into, +{ + let options = options.into(); + let swap_dims = options.indexing == GridIndexing::Cartesian && N > 1; + let dense = options.sparsity == GridSparsity::Dense; + + let grid_shape: [usize; N] = tensors + .iter() + .map(|t| t.dims()[0]) + .collect::>() + .try_into() + .unwrap(); + + tensors + .iter() + .enumerate() + .map(|(i, tensor)| { + let mut coord_tensor_shape = [1; N]; + coord_tensor_shape[i] = grid_shape[i]; + + // Reshape the tensor to have singleton dimensions in all but the i-th dimension + let mut tensor = tensor.clone().reshape(coord_tensor_shape); + + if dense { + tensor = tensor.expand(grid_shape); + } + if swap_dims { + tensor = tensor.swap_dims(0, 1); + } + + tensor + }) + .collect::>() + .try_into() + .unwrap() +} + +/// Return a coordinate matrix for a given set of 1D coordinate tensors. +/// +/// Equivalent to stacking a dense matrix `meshgrid`, +/// where the stack is along the first or last dimension. +/// +/// # Arguments +/// +/// * `tensors`: A slice of 1D tensors. +/// * `index_pos`: The position of the index in the output tensor. +/// +/// # Returns +/// +/// A tensor of either ``(N, ..., |T[i]|, ...)`` or ``(..., |T[i]|, ..., N)``, +/// of coordinates, indexed on the first or last dimension. +pub fn meshgrid_stack( + tensors: &[Tensor<1, K>; D], + index_pos: IndexPos, +) -> Tensor +where + K: Basic, +{ + assert_eq!(D2, D + 1, "D2 ({D2}) != D ({D}) + 1"); + + let xs: Vec> = meshgrid(tensors, GridOptions::default()) + .into_iter() + .collect(); + + let dim = match index_pos { + IndexPos::First => 0, + IndexPos::Last => D, + }; + + Tensor::stack(xs, dim) +} diff --git a/crates/burn-tensor/src/tensor/grid/mod.rs b/crates/burn-tensor/src/tensor/grid/mod.rs new file mode 100644 index 0000000..3750962 --- /dev/null +++ b/crates/burn-tensor/src/tensor/grid/mod.rs @@ -0,0 +1,68 @@ +mod affine_grid; +mod meshgrid; + +pub use meshgrid::*; + +pub use affine_grid::*; + +/// Enum to specify index cardinal layout. +#[derive(Default, Debug, Clone, Copy, PartialEq, Eq)] +pub enum GridIndexing { + /// Dimensions are in the same order as the cardinality of the inputs. + /// Equivalent to "ij" indexing in NumPy and PyTorch. + #[default] + Matrix, + + /// The same as Matrix, but the first two dimensions are swapped. + /// Equivalent to "xy" indexing in NumPy and PyTorch. + Cartesian, +} + +/// Enum to specify grid sparsity mode. +#[derive(Default, Debug, Clone, Copy, PartialEq, Eq)] +pub enum GridSparsity { + /// The grid is fully expanded to the full cartesian product shape. + #[default] + Dense, + + /// The grid is sparse, expanded only at the cardinal dimensions. + Sparse, +} + +/// Grid policy options. +#[derive(new, Default, Debug, Copy, Clone)] +pub struct GridOptions { + /// Indexing mode. + pub indexing: GridIndexing, + + /// Sparsity mode. + pub sparsity: GridSparsity, +} + +impl From for GridOptions { + fn from(value: GridIndexing) -> Self { + Self { + indexing: value, + ..Default::default() + } + } +} +impl From for GridOptions { + fn from(value: GridSparsity) -> Self { + Self { + sparsity: value, + ..Default::default() + } + } +} + +/// Enum to specify the index dimension position. +#[derive(Default, Debug, Copy, Clone)] +pub enum IndexPos { + /// The index is in the first dimension. + #[default] + First, + + /// The index is in the last dimension. + Last, +} diff --git a/crates/burn-tensor/src/tensor/kind.rs b/crates/burn-tensor/src/tensor/kind.rs new file mode 100644 index 0000000..fa50f56 --- /dev/null +++ b/crates/burn-tensor/src/tensor/kind.rs @@ -0,0 +1,33 @@ +// Sealed traits, limited to Bool, Float and Int types which implement the pub(crate) traits. +#![allow(private_bounds)] + +#[cfg(feature = "extension")] +pub use crate::bridge::BridgeTensor; + +pub use crate::bridge::{Bool, Float, Int, Kind, TensorKind}; + +/// The base trait for any tensor kind. +pub trait Basic: crate::ops::BasicOps {} +impl Basic for K {} + +/// Kinds that support numeric operations. +pub trait Numeric: Basic + crate::ops::Numeric {} +impl Numeric for K {} + +/// Kinds that support ordered operations. +pub trait Ordered: Numeric + crate::ops::Ordered {} +impl Ordered for K {} + +/// Kinds that support float math operations. +pub trait FloatMath: Numeric + crate::ops::FloatMathOps {} +impl FloatMath for K {} + +/// Kinds that support transaction operations. +pub trait Transaction: Basic + crate::ops::TransactionOp {} +impl Transaction for K {} + +/// Kinds that support autodiff operations. +// #[cfg(feature = "autodiff")] +pub trait Autodiff: Basic + crate::ops::BasicAutodiffOps {} +// #[cfg(feature = "autodiff")] +impl Autodiff for K {} diff --git a/crates/burn-tensor/src/tensor/linalg/cosine_similarity.rs b/crates/burn-tensor/src/tensor/linalg/cosine_similarity.rs new file mode 100644 index 0000000..4e09feb --- /dev/null +++ b/crates/burn-tensor/src/tensor/linalg/cosine_similarity.rs @@ -0,0 +1,51 @@ +use burn_std::FloatDType; + +use crate::tensor::Tensor; + +use super::l2_norm; + +/// Computes the cosine similarity between two tensors along a specified dimension. +/// +/// Calculates the cosine of the angle between inputs as their dot product divided +/// by the product of their L2 norms. +/// +/// # Arguments +/// +/// * `x1` - First input tensor +/// * `x2` - Second input tensor +/// * `dim` - Dimension along which to compute the similarity +/// (negative indices allowed: -1 for last dimension) +/// * `eps` - Small value to avoid division by zero (default: dtype's smallest positive normal) +/// +/// # Returns +/// +/// Tensor containing the cosine similarity between x1 and x2 +pub fn cosine_similarity( + x1: Tensor, + x2: Tensor, + dim: i32, + eps: Option, +) -> Tensor { + let eps = eps.unwrap_or_else(|| { + x1.dtype() + .finfo() + .unwrap_or(FloatDType::F32.finfo()) + .min_positive + }); + + // Convert negative dimension to positive + let dim_idx = if dim < 0 { D as i32 + dim } else { dim } as usize; + + // Compute dot product: sum(x1 * x2) along the specified dimension + let dot_product = (x1.clone() * x2.clone()).sum_dim(dim_idx); + + // Compute L2 norms: ||x1|| and ||x2|| + let norm_x1 = l2_norm(x1, dim_idx); + let norm_x2 = l2_norm(x2, dim_idx); + + // Calculate the denominator (product of the norms) with epsilon to avoid division by zero + let denominator = norm_x1.clamp_min(eps) * norm_x2.clamp_min(eps); + + // Return the cosine similarity (dot product divided by the product of norms) + dot_product / denominator +} diff --git a/crates/burn-tensor/src/tensor/linalg/det.rs b/crates/burn-tensor/src/tensor/linalg/det.rs new file mode 100644 index 0000000..d49ca31 --- /dev/null +++ b/crates/burn-tensor/src/tensor/linalg/det.rs @@ -0,0 +1,179 @@ +use crate::check::TensorCheck; +use crate::{Tensor, check, linalg}; +use burn_std::{DType, FloatDType}; +#[allow(unused_imports)] +use num_traits::float::Float; + +/// Computes the determinant on the last two dimensions of the input tensor. +/// +/// # Arguments +/// - `tensor` - The input tensor of shape `[..., N, N]`. +/// +/// # Returns +/// - The determinant tensor of shape `[...]` where its rank is less than the +/// input tensor's rank by two. +/// +/// # Generic Parameters +/// - `D`: The rank of the input tensor. +/// - `D1`: Must be set to `D - 1`. +/// - `D2`: Must be set to `D - 2`. +/// +/// # Panics +/// This function will panic if: +/// - The generic parameters do not satisfy `D - 1 == D1`. +/// - The generic parameters do not satisfy `D - 2 == D2`. +/// - The input tensor rank `D` is less than 3. +/// - The last two dimensions of the input tensor are not equal. +/// - The input is a quantized tensor with dtype `DType::QFloat`. +/// +/// # Performance Note +/// The determinant for 1 by 1, 2 by 2, and 3 by 3 matrices are computed using closed-form +/// expressions. For larger matrices (4 by 4 or larger), the determinant function relies on +/// the LU decomposition function under the hood,which is not fully optimized. It will not be +/// as fast as highly tuned specialized libraries, especially for very large matrices or large +/// batch sizes. +/// +/// # Numerical Behavior +/// - If the input tensors have types F16 or BF16, then they are internally upcast to +/// F32 to perform the computations and cast back to the original data type (F16 or BF16) +/// right before the function returns. +/// - In this case, if the determinant values fall outside of the original data type's +/// range, then the cast-back will underflow to zero. +/// +/// # Example +/// ```rust,ignore +/// use burn::tensor::Tensor; +/// use burn::tensor::linalg; +/// +/// fn example() { +/// let device = Default::default(); +/// let tensor = Tensor::<3>::from_data([[[4.0, 3.0], [6.0, 3.0]]], &device); +/// +/// // Compute determinant +/// let result = linalg::det::(tensor); +/// +/// // Expected Output: +/// // result: [-6.0] +/// } +/// +/// fn example2() { +/// let device = Default::default(); +/// let tensor = Tensor::<3>::from_data( +/// [ +/// [[1.0, 2.0], [3.0, 4.0]], // det = -2 +/// [[2.0, 0.0], [0.0, 3.0]], // det = 6 +/// [[5.0, 6.0], [7.0, 8.0]], // det = -2 +/// ], +/// &device, +/// ); +/// +/// // Compute determinant +/// let result = linalg::det::(tensor); +/// +/// // Expected Output: +/// // result: [-2.0, 6.0, -2.0] +/// } +/// ``` +pub fn det(mut tensor: Tensor) -> Tensor { + // Check whether input tensor has valid shape to compute determinant + let dims = tensor.dims(); + let original_dtype = tensor.dtype(); + check!(TensorCheck::det::(dims, original_dtype)); + + // Upcast f16 and bf16 to f32 + let needs_upcast = original_dtype == DType::F16 || original_dtype == DType::BF16; + let working_float_dtype: FloatDType; + if needs_upcast { + working_float_dtype = FloatDType::F32; + tensor = tensor.cast(working_float_dtype); + } else { + working_float_dtype = original_dtype.into() + }; + + // Compute determinant for base cases (1x1, 2x2, and 3x3 matrices) + let rank = D as isize; + if dims[D - 1] == 1 { + let det_tensor = tensor.squeeze_dims::(&[rank - 2, rank - 1]); + if needs_upcast { + return det_tensor.cast(original_dtype); + } + return det_tensor; + } else if dims[D - 1] == 2 { + let a = tensor.clone().slice_dim(D - 2, 0).slice_dim(D - 1, 0); + let b = tensor.clone().slice_dim(D - 2, 0).slice_dim(D - 1, 1); + let c = tensor.clone().slice_dim(D - 2, 1).slice_dim(D - 1, 0); + let d = tensor.clone().slice_dim(D - 2, 1).slice_dim(D - 1, 1); + let det_tensor = (a * d - b * c).squeeze_dims::(&[rank - 2, rank - 1]); + if needs_upcast { + return det_tensor.cast(original_dtype); + } + return det_tensor; + } else if dims[D - 1] == 3 { + let a = tensor.clone().slice_dim(D - 2, 0).slice_dim(D - 1, 0); + let b = tensor.clone().slice_dim(D - 2, 0).slice_dim(D - 1, 1); + let c = tensor.clone().slice_dim(D - 2, 0).slice_dim(D - 1, 2); + let d = tensor.clone().slice_dim(D - 2, 1).slice_dim(D - 1, 0); + let e = tensor.clone().slice_dim(D - 2, 1).slice_dim(D - 1, 1); + let f = tensor.clone().slice_dim(D - 2, 1).slice_dim(D - 1, 2); + let g = tensor.clone().slice_dim(D - 2, 2).slice_dim(D - 1, 0); + let h = tensor.clone().slice_dim(D - 2, 2).slice_dim(D - 1, 1); + let i = tensor.clone().slice_dim(D - 2, 2).slice_dim(D - 1, 2); + let det_tensor = (a * (e.clone() * i.clone() - f.clone() * h.clone()) + - b * (d.clone() * i - f * g.clone()) + + c * (d * h - e * g)) + .squeeze_dims::(&[rank - 2, rank - 1]); + if needs_upcast { + return det_tensor.cast(original_dtype); + } + return det_tensor; + } + + // Compute determinant for general case + // det(A) = det(P) * det(L) * det(U) + // det(A) = det(P) * 1 * det(U) + let (lu, pivots) = linalg::compute_lu_decomposition::(tensor.clone()); + + // Compute the determinant of P + let squeezed_pivots = pivots.squeeze_dim::(D - 1); + let n_pivots = squeezed_pivots.dims()[D1 - 1] as i64; + let range_1d: Tensor<1> = + Tensor::arange(0..n_pivots, &tensor.device()).cast(working_float_dtype); + let mut reshape_dims = [1; D1]; + reshape_dims[D1 - 1] = n_pivots; + let range = range_1d.reshape(reshape_dims); + let expand_dims: [usize; D1] = squeezed_pivots.dims(); + let batched_range_tensor = range.expand(expand_dims); + let n_row_swaps = squeezed_pivots + .not_equal(batched_range_tensor) + .int() + .sum_dim(D1 - 1); + let odd_mask = n_row_swaps.clone().remainder_scalar(2).equal_elem(1); + let p_det = n_row_swaps + .cast(working_float_dtype) + .ones_like() + .mask_fill(odd_mask, -1.0) + .squeeze_dim(D1 - 1); + + // Compute the determinant of U + let u_diag = linalg::diag::(lu); + let mut u_det = u_diag.clone().prod_dim(D1 - 1).squeeze_dim(D1 - 1); + let eps = tensor + .dtype() + .finfo() + .expect("The input tensor to linalg::det should have float dtype.") + .epsilon; + let n = dims[D - 1]; // The input tensor contains n by n matrices + let threshold = u_diag.clone().abs().max_dim(D1 - 1) * (n as f64).sqrt() * eps; + let near_zero = u_diag.abs().lower_equal(threshold); + let singular_mask = near_zero.any_dim(D1 - 1).squeeze_dim::(D1 - 1); + u_det = u_det.mask_fill(singular_mask, 0.0); + + let final_det = p_det * u_det; + + // Cast back to original dtypes + if needs_upcast { + final_det.cast(original_dtype) + } else { + final_det + } +} diff --git a/crates/burn-tensor/src/tensor/linalg/diag.rs b/crates/burn-tensor/src/tensor/linalg/diag.rs new file mode 100644 index 0000000..79adf72 --- /dev/null +++ b/crates/burn-tensor/src/tensor/linalg/diag.rs @@ -0,0 +1,41 @@ +use crate::check; +use crate::check::TensorCheck; +use crate::kind::Basic; +use crate::tensor::{Int, Tensor}; + +/// Returns the diag of a matrix. +/// +/// For batched inputs, returns of each matrix in the batch independently. +/// +/// The diag operation extracts the diagonal elements of the last two dimensions, +/// treating them as the matrix dimensions, while preserving all leading batch dimensions. +/// +/// # Arguments +/// +/// * `tensor` - The input tensor with at least 2 dimensions. +/// +/// # Returns +/// A tensor of rank `D - 1`, where the last dimension contains the diagonal elements of the input. +pub fn diag(tensor: Tensor) -> Tensor +where + K: Basic, +{ + check!(TensorCheck::diag::()); + + let shape = tensor.shape(); + let rows = shape[D - 2]; + let cols = shape[D - 1]; + let diag_len = rows.min(cols); + let device = tensor.device(); + + // create the indices for the diag + let mut flat_shape = shape.clone(); + flat_shape[D - 2] = rows * cols; + flat_shape[D - 1] = 1; + let flat: Tensor = tensor.reshape(flat_shape); + + let range = Tensor::<1, Int>::arange(0..diag_len as i64, &device); + let step_tensor = Tensor::<1, Int>::from_data([cols as i64 + 1], &device); + let indices = range * step_tensor; + flat.take::<1, D>(D - 2, indices).squeeze_dim(D - 1) +} diff --git a/crates/burn-tensor/src/tensor/linalg/lu.rs b/crates/burn-tensor/src/tensor/linalg/lu.rs new file mode 100644 index 0000000..bb20953 --- /dev/null +++ b/crates/burn-tensor/src/tensor/linalg/lu.rs @@ -0,0 +1,521 @@ +use crate::{Bool, Device, ElementConversion, Float, Int, Tensor, check, check::TensorCheck}; +use alloc::vec; +use alloc::vec::Vec; +use burn_std::{DType, FloatDType, IndexingUpdateOp, Slice}; + +/// Computes the LU decomposition of a square or rectangular matrix with partial pivoting. +/// +/// This function decomposes the input tensor A into three tensors P, L, and U +/// such that A = PLU. +/// +/// # Arguments +/// - `tensor` - The input tensor of shape `[..., n_rows, n_cols]`. +/// +/// # Returns +/// A tuple of three tensors `(P, L, U)`: +/// - `P` - The permutation tensor of shape `[..., n_rows, n_rows]`. +/// - `L` - The lower triangular tensor of shape `[..., n_rows, min(n_rows, n_cols)]` +/// with unit diagonal elements. +/// - `U` - The upper triangular tensor of shape `[..., min(n_rows, n_cols), n_cols]`. +/// +/// # Generic Parameters +/// +/// - `D`: The number of dimensions of the input tensor. +/// - `D1`: The number of dimensions of the 1D pivot tensor. Must be exactly `D - 1`. +/// +/// # Panics +/// This function will panic if the tensor checks fail: +/// - The input tensor has less than 2 dimensions (`D < 2`). +/// - The generic parameters do not satisfy `D - 1 == D1`. +/// - The input is a quantized tensor with dtype `DType::QFloat`. +/// +/// # Performance Note +/// The current implementation of LU decomposition is not fully optimized. It will not +/// be as fast as highly tuned specialized libraries, especially for very large +/// matrices or large batch sizes. +/// +/// # Numerical Behavior +/// - If the input tensor has dtype F16 or BF16, it is internally upcast to F32 +/// for the computation and cast back to the original dtype before returning. +/// - In this case, values in L and U that fall outside the original dtype's +/// representable range will saturate or underflow on cast-back. +/// +/// # Example +/// ```rust,ignore +/// use burn::tensor::Tensor; +/// use burn::backend::Flex; +/// use burn::tensor::linalg; +/// +/// fn example() { +/// let device = Default::default(); +/// let tensor = Tensor::<2>::from_data([[4.0, 3.0], [6.0, 3.0]], &device); +/// +/// // Compute P, L, U +/// let (p, l, u) = linalg::lu::< 2, 1>(tensor); +/// +/// // Expected Output: +/// // p: [[0.0, 1.0], +/// // [1.0, 0.0]] +/// // +/// // l: [[1.0, 0.0], +/// // [0.6666667, 1.0]] +/// // +/// // u: [[6.0, 3.0], +/// // [0.0, 1.0]] +/// } +/// ``` +pub fn lu( + mut tensor: Tensor, +) -> (Tensor, Tensor, Tensor) { + let dims = tensor.dims(); + let original_dtype = tensor.dtype(); + check!(TensorCheck::lu_generic_param::("linalg::lu")); + check!(TensorCheck::lu_input_tensor::( + "linalg::lu", + &dims, + original_dtype + )); + + let device = tensor.device(); + let n_rows = dims[D - 2]; + let n_cols = dims[D - 1]; + + // Upcast f16 and bf16 to f32 + let needs_upcast = original_dtype == DType::F16 || original_dtype == DType::BF16; + if needs_upcast { + tensor = tensor.cast(FloatDType::F32) + } + + let (lu_tensor, p_compact) = compute_lu_decomposition::(tensor); + + let u; + let temp_l; + if n_rows < n_cols { + temp_l = lu_tensor.clone().slice_dim(D - 1, 0..n_rows).tril(0); + u = lu_tensor.triu(0); + } else { + temp_l = lu_tensor.clone().tril(0); + u = lu_tensor.slice_dim(D - 2, 0..n_cols).triu(0); + } + let mask = Tensor::::diag_mask(temp_l.shape(), 0, &device).bool_not(); + let l = temp_l.mask_fill(mask, 1.0); + let p = construct_full_permutation_tensor(p_compact, n_rows, &device).transpose(); + + if needs_upcast { + ( + p.cast(original_dtype), + l.cast(original_dtype), + u.cast(original_dtype), + ) + } else { + (p, l, u) + } +} + +/// Dispatches the LU decomposition to either the block or standard algorithm based on +/// the size of the matrix. +pub(super) fn compute_lu_decomposition( + tensor: Tensor, +) -> (Tensor, Tensor) { + let device = tensor.device(); + let dims = tensor.dims(); + let n_rows = dims[D - 2]; + let n_cols = dims[D - 1]; + let size = n_rows.min(n_cols); + if size < 256 { + return standard_lu_with_partial_piv::(tensor, &device); + } + + block_lu_with_partial_piv::(tensor) +} + +/// Performs block LU decomposition with partial pivoting. +/// +/// This algorithm divides the matrix into blocks to maximize matrix-matrix multiplications (GEMM), +/// which are highly optimized on modern hardware, compared to vector-vector operations. +fn block_lu_with_partial_piv( + mut tensor: Tensor, +) -> (Tensor, Tensor) { + let device = tensor.device(); + let dims = tensor.dims(); + let n_rows = dims[D - 2]; + let n_cols = dims[D - 1]; + let piv_nums = n_rows.min(n_cols); + let dtype = tensor.dtype().into(); + let mut global_piv = create_permutation_tensor::(piv_nums, dims, dtype, &device); + let block_size = 128; + + // Computes the total number of blocks including incomplete blocks + // E.g., piv_nums = 100 & block_size = 32 -> n_blocks = 4 (not 3) where + // the 4th block is smaller than other blocks + let n_blocks = piv_nums.div_ceil(block_size); + let mut slices = vec![Slice::full(); D]; // For updating the original tensor in-place + + // k is the current block number + for block_k in 0..n_blocks { + // Determine the index range for the current block column + let k_start = block_k * block_size; + let k_end = (k_start + block_size).min(piv_nums); + let current_block_size = k_end - k_start; + + // Apply standard LU decomposition with partial pivoting to the current block column + let sub_tensor = tensor + .clone() + .slice_dim(D - 2, k_start..) + .slice_dim(D - 1, k_start..k_end); + let (block_column, local_piv) = standard_lu_with_partial_piv::(sub_tensor, &device); + slices[D - 2] = Slice::from(k_start..); + slices[D - 1] = Slice::from(k_start..k_end); + tensor = tensor.slice_assign(&slices, block_column); + + // Update `permutations` to global indices + global_piv = + update_permutations_to_global_idx(global_piv.clone(), local_piv.clone(), k_start); + + // Apply `local_piv` permutations to the left sub-tensor + if block_k != 0 { + let left_sub_tensor = tensor + .clone() + .slice_dim(D - 2, k_start..) + .slice_dim(D - 1, ..k_start); + let permutated_left_sub_tensor = + apply_permutations_to_tensor(left_sub_tensor, local_piv.clone(), &device); + slices[D - 2] = Slice::from(k_start..); + slices[D - 1] = Slice::from(..k_start); + tensor = tensor.slice_assign(&slices, permutated_left_sub_tensor); + } + + // Only update the right side if there are columns left + if k_end < n_cols { + // Apply `local_piv` permutations to the remaining right sub-tensor + let right_sub_tensor = tensor + .clone() + .slice_dim(D - 2, k_start..) + .slice_dim(D - 1, k_end..); + let permutated_right_sub_tensor = + apply_permutations_to_tensor(right_sub_tensor, local_piv, &device); + slices[D - 2] = Slice::from(k_start..); + slices[D - 1] = Slice::from(k_end..); + tensor = tensor.slice_assign(&slices, permutated_right_sub_tensor); + + // Update the cols to the right of the current diagonal block. + // Triangular solve for U blocks. + let diagonal_l_block = tensor + .clone() + .slice_dim(D - 2, k_start..k_end) + .slice_dim(D - 1, k_start..k_end); + let row_blocks = tensor + .clone() + .slice_dim(D - 2, k_start..k_end) + .slice_dim(D - 1, k_end..); + let updated_row_blocks = + solve_for_u_blocks(diagonal_l_block, row_blocks, current_block_size); + slices[D - 2] = Slice::from(k_start..k_end); + slices[D - 1] = Slice::from(k_end..); + tensor = tensor.slice_assign(&slices, updated_row_blocks.clone()); + + // Only update trailing A blocks if there are rows left below + if k_end < n_rows { + // Update the trailing A blocks + let trailing_a_blocks = tensor + .clone() + .slice_dim(D - 2, k_end..) + .slice_dim(D - 1, k_end..); + let l_col_blocks = tensor + .clone() + .slice_dim(D - 2, k_end..) + .slice_dim(D - 1, k_start..k_end); + let outer_prod = l_col_blocks.matmul(updated_row_blocks); + let new_trailing_a_blocks = trailing_a_blocks - outer_prod; + + // Overwrite part of the tensor with new_trailing_a_blocks + slices[D - 2] = Slice::from(k_end..); + slices[D - 1] = Slice::from(k_end..); + tensor = tensor.slice_assign(&slices, new_trailing_a_blocks); + } + } + } + + (tensor, global_piv) +} + +/// Performs standard LU decomposition (outer product LU) with partial pivoting. +/// +/// This is an iterative, unblocked algorithm that processes the matrix column by column. +fn standard_lu_with_partial_piv( + mut tensor: Tensor, + device: &Device, +) -> (Tensor, Tensor) { + let dims = tensor.dims(); + let n_rows = dims[D - 2]; + let n_cols = dims[D - 1]; + let piv_nums = n_rows.min(n_cols); + let dtype = tensor.dtype().into(); + let mut permutations = create_permutation_tensor::(piv_nums, dims, dtype, device); + + for k in 0..piv_nums { + // Find the index of the maximum absolute value in the k-th column (from row k downwards) + // Shape: [B1, ..., BN, 1, 1] + let max_row_indices = tensor + .clone() + .slice_dim(D - 2, k..) + .slice_dim(D - 1, k) + .abs() + .argmax(D - 2) + + (k as i64); + + // Swap current row (k-th row) with the row with maximum absolute value + tensor = swap_tensor_rows(tensor, max_row_indices.clone(), k); + // Store the max row index in the k-th entry of the permutations vector/tensor + permutations = update_permutations(permutations, max_row_indices, k, dtype); + + // If there are rows left under the k-th pivot + if k < n_rows - 1 { + // Update k-th column under the diagonal + tensor = update_kth_column(tensor, k); + + // If there still exists columns to right of the k-th pivot + if k < piv_nums - 1 { + tensor = update_trailing_submatrix::(tensor, k); + } + } + } + + (tensor, permutations) +} + +/// Constructs a full square permutation matrix \( P \) from a compact pivot tensor. +fn construct_full_permutation_tensor( + piv: Tensor, + n_rows: usize, + device: &Device, +) -> Tensor { + let dims = piv.dims(); + let identity_2d_uncasted: Tensor<2, Float> = Tensor::eye(n_rows, device); + let identity_2d = identity_2d_uncasted.cast(piv.dtype()); + + // Reshape the `identity` tensor from 2 dims to D dims + let mut reshape_dims = [1; D]; + reshape_dims[D - 2] = n_rows; + reshape_dims[D - 1] = n_rows; + let reshaped_identity = identity_2d.reshape(reshape_dims); + + // Expand the batch dimensions to match the original input tensor's shape + let mut expand_dims = [n_rows; D]; + expand_dims[..(D - 2)].copy_from_slice(&dims[..(D - 2)]); + let identity = reshaped_identity.expand(expand_dims); + + // Iterate through `piv` and apply rows swap to the `identity` tensor + // to construct the full permutation tensor + + apply_permutations_to_tensor(identity, piv, device) +} + +/// Initializes a permutation tensor representing the identity permutation `[0, 1, 2, ..., piv_nums - 1]`. +fn create_permutation_tensor( + piv_nums: usize, + dims: [usize; D], + dtype: FloatDType, + device: &Device, +) -> Tensor { + let piv = Tensor::arange(0..piv_nums as i64, device).cast(dtype); + + // Reshape the piv tensor from 1 dim to D dims + let mut reshape_dims = [1; D]; + reshape_dims[D - 2] = piv_nums; + let reshaped = piv.reshape(reshape_dims); + + // Expand the batch dimensions to match the original input tensor's shape + let mut expand_dims = [piv_nums; D]; + expand_dims[..(D - 2)].copy_from_slice(&dims[..(D - 2)]); + expand_dims[D - 1] = 1; + + reshaped.expand(expand_dims) +} + +/// Swaps the `k`-th row with the rows specified in `swap_target_row_tensor`. +fn swap_tensor_rows( + tensor: Tensor, + mut swap_target_row_tensor: Tensor, + k: usize, +) -> Tensor { + let mut expand_dims = tensor.dims(); + expand_dims[D - 2] = 1; + swap_target_row_tensor = swap_target_row_tensor.expand(expand_dims); + + let val_k = tensor.clone().slice_dim(D - 2, k); + let val_r = tensor.clone().gather(D - 2, swap_target_row_tensor.clone()); + + let mut slices = vec![Slice::full(); D]; + slices[D - 2] = Slice::from(k); + let tensor = tensor.slice_assign(&slices, val_r.clone()); + + let val_k_minus_r = val_k - val_r; + tensor.scatter( + D - 2, + swap_target_row_tensor, + val_k_minus_r, + IndexingUpdateOp::Add, + ) +} + +/// Updates the permutation tensor by recording the swap at step `k`. +fn update_permutations( + mut permutations: Tensor, + max_row_index_tensor: Tensor, + k: usize, + dtype: FloatDType, +) -> Tensor { + // Store the max row index in the k-th index of the permutations vector/tensor + let mut slices = vec![Slice::full(); D]; + slices[D - 2] = Slice::from(k); + let float_max_row_indices = max_row_index_tensor.cast(dtype); + permutations = permutations.slice_assign(&slices, float_max_row_indices); + + permutations +} + +/// Scales the `k`-th column below the diagonal by the pivot element A_{kk}. +fn update_kth_column(tensor: Tensor, k: usize) -> Tensor { + let a_kk = tensor.clone().slice_dim(D - 2, k).slice_dim(D - 1, k); + let a_rho_k = tensor.clone().slice_dim(D - 2, k + 1..).slice_dim(D - 1, k); + + // A singular matrix will have a pivot of exactly 0. + // Due to partial pivoting, if the pivot is 0, all elements below it are also 0. + // We replace 0 with 1.0 to avoid NaN when dividing 0.0 / 0.0. + let is_zero_mask = a_kk.clone().equal_elem(0.0); + let safe_a_kk = a_kk.mask_fill(is_zero_mask, 1.0); + let updated_column = a_rho_k / safe_a_kk; + + let mut slices = vec![Slice::full(); D]; + slices[D - 2] = Slice::from((k + 1)..); // Rows k+1 to the end + slices[D - 1] = Slice::from(k..(k + 1)); // Column k + + tensor.slice_assign(&slices, updated_column) +} + +/// Updates the trailing submatrix: A_{k+1:, k+1:} -= A_{k+1:, k} * A_{k, k+1:}. +fn update_trailing_submatrix( + tensor: Tensor, + k: usize, +) -> Tensor { + let a_rho_k = tensor.clone().slice_dim(D - 2, k + 1..).slice_dim(D - 1, k); + let a_k_rho = tensor.clone().slice_dim(D - 2, k).slice_dim(D - 1, k + 1..); + let outer_product = a_rho_k.matmul(a_k_rho); + + let a_rho_rho = tensor + .clone() + .slice_dim(D - 2, k + 1..) + .slice_dim(D - 1, k + 1..); + let updated_a_rho_rho = a_rho_rho - outer_product; + + let mut slices = vec![Slice::full(); D]; + slices[D - 2] = Slice::from((k + 1)..); // Rows k+1 to the end + slices[D - 1] = Slice::from((k + 1)..); // Cols k+1 to the end + tensor.slice_assign(&slices, updated_a_rho_rho) +} + +/// Shifts local pivot indices from a block factorization to global indices. +fn update_permutations_to_global_idx( + global_piv: Tensor, + local_piv: Tensor, + k_start: usize, +) -> Tensor { + let n = local_piv.dims()[D - 2]; + let mut slices = vec![Slice::full(); D]; + slices[D - 2] = Slice::from(k_start..(n + k_start)); + + let global_val = local_piv.add_scalar(k_start as f32); + + global_piv.slice_assign(&slices, global_val) +} + +/// Applies the permutations to the entire width of the tensor. +fn apply_permutations_to_tensor( + tensor: Tensor, + piv: Tensor, + device: &Device, +) -> Tensor { + let tensor_dims = tensor.dims(); + let n_rows = tensor_dims[D - 2]; + let n_pivots = piv.dims()[D - 2]; + let piv_data: Vec = piv.into_data().convert::().into_vec::().unwrap(); + + // Compute total batch size (product of all batch dimensions) + let batch_size: usize = tensor_dims[..D - 2].iter().product(); + if batch_size <= 1 { + // No batch dims (or batch size 1) + let mut perm: Vec = (0..n_rows as i64).collect(); + for (i, piv_val) in piv_data.iter().enumerate().take(n_pivots) { + let j = piv_val.elem::() as usize; + perm.swap(i, j); + } + let perm_tensor = Tensor::<1, Int>::from_data(&perm[..], device); + return tensor.select(D - 2, perm_tensor); + } + + // If input tensor has batch dimensions, then flatten batch dims, + // iterate, then reshape back. + // Reshape tensor: [b1, b2, ..., bN, rows, cols] -> [B, rows, cols] + let n_cols = tensor_dims[D - 1]; + let flat_tensor: Tensor<3> = tensor.reshape([batch_size, n_rows, n_cols]); + // Reshape pivot: [b1, b2, ..., bN, n_pivots, 1] -> [B * n_pivots] + let mut results: Vec> = Vec::with_capacity(batch_size); + for b in 0..batch_size { + // Build permutation for this batch element + let mut perm: Vec = (0..n_rows as i64).collect(); + let offset = b * n_pivots; + for i in 0..n_pivots { + let j = piv_data[offset + i].elem::() as usize; + perm.swap(i, j); + } + let perm_tensor = Tensor::<1, Int>::from_data(&perm[..], device); + + // Extract this batch element [1, rows, cols], select rows, collect + let batch_elem = flat_tensor.clone().slice_dim(0, b); // [1, rows, cols] + let permuted = batch_elem.select(1, perm_tensor); // [1, rows, cols] + results.push(permuted); + } + + // Concatenate along batch dim and reshape back to original shape + let concatenated: Tensor<3> = Tensor::cat(results, 0); // [B, rows, cols] + concatenated.reshape(tensor_dims) +} + +/// Solves for the U blocks using forward substitution. +/// +/// Solves the equation L_{kk} U_{k, k+1:} = A_{k, k+1:} for U_{k, k+1:}. +/// +/// # Arguments +/// - `diagonal_l_block`: The L block L_{k, k}, [k_start..k_end, k_start..k_end] +/// - `row_blocks`: The row blocks A_{k, k+1}, ..., A_{k, N}, [k_start..k_end, k_end..] +/// - `block_size`: The size of the current block +fn solve_for_u_blocks( + diagonal_l_block: Tensor, + mut a_row_blocks: Tensor, + block_size: usize, +) -> Tensor { + // The first row requires no computation since the first row of + // diagonal_l_block is [1, 0, 0, ..., 0] + let mut slices = vec![Slice::full(); D]; + + for i in 1..block_size { + // Shape of each matrix: 1 by i + let l_multipliers = diagonal_l_block + .clone() + .slice_dim(D - 2, i) + .slice_dim(D - 1, 0..i); + // Shape of each matrix: i by c where c is the number of cols in a_row_blocks + let u_computed = a_row_blocks.clone().slice_dim(D - 2, 0..i); + let prod = l_multipliers.matmul(u_computed); + + let current_rows = a_row_blocks.clone().slice_dim(D - 2, i); + let updated_rows = current_rows - prod; + + // Update the i-th row of a_row_blocks with the solved values of U blocks + slices[D - 2] = Slice::from(i); + a_row_blocks = a_row_blocks.slice_assign(&slices, updated_rows); + } + + a_row_blocks.clone() +} diff --git a/crates/burn-tensor/src/tensor/linalg/lu_decomposition.rs b/crates/burn-tensor/src/tensor/linalg/lu_decomposition.rs new file mode 100644 index 0000000..4b78ca5 --- /dev/null +++ b/crates/burn-tensor/src/tensor/linalg/lu_decomposition.rs @@ -0,0 +1,78 @@ +use crate::{ + Int, cast::ToElement, check, check::TensorCheck, linalg::swap_slices, s, tensor::Tensor, +}; +/// Performs PLU decomposition of a square matrix. +/// +/// The function decomposes a given square matrix `A` into three matrices: a permutation vector `p`, +/// a lower triangular matrix `L`, and an upper triangular matrix `U`, such that `PA = LU`. +/// The permutation vector `p` represents the row swaps made during the decomposition process. +/// The lower triangular matrix `L` has ones on its diagonal and contains the multipliers used +/// during the elimination process below the diagonal. The upper triangular matrix `U` contains +/// the resulting upper triangular form of the matrix after the elimination process. +/// +/// # Arguments +/// * `tensor` - A square matrix to decompose, represented as a 2D tensor. +/// +/// # Returns +/// A tuple containing: +/// - A 2D tensor representing the combined `L` and `U` matrices. +/// - A 1D tensor representing the permutation vector `p`. +/// +/// # Panics and numerical issues +/// - The function will panic if the input matrix is singular or near-singular. +/// - The function will panic if the input matrix is not square. +/// # Performance note (synchronization / device transfers) +/// This function may involve multiple synchronizations and device transfers, especially +/// when determining pivot elements and performing row swaps. This can impact performance, +pub fn lu_decomposition(tensor: Tensor<2>) -> (Tensor<2>, Tensor<1, Int>) { + check!(TensorCheck::is_square::<2>( + "lu_decomposition", + &tensor.shape() + )); + let dims = tensor.shape().dims::<2>(); + let n = dims[0]; + + let mut permutations = Tensor::arange(0..n as i64, &tensor.device()); + let mut tensor = tensor; + + for k in 0..n { + // Find the pivot row + let p = tensor + .clone() + .slice(s![k.., k]) + .abs() + .argmax(0) + .into_scalar() + .to_usize() + + k; + let max = tensor.clone().slice(s![p, k]).abs(); + + // Avoid division by zero + let pivot = max.into_scalar(); + check!(TensorCheck::lu_decomposition_pivot(pivot)); + + if p != k { + tensor = swap_slices(tensor, s![k, ..], s![p, ..]); + permutations = swap_slices(permutations, s![k], s![p]); + } + + // Normalize k-th column under the diagonal + if k < n - 1 { + let a_kk = tensor.clone().slice(s![k, k]); + let column = tensor.clone().slice(s![(k + 1).., k]) / a_kk; + tensor = tensor.slice_assign(s![(k + 1).., k], column); + } + + // Update the trailing submatrix + for i in (k + 1)..n { + // a[i, k+1..] -= a[i, k] * a[k, k+1..] + let a_ik = tensor.clone().slice(s![i, k]); + let row_k = tensor.clone().slice(s![k, (k + 1)..]); + let update = a_ik * row_k; + let row_i = tensor.clone().slice(s![i, (k + 1)..]); + tensor = tensor.slice_assign(s![i, (k + 1)..], row_i - update); + } + } + + (tensor, permutations) +} diff --git a/crates/burn-tensor/src/tensor/linalg/matvec.rs b/crates/burn-tensor/src/tensor/linalg/matvec.rs new file mode 100644 index 0000000..128a03f --- /dev/null +++ b/crates/burn-tensor/src/tensor/linalg/matvec.rs @@ -0,0 +1,60 @@ +use crate::{ + kind::Numeric, + tensor::{Shape, Tensor}, +}; + +/// Performs matrix-vector multiplication with optional batch dimensions. +/// +/// The `matrix` tensor is expected to have rank `DM` with the last two dimensions representing +/// the matrix rows and columns. The `vector` tensor should have rank `DV = DM - 1`, sharing +/// broadcast-compatible batch dimensions and matching the last dimension of the matrix. +/// +/// # Panics +/// +/// * If the matrix rank is lower than 2. +/// * If the vector rank isn't one less than the matrix rank. +/// * If batch dimensions differ between the operands. +/// * If the inner dimensions are incompatible for multiplication. +pub fn matvec( + matrix: Tensor, + vector: Tensor, +) -> Tensor +where + K: Numeric, +{ + assert!( + DM >= 2, + "matvec expects the matrix to be at least rank 2 (got {DM})" + ); + assert!( + DM == DV + 1, + "matvec expects the vector rank ({DV}) to be exactly one less than the matrix rank ({DM})", + ); + + let matrix_dims = matrix.shape().dims::(); + let vector_dims = vector.shape().dims::(); + + // Validate batch dimensions (all leading dimensions prior to the matrix axes). + let batch_rank = DM.saturating_sub(2); + if batch_rank > 0 { + let matrix_batch = Shape::from(&matrix_dims[..batch_rank]); + let vector_batch = Shape::from(&vector_dims[..batch_rank]); + + assert!( + matrix_batch.broadcast(&vector_batch).is_ok(), + "Batch dimensions are not broadcast-compatible: matrix {:?} vs vector {:?}", + &matrix_dims[..batch_rank], + &vector_dims[..batch_rank] + ); + } + + let matrix_inner = matrix_dims[DM - 1]; + let vector_inner = vector_dims[DV - 1]; + assert!( + matrix_inner == vector_inner, + "Inner dimension mismatch: matrix has {matrix_inner} columns but vector has {vector_inner} entries", + ); + + let vector_expanded = vector.unsqueeze_dim::(DV); + matrix.matmul(vector_expanded).squeeze_dim::(DM - 1) +} diff --git a/crates/burn-tensor/src/tensor/linalg/mod.rs b/crates/burn-tensor/src/tensor/linalg/mod.rs new file mode 100644 index 0000000..38cdd1b --- /dev/null +++ b/crates/burn-tensor/src/tensor/linalg/mod.rs @@ -0,0 +1,17 @@ +mod cosine_similarity; +mod det; +mod diag; +mod lu; +mod matvec; +mod outer; +mod trace; +mod vector_norm; + +pub use cosine_similarity::*; +pub use det::*; +pub use diag::*; +pub use lu::*; +pub use matvec::*; +pub use outer::*; +pub use trace::*; +pub use vector_norm::*; diff --git a/crates/burn-tensor/src/tensor/linalg/outer.rs b/crates/burn-tensor/src/tensor/linalg/outer.rs new file mode 100644 index 0000000..f820f79 --- /dev/null +++ b/crates/burn-tensor/src/tensor/linalg/outer.rs @@ -0,0 +1,75 @@ +use crate::tensor::Tensor; +use crate::{AsIndex, kind::Numeric}; + +/// Computes the outer product for the last columns of 2 tensors. +/// +/// See also: [`outer_dim`]. +/// +/// # Arguments +/// - `lhs`: the "row" tensor, with shape ``[..., i]``. +/// - `rhs`: the "col" tensor, with shape ``[..., j]``. +/// +/// # Returns +/// +/// A tensor of rank `R = D + 1`, where: +/// +/// `` +/// result[..., i, j] = lhs[..., i] * rhs[..., j] +/// `` +pub fn outer( + lhs: Tensor, + rhs: Tensor, +) -> Tensor +where + K: Numeric, +{ + outer_dim(lhs, rhs, -1) +} + +/// Computes the outer product along a specific dimension, broadcasting over others. +/// +/// For the given `dim`, computes the outer product of elements along that dimension, +/// expanding it into two dimensions of size ``M × N`` at positions ``(dim, dim + 1)``. +/// +/// # Arguments +/// +/// - `lhs`: left operand, the "row" tensor, with size `M` at dimension `dim`. +/// - `rhs`: right operand, the "col" tensor, with size `N` at dimension `dim`. +/// - `dim`: dimension to compute the outer product along (supports negative indexing). +/// +/// # Returns +/// +/// A tensor of rank `R = D + 1`, where: +/// +/// `` +/// result[..., i, j, ...] = lhs[..., i, ...] * rhs[..., j, ...] +/// `` +// +// Notes: +// - For large batched inputs, `x_col.matmul(y_row)` *might* be more performant +// than broadcasted elemwise multiply; benchmarking needed to confirm. +pub fn outer_dim( + lhs: Tensor, + rhs: Tensor, + dim: Dim, +) -> Tensor +where + K: Numeric, +{ + assert_eq!( + R, + D + 1, + "`outer` with D={D} expects R={} (got R={R})", + D + 1 + ); + let dim = dim.expect_dim_index(D); + + // (..., i, 1, ...) + let x = lhs.unsqueeze_dim::(dim + 1); + + // (..., 1, j, ...) + let y = rhs.unsqueeze_dim::(dim); + + // (..., i, j, ...) + x * y +} diff --git a/crates/burn-tensor/src/tensor/linalg/trace.rs b/crates/burn-tensor/src/tensor/linalg/trace.rs new file mode 100644 index 0000000..b51d02d --- /dev/null +++ b/crates/burn-tensor/src/tensor/linalg/trace.rs @@ -0,0 +1,23 @@ +use super::diag; +use crate::tensor::Tensor; + +/// Computes the trace of a matrix. +/// +/// For batched inputs, computes the trace of each matrix in the batch independently. +/// +/// The trace operation sums the diagonal elements of the last two dimensions, +/// treating them as the matrix dimensions, while preserving all leading batch dimensions. +/// +/// # Arguments +/// +/// * `tensor` - The input tensor with at least 2 dimensions. +/// +/// # Returns +/// +/// A tensor of rank `D - 1`, where the last dimension contains the sum along the diagonals +/// of the input. +pub fn trace(tensor: Tensor) -> Tensor { + let diag_tensor = diag::(tensor); + + diag_tensor.sum_dim(DO - 1) +} diff --git a/crates/burn-tensor/src/tensor/linalg/vector_norm.rs b/crates/burn-tensor/src/tensor/linalg/vector_norm.rs new file mode 100644 index 0000000..18878eb --- /dev/null +++ b/crates/burn-tensor/src/tensor/linalg/vector_norm.rs @@ -0,0 +1,281 @@ +use crate::tensor::Tensor; +use crate::{ + ElementConversion, + kind::{Numeric, Ordered}, +}; +#[allow(unused_imports)] +use num_traits::float::Float; +/// Specifies the type of norm to compute. +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum Norm { + /// L0 norm (count of non-zero elements) + L0, + + /// L1 norm (sum of absolute values) + L1, + + /// L2 norm (Euclidean norm) + L2, + + /// L:INFINITY norm (maximum absolute value) + LInf, + + /// L:NEG_INFINITY norm (minimum absolute value) + LNegInf, + + /// Lp norm (generalized norm) + Lp(f64), +} + +impl Norm { + /// Get the exponent of the norm. + pub fn to_exponent(self) -> f64 { + use Norm::*; + match self { + L0 => 0.0, + L1 => 1.0, + L2 => 2.0, + LInf => f64::INFINITY, + LNegInf => f64::NEG_INFINITY, + Lp(p) => p, + } + } +} + +impl From for Norm { + fn from(value: u32) -> Self { + use Norm::*; + match value { + 0 => L0, + 1 => L1, + 2 => L2, + u32::MAX => LInf, + _ => Lp(value as f64), + } + } +} + +impl From for Norm { + fn from(value: i32) -> Self { + use Norm::*; + match value { + 0 => L0, + 1 => L1, + 2 => L2, + i32::MAX => LInf, + i32::MIN => LNegInf, + _ => Lp(value as f64), + } + } +} + +impl From for Norm { + fn from(value: f32) -> Self { + use Norm::*; + match value { + 0.0 => L0, + 1.0 => L1, + 2.0 => L2, + f32::INFINITY => LInf, + f32::NEG_INFINITY => LNegInf, + _ => Lp(value as f64), + } + } +} + +impl From for Norm { + fn from(value: f64) -> Self { + use Norm::*; + match value { + 0.0 => L0, + 1.0 => L1, + 2.0 => L2, + f64::INFINITY => LInf, + f64::NEG_INFINITY => LNegInf, + _ => Lp(value), + } + } +} + +/// Computes the vector norm of a tensor along a specified dimension. +/// +/// Generic dispatch wrapper over specialized / optimized norms. +/// +/// See: +/// - [torch.linalg.vector_norm](https://pytorch.org/docs/stable/generated/torch.linalg.vector_norm.html) +/// - [numpy.linalg.vector_norm](https://numpy.org/doc/stable/reference/generated/numpy.linalg.vector_norm.html) +/// +/// # Arguments +/// +/// * `x` - The input tensor. +/// * `norm` - The selected norm. +/// * `dim` - The dimension to compute the norm over. +/// +/// # Returns +/// +/// The vector norm of the input tensor. +pub fn vector_norm(x: Tensor, norm: impl Into, dim: usize) -> Tensor { + lp_norm(x, norm.into().to_exponent(), dim) +} + +/// Computes the general ``L(p)`` norm of a tensor along a specified dimension. +/// +/// Uses the specialized implementations for: +/// * 0.0 +/// * 1.0 +/// * 2.0 +/// * 2 * N for integral N, +/// * f64::INFINITY, +/// * f64::NEG_INFINITY, +/// +/// # Arguments +/// +/// * `x` - The input tensor. +/// * `p` - The exponent of the Lp norm. +/// * `dim` - The dimension to compute the norm over. +/// +/// # Returns +/// +/// The ``L(p)`` norm of the input tensor. +pub fn lp_norm(x: Tensor, p: f64, dim: usize) -> Tensor { + match p { + 0.0 => l0_norm(x, dim), + 1.0 => l1_norm(x, dim), + 2.0 => l2_norm(x, dim), + p if is_even_integer(p) => lp_signed_norm(x, p as u32, dim), + f64::INFINITY => max_abs_norm(x, dim), + f64::NEG_INFINITY => min_abs_norm(x, dim), + _ => lp_norm_base(x, p, dim), + } +} + +/// Normalize a tensor versus its `vector_norm`. +/// +/// Equivalent to ``x.clone() / vector_norm(x, norm, dim).clamp_min(eps)``. +/// +/// # Arguments +/// +/// * `x` - The input tensor. +/// * `norm` - The selected norm. +/// * `dim` - The dimension to compute the norm over. +/// * `eps` - The epsilon for the norm. +/// +/// # Returns +/// +/// The normalized tensor. +pub fn vector_normalize( + x: Tensor, + norm: impl Into, + dim: usize, + eps: E, +) -> Tensor { + let norm = vector_norm(x.clone(), norm, dim).clamp_min(eps); + x / norm +} + +/// Computes the L0 norm of a tensor along a specified dimension. +/// +/// # Arguments +/// +/// * `x` - The input tensor. +/// * `dim` - The dimension to compute the norm over. +/// +/// # Returns +/// +/// The L0 norm of the input tensor. +pub fn l0_norm(x: Tensor, dim: usize) -> Tensor +where + K: Numeric, +{ + x.zeros_like() + .mask_fill(x.not_equal_elem(0), 1) + .sum_dim(dim) +} + +/// Computes the L1 norm of a tensor along a specified dimension. +/// +/// This is a convenience function that wraps `vector_norm` with `p = 1.0`. +/// +/// # Arguments +/// +/// * `x` - The input tensor. +/// * `dim` - The dimension to compute the norm over. +/// +/// # Returns +/// +/// The L1 norm of the input tensor. +pub fn l1_norm(x: Tensor, dim: usize) -> Tensor +where + K: Numeric, +{ + x.abs().sum_dim(dim) +} + +/// Computes the L2 norm of a tensor along a specified dimension. +/// +/// # Arguments +/// +/// * `x` - The input tensor. +/// * `dim` - The dimension to compute the norm over. +/// +/// # Returns +/// +/// The L2 norm of the input tensor. +pub fn l2_norm(x: Tensor, dim: usize) -> Tensor { + x.square().sum_dim(dim).sqrt() +} + +fn is_even_integer(x: f64) -> bool { + x.fract() == 0.0 && (x as i64) % 2 == 0 +} + +/// Computes ``L(2*n)`` for even integer ``n``. +/// +/// This lets us skip the abs. +fn lp_signed_norm(x: Tensor, p: u32, dim: usize) -> Tensor { + x.powi_scalar(p).sum_dim(dim).powf_scalar(1. / (p as f64)) +} + +/// Computes the general ``L(p)`` using the generalized method. +/// +/// This uses no specialized implementations and cannot handle: +/// * 0.0 +/// * f64::INFINITY, +/// * f64::NEG_INFINITY, +fn lp_norm_base(x: Tensor, p: f64, dim: usize) -> Tensor { + x.abs().powf_scalar(p).sum_dim(dim).powf_scalar(1. / p) +} + +/// Computes the L:INFINITY norm of a tensor along a specified dimension. +/// +/// # Arguments +/// +/// * `x` - The input tensor. +/// * `dim` - The dimension to compute the norm over. +/// +/// # Returns +/// +/// The L:INFINITY norm of the input tensor. +pub fn max_abs_norm(x: Tensor, dim: usize) -> Tensor +where + K: Ordered, +{ + x.max_abs_dim(dim) +} + +/// Computes the L:NEG_INFINITY norm of a tensor along a specified dimension. +/// +/// # Arguments +/// +/// * `x` - The input tensor. +/// * `dim` - The dimension to compute the norm over. +/// +/// # Returns +/// +/// The L:NEG_INFINITY norm of the input tensor. +pub fn min_abs_norm(x: Tensor, dim: usize) -> Tensor +where + K: Ordered, +{ + x.abs().min_dim(dim) +} diff --git a/crates/burn-tensor/src/tensor/loss/mod.rs b/crates/burn-tensor/src/tensor/loss/mod.rs new file mode 100644 index 0000000..8976e2c --- /dev/null +++ b/crates/burn-tensor/src/tensor/loss/mod.rs @@ -0,0 +1,22 @@ +use crate::{Tensor, activation}; + +/// Computes the log softmax cross entropy between logits and target probabilities. +/// +/// # Arguments +/// +/// * `logits` - The logits. +/// * `target_probs` - The target probabilities. +/// +/// # Returns +/// +/// The log softmax cross entropy. +pub fn cross_entropy_with_logits( + logits: Tensor, + target_probs: Tensor, +) -> Tensor<1> { + let tensor = activation::log_softmax(logits, D - 1); + let tensor = tensor.mul(target_probs); + let tensor = tensor.sum_dim(D - 1); + + tensor.mean().neg() +} diff --git a/crates/burn-tensor/src/tensor/mod.rs b/crates/burn-tensor/src/tensor/mod.rs new file mode 100644 index 0000000..2f0f209 --- /dev/null +++ b/crates/burn-tensor/src/tensor/mod.rs @@ -0,0 +1,56 @@ +pub(crate) mod stats; + +mod api; + +pub use api::*; + +// Re-exported types +pub use burn_std::{ + BoolDType, BoolStore, DType, DataError, FloatDType, IndexingUpdateOp, IntDType, TensorData, + Tolerance, distribution::*, element::*, indexing::*, s, shape::*, slice::*, +}; + +/// The tensor kind module. +pub mod kind; +pub use kind::{Bool, Float, Int}; + +/// The activation module. +pub mod activation; + +/// The container module. +pub mod container { + pub use burn_std::tensor::container::TensorContainer; +} + +/// The grid module. +pub mod grid; + +/// The linalg module. +pub mod linalg; + +/// The loss module. +pub mod loss; + +/// The neural network module. +pub mod module; + +/// The signal processing module. +pub mod signal; + +/// Operations on tensors module. +pub mod ops { + pub(crate) use crate::bridge::*; + pub use burn_std::ops::*; +} + +/// Tensor quantization module. +pub mod quantization; + +#[cfg(feature = "std")] +pub mod distributed; + +#[cfg(feature = "std")] +pub use report::*; + +#[cfg(feature = "std")] +mod report; diff --git a/crates/burn-tensor/src/tensor/module.rs b/crates/burn-tensor/src/tensor/module.rs new file mode 100644 index 0000000..6d99290 --- /dev/null +++ b/crates/burn-tensor/src/tensor/module.rs @@ -0,0 +1,671 @@ +use burn_backend::ops::ModuleOps; +use burn_dispatch::Dispatch; +use burn_std::{MatmulTransformAction, MatmulTransformAnalysis, MatmulTransformPolicy}; + +use crate::{ + Bool, DType, Int, Tensor, check, + check::TensorCheck, + ops::{ + AttentionModuleOptions, BridgeTensor, ConvOptions, ConvTransposeOptions, DeformConvOptions, + InterpolateOptions, PadMode, PaddedConvOptions, UnfoldOptions, + }, +}; + +/// Computes the [CTC loss](burn_backend::ops::ModuleOps::ctc_loss). +/// +/// # Arguments +/// +/// * `log_probs` - Log-probabilities of shape `[T, N, C]` +/// * `targets` - Target label indices of shape `[N, S]` +/// * `input_lengths` - Actual input sequence lengths per batch element `[N]` +/// * `target_lengths` - Actual target lengths per batch element `[N]` +/// * `blank` - Index of the blank label +/// +/// # Returns +/// +/// Per-sample loss of shape `[N]` +pub fn ctc_loss( + log_probs: Tensor<3>, + targets: Tensor<2, Int>, + input_lengths: Tensor<1, Int>, + target_lengths: Tensor<1, Int>, + blank: usize, +) -> Tensor<1> { + Tensor::new(BridgeTensor::float(Dispatch::ctc_loss( + log_probs.primitive.into_float(), + targets.primitive.into(), + input_lengths.primitive.into(), + target_lengths.primitive.into(), + blank, + ))) +} + +/// Applies the [embedding module](burn_backend::ops::ModuleOps::embedding). +pub fn embedding(weights: Tensor<2>, indices: Tensor<2, Int>) -> Tensor<3> { + Tensor::new(BridgeTensor::float(Dispatch::embedding( + weights.primitive.into_float(), + indices.primitive.into(), + ))) +} + +/// Applies a [1D convolution](burn_backend::ops::ModuleOps::conv1d). +/// +/// Accepts [`ConvOptions`] for symmetric padding, or [`PaddedConvOptions`] for +/// asymmetric padding. When asymmetric padding is specified, an explicit pad +/// operation is applied before the convolution backend op. +pub fn conv1d( + x: Tensor<3>, + weight: Tensor<3>, + bias: Option>, + options: impl Into>, +) -> Tensor<3> { + let padded_options = options.into(); + check!(TensorCheck::conv( + "conv1d", + x.dims(), + weight.dims(), + padded_options.options.groups, + )); + + if let Some(padding_end) = padded_options.padding_end { + let left = padded_options.options.padding[0]; + let right = padding_end[0]; + // For 1D (NCL format), pad the length dimension + let padded = x.pad((left, right, 0, 0), PadMode::Constant(0.0)); + let zero_options = ConvOptions::new( + padded_options.options.stride, + [0], + padded_options.options.dilation, + padded_options.options.groups, + ); + Tensor::new(BridgeTensor::float(Dispatch::conv1d( + padded.primitive.into_float(), + weight.primitive.into_float(), + bias.map(|b| b.primitive.into_float()), + zero_options, + ))) + } else { + Tensor::new(BridgeTensor::float(Dispatch::conv1d( + x.primitive.into_float(), + weight.primitive.into_float(), + bias.map(|b| b.primitive.into_float()), + padded_options.options, + ))) + } +} + +/// Applies a [2D convolution](burn_backend::ops::ModuleOps::conv2d). +/// +/// Accepts [`ConvOptions`] for symmetric padding, or [`PaddedConvOptions`] for +/// asymmetric padding. When asymmetric padding is specified, an explicit pad +/// operation is applied before the convolution backend op. +pub fn conv2d( + x: Tensor<4>, + weight: Tensor<4>, + bias: Option>, + options: impl Into>, +) -> Tensor<4> { + let padded_options = options.into(); + check!(TensorCheck::conv( + "conv2d", + x.dims(), + weight.dims(), + padded_options.options.groups, + )); + + if let Some(padding_end) = padded_options.padding_end { + let top = padded_options.options.padding[0]; + let left = padded_options.options.padding[1]; + let bottom = padding_end[0]; + let right = padding_end[1]; + // For 2D (NCHW format), pad height and width + let padded = x.pad((left, right, top, bottom), PadMode::Constant(0.0)); + let zero_options = ConvOptions::new( + padded_options.options.stride, + [0, 0], + padded_options.options.dilation, + padded_options.options.groups, + ); + Tensor::new(BridgeTensor::float(Dispatch::conv2d( + padded.primitive.into_float(), + weight.primitive.into_float(), + bias.map(|b| b.primitive.into_float()), + zero_options, + ))) + } else { + Tensor::new(BridgeTensor::float(Dispatch::conv2d( + x.primitive.into_float(), + weight.primitive.into_float(), + bias.map(|b| b.primitive.into_float()), + padded_options.options, + ))) + } +} + +/// Applies a [3D convolution](burn_backend::ops::ModuleOps::conv3d). +/// +/// Accepts [`ConvOptions`] for symmetric padding, or [`PaddedConvOptions`] for +/// asymmetric padding. Asymmetric 3D padding is not yet supported. +pub fn conv3d( + x: Tensor<5>, + weight: Tensor<5>, + bias: Option>, + options: impl Into>, +) -> Tensor<5> { + let padded_options = options.into(); + check!(TensorCheck::conv( + "conv3d", + x.dims(), + weight.dims(), + padded_options.options.groups, + )); + + if padded_options.is_asymmetric() { + panic!("Asymmetric padding is not yet supported for conv3d"); + } + + Tensor::new(BridgeTensor::float(Dispatch::conv3d( + x.primitive.into_float(), + weight.primitive.into_float(), + bias.map(|b| b.primitive.into_float()), + padded_options.options, + ))) +} + +/// Applies a [Deformable 2D convolution](burn_backend::ops::ModuleOps::deform_conv2d). +pub fn deform_conv2d( + x: Tensor<4>, + offset: Tensor<4>, + weight: Tensor<4>, + mask: Option>, + bias: Option>, + options: DeformConvOptions<2>, +) -> Tensor<4> { + check!(TensorCheck::conv( + "deform_conv2d", + x.dims(), + weight.dims(), + options.weight_groups, + )); + Tensor::new(BridgeTensor::float(Dispatch::deform_conv2d( + x.primitive.into_float(), + offset.primitive.into_float(), + weight.primitive.into_float(), + mask.map(|m| m.primitive.into_float()), + bias.map(|b| b.primitive.into_float()), + options, + ))) +} + +/// Applies a [1D transposed convolution](burn_backend::ops::ModuleOps::conv_transpose1d). +pub fn conv_transpose1d( + x: Tensor<3>, + weight: Tensor<3>, + bias: Option>, + options: ConvTransposeOptions<1>, +) -> Tensor<3> { + check!(TensorCheck::conv_transpose( + "conv_transpose1d", + x.dims(), + weight.dims(), + )); + Tensor::new(BridgeTensor::float(Dispatch::conv_transpose1d( + x.primitive.into_float(), + weight.primitive.into_float(), + bias.map(|b| b.primitive.into_float()), + options, + ))) +} + +/// Applies a [2D transposed convolution](burn_backend::ops::ModuleOps::conv_transpose2d). +pub fn conv_transpose2d( + x: Tensor<4>, + weight: Tensor<4>, + bias: Option>, + options: ConvTransposeOptions<2>, +) -> Tensor<4> { + check!(TensorCheck::conv_transpose( + "conv_transpose2d", + x.dims(), + weight.dims(), + )); + Tensor::new(BridgeTensor::float(Dispatch::conv_transpose2d( + x.primitive.into_float(), + weight.primitive.into_float(), + bias.map(|b| b.primitive.into_float()), + options, + ))) +} + +/// Applies a 3D transposed convolution](burn_backend::ops::ModuleOps::conv_transpose3d). +pub fn conv_transpose3d( + x: Tensor<5>, + weight: Tensor<5>, + bias: Option>, + options: ConvTransposeOptions<3>, +) -> Tensor<5> { + check!(TensorCheck::conv_transpose( + "conv_transpose3d", + x.dims(), + weight.dims(), + )); + Tensor::new(BridgeTensor::float(Dispatch::conv_transpose3d( + x.primitive.into_float(), + weight.primitive.into_float(), + bias.map(|b| b.primitive.into_float()), + options, + ))) +} + +/// Applies a [4D to 3D unfold](burn_backend::ops::ModuleOps::unfold4d). +pub fn unfold4d(x: Tensor<4>, kernel_size: [usize; 2], options: UnfoldOptions) -> Tensor<3> { + Tensor::new(BridgeTensor::float(Dispatch::unfold4d( + x.primitive.into_float(), + kernel_size, + options, + ))) +} + +/// Applies a [1D max pooling](burn_backend::ops::ModuleOps::max_pool1d). +pub fn max_pool1d( + x: Tensor<3>, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, +) -> Tensor<3> { + Tensor::new(BridgeTensor::float(Dispatch::max_pool1d( + x.primitive.into_float(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + ))) +} + +/// Applies a [2D max pooling](burn_backend::ops::ModuleOps::max_pool2d). +pub fn max_pool2d( + x: Tensor<4>, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, +) -> Tensor<4> { + Tensor::new(BridgeTensor::float(Dispatch::max_pool2d( + x.primitive.into_float(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + ))) +} + +/// Applies a [2D avg pooling](burn_backend::ops::ModuleOps::avg_pool2d). +pub fn avg_pool2d( + x: Tensor<4>, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, +) -> Tensor<4> { + Tensor::new(BridgeTensor::float(Dispatch::avg_pool2d( + x.primitive.into_float(), + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ))) +} + +/// Applies a [1D avg pooling](burn_backend::ops::ModuleOps::avg_pool1d). +pub fn avg_pool1d( + x: Tensor<3>, + kernel_size: usize, + stride: usize, + padding: usize, + count_include_pad: bool, + ceil_mode: bool, +) -> Tensor<3> { + Tensor::new(BridgeTensor::float(Dispatch::avg_pool1d( + x.primitive.into_float(), + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ))) +} + +/// Applies a [1D max pooling](burn_backend::ops::ModuleOps::max_pool1d). +pub fn max_pool1d_with_indices( + x: Tensor<3>, + kernel_size: usize, + stride: usize, + padding: usize, + dilation: usize, + ceil_mode: bool, +) -> (Tensor<3>, Tensor<3, Int>) { + let indices_dtype = x.device().settings().int_dtype; + let output = Dispatch::max_pool1d_with_indices( + x.primitive.into_float(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + indices_dtype, + ); + + ( + Tensor::new(BridgeTensor::float(output.output)), + Tensor::new(BridgeTensor::int(output.indices)), + ) +} + +/// Applies a [2D max pooling with indices](burn_backend::ops::ModuleOps::max_pool2d_with_indices). +pub fn max_pool2d_with_indices( + x: Tensor<4>, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, +) -> (Tensor<4>, Tensor<4, Int>) { + let indices_dtype = x.device().settings().int_dtype; + let output = Dispatch::max_pool2d_with_indices( + x.primitive.into_float(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + indices_dtype, + ); + + ( + Tensor::new(BridgeTensor::float(output.output)), + Tensor::new(BridgeTensor::int(output.indices)), + ) +} + +/// Applies a [2D adaptive avg pooling](burn_backend::ops::ModuleOps::adaptive_avg_pool2d). +pub fn adaptive_avg_pool2d(x: Tensor<4>, output_size: [usize; 2]) -> Tensor<4> { + Tensor::new(BridgeTensor::float(Dispatch::adaptive_avg_pool2d( + x.primitive.into_float(), + output_size, + ))) +} + +/// Applies a [1D adaptive avg pooling](burn_backend::ops::ModuleOps::adaptive_avg_pool1d). +pub fn adaptive_avg_pool1d(x: Tensor<3>, output_size: usize) -> Tensor<3> { + Tensor::new(BridgeTensor::float(Dispatch::adaptive_avg_pool1d( + x.primitive.into_float(), + output_size, + ))) +} + +/// Applies a [2D interpolation](burn_backend::ops::ModuleOps::interpolate). +pub fn interpolate( + x: Tensor<4>, + output_size: [usize; 2], + options: InterpolateOptions, +) -> Tensor<4> { + Tensor::new(BridgeTensor::float(Dispatch::interpolate( + x.primitive.into_float(), + output_size, + options, + ))) +} + +/// Applies a linear transformation to the input tensor using the given weight and bias. +/// +/// ```math +/// y = x @ weight + [bias] +/// ``` +/// +/// # Arguments: +/// +/// - `input` is the input tensor, ``[..., d_input]``. +/// - `weight` is the weight tensor, ``[d_input, d_output]``. +/// - `bias` is the bias tensor (optional), ``[d_output]``. +/// +/// # Returns: +/// +/// The transformed tensor, ``[..., d_output]``. +/// +/// # Compatibility +/// +/// This function differs from PyTorch's ``torch.nn.functional.linear`` in that it does not +/// transpose the weight matrix. In PyTorch, the weight matrix is transposed before +/// multiplication: +/// +/// ```math +/// y = x @ weight^T + [bias] +/// ``` +pub fn linear( + input: Tensor, + weight: Tensor<2>, + bias: Option>, +) -> Tensor { + if D == 1 { + // Insert and remove an extra batch dimension for the batch matmul to work. + let input = input.unsqueeze::<2>(); + let output = linear(input, weight, bias); + return output.squeeze_dim(0); + } + + // A quantized weight must stay quantized: `linear_impl` converts its + // operands to float, which would dequantize (materialize) the whole weight + // matrix on every forward. Route through the quantized matmul instead, which + // streams the packed weight directly — but reuse the same batch-fold policy + // the float `linear` applies, so a decode-shaped call folds its batches into + // the rows for one `[rows, d_in] @ [d_in, d_out]` matmul rather than a + // broadcast batched matmul that re-reads the packed weight per batch. + if let DType::QFloat(_) = weight.dtype() { + let dims = input.dims(); + let analysis = MatmulTransformAnalysis::from_shapes(&input.shape(), &weight.shape()); + + let output = match MatmulTransformPolicy::default().action(&analysis) { + MatmulTransformAction::MergeBatches { rows } => { + let d_in = dims[D - 1]; + let d_out = weight.dims()[1]; + + let folded = input.reshape([rows, d_in]).matmul(weight); + + let mut out_dims = dims; + out_dims[D - 1] = d_out; + folded.reshape(out_dims) + } + MatmulTransformAction::Keep => input.matmul(weight.unsqueeze::()), + }; + + return match bias { + Some(bias) => output + bias.unsqueeze(), + None => output, + }; + } + + Tensor::new(linear_impl( + input.primitive, + weight.primitive, + bias.map(|b| b.primitive), + )) +} + +fn linear_impl( + input: BridgeTensor, + weight: BridgeTensor, + bias: Option, +) -> BridgeTensor { + BridgeTensor::float(Dispatch::linear( + input.into_float(), + weight.into_float(), + bias.map(|b| b.into_float()), + )) +} + +/// Computes scaled dot-product attention: softmax(QKᵗ * scale) · V, +/// where scale defaults to 1/sqrt(head_dim) (configurable via `options.scale`). +/// Optionally applies masking, additive bias, causal masking, and softcap. +/// +/// # Arguments +/// - `query`: Query tensor of shape `[batch_size, num_heads, seq_len_q, head_dim]` +/// - `key`: Key tensor of shape `[batch_size, num_heads, seq_len_k, head_dim]` +/// - `value`: Value tensor of shape `[batch_size, num_heads, seq_len_k, val_dim]` +/// - `mask`: Optional boolean mask of shape `[batch_size, num_heads, seq_len_q, seq_len_k]`, +/// where `true` indicates positions to mask (i.e. set to -inf before softmax). +/// - `attn_bias`: Optional float tensor of shape `[batch_size, num_heads, seq_len_q, seq_len_k]` +/// added to the attention scores before softmax (e.g. ALiBi, relative position biases). +/// - `options`: Additional attention options (custom scale, softcap, causal masking). +/// +/// # Returns +/// A tensor of shape `[batch_size, num_heads, seq_len_q, val_dim]` +/// representing the attended context per head. +/// +/// # Note +/// This implementation does not support dropout and is intended for inference or +/// use cases where dropout is not needed. +pub fn attention( + query: Tensor<4>, + key: Tensor<4>, + value: Tensor<4>, + mask: Option>, + attn_bias: Option>, + options: AttentionModuleOptions, +) -> Tensor<4> { + Tensor::new(BridgeTensor::float(Dispatch::attention( + query.primitive.into_float(), + key.primitive.into_float(), + value.primitive.into_float(), + mask.map(|mask| mask.primitive.into()), + attn_bias.map(|bias| bias.primitive.into_float()), + options, + ))) +} + +/// Exports attention fallback to test backend's attention against. +pub fn attention_fallback( + query: Tensor<4>, + key: Tensor<4>, + value: Tensor<4>, + mask: Option>, + attn_bias: Option>, + options: AttentionModuleOptions, +) -> Tensor<4> { + Tensor::new(BridgeTensor::float( + burn_backend::ops::attention::attention_fallback::( + query.primitive.into_float(), + key.primitive.into_float(), + value.primitive.into_float(), + mask.map(|mask| mask.primitive.into()), + attn_bias.map(|bias| bias.primitive.into_float()), + options, + ), + )) +} + +/// Calculate the [2D convolution](burn_backend::ops::ModuleOps::conv2d) backward pass, returning the gradient for `weight`. +pub fn conv2d_weight_backward( + x: Tensor<4>, + weight: Tensor<4>, + output_grad: Tensor<4>, + options: ConvOptions<2>, +) -> Tensor<4> { + Tensor::new(BridgeTensor::float(Dispatch::conv2d_weight_backward( + x.primitive.into_float(), + weight.primitive.into_float(), + output_grad.primitive.into_float(), + options, + ))) +} + +/// Backward pass for the [avg pooling 2d](ModuleOps::avg_pool2d) operation. +pub fn avg_pool2d_backward( + x: Tensor<4>, + grad: Tensor<4>, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + count_include_pad: bool, + ceil_mode: bool, +) -> Tensor<4> { + Tensor::new(BridgeTensor::float(Dispatch::avg_pool2d_backward( + x.primitive.into_float(), + grad.primitive.into_float(), + kernel_size, + stride, + padding, + count_include_pad, + ceil_mode, + ))) +} + +/// Backward pass for the [max pooling 2d](ModuleOps::max_pool2d_with_indices) operation. +#[allow(clippy::too_many_arguments)] +pub fn max_pool2d_with_indices_backward( + x: Tensor<4>, + kernel_size: [usize; 2], + stride: [usize; 2], + padding: [usize; 2], + dilation: [usize; 2], + ceil_mode: bool, + output_grad: Tensor<4>, + indices: Tensor<4, Int>, +) -> Tensor<4> { + Tensor::new(BridgeTensor::float( + Dispatch::max_pool2d_with_indices_backward( + x.primitive.into_float(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + output_grad.primitive.into_float(), + indices.primitive.into(), + ) + .x_grad, + )) +} + +/// Applies Layer Normalization over the last dimension of the input tensor. +/// +/// Computes `(x - mean) / sqrt(var + epsilon) * gamma + beta`, where `mean` and +/// (biased) `var` are reduced over the last axis. +/// +/// # Shapes +/// +/// - input: `[..., any, d_model]` +/// - output: `[..., any, d_model]` +pub fn layer_norm( + input: Tensor, + gamma: Tensor<1>, + beta: Option>, + epsilon: f64, +) -> Tensor { + Tensor::new(layer_norm_impl( + input.primitive, + gamma.primitive, + beta.map(|b| b.primitive), + epsilon, + )) +} + +fn layer_norm_impl( + input: BridgeTensor, + gamma: BridgeTensor, + beta: Option, + epsilon: f64, +) -> BridgeTensor { + BridgeTensor::float(Dispatch::layer_norm( + input.into_float(), + gamma.into_float(), + beta.map(|b| b.into_float()), + epsilon, + )) +} diff --git a/crates/burn-tensor/src/tensor/quantization.rs b/crates/burn-tensor/src/tensor/quantization.rs new file mode 100644 index 0000000..1edb42c --- /dev/null +++ b/crates/burn-tensor/src/tensor/quantization.rs @@ -0,0 +1,57 @@ +use crate::{ + Tensor, + ops::{BridgeKind, BridgeTensor}, +}; +use burn_backend::quantization; + +// User-facing quantization data types come from burn-std. +use burn_dispatch::Dispatch; +pub use burn_std::quantization::*; + +/// The tensor quantization parameters. +pub type QuantizationParameters = QParams>; + +/// The observed input calibration range. +#[derive(Clone, Debug)] +pub struct CalibrationRange { + /// Minimum observed value(s). + pub min: Tensor<1>, + /// Maximum observed value(s). + pub max: Tensor<1>, +} + +/// Compute the quantization range mapping. +pub fn compute_range( + scheme: &QuantScheme, + tensor: &Tensor, + calibration: &Calibration, +) -> CalibrationRange { + let (kind, inner) = tensor.primitive.as_parts(); + let (min, max) = match kind { + BridgeKind::Float => { + quantization::compute_range::(scheme, inner.clone(), calibration) + } + BridgeKind::QFloat => unreachable!(), + _ => panic!("Should be Float primitive kind"), + }; + + CalibrationRange { + min: Tensor::new(BridgeTensor::float(min)), + max: Tensor::new(BridgeTensor::float(max)), + } +} + +/// Compute the quantization parameters. +pub fn compute_q_params(scheme: &QuantScheme, range: CalibrationRange) -> QuantizationParameters { + let (min_kind, min) = range.min.primitive.into_parts(); + let (max_kind, max) = range.max.primitive.into_parts(); + match (min_kind, max_kind) { + (BridgeKind::Float, BridgeKind::Float) => { + let qparams = quantization::compute_q_params::(scheme, min, max); + QuantizationParameters { + scales: Tensor::new(BridgeTensor::float(qparams.scales)), + } + } + _ => unreachable!(), + } +} diff --git a/crates/burn-tensor/src/tensor/report.rs b/crates/burn-tensor/src/tensor/report.rs new file mode 100644 index 0000000..1d2f76f --- /dev/null +++ b/crates/burn-tensor/src/tensor/report.rs @@ -0,0 +1,105 @@ +use super::Tensor; + +use colored::*; + +/// Checks the closeness of two tensors and prints the results. +/// +/// Compares tensors by checking the absolute difference between each element. +/// Prints the percentage of elements within specified tolerances. +/// +/// # Arguments +/// +/// * `output` - The output tensor. +/// * `expected` - The expected tensor. +/// +/// # Example +/// +/// ```no_run +/// use burn_tensor::{check_closeness, Tensor}; +/// +/// fn example() { +/// let device = Default::default(); +/// let tensor1 = Tensor::<1>::from_floats( +/// [1.0, 2.0, 3.0, 4.0, 5.0, 6.001, 7.002, 8.003, 9.004, 10.1], +/// &device, +/// ); +/// let tensor2 = Tensor::<1>::from_floats( +/// [1.0, 2.0, 3.0, 4.000, 5.0, 6.0, 7.001, 8.002, 9.003, 10.004], +/// &device, +/// ); +/// check_closeness(&tensor1, &tensor2); +///} +/// ``` +/// +/// # Output +/// +/// ```text +/// Tensor Closeness Check Results: +/// =============================== +/// Epsilon: 1e-1 +/// Close elements: 10/10 (100.00%) +/// [PASS] All elements are within tolerance +/// +/// Epsilon: 1e-2 +/// Close elements: 10/10 (100.00%) +/// [PASS] All elements are within tolerance +/// +/// Epsilon: 1e-3 +/// Close elements: 9/10 (90.00%) +/// [WARN] Most elements are within tolerance +/// +/// Epsilon: 1e-4 +/// Close elements: 6/10 (60.00%) +/// [FAIL] Significant differences detected +/// +/// Epsilon: 1e-5 +/// Close elements: 5/10 (50.00%) +/// [FAIL] Significant differences detected +/// +/// Epsilon: 1e-6 +/// Close elements: 5/10 (50.00%) +/// [FAIL] Significant differences detected +/// +/// Epsilon: 1e-7 +/// Close elements: 5/10 (50.00%) +/// [FAIL] Significant differences detected +/// +/// Epsilon: 1e-8 +/// Close elements: 5/10 (50.00%) +/// [FAIL] Significant differences detected +/// +/// Closeness check complete. +/// ``` +pub fn check_closeness(output: &Tensor, expected: &Tensor) { + println!("{}", "Tensor Closeness Check Results:".bold()); + println!("==============================="); + + for epsilon in [1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6, 1e-7, 1e-8].iter() { + println!("{} {:e}", "Epsilon:".bold(), epsilon); + + let close = output + .clone() + .is_close(expected.clone(), Some(*epsilon), Some(*epsilon)); + let data = close.clone().into_data(); + let num_elements = data.num_elements(); + + // Count the number of elements that are close (true) + let count = data.iter::().filter(|x| *x).count(); + + let percentage = (count as f64 / num_elements as f64) * 100.0; + + println!(" Close elements: {count}/{num_elements} ({percentage:.2}%)"); + + if percentage == 100.0 { + println!(" {} All elements are within tolerance", "[PASS]".green()); + } else if percentage >= 90.0 { + println!(" {} Most elements are within tolerance", "[WARN]".yellow()); + } else { + println!(" {} Significant differences detected", "[FAIL]".red()); + } + + println!(); + } + + println!("{}", "Closeness check complete.".bold()); +} diff --git a/crates/burn-tensor/src/tensor/signal/blackman_window.rs b/crates/burn-tensor/src/tensor/signal/blackman_window.rs new file mode 100644 index 0000000..c142f05 --- /dev/null +++ b/crates/burn-tensor/src/tensor/signal/blackman_window.rs @@ -0,0 +1,92 @@ +use crate::{Float, Int, Tensor, TensorCreationOptions, check, check::TensorCheck}; + +/// Creates a 1D Blackman window tensor. +/// +#[cfg_attr( + doc, + doc = r#" +$$w_n = 0.42 - 0.5 \cos\left(\frac{2\pi n}{N}\right) + 0.08 \cos\left(\frac{4\pi n}{N}\right)$$ + +where $N$ = `size` when `periodic` is `true`, or $N$ = `size - 1` when `periodic` is `false`. +"# +)] +#[cfg_attr( + not(doc), + doc = "`w_n = 0.42 - 0.5 * cos(2πn / N) + 0.08 * cos(4πn / N)` where N = size (periodic) or N = size-1 (symmetric)" +)] +/// +/// # Arguments +/// - `size`: Size of the returned 1D window tensor. +/// - `periodic`: If `true`, the window is treated as periodic (i.e., `N = size`). +/// If `false`, the window is symmetric (i.e., `N = size - 1`). +/// - `options`: Controls the output device and optional dtype. Accepts: +/// - `&device` - uses the device's default float dtype +/// - `(&device, DType::F32)` - uses an explicit dtype +/// - `TensorCreationOptions` directly for full control. +/// +/// # Returns +/// - A 1D tensor of shape `[size]` containing the window. +/// +/// # Notes +/// - If `size == 0`, the function returns an empty tensor. +/// - If `size == 1`, the returned window contains a single value 1.0 which overrides the formula. +/// +/// # Panics +/// Panics if `size` exceeds `i64::MAX`. +/// +/// # Example +/// ```rust +/// use burn_tensor::{Device, DType, signal::blackman_window}; +/// +/// fn example() { +/// // Creating a window with default dtype +/// let device = Device::default(); +/// let window_tensor = blackman_window(5, true, &device); +/// // Output: [0.0, 0.20077015, 0.84922993, 0.8492298, 0.2007701] +/// +/// // Creating a window with explicit dtype. +/// // Note that this does not perform the computation at higher precision but it +/// // widens the storage of the returned tensor to F64. +/// let device = Device::default(); +/// let window_tensor_f64 = blackman_window(5, true, (&device, DType::F64)); +/// // Output: [0.0, 0.20077015, 0.84922993, 0.8492298, 0.2007701] +/// } +/// ``` +pub fn blackman_window( + size: usize, + periodic: bool, + options: impl Into, +) -> Tensor<1> { + let opt = options.into(); + let dtype = opt.resolve_dtype::(); + let shape = [size]; + check!(TensorCheck::creation_ops::<1>("BlackmanWindow", &shape)); + + if size == 0 { + return Tensor::<1>::empty(shape, opt).cast(dtype); + } + + if size == 1 { + return Tensor::<1>::ones(shape, opt).cast(dtype); + } + + let size_i64 = i64::try_from(size).expect("BlackmanWindow size doesn't fit in i64 range."); + let denominator = if periodic { size } else { size - 1 }; + let angular_increment = (2.0 * core::f64::consts::PI) / denominator as f64; + let cos_val = Tensor::<1, Int>::arange(0..size_i64, &opt.device) + .float() + .mul_scalar(angular_increment) + .cos(); + + // Using the double angle property of cosine: cos(2θ) = 2cos^2(θ) - 1 + // w[n] = 0.42 - 0.5cos(2πn / N) + 0.08cos(4πn / N) + // w[n] = 0.42 - 0.5cos(2πn / N) + 0.08cos(2 * (2πn / N)) + // w[n] = 0.42 - 0.5cos(2πn / N) + 0.08(2cos^2(2πn / N) - 1) + // w[n] = 0.34 - 0.5cos(2πn / N) + 0.16cos^2(2πn / N) + let first_cos_term = cos_val.clone().mul_scalar(-0.5); + let second_cos_term = cos_val.powi_scalar(2).mul_scalar(0.16); + first_cos_term + .add(second_cos_term) + .add_scalar(0.34) + .cast(dtype) +} diff --git a/crates/burn-tensor/src/tensor/signal/fft.rs b/crates/burn-tensor/src/tensor/signal/fft.rs new file mode 100644 index 0000000..a48d365 --- /dev/null +++ b/crates/burn-tensor/src/tensor/signal/fft.rs @@ -0,0 +1,280 @@ +use alloc::vec; +use burn_backend::ops::ModuleOps; +use burn_dispatch::Dispatch; + +use crate::Tensor; +use crate::check; +use crate::check::TensorCheck; +use crate::ops::BridgeTensor; + +/// Computes the 1-dimensional discrete Fourier Transform of real-valued input. +/// +/// Since the input is real, the Hermitian symmetry is exploited, and only the +/// first non-redundant values are returned ($N/2 + 1$). +/// For now, the autodiff is not yet supported +/// +#[cfg_attr( + doc, + doc = r#" +The mathematical formulation for each element $k$ in the frequency domain is: + +$$X\[k\] = \sum_{n=0}^{N-1} x\[n\] \left\[ \cos\left(\frac{2\pi kn}{N}\right) - i \sin\left(\frac{2\pi kn}{N}\right) \right\]$$ + +where $N$ is the size of the signal along the specified dimension. +"# +)] +#[cfg_attr(not(doc), doc = r"X\[k\] = Σ x\[n\] * exp(-i*2πkn/N)")] +/// +/// # Arguments +/// +/// * `signal` - The input tensor containing the real-valued signal. +/// * `dim` - The dimension along which to take the FFT. +/// * `n` - Optional FFT length. When `None`, the signal must be a power of two along `dim`. +/// When `Some(n)`, `n` must also be a power of two; the signal is truncated or zero-padded +/// to length `n`. Non-power-of-two `n` is rejected with a panic (true arbitrary-size DFT +/// support via Bluestein's algorithm is tracked as a follow-up). +/// +/// # Returns +/// +/// A tuple containing: +/// 1. The real part of the spectrum. Output length along `dim` is `n / 2 + 1` (using `n` or +/// `signal_len` respectively). +/// 2. The imaginary part of the spectrum (same shape). +/// +/// # Example +/// +/// ```rust +/// use burn_tensor::Tensor; +/// +/// fn example() { +/// let device = Default::default(); +/// let signal = Tensor::<1>::from_floats([1.0, 2.0, 3.0, 4.0], &device); +/// let (real, imag) = burn_tensor::signal::rfft(signal, 0, None); +/// } +/// ``` +pub fn rfft( + signal: Tensor, + dim: usize, + n: Option, +) -> (Tensor, Tensor) { + check!(TensorCheck::check_dim::(dim)); + + match n { + None => check!(TensorCheck::check_is_power_of_two::( + &signal.shape(), + dim + )), + Some(n) => { + assert!(n >= 1, "rfft: n must be >= 1, got {n}"); + assert!( + n.is_power_of_two(), + "rfft: n must be a power of two, got {n}. True non-power-of-two \ + DFT support is tracked as a follow-up (Bluestein's algorithm)." + ); + } + } + + let (re, im) = rfft_impl(signal.primitive, dim, n); + (Tensor::new(re), Tensor::new(im)) +} + +fn rfft_impl(signal: BridgeTensor, dim: usize, n: Option) -> (BridgeTensor, BridgeTensor) { + let (re, im) = Dispatch::rfft(signal.into_float(), dim, n); + (BridgeTensor::float(re), BridgeTensor::float(im)) +} + +/// Computes the 1-dimensional inverse discrete Fourier Transform for real-valued signals. +/// +/// This function reconstructs the real-valued time-domain signal from the +/// first non-redundant values ($N/2 + 1$) of the frequency-domain spectrum. +/// For now, the autodiff is not yet supported. +/// +#[cfg_attr( + doc, + doc = r#" +The mathematical formulation for each element $n$ in the time domain is: + +$$x\[n\] = \frac{1}{N} \sum_{k=0}^{N-1} X\[k\] \left\[ \cos\left(\frac{2\pi kn}{N}\right) + i \sin\left(\frac{2\pi kn}{N}\right) \right\]$$ + +where $N$ is the size of the reconstructed signal. +"# +)] +#[cfg_attr(not(doc), doc = r"x\[n\] = (1/N) * Σ X\[k\] * exp(i*2πkn/N)")] +/// +/// # Arguments +/// +/// * `spectrum_re` - The real part of the spectrum. +/// * `spectrum_im` - The imaginary part of the spectrum. +/// * `dim` - The dimension along which to take the inverse FFT. +/// * `n` - Optional output signal length. When `None`, the reconstructed signal length +/// `2 * (size - 1)` must be a power of two. When `Some(n)`, `n` must also be a power of +/// two and the output has exactly `n` samples. Non-power-of-two `n` is rejected. +/// +/// # Returns +/// +/// The reconstructed real-valued signal. +/// +/// # Example +/// +/// ```rust +/// use burn_tensor::Tensor; +/// +/// fn example() { +/// let device = Default::default(); +/// let real = Tensor::<1>::from_floats([10.0, -2.0, 2.0], &device); +/// let imag = Tensor::<1>::from_floats([0.0, 2.0, 0.0], &device); +/// let signal = burn_tensor::signal::irfft(real, imag, 0, None); +/// } +/// ``` +pub fn irfft( + spectrum_re: Tensor, + spectrum_im: Tensor, + dim: usize, + n: Option, +) -> Tensor { + check!(TensorCheck::check_dim::(dim)); + + if let Some(n) = n { + assert!(n >= 1, "irfft: n must be >= 1, got {n}"); + assert!( + n.is_power_of_two(), + "irfft: n must be a power of two, got {n}. True non-power-of-two \ + DFT support is tracked as a follow-up (Bluestein's algorithm)." + ); + } + + Tensor::new(irfft_impl( + spectrum_re.primitive, + spectrum_im.primitive, + dim, + n, + )) +} + +fn irfft_impl( + spectrum_re: BridgeTensor, + spectrum_im: BridgeTensor, + dim: usize, + n: Option, +) -> BridgeTensor { + BridgeTensor::float(Dispatch::irfft( + spectrum_re.into_float(), + spectrum_im.into_float(), + dim, + n, + )) +} + +/// Computes the 1-dimensional discrete Fourier Transform of complex-valued input. +/// +/// Internally calls [`rfft`] on the real and imaginary parts separately, +/// extends each half-spectrum to the full `N`-bin spectrum via Hermitian +/// symmetry. +/// +/// Autodiff is not yet supported. +/// +#[cfg_attr( + doc, + doc = r#" + +Due to the linearity of the Fourier Transform, a complex-valued signal $x\[n\] = x_{re}\[n\] + i x_{im}\[n\]$ can be transformed by applying the FFT to its real and imaginary parts separately: + +$$ \text{FFT}(x\[n\]) = \text{FFT}(x_{re}\[n\]) + i \text{FFT}(x_{im}\[n\]) $$ + +Since $x_{re}\[n\]$ and $x_{im}\[n\]$ are purely real, their transforms can be computed efficiently using the real FFT ([`rfft`]). The full spectrum is then reconstructed by exploiting Hermitian symmetry. +"# +)] +#[cfg_attr(not(doc), doc = r"X\[k\] = Σ x\[n\] * exp(-i*2πkn/N)")] +/// +/// # Arguments +/// +/// * `signal_re` - The real part of the complex input signal. +/// * `signal_im` - The imaginary part of the complex input signal. Must have the +/// same shape as `signal_re`. +/// * `dim` - The dimension along which to take the FFT. +/// * `n` - Optional FFT length. When `None`, the signal must be a power of two +/// along `dim`. When `Some(n)`, `n` must also be a power of two; the signal is +/// truncated or zero-padded to length `n`. +/// +/// # Returns +/// +/// A tuple `(re, im)` representing the full complex spectrum, each with `n` +/// elements along `dim`. +/// +/// # Example +/// +/// ```rust +/// use burn_tensor::Tensor; +/// +/// fn example() { +/// let device = Default::default(); +/// let re = Tensor::<1>::from_floats([1.0, 0.0, -1.0, 0.0], &device); +/// let im = Tensor::<1>::from_floats([0.0, 1.0, 0.0, -1.0], &device); +/// let (spec_re, spec_im) = burn_tensor::signal::cfft(re, im, 0, None); +/// } +/// ``` +pub fn cfft( + signal_re: Tensor, + signal_im: Tensor, + dim: usize, + n: Option, +) -> (Tensor, Tensor) { + assert!( + signal_re.shape() == signal_im.shape(), + "cfft: signal_re and signal_im must have the same shape, \ + got {:?} and {:?}", + signal_re.shape(), + signal_im.shape(), + ); + + check!(TensorCheck::check_dim::(dim)); + let fft_size = n.unwrap_or(signal_re.dims()[dim]); + + // rfft validates power-of-two and n constraints internally + let (xr, xi) = rfft(signal_re, dim, n); + let (yr, yi) = rfft(signal_im, dim, n); + + // Extend half-spectra (N/2+1 bins) to full N-bin spectra via Hermitian symmetry + let (xr, xi) = hermitian_extend(xr, xi, dim, fft_size); + let (yr, yi) = hermitian_extend(yr, yi, dim, fft_size); + + // FFT(z) = FFT(x) + i·FFT(y) + // = (Xr + i·Xi) + i·(Yr + i·Yi) + // = (Xr - Yi) + i·(Xi + Yr) + (xr - yi, xi + yr) +} + +/// Extend a half-spectrum from [`rfft`] (`N/2 + 1` bins) to the full `N`-bin +/// spectrum using Hermitian symmetry: `X[k] = conj(X[N-k])` for `k > N/2`. +pub(super) fn hermitian_extend( + half_re: Tensor, + half_im: Tensor, + dim: usize, + full_len: usize, +) -> (Tensor, Tensor) { + let half_len = half_re.dims()[dim]; // N/2 + 1 + + // For N <= 2, the half-spectrum already covers all bins + if full_len <= half_len { + return (half_re, half_im); + } + + // Mirror bins: reverse of bins 1..N/2-1 (skipping the Nyquist bin), + // with conjugated imaginary part. This produces X[N/2+1], X[N/2+2], ..., X[N-1] + let mirror_len = full_len - half_len; // N/2 - 1 + let mirror_re = half_re + .clone() + .narrow(dim, 1, mirror_len) + .flip([dim as isize]); + let mirror_im = half_im + .clone() + .narrow(dim, 1, mirror_len) + .flip([dim as isize]) + .neg(); + + // Full spectrum = [half_spectrum, conjugate_mirror] + let full_re = Tensor::cat(vec![half_re, mirror_re], dim); + let full_im = Tensor::cat(vec![half_im, mirror_im], dim); + + (full_re, full_im) +} diff --git a/crates/burn-tensor/src/tensor/signal/hamming_window.rs b/crates/burn-tensor/src/tensor/signal/hamming_window.rs new file mode 100644 index 0000000..6987d94 --- /dev/null +++ b/crates/burn-tensor/src/tensor/signal/hamming_window.rs @@ -0,0 +1,67 @@ +use crate::{Float, Int, Tensor, TensorCreationOptions, check, check::TensorCheck}; + +/// Creates a 1D Hamming window. +/// +#[cfg_attr( + doc, + doc = r#" +$$w_n = \alpha - \beta \cos\left(\frac{2\pi n}{N}\right)$$ + +where $\alpha = 25/46$, $\beta = 1 - \alpha$, and $N$ = `size` when `periodic` is `true`, or $N$ = `size - 1` when `periodic` is `false`. +"# +)] +#[cfg_attr( + not(doc), + doc = "`w_n = 25/46 - 21/46 * cos(2πn/N)` where N = size (periodic) or N = size-1 (symmetric)" +)] +/// +/// # Notes +/// +/// - `size == 0` returns an empty tensor. +/// - `size == 1` returns `[1.0]` regardless of `periodic`. +/// +/// # Example +/// +/// ```rust +/// use burn_tensor::Device; +/// use burn_tensor::signal::hamming_window; +/// +/// fn example() { +/// let device = Default::default(); +/// let window = hamming_window(8, true, &device); +/// println!("{window}"); +/// } +/// ``` +pub fn hamming_window( + size: usize, + periodic: bool, + options: impl Into, +) -> Tensor<1> { + let opt = options.into(); + let dtype = opt.resolve_dtype::(); + let shape = [size]; + check!(TensorCheck::creation_ops::<1>("HammingWindow", &shape)); + + if size == 0 { + return Tensor::<1>::empty(shape, opt).cast(dtype); + } + + if size == 1 { + return Tensor::<1>::ones(shape, opt).cast(dtype); + } + + let size_i64 = i64::try_from(size).expect("HammingWindow size doesn't fit in i64 range."); + let denominator = if periodic { size } else { size - 1 }; + let angular_increment = (2.0 * core::f64::consts::PI) / denominator as f64; + + let alpha = 25.0_f64 / 46.0_f64; + let beta = 1.0 - alpha; + + Tensor::<1, Int>::arange(0..size_i64, &opt.device) + .float() + .mul_scalar(angular_increment) + .cos() + .mul_scalar(-beta) + .add_scalar(alpha) + .cast(dtype) +} diff --git a/crates/burn-tensor/src/tensor/signal/hann_window.rs b/crates/burn-tensor/src/tensor/signal/hann_window.rs new file mode 100644 index 0000000..a6eac0b --- /dev/null +++ b/crates/burn-tensor/src/tensor/signal/hann_window.rs @@ -0,0 +1,64 @@ +use crate::{Float, Int, Tensor, TensorCreationOptions, check, check::TensorCheck}; + +/// Creates a 1D Hann window. +/// +#[cfg_attr( + doc, + doc = r#" +$$w_n = 0.5 - 0.5 \cos\left(\frac{2\pi n}{N}\right)$$ + +where $N$ = `size` when `periodic` is `true`, or $N$ = `size - 1` when `periodic` is `false`. +"# +)] +#[cfg_attr( + not(doc), + doc = "`w_n = 0.5 - 0.5 * cos(2πn/N)` where N = size (periodic) or N = size-1 (symmetric)" +)] +/// +/// # Notes +/// +/// - `size == 0` returns an empty tensor. +/// - `size == 1` returns `[1.0]` regardless of `periodic`. +/// +/// # Example +/// +/// ```rust +/// use burn_tensor::Device; +/// use burn_tensor::signal::hann_window; +/// +/// fn example() { +/// let device = Default::default(); +/// let window = hann_window(8, true, &device); +/// println!("{window}"); +/// } +/// ``` +pub fn hann_window( + size: usize, + periodic: bool, + options: impl Into, +) -> Tensor<1> { + let opt = options.into(); + let dtype = opt.resolve_dtype::(); + let shape = [size]; + check!(TensorCheck::creation_ops::<1>("HannWindow", &shape)); + + if size == 0 { + return Tensor::<1>::empty(shape, opt).cast(dtype); + } + + if size == 1 { + return Tensor::<1>::ones(shape, opt).cast(dtype); + } + + let size_i64 = i64::try_from(size).expect("HannWindow size doesn't fit in i64 range."); + let denominator = if periodic { size } else { size - 1 }; + let angular_increment = (2.0 * core::f64::consts::PI) / denominator as f64; + + Tensor::<1, Int>::arange(0..size_i64, &opt.device) + .float() + .mul_scalar(angular_increment) + .cos() + .mul_scalar(-0.5) + .add_scalar(0.5) + .cast(dtype) +} diff --git a/crates/burn-tensor/src/tensor/signal/mod.rs b/crates/burn-tensor/src/tensor/signal/mod.rs new file mode 100644 index 0000000..5cd8ac1 --- /dev/null +++ b/crates/burn-tensor/src/tensor/signal/mod.rs @@ -0,0 +1,11 @@ +mod blackman_window; +mod fft; +mod hamming_window; +mod hann_window; +mod stft; + +pub use blackman_window::*; +pub use fft::*; +pub use hamming_window::*; +pub use hann_window::*; +pub use stft::*; diff --git a/crates/burn-tensor/src/tensor/signal/stft.rs b/crates/burn-tensor/src/tensor/signal/stft.rs new file mode 100644 index 0000000..a29b67f --- /dev/null +++ b/crates/burn-tensor/src/tensor/signal/stft.rs @@ -0,0 +1,300 @@ +use alloc::vec; + +use crate::Tensor; +use crate::ops::PadMode; + +use super::{hermitian_extend, irfft, rfft}; + +/// Configuration shared by [`stft`] and [`istft`]. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct StftOptions { + /// Size of each FFT frame (must be >= 1). + pub n_fft: usize, + /// Stride between successive frames (must be >= 1 and <= effective window length + /// so overlap-add can reconstruct the signal). + pub hop_length: usize, + /// Window length. If `Some(w)`, the window is center-padded to `n_fft`. Defaults to `n_fft`. + pub win_length: Option, + /// If `true`, the signal is reflect-padded by `n_fft / 2` on both sides before framing. + pub center: bool, + /// If `true` (typical for real input), output has `n_fft/2 + 1` frequency bins; otherwise + /// the full `n_fft` bins are returned. + pub onesided: bool, +} + +impl StftOptions { + /// Construct default options for the given FFT size, matching PyTorch defaults + /// (`hop_length = n_fft / 4`, `win_length = None`, `center = true`, `onesided = true`). + pub fn new(n_fft: usize) -> Self { + Self { + n_fft, + hop_length: (n_fft / 4).max(1), + win_length: None, + center: true, + onesided: true, + } + } + + fn effective_win_length(&self) -> usize { + self.win_length.unwrap_or(self.n_fft) + } + + fn assert_valid(&self, op: &'static str) { + let n_fft = self.n_fft; + let hop_length = self.hop_length; + assert!(n_fft >= 1, "{op}: n_fft must be >= 1, got {n_fft}"); + assert!( + n_fft.is_power_of_two(), + "{op}: n_fft must be a power of two, got {n_fft}. True non-power-of-two \ + DFT support is tracked as a follow-up (Bluestein's algorithm)." + ); + assert!( + hop_length >= 1, + "{op}: hop_length must be >= 1, got {hop_length}" + ); + let win_len = self.effective_win_length(); + assert!( + win_len >= 1, + "{op}: effective win_length must be >= 1, got {win_len}" + ); + assert!( + win_len <= n_fft, + "{op}: win_length ({win_len}) must be <= n_fft ({n_fft})" + ); + assert!( + hop_length <= win_len, + "{op}: hop_length ({hop_length}) must be <= effective win_length ({win_len}) \ + so the window/hop combination satisfies the COLA/NOLA condition required \ + for invertibility" + ); + } +} + +impl Default for StftOptions { + fn default() -> Self { + Self::new(400) + } +} + +/// Computes the Short-Time Fourier Transform (STFT). +/// +/// Splits the signal into overlapping windowed frames and computes the FFT on each. +/// +/// # Returns +/// +/// A tensor of shape `[batch, n_frames, n_freqs, 2]` where the last dimension holds +/// `[real, imaginary]`. +pub fn stft(signal: Tensor<2>, window: Option>, options: StftOptions) -> Tensor<4> { + options.assert_valid("stft"); + let n_fft = options.n_fft; + let hop_length = options.hop_length; + let center = options.center; + let onesided = options.onesided; + let win_len = options.effective_win_length(); + + let device = signal.device(); + + let window = match window { + Some(w) => { + assert_eq!( + w.dims()[0], + win_len, + "stft: window length ({}) must match effective win_length ({win_len})", + w.dims()[0], + ); + w + } + None => Tensor::ones([win_len], &device), + }; + + let window = pad_window_to_n_fft(window, win_len, n_fft); + + let [_, raw_sig_len] = signal.dims(); + if center { + assert!( + raw_sig_len > n_fft / 2, + "stft: signal length ({raw_sig_len}) must be > n_fft/2 ({}) for reflect pad with center=true", + n_fft / 2 + ); + } + + let signal = if center { + let pad_amount = n_fft / 2; + signal.pad([(0, 0), (pad_amount, pad_amount)], PadMode::Reflect) + } else { + signal + }; + + let [batch, sig_len] = signal.dims(); + assert!( + sig_len >= n_fft, + "stft: signal length ({sig_len}) must be >= n_fft ({n_fft}) after padding" + ); + + let n_frames = 1 + (sig_len - n_fft) / hop_length; + + // Extract overlapping frames: [batch, n_frames, n_fft] + let frames: Tensor<3> = signal.unfold(1, n_fft, hop_length); + + // Apply window (broadcast over batch and n_frames) + let window_3d: Tensor<3> = window.reshape([1, 1, n_fft]); + let windowed = frames.mul(window_3d); + + // Flatten to [batch * n_frames, n_fft] for rfft + let flat: Tensor<2> = windowed.reshape([batch * n_frames, n_fft]); + + // rfft returns n_fft/2 + 1 bins along dim=1 (n_fft is pow2). + let (re, im) = rfft(flat, 1, Some(n_fft)); + + let (re, im, n_freqs) = if onesided { + (re, im, n_fft / 2 + 1) + } else { + let (re_full, im_full) = hermitian_extend(re, im, 1, n_fft); + (re_full, im_full, n_fft) + }; + + // Reshape to [batch, n_frames, n_freqs] then stack real/imag + let re: Tensor<3> = re.reshape([batch, n_frames, n_freqs]); + let im: Tensor<3> = im.reshape([batch, n_frames, n_freqs]); + + Tensor::stack::<4>(vec![re, im], 3) +} + +/// Center-pad window from `win_len` to `n_fft` with zeros. +fn pad_window_to_n_fft(window: Tensor<1>, win_len: usize, n_fft: usize) -> Tensor<1> { + if win_len < n_fft { + let pad_left = (n_fft - win_len) / 2; + let pad_right = n_fft - win_len - pad_left; + window.pad([(pad_left, pad_right)], PadMode::Constant(0.0)) + } else { + window + } +} + +/// Computes the inverse Short-Time Fourier Transform (ISTFT). +/// +/// Reconstructs a time-domain signal from its STFT representation using overlap-add. +/// +/// # Arguments +/// +/// * `stft_matrix` - Complex STFT tensor of shape `[batch, n_frames, n_freqs, 2]`. +/// * `window` - Window tensor used in the forward STFT. Defaults to rectangular. +/// * `length` - Optional output signal length. If `None`, the length is inferred. +/// * `options` - STFT configuration (must match the forward STFT). +/// +/// # Returns +/// +/// A real-valued tensor of shape `[batch, signal_length]`. +pub fn istft( + stft_matrix: Tensor<4>, + window: Option>, + length: Option, + options: StftOptions, +) -> Tensor<2> { + options.assert_valid("istft"); + let n_fft = options.n_fft; + let hop_length = options.hop_length; + let center = options.center; + let onesided = options.onesided; + let [batch, n_frames, n_freqs_in, two] = stft_matrix.dims(); + assert_eq!( + two, 2, + "istft: last dimension of stft_matrix must be 2 (real, imag), got {two}" + ); + assert!( + n_frames >= 1, + "istft: stft_matrix must contain at least one frame, got n_frames=0" + ); + let expected_n_freqs = if onesided { n_fft / 2 + 1 } else { n_fft }; + assert_eq!( + n_freqs_in, expected_n_freqs, + "istft: n_freqs dimension ({n_freqs_in}) does not match expected for \ + n_fft={n_fft}, onesided={onesided} (expected {expected_n_freqs})" + ); + + let win_len = options.effective_win_length(); + let device = stft_matrix.device(); + + let window = match window { + Some(w) => { + assert_eq!( + w.dims()[0], + win_len, + "istft: window length ({}) must match effective win_length ({win_len})", + w.dims()[0], + ); + w + } + None => Tensor::ones([win_len], &device), + }; + let window = pad_window_to_n_fft(window, win_len, n_fft); + + // Split real and imaginary: each [batch, n_frames, n_freqs] + let re: Tensor<3> = stft_matrix.clone().narrow(3, 0, 1).squeeze_dim(3); + let im: Tensor<3> = stft_matrix.narrow(3, 1, 1).squeeze_dim(3); + + let re: Tensor<2> = re.reshape([batch * n_frames, n_freqs_in]); + let im: Tensor<2> = im.reshape([batch * n_frames, n_freqs_in]); + + // Extract onesided spectrum for irfft + let (re_half, im_half) = if onesided { + (re, im) + } else { + let half = n_fft / 2 + 1; + (re.narrow(1, 0, half), im.narrow(1, 0, half)) + }; + + let frames = irfft(re_half, im_half, 1, Some(n_fft)); + + // Reshape to [batch, n_frames, n_fft] + let frames: Tensor<3> = frames.reshape([batch, n_frames, n_fft]); + + // Apply window to each frame + let window_3d: Tensor<3> = window.reshape([1, 1, n_fft]); + let windowed = frames.mul(window_3d.clone()); + + let expected_len = n_fft + (n_frames - 1) * hop_length; + let mut output = Tensor::<2>::zeros([batch, expected_len], &device); + let mut window_sum = Tensor::<2>::zeros([batch, expected_len], &device); + + let window_1d: Tensor<1> = window_3d.clone().reshape([n_fft]); + let window_sq_1d: Tensor<1> = window_1d.clone().mul(window_1d); + + for f in 0..n_frames { + let start = f * hop_length; + let right_pad = expected_len - n_fft - start; + let frame: Tensor<2> = windowed.clone().narrow(1, f, 1).squeeze_dim(1); + let frame_full = frame.pad([(0, 0), (start, right_pad)], PadMode::Constant(0.0)); + output = output.add(frame_full); + + let win_full: Tensor<1> = window_sq_1d + .clone() + .pad([(start, right_pad)], PadMode::Constant(0.0)); + let win_full: Tensor<2> = win_full.unsqueeze_dim(0); + window_sum = window_sum.add(win_full); + } + + let window_sum = window_sum.clamp_min(1e-10); + output = output.div(window_sum); + + // Remove center padding + let output = if center { + let pad_amount = n_fft / 2; + let trimmed_len = expected_len - 2 * pad_amount; + output.narrow(1, pad_amount, trimmed_len) + } else { + output + }; + + match length { + Some(len) => { + let current_len = output.dims()[1]; + if len <= current_len { + output.narrow(1, 0, len) + } else { + output.pad([(0, 0), (0, len - current_len)], PadMode::Constant(0.0)) + } + } + None => output, + } +} diff --git a/crates/burn-tensor/src/tensor/stats/mod.rs b/crates/burn-tensor/src/tensor/stats/mod.rs new file mode 100644 index 0000000..16f130e --- /dev/null +++ b/crates/burn-tensor/src/tensor/stats/mod.rs @@ -0,0 +1,69 @@ +use crate::{Int, Tensor}; + +pub fn var(tensor: Tensor, dim: usize) -> Tensor { + let mean = tensor.clone().mean_dim(dim); + var_with_mean(tensor, mean, dim) +} + +pub fn var_with_mean(tensor: Tensor, mean: Tensor, dim: usize) -> Tensor { + let n = tensor.shape()[dim] - 1; + var_with_mean_n(tensor, mean, dim, n) +} + +pub fn var_bias(tensor: Tensor, dim: usize) -> Tensor { + let mean = tensor.clone().mean_dim(dim); + var_with_mean_bias(tensor, mean, dim) +} + +pub fn var_with_mean_bias( + tensor: Tensor, + mean: Tensor, + dim: usize, +) -> Tensor { + let n = tensor.shape()[dim]; + var_with_mean_n(tensor, mean, dim, n) +} + +pub fn var_with_mean_n( + tensor: Tensor, + mean: Tensor, + dim: usize, + n: usize, +) -> Tensor { + tensor.sub(mean).square().sum_dim(dim).div_scalar(n as f32) +} + +pub fn median(tensor: Tensor, dim: usize) -> Tensor { + let total_elem_numbers = tensor.dims()[dim]; + let sorted_tensor = tensor.sort(dim); + + // Following the PyTorch behavior: + // - Odd count: the median + // - Even count: the lower of the two median elements + // + // Example: + // - 5 elements: (5 - 1) / 2 = 4 / 2 = 2 + // - 4 elements: (4 - 1) / 2 = 3 / 2 = 1 + let median_index = (total_elem_numbers - 1) / 2; + sorted_tensor.narrow(dim, median_index, 1) +} + +pub fn median_with_indices( + tensor: Tensor, + dim: usize, +) -> (Tensor, Tensor) { + let total_elem_numbers = tensor.dims()[dim]; + let (sorted_tensor, indices) = tensor.sort_with_indices(dim); + + // Following the PyTorch behavior: + // - Odd count: the median + // - Even count: the lower of the two median elements + // + // Example: + // - 5 elements: (5 - 1) / 2 = 4 / 2 = 2 + // - 4 elements: (4 - 1) / 2 = 3 / 2 = 1 + let median_index = (total_elem_numbers - 1) / 2; + let median_values = sorted_tensor.narrow(dim, median_index, 1); + let median_indices = indices.narrow(dim, median_index, 1); + (median_values, median_indices) +} diff --git a/crates/burn-train/Cargo.toml b/crates/burn-train/Cargo.toml new file mode 100644 index 0000000..ba1d9c2 --- /dev/null +++ b/crates/burn-train/Cargo.toml @@ -0,0 +1,72 @@ +[package] +authors = ["nathanielsimard "] +categories = ["science"] +description = "Training crate for the Burn framework" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "tensor", "pytorch", "ndarray"] +license.workspace = true +name = "burn-train" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-train" +documentation = "https://docs.rs/burn-train" +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["sys-metrics", "tui", "rl"] +doc = ["default"] +vision = ["burn-nn", "burn-store/pytorch", "burn-std/network", "dirs"] +tracing = ["burn-core/tracing", "burn-optim/tracing"] + + +sys-metrics = ["nvml-wrapper", "sysinfo", "systemstat"] +tui = ["ratatui", "unicode-width"] +rl = ["burn-rl"] + +[dependencies] +burn-core = { workspace = true, features = [ + "dataset", + "std", + "flex", # RL loop +] } +burn-optim = { workspace = true, features = ["std"] } +burn-rl = { workspace = true, optional = true } +burn-nn = { workspace = true, optional = true, features = ["std"] } +burn-store = { workspace = true, optional = true, features = ["std"] } +burn-std = { workspace = true, default-features = false, features = ["std"] } +dirs = { workspace = true, optional = true } + +log = { workspace = true } +tracing-subscriber = { workspace = true } +tracing-appender = { workspace = true } +tracing-core = { workspace = true } + +# System Metrics +nvml-wrapper = { workspace = true, optional = true } +sysinfo = { workspace = true, optional = true } +systemstat = { workspace = true, optional = true } + +# Text UI +ratatui = { workspace = true, optional = true, features = [ + "all-widgets", + "crossterm", +] } +unicode-width = { workspace = true, optional = true } + +# Utilities +derive-new = { workspace = true } +serde = { workspace = true, features = ["std", "derive"] } +async-channel = { workspace = true } +burn-flex = { workspace = true, features = ["default"] } +rstest.workspace = true +thiserror.workspace = true +rand.workspace = true + +[dev-dependencies] +tempfile.workspace = true + +[package.metadata.docs.rs] +features = ["doc"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/crates/burn-train/LICENSE-APACHE b/crates/burn-train/LICENSE-APACHE new file mode 120000 index 0000000..1cd601d --- /dev/null +++ b/crates/burn-train/LICENSE-APACHE @@ -0,0 +1 @@ +../../LICENSE-APACHE \ No newline at end of file diff --git a/crates/burn-train/LICENSE-MIT b/crates/burn-train/LICENSE-MIT new file mode 120000 index 0000000..b2cfbdc --- /dev/null +++ b/crates/burn-train/LICENSE-MIT @@ -0,0 +1 @@ +../../LICENSE-MIT \ No newline at end of file diff --git a/crates/burn-train/README.md b/crates/burn-train/README.md new file mode 100644 index 0000000..e693c68 --- /dev/null +++ b/crates/burn-train/README.md @@ -0,0 +1,6 @@ +# Burn Train + +This crate should be used with [burn](https://github.com/tracel-ai/burn). + +[![Current Crates.io Version](https://img.shields.io/crates/v/burn-train.svg)](https://crates.io/crates/burn-train) +[![license](https://shields.io/badge/license-MIT%2FApache--2.0-blue)](https://github.com/tracel-ai/burn-train/blob/master/README.md) diff --git a/crates/burn-train/src/checkpoint/async_checkpoint.rs b/crates/burn-train/src/checkpoint/async_checkpoint.rs new file mode 100644 index 0000000..cb3d1b5 --- /dev/null +++ b/crates/burn-train/src/checkpoint/async_checkpoint.rs @@ -0,0 +1,158 @@ +use super::{Checkpoint, Checkpointer, CheckpointerError}; +use crate::Interrupter; +use std::sync::mpsc; + +enum Message { + Restore( + usize, + mpsc::SyncSender>, + Option, + ), + Save(usize, R, Option), + Delete(usize, Option), + End, +} + +#[derive(new)] +struct CheckpointerThread { + checkpointer: C, + receiver: mpsc::Receiver>, +} + +impl CheckpointerThread +where + C: Checkpointer, + R: Checkpoint, +{ + fn run(self) { + for item in self.receiver.iter() { + match item { + Message::Restore(epoch, callback, interrupter) => { + let record = self.checkpointer.restore(epoch); + callback.send(record).unwrap_or_else(|err| { + interrupter.map_or_else( + || { + panic!( + "Error when sending response through callback channel: {err}" + ) + }, + |int| int.stop(Some(&err.to_string())), + ) + }); + } + Message::Save(epoch, state, interrupter) => { + self.checkpointer.save(epoch, state).unwrap_or_else(|err| { + interrupter.map_or_else( + || panic!("Error when saving the state: {err}"), + |int| int.stop(Some(&err.to_string())), + ) + }); + } + Message::Delete(epoch, interrupter) => { + self.checkpointer.delete(epoch).unwrap_or_else(|err| { + interrupter.map_or_else( + || panic!("Error when deleting the state: {err}"), + |int| int.stop(Some(&err.to_string())), + ) + }); + } + + Message::End => { + return; + } + }; + } + } +} + +/// Async checkpointer. +pub struct AsyncCheckpointer { + sender: mpsc::SyncSender>, + handler: Option>, + interrupter: Option, +} + +impl AsyncCheckpointer +where + R: Checkpoint, +{ + /// Create a new async checkpointer. + /// + /// # Arguments + /// + /// * `checkpointer` - The checkpointer. + /// + /// # Returns + /// + /// The async checkpointer. + pub fn new(checkpointer: C) -> Self + where + C: Checkpointer + Send + 'static, + { + // Only on checkpoint can be done in advance. + let (sender, receiver) = mpsc::sync_channel(0); + let thread = CheckpointerThread::new(checkpointer, receiver); + let handler = Some(std::thread::spawn(move || thread.run())); + + Self { + sender, + handler, + interrupter: None, + } + } + + /// Assign a handle used to interrupt training in case of checkpointing error. + pub fn with_interrupter(mut self, interrupter: Interrupter) -> Self { + self.interrupter = Some(interrupter); + self + } +} + +impl Checkpointer for AsyncCheckpointer +where + R: Checkpoint, +{ + fn save(&self, epoch: usize, record: R) -> Result<(), CheckpointerError> { + self.sender + .send(Message::Save(epoch, record, self.interrupter.clone())) + .expect("Can send message to checkpointer thread."); + + Ok(()) + } + + fn restore(&self, epoch: usize) -> Result { + let (sender, receiver) = mpsc::sync_channel(1); + self.sender + .send(Message::Restore(epoch, sender, self.interrupter.clone())) + .map_err(|e| CheckpointerError::Unknown(e.to_string()))?; + + if let Ok(record) = receiver.recv() { + return record; + }; + + Err(CheckpointerError::Unknown("Channel error.".to_string())) + } + + fn delete(&self, epoch: usize) -> Result<(), CheckpointerError> { + self.sender + .send(Message::Delete(epoch, self.interrupter.clone())) + .map_err(|e| CheckpointerError::Unknown(e.to_string()))?; + + Ok(()) + } +} + +impl Drop for AsyncCheckpointer { + fn drop(&mut self) { + self.sender + .send(Message::End) + .expect("Can send the end message to the checkpointer thread."); + let handler = self.handler.take(); + + if let Some(handler) = handler { + handler + .join() + .expect("The checkpointer thread should stop."); + } + } +} diff --git a/crates/burn-train/src/checkpoint/base.rs b/crates/burn-train/src/checkpoint/base.rs new file mode 100644 index 0000000..40d3d3c --- /dev/null +++ b/crates/burn-train/src/checkpoint/base.rs @@ -0,0 +1,123 @@ +use burn_core::store::{ModuleRecord, RecordError}; +use burn_optim::OptimizerRecord; +use burn_optim::lr_scheduler::LrSchedulerRecord; +use burn_std::Bytes; +use std::path::PathBuf; +use thiserror::Error; + +/// The error type for checkpointer. +#[derive(Error, Debug)] +pub enum CheckpointerError { + /// IO error. + #[error("I/O Error: `{0}`")] + IOError(std::io::Error), + + /// Record (burnpack) error. + #[error("Record error: `{0}`")] + Record(RecordError), + + /// Other errors. + #[error("Unknown error: `{0}`")] + Unknown(String), +} + +/// A record that can be saved to and loaded from a burnpack file. +/// +/// Implemented for the burnpack record types used during training: the module +/// ([`ModuleRecord`]), the optimizer ([`OptimizerRecord`]) and the learning rate scheduler +/// ([`LrSchedulerRecord`]). +/// +/// Records are device-free: a checkpoint is just file-backed bytes. Device placement is decided +/// when a record is applied (the module keeps its existing parameter device; optimizer state +/// migrates to each parameter's device on the next step), not when the checkpoint is loaded. +pub trait Checkpoint: Sized + Send + 'static { + /// Save the record to `path`. + fn save(self, path: PathBuf) -> Result<(), CheckpointerError>; + /// Load the record from `path`. + fn load(path: PathBuf) -> Result; + /// Creates a checkpoint from bytes + fn checkpoint_from_bytes(bytes: Bytes) -> Result; + /// Transforms a checkpoint into bytes + fn checkpoint_into_bytes(self) -> Result; +} + +/// A stateless record: nothing to save or load. +impl Checkpoint for () { + fn save(self, _path: PathBuf) -> Result<(), CheckpointerError> { + Ok(()) + } + fn load(_path: PathBuf) -> Result { + Ok(()) + } + fn checkpoint_from_bytes(_bytes: Bytes) -> Result { + Ok(()) + } + fn checkpoint_into_bytes(self) -> Result { + Ok(Bytes::from_bytes_vec(vec![0])) + } +} + +impl Checkpoint for ModuleRecord { + fn save(self, path: PathBuf) -> Result<(), CheckpointerError> { + ModuleRecord::save(self, path).map_err(CheckpointerError::Record) + } + fn load(path: PathBuf) -> Result { + ModuleRecord::load(path).map_err(CheckpointerError::Record) + } + fn checkpoint_into_bytes(self) -> Result { + self.into_bytes() + } + fn checkpoint_from_bytes(bytes: Bytes) -> Result { + ModuleRecord::from_bytes(bytes) + } +} + +impl Checkpoint for OptimizerRecord { + fn save(self, path: PathBuf) -> Result<(), CheckpointerError> { + OptimizerRecord::save(self, path).map_err(CheckpointerError::Record) + } + fn load(path: PathBuf) -> Result { + OptimizerRecord::load(path).map_err(CheckpointerError::Record) + } + fn checkpoint_from_bytes(bytes: Bytes) -> Result { + OptimizerRecord::from_bytes(bytes) + } + fn checkpoint_into_bytes(self) -> Result { + self.into_bytes() + } +} + +impl Checkpoint for LrSchedulerRecord { + fn save(self, path: PathBuf) -> Result<(), CheckpointerError> { + LrSchedulerRecord::save(self, path).map_err(CheckpointerError::Record) + } + fn load(path: PathBuf) -> Result { + LrSchedulerRecord::load(path).map_err(CheckpointerError::Record) + } + fn checkpoint_from_bytes(bytes: Bytes) -> Result { + LrSchedulerRecord::from_bytes(bytes) + } + fn checkpoint_into_bytes(self) -> Result { + self.into_bytes() + } +} + +/// The trait for checkpointer. +pub trait Checkpointer: Send + Sync +where + R: Checkpoint, +{ + /// Save the record. + /// + /// # Arguments + /// + /// * `epoch` - The epoch. + /// * `record` - The record. + fn save(&self, epoch: usize, record: R) -> Result<(), CheckpointerError>; + + /// Delete the checkpoint saved at the given epoch if present. + fn delete(&self, epoch: usize) -> Result<(), CheckpointerError>; + + /// Restore the record from the checkpoint saved at the given epoch. + fn restore(&self, epoch: usize) -> Result; +} diff --git a/crates/burn-train/src/checkpoint/file.rs b/crates/burn-train/src/checkpoint/file.rs new file mode 100644 index 0000000..da52ac2 --- /dev/null +++ b/crates/burn-train/src/checkpoint/file.rs @@ -0,0 +1,67 @@ +use std::path::{Path, PathBuf}; + +use super::{Checkpoint, Checkpointer, CheckpointerError}; + +/// The file checkpointer. +/// +/// Saves each record as a [burnpack](burn_core::store) file in the given directory. +pub struct FileCheckpointer { + directory: PathBuf, + name: String, +} + +impl FileCheckpointer { + /// Creates a new file checkpointer. + /// + /// # Arguments + /// + /// * `directory` - The directory to save the checkpoints. + /// * `name` - The name of the checkpoint. + pub fn new(directory: impl AsRef, name: &str) -> Self { + let directory = directory.as_ref(); + std::fs::create_dir_all(directory).ok(); + + Self { + directory: directory.to_path_buf(), + name: name.to_string(), + } + } + + fn path_for_epoch(&self, epoch: usize) -> PathBuf { + self.directory.join(format!("{}-{}.bpk", self.name, epoch)) + } +} + +impl Checkpointer for FileCheckpointer +where + R: Checkpoint, +{ + fn save(&self, epoch: usize, record: R) -> Result<(), CheckpointerError> { + let file_path = self.path_for_epoch(epoch); + log::trace!("Saving checkpoint {} to {}", epoch, file_path.display()); + + record.save(file_path) + } + + fn restore(&self, epoch: usize) -> Result { + let file_path = self.path_for_epoch(epoch); + log::info!( + "Restoring checkpoint {} from {}", + epoch, + file_path.display() + ); + + R::load(file_path) + } + + fn delete(&self, epoch: usize) -> Result<(), CheckpointerError> { + let file_path = self.path_for_epoch(epoch); + + if file_path.exists() { + log::trace!("Removing checkpoint {}", file_path.display()); + std::fs::remove_file(file_path).map_err(CheckpointerError::IOError)?; + } + + Ok(()) + } +} diff --git a/crates/burn-train/src/checkpoint/mod.rs b/crates/burn-train/src/checkpoint/mod.rs new file mode 100644 index 0000000..14b8f15 --- /dev/null +++ b/crates/burn-train/src/checkpoint/mod.rs @@ -0,0 +1,9 @@ +mod async_checkpoint; +mod base; +mod file; +mod strategy; + +pub use async_checkpoint::*; +pub use base::*; +pub use file::*; +pub use strategy::*; diff --git a/crates/burn-train/src/checkpoint/strategy/base.rs b/crates/burn-train/src/checkpoint/strategy/base.rs new file mode 100644 index 0000000..07d6213 --- /dev/null +++ b/crates/burn-train/src/checkpoint/strategy/base.rs @@ -0,0 +1,34 @@ +use std::ops::DerefMut; + +use crate::metric::store::EventStoreClient; + +/// Action to be taken by a [checkpointer](crate::checkpoint::Checkpointer). +#[derive(Clone, PartialEq, Debug)] +pub enum CheckpointingAction { + /// Delete the given epoch. + Delete(usize), + /// Save the current record. + Save, +} + +/// Define when checkpoint should be saved and deleted. +pub trait CheckpointingStrategy: Send { + /// Based on the epoch, determine if the checkpoint should be saved. + fn checkpointing( + &mut self, + epoch: usize, + collector: &EventStoreClient, + ) -> Vec; +} + +// We make dyn box implement the checkpointing strategy so that it can be used with generic, but +// still be dynamic. +impl CheckpointingStrategy for Box { + fn checkpointing( + &mut self, + epoch: usize, + collector: &EventStoreClient, + ) -> Vec { + self.deref_mut().checkpointing(epoch, collector) + } +} diff --git a/crates/burn-train/src/checkpoint/strategy/composed.rs b/crates/burn-train/src/checkpoint/strategy/composed.rs new file mode 100644 index 0000000..8029c9e --- /dev/null +++ b/crates/burn-train/src/checkpoint/strategy/composed.rs @@ -0,0 +1,146 @@ +use crate::metric::store::EventStoreClient; + +use super::{CheckpointingAction, CheckpointingStrategy}; +use std::collections::HashSet; + +/// Compose multiple checkpointing strategy and only delete checkpoints when both strategy flag an +/// epoch to be deleted. +pub struct ComposedCheckpointingStrategy { + strategies: Vec>, + deleted: Vec>, +} + +/// Help building a [checkpointing strategy](CheckpointingStrategy) by combining multiple ones. +#[derive(Default)] +pub struct ComposedCheckpointingStrategyBuilder { + strategies: Vec>, +} + +impl ComposedCheckpointingStrategyBuilder { + /// Add a new [checkpointing strategy](CheckpointingStrategy). + #[allow(clippy::should_implement_trait)] + pub fn add(mut self, strategy: S) -> Self + where + S: CheckpointingStrategy + 'static, + { + self.strategies.push(Box::new(strategy)); + self + } + + /// Create a new [composed checkpointing strategy](ComposedCheckpointingStrategy). + pub fn build(self) -> ComposedCheckpointingStrategy { + ComposedCheckpointingStrategy::new(self.strategies) + } +} + +impl ComposedCheckpointingStrategy { + fn new(strategies: Vec>) -> Self { + Self { + deleted: strategies.iter().map(|_| HashSet::new()).collect(), + strategies, + } + } + /// Create a new builder which help compose multiple + /// [checkpointing strategies](CheckpointingStrategy). + pub fn builder() -> ComposedCheckpointingStrategyBuilder { + ComposedCheckpointingStrategyBuilder::default() + } +} + +impl CheckpointingStrategy for ComposedCheckpointingStrategy { + fn checkpointing( + &mut self, + epoch: usize, + collector: &EventStoreClient, + ) -> Vec { + let mut saved = false; + let mut actions = Vec::new(); + let mut epochs_to_check = Vec::new(); + + for (i, strategy) in self.strategies.iter_mut().enumerate() { + let actions = strategy.checkpointing(epoch, collector); + // We assume that the strategy would not want the current epoch to be saved. + // So we flag it as deleted. + if actions.is_empty() { + self.deleted + .get_mut(i) + .expect("As many 'deleted' as 'strategies'.") + .insert(epoch); + } + + for action in actions { + match action { + CheckpointingAction::Delete(epoch) => { + self.deleted + .get_mut(i) + .expect("As many 'deleted' as 'strategies'.") + .insert(epoch); + epochs_to_check.push(epoch); + } + CheckpointingAction::Save => saved = true, + } + } + } + + if saved { + actions.push(CheckpointingAction::Save); + } + + for epoch in epochs_to_check.into_iter() { + let mut num_true = 0; + for i in 0..self.strategies.len() { + if self + .deleted + .get(i) + .expect("Ad many 'deleted' as 'strategies'.") + .contains(&epoch) + { + num_true += 1; + } + } + + if num_true == self.strategies.len() { + actions.push(CheckpointingAction::Delete(epoch)); + + for i in 0..self.strategies.len() { + self.deleted + .get_mut(i) + .expect("As many 'deleted' as 'strategies'.") + .remove(&epoch); + } + } + } + + actions + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{checkpoint::KeepLastNCheckpoints, metric::store::LogEventStore}; + + #[test] + fn should_delete_when_both_deletes() { + let store = EventStoreClient::new(LogEventStore::default()); + let mut strategy = ComposedCheckpointingStrategy::builder() + .add(KeepLastNCheckpoints::new(1)) + .add(KeepLastNCheckpoints::new(2)) + .build(); + + assert_eq!( + vec![CheckpointingAction::Save], + strategy.checkpointing(1, &store) + ); + + assert_eq!( + vec![CheckpointingAction::Save], + strategy.checkpointing(2, &store) + ); + + assert_eq!( + vec![CheckpointingAction::Save, CheckpointingAction::Delete(1)], + strategy.checkpointing(3, &store) + ); + } +} diff --git a/crates/burn-train/src/checkpoint/strategy/lastn.rs b/crates/burn-train/src/checkpoint/strategy/lastn.rs new file mode 100644 index 0000000..6e8b46a --- /dev/null +++ b/crates/burn-train/src/checkpoint/strategy/lastn.rs @@ -0,0 +1,56 @@ +use super::CheckpointingStrategy; +use crate::{checkpoint::CheckpointingAction, metric::store::EventStoreClient}; + +/// Keep the last N checkpoints. +/// +/// Very useful when training, minimizing disk space while ensuring that the training can be +/// resumed even if something goes wrong. +#[derive(new)] +pub struct KeepLastNCheckpoints { + num_keep: usize, +} + +impl CheckpointingStrategy for KeepLastNCheckpoints { + fn checkpointing( + &mut self, + epoch: usize, + _store: &EventStoreClient, + ) -> Vec { + let mut actions = vec![CheckpointingAction::Save]; + + if let Some(epoch) = usize::checked_sub(epoch, self.num_keep) + && epoch > 0 + { + actions.push(CheckpointingAction::Delete(epoch)); + } + + actions + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::metric::store::LogEventStore; + + #[test] + fn should_always_delete_lastn_epoch_if_higher_than_one() { + let mut strategy = KeepLastNCheckpoints::new(2); + let store = EventStoreClient::new(LogEventStore::default()); + + assert_eq!( + vec![CheckpointingAction::Save], + strategy.checkpointing(1, &store) + ); + + assert_eq!( + vec![CheckpointingAction::Save], + strategy.checkpointing(2, &store) + ); + + assert_eq!( + vec![CheckpointingAction::Save, CheckpointingAction::Delete(1)], + strategy.checkpointing(3, &store) + ); + } +} diff --git a/crates/burn-train/src/checkpoint/strategy/metric.rs b/crates/burn-train/src/checkpoint/strategy/metric.rs new file mode 100644 index 0000000..8769f83 --- /dev/null +++ b/crates/burn-train/src/checkpoint/strategy/metric.rs @@ -0,0 +1,143 @@ +use super::CheckpointingStrategy; +use crate::{ + checkpoint::CheckpointingAction, + metric::{ + Metric, MetricName, + store::{Aggregate, Direction, EventStoreClient, Split}, + }, +}; + +/// Keep the best checkpoint based on a metric. +pub struct MetricCheckpointingStrategy { + current: Option, + aggregate: Aggregate, + direction: Direction, + split: Split, + name: MetricName, +} + +impl MetricCheckpointingStrategy { + /// Create a new metric checkpointing strategy. + pub fn new(metric: &M, aggregate: Aggregate, direction: Direction, split: Split) -> Self + where + M: Metric, + { + Self { + current: None, + name: metric.name(), + aggregate, + direction, + split, + } + } +} + +impl CheckpointingStrategy for MetricCheckpointingStrategy { + fn checkpointing( + &mut self, + epoch: usize, + store: &EventStoreClient, + ) -> Vec { + let best_epoch = + match store.find_epoch(&self.name, self.aggregate, self.direction, &self.split) { + Some(epoch_best) => epoch_best, + None => epoch, + }; + + let mut actions = Vec::new(); + + if let Some(current) = self.current + && current != best_epoch + { + actions.push(CheckpointingAction::Delete(current)); + } + + if best_epoch == epoch { + actions.push(CheckpointingAction::Save); + } + + self.current = Some(best_epoch); + + actions + } +} + +#[cfg(test)] +mod tests { + use crate::{ + EventProcessorTraining, + logger::InMemoryMetricLogger, + metric::{ + LossMetric, + processor::{ + MetricsTraining, MinimalEventProcessor, + test_utils::{end_epoch, process_train}, + }, + store::LogEventStore, + }, + test_utils::start_epoch, + }; + + use super::*; + use std::sync::Arc; + + #[test] + fn always_keep_the_best_epoch() { + let loss = LossMetric::new(); + let mut store = LogEventStore::default(); + let mut strategy = MetricCheckpointingStrategy::new( + &loss, + Aggregate::Mean, + Direction::Lowest, + Split::Train, + ); + let mut metrics = MetricsTraining::::default(); + // Register an in memory logger. + store.register_logger(InMemoryMetricLogger::default()); + // Register the loss metric. + metrics.register_train_metric_numeric(loss); + let store = Arc::new(EventStoreClient::new(store)); + let mut processor = MinimalEventProcessor::new(metrics, store.clone()); + processor.process_train(crate::LearnerEvent::Start { + total_epochs: 0, + starting_epoch: 0, + }); + + // Two points for the first epoch. Mean 0.75 + let mut epoch = 1; + start_epoch(&mut processor, epoch, 2); + process_train(&mut processor, 1.0, epoch); + process_train(&mut processor, 0.5, epoch); + end_epoch(&mut processor, epoch); + + // Should save the current record. + assert_eq!( + vec![CheckpointingAction::Save], + strategy.checkpointing(epoch, &store) + ); + + // Two points for the second epoch. Mean 0.4 + epoch += 1; + start_epoch(&mut processor, epoch, 2); + process_train(&mut processor, 0.5, epoch); + process_train(&mut processor, 0.3, epoch); + end_epoch(&mut processor, epoch); + + // Should save the current record and delete the previous one. + assert_eq!( + vec![CheckpointingAction::Delete(1), CheckpointingAction::Save], + strategy.checkpointing(epoch, &store) + ); + + // Two points for the last epoch. Mean 2.0 + epoch += 1; + start_epoch(&mut processor, epoch, 2); + process_train(&mut processor, 1.0, epoch); + process_train(&mut processor, 3.0, epoch); + end_epoch(&mut processor, epoch); + + // Should not delete the previous record, since it's the best one, and should not save a + // new one. + assert!(strategy.checkpointing(epoch, &store).is_empty()); + } +} diff --git a/crates/burn-train/src/checkpoint/strategy/mod.rs b/crates/burn-train/src/checkpoint/strategy/mod.rs new file mode 100644 index 0000000..2915a71 --- /dev/null +++ b/crates/burn-train/src/checkpoint/strategy/mod.rs @@ -0,0 +1,9 @@ +mod base; +mod composed; +mod lastn; +mod metric; + +pub use base::*; +pub use composed::*; +pub use lastn::*; +pub use metric::*; diff --git a/crates/burn-train/src/components.rs b/crates/burn-train/src/components.rs new file mode 100644 index 0000000..5c49a34 --- /dev/null +++ b/crates/burn-train/src/components.rs @@ -0,0 +1,22 @@ +use crate::{InferenceStep, TrainStep}; +use burn_core::module::AutodiffModule; + +/// A single keyword trait for a Burn [module](burn_core::module::Module) used for learning. +pub trait LearnerModel: + TrainStep + InferenceStep + AutodiffModule + core::fmt::Display + 'static +{ +} + +impl LearnerModel for T where + T: TrainStep + InferenceStep + AutodiffModule + core::fmt::Display + 'static +{ +} + +/// Type for training input. +pub(crate) type TrainingModelInput = ::Input; +/// Type for inference input. +pub(crate) type InferenceModelInput = ::Input; +/// Type for training output. +pub(crate) type TrainingModelOutput = ::Output; +/// Type for inference output. +pub(crate) type InferenceModelOutput = ::Output; diff --git a/crates/burn-train/src/evaluator/base.rs b/crates/burn-train/src/evaluator/base.rs new file mode 100644 index 0000000..73476a3 --- /dev/null +++ b/crates/burn-train/src/evaluator/base.rs @@ -0,0 +1,93 @@ +use crate::{ + AsyncProcessorEvaluation, EvaluationItem, FullEventProcessorEvaluation, InferenceStep, + Interrupter, LearnerSummaryConfig, + evaluator::components::EvaluatorComponentTypes, + metric::processor::{EvaluatorEvent, EventProcessorEvaluation}, + renderer::{EvaluationName, MetricsRenderer}, +}; +use burn_core::{data::dataloader::DataLoader, module::Module}; +use std::sync::Arc; + +pub(crate) type TestInput = <::Model as InferenceStep>::Input; +pub(crate) type TestOutput = <::Model as InferenceStep>::Output; + +pub(crate) type TestLoader = Arc>>; + +/// Evaluates a model on a specific dataset. +pub struct Evaluator { + pub(crate) model: EC::Model, + pub(crate) interrupter: Interrupter, + pub(crate) event_processor: + AsyncProcessorEvaluation>>, + /// Config for creating a summary of the evaluation + pub summary: Option, +} + +impl Evaluator { + /// Run the evaluation on the given dataset. + /// + /// The data will be stored and displayed under the provided name. + pub fn eval( + self, + name: S, + dataloader: TestLoader, + ) -> Box { + self.eval_all([(name, dataloader)]) + } + + /// Run the evaluation on multiple named datasets sequentially. + /// + /// Prefer this over calling [`eval`](Self::eval) in a loop — the progress logger + /// receives the correct `total_tests` count and `end_test` is called between splits. + pub fn eval_all( + mut self, + splits: impl IntoIterator)>, + ) -> Box { + let splits: Vec<_> = splits.into_iter().collect(); + let total_tests = splits.len(); + + self.event_processor + .process_test(EvaluatorEvent::Start { total_tests }); + + for (name, dataloader) in splits { + let dataloader = dataloader.to_device(self.model.devices().first().unwrap()); + let name = EvaluationName::new(name); + let total_items = dataloader.num_items(); + let mut iterator = dataloader.iter(); + let mut iteration = 0; + + self.event_processor + .process_test(EvaluatorEvent::StartTest(name.clone(), total_items)); + + while let Some(item) = iterator.next() { + let progress = iterator.progress(); + iteration += 1; + + let item = self.model.step(item); + let item = EvaluationItem::new(item, progress, Some(iteration)); + + self.event_processor + .process_test(EvaluatorEvent::ProcessedItem(name.clone(), item)); + + if self.interrupter.should_stop() { + log::info!("Testing interrupted."); + break; + } + } + + self.event_processor.process_test(EvaluatorEvent::EndTest); + } + + let summary = self.summary.and_then(|summary| { + summary + .init() + .map(|summary| summary.with_model(self.model.to_string())) + .ok() + }); + + self.event_processor + .process_test(EvaluatorEvent::End(summary)); + + self.event_processor.renderer() + } +} diff --git a/crates/burn-train/src/evaluator/builder.rs b/crates/burn-train/src/evaluator/builder.rs new file mode 100644 index 0000000..1374f88 --- /dev/null +++ b/crates/burn-train/src/evaluator/builder.rs @@ -0,0 +1,231 @@ +use crate::{ + ApplicationLoggerInstaller, Evaluator, FileApplicationLoggerInstaller, InferenceStep, + Interrupter, LearnerSummaryConfig, TestOutput, + evaluator::components::{EvaluatorComponentTypes, EvaluatorComponentTypesMarker}, + logger::{EvaluationProgressLogger, FileMetricLogger}, + metric::{ + Adaptor, Metric, Numeric, + processor::{AsyncProcessorEvaluation, FullEventProcessorEvaluation, MetricsEvaluation}, + store::{EventStoreClient, LogEventStore}, + }, + renderer::{MetricsRenderer, default_renderer}, +}; +use burn_core::module::Module; +use std::{ + collections::BTreeSet, + path::{Path, PathBuf}, + sync::Arc, +}; + +/// Struct to configure and create an [evaluator](Evaluator). +/// +/// The generics components of the builder should probably not be set manually, as they are +/// optimized for Rust type inference. +pub struct EvaluatorBuilder { + tracing_logger: Option>, + event_store: LogEventStore, + summary_metrics: BTreeSet, + renderer: Option>, + interrupter: Interrupter, + metrics: MetricsEvaluation>, + directory: PathBuf, + summary: bool, + progress_logger: Option>, +} + +impl EvaluatorBuilder> +where + M: Module + InferenceStep + core::fmt::Display + 'static, +{ + /// Creates a new evaluator builder. + /// + /// # Arguments + /// + /// * `directory` - The directory to save the checkpoints. + pub fn new(directory: impl AsRef) -> Self { + let directory = directory.as_ref().to_path_buf(); + let log_file = directory.join("evaluation.log"); + + Self { + tracing_logger: Some(Box::new(FileApplicationLoggerInstaller::new(log_file))), + event_store: LogEventStore::default(), + summary_metrics: Default::default(), + renderer: None, + interrupter: Interrupter::new(), + summary: false, + metrics: MetricsEvaluation::default(), + directory, + progress_logger: None, + } + } +} + +impl EvaluatorBuilder { + /// Registers [numeric](crate::metric::Numeric) test [metrics](Metric). + pub fn metrics>(self, metrics: Me) -> Self { + metrics.register(self) + } + + /// Registers text [metrics](Metric). + pub fn metrics_text>(self, metrics: Me) -> Self { + metrics.register(self) + } + + /// By default, Rust logs are captured and written into + /// `evaluation.log`. If disabled, standard Rust log handling + /// will apply. + pub fn with_application_logger( + mut self, + logger: Option>, + ) -> Self { + self.tracing_logger = logger; + self + } + + /// Register a [numeric](crate::metric::Numeric) test [metric](Metric). + pub fn metric_numeric(mut self, metric: Me) -> Self + where + Me: Metric + Numeric + 'static, + TestOutput: Adaptor, + { + self.summary_metrics.insert(metric.name().to_string()); + self.metrics.register_test_metric_numeric(metric); + self + } + + /// Register a text test [metric](Metric). + pub fn metric(mut self, metric: Me) -> Self + where + Me: Metric + 'static, + TestOutput: Adaptor, + { + self.summary_metrics.insert(metric.name().to_string()); + self.metrics.register_test_metric(metric); + self + } + + /// Replace the default CLI renderer with a custom one. + /// + /// # Arguments + /// + /// * `renderer` - The custom renderer. + pub fn renderer(mut self, renderer: Box) -> Self { + self.renderer = Some(renderer); + self + } + + /// Enable the evaluation summary report. + /// + /// The summary will be displayed at the end of `.eval()`. + pub fn summary(mut self) -> Self { + self.summary = true; + self + } + + /// Register a progress logger to track and store evaluation progress. + /// + /// # Example + /// ```ignore + /// // `MyEvaluationProgressLogger` is a user-defined type that implements + /// // `burn_train::logger::EvaluationProgressLogger`. + /// let evaluator = EvaluatorBuilder::new(...) + /// .with_progress_logger(MyEvaluationProgressLogger); + /// ``` + pub fn with_progress_logger(mut self, logger: PL) -> Self + where + PL: EvaluationProgressLogger + 'static, + { + self.progress_logger = Some(Box::new(logger)); + self + } + + /// Builds the evaluator. + #[allow(clippy::type_complexity)] + pub fn build(mut self, model: EC::Model) -> Evaluator { + let renderer = self + .renderer + .unwrap_or_else(|| default_renderer(self.interrupter.clone(), None)); + + self.event_store + .register_logger(FileMetricLogger::new_eval(self.directory.clone())); + let event_store = Arc::new(EventStoreClient::new(self.event_store)); + + let full_processor = FullEventProcessorEvaluation::new(self.metrics, renderer, event_store); + let full_processor = match self.progress_logger { + Some(logger) => full_processor.with_progress_logger(logger), + None => full_processor, + }; + let event_processor = AsyncProcessorEvaluation::new(full_processor); + + let summary = if self.summary { + Some(LearnerSummaryConfig { + directory: self.directory, + metrics: self.summary_metrics.into_iter().collect::>(), + }) + } else { + None + }; + + Evaluator { + model, + interrupter: self.interrupter, + event_processor, + summary, + } + } +} + +/// Trait to fake variadic generics. +pub trait EvalMetricRegistration: Sized { + /// Register the metrics. + fn register(self, builder: EvaluatorBuilder) -> EvaluatorBuilder; +} + +/// Trait to fake variadic generics. +pub trait EvalTextMetricRegistration: Sized { + /// Register the metrics. + fn register(self, builder: EvaluatorBuilder) -> EvaluatorBuilder; +} + +macro_rules! gen_tuple { + ($($M:ident),*) => { + impl<$($M,)* EC: EvaluatorComponentTypes> EvalTextMetricRegistration for ($($M,)*) + where + $(TestOutput: Adaptor<$M::Input>,)* + $($M: Metric + 'static,)* + { + #[allow(non_snake_case)] + fn register( + self, + builder: EvaluatorBuilder, + ) -> EvaluatorBuilder { + let ($($M,)*) = self; + $(let builder = builder.metric($M);)* + builder + } + } + + impl<$($M,)* EC: EvaluatorComponentTypes> EvalMetricRegistration for ($($M,)*) + where + $(TestOutput: Adaptor<$M::Input>,)* + $($M: Metric + $crate::metric::Numeric + 'static,)* + { + #[allow(non_snake_case)] + fn register( + self, + builder: EvaluatorBuilder, + ) -> EvaluatorBuilder { + let ($($M,)*) = self; + $(let builder = builder.metric_numeric($M);)* + builder + } + } + }; +} + +gen_tuple!(M1); +gen_tuple!(M1, M2); +gen_tuple!(M1, M2, M3); +gen_tuple!(M1, M2, M3, M4); +gen_tuple!(M1, M2, M3, M4, M5); +gen_tuple!(M1, M2, M3, M4, M5, M6); diff --git a/crates/burn-train/src/evaluator/components.rs b/crates/burn-train/src/evaluator/components.rs new file mode 100644 index 0000000..624a2f1 --- /dev/null +++ b/crates/burn-train/src/evaluator/components.rs @@ -0,0 +1,21 @@ +use crate::InferenceStep; +use burn_core::module::Module; +use std::marker::PhantomData; + +/// All components necessary to evaluate a model grouped in one trait. +pub trait EvaluatorComponentTypes { + /// The model to evaluate. + type Model: Module + InferenceStep + core::fmt::Display + 'static; +} + +/// A marker type used to provide [evaluation components](EvaluatorComponentTypes). +pub struct EvaluatorComponentTypesMarker { + _p: PhantomData, +} + +impl EvaluatorComponentTypes for EvaluatorComponentTypesMarker +where + M: Module + InferenceStep + core::fmt::Display + 'static, +{ + type Model = M; +} diff --git a/crates/burn-train/src/evaluator/mod.rs b/crates/burn-train/src/evaluator/mod.rs new file mode 100644 index 0000000..aa20339 --- /dev/null +++ b/crates/burn-train/src/evaluator/mod.rs @@ -0,0 +1,7 @@ +mod base; +mod builder; + +pub(crate) mod components; + +pub use base::*; +pub use builder::*; diff --git a/crates/burn-train/src/learner/application_logger.rs b/crates/burn-train/src/learner/application_logger.rs new file mode 100644 index 0000000..42f725c --- /dev/null +++ b/crates/burn-train/src/learner/application_logger.rs @@ -0,0 +1,69 @@ +use std::path::{Path, PathBuf}; +use tracing_core::{Level, LevelFilter}; +use tracing_subscriber::filter::filter_fn; +use tracing_subscriber::prelude::*; +use tracing_subscriber::{Layer, registry}; + +/// This trait is used to install an application logger. +pub trait ApplicationLoggerInstaller { + /// Install the application logger. + fn install(&self) -> Result<(), String>; +} + +/// This struct is used to install a local file application logger to output logs to a given file path. +pub struct FileApplicationLoggerInstaller { + path: PathBuf, +} + +impl FileApplicationLoggerInstaller { + /// Create a new file application logger. + pub fn new(path: impl AsRef) -> Self { + Self { + path: path.as_ref().to_path_buf(), + } + } +} + +impl ApplicationLoggerInstaller for FileApplicationLoggerInstaller { + fn install(&self) -> Result<(), String> { + let path = Path::new(&self.path); + let writer = tracing_appender::rolling::never( + path.parent().unwrap_or_else(|| Path::new(".")), + path.file_name().unwrap_or_else(|| { + panic!("The path '{}' to point to a file.", self.path.display()) + }), + ); + let layer = tracing_subscriber::fmt::layer() + .with_ansi(false) + .with_writer(writer) + .with_filter(LevelFilter::INFO) + .with_filter(filter_fn(|m| { + if let Some(path) = m.module_path() { + // The wgpu crate is logging too much, so we skip `info` level. + if path.starts_with("wgpu") && *m.level() >= Level::INFO { + return false; + } + } + true + })); + + if registry().with(layer).try_init().is_err() { + return Err("Failed to install the file logger.".to_string()); + } + + let hook = std::panic::take_hook(); + let file_path = self.path.to_owned(); + + std::panic::set_hook(Box::new(move |info| { + log::error!("PANIC => {info}"); + eprintln!( + "=== PANIC ===\nA fatal error happened, you can check the experiment logs here => \ + '{}'\n=============", + file_path.display() + ); + hook(info); + })); + + Ok(()) + } +} diff --git a/crates/burn-train/src/learner/base.rs b/crates/burn-train/src/learner/base.rs new file mode 100644 index 0000000..cf6bc44 --- /dev/null +++ b/crates/burn-train/src/learner/base.rs @@ -0,0 +1,261 @@ +use crate::checkpoint::{ + AsyncCheckpointer, Checkpointer, CheckpointingAction, CheckpointingStrategy, +}; +use crate::metric::store::EventStoreClient; +use crate::{ + CloneEarlyStoppingStrategy, LearnerModel, TrainOutput, TrainStep, TrainingModelInput, + TrainingModelOutput, +}; +use burn_core::store::ModuleRecord; +use burn_core::tensor::Device; +use burn_optim::lr_scheduler::LrSchedulerRecord; +use burn_optim::lr_scheduler::module_lr_scheduler::{ModuleLearningRate, ModuleLrScheduler}; +use burn_optim::{GradientsParams, ModuleOptimizer, MultiGradientsParams, OptimizerRecord}; +use std::marker::PhantomData; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; + +/// Learner struct encapsulating all components necessary to train a Neural Network model. +pub struct Learner { + pub(crate) model: M, + optim: ModuleOptimizer, + lr_scheduler: ModuleLrScheduler, + lr_module: ModuleLearningRate, +} + +impl Clone for Learner { + fn clone(&self) -> Self { + Self { + model: self.model.clone(), + optim: self.optim.clone(), + lr_scheduler: self.lr_scheduler.clone(), + lr_module: self.lr_module.clone(), + } + } +} + +impl Learner { + /// Create a learner. + pub fn new( + model: M, + optim: ModuleOptimizer, + lr_scheduler: impl Into, + ) -> Self { + Self { + model, + optim, + lr_scheduler: lr_scheduler.into(), + lr_module: 0.0.into(), + } + } +} + +impl Learner { + /// Fork the learner's model to the given device. + pub fn fork(&mut self, device: &Device) { + self.model = self.model().fork(device); + } + + /// Returns the current model. + pub fn model(&self) -> M { + self.model.clone() + } + + /// Returns the current learning rate. + pub fn lr_current(&self) -> ModuleLearningRate { + self.lr_module.clone() + } + + /// Executes a step of the learning rate scheduler. + pub fn lr_step(&mut self) { + self.lr_module = self.lr_scheduler.step(); + } + + /// Runs a step of the model for training, which executes the forward and backward passes. + /// + /// # Arguments + /// + /// * `item` - The input for the model. + /// + /// # Returns + /// + /// The output containing the model output and the gradients. + pub fn train_step(&self, item: TrainingModelInput) -> TrainOutput> { + TrainStep::step(&self.model, item) + } + + /// Optimize the current module with the provided gradients and learning rate. + /// + /// # Arguments + /// + /// * `optim`: Optimizer used for learning. + /// * `lr`: The learning rate used for this step. + /// * `grads`: The gradients of each parameter in the current model. + pub fn optimizer_step(&mut self, grads: GradientsParams) { + self.model = self + .model() + .optimize(&mut self.optim, self.lr_module.clone(), grads); + } + + /// Optimize the current module with the provided gradients and learning rate. + /// + /// # Arguments + /// + /// * `optim`: Optimizer used for learning. + /// * `lr`: The learning rate used for this step. + /// * `grads`: Multiple gradients associated to each parameter in the current model. + pub fn optimizer_step_multi(&mut self, grads: MultiGradientsParams) { + self.model = self + .model() + .optimize_multi(&mut self.optim, self.lr_module.clone(), grads); + } + + /// Load the module state from a [record](ModuleRecord). + pub fn load_model(&mut self, record: ModuleRecord) { + self.model = self.model.clone().load_record(record); + } + + /// Load the state of the learner's optimizer from a [record](OptimizerRecord). + /// + /// No device is needed: the optimizer state is migrated to each parameter's device on the next + /// step (see [`ModuleOptimizer::load_record`](burn_optim::ModuleOptimizer::load_record)). + pub fn load_optim(&mut self, record: OptimizerRecord) { + self.optim = self.optim.clone().load_record(record); + } + + /// Load the state of the learner's scheduler from a [record](LrSchedulerRecord). + pub fn load_scheduler(&mut self, record: LrSchedulerRecord) { + self.lr_scheduler = self.lr_scheduler.clone().load_record(record); + } +} + +/// Used to create, delete, or load checkpoints of the training process. +pub struct LearningCheckpointer { + model: AsyncCheckpointer, + optim: AsyncCheckpointer, + lr_scheduler: AsyncCheckpointer, + strategy: Box, + _phantom: PhantomData, +} + +impl LearningCheckpointer { + /// Create a new learning checkpointer. + pub fn new( + model: AsyncCheckpointer, + optim: AsyncCheckpointer, + lr_scheduler: AsyncCheckpointer, + strategy: Box, + ) -> Self { + Self { + model, + optim, + lr_scheduler, + strategy, + _phantom: PhantomData, + } + } + + /// Create checkpoint for the training process. + pub fn checkpoint(&mut self, learner: &Learner, epoch: usize, store: &EventStoreClient) { + let actions = self.strategy.checkpointing(epoch, store); + + for action in actions { + match action { + CheckpointingAction::Delete(epoch) => { + self.model + .delete(epoch) + .expect("Can delete model checkpoint."); + self.optim + .delete(epoch) + .expect("Can delete optimizer checkpoint."); + self.lr_scheduler + .delete(epoch) + .expect("Can delete learning rate scheduler checkpoint."); + } + CheckpointingAction::Save => { + self.model + .save(epoch, learner.model.clone().into_record()) + .expect("Can save model checkpoint."); + self.optim + .save(epoch, learner.optim.to_record()) + .expect("Can save optimizer checkpoint."); + self.lr_scheduler + .save(epoch, learner.lr_scheduler.to_record()) + .expect("Can save learning rate scheduler checkpoint."); + } + } + } + } + + /// Load a training checkpoint. + /// + /// No device is taken: checkpoints are device-free burnpack records (file-backed bytes). On + /// load, the model keeps the device of the learner's existing parameters, and the optimizer + /// state is migrated to each parameter's device on the next step. The training device is fixed + /// earlier, when the learner's model is created/forked. + pub fn load_checkpoint(&self, mut learner: Learner, epoch: usize) -> Learner { + let record = self + .model + .restore(epoch) + .expect("Can load model checkpoint."); + learner.load_model(record); + + let record = self + .optim + .restore(epoch) + .expect("Can load optimizer checkpoint."); + learner.load_optim(record); + + let record = self + .lr_scheduler + .restore(epoch) + .expect("Can load learning rate scheduler checkpoint."); + learner.load_scheduler(record); + + learner + } +} + +/// Cloneable reference to an early stopping strategy +pub(crate) type EarlyStoppingStrategyRef = Box; + +#[derive(Clone, Default)] +/// A handle that allows aborting the training/evaluation process early. +pub struct Interrupter { + state: Arc, + message: Arc>>, +} + +impl Interrupter { + /// Create a new instance. + pub fn new() -> Self { + Self::default() + } + + /// Notify the learner that it should stop. + /// # Arguments + /// * `reason` - A string describing the reason the training was stopped. + pub fn stop(&self, reason: Option<&str>) { + self.state.store(true, Ordering::Relaxed); + reason.inspect(|r| { + let mut message = self.message.lock().unwrap(); + *message = Some(String::from(*r)); + }); + } + + /// Reset the interrupter. + pub fn reset(&self) { + self.state.store(false, Ordering::Relaxed); + } + + /// True if .stop() has been called. + pub fn should_stop(&self) -> bool { + self.state.load(Ordering::Relaxed) + } + + /// Get the message associated with the interrupt. + pub fn get_message(&self) -> Option { + let message = self.message.lock().unwrap(); + message.clone() + } +} diff --git a/crates/burn-train/src/learner/classification.rs b/crates/burn-train/src/learner/classification.rs new file mode 100644 index 0000000..011bb1e --- /dev/null +++ b/crates/burn-train/src/learner/classification.rs @@ -0,0 +1,147 @@ +use crate::metric::{ + AccuracyInput, Adaptor, ConfusionStatsInput, HammingScoreInput, LossInput, PerplexityInput, + TopKAccuracyInput, processor::ItemLazy, +}; +use burn_core::tensor::{Device, Int, Tensor, Transaction}; + +/// Simple classification output adapted for multiple metrics. +/// +/// Supported metrics: +/// - Accuracy +/// - AUROC +/// - TopKAccuracy +/// - Perplexity +/// - Precision (via ConfusionStatsInput) +/// - Recall (via ConfusionStatsInput) +/// - FBetaScore (via ConfusionStatsInput) +/// - Loss. +#[derive(new)] +pub struct ClassificationOutput { + /// The loss. + pub loss: Tensor<1>, + + /// The class logits or probabilities. Shape: \[batch_size, num_classes\]. + pub output: Tensor<2>, + + /// The ground truth class index for each sample. Shape: \[batch_size\]. + pub targets: Tensor<1, Int>, +} + +impl ItemLazy for ClassificationOutput { + fn sync(self) -> Self { + let [output, loss, targets] = Transaction::default() + .register(self.output) + .register(self.loss) + .register(self.targets) + .execute() + .try_into() + .expect("Correct amount of tensor data"); + + let device: Device = Device::flex(); + + ClassificationOutput { + output: Tensor::from_data(output, &device), + loss: Tensor::from_data(loss, &device), + targets: Tensor::from_data(targets, &device), + } + } +} + +impl Adaptor for ClassificationOutput { + fn adapt(&self) -> AccuracyInput { + AccuracyInput::new(self.output.clone(), self.targets.clone()) + } +} + +impl Adaptor for ClassificationOutput { + fn adapt(&self) -> LossInput { + LossInput::new(self.loss.clone()) + } +} + +impl Adaptor for ClassificationOutput { + fn adapt(&self) -> TopKAccuracyInput { + TopKAccuracyInput::new(self.output.clone(), self.targets.clone()) + } +} + +impl Adaptor for ClassificationOutput { + fn adapt(&self) -> PerplexityInput { + PerplexityInput::new(self.output.clone(), self.targets.clone()) + } +} + +impl Adaptor for ClassificationOutput { + fn adapt(&self) -> ConfusionStatsInput { + let [_, num_classes] = self.output.dims(); + if num_classes > 1 { + ConfusionStatsInput::new( + self.output.clone(), + self.targets.clone().one_hot(num_classes).bool(), + ) + } else { + ConfusionStatsInput::new( + self.output.clone(), + self.targets.clone().unsqueeze_dim(1).bool(), + ) + } + } +} + +/// Multi-label classification output adapted for multiple metrics. +/// +/// Supported metrics: +/// - HammingScore +/// - Precision (via ConfusionStatsInput) +/// - Recall (via ConfusionStatsInput) +/// - FBetaScore (via ConfusionStatsInput) +/// - Loss +#[derive(new)] +pub struct MultiLabelClassificationOutput { + /// The loss. + pub loss: Tensor<1>, + + /// The label logits or probabilities. Shape: \[batch_size, num_classes\]. + pub output: Tensor<2>, + + /// The ground truth labels. Shape: \[batch_size, num_classes\]. + pub targets: Tensor<2, Int>, +} + +impl ItemLazy for MultiLabelClassificationOutput { + fn sync(self) -> Self { + let [output, loss, targets] = Transaction::default() + .register(self.output) + .register(self.loss) + .register(self.targets) + .execute() + .try_into() + .expect("Correct amount of tensor data"); + + let device: Device = Device::flex(); + + MultiLabelClassificationOutput { + output: Tensor::from_data(output, &device), + loss: Tensor::from_data(loss, &device), + targets: Tensor::from_data(targets, &device), + } + } +} + +impl Adaptor for MultiLabelClassificationOutput { + fn adapt(&self) -> HammingScoreInput { + HammingScoreInput::new(self.output.clone(), self.targets.clone()) + } +} + +impl Adaptor for MultiLabelClassificationOutput { + fn adapt(&self) -> LossInput { + LossInput::new(self.loss.clone()) + } +} + +impl Adaptor for MultiLabelClassificationOutput { + fn adapt(&self) -> ConfusionStatsInput { + ConfusionStatsInput::new(self.output.clone(), self.targets.clone().bool()) + } +} diff --git a/crates/burn-train/src/learner/early_stopping.rs b/crates/burn-train/src/learner/early_stopping.rs new file mode 100644 index 0000000..3bb8aed --- /dev/null +++ b/crates/burn-train/src/learner/early_stopping.rs @@ -0,0 +1,301 @@ +use crate::metric::{ + Metric, MetricName, + store::{Aggregate, Direction, EventStoreClient, Split}, +}; + +/// The condition that [early stopping strategies](EarlyStoppingStrategy) should follow. +#[derive(Clone)] +pub enum StoppingCondition { + /// When no improvement has happened since the given number of epochs. + NoImprovementSince { + /// The number of epochs allowed to worsen before it gets better. + n_epochs: usize, + }, +} + +/// A strategy that checks if the training should be stopped. +pub trait EarlyStoppingStrategy: Send { + /// Update its current state and returns if the training should be stopped. + fn should_stop(&mut self, epoch: usize, store: &EventStoreClient) -> bool; +} + +/// A helper trait to provide type-erased cloning. +pub trait CloneEarlyStoppingStrategy: EarlyStoppingStrategy + Send { + /// Clone into a boxed trait object. + fn clone_box(&self) -> Box; +} + +/// Blanket-implement `CloneEarlyStoppingStrategy` for any `T` that +/// already implements your strategy + `Clone` + `Send` + `'static`. +impl CloneEarlyStoppingStrategy for T +where + T: EarlyStoppingStrategy + Clone + Send + 'static, +{ + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } +} + +/// Now you can `impl Clone` for the boxed trait object. +impl Clone for Box { + fn clone(&self) -> Box { + self.clone_box() + } +} + +/// An [early stopping strategy](EarlyStoppingStrategy) based on a metrics collected +/// during training or validation. +#[derive(Clone)] +pub struct MetricEarlyStoppingStrategy { + condition: StoppingCondition, + metric_name: MetricName, + aggregate: Aggregate, + direction: Direction, + split: Split, + best_epoch: usize, + best_value: f64, + warmup_epochs: Option, +} + +impl EarlyStoppingStrategy for MetricEarlyStoppingStrategy { + fn should_stop(&mut self, epoch: usize, store: &EventStoreClient) -> bool { + let current_value = + match store.find_metric(&self.metric_name, epoch, self.aggregate, &self.split) { + Some(value) => value, + None => { + log::warn!("Can't find metric for early stopping."); + return false; + } + }; + + let is_best = match self.direction { + Direction::Lowest => current_value < self.best_value, + Direction::Highest => current_value > self.best_value, + }; + + if is_best { + log::info!( + "New best epoch found {} {}: {}", + epoch, + self.metric_name, + current_value + ); + self.best_value = current_value; + self.best_epoch = epoch; + return false; + } + + if let Some(warmup_epochs) = self.warmup_epochs + && epoch <= warmup_epochs + { + return false; + } + + match self.condition { + StoppingCondition::NoImprovementSince { n_epochs } => { + let should_stop = epoch - self.best_epoch >= n_epochs; + + if should_stop { + log::info!( + "Stopping training loop, no improvement since epoch {}, {}: {}, current \ + epoch {}, {}: {}", + self.best_epoch, + self.metric_name, + self.best_value, + epoch, + self.metric_name, + current_value + ); + } + + should_stop + } + } + } +} + +impl MetricEarlyStoppingStrategy { + /// Create a new [early stopping strategy](EarlyStoppingStrategy) based on a metrics collected + /// during training or validation. + /// + /// # Notes + /// + /// The metric should be registered for early stopping to work, otherwise no data is collected. + pub fn new( + metric: &Me, + aggregate: Aggregate, + direction: Direction, + split: Split, + condition: StoppingCondition, + ) -> Self { + let init_value = match direction { + Direction::Lowest => f64::MAX, + Direction::Highest => f64::MIN, + }; + + Self { + metric_name: metric.name(), + condition, + aggregate, + direction, + split, + best_epoch: 1, + best_value: init_value, + warmup_epochs: None, + } + } + + /// Get the warmup period. + /// + /// Early stopping will not trigger during the warmup epochs. + pub fn warmup_epochs(&self) -> Option { + self.warmup_epochs + } + + /// Set the warmup epochs. + /// + /// Early stopping will not trigger during the warmup epochs. + /// + /// # Arguments + /// - `warmup`: the number of warmup epochs, or None. + pub fn with_warmup_epochs(self, warmup: Option) -> Self { + Self { + warmup_epochs: warmup, + ..self + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use crate::{ + EventProcessorTraining, + logger::InMemoryMetricLogger, + metric::{ + LossMetric, + processor::{ + MetricsTraining, MinimalEventProcessor, + test_utils::{end_epoch, process_train}, + }, + store::LogEventStore, + }, + test_utils::start_epoch, + }; + + use super::*; + + #[test] + fn never_early_stop_while_it_is_improving() { + test_early_stopping( + None, + 1, + &[ + (&[0.5, 0.3], false, "Should not stop first epoch"), + (&[0.4, 0.3], false, "Should not stop when improving"), + (&[0.3, 0.3], false, "Should not stop when improving"), + (&[0.2, 0.3], false, "Should not stop when improving"), + ], + ); + } + + #[test] + fn early_stop_when_no_improvement_since_two_epochs() { + test_early_stopping( + None, + 2, + &[ + (&[1.0, 0.5], false, "Should not stop first epoch"), + (&[0.5, 0.3], false, "Should not stop when improving"), + ( + &[1.0, 3.0], + false, + "Should not stop first time it gets worse", + ), + ( + &[1.0, 2.0], + true, + "Should stop since two following epochs didn't improve", + ), + ], + ); + } + + #[test] + fn early_stopping_with_warmup() { + test_early_stopping( + Some(3), + 2, + &[ + (&[1.0, 0.5], false, "Should not stop during warmup"), + (&[1.0, 0.5], false, "Should not stop during warmup"), + (&[1.0, 0.5], false, "Should not stop during warmup"), + ( + &[1.0, 0.5], + true, + "Should stop when not improving after warmup", + ), + ], + ) + } + + #[test] + fn early_stop_when_stays_equal() { + test_early_stopping( + None, + 2, + &[ + (&[0.5, 0.3], false, "Should not stop first epoch"), + ( + &[0.5, 0.3], + false, + "Should not stop first time it stars the same", + ), + ( + &[0.5, 0.3], + true, + "Should stop since two following epochs didn't improve", + ), + ], + ); + } + + fn test_early_stopping(warmup: Option, n_epochs: usize, data: &[(&[f64], bool, &str)]) { + let loss = LossMetric::new(); + let mut early_stopping = MetricEarlyStoppingStrategy::new( + &loss, + Aggregate::Mean, + Direction::Lowest, + Split::Train, + StoppingCondition::NoImprovementSince { n_epochs }, + ) + .with_warmup_epochs(warmup); + let mut store = LogEventStore::default(); + let mut metrics = MetricsTraining::::default(); + + store.register_logger(InMemoryMetricLogger::default()); + metrics.register_train_metric_numeric(loss); + + let store = Arc::new(EventStoreClient::new(store)); + let mut processor = MinimalEventProcessor::new(metrics, store.clone()); + + processor.process_train(crate::LearnerEvent::Start { + total_epochs: 0, + starting_epoch: 0, + }); + for (epoch, (points, should_start, comment)) in (1..).zip(data.iter()) { + start_epoch(&mut processor, epoch, points.len()); + for point in points.iter() { + process_train(&mut processor, *point, epoch); + } + end_epoch(&mut processor, epoch); + + assert_eq!( + *should_start, + early_stopping.should_stop(epoch, &store), + "{comment}" + ); + } + } +} diff --git a/crates/burn-train/src/learner/mod.rs b/crates/burn-train/src/learner/mod.rs new file mode 100644 index 0000000..bdbc573 --- /dev/null +++ b/crates/burn-train/src/learner/mod.rs @@ -0,0 +1,26 @@ +#[cfg(feature = "rl")] +mod rl; +#[cfg(feature = "rl")] +pub use rl::*; + +mod application_logger; +mod base; +mod classification; +mod early_stopping; +mod regression; +mod sequence; +mod sharder; +mod summary; +mod supervised; +mod train_val; + +pub use application_logger::*; +pub use base::*; +pub use classification::*; +pub use early_stopping::*; +pub use regression::*; +pub use sequence::*; +pub use sharder::*; +pub use summary::*; +pub use supervised::*; +pub use train_val::*; diff --git a/crates/burn-train/src/learner/regression.rs b/crates/burn-train/src/learner/regression.rs new file mode 100644 index 0000000..a70e8a8 --- /dev/null +++ b/crates/burn-train/src/learner/regression.rs @@ -0,0 +1,42 @@ +use crate::metric::processor::ItemLazy; +use crate::metric::{Adaptor, LossInput}; +use burn_core::tensor::{Device, Tensor, Transaction}; + +/// Regression output adapted for the loss metric. +#[derive(new)] +pub struct RegressionOutput { + /// The loss. + pub loss: Tensor<1>, + + /// The predicted values. Shape: \[batch_size, num_targets\]. + pub output: Tensor<2>, + + /// The ground truth values. Shape: \[batch_size, num_targets\]. + pub targets: Tensor<2>, +} + +impl Adaptor for RegressionOutput { + fn adapt(&self) -> LossInput { + LossInput::new(self.loss.clone()) + } +} + +impl ItemLazy for RegressionOutput { + fn sync(self) -> Self { + let [output, loss, targets] = Transaction::default() + .register(self.output) + .register(self.loss) + .register(self.targets) + .execute() + .try_into() + .expect("Correct amount of tensor data"); + + let device: Device = Device::flex(); + + RegressionOutput { + output: Tensor::from_data(output, &device), + loss: Tensor::from_data(loss, &device), + targets: Tensor::from_data(targets, &device), + } + } +} diff --git a/crates/burn-train/src/learner/rl/checkpointer.rs b/crates/burn-train/src/learner/rl/checkpointer.rs new file mode 100644 index 0000000..dbe40e8 --- /dev/null +++ b/crates/burn-train/src/learner/rl/checkpointer.rs @@ -0,0 +1,77 @@ +use burn_rl::{Policy, PolicyLearner, PolicyState}; + +use crate::RLAgentRecord; +use crate::{ + RLComponentsTypes, RLPolicyRecord, + checkpoint::Checkpointer, + checkpoint::{AsyncCheckpointer, Checkpoint, CheckpointingAction, CheckpointingStrategy}, + metric::store::EventStoreClient, +}; + +#[derive(new)] +/// Used to create, delete, or load checkpoints of the training process. +pub struct RLCheckpointer { + policy: AsyncCheckpointer>, + learning_agent: AsyncCheckpointer>, + strategy: Box, +} + +impl RLCheckpointer +where + RLPolicyRecord: Checkpoint, + RLAgentRecord: Checkpoint, +{ + /// Create checkpoint for the training process. + pub fn checkpoint( + &mut self, + policy: &RLC::PolicyState, + learning_agent: &RLC::LearningAgent, + epoch: usize, + store: &EventStoreClient, + ) { + let actions = self.strategy.checkpointing(epoch, store); + + for action in actions { + match action { + CheckpointingAction::Delete(epoch) => { + self.policy + .delete(epoch) + .expect("Can delete policy checkpoint."); + self.learning_agent + .delete(epoch) + .expect("Can delete learning agent checkpoint.") + } + CheckpointingAction::Save => { + self.policy + .save(epoch, policy.clone().into_record()) + .expect("Can save policy checkpoint."); + self.learning_agent + .save(epoch, learning_agent.record()) + .expect("Can save learning agent checkpoint."); + } + } + } + } + + /// Load a training checkpoint. + pub fn load_checkpoint( + &self, + learning_agent: RLC::LearningAgent, + epoch: usize, + ) -> RLC::LearningAgent { + let record = self + .policy + .restore(epoch) + .expect("Can load model checkpoint."); + let policy = learning_agent.policy().load_record(record); + + let record = self + .learning_agent + .restore(epoch) + .expect("Can load learning agent checkpoint."); + let mut learning_agent = learning_agent.load_record(record); + learning_agent.update_policy(policy); + + learning_agent + } +} diff --git a/crates/burn-train/src/learner/rl/components.rs b/crates/burn-train/src/learner/rl/components.rs new file mode 100644 index 0000000..5164e6e --- /dev/null +++ b/crates/burn-train/src/learner/rl/components.rs @@ -0,0 +1,110 @@ +use std::marker::PhantomData; + +use burn_rl::{ + Batchable, Environment, EnvironmentInit, Policy, PolicyLearner, PolicyState, ToAction, + ToObservation, +}; + +use crate::checkpoint::Checkpoint; +use crate::{AgentEvaluationEvent, AsyncProcessorTraining, ItemLazy, RLEvent}; + +/// All components used by the reinforcement learning paradigm, grouped in one trait. +pub trait RLComponentsTypes { + /// The learning environment. + type Env: Environment + 'static; + /// Specifies how to initialize the environment. + type EnvInit: EnvironmentInit + Send + 'static; + /// The type of the environment state. + type State: ToObservation<::Observation> + Clone + Send + 'static; + /// The type of the environment action. + type Action: From<::Action> + + ToAction<::Action> + + Clone + + Send + + 'static; + + /// The policy used to take actions in the environment. + type Policy: Policy< + Observation = Self::PolicyObs, + ActionDistribution = Self::PolicyAD, + Action = Self::PolicyAction, + ActionContext = Self::ActionContext, + PolicyState = Self::PolicyState, + > + Send + + 'static; + /// The policy's observation type. + type PolicyObs: Clone + Send + Batchable + 'static; + /// The policy's action distribution type. + type PolicyAD: Clone + Send + Batchable; + /// The policy's action type. + type PolicyAction: Clone + Send + Batchable; + /// Additional data as context for an agent's action. + type ActionContext: ItemLazy + Clone + Send + 'static; + /// The state of the parameterized policy. + type PolicyState: Clone + Send + PolicyState + 'static; + + /// The learning agent. + type LearningAgent: PolicyLearner< + TrainContext = Self::TrainingOutput, + InnerPolicy = Self::Policy, + Record: Checkpoint, + > + Send + + 'static; + /// The output data of a training step. + type TrainingOutput: ItemLazy + Clone + Send; +} + +/// Concrete type that implements the [RLComponentsTypes](RLComponentsTypes) trait. +pub struct RLComponentsMarker { + _env: PhantomData, + _env_init: PhantomData, + _agent: PhantomData, +} + +impl RLComponentsTypes for RLComponentsMarker +where + E: Environment + 'static, + EI: EnvironmentInit + Send + 'static, + A: PolicyLearner + Send + 'static, + ::Record: Checkpoint, + A::TrainContext: ItemLazy + Clone + Send, + A::InnerPolicy: Policy + Send, + ::Observation: Batchable + Clone + Send, + ::ActionDistribution: Batchable + Clone + Send, + ::Action: Batchable + Clone + Send, + ::ActionContext: ItemLazy + Clone + Send + 'static, + ::PolicyState: Clone + Send, + <::PolicyState as PolicyState>::Record: Checkpoint, + E::State: ToObservation<::Observation> + Clone + Send + 'static, + E::Action: From<::Action> + + ToAction<::Action> + + Clone + + Send + + 'static, +{ + type Env = E; + type EnvInit = EI; + type LearningAgent = A; + type Policy = A::InnerPolicy; + type PolicyObs = ::Observation; + type PolicyAD = ::ActionDistribution; + type PolicyAction = ::Action; + type ActionContext = ::ActionContext; + type PolicyState = ::PolicyState; + type TrainingOutput = A::TrainContext; + type State = E::State; + type Action = E::Action; +} + +pub(crate) type RlPolicy = + <::LearningAgent as PolicyLearner>::InnerPolicy; +/// The event processor type for reinforcement learning. +pub type RLEventProcessorType = AsyncProcessorTraining< + RLEvent<::TrainingOutput, ::ActionContext>, + AgentEvaluationEvent<::ActionContext>, +>; +/// The record of the policy. +pub type RLPolicyRecord = + <<::Policy as Policy>::PolicyState as PolicyState>::Record; +/// The record of the learning agent. +pub type RLAgentRecord = <::LearningAgent as PolicyLearner>::Record; diff --git a/crates/burn-train/src/learner/rl/env_runner/async_runner.rs b/crates/burn-train/src/learner/rl/env_runner/async_runner.rs new file mode 100644 index 0000000..93a2e62 --- /dev/null +++ b/crates/burn-train/src/learner/rl/env_runner/async_runner.rs @@ -0,0 +1,735 @@ +use rand::prelude::SliceRandom; +use std::{ + sync::mpsc::{Receiver, Sender}, + thread::spawn, +}; + +use burn_core::{Tensor, data::dataloader::Progress, tensor::Device}; +use burn_rl::Policy; +use burn_rl::Transition; +use burn_rl::{AsyncPolicy, Environment}; +use burn_rl::{EnvironmentInit, ToObservation}; + +use crate::{ + AgentEnvLoop, AgentEvaluationEvent, EpisodeSummary, EvaluationItem, EventProcessorTraining, + Interrupter, RLComponentsTypes, RLEvent, RLEventProcessorType, RLTimeStep, RLTrajectory, + RlPolicy, TimeStep, Trajectory, +}; + +enum RequestMessage { + Step(), + Episode(), +} + +/// Configuration for an async agent/environment loop. +pub struct AsyncAgentEnvLoopConfig { + /// If the loop is used for evaluation (as opposed to training). + pub eval: bool, + /// If the agent should take action deterministically. + pub deterministic: bool, + /// An arbitrary ID for the loop. + pub id: usize, +} + +/// An asynchronous agent/environement interface. +pub struct AgentEnvAsyncLoop { + eval: bool, + agent: AsyncPolicy>, + transition_receiver: Receiver>, + trajectory_receiver: Receiver>, + request_sender: Sender, + device: Device, +} + +impl AgentEnvAsyncLoop { + /// Create a new asynchronous runner. + /// + /// # Arguments + /// + /// * `env_init` - A function returning an environment instance. + /// * `agent` - An [AsyncPolicy](AsyncPolicy) taking actions in the loop. + /// * `config` - An [AsyncAgentEnvLoopConfig](AsyncAgentEnvLoopConfig). + /// * `transition_sender` - Optional sender for transitions if you want to drive the requests from outside of the loop instance. + /// * `trajectory_sender` - Optional sender for trajectories if you want to drive the requests from outside of the loop instance. + /// + /// # Returns + /// + /// An async Agent/Environement loop. + pub fn new( + env_init: RLC::EnvInit, + agent: AsyncPolicy>, + config: AsyncAgentEnvLoopConfig, + transition_device: &Device, + transition_sender: Option>>, + trajectory_sender: Option>>, + ) -> Self { + let (loop_transition_sender, transition_receiver) = std::sync::mpsc::channel(); + let (loop_trajectory_sender, trajectory_receiver) = std::sync::mpsc::channel(); + let (request_sender, request_receiver) = std::sync::mpsc::channel(); + let loop_transition_sender = transition_sender.unwrap_or(loop_transition_sender); + let loop_trajectory_sender = trajectory_sender.unwrap_or(loop_trajectory_sender); + + let device = transition_device.clone(); + let mut loop_agent = agent.clone().to_device(transition_device); + let eval = config.eval; + + let mut current_steps = vec![]; + let mut current_reward = 0.0; + let mut step_num = 0; + spawn(move || { + let mut env = env_init.init(); + env.reset(); + + let mut request_episode = false; + loop { + let state = env.state(); + let (action, context) = + loop_agent.action(state.clone().to_observation(&device), config.deterministic); + + let env_action = RLC::Action::from(action); + let step_result = env.step(env_action.clone()); + + current_reward += step_result.reward; + step_num += 1; + + let transition = Transition::new( + state.clone(), + step_result.next_state, + env_action, + Tensor::from_data([step_result.reward], &device), + Tensor::from_data( + [(step_result.done || step_result.truncated) as i32 as f64], + &device, + ), + ); + + if !request_episode { + loop_agent.decrement_agents(1); + let request = match request_receiver.recv() { + Ok(req) => req, + Err(err) => { + log::error!("Error in env runner : {}", err); + break; + } + }; + loop_agent.increment_agents(1); + + match request { + RequestMessage::Step() => (), + RequestMessage::Episode() => request_episode = true, + } + } + + let time_step = TimeStep { + env_id: config.id, + transition, + done: step_result.done, + ep_len: step_num, + cum_reward: current_reward, + action_context: context[0].clone(), + }; + current_steps.push(time_step.clone()); + + if !request_episode && let Err(err) = loop_transition_sender.send(time_step) { + log::error!("Error in env runner : {}", err); + break; + } + + if step_result.done || step_result.truncated { + if request_episode { + request_episode = false; + loop_trajectory_sender + .send(Trajectory { + timesteps: current_steps.clone(), + }) + .expect("Can send trajectory to main thread."); + } + current_steps.clear(); + + env.reset(); + current_reward = 0.; + step_num = 0; + } + } + }); + + Self { + eval, + agent, + transition_receiver, + trajectory_receiver, + request_sender, + device: transition_device.clone(), + } + } +} + +impl AgentEnvLoop for AgentEnvAsyncLoop +where + RLC: RLComponentsTypes, +{ + fn run_steps( + &mut self, + num_steps: usize, + processor: &mut RLEventProcessorType, + interrupter: &Interrupter, + progress: &mut Progress, + ) -> Vec> { + let mut items = vec![]; + for _ in 0..num_steps { + self.request_sender + .send(RequestMessage::Step()) + .expect("Can request transitions."); + let transition = self + .transition_receiver + .recv() + .expect("Can receive transitions."); + items.push(transition.clone()); + + if !self.eval { + progress.items_processed += 1; + processor.process_train(RLEvent::EnvStep(EvaluationItem::new( + transition.action_context, + progress.clone(), + None, + ))); + + if transition.done { + processor.process_train(RLEvent::EpisodeEnd(EvaluationItem::new( + EpisodeSummary { + episode_length: transition.ep_len, + cum_reward: transition.cum_reward, + }, + progress.clone(), + None, + ))); + } + } + + if interrupter.should_stop() { + break; + } + } + items + } + + fn run_episodes( + &mut self, + num_episodes: usize, + processor: &mut RLEventProcessorType, + interrupter: &Interrupter, + _progress: &mut Progress, + ) -> Vec> { + let mut items = vec![]; + self.agent.increment_agents(1); + for episode_num in 0..num_episodes { + self.request_sender + .send(RequestMessage::Episode()) + .expect("Can request episodes."); + let trajectory = self + .trajectory_receiver + .recv() + .expect("Main thread can receive trajectory."); + + for (i, step) in trajectory.timesteps.iter().enumerate() { + // TODO : clean this. + if self.eval { + processor.process_valid(AgentEvaluationEvent::EnvStep(EvaluationItem::new( + step.action_context.clone(), + Progress::new(i, i, Some("steps".to_string())), + None, + ))); + + if step.done { + processor.process_valid(AgentEvaluationEvent::EpisodeEnd( + EvaluationItem::new( + EpisodeSummary { + episode_length: step.ep_len, + cum_reward: step.cum_reward, + }, + Progress::new( + episode_num + 1, + num_episodes, + Some("episodes".to_string()), + ), + None, + ), + )); + } + } else { + processor.process_train(RLEvent::EnvStep(EvaluationItem::new( + step.action_context.clone(), + Progress::new(i, i, Some("steps".to_string())), + None, + ))); + + if step.done { + processor.process_train(RLEvent::EpisodeEnd(EvaluationItem::new( + EpisodeSummary { + episode_length: step.ep_len, + cum_reward: step.cum_reward, + }, + Progress::new( + episode_num + 1, + num_episodes, + Some("episodes".to_string()), + ), + None, + ))); + } + } + } + + items.push(trajectory); + if interrupter.should_stop() { + break; + } + } + self.agent.decrement_agents(1); + items + } + + fn update_policy(&mut self, update: RLC::PolicyState) { + self.agent.update(update); + self.agent.clone().to_device(&self.device); + } + + fn policy(&self) -> RLC::PolicyState { + self.agent.state() + } + + fn device(&self) -> Device { + self.device.clone() + } +} + +/// An asynchronous runner for multiple agent/environement interfaces. +pub struct MultiAgentEnvLoop { + num_envs: usize, + eval: bool, + agent: AsyncPolicy, + transition_receiver: Receiver>, + trajectory_receiver: Receiver>, + request_senders: Vec>, + device: Device, +} + +impl MultiAgentEnvLoop { + /// Create a new asynchronous runner for multiple agent/environement interfaces. + pub fn new( + num_envs: usize, + env_init: RLC::EnvInit, + agent: AsyncPolicy, + eval: bool, + deterministic: bool, + device: &Device, + ) -> Self { + let (transition_sender, transition_receiver) = std::sync::mpsc::channel(); + let (trajectory_sender, trajectory_receiver) = std::sync::mpsc::channel(); + let mut request_senders = vec![]; + + // Double batching : The environments are always one step ahead of requests. This allows inference for the first batch of steps. + agent.increment_agents(num_envs); + + for i in 0..num_envs { + let config = AsyncAgentEnvLoopConfig { + eval, + deterministic, + id: i, + }; + let runner = AgentEnvAsyncLoop::::new( + env_init.clone(), + agent.clone(), + config, + &device.clone(), + Some(transition_sender.clone()), + Some(trajectory_sender.clone()), + ); + request_senders.push(runner.request_sender.clone()); + } + + // Double batching : The environments are always one step ahead. + request_senders.iter().for_each(|s| { + s.send(RequestMessage::Step()) + .expect("Main thread can send step requests.") + }); + + Self { + num_envs, + eval, + agent: agent.clone(), + transition_receiver, + trajectory_receiver, + request_senders, + device: device.clone(), + } + } +} + +impl AgentEnvLoop for MultiAgentEnvLoop +where + RLC: RLComponentsTypes, +{ + fn run_steps( + &mut self, + num_steps: usize, + processor: &mut RLEventProcessorType, + interrupter: &Interrupter, + progress: &mut Progress, + ) -> Vec> { + let mut items = vec![]; + for _ in 0..num_steps { + let transition = self + .transition_receiver + .recv() + .expect("Can receive transitions."); + items.push(transition.clone()); + + self.request_senders[transition.env_id] + .send(RequestMessage::Step()) + .expect("Main thread can request steps."); + + if !self.eval { + progress.items_processed += 1; + processor.process_train(RLEvent::EnvStep(EvaluationItem::new( + transition.action_context, + progress.clone(), + None, + ))); + + if transition.done { + processor.process_train(RLEvent::EpisodeEnd(EvaluationItem::new( + EpisodeSummary { + episode_length: transition.ep_len, + cum_reward: transition.cum_reward, + }, + progress.clone(), + None, + ))); + } + } + + if interrupter.should_stop() { + break; + } + } + items + } + + fn update_policy(&mut self, update: RLC::PolicyState) { + self.agent.update(update); + self.agent.clone().to_device(&self.device); + } + + fn run_episodes( + &mut self, + num_episodes: usize, + processor: &mut RLEventProcessorType, + interrupter: &Interrupter, + _progress: &mut Progress, + ) -> Vec> { + // Send `num_episodes` initial requests. + let mut idx = vec![]; + if num_episodes < self.num_envs { + let mut rng = rand::rng(); + let mut vec: Vec = (0..self.num_envs).collect(); + vec.shuffle(&mut rng); + idx = vec.into_iter().take(num_episodes).collect(); + } else { + idx = (0..self.num_envs).collect(); + } + let num_requests = self.num_envs.min(num_episodes); + idx.into_iter().for_each(|i| { + self.request_senders[i] + .send(RequestMessage::Episode()) + .expect("Main thread can request steps."); + }); + + let mut items = vec![]; + for episode_num in 0..num_episodes { + let trajectory = self + .trajectory_receiver + .recv() + .expect("Can receive trajectory."); + items.push(trajectory.clone()); + if items.len() + num_requests <= num_episodes { + self.request_senders[trajectory.timesteps[0].env_id] + .send(RequestMessage::Episode()) + .expect("Main thread can request steps."); + } + for (i, step) in trajectory.timesteps.iter().enumerate() { + if self.eval { + processor.process_valid(AgentEvaluationEvent::EnvStep(EvaluationItem::new( + step.action_context.clone(), + Progress::new(i, i, Some("steps".to_string())), + None, + ))); + + if step.done { + processor.process_valid(AgentEvaluationEvent::EpisodeEnd( + EvaluationItem::new( + EpisodeSummary { + episode_length: step.ep_len, + cum_reward: step.cum_reward, + }, + Progress::new( + episode_num + 1, + num_episodes, + Some("episodes".to_string()), + ), + None, + ), + )); + } + } else { + processor.process_train(RLEvent::EnvStep(EvaluationItem::new( + step.action_context.clone(), + Progress::new(i, i, Some("steps".to_string())), + None, + ))); + + if step.done { + processor.process_train(RLEvent::EpisodeEnd(EvaluationItem::new( + EpisodeSummary { + episode_length: step.ep_len, + cum_reward: step.cum_reward, + }, + Progress::new( + episode_num + 1, + num_episodes, + Some("episodes".to_string()), + ), + None, + ))); + } + } + } + + if interrupter.should_stop() { + break; + } + } + + items + } + + fn policy(&self) -> RLC::PolicyState { + self.agent.state() + } + + fn device(&self) -> Device { + self.device.clone() + } +} + +#[cfg(test)] +#[allow(clippy::needless_range_loop)] +mod tests { + use burn_core::data::dataloader::Progress; + use burn_rl::AsyncPolicy; + + use crate::learner::rl::env_runner::async_runner::AsyncAgentEnvLoopConfig; + use crate::learner::rl::env_runner::base::AgentEnvLoop; + use crate::learner::tests::{MockPolicyState, MockProcessor}; + use crate::{ + AgentEnvAsyncLoop, + learner::tests::{MockEnvInit, MockPolicy, MockRLComponents}, + }; + use crate::{AsyncProcessorTraining, Interrupter, MultiAgentEnvLoop}; + + fn setup_async_loop( + state: usize, + eval: bool, + deterministic: bool, + ) -> AgentEnvAsyncLoop { + let env_init = MockEnvInit; + let agent = MockPolicy(state); + let config = AsyncAgentEnvLoopConfig { + eval, + deterministic, + id: 0, + }; + AgentEnvAsyncLoop::::new( + env_init, + AsyncPolicy::new(1, agent), + config, + &Default::default(), + None, + None, + ) + } + + fn setup_multi_loop( + num_envs: usize, + autobatch_size: usize, + state: usize, + eval: bool, + deterministic: bool, + ) -> MultiAgentEnvLoop { + let env_init = MockEnvInit; + let agent = MockPolicy(state); + MultiAgentEnvLoop::::new( + num_envs, + env_init, + AsyncPolicy::new(autobatch_size, agent), + eval, + deterministic, + &Default::default(), + ) + } + + #[test] + fn test_policy_async_loop() { + let runner = setup_async_loop(1000, false, false); + let policy_state = runner.policy(); + assert_eq!(policy_state.0, 1000); + } + + #[test] + fn test_update_policy_async_loop() { + let mut runner = setup_async_loop(0, false, false); + + runner.update_policy(MockPolicyState(1)); + assert_eq!(runner.policy().0, 1); + } + + #[test] + fn run_steps_returns_requested_number_async_loop() { + let mut runner = setup_async_loop(0, false, false); + let mut processor = AsyncProcessorTraining::new(MockProcessor); + let interrupter = Interrupter::new(); + let mut progress = Progress { + items_processed: 0, + items_total: 1, + unit: Some("steps".to_string()), + }; + + let steps = runner.run_steps(1, &mut processor, &interrupter, &mut progress); + assert_eq!(steps.len(), 1); + let steps = runner.run_steps(8, &mut processor, &interrupter, &mut progress); + assert_eq!(steps.len(), 8); + } + + #[test] + fn run_episodes_returns_requested_number_async_loop() { + let mut runner = setup_async_loop(0, false, false); + let mut processor = AsyncProcessorTraining::new(MockProcessor); + let interrupter = Interrupter::new(); + let mut progress = Progress { + items_processed: 0, + items_total: 1, + unit: Some("steps".to_string()), + }; + + let trajectories = runner.run_episodes(1, &mut processor, &interrupter, &mut progress); + assert_eq!(trajectories.len(), 1); + assert_ne!(trajectories[0].timesteps.len(), 0); + let trajectories = runner.run_episodes(8, &mut processor, &interrupter, &mut progress); + assert_eq!(trajectories.len(), 8); + for i in 0..8 { + assert_ne!(trajectories[i].timesteps.len(), 0); + } + } + + #[test] + fn test_policy_multi_loop() { + let runner = setup_multi_loop(4, 4, 1000, false, false); + let policy_state = runner.policy(); + assert_eq!(policy_state.0, 1000); + } + + #[test] + fn test_update_policy_multi_loop() { + let mut runner = setup_multi_loop(4, 4, 0, false, false); + + runner.update_policy(MockPolicyState(1)); + assert_eq!(runner.policy().0, 1); + } + + #[test] + fn run_steps_returns_requested_number_multi_loop() { + fn run_test(num_envs: usize, autobatch_size: usize) { + let mut runner = setup_multi_loop(num_envs, autobatch_size, 0, false, false); + let mut processor = AsyncProcessorTraining::new(MockProcessor); + let interrupter = Interrupter::new(); + let mut progress = Progress { + items_processed: 0, + items_total: 1, + unit: Some("steps".to_string()), + }; + + // Kickstart tests by running some steps to make sure it's not a double batching edge case success. + let steps = runner.run_steps(8, &mut processor, &interrupter, &mut progress); + assert_eq!(steps.len(), 8); + + for i in 0..16 { + let steps = runner.run_steps(i, &mut processor, &interrupter, &mut progress); + assert_eq!(steps.len(), i); + } + } + + // num_envs == autobatch_size + run_test(1, 1); + run_test(4, 4); + // num_envs < autobatch_size + run_test(1, 2); + run_test(1, 3); + run_test(2, 3); + run_test(2, 4); + run_test(5, 19); + // num_envs > autobatch_size + run_test(2, 1); + run_test(8, 1); + run_test(3, 2); + run_test(8, 2); + run_test(8, 3); + run_test(8, 7); + } + + #[test] + fn run_episodes_returns_requested_number_multi_loop() { + fn run_test(num_envs: usize, autobatch_size: usize) { + let mut runner = setup_multi_loop(num_envs, autobatch_size, 0, false, false); + let mut processor = AsyncProcessorTraining::new(MockProcessor); + let interrupter = Interrupter::new(); + let mut progress = Progress { + items_processed: 0, + items_total: 1, + unit: Some("steps".to_string()), + }; + + // Kickstart tests by running some episodes to make sure it's not a double batching edge case success. + let trajectories = runner.run_episodes(8, &mut processor, &interrupter, &mut progress); + assert_eq!(trajectories.len(), 8); + for j in 0..8 { + assert_ne!(trajectories[j].timesteps.len(), 0); + } + + for i in 0..16 { + let trajectories = + runner.run_episodes(i, &mut processor, &interrupter, &mut progress); + assert_eq!(trajectories.len(), i); + for j in 0..i { + assert_ne!(trajectories[j].timesteps.len(), 0); + } + } + } + + // num_envs == autobatch_size + run_test(1, 1); + run_test(4, 4); + // num_envs < autobatch_size + run_test(1, 2); + run_test(1, 3); + run_test(2, 3); + run_test(2, 4); + run_test(5, 19); + // num_envs > autobatch_size + run_test(2, 1); + run_test(8, 1); + run_test(3, 2); + run_test(8, 2); + run_test(8, 3); + run_test(8, 7); + } +} diff --git a/crates/burn-train/src/learner/rl/env_runner/base.rs b/crates/burn-train/src/learner/rl/env_runner/base.rs new file mode 100644 index 0000000..3c3ac4f --- /dev/null +++ b/crates/burn-train/src/learner/rl/env_runner/base.rs @@ -0,0 +1,355 @@ +use burn_core::Tensor; +use burn_core::data::dataloader::Progress; +use burn_core::tensor::Device; +use burn_rl::Policy; +use burn_rl::ToObservation; +use burn_rl::Transition; +use burn_rl::{Environment, EnvironmentInit}; + +use crate::RLEvent; +use crate::{ + AgentEvaluationEvent, EpisodeSummary, EvaluationItem, EventProcessorTraining, + RLEventProcessorType, +}; +use crate::{Interrupter, RLComponentsTypes}; + +/// A trajectory, i.e. a list of ordered [TimeStep](TimeStep). +#[derive(Clone, new)] +pub struct Trajectory { + /// A list of ordered [TimeStep](TimeStep)s. + pub timesteps: Vec>, +} + +/// A timestep debscribing an iteration of the state/decision process. +#[derive(Clone)] +pub struct TimeStep { + /// The environment id. + pub env_id: usize, + /// The [burn_rl::Transition](burn_rl::Transition). + pub transition: Transition, + /// True if the environment reaches a terminal state. + pub done: bool, + /// The running length of the current episode. + pub ep_len: usize, + /// The running cumulative reward. + pub cum_reward: f64, + /// The action's context for this timestep. + pub action_context: C, +} + +pub(crate) type RLTimeStep = TimeStep< + ::State, + ::Action, + ::ActionContext, +>; + +pub(crate) type RLTrajectory = Trajectory< + ::State, + ::Action, + ::ActionContext, +>; + +/// Trait for a structure that implements an agent/environement interface. +pub trait AgentEnvLoop { + /// Run a certain number of timesteps. + /// + /// # Arguments + /// + /// * `num_steps` - The number of time_steps to run. + /// * `processor` - An [crate::EventProcessorTraining](crate::EventProcessorTraining). + /// * `interrupter` - An [crate::Interrupter](crate::Interrupter). + /// * `num_steps` - The number of time_steps to run. + /// * `progress` - A mutable reference to the learning progress. + /// + /// # Returns + /// + /// A list of ordered timesteps. + fn run_steps( + &mut self, + num_steps: usize, + processor: &mut RLEventProcessorType, + interrupter: &Interrupter, + progress: &mut Progress, + ) -> Vec>; + /// Run a certain number of episodes. + /// + /// # Arguments + /// + /// * `num_episodes` - The number of episodes to run. + /// * `processor` - An [crate::EventProcessorTraining](crate::EventProcessorTraining). + /// * `interrupter` - An [crate::Interrupter](crate::Interrupter). + /// * `progress` - A mutable reference to the learning progress. + /// + /// # Returns + /// + /// A list of ordered timesteps. + fn run_episodes( + &mut self, + num_episodes: usize, + processor: &mut RLEventProcessorType, + interrupter: &Interrupter, + progress: &mut Progress, + ) -> Vec>; + /// Update the runner's agent. + fn update_policy(&mut self, update: RLC::PolicyState); + /// Get the state of the runner's agent. + fn policy(&self) -> RLC::PolicyState; + /// Returns the device on which the runner's agent runs. + fn device(&self) -> Device; +} + +/// A simple, synchronized agent/environement interface. +pub struct AgentEnvBaseLoop { + env: RLC::Env, + eval: bool, + agent: RLC::Policy, + deterministic: bool, + current_reward: f64, + run_num: usize, + step_num: usize, + device: Device, +} + +impl AgentEnvBaseLoop { + /// Create a new base runner. + pub fn new( + env_init: RLC::EnvInit, + agent: RLC::Policy, + eval: bool, + deterministic: bool, + device: &Device, + ) -> Self { + let agent = agent.to_device(device); + let mut env = env_init.init(); + env.reset(); + + Self { + env, + eval, + agent: agent.clone(), + deterministic, + current_reward: 0.0, + run_num: 0, + step_num: 0, + device: device.clone(), + } + } +} + +impl AgentEnvLoop for AgentEnvBaseLoop +where + RLC: RLComponentsTypes, +{ + fn run_steps( + &mut self, + num_steps: usize, + processor: &mut RLEventProcessorType, + interrupter: &Interrupter, + progress: &mut Progress, + ) -> Vec> { + let mut items = vec![]; + let device = Default::default(); + for _ in 0..num_steps { + let state = self.env.state(); + let (action, context) = self.agent.action( + state.clone().to_observation(&self.device), + self.deterministic, + ); + + let step_result = self.env.step(RLC::Action::from(action.clone())); + + self.current_reward += step_result.reward; + self.step_num += 1; + + let transition = Transition::new( + state.clone(), + step_result.next_state, + RLC::Action::from(action), + Tensor::from_data([step_result.reward], &device), + Tensor::from_data( + [(step_result.done || step_result.truncated) as i32 as f64], + &device, + ), + ); + items.push(TimeStep { + env_id: 0, + transition, + done: step_result.done, + ep_len: self.step_num, + cum_reward: self.current_reward, + action_context: context[0].clone(), + }); + + if !self.eval { + progress.items_processed += 1; + processor.process_train(RLEvent::EnvStep(EvaluationItem::new( + context[0].clone(), + progress.clone(), + None, + ))); + + if step_result.done { + processor.process_train(RLEvent::EpisodeEnd(EvaluationItem::new( + EpisodeSummary { + episode_length: self.step_num, + cum_reward: self.current_reward, + }, + progress.clone(), + None, + ))); + } + } + + if interrupter.should_stop() { + break; + } + + if step_result.done || step_result.truncated { + self.env.reset(); + self.current_reward = 0.; + self.step_num = 0; + self.run_num += 1; + } + } + items + } + + fn update_policy(&mut self, update: RLC::PolicyState) { + self.agent.update(update); + } + + fn run_episodes( + &mut self, + num_episodes: usize, + processor: &mut RLEventProcessorType, + interrupter: &Interrupter, + progress: &mut Progress, + ) -> Vec> { + self.env.reset(); + + let mut items = vec![]; + for ep in 0..num_episodes { + let mut steps = vec![]; + loop { + let step = self.run_steps(1, processor, interrupter, progress)[0].clone(); + steps.push(step.clone()); + + if self.eval { + processor.process_valid(AgentEvaluationEvent::EnvStep(EvaluationItem::new( + step.action_context.clone(), + Progress::new(steps.len() + 1, steps.len() + 1, Some("steps".to_string())), + None, + ))); + + if step.done { + processor.process_valid(AgentEvaluationEvent::EpisodeEnd( + EvaluationItem::new( + EpisodeSummary { + episode_length: step.ep_len, + cum_reward: step.cum_reward, + }, + Progress::new(ep + 1, num_episodes, Some("episodes".to_string())), + None, + ), + )); + } + } + + if interrupter.should_stop() || step.done { + break; + } + } + items.push(Trajectory::new(steps)); + + if interrupter.should_stop() { + break; + } + } + items + } + + fn policy(&self) -> RLC::PolicyState { + self.agent.state() + } + + fn device(&self) -> Device { + self.device.clone() + } +} + +#[cfg(test)] +#[allow(clippy::needless_range_loop)] +mod tests { + use crate::AsyncProcessorTraining; + + use crate::learner::tests::{ + MockEnvInit, MockPolicy, MockPolicyState, MockProcessor, MockRLComponents, + }; + + use super::*; + + fn setup(state: usize, eval: bool, deterministic: bool) -> AgentEnvBaseLoop { + let env_init = MockEnvInit; + let agent = MockPolicy(state); + AgentEnvBaseLoop::::new( + env_init, + agent, + eval, + deterministic, + &Default::default(), + ) + } + + #[test] + fn test_policy_returns_agent_state() { + let runner = setup(1000, false, false); + let policy_state = runner.policy(); + assert_eq!(policy_state.0, 1000); + } + + #[test] + fn test_update_policy() { + let mut runner = setup(0, false, false); + + runner.update_policy(MockPolicyState(1)); + assert_eq!(runner.policy().0, 1); + } + + #[test] + fn run_steps_returns_requested_number() { + let mut runner = setup(0, false, false); + let mut processor = AsyncProcessorTraining::new(MockProcessor); + let interrupter = Interrupter::new(); + let mut progress = Progress { + items_processed: 0, + items_total: 1, + unit: Some("steps".to_string()), + }; + + let steps = runner.run_steps(1, &mut processor, &interrupter, &mut progress); + assert_eq!(steps.len(), 1); + let steps = runner.run_steps(8, &mut processor, &interrupter, &mut progress); + assert_eq!(steps.len(), 8); + } + + #[test] + fn run_episodes_returns_requested_number() { + let mut runner = setup(0, false, false); + let mut processor = AsyncProcessorTraining::new(MockProcessor); + let interrupter = Interrupter::new(); + let mut progress = Progress { + items_processed: 0, + items_total: 1, + unit: Some("steps".to_string()), + }; + + let trajectories = runner.run_episodes(1, &mut processor, &interrupter, &mut progress); + assert_eq!(trajectories.len(), 1); + assert_ne!(trajectories[0].timesteps.len(), 0); + let trajectories = runner.run_episodes(8, &mut processor, &interrupter, &mut progress); + assert_eq!(trajectories.len(), 8); + for i in 0..8 { + assert_ne!(trajectories[i].timesteps.len(), 0); + } + } +} diff --git a/crates/burn-train/src/learner/rl/env_runner/mod.rs b/crates/burn-train/src/learner/rl/env_runner/mod.rs new file mode 100644 index 0000000..8eacd49 --- /dev/null +++ b/crates/burn-train/src/learner/rl/env_runner/mod.rs @@ -0,0 +1,302 @@ +mod async_runner; +mod base; + +pub use async_runner::*; +pub use base::*; + +#[cfg(test)] +pub(crate) mod tests { + use burn_rl::{ + Batchable, Environment, EnvironmentInit, Policy, PolicyState, ToAction, ToObservation, + }; + + use crate::{ + AgentEvaluationEvent, EventProcessorTraining, ItemLazy, RLComponentsTypes, RLEvent, + }; + use burn_rl::{LearnerTransitionBatch, PolicyLearner, RLTrainOutput, StepResult}; + + /// Mock policy for testing + /// + /// Calling `forward()` with a [MockObservation](MockObservation) (list of f32) returns a [MockActionDistribution](MockActionDistribution) + /// containing a list of 0s of the same length as the observation. + /// + /// Calling `action()` with a [MockObservation](MockObservation) (list of f32) returns a [MockPolicyAction](MockPolicyAction) with a list of actions of the same length as the observation. + /// The actions are all 1 if the call is requested as deterministic, or else 0. + + #[derive(Clone)] + pub(crate) struct MockPolicy(pub usize); + + impl Policy for MockPolicy { + type Observation = MockObservation; + type ActionDistribution = MockActionDistribution; + type Action = MockPolicyAction; + type ActionContext = MockActionContext; + type PolicyState = MockPolicyState; + + fn forward(&mut self, obs: Self::Observation) -> Self::ActionDistribution { + let mut dists = vec![]; + for _ in obs.0 { + dists.push(MockActionDistribution(vec![0.])); + } + MockActionDistribution::batch(dists) + } + + fn action( + &mut self, + obs: Self::Observation, + deterministic: bool, + ) -> (Self::Action, Vec) { + let mut actions = vec![]; + let mut contexts = vec![]; + + for _ in obs.0 { + if deterministic { + actions.push(MockPolicyAction(vec![1])); + } else { + actions.push(MockPolicyAction(vec![0])); + } + contexts.push(MockActionContext); + } + + (MockPolicyAction::batch(actions), contexts) + } + + fn update(&mut self, update: Self::PolicyState) { + self.0 = update.0; + } + + fn state(&self) -> Self::PolicyState { + MockPolicyState(self.0) + } + + fn to_device(self, _device: &burn_core::prelude::Device) -> Self { + self + } + + fn load_record(self, _record: ::Record) -> Self { + self + } + } + + /// Mock observation for testing represented as a vector of f32. Can call `batch()` and `unbatch` on it. + #[derive(Clone)] + pub(crate) struct MockObservation(pub Vec); + + /// Mock action for testing represented as a vector of i32. Can call `batch()` and `unbatch` on it. + #[derive(Clone)] + pub(crate) struct MockPolicyAction(pub Vec); + + /// Mock action distribution for testing represented as a vector of i32. Can call `batch()` and `unbatch` on it. + #[derive(Clone)] + pub(crate) struct MockActionDistribution(Vec); + + #[derive(Clone)] + pub(crate) struct MockActionContext; + + /// Mock policy state for testing represented as an arbitrary `usize` that has no effect on the policy. + #[derive(Clone)] + pub(crate) struct MockPolicyState(pub usize); + + impl PolicyState for MockPolicyState { + type Record = (); + + fn into_record(self) -> Self::Record {} + + fn load_record(&self, _record: Self::Record) -> Self { + self.clone() + } + } + + impl Batchable for MockObservation { + fn batch(items: Vec) -> Self { + MockObservation(items.iter().flat_map(|m| m.0.clone()).collect()) + } + + fn unbatch(self) -> Vec { + vec![MockObservation(self.0)] + } + } + + impl Batchable for MockPolicyAction { + fn batch(items: Vec) -> Self { + MockPolicyAction(items.iter().flat_map(|m| m.0.clone()).collect()) + } + + fn unbatch(self) -> Vec { + let mut actions = vec![]; + for a in self.0 { + actions.push(MockPolicyAction(vec![a])); + } + actions + } + } + + impl Batchable for MockActionDistribution { + fn batch(items: Vec) -> Self { + MockActionDistribution(items.iter().flat_map(|m| m.0.clone()).collect()) + } + + fn unbatch(self) -> Vec { + let mut dists = vec![]; + for _ in self.0 { + dists.push(MockActionDistribution(vec![0.])); + } + dists + } + } + + /// Mock environment for testing + #[derive(Clone)] + pub(crate) struct MockEnv { + counter: usize, + } + + #[derive(Clone, Debug)] + pub(crate) struct MockState; + + #[derive(Clone, Debug)] + pub(crate) struct MockAction(pub i32); + + impl ToObservation for MockState { + fn to_observation(&self, _device: &burn_core::prelude::Device) -> MockObservation { + MockObservation(vec![0.]) + } + } + + impl From for MockAction { + fn from(value: MockPolicyAction) -> Self { + MockAction(value.0[0]) + } + } + + impl ToAction for MockAction { + fn to_action(&self, _device: &burn_core::prelude::Device) -> MockPolicyAction { + MockPolicyAction(vec![self.0]) + } + } + + impl ItemLazy for MockActionContext { + fn sync(self) -> Self { + self + } + } + + impl MockEnv { + fn new() -> Self { + Self { counter: 0 } + } + } + + impl Environment for MockEnv { + type State = MockState; + type Action = MockAction; + const MAX_STEPS: usize = 5; + + fn reset(&mut self) { + self.counter = 0; + } + + fn step(&mut self, _action: Self::Action) -> StepResult { + self.counter += 1; + let done = self.counter >= Self::MAX_STEPS; + + burn_rl::StepResult { + next_state: MockState, + reward: 1.0, + done, + truncated: false, + } + } + + fn state(&self) -> Self::State { + MockState + } + } + + /// Mock environment init for testing + #[derive(Clone)] + pub(crate) struct MockEnvInit; + + impl EnvironmentInit for MockEnvInit { + fn init(&self) -> MockEnv { + MockEnv::new() + } + } + + // Mock RLComponentsTypes for testing + pub(crate) struct MockRLComponents; + + impl RLComponentsTypes for MockRLComponents { + type Env = MockEnv; + type EnvInit = MockEnvInit; + type State = MockState; + type Action = MockAction; + type Policy = MockPolicy; + type PolicyObs = MockObservation; + type PolicyAD = MockActionDistribution; + type PolicyAction = MockPolicyAction; + type ActionContext = MockActionContext; + type PolicyState = MockPolicyState; + type LearningAgent = MockLearningAgent; + type TrainingOutput = (); + } + + // Mock learning agent for testing + #[derive(Clone)] + pub(crate) struct MockLearningAgent; + + impl PolicyLearner for MockLearningAgent { + type InnerPolicy = MockPolicy; + type TrainContext = (); + type Record = (); + + fn train( + &mut self, + _input: LearnerTransitionBatch, + ) -> RLTrainOutput::PolicyState> { + unimplemented!() + } + + fn policy(&self) -> Self::InnerPolicy { + unimplemented!() + } + + fn update_policy(&mut self, _update: Self::InnerPolicy) { + unimplemented!() + } + + fn record(&self) -> Self::Record { + unimplemented!() + } + + fn load_record(self, _record: Self::Record) -> Self { + unimplemented!() + } + + fn device(&self) -> burn_core::prelude::Device { + unimplemented!() + } + } + + // Mock event processor for testing + pub(crate) struct MockProcessor; + + impl + EventProcessorTraining< + RLEvent<(), MockActionContext>, + AgentEvaluationEvent, + > for MockProcessor + { + fn process_train(&mut self, _event: RLEvent<(), MockActionContext>) { + // Mock process train + } + + fn process_valid(&mut self, _event: AgentEvaluationEvent) { + // Mock process valid + } + + fn renderer(self) -> Box { + unimplemented!() + } + } +} diff --git a/crates/burn-train/src/learner/rl/mod.rs b/crates/burn-train/src/learner/rl/mod.rs new file mode 100644 index 0000000..34874dc --- /dev/null +++ b/crates/burn-train/src/learner/rl/mod.rs @@ -0,0 +1,15 @@ +mod checkpointer; +mod components; +mod env_runner; +mod off_policy; +mod output; +mod paradigm; +mod strategy; + +pub use checkpointer::*; +pub use components::*; +pub use env_runner::*; +pub use off_policy::*; +pub use output::*; +pub use paradigm::*; +pub use strategy::*; diff --git a/crates/burn-train/src/learner/rl/off_policy.rs b/crates/burn-train/src/learner/rl/off_policy.rs new file mode 100644 index 0000000..f480b7e --- /dev/null +++ b/crates/burn-train/src/learner/rl/off_policy.rs @@ -0,0 +1,190 @@ +use crate::{ + AgentEnvAsyncLoop, AgentEnvLoop, AsyncAgentEnvLoopConfig, EvaluationItem, + EventProcessorTraining, MultiAgentEnvLoop, RLComponents, RLComponentsTypes, RLEvent, + RLEventProcessorType, RLStrategy, +}; +use burn_core::{self as burn}; +use burn_core::{config::Config, data::dataloader::Progress}; +use burn_rl::{ + AsyncPolicy, Policy, PolicyLearner, SliceAccess, ToAction, ToObservation, TransitionBuffer, +}; + +/// Parameters of an on policy training with multi environments and double-batching. +#[derive(Config, Debug)] +pub struct OffPolicyConfig { + /// The number of environments to run simultaneously for experience collection. + #[config(default = 1)] + pub num_envs: usize, + /// Number of environment state to accumulate before running one step of inference with the policy. + /// Must be equal or less than the number of simultaneous environments. + #[config(default = 1)] + pub autobatch_size: usize, + /// Max number of transitions stored in the replay buffer. + #[config(default = 1024)] + pub replay_buffer_size: usize, + /// The number of steps to collect between each step of training. + #[config(default = 1)] + pub train_interval: usize, + /// Number of optimization steps done each `train_interval`. + #[config(default = 1)] + pub train_steps: usize, + /// The number of steps to collect between each evaluation. + #[config(default = 10_000)] + pub eval_interval: usize, + /// The number of episodes to run for each evaluation. + #[config(default = 1)] + pub eval_episodes: usize, + /// The number of transition to train on. + #[config(default = 32)] + pub train_batch_size: usize, + /// Number of steps to collect before starting to train. + #[config(default = 0)] + pub warmup_steps: usize, +} + +/// Off-policy reinforcement learning strategy with multi-env experience collection and double-batching. +pub struct OffPolicyStrategy { + config: OffPolicyConfig, +} +impl OffPolicyStrategy { + /// Create a new off-policy base strategy. + pub fn new(config: OffPolicyConfig) -> Self { + Self { config } + } +} + +impl RLStrategy for OffPolicyStrategy +where + RLC: RLComponentsTypes, + RLC::PolicyObs: SliceAccess, + RLC::PolicyAction: SliceAccess, +{ + fn train_loop( + &self, + training_components: RLComponents, + learner_agent: &mut RLC::LearningAgent, + starting_epoch: usize, + env_init: RLC::EnvInit, + ) -> (RLC::Policy, RLEventProcessorType) { + let mut event_processor = training_components.event_processor; + let mut checkpointer = training_components.checkpointer; + let num_steps_total = training_components.num_steps; + + let inference_device = training_components.inference_device; + let mut env_runner = MultiAgentEnvLoop::::new( + self.config.num_envs, + env_init.clone(), + AsyncPolicy::new( + self.config.num_envs.min(self.config.autobatch_size), + learner_agent.policy(), + ), + false, + false, + &inference_device, + ); + let runner_config = AsyncAgentEnvLoopConfig { + eval: true, + deterministic: true, + id: 0, + }; + let mut env_runner_valid = AgentEnvAsyncLoop::::new( + env_init, + AsyncPolicy::new(1, learner_agent.policy()), + runner_config, + &inference_device, + None, + None, + ); + + let mut transition_buffer = TransitionBuffer::::new( + self.config.replay_buffer_size, + &learner_agent.device(), + ); + + let mut valid_next = self.config.eval_interval + starting_epoch - 1; + let mut progress = Progress { + items_processed: starting_epoch, + items_total: num_steps_total, + unit: Some("steps".to_string()), + }; + + let mut intermediary_update: Option<::PolicyState> = None; + while progress.items_processed < num_steps_total { + if training_components.interrupter.should_stop() { + let reason = training_components + .interrupter + .get_message() + .unwrap_or(String::from("Reason unknown")); + log::info!("Training interrupted: {reason}"); + break; + } + + let previous_steps = progress.items_processed; + let items = env_runner.run_steps( + self.config.train_interval, + &mut event_processor, + &training_components.interrupter, + &mut progress, + ); + + for item in &items { + let t = &item.transition; + let state: RLC::PolicyObs = t.state.clone().to_observation(&env_runner.device()); + let next_state: RLC::PolicyObs = + t.next_state.clone().to_observation(&env_runner.device()); + let action: RLC::PolicyAction = t.action.clone().to_action(&env_runner.device()); + let reward = t.reward.to_data().to_vec::().unwrap()[0]; + let done = t.done.to_data().to_vec::().unwrap()[0] > 0.5; + transition_buffer.push(state, next_state, action, reward, done); + } + + if transition_buffer.len() >= self.config.train_batch_size + && progress.items_processed >= self.config.warmup_steps + { + if let Some(ref u) = intermediary_update { + env_runner.update_policy(u.clone()); + } + for _ in 0..self.config.train_steps { + let batch = transition_buffer.sample(self.config.train_batch_size); + let train_item = learner_agent.train(batch); + intermediary_update = Some(learner_agent.policy().state()); + + event_processor.process_train(RLEvent::TrainStep(EvaluationItem::new( + train_item.item, + progress.clone(), + None, + ))); + } + } + + if valid_next > previous_steps && valid_next <= progress.items_processed { + event_processor.process_valid(crate::AgentEvaluationEvent::Start( + self.config.eval_episodes, + )); + + env_runner_valid.update_policy(learner_agent.policy().state()); + env_runner_valid.run_episodes( + self.config.eval_episodes, + &mut event_processor, + &training_components.interrupter, + &mut progress, + ); + + if let Some(checkpointer) = &mut checkpointer { + checkpointer.checkpoint( + &env_runner.policy(), + learner_agent, + valid_next, + &training_components.event_store, + ); + } + + valid_next += self.config.eval_interval; + + event_processor.process_valid(crate::AgentEvaluationEvent::End); + } + } + + (learner_agent.policy(), event_processor) + } +} diff --git a/crates/burn-train/src/learner/rl/output.rs b/crates/burn-train/src/learner/rl/output.rs new file mode 100644 index 0000000..238f6c0 --- /dev/null +++ b/crates/burn-train/src/learner/rl/output.rs @@ -0,0 +1,30 @@ +use crate::{ + ItemLazy, + metric::{Adaptor, CumulativeRewardInput, EpisodeLengthInput}, +}; + +/// Summary of an episode. +pub struct EpisodeSummary { + /// The total length of the episode. + pub episode_length: usize, + /// The final cumulative reward. + pub cum_reward: f64, +} + +impl ItemLazy for EpisodeSummary { + fn sync(self) -> Self { + self + } +} + +impl Adaptor for EpisodeSummary { + fn adapt(&self) -> EpisodeLengthInput { + EpisodeLengthInput::new(self.episode_length as f64) + } +} + +impl Adaptor for EpisodeSummary { + fn adapt(&self) -> CumulativeRewardInput { + CumulativeRewardInput::new(self.cum_reward) + } +} diff --git a/crates/burn-train/src/learner/rl/paradigm.rs b/crates/burn-train/src/learner/rl/paradigm.rs new file mode 100644 index 0000000..4f7384e --- /dev/null +++ b/crates/burn-train/src/learner/rl/paradigm.rs @@ -0,0 +1,542 @@ +use crate::checkpoint::{ + AsyncCheckpointer, Checkpoint, CheckpointingStrategy, ComposedCheckpointingStrategy, + FileCheckpointer, KeepLastNCheckpoints, MetricCheckpointingStrategy, +}; +use crate::learner::base::Interrupter; +use crate::logger::{FileMetricLogger, MetricLogger}; +use crate::metric::store::{Aggregate, Direction, EventStoreClient, LogEventStore, Split}; +use crate::metric::{Adaptor, EpisodeLengthMetric, Metric, Numeric}; +use crate::renderer::{MetricsRenderer, default_renderer}; +use crate::{ + ApplicationLoggerInstaller, AsyncProcessorTraining, FileApplicationLoggerInstaller, ItemLazy, + LearnerSummaryConfig, OffPolicyConfig, OffPolicyStrategy, RLAgentRecord, RLCheckpointer, + RLComponents, RLComponentsMarker, RLComponentsTypes, RLEventProcessor, RLMetrics, + RLPolicyRecord, RLStrategy, +}; +use crate::{EpisodeSummary, RLStrategies}; +use burn_core::tensor::Device; +use burn_rl::{ + Batchable, Environment, EnvironmentInit, Policy, PolicyLearner, PolicyState, SliceAccess, + ToAction, ToObservation, +}; +use std::collections::BTreeSet; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +/// Structure to configure and launch reinforcement learning trainings. +pub struct RLTraining { + // Not that complex. Extracting into yet another type would only make it more confusing. + #[allow(clippy::type_complexity)] + checkpointers: Option<( + AsyncCheckpointer>, + AsyncCheckpointer>, + )>, + num_steps: usize, + checkpoint: Option, + directory: PathBuf, + grad_accumulation: Option, + renderer: Option>, + metrics: RLMetrics, + event_store: LogEventStore, + interrupter: Interrupter, + tracing_logger: Option>, + checkpointer_strategy: Box, + learning_strategy: RLStrategies, + // Use BTreeSet instead of HashSet for consistent (alphabetical) iteration order + summary_metrics: BTreeSet, + summary: bool, + env_initializer: RLC::EnvInit, + inference_device: Device, +} + +impl RLTraining> +where + E: Environment + 'static, + EI: EnvironmentInit + Send + 'static, + A: PolicyLearner + Send + 'static, + ::Record: Checkpoint, + A::TrainContext: ItemLazy + Clone + Send, + A::InnerPolicy: Policy + Send, + ::Observation: Batchable + Clone + Send, + ::ActionDistribution: Batchable + Clone + Send, + ::Action: Batchable + Clone + Send, + ::ActionContext: ItemLazy + Clone + Send + 'static, + ::PolicyState: Clone + Send, + <::PolicyState as PolicyState>::Record: Checkpoint, + E::State: ToObservation<::Observation> + Clone + Send + 'static, + E::Action: From<::Action> + + ToAction<::Action> + + Clone + + Send + + 'static, +{ + /// Creates a new runner for reinforcement learning. + /// + /// # Arguments + /// + /// * `directory` - The directory to save the checkpoints. + /// * `env_init` - Specifies how to initialize the environment. + pub fn new(directory: impl AsRef, env_initializer: EI) -> Self { + let directory = directory.as_ref().to_path_buf(); + let experiment_log_file = directory.join("experiment.log"); + Self { + num_steps: 1, + checkpoint: None, + checkpointers: None, + directory, + grad_accumulation: None, + metrics: RLMetrics::default(), + event_store: LogEventStore::default(), + renderer: None, + interrupter: Interrupter::new(), + tracing_logger: Some(Box::new(FileApplicationLoggerInstaller::new( + experiment_log_file, + ))), + checkpointer_strategy: Box::new( + ComposedCheckpointingStrategy::builder() + .add(KeepLastNCheckpoints::new(2)) + .add(MetricCheckpointingStrategy::new( + &EpisodeLengthMetric::new(), // default to evaluations' cumulative reward. + Aggregate::Mean, + Direction::Lowest, + Split::Valid, + )) + .build(), + ), + learning_strategy: RLStrategies::OffPolicyStrategy(OffPolicyConfig::new()), + summary_metrics: BTreeSet::new(), + summary: false, + env_initializer, + inference_device: Default::default(), + } + } +} + +impl RLTraining { + /// Replace the default learning strategy (Off Policy learning) with the provided one. + /// + /// # Arguments + /// + /// * `training_strategy` - The training strategy. + pub fn with_learning_strategy(mut self, learning_strategy: RLStrategies) -> Self { + self.learning_strategy = learning_strategy; + self + } + + /// Replace the default metric loggers with the provided ones. + /// + /// # Arguments + /// + /// * `logger` - The training logger. + pub fn with_metric_logger(mut self, logger: ML) -> Self + where + ML: MetricLogger + 'static, + { + self.event_store.register_logger(logger); + self + } + + /// Update the checkpointing_strategy. + pub fn with_checkpointing_strategy( + mut self, + strategy: CS, + ) -> Self { + self.checkpointer_strategy = Box::new(strategy); + self + } + + /// Replace the default CLI renderer with a custom one. + /// + /// # Arguments + /// + /// * `renderer` - The custom renderer. + pub fn renderer(mut self, renderer: MR) -> Self + where + MR: MetricsRenderer + 'static, + { + self.renderer = Some(Box::new(renderer)); + self + } + + /// Register numerical metrics for a training step of the agent. + pub fn metrics_train>(self, metrics: Me) -> Self { + metrics.register(self) + } + + /// Register textual metrics for a training step of the agent. + pub fn text_metrics_train>(self, metrics: Me) -> Self { + metrics.register(self) + } + + /// Register numerical metrics for each action of the agent. + pub fn metrics_agent>(self, metrics: Me) -> Self { + metrics.register(self) + } + + /// Register textual metrics for each action of the agent. + pub fn text_metrics_agent>(self, metrics: Me) -> Self { + metrics.register(self) + } + + /// Register numerical metrics for a completed episode. + pub fn metrics_episode>(self, metrics: Me) -> Self { + metrics.register(self) + } + + /// Register textual metrics for a completed episode. + pub fn text_metrics_episode>(self, metrics: Me) -> Self { + metrics.register(self) + } + + /// Register a textual metric for a training step. + pub fn text_metric_train(mut self, metric: Me) -> Self + where + RLC::TrainingOutput: Adaptor, + { + self.metrics.register_text_metric_train(metric); + self + } + + /// Register a [numeric](crate::metric::Numeric) [metric](Metric) for a training step. + pub fn metric_train(mut self, metric: Me) -> Self + where + Me: Metric + Numeric + 'static, + RLC::TrainingOutput: Adaptor, + { + self.summary_metrics.insert(metric.name().to_string()); + self.metrics.register_metric_train(metric); + self + } + + /// Register a textual metric for each action taken by the agent. + pub fn text_metric_agent(mut self, metric: Me) -> Self + where + RLC::ActionContext: Adaptor, + { + self.metrics.register_text_metric_agent(metric.clone()); + self.metrics.register_text_metric_agent_valid(metric); + self + } + + /// Register a [numeric](crate::metric::Numeric) [metric](Metric) for each action taken by the agent. + pub fn metric_agent(mut self, metric: Me) -> Self + where + Me: Metric + Numeric + 'static, + RLC::ActionContext: Adaptor, + { + self.summary_metrics.insert(metric.name().to_string()); + self.metrics.register_agent_metric(metric.clone()); + self.metrics.register_agent_metric_valid(metric); + self + } + + /// Register a textual metric for a completed episode. + pub fn text_metric_episode(mut self, metric: Me) -> Self + where + EpisodeSummary: Adaptor + 'static, + { + self.metrics.register_text_metric_episode(metric.clone()); + self.metrics.register_text_metric_episode_valid(metric); + self + } + + /// Register a [numeric](crate::metric::Numeric) [metric](Metric) for a completed episode. + pub fn metric_episode(mut self, metric: Me) -> Self + where + Me: Metric + Numeric + 'static, + EpisodeSummary: Adaptor + 'static, + { + self.summary_metrics.insert(metric.name().to_string()); + self.metrics.register_episode_metric(metric.clone()); + self.metrics.register_episode_metric_valid(metric); + self + } + + /// The number of environment steps to train for. + pub fn num_steps(mut self, num_steps: usize) -> Self { + self.num_steps = num_steps; + self + } + + /// The step from which the training must resume. + pub fn checkpoint(mut self, checkpoint: usize) -> Self { + self.checkpoint = Some(checkpoint); + self + } + + /// Provides a handle that can be used to interrupt training. + pub fn interrupter(&self) -> Interrupter { + self.interrupter.clone() + } + + /// Override the handle for stopping training with an externally provided handle + pub fn with_interrupter(mut self, interrupter: Interrupter) -> Self { + self.interrupter = interrupter; + self + } + + /// By default, Rust logs are captured and written into + /// `experiment.log`. If disabled, standard Rust log handling + /// will apply. + pub fn with_application_logger( + mut self, + logger: Option>, + ) -> Self { + self.tracing_logger = logger; + self + } + + /// Register a checkpointer that will save the environment runner's [policy](Policy) + /// and the [PolicyLearner](PolicyLearner) state to different files. + pub fn with_checkpointer(mut self) -> Self + where + RLPolicyRecord: Checkpoint, + RLAgentRecord: Checkpoint, + { + let checkpoint_dir = self.directory.join("checkpoint"); + let checkpointer_policy = FileCheckpointer::new(&checkpoint_dir, "policy"); + let checkpointer_learning = FileCheckpointer::new(&checkpoint_dir, "learning-agent"); + + self.checkpointers = Some(( + AsyncCheckpointer::new(checkpointer_policy), + AsyncCheckpointer::new(checkpointer_learning), + )); + + self + } + + /// The device on which to run inference during rollout collection and validation. + pub fn with_inference_device(mut self, device: Device) -> Self { + self.inference_device = device; + self + } + + /// Enable the training summary report. + /// + /// The summary will be displayed after `.launch()`, when the renderer is dropped. + pub fn summary(mut self) -> Self { + self.summary = true; + self + } + + /// Launch the training with the specified [PolicyLearner](PolicyLearner) on the specified environment. + pub fn launch(mut self, learner_agent: RLC::LearningAgent) -> RLResult + where + RLC::PolicyObs: SliceAccess, + RLC::PolicyAction: SliceAccess, + { + if self.tracing_logger.is_some() + && let Err(e) = self.tracing_logger.as_ref().unwrap().install() + { + log::warn!("Failed to install the experiment logger: {e}"); + } + let renderer = self + .renderer + .unwrap_or_else(|| default_renderer(self.interrupter.clone(), self.checkpoint)); + + if !self.event_store.has_loggers() { + self.event_store + .register_logger(FileMetricLogger::new(self.directory.clone())); + } + + let event_store = Arc::new(EventStoreClient::new(self.event_store)); + let event_processor = AsyncProcessorTraining::new(RLEventProcessor::new( + self.metrics, + renderer, + event_store.clone(), + )); + + let checkpointer = self.checkpointers.map(|(policy, learning_agent)| { + RLCheckpointer::new(policy, learning_agent, self.checkpointer_strategy) + }); + + let summary = if self.summary { + Some(LearnerSummaryConfig { + directory: self.directory, + metrics: self.summary_metrics.into_iter().collect::>(), + }) + } else { + None + }; + + let components = RLComponents:: { + checkpoint: self.checkpoint, + checkpointer, + interrupter: self.interrupter, + event_processor, + event_store, + num_steps: self.num_steps, + grad_accumulation: self.grad_accumulation, + summary, + inference_device: self.inference_device, + }; + + let mut learner_agent = learner_agent; + if let Some(checkpoint) = components.checkpoint + && let Some(checkpointer) = &components.checkpointer + { + learner_agent = checkpointer.load_checkpoint(learner_agent, checkpoint); + } + + match self.learning_strategy { + RLStrategies::OffPolicyStrategy(config) => { + let strategy = OffPolicyStrategy::new(config); + strategy.train(learner_agent, components, self.env_initializer) + } + RLStrategies::Custom(strategy) => { + strategy.train(learner_agent, components, self.env_initializer) + } + } + } +} + +/// The result of reinforcement learning, containing the final policy along with the [renderer](MetricsRenderer). +pub struct RLResult

{ + /// The learned policy. + pub policy: P, + /// The renderer that can be used for follow up training and evaluation. + pub renderer: Box, +} + +/// Trait to fake variadic generics for train step metrics. +pub trait AgentMetricRegistration: Sized { + /// Register the metrics. + fn register(self, builder: RLTraining) -> RLTraining; +} + +/// Trait to fake variadic generics for train step text metrics. +pub trait AgentTextMetricRegistration: Sized { + /// Register the metrics. + fn register(self, builder: RLTraining) -> RLTraining; +} + +/// Trait to fake variadic generics for env step metrics. +pub trait TrainMetricRegistration: Sized { + /// Register the metrics. + fn register(self, builder: RLTraining) -> RLTraining; +} + +/// Trait to fake variadic generics for env step text metrics. +pub trait TrainTextMetricRegistration: Sized { + /// Register the metrics. + fn register(self, builder: RLTraining) -> RLTraining; +} + +/// Trait to fake variadic generics for episode metrics. +pub trait EpisodeMetricRegistration: Sized { + /// Register the metrics. + fn register(self, builder: RLTraining) -> RLTraining; +} + +/// Trait to fake variadic generics for episode text metrics. +pub trait EpisodeTextMetricRegistration: Sized { + /// Register the metrics. + fn register(self, builder: RLTraining) -> RLTraining; +} + +macro_rules! gen_tuple { + ($($M:ident),*) => { + impl<$($M,)* RLC: RLComponentsTypes + 'static> TrainTextMetricRegistration for ($($M,)*) + where + $(RLC::TrainingOutput: Adaptor<$M::Input>,)* + $($M: Metric + 'static,)* + { + #[allow(non_snake_case)] + fn register( + self, + builder: RLTraining, + ) -> RLTraining { + let ($($M,)*) = self; + $(let builder = builder.text_metric_train($M.clone());)* + builder + } + } + + impl<$($M,)* RLC: RLComponentsTypes + 'static> TrainMetricRegistration for ($($M,)*) + where + $(RLC::TrainingOutput: Adaptor<$M::Input>,)* + $($M: Metric + Numeric + 'static,)* + { + #[allow(non_snake_case)] + fn register( + self, + builder: RLTraining, + ) -> RLTraining { + let ($($M,)*) = self; + $(let builder = builder.metric_train($M.clone());)* + builder + } + } + + impl<$($M,)* RLC: RLComponentsTypes + 'static> AgentTextMetricRegistration for ($($M,)*) + where + $(RLC::ActionContext: Adaptor<$M::Input>,)* + $($M: Metric + 'static,)* + { + #[allow(non_snake_case)] + fn register( + self, + builder: RLTraining, + ) -> RLTraining { + let ($($M,)*) = self; + $(let builder = builder.text_metric_agent($M.clone());)* + builder + } + } + + impl<$($M,)* RLC: RLComponentsTypes + 'static> AgentMetricRegistration for ($($M,)*) + where + $(RLC::ActionContext: Adaptor<$M::Input>,)* + $($M: Metric + Numeric + 'static,)* + { + #[allow(non_snake_case)] + fn register( + self, + builder: RLTraining, + ) -> RLTraining { + let ($($M,)*) = self; + $(let builder = builder.metric_agent($M.clone());)* + builder + } + } + + impl<$($M,)* RLC: RLComponentsTypes + 'static> EpisodeTextMetricRegistration for ($($M,)*) + where + $(EpisodeSummary: Adaptor<$M::Input> + 'static,)* + $($M: Metric + 'static,)* + { + #[allow(non_snake_case)] + fn register( + self, + builder: RLTraining, + ) -> RLTraining { + let ($($M,)*) = self; + $(let builder = builder.text_metric_episode($M.clone());)* + builder + } + } + + impl<$($M,)* RLC: RLComponentsTypes + 'static> EpisodeMetricRegistration for ($($M,)*) + where + $(EpisodeSummary: Adaptor<$M::Input> + 'static,)* + $($M: Metric + Numeric + 'static,)* + { + #[allow(non_snake_case)] + fn register( + self, + builder: RLTraining, + ) -> RLTraining { + let ($($M,)*) = self; + $(let builder = builder.metric_episode($M.clone());)* + builder + } + } + }; +} + +gen_tuple!(M1); +gen_tuple!(M1, M2); +gen_tuple!(M1, M2, M3); +gen_tuple!(M1, M2, M3, M4); +gen_tuple!(M1, M2, M3, M4, M5); +gen_tuple!(M1, M2, M3, M4, M5, M6); diff --git a/crates/burn-train/src/learner/rl/strategy.rs b/crates/burn-train/src/learner/rl/strategy.rs new file mode 100644 index 0000000..0f4a6a0 --- /dev/null +++ b/crates/burn-train/src/learner/rl/strategy.rs @@ -0,0 +1,92 @@ +use std::sync::Arc; + +use burn_core::tensor::Device; + +use crate::{ + Interrupter, LearnerSummaryConfig, OffPolicyConfig, RLCheckpointer, RLComponentsTypes, RLEvent, + RLEventProcessorType, RLResult, + metric::{processor::EventProcessorTraining, store::EventStoreClient}, +}; + +/// Struct to minimise parameters passed to [RLStrategy::train]. +pub struct RLComponents { + /// The total number of environment steps. + pub num_steps: usize, + /// The step number from which to continue the training. + pub checkpoint: Option, + /// A checkpointer used to load and save learning checkpoints. + pub checkpointer: Option>, + /// Enables gradients accumulation. + pub grad_accumulation: Option, + /// An [Interupter](Interrupter) that allows aborting the training/evaluation process early. + pub interrupter: Interrupter, + /// An [EventProcessor](crate::EventProcessorTraining) that processes events happening during training and evaluation. + pub event_processor: RLEventProcessorType, + /// A reference to an [EventStoreClient](EventStoreClient). + pub event_store: Arc, + /// Config for creating a summary of the learning + pub summary: Option, + /// Device used for running inference during environmment sampling or validation. + pub inference_device: Device, +} + +/// The strategy for reinforcement learning. +#[derive(Clone)] +pub enum RLStrategies { + /// Training on one device + OffPolicyStrategy(OffPolicyConfig), + /// Training using a custom learning strategy + Custom(CustomRLStrategy), +} + +/// A reference to an implementation of [RLStrategy]. +pub type CustomRLStrategy = Arc>; + +/// Provides the `fit` function for any learning strategy +pub trait RLStrategy { + /// Train the learner agent with this strategy. + fn train( + &self, + mut learner_agent: RLC::LearningAgent, + mut training_components: RLComponents, + env_init: RLC::EnvInit, + ) -> RLResult { + let starting_epoch = training_components.checkpoint.unwrap_or(0) + 1; + let summary_config = training_components.summary.clone(); + + // Event processor start training + training_components + .event_processor + .process_train(RLEvent::Start { + total_items: training_components.num_steps, + }); + + // Training loop + let (policy, mut event_processor) = self.train_loop( + training_components, + &mut learner_agent, + starting_epoch, + env_init, + ); + + let summary = summary_config.and_then(|summary| summary.init().ok()); + + // Signal training end. For the TUI renderer, this handles the exit & return to main screen. + // TODO: summary makes sense for RL? + event_processor.process_train(RLEvent::End(summary)); + + // let model = model.valid(); + let renderer = event_processor.renderer(); + + RLResult { policy, renderer } + } + + /// Training loop for this strategy + fn train_loop( + &self, + training_components: RLComponents, + learner_agent: &mut RLC::LearningAgent, + starting_epoch: usize, + env_init: RLC::EnvInit, + ) -> (RLC::Policy, RLEventProcessorType); +} diff --git a/crates/burn-train/src/learner/sequence.rs b/crates/burn-train/src/learner/sequence.rs new file mode 100644 index 0000000..cb15608 --- /dev/null +++ b/crates/burn-train/src/learner/sequence.rs @@ -0,0 +1,127 @@ +use crate::metric::{AccuracyInput, PerplexityInput, TopKAccuracyInput}; +use crate::metric::{Adaptor, CerInput, LossInput, WerInput, processor::ItemLazy}; +use burn_core::tensor::{Device, Int, Tensor, Transaction}; + +/// Sequence prediction output adapted for multiple metrics. +/// +/// Supported metrics: +/// - Accuracy +/// - TopKAccuracy +/// - Perplexity +/// - Loss +/// - CER +/// - WER +#[derive(new)] +pub struct SequenceOutput { + /// The loss. + pub loss: Tensor<1>, + + /// Raw logits. Shape: `[batch_size, seq_len, vocab_size]` + pub logits: Tensor<3>, + + /// Optional predicted token indices. Shape: `[batch_size, seq_length]`. + /// If not provided, predictions default to argmax of `logits` along the last dimension. + pub predictions: Option>, + + /// The target token indices. Shape: `[batch_size, seq_length]` + pub targets: Tensor<2, Int>, +} + +impl SequenceOutput { + fn predicted_tokens(&self) -> Tensor<2, Int> { + match &self.predictions { + Some(preds) => preds.clone(), + None => self.logits.clone().argmax(2).squeeze_dim::<2>(2), + } + } + + fn flat_logits(&self) -> Tensor<2> { + let [batch_size, seq_len, vocab_size] = self.logits.dims(); + self.logits + .clone() + .reshape([batch_size * seq_len, vocab_size]) + } + + fn flat_targets(&self) -> Tensor<1, Int> { + let [batch_size, seq_len] = self.targets.dims(); + self.targets.clone().reshape([batch_size * seq_len]) + } +} + +impl ItemLazy for SequenceOutput { + fn sync(self) -> Self { + let device: Device = Device::flex(); + + match self.predictions { + Some(preds) => { + let [logits, loss, targets, predictions] = Transaction::default() + .register(self.logits) + .register(self.loss) + .register(self.targets) + .register(preds) + .execute() + .try_into() + .expect("Correct amount of tensor data"); + + SequenceOutput { + logits: Tensor::from_data(logits, &device), + loss: Tensor::from_data(loss, &device), + targets: Tensor::from_data(targets, &device), + predictions: Some(Tensor::from_data(predictions, &device)), + } + } + None => { + let [logits, loss, targets] = Transaction::default() + .register(self.logits) + .register(self.loss) + .register(self.targets) + .execute() + .try_into() + .expect("Correct amount of tensor data"); + + SequenceOutput { + logits: Tensor::from_data(logits, &device), + loss: Tensor::from_data(loss, &device), + targets: Tensor::from_data(targets, &device), + predictions: None, + } + } + } + } +} + +impl Adaptor for SequenceOutput { + fn adapt(&self) -> LossInput { + LossInput::new(self.loss.clone()) + } +} + +impl Adaptor for SequenceOutput { + fn adapt(&self) -> CerInput { + CerInput::new(self.predicted_tokens(), self.targets.clone()) + } +} + +impl Adaptor for SequenceOutput { + fn adapt(&self) -> WerInput { + WerInput::new(self.predicted_tokens(), self.targets.clone()) + } +} + +impl Adaptor for SequenceOutput { + fn adapt(&self) -> AccuracyInput { + AccuracyInput::new(self.flat_logits(), self.flat_targets()) + } +} + +impl Adaptor for SequenceOutput { + fn adapt(&self) -> TopKAccuracyInput { + TopKAccuracyInput::new(self.flat_logits(), self.flat_targets()) + } +} + +impl Adaptor for SequenceOutput { + fn adapt(&self) -> PerplexityInput { + PerplexityInput::new(self.flat_logits(), self.flat_targets()) + } +} diff --git a/crates/burn-train/src/learner/sharder.rs b/crates/burn-train/src/learner/sharder.rs new file mode 100644 index 0000000..329e525 --- /dev/null +++ b/crates/burn-train/src/learner/sharder.rs @@ -0,0 +1,24 @@ +use burn_core::{ + Tensor, + module::{ModuleMapper, Param}, +}; + +use crate::{Learner, LearnerModel}; + +/// Describes how the module is distributed across multiple devices. +pub struct ModuleSharder; + +impl ModuleMapper for ModuleSharder { + fn map_float(&mut self, param: Param>) -> Param> { + let (id, tensor, mapper) = param.consume(); + let tensor = tensor.set_distributed(id); + Param::from_mapped_value(id, tensor, mapper) + } +} + +impl Learner { + /// Mark the model as sharded across multiple devices. + pub fn grad_sharded(&mut self) { + self.model = self.model.clone().map(&mut ModuleSharder); + } +} diff --git a/crates/burn-train/src/learner/summary.rs b/crates/burn-train/src/learner/summary.rs new file mode 100644 index 0000000..478fbb0 --- /dev/null +++ b/crates/burn-train/src/learner/summary.rs @@ -0,0 +1,475 @@ +use core::cmp::Ordering; +use std::{ + collections::{HashMap, hash_map::Entry}, + fmt::Display, + path::{Path, PathBuf}, +}; + +use crate::{ + logger::FileMetricLogger, + metric::store::{Aggregate, EventStore, LogEventStore, Split}, +}; + +/// Contains the metric value at a given time. +#[derive(Debug)] +pub struct MetricEntry { + /// The step at which the metric was recorded (i.e., epoch). + pub step: usize, + /// The metric value. + pub value: f64, +} + +/// Contains the summary of recorded values for a given metric. +#[derive(Debug)] +pub struct MetricSummary { + /// The metric name. + pub name: String, + /// The metric entries. + pub entries: Vec, +} + +impl MetricSummary { + fn collect( + event_store: &mut E, + metric: &str, + split: &Split, + num_epochs: usize, + ) -> Option { + let entries = (1..=num_epochs) + .filter_map(|epoch| { + event_store + .find_metric(metric, epoch, Aggregate::Mean, split) + .map(|value| MetricEntry { step: epoch, value }) + }) + .collect::>(); + + if entries.is_empty() { + None + } else { + Some(Self { + name: metric.to_string(), + entries, + }) + } + } +} + +/// Contains the summary of recorded metrics for the training and validation steps. +pub struct SummaryMetrics { + /// Training metrics summary. + pub train: Vec, + /// Validation metrics summary. + pub valid: Vec, + /// Test metrics summary per test split tag. + /// + /// Each key corresponds to a `Split::Test(Some(tag))`. + /// The empty string represents `Split::Test(None)`. + pub test: HashMap>, +} + +/// Detailed training summary. +pub struct LearnerSummary { + /// The number of epochs completed. + pub epochs: usize, + /// The summary of recorded metrics during training. + pub metrics: SummaryMetrics, + /// The model name (only recorded within the learner). + pub(crate) model: Option, +} + +impl LearnerSummary { + /// Creates a new learner summary for the specified metrics. + /// + /// # Arguments + /// + /// * `directory` - The directory containing the training artifacts (checkpoints and logs). + /// * `metrics` - The list of metrics to collect for the summary. + pub fn new>(directory: impl AsRef, metrics: &[S]) -> Result { + let directory = directory.as_ref(); + if !directory.exists() { + return Err(format!( + "Artifact directory does not exist at: {}", + directory.display() + )); + } + + let mut event_store = LogEventStore::default(); + let train_split = Split::Train; + let valid_split = Split::Valid; + + let logger = FileMetricLogger::new(directory); + let test_split_root = logger.split_dir(&Split::Test(None)); + if !logger.split_exists(&train_split) + && !logger.split_exists(&valid_split) + && test_split_root.is_none() + { + return Err(format!( + "No training, validation or test artifacts found at: {}", + directory.display() + )); + } + + // Number of recorded epochs + let epochs = logger.epochs(); + + event_store.register_logger(logger); + + let train_summary = metrics + .iter() + .filter_map(|metric| { + MetricSummary::collect(&mut event_store, metric.as_ref(), &train_split, epochs) + }) + .collect::>(); + + let valid_summary = metrics + .iter() + .filter_map(|metric| { + MetricSummary::collect(&mut event_store, metric.as_ref(), &valid_split, epochs) + }) + .collect::>(); + + let test_summary = match test_split_root { + Some(root) => collect_test_split_metrics(root, metrics, &mut event_store, epochs), + None => Default::default(), + }; + + Ok(Self { + epochs, + metrics: SummaryMetrics { + train: train_summary, + valid: valid_summary, + test: test_summary, + }, + model: None, + }) + } + + pub(crate) fn with_model(mut self, name: String) -> Self { + self.model = Some(name); + self + } + + /// Merges another summary into this one, combining all metric entries. + pub(crate) fn merge(mut self, other: LearnerSummary) -> Self { + fn merge_metrics( + base: Vec, + incoming: Vec, + ) -> Vec { + let mut map: HashMap = + base.into_iter().map(|m| (m.name.clone(), m)).collect(); + + for metric in incoming { + match map.entry(metric.name.clone()) { + Entry::Occupied(mut entry) => { + entry.get_mut().entries.extend(metric.entries); + } + Entry::Vacant(entry) => { + entry.insert(metric); + } + } + } + map.into_values().collect() + } + + self.metrics.train = merge_metrics(self.metrics.train, other.metrics.train); + self.metrics.valid = merge_metrics(self.metrics.valid, other.metrics.valid); + + for (tag, metrics) in other.metrics.test { + match self.metrics.test.entry(tag) { + Entry::Occupied(mut entry) => { + let current = std::mem::take(entry.get_mut()); + let merged = merge_metrics(current, metrics); + *entry.get_mut() = merged; + } + Entry::Vacant(entry) => { + entry.insert(metrics); + } + } + } + + if self.model != other.model { + self.model = None; + } + + self + } +} + +fn collect_test_split_metrics, S: AsRef>( + root: P, + metrics: &[S], + event_store: &mut LogEventStore, + epochs: usize, +) -> HashMap> { + // Collect immediate child directories + let dirs = match std::fs::read_dir(root) { + Ok(entries) => entries + .filter_map(|entry| { + let entry = entry.ok()?; + let file_type = entry.file_type().ok()?; + if file_type.is_dir() { + Some(entry.file_name().to_string_lossy().to_string()) + } else { + None + } + }) + .collect::>(), + Err(_) => Vec::new(), + }; + + let mut map = HashMap::new(); + + if dirs.is_empty() { + return map; + } + + // Detect if all directories are epoch directories + let all_epochs = dirs.iter().all(FileMetricLogger::is_epoch_dir); + + if all_epochs { + // Single untagged test split + let split = Split::Test(None); + + let summaries = metrics + .iter() + .filter_map(|metric| { + MetricSummary::collect(event_store, metric.as_ref(), &split, epochs) + }) + .collect::>(); + + // Untagged marked with empty string + map.insert("".to_string(), summaries); + } else { + // Tagged splits + for tag in dirs { + let split = Split::Test(Some(tag.clone().into())); + + let summaries = metrics + .iter() + .filter_map(|metric| { + MetricSummary::collect(event_store, metric.as_ref(), &split, epochs) + }) + .collect::>(); + + map.insert(tag, summaries); + } + } + + map +} + +impl Display for LearnerSummary { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // Compute the max length for each column + let mut max_split_len = 5; // "Train" + let mut max_metric_len = "Metric".len(); + for metric in self.metrics.train.iter() { + max_metric_len = max_metric_len.max(metric.name.len()); + } + for metric in self.metrics.valid.iter() { + max_metric_len = max_metric_len.max(metric.name.len()); + } + for (tag, metrics) in self.metrics.test.iter() { + let split_name = if tag.is_empty() { + "Test".to_string() + } else { + format!("Test ({tag})") + }; + + max_split_len = max_split_len.max(split_name.len()); + + for metric in metrics { + max_metric_len = max_metric_len.max(metric.name.len()); + } + } + + // Summary header + writeln!( + f, + "{:=>width_symbol$} Learner Summary {:=>width_symbol$}", + "", + "", + width_symbol = 24, + )?; + + if let Some(model) = &self.model { + writeln!(f, "Model:\n{model}")?; + } + writeln!(f, "Total Epochs: {epochs}\n\n", epochs = self.epochs)?; + + // Metrics table header + writeln!( + f, + "| {:width_split$}--|{:->width_metric$}--|----------|----------|----------|----------|", + "Split", + "Metric", + "", + "", + width_split = max_split_len, + width_metric = max_metric_len, + )?; + + // Table entries + fn cmp_f64(a: &f64, b: &f64) -> Ordering { + match (a.is_nan(), b.is_nan()) { + (true, true) => Ordering::Equal, + (true, false) => Ordering::Greater, + (false, true) => Ordering::Less, + _ => a.partial_cmp(b).unwrap(), + } + } + + fn fmt_val(val: f64) -> String { + if val < 1e-2 { + // Use scientific notation for small values which would otherwise be truncated + format!("{val:<9.3e}") + } else { + format!("{val:<9.3}") + } + } + + let mut write_metrics_summary = + |metrics: &[MetricSummary], split: String| -> std::fmt::Result { + for metric in metrics.iter() { + if metric.entries.is_empty() { + continue; // skip metrics with no recorded values + } + + // Compute the min & max for each metric + let metric_min = metric + .entries + .iter() + .min_by(|a, b| cmp_f64(&a.value, &b.value)) + .unwrap(); + let metric_max = metric + .entries + .iter() + .max_by(|a, b| cmp_f64(&a.value, &b.value)) + .unwrap(); + + writeln!( + f, + "| {:, +} + +impl LearnerSummaryConfig { + /// Create the learning summary. + pub fn init(&self) -> Result { + LearnerSummary::new(&self.directory, &self.metrics[..]) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + #[should_panic = "Summary artifacts should exist"] + fn test_artifact_dir_should_exist() { + let dir = "/tmp/learner-summary-not-found"; + let _summary = LearnerSummary::new(dir, &["Loss"]).expect("Summary artifacts should exist"); + } + + #[test] + #[should_panic = "Summary artifacts should exist"] + fn test_train_valid_artifacts_should_exist() { + let dir = "/tmp/test-learner-summary-empty"; + std::fs::create_dir_all(dir).ok(); + let _summary = LearnerSummary::new(dir, &["Loss"]).expect("Summary artifacts should exist"); + } + + #[test] + fn test_summary_should_be_empty() { + let dir = Path::new("/tmp/test-learner-summary-empty-metrics"); + std::fs::create_dir_all(dir).unwrap(); + std::fs::create_dir_all(dir.join("train/epoch-1")).unwrap(); + std::fs::create_dir_all(dir.join("valid/epoch-1")).unwrap(); + let summary = LearnerSummary::new(dir.to_str().unwrap(), &["Loss"]) + .expect("Summary artifacts should exist"); + + assert_eq!(summary.epochs, 1); + + assert_eq!(summary.metrics.train.len(), 0); + assert_eq!(summary.metrics.valid.len(), 0); + + std::fs::remove_dir_all(dir).unwrap(); + } + + #[test] + fn test_summary_should_be_collected() { + let dir = Path::new("/tmp/test-learner-summary"); + let train_dir = dir.join("train/epoch-1"); + let valid_dir = dir.join("valid/epoch-1"); + std::fs::create_dir_all(dir).unwrap(); + std::fs::create_dir_all(&train_dir).unwrap(); + std::fs::create_dir_all(&valid_dir).unwrap(); + + std::fs::write(train_dir.join("Loss.log"), "1.0\n2.0").expect("Unable to write file"); + std::fs::write(valid_dir.join("Loss.log"), "1.0").expect("Unable to write file"); + + let summary = LearnerSummary::new(dir.to_str().unwrap(), &["Loss"]) + .expect("Summary artifacts should exist"); + + assert_eq!(summary.epochs, 1); + + // Only Loss metric + assert_eq!(summary.metrics.train.len(), 1); + assert_eq!(summary.metrics.valid.len(), 1); + + // Aggregated train metric entries for 1 epoch + let train_metric = &summary.metrics.train[0]; + assert_eq!(train_metric.name, "Loss"); + assert_eq!(train_metric.entries.len(), 1); + let entry = &train_metric.entries[0]; + assert_eq!(entry.step, 1); // epoch = 1 + assert_eq!(entry.value, 1.5); // (1 + 2) / 2 + + // Aggregated valid metric entries for 1 epoch + let valid_metric = &summary.metrics.valid[0]; + assert_eq!(valid_metric.name, "Loss"); + assert_eq!(valid_metric.entries.len(), 1); + let entry = &valid_metric.entries[0]; + assert_eq!(entry.step, 1); // epoch = 1 + assert_eq!(entry.value, 1.0); + + std::fs::remove_dir_all(dir).unwrap(); + } +} diff --git a/crates/burn-train/src/learner/supervised/mod.rs b/crates/burn-train/src/learner/supervised/mod.rs new file mode 100644 index 0000000..56b2f0f --- /dev/null +++ b/crates/burn-train/src/learner/supervised/mod.rs @@ -0,0 +1,7 @@ +mod paradigm; +mod step; +mod strategies; + +pub use paradigm::*; +pub use step::*; +pub use strategies::*; diff --git a/crates/burn-train/src/learner/supervised/paradigm.rs b/crates/burn-train/src/learner/supervised/paradigm.rs new file mode 100644 index 0000000..2ad5124 --- /dev/null +++ b/crates/burn-train/src/learner/supervised/paradigm.rs @@ -0,0 +1,560 @@ +use crate::checkpoint::{ + AsyncCheckpointer, Checkpointer, CheckpointingStrategy, ComposedCheckpointingStrategy, + FileCheckpointer, KeepLastNCheckpoints, MetricCheckpointingStrategy, +}; +use crate::learner::EarlyStoppingStrategy; +use crate::learner::base::Interrupter; +use crate::logger::{FileMetricLogger, MetricLogger, TrainingProgressLogger}; +use crate::metric::processor::{ + AsyncProcessorTraining, FullEventProcessorTraining, MetricsTraining, +}; +use crate::metric::store::{Aggregate, Direction, EventStoreClient, LogEventStore, Split}; +use crate::metric::{Adaptor, LossMetric, Metric, Numeric}; +use crate::multi::MultiDeviceLearningStrategy; +use crate::renderer::{MetricsRenderer, default_renderer}; +use crate::single::SingleDeviceTrainingStrategy; +use crate::{ + ApplicationLoggerInstaller, EarlyStoppingStrategyRef, ExecutionStrategy, + FileApplicationLoggerInstaller, InferenceModelInput, InferenceModelOutput, InferenceStep, + LearnerEvent, LearnerModel, LearnerSummaryConfig, LearningCheckpointer, LearningResult, + TrainStep, TrainingComponents, TrainingModelInput, TrainingModelOutput, TrainingStrategy, +}; +use crate::{Learner, SupervisedLearningStrategy}; +use burn_core::data::dataloader::DataLoader; +use burn_core::store::ModuleRecord; +use burn_core::tensor::Device; +use burn_optim::OptimizerRecord; +use burn_optim::lr_scheduler::LrSchedulerRecord; +use std::collections::BTreeSet; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +/// A reference to the training split [DataLoader](DataLoader). +pub type TrainLoader = Arc>>; +/// A reference to the validation split [DataLoader](DataLoader). +pub type ValidLoader = Arc>>; +/// The event processor type for supervised learning. +pub type SupervisedTrainingEventProcessor = AsyncProcessorTraining< + LearnerEvent>, + LearnerEvent>, +>; + +/// Structure to configure and launch supervised learning trainings. +pub struct SupervisedTraining { + // Not that complex. Extracting into another type would only make it more confusing. + #[allow(clippy::type_complexity)] + checkpointers: Option<( + AsyncCheckpointer, + AsyncCheckpointer, + AsyncCheckpointer, + )>, + num_epochs: usize, + checkpoint: Option, + directory: PathBuf, + grad_accumulation: Option, + grad_checkpointing: bool, + renderer: Option>, + metrics: MetricsTraining, InferenceModelOutput>, + event_store: LogEventStore, + interrupter: Interrupter, + tracing_logger: Option>, + checkpointer_strategy: Box, + early_stopping: Option, + training_strategy: Option>, + dataloader_train: TrainLoader, + dataloader_valid: ValidLoader, + // Use BTreeSet instead of HashSet for consistent (alphabetical) iteration order + summary_metrics: BTreeSet, + summary: bool, + progress_logger: Option>, +} + +impl SupervisedTraining { + /// Creates a new runner for a supervised training. + /// + /// # Arguments + /// + /// * `directory` - The directory to save the checkpoints. + /// * `dataloader_train` - The dataloader for the training split. + /// * `dataloader_valid` - The dataloader for the validation split. + pub fn new( + directory: impl AsRef, + dataloader_train: Arc::Input>>, + dataloader_valid: Arc::Input>>, + ) -> Self { + let directory = directory.as_ref().to_path_buf(); + let experiment_log_file = directory.join("experiment.log"); + Self { + num_epochs: 1, + checkpoint: None, + checkpointers: None, + directory, + grad_accumulation: None, + grad_checkpointing: false, + metrics: MetricsTraining::default(), + event_store: LogEventStore::default(), + renderer: None, + interrupter: Interrupter::new(), + tracing_logger: Some(Box::new(FileApplicationLoggerInstaller::new( + experiment_log_file, + ))), + checkpointer_strategy: Box::new( + ComposedCheckpointingStrategy::builder() + .add(KeepLastNCheckpoints::new(2)) + .add(MetricCheckpointingStrategy::new( + &LossMetric::new(), // default to valid loss + Aggregate::Mean, + Direction::Lowest, + Split::Valid, + )) + .build(), + ), + early_stopping: None, + training_strategy: None, + summary_metrics: BTreeSet::new(), + summary: false, + dataloader_train, + dataloader_valid, + progress_logger: None, + } + } +} + +impl SupervisedTraining { + /// Replace the default training strategy (SingleDeviceTrainingStrategy) with the provided one. + /// + /// # Arguments + /// + /// * `training_strategy` - The training strategy. + pub fn with_training_strategy(mut self, training_strategy: TrainingStrategy) -> Self { + self.training_strategy = Some(training_strategy); + self + } + + /// Replace the default metric loggers with the provided ones. + /// + /// # Arguments + /// + /// * `logger` - The training logger. + pub fn with_metric_logger(mut self, logger: ML) -> Self + where + ML: MetricLogger + 'static, + { + self.event_store.register_logger(logger); + self + } + + /// Register a progress logger to track and store training progress. + /// + /// # Example + /// + /// ```ignore + /// // `MyTrainingProgressLogger` is a user-defined type that implements + /// // `burn_train::logger::TrainingProgressLogger`. + /// let learner = SupervisedTraining::new(...) + /// .with_progress_logger(MyTrainingProgressLogger); + /// ``` + pub fn with_progress_logger(mut self, logger: PL) -> Self + where + PL: TrainingProgressLogger + 'static, + { + self.progress_logger = Some(Box::new(logger)); + self + } + + /// Update the checkpointing_strategy. + pub fn with_checkpointing_strategy( + mut self, + strategy: CS, + ) -> Self { + self.checkpointer_strategy = Box::new(strategy); + self + } + + /// Replace the default CLI renderer with a custom one. + /// + /// # Arguments + /// + /// * `renderer` - The custom renderer. + pub fn renderer(mut self, renderer: MR) -> Self + where + MR: MetricsRenderer + 'static, + { + self.renderer = Some(Box::new(renderer)); + self + } + + /// Register all metrics as numeric for the training and validation set. + pub fn metrics>(self, metrics: Me) -> Self { + metrics.register(self) + } + + /// Register all metrics as text for the training and validation set. + pub fn metrics_text>(self, metrics: Me) -> Self { + metrics.register(self) + } + + /// Register a training metric. + pub fn metric_train(mut self, metric: Me) -> Self + where + TrainingModelOutput: Adaptor, + { + self.metrics.register_train_metric(metric); + self + } + + /// Register a validation metric. + pub fn metric_valid(mut self, metric: Me) -> Self + where + InferenceModelOutput: Adaptor, + { + self.metrics.register_valid_metric(metric); + self + } + + /// Enable gradients accumulation. + /// + /// # Notes + /// + /// When you enable gradients accumulation, the gradients object used by the optimizer will be + /// the sum of all gradients generated by each backward pass. It might be a good idea to + /// reduce the learning to compensate. + /// + /// The effect is similar to increasing the `batch size` and the `learning rate` by the `accumulation` + /// amount. + pub fn grads_accumulation(mut self, accumulation: usize) -> Self { + self.grad_accumulation = Some(accumulation); + self + } + + /// Enables autodiff checkpointing. + /// + /// # Notes + /// Gradient checkpointing recomputes activations during backpropagation for operations + /// marked as memory-bound, while compute-bound operations still cache their + /// output. This reduces peak memory usage at the cost of additional computation + /// for memory-bound ops. + pub fn gradient_checkpointing(mut self) -> Self { + self.grad_checkpointing = true; + self + } + + /// Register a [numeric](crate::metric::Numeric) training [metric](Metric). + pub fn metric_train_numeric(mut self, metric: Me) -> Self + where + Me: Metric + Numeric + 'static, + TrainingModelOutput: Adaptor, + { + self.summary_metrics.insert(metric.name().to_string()); + self.metrics.register_train_metric_numeric(metric); + self + } + + /// Register a [numeric](crate::metric::Numeric) validation [metric](Metric). + pub fn metric_valid_numeric(mut self, metric: Me) -> Self + where + InferenceModelOutput: Adaptor, + { + self.summary_metrics.insert(metric.name().to_string()); + self.metrics.register_valid_metric_numeric(metric); + self + } + + /// The number of epochs the training should last. + pub fn num_epochs(mut self, num_epochs: usize) -> Self { + self.num_epochs = num_epochs; + self + } + + /// The epoch from which the training must resume. + pub fn checkpoint(mut self, checkpoint: usize) -> Self { + self.checkpoint = Some(checkpoint); + self + } + + /// Provides a handle that can be used to interrupt training. + pub fn interrupter(&self) -> Interrupter { + self.interrupter.clone() + } + + /// Override the handle for stopping training with an externally provided handle + pub fn with_interrupter(mut self, interrupter: Interrupter) -> Self { + self.interrupter = interrupter; + self + } + + /// Register an [early stopping strategy](EarlyStoppingStrategy) to stop the training when the + /// conditions are meet. + pub fn early_stopping(mut self, strategy: Strategy) -> Self + where + Strategy: EarlyStoppingStrategy + Clone + Send + Sync + 'static, + { + self.early_stopping = Some(Box::new(strategy)); + self + } + + /// By default, Rust logs are captured and written into + /// `experiment.log`. If disabled, standard Rust log handling + /// will apply. + pub fn with_application_logger( + mut self, + logger: Option>, + ) -> Self { + self.tracing_logger = logger; + self + } + + /// Register a checkpointer that will save the [optimizer](burn_optim::ModuleOptimizer), the + /// [model](LearnerModel) and the [learning rate scheduler](burn_optim::lr_scheduler::module_lr_scheduler::ModuleLrScheduler) to separate burnpack files. + pub fn with_default_checkpointers(mut self) -> Self { + let checkpoint_dir = self.directory.join("checkpoint"); + let checkpointer_model = FileCheckpointer::new(&checkpoint_dir, "model"); + let checkpointer_optimizer = FileCheckpointer::new(&checkpoint_dir, "optim"); + let checkpointer_scheduler = FileCheckpointer::new(&checkpoint_dir, "scheduler"); + + self.checkpointers = Some(( + AsyncCheckpointer::new(checkpointer_model), + AsyncCheckpointer::new(checkpointer_optimizer), + AsyncCheckpointer::new(checkpointer_scheduler), + )); + + self + } + + /// Register your own checkpointers that will save the [optimizer](burn_optim::ModuleOptimizer), the + /// [model](LearnerModel) and the [learning rate scheduler](burn_optim::lr_scheduler::module_lr_scheduler::ModuleLrScheduler) to separate burnpack files. + pub fn with_custom_checkpointers( + mut self, + module_checkpointer: CM, + optimizer_checkpointer: CO, + lr_checkpointer: CL, + ) -> Self + where + CM: Checkpointer + 'static, + CO: Checkpointer + 'static, + CL: Checkpointer + 'static, + { + self.checkpointers = Some(( + AsyncCheckpointer::new(module_checkpointer), + AsyncCheckpointer::new(optimizer_checkpointer), + AsyncCheckpointer::new(lr_checkpointer), + )); + + self + } + + /// Enable the training summary report. + /// + /// The summary will be displayed after `.fit()`, when the renderer is dropped. + pub fn summary(mut self) -> Self { + self.summary = true; + self + } +} + +impl SupervisedTraining { + /// Launch this training with the given [Learner](Learner). + pub fn launch(mut self, learner: Learner) -> LearningResult { + if self.tracing_logger.is_some() + && let Err(e) = self.tracing_logger.as_ref().unwrap().install() + { + log::warn!("Failed to install the experiment logger: {e}"); + } + let renderer = self + .renderer + .unwrap_or_else(|| default_renderer(self.interrupter.clone(), self.checkpoint)); + + if !self.event_store.has_loggers() { + self.event_store + .register_logger(FileMetricLogger::new(self.directory.clone())); + } + + let event_store = Arc::new(EventStoreClient::new(self.event_store)); + let full_processor = + FullEventProcessorTraining::new(self.metrics, renderer, event_store.clone()); + let full_processor = match self.progress_logger { + Some(logger) => full_processor.with_progress_logger(logger), + None => full_processor, + }; + let event_processor = AsyncProcessorTraining::new(full_processor); + + let checkpointer = self.checkpointers.map(|(model, optim, scheduler)| { + LearningCheckpointer::new( + model.with_interrupter(self.interrupter.clone()), + optim.with_interrupter(self.interrupter.clone()), + scheduler.with_interrupter(self.interrupter.clone()), + self.checkpointer_strategy, + ) + }); + + let summary = if self.summary { + Some(LearnerSummaryConfig { + directory: self.directory, + metrics: self.summary_metrics.into_iter().collect::>(), + }) + } else { + None + }; + + let components = TrainingComponents { + checkpoint: self.checkpoint, + checkpointer, + interrupter: self.interrupter, + early_stopping: self.early_stopping, + event_processor, + event_store, + num_epochs: self.num_epochs, + grad_accumulation: self.grad_accumulation, + summary, + }; + + // Default to single device based on model + let training_strategy = self.training_strategy.unwrap_or(TrainingStrategy::Default( + ExecutionStrategy::SingleDevice(autodiff_device( + learner.model.devices()[0].clone(), + self.grad_checkpointing, + )), + )); + + let mut learner = learner; + if let Some(checkpoint) = components.checkpoint + && let Some(checkpointer) = &components.checkpointer + { + learner = checkpointer.load_checkpoint(learner, checkpoint); + } + + match training_strategy { + TrainingStrategy::Custom(learning_paradigm) => learning_paradigm.train( + learner, + self.dataloader_train, + self.dataloader_valid, + components, + ), + TrainingStrategy::Default(strategy) => match strategy { + ExecutionStrategy::SingleDevice(device) => { + let single_device = SingleDeviceTrainingStrategy::new(autodiff_device( + device, + self.grad_checkpointing, + )); + single_device.train( + learner, + self.dataloader_train, + self.dataloader_valid, + components, + ) + } + ExecutionStrategy::MultiDevice(devices, multi_device_optim) => { + let strategy: Box> = match devices.len() == 1 + { + true => Box::new(SingleDeviceTrainingStrategy::new(autodiff_device( + devices[0].clone(), + self.grad_checkpointing, + ))), + false => Box::new(MultiDeviceLearningStrategy::new( + devices + .into_iter() + .map(|d| autodiff_device(d, self.grad_checkpointing)) + .collect(), + multi_device_optim, + )), + }; + strategy.train( + learner, + self.dataloader_train, + self.dataloader_valid, + components, + ) + } + ExecutionStrategy::DistributedDataParallel { devices, context } => { + use crate::ddp::DdpTrainingStrategy; + + let ddp = DdpTrainingStrategy::new( + devices + .into_iter() + .map(|d| autodiff_device(d, self.grad_checkpointing)) + .collect(), + context, + ); + ddp.train( + learner, + self.dataloader_train, + self.dataloader_valid, + components, + ) + } + }, + } + } +} + +// Validate the autodiff device property and enable grad checkpointing. +fn autodiff_device(mut device: Device, grad_checkpointing: bool) -> Device { + if !device.is_autodiff() { + device = device.autodiff(); + } + + if grad_checkpointing { + device = device.gradient_checkpointing(); + } + + device +} + +/// Trait to fake variadic generics. +pub trait MetricRegistration: Sized { + /// Register the metrics. + fn register(self, builder: SupervisedTraining) -> SupervisedTraining; +} + +/// Trait to fake variadic generics. +pub trait TextMetricRegistration: Sized { + /// Register the metrics. + fn register(self, builder: SupervisedTraining) -> SupervisedTraining; +} + +macro_rules! gen_tuple { + ($($M:ident),*) => { + impl<$($M,)* M: LearnerModel> TextMetricRegistration for ($($M,)*) + where + $(TrainingModelOutput: Adaptor<$M::Input>,)* + $(InferenceModelOutput: Adaptor<$M::Input>,)* + $($M: Metric + 'static,)* + { + #[allow(non_snake_case)] + fn register( + self, + builder: SupervisedTraining, + ) -> SupervisedTraining { + let ($($M,)*) = self; + $(let builder = builder.metric_train($M.clone());)* + $(let builder = builder.metric_valid($M);)* + builder + } + } + + impl<$($M,)* M: LearnerModel> MetricRegistration for ($($M,)*) + where + $(TrainingModelOutput: Adaptor<$M::Input>,)* + $(InferenceModelOutput: Adaptor<$M::Input>,)* + $($M: Metric + Numeric + 'static,)* + { + #[allow(non_snake_case)] + fn register( + self, + builder: SupervisedTraining, + ) -> SupervisedTraining { + let ($($M,)*) = self; + $(let builder = builder.metric_train_numeric($M.clone());)* + $(let builder = builder.metric_valid_numeric($M);)* + builder + } + } + }; +} + +gen_tuple!(M1); +gen_tuple!(M1, M2); +gen_tuple!(M1, M2, M3); +gen_tuple!(M1, M2, M3, M4); +gen_tuple!(M1, M2, M3, M4, M5); +gen_tuple!(M1, M2, M3, M4, M5, M6); diff --git a/crates/burn-train/src/learner/supervised/step/mod.rs b/crates/burn-train/src/learner/supervised/step/mod.rs new file mode 100644 index 0000000..066b064 --- /dev/null +++ b/crates/burn-train/src/learner/supervised/step/mod.rs @@ -0,0 +1,2 @@ +/// The trainer module. +pub mod train; diff --git a/crates/burn-train/src/learner/supervised/step/train.rs b/crates/burn-train/src/learner/supervised/step/train.rs new file mode 100644 index 0000000..72e5e2e --- /dev/null +++ b/crates/burn-train/src/learner/supervised/step/train.rs @@ -0,0 +1,150 @@ +use crate::LearnerModel; +use crate::{TrainOutput, TrainStep, TrainingModelInput, TrainingModelOutput}; +use burn_core::data::dataloader::DataLoaderIterator; +use burn_core::data::dataloader::Progress; +use burn_core::tensor::Device; +use std::sync::mpsc::{Receiver, Sender}; +use std::thread::spawn; + +/// Multi devices train step. +pub struct MultiDevicesTrainStep { + workers: Vec>, + receiver: Receiver>>, +} + +struct Message { + item: TI, + model: M, +} + +struct Worker { + // Not that complex. Extracting into another type would only make it more confusing. + // #[allow(clippy::type_complexity)] + sender_input: Sender>>, + device: Device, + device_id: usize, +} + +impl Worker { + fn register(&self, item: TrainingModelInput, model: &M) { + let message = Message { + item, + model: model.clone(), + }; + self.sender_input.send(message).unwrap(); + } + + // Not that complex. Extracting into another type would only make it more confusing. + // #[allow(clippy::type_complexity)] + fn start( + &self, + sender_output: Sender>>, + receiver_input: Receiver>>, + ) { + let device = self.device.clone(); + let device_id = self.device_id; + + spawn(move || { + loop { + match receiver_input.recv() { + Ok(item) => { + let model = item.model.fork(&device); + let output = TrainStep::step(&model, item.item); + let item = MultiTrainOutput { output, device_id }; + + sender_output.send(item).unwrap(); + } + Err(_err) => { + log::info!("Closing thread on device {device:?}"); + break; + } + } + } + }); + } +} + +/// Multiple output items. +pub struct MultiTrainOutput { + /// The training output. + pub output: TrainOutput, + /// The worker/device on which the computing happened. + pub(crate) device_id: usize, +} + +impl MultiDevicesTrainStep { + /// Create a new multi devices train step. + /// + /// # Arguments + /// + /// * `devices` - Devices. + /// + /// # Returns + /// + /// MultiDevicesTrainStep instance. + pub fn new(devices: &[Device]) -> Self { + let (sender_output, receiver_output) = std::sync::mpsc::channel(); + let workers = devices + .iter() + .enumerate() + .map(|(device_id, device)| { + let (sender_input, receiver_input) = std::sync::mpsc::channel(); + let worker = Worker { + sender_input, + device: device.clone(), + device_id, + }; + + worker.start(sender_output.clone(), receiver_input); + worker + }) + .collect(); + + Self { + workers, + receiver: receiver_output, + } + } + + /// Collect outputs from workers for one step. + /// + /// # Arguments + /// + /// * `model` - Model. + /// * `dataloaders` - The data loader for each worker. + /// + /// # Returns + /// + /// Outputs. + pub fn step<'a>( + &self, + dataloaders: &mut [Box> + 'a>], + model: &M, + ) -> (Vec>>, Progress) { + let mut num_send = 0; + + let mut items_total = 0; + let mut items_processed = 0; + let unit: Option = Some("items".to_string()); + + for (i, worker) in self.workers.iter().enumerate() { + let dataloader = &mut dataloaders[i]; + if let Some(item) = dataloader.next() { + worker.register(item, model); + num_send += 1; + let progress = dataloader.progress(); + items_total += progress.items_total; + items_processed += progress.items_processed; + } + } + + let mut outputs = Vec::with_capacity(num_send); + + for _ in 0..num_send { + let output = self.receiver.recv().unwrap(); + outputs.push(output); + } + + (outputs, Progress::new(items_processed, items_total, unit)) + } +} diff --git a/crates/burn-train/src/learner/supervised/strategies/base.rs b/crates/burn-train/src/learner/supervised/strategies/base.rs new file mode 100644 index 0000000..7e09987 --- /dev/null +++ b/crates/burn-train/src/learner/supervised/strategies/base.rs @@ -0,0 +1,172 @@ +use crate::{ + EarlyStoppingStrategyRef, Interrupter, Learner, LearnerModel, LearnerSummaryConfig, + LearningCheckpointer, LearningResult, SupervisedTrainingEventProcessor, TrainLoader, + ValidLoader, + metric::{ + processor::{EventProcessorTraining, LearnerEvent}, + store::EventStoreClient, + }, +}; +use burn_core::prelude::Device; +use burn_core::tensor::distributed::{DistributedConfig, DistributedContext}; +use std::sync::Arc; + +/// A reference to an implementation of SupervisedLearningStrategy. +pub type CustomLearningStrategy = Arc>; + +#[derive(Clone, Copy, Debug)] +/// Determine how the optimization is performed when training with multiple devices. +pub enum MultiDeviceOptim { + /// The optimization is done on an elected device. + OptimMainDevice, + /// The optimization is sharded across all devices. + OptimSharded, +} + +/// Describes where training runs. +pub enum ExecutionStrategy { + /// Training on one device + SingleDevice(Device), + /// Performs data-parallel distributed training where the optimization is + /// done on an elected master device. + MultiDevice(Vec, MultiDeviceOptim), + /// Training with input distributed across devices, each device has its own copy of the model. + /// Collective ops are used to sync the gradients after each pass. + DistributedDataParallel { + /// Devices on this node for the DDP + devices: Vec, + /// The distributed runtime. + context: DistributedContext, + }, +} + +impl ExecutionStrategy { + /// Returns the primary device responsible for coordination. + pub fn main_device(&self) -> &Device { + match self { + ExecutionStrategy::SingleDevice(device) => device, + ExecutionStrategy::MultiDevice(devices, _optim) => &devices[0], + ExecutionStrategy::DistributedDataParallel { + devices, + context: _, + } => &devices[0], + } + } + + /// Creates a strategy for a single device. + pub fn single(device: Device) -> Self { + Self::SingleDevice(device) + } + + /// Creates a multi-device strategy. + pub fn multi(devices: Vec, optim: MultiDeviceOptim) -> Self { + Self::MultiDevice(devices, optim) + } +} + +impl ExecutionStrategy { + /// Creates a distributed data parallel (DDP) strategy. + pub fn ddp(devices: Vec, config: DistributedConfig) -> Self { + let context = DistributedContext::init(devices.clone(), config); + Self::DistributedDataParallel { devices, context } + } +} + +/// How should the learner run the learning for the model +pub enum TrainingStrategy { + /// Default training loop with specified device strategy. + Default(ExecutionStrategy), + /// Training using a custom learning strategy + Custom(CustomLearningStrategy), +} + +impl From for TrainingStrategy { + fn from(value: ExecutionStrategy) -> Self { + Self::Default(value) + } +} + +impl Default for TrainingStrategy { + fn default() -> Self { + Self::Default(ExecutionStrategy::SingleDevice(Default::default())) + } +} + +/// Struct to minimise parameters passed to [SupervisedLearningStrategy::train]. +/// These components are used during training. +pub struct TrainingComponents { + /// The total number of epochs + pub num_epochs: usize, + /// The epoch number from which to continue the training. + pub checkpoint: Option, + /// A checkpointer used to load and save learner checkpoints. + pub checkpointer: Option>, + /// Enables gradients accumulation. + pub grad_accumulation: Option, + /// An [Interupter](Interrupter) that allows aborting the training/evaluation process early. + pub interrupter: Interrupter, + /// Cloneable reference to an early stopping strategy. + pub early_stopping: Option, + /// An [EventProcessor](crate::EventProcessorTraining) that processes events happening during training and validation. + pub event_processor: SupervisedTrainingEventProcessor, + /// A reference to an [EventStoreClient](EventStoreClient). + pub event_store: Arc, + /// Config for creating a summary of the learning + pub summary: Option, +} + +/// Provides the `fit` function for any learning strategy +pub trait SupervisedLearningStrategy { + /// Train the learner's model with this strategy. + fn train( + &self, + learner: Learner, + dataloader_train: TrainLoader, + dataloader_valid: ValidLoader, + mut training_components: TrainingComponents, + ) -> LearningResult { + let starting_epoch = training_components.checkpoint.unwrap_or(0) + 1; + let summary_config = training_components.summary.clone(); + + // Event processor start training + training_components + .event_processor + .process_train(LearnerEvent::Start { + total_epochs: training_components.num_epochs, + starting_epoch, + }); + // Training loop + let (model, mut event_processor) = self.fit( + training_components, + learner, + dataloader_train, + dataloader_valid, + starting_epoch, + ); + + let summary = summary_config.and_then(|summary| { + summary + .init() + .map(|summary| summary.with_model(model.to_string())) + .ok() + }); + + // Signal training end. For the TUI renderer, this handles the exit & return to main screen. + event_processor.process_train(LearnerEvent::End(summary)); + + let model = model.valid(); + let renderer = event_processor.renderer(); + + LearningResult:: { model, renderer } + } + + /// Training loop for this strategy + fn fit( + &self, + training_components: TrainingComponents, + learner: Learner, + dataloader_train: TrainLoader, + dataloader_valid: ValidLoader, + starting_epoch: usize, + ) -> (M, SupervisedTrainingEventProcessor); +} diff --git a/crates/burn-train/src/learner/supervised/strategies/ddp/README.md b/crates/burn-train/src/learner/supervised/strategies/ddp/README.md new file mode 100644 index 0000000..b38c9c7 --- /dev/null +++ b/crates/burn-train/src/learner/supervised/strategies/ddp/README.md @@ -0,0 +1,17 @@ +## DDP +Distributed Data Parallel + +The DDP is a learning strategy that trains a replica of the model on each device. + +The DDP launches threads for each local device. Each thread on each node will run the model. +After the forward and backward passes, the gradients are synced between all peers on all nodes +with an `all-reduce` operation. + +While the DDP launches threads for each local device, it is the user's responsibility to launch the +DDP on each node, and assure the collective configuration matches. + +## Main device vs secondary devices + +The main device is responsible for validation, as well as event processing, which is used in the UI. + +The first device is chosen as the main device. diff --git a/crates/burn-train/src/learner/supervised/strategies/ddp/epoch.rs b/crates/burn-train/src/learner/supervised/strategies/ddp/epoch.rs new file mode 100644 index 0000000..f2120a9 --- /dev/null +++ b/crates/burn-train/src/learner/supervised/strategies/ddp/epoch.rs @@ -0,0 +1,138 @@ +use burn_core::data::dataloader::Progress; +use burn_optim::GradientsAccumulator; +use std::sync::{Arc, Mutex}; + +use crate::SupervisedTrainingEventProcessor; +use crate::learner::base::Interrupter; +use crate::metric::processor::{EventProcessorTraining, LearnerEvent, TrainingItem}; +use crate::{InferenceStep, Learner, LearnerModel, TrainLoader, ValidLoader}; + +/// A validation epoch. +#[derive(new)] +pub struct DdpValidEpoch { + dataloader: ValidLoader, +} + +/// A training epoch. +#[derive(new)] +pub struct DdpTrainEpoch { + dataloader: TrainLoader, + grad_accumulation: Option, +} + +impl DdpValidEpoch { + /// Runs the validation epoch. + /// + /// # Arguments + /// + /// * `model` - The model to validate. + /// * `processor` - The event processor to use. + pub fn run( + &self, + model: &M, + global_progress: &Progress, + processor: &mut SupervisedTrainingEventProcessor, + interrupter: &Interrupter, + ) { + let epoch = global_progress.items_processed; + log::info!("Executing validation step for epoch {}", epoch); + let model = model.valid(); + + let mut iterator = self.dataloader.iter(); + let mut iteration = 0; + + while let Some(item) = iterator.next() { + let progress = iterator.progress(); + iteration += 1; + + let item = InferenceStep::step(&model, item); + let item = TrainingItem::new(item, progress, Some(iteration), None); + + processor.process_valid(LearnerEvent::ProcessedItem(item)); + + if interrupter.should_stop() { + log::info!("Training interrupted."); + break; + } + } + } +} + +impl DdpTrainEpoch { + /// Runs the training epoch. + /// + /// # Arguments + /// + /// * `model` - The model to train. + /// * `optim` - The optimizer to use. + /// * `scheduler` - The learning rate scheduler to use. + /// * `processor` - The event processor to use. + /// + /// # Returns + /// + /// The trained model and the optimizer. + pub fn run( + &self, + learner: &mut Learner, + global_progress: &Progress, + processor: Arc>>, + interrupter: &Interrupter, + peer_count: usize, + ) { + let epoch = global_progress.items_processed; + log::info!("Executing training step for epoch {}", epoch,); + + let mut iterator = self.dataloader.iter(); + let mut iteration = 0; + let mut accumulator = GradientsAccumulator::new(); + let mut accumulation_current = 0; + + while let Some(item) = iterator.next() { + for _ in 0..peer_count { + iteration += 1; + learner.lr_step(); + } + log::info!("Iteration {iteration}"); + + let mut progress = iterator.progress(); + progress.items_processed *= peer_count; + progress.items_total *= peer_count; + + let item = learner.train_step(item); + + match self.grad_accumulation { + Some(accumulation) => { + accumulator.accumulate(&learner.model(), item.grads); + accumulation_current += 1; + + if accumulation <= accumulation_current { + let grads = accumulator.grads(); + + learner.optimizer_step(grads); + accumulation_current = 0; + } + } + None => { + learner.optimizer_step(item.grads); + } + } + + let item = TrainingItem::new( + item.item, + progress, + Some(iteration), + Some(learner.lr_current()), + ); + + { + let mut processor = processor.lock().unwrap(); + processor.process_train(LearnerEvent::ProcessedItem(item)); + } + + if interrupter.should_stop() { + log::info!("Training interrupted."); + break; + } + } + } +} diff --git a/crates/burn-train/src/learner/supervised/strategies/ddp/mod.rs b/crates/burn-train/src/learner/supervised/strategies/ddp/mod.rs new file mode 100644 index 0000000..1828753 --- /dev/null +++ b/crates/burn-train/src/learner/supervised/strategies/ddp/mod.rs @@ -0,0 +1,5 @@ +mod epoch; +mod strategy; +mod worker; + +pub use strategy::*; diff --git a/crates/burn-train/src/learner/supervised/strategies/ddp/strategy.rs b/crates/burn-train/src/learner/supervised/strategies/ddp/strategy.rs new file mode 100644 index 0000000..bf90292 --- /dev/null +++ b/crates/burn-train/src/learner/supervised/strategies/ddp/strategy.rs @@ -0,0 +1,146 @@ +use core::panic; +use std::sync::{Arc, Mutex}; + +use crate::ddp::worker::DdpWorker; +use crate::metric::store::EventStoreClient; +use crate::{ + EarlyStoppingStrategyRef, Interrupter, Learner, LearnerModel, SupervisedLearningStrategy, + SupervisedTrainingEventProcessor, TrainLoader, TrainingComponents, ValidLoader, +}; +use burn_core::data::dataloader::split::split_dataloader; +use burn_core::tensor::Device; +use burn_core::tensor::distributed::DistributedContext; + +#[derive(Clone)] +pub(crate) struct WorkerComponents { + /// The total number of epochs + pub num_epochs: usize, + /// Enables gradients accumulation. + pub grad_accumulation: Option, + /// An [Interupter](Interrupter) that allows aborting the training/evaluation process early. + pub interrupter: Interrupter, + /// Cloneable reference to an early stopping strategy. + pub early_stopping: Option, + /// A reference to an [EventStoreClient](EventStoreClient). + pub event_store: Arc, + /// The total number of items in the training dataset. + pub train_total_items: usize, + /// The total number of items in the validation dataset. + pub valid_total_items: usize, +} + +/// A training strategy for Distributed Data Parallel (DDP) training. +/// +/// This strategy manages multiple workers and coordinates cross-device +/// gradient synchronization using the provided [`DistributedContext`]. +pub struct DdpTrainingStrategy { + devices: Vec, + /// Kept alive to anchor the lifetime of the underlying distributed server. + /// Spawns communication servers on creation, automatically tears them down on drop. + _context: DistributedContext, +} + +impl DdpTrainingStrategy { + pub fn new(devices: Vec, context: DistributedContext) -> Self { + Self { + devices, + _context: context, + } + } +} + +impl SupervisedLearningStrategy for DdpTrainingStrategy { + fn fit( + &self, + training_components: TrainingComponents, + learner: Learner, + dataloader_train: TrainLoader, + dataloader_valid: ValidLoader, + starting_epoch: usize, + ) -> (M, SupervisedTrainingEventProcessor) { + // The reference model is always on the first device provided. + let main_device = self.devices.first().unwrap(); + let train_total_items = dataloader_train.num_items(); + let valid_total_items = dataloader_valid.num_items(); + // One worker per device, so we use a fixed device strategy + // for each (worker) data loader. This matches the expected device on the worker, so we + // don't have to move the data between devices. + let mut dataloaders_train = split_dataloader(dataloader_train, &self.devices); + let dataloader_valid = dataloader_valid.to_device(&main_device.clone().inner()); + + let main_device = self.devices[0].clone(); + let peer_count = self.devices.len(); + let event_processor = Arc::new(Mutex::new(training_components.event_processor)); + + let interrupter = training_components.interrupter; + let worker_components = WorkerComponents { + num_epochs: training_components.num_epochs, + grad_accumulation: training_components.grad_accumulation, + interrupter: interrupter.clone(), + early_stopping: training_components.early_stopping, + event_store: training_components.event_store, + train_total_items, + valid_total_items, + }; + + // Start worker for main device + // First training dataloader corresponds to main device + let main_handle = DdpWorker::::start( + main_device.clone(), + learner.clone(), + event_processor.clone(), + worker_components.clone(), + training_components.checkpointer, + dataloaders_train.remove(0), + Some(dataloader_valid), + starting_epoch, + peer_count, + true, + ); + + // Spawn other workers for the other devices, starting with peer id 1 + let mut secondary_workers = vec![]; + for device in &self.devices[1..] { + let handle = DdpWorker::::start( + device.clone(), + learner.clone(), + event_processor.clone(), + worker_components.clone(), + None, + dataloaders_train.remove(0), + None, + starting_epoch, + peer_count, + false, + ); + + secondary_workers.push(handle); + } + + // Wait for all devices to finish + for worker in secondary_workers { + worker + .join() + .expect("Distributed data parallel worker failed"); + } + + // Main worker had the event processor + let model = main_handle + .join() + .expect("Distributed data parallel main worker failed"); + + if interrupter.should_stop() { + let reason = interrupter + .get_message() + .unwrap_or(String::from("Reason unknown")); + log::info!("Training interrupted: {reason}"); + } + let Ok(event_processor) = Arc::try_unwrap(event_processor) else { + panic!("Event processor still held!"); + }; + let Ok(event_processor) = event_processor.into_inner() else { + panic!("Event processor lock poisoned"); + }; + (model, event_processor) + } +} diff --git a/crates/burn-train/src/learner/supervised/strategies/ddp/worker.rs b/crates/burn-train/src/learner/supervised/strategies/ddp/worker.rs new file mode 100644 index 0000000..be36300 --- /dev/null +++ b/crates/burn-train/src/learner/supervised/strategies/ddp/worker.rs @@ -0,0 +1,142 @@ +use crate::ddp::epoch::{DdpTrainEpoch, DdpValidEpoch}; +use crate::ddp::strategy::WorkerComponents; +use crate::metric::processor::{EventProcessorTraining, LearnerEvent}; +use crate::single::TrainingLoop; +use crate::{ + Learner, LearnerModel, LearningCheckpointer, SupervisedTrainingEventProcessor, TrainLoader, + ValidLoader, +}; +use burn_core::tensor::Device; +use std::sync::{Arc, Mutex}; +use std::thread::JoinHandle; + +/// A worker runs the model, syncing gradients using collective operations. +/// Event processing and validation is optional too. +pub(crate) struct DdpWorker { + device: Device, + learner: Learner, + event_processor: Arc>>, + components: WorkerComponents, + checkpointer: Option>, + dataloader_train: TrainLoader, + dataloader_valid: Option>, + starting_epoch: usize, + peer_count: usize, + is_main: bool, +} + +impl DdpWorker { + /// Starts a worker that runs the model in a data distributed parallel + #[allow(clippy::too_many_arguments)] + pub fn start( + device: Device, + learner: Learner, + event_processor: Arc>>, + components: WorkerComponents, + checkpointer: Option>, + dataloader_train: TrainLoader, + dataloader_valid: Option>, + starting_epoch: usize, + peer_count: usize, + is_main: bool, + ) -> JoinHandle { + let worker = Self { + device, + learner, + event_processor, + components, + checkpointer, + dataloader_train, + dataloader_valid, + starting_epoch, + peer_count, + is_main, + }; + + std::thread::spawn(|| worker.fit()) + } + + /// Fits the model, + pub fn fit(mut self) -> M { + let num_epochs = self.components.num_epochs; + let interrupter = self.components.interrupter; + + // Changed the train epoch to keep the dataloaders + let epoch_train = DdpTrainEpoch::::new( + self.dataloader_train.clone(), + self.components.grad_accumulation, + ); + let epoch_valid = self + .dataloader_valid + .map(|dataloader| DdpValidEpoch::::new(dataloader)); + self.learner.fork(&self.device); + self.learner.grad_sharded(); + + for training_progress in TrainingLoop::new(self.starting_epoch, num_epochs) { + let epoch = training_progress.items_processed; + + if self.is_main { + self.event_processor + .lock() + .unwrap() + .process_train(LearnerEvent::StartSplit { + epoch_number: epoch, + total_items: self.components.train_total_items, + }); + } + + epoch_train.run( + &mut self.learner, + &training_progress, + self.event_processor.clone(), + &interrupter, + self.peer_count, + ); + + if self.is_main { + self.event_processor + .lock() + .unwrap() + .process_train(LearnerEvent::EndSplit(epoch)); + } + + if interrupter.should_stop() { + break; + } + + // Validation + if let Some(runner) = &epoch_valid { + { + self.event_processor + .lock() + .unwrap() + .process_valid(LearnerEvent::StartSplit { + epoch_number: epoch, + total_items: self.components.valid_total_items, + }); + } + let mut event_processor = self.event_processor.lock().unwrap(); + runner.run( + &self.learner.model(), + &training_progress, + &mut event_processor, + &interrupter, + ); + event_processor.process_valid(LearnerEvent::EndSplit(epoch)); + event_processor.process_train(LearnerEvent::EndEpoch(epoch)); + } + + if let Some(checkpointer) = &mut self.checkpointer { + checkpointer.checkpoint(&self.learner, epoch, &self.components.event_store); + } + + if let Some(early_stopping) = &mut self.components.early_stopping + && early_stopping.should_stop(epoch, &self.components.event_store) + { + break; + } + } + + self.learner.model() + } +} diff --git a/crates/burn-train/src/learner/supervised/strategies/mod.rs b/crates/burn-train/src/learner/supervised/strategies/mod.rs new file mode 100644 index 0000000..def5b8d --- /dev/null +++ b/crates/burn-train/src/learner/supervised/strategies/mod.rs @@ -0,0 +1,7 @@ +mod base; + +pub(crate) mod ddp; +pub(crate) mod multi; +pub(crate) mod single; + +pub use base::*; diff --git a/crates/burn-train/src/learner/supervised/strategies/multi/epoch.rs b/crates/burn-train/src/learner/supervised/strategies/multi/epoch.rs new file mode 100644 index 0000000..e6cc1fb --- /dev/null +++ b/crates/burn-train/src/learner/supervised/strategies/multi/epoch.rs @@ -0,0 +1,205 @@ +use crate::learner::base::Interrupter; +use crate::metric::processor::{EventProcessorTraining, LearnerEvent, TrainingItem}; +use crate::train::MultiDevicesTrainStep; +use crate::{ + Learner, LearnerModel, MultiDeviceOptim, SupervisedTrainingEventProcessor, TrainLoader, +}; +use burn_core::data::dataloader::Progress; +use burn_core::tensor::Device; +use burn_optim::GradientsAccumulator; +use burn_optim::MultiGradientsParams; + +/// A training epoch. +#[derive(new)] +pub struct MultiDeviceTrainEpoch { + dataloaders: Vec>, + grad_accumulation: Option, +} + +impl MultiDeviceTrainEpoch { + /// Runs the training epoch on multiple devices. + /// + /// # Arguments + /// + /// * `model` - The model to train. + /// * `optim` - The optimizer to use. + /// * `lr_scheduler` - The learning rate scheduler to use. + /// * `processor` - The event processor to use. + /// * `devices` - The devices to use. + /// + /// # Returns + /// + /// The trained model and the optimizer. + #[allow(clippy::too_many_arguments)] + pub fn run( + &self, + learner: &mut Learner, + global_progress: &Progress, + event_processor: &mut SupervisedTrainingEventProcessor, + interrupter: &Interrupter, + devices: Vec, + strategy: MultiDeviceOptim, + ) { + match strategy { + MultiDeviceOptim::OptimMainDevice => self.run_optim_main( + learner, + global_progress, + event_processor, + interrupter, + devices, + ), + MultiDeviceOptim::OptimSharded => self.run_optim_distr( + learner, + global_progress, + event_processor, + interrupter, + devices, + ), + } + } + + fn run_optim_main( + &self, + learner: &mut Learner, + global_progress: &Progress, + event_processor: &mut SupervisedTrainingEventProcessor, + interrupter: &Interrupter, + devices: Vec, + ) { + let epoch = global_progress.items_processed; + log::info!( + "Executing training step for epoch {} on devices {:?}", + epoch, + devices + ); + + let mut iterators = self + .dataloaders + .iter() + .map(|d| d.iter()) + .collect::>(); + let mut iteration = 0; + let mut accumulator = GradientsAccumulator::new(); + let mut accumulation_current = 0; + + let accumulation = self.grad_accumulation.unwrap_or(1); + let step = MultiDevicesTrainStep::::new(&devices); + + // The main device is always the first in the list. + let device_main = devices.first().expect("A minimum of one device.").clone(); + + loop { + let (items, progress) = step.step(iterators.as_mut_slice(), &learner.model()); + if items.is_empty() { + break; + } + + learner.lr_step(); + + let mut progress_items = Vec::with_capacity(items.len()); + for item in items.into_iter() { + let grads = item.output.grads.to_device(&device_main, &learner.model()); + accumulator.accumulate(&learner.model(), grads); + progress_items.push(item.output.item); + } + + accumulation_current += 1; + + if accumulation <= accumulation_current { + let grads = accumulator.grads(); + learner.optimizer_step(grads); + accumulation_current = 0; + } + + for item in progress_items { + iteration += 1; + let item = TrainingItem::new( + item, + progress.clone(), + Some(iteration), + Some(learner.lr_current()), + ); + + event_processor.process_train(LearnerEvent::ProcessedItem(item)); + } + + if interrupter.should_stop() { + break; + } + } + } + + fn run_optim_distr( + &self, + learner: &mut Learner, + global_progress: &Progress, + event_processor: &mut SupervisedTrainingEventProcessor, + interrupter: &Interrupter, + devices: Vec, + ) { + let epoch = global_progress.items_processed; + log::info!( + "Executing training step for epoch {} on devices {:?}", + epoch, + devices + ); + + let mut iterators = self + .dataloaders + .iter() + .map(|d| d.iter()) + .collect::>(); + let mut iteration = 0; + let mut accumulators: Vec> = (0..devices.len()) + .map(|_| GradientsAccumulator::new()) + .collect(); + let mut accumulation_current = 0; + + let accumulation = self.grad_accumulation.unwrap_or(1); + let step = MultiDevicesTrainStep::::new(&devices); + + loop { + let (items, progress) = step.step(iterators.as_mut_slice(), &learner.model()); + if items.is_empty() { + break; + } + + learner.lr_step(); + + let mut progress_items = Vec::with_capacity(items.len()); + for item in items.into_iter() { + let accumulator = &mut accumulators[item.device_id]; + accumulator.accumulate(&learner.model(), item.output.grads); + progress_items.push(item.output.item); + } + + accumulation_current += 1; + + if accumulation <= accumulation_current { + let mut grads = MultiGradientsParams::default(); + for (device_id, accumulator) in accumulators.iter_mut().enumerate() { + let grad = accumulator.grads(); + grads.grads.push((grad, devices[device_id].clone())); + } + learner.optimizer_step_multi(grads); + accumulation_current = 0; + } + + for item in progress_items { + iteration += 1; + let item = TrainingItem::new( + item, + progress.clone(), + Some(iteration), + Some(learner.lr_current()), + ); + + event_processor.process_train(LearnerEvent::ProcessedItem(item)); + } + + if interrupter.should_stop() { + break; + } + } + } +} diff --git a/crates/burn-train/src/learner/supervised/strategies/multi/mod.rs b/crates/burn-train/src/learner/supervised/strategies/multi/mod.rs new file mode 100644 index 0000000..b141f9d --- /dev/null +++ b/crates/burn-train/src/learner/supervised/strategies/multi/mod.rs @@ -0,0 +1,4 @@ +pub(crate) mod epoch; +mod strategy; + +pub use strategy::*; diff --git a/crates/burn-train/src/learner/supervised/strategies/multi/strategy.rs b/crates/burn-train/src/learner/supervised/strategies/multi/strategy.rs new file mode 100644 index 0000000..c7b626d --- /dev/null +++ b/crates/burn-train/src/learner/supervised/strategies/multi/strategy.rs @@ -0,0 +1,109 @@ +use crate::{ + Learner, LearnerEvent, LearnerModel, MultiDeviceOptim, SupervisedLearningStrategy, + SupervisedTrainingEventProcessor, TrainLoader, TrainingComponents, ValidLoader, + metric::processor::EventProcessorTraining, + multi::epoch::MultiDeviceTrainEpoch, + single::{TrainingLoop, epoch::SingleDeviceValidEpoch}, +}; +use burn_core::{data::dataloader::split::split_dataloader, tensor::Device}; + +pub struct MultiDeviceLearningStrategy { + devices: Vec, + optim: MultiDeviceOptim, +} +impl MultiDeviceLearningStrategy { + pub fn new(devices: Vec, optim: MultiDeviceOptim) -> Self { + Self { devices, optim } + } +} + +impl SupervisedLearningStrategy for MultiDeviceLearningStrategy { + fn fit( + &self, + training_components: TrainingComponents, + mut learner: Learner, + dataloader_train: TrainLoader, + dataloader_valid: ValidLoader, + starting_epoch: usize, + ) -> (M, SupervisedTrainingEventProcessor) { + let main_device = self.devices.first().unwrap(); + + // `MultiDevicesTrainStep` has one worker per device, so we use a fixed device strategy + // for each (worker) data loader. This matches the expected device on the worker, so we + // don't have to move the data between devices. + let train_total_items = dataloader_train.num_items(); + let dataloader_train = split_dataloader(dataloader_train, &self.devices); + let dataloader_valid = dataloader_valid.to_device(&main_device.clone().inner()); + let valid_total_items = dataloader_valid.num_items(); + + learner.fork(main_device); + let mut event_processor = training_components.event_processor; + let mut checkpointer = training_components.checkpointer; + let mut early_stopping = training_components.early_stopping; + + let epoch_train = MultiDeviceTrainEpoch::::new( + dataloader_train.clone(), + training_components.grad_accumulation, + ); + let epoch_valid: SingleDeviceValidEpoch = + SingleDeviceValidEpoch::new(dataloader_valid.clone()); + + for training_progress in TrainingLoop::new(starting_epoch, training_components.num_epochs) { + let epoch = training_progress.items_processed; + + event_processor.process_train(LearnerEvent::StartSplit { + epoch_number: epoch, + total_items: train_total_items, + }); + epoch_train.run( + &mut learner, + &training_progress, + &mut event_processor, + &training_components.interrupter, + self.devices.to_vec(), + self.optim, + ); + event_processor.process_train(LearnerEvent::EndSplit(epoch)); + + if training_components.interrupter.should_stop() { + let reason = training_components + .interrupter + .get_message() + .unwrap_or(String::from("Reason unknown")); + log::info!("Training interrupted: {reason}"); + break; + } + + // After OptimSharded training, model parameters are scattered across + // devices. Fork back to main_device before single-device validation. + if matches!(self.optim, MultiDeviceOptim::OptimSharded) { + learner.fork(main_device); + } + + event_processor.process_valid(LearnerEvent::StartSplit { + epoch_number: epoch, + total_items: valid_total_items, + }); + epoch_valid.run( + &learner, + &training_progress, + &mut event_processor, + &training_components.interrupter, + ); + event_processor.process_valid(LearnerEvent::EndSplit(epoch)); + event_processor.process_train(LearnerEvent::EndEpoch(epoch)); + + if let Some(checkpointer) = &mut checkpointer { + checkpointer.checkpoint(&learner, epoch, &training_components.event_store); + } + + if let Some(early_stopping) = &mut early_stopping + && early_stopping.should_stop(epoch, &training_components.event_store) + { + break; + } + } + + (learner.model(), event_processor) + } +} diff --git a/crates/burn-train/src/learner/supervised/strategies/single/epoch.rs b/crates/burn-train/src/learner/supervised/strategies/single/epoch.rs new file mode 100644 index 0000000..127f767 --- /dev/null +++ b/crates/burn-train/src/learner/supervised/strategies/single/epoch.rs @@ -0,0 +1,126 @@ +use crate::learner::base::Interrupter; +use crate::metric::processor::{EventProcessorTraining, LearnerEvent, TrainingItem}; +use crate::{ + InferenceStep, Learner, LearnerModel, SupervisedTrainingEventProcessor, TrainLoader, + ValidLoader, +}; +use burn_core::data::dataloader::Progress; +use burn_optim::GradientsAccumulator; + +/// A validation epoch. +#[derive(new)] +pub struct SingleDeviceValidEpoch { + dataloader: ValidLoader, +} + +/// A training epoch. +#[derive(new)] +pub struct SingleDeviceTrainEpoch { + dataloader: TrainLoader, + grad_accumulation: Option, +} + +impl SingleDeviceValidEpoch { + /// Runs the validation epoch. + /// + /// # Arguments + /// + /// * `model` - The model to validate. + /// * `processor` - The event processor to use. + pub fn run( + &self, + learner: &Learner, + global_progress: &Progress, + processor: &mut SupervisedTrainingEventProcessor, + interrupter: &Interrupter, + ) { + let epoch = global_progress.items_processed; + log::info!("Executing validation step for epoch {}", epoch); + let model = learner.model().valid(); + + let mut iterator = self.dataloader.iter(); + let mut iteration = 0; + + while let Some(item) = iterator.next() { + let progress = iterator.progress(); + iteration += 1; + + let item = InferenceStep::step(&model, item); + let item = TrainingItem::new(item, progress, Some(iteration), None); + + processor.process_valid(LearnerEvent::ProcessedItem(item)); + + if interrupter.should_stop() { + break; + } + } + } +} + +impl SingleDeviceTrainEpoch { + /// Runs the training epoch. + /// + /// # Arguments + /// + /// * `model` - The model to train. + /// * `optim` - The optimizer to use. + /// * `scheduler` - The learning rate scheduler to use. + /// * `processor` - The event processor to use. + /// + /// # Returns + /// + /// The trained model and the optimizer. + pub fn run( + &self, + learner: &mut Learner, + global_progress: &Progress, + processor: &mut SupervisedTrainingEventProcessor, + interrupter: &Interrupter, + ) { + let epoch = global_progress.items_processed; + log::info!("Executing training step for epoch {}", epoch,); + + // Single device / dataloader + let mut iterator = self.dataloader.iter(); + let mut iteration = 0; + let mut accumulator = GradientsAccumulator::new(); + let mut accumulation_current = 0; + + while let Some(item) = iterator.next() { + iteration += 1; + learner.lr_step(); + log::info!("Iteration {iteration}"); + + let progress = iterator.progress(); + let item = learner.train_step(item); + + match self.grad_accumulation { + Some(accumulation) => { + accumulator.accumulate(&learner.model(), item.grads); + accumulation_current += 1; + + if accumulation <= accumulation_current { + let grads = accumulator.grads(); + + learner.optimizer_step(grads); + accumulation_current = 0; + } + } + None => learner.optimizer_step(item.grads), + } + + let item = TrainingItem::new( + item.item, + progress, + Some(iteration), + Some(learner.lr_current()), + ); + + processor.process_train(LearnerEvent::ProcessedItem(item)); + + if interrupter.should_stop() { + break; + } + } + } +} diff --git a/crates/burn-train/src/learner/supervised/strategies/single/mod.rs b/crates/burn-train/src/learner/supervised/strategies/single/mod.rs new file mode 100644 index 0000000..b141f9d --- /dev/null +++ b/crates/burn-train/src/learner/supervised/strategies/single/mod.rs @@ -0,0 +1,4 @@ +pub(crate) mod epoch; +mod strategy; + +pub use strategy::*; diff --git a/crates/burn-train/src/learner/supervised/strategies/single/strategy.rs b/crates/burn-train/src/learner/supervised/strategies/single/strategy.rs new file mode 100644 index 0000000..e85561d --- /dev/null +++ b/crates/burn-train/src/learner/supervised/strategies/single/strategy.rs @@ -0,0 +1,118 @@ +use crate::{ + EventProcessorTraining, Learner, LearnerEvent, LearnerModel, SupervisedLearningStrategy, + SupervisedTrainingEventProcessor, TrainLoader, TrainingComponents, ValidLoader, + single::epoch::{SingleDeviceTrainEpoch, SingleDeviceValidEpoch}, +}; +use burn_core::{data::dataloader::Progress, tensor::Device}; + +/// Simplest learning strategy possible, with only a single devices doing both the training and +/// validation. +pub struct SingleDeviceTrainingStrategy { + device: Device, +} +impl SingleDeviceTrainingStrategy { + pub fn new(device: Device) -> Self { + Self { device } + } +} + +#[derive(new)] +pub(crate) struct TrainingLoop { + next_iteration: usize, + total_iteration: usize, +} + +/// An iterator that returns the progress of the training. +impl Iterator for TrainingLoop { + type Item = Progress; + + fn next(&mut self) -> Option { + if self.next_iteration > self.total_iteration { + return None; + } + + let progress = Progress { + items_processed: self.next_iteration, + items_total: self.total_iteration, + unit: Some("epochs".to_string()), + }; + + self.next_iteration += 1; + Some(progress) + } +} + +impl SupervisedLearningStrategy for SingleDeviceTrainingStrategy { + fn fit( + &self, + training_components: TrainingComponents, + mut learner: Learner, + dataloader_train: TrainLoader, + dataloader_valid: ValidLoader, + starting_epoch: usize, + ) -> (M, SupervisedTrainingEventProcessor) { + let dataloader_train = dataloader_train.to_device(&self.device); + let train_total_items = dataloader_train.num_items(); + let dataloader_valid = dataloader_valid.to_device(&self.device.clone().inner()); + let valid_total_items = dataloader_valid.num_items(); + learner.fork(&self.device); + let mut event_processor = training_components.event_processor; + let mut checkpointer = training_components.checkpointer; + let mut early_stopping = training_components.early_stopping; + + let epoch_train: SingleDeviceTrainEpoch = + SingleDeviceTrainEpoch::new(dataloader_train, training_components.grad_accumulation); + let epoch_valid: SingleDeviceValidEpoch = + SingleDeviceValidEpoch::new(dataloader_valid.clone()); + + for training_progress in TrainingLoop::new(starting_epoch, training_components.num_epochs) { + let epoch = training_progress.items_processed; + + event_processor.process_train(LearnerEvent::StartSplit { + epoch_number: epoch, + total_items: train_total_items, + }); + epoch_train.run( + &mut learner, + &training_progress, + &mut event_processor, + &training_components.interrupter, + ); + event_processor.process_train(LearnerEvent::EndSplit(epoch)); + + if training_components.interrupter.should_stop() { + let reason = training_components + .interrupter + .get_message() + .unwrap_or(String::from("Reason unknown")); + log::info!("Training interrupted: {reason}"); + break; + } + + event_processor.process_valid(LearnerEvent::StartSplit { + epoch_number: epoch, + total_items: valid_total_items, + }); + epoch_valid.run( + &learner, + &training_progress, + &mut event_processor, + &training_components.interrupter, + ); + event_processor.process_valid(LearnerEvent::EndSplit(epoch)); + event_processor.process_train(LearnerEvent::EndEpoch(epoch)); + + if let Some(checkpointer) = &mut checkpointer { + checkpointer.checkpoint(&learner, epoch, &training_components.event_store); + } + + if let Some(early_stopping) = &mut early_stopping + && early_stopping.should_stop(epoch, &training_components.event_store) + { + break; + } + } + + (learner.model(), event_processor) + } +} diff --git a/crates/burn-train/src/learner/train_val.rs b/crates/burn-train/src/learner/train_val.rs new file mode 100644 index 0000000..e588fc8 --- /dev/null +++ b/crates/burn-train/src/learner/train_val.rs @@ -0,0 +1,133 @@ +use crate::{ItemLazy, renderer::MetricsRenderer}; +use burn_core::{module::AutodiffModule, tensor::Gradients}; +use burn_optim::{ + GradientsParams, ModuleOptimizer, MultiGradientsParams, + lr_scheduler::module_lr_scheduler::ModuleLearningRate, +}; + +/// A training output. +pub struct TrainOutput { + /// The gradients. + pub grads: GradientsParams, + + /// The item. + pub item: TO, +} + +impl TrainOutput { + /// Creates a new training output. + /// + /// # Arguments + /// + /// * `module` - The module. + /// * `grads` - The gradients. + /// * `item` - The item. + /// + /// # Returns + /// + /// A new training output. + pub fn new(module: &M, grads: Gradients, item: TO) -> Self { + let grads = GradientsParams::from_grads(grads, module); + Self { grads, item } + } +} + +/// Trait to be implemented for models to be able to be trained. +/// +/// The [step](TrainStep::step) method needs to be manually implemented for all structs. +/// +/// The [optimize](TrainStep::optimize) method can be overridden if you want to control how the +/// optimizer is used to update the model. This can be useful if you want to call custom mutable +/// functions on your model (e.g., clipping the weights) before or after the optimizer is used. +/// +/// # Notes +/// +/// To be used with the [Learner](crate::Learner) struct, the struct which implements this trait must +/// also implement the [AutodiffModule] trait, which is done automatically with the +/// [Module](burn_core::module::Module) derive. +pub trait TrainStep { + /// Type of input for a step of the training stage. + type Input: Send + 'static; + /// Type of output for a step of the training stage. + type Output: ItemLazy + 'static; + /// Runs a step for training, which executes the forward and backward passes. + /// + /// # Arguments + /// + /// * `item` - The input for the model. + /// + /// # Returns + /// + /// The output containing the model output and the gradients. + fn step(&self, item: Self::Input) -> TrainOutput; + /// Optimize the current module with the provided gradients and learning rate. + /// + /// # Arguments + /// + /// * `optim`: Optimizer used for learning. + /// * `lr`: The learning rate used for this step. + /// * `grads`: The gradients of each parameter in the current model. + /// + /// # Returns + /// + /// The updated model. + fn optimize( + self, + optim: &mut ModuleOptimizer, + lr_module: ModuleLearningRate, + grads: GradientsParams, + ) -> Self + where + Self: AutodiffModule + Sized, + { + optim.step(lr_module, self, grads) + } + /// Optimize the current module with the provided gradients and learning rate. + /// + /// # Arguments + /// + /// * `optim`: Optimizer used for learning. + /// * `lr`: The learning rate used for this step. + /// * `grads`: Multiple gradients associated to each parameter in the current model. + /// + /// # Returns + /// + /// The updated model. + fn optimize_multi( + self, + optim: &mut ModuleOptimizer, + lr_module: ModuleLearningRate, + grads: MultiGradientsParams, + ) -> Self + where + Self: AutodiffModule + Sized, + { + optim.step_multi(lr_module, self, grads) + } +} + +/// Trait to be implemented for validating models. +pub trait InferenceStep { + /// Type of input for an inference step. + type Input: Send + 'static; + /// Type of output for an inference step. + type Output: ItemLazy + 'static; + /// Runs a validation step. + /// + /// # Arguments + /// + /// * `item` - The item to validate on. + /// + /// # Returns + /// + /// The validation output. + fn step(&self, item: Self::Input) -> Self::Output; +} + +/// The result of a training, containing the model along with the [renderer](MetricsRenderer). +pub struct LearningResult { + /// The model with the learned weights. + pub model: M, + /// The renderer that can be used for follow up training and evaluation. + pub renderer: Box, +} diff --git a/crates/burn-train/src/lib.rs b/crates/burn-train/src/lib.rs new file mode 100644 index 0000000..616dcaf --- /dev/null +++ b/crates/burn-train/src/lib.rs @@ -0,0 +1,112 @@ +#![warn(missing_docs)] +#![cfg_attr(docsrs, feature(doc_cfg))] + +//! A library for training neural networks using the burn crate. + +#[macro_use] +extern crate derive_new; + +/// The checkpoint module. +pub mod checkpoint; + +pub(crate) mod components; + +/// Renderer modules to display metrics and training information. +pub mod renderer; + +/// The logger module. +pub mod logger; + +/// The metric module. +pub mod metric; + +pub use metric::processor::*; + +mod learner; + +pub use learner::*; + +mod evaluator; + +pub use evaluator::*; + +pub use components::*; + +#[cfg(test)] +pub(crate) mod tests { + use burn_core::{prelude::Tensor, tensor::Bool}; + use std::default::Default; + + /// Probability of tp before adding errors + pub const THRESHOLD: f64 = 0.5; + + #[derive(Debug, Default)] + pub enum ClassificationType { + #[default] + Binary, + Multiclass, + Multilabel, + } + + /// Sample x Class shaped matrix for use in + /// classification metrics testing + pub fn dummy_classification_input( + classification_type: &ClassificationType, + ) -> (Tensor<2>, Tensor<2, Bool>) { + match classification_type { + ClassificationType::Binary => { + ( + Tensor::from_data([[0.3], [0.2], [0.7], [0.1], [0.55]], &Default::default()), + // targets + Tensor::from_data([[0], [1], [0], [0], [1]], &Default::default()), + // predictions @ threshold=0.5 + // [[0], [0], [1], [0], [1]] + ) + } + ClassificationType::Multiclass => { + ( + Tensor::from_data( + [ + [0.2, 0.8, 0.0], + [0.3, 0.6, 0.1], + [0.7, 0.25, 0.05], + [0.1, 0.15, 0.8], + [0.9, 0.03, 0.07], + ], + &Default::default(), + ), + Tensor::from_data( + // targets + [[0, 1, 0], [1, 0, 0], [0, 0, 1], [0, 0, 1], [1, 0, 0]], + // predictions @ top_k=1 + // [[0, 1, 0], [0, 1, 0], [1, 0, 0], [0, 0, 1], [1, 0, 0]] + // predictions @ top_k=2 + // [[1, 1, 0], [1, 1, 0], [1, 1, 0], [0, 1, 1], [1, 0, 1]] + &Default::default(), + ), + ) + } + ClassificationType::Multilabel => { + ( + Tensor::from_data( + [ + [0.1, 0.7, 0.6], + [0.3, 0.9, 0.05], + [0.8, 0.9, 0.4], + [0.7, 0.5, 0.9], + [1.0, 0.3, 0.2], + ], + &Default::default(), + ), + // targets + Tensor::from_data( + [[1, 1, 0], [1, 0, 1], [1, 1, 1], [0, 0, 1], [1, 0, 0]], + // predictions @ threshold=0.5 + // [[0, 1, 1], [0, 1, 0], [1, 1, 0], [1, 0, 1], [1, 0, 0]] + &Default::default(), + ), + ) + } + } + } +} diff --git a/crates/burn-train/src/logger/async_logger.rs b/crates/burn-train/src/logger/async_logger.rs new file mode 100644 index 0000000..c659098 --- /dev/null +++ b/crates/burn-train/src/logger/async_logger.rs @@ -0,0 +1,91 @@ +use super::Logger; +use std::sync::mpsc; + +enum Message { + Log(T), + End, + Sync(mpsc::Sender<()>), +} +/// Async logger. +pub struct AsyncLogger { + sender: mpsc::Sender>, + handler: Option>, +} + +#[derive(new)] +struct LoggerThread> { + logger: L, + receiver: mpsc::Receiver>, +} + +impl LoggerThread +where + L: Logger, +{ + fn run(mut self) { + for item in self.receiver.iter() { + match item { + Message::Log(item) => { + self.logger.log(item); + } + Message::End => { + return; + } + Message::Sync(callback) => { + callback + .send(()) + .expect("Can return result with the callback channel."); + } + } + } + } +} + +impl AsyncLogger { + /// Create a new async logger. + pub fn new(logger: L) -> Self + where + L: Logger + 'static, + { + let (sender, receiver) = mpsc::channel(); + let thread = LoggerThread::new(logger, receiver); + + let handler = Some(std::thread::spawn(move || thread.run())); + + Self { sender, handler } + } + + /// Sync the async logger. + pub(crate) fn sync(&self) { + let (sender, receiver) = mpsc::channel(); + + self.sender + .send(Message::Sync(sender)) + .expect("Can send message to logger thread."); + + receiver + .recv() + .expect("Should sync, otherwise the thread is dead."); + } +} + +impl Logger for AsyncLogger { + fn log(&mut self, item: T) { + self.sender + .send(Message::Log(item)) + .expect("Can log using the logger thread."); + } +} + +impl Drop for AsyncLogger { + fn drop(&mut self) { + self.sender + .send(Message::End) + .expect("Can send the end message to the logger thread."); + let handler = self.handler.take(); + + if let Some(handler) = handler { + handler.join().expect("The logger thread should stop."); + } + } +} diff --git a/crates/burn-train/src/logger/base.rs b/crates/burn-train/src/logger/base.rs new file mode 100644 index 0000000..5bb98ff --- /dev/null +++ b/crates/burn-train/src/logger/base.rs @@ -0,0 +1,9 @@ +/// The logger trait. +pub trait Logger: Send { + /// Logs an item. + /// + /// # Arguments + /// + /// * `item` - The item. + fn log(&mut self, item: T); +} diff --git a/crates/burn-train/src/logger/file.rs b/crates/burn-train/src/logger/file.rs new file mode 100644 index 0000000..c852c21 --- /dev/null +++ b/crates/burn-train/src/logger/file.rs @@ -0,0 +1,46 @@ +use super::Logger; +use std::{fs::File, io::Write, path::Path}; + +/// File logger. +pub struct FileLogger { + file: File, +} + +impl FileLogger { + /// Create a new file logger. + /// + /// # Arguments + /// + /// * `path` - The path. + /// + /// # Returns + /// + /// The file logger. + pub fn new(path: impl AsRef) -> Self { + let path = path.as_ref(); + let mut options = std::fs::File::options(); + let file = options + .write(true) + .truncate(true) + .create(true) + .open(path) + .unwrap_or_else(|err| { + panic!( + "Should be able to create the new file '{}': {}", + path.display(), + err + ) + }); + + Self { file } + } +} + +impl Logger for FileLogger +where + T: std::fmt::Display, +{ + fn log(&mut self, item: T) { + writeln!(&mut self.file, "{item}").expect("Can log an item."); + } +} diff --git a/crates/burn-train/src/logger/in_memory.rs b/crates/burn-train/src/logger/in_memory.rs new file mode 100644 index 0000000..31cf3f1 --- /dev/null +++ b/crates/burn-train/src/logger/in_memory.rs @@ -0,0 +1,16 @@ +use super::Logger; + +/// In memory logger. +#[derive(Default)] +pub struct InMemoryLogger { + pub(crate) values: Vec, +} + +impl Logger for InMemoryLogger +where + T: std::fmt::Display, +{ + fn log(&mut self, item: T) { + self.values.push(item.to_string()); + } +} diff --git a/crates/burn-train/src/logger/metric.rs b/crates/burn-train/src/logger/metric.rs new file mode 100644 index 0000000..22e59c4 --- /dev/null +++ b/crates/burn-train/src/logger/metric.rs @@ -0,0 +1,385 @@ +use super::{AsyncLogger, FileLogger, InMemoryLogger, Logger}; +use crate::metric::{ + MetricDefinition, MetricEntry, MetricId, NumericEntry, + store::{EpochSummary, MetricsUpdate, Split}, +}; +use std::{ + collections::HashMap, + fs, + path::{Path, PathBuf}, +}; + +const EPOCH_PREFIX: &str = "epoch-"; + +/// Metric logger. +pub trait MetricLogger: Send { + /// Logs an item. + /// + /// # Arguments + /// + /// * `update` - Update information for all registered metrics. + /// * `epoch` - Current epoch. + /// * `split` - Current dataset split. + fn log(&mut self, update: MetricsUpdate, epoch: usize, split: &Split); + + /// Read the logs for an epoch. + fn read_numeric( + &mut self, + name: &str, + epoch: usize, + split: &Split, + ) -> Result, String>; + + /// Logs the metric definition information (name, description, unit, etc.) + fn log_metric_definition(&mut self, definition: MetricDefinition); + + /// Logs summary at the end of the epoch. + fn log_epoch_summary(&mut self, summary: EpochSummary); +} + +/// The file metric logger. +pub struct FileMetricLogger { + loggers: HashMap>, + directory: PathBuf, + metric_definitions: HashMap, + is_eval: bool, + last_epoch: Option, +} + +impl FileMetricLogger { + /// Create a new file metric logger. + /// + /// # Arguments + /// + /// * `directory` - The directory. + /// + /// # Returns + /// + /// The file metric logger. + pub fn new(directory: impl AsRef) -> Self { + Self { + loggers: HashMap::new(), + directory: directory.as_ref().to_path_buf(), + metric_definitions: HashMap::default(), + is_eval: false, + last_epoch: None, + } + } + + /// Create a new file metric logger. + /// + /// # Arguments + /// + /// * `directory` - The directory. + /// + /// # Returns + /// + /// The file metric logger. + pub fn new_eval(directory: impl AsRef) -> Self { + Self { + loggers: HashMap::new(), + directory: directory.as_ref().to_path_buf(), + metric_definitions: HashMap::default(), + is_eval: true, + last_epoch: None, + } + } + + pub(crate) fn split_exists(&self, split: &Split) -> bool { + self.split_dir(split).is_some() + } + + pub(crate) fn split_dir(&self, split: &Split) -> Option { + let split_path = match split { + Split::Test(Some(tag)) => self.directory.join(split.to_string()).join(tag.as_str()), + other => self.directory.join(other.to_string()), + }; + (split_path.exists() && split_path.is_dir()).then_some(split_path) + } + + pub(crate) fn is_epoch_dir>(dirname: P) -> bool { + dirname.as_ref().starts_with(EPOCH_PREFIX) + } + + /// Number of epochs recorded. + pub(crate) fn epochs(&self) -> usize { + if self.is_eval { + log::warn!("Number of epochs not available when testing."); + return 0; + } + + let mut max_epoch = 0; + + // with split + for path in fs::read_dir(&self.directory).unwrap() { + let path = path.unwrap(); + + if fs::metadata(path.path()).unwrap().is_dir() { + for split_path in fs::read_dir(path.path()).unwrap() { + let split_path = split_path.unwrap(); + + if fs::metadata(split_path.path()).unwrap().is_dir() { + let dir_name = split_path.file_name().into_string().unwrap(); + + if !dir_name.starts_with(EPOCH_PREFIX) { + continue; + } + + let epoch = dir_name.replace(EPOCH_PREFIX, "").parse::().ok(); + + if let Some(epoch) = epoch + && epoch > max_epoch + { + max_epoch = epoch; + } + } + } + } + } + + max_epoch + } + + fn train_directory(&self, epoch: usize, split: &Split) -> PathBuf { + let name = format!("{EPOCH_PREFIX}{epoch}"); + + match split { + Split::Train | Split::Valid | Split::Test(None) => { + self.directory.join(split.to_string()).join(name) + } + Split::Test(Some(tag)) => { + let tag = format_tag(tag); + self.directory.join(split.to_string()).join(tag).join(name) + } + } + } + + fn eval_directory(&self, split: &Split) -> PathBuf { + match split { + Split::Train | Split::Valid | Split::Test(None) => self.directory.clone(), + Split::Test(Some(tag)) => self.directory.join(split.to_string()).join(format_tag(tag)), + } + } + + fn file_path(&self, name: &str, epoch: Option, split: &Split) -> PathBuf { + let directory = match epoch { + Some(epoch) => self.train_directory(epoch, split), + None => self.eval_directory(split), + }; + let name = name.replace(' ', "_"); + let name = format!("{name}.log"); + directory.join(name) + } + + fn create_directory(&self, epoch: Option, split: &Split) { + let directory = match epoch { + Some(epoch) => self.train_directory(epoch, split), + None => self.eval_directory(split), + }; + std::fs::create_dir_all(directory).ok(); + } +} + +impl FileMetricLogger { + fn log_item(&mut self, item: &MetricEntry, epoch: Option, split: &Split) { + let name = &self.metric_definitions.get(&item.metric_id).unwrap().name; + let key = logger_key(name, split); + let value = &item.serialized_entry.serialized; + + let logger = match self.loggers.get_mut(&key) { + Some(val) => val, + None => { + self.create_directory(epoch, split); + + let file_path = self.file_path(name, epoch, split); + let logger = FileLogger::new(file_path); + let logger = AsyncLogger::new(logger); + + self.loggers.insert(key.clone(), logger); + self.loggers + .get_mut(&key) + .expect("Can get the previously saved logger.") + } + }; + + logger.log(value.clone()); + } +} + +fn format_tag(tag: &str) -> String { + tag.trim().replace(' ', "-").to_lowercase() +} + +impl MetricLogger for FileMetricLogger { + fn log(&mut self, update: MetricsUpdate, epoch: usize, split: &Split) { + if !self.is_eval && self.last_epoch != Some(epoch) { + self.loggers.clear(); + self.last_epoch = Some(epoch); + } + + let entries: Vec<_> = update + .entries + .iter() + .chain( + update + .entries_numeric + .iter() + .map(|numeric_update| &numeric_update.entry), + ) + .cloned() + .collect(); + + let epoch = if !self.is_eval { Some(epoch) } else { None }; + for item in entries.iter() { + self.log_item(item, epoch, split); + } + } + + fn read_numeric( + &mut self, + name: &str, + epoch: usize, + split: &Split, + ) -> Result, String> { + if let Some(value) = self.loggers.get(name) { + value.sync() + } + + let file_path = self.file_path(name, Some(epoch), split); + + let mut errors = false; + + let data = std::fs::read_to_string(file_path) + .unwrap_or_default() + .split('\n') + .filter_map(|value| { + if value.is_empty() { + None + } else { + match NumericEntry::deserialize(value) { + Ok(value) => Some(value), + Err(err) => { + log::error!("{err}"); + errors = true; + None + } + } + } + }) + .collect(); + + if errors { + Err("Parsing numeric entry errors".to_string()) + } else { + Ok(data) + } + } + + fn log_metric_definition(&mut self, definition: MetricDefinition) { + self.metric_definitions + .insert(definition.metric_id.clone(), definition); + } + + fn log_epoch_summary(&mut self, _summary: EpochSummary) { + if !self.is_eval { + self.loggers.clear(); + } + } +} + +fn logger_key(name: &str, split: &Split) -> String { + match split { + Split::Train | Split::Valid => format!("{name}_{split}"), + Split::Test(tag) => { + if let Some(t) = tag { + format!("{name}_{split}_{}", *t) + } else { + format!("{name}_{split}") + } + } + } +} + +/// In memory metric logger, useful when testing and debugging. +#[derive(Default)] +pub struct InMemoryMetricLogger { + values: HashMap>, + last_epoch: Option, + metric_definitions: HashMap, +} + +impl InMemoryMetricLogger { + /// Create a new in-memory metric logger. + pub fn new() -> Self { + Self::default() + } +} + +impl MetricLogger for InMemoryMetricLogger { + fn log(&mut self, update: MetricsUpdate, epoch: usize, split: &Split) { + if self.last_epoch != Some(epoch) { + self.values + .values_mut() + .for_each(|loggers| loggers.push(InMemoryLogger::default())); + self.last_epoch = Some(epoch); + } + + let entries: Vec<_> = update + .entries + .iter() + .chain( + update + .entries_numeric + .iter() + .map(|numeric_update| &numeric_update.entry), + ) + .cloned() + .collect(); + + for item in entries.iter() { + let name = &self.metric_definitions.get(&item.metric_id).unwrap().name; + let key = logger_key(name, split); + + if !self.values.contains_key(&key) { + self.values + .insert(key.to_string(), vec![InMemoryLogger::default()]); + } + + let values = self.values.get_mut(&key).unwrap(); + + values + .last_mut() + .unwrap() + .log(item.serialized_entry.serialized.clone()); + } + } + + fn read_numeric( + &mut self, + name: &str, + epoch: usize, + split: &Split, + ) -> Result, String> { + let key = logger_key(name, split); + let values = match self.values.get(&key) { + Some(values) => values, + None => return Ok(Vec::new()), + }; + + match values.get(epoch - 1) { + Some(logger) => Ok(logger + .values + .iter() + .filter_map(|value| NumericEntry::deserialize(value).ok()) + .collect()), + None => Ok(Vec::new()), + } + } + + fn log_metric_definition(&mut self, definition: MetricDefinition) { + self.metric_definitions + .insert(definition.metric_id.clone(), definition); + } + + fn log_epoch_summary(&mut self, _summary: EpochSummary) {} +} diff --git a/crates/burn-train/src/logger/mod.rs b/crates/burn-train/src/logger/mod.rs new file mode 100644 index 0000000..1d41ced --- /dev/null +++ b/crates/burn-train/src/logger/mod.rs @@ -0,0 +1,13 @@ +mod async_logger; +mod base; +mod file; +mod in_memory; +mod metric; +mod progress; + +pub use async_logger::*; +pub use base::*; +pub use file::*; +pub use in_memory::*; +pub use metric::*; +pub use progress::*; diff --git a/crates/burn-train/src/logger/progress.rs b/crates/burn-train/src/logger/progress.rs new file mode 100644 index 0000000..a2f2a8d --- /dev/null +++ b/crates/burn-train/src/logger/progress.rs @@ -0,0 +1,110 @@ +use burn_core::data::dataloader::Progress; + +/// Trait for logging training progress at each step and end of epoch. +/// +/// # Call sequence +/// +/// Implementors can expect the following sequence of calls for a complete training run: +/// +/// ```text +/// start(total_epochs, total_items) +/// for each epoch: +/// start_split("train", total_items_train) +/// update_split(items_processed) // called once per batch +/// ... +/// end_split() +/// start_split("valid", total_items_valid) +/// update_split(items_processed) // called once per batch +/// ... +/// end_split() +/// update_epoch(epoch) +/// end() +/// ``` +/// +/// `end()` is called whether training completes normally or is interrupted early. +/// +/// Implementors are responsible for tracking `total_items` and epoch state in order +/// to reconstruct the full progress picture when `update_split` is called. +pub trait TrainingProgressLogger: Send { + /// Called once at the start of training, providing the total number of epochs. + /// + /// The total number of items of the training can optionally be provided if it is known. + fn start(&mut self, total_epochs: usize, starting_epoch: usize, total_items: Option); + + /// Called at the end of each full epoch (after both train and valid splits complete). + fn update_epoch(&mut self, epoch: usize); + + /// Called at the start of a training split, providing the split name and total number of items. + fn start_split(&mut self, split: &str, total_items: usize); + + /// Log the progress of the current training step. + fn update_split(&mut self, items_processed: usize); + + /// Called at the end of a training split. + fn end_split(&mut self); + + /// Called at the end of training, whether it completed successfully or was interrupted. + fn end(&mut self); + + /// Log a custom named event that falls outside the standard training lifecycle callbacks. + fn log_event_training(&mut self, event: String); +} + +/// Trait for logging evaluation progress at each step and end of evaluation. +/// +/// # Call sequence +/// +/// Implementors can expect the following sequence of calls for a complete evaluation run: +/// +/// ```text +/// start_global_progress(total_tests) +/// for each test split: +/// start_test(name, total_items) +/// update_test_progress(items_processed) // called once per batch +/// ... +/// end_test() +/// end_global_progress() +/// ``` +/// +/// `end_global_progress()` is called whether evaluation completes normally or is interrupted early. +/// +/// Implementors are responsible for tracking `total_tests` and `total_items` (stored from +/// `start_global_progress` and `start_test`) to reconstruct the full progress picture. +pub trait EvaluationProgressLogger: Send { + /// Called once at the start of evaluation, providing the total number of test splits. + fn start_global_progress(&mut self, total_tests: usize); + + /// Called at the start of a test split, providing the split name and total number of items. + fn start_test(&mut self, name: &str, total_items: usize); + + /// Log the progress of the current test step. + fn update_test_progress(&mut self, items_processed: usize); + + /// Called at the end of a test split. + fn end_test(&mut self); + + /// Called at the end of evaluation. + fn end_global_progress(&mut self); + + /// Log a custom named event that falls outside the standard evaluation lifecycle callbacks. + fn log_event_evaluation(&mut self, event: String); +} + +/// Two-level progress snapshot combining run-level and phase-level tracking. +/// +/// `global_progress` spans the full training run (e.g., epochs completed out of total), +/// while `split_progress` tracks the current phase (e.g., batches within the current epoch). +#[derive(Debug, Clone)] +pub struct ProgressSnapshot { + /// Progress across the entire training run. + pub global: Progress, + /// Progress within the current phase (epoch or evaluation split). + pub split: Progress, +} + +impl ProgressSnapshot { + /// Create a new overall progress snapshot. + pub fn new(global: Progress, split: Progress) -> Self { + Self { global, split } + } +} diff --git a/crates/burn-train/src/metric/acc.rs b/crates/burn-train/src/metric/acc.rs new file mode 100644 index 0000000..68fc751 --- /dev/null +++ b/crates/burn-train/src/metric/acc.rs @@ -0,0 +1,151 @@ +use super::MetricMetadata; +use super::state::{FormatOptions, NumericMetricState}; +use crate::metric::{Metric, MetricAttributes, MetricName, Numeric, SerializedEntry}; +use burn_core::tensor::{Int, Tensor}; + +/// The accuracy metric. +#[derive(Clone)] +pub struct AccuracyMetric { + name: MetricName, + state: NumericMetricState, + pad_token: Option, +} + +/// The [accuracy metric](AccuracyMetric) input type. +#[derive(new)] +pub struct AccuracyInput { + outputs: Tensor<2>, + targets: Tensor<1, Int>, +} + +impl Default for AccuracyMetric { + fn default() -> Self { + Self::new() + } +} + +impl AccuracyMetric { + /// Creates the metric. + pub fn new() -> Self { + Self { + name: MetricName::new("Accuracy".to_string()), + state: Default::default(), + pad_token: Default::default(), + } + } + + /// Sets the pad token. + pub fn with_pad_token(mut self, index: usize) -> Self { + self.pad_token = Some(index); + self + } +} + +impl Metric for AccuracyMetric { + type Input = AccuracyInput; + + fn update(&mut self, input: &AccuracyInput, _metadata: &MetricMetadata) -> SerializedEntry { + let targets = input.targets.clone(); + let outputs = input.outputs.clone(); + + let [batch_size, _n_classes] = outputs.dims(); + + let outputs = outputs.argmax(1).reshape([batch_size]); + + let accuracy = match self.pad_token { + Some(pad_token) => { + let mask = targets.clone().equal_elem(pad_token as i64); + let matches = outputs.equal(targets).float().mask_fill(mask.clone(), 0); + let num_pad = mask.float().sum(); + + let acc = matches.sum() / (num_pad.neg() + batch_size as f32); + + acc.into_scalar::() + } + None => outputs.equal(targets).int().sum().into_scalar::() / batch_size as f64, + }; + + self.state.update( + 100.0 * accuracy, + batch_size, + FormatOptions::new(self.name()).unit("%").precision(2), + ) + } + + fn clear(&mut self) { + self.state.reset() + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + super::NumericAttributes { + unit: Some("%".to_string()), + higher_is_better: true, + ..Default::default() + } + .into() + } +} + +impl Numeric for AccuracyMetric { + fn value(&self) -> super::NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> super::NumericEntry { + self.state.running_value() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_accuracy_without_padding() { + let device = Default::default(); + let mut metric = AccuracyMetric::new(); + let input = AccuracyInput::new( + Tensor::from_data( + [ + [0.0, 0.2, 0.8], // 2 + [1.0, 2.0, 0.5], // 1 + [0.4, 0.1, 0.2], // 0 + [0.6, 0.7, 0.2], // 1 + ], + &device, + ), + Tensor::from_data([2, 2, 1, 1], &device), + ); + + let _entry = metric.update(&input, &MetricMetadata::fake()); + assert_eq!(50.0, metric.value().current()); + } + + #[test] + fn test_accuracy_with_padding() { + let device = Default::default(); + let mut metric = AccuracyMetric::new().with_pad_token(3); + let input = AccuracyInput::new( + Tensor::from_data( + [ + [0.0, 0.2, 0.8, 0.0], // 2 + [1.0, 2.0, 0.5, 0.0], // 1 + [0.4, 0.1, 0.2, 0.0], // 0 + [0.6, 0.7, 0.2, 0.0], // 1 + [0.0, 0.1, 0.2, 5.0], // Predicted padding should not count + [0.0, 0.1, 0.2, 0.0], // Error on padding should not count + [0.6, 0.0, 0.2, 0.0], // Error on padding should not count + ], + &device, + ), + Tensor::from_data([2, 2, 1, 1, 3, 3, 3], &device), + ); + + let _entry = metric.update(&input, &MetricMetadata::fake()); + assert_eq!(50.0, metric.value().current()); + } +} diff --git a/crates/burn-train/src/metric/auc_pr.rs b/crates/burn-train/src/metric/auc_pr.rs new file mode 100644 index 0000000..e1744a1 --- /dev/null +++ b/crates/burn-train/src/metric/auc_pr.rs @@ -0,0 +1,310 @@ +use super::MetricMetadata; +use super::state::{FormatOptions, PredictionAccumulatorState}; +use crate::metric::{ + ClassReduction, ConfusionStatsInput, Metric, MetricAttributes, MetricName, Numeric, + NumericAggregation, NumericAttributes, SerializedEntry, +}; +use burn_core::tensor::{Int, Tensor}; +use std::sync::Arc; + +/// The Area Under the Precision-Recall Curve (AUC-PR). +/// +/// Computed as **Average Precision** — `AP = Σ (Rₙ − Rₙ₋₁) · Pₙ` — the +/// standard non-interpolated estimator of the area under the +/// precision-recall curve (equivalent to scikit-learn's +/// `average_precision_score`), not the (biased) trapezoidal integration. +/// +/// Supports binary, multiclass and multi-label classification through a +/// One-vs-Rest decomposition, aggregated with the configured +/// [class reduction](ClassReduction). +#[derive(Clone)] +pub struct AucPrMetric { + name: MetricName, + state: PredictionAccumulatorState, + class_reduction: ClassReduction, +} + +impl Default for AucPrMetric { + fn default() -> Self { + Self::new(Default::default()) + } +} + +impl AucPrMetric { + fn new(class_reduction: ClassReduction) -> Self { + let state = Default::default(); + let name = Arc::new(format!("AUC-PR [{:?}]", class_reduction)); + + Self { + state, + class_reduction, + name, + } + } + + /// AUC-PR metric for binary classification. + #[allow(dead_code)] + pub fn binary() -> Self { + Self::new(ClassReduction::default()) + } + + /// AUC-PR metric for multiclass classification. + /// + /// # Arguments + /// + /// * `class_reduction` - [Class reduction](ClassReduction) type. + #[allow(dead_code)] + pub fn multiclass(class_reduction: ClassReduction) -> Self { + Self::new(class_reduction) + } + + /// AUC-PR metric for multi-label classification. + /// + /// # Arguments + /// + /// * `class_reduction` - [Class reduction](ClassReduction) type. + #[allow(dead_code)] + pub fn multilabel(class_reduction: ClassReduction) -> Self { + Self::new(class_reduction) + } + + /// Per-column Average Precision via the step-wise estimator + /// `AP = (1/P) · Σ_{positives, score desc} (cumulative positives / rank)`. + /// + /// `scores` and `targets` are `[n, c]` (`targets` as 0./1.); a column + /// with no positive (`P = 0`) yields `NaN` (handled by the caller). + fn average_precision(scores: Tensor<2>, targets: Tensor<2>) -> Tensor<1> { + let [n, _c] = scores.dims(); + let device = scores.device(); + + let order = scores.argsort_descending(0); + let sorted_targets = targets.clone().gather(0, order); + + let tp = sorted_targets.clone().cumsum(0); + + let ranks = Tensor::<1, Int>::arange(1..n as i64 + 1, &device) + .float() + .reshape([n, 1]); + + let precision = tp / ranks; + + let p_total = targets.sum_dim(0); + let delta_recall = sorted_targets / p_total; + + (precision * delta_recall) + .sum_dim(0) + .squeeze_dims::<1>(&[0]) + } +} + +impl Metric for AucPrMetric { + type Input = ConfusionStatsInput; + + fn update( + &mut self, + input: &ConfusionStatsInput, + _metadata: &MetricMetadata, + ) -> SerializedEntry { + self.state + .accumulate(input.predictions.clone(), input.targets.clone()); + + // Recompute over the whole epoch: AP is rank-based, not a per-batch mean. + let (predictions, targets) = self.state.tensors(); + let [n, c] = predictions.dims(); + + let (scores, targets) = match self.class_reduction { + ClassReduction::Macro => (predictions, targets.float()), + ClassReduction::Micro => ( + predictions.reshape([n * c, 1]), + targets.float().reshape([n * c, 1]), + ), + }; + + let ap = Self::average_precision(scores, targets); + + let keep = ap + .clone() + .is_nan() + .bool_not() + .argwhere() + .squeeze_dim::<1>(1); + + let metric = if keep.dims()[0] == 0 { + log::warn!( + "AUC-PR is undefined (no class has positive samples in the epoch); reporting \ + 0.5 as a neutral fallback." + ); + 0.5 + } else { + ap.select(0, keep).mean().into_scalar() + }; + + self.state.update( + 100.0 * metric, + FormatOptions::new(self.name()).unit("%").precision(2), + ) + } + + fn clear(&mut self) { + self.state.reset() + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: Some("%".to_string()), + higher_is_better: true, + aggregation: NumericAggregation::Last, + } + .into() + } +} + +impl Numeric for AucPrMetric { + fn value(&self) -> super::NumericEntry { + self.state.value() + } + + fn running_value(&self) -> super::NumericEntry { + self.state.value() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::metric::ClassReduction::{self, *}; + use burn_core::tensor::{TensorData, Tolerance}; + use rstest::rstest; + + /// Inputs and expected Average Precision computed with an independent + /// reference equivalent to scikit-learn's `average_precision_score` + /// (step-wise `AP = Σ (Rₙ−Rₙ₋₁)·Pₙ`). Scores are distinct so the + /// estimator is unambiguous and matches sklearn exactly. + #[derive(Clone, Copy)] + enum Data { + Binary, + Multiclass, + Multilabel, + } + + fn input(data: Data) -> ConfusionStatsInput { + let dev = Default::default(); + match data { + Data::Binary => ConfusionStatsInput::new( + Tensor::from_data([[0.63], [0.25], [0.71], [0.3], [0.07], [0.66]], &dev), + Tensor::from_data([[0], [1], [0], [0], [0], [0]], &dev), + ), + Data::Multiclass => ConfusionStatsInput::new( + Tensor::from_data( + [ + [0.45, 0.3, 0.36], + [0.83, 0.24, 0.09], + [0.19, 0.39, 0.29], + [0.3, 0.14, 0.46], + [0.73, 0.74, 0.16], + [0.43, 0.37, 0.88], + ], + &dev, + ), + Tensor::from_data( + [ + [0, 0, 1], + [0, 0, 1], + [0, 0, 1], + [0, 0, 1], + [0, 1, 0], + [1, 0, 0], + ], + &dev, + ), + ), + Data::Multilabel => ConfusionStatsInput::new( + Tensor::from_data( + [ + [0.1, 0.73, 0.84], + [0.84, 0.74, 0.24], + [0.13, 0.54, 0.54], + [0.49, 0.48, 0.71], + [0.9, 0.17, 0.43], + [0.11, 0.29, 0.23], + ], + &dev, + ), + Tensor::from_data( + [ + [1, 0, 1], + [0, 0, 1], + [0, 0, 1], + [0, 0, 1], + [1, 0, 0], + [1, 1, 0], + ], + &dev, + ), + ), + } + } + + #[rstest] + // Binary is a single column -> Macro == Micro. + #[case::binary_macro(Data::Binary, Macro, 0.2)] + #[case::binary_micro(Data::Binary, Micro, 0.2)] + #[case::multiclass_macro(Data::Multiclass, Macro, 0.6319444444444444)] + #[case::multiclass_micro(Data::Multiclass, Micro, 0.379975579975580)] + #[case::multilabel_macro(Data::Multilabel, Macro, 0.5944444444444444)] + #[case::multilabel_micro(Data::Multilabel, Micro, 0.5918017848017848)] + fn test_auc_pr( + #[case] data: Data, + #[case] class_reduction: ClassReduction, + #[case] expected: f64, + ) { + let mut metric = AucPrMetric::new(class_reduction); + + let _entry = metric.update(&input(data), &MetricMetadata::fake()); + + TensorData::from([metric.value().current()]) + .assert_approx_eq::(&TensorData::from([expected * 100.0]), Tolerance::default()); + } + + #[test] + fn test_auc_pr_accumulates_across_batches() { + let dev = Default::default(); + + // Whole dataset as a single batch. + let mut single = AucPrMetric::binary(); + single.update( + &ConfusionStatsInput::new( + Tensor::from_data([[0.9], [0.4], [0.8], [0.2], [0.6], [0.1]], &dev), + Tensor::from_data([[1], [0], [1], [0], [1], [0]], &dev), + ), + &MetricMetadata::fake(), + ); + + // Same dataset split across two batches. + let mut split = AucPrMetric::binary(); + split.update( + &ConfusionStatsInput::new( + Tensor::from_data([[0.9], [0.4], [0.8]], &dev), + Tensor::from_data([[1], [0], [1]], &dev), + ), + &MetricMetadata::fake(), + ); + split.update( + &ConfusionStatsInput::new( + Tensor::from_data([[0.2], [0.6], [0.1]], &dev), + Tensor::from_data([[0], [1], [0]], &dev), + ), + &MetricMetadata::fake(), + ); + + // Epoch value is independent of batching. + TensorData::from([split.value().current()]).assert_approx_eq::( + &TensorData::from([single.value().current()]), + Tolerance::default(), + ); + } +} diff --git a/crates/burn-train/src/metric/auroc.rs b/crates/burn-train/src/metric/auroc.rs new file mode 100644 index 0000000..ad1f9b0 --- /dev/null +++ b/crates/burn-train/src/metric/auroc.rs @@ -0,0 +1,335 @@ +use core::f64; + +use super::MetricMetadata; +use super::state::{FormatOptions, NumericMetricState}; +use crate::metric::{ + ClassReduction, ConfusionStatsInput, Metric, MetricName, Numeric, SerializedEntry, +}; +use burn_core::tensor::{Bool, Tensor}; +use std::sync::Arc; + +/// The Area Under the Receiver Operating Characteristic Curve (AUROC, also +/// referred to as [ROC AUC](https://en.wikipedia.org/wiki/Receiver_operating_characteristic)). +/// +/// Supports binary, multiclass and multi-label classification through a +/// One-vs-Rest decomposition, aggregated with the configured +/// [class reduction](ClassReduction). +#[derive(Clone)] +pub struct AurocMetric { + name: MetricName, + state: NumericMetricState, + class_reduction: ClassReduction, +} + +impl Default for AurocMetric { + fn default() -> Self { + Self::new(Default::default()) + } +} + +impl AurocMetric { + fn new(class_reduction: ClassReduction) -> Self { + let state = Default::default(); + let name = Arc::new(format!("AUROC [{:?}]", class_reduction)); + + Self { + state, + class_reduction, + name, + } + } + + /// AUROC metric for binary classification. + #[allow(dead_code)] + pub fn binary() -> Self { + Self::new(ClassReduction::default()) + } + + /// AUROC metric for multiclass classification. + /// + /// # Arguments + /// + /// * `class_reduction` - [Class reduction](ClassReduction) type. + #[allow(dead_code)] + pub fn multiclass(class_reduction: ClassReduction) -> Self { + Self::new(class_reduction) + } + + /// AUROC metric for multi-label classification. + /// + /// # Arguments + /// + /// * `class_reduction` - [Class reduction](ClassReduction) type. + #[allow(dead_code)] + pub fn multilabel(class_reduction: ClassReduction) -> Self { + Self::new(class_reduction) + } + + fn pairwise_auc(scores: Tensor<2>, targets: Tensor<2>) -> Tensor<1> { + let [n, c] = scores.dims(); + + let si = scores.clone().reshape([n, 1, c]); + let sj = scores.reshape([1, n, c]); + + let yi = targets.clone().reshape([n, 1, c]); + let yj = targets.reshape([1, n, c]); + + let valid: Tensor<3> = yi * (1.0 - yj); + + let reduce = |t: Tensor<3>| t.sum_dim(0).sum_dim(1).squeeze_dims::<1>(&[0, 1]); + + let num_pairs = reduce(valid.clone()); + let correct_pairs = reduce(si.clone().greater(sj.clone()).float() * valid.clone()); + let tied_pairs = reduce(si.equal(sj).float() * valid); + + (correct_pairs + 0.5 * tied_pairs) / num_pairs + } + + fn compute_auc(&self, predictions: &Tensor<2>, targets: &Tensor<2, Bool>) -> f64 { + let [n, c] = predictions.dims(); + + let (scores, targets) = match self.class_reduction { + ClassReduction::Macro => (predictions.clone(), targets.clone().float()), + ClassReduction::Micro => ( + predictions.clone().reshape([n * c, 1]), + targets.clone().float().reshape([n * c, 1]), + ), + }; + + let auc = Self::pairwise_auc(scores, targets); + + let keep = auc + .clone() + .is_nan() + .bool_not() + .argwhere() + .squeeze_dim::<1>(1); + + if keep.dims()[0] == 0 { + log::warn!( + "AUROC is undefined (no class has both positive and negative samples in the \ + batch); reporting 0.5 (chance level)." + ); + return 0.5; + } + + auc.select(0, keep).mean().into_scalar() + } +} + +impl Metric for AurocMetric { + type Input = ConfusionStatsInput; + + fn update( + &mut self, + input: &ConfusionStatsInput, + _metadata: &MetricMetadata, + ) -> SerializedEntry { + let [sample_size, _] = input.predictions.dims(); + + let metric = self.compute_auc(&input.predictions, &input.targets); + + self.state.update( + 100.0 * metric, + sample_size, + FormatOptions::new(self.name()).unit("%").precision(2), + ) + } + + fn clear(&mut self) { + self.state.reset() + } + + fn name(&self) -> MetricName { + self.name.clone() + } +} + +impl Numeric for AurocMetric { + fn value(&self) -> super::NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> super::NumericEntry { + self.state.running_value() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::metric::ClassReduction::{self, *}; + use burn_core::tensor::{TensorData, Tolerance}; + use rstest::rstest; + + /// Inputs and expected AUROC computed with an independent reference + /// equivalent to scikit-learn's `roc_auc_score` (Mann-Whitney U: + /// `(#pos>neg + 0.5·ties) / (P·N)`, One-vs-Rest, macro/micro). Scores + /// are distinct so the statistic is unambiguous and matches sklearn. + #[derive(Clone, Copy)] + enum Data { + Binary, + Multiclass, + Multilabel, + } + + fn input(data: Data) -> ConfusionStatsInput { + let dev = Default::default(); + match data { + Data::Binary => ConfusionStatsInput::new( + Tensor::from_data([[0.34], [0.64], [0.12], [0.19], [0.53], [0.38]], &dev), + Tensor::from_data([[0], [0], [0], [0], [1], [1]], &dev), + ), + Data::Multiclass => ConfusionStatsInput::new( + Tensor::from_data( + [ + [0.79, 0.41, 0.16], + [0.25, 0.93, 0.78], + [0.61, 0.09, 0.21], + [0.9, 0.31, 0.33], + [0.16, 0.82, 0.57], + [0.57, 0.18, 0.63], + ], + &dev, + ), + Tensor::from_data( + [ + [1, 0, 0], + [1, 0, 0], + [1, 0, 0], + [0, 0, 1], + [0, 1, 0], + [1, 0, 0], + ], + &dev, + ), + ), + Data::Multilabel => ConfusionStatsInput::new( + Tensor::from_data( + [ + [0.11, 0.57, 0.9], + [0.13, 0.66, 0.37], + [0.71, 0.85, 0.6], + [0.29, 0.69, 0.49], + [0.68, 0.45, 0.25], + [0.33, 0.36, 0.31], + ], + &dev, + ), + Tensor::from_data( + [ + [1, 1, 1], + [0, 0, 1], + [0, 1, 0], + [1, 1, 0], + [0, 1, 1], + [1, 1, 1], + ], + &dev, + ), + ), + } + } + + #[rstest] + // Binary is a single column -> Macro == Micro. + #[case::binary_macro(Data::Binary, Macro, 0.75)] + #[case::binary_micro(Data::Binary, Micro, 0.75)] + #[case::multiclass_macro(Data::Multiclass, Macro, 0.5666666666666667)] + #[case::multiclass_micro(Data::Multiclass, Micro, 0.6458333333333333)] + #[case::multilabel_macro(Data::Multilabel, Macro, 0.2907407407407407)] + #[case::multilabel_micro(Data::Multilabel, Micro, 0.3611111111111111)] + fn test_auroc( + #[case] data: Data, + #[case] class_reduction: ClassReduction, + #[case] expected: f64, + ) { + let mut metric = AurocMetric::new(class_reduction); + + let _entry = metric.update(&input(data), &MetricMetadata::fake()); + + TensorData::from([metric.value().current()]) + .assert_approx_eq::(&TensorData::from([expected * 100.0]), Tolerance::default()); + } + + #[rstest] + #[case::macro_reduction(Macro)] + #[case::micro_reduction(Micro)] + fn test_auroc_perfect_separation(#[case] class_reduction: ClassReduction) { + let device = Default::default(); + let mut metric = AurocMetric::new(class_reduction); + + let input = ConfusionStatsInput::new( + Tensor::from_data([[0.0, 1.0], [1.0, 0.0], [1.0, 0.0], [0.0, 1.0]], &device), + Tensor::from_data([[0, 1], [1, 0], [1, 0], [0, 1]], &device), + ); + + let _entry = metric.update(&input, &MetricMetadata::fake()); + assert_eq!(metric.value().current(), 100.0); + } + + #[rstest] + #[case::macro_reduction(Macro)] + #[case::micro_reduction(Micro)] + fn test_auroc_chance_level(#[case] class_reduction: ClassReduction) { + let device = Default::default(); + let mut metric = AurocMetric::new(class_reduction); + + // All scores tied -> every pair is a tie -> AUROC = 0.5. + let input = ConfusionStatsInput::new( + Tensor::from_data([[0.5, 0.5], [0.5, 0.5], [0.5, 0.5], [0.5, 0.5]], &device), + Tensor::from_data([[0, 1], [1, 0], [1, 0], [0, 1]], &device), + ); + + let _entry = metric.update(&input, &MetricMetadata::fake()); + assert_eq!(metric.value().current(), 50.0); + } + + #[test] + fn test_auroc_macro_drops_degenerate_class() { + let device = Default::default(); + let mut metric = AurocMetric::new(Macro); + + // Class 2 never appears (column all-negative) -> its AUROC is undefined + // and must be dropped, leaving the two well-separated classes at 1.0. + let input = ConfusionStatsInput::new( + Tensor::from_data( + [ + [0.9, 0.1, 0.0], + [0.2, 0.8, 0.0], + [0.7, 0.3, 0.0], + [0.1, 0.6, 0.0], + ], + &device, + ), + Tensor::from_data([[1, 0, 0], [0, 1, 0], [1, 0, 0], [0, 1, 0]], &device), + ); + + let _entry = metric.update(&input, &MetricMetadata::fake()); + assert_eq!(metric.value().current(), 100.0); + } + + #[test] + fn test_auroc_all_degenerate_is_chance() { + let device = Default::default(); + let mut metric = AurocMetric::binary(); + + // Only positives -> no valid pair in any column -> undefined -> + // reported as chance level (0.5). + let input = ConfusionStatsInput::new( + Tensor::from_data([[0.9], [0.8], [0.7], [0.6]], &device), + Tensor::from_data([[1], [1], [1], [1]], &device), + ); + + let _entry = metric.update(&input, &MetricMetadata::fake()); + assert_eq!(metric.value().current(), 50.0); + } + + #[test] + fn test_auroc_reduction_changes_name() { + let macro_metric = AurocMetric::new(Macro); + let micro_metric = AurocMetric::new(Micro); + + assert_ne!(macro_metric.name(), micro_metric.name()); + } +} diff --git a/crates/burn-train/src/metric/base.rs b/crates/burn-train/src/metric/base.rs new file mode 100644 index 0000000..064399c --- /dev/null +++ b/crates/burn-train/src/metric/base.rs @@ -0,0 +1,285 @@ +use std::sync::Arc; + +use burn_core::data::dataloader::Progress; +use burn_optim::lr_scheduler::module_lr_scheduler::ModuleLearningRate; + +/// Metric metadata that can be used when computing metrics. +pub struct MetricMetadata { + /// The current progress. + pub progress: Progress, + + /// The current iteration. + pub iteration: Option, + + /// The current learning rate. + pub lr: Option, +} + +impl MetricMetadata { + /// Fake metric metadata + #[cfg(test)] + pub fn fake() -> Self { + Self { + progress: Progress { + items_processed: 1, + items_total: 1, + unit: Some("items".to_string()), + }, + iteration: Some(0), + lr: None, + } + } +} + +/// Metric id that can be used to compare metrics and retrieve entries of the same metric. +/// For now we take the name as id to make sure that the same metric has the same id across different runs. +#[derive(Debug, Clone, new, PartialEq, Eq, Hash)] +pub struct MetricId { + /// The metric id. + id: Arc, +} + +/// Metric attributes define the properties intrinsic to different types of metric. +#[derive(Clone, Debug)] +pub enum MetricAttributes { + /// Numeric attributes. + Numeric(NumericAttributes), + /// No attributes. + None, +} + +/// Definition of a metric. +#[derive(Clone, Debug)] +pub struct MetricDefinition { + /// The metric's id. + pub metric_id: MetricId, + /// The name of the metric. + pub name: String, + /// The description of the metric. + pub description: Option, + /// The attributes of the metric. + pub attributes: MetricAttributes, +} + +impl MetricDefinition { + /// Create a new metric definition given the metric and a unique id. + pub fn new(metric_id: MetricId, metric: &Me) -> Self { + Self { + metric_id, + name: metric.name().to_string(), + description: metric.description(), + attributes: metric.attributes(), + } + } +} + +/// Metric trait. +/// +/// # Notes +/// +/// Implementations should define their own input type only used by the metric. +/// This is important since some conflict may happen when the model output is adapted for each +/// metric's input type. +pub trait Metric: Send + Sync + Clone { + /// The input type of the metric. + type Input; + + /// The parameterized name of the metric. + /// + /// This should be unique, so avoid using short generic names, prefer using the long name. + /// + /// For a metric that can exist at different parameters (e.g., top-k accuracy for different + /// values of k), the name should be unique for each instance. + fn name(&self) -> MetricName; + + /// A short description of the metric. + fn description(&self) -> Option { + None + } + + /// Attributes of the metric. + /// + /// By default, metrics have no attributes. + fn attributes(&self) -> MetricAttributes { + MetricAttributes::None + } + + /// Update the metric state and returns the current metric entry. + fn update(&mut self, item: &Self::Input, metadata: &MetricMetadata) -> SerializedEntry; + + /// Clear the metric state. + fn clear(&mut self); +} + +/// Type used to store metric names efficiently. +pub type MetricName = Arc; + +/// Adaptor are used to transform types so that they can be used by metrics. +/// +/// This should be implemented by a model's output type for all [metric inputs](Metric::Input) that are +/// registered with the specific learning paradigm (i.e. [SupervisedTraining](crate::SupervisedTraining)). +pub trait Adaptor { + /// Adapt the type to be passed to a [metric](Metric). + fn adapt(&self) -> T; +} + +impl Adaptor<()> for T { + fn adapt(&self) {} +} + +/// How a numeric metric's per-batch values are reduced into an epoch value. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum NumericAggregation { + /// Sample-weighted mean of the per-batch values. + #[default] + Mean, + /// Last logged value, already computed over the whole epoch. + Last, +} + +/// Attributes that describe intrinsic properties of a numeric metric. +#[derive(Clone, Debug)] +pub struct NumericAttributes { + /// Optional unit (e.g. "%", "ms", "pixels") + pub unit: Option, + /// Whether larger values are better (true) or smaller are better (false). + pub higher_is_better: bool, + /// How per-batch values are reduced into the epoch value. + pub aggregation: NumericAggregation, +} + +impl From for MetricAttributes { + fn from(attr: NumericAttributes) -> Self { + MetricAttributes::Numeric(attr) + } +} + +impl Default for NumericAttributes { + fn default() -> Self { + Self { + unit: None, + higher_is_better: true, + aggregation: NumericAggregation::default(), + } + } +} + +/// Declare a metric to be numeric. +/// +/// This is useful to plot the values of a metric during training. +pub trait Numeric { + /// Returns the numeric value of the metric. + fn value(&self) -> NumericEntry; + /// Returns the current aggregated value of the metric over the global step (epoch). + fn running_value(&self) -> NumericEntry; +} + +/// Serialized form of a metric entry. +#[derive(Debug, Clone, new)] +pub struct SerializedEntry { + /// The string to be displayed. + pub formatted: String, + /// The string to be saved. + pub serialized: String, +} + +/// Data type that contains the current state of a metric at a given time. +#[derive(Debug, Clone)] +pub struct MetricEntry { + /// Id of the entry's metric. + pub metric_id: MetricId, + /// The serialized form of the entry. + pub serialized_entry: SerializedEntry, +} + +impl MetricEntry { + /// Create a new metric. + pub fn new(metric_id: MetricId, serialized_entry: SerializedEntry) -> Self { + Self { + metric_id, + serialized_entry, + } + } +} + +/// Numeric metric entry. +#[derive(Debug, Clone)] +pub enum NumericEntry { + /// Single numeric value. + Value(f64), + /// Aggregated numeric (value, number of elements). + Aggregated { + /// The aggregated value of all entries. + aggregated_value: f64, + /// The number of entries present in the aggregated value. + count: usize, + }, +} + +impl NumericEntry { + /// Gets the current aggregated value of the metric. + pub fn current(&self) -> f64 { + match self { + NumericEntry::Value(val) => *val, + NumericEntry::Aggregated { + aggregated_value, .. + } => *aggregated_value, + } + } + + /// Returns a String representing the NumericEntry + pub fn serialize(&self) -> String { + match self { + Self::Value(v) => v.to_string(), + Self::Aggregated { + aggregated_value, + count, + } => format!("{aggregated_value},{count}"), + } + } + + /// De-serializes a string representing a NumericEntry and returns a Result containing the corresponding NumericEntry. + pub fn deserialize(entry: &str) -> Result { + // Check for comma separated values + let values = entry.split(',').collect::>(); + let num_values = values.len(); + + if num_values == 1 { + // Numeric value + match values[0].parse::() { + Ok(value) => Ok(NumericEntry::Value(value)), + Err(err) => Err(err.to_string()), + } + } else if num_values == 2 { + // Aggregated numeric (value, number of elements) + let (value, numel) = (values[0], values[1]); + match value.parse::() { + Ok(value) => match numel.parse::() { + Ok(numel) => Ok(NumericEntry::Aggregated { + aggregated_value: value, + count: numel, + }), + Err(err) => Err(err.to_string()), + }, + Err(err) => Err(err.to_string()), + } + } else { + Err("Invalid number of values for numeric entry".to_string()) + } + } + + /// Compare this numeric metric's value with another one using the specified direction. + pub fn better_than(&self, other: &NumericEntry, higher_is_better: bool) -> bool { + (self.current() > other.current()) == higher_is_better + } +} + +/// Format a float with the given precision. Will use scientific notation if necessary. +pub fn format_float(float: f64, precision: usize) -> String { + let scientific_notation_threshold = 0.1_f64.powf(precision as f64 - 1.0); + + match scientific_notation_threshold >= float { + true => format!("{float:.precision$e}"), + false => format!("{float:.precision$}"), + } +} diff --git a/crates/burn-train/src/metric/bleu.rs b/crates/burn-train/src/metric/bleu.rs new file mode 100644 index 0000000..f919a50 --- /dev/null +++ b/crates/burn-train/src/metric/bleu.rs @@ -0,0 +1,538 @@ +use super::state::{FormatOptions, NumericMetricState}; +use super::{MetricMetadata, SerializedEntry}; +use crate::metric::{ + Metric, MetricAttributes, MetricName, Numeric, NumericAttributes, NumericEntry, +}; +use burn_core::tensor::{Int, Tensor}; +use std::collections::HashMap; +use std::sync::Arc; + +/// Smoothing method for BLEU score computation. +/// +/// Sentence-level BLEU often produces zero scores when higher-order n-grams +/// have no matches. Smoothing techniques address this by assigning small +/// non-zero values to zero-count precisions. +/// +/// # References +/// +/// Chen & Cherry, "A Systematic Comparison of Smoothing Techniques for +/// Sentence-Level BLEU", WMT 2014. +#[derive(Clone, Debug, Default)] +pub enum BleuSmoothing { + /// No smoothing. Zero precision at any n-gram order produces a zero + /// overall score (standard corpus-level BLEU behaviour). + #[default] + None, + + /// Add a small constant (`epsilon`, default 0.1) to zero-count + /// precisions. Corresponds to method 1 in Chen & Cherry (2014). + AddEpsilon(f64), + + /// Exponential decay: for each n-gram order with zero matches, double a + /// running multiplier `k` (starting at 1 and doubling on every zero) and + /// replace the precision with `1 / (k * total_n)`. Corresponds to + /// method 3 in Chen & Cherry (2014) and the default smoothing in + /// SacreBLEU for sentence-level BLEU. + Exponential, +} + +/// Computes the BLEU (Bilingual Evaluation Understudy) score between predicted +/// and reference token sequences. +/// +/// BLEU measures the quality of machine-translated text by comparing n-gram +/// overlap between the candidate (prediction) and reference sequences. The +/// score combines modified n-gram precision for n = 1..max_n with a brevity +/// penalty that discourages overly short translations. +/// +/// The metric operates on integer token IDs (not raw text), matching the +/// convention used by [`CharErrorRate`](super::CharErrorRate) and +/// [`WordErrorRate`](super::WordErrorRate). +/// +/// # Batch-level scoring +/// +/// Within each batch the metric accumulates n-gram counts across all +/// sentences and computes a single corpus-style BLEU score, following the +/// same pattern CER/WER use for edit-distance aggregation. +/// +/// Epoch-level (running) aggregation averages these batch scores, which is +/// slightly inaccurate compared to true corpus BLEU. Correct corpus-level +/// accumulation would require a custom metric state; a TODO is left for +/// future work. +/// +/// # References +/// +/// Papineni et al., "BLEU: a Method for Automatic Evaluation of Machine +/// Translation", ACL 2002. +/// +/// Chen & Cherry, "A Systematic Comparison of Smoothing Techniques for +/// Sentence-Level BLEU", WMT 2014. +#[derive(Clone)] +pub struct BleuScore { + name: MetricName, + state: NumericMetricState, + max_n: usize, + pad_token: Option, + smoothing: BleuSmoothing, +} + +/// The [BLEU score metric](BleuScore) input type. +#[derive(new)] +pub struct BleuInput { + /// The predicted token sequences (2-D tensor of token indices). + pub outputs: Tensor<2, Int>, + /// The reference token sequences (2-D tensor of token indices). + pub targets: Tensor<2, Int>, +} + +impl Default for BleuScore { + fn default() -> Self { + Self::with_max_n(4) + } +} + +impl BleuScore { + /// Creates a BLEU metric with the given maximum n-gram order. + /// + /// Common values: 1 (BLEU-1), 2 (BLEU-2), 4 (BLEU-4). + /// + /// # Panics + /// + /// Panics if `max_n` is zero. + pub fn with_max_n(max_n: usize) -> Self { + assert!(max_n >= 1, "max_n must be at least 1"); + Self { + name: Arc::new(format!("BLEU-{max_n}")), + state: NumericMetricState::default(), + max_n, + pad_token: None, + smoothing: BleuSmoothing::default(), + } + } + + /// Creates a BLEU-4 metric (the most common configuration). + pub fn new() -> Self { + Self::default() + } + + /// Sets the pad token index. Tokens matching this value are stripped from + /// the right of each sequence before scoring. + pub fn with_pad_token(mut self, index: usize) -> Self { + self.pad_token = Some(index); + self + } + + /// Sets the smoothing method for handling zero-count n-gram precisions. + /// + /// Smoothing is recommended when evaluating short sentences where + /// higher-order n-gram matches are sparse. + pub fn with_smoothing(mut self, smoothing: BleuSmoothing) -> Self { + self.smoothing = smoothing; + self + } +} + +/// Extracts n-grams of order `n` from a slice and returns their counts. +fn ngram_counts(tokens: &[i32], n: usize) -> HashMap, usize> { + let mut counts = HashMap::new(); + if tokens.len() >= n { + for window in tokens.windows(n) { + *counts.entry(window.to_vec()).or_insert(0) += 1; + } + } + counts +} + +/// Computes corpus-style BLEU score from accumulated n-gram statistics. +/// +/// `clipped_counts[n]` and `total_counts[n]` hold the clipped and total +/// n-gram counts for order `n+1` across all sentences. +/// `candidate_len` and `reference_len` are the total token counts. +/// +/// Returns a value in [0, 100]. +fn corpus_bleu( + clipped_counts: &[usize], + total_counts: &[usize], + candidate_len: usize, + reference_len: usize, + max_n: usize, + smoothing: &BleuSmoothing, +) -> f64 { + if candidate_len == 0 { + return 0.0; + } + + // Brevity penalty + let bp = if candidate_len < reference_len { + (1.0 - reference_len as f64 / candidate_len as f64).exp() + } else { + 1.0 + }; + + // Modified n-gram precisions (log-space) + let mut log_avg = 0.0; + let mut counted_orders = 0; + // Stateful multiplier for exponential smoothing (Chen & Cherry 2014, method 3; + // also used by SacreBLEU). Doubles on every n-gram order with zero matches. + let mut smooth_mult = 1.0_f64; + + for n in 0..max_n { + let total = total_counts[n]; + let clipped = clipped_counts[n]; + + if total == 0 { + // Candidate has no n-grams of this order (too short). + // Standard BLEU: this order contributes 0 to the geometric mean, + // which collapses the entire score to 0. + return 0.0; + } + + let precision = if clipped == 0 { + // Apply smoothing to zero-match precisions. + match smoothing { + BleuSmoothing::None => return 0.0, + BleuSmoothing::AddEpsilon(eps) => *eps / total as f64, + BleuSmoothing::Exponential => { + smooth_mult *= 2.0; + 1.0 / (smooth_mult * total as f64) + } + } + } else { + clipped as f64 / total as f64 + }; + + log_avg += precision.ln(); + counted_orders += 1; + } + + if counted_orders == 0 { + return 0.0; + } + + let score = bp * (log_avg / counted_orders as f64).exp(); + score * 100.0 +} + +impl Metric for BleuScore { + type Input = BleuInput; + + fn update(&mut self, input: &BleuInput, _metadata: &MetricMetadata) -> SerializedEntry { + let outputs = &input.outputs; + let targets = &input.targets; + let [batch_size, seq_len] = targets.dims(); + + let outputs_data = outputs.to_data().iter::().collect::>(); + let targets_data = targets.to_data().iter::().collect::>(); + + let pad_token = self.pad_token.map(|p| p as i32); + + // Accumulate n-gram counts across the batch (corpus-style). + let mut clipped_counts = vec![0usize; self.max_n]; + let mut total_counts = vec![0usize; self.max_n]; + let mut total_candidate_len = 0usize; + let mut total_reference_len = 0usize; + + for i in 0..batch_size { + let start = i * seq_len; + let end = (i + 1) * seq_len; + + let output_seq = &outputs_data[start..end]; + let target_seq = &targets_data[start..end]; + + // Strip right-padding if configured. + let output_seq = match pad_token { + Some(pad) => { + let len = output_seq + .iter() + .position(|&x| x == pad) + .unwrap_or(output_seq.len()); + &output_seq[..len] + } + None => output_seq, + }; + let target_seq = match pad_token { + Some(pad) => { + let len = target_seq + .iter() + .position(|&x| x == pad) + .unwrap_or(target_seq.len()); + &target_seq[..len] + } + None => target_seq, + }; + + total_candidate_len += output_seq.len(); + total_reference_len += target_seq.len(); + + for n in 1..=self.max_n { + let cand_ngrams = ngram_counts(output_seq, n); + let ref_ngrams = ngram_counts(target_seq, n); + + for (ngram, &count) in &cand_ngrams { + let ref_count = ref_ngrams.get(ngram).copied().unwrap_or(0); + clipped_counts[n - 1] += count.min(ref_count); + total_counts[n - 1] += count; + } + } + } + + let value = corpus_bleu( + &clipped_counts, + &total_counts, + total_candidate_len, + total_reference_len, + self.max_n, + &self.smoothing, + ); + + // TODO: Epoch-level aggregation averages batch BLEU scores, which is + // slightly inaccurate compared to true corpus BLEU. Correct + // accumulation would require a custom metric state that tracks raw + // n-gram counts across batches. + self.state.update( + value, + batch_size, + FormatOptions::new(self.name()).unit("%").precision(2), + ) + } + + fn clear(&mut self) { + self.state.reset(); + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: Some("%".to_string()), + higher_is_better: true, + ..Default::default() + } + .into() + } +} + +impl Numeric for BleuScore { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Perfect match => BLEU = 100 %. + #[test] + fn test_bleu_perfect_match() { + let device = Default::default(); + let mut metric = BleuScore::new(); + + let preds = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + let tgts = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + + metric.update(&BleuInput::new(preds, tgts), &MetricMetadata::fake()); + + assert!((metric.value().current() - 100.0).abs() < 1e-6); + } + + /// Completely different sequences => BLEU = 0 %. + #[test] + fn test_bleu_no_match() { + let device = Default::default(); + let mut metric = BleuScore::new(); + + let preds = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + let tgts = Tensor::from_data([[6, 7, 8, 9, 10]], &device); + + metric.update(&BleuInput::new(preds, tgts), &MetricMetadata::fake()); + + assert_eq!(0.0, metric.value().current()); + } + + /// Partial overlap with BLEU-1 (unigram precision only). + #[test] + fn test_bleu1_partial_match() { + let device = Default::default(); + let mut metric = BleuScore::with_max_n(1); + + // 3 out of 5 unigrams match, same length so BP = 1 + let preds = Tensor::from_data([[1, 2, 3, 6, 7]], &device); + let tgts = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + + metric.update(&BleuInput::new(preds, tgts), &MetricMetadata::fake()); + + // BLEU-1 = 3/5 * 100 = 60.0 + assert!((metric.value().current() - 60.0).abs() < 1e-6); + } + + /// Brevity penalty applied when candidate is shorter than reference. + #[test] + fn test_bleu_brevity_penalty() { + let device = Default::default(); + let pad = 0_i64; + let mut metric = BleuScore::with_max_n(1).with_pad_token(pad as usize); + + // Candidate has 3 tokens, reference has 5 tokens. + // Unigram precision = 3/3 = 1.0, BP = exp(1 - 5/3) ~= 0.5134 + let preds = Tensor::from_data([[1, 2, 3, pad, pad]], &device); + let tgts = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + + metric.update(&BleuInput::new(preds, tgts), &MetricMetadata::fake()); + + let expected = 100.0 * (1.0 - 5.0 / 3.0_f64).exp(); + assert!((metric.value().current() - expected).abs() < 0.1); + } + + /// With padding, padding tokens should be stripped. + #[test] + fn test_bleu_with_padding() { + let device = Default::default(); + let pad = 99_i64; + let mut metric = BleuScore::new().with_pad_token(pad as usize); + + // Same non-pad content => should be 100% + let preds = Tensor::from_data([[1, 2, 3, 4, 5, pad, pad]], &device); + let tgts = Tensor::from_data([[1, 2, 3, 4, 5, pad, pad]], &device); + + metric.update(&BleuInput::new(preds, tgts), &MetricMetadata::fake()); + + assert!((metric.value().current() - 100.0).abs() < 1e-6); + } + + /// Batch of two sequences: corpus-style accumulation. + #[test] + fn test_bleu_batch_corpus_style() { + let device = Default::default(); + let mut metric = BleuScore::with_max_n(1); + + // Sequence 1: perfect match [1,2,3,4,5] + // Sequence 2: no match [6,7,8,9,10] vs [11,12,13,14,15] + // Corpus unigram: clipped = 5, total = 10, precision = 0.5 + // BP: candidate_len = 10, ref_len = 10, BP = 1.0 + // BLEU-1 = 50.0 + let preds = Tensor::from_data([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], &device); + let tgts = Tensor::from_data([[1, 2, 3, 4, 5], [11, 12, 13, 14, 15]], &device); + + metric.update(&BleuInput::new(preds, tgts), &MetricMetadata::fake()); + + assert!((metric.value().current() - 50.0).abs() < 1e-6); + } + + /// `clear()` must reset the running statistics. + #[test] + fn test_clear_resets_state() { + let device = Default::default(); + let mut metric = BleuScore::new(); + + let preds = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + let tgts = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + + metric.update(&BleuInput::new(preds, tgts), &MetricMetadata::fake()); + assert!(metric.value().current() > 0.0); + + metric.clear(); + assert!(metric.value().current().is_nan()); + } + + /// BLEU-2 on a partial match verifies bigram counting. + #[test] + fn test_bleu2_bigrams() { + let device = Default::default(); + let mut metric = BleuScore::with_max_n(2); + + // Candidate: [1, 2, 3, 4] Reference: [1, 2, 5, 6] + // Unigrams: candidate {1,2,3,4}, reference {1,2,5,6} + // matches: 1,2 => clipped 2/4 + // Bigrams: candidate {(1,2),(2,3),(3,4)}, reference {(1,2),(2,5),(5,6)} + // matches: (1,2) => clipped 1/3 + // BP = 1.0 (same length) + // BLEU-2 = exp((ln(2/4) + ln(1/3)) / 2) * 100 + let preds = Tensor::from_data([[1, 2, 3, 4]], &device); + let tgts = Tensor::from_data([[1, 2, 5, 6]], &device); + + metric.update(&BleuInput::new(preds, tgts), &MetricMetadata::fake()); + + let expected = 100.0 * ((0.5_f64.ln() + (1.0 / 3.0_f64).ln()) / 2.0).exp(); + assert!((metric.value().current() - expected).abs() < 0.1); + } + + /// BLEU with custom name reflects max_n. + #[test] + fn test_bleu_custom_name() { + let metric = BleuScore::with_max_n(2); + assert_eq!(*metric.name(), "BLEU-2"); + } + + /// Default name is BLEU-4. + #[test] + fn test_bleu_default_name() { + let metric = BleuScore::new(); + assert_eq!(*metric.name(), "BLEU-4"); + } + + /// Short candidate with BLEU-4 returns 0 without smoothing + /// (too few tokens for 4-grams). + #[test] + fn test_bleu_short_candidate_no_smoothing() { + let device = Default::default(); + let pad = 0_i64; + let mut metric = BleuScore::new().with_pad_token(pad as usize); + + // Only 3 tokens — no 4-grams exist, score must be 0. + let preds = Tensor::from_data([[1, 2, 3, pad, pad]], &device); + let tgts = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + + metric.update(&BleuInput::new(preds, tgts), &MetricMetadata::fake()); + assert_eq!(0.0, metric.value().current()); + } + + /// Exponential smoothing produces a non-zero score for short candidates. + #[test] + fn test_bleu_exponential_smoothing() { + let device = Default::default(); + let mut metric = BleuScore::with_max_n(2).with_smoothing(BleuSmoothing::Exponential); + + // Unigrams: {1,3,5,7,9} vs {1,2,3,4,5} — clipped = 2/5 + // Bigrams: {(1,3),(3,5),(5,7),(7,9)} vs {(1,2),(2,3),(3,4),(4,5)} — clipped = 0/4 + let preds = Tensor::from_data([[1, 3, 5, 7, 9]], &device); + let tgts = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + + // Verify without smoothing this is 0. + let mut metric_no_smooth = BleuScore::with_max_n(2); + metric_no_smooth.update( + &BleuInput::new(preds.clone(), tgts.clone()), + &MetricMetadata::fake(), + ); + assert_eq!(0.0, metric_no_smooth.value().current()); + + metric.update(&BleuInput::new(preds, tgts), &MetricMetadata::fake()); + assert!( + metric.value().current() > 0.0, + "smoothing should produce non-zero score" + ); + } + + /// Add-epsilon smoothing produces a non-zero score for zero-precision orders. + #[test] + fn test_bleu_add_epsilon_smoothing() { + let device = Default::default(); + let mut metric = BleuScore::with_max_n(2).with_smoothing(BleuSmoothing::AddEpsilon(0.1)); + + // Unigrams: {1,3,5,7,9} vs {1,2,3,4,5} — clipped 2, total 5 + // Bigrams: {(1,3),(3,5),(5,7),(7,9)} vs {(1,2),(2,3),(3,4),(4,5)} — clipped 0, total 4 + let preds = Tensor::from_data([[1, 3, 5, 7, 9]], &device); + let tgts = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + + metric.update(&BleuInput::new(preds, tgts), &MetricMetadata::fake()); + assert!( + metric.value().current() > 0.0, + "epsilon smoothing should produce non-zero score" + ); + } +} diff --git a/crates/burn-train/src/metric/cer.rs b/crates/burn-train/src/metric/cer.rs new file mode 100644 index 0000000..79a0bb5 --- /dev/null +++ b/crates/burn-train/src/metric/cer.rs @@ -0,0 +1,241 @@ +use super::state::{FormatOptions, NumericMetricState}; +use super::{MetricMetadata, SerializedEntry}; +use crate::metric::{Metric, MetricAttributes, MetricName, Numeric, NumericEntry}; +use burn_core::tensor::{Int, Tensor}; +use std::sync::Arc; + +/// Computes the edit distance (Levenshtein distance) between two sequences of integers. +/// +/// The edit distance is defined as the minimum number of single-element edits (insertions, +/// deletions, or substitutions) required to change one sequence into the other. This +/// implementation is optimized for space, using only two rows of the dynamic programming table. +/// +pub(crate) fn edit_distance(reference: &[i32], prediction: &[i32]) -> usize { + let mut prev = (0..=prediction.len()).collect::>(); + let mut curr = vec![0; prediction.len() + 1]; + + for (i, &r) in reference.iter().enumerate() { + curr[0] = i + 1; + for (j, &p) in prediction.iter().enumerate() { + curr[j + 1] = if r == p { + prev[j] // no operation needed + } else { + 1 + prev[j].min(prev[j + 1]).min(curr[j]) // substitution, insertion, deletion + }; + } + core::mem::swap(&mut prev, &mut curr); + } + prev[prediction.len()] +} + +/// Character error rate (CER) is defined as the edit distance (e.g. Levenshtein distance) between the predicted +/// and reference character sequences, divided by the total number of characters in the reference. +/// This metric is commonly used in tasks such as speech recognition, OCR, or text generation +/// to quantify how closely the predicted output matches the ground truth at a character level. +/// +#[derive(Clone)] +pub struct CharErrorRate { + name: MetricName, + state: NumericMetricState, + pad_token: Option, +} + +/// The [character error rate metric](CharErrorRate) input type. +#[derive(new)] +pub struct CerInput { + /// The predicted token sequences (as a 2-D tensor of token indices). + pub outputs: Tensor<2, Int>, + /// The target token sequences (as a 2-D tensor of token indices). + pub targets: Tensor<2, Int>, +} + +impl Default for CharErrorRate { + fn default() -> Self { + Self::new() + } +} + +impl CharErrorRate { + /// Creates the metric. + pub fn new() -> Self { + Self { + name: Arc::new("CER".to_string()), + state: NumericMetricState::default(), + pad_token: None, + } + } + + /// Sets the pad token. + pub fn with_pad_token(mut self, index: usize) -> Self { + self.pad_token = Some(index); + self + } +} + +/// The [character error rate metric](CharErrorRate) implementation. +impl Metric for CharErrorRate { + type Input = CerInput; + + fn update(&mut self, input: &CerInput, _metadata: &MetricMetadata) -> SerializedEntry { + let outputs = &input.outputs; + let targets = &input.targets; + let [batch_size, seq_len] = targets.dims(); + + let (output_lengths, target_lengths) = if let Some(pad) = self.pad_token { + // Create boolean masks for non-padding tokens. + let output_mask = outputs.clone().not_equal_elem(pad as i64); + let target_mask = targets.clone().not_equal_elem(pad as i64); + + let output_lengths_tensor = output_mask.int().sum_dim(1); + let target_lengths_tensor = target_mask.int().sum_dim(1); + + ( + output_lengths_tensor + .into_data() + .convert::() + .to_vec() + .unwrap(), + target_lengths_tensor + .into_data() + .convert::() + .to_vec() + .unwrap(), + ) + } else { + // If there's no padding, all sequences have the full length. + ( + vec![seq_len as i32; batch_size], + vec![seq_len as i32; batch_size], + ) + }; + + let outputs_data = outputs.to_data().convert::().to_vec().unwrap(); + let targets_data = targets.to_data().convert::().to_vec().unwrap(); + + let total_edit_distance: usize = (0..batch_size) + .map(|i| { + let start = i * seq_len; + + // Get pre-calculated lengths for the current sequence. + let output_len = output_lengths[i] as usize; + let target_len = target_lengths[i] as usize; + + let output_seq_slice = &outputs_data[start..(start + output_len)]; + let target_seq_slice = &targets_data[start..(start + target_len)]; + + edit_distance(target_seq_slice, output_seq_slice) + }) + .sum(); + + let total_target_length = target_lengths.iter().map(|&x| x as f64).sum::(); + + let value = if total_target_length > 0.0 { + 100.0 * total_edit_distance as f64 / total_target_length + } else { + 0.0 + }; + + self.state.update( + value, + batch_size, + FormatOptions::new(self.name()).unit("%").precision(2), + ) + } + + fn clear(&mut self) { + self.state.reset(); + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + super::NumericAttributes { + unit: Some("%".to_string()), + higher_is_better: false, + ..Default::default() + } + .into() + } +} + +impl Numeric for CharErrorRate { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Perfect match ⇒ CER = 0 %. + #[test] + fn test_cer_without_padding() { + let device = Default::default(); + let mut metric = CharErrorRate::new(); + + // Batch size = 2, sequence length = 2 + let preds = Tensor::from_data([[1, 2], [3, 4]], &device); + let tgts = Tensor::from_data([[1, 2], [3, 4]], &device); + + metric.update(&CerInput::new(preds, tgts), &MetricMetadata::fake()); + + assert_eq!(0.0, metric.value().current()); + } + + /// Two edits in four target tokens ⇒ 50 %. + #[test] + fn test_cer_without_padding_two_errors() { + let device = Default::default(); + let mut metric = CharErrorRate::new(); + + // One substitution in each sequence. + let preds = Tensor::from_data([[1, 2], [3, 5]], &device); + let tgts = Tensor::from_data([[1, 3], [3, 4]], &device); + + metric.update(&CerInput::new(preds, tgts), &MetricMetadata::fake()); + + // 2 edits / 4 tokens = 50 % + assert_eq!(50.0, metric.value().current()); + } + + /// Same scenario as above, but with right-padding (token 9) ignored. + #[test] + fn test_cer_with_padding() { + let device = Default::default(); + let pad = 9_i64; + let mut metric = CharErrorRate::new().with_pad_token(pad as usize); + + // Each row has three columns, last one is the pad token. + let preds = Tensor::from_data([[1, 2, pad], [3, 5, pad]], &device); + let tgts = Tensor::from_data([[1, 3, pad], [3, 4, pad]], &device); + + metric.update(&CerInput::new(preds, tgts), &MetricMetadata::fake()); + assert_eq!(50.0, metric.value().current()); + } + + /// `clear()` must reset the running statistics to zero. + #[test] + fn test_clear_resets_state() { + let device = Default::default(); + let mut metric = CharErrorRate::new(); + + let preds = Tensor::from_data([[1, 2]], &device); + let tgts = Tensor::from_data([[1, 3]], &device); // one error + + metric.update( + &CerInput::new(preds.clone(), tgts.clone()), + &MetricMetadata::fake(), + ); + assert!(metric.value().current() > 0.0); + + metric.clear(); + assert!(metric.value().current().is_nan()); + } +} diff --git a/crates/burn-train/src/metric/classification.rs b/crates/burn-train/src/metric/classification.rs new file mode 100644 index 0000000..6d66c63 --- /dev/null +++ b/crates/burn-train/src/metric/classification.rs @@ -0,0 +1,33 @@ +use std::num::NonZeroUsize; + +/// Necessary data for classification metrics. +#[derive(Default, Debug, Clone)] +pub struct ClassificationMetricConfig { + pub decision_rule: DecisionRule, + pub class_reduction: ClassReduction, +} + +/// The prediction decision rule for classification metrics. +#[derive(Debug, Clone)] +pub enum DecisionRule { + /// Consider a class predicted if its probability exceeds the threshold. + Threshold(f64), + /// Consider a class predicted correctly if it is within the top k predicted classes based on scores. + TopK(NonZeroUsize), +} + +impl Default for DecisionRule { + fn default() -> Self { + Self::Threshold(0.5) + } +} + +/// The reduction strategy for classification metrics. +#[derive(Copy, Clone, Default, Debug)] +pub enum ClassReduction { + /// Computes the statistics over all classes before averaging + Micro, + /// Computes the statistics independently for each class before averaging + #[default] + Macro, +} diff --git a/crates/burn-train/src/metric/confusion_stats.rs b/crates/burn-train/src/metric/confusion_stats.rs new file mode 100644 index 0000000..214ddc1 --- /dev/null +++ b/crates/burn-train/src/metric/confusion_stats.rs @@ -0,0 +1,340 @@ +use super::classification::{ClassReduction, ClassificationMetricConfig, DecisionRule}; +use burn_core::{ + prelude::{Bool, Int, Tensor}, + tensor::IndexingUpdateOp, +}; +use std::fmt::{self, Debug}; + +/// Input for confusion statistics error types. +#[derive(new, Debug, Clone)] +pub struct ConfusionStatsInput { + /// Sample x Class Non thresholded normalized predictions. + pub predictions: Tensor<2>, + /// Sample x Class one-hot encoded target. + pub targets: Tensor<2, Bool>, +} + +impl From for (Tensor<2>, Tensor<2, Bool>) { + fn from(input: ConfusionStatsInput) -> Self { + (input.predictions, input.targets) + } +} + +impl From<(Tensor<2>, Tensor<2, Bool>)> for ConfusionStatsInput { + fn from(value: (Tensor<2>, Tensor<2, Bool>)) -> Self { + Self::new(value.0, value.1) + } +} + +#[derive(Clone)] +pub struct ConfusionStats { + confusion_classes: Tensor<2, Int>, + class_reduction: ClassReduction, +} + +impl Debug for ConfusionStats { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let to_vec = |tensor_data: Tensor<1>| { + tensor_data + .to_data() + .to_vec::() + .expect("A vector representation of the input Tensor is expected") + }; + let ratio_of_support_vec = + |metric: Tensor<1>| to_vec(self.clone().ratio_of_support(metric)); + f.debug_struct("ConfusionStats") + .field("tp", &ratio_of_support_vec(self.clone().true_positive())) + .field("fp", &ratio_of_support_vec(self.clone().false_positive())) + .field("tn", &ratio_of_support_vec(self.clone().true_negative())) + .field("fn", &ratio_of_support_vec(self.clone().false_negative())) + .field("support", &to_vec(self.clone().support())) + .finish() + } +} + +impl ConfusionStats { + /// Expects `predictions` to be normalized. + pub fn new(input: &ConfusionStatsInput, config: &ClassificationMetricConfig) -> Self { + let prediction_mask = match config.decision_rule { + DecisionRule::Threshold(threshold) => input.predictions.clone().greater_elem(threshold), + DecisionRule::TopK(top_k) => { + let mask = input.predictions.zeros_like(); + let indexes = + input + .predictions + .clone() + .argsort_descending(1) + .narrow(1, 0, top_k.get()); + let values = indexes.ones_like().float(); + mask.scatter(1, indexes, values, IndexingUpdateOp::Add) + .bool() + } + }; + Self { + confusion_classes: prediction_mask.int() + input.targets.clone().int() * 2, + class_reduction: config.class_reduction, + } + } + + /// sum over samples + fn aggregate(sample_class_mask: Tensor<2, Bool>, class_reduction: ClassReduction) -> Tensor<1> { + use ClassReduction::{Macro, Micro}; + match class_reduction { + Micro => sample_class_mask.float().sum(), + Macro => sample_class_mask.float().sum_dim(0).squeeze_dim(0), + } + } + + pub fn true_positive(self) -> Tensor<1> { + Self::aggregate(self.confusion_classes.equal_elem(3), self.class_reduction) + } + + pub fn true_negative(self) -> Tensor<1> { + Self::aggregate(self.confusion_classes.equal_elem(0), self.class_reduction) + } + + pub fn false_positive(self) -> Tensor<1> { + Self::aggregate(self.confusion_classes.equal_elem(1), self.class_reduction) + } + + pub fn false_negative(self) -> Tensor<1> { + Self::aggregate(self.confusion_classes.equal_elem(2), self.class_reduction) + } + + pub fn positive(self) -> Tensor<1> { + self.clone().true_positive() + self.false_negative() + } + + pub fn negative(self) -> Tensor<1> { + self.clone().true_negative() + self.false_positive() + } + + pub fn predicted_positive(self) -> Tensor<1> { + self.clone().true_positive() + self.false_positive() + } + + pub fn support(self) -> Tensor<1> { + self.clone().positive() + self.negative() + } + + pub fn ratio_of_support(self, metric: Tensor<1>) -> Tensor<1> { + metric / self.clone().support() + } +} + +#[cfg(test)] +mod tests { + use super::{ConfusionStats, ConfusionStatsInput}; + use crate::{ + metric::classification::{ClassReduction, ClassificationMetricConfig, DecisionRule}, + tests::{ClassificationType, THRESHOLD, dummy_classification_input}, + }; + use burn_core::prelude::TensorData; + use rstest::{fixture, rstest}; + use std::num::NonZeroUsize; + + fn top_k_config( + top_k: NonZeroUsize, + class_reduction: ClassReduction, + ) -> ClassificationMetricConfig { + ClassificationMetricConfig { + decision_rule: DecisionRule::TopK(top_k), + class_reduction, + } + } + #[fixture] + #[once] + fn top_k_config_k1_micro() -> ClassificationMetricConfig { + top_k_config(NonZeroUsize::new(1).unwrap(), ClassReduction::Micro) + } + + #[fixture] + #[once] + fn top_k_config_k1_macro() -> ClassificationMetricConfig { + top_k_config(NonZeroUsize::new(1).unwrap(), ClassReduction::Macro) + } + #[fixture] + #[once] + fn top_k_config_k2_micro() -> ClassificationMetricConfig { + top_k_config(NonZeroUsize::new(2).unwrap(), ClassReduction::Micro) + } + #[fixture] + #[once] + fn top_k_config_k2_macro() -> ClassificationMetricConfig { + top_k_config(NonZeroUsize::new(2).unwrap(), ClassReduction::Macro) + } + + fn threshold_config( + threshold: f64, + class_reduction: ClassReduction, + ) -> ClassificationMetricConfig { + ClassificationMetricConfig { + decision_rule: DecisionRule::Threshold(threshold), + class_reduction, + } + } + #[fixture] + #[once] + fn threshold_config_micro() -> ClassificationMetricConfig { + threshold_config(THRESHOLD, ClassReduction::Micro) + } + #[fixture] + #[once] + fn threshold_config_macro() -> ClassificationMetricConfig { + threshold_config(THRESHOLD, ClassReduction::Macro) + } + + #[rstest] + #[case::binary_micro(ClassificationType::Binary, threshold_config_micro(), [1].into())] + #[case::binary_macro(ClassificationType::Binary, threshold_config_macro(), [1].into())] + #[case::multiclass_micro(ClassificationType::Multiclass, top_k_config_k1_micro(), [3].into())] + #[case::multiclass_macro(ClassificationType::Multiclass, top_k_config_k1_macro(), [1, 1, 1].into())] + #[case::multiclass_micro(ClassificationType::Multiclass, top_k_config_k2_micro(), [4].into())] + #[case::multiclass_macro(ClassificationType::Multiclass, top_k_config_k2_macro(), [2, 1, 1].into())] + #[case::multilabel_micro(ClassificationType::Multilabel, threshold_config_micro(), [5].into())] + #[case::multilabel_macro(ClassificationType::Multilabel, threshold_config_macro(), [2, 2, 1].into())] + fn test_true_positive( + #[case] classification_type: ClassificationType, + #[case] config: ClassificationMetricConfig, + #[case] expected: Vec, + ) { + let input: ConfusionStatsInput = dummy_classification_input(&classification_type).into(); + ConfusionStats::new(&input, &config) + .true_positive() + .int() + .into_data() + .assert_eq(&TensorData::from(expected.as_slice()), false); + } + + #[rstest] + #[case::binary_micro(ClassificationType::Binary, threshold_config_micro(), [2].into())] + #[case::binary_macro(ClassificationType::Binary, threshold_config_macro(), [2].into())] + #[case::multiclass_micro(ClassificationType::Multiclass, top_k_config_k1_micro(), [8].into())] + #[case::multiclass_macro(ClassificationType::Multiclass, top_k_config_k1_macro(), [2, 3, 3].into())] + #[case::multiclass_micro(ClassificationType::Multiclass, top_k_config_k2_micro(), [4].into())] + #[case::multiclass_macro(ClassificationType::Multiclass, top_k_config_k2_macro(), [1, 1, 2].into())] + #[case::multilabel_micro(ClassificationType::Multilabel, threshold_config_micro(), [3].into())] + #[case::multilabel_macro(ClassificationType::Multilabel, threshold_config_macro(), [0, 2, 1].into())] + fn test_true_negative( + #[case] classification_type: ClassificationType, + #[case] config: ClassificationMetricConfig, + #[case] expected: Vec, + ) { + let input: ConfusionStatsInput = dummy_classification_input(&classification_type).into(); + ConfusionStats::new(&input, &config) + .true_negative() + .int() + .into_data() + .assert_eq(&TensorData::from(expected.as_slice()), false); + } + + #[rstest] + #[case::binary_micro(ClassificationType::Binary, threshold_config_micro(), [1].into())] + #[case::binary_macro(ClassificationType::Binary, threshold_config_macro(), [1].into())] + #[case::multiclass_micro(ClassificationType::Multiclass, top_k_config_k1_micro(), [2].into())] + #[case::multiclass_macro(ClassificationType::Multiclass, top_k_config_k1_macro(), [1, 1, 0].into())] + #[case::multiclass_micro(ClassificationType::Multiclass, top_k_config_k2_micro(), [6].into())] + #[case::multiclass_macro(ClassificationType::Multiclass, top_k_config_k2_macro(), [2, 3, 1].into())] + #[case::multilabel_micro(ClassificationType::Multilabel, threshold_config_micro(), [3].into())] + #[case::multilabel_macro(ClassificationType::Multilabel, threshold_config_macro(), [1, 1, 1].into())] + fn test_false_positive( + #[case] classification_type: ClassificationType, + #[case] config: ClassificationMetricConfig, + #[case] expected: Vec, + ) { + let input: ConfusionStatsInput = dummy_classification_input(&classification_type).into(); + ConfusionStats::new(&input, &config) + .false_positive() + .int() + .into_data() + .assert_eq(&TensorData::from(expected.as_slice()), false); + } + + #[rstest] + #[case::binary_micro(ClassificationType::Binary, threshold_config_micro(), [1].into())] + #[case::binary_macro(ClassificationType::Binary, threshold_config_macro(), [1].into())] + #[case::multiclass_micro(ClassificationType::Multiclass, top_k_config_k1_micro(), [2].into())] + #[case::multiclass_macro(ClassificationType::Multiclass, top_k_config_k1_macro(), [1, 0, 1].into())] + #[case::multiclass_micro(ClassificationType::Multiclass, top_k_config_k2_micro(), [1].into())] + #[case::multiclass_macro(ClassificationType::Multiclass, top_k_config_k2_macro(), [0, 0, 1].into())] + #[case::multilabel_micro(ClassificationType::Multilabel, threshold_config_micro(), [4].into())] + #[case::multilabel_macro(ClassificationType::Multilabel, threshold_config_macro(), [2, 0, 2].into())] + fn test_false_negatives( + #[case] classification_type: ClassificationType, + #[case] config: ClassificationMetricConfig, + #[case] expected: Vec, + ) { + let input: ConfusionStatsInput = dummy_classification_input(&classification_type).into(); + ConfusionStats::new(&input, &config) + .false_negative() + .int() + .into_data() + .assert_eq(&TensorData::from(expected.as_slice()), false); + } + + #[rstest] + #[case::binary_micro(ClassificationType::Binary, threshold_config_micro(), [2].into())] + #[case::binary_macro(ClassificationType::Binary, threshold_config_macro(), [2].into())] + #[case::multiclass_micro(ClassificationType::Multiclass, top_k_config_k1_micro(), [5].into())] + #[case::multiclass_macro(ClassificationType::Multiclass, top_k_config_k1_macro(), [2, 1, 2].into())] + #[case::multiclass_micro(ClassificationType::Multiclass, top_k_config_k2_micro(), [5].into())] + #[case::multiclass_macro(ClassificationType::Multiclass, top_k_config_k2_macro(), [2, 1, 2].into())] + #[case::multilabel_micro(ClassificationType::Multilabel, threshold_config_micro(), [9].into())] + #[case::multilabel_macro(ClassificationType::Multilabel, threshold_config_macro(), [4, 2, 3].into())] + fn test_positive( + #[case] classification_type: ClassificationType, + #[case] config: ClassificationMetricConfig, + #[case] expected: Vec, + ) { + let input: ConfusionStatsInput = dummy_classification_input(&classification_type).into(); + ConfusionStats::new(&input, &config) + .positive() + .int() + .into_data() + .assert_eq(&TensorData::from(expected.as_slice()), false); + } + + #[rstest] + #[case::binary_micro(ClassificationType::Binary, threshold_config_micro(), [3].into())] + #[case::binary_macro(ClassificationType::Binary, threshold_config_macro(), [3].into())] + #[case::multiclass_micro(ClassificationType::Multiclass, top_k_config_k1_micro(), [10].into())] + #[case::multiclass_macro(ClassificationType::Multiclass, top_k_config_k1_macro(), [3, 4, 3].into())] + #[case::multiclass_micro(ClassificationType::Multiclass, top_k_config_k2_micro(), [10].into())] + #[case::multiclass_macro(ClassificationType::Multiclass, top_k_config_k2_macro(), [3, 4, 3].into())] + #[case::multilabel_micro(ClassificationType::Multilabel, threshold_config_micro(), [6].into())] + #[case::multilabel_macro(ClassificationType::Multilabel, threshold_config_macro(), [1, 3, 2].into())] + fn test_negative( + #[case] classification_type: ClassificationType, + #[case] config: ClassificationMetricConfig, + #[case] expected: Vec, + ) { + let input: ConfusionStatsInput = dummy_classification_input(&classification_type).into(); + ConfusionStats::new(&input, &config) + .negative() + .int() + .into_data() + .assert_eq(&TensorData::from(expected.as_slice()), false); + } + + #[rstest] + #[case::binary_micro(ClassificationType::Binary, threshold_config_micro(), [2].into())] + #[case::binary_macro(ClassificationType::Binary, threshold_config_macro(), [2].into())] + #[case::multiclass_micro(ClassificationType::Multiclass, top_k_config_k1_micro(), [5].into())] + #[case::multiclass_macro(ClassificationType::Multiclass, top_k_config_k1_macro(), [2, 2, 1].into())] + #[case::multiclass_micro(ClassificationType::Multiclass, top_k_config_k2_micro(), [10].into())] + #[case::multiclass_macro(ClassificationType::Multiclass, top_k_config_k2_macro(), [4, 4, 2].into())] + #[case::multilabel_micro(ClassificationType::Multilabel, threshold_config_micro(), [8].into())] + #[case::multilabel_macro(ClassificationType::Multilabel, threshold_config_macro(), [3, 3, 2].into())] + fn test_predicted_positive( + #[case] classification_type: ClassificationType, + #[case] config: ClassificationMetricConfig, + #[case] expected: Vec, + ) { + let input: ConfusionStatsInput = dummy_classification_input(&classification_type).into(); + ConfusionStats::new(&input, &config) + .predicted_positive() + .int() + .into_data() + .assert_eq(&TensorData::from(expected.as_slice()), false); + } +} diff --git a/crates/burn-train/src/metric/cpu_temp.rs b/crates/burn-train/src/metric/cpu_temp.rs new file mode 100644 index 0000000..ddc2e65 --- /dev/null +++ b/crates/burn-train/src/metric/cpu_temp.rs @@ -0,0 +1,77 @@ +use std::sync::Arc; + +/// CPU Temperature metric +use super::MetricMetadata; +use crate::metric::{Metric, MetricAttributes, MetricName, Numeric, NumericEntry, SerializedEntry}; +use systemstat::{Platform, System}; + +/// CPU Temperature in celsius degrees +#[derive(Clone)] +pub struct CpuTemperature { + name: MetricName, + temp_celsius: f32, + sys: Arc, +} + +impl CpuTemperature { + /// Creates a new CPU temp metric + pub fn new() -> Self { + let name = Arc::new("CPU Temperature".to_string()); + + Self { + name, + temp_celsius: 0., + sys: Arc::new(System::new()), + } + } +} + +impl Default for CpuTemperature { + fn default() -> Self { + CpuTemperature::new() + } +} + +impl Metric for CpuTemperature { + type Input = (); + + fn update(&mut self, _item: &Self::Input, _metadata: &MetricMetadata) -> SerializedEntry { + match self.sys.cpu_temp() { + Ok(temp) => self.temp_celsius = temp, + Err(_) => self.temp_celsius = f32::NAN, + } + + let formatted = match self.temp_celsius.is_nan() { + true => format!("{}: NaN °C", self.name()), + false => format!("{}: {:.2} °C", self.name(), self.temp_celsius), + }; + let raw = format!("{:.2}", self.temp_celsius); + + SerializedEntry::new(formatted, raw) + } + + fn clear(&mut self) {} + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + super::NumericAttributes { + unit: Some("°C".to_string()), + higher_is_better: false, + ..Default::default() + } + .into() + } +} + +impl Numeric for CpuTemperature { + fn value(&self) -> NumericEntry { + NumericEntry::Value(self.temp_celsius as f64) + } + + fn running_value(&self) -> NumericEntry { + NumericEntry::Value(self.temp_celsius as f64) + } +} diff --git a/crates/burn-train/src/metric/cpu_use.rs b/crates/burn-train/src/metric/cpu_use.rs new file mode 100644 index 0000000..693a4fb --- /dev/null +++ b/crates/burn-train/src/metric/cpu_use.rs @@ -0,0 +1,104 @@ +use super::MetricMetadata; +use crate::metric::{Metric, MetricAttributes, MetricName, Numeric, NumericEntry, SerializedEntry}; +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; +use sysinfo::{CpuRefreshKind, RefreshKind, System}; + +/// General CPU Usage metric +pub struct CpuUse { + name: MetricName, + last_refresh: Instant, + refresh_frequency: Duration, + sys: System, + current: f64, +} + +impl Clone for CpuUse { + fn clone(&self) -> Self { + Self { + name: self.name.clone(), + last_refresh: self.last_refresh, + refresh_frequency: self.refresh_frequency, + sys: System::new(), + current: self.current, + } + } +} + +impl CpuUse { + /// Creates a new CPU metric + pub fn new() -> Self { + let mut sys = System::new(); + let current = Self::refresh(&mut sys); + let name = "CPU Usage".to_string(); + + Self { + name: Arc::new(name), + last_refresh: Instant::now(), + refresh_frequency: Duration::from_millis(200), + sys, + current, + } + } + + fn refresh(sys: &mut System) -> f64 { + sys.refresh_specifics( + RefreshKind::nothing().with_cpu(CpuRefreshKind::nothing().with_cpu_usage()), + ); + + let cpus = sys.cpus(); + let num_cpus = cpus.len(); + let use_percentage = cpus.iter().fold(0.0, |acc, cpu| acc + cpu.cpu_usage()) as f64; + + use_percentage / num_cpus as f64 + } +} + +impl Default for CpuUse { + fn default() -> Self { + CpuUse::new() + } +} + +impl Metric for CpuUse { + type Input = (); + + fn update(&mut self, _item: &Self::Input, _metadata: &MetricMetadata) -> SerializedEntry { + if self.last_refresh.elapsed() >= self.refresh_frequency { + self.current = Self::refresh(&mut self.sys); + self.last_refresh = Instant::now(); + } + + let formatted = format!("{}: {:.2} %", self.name(), self.current); + let raw = format!("{:.2}", self.current); + + SerializedEntry::new(formatted, raw) + } + + fn clear(&mut self) {} + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + super::NumericAttributes { + unit: Some("%".to_string()), + higher_is_better: false, + ..Default::default() + } + .into() + } +} + +impl Numeric for CpuUse { + fn value(&self) -> NumericEntry { + NumericEntry::Value(self.current) + } + + fn running_value(&self) -> NumericEntry { + NumericEntry::Value(self.current) + } +} diff --git a/crates/burn-train/src/metric/cuda.rs b/crates/burn-train/src/metric/cuda.rs new file mode 100644 index 0000000..3bbcf08 --- /dev/null +++ b/crates/burn-train/src/metric/cuda.rs @@ -0,0 +1,108 @@ +use std::sync::Arc; + +use super::MetricMetadata; +use crate::metric::{Metric, MetricName, SerializedEntry}; +use nvml_wrapper::Nvml; + +/// Track basic cuda infos. +#[derive(Clone)] +pub struct CudaMetric { + name: MetricName, + nvml: Arc>, +} + +impl CudaMetric { + /// Creates a new metric for CUDA. + pub fn new() -> Self { + Self { + name: Arc::new("Cuda".to_string()), + nvml: Arc::new(Nvml::init().map(Some).unwrap_or_else(|err| { + log::warn!("Unable to initialize CUDA Metric: {err}"); + None + })), + } + } +} + +impl Default for CudaMetric { + fn default() -> Self { + Self::new() + } +} + +impl Metric for CudaMetric { + type Input = (); + + fn update(&mut self, _item: &(), _metadata: &MetricMetadata) -> SerializedEntry { + let not_available = + || SerializedEntry::new("Unavailable".to_string(), "Unavailable".to_string()); + + let available = |nvml: &Nvml| { + let mut formatted = String::new(); + let mut raw_running = String::new(); + + let device_count = match nvml.device_count() { + Ok(val) => val, + Err(err) => { + log::warn!("Unable to get the number of cuda devices: {err}"); + return not_available(); + } + }; + + for index in 0..device_count { + let device = match nvml.device_by_index(index) { + Ok(val) => val, + Err(err) => { + log::warn!("Unable to get device {index}: {err}"); + return not_available(); + } + }; + let memory_info = match device.memory_info() { + Ok(info) => info, + Err(err) => { + log::warn!("Unable to get memory info from device {index}: {err}"); + return not_available(); + } + }; + + let used_gb = memory_info.used as f64 * 1e-9; + let total_gb = memory_info.total as f64 * 1e-9; + + let memory_info_formatted = format!("{used_gb:.2}/{total_gb:.2} Gb"); + let memory_info_raw = format!("{used_gb}/{total_gb}"); + + formatted = format!("{formatted} GPU #{index} - Memory {memory_info_formatted}"); + raw_running = format!("{memory_info_raw} "); + + let utilization_rates = match device.utilization_rates() { + Ok(rate) => rate, + Err(err) => { + log::warn!("Unable to get utilization rates from device {index}: {err}"); + return not_available(); + } + }; + let utilization_rate_formatted = format!("{}%", utilization_rates.gpu); + formatted = format!("{formatted} - Usage {utilization_rate_formatted}"); + + // Power is the currency for perf/W. NVML reports milliwatts. + if let Ok(power_mw) = device.power_usage() { + let power_w = power_mw as f64 / 1000.0; + formatted = format!("{formatted} - Power {power_w:.1} W"); + } + } + + SerializedEntry::new(formatted, raw_running) + }; + + match self.nvml.as_ref() { + Some(nvml) => available(nvml), + None => not_available(), + } + } + + fn clear(&mut self) {} + + fn name(&self) -> MetricName { + self.name.clone() + } +} diff --git a/crates/burn-train/src/metric/fbetascore.rs b/crates/burn-train/src/metric/fbetascore.rs new file mode 100644 index 0000000..1602206 --- /dev/null +++ b/crates/burn-train/src/metric/fbetascore.rs @@ -0,0 +1,245 @@ +use crate::metric::{MetricName, Numeric}; + +use super::{ + Metric, MetricAttributes, MetricMetadata, NumericAttributes, NumericEntry, SerializedEntry, + classification::{ClassReduction, ClassificationMetricConfig, DecisionRule}, + confusion_stats::{ConfusionStats, ConfusionStatsInput}, + state::{FormatOptions, NumericMetricState}, +}; +use burn_core::prelude::Tensor; +use std::{num::NonZeroUsize, sync::Arc}; + +/// The [F-beta score](https://en.wikipedia.org/wiki/F-score) metric. +/// +/// The `beta` parameter represents the ratio of recall importance to precision importance. +/// `beta > 1` gives more weight to recall, while `beta < 1` favors precision. +#[derive(Clone)] +pub struct FBetaScoreMetric { + name: MetricName, + state: NumericMetricState, + config: ClassificationMetricConfig, + beta: f64, +} + +impl Default for FBetaScoreMetric { + fn default() -> Self { + Self::new(Default::default(), Default::default()) + } +} + +impl FBetaScoreMetric { + #[allow(dead_code)] + fn new(config: ClassificationMetricConfig, beta: f64) -> Self { + let name = Arc::new(format!( + "FBetaScore ({}) @ {:?} [{:?}]", + beta, config.decision_rule, config.class_reduction + )); + Self { + name, + config, + beta, + state: Default::default(), + } + } + + /// F-beta score metric for binary classification. + /// + /// # Arguments + /// + /// * `beta` - Positive real factor to weight recall's importance. + /// * `threshold` - The threshold to transform a probability into a binary prediction. + #[allow(dead_code)] + pub fn binary(beta: f64, threshold: f64) -> Self { + Self::new( + ClassificationMetricConfig { + decision_rule: DecisionRule::Threshold(threshold), + // binary classification results are the same independently of class_reduction + ..Default::default() + }, + beta, + ) + } + + /// F-beta score metric for multiclass classification. + /// + /// # Arguments + /// + /// * `beta` - Positive real factor to weight recall's importance. + /// * `top_k` - The number of highest predictions considered to find the correct label (typically `1`). + /// * `class_reduction` - [Class reduction](ClassReduction) type. + #[allow(dead_code)] + pub fn multiclass(beta: f64, top_k: usize, class_reduction: ClassReduction) -> Self { + Self::new( + ClassificationMetricConfig { + decision_rule: DecisionRule::TopK( + NonZeroUsize::new(top_k).expect("top_k must be non-zero"), + ), + class_reduction, + }, + beta, + ) + } + + /// F-beta score metric for multi-label classification. + /// + /// # Arguments + /// + /// * `beta` - Positive real factor to weight recall's importance. + /// * `threshold` - The threshold to transform a probability into a binary prediction. + /// * `class_reduction` - [Class reduction](ClassReduction) type. + #[allow(dead_code)] + pub fn multilabel(beta: f64, threshold: f64, class_reduction: ClassReduction) -> Self { + Self::new( + ClassificationMetricConfig { + decision_rule: DecisionRule::Threshold(threshold), + class_reduction, + }, + beta, + ) + } + + fn class_average(&self, mut aggregated_metric: Tensor<1>) -> f64 { + use ClassReduction::{Macro, Micro}; + let avg_tensor = match self.config.class_reduction { + Micro => aggregated_metric, + Macro => { + if aggregated_metric.clone().contains_nan().any().into_scalar() { + let nan_mask = aggregated_metric.clone().is_nan(); + aggregated_metric = aggregated_metric + .clone() + .select(0, nan_mask.bool_not().argwhere().squeeze_dim(1)) + } + aggregated_metric.mean() + } + }; + avg_tensor.into_scalar() + } +} + +impl Metric for FBetaScoreMetric { + type Input = ConfusionStatsInput; + + fn update(&mut self, input: &Self::Input, _metadata: &MetricMetadata) -> SerializedEntry { + let [sample_size, _] = input.predictions.dims(); + + let cf_stats = ConfusionStats::new(input, &self.config); + let scaled_true_positive = cf_stats.clone().true_positive() * (1.0 + self.beta.powi(2)); + let metric = self.class_average( + scaled_true_positive.clone() + / (scaled_true_positive + + cf_stats.clone().false_negative() * self.beta.powi(2) + + cf_stats.false_positive()), + ); + + self.state.update( + 100.0 * metric, + sample_size, + FormatOptions::new(self.name()).unit("%").precision(2), + ) + } + + fn clear(&mut self) { + self.state.reset() + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: Some("%".to_string()), + higher_is_better: true, + ..Default::default() + } + .into() + } +} + +impl Numeric for FBetaScoreMetric { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} + +#[cfg(test)] +mod tests { + use super::{ + ClassReduction::{self, *}, + FBetaScoreMetric, Metric, MetricMetadata, + }; + use crate::metric::Numeric; + use crate::tests::{ClassificationType, THRESHOLD, dummy_classification_input}; + use burn_core::tensor::TensorData; + use burn_core::tensor::Tolerance; + use rstest::rstest; + + #[rstest] + #[case::binary_b1(1.0, THRESHOLD, 0.5)] + #[case::binary_b2(2.0, THRESHOLD, 0.5)] + fn test_binary_fscore(#[case] beta: f64, #[case] threshold: f64, #[case] expected: f64) { + let input = dummy_classification_input(&ClassificationType::Binary).into(); + let mut metric = FBetaScoreMetric::binary(beta, threshold); + let _entry = metric.update(&input, &MetricMetadata::fake()); + TensorData::from([metric.value().current()]) + .assert_approx_eq::(&TensorData::from([expected * 100.0]), Tolerance::default()) + } + + #[rstest] + #[case::multiclass_b1_micro_k1(1.0, Micro, 1, 3.0/5.0)] + #[case::multiclass_b1_micro_k2(1.0, Micro, 2, 2.0/(5.0/4.0 + 10.0/4.0))] + #[case::multiclass_b1_macro_k1(1.0, Macro, 1, (0.5 + 2.0/(1.0 + 2.0) + 2.0/(2.0 + 1.0))/3.0)] + #[case::multiclass_b1_macro_k2(1.0, Macro, 2, (2.0/(1.0 + 2.0) + 2.0/(1.0 + 4.0) + 0.5)/3.0)] + #[case::multiclass_b2_micro_k1(2.0, Micro, 1, 3.0/5.0)] + #[case::multiclass_b2_micro_k2(2.0, Micro, 2, 5.0*4.0/(4.0*5.0 + 10.0))] + #[case::multiclass_b2_macro_k1(2.0, Macro, 1, (0.5 + 5.0/(4.0 + 2.0) + 5.0/(8.0 + 1.0))/3.0)] + #[case::multiclass_b2_macro_k2(2.0, Macro, 2, (5.0/(4.0 + 2.0) + 5.0/(4.0 + 4.0) + 0.5)/3.0)] + fn test_multiclass_fscore( + #[case] beta: f64, + #[case] class_reduction: ClassReduction, + #[case] top_k: usize, + #[case] expected: f64, + ) { + let input = dummy_classification_input(&ClassificationType::Multiclass).into(); + let mut metric = FBetaScoreMetric::multiclass(beta, top_k, class_reduction); + let _entry = metric.update(&input, &MetricMetadata::fake()); + TensorData::from([metric.value().current()]) + .assert_approx_eq::(&TensorData::from([expected * 100.0]), Tolerance::default()) + } + + #[rstest] + #[case::multilabel_micro(1.0, Micro, THRESHOLD, 2.0/(9.0/5.0 + 8.0/5.0))] + #[case::multilabel_macro(1.0, Macro, THRESHOLD, (2.0/(2.0 + 3.0/2.0) + 2.0/(1.0 + 3.0/2.0) + 2.0/(3.0+2.0))/3.0)] + #[case::multilabel_micro(2.0, Micro, THRESHOLD, 5.0/(4.0*9.0/5.0 + 8.0/5.0))] + #[case::multilabel_macro(2.0, Macro, THRESHOLD, (5.0/(8.0 + 3.0/2.0) + 5.0/(4.0 + 3.0/2.0) + 5.0/(12.0+2.0))/3.0)] + fn test_multilabel_fscore( + #[case] beta: f64, + #[case] class_reduction: ClassReduction, + #[case] threshold: f64, + #[case] expected: f64, + ) { + let input = dummy_classification_input(&ClassificationType::Multilabel).into(); + let mut metric = FBetaScoreMetric::multilabel(beta, threshold, class_reduction); + let _entry = metric.update(&input, &MetricMetadata::fake()); + TensorData::from([metric.value().current()]) + .assert_approx_eq::(&TensorData::from([expected * 100.0]), Tolerance::default()) + } + + #[test] + fn test_parameterized_unique_name() { + let metric_a = FBetaScoreMetric::multiclass(0.5, 1, ClassReduction::Macro); + let metric_b = FBetaScoreMetric::multiclass(0.5, 2, ClassReduction::Macro); + let metric_c = FBetaScoreMetric::multiclass(0.5, 1, ClassReduction::Macro); + + assert_ne!(metric_a.name(), metric_b.name()); + assert_eq!(metric_a.name(), metric_c.name()); + + let metric_a = FBetaScoreMetric::binary(0.5, 0.5); + let metric_b = FBetaScoreMetric::binary(0.75, 0.5); + assert_ne!(metric_a.name(), metric_b.name()); + } +} diff --git a/crates/burn-train/src/metric/hamming.rs b/crates/burn-train/src/metric/hamming.rs new file mode 100644 index 0000000..31a472b --- /dev/null +++ b/crates/burn-train/src/metric/hamming.rs @@ -0,0 +1,190 @@ +use std::sync::Arc; + +use super::state::{FormatOptions, NumericMetricState}; +use super::{MetricMetadata, SerializedEntry}; +use crate::metric::{ + Metric, MetricAttributes, MetricName, Numeric, NumericAttributes, NumericEntry, +}; +use burn_core::tensor::{Int, Tensor, activation::sigmoid}; + +/// The hamming score, sometimes referred to as multi-label or label-based accuracy. +#[derive(Clone)] +pub struct HammingScore { + name: MetricName, + state: NumericMetricState, + threshold: f32, + sigmoid: bool, +} + +/// The [hamming score](HammingScore) input type. +#[derive(new)] +pub struct HammingScoreInput { + outputs: Tensor<2>, + targets: Tensor<2, Int>, +} + +impl HammingScore { + /// Creates the metric. + pub fn new() -> Self { + Self::default() + } + + fn update_name(&mut self) { + self.name = Arc::new(format!("Hamming Score @ Threshold({})", self.threshold)); + } + + /// Sets the threshold. + pub fn with_threshold(mut self, threshold: f32) -> Self { + self.threshold = threshold; + self.update_name(); + self + } + + /// Sets the sigmoid activation function usage. + pub fn with_sigmoid(mut self, sigmoid: bool) -> Self { + self.sigmoid = sigmoid; + self.update_name(); + self + } +} + +impl Default for HammingScore { + /// Creates a new metric instance with default values. + fn default() -> Self { + let threshold = 0.5; + let name = Arc::new(format!("Hamming Score @ Threshold({})", threshold)); + + Self { + name, + state: NumericMetricState::default(), + threshold, + sigmoid: false, + } + } +} + +impl Metric for HammingScore { + type Input = HammingScoreInput; + + fn update(&mut self, input: &HammingScoreInput, _metadata: &MetricMetadata) -> SerializedEntry { + let [batch_size, _n_classes] = input.outputs.dims(); + + let targets = input.targets.clone(); + + let mut outputs = input.outputs.clone(); + + if self.sigmoid { + outputs = sigmoid(outputs); + } + + let score = outputs + .greater_elem(self.threshold) + .equal(targets.bool()) + .float() + .mean() + .into_scalar::(); + + self.state.update( + 100.0 * score, + batch_size, + FormatOptions::new(self.name()).unit("%").precision(2), + ) + } + + fn clear(&mut self) { + self.state.reset() + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: Some("%".to_string()), + higher_is_better: true, + ..Default::default() + } + .into() + } +} + +impl Numeric for HammingScore { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_hamming_score() { + let device = Default::default(); + let mut metric = HammingScore::new(); + + let x = Tensor::from_data( + [ + [0.32, 0.52, 0.38, 0.68, 0.61], // with x > 0.5: [0, 1, 0, 1, 1] + [0.43, 0.31, 0.21, 0.63, 0.53], // [0, 0, 0, 1, 1] + [0.44, 0.25, 0.71, 0.39, 0.73], // [0, 0, 1, 0, 1] + [0.49, 0.37, 0.68, 0.39, 0.31], // [0, 0, 1, 0, 0] + ], + &device, + ); + let y = Tensor::from_data( + [ + [0, 1, 0, 1, 1], + [0, 0, 0, 1, 1], + [0, 0, 1, 0, 1], + [0, 0, 1, 0, 0], + ], + &device, + ); + + let _entry = metric.update( + &HammingScoreInput::new(x.clone(), y.clone()), + &MetricMetadata::fake(), + ); + assert_eq!(100.0, metric.value().current()); + + // Invert all targets: y = (1 - y) + let y = y.neg().add_scalar(1); + let _entry = metric.update( + &HammingScoreInput::new(x.clone(), y), // invert targets (1 - y) + &MetricMetadata::fake(), + ); + assert_eq!(0.0, metric.value().current()); + + // Invert 5 target values -> 1 - (5/20) = 0.75 + let y = Tensor::from_data( + [ + [0, 1, 1, 0, 1], + [0, 0, 0, 0, 1], + [0, 0, 0, 0, 1], + [0, 1, 1, 0, 0], + ], + &device, + ); + let _entry = metric.update( + &HammingScoreInput::new(x, y), // invert targets (1 - y) + &MetricMetadata::fake(), + ); + assert_eq!(75.0, metric.value().current()); + } + + #[test] + fn test_parameterized_unique_name() { + let metric_a = HammingScore::new().with_threshold(0.5); + let metric_b = HammingScore::new().with_threshold(0.75); + let metric_c = HammingScore::new().with_threshold(0.5); + + assert_ne!(metric_a.name(), metric_b.name()); + assert_eq!(metric_a.name(), metric_c.name()); + } +} diff --git a/crates/burn-train/src/metric/iteration.rs b/crates/burn-train/src/metric/iteration.rs new file mode 100644 index 0000000..eec8cb9 --- /dev/null +++ b/crates/burn-train/src/metric/iteration.rs @@ -0,0 +1,90 @@ +use std::sync::Arc; + +use super::MetricMetadata; +use super::SerializedEntry; +use super::state::FormatOptions; +use super::state::NumericMetricState; +use crate::metric::MetricName; +use crate::metric::Numeric; +use crate::metric::{Metric, MetricAttributes, NumericAttributes, NumericEntry}; + +/// The loss metric. +#[derive(Clone)] +pub struct IterationSpeedMetric { + name: MetricName, + state: NumericMetricState, + instant: Option, +} + +impl Default for IterationSpeedMetric { + fn default() -> Self { + Self::new() + } +} + +impl IterationSpeedMetric { + /// Create the metric. + pub fn new() -> Self { + Self { + name: Arc::new("Iteration Speed".to_string()), + state: Default::default(), + instant: Default::default(), + } + } +} + +impl Metric for IterationSpeedMetric { + type Input = (); + + fn update(&mut self, _: &Self::Input, metadata: &MetricMetadata) -> SerializedEntry { + let raw = match self.instant { + Some(val) => { + // If iteration is not logged, compute the speed over the number of items processed. + // 1 iteration should equal 1 item when iteration is not logged. + metadata + .iteration + .unwrap_or(metadata.progress.items_processed) as f64 + / val.elapsed().as_secs_f64() + } + None => { + self.instant = Some(std::time::Instant::now()); + 0.0 + } + }; + + self.state.update( + raw, + 1, + FormatOptions::new(self.name()) + .unit("iter/sec") + .precision(2), + ) + } + + fn clear(&mut self) { + self.instant = None; + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: Some("iter/sec".to_string()), + higher_is_better: true, + ..Default::default() + } + .into() + } +} + +impl Numeric for IterationSpeedMetric { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} diff --git a/crates/burn-train/src/metric/learning_rate.rs b/crates/burn-train/src/metric/learning_rate.rs new file mode 100644 index 0000000..5809e83 --- /dev/null +++ b/crates/burn-train/src/metric/learning_rate.rs @@ -0,0 +1,69 @@ +use std::sync::Arc; + +use super::{ + MetricAttributes, MetricMetadata, NumericAttributes, NumericEntry, + state::{FormatOptions, NumericMetricState}, +}; +use crate::metric::{Metric, MetricName, Numeric, SerializedEntry}; + +/// Track the learning rate across iterations. +#[derive(Clone)] +pub struct LearningRateMetric { + name: MetricName, + state: NumericMetricState, +} + +impl LearningRateMetric { + /// Creates a new learning rate metric. + pub fn new() -> Self { + Self { + name: Arc::new("Learning Rate".to_string()), + state: NumericMetricState::new(), + } + } +} + +impl Default for LearningRateMetric { + fn default() -> Self { + Self::new() + } +} + +impl Metric for LearningRateMetric { + type Input = (); + + fn update(&mut self, _item: &(), metadata: &MetricMetadata) -> SerializedEntry { + // TODO: We only log the default learning rate. Yet another motivation to introduce metric groups. + let lr = metadata.lr.as_ref().map(|val| val.base()).unwrap_or(0.0); + + self.state + .update(lr, 1, FormatOptions::new(self.name()).precision(2)) + } + + fn clear(&mut self) { + self.state.reset() + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: None, + higher_is_better: false, + ..Default::default() + } + .into() + } +} + +impl Numeric for LearningRateMetric { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} diff --git a/crates/burn-train/src/metric/loss.rs b/crates/burn-train/src/metric/loss.rs new file mode 100644 index 0000000..a0b2ef4 --- /dev/null +++ b/crates/burn-train/src/metric/loss.rs @@ -0,0 +1,87 @@ +use std::sync::Arc; + +use super::MetricMetadata; +use super::SerializedEntry; +use super::state::FormatOptions; +use super::state::NumericMetricState; +use crate::metric::MetricName; +use crate::metric::{Metric, MetricAttributes, Numeric, NumericAttributes, NumericEntry}; +use burn_core::tensor::Tensor; + +/// The loss metric. +#[derive(Clone)] +pub struct LossMetric { + name: Arc, + state: NumericMetricState, +} + +/// The [loss metric](LossMetric) input type. +#[derive(new)] +pub struct LossInput { + tensor: Tensor<1>, +} + +impl Default for LossMetric { + fn default() -> Self { + Self::new() + } +} + +impl LossMetric { + /// Create the metric. + pub fn new() -> Self { + Self { + name: Arc::new("Loss".to_string()), + state: NumericMetricState::default(), + } + } +} + +impl Metric for LossMetric { + type Input = LossInput; + + fn update(&mut self, loss: &Self::Input, _metadata: &MetricMetadata) -> SerializedEntry { + let [batch_size] = loss.tensor.dims(); + let loss = loss + .tensor + .clone() + .mean() + .into_data() + .iter::() + .next() + .unwrap(); + + self.state.update( + loss, + batch_size, + FormatOptions::new(self.name()).precision(2), + ) + } + + fn clear(&mut self) { + self.state.reset() + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: None, + higher_is_better: false, + ..Default::default() + } + .into() + } +} + +impl Numeric for LossMetric { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} diff --git a/crates/burn-train/src/metric/memory_use.rs b/crates/burn-train/src/metric/memory_use.rs new file mode 100644 index 0000000..8b4b781 --- /dev/null +++ b/crates/burn-train/src/metric/memory_use.rs @@ -0,0 +1,112 @@ +/// RAM use metric +use super::{MetricAttributes, MetricMetadata, NumericAttributes}; +use crate::metric::{Metric, Numeric, NumericEntry, SerializedEntry}; +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; +use sysinfo::System; + +/// Memory information +pub struct CpuMemory { + name: Arc, + last_refresh: Instant, + refresh_frequency: Duration, + sys: System, + ram_bytes_total: u64, + ram_bytes_used: u64, +} + +impl Clone for CpuMemory { + fn clone(&self) -> Self { + Self { + name: self.name.clone(), + last_refresh: self.last_refresh, + refresh_frequency: self.refresh_frequency, + sys: System::new(), + ram_bytes_total: self.ram_bytes_total, + ram_bytes_used: self.ram_bytes_used, + } + } +} + +impl CpuMemory { + /// Creates a new memory metric + pub fn new() -> Self { + let mut metric = Self { + name: Arc::new("CPU Memory".into()), + last_refresh: Instant::now(), + refresh_frequency: Duration::from_millis(200), + sys: System::new(), + ram_bytes_total: 0, + ram_bytes_used: 0, + }; + metric.refresh(); + metric + } + + fn refresh(&mut self) { + self.sys.refresh_memory(); + self.last_refresh = Instant::now(); + + // bytes of RAM available + self.ram_bytes_total = self.sys.total_memory(); + + // bytes of RAM in use + self.ram_bytes_used = self.sys.used_memory(); + } +} + +impl Default for CpuMemory { + fn default() -> Self { + CpuMemory::new() + } +} + +impl Metric for CpuMemory { + type Input = (); + + fn update(&mut self, _item: &Self::Input, _metadata: &MetricMetadata) -> SerializedEntry { + if self.last_refresh.elapsed() >= self.refresh_frequency { + self.refresh(); + } + + let raw = bytes2gb(self.ram_bytes_used); + let formatted = format!( + "RAM Used: {:.2} / {:.2} Gb", + raw, + bytes2gb(self.ram_bytes_total), + ); + + SerializedEntry::new(formatted, raw.to_string()) + } + + fn clear(&mut self) {} + + fn name(&self) -> Arc { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: Some("Gb".to_string()), + higher_is_better: false, + ..Default::default() + } + .into() + } +} + +impl Numeric for CpuMemory { + fn value(&self) -> NumericEntry { + NumericEntry::Value(bytes2gb(self.ram_bytes_used)) + } + + fn running_value(&self) -> NumericEntry { + NumericEntry::Value(bytes2gb(self.ram_bytes_used)) + } +} + +fn bytes2gb(bytes: u64) -> f64 { + bytes as f64 / 1e9 +} diff --git a/crates/burn-train/src/metric/mod.rs b/crates/burn-train/src/metric/mod.rs new file mode 100644 index 0000000..ca8c574 --- /dev/null +++ b/crates/burn-train/src/metric/mod.rs @@ -0,0 +1,77 @@ +/// State module. +pub mod state; +/// Module responsible to save and exposes data collected during training. +pub mod store; +/// Metrics module for vision tasks. +#[cfg(feature = "vision")] +pub mod vision; + +//Metrics for reinforcement learning. +#[cfg(feature = "rl")] +mod rl; +#[cfg(feature = "rl")] +pub use rl::*; + +// System metrics +#[cfg(feature = "sys-metrics")] +mod cpu_temp; +#[cfg(feature = "sys-metrics")] +mod cpu_use; +#[cfg(feature = "sys-metrics")] +mod cuda; +#[cfg(feature = "sys-metrics")] +mod memory_use; +#[cfg(feature = "sys-metrics")] +pub use cpu_temp::*; +#[cfg(feature = "sys-metrics")] +pub use cpu_use::*; +#[cfg(feature = "sys-metrics")] +pub use cuda::*; +#[cfg(feature = "sys-metrics")] +pub use memory_use::*; + +// Training metrics +mod acc; +mod auc_pr; +mod auroc; +mod base; +mod bleu; +mod cer; +mod confusion_stats; +mod fbetascore; +mod hamming; +mod iteration; +mod learning_rate; +mod loss; +mod perplexity; +mod precision; +mod recall; +mod rouge; +mod top_k_acc; +mod wer; + +pub use acc::*; +pub use auc_pr::*; +pub use auroc::*; +pub use base::*; +pub use bleu::*; +pub use cer::*; +pub use confusion_stats::ConfusionStatsInput; +pub use fbetascore::*; +pub use hamming::*; +pub use iteration::*; +pub use learning_rate::*; +pub use loss::*; +pub use perplexity::*; +pub use precision::*; +pub use recall::*; +pub use rouge::*; +pub use top_k_acc::*; +pub use wer::*; + +pub(crate) mod classification; +pub(crate) mod processor; + +pub use crate::metric::classification::ClassReduction; +// Expose `ItemLazy` so it can be implemented for custom types +pub use processor::ItemLazy; diff --git a/crates/burn-train/src/metric/perplexity.rs b/crates/burn-train/src/metric/perplexity.rs new file mode 100644 index 0000000..78ac96a --- /dev/null +++ b/crates/burn-train/src/metric/perplexity.rs @@ -0,0 +1,429 @@ +use super::state::FormatOptions; +use super::{MetricMetadata, NumericEntry, SerializedEntry, format_float}; +use crate::metric::{Metric, MetricAttributes, MetricName, Numeric, NumericAttributes}; +use burn_core::tensor::{Int, Tensor}; + +/// Custom state for perplexity metric that correctly accumulates negative log-likelihood. +/// +/// Unlike other metrics that can be averaged, perplexity requires special handling: +/// - Accumulate total negative log-likelihood across all tokens +/// - Accumulate total number of effective tokens +/// - Compute perplexity as exp(total_nll / total_tokens) only at the end +#[derive(Clone)] +struct PerplexityState { + /// Sum of negative log-likelihood across all tokens + sum_nll: f64, + /// Total number of effective tokens (excluding padding) + total_tokens: usize, + /// Current batch perplexity (for display purposes) + current: f64, +} + +impl PerplexityState { + fn new() -> Self { + Self { + sum_nll: 0.0, + total_tokens: 0, + current: f64::NAN, + } + } + + fn reset(&mut self) { + self.sum_nll = 0.0; + self.total_tokens = 0; + self.current = f64::NAN; + } + + /// Update state with negative log-likelihood and token count from current batch + fn update( + &mut self, + sum_log_prob: f64, + effective_tokens: usize, + format: FormatOptions, + ) -> SerializedEntry { + // sum_log_prob is already the sum of log probabilities (negative values) + // We need to negate it to get negative log-likelihood + let batch_nll = -sum_log_prob; + + // Accumulate across batches + self.sum_nll += batch_nll; + self.total_tokens += effective_tokens; + + // Compute current batch perplexity for display + let batch_perplexity = if effective_tokens > 0 { + (batch_nll / effective_tokens as f64).exp() + } else { + f64::INFINITY + }; + self.current = batch_perplexity; + + // Compute running epoch perplexity + let epoch_perplexity = if self.total_tokens > 0 { + (self.sum_nll / self.total_tokens as f64).exp() + } else { + f64::INFINITY + }; + + // Format for display + let (formatted_current, formatted_running) = match format.precision_value() { + Some(precision) => ( + format_float(batch_perplexity, precision), + format_float(epoch_perplexity, precision), + ), + None => (format!("{batch_perplexity}"), format!("{epoch_perplexity}")), + }; + + let formatted = match format.unit_value() { + Some(unit) => { + format!("epoch {formatted_running} {unit} - batch {formatted_current} {unit}") + } + None => format!("epoch {formatted_running} - batch {formatted_current}"), + }; + + // Serialize the state for aggregation + let serialized = NumericEntry::Aggregated { + aggregated_value: epoch_perplexity, + count: self.total_tokens, + } + .serialize(); + + SerializedEntry::new(formatted, serialized) + } + + fn value(&self) -> NumericEntry { + let perplexity = if self.total_tokens > 0 { + (self.sum_nll / self.total_tokens as f64).exp() + } else { + f64::INFINITY + }; + + NumericEntry::Aggregated { + aggregated_value: perplexity, + count: self.total_tokens, + } + } + + fn running_value(&self) -> NumericEntry { + self.value() + } +} + +/// The perplexity metric. +/// +/// Perplexity is a measure of how well a probability distribution or probability model +/// predicts a sample. It's commonly used to evaluate language models. A lower perplexity +/// indicates that the model is more confident in its predictions. +/// +/// Mathematically, perplexity is defined as the exponentiation of the cross-entropy loss: +/// PPL = exp(H(p, q)) = exp(-1/N * Σ log(p(x_i))) +/// +/// where: +/// - H(p, q) is the cross-entropy between the true distribution p and predicted distribution q +/// - N is the number of tokens +/// - p(x_i) is the predicted probability of the i-th token +/// +/// # Aggregation +/// Unlike other metrics, perplexity cannot be simply averaged across batches. +/// This implementation correctly accumulates the total negative log-likelihood and +/// total token count across batches, then computes perplexity as exp(total_nll / total_tokens). +#[derive(Clone)] +pub struct PerplexityMetric { + name: MetricName, + state: PerplexityState, + pad_token: Option, +} + +/// The [perplexity metric](PerplexityMetric) input type. +#[derive(new)] +pub struct PerplexityInput { + /// Logits tensor of shape [batch_size * sequence_length, vocab_size] + outputs: Tensor<2>, + /// Target tokens tensor of shape [batch_size * sequence_length] + targets: Tensor<1, Int>, +} + +impl Default for PerplexityMetric { + fn default() -> Self { + Self::new() + } +} + +impl PerplexityMetric { + /// Creates the metric. + pub fn new() -> Self { + Self { + name: MetricName::new("Perplexity".to_string()), + state: PerplexityState::new(), + pad_token: Default::default(), + } + } + + /// Sets the pad token to exclude from perplexity calculation. + /// + /// When a pad token is set, predictions for padding tokens are masked out + /// and do not contribute to the perplexity calculation. This is important + /// for variable-length sequences where padding is used. + pub fn with_pad_token(mut self, index: usize) -> Self { + self.pad_token = Some(index); + self + } +} + +impl Metric for PerplexityMetric { + type Input = PerplexityInput; + + fn update(&mut self, input: &PerplexityInput, _metadata: &MetricMetadata) -> SerializedEntry { + let targets = input.targets.clone(); + let outputs = input.outputs.clone(); + + let [total_tokens, _vocab_size] = outputs.dims(); + + // Convert logits to log probabilities using log_softmax for numerical stability + let log_probs = burn_core::tensor::activation::log_softmax(outputs, 1); + + // Gather the log probabilities for the target tokens + let target_log_probs = log_probs + .gather(1, targets.clone().unsqueeze_dim(1)) + .squeeze_dim(1); + + let (sum_log_prob, effective_tokens) = match self.pad_token { + Some(pad_token) => { + // Create a mask for non-padding tokens + let mask = targets.clone().not_equal_elem(pad_token as i64); + + // Apply mask to log probabilities (set padding log probs to 0) + let masked_log_probs = target_log_probs.mask_fill(mask.clone().bool_not(), 0.0); + + // Sum the log probabilities and count effective tokens + let sum_log_prob = masked_log_probs.sum().into_scalar::(); + let effective_tokens = mask.int().sum().into_scalar::() as usize; + + (sum_log_prob, effective_tokens) + } + None => { + // No padding, use all tokens + let sum_log_prob = target_log_probs.sum().into_scalar::(); + (sum_log_prob, total_tokens) + } + }; + + // Pass the sum_log_prob and effective_tokens to the state + // The state will handle the correct accumulation and perplexity calculation + self.state.update( + sum_log_prob, + effective_tokens, + FormatOptions::new(self.name()).precision(2), + ) + } + + fn clear(&mut self) { + self.state.reset() + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: None, + higher_is_better: false, + ..Default::default() + } + .into() + } +} + +impl Numeric for PerplexityMetric { + fn value(&self) -> NumericEntry { + self.state.value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_perplexity_perfect_prediction() { + let device = Default::default(); + let mut metric = PerplexityMetric::new(); + + // Perfect prediction: target is always the highest probability class + let input = PerplexityInput::new( + Tensor::from_data( + [ + [10.0, 0.0, 0.0], // Very confident prediction for class 0 + [0.0, 10.0, 0.0], // Very confident prediction for class 1 + [0.0, 0.0, 10.0], // Very confident prediction for class 2 + ], + &device, + ), + Tensor::from_data([0, 1, 2], &device), + ); + + let _entry = metric.update(&input, &MetricMetadata::fake()); + let perplexity = metric.value().current(); + + // Perfect predictions should result in very low perplexity (close to 1.0) + assert!( + perplexity < 1.1, + "Perfect predictions should have low perplexity, got {}", + perplexity + ); + } + + #[test] + fn test_perplexity_uniform_prediction() { + let device = Default::default(); + let mut metric = PerplexityMetric::new(); + + // Uniform prediction: all classes have equal probability + let input = PerplexityInput::new( + Tensor::from_data( + [ + [0.0, 0.0, 0.0], // Uniform distribution (after softmax) + [0.0, 0.0, 0.0], // Uniform distribution (after softmax) + [0.0, 0.0, 0.0], // Uniform distribution (after softmax) + ], + &device, + ), + Tensor::from_data([0, 1, 2], &device), + ); + + let _entry = metric.update(&input, &MetricMetadata::fake()); + let perplexity = metric.value().current(); + + // Uniform distribution over 3 classes should have perplexity ≈ 3.0 + assert!( + (perplexity - 3.0).abs() < 0.1, + "Uniform distribution perplexity should be ~3.0, got {}", + perplexity + ); + } + + #[test] + fn test_perplexity_with_padding() { + let device = Default::default(); + let mut metric = PerplexityMetric::new().with_pad_token(3); + + let input = PerplexityInput::new( + Tensor::from_data( + [ + [10.0, 0.0, 0.0, 0.0], // Good prediction for class 0 + [0.0, 10.0, 0.0, 0.0], // Good prediction for class 1 + [0.0, 0.0, 0.0, 1.0], // This is padding - should be ignored + [0.0, 0.0, 0.0, 1.0], // This is padding - should be ignored + ], + &device, + ), + Tensor::from_data([0, 1, 3, 3], &device), // 3 is pad token + ); + + let _entry = metric.update(&input, &MetricMetadata::fake()); + let perplexity = metric.value().current(); + + // Should only consider the first two predictions, both of which are confident + assert!( + perplexity < 1.1, + "Good predictions with padding should have low perplexity, got {}", + perplexity + ); + } + + #[test] + fn test_perplexity_wrong_prediction() { + let device = Default::default(); + let mut metric = PerplexityMetric::new(); + + // Wrong predictions: target class has very low probability + let input = PerplexityInput::new( + Tensor::from_data( + [ + [0.0, 10.0, 0.0], // Predicts class 1, but target is 0 + [10.0, 0.0, 0.0], // Predicts class 0, but target is 1 + [0.0, 0.0, 10.0], // Predicts class 2, but target is 0 + ], + &device, + ), + Tensor::from_data([0, 1, 0], &device), + ); + + let _entry = metric.update(&input, &MetricMetadata::fake()); + let perplexity = metric.value().current(); + + // Wrong predictions should result in high perplexity + assert!( + perplexity > 10.0, + "Wrong predictions should have high perplexity, got {}", + perplexity + ); + } + + #[test] + fn test_perplexity_multi_batch_aggregation() { + let device = Default::default(); + let mut metric = PerplexityMetric::new(); + + // First batch: 2 tokens with uniform distribution (log_prob ≈ -1.0986 each) + let input1 = PerplexityInput::new( + Tensor::from_data( + [ + [0.0, 0.0, 0.0], // Uniform distribution (log_prob ≈ -1.0986) + [0.0, 0.0, 0.0], // Uniform distribution (log_prob ≈ -1.0986) + ], + &device, + ), + Tensor::from_data([0, 1], &device), + ); + + // Second batch: 1 token with uniform distribution + let input2 = PerplexityInput::new( + Tensor::from_data( + [ + [0.0, 0.0, 0.0], // Uniform distribution (log_prob ≈ -1.0986) + ], + &device, + ), + Tensor::from_data([2], &device), + ); + + // Update with both batches + let _entry1 = metric.update(&input1, &MetricMetadata::fake()); + let _entry2 = metric.update(&input2, &MetricMetadata::fake()); + + let aggregated_perplexity = metric.value().current(); + + // For uniform distribution over 3 classes: log_prob ≈ -log(3) ≈ -1.0986 + // Total negative log-likelihood: 3 * 1.0986 ≈ 3.2958 + // Total tokens: 3 + // Expected perplexity: exp(3.2958 / 3) = exp(1.0986) ≈ 3.0 + assert!( + (aggregated_perplexity - 3.0).abs() < 0.1, + "Multi-batch aggregated perplexity should be ~3.0, got {}", + aggregated_perplexity + ); + + // Compare with single batch containing all data + let mut single_batch_metric = PerplexityMetric::new(); + let single_input = PerplexityInput::new( + Tensor::from_data([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], &device), + Tensor::from_data([0, 1, 2], &device), + ); + + let _single_entry = single_batch_metric.update(&single_input, &MetricMetadata::fake()); + let single_batch_perplexity = single_batch_metric.value().current(); + + // Multi-batch and single-batch should give the same result + assert!( + (aggregated_perplexity - single_batch_perplexity).abs() < 0.01, + "Multi-batch ({}) and single-batch ({}) perplexity should match", + aggregated_perplexity, + single_batch_perplexity + ); + } +} diff --git a/crates/burn-train/src/metric/precision.rs b/crates/burn-train/src/metric/precision.rs new file mode 100644 index 0000000..6eef8d4 --- /dev/null +++ b/crates/burn-train/src/metric/precision.rs @@ -0,0 +1,217 @@ +use crate::metric::{MetricName, Numeric}; + +use super::{ + Metric, MetricAttributes, MetricMetadata, NumericAttributes, NumericEntry, SerializedEntry, + classification::{ClassReduction, ClassificationMetricConfig, DecisionRule}, + confusion_stats::{ConfusionStats, ConfusionStatsInput}, + state::{FormatOptions, NumericMetricState}, +}; +use burn_core::prelude::Tensor; +use std::{num::NonZeroUsize, sync::Arc}; + +/// The Precision Metric +#[derive(Clone)] +pub struct PrecisionMetric { + name: MetricName, + state: NumericMetricState, + config: ClassificationMetricConfig, +} + +impl Default for PrecisionMetric { + fn default() -> Self { + Self::new(Default::default()) + } +} + +impl PrecisionMetric { + fn new(config: ClassificationMetricConfig) -> Self { + let state = Default::default(); + let name = Arc::new(format!( + "Precision @ {:?} [{:?}]", + config.decision_rule, config.class_reduction + )); + + Self { + state, + config, + name, + } + } + /// Precision metric for binary classification. + /// + /// # Arguments + /// + /// * `threshold` - The threshold to transform a probability into a binary prediction. + #[allow(dead_code)] + pub fn binary(threshold: f64) -> Self { + Self::new(ClassificationMetricConfig { + decision_rule: DecisionRule::Threshold(threshold), + // binary classification results are the same independently of class_reduction + ..Default::default() + }) + } + + /// Precision metric for multiclass classification. + /// + /// # Arguments + /// + /// * `top_k` - The number of highest predictions considered to find the correct label (typically `1`). + /// * `class_reduction` - [Class reduction](ClassReduction) type. + #[allow(dead_code)] + pub fn multiclass(top_k: usize, class_reduction: ClassReduction) -> Self { + Self::new(ClassificationMetricConfig { + decision_rule: DecisionRule::TopK( + NonZeroUsize::new(top_k).expect("top_k must be non-zero"), + ), + class_reduction, + }) + } + + /// Precision metric for multi-label classification. + /// + /// # Arguments + /// + /// * `threshold` - The threshold to transform a probability into a binary value. + /// * `class_reduction` - [Class reduction](ClassReduction) type. + #[allow(dead_code)] + pub fn multilabel(threshold: f64, class_reduction: ClassReduction) -> Self { + Self { + config: ClassificationMetricConfig { + decision_rule: DecisionRule::Threshold(threshold), + class_reduction, + }, + ..Default::default() + } + } + + fn class_average(&self, mut aggregated_metric: Tensor<1>) -> f64 { + use ClassReduction::{Macro, Micro}; + let avg_tensor = match self.config.class_reduction { + Micro => aggregated_metric, + Macro => { + if aggregated_metric.clone().contains_nan().any().into_scalar() { + let nan_mask = aggregated_metric.clone().is_nan(); + aggregated_metric = aggregated_metric + .clone() + .select(0, nan_mask.bool_not().argwhere().squeeze_dim(1)) + } + aggregated_metric.mean() + } + }; + avg_tensor.into_scalar() + } +} + +impl Metric for PrecisionMetric { + type Input = ConfusionStatsInput; + + fn update(&mut self, input: &Self::Input, _metadata: &MetricMetadata) -> SerializedEntry { + let [sample_size, _] = input.predictions.dims(); + + let cf_stats = ConfusionStats::new(input, &self.config); + let metric = + self.class_average(cf_stats.clone().true_positive() / cf_stats.predicted_positive()); + + self.state.update( + 100.0 * metric, + sample_size, + FormatOptions::new(self.name()).unit("%").precision(2), + ) + } + + fn clear(&mut self) { + self.state.reset() + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: Some("%".to_string()), + higher_is_better: true, + ..Default::default() + } + .into() + } +} + +impl Numeric for PrecisionMetric { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} + +#[cfg(test)] +mod tests { + use super::{ + ClassReduction::{self, *}, + Metric, MetricMetadata, PrecisionMetric, + }; + use crate::metric::Numeric; + use crate::tests::{ClassificationType, THRESHOLD, dummy_classification_input}; + use burn_core::tensor::TensorData; + use burn_core::tensor::Tolerance; + use rstest::rstest; + + #[rstest] + #[case::binary(THRESHOLD, 0.5)] + fn test_binary_precision(#[case] threshold: f64, #[case] expected: f64) { + let input = dummy_classification_input(&ClassificationType::Binary).into(); + let mut metric = PrecisionMetric::binary(threshold); + let _entry = metric.update(&input, &MetricMetadata::fake()); + TensorData::from([metric.value().current()]) + .assert_approx_eq::(&TensorData::from([expected * 100.0]), Tolerance::default()) + } + + #[rstest] + #[case::multiclass_micro_k1(Micro, 1, 3.0/5.0)] + #[case::multiclass_micro_k2(Micro, 2, 4.0/10.0)] + #[case::multiclass_macro_k1(Macro, 1, (0.5 + 0.5 + 1.0)/3.0)] + #[case::multiclass_macro_k2(Macro, 2, (0.5 + 1.0/4.0 + 0.5)/3.0)] + fn test_multiclass_precision( + #[case] class_reduction: ClassReduction, + #[case] top_k: usize, + #[case] expected: f64, + ) { + let input = dummy_classification_input(&ClassificationType::Multiclass).into(); + let mut metric = PrecisionMetric::multiclass(top_k, class_reduction); + let _entry = metric.update(&input, &MetricMetadata::fake()); + TensorData::from([metric.value().current()]) + .assert_approx_eq::(&TensorData::from([expected * 100.0]), Tolerance::default()) + } + + #[rstest] + #[case::multilabel_micro(Micro, THRESHOLD, 5.0/8.0)] + #[case::multilabel_macro(Macro, THRESHOLD, (2.0/3.0 + 2.0/3.0 + 0.5)/3.0)] + fn test_multilabel_precision( + #[case] class_reduction: ClassReduction, + #[case] threshold: f64, + #[case] expected: f64, + ) { + let input = dummy_classification_input(&ClassificationType::Multilabel).into(); + let mut metric = PrecisionMetric::multilabel(threshold, class_reduction); + let _entry = metric.update(&input, &MetricMetadata::fake()); + TensorData::from([metric.value().current()]) + .assert_approx_eq::(&TensorData::from([expected * 100.0]), Tolerance::default()) + } + + #[test] + fn test_parameterized_unique_name() { + let metric_a = PrecisionMetric::multiclass(1, ClassReduction::Macro); + let metric_b = PrecisionMetric::multiclass(2, ClassReduction::Macro); + let metric_c = PrecisionMetric::multiclass(1, ClassReduction::Macro); + + assert_ne!(metric_a.name(), metric_b.name()); + assert_eq!(metric_a.name(), metric_c.name()); + + let metric_a = PrecisionMetric::binary(0.5); + let metric_b = PrecisionMetric::binary(0.75); + assert_ne!(metric_a.name(), metric_b.name()); + } +} diff --git a/crates/burn-train/src/metric/processor/async_wrapper.rs b/crates/burn-train/src/metric/processor/async_wrapper.rs new file mode 100644 index 0000000..4802c1c --- /dev/null +++ b/crates/burn-train/src/metric/processor/async_wrapper.rs @@ -0,0 +1,142 @@ +use crate::metric::processor::{EvaluatorEvent, EventProcessorEvaluation}; + +use super::EventProcessorTraining; +use async_channel::{Receiver, Sender}; + +/// Event processor for the training process. +pub struct AsyncProcessorTraining { + sender: Sender>, +} + +/// Event processor for the model evaluation. +pub struct AsyncProcessorEvaluation { + sender: Sender>, +} + +struct WorkerTraining> { + processor: P, + rec: Receiver>, +} + +struct WorkerEvaluation { + processor: P, + rec: Receiver>, +} + +impl + 'static> + WorkerTraining +{ + pub fn start(processor: P, rec: Receiver>) { + let mut worker = Self { processor, rec }; + std::thread::Builder::new() + .name("train-worker".into()) + .spawn(move || { + while let Ok(msg) = worker.rec.recv_blocking() { + match msg { + Message::Train(event) => worker.processor.process_train(event), + Message::Valid(event) => worker.processor.process_valid(event), + Message::Renderer(callback) => { + callback.send_blocking(worker.processor.renderer()).unwrap(); + return; + } + } + } + }) + .unwrap(); + } +} +impl WorkerEvaluation

{ + pub fn start(processor: P, rec: Receiver>) { + let mut worker = Self { processor, rec }; + + std::thread::Builder::new() + .name("evel-worker".into()) + .spawn(move || { + while let Ok(event) = worker.rec.recv_blocking() { + match event { + EvalMessage::Test(event) => worker.processor.process_test(event), + EvalMessage::Renderer(sender) => { + sender.send_blocking(worker.processor.renderer()).unwrap(); + return; + } + } + } + }) + .unwrap(); + } +} + +impl AsyncProcessorTraining { + /// Create an event processor for training. + pub fn new + 'static>(processor: P) -> Self { + let (sender, rec) = async_channel::bounded(1); + + WorkerTraining::start(processor, rec); + + Self { sender } + } +} + +impl AsyncProcessorEvaluation

{ + /// Create an event processor for model evaluation. + pub fn new(processor: P) -> Self { + let (sender, rec) = async_channel::bounded(1); + + WorkerEvaluation::start(processor, rec); + + Self { sender } + } +} + +enum Message { + Train(EventTrain), + Valid(EventValid), + Renderer(Sender>), +} + +enum EvalMessage { + Test(EvaluatorEvent), + Renderer(Sender>), +} + +impl EventProcessorTraining for AsyncProcessorTraining { + fn process_train(&mut self, event: ET) { + self.sender.send_blocking(Message::Train(event)).unwrap(); + } + + fn process_valid(&mut self, event: EV) { + self.sender.send_blocking(Message::Valid(event)).unwrap(); + } + + fn renderer(self) -> Box { + let (sender, rec) = async_channel::bounded(1); + self.sender + .send_blocking(Message::Renderer(sender)) + .unwrap(); + + match rec.recv_blocking() { + Ok(value) => value, + Err(err) => panic!("{err:?}"), + } + } +} + +impl EventProcessorEvaluation for AsyncProcessorEvaluation

{ + type ItemTest = P::ItemTest; + + fn process_test(&mut self, event: EvaluatorEvent) { + self.sender.send_blocking(EvalMessage::Test(event)).unwrap(); + } + + fn renderer(self) -> Box { + let (sender, rec) = async_channel::bounded(1); + self.sender + .send_blocking(EvalMessage::Renderer(sender)) + .unwrap(); + + match rec.recv_blocking() { + Ok(value) => value, + Err(err) => panic!("{err:?}"), + } + } +} diff --git a/crates/burn-train/src/metric/processor/base.rs b/crates/burn-train/src/metric/processor/base.rs new file mode 100644 index 0000000..66098da --- /dev/null +++ b/crates/burn-train/src/metric/processor/base.rs @@ -0,0 +1,134 @@ +use burn_core::data::dataloader::Progress; +use burn_optim::lr_scheduler::module_lr_scheduler::ModuleLearningRate; + +use crate::{ + LearnerSummary, + renderer::{EvaluationName, MetricsRenderer}, +}; + +/// Event happening during the training/validation process. +pub enum LearnerEvent { + /// Signal the start of the process (e.g., training start). + Start { + /// The total number of training epochs. + total_epochs: usize, + /// The starting epoch. + starting_epoch: usize, + }, + /// Signal that an item have been processed. + ProcessedItem(TrainingItem), + /// Signal the start of a split. + StartSplit { + /// The epoch number. + epoch_number: usize, + /// The total number of items to be processed during this split. + total_items: usize, + }, + /// Signal the end of a split, carrying the current epoch number. + EndSplit(usize), + /// Signal the end of a full epoch. + EndEpoch(usize), + /// Signal the end of the process (e.g., training end). + End(Option), +} + +/// Event happening during the evaluation process. +pub enum EvaluatorEvent { + /// Signal the start of the process (e.g., evaluation start) + Start { + /// The total number of items to evaluate. + total_tests: usize, + }, + /// Signal the start of a test split, carrying the split name and total number of items. + StartTest(EvaluationName, usize), + /// Signal that an item have been processed. + ProcessedItem(EvaluationName, EvaluationItem), + /// Signal the end of a single test split. + EndTest, + /// Signal the end of the process (e.g., evaluation end). + End(Option), +} + +/// Items that are lazy are not ready to be processed by metrics. +/// +/// We want to sync them on a different thread to avoid blocking training. +pub trait ItemLazy: Send { + /// Sync the item. + fn sync(self) -> Self; +} + +/// Process events happening during training and validation. +pub trait EventProcessorTraining: Send { + /// Collect a training event. + fn process_train(&mut self, event: TrainEvent); + /// Collect a validation event. + fn process_valid(&mut self, event: ValidEvent); + /// Returns the renderer used for training. + fn renderer(self) -> Box; +} + +/// Process events happening during evaluation. +pub trait EventProcessorEvaluation: Send { + /// The test item. + type ItemTest: ItemLazy; + + /// Collect a test event. + fn process_test(&mut self, event: EvaluatorEvent); + + /// Returns the renderer used for evaluation. + fn renderer(self) -> Box; +} + +/// A learner item. +#[derive(new)] +pub struct TrainingItem { + /// The item. + pub item: T, + + /// The progress. + pub progress: Progress, + + /// The iteration, if it it different from the items processed. + pub iteration: Option, + + /// The learning rate for a module's parameters. + pub lr: Option, +} + +impl ItemLazy for TrainingItem { + fn sync(self) -> Self { + TrainingItem { + item: self.item.sync(), + progress: self.progress, + iteration: self.iteration, + lr: self.lr, + } + } +} + +/// An evaluation item. +#[derive(new)] +pub struct EvaluationItem { + /// The item. + pub item: T, + + /// The progress. + pub progress: Progress, + + /// The iteration, if it it different from the items processed. + pub iteration: Option, +} + +impl ItemLazy for EvaluationItem { + fn sync(self) -> Self { + EvaluationItem { + item: self.item.sync(), + progress: self.progress, + iteration: self.iteration, + } + } +} + +impl ItemLazy for () { + fn sync(self) -> Self {} +} diff --git a/crates/burn-train/src/metric/processor/full.rs b/crates/burn-train/src/metric/processor/full.rs new file mode 100644 index 0000000..d259c18 --- /dev/null +++ b/crates/burn-train/src/metric/processor/full.rs @@ -0,0 +1,333 @@ +use super::{EventProcessorTraining, ItemLazy, LearnerEvent, MetricsTraining}; +use crate::logger::{EvaluationProgressLogger, TrainingProgressLogger}; +use crate::metric::MetricMetadata; +use crate::metric::processor::{EvaluatorEvent, EventProcessorEvaluation, MetricsEvaluation}; +use crate::metric::store::{EpochSummary, EventStoreClient, Split}; +use crate::renderer::{MetricState, MetricsRenderer}; +use std::sync::Arc; + +/// An [event processor](EventProcessorTraining) that handles: +/// - Computing and storing metrics in an [event store](crate::metric::store::EventStore). +/// - Render metrics using a [metrics renderer](MetricsRenderer). +pub struct FullEventProcessorTraining { + metrics: MetricsTraining, + renderer: Box, + store: Arc, + progress_logger: Option>, + current_epoch: usize, + total_epochs: usize, +} + +/// An [event processor](EventProcessorEvaluation) that handles: +/// - Computing and storing metrics in an [event store](crate::metric::store::EventStore). +/// - Render metrics using a [metrics renderer](MetricsRenderer). +pub struct FullEventProcessorEvaluation { + metrics: MetricsEvaluation, + renderer: Box, + store: Arc, + progress_logger: Option>, + total_tests: usize, + current_test: usize, +} + +impl FullEventProcessorTraining { + pub(crate) fn new( + metrics: MetricsTraining, + renderer: Box, + store: Arc, + ) -> Self { + Self { + metrics, + renderer, + store, + progress_logger: None, + current_epoch: 1, + total_epochs: 0, + } + } + + pub(crate) fn with_progress_logger(mut self, logger: Box) -> Self { + self.progress_logger = Some(logger); + self + } +} + +impl FullEventProcessorEvaluation { + pub(crate) fn new( + metrics: MetricsEvaluation, + renderer: Box, + store: Arc, + ) -> Self { + Self { + metrics, + renderer, + store, + progress_logger: None, + total_tests: 0, + current_test: 0, + } + } + + pub(crate) fn with_progress_logger( + mut self, + logger: Box, + ) -> Self { + self.progress_logger = Some(logger); + self + } +} + +impl EventProcessorEvaluation for FullEventProcessorEvaluation { + type ItemTest = T; + + fn process_test(&mut self, event: EvaluatorEvent) { + match event { + EvaluatorEvent::Start { total_tests } => { + let definitions = self.metrics.metric_definitions(); + self.store + .add_event_train(crate::metric::store::Event::MetricsInit( + definitions.clone(), + )); + definitions + .iter() + .for_each(|definition| self.renderer.register_metric(definition.clone())); + self.total_tests = total_tests; + self.current_test = 0; + if let Some(logger) = &mut self.progress_logger { + logger.start_global_progress(total_tests); + } + self.renderer.start_global_progress(total_tests); + } + EvaluatorEvent::StartTest(name, total_items) => { + self.current_test += 1; + self.renderer.start_test(name.as_str(), total_items); + if let Some(logger) = &mut self.progress_logger { + logger.start_test(name.as_str(), total_items); + } + } + EvaluatorEvent::ProcessedItem(name, item) => { + let item = item.sync(); + let metadata = (&item).into(); + + let update = self.metrics.update_test(&item, &metadata); + + self.store.add_event_test( + crate::metric::store::Event::MetricsUpdate(update.clone()), + name.name.clone(), + ); + + update.entries.into_iter().for_each(|entry| { + self.renderer + .update_test(name.clone(), MetricState::Generic(entry)) + }); + + update + .entries_numeric + .into_iter() + .for_each(|numeric_update| { + self.renderer.update_test( + name.clone(), + MetricState::Numeric( + numeric_update.entry, + numeric_update.numeric_entry, + ), + ) + }); + + if let Some(logger) = &mut self.progress_logger { + logger.update_test_progress(item.progress.items_processed); + logger.log_event_evaluation("Iteration".to_string()); + } + self.renderer + .update_test_progress(item.progress.items_processed); + self.renderer.log_event_evaluation("Iteration".to_string()); + } + EvaluatorEvent::EndTest => { + if let Some(logger) = &mut self.progress_logger { + logger.end_test(); + } + self.renderer.end_test(); + } + EvaluatorEvent::End(summary) => { + if let Some(logger) = &mut self.progress_logger { + logger.end_global_progress(); + } + self.renderer.end_global_progress(); + self.renderer.on_test_end(summary).ok(); + } + } + } + + fn renderer(self) -> Box { + self.renderer + } +} + +impl EventProcessorTraining, LearnerEvent> + for FullEventProcessorTraining +{ + fn process_train(&mut self, event: LearnerEvent) { + match event { + LearnerEvent::Start { + total_epochs, + starting_epoch, + } => { + self.total_epochs = total_epochs; + self.current_epoch = 1; + let definitions = self.metrics.metric_definitions(); + self.store + .add_event_train(crate::metric::store::Event::MetricsInit( + definitions.clone(), + )); + definitions + .iter() + .for_each(|definition| self.renderer.register_metric(definition.clone())); + if let Some(logger) = &mut self.progress_logger { + logger.start(total_epochs, starting_epoch, None); + } + self.renderer.start(total_epochs, starting_epoch, None); + } + LearnerEvent::StartSplit { + epoch_number, + total_items, + } => { + self.store + .add_event_train(crate::metric::store::Event::StartSplit(epoch_number)); + self.renderer.start_split(Split::Train.into(), total_items); + if let Some(logger) = &mut self.progress_logger { + logger.start_split(Split::Train.into(), total_items); + } + } + LearnerEvent::ProcessedItem(item) => { + let item = item.sync(); + let metadata = MetricMetadata { + progress: item.progress.clone(), + iteration: item.iteration, + lr: item.lr.clone(), + }; + + let update = self.metrics.update_train(&item, &metadata); + + self.store + .add_event_train(crate::metric::store::Event::MetricsUpdate(update.clone())); + + update + .entries + .into_iter() + .for_each(|entry| self.renderer.update_train(MetricState::Generic(entry))); + + update + .entries_numeric + .into_iter() + .for_each(|numeric_update| { + self.renderer.update_train(MetricState::Numeric( + numeric_update.entry, + numeric_update.numeric_entry, + )) + }); + + if let Some(logger) = &mut self.progress_logger { + logger.update_split(item.progress.items_processed); + logger.log_event_training("Iteration".to_string()); + } + self.renderer.update_split(item.progress.items_processed); + self.renderer.log_event_training("Iteration".to_string()); + } + LearnerEvent::EndSplit(epoch) => { + self.store + .add_event_train(crate::metric::store::Event::EndEpoch(EpochSummary::new( + epoch, + Split::Train, + ))); + if let Some(logger) = &mut self.progress_logger { + logger.end_split(); + } + self.renderer.end_split(); + self.metrics.end_epoch_train(); + } + LearnerEvent::EndEpoch(epoch) => { + self.current_epoch = epoch + 1; + if let Some(logger) = &mut self.progress_logger { + logger.update_epoch(epoch); + } + self.renderer.update_epoch(epoch) + } + LearnerEvent::End(summary) => { + if let Some(logger) = &mut self.progress_logger { + logger.end(); + } + self.renderer.end(); + self.renderer.on_train_end(summary).ok(); + } + } + } + + fn process_valid(&mut self, event: LearnerEvent) { + match event { + LearnerEvent::Start { .. } => {} // no-op: valid has no separate start event + LearnerEvent::StartSplit { + epoch_number, + total_items, + } => { + self.store + .add_event_valid(crate::metric::store::Event::StartSplit(epoch_number)); + if let Some(logger) = &mut self.progress_logger { + logger.start_split(Split::Valid.into(), total_items); + } + self.renderer.start_split(Split::Valid.into(), total_items); + } + LearnerEvent::ProcessedItem(item) => { + let item = item.sync(); + let metadata = MetricMetadata { + progress: item.progress.clone(), + iteration: item.iteration, + lr: item.lr.clone(), + }; + + let update = self.metrics.update_valid(&item, &metadata); + + self.store + .add_event_valid(crate::metric::store::Event::MetricsUpdate(update.clone())); + + update + .entries + .into_iter() + .for_each(|entry| self.renderer.update_valid(MetricState::Generic(entry))); + + update + .entries_numeric + .into_iter() + .for_each(|numeric_update| { + self.renderer.update_valid(MetricState::Numeric( + numeric_update.entry, + numeric_update.numeric_entry, + )) + }); + + if let Some(logger) = &mut self.progress_logger { + logger.update_split(item.progress.items_processed); + logger.log_event_training("Iteration".to_string()); + } + self.renderer.update_split(item.progress.items_processed); + self.renderer.log_event_training("Iteration".to_string()); + } + LearnerEvent::EndSplit(epoch) => { + self.store + .add_event_valid(crate::metric::store::Event::EndEpoch(EpochSummary::new( + epoch, + Split::Valid, + ))); + if let Some(logger) = &mut self.progress_logger { + logger.end_split(); + } + self.renderer.end_split(); + self.metrics.end_epoch_valid(); + } + LearnerEvent::EndEpoch(_) => {} // update_epoch is handled in process_train(EndEpoch) + LearnerEvent::End(_) => {} // no-op + } + } + fn renderer(self) -> Box { + self.renderer + } +} diff --git a/crates/burn-train/src/metric/processor/metrics.rs b/crates/burn-train/src/metric/processor/metrics.rs new file mode 100644 index 0000000..2f4fe57 --- /dev/null +++ b/crates/burn-train/src/metric/processor/metrics.rs @@ -0,0 +1,309 @@ +use std::collections::HashMap; + +use super::{ItemLazy, TrainingItem}; +use crate::{ + EvaluationItem, + metric::{ + Adaptor, Metric, MetricDefinition, MetricEntry, MetricId, MetricMetadata, Numeric, + store::{MetricsUpdate, NumericMetricUpdate}, + }, +}; + +pub(crate) struct MetricsTraining { + train: Vec>>, + valid: Vec>>, + train_numeric: Vec>>, + valid_numeric: Vec>>, + metric_definitions: HashMap, +} + +pub(crate) struct MetricsEvaluation { + test: Vec>>, + test_numeric: Vec>>, + metric_definitions: HashMap, +} + +impl Default for MetricsEvaluation { + fn default() -> Self { + Self { + test: Default::default(), + test_numeric: Default::default(), + metric_definitions: HashMap::default(), + } + } +} + +impl Default for MetricsTraining { + fn default() -> Self { + Self { + train: Vec::default(), + valid: Vec::default(), + train_numeric: Vec::default(), + valid_numeric: Vec::default(), + metric_definitions: HashMap::default(), + } + } +} + +impl MetricsEvaluation { + /// Register a testing metric. + pub(crate) fn register_test_metric(&mut self, metric: Me) + where + T: Adaptor + 'static, + { + let metric = MetricWrapper::new(metric); + self.register_definition(&metric); + self.test.push(Box::new(metric)) + } + + /// Register a numeric testing metric. + pub(crate) fn register_test_metric_numeric( + &mut self, + metric: Me, + ) where + T: Adaptor + 'static, + { + let metric = MetricWrapper::new(metric); + self.register_definition(&metric); + self.test_numeric.push(Box::new(metric)) + } + + fn register_definition(&mut self, metric: &MetricWrapper) { + self.metric_definitions.insert( + metric.id.clone(), + MetricDefinition::new(metric.id.clone(), &metric.metric), + ); + } + + /// Get metric definitions. + pub(crate) fn metric_definitions(&mut self) -> Vec { + self.metric_definitions.values().cloned().collect() + } + + /// Update the testing information from the testing item. + pub(crate) fn update_test( + &mut self, + item: &EvaluationItem, + metadata: &MetricMetadata, + ) -> MetricsUpdate { + let mut entries = Vec::with_capacity(self.test.len()); + let mut entries_numeric = Vec::with_capacity(self.test_numeric.len()); + + for metric in self.test.iter_mut() { + let state = metric.update(&item.item, metadata); + entries.push(state); + } + + for metric in self.test_numeric.iter_mut() { + let numeric_update = metric.update(&item.item, metadata); + entries_numeric.push(numeric_update); + } + + MetricsUpdate::new(entries, entries_numeric) + } +} + +impl MetricsTraining { + /// Register a training metric. + pub(crate) fn register_train_metric(&mut self, metric: Me) + where + T: Adaptor + 'static, + { + let metric = MetricWrapper::new(metric); + self.register_definition(&metric); + self.train.push(Box::new(metric)) + } + + /// Register a validation metric. + pub(crate) fn register_valid_metric(&mut self, metric: Me) + where + V: Adaptor + 'static, + { + let metric = MetricWrapper::new(metric); + self.register_definition(&metric); + self.valid.push(Box::new(metric)) + } + + /// Register a numeric training metric. + pub(crate) fn register_train_metric_numeric( + &mut self, + metric: Me, + ) where + T: Adaptor + 'static, + { + let metric = MetricWrapper::new(metric); + self.register_definition(&metric); + self.train_numeric.push(Box::new(metric)) + } + + /// Register a numeric validation metric. + pub(crate) fn register_valid_metric_numeric(&mut self, metric: Me) + where + V: Adaptor + 'static, + Me: Metric + Numeric + 'static, + { + let metric = MetricWrapper::new(metric); + self.register_definition(&metric); + self.valid_numeric.push(Box::new(metric)) + } + + fn register_definition(&mut self, metric: &MetricWrapper) { + self.metric_definitions.insert( + metric.id.clone(), + MetricDefinition::new(metric.id.clone(), &metric.metric), + ); + } + + /// Get metric definitions for all splits + pub(crate) fn metric_definitions(&mut self) -> Vec { + self.metric_definitions.values().cloned().collect() + } + + /// Update the training information from the training item. + pub(crate) fn update_train( + &mut self, + item: &TrainingItem, + metadata: &MetricMetadata, + ) -> MetricsUpdate { + let mut entries = Vec::with_capacity(self.train.len()); + let mut entries_numeric = Vec::with_capacity(self.train_numeric.len()); + + for metric in self.train.iter_mut() { + let state = metric.update(&item.item, metadata); + entries.push(state); + } + + for metric in self.train_numeric.iter_mut() { + let numeric_update = metric.update(&item.item, metadata); + entries_numeric.push(numeric_update); + } + + MetricsUpdate::new(entries, entries_numeric) + } + + /// Update the training information from the validation item. + pub(crate) fn update_valid( + &mut self, + item: &TrainingItem, + metadata: &MetricMetadata, + ) -> MetricsUpdate { + let mut entries = Vec::with_capacity(self.valid.len()); + let mut entries_numeric = Vec::with_capacity(self.valid_numeric.len()); + + for metric in self.valid.iter_mut() { + let state = metric.update(&item.item, metadata); + entries.push(state); + } + + for metric in self.valid_numeric.iter_mut() { + let numeric_update = metric.update(&item.item, metadata); + entries_numeric.push(numeric_update); + } + + MetricsUpdate::new(entries, entries_numeric) + } + + /// Signal the end of a training epoch. + pub(crate) fn end_epoch_train(&mut self) { + for metric in self.train.iter_mut() { + metric.clear(); + } + for metric in self.train_numeric.iter_mut() { + metric.clear(); + } + } + + /// Signal the end of a validation epoch. + pub(crate) fn end_epoch_valid(&mut self) { + for metric in self.valid.iter_mut() { + metric.clear(); + } + for metric in self.valid_numeric.iter_mut() { + metric.clear(); + } + } +} + +impl From<&TrainingItem> for MetricMetadata { + fn from(item: &TrainingItem) -> Self { + Self { + progress: item.progress.clone(), + iteration: item.iteration, + lr: item.lr.clone(), + } + } +} + +impl From<&EvaluationItem> for MetricMetadata { + fn from(item: &EvaluationItem) -> Self { + Self { + progress: item.progress.clone(), + iteration: item.iteration, + lr: None, + } + } +} + +pub(crate) trait NumericMetricUpdater: Send + Sync { + fn update(&mut self, item: &T, metadata: &MetricMetadata) -> NumericMetricUpdate; + fn clear(&mut self); +} + +pub(crate) trait MetricUpdater: Send + Sync { + fn update(&mut self, item: &T, metadata: &MetricMetadata) -> MetricEntry; + fn clear(&mut self); +} + +pub(crate) struct MetricWrapper { + pub id: MetricId, + pub metric: M, +} + +impl MetricWrapper { + pub fn new(metric: M) -> Self { + Self { + id: MetricId::new(metric.name()), + metric, + } + } +} + +impl NumericMetricUpdater for MetricWrapper +where + T: 'static, + M: Metric + Numeric + 'static, + T: Adaptor, +{ + fn update(&mut self, item: &T, metadata: &MetricMetadata) -> NumericMetricUpdate { + let serialized_entry = self.metric.update(&item.adapt(), metadata); + let update = MetricEntry::new(self.id.clone(), serialized_entry); + let numeric = self.metric.value(); + let running = self.metric.running_value(); + + NumericMetricUpdate { + entry: update, + numeric_entry: numeric, + running_entry: running, + } + } + + fn clear(&mut self) { + self.metric.clear() + } +} + +impl MetricUpdater for MetricWrapper +where + T: 'static, + M: Metric + 'static, + T: Adaptor, +{ + fn update(&mut self, item: &T, metadata: &MetricMetadata) -> MetricEntry { + let serialized_entry = self.metric.update(&item.adapt(), metadata); + MetricEntry::new(self.id.clone(), serialized_entry) + } + + fn clear(&mut self) { + self.metric.clear() + } +} diff --git a/crates/burn-train/src/metric/processor/minimal.rs b/crates/burn-train/src/metric/processor/minimal.rs new file mode 100644 index 0000000..5060d89 --- /dev/null +++ b/crates/burn-train/src/metric/processor/minimal.rs @@ -0,0 +1,141 @@ +use super::{EventProcessorTraining, ItemLazy, LearnerEvent, MetricsTraining}; +use crate::{ + logger::TrainingProgressLogger, + metric::store::{EpochSummary, EventStoreClient, Split}, + renderer::cli::CliMetricsRenderer, +}; +use std::sync::Arc; + +/// An [event processor](EventProcessor) that handles: +/// - Computing and storing metrics in an [event store](crate::metric::store::EventStore). +/// - Optionally logging training progress via a [TrainingProgressLogger]. +#[allow(dead_code)] +pub(crate) struct MinimalEventProcessor { + metrics: MetricsTraining, + store: Arc, + progress_logger: Option>, +} + +#[allow(dead_code)] +impl MinimalEventProcessor { + pub(crate) fn new(metrics: MetricsTraining, store: Arc) -> Self { + Self { + metrics, + store, + progress_logger: None, + } + } + + pub(crate) fn with_progress_logger(mut self, logger: Box) -> Self { + self.progress_logger = Some(logger); + self + } +} + +impl EventProcessorTraining, LearnerEvent> + for MinimalEventProcessor +{ + fn process_train(&mut self, event: LearnerEvent) { + match event { + LearnerEvent::Start { + total_epochs, + starting_epoch, + } => { + let definitions = self.metrics.metric_definitions(); + self.store + .add_event_train(crate::metric::store::Event::MetricsInit(definitions)); + if let Some(logger) = &mut self.progress_logger { + logger.start(total_epochs, starting_epoch, None); + } + } + LearnerEvent::StartSplit { + epoch_number, + total_items, + } => { + self.store + .add_event_train(crate::metric::store::Event::StartSplit(epoch_number)); + if let Some(logger) = &mut self.progress_logger { + logger.start_split(Split::Train.into(), total_items); + } + } + LearnerEvent::ProcessedItem(item) => { + let item = item.sync(); + let metadata = (&item).into(); + + let update = self.metrics.update_train(&item, &metadata); + + self.store + .add_event_train(crate::metric::store::Event::MetricsUpdate(update)); + if let Some(logger) = &mut self.progress_logger { + logger.update_split(item.progress.items_processed); + } + } + LearnerEvent::EndSplit(epoch) => { + self.metrics.end_epoch_train(); + self.store + .add_event_train(crate::metric::store::Event::EndEpoch(EpochSummary::new( + epoch, + Split::Train, + ))); + if let Some(logger) = &mut self.progress_logger { + logger.end_split(); + } + } + LearnerEvent::EndEpoch(epoch) => { + if let Some(logger) = &mut self.progress_logger { + logger.update_epoch(epoch); + } + } + LearnerEvent::End(_summary) => { + if let Some(logger) = &mut self.progress_logger { + logger.end(); + } + } + } + } + + fn process_valid(&mut self, event: LearnerEvent) { + match event { + LearnerEvent::Start { .. } => {} // no-op + LearnerEvent::StartSplit { + epoch_number, + total_items, + } => { + self.store + .add_event_valid(crate::metric::store::Event::StartSplit(epoch_number)); + if let Some(logger) = &mut self.progress_logger { + logger.start_split(Split::Valid.into(), total_items); + } + } + LearnerEvent::ProcessedItem(item) => { + let item = item.sync(); + let metadata = (&item).into(); + + let update = self.metrics.update_valid(&item, &metadata); + + self.store + .add_event_valid(crate::metric::store::Event::MetricsUpdate(update)); + if let Some(logger) = &mut self.progress_logger { + logger.update_split(item.progress.items_processed); + } + } + LearnerEvent::EndSplit(epoch) => { + self.metrics.end_epoch_valid(); + self.store + .add_event_valid(crate::metric::store::Event::EndEpoch(EpochSummary::new( + epoch, + Split::Valid, + ))); + if let Some(logger) = &mut self.progress_logger { + logger.end_split(); + } + } + LearnerEvent::EndEpoch(_) => {} // update_epoch handled in process_train(EndEpoch) + LearnerEvent::End(_) => {} // no-op: End is only emitted on process_train + } + } + fn renderer(self) -> Box { + // TODO: Check for another default. + Box::new(CliMetricsRenderer::new()) + } +} diff --git a/crates/burn-train/src/metric/processor/mod.rs b/crates/burn-train/src/metric/processor/mod.rs new file mode 100644 index 0000000..860aa74 --- /dev/null +++ b/crates/burn-train/src/metric/processor/mod.rs @@ -0,0 +1,85 @@ +mod async_wrapper; +mod base; +mod full; +mod metrics; +mod minimal; +#[cfg(feature = "rl")] +mod rl_metrics; +#[cfg(feature = "rl")] +mod rl_processor; + +pub use base::*; +pub(crate) use full::*; +pub(crate) use metrics::*; +#[cfg(feature = "rl")] +pub(crate) use rl_metrics::*; +#[cfg(feature = "rl")] +pub use rl_processor::*; + +#[cfg(test)] +pub(crate) use minimal::*; + +pub use async_wrapper::{AsyncProcessorEvaluation, AsyncProcessorTraining}; + +#[cfg(test)] +pub(crate) mod test_utils { + use crate::metric::{ + Adaptor, LossInput, + processor::{EventProcessorTraining, LearnerEvent, MinimalEventProcessor, TrainingItem}, + }; + use burn_core::tensor::Tensor; + + use super::ItemLazy; + + impl ItemLazy for f64 { + fn sync(self) -> Self { + self + } + } + + impl Adaptor for f64 { + fn adapt(&self) -> LossInput { + LossInput::new(Tensor::from_data([*self], &Default::default())) + } + } + + pub(crate) fn process_train( + processor: &mut MinimalEventProcessor, + value: f64, + epoch: usize, + ) { + let dummy_progress = burn_core::data::dataloader::Progress { + items_processed: epoch, + items_total: 3, + unit: Some("items".to_string()), + }; + let dummy_iteration = Some(1); + + processor.process_train(LearnerEvent::ProcessedItem(TrainingItem::new( + value, + dummy_progress, + dummy_iteration, + None, + ))); + } + + pub(crate) fn start_epoch( + processor: &mut MinimalEventProcessor, + epoch: usize, + num_items: usize, + ) { + processor.process_train(LearnerEvent::StartSplit { + epoch_number: epoch, + total_items: num_items, + }); + processor.process_valid(LearnerEvent::StartSplit { + epoch_number: epoch, + total_items: num_items, + }); + } + + pub(crate) fn end_epoch(processor: &mut MinimalEventProcessor, epoch: usize) { + processor.process_train(LearnerEvent::EndSplit(epoch)); + processor.process_valid(LearnerEvent::EndSplit(epoch)); + } +} diff --git a/crates/burn-train/src/metric/processor/rl_metrics.rs b/crates/burn-train/src/metric/processor/rl_metrics.rs new file mode 100644 index 0000000..c3c3c3e --- /dev/null +++ b/crates/burn-train/src/metric/processor/rl_metrics.rs @@ -0,0 +1,268 @@ +use std::collections::HashMap; + +use crate::{ + EpisodeSummary, EvaluationItem, ItemLazy, MetricUpdater, MetricWrapper, NumericMetricUpdater, + metric::{ + Adaptor, Metric, MetricDefinition, MetricId, MetricMetadata, Numeric, store::MetricsUpdate, + }, +}; + +pub(crate) struct RLMetrics { + train_step: Vec>>, + env_step: Vec>>, + env_step_valid: Vec>>, + episode_end: Vec>>, + episode_end_valid: Vec>>, + + train_step_numeric: Vec>>, + env_step_numeric: Vec>>, + env_step_valid_numeric: Vec>>, + episode_end_numeric: Vec>>, + episode_end_valid_numeric: Vec>>, + + metric_definitions: HashMap, +} + +impl Default for RLMetrics { + fn default() -> Self { + Self { + train_step: Vec::default(), + env_step: Vec::default(), + env_step_valid: Vec::default(), + episode_end: Vec::default(), + episode_end_valid: Vec::default(), + train_step_numeric: Vec::default(), + env_step_numeric: Vec::default(), + env_step_valid_numeric: Vec::default(), + episode_end_numeric: Vec::default(), + episode_end_valid_numeric: Vec::default(), + metric_definitions: HashMap::default(), + } + } +} + +impl RLMetrics { + /// Register a training metric. + pub(crate) fn register_text_metric_agent(&mut self, metric: Me) + where + ES: Adaptor + 'static, + { + let metric = MetricWrapper::new(metric); + self.register_definition(&metric); + self.env_step.push(Box::new(metric)) + } + + /// Register a training metric. + pub(crate) fn register_agent_metric(&mut self, metric: Me) + where + ES: Adaptor + 'static, + { + let metric = MetricWrapper::new(metric); + self.register_definition(&metric); + self.env_step_numeric.push(Box::new(metric)) + } + + /// Register a training metric. + pub(crate) fn register_text_metric_train(&mut self, metric: Me) + where + TS: Adaptor + 'static, + { + let metric = MetricWrapper::new(metric); + self.register_definition(&metric); + self.train_step.push(Box::new(metric)) + } + + /// Register a training metric. + pub(crate) fn register_metric_train(&mut self, metric: Me) + where + TS: Adaptor + 'static, + { + let metric = MetricWrapper::new(metric); + self.register_definition(&metric); + self.train_step_numeric.push(Box::new(metric)) + } + + /// Register a validation env-step metric. + pub(crate) fn register_text_metric_agent_valid(&mut self, metric: Me) + where + ES: Adaptor + 'static, + { + let metric = MetricWrapper::new(metric); + self.register_definition(&metric); + self.env_step_valid.push(Box::new(metric)) + } + + /// Register a validation env-step numeric metric. + pub(crate) fn register_agent_metric_valid(&mut self, metric: Me) + where + ES: Adaptor + 'static, + { + let metric = MetricWrapper::new(metric); + self.register_definition(&metric); + self.env_step_valid_numeric.push(Box::new(metric)) + } + + /// Register an episode-end metric. + pub(crate) fn register_text_metric_episode(&mut self, metric: Me) + where + EpisodeSummary: Adaptor + 'static, + { + let metric = MetricWrapper::new(metric); + self.register_definition(&metric); + self.episode_end.push(Box::new(metric)) + } + + /// Register an episode-end numeric metric. + pub(crate) fn register_episode_metric(&mut self, metric: Me) + where + EpisodeSummary: Adaptor + 'static, + { + let metric = MetricWrapper::new(metric); + self.register_definition(&metric); + self.episode_end_numeric.push(Box::new(metric)) + } + + /// Register an episode-end metric for validation. + pub(crate) fn register_text_metric_episode_valid(&mut self, metric: Me) + where + EpisodeSummary: Adaptor + 'static, + { + let metric = MetricWrapper::new(metric); + self.register_definition(&metric); + self.episode_end_valid.push(Box::new(metric)) + } + + /// Register an episode-end numeric metric for validation. + pub(crate) fn register_episode_metric_valid( + &mut self, + metric: Me, + ) where + EpisodeSummary: Adaptor + 'static, + { + let metric = MetricWrapper::new(metric); + self.register_definition(&metric); + self.episode_end_valid_numeric.push(Box::new(metric)) + } + + fn register_definition(&mut self, metric: &MetricWrapper) { + self.metric_definitions.insert( + metric.id.clone(), + MetricDefinition::new(metric.id.clone(), &metric.metric), + ); + } + + /// Get metric definitions for all splits + pub(crate) fn metric_definitions(&mut self) -> Vec { + self.metric_definitions.values().cloned().collect() + } + + /// Update the training information from the training item. + pub(crate) fn update_train_step( + &mut self, + item: &EvaluationItem, + metadata: &MetricMetadata, + ) -> MetricsUpdate { + let mut entries = Vec::with_capacity(self.train_step.len()); + let mut entries_numeric = Vec::with_capacity(self.train_step_numeric.len()); + + for metric in self.train_step.iter_mut() { + let state = metric.update(&item.item, metadata); + entries.push(state); + } + + for metric in self.train_step_numeric.iter_mut() { + let numeric_update = metric.update(&item.item, metadata); + entries_numeric.push(numeric_update); + } + + MetricsUpdate::new(entries, entries_numeric) + } + + /// Update the env-step metrics from an environment step item. + pub(crate) fn update_env_step( + &mut self, + item: &EvaluationItem, + metadata: &MetricMetadata, + ) -> MetricsUpdate { + let mut entries = Vec::with_capacity(self.env_step.len()); + let mut entries_numeric = Vec::with_capacity(self.env_step_numeric.len()); + + for metric in self.env_step.iter_mut() { + let state = metric.update(&item.item, metadata); + entries.push(state); + } + + for metric in self.env_step_numeric.iter_mut() { + let numeric_update = metric.update(&item.item, metadata); + entries_numeric.push(numeric_update); + } + + MetricsUpdate::new(entries, entries_numeric) + } + + /// Update the env-step metrics for validation from an environment step item. + pub(crate) fn update_env_step_valid( + &mut self, + item: &EvaluationItem, + metadata: &MetricMetadata, + ) -> MetricsUpdate { + let mut entries = Vec::with_capacity(self.env_step_valid.len()); + let mut entries_numeric = Vec::with_capacity(self.env_step_valid_numeric.len()); + + for metric in self.env_step_valid.iter_mut() { + let state = metric.update(&item.item, metadata); + entries.push(state); + } + + for metric in self.env_step_valid_numeric.iter_mut() { + let numeric_update = metric.update(&item.item, metadata); + entries_numeric.push(numeric_update); + } + + MetricsUpdate::new(entries, entries_numeric) + } + + /// Update the episode-end metrics from an episode summary. + pub(crate) fn update_episode_end( + &mut self, + item: &EvaluationItem, + metadata: &MetricMetadata, + ) -> MetricsUpdate { + let mut entries = Vec::with_capacity(self.episode_end.len()); + let mut entries_numeric = Vec::with_capacity(self.episode_end_numeric.len()); + + for metric in self.episode_end.iter_mut() { + let state = metric.update(&item.item, metadata); + entries.push(state); + } + + for metric in self.episode_end_numeric.iter_mut() { + let numeric_update = metric.update(&item.item, metadata); + entries_numeric.push(numeric_update); + } + + MetricsUpdate::new(entries, entries_numeric) + } + + /// Update the episode-end metrics for validation from an episode summary. + pub(crate) fn update_episode_end_valid( + &mut self, + item: &EvaluationItem, + metadata: &MetricMetadata, + ) -> MetricsUpdate { + let mut entries = Vec::with_capacity(self.episode_end_valid.len()); + let mut entries_numeric = Vec::with_capacity(self.episode_end_valid_numeric.len()); + + for metric in self.episode_end_valid.iter_mut() { + let state = metric.update(&item.item, metadata); + entries.push(state); + } + + for metric in self.episode_end_valid_numeric.iter_mut() { + let numeric_update = metric.update(&item.item, metadata); + entries_numeric.push(numeric_update); + } + + MetricsUpdate::new(entries, entries_numeric) + } +} diff --git a/crates/burn-train/src/metric/processor/rl_processor.rs b/crates/burn-train/src/metric/processor/rl_processor.rs new file mode 100644 index 0000000..acc8c7f --- /dev/null +++ b/crates/burn-train/src/metric/processor/rl_processor.rs @@ -0,0 +1,215 @@ +use std::sync::Arc; + +use crate::{ + EpisodeSummary, EvaluationItem, EventProcessorTraining, ItemLazy, LearnerSummary, RLMetrics, + logger::TrainingProgressLogger, + metric::store::{Event, EventStoreClient, MetricsUpdate}, + renderer::{MetricState, MetricsRenderer}, +}; + +/// Event happening during reinforcement learning. +pub enum RLEvent { + /// Signal the start of the process (e.g., learning starts). + Start { + /// The total number of items to process during training (e.g., total number of environment steps). + total_items: usize, + }, + /// Signal an agent's training step. + TrainStep(EvaluationItem), + /// Signal a timestep of the agent-environment interface. + EnvStep(EvaluationItem), + /// Signal an episode end. + EpisodeEnd(EvaluationItem), + /// Signal the end of the process (e.g., learning ends). + End(Option), +} + +/// Event happening during evaluation of a reinforcement learning's agent. +pub enum AgentEvaluationEvent { + /// Signal the start of the process (e.g., training start) + Start(usize), + /// Signal a timestep of the agent-environment interface. + EnvStep(EvaluationItem), + /// Signal an episode end. + EpisodeEnd(EvaluationItem), + /// Signal the end of the process (e.g., training end). + End, +} + +/// An [event processor](EventProcessorTraining) that handles: +/// - Computing and storing metrics in an [event store](crate::metric::store::EventStore). +/// - Render metrics using a [metrics renderer](MetricsRenderer). +pub struct RLEventProcessor { + metrics: RLMetrics, + renderer: Box, + store: Arc, + training_progress_logger: Option>, +} + +impl RLEventProcessor { + pub(crate) fn new( + metrics: RLMetrics, + renderer: Box, + store: Arc, + ) -> Self { + Self { + metrics, + renderer, + store, + training_progress_logger: None, + } + } + + fn process_update_train(&mut self, update: MetricsUpdate) { + self.store + .add_event_train(crate::metric::store::Event::MetricsUpdate(update.clone())); + + update + .entries + .into_iter() + .for_each(|entry| self.renderer.update_train(MetricState::Generic(entry))); + + update + .entries_numeric + .into_iter() + .for_each(|numeric_update| { + self.renderer.update_train(MetricState::Numeric( + numeric_update.entry, + numeric_update.numeric_entry, + )) + }); + } + + fn process_update_valid(&mut self, update: MetricsUpdate) { + self.store + .add_event_valid(crate::metric::store::Event::MetricsUpdate(update.clone())); + + update + .entries + .into_iter() + .for_each(|entry| self.renderer.update_valid(MetricState::Generic(entry))); + + update + .entries_numeric + .into_iter() + .for_each(|numeric_update| { + self.renderer.update_valid(MetricState::Numeric( + numeric_update.entry, + numeric_update.numeric_entry, + )) + }); + } +} + +impl EventProcessorTraining, AgentEvaluationEvent> + for RLEventProcessor +{ + fn process_train(&mut self, event: RLEvent) { + match event { + RLEvent::Start { total_items } => { + let definitions = self.metrics.metric_definitions(); + self.store + .add_event_train(Event::MetricsInit(definitions.clone())); + definitions + .iter() + .for_each(|definition| self.renderer.register_metric(definition.clone())); + if let Some(logger) = &mut self.training_progress_logger { + logger.start(0, 0, Some(total_items)); + } + self.renderer.start(0, 0, Some(total_items)); + } + RLEvent::TrainStep(item) => { + let item = item.sync(); + let metadata = (&item).into(); + + let update = self.metrics.update_train_step(&item, &metadata); + self.process_update_train(update); + + if let Some(logger) = &mut self.training_progress_logger { + logger.log_event_training("TrainStep".to_string()); + } + self.renderer.log_event_training("TrainStep".to_string()); + } + RLEvent::EnvStep(item) => { + let item = item.sync(); + let metadata = (&item).into(); + + let update = self.metrics.update_env_step(&item, &metadata); + self.process_update_train(update); + + if let Some(logger) = &mut self.training_progress_logger { + logger.update_split(item.progress.items_processed); + logger.log_event_training("EnvStep".to_string()); + } + self.renderer.update_split(item.progress.items_processed); + self.renderer.log_event_training("EnvStep".to_string()); + } + RLEvent::EpisodeEnd(item) => { + let item = item.sync(); + let metadata = (&item).into(); + + let update = self.metrics.update_episode_end(&item, &metadata); + self.process_update_train(update); + + if let Some(logger) = &mut self.training_progress_logger { + logger.log_event_training("EpisodeEnd".to_string()); + } + self.renderer.log_event_training("EpisodeEnd".to_string()); + } + RLEvent::End(learner_summary) => { + if let Some(logger) = &mut self.training_progress_logger { + logger.end(); + } + self.renderer.end(); + self.renderer.on_train_end(learner_summary).ok(); + } + } + } + + fn process_valid(&mut self, event: AgentEvaluationEvent) { + match event { + AgentEvaluationEvent::Start(num_episodes) => { + if let Some(logger) = &mut self.training_progress_logger { + logger.start_split("valid", num_episodes); + } + self.renderer.start_split("valid", num_episodes); + } + AgentEvaluationEvent::EnvStep(item) => { + let item = item.sync(); + let metadata = (&item).into(); + + let update = self.metrics.update_env_step_valid(&item, &metadata); + self.process_update_valid(update); + + if let Some(logger) = &mut self.training_progress_logger { + logger.log_event_training("EnvStep".to_string()); + } + self.renderer.log_event_training("EnvStep".to_string()); + } + AgentEvaluationEvent::EpisodeEnd(item) => { + let item = item.sync(); + let metadata = (&item).into(); + + let update = self.metrics.update_episode_end_valid(&item, &metadata); + self.process_update_valid(update); + + if let Some(logger) = &mut self.training_progress_logger { + logger.update_split(item.progress.items_processed); + logger.log_event_training("EpisodeEnd".to_string()); + } + self.renderer.update_split(item.progress.items_processed); + self.renderer.log_event_training("EpisodeEnd".to_string()); + } + AgentEvaluationEvent::End => { + if let Some(logger) = &mut self.training_progress_logger { + logger.end_split(); + } + self.renderer.end_split(); + } + } + } + + fn renderer(self) -> Box { + self.renderer + } +} diff --git a/crates/burn-train/src/metric/recall.rs b/crates/burn-train/src/metric/recall.rs new file mode 100644 index 0000000..0cd1fe0 --- /dev/null +++ b/crates/burn-train/src/metric/recall.rs @@ -0,0 +1,212 @@ +use crate::metric::{MetricName, Numeric}; + +use super::{ + Metric, MetricAttributes, MetricMetadata, NumericAttributes, NumericEntry, SerializedEntry, + classification::{ClassReduction, ClassificationMetricConfig, DecisionRule}, + confusion_stats::{ConfusionStats, ConfusionStatsInput}, + state::{FormatOptions, NumericMetricState}, +}; +use burn_core::prelude::Tensor; +use std::{num::NonZeroUsize, sync::Arc}; + +///The Recall Metric +#[derive(Clone)] +pub struct RecallMetric { + name: MetricName, + state: NumericMetricState, + config: ClassificationMetricConfig, +} + +impl Default for RecallMetric { + fn default() -> Self { + Self::new(Default::default()) + } +} + +impl RecallMetric { + fn new(config: ClassificationMetricConfig) -> Self { + let state = Default::default(); + let name = Arc::new(format!( + "Recall @ {:?} [{:?}]", + config.decision_rule, config.class_reduction + )); + + Self { + state, + config, + name, + } + } + /// Recall metric for binary classification. + /// + /// # Arguments + /// + /// * `threshold` - The threshold to transform a probability into a binary prediction. + #[allow(dead_code)] + pub fn binary(threshold: f64) -> Self { + Self::new(ClassificationMetricConfig { + decision_rule: DecisionRule::Threshold(threshold), + // binary classification results are the same independently of class_reduction + ..Default::default() + }) + } + + /// Recall metric for multiclass classification. + /// + /// # Arguments + /// + /// * `top_k` - The number of highest predictions considered to find the correct label (typically `1`). + /// * `class_reduction` - [Class reduction](ClassReduction) type. + #[allow(dead_code)] + pub fn multiclass(top_k: usize, class_reduction: ClassReduction) -> Self { + Self::new(ClassificationMetricConfig { + decision_rule: DecisionRule::TopK( + NonZeroUsize::new(top_k).expect("top_k must be non-zero"), + ), + class_reduction, + }) + } + + /// Recall metric for multi-label classification. + /// + /// # Arguments + /// + /// * `threshold` - The threshold to transform a probability into a binary prediction. + /// * `class_reduction` - [Class reduction](ClassReduction) type. + #[allow(dead_code)] + pub fn multilabel(threshold: f64, class_reduction: ClassReduction) -> Self { + Self::new(ClassificationMetricConfig { + decision_rule: DecisionRule::Threshold(threshold), + class_reduction, + }) + } + + fn class_average(&self, mut aggregated_metric: Tensor<1>) -> f64 { + use ClassReduction::{Macro, Micro}; + let avg_tensor = match self.config.class_reduction { + Micro => aggregated_metric, + Macro => { + if aggregated_metric.clone().contains_nan().any().into_scalar() { + let nan_mask = aggregated_metric.clone().is_nan(); + aggregated_metric = aggregated_metric + .clone() + .select(0, nan_mask.bool_not().argwhere().squeeze_dim(1)) + } + aggregated_metric.mean() + } + }; + avg_tensor.into_scalar() + } +} + +impl Metric for RecallMetric { + type Input = ConfusionStatsInput; + + fn update(&mut self, input: &Self::Input, _metadata: &MetricMetadata) -> SerializedEntry { + let [sample_size, _] = input.predictions.dims(); + + let cf_stats = ConfusionStats::new(input, &self.config); + let metric = self.class_average(cf_stats.clone().true_positive() / cf_stats.positive()); + + self.state.update( + 100.0 * metric, + sample_size, + FormatOptions::new(self.name()).unit("%").precision(2), + ) + } + + fn clear(&mut self) { + self.state.reset() + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: Some("%".to_string()), + higher_is_better: true, + ..Default::default() + } + .into() + } +} + +impl Numeric for RecallMetric { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} + +#[cfg(test)] +mod tests { + use super::{ + ClassReduction::{self, *}, + Metric, MetricMetadata, RecallMetric, + }; + use crate::metric::Numeric; + use crate::tests::{ClassificationType, THRESHOLD, dummy_classification_input}; + use burn_core::tensor::{TensorData, Tolerance}; + use rstest::rstest; + + #[rstest] + #[case::binary(THRESHOLD, 0.5)] + fn test_binary_recall(#[case] threshold: f64, #[case] expected: f64) { + let input = dummy_classification_input(&ClassificationType::Binary).into(); + let mut metric = RecallMetric::binary(threshold); + let _entry = metric.update(&input, &MetricMetadata::fake()); + TensorData::from([metric.value().current()]) + .assert_approx_eq::(&TensorData::from([expected * 100.0]), Tolerance::default()) + } + + #[rstest] + #[case::multiclass_micro_k1(Micro, 1, 3.0/5.0)] + #[case::multiclass_micro_k2(Micro, 2, 4.0/5.0)] + #[case::multiclass_macro_k1(Macro, 1, (0.5 + 1.0 + 0.5)/3.0)] + #[case::multiclass_macro_k2(Macro, 2, (1.0 + 1.0 + 0.5)/3.0)] + fn test_multiclass_recall( + #[case] class_reduction: ClassReduction, + #[case] top_k: usize, + #[case] expected: f64, + ) { + let input = dummy_classification_input(&ClassificationType::Multiclass).into(); + let mut metric = RecallMetric::multiclass(top_k, class_reduction); + let _entry = metric.update(&input, &MetricMetadata::fake()); + TensorData::from([metric.value().current()]) + .assert_approx_eq::(&TensorData::from([expected * 100.0]), Tolerance::default()) + } + + #[rstest] + #[case::multilabel_micro(Micro, THRESHOLD, 5.0/9.0)] + #[case::multilabel_macro(Macro, THRESHOLD, (0.5 + 1.0 + 1.0/3.0)/3.0)] + fn test_multilabel_recall( + #[case] class_reduction: ClassReduction, + #[case] threshold: f64, + #[case] expected: f64, + ) { + let input = dummy_classification_input(&ClassificationType::Multilabel).into(); + let mut metric = RecallMetric::multilabel(threshold, class_reduction); + let _entry = metric.update(&input, &MetricMetadata::fake()); + TensorData::from([metric.value().current()]) + .assert_approx_eq::(&TensorData::from([expected * 100.0]), Tolerance::default()) + } + + #[test] + fn test_parameterized_unique_name() { + let metric_a = RecallMetric::multiclass(1, ClassReduction::Macro); + let metric_b = RecallMetric::multiclass(2, ClassReduction::Macro); + let metric_c = RecallMetric::multiclass(1, ClassReduction::Macro); + + assert_ne!(metric_a.name(), metric_b.name()); + assert_eq!(metric_a.name(), metric_c.name()); + + let metric_a = RecallMetric::binary(0.5); + let metric_b = RecallMetric::binary(0.75); + assert_ne!(metric_a.name(), metric_b.name()); + } +} diff --git a/crates/burn-train/src/metric/rl/cum_reward.rs b/crates/burn-train/src/metric/rl/cum_reward.rs new file mode 100644 index 0000000..b20ed96 --- /dev/null +++ b/crates/burn-train/src/metric/rl/cum_reward.rs @@ -0,0 +1,79 @@ +use std::sync::Arc; + +use super::super::{ + MetricAttributes, MetricMetadata, NumericAttributes, NumericEntry, + state::{FormatOptions, NumericMetricState}, +}; +use crate::metric::{Metric, MetricName, Numeric, SerializedEntry}; + +/// Metric for the cumulative reward of the last completed episode. +#[derive(Clone)] +pub struct CumulativeRewardMetric { + name: MetricName, + state: NumericMetricState, +} + +impl CumulativeRewardMetric { + /// Creates a new episode length metric. + pub fn new() -> Self { + Self { + name: Arc::new("Cum. Reward".to_string()), + state: NumericMetricState::new(), + } + } +} + +impl Default for CumulativeRewardMetric { + fn default() -> Self { + Self::new() + } +} + +/// The [CumulativeRewardMetric](CumulativeRewardMetric) input type. +#[derive(new)] +pub struct CumulativeRewardInput { + cum_reward: f64, +} + +impl Metric for CumulativeRewardMetric { + type Input = CumulativeRewardInput; + + fn update( + &mut self, + item: &CumulativeRewardInput, + _metadata: &MetricMetadata, + ) -> SerializedEntry { + self.state.update( + item.cum_reward, + 1, + FormatOptions::new(self.name()).precision(2), + ) + } + + fn clear(&mut self) { + self.state.reset() + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: None, + higher_is_better: true, + ..Default::default() + } + .into() + } +} + +impl Numeric for CumulativeRewardMetric { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} diff --git a/crates/burn-train/src/metric/rl/ep_len.rs b/crates/burn-train/src/metric/rl/ep_len.rs new file mode 100644 index 0000000..177d03e --- /dev/null +++ b/crates/burn-train/src/metric/rl/ep_len.rs @@ -0,0 +1,72 @@ +use std::sync::Arc; + +use super::super::{ + MetricAttributes, MetricMetadata, NumericAttributes, NumericEntry, + state::{FormatOptions, NumericMetricState}, +}; +use crate::metric::{Metric, MetricName, Numeric, SerializedEntry}; + +/// Metric for the length of the last completed episode. +#[derive(Clone)] +pub struct EpisodeLengthMetric { + name: MetricName, + state: NumericMetricState, +} + +impl EpisodeLengthMetric { + /// Creates a new episode length metric. + pub fn new() -> Self { + Self { + name: Arc::new("Episode length".to_string()), + state: NumericMetricState::new(), + } + } +} + +impl Default for EpisodeLengthMetric { + fn default() -> Self { + Self::new() + } +} + +/// The [EpisodeLengthMetric](EpisodeLengthMetric) input type. +#[derive(new)] +pub struct EpisodeLengthInput { + ep_len: f64, +} + +impl Metric for EpisodeLengthMetric { + type Input = EpisodeLengthInput; + + fn update(&mut self, item: &EpisodeLengthInput, _metadata: &MetricMetadata) -> SerializedEntry { + self.state + .update(item.ep_len, 1, FormatOptions::new(self.name()).precision(0)) + } + + fn clear(&mut self) { + self.state.reset() + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: Some(String::from("steps")), + higher_is_better: true, + ..Default::default() + } + .into() + } +} + +impl Numeric for EpisodeLengthMetric { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} diff --git a/crates/burn-train/src/metric/rl/exploration_rate.rs b/crates/burn-train/src/metric/rl/exploration_rate.rs new file mode 100644 index 0000000..a57a4ec --- /dev/null +++ b/crates/burn-train/src/metric/rl/exploration_rate.rs @@ -0,0 +1,79 @@ +use std::sync::Arc; + +use super::super::{ + MetricAttributes, MetricMetadata, NumericAttributes, NumericEntry, + state::{FormatOptions, NumericMetricState}, +}; +use crate::metric::{Metric, MetricName, Numeric, SerializedEntry}; + +/// Metric for the length of the last completed episode. +#[derive(Clone)] +pub struct ExplorationRateMetric { + name: MetricName, + state: NumericMetricState, +} + +impl ExplorationRateMetric { + /// Creates a new episode length metric. + pub fn new() -> Self { + Self { + name: Arc::new("Exploration rate".to_string()), + state: NumericMetricState::new(), + } + } +} + +impl Default for ExplorationRateMetric { + fn default() -> Self { + Self::new() + } +} + +/// The [ExplorationRateMetric](ExplorationRateMetric) input type. +#[derive(new)] +pub struct ExplorationRateInput { + exploration_rate: f64, +} + +impl Metric for ExplorationRateMetric { + type Input = ExplorationRateInput; + + fn update( + &mut self, + item: &ExplorationRateInput, + _metadata: &MetricMetadata, + ) -> SerializedEntry { + self.state.update( + item.exploration_rate, + 1, + FormatOptions::new(self.name()).precision(3), + ) + } + + fn clear(&mut self) { + self.state.reset() + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: Some(String::from("%")), + higher_is_better: false, + ..Default::default() + } + .into() + } +} + +impl Numeric for ExplorationRateMetric { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} diff --git a/crates/burn-train/src/metric/rl/mod.rs b/crates/burn-train/src/metric/rl/mod.rs new file mode 100644 index 0000000..f2a2d22 --- /dev/null +++ b/crates/burn-train/src/metric/rl/mod.rs @@ -0,0 +1,7 @@ +mod cum_reward; +mod ep_len; +mod exploration_rate; + +pub use cum_reward::*; +pub use ep_len::*; +pub use exploration_rate::*; diff --git a/crates/burn-train/src/metric/rouge.rs b/crates/burn-train/src/metric/rouge.rs new file mode 100644 index 0000000..f970573 --- /dev/null +++ b/crates/burn-train/src/metric/rouge.rs @@ -0,0 +1,282 @@ +use super::state::{FormatOptions, NumericMetricState}; +use super::{MetricMetadata, SerializedEntry}; +use crate::metric::{ + Metric, MetricAttributes, MetricName, Numeric, NumericAttributes, NumericEntry, +}; +use burn_core::tensor::{Int, Tensor}; +use std::sync::Arc; + +fn lcs_length(a: &[i32], b: &[i32]) -> usize { + let (shorter, longer) = if a.len() <= b.len() { (a, b) } else { (b, a) }; + let mut prev = vec![0usize; shorter.len() + 1]; + let mut curr = vec![0usize; shorter.len() + 1]; + + for &x in longer { + for (j, &y) in shorter.iter().enumerate() { + if x == y { + curr[j + 1] = prev[j] + 1; + } else { + curr[j + 1] = curr[j].max(prev[j + 1]); + } + } + core::mem::swap(&mut prev, &mut curr); + } + prev[shorter.len()] +} + +/// ROUGE-L metric based on longest common subsequence. +#[derive(Clone)] +pub struct RougeLScore { + name: MetricName, + state: NumericMetricState, + pad_token: Option, +} + +/// Input for [RougeLScore]. +#[derive(new)] +pub struct RougeLInput { + /// Predicted token sequences. + pub outputs: Tensor<2, Int>, + /// Reference token sequences. + pub targets: Tensor<2, Int>, +} + +impl Default for RougeLScore { + fn default() -> Self { + Self::new() + } +} + +impl RougeLScore { + /// Creates a new ROUGE-L metric. + pub fn new() -> Self { + Self { + name: Arc::new("ROUGE-L".to_string()), + state: NumericMetricState::default(), + pad_token: None, + } + } + + /// Sets the pad token index. Tokens matching this value are stripped + /// from the right of each sequence before scoring. + pub fn with_pad_token(mut self, index: usize) -> Self { + self.pad_token = Some(index); + self + } +} + +impl Metric for RougeLScore { + type Input = RougeLInput; + + fn update(&mut self, input: &RougeLInput, _metadata: &MetricMetadata) -> SerializedEntry { + let outputs = &input.outputs; + let targets = &input.targets; + let [batch_size, seq_len] = targets.dims(); + + let outputs_data = outputs.to_data().iter::().collect::>(); + let targets_data = targets.to_data().iter::().collect::>(); + + let pad_token = self.pad_token.map(|p| p as i32); + + let mut total_f1 = 0.0_f64; + + for i in 0..batch_size { + let start = i * seq_len; + let end = (i + 1) * seq_len; + + let output_seq = &outputs_data[start..end]; + let target_seq = &targets_data[start..end]; + + let output_seq = match pad_token { + Some(pad) => { + let len = output_seq + .iter() + .position(|&x| x == pad) + .unwrap_or(output_seq.len()); + &output_seq[..len] + } + None => output_seq, + }; + let target_seq = match pad_token { + Some(pad) => { + let len = target_seq + .iter() + .position(|&x| x == pad) + .unwrap_or(target_seq.len()); + &target_seq[..len] + } + None => target_seq, + }; + + let lcs_len = lcs_length(target_seq, output_seq) as f64; + let ref_len = target_seq.len() as f64; + let cand_len = output_seq.len() as f64; + + if ref_len == 0.0 && cand_len == 0.0 { + total_f1 += 100.0; + continue; + } + + if ref_len == 0.0 || cand_len == 0.0 { + continue; + } + + let precision = lcs_len / cand_len; + let recall = lcs_len / ref_len; + + let f1 = if precision + recall > 0.0 { + 2.0 * precision * recall / (precision + recall) + } else { + 0.0 + }; + + total_f1 += f1 * 100.0; + } + + let value = total_f1 / batch_size as f64; + + self.state.update( + value, + batch_size, + FormatOptions::new(self.name()).unit("%").precision(2), + ) + } + + fn clear(&mut self) { + self.state.reset(); + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: Some("%".to_string()), + higher_is_better: true, + ..Default::default() + } + .into() + } +} + +impl Numeric for RougeLScore { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn_core::tensor::TensorData; + + #[test] + fn test_rouge_l_perfect_match() { + let device = Default::default(); + let mut metric = RougeLScore::new(); + + let preds = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + let tgts = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + + metric.update(&RougeLInput::new(preds, tgts), &MetricMetadata::fake()); + assert!((metric.value().current() - 100.0).abs() < 1e-6); + } + + #[test] + fn test_rouge_l_no_match() { + let device = Default::default(); + let mut metric = RougeLScore::new(); + + let preds = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + let tgts = Tensor::from_data([[6, 7, 8, 9, 10]], &device); + + metric.update(&RougeLInput::new(preds, tgts), &MetricMetadata::fake()); + assert_eq!(0.0, metric.value().current()); + } + + #[test] + fn test_rouge_l_partial_match() { + let device = Default::default(); + let mut metric = RougeLScore::new(); + + let preds = Tensor::from_data([[1, 3, 5, 7, 9]], &device); + let tgts = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + + metric.update(&RougeLInput::new(preds, tgts), &MetricMetadata::fake()); + + let expected = 60.0; + assert!((metric.value().current() - expected).abs() < 1e-6); + } + + #[test] + fn test_rouge_l_shorter_candidate() { + let device = Default::default(); + let pad = 99_i64; + let mut metric = RougeLScore::new().with_pad_token(pad as usize); + + let preds = Tensor::from_data([[1, 2, 3, pad, pad]], &device); + let tgts = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + + metric.update(&RougeLInput::new(preds, tgts), &MetricMetadata::fake()); + + let expected = 75.0; + assert!((metric.value().current() - expected).abs() < 1e-6); + } + + #[test] + fn test_rouge_l_with_padding() { + let device = Default::default(); + let pad = 99_i64; + let mut metric = RougeLScore::new().with_pad_token(pad as usize); + + let preds = Tensor::from_data([[1, 2, 3, 4, 5, pad, pad]], &device); + let tgts = Tensor::from_data([[1, 2, 3, 4, 5, pad, pad]], &device); + + metric.update(&RougeLInput::new(preds, tgts), &MetricMetadata::fake()); + assert!((metric.value().current() - 100.0).abs() < 1e-6); + } + + #[test] + fn test_rouge_l_batch() { + let device = Default::default(); + let mut metric = RougeLScore::new(); + + let preds = Tensor::from_data([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], &device); + let tgts = Tensor::from_data([[1, 2, 3, 4, 5], [11, 12, 13, 14, 15]], &device); + + metric.update(&RougeLInput::new(preds, tgts), &MetricMetadata::fake()); + assert!((metric.value().current() - 50.0).abs() < 1e-6); + } + + #[test] + fn test_rouge_l_both_empty() { + let device = Default::default(); + let mut metric = RougeLScore::new(); + + let preds = Tensor::<2, Int>::from_data(TensorData::from([[0i32; 0]; 1]), &device); + let tgts = Tensor::<2, Int>::from_data(TensorData::from([[0i32; 0]; 1]), &device); + + metric.update(&RougeLInput::new(preds, tgts), &MetricMetadata::fake()); + assert!((metric.value().current() - 100.0).abs() < 1e-6); + } + + #[test] + fn test_clear_resets_state() { + let device = Default::default(); + let mut metric = RougeLScore::new(); + + let preds = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + let tgts = Tensor::from_data([[1, 2, 3, 4, 5]], &device); + + metric.update(&RougeLInput::new(preds, tgts), &MetricMetadata::fake()); + assert!(metric.value().current() > 0.0); + + metric.clear(); + assert!(metric.value().current().is_nan()); + } +} diff --git a/crates/burn-train/src/metric/state.rs b/crates/burn-train/src/metric/state.rs new file mode 100644 index 0000000..4353359 --- /dev/null +++ b/crates/burn-train/src/metric/state.rs @@ -0,0 +1,219 @@ +use std::sync::Arc; + +use burn_core::tensor::{Bool, Tensor}; + +use crate::metric::{MetricName, NumericEntry, SerializedEntry, format_float}; + +/// Useful utility to implement numeric metrics. +/// +/// # Notes +/// +/// The numeric metric store values inside floats. +/// Even if some metric are integers, their mean are floats. +#[derive(Clone)] +pub struct NumericMetricState { + sum: f64, + count: usize, + current: f64, + current_count: usize, +} +/// Accumulates raw predictions and targets across batches. +/// +/// Used by rank-based metrics (AUROC, AUC-PR) that must recompute over the +/// whole epoch. Buffers are freed on [`reset`](Self::reset). +#[derive(Clone)] +pub struct PredictionAccumulatorState { + predictions: Vec>, + targets: Vec>, + current: f64, +} + +/// Formatting options for the [numeric metric state](NumericMetricState). +pub struct FormatOptions { + name: Arc, + unit: Option, + precision: Option, +} + +impl PredictionAccumulatorState { + /// Create a new [prediction accumulator state](PredictionAccumulatorState). + pub fn new() -> Self { + Self { + predictions: vec![], + targets: vec![], + current: f64::NAN, + } + } + + /// Accumulate a batch of predictions and targets. + pub fn accumulate(&mut self, preds: Tensor<2>, targets: Tensor<2, Bool>) { + self.predictions.push(preds); + self.targets.push(targets); + } + + /// All accumulated predictions and targets, concatenated along the samples. + pub fn tensors(&self) -> (Tensor<2>, Tensor<2, Bool>) { + ( + Tensor::cat(self.predictions.clone(), 0), + Tensor::cat(self.targets.clone(), 0), + ) + } + + /// Record the value computed over the accumulated set and return the entry + /// to log. Metrics using this state must declare + /// [`NumericAggregation::Last`](crate::metric::NumericAggregation). + pub fn update(&mut self, value: f64, format: FormatOptions) -> SerializedEntry { + self.current = value; + + let serialized = NumericEntry::Value(value).serialize(); + + let formatted_value = match format.precision { + Some(precision) => format_float(value, precision), + None => format!("{value}"), + }; + let formatted = match format.unit { + Some(unit) => format!("epoch {formatted_value} {unit}"), + None => format!("epoch {formatted_value}"), + }; + + SerializedEntry::new(formatted, serialized) + } + + /// Current value (computed over the accumulated set). + pub fn value(&self) -> NumericEntry { + NumericEntry::Value(self.current) + } + + /// Reset the state, freeing the accumulated tensors. + pub fn reset(&mut self) { + self.predictions.clear(); + self.targets.clear(); + self.current = f64::NAN; + } +} + +impl Default for PredictionAccumulatorState { + fn default() -> Self { + Self::new() + } +} + +impl FormatOptions { + /// Create the [formatting options](FormatOptions) with a name. + pub fn new(name: MetricName) -> Self { + Self { + name: name.clone(), + unit: None, + precision: None, + } + } + + /// Specify the metric unit. + pub fn unit(mut self, unit: &str) -> Self { + self.unit = Some(unit.to_string()); + self + } + + /// Specify the floating point precision. + pub fn precision(mut self, precision: usize) -> Self { + self.precision = Some(precision); + self + } + + /// Get the metric name. + pub fn name(&self) -> &Arc { + &self.name + } + + /// Get the metric unit. + pub fn unit_value(&self) -> &Option { + &self.unit + } + + /// Get the precision. + pub fn precision_value(&self) -> Option { + self.precision + } +} + +impl NumericMetricState { + /// Create a new [numeric metric state](NumericMetricState). + pub fn new() -> Self { + Self { + sum: 0.0, + count: 0, + current: f64::NAN, + current_count: 0, + } + } + + /// Reset the state. + pub fn reset(&mut self) { + self.sum = 0.0; + self.count = 0; + self.current = f64::NAN; + self.current_count = 0; + } + + /// Update the state. + pub fn update( + &mut self, + value: f64, + batch_size: usize, + format: FormatOptions, + ) -> SerializedEntry { + self.sum += value * batch_size as f64; + self.count += batch_size; + self.current = value; + self.current_count = batch_size; + + let value_current = value; + let value_running = self.sum / self.count as f64; + // Numeric metric state is an aggregated value + let serialized = NumericEntry::Aggregated { + aggregated_value: value_current, + count: batch_size, + } + .serialize(); + + let (formatted_current, formatted_running) = match format.precision { + Some(precision) => ( + format_float(value_current, precision), + format_float(value_running, precision), + ), + None => (format!("{value_current}"), format!("{value_running}")), + }; + + // TODO: naming inconsistent with RL. + let formatted = match format.unit { + Some(unit) => { + format!("epoch {formatted_running} {unit} - batch {formatted_current} {unit}") + } + None => format!("epoch {formatted_running} - batch {formatted_current}"), + }; + + SerializedEntry::new(formatted, serialized) + } + + /// Get the numeric value. + pub fn current_value(&self) -> NumericEntry { + NumericEntry::Aggregated { + aggregated_value: self.current, + count: self.current_count, + } + } + + /// Get the running aggregated value. + pub fn running_value(&self) -> NumericEntry { + NumericEntry::Aggregated { + aggregated_value: self.sum / self.count as f64, + count: self.count, + } + } +} + +impl Default for NumericMetricState { + fn default() -> Self { + Self::new() + } +} diff --git a/crates/burn-train/src/metric/store/aggregate.rs b/crates/burn-train/src/metric/store/aggregate.rs new file mode 100644 index 0000000..051bef7 --- /dev/null +++ b/crates/burn-train/src/metric/store/aggregate.rs @@ -0,0 +1,312 @@ +use crate::{ + logger::MetricLogger, + metric::{MetricAttributes, MetricDefinition, NumericAggregation, NumericEntry, store::Split}, +}; +use std::collections::HashMap; + +use super::{Aggregate, Direction}; + +/// Type that can be used to fetch and use numeric metric aggregates. +#[derive(Default, Debug)] +pub(crate) struct NumericMetricsAggregate { + value_for_each_epoch: HashMap, + aggregations: HashMap, +} + +#[derive(new, Hash, PartialEq, Eq, Debug)] +struct Key { + name: String, + epoch: usize, + split: Split, + aggregate: Aggregate, +} + +impl NumericMetricsAggregate { + /// Record each numeric metric's epoch-aggregation strategy from its definition. + pub(crate) fn register_definitions(&mut self, definitions: &[MetricDefinition]) { + for def in definitions { + if let MetricAttributes::Numeric(numeric) = &def.attributes { + self.aggregations + .insert(def.name.clone(), numeric.aggregation); + } + } + } + + pub(crate) fn aggregate( + &mut self, + name: &str, + epoch: usize, + split: &Split, + aggregate: Aggregate, + loggers: &mut [Box], + ) -> Option { + let key = Key::new(name.to_string(), epoch, split.clone(), aggregate); + + if let Some(value) = self.value_for_each_epoch.get(&key) { + return Some(*value); + } + + // How this metric reduces its per-batch points (defaults to a weighted mean). + let aggregation = self.aggregations.get(name).copied().unwrap_or_default(); + + let points = || { + let mut errors = Vec::new(); + for logger in loggers { + match logger.read_numeric(name, epoch, split) { + Ok(points) => return Ok(points), + Err(err) => errors.push(err), + }; + } + + Err(errors.join(" ")) + }; + + let points = points().expect("Can read values"); + + if points.is_empty() { + return None; + } + + let value = match aggregation { + // Already computed over the whole epoch; do not average. + NumericAggregation::Last => points.last().expect("Points are not empty").current(), + // Weighted by the *actual* number of points since mini-batches may differ in size. + NumericAggregation::Mean => { + let (sum, num_points) = points + .into_iter() + .map(|entry| match entry { + NumericEntry::Value(v) => (v, 1), + // Right now the mean is the only aggregate available, so we can assume that the sum + // of an entry corresponds to (value * number of elements) + NumericEntry::Aggregated { + aggregated_value, + count, + } => (aggregated_value * count as f64, count), + }) + .reduce(|(acc_v, acc_n), (v, n)| (acc_v + v, acc_n + n)) + .unwrap(); + match aggregate { + Aggregate::Mean => sum / num_points as f64, + } + } + }; + + self.value_for_each_epoch.insert(key, value); + Some(value) + } + + pub(crate) fn find_epoch( + &mut self, + name: &str, + split: &Split, + aggregate: Aggregate, + direction: Direction, + loggers: &mut [Box], + ) -> Option { + let mut data = Vec::new(); + let mut current_epoch = 1; + + while let Some(value) = self.aggregate(name, current_epoch, split, aggregate, loggers) { + data.push(value); + current_epoch += 1; + } + + if data.is_empty() { + return None; + } + + let mut current_value = match &direction { + Direction::Lowest => f64::MAX, + Direction::Highest => f64::MIN, + }; + + for (i, value) in data.into_iter().enumerate() { + match &direction { + Direction::Lowest => { + if value < current_value { + current_value = value; + current_epoch = i + 1; + } + } + Direction::Highest => { + if value > current_value { + current_value = value; + current_epoch = i + 1; + } + } + } + } + + Some(current_epoch) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use crate::{ + logger::{FileMetricLogger, InMemoryMetricLogger}, + metric::{MetricDefinition, MetricEntry, MetricId, SerializedEntry, store::MetricsUpdate}, + }; + + use super::*; + + struct TestLogger { + logger: FileMetricLogger, + epoch: usize, + } + const NAME: &str = "test-logger"; + + impl TestLogger { + fn new() -> Self { + Self { + logger: FileMetricLogger::new("/tmp"), + epoch: 1, + } + } + fn log(&mut self, num: f64) { + let entry = MetricEntry::new( + MetricId::new(Arc::new(NAME.into())), + SerializedEntry::new(num.to_string(), num.to_string()), + ); + let entries = Vec::from([entry]); + let metrics_update = MetricsUpdate::new(entries, vec![]); + self.logger.log(metrics_update, self.epoch, &Split::Train); + } + fn log_definition(&mut self) { + let definition = MetricDefinition { + metric_id: MetricId::new(Arc::new(NAME.into())), + name: NAME.into(), + attributes: crate::metric::MetricAttributes::None, + description: None, + }; + self.logger.log_metric_definition(definition); + } + fn new_epoch(&mut self) { + self.epoch += 1; + } + } + + #[test] + fn should_find_epoch() { + let mut logger = TestLogger::new(); + let mut aggregate = NumericMetricsAggregate::default(); + logger.log_definition(); + + logger.log(500.); // Epoch 1 + logger.log(1000.); // Epoch 1 + logger.new_epoch(); + logger.log(200.); // Epoch 2 + logger.log(1000.); // Epoch 2 + logger.new_epoch(); + logger.log(10000.); // Epoch 3 + + let value = aggregate + .find_epoch( + NAME, + &Split::Train, + Aggregate::Mean, + Direction::Lowest, + &mut [Box::new(logger.logger)], + ) + .unwrap(); + + assert_eq!(value, 2); + } + + #[test] + fn should_aggregate_numeric_entry() { + let mut logger = InMemoryMetricLogger::default(); + let mut aggregate = NumericMetricsAggregate::default(); + let metric_name = Arc::new("Loss".to_string()); + let metric_id = MetricId::new(metric_name.clone()); + let definition = MetricDefinition { + metric_id: metric_id.clone(), + name: metric_name.to_string(), + attributes: crate::metric::MetricAttributes::None, + description: None, + }; + logger.log_metric_definition(definition); + + // Epoch 1 + let loss_1 = 0.5; + let loss_2 = 1.25; // (1.5 + 1.0) / 2 = 2.5 / 2 + let entry = MetricEntry::new( + metric_id.clone(), + SerializedEntry::new(loss_1.to_string(), NumericEntry::Value(loss_1).serialize()), + ); + let entries = Vec::from([entry]); + let metrics_update = MetricsUpdate::new(entries, vec![]); + logger.log(metrics_update, 1, &Split::Train); + let entry = MetricEntry::new( + metric_id.clone(), + SerializedEntry::new( + loss_2.to_string(), + NumericEntry::Aggregated { + aggregated_value: loss_2, + count: 2, + } + .serialize(), + ), + ); + let entries = Vec::from([entry]); + let metrics_update = MetricsUpdate::new(entries, vec![]); + logger.log(metrics_update, 1, &Split::Train); + + let value = aggregate + .aggregate( + &metric_name, + 1, + &Split::Train, + Aggregate::Mean, + &mut [Box::new(logger)], + ) + .unwrap(); + + // Average should be (0.5 + 1.25 * 2) / 3 = 1.0, not (0.5 + 1.25) / 2 = 0.875 + assert_eq!(value, 1.0); + } + + #[test] + fn should_take_last_value_for_last_aggregation() { + let mut logger = InMemoryMetricLogger::default(); + let mut aggregate = NumericMetricsAggregate::default(); + let metric_name = Arc::new("AUROC".to_string()); + let metric_id = MetricId::new(metric_name.clone()); + let definition = MetricDefinition { + metric_id: metric_id.clone(), + name: metric_name.to_string(), + attributes: crate::metric::NumericAttributes { + unit: None, + higher_is_better: true, + aggregation: NumericAggregation::Last, + } + .into(), + description: None, + }; + logger.log_metric_definition(definition.clone()); + aggregate.register_definitions(&[definition]); + + for value in [0.6, 0.7, 0.85] { + let entry = MetricEntry::new( + metric_id.clone(), + SerializedEntry::new(value.to_string(), NumericEntry::Value(value).serialize()), + ); + logger.log(MetricsUpdate::new(vec![entry], vec![]), 1, &Split::Train); + } + + let value = aggregate + .aggregate( + &metric_name, + 1, + &Split::Train, + Aggregate::Mean, + &mut [Box::new(logger)], + ) + .unwrap(); + + // Last point, not the mean (~0.717). + assert_eq!(value, 0.85); + } +} diff --git a/crates/burn-train/src/metric/store/base.rs b/crates/burn-train/src/metric/store/base.rs new file mode 100644 index 0000000..6f18fe5 --- /dev/null +++ b/crates/burn-train/src/metric/store/base.rs @@ -0,0 +1,117 @@ +use std::sync::Arc; + +use crate::metric::{MetricDefinition, MetricEntry, NumericEntry}; + +/// Event happening during the training/validation process. +pub enum Event { + /// Signal the start of an epoch with the wpoch number. + StartSplit(usize), + /// Signal the initialization of the metrics + MetricsInit(Vec), + /// Signal that metrics have been updated. + MetricsUpdate(MetricsUpdate), + /// Signal the end of an epoch. + EndEpoch(EpochSummary), +} + +/// Contains all metric information. +#[derive(new, Clone, Debug)] +pub struct NumericMetricUpdate { + /// Generic metric information. + pub entry: MetricEntry, + /// The numeric information. + pub numeric_entry: NumericEntry, + /// Numeric value averaged over the global step (epoch). + pub running_entry: NumericEntry, +} + +/// Contains all metric information. +#[derive(new, Clone, Debug)] +pub struct MetricsUpdate { + /// Metrics information related to non-numeric metrics. + pub entries: Vec, + /// Metrics information related to numeric metrics. + pub entries_numeric: Vec, +} + +/// Summary information about a given epoch +#[derive(new, Clone, Debug)] +pub struct EpochSummary { + /// Epoch number. + pub epoch_number: usize, + /// Dataset split (train, valid, test). + pub split: Split, +} + +/// Defines how training and validation events are collected and searched. +/// +/// This trait also exposes methods that uses the collected data to compute useful information. +pub trait EventStore: Send { + /// Collect a training/validation event. + fn add_event(&mut self, event: Event, split: Split); + + /// Find the epoch following the given criteria from the collected data. + fn find_epoch( + &mut self, + name: &str, + aggregate: Aggregate, + direction: Direction, + split: &Split, + ) -> Option; + + /// Find the metric value for the current epoch following the given criteria. + fn find_metric( + &mut self, + name: &str, + epoch: usize, + aggregate: Aggregate, + split: &Split, + ) -> Option; +} + +#[derive(Copy, Clone, Hash, PartialEq, Eq, Debug)] +/// How to aggregate the metric. +pub enum Aggregate { + /// Compute the average. + Mean, +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +/// The split to use. +pub enum Split { + /// The training split. + Train, + /// The validation split. + Valid, + /// The testing split, which might be tagged. + Test(Option>), +} + +impl std::fmt::Display for Split { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Split::Train => write!(f, "train"), + Split::Valid => write!(f, "valid"), + Split::Test(_) => write!(f, "test"), + } + } +} + +impl From for &'static str { + fn from(value: Split) -> Self { + match value { + Split::Train => "train", + Split::Valid => "valid", + Split::Test(_) => "test", + } + } +} + +#[derive(Copy, Clone)] +/// The direction of the query. +pub enum Direction { + /// Lower is better. + Lowest, + /// Higher is better. + Highest, +} diff --git a/crates/burn-train/src/metric/store/client.rs b/crates/burn-train/src/metric/store/client.rs new file mode 100644 index 0000000..3135cf4 --- /dev/null +++ b/crates/burn-train/src/metric/store/client.rs @@ -0,0 +1,171 @@ +use super::EventStore; +use super::{Aggregate, Direction, Event, Split}; +use std::sync::Arc; +use std::{sync::mpsc, thread::JoinHandle}; + +/// Type that allows to communicate with an [event store](EventStore). +pub struct EventStoreClient { + sender: mpsc::Sender, + handler: Option>, +} + +impl EventStoreClient { + /// Create a new [event store](EventStore) client. + pub(crate) fn new(store: C) -> Self + where + C: EventStore + 'static, + { + let (sender, receiver) = mpsc::channel(); + let thread = WorkerThread::new(store, receiver); + + let handler = std::thread::spawn(move || thread.run()); + let handler = Some(handler); + + Self { sender, handler } + } +} + +impl EventStoreClient { + /// Add a training event to the [event store](EventStore). + pub(crate) fn add_event_train(&self, event: Event) { + self.sender + .send(Message::OnEventTrain(event)) + .expect("Can send event to event store thread."); + } + + /// Add a validation event to the [event store](EventStore). + pub(crate) fn add_event_valid(&self, event: Event) { + self.sender + .send(Message::OnEventValid(event)) + .expect("Can send event to event store thread."); + } + + /// Add a testing event to the [event store](EventStore). + pub(crate) fn add_event_test(&self, event: Event, tag: Arc) { + self.sender + .send(Message::OnEventTest(event, tag)) + .expect("Can send event to event store thread."); + } + + /// Find the epoch following the given criteria from the collected data. + pub fn find_epoch( + &self, + name: &str, + aggregate: Aggregate, + direction: Direction, + split: &Split, + ) -> Option { + let (sender, receiver) = mpsc::sync_channel(1); + self.sender + .send(Message::FindEpoch( + name.to_string(), + aggregate, + direction, + split.clone(), + sender, + )) + .expect("Can send event to event store thread."); + + match receiver.recv() { + Ok(value) => value, + Err(err) => panic!("Event store thread crashed: {err:?}"), + } + } + + /// Find the metric value for the current epoch following the given criteria. + pub fn find_metric( + &self, + name: &str, + epoch: usize, + aggregate: Aggregate, + split: &Split, + ) -> Option { + let (sender, receiver) = mpsc::sync_channel(1); + self.sender + .send(Message::FindMetric( + name.to_string(), + epoch, + aggregate, + split.clone(), + sender, + )) + .expect("Can send event to event store thread."); + + match receiver.recv() { + Ok(value) => value, + Err(err) => panic!("Event store thread crashed: {err:?}"), + } + } +} + +#[derive(new)] +struct WorkerThread { + store: S, + receiver: mpsc::Receiver, +} + +impl WorkerThread +where + C: EventStore, +{ + fn run(mut self) { + for item in self.receiver.iter() { + match item { + Message::End => { + return; + } + Message::FindEpoch(name, aggregate, direction, split, callback) => { + let response = self.store.find_epoch(&name, aggregate, direction, &split); + callback + .send(response) + .expect("Can send response using callback channel."); + } + Message::FindMetric(name, epoch, aggregate, split, callback) => { + let response = self.store.find_metric(&name, epoch, aggregate, &split); + callback + .send(response) + .expect("Can send response using callback channel."); + } + Message::OnEventTrain(event) => self.store.add_event(event, Split::Train), + Message::OnEventValid(event) => self.store.add_event(event, Split::Valid), + Message::OnEventTest(event, tag) => { + self.store.add_event(event, Split::Test(Some(tag))) + } + } + } + } +} + +enum Message { + OnEventTest(Event, Arc), + OnEventTrain(Event), + OnEventValid(Event), + End, + FindEpoch( + String, + Aggregate, + Direction, + Split, + mpsc::SyncSender>, + ), + FindMetric( + String, + usize, + Aggregate, + Split, + mpsc::SyncSender>, + ), +} + +impl Drop for EventStoreClient { + fn drop(&mut self) { + self.sender + .send(Message::End) + .expect("Can send the end message to the event store thread."); + let handler = self.handler.take(); + + if let Some(handler) = handler { + handler.join().expect("The event store thread should stop."); + } + } +} diff --git a/crates/burn-train/src/metric/store/log.rs b/crates/burn-train/src/metric/store/log.rs new file mode 100644 index 0000000..153cc1f --- /dev/null +++ b/crates/burn-train/src/metric/store/log.rs @@ -0,0 +1,75 @@ +use std::collections::HashMap; + +use super::{Aggregate, Direction, Event, EventStore, Split, aggregate::NumericMetricsAggregate}; +use crate::logger::MetricLogger; + +#[derive(Default)] +pub(crate) struct LogEventStore { + loggers: Vec>, + aggregate: NumericMetricsAggregate, + epochs: HashMap, +} + +impl EventStore for LogEventStore { + fn add_event(&mut self, event: Event, split: Split) { + let epoch = *self.epochs.entry(split.clone()).or_insert(1); + + match event { + Event::StartSplit(epoch) => { + self.epochs.insert(split, epoch); + } + Event::MetricsInit(definitions) => { + self.aggregate.register_definitions(&definitions); + definitions.iter().for_each(|def| { + self.loggers + .iter_mut() + .for_each(|logger| logger.log_metric_definition(def.clone())); + }); + } + Event::MetricsUpdate(update) => { + self.loggers + .iter_mut() + .for_each(|logger| logger.log(update.clone(), epoch, &split)); + } + Event::EndEpoch(summary) => { + self.loggers + .iter_mut() + .for_each(|logger| logger.log_epoch_summary(summary.clone())); + } + } + } + + fn find_epoch( + &mut self, + name: &str, + aggregate: Aggregate, + direction: Direction, + split: &Split, + ) -> Option { + self.aggregate + .find_epoch(name, split, aggregate, direction, &mut self.loggers) + } + + fn find_metric( + &mut self, + name: &str, + epoch: usize, + aggregate: Aggregate, + split: &Split, + ) -> Option { + self.aggregate + .aggregate(name, epoch, split, aggregate, &mut self.loggers) + } +} + +impl LogEventStore { + /// Register a logger for metrics. + pub(crate) fn register_logger(&mut self, logger: ML) { + self.loggers.push(Box::new(logger)); + } + + /// Returns whether any loggers are registered. + pub(crate) fn has_loggers(&self) -> bool { + !self.loggers.is_empty() + } +} diff --git a/crates/burn-train/src/metric/store/mod.rs b/crates/burn-train/src/metric/store/mod.rs new file mode 100644 index 0000000..b86c0f4 --- /dev/null +++ b/crates/burn-train/src/metric/store/mod.rs @@ -0,0 +1,9 @@ +pub(crate) mod aggregate; + +mod base; +mod client; +mod log; + +pub(crate) use self::log::*; +pub use base::*; +pub use client::*; diff --git a/crates/burn-train/src/metric/top_k_acc.rs b/crates/burn-train/src/metric/top_k_acc.rs new file mode 100644 index 0000000..bfd7252 --- /dev/null +++ b/crates/burn-train/src/metric/top_k_acc.rs @@ -0,0 +1,176 @@ +use std::sync::Arc; + +use super::state::{FormatOptions, NumericMetricState}; +use super::{MetricMetadata, SerializedEntry}; +use crate::metric::{ + Metric, MetricAttributes, MetricName, Numeric, NumericAttributes, NumericEntry, +}; +use burn_core::tensor::{Int, Tensor}; + +/// The Top-K accuracy metric. +/// +/// For K=1, this is equivalent to the [accuracy metric](`super::acc::AccuracyMetric`). +#[derive(Default, Clone)] +pub struct TopKAccuracyMetric { + name: Arc, + k: usize, + state: NumericMetricState, + /// If specified, targets equal to this value will be considered padding and will not count + /// towards the metric + pad_token: Option, +} + +/// The [top-k accuracy metric](TopKAccuracyMetric) input type. +#[derive(new)] +pub struct TopKAccuracyInput { + /// The outputs (batch_size, num_classes) + outputs: Tensor<2>, + /// The labels (batch_size) + targets: Tensor<1, Int>, +} + +impl TopKAccuracyMetric { + /// Creates the metric. + pub fn new(k: usize) -> Self { + Self { + name: Arc::new(format!("Top-K Accuracy @ TopK({})", k)), + k, + ..Default::default() + } + } + + /// Sets the pad token. + pub fn with_pad_token(mut self, index: usize) -> Self { + self.pad_token = Some(index); + self + } +} + +impl Metric for TopKAccuracyMetric { + type Input = TopKAccuracyInput; + + fn update(&mut self, input: &TopKAccuracyInput, _metadata: &MetricMetadata) -> SerializedEntry { + let [batch_size, _n_classes] = input.outputs.dims(); + + let targets = input.targets.clone(); + + let outputs = input + .outputs + .clone() + .argsort_descending(1) + .narrow(1, 0, self.k) + .reshape([batch_size, self.k]); + + let (targets, num_pad) = match self.pad_token { + Some(pad_token) => { + // we ignore the samples where the target is equal to the pad token + let mask = targets.clone().equal_elem(pad_token as i64); + let num_pad = mask.clone().int().sum().into_scalar::(); + (targets.clone().mask_fill(mask, -1_i64), num_pad) + } + None => (targets.clone(), 0_f64), + }; + + let accuracy = targets + .reshape([batch_size, 1]) + .repeat_dim(1, self.k) + .equal(outputs) + .int() + .sum() + .into_scalar::() + / (batch_size as f64 - num_pad); + + self.state.update( + 100.0 * accuracy, + batch_size, + FormatOptions::new(self.name()).unit("%").precision(2), + ) + } + + fn clear(&mut self) { + self.state.reset() + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: Some("%".to_string()), + higher_is_better: true, + ..Default::default() + } + .into() + } +} + +impl Numeric for TopKAccuracyMetric { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_accuracy_without_padding() { + let device = Default::default(); + let mut metric = TopKAccuracyMetric::new(2); + let input = TopKAccuracyInput::new( + Tensor::from_data( + [ + [0.0, 0.2, 0.8], // 2, 1 + [1.0, 2.0, 0.5], // 1, 0 + [0.4, 0.1, 0.2], // 0, 2 + [0.6, 0.7, 0.2], // 1, 0 + ], + &device, + ), + Tensor::from_data([2, 2, 1, 1], &device), + ); + + let _entry = metric.update(&input, &MetricMetadata::fake()); + assert_eq!(50.0, metric.value().current()); + } + + #[test] + fn test_accuracy_with_padding() { + let device = Default::default(); + let mut metric = TopKAccuracyMetric::new(2).with_pad_token(3); + let input = TopKAccuracyInput::new( + Tensor::from_data( + [ + [0.0, 0.2, 0.8, 0.0], // 2, 1 + [1.0, 2.0, 0.5, 0.0], // 1, 0 + [0.4, 0.1, 0.2, 0.0], // 0, 2 + [0.6, 0.7, 0.2, 0.0], // 1, 0 + [0.0, 0.1, 0.2, 5.0], // Predicted padding should not count + [0.0, 0.1, 0.2, 0.0], // Error on padding should not count + [0.6, 0.0, 0.2, 0.0], // Error on padding should not count + ], + &device, + ), + Tensor::from_data([2, 2, 1, 1, 3, 3, 3], &device), + ); + + let _entry = metric.update(&input, &MetricMetadata::fake()); + assert_eq!(50.0, metric.value().current()); + } + + #[test] + fn test_parameterized_unique_name() { + let metric_a = TopKAccuracyMetric::new(2); + let metric_b = TopKAccuracyMetric::new(1); + let metric_c = TopKAccuracyMetric::new(2); + + assert_ne!(metric_a.name(), metric_b.name()); + assert_eq!(metric_a.name(), metric_c.name()); + } +} diff --git a/crates/burn-train/src/metric/vision/afine/calibrators.rs b/crates/burn-train/src/metric/vision/afine/calibrators.rs new file mode 100644 index 0000000..38a725d --- /dev/null +++ b/crates/burn-train/src/metric/vision/afine/calibrators.rs @@ -0,0 +1,246 @@ +//! Calibrators, adapter, and final-score mapping for A-FINE. +//! +//! Four small modules sit between the heads and the final per-sample +//! quality score: +//! +//! - [`NrCalibrator`] — logistic mapping of the naturalness head's raw +//! output into `(-2, 2)`. Two learnable scalars. +//! - [`FrCalibratorWithLimit`] — logistic mapping of the fidelity head's +//! raw output into `(-2, 2)`, with `yita3` clamped to `[0.05, 0.95]` +//! and `yita4` to `[0.01, 0.70]` on every forward. +//! - [`AfineAdapter`] — `D = exp(softplus(k) * (N_ref - N_dis)) * N_dis + F`. +//! Single learnable scalar `k`. +//! - [`scale_finalscore`] — fixed logistic into `(0, 100)` with the +//! paper-reported constants. +//! +//! All three calibrators implement the same logistic shape: +//! `out = (yita1 - yita2) * sigmoid((x - yita3) / (|yita4| + eps)) + yita2`. +//! This is the algebraic equivalent of PyIQA's two-branch +//! `if exp_pow >= 10` formulation, rewritten as a single expression so +//! it batches correctly. PyIQA's branch only works on 0-D scalar +//! tensors. + +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Module, Param}; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn::tensor::activation::{sigmoid, softplus}; + +const NR_YITA1: f64 = 2.0; +const NR_YITA2: f64 = -2.0; +const NR_YITA3_INIT: f32 = 4.9592; +const NR_YITA4_INIT: f32 = 21.5968; + +const FR_YITA1: f64 = 2.0; +const FR_YITA2: f64 = -2.0; +const FR_YITA3_INIT: f32 = 0.5; +const FR_YITA4_INIT: f32 = 0.15; +const FR_YITA3_MIN: f32 = 0.05; +const FR_YITA3_MAX: f32 = 0.95; +const FR_YITA4_MIN: f32 = 0.01; +const FR_YITA4_MAX: f32 = 0.70; + +const ADAPTER_K_INIT: f32 = 5.0; + +const SCALE_YITA1: f64 = 100.0; +const SCALE_YITA2: f64 = 0.0; +const SCALE_YITA3: f64 = -1.971_0; +const SCALE_YITA4: f64 = -2.373_4; + +/// Numerical-stability epsilon in the logistic denominator. Matches +/// PyIQA exactly; do not change without coordinating a parity-test +/// re-capture. +const EPS: f64 = 1e-10; + +/// Apply `(yita1 - yita2) * sigmoid((x - yita3) / (|yita4| + eps)) + yita2` +/// element-wise, broadcasting the 1-D scalar parameters over the input. +fn logistic_calibrate( + x: Tensor<2>, + yita3: Tensor<1>, + yita4_abs: Tensor<1>, + yita1: f64, + yita2: f64, +) -> Tensor<2> { + let yita3 = yita3.reshape([1, 1]); + let denom = yita4_abs.reshape([1, 1]).add_scalar(EPS); + let inner = (x - yita3) / denom; + sigmoid(inner).mul_scalar(yita1 - yita2).add_scalar(yita2) +} + +/// Configuration for [`NrCalibrator`]. +#[derive(Config, Debug)] +pub(crate) struct NrCalibratorConfig {} + +impl NrCalibratorConfig { + pub(crate) fn init(&self, device: &Device) -> NrCalibrator { + NrCalibrator { + yita3: Param::from_tensor(Tensor::from_floats([NR_YITA3_INIT], device)), + yita4: Param::from_tensor(Tensor::from_floats([NR_YITA4_INIT], device)), + } + } +} + +/// Naturalness logistic calibrator. Maps `[B, 1]` into `(-2, 2)`. +#[derive(Module, Debug)] +pub(crate) struct NrCalibrator { + pub(crate) yita3: Param>, + pub(crate) yita4: Param>, +} + +impl NrCalibrator { + pub(crate) fn forward(&self, x: Tensor<2>) -> Tensor<2> { + logistic_calibrate( + x, + self.yita3.val(), + self.yita4.val().abs(), + NR_YITA1, + NR_YITA2, + ) + } +} + +/// Configuration for [`FrCalibratorWithLimit`]. +#[derive(Config, Debug)] +pub(crate) struct FrCalibratorWithLimitConfig {} + +impl FrCalibratorWithLimitConfig { + pub(crate) fn init(&self, device: &Device) -> FrCalibratorWithLimit { + FrCalibratorWithLimit { + yita3: Param::from_tensor(Tensor::from_floats([FR_YITA3_INIT], device)), + yita4: Param::from_tensor(Tensor::from_floats([FR_YITA4_INIT], device)), + } + } +} + +/// Fidelity logistic calibrator with on-forward clamping of `yita3` and +/// `yita4`. PyIQA clamps the values used in the formula on every call; +/// the stored parameter is unchanged. +#[derive(Module, Debug)] +pub(crate) struct FrCalibratorWithLimit { + pub(crate) yita3: Param>, + pub(crate) yita4: Param>, +} + +impl FrCalibratorWithLimit { + pub(crate) fn forward(&self, x: Tensor<2>) -> Tensor<2> { + // Match PyIQA semantics exactly: clamp first, then abs. The + // clamp range is positive so the abs is a no-op for in-range + // values, but for an out-of-range checkpoint or a parameter + // that drifts negative during training the order matters. + let yita3 = self.yita3.val().clamp(FR_YITA3_MIN, FR_YITA3_MAX); + let yita4 = self.yita4.val().clamp(FR_YITA4_MIN, FR_YITA4_MAX); + logistic_calibrate(x, yita3, yita4.abs(), FR_YITA1, FR_YITA2) + } +} + +/// Configuration for [`AfineAdapter`]. +#[derive(Config, Debug)] +pub(crate) struct AfineAdapterConfig {} + +impl AfineAdapterConfig { + pub(crate) fn init(&self, device: &Device) -> AfineAdapter { + AfineAdapter { + k: Param::from_tensor(Tensor::from_floats([ADAPTER_K_INIT], device)), + } + } +} + +/// Fuses the calibrated naturalness and fidelity scores into a single +/// raw `D` value. +/// +/// `D = exp(softplus(k) * (N_ref - N_dis)) * N_dis + F`. The `softplus` +/// wrapper enforces `k > 0` without constraining the stored parameter. +#[derive(Module, Debug)] +pub(crate) struct AfineAdapter { + pub(crate) k: Param>, +} + +impl AfineAdapter { + pub(crate) fn forward( + &self, + x_nr: Tensor<2>, + ref_nr: Tensor<2>, + xref_fr: Tensor<2>, + ) -> Tensor<2> { + let k_pos = softplus(self.k.val(), 1.0).reshape([1, 1]); + let weight = (k_pos * (ref_nr - x_nr.clone())).exp(); + weight * x_nr + xref_fr + } +} + +/// Map a raw adapter score into `(0, 100)` via a fixed 4-parameter +/// logistic. Constants are the paper-reported defaults. +pub(crate) fn scale_finalscore(score: Tensor<2>) -> Tensor<2> { + let denom = SCALE_YITA4.abs() + EPS; + let inner = score.sub_scalar(SCALE_YITA3).div_scalar(denom); + sigmoid(inner) + .mul_scalar(SCALE_YITA1 - SCALE_YITA2) + .add_scalar(SCALE_YITA2) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn nr_calibrator_maps_to_bounded_range() { + let device = Default::default(); + let calibrator = NrCalibratorConfig::new().init(&device); + + let extremes = Tensor::<2>::from_floats([[-1000.0], [0.0], [1000.0]], &device); + let out = calibrator.forward(extremes); + let values = out.into_data().to_vec::().unwrap(); + + for v in &values { + assert!(*v >= -2.0 && *v <= 2.0, "out-of-range value: {v}"); + } + // Monotonic increasing. + assert!(values[0] < values[1]); + assert!(values[1] < values[2]); + } + + #[test] + fn fr_calibrator_clamp_does_not_panic() { + let device = Default::default(); + let calibrator = FrCalibratorWithLimitConfig::new().init(&device); + + let input = Tensor::<2>::from_floats([[0.5], [1.5], [-0.5]], &device); + let out = calibrator.forward(input); + + assert_eq!(out.dims(), [3, 1]); + } + + #[test] + fn adapter_forward_propagates_shape() { + let device = Default::default(); + let adapter = AfineAdapterConfig::new().init(&device); + + let nr_dis = Tensor::<2>::from_floats([[0.5], [-0.3]], &device); + let nr_ref = Tensor::<2>::from_floats([[0.7], [-0.1]], &device); + let fr = Tensor::<2>::from_floats([[0.2], [0.4]], &device); + + let out = adapter.forward(nr_dis, nr_ref, fr); + assert_eq!(out.dims(), [2, 1]); + } + + #[test] + fn scale_finalscore_maps_to_0_100_range() { + let device = Default::default(); + + let scores = Tensor::<2>::from_floats([[-1000.0], [-1.971], [1000.0]], &device); + let out = scale_finalscore(scores); + let values = out.into_data().to_vec::().unwrap(); + + assert!(values[0] >= 0.0 && values[0] <= 100.0); + assert!(values[2] >= 0.0 && values[2] <= 100.0); + // At yita3 = -1.971 the sigmoid argument is 0, so the output is + // 100 * 0.5 = 50. + assert!( + (values[1] - 50.0).abs() < 0.5, + "midpoint should be ~50, got {}", + values[1] + ); + } +} diff --git a/crates/burn-train/src/metric/vision/afine/clip_attention.rs b/crates/burn-train/src/metric/vision/afine/clip_attention.rs new file mode 100644 index 0000000..4b36b3d --- /dev/null +++ b/crates/burn-train/src/metric/vision/afine/clip_attention.rs @@ -0,0 +1,134 @@ +//! Self-attention block matching PyTorch's `nn.MultiheadAttention` wire +//! format. +//! +//! PyTorch stores Q/K/V as a single fused `in_proj_weight` of shape +//! `[3 * d_model, d_model]` and `in_proj_bias` of shape `[3 * d_model]`. +//! Burn's [`burn_nn::attention::MultiHeadAttention`] uses three separate +//! Linear layers, so the CLIP checkpoint cannot map to it without +//! pre-splitting the weights at load time. +//! +//! This module keeps the fused layout: a single `qkv_proj` Linear of +//! shape `(d_model -> 3 * d_model)` and a `chunk(3, -1)` at forward, +//! giving a one-to-one mapping with the checkpoint. +//! +//! No attention mask is supported. CLIP's image encoder runs +//! unconditional self-attention; the text encoder (which uses a causal +//! mask) is not ported. + +use burn_core as burn; + +use burn::config::Config; +use burn::module::Module; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn::tensor::activation::softmax; +use burn_nn::{Linear, LinearConfig}; + +/// Configuration for [`ClipQkvAttention`]. +#[derive(Config, Debug)] +pub(crate) struct ClipQkvAttentionConfig { + /// Embedding dimension. Must be divisible by `n_heads`. + pub d_model: usize, + /// Number of attention heads. + pub n_heads: usize, +} + +impl ClipQkvAttentionConfig { + /// Initialize a [`ClipQkvAttention`] block. + pub(crate) fn init(&self, device: &Device) -> ClipQkvAttention { + assert_eq!( + self.d_model % self.n_heads, + 0, + "d_model ({}) must be divisible by n_heads ({})", + self.d_model, + self.n_heads + ); + let head_dim = self.d_model / self.n_heads; + ClipQkvAttention { + qkv_proj: LinearConfig::new(self.d_model, 3 * self.d_model) + .with_bias(true) + .init(device), + out_proj: LinearConfig::new(self.d_model, self.d_model) + .with_bias(true) + .init(device), + d_model: self.d_model, + n_heads: self.n_heads, + head_dim, + } + } +} + +/// Self-attention with a fused QKV projection, matching CLIP's checkpoint +/// layout one-to-one. +#[derive(Module, Debug)] +pub(crate) struct ClipQkvAttention { + /// Fused projection `d_model -> 3 * d_model` for Q, K, V. + pub(crate) qkv_proj: Linear, + /// Output projection `d_model -> d_model`. + pub(crate) out_proj: Linear, + /// Embedding dimension. + pub(crate) d_model: usize, + /// Number of attention heads. + pub(crate) n_heads: usize, + /// Per-head dimension (`d_model / n_heads`). + pub(crate) head_dim: usize, +} + +impl ClipQkvAttention { + /// Apply self-attention. Input and output shape: `[batch, seq, d_model]`. + pub(crate) fn forward(&self, x: Tensor<3>) -> Tensor<3> { + let [batch, seq, _] = x.dims(); + + let qkv = self.qkv_proj.forward(x); + let mut chunks = qkv.chunk(3, 2); + let value = chunks.remove(2); + let key = chunks.remove(1); + let query = chunks.remove(0); + + let to_heads = |t: Tensor<3>| { + t.reshape([batch, seq, self.n_heads, self.head_dim]) + .swap_dims(1, 2) + }; + let query = to_heads(query); + let key = to_heads(key); + let value = to_heads(value); + + let scale = (self.head_dim as f32).sqrt(); + let scores = query.matmul(key.transpose()).div_scalar(scale); + let weights = softmax(scores, 3); + + let context = weights + .matmul(value) + .swap_dims(1, 2) + .reshape([batch, seq, self.d_model]); + self.out_proj.forward(context) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Distribution; + + #[test] + fn clip_qkv_attention_preserves_shape() { + let device = Default::default(); + let attn = ClipQkvAttentionConfig::new(768, 12).init(&device); + + let input = Tensor::<3>::random([1, 50, 768], Distribution::Default, &device); + let output = attn.forward(input); + + assert_eq!(output.dims(), [1, 50, 768]); + } + + #[test] + fn clip_qkv_attention_handles_batch() { + let device = Default::default(); + let attn = ClipQkvAttentionConfig::new(64, 4).init(&device); + + let input = Tensor::<3>::random([3, 16, 64], Distribution::Default, &device); + let output = attn.forward(input); + + assert_eq!(output.dims(), [3, 16, 64]); + } +} diff --git a/crates/burn-train/src/metric/vision/afine/clip_vit.rs b/crates/burn-train/src/metric/vision/afine/clip_vit.rs new file mode 100644 index 0000000..8f4b35c --- /dev/null +++ b/crates/burn-train/src/metric/vision/afine/clip_vit.rs @@ -0,0 +1,362 @@ +//! CLIP ViT-B/32 image encoder. +//! +//! Structurally a port of OpenAI CLIP's `VisionTransformer` minus the +//! text-side and joint-embedding projection. A-FINE consumes the +//! per-layer patch features (twelve `[batch, num_patches, 768]` tensors), +//! not the final CLIP embedding, so [`Self::forward_with_features`] is +//! the entry point used by the metric. [`Self::forward`] returns just the +//! `ln_post`-normalized class-token embedding and is provided for +//! completeness. +//! +//! Layout: NLD (`[batch, seq, embed]`) end-to-end. The PyTorch reference +//! permutes to LND (`[seq, batch, embed]`) inside the transformer stack +//! and back again before `ln_post`; attention is invariant to that swap, +//! so we stay in NLD throughout. + +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Initializer, Module, Param}; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn_nn::conv::{Conv2d, Conv2dConfig}; +use burn_nn::{LayerNorm, LayerNormConfig, Linear, LinearConfig}; + +use super::clip_attention::{ClipQkvAttention, ClipQkvAttentionConfig}; +use super::quick_gelu::QuickGelu; + +/// Configuration for [`ClipVisualEncoder`]. +#[derive(Config, Debug)] +pub struct ClipVisualEncoderConfig { + /// Input image channels. Defaults to 3 (RGB). + #[config(default = "3")] + pub in_channels: usize, + /// Token embedding dimension. Defaults to 768 (CLIP ViT-B/32). + #[config(default = "768")] + pub embed_dim: usize, + /// Patch (and conv stride) size. Defaults to 32 (CLIP ViT-B/32). + #[config(default = "32")] + pub patch_size: usize, + /// Number of transformer blocks. Defaults to 12. + #[config(default = "12")] + pub num_layers: usize, + /// Number of attention heads per block. Defaults to 12. + #[config(default = "12")] + pub num_heads: usize, + /// Hidden dimension of the per-block MLP. Defaults to 3072 (4 * 768). + #[config(default = "3072")] + pub mlp_dim: usize, + /// Expected input resolution. Defaults to 256. Must be a multiple of `patch_size`. + #[config(default = "256")] + pub image_size: usize, +} + +impl ClipVisualEncoderConfig { + /// Initialize a [`ClipVisualEncoder`] with random weights. + pub fn init(&self, device: &Device) -> ClipVisualEncoder { + assert_eq!( + self.image_size % self.patch_size, + 0, + "image_size ({}) must be a multiple of patch_size ({})", + self.image_size, + self.patch_size + ); + + let num_patches = (self.image_size / self.patch_size).pow(2); + let seq_len = num_patches + 1; // +1 for the prepended CLS token. + + let patch_embed = Conv2dConfig::new( + [self.in_channels, self.embed_dim], + [self.patch_size, self.patch_size], + ) + .with_stride([self.patch_size, self.patch_size]) + .with_bias(false) + .init(device); + + // Standard transformer init magnitude. Final values get overwritten + // when pretrained weights load; this just keeps random-init forward + // passes well-conditioned for tests. + let init = Initializer::Normal { + mean: 0.0, + std: 0.02, + }; + let class_token = init.init([self.embed_dim], device); + let positional_embedding = init.init([seq_len, self.embed_dim], device); + + let blocks = (0..self.num_layers) + .map(|_| { + TransformerBlockConfig::new(self.embed_dim, self.num_heads, self.mlp_dim) + .init(device) + }) + .collect(); + + ClipVisualEncoder { + patch_embed, + class_token, + positional_embedding, + ln_pre: LayerNormConfig::new(self.embed_dim).init(device), + blocks, + ln_post: LayerNormConfig::new(self.embed_dim).init(device), + embed_dim: self.embed_dim, + patch_size: self.patch_size, + image_size: self.image_size, + } + } +} + +/// Output of [`ClipVisualEncoder::forward_with_features`]. +/// +/// `features` are the per-block patch tokens A-FINE consumes. `cls` is +/// the post-`ln_post` class-token embedding, present only when the +/// caller requests it; the metric path skips the slice + LayerNorm to +/// avoid wasted work. +#[derive(Debug)] +pub struct ClipOutput { + /// Per-block patch tokens, length `num_layers`, each + /// `[batch, num_patches, embed_dim]`. + pub features: Vec>, + /// Post-`ln_post` class-token embedding, `[batch, embed_dim]`. `Some` + /// when `forward_with_features` is called with `return_cls = true`. + pub cls: Option>, +} + +/// CLIP ViT-B/32 image encoder. +/// +/// Consumes a normalized RGB image and returns either the class-token +/// embedding ([`Self::forward`]) or the per-layer patch-token features +/// used by A-FINE ([`Self::forward_with_features`]). +#[derive(Module, Debug)] +pub struct ClipVisualEncoder { + /// Patch-embedding convolution: `(in_channels, embed_dim)`, kernel and + /// stride both `patch_size`, no bias (matches CLIP). + pub(crate) patch_embed: Conv2d, + /// Learnable class token, shape `[embed_dim]`. Stored 1-D to match the + /// CLIP checkpoint key `class_embedding`; reshaped to `[1, 1, embed]` + /// at forward time. + pub(crate) class_token: Param>, + /// Learnable positional embedding, shape `[seq_len, embed_dim]`. + pub(crate) positional_embedding: Param>, + /// Pre-stack LayerNorm. + pub(crate) ln_pre: LayerNorm, + /// Stack of transformer blocks. + pub(crate) blocks: Vec, + /// Post-stack LayerNorm applied to the class token only. + pub(crate) ln_post: LayerNorm, + + pub(crate) embed_dim: usize, + pub(crate) patch_size: usize, + pub(crate) image_size: usize, +} + +impl ClipVisualEncoder { + /// Encode an image to its class-token embedding. + /// + /// # Shapes + /// - input: `[batch, in_channels, height, width]` + /// - output: `[batch, embed_dim]` + pub fn forward(&self, image: Tensor<4>) -> Tensor<2> { + self.forward_with_features(image, true) + .cls + .expect("cls requested") + } + + /// Encode an image and return the patch-token features after each + /// transformer block. The class token is stripped from each level. + /// The post-`ln_post` cls embedding is computed only when + /// `return_cls` is true; A-FINE itself does not use it. + /// + /// # Shapes + /// - input: `[batch, in_channels, height, width]` + /// - features: `Vec` of length `num_layers`, each + /// `[batch, num_patches, embed_dim]` + /// - cls (when present): `[batch, embed_dim]` + pub fn forward_with_features(&self, image: Tensor<4>, return_cls: bool) -> ClipOutput { + let [batch, _, height, width] = image.dims(); + assert_eq!( + height % self.patch_size, + 0, + "image height ({}) must be a multiple of patch_size ({})", + height, + self.patch_size + ); + assert_eq!( + width % self.patch_size, + 0, + "image width ({}) must be a multiple of patch_size ({})", + width, + self.patch_size + ); + + let embed = self.embed_dim; + + // Patch-embed: [B, C, H, W] -> [B, embed, H/p, W/p]. + let x = self.patch_embed.forward(image); + let [_, _, h_out, w_out] = x.dims(); + let num_patches = h_out * w_out; + let seq_len = num_patches + 1; + + // Flatten spatial axes, then put the patch axis before the channel + // axis: [B, embed, N] -> [B, N, embed]. + let x = x.reshape([batch, embed, num_patches]).swap_dims(1, 2); + + // Prepend the CLS token. Reshape [embed] -> [1, 1, embed], then + // broadcast-expand to [B, 1, embed] before concatenating along the + // sequence axis. + let cls = self + .class_token + .val() + .reshape([1, 1, embed]) + .expand([batch, 1, embed]); + let x = Tensor::cat(vec![cls, x], 1); + + // Add positional embedding. The pos-embed param is initialised to + // match `seq_len` for the configured image_size; runtime mismatches + // are caught here. Bicubic resize of a checkpointed pos-embed (e.g. + // 50 -> 65 when loading a 224-trained checkpoint at 256) is handled + // in the weights loader, not here. + let pos_seq = self.positional_embedding.dims()[0]; + assert_eq!( + pos_seq, seq_len, + "positional_embedding length {} does not match runtime sequence length {}", + pos_seq, seq_len + ); + let pos = self.positional_embedding.val().reshape([1, seq_len, embed]); + let x = x + pos; + + let mut x = self.ln_pre.forward(x); + + // Run the transformer stack, collecting patch features (CLS + // dropped) after each block. A-FINE consumes all twelve. + let mut features = Vec::with_capacity(self.blocks.len()); + for block in &self.blocks { + x = block.forward(x); + let patch_features = x.clone().slice([0..batch, 1..seq_len, 0..embed]); + features.push(patch_features); + } + + // ln_post is applied only to the class-token output. Skip both + // the slice and the LayerNorm when the caller doesn't need it. + let cls = return_cls.then(|| { + let cls = x.slice([0..batch, 0..1, 0..embed]).reshape([batch, embed]); + self.ln_post.forward(cls) + }); + + ClipOutput { features, cls } + } +} + +/// Configuration for a single transformer block. +#[derive(Config, Debug)] +pub(crate) struct TransformerBlockConfig { + pub d_model: usize, + pub n_heads: usize, + pub mlp_dim: usize, +} + +impl TransformerBlockConfig { + pub(crate) fn init(&self, device: &Device) -> TransformerBlock { + TransformerBlock { + ln_1: LayerNormConfig::new(self.d_model).init(device), + attn: ClipQkvAttentionConfig::new(self.d_model, self.n_heads).init(device), + ln_2: LayerNormConfig::new(self.d_model).init(device), + mlp: TransformerMlpConfig::new(self.d_model, self.mlp_dim).init(device), + } + } +} + +/// Pre-norm transformer block: `x = x + attn(ln_1(x))`, then +/// `x = x + mlp(ln_2(x))`. Matches CLIP's encoder block exactly. +#[derive(Module, Debug)] +pub(crate) struct TransformerBlock { + pub(crate) ln_1: LayerNorm, + pub(crate) attn: ClipQkvAttention, + pub(crate) ln_2: LayerNorm, + pub(crate) mlp: TransformerMlp, +} + +impl TransformerBlock { + pub(crate) fn forward(&self, x: Tensor<3>) -> Tensor<3> { + let attn_out = self.attn.forward(self.ln_1.forward(x.clone())); + let x = x + attn_out; + let mlp_out = self.mlp.forward(self.ln_2.forward(x.clone())); + x + mlp_out + } +} + +/// Configuration for [`TransformerMlp`]. +#[derive(Config, Debug)] +pub(crate) struct TransformerMlpConfig { + pub d_model: usize, + pub mlp_dim: usize, +} + +impl TransformerMlpConfig { + pub(crate) fn init(&self, device: &Device) -> TransformerMlp { + TransformerMlp { + c_fc: LinearConfig::new(self.d_model, self.mlp_dim) + .with_bias(true) + .init(device), + activation: QuickGelu, + c_proj: LinearConfig::new(self.mlp_dim, self.d_model) + .with_bias(true) + .init(device), + } + } +} + +/// Two-layer MLP with [`QuickGelu`] in between. CLIP's `c_fc` and `c_proj` +/// names are preserved so the checkpoint maps without remapping. +#[derive(Module, Debug)] +pub(crate) struct TransformerMlp { + pub(crate) c_fc: Linear, + pub(crate) activation: QuickGelu, + pub(crate) c_proj: Linear, +} + +impl TransformerMlp { + pub(crate) fn forward(&self, x: Tensor<3>) -> Tensor<3> { + let x = self.c_fc.forward(x); + let x = self.activation.forward(x); + self.c_proj.forward(x) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Distribution; + + #[test] + fn clip_visual_encoder_forward_shape() { + let device = Default::default(); + let encoder = ClipVisualEncoderConfig::new() + .with_image_size(64) + .with_num_layers(2) + .init(&device); + + let image = Tensor::<4>::random([1, 3, 64, 64], Distribution::Default, &device); + let cls = encoder.forward(image); + + assert_eq!(cls.dims(), [1, 768]); + } + + #[test] + fn clip_visual_encoder_forward_with_features_shape() { + let device = Default::default(); + let encoder = ClipVisualEncoderConfig::new() + .with_image_size(64) + .with_num_layers(3) + .init(&device); + + let image = Tensor::<4>::random([2, 3, 64, 64], Distribution::Default, &device); + let ClipOutput { features, cls } = encoder.forward_with_features(image, true); + let cls = cls.expect("cls requested"); + + assert_eq!(cls.dims(), [2, 768]); + assert_eq!(features.len(), 3); + // 64x64 image at patch_size=32 -> 2x2 = 4 patches. + for level in &features { + assert_eq!(level.dims(), [2, 4, 768]); + } + } +} diff --git a/crates/burn-train/src/metric/vision/afine/heads.rs b/crates/burn-train/src/metric/vision/afine/heads.rs new file mode 100644 index 0000000..43eb55d --- /dev/null +++ b/crates/burn-train/src/metric/vision/afine/heads.rs @@ -0,0 +1,377 @@ +//! Naturalness and fidelity heads. +//! +//! Both heads consume the 12 per-block CLIP feature maps plus the raw +//! RGB image (re-introduced as a "level 0" feature). The heads' `mean` +//! and `std` buffers are CLIP's normalization constants, used here as +//! the **inverse** of preprocessing — they de-normalize the input back +//! into pixel space before computing raw-RGB statistics. +//! +//! - [`AfineQHead`] — naturalness, single-image, returns `[B, 1]`. +//! - [`AfineDHead`] — fidelity, two-image, returns `[B, 1]`. +//! +//! The `chns` tuple has 13 entries: `[3, 768, 768, ..., 768]`. The +//! leading 3 is the raw image; the remaining twelve 768s are the CLIP +//! transformer outputs after each block. + +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Module, Param}; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn::tensor::activation::{relu, softplus}; +use burn_nn::{Gelu, Linear, LinearConfig}; + +/// CLIP RGB normalization mean (per-channel). Used as a de-normalization +/// constant inside the heads. +const CLIP_MEAN: [f32; 3] = [0.481_454_66, 0.457_827_5, 0.408_210_73]; +/// CLIP RGB normalization std. +const CLIP_STD: [f32; 3] = [0.268_629_54, 0.261_302_58, 0.275_777_11]; + +/// Per-level channel counts: raw RGB (3) plus 12 CLIP layers (768 each). +const CHNS: [usize; 13] = [ + 3, 768, 768, 768, 768, 768, 768, 768, 768, 768, 768, 768, 768, +]; +const NUM_LEVELS: usize = 13; + +/// Hidden dim of the shared per-CLIP-level projection (`proj_feat`). +const Q_HIDDEN_DIM: usize = 128; +/// Output dim of the first MLP layer in `proj_head`. +const Q_PROJ_HEAD_OUT: usize = 768; +/// Input dim of `proj_head`: 6 (raw, k=0) + 128 * 12 (projected, k=1..=12). +const Q_PROJ_HEAD_IN: usize = 6 + Q_HIDDEN_DIM * 12; + +/// Sum of `CHNS`. `alpha` and `beta` are length `D_CHNS_SUM` along the +/// last axis. +const D_CHNS_SUM: usize = 9219; + +/// Numerical-stability epsilon used in the SSIM-like ratios. **Must be +/// 1e-10**, not the larger 1e-6 that DISTS uses; the feature magnitudes +/// after CLIP + ReLU are unit-scale-ish, so DISTS's eps would overwhelm +/// the signal. +const EPS: f64 = 1e-10; + +fn build_clip_mean_std(device: &Device) -> (Tensor<4>, Tensor<4>) { + let mean = Tensor::from_floats( + [[[[CLIP_MEAN[0]]], [[CLIP_MEAN[1]]], [[CLIP_MEAN[2]]]]], + device, + ); + let std = Tensor::from_floats( + [[[[CLIP_STD[0]]], [[CLIP_STD[1]]], [[CLIP_STD[2]]]]], + device, + ); + (mean, std) +} + +/// Mean and biased variance over the token axis, returned as a 2-D +/// tensor `[B, 2 * channels]` (mean followed by var, both flattened). +fn level_mean_var(feat: Tensor<3>) -> (Tensor<3>, Tensor<3>, Tensor<2>) { + let mean = feat.clone().mean_dim(1); + let centered = feat - mean.clone(); + let var = centered.powi_scalar(2).mean_dim(1); + let descriptor = Tensor::cat( + vec![ + mean.clone().flatten::<2>(1, 2), + var.clone().flatten::<2>(1, 2), + ], + 1, + ); + (mean, var, descriptor) +} + +/// Configuration for [`AfineQHead`]. +#[derive(Config, Debug)] +pub struct AfineQHeadConfig {} + +impl AfineQHeadConfig { + /// Initialize the naturalness head with random weights. + pub fn init(&self, device: &Device) -> AfineQHead { + let (mean, std) = build_clip_mean_std(device); + AfineQHead { + mean, + std, + proj_feat: LinearConfig::new(2 * 768, Q_HIDDEN_DIM) + .with_bias(true) + .init(device), + proj_head_fc1: LinearConfig::new(Q_PROJ_HEAD_IN, Q_PROJ_HEAD_OUT) + .with_bias(true) + .init(device), + proj_head_fc2: LinearConfig::new(Q_PROJ_HEAD_OUT, 1) + .with_bias(true) + .init(device), + activation: Gelu::new(), + } + } +} + +/// A-FINE naturalness head. Per-level mean+variance over CLIP tokens, +/// shared projection on the 12 CLIP levels, then a small MLP to a +/// scalar. The level-0 raw-image descriptor (`[B, 6]`) is concatenated +/// directly without going through `proj_feat`. +/// +/// Activation is the erf-based [`burn_nn::Gelu`], not the QuickGELU used +/// elsewhere in this crate. PyIQA's reference is explicit on this. +#[derive(Module, Debug)] +pub struct AfineQHead { + pub(crate) mean: Tensor<4>, + pub(crate) std: Tensor<4>, + pub(crate) proj_feat: Linear, + pub(crate) proj_head_fc1: Linear, + pub(crate) proj_head_fc2: Linear, + pub(crate) activation: Gelu, +} + +impl AfineQHead { + /// Compute the naturalness score for one image plus its CLIP + /// features. + /// + /// # Shapes + /// - `image`: `[B, 3, H, W]`, **already CLIP-normalized**. + /// - `clip_features`: 12 levels of `[B, num_patches, 768]`. + /// - returns: `[B, 1]`. + pub fn forward(&self, image: Tensor<4>, clip_features: &[Tensor<3>]) -> Tensor<2> { + assert_eq!( + clip_features.len(), + 12, + "AfineQHead expects 12 CLIP feature maps, got {}", + clip_features.len() + ); + let [batch, channels, height, width] = image.dims(); + + // De-normalize: x = x * std + mean. + let img = image * self.std.clone() + self.mean.clone(); + + // [B, 3, H, W] -> [B, H*W, 3]: token-major raw-image features. + let img_feat = img + .reshape([batch, channels, height * width]) + .swap_dims(1, 2); + + let mut level_descriptors: Vec> = Vec::with_capacity(NUM_LEVELS); + + // Level 0: raw image, no projection. + let (_, _, raw_descriptor) = level_mean_var(img_feat); + level_descriptors.push(raw_descriptor); + + // Levels 1..=12: ReLU'd CLIP features, then shared `proj_feat`. + for h in clip_features { + let activated = relu(h.clone()); + let (_, _, descriptor) = level_mean_var(activated); + level_descriptors.push(self.proj_feat.forward(descriptor)); + } + + let concat_all = Tensor::cat(level_descriptors, 1); + let hidden = self + .activation + .forward(self.proj_head_fc1.forward(concat_all)); + self.proj_head_fc2.forward(hidden) + } +} + +/// Configuration for [`AfineDHead`]. +#[derive(Config, Debug)] +pub struct AfineDHeadConfig {} + +impl AfineDHeadConfig { + /// Initialize the fidelity head with random weights. + pub fn init(&self, device: &Device) -> AfineDHead { + let (mean, std) = build_clip_mean_std(device); + // PyIQA initializes alpha/beta with `.normal_(0.1, 0.01)`. That + // distribution doesn't really matter under random init — values + // are always overwritten by the checkpoint at `init_pretrained` + // time — but we mirror the magnitude here so the random-init + // forward stays well-conditioned. + let alpha = Tensor::random( + [1, 1, D_CHNS_SUM], + burn::tensor::Distribution::Normal(0.1, 0.01), + device, + ); + let beta = Tensor::random( + [1, 1, D_CHNS_SUM], + burn::tensor::Distribution::Normal(0.1, 0.01), + device, + ); + AfineDHead { + mean, + std, + alpha: Param::from_tensor(alpha), + beta: Param::from_tensor(beta), + } + } +} + +/// A-FINE fidelity head. SSIM-like luminance and contrast statistics +/// across 13 levels, weighted by globally-normalized softplus(alpha) and +/// softplus(beta). +#[derive(Module, Debug)] +pub struct AfineDHead { + pub(crate) mean: Tensor<4>, + pub(crate) std: Tensor<4>, + /// Luminance weights, shape `[1, 1, sum(chns)] = [1, 1, 9219]`. + pub(crate) alpha: Param>, + /// Contrast weights, same shape as `alpha`. + pub(crate) beta: Param>, +} + +impl AfineDHead { + /// Compute the fidelity score between distorted and reference + /// images. + /// + /// # Shapes + /// - `distorted`, `reference`: `[B, 3, H, W]`, both CLIP-normalized. + /// - `feat_dis`, `feat_ref`: 12 levels of `[B, num_patches, 768]`. + /// - returns: `[B, 1]`. + pub fn forward( + &self, + distorted: Tensor<4>, + reference: Tensor<4>, + feat_dis: &[Tensor<3>], + feat_ref: &[Tensor<3>], + ) -> Tensor<2> { + assert_eq!(feat_dis.len(), 12); + assert_eq!(feat_ref.len(), 12); + let [batch, channels, height, width] = distorted.dims(); + + // De-normalize both images and lay them out as token sequences. + let raw_x = (distorted * self.std.clone() + self.mean.clone()) + .reshape([batch, channels, height * width]) + .swap_dims(1, 2); + let raw_y = (reference * self.std.clone() + self.mean.clone()) + .reshape([batch, channels, height * width]) + .swap_dims(1, 2); + + // 13-entry feature lists: raw RGB at level 0, ReLU'd CLIP + // features at levels 1..=12. + let mut feat_x: Vec> = Vec::with_capacity(NUM_LEVELS); + let mut feat_y: Vec> = Vec::with_capacity(NUM_LEVELS); + feat_x.push(raw_x); + feat_y.push(raw_y); + for h in feat_dis { + feat_x.push(relu(h.clone())); + } + for h in feat_ref { + feat_y.push(relu(h.clone())); + } + + // Global softplus normalization. `alpha_/w_sum + beta_/w_sum` + // sums to ~1 across all 13 levels and both terms. We keep + // `w_sum` as a tensor and broadcast-divide so the autodiff graph + // stays intact for any future training-mode caller. + let alpha_sp = softplus(self.alpha.val(), 1.0); + let beta_sp = softplus(self.beta.val(), 1.0); + let w_sum = (alpha_sp.clone().sum() + beta_sp.clone().sum()) + .add_scalar(EPS) + .reshape([1, 1, 1]); + let alpha_norm = alpha_sp / w_sum.clone(); + let beta_norm = beta_sp / w_sum; + + let mut dist1: Option> = None; + let mut dist2: Option> = None; + let mut offset: usize = 0; + + for k in 0..NUM_LEVELS { + let cn = CHNS[k]; + let alpha_k = alpha_norm.clone().slice([0..1, 0..1, offset..offset + cn]); + let beta_k = beta_norm.clone().slice([0..1, 0..1, offset..offset + cn]); + + let xm = feat_x[k].clone().mean_dim(1); // [B, 1, cn] + let ym = feat_y[k].clone().mean_dim(1); + + // S1 (luminance): (2 * x_mean * y_mean + eps) + // / (x_mean^2 + y_mean^2 + eps). + let s1_num = (xm.clone() * ym.clone()).mul_scalar(2.0).add_scalar(EPS); + let s1_den = xm.clone().powi_scalar(2) + ym.clone().powi_scalar(2); + let s1 = s1_num / s1_den.add_scalar(EPS); + let term1 = (alpha_k * s1).sum_dim(2); // [B, 1, 1] + dist1 = Some(match dist1 { + None => term1, + Some(d) => d + term1, + }); + + // S2 (contrast/structure): + // var_x = mean((x - x_mean)^2) (biased), + // var_y = mean((y - y_mean)^2), + // cov_xy = mean(x*y) - x_mean * y_mean, + // S2 = (2 * cov_xy + eps) / (var_x + var_y + eps). + let xc = feat_x[k].clone() - xm.clone(); + let yc = feat_y[k].clone() - ym.clone(); + let x_var = xc.powi_scalar(2).mean_dim(1); + let y_var = yc.powi_scalar(2).mean_dim(1); + let xy_mean = (feat_x[k].clone() * feat_y[k].clone()).mean_dim(1); + let xy_cov = xy_mean - xm * ym; + let s2_num = xy_cov.mul_scalar(2.0).add_scalar(EPS); + let s2_den = (x_var + y_var).add_scalar(EPS); + let s2 = s2_num / s2_den; + let term2 = (beta_k * s2).sum_dim(2); + dist2 = Some(match dist2 { + None => term2, + Some(d) => d + term2, + }); + + offset += cn; + } + + let total = dist1.unwrap() + dist2.unwrap(); // [B, 1, 1] + let score = total.ones_like() - total; // 1 - (dist1 + dist2) + score.squeeze_dim::<2>(2) // [B, 1] + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Distribution; + + #[test] + fn afine_q_head_forward_shape() { + let device = Default::default(); + let head = AfineQHeadConfig::new().init(&device); + + let image = Tensor::<4>::random([2, 3, 64, 64], Distribution::Default, &device); + let features: Vec<_> = (0..12) + .map(|_| Tensor::<3>::random([2, 4, 768], Distribution::Default, &device)) + .collect(); + + let out = head.forward(image, &features); + assert_eq!(out.dims(), [2, 1]); + } + + #[test] + fn afine_d_head_forward_shape() { + let device = Default::default(); + let head = AfineDHeadConfig::new().init(&device); + + let dis = Tensor::<4>::random([2, 3, 64, 64], Distribution::Default, &device); + let reference = Tensor::<4>::random([2, 3, 64, 64], Distribution::Default, &device); + let feat_dis: Vec<_> = (0..12) + .map(|_| Tensor::<3>::random([2, 4, 768], Distribution::Default, &device)) + .collect(); + let feat_ref: Vec<_> = (0..12) + .map(|_| Tensor::<3>::random([2, 4, 768], Distribution::Default, &device)) + .collect(); + + let out = head.forward(dis, reference, &feat_dis, &feat_ref); + assert_eq!(out.dims(), [2, 1]); + } + + #[test] + fn afine_d_head_identical_inputs_unit_score() { + // When dis == ref and CLIP features match, S1 and S2 should both + // be ≈ 1, so dist1 + dist2 ≈ 1 and the final score 1 - (...) ≈ 0. + // Random init alpha/beta sum to 1 so this is a property-test + // sanity check on the global normalization. + let device = Default::default(); + let head = AfineDHeadConfig::new().init(&device); + + let image = Tensor::<4>::random([1, 3, 32, 32], Distribution::Default, &device); + let features: Vec<_> = (0..12) + .map(|_| Tensor::<3>::random([1, 1, 768], Distribution::Default, &device)) + .collect(); + + let out = head.forward(image.clone(), image, &features.clone(), &features); + let value = out.into_data().to_vec::().unwrap()[0]; + assert!( + value.abs() < 0.1, + "fidelity head on identical inputs should yield ~0, got {value}" + ); + } +} diff --git a/crates/burn-train/src/metric/vision/afine/metric.rs b/crates/burn-train/src/metric/vision/afine/metric.rs new file mode 100644 index 0000000..2239713 --- /dev/null +++ b/crates/burn-train/src/metric/vision/afine/metric.rs @@ -0,0 +1,361 @@ +//! A-FINE metric: top-level module, configuration, end-to-end forward +//! pass, `ModuleDisplay`, and property tests. + +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Module, ModuleDisplay}; +use burn::tensor::Device; +use burn::tensor::Tensor; + +use super::calibrators::{ + AfineAdapter, AfineAdapterConfig, FrCalibratorWithLimit, FrCalibratorWithLimitConfig, + NrCalibrator, NrCalibratorConfig, scale_finalscore, +}; +use super::clip_vit::{ClipVisualEncoder, ClipVisualEncoderConfig}; +use super::heads::{AfineDHead, AfineDHeadConfig, AfineQHead, AfineQHeadConfig}; + +/// CLIP RGB normalization mean (per-channel). Applied at the top of +/// `forward` when `normalize_input` is true. +const CLIP_MEAN: [f32; 3] = [0.481_454_66, 0.457_827_5, 0.408_210_73]; +/// CLIP RGB normalization std. +const CLIP_STD: [f32; 3] = [0.268_629_54, 0.261_302_58, 0.275_777_11]; + +/// Configuration for the A-FINE metric. +#[derive(Config, Debug)] +pub struct AfineConfig { + /// Expected input resolution. Must be a multiple of 32. Defaults to 256. + #[config(default = "256")] + pub image_size: usize, + + /// Apply CLIP RGB normalization inside the forward pass. Set to false + /// if the caller has already normalized the input. Defaults to true. + #[config(default = "true")] + pub normalize_input: bool, +} + +impl AfineConfig { + /// Initialize an A-FINE module with default (random) weights. + /// + /// All six learnable submodules are initialized fresh; useful for + /// shape/property tests but produces meaningless quality scores + /// until [`Self::init_pretrained`] is called. + pub fn init(&self, device: &Device) -> Afine { + assert_eq!( + self.image_size % 32, + 0, + "A-FINE image_size must be a multiple of 32, got {}", + self.image_size + ); + + Afine { + clip_visual: ClipVisualEncoderConfig::new() + .with_image_size(self.image_size) + .init(device), + qhead: AfineQHeadConfig::new().init(device), + dhead: AfineDHeadConfig::new().init(device), + nr_calibrator: NrCalibratorConfig::new().init(device), + fr_calibrator: FrCalibratorWithLimitConfig::new().init(device), + adapter: AfineAdapterConfig::new().init(device), + image_size: self.image_size, + normalize_input: self.normalize_input, + } + } + + /// Initialize an A-FINE module and load pretrained PyIQA weights. + /// + /// Downloads `afine.pth` (~600 MB) from the PyIQA Hugging Face + /// mirror on first call, caches it under + /// `~/.cache/burn-dataset/afine/`, and loads all six shards into the + /// matching submodules. + pub fn init_pretrained(&self, device: &Device) -> Afine { + let afine = self.init(device); + super::weights::load_pretrained_weights(afine) + } +} + +/// A-FINE (Adaptive Fidelity-Naturalness Evaluator) full-reference image +/// quality metric built on CLIP ViT-B/32 features. +/// +/// `forward(distorted, reference)` returns a per-sample quality score in +/// `(0, 100)` — higher means worse perceptual quality. The metric is +/// **not symmetric**: `forward(a, b) != forward(b, a)` in general, +/// because the adapter term `exp(k * (N_ref - N_dis))` weights the two +/// inputs differently. +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Afine { + pub(crate) clip_visual: ClipVisualEncoder, + pub(crate) qhead: AfineQHead, + pub(crate) dhead: AfineDHead, + pub(crate) nr_calibrator: NrCalibrator, + pub(crate) fr_calibrator: FrCalibratorWithLimit, + pub(crate) adapter: AfineAdapter, + + pub(crate) image_size: usize, + pub(crate) normalize_input: bool, +} + +impl Afine { + /// Compute the A-FINE quality score. + /// + /// # Shapes + /// - `distorted`, `reference`: `[batch, 3, H, W]` with `H` and `W` + /// both multiples of 32. Values in `[0, 1]` (RGB) when + /// `normalize_input` is true; already-normalized when false. + /// - returns: `[batch]` per-sample score in `(0, 100)`. + pub fn forward(&self, distorted: Tensor<4>, reference: Tensor<4>) -> Tensor<1> { + let [_, _, height, width] = distorted.dims(); + assert_eq!( + height % 32, + 0, + "A-FINE input height ({height}) must be a multiple of 32" + ); + assert_eq!( + width % 32, + 0, + "A-FINE input width ({width}) must be a multiple of 32" + ); + + let device = distorted.device(); + let (dis_norm, ref_norm) = if self.normalize_input { + let mean = Tensor::<4>::from_floats( + [[[[CLIP_MEAN[0]]], [[CLIP_MEAN[1]]], [[CLIP_MEAN[2]]]]], + &device, + ); + let std = Tensor::<4>::from_floats( + [[[[CLIP_STD[0]]], [[CLIP_STD[1]]], [[CLIP_STD[2]]]]], + &device, + ); + ( + (distorted - mean.clone()) / std.clone(), + (reference - mean) / std, + ) + } else { + (distorted, reference) + }; + + // CLIP gives us the 12 per-block patch-feature maps the heads + // consume; the class-token vector is unused by A-FINE so we ask + // the encoder to skip it. + let feat_dis = self + .clip_visual + .forward_with_features(dis_norm.clone(), false) + .features; + let feat_ref = self + .clip_visual + .forward_with_features(ref_norm.clone(), false) + .features; + + let natural_dis = self.qhead.forward(dis_norm.clone(), &feat_dis); + let natural_ref = self.qhead.forward(ref_norm.clone(), &feat_ref); + let natural_dis_scaled = self.nr_calibrator.forward(natural_dis); + let natural_ref_scaled = self.nr_calibrator.forward(natural_ref); + + let fidelity = self.dhead.forward(dis_norm, ref_norm, &feat_dis, &feat_ref); + let fidelity_scaled = self.fr_calibrator.forward(fidelity); + + let d = self + .adapter + .forward(natural_dis_scaled, natural_ref_scaled, fidelity_scaled); + let score = scale_finalscore(d); // [B, 1] + score.squeeze_dim::<1>(1) // [B] + } +} + +impl ModuleDisplay for Afine { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("backbone", &"CLIP ViT-B/32".to_string()) + .add("image_size", &self.image_size.to_string()) + .add("normalize_input", &self.normalize_input.to_string()) + .optional() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn::tensor::Distribution; + + fn small_metric() -> Afine { + // A small image_size keeps tests fast — 12-layer CLIP at the + // default 256x256 takes a noticeable fraction of a second per + // forward. + let device = Default::default(); + AfineConfig::new().with_image_size(64).init(&device) + } + + #[test] + fn afine_forward_shape() { + let device = Default::default(); + let metric = small_metric(); + let dis = Tensor::<4>::random([2, 3, 64, 64], Distribution::Default, &device); + let reference = Tensor::<4>::random([2, 3, 64, 64], Distribution::Default, &device); + let score = metric.forward(dis, reference); + assert_eq!(score.dims(), [2]); + } + + #[test] + fn afine_batch_processing() { + let device = Default::default(); + let metric = small_metric(); + let dis = Tensor::<4>::random([4, 3, 64, 64], Distribution::Default, &device); + let reference = Tensor::<4>::random([4, 3, 64, 64], Distribution::Default, &device); + let score = metric.forward(dis, reference); + let values = score.into_data().to_vec::().unwrap(); + assert_eq!(values.len(), 4); + for v in values { + assert!(v.is_finite(), "A-FINE produced non-finite value: {v}"); + } + } + + #[test] + fn afine_finite_on_constant_inputs() { + // Guards against the c1=c2=1e-10 epsilon failing to keep the + // SSIM ratios finite when feature variance is zero. + let device = Default::default(); + let metric = small_metric(); + let zeros = Tensor::<4>::zeros([1, 3, 64, 64], &device); + let ones = Tensor::<4>::ones([1, 3, 64, 64], &device); + let score = metric.forward(zeros, ones); + let value = score.into_data().to_vec::().unwrap()[0]; + assert!(value.is_finite(), "got non-finite value {value}"); + } + + #[test] + fn afine_asymmetry() { + // D(a, b) != D(b, a) — the adapter weights N_dis by an + // exponential of (N_ref - N_dis), so swapping inputs flips the + // sign in the exponent. + let device = Default::default(); + let metric = small_metric(); + let a = Tensor::<4>::random([1, 3, 64, 64], Distribution::Default, &device); + let b = Tensor::<4>::random([1, 3, 64, 64], Distribution::Default, &device); + let forward = metric + .forward(a.clone(), b.clone()) + .into_data() + .to_vec::() + .unwrap()[0]; + let reverse = metric.forward(b, a).into_data().to_vec::().unwrap()[0]; + assert!( + (forward - reverse).abs() > 1e-6, + "expected asymmetric output, got fwd={forward}, rev={reverse}" + ); + } + + #[test] + fn afine_image_size_must_be_multiple_of_32() { + let result = std::panic::catch_unwind(|| { + AfineConfig::new() + .with_image_size(33) + .init(&Default::default()); + }); + assert!(result.is_err(), "expected init to panic on bad image_size"); + } + + #[test] + fn display_afine() { + let metric = small_metric(); + let formatted = format!("{metric}"); + assert!( + formatted.contains("CLIP ViT-B/32"), + "expected backbone name in display output: {formatted}" + ); + } + + #[test] + #[ignore = "downloads pre-trained weights"] + fn test_afine_pretrained_parity() { + // Numerical parity against PyIQA on a deterministic input pair. + // Catches silent QuickGELU-vs-GELU swaps, fused-QKV transpose- + // direction bugs, channel-order mistakes, `c1=c2=1e-10` epsilon + // drift, and any partially loaded shard — all of which the + // property tests miss. + // + // Expected value captured 2026-04-28 via + // `Notes/capture_afine_fixtures.py` (torch=2.11.0, + // pyiqa=0.1.15.post2). Input pair: `arange/(N-1)` vs `arange/N` + // reshaped to `[1, 3, 224, 224]`. Re-capture if the upstream + // PyIQA checkpoint or normalization conventions change. + const D_EXPECTED: f32 = 43.15711594_f32; + + let device = Default::default(); + let metric = AfineConfig::new() + .with_image_size(224) + .init_pretrained(&device); + + let total: i64 = 1 * 3 * 224 * 224; + let dis = Tensor::<1, burn::tensor::Int>::arange(0..total, &device) + .float() + .div_scalar((total - 1) as f64) + .reshape([1, 3, 224, 224]); + let reference = Tensor::<1, burn::tensor::Int>::arange(0..total, &device) + .float() + .div_scalar(total as f64) + .reshape([1, 3, 224, 224]); + + let value = metric + .forward(dis, reference) + .into_data() + .to_vec::() + .unwrap()[0]; + + // Tolerance: relative 5e-3 covers fp32 attention drift across + // 12 transformer layers. Absolute floor of 5e-3 prevents a + // captured value near zero from forcing absurd precision. In + // practice the burn output matches the captured PyIQA scalar to + // within ~5e-5 absolute, well inside this budget. + let tolerance = D_EXPECTED.abs() * 5e-3 + 5e-3; + assert!( + (value - D_EXPECTED).abs() < tolerance, + "expected {D_EXPECTED}, got {value} (tolerance {tolerance})" + ); + } + + #[test] + #[ignore = "downloads pre-trained weights"] + fn test_afine_pretrained() { + // Loads the real PyIQA checkpoint (~600 MB on first run, cached + // afterwards) and verifies the forward pass produces a finite + // score that meaningfully differs from the random-init metric. + // The differs-from-random check guards against a silent + // `allow_partial(true)` skip — if a regex remap is wrong the + // weights stay at random init and the property tests still + // pass, but the pretrained output should land in a very + // different region of the score range. + let device = Default::default(); + let dis = Tensor::<4>::random([1, 3, 256, 256], Distribution::Default, &device); + let reference = Tensor::<4>::random([1, 3, 256, 256], Distribution::Default, &device); + + let random_metric = AfineConfig::new().init(&device); + let random_score = random_metric + .forward(dis.clone(), reference.clone()) + .into_data() + .to_vec::() + .unwrap()[0]; + + let pretrained_metric = AfineConfig::new().init_pretrained(&device); + let pretrained_score = pretrained_metric + .forward(dis, reference) + .into_data() + .to_vec::() + .unwrap()[0]; + + assert!( + pretrained_score.is_finite(), + "pretrained A-FINE produced non-finite output: {pretrained_score}" + ); + assert!( + (pretrained_score - random_score).abs() > 1e-3, + "pretrained output ({pretrained_score}) too close to random-init ({random_score}) — \ + a shard may have failed to load silently" + ); + } +} diff --git a/crates/burn-train/src/metric/vision/afine/mod.rs b/crates/burn-train/src/metric/vision/afine/mod.rs new file mode 100644 index 0000000..dab543d --- /dev/null +++ b/crates/burn-train/src/metric/vision/afine/mod.rs @@ -0,0 +1,19 @@ +//! A-FINE (Adaptive Fidelity-Naturalness Evaluator) metric. +//! +//! Full-reference perceptual image quality metric that combines a naturalness +//! branch and a fidelity branch over CLIP ViT-B/32 features. +//! +//! Reference: "A Novel Fidelity-Naturalness Evaluator for Image Quality Assessment" +//! + +mod calibrators; +mod clip_attention; +mod clip_vit; +mod heads; +mod metric; +mod quick_gelu; +mod weights; + +pub use clip_vit::{ClipOutput, ClipVisualEncoder, ClipVisualEncoderConfig}; +pub use heads::{AfineDHead, AfineDHeadConfig, AfineQHead, AfineQHeadConfig}; +pub use metric::{Afine, AfineConfig}; diff --git a/crates/burn-train/src/metric/vision/afine/quick_gelu.rs b/crates/burn-train/src/metric/vision/afine/quick_gelu.rs new file mode 100644 index 0000000..22ef68a --- /dev/null +++ b/crates/burn-train/src/metric/vision/afine/quick_gelu.rs @@ -0,0 +1,65 @@ +//! QuickGELU activation used by OpenAI's CLIP. +//! +//! Defined as `x * sigmoid(1.702 * x)`. Distinct from the erf-based +//! [`burn_nn::Gelu`]: CLIP was trained with QuickGELU and substituting +//! standard GELU silently shifts every transformer-block MLP output. + +use burn_core as burn; + +use burn::module::Module; +use burn::tensor::Tensor; +use burn::tensor::activation::sigmoid; + +/// Approximation coefficient. Must be exactly 1.702; CLIP weights are +/// trained with this value and changing it breaks numerical parity. +const COEFFICIENT: f64 = 1.702; + +/// QuickGELU activation: `x * sigmoid(1.702 * x)`. +#[derive(Module, Debug, Default)] +pub(crate) struct QuickGelu; + +impl QuickGelu { + /// Apply the activation element-wise. Shape-preserving. + pub(crate) fn forward(&self, x: Tensor) -> Tensor { + let scaled = x.clone().mul_scalar(COEFFICIENT); + x * sigmoid(scaled) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn_core::tensor::Tolerance; + + type FT = f32; + + #[test] + fn quick_gelu_matches_formula() { + let device = Default::default(); + let activation = QuickGelu; + + let input = Tensor::<2>::from_floats([[-1.0, 0.0, 1.0, 2.0]], &device); + let output = activation.forward(input); + + // Hand-computed: x * sigmoid(1.702 * x). + // sigmoid(-1.702) = 1 / (1 + exp(1.702)) ≈ 0.154160 + // sigmoid( 1.702) ≈ 0.845840 + // sigmoid( 3.404) ≈ 0.967812 + let expected = Tensor::<2>::from_floats([[-0.154_160, 0.0, 0.845_840, 1.935_624]], &device); + + output + .into_data() + .assert_approx_eq::(&expected.into_data(), Tolerance::default()); + } + + #[test] + fn quick_gelu_preserves_shape() { + let device = Default::default(); + let activation = QuickGelu; + + let input = Tensor::<3>::zeros([2, 5, 8], &device); + let output = activation.forward(input); + + assert_eq!(output.dims(), [2, 5, 8]); + } +} diff --git a/crates/burn-train/src/metric/vision/afine/weights.rs b/crates/burn-train/src/metric/vision/afine/weights.rs new file mode 100644 index 0000000..ceb8b9d --- /dev/null +++ b/crates/burn-train/src/metric/vision/afine/weights.rs @@ -0,0 +1,257 @@ +//! Pretrained weights loading for A-FINE. +//! +//! `afine.pth` is published by the PyIQA project on Hugging Face and +//! bundles all six A-FINE shards into a single file: a fine-tuned CLIP +//! ViT-B/32 (`finetuned_clip`), the naturalness head (`natural`), the +//! fidelity head (`fidelity`), the two per-head logistic calibrators +//! (`natural_scale`, `fidelity_scale`), and the adapter (`adapter`). +//! Each shard is loaded into the corresponding submodule of [`Afine`] +//! via a separate [`PytorchStore`] with `with_top_level_key`. +//! +//! The OpenAI CLIP `ViT-B-32.pt` checkpoint is **not** used. That file +//! is a TorchScript archive that `burn-store` cannot read, and PyIQA's +//! `finetuned_clip` shard is the variant A-FINE was actually trained +//! with — substituting stock CLIP would silently produce wrong scores. +//! +//! ## Hosting +//! +//! The URL points at the PyIQA author's personal HF account, which is +//! the canonical source for these weights. There is no organization- +//! hosted mirror anywhere (HF orgs, Zenodo, figshare, ModelScope, OSF, +//! Kaggle, GitHub releases — all checked). The hosting question is +//! flagged in the PR description for the maintainer; if a different +//! mirror is preferred, swap the URL constant below. + +use burn_core as burn; + +use burn::module::Param; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn_std::network::downloader::download_file_as_bytes; +use burn_store::pytorch::PytorchReader; +use burn_store::{ModuleSnapshot, PytorchStore}; +use std::fs::{File, create_dir_all}; +use std::io::Write; +use std::path::Path; +use std::path::PathBuf; + +use super::calibrators::{AfineAdapter, FrCalibratorWithLimit, NrCalibrator}; +use super::metric::Afine; + +/// Source URL for `afine.pth` on Hugging Face. +const AFINE_URL: &str = + "https://huggingface.co/chaofengc/IQA-PyTorch-Weights/resolve/main/afine.pth"; + +/// Cached filename inside `~/.cache/burn-dataset/afine/`. +const CACHE_FILENAME: &str = "afine.pth"; + +fn get_cache_dir() -> PathBuf { + let cache_dir = dirs::cache_dir() + .expect("Could not get cache directory") + .join("burn-dataset") + .join("afine"); + if !cache_dir.exists() { + create_dir_all(&cache_dir).expect("Failed to create cache directory"); + } + cache_dir +} + +fn download_if_needed(url: &str, cache_path: &PathBuf, message: &str) { + if !cache_path.exists() { + let bytes = download_file_as_bytes(url, message); + let mut file = File::create(cache_path).expect("Failed to create cache file"); + file.write_all(&bytes).expect("Failed to write weights"); + } +} + +/// Download `afine.pth` (if not cached) and load all six shards into +/// the matching submodules of an `Afine` previously produced by +/// `AfineConfig::init`. +/// +/// Errors during loading are logged via `log::warn!` and the function +/// returns the module unconditionally — `allow_partial(true)` plus +/// per-shard regex remapping mean unmapped or unknown checkpoint keys +/// are silently dropped, matching the behaviour of LPIPS, DISTS, and +/// FID's pretrained loaders. +pub(crate) fn load_pretrained_weights(mut afine: Afine) -> Afine { + let cache_dir = get_cache_dir(); + let cache_path = cache_dir.join(CACHE_FILENAME); + download_if_needed(AFINE_URL, &cache_path, "Downloading A-FINE weights..."); + + afine.clip_visual = load_clip_shard(afine.clip_visual, &cache_path); + afine.qhead = load_qhead_shard(afine.qhead, &cache_path); + afine.dhead = load_simple_shard(afine.dhead, &cache_path, "fidelity", "fidelity head"); + + // The calibrator and adapter shards store yita3, yita4, and k as + // 0-D scalar tensors (`shape=()`). Burn's `Param>` of + // length 1 expects shape `(1,)`, and PytorchStore drops keys whose + // shape doesn't match the destination — silently leaving the + // params at random init. Load these manually via PytorchReader so + // the 0-D scalars become 1-element 1-D tensors. + let device = afine.adapter.k.val().device(); + afine.nr_calibrator = load_nr_calibrator_scalars(&cache_path, &device); + afine.fr_calibrator = load_fr_calibrator_scalars(&cache_path, &device); + afine.adapter = load_adapter_scalar(&cache_path, &device); + + afine +} + +/// Read a 0-D scalar value from a checkpoint shard, returning it as an +/// `f32`. Logs a warning and returns `None` on failure so the caller can +/// fall back to a default. +fn read_scalar_value>( + cache_path: P, + top_level_key: &str, + field_name: &str, +) -> Option { + let reader = PytorchReader::with_top_level_key(cache_path.as_ref(), top_level_key) + .map_err(|e| log::warn!("Failed to open shard '{top_level_key}': {e:?}")) + .ok()?; + let snapshot = reader.get(field_name)?; + let data = snapshot + .to_data() + .map_err(|e| log::warn!("Failed to read '{top_level_key}.{field_name}' tensor data: {e:?}")) + .ok()?; + let values = data + .to_vec::() + .map_err(|e| log::warn!("Failed to convert '{top_level_key}.{field_name}' to f32: {e:?}")) + .ok()?; + values.first().copied() +} + +fn scalar_param(value: f32, device: &Device) -> Param> { + Param::from_tensor(Tensor::from_floats([value], device)) +} + +fn load_nr_calibrator_scalars(cache_path: &PathBuf, device: &Device) -> NrCalibrator { + // PyIQA defaults if reading fails — preserves random-init fallback. + const NR_YITA3_FALLBACK: f32 = 4.9592; + const NR_YITA4_FALLBACK: f32 = 21.5968; + let yita3 = + read_scalar_value(cache_path, "natural_scale", "yita3").unwrap_or(NR_YITA3_FALLBACK); + let yita4 = + read_scalar_value(cache_path, "natural_scale", "yita4").unwrap_or(NR_YITA4_FALLBACK); + NrCalibrator { + yita3: scalar_param(yita3, device), + yita4: scalar_param(yita4, device), + } +} + +fn load_fr_calibrator_scalars(cache_path: &PathBuf, device: &Device) -> FrCalibratorWithLimit { + const FR_YITA3_FALLBACK: f32 = 0.5; + const FR_YITA4_FALLBACK: f32 = 0.15; + let yita3 = + read_scalar_value(cache_path, "fidelity_scale", "yita3").unwrap_or(FR_YITA3_FALLBACK); + let yita4 = + read_scalar_value(cache_path, "fidelity_scale", "yita4").unwrap_or(FR_YITA4_FALLBACK); + FrCalibratorWithLimit { + yita3: scalar_param(yita3, device), + yita4: scalar_param(yita4, device), + } +} + +fn load_adapter_scalar(cache_path: &PathBuf, device: &Device) -> AfineAdapter { + const K_FALLBACK: f32 = 5.0; + let k = read_scalar_value(cache_path, "adapter", "k").unwrap_or(K_FALLBACK); + AfineAdapter { + k: scalar_param(k, device), + } +} + +/// Load the `finetuned_clip` shard into the visual encoder. +/// +/// Key-remap rules strip the `visual.` prefix used in the CLIP state +/// dict and rename the few fields where the burn module diverges from +/// PyTorch's CLIP layout: `conv1` → `patch_embed`, +/// `class_embedding` → `class_token`, `transformer.resblocks` → `blocks`, +/// and the fused `attn.in_proj_*` → `attn.qkv_proj.*`. The +/// `weight`/`bias` ↔ `gamma`/`beta` rename for LayerNorm is handled +/// automatically by `PyTorchToBurnAdapter`. +/// +/// `with_regex(r"^visual\.")` filters the load to only the visual +/// encoder; the CLIP text encoder (~150 keys) is dropped at the filter +/// stage. `visual.proj` (the unused joint-embedding projection) has +/// no corresponding burn field and is silently dropped via +/// `allow_partial(true)`. +fn load_clip_shard( + mut clip: super::clip_vit::ClipVisualEncoder, + cache_path: &PathBuf, +) -> super::clip_vit::ClipVisualEncoder { + let mut store = PytorchStore::from_file(cache_path) + .with_top_level_key("finetuned_clip") + .allow_partial(true) + .skip_enum_variants(true) + // No `with_regex` filter: PathFilter matches against burn-side + // destination paths, not the source PyTorch keys, so a regex on + // `^visual\.` rejects every remapped key. Text-encoder keys + // (token_embedding, transformer.resblocks.*) have no + // corresponding burn fields and are dropped by + // `allow_partial(true)`, which is the same end result. + // + // Special case: the text encoder also has a top-level + // `positional_embedding` key (shape `[77, 512]`) which collides + // with the visual one after our `^visual\.positional_embedding$` + // rename targets `positional_embedding`. Without this first + // rule, HashMap iteration order decides which one wins, giving + // a flaky load. Rename the text encoder's first to a key that + // matches no burn field. + .with_key_remapping(r"^positional_embedding$", "_text_positional_embedding_drop") + .with_key_remapping(r"^visual\.conv1\.", "patch_embed.") + .with_key_remapping(r"^visual\.class_embedding$", "class_token") + .with_key_remapping(r"^visual\.positional_embedding$", "positional_embedding") + .with_key_remapping(r"^visual\.ln_pre\.", "ln_pre.") + .with_key_remapping(r"^visual\.ln_post\.", "ln_post.") + .with_key_remapping(r"^visual\.transformer\.resblocks\.", "blocks.") + .with_key_remapping(r"\.attn\.in_proj_weight$", ".attn.qkv_proj.weight") + .with_key_remapping(r"\.attn\.in_proj_bias$", ".attn.qkv_proj.bias"); + if let Err(e) = clip.load_from(&mut store) { + log::warn!( + "Some CLIP visual encoder weights could not be loaded: {:?}", + e + ); + } + clip +} + +/// Load the `natural` shard into the naturalness head. +/// +/// PyIQA stores the proj_head as a `nn.Sequential` with index-2 GELU, +/// so the FC layers come out as `proj_head.0.*` and `proj_head.2.*`. +/// We named them explicitly to avoid the GELU-index gap, hence the +/// remap. +fn load_qhead_shard( + mut qhead: super::heads::AfineQHead, + cache_path: &PathBuf, +) -> super::heads::AfineQHead { + let mut store = PytorchStore::from_file(cache_path) + .with_top_level_key("natural") + .allow_partial(true) + .skip_enum_variants(true) + .with_key_remapping(r"^proj_head\.0\.", "proj_head_fc1.") + .with_key_remapping(r"^proj_head\.2\.", "proj_head_fc2."); + if let Err(e) = qhead.load_from(&mut store) { + log::warn!("Some naturalness head weights could not be loaded: {:?}", e); + } + qhead +} + +/// Load any shard whose keys map directly to burn-side field names. +/// Used for the fidelity head, both calibrators, and the adapter. +fn load_simple_shard( + mut module: M, + cache_path: &PathBuf, + top_level_key: &'static str, + description: &str, +) -> M +where + M: ModuleSnapshot, +{ + let mut store = PytorchStore::from_file(cache_path) + .with_top_level_key(top_level_key) + .allow_partial(true) + .skip_enum_variants(true); + if let Err(e) = module.load_from(&mut store) { + log::warn!("Some {} weights could not be loaded: {:?}", description, e); + } + module +} diff --git a/crates/burn-train/src/metric/vision/dice.rs b/crates/burn-train/src/metric/vision/dice.rs new file mode 100644 index 0000000..19f910a --- /dev/null +++ b/crates/burn-train/src/metric/vision/dice.rs @@ -0,0 +1,341 @@ +use crate::metric::{MetricAttributes, MetricName, SerializedEntry}; + +use super::super::{ + Metric, MetricMetadata, + state::{FormatOptions, NumericMetricState}, +}; +use burn_core::{ + prelude::Tensor, + tensor::{ElementConversion, Int, s}, +}; + +/// Input type for the [DiceMetric]. +/// +/// # Type Parameters +/// - `D`: Number of dimensions. Should be more than, or equal to 3 (default 4). +pub struct DiceInput { + /// Model outputs (predictions), as a tensor. + outputs: Tensor, + /// Ground truth targets, as a tensor. + targets: Tensor, +} + +impl DiceInput { + /// Creates a new DiceInput with the given outputs and targets. + /// + /// Inputs are expected to have the dimensions `[B, C, ...]` + /// where `B` is the batch size, `C` is the number of classes, + /// and `...` represents additional dimensions (e.g., height, width for images). + /// + /// If `C` is more than 1, the first class (index 0) is considered the background. + /// Additionally, one-hot encoding is the responsibility of the caller. + /// + /// # Arguments + /// - `outputs`: The model outputs as a tensor. + /// - `targets`: The ground truth targets as a tensor. + /// + /// # Returns + /// A new instance of `DiceInput`. + /// + /// # Panics + /// - If `D` is less than 3. + /// - If `outputs` and `targets` do not have the same dimensions. + /// - If `outputs` or `targets` do not have exactly `D` dimensions. + /// - If `outputs` and `targets` do not have the same shape. + pub fn new(outputs: Tensor, targets: Tensor) -> Self { + assert!(D >= 3, "DiceInput requires at least 3 dimensions."); + assert!( + outputs.dims() == targets.dims(), + "Outputs and targets must have the same dimensions. Got {:?} and {:?}", + outputs.dims(), + targets.dims() + ); + Self { outputs, targets } + } +} + +/// Configuration for the [DiceMetric]. +#[derive(Debug, Clone, Copy)] +pub struct DiceMetricConfig { + /// Epsilon value to avoid division by zero. + pub epsilon: f64, + /// Whether to include the background class in the metric calculation. + /// The background is assumed to be the first class (index 0). + /// if `true`, will panic if there are fewer than 2 classes. + pub include_background: bool, +} + +impl Default for DiceMetricConfig { + fn default() -> Self { + Self { + epsilon: 1e-7, + include_background: false, + } + } +} + +/// The Dice-Sorenson coefficient (DSC) for evaluating overlap between two binary masks. +/// The DSC is defined as: +/// `DSC = 2 * (|X ∩ Y|) / (|X| + |Y|)` +/// where `X` is the model output and `Y` is the ground truth target. +/// +/// # Type Parameters +/// - `D`: Number of dimensions. Should be more than, or equal to 3 (default 4). +#[derive(Default, Clone)] +pub struct DiceMetric { + name: MetricName, + /// Internal state for numeric metric aggregation. + state: NumericMetricState, + /// Configuration for the metric. + config: DiceMetricConfig, +} + +impl DiceMetric { + /// Creates a new Dice metric instance with default config. + pub fn new() -> Self { + Self::with_config(DiceMetricConfig::default()) + } + + /// Creates a new Dice metric with a custom config. + pub fn with_config(config: DiceMetricConfig) -> Self { + let name = MetricName::new(format!("{D}D Dice Metric")); + assert!(D >= 3, "DiceMetric requires at least 3 dimensions."); + Self { + name, + config, + ..Default::default() + } + } +} + +impl Metric for DiceMetric { + type Input = DiceInput; + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn update(&mut self, item: &Self::Input, _metadata: &MetricMetadata) -> SerializedEntry { + // Dice coefficient: 2 * (|X ∩ Y|) / (|X| + |Y|) + if item.outputs.dims() != item.targets.dims() { + panic!( + "Outputs and targets must have the same dimensions. Got {:?} and {:?}", + item.outputs.dims(), + item.targets.dims() + ); + } + + let dims = item.outputs.dims(); + let batch_size = dims[0]; + let n_classes = dims[1]; + + let mut outputs = item.outputs.clone(); + let mut targets = item.targets.clone(); + + if !self.config.include_background && n_classes > 1 { + // If not including background, we can ignore the first class + outputs = outputs.slice(s![.., 1..]); + targets = targets.slice(s![.., 1..]); + } else if self.config.include_background && n_classes < 2 { + // If including background, we need at least 2 classes + panic!("Dice metric requires at least 2 classes when including background."); + } + + let intersection = (outputs.clone() * targets.clone()).sum(); + let outputs_sum = outputs.sum(); + let targets_sum = targets.sum(); + + // Convert to f64 + let intersection_val = intersection.into_scalar::(); + let outputs_sum_val = outputs_sum.into_scalar::(); + let targets_sum_val = targets_sum.into_scalar::(); + + // Use epsilon from config + let epsilon = self.config.epsilon; + let dice = + (2.0 * intersection_val + epsilon) / (outputs_sum_val + targets_sum_val + epsilon); + + self.state.update( + dice, + batch_size, + FormatOptions::new(self.name()).precision(4), + ) + } + + /// Clears the metric state. + fn clear(&mut self) { + self.state.reset(); + } + + fn attributes(&self) -> MetricAttributes { + crate::metric::NumericAttributes { + unit: None, + higher_is_better: true, + ..Default::default() + } + .into() + } +} + +impl crate::metric::Numeric for DiceMetric { + fn value(&self) -> crate::metric::NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> crate::metric::NumericEntry { + self.state.running_value() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::metric::Numeric; + use burn_core::tensor::{Shape, Tensor}; + + #[test] + fn test_dice_perfect_overlap() { + let device = Default::default(); + let mut metric = DiceMetric::<4>::new(); + let input = DiceInput::new( + Tensor::from_data([[[[1, 0], [1, 0]]]], &device), + Tensor::from_data([[[[1, 0], [1, 0]]]], &device), + ); + let _entry = metric.update(&input, &MetricMetadata::fake()); + assert!((metric.value().current() - 1.0).abs() < 1e-6); + } + + #[test] + fn test_dice_no_overlap() { + let device = Default::default(); + let mut metric = DiceMetric::<4>::new(); + let input = DiceInput::new( + Tensor::from_data([[[[1, 0], [1, 0]]]], &device), + Tensor::from_data([[[[0, 1], [0, 1]]]], &device), + ); + let _entry = metric.update(&input, &MetricMetadata::fake()); + assert!(metric.value().current() < 1e-6); + } + + #[test] + fn test_dice_partial_overlap() { + let device = Default::default(); + let mut metric = DiceMetric::<4>::new(); + let input = DiceInput::new( + Tensor::from_data([[[[1, 1], [0, 0]]]], &device), + Tensor::from_data([[[[1, 0], [1, 0]]]], &device), + ); + let _entry = metric.update(&input, &MetricMetadata::fake()); + // intersection = 1, sum = 2+2=4, dice = 2*1/4 = 0.5 + assert!((metric.value().current() - 0.5).abs() < 1e-6); + } + + #[test] + fn test_dice_empty_masks() { + let device = Default::default(); + let mut metric = DiceMetric::<4>::new(); + let input = DiceInput::new( + Tensor::from_data([[[[0, 0], [0, 0]]]], &device), + Tensor::from_data([[[[0, 0], [0, 0]]]], &device), + ); + let _entry = metric.update(&input, &MetricMetadata::fake()); + assert!((metric.value().current() - 1.0).abs() < 1e-6); + } + + #[test] + fn test_dice_no_background() { + let device = Default::default(); + let mut metric = DiceMetric::<4>::new(); + let input = DiceInput::new( + Tensor::ones(Shape::new([1, 1, 2, 2]), &device), + Tensor::ones(Shape::new([1, 1, 2, 2]), &device), + ); + let _entry = metric.update(&input, &MetricMetadata::fake()); + assert!((metric.value().current() - 1.0).abs() < 1e-6); + } + + #[test] + fn test_dice_with_background() { + let device = Default::default(); + let config = DiceMetricConfig { + epsilon: 1e-7, + include_background: true, + }; + let mut metric = DiceMetric::<4>::with_config(config); + let input = DiceInput::new( + Tensor::ones(Shape::new([1, 2, 2, 2]), &device), + Tensor::ones(Shape::new([1, 2, 2, 2]), &device), + ); + let _entry = metric.update(&input, &MetricMetadata::fake()); + assert!((metric.value().current() - 1.0).abs() < 1e-6); + } + + #[test] + fn test_dice_ignored_background() { + let device = Default::default(); + let config = DiceMetricConfig { + epsilon: 1e-7, + include_background: false, + }; + let mut metric = DiceMetric::<4>::with_config(config); + let input = DiceInput::new( + Tensor::ones(Shape::new([1, 2, 2, 2]), &device), + Tensor::ones(Shape::new([1, 2, 2, 2]), &device), + ); + let _entry = metric.update(&input, &MetricMetadata::fake()); + assert!((metric.value().current() - 1.0).abs() < 1e-6); + } + + #[test] + #[should_panic(expected = "DiceInput requires at least 3 dimensions.")] + fn test_invalid_input_dimensions() { + let device = Default::default(); + // D = 2, should panic + let _ = DiceInput::<2>::new( + Tensor::from_data([[0.0, 0.0]], &device), + Tensor::from_data([[0.0, 0.0]], &device), + ); + } + + #[test] + #[should_panic( + expected = "Outputs and targets must have the same dimensions. Got [1, 1, 2, 2] and [1, 1, 2, 3]" + )] + fn test_mismatched_shape() { + let device = Default::default(); + // shapes differ + let _ = DiceInput::<4>::new( + Tensor::from_data([[[[0.0; 2]; 2]; 1]; 1], &device), + Tensor::from_data([[[[0.0; 3]; 2]; 1]; 1], &device), + ); + } + + #[test] + #[should_panic(expected = "Dice metric requires at least 2 classes when including background.")] + fn test_include_background_panic() { + let device = Default::default(); + let config = DiceMetricConfig { + epsilon: 1e-7, + include_background: true, + }; + let mut metric = DiceMetric::<4>::with_config(config); + let input = DiceInput::new( + Tensor::from_data([[[[1.0; 2]; 1]; 1]; 1], &device), + Tensor::from_data([[[[1.0; 2]; 1]; 1]; 1], &device), + ); + // n_classes = 2, should not panic + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let config = DiceMetricConfig { + epsilon: 1e-7, + include_background: true, + }; + let mut metric = DiceMetric::<4>::with_config(config); + let input = DiceInput::new( + Tensor::from_data([[[[1.0; 1]; 1]; 1]; 1], &device), + Tensor::from_data([[[[1.0; 1]; 1]; 1]; 1], &device), + ); + // n_classes = 1, should panic + let _entry = metric.update(&input, &MetricMetadata::fake()); + } +} diff --git a/crates/burn-train/src/metric/vision/dists/l2pool.rs b/crates/burn-train/src/metric/vision/dists/l2pool.rs new file mode 100644 index 0000000..794b941 --- /dev/null +++ b/crates/burn-train/src/metric/vision/dists/l2pool.rs @@ -0,0 +1,160 @@ +//! L2 Pooling layer for DISTS. +//! +//! L2 Pooling applies a Hanning window filter and computes the L2 norm +//! across the pooling window. This is used in DISTS instead of MaxPooling. + +use burn_core as burn; + +use burn::module::Module; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn_nn::PaddingConfig2d; +use burn_nn::conv::{Conv2d, Conv2dConfig}; + +/// L2 Pooling layer configuration. +#[derive(Debug, Clone)] +pub struct L2Pool2dConfig { + /// Kernel size for pooling + pub kernel_size: usize, + /// Stride for pooling + pub stride: usize, + /// Padding for pooling + pub padding: usize, +} + +impl Default for L2Pool2dConfig { + fn default() -> Self { + Self { + kernel_size: 5, + stride: 2, + padding: 2, + } + } +} + +impl L2Pool2dConfig { + /// Create a new L2Pool2d configuration. + #[allow(dead_code)] + pub fn new(kernel_size: usize, stride: usize, padding: usize) -> Self { + Self { + kernel_size, + stride, + padding, + } + } + + /// Initialize the L2Pool2d layer. + pub fn init(&self, channels: usize, device: &Device) -> L2Pool2d { + L2Pool2d::new( + channels, + self.kernel_size, + self.stride, + self.padding, + device, + ) + } +} + +/// L2 Pooling layer. +/// +/// Applies a 2D Hanning window filter followed by L2 normalization. +/// This provides smoother downsampling compared to MaxPooling. +#[derive(Module, Debug)] +pub struct L2Pool2d { + /// Depthwise convolution with Hanning kernel + conv: Conv2d, +} + +impl L2Pool2d { + /// Create a new L2Pool2d layer with Hanning window kernel. + pub fn new( + channels: usize, + kernel_size: usize, + stride: usize, + padding: usize, + device: &Device, + ) -> Self { + // Create Hanning kernel + let kernel = Self::create_hanning_kernel(channels, kernel_size, device); + + // Create depthwise convolution (groups = channels) + let mut conv = Conv2dConfig::new([channels, channels], [kernel_size, kernel_size]) + .with_stride([stride, stride]) + .with_padding(PaddingConfig2d::Explicit( + padding, padding, padding, padding, + )) + .with_groups(channels) + .with_bias(false) + .init(device); + + // Set the kernel weights to Hanning window + conv.weight = burn::module::Param::from_tensor(kernel); + + Self { conv } + } + + /// Create a Hanning kernel for depthwise convolution. + /// Output shape: [channels, 1, kernel_size, kernel_size] + fn create_hanning_kernel(channels: usize, kernel_size: usize, device: &Device) -> Tensor<4> { + // Create 1D Hanning window + let mut hanning_1d = Vec::with_capacity(kernel_size); + for i in 0..kernel_size { + let n = i as f32; + let n_minus_1 = (kernel_size - 1) as f32; + let value = if n_minus_1 == 0.0 { + 1.0 + } else { + 0.5 * (1.0 - (2.0 * std::f32::consts::PI * n / n_minus_1).cos()) + }; + hanning_1d.push(value); + } + + // Create 2D Hanning window by outer product + let mut hanning_2d = Vec::with_capacity(kernel_size * kernel_size); + let mut sum = 0.0; + for i in 0..kernel_size { + for j in 0..kernel_size { + let value = hanning_1d[i] * hanning_1d[j]; + hanning_2d.push(value); + sum += value; + } + } + + // Normalize + for v in hanning_2d.iter_mut() { + *v /= sum; + } + + // Create tensor of shape [1, 1, kernel_size, kernel_size] + let kernel_single = Tensor::<1>::from_floats(hanning_2d.as_slice(), device).reshape([ + 1, + 1, + kernel_size, + kernel_size, + ]); + + // Expand to [channels, 1, kernel_size, kernel_size] + kernel_single.repeat_dim(0, channels) + } + + /// Apply L2 pooling to the input tensor. + /// + /// # Arguments + /// + /// * `x` - Input tensor of shape `[batch, channels, height, width]` + /// + /// # Returns + /// + /// Pooled tensor with reduced spatial dimensions. + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + // Square the input + let x_sq = x.clone().mul(x); + + // Apply depthwise convolution with Hanning kernel + let pooled = self.conv.forward(x_sq); + + // Take square root for L2 norm + // Add small epsilon to avoid sqrt of negative numbers due to numerical errors + pooled.clamp_min(1e-10).sqrt() + } +} diff --git a/crates/burn-train/src/metric/vision/dists/metric.rs b/crates/burn-train/src/metric/vision/dists/metric.rs new file mode 100644 index 0000000..23b6e29 --- /dev/null +++ b/crates/burn-train/src/metric/vision/dists/metric.rs @@ -0,0 +1,487 @@ +//! DISTS (Deep Image Structure and Texture Similarity) metric. +//! +//! DISTS is a full-reference image quality assessment metric that combines +//! structure and texture similarity using deep features from VGG16. +//! +//! Reference: "Image Quality Assessment: Unifying Structure and Texture Similarity" +//! https://arxiv.org/abs/2004.07728 + +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Module, ModuleDisplay, Param}; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn_nn::loss::Reduction; + +use super::vgg16_l2pool::Vgg16L2PoolExtractor; + +/// Channel counts for each stage: [input, stage1, stage2, stage3, stage4, stage5] +const CHANNELS: [usize; 6] = [3, 64, 128, 256, 512, 512]; + +/// Small constant for numerical stability in structure similarity. +const C1: f32 = 1e-6; + +/// Small constant for numerical stability in texture similarity. +const C2: f32 = 1e-6; + +/// ImageNet normalization constants. +const IMAGENET_MEAN: [f32; 3] = [0.485, 0.456, 0.406]; +const IMAGENET_STD: [f32; 3] = [0.229, 0.224, 0.225]; + +/// Image normalizer with pre-initialized mean and std tensors. +/// +/// This struct holds the mean and std tensors for normalization, +/// avoiding the need to create them on each forward pass. +#[derive(Module, Debug)] +pub struct Normalizer { + /// Mean tensor of shape [1, 3, 1, 1] for broadcasting. + pub mean: Tensor<4>, + /// Std tensor of shape [1, 3, 1, 1] for broadcasting. + pub std: Tensor<4>, +} + +impl Normalizer { + /// Create a new ImageNet normalizer. + pub fn imagenet(device: &Device) -> Self { + // Shape: [1, 3, 1, 1] for broadcasting over [batch, channels, height, width] + let mean = Tensor::from_floats( + [[ + [[IMAGENET_MEAN[0]]], + [[IMAGENET_MEAN[1]]], + [[IMAGENET_MEAN[2]]], + ]], + device, + ); + let std = Tensor::from_floats( + [[ + [[IMAGENET_STD[0]]], + [[IMAGENET_STD[1]]], + [[IMAGENET_STD[2]]], + ]], + device, + ); + Self { mean, std } + } + + /// Normalize a tensor: (x - mean) / std + pub fn normalize(&self, x: Tensor<4>) -> Tensor<4> { + x.sub(self.mean.clone()).div(self.std.clone()) + } +} + +/// Configuration for DISTS metric. +#[derive(Config, Debug)] +pub struct DistsConfig { + /// Whether to apply ImageNet normalization to input images. + #[config(default = true)] + pub normalize: bool, +} + +impl DistsConfig { + /// Initialize a DISTS module with default weights. + pub fn init(&self, device: &Device) -> Dists { + let total_channels: usize = CHANNELS.iter().sum(); + + // Initialize alpha and beta with constant value 0.1 for all channels + let alpha_data: Vec = (0..total_channels).map(|_| 0.1).collect(); + let beta_data: Vec = (0..total_channels).map(|_| 0.1).collect(); + + let normalizer = if self.normalize { + Some(Normalizer::imagenet(device)) + } else { + None + }; + + Dists { + extractor: Vgg16L2PoolExtractor::new(device), + alpha: Param::from_tensor(Tensor::from_floats(alpha_data.as_slice(), device)), + beta: Param::from_tensor(Tensor::from_floats(beta_data.as_slice(), device)), + normalizer, + } + } + + /// Initialize a DISTS module with pretrained weights. + pub fn init_pretrained(&self, device: &Device) -> Dists { + let dists = self.init(device); + super::weights::load_pretrained_weights(dists) + } +} + +/// DISTS (Deep Image Structure and Texture Similarity) metric. +/// +/// Computes perceptual similarity between two images by combining +/// structure similarity (based on spatial means) and texture similarity +/// (based on variances and covariances) across VGG16 feature maps. +/// +/// # Example +/// +/// ```ignore +/// use burn_train::metric::vision::{DistsConfig, Reduction}; +/// +/// let device = Default::default(); +/// let dists = DistsConfig::new().init_pretrained(&device); +/// +/// let img1: Tensor<4> = /* [batch, 3, H, W] */; +/// let img2: Tensor<4> = /* [batch, 3, H, W] */; +/// +/// let distance = dists.forward(img1, img2, Reduction::Mean); +/// ``` +#[derive(Module, Debug)] +#[module(custom_display)] +pub struct Dists { + /// VGG16 feature extractor with L2 pooling + pub(crate) extractor: Vgg16L2PoolExtractor, + /// Learned weights for structure similarity (per channel) + pub(crate) alpha: Param>, + /// Learned weights for texture similarity (per channel) + pub(crate) beta: Param>, + /// Optional normalizer for input preprocessing + pub(crate) normalizer: Option, +} + +impl ModuleDisplay for Dists { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + content + .add("backbone", &"VGG16-L2Pool".to_string()) + .add("normalize", &self.normalizer.is_some().to_string()) + .optional() + } +} + +impl Dists { + /// Compute DISTS distance with reduction. + /// + /// # Arguments + /// + /// * `input` - First image tensor of shape `[batch, 3, H, W]` + /// * `target` - Second image tensor of shape `[batch, 3, H, W]` + /// * `reduction` - How to reduce the output (Mean, Sum, or Auto) + /// + /// # Returns + /// + /// Scalar tensor of shape `[1]`. + pub fn forward(&self, input: Tensor<4>, target: Tensor<4>, reduction: Reduction) -> Tensor<1> { + let distance = self.forward_no_reduction(input, target); + + match reduction { + Reduction::Mean | Reduction::Auto | Reduction::BatchMean => distance.mean(), + Reduction::Sum => distance.sum(), + } + } + + /// Compute DISTS distance without reduction. + /// + /// # Arguments + /// + /// * `input` - First image tensor of shape `[batch, 3, H, W]` + /// * `target` - Second image tensor of shape `[batch, 3, H, W]` + /// + /// # Returns + /// + /// Per-sample distance tensor of shape `[batch]`. + pub fn forward_no_reduction(&self, input: Tensor<4>, target: Tensor<4>) -> Tensor<1> { + let [batch, _, _, _] = input.dims(); + + // Preprocess inputs + let (input, target) = self.preprocess(input, target); + + // Extract features from both images + let feats_x = self.extractor.forward(input); + let feats_y = self.extractor.forward(target); + + // Get alpha and beta weights + let alpha = self.alpha.val(); + let beta = self.beta.val(); + + // Compute weighted sum of alpha and beta for normalization + let alpha_sum = alpha.clone().sum(); + let beta_sum = beta.clone().sum(); + + let device = feats_x[0].device(); + + // Initialize accumulators + let mut structure_dist = Tensor::<1>::zeros([batch], &device); + let mut texture_dist = Tensor::<1>::zeros([batch], &device); + + let mut channel_offset = 0; + + // Compute similarity for each stage + for (feat_x, feat_y) in feats_x.iter().zip(feats_y.iter()) { + let [_b, c, _h, _w] = feat_x.dims(); + + // Get alpha and beta for this stage + let alpha_stage = alpha.clone().narrow(0, channel_offset, c); + let beta_stage = beta.clone().narrow(0, channel_offset, c); + + // Compute structure and texture similarity for this stage + let (s_dist, t_dist) = self.compute_stage_similarity( + feat_x.clone(), + feat_y.clone(), + alpha_stage, + beta_stage, + ); + + structure_dist = structure_dist.add(s_dist); + texture_dist = texture_dist.add(t_dist); + + channel_offset += c; + } + + // Normalize by sum of weights + structure_dist = structure_dist.div(alpha_sum); + texture_dist = texture_dist.div(beta_sum); + + // DISTS = 1 - (structure_similarity + texture_similarity) + // Since we computed distances (1 - similarity), we return the sum + structure_dist.add(texture_dist) + } + + /// Compute structure and texture similarity for a single stage. + fn compute_stage_similarity( + &self, + feat_x: Tensor<4>, + feat_y: Tensor<4>, + alpha: Tensor<1>, + beta: Tensor<1>, + ) -> (Tensor<1>, Tensor<1>) { + let [batch, channels, height, width] = feat_x.dims(); + let device = feat_x.device(); + + // Reshape to [batch, channels, H*W] for easier computation + let x = feat_x.reshape([batch, channels, height * width]); + let y = feat_y.reshape([batch, channels, height * width]); + + // Compute means: [batch, channels] (squeeze after mean_dim to remove the reduced dimension) + let mean_x = x.clone().mean_dim(2).squeeze_dim::<2>(2); + let mean_y = y.clone().mean_dim(2).squeeze_dim::<2>(2); + + // Compute structure similarity: (2*mean_x*mean_y + c1) / (mean_x^2 + mean_y^2 + c1) + let c1 = Tensor::<2>::full([batch, channels], C1, &device); + let structure_sim = mean_x + .clone() + .mul(mean_y.clone()) + .mul_scalar(2.0) + .add(c1.clone()) + .div( + mean_x + .clone() + .mul(mean_x.clone()) + .add(mean_y.clone().mul(mean_y.clone())) + .add(c1), + ); + + // Compute variances and covariance + // var_x = E[x^2] - E[x]^2, clamped at 0 for numerical stability + let var_x = x + .clone() + .mul(x.clone()) + .mean_dim(2) + .squeeze_dim::<2>(2) + .sub(mean_x.clone().mul(mean_x.clone())) + .clamp_min(0.0); + let var_y = y + .clone() + .mul(y.clone()) + .mean_dim(2) + .squeeze_dim::<2>(2) + .sub(mean_y.clone().mul(mean_y.clone())) + .clamp_min(0.0); + + // cov_xy = E[xy] - E[x]E[y] + let cov_xy = x + .mul(y) + .mean_dim(2) + .squeeze_dim::<2>(2) + .sub(mean_x.clone().mul(mean_y.clone())); + + // Compute texture similarity: (2*cov_xy + c2) / (var_x + var_y + c2) + let c2 = Tensor::<2>::full([batch, channels], C2, &device); + let texture_sim = cov_xy + .mul_scalar(2.0) + .add(c2.clone()) + .div(var_x.add(var_y).add(c2)); + + // Convert similarity to distance: 1 - similarity + let structure_dist = Tensor::<2>::ones([batch, channels], &device).sub(structure_sim); + let texture_dist = Tensor::<2>::ones([batch, channels], &device).sub(texture_sim); + + // Apply weights: [batch, channels] * [channels] -> [batch, channels] + // Then sum over channels -> [batch] + let weighted_structure = structure_dist + .mul(alpha.unsqueeze_dim::<2>(0)) + .sum_dim(1) + .squeeze_dim::<1>(1); + let weighted_texture = texture_dist + .mul(beta.unsqueeze_dim::<2>(0)) + .sum_dim(1) + .squeeze_dim::<1>(1); + + (weighted_structure, weighted_texture) + } + + /// Preprocess input images using the configured normalizer. + fn preprocess(&self, input: Tensor<4>, target: Tensor<4>) -> (Tensor<4>, Tensor<4>) { + match &self.normalizer { + Some(normalizer) => { + let input = normalizer.normalize(input); + let target = normalizer.normalize(target); + (input, target) + } + None => (input, target), + } + } +} + +// ============================================================================= +// Tests +// ============================================================================= + +#[cfg(test)] +mod tests { + use super::*; + use burn_core::tensor::{TensorData, Tolerance}; + + type FT = f32; + + #[test] + fn test_dists_identical_images_zero_distance() { + let device = Default::default(); + // Use random image instead of constant to avoid numerical edge cases + let image = Tensor::<4>::random( + [1, 3, 64, 64], + burn_core::tensor::Distribution::Uniform(0.0, 1.0), + &device, + ); + + let dists = DistsConfig::new().init(&device); + let distance = dists.forward(image.clone(), image, Reduction::Mean); + + let expected = TensorData::from([0.0]); + distance + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + #[test] + fn test_dists_different_images_nonzero_distance() { + let device = Default::default(); + + let image1 = Tensor::<4>::zeros([1, 3, 64, 64], &device); + let image2 = Tensor::<4>::ones([1, 3, 64, 64], &device); + + let dists = DistsConfig::new().init(&device); + let distance = dists.forward(image1, image2, Reduction::Mean); + + let distance_value = distance.into_data().to_vec::().unwrap()[0]; + assert!( + distance_value.abs() > 1e-6, + "DISTS should be != 0 for different images" + ); + } + + #[test] + fn test_dists_symmetry() { + let device = Default::default(); + + let image1 = Tensor::<4>::zeros([1, 3, 32, 32], &device); + let image2 = Tensor::<4>::ones([1, 3, 32, 32], &device); + + let dists = DistsConfig::new().init(&device); + let distance_forward = dists.forward(image1.clone(), image2.clone(), Reduction::Mean); + let distance_reverse = dists.forward(image2, image1, Reduction::Mean); + + distance_forward + .into_data() + .assert_approx_eq::(&distance_reverse.into_data(), Tolerance::default()); + } + + #[test] + fn test_dists_batch_processing() { + let device = Default::default(); + + let image1 = Tensor::<4>::zeros([2, 3, 32, 32], &device); + let image2 = Tensor::<4>::ones([2, 3, 32, 32], &device); + + let dists = DistsConfig::new().init(&device); + let distance = dists.forward(image1, image2, Reduction::Mean); + + assert_eq!(distance.dims(), [1]); + } + + #[test] + fn test_dists_no_reduction() { + let device = Default::default(); + + let batch_size = 4; + let image1 = Tensor::<4>::zeros([batch_size, 3, 32, 32], &device); + let image2 = Tensor::<4>::ones([batch_size, 3, 32, 32], &device); + + let dists = DistsConfig::new().init(&device); + let distance = dists.forward_no_reduction(image1, image2); + + assert_eq!(distance.dims(), [batch_size]); + } + + #[test] + fn display_dists() { + let device = Default::default(); + let dists = DistsConfig::new().init(&device); + + let display_str = format!("{dists}"); + assert!(display_str.contains("Dists")); + assert!(display_str.contains("VGG16-L2Pool")); + } + + // ========================================================================= + // Pretrained Weights Tests (requires network) + // ========================================================================= + + /// Test DISTS pretrained weights download and loading. + #[test] + #[ignore = "downloads pre-trained weights"] + fn test_dists_pretrained() { + let device = Default::default(); + + let dists = DistsConfig::new().init_pretrained(&device); + + // Test with identical images - should be ~0 + // Use random image to avoid numerical edge cases with constant images + let image = Tensor::<4>::random( + [1, 3, 64, 64], + burn_core::tensor::Distribution::Uniform(0.0, 1.0), + &device, + ); + let distance = dists.forward(image.clone(), image, Reduction::Mean); + let distance_value = distance.into_data().to_vec::().unwrap()[0]; + assert!( + distance_value.abs() < 1e-5, + "Pretrained DISTS should be ~0 for identical images, got {}", + distance_value + ); + + // Test with different images - should be positive + let image1 = Tensor::<4>::random( + [1, 3, 64, 64], + burn_core::tensor::Distribution::Uniform(0.0, 0.3), + &device, + ); + let image2 = Tensor::<4>::random( + [1, 3, 64, 64], + burn_core::tensor::Distribution::Uniform(0.7, 1.0), + &device, + ); + let distance = dists.forward(image1, image2, Reduction::Mean); + let distance_value = distance.into_data().to_vec::().unwrap()[0]; + assert!( + distance_value > 0.0, + "Pretrained DISTS should be > 0 for different images" + ); + } +} diff --git a/crates/burn-train/src/metric/vision/dists/mod.rs b/crates/burn-train/src/metric/vision/dists/mod.rs new file mode 100644 index 0000000..4c031ea --- /dev/null +++ b/crates/burn-train/src/metric/vision/dists/mod.rs @@ -0,0 +1,14 @@ +//! DISTS (Deep Image Structure and Texture Similarity) metric. +//! +//! This module implements DISTS, a full-reference image quality assessment metric +//! that combines structure and texture similarity using deep features. +//! +//! Reference: "Image Quality Assessment: Unifying Structure and Texture Similarity" +//! https://arxiv.org/abs/2004.07728 + +mod l2pool; +mod metric; +mod vgg16_l2pool; +mod weights; + +pub use metric::{Dists, DistsConfig}; diff --git a/crates/burn-train/src/metric/vision/dists/vgg16_l2pool.rs b/crates/burn-train/src/metric/vision/dists/vgg16_l2pool.rs new file mode 100644 index 0000000..3429de6 --- /dev/null +++ b/crates/burn-train/src/metric/vision/dists/vgg16_l2pool.rs @@ -0,0 +1,169 @@ +//! VGG16 feature extractor with L2 Pooling for DISTS. +//! +//! This module implements the VGG16 backbone used in DISTS, +//! with L2Pooling replacing MaxPooling for smoother feature extraction. + +use burn_core as burn; + +use burn::module::Module; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn::tensor::activation::relu; +use burn_nn::PaddingConfig2d; +use burn_nn::conv::{Conv2d, Conv2dConfig}; + +use super::l2pool::{L2Pool2d, L2Pool2dConfig}; + +/// VGG16 feature extractor with L2 Pooling for DISTS. +/// +/// Extracts features from 5 stages of VGG16, using L2Pooling +/// instead of MaxPooling for smoother downsampling. +/// +/// Output channels per stage: [64, 128, 256, 512, 512] +#[derive(Module, Debug)] +pub struct Vgg16L2PoolExtractor { + // Stage 1: 2 conv layers, 64 channels + pub(crate) conv1_1: Conv2d, + pub(crate) conv1_2: Conv2d, + pub(crate) pool1: L2Pool2d, + + // Stage 2: 2 conv layers, 128 channels + pub(crate) conv2_1: Conv2d, + pub(crate) conv2_2: Conv2d, + pub(crate) pool2: L2Pool2d, + + // Stage 3: 3 conv layers, 256 channels + pub(crate) conv3_1: Conv2d, + pub(crate) conv3_2: Conv2d, + pub(crate) conv3_3: Conv2d, + pub(crate) pool3: L2Pool2d, + + // Stage 4: 3 conv layers, 512 channels + pub(crate) conv4_1: Conv2d, + pub(crate) conv4_2: Conv2d, + pub(crate) conv4_3: Conv2d, + pub(crate) pool4: L2Pool2d, + + // Stage 5: 3 conv layers, 512 channels + pub(crate) conv5_1: Conv2d, + pub(crate) conv5_2: Conv2d, + pub(crate) conv5_3: Conv2d, +} + +impl Vgg16L2PoolExtractor { + /// Create a new VGG16 feature extractor with L2 Pooling. + pub fn new(device: &Device) -> Self { + let pool_config = L2Pool2dConfig::default(); + + Self { + // Stage 1 + conv1_1: Conv2dConfig::new([3, 64], [3, 3]) + .with_padding(PaddingConfig2d::Same) + .init(device), + conv1_2: Conv2dConfig::new([64, 64], [3, 3]) + .with_padding(PaddingConfig2d::Same) + .init(device), + pool1: pool_config.init(64, device), + + // Stage 2 + conv2_1: Conv2dConfig::new([64, 128], [3, 3]) + .with_padding(PaddingConfig2d::Same) + .init(device), + conv2_2: Conv2dConfig::new([128, 128], [3, 3]) + .with_padding(PaddingConfig2d::Same) + .init(device), + pool2: pool_config.init(128, device), + + // Stage 3 + conv3_1: Conv2dConfig::new([128, 256], [3, 3]) + .with_padding(PaddingConfig2d::Same) + .init(device), + conv3_2: Conv2dConfig::new([256, 256], [3, 3]) + .with_padding(PaddingConfig2d::Same) + .init(device), + conv3_3: Conv2dConfig::new([256, 256], [3, 3]) + .with_padding(PaddingConfig2d::Same) + .init(device), + pool3: pool_config.init(256, device), + + // Stage 4 + conv4_1: Conv2dConfig::new([256, 512], [3, 3]) + .with_padding(PaddingConfig2d::Same) + .init(device), + conv4_2: Conv2dConfig::new([512, 512], [3, 3]) + .with_padding(PaddingConfig2d::Same) + .init(device), + conv4_3: Conv2dConfig::new([512, 512], [3, 3]) + .with_padding(PaddingConfig2d::Same) + .init(device), + pool4: pool_config.init(512, device), + + // Stage 5 + conv5_1: Conv2dConfig::new([512, 512], [3, 3]) + .with_padding(PaddingConfig2d::Same) + .init(device), + conv5_2: Conv2dConfig::new([512, 512], [3, 3]) + .with_padding(PaddingConfig2d::Same) + .init(device), + conv5_3: Conv2dConfig::new([512, 512], [3, 3]) + .with_padding(PaddingConfig2d::Same) + .init(device), + } + } + + /// Extract features from all 5 stages. + /// + /// # Arguments + /// + /// * `x` - Input tensor of shape `[batch, 3, height, width]` + /// + /// # Returns + /// + /// Vector of 6 feature tensors: + /// - Stage 0: Input image [batch, 3, H, W] + /// - Stage 1: After conv1 [batch, 64, H/2, W/2] + /// - Stage 2: After conv2 [batch, 128, H/4, W/4] + /// - Stage 3: After conv3 [batch, 256, H/8, W/8] + /// - Stage 4: After conv4 [batch, 512, H/16, W/16] + /// - Stage 5: After conv5 [batch, 512, H/32, W/32] + pub fn forward(&self, x: Tensor<4>) -> Vec> { + let mut features = Vec::with_capacity(6); + + // Stage 0: Input image + features.push(x.clone()); + + // Stage 1 + let x = relu(self.conv1_1.forward(x)); + let x = relu(self.conv1_2.forward(x)); + features.push(x.clone()); + let x = self.pool1.forward(x); + + // Stage 2 + let x = relu(self.conv2_1.forward(x)); + let x = relu(self.conv2_2.forward(x)); + features.push(x.clone()); + let x = self.pool2.forward(x); + + // Stage 3 + let x = relu(self.conv3_1.forward(x)); + let x = relu(self.conv3_2.forward(x)); + let x = relu(self.conv3_3.forward(x)); + features.push(x.clone()); + let x = self.pool3.forward(x); + + // Stage 4 + let x = relu(self.conv4_1.forward(x)); + let x = relu(self.conv4_2.forward(x)); + let x = relu(self.conv4_3.forward(x)); + features.push(x.clone()); + let x = self.pool4.forward(x); + + // Stage 5 + let x = relu(self.conv5_1.forward(x)); + let x = relu(self.conv5_2.forward(x)); + let x = relu(self.conv5_3.forward(x)); + features.push(x); + + features + } +} diff --git a/crates/burn-train/src/metric/vision/dists/weights.rs b/crates/burn-train/src/metric/vision/dists/weights.rs new file mode 100644 index 0000000..1c32297 --- /dev/null +++ b/crates/burn-train/src/metric/vision/dists/weights.rs @@ -0,0 +1,125 @@ +//! Pretrained weights loading for DISTS. + +use burn_std::network::downloader::download_file_as_bytes; +use burn_store::{ModuleSnapshot, PytorchStore}; +use std::fs::{File, create_dir_all}; +use std::io::Write; +use std::path::PathBuf; + +use super::metric::Dists; + +/// URL for pretrained DISTS alpha/beta weights from the official repository. +/// Reference: https://github.com/dingkeyan93/DISTS +const DISTS_WEIGHTS_URL: &str = + "https://github.com/dingkeyan93/DISTS/raw/master/DISTS_pytorch/weights.pt"; + +/// URL for ImageNet pretrained VGG16 backbone weights from PyTorch. +const VGG16_IMAGENET_URL: &str = "https://download.pytorch.org/models/vgg16-397923af.pth"; + +/// Get the cache directory for DISTS weights. +fn get_cache_dir() -> PathBuf { + let cache_dir = dirs::cache_dir() + .expect("Could not get cache directory") + .join("burn-dataset") + .join("dists"); + + if !cache_dir.exists() { + create_dir_all(&cache_dir).expect("Failed to create cache directory"); + } + + cache_dir +} + +/// Download file if not cached. +fn download_if_needed(url: &str, cache_path: &PathBuf, message: &str) { + if !cache_path.exists() { + let bytes = download_file_as_bytes(url, message); + let mut file = File::create(cache_path).expect("Failed to create cache file"); + file.write_all(&bytes).expect("Failed to write weights"); + } +} + +/// Download and load pretrained weights into a DISTS module. +/// +/// This loads both: +/// 1. ImageNet pretrained VGG16 backbone weights +/// 2. DISTS trained alpha/beta weights +/// +/// Weights are cached in the user's cache directory to avoid re-downloading. +/// +/// # Arguments +/// +/// * `dists` - The DISTS module to load weights into. +/// +/// # Returns +/// +/// The DISTS module with loaded pretrained weights. +pub fn load_pretrained_weights(mut dists: Dists) -> Dists { + let cache_dir = get_cache_dir(); + + // Step 1: Download and load VGG16 ImageNet backbone weights + let vgg_cache_path = cache_dir.join("vgg16_backbone.pth"); + download_if_needed( + VGG16_IMAGENET_URL, + &vgg_cache_path, + "Downloading VGG16 ImageNet weights for DISTS...", + ); + + // Step 2: Download DISTS alpha/beta weights + let dists_cache_path = cache_dir.join("dists_weights.pt"); + download_if_needed( + DISTS_WEIGHTS_URL, + &dists_cache_path, + "Downloading DISTS alpha/beta weights...", + ); + + // Load VGG16 backbone weights first + dists = load_vgg16_backbone_weights(dists, &vgg_cache_path); + + // Then load DISTS alpha/beta weights + dists = load_dists_weights(dists, &dists_cache_path); + + dists +} + +/// Load VGG16 ImageNet pretrained backbone weights. +fn load_vgg16_backbone_weights(mut dists: Dists, cache_path: &PathBuf) -> Dists { + let mut store = PytorchStore::from_file(cache_path) + .allow_partial(true) + .skip_enum_variants(true) + // VGG16 features.X -> extractor.convY_Z + .with_key_remapping(r"^features\.0\.", "extractor.conv1_1.") + .with_key_remapping(r"^features\.2\.", "extractor.conv1_2.") + .with_key_remapping(r"^features\.5\.", "extractor.conv2_1.") + .with_key_remapping(r"^features\.7\.", "extractor.conv2_2.") + .with_key_remapping(r"^features\.10\.", "extractor.conv3_1.") + .with_key_remapping(r"^features\.12\.", "extractor.conv3_2.") + .with_key_remapping(r"^features\.14\.", "extractor.conv3_3.") + .with_key_remapping(r"^features\.17\.", "extractor.conv4_1.") + .with_key_remapping(r"^features\.19\.", "extractor.conv4_2.") + .with_key_remapping(r"^features\.21\.", "extractor.conv4_3.") + .with_key_remapping(r"^features\.24\.", "extractor.conv5_1.") + .with_key_remapping(r"^features\.26\.", "extractor.conv5_2.") + .with_key_remapping(r"^features\.28\.", "extractor.conv5_3."); + + let result = dists.load_from(&mut store); + if let Err(e) = result { + log::warn!("Some VGG16 backbone weights could not be loaded: {:?}", e); + } + + dists +} + +/// Load DISTS trained alpha/beta weights. +fn load_dists_weights(mut dists: Dists, cache_path: &PathBuf) -> Dists { + let mut store = PytorchStore::from_file(cache_path) + .allow_partial(true) + .skip_enum_variants(true); + + let result = dists.load_from(&mut store); + if let Err(e) = result { + log::warn!("Some DISTS weights could not be loaded: {:?}", e); + } + + dists +} diff --git a/crates/burn-train/src/metric/vision/fid/inception.rs b/crates/burn-train/src/metric/vision/fid/inception.rs new file mode 100644 index 0000000..d4a0ea5 --- /dev/null +++ b/crates/burn-train/src/metric/vision/fid/inception.rs @@ -0,0 +1,447 @@ +//! InceptionV3 feature extractor for FID (pytorch-fid variant with TF-ported weights). +//! +//! Reference: + +use burn_core as burn; + +use burn::module::Module; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn::tensor::activation::relu; +use burn::tensor::ops::{InterpolateMode, InterpolateOptions}; +use burn_nn::conv::{Conv2d, Conv2dConfig}; +use burn_nn::{BatchNorm, BatchNormConfig, PaddingConfig2d}; + +/// Conv2d + BatchNorm + ReLU building block. +#[derive(Module, Debug)] +pub struct BasicConv2d { + conv: Conv2d, + bn: BatchNorm, +} + +impl BasicConv2d { + pub fn new(conv_config: Conv2dConfig, device: &Device) -> Self { + let out_channels = conv_config.channels[1]; + Self { + conv: conv_config.with_bias(false).init(device), + bn: BatchNormConfig::new(out_channels) + .with_epsilon(0.001) + .init(device), + } + } + + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + relu(self.bn.forward(self.conv.forward(x))) + } +} + +#[derive(Module, Debug)] +pub struct InceptionA { + branch1x1: BasicConv2d, + branch5x5_1: BasicConv2d, + branch5x5_2: BasicConv2d, + branch3x3dbl_1: BasicConv2d, + branch3x3dbl_2: BasicConv2d, + branch3x3dbl_3: BasicConv2d, + branch_pool: BasicConv2d, +} + +impl InceptionA { + pub fn new(in_channels: usize, pool_features: usize, device: &Device) -> Self { + Self { + branch1x1: BasicConv2d::new(Conv2dConfig::new([in_channels, 64], [1, 1]), device), + branch5x5_1: BasicConv2d::new(Conv2dConfig::new([in_channels, 48], [1, 1]), device), + branch5x5_2: BasicConv2d::new( + Conv2dConfig::new([48, 64], [5, 5]) + .with_padding(PaddingConfig2d::Explicit(2, 2, 2, 2)), + device, + ), + branch3x3dbl_1: BasicConv2d::new(Conv2dConfig::new([in_channels, 64], [1, 1]), device), + branch3x3dbl_2: BasicConv2d::new( + Conv2dConfig::new([64, 96], [3, 3]) + .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1)), + device, + ), + branch3x3dbl_3: BasicConv2d::new( + Conv2dConfig::new([96, 96], [3, 3]) + .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1)), + device, + ), + branch_pool: BasicConv2d::new( + Conv2dConfig::new([in_channels, pool_features], [1, 1]), + device, + ), + } + } + + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + let branch1x1 = self.branch1x1.forward(x.clone()); + + let branch5x5 = self.branch5x5_1.forward(x.clone()); + let branch5x5 = self.branch5x5_2.forward(branch5x5); + + let branch3x3dbl = self.branch3x3dbl_1.forward(x.clone()); + let branch3x3dbl = self.branch3x3dbl_2.forward(branch3x3dbl); + let branch3x3dbl = self.branch3x3dbl_3.forward(branch3x3dbl); + + let branch_pool = + burn_core::tensor::module::avg_pool2d(x, [3, 3], [1, 1], [1, 1], false, false); + let branch_pool = self.branch_pool.forward(branch_pool); + + Tensor::cat(vec![branch1x1, branch5x5, branch3x3dbl, branch_pool], 1) + } +} + +#[derive(Module, Debug)] +pub struct InceptionB { + branch3x3: BasicConv2d, + branch3x3dbl_1: BasicConv2d, + branch3x3dbl_2: BasicConv2d, + branch3x3dbl_3: BasicConv2d, +} + +impl InceptionB { + pub fn new(in_channels: usize, device: &Device) -> Self { + Self { + branch3x3: BasicConv2d::new( + Conv2dConfig::new([in_channels, 384], [3, 3]).with_stride([2, 2]), + device, + ), + branch3x3dbl_1: BasicConv2d::new(Conv2dConfig::new([in_channels, 64], [1, 1]), device), + branch3x3dbl_2: BasicConv2d::new( + Conv2dConfig::new([64, 96], [3, 3]) + .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1)), + device, + ), + branch3x3dbl_3: BasicConv2d::new( + Conv2dConfig::new([96, 96], [3, 3]).with_stride([2, 2]), + device, + ), + } + } + + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + let branch3x3 = self.branch3x3.forward(x.clone()); + + let branch3x3dbl = self.branch3x3dbl_1.forward(x.clone()); + let branch3x3dbl = self.branch3x3dbl_2.forward(branch3x3dbl); + let branch3x3dbl = self.branch3x3dbl_3.forward(branch3x3dbl); + + let branch_pool = + burn_core::tensor::module::max_pool2d(x, [3, 3], [2, 2], [0, 0], [1, 1], false); + + Tensor::cat(vec![branch3x3, branch3x3dbl, branch_pool], 1) + } +} + +#[derive(Module, Debug)] +pub struct InceptionC { + branch1x1: BasicConv2d, + branch7x7_1: BasicConv2d, + branch7x7_2: BasicConv2d, + branch7x7_3: BasicConv2d, + branch7x7dbl_1: BasicConv2d, + branch7x7dbl_2: BasicConv2d, + branch7x7dbl_3: BasicConv2d, + branch7x7dbl_4: BasicConv2d, + branch7x7dbl_5: BasicConv2d, + branch_pool: BasicConv2d, +} + +impl InceptionC { + pub fn new(in_channels: usize, channels_7x7: usize, device: &Device) -> Self { + let c7 = channels_7x7; + Self { + branch1x1: BasicConv2d::new(Conv2dConfig::new([in_channels, 192], [1, 1]), device), + branch7x7_1: BasicConv2d::new(Conv2dConfig::new([in_channels, c7], [1, 1]), device), + branch7x7_2: BasicConv2d::new( + Conv2dConfig::new([c7, c7], [1, 7]) + .with_padding(PaddingConfig2d::Explicit(0, 3, 0, 3)), + device, + ), + branch7x7_3: BasicConv2d::new( + Conv2dConfig::new([c7, 192], [7, 1]) + .with_padding(PaddingConfig2d::Explicit(3, 0, 3, 0)), + device, + ), + branch7x7dbl_1: BasicConv2d::new(Conv2dConfig::new([in_channels, c7], [1, 1]), device), + branch7x7dbl_2: BasicConv2d::new( + Conv2dConfig::new([c7, c7], [7, 1]) + .with_padding(PaddingConfig2d::Explicit(3, 0, 3, 0)), + device, + ), + branch7x7dbl_3: BasicConv2d::new( + Conv2dConfig::new([c7, c7], [1, 7]) + .with_padding(PaddingConfig2d::Explicit(0, 3, 0, 3)), + device, + ), + branch7x7dbl_4: BasicConv2d::new( + Conv2dConfig::new([c7, c7], [7, 1]) + .with_padding(PaddingConfig2d::Explicit(3, 0, 3, 0)), + device, + ), + branch7x7dbl_5: BasicConv2d::new( + Conv2dConfig::new([c7, 192], [1, 7]) + .with_padding(PaddingConfig2d::Explicit(0, 3, 0, 3)), + device, + ), + branch_pool: BasicConv2d::new(Conv2dConfig::new([in_channels, 192], [1, 1]), device), + } + } + + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + let branch1x1 = self.branch1x1.forward(x.clone()); + + let branch7x7 = self.branch7x7_1.forward(x.clone()); + let branch7x7 = self.branch7x7_2.forward(branch7x7); + let branch7x7 = self.branch7x7_3.forward(branch7x7); + + let branch7x7dbl = self.branch7x7dbl_1.forward(x.clone()); + let branch7x7dbl = self.branch7x7dbl_2.forward(branch7x7dbl); + let branch7x7dbl = self.branch7x7dbl_3.forward(branch7x7dbl); + let branch7x7dbl = self.branch7x7dbl_4.forward(branch7x7dbl); + let branch7x7dbl = self.branch7x7dbl_5.forward(branch7x7dbl); + + let branch_pool = + burn_core::tensor::module::avg_pool2d(x, [3, 3], [1, 1], [1, 1], false, false); + let branch_pool = self.branch_pool.forward(branch_pool); + + Tensor::cat(vec![branch1x1, branch7x7, branch7x7dbl, branch_pool], 1) + } +} + +#[derive(Module, Debug)] +pub struct InceptionD { + branch3x3_1: BasicConv2d, + branch3x3_2: BasicConv2d, + branch7x7x3_1: BasicConv2d, + branch7x7x3_2: BasicConv2d, + branch7x7x3_3: BasicConv2d, + branch7x7x3_4: BasicConv2d, +} + +impl InceptionD { + pub fn new(in_channels: usize, device: &Device) -> Self { + Self { + branch3x3_1: BasicConv2d::new(Conv2dConfig::new([in_channels, 192], [1, 1]), device), + branch3x3_2: BasicConv2d::new( + Conv2dConfig::new([192, 320], [3, 3]).with_stride([2, 2]), + device, + ), + branch7x7x3_1: BasicConv2d::new(Conv2dConfig::new([in_channels, 192], [1, 1]), device), + branch7x7x3_2: BasicConv2d::new( + Conv2dConfig::new([192, 192], [1, 7]) + .with_padding(PaddingConfig2d::Explicit(0, 3, 0, 3)), + device, + ), + branch7x7x3_3: BasicConv2d::new( + Conv2dConfig::new([192, 192], [7, 1]) + .with_padding(PaddingConfig2d::Explicit(3, 0, 3, 0)), + device, + ), + branch7x7x3_4: BasicConv2d::new( + Conv2dConfig::new([192, 192], [3, 3]).with_stride([2, 2]), + device, + ), + } + } + + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + let branch3x3 = self.branch3x3_1.forward(x.clone()); + let branch3x3 = self.branch3x3_2.forward(branch3x3); + + let branch7x7x3 = self.branch7x7x3_1.forward(x.clone()); + let branch7x7x3 = self.branch7x7x3_2.forward(branch7x7x3); + let branch7x7x3 = self.branch7x7x3_3.forward(branch7x7x3); + let branch7x7x3 = self.branch7x7x3_4.forward(branch7x7x3); + + let branch_pool = + burn_core::tensor::module::max_pool2d(x, [3, 3], [2, 2], [0, 0], [1, 1], false); + + Tensor::cat(vec![branch3x3, branch7x7x3, branch_pool], 1) + } +} + +#[derive(Module, Debug)] +pub struct InceptionE { + branch1x1: BasicConv2d, + branch3x3_1: BasicConv2d, + branch3x3_2a: BasicConv2d, + branch3x3_2b: BasicConv2d, + branch3x3dbl_1: BasicConv2d, + branch3x3dbl_2: BasicConv2d, + branch3x3dbl_3a: BasicConv2d, + branch3x3dbl_3b: BasicConv2d, + branch_pool: BasicConv2d, + #[module(skip)] + use_max_pool: bool, +} + +impl InceptionE { + pub fn new(in_channels: usize, use_max_pool: bool, device: &Device) -> Self { + Self { + branch1x1: BasicConv2d::new(Conv2dConfig::new([in_channels, 320], [1, 1]), device), + branch3x3_1: BasicConv2d::new(Conv2dConfig::new([in_channels, 384], [1, 1]), device), + branch3x3_2a: BasicConv2d::new( + Conv2dConfig::new([384, 384], [1, 3]) + .with_padding(PaddingConfig2d::Explicit(0, 1, 0, 1)), + device, + ), + branch3x3_2b: BasicConv2d::new( + Conv2dConfig::new([384, 384], [3, 1]) + .with_padding(PaddingConfig2d::Explicit(1, 0, 1, 0)), + device, + ), + branch3x3dbl_1: BasicConv2d::new(Conv2dConfig::new([in_channels, 448], [1, 1]), device), + branch3x3dbl_2: BasicConv2d::new( + Conv2dConfig::new([448, 384], [3, 3]) + .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1)), + device, + ), + branch3x3dbl_3a: BasicConv2d::new( + Conv2dConfig::new([384, 384], [1, 3]) + .with_padding(PaddingConfig2d::Explicit(0, 1, 0, 1)), + device, + ), + branch3x3dbl_3b: BasicConv2d::new( + Conv2dConfig::new([384, 384], [3, 1]) + .with_padding(PaddingConfig2d::Explicit(1, 0, 1, 0)), + device, + ), + branch_pool: BasicConv2d::new(Conv2dConfig::new([in_channels, 192], [1, 1]), device), + use_max_pool, + } + } + + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + let branch1x1 = self.branch1x1.forward(x.clone()); + + let branch3x3 = self.branch3x3_1.forward(x.clone()); + let branch3x3_a = self.branch3x3_2a.forward(branch3x3.clone()); + let branch3x3_b = self.branch3x3_2b.forward(branch3x3); + let branch3x3 = Tensor::cat(vec![branch3x3_a, branch3x3_b], 1); + + let branch3x3dbl = self.branch3x3dbl_1.forward(x.clone()); + let branch3x3dbl = self.branch3x3dbl_2.forward(branch3x3dbl); + let branch3x3dbl_a = self.branch3x3dbl_3a.forward(branch3x3dbl.clone()); + let branch3x3dbl_b = self.branch3x3dbl_3b.forward(branch3x3dbl); + let branch3x3dbl = Tensor::cat(vec![branch3x3dbl_a, branch3x3dbl_b], 1); + + let branch_pool = if self.use_max_pool { + burn_core::tensor::module::max_pool2d(x, [3, 3], [1, 1], [1, 1], [1, 1], false) + } else { + burn_core::tensor::module::avg_pool2d(x, [3, 3], [1, 1], [1, 1], false, false) + }; + let branch_pool = self.branch_pool.forward(branch_pool); + + Tensor::cat(vec![branch1x1, branch3x3, branch3x3dbl, branch_pool], 1) + } +} + +/// InceptionV3 feature extractor for FID computation. +/// +/// Outputs a 2048-dimensional feature vector per image, matching the +/// pytorch-fid variant (TF-ported weights). +#[derive(Module, Debug)] +pub struct InceptionV3FeatureExtractor { + // Stem + conv2d_1a: BasicConv2d, + conv2d_2a: BasicConv2d, + conv2d_2b: BasicConv2d, + conv2d_3b: BasicConv2d, + conv2d_4a: BasicConv2d, + // Inception blocks + mixed_5b: InceptionA, + mixed_5c: InceptionA, + mixed_5d: InceptionA, + mixed_6a: InceptionB, + mixed_6b: InceptionC, + mixed_6c: InceptionC, + mixed_6d: InceptionC, + mixed_6e: InceptionC, + mixed_7a: InceptionD, + mixed_7b: InceptionE, + mixed_7c: InceptionE, +} + +impl InceptionV3FeatureExtractor { + /// Creates a new feature extractor with random weights. + pub fn new(device: &Device) -> Self { + Self { + // Stem: 3 -> 32 -> 32 -> 64 -> 80 -> 192 + conv2d_1a: BasicConv2d::new( + Conv2dConfig::new([3, 32], [3, 3]).with_stride([2, 2]), + device, + ), + conv2d_2a: BasicConv2d::new(Conv2dConfig::new([32, 32], [3, 3]), device), + conv2d_2b: BasicConv2d::new( + Conv2dConfig::new([32, 64], [3, 3]) + .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1)), + device, + ), + conv2d_3b: BasicConv2d::new(Conv2dConfig::new([64, 80], [1, 1]), device), + conv2d_4a: BasicConv2d::new(Conv2dConfig::new([80, 192], [3, 3]), device), + mixed_5b: InceptionA::new(192, 32, device), + mixed_5c: InceptionA::new(256, 64, device), + mixed_5d: InceptionA::new(288, 64, device), + mixed_6a: InceptionB::new(288, device), + mixed_6b: InceptionC::new(768, 128, device), + mixed_6c: InceptionC::new(768, 160, device), + mixed_6d: InceptionC::new(768, 160, device), + mixed_6e: InceptionC::new(768, 192, device), + mixed_7a: InceptionD::new(768, device), + mixed_7b: InceptionE::new(1280, false, device), + mixed_7c: InceptionE::new(2048, true, device), + } + } + + /// Extract 2048-dim features. Input is resized to 299x299 via bilinear + /// interpolation to match the pytorch-fid reference. + pub fn forward(&self, x: Tensor<4>) -> Tensor<2> { + let [batch, _, h, w] = x.dims(); + + let x = if h != 299 || w != 299 { + burn_core::tensor::module::interpolate( + x, + [299, 299], + InterpolateOptions::new(InterpolateMode::Bilinear), + ) + } else { + x + }; + + // Stem + let x = self.conv2d_1a.forward(x); + let x = self.conv2d_2a.forward(x); + let x = self.conv2d_2b.forward(x); + let x = burn_core::tensor::module::max_pool2d(x, [3, 3], [2, 2], [0, 0], [1, 1], false); + let x = self.conv2d_3b.forward(x); + let x = self.conv2d_4a.forward(x); + let x = burn_core::tensor::module::max_pool2d(x, [3, 3], [2, 2], [0, 0], [1, 1], false); + + // InceptionA + let x = self.mixed_5b.forward(x); + let x = self.mixed_5c.forward(x); + let x = self.mixed_5d.forward(x); + + // InceptionB (reduction) + let x = self.mixed_6a.forward(x); + + // InceptionC + let x = self.mixed_6b.forward(x); + let x = self.mixed_6c.forward(x); + let x = self.mixed_6d.forward(x); + let x = self.mixed_6e.forward(x); + + // InceptionD (reduction) + let x = self.mixed_7a.forward(x); + + // InceptionE + let x = self.mixed_7b.forward(x); + let x = self.mixed_7c.forward(x); + + // Global average pool -> [N, 2048] + let x = burn_core::tensor::module::adaptive_avg_pool2d(x, [1, 1]); + x.reshape([batch, 2048]) + } +} diff --git a/crates/burn-train/src/metric/vision/fid/metric.rs b/crates/burn-train/src/metric/vision/fid/metric.rs new file mode 100644 index 0000000..cfec4e0 --- /dev/null +++ b/crates/burn-train/src/metric/vision/fid/metric.rs @@ -0,0 +1,345 @@ +use burn_core as burn; + +use burn::config::Config; +use burn::module::Module; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn::tensor::linalg; + +use super::inception::InceptionV3FeatureExtractor; + +const IMAGENET_MEAN: [f32; 3] = [0.485, 0.456, 0.406]; +const IMAGENET_STD: [f32; 3] = [0.229, 0.224, 0.225]; +const EPS: f64 = 1e-6; + +/// Configuration for [Fid]. +/// +/// ```ignore +/// let fid = FidConfig::new().init_pretrained(&device); +/// let score = fid.forward(real_images, generated_images); +/// ``` +#[derive(Config, Debug)] +pub struct FidConfig { + /// Normalize input images from [0,1] to ImageNet range. + #[config(default = true)] + pub normalize: bool, + + /// Number of Newton-Schulz iterations for matrix square root. + #[config(default = 50)] + pub num_iterations: usize, +} + +impl FidConfig { + /// Initialize with pretrained InceptionV3 weights from pytorch-fid. + pub fn init_pretrained(&self, device: &Device) -> Fid { + let fid = self.init(device); + super::weights::load_pretrained_weights(fid) + } + + /// Initialize with random weights. + pub fn init(&self, device: &Device) -> Fid { + Fid { + extractor: InceptionV3FeatureExtractor::new(device), + normalize: self.normalize, + num_iterations: self.num_iterations, + } + } +} + +/// Frechet Inception Distance metric. +/// +/// Computes the Frechet distance between feature distributions of real and +/// generated images using an InceptionV3 feature extractor. +/// +/// ```ignore +/// let fid = FidConfig::new().init_pretrained(&device); +/// let feats_real = fid.extract_features(real_images); +/// let feats_gen = fid.extract_features(generated_images); +/// let score = fid.compute_fid(feats_real, feats_gen); +/// ``` +#[derive(Module, Debug)] +pub struct Fid { + extractor: InceptionV3FeatureExtractor, + normalize: bool, + num_iterations: usize, +} + +impl Fid { + /// Extract 2048-dim InceptionV3 features from images of shape `[batch, 3, H, W]`. + pub fn extract_features(&self, images: Tensor<4>) -> Tensor<2> { + let images = if self.normalize { + imagenet_normalize(images) + } else { + images + }; + self.extractor.forward(images) + } + + /// Compute FID from pre-extracted feature tensors of shape `[N, 2048]`. + pub fn compute_fid(&self, features_real: Tensor<2>, features_gen: Tensor<2>) -> Tensor<1> { + let (mu1, sigma1) = compute_statistics(features_real); + let (mu2, sigma2) = compute_statistics(features_gen); + frechet_distance(mu1, sigma1, mu2, sigma2, self.num_iterations) + } + + /// Compute FID end-to-end from image tensors of shape `[N, 3, H, W]` in [0,1]. + pub fn forward(&self, images_real: Tensor<4>, images_gen: Tensor<4>) -> Tensor<1> { + let features_real = self.extract_features(images_real); + let features_gen = self.extract_features(images_gen); + self.compute_fid(features_real, features_gen) + } +} + +fn imagenet_normalize(x: Tensor<4>) -> Tensor<4> { + let device = x.device(); + + let mean = Tensor::<1>::from_floats(IMAGENET_MEAN, &device).reshape([1, 3, 1, 1]); + let std = Tensor::<1>::from_floats(IMAGENET_STD, &device).reshape([1, 3, 1, 1]); + + x.sub(mean).div(std) +} + +/// Mean vector `[D]` and unbiased covariance matrix `[D, D]` from feature rows `[N, D]`. +fn compute_statistics(features: Tensor<2>) -> (Tensor<1>, Tensor<2>) { + let [n, d] = features.dims(); + let n_f = n as f64; + + let mean = features.clone().mean_dim(0).squeeze_dim::<1>(0); + let centered = features.sub(mean.clone().unsqueeze_dim::<2>(0).expand([n, d])); + + let cov = centered + .clone() + .transpose() + .matmul(centered) + .div_scalar(n_f - 1.0); + + (mean, cov) +} + +/// Newton-Schulz iteration for the square root of a symmetric PD matrix. +/// Input must be symmetric positive-definite for convergence. +fn matrix_sqrt_newton_schulz(a: Tensor<2>, num_iterations: usize) -> Tensor<2> { + let [d, _] = a.dims(); + let device = a.device(); + + // Clamp to avoid division by near-zero norms (also avoids a GPU sync). + let norm_a = a.clone().mul(a.clone()).sum().sqrt().clamp_min(EPS); + + let identity = Tensor::<2>::eye(d, &device); + let mut y = a.div(norm_a.clone().unsqueeze_dim::<2>(0).expand([d, d])); + let mut z = identity.clone(); + let three_i = identity.clone().mul_scalar(3.0); + + for _ in 0..num_iterations { + let t = three_i + .clone() + .sub(z.clone().matmul(y.clone())) + .mul_scalar(0.5); + y = y.matmul(t.clone()); + z = t.matmul(z); + } + + let sqrt_norm = norm_a.sqrt().unsqueeze_dim::<2>(0).expand([d, d]); + y.mul(sqrt_norm) +} + +/// Frechet distance between two multivariate Gaussians. +/// +/// Uses the symmetric form (S @ sigma2 @ S where S = sqrtm(sigma1)) so that +/// Newton-Schulz converges — the naive sqrtm(sigma1 @ sigma2) is non-symmetric. +fn frechet_distance( + mu1: Tensor<1>, + sigma1: Tensor<2>, + mu2: Tensor<1>, + sigma2: Tensor<2>, + num_iterations: usize, +) -> Tensor<1> { + let [d, _] = sigma1.dims(); + let device = sigma1.device(); + + let diff = mu1.sub(mu2); + let mean_term = diff.clone().mul(diff).sum(); + + // Small regularization (eps · I) scaled to the average variance for numerical + // stability with near-singular covariances. Done entirely with tensor ops to + // avoid forcing a GPU sync. + let tr_sum = linalg::trace::<2, 1>(sigma1.clone()).add(linalg::trace::<2, 1>(sigma2.clone())); + let avg_variance = tr_sum.div_scalar(2.0 * d as f64).clamp_min(EPS); + let reg = Tensor::<2>::eye(d, &device).mul(avg_variance.mul_scalar(EPS).unsqueeze_dim::<2>(0)); + let sigma1 = sigma1.add(reg.clone()); + let sigma2 = sigma2.add(reg); + + let sqrt_sigma1 = matrix_sqrt_newton_schulz(sigma1.clone(), num_iterations); + let m = sqrt_sigma1 + .clone() + .matmul(sigma2.clone()) + .matmul(sqrt_sigma1); + let sqrt_m = matrix_sqrt_newton_schulz(m, num_iterations); + + let cov_term = sigma1.add(sigma2).sub(sqrt_m.mul_scalar(2.0)); + let trace = linalg::trace::<2, 1>(cov_term); + + mean_term.add(trace).reshape([1]) +} + +#[cfg(test)] +mod tests { + use super::*; + use burn_core::tensor::{TensorData, Tolerance}; + + type FT = f32; + + #[test] + fn test_newton_schulz_identity() { + let device = Default::default(); + let identity = Tensor::<2>::eye(3, &device); + let sqrt_i = matrix_sqrt_newton_schulz(identity.clone(), 50); + + sqrt_i + .into_data() + .assert_approx_eq::(&identity.into_data(), Tolerance::relative(1e-4)); + } + + #[test] + fn test_newton_schulz_diagonal() { + // sqrt(diag(4, 9, 16)) should be diag(2, 3, 4) + let device = Default::default(); + let a = Tensor::<2>::from_floats( + [[4.0, 0.0, 0.0], [0.0, 9.0, 0.0], [0.0, 0.0, 16.0]], + &device, + ); + let expected = + Tensor::<2>::from_floats([[2.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 4.0]], &device); + + let sqrt_a = matrix_sqrt_newton_schulz(a, 50); + + sqrt_a + .into_data() + .assert_approx_eq::(&expected.into_data(), Tolerance::relative(1e-3)); + } + + #[test] + fn test_compute_statistics() { + // [[1,2],[3,4],[5,6]] -> mean [3,4], centered [[-2,-2],[0,0],[2,2]] + // cov = [[8,8],[8,8]] / 2 = [[4,4],[4,4]] + let device = Default::default(); + let features = Tensor::<2>::from_floats([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], &device); + + let (mean, cov) = compute_statistics(features); + + mean.into_data() + .assert_approx_eq::(&TensorData::from([3.0_f32, 4.0]), Tolerance::default()); + cov.into_data().assert_approx_eq::( + &TensorData::from([[4.0_f32, 4.0], [4.0, 4.0]]), + Tolerance::default(), + ); + } + + #[test] + fn test_fid_same_features_is_zero() { + let device = Default::default(); + let features = Tensor::<2>::from_floats( + [ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [1.0, 1.0, 0.0], + [0.0, 1.0, 1.0], + ], + &device, + ); + + let (mu, sigma) = compute_statistics(features.clone()); + let fid = frechet_distance(mu.clone(), sigma.clone(), mu, sigma, 50); + + assert!(fid.into_data().to_vec::().unwrap()[0].abs() < 0.1); + } + + #[test] + fn test_fid_shifted_mean() { + // Shift one distribution by [2, 0, 0] — FID should be ~||shift||² = 4.0 + let device = Default::default(); + + // Handcrafted feature rows with some spread + let base = Tensor::<2>::from_floats( + [ + [-0.3, 0.1, 0.4], + [0.2, -0.5, 0.1], + [0.5, 0.3, -0.2], + [-0.1, 0.4, -0.3], + [0.0, -0.2, 0.5], + [0.3, 0.0, -0.4], + [-0.4, 0.2, 0.3], + [0.1, -0.1, 0.0], + [0.4, 0.5, -0.1], + [-0.2, -0.3, 0.2], + ], + &device, + ); + + // Same data but shifted by [2, 0, 0] + let shift = Tensor::<2>::from_floats([[2.0, 0.0, 0.0]], &device).expand([10, 3]); + let shifted = base.clone().add(shift); + + let (mu1, sigma1) = compute_statistics(base); + let (mu2, sigma2) = compute_statistics(shifted); + let fid_val = frechet_distance(mu1, sigma1, mu2, sigma2, 50) + .into_data() + .to_vec::() + .unwrap()[0]; + + assert!((fid_val - 4.0).abs() < 0.5); + } + + #[test] + fn test_fid_symmetry() { + let device = Default::default(); + let features1 = + Tensor::<2>::from_floats([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [0.5, 0.5]], &device); + let features2 = + Tensor::<2>::from_floats([[2.0, 1.0], [1.0, 2.0], [2.0, 2.0], [1.5, 1.5]], &device); + + let (mu1, sigma1) = compute_statistics(features1.clone()); + let (mu2, sigma2) = compute_statistics(features2.clone()); + + let fid_forward = + frechet_distance(mu1.clone(), sigma1.clone(), mu2.clone(), sigma2.clone(), 50); + let fid_reverse = frechet_distance(mu2, sigma2, mu1, sigma1, 50); + + fid_forward + .into_data() + .assert_approx_eq::(&fid_reverse.into_data(), Tolerance::relative(1e-3)); + } + + #[test] + fn test_inception_output_shape() { + let device = Default::default(); + let extractor = InceptionV3FeatureExtractor::new(&device); + let input = Tensor::<4>::zeros([1, 3, 299, 299], &device); + assert_eq!(extractor.forward(input).dims(), [1, 2048]); + } + + #[test] + fn test_fid_extract_features_shape() { + let device = Default::default(); + let fid = FidConfig::new().init(&device); + let images = Tensor::<4>::zeros([2, 3, 299, 299], &device); + assert_eq!(fid.extract_features(images).dims(), [2, 2048]); + } + + #[test] + #[ignore = "downloads pre-trained weights"] + fn test_fid_pretrained_features() { + let device = Default::default(); + let fid = FidConfig::new().init_pretrained(&device); + + let images = Tensor::<4>::ones([1, 3, 299, 299], &device).mul_scalar(0.5); + let features = fid.extract_features(images); + + assert_eq!(features.dims(), [1, 2048]); + let feat_data = features.into_data().to_vec::().unwrap(); + assert!(feat_data.iter().all(|v| v.is_finite())); + let norm: f32 = feat_data.iter().map(|v| v * v).sum(); + assert!(norm > 0.0); + } +} diff --git a/crates/burn-train/src/metric/vision/fid/mod.rs b/crates/burn-train/src/metric/vision/fid/mod.rs new file mode 100644 index 0000000..80ad173 --- /dev/null +++ b/crates/burn-train/src/metric/vision/fid/mod.rs @@ -0,0 +1,13 @@ +//! Frechet Inception Distance (FID) metric. +//! +//! Measures the distance between distributions of generated and real images +//! using InceptionV3 features. Lower FID = higher quality and diversity. +//! +//! Reference: + +mod inception; +mod metric; +mod weights; + +pub use inception::InceptionV3FeatureExtractor; +pub use metric::{Fid, FidConfig}; diff --git a/crates/burn-train/src/metric/vision/fid/weights.rs b/crates/burn-train/src/metric/vision/fid/weights.rs new file mode 100644 index 0000000..5ce06e2 --- /dev/null +++ b/crates/burn-train/src/metric/vision/fid/weights.rs @@ -0,0 +1,70 @@ +use burn_std::network::downloader::download_file_as_bytes; +use burn_store::{ModuleSnapshot, PytorchStore}; +use std::fs::{File, create_dir_all}; +use std::io::Write; +use std::path::PathBuf; + +use super::metric::Fid; + +const INCEPTION_WEIGHTS_URL: &str = "https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt-inception-2015-12-05-6726825d.pth"; + +fn get_cache_dir() -> PathBuf { + let cache_dir = dirs::cache_dir() + .expect("Could not get cache directory") + .join("burn-dataset") + .join("fid"); + + if !cache_dir.exists() { + create_dir_all(&cache_dir).expect("Failed to create cache directory"); + } + + cache_dir +} + +fn download_if_needed(url: &str, cache_path: &PathBuf, message: &str) { + if !cache_path.exists() { + let bytes = download_file_as_bytes(url, message); + let mut file = File::create(cache_path).expect("Failed to create cache file"); + file.write_all(&bytes).expect("Failed to write weights"); + } +} + +/// Download and load pretrained pytorch-fid InceptionV3 weights. +/// +/// Weights are cached in `~/.cache/burn-dataset/fid/`. +pub fn load_pretrained_weights(mut fid: Fid) -> Fid { + let cache_dir = get_cache_dir(); + let cache_path = cache_dir.join("pt-inception-2015-12-05-6726825d.pth"); + + download_if_needed( + INCEPTION_WEIGHTS_URL, + &cache_path, + "Downloading InceptionV3 weights for FID...", + ); + + // PyTorch keys use "Conv2d_1a_3x3.*" etc, our model nests them under "extractor.*" + let mut store = PytorchStore::from_file(&cache_path) + .allow_partial(true) + .with_key_remapping(r"^Conv2d_1a_3x3\.", "extractor.conv2d_1a.") + .with_key_remapping(r"^Conv2d_2a_3x3\.", "extractor.conv2d_2a.") + .with_key_remapping(r"^Conv2d_2b_3x3\.", "extractor.conv2d_2b.") + .with_key_remapping(r"^Conv2d_3b_1x1\.", "extractor.conv2d_3b.") + .with_key_remapping(r"^Conv2d_4a_3x3\.", "extractor.conv2d_4a.") + .with_key_remapping(r"^Mixed_5b\.", "extractor.mixed_5b.") + .with_key_remapping(r"^Mixed_5c\.", "extractor.mixed_5c.") + .with_key_remapping(r"^Mixed_5d\.", "extractor.mixed_5d.") + .with_key_remapping(r"^Mixed_6a\.", "extractor.mixed_6a.") + .with_key_remapping(r"^Mixed_6b\.", "extractor.mixed_6b.") + .with_key_remapping(r"^Mixed_6c\.", "extractor.mixed_6c.") + .with_key_remapping(r"^Mixed_6d\.", "extractor.mixed_6d.") + .with_key_remapping(r"^Mixed_6e\.", "extractor.mixed_6e.") + .with_key_remapping(r"^Mixed_7a\.", "extractor.mixed_7a.") + .with_key_remapping(r"^Mixed_7b\.", "extractor.mixed_7b.") + .with_key_remapping(r"^Mixed_7c\.", "extractor.mixed_7c."); + + if let Err(e) = fid.load_from(&mut store) { + log::warn!("Some InceptionV3 weights could not be loaded: {:?}", e); + } + + fid +} diff --git a/crates/burn-train/src/metric/vision/lpips/alexnet.rs b/crates/burn-train/src/metric/vision/lpips/alexnet.rs new file mode 100644 index 0000000..908a210 --- /dev/null +++ b/crates/burn-train/src/metric/vision/lpips/alexnet.rs @@ -0,0 +1,100 @@ +//! AlexNet feature extractor for LPIPS. + +use burn_core as burn; + +use burn::module::Module; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn::tensor::activation::relu; +use burn_nn::PaddingConfig2d; +use burn_nn::conv::{Conv2d, Conv2dConfig}; + +/// AlexNet feature extractor for LPIPS. +/// +/// Extracts features from 5 layers: +/// - conv1: 64 channels (after ReLU) +/// - conv2: 192 channels (after ReLU) +/// - conv3: 384 channels (after ReLU) +/// - conv4: 256 channels (after ReLU) +/// - conv5: 256 channels (after ReLU) +#[derive(Module, Debug)] +pub struct AlexFeatureExtractor { + /// Conv1: 3 -> 64, kernel 11x11, stride 4, padding 2 + conv1: Conv2d, + /// Conv2: 64 -> 192, kernel 5x5, stride 1, padding 2 + conv2: Conv2d, + /// Conv3: 192 -> 384, kernel 3x3, stride 1, padding 1 + conv3: Conv2d, + /// Conv4: 384 -> 256, kernel 3x3, stride 1, padding 1 + conv4: Conv2d, + /// Conv5: 256 -> 256, kernel 3x3, stride 1, padding 1 + conv5: Conv2d, +} + +impl AlexFeatureExtractor { + /// Create a new AlexNet feature extractor. + pub fn new(device: &Device) -> Self { + Self { + // Conv1: 3 -> 64, 11x11, stride 4, padding 2 + conv1: Conv2dConfig::new([3, 64], [11, 11]) + .with_stride([4, 4]) + .with_padding(PaddingConfig2d::Explicit(2, 2, 2, 2)) + .with_bias(true) + .init(device), + // Conv2: 64 -> 192, 5x5, stride 1, padding 2 + conv2: Conv2dConfig::new([64, 192], [5, 5]) + .with_padding(PaddingConfig2d::Explicit(2, 2, 2, 2)) + .with_bias(true) + .init(device), + // Conv3: 192 -> 384, 3x3, stride 1, padding 1 + conv3: Conv2dConfig::new([192, 384], [3, 3]) + .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1)) + .with_bias(true) + .init(device), + // Conv4: 384 -> 256, 3x3, stride 1, padding 1 + conv4: Conv2dConfig::new([384, 256], [3, 3]) + .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1)) + .with_bias(true) + .init(device), + // Conv5: 256 -> 256, 3x3, stride 1, padding 1 + conv5: Conv2dConfig::new([256, 256], [3, 3]) + .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1)) + .with_bias(true) + .init(device), + } + } + + /// Extract features from 5 AlexNet layers. + pub fn forward(&self, x: Tensor<4>) -> Vec> { + let mut features = Vec::with_capacity(5); + + // Slice 1: Conv1 + ReLU + let x = relu(self.conv1.forward(x)); + features.push(x.clone()); + + // Slice 2: MaxPool + Conv2 + ReLU + let x = max_pool2d_alex(x); + let x = relu(self.conv2.forward(x)); + features.push(x.clone()); + + // Slice 3: MaxPool + Conv3 + ReLU + let x = max_pool2d_alex(x); + let x = relu(self.conv3.forward(x)); + features.push(x.clone()); + + // Slice 4: Conv4 + ReLU (no pooling) + let x = relu(self.conv4.forward(x)); + features.push(x.clone()); + + // Slice 5: Conv5 + ReLU (no pooling) + let x = relu(self.conv5.forward(x)); + features.push(x); + + features + } +} + +/// 3x3 max pooling with stride 2 (for AlexNet). +fn max_pool2d_alex(x: Tensor<4>) -> Tensor<4> { + burn_core::tensor::module::max_pool2d(x, [3, 3], [2, 2], [0, 0], [1, 1], false) +} diff --git a/crates/burn-train/src/metric/vision/lpips/metric.rs b/crates/burn-train/src/metric/vision/lpips/metric.rs new file mode 100644 index 0000000..712fceb --- /dev/null +++ b/crates/burn-train/src/metric/vision/lpips/metric.rs @@ -0,0 +1,790 @@ +//! LPIPS (Learned Perceptual Image Patch Similarity) metric module. +//! +//! LPIPS measures perceptual similarity between images using deep features. +//! Supports VGG16, AlexNet, and SqueezeNet as backbone networks. +//! +//! Reference: "The Unreasonable Effectiveness of Deep Features as a Perceptual Metric" +//! + +use burn_core as burn; + +use burn::config::Config; +use burn::module::{Content, DisplaySettings, Module, ModuleDisplay}; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn_nn::conv::{Conv2d, Conv2dConfig}; +use burn_nn::loss::Reduction; + +use super::alexnet::AlexFeatureExtractor; +use super::squeezenet::SqueezeFeatureExtractor; +use super::vgg::VggFeatureExtractor; + +/// Network type for LPIPS. +#[derive(Config, Debug, Copy, PartialEq, Eq)] +pub enum LpipsNet { + /// VGG16 network (default) + Vgg, + /// AlexNet network + Alex, + /// SqueezeNet network + Squeeze, +} + +/// Configuration for [Lpips](Lpips) metric module. +/// +/// # Example +/// +/// ```ignore +/// use burn_train::metric::vision::{LpipsConfig, LpipsNet}; +/// +/// // VGG (default) +/// let lpips_vgg = LpipsConfig::new().init(&device); +/// +/// // AlexNet +/// let lpips_alex = LpipsConfig::new() +/// .with_net(LpipsNet::Alex) +/// .init(&device); +/// +/// // SqueezeNet +/// let lpips_squeeze = LpipsConfig::new() +/// .with_net(LpipsNet::Squeeze) +/// .init(&device); +/// ``` +#[derive(Config, Debug)] +pub struct LpipsConfig { + /// Network type for feature extraction. + #[config(default = "LpipsNet::Vgg")] + pub net: LpipsNet, + + /// Whether to normalize input images to [-1, 1] range. + /// Set to true if input is in [0, 1] range. + #[config(default = true)] + pub normalize: bool, +} + +impl LpipsConfig { + /// Initialize a new [Lpips](Lpips) module with pretrained weights. + /// + /// Downloads and loads official LPIPS pretrained weights from the + /// PerceptualSimilarity repository. + /// + /// # Arguments + /// + /// * `device` - Device to create the module on. + /// + /// # Returns + /// + /// A new LPIPS module with pretrained weights loaded. + /// + /// # Example + /// + /// ```ignore + /// use burn_train::metric::vision::{LpipsConfig, LpipsNet}; + /// + /// let lpips = LpipsConfig::new() + /// .with_net(LpipsNet::Vgg) + /// .init_pretrained(&device); + /// ``` + pub fn init_pretrained(&self, device: &Device) -> Lpips { + let lpips = self.init(device); + super::weights::load_pretrained_weights(lpips, self.net) + } + + /// Initialize a new [Lpips](Lpips) module with random weights. + /// + /// # Arguments + /// + /// * `device` - Device to create the module on. + /// + /// # Returns + /// + /// A new LPIPS module with random weights. Use `init_pretrained` for accurate results. + pub fn init(&self, device: &Device) -> Lpips { + match self.net { + LpipsNet::Vgg => { + // Channel sizes for VGG16: [64, 128, 256, 512, 512] + Lpips::Vgg(LpipsVgg { + extractor: VggFeatureExtractor::new(device), + lin0: Conv2dConfig::new([64, 1], [1, 1]) + .with_bias(false) + .init(device), + lin1: Conv2dConfig::new([128, 1], [1, 1]) + .with_bias(false) + .init(device), + lin2: Conv2dConfig::new([256, 1], [1, 1]) + .with_bias(false) + .init(device), + lin3: Conv2dConfig::new([512, 1], [1, 1]) + .with_bias(false) + .init(device), + lin4: Conv2dConfig::new([512, 1], [1, 1]) + .with_bias(false) + .init(device), + normalize: self.normalize, + }) + } + LpipsNet::Alex => { + // Channel sizes for AlexNet: [64, 192, 384, 256, 256] + Lpips::Alex(LpipsAlex { + extractor: AlexFeatureExtractor::new(device), + lin0: Conv2dConfig::new([64, 1], [1, 1]) + .with_bias(false) + .init(device), + lin1: Conv2dConfig::new([192, 1], [1, 1]) + .with_bias(false) + .init(device), + lin2: Conv2dConfig::new([384, 1], [1, 1]) + .with_bias(false) + .init(device), + lin3: Conv2dConfig::new([256, 1], [1, 1]) + .with_bias(false) + .init(device), + lin4: Conv2dConfig::new([256, 1], [1, 1]) + .with_bias(false) + .init(device), + normalize: self.normalize, + }) + } + LpipsNet::Squeeze => { + // Channel sizes for SqueezeNet: [64, 128, 256, 384, 384, 512, 512] + Lpips::Squeeze(LpipsSqueeze { + extractor: SqueezeFeatureExtractor::new(device), + lin0: Conv2dConfig::new([64, 1], [1, 1]) + .with_bias(false) + .init(device), + lin1: Conv2dConfig::new([128, 1], [1, 1]) + .with_bias(false) + .init(device), + lin2: Conv2dConfig::new([256, 1], [1, 1]) + .with_bias(false) + .init(device), + lin3: Conv2dConfig::new([384, 1], [1, 1]) + .with_bias(false) + .init(device), + lin4: Conv2dConfig::new([384, 1], [1, 1]) + .with_bias(false) + .init(device), + lin5: Conv2dConfig::new([512, 1], [1, 1]) + .with_bias(false) + .init(device), + lin6: Conv2dConfig::new([512, 1], [1, 1]) + .with_bias(false) + .init(device), + normalize: self.normalize, + }) + } + } + } +} + +/// LPIPS (Learned Perceptual Image Patch Similarity) metric module. +/// +/// Computes perceptual distance between two images using deep features. +/// Supports VGG16, AlexNet, and SqueezeNet as backbone networks. +/// +/// # Example +/// +/// ```ignore +/// use burn_train::metric::vision::{LpipsConfig, LpipsNet, Reduction}; +/// +/// let device = Default::default(); +/// let lpips = LpipsConfig::new().init(&device); +/// +/// let img1: Tensor<4> = /* [batch, 3, H, W] */; +/// let img2: Tensor<4> = /* [batch, 3, H, W] */; +/// +/// // Compute LPIPS distance +/// let distance = lpips.forward(img1, img2, Reduction::Mean); +/// ``` +#[derive(Module, Debug)] +#[allow(clippy::large_enum_variant)] +#[module(custom_display)] +pub enum Lpips { + /// VGG16 backbone (5 feature layers) + Vgg(LpipsVgg), + /// AlexNet backbone (5 feature layers) + Alex(LpipsAlex), + /// SqueezeNet backbone (7 feature layers) + Squeeze(LpipsSqueeze), +} + +/// LPIPS with VGG16 backbone. +#[derive(Module, Debug)] +pub struct LpipsVgg { + /// VGG feature extractor + pub(crate) extractor: VggFeatureExtractor, + /// Linear layers for each feature level + pub(crate) lin0: Conv2d, + pub(crate) lin1: Conv2d, + pub(crate) lin2: Conv2d, + pub(crate) lin3: Conv2d, + pub(crate) lin4: Conv2d, + /// Whether to normalize input + pub(crate) normalize: bool, +} + +/// LPIPS with AlexNet backbone. +#[derive(Module, Debug)] +pub struct LpipsAlex { + /// AlexNet feature extractor + pub(crate) extractor: AlexFeatureExtractor, + /// Linear layers for each feature level + pub(crate) lin0: Conv2d, + pub(crate) lin1: Conv2d, + pub(crate) lin2: Conv2d, + pub(crate) lin3: Conv2d, + pub(crate) lin4: Conv2d, + /// Whether to normalize input + pub(crate) normalize: bool, +} + +/// LPIPS with SqueezeNet backbone. +#[derive(Module, Debug)] +pub struct LpipsSqueeze { + /// SqueezeNet feature extractor + pub(crate) extractor: SqueezeFeatureExtractor, + /// Linear layers for each feature level + pub(crate) lin0: Conv2d, + pub(crate) lin1: Conv2d, + pub(crate) lin2: Conv2d, + pub(crate) lin3: Conv2d, + pub(crate) lin4: Conv2d, + pub(crate) lin5: Conv2d, + pub(crate) lin6: Conv2d, + /// Whether to normalize input + pub(crate) normalize: bool, +} + +impl LpipsVgg { + /// Compute LPIPS distance without reduction using VGG backbone. + pub fn forward_no_reduction(&self, input: Tensor<4>, target: Tensor<4>) -> Tensor<1> { + // Preprocess inputs + let (input, target) = preprocess_inputs(input, target, self.normalize); + + // Extract features from both images + let feats0 = self.extractor.forward(input); + let feats1 = self.extractor.forward(target); + + // Compute distance for each layer using stack + sum + let layer_distances: Vec> = vec![ + compute_layer_distance(&feats0[0], &feats1[0], &self.lin0).unsqueeze_dim(1), + compute_layer_distance(&feats0[1], &feats1[1], &self.lin1).unsqueeze_dim(1), + compute_layer_distance(&feats0[2], &feats1[2], &self.lin2).unsqueeze_dim(1), + compute_layer_distance(&feats0[3], &feats1[3], &self.lin3).unsqueeze_dim(1), + compute_layer_distance(&feats0[4], &feats1[4], &self.lin4).unsqueeze_dim(1), + ]; + + Tensor::cat(layer_distances, 1) + .sum_dim(1) + .squeeze_dim::<1>(1) + } +} + +impl LpipsAlex { + /// Compute LPIPS distance without reduction using AlexNet backbone. + pub fn forward_no_reduction(&self, input: Tensor<4>, target: Tensor<4>) -> Tensor<1> { + // Preprocess inputs + let (input, target) = preprocess_inputs(input, target, self.normalize); + + // Extract features from both images + let feats0 = self.extractor.forward(input); + let feats1 = self.extractor.forward(target); + + // Compute distance for each layer using stack + sum + let layer_distances: Vec> = vec![ + compute_layer_distance(&feats0[0], &feats1[0], &self.lin0).unsqueeze_dim(1), + compute_layer_distance(&feats0[1], &feats1[1], &self.lin1).unsqueeze_dim(1), + compute_layer_distance(&feats0[2], &feats1[2], &self.lin2).unsqueeze_dim(1), + compute_layer_distance(&feats0[3], &feats1[3], &self.lin3).unsqueeze_dim(1), + compute_layer_distance(&feats0[4], &feats1[4], &self.lin4).unsqueeze_dim(1), + ]; + + Tensor::cat(layer_distances, 1) + .sum_dim(1) + .squeeze_dim::<1>(1) + } +} + +impl LpipsSqueeze { + /// Compute LPIPS distance without reduction using SqueezeNet backbone. + pub fn forward_no_reduction(&self, input: Tensor<4>, target: Tensor<4>) -> Tensor<1> { + // Preprocess inputs + let (input, target) = preprocess_inputs(input, target, self.normalize); + + // Extract features from both images + let feats0 = self.extractor.forward(input); + let feats1 = self.extractor.forward(target); + + // Compute distance for each layer using stack + sum (7 layers for SqueezeNet) + let layer_distances: Vec> = vec![ + compute_layer_distance(&feats0[0], &feats1[0], &self.lin0).unsqueeze_dim(1), + compute_layer_distance(&feats0[1], &feats1[1], &self.lin1).unsqueeze_dim(1), + compute_layer_distance(&feats0[2], &feats1[2], &self.lin2).unsqueeze_dim(1), + compute_layer_distance(&feats0[3], &feats1[3], &self.lin3).unsqueeze_dim(1), + compute_layer_distance(&feats0[4], &feats1[4], &self.lin4).unsqueeze_dim(1), + compute_layer_distance(&feats0[5], &feats1[5], &self.lin5).unsqueeze_dim(1), + compute_layer_distance(&feats0[6], &feats1[6], &self.lin6).unsqueeze_dim(1), + ]; + + Tensor::cat(layer_distances, 1) + .sum_dim(1) + .squeeze_dim::<1>(1) + } +} + +impl ModuleDisplay for Lpips { + fn custom_settings(&self) -> Option { + DisplaySettings::new() + .with_new_line_after_attribute(false) + .optional() + } + + fn custom_content(&self, content: Content) -> Option { + let (net_name, normalize) = match self { + Lpips::Vgg(inner) => ("Vgg", inner.normalize), + Lpips::Alex(inner) => ("Alex", inner.normalize), + Lpips::Squeeze(inner) => ("Squeeze", inner.normalize), + }; + content + .add("net", &net_name.to_string()) + .add("normalize", &normalize.to_string()) + .optional() + } +} + +impl Lpips { + /// Compute LPIPS distance with reduction. + /// + /// # Arguments + /// + /// * `input` - First image tensor of shape `[batch, 3, H, W]` + /// * `target` - Second image tensor of shape `[batch, 3, H, W]` + /// * `reduction` - How to reduce the output (Mean, Sum, or Auto) + /// + /// # Returns + /// + /// Scalar tensor of shape `[1]`. + /// + /// # Shapes + /// + /// - input: `[batch, 3, H, W]` + /// - target: `[batch, 3, H, W]` + /// - output: `[1]` + pub fn forward(&self, input: Tensor<4>, target: Tensor<4>, reduction: Reduction) -> Tensor<1> { + let distance = self.forward_no_reduction(input, target); + + match reduction { + Reduction::Mean | Reduction::Auto | Reduction::BatchMean => distance.mean(), + Reduction::Sum => distance.sum(), + } + } + + /// Compute LPIPS distance without reduction. + /// + /// # Arguments + /// + /// * `input` - First image tensor of shape `[batch, 3, H, W]` + /// * `target` - Second image tensor of shape `[batch, 3, H, W]` + /// + /// # Returns + /// + /// Per-sample distance tensor of shape `[batch]`. + /// + /// # Shapes + /// + /// - input: `[batch, 3, H, W]` + /// - target: `[batch, 3, H, W]` + /// - output: `[batch]` + pub fn forward_no_reduction(&self, input: Tensor<4>, target: Tensor<4>) -> Tensor<1> { + match self { + Lpips::Vgg(inner) => inner.forward_no_reduction(input, target), + Lpips::Alex(inner) => inner.forward_no_reduction(input, target), + Lpips::Squeeze(inner) => inner.forward_no_reduction(input, target), + } + } +} + +// ============================================================================= +// Helper Functions +// ============================================================================= + +/// Normalize tensor to unit norm along channel dimension. +fn normalize_tensor(x: Tensor<4>) -> Tensor<4> { + let norm = x.clone().mul(x.clone()).sum_dim(1).sqrt().clamp_min(1e-10); + x.div(norm) +} + +/// Apply ImageNet normalization used by PyTorch lpips. +/// shift = [-.030, -.088, -.188], scale = [.458, .448, .450] +/// output = (input - shift) / scale +fn scaling_layer(x: Tensor<4>) -> Tensor<4> { + let device = x.device(); + let [batch, _, h, w] = x.dims(); + + // Create shift and scale tensors [1, 3, 1, 1] and broadcast + let shift = Tensor::<2>::from_floats([[-0.030], [-0.088], [-0.188]], &device) + .reshape([1, 3, 1, 1]) + .expand([batch, 3, h, w]); + let scale = Tensor::<2>::from_floats([[0.458], [0.448], [0.450]], &device) + .reshape([1, 3, 1, 1]) + .expand([batch, 3, h, w]); + + x.sub(shift).div(scale) +} + +/// Compute normalized L2 distance for a single layer. +fn compute_layer_distance(feat0: &Tensor<4>, feat1: &Tensor<4>, lin: &Conv2d) -> Tensor<1> { + // Normalize features (unit norm along channel dimension) + let feat0_norm = normalize_tensor(feat0.clone()); + let feat1_norm = normalize_tensor(feat1.clone()); + + // Compute squared difference + let diff = feat0_norm.sub(feat1_norm); + let diff_sq = diff.clone().mul(diff); + + // Apply linear layer (learned weights) + // Shape: [batch, C, H, W] -> [batch, 1, H, W] + let weighted = lin.forward(diff_sq); + + // Spatial average: compute mean over C, H, W dimensions + // Shape: [batch, 1, H, W] -> [batch] + let [batch, c, h, w] = weighted.dims(); + + // Reshape to [batch, c*h*w] then take mean over last dimension + weighted + .reshape([batch, c * h * w]) + .mean_dim(1) + .squeeze_dim::<1>(1) +} + +/// Preprocess input images for LPIPS computation. +fn preprocess_inputs( + input: Tensor<4>, + target: Tensor<4>, + normalize: bool, +) -> (Tensor<4>, Tensor<4>) { + // Normalize to [-1, 1] if needed + let (input, target) = if normalize { + ( + input.mul_scalar(2.0).sub_scalar(1.0), + target.mul_scalar(2.0).sub_scalar(1.0), + ) + } else { + (input, target) + }; + + // Apply ImageNet normalization (same as PyTorch lpips scaling_layer) + (scaling_layer(input), scaling_layer(target)) +} + +// ============================================================================= +// Tests +// ============================================================================= + +#[cfg(test)] +mod tests { + use super::*; + use burn_core::tensor::{TensorData, Tolerance}; + + type FT = f32; + + // ========================================================================= + // Basic Functionality Tests + // ========================================================================= + + /// Identical images should have LPIPS distance of 0. + #[test] + fn test_lpips_identical_images_zero_distance() { + let device = Default::default(); + let image = Tensor::<4>::ones([1, 3, 32, 32], &device); + + let lpips = LpipsConfig::new().init(&device); + let distance = lpips.forward(image.clone(), image, Reduction::Mean); + + // Identical images → distance = 0 + let expected = TensorData::from([0.0]); + distance + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + /// Different images should have LPIPS distance != 0. + /// Note: With random weights, distance can be negative, so we only check != 0. + /// Non-negativity is tested with pretrained weights. + #[test] + fn test_lpips_different_images_nonzero_distance() { + let device = Default::default(); + + let image1 = Tensor::<4>::zeros([1, 3, 32, 32], &device); + let image2 = Tensor::<4>::ones([1, 3, 32, 32], &device); + + let lpips = LpipsConfig::new().init(&device); + let distance = lpips.forward(image1, image2, Reduction::Mean); + + let distance_value = distance.into_data().to_vec::().unwrap()[0]; + assert!( + distance_value.abs() > 1e-6, + "LPIPS should be != 0 for different images" + ); + } + + /// Test symmetry: LPIPS(a, b) == LPIPS(b, a). + #[test] + fn test_lpips_symmetry() { + let device = Default::default(); + + let image1 = Tensor::<4>::zeros([1, 3, 32, 32], &device); + let image2 = Tensor::<4>::ones([1, 3, 32, 32], &device); + + let lpips = LpipsConfig::new().init(&device); + let distance_forward = lpips.forward(image1.clone(), image2.clone(), Reduction::Mean); + let distance_reverse = lpips.forward(image2, image1, Reduction::Mean); + + distance_forward + .into_data() + .assert_approx_eq::(&distance_reverse.into_data(), Tolerance::default()); + } + + // ========================================================================= + // Reduction Tests + // ========================================================================= + + #[test] + fn test_lpips_forward_mean_reduction() { + let device = Default::default(); + + let image1 = Tensor::<4>::zeros([2, 3, 32, 32], &device); + let image2 = Tensor::<4>::ones([2, 3, 32, 32], &device); + + let lpips = LpipsConfig::new().init(&device); + let distance = lpips.forward(image1, image2, Reduction::Mean); + + assert_eq!(distance.dims(), [1]); + } + + #[test] + fn test_lpips_forward_no_reduction() { + let device = Default::default(); + + let batch_size = 4; + let image1 = Tensor::<4>::zeros([batch_size, 3, 32, 32], &device); + let image2 = Tensor::<4>::ones([batch_size, 3, 32, 32], &device); + + let lpips = LpipsConfig::new().init(&device); + let distance = lpips.forward_no_reduction(image1, image2); + + assert_eq!(distance.dims(), [batch_size]); + } + + // ========================================================================= + // AlexNet Tests + // ========================================================================= + + /// Test AlexNet LPIPS with identical images. + #[test] + fn test_lpips_alex_identical_images_zero_distance() { + let device = Default::default(); + let image = Tensor::<4>::ones([1, 3, 64, 64], &device); + + let lpips = LpipsConfig::new().with_net(LpipsNet::Alex).init(&device); + let distance = lpips.forward(image.clone(), image, Reduction::Mean); + + let expected = TensorData::from([0.0]); + distance + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + /// Test AlexNet LPIPS with different images produces non-zero distance. + #[test] + fn test_lpips_alex_different_images_nonzero_distance() { + let device = Default::default(); + + let image1 = Tensor::<4>::zeros([1, 3, 64, 64], &device); + let image2 = Tensor::<4>::ones([1, 3, 64, 64], &device); + + let lpips = LpipsConfig::new().with_net(LpipsNet::Alex).init(&device); + let distance = lpips.forward(image1, image2, Reduction::Mean); + + let distance_value = distance.into_data().to_vec::().unwrap()[0]; + // Note: With random weights, non-negativity is not guaranteed. + // We only check that different images produce a non-zero distance. + assert!( + distance_value.abs() > 1e-6, + "LPIPS (Alex) should be != 0 for different images" + ); + } + + // ========================================================================= + // SqueezeNet Tests + // ========================================================================= + + /// Test SqueezeNet LPIPS with identical images. + #[test] + fn test_lpips_squeeze_identical_images_zero_distance() { + let device = Default::default(); + let image = Tensor::<4>::ones([1, 3, 64, 64], &device); + + let lpips = LpipsConfig::new().with_net(LpipsNet::Squeeze).init(&device); + let distance = lpips.forward(image.clone(), image, Reduction::Mean); + + let expected = TensorData::from([0.0]); + distance + .into_data() + .assert_approx_eq::(&expected, Tolerance::default()); + } + + /// Test SqueezeNet LPIPS with different images produces non-zero distance. + #[test] + fn test_lpips_squeeze_different_images_nonzero_distance() { + let device = Default::default(); + + let image1 = Tensor::<4>::zeros([1, 3, 64, 64], &device); + let image2 = Tensor::<4>::ones([1, 3, 64, 64], &device); + + let lpips = LpipsConfig::new().with_net(LpipsNet::Squeeze).init(&device); + let distance = lpips.forward(image1, image2, Reduction::Mean); + + let distance_value = distance.into_data().to_vec::().unwrap()[0]; + // Note: With random weights, non-negativity is not guaranteed. + // We only check that different images produce a non-zero distance. + assert!( + distance_value.abs() > 1e-6, + "LPIPS (Squeeze) should be != 0 for different images" + ); + } + + // ========================================================================= + // Display Tests + // ========================================================================= + + #[test] + fn display_vgg() { + let device = Default::default(); + let lpips = LpipsConfig::new().init(&device); + + let display_str = format!("{lpips}"); + assert!(display_str.contains("Lpips")); + assert!(display_str.contains("Vgg")); + } + + #[test] + fn display_alex() { + let device = Default::default(); + let lpips = LpipsConfig::new().with_net(LpipsNet::Alex).init(&device); + + let display_str = format!("{lpips}"); + assert!(display_str.contains("Lpips")); + assert!(display_str.contains("Alex")); + } + + #[test] + fn display_squeeze() { + let device = Default::default(); + let lpips = LpipsConfig::new().with_net(LpipsNet::Squeeze).init(&device); + + let display_str = format!("{lpips}"); + assert!(display_str.contains("Lpips")); + assert!(display_str.contains("Squeeze")); + } + + // ========================================================================= + // Pretrained Weights Tests (requires network) + // ========================================================================= + + /// Test VGG pretrained weights download and loading. + #[test] + #[ignore = "downloads pre-trained weights"] + fn test_lpips_pretrained_vgg() { + let device = Default::default(); + + // This will download ~60MB of weights + let lpips = LpipsConfig::new() + .with_net(LpipsNet::Vgg) + .init_pretrained(&device); + + // Test with identical images - should be 0 + let image = Tensor::<4>::ones([1, 3, 64, 64], &device); + let distance = lpips.forward(image.clone(), image, Reduction::Mean); + let distance_value = distance.into_data().to_vec::().unwrap()[0]; + assert!( + distance_value.abs() < 1e-5, + "Pretrained LPIPS (VGG) should be ~0 for identical images, got {}", + distance_value + ); + + // Test with different images - should be positive + let image1 = Tensor::<4>::zeros([1, 3, 64, 64], &device); + let image2 = Tensor::<4>::ones([1, 3, 64, 64], &device); + let distance = lpips.forward(image1, image2, Reduction::Mean); + let distance_value = distance.into_data().to_vec::().unwrap()[0]; + assert!( + distance_value > 0.0, + "Pretrained LPIPS (VGG) should be > 0 for different images, got {}", + distance_value + ); + } + + /// Test AlexNet pretrained weights download and loading. + #[test] + #[ignore = "downloads pre-trained weights"] + fn test_lpips_pretrained_alex() { + let device = Default::default(); + + let lpips = LpipsConfig::new() + .with_net(LpipsNet::Alex) + .init_pretrained(&device); + + // Test with identical images + let image = Tensor::<4>::ones([1, 3, 64, 64], &device); + let distance = lpips.forward(image.clone(), image, Reduction::Mean); + let distance_value = distance.into_data().to_vec::().unwrap()[0]; + assert!( + distance_value.abs() < 1e-5, + "Pretrained LPIPS (Alex) should be ~0 for identical images, got {}", + distance_value + ); + + // Test with different images + let image1 = Tensor::<4>::zeros([1, 3, 64, 64], &device); + let image2 = Tensor::<4>::ones([1, 3, 64, 64], &device); + let distance = lpips.forward(image1, image2, Reduction::Mean); + let distance_value = distance.into_data().to_vec::().unwrap()[0]; + assert!( + distance_value > 0.0, + "Pretrained LPIPS (Alex) should be > 0 for different images" + ); + } + + /// Test SqueezeNet pretrained weights download and loading. + #[test] + #[ignore = "downloads pre-trained weights"] + fn test_lpips_pretrained_squeeze() { + let device = Default::default(); + + let lpips = LpipsConfig::new() + .with_net(LpipsNet::Squeeze) + .init_pretrained(&device); + + // Test with identical images + let image = Tensor::<4>::ones([1, 3, 64, 64], &device); + let distance = lpips.forward(image.clone(), image, Reduction::Mean); + let distance_value = distance.into_data().to_vec::().unwrap()[0]; + assert!( + distance_value.abs() < 1e-5, + "Pretrained LPIPS (Squeeze) should be ~0 for identical images, got {}", + distance_value + ); + + // Test with different images + let image1 = Tensor::<4>::zeros([1, 3, 64, 64], &device); + let image2 = Tensor::<4>::ones([1, 3, 64, 64], &device); + let distance = lpips.forward(image1, image2, Reduction::Mean); + let distance_value = distance.into_data().to_vec::().unwrap()[0]; + assert!( + distance_value > 0.0, + "Pretrained LPIPS (Squeeze) should be > 0 for different images, got {}", + distance_value + ); + } +} diff --git a/crates/burn-train/src/metric/vision/lpips/mod.rs b/crates/burn-train/src/metric/vision/lpips/mod.rs new file mode 100644 index 0000000..9e5e6a5 --- /dev/null +++ b/crates/burn-train/src/metric/vision/lpips/mod.rs @@ -0,0 +1,21 @@ +//! LPIPS (Learned Perceptual Image Patch Similarity) metric module. +//! +//! LPIPS measures perceptual similarity between images using deep features. +//! Supports VGG16, AlexNet, and SqueezeNet as backbone networks. +//! +//! Reference: "The Unreasonable Effectiveness of Deep Features as a Perceptual Metric" +//! + +mod alexnet; +mod metric; +mod squeezenet; +mod vgg; +mod weights; + +pub use metric::{Lpips, LpipsAlex, LpipsConfig, LpipsNet, LpipsSqueeze, LpipsVgg}; +pub use weights::{get_backbone_weights_url, get_lpips_weights_url, load_pretrained_weights}; + +// Re-export feature extractors for advanced use cases +pub use alexnet::AlexFeatureExtractor; +pub use squeezenet::{FireModule, SqueezeFeatureExtractor}; +pub use vgg::VggFeatureExtractor; diff --git a/crates/burn-train/src/metric/vision/lpips/squeezenet.rs b/crates/burn-train/src/metric/vision/lpips/squeezenet.rs new file mode 100644 index 0000000..7dde042 --- /dev/null +++ b/crates/burn-train/src/metric/vision/lpips/squeezenet.rs @@ -0,0 +1,157 @@ +//! SqueezeNet feature extractor for LPIPS. + +use burn_core as burn; + +use burn::module::Module; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn::tensor::activation::relu; +use burn_nn::PaddingConfig2d; +use burn_nn::conv::{Conv2d, Conv2dConfig}; + +/// Fire module for SqueezeNet. +/// +/// A fire module consists of: +/// - Squeeze layer: 1x1 conv to reduce channels +/// - Expand layers: parallel 1x1 and 3x3 convs, concatenated +#[derive(Module, Debug)] +pub struct FireModule { + /// Squeeze layer: 1x1 conv + squeeze: Conv2d, + /// Expand 1x1 conv + expand1x1: Conv2d, + /// Expand 3x3 conv + expand3x3: Conv2d, +} + +impl FireModule { + /// Create a new Fire module. + pub fn new( + in_channels: usize, + squeeze_channels: usize, + expand1x1_channels: usize, + expand3x3_channels: usize, + device: &Device, + ) -> Self { + Self { + squeeze: Conv2dConfig::new([in_channels, squeeze_channels], [1, 1]) + .with_bias(true) + .init(device), + expand1x1: Conv2dConfig::new([squeeze_channels, expand1x1_channels], [1, 1]) + .with_bias(true) + .init(device), + expand3x3: Conv2dConfig::new([squeeze_channels, expand3x3_channels], [3, 3]) + .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1)) + .with_bias(true) + .init(device), + } + } + + /// Forward pass through fire module. + pub fn forward(&self, x: Tensor<4>) -> Tensor<4> { + let squeezed = relu(self.squeeze.forward(x)); + let e1 = relu(self.expand1x1.forward(squeezed.clone())); + let e3 = relu(self.expand3x3.forward(squeezed)); + // Concatenate along channel dimension + Tensor::cat(vec![e1, e3], 1) + } +} + +/// SqueezeNet 1.1 feature extractor for LPIPS. +/// +/// Extracts features from 7 layers: +/// - After conv1+relu: 64 channels +/// - After fire1+fire2: 128 channels +/// - After fire3+fire4: 256 channels +/// - After fire5: 384 channels +/// - After fire6: 384 channels +/// - After fire7: 512 channels +/// - After fire8: 512 channels +#[derive(Module, Debug)] +pub struct SqueezeFeatureExtractor { + /// Conv1: 3 -> 64, kernel 3x3, stride 2 + conv1: Conv2d, + /// Fire1: 64 -> 128 (squeeze=16, expand=64+64) + fire1: FireModule, + /// Fire2: 128 -> 128 (squeeze=16, expand=64+64) + fire2: FireModule, + /// Fire3: 128 -> 256 (squeeze=32, expand=128+128) + fire3: FireModule, + /// Fire4: 256 -> 256 (squeeze=32, expand=128+128) + fire4: FireModule, + /// Fire5: 256 -> 384 (squeeze=48, expand=192+192) + fire5: FireModule, + /// Fire6: 384 -> 384 (squeeze=48, expand=192+192) + fire6: FireModule, + /// Fire7: 384 -> 512 (squeeze=64, expand=256+256) + fire7: FireModule, + /// Fire8: 512 -> 512 (squeeze=64, expand=256+256) + fire8: FireModule, +} + +impl SqueezeFeatureExtractor { + /// Create a new SqueezeNet feature extractor. + pub fn new(device: &Device) -> Self { + Self { + // Conv1: 3 -> 64, 3x3, stride 2 + conv1: Conv2dConfig::new([3, 64], [3, 3]) + .with_stride([2, 2]) + .with_bias(true) + .init(device), + // Fire modules (SqueezeNet 1.1 configuration) + fire1: FireModule::new(64, 16, 64, 64, device), // -> 128 + fire2: FireModule::new(128, 16, 64, 64, device), // -> 128 + fire3: FireModule::new(128, 32, 128, 128, device), // -> 256 + fire4: FireModule::new(256, 32, 128, 128, device), // -> 256 + fire5: FireModule::new(256, 48, 192, 192, device), // -> 384 + fire6: FireModule::new(384, 48, 192, 192, device), // -> 384 + fire7: FireModule::new(384, 64, 256, 256, device), // -> 512 + fire8: FireModule::new(512, 64, 256, 256, device), // -> 512 + } + } + + /// Extract features from 7 SqueezeNet layers. + pub fn forward(&self, x: Tensor<4>) -> Vec> { + let mut features = Vec::with_capacity(7); + + // Slice 1: Conv1 + ReLU (64 channels) + let x = relu(self.conv1.forward(x)); + features.push(x.clone()); + + // Slice 2: MaxPool + Fire1 + Fire2 (128 channels) + let x = max_pool2d_squeeze(x); + let x = self.fire1.forward(x); + let x = self.fire2.forward(x); + features.push(x.clone()); + + // Slice 3: MaxPool + Fire3 + Fire4 (256 channels) + let x = max_pool2d_squeeze(x); + let x = self.fire3.forward(x); + let x = self.fire4.forward(x); + features.push(x.clone()); + + // Slice 4: MaxPool + Fire5 (384 channels) + let x = max_pool2d_squeeze(x); + let x = self.fire5.forward(x); + features.push(x.clone()); + + // Slice 5: Fire6 (384 channels) + let x = self.fire6.forward(x); + features.push(x.clone()); + + // Slice 6: Fire7 (512 channels) + let x = self.fire7.forward(x); + features.push(x.clone()); + + // Slice 7: Fire8 (512 channels) + let x = self.fire8.forward(x); + features.push(x); + + features + } +} + +/// 3x3 max pooling with stride 2, ceil mode (for SqueezeNet). +fn max_pool2d_squeeze(x: Tensor<4>) -> Tensor<4> { + burn_core::tensor::module::max_pool2d(x, [3, 3], [2, 2], [0, 0], [1, 1], true) +} diff --git a/crates/burn-train/src/metric/vision/lpips/vgg.rs b/crates/burn-train/src/metric/vision/lpips/vgg.rs new file mode 100644 index 0000000..0d3d503 --- /dev/null +++ b/crates/burn-train/src/metric/vision/lpips/vgg.rs @@ -0,0 +1,116 @@ +//! VGG16 feature extractor for LPIPS. + +use burn_core as burn; + +use burn::module::Module; +use burn::tensor::Device; +use burn::tensor::Tensor; +use burn::tensor::activation::relu; +use burn_nn::PaddingConfig2d; +use burn_nn::conv::{Conv2d, Conv2dConfig}; + +/// VGG16 feature extractor for LPIPS. +/// +/// Extracts features from 5 layers: +/// - conv1_2: 64 channels +/// - conv2_2: 128 channels +/// - conv3_3: 256 channels +/// - conv4_3: 512 channels +/// - conv5_3: 512 channels +#[derive(Module, Debug)] +pub struct VggFeatureExtractor { + // Block 1 + conv1_1: Conv2d, + conv1_2: Conv2d, + // Block 2 + conv2_1: Conv2d, + conv2_2: Conv2d, + // Block 3 + conv3_1: Conv2d, + conv3_2: Conv2d, + conv3_3: Conv2d, + // Block 4 + conv4_1: Conv2d, + conv4_2: Conv2d, + conv4_3: Conv2d, + // Block 5 + conv5_1: Conv2d, + conv5_2: Conv2d, + conv5_3: Conv2d, +} + +impl VggFeatureExtractor { + /// Create a new VGG16 feature extractor. + pub fn new(device: &Device) -> Self { + let conv_config = |in_ch, out_ch| { + Conv2dConfig::new([in_ch, out_ch], [3, 3]) + .with_padding(PaddingConfig2d::Same) + .with_bias(true) + }; + + Self { + // Block 1: 3 -> 64 + conv1_1: conv_config(3, 64).init(device), + conv1_2: conv_config(64, 64).init(device), + // Block 2: 64 -> 128 + conv2_1: conv_config(64, 128).init(device), + conv2_2: conv_config(128, 128).init(device), + // Block 3: 128 -> 256 + conv3_1: conv_config(128, 256).init(device), + conv3_2: conv_config(256, 256).init(device), + conv3_3: conv_config(256, 256).init(device), + // Block 4: 256 -> 512 + conv4_1: conv_config(256, 512).init(device), + conv4_2: conv_config(512, 512).init(device), + conv4_3: conv_config(512, 512).init(device), + // Block 5: 512 -> 512 + conv5_1: conv_config(512, 512).init(device), + conv5_2: conv_config(512, 512).init(device), + conv5_3: conv_config(512, 512).init(device), + } + } + + /// Extract features from 5 VGG layers. + pub fn forward(&self, x: Tensor<4>) -> Vec> { + let mut features = Vec::with_capacity(5); + + // Block 1 + let x = relu(self.conv1_1.forward(x)); + let x = relu(self.conv1_2.forward(x)); + features.push(x.clone()); + let x = max_pool2d(x); + + // Block 2 + let x = relu(self.conv2_1.forward(x)); + let x = relu(self.conv2_2.forward(x)); + features.push(x.clone()); + let x = max_pool2d(x); + + // Block 3 + let x = relu(self.conv3_1.forward(x)); + let x = relu(self.conv3_2.forward(x)); + let x = relu(self.conv3_3.forward(x)); + features.push(x.clone()); + let x = max_pool2d(x); + + // Block 4 + let x = relu(self.conv4_1.forward(x)); + let x = relu(self.conv4_2.forward(x)); + let x = relu(self.conv4_3.forward(x)); + features.push(x.clone()); + let x = max_pool2d(x); + + // Block 5 + let x = relu(self.conv5_1.forward(x)); + let x = relu(self.conv5_2.forward(x)); + let x = relu(self.conv5_3.forward(x)); + features.push(x); + + features + } +} + +/// 2x2 max pooling with stride 2. +fn max_pool2d(x: Tensor<4>) -> Tensor<4> { + burn_core::tensor::module::max_pool2d(x, [2, 2], [2, 2], [0, 0], [1, 1], false) +} diff --git a/crates/burn-train/src/metric/vision/lpips/weights.rs b/crates/burn-train/src/metric/vision/lpips/weights.rs new file mode 100644 index 0000000..6dfc1dc --- /dev/null +++ b/crates/burn-train/src/metric/vision/lpips/weights.rs @@ -0,0 +1,225 @@ +//! Pretrained weights loading for LPIPS. + +use burn_std::network::downloader::download_file_as_bytes; +use burn_store::{ModuleSnapshot, PytorchStore}; +use std::fs::{File, create_dir_all}; +use std::io::Write; +use std::path::PathBuf; + +use super::metric::{Lpips, LpipsNet}; + +/// URLs for pretrained LPIPS linear layer weights from the official repository. +/// Reference: https://github.com/richzhang/PerceptualSimilarity +const LPIPS_VGG_URL: &str = + "https://github.com/richzhang/PerceptualSimilarity/raw/master/lpips/weights/v0.1/vgg.pth"; +const LPIPS_ALEX_URL: &str = + "https://github.com/richzhang/PerceptualSimilarity/raw/master/lpips/weights/v0.1/alex.pth"; +const LPIPS_SQUEEZE_URL: &str = + "https://github.com/richzhang/PerceptualSimilarity/raw/master/lpips/weights/v0.1/squeeze.pth"; + +/// URLs for ImageNet pretrained backbone weights from PyTorch. +const VGG16_IMAGENET_URL: &str = "https://download.pytorch.org/models/vgg16-397923af.pth"; +const ALEXNET_IMAGENET_URL: &str = "https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth"; +const SQUEEZENET_IMAGENET_URL: &str = + "https://download.pytorch.org/models/squeezenet1_1-f364aa15.pth"; + +/// Get the download URL for LPIPS linear layer weights. +pub fn get_lpips_weights_url(net: LpipsNet) -> &'static str { + match net { + LpipsNet::Vgg => LPIPS_VGG_URL, + LpipsNet::Alex => LPIPS_ALEX_URL, + LpipsNet::Squeeze => LPIPS_SQUEEZE_URL, + } +} + +/// Get the download URL for backbone ImageNet weights. +pub fn get_backbone_weights_url(net: LpipsNet) -> &'static str { + match net { + LpipsNet::Vgg => VGG16_IMAGENET_URL, + LpipsNet::Alex => ALEXNET_IMAGENET_URL, + LpipsNet::Squeeze => SQUEEZENET_IMAGENET_URL, + } +} + +/// Get the cache directory for LPIPS weights. +fn get_cache_dir() -> PathBuf { + let cache_dir = dirs::cache_dir() + .expect("Could not get cache directory") + .join("burn-dataset") + .join("lpips"); + + if !cache_dir.exists() { + create_dir_all(&cache_dir).expect("Failed to create cache directory"); + } + + cache_dir +} + +/// Download file if not cached and return the cache path. +fn download_if_needed(url: &str, cache_path: &PathBuf, message: &str) { + if !cache_path.exists() { + let bytes = download_file_as_bytes(url, message); + let mut file = File::create(cache_path).expect("Failed to create cache file"); + file.write_all(&bytes).expect("Failed to write weights"); + } +} + +/// Download and load pretrained weights into an LPIPS module. +/// +/// This loads both: +/// 1. ImageNet pretrained backbone weights (VGG16/AlexNet/SqueezeNet) +/// 2. LPIPS trained linear layer weights +/// +/// Weights are cached in the user's cache directory to avoid re-downloading. +/// +/// # Arguments +/// +/// * `lpips` - The LPIPS module to load weights into. +/// * `net` - The network type (determines which weights to download). +/// +/// # Returns +/// +/// The LPIPS module with loaded pretrained weights. +pub fn load_pretrained_weights(mut lpips: Lpips, net: LpipsNet) -> Lpips { + let cache_dir = get_cache_dir(); + + // Step 1: Load backbone ImageNet weights + let backbone_url = get_backbone_weights_url(net); + let backbone_cache_path = cache_dir.join(format!("{:?}_backbone.pth", net).to_lowercase()); + let backbone_message = match net { + LpipsNet::Vgg => "Downloading VGG16 ImageNet weights...", + LpipsNet::Alex => "Downloading AlexNet ImageNet weights...", + LpipsNet::Squeeze => "Downloading SqueezeNet ImageNet weights...", + }; + download_if_needed(backbone_url, &backbone_cache_path, backbone_message); + + // Step 2: Load LPIPS linear layer weights + let lpips_url = get_lpips_weights_url(net); + let lpips_cache_path = cache_dir.join(format!("{:?}_lpips.pth", net).to_lowercase()); + let lpips_message = match net { + LpipsNet::Vgg => "Downloading LPIPS VGG weights...", + LpipsNet::Alex => "Downloading LPIPS AlexNet weights...", + LpipsNet::Squeeze => "Downloading LPIPS SqueezeNet weights...", + }; + download_if_needed(lpips_url, &lpips_cache_path, lpips_message); + + // Load backbone weights first + lpips = load_backbone_weights(lpips, &backbone_cache_path); + + // Then load LPIPS linear layer weights + lpips = load_lpips_weights(lpips, &lpips_cache_path); + + lpips +} + +/// Load ImageNet pretrained backbone weights. +fn load_backbone_weights(lpips: Lpips, cache_path: &PathBuf) -> Lpips { + // Load directly into the inner struct to avoid enum variant issues + match lpips { + Lpips::Vgg(mut inner) => { + let mut store = PytorchStore::from_file(cache_path) + .allow_partial(true) + // VGG16 features.X -> extractor.convY_Z + .with_key_remapping(r"^features\.0\.", "extractor.conv1_1.") + .with_key_remapping(r"^features\.2\.", "extractor.conv1_2.") + .with_key_remapping(r"^features\.5\.", "extractor.conv2_1.") + .with_key_remapping(r"^features\.7\.", "extractor.conv2_2.") + .with_key_remapping(r"^features\.10\.", "extractor.conv3_1.") + .with_key_remapping(r"^features\.12\.", "extractor.conv3_2.") + .with_key_remapping(r"^features\.14\.", "extractor.conv3_3.") + .with_key_remapping(r"^features\.17\.", "extractor.conv4_1.") + .with_key_remapping(r"^features\.19\.", "extractor.conv4_2.") + .with_key_remapping(r"^features\.21\.", "extractor.conv4_3.") + .with_key_remapping(r"^features\.24\.", "extractor.conv5_1.") + .with_key_remapping(r"^features\.26\.", "extractor.conv5_2.") + .with_key_remapping(r"^features\.28\.", "extractor.conv5_3."); + if let Err(e) = inner.load_from(&mut store) { + log::warn!("Some VGG backbone weights could not be loaded: {:?}", e); + } + Lpips::Vgg(inner) + } + Lpips::Alex(mut inner) => { + let mut store = PytorchStore::from_file(cache_path) + .allow_partial(true) + // AlexNet features.X -> extractor.convY + .with_key_remapping(r"^features\.0\.", "extractor.conv1.") + .with_key_remapping(r"^features\.3\.", "extractor.conv2.") + .with_key_remapping(r"^features\.6\.", "extractor.conv3.") + .with_key_remapping(r"^features\.8\.", "extractor.conv4.") + .with_key_remapping(r"^features\.10\.", "extractor.conv5."); + if let Err(e) = inner.load_from(&mut store) { + log::warn!("Some AlexNet backbone weights could not be loaded: {:?}", e); + } + Lpips::Alex(inner) + } + Lpips::Squeeze(mut inner) => { + let mut store = PytorchStore::from_file(cache_path) + .allow_partial(true) + // SqueezeNet features.X -> extractor.* + .with_key_remapping(r"^features\.0\.", "extractor.conv1.") + .with_key_remapping(r"^features\.3\.", "extractor.fire1.") + .with_key_remapping(r"^features\.4\.", "extractor.fire2.") + .with_key_remapping(r"^features\.6\.", "extractor.fire3.") + .with_key_remapping(r"^features\.7\.", "extractor.fire4.") + .with_key_remapping(r"^features\.9\.", "extractor.fire5.") + .with_key_remapping(r"^features\.10\.", "extractor.fire6.") + .with_key_remapping(r"^features\.11\.", "extractor.fire7.") + .with_key_remapping(r"^features\.12\.", "extractor.fire8."); + if let Err(e) = inner.load_from(&mut store) { + log::warn!( + "Some SqueezeNet backbone weights could not be loaded: {:?}", + e + ); + } + Lpips::Squeeze(inner) + } + } +} + +/// Load LPIPS trained linear layer weights. +fn load_lpips_weights(lpips: Lpips, cache_path: &PathBuf) -> Lpips { + // Load directly into the inner struct to avoid enum variant issues + match lpips { + Lpips::Vgg(mut inner) => { + let mut store = PytorchStore::from_file(cache_path) + .allow_partial(true) + .with_key_remapping(r"^lin0\.model\.1\.", "lin0.") + .with_key_remapping(r"^lin1\.model\.1\.", "lin1.") + .with_key_remapping(r"^lin2\.model\.1\.", "lin2.") + .with_key_remapping(r"^lin3\.model\.1\.", "lin3.") + .with_key_remapping(r"^lin4\.model\.1\.", "lin4."); + if let Err(e) = inner.load_from(&mut store) { + log::warn!("Some VGG LPIPS weights could not be loaded: {:?}", e); + } + Lpips::Vgg(inner) + } + Lpips::Alex(mut inner) => { + let mut store = PytorchStore::from_file(cache_path) + .allow_partial(true) + .with_key_remapping(r"^lin0\.model\.1\.", "lin0.") + .with_key_remapping(r"^lin1\.model\.1\.", "lin1.") + .with_key_remapping(r"^lin2\.model\.1\.", "lin2.") + .with_key_remapping(r"^lin3\.model\.1\.", "lin3.") + .with_key_remapping(r"^lin4\.model\.1\.", "lin4."); + if let Err(e) = inner.load_from(&mut store) { + log::warn!("Some AlexNet LPIPS weights could not be loaded: {:?}", e); + } + Lpips::Alex(inner) + } + Lpips::Squeeze(mut inner) => { + let mut store = PytorchStore::from_file(cache_path) + .allow_partial(true) + .with_key_remapping(r"^lin0\.model\.1\.", "lin0.") + .with_key_remapping(r"^lin1\.model\.1\.", "lin1.") + .with_key_remapping(r"^lin2\.model\.1\.", "lin2.") + .with_key_remapping(r"^lin3\.model\.1\.", "lin3.") + .with_key_remapping(r"^lin4\.model\.1\.", "lin4.") + .with_key_remapping(r"^lin5\.model\.1\.", "lin5.") + .with_key_remapping(r"^lin6\.model\.1\.", "lin6."); + if let Err(e) = inner.load_from(&mut store) { + log::warn!("Some SqueezeNet LPIPS weights could not be loaded: {:?}", e); + } + Lpips::Squeeze(inner) + } + } +} diff --git a/crates/burn-train/src/metric/vision/mod.rs b/crates/burn-train/src/metric/vision/mod.rs new file mode 100644 index 0000000..9fbd425 --- /dev/null +++ b/crates/burn-train/src/metric/vision/mod.rs @@ -0,0 +1,17 @@ +mod afine; +mod dice; +mod dists; +mod fid; +mod lpips; +mod ms_ssim; +mod psnr; +mod ssim; + +pub use afine::*; +pub use dice::*; +pub use dists::*; +pub use fid::*; +pub use lpips::*; +pub use ms_ssim::*; +pub use psnr::*; +pub use ssim::*; diff --git a/crates/burn-train/src/metric/vision/ms_ssim.rs b/crates/burn-train/src/metric/vision/ms_ssim.rs new file mode 100644 index 0000000..518503f --- /dev/null +++ b/crates/burn-train/src/metric/vision/ms_ssim.rs @@ -0,0 +1,863 @@ +use crate::metric::{ + Metric, MetricAttributes, MetricMetadata, MetricName, Numeric, NumericAttributes, NumericEntry, + SerializedEntry, + state::{FormatOptions, NumericMetricState}, +}; +use burn_core::{ + prelude::{Device, Int, Tensor}, + tensor::{ + ElementConversion, + module::{avg_pool2d, conv2d}, + ops::{ConvOptions, PadMode}, + }, +}; + +/// Input type for the [MsSsimMetric]. +/// +/// Both tensors must have shape `[N, C, H, W]`: +/// - `N`: Batch size +/// - `C`: Number of channels (1 for grayscale, 3 for RGB, etc.) +/// - `H`: Height +/// - `W`: Width +/// +/// # Important +/// The image dimensions must be sufficiently large to accommodate the multi-scale +/// computation. Specifically, for the default 5 scales used by Burn, the image dimensions +/// should be at least `kernel_size * 2^(scales-1)` (e.g., 11 × 2^4 = 11 * 16 = 176 for default kernel size). +/// If your images are smaller, reduce the kernel size or number of scales. +/// +/// # Example +/// ```rust,ignore +/// // Create input for RGB images +/// let outputs: Tensor<4> = /* tensor */; +/// let targets: Tensor<4> = /* tensor */; +/// let input = MsSsimInput::new(outputs, targets); +/// ``` +pub struct MsSsimInput { + /// Model outputs with shape [N, C, H, W]. + outputs: Tensor<4>, + /// Ground truth targets with shape [N, C, H, W]. + targets: Tensor<4>, +} + +impl MsSsimInput { + /// Creates a new MsSsimInput with the given outputs and targets. + /// + /// # Arguments + /// - `outputs`: The model output images with shape [N, C, H, W]. + /// - `targets`: The ground truth images with shape [N, C, H, W]. + /// + /// # Returns + /// A new instance of `MsSsimInput`. + /// + /// # Panics + /// - If `outputs` and `targets` do not have the same shape. + pub fn new(outputs: Tensor<4>, targets: Tensor<4>) -> Self { + assert!( + outputs.dims() == targets.dims(), + "Shape mismatch: outputs {:?} targets {:?}", + outputs.dims(), + targets.dims() + ); + Self { outputs, targets } + } +} + +/// Configuration for the [MsSsimMetric]. +#[derive(Debug, Clone)] +pub struct MsSsimMetricConfig { + /// A parameter of SSIM used to stabilize the luminance comparison. + /// Default is 0.01. + pub k1: f32, + /// A parameter of SSIM used to stabilize the contrast comparison. + /// Default is 0.03. + pub k2: f32, + /// The range of the pixel values in images which can be computed as following: + /// `let pixel_range = max_pixel_val - min_pixel_val;` + /// where `max_pixel_val` is the maximum possible pixel value and `min_pixel_val` + /// is the minimum possible pixel value. + /// + /// - For normalized images in range [0, 1], it should be set to `1.0 - 0.0 = 1.0` + /// - For normalized images in range [-1, 1], it should be set to `1.0 - (-1.0) = 2.0` + /// - For 8-bit images in range [0, 255], it should be set to `255.0 - 0.0 = 255.0` + pub pixel_range: f32, + /// The MS-SSIM metric involves applying convolution to the input tensors using a Gaussian kernel. + /// This is the kernel size of the Gaussian kernel. Default is 11. + pub kernel_size: usize, + /// The MS-SSIM metric involves applying convolution to the input tensors using a Gaussian kernel. + /// This is the standard deviation of the Gaussian kernel. Default is 1.5. + pub sigma: f32, + /// The number of channels in the input images (e.g., 1 for grayscale, 3 for RGB). + /// This is used to create the appropriate convolution kernels. Default is 3. + pub channels: usize, + /// The weights/betas for each scale in the MS-SSIM computation. + /// The length of this vector determines the number of scales. + /// Default is \[0.0448, 0.2856, 0.3001, 0.2363, 0.1333\] (5 scales). + pub betas: Vec, +} + +impl MsSsimMetricConfig { + /// Creates a configuration with the specified data range and default parameters. + /// + /// # Default parameters + /// - k1: 0.01 + /// - k2: 0.03 + /// - kernel_size: 11 + /// - sigma: 1.5 + /// - channels: 3 + /// + /// # Panics + /// - If `pixel_range` is not positive. + /// + /// # Example + /// ```rust,ignore + /// // For normalized RGB images [0, 1] + /// let config1 = MsSsimMetricConfig::new(1.0); + /// + /// // For 8-bit images [0, 255] + /// let config2 = MsSsimMetricConfig::new(255.0); + /// + /// // For grayscale with custom kernel + /// let config3 = MsSsimMetricConfig::new(1.0) + /// .with_channels(1) + /// .with_kernel_size(7); + /// ``` + pub fn new(pixel_range: f32) -> Self { + assert!(pixel_range > 0.0, "pixel_range must be positive"); + Self { + k1: 0.01, + k2: 0.03, + pixel_range, + kernel_size: 11, + sigma: 1.5, + channels: 3, + betas: vec![0.0448, 0.2856, 0.3001, 0.2363, 0.1333], + } + } + + /// Sets custom values for the k1 and k2 parameters of MS-SSIM which are + /// used for numerical stability. + /// + /// # Default values + /// - k1: 0.01 + /// - k2: 0.03 + /// + /// # Panics + /// - If `k1` or `k2` is not positive. + pub fn with_k1_k2(mut self, k1: f32, k2: f32) -> Self { + assert!(k1 > 0.0, "k1 must be positive"); + assert!(k2 > 0.0, "k2 must be positive"); + self.k1 = k1; + self.k2 = k2; + self + } + + /// Sets a custom kernel size for the Gaussian kernel used in MS-SSIM. The + /// kernel size must be a positive odd number. + /// + /// # Default value + /// - kernel_size: 11 + /// + /// # Panics + /// - If `kernel_size` is not a positive odd number. + pub fn with_kernel_size(mut self, kernel_size: usize) -> Self { + assert!( + kernel_size > 0 && kernel_size % 2 == 1, + "kernel_size must be positive and an odd number" + ); + self.kernel_size = kernel_size; + self + } + + /// Sets a custom sigma (standard deviation) for the Gaussian kernel used in MS-SSIM. + /// + /// # Default value + /// - sigma: 1.5 + /// + /// # Panics + /// - If `sigma` is not positive. + pub fn with_sigma(mut self, sigma: f32) -> Self { + assert!(sigma > 0.0, "sigma must be a positive number"); + self.sigma = sigma; + self + } + + /// Sets the number of channels for the input images. + /// + /// This affects the shape of the pre-computed convolution kernels. + /// Change this if working with grayscale (1) or multispectral images (>3). + /// + /// # Default value + /// - channels: 3 + /// + /// # Panics + /// - If `channels` is 0. + pub fn with_channels(mut self, channels: usize) -> Self { + assert!(channels > 0, "channels must be a positive number"); + self.channels = channels; + self + } + + /// Sets custom betas for the scales. The length of the betas vector + /// determines the number of scales used in the MS-SSIM computation. + /// If you want to make different parameter settings comparable, the betas + /// vector should sum to 1 as per the original paper. However, note + /// that this is not a strict requirement. + /// + /// # Default value + /// - betas: `[0.0448, 0.2856, 0.3001, 0.2363, 0.1333]` (5 scales) + /// + /// # Panics + /// - If `betas` is empty. + /// - If not all values in `betas` are positive. + pub fn with_betas(mut self, betas: Vec) -> Self { + assert!(!betas.is_empty(), "betas vector cannot be empty"); + + assert!( + betas.iter().all(|&b| b >= 0.0), + "All beta values must be non-negative" + ); + + self.betas = betas; + self + } +} + +/// Multi-Scale Structural Similarity Index (MS-SSIM) metric for image quality assessment. +/// +/// MS-SSIM extends the single-scale [SSIM](crate::metric::vision::SsimMetric) by computing +/// the index at multiple resolutions (scales) and combining them using weighted averaging. +/// This approach better correlates with human visual perception, especially for +/// high-resolution images where fine details and texture variations are important. +/// +/// # Algorithm Overview +/// +/// MS-SSIM computes structural similarity across M scales (M=5 in Burn): +/// +/// 1. **Contrast** and **Structure** components are computed at every scale +/// 2. **Luminance** is computed only at the coarsest (last) scale +/// 3. Between scales, images are downsampled by a factor of 2 using average pooling +/// +/// The final metric is computed as: +/// ```text +/// MS-SSIM = L_M^{α_M} × ∏_{j=1}^M (C_j^{β_j} × S_j^{γ_j}) +/// ``` +/// +/// Where: +/// - `L_M` is luminance at the last scale (M) +/// - `C_j` is contrast at scale j: `(2σ_xσ_y + C2) / (σ_x² + σ_y² + C2)` +/// - `S_j` is structure at scale j: `(σ_xy + C3) / (σ_xσ_y + C3)` +/// - `α_M, β_j, γ_j` are weights from Wang et al. (\[0.0448, 0.2856, 0.3001, 0.2363, 0.1333\]) +/// +/// # Notes +/// +/// - This implementation uses separable Gaussian convolution for efficiency (reduces complexity from O(K^2) to O(2K) per pixel) +/// - Gaussian kernels are pre-computed during initialization to avoid redundant computation +/// - The metric requires images to be large enough to survive the downsampling operations +/// +/// # Value Range +/// +/// MS-SSIM values typically range from 0 to 1, where: +/// - 1.0 indicates perfect structural similarity (identical images) +/// - 0.0 indicates no structural similarity +/// - Values are usually positive due to the stability constants (C1, C2, C3) +/// +/// # References +/// +/// [Multi-scale Structural Similarity for Image Quality Assessment](https://www.cns.nyu.edu/pub/eero/wang03b.pdf) +#[derive(Clone)] +pub struct MsSsimMetric { + name: MetricName, + /// Internal state for numeric metric aggregation. + state: NumericMetricState, + /// Configuration for the metric. + config: MsSsimMetricConfig, + /// Pre-computed horizontal Gaussian kernel with shape [C, 1, 1, K] + horizontal_kernel: Tensor<4>, + /// Pre-computed vertical Gaussian kernel with shape [C, 1, K, 1] + vertical_kernel: Tensor<4>, +} + +impl MsSsimMetric { + /// Creates a new MS-SSIM metric with the given configuration. + /// + /// # Arguments + /// - `config`: Configuration for the metric (data range, kernel size, etc.) + /// - `device`: Device to place the Gaussian kernels on + /// + /// # Note + /// The default metric name format is "MS-SSIM (pr={}, k={}, σ={})" + /// where pr is the pixel range, k is the kernel size, and σ is the + /// standard deviation. + /// + /// # Example + /// ```ignore + /// let config = MsSsimMetricConfig::new(1.0).with_channels(1); // Grayscale + /// let metric = MsSsimMetric::new(config, &device); + /// ``` + pub fn new(config: MsSsimMetricConfig, device: &Device) -> Self { + let kernel = Self::create_1d_gaussian_kernel(&config, device); + let size = config.kernel_size; + + // Create horizontal kernel: shape [C, 1, 1, K] for depthwise conv + let horizontal_kernel = kernel + .clone() + .reshape([1, 1, 1, size]) + .repeat_dim(0, config.channels); + + // Create vertical kernel: shape [C, 1, K, 1] for depthwise conv + let vertical_kernel = kernel + .reshape([1, 1, size, 1]) + .repeat_dim(0, config.channels); + + Self { + name: MetricName::new(format!( + "MS-SSIM (pr={}, k={}, σ={})", + config.pixel_range, config.kernel_size, config.sigma + )), + state: NumericMetricState::default(), + config, + horizontal_kernel, + vertical_kernel, + } + } + + /// Overrides the default metric name. + /// + /// # Example + /// ```ignore + /// let metric = MsSsimMetric::new(config, &device) + /// .with_name("Custom MS-SSIM Name"); + /// ``` + pub fn with_name(mut self, name: &str) -> Self { + self.name = MetricName::new(name.to_string()); + self + } + + /// Creates a normalized 1D Gaussian kernel as a tensor where the kernel values sum to 1.0. + fn create_1d_gaussian_kernel(config: &MsSsimMetricConfig, device: &Device) -> Tensor<1> { + let size = config.kernel_size as i64; + let sigma = config.sigma; + let center = (size / 2) as f32; + + let one_to_size_tensor = Tensor::<1, Int>::arange(0..size, device).float(); + let x_vals = one_to_size_tensor.sub_scalar(center); + + // Gaussian: exp(-x² / 2σ²) + let x_squared = x_vals.clone().mul(x_vals); + let x_squared_div_2_sigma_squared = x_squared.div_scalar(2.0 * sigma * sigma); + let unnormalized_kernel = x_squared_div_2_sigma_squared.neg().exp(); + let kernel_vals_sum = unnormalized_kernel.clone().sum(); + unnormalized_kernel.div(kernel_vals_sum) + } + + /// Applies separable Gaussian convolution using pre-computed kernels. + /// + /// Performs two 1D convolutions (horizontal then vertical) which is + /// computationally cheaper than a single 2D convolution. + /// + /// # Arguments + /// - `input`: Tensor of shape [N, C, H, W] + fn gaussian_separable_conv(&self, input: Tensor<4>) -> Tensor<4> { + let padding = self.config.kernel_size / 2; + let h_kernel = self.horizontal_kernel.clone(); + let v_kernel = self.vertical_kernel.clone(); + + // Apply reflect padding to all 4 sides of the input tensor before convolution + // Format: (left, right, top, bottom) + let padded_input = input.pad((padding, padding, padding, padding), PadMode::Reflect); + + let h_conv_options = ConvOptions::new([1, 1], [0, 0], [1, 1], self.config.channels); + let v_conv_options = ConvOptions::new([1, 1], [0, 0], [1, 1], self.config.channels); + + let input_after_h_conv = conv2d(padded_input, h_kernel, None, h_conv_options); + conv2d(input_after_h_conv, v_kernel, None, v_conv_options) + } +} + +impl Metric for MsSsimMetric { + type Input = MsSsimInput; + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn update(&mut self, item: &Self::Input, _metadata: &MetricMetadata) -> SerializedEntry { + let dims = item.outputs.dims(); + let scales = self.config.betas.len(); + + assert_eq!( + dims[1], self.config.channels, + "Input has {} channels but metric was configured for {}", + dims[1], self.config.channels + ); + + // Verify minimum size for the given number of scales + // After (scales - 1) downsamples, size is original / 2^(scales-1) + // We need kernel_size at that scale + let downsample_ops_num = scales.saturating_sub(1) as u32; + let min_size = self.config.kernel_size * (2usize.pow(downsample_ops_num)); + let h = dims[2]; + let w = dims[3]; + assert!( + h >= min_size && w >= min_size, + "Image dimensions (H={}, W={}) must be at least {} to support {} scales of MS-SSIM \ + with kernel_size={}. Either increase image size, reduce kernel_size, or reduce the number of scales (betas).", + h, + w, + min_size, + scales, + self.config.kernel_size + ); + + let mut x = item.outputs.clone(); + let mut y = item.targets.clone(); + let betas = &self.config.betas; + + // Compute c1 = (k1 * L)^2 and c2 = (k2 * L)^2, c3 = c2/2 + let c1 = (self.config.k1 * self.config.pixel_range).powi(2); + let c2 = (self.config.k2 * self.config.pixel_range).powi(2); + + // Initialize accumulator to 1 for update via multiplication + // Shape: [N, C] + let batch_size = dims[0]; + let channels = dims[1]; + let mut ms_ssim_tensor = Tensor::<2>::ones([batch_size, channels], &item.outputs.device()); + + for (j, beta_j) in betas.iter().enumerate() { + // Compute mu_x and mu_y + let mu_x = self.gaussian_separable_conv(x.clone()); + let mu_y = self.gaussian_separable_conv(y.clone()); + let square_of_mu_x = mu_x.clone() * mu_x.clone(); + let square_of_mu_y = mu_y.clone() * mu_y.clone(); + + // Var(X) = E(X^2) - E(X)^2 + let mu_of_x_squared = self.gaussian_separable_conv(x.clone() * x.clone()); + let mu_of_y_squared = self.gaussian_separable_conv(y.clone() * y.clone()); + let var_x = (mu_of_x_squared - square_of_mu_x.clone()).clamp_min(0.0); + let var_y = (mu_of_y_squared - square_of_mu_y.clone()).clamp_min(0.0); + + // Cov(X, Y) = E(XY) - E(X)E(Y) + let mu_of_xy = self.gaussian_separable_conv(x.clone() * y.clone()); + let cov_xy = mu_of_xy - (mu_x.clone() * mu_y.clone()); + + // Compute cs_map = (2σxy + C2) / (σx² + σy² + C2) + // This is mathematically equivalent to c(x,y) * s(x,y) when C3 = C2 / 2 + let contrast_structure = (cov_xy * 2.0 + c2) / (var_x + var_y + c2); + + // Include luminance at the last scale + if j == betas.len() - 1 { + // Compute l(x, y) = (2μxμy + C1) / (μx² + μy² + C1) + let luminance: Tensor<4> = + (2 * mu_x * mu_y + c1) / (square_of_mu_x + square_of_mu_y + c1); + let ssim = luminance * contrast_structure; + let ssim_spatial_mean = ssim.mean_dims(&[2, 3]).reshape([batch_size, channels]); + // Clamp to avoid negative values before raising to power (prevents NaNs) + let ssim_mean_clamped = ssim_spatial_mean.clamp_min(0.0); + ms_ssim_tensor = ms_ssim_tensor * ssim_mean_clamped.powf_scalar(*beta_j); + } else { + let contrast_structure_spatial_mean = contrast_structure + .mean_dims(&[2, 3]) + .reshape([batch_size, channels]); + // Clamp to avoid negative values before raising to power (prevents NaNs) + let c_s_mean_clamped = contrast_structure_spatial_mean.clamp_min(0.0); + ms_ssim_tensor = ms_ssim_tensor * c_s_mean_clamped.powf_scalar(*beta_j); + + x = avg_pool2d(x, [2, 2], [2, 2], [0, 0], false, false); + y = avg_pool2d(y, [2, 2], [2, 2], [0, 0], false, false); + } + } + + let ms_ssim_per_image = ms_ssim_tensor.mean_dim(1); + let avg_ms_ssim = ms_ssim_per_image.mean().into_scalar::(); + + self.state.update( + avg_ms_ssim, + batch_size, + FormatOptions::new(self.name()).precision(4), + ) + } + + /// Clears the metric state. + fn clear(&mut self) { + self.state.reset(); + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: None, + higher_is_better: true, + ..Default::default() + } + .into() + } +} + +impl Numeric for MsSsimMetric { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::metric::Numeric; + use burn_core::tensor::Distribution; + + fn test_config() -> MsSsimMetricConfig { + // Use small kernel and single channel for testing + // With kernel_size=3, we need images >= 3*16=48 + MsSsimMetricConfig::new(1.0) + .with_kernel_size(3) + .with_sigma(1.0) + .with_channels(1) + } + + #[test] + fn test_ms_ssim_perfect_similarity() { + // Identical images should give MS-SSIM = 1.0 + let device = Default::default(); + let outputs = Tensor::<4>::from_data( + [[[ + [0.5_f32; 64]; 64 // 64x64 constant image + ]]], + &device, + ); + let targets = outputs.clone(); + + let mut metric = MsSsimMetric::new(test_config(), &device); + let input = MsSsimInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let ms_ssim = metric.value().current(); + assert!( + ms_ssim > 0.99, + "MS-SSIM for identical images should be 1.0, got {}", + ms_ssim + ); + } + + #[test] + fn test_ms_ssim_completely_different() { + // Black vs white images should give very low MS-SSIM (close to 0.0) + let device = Default::default(); + let outputs = Tensor::<4>::zeros([1, 1, 256, 256], &device); + let targets = Tensor::<4>::ones([1, 1, 256, 256], &device); + + let mut metric = MsSsimMetric::new(test_config(), &device); + let input = MsSsimInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let ms_ssim = metric.value().current(); + assert!( + (ms_ssim - 0.3).abs() < 0.01, + "MS-SSIM for black vs white should be low (around 0.3), got {}", + ms_ssim + ); + } + + #[test] + fn test_ms_ssim_similar_images() { + // Small perturbation should give high MS-SSIM (close to 1.0) + let device = Default::default(); + let outputs = Tensor::<4>::full([1, 1, 64, 64], 0.5, &device); + let targets = Tensor::<4>::full([1, 1, 64, 64], 0.52, &device); + + let mut metric = MsSsimMetric::new(test_config(), &device); + let input = MsSsimInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let ms_ssim = metric.value().current(); + assert!( + ms_ssim > 0.95, + "MS-SSIM for very similar images should be close to 1.0, got {}", + ms_ssim + ); + } + + #[test] + fn test_ms_ssim_batch_averaging() { + let device = Default::default(); + // Batch of 2: one identical, one different + let outputs = Tensor::<4>::from_data( + [ + [[[0.5_f32; 64]; 64]], // Image 1: constant 0.5 + [[[0.0_f32; 64]; 64]], // Image 2: constant 0.0 (black) + ], + &device, + ); + let targets = Tensor::<4>::from_data( + [ + [[[0.5_f32; 64]; 64]], // Image 1: identical + [[[1.0_f32; 64]; 64]], // Image 2: white (opposite) + ], + &device, + ); + + let mut metric = MsSsimMetric::new(test_config(), &device); + let input = MsSsimInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let ms_ssim = metric.value().current(); + // Average of ~1.0 and ~0.292 should be around 0.64 + assert!( + (ms_ssim - 0.64).abs() < 0.02, + "Average MS-SSIM should be around 0.64, got {}", + ms_ssim + ); + } + + #[test] + fn test_ms_ssim_multichannel() { + let device = Default::default(); + // Test with 3 channels (RGB) + let config = MsSsimMetricConfig::new(1.0) + .with_kernel_size(3) + .with_sigma(1.0) + .with_channels(3); + + let outputs = Tensor::<4>::random([2, 3, 64, 64], Distribution::Uniform(0.0, 1.0), &device); + let targets = outputs.clone(); + + let mut metric = MsSsimMetric::new(config, &device); + let input = MsSsimInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let ms_ssim = metric.value().current(); + assert!( + ms_ssim > 0.99, + "MS-SSIM for identical RGB images should be 1.0, got {}", + ms_ssim + ); + } + + #[test] + fn test_ms_ssim_running_average() { + let device = Default::default(); + let mut metric = MsSsimMetric::new(test_config(), &device); + + // First update: identical (1.0) + let img1 = Tensor::<4>::full([1, 1, 64, 64], 0.5, &device); + let input1 = MsSsimInput::new(img1.clone(), img1); + metric.update(&input1, &MetricMetadata::fake()); + + assert!( + metric.value().current() > 0.99, + "First update should be approximately 1.0" + ); + + // Second update: different (~0.29) + let black = Tensor::<4>::zeros([1, 1, 64, 64], &device); + let white = Tensor::<4>::ones([1, 1, 64, 64], &device); + let input2 = MsSsimInput::new(black, white); + metric.update(&input2, &MetricMetadata::fake()); + + let running = metric.running_value().current(); + assert!( + (running - 0.64).abs() < 0.02, + "Running average should be approximately 0.64, got {}", + running + ); + } + + #[test] + fn test_ms_ssim_single_scale_small_image() { + let device = Default::default(); + // Default 5 scales with kernel_size=11 requires a 176x176 image. + // With a single scale, the minimum required size drops to + // just 11x11 (kernel_size * 2^0). + let config = MsSsimMetricConfig::new(1.0) + .with_channels(1) + .with_betas(vec![1.0]); // 1 scale + + let mut metric = MsSsimMetric::new(config, &device); + + // Create a 16x16 image. This would normally panic with 5 scales, + // but should succeed with 1 scale. + let outputs = Tensor::<4>::zeros([1, 1, 16, 16], &device); + let targets = outputs.clone(); + let input = MsSsimInput::new(outputs, targets); + + // This should not panic + let _ = metric.update(&input, &MetricMetadata::fake()); + + // Identical images should still yield ~1.0 + let ms_ssim = metric.value().current(); + assert!( + ms_ssim > 0.99, + "1-scale MS-SSIM for identical images should be 1.0, got {}", + ms_ssim + ); + } + + #[test] + fn test_ssim_symmetry() { + // MS-SSIM(x, y) should equal MS-SSIM(y, x) + // Symmetry is one of the mathematical properties of MS-SSIM + let device = Default::default(); + let config = MsSsimMetricConfig::new(1.0) + .with_kernel_size(3) + .with_sigma(1.0) + .with_channels(3); + + let img1 = Tensor::<4>::random([2, 3, 64, 64], Distribution::Uniform(0.0, 1.0), &device); + let img2 = Tensor::<4>::random([2, 3, 64, 64], Distribution::Uniform(0.0, 1.0), &device); + + let mut metric1 = MsSsimMetric::new(config.clone(), &device); + let input1 = MsSsimInput::new(img1.clone(), img2.clone()); + let _entry = metric1.update(&input1, &MetricMetadata::fake()); + let ms_ssim1 = metric1.value().current(); + + let mut metric2 = MsSsimMetric::new(config, &device); + let input2 = MsSsimInput::new(img2, img1); + let _entry = metric2.update(&input2, &MetricMetadata::fake()); + let ms_ssim2 = metric2.value().current(); + + assert!( + (ms_ssim1 - ms_ssim2).abs() < 0.001, + "MS-SSIM should be symmetric: MS-SSIM(x,y)={} vs MS-SSIM(y,x)={}", + ms_ssim1, + ms_ssim2 + ); + } + + #[test] + fn test_ms_ssim_clear() { + let device = Default::default(); + let mut metric = MsSsimMetric::new(test_config(), &device); + + let img = Tensor::<4>::full([1, 1, 64, 64], 0.5, &device); + let input = MsSsimInput::new(img.clone(), img); + metric.update(&input, &MetricMetadata::fake()); + + assert!(metric.value().current() > 0.99); + + metric.clear(); + assert!(metric.running_value().current().is_nan()); + } + + #[test] + fn test_ms_ssim_custom_name() { + let device = Default::default(); + let config = MsSsimMetricConfig::new(1.0); + let metric = MsSsimMetric::new(config, &device).with_name("CustomMS-SSIM"); + assert_eq!(metric.name().to_string(), "CustomMS-SSIM"); + } + + #[test] + fn test_ms_ssim_default_name() { + let device = Default::default(); + let config = MsSsimMetricConfig::new(255.0); + let metric = MsSsimMetric::new(config, &device); + assert_eq!(metric.name().to_string(), "MS-SSIM (pr=255, k=11, σ=1.5)"); + } + + #[test] + fn test_ms_ssim_attributes() { + let device = Default::default(); + let config = MsSsimMetricConfig::new(1.0); + let metric = MsSsimMetric::new(config, &device); + + match metric.attributes() { + MetricAttributes::Numeric(attrs) => { + assert!(attrs.higher_is_better); + assert_eq!(attrs.unit, None); + } + _ => panic!("Expected numeric attributes"), + } + } + + #[test] + #[should_panic(expected = "Shape mismatch")] + fn test_ms_ssim_shape_mismatch() { + let device = Default::default(); + let outputs = Tensor::<4>::zeros([1, 1, 64, 64], &device); + let targets = Tensor::<4>::zeros([1, 1, 32, 32], &device); + let _ = MsSsimInput::new(outputs, targets); + } + + #[test] + #[should_panic(expected = "k1 must be positive")] + fn test_ms_ssim_negative_k1() { + let _ = MsSsimMetricConfig::new(1.0).with_k1_k2(-0.01, 0.03); + } + + #[test] + #[should_panic(expected = "k2 must be positive")] + fn test_ms_ssim_negative_k2() { + let _ = MsSsimMetricConfig::new(1.0).with_k1_k2(0.01, -0.03); + } + + #[test] + #[should_panic(expected = "pixel_range must be positive")] + fn test_ms_ssim_negative_data_range() { + let _ = MsSsimMetricConfig::new(-1.0); + } + + #[test] + #[should_panic(expected = "pixel_range must be positive")] + fn test_ms_ssim_zero_data_range() { + let _ = MsSsimMetricConfig::new(0.0); + } + + #[test] + #[should_panic(expected = "kernel_size must be positive and an odd number")] + fn test_ms_ssim_even_kernel_size() { + let _ = MsSsimMetricConfig::new(1.0).with_kernel_size(10); + } + + #[test] + #[should_panic(expected = "kernel_size must be positive and an odd number")] + fn test_ms_ssim_zero_kernel_size() { + let _ = MsSsimMetricConfig::new(1.0).with_kernel_size(0); + } + + #[test] + #[should_panic(expected = "sigma must be a positive number")] + fn test_ms_ssim_negative_sigma() { + let _ = MsSsimMetricConfig::new(1.0).with_sigma(-1.5); + } + + #[test] + #[should_panic(expected = "sigma must be a positive number")] + fn test_ms_ssim_zero_sigma() { + let _ = MsSsimMetricConfig::new(1.0).with_sigma(0.0); + } + + #[test] + #[should_panic(expected = "channels must be a positive number")] + fn test_ms_ssim_zero_channels() { + let _ = MsSsimMetricConfig::new(1.0).with_channels(0); + } + + #[test] + #[should_panic(expected = "betas vector cannot be empty")] + fn test_ms_ssim_empty_betas() { + let _ = MsSsimMetricConfig::new(1.0).with_betas(vec![]); + } + + #[test] + #[should_panic(expected = "All beta values must be non-negative")] + fn test_ms_ssim_negative_betas() { + let _ = MsSsimMetricConfig::new(1.0).with_betas(vec![0.3, 0.3, -0.1, 0.5]); + } + + #[test] + #[should_panic(expected = "Image dimensions")] + fn test_ms_ssim_image_too_small() { + let device = Default::default(); + // 3 scales with kernel_size=11 requires 44x44 minimum (11 * 2^2) + let config = MsSsimMetricConfig::new(1.0).with_betas(vec![0.5, 0.3, 0.2]); + let mut metric = MsSsimMetric::new(config, &device); + + let outputs = Tensor::<4>::zeros([1, 3, 32, 32], &device); // Too small (32 < 44) + let targets = outputs.clone(); + let input = MsSsimInput::new(outputs, targets); + let _ = metric.update(&input, &MetricMetadata::fake()); + } +} diff --git a/crates/burn-train/src/metric/vision/psnr.rs b/crates/burn-train/src/metric/vision/psnr.rs new file mode 100644 index 0000000..86a4b2e --- /dev/null +++ b/crates/burn-train/src/metric/vision/psnr.rs @@ -0,0 +1,590 @@ +use crate::metric::{ + Metric, MetricAttributes, MetricMetadata, MetricName, Numeric, NumericAttributes, NumericEntry, + SerializedEntry, + state::{FormatOptions, NumericMetricState}, +}; +use burn_core::{prelude::Tensor, tensor::ElementConversion}; +use std::f64::consts::LN_10; + +/// Input type for the [PsnrMetric]. +/// +/// Both tensors must have shape `[N, C, H, W]`: +/// - `N`: Batch size +/// - `C`: Number of channels (1 for grayscale, 3 for RGB, etc.) +/// - `H`: Height +/// - `W`: Width +pub struct PsnrInput { + /// Model output (predictions/reconstructions) images with shape `[N, C, H, W]`. + outputs: Tensor<4>, + /// Ground truth images with shape `[N, C, H, W]`. + targets: Tensor<4>, +} + +impl PsnrInput { + /// Creates a new PsnrInput with the given outputs and targets. + /// + /// Inputs are expected to have the dimensions `[N, C, H, W]` + /// where `N` is the batch size, `C` is the number of channels, + /// `H` is the height of the image, and `W` is the width of the image. + /// + /// # Arguments + /// - `outputs`: The model output images with shape `[N, C, H, W]`. + /// - `targets`: The ground truth images with shape `[N, C, H, W]`. + /// + /// # Returns + /// A new instance of `PsnrInput`. + /// + /// # Panics + /// - If `outputs` and `targets` do not have the same shape. + pub fn new(outputs: Tensor<4>, targets: Tensor<4>) -> Self { + assert!( + outputs.dims() == targets.dims(), + "Shape mismatch: outputs {:?}, targets {:?}", + outputs.dims(), + targets.dims() + ); + Self { outputs, targets } + } +} + +/// Configuration for the [PsnrMetric]. +#[derive(Debug, Clone, Copy)] +pub struct PsnrMetricConfig { + /// Maximum possible pixel value. + /// - Use `1.0` for normalized images in range \[0, 1\] + /// - Use `255.0` for 8-bit images in range \[0, 255\] + pub max_pixel_val: f64, + /// Epsilon value for numerical stability when MSE is very small or zero. + /// + /// When MSE falls below this threshold, it is clamped to `epsilon`, + /// resulting in a maximum PSNR of approximately `10 * log10(max_pixel_val² / epsilon)` dB. + /// + /// Default is `1e-10`, which yields ~100 dB for perfect reconstruction with `max_pixel_val = 1.0`. + pub epsilon: f64, +} + +impl PsnrMetricConfig { + /// Creates a configuration with the specified maximum pixel value. + /// + /// # Example + /// ```ignore + /// // Normalized images [0, 1] + /// let config = PsnrMetricConfig::new(1.0); + /// + /// // 8-bit images [0, 255] + /// let config = PsnrMetricConfig::new(255.0); + /// // Also set a custom epsilon value + /// let config = PsnrMetricConfig::new(255.0).with_epsilon(1e-8); + /// ``` + pub fn new(max_pixel_val: f64) -> Self { + assert!(max_pixel_val > 0.0, "max_pixel_val must be positive"); + Self { + max_pixel_val, + epsilon: 1e-10, + } + } + + /// Sets a custom epsilon for numerical stability near zero MSE + pub fn with_epsilon(mut self, epsilon: f64) -> Self { + assert!(epsilon > 0.0, "epsilon must be positive"); + self.epsilon = epsilon; + self + } +} + +/// The peak signal-to-noise ratio (PSNR) metric for image quality assessment. +/// +/// PSNR is commonly used to measure the quality of reconstructed images +/// compared to the original. Higher values (in dB) indicate better quality. +/// +/// # Formula +/// ```text +/// PSNR = 10 * log10(MAX^2 / MSE) +/// ``` +/// where MAX is the maximum possible pixel value and MSE is the mean squared error. +/// +/// # Note +/// - PSNR is computed for each image first, and then it is averaged across all the images in the batch. +/// - For perfect reconstruction (MSE = 0), the MSE is clamped to `epsilon` to avoid division by zero, +/// yielding a maximum PSNR of `10 * log10(MAX^2 / epsilon)` dB. +#[derive(Clone)] +pub struct PsnrMetric { + name: MetricName, + /// Internal state for numeric metric aggregation. + state: NumericMetricState, + /// Configuration for the metric. + config: PsnrMetricConfig, +} + +impl PsnrMetric { + /// Creates a new PSNR metric with the given configuration. + /// + /// # Example + /// ```ignore + /// let config = PsnrMetricConfig::new(1.0); + /// let metric = PsnrMetric::new(config); + /// ``` + pub fn new(config: PsnrMetricConfig) -> Self { + Self { + name: MetricName::new(format!("PSNR@{}", config.max_pixel_val)), + state: NumericMetricState::default(), + config, + } + } + + /// Overrides the default metric name which is `PSNR@{max_pixel_val}`. + /// + /// Examples names: + /// - `PSNR@1.0` + /// - `PSNR@255.0` + /// + /// Use this method to provide a custom name. + pub fn with_name(mut self, name: &str) -> Self { + self.name = MetricName::new(name.to_string()); + self + } +} + +impl Metric for PsnrMetric { + type Input = PsnrInput; + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn update(&mut self, item: &Self::Input, _metadata: &MetricMetadata) -> SerializedEntry { + let dims = item.outputs.dims(); + let batch_size = dims[0]; + let outputs = item.outputs.clone(); + let targets = item.targets.clone(); + + // Compute per-image MSE by reducing over all dimensions except batch (dims 1, 2, 3) + // Resulting shape: [N, 1, 1, 1] + let diff = outputs.sub(targets); + let mse_per_image = diff.powi_scalar(2).mean_dims(&[1, 2, 3]); + // Flatten to shape: [N] + let mse_flat = mse_per_image.flatten::<1>(0, 3); + // Clamp MSE to avoid division by 0 in the expression (MAX^2 / MSE) + let mse_clamped = mse_flat.clamp_min(self.config.epsilon); + let max_squared = self.config.max_pixel_val * self.config.max_pixel_val; + + // Compute PSNR for each image and accumulate + // PSNR value in dB (using the change of base formula): + // 10 * log10(MAX^2 / MSE) = 10 * ln(MAX^2 / MSE) / ln(10) + // = ln(MAX^2 / MSE) * (10 / ln(10)) + let psnr_per_image = mse_clamped + .recip() + .mul_scalar(max_squared) + .log() + .mul_scalar(10.0 / LN_10); + let avg_psnr = psnr_per_image.mean().into_scalar::(); + + self.state.update( + avg_psnr, + batch_size, + FormatOptions::new(self.name()).unit("dB").precision(2), + ) + } + + /// Clears the metric state. + fn clear(&mut self) { + self.state.reset(); + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: Some("dB".to_string()), + higher_is_better: true, + ..Default::default() + } + .into() + } +} + +impl Numeric for PsnrMetric { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::metric::Numeric; + use burn_core::tensor::TensorData; + + #[test] + fn test_psnr_perfect_reconstruction() { + // When outputs exactly match targets, PSNR should be very high + // (limited by epsilon clamping to ~100 dB with default epsilon=1e-10) + let device = Default::default(); + let outputs = Tensor::<4>::from_data( + TensorData::from([[[[1.0_f32, 0.5], [0.25, 0.75]]]]), + &device, + ); + let targets = outputs.clone(); + + let config = PsnrMetricConfig::new(1.0); + let mut metric = PsnrMetric::new(config); + let input = PsnrInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + // With epsilon = 1e-10 and max=1.0: + // PSNR = 10 * log10(1.0 / 1e-10) = 100 dB + let psnr = metric.value().current(); + assert!( + psnr >= 99.0, + "PSNR for perfect reconstruction should be ~100 dB, got {} dB", + psnr + ); + } + + #[test] + fn test_psnr_constant_error() { + // Constant error of 0.1 across all pixels + // MSE = 0.01, PSNR = 10 * log10(1.0 / 0.01) = 20 dB + let device = Default::default(); + let outputs = + Tensor::<4>::from_data(TensorData::from([[[[0.1_f32, 0.1], [0.1, 0.1]]]]), &device); + let targets = + Tensor::<4>::from_data(TensorData::from([[[[0.0_f32, 0.0], [0.0, 0.0]]]]), &device); + + let config = PsnrMetricConfig::new(1.0); + let mut metric = PsnrMetric::new(config); + let input = PsnrInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let psnr = metric.value().current(); + assert!( + (psnr - 20.0).abs() < 0.01, + "Expected PSNR ~20 dB, got {} dB", + psnr + ); + } + + #[test] + fn test_psnr_varying_error() { + // Errors: 0.1, 0.2, 0.3, 0.4 → squared: 0.01, 0.04, 0.09, 0.16 + // MSE = 0.075, PSNR = 10 * log10(1.0 / 0.075) ≈ 11.249 dB + let device = Default::default(); + let outputs = + Tensor::<4>::from_data(TensorData::from([[[[0.1_f32, 0.2], [0.3, 0.4]]]]), &device); + let targets = + Tensor::<4>::from_data(TensorData::from([[[[0.0_f32, 0.0], [0.0, 0.0]]]]), &device); + + let config = PsnrMetricConfig::new(1.0); + let mut metric = PsnrMetric::new(config); + let input = PsnrInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let psnr = metric.value().current(); + let expected_psnr = 10.0 * (1.0_f64 / 0.075).log10(); + assert!( + (psnr - expected_psnr).abs() < 0.01, + "Expected PSNR ~{:.3} dB, got {} dB", + expected_psnr, + psnr + ); + } + + #[test] + fn test_psnr_max_pixel_255() { + // Test with 8-bit image range [0, 255] + // Error = 10 everywhere, MSE = 100 + // PSNR = 10 * log10(255^2 / 100) ≈ 28.13 dB + let device = Default::default(); + let outputs = Tensor::<4>::from_data( + TensorData::from([[[[10.0_f32, 10.0], [10.0, 10.0]]]]), + &device, + ); + let targets = + Tensor::<4>::from_data(TensorData::from([[[[0.0_f32, 0.0], [0.0, 0.0]]]]), &device); + + let config = PsnrMetricConfig::new(255.0); + let mut metric = PsnrMetric::new(config); + let input = PsnrInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let psnr = metric.value().current(); + let expected_psnr = 10.0 * (255.0_f64 * 255.0 / 100.0).log10(); + assert!( + (psnr - expected_psnr).abs() < 0.01, + "Expected PSNR ~{:.3} dB, got {} dB", + expected_psnr, + psnr + ); + } + + #[test] + fn test_psnr_batch_averaging() { + // Batch of 2 images with different MSEs + // Image 1: error 0.1 → MSE = 0.01 → PSNR = 20 dB + // Image 2: error 0.01 → MSE = 0.0001 → PSNR = 40 dB + // Average PSNR = 30 dB + let device = Default::default(); + let outputs = Tensor::<4>::from_data( + TensorData::from([ + [[[0.1_f32, 0.1], [0.1, 0.1]]], + [[[0.01_f32, 0.01], [0.01, 0.01]]], + ]), + &device, + ); + let targets = Tensor::<4>::from_data( + TensorData::from([ + [[[0.0_f32, 0.0], [0.0, 0.0]]], + [[[0.0_f32, 0.0], [0.0, 0.0]]], + ]), + &device, + ); + + let config = PsnrMetricConfig::new(1.0); + let mut metric = PsnrMetric::new(config); + let input = PsnrInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let psnr = metric.value().current(); + let expected_psnr = 30.0; + assert!( + (psnr - expected_psnr).abs() < 0.01, + "Expected average PSNR ~{} dB, got {} dB", + expected_psnr, + psnr + ); + } + + #[test] + fn test_psnr_multichannel() { + // Test with 3 channels (RGB-like) + // All channels have constant error 0.1 → MSE = 0.01 → PSNR = 20 dB + let device = Default::default(); + let outputs = Tensor::<4>::from_data( + TensorData::from([[ + [[0.1_f32, 0.1], [0.1, 0.1]], + [[0.1_f32, 0.1], [0.1, 0.1]], + [[0.1_f32, 0.1], [0.1, 0.1]], + ]]), + &device, + ); + let targets = Tensor::<4>::zeros([1, 3, 2, 2], &device); + + let config = PsnrMetricConfig::new(1.0); + let mut metric = PsnrMetric::new(config); + let input = PsnrInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let psnr = metric.value().current(); + let expected_psnr = 20.0; + assert!( + (psnr - expected_psnr).abs() < 0.01, + "Expected PSNR ~{} dB, got {} dB", + expected_psnr, + psnr + ); + } + + #[test] + fn test_psnr_running_average() { + // Test running average across multiple updates + let device = Default::default(); + let config = PsnrMetricConfig::new(1.0); + let mut metric = PsnrMetric::new(config); + + // First update: error 0.1 → MSE = 0.01 → PSNR = 20 dB + let outputs1 = + Tensor::<4>::from_data(TensorData::from([[[[0.1_f32, 0.1], [0.1, 0.1]]]]), &device); + let targets1 = Tensor::<4>::zeros([1, 1, 2, 2], &device); + let input1 = PsnrInput::new(outputs1, targets1); + let _entry = metric.update(&input1, &MetricMetadata::fake()); + + let psnr1 = metric.value().current(); + let expected_psnr1 = 20.0; + assert!( + (psnr1 - expected_psnr1).abs() < 0.01, + "First update PSNR should be ~{} dB, got {} dB", + expected_psnr1, + psnr1 + ); + + // Second update: error 0.01 → MSE = 0.0001 → PSNR = 40 dB + let outputs2 = Tensor::<4>::from_data( + TensorData::from([[[[0.01_f32, 0.01], [0.01, 0.01]]]]), + &device, + ); + let targets2 = Tensor::<4>::zeros([1, 1, 2, 2], &device); + let input2 = PsnrInput::new(outputs2, targets2); + let _entry = metric.update(&input2, &MetricMetadata::fake()); + + // Running average: (20 + 40) / 2 = 30 dB + let running_avg_psnr = metric.running_value().current(); + let expected_running_avg_psnr = 30.0; + assert!( + (running_avg_psnr - expected_running_avg_psnr).abs() < 0.01, + "Running average should be ~{} dB, got {} dB", + expected_running_avg_psnr, + running_avg_psnr + ); + } + + #[test] + fn test_psnr_clear() { + // Error 0.1 → MSE = 0.01 → PSNR = 20 dB + let device = Default::default(); + let config = PsnrMetricConfig::new(1.0); + let mut metric = PsnrMetric::new(config); + + let outputs = + Tensor::<4>::from_data(TensorData::from([[[[0.1_f32, 0.1], [0.1, 0.1]]]]), &device); + let targets = Tensor::<4>::zeros([1, 1, 2, 2], &device); + let input = PsnrInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let psnr = metric.value().current(); + let expected_psnr = 20.0; + assert!( + (psnr - expected_psnr).abs() < 0.01, + "Expected PSNR ~{} dB, got {} dB", + expected_psnr, + psnr + ); + + // Clear and verify reset + metric.clear(); + let psnr = metric.running_value().current(); + assert!(psnr.is_nan(), "Expected NaN after clear, got {} dB", psnr) + } + + #[test] + fn test_psnr_custom_name() { + let config = PsnrMetricConfig::new(1.0); + let metric = PsnrMetric::new(config).with_name("CustomPSNR"); + + assert_eq!(metric.name().to_string(), "CustomPSNR"); + } + + #[test] + fn test_psnr_custom_epsilon() { + let device = Default::default(); + // With a larger epsilon, perfect reconstruction gives lower PSNR + let config = PsnrMetricConfig::new(1.0).with_epsilon(0.01); + let mut metric = PsnrMetric::new(config); + + let outputs = + Tensor::<4>::from_data(TensorData::from([[[[0.5_f32, 0.5], [0.5, 0.5]]]]), &device); + let targets = outputs.clone(); + let input = PsnrInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + // With epsilon = 0.01, PSNR = 10 * log10(1.0 / 0.01) = 20 dB + let psnr = metric.value().current(); + let expected_psnr = 20.0; + assert!( + (psnr - expected_psnr).abs() < 0.01, + "Expected PSNR ~{} dB with epsilon=0.01, got {}", + expected_psnr, + psnr + ); + } + + #[test] + fn test_psnr_negative_errors() { + // Test that negative differences (target > output) work correctly + let device = Default::default(); + let outputs = + Tensor::<4>::from_data(TensorData::from([[[[0.0_f32, 0.0], [0.0, 0.0]]]]), &device); + let targets = + Tensor::<4>::from_data(TensorData::from([[[[0.1_f32, 0.1], [0.1, 0.1]]]]), &device); + + let config = PsnrMetricConfig::new(1.0); + let mut metric = PsnrMetric::new(config); + let input = PsnrInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + // Same MSE as positive errors (0.01), so PSNR = 20 dB + let psnr = metric.value().current(); + let expected_psnr = 20.0; + assert!( + (psnr - expected_psnr).abs() < 0.01, + "Expected PSNR ~{} dB, got {}", + expected_psnr, + psnr + ); + } + + #[test] + fn test_psnr_large_batch() { + // Test with a larger batch to verify batch dimension handling + let device = Default::default(); + let batch_size = 8; + + // All images have constant error 0.1 → MSE = 0.01 → PSNR = 20 dB + let outputs = Tensor::<4>::full([batch_size, 3, 4, 4], 0.1, &device); + let targets = Tensor::<4>::zeros([batch_size, 3, 4, 4], &device); + + let config = PsnrMetricConfig::new(1.0); + let mut metric = PsnrMetric::new(config); + let input = PsnrInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let psnr = metric.value().current(); + let expected_psnr = 20.0; + assert!( + (psnr - expected_psnr).abs() < 0.01, + "Expected PSNR ~{} dB, got {}", + expected_psnr, + psnr + ); + } + + #[test] + fn test_psnr_attributes() { + let config = PsnrMetricConfig::new(1.0); + let metric = PsnrMetric::new(config); + let attrs = metric.attributes(); + + match attrs { + MetricAttributes::Numeric(numeric_attrs) => { + assert_eq!(numeric_attrs.unit, Some("dB".to_string())); + assert!(numeric_attrs.higher_is_better); + } + _ => panic!("Expected numeric attributes"), + } + } + + #[test] + #[should_panic(expected = "Shape mismatch")] + fn test_psnr_shape_mismatch() { + let device = Default::default(); + let outputs = Tensor::<4>::zeros([1, 1, 2, 2], &device); + let targets = Tensor::<4>::zeros([1, 1, 3, 3], &device); + + let _ = PsnrInput::new(outputs, targets); + } + + #[test] + #[should_panic(expected = "max_pixel_val must be positive")] + fn test_psnr_negative_max_pixel_val() { + let _ = PsnrMetricConfig::new(-1.0); + } + + #[test] + #[should_panic(expected = "max_pixel_val must be positive")] + fn test_psnr_zero_max_pixel_val() { + let _ = PsnrMetricConfig::new(0.0); + } + + #[test] + #[should_panic(expected = "epsilon must be positive")] + fn test_psnr_negative_epsilon() { + let _ = PsnrMetricConfig::new(1.0).with_epsilon(-1e-10); + } + + #[test] + #[should_panic(expected = "epsilon must be positive")] + fn test_psnr_zero_epsilon() { + let _ = PsnrMetricConfig::new(1.0).with_epsilon(0.0); + } +} diff --git a/crates/burn-train/src/metric/vision/ssim.rs b/crates/burn-train/src/metric/vision/ssim.rs new file mode 100644 index 0000000..0b38fbb --- /dev/null +++ b/crates/burn-train/src/metric/vision/ssim.rs @@ -0,0 +1,872 @@ +use crate::metric::{ + Metric, MetricAttributes, MetricMetadata, MetricName, Numeric, NumericAttributes, NumericEntry, + SerializedEntry, + state::{FormatOptions, NumericMetricState}, +}; +use burn_core::{ + prelude::{Device, Tensor}, + tensor::{ElementConversion, module::conv2d, ops::ConvOptions}, +}; + +/// Input type for the [SsimMetric]. +/// +/// Both tensors must have shape `[N, C, H, W]`: +/// - `N`: Batch size +/// - `C`: Number of channels (1 for grayscale, 3 for RGB, etc.) +/// - `H`: Height +/// - `W`: Width +pub struct SsimInput { + /// Model output (predictions/reconstructions) images with shape [N, C, H, W]. + outputs: Tensor<4>, + /// Ground truth images with shape [N, C, H, W]. + targets: Tensor<4>, +} + +impl SsimInput { + /// Creates a new SsimInput with the given outputs and targets. + /// + /// Inputs are expected to have the dimensions `[N, C, H, W]` + /// where `N` is the batch size, `C` is the number of channels, + /// `H` is the height of the image, and `W` is the width of the image. + /// + /// # Arguments + /// - `outputs`: The model output images with shape [N, C, H, W]. + /// - `targets`: The ground truth images with shape [N, C, H, W]. + /// + /// # Returns + /// A new instance of `SsimInput`. + /// + /// # Panics + /// - If `outputs` and `targets` do not have the same shape. + pub fn new(outputs: Tensor<4>, targets: Tensor<4>) -> Self { + assert!( + outputs.dims() == targets.dims(), + "Shape mismatch: outputs {:?}, targets {:?}", + outputs.dims(), + targets.dims() + ); + Self { outputs, targets } + } +} + +/// Configuration for the [SsimMetric]. +#[derive(Debug, Clone, Copy)] +pub struct SsimMetricConfig { + /// The range of the pixel values in images which can be computed as following: + /// `let pixel_range = max_pixel_val - min_pixel_val;` + /// where `max_pixel_val` is the maximum possible pixel value and `min_pixel_val` + /// is the minimum possible pixel value. + /// + /// - For normalized images in range [0, 1], it should be set to `1.0 - 0.0 = 1.0` + /// - For normalized images in range [-1, 1], it should be set to `1.0 - (-1.0) = 2.0` + /// - For 8-bit images in range [0, 255], it should be set to `255.0 - 0.0 = 255.0` + pub pixel_range: f32, + /// A parameter of SSIM used to stabilize the luminance comparison. + /// Default is 0.01. + pub k1: f32, + /// A parameter of SSIM used to stabilize the contrast comparison. + /// Default is 0.03. + pub k2: f32, + /// The SSIM metric involves applying convolution to the input tensors using a Gaussian kernel. + /// This is the kernel size of the Gaussian kernel. Default is 11. + pub kernel_size: usize, + /// The SSIM metric involves applying convolution to the input tensors using a Gaussian kernel. + /// This is the standard deviation of the Gaussian kernel. Default is 1.5. + pub sigma: f32, +} + +impl SsimMetricConfig { + /// Creates a configuration with the specified data range and default parameters. + /// + /// # Default parameters + /// - k1: 0.01 + /// - k2: 0.03 + /// - kernel_size: 11 + /// - sigma: 1.5 + /// + /// # Panics + /// - If `pixel_range` is not positive. + /// + /// # Example + /// ```ignore + /// // Normalized images [0, 1] + /// let config1 = SsimMetricConfig::new(1.0); + /// + /// // 8-bit images [0, 255] + /// let config2 = SsimMetricConfig::new(255.0); + /// + /// // Also set custom values for k1 and k2 + /// let config3 = SsimMetricConfig::new(1.0).with_k1_k2(0.015, 0.025); + /// + /// // Also set a custom value for window size + /// config3.with_kernel_size(13); + /// ``` + pub fn new(pixel_range: f32) -> Self { + assert!(pixel_range > 0.0, "pixel_range must be positive"); + Self { + pixel_range: pixel_range, + k1: 0.01, + k2: 0.03, + kernel_size: 11, + sigma: 1.5, + } + } + + /// Sets a custom value for the k1 and k2 parameters of SSIM which are + /// used for numerical stability. + /// + /// # Default values + /// - k1: 0.01 + /// - k2: 0.03 + /// + /// # Panics + /// - If `k1` or `k2` is not positive. + pub fn with_k1_k2(mut self, k1: f32, k2: f32) -> Self { + assert!(k1 > 0.0, "k1 must be positive"); + assert!(k2 > 0.0, "k2 must be positive"); + self.k1 = k1; + self.k2 = k2; + self + } + + /// Sets a custom window size for the Gaussian kernel used in SSIM. The + /// window size must be a positive odd number. + /// + /// # Default value + /// - kernel_size: 11 + /// + /// # Panics + /// - If `kernel_size` is not a positive odd number. + pub fn with_kernel_size(mut self, kernel_size: usize) -> Self { + assert!( + kernel_size > 0 && kernel_size % 2 == 1, + "kernel_size must be positive and an odd number" + ); + self.kernel_size = kernel_size; + self + } + + /// Sets a custom sigma (standard deviation) for the Gaussian kernel used in SSIM. + /// + /// # Default value + /// - sigma: 1.5 + /// + /// # Panics + /// - If `sigma` is not positive. + pub fn with_sigma(mut self, sigma: f32) -> Self { + assert!(sigma > 0.0, "sigma must be positive"); + self.sigma = sigma; + self + } +} + +/// The SSIM (structural similarity index measure) metric for image quality assessment. +/// +/// SSIM measures the perceived quality of images by comparing luminance, +/// contrast, and structure. Values range from -1 to 1, where 1 indicates +/// perfect structural similarity. +/// +/// # Formula +/// ```text +/// SSIM(x, y) = (2μxμy + C1)(2σxy + C2) / (μx² + μy² + C1)(σx² + σy² + C2) +/// ``` +/// +/// # Note +/// - This implementation uses separable Gaussian convolution for efficiency. Instead of a +/// single 2D convolution with a K by K kernel, it applies two 1D convolutions (horizontal +/// then vertical). This reduces the computational complexity from O(K^2) to O(2K) per pixel. +/// - SSIM is computed for each image first, and then it is averaged across all the images in the batch. +#[derive(Clone)] +pub struct SsimMetric { + name: MetricName, + /// Internal state for numeric metric aggregation. + state: NumericMetricState, + /// Configuration for the metric. + config: SsimMetricConfig, +} + +impl SsimMetric { + /// Creates a new SSIM metric with the given configuration. + /// + /// # Note + /// The metric name format is "SSIM (dr={}, w={}, σ={})" + /// where dr is the data range, w is the window size, sigma is the + /// standard deviation. For example, the metric name might be + /// "SSIM (dr=1.0, w=11, σ=1.5)". + /// + /// # Example + /// ```ignore + /// let ssim_config = SsimMetricConfig::new(1.0); + /// let ssim_metric = SsimMetric::new(ssim_config); + /// ``` + pub fn new(config: SsimMetricConfig) -> Self { + Self { + name: MetricName::new(format!( + "SSIM (dr={}, w={}, σ={})", + config.pixel_range, config.kernel_size, config.sigma, + )), + state: NumericMetricState::default(), + config, + } + } + + /// Overrides the default metric name which is "SSIM". + pub fn with_name(mut self, name: &str) -> Self { + self.name = MetricName::new(name.to_string()); + self + } + + /// Creates a 1D Gaussian kernel as a tensor. + /// + /// Returns a normalized kernel where all values sum to 1. + /// The returned kernel will be reshaped by the `gaussian_conv_separable` + /// associated function later. + fn create_1d_gaussian_kernel(&self) -> Vec { + let size = self.config.kernel_size; + let sigma = self.config.sigma; + let center = (size / 2) as f32; + + let mut kernel = vec![0.0f32; size]; + let mut sum = 0.0f32; + + for (i, v) in kernel.iter_mut().enumerate() { + let x = i as f32 - center; + let value = (-(x * x) / (2.0 * sigma * sigma)).exp(); + *v = value; + sum += value; + } + + // Normalize so values sum to 1 + for v in kernel.iter_mut() { + *v /= sum; + } + + kernel + } + + /// Applies separable convolution using two 1D Gaussian kernels. + /// + /// # Arguments + /// - `inputs`: Tensor of shape [N, C, H, W] + /// - `kernel_1d`: The 1D Gaussian kernel values + /// - `channels`: Number of channels for depthwise convolution. + fn gaussian_conv_separable( + &self, + input: Tensor<4>, + kernel_1d: &[f32], + channels: usize, + device: &Device, + ) -> Tensor<4> { + let size = self.config.kernel_size; + let padding = size / 2; + + // Create horizontal kernel: shape [C, 1, 1, K] + let horizontal_kernel = Tensor::<1>::from_floats(kernel_1d, device) + .reshape([1, 1, 1, size]) // [1, 1, 1, K] + .repeat_dim(0, channels); // [C, 1, 1, K] + + let vertical_kernel = Tensor::<1>::from_floats(kernel_1d, device) + .reshape([1, 1, size, 1]) // [1, 1, K, 1] + .repeat_dim(0, channels); // [C, 1, K, 1] + + // Apply horizontal convolution + let horizontal_conv_options = ConvOptions::new([1, 1], [0, padding], [1, 1], channels); + let input_after_horizontal_conv = + conv2d(input, horizontal_kernel, None, horizontal_conv_options); + + // Apply vertical convolution + let vertical_conv_options = ConvOptions::new([1, 1], [padding, 0], [1, 1], channels); + conv2d( + input_after_horizontal_conv, + vertical_kernel, + None, + vertical_conv_options, + ) + } +} + +impl Metric for SsimMetric { + type Input = SsimInput; + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn update(&mut self, item: &Self::Input, _metadata: &MetricMetadata) -> SerializedEntry { + let dims = item.outputs.dims(); + let batch_size = dims[0]; + let channels = dims[1]; + let device = item.outputs.device(); + + let img_height = dims[2]; + let img_width = dims[3]; + assert!( + img_height >= self.config.kernel_size && img_width >= self.config.kernel_size, + "Image dimensions (H={}, W={}) must be >= kernel_size ({})", + img_height, + img_width, + self.config.kernel_size + ); + + // Constants in SSIM formula used for numerical stability + let c1 = (self.config.k1 * self.config.pixel_range).powi(2); + let c2 = (self.config.k2 * self.config.pixel_range).powi(2); + + // Create 1D Gaussian kernel to apply separable convolutions twice (horizontally and vertically) + let kernel_1d = self.create_1d_gaussian_kernel(); + + // Compute mu_x and mu_y, their product and squares + let x = item.outputs.clone(); + let y = item.targets.clone(); + let mu_x = self.gaussian_conv_separable(x.clone(), &kernel_1d, channels, &device); + let mu_y = self.gaussian_conv_separable(y.clone(), &kernel_1d, channels, &device); + let mu_x_mu_y = mu_x.clone() * mu_y.clone(); + let square_of_mu_x = mu_x.clone() * mu_x.clone(); + let square_of_mu_y = mu_y.clone() * mu_y.clone(); + + // Compute var_x, var_y (which are the same as (sigma_x)^2 and (sigma_y)^2): + // Var(X) = E[X^2] - E[X]^2 + // var_x = mu_of_x_squared - (mu_x * mu_x) + let mu_of_x_squared = + self.gaussian_conv_separable(x.clone() * x.clone(), &kernel_1d, channels, &device); + let mu_of_y_squared = + self.gaussian_conv_separable(y.clone() * y.clone(), &kernel_1d, channels, &device); + let var_x = (mu_of_x_squared - square_of_mu_x.clone()).clamp_min(0.0); + let var_y = (mu_of_y_squared - square_of_mu_y.clone()).clamp_min(0.0); + + // Compute the sample covariance of x and y: sigma_xy + // Cov(X, Y) = E[XY] - E[X]E[Y] + // sigma_xy = mu_xy - (mu_x * mu_y) + let mu_xy = self.gaussian_conv_separable(x * y, &kernel_1d, channels, &device); + let sigma_xy = mu_xy - mu_x_mu_y.clone(); + + // Compute SSIM: + // SSIM(x, y) = (2μxμy + C1)(2σxy + C2) / (μx² + μy² + C1)(σx² + σy² + C2) + let numerator = (mu_x_mu_y.mul_scalar(2.0_f32) + c1) * (sigma_xy.mul_scalar(2.0_f32) + c2); + let denominator = (square_of_mu_x + square_of_mu_y + c1) * (var_x + var_y + c2); + let ssim_tensor = numerator / denominator; + + // Average SSIM across all dimensions to get a single scalar value + let ssim_per_image = ssim_tensor.mean_dims(&[1, 2, 3]); + let avg_ssim = ssim_per_image.mean().into_scalar::(); + + self.state.update( + avg_ssim, + batch_size, + FormatOptions::new(self.name()).precision(4), + ) + } + + /// Clears the metric state. + fn clear(&mut self) { + self.state.reset(); + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: None, + higher_is_better: true, + ..Default::default() + } + .into() + } +} + +impl Numeric for SsimMetric { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} + +#[cfg(test)] +#[allow(clippy::manual_range_contains)] +mod tests { + use super::*; + use crate::metric::Numeric; + use burn_core::tensor::{Distribution, Shape, TensorData}; + + fn test_config() -> SsimMetricConfig { + SsimMetricConfig::new(1.0) + .with_kernel_size(3) + .with_sigma(1.0) + } + + #[test] + fn test_ssim_perfect_similarity() { + // When outputs exactly match targets, SSIM should be 1.0 + let device = Default::default(); + let outputs = Tensor::<4>::from_data( + TensorData::from([[[ + [0.1_f32, 0.2, 0.3, 0.4], + [0.5, 0.6, 0.7, 0.8], + [0.2, 0.3, 0.4, 0.5], + [0.6, 0.7, 0.8, 0.9], + ]]]), + &device, + ); + let targets = outputs.clone(); + + let mut metric = SsimMetric::new(test_config()); + let input = SsimInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let ssim = metric.value().current(); + assert!( + (ssim - 1.0).abs() < 0.001, + "SSIM for identical images should be 1.0, got {}", + ssim + ); + } + + #[test] + fn test_ssim_completely_different() { + // Constant black vs constant white + // With constant images: SSIM = (2*mu_x*mu_y + C1) / (mu_x^2 + mu_y^2 + C1) + // For x=0, y=1 with C1=(0.01)^2=0.0001: SSIM ≈ 0.0001 / (1 + 0.00001) = 0.00009999 + let device = Default::default(); + let outputs = Tensor::<4>::zeros([1, 1, 4, 4], &device); + let targets = Tensor::<4>::ones([1, 1, 4, 4], &device); + + let mut metric = SsimMetric::new(test_config()); + let input = SsimInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let ssim = metric.value().current(); + assert!( + ssim < 0.0001, + "SSIM for black vs white images should be very low, got {}", + ssim + ); + } + + #[test] + fn test_ssim_similar_images() { + // Small perturbation should give high SSIM + let device = Default::default(); + let outputs = Tensor::<4>::full([1, 1, 4, 4], 0.5, &device); + let targets = Tensor::<4>::full([1, 1, 4, 4], 0.51, &device); + + let mut metric = SsimMetric::new(test_config()); + let input = SsimInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let ssim = metric.value().current(); + assert!( + ssim > 0.99, + "SSIM for very similar images should be close to 1.0, got {}", + ssim + ); + } + + #[test] + fn test_ssim_batch_averaging() { + // Batch of 2 images: + // Image 1: identical (SSIM = 1.0) + // Image 2: black vs white (SSIM ≈ 0) + let device = Default::default(); + let outputs = Tensor::<4>::from_data( + TensorData::from([ + [[ + [0.5_f32, 0.5, 0.5, 0.5], + [0.5, 0.5, 0.5, 0.5], + [0.5, 0.5, 0.5, 0.5], + [0.5, 0.5, 0.5, 0.5], + ]], + [[ + [0.0_f32, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + ]], + ]), + &device, + ); + let targets = Tensor::<4>::from_data( + TensorData::from([ + [[ + [0.5_f32, 0.5, 0.5, 0.5], + [0.5, 0.5, 0.5, 0.5], + [0.5, 0.5, 0.5, 0.5], + [0.5, 0.5, 0.5, 0.5], + ]], + [[ + [1.0_f32, 1.0, 1.0, 1.0], + [1.0, 1.0, 1.0, 1.0], + [1.0, 1.0, 1.0, 1.0], + [1.0, 1.0, 1.0, 1.0], + ]], + ]), + &device, + ); + + let mut metric = SsimMetric::new(test_config()); + let input = SsimInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let ssim = metric.value().current(); + // Average of ~1.0 and ~0.0 should be around 0.5 + assert!( + ssim > 0.49 && ssim < 0.51, + "Average SSIM should be around 0.5, got {}", + ssim + ); + } + + #[test] + fn test_ssim_multichannel() { + // Test with 3 channels (e.g., RGB) + let device = Default::default(); + let outputs = Tensor::<4>::from_data( + TensorData::from([[ + [ + [0.5_f32, 0.6, 0.7, 0.8], + [0.4, 0.5, 0.6, 0.7], + [0.3, 0.4, 0.5, 0.6], + [0.2, 0.3, 0.4, 0.5], + ], + [ + [0.3_f32, 0.4, 0.5, 0.6], + [0.2, 0.3, 0.4, 0.5], + [0.1, 0.2, 0.3, 0.4], + [0.0, 0.1, 0.2, 0.3], + ], + [ + [0.7_f32, 0.8, 0.9, 1.0], + [0.6, 0.7, 0.8, 0.9], + [0.5, 0.6, 0.7, 0.8], + [0.4, 0.5, 0.6, 0.7], + ], + ]]), + &device, + ); + let targets = outputs.clone(); + + let mut metric = SsimMetric::new(test_config()); + let input = SsimInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let ssim = metric.value().current(); + assert!( + (ssim - 1.0).abs() < 0.001, + "SSIM for identical RGB images should be 1.0, got {}", + ssim + ); + } + + #[test] + fn test_ssim_symmetry() { + // SSIM(x, y) should equal SSIM(y, x) + // Symmetry is one of the mathematical properties of SSIM + let device = Default::default(); + let img1 = Tensor::<4>::from_data( + TensorData::from([[[ + [0.1_f32, 0.2, 0.3, 0.4], + [0.5, 0.6, 0.7, 0.8], + [0.2, 0.3, 0.4, 0.5], + [0.6, 0.7, 0.8, 0.9], + ]]]), + &device, + ); + let img2 = Tensor::<4>::from_data( + TensorData::from([[[ + [0.2_f32, 0.3, 0.4, 0.5], + [0.6, 0.7, 0.8, 0.9], + [0.3, 0.4, 0.5, 0.6], + [0.7, 0.8, 0.9, 1.0], + ]]]), + &device, + ); + + let config = test_config(); + + let mut metric1 = SsimMetric::new(config); + let input1 = SsimInput::new(img1.clone(), img2.clone()); + let _entry = metric1.update(&input1, &MetricMetadata::fake()); + let ssim1 = metric1.value().current(); + + let mut metric2 = SsimMetric::new(config); + let input2 = SsimInput::new(img2, img1); + let _entry = metric2.update(&input2, &MetricMetadata::fake()); + let ssim2 = metric2.value().current(); + + assert!( + (ssim1 - ssim2).abs() < 0.001, + "SSIM should be symmetric: SSIM(x,y)={} vs SSIM(y,x)={}", + ssim1, + ssim2 + ); + } + + #[test] + fn test_ssim_range() { + // SSIM values should be in [-1, 1] range + let device = Default::default(); + let shape = Shape::new([1, 1, 11, 11]); + let distribution = Distribution::Uniform(0.0, 1.0); + let outputs = Tensor::<4>::random(shape.clone(), distribution, &device); + let targets = Tensor::<4>::random(shape, distribution, &device); + + let mut metric = SsimMetric::new(test_config()); + let input = SsimInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let ssim = metric.value().current(); + assert!( + ssim >= -1.0 && ssim <= 1.0, + "SSIM should be in range [-1, 1], got {}", + ssim + ); + } + + #[test] + fn test_ssim_running_average() { + let device = Default::default(); + let mut metric = SsimMetric::new(test_config()); + + // First update: identical images (SSIM = 1.0) + let outputs1 = Tensor::<4>::from_data( + TensorData::from([[[ + [0.5_f32, 0.6, 0.7, 0.8], + [0.4, 0.5, 0.6, 0.7], + [0.3, 0.4, 0.5, 0.6], + [0.2, 0.3, 0.4, 0.5], + ]]]), + &device, + ); + let targets1 = outputs1.clone(); + let input1 = SsimInput::new(outputs1, targets1); + let _entry = metric.update(&input1, &MetricMetadata::fake()); + + let ssim1 = metric.value().current(); + assert!( + (ssim1 - 1.0).abs() < 0.001, + "First update SSIM should be ~1.0, got {}", + ssim1 + ); + + // Second update: very different images (SSIM close to 0) + let outputs2 = Tensor::<4>::zeros([1, 1, 4, 4], &device); + let targets2 = Tensor::<4>::ones([1, 1, 4, 4], &device); + let input2 = SsimInput::new(outputs2, targets2); + let _entry = metric.update(&input2, &MetricMetadata::fake()); + + // Running average should be around 0.5 + let running_avg = metric.running_value().current(); + assert!( + running_avg > 0.49 && running_avg < 0.51, + "Running average should be around 0.5, got {}", + running_avg + ); + } + + #[test] + fn test_ssim_clear() { + let device = Default::default(); + let mut metric = SsimMetric::new(test_config()); + + let outputs = Tensor::<4>::from_data( + TensorData::from([[[ + [0.5_f32, 0.6, 0.7, 0.8], + [0.4, 0.5, 0.6, 0.7], + [0.3, 0.4, 0.5, 0.6], + [0.2, 0.3, 0.4, 0.5], + ]]]), + &device, + ); + let targets = outputs.clone(); + let input = SsimInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let ssim = metric.value().current(); + assert!( + (ssim - 1.0).abs() < 0.001, + "Expected SSIM ~1.0, got {}", + ssim + ); + + // Clear and verify reset + metric.clear(); + let ssim = metric.running_value().current(); + assert!(ssim.is_nan(), "Expected NaN after clear, got {}", ssim); + } + + #[test] + fn test_ssim_custom_name() { + let config = SsimMetricConfig::new(1.0); + let metric = SsimMetric::new(config).with_name("CustomSSIM"); + assert_eq!(metric.name().to_string(), "CustomSSIM"); + + let metric = SsimMetric::new(test_config()); + assert_eq!(metric.name().to_string(), "SSIM (dr=1, w=3, σ=1)"); + + let config = SsimMetricConfig::new(255.0); + let metric = SsimMetric::new(config); + assert_eq!(metric.name().to_string(), "SSIM (dr=255, w=11, σ=1.5)"); + } + + #[test] + fn test_ssim_pixel_range_255() { + // Test with 8-bit image range [0, 255] + let device = Default::default(); + let shape = Shape::new([1, 1, 10, 10]); + let distribution = Distribution::Uniform(0.0, 255.0); + let outputs = Tensor::<4>::random(shape.clone(), distribution, &device); + let targets = outputs.clone(); + + let config = SsimMetricConfig::new(255.0).with_kernel_size(3); + let mut metric = SsimMetric::new(config); + let input = SsimInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let ssim = metric.value().current(); + assert!( + (ssim - 1.0).abs() < 0.001, + "SSIM for identical 8-bit images should be 1.0, got {}", + ssim + ); + } + + #[test] + fn test_ssim_large_batch() { + let device = Default::default(); + let shape = Shape::new([20, 3, 30, 30]); + let distribution = Distribution::Uniform(0.0, 1.0); + let outputs = Tensor::<4>::random(shape, distribution, &device); + let targets = outputs.clone(); + + let mut metric = SsimMetric::new(test_config()); + let input = SsimInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let ssim = metric.value().current(); + assert!( + (ssim - 1.0).abs() < 0.001, + "SSIM for identical batch should be 1.0, got {}", + ssim + ); + } + + #[test] + fn test_ssim_default_kernel_size() { + // Test with default kernel_size=11, need images >= 11x11 + let device = Default::default(); + let shape = Shape::new([1, 1, 1080, 1920]); + let distribution = Distribution::Uniform(0.0, 1.0); + let outputs = Tensor::<4>::random(shape, distribution, &device); + let targets = outputs.clone(); + + let config = SsimMetricConfig::new(1.0); // default kernel_size=11 + let mut metric = SsimMetric::new(config); + let input = SsimInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + + let ssim = metric.value().current(); + assert!( + (ssim - 1.0).abs() < 0.001, + "SSIM with default window size should work and SSIM should be ~0.0, got {}", + ssim + ); + } + + #[test] + fn test_ssim_attributes() { + let config = SsimMetricConfig::new(1.0); + let metric = SsimMetric::new(config); + let attrs = metric.attributes(); + + match attrs { + MetricAttributes::Numeric(numeric_attrs) => { + assert_eq!(numeric_attrs.unit, None); + assert!(numeric_attrs.higher_is_better); + } + _ => panic!("Expected numeric attributes"), + } + } + + #[test] + #[should_panic(expected = "Shape mismatch")] + fn test_ssim_shape_mismatch() { + let device = Default::default(); + let outputs = Tensor::<4>::zeros([1, 1, 4, 4], &device); + let targets = Tensor::<4>::zeros([1, 1, 5, 5], &device); + + let _ = SsimInput::new(outputs, targets); + } + + #[test] + #[should_panic(expected = "Image dimensions (H=4, W=4) must be >= kernel_size (11)")] + fn test_ssim_image_too_small() { + let device = Default::default(); + let outputs = Tensor::<4>::zeros([1, 1, 4, 4], &device); + let targets = outputs.clone(); + + // Default kernel_size=11, but image is only 4x4 + let config = SsimMetricConfig::new(1.0); + let mut metric = SsimMetric::new(config); + let input = SsimInput::new(outputs, targets); + let _entry = metric.update(&input, &MetricMetadata::fake()); + } + + #[test] + fn test_ssim_valid_k1_k2() { + let config = SsimMetricConfig::new(1.0).with_k1_k2(0.015, 0.035); + assert!( + config.k1 == 0.015 && config.k2 == 0.035, + "Expected k1=0.015 and k2=0.035, got k1={} and k2={}", + config.k1, + config.k2 + ); + } + + #[test] + #[should_panic(expected = "pixel_range must be positive")] + fn test_ssim_negative_pixel_range() { + let _ = SsimMetricConfig::new(-1.0); + } + + #[test] + #[should_panic(expected = "pixel_range must be positive")] + fn test_ssim_zero_pixel_range() { + let _ = SsimMetricConfig::new(0.0); + } + + #[test] + #[should_panic(expected = "k1 must be positive")] + fn test_ssim_negative_k1() { + let _ = SsimMetricConfig::new(1.0).with_k1_k2(-0.01, 0.03); + } + + #[test] + #[should_panic(expected = "k2 must be positive")] + fn test_ssim_negative_k2() { + let _ = SsimMetricConfig::new(1.0).with_k1_k2(0.01, -0.03); + } + + #[test] + #[should_panic(expected = "kernel_size must be positive and an odd number")] + fn test_ssim_even_kernel_size() { + let _ = SsimMetricConfig::new(1.0).with_kernel_size(10); + } + + #[test] + #[should_panic(expected = "kernel_size must be positive and an odd number")] + fn test_ssim_zero_kernel_size() { + let _ = SsimMetricConfig::new(1.0).with_kernel_size(0); + } + + #[test] + #[should_panic(expected = "sigma must be positive")] + fn test_ssim_negative_sigma() { + let _ = SsimMetricConfig::new(1.0).with_sigma(-1.5); + } + + #[test] + #[should_panic(expected = "sigma must be positive")] + fn test_ssim_zero_sigma() { + let _ = SsimMetricConfig::new(1.0).with_sigma(0.0); + } +} diff --git a/crates/burn-train/src/metric/wer.rs b/crates/burn-train/src/metric/wer.rs new file mode 100644 index 0000000..79fc051 --- /dev/null +++ b/crates/burn-train/src/metric/wer.rs @@ -0,0 +1,226 @@ +use super::cer::edit_distance; +use super::state::{FormatOptions, NumericMetricState}; +use super::{MetricMetadata, SerializedEntry}; +use crate::metric::{ + Metric, MetricAttributes, MetricName, Numeric, NumericAttributes, NumericEntry, +}; +use burn_core::tensor::{Int, Tensor}; +use std::sync::Arc; + +// The edit_distance function remains the same as it calculates the Levenshtein distance +// between two sequences. The "units" within the sequences will now be treated as words. +/// The word error rate (WER) metric, similar to the CER, is defined as the edit distance (e.g. Levenshtein distance) between the predicted +/// and reference word sequences, divided by the total number of words in the reference. Here, the "units" within the sequences are words. +/// +#[derive(Clone)] +pub struct WordErrorRate { + name: MetricName, + state: NumericMetricState, + pad_token: Option, +} + +/// The [word error rate metric](WordErrorRate) input type. +#[derive(new)] +pub struct WerInput { + /// The predicted token sequences (as a 2-D tensor of token indices). + pub outputs: Tensor<2, Int>, + /// The target token sequences (as a 2-D tensor of token indices). + pub targets: Tensor<2, Int>, +} +impl Default for WordErrorRate { + fn default() -> Self { + Self::new() + } +} + +impl WordErrorRate { + /// Creates the metric. + pub fn new() -> Self { + Self { + name: Arc::new("WER".to_string()), + state: NumericMetricState::default(), + pad_token: None, + } + } + + /// Sets the pad token. + pub fn with_pad_token(mut self, index: usize) -> Self { + self.pad_token = Some(index); + self + } +} + +impl Metric for WordErrorRate { + type Input = WerInput; + + fn update(&mut self, input: &WerInput, _metadata: &MetricMetadata) -> SerializedEntry { + let outputs = input.outputs.clone(); + let targets = input.targets.clone(); + let [batch_size, seq_len] = targets.dims(); + + let outputs_data = outputs + .to_data() + .convert::() + .to_vec() + .expect("Failed to convert outputs to Vec"); + let targets_data = targets + .to_data() + .convert::() + .to_vec() + .expect("Failed to convert targets to Vec"); + + let pad_token = self.pad_token.map(|p| p as i32); + + let mut total_edit_distance = 0.0; + let mut total_target_length = 0.0; + + // Process each sequence in the batch + for i in 0..batch_size { + let start = i * seq_len; + let end = (i + 1) * seq_len; + let output_seq = &outputs_data[start..end]; + let target_seq = &targets_data[start..end]; + + // Handle padding and map elements to i32. + // These sequences now represent "words" (token IDs). + let (ed, target_len) = match pad_token { + Some(pad) => { + let output_seq_no_pad = output_seq + .iter() + .take_while(|&&x| x != pad) + .copied() + .collect::>(); + + let target_seq_no_pad = target_seq + .iter() + .take_while(|&&x| x != pad) + .copied() + .collect::>(); + + ( + edit_distance(&target_seq_no_pad, &output_seq_no_pad), + target_seq_no_pad.len(), + ) + } + None => (edit_distance(target_seq, output_seq), target_seq.len()), + }; + + total_edit_distance += ed as f64; + total_target_length += target_len as f64; + } + + // Compute current WER value as a percentage + let value = if total_target_length > 0.0 { + 100.0 * total_edit_distance / total_target_length + } else { + 0.0 + }; + + self.state.update( + value, + batch_size, + FormatOptions::new(self.name()).unit("%").precision(2), + ) + } + + fn name(&self) -> MetricName { + self.name.clone() + } + + fn clear(&mut self) { + self.state.reset(); + } + + fn attributes(&self) -> MetricAttributes { + NumericAttributes { + unit: Some("%".to_string()), + higher_is_better: false, + ..Default::default() + } + .into() + } +} + +impl Numeric for WordErrorRate { + fn value(&self) -> NumericEntry { + self.state.current_value() + } + + fn running_value(&self) -> NumericEntry { + self.state.running_value() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Perfect match => WER = 0 %. + #[test] + fn test_wer_without_padding() { + let device = Default::default(); + let mut metric = WordErrorRate::new(); + + // Batch size = 2, sequence length = 2 + let preds = Tensor::from_data([[1, 2], [3, 4]], &device); + let tgts = Tensor::from_data([[1, 2], [3, 4]], &device); + + metric.update(&WerInput::new(preds, tgts), &MetricMetadata::fake()); + + assert_eq!(0.0, metric.value().current()); + } + + /// Two word edits in four target words => 50 %. + #[test] + fn test_wer_without_padding_two_errors() { + let device = Default::default(); + let mut metric = WordErrorRate::new(); + + // One substitution in each sequence. + // Sequence 1: target [1, 3], pred [1, 2] -> 1 error (3 vs 2) + // Sequence 2: target [3, 4], pred [3, 5] -> 1 error (4 vs 5) + let preds = Tensor::from_data([[1, 2], [3, 5]], &device); + let tgts = Tensor::from_data([[1, 3], [3, 4]], &device); + + metric.update(&WerInput::new(preds, tgts), &MetricMetadata::fake()); + + // Total errors = 2, Total target words = 4. WER = (2/4) * 100 = 50 % + assert_eq!(50.0, metric.value().current()); + } + + /// Same scenario as above, but with right-padding (token 9) ignored. + #[test] + fn test_wer_with_padding() { + let device = Default::default(); + let pad = 9_i64; + let mut metric = WordErrorRate::new().with_pad_token(pad as usize); + + // Each row has three columns, last one is the pad token. + // Target sequences after removing pad: [1, 3] and [3, 4] (total length 4) + // Predicted sequences after removing pad: [1, 2] and [3, 5] + let preds = Tensor::from_data([[1, 2, pad], [3, 5, pad]], &device); + let tgts = Tensor::from_data([[1, 3, pad], [3, 4, pad]], &device); + + metric.update(&WerInput::new(preds, tgts), &MetricMetadata::fake()); + assert_eq!(50.0, metric.value().current()); + } + + /// `clear()` must reset the running statistics to NaN. + #[test] + fn test_clear_resets_state() { + let device = Default::default(); + let mut metric = WordErrorRate::new(); + + let preds = Tensor::from_data([[1, 2]], &device); + let tgts = Tensor::from_data([[1, 3]], &device); // one error + + metric.update( + &WerInput::new(preds.clone(), tgts.clone()), + &MetricMetadata::fake(), + ); + assert!(metric.value().current() > 0.0); + + metric.clear(); + assert!(metric.value().current().is_nan()); + } +} diff --git a/crates/burn-train/src/renderer/base.rs b/crates/burn-train/src/renderer/base.rs new file mode 100644 index 0000000..4294bbe --- /dev/null +++ b/crates/burn-train/src/renderer/base.rs @@ -0,0 +1,118 @@ +use std::sync::Arc; + +use crate::{ + LearnerSummary, + logger::{EvaluationProgressLogger, TrainingProgressLogger}, + metric::{MetricDefinition, MetricEntry, NumericEntry}, +}; + +/// Trait for rendering metrics. +pub trait MetricsRendererTraining: Send + Sync { + /// Updates the training metric state. + /// + /// # Arguments + /// + /// * `state` - The metric state. + fn update_train(&mut self, state: MetricState); + + /// Updates the validation metric state. + /// + /// # Arguments + /// + /// * `state` - The metric state. + fn update_valid(&mut self, state: MetricState); + + /// Callback method invoked when training ends, whether it + /// completed successfully or was interrupted. + /// + /// # Returns + /// + /// A result indicating whether the end-of-training actions were successful. + fn on_train_end( + &mut self, + summary: Option, + ) -> Result<(), Box> { + default_summary_action(summary); + Ok(()) + } +} + +/// A renderer that can be used for both training and evaluation. +pub trait MetricsRenderer: + MetricsRendererEvaluation + + MetricsRendererTraining + + TrainingProgressLogger + + EvaluationProgressLogger +{ + /// Keep the renderer from automatically closing, requiring manual action to close it. + fn manual_close(&mut self); + /// Register a new metric. + fn register_metric(&mut self, definition: MetricDefinition); +} + +#[derive(Clone)] +/// The name of an evaluation. +/// +/// This is going to group metrics together for easier analysis. +pub struct EvaluationName { + pub(crate) name: Arc, +} + +impl EvaluationName { + /// Creates a new evaluation name. + pub fn new(s: S) -> Self { + Self { + name: Arc::new(format!("{s}")), + } + } + + /// Returns the evaluation name. + pub fn as_str(&self) -> &str { + &self.name + } +} + +impl core::fmt::Display for EvaluationName { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_str(&self.name) + } +} + +/// Trait for rendering metrics. +pub trait MetricsRendererEvaluation: Send + Sync { + /// Updates the testing metric state. + /// + /// # Arguments + /// + /// * `state` - The metric state. + fn update_test(&mut self, name: EvaluationName, state: MetricState); + + /// Callback method invoked when testing ends, whether it + /// completed successfully or was interrupted. + /// + /// # Returns + /// + /// A result indicating whether the end-of-testing actions were successful. + fn on_test_end( + &mut self, + summary: Option, + ) -> Result<(), Box> { + default_summary_action(summary); + Ok(()) + } +} + +/// The state of a metric. +#[derive(Debug)] +pub enum MetricState { + /// A generic metric. + Generic(MetricEntry), + /// A numeric metric. + Numeric(MetricEntry, NumericEntry), +} + +fn default_summary_action(summary: Option) { + if let Some(summary) = summary { + println!("{summary}"); + } +} diff --git a/crates/burn-train/src/renderer/cli.rs b/crates/burn-train/src/renderer/cli.rs new file mode 100644 index 0000000..5d997fc --- /dev/null +++ b/crates/burn-train/src/renderer/cli.rs @@ -0,0 +1,112 @@ +use burn_core::data::dataloader::Progress; + +use crate::{ + logger::{EvaluationProgressLogger, ProgressSnapshot, TrainingProgressLogger}, + renderer::{MetricState, MetricsRenderer, MetricsRendererEvaluation, MetricsRendererTraining}, +}; + +/// A simple renderer for when the cli feature is not enabled. +pub struct CliMetricsRenderer { + training_progress: ProgressSnapshot, + eval_progress: ProgressSnapshot, +} + +#[allow(clippy::new_without_default)] +impl CliMetricsRenderer { + /// Create a new instance. + pub fn new() -> Self { + let init = Progress::new(0, 0, Some(String::new())); + Self { + training_progress: ProgressSnapshot::new(init.clone(), init.clone()), + eval_progress: ProgressSnapshot::new(init.clone(), init), + } + } +} + +impl MetricsRendererTraining for CliMetricsRenderer { + fn update_train(&mut self, _state: MetricState) {} + + fn update_valid(&mut self, _state: MetricState) {} +} + +impl TrainingProgressLogger for CliMetricsRenderer { + fn start(&mut self, total_epochs: usize, starting_epoch: usize, total_items: Option) { + self.training_progress.global = + Progress::new(starting_epoch, total_epochs, Some("epochs".to_string())); + if let Some(items) = total_items { + self.training_progress.split = Progress::new(0, items, Some("items".to_string())); + } + println!("Starting training for {total_epochs} epochs."); + } + + fn start_split(&mut self, split_name: &str, total_items: usize) { + self.training_progress.split = Progress::new(0, total_items, Some("items".to_string())); + println!("Starting split '{split_name}' with {total_items} items."); + } + + fn update_split(&mut self, items_processed: usize) { + let total = self.training_progress.split.items_total; + let unit = self.training_progress.split.unit.clone(); + self.training_progress.split = Progress::new(items_processed, total, unit); + + // For RL: global_progress.items_total == 0 means no epoch concept; mirror split. + if self.training_progress.global.items_total == 0 { + self.training_progress.global = self.training_progress.split.clone(); + } + println!("{:?}", self.training_progress); + } + + fn update_epoch(&mut self, epoch: usize) { + let total = self.training_progress.global.items_total; + let unit = self.training_progress.global.unit.clone(); + self.training_progress.global = Progress::new(epoch + 1, total, unit); + } + + fn end_split(&mut self) { + println!("Split ended."); + } + + fn end(&mut self) { + println!("Training ended."); + } + + fn log_event_training(&mut self, _event: String) {} +} + +impl EvaluationProgressLogger for CliMetricsRenderer { + fn start_global_progress(&mut self, total_tests: usize) { + self.eval_progress.global = Progress::new(0, total_tests, Some("tests".to_string())); + println!("Starting evaluation with {total_tests} test(s)."); + } + + fn start_test(&mut self, name: &str, total_items: usize) { + let current = self.eval_progress.global.items_processed + 1; + let total = self.eval_progress.global.items_total; + self.eval_progress.global = Progress::new(current, total, Some("tests".to_string())); + self.eval_progress.split = Progress::new(0, total_items, Some("items".to_string())); + println!("Starting test '{name}' with {total_items} items."); + } + + fn update_test_progress(&mut self, items_processed: usize) { + let total = self.eval_progress.split.items_total; + let unit = self.eval_progress.split.unit.clone(); + self.eval_progress.split = Progress::new(items_processed, total, unit); + println!("{:?}", self.eval_progress); + } + + fn end_test(&mut self) {} + + fn end_global_progress(&mut self) {} + + fn log_event_evaluation(&mut self, _event: String) {} +} + +impl MetricsRendererEvaluation for CliMetricsRenderer { + fn update_test(&mut self, _name: super::EvaluationName, _state: MetricState) {} +} + +impl MetricsRenderer for CliMetricsRenderer { + fn manual_close(&mut self) {} + + fn register_metric(&mut self, _definition: crate::metric::MetricDefinition) {} +} diff --git a/crates/burn-train/src/renderer/mod.rs b/crates/burn-train/src/renderer/mod.rs new file mode 100644 index 0000000..6e441fa --- /dev/null +++ b/crates/burn-train/src/renderer/mod.rs @@ -0,0 +1,34 @@ +#[cfg(feature = "tui")] +use std::io::IsTerminal; + +mod base; +pub use base::*; + +pub(crate) mod cli; + +pub use cli::*; + +/// The tui renderer +#[cfg(feature = "tui")] +pub mod tui; +use crate::Interrupter; + +/// Return the default metrics renderer. +/// +/// This can be either: +/// - `TuiMetricsRenderer`, when the `tui` feature is enabled and `stdout` is +/// a terminal, or +/// - `CliMetricsRenderer`, when the `tui` feature is not enabled, or `stdout` +/// is not a terminal. +#[allow(unused_variables)] +pub(crate) fn default_renderer( + interuptor: Interrupter, + checkpoint: Option, +) -> Box { + #[cfg(feature = "tui")] + if std::io::stdout().is_terminal() { + return Box::new(tui::TuiMetricsRendererWrapper::new(interuptor, checkpoint)); + } + + Box::new(CliMetricsRenderer::new()) +} diff --git a/crates/burn-train/src/renderer/tui/base.rs b/crates/burn-train/src/renderer/tui/base.rs new file mode 100644 index 0000000..12e40e3 --- /dev/null +++ b/crates/burn-train/src/renderer/tui/base.rs @@ -0,0 +1,106 @@ +use std::sync::Arc; + +use super::{ + ControlsView, NumericMetricView, ProgressBarView, StatusView, TerminalFrame, TextMetricView, +}; +use ratatui::{ + prelude::{Constraint, Direction, Layout, Rect}, + style::Color, +}; + +#[derive(new)] +pub(crate) struct MetricsView<'a> { + metric_numeric: NumericMetricView<'a>, + metric_text: TextMetricView<'a>, + progress: ProgressBarView, + controls: ControlsView, + status: StatusView, +} + +impl MetricsView<'_> { + pub(crate) fn render(self, frame: &mut TerminalFrame<'_>, size: Rect) { + let chunks = Layout::default() + .direction(Direction::Vertical) + .constraints([Constraint::Min(16), Constraint::Max(4)].as_ref()) + .split(size); + let size_other = chunks[0]; + let size_progress = chunks[1]; + + let chunks = Layout::default() + .direction(Direction::Horizontal) + .constraints([Constraint::Percentage(38), Constraint::Percentage(62)].as_ref()) + .split(size_other); + let size_other = chunks[0]; + let size_metric_numeric = chunks[1]; + + let chunks = Layout::default() + .direction(Direction::Vertical) + .constraints([Constraint::Max(5), Constraint::Min(6), Constraint::Max(6)].as_ref()) + .split(size_other); + let size_controls = chunks[0]; + let size_metric_text = chunks[1]; + let size_status = chunks[2]; + + self.metric_numeric.render(frame, size_metric_numeric); + self.metric_text.render(frame, size_metric_text); + self.controls.render(frame, size_controls); + self.progress.render(frame, size_progress); + self.status.render(frame, size_status); + } +} + +#[derive(Hash, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub(crate) enum TuiSplit { + Train, + Valid, + Test, +} + +#[derive(Hash, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub(crate) enum TuiGroup { + Default, + Named(Arc), +} + +#[derive(new, Hash, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub(crate) struct TuiTag { + pub(crate) split: TuiSplit, + pub(crate) group: TuiGroup, +} + +impl core::fmt::Display for TuiTag { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match &self.group { + TuiGroup::Default => f.write_fmt(format_args!("{}", self.split)), + TuiGroup::Named(group) => f.write_fmt(format_args!("{} - {}", self.split, group)), + } + } +} +impl core::fmt::Display for TuiGroup { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + TuiGroup::Default => f.write_str(""), + TuiGroup::Named(group) => f.write_fmt(format_args!("{group} ")), + } + } +} + +impl core::fmt::Display for TuiSplit { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + TuiSplit::Train => f.write_str("Train"), + TuiSplit::Valid => f.write_str("Valid"), + TuiSplit::Test => f.write_str("Test"), + } + } +} + +impl TuiSplit { + pub(crate) fn color(&self) -> Color { + match self { + TuiSplit::Train => Color::LightRed, + TuiSplit::Valid => Color::LightBlue, + TuiSplit::Test => Color::LightGreen, + } + } +} diff --git a/crates/burn-train/src/renderer/tui/controls.rs b/crates/burn-train/src/renderer/tui/controls.rs new file mode 100644 index 0000000..e487780 --- /dev/null +++ b/crates/burn-train/src/renderer/tui/controls.rs @@ -0,0 +1,46 @@ +use super::TerminalFrame; +use ratatui::{ + prelude::{Alignment, Rect}, + style::{Color, Style, Stylize}, + text::{Line, Span}, + widgets::{Block, Borders, Paragraph, Wrap}, +}; + +/// Controls view. +pub(crate) struct ControlsView; + +impl ControlsView { + /// Render the view. + pub(crate) fn render(self, frame: &mut TerminalFrame<'_>, size: Rect) { + let lines = vec![ + vec![ + Span::from(" Quit : ").yellow().bold(), + Span::from("q ").bold(), + Span::from(" Stop the training.").italic(), + ], + vec![ + Span::from(" Plots Metrics : ").yellow().bold(), + Span::from("⬅ ➡").bold(), + Span::from(" Switch between metrics.").italic(), + ], + vec![ + Span::from(" Plots Type : ").yellow().bold(), + Span::from("⬆ ⬇").bold(), + Span::from(" Switch between types.").italic(), + ], + ]; + let paragraph = Paragraph::new(lines.into_iter().map(Line::from).collect::>()) + .alignment(Alignment::Left) + .wrap(Wrap { trim: false }) + .style(Style::default().fg(Color::Gray)) + .block( + Block::default() + .borders(Borders::ALL) + .style(Style::default().fg(Color::Gray)) + .title_alignment(Alignment::Left) + .title("Controls"), + ); + + frame.render_widget(paragraph, size); + } +} diff --git a/crates/burn-train/src/renderer/tui/full_history.rs b/crates/burn-train/src/renderer/tui/full_history.rs new file mode 100644 index 0000000..437c5ce --- /dev/null +++ b/crates/burn-train/src/renderer/tui/full_history.rs @@ -0,0 +1,295 @@ +use super::PlotAxes; +use crate::{ + metric::NumericEntry, + renderer::tui::{TuiSplit, TuiTag}, +}; +use ratatui::{ + style::{Color, Style}, + symbols, + widgets::{Bar, Dataset, GraphType}, +}; +use std::collections::BTreeMap; + +/// A plot that shows the full history at a reduced resolution. +pub(crate) struct FullHistoryPlot { + pub(crate) axes: PlotAxes, + points: BTreeMap, + max_samples: usize, + max_samples_ratio: BTreeMap, + next_x_state: usize, +} + +struct FullHistoryPoints { + min_x: f64, + max_x: f64, + min_y: f64, + max_y: f64, + avg_sum: f64, + avg_counter: f64, + points: Vec<(f64, f64)>, + max_samples: usize, + step_size: usize, +} + +impl FullHistoryPlot { + /// Create a new history plot. + pub(crate) fn new(max_samples: usize) -> Self { + Self { + points: BTreeMap::default(), + axes: PlotAxes::default(), + max_samples, + max_samples_ratio: BTreeMap::default(), + next_x_state: 0, + } + } + + /// Update the maximum amount of sample to display for the validation points. + /// + /// This is necessary if we want the validation line to have the same point density as the + /// training line. + pub(crate) fn update_max_sample(&mut self, split: TuiSplit, ratio: f64) { + self.max_samples_ratio.insert(split, ratio); + + self.points + .iter_mut() + .filter(|(tag, _)| tag.split == split) + .for_each(|(_, points)| { + points.max_samples = (self.max_samples as f64 * ratio) as usize; + }); + } + + /// Register a training data point. + pub(crate) fn push(&mut self, tag: TuiTag, data: NumericEntry) { + let x_current = self.next_x(); + let points = match self.points.get_mut(&tag) { + Some(val) => val, + None => { + let max_samples = self + .max_samples_ratio + .get(&tag.split) + .map(|ratio| (*ratio * self.max_samples as f64) as usize) + .unwrap_or(self.max_samples); + self.points + .insert(tag.clone(), FullHistoryPoints::new(max_samples)); + self.points.get_mut(&tag).unwrap() + } + }; + + points.push((x_current, data)); + + self.update_bounds(); + } + + pub(crate) fn datasets(&self) -> Vec> { + let mut datasets = Vec::with_capacity(2); + + for (tag, points) in self.points.iter() { + datasets.push(points.dataset(format!("{tag}"), tag.split.color())); + } + + datasets + } + + pub(crate) fn bars(&self, max: u64, bar_width: &mut usize) -> Vec> { + let mut bars = Vec::new(); + + for (tag, points) in self.points.iter() { + if let Some((bar, width)) = points.bar(tag, max) { + *bar_width = usize::max(*bar_width, width); + bars.push(bar); + } + } + + bars + } + + fn next_x(&mut self) -> f64 { + let value = self.next_x_state; + self.next_x_state += 1; + value as f64 + } + + fn update_bounds(&mut self) { + let (mut x_min, mut x_max) = (f64::MAX, f64::MIN); + let (mut y_min, mut y_max) = (f64::MAX, f64::MIN); + + for points in self.points.values() { + x_min = f64::min(x_min, points.min_x); + x_max = f64::max(x_max, points.max_x); + y_min = f64::min(y_min, points.min_y); + y_max = f64::max(y_max, points.max_y); + } + + self.axes.update_bounds((x_min, x_max), (y_min, y_max)); + } +} + +impl FullHistoryPoints { + fn new(max_samples: usize) -> Self { + Self { + min_x: 0., + max_x: 0., + min_y: f64::MAX, + max_y: f64::MIN, + avg_sum: 0.0, + avg_counter: 0.0, + points: Vec::with_capacity(max_samples), + max_samples, + step_size: 1, + } + } + + fn push(&mut self, (x, y): (f64, NumericEntry)) { + if !(x as usize).is_multiple_of(self.step_size) { + return; + } + + let y = match y { + NumericEntry::Value(val) => { + self.avg_sum += val; + self.avg_counter += 1.0; + val + } + NumericEntry::Aggregated { + aggregated_value, + count, + } => { + self.avg_sum += aggregated_value * count as f64; + self.avg_counter += count as f64; + aggregated_value + } + }; + + if x > self.max_x { + self.max_x = x; + } + if x < self.min_x { + self.min_x = x; + } + if y > self.max_y { + self.max_y = y; + } + if y < self.min_y { + self.min_y = y + } + + self.points.push((x, y)); + + if self.points.len() > self.max_samples { + self.resize(); + } + } + + /// We keep only half the points and we double the step size. + /// + /// This ensure that we have the same amount of points across the X axis. + fn resize(&mut self) { + let mut points = Vec::with_capacity(self.max_samples / 2); + let mut max_x = f64::MIN; + let mut max_y = f64::MIN; + let mut min_x = f64::MAX; + let mut min_y = f64::MAX; + + for (i, (x, y)) in self.points.drain(0..self.points.len()).enumerate() { + if i % 2 == 0 { + if x > max_x { + max_x = x; + } + if x < min_x { + min_x = x; + } + if y > max_y { + max_y = y; + } + if y < min_y { + min_y = y; + } + + points.push((x, y)); + } + } + + self.points = points; + self.step_size *= 2; + + self.min_x = min_x; + self.max_x = max_x; + self.min_y = min_y; + self.max_y = max_y; + } + + fn dataset<'a>(&'a self, name: String, color: Color) -> Dataset<'a> { + Dataset::default() + .name(name) + .marker(symbols::Marker::Braille) + .style(Style::default().fg(color).bold()) + .graph_type(GraphType::Line) + .data(&self.points) + } + + fn bar<'a>(&'a self, tag: &TuiTag, max: u64) -> Option<(Bar<'a>, usize)> { + if self.avg_sum == 0.0 { + return None; + } + + let label = format!("{tag}"); + let width = usize::max(label.len(), 7); // 7 min width + + let factor = max as f64; + + let avg = self.avg_sum / self.avg_counter; + + Some(( + Bar::default() + .value((avg * factor) as u64) + .style(tag.split.color()) + .text_value(format!("{:.2}", avg)) + .label(label), + width, + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::renderer::tui::{TuiGroup, TuiSplit}; + + #[test] + fn test_points() { + let mut chart = FullHistoryPlot::new(10); + let tag_train = TuiTag::new(TuiSplit::Train, TuiGroup::Default); + let tag_valid = TuiTag::new(TuiSplit::Valid, TuiGroup::Default); + chart.update_max_sample(tag_valid.split, 0.6); + + for i in 0..100 { + chart.push(tag_train.clone(), NumericEntry::Value(i as f64)); + } + for i in 0..60 { + chart.push(tag_valid.clone(), NumericEntry::Value(i as f64)); + } + + let expected_train = vec![ + (0.0, 0.0), + (16.0, 16.0), + (32.0, 32.0), + (48.0, 48.0), + (64.0, 64.0), + (80.0, 80.0), + (96.0, 96.0), + ]; + + let expected_valid = vec![(100.0, 0.0), (116.0, 16.0), (128.0, 28.0), (144.0, 44.0)]; + + assert_eq!( + chart.points.get(&tag_train).unwrap().points, + expected_train, + "Expected train data points" + ); + assert_eq!( + chart.points.get(&tag_valid).unwrap().points, + expected_valid, + "Expected valid data points" + ); + } +} diff --git a/crates/burn-train/src/renderer/tui/metric_numeric.rs b/crates/burn-train/src/renderer/tui/metric_numeric.rs new file mode 100644 index 0000000..5344ffc --- /dev/null +++ b/crates/burn-train/src/renderer/tui/metric_numeric.rs @@ -0,0 +1,819 @@ +use crate::{ + logger::ProgressSnapshot, + metric::{MetricName, NumericEntry}, + renderer::tui::TuiTag, +}; + +use super::{FullHistoryPlot, RecentHistoryPlot, TerminalFrame, TuiSplit}; +use ratatui::{ + crossterm::event::{ + Event, KeyCode, KeyEvent, KeyEventKind, MouseButton, MouseEvent, MouseEventKind, + }, + prelude::{Alignment, Constraint, Direction, Layout, Position, Rect}, + style::{Color, Modifier, Style, Stylize}, + text::Line, + widgets::{ + Axis, BarChart, BarGroup, Block, Borders, Chart, LegendPosition, Padding, Paragraph, Tabs, + Widget, + }, +}; +use std::collections::BTreeMap; +use unicode_width::UnicodeWidthStr; + +/// 1 cell of padding on each side of a tab title, matching ratatui's default `Tabs` widget. +const TAB_PADDING: u16 = 2; +/// 1-cell `│` divider between adjacent tabs in ratatui's default `Tabs` widget. +const TAB_DIVIDER: u16 = 1; + +/// 1000 seems to be required to see some improvement. +const MAX_NUM_SAMPLES_RECENT: usize = 1000; +/// 250 seems to be the right resolution when plotting all history. +/// Otherwise, there is too much points and the lines arent't smooth enough. +const MAX_NUM_SAMPLES_FULL: usize = 250; + +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +enum ChevronSide { + Left, + Right, +} + +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +enum HoverTarget { + Tab(usize), + Chevron(ChevronSide), +} + +/// Hit-test geometry and hover state for the tab strip, populated by `render_tab_strip` +/// on every frame and consumed by `on_mouse_event`. +#[derive(Default)] +pub(crate) struct TabStripState { + hovered: Option, + tab_rects: Vec, + chevron_left: Option, + chevron_right: Option, +} + +/// Numeric metrics state that handles creating plots. +#[derive(Default)] +pub(crate) struct NumericMetricsState { + data: BTreeMap, + names: Vec, + selected: usize, + kind: PlotKind, + num_samples_train: Option, + num_samples_valid: Option, + num_samples_test: Option, + epoch: usize, + strip: TabStripState, +} + +/// The kind of plot to display. +#[derive(Default, Clone, Copy)] +pub(crate) enum PlotKind { + /// Display the full history of the metric with reduced resolution. + #[default] + Full, + /// Display only the recent history of the metric, but with more resolution. + Recent, + Summary, +} + +impl NumericMetricsState { + /// Register a new training value for the metric with the given name. + pub(crate) fn push(&mut self, tag: TuiTag, name: MetricName, data: NumericEntry) { + if let Some((recent, full)) = self.data.get_mut(name.as_ref()) { + recent.push(tag.clone(), data.current()); + full.push(tag, data); + } else { + let mut recent = RecentHistoryPlot::new(MAX_NUM_SAMPLES_RECENT); + let mut full = FullHistoryPlot::new(MAX_NUM_SAMPLES_FULL); + + recent.push(tag.clone(), data.current()); + full.push(tag, data); + + self.names.push(name.clone()); + self.data.insert(name, (recent, full)); + } + } + + /// Update the state with the training progress. + pub(crate) fn update_progress_train(&mut self, progress: &ProgressSnapshot) { + self.epoch = progress.global.items_processed; + + if self.num_samples_train.is_some() { + return; + } + + self.num_samples_train = Some(progress.split.items_total); + } + + /// Update the state with the validation progress. + pub(crate) fn update_progress_valid(&mut self, progress: &ProgressSnapshot) { + if self.num_samples_valid.is_some() { + return; + } + + // If num_samples_train is None, keep the default max_samples for validation. + if let Some(num_sample_train) = self.num_samples_train { + for (_recent, full) in self.data.values_mut() { + let ratio = progress.split.items_total as f64 / num_sample_train as f64; + full.update_max_sample(TuiSplit::Valid, ratio); + } + } + + self.epoch = progress.global.items_processed; + self.num_samples_valid = Some(progress.split.items_total); + } + + /// Update the state with the testing progress. + pub(crate) fn update_progress_test(&mut self, progress: &ProgressSnapshot) { + if self.num_samples_test.is_some() { + return; + } + + if let Some(num_sample_train) = self.num_samples_train { + for (_recent, full) in self.data.values_mut() { + let ratio = progress.split.items_total as f64 / num_sample_train as f64; + full.update_max_sample(TuiSplit::Test, ratio); + } + } + + self.num_samples_test = Some(progress.split.items_total); + } + + /// Create a view to display the numeric metrics. + pub(crate) fn view(&mut self) -> NumericMetricView<'_> { + if self.names.is_empty() { + return NumericMetricView::None; + } + match self.kind { + PlotKind::Summary => { + let chart = Self::bar_chart(&self.names, &self.data, self.selected); + NumericMetricView::BarPlots { + titles: &self.names, + selected: self.selected, + chart, + strip: &mut self.strip, + } + } + kind => { + let chart = Self::line_chart(&self.names, &self.data, self.selected, kind); + NumericMetricView::LinePlots { + titles: &self.names, + selected: self.selected, + chart, + kind, + strip: &mut self.strip, + } + } + } + } + + /// Handle the current event. Returns `true` when visible state changed and a + /// redraw is warranted, `false` for events that produced no observable change + /// (so the caller can skip an unnecessary redraw). + pub(crate) fn on_event(&mut self, event: &Event) -> bool { + match event { + Event::Key(key) => self.on_key_event(key), + Event::Mouse(mouse) => self.on_mouse_event(mouse), + _ => false, + } + } + + fn on_key_event(&mut self, key: &KeyEvent) -> bool { + match key.kind { + KeyEventKind::Release | KeyEventKind::Repeat => (), + #[cfg(target_os = "windows")] // Fix the double toggle on Windows. + KeyEventKind::Press => return false, + #[cfg(not(target_os = "windows"))] + KeyEventKind::Press => (), + } + match key.code { + KeyCode::Right => { + self.next_metric(); + true + } + KeyCode::Left => { + self.previous_metric(); + true + } + KeyCode::Up | KeyCode::Down => { + self.switch_kind(); + true + } + _ => false, + } + } + + fn on_mouse_event(&mut self, mouse: &MouseEvent) -> bool { + let pos = Position::new(mouse.column, mouse.row); + let target = hover_target_at(&self.strip, pos); + match mouse.kind { + MouseEventKind::Moved => { + if self.strip.hovered == target { + false + } else { + self.strip.hovered = target; + true + } + } + MouseEventKind::Down(MouseButton::Left) => match target { + Some(HoverTarget::Tab(idx)) => { + self.selected = idx; + true + } + Some(HoverTarget::Chevron(ChevronSide::Left)) => { + self.previous_metric(); + true + } + Some(HoverTarget::Chevron(ChevronSide::Right)) => { + self.next_metric(); + true + } + None => false, + }, + _ => false, + } + } + + pub(crate) fn select_by_name(&mut self, name: &MetricName) { + if let Some(idx) = self.names.iter().position(|n| n == name) { + self.selected = idx; + } + } + + fn switch_kind(&mut self) { + self.kind = match self.kind { + PlotKind::Full => PlotKind::Recent, + PlotKind::Recent => PlotKind::Summary, + PlotKind::Summary => PlotKind::Full, + }; + } + + fn next_metric(&mut self) { + let len = self.data.len(); + if len == 0 { + return; + } + self.selected = (self.selected + 1) % len; + } + + fn previous_metric(&mut self) { + let len = self.data.len(); + if len == 0 { + return; + } + if self.selected > 0 { + self.selected -= 1; + } else { + self.selected = len - 1; + } + } + + fn line_chart<'a>( + names: &'a [MetricName], + data: &'a BTreeMap, + selected: usize, + kind: PlotKind, + ) -> Chart<'a> { + let name = names.get(selected).unwrap(); + let (recent, full) = data.get(name).unwrap(); + + let (datasets, axes) = match kind { + PlotKind::Full => (full.datasets(), &full.axes), + PlotKind::Recent => (recent.datasets(), &recent.axes), + _ => unreachable!(), + }; + + Chart::<'a>::new(datasets) + .block(Block::default()) + .x_axis( + Axis::default() + .style(Style::default().fg(Color::DarkGray)) + .labels(axes.labels_x.clone().into_iter().map(|s| s.bold())) + .bounds(axes.bounds_x), + ) + .y_axis( + Axis::default() + .style(Style::default().fg(Color::DarkGray)) + .labels(axes.labels_y.clone().into_iter().map(|s| s.bold())) + .bounds(axes.bounds_y), + ) + .legend_position(Some(LegendPosition::Right)) + } + + fn bar_chart<'a>( + names: &'a [MetricName], + data: &'a BTreeMap, + selected: usize, + ) -> BarChart<'a> { + let name = names.get(selected).unwrap(); + let (_recent, full) = data.get(name).unwrap(); + let mut bar_width = 0; + let bars = full.bars(100, &mut bar_width); + + let data = BarGroup::default().bars(&bars); + BarChart::default() + .block(Block::default().padding(Padding::new(2, 2, 2, 0))) + .bar_width(bar_width as u16) + .bar_gap(2) + .data(data) + } +} + +#[allow(clippy::large_enum_variant)] +pub(crate) enum NumericMetricView<'a> { + LinePlots { + titles: &'a [MetricName], + selected: usize, + chart: Chart<'a>, + kind: PlotKind, + strip: &'a mut TabStripState, + }, + BarPlots { + titles: &'a [MetricName], + selected: usize, + chart: BarChart<'a>, + strip: &'a mut TabStripState, + }, + None, +} + +impl NumericMetricView<'_> { + pub(crate) fn render(self, frame: &mut TerminalFrame<'_>, size: Rect) { + match self { + Self::LinePlots { + titles, + selected, + chart, + kind, + strip, + } => { + let plot_title = match kind { + PlotKind::Full => "Full History", + PlotKind::Recent => "Recent History", + _ => unreachable!(), + }; + render_plot_panel( + frame, size, "Plots", plot_title, titles, selected, strip, chart, + ); + } + Self::BarPlots { + titles, + selected, + chart, + strip, + } => { + render_plot_panel( + frame, size, "Summary", "Summary", titles, selected, strip, chart, + ); + } + Self::None => {} + } + } +} + +/// Draw the bordered plot panel: tab strip on top, centered plot title, then the chart. +#[allow(clippy::too_many_arguments)] +fn render_plot_panel( + frame: &mut TerminalFrame<'_>, + size: Rect, + block_title: &str, + plot_title: &str, + titles: &[MetricName], + selected: usize, + strip: &mut TabStripState, + chart: W, +) { + let block = Block::default() + .borders(Borders::ALL) + .title(block_title) + .title_alignment(Alignment::Left); + let inner = block.inner(size); + frame.render_widget(block, size); + + let chunks = Layout::default() + .direction(Direction::Vertical) + .constraints([ + Constraint::Length(2), + Constraint::Length(1), + Constraint::Min(0), + ]) + .split(inner); + + render_tab_strip(frame, chunks[0], titles, selected, strip); + let title = Paragraph::new(Line::from(plot_title.bold())).alignment(Alignment::Center); + frame.render_widget(title, chunks[1]); + frame.render_widget(chart, chunks[2]); +} + +/// Render the metric tabs in `area`, scrolling horizontally so the `selected` tab is always +/// visible. A `‹` / `›` indicator is drawn in a reserved cell on each side when tabs are +/// hidden off that edge. The hovered tab gets an extra underline, a hovered chevron gets a +/// brighter foreground. Hit-test rects for the visible tabs and the two chevrons are written +/// back into `strip` so `on_mouse_event` can route clicks. +fn render_tab_strip( + frame: &mut TerminalFrame<'_>, + area: Rect, + titles: &[MetricName], + selected: usize, + strip: &mut TabStripState, +) { + strip.tab_rects.clear(); + strip.tab_rects.resize(titles.len(), Rect::default()); + strip.chevron_left = None; + strip.chevron_right = None; + + if titles.is_empty() || area.width == 0 { + return; + } + + let titles_str: Vec = titles.iter().map(|t| t.to_string()).collect(); + let widths: Vec = titles_str.iter().map(|s| tab_cell_width(s)).collect(); + + let inner_width = area.width.saturating_sub(2); + let (start, end) = visible_tab_window(&widths, selected, inner_width); + + if start > 0 { + let left = Rect { width: 1, ..area }; + let color = chevron_color(strip.hovered, ChevronSide::Left); + frame.render_widget(Paragraph::new("‹").style(Style::default().fg(color)), left); + strip.chevron_left = Some(left); + } + if end < titles.len() { + let right = Rect { + x: area.x + area.width - 1, + width: 1, + ..area + }; + let color = chevron_color(strip.hovered, ChevronSide::Right); + frame.render_widget(Paragraph::new("›").style(Style::default().fg(color)), right); + strip.chevron_right = Some(right); + } + + let tabs_area = Rect { + x: area.x + 1, + width: inner_width, + ..area + }; + + // Hit-test rects mirror the on-screen layout of the visible window. Tabs scrolled + // off either edge keep the zero-sized default from `resize` above so the mouse test + // misses them. + let mut x = tabs_area.x; + let tabs_end = tabs_area.x.saturating_add(tabs_area.width); + for (i, &w) in (start..end).zip(&widths[start..end]) { + let remaining = tabs_end.saturating_sub(x); + strip.tab_rects[i] = Rect { + x: x.min(tabs_end), + y: tabs_area.y, + width: w.min(remaining), + height: tabs_area.height, + }; + x = x.saturating_add(w).saturating_add(TAB_DIVIDER); + } + + let tabs = Tabs::new(titles_str[start..end].iter().enumerate().map(|(local, s)| { + let span = s.clone().yellow(); + let span = if strip.hovered == Some(HoverTarget::Tab(start + local)) { + span.underlined() + } else { + span + }; + Line::from(vec![span]) + })) + .select(selected - start) + .highlight_style( + Style::default() + .add_modifier(Modifier::BOLD | Modifier::UNDERLINED) + .fg(Color::LightYellow), + ); + frame.render_widget(tabs, tabs_area); +} + +fn chevron_color(hovered: Option, side: ChevronSide) -> Color { + if hovered == Some(HoverTarget::Chevron(side)) { + Color::Gray + } else { + Color::DarkGray + } +} + +fn hover_target_at(strip: &TabStripState, pos: Position) -> Option { + if strip.chevron_left.is_some_and(|r| r.contains(pos)) { + return Some(HoverTarget::Chevron(ChevronSide::Left)); + } + if strip.chevron_right.is_some_and(|r| r.contains(pos)) { + return Some(HoverTarget::Chevron(ChevronSide::Right)); + } + strip + .tab_rects + .iter() + .position(|r| r.contains(pos)) + .map(HoverTarget::Tab) +} + +/// Cells consumed by one tab. Title display width plus ratatui's default padding. +fn tab_cell_width(title: &str) -> u16 { + u16::try_from(UnicodeWidthStr::width(title) + TAB_PADDING as usize).unwrap_or(u16::MAX) +} + +/// Pick the `[start, end)` slice of `widths` to render so the tab at `selected` is visible +/// inside `available` cells. The selected tab is pinned as far right as fits. `end` is then +/// grown rightward as far as the remaining space allows. Always returns +/// `start <= selected < end` when `widths` is non-empty. If a single tab exceeds +/// `available`, clipping is delegated to ratatui's `Tabs`. +fn visible_tab_window(widths: &[u16], selected: usize, available: u16) -> (usize, usize) { + if widths.is_empty() { + return (0, 0); + } + let selected = selected.min(widths.len() - 1); + let available = available as u32; + let divider = TAB_DIVIDER as u32; + + // Width of titles[start..=selected] including dividers between them. Maintained + // incrementally so the windowing loop is O(N) over the full title list. + let mut width: u32 = + widths[..=selected].iter().map(|&w| w as u32).sum::() + selected as u32 * divider; + + let mut start = 0; + while width > available && start < selected { + width -= widths[start] as u32 + divider; + start += 1; + } + + let mut end = selected + 1; + while end < widths.len() { + let next = width + widths[end] as u32 + divider; + if next > available { + break; + } + width = next; + end += 1; + } + + (start, end) +} + +#[cfg(test)] +mod tests { + use super::*; + use ratatui::crossterm::event::{KeyModifiers, MouseEvent}; + use std::sync::Arc; + + fn name(s: &str) -> MetricName { + Arc::new(s.to_string()) + } + + fn mouse(kind: MouseEventKind, column: u16, row: u16) -> Event { + Event::Mouse(MouseEvent { + kind, + column, + row, + modifiers: KeyModifiers::NONE, + }) + } + + #[test] + fn metric_navigation_on_empty_state_is_a_no_op() { + // Pressing Right or Left before any metric has been recorded must not + // panic. Previously this triggered `(0 + 1) % 0` and `0usize - 1` on + // the empty `data` map. + let mut state = NumericMetricsState::default(); + + state.next_metric(); + state.previous_metric(); + + assert_eq!(state.selected, 0); + assert!(state.data.is_empty()); + } + + #[test] + fn select_by_name_sets_index_on_match_and_no_ops_otherwise() { + let mut state = NumericMetricsState { + names: vec![name("loss"), name("acc"), name("lr")], + ..NumericMetricsState::default() + }; + + state.select_by_name(&name("acc")); + assert_eq!(state.selected, 1); + + state.select_by_name(&name("missing")); + assert_eq!(state.selected, 1); + } + + fn strip_with_tabs(tabs: Vec) -> TabStripState { + TabStripState { + tab_rects: tabs, + ..TabStripState::default() + } + } + + #[test] + fn mouse_click_on_tab_selects_it() { + let mut state = NumericMetricsState { + names: vec![name("loss"), name("acc")], + strip: strip_with_tabs(vec![Rect::new(0, 0, 6, 2), Rect::new(7, 0, 5, 2)]), + ..NumericMetricsState::default() + }; + + state.on_event(&mouse(MouseEventKind::Down(MouseButton::Left), 8, 0)); + + assert_eq!(state.selected, 1); + } + + #[test] + fn mouse_click_outside_tabs_does_not_change_selection() { + let mut state = NumericMetricsState { + names: vec![name("loss"), name("acc")], + selected: 0, + strip: strip_with_tabs(vec![Rect::new(0, 0, 6, 2), Rect::new(7, 0, 5, 2)]), + ..NumericMetricsState::default() + }; + + state.on_event(&mouse(MouseEventKind::Down(MouseButton::Left), 50, 50)); + + assert_eq!(state.selected, 0); + } + + #[test] + fn mouse_moved_updates_hovered() { + let mut state = NumericMetricsState { + names: vec![name("loss"), name("acc")], + strip: strip_with_tabs(vec![Rect::new(0, 0, 6, 2), Rect::new(7, 0, 5, 2)]), + ..NumericMetricsState::default() + }; + + state.on_event(&mouse(MouseEventKind::Moved, 2, 0)); + assert_eq!(state.strip.hovered, Some(HoverTarget::Tab(0))); + + state.on_event(&mouse(MouseEventKind::Moved, 50, 50)); + assert_eq!(state.strip.hovered, None); + } + + #[test] + fn mouse_click_on_left_chevron_goes_to_previous_metric() { + // Three metrics with the second one selected. A click on the left chevron should + // wrap-decrement the selection just like KeyCode::Left does. + let mut state = NumericMetricsState { + names: vec![name("a"), name("b"), name("c")], + data: BTreeMap::from_iter([ + ( + name("a"), + (RecentHistoryPlot::new(1), FullHistoryPlot::new(1)), + ), + ( + name("b"), + (RecentHistoryPlot::new(1), FullHistoryPlot::new(1)), + ), + ( + name("c"), + (RecentHistoryPlot::new(1), FullHistoryPlot::new(1)), + ), + ]), + selected: 1, + strip: TabStripState { + chevron_left: Some(Rect::new(0, 0, 1, 2)), + ..TabStripState::default() + }, + ..NumericMetricsState::default() + }; + + state.on_event(&mouse(MouseEventKind::Down(MouseButton::Left), 0, 0)); + assert_eq!(state.selected, 0); + } + + #[test] + fn mouse_click_on_right_chevron_goes_to_next_metric() { + let mut state = NumericMetricsState { + names: vec![name("a"), name("b"), name("c")], + data: BTreeMap::from_iter([ + ( + name("a"), + (RecentHistoryPlot::new(1), FullHistoryPlot::new(1)), + ), + ( + name("b"), + (RecentHistoryPlot::new(1), FullHistoryPlot::new(1)), + ), + ( + name("c"), + (RecentHistoryPlot::new(1), FullHistoryPlot::new(1)), + ), + ]), + selected: 1, + strip: TabStripState { + chevron_right: Some(Rect::new(20, 0, 1, 2)), + ..TabStripState::default() + }, + ..NumericMetricsState::default() + }; + + state.on_event(&mouse(MouseEventKind::Down(MouseButton::Left), 20, 0)); + assert_eq!(state.selected, 2); + } + + #[test] + fn mouse_moved_over_chevron_updates_hover_target() { + let mut state = NumericMetricsState { + strip: TabStripState { + chevron_left: Some(Rect::new(0, 0, 1, 2)), + chevron_right: Some(Rect::new(20, 0, 1, 2)), + ..TabStripState::default() + }, + ..NumericMetricsState::default() + }; + + state.on_event(&mouse(MouseEventKind::Moved, 0, 0)); + assert_eq!( + state.strip.hovered, + Some(HoverTarget::Chevron(ChevronSide::Left)) + ); + + state.on_event(&mouse(MouseEventKind::Moved, 20, 0)); + assert_eq!( + state.strip.hovered, + Some(HoverTarget::Chevron(ChevronSide::Right)) + ); + } + + #[test] + fn on_event_returns_false_for_mouse_move_that_does_not_change_hover() { + let mut state = NumericMetricsState { + names: vec![name("loss"), name("acc")], + strip: TabStripState { + tab_rects: vec![Rect::new(0, 0, 6, 2), Rect::new(7, 0, 5, 2)], + ..TabStripState::default() + }, + ..NumericMetricsState::default() + }; + + // First move onto tab 0: hover changes from None to Some(Tab(0)). + assert!(state.on_event(&mouse(MouseEventKind::Moved, 2, 0))); + // Second move still on tab 0: no change, no redraw needed. + assert!(!state.on_event(&mouse(MouseEventKind::Moved, 3, 0))); + // Move off any tab: hover changes back to None. + assert!(state.on_event(&mouse(MouseEventKind::Moved, 50, 50))); + // Move outside again: still None, no change. + assert!(!state.on_event(&mouse(MouseEventKind::Moved, 60, 60))); + } + + #[test] + fn on_event_returns_true_for_tab_click_false_for_missed_click() { + let mut state = NumericMetricsState { + names: vec![name("loss"), name("acc")], + strip: TabStripState { + tab_rects: vec![Rect::new(0, 0, 6, 2), Rect::new(7, 0, 5, 2)], + ..TabStripState::default() + }, + ..NumericMetricsState::default() + }; + + assert!(state.on_event(&mouse(MouseEventKind::Down(MouseButton::Left), 8, 0))); + assert!(!state.on_event(&mouse(MouseEventKind::Down(MouseButton::Left), 50, 50))); + } + + fn cells(titles: &[&str]) -> Vec { + titles.iter().copied().map(tab_cell_width).collect() + } + + #[test] + fn tab_cell_width_includes_padding() { + assert_eq!(tab_cell_width("Loss"), 4 + TAB_PADDING); + assert_eq!(tab_cell_width(""), TAB_PADDING); + } + + #[test] + fn visible_window_is_empty_when_no_tabs() { + assert_eq!(visible_tab_window(&[], 0, 80), (0, 0)); + } + + #[test] + fn visible_window_returns_full_range_when_everything_fits() { + let widths = cells(&["Loss", "Acc", "F1"]); + assert_eq!(visible_tab_window(&widths, 1, 80), (0, widths.len())); + } + + #[test] + fn visible_window_pins_selected_to_right_edge_when_growing_from_left() { + // Each tab cell is 2 + TAB_PADDING = 4. One divider between = 5 per added tab. + // available = 9 means exactly two tabs fit ([start..=selected] = 4 + 1 + 4 = 9). + let widths = cells(&["AA", "BB", "CC", "DD", "EE"]); + assert_eq!(visible_tab_window(&widths, 3, 9), (2, 4)); + } + + #[test] + fn visible_window_does_not_panic_when_single_tab_is_wider_than_available() { + let widths = cells(&["this title is much wider than the cell", "B"]); + let (start, end) = visible_tab_window(&widths, 0, 4); + assert_eq!(start, 0); + assert!(end >= 1); + } + + #[test] + fn visible_window_clamps_out_of_range_selected() { + let widths = cells(&["A", "B", "C"]); + let (start, end) = visible_tab_window(&widths, 99, 80); + assert!(start <= 2 && 2 < end); + } +} diff --git a/crates/burn-train/src/renderer/tui/metric_text.rs b/crates/burn-train/src/renderer/tui/metric_text.rs new file mode 100644 index 0000000..1a016a0 --- /dev/null +++ b/crates/burn-train/src/renderer/tui/metric_text.rs @@ -0,0 +1,290 @@ +use super::TerminalFrame; +use crate::{ + metric::{MetricEntry, MetricName}, + renderer::tui::{TuiGroup, TuiSplit}, +}; +use ratatui::{ + crossterm::event::{Event, MouseButton, MouseEventKind}, + prelude::{Alignment, Position, Rect}, + style::{Color, Style, Stylize}, + text::{Line, Span}, + widgets::{Block, Borders, Paragraph, Wrap}, +}; +use std::{collections::BTreeMap, ops::Range, sync::Arc}; + +/// Hit-test geometry and hover state for the metrics pane, populated by +/// `TextMetricView::render` on every frame and consumed by `on_event`. +#[derive(Default)] +pub(crate) struct TextHitState { + hovered: Option, + rect: Option, + header_rows: Vec<(MetricName, Range)>, +} + +#[derive(Default)] +pub(crate) struct TextMetricsState { + data: BTreeMap, + names: Vec, + pane: TextHitState, +} + +/// What a mouse event meant for the left pane. Drives both selection routing +/// (the `Clicked` arm carries the metric name to switch to) and redraw gating +/// (anything other than `Ignored` should cause a redraw). +pub(crate) enum TextEventOutcome { + Clicked(MetricName), + HoverChanged, + Ignored, +} + +struct MetricGroup { + groups: BTreeMap, +} + +impl MetricGroup { + fn new(group: TuiGroup, metric: MetricSplits) -> Self { + Self { + groups: BTreeMap::from_iter(Some((group, metric))), + } + } + fn update(&mut self, split: TuiSplit, group: TuiGroup, metric: MetricEntry) { + match self.groups.get_mut(&group) { + Some(value) => value.update(split, metric), + None => { + let value = MetricSplits::new(split, metric); + + self.groups.insert(group, value); + } + } + } +} + +struct MetricSplits { + splits: BTreeMap, +} + +impl MetricSplits { + fn new(split: TuiSplit, metric: MetricEntry) -> Self { + Self { + splits: BTreeMap::from_iter(Some((split, metric))), + } + } + + fn update(&mut self, split: TuiSplit, metric: MetricEntry) { + self.splits.insert(split, metric); + } +} + +impl TextMetricsState { + pub(crate) fn update( + &mut self, + split: TuiSplit, + group: TuiGroup, + metric: MetricEntry, + name: Arc, + ) { + if let Some(existing) = self.data.get_mut(name.as_ref()) { + existing.update(split, group, metric); + } else { + let key = name.clone(); + let value = MetricSplits::new(split, metric); + + self.names.push(key.clone()); + self.data + .insert(key.to_string(), MetricGroup::new(group, value)); + } + } + + pub(crate) fn view(&mut self) -> TextMetricView<'_> { + TextMetricView::new(&self.names, &self.data, &mut self.pane) + } + + /// Updates hover state and reports what the event meant for the left pane. + pub(crate) fn on_event(&mut self, event: &Event) -> TextEventOutcome { + let Event::Mouse(mouse) = event else { + return TextEventOutcome::Ignored; + }; + let pos = Position::new(mouse.column, mouse.row); + let hit = if self.pane.rect.is_some_and(|pane| pane.contains(pos)) { + self.pane + .header_rows + .iter() + .find(|(_, rows)| rows.contains(&pos.y)) + .map(|(name, _)| name.clone()) + } else { + None + }; + + match mouse.kind { + MouseEventKind::Moved => { + if self.pane.hovered == hit { + TextEventOutcome::Ignored + } else { + self.pane.hovered = hit; + TextEventOutcome::HoverChanged + } + } + MouseEventKind::Down(MouseButton::Left) => match hit { + Some(name) => TextEventOutcome::Clicked(name), + None => TextEventOutcome::Ignored, + }, + _ => TextEventOutcome::Ignored, + } + } +} + +pub(crate) struct TextMetricView<'a> { + lines: Vec>>, + /// Index into `lines` of each metric's header row, in display order. + header_line_indices: Vec<(MetricName, usize)>, + pane: &'a mut TextHitState, +} + +impl<'a> TextMetricView<'a> { + fn new( + names: &[MetricName], + data: &BTreeMap, + pane: &'a mut TextHitState, + ) -> Self { + let mut lines = Vec::with_capacity(names.len() * 4); + let mut header_line_indices = Vec::with_capacity(names.len()); + + let hovered = pane.hovered.as_ref(); + let start_line = |title: &str, is_hovered: bool| { + let span = Span::from(format!(" {title} ")).bold().yellow(); + let span = if is_hovered { span.underlined() } else { span }; + vec![span] + }; + let format_line = |group: &TuiGroup, split: &TuiSplit, formatted: &str| { + vec![ + Span::from(format!(" {group}{split} ")).bold(), + Span::from(formatted.to_string()).italic(), + ] + }; + + for name in names { + let is_hovered = hovered.is_some_and(|h| h.as_ref() == name.as_ref()); + header_line_indices.push((name.clone(), lines.len())); + lines.push(start_line(name, is_hovered)); + + let entry = data.get(name.as_ref()).unwrap(); + + for (name, group) in entry.groups.iter() { + for (split, entry) in group.splits.iter() { + lines.push(format_line(name, split, &entry.serialized_entry.formatted)); + } + } + + lines.push(vec![Span::from("")]); + } + + Self { + lines, + header_line_indices, + pane, + } + } + + pub(crate) fn render(self, frame: &mut TerminalFrame<'_>, size: Rect) { + let Self { + lines, + header_line_indices, + pane, + } = self; + + // Skip the 1-cell top border. Header lines longer than the pane will + // wrap and misplace the hit zone, accepted as a known limitation. + let text_origin_y = size.y.saturating_add(1); + pane.rect = Some(size); + pane.header_rows = header_line_indices + .into_iter() + .map(|(name, line_idx)| { + let row = text_origin_y.saturating_add(line_idx as u16); + (name, row..row.saturating_add(1)) + }) + .collect(); + + let paragraph = Paragraph::new(lines.into_iter().map(Line::from).collect::>()) + .alignment(Alignment::Left) + .wrap(Wrap { trim: false }) + .block(Block::default().borders(Borders::ALL).title("Metrics")) + .style(Style::default().fg(Color::Gray)); + + frame.render_widget(paragraph, size); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use ratatui::crossterm::event::{KeyModifiers, MouseEvent}; + + fn name(s: &str) -> MetricName { + Arc::new(s.to_string()) + } + + fn mouse(kind: MouseEventKind, column: u16, row: u16) -> Event { + Event::Mouse(MouseEvent { + kind, + column, + row, + modifiers: KeyModifiers::NONE, + }) + } + + fn state_with(headers: Vec<(MetricName, Range)>) -> TextMetricsState { + TextMetricsState { + pane: TextHitState { + rect: Some(Rect::new(0, 0, 20, 10)), + header_rows: headers, + ..TextHitState::default() + }, + ..TextMetricsState::default() + } + } + + #[test] + fn click_on_header_row_returns_clicked() { + let mut state = state_with(vec![(name("Loss"), 1..2), (name("Accuracy"), 5..6)]); + + let outcome = state.on_event(&mouse(MouseEventKind::Down(MouseButton::Left), 3, 5)); + + match outcome { + TextEventOutcome::Clicked(n) => assert_eq!(n.as_str(), "Accuracy"), + _ => panic!("expected Clicked"), + } + } + + #[test] + fn click_off_any_header_returns_ignored() { + let mut state = state_with(vec![(name("Loss"), 1..2)]); + + let on_data_row = state.on_event(&mouse(MouseEventKind::Down(MouseButton::Left), 3, 3)); + let outside_pane = state.on_event(&mouse(MouseEventKind::Down(MouseButton::Left), 50, 50)); + + assert!(matches!(on_data_row, TextEventOutcome::Ignored)); + assert!(matches!(outside_pane, TextEventOutcome::Ignored)); + } + + #[test] + fn moved_event_signals_hover_change_only_when_target_changes() { + let mut state = state_with(vec![(name("Loss"), 1..2)]); + + // First move onto Loss: hover changes from None to Some(Loss). + let r1 = state.on_event(&mouse(MouseEventKind::Moved, 3, 1)); + assert!(matches!(r1, TextEventOutcome::HoverChanged)); + assert_eq!( + state.pane.hovered.as_deref().map(|s| s.as_str()), + Some("Loss") + ); + + // Second move still on Loss: no change, ignored (no redraw needed). + let r2 = state.on_event(&mouse(MouseEventKind::Moved, 4, 1)); + assert!(matches!(r2, TextEventOutcome::Ignored)); + + // Move off the header: hover changes to None. + let r3 = state.on_event(&mouse(MouseEventKind::Moved, 3, 8)); + assert!(matches!(r3, TextEventOutcome::HoverChanged)); + assert!(state.pane.hovered.is_none()); + } +} diff --git a/crates/burn-train/src/renderer/tui/mod.rs b/crates/burn-train/src/renderer/tui/mod.rs new file mode 100644 index 0000000..b69d0c5 --- /dev/null +++ b/crates/burn-train/src/renderer/tui/mod.rs @@ -0,0 +1,23 @@ +mod base; +mod controls; +mod full_history; +mod metric_numeric; +mod metric_text; +mod plot_utils; +mod popup; +mod progress; +mod recent_history; +mod renderer; +mod status; + +pub(crate) use base::*; +pub(crate) use controls::*; +pub(crate) use full_history::*; +pub(crate) use metric_numeric::*; +pub(crate) use metric_text::*; +pub(crate) use plot_utils::*; +pub(crate) use popup::*; +pub(crate) use progress::*; +pub(crate) use recent_history::*; +pub use renderer::*; +pub(crate) use status::*; diff --git a/crates/burn-train/src/renderer/tui/plot_utils.rs b/crates/burn-train/src/renderer/tui/plot_utils.rs new file mode 100644 index 0000000..7a8a6d4 --- /dev/null +++ b/crates/burn-train/src/renderer/tui/plot_utils.rs @@ -0,0 +1,37 @@ +use crate::metric::format_float; + +const AXIS_TITLE_PRECISION: usize = 2; + +/// The data describing both X and Y axes. +pub(crate) struct PlotAxes { + pub(crate) labels_x: Vec, + pub(crate) labels_y: Vec, + pub(crate) bounds_x: [f64; 2], + pub(crate) bounds_y: [f64; 2], +} + +impl Default for PlotAxes { + fn default() -> Self { + Self { + bounds_x: [f64::MAX, f64::MIN], + bounds_y: [f64::MAX, f64::MIN], + labels_x: Vec::new(), + labels_y: Vec::new(), + } + } +} + +impl PlotAxes { + /// Update the bounds based on the min max of each X and Y axes with both train and valid data. + pub(crate) fn update_bounds(&mut self, (x_min, x_max): (f64, f64), (y_min, y_max): (f64, f64)) { + self.bounds_x = [x_min, x_max]; + self.bounds_y = [y_min, y_max]; + + // We know x are integers. + self.labels_x = vec![format!("{x_min}"), format!("{x_max}")]; + self.labels_y = vec![ + format_float(y_min, AXIS_TITLE_PRECISION), + format_float(y_max, AXIS_TITLE_PRECISION), + ]; + } +} diff --git a/crates/burn-train/src/renderer/tui/popup.rs b/crates/burn-train/src/renderer/tui/popup.rs new file mode 100644 index 0000000..e3092c5 --- /dev/null +++ b/crates/burn-train/src/renderer/tui/popup.rs @@ -0,0 +1,146 @@ +use ratatui::{ + crossterm::event::{Event, KeyCode}, + prelude::{Alignment, Constraint, Direction, Layout, Rect}, + style::{Color, Modifier, Style, Stylize}, + text::{Line, Span}, + widgets::{Block, Borders, Paragraph, Wrap}, +}; + +use super::TerminalFrame; + +/// Popup callback function. +pub(crate) trait CallbackFn: Send + Sync { + /// Call the function and return if the popup state should be reset. + fn call(&self) -> bool; +} + +/// Popup callback. +pub(crate) struct Callback { + title: String, + description: String, + trigger: char, + callback: Box, +} + +impl Callback { + /// Create a new popup. + pub(crate) fn new(title: T, description: D, trigger: char, callback: C) -> Self + where + T: Into, + D: Into, + C: CallbackFn + 'static, + { + Self { + title: title.into(), + description: description.into(), + trigger, + callback: Box::new(callback), + } + } +} + +/// Popup state. +pub(crate) enum PopupState { + Empty, + Full(String, Vec), +} + +impl PopupState { + /// If the popup is empty. + pub(crate) fn is_empty(&self) -> bool { + matches!(&self, PopupState::Empty) + } + /// Handle popup events. Returns `true` when the popup state changed and a + /// redraw is warranted, `false` for events that left the popup untouched. + pub(crate) fn on_event(&mut self, event: &Event) -> bool { + let mut reset = false; + + match self { + PopupState::Empty => {} + PopupState::Full(_, callbacks) => { + for callback in callbacks.iter() { + if let Event::Key(key) = event + && let KeyCode::Char(key) = &key.code + && &callback.trigger == key + && callback.callback.call() + { + reset = true; + } + } + } + }; + + if reset { + *self = Self::Empty; + } + reset + } + /// Create the popup view. + pub(crate) fn view(&self) -> Option> { + match self { + PopupState::Empty => None, + PopupState::Full(title, callbacks) => Some(PopupView::new(title, callbacks)), + } + } +} + +#[derive(new)] +pub(crate) struct PopupView<'a> { + title: &'a String, + callbacks: &'a [Callback], +} + +impl<'a> PopupView<'a> { + /// Render the view. + pub(crate) fn render<'b>(&'a self, frame: &mut TerminalFrame<'b>, size: Rect) { + let lines = self + .callbacks + .iter() + .flat_map(|callback| { + vec![ + Line::from(vec![ + Span::from(format!("[{}] ", callback.trigger)).bold(), + Span::from(format!("{} ", callback.title)).yellow().bold(), + ]), + Line::from(Span::from("")), + Line::from(Span::from(callback.description.to_string()).italic()), + Line::from(Span::from("")), + ] + }) + .collect::>(); + + let paragraph = Paragraph::new(lines) + .alignment(Alignment::Left) + .wrap(Wrap { trim: false }) + .style(Style::default().fg(Color::Gray)) + .block( + Block::default() + .borders(Borders::ALL) + .title_alignment(Alignment::Center) + .style(Style::default().fg(Color::Gray)) + .title(Span::styled( + self.title, + Style::default().add_modifier(Modifier::BOLD), + )), + ); + + let area = centered_percent(20, size, Direction::Horizontal); + let area = centered_percent(20, area, Direction::Vertical); + + frame.render_widget(paragraph, area); + } +} + +/// The percent represents the amount of space that will be taken by each side. +fn centered_percent(percent: u16, size: Rect, direction: Direction) -> Rect { + let center = 100 - (percent * 2); + + Layout::default() + .direction(direction) + .constraints([ + Constraint::Percentage(percent), + Constraint::Percentage(center), + Constraint::Percentage(percent), + ]) + .split(size)[1] +} diff --git a/crates/burn-train/src/renderer/tui/progress.rs b/crates/burn-train/src/renderer/tui/progress.rs new file mode 100644 index 0000000..f9ff19e --- /dev/null +++ b/crates/burn-train/src/renderer/tui/progress.rs @@ -0,0 +1,296 @@ +use super::TerminalFrame; +use crate::{logger::ProgressSnapshot, renderer::tui::TuiSplit}; +use ratatui::{ + prelude::{Alignment, Constraint, Direction, Layout, Rect}, + style::{Color, Style, Stylize}, + text::{Line, Span}, + widgets::{Block, Borders, Gauge, Paragraph}, +}; +use std::time::{Duration, Instant}; + +/// Simple progress bar for the training. +/// +/// We currently ignore the time taken for the validation part. +pub(crate) struct ProgressBarState { + progress_total: f64, // Progress for total execution. + progress_task: f64, // Progress for current task. + split: TuiSplit, + starting_epoch: usize, + estimate: ProgressEstimate, +} + +const MINUTE: u64 = 60; +const HOUR: u64 = 60 * 60; +const DAY: u64 = 24 * 60 * 60; + +impl ProgressBarState { + pub fn new(checkpoint: Option) -> Self { + Self { + progress_total: 0.0, + progress_task: 0.0, + split: TuiSplit::Train, + estimate: ProgressEstimate::new(), + starting_epoch: checkpoint.unwrap_or(0), + } + } + + /// Update the training progress. + pub(crate) fn update_train(&mut self, progress: &ProgressSnapshot) { + self.progress_total = calculate_progress(progress, 0, 0); + self.progress_task = + progress.split.items_processed as f64 / progress.split.items_total as f64; + self.estimate.update(progress, self.starting_epoch); + self.split = TuiSplit::Train; + } + + /// Update the validation progress. + pub(crate) fn update_valid(&mut self, progress: &ProgressSnapshot) { + // We don't use the validation for the total progress yet. + self.progress_task = + progress.split.items_processed as f64 / progress.split.items_total as f64; + self.split = TuiSplit::Valid; + } + + /// Update the testing progress. + pub(crate) fn update_test(&mut self, progress: &ProgressSnapshot) { + // We don't use the testing for the total progress yet. + self.progress_task = + progress.split.items_processed as f64 / progress.split.items_total as f64; + self.split = TuiSplit::Test; + } + + /// Create a view for the current progress. + pub(crate) fn view(&self) -> ProgressBarView { + const NO_ETA: &str = "---"; + + let eta = match self.estimate.secs() { + Some(eta) => format_eta(eta), + None => NO_ETA.to_string(), + }; + ProgressBarView::new( + self.progress_total, + self.progress_task, + self.split.color(), + eta, + ) + } +} + +#[derive(new)] +pub(crate) struct ProgressBarView { + progress: f64, + progress_task: f64, + color_task: Color, + eta: String, +} + +impl ProgressBarView { + /// Render the view. + pub(crate) fn render(self, frame: &mut TerminalFrame<'_>, size: Rect) { + let block = Block::default() + .borders(Borders::ALL) + .title("Progress") + .title_alignment(Alignment::Left); + let size_new = block.inner(size); + frame.render_widget(block, size); + let size = size_new; + + let chunks = Layout::default() + .direction(Direction::Vertical) + .constraints([Constraint::Ratio(1, 2), Constraint::Ratio(1, 2)].as_ref()) + .split(size); + + let size_task = chunks[0]; + let size_total = chunks[1]; + + let calculate_size = |size: Rect| { + Layout::default() + .direction(Direction::Horizontal) + .constraints( + [ + Constraint::Length(1), // Empty space + Constraint::Min(0), + Constraint::Length(self.eta.len() as u16 + 4), + ] + .as_ref(), + ) + .split(size) + }; + + let chunks = calculate_size(size_total); + let size_gauge_total = chunks[1]; + let size_eta = chunks[2]; + let chunks = calculate_size(size_task); + let size_gauge_task = chunks[1]; + + let progress_total = Gauge::default() + .gauge_style(Style::default().fg(Color::Yellow)) + .ratio(self.progress.min(1.0)); + let progress_task = Gauge::default() + .gauge_style(Style::default().fg(self.color_task)) + .ratio(self.progress_task.min(1.0)); + + let eta = Paragraph::new(Line::from(vec![ + Span::from(" ("), + Span::from(self.eta).italic(), + Span::from(") "), + ])); + + frame.render_widget(progress_task, size_gauge_task); + frame.render_widget(progress_total, size_gauge_total); + frame.render_widget(eta, size_eta); + } +} + +struct ProgressEstimate { + started: Instant, + started_after_warmup: Option, + warmup_num_items: usize, + progress: f64, +} + +impl ProgressEstimate { + fn new() -> Self { + Self { + started: Instant::now(), + started_after_warmup: None, + warmup_num_items: 0, + progress: 0.0, + } + } + + fn secs(&self) -> Option { + let eta = self.started_after_warmup?.elapsed(); + + let total_estimated = (eta.as_secs() as f64) / self.progress; + + if total_estimated.is_normal() { + let remaining = 1.0 - self.progress; + let eta = (total_estimated * remaining) as u64; + Some(eta) + } else { + None + } + } + + fn update(&mut self, progress: &ProgressSnapshot, starting_epoch: usize) { + if self.started_after_warmup.is_some() { + self.progress = calculate_progress(progress, starting_epoch, self.warmup_num_items); + return; + } + + const WARMUP_NUM_ITERATION: usize = 10; + + // When the training has started since 30 seconds. + if self.started.elapsed() > Duration::from_secs(30) { + self.init(progress, starting_epoch); + return; + } + + // When the training has started since at least 10 seconds and completed 10 iterations. + if progress.split.items_processed >= WARMUP_NUM_ITERATION + && self.started.elapsed() > Duration::from_secs(10) + { + self.init(progress, starting_epoch); + } + } + + fn init(&mut self, progress: &ProgressSnapshot, starting_epoch: usize) { + let epoch = progress.global.items_processed - starting_epoch; + let local = &progress.split; + + let epoch_items = (epoch - 1) * local.items_total; + let iteration_items = local.items_processed; + self.warmup_num_items = epoch_items + iteration_items; + + self.started_after_warmup = Some(Instant::now()); + self.progress = calculate_progress(progress, starting_epoch, self.warmup_num_items); + } +} + +fn calculate_progress( + progress: &ProgressSnapshot, + starting_epoch: usize, + ignore_num_items: usize, +) -> f64 { + let epoch_total = progress.global.items_total - starting_epoch; + let epoch = progress.global.items_processed - starting_epoch; + let local = &progress.split; + + let total_items = local.items_total * epoch_total; + let epoch_items = (epoch - 1) * local.items_total; + let iteration_items = local.items_processed; + let num_items = epoch_items + iteration_items - ignore_num_items; + + num_items as f64 / total_items as f64 +} + +fn format_eta(eta_secs: u64) -> String { + let seconds = eta_secs % 60; + let minutes = eta_secs / MINUTE % 60; + let hours = eta_secs / HOUR % 24; + let days = eta_secs / DAY; + + if days > 1 { + format!("{days} days") + } else if days == 1 { + "1 day".to_string() + } else if hours > 1 { + format!("{hours} hours") + } else if hours == 1 { + "1 hour".to_string() + } else if minutes > 1 { + format!("{minutes} mins") + } else if minutes == 1 { + "1 min".to_string() + } else if seconds > 1 { + format!("{seconds} secs") + } else { + "1 sec".to_string() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use burn_core::data::dataloader::Progress; + + #[test] + fn test_format_eta() { + assert_eq!("55 secs", format_eta(55), "Less than 1 minutes"); + assert_eq!("1 min", format_eta(61), "More than 1 minutes"); + assert_eq!("2 mins", format_eta(2 * 61), "More than 2 minutes"); + assert_eq!("1 hour", format_eta(3601), "More than 1 hour"); + assert_eq!("2 hours", format_eta(2 * 3601), "More than 2 hour"); + assert_eq!("1 day", format_eta(24 * 3601), "More than 1 day"); + assert_eq!("2 days", format_eta(48 * 3601), "More than 2 day"); + } + + #[test] + fn calculate_progress_for_eta() { + let progress = ProgressSnapshot::new( + Progress::new(9, 10, Some("epochs".to_string())), + Progress::new(5, 10, Some("items".to_string())), + ); + + let starting_epoch = 8; + let result = calculate_progress(&progress, starting_epoch, 0); + + // Two epochs remaining while the first is half done. + assert_eq!(0.25, result); + } + + #[test] + fn calculate_progress_for_eta_with_warmup() { + let progress = ProgressSnapshot::new( + Progress::new(9, 10, Some("epochs".to_string())), + Progress::new(110, 1000, Some("items".to_string())), + ); + + let starting_epoch = 8; + let result = calculate_progress(&progress, starting_epoch, 10); + + // Two epochs remaining while the first is 11% done, minus 10 warmup items. + assert_eq!(0.05, result); + } +} diff --git a/crates/burn-train/src/renderer/tui/recent_history.rs b/crates/burn-train/src/renderer/tui/recent_history.rs new file mode 100644 index 0000000..909aff5 --- /dev/null +++ b/crates/burn-train/src/renderer/tui/recent_history.rs @@ -0,0 +1,249 @@ +use super::PlotAxes; +use crate::renderer::tui::TuiTag; +use ratatui::{ + style::{Color, Style}, + symbols, + widgets::{Dataset, GraphType}, +}; +use std::collections::BTreeMap; + +const FACTOR_BEFORE_RESIZE: usize = 2; + +/// A plot that shows the recent history at full resolution. +pub(crate) struct RecentHistoryPlot { + pub(crate) axes: PlotAxes, + points: BTreeMap, + max_samples: usize, +} + +struct RecentHistoryPoints { + min_x: f64, + max_x: f64, + min_y: f64, + max_y: f64, + cursor: usize, + points: Vec<(f64, f64)>, + max_samples: usize, + factor_before_resize: usize, +} + +impl RecentHistoryPlot { + pub(crate) fn new(max_samples: usize) -> Self { + Self { + axes: PlotAxes::default(), + points: BTreeMap::default(), + max_samples, + } + } + + pub(crate) fn push(&mut self, tag: TuiTag, data: f64) { + if !self.points.contains_key(&tag) { + self.points + .insert(tag.clone(), RecentHistoryPoints::new(self.max_samples)); + } + + let (x_min, x_current) = self.point_x(); + + for (s, entry) in self.points.iter_mut() { + if s == &tag { + entry.push((x_current, data)); + } + entry.update_cursor(x_min); + } + + self.update_bounds(); + } + + pub(crate) fn datasets(&self) -> Vec> { + let mut datasets = Vec::new(); + + for (tag, points) in self.points.iter() { + datasets.push(points.dataset(format!("{tag}"), tag.split.color())); + } + + datasets + } + + fn point_x(&mut self) -> (f64, f64) { + let mut x_current = f64::MIN; + let mut x_min = f64::MAX; + + for point in self.points.values() { + x_current = f64::max(x_current, point.max_x); + x_min = f64::min(x_min, point.min_x); + } + + if x_current - x_min >= self.max_samples as f64 { + x_min += 1.0; + } + + (x_min, x_current + 1.0) + } + + fn update_bounds(&mut self) { + let (mut x_min, mut x_max) = (f64::MAX, f64::MIN); + let (mut y_min, mut y_max) = (f64::MAX, f64::MIN); + + for points in self.points.values() { + x_min = f64::min(x_min, points.min_x); + x_max = f64::max(x_max, points.max_x); + y_min = f64::min(y_min, points.min_y); + y_max = f64::max(y_max, points.max_y); + } + + self.axes.update_bounds((x_min, x_max), (y_min, y_max)); + } +} + +impl RecentHistoryPoints { + fn new(max_samples: usize) -> Self { + let factor_before_resize = FACTOR_BEFORE_RESIZE; + + Self { + min_x: 0., + max_x: 0., + min_y: f64::MAX, + max_y: f64::MIN, + points: Vec::with_capacity(factor_before_resize * max_samples), + cursor: 0, + max_samples, + factor_before_resize, + } + } + + fn push(&mut self, (x, y): (f64, f64)) { + if x > self.max_x { + self.max_x = x; + } + if x < self.min_x { + self.min_x = x; + } + if y > self.max_y { + self.max_y = y; + } + if y < self.min_y { + self.min_y = y + } + self.points.push((x, y)); + } + + fn update_cursor(&mut self, min_x: f64) { + if self.min_x >= min_x { + return; + } + self.min_x = min_x; + + let mut update_y_max = false; + let mut update_y_min = false; + + while let Some((x, y)) = self.points.get(self.cursor) { + if *x >= self.min_x { + break; + } + + if *y == self.max_y { + update_y_max = true + } + if *y == self.min_y { + update_y_min = true; + } + + self.cursor += 1; + } + + if update_y_max { + self.max_y = self.calculate_max_y(); + } + + if update_y_min { + self.min_y = self.calculate_min_y(); + } + + if self.points.len() >= self.max_samples * self.factor_before_resize { + self.resize(); + } + } + + fn slice(&self) -> &[(f64, f64)] { + &self.points[self.cursor..self.points.len()] + } + + fn calculate_max_y(&self) -> f64 { + let mut max_y = f64::MIN; + + for (_x, y) in self.slice() { + max_y = f64::max(max_y, *y); + } + + max_y + } + + fn calculate_min_y(&self) -> f64 { + let mut min_y = f64::MAX; + + for (_x, y) in self.slice() { + if *y < min_y { + min_y = *y; + } + } + + min_y + } + + fn resize(&mut self) { + let mut points = Vec::with_capacity(self.max_samples * self.factor_before_resize); + + for i in self.cursor..self.points.len() { + points.push(self.points[i]); + } + + self.points = points; + self.cursor = 0; + } + + fn dataset<'a>(&'a self, name: String, color: Color) -> Dataset<'a> { + let data = &self.points[self.cursor..self.points.len()]; + + Dataset::default() + .name(name) + .marker(symbols::Marker::Braille) + .style(Style::default().fg(color).bold()) + .graph_type(GraphType::Scatter) + .data(data) + } +} + +#[cfg(test)] +mod tests { + use crate::renderer::tui::{TuiGroup, TuiSplit}; + + use super::*; + + #[test] + fn test_push_update_bounds_max_y() { + let mut chart = RecentHistoryPlot::new(2); + let tag = TuiTag::new(TuiSplit::Train, TuiGroup::Default); + + chart.push(tag.clone(), 15.0); + chart.push(tag.clone(), 10.0); + chart.push(tag.clone(), 14.0); + + assert_eq!(chart.axes.bounds_y[1], 15.); + chart.push(tag, 10.0); + assert_eq!(chart.axes.bounds_y[1], 14.); + } + + #[test] + fn test_push_update_bounds_min_y() { + let mut chart = RecentHistoryPlot::new(2); + let tag = TuiTag::new(TuiSplit::Train, TuiGroup::Default); + + chart.push(tag.clone(), 5.0); + chart.push(tag.clone(), 10.0); + chart.push(tag.clone(), 14.0); + + assert_eq!(chart.axes.bounds_y[0], 5.); + chart.push(tag, 10.0); + assert_eq!(chart.axes.bounds_y[0], 10.); + } +} diff --git a/crates/burn-train/src/renderer/tui/renderer.rs b/crates/burn-train/src/renderer/tui/renderer.rs new file mode 100644 index 0000000..ca1c9d9 --- /dev/null +++ b/crates/burn-train/src/renderer/tui/renderer.rs @@ -0,0 +1,640 @@ +use crate::logger::{EvaluationProgressLogger, ProgressSnapshot, TrainingProgressLogger}; +use crate::metric::{MetricDefinition, MetricId}; +use crate::renderer::tui::TuiSplit; +use crate::renderer::{EvaluationName, MetricState, MetricsRenderer, MetricsRendererEvaluation}; +use crate::renderer::{MetricsRendererTraining, tui::NumericMetricsState}; +use crate::{Interrupter, LearnerSummary}; +use burn_core::data::dataloader::Progress; +use ratatui::{ + Terminal, + crossterm::{ + event::{self, DisableMouseCapture, EnableMouseCapture, Event, KeyCode}, + execute, + terminal::{EnterAlternateScreen, LeaveAlternateScreen, disable_raw_mode, enable_raw_mode}, + }, + prelude::*, +}; +use std::collections::HashMap; +use std::panic::{set_hook, take_hook}; +use std::sync::mpsc::{Receiver, Sender}; +use std::sync::{Arc, Mutex, mpsc}; +use std::thread::JoinHandle; +use std::{ + error::Error, + io::{self, Stdout}, + time::{Duration, Instant}, +}; + +use super::{ + Callback, CallbackFn, ControlsView, MetricsView, PopupState, ProgressBarState, StatusState, + TextEventOutcome, TextMetricsState, TuiGroup, TuiTag, +}; + +/// The current terminal backend. +pub(crate) type TerminalBackend = CrosstermBackend; +/// The current terminal frame. +pub(crate) type TerminalFrame<'a> = ratatui::Frame<'a>; + +type PanicHook = Box) + 'static + Sync + Send>; + +const MAX_REFRESH_RATE_MILLIS: u64 = 100; + +enum TuiRendererEvent { + MetricRegistration(MetricDefinition), + MetricsUpdate((TuiSplit, TuiGroup, MetricState)), + StatusUpdateTrain((TuiSplit, ProgressSnapshot)), + StatusUpdateTest(ProgressSnapshot), + ProcessEnd { + summary: Option, + /// Interrupter reset. + reset: bool, + }, + CounterUpdate(String), + SplitEnd, + ManualClose, + Close, + Persistent, +} + +/// The terminal UI metrics renderer. +pub struct TuiMetricsRendererWrapper { + sender: mpsc::Sender, + interrupter: Interrupter, + handle_join: Option>, + kill_signal: Arc>>, + current_split: TuiSplit, + training_progress: ProgressSnapshot, + eval_progress: ProgressSnapshot, +} + +impl TuiMetricsRendererWrapper { + /// Create a new terminal UI renderer. + pub fn new(interrupter: Interrupter, checkpoint: Option) -> Self { + let (sender, receiver) = mpsc::channel(); + let (kill_signal_sender, kill_signal_receiver) = mpsc::channel(); + + let interrupter_clone = interrupter.clone(); + let handle_join = std::thread::Builder::new() + .name("train-renderer".into()) + .spawn(move || { + let mut renderer = + TuiMetricsRenderer::new(interrupter_clone, checkpoint, kill_signal_sender); + + let tick_rate = Duration::from_millis(MAX_REFRESH_RATE_MILLIS); + loop { + let remaining_time = tick_rate.saturating_sub(renderer.last_update.elapsed()); + match receiver.recv_timeout(remaining_time) { + Ok(event) => renderer.handle_event(event), + Err(mpsc::RecvTimeoutError::Timeout) => (), + Err(mpsc::RecvTimeoutError::Disconnected) => { + log::error!("Renderer thread disconnected."); + break; + } + } + + // Render + if renderer.last_update.elapsed() >= tick_rate + && let Err(err) = renderer.render() + { + log::error!("Render error: {err}"); + break; + } + + if (renderer.manual_close && renderer.interrupter.should_stop()) + || renderer.close + { + break; + } + } + }) + .unwrap(); + + let init = Progress::new(0, 0, None); + Self { + sender, + interrupter, + handle_join: Some(handle_join), + kill_signal: Arc::new(Mutex::new(kill_signal_receiver)), + current_split: TuiSplit::Train, + training_progress: ProgressSnapshot::new(init.clone(), init.clone()), + eval_progress: ProgressSnapshot::new(init.clone(), init), + } + } + + fn send_event(&self, event: TuiRendererEvent) { + if self.kill_signal.lock().unwrap().try_recv().is_ok() { + panic!("Killing training from user input.") + } + if let Err(e) = self.sender.send(event) { + log::warn!("Failed to send TUI event: {e}"); + } + } + + /// Set the renderer to persistent mode. + pub fn persistent(self) -> Self { + self.send_event(TuiRendererEvent::Persistent); + self + } +} + +struct TuiMetricsRenderer { + terminal: Terminal, + last_update: std::time::Instant, + progress: ProgressBarState, + metric_definitions: HashMap, + metrics_numeric: NumericMetricsState, + metrics_text: TextMetricsState, + status: StatusState, + interrupter: Interrupter, + popup: PopupState, + previous_panic_hook: Option>, + persistent: bool, + manual_close: bool, + close: bool, + summary: Option, + kill_signal: Sender<()>, +} + +impl MetricsRendererEvaluation for TuiMetricsRendererWrapper { + fn update_test(&mut self, name: EvaluationName, state: MetricState) { + self.send_event(TuiRendererEvent::MetricsUpdate(( + TuiSplit::Test, + TuiGroup::Named(name.name), + state, + ))); + } + + fn on_test_end(&mut self, summary: Option) -> Result<(), Box> { + // Update the summary + self.send_event(TuiRendererEvent::ProcessEnd { + summary, + reset: false, + }); + Ok(()) + } +} + +impl MetricsRenderer for TuiMetricsRendererWrapper { + fn manual_close(&mut self) { + self.send_event(TuiRendererEvent::ManualClose); + let _ = self.handle_join.take().unwrap().join(); + } + + fn register_metric(&mut self, definition: MetricDefinition) { + self.send_event(TuiRendererEvent::MetricRegistration(definition)); + } +} + +impl MetricsRendererTraining for TuiMetricsRendererWrapper { + fn update_train(&mut self, state: MetricState) { + self.send_event(TuiRendererEvent::MetricsUpdate(( + TuiSplit::Train, + TuiGroup::Default, + state, + ))); + } + + fn update_valid(&mut self, state: MetricState) { + self.send_event(TuiRendererEvent::MetricsUpdate(( + TuiSplit::Valid, + TuiGroup::Default, + state, + ))); + } + + fn on_train_end(&mut self, summary: Option) -> Result<(), Box> { + // Reset for following steps. + self.interrupter.reset(); + // Update the summary + self.send_event(TuiRendererEvent::ProcessEnd { + summary, + reset: true, + }); + Ok(()) + } +} + +impl TrainingProgressLogger for TuiMetricsRendererWrapper { + fn start(&mut self, total_epochs: usize, starting_epoch: usize, total_items: Option) { + self.training_progress.global = + Progress::new(starting_epoch, total_epochs, Some("epochs".to_string())); + if let Some(items) = total_items { + self.training_progress.split = Progress::new(0, items, Some("items".to_string())); + } + } + + fn update_epoch(&mut self, epoch: usize) { + let total = self.training_progress.global.items_total; + let unit = self.training_progress.global.unit.clone(); + self.training_progress.global = Progress::new(epoch + 1, total, unit); + } + + fn start_split(&mut self, split: &str, total_items: usize) { + self.training_progress.split = Progress::new(0, total_items, Some("items".to_string())); + self.current_split = if split == "train" { + TuiSplit::Train + } else { + TuiSplit::Valid + }; + } + + fn update_split(&mut self, items_processed: usize) { + let total = self.training_progress.split.items_total; + let unit = self.training_progress.split.unit.clone(); + self.training_progress.split = Progress::new(items_processed, total, unit); + if self.training_progress.global.items_total == 0 { + self.training_progress.global = self.training_progress.split.clone(); + } + self.send_event(TuiRendererEvent::StatusUpdateTrain(( + self.current_split, + self.training_progress.clone(), + ))); + } + + fn end_split(&mut self) { + self.send_event(TuiRendererEvent::SplitEnd); + self.current_split = TuiSplit::Train; + } + + fn end(&mut self) {} + + fn log_event_training(&mut self, event: String) { + self.send_event(TuiRendererEvent::CounterUpdate(event)); + } +} + +impl EvaluationProgressLogger for TuiMetricsRendererWrapper { + fn start_global_progress(&mut self, total_tests: usize) { + self.eval_progress.global = Progress::new(0, total_tests, Some("tests".to_string())); + } + + fn start_test(&mut self, _name: &str, total_items: usize) { + let current = self.eval_progress.global.items_processed + 1; + let total = self.eval_progress.global.items_total; + self.eval_progress.global = Progress::new(current, total, Some("tests".to_string())); + self.eval_progress.split = Progress::new(0, total_items, Some("items".to_string())); + } + + fn update_test_progress(&mut self, items_processed: usize) { + let total = self.eval_progress.split.items_total; + let unit = self.eval_progress.split.unit.clone(); + self.eval_progress.split = Progress::new(items_processed, total, unit); + self.send_event(TuiRendererEvent::StatusUpdateTest( + self.eval_progress.clone(), + )); + } + + fn end_test(&mut self) { + self.send_event(TuiRendererEvent::SplitEnd); + } + + fn end_global_progress(&mut self) {} + + fn log_event_evaluation(&mut self, event: String) { + self.send_event(TuiRendererEvent::CounterUpdate(event)); + } +} + +impl Drop for TuiMetricsRendererWrapper { + fn drop(&mut self) { + if !std::thread::panicking() { + self.send_event(TuiRendererEvent::Close); + if let Some(handle) = self.handle_join.take() { + let _ = handle.join(); + } + } + } +} + +impl TuiMetricsRenderer { + fn update_metric(&mut self, split: TuiSplit, group: TuiGroup, state: MetricState) { + match state { + MetricState::Generic(entry) => { + let name = self + .metric_definitions + .get(&entry.metric_id) + .unwrap() + .name + .clone() + .into(); + self.metrics_text.update(split, group, entry, name); + } + MetricState::Numeric(entry, value) => { + let name: Arc = self + .metric_definitions + .get(&entry.metric_id) + .unwrap() + .name + .clone() + .into(); + self.metrics_numeric + .push(TuiTag::new(split, group.clone()), name.clone(), value); + self.metrics_text.update(split, group, entry, name); + } + }; + } + + pub fn new( + interrupter: Interrupter, + checkpoint: Option, + kill_signal: Sender<()>, + ) -> Self { + let mut stdout = io::stdout(); + execute!(stdout, EnterAlternateScreen, EnableMouseCapture).unwrap(); + enable_raw_mode().unwrap(); + let terminal = Terminal::new(CrosstermBackend::new(stdout)).unwrap(); + + // Reset the terminal to raw mode on panic before running the panic handler + // This prevents that the panic message is not visible for the user. + let previous_panic_hook = Arc::new(take_hook()); + set_hook(Box::new({ + let previous_panic_hook = previous_panic_hook.clone(); + move |panic_info| { + let _ = disable_raw_mode(); + let _ = execute!(io::stdout(), DisableMouseCapture, LeaveAlternateScreen); + previous_panic_hook(panic_info); + } + })); + + Self { + terminal, + last_update: Instant::now(), + progress: ProgressBarState::new(checkpoint), + metric_definitions: HashMap::default(), + metrics_numeric: NumericMetricsState::default(), + metrics_text: TextMetricsState::default(), + status: StatusState::default(), + interrupter, + popup: PopupState::Empty, + previous_panic_hook: Some(previous_panic_hook), + persistent: false, + manual_close: false, + close: false, + summary: None, + kill_signal, + } + } + + fn handle_event(&mut self, event: TuiRendererEvent) { + match event { + TuiRendererEvent::MetricRegistration(definition) => { + self.metric_definitions + .insert(definition.metric_id.clone(), definition); + } + TuiRendererEvent::MetricsUpdate((split, group, state)) => { + self.update_metric(split, group, state); + } + TuiRendererEvent::StatusUpdateTrain((split, item)) => match split { + TuiSplit::Train => { + self.progress.update_train(&item); + self.metrics_numeric.update_progress_train(&item); + self.status.update_train(&item); + } + TuiSplit::Valid => { + self.progress.update_valid(&item); + self.metrics_numeric.update_progress_valid(&item); + self.status.update_valid(&item); + } + _ => (), + }, + TuiRendererEvent::StatusUpdateTest(item) => { + self.progress.update_test(&item); + self.metrics_numeric.update_progress_test(&item); + self.status.update_test(&item); + } + TuiRendererEvent::ProcessEnd { summary, reset } => { + match (self.summary.take(), summary) { + (None, Some(summary)) => { + self.summary = Some(summary); + } + (Some(current), Some(other)) => self.summary = Some(current.merge(other)), + (_, _) => { /* nothing to update */ } + } + + if reset { + self.interrupter.reset(); + } + } + TuiRendererEvent::CounterUpdate(event) => { + self.status.update_counter(event); + } + TuiRendererEvent::SplitEnd => { + self.status.reset_counters(); + } + TuiRendererEvent::ManualClose => self.manual_close = true, + TuiRendererEvent::Persistent => self.persistent = true, + TuiRendererEvent::Close => self.close = true, + } + } + + fn render(&mut self) -> Result<(), Box> { + self.draw()?; + self.handle_user_input()?; + + self.last_update = Instant::now(); + + Ok(()) + } + + fn draw(&mut self) -> Result<(), Box> { + self.terminal.draw(|frame| { + let size = frame.area(); + + match self.popup.view() { + Some(view) => view.render(frame, size), + None => { + let view = MetricsView::new( + self.metrics_numeric.view(), + self.metrics_text.view(), + self.progress.view(), + ControlsView, + self.status.view(), + ); + + view.render(frame, size); + } + }; + })?; + + Ok(()) + } + + /// Dispatch a single user event to the popup / numeric / text components. + /// Returns `true` when something visible changed and a redraw is warranted. + /// The training loop ignores this (it is tick-gated); the post-training loop + /// uses it to skip redraws on inert events like mouse jitter. + fn dispatch_user_event(&mut self, event: &Event) -> bool { + let mut redraw = self.popup.on_event(event); + if self.popup.is_empty() { + redraw |= self.metrics_numeric.on_event(event); + redraw |= match self.metrics_text.on_event(event) { + TextEventOutcome::Clicked(name) => { + self.metrics_numeric.select_by_name(&name); + true + } + TextEventOutcome::HoverChanged => true, + TextEventOutcome::Ignored => false, + }; + } + redraw + } + + fn handle_user_input(&mut self) -> Result<(), Box> { + while event::poll(Duration::from_secs(0))? { + let event = event::read()?; + let _ = self.dispatch_user_event(&event); + + if self.popup.is_empty() + && let Event::Key(key) = event + && let KeyCode::Char('q') = key.code + { + self.popup = PopupState::Full( + "Quit".to_string(), + vec![ + Callback::new( + "Stop the training.", + "Stop the training immediately. This will break from the \ + training loop, but any remaining code after the loop will be \ + executed.", + 's', + QuitPopupAccept(self.interrupter.clone()), + ), + Callback::new( + "Stop the training immediately.", + "Kill the program. This will create a panic! which will make \ + the current training fails. Any code following the training \ + won't be executed.", + 'k', + KillPopupAccept(self.kill_signal.clone()), + ), + Callback::new( + "Cancel", + "Cancel the action, continue the training.", + 'c', + PopupCancel, + ), + ], + ); + } + } + + Ok(()) + } + + fn handle_post_training(&mut self) -> Result<(), Box> { + self.popup = PopupState::Full( + "Training is done".to_string(), + vec![Callback::new( + "Training Done", + "Press 'x' to close this popup. Press 'q' to exit the application after the \ + popup is closed.", + 'x', + PopupCancel, + )], + ); + + self.draw().ok(); + + loop { + if let Ok(true) = event::poll(Duration::from_millis(MAX_REFRESH_RATE_MILLIS)) { + match event::read() { + Ok(event @ Event::Key(key)) => { + let redraw = self.dispatch_user_event(&event); + if self.popup.is_empty() + && let KeyCode::Char('q') = key.code + { + break; + } + if redraw { + self.draw().ok(); + } + } + + Ok(event @ Event::Mouse(_)) => { + if self.dispatch_user_event(&event) { + self.draw().ok(); + } + } + + Ok(Event::Resize(..)) => { + self.draw().ok(); + } + Err(err) => { + eprintln!("Error reading event: {err}"); + break; + } + _ => continue, + } + } + } + Ok(()) + } + + // Reset the terminal back to raw mode. + fn reset(&mut self) -> Result<(), Box> { + // If previous panic hook has already been re-instated, then the terminal was already reset. + if self.previous_panic_hook.is_some() { + if self.persistent + && let Err(err) = self.handle_post_training() + { + eprintln!("Error in post-training handling: {err}"); + } + + disable_raw_mode()?; + execute!( + self.terminal.backend_mut(), + DisableMouseCapture, + LeaveAlternateScreen + )?; + self.terminal.show_cursor()?; + + // Reinstall the previous panic hook + let _ = take_hook(); + if let Some(previous_panic_hook) = + Arc::into_inner(self.previous_panic_hook.take().unwrap()) + { + set_hook(previous_panic_hook); + } + } + Ok(()) + } +} + +struct QuitPopupAccept(Interrupter); +struct KillPopupAccept(Sender<()>); +struct PopupCancel; + +impl CallbackFn for KillPopupAccept { + fn call(&self) -> bool { + self.0.send(()).unwrap(); + panic!("Killing training from user input."); + } +} + +impl CallbackFn for QuitPopupAccept { + fn call(&self) -> bool { + self.0.stop(Some("Stopping training from user input.")); + true + } +} + +impl CallbackFn for PopupCancel { + fn call(&self) -> bool { + true + } +} + +impl Drop for TuiMetricsRenderer { + fn drop(&mut self) { + // Reset the terminal back to raw mode. This can be skipped during + // panicking because the panic hook has already reset the terminal + if !std::thread::panicking() { + self.reset().unwrap(); + + if let Some(summary) = &self.summary { + println!("{summary}"); + log::info!("{summary}"); + } + } + } +} diff --git a/crates/burn-train/src/renderer/tui/status.rs b/crates/burn-train/src/renderer/tui/status.rs new file mode 100644 index 0000000..6e3ed9f --- /dev/null +++ b/crates/burn-train/src/renderer/tui/status.rs @@ -0,0 +1,151 @@ +use std::collections::BTreeMap; + +use crate::logger::ProgressSnapshot; + +use super::TerminalFrame; +use ratatui::{ + prelude::{Alignment, Rect}, + style::{Color, Style, Stylize}, + text::{Line, Span}, + widgets::{Block, Borders, Paragraph, Wrap}, +}; + +/// Show the training status with various information. +pub(crate) struct StatusState { + progress: Option, + mode: Mode, + event_counters: BTreeMap, +} + +enum Mode { + Valid, + Train, + Evaluation, +} + +impl Default for StatusState { + fn default() -> Self { + Self { + progress: None, + mode: Mode::Train, + event_counters: BTreeMap::new(), + } + } +} + +impl StatusState { + /// Update the training information. + pub(crate) fn update_train(&mut self, progress: &ProgressSnapshot) { + self.progress = Some(progress.clone()); + self.mode = Mode::Train; + } + /// Update the validation information. + pub(crate) fn update_valid(&mut self, progress: &ProgressSnapshot) { + self.progress = Some(progress.clone()); + self.mode = Mode::Valid; + } + /// Update the testing information. + pub(crate) fn update_test(&mut self, progress: &ProgressSnapshot) { + self.progress = Some(progress.clone()); + self.mode = Mode::Evaluation; + } + /// Update counters from a progress event. + pub(crate) fn update_counter(&mut self, event: String) { + *self.event_counters.entry(event).or_insert(0) += 1; + } + + /// Reset all counters at the end of a split. + pub(crate) fn reset_counters(&mut self) { + for val in self.event_counters.values_mut() { + *val = 0; + } + } + + /// Create a view. + pub(crate) fn view(&self) -> StatusView { + StatusView::new(self.progress.as_ref(), &self.mode, &self.event_counters) + } +} + +pub(crate) struct StatusView { + lines: Vec>>, +} + +fn capitalize(s: &str) -> String { + let mut chars = s.chars(); + match chars.next() { + None => String::new(), + Some(first) => first.to_uppercase().collect::() + chars.as_str(), + } +} + +impl StatusView { + fn new( + progress: Option<&ProgressSnapshot>, + mode: &Mode, + event_counters: &BTreeMap, + ) -> Self { + let title = |title: &str| Span::from(format!(" {title} ")).bold().yellow(); + let value = |value: String| Span::from(value).italic(); + let mode_str = match mode { + Mode::Valid => "Validating", + Mode::Train => "Training", + Mode::Evaluation => "Evaluation", + }; + + let width = progress + .map(|p| { + p.global + .unit + .as_deref() + .map_or(0, |s| s.len()) + .max(p.split.unit.as_deref().map_or(0, |s| s.len())) + }) + .unwrap_or(0) + .max("Mode".len()) + .max(event_counters.keys().map(|k| k.len()).max().unwrap_or(0)); + + let mut lines = vec![vec![ + title(&format!("{: , size: Rect) { + let paragraph = Paragraph::new(self.lines.into_iter().map(Line::from).collect::>()) + .alignment(Alignment::Left) + .block(Block::default().borders(Borders::ALL).title("Status")) + .wrap(Wrap { trim: false }) + .style(Style::default().fg(Color::Gray)); + + frame.render_widget(paragraph, size); + } +} diff --git a/crates/burn-train/tests/checkpoint.rs b/crates/burn-train/tests/checkpoint.rs new file mode 100644 index 0000000..fb80179 --- /dev/null +++ b/crates/burn-train/tests/checkpoint.rs @@ -0,0 +1,261 @@ +//! Integration tests for checkpoint saving/loading in a minimal supervised training loop. + +mod common; +use std::fs; + +use common::*; + +use burn_core::tensor::Device; +use burn_train::{ + SupervisedTraining, checkpoint::KeepLastNCheckpoints, logger::InMemoryMetricLogger, + metric::LossMetric, +}; + +#[test] +fn checkpoint_saves_bpk_files() { + let dir = tempfile::tempdir().expect("create temp dir"); + let dir_path = dir.path().to_path_buf(); + let checkpoint_dir = dir_path.join("checkpoint"); + + let device = Device::flex().autodiff(); + let (dl_train, dl_valid) = make_dataloaders(); + + SupervisedTraining::new(&dir_path, dl_train, dl_valid) + .num_epochs(3) + .with_default_checkpointers() + .with_checkpointing_strategy(KeepLastNCheckpoints::new(3)) + .with_metric_logger(InMemoryMetricLogger::new()) + .metric_train_numeric(LossMetric::new()) + .metric_valid_numeric(LossMetric::new()) + .with_application_logger(None) + .launch(make_learner(&device)); + + for epoch in 1..=3 { + let model_path = checkpoint_dir.join(format!("model-{epoch}.bpk")); + assert!( + model_path.exists(), + "expected checkpoint file: {}", + model_path.display() + ); + } +} + +/// Training can resume from a saved checkpoint: `SupervisedTraining::checkpoint(epoch)` loads +/// the model/optimizer/scheduler and continues from the next epoch. +#[test] +fn checkpoint_resume_continues_training() { + let dir = tempfile::tempdir().expect("create temp dir"); + let dir_path = dir.path().to_path_buf(); + + let device = Device::flex().autodiff(); + + { + let (dl_train, dl_valid) = make_dataloaders(); + SupervisedTraining::new(&dir_path, dl_train, dl_valid) + .num_epochs(2) + .with_default_checkpointers() + .with_checkpointing_strategy(KeepLastNCheckpoints::new(2)) + .with_metric_logger(InMemoryMetricLogger::new()) + .metric_train_numeric(LossMetric::new()) + .metric_valid_numeric(LossMetric::new()) + .with_application_logger(None) + .launch(make_learner(&device)); + } + + let ckpt = dir_path.join("checkpoint").join("model-2.bpk"); + assert!(ckpt.exists(), "epoch-2 checkpoint must exist before resume"); + + // Resume from epoch 2 and train up to epoch 4. + { + let (dl_train, dl_valid) = make_dataloaders(); + SupervisedTraining::new(&dir_path, dl_train, dl_valid) + .num_epochs(4) + .checkpoint(2) + .with_default_checkpointers() + .with_checkpointing_strategy(KeepLastNCheckpoints::new(4)) + .with_metric_logger(InMemoryMetricLogger::new()) + .metric_train_numeric(LossMetric::new()) + .metric_valid_numeric(LossMetric::new()) + .with_application_logger(None) + .launch(make_learner(&device)); + } + + for epoch in 1..=4 { + let path = dir_path + .join("checkpoint") + .join(format!("model-{epoch}.bpk")); + assert!( + path.exists(), + "expected checkpoint file for epoch {epoch}: {}", + path.display() + ); + } +} + +#[test] +fn checkpoint_restores_model_weights() { + let dir = tempfile::tempdir().expect("create temp dir"); + let dir_path = dir.path().to_path_buf(); + + let device = Device::flex().autodiff(); + let (dl_train, dl_valid) = make_dataloaders(); + + let result = SupervisedTraining::new(&dir_path, dl_train, dl_valid) + .num_epochs(1) + .with_default_checkpointers() + .with_checkpointing_strategy(KeepLastNCheckpoints::new(1)) + .with_metric_logger(InMemoryMetricLogger::new()) + .metric_train_numeric(LossMetric::new()) + .metric_valid_numeric(LossMetric::new()) + .with_application_logger(None) + .launch(make_learner(&device)); + + let trained_weights = result + .model + .weight + .val() + .into_data() + .to_vec::() + .unwrap(); + + // Load the checkpoint into a fresh learner + use burn_train::{ + LearningCheckpointer, + checkpoint::{AsyncCheckpointer, FileCheckpointer}, + }; + + let ckpt_dir = dir_path.join("checkpoint"); + let checkpointer = LearningCheckpointer::new( + AsyncCheckpointer::new(FileCheckpointer::new(&ckpt_dir, "model")), + AsyncCheckpointer::new(FileCheckpointer::new(&ckpt_dir, "optim")), + AsyncCheckpointer::new(FileCheckpointer::new(&ckpt_dir, "scheduler")), + Box::new(KeepLastNCheckpoints::new(1)), + ); + + let fresh = make_learner(&device); + let restored = checkpointer.load_checkpoint(fresh, 1); + + let restored_weights = restored + .model() + .weight + .val() + .into_data() + .to_vec::() + .unwrap(); + + assert_eq!( + trained_weights, restored_weights, + "restored weights must match the saved checkpoint" + ); +} + +#[test] +fn file_metric_logger_creates_log_directories() { + let dir = tempfile::tempdir().expect("create temp dir"); + let dir_path = dir.path().to_path_buf(); + + let device = Device::flex().autodiff(); + let (dl_train, dl_valid) = make_dataloaders(); + + use burn_train::logger::FileMetricLogger; + + SupervisedTraining::new(&dir_path, dl_train, dl_valid) + .num_epochs(2) + .with_metric_logger(FileMetricLogger::new(&dir_path)) + .metric_train_numeric(LossMetric::new()) + .metric_valid_numeric(LossMetric::new()) + .with_application_logger(None) + .launch(make_learner(&device)); + + for epoch in 1..=2 { + let train_dir = dir_path.join("train").join(format!("epoch-{epoch}")); + let valid_dir = dir_path.join("valid").join(format!("epoch-{epoch}")); + + assert!( + train_dir.exists(), + "missing train log dir for epoch {epoch}: {}", + train_dir.display() + ); + assert!( + valid_dir.exists(), + "missing valid log dir for epoch {epoch}: {}", + valid_dir.display() + ); + } +} + +#[test] +fn file_metric_logger_resumes_logging_at_checkpoint() { + let dir = tempfile::tempdir().expect("create temp dir"); + let dir_path = dir.path().to_path_buf(); + + let device = Device::flex().autodiff(); + let (dl_train, dl_valid) = make_dataloaders(); + + use burn_train::logger::FileMetricLogger; + + { + SupervisedTraining::new(&dir_path, dl_train, dl_valid) + .num_epochs(1) + .with_metric_logger(FileMetricLogger::new(&dir_path)) + .metric_train_numeric(LossMetric::new()) + .metric_valid_numeric(LossMetric::new()) + .with_default_checkpointers() + .with_application_logger(None) + .launch(make_learner(&device)); + } + + let loss_file_train = dir_path.join("train").join("epoch-1/Loss.log"); + let loss_file_valid = dir_path.join("valid").join("epoch-1/Loss.log"); + + let content_train_expected = fs::read(loss_file_train).expect("Cannot read file."); + let content_valid_expected = fs::read(loss_file_valid).expect("Cannot read file."); + + // Resume from epoch 1 and train up to epoch 2. + { + let (dl_train, dl_valid) = make_dataloaders(); + SupervisedTraining::new(&dir_path, dl_train, dl_valid) + .num_epochs(2) + .checkpoint(1) + .with_default_checkpointers() + .with_metric_logger(FileMetricLogger::new(&dir_path)) + .metric_train_numeric(LossMetric::new()) + .metric_valid_numeric(LossMetric::new()) + .with_application_logger(None) + .launch(make_learner(&device)); + } + + for epoch in 1..=2 { + let train_dir = dir_path.join("train").join(format!("epoch-{epoch}")); + let valid_dir = dir_path.join("valid").join(format!("epoch-{epoch}")); + + assert!( + train_dir.exists(), + "missing train log dir for epoch {epoch}: {}", + train_dir.display() + ); + assert!( + valid_dir.exists(), + "missing valid log dir for epoch {epoch}: {}", + valid_dir.display() + ); + + // First epoch's logs remain unchanged. + if epoch == 1 { + let loss_file_train = train_dir.join("Loss.log"); + let loss_file_valid = valid_dir.join("Loss.log"); + + let content_train = fs::read(loss_file_train).expect("Cannot read file."); + let content_valid = fs::read(loss_file_valid).expect("Cannot read file."); + + assert_eq!( + content_train, content_train_expected, + "First epoch's logs should remain unchanged", + ); + assert_eq!( + content_valid, content_valid_expected, + "First epoch's logs should remain unchanged", + ); + } + } +} diff --git a/crates/burn-train/tests/common/mod.rs b/crates/burn-train/tests/common/mod.rs new file mode 100644 index 0000000..b1f9768 --- /dev/null +++ b/crates/burn-train/tests/common/mod.rs @@ -0,0 +1,164 @@ +use std::sync::Arc; + +use burn_core as burn; + +use burn_core::{ + data::{ + dataloader::{DataLoaderBuilder, batcher::Batcher}, + dataset::InMemDataset, + }, + module::{Module, Param}, + tensor::{Device, Distribution, Tensor}, +}; +use burn_optim::{ModuleOptimizer, SgdConfig, lr_scheduler::constant::ConstantLr}; +use burn_train::{InferenceStep, Learner, RegressionOutput, TrainOutput, TrainStep}; + +// Minimal toy model +#[derive(Module, Debug)] +pub struct ToyModel { + pub weight: Param>, +} + +impl ToyModel { + #[allow(unused)] + pub fn new(device: &Device) -> Self { + Self { + weight: Param::from_tensor(Tensor::random([1, 2], Distribution::Default, device)), + } + } +} + +#[derive(Clone, Debug)] +pub struct DummyBatch { + target: Tensor<2>, +} + +pub struct DummyBatcher; + +impl Batcher<(), DummyBatch> for DummyBatcher { + fn batch(&self, _items: Vec<()>, device: &Device) -> DummyBatch { + DummyBatch { + target: Tensor::zeros([1, 2], device), + } + } +} + +impl TrainStep for ToyModel { + type Input = DummyBatch; + type Output = RegressionOutput; + + fn step(&self, item: DummyBatch) -> TrainOutput { + let output = self.weight.val(); + let loss = output + .clone() + .sub(item.target.clone()) + .powi_scalar(2) + .mean(); + let regression = RegressionOutput::new(loss.clone(), output, item.target); + TrainOutput::new(self, loss.backward(), regression) + } +} + +impl InferenceStep for ToyModel { + type Input = DummyBatch; + type Output = RegressionOutput; + + fn step(&self, item: DummyBatch) -> RegressionOutput { + let output = self.weight.val(); + let loss = output + .clone() + .sub(item.target.clone()) + .powi_scalar(2) + .mean(); + RegressionOutput::new(loss, output, item.target) + } +} + +#[allow(unused)] +pub type ToyLearner = Learner; + +// Two-parameter model: `frozen` is a named sub-module so the optimizer mapper +// builds path "frozen" for its weight, making it targetable by ParamGroup::from_predicate. +#[allow(unused)] +#[derive(Module, Debug)] +pub struct FrozenLayer { + pub weight: Param>, +} + +impl FrozenLayer { + #[allow(unused)] + pub fn new(device: &Device) -> Self { + Self { + weight: Param::from_tensor(Tensor::random([1, 2], Distribution::Default, device)), + } + } +} + +#[allow(unused)] +#[derive(Module, Debug)] +pub struct TwoLayerModel { + pub frozen: FrozenLayer, + pub active: Param>, +} + +impl TwoLayerModel { + #[allow(unused)] + pub fn new(device: &Device) -> Self { + Self { + frozen: FrozenLayer::new(device), + active: Param::from_tensor(Tensor::random([1, 2], Distribution::Default, device)), + } + } +} + +impl TrainStep for TwoLayerModel { + type Input = DummyBatch; + type Output = RegressionOutput; + + fn step(&self, item: DummyBatch) -> TrainOutput { + let output = self.frozen.weight.val() + self.active.val(); + let loss = output + .clone() + .sub(item.target.clone()) + .powi_scalar(2) + .mean(); + let regression = RegressionOutput::new(loss.clone(), output, item.target); + TrainOutput::new(self, loss.backward(), regression) + } +} + +impl InferenceStep for TwoLayerModel { + type Input = DummyBatch; + type Output = RegressionOutput; + + fn step(&self, item: DummyBatch) -> RegressionOutput { + let output = self.frozen.weight.val() + self.active.val(); + let loss = output + .clone() + .sub(item.target.clone()) + .powi_scalar(2) + .mean(); + RegressionOutput::new(loss, output, item.target) + } +} + +pub fn make_dataloaders() -> ( + Arc>, + Arc>, +) { + let dl_train = DataLoaderBuilder::new(DummyBatcher) + .batch_size(2) + .build(InMemDataset::new(vec![(); 4])); + let dl_valid = DataLoaderBuilder::new(DummyBatcher) + .batch_size(2) + .build(InMemDataset::new(vec![(); 4])); + (dl_train, dl_valid) +} + +#[allow(unused)] +pub fn make_learner(device: &Device) -> ToyLearner { + let model = ToyModel::new(device); + let optim: ModuleOptimizer = SgdConfig::new().init(); + let scheduler = ConstantLr::new(1e-3); + Learner::new(model, optim, scheduler) +} diff --git a/crates/burn-train/tests/lr_policy.rs b/crates/burn-train/tests/lr_policy.rs new file mode 100644 index 0000000..42b3f61 --- /dev/null +++ b/crates/burn-train/tests/lr_policy.rs @@ -0,0 +1,70 @@ +//! Integration tests verifying that LrPolicyScheduler correctly freezes parameter groups. + +mod common; + +use common::*; + +use burn_core::{module::ParamGroup, tensor::Device}; +use burn_optim::{SgdConfig, lr_scheduler::module_lr_scheduler::ModuleLrSchedulerConfig}; +use burn_train::{Learner, SupervisedTraining, logger::InMemoryMetricLogger, metric::LossMetric}; + +/// Test the integration of [LrPolicyScheduler](burn_optim::lr_scheduler::policy::LrPolicyScheduler) in burn-train. +/// A parameter group with LR=0 must not change after training, while parameters using +/// the default LR must change (given a non-zero gradient). +#[test] +fn frozen_group_param_unchanged_after_training() { + let device = Device::flex().autodiff(); + let model = TwoLayerModel::new(&device); + + let before_frozen = model + .frozen + .weight + .val() + .into_data() + .to_vec::() + .unwrap(); + let before_active = model.active.val().into_data().to_vec::().unwrap(); + + let optim = SgdConfig::new().init(); + + let scheduler = ModuleLrSchedulerConfig::new(1e-2.into()) + .with_group(ParamGroup::from_predicate("frozen"), 0.0_f64) + .init() + .unwrap(); + + let (dl_train, dl_valid) = make_dataloaders(); + + let dir = tempfile::tempdir().unwrap(); + let result = SupervisedTraining::new(dir.path(), dl_train, dl_valid) + .num_epochs(1) + .with_metric_logger(InMemoryMetricLogger::new()) + .metric_train_numeric(LossMetric::new()) + .metric_valid_numeric(LossMetric::new()) + .with_application_logger(None) + .launch(Learner::new(model, optim, scheduler)); + + let after_frozen = result + .model + .frozen + .weight + .val() + .into_data() + .to_vec::() + .unwrap(); + let after_active = result + .model + .active + .val() + .into_data() + .to_vec::() + .unwrap(); + + assert_eq!( + before_frozen, after_frozen, + "frozen param (LR=0) must not change after training" + ); + assert_ne!( + before_active, after_active, + "active param (LR=1e-2) must change after training" + ); +} diff --git a/crates/burn-vision/Cargo.toml b/crates/burn-vision/Cargo.toml new file mode 100644 index 0000000..8273914 --- /dev/null +++ b/crates/burn-vision/Cargo.toml @@ -0,0 +1,75 @@ +[package] +authors = [ + "nathanielsimard ", + "wingertge ", +] +categories = ["science"] +description = "Vision processing operations for burn tensors" +documentation = "https://docs.rs/burn-vision" +edition.workspace = true +keywords = ["deep-learning", "machine-learning", "gpu"] +license.workspace = true +name = "burn-vision" +readme.workspace = true +repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-vision" +version.workspace = true + +[lints] +workspace = true + + +[features] +default = ["burn-core/default", "flex", "cubecl-backend", "fusion", "std"] +std = ["aligned-vec/std", "burn-core/std", "burn-store?/std"] +tracing = [ + "burn-cubecl?/tracing", + "burn-fusion?/tracing", + "burn-ir/tracing", + "burn-core/tracing", + "cubecl/tracing", +] + +cubecl-backend = ["cubecl", "burn-cubecl"] +fusion = ["burn-fusion", "burn-core/fusion"] + +cuda = ["cubecl-backend", "burn-core/cuda"] +rocm = ["cubecl-backend", "burn-core/rocm"] +wgpu = ["cubecl-backend", "burn-core/wgpu"] +vulkan = ["cubecl-backend", "burn-core/vulkan"] +webgpu = ["cubecl-backend", "burn-core/webgpu"] +metal = ["cubecl-backend", "burn-core/metal"] +cpu = ["cubecl-backend", "burn-core/cpu"] +flex = ["burn-core/flex"] +tch = ["burn-core/tch"] + +# TODO: move somewhere else +loss = ["std", "burn-store/pytorch", "burn-core/network", "dirs"] + +[dependencies] +burn-cubecl = { workspace = true, optional = true, features = ["default"]} +burn-fusion = { workspace = true, optional = true, features = ["default"]} +burn-ir = { workspace = true, features = ["default"]} +burn-core = { workspace = true, features = [ + "extension", +] } + +# TODO: requires burn-backend. The `extension` feature should expose this from `burn-core`. + +aligned-vec = { version = "0.6", default-features = false } +bon = { workspace = true } + +# For loss functions requiring pretrained models (e.g., Gram Matrix Loss) +burn-store = { workspace = true, optional = true } +dirs = { workspace = true, optional = true } + +bytemuck = { workspace = true } +cubecl = { workspace = true, optional = true } +derive-new = { workspace = true } +half = { workspace = true } +image = { version = "0.25" } +macerator = { workspace = true } +ndarray = { workspace = true } +num-traits = { workspace = true } +paste = { workspace = true } +serde = { workspace = true } + diff --git a/crates/burn-vision/src/backends/cpu/base.rs b/crates/burn-vision/src/backends/cpu/base.rs new file mode 100644 index 0000000..c295360 --- /dev/null +++ b/crates/burn-vision/src/backends/cpu/base.rs @@ -0,0 +1,42 @@ +pub trait MinMax { + fn min(self, other: Self) -> Self; + fn max(self, other: Self) -> Self; +} + +macro_rules! impl_minmax { + ($ty: ty) => { + impl MinMax for $ty { + fn min(self, other: Self) -> Self { + Ord::min(self, other) + } + fn max(self, other: Self) -> Self { + Ord::max(self, other) + } + } + }; + ($($ty: ty),*) => { + $(impl_minmax!($ty);)* + } +} + +impl_minmax!(u8, i8, u16, i16, u32, i32, u64, i64); + +impl MinMax for f32 { + fn min(self, other: Self) -> Self { + self.min(other) + } + + fn max(self, other: Self) -> Self { + self.max(other) + } +} + +impl MinMax for f64 { + fn min(self, other: Self) -> Self { + self.min(other) + } + + fn max(self, other: Self) -> Self { + self.max(other) + } +} diff --git a/crates/burn-vision/src/backends/cpu/connected_components.rs b/crates/burn-vision/src/backends/cpu/connected_components.rs new file mode 100644 index 0000000..8aac8ac --- /dev/null +++ b/crates/burn-vision/src/backends/cpu/connected_components.rs @@ -0,0 +1,249 @@ +use std::{cmp::Ordering, marker::PhantomData}; + +use alloc::vec::Vec; +use burn_core::backend::{ + Backend, TensorMetadata, tensor::{BoolTensor, Device} +}; +use burn_core::tensor::{ + Element, ElementConversion, ElementLimits, ElementOrdered, IntDType, Shape, TensorData, + read_sync, +}; +use ndarray::Array2; + +use crate::{ + ConnectedStatsOptions, ConnectedStatsPrimitive, Connectivity, dispatch_bool_dtype, + dispatch_int_dtype, +}; + +mod spaghetti; +mod spaghetti_4c; + +pub fn connected_components( + img: BoolTensor, + connectivity: Connectivity, + out_dtype: IntDType, +) -> TensorData { + let img = read_sync(B::bool_into_data(img)).expect("Should read data."); + dispatch_bool_dtype!(img.dtype.into(), |B| { + dispatch_int_dtype!(out_dtype, |I| run::>( + img, + connectivity, + NoOp::default + ) + .0) + }) +} + +pub fn connected_components_with_stats( + img: BoolTensor, + connectivity: Connectivity, + _options: ConnectedStatsOptions, + out_dtype: IntDType, +) -> (TensorData, ConnectedStatsPrimitive) { + let device = &img.device(); + let img = read_sync(B::bool_into_data(img)).expect("Should read data."); + dispatch_bool_dtype!(img.dtype.into(), |BT| { + dispatch_int_dtype!(out_dtype, |I| { + let (labels, stats) = + run::>(img, connectivity, ConnectedStatsOp::default); + let stats = finalize_stats::(device, stats); + (labels, stats) + }) + }) +} + +fn run>( + img: TensorData, + connectivity: Connectivity, + stats: impl Fn() -> Stats, +) -> (TensorData, Stats) { + let [height, width] = img.shape.dims(); + let img = img.into_vec::().unwrap(); + + let mut stats = stats(); + + let out = match connectivity { + Connectivity::Four => { + spaghetti_4c::process::>(img, height, width, &mut stats) + } + Connectivity::Eight => { + // SAFETY: This is validated by `TensorData` + let img = unsafe { Array2::from_shape_vec_unchecked((height, width), img) }; + spaghetti::process::>(img, &mut stats) + } + }; + + let (data, _) = out.into_raw_vec_and_offset(); + let labels = TensorData::new(data, Shape::new([height, width])); + (labels, stats) +} + +pub trait Solver { + type Label: ElementOrdered; + + fn init(max_labels: usize) -> Self; + /// Hack to get around mutable borrow limitations on methods + fn merge(label_1: Self::Label, label_2: Self::Label, solver: &mut Self) -> Self::Label; + fn new_label(&mut self) -> Self::Label; + fn flatten(&mut self) -> Self::Label; + fn get_label(&self, i_label: Self::Label) -> Self::Label; +} + +pub(crate) struct UnionFind { + labels: Vec, +} + +impl Solver for UnionFind { + type Label = I; + + fn init(max_labels: usize) -> Self { + let mut labels = Vec::with_capacity(max_labels); + labels.push(0.elem()); + Self { labels } + } + + fn merge(mut label_1: I, mut label_2: I, solver: &mut Self) -> I { + use Ordering::Less; + + while matches!(solver.labels[label_1.to_usize()].cmp(&label_1), Less) { + label_1 = solver.labels[label_1.to_usize()]; + } + + while matches!(solver.labels[label_2.to_usize()].cmp(&label_2), Less) { + label_2 = solver.labels[label_2.to_usize()]; + } + + if matches!(label_1.cmp(&label_2), Less) { + solver.labels[label_2.to_usize()] = label_1; + label_1 + } else { + solver.labels[label_1.to_usize()] = label_2; + label_2 + } + } + + fn new_label(&mut self) -> I { + let len = I::from_elem(self.labels.len()); + self.labels.push(len); + len + } + + fn flatten(&mut self) -> I { + let mut k = 1; + for i in 1..self.labels.len() { + if matches!(self.labels[i].cmp(&I::from_elem(i)), Ordering::Less) { + self.labels[i] = self.labels[self.labels[i].to_usize()]; + } else { + self.labels[i] = k.elem(); + k += 1; + } + } + k.elem() + } + + fn get_label(&self, i_label: I) -> I { + self.labels[i_label.to_usize()] + } +} + +pub trait StatsOp { + type Label; + + fn init(&mut self, num_labels: usize); + fn update(&mut self, row: usize, column: usize, label: Self::Label); + fn finish(&mut self); +} + +#[derive(Default)] +struct NoOp { + _i: PhantomData, +} + +impl StatsOp for NoOp { + type Label = I; // placeholder still required + + fn init(&mut self, _num_labels: usize) {} + + fn update(&mut self, _row: usize, _column: usize, _label: Self::Label) {} + + fn finish(&mut self) {} +} + +#[derive(Default, Debug)] +struct ConnectedStatsOp { + pub area: Vec, + pub left: Vec, + pub top: Vec, + pub right: Vec, + pub bottom: Vec, +} + +impl StatsOp for ConnectedStatsOp { + type Label = I; + + fn init(&mut self, num_labels: usize) { + self.area = vec![0.elem(); num_labels]; + self.left = vec![I::MAX; num_labels]; + self.top = vec![I::MAX; num_labels]; + self.right = vec![0.elem(); num_labels]; + self.bottom = vec![0.elem(); num_labels]; + } + + fn update(&mut self, row: usize, column: usize, label: I) { + let l = label.to_usize(); + unsafe { + *self.area.get_unchecked_mut(l) = + I::from_elem((*self.area.get_unchecked(l)).to_usize() + 1); + *self.left.get_unchecked_mut(l) = + I::from_elem((*self.left.get_unchecked(l)).to_usize().min(column)); + *self.top.get_unchecked_mut(l) = + I::from_elem((*self.top.get_unchecked(l)).to_usize().min(row)); + *self.right.get_unchecked_mut(l) = + I::from_elem((*self.right.get_unchecked(l)).to_usize().max(column)); + *self.bottom.get_unchecked_mut(l) = + I::from_elem((*self.bottom.get_unchecked(l)).to_usize().max(row)); + } + } + + fn finish(&mut self) { + // Background shouldn't have stats + self.area[0] = 0.elem(); + self.left[0] = 0.elem(); + self.right[0] = 0.elem(); + self.top[0] = 0.elem(); + self.bottom[0] = 0.elem(); + } +} + +fn finalize_stats( + device: &Device, + stats: ConnectedStatsOp, +) -> ConnectedStatsPrimitive { + let labels = stats.area.len(); + + let into_prim = |data: Vec| { + let data = TensorData::new(data, Shape::new([labels])); + B::int_from_data(data, device) + }; + + let max_label = { + let data = TensorData::new(vec![I::from_elem(labels - 1)], Shape::new([1])); + B::int_from_data(data, device) + }; + + ConnectedStatsPrimitive { + area: into_prim(stats.area), + left: into_prim(stats.left), + top: into_prim(stats.top), + right: into_prim(stats.right), + bottom: into_prim(stats.bottom), + max_label, + } +} + +pub fn max_labels(h: usize, w: usize, conn: Connectivity) -> usize { + match conn { + Connectivity::Four => (h * w).div_ceil(2) + 1, + Connectivity::Eight => h.div_ceil(2) * w.div_ceil(2) + 1, + } +} diff --git a/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/Spaghetti_center_line_forest_code.rs b/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/Spaghetti_center_line_forest_code.rs new file mode 100644 index 0000000..6b6a796 --- /dev/null +++ b/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/Spaghetti_center_line_forest_code.rs @@ -0,0 +1,1954 @@ +no_analyze!{{ +use centerLabels::*;let mut label = entry; +while let Some(next) = (|label| -> Option { match label { + NODE_1=> { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_2); + } + else { + return Some(NODE_3); + } + } + NODE_3=> { + if (*img_row01.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(cl_tree_2); + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + return Some(cl_tree_1); + } + } + NODE_4=> { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + return Some(NODE_5); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(cl_tree_5); + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_4); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(cl_tree_3); + } + } + } + NODE_6=> { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_2); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(cl_tree_7); + } + } + else { + return Some(NODE_1); + } + } + NODE_2=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_6); + } + else { + return Some(NODE_4); + } + } + NODE_7=> { + if (*img_row12.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(cl_tree_5); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row12.add((c + 2) as usize), *img_labels_row12.add((c - 2) as usize), solver); + return Some(cl_tree_5); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row12.add((c + 2) as usize), *img_labels_row12.add((c - 2) as usize), solver); + return Some(cl_tree_5); + } + } + NODE_5=> { + if (*img_row12.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(cl_tree_5); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row12.add((c + 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_5); + } + } + NODE_8=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_6); + } + else { + return Some(NODE_9); + } + } + NODE_10=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_11); + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_11); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver), *img_labels_row12.add((c - 2) as usize), solver); + return Some(cl_tree_11); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_11); + } + } + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + if (*img_row12.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_5); + } + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + if (*img_row12.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver), *img_labels_row12.add((c - 2) as usize), solver); + return Some(cl_tree_5); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver), *img_labels_row12.add((c - 2) as usize), solver); + return Some(cl_tree_5); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + } + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c - 2) as usize), solver); + return Some(cl_tree_8); + } + else { + return Some(NODE_11); + } + } + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c - 2) as usize), solver); + return Some(cl_tree_12); + } + else { + return Some(NODE_12); + } + } + } + } + NODE_11=> { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_4); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_3); + } + } + NODE_13=> { + if (*img_row12.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_11); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row12.add((c) as usize), *img_labels_row12.add((c - 2) as usize), solver); + return Some(cl_tree_11); + } + } + NODE_9=> { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row12.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_5); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + } + } + else { + return Some(NODE_11); + } + } + NODE_12=> { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_10); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_9); + } + } + NODE_14=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_11); + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_4); + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_10); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(cl_tree_9); + } + } + } + } + NODE_15=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_11); + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + return Some(NODE_13); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_11); + } + } + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + return Some(NODE_5); + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + return Some(NODE_7); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(cl_tree_5); + } + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_4); + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c - 2) as usize); + return Some(cl_tree_3); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(cl_tree_3); + } + } + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_10); + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c - 2) as usize); + return Some(cl_tree_9); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(cl_tree_9); + } + } + } + } + } + NODE_16=> { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row01.add((c - 1) as usize)).to_bool() { + return Some(NODE_8); + } + else { + return Some(NODE_2); + } + } + else { + return Some(NODE_17); + } + } + else { + return Some(NODE_1); + } + } + NODE_18=> { + if (*img_row12.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(cl_tree_5); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + } + NODE_19=> { + if (*img_row11.add((c + 2) as usize)).to_bool() { + return Some(NODE_20); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_8); + } + } + NODE_21=> { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_6); + } + else { + if (*img_row11.add((c + 2) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(cl_tree_5); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(cl_tree_3); + } + } + } + else { + return Some(NODE_3); + } + } + NODE_22=> { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_6); + } + else { + if (*img_row12.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_6); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_6); + } + } + } + NODE_23=> { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_11); + } + else { + if (*img_row12.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_11); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_11); + } + } + } + NODE_24=> { + if (*img_row12.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_6); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_6); + } + } + NODE_17=> { + if (*img_row01.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_7); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(cl_tree_7); + } + } + NODE_25=> { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row12.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + return Some(NODE_18); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_8); + } + } + NODE_20=> { + if (*img_row12.add((c + 1) as usize)).to_bool() { + return Some(NODE_26); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + } + NODE_27=> { + if (*img_row12.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_11); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_11); + } + } + NODE_28=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + return Some(NODE_22); + } + else { + return Some(NODE_19); + } + } + NODE_26=> { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(cl_tree_5); + } + else { + if (*img_row12.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(cl_tree_5); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + } + } + NODE_29=> { + if (*img_row11.add((c + 2) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_8); + } + } + NODE_30=> { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row12.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(cl_tree_5); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_8); + } + } + NODE_31=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + return Some(NODE_23); + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_19); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_12); + } + } + } + NODE_32=> { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row12.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c - 2) as usize)).to_bool() { + return Some(NODE_33); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + if (*img_row11.add((c - 2) as usize)).to_bool() { + return Some(NODE_34); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_5); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + if (*img_row11.add((c - 2) as usize)).to_bool() { + return Some(NODE_35); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_4); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_3); + } + } + } + NODE_36=> { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row12.add((c + 1) as usize)).to_bool() { + return Some(NODE_33); + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + return Some(NODE_34); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + return Some(NODE_35); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_3); + } + } + } + NODE_37=> { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row12.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + if (*img_row11.add((c - 2) as usize)).to_bool() { + return Some(NODE_18); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_8); + } + } + NODE_33=> { + if (*img_row12.add((c - 1) as usize)).to_bool() { + return Some(NODE_26); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(cl_tree_5); + } + } + NODE_38=> { + if (*img_row12.add((c - 1) as usize)).to_bool() { + return Some(NODE_22); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_6); + } + } + NODE_39=> { + if (*img_row12.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_10); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_10); + } + } + NODE_35=> { + if (*img_row12.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_4); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_4); + } + } + NODE_40=> { + if (*img_row12.add((c - 1) as usize)).to_bool() { + return Some(NODE_23); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_11); + } + } + NODE_34=> { + if (*img_row12.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row12.add((c + 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_5); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_5); + } + } +cl_tree_0 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(cl_break_0_0); } else { return Some(cl_break_1_0); } } + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_14); + } + else { + return Some(NODE_6); + } +} +cl_tree_1 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(cl_break_0_1); } else { return Some(cl_break_1_1); } } + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_15); + } + else { + return Some(NODE_6); + } +} +cl_tree_2 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(cl_break_0_2); } else { return Some(cl_break_1_2); } } + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_10); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_8); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_7); + } + } + else { + return Some(NODE_1); + } + } +} +cl_tree_3 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(cl_break_0_3); } else { return Some(cl_break_1_3); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_11); + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_29); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_12); + } + } + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_6); + } + else { + return Some(NODE_29); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_7); + } + } + else { + return Some(NODE_21); + } + } +} +cl_tree_4 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(cl_break_0_3); } else { return Some(cl_break_1_4); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + return Some(NODE_27); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_11); + } + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_25); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_12); + } + } + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + return Some(NODE_24); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_6); + } + } + else { + return Some(NODE_25); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_7); + } + } + else { + return Some(NODE_21); + } + } +} +cl_tree_5 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(cl_break_0_3); } else { return Some(cl_break_1_5); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_11); + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_30); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_12); + } + } + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_6); + } + else { + return Some(NODE_30); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_7); + } + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_6); + } + else { + if (*img_row11.add((c + 2) as usize)).to_bool() { + return Some(NODE_5); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_4); + } + } + } + else { + return Some(NODE_3); + } + } + } +} +cl_tree_6 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(cl_break_0_3); } else { return Some(cl_break_1_6); } } + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_31); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_28); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_7); + } + } + else { + return Some(NODE_1); + } + } +} +cl_tree_7 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(cl_break_0_4); } else { return Some(cl_break_1_7); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row01.add((c - 1) as usize)).to_bool() { + return Some(NODE_10); + } + else { + return Some(NODE_15); + } + } + else { + return Some(NODE_16); + } +} +cl_tree_8 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(cl_break_0_3); } else { return Some(cl_break_1_8); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + if (*img_row11.add((c - 2) as usize)).to_bool() { + return Some(NODE_27); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_11); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_11); + } + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_37); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_12); + } + } + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + if (*img_row11.add((c - 2) as usize)).to_bool() { + return Some(NODE_24); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_6); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_6); + } + } + else { + return Some(NODE_37); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_7); + } + } + else { + return Some(NODE_21); + } + } +} +cl_tree_9 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(cl_break_0_5); } else { return Some(cl_break_1_9); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row01.add((c - 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_11); + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_9); + } + else { + return Some(NODE_12); + } + } + } + else { + return Some(NODE_14); + } + } + else { + return Some(NODE_16); + } +} +cl_tree_10 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(cl_break_0_6); } else { return Some(cl_break_1_10); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row01.add((c - 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + return Some(NODE_40); + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_36); + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + return Some(NODE_39); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_9); + } + } + } + } + else { + return Some(NODE_14); + } + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row01.add((c - 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + return Some(NODE_38); + } + else { + return Some(NODE_36); + } + } + else { + return Some(NODE_2); + } + } + else { + return Some(NODE_17); + } + } + else { + return Some(NODE_1); + } + } +} +cl_tree_11 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(cl_break_0_7); } else { return Some(cl_break_1_11); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row00.add((c - 1) as usize)).to_bool() { + return Some(NODE_31); + } + else { + if (*img_row01.add((c - 1) as usize)).to_bool() { + return Some(NODE_31); + } + else { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_11); + } + else { + return Some(NODE_13); + } + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + return Some(NODE_5); + } + else { + return Some(NODE_7); + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_4); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c - 2) as usize); + return Some(cl_tree_3); + } + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_10); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c - 2) as usize); + return Some(cl_tree_9); + } + } + } + } + } + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row00.add((c - 1) as usize)).to_bool() { + return Some(NODE_28); + } + else { + if (*img_row01.add((c - 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + return Some(NODE_22); + } + else { + if (*img_row11.add((c + 2) as usize)).to_bool() { + return Some(NODE_20); + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(cl_tree_4); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_3); + } + } + } + } + else { + return Some(NODE_2); + } + } + } + else { + if (*img_row01.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_7); + } + else { + if (*img_row00.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_7); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(cl_tree_7); + } + } + } + } + else { + return Some(NODE_1); + } + } +} +cl_tree_12 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(cl_break_0_8); } else { return Some(cl_break_1_12); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row01.add((c - 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c - 2) as usize)).to_bool() { + return Some(NODE_40); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_11); + } + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_32); + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + if (*img_row11.add((c - 2) as usize)).to_bool() { + return Some(NODE_39); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_10); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(cl_tree_9); + } + } + } + } + else { + return Some(NODE_14); + } + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row01.add((c - 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c - 2) as usize)).to_bool() { + return Some(NODE_38); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(cl_tree_6); + } + } + else { + return Some(NODE_32); + } + } + else { + return Some(NODE_2); + } + } + else { + return Some(NODE_17); + } + } + else { + return Some(NODE_1); + } + } +} + NODE_41=> { + if (*img_row11.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c - 2) as usize), solver); + } + else { + return Some(NODE_42); + } + } + NODE_43=> { + if (*img_row01.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + } + NODE_42=> { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + NODE_44=> { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c - 2) as usize); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } + } + NODE_45=> { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row01.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + } + NODE_46=> { + if (*img_row12.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + } + NODE_47=> { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } + NODE_48=> { + if (*img_row01.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + } +cl_break_0_0 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_47); + } + else { + return Some(NODE_48); + } + return None;} +cl_break_0_1 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_44); + } + else { + return Some(NODE_48); + } + return None;} +cl_break_0_2 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_41); + } + else { + return Some(NODE_43); + } + return None;} +cl_break_0_3 => { + if (*img_row00.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + return Some(NODE_43); + } + return None;} +cl_break_0_4 => { + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row01.add((c - 1) as usize)).to_bool() { + return Some(NODE_41); + } + else { + return Some(NODE_44); + } + } + else { + return Some(NODE_45); + } + return None;} +cl_break_0_5 => { + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row01.add((c - 1) as usize)).to_bool() { + return Some(NODE_42); + } + else { + return Some(NODE_47); + } + } + else { + return Some(NODE_45); + } + return None;} +cl_break_0_6 => { + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row01.add((c - 1) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + return Some(NODE_46); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + else { + return Some(NODE_47); + } + } + else { + return Some(NODE_45); + } + return None;} +cl_break_0_7 => { + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row00.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + if (*img_row01.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c - 2) as usize); + } + } + } + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row01.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + if (*img_row00.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + } + return None;} +cl_break_0_8 => { + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row01.add((c - 1) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + if (*img_row11.add((c - 2) as usize)).to_bool() { + return Some(NODE_46); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + else { + return Some(NODE_47); + } + } + else { + return Some(NODE_45); + } + return None;} + NODE_49=> { + if (*img_row12.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + } + NODE_50=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + return Some(NODE_49); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + NODE_51=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + return Some(NODE_52); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + NODE_52=> { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + if (*img_row12.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + } + } + NODE_53=> { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_54); + } + else { + return Some(NODE_55); + } + } + else { + return Some(NODE_56); + } + } + NODE_55=> { + if (*img_row01.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } + NODE_54=> { + if (*img_row01.add((c - 1) as usize)).to_bool() { + return Some(NODE_57); + } + else { + return Some(NODE_58); + } + } + NODE_59=> { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } + else { + return Some(NODE_60); + } + } + NODE_61=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + NODE_62=> { + if (*img_row01.add((c - 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + return Some(NODE_63); + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + return Some(NODE_49); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + } + else { + return Some(NODE_58); + } + } + NODE_63=> { + if (*img_row12.add((c - 1) as usize)).to_bool() { + return Some(NODE_52); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + } + NODE_64=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + return Some(NODE_65); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + NODE_65=> { + if (*img_row11.add((c - 2) as usize)).to_bool() { + return Some(NODE_49); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + } + NODE_66=> { + if (*img_row01.add((c - 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c - 2) as usize)).to_bool() { + return Some(NODE_63); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + return Some(NODE_65); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + } + else { + return Some(NODE_58); + } + } + NODE_67=> { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_58); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } + else { + return Some(NODE_56); + } + } + NODE_56=> { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_58); + } + else { + return Some(NODE_60); + } + } + NODE_58=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } + } + NODE_68=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver), *img_labels_row12.add((c - 2) as usize), solver); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + } + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c - 2) as usize), solver); + } + else { + return Some(NODE_69); + } + } + } + NODE_70=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + return Some(NODE_71); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c - 2) as usize); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } + } + } + NODE_57=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + else { + return Some(NODE_69); + } + } + NODE_60=> { + if (*img_row01.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + } + NODE_71=> { + if (*img_row12.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row12.add((c) as usize), *img_labels_row12.add((c - 2) as usize), solver); + } + } + NODE_69=> { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } +cl_break_1_0 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_58); + } + else { + return Some(NODE_67); + } + return None;} +cl_break_1_1 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_70); + } + else { + return Some(NODE_67); + } + return None;} +cl_break_1_2 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_68); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_57); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + else { + return Some(NODE_56); + } + } + return None;} +cl_break_1_3 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_61); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_61); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + else { + return Some(NODE_59); + } + } + return None;} +cl_break_1_4 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_50); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_50); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + else { + return Some(NODE_59); + } + } + return None;} +cl_break_1_5 => { + if (*img_row00.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + return Some(NODE_60); + } + } + } + return None;} +cl_break_1_6 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_51); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_51); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + else { + return Some(NODE_56); + } + } + return None;} +cl_break_1_7 => { + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row01.add((c - 1) as usize)).to_bool() { + return Some(NODE_68); + } + else { + return Some(NODE_70); + } + } + else { + return Some(NODE_53); + } + return None;} +cl_break_1_8 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_64); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_64); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + else { + return Some(NODE_59); + } + } + return None;} +cl_break_1_9 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_54); + } + else { + return Some(NODE_53); + } + return None;} +cl_break_1_10 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_62); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_62); + } + else { + return Some(NODE_55); + } + } + else { + return Some(NODE_56); + } + } + return None;} +cl_break_1_11 => { + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row00.add((c - 1) as usize)).to_bool() { + return Some(NODE_51); + } + else { + if (*img_row01.add((c - 1) as usize)).to_bool() { + return Some(NODE_51); + } + else { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + return Some(NODE_71); + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c - 2) as usize); + } + } + } + } + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row00.add((c - 1) as usize)).to_bool() { + return Some(NODE_51); + } + else { + if (*img_row01.add((c - 1) as usize)).to_bool() { + return Some(NODE_51); + } + else { + return Some(NODE_58); + } + } + } + else { + if (*img_row01.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + if (*img_row00.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } + } + } + else { + return Some(NODE_56); + } + } + return None;} +cl_break_1_12 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_66); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_66); + } + else { + return Some(NODE_55); + } + } + else { + return Some(NODE_56); + } + } + return None;} + }; None})(label) +{ +label = next; +} +}} diff --git a/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/Spaghetti_first_line_forest_code.rs b/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/Spaghetti_first_line_forest_code.rs new file mode 100644 index 0000000..6524b28 --- /dev/null +++ b/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/Spaghetti_first_line_forest_code.rs @@ -0,0 +1,223 @@ +no_analyze!{{ +use firstLabels::*;let mut label = entry; +while let Some(next) = (|label| -> Option { match label { + NODE_72=> { + if (*img_row00.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(fl_tree_1); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(fl_tree_2); + } + } + NODE_73=> { + if (*img_row00.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(fl_tree_1); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(fl_tree_2); + } + } + NODE_74=> { + if (*img_row00.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(fl_tree_1); + } + else { + if (*img_row01.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(fl_tree_1); + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + return Some(fl_tree_0); + } + } + } +fl_tree_0 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(fl_break_0_0); } else { return Some(fl_break_1_0); } } + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_72); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + return Some(NODE_72); + } + else { + return Some(NODE_74); + } + } +} +fl_tree_1 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(fl_break_0_1); } else { return Some(fl_break_1_1); } } + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_73); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + return Some(NODE_73); + } + else { + return Some(NODE_74); + } + } +} +fl_tree_2 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(fl_break_0_2); } else { return Some(fl_break_1_2); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row01.add((c - 1) as usize)).to_bool() { + return Some(NODE_73); + } + else { + return Some(NODE_72); + } + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row01.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(fl_tree_1); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(fl_tree_1); + } + } + else { + if (*img_row01.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(fl_tree_2); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(fl_tree_2); + } + } + } + else { + return Some(NODE_74); + } + } +} + NODE_75=> { + if (*img_row01.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } +fl_break_0_0 => { + if (*img_row00.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + } + return None;} +fl_break_0_1 => { + if (*img_row00.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + } + return None;} +fl_break_0_2 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_75); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + return Some(NODE_75); + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + } + return None;} + NODE_76=> { + if (*img_row00.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + else { + if (*img_row01.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + } + } + NODE_77=> { + if (*img_row01.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } +fl_break_1_0 => { + if (*img_row00.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + else { + return Some(NODE_76); + } + } + return None;} +fl_break_1_1 => { + if (*img_row00.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + return Some(NODE_76); + } + } + return None;} +fl_break_1_2 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_77); + } + else { + if (*img_row01.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_77); + } + else { + return Some(NODE_77); + } + } + else { + return Some(NODE_76); + } + } + return None;} +fl_ => {}, + }; None})(label) +{ +label = next; +} +}} diff --git a/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/Spaghetti_forest_labels.rs b/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/Spaghetti_forest_labels.rs new file mode 100644 index 0000000..6c994fc --- /dev/null +++ b/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/Spaghetti_forest_labels.rs @@ -0,0 +1,191 @@ +/// Workaround for rust-analyzer bug that causes invalid errors on the `include!`. +macro_rules! no_analyze { + ($tokens:tt) => { + $tokens + }; +} + +pub(crate) use no_analyze; + +#[allow(non_snake_case, non_camel_case_types, unused)] +pub enum centerLabels { + NODE_1, + NODE_2, + NODE_3, + NODE_4, + NODE_5, + NODE_6, + NODE_7, + NODE_8, + NODE_9, + NODE_10, + NODE_11, + NODE_12, + NODE_13, + NODE_14, + NODE_15, + NODE_16, + NODE_17, + NODE_18, + NODE_19, + NODE_20, + NODE_21, + NODE_22, + NODE_23, + NODE_24, + NODE_25, + NODE_26, + NODE_27, + NODE_28, + NODE_29, + NODE_30, + NODE_31, + NODE_32, + NODE_33, + NODE_34, + NODE_35, + NODE_36, + NODE_37, + NODE_38, + NODE_39, + NODE_40, + NODE_41, + NODE_42, + NODE_43, + NODE_44, + NODE_45, + NODE_46, + NODE_47, + NODE_48, + NODE_49, + NODE_50, + NODE_51, + NODE_52, + NODE_53, + NODE_54, + NODE_55, + NODE_56, + NODE_57, + NODE_58, + NODE_59, + NODE_60, + NODE_61, + NODE_62, + NODE_63, + NODE_64, + NODE_65, + NODE_66, + NODE_67, + NODE_68, + NODE_69, + NODE_70, + NODE_71, + cl_tree_0, + cl_tree_1, + cl_tree_2, + cl_tree_3, + cl_tree_4, + cl_tree_5, + cl_tree_6, + cl_tree_7, + cl_tree_8, + cl_tree_9, + cl_tree_10, + cl_tree_11, + cl_tree_12, + cl_break_0_0, + cl_break_0_1, + cl_break_0_2, + cl_break_0_3, + cl_break_0_4, + cl_break_0_5, + cl_break_0_6, + cl_break_0_7, + cl_break_0_8, + cl_break_1_0, + cl_break_1_1, + cl_break_1_2, + cl_break_1_3, + cl_break_1_4, + cl_break_1_5, + cl_break_1_6, + cl_break_1_7, + cl_break_1_8, + cl_break_1_9, + cl_break_1_10, + cl_break_1_11, + cl_break_1_12, +} + +#[allow(non_snake_case, non_camel_case_types, unused)] +pub enum firstLabels { + NODE_72, + NODE_73, + NODE_74, + NODE_75, + NODE_76, + NODE_77, + fl_tree_0, + fl_tree_1, + fl_tree_2, + fl_break_0_0, + fl_break_0_1, + fl_break_0_2, + fl_break_1_0, + fl_break_1_1, + fl_break_1_2, + fl_, +} + +#[allow(non_snake_case, non_camel_case_types, unused)] +pub enum lastLabels { + NODE_78, + NODE_79, + NODE_80, + NODE_81, + NODE_82, + NODE_83, + NODE_84, + NODE_85, + NODE_86, + NODE_87, + NODE_88, + NODE_89, + NODE_90, + NODE_91, + NODE_92, + ll_tree_0, + ll_tree_1, + ll_tree_2, + ll_tree_3, + ll_tree_4, + ll_tree_5, + ll_tree_6, + ll_tree_7, + ll_break_0_0, + ll_break_0_1, + ll_break_0_2, + ll_break_0_3, + ll_break_1_0, + ll_break_1_1, + ll_break_1_2, + ll_break_1_3, + ll_break_1_4, + ll_break_1_5, + ll_break_1_6, + ll_break_1_7, + ll_, +} + +#[allow(non_snake_case, non_camel_case_types, unused)] +pub enum singleLabels { + NODE_93, + NODE_94, + sl_tree_0, + sl_tree_1, + sl_break_0_0, + sl_break_0_1, + sl_break_1_0, + sl_break_1_1, + sl_, +} diff --git a/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/Spaghetti_last_line_forest_code.rs b/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/Spaghetti_last_line_forest_code.rs new file mode 100644 index 0000000..9187a4c --- /dev/null +++ b/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/Spaghetti_last_line_forest_code.rs @@ -0,0 +1,787 @@ +no_analyze!{{ +use lastLabels::*;let mut label = entry; +while let Some(next) = (|label| -> Option { match label { + NODE_78=> { + if (*img_row12.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(ll_tree_4); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row12.add((c + 2) as usize), *img_labels_row12.add((c - 2) as usize), solver); + return Some(ll_tree_4); + } + } + NODE_79=> { + if (*img_row12.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_6); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row12.add((c) as usize), *img_labels_row12.add((c - 2) as usize), solver); + return Some(ll_tree_6); + } + } + NODE_80=> { + if (*img_row12.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_6); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row12.add((c) as usize), *img_labels_row12.add((c - 2) as usize), solver); + return Some(ll_tree_6); + } + } + NODE_81=> { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + return Some(NODE_82); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(ll_tree_4); + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_3); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(ll_tree_2); + } + } + } + NODE_83=> { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_5); + } + else { + if (*img_row11.add((c + 2) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(ll_tree_4); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(ll_tree_2); + } + } + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + return Some(ll_tree_1); + } + } + NODE_84=> { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_5); + } + else { + return Some(NODE_81); + } + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + return Some(ll_tree_1); + } + } + NODE_82=> { + if (*img_row12.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(ll_tree_4); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row12.add((c + 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(ll_tree_4); + } + } + NODE_85=> { + if (*img_row12.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(ll_tree_4); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row12.add((c + 2) as usize), *img_labels_row12.add((c - 2) as usize), solver); + return Some(ll_tree_4); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row12.add((c + 2) as usize), *img_labels_row12.add((c - 2) as usize), solver); + return Some(ll_tree_4); + } + } + NODE_86=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_6); + } + else { + if (*img_row12.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_6); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(ll_tree_6); + } + } + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row12.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(ll_tree_4); + } + else { + if (*img_row12.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(ll_tree_4); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(ll_tree_4); + } + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(ll_tree_4); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(ll_tree_7); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(ll_tree_0); + } + } + } +ll_tree_0 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(ll_break_0_0); } else { return Some(ll_break_1_0); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_6); + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_81); + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_0); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(ll_tree_0); + } + } + } + } + else { + return Some(NODE_84); + } +} +ll_tree_1 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(ll_break_0_1); } else { return Some(ll_break_1_1); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_6); + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + return Some(NODE_80); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_6); + } + } + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + return Some(NODE_82); + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + return Some(NODE_85); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(ll_tree_4); + } + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_3); + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c - 2) as usize); + return Some(ll_tree_2); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(ll_tree_2); + } + } + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_0); + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c - 2) as usize); + return Some(ll_tree_0); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(ll_tree_0); + } + } + } + } + } + else { + return Some(NODE_84); + } +} +ll_tree_2 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(ll_break_0_2); } else { return Some(ll_break_1_2); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(ll_tree_6); + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 2) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(ll_tree_4); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(ll_tree_7); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(ll_tree_0); + } + } + } + else { + return Some(NODE_83); + } +} +ll_tree_3 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(ll_break_0_2); } else { return Some(ll_break_1_3); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + return Some(NODE_79); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(ll_tree_6); + } + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row12.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + return Some(NODE_78); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(ll_tree_4); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(ll_tree_4); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(ll_tree_7); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(ll_tree_0); + } + } + } + else { + return Some(NODE_83); + } +} +ll_tree_4 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(ll_break_0_2); } else { return Some(ll_break_1_4); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_6); + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row12.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c + 2) as usize); + return Some(ll_tree_4); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(ll_tree_4); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(ll_tree_7); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(ll_tree_0); + } + } + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_5); + } + else { + if (*img_row11.add((c + 2) as usize)).to_bool() { + return Some(NODE_82); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_3); + } + } + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + return Some(ll_tree_1); + } + } +} +ll_tree_5 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(ll_break_0_2); } else { return Some(ll_break_1_5); } } + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_86); + } + else { + return Some(NODE_84); + } +} +ll_tree_6 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(ll_break_0_3); } else { return Some(ll_break_1_6); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row00.add((c - 1) as usize)).to_bool() { + return Some(NODE_86); + } + else { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_6); + } + else { + return Some(NODE_80); + } + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + return Some(NODE_82); + } + else { + return Some(NODE_85); + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_3); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c - 2) as usize); + return Some(ll_tree_2); + } + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + return Some(ll_tree_0); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c - 2) as usize); + return Some(ll_tree_0); + } + } + } + } + } + else { + return Some(NODE_84); + } +} +ll_tree_7 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(ll_break_0_2); } else { return Some(ll_break_1_7); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + if (*img_row11.add((c - 2) as usize)).to_bool() { + return Some(NODE_79); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(ll_tree_6); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + return Some(ll_tree_6); + } + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 2) as usize)).to_bool() { + if (*img_row12.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + if (*img_row11.add((c - 2) as usize)).to_bool() { + return Some(NODE_78); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(ll_tree_4); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(ll_tree_4); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c + 2) as usize), solver); + return Some(ll_tree_4); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(ll_tree_7); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(ll_tree_0); + } + } + } + else { + return Some(NODE_83); + } +} +ll_break_0_0 => { + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + return None;} +ll_break_0_1 => { + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c - 2) as usize); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + return None;} +ll_break_0_2 => { + if (*img_row00.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + return None;} +ll_break_0_3 => { + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row00.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c - 2) as usize); + } + } + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + return None;} + NODE_87=> { + if (*img_row00.add((c + 1) as usize)).to_bool() { + return Some(NODE_88); + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + } + NODE_88=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } + } + NODE_89=> { + if (*img_row12.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row12.add((c) as usize), *img_labels_row12.add((c - 2) as usize), solver); + } + } + NODE_90=> { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + if (*img_row12.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + NODE_91=> { + if (*img_row00.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + } + NODE_92=> { + if (*img_row12.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row12.add((c) as usize), *img_labels_row12.add((c - 2) as usize), solver); + } + } +ll_break_1_0 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_88); + } + else { + return Some(NODE_87); + } + return None;} +ll_break_1_1 => { + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + return Some(NODE_92); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + if (*img_row11.add((c - 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c - 2) as usize); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + } + } + } + else { + return Some(NODE_87); + } + return None;} +ll_break_1_2 => { + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + else { + return Some(NODE_91); + } + return None;} +ll_break_1_3 => { + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + return Some(NODE_89); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + else { + return Some(NODE_91); + } + return None;} +ll_break_1_4 => { + if (*img_row00.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + if (*img_row00.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + } + return None;} +ll_break_1_5 => { + if (*img_row00.add((c) as usize)).to_bool() { + return Some(NODE_90); + } + else { + return Some(NODE_87); + } + return None;} +ll_break_1_6 => { + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row00.add((c - 1) as usize)).to_bool() { + return Some(NODE_90); + } + else { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + return Some(NODE_92); + } + } + else { + if (*img_row11.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c) as usize); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row12.add((c - 2) as usize); + } + } + } + } + else { + return Some(NODE_87); + } + return None;} +ll_break_1_7 => { + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row11.add((c + 1) as usize)).to_bool() { + if (*img_row12.add((c) as usize)).to_bool() { + if (*img_row11.add((c - 2) as usize)).to_bool() { + return Some(NODE_89); + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + } + else { + *img_labels_row00.add(c as usize) = LabelsSolver::merge(*img_labels_row00.add((c - 2) as usize), *img_labels_row12.add((c) as usize), solver); + } + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + } + else { + return Some(NODE_91); + } + return None;} +ll_ => {}, + }; None})(label) +{ +label = next; +} +}} diff --git a/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/Spaghetti_single_line_forest_code.rs b/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/Spaghetti_single_line_forest_code.rs new file mode 100644 index 0000000..bc36ac7 --- /dev/null +++ b/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/Spaghetti_single_line_forest_code.rs @@ -0,0 +1,91 @@ +no_analyze!{{ +use singleLabels::*;let mut label = entry; +while let Some(next) = (|label| -> Option { match label { + NODE_93=> { + if (*img_row00.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(sl_tree_1); + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + return Some(sl_tree_0); + } + } +sl_tree_0 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(sl_break_0_0); } else { return Some(sl_break_1_0); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(sl_tree_1); + } + else { + *img_labels_row00.add(c as usize) = solver.new_label(); + return Some(sl_tree_0); + } + } + else { + return Some(NODE_93); + } +} +sl_tree_1 => { +if ({c+=2; c}) >= w - 2 { if c > w - 2 { return Some(sl_break_0_1); } else { return Some(sl_break_1_1); } } + if (*img_row00.add((c) as usize)).to_bool() { + if (*img_row00.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(sl_tree_1); + } + else { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + return Some(sl_tree_0); + } + } + else { + return Some(NODE_93); + } +} +sl_break_0_0 => { + if (*img_row00.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + return None;} +sl_break_0_1 => { + if (*img_row00.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + return None;} + NODE_94=> { + if (*img_row00.add((c + 1) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + else { + *img_labels_row00.add(c as usize) = 0.elem(); + } + } +sl_break_1_0 => { + if (*img_row00.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = solver.new_label(); + } + else { + return Some(NODE_94); + } + return None;} +sl_break_1_1 => { + if (*img_row00.add((c) as usize)).to_bool() { + *img_labels_row00.add(c as usize) = *img_labels_row00.add((c - 2) as usize); + } + else { + return Some(NODE_94); + } + return None;} +sl_ => {}, + }; None})(label) +{ +label = next; +} +}} diff --git a/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/mod.rs b/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/mod.rs new file mode 100644 index 0000000..f8db796 --- /dev/null +++ b/crates/burn-vision/src/backends/cpu/connected_components/spaghetti/mod.rs @@ -0,0 +1,273 @@ +//! Spaghetti algorithm for connected component labeling +//! F. Bolelli, S. Allegretti, L. Baraldi, and C. Grana, +//! "Spaghetti Labeling: Directed Acyclic Graphs for Block-Based Bonnected Components Labeling," +//! IEEE Transactions on Image Processing, vol. 29, no. 1, pp. 1999-2012, 2019. +//! +//! Decision forests are generated using a modified [GRAPHGEN](https://github.com/wingertge/GRAPHGEN) +//! as described in +//! +//! F. Bolelli, S. Allegretti, C. Grana. +//! "One DAG to Rule Them All." +//! IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021 + +#![allow( + unreachable_code, + clippy::collapsible_else_if, + clippy::if_same_then_else +)] + +use std::cmp::Ordering; + +use burn_core::tensor::{Element, ElementComparison, ElementConversion, cast::ToElement}; +use ndarray::{Array2, Axis, s}; + +#[allow(non_snake_case)] +mod Spaghetti_forest_labels; +pub(crate) use Spaghetti_forest_labels::*; + +use crate::Connectivity; + +use super::{Solver, StatsOp, max_labels}; + +pub fn process( + img_arr: Array2, + stats: &mut impl StatsOp

{ + // Only the inner policy's parameters are persisted; the exploration `step` counter resets on + // load (it is part of the exploration schedule, not the learned parameters). + type Record = ::Record; + + fn into_record(self) -> Self::Record { + self.inner_state.into_record() + } + + fn load_record(&self, record: Self::Record) -> Self { + Self { + inner_state: self.inner_state.load_record(record), + step: self.step, + } + } +} + +#[derive(Clone, Debug)] +pub struct EpsilonGreedyPolicy { + inner_policy: P, + eps_start: f64, + eps_end: f64, + eps_decay: f64, + step: usize, +} + +impl EpsilonGreedyPolicy

{ + pub fn new(inner_policy: P, eps_start: f64, eps_end: f64, eps_decay: f64) -> Self { + Self { + inner_policy, + eps_start, + eps_end, + eps_decay, + step: 0, + } + } + + pub fn inner_policy(&self) -> P { + self.inner_policy.clone() + } + + pub fn set_inner_policy(&mut self, policy: P) { + self.inner_policy = policy; + } + + fn get_threshold(&self) -> f64 { + self.eps_end + + (self.eps_start - self.eps_end) * f64::exp(-(self.step as f64) / self.eps_decay) + } + + fn step(&mut self) -> f64 { + let thresh = self.get_threshold(); + self.step += 1; + thresh + } +} + +impl

Policy for EpsilonGreedyPolicy

+where + P: Policy, Action = DiscreteActionTensor<2>>, +{ + type ActionContext = EpsilonGreedyPolicyOutput; + type PolicyState = EpsilonGreedyPolicyState

; + + type Observation = P::Observation; + type ActionDistribution = DiscreteLogitsTensor<2>; + type Action = DiscreteActionTensor<2>; + + fn forward(&mut self, states: Self::Observation) -> Self::ActionDistribution { + self.inner_policy.forward(states) + } + + fn action( + &mut self, + states: Self::Observation, + deterministic: bool, + ) -> (Self::Action, Vec) { + let logits = self.inner_policy.forward(states).logits; + let greedy_actions = logits.argmax(1); + let greedy_actions = greedy_actions.split(1, 0); + + let mut contexts = vec![]; + let mut actions = vec![]; + for a in greedy_actions { + let threshold = self.step(); + let threshold = if deterministic { 0.0 } else { threshold }; + contexts.push(EpsilonGreedyPolicyOutput { epsilon: threshold }); + if random::() > threshold { + actions.push(a.clone().float().inner()); + } else { + actions + .push(Tensor::<1>::from_floats([random_range(0..2)], &a.device()).unsqueeze()); + } + } + + let output = Tensor::cat(actions, 0); + (DiscreteActionTensor { actions: output }, contexts) + } + + fn update(&mut self, update: Self::PolicyState) { + // Note : updating an epsilon greedy policy doesn't change the step. + self.inner_policy.update(update.inner_state); + } + + fn state(&self) -> Self::PolicyState { + EpsilonGreedyPolicyState { + inner_state: self.inner_policy.state(), + step: self.step, + } + } + + fn to_device(self, device: &Device) -> Self { + let mut policy = self.clone(); + let inner = policy.inner_policy().to_device(device); + policy.set_inner_policy(inner); + policy + } + + fn load_record(self, record: ::Record) -> Self { + let state = self.state().load_record(record); + let inner_policy = self + .inner_policy + .load_record(state.inner_state.into_record()); + EpsilonGreedyPolicy { + inner_policy, + eps_start: self.eps_start, + eps_end: self.eps_end, + eps_decay: self.eps_decay, + step: state.step, + } + } +} diff --git a/examples/guide/Cargo.toml b/examples/guide/Cargo.toml new file mode 100644 index 0000000..6797d66 --- /dev/null +++ b/examples/guide/Cargo.toml @@ -0,0 +1,27 @@ +[package] +authors = ["nathanielsimard "] +edition.workspace = true +license.workspace = true +name = "guide" +publish = false +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["burn/default", "burn/tui"] +# Opt-in for macOS users with the Vulkan SDK installed (enabled by default on other platforms) +vulkan = ["burn/vulkan"] + +[dependencies] +burn = { workspace = true, features = [ + "wgpu", + "train", + "vision", +] } + +[target.'cfg(not(target_os = "macos"))'.dependencies] +burn = { workspace = true, features = ["default", "vulkan"] } + +log = { workspace = true } diff --git a/examples/guide/README.md b/examples/guide/README.md new file mode 100644 index 0000000..6e107e6 --- /dev/null +++ b/examples/guide/README.md @@ -0,0 +1,24 @@ +# Basic Workflow: From Training to Inference + +This example corresponds to the [book's guide](https://burn.dev/books/burn/basic-workflow/). + +## Example Usage + + +### Training + +```sh +cargo run --bin train --release +``` + +### Inference + +```sh +cargo run --bin infer --release +``` + +### Print the model + +```sh +cargo run --bin print --release +``` diff --git a/examples/guide/examples/guide.rs b/examples/guide/examples/guide.rs new file mode 100644 index 0000000..bbac496 --- /dev/null +++ b/examples/guide/examples/guide.rs @@ -0,0 +1,16 @@ +// +// Note: If you are following the Burn Book guide this file can be ignored. +// +// This example file is added only for convenience and consistency so that +// the guide example can be executed like any other examples using: +// +// cargo run --release --example guide +// +use std::process::Command; + +fn main() { + Command::new("cargo") + .args(["run", "--bin", "train", "--release"]) + .status() + .expect("guide example should run"); +} diff --git a/examples/guide/src/bin/infer.rs b/examples/guide/src/bin/infer.rs new file mode 100644 index 0000000..95f5da6 --- /dev/null +++ b/examples/guide/src/bin/infer.rs @@ -0,0 +1,19 @@ +#![recursion_limit = "131"] +use burn::{data::dataset::Dataset, prelude::*}; +use guide::inference; + +fn main() { + let device = Device::wgpu(DeviceKind::DefaultDevice); + + // All the training artifacts are saved in this directory + let artifact_dir = "/tmp/guide"; + + // Infer the model + inference::infer( + artifact_dir, + device, + burn::data::dataset::vision::MnistDataset::test() + .get(42) + .unwrap(), + ); +} diff --git a/examples/guide/src/bin/print.rs b/examples/guide/src/bin/print.rs new file mode 100644 index 0000000..4fac526 --- /dev/null +++ b/examples/guide/src/bin/print.rs @@ -0,0 +1,9 @@ +use burn::prelude::*; +use guide::model::ModelConfig; + +fn main() { + let device = Device::wgpu(DeviceKind::DefaultDevice); + let model = ModelConfig::new(10, 512).init(&device); + + println!("{model}"); +} diff --git a/examples/guide/src/bin/train.rs b/examples/guide/src/bin/train.rs new file mode 100644 index 0000000..0688238 --- /dev/null +++ b/examples/guide/src/bin/train.rs @@ -0,0 +1,31 @@ +#![recursion_limit = "131"] +use burn::{data::dataset::Dataset, optim::AdamConfig, prelude::*}; +use guide::{ + inference, + model::ModelConfig, + training::{self, TrainingConfig}, +}; + +fn main() { + // Create a default Wgpu-backed device. + let device = Device::wgpu(DeviceKind::DefaultDevice); + + // All the training artifacts will be saved in this directory + let artifact_dir = "target/guide"; + + // Train the model + training::train( + artifact_dir, + TrainingConfig::new(ModelConfig::new(10, 512), AdamConfig::new()), + device.clone(), + ); + + // Infer the model + inference::infer( + artifact_dir, + device, + burn::data::dataset::vision::MnistDataset::test() + .get(42) + .unwrap(), + ); +} diff --git a/examples/guide/src/data.rs b/examples/guide/src/data.rs new file mode 100644 index 0000000..a162c26 --- /dev/null +++ b/examples/guide/src/data.rs @@ -0,0 +1,38 @@ +use burn::{ + data::{dataloader::batcher::Batcher, dataset::vision::MnistItem}, + prelude::*, +}; + +#[derive(Clone, Default)] +pub struct MnistBatcher {} + +#[derive(Clone, Debug)] +pub struct MnistBatch { + pub images: Tensor<3>, + pub targets: Tensor<1, Int>, +} + +impl Batcher for MnistBatcher { + fn batch(&self, items: Vec, device: &Device) -> MnistBatch { + let images = items + .iter() + .map(|item| TensorData::from(item.image)) + .map(|data| Tensor::<2>::from_data(data, device)) + .map(|tensor| tensor.reshape([1, 28, 28])) + // Normalize: scale between [0,1] and make the mean=0 and std=1 + // values mean=0.1307,std=0.3081 are from the PyTorch MNIST example + // https://github.com/pytorch/examples/blob/54f4572509891883a947411fd7239237dd2a39c3/mnist/main.py#L122 + .map(|tensor| ((tensor / 255) - 0.1307) / 0.3081) + .collect(); + + let targets = items + .iter() + .map(|item| Tensor::<1, Int>::from_data([item.label as i64], device)) + .collect(); + + let images = Tensor::cat(images, 0); + let targets = Tensor::cat(targets, 0); + + MnistBatch { images, targets } + } +} diff --git a/examples/guide/src/inference.rs b/examples/guide/src/inference.rs new file mode 100644 index 0000000..d63aa46 --- /dev/null +++ b/examples/guide/src/inference.rs @@ -0,0 +1,24 @@ +use crate::{data::MnistBatcher, training::TrainingConfig}; +use burn::{ + data::{dataloader::batcher::Batcher, dataset::vision::MnistItem}, + prelude::*, + store::ModuleRecord, +}; + +pub fn infer(artifact_dir: &str, device: impl Into, item: MnistItem) { + let device = device.into(); + let config = TrainingConfig::load(format!("{artifact_dir}/config.json")) + .expect("Config should exist for the model; run train first"); + let record = ModuleRecord::load(format!("{artifact_dir}/model")) + .expect("Trained model should exist; run train first"); + + let model = config.model.init(&device).load_record(record); + + let label = item.label; + let batcher = MnistBatcher::default(); + let batch = batcher.batch(vec![item], &device); + let output = model.forward(batch.images); + let predicted: u8 = output.argmax(1).flatten::<1>(0, 1).into_scalar(); + + println!("Predicted {predicted} Expected {label}"); +} diff --git a/examples/guide/src/lib.rs b/examples/guide/src/lib.rs new file mode 100644 index 0000000..a56bdbe --- /dev/null +++ b/examples/guide/src/lib.rs @@ -0,0 +1,10 @@ +// +// Note: If you are following the Burn Book guide this file can be ignored. +// +// This lib.rs file is added only for convenience so that the code in this +// guide can be reused. +// +pub mod data; +pub mod inference; +pub mod model; +pub mod training; diff --git a/examples/guide/src/model.rs b/examples/guide/src/model.rs new file mode 100644 index 0000000..adaba2a --- /dev/null +++ b/examples/guide/src/model.rs @@ -0,0 +1,68 @@ +use burn::{ + nn::{ + Dropout, DropoutConfig, Linear, LinearConfig, Relu, + conv::{Conv2d, Conv2dConfig}, + pool::{AdaptiveAvgPool2d, AdaptiveAvgPool2dConfig}, + }, + prelude::*, +}; + +#[derive(Module, Debug)] +pub struct Model { + conv1: Conv2d, + conv2: Conv2d, + pool: AdaptiveAvgPool2d, + dropout: Dropout, + linear1: Linear, + linear2: Linear, + activation: Relu, +} + +#[derive(Config, Debug)] +pub struct ModelConfig { + num_classes: usize, + hidden_size: usize, + #[config(default = "0.5")] + dropout: f64, +} + +impl ModelConfig { + /// Returns the initialized model. + pub fn init(&self, device: &Device) -> Model { + Model { + conv1: Conv2dConfig::new([1, 8], [3, 3]).init(device), + conv2: Conv2dConfig::new([8, 16], [3, 3]).init(device), + pool: AdaptiveAvgPool2dConfig::new([8, 8]).init(), + activation: Relu::new(), + linear1: LinearConfig::new(16 * 8 * 8, self.hidden_size).init(device), + linear2: LinearConfig::new(self.hidden_size, self.num_classes).init(device), + dropout: DropoutConfig::new(self.dropout).init(), + } + } +} + +impl Model { + /// # Shapes + /// - Images [batch_size, height, width] + /// - Output [batch_size, class_prob] + pub fn forward(&self, images: Tensor<3>) -> Tensor<2> { + let [batch_size, height, width] = images.dims(); + + // Create a channel. + let x = images.reshape([batch_size, 1, height, width]); + + let x = self.conv1.forward(x); // [batch_size, 8, _, _] + let x = self.dropout.forward(x); + let x = self.conv2.forward(x); // [batch_size, 16, _, _] + let x = self.dropout.forward(x); + let x = self.activation.forward(x); + + let x = self.pool.forward(x); // [batch_size, 16, 8, 8] + let x = x.reshape([batch_size, 16 * 8 * 8]); + let x = self.linear1.forward(x); + let x = self.dropout.forward(x); + let x = self.activation.forward(x); + + self.linear2.forward(x) // [batch_size, num_classes] + } +} diff --git a/examples/guide/src/training.rs b/examples/guide/src/training.rs new file mode 100644 index 0000000..37c8599 --- /dev/null +++ b/examples/guide/src/training.rs @@ -0,0 +1,116 @@ +use std::path::PathBuf; + +use crate::{ + data::{MnistBatch, MnistBatcher}, + model::{Model, ModelConfig}, +}; +use burn::{ + data::{dataloader::DataLoaderBuilder, dataset::vision::MnistDataset}, + nn::loss::CrossEntropyLossConfig, + optim::AdamConfig, + prelude::*, + train::{ + ClassificationOutput, InferenceStep, Learner, SupervisedTraining, TrainOutput, TrainStep, + metric::{AccuracyMetric, LossMetric}, + }, +}; + +impl Model { + pub fn forward_classification( + &self, + images: Tensor<3>, + targets: Tensor<1, Int>, + ) -> ClassificationOutput { + let output = self.forward(images); + let loss = CrossEntropyLossConfig::new() + .init(&output.device()) + .forward(output.clone(), targets.clone()); + + ClassificationOutput::new(loss, output, targets) + } +} + +impl TrainStep for Model { + type Input = MnistBatch; + type Output = ClassificationOutput; + + fn step(&self, batch: MnistBatch) -> TrainOutput { + let item = self.forward_classification(batch.images, batch.targets); + + TrainOutput::new(self, item.loss.backward(), item) + } +} + +impl InferenceStep for Model { + type Input = MnistBatch; + type Output = ClassificationOutput; + + fn step(&self, batch: MnistBatch) -> ClassificationOutput { + self.forward_classification(batch.images, batch.targets) + } +} + +#[derive(Config, Debug)] +pub struct TrainingConfig { + pub model: ModelConfig, + pub optimizer: AdamConfig, + #[config(default = 10)] + pub num_epochs: usize, + #[config(default = 64)] + pub batch_size: usize, + #[config(default = 4)] + pub num_workers: usize, + #[config(default = 42)] + pub seed: u64, + #[config(default = 1.0e-4)] + pub learning_rate: f64, +} + +fn create_artifact_dir(artifact_dir: &str) { + std::fs::remove_file(PathBuf::from(artifact_dir).join("experiment.log")).ok(); + std::fs::create_dir_all(artifact_dir).ok(); +} + +pub fn train(artifact_dir: &str, config: TrainingConfig, device: impl Into) { + create_artifact_dir(artifact_dir); + config + .save(format!("{artifact_dir}/config.json")) + .expect("Config should be saved successfully"); + + let device = device.into(); + device.seed(config.seed); + let autodiff_device = device.clone().autodiff(); + + let batcher = MnistBatcher::default(); + + let dataloader_train = DataLoaderBuilder::new(batcher.clone()) + .batch_size(config.batch_size) + .shuffle(config.seed) + .num_workers(config.num_workers) + .build(MnistDataset::train()); + + let dataloader_test = DataLoaderBuilder::new(batcher) + .batch_size(config.batch_size) + .shuffle(config.seed) + .num_workers(config.num_workers) + .build(MnistDataset::test()); + + let training = SupervisedTraining::new(artifact_dir, dataloader_train, dataloader_test) + .metrics((AccuracyMetric::new(), LossMetric::new())) + .with_default_checkpointers() + .num_epochs(config.num_epochs) + .summary(); + + let model = config.model.init(&autodiff_device); + let result = training.launch(Learner::new( + model, + config.optimizer.init(), + config.learning_rate, + )); + + result + .model + .into_record() + .save(format!("{artifact_dir}/model")) + .expect("Trained model should be saved successfully"); +} diff --git a/examples/import-model-weights/Cargo.toml b/examples/import-model-weights/Cargo.toml new file mode 100644 index 0000000..f9e388f --- /dev/null +++ b/examples/import-model-weights/Cargo.toml @@ -0,0 +1,26 @@ +[package] +authors = ["Dilshod Tadjibaev (@antimora)"] +edition.workspace = true +license = "MIT OR Apache-2.0" +name = "import-model-weights" +publish = false +version.workspace = true + +[lints] +workspace = true + +[dependencies] + +burn = { workspace = true, features = [ + "default", + "flex", + "dataset", + "vision", +] } + +burn-store = { workspace = true, features = [ + "std", + "pytorch", + "safetensors", + "burnpack", +] } diff --git a/examples/import-model-weights/README.md b/examples/import-model-weights/README.md new file mode 100644 index 0000000..710c8bd --- /dev/null +++ b/examples/import-model-weights/README.md @@ -0,0 +1,110 @@ +# Import Model Weights + +This crate provides examples for importing model weights from different formats to Burn. + +## Examples + +### PyTorch + +Imports weights from a PyTorch `.pt` file using `burn-store`. + +```bash +cargo run --bin pytorch -- +``` + +Example: +```bash +cargo run --bin pytorch -- 15 + +Loading PyTorch model weights from file: weights/mnist.pt +Image index: 15 +Success! +Predicted: 5 +Actual: 5 +See the image online, click the link below: +https://huggingface.co/datasets/ylecun/mnist/viewer/mnist/test?row=15 +``` + +### Safetensors + +Imports weights from a Safetensors file using `burn-store`. + +```bash +cargo run --bin safetensors -- +``` + +Example: +```bash +cargo run --bin safetensors -- 42 + +Loading Safetensors model weights from file: weights/mnist.safetensors +Image index: 42 +Success! +Predicted: 4 +Actual: 4 +See the image online, click the link below: +https://huggingface.co/datasets/ylecun/mnist/viewer/mnist/test?row=42 +``` + +### Convert + +Converts between different weight formats (PyTorch or Safetensors) to Burn's native Burnpack format. + +```bash +cargo run --bin convert -- +``` + +Where: +- ``: Either `pytorch` or `safetensors` +- ``: Path to save the converted model file + +Example with PyTorch: +```bash +cargo run --bin convert -- pytorch /tmp/burn-convert + +Loading PyTorch weights from 'weights/mnist.pt'... +Saving model to '/tmp/burn-convert/mnist.bpk'... +Model successfully saved to '/tmp/burn-convert/mnist.bpk'. +``` + +Example with Safetensors: +```bash +cargo run --bin convert -- safetensors /tmp/burn-convert + +Loading Safetensors weights from 'weights/mnist.safetensors'... +Saving model to '/tmp/burn-convert/mnist.bpk'... +Model successfully saved to '/tmp/burn-convert/mnist.bpk'. +``` + +### Burnpack + +Demonstrates loading and using a model from Burn's native Burnpack format. + +```bash +cargo run --bin burnpack -- +``` + +Where: +- ``: Index of the MNIST test image to classify +- ``: Path to the model file (without extension) + +Example: +```bash +cargo run --bin burnpack -- 35 /tmp/burn-convert/mnist + +Loading model weights from file: /tmp/burn-convert/mnist.bpk +Image index: 35 +Success! +Predicted: 2 +Actual: 2 +See the image online, click the link below: +https://huggingface.co/datasets/ylecun/mnist/viewer/mnist/test?row=35 +``` + +## Workflow + +A typical workflow using these examples: + +1. Start with pre-trained weights in either PyTorch or Safetensors format +2. Use the `convert` example to convert to Burn's native Burnpack format +3. Load and use the converted model with the `burnpack` example diff --git a/examples/import-model-weights/src/bin/burnpack.rs b/examples/import-model-weights/src/bin/burnpack.rs new file mode 100644 index 0000000..d653b4c --- /dev/null +++ b/examples/import-model-weights/src/bin/burnpack.rs @@ -0,0 +1,34 @@ +use std::env; +use std::path::Path; + +use burn_store::{BurnpackStore, ModuleSnapshot}; +use import_model_weights::{Model, infer}; + +pub fn main() { + let args: Vec = env::args().collect(); + + if args.len() < 3 { + eprintln!("Usage: {} ", args[0]); + std::process::exit(1); + } + + let model_path_str = &args[2]; + let model_path = Path::new(model_path_str); + println!( + "Loading model weights from file: {}.bpk", + model_path.display() + ); + + // Initialize a model with default weights + let device = Default::default(); + let mut model = Model::init(&device); + + // Load the model from the Burnpack file + let mut store = BurnpackStore::from_file(model_path); + model + .load_from(&mut store) + .expect("Failed to load model from Burnpack file"); + + // Infer using the loaded model + infer(model); +} diff --git a/examples/import-model-weights/src/bin/convert.rs b/examples/import-model-weights/src/bin/convert.rs new file mode 100644 index 0000000..eeb5a19 --- /dev/null +++ b/examples/import-model-weights/src/bin/convert.rs @@ -0,0 +1,80 @@ +use std::{env, path::Path, process}; + +use burn_store::{ + BurnpackStore, ModuleSnapshot, PyTorchToBurnAdapter, PytorchStore, SafetensorsStore, +}; +use import_model_weights::Model; + +// Path constants +const PYTORCH_WEIGHTS_PATH: &str = "weights/mnist.pt"; +const SAFETENSORS_WEIGHTS_PATH: &str = "weights/mnist.safetensors"; +const MODEL_OUTPUT_NAME: &str = "mnist"; + +pub fn main() { + let args: Vec = env::args().collect(); + + // Check argument count + if args.len() < 3 { + eprintln!( + "Usage: {} ", + args[0] + ); + process::exit(1); + } + + // Get weight format and output directory from arguments + let weight_format = args[1].as_str(); + let output_directory = Path::new(&args[2]); + + // Use the default device (CPU) + let device = Default::default(); + + // Initialize a model with default weights + let mut model = Model::init(&device); + + // Load the model weights based on the specified format + match weight_format { + "pytorch" => { + println!("Loading PyTorch weights from '{PYTORCH_WEIGHTS_PATH}'..."); + let mut store = PytorchStore::from_file(PYTORCH_WEIGHTS_PATH); + model.load_from(&mut store).unwrap_or_else(|e| { + panic!("Failed to load PyTorch model weights from '{PYTORCH_WEIGHTS_PATH}': {e}") + }); + } + "safetensors" => { + println!("Loading Safetensors weights from '{SAFETENSORS_WEIGHTS_PATH}'..."); + let mut store = SafetensorsStore::from_file(SAFETENSORS_WEIGHTS_PATH) + .with_from_adapter(PyTorchToBurnAdapter); + model.load_from(&mut store).unwrap_or_else(|e| { + panic!( + "Failed to load Safetensors model weights from '{SAFETENSORS_WEIGHTS_PATH}': {e}" + ) + }); + } + _ => { + eprintln!( + "Error: Unsupported weight format '{weight_format}'. Please use 'pytorch' or 'safetensors'." + ); + process::exit(1); + } + }; + + // Define the output path for the Burn model file + let output_file_path = output_directory.join(MODEL_OUTPUT_NAME); + + println!("Saving model to '{}.bpk'...", output_file_path.display()); + + // Save the model using BurnpackStore + let mut store = BurnpackStore::from_file(&output_file_path).overwrite(true); + model.save_into(&mut store).unwrap_or_else(|e| { + panic!( + "Failed to save model to '{}.bpk': {e}", + output_file_path.display() + ) + }); + + println!( + "Model successfully saved to '{}.bpk'.", + output_file_path.display() + ); +} diff --git a/examples/import-model-weights/src/bin/pytorch.rs b/examples/import-model-weights/src/bin/pytorch.rs new file mode 100644 index 0000000..251a79a --- /dev/null +++ b/examples/import-model-weights/src/bin/pytorch.rs @@ -0,0 +1,22 @@ +use burn_store::{ModuleSnapshot, PytorchStore}; + +use import_model_weights::{Model, infer}; + +const WEIGHTS_FILE: &str = "weights/mnist.pt"; + +pub fn main() { + println!("Loading PyTorch model weights from file: {WEIGHTS_FILE}"); + + // Initialize a model with default weights + let device = Default::default(); + let mut model = Model::init(&device); + + // Load PyTorch weights into the model + let mut store = PytorchStore::from_file(WEIGHTS_FILE); + model + .load_from(&mut store) + .expect("Failed to load PyTorch model weights"); + + // Infer using the loaded model + infer(model); +} diff --git a/examples/import-model-weights/src/bin/safetensors.rs b/examples/import-model-weights/src/bin/safetensors.rs new file mode 100644 index 0000000..175355a --- /dev/null +++ b/examples/import-model-weights/src/bin/safetensors.rs @@ -0,0 +1,23 @@ +use burn_store::{ModuleSnapshot, PyTorchToBurnAdapter, SafetensorsStore}; + +use import_model_weights::{Model, infer}; + +const WEIGHTS_FILE: &str = "weights/mnist.safetensors"; + +pub fn main() { + println!("Loading Safetensors model weights from file: {WEIGHTS_FILE}"); + + // Initialize a model with default weights + let device = Default::default(); + let mut model = Model::init(&device); + + // Load Safetensors weights into the model (using PyTorch adapter since weights were exported from PyTorch) + let mut store = + SafetensorsStore::from_file(WEIGHTS_FILE).with_from_adapter(PyTorchToBurnAdapter); + model + .load_from(&mut store) + .expect("Failed to load Safetensors model weights"); + + // Infer using the loaded model + infer(model); +} diff --git a/examples/import-model-weights/src/inference.rs b/examples/import-model-weights/src/inference.rs new file mode 100644 index 0000000..75589ec --- /dev/null +++ b/examples/import-model-weights/src/inference.rs @@ -0,0 +1,56 @@ +use burn::prelude::*; + +use std::env::args; + +use burn::data::dataloader::Dataset; +use burn::data::dataset::vision::MnistDataset; + +use crate::model::Model; + +const IMAGE_INX: usize = 42; // <- Change this to test a different image + +pub fn infer(model: Model) { + // Get image index argument (first) from command line + + let image_index = if let Some(image_index) = args().nth(1) { + println!("Image index: {image_index}"); + image_index + .parse::() + .expect("Failed to parse image index") + } else { + println!("No image index provided; Using default image index: {IMAGE_INX}"); + IMAGE_INX + }; + + assert!(image_index < 10000, "Image index must be less than 10000"); + + // Get device from the model + let device = model.devices().into_iter().next().unwrap_or_default(); + + // Load the MNIST dataset and get an item + let dataset = MnistDataset::test(); + let item = dataset.get(image_index).unwrap(); + + // Create a tensor from the image data + let image_data = item.image.iter().copied().flatten().collect::>(); + let mut input = + Tensor::<1>::from_floats(image_data.as_slice(), &device).reshape([1, 1, 28, 28]); + + // Normalize the input + input = ((input / 255) - 0.1307) / 0.3081; + + // Run the model on the input + let output = model.forward(input); + + // Get the index of the maximum value + let arg_max: u8 = output.argmax(1).into_scalar(); + + // Check if the index matches the label + assert!(arg_max == item.label); + + println!("Success!"); + println!("Predicted: {arg_max}"); + println!("Actual: {}", item.label); + println!("See the image online, click the link below:"); + println!("https://huggingface.co/datasets/ylecun/mnist/viewer/mnist/test?row={image_index}"); +} diff --git a/examples/import-model-weights/src/lib.rs b/examples/import-model-weights/src/lib.rs new file mode 100644 index 0000000..f6b8dda --- /dev/null +++ b/examples/import-model-weights/src/lib.rs @@ -0,0 +1,5 @@ +pub mod inference; +pub mod model; + +pub use inference::infer; +pub use model::Model; diff --git a/examples/import-model-weights/src/model.rs b/examples/import-model-weights/src/model.rs new file mode 100644 index 0000000..fec8356 --- /dev/null +++ b/examples/import-model-weights/src/model.rs @@ -0,0 +1,57 @@ +use burn::{ + nn::{ + BatchNorm, BatchNormConfig, Linear, LinearConfig, + conv::{Conv2d, Conv2dConfig}, + }, + prelude::*, + tensor::activation::{log_softmax, relu}, +}; + +#[derive(Module, Debug)] +pub struct Model { + conv1: Conv2d, + conv2: Conv2d, + conv3: Conv2d, + norm1: BatchNorm, + fc1: Linear, + fc2: Linear, + norm2: BatchNorm, +} + +impl Model { + pub fn init(device: &Device) -> Self { + let conv1 = Conv2dConfig::new([1, 8], [3, 3]).init(device); + let conv2 = Conv2dConfig::new([8, 16], [3, 3]).init(device); + let conv3 = Conv2dConfig::new([16, 24], [3, 3]).init(device); + let norm1 = BatchNormConfig::new(24).init(device); + let fc1 = LinearConfig::new(11616, 32).init(device); + let fc2 = LinearConfig::new(32, 10).init(device); + let norm2 = BatchNormConfig::new(10).init(device); + + Self { + conv1, + conv2, + conv3, + norm1, + fc1, + fc2, + norm2, + } + } + + pub fn forward(&self, input1: Tensor<4>) -> Tensor<2> { + let conv1_out1 = self.conv1.forward(input1); + let relu1_out1 = relu(conv1_out1); + let conv2_out1 = self.conv2.forward(relu1_out1); + let relu2_out1 = relu(conv2_out1); + let conv3_out1 = self.conv3.forward(relu2_out1); + let relu3_out1 = relu(conv3_out1); + let norm1_out1 = self.norm1.forward(relu3_out1); + let flatten1_out1 = norm1_out1.flatten(1, 3); + let fc1_out1 = self.fc1.forward(flatten1_out1); + let relu4_out1 = relu(fc1_out1); + let fc2_out1 = self.fc2.forward(relu4_out1); + let norm2_out1 = self.norm2.forward(fc2_out1); + log_softmax(norm2_out1, 1) + } +} diff --git a/examples/import-model-weights/weights/mnist.pt b/examples/import-model-weights/weights/mnist.pt new file mode 100644 index 0000000..96225ff Binary files /dev/null and b/examples/import-model-weights/weights/mnist.pt differ diff --git a/examples/import-model-weights/weights/mnist.safetensors b/examples/import-model-weights/weights/mnist.safetensors new file mode 100644 index 0000000..95814d4 Binary files /dev/null and b/examples/import-model-weights/weights/mnist.safetensors differ diff --git a/examples/import-model-weights/weights/mnist_train_export.py b/examples/import-model-weights/weights/mnist_train_export.py new file mode 100755 index 0000000..38f22f6 --- /dev/null +++ b/examples/import-model-weights/weights/mnist_train_export.py @@ -0,0 +1,223 @@ +#!/usr/bin/env python3 + +# Originally copied and modified from: https://github.com/pytorch/examples/blob/main/mnist/main.py +# under the following license: BSD-3-Clause license + +from __future__ import print_function +import argparse +from safetensors.torch import save_file +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torchvision import datasets, transforms +from torch.optim.lr_scheduler import StepLR + + +class Net(nn.Module): + def __init__(self): + super(Net, self).__init__() + self.conv1 = nn.Conv2d(1, 8, 3) + self.conv2 = nn.Conv2d(8, 16, 3) + self.conv3 = nn.Conv2d(16, 24, 3) + self.norm1 = nn.BatchNorm2d(24) + self.dropout1 = nn.Dropout(0.3) + self.fc1 = nn.Linear(24 * 22 * 22, 32) + self.fc2 = nn.Linear(32, 10) + self.norm2 = nn.BatchNorm1d(10) + + def forward(self, x): + x = self.conv1(x) + x = F.relu(x) + x = self.conv2(x) + x = F.relu(x) + x = self.conv3(x) + x = F.relu(x) + x = self.norm1(x) + x = torch.flatten(x, 1) + x = self.fc1(x) + x = F.relu(x) + x = self.dropout1(x) + x = self.fc2(x) + x = self.norm2(x) + output = F.log_softmax(x, dim=1) + return output + + +def train(args, model, device, train_loader, optimizer, epoch): + model.train() + for batch_idx, (data, target) in enumerate(train_loader): + data, target = data.to(device), target.to(device) + optimizer.zero_grad() + output = model(data) + loss = F.nll_loss(output, target) + loss.backward() + optimizer.step() + if batch_idx % args.log_interval == 0: + print( + "Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}".format( + epoch, + batch_idx * len(data), + len(train_loader.dataset), + 100.0 * batch_idx / len(train_loader), + loss.item(), + ) + ) + if args.dry_run: + break + + +def test(model, device, test_loader): + model.eval() + test_loss = 0 + correct = 0 + with torch.no_grad(): + for data, target in test_loader: + data, target = data.to(device), target.to(device) + output = model(data) + # sum up batch loss + test_loss += F.nll_loss(output, target, reduction="sum").item() + # get the index of the max log-probability + pred = output.argmax(dim=1, keepdim=True) + correct += pred.eq(target.view_as(pred)).sum().item() + + test_loss /= len(test_loader.dataset) + + print( + "\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n".format( + test_loss, + correct, + len(test_loader.dataset), + 100.0 * correct / len(test_loader.dataset), + ) + ) + + +def main(): + # Training settings + parser = argparse.ArgumentParser(description="PyTorch MNIST Example") + parser.add_argument( + "--batch-size", + type=int, + default=64, + metavar="N", + help="input batch size for training (default: 64)", + ) + parser.add_argument( + "--test-batch-size", + type=int, + default=1000, + metavar="N", + help="input batch size for testing (default: 1000)", + ) + parser.add_argument( + "--epochs", + type=int, + default=8, + metavar="N", + help="number of epochs to train (default: 14)", + ) + parser.add_argument( + "--lr", + type=float, + default=1.0, + metavar="LR", + help="learning rate (default: 1.0)", + ) + parser.add_argument( + "--gamma", + type=float, + default=0.7, + metavar="M", + help="Learning rate step gamma (default: 0.7)", + ) + parser.add_argument( + "--no-cuda", action="store_true", default=False, help="disables CUDA training" + ) + parser.add_argument( + "--no-mps", + action="store_true", + default=False, + help="disables macOS GPU training", + ) + parser.add_argument( + "--dry-run", + action="store_true", + default=False, + help="quickly check a single pass", + ) + parser.add_argument( + "--seed", type=int, default=1, metavar="S", help="random seed (default: 1)" + ) + parser.add_argument( + "--log-interval", + type=int, + default=10, + metavar="N", + help="how many batches to wait before logging training status", + ) + parser.add_argument( + "--save-model", + action="store_true", + default=True, + help="For Saving the current Model", + ) + parser.add_argument( + "--export-onnx", + action="store_true", + default=False, + help="For Saving the current Model in ONNX format", + ) + args = parser.parse_args() + use_cuda = not args.no_cuda and torch.cuda.is_available() + use_mps = not args.no_mps and torch.backends.mps.is_available() + + torch.manual_seed(args.seed) + + if use_cuda: + device = torch.device("cuda") + elif use_mps: + device = torch.device("mps") + print("using MPS") + else: + device = torch.device("cpu") + + train_kwargs = {"batch_size": args.batch_size} + test_kwargs = {"batch_size": args.test_batch_size} + if use_cuda: + cuda_kwargs = {"num_workers": 1, "pin_memory": True, "shuffle": True} + train_kwargs.update(cuda_kwargs) + test_kwargs.update(cuda_kwargs) + + transform = transforms.Compose( + [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))] + ) + dataset1 = datasets.MNIST( + "/tmp/mnist-data", train=True, download=True, transform=transform + ) + dataset2 = datasets.MNIST("/tmp/mnist-data", train=False, transform=transform) + train_loader = torch.utils.data.DataLoader(dataset1, **train_kwargs) + test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs) + + model = Net().to(device) + optimizer = optim.Adadelta(model.parameters(), lr=args.lr) + + scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma) + for epoch in range(1, args.epochs + 1): + train(args, model, device, train_loader, optimizer, epoch) + test(model, device, test_loader) + scheduler.step() + + if args.save_model: + torch.save(model.state_dict(), "mnist.pt") + save_file(model.state_dict(), "mnist.safetensors") + + if args.export_onnx: + dummy_input = torch.randn(1, 1, 28, 28, device=device) + torch.onnx.export( + model, dummy_input, "mnist.onnx", verbose=True, opset_version=16 + ) + + +if __name__ == "__main__": + main() diff --git a/examples/mnist-inference-web/Cargo.toml b/examples/mnist-inference-web/Cargo.toml new file mode 100644 index 0000000..379a7cf --- /dev/null +++ b/examples/mnist-inference-web/Cargo.toml @@ -0,0 +1,29 @@ +[package] +authors = ["Dilshod Tadjibaev (@antimora)"] +edition.workspace = true +license = "MIT OR Apache-2.0" +name = "mnist-inference-web" +publish = false +version.workspace = true + +[lints] +workspace = true + +[lib] +crate-type = ["cdylib"] + +[features] +default = ["flex"] + +flex = ["burn/flex"] +wgpu = ["burn/wgpu"] + +[dependencies] +burn = { workspace = true, features = ["extension"]} +serde = { workspace = true } +console_error_panic_hook = { workspace = true } + +# Wasm dependencies +wasm-bindgen = "0.2" +wasm-bindgen-futures = "0.4" +js-sys = "0.3" diff --git a/examples/mnist-inference-web/README.md b/examples/mnist-inference-web/README.md new file mode 100644 index 0000000..3fe742a --- /dev/null +++ b/examples/mnist-inference-web/README.md @@ -0,0 +1,92 @@ +# MNIST Inference on Web + +[![Live Demo](https://img.shields.io/badge/live-demo-brightgreen)](https://burn.dev/demo) + +This crate demonstrates how to run an MNIST-trained model in the browser for inference. + +## Running + +1. Build + + ```shell + ./build-for-web.sh {backend} + ``` + + The backend can either be `flex` or `wgpu`. Note that `wgpu` only works for browsers with support for WebGPU. + +2. Run the server + + ```shell + ./run-server.sh + ``` + +3. Open the [`http://localhost:8000/`](http://localhost:8000/) in the browser. + +## Design + +The inference components of `burn` with the `flex` backend can be built with `#![no_std]`. This +makes it possible to build and run the model with the `wasm32-unknown-unknown` target without a +special system library, such as [WASI](https://wasi.dev/). (See [Cargo.toml](./Cargo.toml) on how to +include burn dependencies without `std`). + +For this demo, we use trained parameters (`model.bpk`, in the burnpack format) and model +(`model.rs`) from the +[`burn` MNIST example](https://github.com/tracel-ai/burn/tree/main/examples/mnist). + +The inference API for JavaScript is exposed with the help of +[`wasm-bindgen`](https://github.com/rustwasm/wasm-bindgen)'s library and tools. + +JavaScript (`index.js`) is used to transform hand-drawn digits to a format that the inference API +accepts. The transformation includes image cropping, scaling down, and converting it to grayscale +values. + +## Model + +Layers: + +1. Input Image (28,28, 1ch) +2. `Conv2d`(3x3, 64ch), `BatchNorm2d`, `Gelu`, `MaxPool`(2x2) +3. `Conv2d`(3x3, 16ch), `BatchNorm2d`, `Gelu`, `MaxPool`(2x2) +4. `Linear`(1600, 128), `Relu` +4. `Linear`(128, 128), `Relu` +5. `Linear`(128, 10) +6. Softmax Output + +The total number of parameters is 260,810. + +The model is trained with 18 epochs and the final test accuracy is 95.83%. + +Random transformations are used for data augmentation. + +The training and hyper parameter information in can be found in +[`burn` MNIST example](https://github.com/tracel-ai/burn/tree/main/examples/mnist). + +## Comparison + +The main differentiating factor of this example's approach (compiling rust model into wasm) and +other popular tools, such as [TensorFlow.js](https://www.tensorflow.org/js), +[ONNX Runtime JS](https://onnxruntime.ai/docs/tutorials/web/) and +[TVM Web](https://github.com/apache/tvm/tree/main/web) is the absence of runtime code. The rust +compiler optimizes and includes only used `burn` routines. 1,509,747 bytes out of Wasm's 1,866,491 +byte file is the model's parameters. The rest of 356,744 bytes contain all the code (including +`burn`'s `nn` components, the data deserialization library, and math operations). + +## Future Improvements + +There are several planned enhancements in place: + +- [#1271](https://github.com/rust-ndarray/ndarray/issues/1271) - + [WASM SIMD](https://github.com/WebAssembly/simd/blob/master/proposals/simd/SIMD.md) support in + NDArray that can speed up computation on CPU. + +## Acknowledgements + +Two online MNIST demos inspired and helped build this demo: +[MNIST Draw](https://mco-mnist-draw-rwpxka3zaa-ue.a.run.app/) by Marc (@mco-gh) and +[MNIST Web Demo](https://ufal.mff.cuni.cz/~straka/courses/npfl129/2223/demos/mnist_web.html) (no +code was copied but helped tremendously with an implementation approach). + +## Resources + +1. [Rust 🦀 and WebAssembly](https://rustwasm.github.io/docs/book/) +2. [wasm-bindgen](https://rustwasm.github.io/wasm-bindgen/) diff --git a/examples/mnist-inference-web/build-for-web.sh b/examples/mnist-inference-web/build-for-web.sh new file mode 100755 index 0000000..eebcef0 --- /dev/null +++ b/examples/mnist-inference-web/build-for-web.sh @@ -0,0 +1,18 @@ +#!/usr/bin/env bash + +# Add wasm32 target for compiler. +rustup target add wasm32-unknown-unknown + +if ! command -v wasm-pack &> /dev/null +then + echo "wasm-pack could not be found. Installing ..." + cargo install wasm-pack +fi + +# Set optimization flags +RUSTFLAGS="-C embed-bitcode=yes -C codegen-units=1 -C opt-level=3 --cfg web_sys_unstable_apis" + +# Run wasm pack tool to build JS wrapper files and copy wasm to pkg directory. +mkdir -p pkg +wasm-pack build --out-dir pkg --release --target web --no-typescript --no-default-features --features $1 + diff --git a/examples/mnist-inference-web/index.html b/examples/mnist-inference-web/index.html new file mode 100644 index 0000000..ed8f2d3 --- /dev/null +++ b/examples/mnist-inference-web/index.html @@ -0,0 +1,156 @@ + + + + + + Burn MNIST Inference Web Demo + + + + + + + + + + + + +

Burn MNIST Inference Demo

+ + + + + + + + + + + + + + + + + +
Draw a digit hereCropped and scaledProbability result
+ + + + + + +
+ +
+ + + + diff --git a/examples/mnist-inference-web/index.js b/examples/mnist-inference-web/index.js new file mode 100644 index 0000000..9896ff5 --- /dev/null +++ b/examples/mnist-inference-web/index.js @@ -0,0 +1,175 @@ +/** + * + * This demo is part of Burn project: https://github.com/tracel-ai/burn + * + * Released under a dual license: + * https://github.com/tracel-ai/burn/blob/main/LICENSE-MIT + * https://github.com/tracel-ai/burn/blob/main/LICENSE-APACHE + * + */ + +/** + * Auto crops the image, scales to 28x28 pixel image, and returns as grayscale image. + * @param {object} mainContext - The 2d context of the source canvas. + * @param {object} cropContext - The 2d context of an intermediate hidden canvas. + * @param {object} scaledContext - The 2d context of the destination 28x28 canvas. + */ +export function cropScaleGetImageData(mainContext, cropContext, scaledContext) { + + const cropEl = cropContext.canvas; + + // Get the auto-cropped image data and put into the intermediate/hidden canvas + cropContext.fillStyle = "rgba(255, 255, 255, 255)"; // white non-transparent color + cropContext.fillRect(0, 0, cropEl.width, cropEl.height); + cropContext.save(); + const [w, h, croppedImage] = cropImageFromCanvas(mainContext); + cropEl.width = Math.max(w, h) * 1.2; + cropEl.height = Math.max(w, h) * 1.2; + const leftPadding = (cropEl.width - w) / 2; + const topPadding = (cropEl.height - h) / 2; + cropContext.putImageData(croppedImage, leftPadding, topPadding); + + // Copy image data to scale 28x28 canvas + scaledContext.save(); + scaledContext.clearRect(0, 0, scaledContext.canvas.height, scaledContext.canvas.width); + scaledContext.fillStyle = "rgba(255, 255, 255, 255)"; // white non-transparent color + scaledContext.fillRect(0, 0, cropEl.width, cropEl.height); + scaledContext.scale(28.0 / cropContext.canvas.width, 28.0 / cropContext.canvas.height); + scaledContext.drawImage(cropEl, 0, 0); + + // Extract image data and convert into single value (greyscale) array + const data = rgba2gray(scaledContext.getImageData(0, 0, 28, 28).data); + scaledContext.restore(); + + return data; +} + +/** + * Converts RGBA image data from canvas to grayscale (0 is white & 255 is black). + * @param {int[]} - Image data. + */ +export function rgba2gray(data) { + let converted = new Float32Array(data.length / 4); + + // Data is stored as [r0,g0,b0,a0, ... r[n],g[n],b[n],a[n]] where n is number of pixels. + for (let i = 0; i < data.length; i += 4) { + let r = 255 - data[i]; // red + let g = 255 - data[i + 1]; // green + let b = 255 - data[i + 2]; // blue + let a = 255 - data[i + 3]; // alpha + + // Use RGB grayscale coefficients (https://imagej.nih.gov/ij/docs/menus/image.html) + let y = 0.299 * r + 0.587 * g + 0.114 * b; + converted[i / 4] = y; // 4 times fewer data points but the same number of pixels. + } + return converted; +} + +/** + * Auto crops a canvas images and returns its image data. + * @param {object} ctx - canvas 2d context. + * src: https://stackoverflow.com/a/22267731 + */ +export function cropImageFromCanvas(ctx) { + let canvas = ctx.canvas, + w = canvas.width, + h = canvas.height, + pix = { x: [], y: [] }, + imageData = ctx.getImageData(0, 0, canvas.width, canvas.height), + x, + y, + index; + for (y = 0; y < h; y++) { + for (x = 0; x < w; x++) { + index = (y * w + x) * 4; + + let r = imageData.data[index]; + let g = imageData.data[index + 1]; + let b = imageData.data[index + 2]; + // On some browsers the canvas has a grey border which prevents cropping if we do min != 255 + if (Math.min(r, g, b) < 240) { + pix.x.push(x); + pix.y.push(y); + } + } + } + pix.x.sort(function (a, b) { + return a - b; + }); + pix.y.sort(function (a, b) { + return a - b; + }); + let n = pix.x.length - 1; + w = 1 + pix.x[n] - pix.x[0]; + h = 1 + pix.y[n] - pix.y[0]; + return [w, h, ctx.getImageData(pix.x[0], pix.y[0], w, h, { willReadFrequently: true })]; +} + +/** + * Truncates number to a given decimal position + * @param {number} num - Number to truncate. + * @param {number} fixed - Decimal positions. + * src: https://stackoverflow.com/a/11818658 + */ +export function toFixed(num, fixed) { + const re = new RegExp('^-?\\d+(?:\.\\d{0,' + (fixed || -1) + '})?'); + return num.toString().match(re)[0]; +} + +/** + * Looks up element by an id. + * @param {string} - Element id. + */ +export function $(id) { + return document.getElementById(id); +} + +/** + * Helper function that builds a chart using Chart.js library. + * @param {object} chartEl - Chart canvas element. + * + * NOTE: Assumes chart.js is loaded into the global. + */ +export function chartConfigBuilder(chartEl) { + Chart.register(ChartDataLabels); + return new Chart(chartEl, { + plugins: [ChartDataLabels], + type: "bar", + data: { + labels: ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"], + datasets: [ + { + data: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + borderWidth: 0, + fill: true, + backgroundColor: "#247ABF", + }, + ], + }, + options: { + responsive: false, + maintainAspectRatio: false, + animation: true, + plugins: { + legend: { + display: false, + }, + tooltip: { + enabled: true, + }, + datalabels: { + color: "white", + formatter: function (value, context) { + return toFixed(value, 2); + }, + }, + }, + scales: { + y: { + beginAtZero: true, + max: 1.0, + }, + }, + }, + }); +} \ No newline at end of file diff --git a/examples/mnist-inference-web/model.bpk b/examples/mnist-inference-web/model.bpk new file mode 100644 index 0000000..0eb0e55 Binary files /dev/null and b/examples/mnist-inference-web/model.bpk differ diff --git a/examples/mnist-inference-web/run-server.sh b/examples/mnist-inference-web/run-server.sh new file mode 100755 index 0000000..0ce038f --- /dev/null +++ b/examples/mnist-inference-web/run-server.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash + +# Opening index.html file directly by a browser does not work because of +# the security restrictions by the browser. Viewing the HTML file will fail with +# this error message: + +# ``` +# Access to script at +# 'file:///Users/user/Projects/burn-mac/examples/mnist-inference-web/pkg/mnist_inference_web.js' +# from origin 'null' has been blocked by CORS policy: +# Cross origin requests are only supported for protocol schemes: +# http, data, isolated-app, chrome-extension, chrome, https, chrome-untrusted. +# ``` +# So that's why running a local HTTP server is needed. + +if ! command -v python3 &> /dev/null +then + echo "python3 could not be found. Running server requires python3." + exit +fi + +echo "Running local python HTTP server on port 8000 ..." +python3 -m http.server 8000 diff --git a/examples/mnist-inference-web/src/lib.rs b/examples/mnist-inference-web/src/lib.rs new file mode 100644 index 0000000..3f192ad --- /dev/null +++ b/examples/mnist-inference-web/src/lib.rs @@ -0,0 +1,7 @@ +#![cfg_attr(not(test), no_std)] + +pub mod model; +pub mod state; +pub mod web; + +extern crate alloc; diff --git a/examples/mnist-inference-web/src/model.rs b/examples/mnist-inference-web/src/model.rs new file mode 100644 index 0000000..4856399 --- /dev/null +++ b/examples/mnist-inference-web/src/model.rs @@ -0,0 +1,104 @@ +use burn::{ + nn::{ + BatchNorm, PaddingConfig2d, + pool::{MaxPool2d, MaxPool2dConfig}, + }, + prelude::*, +}; + +#[derive(Module, Debug)] +pub struct Model { + conv1: ConvBlock, + conv2: ConvBlock, + dropout: nn::Dropout, + fc1: nn::Linear, + fc2: nn::Linear, + fc3: nn::Linear, + activation: nn::Gelu, +} + +const NUM_CLASSES: usize = 10; + +impl Model { + pub fn new(device: &Device) -> Self { + let conv1 = ConvBlock::new([1, 64], [3, 3], device, true); // out: max_pool -> [Batch,32,13,13] + let conv2 = ConvBlock::new([64, 64], [3, 3], device, true); // out: max_pool -> [Batch,64,5,5] + let hidden_size = 64 * 5 * 5; + let fc1 = nn::LinearConfig::new(hidden_size, 128).init(device); + let fc2 = nn::LinearConfig::new(128, 128).init(device); + let fc3 = nn::LinearConfig::new(128, NUM_CLASSES).init(device); + + let dropout = nn::DropoutConfig::new(0.25).init(); + + Self { + conv1, + conv2, + dropout, + fc1, + fc2, + fc3, + activation: nn::Gelu::new(), + } + } + + pub fn forward(&self, input: Tensor<3>) -> Tensor<2> { + let [batch_size, height, width] = input.dims(); + + let x = input.reshape([batch_size, 1, height, width]).detach(); + let x = self.conv1.forward(x); + let x = self.conv2.forward(x); + + let [batch_size, channels, height, width] = x.dims(); + let x = x.reshape([batch_size, channels * height * width]); + + let x = self.fc1.forward(x); + let x = self.activation.forward(x); + let x = self.dropout.forward(x); + let x = self.fc2.forward(x); + let x = self.activation.forward(x); + let x = self.dropout.forward(x); + + self.fc3.forward(x) + } +} + +#[derive(Module, Debug)] +pub struct ConvBlock { + conv: nn::conv::Conv2d, + norm: BatchNorm, + pool: Option, + activation: nn::Relu, +} + +impl ConvBlock { + pub fn new(channels: [usize; 2], kernel_size: [usize; 2], device: &Device, pool: bool) -> Self { + let conv = nn::conv::Conv2dConfig::new(channels, kernel_size) + .with_padding(PaddingConfig2d::Valid) + .init(device); + let norm = nn::BatchNormConfig::new(channels[1]).init(device); + let pool = if pool { + Some(MaxPool2dConfig::new([2, 2]).with_strides([2, 2]).init()) + } else { + None + }; + + Self { + conv, + norm, + pool, + activation: nn::Relu::new(), + } + } + + pub fn forward(&self, input: Tensor<4>) -> Tensor<4> { + let x = self.conv.forward(input); + let x = self.norm.forward(x); + let x = self.activation.forward(x); + + if let Some(pool) = &self.pool { + pool.forward(x) + } else { + x + } + } +} diff --git a/examples/mnist-inference-web/src/state.rs b/examples/mnist-inference-web/src/state.rs new file mode 100644 index 0000000..0b7490a --- /dev/null +++ b/examples/mnist-inference-web/src/state.rs @@ -0,0 +1,38 @@ +use crate::model::Model; +use burn::{module::Module, prelude::Device, store::ModuleRecord, tensor::Bytes}; + +// Trained parameters in the burnpack format, produced by the `mnist` example +// (`model.into_record().save(..)`) and copied here. Regenerate with the same command if the +// model architecture changes. +static STATE_ENCODED: &[u8] = include_bytes!("../model.bpk"); + +/// Builds and loads trained parameters into the model. +pub async fn build_and_load_model() -> Model { + #[cfg(all(feature = "flex", not(feature = "wgpu")))] + let device = Device::flex(); + #[cfg(feature = "wgpu")] + // Calls init_setup_async + let device = Device::wgpu_async(Default::default()).await; + + let model = Model::new(&device); + let record = ModuleRecord::from_bytes(Bytes::from_bytes_vec(STATE_ENCODED.to_vec())) + .expect("Failed to decode state"); + + model.load_record(record) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn embedded_model_decodes_into_architecture() { + let device = Device::flex(); + let model = Model::new(&device); + // `load_record` validates that every model parameter is present with a matching shape; a + // stale/mismatched asset would panic here. + let record = ModuleRecord::from_bytes(Bytes::from_bytes_vec(STATE_ENCODED.to_vec())) + .expect("Embedded model.bpk should decode as burnpack"); + let _model = model.load_record(record); + } +} diff --git a/examples/mnist-inference-web/src/web.rs b/examples/mnist-inference-web/src/web.rs new file mode 100644 index 0000000..57796e0 --- /dev/null +++ b/examples/mnist-inference-web/src/web.rs @@ -0,0 +1,80 @@ +#![allow(clippy::new_without_default)] + +use alloc::string::String; +use js_sys::Array; + +#[cfg(target_family = "wasm")] +use wasm_bindgen::prelude::*; + +use crate::model::Model; +use crate::state::build_and_load_model; + +use burn::tensor::Tensor; + +#[cfg_attr(target_family = "wasm", wasm_bindgen(start))] +pub fn start() { + console_error_panic_hook::set_once(); +} + +/// Mnist structure that corresponds to JavaScript class. +/// See:[exporting-rust-struct](https://rustwasm.github.io/wasm-bindgen/contributing/design/exporting-rust-struct.html) +#[cfg_attr(target_family = "wasm", wasm_bindgen)] +pub struct Mnist { + model: Option, +} + +#[cfg_attr(target_family = "wasm", wasm_bindgen)] +impl Mnist { + /// Constructor called by JavaScripts with the new keyword. + #[cfg_attr(target_family = "wasm", wasm_bindgen(constructor))] + pub fn new() -> Self { + console_error_panic_hook::set_once(); + Self { model: None } + } + + /// Returns the inference results. + /// + /// This method is called from JavaScript via generated wrapper code by wasm-bindgen. + /// + /// # Arguments + /// + /// * `input` - A f32 slice of input 28x28 image + /// + /// See bindgen support types for passing and returning arrays: + /// * [number-slices](https://rustwasm.github.io/wasm-bindgen/reference/types/number-slices.html) + /// * [boxed-number-slices](https://rustwasm.github.io/wasm-bindgen/reference/types/boxed-number-slices.html) + /// + pub async fn inference(&mut self, input: &[f32]) -> Result { + if self.model.is_none() { + self.model = Some(build_and_load_model().await); + } + + let model = self.model.as_ref().unwrap(); + + let device = Default::default(); + // Reshape from the 1D array to 3d tensor [batch, height, width] + let input = Tensor::<1>::from_floats(input, &device).reshape([1, 28, 28]); + + // Normalize input: make between [0,1] and make the mean=0 and std=1 + // values mean=0.1307,std=0.3081 were copied from Pytorch Mist Example + // https://github.com/pytorch/examples/blob/54f4572509891883a947411fd7239237dd2a39c3/mnist/main.py#L122 + + let input = ((input / 255) - 0.1307) / 0.3081; + + // Run the tensor input through the model + let output: Tensor<2> = model.forward(input); + + // Convert the model output into probability distribution using softmax formula + let output = burn::tensor::activation::softmax(output, 1); + + // Flatten output tensor with [1, 10] shape into boxed slice of [f32] + let output = output.into_data_async().await.unwrap(); + + let array = Array::new(); + for value in output.iter::() { + array.push(&value.into()); + } + + Ok(array) + } +} diff --git a/examples/mnist/Cargo.toml b/examples/mnist/Cargo.toml new file mode 100644 index 0000000..1922c45 --- /dev/null +++ b/examples/mnist/Cargo.toml @@ -0,0 +1,34 @@ +[package] +authors = ["nathanielsimard "] +edition.workspace = true +license.workspace = true +name = "mnist" +publish = false +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["burn/default", "burn/tui"] +flex = ["burn/flex"] +tch-cpu = ["burn/tch"] +tch-gpu = ["burn/tch"] +wgpu = ["burn/wgpu"] +metal = ["burn/metal"] +cuda = ["burn/cuda"] +vulkan = ["burn/vulkan"] +rocm = ["burn/rocm"] +remote = ["burn/remote-websocket"] + +[dependencies] +burn = { workspace = true, features = [ + "train", + "vision", + "metrics", + "fusion", + "flex", +] } + +log = { workspace = true } +rand.workspace = true diff --git a/examples/mnist/README.md b/examples/mnist/README.md new file mode 100644 index 0000000..28af009 --- /dev/null +++ b/examples/mnist/README.md @@ -0,0 +1,23 @@ +# MNIST + +The example is showing you how to: + +- Define your own custom module (MLP). +- Create the data pipeline from a raw dataset to a batched multi-threaded fast DataLoader. +- Configure a learner to display and log metrics as well as to keep training checkpoints. + +The example can be run like so: + +```bash +git clone https://github.com/tracel-ai/burn.git +cd burn +# Use the --release flag to really speed up training. +echo "Using flex backend" +cargo run --example mnist --release --features flex # CPU Flex Backend - f32 +echo "Using tch backend" +export TORCH_CUDA_VERSION=cu128 # Set the cuda version +cargo run --example mnist --release --features tch-gpu # GPU Tch Backend - f32 +cargo run --example mnist --release --features tch-cpu # CPU Tch Backend - f32 +echo "Using vulkan backend" +cargo run --example mnist --release --features vulkan +``` diff --git a/examples/mnist/cubecl.toml b/examples/mnist/cubecl.toml new file mode 100644 index 0000000..c289ed4 --- /dev/null +++ b/examples/mnist/cubecl.toml @@ -0,0 +1,15 @@ +[profiling] +logger = { log = "info", level = "disabled" } + +[autotune] +level = "balanced" +cache = "target" +logger = { file = "/tmp/autotune.log", level = "disabled" } + +[compilation] +logger = { level = "disabled" } +cache = "target" + +[memory] +logger = { level = "disabled", file = "/tmp/memory.log" } +persistent_memory = "enabled" diff --git a/examples/mnist/examples/mnist.rs b/examples/mnist/examples/mnist.rs new file mode 100644 index 0000000..f6ec32c --- /dev/null +++ b/examples/mnist/examples/mnist.rs @@ -0,0 +1,42 @@ +#![recursion_limit = "256"] + +use burn::tensor::Device; +use mnist::training; + +#[allow(unreachable_code)] +fn select_device() -> Device { + #[cfg(feature = "flex")] + return Device::flex(); + + #[cfg(all(feature = "tch-gpu", not(target_os = "macos")))] + return Device::libtorch_cuda(burn::tensor::DeviceIndex::Default); + + #[cfg(all(feature = "tch-gpu", target_os = "macos"))] + return Device::libtorch_mps(); + + #[cfg(feature = "tch-cpu")] + return Device::libtorch(); + + #[cfg(feature = "vulkan")] + return Device::vulkan(burn::tensor::DeviceKind::DefaultDevice); + #[cfg(feature = "metal")] + return Device::metal(burn::tensor::DeviceKind::DefaultDevice); + #[cfg(feature = "wgpu")] + return Device::wgpu(burn::tensor::DeviceKind::DefaultDevice); + + #[cfg(feature = "cuda")] + return Device::cuda(burn::tensor::DeviceIndex::Default); + + #[cfg(feature = "rocm")] + return Device::rocm(burn::tensor::DeviceIndex::Default); + + #[cfg(feature = "remote")] + return Device::remote_websocket("ws://localhost:3000", 0); + + unreachable!("At least one backend will be selected.") +} + +fn main() { + let device = select_device(); + training::run(device); +} diff --git a/examples/mnist/src/data.rs b/examples/mnist/src/data.rs new file mode 100644 index 0000000..7f38c0c --- /dev/null +++ b/examples/mnist/src/data.rs @@ -0,0 +1,159 @@ +use std::{f32::consts::FRAC_PI_4, fmt::Display}; + +use burn::{ + data::{ + dataloader::batcher::Batcher, + dataset::{transform::Mapper, vision::MnistItem}, + }, + prelude::*, + vision::Transform2D, +}; +use rand::RngExt; + +#[derive(Clone, Debug, Default)] +pub struct MnistBatcher {} + +#[derive(Clone, Debug)] +pub struct MnistBatch { + pub images: Tensor<3>, + pub targets: Tensor<1, Int>, +} + +impl Batcher for MnistBatcher { + fn batch(&self, items: Vec, device: &Device) -> MnistBatch { + let images = items.iter().map(|item| item.image.clone()).collect(); + + let targets = items + .iter() + .map(|item| { + Tensor::<1, Int>::from_data(TensorData::from([item.label as i64]), &Device::flex()) + }) + .collect(); + + let images = Tensor::cat(images, 0); + let images = Tensor::from_data(images.into_data(), device); + + let targets = Tensor::cat(targets, 0); + let targets = Tensor::from_data(targets.into_data(), device); + + MnistBatch { images, targets } + } +} + +#[derive(Clone, Debug, Copy)] +pub enum Transform { + Translate, + Shear, + Scale, + Rotation, +} + +impl Display for Transform { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Transform::Translate => f.write_str("Tr"), + Transform::Shear => f.write_str("Sr"), + Transform::Scale => f.write_str("Sc"), + Transform::Rotation => f.write_str("Rot"), + } + } +} + +#[derive(Default)] +pub struct MnistMapper { + transforms: Vec, +} + +impl MnistMapper { + pub fn transform(mut self, transforms: &[Transform]) -> Self { + for t in transforms { + self.transforms.push(*t); + } + self + } + pub fn translate(mut self) -> Self { + self.transforms.push(Transform::Translate); + self + } + pub fn shear(mut self) -> Self { + self.transforms.push(Transform::Shear); + self + } + pub fn scale(mut self) -> Self { + self.transforms.push(Transform::Scale); + self + } + pub fn rotation(mut self) -> Self { + self.transforms.push(Transform::Rotation); + self + } +} + +impl Mapper for MnistMapper { + fn map(&self, item: &MnistItem) -> MnistItemPrepared { + prepare_image(&self.transforms, item.clone()) + } +} + +#[derive(Clone, Debug)] +pub struct MnistItemPrepared { + image: Tensor<3>, + label: u8, +} + +fn prepare_image(transforms: &[Transform], item: MnistItem) -> MnistItemPrepared { + let data = TensorData::from(item.image); + let tensor = Tensor::<2>::from_data(data.convert::(), &Device::flex()); + let tensor = tensor.reshape([1, 28, 28]); + + // normalize: make between [0,1] and make the mean = 0 and std = 1 + // values mean=0.1307,std=0.3081 were copied from Pytorch Mist Example + // https://github.com/pytorch/examples/blob/54f4572509891883a947411fd7239237dd2a39c3/mnist/main.py#L122 + let tensor = ((tensor / 255) - 0.1307) / 0.3081; + let tensor = if !transforms.is_empty() { + mangle_image_batch(transforms, tensor) + } else { + tensor + }; + + MnistItemPrepared { + image: tensor, + label: item.label, + } +} + +/// Mange the image by applying small random transformations to augment the dataset. +/// +/// * `images` - The images with shape [batch size, height, width] +/// +/// ## Return +/// +/// The transformed images tensor with shape [batch size, height, width] +fn mangle_image_batch(transforms: &[Transform], images: Tensor<3>) -> Tensor<3> { + let mut rng = rand::rng(); + + let transforms = transforms.iter().map(|transform| match transform { + Transform::Translate => { + Transform2D::translation(rng.random_range(-0.2..0.2), rng.random_range(-0.2..0.2)) + } + Transform::Shear => Transform2D::shear( + rng.random_range(-0.6..0.6), + rng.random_range(-0.6..0.6), + 0.0, + 0.0, + ), + Transform::Scale => Transform2D::scale( + rng.random_range(0.6..1.5), + rng.random_range(0.6..1.5), + 0.0, + 0.0, + ), + Transform::Rotation => { + Transform2D::rotation(rng.random_range(-FRAC_PI_4..FRAC_PI_4), 0.0, 0.0) + } + }); + + Transform2D::composed(transforms) + .transform(images.unsqueeze_dim::<4>(1)) + .squeeze_dims::<3>(&[1]) +} diff --git a/examples/mnist/src/file_progress.rs b/examples/mnist/src/file_progress.rs new file mode 100644 index 0000000..e98124a --- /dev/null +++ b/examples/mnist/src/file_progress.rs @@ -0,0 +1,159 @@ +use std::{ + collections::HashMap, + fs::{File, OpenOptions}, + io::Write, + path::Path, +}; + +use burn::train::logger::{EvaluationProgressLogger, TrainingProgressLogger}; + +/// A progress logger that appends training progress to a file. +/// +/// Each event is written as a new line so the complete history across all +/// epochs and splits is preserved. Safe to use alongside a TUI renderer. +/// +/// # Example +/// +/// ```ignore +/// let training = SupervisedTraining::new(/* ... */) +/// .with_progress_logger( +/// FileTrainingProgressLogger::new("/tmp/my-run/training_progress.log").unwrap() +/// ); +/// ``` +pub struct FileTrainingProgressLogger { + writer: File, + event_counters: HashMap, +} + +impl FileTrainingProgressLogger { + /// Opens (or creates) the file at `path` in append mode. + pub fn new(path: impl AsRef) -> std::io::Result { + let file = OpenOptions::new().create(true).append(true).open(path)?; + Ok(Self { + writer: file, + event_counters: HashMap::new(), + }) + } + + fn write(&mut self, line: &str) { + if let Err(e) = writeln!(self.writer, "{line}") { + log::warn!("FileTrainingProgressLogger write error: {e}"); + } + } +} + +impl TrainingProgressLogger for FileTrainingProgressLogger { + fn start(&mut self, total_epochs: usize, starting_epoch: usize, total_items: Option) { + match total_items { + Some(n) => self.write(&format!( + "[Training] start epochs={total_epochs} total_items={n} starting_epoch={starting_epoch}" + )), + None => self.write(&format!("[Training] start epochs={total_epochs} starting_epoch={starting_epoch}")), + } + } + + fn update_epoch(&mut self, epoch: usize) { + self.write(&format!("[Training] epoch_complete epoch={epoch}")); + } + + fn start_split(&mut self, split: &str, total_items: usize) { + self.write(&format!( + "[Training] split_start split={split} total_items={total_items}" + )); + } + + fn update_split(&mut self, items_processed: usize) { + self.write(&format!( + "[Training] split_update items_processed={items_processed}" + )); + } + + fn end_split(&mut self) { + self.write("[Training] split_end"); + self.event_counters.values_mut().for_each(|v| *v = 0); + } + + fn end(&mut self) { + self.write("[Training] end"); + } + + fn log_event_training(&mut self, event: String) { + let count = { + let c = self.event_counters.entry(event.clone()).or_insert(0); + *c += 1; + *c + }; + self.write(&format!("[event] {event} = {count}")); + } +} + +/// A progress logger that appends evaluation progress to a file. +/// +/// Each event is written as a new line so the complete history across all +/// test splits is preserved. Safe to use alongside a TUI renderer. +/// +/// # Example +/// +/// ```ignore +/// let evaluator = EvaluatorBuilder::new(/* ... */) +/// .with_progress_logger( +/// FileEvaluationProgressLogger::new("/tmp/my-run/evaluation_progress.log").unwrap() +/// ); +/// ``` +pub struct FileEvaluationProgressLogger { + writer: File, + event_counters: HashMap, +} + +impl FileEvaluationProgressLogger { + /// Opens (or creates) the file at `path` in append mode. + pub fn new(path: impl AsRef) -> std::io::Result { + let file = OpenOptions::new().create(true).append(true).open(path)?; + Ok(Self { + writer: file, + event_counters: HashMap::new(), + }) + } + + fn write(&mut self, line: &str) { + if let Err(e) = writeln!(self.writer, "{line}") { + log::warn!("FileEvaluationProgressLogger write error: {e}"); + } + } +} + +impl EvaluationProgressLogger for FileEvaluationProgressLogger { + fn start_global_progress(&mut self, total_tests: usize) { + self.write(&format!("[Evaluation] start total_tests={total_tests}")); + } + + fn start_test(&mut self, name: &str, total_items: usize) { + self.write(&format!( + "[Evaluation] test_start name={name} total_items={total_items}" + )); + } + + fn update_test_progress(&mut self, items_processed: usize) { + self.write(&format!( + "[Evaluation] test_update items_processed={items_processed}" + )); + } + + fn end_test(&mut self) { + self.write("[Evaluation] test_end"); + self.event_counters.values_mut().for_each(|v| *v = 0); + } + + fn end_global_progress(&mut self) { + self.write("[Evaluation] end"); + } + + fn log_event_evaluation(&mut self, event: String) { + let count = { + let c = self.event_counters.entry(event.clone()).or_insert(0); + *c += 1; + *c + }; + self.write(&format!("[event] {event} = {count}")); + } +} diff --git a/examples/mnist/src/lib.rs b/examples/mnist/src/lib.rs new file mode 100644 index 0000000..bac3196 --- /dev/null +++ b/examples/mnist/src/lib.rs @@ -0,0 +1,4 @@ +pub mod data; +pub mod file_progress; +pub mod model; +pub mod training; diff --git a/examples/mnist/src/model.rs b/examples/mnist/src/model.rs new file mode 100644 index 0000000..f4f7678 --- /dev/null +++ b/examples/mnist/src/model.rs @@ -0,0 +1,142 @@ +use crate::data::MnistBatch; +use burn::{ + nn::{ + BatchNorm, PaddingConfig2d, + loss::CrossEntropyLossConfig, + pool::{MaxPool2d, MaxPool2dConfig}, + }, + prelude::*, + train::{ClassificationOutput, InferenceStep, TrainOutput, TrainStep}, +}; + +#[derive(Module, Debug)] +pub struct Model { + conv1: ConvBlock, + conv2: ConvBlock, + dropout: nn::Dropout, + fc1: nn::Linear, + fc2: nn::Linear, + fc3: nn::Linear, + activation: nn::Gelu, +} + +const NUM_CLASSES: usize = 10; + +impl Model { + pub fn new(device: &Device) -> Self { + let conv1 = ConvBlock::new([1, 64], [3, 3], device, true); // out: max_pool -> [Batch,32,13,13] + let conv2 = ConvBlock::new([64, 64], [3, 3], device, true); // out: max_pool -> [Batch,64,5,5] + let hidden_size = 64 * 5 * 5; + let fc1 = nn::LinearConfig::new(hidden_size, 128).init(device); + let fc2 = nn::LinearConfig::new(128, 128).init(device); + let fc3 = nn::LinearConfig::new(128, NUM_CLASSES).init(device); + + let dropout = nn::DropoutConfig::new(0.25).init(); + + Self { + conv1, + conv2, + dropout, + fc1, + fc2, + fc3, + activation: nn::Gelu::new(), + } + } + + pub fn forward(&self, input: Tensor<3>) -> Tensor<2> { + let [batch_size, height, width] = input.dims(); + + let x = input.reshape([batch_size, 1, height, width]).detach(); + let x = self.conv1.forward(x); + let x = self.conv2.forward(x); + + let [batch_size, channels, height, width] = x.dims(); + let x = x.reshape([batch_size, channels * height * width]); + + let x = self.fc1.forward(x); + let x = self.activation.forward(x); + let x = self.dropout.forward(x); + + let x = self.fc2.forward(x); + let x = self.activation.forward(x); + let x = self.dropout.forward(x); + + self.fc3.forward(x) + } + + pub fn forward_classification(&self, item: MnistBatch) -> ClassificationOutput { + let targets = item.targets; + let output = self.forward(item.images); + let loss = CrossEntropyLossConfig::new() + .init(&output.device()) + .forward(output.clone(), targets.clone()); + + ClassificationOutput { + loss, + output, + targets, + } + } +} + +#[derive(Module, Debug)] +pub struct ConvBlock { + conv: nn::conv::Conv2d, + norm: BatchNorm, + pool: Option, + activation: nn::Relu, +} + +impl ConvBlock { + pub fn new(channels: [usize; 2], kernel_size: [usize; 2], device: &Device, pool: bool) -> Self { + let conv = nn::conv::Conv2dConfig::new(channels, kernel_size) + .with_padding(PaddingConfig2d::Valid) + .init(device); + let norm = nn::BatchNormConfig::new(channels[1]).init(device); + let pool = if pool { + Some(MaxPool2dConfig::new([2, 2]).with_strides([2, 2]).init()) + } else { + None + }; + + Self { + conv, + norm, + pool, + activation: nn::Relu::new(), + } + } + + pub fn forward(&self, input: Tensor<4>) -> Tensor<4> { + let x = self.conv.forward(input); + let x = self.norm.forward(x); + let x = self.activation.forward(x); + + if let Some(pool) = &self.pool { + pool.forward(x) + } else { + x + } + } +} + +impl TrainStep for Model { + type Input = MnistBatch; + type Output = ClassificationOutput; + + fn step(&self, item: MnistBatch) -> TrainOutput { + let item = self.forward_classification(item); + + TrainOutput::new(self, item.loss.backward(), item) + } +} + +impl InferenceStep for Model { + type Input = MnistBatch; + type Output = ClassificationOutput; + + fn step(&self, item: MnistBatch) -> ClassificationOutput { + self.forward_classification(item) + } +} diff --git a/examples/mnist/src/training.rs b/examples/mnist/src/training.rs new file mode 100644 index 0000000..57a4734 --- /dev/null +++ b/examples/mnist/src/training.rs @@ -0,0 +1,264 @@ +use std::{path::PathBuf, sync::Arc}; + +use crate::{ + data::{MnistBatcher, MnistItemPrepared, MnistMapper, Transform}, + file_progress::{FileEvaluationProgressLogger, FileTrainingProgressLogger}, + model::Model, +}; + +use burn::{ + data::{ + dataloader::DataLoaderBuilder, + dataset::{ + Dataset, + transform::{ComposedDataset, MapperDataset, PartialDataset, SamplerDataset}, + vision::{MnistDataset, MnistItem}, + }, + }, + lr_scheduler::{ + composed::ComposedLrSchedulerConfig, cosine::CosineAnnealingLrSchedulerConfig, + linear::LinearLrSchedulerConfig, + }, + prelude::*, + train::{ + EvaluatorBuilder, Learner, MetricEarlyStoppingStrategy, StoppingCondition, + metric::{ + AccuracyMetric, LearningRateMetric, LossMetric, + store::{Aggregate, Direction, Split}, + }, + }, +}; +use burn::{optim::AdamWConfig, train::SupervisedTraining}; + +static ARTIFACT_DIR: &str = "/tmp/burn-example-mnist"; + +#[derive(Config, Debug)] +pub struct MnistTrainingConfig { + #[config(default = 5)] + pub num_epochs: usize, + + #[config(default = 256)] + pub batch_size: usize, + + #[config(default = 8)] + pub num_workers: usize, + + #[config(default = 42)] + pub seed: u64, + + pub optimizer: AdamWConfig, +} + +fn create_artifact_dir(artifact_dir: &str) { + std::fs::remove_file(PathBuf::from(artifact_dir).join("experiment.log")).ok(); + std::fs::create_dir_all(artifact_dir).ok(); +} + +pub fn run(device: Device) { + create_artifact_dir(ARTIFACT_DIR); + // Config + let config_optimizer = AdamWConfig::new() + .with_cautious_weight_decay(true) + .with_weight_decay(5e-5); + + let config = MnistTrainingConfig::new(config_optimizer); + + device.seed(config.seed); + let autodiff_device = device.clone().autodiff(); + + let model = Model::new(&autodiff_device); + + let dataset_train_original = Arc::new(MnistDataset::train()); + let dataset_train_plain = PartialDataset::new(dataset_train_original.clone(), 0, 55_000); + let dataset_valid_plain = PartialDataset::new(dataset_train_original.clone(), 55_000, 60_000); + + let ident_trains = generate_idents(Some(10000)); + let ident_valid = generate_idents(None); + let dataset_train = DatasetIdent::compose(ident_trains, dataset_train_plain); + let dataset_valid = DatasetIdent::compose(ident_valid, dataset_valid_plain); + + let dataloader_train = DataLoaderBuilder::new(MnistBatcher::default()) + .batch_size(config.batch_size) + .shuffle(config.seed) + .num_workers(config.num_workers) + .build(dataset_train); + let dataloader_valid = DataLoaderBuilder::new(MnistBatcher::default()) + .batch_size(config.batch_size) + .shuffle(config.seed) + .num_workers(config.num_workers) + .build(dataset_valid); + let lr_scheduler = ComposedLrSchedulerConfig::new() + .cosine(CosineAnnealingLrSchedulerConfig::new(1.0, 2000)) + // Warmup + .linear(LinearLrSchedulerConfig::new(1e-8, 1.0, 2000)) + .linear(LinearLrSchedulerConfig::new(1e-2, 1e-6, 10000)); + + let training = SupervisedTraining::new(ARTIFACT_DIR, dataloader_train, dataloader_valid) + .metrics((AccuracyMetric::new(), LossMetric::new())) + .metric_train_numeric(LearningRateMetric::new()) + .with_default_checkpointers() + .early_stopping(MetricEarlyStoppingStrategy::new( + &LossMetric::new(), + Aggregate::Mean, + Direction::Lowest, + Split::Valid, + StoppingCondition::NoImprovementSince { n_epochs: 5 }, + )) + .with_progress_logger( + FileTrainingProgressLogger::new(format!("{ARTIFACT_DIR}/training_progress.log")) + .expect("Failed to create training progress log"), + ) + .num_epochs(config.num_epochs) + .summary(); + + let result = training.launch(Learner::new( + model, + config.optimizer.init(), + lr_scheduler.init().unwrap(), + )); + + let dataset_test_plain = Arc::new(MnistDataset::test()); + + let splits: Vec<_> = generate_idents(None) + .into_iter() + .map(|(ident, _)| { + let name = ident.to_string(); + let dataset_test = DatasetIdent::prepare(ident, dataset_test_plain.clone()); + let dataloader = DataLoaderBuilder::new(MnistBatcher::default()) + .batch_size(config.batch_size) + .num_workers(2) + .build(dataset_test); + (name, dataloader) + }) + .collect(); + + let mut renderer = EvaluatorBuilder::new(ARTIFACT_DIR) + .renderer(result.renderer) + .metrics((AccuracyMetric::new(), LossMetric::new())) + .with_progress_logger( + FileEvaluationProgressLogger::new(format!("{ARTIFACT_DIR}/evaluation_progress.log")) + .expect("Failed to create evaluation progress log"), + ) + .summary() + .build(result.model.clone()) + .eval_all(splits); + + result + .model + .into_record() + .save(format!("{ARTIFACT_DIR}/model")) + .expect("Failed to save trained model"); + + config + .save(format!("{ARTIFACT_DIR}/config.json").as_str()) + .unwrap(); + + renderer.manual_close(); +} + +enum DatasetIdent { + Plain, + Transformed(Vec), + All, +} + +impl core::fmt::Display for DatasetIdent { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + DatasetIdent::Plain => f.write_str("Plain")?, + DatasetIdent::Transformed(items) => { + for i in 0..items.len() { + f.write_fmt(format_args!("{}", items[i]))?; + if i < items.len() - 1 { + f.write_str(" ")?; + } + } + } + DatasetIdent::All => f.write_str("All")?, + }; + + Ok(()) + } +} + +impl DatasetIdent { + pub fn many(transforms: Vec) -> Self { + Self::Transformed(transforms) + } + + pub fn prepare(self, dataset: impl Dataset) -> impl Dataset { + let items = match self { + DatasetIdent::Plain => Vec::new(), + DatasetIdent::All => { + vec![ + Transform::Translate, + Transform::Shear, + Transform::Scale, + Transform::Rotation, + ] + } + DatasetIdent::Transformed(items) => items.clone(), + }; + MapperDataset::new(dataset, MnistMapper::default().transform(&items)) + } + + pub fn compose( + idents: Vec<(Self, Option)>, + dataset: PartialDataset, MnistItem>, + ) -> impl Dataset { + let datasets = idents + .into_iter() + .map(|(ident, size)| match size { + Some(size) => { + SamplerDataset::with_replacement(ident.prepare(dataset.clone()), size) + } + None => { + let dataset = ident.prepare(dataset.clone()); + let size = dataset.len(); + SamplerDataset::without_replacement(dataset, size) + } + }) + .collect(); + ComposedDataset::new(datasets) + } +} + +fn generate_idents(num_samples_base: Option) -> Vec<(DatasetIdent, Option)> { + let mut current = Vec::new(); + let mut idents = Vec::new(); + + for shear in [None, Some(Transform::Shear)] { + for scale in [None, Some(Transform::Scale)] { + for rotation in [None, Some(Transform::Rotation)] { + for translate in [None, Some(Transform::Translate)] { + if let Some(tr) = shear { + current.push(tr); + } + if let Some(tr) = scale { + current.push(tr); + } + if let Some(tr) = rotation { + current.push(tr); + } + if let Some(tr) = translate { + current.push(tr); + } + + let num_samples = num_samples_base.map(|val| val * current.len()); + + if current.len() == 4 { + idents.push((DatasetIdent::All, num_samples)); + } else if current.is_empty() { + idents.push((DatasetIdent::Plain, num_samples)); + } else { + idents.push((DatasetIdent::many(current.clone()), num_samples)); + } + + current.clear(); + } + } + } + } + + idents +} diff --git a/examples/modern-lstm/Cargo.toml b/examples/modern-lstm/Cargo.toml new file mode 100644 index 0000000..418664b --- /dev/null +++ b/examples/modern-lstm/Cargo.toml @@ -0,0 +1,28 @@ +[package] +edition.workspace = true +name = "modern-lstm" +version = "0.22.0-pre.1" + +[lints] +workspace = true + +[features] +cuda = ["burn/cuda"] +flex = ["burn/flex"] +tch-cpu = ["burn/tch"] +tch-gpu = ["burn/tch"] +wgpu = ["burn/wgpu"] + +[dependencies] +burn = { workspace = true, features = ["default", "train"] } + +# Random number generator +rand = { workspace = true, features = ["thread_rng"] } +rand_distr = { workspace = true } + +# Serialization +serde = { workspace = true, features = ["std", "derive"] } + +# Organise the results in dataframe +planus = { workspace = true } +polars = { workspace = true } diff --git a/examples/modern-lstm/README.md b/examples/modern-lstm/README.md new file mode 100644 index 0000000..8fedb4c --- /dev/null +++ b/examples/modern-lstm/README.md @@ -0,0 +1,44 @@ +# Advanced LSTM Implementation with Burn + +A more advanced implementation of Long Short-Term Memory (LSTM) networks in Burn with combined +weight matrices for the input and hidden states, based on the +[PyTorch implementation](https://github.com/shiv08/Advanced-LSTM-Implementation-with-PyTorch). + +`LstmNetwork` is the top-level module with bidirectional and regularization support. The LSTM +variants differ by `bidirectional` and `num_layers` settings： + +- LSTM: `num_layers = 1` and `bidirectional = false` +- Stacked LSTM: `num_layers > 1` and `bidirectional = false` +- Bidirectional LSTM: `num_layers = 1` and `bidirectional = true` +- Bidirectional Stacked LSTM: `num_layers > 1` and `bidirectional = true` + +This implementation is complementary to Burn's official LSTM, users can choose either one depends on +the project's specific needs. + +## Usage + +## Training + +```sh +# Cuda backend +cargo run --example lstm-train --release --features cuda + +# Wgpu backend +cargo run --example lstm-train --release --features wgpu + +# Tch GPU backend +export TORCH_CUDA_VERSION=cu128 # Set the cuda version +cargo run --example lstm-train --release --features tch-gpu + +# Tch CPU backend +cargo run --example lstm-train --release --features tch-cpu + +# Flex backend (CPU) +cargo run --example lstm-train --release --features flex +``` + +### Inference + +```sh +cargo run --example lstm-infer --release --features cuda +``` diff --git a/examples/modern-lstm/examples/lstm-infer.rs b/examples/modern-lstm/examples/lstm-infer.rs new file mode 100644 index 0000000..ebc875d --- /dev/null +++ b/examples/modern-lstm/examples/lstm-infer.rs @@ -0,0 +1,68 @@ +use burn::tensor::Device; + +pub fn launch(device: Device) { + modern_lstm::inference::infer("/tmp/modern-lstm", device); +} + +#[cfg(feature = "flex")] +mod flex { + use burn::tensor::Device; + + pub fn run() { + crate::launch(Device::flex()); + } +} + +#[cfg(feature = "tch-gpu")] +mod tch_gpu { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + #[cfg(not(target_os = "macos"))] + let device = Device::libtorch_cuda(DeviceIndex::Default); + #[cfg(target_os = "macos")] + let device = Device::libtorch_mps(); + + crate::launch(device); + } +} + +#[cfg(feature = "tch-cpu")] +mod tch_cpu { + use burn::tensor::Device; + + pub fn run() { + crate::launch(Device::libtorch()); + } +} + +#[cfg(feature = "wgpu")] +mod wgpu { + use burn::tensor::{Device, DeviceKind}; + + pub fn run() { + crate::launch(Device::wgpu(DeviceKind::DefaultDevice)); + } +} + +#[cfg(feature = "cuda")] +mod cuda { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + crate::launch(Device::cuda(DeviceIndex::Default)); + } +} + +fn main() { + #[cfg(feature = "flex")] + flex::run(); + #[cfg(feature = "tch-gpu")] + tch_gpu::run(); + #[cfg(feature = "tch-cpu")] + tch_cpu::run(); + #[cfg(feature = "wgpu")] + wgpu::run(); + #[cfg(feature = "cuda")] + cuda::run(); +} diff --git a/examples/modern-lstm/examples/lstm-train.rs b/examples/modern-lstm/examples/lstm-train.rs new file mode 100644 index 0000000..507c283 --- /dev/null +++ b/examples/modern-lstm/examples/lstm-train.rs @@ -0,0 +1,75 @@ +use burn::{grad_clipping::GradientClippingConfig, optim::AdamConfig, tensor::Device}; +use modern_lstm::{model::LstmNetworkConfig, training::TrainingConfig}; + +pub fn launch(device: Device) { + let config = TrainingConfig::new( + LstmNetworkConfig::new(), + // Gradient clipping via optimizer config + AdamConfig::new().with_grad_clipping(Some(GradientClippingConfig::Norm(1.0))), + ); + + modern_lstm::training::train("/tmp/modern-lstm", config, device); +} + +#[cfg(feature = "flex")] +mod flex { + use burn::tensor::Device; + + pub fn run() { + crate::launch(Device::flex()); + } +} + +#[cfg(feature = "tch-gpu")] +mod tch_gpu { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + #[cfg(not(target_os = "macos"))] + let device = Device::libtorch_cuda(DeviceIndex::Default); + #[cfg(target_os = "macos")] + let device = Device::libtorch_mps(); + + crate::launch(device); + } +} + +#[cfg(feature = "tch-cpu")] +mod tch_cpu { + use burn::tensor::Device; + + pub fn run() { + crate::launch(Device::libtorch()); + } +} + +#[cfg(feature = "wgpu")] +mod wgpu { + use burn::tensor::{Device, DeviceKind}; + + pub fn run() { + crate::launch(Device::wgpu(DeviceKind::DefaultDevice)); + } +} + +#[cfg(feature = "cuda")] +mod cuda { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + crate::launch(Device::cuda(DeviceIndex::Default)); + } +} + +fn main() { + #[cfg(feature = "flex")] + flex::run(); + #[cfg(feature = "tch-gpu")] + tch_gpu::run(); + #[cfg(feature = "tch-cpu")] + tch_cpu::run(); + #[cfg(feature = "wgpu")] + wgpu::run(); + #[cfg(feature = "cuda")] + cuda::run(); +} diff --git a/examples/modern-lstm/src/dataset.rs b/examples/modern-lstm/src/dataset.rs new file mode 100644 index 0000000..c2e9ef5 --- /dev/null +++ b/examples/modern-lstm/src/dataset.rs @@ -0,0 +1,102 @@ +use burn::{ + data::{ + dataloader::batcher::Batcher, + dataset::{Dataset, InMemDataset}, + }, + prelude::*, +}; +use rand::RngExt; +use rand_distr::{Distribution, Normal}; +use serde::{Deserialize, Serialize}; + +// Dataset parameters +pub const NUM_SEQUENCES: usize = 1000; +pub const SEQ_LENGTH: usize = 10; +pub const NOISE_LEVEL: f32 = 0.1; +pub const RANDOM_SEED: u64 = 5; + +// Generate a sequence where each number is the sum of previous two numbers plus noise +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct SequenceDatasetItem { + pub sequence: Vec, + pub target: f32, +} + +impl SequenceDatasetItem { + pub fn new(seq_length: usize, noise_level: f32) -> Self { + // Start with two random numbers between 0 and 1 + let mut seq = vec![rand::rng().random(), rand::rng().random()]; + + // Generate sequence + for _i in 0..seq_length { + // Next number is sum of previous two plus noise + let normal = Normal::new(0.0, noise_level).unwrap(); + let next_val = + seq[seq.len() - 2] + seq[seq.len() - 1] + normal.sample(&mut rand::rng()); + seq.push(next_val); + } + + Self { + // Convert to sequence and target + sequence: seq[0..seq.len() - 1].to_vec(), // All but last + target: seq[seq.len() - 1], // Last value + } + } +} + +// Custom Dataset for Sequence Data +pub struct SequenceDataset { + dataset: InMemDataset, +} + +impl SequenceDataset { + pub fn new(num_sequences: usize, seq_length: usize, noise_level: f32) -> Self { + let mut items = vec![]; + for _i in 0..num_sequences { + items.push(SequenceDatasetItem::new(seq_length, noise_level)); + } + let dataset = InMemDataset::new(items); + + Self { dataset } + } +} + +impl Dataset for SequenceDataset { + fn get(&self, index: usize) -> Option { + self.dataset.get(index) + } + + fn len(&self) -> usize { + self.dataset.len() + } +} + +#[derive(Clone, Debug, Default)] +pub struct SequenceBatcher {} + +#[derive(Clone, Debug)] +pub struct SequenceBatch { + pub sequences: Tensor<3>, // [batch_size, seq_length, input_size] + pub targets: Tensor<2>, // [batch_size, 1] +} + +impl Batcher for SequenceBatcher { + fn batch(&self, items: Vec, device: &Device) -> SequenceBatch { + let mut sequences: Vec> = Vec::new(); + + for item in items.iter() { + let seq_tensor = Tensor::<1>::from_floats(item.sequence.as_slice(), device); + // Add feature dimension, the input_size is 1 implicitly. We can change the input_size here with some operations + sequences.push(seq_tensor.unsqueeze_dims(&[-1])); + } + let sequences = Tensor::stack(sequences, 0); + + let targets = items + .iter() + .map(|item| Tensor::<1>::from_floats([item.target], device)) + .collect(); + let targets = Tensor::stack(targets, 0); + + SequenceBatch { sequences, targets } + } +} diff --git a/examples/modern-lstm/src/inference.rs b/examples/modern-lstm/src/inference.rs new file mode 100644 index 0000000..0afbe0a --- /dev/null +++ b/examples/modern-lstm/src/inference.rs @@ -0,0 +1,44 @@ +use crate::{ + dataset::{ + NOISE_LEVEL, NUM_SEQUENCES, SEQ_LENGTH, SequenceBatcher, SequenceDataset, + SequenceDatasetItem, + }, + model::LstmNetwork, + training::TrainingConfig, +}; +use burn::{ + data::{dataloader::batcher::Batcher, dataset::Dataset}, + prelude::*, + store::ModuleRecord, +}; +use polars::prelude::*; + +pub fn infer(artifact_dir: &str, device: Device) { + // Loading model + let config = TrainingConfig::load(format!("{artifact_dir}/config.json")) + .expect("Config should exist for the model; run train first"); + let record = ModuleRecord::load(format!("{artifact_dir}/model")) + .expect("Trained model should exist; run train first"); + + let model: LstmNetwork = config.model.init(&device).load_record(record); + + let dataset = SequenceDataset::new(NUM_SEQUENCES / 5, SEQ_LENGTH, NOISE_LEVEL); + let items: Vec = dataset.iter().collect(); + + let batcher = SequenceBatcher::default(); + // Put all items in one batch + let batch = batcher.batch(items, &device); + let predicted = model.forward(batch.sequences, None); + let targets = batch.targets; + + let predicted = predicted.squeeze_dim::<1>(1).into_data(); + let expected = targets.squeeze_dim::<1>(1).into_data(); + + // Display the predicted vs expected values + let results = df![ + "predicted" => &predicted.to_vec::().unwrap(), + "expected" => &expected.to_vec::().unwrap(), + ] + .unwrap(); + println!("{}", results.head(Some(10))); +} diff --git a/examples/modern-lstm/src/lib.rs b/examples/modern-lstm/src/lib.rs new file mode 100644 index 0000000..1a167ff --- /dev/null +++ b/examples/modern-lstm/src/lib.rs @@ -0,0 +1,4 @@ +pub mod dataset; +pub mod inference; +pub mod model; +pub mod training; diff --git a/examples/modern-lstm/src/model.rs b/examples/modern-lstm/src/model.rs new file mode 100644 index 0000000..d531121 --- /dev/null +++ b/examples/modern-lstm/src/model.rs @@ -0,0 +1,359 @@ +use burn::{ + nn::{ + Dropout, DropoutConfig, Initializer, LayerNorm, LayerNormConfig, Linear, LinearConfig, + LstmState, Sigmoid, Tanh, + }, + prelude::*, +}; + +/// LSTM Cell implementation with layer normalization. +/// +/// Mathematical formulation of LSTM: +/// f_t = σ(W_f · [h_{t-1}, x_t] + b_f) # Forget gate +/// i_t = σ(W_i · [h_{t-1}, x_t] + b_i] # Input gate +/// g_t = tanh(W_g · [h_{t-1}, x_t] + b_g] # Candidate cell state +/// o_t = σ(W_o · [h_{t-1}, x_t] + b_o) # Output gate +/// +/// c_t = f_t ⊙ c_{t-1} + i_t ⊙ g_t # New cell state +/// h_t = o_t ⊙ tanh(c_t) # New hidden state +/// +/// where: +/// - σ is the sigmoid function +/// - ⊙ is the element-wise multiplication +/// - [h_{t-1}, x_t] represents concatenation + +#[derive(Module, Debug)] +pub struct LstmCell { + pub hidden_size: usize, + // Combined weight matrices for efficiency + // weight_ih layer uses combined weights for [i_t, f_t, g_t, o_t] for input x_t + // weight_hh layer uses combined weights for [i_t, f_t, g_t, o_t] for hidden state h_{t-1} + pub weight_ih: Linear, + pub weight_hh: Linear, + // Layer Normalization for better training stability. Don't use BatchNorm because the input distribution is always changing for LSTM. + pub norm_x: LayerNorm, // Normalize gate pre-activations + pub norm_h: LayerNorm, // Normalize hidden state + pub norm_c: LayerNorm, // Normalize cell state + pub dropout: Dropout, +} + +/// Configuration to create a Lstm module using the init function. +#[derive(Config, Debug)] +pub struct LstmCellConfig { + // The size of the input features + pub input_size: usize, + // The size of the hidden state + pub hidden_size: usize, + // The number of hidden layers + pub dropout: f64, +} + +impl LstmCellConfig { + // Initialize parameters using best practices: + // 1. Orthogonal initialization for better gradient flow (here we use Xavier because of the lack of Orthogonal in burn) + // 2. Initialize forget gate bias to 1.0 to prevent forgetting at start of training + #[allow(clippy::single_range_in_vec_init)] + pub fn init(&self, device: &Device) -> LstmCell { + let initializer = Initializer::XavierNormal { gain: 1.0 }; + let init_bias = Tensor::<1>::ones([self.hidden_size], device); + + let mut weight_ih = LinearConfig::new(self.input_size, 4 * self.hidden_size) + .with_initializer(initializer.clone()) + .init(device); + // Set forget gate bias to 1.0 (helps with learning long sequences) + let bias = weight_ih + .bias + .clone() + .unwrap() + .val() + .slice_assign([self.hidden_size..2 * self.hidden_size], init_bias.clone()); + weight_ih.bias = weight_ih.bias.map(|p| p.map(|_t| bias)); + + let mut weight_hh = LinearConfig::new(self.hidden_size, 4 * self.hidden_size) + .with_initializer(initializer) + .init(device); + let bias = weight_hh + .bias + .clone() + .unwrap() + .val() + .slice_assign([self.hidden_size..2 * self.hidden_size], init_bias); + weight_hh.bias = weight_hh.bias.map(|p| p.map(|_t| bias)); + + LstmCell { + hidden_size: self.hidden_size, + weight_ih, + weight_hh, + norm_x: LayerNormConfig::new(4 * self.hidden_size).init(device), + norm_h: LayerNormConfig::new(self.hidden_size).init(device), + norm_c: LayerNormConfig::new(self.hidden_size).init(device), + dropout: DropoutConfig::new(self.dropout).init(), + } + } +} + +impl LstmCell { + /// Forward pass of LSTM cell. + /// Args: + /// x: Input tensor of shape (batch_size, input_size) + /// state: Tuple of (h_{t-1}, c_{t-1}) each of shape (batch_size, hidden_size) + /// Returns: + /// Tuple of (h_t, c_t) representing new hidden and cell states + pub fn forward(&self, x: Tensor<2>, state: LstmState<2>) -> LstmState<2> { + let (h_prev, c_prev) = (state.hidden, state.cell); + + // Combined matrix multiplication for all gates + // Shape: (batch_size, 4 * hidden_size) + let gates_x = self.weight_ih.forward(x); // Transform input + let gates_h = self.weight_hh.forward(h_prev); // Transform previous hidden state + + // Apply layer normalization + let gates_x = self.norm_x.forward(gates_x); + // Combined gate pre-activations + let gates = gates_x + gates_h; + + // Split into individual gates + // Each gate shape: (batch_size, hidden_size) + let gates = gates.chunk(4, 1); + let i_gate = gates[0].clone(); + let f_gate = gates[1].clone(); + let g_gate = gates[2].clone(); + let o_gate = gates[3].clone(); + + // Apply gate non-linearities + let i_t = Sigmoid::new().forward(i_gate); + let f_t = Sigmoid::new().forward(f_gate); + let g_t = Tanh::new().forward(g_gate); + let o_t = Sigmoid::new().forward(o_gate); + + // Update cell state: c_t = f_t ⊙ c_{t-1} + i_t ⊙ g_t + let c_t = f_t * c_prev + i_t * g_t; + let c_t = self.norm_c.forward(c_t); + + // Update cell state: h_t = o_t ⊙ tanh(c_t) + let h_t = o_t * Tanh::new().forward(c_t.clone()); + let h_t = self.norm_h.forward(h_t); + + let h_t = self.dropout.forward(h_t); + + LstmState::new(h_t, c_t) + } + + // Initialize cell state and hidden state if provided or with zeros + pub fn init_state(&self, batch_size: usize, device: &Device) -> LstmState<2> { + let cell = Tensor::zeros([batch_size, self.hidden_size], device); + let hidden = Tensor::zeros([batch_size, self.hidden_size], device); + + LstmState::new(cell, hidden) + } +} + +/// Stacked LSTM implementation supporting multiple layers +/// Each layer processes the output of the previous layer +#[derive(Module, Debug)] +pub struct StackedLstm { + pub layers: Vec, +} + +#[derive(Config, Debug)] +pub struct StackedLstmConfig { + pub input_size: usize, + pub hidden_size: usize, + pub num_layers: usize, + pub dropout: f64, +} + +impl StackedLstmConfig { + pub fn init(&self, device: &Device) -> StackedLstm { + let mut layers: Vec = vec![]; + // Create list of LSTM cells, one for each layer + for i in 0..self.num_layers { + if i == 0 { + if i < self.num_layers - 1 { + layers.push( + LstmCellConfig::new(self.input_size, self.hidden_size, self.dropout) + .init(device), + ); + } else { + // No dropout on last layer + layers.push( + LstmCellConfig::new(self.input_size, self.hidden_size, 0.0).init(device), + ); + } + } else if i < self.num_layers - 1 { + layers.push( + LstmCellConfig::new(self.hidden_size, self.hidden_size, self.dropout) + .init(device), + ); + } else { + // No dropout on last layer + layers.push( + LstmCellConfig::new(self.hidden_size, self.hidden_size, 0.0).init(device), + ); + } + } + StackedLstm { layers } + } +} + +impl StackedLstm { + /// Process input sequence through stacked LSTM layers. + /// + /// Args: + /// x: Input tensor of shape (batch_size, seq_length, input_size) + /// states: Optional initial states for each layer + /// + /// Returns: + /// Tuple of (output, states) where output has shape (batch_size, seq_length, hidden_size) + /// and states is a vector of length num_layers, both cell and hidden state in each element have shape (batch_size, hidden_size) + pub fn forward( + &self, + x: Tensor<3>, + states: Option>>, + ) -> (Tensor<3>, Vec>) { + let [batch_size, seq_length, _] = x.dims(); + let device = x.device(); + + let mut states = match states { + None => { + let mut temp: Vec> = vec![]; + for layer in self.layers.iter() { + temp.push(layer.init_state(batch_size, &device)); + } + temp + } + _ => states.unwrap(), + }; + + let mut layer_outputs = vec![]; + for t in 0..seq_length { + let mut input_t = x.clone().slice(s![.., t..t + 1, ..]).squeeze_dim::<2>(1); + for (i, lstm_cell) in self.layers.iter().enumerate() { + let mut state: LstmState<2> = + LstmState::new(states[i].cell.clone(), states[i].hidden.clone()); + state = lstm_cell.forward(input_t, state); + input_t = state.hidden.clone(); + states[i] = state; + } + layer_outputs.push(input_t); + } + + // Stack output along sequence dimension + let output = Tensor::stack(layer_outputs, 1); + + (output, states) + } +} + +/// Complete LSTM network with bidirectional support. +/// +/// In bidirectional mode: +/// - Forward LSTM processes sequence from left to right +/// - Backward LSTM processes sequence from right to left +/// - Outputs are concatenated for final prediction +#[derive(Module, Debug)] +pub struct LstmNetwork { + // Forward direction LSTM + pub stacked_lstm: StackedLstm, + // Optional backward direction LSTM for bidirectional processing + pub reverse_lstm: Option, + pub dropout: Dropout, + pub fc: Linear, +} + +#[derive(Config, Debug)] +pub struct LstmNetworkConfig { + #[config(default = 1)] + pub input_size: usize, // Single feature (number sequence) + #[config(default = 32)] + pub hidden_size: usize, // Size of LSTM hidden state + #[config(default = 2)] + pub num_layers: usize, // Number of LSTM layers + #[config(default = 1)] + pub output_size: usize, // Predict one number + #[config(default = 0.1)] + pub dropout: f64, + #[config(default = true)] + pub bidirectional: bool, // Use bidirectional LSTM +} + +impl LstmNetworkConfig { + pub fn init(&self, device: &Device) -> LstmNetwork { + // Forward direction LSTM + let stacked_lstm = StackedLstmConfig::new( + self.input_size, + self.hidden_size, + self.num_layers, + self.dropout, + ) + .init(device); + + // Optional backward direction LSTM for bidirectional processing + let (reverse_lstm, hidden_size) = if self.bidirectional { + let lstm = StackedLstmConfig::new( + self.input_size, + self.hidden_size, + self.num_layers, + self.dropout, + ) + .init(device); + (Some(lstm), 2 * self.hidden_size) + } else { + (None, self.hidden_size) + }; + + let fc = LinearConfig::new(hidden_size, self.output_size).init(device); + let dropout = DropoutConfig::new(self.dropout).init(); + + LstmNetwork { + stacked_lstm, + reverse_lstm, + dropout, + fc, + } + } +} + +impl LstmNetwork { + /// Forward pass of the network. + /// + /// For bidirectional processing: + /// 1. Process sequence normally with forward LSTM + /// 2. Process reversed sequence with backward LSTM + /// 3. Concatenate both outputs + /// 4. Apply final linear transformation + /// + /// Args: + /// x: Input tensor of shape (batch_size, seq_length, input_size) + /// states: Optional initial states + /// + /// Returns: + /// Output tensor of shape (batch_size, output_size) + pub fn forward(&self, x: Tensor<3>, states: Option>>) -> Tensor<2> { + let seq_length = x.dims()[1]; + // Forward direction + let (mut output, _states) = self.stacked_lstm.forward(x.clone(), states); + + output = match &self.reverse_lstm { + Some(reverse_lstm) => { + //Process sequence in reverse direction + let (mut reverse_output, _states) = reverse_lstm.forward(x.flip([1]), None); + // Flip back to align with forward sequence + reverse_output = reverse_output.flip([1]); + // Concatenate forward and backward outputs along the feature dimension + output = Tensor::cat(vec![output, reverse_output], 2); + output + } + None => output, + }; + + // Apply dropout before final layer + output = self.dropout.forward(output); + // Use final timestep output for prediction + self.fc.forward( + output + .slice(s![.., seq_length - 1..seq_length, ..]) + .squeeze_dim::<2>(1), + ) + } +} diff --git a/examples/modern-lstm/src/training.rs b/examples/modern-lstm/src/training.rs new file mode 100644 index 0000000..b1fde09 --- /dev/null +++ b/examples/modern-lstm/src/training.rs @@ -0,0 +1,133 @@ +use std::path::PathBuf; + +use crate::dataset::{ + NOISE_LEVEL, NUM_SEQUENCES, RANDOM_SEED, SEQ_LENGTH, SequenceBatcher, SequenceDataset, +}; +use crate::model::LstmNetworkConfig; +use burn::{ + data::dataloader::DataLoaderBuilder, + module::AutodiffModule, + nn::loss::{MseLoss, Reduction::Mean}, + optim::{AdamConfig, GradientsParams}, + prelude::*, +}; + +#[derive(Config, Debug)] +pub struct TrainingConfig { + pub model: LstmNetworkConfig, + pub optimizer: AdamConfig, + + #[config(default = 30)] + pub num_epochs: usize, + #[config(default = 32)] + pub batch_size: usize, + #[config(default = 2)] + pub num_workers: usize, + #[config(default = 1e-3)] + pub lr: f64, +} + +// Create the directory to save the model and model config +fn create_artifact_dir(artifact_dir: &str) { + std::fs::remove_file(PathBuf::from(artifact_dir).join("experiment.log")).ok(); + std::fs::create_dir_all(artifact_dir).ok(); +} + +pub fn train(artifact_dir: &str, config: TrainingConfig, device: Device) { + create_artifact_dir(artifact_dir); + + // Save training config + config + .save(format!("{artifact_dir}/config.json")) + .expect("Config should be saved successfully"); + + device.seed(RANDOM_SEED); + let autodiff_device = device.clone().autodiff(); + + // Create the model and optimizer + let mut model = config.model.init(&autodiff_device); + let mut optim = config.optimizer.init(); + + // Create the batcher + let batcher = SequenceBatcher::default(); + + // Create the dataloaders + let dataloader_train = DataLoaderBuilder::new(batcher.clone()) + .batch_size(config.batch_size) + .shuffle(RANDOM_SEED) + .num_workers(config.num_workers) + .set_device(autodiff_device) + .build(SequenceDataset::new(NUM_SEQUENCES, SEQ_LENGTH, NOISE_LEVEL)); + + let dataloader_valid = DataLoaderBuilder::new(batcher) + .batch_size(config.batch_size) + .shuffle(RANDOM_SEED) + .num_workers(config.num_workers) + // 20% size of training + .build(SequenceDataset::new( + NUM_SEQUENCES / 5, + SEQ_LENGTH, + NOISE_LEVEL, + )); + + let train_num_items = dataloader_train.num_items(); + let valid_num_items = dataloader_valid.num_items(); + + println!("Starting training..."); + // Iterate over our training for X epochs + for epoch in 1..config.num_epochs + 1 { + // Initialize the training and validation metrics at the start of each epoch + let mut train_losses = vec![]; + let mut train_loss = 0.0; + let mut valid_losses = vec![]; + let mut valid_loss = 0.0; + + // Implement our training loop + for batch in dataloader_train.iter() { + let output = model.forward(batch.sequences, None); + let loss = MseLoss::new().forward(output, batch.targets.clone(), Mean); + train_loss += loss.clone().into_scalar::() * batch.targets.dims()[0] as f32; + + // Gradients for the current backward pass + let grads = loss.backward(); + // Gradients linked to each parameter of the model + let grads = GradientsParams::from_grads(grads, &model); + // Update the model using the optimizer + model = optim.step(config.lr.into(), model, grads); + } + + // The averaged train loss per epoch + let avg_train_loss = train_loss / train_num_items as f32; + train_losses.push(avg_train_loss); + + // Get the model without autodiff + let valid_model = model.valid(); + + // Implement our validation loop + for batch in dataloader_valid.iter() { + let output = valid_model.forward(batch.sequences, None); + let loss = MseLoss::new().forward(output, batch.targets.clone(), Mean); + valid_loss += loss.clone().into_scalar::() * batch.targets.dims()[0] as f32; + } + // The averaged train loss per epoch + let avg_valid_loss = valid_loss / valid_num_items as f32; + valid_losses.push(avg_valid_loss); + + // Display the averaged training and validation metrics every 10 epochs + if (epoch + 1) % 5 == 0 { + println!( + "Epoch {}/{}, Avg Loss {:.4}, Avg Val Loss: {:.4}", + epoch + 1, + config.num_epochs, + avg_train_loss, + avg_valid_loss, + ); + } + } + + // Save the trained model + model + .into_record() + .save(format!("{artifact_dir}/model")) + .expect("Trained model should be saved successfully"); +} diff --git a/examples/notebook/README.md b/examples/notebook/README.md new file mode 100644 index 0000000..c141aa7 --- /dev/null +++ b/examples/notebook/README.md @@ -0,0 +1,47 @@ +# Jupyter Notebook Examples with Burn + +This directory includes Jupyter Notebook examples showcasing the usage of the Burn deep learning +framework in Rust through +[Evcxr Jupyter](https://github.com/evcxr/evcxr/blob/main/evcxr_jupyter/README.md). The examples are +systematically organized based on the specific Burn features they illustrate. + +## Viewing Options + +You can explore the examples in different ways: + +- **Notebook Viewer:** If you prefer not to set up the entire crate package, you can view the + examples in a notebook viewer or run them to see images and other media outputs. + +- **Visual Studio Code (vscode):** If you're using vscode, you already have access to a built-in + notebook viewer, enabling you to open and interact with the notebook files directly. + +For other editors, you can utilize the [Jupyter Notebook Viewer](https://nbviewer.jupyter.org/). + +## Getting Started with Rust and Evcxr + +To execute the Rust code within the notebooks, you must install the Evcxr kernel. Here's how to get +started: + +### Install Evcxr Kernel + +1. **Build Evcxr Kernel:** Install the required package with the following command: + + ```shell + cargo install evcxr_jupyter + ``` + +2. **Install and Register the Kernel to Jupyter:** + ```shell + evcxr_jupyter --install + ``` + +### Open and Run Notebooks + +Once the kernel is installed, you can open the notebook files in your preferred editor and run the +code. Ensure that the kernel is set to `Rust` within the notebook for proper execution. + +## Additional Reading Resources + +- [Notebook Special Commands for Evcxr](https://github.com/evcxr/evcxr/blob/main/COMMON.md): Learn + about the unique commands and functionalities offered by Evcxr for a more efficient workflow with + Jupyter Notebooks. diff --git a/examples/notebook/autodiff.ipynb b/examples/notebook/autodiff.ipynb new file mode 100644 index 0000000..1189410 --- /dev/null +++ b/examples/notebook/autodiff.ipynb @@ -0,0 +1,600 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Autodifferentiation and Gradient Descent in Burn\n", + "\n", + "This notebook demonstrates how to use automatic differentiation in Burn to compute gradients and implement gradient descent." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, + "outputs": [], + "source": [ + "// Dependency declarations\n", + ":dep burn = {path = \"../../crates/burn\"}\n", + ":dep burn-flex = {path = \"../../crates/burn-flex\"}\n", + ":dep burn-autodiff = {path = \"../../crates/burn-autodiff\"}\n" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, + "outputs": [], + "source": [ + "// Import packages\n", + "use burn::prelude::*;\n", + "use burn_autodiff::Autodiff;\n", + "use burn_flex::Flex;\n", + "\n", + "// Type alias: Autodiff enables automatic differentiation\n", + "type B = Autodiff;\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Understanding require_grad()\n", + "\n", + "In Burn, tensors can be marked for gradient tracking using `.require_grad()`. This tells the framework to track operations on this tensor so gradients can be computed later." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Regular tensor x: Tensor {\n", + " data:\n", + "[1.0, 2.0, 3.0, 4.0],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"autodiff\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Tensor y with require_grad: Tensor {\n", + " data:\n", + "[1.0, 2.0, 3.0, 4.0],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"autodiff\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "result = sum(y * 2) = Tensor {\n", + " data:\n", + "[20.0],\n", + " shape: [1],\n", + " device: Cpu,\n", + " backend: \"autodiff\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "let device = Default::default();\n", + "\n", + "// Create a regular tensor - no gradient tracking\n", + "let x: Tensor = Tensor::from_floats([1.0, 2.0, 3.0, 4.0], &device);\n", + "println!(\"Regular tensor x: {}\", x);\n", + "\n", + "// Create a tensor that requires gradient computation\n", + "let y: Tensor = Tensor::from_floats([1.0, 2.0, 3.0, 4.0], &device).require_grad();\n", + "println!(\"Tensor y with require_grad: {}\", y);\n", + "\n", + "// Now let's do some operations on y\n", + "let z = y.clone() * 2.0;\n", + "let result = z.sum();\n", + "println!(\"result = sum(y * 2) = {}\", result);\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Computing Gradients with backward()\n", + "\n", + "The `.backward()` method computes the gradients of all tensors that have `require_grad()` set. It returns a gradients object that holds the computed gradients." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "y = Tensor {\n", + " data:\n", + "[1.0, 2.0, 3.0, 4.0],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"autodiff\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "dy/dx = Tensor {\n", + " data:\n", + "[2.0, 2.0, 2.0, 2.0],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "// Example: y = [1, 2, 3, 4]\n", + "// z = y * 2 = [2, 4, 6, 8]\n", + "// result = sum(z) = 20\n", + "//\n", + "// d(result)/d(y) = d(result)/dz * dz/dy = 1 * 2 = [2, 2, 2, 2]\n", + "\n", + "let device = Default::default();\n", + "let y: Tensor = Tensor::from_floats([1.0, 2.0, 3.0, 4.0], &device).require_grad();\n", + "let z = y.clone() * 2.0;\n", + "let result = z.sum();\n", + "\n", + "// Compute gradients\n", + "let grads = result.backward();\n", + "\n", + "// Get gradient for y\n", + "let y_grad = y.grad(&grads).unwrap();\n", + "println!(\"y = {}\", y);\n", + "println!(\"d(result)/dy = {}\", y_grad);\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. More Complex Example: Quadratic Function\n", + "Let's compute the gradient of a more complex function: f(x) = x²\n", + "\n", + "The derivative is: f'(x) = 2x" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "x = Tensor {\n", + " data:\n", + "[1.0, 2.0, 3.0, 4.0],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"autodiff\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "x^2 = Tensor {\n", + " data:\n", + "[1.0, 4.0, 9.0, 16.0],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"autodiff\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "d(x^2)/dx = Tensor {\n", + " data:\n", + "[2.0, 4.0, 6.0, 8.0],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Expected: [2, 4, 6, 8]\n" + ] + } + ], + "source": [ + "// f(x) = x^2\n", + "// f'(x) = 2x\n", + "\n", + "let device = Default::default();\n", + "let x: Tensor = Tensor::from_floats([1.0, 2.0, 3.0, 4.0], &device).require_grad();\n", + "let y = x.clone().powf_scalar(2.0);\n", + "let result = y.clone().sum();\n", + "\n", + "let grads = result.backward();\n", + "let x_grad = x.grad(&grads).unwrap();\n", + "\n", + "println!(\"x = {}\", x);\n", + "println!(\"x^2 = {}\", y);\n", + "println!(\"d(x^2)/dx = {}\", x_grad);\n", + "\n", + "// Verify: d(x^2)/dx should be [2, 4, 6, 8]\n", + "println!(\"Expected: [2, 4, 6, 8]\");\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Chain Rule Example\n", + "\n", + "Let's verify the chain rule: f(g(x))' = f'(g(x)) * g'(x)\n", + "\n", + "Example: y = sin(x²), we want dy/dx\n", + "\n", + "Let u = x², y = sin(u)\n", + "dy/du = cos(u), du/dx = 2x\n", + "dy/dx = cos(x²) * 2x" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "x = Tensor {\n", + " data:\n", + "[0.0, 1.0, 2.0, 3.0],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"autodiff\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "y = sin(x^2) = Tensor {\n", + " data:\n", + "[0.0, 0.84147096, -0.7568025, 0.4121185],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"autodiff\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "dy/dx = Tensor {\n", + " data:\n", + "[0.0, 1.0806046, -2.6145744, -5.4667816],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Expected (cos(x^2) * 2x): Tensor {\n", + " data:\n", + "[0.0, 1.0806046, -2.6145744, -5.4667816],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"autodiff\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "// y = sin(x^2)\n", + "// dy/dx = cos(x^2) * 2x\n", + "\n", + "let device = Default::default();\n", + "let x: Tensor = Tensor::from_floats([0.0, 1.0, 2.0, 3.0], &device).require_grad();\n", + "\n", + "// Forward pass\n", + "let x_squared = x.clone().powf_scalar(2.0);\n", + "let y = x_squared.sin();\n", + "let result = y.clone().sum();\n", + "\n", + "// Backward pass\n", + "let grads = result.backward();\n", + "let x_grad = x.grad(&grads).unwrap();\n", + "\n", + "println!(\"x = {}\", x);\n", + "println!(\"y = sin(x^2) = {}\", y);\n", + "println!(\"dy/dx = {}\", x_grad);\n", + "\n", + "// Verify manually: cos(x^2) * 2x\n", + "let expected_grad = x.clone().powf_scalar(2.0).cos() * (x.clone() * 2.0);\n", + "println!(\"Expected (cos(x^2) * 2x): {}\", expected_grad);\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Gradient Descent from Scratch\n", + "\n", + "Now let's implement the classic gradient descent algorithm to find the minimum of a function.\n", + "\n", + "We'll minimize: f(x) = (x - 3)²\n", + "\n", + "The minimum is at x = 3, where f(x) = 0" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Starting gradient descent to minimize (x - 3)^2\n", + "Expected minimum: x = 3\n", + "---\n", + "Iteration 0: x = 0.0000, loss = 9.0000\n", + "Iteration 1: x = 0.6000, loss = 5.7600\n", + "Iteration 2: x = 1.0800, loss = 3.6864\n", + "Iteration 3: x = 1.4640, loss = 2.3593\n", + "Iteration 4: x = 1.7712, loss = 1.5099\n", + "Iteration 5: x = 2.0170, loss = 0.9664\n", + "Iteration 6: x = 2.2136, loss = 0.6185\n", + "Iteration 7: x = 2.3709, loss = 0.3958\n", + "Iteration 8: x = 2.4967, loss = 0.2533\n", + "Iteration 9: x = 2.5973, loss = 0.1621\n", + "Iteration 10: x = 2.6779, loss = 0.1038\n", + "Iteration 11: x = 2.7423, loss = 0.0664\n", + "Iteration 12: x = 2.7938, loss = 0.0425\n", + "Iteration 13: x = 2.8351, loss = 0.0272\n", + "Iteration 14: x = 2.8681, loss = 0.0174\n", + "Iteration 15: x = 2.8944, loss = 0.0111\n", + "Iteration 16: x = 2.9156, loss = 0.0071\n", + "Iteration 17: x = 2.9324, loss = 0.0046\n", + "Iteration 18: x = 2.9460, loss = 0.0029\n", + "Iteration 19: x = 2.9568, loss = 0.0019\n", + "---\n", + "Final x = 2.9654\n" + ] + } + ], + "source": [ + "// Target: minimize f(x) = (x - 3)^2\n", + "// This has minimum at x = 3\n", + "\n", + "fn loss(x: &Tensor) -> Tensor {\n", + " // f(x) = (x - 3)^2\n", + " (x.clone() - 3.0).powf_scalar(2.0)\n", + "}\n", + "\n", + "let device = Default::default();\n", + "// Start from x = 0\n", + "let mut x_val: f32 = 0.0;\n", + "\n", + "let learning_rate: f32 = 0.1;\n", + "\n", + "println!(\"Starting gradient descent to minimize (x - 3)^2\");\n", + "println!(\"Expected minimum: x = 3\");\n", + "println!(\"---\");\n", + "\n", + "for i in 0..20 {\n", + " // Create a new tensor with current x value and require gradients\n", + " let x = Tensor::::from_floats([x_val], &device).require_grad();\n", + " \n", + " // Forward pass\n", + " let loss_value = loss(&x);\n", + " \n", + " // Get loss as f32 for printing\n", + " let loss_scalar: f32 = loss_value.clone().into_scalar::();\n", + " \n", + " println!(\"Iteration {}: x = {:.4}, loss = {:.4}\", i, x_val, loss_scalar);\n", + "\n", + " // Backward pass\n", + " let grads = loss_value.backward();\n", + " let grad = x.grad(&grads).unwrap();\n", + " \n", + " // Update: x = x - learning_rate * gradient\n", + " let grad_val: f32 = grad.into_scalar::();\n", + " x_val = x_val - grad_val * learning_rate;\n", + "}\n", + "\n", + "println!(\"---\");\n", + "println!(\"Final x = {:.4}\", x_val);\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Linear Regression with Gradient Descent\n", + "\n", + "Let's use gradient descent to fit a simple linear regression model: y = wx + b\n", + "\n", + "We'll generate synthetic data where the true relationship is y = 2x + 1" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Generated 100 data points\n", + "True relationship: y = 2x + 1\n", + "First 5 x values: [0.0, 0.1, 0.2, 0.3, 0.4]\n", + "First 5 y values: [0.87993187, 0.98804677, 1.5366085, 1.7324162, 1.653858]\n" + ] + } + ], + "source": [ + "use burn::tensor::{Distribution, TensorData};\n", + "\n", + "let device = Default::default();\n", + "// Generate synthetic data: y = 2x + 1 + noise\n", + "let num_samples = 100;\n", + "let x_data = TensorData::new((0..num_samples).map(|i| i as f32 / 10.0).collect(), [num_samples, 1]);\n", + "// Generate noise using Burn's random tensor\n", + "let noise = Tensor::::random([num_samples, 1], Distribution::Uniform(-0.25, 0.25), &device);\n", + "\n", + "let x = Tensor::::from(x_data);\n", + "let y: Tensor = 2 * x.clone() + 1 + noise;\n", + "\n", + "println!(\"Generated {} data points\", num_samples);\n", + "println!(\"True relationship: y = 2x + 1\");\n", + "println!(\"First 5 x values: {}\", x.clone().slice([0..5, 0..1]));\n", + "println!(\"First 5 y values: {}\", y.clone().slice([0..5, 0..1]));\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "vscode": { + "languageId": "rust" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Training linear regression with gradient descent...\n", + "Initial w = 0.5000, b = 0.5000\n", + "Epoch 0: loss = 81.7705, w = 1.5358, b = 0.6586\n", + "Epoch 20: loss = 0.0365, w = 2.0384, b = 0.7594\n", + "Epoch 40: loss = 0.0341, w = 2.0351, b = 0.7810\n", + "Epoch 60: loss = 0.0321, w = 2.0322, b = 0.8006\n", + "Epoch 80: loss = 0.0305, w = 2.0295, b = 0.8184\n", + "---\n", + "Final: w = 2.0272, b = 0.8336\n", + "True: w = 2.0, b = 1.0\n" + ] + } + ], + "source": [ + "// Initialize weights randomly\n", + "let device = Default::default();\n", + "let mut w_val: f32 = 0.5; // Start with reasonable initial values\n", + "let mut b_val: f32 = 0.5;\n", + "\n", + "let learning_rate: f32 = 0.01;\n", + "let num_epochs = 100;\n", + "\n", + "println!(\"Training linear regression with gradient descent...\");\n", + "println!(\"Initial w = {:.4}, b = {:.4}\", w_val, b_val);\n", + "\n", + "for epoch in 0..num_epochs {\n", + " // Create tensors with current parameter values\n", + " let w = Tensor::::from_floats([[w_val]], &device).require_grad();\n", + " let b = Tensor::::from_floats([[b_val]], &device).require_grad();\n", + " \n", + " // Forward pass: y_pred = w * x + b\n", + " let y_pred = x.clone().matmul(w.clone()) + b.clone();\n", + " \n", + " // Compute loss: MSE = (1/n) * sum((y_pred - y)^2)\n", + " let loss = (y_pred.clone() - y.clone()).powf_scalar(2.0).mean();\n", + " \n", + " // Backward pass\n", + " let grads = loss.backward();\n", + " let w_grad = w.grad(&grads).unwrap();\n", + " let b_grad = b.grad(&grads).unwrap();\n", + " \n", + " // Update weights\n", + " let w_grad_val: f32 = w_grad.into_scalar::();\n", + " let b_grad_val: f32 = b_grad.into_scalar::();\n", + " w_val = w_val - w_grad_val * learning_rate;\n", + " b_val = b_val - b_grad_val * learning_rate;\n", + " \n", + " if epoch % 20 == 0 {\n", + " let loss_val: f32 = loss.clone().into_scalar::();\n", + " println!(\"Epoch {:3}: loss = {:.4}, w = {:.4}, b = {:.4}\", epoch, loss_val, w_val, b_val);\n", + " }\n", + "}\n", + "\n", + "println!(\"---\");\n", + "println!(\"Final: w = {:.4}, b = {:.4}\", w_val, b_val);\n", + "println!(\"True: w = 2.0, b = 1.0\");\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "In this notebook, we covered:\n", + "\n", + "- **require_grad()**: Mark tensors for gradient tracking\n", + "- **backward()**: Compute gradients automatically using reverse-mode autodiff\n", + "- **grad()**: Retrieve computed gradients\n", + "- **Gradient Descent**: Implemented from scratch to minimize a quadratic function\n", + "- **Linear Regression**: Used gradient descent to fit a linear model to data\n", + "\n", + "These concepts are the foundation of neural network training in Burn!" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Rust", + "language": "rust", + "name": "rust" + }, + "language_info": { + "codemirror_mode": "rust", + "file_extension": ".rs", + "mimetype": "text/rust", + "name": "Rust", + "pygment_lexer": "rust", + "version": "" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/notebook/basic-tensor-op.ipynb b/examples/notebook/basic-tensor-op.ipynb new file mode 100644 index 0000000..7d44d18 --- /dev/null +++ b/examples/notebook/basic-tensor-op.ipynb @@ -0,0 +1,1115 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Tensor Operations in Burn\n", + "\n", + "This notebook demonstrates basic tensor operations in Burn, a deep learning framework written in Rust." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "// Dependency declarations for the notebook.\n", + "// The syntax is similar to Cargo.toml. Just prefix with :dep\n", + "\n", + ":dep burn = {path = \"../../crates/burn\"}\n", + ":dep burn-flex = {path = \"../../crates/burn-flex\"}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "// Import packages\n", + "use burn::prelude::*;\n", + "use burn_flex::Flex;\n", + "\n", + "// Type alias for the backend (using CPU/Flex)\n", + "type B = Flex;" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Tensor Creation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Empty tensor shape: Shape { dims: [2, 3, 4] }\n", + "Zeros tensor: Tensor {\n", + " data:\n", + "[[0.0, 0.0, 0.0],\n", + " [0.0, 0.0, 0.0],\n", + " [0.0, 0.0, 0.0]],\n", + " shape: [3, 3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Ones tensor: Tensor {\n", + " data:\n", + "[[1.0, 1.0, 1.0, 1.0],\n", + " [1.0, 1.0, 1.0, 1.0]],\n", + " shape: [2, 4],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Full tensor (7.0): Tensor {\n", + " data:\n", + "[[7.0, 7.0, 7.0],\n", + " [7.0, 7.0, 7.0]],\n", + " shape: [2, 3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "let device = Default::default();\n", + "\n", + "// Create an empty tensor (uninitialized values)\n", + "let empty: Tensor = Tensor::empty([2, 3, 4], &device);\n", + "println!(\"Empty tensor shape: {:?}\", empty.shape());\n", + "\n", + "// Create a tensor filled with zeros\n", + "let zeros: Tensor = Tensor::zeros([3, 3], &device);\n", + "println!(\"Zeros tensor: {}\", zeros);\n", + "\n", + "// Create a tensor filled with ones\n", + "let ones: Tensor = Tensor::ones([2, 4], &device);\n", + "println!(\"Ones tensor: {}\", ones);\n", + "\n", + "// Create a tensor filled with a specific value\n", + "let full: Tensor = Tensor::full([2, 3], 7.0, &device);\n", + "println!(\"Full tensor (7.0): {}\", full);" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "From slice:\n", + "Tensor {\n", + " data:\n", + "[[1.0, 2.0, 3.0],\n", + " [4.0, 5.0, 6.0]],\n", + " shape: [2, 3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Random tensor: Tensor {\n", + " data:\n", + "[0.32371014, 0.41100568, 0.94457513, 0.8408601, 0.42262083],\n", + " shape: [5],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Normal distribution: Tensor {\n", + " data:\n", + "[-0.22402725, 1.8367178, -1.1049407, -0.6302627, 1.1106112],\n", + " shape: [5],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Uniform [0, 10): Tensor {\n", + " data:\n", + "[8.110331, 7.335061, 9.858947, 6.0834813, 3.6619747],\n", + " shape: [5],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "// Create a tensor from a slice of values\n", + "let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];\n", + "let from_slice = Tensor::::from_floats(data, &device).reshape([2, 3]);\n", + "println!(\"From slice:\\n{}\", from_slice);\n", + "\n", + "// Create a random tensor\n", + "use burn::tensor::Distribution;\n", + "let random: Tensor = Tensor::random([5], Distribution::Default, &device);\n", + "println!(\"Random tensor: {}\", random);\n", + "\n", + "// Create a tensor with normal distribution\n", + "let normal: Tensor = Tensor::random([5], Distribution::Normal(0.0, 1.0), &device);\n", + "println!(\"Normal distribution: {}\", normal);\n", + "\n", + "// Create a tensor with uniform distribution in range [0, 10)\n", + "let uniform: Tensor = Tensor::random([5], Distribution::Uniform(0.0, 10.0), &device);\n", + "println!(\"Uniform [0, 10): {}\", uniform);" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Shape Operations" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Original (2x3):\n", + "Tensor {\n", + " data:\n", + "[[1.0, 2.0, 3.0],\n", + " [4.0, 5.0, 6.0]],\n", + " shape: [2, 3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Reshaped (1x2x3): Tensor {\n", + " data:\n", + "[[[1.0, 2.0, 3.0],\n", + " [4.0, 5.0, 6.0]]],\n", + " shape: [1, 2, 3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Flattened: Tensor {\n", + " data:\n", + "[1.0, 2.0, 3.0, 4.0, 5.0, 6.0],\n", + " shape: [6],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "// Reshape tensor - change the dimensions without changing the data\n", + "let tensor = Tensor::::from_floats([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &device).reshape([2, 3]);\n", + "println!(\"Original (2x3):\\n{}\", tensor);\n", + "\n", + "let reshaped: Tensor = tensor.clone().reshape([1, 2, 3]);\n", + "println!(\"Reshaped (1x2x3): {}\", reshaped);\n", + "\n", + "// Flatten - reshape to 1D\n", + "let flat: Tensor = tensor.flatten(0, 1);\n", + "println!(\"Flattened: {}\", flat);" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Original:\n", + "Tensor {\n", + " data:\n", + "[[1.0, 2.0],\n", + " [3.0, 4.0]],\n", + " shape: [2, 2],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Transposed:\n", + "Tensor {\n", + " data:\n", + "[[1.0, 3.0],\n", + " [2.0, 4.0]],\n", + " shape: [2, 2],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Using .t():\n", + "Tensor {\n", + " data:\n", + "[[1.0, 3.0],\n", + " [2.0, 4.0]],\n", + " shape: [2, 2],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "// Transpose - swap dimensions\n", + "let tensor = Tensor::::from_floats([1.0, 2.0, 3.0, 4.0], &device).reshape([2, 2]);\n", + "println!(\"Original:\\n{}\", tensor);\n", + "\n", + "let transposed = tensor.clone().transpose();\n", + "println!(\"Transposed:\\n{}\", transposed);\n", + "\n", + "// Also .t() works for 2D tensors\n", + "let t = tensor.t();\n", + "println!(\"Using .t():\\n{}\", t);" + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Before squeeze [1,1,2]: shape = Shape { dims: [1, 1, 2] }\n", + "After squeeze: shape = Shape { dims: [2] }\n", + "Before unsqueeze [2,2]: shape = Shape { dims: [2, 2] }\n", + "After unsqueeze: shape = Shape { dims: [1, 2, 2] }\n" + ] + } + ], + "source": [ + "// Squeeze - remove dimensions of size 1\n", + "let tensor = Tensor::::from_floats([1.0, 2.0], &device).reshape([1, 1, 2]);\n", + "println!(\"Before squeeze [1,1,2]: shape = {:?}\", tensor.shape());\n", + "\n", + "let squeezed = tensor.squeeze::<1>();\n", + "println!(\"After squeeze: shape = {:?}\", squeezed.shape());\n", + "\n", + "// Unsqueeze - add a dimension of size 1 at specified position\n", + "let tensor = Tensor::::from_floats([1.0, 2.0, 3.0, 4.0], &device).reshape([2, 2]);\n", + "println!(\"Before unsqueeze [2,2]: shape = {:?}\", tensor.shape());\n", + "\n", + "let unsqueezed = tensor.unsqueeze::<3>();\n", + "println!(\"After unsqueeze: shape = {:?}\", unsqueezed.shape());" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Indexing and Slicing" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Original tensor:\n", + "Tensor {\n", + " data:\n", + "[[1.0, 2.0, 3.0, 4.0],\n", + " [5.0, 6.0, 7.0, 8.0],\n", + " [9.0, 10.0, 11.0, 12.0]],\n", + " shape: [3, 4],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "// Create a tensor for indexing examples\n", + "let tensor = Tensor::::from_floats(\n", + " [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0],\n", + "&device\n", + ").reshape([3, 4]);\n", + "println!(\"Original tensor:\\n{}\", tensor);" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Sliced [1..3, 1..4]:\n", + "Tensor {\n", + " data:\n", + "[[6.0, 7.0, 8.0],\n", + " [10.0, 11.0, 12.0]],\n", + " shape: [2, 3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Row 1: Tensor {\n", + " data:\n", + "[[5.0, 6.0, 7.0, 8.0]],\n", + " shape: [1, 4],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Column 2: Tensor {\n", + " data:\n", + "[[3.0],\n", + " [7.0],\n", + " [11.0]],\n", + " shape: [3, 1],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "// Slice tensor - select a portion using ranges\n", + "// Get rows 1-2 (index 1 to end), columns 1-3 (index 1 to 3)\n", + "let sliced = tensor.clone().slice([1..3, 1..4]);\n", + "println!(\"Sliced [1..3, 1..4]:\\n{}\", sliced);\n", + "\n", + "// Get single row\n", + "let row = tensor.clone().slice([1..2, 0..4]);\n", + "println!(\"Row 1: {}\", row);\n", + "\n", + "// Get single column\n", + "let col = tensor.slice([0..3, 2..3]);\n", + "println!(\"Column 2: {}\", col);" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Basic Math Operations" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "a = Tensor {\n", + " data:\n", + "[[1.0, 2.0],\n", + " [3.0, 4.0]],\n", + " shape: [2, 2],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "b = Tensor {\n", + " data:\n", + "[[5.0, 6.0],\n", + " [7.0, 8.0]],\n", + " shape: [2, 2],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "a + b = Tensor {\n", + " data:\n", + "[[6.0, 8.0],\n", + " [10.0, 12.0]],\n", + " shape: [2, 2],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "a - b = Tensor {\n", + " data:\n", + "[[-4.0, -4.0],\n", + " [-4.0, -4.0]],\n", + " shape: [2, 2],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "a * b = Tensor {\n", + " data:\n", + "[[5.0, 12.0],\n", + " [21.0, 32.0]],\n", + " shape: [2, 2],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "a / b = Tensor {\n", + " data:\n", + "[[0.2, 0.33333334],\n", + " [0.42857143, 0.5]],\n", + " shape: [2, 2],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "let a = Tensor::::from_floats([1.0, 2.0, 3.0, 4.0], &device).reshape([2, 2]);\n", + "let b = Tensor::::from_floats([5.0, 6.0, 7.0, 8.0], &device).reshape([2, 2]);\n", + "\n", + "println!(\"a = {}\", a);\n", + "println!(\"b = {}\", b);\n", + "\n", + "// Addition\n", + "let c = a.clone() + b.clone();\n", + "println!(\"a + b = {}\", c);\n", + "\n", + "// Subtraction\n", + "let c = a.clone() - b.clone();\n", + "println!(\"a - b = {}\", c);\n", + "\n", + "// Multiplication (element-wise)\n", + "let c = a.clone() * b.clone();\n", + "println!(\"a * b = {}\", c);\n", + "\n", + "// Division (element-wise)\n", + "let c = a.clone() / b.clone();\n", + "println!(\"a / b = {}\", c);" + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "a = Tensor {\n", + " data:\n", + "[[1.0, 2.0],\n", + " [3.0, 4.0]],\n", + " shape: [2, 2],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "a + 10 = Tensor {\n", + " data:\n", + "[[11.0, 12.0],\n", + " [13.0, 14.0]],\n", + " shape: [2, 2],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "a * 2 = Tensor {\n", + " data:\n", + "[[2.0, 4.0],\n", + " [6.0, 8.0]],\n", + " shape: [2, 2],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "// Scalar operations\n", + "let a = Tensor::::from_floats([1.0, 2.0, 3.0, 4.0], &device).reshape([2, 2]);\n", + "\n", + "println!(\"a = {}\", a);\n", + "\n", + "// Add scalar\n", + "let c = a.clone() + 10.0;\n", + "println!(\"a + 10 = {}\", c);\n", + "\n", + "// Multiply scalar\n", + "let c = a.clone() * 2.0;\n", + "println!(\"a * 2 = {}\", c);" + ] + }, + { + "cell_type": "code", + "execution_count": 30, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "a = Tensor {\n", + " data:\n", + "[[1.0, 2.0],\n", + " [3.0, 4.0]],\n", + " shape: [2, 2],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "b = Tensor {\n", + " data:\n", + "[[5.0, 6.0],\n", + " [7.0, 8.0]],\n", + " shape: [2, 2],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "a @ b (matmul) = Tensor {\n", + " data:\n", + "[[19.0, 22.0],\n", + " [43.0, 50.0]],\n", + " shape: [2, 2],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "// Matrix multiplication\n", + "let a = Tensor::::from_floats([1.0, 2.0, 3.0, 4.0], &device).reshape([2, 2]);\n", + "let b = Tensor::::from_floats([5.0, 6.0, 7.0, 8.0], &device).reshape([2, 2]);\n", + "\n", + "println!(\"a = {}\", a);\n", + "println!(\"b = {}\", b);\n", + "\n", + "let result = a.matmul(b);\n", + "println!(\"a @ b (matmul) = {}\", result);\n", + "\n", + "// Verify (rows of a · columns of b): row1 [1,2] · col1 [5,7] = 1*5+2*7 = 19, row1 [1,2] · col2 [6,8] = 1*6+2*8 = 22\n", + "// row2 [3,4] · col1 [5,7] = 3*5+4*7 = 43, row2 [3,4] · col2 [6,8] = 3*6+4*8 = 50" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Element-wise Math Functions" + ] + }, + { + "cell_type": "code", + "execution_count": 32, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "a = Tensor {\n", + " data:\n", + "[0.0, 1.0, 2.0],\n", + " shape: [3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "exp(a) = Tensor {\n", + " data:\n", + "[1.0, 2.7182817, 7.389056],\n", + " shape: [3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "log(a + 1) = Tensor {\n", + " data:\n", + "[0.0, 0.6931472, 1.0986123],\n", + " shape: [3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "a.powf(2) = Tensor {\n", + " data:\n", + "[0.0, 1.0, 4.0],\n", + " shape: [3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "a.powf(0.5) = Tensor {\n", + " data:\n", + "[0.0, 1.0, 1.4142135],\n", + " shape: [3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "let a: Tensor = Tensor::from_floats([0.0, 1.0, 2.0], &device);\n", + "\n", + "println!(\"a = {}\", a);\n", + "\n", + "// Exponential\n", + "println!(\"exp(a) = {}\", a.clone().exp());\n", + "\n", + "// Natural logarithm\n", + "println!(\"log(a + 1) = {}\", (a.clone() + 1.0).log());\n", + "\n", + "// Power\n", + "println!(\"a.powf(2) = {}\", a.clone().powf_scalar(2.0));\n", + "println!(\"a.powf(0.5) = {}\", a.clone().powf_scalar(0.5));" + ] + }, + { + "cell_type": "code", + "execution_count": 33, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "angles = Tensor {\n", + " data:\n", + "[0.0, 0.7853982, 1.5707964],\n", + " shape: [3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "sin(angles) = Tensor {\n", + " data:\n", + "[0.0, 0.70710677, 1.0],\n", + " shape: [3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "cos(angles) = Tensor {\n", + " data:\n", + "[1.0, 0.70710677, -4.371139e-8],\n", + " shape: [3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "tan(angles) = Tensor {\n", + " data:\n", + "[0.0, 1.0, -22877332.0],\n", + " shape: [3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "// Trigonometric functions\n", + "let angles: Tensor = Tensor::from_floats([0.0, std::f32::consts::PI / 4.0, std::f32::consts::PI / 2.0], &device);\n", + "\n", + "println!(\"angles = {}\", angles);\n", + "println!(\"sin(angles) = {}\", angles.clone().sin());\n", + "println!(\"cos(angles) = {}\", angles.clone().cos());\n", + "println!(\"tan(angles) = {}\", angles.clone().tan());" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Reduction Operations" + ] + }, + { + "cell_type": "code", + "execution_count": 35, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Tensor:\n", + "Tensor {\n", + " data:\n", + "[[1.0, 2.0, 3.0],\n", + " [4.0, 5.0, 6.0]],\n", + " shape: [2, 3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Sum: Tensor {\n", + " data:\n", + "[21.0],\n", + " shape: [1],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Mean: Tensor {\n", + " data:\n", + "[3.5],\n", + " shape: [1],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Product: Tensor {\n", + " data:\n", + "[720.0],\n", + " shape: [1],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Max: Tensor {\n", + " data:\n", + "[6.0],\n", + " shape: [1],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Min: Tensor {\n", + " data:\n", + "[1.0],\n", + " shape: [1],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "let tensor = Tensor::::from_floats([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &device).reshape([2, 3]);\n", + "println!(\"Tensor:\\n{}\", tensor);\n", + "\n", + "// Sum all elements\n", + "println!(\"Sum: {}\", tensor.clone().sum());\n", + "\n", + "// Mean of all elements\n", + "println!(\"Mean: {}\", tensor.clone().mean());\n", + "\n", + "// Product of all elements\n", + "println!(\"Product: {}\", tensor.clone().prod());\n", + "\n", + "// Maximum and minimum\n", + "println!(\"Max: {}\", tensor.clone().max());\n", + "println!(\"Min: {}\", tensor.clone().min());" + ] + }, + { + "cell_type": "code", + "execution_count": 37, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Tensor:\n", + "Tensor {\n", + " data:\n", + "[[1.0, 2.0, 3.0],\n", + " [4.0, 5.0, 6.0]],\n", + " shape: [2, 3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Sum dim 0: Tensor {\n", + " data:\n", + "[[5.0, 7.0, 9.0]],\n", + " shape: [1, 3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Sum dim 1: Tensor {\n", + " data:\n", + "[[6.0],\n", + " [15.0]],\n", + " shape: [2, 1],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Mean dim 0: Tensor {\n", + " data:\n", + "[[2.5, 3.5, 4.5]],\n", + " shape: [1, 3],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "// Reduce along specific dimensions\n", + "let tensor = Tensor::::from_floats([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &device).reshape([2, 3]);\n", + "println!(\"Tensor:\\n{}\", tensor);\n", + "\n", + "// Sum along dimension 0 (columns)\n", + "println!(\"Sum dim 0: {}\", tensor.clone().sum_dim(0));\n", + "\n", + "// Sum along dimension 1 (rows)\n", + "println!(\"Sum dim 1: {}\", tensor.clone().sum_dim(1));\n", + "\n", + "// Mean along dimension 0\n", + "println!(\"Mean dim 0: {}\", tensor.clone().mean_dim(0));" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Comparison and Selection" + ] + }, + { + "cell_type": "code", + "execution_count": 42, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "a = Tensor {\n", + " data:\n", + "[1.0, 5.0, 3.0, 8.0],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "b = Tensor {\n", + " data:\n", + "[4.0, 2.0, 6.0, 7.0],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "a > b: Tensor {\n", + " data:\n", + "[false, true, false, true],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Bool\",\n", + " dtype: \"bool\",\n", + "}\n", + "a < b: Tensor {\n", + " data:\n", + "[true, false, true, false],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Bool\",\n", + " dtype: \"bool\",\n", + "}\n", + "a == b: Tensor {\n", + " data:\n", + "[false, false, false, false],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Bool\",\n", + " dtype: \"bool\",\n", + "}\n" + ] + } + ], + "source": [ + "let a: Tensor = Tensor::from_floats([1.0, 5.0, 3.0, 8.0], &device);\n", + "let b: Tensor = Tensor::from_floats([4.0, 2.0, 6.0, 7.0], &device);\n", + "\n", + "println!(\"a = {}\", a);\n", + "println!(\"b = {}\", b);\n", + "\n", + "// Element-wise comparison returns a boolean tensor\n", + "let greater = a.clone().greater(b.clone());\n", + "println!(\"a > b: {}\", greater);\n", + "\n", + "let less = a.clone().lower(b.clone());\n", + "println!(\"a < b: {}\", less);\n", + "\n", + "let equal = a.clone().equal(b.clone());\n", + "println!(\"a == b: {}\", equal);" + ] + }, + { + "cell_type": "code", + "execution_count": 41, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Original: Tensor {\n", + " data:\n", + "[1.0, 5.0, 3.0, 8.0],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Where > 4, replace with 0: Tensor {\n", + " data:\n", + "[1.0, 0.0, 3.0, 0.0],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n", + "Where > 4, replace with -1: Tensor {\n", + " data:\n", + "[1.0, -1.0, 3.0, -1.0],\n", + " shape: [4],\n", + " device: Cpu,\n", + " backend: \"flex\",\n", + " kind: \"Float\",\n", + " dtype: \"f32\",\n", + "}\n" + ] + } + ], + "source": [ + "// Conditional selection\n", + "let a: Tensor = Tensor::from_floats([1.0, 5.0, 3.0, 8.0], &device);\n", + "\n", + "// mask_where: where condition is true, use replacement value, else keep original value\n", + "let condition = a.clone().greater_elem(4.0);\n", + "let result = a.clone().mask_where(condition, Tensor::zeros([4], &device));\n", + "println!(\"Original: {}\", a);\n", + "println!(\"Where > 4, replace with 0: {}\", result);\n", + "\n", + "// mask_fill: simpler - just replace values matching condition\n", + "let result = a.clone().mask_fill(a.clone().greater_elem(4.0), -1.0);\n", + "println!(\"Where > 4, replace with -1: {}\", result);" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "In this notebook, we covered:\n", + "- **Tensor Creation**: empty, zeros, ones, full, from_floats, random\n", + "- **Shape Operations**: reshape, transpose, flatten, squeeze, unsqueeze\n", + "- **Indexing and Slicing**: slice operation with ranges\n", + "- **Math Operations**: add, sub, mul, div, matmul\n", + "- **Element-wise Functions**: exp, log, powf_scalar, sin, cos, tan\n", + "- **Reduction Operations**: sum, mean, prod, max, min\n", + "- **Comparison**: greater, lower, equal, mask_where, mask_fill\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Rust", + "language": "rust", + "name": "rust" + }, + "language_info": { + "codemirror_mode": "rust", + "file_extension": ".rs", + "mimetype": "text/rust", + "name": "rust", + "pygment_lexer": "rust", + "version": "" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/p2p-remote-training/Cargo.toml b/examples/p2p-remote-training/Cargo.toml new file mode 100644 index 0000000..6cd8fe0 --- /dev/null +++ b/examples/p2p-remote-training/Cargo.toml @@ -0,0 +1,19 @@ +[package] +authors = ["Jonathan Richard"] +edition.workspace = true +license.workspace = true +name = "p2p-remote-training" +publish = false +version.workspace = true +description = "P2P distributed training via Iroh transport with topic-based discovery" + +[lints] +workspace = true + +[dependencies] +burn = { workspace = true, features = ["remote-server", "flex"] } +tokio = { workspace = true, features = ["rt-multi-thread", "macros"] } +iroh = { workspace = true } +blake3 = "1" +tracing = { workspace = true } +tracing-subscriber = { workspace = true, features = ["env-filter", "fmt"] } diff --git a/examples/p2p-remote-training/examples/p2p-remote-training.rs b/examples/p2p-remote-training/examples/p2p-remote-training.rs new file mode 100644 index 0000000..90c8464 --- /dev/null +++ b/examples/p2p-remote-training/examples/p2p-remote-training.rs @@ -0,0 +1,25 @@ +use p2p_remote_training::{run_client, run_server}; + +#[tokio::main] +async fn main() { + let args: Vec = std::env::args().collect(); + match args.get(1).map(String::as_str) { + Some("server") => { + let topic = args.get(2).map(String::as_str).unwrap_or("burn-default"); + run_server(topic).await; + } + Some("client") => { + let topic = args.get(2).map(String::as_str).unwrap_or("burn-default"); + run_client(topic).await; + } + _ => { + eprintln!("usage:"); + eprintln!(" server [topic] start compute server"); + eprintln!(" client [topic] connect and train"); + eprintln!(); + eprintln!("both sides must use the same topic string"); + eprintln!("example: cargo run --example p2p-remote-training -- server my-cluster"); + std::process::exit(1); + } + } +} diff --git a/examples/p2p-remote-training/src/lib.rs b/examples/p2p-remote-training/src/lib.rs new file mode 100644 index 0000000..20b1b87 --- /dev/null +++ b/examples/p2p-remote-training/src/lib.rs @@ -0,0 +1,90 @@ +use burn::server::{Channel, RemoteSecret}; +use burn::tensor::{Device, Distribution, Tensor}; +use iroh::{Endpoint, EndpointId, endpoint::presets}; +use tracing_subscriber::{EnvFilter, fmt}; + +fn init_logging() { + let filter = EnvFilter::try_from_default_env() + .unwrap_or_else(|_| EnvFilter::new("info,burn_remote=debug")); + fmt().with_env_filter(filter).init(); +} + +/// Derive a stable server identity from a human-friendly topic, so both ends agree on the address +/// without exchanging keys. The topic acts as a shared secret here (anyone who knows it can host as +/// this identity), which suits a demo; a real deployment would use `RemoteSecret::random()` and +/// share its `id()`. +fn topic_secret(topic: &str) -> RemoteSecret { + let hash = blake3::hash(format!("burn-p2p:{topic}").as_bytes()); + RemoteSecret::from_bytes(*hash.as_bytes()) +} + +pub async fn run_server(topic: &str) { + init_logging(); + let secret = topic_secret(topic); + tracing::info!(topic, server_id = %secret.id(), "server ready"); + tracing::info!("waiting for clients (press Ctrl-C to stop)"); + burn::server::start_async( + Device::flex(), + Channel::Iroh { + secret: Box::new(secret), + }, + ) + .await; + tracing::info!("server stopped"); +} + +pub async fn run_client(topic: &str) { + let server_id: EndpointId = topic_secret(topic).id(); + + println!("topic : {topic}"); + println!("server id : {server_id}"); + println!("connecting..."); + + let endpoint = Endpoint::builder(presets::N0) + .bind() + .await + .expect("bind failed"); + let device = Device::remote_iroh(&endpoint, server_id, 0); + + println!("connected\n"); + train(&device); +} + +fn train(device: &Device) { + const N: usize = 512; + const STEPS: usize = 80; + const LR: f32 = 0.08; + + println!("target: y = 2.5 * x + 0.5"); + println!("steps : {STEPS} samples: {N}\n"); + + let x = Tensor::<1>::random([N], Distribution::Default, device) * 2.0 - 1.0; + let y_true = x.clone() * 2.5 + 0.5; + + let mut w = Tensor::<1>::from_floats([0.0_f32], device); + let mut b = Tensor::<1>::from_floats([0.0_f32], device); + + println!("{:>5} {:>10}", "step", "loss"); + + for step in 0..STEPS { + let y_pred = x.clone() * w.clone().expand([N]) + b.clone().expand([N]); + let error = y_pred - y_true.clone(); + let loss = (error.clone() * error.clone()).mean(); + let dw = (x.clone() * error.clone()).mean() * 2.0_f32; + let db = error.mean() * 2.0_f32; + + w = w - dw * LR; + b = b - db * LR; + + if step % 10 == 0 || step == STEPS - 1 { + let loss_val = loss.to_data().to_vec::().unwrap()[0]; + println!("{:>5} {:>10.6}", step + 1, loss_val); + } + } + + let w_val = w.to_data().to_vec::().unwrap()[0]; + let b_val = b.to_data().to_vec::().unwrap()[0]; + + println!("\nlearned: y = {w_val:.4} * x + {b_val:.4}"); + println!("target : y = 2.5000 * x + 0.5000"); +} diff --git a/examples/remote-inference-web/Cargo.toml b/examples/remote-inference-web/Cargo.toml new file mode 100644 index 0000000..e0698cd --- /dev/null +++ b/examples/remote-inference-web/Cargo.toml @@ -0,0 +1,31 @@ +[package] +authors = ["Jonathan Richard"] +edition.workspace = true +license = "MIT OR Apache-2.0" +name = "remote-inference-web" +publish = false +version.workspace = true +description = "Run a Burn model in the browser with computation executed on a remote Iroh GPU peer" + +[lints] +workspace = true + +[lib] +crate-type = ["cdylib"] + +[features] +default = ["flex"] + +# Local fallback backend, only for a sensible `Device::default()`; compute runs on the remote. +flex = ["burn/flex"] + +[dependencies] +burn = { workspace = true, features = ["extension", "remote"] } +blake3 = { version = "1", default-features = false } +console_error_panic_hook = { workspace = true } +iroh = { workspace = true } + +# Wasm dependencies +wasm-bindgen = { workspace = true } +wasm-bindgen-futures = { workspace = true } +js-sys = { workspace = true } diff --git a/examples/remote-inference-web/README.md b/examples/remote-inference-web/README.md new file mode 100644 index 0000000..08f1d8a --- /dev/null +++ b/examples/remote-inference-web/README.md @@ -0,0 +1,74 @@ +# Remote MNIST Inference on the Web + +This example runs a Burn model **in the browser** while executing every tensor operation on a +**remote compute peer** reached over [Iroh](https://iroh.computer/). The browser holds only the +model definition and weights; convolutions, matrix multiplies and everything else run on the +peer's backend (CPU or GPU). Only the 28x28 input and the 10 output probabilities cross the wire. + +It mirrors the [`mnist-inference-web`](../mnist-inference-web) demo, but swaps the local WebAssembly +backend for the `burn-remote` Iroh client. + +## Why this is interesting + +- **Web scripting and experimentation against real GPUs.** A browser tab can drive a CUDA/Metal/Vulkan + backend without shipping a native build to the user. +- **Models larger than the browser sandbox.** Memory and compute live on the peer, so the browser + is not bound by WebGPU limits or tab memory. +- **One model, many backends.** The same client code targets whatever backend the peer hosts. + +## How it works + +The browser and the compute peer find each other through a shared **topic** string. Both sides hash +the topic to derive the peer's Iroh identity, so no node id has to be copied around: + +``` +topic --blake3--> secret key --> peer endpoint identity + (peer binds the secret key; the browser derives its public half) +``` + +The client then binds its own Iroh endpoint, opens an authenticated QUIC session to the peer +(through a relay when a direct path is not available), and ships operations as they are submitted. + +## Running it + +### 1. Start a compute peer (native) + +CPU backend: + +```sh +cargo run -p remote-compute-peer -- burn-web +``` + +GPU backend (wgpu): + +```sh +cargo run -p remote-compute-peer --features wgpu -- burn-web +``` + +The argument (`burn-web`) is the topic; it must match what you type in the browser. + +### 2. Build the web client + +```sh +cd examples/remote-inference-web +./build-for-web.sh +``` + +### 3. Serve and open + +```sh +./run-server.sh +``` + +Open , enter the same topic, click **Connect**, then draw a digit. The +probabilities are computed on the peer and streamed back. + +## Notes + +- The compute peer is model-agnostic; it just executes operations. Swapping the model on the + client side requires no change to the peer. +- `model.bpk` is the trained MNIST model from the [`mnist`](../mnist) example, identical to the one + used by `mnist-inference-web`. +- Connecting through public relays requires outbound network access from both the browser and the + peer. For a fully local setup, configure both endpoints with a self-hosted relay or direct + addressing through the Iroh `Endpoint` builder. diff --git a/examples/remote-inference-web/build-for-web.sh b/examples/remote-inference-web/build-for-web.sh new file mode 100755 index 0000000..59df1c7 --- /dev/null +++ b/examples/remote-inference-web/build-for-web.sh @@ -0,0 +1,16 @@ +#!/usr/bin/env bash + +# Add wasm32 target for compiler. +rustup target add wasm32-unknown-unknown + +if ! command -v wasm-pack &> /dev/null +then + echo "wasm-pack could not be found. Installing ..." + cargo install wasm-pack +fi + +# Iroh's wasm randomness backend is selected through getrandom's wasm_js cfg. +RUSTFLAGS='-C embed-bitcode=yes -C codegen-units=1 -C opt-level=3 --cfg web_sys_unstable_apis --cfg getrandom_backend="wasm_js"' + +mkdir -p pkg +wasm-pack build --out-dir pkg --release --target web --no-typescript diff --git a/examples/remote-inference-web/index.html b/examples/remote-inference-web/index.html new file mode 100644 index 0000000..6bb06f7 --- /dev/null +++ b/examples/remote-inference-web/index.html @@ -0,0 +1,179 @@ + + + + + + Burn Remote MNIST Inference (WebAssembly + Iroh) + + + + + + + + + + + + +

Burn Remote MNIST Inference

+ +
+ Compute peer topic: + + + not connected +
+ + + + + + + + + + + + + + + + + +
Draw a digit hereCropped and scaledProbability result
+ + + + + + +
+ + + + diff --git a/examples/remote-inference-web/index.js b/examples/remote-inference-web/index.js new file mode 100644 index 0000000..9896ff5 --- /dev/null +++ b/examples/remote-inference-web/index.js @@ -0,0 +1,175 @@ +/** + * + * This demo is part of Burn project: https://github.com/tracel-ai/burn + * + * Released under a dual license: + * https://github.com/tracel-ai/burn/blob/main/LICENSE-MIT + * https://github.com/tracel-ai/burn/blob/main/LICENSE-APACHE + * + */ + +/** + * Auto crops the image, scales to 28x28 pixel image, and returns as grayscale image. + * @param {object} mainContext - The 2d context of the source canvas. + * @param {object} cropContext - The 2d context of an intermediate hidden canvas. + * @param {object} scaledContext - The 2d context of the destination 28x28 canvas. + */ +export function cropScaleGetImageData(mainContext, cropContext, scaledContext) { + + const cropEl = cropContext.canvas; + + // Get the auto-cropped image data and put into the intermediate/hidden canvas + cropContext.fillStyle = "rgba(255, 255, 255, 255)"; // white non-transparent color + cropContext.fillRect(0, 0, cropEl.width, cropEl.height); + cropContext.save(); + const [w, h, croppedImage] = cropImageFromCanvas(mainContext); + cropEl.width = Math.max(w, h) * 1.2; + cropEl.height = Math.max(w, h) * 1.2; + const leftPadding = (cropEl.width - w) / 2; + const topPadding = (cropEl.height - h) / 2; + cropContext.putImageData(croppedImage, leftPadding, topPadding); + + // Copy image data to scale 28x28 canvas + scaledContext.save(); + scaledContext.clearRect(0, 0, scaledContext.canvas.height, scaledContext.canvas.width); + scaledContext.fillStyle = "rgba(255, 255, 255, 255)"; // white non-transparent color + scaledContext.fillRect(0, 0, cropEl.width, cropEl.height); + scaledContext.scale(28.0 / cropContext.canvas.width, 28.0 / cropContext.canvas.height); + scaledContext.drawImage(cropEl, 0, 0); + + // Extract image data and convert into single value (greyscale) array + const data = rgba2gray(scaledContext.getImageData(0, 0, 28, 28).data); + scaledContext.restore(); + + return data; +} + +/** + * Converts RGBA image data from canvas to grayscale (0 is white & 255 is black). + * @param {int[]} - Image data. + */ +export function rgba2gray(data) { + let converted = new Float32Array(data.length / 4); + + // Data is stored as [r0,g0,b0,a0, ... r[n],g[n],b[n],a[n]] where n is number of pixels. + for (let i = 0; i < data.length; i += 4) { + let r = 255 - data[i]; // red + let g = 255 - data[i + 1]; // green + let b = 255 - data[i + 2]; // blue + let a = 255 - data[i + 3]; // alpha + + // Use RGB grayscale coefficients (https://imagej.nih.gov/ij/docs/menus/image.html) + let y = 0.299 * r + 0.587 * g + 0.114 * b; + converted[i / 4] = y; // 4 times fewer data points but the same number of pixels. + } + return converted; +} + +/** + * Auto crops a canvas images and returns its image data. + * @param {object} ctx - canvas 2d context. + * src: https://stackoverflow.com/a/22267731 + */ +export function cropImageFromCanvas(ctx) { + let canvas = ctx.canvas, + w = canvas.width, + h = canvas.height, + pix = { x: [], y: [] }, + imageData = ctx.getImageData(0, 0, canvas.width, canvas.height), + x, + y, + index; + for (y = 0; y < h; y++) { + for (x = 0; x < w; x++) { + index = (y * w + x) * 4; + + let r = imageData.data[index]; + let g = imageData.data[index + 1]; + let b = imageData.data[index + 2]; + // On some browsers the canvas has a grey border which prevents cropping if we do min != 255 + if (Math.min(r, g, b) < 240) { + pix.x.push(x); + pix.y.push(y); + } + } + } + pix.x.sort(function (a, b) { + return a - b; + }); + pix.y.sort(function (a, b) { + return a - b; + }); + let n = pix.x.length - 1; + w = 1 + pix.x[n] - pix.x[0]; + h = 1 + pix.y[n] - pix.y[0]; + return [w, h, ctx.getImageData(pix.x[0], pix.y[0], w, h, { willReadFrequently: true })]; +} + +/** + * Truncates number to a given decimal position + * @param {number} num - Number to truncate. + * @param {number} fixed - Decimal positions. + * src: https://stackoverflow.com/a/11818658 + */ +export function toFixed(num, fixed) { + const re = new RegExp('^-?\\d+(?:\.\\d{0,' + (fixed || -1) + '})?'); + return num.toString().match(re)[0]; +} + +/** + * Looks up element by an id. + * @param {string} - Element id. + */ +export function $(id) { + return document.getElementById(id); +} + +/** + * Helper function that builds a chart using Chart.js library. + * @param {object} chartEl - Chart canvas element. + * + * NOTE: Assumes chart.js is loaded into the global. + */ +export function chartConfigBuilder(chartEl) { + Chart.register(ChartDataLabels); + return new Chart(chartEl, { + plugins: [ChartDataLabels], + type: "bar", + data: { + labels: ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"], + datasets: [ + { + data: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + borderWidth: 0, + fill: true, + backgroundColor: "#247ABF", + }, + ], + }, + options: { + responsive: false, + maintainAspectRatio: false, + animation: true, + plugins: { + legend: { + display: false, + }, + tooltip: { + enabled: true, + }, + datalabels: { + color: "white", + formatter: function (value, context) { + return toFixed(value, 2); + }, + }, + }, + scales: { + y: { + beginAtZero: true, + max: 1.0, + }, + }, + }, + }); +} \ No newline at end of file diff --git a/examples/remote-inference-web/model.bpk b/examples/remote-inference-web/model.bpk new file mode 100644 index 0000000..0eb0e55 Binary files /dev/null and b/examples/remote-inference-web/model.bpk differ diff --git a/examples/remote-inference-web/run-server.sh b/examples/remote-inference-web/run-server.sh new file mode 100755 index 0000000..0ce038f --- /dev/null +++ b/examples/remote-inference-web/run-server.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash + +# Opening index.html file directly by a browser does not work because of +# the security restrictions by the browser. Viewing the HTML file will fail with +# this error message: + +# ``` +# Access to script at +# 'file:///Users/user/Projects/burn-mac/examples/mnist-inference-web/pkg/mnist_inference_web.js' +# from origin 'null' has been blocked by CORS policy: +# Cross origin requests are only supported for protocol schemes: +# http, data, isolated-app, chrome-extension, chrome, https, chrome-untrusted. +# ``` +# So that's why running a local HTTP server is needed. + +if ! command -v python3 &> /dev/null +then + echo "python3 could not be found. Running server requires python3." + exit +fi + +echo "Running local python HTTP server on port 8000 ..." +python3 -m http.server 8000 diff --git a/examples/remote-inference-web/src/lib.rs b/examples/remote-inference-web/src/lib.rs new file mode 100644 index 0000000..6405bf1 --- /dev/null +++ b/examples/remote-inference-web/src/lib.rs @@ -0,0 +1,8 @@ +#![cfg_attr(not(test), no_std)] + +extern crate alloc; + +pub mod model; + +#[cfg(target_family = "wasm")] +pub mod web; diff --git a/examples/remote-inference-web/src/model.rs b/examples/remote-inference-web/src/model.rs new file mode 100644 index 0000000..4856399 --- /dev/null +++ b/examples/remote-inference-web/src/model.rs @@ -0,0 +1,104 @@ +use burn::{ + nn::{ + BatchNorm, PaddingConfig2d, + pool::{MaxPool2d, MaxPool2dConfig}, + }, + prelude::*, +}; + +#[derive(Module, Debug)] +pub struct Model { + conv1: ConvBlock, + conv2: ConvBlock, + dropout: nn::Dropout, + fc1: nn::Linear, + fc2: nn::Linear, + fc3: nn::Linear, + activation: nn::Gelu, +} + +const NUM_CLASSES: usize = 10; + +impl Model { + pub fn new(device: &Device) -> Self { + let conv1 = ConvBlock::new([1, 64], [3, 3], device, true); // out: max_pool -> [Batch,32,13,13] + let conv2 = ConvBlock::new([64, 64], [3, 3], device, true); // out: max_pool -> [Batch,64,5,5] + let hidden_size = 64 * 5 * 5; + let fc1 = nn::LinearConfig::new(hidden_size, 128).init(device); + let fc2 = nn::LinearConfig::new(128, 128).init(device); + let fc3 = nn::LinearConfig::new(128, NUM_CLASSES).init(device); + + let dropout = nn::DropoutConfig::new(0.25).init(); + + Self { + conv1, + conv2, + dropout, + fc1, + fc2, + fc3, + activation: nn::Gelu::new(), + } + } + + pub fn forward(&self, input: Tensor<3>) -> Tensor<2> { + let [batch_size, height, width] = input.dims(); + + let x = input.reshape([batch_size, 1, height, width]).detach(); + let x = self.conv1.forward(x); + let x = self.conv2.forward(x); + + let [batch_size, channels, height, width] = x.dims(); + let x = x.reshape([batch_size, channels * height * width]); + + let x = self.fc1.forward(x); + let x = self.activation.forward(x); + let x = self.dropout.forward(x); + let x = self.fc2.forward(x); + let x = self.activation.forward(x); + let x = self.dropout.forward(x); + + self.fc3.forward(x) + } +} + +#[derive(Module, Debug)] +pub struct ConvBlock { + conv: nn::conv::Conv2d, + norm: BatchNorm, + pool: Option, + activation: nn::Relu, +} + +impl ConvBlock { + pub fn new(channels: [usize; 2], kernel_size: [usize; 2], device: &Device, pool: bool) -> Self { + let conv = nn::conv::Conv2dConfig::new(channels, kernel_size) + .with_padding(PaddingConfig2d::Valid) + .init(device); + let norm = nn::BatchNormConfig::new(channels[1]).init(device); + let pool = if pool { + Some(MaxPool2dConfig::new([2, 2]).with_strides([2, 2]).init()) + } else { + None + }; + + Self { + conv, + norm, + pool, + activation: nn::Relu::new(), + } + } + + pub fn forward(&self, input: Tensor<4>) -> Tensor<4> { + let x = self.conv.forward(input); + let x = self.norm.forward(x); + let x = self.activation.forward(x); + + if let Some(pool) = &self.pool { + pool.forward(x) + } else { + x + } + } +} diff --git a/examples/remote-inference-web/src/web.rs b/examples/remote-inference-web/src/web.rs new file mode 100644 index 0000000..411ce7c --- /dev/null +++ b/examples/remote-inference-web/src/web.rs @@ -0,0 +1,75 @@ +//! Browser entry point: a MNIST classifier whose tensor operations run on a remote Iroh compute +//! peer. The model is defined client-side, but every op executes on the peer; only the input and +//! output probabilities cross the wire. + +use alloc::format; +use alloc::string::{String, ToString}; +use alloc::vec::Vec; + +use wasm_bindgen::prelude::*; + +use burn::backend::remote::{EndpointId, RemoteSecret}; +use burn::module::Module; +use burn::store::ModuleRecord; +use burn::tensor::{Bytes, Device, Tensor, activation::softmax}; +use iroh::{Endpoint, endpoint::presets}; + +use crate::model::Model; + +/// Trained MNIST parameters in the burnpack format, produced by the `mnist` example. +static STATE_ENCODED: &[u8] = include_bytes!("../model.bpk"); + +#[wasm_bindgen(start)] +pub fn start() { + console_error_panic_hook::set_once(); +} + +/// Derive the server's identity from the shared topic; both ends compute the same id from the same +/// string (a demo convenience; see the native example for the security note). +fn server_id(topic: &str) -> EndpointId { + let hash = blake3::hash(format!("burn-p2p:{topic}").as_bytes()); + RemoteSecret::from_bytes(*hash.as_bytes()).id() +} + +#[wasm_bindgen] +pub struct RemoteMnist { + device: Device, + model: Model, +} + +#[wasm_bindgen] +impl RemoteMnist { + /// Connect to the peer reachable under `topic` and load the model onto it. + pub async fn connect(topic: String) -> Result { + console_error_panic_hook::set_once(); + + let endpoint = Endpoint::builder(presets::N0) + .bind() + .await + .map_err(|err| err.to_string())?; + let device = Device::remote_iroh_async(&endpoint, server_id(&topic), 0).await; + + let record = ModuleRecord::from_bytes(Bytes::from_bytes_vec(STATE_ENCODED.to_vec())) + .map_err(|err| format!("Failed to decode model weights: {err}"))?; + let model = Model::new(&device).load_record(record); + + Ok(Self { device, model }) + } + + /// Classify a 28x28 grayscale image (row-major length-784 `f32`, pixels in `[0, 255]`), + /// returning the 10 class probabilities. + pub async fn inference(&self, input: &[f32]) -> Result, String> { + let input = Tensor::<1>::from_floats(input, &self.device).reshape([1, 28, 28]); + + // MNIST training mean/std (from the PyTorch example). + let input = ((input / 255) - 0.1307) / 0.3081; + + let output = softmax(self.model.forward(input), 1); + let data = output + .into_data_async() + .await + .map_err(|err| format!("Failed to read inference result: {err:?}"))?; + + Ok(data.iter::().collect()) + } +} diff --git a/examples/server/Cargo.toml b/examples/server/Cargo.toml new file mode 100644 index 0000000..f9e44e1 --- /dev/null +++ b/examples/server/Cargo.toml @@ -0,0 +1,21 @@ +[package] +authors = ["nathanielsimard "] +edition.workspace = true +license.workspace = true +name = "server" +publish = false +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["webgpu"] +cuda = ["burn/cuda"] +webgpu = ["burn/webgpu"] +vulkan = ["burn/vulkan"] +rocm = ["burn/rocm"] +flex = ["burn/flex"] + +[dependencies] +burn = { workspace = true, features = ["default", "remote-server", "remote-websocket"] } diff --git a/examples/server/burn.toml b/examples/server/burn.toml new file mode 100644 index 0000000..723a097 --- /dev/null +++ b/examples/server/burn.toml @@ -0,0 +1,9 @@ +[cubecl.profiling] +logger = { log = "info", level = "disabled" } + +[cubecl.autotune] +logger = { log = "info", level = "minimal" } + +[cubecl.compilation] +logger = { log = "info", level = "disabled" } +cache = "target" diff --git a/examples/server/examples/server.rs b/examples/server/examples/server.rs new file mode 100644 index 0000000..8d79ebb --- /dev/null +++ b/examples/server/examples/server.rs @@ -0,0 +1,3 @@ +fn main() { + server::start(); +} diff --git a/examples/server/src/lib.rs b/examples/server/src/lib.rs new file mode 100644 index 0000000..1e725fa --- /dev/null +++ b/examples/server/src/lib.rs @@ -0,0 +1,14 @@ +#![recursion_limit = "141"] + +use burn::{server::Channel, tensor::Device}; + +pub fn start() { + let port = std::env::var("REMOTE_BACKEND_PORT") + .map(|port| match port.parse::() { + Ok(val) => val, + Err(err) => panic!("Invalid port, got {port} with error {err}"), + }) + .unwrap_or(3000); + + burn::server::start(Device::default(), Channel::WebSocket { port }); +} diff --git a/examples/simple-regression/Cargo.toml b/examples/simple-regression/Cargo.toml new file mode 100644 index 0000000..55f590a --- /dev/null +++ b/examples/simple-regression/Cargo.toml @@ -0,0 +1,29 @@ +[package] +authors = ["aasheeshsingh "] +edition.workspace = true +license.workspace = true +name = "simple-regression" +publish = false +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["burn/dataset", "burn/sqlite-bundled"] +flex = ["burn/flex"] +tch-cpu = ["burn/tch"] +tch-gpu = ["burn/tch"] +wgpu = ["burn/wgpu"] +remote = ["burn/remote"] + +[dependencies] +burn = { workspace = true, features = ["default", "train"] } + +# Serialization +log = { workspace = true } +serde = { workspace = true, features = ["std", "derive"] } + +# Displaying results +textplots = "0.8.7" +rgb = "0.8.52" diff --git a/examples/simple-regression/README.md b/examples/simple-regression/README.md new file mode 100644 index 0000000..7d8deb5 --- /dev/null +++ b/examples/simple-regression/README.md @@ -0,0 +1,32 @@ +# Regression + +The example shows you how to: + +- Define a custom dataset for regression problems. We implement the + [California Housing Dataset](https://huggingface.co/datasets/gvlassis/california_housing) from + HuggingFace hub. The dataset is also available as part of toy regression datasets in + sklearn[datasets](https://scikit-learn.org/stable/datasets/real_world.html#california-housing-dataset). +- Create a data pipeline from a raw dataset to a batched fast DataLoader with min-max feature + scaling. +- Define a Simple NN model for regression using Burn Modules. + +> **Note** +> This example makes use of the HuggingFace [`datasets`](https://huggingface.co/docs/datasets/index) +> library to download the datasets. Make sure you have [Python](https://www.python.org/downloads/) +> installed on your computer. + +The example can be run like so: + +```bash +git clone https://github.com/tracel-ai/burn.git +cd burn +# Use the --release flag to really speed up training. +echo "Using flex backend" +cargo run --example regression --release --features flex # CPU Flex Backend - f32 +echo "Using tch backend" +export TORCH_CUDA_VERSION=cu128 # Set the cuda version +cargo run --example regression --release --features tch-gpu # GPU Tch Backend - f32 +cargo run --example regression --release --features tch-cpu # CPU Tch Backend - f32 +echo "Using wgpu backend" +cargo run --example regression --release --features wgpu +``` diff --git a/examples/simple-regression/examples/regression.rs b/examples/simple-regression/examples/regression.rs new file mode 100644 index 0000000..20b4f39 --- /dev/null +++ b/examples/simple-regression/examples/regression.rs @@ -0,0 +1,73 @@ +use burn::tensor::Device; +use simple_regression::{inference, training}; + +static ARTIFACT_DIR: &str = "/tmp/burn-example-regression"; + +#[cfg(feature = "flex")] +mod flex { + use burn::tensor::Device; + + pub fn run() { + super::run(Device::flex()); + } +} + +#[cfg(feature = "tch-gpu")] +mod tch_gpu { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + #[cfg(not(target_os = "macos"))] + let device = Device::libtorch_cuda(DeviceIndex::Default); + #[cfg(target_os = "macos")] + let device = Device::libtorch_mps(); + + super::run(device); + } +} + +#[cfg(feature = "wgpu")] +mod wgpu { + use burn::tensor::{Device, DeviceKind}; + + pub fn run() { + super::run(Device::wgpu(DeviceKind::DefaultDevice)); + } +} + +#[cfg(feature = "tch-cpu")] +mod tch_cpu { + use burn::tensor::Device; + pub fn run() { + super::run(Device::libtorch()); + } +} + +// #[cfg(feature = "remote")] +// mod remote { +// use burn::backend::{RemoteBackend, remote::RemoteDevice}; + +// pub fn run() { +// let device = RemoteDevice::default(); +// super::run::(device); +// } +// } + +/// Train a regression model and predict results on a number of samples. +pub fn run(device: Device) { + training::run(ARTIFACT_DIR, device.clone()); + inference::infer(ARTIFACT_DIR, device) +} + +fn main() { + #[cfg(feature = "flex")] + flex::run(); + #[cfg(feature = "tch-gpu")] + tch_gpu::run(); + #[cfg(feature = "tch-cpu")] + tch_cpu::run(); + #[cfg(feature = "wgpu")] + wgpu::run(); + #[cfg(feature = "remote")] + remote::run(); +} diff --git a/examples/simple-regression/src/dataset.rs b/examples/simple-regression/src/dataset.rs new file mode 100644 index 0000000..d7da6e0 --- /dev/null +++ b/examples/simple-regression/src/dataset.rs @@ -0,0 +1,175 @@ +use burn::{ + data::{ + dataloader::batcher::Batcher, + dataset::{Dataset, HuggingfaceDatasetLoader, SqliteDataset}, + }, + prelude::*, +}; + +pub const NUM_FEATURES: usize = 8; + +// Pre-computed statistics for the housing dataset features +const FEATURES_MIN: [f32; NUM_FEATURES] = [0.4999, 1., 0.8461, 0.375, 3., 0.6923, 32.54, -124.35]; +const FEATURES_MAX: [f32; NUM_FEATURES] = [ + 15., 52., 141.9091, 34.0667, 35682., 1243.3333, 41.95, -114.31, +]; + +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +pub struct HousingDistrictItem { + /// Median income + #[serde(rename = "MedInc")] + pub median_income: f32, + + /// Median house age + #[serde(rename = "HouseAge")] + pub house_age: f32, + + /// Average number of rooms per household + #[serde(rename = "AveRooms")] + pub avg_rooms: f32, + + /// Average number of bedrooms per household + #[serde(rename = "AveBedrms")] + pub avg_bedrooms: f32, + + /// Block group population + #[serde(rename = "Population")] + pub population: f32, + + /// Average number of household members + #[serde(rename = "AveOccup")] + pub avg_occupancy: f32, + + /// Block group latitude + #[serde(rename = "Latitude")] + pub latitude: f32, + + /// Block group longitude + #[serde(rename = "Longitude")] + pub longitude: f32, + + /// Median house value (in 100 000$) + #[serde(rename = "MedHouseVal")] + pub median_house_value: f32, +} + +pub struct HousingDataset { + dataset: SqliteDataset, +} + +impl Dataset for HousingDataset { + fn get(&self, index: usize) -> Option { + self.dataset.get(index) + } + + fn len(&self) -> usize { + self.dataset.len() + } +} + +impl HousingDataset { + pub fn train() -> Self { + Self::new("train") + } + + pub fn validation() -> Self { + Self::new("validation") + } + + pub fn test() -> Self { + Self::new("test") + } + + pub fn new(split: &str) -> Self { + let dataset: SqliteDataset = + HuggingfaceDatasetLoader::new("gvlassis/california_housing") + .dataset(split) + .unwrap(); + + Self { dataset } + } +} + +/// Normalizer for the housing dataset. +#[derive(Clone, Debug)] +pub struct Normalizer { + pub min: Tensor<2>, + pub max: Tensor<2>, +} + +impl Normalizer { + /// Creates a new normalizer. + pub fn new(device: &Device, min: &[f32], max: &[f32]) -> Self { + let min = Tensor::<1>::from_floats(min, device).unsqueeze(); + let max = Tensor::<1>::from_floats(max, device).unsqueeze(); + Self { min, max } + } + + /// Normalizes the input image according to the housing dataset min/max. + pub fn normalize(&self, input: Tensor<2>) -> Tensor<2> { + (input - self.min.clone()) / (self.max.clone() - self.min.clone()) + } + + /// Returns a new normalizer on the given device. + pub fn to_device(&self, device: &Device) -> Self { + Self { + min: self.min.clone().to_device(device), + max: self.max.clone().to_device(device), + } + } +} + +#[derive(Clone, Debug)] +pub struct HousingBatcher { + normalizer: Normalizer, +} + +#[derive(Clone, Debug)] +pub struct HousingBatch { + pub inputs: Tensor<2>, + pub targets: Tensor<1>, +} + +impl HousingBatcher { + pub fn new(device: &Device) -> Self { + Self { + normalizer: Normalizer::new(device, &FEATURES_MIN, &FEATURES_MAX), + } + } +} + +impl Batcher for HousingBatcher { + fn batch(&self, items: Vec, device: &Device) -> HousingBatch { + let mut inputs: Vec> = Vec::new(); + + for item in items.iter() { + let input_tensor = Tensor::<1>::from_floats( + [ + item.median_income, + item.house_age, + item.avg_rooms, + item.avg_bedrooms, + item.population, + item.avg_occupancy, + item.latitude, + item.longitude, + ], + device, + ); + + inputs.push(input_tensor.unsqueeze()); + } + + let inputs = Tensor::cat(inputs, 0); + let inputs = self.normalizer.to_device(device).normalize(inputs); + + let targets = items + .iter() + .map(|item| Tensor::<1>::from_floats([item.median_house_value], device)) + .collect(); + + let targets = Tensor::cat(targets, 0); + + HousingBatch { inputs, targets } + } +} diff --git a/examples/simple-regression/src/inference.rs b/examples/simple-regression/src/inference.rs new file mode 100644 index 0000000..b4e18e6 --- /dev/null +++ b/examples/simple-regression/src/inference.rs @@ -0,0 +1,56 @@ +use burn::{ + data::{dataloader::batcher::Batcher, dataset::Dataset}, + module::Module, + store::ModuleRecord, + tensor::Device, +}; +use rgb::RGB8; +use textplots::{Chart, ColorPlot, Shape}; + +use crate::{ + dataset::{HousingBatcher, HousingDataset, HousingDistrictItem}, + model::RegressionModelConfig, +}; + +pub fn infer(artifact_dir: &str, device: impl Into) { + let device = device.into(); + let record = ModuleRecord::load(format!("{artifact_dir}/model")) + .expect("Trained model should exist; run train first"); + + let model = RegressionModelConfig::new() + .init(&device) + .load_record(record); + + // Use a sample of 1000 items from the test split + let dataset = HousingDataset::test(); + let items: Vec = dataset.iter().take(1000).collect(); + + let batcher = HousingBatcher::new(&device); + let batch = batcher.batch(items.clone(), &device); + let predicted = model.forward(batch.inputs); + let targets = batch.targets; + + // Display the predicted vs expected values + let predicted = predicted.squeeze_dim::<1>(1).into_data(); + let expected = targets.into_data(); + + let points = predicted + .iter::() + .zip(expected.iter::()) + .collect::>(); + + println!("Predicted vs. Expected Median House Value (in 100,000$)"); + Chart::new_with_y_range(120, 60, 0., 5., 0., 5.) + .linecolorplot( + &Shape::Points(&points), + RGB8 { + r: 255, + g: 85, + b: 85, + }, + ) + .display(); + + // Print a single numeric value as an example + println!("Predicted {} Expected {}", points[0].0, points[0].1); +} diff --git a/examples/simple-regression/src/lib.rs b/examples/simple-regression/src/lib.rs new file mode 100644 index 0000000..1a167ff --- /dev/null +++ b/examples/simple-regression/src/lib.rs @@ -0,0 +1,4 @@ +pub mod dataset; +pub mod inference; +pub mod model; +pub mod training; diff --git a/examples/simple-regression/src/model.rs b/examples/simple-regression/src/model.rs new file mode 100644 index 0000000..dd4d5d7 --- /dev/null +++ b/examples/simple-regression/src/model.rs @@ -0,0 +1,80 @@ +use crate::dataset::{HousingBatch, NUM_FEATURES}; +use burn::{ + nn::{ + Linear, LinearConfig, Relu, + loss::{MseLoss, Reduction::Mean}, + }, + prelude::*, + train::{InferenceStep, RegressionOutput, TrainOutput, TrainStep}, +}; + +#[derive(Module, Debug)] +pub struct RegressionModel { + input_layer: Linear, + output_layer: Linear, + activation: Relu, +} + +#[derive(Config, Debug)] +pub struct RegressionModelConfig { + #[config(default = 64)] + pub hidden_size: usize, +} + +impl RegressionModelConfig { + pub fn init(&self, device: &Device) -> RegressionModel { + let input_layer = LinearConfig::new(NUM_FEATURES, self.hidden_size) + .with_bias(true) + .init(device); + let output_layer = LinearConfig::new(self.hidden_size, 1) + .with_bias(true) + .init(device); + + RegressionModel { + input_layer, + output_layer, + activation: Relu::new(), + } + } +} + +impl RegressionModel { + pub fn forward(&self, input: Tensor<2>) -> Tensor<2> { + let x = self.input_layer.forward(input); + let x = self.activation.forward(x); + self.output_layer.forward(x) + } + + pub fn forward_step(&self, item: HousingBatch) -> RegressionOutput { + let targets: Tensor<2> = item.targets.unsqueeze_dim(1); + let output: Tensor<2> = self.forward(item.inputs); + + let loss = MseLoss::new().forward(output.clone(), targets.clone(), Mean); + + RegressionOutput { + loss, + output, + targets, + } + } +} + +impl TrainStep for RegressionModel { + type Input = HousingBatch; + type Output = RegressionOutput; + + fn step(&self, item: HousingBatch) -> TrainOutput { + let item = self.forward_step(item); + + TrainOutput::new(self, item.loss.backward(), item) + } +} + +impl InferenceStep for RegressionModel { + type Input = HousingBatch; + type Output = RegressionOutput; + + fn step(&self, item: HousingBatch) -> RegressionOutput { + self.forward_step(item) + } +} diff --git a/examples/simple-regression/src/training.rs b/examples/simple-regression/src/training.rs new file mode 100644 index 0000000..7f83c64 --- /dev/null +++ b/examples/simple-regression/src/training.rs @@ -0,0 +1,89 @@ +use std::path::PathBuf; + +use crate::dataset::{HousingBatcher, HousingDataset}; +use crate::model::RegressionModelConfig; +use burn::optim::AdamConfig; +use burn::train::{Learner, SupervisedTraining}; +use burn::{ + data::{dataloader::DataLoaderBuilder, dataset::Dataset}, + prelude::*, + train::metric::LossMetric, +}; + +#[derive(Config, Debug)] +pub struct ExpConfig { + #[config(default = 100)] + pub num_epochs: usize, + + #[config(default = 2)] + pub num_workers: usize, + + #[config(default = 1337)] + pub seed: u64, + + pub optimizer: AdamConfig, + + #[config(default = 256)] + pub batch_size: usize, +} + +fn create_artifact_dir(artifact_dir: &str) { + std::fs::remove_file(PathBuf::from(artifact_dir).join("experiment.log")).ok(); + std::fs::create_dir_all(artifact_dir).ok(); +} + +pub fn run(artifact_dir: &str, device: impl Into) { + create_artifact_dir(artifact_dir); + + // Config + let optimizer = AdamConfig::new(); + let config = ExpConfig::new(optimizer); + + let device = device.into(); + device.seed(config.seed); + let autodiff_device = device.clone().autodiff(); + + let model = RegressionModelConfig::new().init(&autodiff_device); + + // Define train/valid datasets and dataloaders + let train_dataset = HousingDataset::train(); + let valid_dataset = HousingDataset::validation(); + + println!("Train Dataset Size: {}", train_dataset.len()); + println!("Valid Dataset Size: {}", valid_dataset.len()); + + let batcher_train = HousingBatcher::new(&autodiff_device); + let batcher_test = HousingBatcher::new(&device); + + let dataloader_train = DataLoaderBuilder::new(batcher_train) + .batch_size(config.batch_size) + .shuffle(config.seed) + .num_workers(config.num_workers) + .build(train_dataset); + + let dataloader_test = DataLoaderBuilder::new(batcher_test) + .batch_size(config.batch_size) + .shuffle(config.seed) + .num_workers(config.num_workers) + .build(valid_dataset); + + // Model + let training = SupervisedTraining::new(artifact_dir, dataloader_train, dataloader_test) + .metric_train_numeric(LossMetric::new()) + .metric_valid_numeric(LossMetric::new()) + .with_default_checkpointers() + .num_epochs(config.num_epochs) + .summary(); + + let result = training.launch(Learner::new(model, config.optimizer.init(), 1e-3)); + + config + .save(format!("{artifact_dir}/config.json").as_str()) + .unwrap(); + + result + .model + .into_record() + .save(format!("{artifact_dir}/model")) + .expect("Failed to save trained model"); +} diff --git a/examples/text-classification/Cargo.toml b/examples/text-classification/Cargo.toml new file mode 100644 index 0000000..a5b1e28 --- /dev/null +++ b/examples/text-classification/Cargo.toml @@ -0,0 +1,52 @@ +[package] +authors = ["nathanielsimard "] +edition.workspace = true +license.workspace = true +name = "text-classification" +publish = false +version.workspace = true + +[lints] +workspace = true + +[features] +default = [] +f16 = [] +flex32 = [] +tch-cpu = ["burn/tch"] +tch-gpu = ["burn/tch"] +wgpu = ["burn/wgpu"] +vulkan = ["burn/vulkan"] +flex = ["burn/flex"] +remote = ["burn/remote-websocket"] +cuda = ["burn/cuda"] +rocm = ["burn/rocm"] +metal = ["burn/metal"] +# DDP (distributed data-parallel) training. The distributed runtime is always available with +# `std`, so this feature is just a marker that selects the DDP code paths in the example. +ddp = [] + +[dependencies] +# Burn +burn = { workspace = true, features = [ + "train", + "optim", + "tui", + "sqlite-bundled", + "metrics", + "flex", + "autotune", + "fusion", + "std", +] } +log = { workspace = true } + +# Tokenizer +tokenizers = { version = "0.22.2", default-features = false, features = [ + "onig", + "http", +] } + +# Utils +derive-new = { workspace = true } +serde = { workspace = true, features = ["std", "derive"] } diff --git a/examples/text-classification/README.md b/examples/text-classification/README.md new file mode 100644 index 0000000..c4734fe --- /dev/null +++ b/examples/text-classification/README.md @@ -0,0 +1,120 @@ +# Text Classification + +This project provides an example implementation for training and inferencing text classification +models on AG News and DbPedia datasets using the Rust-based Burn Deep Learning Library. + +> **Note** +> This example makes use of the HuggingFace [`datasets`](https://huggingface.co/docs/datasets/index) +> library to download the datasets. Make sure you have [Python](https://www.python.org/downloads/) +> installed on your computer. + +## Dataset Details + +- AG News: The AG News dataset is a collection of news articles from more than 2000 news sources. + This library helps you load and process this dataset, categorizing articles into four classes: + "World", "Sports", "Business", and "Technology". + +- DbPedia: The DbPedia dataset is a large multi-class text classification dataset extracted from + Wikipedia. This library helps you load and process this dataset, categorizing articles into 14 + classes including "Company", "Educational Institution", "Artist", among others. + +# Usage + +## Torch GPU backend + +```bash +git clone https://github.com/tracel-ai/burn.git +cd burn + +# Use the --release flag to really speed up training. +# Use the f16 feature if your CUDA device supports FP16 (half precision) operations. May not work well on every device. + +export TORCH_CUDA_VERSION=cu128 # Set the cuda version (CUDA users) + +# AG News +cargo run --example ag-news-train --release --features tch-gpu # Train on the ag news dataset +cargo run --example ag-news-infer --release --features tch-gpu # Run inference on the ag news dataset + +# DbPedia +cargo run --example db-pedia-train --release --features tch-gpu # Train on the db pedia dataset +cargo run --example db-pedia-infer --release --features tch-gpu # Run inference db pedia dataset +``` + +## Torch CPU backend + +```bash +git clone https://github.com/tracel-ai/burn.git +cd burn + +# Use the --release flag to really speed up training. + +# AG News +cargo run --example ag-news-train --release --features tch-cpu # Train on the ag news dataset +cargo run --example ag-news-infer --release --features tch-cpu # Run inference on the ag news dataset + +# DbPedia +cargo run --example db-pedia-train --release --features tch-cpu # Train on the db pedia dataset +cargo run --example db-pedia-infer --release --features tch-cpu # Run inference db pedia dataset +``` + +## Flex backend + +```bash +git clone https://github.com/tracel-ai/burn.git +cd burn + +# Use the --release flag to really speed up training. + +# AG News +cargo run --example ag-news-train --release --features flex # Train on the ag news dataset +cargo run --example ag-news-infer --release --features flex # Run inference on the ag news dataset + +# DbPedia +cargo run --example db-pedia-train --release --features flex # Train on the db pedia dataset +cargo run --example db-pedia-infer --release --features flex # Run inference db pedia dataset +``` + +## WGPU backend + +```bash +git clone https://github.com/tracel-ai/burn.git +cd burn + +# Use the --release flag to really speed up training. + +# AG News +cargo run --example ag-news-train --release --features wgpu # Train on the ag news dataset +cargo run --example ag-news-infer --release --features wgpu # Run inference on the ag news dataset + +# DbPedia +cargo run --example db-pedia-train --release --features wgpu # Train on the db pedia dataset +cargo run --example db-pedia-infer --release --features wgpu # Run inference db pedia dataset +``` + +## CUDA backend + +```bash +git clone https://github.com/tracel-ai/burn.git +cd burn + +# Use the --release flag to really speed up training. +# Add the f16 feature to run in f16. + +# AG News +cargo run --example ag-news-train --release --features cuda # Train on the ag news dataset +cargo run --example ag-news-infer --release --features cuda # Run inference on the ag news dataset +``` + +## Metal backend + +```bash +git clone https://github.com/tracel-ai/burn.git +cd burn + +# Use the --release flag to really speed up training. +# Add the f16 feature to run in f16. + +# AG News +cargo run --example ag-news-train --release --features metal # Train on the ag news dataset +cargo run --example ag-news-infer --release --features metal # Run inference on the ag news dataset +``` diff --git a/examples/text-classification/burn.toml b/examples/text-classification/burn.toml new file mode 100644 index 0000000..f15d577 --- /dev/null +++ b/examples/text-classification/burn.toml @@ -0,0 +1,22 @@ +[fusion] +logger = { file = "/tmp/fusion.log", level = "full" } + +[cubecl.profiling] +logger = { log = "info", level = "disabled" } + +[cubecl.autotune] +level = "balanced" +cache = "target" +logger = { info = true, level = "full" } + +[cubecl.compilation] +logger = { level = "disabled" } +cache = "target" +check_mode = "auto" + +[cubecl.memory] +logger = { level = "disabled", file = "/tmp/memory.log" } +persistent_memory = "enabled" + +[cubecl.streaming] +max_streams = 8 diff --git a/examples/text-classification/examples/ag-news-infer.rs b/examples/text-classification/examples/ag-news-infer.rs new file mode 100644 index 0000000..d89fbb8 --- /dev/null +++ b/examples/text-classification/examples/ag-news-infer.rs @@ -0,0 +1,109 @@ +#![recursion_limit = "256"] + +use burn::tensor::{Device, DeviceConfig, Element}; +use text_classification::AgNewsDataset; + +#[cfg(not(feature = "f16"))] +#[allow(dead_code)] +type ElemType = f32; +#[cfg(feature = "f16")] +type ElemType = burn::tensor::f16; + +pub fn launch(mut device: Device) { + device + .configure(DeviceConfig::default().float_dtype(ElemType::dtype())) + .unwrap(); + + text_classification::inference::infer::( + device, + "/tmp/text-classification-ag-news", + // Samples from the test dataset, but you are free to test with your own text. + vec![ + "Jays power up to take finale Contrary to popular belief, the power never really \ + snapped back at SkyDome on Sunday. The lights came on after an hour delay, but it \ + took some extra time for the batting orders to provide some extra wattage." + .to_string(), + "Yemen Sentences 15 Militants on Terror Charges A court in Yemen has sentenced one \ + man to death and 14 others to prison terms for a series of attacks and terrorist \ + plots in 2002, including the bombing of a French oil tanker." + .to_string(), + "IBM puts grids to work at U.S. Open IBM will put a collection of its On \ + Demand-related products and technologies to this test next week at the U.S. Open \ + tennis championships, implementing a grid-based infrastructure capable of running \ + multiple workloads including two not associated with the tournament." + .to_string(), + ], + ); +} + +#[cfg(feature = "flex")] +mod flex { + use burn::tensor::Device; + + pub fn run() { + crate::launch(Device::flex()); + } +} + +#[cfg(feature = "tch-gpu")] +mod tch_gpu { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + #[cfg(not(target_os = "macos"))] + let device = Device::libtorch_cuda(DeviceIndex::Default); + #[cfg(target_os = "macos")] + let device = Device::libtorch_mps(); + + crate::launch(device); + } +} + +#[cfg(feature = "tch-cpu")] +mod tch_cpu { + use burn::tensor::Device; + + pub fn run() { + crate::launch(Device::libtorch()); + } +} + +#[cfg(feature = "wgpu")] +mod wgpu { + use burn::tensor::{Device, DeviceKind}; + + pub fn run() { + crate::launch(Device::wgpu(DeviceKind::DefaultDevice)); + } +} + +#[cfg(feature = "metal")] +mod metal { + use burn::tensor::{Device, DeviceKind}; + + pub fn run() { + crate::launch(Device::wgpu(DeviceKind::DefaultDevice)); + } +} + +#[cfg(feature = "cuda")] +mod cuda { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + crate::launch(Device::cuda(DeviceIndex::Default)); + } +} + +fn main() { + #[cfg(feature = "flex")] + flex::run(); + #[cfg(feature = "tch-gpu")] + tch_gpu::run(); + #[cfg(feature = "tch-cpu")] + tch_cpu::run(); + #[cfg(feature = "wgpu")] + wgpu::run(); + #[cfg(feature = "cuda")] + cuda::run(); +} diff --git a/examples/text-classification/examples/ag-news-train.rs b/examples/text-classification/examples/ag-news-train.rs new file mode 100644 index 0000000..4776bf3 --- /dev/null +++ b/examples/text-classification/examples/ag-news-train.rs @@ -0,0 +1,192 @@ +#![recursion_limit = "256"] + +// Only the CUDA `launch_multi` path uses these at the top level; the remote module imports +// them locally. Gating on both features avoids an unused-import warning for `remote,ddp`. +#[cfg(all(feature = "ddp", feature = "cuda"))] +use burn::tensor::distributed::{DistributedConfig, ReduceOperation}; +use burn::{ + nn::transformer::TransformerEncoderConfig, + optim::{AdamConfig, decay::WeightDecayConfig}, + tensor::{Device, DeviceConfig, Element}, + train::ExecutionStrategy, +}; + +use text_classification::{AgNewsDataset, training::ExperimentConfig}; + +#[cfg(not(any(feature = "f16", feature = "flex32")))] +#[allow(unused)] +type ElemType = f32; +#[cfg(feature = "f16")] +type ElemType = burn::tensor::f16; +#[cfg(feature = "flex32")] +type ElemType = burn::tensor::flex32; + +#[cfg(all(feature = "cuda", not(feature = "ddp")))] +pub fn launch_multi() { + let mut devices = Device::enumerate(burn::tensor::DeviceType::Cuda); + devices + .configure(DeviceConfig::default().float_dtype(ElemType::dtype())) + .unwrap(); + + launch(ExecutionStrategy::MultiDevice( + devices.into_vec(), + burn::train::MultiDeviceOptim::OptimSharded, + )) +} + +#[cfg(all(feature = "cuda", feature = "ddp"))] +pub fn launch_multi() { + let mut devices = Device::enumerate(burn::tensor::DeviceType::Cuda); + devices + .configure(DeviceConfig::default().float_dtype(ElemType::dtype())) + .unwrap(); + + launch(ExecutionStrategy::ddp( + devices.into_vec(), + DistributedConfig { + all_reduce_op: ReduceOperation::Mean, + }, + )) +} + +pub fn launch_single(mut device: Device) { + device + .configure(DeviceConfig::default().float_dtype(ElemType::dtype())) + .unwrap(); + + launch(ExecutionStrategy::SingleDevice(device)) +} + +pub fn launch(strategy: ExecutionStrategy) { + let config = ExperimentConfig::new( + TransformerEncoderConfig::new(128, 512, 4, 4) + .with_norm_first(true) + .with_quiet_softmax(true), + AdamConfig::new().with_weight_decay(Some(WeightDecayConfig::new(5e-5))), + ); + + text_classification::training::train::( + strategy, + AgNewsDataset::train(), + AgNewsDataset::test(), + config, + "/tmp/text-classification-ag-news", + ); +} + +#[cfg(feature = "tch-gpu")] +mod tch_gpu { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + #[cfg(not(target_os = "macos"))] + let device = Device::libtorch_cuda(DeviceIndex::Default); + #[cfg(target_os = "macos")] + let device = Device::libtorch_mps(); + + crate::launch_single(device); + } +} + +#[cfg(feature = "tch-cpu")] +mod tch_cpu { + use burn::tensor::Device; + + pub fn run() { + crate::launch_single(Device::libtorch()); + } +} + +#[cfg(any(feature = "wgpu", feature = "vulkan", feature = "metal"))] +mod wgpu { + use burn::tensor::{Device, DeviceKind}; + + pub fn run() { + crate::launch_single(Device::wgpu(DeviceKind::DefaultDevice)); + } +} + +#[cfg(feature = "remote")] +mod remote { + #[cfg(feature = "ddp")] + use crate::ElemType; + #[cfg(feature = "ddp")] + use burn::tensor::distributed::{DistributedConfig, ReduceOperation}; + use burn::tensor::{Device, DeviceType}; + #[cfg(feature = "ddp")] + use burn::tensor::{DeviceConfig, Element}; + #[cfg(feature = "ddp")] + use burn::train::ExecutionStrategy; + + /// Address of the `burn-remote` server to train against. + const ADDRESS: &str = "ws://localhost:3000"; + + /// Train on a single one of the devices the remote server hosts. + /// + /// `launch_single` configures the device it receives, so don't configure the enumerated + /// set here too — doing both locks the device's settings twice and returns + /// [`DeviceError::AlreadyInitialized`](burn::tensor::DeviceError::AlreadyInitialized). + #[cfg(not(feature = "ddp"))] + pub fn run() { + let devices = Device::enumerate(DeviceType::remote_websocket(ADDRESS)); + crate::launch_single(devices.into_vec().pop().unwrap()); + } + + /// Same enumeration, but drive the devices with distributed data-parallel training. + #[cfg(feature = "ddp")] + pub fn run() { + let mut devices = Device::enumerate(DeviceType::remote_websocket(ADDRESS)); + devices + .configure(DeviceConfig::default().float_dtype(ElemType::dtype())) + .unwrap(); + + crate::launch(ExecutionStrategy::ddp( + devices.into_vec(), + DistributedConfig { + all_reduce_op: ReduceOperation::Mean, + }, + )); + } +} + +#[cfg(feature = "cuda")] +mod cuda { + pub fn run() { + crate::launch_multi(); + } +} + +#[cfg(feature = "rocm")] +mod rocm { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + crate::launch_single(Device::rocm(DeviceIndex::Default)); + } +} + +#[cfg(feature = "flex")] +mod flex { + use burn::tensor::Device; + + pub fn run() { + crate::launch_single(Device::flex()); + } +} + +fn main() { + #[cfg(feature = "tch-gpu")] + tch_gpu::run(); + #[cfg(feature = "tch-cpu")] + tch_cpu::run(); + #[cfg(any(feature = "wgpu", feature = "vulkan", feature = "metal"))] + wgpu::run(); + #[cfg(feature = "cuda")] + cuda::run(); + #[cfg(feature = "rocm")] + rocm::run(); + #[cfg(feature = "remote")] + remote::run(); + #[cfg(feature = "flex")] + flex::run(); +} diff --git a/examples/text-classification/examples/db-pedia-infer.rs b/examples/text-classification/examples/db-pedia-infer.rs new file mode 100644 index 0000000..45e0ed1 --- /dev/null +++ b/examples/text-classification/examples/db-pedia-infer.rs @@ -0,0 +1,104 @@ +use text_classification::DbPediaDataset; + +use burn::tensor::{Device, DeviceConfig, Element}; + +#[cfg(not(feature = "f16"))] +#[allow(dead_code)] +type ElemType = f32; +#[cfg(feature = "f16")] +type ElemType = burn::tensor::f16; + +pub fn launch(mut device: Device) { + device + .configure(DeviceConfig::default().float_dtype(ElemType::dtype())) + .unwrap(); + + text_classification::inference::infer::( + device, + "/tmp/text-classification-db-pedia", + // Samples from the test dataset, but you are free to test with your own text. + vec![ + " Magnus Eriksson is a Swedish former footballer who played as a forward.".to_string(), + "Crossbeam Systems is headquartered in Boxborough Massachusetts and has offices in \ + Europe Latin America and Asia Pacific. Crossbeam Systems was acquired by Blue Coat \ + Systems in December 2012 and the Crossbeam brand has been fully absorbed into Blue \ + Coat." + .to_string(), + " Zia is the sequel to the award-winning Island of the Blue Dolphins by Scott O'Dell. \ + It was published in 1976 sixteen years after the publication of the first novel." + .to_string(), + ], + ); +} + +#[cfg(feature = "flex")] +mod flex { + use burn::tensor::Device; + + pub fn run() { + crate::launch(Device::flex()); + } +} + +#[cfg(feature = "tch-gpu")] +mod tch_gpu { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + #[cfg(not(target_os = "macos"))] + let device = Device::libtorch_cuda(DeviceIndex::Default); + #[cfg(target_os = "macos")] + let device = Device::libtorch_mps(); + + crate::launch(device); + } +} + +#[cfg(feature = "tch-cpu")] +mod tch_cpu { + use burn::tensor::Device; + + pub fn run() { + crate::launch(Device::libtorch()); + } +} + +#[cfg(feature = "wgpu")] +mod wgpu { + use burn::tensor::{Device, DeviceKind}; + + pub fn run() { + crate::launch(Device::wgpu(DeviceKind::DefaultDevice)); + } +} + +#[cfg(feature = "metal")] +mod metal { + use burn::tensor::{Device, DeviceKind}; + + pub fn run() { + crate::launch(Device::wgpu(DeviceKind::DefaultDevice)); + } +} + +#[cfg(feature = "cuda")] +mod cuda { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + crate::launch(Device::cuda(DeviceIndex::Default)); + } +} + +fn main() { + #[cfg(feature = "flex")] + flex::run(); + #[cfg(feature = "tch-gpu")] + tch_gpu::run(); + #[cfg(feature = "tch-cpu")] + tch_cpu::run(); + #[cfg(feature = "wgpu")] + wgpu::run(); + #[cfg(feature = "cuda")] + cuda::run(); +} diff --git a/examples/text-classification/examples/db-pedia-train.rs b/examples/text-classification/examples/db-pedia-train.rs new file mode 100644 index 0000000..4df5a1e --- /dev/null +++ b/examples/text-classification/examples/db-pedia-train.rs @@ -0,0 +1,179 @@ +use burn::{ + nn::transformer::TransformerEncoderConfig, + optim::{AdamConfig, decay::WeightDecayConfig}, + tensor::{Device, DeviceConfig, Element}, + train::ExecutionStrategy, +}; + +use text_classification::{DbPediaDataset, training::ExperimentConfig}; + +#[cfg(not(feature = "f16"))] +#[allow(dead_code)] +type ElemType = f32; +#[cfg(feature = "f16")] +type ElemType = burn::tensor::f16; + +pub fn launch(strategy: ExecutionStrategy) { + let config = ExperimentConfig::new( + TransformerEncoderConfig::new(256, 1024, 8, 4).with_norm_first(true), + AdamConfig::new().with_weight_decay(Some(WeightDecayConfig::new(5e-5))), + ); + + text_classification::training::train::( + strategy, + DbPediaDataset::train(), + DbPediaDataset::test(), + config, + "/tmp/text-classification-db-pedia", + ); +} + +pub fn launch_single(mut device: Device) { + device + .configure(DeviceConfig::default().float_dtype(ElemType::dtype())) + .unwrap(); + + launch(ExecutionStrategy::SingleDevice(device)) +} + +#[cfg(all(feature = "cuda", not(feature = "ddp")))] +pub fn launch_multi() { + let mut devices = Device::enumerate(burn::tensor::DeviceType::Cuda); + devices + .configure(DeviceConfig::default().float_dtype(ElemType::dtype())) + .unwrap(); + + launch(ExecutionStrategy::MultiDevice( + devices.into_vec(), + burn::train::MultiDeviceOptim::OptimSharded, + )) +} + +#[cfg(all(feature = "cuda", feature = "ddp"))] +pub fn launch_multi() { + let mut devices = Device::enumerate(burn::tensor::DeviceType::Cuda); + devices + .configure(DeviceConfig::default().float_dtype(ElemType::dtype())) + .unwrap(); + + launch(ExecutionStrategy::ddp( + devices.into_vec(), + DistributedConfig { + all_reduce_op: ReduceOperation::Mean, + }, + )) +} + +#[cfg(feature = "flex")] +mod flex { + use burn::tensor::Device; + + pub fn run() { + crate::launch_single(Device::flex()); + } +} + +#[cfg(feature = "tch-gpu")] +mod tch_gpu { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + #[cfg(not(target_os = "macos"))] + let device = Device::libtorch_cuda(DeviceIndex::Default); + #[cfg(target_os = "macos")] + let device = Device::libtorch_mps(); + + crate::launch_single(device); + } +} + +#[cfg(feature = "tch-cpu")] +mod tch_cpu { + use burn::tensor::Device; + + pub fn run() { + crate::launch_single(Device::libtorch()); + } +} + +#[cfg(any(feature = "wgpu", feature = "vulkan", feature = "metal"))] +mod wgpu { + use burn::tensor::{Device, DeviceKind}; + + pub fn run() { + crate::launch_single(Device::wgpu(DeviceKind::DefaultDevice)); + } +} + +#[cfg(feature = "remote")] +mod remote { + use crate::ElemType; + #[cfg(feature = "ddp")] + use burn::tensor::distributed::{DistributedConfig, ReduceOperation}; + use burn::tensor::{Device, DeviceConfig, DeviceType, Element}; + #[cfg(feature = "ddp")] + use burn::train::ExecutionStrategy; + + /// Address of the `burn-remote` server to train against. + const ADDRESS: &str = "ws://localhost:3000"; + + /// List every device the remote server hosts and train across all of them. + #[cfg(not(feature = "ddp"))] + pub fn run() { + let mut devices = Device::enumerate(DeviceType::remote(ADDRESS)); + devices + .configure(DeviceConfig::default().float_dtype(ElemType::dtype())) + .unwrap(); + + crate::launch_single(devices.into_vec().pop().unwrap()); + } + + /// Same enumeration, but drive the devices with distributed data-parallel training. + #[cfg(feature = "ddp")] + pub fn run() { + let mut devices = Device::enumerate(DeviceType::remote(ADDRESS)); + devices + .configure(DeviceConfig::default().float_dtype(ElemType::dtype())) + .unwrap(); + + crate::launch_single(ExecutionStrategy::ddp( + devices.into_vec(), + DistributedConfig { + all_reduce_op: ReduceOperation::Mean, + }, + )); + } +} + +#[cfg(feature = "cuda")] +mod cuda { + pub fn run() { + crate::launch_multi(); + } +} + +#[cfg(feature = "rocm")] +mod rocm { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + crate::launch_single(Device::rocm(DeviceIndex::Default)); + } +} + +fn main() { + #[cfg(feature = "flex")] + flex::run(); + #[cfg(feature = "tch-gpu")] + tch_gpu::run(); + #[cfg(feature = "tch-cpu")] + tch_cpu::run(); + #[cfg(any(feature = "wgpu", feature = "vulkan", feature = "metal"))] + wgpu::run(); + #[cfg(feature = "cuda")] + cuda::run(); + #[cfg(feature = "rocm")] + rocm::run(); + #[cfg(feature = "remote")] + remote::run(); +} diff --git a/examples/text-classification/src/data/batcher.rs b/examples/text-classification/src/data/batcher.rs new file mode 100644 index 0000000..34c1db3 --- /dev/null +++ b/examples/text-classification/src/data/batcher.rs @@ -0,0 +1,107 @@ +// The module defines two structs TextClassificationTrainingBatch and TextClassificationInferenceBatch +// to handle batches of data during training and inference respectively. The TextClassificationBatcher +// struct is implemented for creating these batches. It is parameterized on the type B: Backend to +// support different computation backends (e.g., CPU, CUDA). + +// Two implementations of the Batcher trait are provided for TextClassificationBatcher, one for creating +// training batches and one for creating inference batches. In each implementation, the batch function is +// defined to convert a vector of items into a batch. For training, the items are instances of +// TextClassificationItem and include both the text and the corresponding label. +// For inference, the items are simply strings without labels. The function tokenizes the text, +// generates a padding mask, and returns a batch object. + +use super::{dataset::TextClassificationItem, tokenizer::Tokenizer}; +use burn::{ + data::dataloader::batcher::Batcher, + nn::attention::{SeqLengthOption, generate_padding_mask}, + prelude::*, +}; +use std::sync::Arc; + +/// Struct for batching text classification items +#[derive(Clone, new)] +pub struct TextClassificationBatcher { + tokenizer: Arc, // Tokenizer for converting text to token IDs + seq_length: SeqLengthOption, // Sequence length option for tokenized text +} + +/// Struct for training batch in text classification task +#[derive(Debug, Clone, new)] +pub struct TextClassificationTrainingBatch { + pub tokens: Tensor<2, Int>, // Tokenized text + pub labels: Tensor<1, Int>, // Labels of the text + pub mask_pad: Tensor<2, Bool>, // Padding mask for the tokenized text +} + +/// Struct for inference batch in text classification task +#[derive(Debug, Clone, new)] +pub struct TextClassificationInferenceBatch { + pub tokens: Tensor<2, Int>, // Tokenized text + pub mask_pad: Tensor<2, Bool>, // Padding mask for the tokenized text +} + +/// Implement Batcher trait for TextClassificationBatcher struct for training +impl Batcher + for TextClassificationBatcher +{ + /// Batches a vector of text classification items into a training batch + fn batch( + &self, + items: Vec, + device: &Device, + ) -> TextClassificationTrainingBatch { + let mut tokens_list = Vec::with_capacity(items.len()); + let mut labels_list = Vec::with_capacity(items.len()); + + // Tokenize text and create label tensor for each item + for item in items { + tokens_list.push(self.tokenizer.encode(&item.text)); + labels_list.push(Tensor::from_data( + TensorData::from([item.label as i64]), + device, + )); + } + + // Generate padding mask for tokenized text + let mask = generate_padding_mask( + self.tokenizer.pad_token(), + tokens_list, + self.seq_length, + device, + ); + + // Create and return training batch + TextClassificationTrainingBatch { + tokens: mask.tensor, + labels: Tensor::cat(labels_list, 0), + mask_pad: mask.mask, + } + } +} + +/// Implement Batcher trait for TextClassificationBatcher struct for inference +impl Batcher for TextClassificationBatcher { + /// Batches a vector of strings into an inference batch + fn batch(&self, items: Vec, device: &Device) -> TextClassificationInferenceBatch { + let mut tokens_list = Vec::with_capacity(items.len()); + + // Tokenize each string + for item in items { + tokens_list.push(self.tokenizer.encode(&item)); + } + + // Generate padding mask for tokenized text + let mask = generate_padding_mask( + self.tokenizer.pad_token(), + tokens_list, + self.seq_length, + device, + ); + + // Create and return inference batch + TextClassificationInferenceBatch { + tokens: mask.tensor.to_device(device), + mask_pad: mask.mask.to_device(device), + } + } +} diff --git a/examples/text-classification/src/data/dataset.rs b/examples/text-classification/src/data/dataset.rs new file mode 100644 index 0000000..9701408 --- /dev/null +++ b/examples/text-classification/src/data/dataset.rs @@ -0,0 +1,172 @@ +// The AgNewsDataset and DbPediaDataset structs are examples of specific text +// classification datasets. Each dataset struct has a field for the underlying +// SQLite dataset and implements methods for accessing and processing the data. +// Each dataset is also provided with specific information about its classes via +// the TextClassificationDataset trait. These implementations are designed to be used +// with a machine learning framework for tasks such as training a text classification model. + +use burn::data::dataset::{Dataset, SqliteDataset, source::huggingface::HuggingfaceDatasetLoader}; + +// Define a struct for text classification items +#[derive(new, Clone, Debug)] +pub struct TextClassificationItem { + pub text: String, // The text for classification + pub label: usize, // The label of the text (classification category) +} + +// Trait for text classification datasets +pub trait TextClassificationDataset: Dataset { + fn num_classes() -> usize; // Returns the number of unique classes in the dataset + fn class_name(label: usize) -> String; // Returns the name of the class given its label +} + +// Struct for items in the AG News dataset +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +pub struct AgNewsItem { + pub text: String, // The text for classification + pub label: usize, // The label of the text (classification category) +} + +// Struct for the AG News dataset +pub struct AgNewsDataset { + dataset: SqliteDataset, // Underlying SQLite dataset +} + +// Implement the Dataset trait for the AG News dataset +impl Dataset for AgNewsDataset { + /// Returns a specific item from the dataset + fn get(&self, index: usize) -> Option { + self.dataset + .get(index) + .map(|item| TextClassificationItem::new(item.text, item.label)) // Map AgNewsItems to TextClassificationItems + } + + /// Returns the length of the dataset + fn len(&self) -> usize { + self.dataset.len() + } +} + +// Implement methods for constructing the AG News dataset +impl AgNewsDataset { + /// Returns the training portion of the dataset + pub fn train() -> Self { + Self::new("train") + } + + /// Returns the testing portion of the dataset + pub fn test() -> Self { + Self::new("test") + } + + /// Constructs the dataset from a split (either "train" or "test") + pub fn new(split: &str) -> Self { + let dataset: SqliteDataset = HuggingfaceDatasetLoader::new("fancyzhx/ag_news") + .dataset(split) + .unwrap(); + Self { dataset } + } +} + +/// Implements the TextClassificationDataset trait for the AG News dataset +impl TextClassificationDataset for AgNewsDataset { + /// Returns the number of unique classes in the dataset + fn num_classes() -> usize { + 4 + } + + /// Returns the name of a class given its label + fn class_name(label: usize) -> String { + match label { + 0 => "World", + 1 => "Sports", + 2 => "Business", + 3 => "Technology", + _ => panic!("invalid class"), + } + .to_string() + } +} + +/// Struct for items in the DbPedia dataset +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +pub struct DbPediaItem { + pub title: String, // The title of the item + pub content: String, // The content of the item + pub label: usize, // The label of the item (classification category) +} + +/// Struct for the DbPedia dataset +pub struct DbPediaDataset { + dataset: SqliteDataset, // Underlying SQLite dataset +} + +/// Implements the Dataset trait for the DbPedia dataset +impl Dataset for DbPediaDataset { + /// Returns a specific item from the dataset + fn get(&self, index: usize) -> Option { + self.dataset.get(index).map(|item| { + TextClassificationItem::new( + format!("Title: {} - Content: {}", item.title, item.content), + item.label, + ) + }) + } + + /// Returns the length of the dataset + fn len(&self) -> usize { + self.dataset.len() + } +} + +/// Implement methods for constructing the DbPedia dataset +impl DbPediaDataset { + /// Returns the training portion of the dataset + pub fn train() -> Self { + Self::new("train") + } + + /// Returns the testing portion of the dataset + pub fn test() -> Self { + Self::new("test") + } + + /// Constructs the dataset from a split (either "train" or "test") + pub fn new(split: &str) -> Self { + let dataset: SqliteDataset = + HuggingfaceDatasetLoader::new("fancyzhx/dbpedia_14") + .dataset(split) + .unwrap(); + Self { dataset } + } +} + +/// Implement the TextClassificationDataset trait for the DbPedia dataset +impl TextClassificationDataset for DbPediaDataset { + /// Returns the number of unique classes in the dataset + fn num_classes() -> usize { + 14 + } + + /// Returns the name of a class given its label + fn class_name(label: usize) -> String { + match label { + 0 => "Company", + 1 => "EducationalInstitution", + 2 => "Artist", + 3 => "Athlete", + 4 => "OfficeHolder", + 5 => "MeanOfTransportation", + 6 => "Building", + 7 => "NaturalPlace", + 8 => "Village", + 9 => "Animal", + 10 => "Plant", + 11 => "Album", + 12 => "Film", + 13 => "WrittenWork", + _ => panic!("invalid class"), + } + .to_string() + } +} diff --git a/examples/text-classification/src/data/mod.rs b/examples/text-classification/src/data/mod.rs new file mode 100644 index 0000000..8f6faeb --- /dev/null +++ b/examples/text-classification/src/data/mod.rs @@ -0,0 +1,7 @@ +mod batcher; +mod dataset; +mod tokenizer; + +pub use batcher::*; +pub use dataset::*; +pub use tokenizer::*; diff --git a/examples/text-classification/src/data/tokenizer.rs b/examples/text-classification/src/data/tokenizer.rs new file mode 100644 index 0000000..4b1c8ad --- /dev/null +++ b/examples/text-classification/src/data/tokenizer.rs @@ -0,0 +1,68 @@ +// This module defines a trait `Tokenizer` that represents a common interface for all tokenizer +// types used in the text classification library. A specific implementation of this trait, +// `BertCasedTokenizer`, uses the BERT cased tokenization strategy provided by the `tokenizers` library. + +// This trait represents the common interface for all tokenizer types. +// The `Send + Sync` bounds are necessary for allowing these operations +// to work across thread boundaries. +#[allow(dead_code)] +pub trait Tokenizer: Send + Sync { + /// Converts a text string into a sequence of tokens. + fn encode(&self, value: &str) -> Vec; + + /// Converts a sequence of tokens back into a text string. + fn decode(&self, tokens: &[usize]) -> String; + + /// Gets the size of the tokenizer's vocabulary. + fn vocab_size(&self) -> usize; + + /// Gets the token used for padding sequences to a consistent length. + fn pad_token(&self) -> usize; + + /// Gets the string representation of the padding token. + /// The default implementation uses `decode` on the padding token. + fn pad_token_value(&self) -> String { + self.decode(&[self.pad_token()]) + } +} + +/// Struct represents a specific tokenizer using the BERT cased tokenization strategy. +pub struct BertCasedTokenizer { + // The underlying tokenizer from the `tokenizers` library. + tokenizer: tokenizers::Tokenizer, +} + +// Default implementation for creating a new BertCasedTokenizer. +// This uses a pretrained BERT cased tokenizer model. +impl Default for BertCasedTokenizer { + fn default() -> Self { + Self { + tokenizer: tokenizers::Tokenizer::from_pretrained("bert-base-cased", None).unwrap(), + } + } +} + +// Implementation of the Tokenizer trait for BertCasedTokenizer. +impl Tokenizer for BertCasedTokenizer { + // Convert a text string into a sequence of tokens using the BERT cased tokenization strategy. + fn encode(&self, value: &str) -> Vec { + let tokens = self.tokenizer.encode(value, true).unwrap(); + tokens.get_ids().iter().map(|t| *t as usize).collect() + } + + /// Converts a sequence of tokens back into a text string. + fn decode(&self, tokens: &[usize]) -> String { + let tokens = tokens.iter().map(|t| *t as u32).collect::>(); + self.tokenizer.decode(&tokens, false).unwrap() + } + + /// Gets the size of the BERT cased tokenizer's vocabulary. + fn vocab_size(&self) -> usize { + self.tokenizer.get_vocab_size(true) + } + + /// Gets the token used for padding sequences to a consistent length. + fn pad_token(&self) -> usize { + self.tokenizer.token_to_id("[PAD]").unwrap() as usize + } +} diff --git a/examples/text-classification/src/inference.rs b/examples/text-classification/src/inference.rs new file mode 100644 index 0000000..a8beeb5 --- /dev/null +++ b/examples/text-classification/src/inference.rs @@ -0,0 +1,72 @@ +// This module defines the inference process for a text classification model. +// It loads a model and its configuration from a directory, and uses a tokenizer +// and a batcher to prepare the input data. The model is then used to make predictions +// on the input samples, and the results are printed out for each sample. +// Import required modules and types + +use crate::{ + data::{BertCasedTokenizer, TextClassificationBatcher, TextClassificationDataset, Tokenizer}, + model::TextClassificationModelConfig, + training::ExperimentConfig, +}; +use burn::{data::dataloader::batcher::Batcher, prelude::*, store::ModuleRecord}; +use std::sync::Arc; + +// Define inference function +pub fn infer( + device: Device, // Device on which to perform computation (e.g., CPU or CUDA device) + artifact_dir: &str, // Directory containing model and config files + samples: Vec, // Text samples for inference +) { + // Load experiment configuration + let config = ExperimentConfig::load(format!("{artifact_dir}/config.json").as_str()) + .expect("Config file present"); + + // Initialize tokenizer + let tokenizer = Arc::new(BertCasedTokenizer::default()); + + // Get number of classes from dataset + let n_classes = D::num_classes(); + + // Initialize batcher for batching samples + let batcher = Arc::new(TextClassificationBatcher::new( + tokenizer.clone(), + config.seq_length, + )); + + // Load pre-trained model weights + println!("Loading weights ..."); + let record = + ModuleRecord::load(format!("{artifact_dir}/model")).expect("Trained model weights tb"); + + // Create model using loaded weights + println!("Creating model ..."); + let model = TextClassificationModelConfig::new( + config.transformer, + n_classes, + tokenizer.vocab_size(), + config.seq_length, + ) + .init(&device) + .load_record(record); // Initialize model with loaded weights + + // Run inference on the given text samples + println!("Running inference ..."); + let item = batcher.batch(samples.clone(), &device); // Batch samples using the batcher + let predictions = model.infer(item); // Get model predictions + + // Print out predictions for each sample + for (i, text) in samples.into_iter().enumerate() { + #[allow(clippy::single_range_in_vec_init)] + let prediction = predictions.clone().slice([i..i + 1]); // Get prediction for current sample + let logits = prediction.to_data(); // Convert prediction tensor to data + let class_index: i32 = prediction.argmax(1).squeeze_dim::<1>(1).into_scalar(); // Get class index with the highest value + let class = D::class_name(class_index as usize); // Get class name + + // Print sample text, predicted logits and predicted class + println!( + "\n=== Item {i} ===\n- Text: {text}\n- Logits: {logits}\n- Prediction: \ + {class}\n================" + ); + } +} diff --git a/examples/text-classification/src/lib.rs b/examples/text-classification/src/lib.rs new file mode 100644 index 0000000..66e1e4f --- /dev/null +++ b/examples/text-classification/src/lib.rs @@ -0,0 +1,9 @@ +#[macro_use] +extern crate derive_new; + +pub mod data; +pub mod inference; +pub mod model; +pub mod training; + +pub use data::{AgNewsDataset, DbPediaDataset, TextClassificationDataset}; diff --git a/examples/text-classification/src/model.rs b/examples/text-classification/src/model.rs new file mode 100644 index 0000000..8a151a8 --- /dev/null +++ b/examples/text-classification/src/model.rs @@ -0,0 +1,162 @@ +// This is a basic text classification model implemented in Rust using the Burn framework. +// It uses a Transformer as the base model and applies Linear and Embedding layers. +// The model is then trained using Cross-Entropy loss. It contains methods for model initialization +// (both with and without pre-trained weights), forward pass, inference, training, and validation. + +use crate::data::{TextClassificationInferenceBatch, TextClassificationTrainingBatch}; +use burn::{ + nn::{ + Embedding, EmbeddingConfig, Linear, LinearConfig, + attention::SeqLengthOption, + loss::CrossEntropyLossConfig, + transformer::{TransformerEncoder, TransformerEncoderConfig, TransformerEncoderInput}, + }, + prelude::*, + tensor::activation::softmax, + train::{ClassificationOutput, InferenceStep, TrainOutput, TrainStep}, +}; + +// Define the model configuration +#[derive(Config, Debug)] +pub struct TextClassificationModelConfig { + transformer: TransformerEncoderConfig, + n_classes: usize, + vocab_size: usize, + seq_length: SeqLengthOption, +} + +// Define the model structure +#[derive(Module, Debug)] +pub struct TextClassificationModel { + transformer: TransformerEncoder, + embedding_token: Embedding, + embedding_pos: Embedding, + output: Linear, + n_classes: usize, +} + +// Define functions for model initialization +impl TextClassificationModelConfig { + /// Initializes a model with default weights + pub fn init(&self, device: &Device) -> TextClassificationModel { + let output = LinearConfig::new(self.transformer.d_model, self.n_classes).init(device); + let transformer = self.transformer.init(device); + let embedding_token = + EmbeddingConfig::new(self.vocab_size, self.transformer.d_model).init(device); + let max_seq_length = match self.seq_length { + SeqLengthOption::Fixed(max) | SeqLengthOption::Max(max) => max, + SeqLengthOption::NoMax => panic!( + "Text classification requires a max sequence length because of the embedding strategy." + ), + }; + let embedding_pos = + EmbeddingConfig::new(max_seq_length, self.transformer.d_model).init(device); + + TextClassificationModel { + transformer, + embedding_token, + embedding_pos, + output, + n_classes: self.n_classes, + } + } +} + +/// Define model behavior +impl TextClassificationModel { + // Defines forward pass for training + pub fn forward(&self, item: TextClassificationTrainingBatch) -> ClassificationOutput { + // Get batch and sequence length, and the device + let [batch_size, seq_length] = item.tokens.dims(); + let device = &self.embedding_token.devices()[0]; + + // Move tensors to the correct device + let tokens = item.tokens.to_device(device); + let labels = item.labels.to_device(device); + let mask_pad = item.mask_pad.to_device(device); + + // Calculate token and position embeddings, and combine them + let index_positions = Tensor::arange(0..seq_length as i64, device) + .reshape([1, seq_length]) + .repeat_dim(0, batch_size); + let embedding_positions = self.embedding_pos.forward(index_positions); + let embedding_tokens = self.embedding_token.forward(tokens); + let embedding = (embedding_positions + embedding_tokens) / 2; + + // Perform transformer encoding, calculate output and loss + let encoded = self + .transformer + .forward(TransformerEncoderInput::new(embedding).mask_pad(mask_pad)); + let output = self.output.forward(encoded); + + let output_classification = output + .slice([0..batch_size, 0..1]) + .reshape([batch_size, self.n_classes]); + + let loss = CrossEntropyLossConfig::new() + .init(&output_classification.device()) + .forward(output_classification.clone(), labels.clone()); + + // Return the output and loss + ClassificationOutput { + loss, + output: output_classification, + targets: labels, + } + } + + /// Defines forward pass for inference + pub fn infer(&self, item: TextClassificationInferenceBatch) -> Tensor<2> { + // Get batch and sequence length, and the device + let [batch_size, seq_length] = item.tokens.dims(); + let device = &self.embedding_token.devices()[0]; + + // Move tensors to the correct device + let tokens = item.tokens.to_device(device); + let mask_pad = item.mask_pad.to_device(device); + + // Calculate token and position embeddings, and combine them + let index_positions = Tensor::arange(0..seq_length as i64, device) + .reshape([1, seq_length]) + .repeat_dim(0, batch_size); + let embedding_positions = self.embedding_pos.forward(index_positions); + let embedding_tokens = self.embedding_token.forward(tokens); + let embedding = (embedding_positions + embedding_tokens) / 2; + + // Perform transformer encoding, calculate output and apply softmax for prediction + let encoded = self + .transformer + .forward(TransformerEncoderInput::new(embedding).mask_pad(mask_pad)); + let output = self.output.forward(encoded); + let output = output + .slice([0..batch_size, 0..1]) + .reshape([batch_size, self.n_classes]); + + softmax(output, 1) + } +} + +/// Define training step +impl TrainStep for TextClassificationModel { + type Input = TextClassificationTrainingBatch; + type Output = ClassificationOutput; + + fn step(&self, item: TextClassificationTrainingBatch) -> TrainOutput { + // Run forward pass, calculate gradients and return them along with the output + let item = self.forward(item); + let grads = item.loss.backward(); + + TrainOutput::new(self, grads, item) + } +} + +/// Define validation step +impl InferenceStep for TextClassificationModel { + type Input = TextClassificationTrainingBatch; + type Output = ClassificationOutput; + + fn step(&self, item: TextClassificationTrainingBatch) -> ClassificationOutput { + // Run forward pass and return the output + self.forward(item) + } +} diff --git a/examples/text-classification/src/training.rs b/examples/text-classification/src/training.rs new file mode 100644 index 0000000..87d9a89 --- /dev/null +++ b/examples/text-classification/src/training.rs @@ -0,0 +1,114 @@ +// This module trains a text classification model using the provided training and testing datasets, +// as well as the provided configuration. It first initializes a tokenizer and batchers for the datasets, +// then initializes the model and data loaders for the datasets. The function then initializes +// an optimizer and a learning rate scheduler, and uses them along with the model and datasets +// to build a learner, which is used to train the model. The trained model and the configuration are +// then saved to the specified directory. + +use crate::{ + data::{BertCasedTokenizer, TextClassificationBatcher, TextClassificationDataset, Tokenizer}, + model::TextClassificationModelConfig, +}; + +use burn::train::{ExecutionStrategy, Learner, SupervisedTraining}; +use burn::{ + data::{dataloader::DataLoaderBuilder, dataset::transform::SamplerDataset}, + lr_scheduler::noam::NoamLrSchedulerConfig, + nn::{attention::SeqLengthOption, transformer::TransformerEncoderConfig}, + optim::AdamConfig, + prelude::*, + train::metric::{ + AccuracyMetric, CudaMetric, IterationSpeedMetric, LearningRateMetric, LossMetric, + }, +}; +use std::{path::PathBuf, sync::Arc}; + +// Define configuration struct for the experiment +#[derive(Config, Debug)] +pub struct ExperimentConfig { + pub transformer: TransformerEncoderConfig, + pub optimizer: AdamConfig, + #[config(default = "SeqLengthOption::Fixed(256)")] + pub seq_length: SeqLengthOption, + #[config(default = 32)] + pub batch_size: usize, + #[config(default = 5)] + pub num_epochs: usize, +} + +fn create_artifact_dir(artifact_dir: &str) { + std::fs::remove_file(PathBuf::from(artifact_dir).join("experiment.log")).ok(); + std::fs::create_dir_all(artifact_dir).ok(); +} + +// Define train function +pub fn train( + strategy: ExecutionStrategy, + dataset_train: D, // Training dataset + dataset_test: D, // Testing dataset + config: ExperimentConfig, // Experiment configuration + artifact_dir: &str, // Directory to save model and config files +) { + create_artifact_dir(artifact_dir); + + // Initialize tokenizer + let tokenizer = Arc::new(BertCasedTokenizer::default()); + + // Initialize batcher + let batcher = TextClassificationBatcher::new(tokenizer.clone(), config.seq_length); + + // Initialize model + let model = TextClassificationModelConfig::new( + config.transformer.clone(), + D::num_classes(), + tokenizer.vocab_size(), + config.seq_length, + ) + .init(&strategy.main_device().clone().autodiff()); + + // Initialize data loaders for training and testing data + let dataloader_train = DataLoaderBuilder::new(batcher.clone()) + .batch_size(config.batch_size) + .num_workers(1) + .build(SamplerDataset::new(dataset_train, 25_000)); + let dataloader_test = DataLoaderBuilder::new(batcher) + .batch_size(config.batch_size) + .num_workers(1) + .build(SamplerDataset::new(dataset_test, 2500)); + + // Initialize optimizer + let optim = config.optimizer.init(); + + // Initialize learning rate scheduler + let lr_scheduler = NoamLrSchedulerConfig::new(1e-2) + .with_warmup_steps(1000) + .with_model_size(config.transformer.d_model) + .init() + .unwrap(); + + // Initialize learner + let training = SupervisedTraining::new(artifact_dir, dataloader_train, dataloader_test) + .metric_train(CudaMetric::new()) + .metric_valid(CudaMetric::new()) + .metric_train(IterationSpeedMetric::new()) + .metric_train_numeric(LossMetric::new()) + .metric_valid_numeric(LossMetric::new()) + .metric_train_numeric(AccuracyMetric::new()) + .metric_valid_numeric(AccuracyMetric::new()) + .metric_train_numeric(LearningRateMetric::new()) + .with_default_checkpointers() + .with_training_strategy(strategy.into()) + .num_epochs(config.num_epochs) + .summary(); + + // Train the model + let result = training.launch(Learner::new(model, optim, lr_scheduler)); + + // Save the configuration and the trained model + config.save(format!("{artifact_dir}/config.json")).unwrap(); + result + .model + .into_record() + .save(format!("{artifact_dir}/model")) + .unwrap(); +} diff --git a/examples/text-generation/Cargo.toml b/examples/text-generation/Cargo.toml new file mode 100644 index 0000000..3862c36 --- /dev/null +++ b/examples/text-generation/Cargo.toml @@ -0,0 +1,29 @@ +[package] +authors = ["nathanielsimard "] +edition.workspace = true +license.workspace = true +name = "text-generation" +publish = false +version.workspace = true + +[lints] +workspace = true + +[features] +default = ["burn/dataset", "burn/sqlite-bundled"] +f16 = [] + +[dependencies] +# Burn +burn = { workspace = true, features = ["default", "train", "tch"] } + +# Tokenizer +tokenizers = { version = "0.22.2", default-features = false, features = [ + "onig", + "http", +] } + +# Utils +derive-new = { workspace = true } +log = { workspace = true } +serde = { workspace = true, features = ["std", "derive"] } diff --git a/examples/text-generation/README.md b/examples/text-generation/README.md new file mode 100644 index 0000000..b5c428f --- /dev/null +++ b/examples/text-generation/README.md @@ -0,0 +1,29 @@ +# Text Generation + +> **Note** +> This example makes use of the HuggingFace [`datasets`](https://huggingface.co/docs/datasets/index) +> library to download the datasets. Make sure you have [Python](https://www.python.org/downloads/) +> installed on your computer. + +The example can be run like so: + +## CUDA users + +```bash +git clone https://github.com/tracel-ai/burn.git +cd burn + +# Use the --release flag to really speed up training. +export TORCH_CUDA_VERSION=cu128 +cargo run --example text-generation --release +``` + +## Mac users + +```bash +git clone https://github.com/tracel-ai/burn.git +cd burn + +# Use the --release flag to really speed up training. +cargo run --example text-generation --release +``` diff --git a/examples/text-generation/examples/text-generation.rs b/examples/text-generation/examples/text-generation.rs new file mode 100644 index 0000000..ec8573b --- /dev/null +++ b/examples/text-generation/examples/text-generation.rs @@ -0,0 +1,36 @@ +use burn::{ + optim::decay::WeightDecayConfig, + tensor::{Device, DeviceConfig, DeviceIndex, Element}, +}; +use text_generation::{DbPediaDataset, training::ExperimentConfig}; + +#[cfg(feature = "f16")] +type Elem = burn::tensor::f16; +#[cfg(not(feature = "f16"))] +type Elem = f32; + +fn main() { + let config = ExperimentConfig::new( + burn::nn::transformer::TransformerEncoderConfig::new(384, 1536, 12, 6) + .with_norm_first(true), + burn::optim::AdamConfig::new().with_weight_decay(Some(WeightDecayConfig::new(1.0e-6))), + ); + + let mut device: Device = if cfg!(target_os = "macos") { + Device::libtorch_mps() + } else { + Device::libtorch_cuda(DeviceIndex::Default) + }; + + device + .configure(DeviceConfig::default().float_dtype(Elem::dtype())) + .unwrap(); + + text_generation::training::train::( + device, + DbPediaDataset::train(), + DbPediaDataset::test(), + config, + "/tmp/text-generation", + ); +} diff --git a/examples/text-generation/src/data/batcher.rs b/examples/text-generation/src/data/batcher.rs new file mode 100644 index 0000000..40dd4f8 --- /dev/null +++ b/examples/text-generation/src/data/batcher.rs @@ -0,0 +1,64 @@ +use super::{dataset::TextGenerationItem, tokenizer::Tokenizer}; +use burn::{data::dataloader::batcher::Batcher, nn::attention::generate_padding_mask, prelude::*}; +use std::sync::Arc; + +#[derive(Clone, new)] +pub struct TextGenerationBatcher { + tokenizer: Arc, + max_seq_length: usize, +} + +#[derive(Debug, Clone, new)] +pub struct TextGenerationBatch { + pub tokens: Tensor<2, Int>, + pub mask_pad: Tensor<2, Bool>, +} + +#[derive(Debug, Clone, new)] +pub struct TrainingTextGenerationBatch { + pub tokens_inputs: Tensor<2, Int>, + pub targets: Tensor<2, Int>, + pub mask_pad: Tensor<2, Bool>, +} + +impl Batcher for TextGenerationBatcher { + fn batch(&self, items: Vec, device: &Device) -> TextGenerationBatch { + let mut tokens_list = Vec::with_capacity(items.len()); + + for item in items { + tokens_list.push(self.tokenizer.encode(&item.text, true)); + } + + let mask = generate_padding_mask( + self.tokenizer.pad_token(), + tokens_list, + Some(self.max_seq_length), + device, + ); + + TextGenerationBatch { + tokens: mask.tensor, + mask_pad: mask.mask, + } + } +} + +impl Batcher for TextGenerationBatcher { + fn batch( + &self, + items: Vec, + device: &Device, + ) -> TrainingTextGenerationBatch { + let item: TextGenerationBatch = self.batch(items, device); + let [batch_size, seq_length] = item.tokens.dims(); + + let inputs = item + .tokens + .clone() + .slice([0..batch_size, 0..seq_length - 1]); + let targets = item.tokens.slice([0..batch_size, 1..seq_length]); + let mask_pad = item.mask_pad.slice([0..batch_size, 0..seq_length - 1]); + + TrainingTextGenerationBatch::new(inputs, targets, mask_pad) + } +} diff --git a/examples/text-generation/src/data/dataset.rs b/examples/text-generation/src/data/dataset.rs new file mode 100644 index 0000000..13e6b66 --- /dev/null +++ b/examples/text-generation/src/data/dataset.rs @@ -0,0 +1,44 @@ +use burn::data::dataset::{Dataset, SqliteDataset, source::huggingface::HuggingfaceDatasetLoader}; + +#[derive(new, Clone, Debug)] +pub struct TextGenerationItem { + pub text: String, +} + +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +pub struct DbPediaItem { + pub content: String, +} + +pub struct DbPediaDataset { + dataset: SqliteDataset, +} + +impl Dataset for DbPediaDataset { + fn get(&self, index: usize) -> Option { + self.dataset + .get(index) + .map(|item| TextGenerationItem::new(item.content)) + } + + fn len(&self) -> usize { + self.dataset.len() + } +} + +impl DbPediaDataset { + pub fn train() -> Self { + Self::new("train") + } + + pub fn test() -> Self { + Self::new("test") + } + pub fn new(split: &str) -> Self { + let dataset: SqliteDataset = + HuggingfaceDatasetLoader::new("fancyzhx/dbpedia_14") + .dataset(split) + .unwrap(); + Self { dataset } + } +} diff --git a/examples/text-generation/src/data/mod.rs b/examples/text-generation/src/data/mod.rs new file mode 100644 index 0000000..8f6faeb --- /dev/null +++ b/examples/text-generation/src/data/mod.rs @@ -0,0 +1,7 @@ +mod batcher; +mod dataset; +mod tokenizer; + +pub use batcher::*; +pub use dataset::*; +pub use tokenizer::*; diff --git a/examples/text-generation/src/data/tokenizer.rs b/examples/text-generation/src/data/tokenizer.rs new file mode 100644 index 0000000..53b294b --- /dev/null +++ b/examples/text-generation/src/data/tokenizer.rs @@ -0,0 +1,94 @@ +#[allow(dead_code)] +pub trait Tokenizer: Send + Sync { + fn encode(&self, value: &str, special_tokens: bool) -> Vec; + fn decode(&self, tokens: &[usize]) -> String; + fn vocab_size(&self) -> usize; + fn pad_token(&self) -> usize; + fn start_token(&self) -> usize; + fn end_token(&self) -> usize; + fn pad_token_value(&self) -> String { + self.decode(&[self.pad_token()]) + } + fn start_token_value(&self) -> String { + self.decode(&[self.start_token()]) + } + fn end_token_value(&self) -> String { + self.decode(&[self.end_token()]) + } +} + +pub struct Gpt2Tokenizer { + tokenizer: tokenizers::Tokenizer, +} + +impl Default for Gpt2Tokenizer { + fn default() -> Self { + let mut tokenizer = tokenizers::Tokenizer::from_pretrained("gpt2", None).unwrap(); + tokenizer.add_special_tokens(&[ + tokenizers::AddedToken::from("[START]", true), + tokenizers::AddedToken::from("[END]", true), + tokenizers::AddedToken::from("[PAD]", true), + ]); + + Self { tokenizer } + } +} + +impl Tokenizer for Gpt2Tokenizer { + fn encode(&self, value: &str, special_tokens: bool) -> Vec { + let text = match special_tokens { + true => "[START]".to_owned() + value + "[END]", + false => value.to_string(), + }; + let tokens = self.tokenizer.encode(text, true).unwrap(); + tokens.get_ids().iter().map(|t| *t as usize).collect() + } + + fn decode(&self, tokens: &[usize]) -> String { + let tokens = tokens.iter().map(|t| *t as u32).collect::>(); + self.tokenizer.decode(&tokens, false).unwrap() + } + + fn vocab_size(&self) -> usize { + self.tokenizer.get_vocab_size(true) + } + + fn pad_token(&self) -> usize { + self.tokenizer.token_to_id("[PAD]").unwrap() as usize + } + + fn start_token(&self) -> usize { + self.tokenizer.token_to_id("[START]").unwrap() as usize + } + + fn end_token(&self) -> usize { + self.tokenizer.token_to_id("[END]").unwrap() as usize + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_encode_decode() { + let tokenizer = Gpt2Tokenizer::default(); + let text = "A sentence"; + + let tokens = tokenizer.encode(text, false); + let decoded = tokenizer.decode(&tokens); + + assert_eq!(decoded, text); + } + + #[test] + fn test_add_start_end_token() { + let tokenizer = Gpt2Tokenizer::default(); + let text = "A sentence"; + + let tokens_without = tokenizer.encode(text, false); + let tokens_with = tokenizer.encode(text, true); + + assert_eq!(tokens_with.len() - 2, tokens_without.len()); + } +} diff --git a/examples/text-generation/src/lib.rs b/examples/text-generation/src/lib.rs new file mode 100644 index 0000000..881bb96 --- /dev/null +++ b/examples/text-generation/src/lib.rs @@ -0,0 +1,8 @@ +#[macro_use] +extern crate derive_new; + +mod data; +mod model; + +pub mod training; +pub use data::DbPediaDataset; diff --git a/examples/text-generation/src/model.rs b/examples/text-generation/src/model.rs new file mode 100644 index 0000000..cef53ac --- /dev/null +++ b/examples/text-generation/src/model.rs @@ -0,0 +1,112 @@ +use crate::data::TrainingTextGenerationBatch; +use burn::{ + nn::{ + Embedding, EmbeddingConfig, Linear, LinearConfig, + attention::generate_autoregressive_mask, + loss::CrossEntropyLossConfig, + transformer::{TransformerEncoder, TransformerEncoderConfig, TransformerEncoderInput}, + }, + prelude::*, + train::{ClassificationOutput, InferenceStep, TrainOutput, TrainStep}, +}; + +#[derive(Config, Debug)] +pub struct TextGenerationModelConfig { + transformer: TransformerEncoderConfig, + vocab_size: usize, + pad_token: usize, + max_seq_length: usize, +} + +#[derive(Module, Debug)] +pub struct TextGenerationModel { + transformer: TransformerEncoder, + embedding_token: Embedding, + embedding_pos: Embedding, + output: Linear, + vocab_size: usize, + pad_token: usize, + max_seq_length: usize, +} + +impl TextGenerationModelConfig { + pub fn init(&self, device: &Device) -> TextGenerationModel { + let output = LinearConfig::new(self.transformer.d_model, self.vocab_size).init(device); + let transformer = self.transformer.init(device); + let embedding_token = + EmbeddingConfig::new(self.vocab_size, self.transformer.d_model).init(device); + let embedding_pos = + EmbeddingConfig::new(self.max_seq_length, self.transformer.d_model).init(device); + + TextGenerationModel { + transformer, + embedding_token, + embedding_pos, + output, + vocab_size: self.vocab_size, + pad_token: self.pad_token, + max_seq_length: self.max_seq_length, + } + } +} +impl TextGenerationModel { + pub fn forward_training(&self, item: TrainingTextGenerationBatch) -> ClassificationOutput { + let [batch_size, seq_length] = item.tokens_inputs.dims(); + let device = &self.devices()[0]; + + let inputs = item.tokens_inputs.to_device(device); + let targets = item.targets.to_device(device); + let mask_pad = item.mask_pad.to_device(device); + + let index_positions = Tensor::arange(0..seq_length as i64, device) + .reshape([1, seq_length]) + .repeat_dim(0, batch_size); + + let embedding_positions = self.embedding_pos.forward(index_positions); + let embedding_tokens = self.embedding_token.forward(inputs); + let embedding = (embedding_positions + embedding_tokens) / 2; + + let mask_attn = generate_autoregressive_mask(batch_size, seq_length, device); + let encoded = self.transformer.forward( + TransformerEncoderInput::new(embedding) + .mask_pad(mask_pad) + .mask_attn(mask_attn), + ); + + let output = self.output.forward(encoded); + let output_flatten = output.reshape([batch_size * seq_length, self.vocab_size]); + let targets_flatten = targets.reshape([batch_size * seq_length]); + + let loss = CrossEntropyLossConfig::new() + .with_pad_tokens(Some(vec![self.pad_token])) + .init(&output_flatten.device()); + let loss = loss.forward(output_flatten.clone(), targets_flatten.clone()); + + ClassificationOutput { + loss, + output: output_flatten, + targets: targets_flatten, + } + } +} + +impl TrainStep for TextGenerationModel { + type Input = TrainingTextGenerationBatch; + type Output = ClassificationOutput; + + fn step(&self, item: TrainingTextGenerationBatch) -> TrainOutput { + let item = self.forward_training(item); + let grads = item.loss.backward(); + + TrainOutput::new(self, grads, item) + } +} + +impl InferenceStep for TextGenerationModel { + type Input = TrainingTextGenerationBatch; + type Output = ClassificationOutput; + + fn step(&self, item: TrainingTextGenerationBatch) -> ClassificationOutput { + self.forward_training(item) + } +} diff --git a/examples/text-generation/src/training.rs b/examples/text-generation/src/training.rs new file mode 100644 index 0000000..7e6b432 --- /dev/null +++ b/examples/text-generation/src/training.rs @@ -0,0 +1,94 @@ +use crate::{ + data::{Gpt2Tokenizer, TextGenerationBatcher, TextGenerationItem, Tokenizer}, + model::TextGenerationModelConfig, +}; +use burn::{ + data::{ + dataloader::DataLoaderBuilder, + dataset::{Dataset, transform::SamplerDataset}, + }, + lr_scheduler::noam::NoamLrSchedulerConfig, + nn::transformer::TransformerEncoderConfig, + optim::AdamConfig, + prelude::*, + train::{ + Learner, SupervisedTraining, + metric::{AccuracyMetric, CudaMetric, LearningRateMetric, LossMetric, PerplexityMetric}, + }, +}; +use std::sync::Arc; + +#[derive(Config, Debug)] +pub struct ExperimentConfig { + transformer: TransformerEncoderConfig, + optimizer: AdamConfig, + #[config(default = 512)] + max_seq_length: usize, + #[config(default = 6)] + batch_size: usize, + #[config(default = 50)] + num_epochs: usize, +} + +pub fn train + 'static>( + device: Device, + dataset_train: D, + dataset_test: D, + config: ExperimentConfig, + artifact_dir: &str, +) { + let tokenizer = Arc::new(Gpt2Tokenizer::default()); + let batcher = TextGenerationBatcher::new(tokenizer.clone(), config.max_seq_length); + let device = device.autodiff(); + + let model = TextGenerationModelConfig::new( + config.transformer.clone(), + tokenizer.vocab_size(), + tokenizer.pad_token(), + config.max_seq_length, + ) + .init(&device); + + let dataloader_train = DataLoaderBuilder::new(batcher.clone()) + .batch_size(config.batch_size) + .num_workers(4) + .build(SamplerDataset::new(dataset_train, 10_000)); + + let dataloader_test = DataLoaderBuilder::new(batcher) + .batch_size(config.batch_size) + .num_workers(4) + .build(SamplerDataset::new(dataset_test, 1000)); + + let accum = 6; // Effective batch size = 6 * 6 = 32. + let optim = config.optimizer.init(); + let lr_scheduler = NoamLrSchedulerConfig::new(0.01 / accum as f64) + .with_warmup_steps(6000) + .with_model_size(config.transformer.d_model) + .init() + .unwrap(); + + let training = SupervisedTraining::new(artifact_dir, dataloader_train, dataloader_test) + .metric_train(CudaMetric::new()) + .metric_valid(CudaMetric::new()) + .metric_train_numeric(AccuracyMetric::new().with_pad_token(tokenizer.pad_token())) + .metric_valid_numeric(AccuracyMetric::new().with_pad_token(tokenizer.pad_token())) + .metric_train_numeric(PerplexityMetric::new().with_pad_token(tokenizer.pad_token())) + .metric_valid_numeric(PerplexityMetric::new().with_pad_token(tokenizer.pad_token())) + .metric_train(LossMetric::new()) + .metric_valid(LossMetric::new()) + .metric_train_numeric(LearningRateMetric::new()) + .with_default_checkpointers() + .grads_accumulation(accum) + .num_epochs(config.num_epochs) + .summary(); + + let result = training.launch(Learner::new(model, optim, lr_scheduler)); + + config.save(format!("{artifact_dir}/config.json")).unwrap(); + + result + .model + .into_record() + .save(format!("{artifact_dir}/model")) + .unwrap(); +} diff --git a/examples/wgan/Cargo.toml b/examples/wgan/Cargo.toml new file mode 100644 index 0000000..0212ef0 --- /dev/null +++ b/examples/wgan/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "wgan" +version = "0.22.0-pre.1" +edition.workspace = true + +[lints] +workspace = true + +[features] +flex = ["burn/flex"] +tch-cpu = ["burn/tch"] +tch-gpu = ["burn/tch"] +wgpu = ["burn/wgpu"] +cuda = ["burn/cuda"] + +[dependencies] +burn = { workspace = true, features=["default", "train", "vision"] } +image = { workspace = true } diff --git a/examples/wgan/README.md b/examples/wgan/README.md new file mode 100644 index 0000000..a85f494 --- /dev/null +++ b/examples/wgan/README.md @@ -0,0 +1,36 @@ +# Wasserstein Generative Adversarial Network + +A burn implementation of an example WGAN model to generate MNIST digits inspired by +[the PyTorch implementation](https://bytepawn.com/training-a-pytorch-wasserstain-mnist-gan-on-google-colab.html). +Please note that better performance maybe gained by adopting a convolution layer in +[some other models](https://github.com/Lornatang/WassersteinGAN-PyTorch). + +## Usage + +## Training + +```sh +# Cuda backend +cargo run --example wgan-mnist --release --features cuda + +# Wgpu backend +cargo run --example wgan-mnist --release --features wgpu + +# Tch GPU backend +export TORCH_CUDA_VERSION=cu128 # Set the cuda version +cargo run --example wgan-mnist --release --features tch-gpu + +# Tch CPU backend +cargo run --example wgan-mnist --release --features tch-cpu + +# Flex backend (CPU) +cargo run --example wgan-mnist --release --features flex # f32 +``` + +### Generating + +To generate a sample of images, you can use `wgan-generate`. The same feature flags are used to select a backend. + +```sh +cargo run --example wgan-generate --release --features cuda +``` diff --git a/examples/wgan/examples/wgan-generate.rs b/examples/wgan/examples/wgan-generate.rs new file mode 100644 index 0000000..6b0809b --- /dev/null +++ b/examples/wgan/examples/wgan-generate.rs @@ -0,0 +1,68 @@ +use burn::tensor::Device; + +pub fn launch(device: Device) { + wgan::infer::generate("/tmp/wgan-mnist", device); +} + +#[cfg(feature = "flex")] +mod flex { + use burn::tensor::Device; + + pub fn run() { + crate::launch(Device::flex()); + } +} + +#[cfg(feature = "tch-gpu")] +mod tch_gpu { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + #[cfg(not(target_os = "macos"))] + let device = Device::libtorch_cuda(DeviceIndex::Default); + #[cfg(target_os = "macos")] + let device = Device::libtorch_mps(); + + crate::launch(device); + } +} + +#[cfg(feature = "tch-cpu")] +mod tch_cpu { + use burn::tensor::Device; + + pub fn run() { + crate::launch(Device::libtorch()); + } +} + +#[cfg(feature = "wgpu")] +mod wgpu { + use burn::tensor::{Device, DeviceKind}; + + pub fn run() { + crate::launch(Device::wgpu(DeviceKind::DefaultDevice)); + } +} + +#[cfg(feature = "cuda")] +mod cuda { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + crate::launch(Device::cuda(DeviceIndex::Default)); + } +} + +fn main() { + #[cfg(feature = "flex")] + flex::run(); + #[cfg(feature = "tch-gpu")] + tch_gpu::run(); + #[cfg(feature = "tch-cpu")] + tch_cpu::run(); + #[cfg(feature = "wgpu")] + wgpu::run(); + #[cfg(feature = "cuda")] + cuda::run(); +} diff --git a/examples/wgan/examples/wgan-mnist.rs b/examples/wgan/examples/wgan-mnist.rs new file mode 100644 index 0000000..bf57d06 --- /dev/null +++ b/examples/wgan/examples/wgan-mnist.rs @@ -0,0 +1,80 @@ +use burn::{optim::RmsPropConfig, tensor::Device}; + +use wgan::{model::ModelConfig, training::TrainingConfig}; + +pub fn launch(device: Device) { + let config = TrainingConfig::new( + ModelConfig::new(), + RmsPropConfig::new() + .with_alpha(0.99) + .with_momentum(0.0) + .with_epsilon(0.00000008) + .with_weight_decay(None) + .with_centered(false), + ); + + wgan::training::train("/tmp/wgan-mnist", config, device); +} + +#[cfg(feature = "flex")] +mod flex { + use burn::tensor::Device; + + pub fn run() { + crate::launch(Device::flex()); + } +} + +#[cfg(feature = "tch-gpu")] +mod tch_gpu { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + #[cfg(not(target_os = "macos"))] + let device = Device::libtorch_cuda(DeviceIndex::Default); + #[cfg(target_os = "macos")] + let device = Device::libtorch_mps(); + + crate::launch(device); + } +} + +#[cfg(feature = "tch-cpu")] +mod tch_cpu { + use burn::tensor::Device; + + pub fn run() { + crate::launch(Device::libtorch()); + } +} + +#[cfg(feature = "wgpu")] +mod wgpu { + use burn::tensor::{Device, DeviceKind}; + + pub fn run() { + crate::launch(Device::wgpu(DeviceKind::DefaultDevice)); + } +} + +#[cfg(feature = "cuda")] +mod cuda { + use burn::tensor::{Device, DeviceIndex}; + + pub fn run() { + crate::launch(Device::cuda(DeviceIndex::Default)); + } +} + +fn main() { + #[cfg(feature = "flex")] + flex::run(); + #[cfg(feature = "tch-gpu")] + tch_gpu::run(); + #[cfg(feature = "tch-cpu")] + tch_cpu::run(); + #[cfg(feature = "wgpu")] + wgpu::run(); + #[cfg(feature = "cuda")] + cuda::run(); +} diff --git a/examples/wgan/src/dataset.rs b/examples/wgan/src/dataset.rs new file mode 100644 index 0000000..aa97f23 --- /dev/null +++ b/examples/wgan/src/dataset.rs @@ -0,0 +1,36 @@ +use burn::{ + data::{dataloader::batcher::Batcher, dataset::vision::MnistItem}, + prelude::*, +}; + +#[derive(Clone, Debug, Default)] +pub struct MnistBatcher {} + +#[derive(Clone, Debug)] +pub struct MnistBatch { + pub images: Tensor<4>, + pub targets: Tensor<1, Int>, +} + +impl Batcher for MnistBatcher { + fn batch(&self, items: Vec, device: &Device) -> MnistBatch { + let images = items + .iter() + .map(|item| TensorData::from(item.image)) + .map(|data| Tensor::<2>::from_data(data, device)) + .map(|tensor| tensor.reshape([1, 28, 28])) + // Set std=0.5 and mean=0.5 to keep consistent with pytorch WGAN example + .map(|tensor| ((tensor / 255) - 0.5) / 0.5) + .collect(); + + let targets = items + .iter() + .map(|item| Tensor::<1, Int>::from_data(TensorData::from([item.label as i64]), device)) + .collect(); + + let images = Tensor::stack(images, 0); + let targets = Tensor::cat(targets, 0); + + MnistBatch { images, targets } + } +} diff --git a/examples/wgan/src/infer.rs b/examples/wgan/src/infer.rs new file mode 100644 index 0000000..75eaf8d --- /dev/null +++ b/examples/wgan/src/infer.rs @@ -0,0 +1,36 @@ +use crate::training::{TrainingConfig, save_image}; +use burn::{prelude::*, store::ModuleRecord, tensor::Distribution}; + +pub fn generate(artifact_dir: &str, device: Device) { + // Loading model + let config = TrainingConfig::load(format!("{artifact_dir}/config.json")) + .expect("Config should exist for the model; run train first"); + let record = ModuleRecord::load(format!("{artifact_dir}/generator")) + .expect("Trained model should exist; run train first"); + let (mut generator, _) = config.model.init(&device); + generator = generator.load_record(record); + + // Get a batch of noise + let noise = Tensor::<2>::random( + [config.batch_size, config.model.latent_dim], + Distribution::Normal(0.0, 1.0), + &device, + ); + let fake_images = generator.forward(noise); // [batch_size, channesl*height*width] + let fake_images = fake_images.reshape([ + config.batch_size, + config.model.channels, + config.model.image_size, + config.model.image_size, + ]); + // [B, C, H, W] to [B, H, C, W] to [B, H, W, C] + let fake_images = fake_images.swap_dims(2, 1).swap_dims(3, 2).slice(0..25); + // Normalize the images. The Rgb32 images should be in range 0.0-1.0 + let fake_images = (fake_images.clone() - fake_images.clone().min().reshape([1, 1, 1, 1])) + / (fake_images.clone().max().reshape([1, 1, 1, 1]) + - fake_images.clone().min().reshape([1, 1, 1, 1])); + // Add 0.5 after unnormalizing to [0, 255] to round to the nearest integer, refer to pytorch save_image source + let fake_images = (fake_images + 0.5 / 255.0).clamp(0.0, 1.0); + // Save images in artifact directory + save_image(fake_images, 5, format!("{artifact_dir}/fake_image.png")).unwrap(); +} diff --git a/examples/wgan/src/lib.rs b/examples/wgan/src/lib.rs new file mode 100644 index 0000000..021f622 --- /dev/null +++ b/examples/wgan/src/lib.rs @@ -0,0 +1,4 @@ +pub mod dataset; +pub mod infer; +pub mod model; +pub mod training; diff --git a/examples/wgan/src/model.rs b/examples/wgan/src/model.rs new file mode 100644 index 0000000..521a559 --- /dev/null +++ b/examples/wgan/src/model.rs @@ -0,0 +1,156 @@ +use burn::{ + module::{Module, ModuleMapper, Param}, + prelude::*, +}; + +/// Layer block of generator model +#[derive(Module, Debug)] +pub struct LayerBlock { + fc: nn::Linear, + bn: nn::BatchNorm, + leakyrelu: nn::LeakyRelu, +} + +impl LayerBlock { + pub fn new(input: usize, output: usize, device: &Device) -> Self { + let fc = nn::LinearConfig::new(input, output) + .with_bias(true) + .init(device); + let bn: nn::BatchNorm = nn::BatchNormConfig::new(output) + .with_epsilon(0.8) + .init(device); + let leakyrelu = nn::LeakyReluConfig::new().with_negative_slope(0.2).init(); + + Self { fc, bn, leakyrelu } + } + + pub fn forward(&self, input: Tensor<2>) -> Tensor<2> { + let output = self.fc.forward(input); // output: [Batch, x] + let output = self.bn.forward(output); // output: [Batch, x] + + self.leakyrelu.forward(output) // output: [Batch, x] + } +} + +/// Generator model +#[derive(Module, Debug)] +pub struct Generator { + layer1: LayerBlock, + layer2: LayerBlock, + layer3: LayerBlock, + layer4: LayerBlock, + fc: nn::Linear, + tanh: nn::Tanh, +} + +impl Generator { + /// Applies the forward pass on the input tensor by specified order + pub fn forward(&self, noise: Tensor<2>) -> Tensor<2> { + let output = self.layer1.forward(noise); + let output = self.layer2.forward(output); + let output = self.layer3.forward(output); + let output = self.layer4.forward(output); + let output = self.fc.forward(output); + + self.tanh.forward(output) // [batch_size, channels*height*width] + } +} + +/// Discriminator model +#[derive(Module, Debug)] +pub struct Discriminator { + fc1: nn::Linear, + leakyrelu1: nn::LeakyRelu, + fc2: nn::Linear, + leakyrelu2: nn::LeakyRelu, + fc3: nn::Linear, +} + +impl Discriminator { + /// Applies the forward pass on the input tensor by specified order. + /// The input image shape is [batch, channels, height, width] + pub fn forward(&self, images: Tensor<4>) -> Tensor<2> { + // Full connection for each batch + let output = images.flatten(1, 3); // output: [batch, channels*height*width] + let output = self.fc1.forward(output); // output: [batch, 512] + let output = self.leakyrelu1.forward(output); // output: [batch, 512] + let output = self.fc2.forward(output); // output: [batch, 256] + let output = self.leakyrelu2.forward(output); // output: [batch, 256] + + self.fc3.forward(output) // output: [batch, 1] + } +} + +// Use model config to construct a generative and adversarial model +#[derive(Config, Debug)] +pub struct ModelConfig { + /// Dimensionality of the latent space + #[config(default = 100)] + pub latent_dim: usize, + #[config(default = 28)] + pub image_size: usize, + #[config(default = 1)] + pub channels: usize, +} + +impl ModelConfig { + /// Initialize the generator and discriminator models based on the config. + pub fn init(&self, device: &Device) -> (Generator, Discriminator) { + // Construct the initialized generator + let layer1 = LayerBlock::new(self.latent_dim, 128, device); + let layer2 = LayerBlock::new(128, 256, device); + let layer3 = LayerBlock::new(256, 512, device); + let layer4 = LayerBlock::new(512, 1024, device); + let fc = nn::LinearConfig::new(1024, self.channels * self.image_size * self.image_size) + .with_bias(true) + .init(device); + + let generator = Generator { + layer1, + layer2, + layer3, + layer4, + fc, + tanh: nn::Tanh::new(), + }; + + // Construct the initialized discriminator + let fc1 = nn::LinearConfig::new(self.channels * self.image_size * self.image_size, 512) + .init(device); + let leakyrelu1 = nn::LeakyReluConfig::new().with_negative_slope(0.2).init(); + let fc2 = nn::LinearConfig::new(512, 256).init(device); + let leakyrelu2 = nn::LeakyReluConfig::new().with_negative_slope(0.2).init(); + let fc3 = nn::LinearConfig::new(256, 1).init(device); + + let discriminator = Discriminator { + fc1, + leakyrelu1, + fc2, + leakyrelu2, + fc3, + }; + + (generator, discriminator) + } +} + +/// Clip module mapper to clip all module parameters between a range of values +#[derive(Module, Debug)] +pub struct Clip { + pub min: f32, + pub max: f32, +} + +impl ModuleMapper for Clip { + fn map_float(&mut self, param: Param>) -> Param> { + let (id, tensor, mapper) = param.consume(); + let is_require_grad = tensor.is_require_grad(); + + let mut tensor = Tensor::from_inner(tensor.inner().clamp(self.min, self.max)); + + if is_require_grad { + tensor = tensor.require_grad(); + } + Param::from_mapped_value(id, tensor, mapper) + } +} diff --git a/examples/wgan/src/training.rs b/examples/wgan/src/training.rs new file mode 100644 index 0000000..233ad16 --- /dev/null +++ b/examples/wgan/src/training.rs @@ -0,0 +1,206 @@ +use crate::dataset::MnistBatcher; +use crate::model::{Clip, ModelConfig}; +use burn::optim::{GradientsParams, RmsPropConfig}; +use burn::{ + data::{dataloader::DataLoaderBuilder, dataset::vision::MnistDataset}, + prelude::*, + tensor::Distribution, +}; +use image::{Rgb32FImage, RgbImage, buffer::ConvertBuffer, error::ImageResult}; +use std::path::{Path, PathBuf}; + +#[derive(Config, Debug)] +pub struct TrainingConfig { + pub model: ModelConfig, + pub optimizer: RmsPropConfig, + + #[config(default = 200)] + pub num_epochs: usize, + #[config(default = 512)] + pub batch_size: usize, + #[config(default = 8)] + pub num_workers: usize, + #[config(default = 5)] + pub seed: u64, + #[config(default = 3e-4)] + pub lr: f64, + + /// Number of training steps for discriminator before generator is trained per iteration + #[config(default = 5)] + pub num_critic: usize, + /// Lower and upper clip value for disc. weights + #[config(default = 0.01)] + pub clip_value: f32, + /// Save a sample of images every `sample_interval` epochs + #[config(default = 10)] + pub sample_interval: usize, +} + +// Create the directory to save the model and model config +fn create_artifact_dir(artifact_dir: &str) { + std::fs::remove_file(PathBuf::from(artifact_dir).join("experiment.log")).ok(); + std::fs::create_dir_all(artifact_dir).ok(); +} + +/// Save the generated images +// The images format is [B, H, W, C] +pub fn save_image>(images: Tensor<4>, nrow: u32, path: Q) -> ImageResult<()> { + let ncol = (images.dims()[0] as f32 / nrow as f32).ceil() as u32; + + let width = images.dims()[2] as u32; + let height = images.dims()[1] as u32; + + // Supports both 1 and 3 channels image + let channels = match images.dims()[3] { + 1 => 3, + 3 => 1, + _ => panic!("Wrong channels number"), + }; + + let mut imgbuf = RgbImage::new(nrow * width, ncol * height); + // Write images into a nrow*ncol grid layout + for row in 0..nrow { + for col in 0..ncol { + let image: Tensor<3> = images + .clone() + .slice((row * nrow + col) as usize..(row * nrow + col + 1) as usize) + .squeeze_dim(0); + // The Rgb32 should be in range 0.0-1.0 + let image = image.into_data().iter::().collect::>(); + // Supports both 1 and 3 channels image + let image = image + .into_iter() + .flat_map(|n| std::iter::repeat_n(n, channels)) + .collect(); + + let image = Rgb32FImage::from_vec(width, height, image).unwrap(); + let image: RgbImage = image.convert(); + for (x, y, pixel) in image.enumerate_pixels() { + imgbuf.put_pixel(row * width + x, col * height + y, *pixel); + } + } + } + imgbuf.save(path) +} + +pub fn train(artifact_dir: &str, config: TrainingConfig, device: Device) { + create_artifact_dir(artifact_dir); + + // Create the Clip module mapper + let mut clip = Clip { + min: -config.clip_value, + max: config.clip_value, + }; + + // Save training config + config + .save(format!("{artifact_dir}/config.json")) + .expect("Config should be saved successfully"); + + device.seed(config.seed); + let device = device.autodiff(); + + // Create the model and optimizer + let (mut generator, mut discriminator) = config.model.init(&device); + let mut optimizer_g = config.optimizer.init(); + let mut optimizer_d = config.optimizer.init(); + + // Create the dataset batcher + let batcher_train = MnistBatcher::default(); + + // Create the dataloaders + let dataloader_train = DataLoaderBuilder::new(batcher_train) + .batch_size(config.batch_size) + .shuffle(config.seed) + .num_workers(config.num_workers) + .build(MnistDataset::train()); + + // Iterate over our training for X epochs + for epoch in 0..config.num_epochs { + // Implement our training loop + for (iteration, batch) in dataloader_train.iter().enumerate() { + // Generate a batch of fake images from noise (standarded normal distribution) + let noise = Tensor::<2>::random( + [config.batch_size, config.model.latent_dim], + Distribution::Normal(0.0, 1.0), + &device, + ); + // datach: do not update generator, only discriminator is updated + let fake_images = generator.forward(noise.clone()).detach(); // [batch_size, channels*height*width] + let fake_images = fake_images.reshape([ + config.batch_size, + config.model.channels, + config.model.image_size, + config.model.image_size, + ]); + // Adversarial loss + let loss_d = -discriminator.forward(batch.images).mean() + + discriminator.forward(fake_images.clone()).mean(); + + // Gradients for the current backward pass + let grads = loss_d.backward(); + // Gradients linked to each parameter of the discriminator + let grads = GradientsParams::from_grads(grads, &discriminator); + // Update the discriminator using the optimizer + discriminator = optimizer_d.step(config.lr.into(), discriminator, grads); + // Clip parameters (weights) of discriminator + discriminator = discriminator.map(&mut clip); + + // Train the generator every num_critic iterations + if iteration % config.num_critic == 0 { + // Generate a batch of images again without detaching + let critic_fake_images = generator.forward(noise.clone()); + let critic_fake_images = critic_fake_images.reshape([ + config.batch_size, + config.model.channels, + config.model.image_size, + config.model.image_size, + ]); + // Adversarial loss. Minimize it to make the fake images as truth + let loss_g = -discriminator.forward(critic_fake_images).mean(); + + let grads = loss_g.backward(); + let grads = GradientsParams::from_grads(grads, &generator); + generator = optimizer_g.step(config.lr.into(), generator, grads); + + // Print the progression + let batch_num = (dataloader_train.num_items() as f32 / config.batch_size as f32) + .ceil() as usize; + println!( + "[Epoch {}/{}] [Batch {}/{}] [D loss: {}] [G loss: {}]", + epoch + 1, + config.num_epochs, + iteration, + batch_num, + loss_d.into_scalar::(), + loss_g.into_scalar::() + ); + } + // If at save interval => save the first 25 generated images + if epoch % config.sample_interval == 0 && iteration == 0 { + // [B, C, H, W] to [B, H, C, W] to [B, H, W, C] + let fake_images = fake_images.swap_dims(2, 1).swap_dims(3, 2).slice(0..25); + // Normalize the images. The Rgb32 images should be in range 0.0-1.0 + let fake_images = (fake_images.clone() + - fake_images.clone().min().reshape([1, 1, 1, 1])) + / (fake_images.clone().max().reshape([1, 1, 1, 1]) + - fake_images.clone().min().reshape([1, 1, 1, 1])); + // Add 0.5/255.0 to the images, refer to pytorch save_image source + let fake_images = (fake_images + 0.5 / 255.0).clamp(0.0, 1.0); + // Save images in artifact directory + let path = format!("{artifact_dir}/image-{epoch}.png"); + save_image(fake_images, 5, path).unwrap(); + } + } + } + + // Save the trained models + generator + .into_record() + .save(format!("{artifact_dir}/generator")) + .expect("Generator should be saved successfully"); + discriminator + .into_record() + .save(format!("{artifact_dir}/discriminator")) + .expect("Discriminator should be saved successfully"); +} diff --git a/rustfmt.toml b/rustfmt.toml new file mode 100644 index 0000000..8e8f959 --- /dev/null +++ b/rustfmt.toml @@ -0,0 +1,5 @@ +max_width = 100 + +# uncomment and run `cargo +nightly fmt --all` to find and fix lines that are too long (and therefore break autoformatting) +# error_on_line_overflow = true +# format_strings = true diff --git a/xtask/Cargo.toml b/xtask/Cargo.toml new file mode 100644 index 0000000..6be979e --- /dev/null +++ b/xtask/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "xtask" +version = "4.10.0" +edition.workspace = true +license = "MIT OR Apache-2.0" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[lints] +workspace = true + +[dependencies] +log = { workspace = true } +strum = { workspace = true } +cargo_metadata = "0.23.1" + +tracel-xtask = "=4.15.1" +# ### For local development. ### +# tracel-xtask = { path = "../xtask/crates/tracel-xtask", default-features = false } + +[dev-dependencies] +rstest = { workspace = true } diff --git a/xtask/src/commands/books.rs b/xtask/src/commands/books.rs new file mode 100644 index 0000000..99e3dc7 --- /dev/null +++ b/xtask/src/commands/books.rs @@ -0,0 +1,109 @@ +use std::path::Path; + +use tracel_xtask::prelude::*; + +#[derive(clap::Args)] +pub struct BooksArgs { + #[command(subcommand)] + book: BookKind, +} + +#[derive(clap::Subcommand)] +pub(crate) enum BookKind { + /// Burn Book, a.k.a. the guide, made for the Burn users. + Burn(BookKindArgs), + /// Contributor book, made for people willing to get all the technical understanding and advice to contribute actively to the project. + Contributor(BookKindArgs), +} + +#[derive(clap::Args)] +pub(crate) struct BookKindArgs { + #[command(subcommand)] + command: BookSubCommand, +} + +#[derive(clap::Subcommand, strum::Display)] +pub(crate) enum BookSubCommand { + /// Build the book + Build, + /// Open the book on the specified port or random port and rebuild it automatically upon changes + Open(OpenArgs), +} + +#[derive(clap::Args)] +pub(crate) struct OpenArgs { + /// Specify the port to open the book on (defaults to a random port if not specified) + #[clap(long, default_value_t = random_port())] + port: u16, +} + +/// Book information +pub(crate) struct Book { + name: &'static str, + path: &'static Path, +} + +impl BooksArgs { + pub(crate) fn parse(&self) -> anyhow::Result<()> { + Book::run(&self.book) + } +} + +impl Book { + const BURN_BOOK_NAME: &'static str = "Burn Book"; + const BURN_BOOK_PATH: &'static str = "./burn-book"; + + const CONTRIBUTOR_BOOK_NAME: &'static str = "Contributor Book"; + const CONTRIBUTOR_BOOK_PATH: &'static str = "./contributor-book"; + + pub(crate) fn run(book_arg: &BookKind) -> anyhow::Result<()> { + let (book, command) = match book_arg { + BookKind::Burn(args) => ( + Self { + name: Self::BURN_BOOK_NAME, + path: Path::new(Self::BURN_BOOK_PATH), + }, + &args.command, + ), + BookKind::Contributor(args) => ( + Self { + name: Self::CONTRIBUTOR_BOOK_NAME, + path: Path::new(Self::CONTRIBUTOR_BOOK_PATH), + }, + &args.command, + ), + }; + book.execute(command) + } + + fn execute(&self, command: &BookSubCommand) -> anyhow::Result<()> { + ensure_cargo_crate_is_installed("mdbook", None, None, false)?; + group!("{}: {}", self.name, command); + match command { + BookSubCommand::Build => self.build(), + BookSubCommand::Open(args) => self.open(args), + }?; + endgroup!(); + Ok(()) + } + + fn build(&self) -> anyhow::Result<()> { + run_process( + "mdbook", + &["build"], + None, + Some(self.path), + "mdbook should build the book successfully", + ) + } + + fn open(&self, args: &OpenArgs) -> anyhow::Result<()> { + run_process( + "mdbook", + &["serve", "--open", "--port", &args.port.to_string()], + None, + Some(self.path), + "mdbook should open the book successfully", + ) + } +} diff --git a/xtask/src/commands/build.rs b/xtask/src/commands/build.rs new file mode 100644 index 0000000..7815bc5 --- /dev/null +++ b/xtask/src/commands/build.rs @@ -0,0 +1,130 @@ +use std::collections::HashMap; + +use tracel_xtask::prelude::{clap::ValueEnum, *}; + +use crate::{ARM_NO_ATOMIC_PTR_TARGET, ARM_TARGET, NO_STD_CRATES, WASM32_TARGET}; + +#[macros::extend_command_args(BuildCmdArgs, Target, None)] +pub struct BurnBuildCmdArgs { + /// Build in CI mode which excludes unsupported crates. + #[arg(long)] + pub ci: bool, +} + +pub(crate) fn handle_command( + mut args: BurnBuildCmdArgs, + env: Environment, + context: Context, +) -> anyhow::Result<()> { + match context { + Context::NoStd => { + [ + "Default", + WASM32_TARGET, + ARM_TARGET, + ARM_NO_ATOMIC_PTR_TARGET, + ] + .iter() + .try_for_each(|build_target| { + let mut build_args = vec!["--no-default-features"]; + let mut env_vars = HashMap::new(); + if *build_target != "Default" { + build_args.extend(vec!["--target", *build_target]); + } + + let mut crates = NO_STD_CRATES.to_vec(); + + if *build_target == ARM_NO_ATOMIC_PTR_TARGET { + // Only build a subset of crates which require `portable_atomic_unsafe_assume_single_core`. + // Other crates build with `burn-flex` (default_backend) which requires the `critical-section` + // feature (automatically enabled via `burn-dispatch` when `not(target_has_atomic = "ptr")`), + // which is mutually exclusive with `portable_atomic_unsafe_assume_single_core` cfg. + crates = vec!["burn-std", "burn-backend", "burn-ndarray"]; + env_vars.insert( + "RUSTFLAGS", + "--cfg portable_atomic_unsafe_assume_single_core", + ); + } + helpers::custom_crates_build( + crates, + build_args.clone(), + Some(env_vars.clone()), + None, + &format!("no-std with target {}", *build_target), + )?; + + // Second pass for `thumbv6m-none-eabi`: crates with `critical-section` feature + // enabled so `once_cell` (transitively pulled from `burn-flex` -> `gemm`) uses + // portable-atomic for CAS emulation. + if *build_target == ARM_NO_ATOMIC_PTR_TARGET { + crates = NO_STD_CRATES.to_vec(); + // Remove `burn-autodiff` from building with the + // target `thumbv6m-none-eabi` as it requires enabling the + // `arbitrary_self_types` feature for the + // `clone_if_require_grad` method of + // `burn-autodiff::graph::Node`. + crates.retain(|&v| { + v != "burn-autodiff" + && v != "burn-std" + && v != "burn-ndarray" + && v != "burn-backend" + }); + + helpers::custom_crates_build( + crates, + build_args, + None, + None, + &format!("no-std with target {} (critical-section)", *build_target), + )?; + } + + anyhow::Ok(()) + })?; + Ok(()) + } + Context::Std => { + if args.ci { + // Exclude crates that are not supported on CI + args.exclude.extend(vec![ + "burn-cuda".to_string(), + "burn-rocm".to_string(), + "burn-tch".to_string(), + ]); + if std::env::var("DISABLE_WGPU").is_ok() { + args.exclude.extend(vec!["burn-wgpu".to_string()]); + }; + } + // Build workspace + base_commands::build::handle_command(args.try_into().unwrap(), env, context)?; + // Specific additional commands to test specific features + // burn-dataset + helpers::custom_crates_build( + vec!["burn-dataset"], + vec!["--all-features"], + None, + None, + "std with all features", + )?; + Ok(()) + } + Context::All => Context::value_variants() + .iter() + .filter(|ctx| **ctx != Context::All) + .try_for_each(|ctx| { + handle_command( + BurnBuildCmdArgs { + target: args.target.clone(), + exclude: args.exclude.clone(), + only: args.only.clone(), + ci: args.ci, + release: args.release, + features: args.features.clone(), + no_default_features: args.no_default_features, + }, + env.clone(), + ctx.clone(), + ) + }), + } +} diff --git a/xtask/src/commands/doc.rs b/xtask/src/commands/doc.rs new file mode 100644 index 0000000..97eb72e --- /dev/null +++ b/xtask/src/commands/doc.rs @@ -0,0 +1,28 @@ +use tracel_xtask::prelude::*; + +pub(crate) fn handle_command( + mut args: DocCmdArgs, + env: Environment, + ctx: Context, +) -> anyhow::Result<()> { + if args.get_command() == DocSubCommand::Build { + args.exclude + .extend(vec!["burn-cuda".to_string(), "burn-rocm".to_string()]); + } + + // Execute documentation command on workspace + base_commands::doc::handle_command(args.clone(), env, ctx)?; + + // Specific additional commands to build other docs + if args.get_command() == DocSubCommand::Build { + // burn-dataset + helpers::custom_crates_doc_build( + vec!["burn-dataset"], + vec!["--all-features"], + None, + None, + "All features", + )?; + } + Ok(()) +} diff --git a/xtask/src/commands/mod.rs b/xtask/src/commands/mod.rs new file mode 100644 index 0000000..e4220a8 --- /dev/null +++ b/xtask/src/commands/mod.rs @@ -0,0 +1,6 @@ +pub(crate) mod books; +pub(crate) mod build; +pub(crate) mod doc; +pub(crate) mod remote; +pub(crate) mod test; +pub(crate) mod validate; diff --git a/xtask/src/commands/remote.rs b/xtask/src/commands/remote.rs new file mode 100644 index 0000000..455fd49 --- /dev/null +++ b/xtask/src/commands/remote.rs @@ -0,0 +1,188 @@ +//! End-to-end validation for the remote backend. +//! +//! Builds the `server` example with a chosen backend, spawns it as a child process, waits +//! for it to bind, then runs `burn-backend-tests --features remote` against it with +//! `BURN_DEVICE=remote`. The child server is killed on exit regardless of how the test run +//! finishes. + +use std::{ + net::{SocketAddr, TcpStream}, + path::PathBuf, + process::{Child, Command, Stdio}, + time::{Duration, Instant}, +}; + +use tracel_xtask::prelude::*; + +#[derive(clap::Args)] +pub struct RemoteCmdArgs { + /// Backend the remote server should use. + #[arg(long, value_enum, default_value_t = ServerBackend::Flex)] + pub backend: ServerBackend, + + /// Port the server should listen on (also used by the client). + #[arg(long, default_value_t = 7373)] + pub port: u16, + + /// Build the server (and run tests) in release mode. + #[arg(long)] + pub release: bool, + + /// Maximum seconds to wait for the server to bind before giving up. + #[arg(long, default_value_t = 30)] + pub startup_timeout: u64, + + /// Test filter passed through to `cargo test` (positional, after `--`). + #[arg(long)] + pub test: Option, +} + +#[derive(Debug, Clone, clap::ValueEnum, strum::Display)] +pub enum ServerBackend { + #[strum(to_string = "flex")] + Flex, + #[strum(to_string = "vulkan")] + Vulkan, + #[strum(to_string = "webgpu")] + Webgpu, + #[strum(to_string = "cuda")] + Cuda, +} + +pub fn handle_command( + args: RemoteCmdArgs, + _env: Environment, + _context: Context, +) -> anyhow::Result<()> { + let backend = args.backend.to_string(); + + info!( + "Building server example (backend = {backend}, release = {}) ...", + args.release + ); + build_server(&backend, args.release)?; + + let server_bin = server_binary_path(args.release); + info!( + "Spawning server at {} on port {} ...", + server_bin.display(), + args.port + ); + let _server = ServerGuard::spawn(&server_bin, args.port)?; + + wait_for_bind(args.port, args.startup_timeout)?; + info!("Server ready on 127.0.0.1:{}", args.port); + + run_remote_tests(args.port, args.release, args.test.as_deref())?; + + info!("Remote validation succeeded — tearing down the server."); + Ok(()) +} + +fn build_server(backend: &str, release: bool) -> anyhow::Result<()> { + // `--no-default-features` so we don't pull in the example's `webgpu` default on top of + // whatever backend the caller picked (the example's `cfg_if!` would then build twice as + // many backends for nothing). + let mut cargo_args = vec![ + "build", + "--example", + "server", + "--no-default-features", + "--features", + backend, + ]; + if release { + cargo_args.push("--release"); + } + run_process( + "cargo", + &cargo_args, + None, + None, + "Failed to build the `server` example", + ) +} + +fn server_binary_path(release: bool) -> PathBuf { + let profile = if release { "release" } else { "debug" }; + PathBuf::from("target") + .join(profile) + .join("examples") + .join("server") +} + +fn run_remote_tests(port: u16, release: bool, filter: Option<&str>) -> anyhow::Result<()> { + let address = format!("ws://127.0.0.1:{port}"); + + // Bypass `run_process`'s `&str` borrowing by owning the formatted strings here. + let mut owned_args: Vec = vec![ + "test".into(), + "-p".into(), + "burn-backend-tests".into(), + "--features".into(), + "remote".into(), + ]; + if release { + owned_args.push("--release".into()); + } + // Tests share a single client connection — keep them single-threaded so a hang in + // one test doesn't deadlock the whole suite. + owned_args.push("--".into()); + owned_args.push("--test-threads=1".into()); + if let Some(f) = filter { + owned_args.push(f.to_string()); + } + let args: Vec<&str> = owned_args.iter().map(String::as_str).collect(); + + let envs = std::collections::HashMap::from([ + ("BURN_DEVICE", "remote"), + ("BURN_REMOTE_ADDRESS", address.as_str()), + ]); + + run_process( + "cargo", + &args, + Some(envs), + None, + "burn-backend-tests failed against the remote backend", + ) +} + +fn wait_for_bind(port: u16, timeout_secs: u64) -> anyhow::Result<()> { + let addr: SocketAddr = format!("127.0.0.1:{port}").parse()?; + let deadline = Instant::now() + Duration::from_secs(timeout_secs); + while Instant::now() < deadline { + if TcpStream::connect_timeout(&addr, Duration::from_millis(200)).is_ok() { + return Ok(()); + } + std::thread::sleep(Duration::from_millis(100)); + } + anyhow::bail!( + "Server failed to start on port {port} within {timeout_secs}s — check the server logs above." + ) +} + +/// Owns the server `Child`. On drop, kills the process so a panicking / failing +/// test run still leaves no orphaned server. +struct ServerGuard(Child); + +impl ServerGuard { + fn spawn(bin: &std::path::Path, port: u16) -> anyhow::Result { + let child = Command::new(bin) + .env("REMOTE_BACKEND_PORT", port.to_string()) + // Inherit stdio so the server's logs show up in the xtask output. This is the + // first thing you want when something goes wrong (e.g. "device not found"). + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()) + .spawn() + .map_err(|e| anyhow::anyhow!("Failed to spawn server binary {}: {e}", bin.display()))?; + Ok(Self(child)) + } +} + +impl Drop for ServerGuard { + fn drop(&mut self) { + let _ = self.0.kill(); + let _ = self.0.wait(); + } +} diff --git a/xtask/src/commands/test.rs b/xtask/src/commands/test.rs new file mode 100644 index 0000000..d1d2c21 --- /dev/null +++ b/xtask/src/commands/test.rs @@ -0,0 +1,480 @@ +use tracel_xtask::{ + prelude::{clap::ValueEnum, *}, + utils::{ + process::{ExitSignal, ProcessExitError}, + workspace::WorkspaceMember, + }, +}; + +use crate::NO_STD_CRATES; + +#[cfg(unix)] +use std::os::unix::process::ExitStatusExt; + +#[macros::extend_command_args(TestCmdArgs, Target, TestSubCommand)] +pub struct BurnTestCmdArgs { + /// Test in CI mode which excludes unsupported crates. + #[arg(long)] + pub ci: CiTestType, +} + +// `cargo check` for examples +impl std::convert::TryInto for BurnTestCmdArgs { + type Error = anyhow::Error; + fn try_into(self) -> Result { + Ok(CompileCmdArgs { + target: self.target, + exclude: self.exclude, + only: self.only, + }) + } +} + +#[allow(clippy::enum_variant_names)] +#[derive(Debug, Clone, ValueEnum, PartialEq)] +pub enum CiTestType { + // Github runner shards + Backends, + Crates, + Examples, + // Other runners + GithubRunner, + GithubMacRunner, + GcpCudaRunner, + GcpVulkanRunner, + GcpWgpuRunner, +} + +#[derive(strum::Display)] +enum TestBackend { + #[strum(to_string = "cuda")] + Cuda, + #[strum(to_string = "metal")] + Metal, + #[strum(to_string = "vulkan")] + Vulkan, + #[strum(to_string = "wgpu")] + Wgpu, + #[allow(unused)] + #[strum(to_string = "rocm")] + Rocm, + #[strum(to_string = "flex")] + Flex, + #[strum(to_string = "ndarray")] + Ndarray, +} + +fn set_burn_device(device: &str) { + // SAFETY: This is called in a single-threaded context within the xtask before spawning child processes. + unsafe { + std::env::set_var("BURN_DEVICE", device); + } +} + +fn handle_backend_tests( + args: TestCmdArgs, + backend: TestBackend, + context: Context, +) -> anyhow::Result<()> { + let backend_name = backend.to_string(); + set_burn_device(&backend_name); // default device + + let mut test_args = vec!["--no-default-features", "--features", &backend_name]; + if !matches!(context, Context::NoStd) { + test_args.extend(["--features", "std"]) + } + + if matches!(backend, TestBackend::Cuda) { + // Collective (all-reduce) tests require a CUDA build with NCCL, which the CI runner + // provides. Kept behind its own feature so plain `--features cuda` still works without it. + test_args.extend(["--features", "distributed"]); + } + + if !matches!(backend, TestBackend::Ndarray | TestBackend::Flex) { + // Fusion enabled tests first + let mut fusion_args = test_args.clone(); + fusion_args.extend(["--features", "fusion"]); + + helpers::custom_crates_tests( + vec!["burn-backend-tests"], + handle_test_args(&fusion_args, args.release), + None, + None, + "fusion backend tests", + )?; + // base_commands::test::handle_command(fusion_args, env.clone(), context.clone())?; + } + + // base_commands::test::handle_command(args, env, context) + helpers::custom_crates_tests( + vec!["burn-backend-tests"], + handle_test_args(&test_args, args.release), + None, + None, + "backend tests", + ) +} + +fn handle_wgpu_test(member: &str, args: &TestCmdArgs) -> anyhow::Result<()> { + #[cfg(unix)] + let filter_err = |e: &&ProcessExitError| { + e.status.signal() == Some(11) || matches!(e.signal, Some(ExitSignal { code: 11, .. })) + }; + #[cfg(not(unix))] + let filter_err = |e: &&ProcessExitError| matches!(e.signal, Some(ExitSignal { code: 11, .. })); + + let workspace_member = WorkspaceMember { + name: member.into(), + path: "".into(), // unused + }; + + if let Err(err) = base_commands::test::run_unit_test(&workspace_member, args) { + let should_ignore = err + .downcast_ref::() + .filter(filter_err) + // Failed to execute unit test for '{member}' + .map(|e| e.message.contains(member)) + .unwrap_or(false); + + if should_ignore { + // Ignore intermittent successful failures + // https://github.com/gfx-rs/wgpu/issues/2949 + // https://github.com/KhronosGroup/Vulkan-ValidationLayers/issues/4391 + eprintln!("⚠️ Ignored SIGSEGV in wgpu test"); + } else { + return Err(err); + } + } + Ok(()) +} + +const EXCLUDE_CRATES: &[&str] = &[ + "burn-cpu", + "burn-cuda", + "burn-rocm", + // "burn-router" uses "burn-wgpu" for the tests. + "burn-router", + "burn-tch", + "burn-wgpu", + // Requires wgpu runtime + "burn-cubecl-fusion", + // Backends are tested individually + "burn-backend-tests", + "burn-ndarray", + "burn-flex", +]; + +fn enumerate_examples() -> anyhow::Result> { + let metadata = cargo_metadata::MetadataCommand::new().exec()?; + + let workspace_root = metadata.workspace_root.as_std_path(); + let examples_dir = workspace_root.join("examples"); + + Ok(metadata + .workspace_packages() + .into_iter() + .filter(|package| { + // Check if the package's Cargo.toml lives inside the examples/ folder + package.manifest_path.starts_with(&examples_dir) + }) + .map(|package| package.name.to_string()) + .collect()) +} + +pub(crate) fn handle_command( + mut args: BurnTestCmdArgs, + env: Environment, + context: Context, +) -> anyhow::Result<()> { + match context { + Context::NoStd => { + // burn-flex's unit tests use `std::f32::consts` and bare `vec!` + // macros directly in test modules, so they only compile under std. + // The build step (`xtask build --no-std`) still validates that + // the crate itself compiles as no_std via `cargo build`, which + // does not pull in test modules. + let no_std_test_crates: Vec<&str> = NO_STD_CRATES + .iter() + .copied() + .filter(|&c| c != "burn-flex") + .collect(); + ["Default"].iter().try_for_each(|test_target| { + let mut test_args = vec!["--no-default-features"]; + if *test_target != "Default" { + test_args.extend(vec!["--target", *test_target]); + } + helpers::custom_crates_tests( + no_std_test_crates.clone(), + handle_test_args(&test_args, args.release), + None, + None, + "no-std", + ) + })?; + handle_backend_tests( + args.clone().try_into().unwrap(), + TestBackend::Ndarray, + context, + )?; + + Ok(()) + } + Context::Std => { + // 1) Tests with default features + // ------------------------------ + match args.ci { + CiTestType::Backends | CiTestType::GithubRunner => { + // Backend ops + handle_backend_tests( + args.clone().try_into().unwrap(), + TestBackend::Ndarray, + context.clone(), + )?; + + handle_backend_tests( + args.clone().try_into().unwrap(), + TestBackend::Flex, + context.clone(), + )?; + + // Backend crates + args.target = Target::AllPackages; + args.only + .extend(["burn-ndarray".to_string(), "burn-flex".to_string()]); + base_commands::test::handle_command( + args.clone().try_into().unwrap(), + env.clone(), + context, + )?; + } + CiTestType::Crates => { + // Default `Target::Workspace` + // Exclude crates that are not supported on CI + args.exclude + .extend(EXCLUDE_CRATES.iter().map(|&s| s.to_string())); + // Exclude examples + // workspace feature unification will cause binary bloat with examples default features + args.exclude.extend(enumerate_examples()?); + + // Burn remote tests don't work on windows for now + #[cfg(target_os = "windows")] + { + args.exclude.extend(vec!["burn-remote".to_string()]); + }; + + set_burn_device("flex"); // default device for base tests + base_commands::test::handle_command( + args.clone().try_into().unwrap(), + env.clone(), + context.clone(), + )?; + } + CiTestType::Examples => { + // NOTE: for the examples we simply run `cargo checks` (no tests, faster validation) + // TODO: switch to `cargo xtask build` or `check` eventually instead of including this in the tests + args.target = Target::AllPackages; + args.only.extend(enumerate_examples()?); + base_commands::compile::handle_command( + args.clone().try_into().unwrap(), + env.clone(), + context.clone(), + )?; + } + CiTestType::GithubMacRunner => { + handle_backend_tests( + args.clone().try_into().unwrap(), + TestBackend::Metal, + context.clone(), + )?; + + args.target = Target::AllPackages; + args.only.push("burn-wgpu".to_string()); + args.features + .get_or_insert_with(Vec::new) + .push("metal".to_string()); + + base_commands::test::handle_command( + args.clone().try_into().unwrap(), + env, + context, + )?; + } + CiTestType::GcpCudaRunner => { + handle_backend_tests( + args.clone().try_into().unwrap(), + TestBackend::Cuda, + context, + )?; + } + CiTestType::GcpVulkanRunner => { + handle_backend_tests( + args.clone().try_into().unwrap(), + TestBackend::Vulkan, + context, + )?; + + args.target = Target::AllPackages; + let mut args_vulkan = args.clone(); + args_vulkan + .features + .get_or_insert_with(Vec::new) + .push("vulkan".to_string()); + + let args_vulkan = args_vulkan.try_into().unwrap(); + handle_wgpu_test("burn-wgpu", &args_vulkan)?; + handle_wgpu_test("burn-core", &args_vulkan)?; + handle_wgpu_test("burn-vision", &args_vulkan)?; + + // Enable burn-core/vulkan + args.features + .get_or_insert_with(Vec::new) + .push("burn-core/vulkan".to_string()); + let args_vulkan = args.clone().try_into().unwrap(); + handle_wgpu_test("burn-optim", &args_vulkan)?; + handle_wgpu_test("burn-nn", &args_vulkan)?; + } + CiTestType::GcpWgpuRunner => { + handle_backend_tests( + args.clone().try_into().unwrap(), + TestBackend::Wgpu, + context, + )?; + args.target = Target::AllPackages; + handle_wgpu_test("burn-cubecl-fusion", &args.clone().try_into().unwrap())?; + + let mut args_wgpu = args.clone(); + args_wgpu + .features + .get_or_insert_with(Vec::new) + .push("webgpu".to_string()); + + let args_wgpu = args.clone().try_into().unwrap(); + handle_wgpu_test("burn-wgpu", &args_wgpu)?; + handle_wgpu_test("burn-core", &args_wgpu)?; + handle_wgpu_test("burn-vision", &args_wgpu)?; + + // Enable burn-core/webgpu + args.features + .get_or_insert_with(Vec::new) + .push("burn-core/webgpu".to_string()); + let args_wgpu = args.clone().try_into().unwrap(); + handle_wgpu_test("burn-optim", &args_wgpu)?; + handle_wgpu_test("burn-nn", &args_wgpu)?; + } + } + + // 2) Specific additional commands to test specific features + // --------------------------------------------------------- + match args.ci { + CiTestType::Backends | CiTestType::GithubRunner => (), + CiTestType::Examples => (), + CiTestType::Crates => { + // burn-dataset + helpers::custom_crates_tests( + vec!["burn-dataset"], + handle_test_args(&["--all-features"], args.release), + None, + None, + "std all features", + )?; + + // burn-core + set_burn_device("tch"); // test-tch + helpers::custom_crates_tests( + vec!["burn-core"], + handle_test_args(&["--features", "tch"], args.release), + None, + None, + "std with features: tch", + )?; + + // burn-nn (pretrained and local tests) + // If the "CI" environment variable is missing, we are running locally. + // if std::env::var("CI").is_err() { + // nn_features.push_str(",test-local"); + // } + // burn-vision + set_burn_device("flex"); + helpers::custom_crates_tests( + vec!["burn-vision"], + handle_test_args(&["--features", "flex", "loss"], args.release), + None, + None, + "std cpu (flex)", + )?; + + // burn-train vision (LPIPS, DISTS metrics) + helpers::custom_crates_tests( + vec!["burn-train"], + handle_test_args(&["--features", "vision"], args.release), + None, + None, + "std vision", + )?; + } + CiTestType::GcpCudaRunner => (), + CiTestType::GcpVulkanRunner | CiTestType::GcpWgpuRunner => (), // handled in tests above + CiTestType::GithubMacRunner => { + // burn-ndarray + helpers::custom_crates_tests( + vec!["burn-ndarray"], + handle_test_args(&["--features", "blas-accelerate"], args.release), + None, + None, + "std blas-accelerate", + )?; + + set_burn_device("metal"); + helpers::custom_crates_tests( + vec!["burn-core"], + handle_test_args(&["--features", "metal"], args.release), + None, + None, + "std metal", + )?; + helpers::custom_crates_tests( + vec!["burn-vision"], + handle_test_args(&["--features", "metal"], args.release), + None, + None, + "std metal", + )?; + } + } + Ok(()) + } + Context::All => Context::value_variants() + .iter() + .filter(|ctx| **ctx != Context::All) + .try_for_each(|ctx| { + handle_command( + BurnTestCmdArgs { + command: args.command.clone(), + target: args.target.clone(), + exclude: args.exclude.clone(), + only: args.only.clone(), + threads: args.threads, + jobs: args.jobs, + ci: args.ci.clone(), + features: args.features.clone(), + no_default_features: args.no_default_features, + release: args.release, + test: args.test.clone(), + force: args.force, + no_capture: args.no_capture, + miri: args.miri, + }, + env.clone(), + ctx.clone(), + ) + }), + } +} + +fn handle_test_args<'a>(args: &'a [&'a str], release: bool) -> Vec<&'a str> { + let mut args = args.to_vec(); + if release { + args.push("--release"); + } + args +} diff --git a/xtask/src/commands/validate.rs b/xtask/src/commands/validate.rs new file mode 100644 index 0000000..33a3d14 --- /dev/null +++ b/xtask/src/commands/validate.rs @@ -0,0 +1,154 @@ +use tracel_xtask::prelude::*; + +use crate::commands::{ + build::BurnBuildCmdArgs, + test::{BurnTestCmdArgs, CiTestType}, +}; + +pub fn handle_command( + args: &ValidateCmdArgs, + env: Environment, + context: Context, +) -> anyhow::Result<()> { + let target = Target::Workspace; + let exclude = vec![]; + let only = vec![]; + + if context == Context::NoStd || context == Context::All { + // ================= + // no-std validation + // ================= + info!("Run validation for no-std execution environment..."); + + #[cfg(target_os = "linux")] + { + // build + super::build::handle_command( + BurnBuildCmdArgs { + target: target.clone(), + exclude: exclude.clone(), + only: only.clone(), + ci: true, + release: args.release, + features: args.features.clone(), + no_default_features: args.no_default_features, + }, + env.clone(), + Context::NoStd, + )?; + + // tests + super::test::handle_command( + BurnTestCmdArgs { + target: target.clone(), + exclude: exclude.clone(), + only: only.clone(), + threads: None, + jobs: None, + command: Some(TestSubCommand::All), + ci: CiTestType::GithubRunner, + features: None, + no_default_features: false, + force: false, + no_capture: false, + release: args.release, + test: None, + miri: false, + }, + env.clone(), + Context::NoStd, + )?; + } + } + + if context == Context::Std || context == Context::All { + // ============== + // std validation + // ============== + info!("Run validation for std execution environment..."); + + // checks + [ + CheckSubCommand::Audit, + CheckSubCommand::Format, + CheckSubCommand::Lint, + CheckSubCommand::Typos, + ] + .iter() + .try_for_each(|c| { + base_commands::check::handle_command( + CheckCmdArgs { + target: target.clone(), + exclude: exclude.clone(), + only: only.clone(), + command: Some(c.clone()), + ignore_audit: args.ignore_audit, + features: args.features.clone(), + no_default_features: args.no_default_features, + ignore_typos: args.ignore_typos, + }, + env.clone(), + context.clone(), + ) + })?; + + // build + super::build::handle_command( + BurnBuildCmdArgs { + target: target.clone(), + exclude: exclude.clone(), + only: only.clone(), + ci: true, + release: args.release, + features: args.features.clone(), + no_default_features: args.no_default_features, + }, + env.clone(), + Context::Std, + )?; + + // tests + super::test::handle_command( + BurnTestCmdArgs { + target: target.clone(), + exclude: exclude.clone(), + only: only.clone(), + threads: None, + jobs: None, + command: Some(TestSubCommand::All), + // NOTE: this will only run `CiTestType::Backends` (burn-backend-tests) + // for local sanity checks this is a quick approach to validating after no-std tests + cargo build + ci: CiTestType::GithubRunner, + features: None, + no_default_features: false, + release: args.release, + test: None, + force: false, + no_capture: false, + miri: false, + }, + env.clone(), + Context::Std, + )?; + + // documentation + [DocSubCommand::Build, DocSubCommand::Tests] + .iter() + .try_for_each(|c| { + super::doc::handle_command( + DocCmdArgs { + target: target.clone(), + exclude: exclude.clone(), + only: only.clone(), + command: Some(c.clone()), + features: args.features.clone(), + no_default_features: args.no_default_features, + }, + env.clone(), + context.clone(), + ) + })?; + } + + Ok(()) +} diff --git a/xtask/src/main.rs b/xtask/src/main.rs new file mode 100644 index 0000000..5f04b37 --- /dev/null +++ b/xtask/src/main.rs @@ -0,0 +1,86 @@ +mod commands; + +#[macro_use] +extern crate log; + +use std::time::Instant; +use tracel_xtask::prelude::*; + +// no-std +const WASM32_TARGET: &str = "wasm32-unknown-unknown"; +const ARM_TARGET: &str = "thumbv7m-none-eabi"; +const ARM_NO_ATOMIC_PTR_TARGET: &str = "thumbv6m-none-eabi"; +const NO_STD_CRATES: &[&str] = &[ + "burn", + "burn-autodiff", + "burn-core", + "burn-std", + "burn-backend", + "burn-tensor", + "burn-ndarray", + "burn-no-std-tests", +]; + +#[macros::base_commands( + Bump, + Check, + Compile, + Coverage, + Doc, + Dependencies, + Fix, + Publish, + Validate, + Vulnerabilities +)] +pub enum Command { + /// Run commands to manage Burn Books. + Books(commands::books::BooksArgs), + /// Build Burn in different modes. + Build(commands::build::BurnBuildCmdArgs), + /// Validate the remote backend end-to-end: spin up the `server` example, point + /// `burn-backend-tests` at it via `BURN_DEVICE=remote`, tear it down on exit. + Remote(commands::remote::RemoteCmdArgs), + /// Test Burn. + Test(commands::test::BurnTestCmdArgs), +} + +fn main() -> anyhow::Result<()> { + let start = Instant::now(); + let (args, environment) = init_xtask::(parse_args::()?)?; + + if args.context == Context::NoStd { + // Install additional targets for no-std execution environments + rustup_add_target(WASM32_TARGET)?; + rustup_add_target(ARM_TARGET)?; + rustup_add_target(ARM_NO_ATOMIC_PTR_TARGET)?; + } + + match args.command { + Command::Books(cmd_args) => cmd_args.parse(), + Command::Build(cmd_args) => { + commands::build::handle_command(cmd_args, environment, args.context) + } + Command::Remote(cmd_args) => { + commands::remote::handle_command(cmd_args, environment, args.context) + } + Command::Doc(cmd_args) => { + commands::doc::handle_command(cmd_args, environment, args.context) + } + Command::Test(cmd_args) => { + commands::test::handle_command(cmd_args, environment, args.context) + } + Command::Validate(cmd_args) => { + commands::validate::handle_command(&cmd_args, environment, args.context) + } + _ => dispatch_base_commands(args, environment), + }?; + + let duration = start.elapsed(); + info!( + "\x1B[32;1mTime elapsed for the current execution: {}\x1B[0m", + format_duration(&duration) + ); + + Ok(()) +}